Keep running maximum and minimum updates as branches (#1826)

packages/compiler/src/backend/llvm/blocks.ts

@@ -95,8 +95,8 @@ export class BlockBuilder {
     this.terminate(`br label %${label}`);
   }
 
-  condBr(v: string, t: string, f: string): void {
-    this.terminate(`br i1 ${v}, label %${t}, label %${f}`);
+  condBr(v: string, t: string, f: string, attachment = ""): void {
+    this.terminate(`br i1 ${v}, label %${t}, label %${f}${attachment}`);
   }
 
   /** Emit a double-indexed `for (i = 0; i < len; i++)` loop. `next` is

packages/compiler/src/backend/llvm/emit-function.ts

@@ -55,6 +55,7 @@ import { INT32_ABI_SUFFIX, SCALAR_UNION_RESULT } from "./int32-abi.js";
 import { f64Lit } from "./common.js";
 import { RESULT_ENTRY_PREFIX } from "./record-results.js";
 import { emitIntegerClone, emitIntegerDispatcher } from "./integer-clones.js";
+import { runningExtremumUpdates } from "./running-extrema.js";
 
 export function emitFunction(emitter: LlEmitter, fn: IrFunction): string {
   const B = new BlockBuilder();
@@ -81,6 +82,7 @@ export function emitFunction(emitter: LlEmitter, fn: IrFunction): string {
     emitter.currentLocals.set(id, plain);
     emitter.lazyCaptureBoxes.set(id, "");
   }
+  emitter.runningExtrema = emitter.debug === null ? runningExtremumUpdates(fn) : new Set();
   emitter.currentConstantCallbacks =
     emitter.debug === null ? (emitter.constantCallbacks.get(fn.name) ?? new Map()) : new Map();
   emitter.knownCallbacks.clear();

packages/compiler/src/backend/llvm/emit-statements.ts

@@ -786,7 +786,12 @@ export function emitStmtBody(emitter: LlEmitter, s: IrStmt): void {
       const trueLabel = B.newLabel("if.t");
       const joinLabel = B.newLabel("if.j");
       const falseLabel = s.else_ ? B.newLabel("if.f") : joinLabel;
-      B.condBr(cond, trueLabel, falseLabel);
+      // A running maximum or minimum is rarely updated; the weights keep
+      // the branch, where a select would make each element wait on the last.
+      const weights = emitter.runningExtrema.has(s)
+        ? `, !prof !{!"branch_weights", i32 1, i32 1000}`
+        : "";
+      B.condBr(cond, trueLabel, falseLabel, weights);
       B.startBlock(trueLabel);
       emitter.emitBlock(s.then);
       B.br(joinLabel);

packages/compiler/src/backend/llvm/emitter.ts

@@ -846,6 +846,9 @@ export class LlEmitter {
   readonly callLifetimes: CallLifetimes;
   readonly constantCallbacks: Map<string, Map<string, string>>;
   currentConstantCallbacks: ReadonlyMap<string, string> = new Map();
+  /** The current function's running maximum and minimum updates, whose
+   * branches are weighted as rarely taken (running-extrema.ts). */
+  runningExtrema: ReadonlySet<IrStmt> = new Set();
   /** Callback bindings whose closures an enclosing callback loop version
    * proved to hold their known function: calls take the direct entry with
    * no fallback (emitCallbackVersionedLoop). */

packages/compiler/src/backend/llvm/running-extrema.test.ts

@@ -0,0 +1,97 @@
+import { expect, test } from "vitest";
+import { BOOL, F64, type IrExpr, type IrFunction, type IrStmt } from "../../ir/ir.js";
+import { runningExtremumUpdates } from "./running-extrema.js";
+
+const loc = { file: "extrema.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 assign = (localId: string, value: IrExpr): IrStmt => ({
+  kind: "assign",
+  localId,
+  value,
+  loc,
+});
+const decl = (localId: string): IrStmt => ({ kind: "varDecl", localId, init: num(0), loc });
+const compare = (op: "<" | ">", left: IrExpr, right: IrExpr): IrExpr => ({
+  kind: "bin",
+  op,
+  left,
+  right,
+  type: BOOL,
+  loc,
+});
+/** `if (v > max) max = v;` */
+const update = (): IrStmt & { kind: "if" } => ({
+  kind: "if",
+  cond: compare(">", ref("v"), ref("max")),
+  then: [assign("max", ref("v"))],
+  else_: null,
+  loc,
+});
+const loop = (body: IrStmt[]): IrStmt => ({
+  kind: "while",
+  cond: compare(">", ref("n"), num(0)),
+  body,
+  loc,
+});
+const fn = (body: IrStmt[], captures: string[] = []): IrFunction => ({
+  name: "scan",
+  params: [],
+  locals: ["v", "max", "n"].map((id) => ({ id, name: id, type: F64, mutable: true })),
+  captures: captures.map((localId) => ({ localId, name: localId, type: F64 })),
+  returnType: F64,
+  body,
+  loc,
+});
+
+test("a loop's update of an accumulator declared outside it qualifies", () => {
+  const s = update();
+  expect(runningExtremumUpdates(fn([decl("max"), loop([decl("v"), s])])).has(s)).toBe(true);
+});
+
+test("the local may be on either side of the comparison", () => {
+  const s = update();
+  s.cond = compare("<", ref("max"), ref("v"));
+  expect(runningExtremumUpdates(fn([decl("max"), loop([s])])).has(s)).toBe(true);
+});
+
+test("a local declared in the loop is not an accumulator", () => {
+  const s = update();
+  expect(runningExtremumUpdates(fn([loop([decl("max"), s])])).size).toBe(0);
+});
+
+test("another write in the loop disqualifies the update", () => {
+  const s = update();
+  const body = [decl("max"), loop([assign("max", num(1)), s])];
+  expect(runningExtremumUpdates(fn(body)).size).toBe(0);
+});
+
+test("the assigned value must be the compared expression", () => {
+  const s = update();
+  s.then = [assign("max", num(5))];
+  expect(runningExtremumUpdates(fn([decl("max"), loop([s])])).size).toBe(0);
+});
+
+test("outside a loop only a captured binding qualifies", () => {
+  const captured = update();
+  expect(runningExtremumUpdates(fn([captured], ["max"])).has(captured)).toBe(true);
+  expect(runningExtremumUpdates(fn([update()])).size).toBe(0);
+});
+
+test("an else branch disqualifies the update", () => {
+  const s = update();
+  s.else_ = [assign("v", num(0))];
+  expect(runningExtremumUpdates(fn([decl("max"), loop([s])])).size).toBe(0);
+});
+
+test("an update nested in another branch does not qualify", () => {
+  const s = update();
+  const guard: IrStmt = {
+    kind: "if",
+    cond: compare(">", ref("v"), num(0)),
+    then: [s],
+    else_: null,
+    loc,
+  };
+  expect(runningExtremumUpdates(fn([decl("max"), loop([guard])])).size).toBe(0);
+});

packages/compiler/src/backend/llvm/running-extrema.ts

@@ -0,0 +1,110 @@
+// Running maximum and minimum updates, `if (v > max) { max = v; ... }`.
+// With the accumulator in a register LLVM turns such an update into a
+// select (`maxsd`, or a compare and mask), so every element waits on the
+// previous one's. The update gets rarer as a scan goes on, so a branch
+// predicts well; emitStmtBody gives these `if`s branch weights that say the
+// update is rarely taken, and LLVM keeps the branch. A clamp of a value read
+// in the same iteration is not an accumulator and keeps its select: its
+// branch would go whichever way the data goes. An update nested in another
+// branch keeps its select too: kept, its branch can stop LLVM from turning
+// the enclosing branch into selects, and that branch may follow the data.
+
+import type { IrExpr, IrFunction, IrStmt } from "../../ir/ir.js";
+import { everyExprChild, everyStmtChild } from "../../ir/traverse.js";
+
+type Loop = IrStmt & { kind: "for" | "while" | "doWhile" | "forOf" };
+
+const RELATIONAL = new Set(["<", "<=", ">", ">="]);
+
+/** The `if` statements of `fn` that update a running maximum or minimum:
+ * no `else`, a numeric `<`, `<=`, `>` or
+ * `>=` between a local and another expression, and a then-block that
+ * assigns that expression to the local. The local must carry across
+ * iterations: inside a loop the `if` is a statement of the innermost loop's
+ * body, the local is declared outside that loop, and the loop writes it only
+ * in this `if`; outside any loop the `if` is a statement of the function
+ * body and the local is a binding the function captures and writes only in
+ * this `if` (a callback that a loop elsewhere calls). */
+export function runningExtremumUpdates(fn: IrFunction): Set<IrStmt> {
+  const found = new Set<IrStmt>();
+  const captured = new Set((fn.captures ?? []).map((c) => c.localId));
+  const loops: Loop[] = [];
+  const candidates: { s: IrStmt; local: string; loop: Loop | undefined }[] = [];
+  const expr = (e: IrExpr): boolean => everyExprChild(e, expr, stmt);
+  const stmt = (s: IrStmt): boolean => {
+    if (s.kind === "if") {
+      const local = updatedExtremum(s);
+      if (local !== undefined) candidates.push({ s, local, loop: loops[loops.length - 1] });
+    }
+    const loop =
+      s.kind === "for" || s.kind === "while" || s.kind === "doWhile" || s.kind === "forOf";
+    if (loop) loops.push(s);
+    everyStmtChild(s, expr, stmt);
+    if (loop) loops.pop();
+    return true;
+  };
+  for (const s of fn.body) stmt(s);
+  for (const { s, local, loop } of candidates) {
+    const scope = loop ?? (captured.has(local) ? fn.body : undefined);
+    const body = loop ? loop.body : fn.body;
+    if (scope !== undefined && body.includes(s) && writes(scope, local) === 1) found.add(s);
+  }
+  return found;
+}
+
+/** The local an `if` updates to a new maximum or minimum, if it is one. */
+function updatedExtremum(s: IrStmt & { kind: "if" }): string | undefined {
+  if (s.else_ !== null && s.else_.length > 0) return undefined;
+  const c = s.cond;
+  if (c.kind !== "bin" || !RELATIONAL.has(c.op)) return undefined;
+  if (c.left.type.kind !== "f64" || c.right.type.kind !== "f64") return undefined;
+  for (const [local, other] of [
+    [c.left, c.right],
+    [c.right, c.left],
+  ]) {
+    if (local!.kind !== "varRef") continue;
+    const id = local!.localId;
+    const assigns = s.then.some(
+      (x) => x.kind === "assign" && x.localId === id && sameExpr(x.value, other!),
+    );
+    if (assigns) return id;
+  }
+  return undefined;
+}
+
+/** Declarations and writes of `local` under `scope`. A declaration counts
+ * as a write too, so a local declared inside a loop does not qualify. */
+function writes(scope: IrStmt | IrStmt[], local: string): number {
+  let n = 0;
+  const expr = (e: IrExpr): boolean => {
+    if ((e.kind === "assignExpr" || e.kind === "incDec") && e.localId === local) n++;
+    return everyExprChild(e, expr, stmt);
+  };
+  const stmt = (s: IrStmt): boolean => {
+    if ((s.kind === "assign" || s.kind === "varDecl" || s.kind === "forOf") && s.localId === local)
+      n++;
+    else if (s.kind === "tryCatch" && s.catchLocalId === local) n++;
+    return everyStmtChild(s, expr, stmt);
+  };
+  if (Array.isArray(scope)) for (const s of scope) stmt(s);
+  else stmt(scope);
+  return n;
+}
+
+/** Structural equality of two small expressions (`v`, `a[i]`), ignoring
+ * source locations and types; larger expressions do not compare equal. */
+function sameExpr(a: IrExpr, b: IrExpr): boolean {
+  let budget = 64;
+  const same = (x: unknown, y: unknown): boolean => {
+    if (--budget < 0) return false;
+    if (typeof x !== "object" || x === null || typeof y !== "object" || y === null) return x === y;
+    if (Array.isArray(x) !== Array.isArray(y)) return false;
+    const keys = Object.keys(x).filter((k) => k !== "loc" && k !== "type");
+    if (keys.length !== Object.keys(y).filter((k) => k !== "loc" && k !== "type").length)
+      return false;
+    return keys.every((k) =>
+      same((x as Record<string, unknown>)[k], (y as Record<string, unknown>)[k]),
+    );
+  };
+  return same(a, b);
+}

packages/compiler/test/running-extrema-branches.test.ts

@@ -0,0 +1,75 @@
+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-running-extrema-"));
+  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.message).join("\n"));
+    return await readFile(outPath, "utf8");
+  } finally {
+    await rm(dir, { recursive: true, force: true });
+  }
+}
+
+const RARE =
+  /br i1 %\S+, label %if\.t\d+, label %if\.j\d+, !prof !\{!"branch_weights", i32 1, i32 1000\}/;
+
+test("a running maximum's update is weighted as rarely taken", async () => {
+  const llvm = await emitLlvm(`
+    function argmax(a: Float64Array): number {
+      let max = -Infinity;
+      let at = -1;
+      for (let i = 0; i < a.length; i++) {
+        const v = a[i]!;
+        if (v > max) { max = v; at = i; }
+      }
+      return at;
+    }
+    console.log(argmax(new Float64Array([1, 3, 2])));
+  `);
+  expect(llvm).toMatch(RARE);
+});
+
+test("a callback's update of a captured minimum is weighted", async () => {
+  const llvm = await emitLlvm(`
+    function least(values: number[]): number {
+      let min = Infinity;
+      values.forEach((v) => { if (min > v) min = v; });
+      return min;
+    }
+    console.log(least([3, 1, 2]));
+  `);
+  expect(llvm).toMatch(RARE);
+});
+
+test("clamps and saturating sums keep their selects", async () => {
+  const llvm = await emitLlvm(`
+    function clampSum(a: Float64Array, hi: number): number {
+      let s = 0;
+      for (let i = 0; i < a.length; i++) {
+        let v = a[i]!;
+        if (hi < v) v = hi;
+        s += v;
+      }
+      return s;
+    }
+    function saturate(a: Float64Array, hi: number): number {
+      let v = 0;
+      for (let i = 0; i < a.length; i++) {
+        v += a[i]!;
+        if (hi < v) v = hi;
+      }
+      return v;
+    }
+    const a = new Float64Array([0.5, 2, 0.25]);
+    console.log(clampSum(a, 1), saturate(a, 1));
+  `);
+  expect(llvm).not.toContain("branch_weights");
+});

packages/compiler/test/ts7/baselines/order-parity/tests/corpus/7301-running-extrema-branches.ts.json

@@ -0,0 +1,6 @@
+{
+ "order": [
+  "<repo>/tests/corpus/7301-running-extrema-branches.ts"
+ ],
+ "diags": []
+}

tests/corpus/7301-running-extrema-branches.ts

@@ -0,0 +1,169 @@
+// Running maximum and minimum scans over ascending, descending, alternating,
+// pseudo-random and NaN-holding data: plain for, while, do-while and for-of
+// loops, `<=` and `>=` ties, an argmax that records its index, typed array
+// and array element reads, callbacks that update captured bindings, a scan
+// nested in an outer loop, an update nested in another branch, and the clamp
+// and saturating updates that are not running extrema.
+
+function shapes(n: number): [string, Float64Array][] {
+  const ascending = new Float64Array(n);
+  const descending = new Float64Array(n);
+  const alternating = new Float64Array(n);
+  const random = new Float64Array(n);
+  const holes = new Float64Array(n);
+  let x = 12345;
+  for (let i = 0; i < n; i++) {
+    ascending[i] = i;
+    descending[i] = n - i;
+    alternating[i] = i % 2 === 0 ? 0 : i;
+    x = (x * 1103515245 + 12345) % 2147483648;
+    random[i] = (x % 1000) - 500;
+    holes[i] = i % 7 === 3 ? NaN : (x % 97) - 48;
+  }
+  return [
+    ["ascending", ascending],
+    ["descending", descending],
+    ["alternating", alternating],
+    ["random", random],
+    ["nan", holes],
+  ];
+}
+
+function argmax(a: Float64Array): string {
+  let max = -Infinity;
+  let at = -1;
+  for (let i = 0; i < a.length; i++) {
+    const v = a[i]!;
+    if (v > max) {
+      max = v;
+      at = i;
+    }
+  }
+  return `${max}@${at}`;
+}
+
+function lastMin(a: Float64Array): string {
+  let min = Infinity;
+  let at = -1;
+  let i = 0;
+  while (i < a.length) {
+    if (a[i]! <= min) {
+      min = a[i]!;
+      at = i;
+    }
+    i++;
+  }
+  return `${min}@${at}`;
+}
+
+function bothEnds(values: number[]): string {
+  let lo = Infinity;
+  let hi = -Infinity;
+  for (const v of values) {
+    if (lo > v) lo = v;
+    if (v >= hi) hi = v;
+  }
+  return `${lo}..${hi}`;
+}
+
+function firstDrop(a: Float64Array): number {
+  let floor = a[0]!;
+  let i = 0;
+  do {
+    if (a[i]! < floor) floor = a[i]!;
+    i++;
+  } while (i < a.length);
+  return floor;
+}
+
+function callbackPeak(a: Float64Array): string {
+  let max = -Infinity;
+  let at = -1;
+  a.forEach((v, i) => {
+    if (v > max) {
+      max = v;
+      at = i;
+    }
+  });
+  return `${max}@${at}`;
+}
+
+function arrayCallbackMin(values: number[]): number {
+  let min = Infinity;
+  values.forEach((v) => {
+    if (v < min) min = v;
+  });
+  return min;
+}
+
+function rowMaxima(a: Float64Array, width: number): string {
+  const out: number[] = [];
+  for (let row = 0; row * width < a.length; row++) {
+    let best = -Infinity;
+    for (let j = row * width; j < Math.min(a.length, (row + 1) * width); j++) {
+      if (a[j]! > best) best = a[j]!;
+    }
+    out.push(best);
+  }
+  return out.join(",");
+}
+
+function longestNonNegativeRun(a: Float64Array): number {
+  let longest = 0;
+  let run = 0;
+  for (let i = 0; i < a.length; i++) {
+    if (a[i]! >= 0) {
+      run++;
+      if (run > longest) longest = run;
+    } else run = 0;
+  }
+  return longest;
+}
+
+function clampSum(a: Float64Array, hi: number): number {
+  let s = 0;
+  for (let i = 0; i < a.length; i++) {
+    let v = a[i]!;
+    if (hi < v) v = hi;
+    s += v;
+  }
+  return s;
+}
+
+function saturate(a: Float64Array, hi: number): number {
+  let v = 0;
+  for (let i = 0; i < a.length; i++) {
+    v += a[i]!;
+    if (hi < v) v = hi;
+  }
+  return v;
+}
+
+for (const [name, a] of shapes(40)) {
+  const values = Array.from(a);
+  console.log(
+    name,
+    argmax(a),
+    lastMin(a),
+    bothEnds(values),
+    firstDrop(a),
+    callbackPeak(a),
+    arrayCallbackMin(values),
+  );
+  console.log(
+    " rows",
+    rowMaxima(a, 7),
+    "run",
+    longestNonNegativeRun(a),
+    "clamp",
+    clampSum(a, 5),
+    "saturate",
+    saturate(a, 30),
+  );
+}
+
+const signed = new Float64Array([0, -0, 0, -0]);
+let top = -Infinity;
+for (let i = 0; i < signed.length; i++) if (signed[i]! >= top) top = signed[i]!;
+console.log("signed zero", Object.is(top, -0), argmax(signed), lastMin(signed));
+console.log("empty", argmax(new Float64Array(0)), bothEnds([]), arrayCallbackMin([]));

tests/harness/corpus-costs/7301-running-extrema-branches.ts.json

@@ -0,0 +1 @@
+{ "plain": 5300, "san": 5800 }