Test a jammed loop's exit every iteration when its counter may not advance

The jammed loop tests its exit only when `(j & 3) === 3`. When the inner
counter starts at NaN, an infinity, or a value that `j + 1` leaves unchanged,
`j & 3` never becomes 3, and the loop ran past the point where every
accumulator was full, without end when its bound never ends it.

The every-fourth test is now used only when the inner loop starts at an
int32 literal: the counter then takes consecutive integer values, so its low
bits reach 3 within four iterations. From any other start, the jammed loop
tests its exit every iteration.

The unit test checks which starts get each test, and the corpus program adds
loops starting at -Infinity and -(2 ** 54) that end once both pixels escape.

packages/compiler/src/ir/saturating-loop-jams.test.ts

@@ -187,6 +187,27 @@ test("jams a constant-count loop whose inner loop exits once its accumulator is
   expect(posts.every((s) => s.kind === "assign" && s.localId === "res")).toBe(true);
 });
 
+test("tests the exit every iteration unless the inner loop starts at an int32 literal", () => {
+  const exitTest = (innerStart: IrExpr): IrExpr => {
+    const out = jamSaturatingLoops(moduleOf(nest({ innerStart })));
+    const inner = loops(out.functions[0]!.body)[0]!;
+    const last = inner.body[inner.body.length - 1]!;
+    if (last.kind !== "if") throw new Error("expected an exit test");
+    return last.cond;
+  };
+  const everyFourth = (cond: IrExpr): boolean => {
+    if (cond.kind !== "logical" || cond.left.kind !== "bin") return false;
+    const low = cond.left.left;
+    return low.kind === "bin" && low.op === "&";
+  };
+  expect(everyFourth(exitTest(num(-5)))).toBe(true);
+  expect(everyFourth(exitTest(num(0x7fffffff)))).toBe(true);
+  // `j` may be NaN, an infinity or a value `j + 1` leaves unchanged.
+  expect(everyFourth(exitTest(ref("start")))).toBe(false);
+  expect(everyFourth(exitTest(num(2 ** 31)))).toBe(false);
+  expect(everyFourth(exitTest(num(-(2 ** 31) - 1)))).toBe(false);
+});
+
 test("keeps loops whose exits or values do not allow running iterations together", () => {
   const reject = (body: IrStmt[]): void => {
     const mod = moduleOf(body);

packages/compiler/src/ir/saturating-loop-jams.ts

@@ -332,8 +332,9 @@ function matchJam(
  * overlap their dependence chains.
  *
  * For the same reason the jammed loop may run past the iteration where every
- * accumulator is full: it tests that every fourth iteration (`j & 3` is 3)
- * and keeps the inner loop's own bound. The copies' statements alternate,
+ * accumulator is full: when the inner loop starts at an int32 literal, it
+ * tests that every fourth iteration (`j & 3` is 3), and it keeps the inner
+ * loop's own bound. The copies' statements alternate,
  * and each local's copies are added next to each other, so the vectorizer
  * packs one copy per lane. */
 export function jamSaturatingLoops(mod: IrModule): IrModule {
@@ -382,7 +383,8 @@ function emitJam(
   const { inner, pre, post, values, counter } = jam;
   const loc = inner.loc;
   const own = declared([...pre, ...inner.body]);
-  const j = (inner.init as IrStmt & { kind: "varDecl" }).localId;
+  const init = inner.init as IrStmt & { kind: "varDecl" };
+  const j = init.localId;
   own.delete(j);
   // Each local's copies are added next to each other: the vectorizer pairs
   // neighboring values of the loop's header, so one copy per lane.
@@ -449,15 +451,23 @@ function emitJam(
   };
   // `(j & 3) === 3`
   const fourth: IrExpr = { kind: "bin", op: "===", left: counter3, right: three, type: BOOL, loc };
+  // `j` reaches a value whose low bits are 3 within four iterations when it
+  // starts at an int32 literal: it then takes consecutive integer values,
+  // which stay exact for more than 2^52 iterations. From another start (NaN,
+  // an infinity, a value where `j + 1` is `j`), the exit is tested every
+  // iteration.
+  const start = intLit(init.init);
+  const everyFourth = start !== null && start >= -0x80000000 && start <= 0x7fffffff;
   const jammed: IrStmt = {
     ...inner,
     body: [
-      // The copies take turns statement by statement; the exit is tested
-      // every fourth iteration.
+      // The copies take turns statement by statement.
       ...copies[0]!.body.flatMap((_, k) => copies.map((copy) => copy.body[k]!)),
       {
         kind: "if",
-        cond: { kind: "logical", op: "&&", left: fourth, right: allFull!, type: BOOL, loc },
+        cond: everyFourth
+          ? { kind: "logical", op: "&&", left: fourth, right: allFull!, type: BOOL, loc }
+          : allFull!,
         then: [{ kind: "break", loc }],
         else_: null,
         loc,

tests/corpus/saturating-loop-jams.ts

@@ -142,6 +142,31 @@ function shifted(initial: Float64Array): number {
   return res;
 }
 
+// The inner counter starts where `j++` leaves it unchanged or where its
+// low bits never reach 3: the loop ends once both pixels escape.
+function stuck(initial: Float64Array, start: number): number {
+  let res = 0;
+  for (let i = 0; i < 4; i += 2) {
+    let zr1 = initial[i]!;
+    let zr2 = initial[i + 1]!;
+    let b = 0;
+    for (let j = start; j < 40; j++) {
+      zr1 = zr1 * zr1 + 0.2;
+      if (zr1 > 2) {
+        b |= 2;
+        if (b === 3) break;
+      }
+      zr2 = zr2 * zr2 + 0.2;
+      if (zr2 > 2) {
+        b |= 1;
+        if (b === 3) break;
+      }
+    }
+    res = res * 4 + b;
+  }
+  return res;
+}
+
 const size = 64;
 const initial = new Float64Array(size);
 for (let x = 0; x < size; x++) initial[x] = (2 * x) / size - 1.5;
@@ -162,3 +187,5 @@ console.log(row(initial, 0.5, (size >> 3) + 1, 50).slice(-2));
 const r = new Float64Array([1.2, 0.5, 1.3, 0.9]);
 const s = new Float64Array([0.1, 0.1, 0.73, 0.74]);
 console.log(resets(q), resets(r), skips(q), skips(s), shifted(q), shifted(r));
+const e = new Float64Array([1.2, 1.5, 3, 0.9]);
+console.log(stuck(e, -Infinity), stuck(e, -(2 ** 54)), stuck(e, 0.5), stuck(e, -7));