Version loops around nested loops only for accesses inside them (#2364)

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))),

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<number, Node[]>();
+  // Accesses inside the loop's nested loops, and inside its nested while
+  // loops.
+  const nested = new Set<Node>();
+  const nestedWhile = new Set<Node>();
+  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)!;

packages/compiler/test/ts7/baselines/order-parity/tests/corpus/typed-array-nested-loop-versions.ts.json

@@ -0,0 +1,6 @@
+{
+ "order": [
+  "<repo>/tests/corpus/typed-array-nested-loop-versions.ts"
+ ],
+ "diags": []
+}

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);
+}

tests/harness/corpus-costs/typed-array-nested-loop-versions.ts.json

@@ -0,0 +1 @@
+{ "plain": 6500, "san": 8800 }