From ca1a6b5bed3ffd8c195e7d7c525b589f6e2b1a2c Mon Sep 17 00:00:00 2001 From: Chris Tate Date: Mon, 26 Oct 2026 19:05:12 -0500 Subject: [PATCH] Version loops around nested loops only for accesses inside them A typed-array loop version runs a checked copy and one or two fast copies of the loop, and each copy repeats every nested loop together with the nested loop's own version. A loop whose guard proves only accesses outside its nested loops, which run once per pass, now runs as written when a nested loop accesses typed arrays; the nested loops still version their own accesses. The loop versions as before when the guard proves the bounds of an access inside a nested loop, or, in a worker program, the ownership of an access inside a nested while loop, which never versions for ownership alone. Spectral-norm's paired row sums no longer copy their column loop for the row loop's two stores. The worker program's au and atu functions are 102 KB of LLVM text instead of 236 KB each, and the program 519 KB instead of 786 KB. Corpus typed-array-nested-loop-versions covers row loops that store past the end or at negative, fractional and NaN rows, views the nested loop writes, and a nested while loop. --- .../backend/llvm/bytes-loop-versions.test.ts | 49 +++++++++++++ .../src/backend/llvm/bytes-loop-versions.ts | 37 ++++++++-- .../typed-array-nested-loop-versions.ts.json | 6 ++ .../typed-array-nested-loop-versions.ts | 68 +++++++++++++++++++ .../typed-array-nested-loop-versions.ts.json | 1 + 5 files changed, 157 insertions(+), 4 deletions(-) create mode 100644 packages/compiler/test/ts7/baselines/order-parity/tests/corpus/typed-array-nested-loop-versions.ts.json create mode 100644 tests/corpus/typed-array-nested-loop-versions.ts create mode 100644 tests/harness/corpus-costs/typed-array-nested-loop-versions.ts.json diff --git a/packages/compiler/src/backend/llvm/bytes-loop-versions.test.ts b/packages/compiler/src/backend/llvm/bytes-loop-versions.test.ts index 742d256b7..76f229215 100644 --- a/packages/compiler/src/backend/llvm/bytes-loop-versions.test.ts +++ b/packages/compiler/src/backend/llvm/bytes-loop-versions.test.ts @@ -222,6 +222,55 @@ test("negative factors, rebinding, captures and outer loops are not versioned", expect(match(fn([outer]), inner)).not.toBeNull(); }); +test("loops with nested loops version only for accesses inside them", () => { + // for (let i = 0; i < 64; i++) { + // let t = 0; + // for (let j = 0; j < n; j++) t = t + src[j]; + // out[i] = t; + // } + const sum = (): IrStmt[] => [ + { kind: "varDecl", localId: "t", init: num(0), loc }, + counted( + [ + { + kind: "assign", + localId: "t", + value: bin("+", ref("t"), get(ref("src", f64s), ref("j"))), + loc, + }, + ], + ref("n"), + "j", + ), + ]; + const extra = [local("t", F64, true)]; + // The guard would prove only `out[i]`, which runs once per row. + const rows = counted([...sum(), set(ref("out", f64s), ref("i"), ref("t"))], num(64)); + expect(match(fn([rows], extra), rows)).toBeNull(); + expect(match(fn([rows], extra), rows.body[1] as IrStmt & { kind: "for" })).not.toBeNull(); + // The guard proves `out[i]` inside a nested loop too. + const inside = counted( + [...sum(), counted([set(ref("out", f64s), ref("i"), ref("t"))], num(2), "j")], + num(64), + ); + expect(match(fn([inside], extra), inside)?.bounds.size).toBe(1); + // A loop whose nested loops access no typed array versions for its own + // accesses. + const scalar = counted( + [ + { kind: "varDecl", localId: "t", init: num(0), loc }, + counted( + [{ kind: "assign", localId: "t", value: bin("+", ref("t"), ref("j")), loc }], + ref("n"), + "j", + ), + set(ref("out", f64s), ref("i"), ref("t")), + ], + num(64), + ); + expect(match(fn([scalar], extra), scalar)?.bounds.size).toBe(1); +}); + test("emitted loops version the body and tag typed-array memory", () => { const loop = counted([ set(ref("out", f64s), ref("i"), bin("*", get(ref("src", f64s), ref("i")), num(2))), diff --git a/packages/compiler/src/backend/llvm/bytes-loop-versions.ts b/packages/compiler/src/backend/llvm/bytes-loop-versions.ts index ccf217f9f..d0bef7fd6 100644 --- a/packages/compiler/src/backend/llvm/bytes-loop-versions.ts +++ b/packages/compiler/src/backend/llvm/bytes-loop-versions.ts @@ -196,9 +196,10 @@ const INVARIANT_OPS = new Set(["+", "-", "*", "&", "|", "^", "<<", ">>", ">>>"]) * and writes thread-owned views directly. * * Only small synchronous loops are versioned (nested loops included), and - * the emitter bounds how deeply versions nest. The checked loop runs - * whenever the guard fails, and keeps every out-of-bounds read, ignored - * write and NaN-index behavior. */ + * the emitter bounds how deeply versions nest. A loop whose nested loops + * access typed arrays versions only for an access inside them. The checked + * loop runs whenever the guard fails, and keeps every out-of-bounds read, + * ignored write and NaN-index behavior. */ export function matchBytesLoopVersion( loop: IrStmt & { kind: "for" | "while" }, counted: IntegerCountedForLoop | null, @@ -492,7 +493,15 @@ export function matchBytesLoopVersion( const pairedCounters = paired !== null && counted === null ? new Set([paired.low.localId, paired.high.localId]) : null; const pairedNodes = new Map(); + // Accesses inside the loop's nested loops, and inside its nested while + // loops. + const nested = new Set(); + const nestedWhile = new Set(); + let loops = 0, + whiles = 0; const access = (node: Node, array: IrExpr, index: IrExpr, width: number): void => { + if (loops > 0) nested.add(node); + if (whiles > 0) nestedWhile.add(node); const ref = root(array); if (!ref) return; if (facts.workers && ref.type.kind === "bytes" && !facts.owned.has(node)) { @@ -603,7 +612,15 @@ export function matchBytesLoopVersion( access(s, s.arr, s.index, 1); scope(s, s.arr, true); } - return everyStmtChild(s, visitExpr, visitStmt); + const isLoop = + s.kind === "for" || s.kind === "forOf" || s.kind === "while" || s.kind === "doWhile"; + const isWhile = s.kind === "while" || s.kind === "doWhile"; + if (isLoop) loops++; + if (isWhile) whiles++; + const result = everyStmtChild(s, visitExpr, visitStmt); + if (isLoop) loops--; + if (isWhile) whiles--; + return result; }; loop.body.every(visitStmt); if (loop.kind === "while") visitExpr(loop.cond); @@ -714,6 +731,18 @@ export function matchBytesLoopVersion( used.add(receiver); for (const node of nodes) bounds.add(node); } + // A loop whose nested loops access typed arrays versions only when the + // guard proves the bounds of one of those accesses (or, in a worker + // program, the ownership of one inside a nested while loop, which never + // versions for ownership alone). Otherwise the guard covers accesses that + // run once per pass, while each copy of the body repeats every nested + // loop and each nested loop's own version. + if ( + nested.size > 0 && + ![...bounds].some((node) => nested.has(node)) && + ![...ownedNodes].some((node) => nestedWhile.has(node)) + ) + return null; const counterChecks: CounterCheck[] = []; for (const [id, min] of counterMins) { const info = whileAccesses!.counters.get(id)!; diff --git a/packages/compiler/test/ts7/baselines/order-parity/tests/corpus/typed-array-nested-loop-versions.ts.json b/packages/compiler/test/ts7/baselines/order-parity/tests/corpus/typed-array-nested-loop-versions.ts.json new file mode 100644 index 000000000..66c967b58 --- /dev/null +++ b/packages/compiler/test/ts7/baselines/order-parity/tests/corpus/typed-array-nested-loop-versions.ts.json @@ -0,0 +1,6 @@ +{ + "order": [ + "/tests/corpus/typed-array-nested-loop-versions.ts" + ], + "diags": [] +} diff --git a/tests/corpus/typed-array-nested-loop-versions.ts b/tests/corpus/typed-array-nested-loop-versions.ts new file mode 100644 index 000000000..19e408d30 --- /dev/null +++ b/tests/corpus/typed-array-nested-loop-versions.ts @@ -0,0 +1,68 @@ +// Loops whose own typed-array accesses run once per pass around nested +// loops that read typed arrays: row stores past the end, reads at NaN, +// negative and fractional rows, rows that read a typed array the nested +// loop writes, and a nested loop that reads the row's own result. +function rowSums(u: Float64Array, out: Float64Array, start: number, end: number): void { + for (let i = start; i < end; i++) { + let t = out[i] ?? -1; + for (let j = 0; j < u.length; j++) t += u[j]! * (i + 1); + out[i] = t; + } +} + +function scaled(u: Float64Array, scale: Float64Array, out: Float64Array, rows: number): void { + for (let i = 0; i < rows; i++) { + const s = scale[i * 2] ?? 0.5; + for (let j = 0; j < u.length; j++) out[j] = out[j]! + u[j]! * s; + out[i] = out[i]! - 1; + } +} + +function chained(v: Int32Array, rows: number): number { + let total = 0; + for (let i = 0; i < rows; i++) { + let k = 0; + while (k < v.length) { + v[k] = v[k]! + (v[i] ?? 7); + k++; + } + total += v[i + 1] ?? -3; + } + return total; +} + +const show = (label: string, v: Float64Array | Int32Array): void => { + const parts: string[] = []; + for (let k = 0; k < v.length; k++) parts.push(Object.is(v[k], -0) ? "-0" : String(v[k])); + console.log(label, parts.join(" ")); +}; + +const u = new Float64Array([0.5, -1.25, 3, 7.75]); +for (const [start, end] of [ + [0, 4], + [0, 7], + [2, 9], + [-2, 3], + [0.5, 3.5], + [NaN, 3], + [3, 1], +]) { + const out = new Float64Array(5).fill(2); + rowSums(u, out, start!, end!); + show(`rows ${start} ${end}`, out); +} + +for (const rows of [0, 1, 3, 6]) { + const out = new Float64Array(4).fill(1); + scaled(u, new Float64Array([2, -1, 0.25, 4, 8]), out, rows); + show(`scaled ${rows}`, out); + const same = new Float64Array([1, 2, 3, 4]); + scaled(same, same, same, rows); + show(`aliased ${rows}`, same); +} + +for (const rows of [0, 2, 4, 9]) { + const v = new Int32Array([1, -2, 3, 2147483647]); + const total = chained(v, rows); + show(`chained ${rows} ${total}`, v); +} diff --git a/tests/harness/corpus-costs/typed-array-nested-loop-versions.ts.json b/tests/harness/corpus-costs/typed-array-nested-loop-versions.ts.json new file mode 100644 index 000000000..1bee89d1a --- /dev/null +++ b/tests/harness/corpus-costs/typed-array-nested-loop-versions.ts.json @@ -0,0 +1 @@ +{ "plain": 6500, "san": 8800 }