From be293e3dfbbee4ce16d1c029249f48cae173d247 Mon Sep 17 00:00:00 2001 From: Chris Tate Date: Mon, 26 Oct 2026 10:19:25 -0500 Subject: [PATCH] Unroll counted loops by the period of a guard on the counter's low bits A counted loop whose body branches on `(i & M) === C`, with M one of 1, 3 and 7 (a check every second, fourth or eighth iteration), now asks LLVM to unroll it by M + 1 (llvm.loop.unroll.count on its back edge). In each unrolled copy the counter's low bits are known, so the guard folds to a constant and only one copy per period keeps the check. Before, LLVM tested the guard on every iteration and, when the guarded condition was cheap, evaluated that condition on every iteration too. The loop must step by 1 from a multiple of the period, have no nested loop, and stay under 4096 IR nodes once unrolled. Loops with the same period share one loop ID node, which the module emits only when a back edge refers to it. --- packages/compiler/src/backend/llvm/blocks.ts | 4 +- .../compiler/src/backend/llvm/emit-module.ts | 3 + .../src/backend/llvm/emit-statements.ts | 12 ++- .../src/backend/llvm/periodic-guards.test.ts | 80 +++++++++++++++ .../src/backend/llvm/periodic-guards.ts | 97 +++++++++++++++++++ .../test/periodic-guard-emission.test.ts | 53 ++++++++++ 6 files changed, 245 insertions(+), 4 deletions(-) create mode 100644 packages/compiler/src/backend/llvm/periodic-guards.test.ts create mode 100644 packages/compiler/src/backend/llvm/periodic-guards.ts create mode 100644 packages/compiler/test/periodic-guard-emission.test.ts diff --git a/packages/compiler/src/backend/llvm/blocks.ts b/packages/compiler/src/backend/llvm/blocks.ts index eee4d00b9..3ae65cb0f 100644 --- a/packages/compiler/src/backend/llvm/blocks.ts +++ b/packages/compiler/src/backend/llvm/blocks.ts @@ -103,8 +103,8 @@ export class BlockBuilder { } } - br(label: string): void { - this.terminate(`br label %${label}`); + br(label: string, attachment = ""): void { + this.terminate(`br label %${label}${attachment}`); } condBr(v: string, t: string, f: string, attachment = ""): void { diff --git a/packages/compiler/src/backend/llvm/emit-module.ts b/packages/compiler/src/backend/llvm/emit-module.ts index d25427e4f..15efa88c7 100644 --- a/packages/compiler/src/backend/llvm/emit-module.ts +++ b/packages/compiler/src/backend/llvm/emit-module.ts @@ -1,4 +1,5 @@ /* LlEmitter's module-level output: module parts, helper definitions, library declarations and FFI callback definitions. */ +import { loopUnrollMetadata } from "./periodic-guards.js"; import { relieveHeapPressure } from "../../heap-pressure.js"; import { emitWasmModules } from "./wasm-modules.js"; import { type LlEmitter, ffiCallbackDummyLl, llStrBytes } from "./emitter.js"; @@ -1099,6 +1100,7 @@ export function emitModuleParts(emitter: LlEmitter): string[] { out.push(...functionAttributes(emitter, hasNoInlineRecordClone)); out.push(``); out.push(...emitter.fieldAliasMetadata()); + out.push(...loopUnrollMetadata(out)); if (emitter.debug !== null) out.push(emitter.debug.render()); return out; } @@ -1274,6 +1276,7 @@ export function emitModuleParts(emitter: LlEmitter): string[] { ); } out.push(...emitter.fieldAliasMetadata()); + out.push(...loopUnrollMetadata(out)); if (emitter.debug !== null) out.push(emitter.debug.render()); return out; } diff --git a/packages/compiler/src/backend/llvm/emit-statements.ts b/packages/compiler/src/backend/llvm/emit-statements.ts index c1f3cf18f..45d341853 100644 --- a/packages/compiler/src/backend/llvm/emit-statements.ts +++ b/packages/compiler/src/backend/llvm/emit-statements.ts @@ -95,6 +95,7 @@ import { } from "./string-construction.js"; import { emitStringInputs } from "./string-lifetimes.js"; import { LOOPS } from "./loop-state.js"; +import { loopUnrollAttachment, periodicGuardPeriod } from "./periodic-guards.js"; import { loopAccessesTypedArrays } from "../loop-polls.js"; import { emitConstantChoice, matchConstantChoice } from "./constant-choices.js"; @@ -1315,6 +1316,13 @@ export function emitStmtBody(emitter: LlEmitter, s: IrStmt): void { const lb = B.newLabel("loop.b"); const lu = s.update || capturedLets.length > 0 || workerBudget !== null ? B.newLabel("loop.u") : lc; + // A periodic guard on the counter (periodic-guards.ts): the back + // edge asks LLVM to unroll the loop by the guard's period. + const period = + countedLoop && integerSlot + ? periodicGuardPeriod(s, countedLoop.localId, countedLoop.start, countedLoop.step) + : null; + const backEdge = period === null ? "" : loopUnrollAttachment(period); B.br(lc); B.startBlock(lc); enterLoopHeader(emitter, lc); @@ -1375,7 +1383,7 @@ export function emitStmtBody(emitter: LlEmitter, s: IrStmt): void { }); emitter.emitBlock(s.body); emitter.jumpTargets.pop(); - B.br(lu); + B.br(lu, lu === lc ? backEdge : ""); if (lu !== lc) { B.startBlock(lu); emitter.emitWorkerLoopCheck(workerBudget); @@ -1408,7 +1416,7 @@ export function emitStmtBody(emitter: LlEmitter, s: IrStmt): void { B.line(`store i64 ${next}, ptr ${pairedSlots[i]}`); }); } else if (s.update) emitter.emitStmt(s.update); - B.br(lc); + B.br(lc, backEdge); } }); B.startBlock(le); diff --git a/packages/compiler/src/backend/llvm/periodic-guards.test.ts b/packages/compiler/src/backend/llvm/periodic-guards.test.ts new file mode 100644 index 000000000..d1df4e7ab --- /dev/null +++ b/packages/compiler/src/backend/llvm/periodic-guards.test.ts @@ -0,0 +1,80 @@ +import { expect, test } from "vitest"; +import { countedFor } from "../../ir/build.js"; +import { BOOL, F64, type IrExpr, type IrStmt } from "../../ir/ir.js"; +import { + loopUnrollAttachment, + loopUnrollMetadata, + periodicGuardPeriod, +} from "./periodic-guards.js"; + +const loc = { file: "guards.ts", start: 0, end: 0 }; +const ref = (localId: string): IrExpr => ({ kind: "varRef", localId, type: F64, loc }); +const num = (value: number): IrExpr => ({ kind: "numLit", value, type: F64, loc }); +const bin = (op: string, left: IrExpr, right: IrExpr, type = F64): IrExpr => + ({ kind: "bin", op, left, right, type, loc }) as IrExpr; +const and = (left: IrExpr, right: IrExpr): IrExpr => ({ + kind: "logical", + op: "&&", + left, + right, + type: BOOL, + loc, +}); +const guard = (cond: IrExpr): IrStmt => ({ + kind: "if", + cond, + then: [{ kind: "break", loc } as IrStmt], + else_: null, + loc, +}); +/** `for (i = 0; i < 50; i++) { ...body }` */ +const loop = (body: IrStmt[]) => countedFor(loc, num(50), () => body) as IrStmt & { kind: "for" }; +const period = (body: IrStmt[], start = 0, step = 1) => + periodicGuardPeriod(loop(body), "i.0", num(start), step); + +test("a guard on the counter's low bits gives the loop its period", () => { + // if ((i & 3) === 3 && x > 4) break; + const masked = bin("===", bin("&", ref("i.0"), num(3)), num(3), BOOL); + expect(period([guard(and(masked, bin(">", ref("x"), num(4), BOOL)))])).toBe(4); + // if (0 == (7 & i)) ... + expect(period([guard(bin("==", num(0), bin("&", num(7), ref("i.0")), BOOL))])).toBe(8); + expect(period([guard(bin("!==", bin("&", ref("i.0"), num(1)), num(0), BOOL))])).toBe(2); +}); + +test("other guards and loops get no period", () => { + const masked = (mask: number, value: number, counter = "i.0") => + guard(bin("===", bin("&", ref(counter), num(mask)), num(value), BOOL)); + // Masks that are not 1, 3 or 7, values outside the mask, another variable. + expect(period([masked(5, 1)])).toBeNull(); + expect(period([masked(15, 15)])).toBeNull(); + expect(period([masked(3, 4)])).toBeNull(); + expect(period([masked(3, 3, "j")])).toBeNull(); + // The masked test is not the condition's first conjunct. + const later = and( + bin(">", ref("x"), num(4), BOOL), + bin("===", bin("&", ref("i.0"), num(3)), num(3), BOOL), + ); + expect(period([guard(later)])).toBeNull(); + // A start that is not a multiple of the period, a step other than 1. + expect(period([masked(3, 3)], 1)).toBeNull(); + expect(period([masked(3, 3)], 0, 2)).toBeNull(); + // A nested loop. + expect(period([masked(3, 3), loop([])])).toBeNull(); +}); + +test("a large body is not unrolled", () => { + const work: IrStmt[] = []; + for (let k = 0; k < 400; k++) + work.push({ kind: "assign", localId: "x", value: bin("+", ref("x"), num(k)), loc } as IrStmt); + expect( + period([guard(bin("===", bin("&", ref("i.0"), num(3)), num(3), BOOL)), ...work]), + ).toBeNull(); +}); + +test("the module carries the loop ID nodes its back edges refer to", () => { + const lines = [` br label %loop.c1${loopUnrollAttachment(4)}`, " ret void"]; + expect(loopUnrollMetadata(lines)).toEqual([ + `!1000004 = distinct !{!1000004, !{!"llvm.loop.unroll.count", i32 4}}`, + ]); + expect(loopUnrollMetadata([" ret void"])).toEqual([]); +}); diff --git a/packages/compiler/src/backend/llvm/periodic-guards.ts b/packages/compiler/src/backend/llvm/periodic-guards.ts new file mode 100644 index 000000000..3a07126db --- /dev/null +++ b/packages/compiler/src/backend/llvm/periodic-guards.ts @@ -0,0 +1,97 @@ +/* Counted loops with a periodic guard: a branch on the counter's low bits, + * such as `if ((j & 3) === 3 && ...)` (a check every fourth iteration). + * LLVM checks such a guard on every iteration, and when the guarded + * condition is cheap it speculates the condition into the loop. The loop's + * back edge asks LLVM to unroll it by the guard's period + * (`llvm.loop.unroll.count`): in each unrolled copy the counter's low bits + * are known, so the guard folds to a constant, and only one copy in each + * period keeps the check. The metadata changes how LLVM compiles the loop, + * not what it computes. */ +import type { IrExpr, IrStmt } from "../../ir/ir.js"; +import { everyExprChild, everyStmtChild } from "../../ir/traverse.js"; + +/** The most IR nodes a loop body may have, times the period, for the loop + * to be unrolled: a larger body gains less than its unrolled copies cost. */ +const UNROLLED_BODY_NODES = 4096; + +/** The loop IDs, one per period (2, 4 and 8). A loop ID is a distinct + * node whose first operand is itself; loops with the same period share + * one, as the copies of an inlined loop do. */ +const loopIds = new Map([ + [2, 1_000_002], + [4, 1_000_004], + [8, 1_000_008], +]); + +/** The guard period of `loop`, a counted loop on `counter` that starts at + * `start` and steps by `step`, or null. The period is M + 1 where the body + * has an `if` whose condition is, or starts with an `&&` of, + * `(counter & M) === C` (or `==`, `!==`, `!=`, either operand order) with + * M one of 1, 3 and 7 and C an integer from 0 to M. The loop steps by 1 + * from a multiple of the period, so that the counter's low bits are known + * in each unrolled copy, and its body has no nested loop. */ +export function periodicGuardPeriod( + loop: IrStmt & { kind: "for" }, + counter: string, + start: IrExpr, + step: number, +): number | null { + if (step !== 1 || start.kind !== "numLit") return null; + let period: number | null = null; + let nodes = 0; + let plain = true; + const expr = (e: IrExpr): boolean => { + nodes++; + return everyExprChild(e, expr, stmt); + }; + const stmt = (s: IrStmt): boolean => { + nodes++; + if (s.kind === "for" || s.kind === "while" || s.kind === "doWhile" || s.kind === "forOf") { + plain = false; + return false; + } + if (s.kind === "if" && period === null) period = guardPeriod(s.cond, counter); + return everyStmtChild(s, expr, stmt); + }; + for (const s of loop.body) if (!stmt(s)) break; + if (!plain || period === null) return null; + if (start.value % period !== 0 || nodes * period > UNROLLED_BODY_NODES) return null; + return period; +} + +function guardPeriod(cond: IrExpr, counter: string): number | null { + let test = cond; + while (test.kind === "logical" && test.op === "&&") test = test.left; + if (test.kind !== "bin" || !["===", "==", "!==", "!="].includes(test.op)) return null; + const masked = (e: IrExpr, other: IrExpr): number | null => { + if (e.kind !== "bin" || e.op !== "&" || other.kind !== "numLit") return null; + const [ref, mask] = e.left.kind === "varRef" ? [e.left, e.right] : [e.right, e.left]; + if (ref.kind !== "varRef" || ref.localId !== counter || mask.kind !== "numLit") return null; + const m = mask.value; + if (m !== 1 && m !== 3 && m !== 7) return null; + const c = other.value; + return Number.isInteger(c) && c >= 0 && c <= m ? m + 1 : null; + }; + return masked(test.left, test.right) ?? masked(test.right, test.left); +} + +/** The attachment that asks LLVM to unroll a loop by `period`, for its + * back edge. */ +export function loopUnrollAttachment(period: number): string { + return `, !llvm.loop !${loopIds.get(period)!}`; +} + +/** The loop ID nodes that the module's text `lines` refer to. */ +export function loopUnrollMetadata(lines: readonly string[]): string[] { + const used = new Set(); + for (const line of lines) { + if (!line.includes("!llvm.loop !")) continue; + for (const [period, id] of loopIds) if (line.includes(`!llvm.loop !${id}`)) used.add(period); + } + return [...used] + .sort((a, b) => a - b) + .map((period) => { + const id = loopIds.get(period)!; + return `!${id} = distinct !{!${id}, !{!"llvm.loop.unroll.count", i32 ${period}}}`; + }); +} diff --git a/packages/compiler/test/periodic-guard-emission.test.ts b/packages/compiler/test/periodic-guard-emission.test.ts new file mode 100644 index 000000000..dfb8c60ae --- /dev/null +++ b/packages/compiler/test/periodic-guard-emission.test.ts @@ -0,0 +1,53 @@ +import { mkdtemp, readFile, rm, writeFile } from "node:fs/promises"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { expect, test } from "vitest"; +import { compile } from "../src/index.js"; + +async function emitLlvm(source: string): Promise { + const dir = await mkdtemp(join(tmpdir(), "scriptc-periodic-guards-")); + try { + const entry = join(dir, "main.ts"); + const outPath = join(dir, "main.ll"); + await writeFile(entry, source); + const result = await compile(entry, { outDir: dir, outPath, outputKind: "llvm" }); + if (!result.ok) + throw new Error(result.diagnostics.map((d) => `${d.code}: ${d.message}`).join("\n")); + return await readFile(outPath, "utf8"); + } finally { + await rm(dir, { recursive: true, force: true }); + } +} + +test("a counted loop with a check every fourth iteration asks LLVM to unroll it by 4", async () => { + const llvm = await emitLlvm(` +function escape(c: number): number { + let z = 0; + let m = 0; + for (let j = 0; j < 50; j++) { + z = z * z + c; + m = z * z; + if ((j & 3) === 3 && !(m <= 4)) break; + } + return m; +} +console.log(escape(Number(process.argv[2] ?? "0.25"))); +`); + expect(llvm).toMatch(/br label %loop\.c\d+, !llvm\.loop !1000004/); + expect(llvm).toContain(`!1000004 = distinct !{!1000004, !{!"llvm.loop.unroll.count", i32 4}}`); +}); + +test("a loop without a periodic guard has no loop metadata", async () => { + const llvm = await emitLlvm(` +function sum(n: number): number { + let s = 0; + for (let j = 0; j < n; j++) { + s += j; + if (s > 1000) break; + } + return s; +} +console.log(sum(Number(process.argv[2] ?? "10"))); +`); + expect(llvm).not.toContain("!llvm.loop"); +});