Finish bailed speculative iterations outside the ordinary loop copy

A speculative integer loop's failed check used to continue at the failed
statement inside the ordinary copy. That gave the ordinary copy a second
entry, so LLVM saw an irreducible cycle and did not hoist its invariant
loads or keep its values in registers. A loop whose check fails on every
entry, such as an inner loop over fractional fields, then ran every
iteration in that unoptimized copy.

A failed check now continues in a third copy that is entered only from the
fast copy: it re-executes the failed statement, the rest of the body and
the update, then jumps to the ordinary copy's header. The ordinary copy has
a single entry and stays a natural loop.

Each loop site also counts its entries in a row whose fast copy bailed. A
fast copy that completes a speculated write and leaves the loop clears the
count, and after 8 bailing entries in a row the site runs the ordinary copy
only. The count is a hint, read and written with unordered atomics: both
copies give the same results.

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

@@ -803,6 +803,7 @@ export function emitStmtBody(emitter: LlEmitter, s: IrStmt): void {
         const lp = workerBudget !== null ? B.newLabel("loop.p") : lc;
         B.br(lc);
         B.startBlock(lc);
+        enterLoopHeader(emitter, lc);
         B.condBr(emitter.emitCondition(s.cond), lb, exit);
         B.startBlock(lb);
         emitter.jumpTargets.push({
@@ -835,6 +836,7 @@ export function emitStmtBody(emitter: LlEmitter, s: IrStmt): void {
         const lc = B.newLabel("loop.c");
         B.br(lb);
         B.startBlock(lb);
+        enterLoopHeader(emitter, lb);
         emitter.jumpTargets.push({
           kind: "loop",
           brkLabel: exit,
@@ -1135,6 +1137,7 @@ export function emitStmtBody(emitter: LlEmitter, s: IrStmt): void {
           s.update || capturedLets.length > 0 || workerBudget !== null ? B.newLabel("loop.u") : lc;
         B.br(lc);
         B.startBlock(lc);
+        enterLoopHeader(emitter, lc);
         if (integerLoop && integerSlot) {
           const receiver = emitter.emitStableReceiver(integerLoop.limitReceiver, []);
           const lenPtr = B.tmp();
@@ -1885,6 +1888,10 @@ export function emitSwitch(emitter: LlEmitter, s: IrStmt & { kind: "switch" }):
   B.startBlock(end);
 }
 
+/** Entries in a row whose fast copy bailed before a speculative loop site
+ * stops trying it. */
+const SPECULATION_BAILS = 8;
+
 /** Emit a loop's condition, body and update through `emitCopy`, twice
  * when speculative-integers.ts plans a fast integer copy. Both copies
  * share the loop's setup and exit; the ordinary copy comes first so the
@@ -1939,23 +1946,66 @@ export function emitLoopCopies(
       ok = both;
     }
   }
+  // Each entry whose fast copy bails counts at the loop site, and a fast
+  // copy that completes a speculated write and leaves the loop clears the
+  // count. After SPECULATION_BAILS bailing entries in a row, the site runs
+  // the ordinary copy only. The count is a hint: both copies give the same
+  // results, and a race between workers loses at most an update.
+  const bails = `@scr.spec.bails.${emitter.speculativeLoopSites++}`;
+  emitter.declare(`${bails} = internal global i8 0`);
+  const count = B.tmp(),
+    trying = B.tmp();
+  B.line(`${count} = load atomic i8, ptr ${bails} unordered, align 1`);
+  B.line(`${trying} = icmp ult i8 ${count}, ${SPECULATION_BAILS}`);
+  if (ok === "true") ok = trying;
+  else {
+    const both = B.tmp();
+    B.line(`${both} = and i1 ${ok}, ${trying}`);
+    ok = both;
+  }
   const fast = B.newLabel("spec.fast"),
     slow = B.newLabel("spec.slow");
   B.condBr(ok, fast, slow);
   const resume = new Map<IrStmt, { label: string; placed: boolean }>();
   for (const s of plan.writes) resume.set(s, { label: B.newLabel("spec.resume"), placed: false });
+  // A failed check finishes its iteration in a third copy, entered only
+  // there: it re-executes the statement and the rest of the body and the
+  // update, then enters the ordinary copy through its header. Jumping
+  // into the middle of the ordinary copy instead would give that loop a
+  // second entry, and LLVM does not optimize an irreducible loop.
+  const points = [...resume.values()];
+  let header: string | null = null;
   B.startBlock(slow);
+  emitter.speculationHeader = (label) => {
+    header = label;
+  };
+  try {
+    emitCopy(exit);
+  } finally {
+    emitter.speculationHeader = null;
+  }
+  if (header === null)
+    throw new InternalCompilerError("llvm emitter bug: speculative loop copy without a header");
+  B.startBlock(B.newLabel("spec.tail"));
   emitter.speculationResume = resume;
+  emitter.speculationHeader = () => {
+    B.br(header!);
+    B.startBlock(B.newLabel("spec.unused"));
+  };
   try {
     emitCopy(exit);
   } finally {
     emitter.speculationResume = null;
+    emitter.speculationHeader = null;
   }
   B.startBlock(fast);
-  if (![...resume.values()].every((r) => r.placed)) {
+  if (!points.every((r) => r.placed)) {
     B.br(slow);
     return;
   }
+  const ran = B.slot();
+  B.entryAllocas.push(`${ran} = alloca i1 ; speculated write completed`);
+  B.line(`store i1 false, ptr ${ran}`);
   for (const id of plan.locals.keys()) {
     const shadow = B.slot();
     B.entryAllocas.push(
@@ -1992,22 +2042,43 @@ export function emitLoopCopies(
     writes: plan.writes,
     resume: new Map([...resume].map(([s, r]) => [s, r.label])),
     deferred,
+    bails,
+    ran,
   };
   try {
-    if (deferred.size === 0) emitCopy(exit);
-    else {
-      const leave = B.newLabel("spec.leave");
-      emitCopy(leave);
-      B.startBlock(leave);
-      emitter.writeBackDeferred(deferred);
-      B.br(exit);
-    }
+    const leave = B.newLabel("spec.leave");
+    emitCopy(leave);
+    B.startBlock(leave);
+    emitter.writeBackDeferred(deferred);
+    // An entry that never ran a speculated write (a loop of no iterations)
+    // says nothing about the site.
+    const wrote = B.tmp(),
+      left = B.tmp(),
+      counted = B.tmp(),
+      both = B.tmp(),
+      clear = B.newLabel("spec.clear");
+    B.line(`${wrote} = load i1, ptr ${ran}`);
+    B.line(`${left} = load atomic i8, ptr ${bails} unordered, align 1`);
+    B.line(`${counted} = icmp ne i8 ${left}, 0`);
+    B.line(`${both} = and i1 ${wrote}, ${counted}`);
+    B.condBr(both, clear, exit);
+    B.startBlock(clear);
+    B.line(`store atomic i8 0, ptr ${bails} unordered, align 1`);
+    B.br(exit);
   } finally {
     emitter.speculation = null;
     for (const id of plan.locals.keys()) emitter.integerLoopBindings.delete(id);
   }
 }
 
+/** At the top of a loop copy's header: lets a speculative loop find its
+ * ordinary copy's header, and keeps its third copy out of the loop. */
+function enterLoopHeader(emitter: LlEmitter, header: string): void {
+  const dispatch = emitter.speculationHeader;
+  emitter.speculationHeader = null;
+  dispatch?.(header);
+}
+
 /** A proven integer loop (speculative-integers.ts) for this loop, or null.
  * A shadowed local's double is stale until the loop is left, which only
  * the loop's exit, a return and an exception do: the loop is planned only
@@ -2134,7 +2205,8 @@ export function emitSpeculativeWrite(
   if (emitter.frames.pop()!.length > 0)
     throw new InternalCompilerError("llvm emitter bug: speculative write owns a temporary");
   B.line(`store i64 ${result}, ptr ${binding.slot}`);
-  const deferred = emitter.speculation!.deferred;
+  const { deferred, bails: site, ran } = emitter.speculation!;
+  if (ran !== undefined) B.line(`store i1 true, ptr ${ran}`);
   if (!deferred.has(target.localId)) {
     const number = B.tmp();
     B.line(`${number} = sitofp i64 ${result} to double`);
@@ -2145,6 +2217,14 @@ export function emitSpeculativeWrite(
   B.br(done);
   B.startBlock(failed);
   emitter.writeBackDeferred(deferred);
+  if (site !== undefined) {
+    const count = B.tmp(),
+      next = B.tmp();
+    emitter.declare("declare i8 @llvm.uadd.sat.i8(i8, i8)");
+    B.line(`${count} = load atomic i8, ptr ${site} unordered, align 1`);
+    B.line(`${next} = call i8 @llvm.uadd.sat.i8(i8 ${count}, i8 1)`);
+    B.line(`store atomic i8 ${next}, ptr ${site} unordered, align 1`);
+  }
   B.br(resume!);
   B.startBlock(done);
 }

packages/compiler/src/backend/llvm/emitter.ts

@@ -1102,9 +1102,20 @@ export class LlEmitter {
     resume: Map<IrStmt, string>;
     /** Globals whose double is written back only when the copy is left. */
     deferred: Set<string>;
+    /** The loop site's count of bailing entries in a row (a global), and
+     * whether this entry completed a speculated write (a slot). A proven
+     * integer loop has neither. */
+    bails?: string;
+    ran?: string;
   } | null = null;
+  /** Speculative loop sites numbered so far in the module. */
+  speculativeLoopSites = 0;
   /** The ordinary copy being emitted: resume labels still to place. */
   speculationResume: Map<IrStmt, { label: string; placed: boolean }> | null = null;
+  /** Run once at the top of the next loop copy's header
+   * (`enterLoopHeader`): records the ordinary copy's header, or sends the
+   * third copy's header there. */
+  speculationHeader: ((header: string) => void) | null = null;
   /** The function being emitted, for speculative globals. */
   speculativeFunction = "";
   /** Every function that writes each number global (computed on demand). */

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

@@ -0,0 +1,85 @@
+import { mkdtemp, readFile, rm, writeFile } from "node:fs/promises";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
+import { expect, test } from "vitest";
+import { compile } from "../src/index.js";
+
+async function emitLlvm(source: string): Promise<string> {
+  const dir = await mkdtemp(join(tmpdir(), "scriptc-speculative-loop-"));
+  try {
+    const entry = join(dir, "main.ts");
+    const outPath = join(dir, "main.ll");
+    await writeFile(entry, source);
+    const result = await compile(entry, { outDir: dir, outPath, outputKind: "llvm" });
+    expect(result.ok, JSON.stringify(result.diagnostics)).toBe(true);
+    return await readFile(outPath, "utf8");
+  } finally {
+    await rm(dir, { recursive: true, force: true });
+  }
+}
+
+const source = `
+class Point {
+  x: number;
+  w = 0;
+  constructor(x: number) {
+    this.x = x;
+  }
+}
+function scaled(points: Point[], scale: number): number {
+  let s = 0;
+  for (let i = 0; i < points.length; i++) {
+    const p = points[i]!;
+    p.w += 1;
+    const k = p.x * scale;
+    s += k * k;
+  }
+  return s;
+}
+console.log(scaled([new Point(1), new Point(2)], 2), scaled([new Point(1)], 0.5));
+`;
+
+function labels(fn: string): string[] {
+  return [...fn.matchAll(/^([\w.]+):/gm)].map((m) => m[1]!);
+}
+
+test("a failed check finishes its iteration outside the ordinary copy", async () => {
+  const llvm = await emitLlvm(source);
+  const fn = /^define internal double @sc_bf_scaled\([^]*?^}/m.exec(llvm)?.[0] ?? "";
+  const order = labels(fn);
+  const slow = order.findIndex((l) => l.startsWith("spec.slow"));
+  const tail = order.findIndex((l) => l.startsWith("spec.tail"));
+  const fast = order.findIndex((l) => l.startsWith("spec.fast"));
+  expect(slow).toBeGreaterThanOrEqual(0);
+  expect(tail).toBeGreaterThan(slow);
+  expect(fast).toBeGreaterThan(tail);
+  // Every resume point is in the third copy, between its start and the
+  // fast copy, so the ordinary copy has one entry.
+  const resumes = order.flatMap((l, i) => (l.startsWith("spec.resume") ? [i] : []));
+  expect(resumes.length).toBeGreaterThan(0);
+  for (const i of resumes) {
+    expect(i).toBeGreaterThan(tail);
+    expect(i).toBeLessThan(fast);
+  }
+  // The third copy's header leaves for the ordinary copy's header.
+  const header = order[slow + 1]!;
+  const third = fn.slice(fn.indexOf(`\n${order[tail]}:`), fn.indexOf(`\n${order[fast]}:`));
+  expect(third).toContain(`br label %${header}`);
+});
+
+test("a loop site stops entering its fast copy after repeated bails", async () => {
+  const llvm = await emitLlvm(source);
+  const site = /^(@scr\.spec\.bails\.\d+) = internal global i8 0$/m.exec(llvm)?.[1];
+  expect(site).toBeDefined();
+  const fn = /^define internal double @sc_bf_scaled\([^]*?^}/m.exec(llvm)?.[0] ?? "";
+  const escaped = site!.replace(/\./g, "\\.");
+  // The entry test reads the count, a bail adds one, and a completed fast
+  // copy clears it.
+  expect(fn).toMatch(
+    new RegExp(
+      `(%\\S+) = load atomic i8, ptr ${escaped} unordered, align 1\\n\\s*%\\S+ = icmp ult i8 \\1, 8`,
+    ),
+  );
+  expect(fn).toMatch(new RegExp(`call i8 @llvm\\.uadd\\.sat\\.i8\\(i8 %\\S+, i8 1\\)`));
+  expect(fn).toContain(`store atomic i8 0, ptr ${site} unordered, align 1`);
+});

packages/compiler/test/ts7/baselines/order-parity/tests/corpus/7321-speculative-loop-bail-sites.ts.json

@@ -0,0 +1,6 @@
+{
+ "order": [
+  "<repo>/tests/corpus/7321-speculative-loop-bail-sites.ts"
+ ],
+ "diags": []
+}

tests/corpus/7321-speculative-loop-bail-sites.ts

@@ -0,0 +1,111 @@
+// Speculative integer loops whose checks fail: on the first statement of
+// every entry (fractional field data), after an effect earlier in the same
+// iteration, before a `continue`, `break` or `return`, and on entries
+// that alternate between integer and fractional data. A site that keeps
+// failing stops trying its integer copy; every result is Node's.
+
+class Point {
+  x: number;
+  y: number;
+  w = 0;
+
+  constructor(x: number, y: number) {
+    this.x = x;
+    this.y = y;
+  }
+}
+
+function show(label: string, value: number): void {
+  console.log(label, Object.is(value, -0) ? "-0" : String(value));
+}
+
+// Each inner entry fails at its first statement when the fields hold
+// fractions, and runs no iteration at all for the last point.
+function pairs(points: Point[]): number {
+  let total = 0;
+  for (let i = 0; i < points.length; i++) {
+    const p = points[i]!;
+    let acc = p.w;
+    for (let j = i + 1; j < points.length; j++) {
+      const q = points[j]!;
+      const dx = p.x - q.x;
+      const dy = p.y - q.y;
+      const d2 = dx * dx + dy * dy;
+      acc += d2;
+      q.w += d2 * 0.5;
+    }
+    p.w = acc;
+    total += acc;
+  }
+  return total;
+}
+
+const fractional = [
+  new Point(0.5, 1),
+  new Point(1.25, -2),
+  new Point(3, 4.75),
+  new Point(-1.5, 0),
+];
+const integral = [new Point(0, 1), new Point(2, -2), new Point(3, 4), new Point(-1, 0)];
+for (let round = 0; round < 12; round++) {
+  show(`pairs fractional ${round}`, pairs(fractional));
+}
+// The same site, now with integer data after it stopped speculating.
+show("pairs integral", pairs(integral));
+show("pairs integral again", pairs(integral));
+
+// A field write comes before the failing statement: the iteration resumes
+// after it and must not repeat it.
+function afterEffect(points: Point[], scale: number): number {
+  let s = 0;
+  for (let i = 0; i < points.length; i++) {
+    const p = points[i]!;
+    p.w += 1;
+    const k = p.x * scale;
+    s += k * k;
+  }
+  return s;
+}
+const marks = [new Point(1, 0), new Point(2, 0), new Point(3, 0)];
+show("after effect 2", afterEffect(marks, 2));
+show("after effect 0.5", afterEffect(marks, 0.5));
+show("after effect 3", afterEffect(marks, 3));
+console.log(marks.map((p) => p.w).join(" "));
+
+// The failing statement is followed by `continue`, `break` and `return`.
+function control(n: number, step: number): number {
+  let acc = 0;
+  let i = 0;
+  while (i < n) {
+    const v = i * step;
+    i++;
+    if (v * v > 400) break;
+    if ((i & 1) === 0) continue;
+    acc += v * 3;
+    if (acc > 1000) return -acc;
+  }
+  return acc;
+}
+for (const step of [1, 0.5, 2, 0.75, 7, 1.5, 3, 0.25, 5, 1.75, 4]) {
+  show(`control ${step}`, control(40, step));
+}
+
+// do-while: the check fails part-way through, on some entries only.
+function steps(start: number, factor: number): number {
+  let v = start;
+  let k = 0;
+  do {
+    v = v * factor - k;
+    k++;
+  } while (k < 6);
+  return v;
+}
+for (let r = 0; r < 20; r++) show(`steps ${r}`, steps(r % 3 === 0 ? r + 0.5 : r, 3));
+
+// Overflow past 2^53 on later entries only.
+function powers(base: number, n: number): number {
+  let p = 1;
+  for (let i = 0; i < n; i++) p = p * base;
+  return p;
+}
+for (const n of [3, 10, 40, 80, 5, 60, 2, 70, 1, 90, 0, 4]) show(`powers ${n}`, powers(3, n));

tests/harness/corpus-costs/7321-speculative-loop-bail-sites.ts.json

@@ -0,0 +1 @@
+{ "plain": 2000, "san": 2600 }