Pack jammed saturating loops one copy per lane and test their exit every fourth iteration

The jammed loop of jamSaturatingLoops now tests whether every accumulator
is full only when `j & 3` is 3, and keeps the inner loop's own bound. Past
the iteration where every accumulator is full, the copies' ORs leave the
accumulators full and every other value they compute is dead, so running
up to three more iterations changes no result.

The copies' statements now alternate in the jammed loop, and each local's
copies are added next to each other. The vectorizer pairs neighboring
values of the loop's header, which were a copy's real and imaginary parts
in mandelbrot; it now packs the same value of two copies into one vector,
without lane shuffles.

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

@@ -162,9 +162,24 @@ test("jams a constant-count loop whose inner loop exits once its accumulator is
   expect(breaks).toBe(1);
   const last = inner.body[inner.body.length - 1]!;
   expect(last.kind).toBe("if");
-  // Each copy keeps its own accumulator and lanes; `post` runs in order.
+  // It is tested every fourth iteration: `(j & 3) === 3 && ...`.
+  if (last.kind !== "if" || last.cond.kind !== "logical") throw new Error("expected an && exit");
+  const fourth = last.cond.left;
+  expect(fourth.kind === "bin" && fourth.op === "===" && fourth.left.kind === "bin").toBe(true);
+  if (fourth.kind === "bin" && fourth.left.kind === "bin") {
+    expect(fourth.left.op).toBe("&");
+    expect(fourth.left.left.kind === "varRef" && fourth.left.left.localId).toBe("j");
+  }
+  // The copies take turns statement by statement.
+  const firstDecls = inner.body.slice(0, 4);
+  expect(firstDecls.every((s) => s.kind === "varDecl")).toBe(true);
+  expect(new Set(firstDecls.map((s) => (s.kind === "varDecl" ? s.localId : ""))).size).toBe(4);
+  // Each copy keeps its own accumulator and lanes; each local's copies sit
+  // next to each other; `post` runs in order.
   const added = out.functions[0]!.locals.filter((local) => local.id.startsWith("%jam."));
   expect(added.filter((local) => local.name === "%jam.b").length).toBe(4);
+  for (let at = 0; at < added.length; at += 4)
+    expect(new Set(added.slice(at, at + 4).map((local) => local.name)).size).toBe(1);
   const block = body[1]!;
   expect(block.kind).toBe("block");
   if (block.kind !== "block") return;

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

@@ -329,7 +329,13 @@ function matchJam(
  * elements that answer NaN out of range, and write their own locals. `pre`
  * and the inner loop read nothing `post` writes, so running later `pre`s
  * and inner loops first changes no value. Independent iterations then
- * overlap their dependence chains. */
+ * 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,
+ * 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 {
   let changed = false;
   const functions = mod.functions.map((fn): IrFunction => {
@@ -378,9 +384,13 @@ function emitJam(
   const own = declared([...pre, ...inner.body]);
   const j = (inner.init as IrStmt & { kind: "varDecl" }).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.
+  const renames = values.map(() => new Map<string, string>());
+  for (const id of own)
+    for (const [copy, rename] of renames.entries()) rename.set(id, fresh(locals.get(id)!, copy));
   const copies = values.map((value, copy) => {
-    const rename = new Map<string, string>();
-    for (const id of own) rename.set(id, fresh(locals.get(id)!, copy));
+    const rename = renames[copy]!;
     const rewrite = (list: readonly IrStmt[]): IrStmt[] =>
       transformStmtList(list as IrStmt[], {
         expr: (e) => {
@@ -428,11 +438,30 @@ function emitJam(
         ? full
         : { kind: "logical", op: "&&", left: allFull, right: full, type: BOOL, loc };
   }
+  const three: IrExpr = { kind: "numLit", value: 3, type: F64, loc };
+  const counter3: IrExpr = {
+    kind: "bin",
+    op: "&",
+    left: { kind: "varRef", localId: j, type: F64, loc },
+    right: three,
+    type: F64,
+    loc,
+  };
+  // `(j & 3) === 3`
+  const fourth: IrExpr = { kind: "bin", op: "===", left: counter3, right: three, type: BOOL, loc };
   const jammed: IrStmt = {
     ...inner,
     body: [
-      ...copies.flatMap((copy) => copy.body),
-      { kind: "if", cond: allFull!, then: [{ kind: "break", loc }], else_: null, loc },
+      // The copies take turns statement by statement; the exit is tested
+      // every fourth iteration.
+      ...copies[0]!.body.flatMap((_, k) => copies.map((copy) => copy.body[k]!)),
+      {
+        kind: "if",
+        cond: { kind: "logical", op: "&&", left: fourth, right: allFull!, type: BOOL, loc },
+        then: [{ kind: "break", loc }],
+        else_: null,
+        loc,
+      },
     ],
   };
   return {