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 }