diff --git a/packages/compiler/src/backend/llvm/blocks.ts b/packages/compiler/src/backend/llvm/blocks.ts index 6d22bf3bf..354e631c6 100644 --- a/packages/compiler/src/backend/llvm/blocks.ts +++ b/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 diff --git a/packages/compiler/src/backend/llvm/emit-function.ts b/packages/compiler/src/backend/llvm/emit-function.ts index 9e8f7d701..135ebc706 100644 --- a/packages/compiler/src/backend/llvm/emit-function.ts +++ b/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(); diff --git a/packages/compiler/src/backend/llvm/emit-statements.ts b/packages/compiler/src/backend/llvm/emit-statements.ts index 7b4c38ccd..39ff5dcd3 100644 --- a/packages/compiler/src/backend/llvm/emit-statements.ts +++ b/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); diff --git a/packages/compiler/src/backend/llvm/emitter.ts b/packages/compiler/src/backend/llvm/emitter.ts index 7b7c742c5..7adf2a89f 100644 --- a/packages/compiler/src/backend/llvm/emitter.ts +++ b/packages/compiler/src/backend/llvm/emitter.ts @@ -846,6 +846,9 @@ export class LlEmitter { readonly callLifetimes: CallLifetimes; readonly constantCallbacks: Map>; currentConstantCallbacks: ReadonlyMap = new Map(); + /** The current function's running maximum and minimum updates, whose + * branches are weighted as rarely taken (running-extrema.ts). */ + runningExtrema: ReadonlySet = 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). */ diff --git a/packages/compiler/src/backend/llvm/running-extrema.test.ts b/packages/compiler/src/backend/llvm/running-extrema.test.ts new file mode 100644 index 000000000..4ce307c5c --- /dev/null +++ b/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); +}); diff --git a/packages/compiler/src/backend/llvm/running-extrema.ts b/packages/compiler/src/backend/llvm/running-extrema.ts new file mode 100644 index 000000000..1ddea603d --- /dev/null +++ b/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 { + const found = new Set(); + 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)[k], (y as Record)[k]), + ); + }; + return same(a, b); +} diff --git a/packages/compiler/test/running-extrema-branches.test.ts b/packages/compiler/test/running-extrema-branches.test.ts new file mode 100644 index 000000000..e1f5a7837 --- /dev/null +++ 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 { + 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"); +}); diff --git a/packages/compiler/test/ts7/baselines/order-parity/tests/corpus/7301-running-extrema-branches.ts.json b/packages/compiler/test/ts7/baselines/order-parity/tests/corpus/7301-running-extrema-branches.ts.json new file mode 100644 index 000000000..6dcb4b57d --- /dev/null +++ b/packages/compiler/test/ts7/baselines/order-parity/tests/corpus/7301-running-extrema-branches.ts.json @@ -0,0 +1,6 @@ +{ + "order": [ + "/tests/corpus/7301-running-extrema-branches.ts" + ], + "diags": [] +} diff --git a/tests/corpus/7301-running-extrema-branches.ts b/tests/corpus/7301-running-extrema-branches.ts new file mode 100644 index 000000000..9941ce538 --- /dev/null +++ b/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([])); diff --git a/tests/harness/corpus-costs/7301-running-extrema-branches.ts.json b/tests/harness/corpus-costs/7301-running-extrema-branches.ts.json new file mode 100644 index 000000000..74df88f0c --- /dev/null +++ b/tests/harness/corpus-costs/7301-running-extrema-branches.ts.json @@ -0,0 +1 @@ +{ "plain": 5300, "san": 5800 }