From f1e27f2e44d018d5b008a8038a1dfd7c4e5c6307 Mon Sep 17 00:00:00 2001 From: Chris Tate Date: Mon, 26 Oct 2026 14:31:32 -0500 Subject: [PATCH] Poll the loops of a nest on every iteration in ELF worker executables --- .../src/backend/llvm/emit-statements.ts | 25 +---- packages/compiler/src/backend/llvm/emitter.ts | 29 +++++- packages/compiler/src/backend/loop-polls.ts | 15 +++ .../test/counted-loop-emission.test.ts | 53 ++++++++++ ...worker-terminate-short-inner-loops.ts.json | 6 ++ ...tive-worker-terminate-short-inner-loops.ts | 96 +++++++++++++++++++ ...worker-terminate-short-inner-loops.ts.json | 1 + 7 files changed, 201 insertions(+), 24 deletions(-) create mode 100644 packages/compiler/test/ts7/baselines/order-parity/tests/corpus/native-worker-terminate-short-inner-loops.ts.json create mode 100644 tests/corpus/native-worker-terminate-short-inner-loops.ts create mode 100644 tests/harness/corpus-costs/native-worker-terminate-short-inner-loops.ts.json diff --git a/packages/compiler/src/backend/llvm/emit-statements.ts b/packages/compiler/src/backend/llvm/emit-statements.ts index 45d341853..1ea18d8e1 100644 --- a/packages/compiler/src/backend/llvm/emit-statements.ts +++ b/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; diff --git a/packages/compiler/src/backend/llvm/emitter.ts b/packages/compiler/src/backend/llvm/emitter.ts index 2f3a0f212..4e6f8af0b 100644 --- a/packages/compiler/src/backend/llvm/emitter.ts +++ b/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 = 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, ): 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); } diff --git a/packages/compiler/src/backend/loop-polls.ts b/packages/compiler/src/backend/loop-polls.ts index 866711fb7..b49026fa2 100644 --- a/packages/compiler/src/backend/loop-polls.ts +++ b/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); +} diff --git a/packages/compiler/test/counted-loop-emission.test.ts b/packages/compiler/test/counted-loop-emission.test.ts index 66d31e14b..914dc81d8 100644 --- a/packages/compiler/test/counted-loop-emission.test.ts +++ b/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 { diff --git a/packages/compiler/test/ts7/baselines/order-parity/tests/corpus/native-worker-terminate-short-inner-loops.ts.json b/packages/compiler/test/ts7/baselines/order-parity/tests/corpus/native-worker-terminate-short-inner-loops.ts.json new file mode 100644 index 000000000..8332ee20d --- /dev/null +++ b/packages/compiler/test/ts7/baselines/order-parity/tests/corpus/native-worker-terminate-short-inner-loops.ts.json @@ -0,0 +1,6 @@ +{ + "order": [ + "/tests/corpus/native-worker-terminate-short-inner-loops.ts" + ], + "diags": [] +} diff --git a/tests/corpus/native-worker-terminate-short-inner-loops.ts b/tests/corpus/native-worker-terminate-short-inner-loops.ts new file mode 100644 index 000000000..33a213ea8 --- /dev/null +++ b/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((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); +} diff --git a/tests/harness/corpus-costs/native-worker-terminate-short-inner-loops.ts.json b/tests/harness/corpus-costs/native-worker-terminate-short-inner-loops.ts.json new file mode 100644 index 000000000..c85984a8b --- /dev/null +++ b/tests/harness/corpus-costs/native-worker-terminate-short-inner-loops.ts.json @@ -0,0 +1 @@ +{ "plain": 4906, "san": 6969 }