Run short constant-count loops together when their inner exits saturate (#2493)

packages/compiler/src/backend/llvm/emitter.ts

@@ -96,6 +96,7 @@ import {
 } from "../../ir/inline-arrays.js";
 import { inlineArraysApply, programLifetimeClassOf } from "./inline-arrays.js";
 import { pairRowSums } from "../../ir/row-pairs.js";
+import { jamSaturatingLoops } from "../../ir/saturating-loop-jams.js";
 import {
   allocatedLiteral,
   fillAllocationLength,
@@ -763,16 +764,18 @@ function moduleEmitter(input: IrModule, options: LlvmTargetOptions): LlEmitter {
           specializeCallbackCalls(
             specializeNumericCalls(
               pairRowSums(
-                inlineArraysPass(
-                  options,
-                  inlineExpressionFunctions(
-                    foldConstantGlobals(
-                      copyStableCaptures(
-                        resolveSelfCaptures(
-                          hoistFinalFieldReads(
-                            freezeInitOnlyGlobals(
-                              holdCharLocalsAsUnits(
-                                dropUnwrittenPropertyBags(runPairedWhileLoopsAsFor(mod)),
+                jamSaturatingLoops(
+                  inlineArraysPass(
+                    options,
+                    inlineExpressionFunctions(
+                      foldConstantGlobals(
+                        copyStableCaptures(
+                          resolveSelfCaptures(
+                            hoistFinalFieldReads(
+                              freezeInitOnlyGlobals(
+                                holdCharLocalsAsUnits(
+                                  dropUnwrittenPropertyBags(runPairedWhileLoopsAsFor(mod)),
+                                ),
                               ),
                             ),
                           ),

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

@@ -0,0 +1,220 @@
+import { expect, test } from "vitest";
+import { BOOL, F64, VOID, type IrExpr, type IrFunction, type IrModule, type IrStmt } from "./ir.js";
+import { everyStmtList } from "./traverse.js";
+import { jamSaturatingLoops } from "./saturating-loop-jams.js";
+
+const loc = { file: "saturating-loop-jams.ts", start: 0, end: 0 };
+const BYTES = { kind: "bytes", elem: "f64" } as const;
+const num = (value: number): IrExpr => ({ loc, kind: "numLit", value, type: F64 });
+const ref = (localId: string): IrExpr => ({ loc, kind: "varRef", localId, type: F64 });
+type Op = "+" | "*" | "-" | "<" | ">" | "|" | "===" | "<<";
+const bin = (op: Op, left: IrExpr, right: IrExpr): IrExpr => ({
+  loc,
+  kind: "bin",
+  op,
+  left,
+  right,
+  type: op === "<" || op === ">" || op === "===" ? BOOL : F64,
+});
+const decl = (id: string, init: IrExpr): IrStmt => ({ loc, kind: "varDecl", localId: id, init });
+const set = (id: string, value: IrExpr): IrStmt => ({ loc, kind: "assign", localId: id, value });
+const stepBy = (id: string, by: number): IrStmt => set(id, bin("+", ref(id), num(by)));
+const loop = (
+  id: string,
+  end: IrExpr,
+  by: number,
+  body: IrStmt[],
+  start: IrExpr = num(0),
+): IrStmt & { kind: "for" } => ({
+  loc,
+  kind: "for",
+  init: decl(id, start),
+  cond: bin("<", ref(id), end),
+  update: stepBy(id, by),
+  body,
+});
+const exitWhen = (acc: string, full: number): IrStmt => ({
+  loc,
+  kind: "if",
+  cond: bin("===", ref(acc), num(full)),
+  then: [{ loc, kind: "break" }],
+  else_: null,
+});
+const get = (index: IrExpr): IrExpr => ({
+  loc,
+  kind: "bytesIntrinsic",
+  method: "get",
+  receiver: { loc, kind: "varRef", localId: "cr", type: BYTES },
+  args: [index],
+  type: F64,
+  invalidNaN: true,
+});
+
+/** One lane of mandelbrot's escape test: `z = z * z + c`, `b |= bit` and
+ * an exit once `b` is full. */
+function lane(z: string, c: string, bit: number, full = 3, extra: IrStmt[] = []): IrStmt[] {
+  return [
+    decl(`${z}.t`, bin("*", ref(z), ref(z))),
+    set(z, bin("+", ref(`${z}.t`), ref(c))),
+    {
+      loc,
+      kind: "if",
+      cond: bin(">", ref(`${z}.t`), num(4)),
+      then: [set("b", bin("|", ref("b"), num(bit))), exitWhen("b", full), ...extra],
+      else_: null,
+    },
+  ];
+}
+
+function nest(
+  options: {
+    full?: number;
+    post?: IrStmt[];
+    preExtra?: IrStmt[];
+    innerExtra?: IrStmt[];
+    innerStart?: IrExpr;
+    innerEnd?: IrExpr;
+    copies?: number;
+  } = {},
+): IrStmt[] {
+  const post = options.post ?? [set("res", bin("|", bin("<<", ref("res"), num(2)), ref("b")))];
+  return [
+    decl("res", num(0)),
+    loop("i", num(2 * (options.copies ?? 4)), 2, [
+      decl("c1", get(ref("i"))),
+      decl("c2", get(bin("+", ref("i"), num(1)))),
+      decl("z1", ref("c1")),
+      decl("z2", ref("c2")),
+      decl("b", num(0)),
+      ...(options.preExtra ?? []),
+      loop(
+        "j",
+        options.innerEnd ?? num(50),
+        1,
+        [
+          ...lane("z1", "c1", 2, options.full),
+          ...lane("z2", "c2", 1, options.full),
+          ...(options.innerExtra ?? []),
+        ],
+        options.innerStart,
+      ),
+      ...post,
+    ]),
+  ];
+}
+
+function moduleOf(body: IrStmt[]): IrModule {
+  const ids = new Set<string>();
+  everyStmtList(body, {
+    expr: (e) => {
+      if (e.kind === "varRef" && e.localId !== "cr") ids.add(e.localId);
+      return true;
+    },
+    stmt: (s) => {
+      if (s.kind === "varDecl" || s.kind === "assign") ids.add(s.localId);
+      return true;
+    },
+  });
+  const fn: IrFunction = {
+    name: "render",
+    params: [{ localId: "cr", name: "cr", type: BYTES }],
+    locals: [
+      { id: "cr", name: "cr", type: BYTES, mutable: false },
+      ...[...ids].map((id) => ({ id, name: id, type: F64, mutable: true })),
+    ],
+    body,
+    returnType: VOID,
+    loc,
+  };
+  return { irVersion: 17, sourceFile: loc.file, entry: "render", functions: [fn] } as IrModule;
+}
+
+function loops(body: IrStmt[]): (IrStmt & { kind: "for" })[] {
+  const found: (IrStmt & { kind: "for" })[] = [];
+  everyStmtList(body, {
+    expr: () => true,
+    stmt: (s) => {
+      if (s.kind === "for") found.push(s);
+      return true;
+    },
+  });
+  return found;
+}
+
+test("jams a constant-count loop whose inner loop exits once its accumulator is full", () => {
+  const mod = moduleOf(nest());
+  const out = jamSaturatingLoops(mod);
+  expect(out).not.toBe(mod);
+  const body = out.functions[0]!.body;
+  // The outer loop is gone: one inner loop runs the four iterations.
+  const remaining = loops(body);
+  expect(remaining.length).toBe(1);
+  const inner = remaining[0]!;
+  // Its only exit tests every copy's accumulator, at the end of the body.
+  let breaks = 0;
+  everyStmtList(inner.body, {
+    expr: () => true,
+    stmt: (s) => {
+      if (s.kind === "break") breaks++;
+      return true;
+    },
+  });
+  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.
+  const added = out.functions[0]!.locals.filter((local) => local.id.startsWith("%jam."));
+  expect(added.filter((local) => local.name === "%jam.b").length).toBe(4);
+  const block = body[1]!;
+  expect(block.kind).toBe("block");
+  if (block.kind !== "block") return;
+  const posts = block.body.slice(-4);
+  expect(posts.every((s) => s.kind === "assign" && s.localId === "res")).toBe(true);
+});
+
+test("keeps loops whose exits or values do not allow running iterations together", () => {
+  const reject = (body: IrStmt[]): void => {
+    const mod = moduleOf(body);
+    expect(jamSaturatingLoops(mod)).toBe(mod);
+  };
+  // An exit before the accumulator is full: later ORs would change it.
+  reject(nest({ full: 2 }));
+  // `post` reads a lane value the copies keep computing past their exit.
+  reject(nest({ post: [set("res", bin("+", ref("res"), ref("z1")))] }));
+  // `pre` reads what `post` writes.
+  reject(nest({ preExtra: [decl("seed", ref("res"))] }));
+  // The inner loop writes the accumulator other than by an OR: a later
+  // iteration can clear bits after an exit.
+  reject(nest({ innerExtra: [set("b", num(0))] }));
+  reject(nest({ innerExtra: [set("b", bin("+", ref("b"), num(4)))] }));
+  // The inner loop's body advances its counter, which the copies share.
+  reject(nest({ innerExtra: [stepBy("j", 1)] }));
+  // The inner loop's start or bound reads the counter or a local of the
+  // iteration.
+  reject(nest({ innerStart: ref("i") }));
+  reject(nest({ innerEnd: bin("+", num(10), ref("i")) }));
+  reject(nest({ preExtra: [decl("start", num(3))], innerStart: ref("start") }));
+  // One iteration, or more than eight.
+  reject(nest({ copies: 1 }));
+  reject(nest({ copies: 9 }));
+  // A store in the inner loop is observable past an exit.
+  reject([
+    decl("res", num(0)),
+    loop("i", num(8), 2, [
+      decl("b", num(0)),
+      decl("z1", num(0)),
+      loop("j", num(50), 1, [
+        ...lane("z1", "z1", 2, 3, [
+          {
+            loc,
+            kind: "bytesSet",
+            arr: { loc, kind: "varRef", localId: "cr", type: BYTES },
+            index: num(0),
+            value: num(1),
+          },
+        ]),
+      ]),
+      set("res", bin("|", ref("res"), ref("b"))),
+    ]),
+  ]);
+});

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

@@ -0,0 +1,443 @@
+import {
+  BOOL,
+  F64,
+  type IrExpr,
+  type IrFunction,
+  type IrLocal,
+  type IrModule,
+  type IrStmt,
+} from "./ir.js";
+import { everyStmtList, transformStmtList } from "./traverse.js";
+
+/** The most iterations of an outer loop that run jammed. */
+const MAX_COPIES = 8;
+
+/** The most statement and expression nodes of one outer iteration: every
+ * copy duplicates it. */
+const MAX_BODY_NODES = 160;
+
+/** Binary operators of numbers and booleans that run no other code and
+ * cannot throw. */
+const PURE_BINARY = new Set<string>([
+  "+",
+  "-",
+  "*",
+  "/",
+  "%",
+  "&",
+  "|",
+  "^",
+  "<<",
+  ">>",
+  ">>>",
+  "<",
+  "<=",
+  ">",
+  ">=",
+  "===",
+  "!==",
+]);
+
+const PURE_UNARY = new Set<string>(["-", "+", "~", "!"]);
+
+interface Jam {
+  counter: string;
+  /** The counter's value on each iteration. */
+  values: number[];
+  pre: IrStmt[];
+  inner: IrStmt & { kind: "for" };
+  post: IrStmt[];
+  /** The accumulator the inner loop's exits test, and its full value. */
+  acc: string;
+  full: number;
+}
+
+function isRef(e: IrExpr | null | undefined, id: string): boolean {
+  return e?.kind === "varRef" && e.localId === id;
+}
+
+function intLit(e: IrExpr | null | undefined): number | null {
+  return e?.kind === "numLit" && Number.isSafeInteger(e.value) && !Object.is(e.value, -0)
+    ? e.value
+    : null;
+}
+
+/** `id = id + step` with a positive integer literal step. */
+function constantStep(update: IrStmt | null, id: string): number | null {
+  if (
+    update?.kind !== "assign" ||
+    update.localId !== id ||
+    update.value.kind !== "bin" ||
+    update.value.op !== "+" ||
+    !isRef(update.value.left, id)
+  )
+    return null;
+  const step = intLit(update.value.right);
+  return step !== null && step > 0 ? step : null;
+}
+
+/** A number or boolean expression that runs no other code and cannot throw:
+ * literals, reads of locals and module bindings, arithmetic, comparisons
+ * and typed-array element reads that answer NaN out of range. */
+function pureExpr(e: IrExpr, locals: ReadonlyMap<string, IrLocal>): boolean {
+  if (e.type.kind !== "f64" && e.type.kind !== "bool") return false;
+  switch (e.kind) {
+    case "numLit":
+    case "boolLit":
+      return true;
+    case "varRef": {
+      const local = locals.get(e.localId);
+      if (local) return !local.boxed && !local.tdz;
+      return e.localId.startsWith("%g.");
+    }
+    case "bin":
+      return PURE_BINARY.has(e.op) && pureExpr(e.left, locals) && pureExpr(e.right, locals);
+    case "unary":
+      return PURE_UNARY.has(e.op) && pureExpr(e.operand, locals);
+    case "bytesIntrinsic":
+      return (
+        e.method === "get" &&
+        e.invalidNaN === true &&
+        e.receiver.kind === "varRef" &&
+        e.args.length === 1 &&
+        pureExpr(e.args[0]!, locals)
+      );
+    default:
+      return false;
+  }
+}
+
+/** A local written by `s` (a declaration or assignment), or null. */
+function writtenLocal(s: IrStmt): string | null {
+  return s.kind === "varDecl" || s.kind === "assign" ? s.localId : null;
+}
+
+/** Straight-line declarations and assignments of pure values. */
+function pureStatements(list: readonly IrStmt[], locals: ReadonlyMap<string, IrLocal>): boolean {
+  return list.every(
+    (s) =>
+      (s.kind === "varDecl" && s.init !== null && pureExpr(s.init, locals)) ||
+      (s.kind === "assign" && locals.has(s.localId) && pureExpr(s.value, locals)),
+  );
+}
+
+/** The bits `C` of `id = id | C`, with `C` a non-negative int32 literal, or
+ * null. */
+function orBits(s: IrStmt): number | null {
+  if (
+    s.kind !== "assign" ||
+    s.value.kind !== "bin" ||
+    s.value.op !== "|" ||
+    !isRef(s.value.left, s.localId)
+  )
+    return null;
+  const bits = intLit(s.value.right);
+  return bits !== null && bits >= 0 && bits <= 0x7fffffff ? bits : null;
+}
+
+/** The statements of the inner loop's body: pure declarations and
+ * assignments, `if` statements over them, `acc = acc | C`, and
+ * `if (acc === full) break` (collected into `breaks`). */
+function innerStatements(
+  list: readonly IrStmt[],
+  locals: ReadonlyMap<string, IrLocal>,
+  acc: { id: string | null; bits: number; breaks: number[] },
+): boolean {
+  return list.every((s) => {
+    if (s.kind === "if") {
+      // `if (acc === full) break`
+      if (
+        s.then.length === 1 &&
+        s.then[0]!.kind === "break" &&
+        s.then[0]!.label === undefined &&
+        (s.else_ === null || s.else_.length === 0) &&
+        s.cond.kind === "bin" &&
+        s.cond.op === "===" &&
+        s.cond.left.kind === "varRef" &&
+        intLit(s.cond.right) !== null
+      ) {
+        if (acc.id !== null && acc.id !== s.cond.left.localId) return false;
+        acc.id = s.cond.left.localId;
+        acc.breaks.push(intLit(s.cond.right)!);
+        return true;
+      }
+      return (
+        pureExpr(s.cond, locals) &&
+        innerStatements(s.then, locals, acc) &&
+        innerStatements(s.else_ ?? [], locals, acc)
+      );
+    }
+    const bits = orBits(s);
+    if (s.kind === "assign" && bits !== null) {
+      if (acc.id !== null && acc.id !== s.localId) return false;
+      acc.id = s.localId;
+      acc.bits |= bits;
+      return true;
+    }
+    return pureStatements([s], locals);
+  });
+}
+
+function countNodes(list: readonly IrStmt[]): number {
+  let n = 0;
+  everyStmtList(list as IrStmt[], {
+    expr: () => ++n <= MAX_BODY_NODES,
+    stmt: () => ++n <= MAX_BODY_NODES,
+  });
+  return n;
+}
+
+/** Reads of local `id` anywhere in `list`. */
+function reads(list: readonly IrStmt[], id: string): boolean {
+  return !everyStmtList(list as IrStmt[], { expr: (e) => !isRef(e, id), stmt: () => true });
+}
+
+/** Locals written anywhere in `list`. */
+function writes(list: readonly IrStmt[]): Set<string> {
+  const out = new Set<string>();
+  everyStmtList(list as IrStmt[], {
+    expr: () => true,
+    stmt: (s) => {
+      const id = writtenLocal(s);
+      if (id !== null) out.add(id);
+      return true;
+    },
+  });
+  return out;
+}
+
+/** Locals declared anywhere in `list`. */
+function declared(list: readonly IrStmt[]): Set<string> {
+  const out = new Set<string>();
+  everyStmtList(list as IrStmt[], {
+    expr: () => true,
+    stmt: (s) => {
+      if (s.kind === "varDecl") out.add(s.localId);
+      return true;
+    },
+  });
+  return out;
+}
+
+/** The jam of `loop`, or null. */
+function matchJam(
+  loop: IrStmt & { kind: "for" },
+  locals: ReadonlyMap<string, IrLocal>,
+): Jam | null {
+  if (loop.labels?.length || loop.init?.kind !== "varDecl") return null;
+  const counter = loop.init.localId;
+  const local = locals.get(counter);
+  if (local?.type.kind !== "f64" || local.boxed || local.tdz) return null;
+  const start = intLit(loop.init.init);
+  const step = constantStep(loop.update, counter);
+  const cond = loop.cond;
+  if (
+    start === null ||
+    step === null ||
+    cond?.kind !== "bin" ||
+    cond.op !== "<" ||
+    !isRef(cond.left, counter)
+  )
+    return null;
+  const end = intLit(cond.right);
+  if (end === null) return null;
+  const values: number[] = [];
+  for (let v = start; v < end && values.length <= MAX_COPIES; v += step) values.push(v);
+  if (values.length < 2 || values.length > MAX_COPIES) return null;
+  // One inner loop, with straight-line code before and after it.
+  const at = loop.body.findIndex((s) => s.kind === "for");
+  if (at < 0) return null;
+  const inner = loop.body[at] as IrStmt & { kind: "for" };
+  const pre = loop.body.slice(0, at),
+    post = loop.body.slice(at + 1);
+  if (!pureStatements(pre, locals) || !pureStatements(post, locals)) return null;
+  if (countNodes(loop.body) > MAX_BODY_NODES) return null;
+  if (inner.labels?.length || inner.init?.kind !== "varDecl") return null;
+  const j = inner.init.localId;
+  const jLocal = locals.get(j);
+  if (jLocal?.type.kind !== "f64" || jLocal.boxed || jLocal.tdz) return null;
+  if (inner.init.init === null || !pureExpr(inner.init.init, locals)) return null;
+  if (inner.cond === null || !pureExpr(inner.cond, locals) || constantStep(inner.update, j) !== 1)
+    return null;
+  const acc = { id: null as string | null, bits: 0, breaks: [] as number[] };
+  if (!innerStatements(inner.body, locals, acc) || acc.id === null) return null;
+  const full = acc.bits;
+  // Every exit tests the accumulator against the value every OR leaves
+  // unchanged once it is reached.
+  if (acc.breaks.length === 0 || acc.breaks.some((value) => value !== full)) return null;
+  const accId = acc.id;
+  // The accumulator starts at a literal in `pre`, and the inner loop changes
+  // it only by its ORs.
+  const accDecl = pre.find((s) => s.kind === "varDecl" && s.localId === accId);
+  if (accDecl?.kind !== "varDecl" || intLit(accDecl.init) === null) return null;
+  const onlyOrs = everyStmtList(inner.body, {
+    expr: () => true,
+    stmt: (s) => writtenLocal(s) !== accId || orBits(s) !== null,
+  });
+  if (!onlyOrs) return null;
+  // Every copy shares the inner counter: only the inner loop's update writes
+  // it.
+  if (writes(inner.body).has(j)) return null;
+  // The counter, the inner bound and every value the inner loop reads from
+  // outside the iteration are the same for every copy: only `pre` and the
+  // counter vary.
+  const own = declared([...pre, inner]);
+  own.delete(counter);
+  const innerWrites = writes([inner]);
+  innerWrites.delete(j);
+  for (const id of innerWrites) if (!own.has(id) && id !== accId) return null;
+  if (writes([inner]).has(counter) || writes(pre).has(counter) || writes(post).has(counter))
+    return null;
+  // Copies run `pre` and the inner loop of later iterations before `post`
+  // of earlier ones: `pre` and the inner loop must not read what `post`
+  // writes, and `pre` writes only its own declarations.
+  const postWrites = writes(post);
+  for (const id of postWrites) if (reads([...pre, inner], id) || own.has(id)) return null;
+  for (const id of writes(pre)) if (!own.has(id)) return null;
+  // After the inner loop, iterations past a copy's exit are dead: `post`
+  // reads no value the inner loop writes besides the accumulator.
+  for (const id of innerWrites) if (id !== accId && reads(post, id)) return null;
+  // The copies share the inner loop's start and bound: they read neither the
+  // counter nor a local of the iteration.
+  const controls: IrStmt[] = [
+    { kind: "exprStmt", expr: inner.init.init, loc: inner.loc },
+    { kind: "exprStmt", expr: inner.cond, loc: inner.loc },
+  ];
+  if (reads(controls, counter)) return null;
+  for (const id of own) if (id !== j && reads(controls, id)) return null;
+  return { counter, values, pre, inner, post, acc: accId, full };
+}
+
+/** Runs the iterations of a short constant-count loop together when its body
+ * is an inner loop that exits once an accumulator reaches the value its ORs
+ * saturate at (mandelbrot's pairs of pixels per output byte):
+ *
+ *   for (let i = 0; i < 8; i += 2) {
+ *     <pre: pure declarations, including `let b = 0`>
+ *     for (let j = 0; j < M; j++) {
+ *       <pure arithmetic>; if (c) { b |= 2; if (b === 3) break; } ...
+ *     }
+ *     <post: pure assignments reading b>
+ *   }
+ *
+ * becomes every iteration's `pre` (with `i` a literal), one inner loop
+ * running every iteration's body in turn that exits when every accumulator
+ * is full, and every iteration's `post` in order. Once an accumulator is
+ * full its ORs leave it full, so the iterations a copy runs past its own
+ * exit leave its accumulator as its own loop left it, and every other value
+ * they compute is dead: the bodies only compute numbers, read typed-array
+ * 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. */
+export function jamSaturatingLoops(mod: IrModule): IrModule {
+  let changed = false;
+  const functions = mod.functions.map((fn): IrFunction => {
+    if (fn.async || fn.generator) return fn;
+    const locals = new Map(fn.locals.map((local) => [local.id, local]));
+    let found = false;
+    everyStmtList(fn.body, {
+      expr: () => !found,
+      stmt: (s) => {
+        if (s.kind === "for" && matchJam(s, locals) !== null) found = true;
+        return !found;
+      },
+    });
+    if (!found) return fn;
+    const added: IrLocal[] = [];
+    const ids = new Set(fn.locals.map((local) => local.id));
+    const fresh = (of: IrLocal, copy: number): string => {
+      let id = `%jam.${copy}.${of.id}`;
+      while (ids.has(id)) id += "_";
+      ids.add(id);
+      const { source: _source, ...rest } = of;
+      added.push({ ...rest, id, name: `%jam.${of.name}` });
+      return id;
+    };
+    const body = transformStmtList(fn.body, {
+      expr: (e) => e,
+      stmt: (s) => {
+        if (s.kind !== "for") return s;
+        const jam = matchJam(s, locals);
+        return jam === null ? s : emitJam(jam, locals, fresh);
+      },
+    });
+    changed = true;
+    return { ...fn, locals: [...fn.locals, ...added], body };
+  });
+  return changed ? { ...mod, functions } : mod;
+}
+
+function emitJam(
+  jam: Jam,
+  locals: ReadonlyMap<string, IrLocal>,
+  fresh: (of: IrLocal, copy: number) => string,
+): IrStmt {
+  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;
+  own.delete(j);
+  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 rewrite = (list: readonly IrStmt[]): IrStmt[] =>
+      transformStmtList(list as IrStmt[], {
+        expr: (e) => {
+          if (e.kind !== "varRef") return e;
+          if (e.localId === counter) return { kind: "numLit", value, type: F64, loc: e.loc };
+          const to = rename.get(e.localId);
+          return to === undefined ? e : { ...e, localId: to };
+        },
+        stmt: (s) => {
+          if (s.kind === "varDecl" || s.kind === "assign") {
+            const to = rename.get(s.localId);
+            if (to !== undefined) return { ...s, localId: to };
+          }
+          return s;
+        },
+      });
+    // The copy's exits become the shared test at the end of the iteration.
+    const withoutBreaks = (list: readonly IrStmt[]): IrStmt[] =>
+      list.flatMap((s): IrStmt[] => {
+        if (s.kind !== "if") return [s];
+        if (s.then.length === 1 && s.then[0]!.kind === "break") return [];
+        return [
+          { ...s, then: withoutBreaks(s.then), else_: s.else_ ? withoutBreaks(s.else_) : null },
+        ];
+      });
+    return {
+      pre: rewrite(pre),
+      body: rewrite(withoutBreaks(inner.body)),
+      post: rewrite(post),
+      acc: rename.get(jam.acc)!,
+    };
+  });
+  let allFull: IrExpr | null = null;
+  for (const copy of copies) {
+    const full: IrExpr = {
+      kind: "bin",
+      op: "===",
+      left: { kind: "varRef", localId: copy.acc, type: F64, loc },
+      right: { kind: "numLit", value: jam.full, type: F64, loc },
+      type: BOOL,
+      loc,
+    };
+    allFull =
+      allFull === null
+        ? full
+        : { kind: "logical", op: "&&", left: allFull, right: full, type: BOOL, loc };
+  }
+  const jammed: IrStmt = {
+    ...inner,
+    body: [
+      ...copies.flatMap((copy) => copy.body),
+      { kind: "if", cond: allFull!, then: [{ kind: "break", loc }], else_: null, loc },
+    ],
+  };
+  return {
+    kind: "block",
+    body: [...copies.flatMap((copy) => copy.pre), jammed, ...copies.flatMap((copy) => copy.post)],
+    loc,
+  };
+}

packages/compiler/test/ts7/baselines/order-parity/tests/corpus/saturating-loop-jams.ts.json

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

tests/corpus/saturating-loop-jams.ts

@@ -0,0 +1,164 @@
+// Escape-time loops in the shape of mandelbrot's rows: a short constant-count
+// loop over pixel pairs whose inner loop exits once both pixels of the pair
+// have escaped (`b |= 2`, `b |= 1`, `if (b === 3) break`). Pixels escape at
+// different iterations or never; each byte packs the pairs' bits in order.
+const LIMIT = 4;
+
+function row(initial: Float64Array, ci: number, bytes: number, iterations: number): number[] {
+  const out: number[] = [];
+  for (let xByte = 0; xByte < bytes; xByte++) {
+    let res = 0;
+    for (let i = 0; i < 8; i += 2) {
+      const x = xByte << 3;
+      const cr1 = initial[x + i]!;
+      const cr2 = initial[x + i + 1]!;
+      let zr1 = cr1;
+      let zi1 = ci;
+      let zr2 = cr2;
+      let zi2 = ci;
+      let b = 0;
+      for (let j = 0; j < iterations; j++) {
+        const tr1 = zr1 * zr1;
+        const ti1 = zi1 * zi1;
+        zi1 = 2 * zr1 * zi1 + ci;
+        zr1 = tr1 - ti1 + cr1;
+        if (tr1 + ti1 > LIMIT) {
+          b |= 2;
+          if (b === 3) break;
+        }
+        const tr2 = zr2 * zr2;
+        const ti2 = zi2 * zi2;
+        zi2 = 2 * zr2 * zi2 + ci;
+        zr2 = tr2 - ti2 + cr2;
+        if (tr2 + ti2 > LIMIT) {
+          b |= 1;
+          if (b === 3) break;
+        }
+      }
+      res = (res << 2) | b;
+    }
+    out.push(~res & 255);
+  }
+  return out;
+}
+
+/** Two pairs per value: the loop runs two iterations. */
+function quads(initial: Float64Array, ci: 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 = 0; j < 30; j++) {
+      zr1 = zr1 * zr1 + ci;
+      if (zr1 > 2) {
+        b |= 2;
+        if (b === 3) break;
+      }
+      zr2 = zr2 * zr2 + ci;
+      if (zr2 > 2) {
+        b |= 1;
+        if (b === 3) break;
+      }
+    }
+    res = res * 4 + b;
+  }
+  return res;
+}
+
+// Loops the jam keeps as they are. The inner loop clears the accumulator.
+function resets(initial: Float64Array): 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 = 0; j < 30; 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;
+      }
+      if (j > 20) b = 0;
+    }
+    res = res * 4 + b;
+  }
+  return res;
+}
+
+// The inner loop's body advances its counter.
+function skips(initial: Float64Array): number {
+  let res = 0;
+  for (let i = 0; i < 4; i += 2) {
+    let zr1 = initial[i]!;
+    let zr2 = initial[i + 1]!;
+    let b = 0;
+    let steps = 0;
+    for (let j = 0; j < 16; 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;
+      }
+      if (zr1 < 0.5) j = j + 1;
+    }
+    res = res * 4 + b;
+  }
+  return res;
+}
+
+// The inner loop's start and bound depend on the outer iteration.
+function shifted(initial: Float64Array): number {
+  let res = 0;
+  for (let i = 0; i < 4; i += 2) {
+    let zr1 = initial[i]!;
+    let zr2 = initial[i + 1]!;
+    let b = 0;
+    const start = i * 3;
+    for (let j = start; j < 10 + i; 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;
+let checksum = 0;
+const lines: string[] = [];
+for (let y = 0; y < size; y += 8) {
+  const ci = (2 * y) / size - 1;
+  const bytes = row(initial, ci, size >> 3, y % 16 === 0 ? 50 : 3);
+  lines.push(bytes.map((v) => v.toString(16).padStart(2, "0")).join(""));
+  for (const v of bytes) checksum = (checksum * 31 + v) % 1000003;
+}
+console.log(lines.join("\n"));
+console.log(checksum);
+const q = new Float64Array([0.1, 1.5, -0.3, 3]);
+console.log(quads(q, 0.2), quads(q, -1.9), quads(new Float64Array(4), 0));
+// A row past the end of the coordinates reads NaN, which never escapes.
+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));

tests/harness/corpus-costs/saturating-loop-jams.ts.json

@@ -0,0 +1 @@
+{ "plain": 5955, "san": 72997 }