diff --git a/packages/compiler/src/backend/llvm/emit-function.ts b/packages/compiler/src/backend/llvm/emit-function.ts index 09e6551c3..9e8f7d701 100644 --- a/packages/compiler/src/backend/llvm/emit-function.ts +++ b/packages/compiler/src/backend/llvm/emit-function.ts @@ -83,6 +83,8 @@ export function emitFunction(emitter: LlEmitter, fn: IrFunction): string { } emitter.currentConstantCallbacks = emitter.debug === null ? (emitter.constantCallbacks.get(fn.name) ?? new Map()) : new Map(); + emitter.knownCallbacks.clear(); + emitter.callbackVersionsEnabled = emitter.debug === null && !fn.async && !fn.generator; // Preserve concrete source bindings for debugger inspection. Suspended // functions keep the established array lifetime across continuations. emitter.streamingSplitsEnabled = emitter.debug === null && !fn.async && !fn.generator; diff --git a/packages/compiler/src/backend/llvm/emit-statements.ts b/packages/compiler/src/backend/llvm/emit-statements.ts index a22de3445..4fe1a27bd 100644 --- a/packages/compiler/src/backend/llvm/emit-statements.ts +++ b/packages/compiler/src/backend/llvm/emit-statements.ts @@ -55,6 +55,7 @@ import { emitBytesLoopVersionGuard } from "./bytes-loop-versions.js"; import { mangleLocal } from "../mangle.js"; import { LlvmUnsupportedError } from "./unsupported.js"; import type { LlValue } from "./expr-context.js"; +import { everyExprChild, everyStmtChild } from "../../ir/traverse.js"; import { widenInteger } from "./integer-values.js"; import { SCALAR_UNION_RESULT } from "./int32-abi.js"; import { emitLiteralSwitch } from "./switch-dispatch.js"; @@ -76,6 +77,9 @@ import { } from "./string-construction.js"; import { emitStringInputs } from "./string-lifetimes.js"; +/** A callback loop version copies the loop; keep the copies small. */ +const MAX_CALLBACK_VERSION_COST = 400; + export function emitStmtBody(emitter: LlEmitter, s: IrStmt): void { const B = emitter.B; emitter.frames.push([]); @@ -795,6 +799,11 @@ export function emitStmtBody(emitter: LlEmitter, s: IrStmt): void { break; } case "while": { + const callbacks = callbackVersionBindings(emitter, s); + if (callbacks) { + emitCallbackVersionedLoop(emitter, s, callbacks); + break; + } const workerBudget = emitter.workerLoopBudget(s); const le = B.newLabel("loop.e"); emitter.emitLoopCopies(s, true, le, (exit) => { @@ -828,6 +837,11 @@ export function emitStmtBody(emitter: LlEmitter, s: IrStmt): void { break; } case "doWhile": { + const callbacks = callbackVersionBindings(emitter, s); + if (callbacks) { + emitCallbackVersionedLoop(emitter, s, callbacks); + break; + } const workerBudget = emitter.workerLoopBudget(s); // Body first (runs at least once); continue jumps to the CONDITION. const le = B.newLabel("loop.e"); @@ -857,6 +871,11 @@ export function emitStmtBody(emitter: LlEmitter, s: IrStmt): void { break; } case "for": { + const callbacks = callbackVersionBindings(emitter, s); + if (callbacks) { + emitCallbackVersionedLoop(emitter, s, callbacks); + break; + } // The init's scope wraps the whole loop (break/continue must NOT // release it — scopeDepth captured after the push, C parity). const integerArrayLoop = @@ -1888,6 +1907,101 @@ export function emitSwitch(emitter: LlEmitter, s: IrStmt & { kind: "switch" }): B.startBlock(end); } +/** The loop copies a callback loop version needs: the bindings the loop + * calls that the constant-callback census resolves to one function, each + * with that function's direct entry. Only small loops of optimized + * synchronous functions, and only bindings declared before the loop, + * qualify: the guard reads each binding before the loop runs. Versions do + * not nest: a loop inside either copy of a version (including the + * versioned loop itself, emitted again for each copy) keeps its calls as + * they are, so a loop nest is emitted at most twice. */ +function callbackVersionBindings( + emitter: LlEmitter, + s: IrStmt & { kind: "for" | "while" | "doWhile" }, +): { ref: IrExpr & { kind: "varRef" }; direct: string }[] | null { + if ( + !emitter.callbackVersionsEnabled || + emitter.currentConstantCallbacks.size === 0 || + emitter.callbackVersioning.size > 0 + ) + return null; + const found = new Map(); + // Bindings the loop declares or writes (a callback declared in its body). + const bound = new Set(); + let cost = 0; + const expr = (e: IrExpr): boolean => { + if (++cost > MAX_CALLBACK_VERSION_COST) return false; + if (e.kind === "assignExpr") bound.add(e.localId); + if (e.kind === "callValue" && e.callee.kind === "varRef" && !found.has(e.callee.localId)) { + const local = emitter.currentLocals.get(e.callee.localId); + const target = + local && local.tdz !== true ? emitter.knownCallbackTarget(e.callee) : undefined; + if (target) found.set(e.callee.localId, { ref: e.callee, direct: target.direct }); + } + return everyExprChild(e, expr, stmt); + }; + const stmt = (x: IrStmt): boolean => { + if (++cost > MAX_CALLBACK_VERSION_COST) return false; + if (x.kind === "varDecl" || x.kind === "assign" || x.kind === "forOf") bound.add(x.localId); + else if (x.kind === "tryCatch" && x.catchLocalId !== null) bound.add(x.catchLocalId); + return everyStmtChild(x, expr, stmt); + }; + if (!stmt(s)) return null; + const bindings = [...found.values()].filter(({ ref }) => !bound.has(ref.localId)); + return bindings.length > 0 ? bindings : null; +} + +/** A callback loop version. The guard tests, once before the loop, that + * each binding's closure holds its known function: a closure's code + * pointer is written only when the closure is built, and the census + * resolves only bindings the function never rebinds. The copy where the + * guard holds calls those functions directly, with no indirect fallback + * that could rewrite the closure or what it captures, so LLVM can keep a + * captured scalar the callback updates in a register. The other copy is the + * ordinary loop. */ +function emitCallbackVersionedLoop( + emitter: LlEmitter, + s: IrStmt & { kind: "for" | "while" | "doWhile" }, + bindings: { ref: IrExpr & { kind: "varRef" }; direct: string }[], +): void { + const B = emitter.B; + let ok = "true"; + for (const { ref, direct } of bindings) { + // A borrowed read: the guard only compares the code pointer. + const closure = emitter.emitReadReceiver(ref).name; + const fnp = B.tmp(), + fn = B.tmp(), + same = B.tmp(); + B.line(`${fnp} = getelementptr inbounds %ScrClosure, ptr ${closure}, i64 0, i32 1`); + B.line(`${fn} = load ptr, ptr ${fnp}`); + B.line(`${same} = icmp eq ptr ${fn}, ${direct}`); + if (ok === "true") ok = same; + else { + const both = B.tmp(); + B.line(`${both} = and i1 ${ok}, ${same}`); + ok = both; + } + } + const fast = B.newLabel("callback.version.fast"), + slow = B.newLabel("callback.version.slow"), + done = B.newLabel("callback.version.done"); + B.condBr(ok, fast, slow); + emitter.callbackVersioning.add(s); + try { + B.startBlock(slow); + emitter.emitStmt(s); + B.br(done); + B.startBlock(fast); + for (const { ref } of bindings) emitter.knownCallbacks.add(ref.localId); + emitter.emitStmt(s); + B.br(done); + } finally { + for (const { ref } of bindings) emitter.knownCallbacks.delete(ref.localId); + emitter.callbackVersioning.delete(s); + } + B.startBlock(done); +} + /** Entries in a row whose fast copy bailed before a speculative loop site * stops trying it. */ const SPECULATION_BAILS = 8; diff --git a/packages/compiler/src/backend/llvm/emitter.ts b/packages/compiler/src/backend/llvm/emitter.ts index c1cf66af4..bfde59a30 100644 --- a/packages/compiler/src/backend/llvm/emitter.ts +++ b/packages/compiler/src/backend/llvm/emitter.ts @@ -831,6 +831,12 @@ export class LlEmitter { readonly callLifetimes: CallLifetimes; readonly constantCallbacks: Map>; currentConstantCallbacks: ReadonlyMap = new Map(); + /** Callback bindings whose closures an enclosing callback loop version + * proved to hold their known function: calls take the direct entry with + * no fallback (emitCallbackVersionedLoop). */ + knownCallbacks = new Set(); + readonly callbackVersioning = new Set(); + callbackVersionsEnabled = false; readonly stackCallbacks: StackCallbacks; readonly receiverReaders: AmbientReceiverReaders; readonly stackCaptures: StackCaptures; @@ -3912,6 +3918,25 @@ export class LlEmitter { return b; } + /** The function a call through `callee` is known to reach (a literal + * closure, or a binding the constant-callback census resolves) and the + * entry it calls directly. Async and generator functions stay indirect. */ + knownCallbackTarget(callee: IrExpr): { known: IrFunction; direct: string } | undefined { + const knownName = + callee.kind === "closure" + ? callee.fnName + : callee.kind === "varRef" + ? this.currentConstantCallbacks.get(callee.localId) + : undefined; + const known = knownName ? this.fnByName.get(knownName) : undefined; + if (!known || known.async || known.generator) return undefined; + if (known.captures === undefined) { + this.fnValues.add(known.name); + return { known, direct: `@${mangleWrapper(known.name)}` }; + } + return { known, direct: `@${this.callTarget(known.name)}` }; + } + /** The payload slot of a capture box (ScrBox.slot, field 5). */ boxSlot(box: string): string { const slot = this.B.tmp(); @@ -3925,12 +3950,18 @@ export class LlEmitter { const r = B.tmp(); // Scalar payloads are read in place (scr_box_get_f64/bool are a // memcpy and a nonzero test of the same slot), so LLVM can keep a - // captured number in a register across a loop of calls. + // captured number in a register across a loop of calls. They have an + // alias category of their own: every other access to a box's payload + // is untagged or a runtime call. if (acc === "f64") { - B.line(`${r} = load double, ptr ${this.boxSlot(box)}`); + B.line( + `${r} = load double, ptr ${this.boxSlot(box)}${this.memoryAliasAttachment("box:scalar")}`, + ); } else if (acc === "bool") { const raw = B.tmp(); - B.line(`${raw} = load i64, ptr ${this.boxSlot(box)}`); + B.line( + `${raw} = load i64, ptr ${this.boxSlot(box)}${this.memoryAliasAttachment("box:scalar")}`, + ); B.line(`${r} = icmp ne i64 ${raw}, 0`); } else { this.declare(`declare ptr @scr_box_get_ref(ptr)`); @@ -3944,11 +3975,15 @@ export class LlEmitter { const B = this.B; const acc = boxAccess(t); if (acc === "f64") { - B.line(`store double ${value}, ptr ${this.boxSlot(box)}`); + B.line( + `store double ${value}, ptr ${this.boxSlot(box)}${this.memoryAliasAttachment("box:scalar")}`, + ); } else if (acc === "bool") { const raw = B.tmp(); B.line(`${raw} = zext i1 ${value} to i64`); - B.line(`store i64 ${raw}, ptr ${this.boxSlot(box)}`); + B.line( + `store i64 ${raw}, ptr ${this.boxSlot(box)}${this.memoryAliasAttachment("box:scalar")}`, + ); } else { this.declare(`declare void @scr_box_set_ref(ptr, ptr)`); B.line(`call void @scr_box_set_ref(ptr ${box}, ptr ${value})`); diff --git a/packages/compiler/src/backend/llvm/expr-calls.ts b/packages/compiler/src/backend/llvm/expr-calls.ts index f0aa3b276..4bcf2a3d5 100644 --- a/packages/compiler/src/backend/llvm/expr-calls.ts +++ b/packages/compiler/src/backend/llvm/expr-calls.ts @@ -29,7 +29,6 @@ import { mangleClassStruct, mangleFnClosure, mangleFunction, - mangleWrapper, mangleVtStruct, mangleAsyncDirectSpawn, } from "../mangle.js"; @@ -839,21 +838,9 @@ export function emitCallExpr( host.declare("declare void @scr_dyn_this_push_dyn(ptr)"); B.line(`call void @scr_dyn_this_push_dyn(ptr ${receiver?.name ?? "null"})`); }; - const knownName = - e.callee.kind === "closure" - ? e.callee.fnName - : e.callee.kind === "varRef" - ? host.currentConstantCallbacks.get(e.callee.localId) - : undefined; - const candidate = knownName ? host.fnByName.get(knownName) : undefined; - const known = candidate && !candidate.async && !candidate.generator ? candidate : undefined; - let direct: string | undefined; - if (known) { - if (known.captures === undefined) { - host.fnValues.add(known.name); - direct = `@${mangleWrapper(known.name)}`; - } else direct = `@${host.callTarget(known.name)}`; - } + const target = host.knownCallbackTarget(e.callee); + const known = target?.known; + const direct = target?.direct; // A known closure body that borrows reference parameters is entered // directly with lent arguments, as a direct call to a borrowing body // is: an argument is read without a reference when neither a later @@ -939,7 +926,13 @@ export function emitCallExpr( let pending = host.indirectMayThrow; let pendingGuard: string | undefined; let t: string; - if (direct && e.callee.kind === "closure") { + // A literal closure, or a binding an enclosing callback loop version + // proved to hold its known function, enters it directly. + if ( + direct && + (e.callee.kind === "closure" || + (e.callee.kind === "varRef" && host.knownCallbacks.has(e.callee.localId))) + ) { t = invoke(lending ?? direct, !host.receiverReaders.ignores(known!)); if (quiet) pending = false; } else { diff --git a/packages/compiler/test/callback-loop-versions.test.ts b/packages/compiler/test/callback-loop-versions.test.ts new file mode 100644 index 000000000..63a669fd4 --- /dev/null +++ b/packages/compiler/test/callback-loop-versions.test.ts @@ -0,0 +1,114 @@ +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-callback-versions-")); + 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 body = (llvm: string, marker: string): string => { + const match = new RegExp(`^define internal [^\\n]*${marker}[^]*?^}`, "m").exec(llvm); + if (!match) throw new Error(`no function ${marker}`); + return match[0]; +}; + +test("loops that call a known callback test its closure once and call it directly", async () => { + const llvm = await emitLlvm(` + function each(values: number[], visit: (v: number) => void): void { + for (let i = 0; i < values.length; i++) visit(values[i]!); + } + let sum = 0; + each([1, 2, 3], (v) => { sum += v; }); + console.log(sum); + `); + const fn = body(llvm, "@sc_bf_each"); + const fast = /^callback\.version\.fast\d+:$[^]*?^callback\.version\.done\d+:$/m.exec(fn)![0]; + // The guard compares the closure's code pointer once, before the loop. + expect(fn.slice(0, fn.indexOf(fast))).toMatch(/icmp eq ptr %\S+, @/); + expect(fast).not.toContain("callback.indirect"); + expect(fast).not.toContain("scr_dyn_this_push_dyn"); + // The ordinary copy keeps the fallback through the closure's pointer. + expect(fn).toContain("callback.indirect"); +}); + +test("captured number payloads carry their own alias tag", async () => { + const llvm = await emitLlvm(` + function total(values: number[]): number { + let sum = 0; + values.forEach((v) => { sum += v; }); + return sum; + } + console.log(total([1, 2, 3])); + `); + // Every read and write through a box's payload slot (field 5) is tagged. + const slots = [ + ...llvm.matchAll(/(%\S+) = getelementptr inbounds %ScrBox, ptr %\S+, i32 0, i32 5/g), + ]; + expect(slots.length).toBeGreaterThan(0); + for (const [, slot] of slots) { + for (const access of llvm.matchAll( + new RegExp(`(load double, ptr|store double %\\S+, ptr) \\${slot}\\b[^\\n]*`, "g"), + )) + expect(access[0]).toMatch(/!tbaa !\d+$/); + } + expect(llvm).toMatch(/load double, ptr %\S+, !tbaa !\d+/); +}); + +test("a callback declared in the loop body is not read before the loop", async () => { + const llvm = await emitLlvm(` + function strides(): number { + let calls = 0; + for (let i = 0; i < 10; i++) { + const bump = (): void => { i += 2; }; + bump(); + calls++; + } + return calls; + } + console.log(strides()); + `); + expect(llvm).not.toContain("callback.version"); +}); + +test("loops inside a callback version are not versioned again", async () => { + // Each inner loop calls a callback a factory builds just before it; without + // the limit, every copy of a version would version the next loop again. + const llvm = await emitLlvm(` + function adder(k: number): (v: number) => number { + return (v: number) => v + k; + } + function nest(n: number, f: (v: number) => number): number { + let s = 0; + for (let i = 0; i < n; i++) { + s += f(i); + const g = adder(i); + for (let j = 0; j < n; j++) { + s += g(j); + const h = adder(j); + for (let k = 0; k < n; k++) s += h(k); + } + } + return s; + } + console.log(nest(3, (v) => v - 1), nest(2, (v) => v * 3)); + `); + // One version for the outer loop of each compiled `nest`; its two copies + // keep the inner loops whole. + const versions = [...llvm.matchAll(/^define [^\n]*@sc_bf_nest[^]*?^}/gm)].map( + ([body]) => body.match(/^callback\.version\.fast\d+:$/gm)?.length ?? 0, + ); + expect(versions.every((count) => count <= 1)).toBe(true); + expect(versions.some((count) => count === 1)).toBe(true); +}); diff --git a/packages/compiler/test/ts7/baselines/order-parity/tests/corpus/7300-callback-captured-scalars.ts.json b/packages/compiler/test/ts7/baselines/order-parity/tests/corpus/7300-callback-captured-scalars.ts.json new file mode 100644 index 000000000..dc1158934 --- /dev/null +++ b/packages/compiler/test/ts7/baselines/order-parity/tests/corpus/7300-callback-captured-scalars.ts.json @@ -0,0 +1,6 @@ +{ + "order": [ + "/tests/corpus/7300-callback-captured-scalars.ts" + ], + "diags": [] +} diff --git a/tests/corpus/7300-callback-captured-scalars.ts b/tests/corpus/7300-callback-captured-scalars.ts new file mode 100644 index 000000000..c8e8b28bd --- /dev/null +++ b/tests/corpus/7300-callback-captured-scalars.ts @@ -0,0 +1,147 @@ +// Callbacks that update captured `let` numbers and booleans from loops: typed +// array and array methods, an authored higher-order function called with one +// function at several calls and with two functions at one, closures of one +// function over different bindings, a callback that throws part way, nested +// callbacks, callbacks declared in the loop body, and reads after the loop. + +function each(values: number[], visit: (v: number, i: number) => void): void { + for (let i = 0; i < values.length; i++) visit(values[i]!, i); +} + +function until(limit: number, step: (i: number) => boolean): number { + let i = 0; + while (i < limit) { + if (!step(i)) break; + i++; + } + return i; +} + +function twice(values: Int16Array, f: (v: number) => void): void { + let i = 0; + do { + for (let j = 0; j < values.length; j++) f(values[j]!); + i++; + } while (i < 2); +} + +const ints = new Int32Array([5, -3, 8, 0, 12, -7, 4]); +let total = 0; +let negatives = 0; +let sawZero = false; +ints.forEach((v) => { + total = (total + v) | 0; + if (v < 0) negatives++; + if (v === 0) sawZero = true; +}); +console.log("forEach", total, negatives, sawZero); + +let max = -Infinity; +let maxAt = -1; +const floats = new Float64Array([0.5, 9.25, -2, 9.25, 3]); +floats.forEach((v, i) => { + if (v > max) { + max = v; + maxAt = i; + } +}); +console.log("max", max, maxAt); + +let calls = 0; +const doubled = ints.map((v) => { + calls++; + return v * 2 + calls; +}); +console.log("map", doubled.join(","), calls); + +let seen = 0; +const big = [3, 14, 15, 92, 65].filter((v) => { + seen += v; + return v > 10; +}); +console.log("filter", big.join(","), seen); + +let sum = 0; +let count = 0; +const summing = (v: number): void => { + sum += v; + count++; +}; +each([1, 2, 3], summing); +each([10, 20], summing); +console.log("each shared", sum, count); + +let a = 0; +let b = 0; +each([1, 2], (v) => { + a += v; +}); +each([1, 2], (v, i) => { + b += v * (i + 1); +}); +console.log("each two", a, b); + +const counters: (() => number)[] = []; +for (let k = 0; k < 3; k++) { + let local = k * 100; + each([1, 1, 1], () => { + local++; + }); + counters.push(() => local); +} +console.log("closures", counters.map((c) => c()).join(",")); + +let steps = 0; +console.log("until", until(10, (i) => (steps += i, steps < 12)), steps); + +let shorts = 0; +twice(new Int16Array([7, -8, 9]), (v) => { + shorts = shorts * 2 + v; +}); +console.log("twice", shorts); + +let partial = 0; +try { + floats.forEach((v) => { + if (v < 0) throw new RangeError(`at ${partial}`); + partial += v; + }); +} catch (error) { + console.log("throw", (error as Error).message, partial); +} + +let outer = 0; +let inner = 0; +[1, 2, 3].forEach((x) => { + outer += x; + ints.forEach((y) => { + if (y > x) inner++; + }); +}); +console.log("nested", outer, inner); + +let flag = false; +let flips = 0; +new Uint8Array([1, 0, 1, 1]).forEach((v) => { + if (v === 1) flag = !flag; + if (flag) flips++; +}); +console.log("bool", flag, flips); + +// Callbacks declared in the loop body: each iteration binds its own. +let strides = 0; +for (let i = 0; i < 10; i++) { + const bump = (): void => { + i += 2; + strides++; + }; + bump(); +} +let countdown = 0; +let k = 3; +while (k > 0) { + const tick = (n: number): number => (n > 0 ? tick(n - 1) + 1 : 0); + countdown += tick(k); + k--; +} +console.log("declared in loop", strides, countdown); diff --git a/tests/harness/corpus-costs/7300-callback-captured-scalars.ts.json b/tests/harness/corpus-costs/7300-callback-captured-scalars.ts.json new file mode 100644 index 000000000..18495cd4e --- /dev/null +++ b/tests/harness/corpus-costs/7300-callback-captured-scalars.ts.json @@ -0,0 +1 @@ +{ "plain": 6000, "san": 7900 }