Emit each loop once in the body a parameter version leaves (#2587)

packages/compiler/src/backend/llvm/emit-function.ts

@@ -148,7 +148,7 @@ export function emitFunction(emitter: LlEmitter, fn: IrFunction): string {
   emitter.integerTempsEnabled = emitter.debug === null && !fn.async && !fn.generator;
   beginFunctionLoops(emitter, fn, numericFn);
   emitter.speculativeLoopsEnabled = LOOPS.of(emitter).countedLoopsEnabled;
-  emitter.loopCopiesEnabled = true;
+  emitter.loopCopiesEnabled = fn.plainLoops !== true;
   emitter.speculativeFunction = fn.name;
   emitter.speculation = null;
   emitter.speculationResume = null;

packages/compiler/src/backend/llvm/emit-statements.ts

@@ -1147,6 +1147,14 @@ export function emitStmtBody(emitter: LlEmitter, s: IrStmt): void {
         (countedLoop?.guarded || pairedLoop?.guarded) &&
         !LOOPS.of(emitter).provenCountedGuards.has(s)
       ) {
+        // A function that emits each loop once takes only the copy without
+        // the integer form, which needs no safety test.
+        if (LOOPS.of(emitter).plainLoops) {
+          LOOPS.of(emitter).countedLoopsEnabled = false;
+          emitSlowCopy(emitter, s);
+          LOOPS.of(emitter).countedLoopsEnabled = true;
+          break;
+        }
         const safe = countedLoop
           ? emitCountedLoopSafety(emitter, countedLoop)
           : emitPairedLoopSafety(emitter, pairedLoop!);
@@ -1318,6 +1326,7 @@ export function emitStmtBody(emitter: LlEmitter, s: IrStmt): void {
       // longer runs take the loop as it is emitted otherwise.
       const shortRuns: { stride: number; left: () => string } | null =
         !innermost ||
+        LOOPS.of(emitter).plainLoops ||
         bounded ||
         capturedLets.length > 0 ||
         !emitter.workerLoopPolls(s) ||

packages/compiler/src/backend/llvm/loop-state.ts

@@ -27,10 +27,11 @@ export function beginFunctionLoops(
   loops.unaliasedArrays = findUnaliasedLocalArrays(fn);
   loops.countedLoopsEnabled = emitter.debug === null && !fn.async && !fn.generator;
   loops.byteWindowsEnabled = loops.countedLoopsEnabled;
-  loops.bytesLoopVersionsEnabled = loops.countedLoopsEnabled;
+  loops.plainLoops = fn.plainLoops === true;
+  loops.bytesLoopVersionsEnabled = loops.countedLoopsEnabled && !loops.plainLoops;
   loops.bytesLoopEntries = null;
   loops.bytesLoopFunction = fn;
-  loops.denseLoopsEnabled = loops.countedLoopsEnabled;
+  loops.denseLoopsEnabled = loops.countedLoopsEnabled && !loops.plainLoops;
   loops.denseLoop = null;
   loops.denseLoopEntries =
     loops.denseLoopsEnabled && numericFn.locals.some((l) => l.type.kind === "array")
@@ -65,8 +66,11 @@ export interface LoopState {
   byteWindowsEnabled: boolean;
   byteWindowEntries: ReturnType<typeof findInitializedByteLoopBindings>;
   /** Typed-array loop versioning (bytes-loop-versions.ts) for the current
-   * function: off in debug and suspending functions and inside a copy. */
+   * function: off in debug and suspending functions, in a function that
+   * emits each loop once and inside a copy. */
   bytesLoopVersionsEnabled: boolean;
+  /** The current function emits each loop once (IrFunction.plainLoops). */
+  plainLoops: boolean;
   /** Loops being versioned around the statement being emitted: a loop is
    * never versioned inside its own copies, and versions nest boundedly. */
   readonly bytesLoopVersioning: Set<IrStmt>;
@@ -89,6 +93,7 @@ export const LOOPS = defineModuleState<LoopState>(() => ({
   byteWindowsEnabled: true,
   byteWindowEntries: new Map(),
   bytesLoopVersionsEnabled: false,
+  plainLoops: false,
   bytesLoopVersioning: new Set<IrStmt>(),
   bytesLoopEntries: null,
   bytesLoopFunction: null,

packages/compiler/src/ir/codec.generated.ts

@@ -5,7 +5,7 @@ import type { IrReader, IrWriter } from "./codec.js";
 
 /** A digest of the coders below: an encoding written with other coders is
  * refused instead of misread, also after a change that keeps IR_VERSION. */
-export const IR_CODEC_LAYOUT = "b964edeec6535396";
+export const IR_CODEC_LAYOUT = "c8e9ba1a465defc2";
 
 /** Writes `mod` field by field (codec.ts describes the bytes). */
 export function encodeIrModule(w: IrWriter, mod: IrModule): void {
@@ -25,7 +25,7 @@ export function decodeIrFunction(r: IrReader): IrFunction {
 }
 
 const KEYS0: readonly string[] = ["irVersion","sourceFile","functions","classes","globals","records","unions","entry","workers","declBoundary","inferredFactFiles","declOrigins","lifetimeRegions","foreignElementKeys","threadSlotKeys","freezableArrays","sloppyDynWrites","untypedIntegrity","commonJsEntry","nodeIsland","ffiImports","lib","contracts","int32ArrayCandidates","wasmModules","wasmUntranslated","wasmCustomSections"];
-const KEYS1: readonly string[] = ["name","sourceName","ownsPrototype","params","returnType","locals","captures","classCaptures","body","loc","jsEntry","tier","speculativeDispatch","fastRead","async","generator","asyncCacheGlobal","asyncCycleCacheGlobal"];
+const KEYS1: readonly string[] = ["name","sourceName","ownsPrototype","params","returnType","locals","captures","classCaptures","body","loc","jsEntry","tier","speculativeDispatch","fastRead","plainLoops","async","generator","asyncCacheGlobal","asyncCycleCacheGlobal"];
 const KEYS2: readonly string[] = ["libnode","nodeVersion","modules"];
 const KEYS3: readonly string[] = ["profileName","prefix","initSymbol","sinkRegisterSymbol","collectSymbol","resultResetSymbol","threadInstances","callbackRegisterSymbol","callbacks","exports","trapOverlays","identity"];
 const KEYS4: readonly string[] = ["int32Fields","int32Params","int32Returns","int32Locals","int32Elements","noThrow","acyclic","immortalGlobals","programLifetime","inlineArrays"];
@@ -444,6 +444,11 @@ function encIrFunction(w: IrWriter, v: IrFunction): void {
     w.byte(1);
     w.str(v.fastRead);
   }
+  if (v.plainLoops === undefined) w.byte(0);
+  else {
+    w.byte(1);
+    w.bool(v.plainLoops);
+  }
   if (v.async === undefined) w.byte(0);
   else {
     w.byte(1);
@@ -480,11 +485,12 @@ function decIrFunction(r: IrReader): IrFunction {
   const f11 = r.tag(2, "IrFunction.tier") === 0 ? undefined : (r.str() as NonNullable<IrFunction["tier"]>);
   const f12 = r.tag(2, "IrFunction.speculativeDispatch") === 0 ? undefined : r.bool();
   const f13 = r.tag(2, "IrFunction.fastRead") === 0 ? undefined : (r.str() as NonNullable<IrFunction["fastRead"]>);
-  const f14 = r.tag(2, "IrFunction.async") === 0 ? undefined : (r.bool() as NonNullable<IrFunction["async"]>);
-  const f15 = r.tag(2, "IrFunction.generator") === 0 ? undefined : decT21(r);
-  const f16 = r.tag(2, "IrFunction.asyncCacheGlobal") === 0 ? undefined : r.str();
-  const f17 = r.tag(2, "IrFunction.asyncCycleCacheGlobal") === 0 ? undefined : r.str();
-  const o: IrFunction = { name: f0, ...(f1 !== undefined ? { sourceName: f1 } : {}), ...(f2 !== undefined ? { ownsPrototype: f2 } : {}), params: f3, returnType: f4, locals: f5, ...(f6 !== undefined ? { captures: f6 } : {}), ...(f7 !== undefined ? { classCaptures: f7 } : {}), body: f8, loc: f9, ...(f10 !== undefined ? { jsEntry: f10 } : {}), ...(f11 !== undefined ? { tier: f11 } : {}), ...(f12 !== undefined ? { speculativeDispatch: f12 } : {}), ...(f13 !== undefined ? { fastRead: f13 } : {}), ...(f14 !== undefined ? { async: f14 } : {}), ...(f15 !== undefined ? { generator: f15 } : {}), ...(f16 !== undefined ? { asyncCacheGlobal: f16 } : {}), ...(f17 !== undefined ? { asyncCycleCacheGlobal: f17 } : {}) };
+  const f14 = r.tag(2, "IrFunction.plainLoops") === 0 ? undefined : (r.bool() as NonNullable<IrFunction["plainLoops"]>);
+  const f15 = r.tag(2, "IrFunction.async") === 0 ? undefined : (r.bool() as NonNullable<IrFunction["async"]>);
+  const f16 = r.tag(2, "IrFunction.generator") === 0 ? undefined : decT21(r);
+  const f17 = r.tag(2, "IrFunction.asyncCacheGlobal") === 0 ? undefined : r.str();
+  const f18 = r.tag(2, "IrFunction.asyncCycleCacheGlobal") === 0 ? undefined : r.str();
+  const o: IrFunction = { name: f0, ...(f1 !== undefined ? { sourceName: f1 } : {}), ...(f2 !== undefined ? { ownsPrototype: f2 } : {}), params: f3, returnType: f4, locals: f5, ...(f6 !== undefined ? { captures: f6 } : {}), ...(f7 !== undefined ? { classCaptures: f7 } : {}), body: f8, loc: f9, ...(f10 !== undefined ? { jsEntry: f10 } : {}), ...(f11 !== undefined ? { tier: f11 } : {}), ...(f12 !== undefined ? { speculativeDispatch: f12 } : {}), ...(f13 !== undefined ? { fastRead: f13 } : {}), ...(f14 !== undefined ? { plainLoops: f14 } : {}), ...(f15 !== undefined ? { async: f15 } : {}), ...(f16 !== undefined ? { generator: f16 } : {}), ...(f17 !== undefined ? { asyncCacheGlobal: f17 } : {}), ...(f18 !== undefined ? { asyncCycleCacheGlobal: f18 } : {}) };
   return o;
 }
 

packages/compiler/src/ir/ir.ts

@@ -2004,6 +2004,11 @@ export interface IrFunction {
    * `index`: a computed read `(value, key)` whose callers pass a number
    * converted to a string as the key. */
   fastRead?: "index";
+  /** The body of a function whose parameter version takes its calls with
+   * int32 arguments (parameter-versions.ts). It runs for other arguments
+   * only, so the backend emits each of its loops once, as in the copy a
+   * failed version guard runs. */
+  plainLoops?: true;
   /** Async: the body runs on a fiber; `returnType` is the INNER type T (a
    * `return v` fulfills with v) while ordinary call sites receive
    * Promise<T>. With `generator` present, call sites receive the lazy

packages/compiler/src/ir/parameter-versions.ts

@@ -171,7 +171,8 @@ function loopBounds(fn: IrFunction, functions: ReadonlyMap<string, IrFunction>):
  * -0 (`1 / p` is not -Infinity). The copy then receives `p | 0`, which is
  * the same double, so it computes what the body would; its only caller
  * proves the parameter an int32, so it gets the i32 entry. Other values
- * run the body as before.
+ * run the body as before, whose loops the backend then emits once
+ * (`plainLoops`).
  *
  * Only plain synchronous functions qualify: not class members, no
  * captures, not the module entry, not dynamic-tier, at most
@@ -301,6 +302,7 @@ export function versionLoopBoundParameters(
         : [{ kind: "return", value: call, loc }];
     replaced.set(fn, {
       ...fn,
+      plainLoops: true,
       body: [{ kind: "if", cond: test!, then: run, else_: null, loc }, ...fn.body],
     });
   }

packages/compiler/test/counted-loop-emission.test.ts

@@ -299,9 +299,10 @@ test("worker programs poll no loop whose proven ranges bound its runs inside ano
     expect(rows).not.toMatch(/^loop\.poll\d+:$/m);
     expect(polls(rows)).toBeGreaterThan(0);
     // With a limit known only at run time the inner loops keep their polls.
-    // The innermost one is strip-mined, and a block that ends at the limit
+    // `open`'s int32 parameter version takes calls with int32 arguments. Its
+    // innermost loop is strip-mined, and a block that ends at the limit
     // leaves it without polling: only a block that ends short of it polls.
-    const open = body("open");
+    const open = body("open_int32Params");
     expect(polls(open)).toBeGreaterThan(polls(rows));
     expect(open).toMatch(/^loop\.poll\d+:$/m);
     const ends = [...open.matchAll(/^loop\.end\d+:$[\s\S]*?^\S/gm)].map((match) => match[0]);
@@ -310,6 +311,12 @@ test("worker programs poll no loop whose proven ranges bound its runs inside ano
     // A run of up to 16384 iterations is one block, polled as it leaves when
     // it took more than 64.
     expect(open).toMatch(/^loop\.poll\.run\d+:$/m);
+    // The body the version leaves runs for other arguments and emits each
+    // loop once: its loops poll without strip-mining.
+    const generic = body("open");
+    expect(polls(generic)).toBeGreaterThan(polls(rows));
+    expect(generic).not.toMatch(/^loop\.end\d+:$/m);
+    expect(generic).not.toMatch(/^loop\.poll\.run\d+:$/m);
     // Steps of 1 below 2^53 are exact, so 50 runs near it are bounded. From
     // 2^53, `j + 1 === j`: a counter ranging up to it can run without end
     // and keeps its poll.

packages/compiler/test/parameter-version-emission.test.ts

@@ -65,6 +65,27 @@ test("a loop-bounding parameter some caller passes unproven gets an int32 versio
   expect(body).toMatch(/fdiv double 0x3FF0000000000000, %\S+/);
 });
 
+test("the body a version leaves emits each loop once", async () => {
+  const llvm = await emitLlvm(`
+const fact: number[] = [];
+function fillFact(n: number): void {
+  fact[0] = 1;
+  for (let i = 1; i <= n; i++) fact[i] = fact[i - 1]! * i;
+}
+fillFact(Number(process.argv.length > 5 ? "7.5" : "6"));
+console.log(fact.join(" "));
+`);
+  expect(versions(llvm)).toEqual(["fillFact"]);
+  const loopBodies = (name: string): number => {
+    const body = new RegExp(`^define [^@]*@sc_\\w+?_${name}\\([^]*?^}`, "m").exec(llvm)?.[0] ?? "";
+    return body.match(/^loop\.b\d+:/gm)?.length ?? 0;
+  };
+  // The body runs only for numbers the version does not take: its loop gets
+  // no speculative integer copy beside the ordinary one.
+  expect(loopBodies("fillFact")).toBe(1);
+  expect(loopBodies("fillFact_int32Params\\.i32")).toBe(1);
+});
+
 test("a module that reads stacks keeps a single body", async () => {
   const llvm = await emitLlvm(`${loops}\nconsole.log(typeof new Error("x").stack);\n`);
   expect(versions(llvm)).toEqual([]);

packages/compiler/test/ts7/baselines/order-parity/tests/corpus/loop-bound-parameter-plain-loops.ts.json

@@ -0,0 +1,6 @@
+{
+ "order": [
+  "<repo>/tests/corpus/loop-bound-parameter-plain-loops.ts"
+ ],
+ "diags": []
+}

tests/corpus/loop-bound-parameter-plain-loops.ts

@@ -0,0 +1,74 @@
+import { Worker, isMainThread, parentPort, workerData } from "node:worker_threads";
+
+// Functions whose count parameters bound array loops: nested, counted,
+// two-counter and dense array loops, called from the main thread and a
+// worker with counts no caller proves to be int32s: integers, fractions,
+// NaN, -0, negatives and runs longer than a worker loop's unpolled run.
+// Counts past an array's end read undefined.
+
+const LENGTH = 40000;
+
+function rowSums(u: number[], v: number[], n: number, rows: number): void {
+  for (let i = 0; i < rows; i++) {
+    let t = 0;
+    for (let j = 0; j < n; j++) t += u[j]! / (i + j + 1);
+    v[i] = t;
+  }
+}
+
+function scale(src: number[], dst: number[], n: number, k: number): number {
+  let sum = 0;
+  for (let i = 0; i < n && i < dst.length; i++) {
+    dst[i] = (src[i]! * k) | 0;
+    sum += dst[i]!;
+  }
+  return sum;
+}
+
+function reverseRange(values: number[], n: number): number {
+  for (let lo = 0, hi = (n | 0) - 1; lo < hi; lo++, hi--) {
+    const t = values[lo]!;
+    values[lo] = values[hi]!;
+    values[hi] = t;
+  }
+  let check = 0;
+  for (let i = 0; i < 40; i++) check = (check * 31 + values[i]!) % 1000003;
+  return check;
+}
+
+function countSevens(values: number[], n: number, step: number): number {
+  let seen = 0;
+  for (let i = 0; i < n; i += step) if (values[i] === 7) seen++;
+  return seen;
+}
+
+function show(value: number): string {
+  return Object.is(value, -0) ? "-0" : Number.isInteger(value) ? String(value) : value.toFixed(6);
+}
+
+function run(label: string, n: number): string {
+  const u: number[] = [];
+  const ints: number[] = [];
+  for (let i = 0; i < LENGTH; i++) {
+    u.push(1 + (i % 13) * 0.25);
+    ints.push((i % 101) * 7 - 300);
+  }
+  const v = [0, 0, 0, 0];
+  const out = new Array<number>(LENGTH).fill(0);
+  rowSums(u, v, n, 4);
+  const scaled = scale(ints, out, n, 3);
+  const reversed = reverseRange(ints.map((x) => x % 10), n);
+  const sevens = countSevens(ints.map((x) => (x + 300) % 11), n, 3);
+  return `${label} ${show(n)}: rows ${v.map(show).join(" ")}; scaled ${scaled}; reversed ${reversed}; sevens ${sevens}`;
+}
+
+const counts: number[] = JSON.parse("[12, 12.5, 0, -0, -3, 7.25, 16385, 39999.5, 40000.5, 50000]");
+counts.push(NaN, Number("9"));
+
+if (isMainThread) {
+  for (const n of counts) console.log(run("main", n));
+  const worker = new Worker(new URL(import.meta.url), { workerData: counts });
+  worker.on("message", (line: string) => console.log(line));
+} else {
+  for (const n of workerData as number[]) parentPort!.postMessage(run("worker", n));
+}

tests/harness/corpus-costs/loop-bound-parameter-plain-loops.ts.json

@@ -0,0 +1 @@
+{ "plain": 4906, "san": 58104 }