From 500dd3e9b4a80124260244044d22c7cd60f4e78f Mon Sep 17 00:00:00 2001 From: Chris Tate Date: Wed, 28 Oct 2026 18:52:05 -0500 Subject: [PATCH 1/2] Run short constant-count loops together when their inner exits saturate A loop of two to eight iterations with a constant count, whose body is an inner loop that exits once an accumulator reaches the value its ORs saturate at (mandelbrot's pixel pairs: `b |= 2`, `b |= 1`, `if (b === 3) break`), now runs its iterations in one inner loop: every iteration's declarations, then one loop running each iteration's body in turn that exits once every accumulator is full, then every iteration's trailing assignments in order. Once an accumulator is full its ORs leave it full, so an iteration that runs past its own exit leaves its accumulator as its own loop did, and every other value it computes is dead. The bodies only compute numbers, read typed-array elements that answer NaN out of range and write their own locals; the declarations and the inner loop read nothing the trailing assignments write. The independent iterations then overlap their dependence chains. --- packages/compiler/src/backend/llvm/emitter.ts | 23 +- .../src/ir/saturating-loop-jams.test.ts | 188 ++++++++ .../compiler/src/ir/saturating-loop-jams.ts | 424 ++++++++++++++++++ .../tests/corpus/saturating-loop-jams.ts.json | 6 + tests/corpus/saturating-loop-jams.ts | 85 ++++ .../corpus-costs/saturating-loop-jams.ts.json | 1 + 6 files changed, 717 insertions(+), 10 deletions(-) create mode 100644 packages/compiler/src/ir/saturating-loop-jams.test.ts create mode 100644 packages/compiler/src/ir/saturating-loop-jams.ts create mode 100644 packages/compiler/test/ts7/baselines/order-parity/tests/corpus/saturating-loop-jams.ts.json create mode 100644 tests/corpus/saturating-loop-jams.ts create mode 100644 tests/harness/corpus-costs/saturating-loop-jams.ts.json diff --git a/packages/compiler/src/backend/llvm/emitter.ts b/packages/compiler/src/backend/llvm/emitter.ts index 9d85bb8da..eeb418512 100644 --- a/packages/compiler/src/backend/llvm/emitter.ts +++ b/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, @@ -759,16 +760,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)), + ), ), ), ), diff --git a/packages/compiler/src/ir/saturating-loop-jams.test.ts b/packages/compiler/src/ir/saturating-loop-jams.test.ts new file mode 100644 index 000000000..b3d1d0fe9 --- /dev/null +++ b/packages/compiler/src/ir/saturating-loop-jams.test.ts @@ -0,0 +1,188 @@ +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[]): IrStmt & { kind: "for" } => ({ + loc, + kind: "for", + init: decl(id, num(0)), + 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[]; 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", num(50), 1, [ + ...lane("z1", "c1", 2, options.full), + ...lane("z2", "c2", 1, options.full), + ]), + ...post, + ]), + ]; +} + +function moduleOf(body: IrStmt[]): IrModule { + const ids = new Set(); + 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"))] })); + // 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"))), + ]), + ]); +}); diff --git a/packages/compiler/src/ir/saturating-loop-jams.ts b/packages/compiler/src/ir/saturating-loop-jams.ts new file mode 100644 index 000000000..ae4bd38c8 --- /dev/null +++ b/packages/compiler/src/ir/saturating-loop-jams.ts @@ -0,0 +1,424 @@ +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([ + "+", + "-", + "*", + "/", + "%", + "&", + "|", + "^", + "<<", + ">>", + ">>>", + "<", + "<=", + ">", + ">=", + "===", + "!==", +]); + +const PURE_UNARY = new Set(["-", "+", "~", "!"]); + +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): 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): 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 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, + 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) + ); + } + if ( + s.kind === "assign" && + s.value.kind === "bin" && + s.value.op === "|" && + isRef(s.value.left, s.localId) + ) { + const bits = intLit(s.value.right); + if (bits !== null && bits >= 0 && bits <= 0x7fffffff) { + 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 { + const out = new Set(); + 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 { + const out = new Set(); + 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, +): 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; + // 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, and the + // inner bound reads nothing the copies write. + for (const id of innerWrites) if (id !== accId && reads(post, id)) return null; + for (const id of own) + if (id !== j && reads([{ kind: "exprStmt", expr: inner.cond, loc: inner.loc }], 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) { + * + * for (let j = 0; j < M; j++) { + * ; if (c) { b |= 2; if (b === 3) break; } ... + * } + * + * } + * + * 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, + 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(); + 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, + }; +} diff --git a/packages/compiler/test/ts7/baselines/order-parity/tests/corpus/saturating-loop-jams.ts.json b/packages/compiler/test/ts7/baselines/order-parity/tests/corpus/saturating-loop-jams.ts.json new file mode 100644 index 000000000..34cd32574 --- /dev/null +++ b/packages/compiler/test/ts7/baselines/order-parity/tests/corpus/saturating-loop-jams.ts.json @@ -0,0 +1,6 @@ +{ + "order": [ + "/tests/corpus/saturating-loop-jams.ts" + ], + "diags": [] +} diff --git a/tests/corpus/saturating-loop-jams.ts b/tests/corpus/saturating-loop-jams.ts new file mode 100644 index 000000000..fdc066eb8 --- /dev/null +++ b/tests/corpus/saturating-loop-jams.ts @@ -0,0 +1,85 @@ +// 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; +} + +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)); diff --git a/tests/harness/corpus-costs/saturating-loop-jams.ts.json b/tests/harness/corpus-costs/saturating-loop-jams.ts.json new file mode 100644 index 000000000..e615532ed --- /dev/null +++ b/tests/harness/corpus-costs/saturating-loop-jams.ts.json @@ -0,0 +1 @@ +{ "plain": 5955, "san": 72997 } From a64b5b082050b6879bbd0dc1b218c21122490838 Mon Sep 17 00:00:00 2001 From: Chris Tate Date: Wed, 28 Oct 2026 23:01:16 -0500 Subject: [PATCH 2/2] Keep saturating loops whose copies cannot share the inner loop The jam of a short constant-count loop now leaves the loop as it is when: - the inner loop's body assigns the accumulator other than by `acc = acc | C`. An iteration that runs past its exit could clear bits after the accumulator was full. - the inner loop's body assigns the inner counter. The copies share the counter, so one copy's write changed the iterations of the others. - the inner loop's start or bound reads the outer counter or a local of the outer iteration. The copies share the inner loop's start and bound, which read neither the counter's value nor the copies' renamed locals. The unit test rejects each shape, and the corpus program adds a function for each that matches Node's output. --- .../src/ir/saturating-loop-jams.test.ts | 46 +++++++++-- .../compiler/src/ir/saturating-loop-jams.ts | 55 ++++++++----- tests/corpus/saturating-loop-jams.ts | 79 +++++++++++++++++++ 3 files changed, 155 insertions(+), 25 deletions(-) diff --git a/packages/compiler/src/ir/saturating-loop-jams.test.ts b/packages/compiler/src/ir/saturating-loop-jams.test.ts index b3d1d0fe9..e25b17e80 100644 --- a/packages/compiler/src/ir/saturating-loop-jams.test.ts +++ b/packages/compiler/src/ir/saturating-loop-jams.test.ts @@ -19,10 +19,16 @@ const bin = (op: Op, left: IrExpr, right: IrExpr): IrExpr => ({ 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[]): IrStmt & { kind: "for" } => ({ +const loop = ( + id: string, + end: IrExpr, + by: number, + body: IrStmt[], + start: IrExpr = num(0), +): IrStmt & { kind: "for" } => ({ loc, kind: "for", - init: decl(id, num(0)), + init: decl(id, start), cond: bin("<", ref(id), end), update: stepBy(id, by), body, @@ -61,7 +67,15 @@ function lane(z: string, c: string, bit: number, full = 3, extra: IrStmt[] = []) } function nest( - options: { full?: number; post?: IrStmt[]; preExtra?: IrStmt[]; copies?: number } = {}, + 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 [ @@ -73,10 +87,17 @@ function nest( decl("z2", ref("c2")), decl("b", num(0)), ...(options.preExtra ?? []), - loop("j", num(50), 1, [ - ...lane("z1", "c1", 2, options.full), - ...lane("z2", "c2", 1, options.full), - ]), + loop( + "j", + options.innerEnd ?? num(50), + 1, + [ + ...lane("z1", "c1", 2, options.full), + ...lane("z2", "c2", 1, options.full), + ...(options.innerExtra ?? []), + ], + options.innerStart, + ), ...post, ]), ]; @@ -162,6 +183,17 @@ test("keeps loops whose exits or values do not allow running iterations together 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 })); diff --git a/packages/compiler/src/ir/saturating-loop-jams.ts b/packages/compiler/src/ir/saturating-loop-jams.ts index ae4bd38c8..96973a483 100644 --- a/packages/compiler/src/ir/saturating-loop-jams.ts +++ b/packages/compiler/src/ir/saturating-loop-jams.ts @@ -121,6 +121,20 @@ function pureStatements(list: readonly IrStmt[], locals: ReadonlyMap= 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`). */ @@ -153,19 +167,12 @@ function innerStatements( innerStatements(s.else_ ?? [], locals, acc) ); } - if ( - s.kind === "assign" && - s.value.kind === "bin" && - s.value.op === "|" && - isRef(s.value.left, s.localId) - ) { - const bits = intLit(s.value.right); - if (bits !== null && bits >= 0 && bits <= 0x7fffffff) { - if (acc.id !== null && acc.id !== s.localId) return false; - acc.id = s.localId; - acc.bits |= bits; - return true; - } + 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); }); @@ -263,6 +270,14 @@ function matchJam( // 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. @@ -280,12 +295,16 @@ function matchJam( 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, and the - // inner bound reads nothing the copies write. + // reads no value the inner loop writes besides the accumulator. for (const id of innerWrites) if (id !== accId && reads(post, id)) return null; - for (const id of own) - if (id !== j && reads([{ kind: "exprStmt", expr: inner.cond, loc: inner.loc }], 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 }; } diff --git a/tests/corpus/saturating-loop-jams.ts b/tests/corpus/saturating-loop-jams.ts index fdc066eb8..fb0fd5bb3 100644 --- a/tests/corpus/saturating-loop-jams.ts +++ b/tests/corpus/saturating-loop-jams.ts @@ -66,6 +66,82 @@ function quads(initial: Float64Array, ci: number): number { 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; @@ -83,3 +159,6 @@ 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));