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"); +});