diff --git a/packages/compiler/src/ir/saturating-loop-jams.test.ts b/packages/compiler/src/ir/saturating-loop-jams.test.ts index 25320b7a6..7593af5a8 100644 --- a/packages/compiler/src/ir/saturating-loop-jams.test.ts +++ b/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); diff --git a/packages/compiler/src/ir/saturating-loop-jams.ts b/packages/compiler/src/ir/saturating-loop-jams.ts index 74c204fea..3773f713f 100644 --- a/packages/compiler/src/ir/saturating-loop-jams.ts +++ b/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, diff --git a/tests/corpus/saturating-loop-jams.ts b/tests/corpus/saturating-loop-jams.ts index fb0fd5bb3..d4fd5e94d 100644 --- a/tests/corpus/saturating-loop-jams.ts +++ b/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));