diff --git a/packages/compiler/src/ir/saturating-loop-jams.test.ts b/packages/compiler/src/ir/saturating-loop-jams.test.ts index e25b17e80..25320b7a6 100644 --- a/packages/compiler/src/ir/saturating-loop-jams.test.ts +++ b/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; diff --git a/packages/compiler/src/ir/saturating-loop-jams.ts b/packages/compiler/src/ir/saturating-loop-jams.ts index 96973a483..74c204fea 100644 --- a/packages/compiler/src/ir/saturating-loop-jams.ts +++ b/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()); + 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(); - 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 {