diff --git a/packages/compiler/src/backend/llvm/emit-statements.ts b/packages/compiler/src/backend/llvm/emit-statements.ts index cbdeb0558..2efe2cf3a 100644 --- a/packages/compiler/src/backend/llvm/emit-statements.ts +++ b/packages/compiler/src/backend/llvm/emit-statements.ts @@ -91,6 +91,7 @@ import { } from "./string-construction.js"; import { emitStringInputs } from "./string-lifetimes.js"; import { LOOPS } from "./loop-state.js"; +import { loopAccessesTypedArrays } from "../loop-polls.js"; /** A callback loop version copies the loop; keep the copies small. */ const MAX_CALLBACK_VERSION_COST = 400; @@ -1211,9 +1212,67 @@ export function emitStmtBody(emitter: LlEmitter, s: IrStmt): void { B.line(`store ptr ${fresh}, ptr ${slot}`); } }; - const workerBudget = emitter.workerLoopBudget(s); + // 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 + // array's length is fixed, so each block can take it as its limit. + const stripMined = + capturedLets.length === 0 && + integerSlot !== null && + ((countedLoop !== null && + countedLimit !== null && + Math.abs(countedLoop.step) * WORKER_POLL_BLOCK <= 2 ** 52) || + (integerLoop !== null && integerArrayLoop === null)) && + emitter.workerLoopPolls(s) && + loopAccessesTypedArrays(s); + const workerBudget = stripMined ? null : emitter.workerLoopBudget(s); + // One slot for every copy of the loop: a speculative copy resumes an + // iteration in another copy's body, which then reads this bound. + let blockBound: string | null = null; + if (stripMined) { + blockBound = B.slot(); + B.entryAllocas.push( + `${blockBound} = alloca ${countedLoop ? "i64" : emitter.sizeType} ; cancellation poll block bound`, + ); + } freshenBindings(); emitter.emitLoopCopies(s, capturedLets.length === 0, le, (exit) => { + if (stripMined) { + const counter: StripMinedCounter = countedLoop + ? { + slot: integerSlot!, + type: "i64", + signed: true, + step: countedLoop.step, + followers: countedLoop.followers.map((f, i) => ({ + slot: followerSlots[i]!, + step: f.step, + })), + limit: () => countedLimit!, + } + : { + slot: integerSlot!, + type: emitter.sizeType, + signed: false, + step: 1, + followers: [], + limit: () => { + const receiver = emitter.emitStableReceiver(integerLoop!.limitReceiver, []); + const lenPtr = B.tmp(); + const len = B.tmp(); + B.line( + `${lenPtr} = getelementptr inbounds %ScrBytes, ptr ${receiver.name}, i64 0, i32 1`, + ); + B.line( + `${len} = load ${emitter.sizeType}, ptr ${lenPtr}${bytesHeaderAlias(emitter)}`, + ); + return len; + }, + }; + emitStripMinedCountedLoop(emitter, s, counter, blockBound!, exit); + return; + } const lc = B.newLabel("loop.c"); const lb = B.newLabel("loop.b"); const lu = @@ -2549,6 +2608,102 @@ function emitInt32KindLoop( /** 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. */ +/** Iterations between cancellation polls of a strip-mined counted loop. */ +const WORKER_POLL_BLOCK = 64; + +/** The counter of a strip-mined loop: its slot, the LLVM integer type and + * signedness of its values, its step, the counters stepped beside it, and + * the limit, emitted in the outer header before each block. */ +interface StripMinedCounter { + slot: string; + type: string; + signed: boolean; + step: number; + followers: { slot: string; step: number }[]; + limit: () => string; +} + +/** A counted `for` loop of a worker program as blocks of WORKER_POLL_BLOCK + * 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. */ +function emitStripMinedCountedLoop( + emitter: LlEmitter, + s: Extract, + counter: StripMinedCounter, + boundSlot: string, + exit: string, +): void { + const B = emitter.B; + const { slot, type, signed, step } = counter; + const ascending = step > 0; + const compare = !signed ? "ult" : ascending ? "slt" : "sgt"; + const saturating = `llvm.${signed ? "sadd" : "uadd"}.sat.${type}`; + const nearer = `llvm.${!signed ? "umin" : ascending ? "smin" : "smax"}.${type}`; + const wrap = signed ? "nsw" : "nuw"; + const lc = B.newLabel("loop.c"); + const block = B.newLabel("loop.block"); + const bc = B.newLabel("loop.bc"); + const lb = B.newLabel("loop.b"); + const lu = B.newLabel("loop.u"); + const poll = B.newLabel("loop.poll"); + emitter.declare(`declare ${type} @${saturating}(${type}, ${type})`); + emitter.declare(`declare ${type} @${nearer}(${type}, ${type})`); + B.br(lc); + B.startBlock(lc); + enterLoopHeader(emitter, lc); + const start = B.tmp(); + const open = B.tmp(); + B.line(`${start} = load ${type}, ptr ${slot}`); + const limit = counter.limit(); + B.line(`${open} = icmp ${compare} ${type} ${start}, ${limit}`); + B.condBr(open, block, exit); + B.startBlock(block); + const stepped = B.tmp(); + const bound = B.tmp(); + B.line( + `${stepped} = call ${type} @${saturating}(${type} ${start}, ${type} ${step * WORKER_POLL_BLOCK})`, + ); + B.line(`${bound} = call ${type} @${nearer}(${type} ${stepped}, ${type} ${limit})`); + B.line(`store ${type} ${bound}, ptr ${boundSlot}`); + B.br(bc); + B.startBlock(bc); + const index = B.tmp(); + const end = B.tmp(); + const inBlock = B.tmp(); + B.line(`${index} = load ${type}, ptr ${slot}`); + B.line(`${end} = load ${type}, ptr ${boundSlot}`); + B.line(`${inBlock} = icmp ${compare} ${type} ${index}, ${end}`); + B.condBr(inBlock, lb, poll); + B.startBlock(lb); + emitter.jumpTargets.push({ + kind: "loop", + brkLabel: exit, + contLabel: lu, + labels: s.labels, + frameDepth: emitter.frames.length, + scopeDepth: emitter.scopes.length, + finallyDepth: emitter.finallyStack.length, + }); + emitter.emitBlock(s.body); + emitter.jumpTargets.pop(); + B.br(lu); + B.startBlock(lu); + for (const { slot: stepSlot, step: by } of [{ slot, step }, ...counter.followers]) { + const old = B.tmp(); + const next = B.tmp(); + B.line(`${old} = load ${type}, ptr ${stepSlot}`); + B.line(`${next} = add ${wrap} ${type} ${old}, ${by}`); + B.line(`store ${type} ${next}, ptr ${stepSlot}`); + } + B.br(bc); + B.startBlock(poll); + emitter.emitWorkerPoll(); + B.br(lc); +} + 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 d1f4ebcca..c1ea7bf35 100644 --- a/packages/compiler/src/backend/llvm/emitter.ts +++ b/packages/compiler/src/backend/llvm/emitter.ts @@ -3932,13 +3932,19 @@ export class LlEmitter { } } + /** Whether `loop` needs a cancellation poll: in a worker program, unless + * the loop already runs a pending test on every iteration, which observes + * a terminate request itself (a separate poll would only add work). */ + workerLoopPolls(loop: Extract): boolean { + return ( + this.mod.workers === true && !loopTestsEveryIteration(loop, (e) => this.callTestsPending(e)) + ); + } + workerLoopBudget( loop: Extract, ): string | null { - if (!this.mod.workers) return null; - // A loop that runs a pending test on every iteration observes a - // terminate request through it; a separate poll would only add work. - if (loopTestsEveryIteration(loop, (e) => this.callTestsPending(e))) return null; + if (!this.workerLoopPolls(loop)) return null; const slot = this.B.slot(); this.B.entryAllocas.push(`${slot} = alloca i32 ; worker cancellation budget`); this.B.line(`store i32 0, ptr ${slot}`); @@ -4056,14 +4062,19 @@ export class LlEmitter { B.condBr(due, poll, next); B.startBlock(poll); B.line(`store i32 63, ptr ${slot}`); - // Calls and throwing operations retain their own checks. Amortize the - // extra cancellation poll across bounded loop batches so scalar work - // does not cross the runtime ABI on every iteration. - this.emitPendingCheck(undefined, true); + this.emitWorkerPoll(); B.br(next); B.startBlock(next); } + /** 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. */ + emitWorkerPoll(): void { + this.emitPendingCheck(undefined, true); + } + /** Moves an already-evaluated value into the runtime's exception cell — * the `throw` statement's kind dispatch (stmts.ts's), shared with * every synthetic thrower. Ownership of a refcounted payload must have diff --git a/packages/compiler/src/backend/loop-polls.ts b/packages/compiler/src/backend/loop-polls.ts index 14c2b049c..c582a4974 100644 --- a/packages/compiler/src/backend/loop-polls.ts +++ b/packages/compiler/src/backend/loop-polls.ts @@ -106,3 +106,34 @@ export function loopTestsEveryIteration(loop: Loop, tested: (call: IrExpr) => bo } return false; } + +/** True when the body of `loop`, outside its nested loops and closures, + * reads or writes a typed-array element: the loops whose cancellation poll + * the emitter moves between blocks of iterations, so that LLVM can + * vectorize them. Other loops keep the budget in their body, which costs + * less in a loop of a few iterations. */ +export function loopAccessesTypedArrays(loop: Loop): boolean { + const inExpr = (e: IrExpr): boolean => { + if (e.kind === "bytesIntrinsic" && e.method === "get") return true; + let found = false; + everyExprChild( + e, + (child) => !(found = inExpr(child)), + (child) => !(found = inStmt(child)), + ); + return found; + }; + const inStmt = (s: IrStmt): boolean => { + if (s.kind === "bytesSet") return true; + if (s.kind === "while" || s.kind === "doWhile" || s.kind === "for" || s.kind === "forOf") + return false; + let found = false; + everyStmtChild( + s, + (child) => !(found = inExpr(child)), + (child) => !(found = inStmt(child)), + ); + return found; + }; + return loop.body.some(inStmt); +} diff --git a/packages/compiler/test/counted-loop-emission.test.ts b/packages/compiler/test/counted-loop-emission.test.ts index 7a3fc1f44..206a3c57e 100644 --- a/packages/compiler/test/counted-loop-emission.test.ts +++ b/packages/compiler/test/counted-loop-emission.test.ts @@ -1,4 +1,4 @@ -import { mkdtemp, readFile, rm } from "node:fs/promises"; +import { mkdtemp, readFile, rm, writeFile } from "node:fs/promises"; import { tmpdir } from "node:os"; import { join } from "node:path"; import { expect, test } from "vitest"; @@ -45,3 +45,81 @@ test("counted loops guard exact induction and preserve wide remainders on both p await rm(dir, { recursive: true, force: true }); } }); + +test("worker programs poll counted loops over typed arrays between blocks of iterations", async () => { + const dir = await mkdtemp(join(tmpdir(), "scriptc-counted-poll-")); + try { + const source = (workers: boolean): string => + [ + workers ? 'import { isMainThread, Worker } from "node:worker_threads";' : "", + "export function scale(a: Float64Array, b: Float64Array, n: number): void {", + " for (let i = 0; i < n; i++) b[i] = a[i]! * 2;", + "}", + "export function codes(a: Uint8Array): void {", + " for (let i = 0; i < a.length; ++i) a[i] = a[i] === 97 ? 0 : 3;", + "}", + "export function mix(n: number): number {", + " let t = 0;", + " for (let i = 0; i < n; i++) t = (t * 31 + i) % 1000003;", + " return t;", + "}", + "const a = new Float64Array(8).fill(1);", + "const b = new Float64Array(8);", + "scale(a, b, 8);", + "const c = new Uint8Array(8).fill(97);", + "codes(c);", + "console.log(b[7], mix(8), c[0]);", + workers ? "if (!isMainThread) new Worker(new URL(import.meta.url));" : "", + ].join("\n"); + const emit = async (workers: boolean): Promise => { + const entry = join(dir, `${workers ? "worker" : "plain"}.ts`); + await writeFile(entry, source(workers)); + const outPath = join(dir, `${workers ? "worker" : "plain"}.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")); + return emitLlvmModule(deserializeModule(await readFile(outPath, "utf8"))); + }; + // Every emitted version of a function (the number and int32 entries). + const bodies = (ll: string, name: string): string => + [ + ...ll.matchAll( + new RegExp( + `define internal [^\\n]+ @sc_b?f_${name}(?:\\.i32)?\\([^\\n]*\\n([\\s\\S]*?)\\n\\}`, + "g", + ), + ), + ] + .map((match) => match[1]!) + .join("\n"); + const workerModule = await emit(true); + const worker = bodies(workerModule, "scale"); + const blocks = worker.split(/\n(?=[\w.]+:)/); + // No budget in the body: the inner loop's latch only steps the counter, + // the poll runs after each block, and the block's bound is the nearer of + // the limit and 64 steps on. + expect(worker).not.toContain("worker cancellation budget"); + expect(worker).toContain("@llvm.sadd.sat.i64(i64 %"); + expect(worker).toMatch(/@llvm\.smin\.i64\(i64 %\w+, i64 %\w+\)/); + for (const latch of blocks.filter((block) => block.startsWith("loop.u"))) + expect(latch).not.toContain("scr_exc_flag"); + const polls = blocks.filter((block) => block.startsWith("loop.poll")); + expect(polls.length).toBeGreaterThan(0); + for (const poll of polls) expect(poll).toContain("@scr_exc_flag"); + // A counted loop without typed-array accesses keeps the budget. + // A loop to a typed array's length is strip-mined the same way, with the + // length as each block's limit. + const codes = bodies(workerModule, "codes"); + expect(codes).not.toContain("worker cancellation budget"); + expect(codes).toMatch(/@llvm\.uadd\.sat\.i64\(i64 %\w+, i64 64\)/); + expect(codes).toMatch(/@llvm\.umin\.i64\(i64 %\w+, i64 %\w+\)/); + expect(codes).toMatch(/^loop\.poll\d+:$/m); + const mix = bodies(workerModule, "mix"); + expect(mix).toContain("worker cancellation budget"); + expect(mix).not.toContain("loop.poll"); + const plain = bodies(await emit(false), "scale"); + expect(plain).not.toContain("loop.poll"); + expect(plain).not.toContain("@scr_exc_flag"); + } finally { + await rm(dir, { recursive: true, force: true }); + } +}); diff --git a/packages/compiler/test/ts7/baselines/order-parity/tests/corpus/native-worker-terminate-counted-loops.ts.json b/packages/compiler/test/ts7/baselines/order-parity/tests/corpus/native-worker-terminate-counted-loops.ts.json new file mode 100644 index 000000000..5f7c54075 --- /dev/null +++ b/packages/compiler/test/ts7/baselines/order-parity/tests/corpus/native-worker-terminate-counted-loops.ts.json @@ -0,0 +1,6 @@ +{ + "order": [ + "/tests/corpus/native-worker-terminate-counted-loops.ts" + ], + "diags": [] +} diff --git a/tests/corpus/native-worker-terminate-counted-loops.ts b/tests/corpus/native-worker-terminate-counted-loops.ts new file mode 100644 index 000000000..55972c4f6 --- /dev/null +++ b/tests/corpus/native-worker-terminate-counted-loops.ts @@ -0,0 +1,66 @@ +// worker.terminate() stops a worker in counted loops over a typed array +// with long bounds: ascending and descending counters, a step above one, a +// second counter that steps with the first, and loops to a typed array's +// length run again and again. +import { Worker, isMainThread, parentPort, workerData } from "node:worker_threads"; + +function ascending(values: Float64Array, limit: number): number { + let acc = 0; + for (let i = 0; i < limit; i++) { + acc = (acc + values[i & 1023]!) % 1000003; + if (i === 5000) parentPort!.postMessage("ready"); + } + return acc; +} + +function descending(values: Float64Array, limit: number): number { + let acc = 0; + for (let i = limit; i > 0; i -= 3) { + acc = (acc * 31 + i + values[i & 1023]!) % 1000003; + if (i === limit - 15000) parentPort!.postMessage("ready"); + } + return acc; +} + +function paired(values: Float64Array, limit: number): number { + let acc = 0; + for (let i = 0, j = 7; i < limit; i++, j += 2) { + acc = (acc + i * j + values[j & 1023]!) % 1000003; + if (i === 5000) parentPort!.postMessage("ready"); + } + return acc; +} + +function lengths(bytes: Uint8Array): number { + let acc = 0; + for (let round = 0; ; round++) { + for (let i = 0; i < bytes.length; ++i) bytes[i] = bytes[i] === 97 ? 0 : (bytes[i]! + round) & 127; + for (let i = 0; i < bytes.length; i++) acc = (acc + bytes[i]!) % 1000003; + if (round === 50) parentPort!.postMessage("ready"); + } + return acc; +} + +if (isMainThread) { + for (const kind of ["ascending", "descending", "paired", "lengths"]) { + await new Promise((resolve) => { + const worker = new Worker(new URL(import.meta.url), { workerData: kind }); + worker.once("message", () => { + void worker.terminate().then((code: number) => console.log(kind, "terminate", code)); + }); + worker.once("exit", (code: number) => { + console.log(kind, "exit", code); + resolve(); + }); + }); + } +} else { + const values = new Float64Array(1024).map((_, i) => i); + const limit = 1e15; + let result = 0; + if (workerData === "ascending") result = ascending(values, limit); + else if (workerData === "descending") result = descending(values, limit); + else if (workerData === "paired") result = paired(values, limit); + else result = lengths(new Uint8Array(1 << 16).fill(97)); + console.log("unexpected result", result); +} diff --git a/tests/harness/corpus-costs/native-worker-terminate-counted-loops.ts.json b/tests/harness/corpus-costs/native-worker-terminate-counted-loops.ts.json new file mode 100644 index 000000000..66a146275 --- /dev/null +++ b/tests/harness/corpus-costs/native-worker-terminate-counted-loops.ts.json @@ -0,0 +1 @@ +{ "plain": 3594, "san": 23593 }