Poll the loops of a nest on every iteration in ELF worker executables

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

@@ -96,7 +96,7 @@ import {
 import { emitStringInputs } from "./string-lifetimes.js";
 import { LOOPS } from "./loop-state.js";
 import { loopUnrollAttachment, periodicGuardPeriod } from "./periodic-guards.js";
-import { loopAccessesTypedArrays } from "../loop-polls.js";
+import { containsLoop, loopAccessesTypedArrays } from "../loop-polls.js";
 import { emitConstantChoice, matchConstantChoice } from "./constant-choices.js";
 
 /** A callback loop version copies the loop; keep the copies small. */
@@ -1228,8 +1228,8 @@ export function emitStmtBody(emitter: LlEmitter, s: IrStmt): void {
       };
       // A counted loop over typed arrays, or a loop to a typed array's
       // length, polls between blocks of WORKER_POLL_BLOCK iterations instead
-      // of counting a budget in its body: the inner loop over a block then
-      // has no exit besides its own bound, which LLVM can vectorize. A typed
+      // of in its body: the inner loop over a block then has no exit
+      // besides its own bound, which LLVM can vectorize. A typed
       // array's length is fixed, so each block can take it as its limit when
       // the loop does not assign the array's binding.
       //
@@ -2760,8 +2760,8 @@ interface StripMinedCounter {
  * iterations, with the cancellation poll between blocks: the outer header
  * tests the counter against the loop's limit, each block runs the body up
  * to the nearer of the limit and WORKER_POLL_BLOCK steps on, and the poll
- * follows the block. It polls as often as the budget it replaces, and the
- * inner loop's only exit besides `break` and `return` is its bound.
+ * follows the block. It polls once per block, and the inner loop's only
+ * exit besides `break` and `return` is its bound.
  *
  * With `runSlot` (for a loop that another loop of the function encloses), a
  * block that ends at the limit leaves the loop without polling, and a run
@@ -2895,21 +2895,6 @@ function emitStripMinedCountedLoop(
   B.br(lc);
 }
 
-/** Whether `body` holds a loop outside its closures. */
-function containsLoop(body: IrStmt[]): boolean {
-  const inStmt = (s: IrStmt): boolean =>
-    s.kind === "while" ||
-    s.kind === "doWhile" ||
-    s.kind === "for" ||
-    s.kind === "forOf" ||
-    !everyStmtChild(
-      s,
-      () => true,
-      (child) => !inStmt(child),
-    );
-  return body.some(inStmt);
-}
-
 function enterLoopHeader(emitter: LlEmitter, header: string): void {
   const dispatch = emitter.speculationHeader;
   emitter.speculationHeader = null;

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

@@ -27,7 +27,7 @@ import {
   streamTypedRefEligible,
   undefinedArmTag,
 } from "../../ir/analysis.js";
-import { loopTestsEveryIteration } from "../loop-polls.js";
+import { containsLoop, loopTestsEveryIteration } from "../loop-polls.js";
 import type {
   IrBytesElem,
   IrExpr,
@@ -826,6 +826,11 @@ const NO_BORROWED: ReadonlySet<number> = new Set();
 
 /** The nodes that write a binding by its id (numberGlobalWriters). */
 const GLOBAL_WRITE_KINDS: readonly IrNodeKind[] = ["assign", "varDecl", "assignExpr", "incDec"];
+
+/** workerLoopBudget's answer for a loop that polls on every iteration
+ * instead of counting down in a slot. */
+const POLL_EVERY_ITERATION = "%worker.poll.every";
+
 /** Runtime library calls (readsLoopMetrics). */
 const LIB_CALL_KINDS: readonly IrNodeKind[] = ["libCall"];
 
@@ -4013,10 +4018,22 @@ export class LlEmitter {
     );
   }
 
+  /** The countdown slot of `loop`'s cancellation poll (emitWorkerLoopCheck),
+   * or null when the loop does not poll. An ELF worker executable's poll is
+   * one ordered load of the thread's own pending word. A loop nest keeps
+   * one countdown per loop, which its register pressure often leaves in
+   * memory, so the loops of a nest there poll on every iteration and get
+   * POLL_EVERY_ITERATION instead of a slot. A loop on its own keeps its
+   * countdown in a register, which costs less than the load. */
   workerLoopBudget(
     loop: Extract<IrStmt, { kind: "while" | "doWhile" | "for" | "forOf" }>,
   ): string | null {
     if (!this.workerLoopPolls(loop)) return null;
+    if (
+      this.executableTls &&
+      (this.jumpTargets.some((target) => target.kind === "loop") || containsLoop(loop.body))
+    )
+      return POLL_EVERY_ITERATION;
     const slot = this.B.slot();
     this.B.entryAllocas.push(`${slot} = alloca i32 ; worker cancellation budget`);
     this.B.line(`store i32 0, ptr ${slot}`);
@@ -4121,6 +4138,10 @@ export class LlEmitter {
    * stay in the loop. */
   emitWorkerLoopCheck(slot: string | null): void {
     if (slot === null) return;
+    if (slot === POLL_EVERY_ITERATION) {
+      this.emitWorkerPoll();
+      return;
+    }
     const B = this.B;
     const budget = B.tmp(),
       due = B.tmp(),
@@ -4140,9 +4161,9 @@ export class LlEmitter {
   }
 
   /** The cancellation poll itself. Calls and throwing operations retain
-   * their own checks. Loops amortize this extra poll across bounded
-   * batches of iterations so scalar work does not cross the runtime ABI on
-   * every iteration. */
+   * their own checks. Outside ELF worker executables, loops amortize this
+   * extra poll across bounded batches of iterations so scalar work does not
+   * cross the runtime ABI on every iteration. */
   emitWorkerPoll(): void {
     this.emitPendingCheck(undefined, true);
   }

packages/compiler/src/backend/loop-polls.ts

@@ -137,3 +137,18 @@ export function loopAccessesTypedArrays(loop: Loop): boolean {
   };
   return loop.body.some(inStmt);
 }
+
+/** Whether `body` holds a loop outside its closures. */
+export function containsLoop(body: readonly IrStmt[]): boolean {
+  const inStmt = (s: IrStmt): boolean =>
+    s.kind === "while" ||
+    s.kind === "doWhile" ||
+    s.kind === "for" ||
+    s.kind === "forOf" ||
+    !everyStmtChild(
+      s,
+      () => true,
+      (child) => !inStmt(child),
+    );
+  return body.some(inStmt);
+}

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

@@ -137,6 +137,59 @@ test("worker programs poll counted loops over typed arrays between blocks of ite
   }
 });
 
+test("ELF worker executables poll the loops of a nest on every iteration", async () => {
+  const dir = await mkdtemp(join(tmpdir(), "scriptc-nest-poll-"));
+  try {
+    const entry = join(dir, "nest.ts");
+    await writeFile(
+      entry,
+      [
+        'import { isMainThread, Worker } from "node:worker_threads";',
+        "export function steps(limit: number): number {",
+        "  let total = 0;",
+        "  for (let start = 1; start < limit; start++) {",
+        "    let x = start;",
+        "    while (x !== 1) {",
+        "      x = x % 2 === 0 ? x / 2 : 3 * x + 1;",
+        "      total++;",
+        "    }",
+        "  }",
+        "  return total;",
+        "}",
+        "export function mix(n: number): number {",
+        "  let t = 0;",
+        "  for (let i = 0; i < n; i++) t = (t * 31 + i) % 1000003;",
+        "  return t;",
+        "}",
+        "console.log(steps(8), mix(8));",
+        "if (!isMainThread) new Worker(new URL(import.meta.url));",
+      ].join("\n"),
+    );
+    const outPath = join(dir, "nest.ir.json");
+    const result = await compile(entry, { outPath, outDir: dir, outputKind: "ir" });
+    if (!result.ok) throw new Error(result.diagnostics.map((d) => d.message).join("\n"));
+    const mod = deserializeModule(await readFile(outPath, "utf8"));
+    const body = (ll: string, name: string): string =>
+      ll.match(
+        new RegExp(
+          `define internal [^\\n]+ @sc_b?f_${name}(?:\\.i32)?\\([^\\n]*\\n([\\s\\S]*?)\\n\\}`,
+        ),
+      )![1]!;
+    // Both loops of the nest test the thread's pending word on every
+    // iteration and keep no countdown; a loop on its own keeps one.
+    const elf = emitLlvmModule(mod, { targetTriple: "x86_64-unknown-linux-gnu" });
+    const steps = body(elf, "steps");
+    expect(steps).not.toContain("worker cancellation budget");
+    expect(steps.match(/= load atomic i32, ptr @scr_exc_flag monotonic/g)?.length).toBe(2);
+    expect(body(elf, "mix")).toContain("worker cancellation budget");
+    // Elsewhere the poll reads the process-wide alert word: the nest counts.
+    const darwin = emitLlvmModule(mod, { targetTriple: "arm64-apple-macosx14.0.0" });
+    expect(body(darwin, "steps").match(/worker cancellation budget/g)?.length).toBe(2);
+  } finally {
+    await rm(dir, { recursive: true, force: true });
+  }
+});
+
 test("worker programs poll no loop whose proven ranges bound its runs inside another loop", async () => {
   const dir = await mkdtemp(join(tmpdir(), "scriptc-bounded-loop-polls-"));
   try {

packages/compiler/test/ts7/baselines/order-parity/tests/corpus/native-worker-terminate-short-inner-loops.ts.json

@@ -0,0 +1,6 @@
+{
+ "order": [
+  "<repo>/tests/corpus/native-worker-terminate-short-inner-loops.ts"
+ ],
+ "diags": []
+}

tests/corpus/native-worker-terminate-short-inner-loops.ts

@@ -0,0 +1,96 @@
+// worker.terminate() stops a worker that spends its time in short runs of
+// inner loops: counted and paired loops of a few iterations inside loops
+// without a condition, runs at the length where an inner loop starts to
+// poll on its own, and a short inner loop whose body calls a slow function.
+// Each terminate resolves promptly.
+import { Worker, isMainThread, parentPort, workerData } from "node:worker_threads";
+
+function flips(n: number): number {
+  const p = new Int32Array(n),
+    q = new Int32Array(n),
+    count = new Int32Array(n);
+  for (let i = 0; i < n; i++) p[i] = i;
+  let total = 0;
+  for (let round = 0; ; round++) {
+    for (let i = 0; i < n; i++) q[i] = p[i]!;
+    let first = q[0]!;
+    for (;;) {
+      if (first === 0) break;
+      for (let i = 0, j = first; i < j; i++, j--) {
+        const t = q[i]!;
+        q[i] = q[j]!;
+        q[j] = t;
+      }
+      total++;
+      first = q[0]!;
+    }
+    // The next permutation: rotate the first k + 1 elements for the
+    // smallest k whose count has not run out.
+    for (let k = 1; ; k++) {
+      if (k === n) {
+        for (let i = 0; i < n; i++) p[i] = i;
+        count.fill(0);
+        break;
+      }
+      const head = p[0]!;
+      for (let i = 0; i < k; i++) p[i] = p[i + 1]!;
+      p[k] = head;
+      if (++count[k]! <= k) break;
+      count[k] = 0;
+    }
+    if (round === 2000) parentPort!.postMessage("ready");
+  }
+  return total;
+}
+
+function thresholds(values: number[]): number {
+  let acc = 0;
+  for (let round = 0; ; round++) {
+    for (const run of [63, 64, 65]) {
+      for (let i = 0; i < run; i++) acc = (acc + values[i]! * round) % 1000003;
+      for (let i = 0, j = run; i < j; i++, j--) acc = (acc + values[i]! + values[j]!) % 1000003;
+    }
+    if (round === 2000) parentPort!.postMessage("ready");
+  }
+  return acc;
+}
+
+function slow(seed: number): number {
+  let acc = seed;
+  for (let i = 0; i < 20000; i++) acc = (acc * 31 + i) % 1000003;
+  return acc;
+}
+
+function slowBodies(): number {
+  let acc = 0;
+  for (let round = 0; ; round++) {
+    for (let i = 0; i < 4; i++) acc = (acc + slow(acc + i)) % 1000003;
+    if (round === 50) parentPort!.postMessage("ready");
+  }
+  return acc;
+}
+
+if (isMainThread) {
+  for (const kind of ["flips", "thresholds", "slow"]) {
+    await new Promise<void>((resolve) => {
+      const worker = new Worker(new URL(import.meta.url), { workerData: kind });
+      worker.once("message", () => {
+        const asked = Date.now();
+        void worker.terminate().then((code: number) => {
+          console.log(kind, "terminate", code, Date.now() - asked < 1000 ? "prompt" : "late");
+        });
+      });
+      worker.once("exit", (code: number) => {
+        console.log(kind, "exit", code);
+        resolve();
+      });
+    });
+  }
+} else {
+  let result = 0;
+  if (workerData === "flips") result = flips(9);
+  else if (workerData === "thresholds")
+    result = thresholds(Array.from({ length: 70 }, (_, i) => i));
+  else result = slowBodies();
+  console.log("unexpected result", result);
+}

tests/harness/corpus-costs/native-worker-terminate-short-inner-loops.ts.json

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