Unroll counted loops by the period of a guard on the counter's low bits (#2341)

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 {

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;
 }

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);

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([]);
+});

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<number, number>([
+  [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<number>();
+  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}}}`;
+    });
+}

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