diff --git a/packages/compiler/src/backend/llvm/emit-statements.ts b/packages/compiler/src/backend/llvm/emit-statements.ts index 2efe2cf3a..d7fccb2d4 100644 --- a/packages/compiler/src/backend/llvm/emit-statements.ts +++ b/packages/compiler/src/backend/llvm/emit-statements.ts @@ -1216,14 +1216,15 @@ export function emitStmtBody(emitter: LlEmitter, s: IrStmt): void { // 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. + // array's length is fixed, so each block can take it as its limit when + // the loop does not assign the array's binding. const stripMined = capturedLets.length === 0 && integerSlot !== null && ((countedLoop !== null && countedLimit !== null && Math.abs(countedLoop.step) * WORKER_POLL_BLOCK <= 2 ** 52) || - (integerLoop !== null && integerArrayLoop === null)) && + (integerLoop !== null && integerArrayLoop === null && integerLoop.stableReceiver)) && emitter.workerLoopPolls(s) && loopAccessesTypedArrays(s); const workerBudget = stripMined ? null : emitter.workerLoopBudget(s); diff --git a/packages/compiler/src/ir/integer-loops.ts b/packages/compiler/src/ir/integer-loops.ts index 7cc49b7a0..ea866a134 100644 --- a/packages/compiler/src/ir/integer-loops.ts +++ b/packages/compiler/src/ir/integer-loops.ts @@ -12,6 +12,9 @@ import { preservesArrayLength } from "./array-length-stability.js"; export interface IntegerBytesForLoop { localId: string; limitReceiver: IrExpr; + /** The receiver is a local of the function that the body never assigns + * (a typed array's length is then the same on every iteration). */ + stableReceiver: boolean; } export interface IntegerCountedForLoop { @@ -682,7 +685,12 @@ export function matchIntegerBytesForLoop( if (!isUnitIncrement(stmt.update, init.localId)) return null; if (writesLocal(stmt.body, init.localId)) return null; - return { localId: init.localId, limitReceiver: cond.right.receiver }; + return { + localId: init.localId, + limitReceiver: cond.right.receiver, + stableReceiver: + receiverLocal !== undefined && !writesLocal(stmt.body, cond.right.receiver.localId), + }; } /** Array lengths are uint32, so zero-based loops and nested `j = i + 1` @@ -729,5 +737,11 @@ export function matchIntegerArrayForLoop( writesLocal(stmt.body, local.id) ) return null; - return { localId: local.id, limitReceiver: cond.right.receiver }; + return { + localId: local.id, + limitReceiver: cond.right.receiver, + stableReceiver: + locals.has(cond.right.receiver.localId) && + !writesLocal(stmt.body, cond.right.receiver.localId), + }; } diff --git a/packages/compiler/test/counted-loop-emission.test.ts b/packages/compiler/test/counted-loop-emission.test.ts index 206a3c57e..f617c6ab4 100644 --- a/packages/compiler/test/counted-loop-emission.test.ts +++ b/packages/compiler/test/counted-loop-emission.test.ts @@ -58,6 +58,14 @@ test("worker programs poll counted loops over typed arrays between blocks of ite "export function codes(a: Uint8Array): void {", " for (let i = 0; i < a.length; ++i) a[i] = a[i] === 97 ? 0 : 3;", "}", + "export function shrink(a: Uint8Array): number {", + " let t = 0;", + " for (let i = 0; i < a.length; ++i) {", + " t += a[i]!;", + " if (i === 2) a = a.subarray(0, 3);", + " }", + " return t;", + "}", "export function mix(n: number): number {", " let t = 0;", " for (let i = 0; i < n; i++) t = (t * 31 + i) % 1000003;", @@ -68,7 +76,7 @@ test("worker programs poll counted loops over typed arrays between blocks of ite "scale(a, b, 8);", "const c = new Uint8Array(8).fill(97);", "codes(c);", - "console.log(b[7], mix(8), c[0]);", + "console.log(b[7], mix(8), c[0], shrink(new Uint8Array(8).fill(1)));", workers ? "if (!isMainThread) new Worker(new URL(import.meta.url));" : "", ].join("\n"); const emit = async (workers: boolean): Promise => { @@ -113,6 +121,11 @@ test("worker programs poll counted loops over typed arrays between blocks of ite 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); + // A loop that assigns its array reads the new length on the next + // iteration: it keeps the budget. + const shrink = bodies(workerModule, "shrink"); + expect(shrink).toContain("worker cancellation budget"); + expect(shrink).not.toMatch(/^loop\.poll\d+:$/m); const mix = bodies(workerModule, "mix"); expect(mix).toContain("worker cancellation budget"); expect(mix).not.toContain("loop.poll"); diff --git a/packages/compiler/test/ts7/baselines/order-parity/tests/corpus/native-worker-length-loop-reassign.ts.json b/packages/compiler/test/ts7/baselines/order-parity/tests/corpus/native-worker-length-loop-reassign.ts.json new file mode 100644 index 000000000..f5230bace --- /dev/null +++ b/packages/compiler/test/ts7/baselines/order-parity/tests/corpus/native-worker-length-loop-reassign.ts.json @@ -0,0 +1,6 @@ +{ + "order": [ + "/tests/corpus/native-worker-length-loop-reassign.ts" + ], + "diags": [] +} diff --git a/tests/corpus/native-worker-length-loop-reassign.ts b/tests/corpus/native-worker-length-loop-reassign.ts new file mode 100644 index 000000000..7ec152b16 --- /dev/null +++ b/tests/corpus/native-worker-length-loop-reassign.ts @@ -0,0 +1,58 @@ +// Loops to a typed array's length in a program that starts workers, where +// the body assigns the loop's array: the condition reads the new array's +// length on the next iteration, shorter or longer. +import { Worker, isMainThread, parentPort } from "node:worker_threads"; + +function shrink(values: Float64Array): string { + let total = 0; + let steps = 0; + for (let i = 0; i < values.length; i++) { + total += values[i]!; + steps++; + if (i === 100) values = values.subarray(0, 50); + } + return `shrink ${steps} ${total}`; +} + +function grow(values: Uint8Array, longer: Uint8Array): string { + let total = 0; + let steps = 0; + for (let i = 0; i < values.length; i++) { + total += values[i]!; + steps++; + if (i === 10) values = longer; + } + return `grow ${steps} ${total}`; +} + +function swap(a: Int32Array, b: Int32Array): string { + let total = 0; + let steps = 0; + for (let i = 0; i < a.length; i++) { + total += a[i]!; + steps++; + if (i % 70 === 69) { + const t = a; + a = b; + b = t; + } + } + return `swap ${steps} ${total}`; +} + +function run(): string[] { + const values = new Float64Array(300).map((_, i) => i * 0.5); + const short = new Uint8Array(20).fill(3); + const long = new Uint8Array(200).map((_, i) => i & 7); + const a = new Int32Array(100).map((_, i) => i); + const b = new Int32Array(160).map((_, i) => -i); + return [shrink(values), grow(short, long), swap(a, b)]; +} + +if (isMainThread) { + for (const line of run()) console.log(line); + const worker = new Worker(new URL(import.meta.url)); + worker.on("message", (lines: string[]) => { + for (const line of lines) console.log("worker", line); + }); +} else parentPort!.postMessage(run()); diff --git a/tests/harness/corpus-costs/native-worker-length-loop-reassign.ts.json b/tests/harness/corpus-costs/native-worker-length-loop-reassign.ts.json new file mode 100644 index 000000000..5ee17db16 --- /dev/null +++ b/tests/harness/corpus-costs/native-worker-length-loop-reassign.ts.json @@ -0,0 +1 @@ +{ "plain": 3158, "san": 20404 }