diff --git a/packages/compiler/src/backend/llvm/emit-statements.ts b/packages/compiler/src/backend/llvm/emit-statements.ts index 9b758e834..3d3aa6781 100644 --- a/packages/compiler/src/backend/llvm/emit-statements.ts +++ b/packages/compiler/src/backend/llvm/emit-statements.ts @@ -2179,6 +2179,7 @@ export function emitLoopCopies( ); }, (id) => emitter.speculativeLocal(id) || emitter.speculativeInvariantGlobal(id), + emitter.speculationContext(), ) : null; if (!plan) { @@ -2237,6 +2238,20 @@ export function emitLoopCopies( B.line(`${both} = and i1 ${ok}, ${trying}`); ok = both; } + // Hoisted leaves live in locals of their own, which the resume copy of + // their write reads (see SpeculativeValue). + for (const { leaves } of plan.values.values()) + for (const leaf of leaves) { + emitter.currentLocals.set(leaf.localId, { + id: leaf.localId, + name: "speculated leaf", + type: leaf.expr.type, + mutable: true, + }); + B.entryAllocas.push( + `${emitter.binding(leaf.localId).slot} = alloca double ; speculated leaf`, + ); + } const fast = B.newLabel("spec.fast"), slow = B.newLabel("spec.slow"); B.condBr(ok, fast, slow); @@ -2262,6 +2277,7 @@ export function emitLoopCopies( throw new InternalCompilerError("llvm emitter bug: speculative loop copy without a header"); B.startBlock(B.newLabel("spec.tail")); emitter.speculationResume = resume; + emitter.speculationResumeValues = plan.values; emitter.speculationHeader = () => { B.br(header!); B.startBlock(B.newLabel("spec.unused")); @@ -2270,6 +2286,7 @@ export function emitLoopCopies( emitCopy(exit); } finally { emitter.speculationResume = null; + emitter.speculationResumeValues = null; emitter.speculationHeader = null; } B.startBlock(fast); @@ -2318,6 +2335,7 @@ export function emitLoopCopies( deferred, bails, ran, + values: plan.values, }; try { const leave = B.newLabel("spec.leave"); @@ -2575,9 +2593,22 @@ export function emitSpeculativeWrite( // at once: a bailing path then owns nothing, and reads have no effect // whose reordering before the arithmetic could be observed. const reads = new Map(); - if (target.value !== null) { + const rewritten = emitter.speculation!.values?.get(s); + const value = rewritten?.value ?? target.value; + if (rewritten !== undefined) { + // Hoisted leaves run first, in the write's order, and keep their + // doubles for the resume copy. emitter.frames.push([]); - for (const read of speculativeReads(target.value)) reads.set(read, emitter.emitExpr(read)); + for (const leaf of rewritten.leaves) { + const evaluated = emitter.emitExpr(leaf.expr); + B.line(`store double ${evaluated.name}, ptr ${emitter.binding(leaf.localId).slot}`); + reads.set(leaf.ref, evaluated); + } + emitter.releaseTemps(emitter.frames.pop()!); + } + if (value !== null) { + emitter.frames.push([]); + for (const read of speculativeReads(value)) reads.set(read, emitter.emitExpr(read)); emitter.releaseTemps(emitter.frames.pop()!); } emitter.frames.push([]); @@ -2592,7 +2623,7 @@ export function emitSpeculativeWrite( B.line(`${biased} = add i64 ${result}, ${Number.MAX_SAFE_INTEGER}`); B.line(`${outside} = icmp ugt i64 ${biased}, ${2 * Number.MAX_SAFE_INTEGER}`); bail(outside); - } else result = emitSpeculativeInteger(emitter, target.value, bail, reads, resume === undefined); + } else result = emitSpeculativeInteger(emitter, value!, bail, reads, resume === undefined); // Speculable values read only numbers past this point. if (emitter.frames.pop()!.length > 0) throw new InternalCompilerError("llvm emitter bug: speculative write owns a temporary"); diff --git a/packages/compiler/src/backend/llvm/emitter.ts b/packages/compiler/src/backend/llvm/emitter.ts index 586517356..f7baf95f9 100644 --- a/packages/compiler/src/backend/llvm/emitter.ts +++ b/packages/compiler/src/backend/llvm/emitter.ts @@ -357,6 +357,13 @@ import { } from "./shapes.js"; import type { ExprOf, LibCallExpr, LlStreamTypedRefAdapter, LlValue } from "./expr-context.js"; import { INLINE_LITERAL_BYTES } from "./string-literal-equality.js"; +import { + inlinableNumberFunction, + withSpeculativeValue, + type InlinableFunction, + type SpeculationContext, + type SpeculativeValue, +} from "./speculative-integers.js"; import { emitFfiCallbackDefs, emitLibDefs, @@ -1524,11 +1531,47 @@ export class LlEmitter { * integer loop has neither. */ bails?: string; ran?: string; + /** Writes that speculate a rewritten value (rewriteSpeculable). */ + values?: ReadonlyMap; } | null = null; /** Speculative loop sites numbered so far in the module. */ speculativeLoopSites = 0; /** The ordinary copy being emitted: resume labels still to place. */ speculationResume: Map | null = null; + /** The rewritten values of the writes `speculationResume` resumes at. */ + speculationResumeValues: ReadonlyMap | null = null; + /** Hoisted speculative leaves numbered so far in the module. */ + private speculativeLeaves = 0; + private readonly inlinableFunctions = new Map(); + + /** What speculative loop planning may assume here (SpeculationContext). */ + speculationContext(): SpeculationContext { + const inlinable = (callee: string): InlinableFunction | undefined => { + if (!this.inlinableFunctions.has(callee)) { + // Unset while its body is checked: a recursive helper is not inlinable. + this.inlinableFunctions.set(callee, undefined); + const fn = this.fnByName.get(callee); + this.inlinableFunctions.set( + callee, + fn ? inlinableNumberFunction(fn, inlinable) : undefined, + ); + } + return this.inlinableFunctions.get(callee); + }; + return { + inlinable, + unaliased: (id) => { + const local = this.currentLocals.get(id); + return ( + local !== undefined && + !local.boxed && + !this.captureIds.has(id) && + !this.lazyCaptureBoxes.has(id) + ); + }, + leafId: () => `%spec.leaf.${this.speculativeLeaves++}`, + }; + } /** Run once at the top of the next loop copy's header * (`enterLoopHeader`): records the ordinary copy's header, or sends the * third copy's header there. */ @@ -5799,6 +5842,22 @@ export class LlEmitter { if (resume && !resume.placed && !this.B.isTerminated()) { // The fast copy's failed checks re-execute this statement here. resume.placed = true; + // A write with hoisted leaves resumed from its own failed check + // re-executes its arithmetic over the leaves' stored doubles: each + // leaf ran once, in the fast copy. Reached from an earlier statement + // of this copy (a failed check before it), its leaves never ran, so + // it runs as written. + const rewritten = this.speculationResumeValues?.get(s); + if (rewritten !== undefined && rewritten.leaves.length > 0) { + const join = this.B.newLabel("spec.resume.join"); + this.emitStmt(s); + if (!this.B.isTerminated()) this.B.br(join); + this.B.startBlock(resume.label); + this.emitStmt(withSpeculativeValue(s, rewritten.value)); + if (!this.B.isTerminated()) this.B.br(join); + this.B.startBlock(join); + return; + } this.B.br(resume.label); this.B.startBlock(resume.label); } diff --git a/packages/compiler/src/backend/llvm/speculative-integers.test.ts b/packages/compiler/src/backend/llvm/speculative-integers.test.ts index 3acfec0fb..1b82768d5 100644 --- a/packages/compiler/src/backend/llvm/speculative-integers.test.ts +++ b/packages/compiler/src/backend/llvm/speculative-integers.test.ts @@ -1,9 +1,11 @@ import { expect, test } from "vitest"; import { F64, type IrExpr, type IrStmt } from "../../ir/ir.js"; import { + inlinableNumberFunction, isSpeculableValue, planProvenIntegerLoop, planSpeculativeLoop, + rewriteSpeculable, speculativeReads, } from "./speculative-integers.js"; @@ -266,3 +268,136 @@ test("bounded cursors with proven writes are shadowed in a proven integer loop", planProvenIntegerLoop(forLoop([...body, jump]) as never, cursors, all, () => false, all), ).toBeNull(); }); + +const libCall = (fn: string, args: IrExpr[]): IrExpr => + ({ kind: "libCall", fn, args, type: F64, loc }) as never; +const call = (callee: string, args: IrExpr[]): IrExpr => ({ + kind: "call", + callee, + args, + type: F64, + loc, +}); +const pow2 = (n: IrExpr): IrExpr => libCall("math.pow", [num(2), n]); +/** `function trim(v, n) { return v - Math.floor(v / Math.pow(2, n)) * Math.pow(2, n); }` */ +const trim = { + name: "trim", + params: [ + { localId: "v", name: "v", type: F64 }, + { localId: "n", name: "n", type: F64 }, + ], + returnType: F64, + locals: [], + body: [ + { + kind: "return", + value: bin( + "-", + ref("v"), + bin("*", libCall("math.floor", [bin("/", ref("v"), pow2(ref("n")))]), pow2(ref("n"))), + ), + loc, + }, + ], + loc, +} as never; + +test("calls of small pure number functions expand into a speculated write", () => { + const inlinable = inlinableNumberFunction(trim)!; + expect(inlinable.params).toEqual(["v", "n"]); + const context = { inlinable: (callee: string) => (callee === "trim" ? inlinable : undefined) }; + // key = trim(key, m) * 4 + c + const value = bin("+", bin("*", call("trim", [ref("key"), ref("m")]), num(4)), ref("c")); + expect(isSpeculableValue(value)).toBe(false); + const rewritten = rewriteSpeculable(value, context)!; + expect(rewritten.leaves).toEqual([]); + expect(JSON.stringify(rewritten.value)).not.toContain('"call"'); + const write = assign("key", value); + const plan = planSpeculativeLoop( + loop([write, increment("i")]) as never, + all, + () => false, + () => false, + all, + context, + )!; + expect(plan.writes.has(write)).toBe(true); + expect(plan.values.get(write)).toEqual(rewritten); + // `m`, an input of the expanded body, is checked on entry with the others. + expect([...plan.locals.keys()].sort()).toEqual(["c", "i", "key", "m", "n"]); +}); + +test("Math.floor needs a power-of-two divisor, and expansion pure arguments", () => { + const context = { inlinable: () => inlinableNumberFunction(trim) }; + expect(rewriteSpeculable(libCall("math.floor", [bin("/", ref("a"), num(8))]), {})).not.toBeNull(); + expect(rewriteSpeculable(libCall("math.floor", [bin("/", ref("a"), num(6))]), {})).toBeNull(); + expect(rewriteSpeculable(libCall("math.floor", [bin("/", ref("a"), ref("b"))]), {})).toBeNull(); + // An argument that is itself a call is evaluated once: no expansion. + expect(rewriteSpeculable(call("trim", [call("other", []), ref("m")]), context)).toBeNull(); + const effectful = { + ...(trim as object), + body: [{ kind: "exprStmt", expr: call("log", []), loc }], + }; + expect(inlinableNumberFunction(effectful as never)).toBeUndefined(); + // A divisor next to a power of two is not one, though Math.log2 rounds it to 52. + const near = (d: number): IrExpr => libCall("math.floor", [bin("/", ref("a"), num(d))]); + expect(rewriteSpeculable(near(2 ** 52), {})).not.toBeNull(); + expect(rewriteSpeculable(near(2 ** 52 + 1), {})).toBeNull(); + expect(rewriteSpeculable(near(2 ** 53), {})).toBeNull(); + expect(rewriteSpeculable(near(0.5), {})).toBeNull(); +}); + +test("a helper's own calls must be of helpers that expand too", () => { + // `function outer(v, n) { return trim(v, n) + other(v); }` + const outer = { + ...(trim as object), + name: "outer", + body: [ + { + kind: "return", + value: bin("+", call("trim", [ref("v"), ref("n")]), call("other", [ref("v")])), + loc, + }, + ], + }; + const helpers = (callee: string) => + callee === "trim" ? inlinableNumberFunction(trim) : undefined; + expect(inlinableNumberFunction(outer as never)).toBeUndefined(); + expect(inlinableNumberFunction(outer as never, helpers)).toBeUndefined(); + const pure = { + ...(outer as object), + body: [{ kind: "return", value: bin("*", call("trim", [ref("v"), ref("n")]), num(2)), loc }], + }; + expect(inlinableNumberFunction(pure as never, helpers)).toBeDefined(); +}); + +test("other leaves are hoisted only when no call can change the write's locals", () => { + let next = 0; + const leafId = (): string => `%leaf.${next++}`; + const value = bin( + "%", + bin("+", bin("*", ref("acc"), num(7)), call("noisy", [ref("i")])), + num(1000003), + ); + const hoisted = rewriteSpeculable(value, { unaliased: () => true, leafId })!; + expect(hoisted.leaves.map((leaf) => leaf.expr)).toEqual([call("noisy", [ref("i")])]); + expect(rewriteSpeculable(value, { unaliased: (id) => id !== "acc", leafId })).toBeNull(); + // A call whose argument writes a local stays where it is. + const bump: IrExpr = { kind: "incDec", localId: "i", op: "+", prefix: false, type: F64, loc }; + expect( + rewriteSpeculable(bin("+", ref("i"), call("noisy", [bump])), { unaliased: () => true, leafId }), + ).toBeNull(); + // Without leaf ids nothing is hoisted. + expect(rewriteSpeculable(value, { unaliased: () => true })).toBeNull(); + const write = assign("acc", value); + const plan = planSpeculativeLoop( + loop([write, increment("i")]) as never, + (id) => !id.startsWith("%leaf"), + () => false, + () => false, + (id) => !id.startsWith("%leaf"), + { unaliased: () => true, leafId }, + )!; + expect(plan.values.get(write)!.leaves).toHaveLength(1); + expect([...plan.locals.keys()].some((id) => id.startsWith("%leaf"))).toBe(false); +}); diff --git a/packages/compiler/src/backend/llvm/speculative-integers.ts b/packages/compiler/src/backend/llvm/speculative-integers.ts index 00fabb788..07e5e3721 100644 --- a/packages/compiler/src/backend/llvm/speculative-integers.ts +++ b/packages/compiler/src/backend/llvm/speculative-integers.ts @@ -1,5 +1,6 @@ -import type { IrExpr, IrStmt } from "../../ir/ir.js"; -import { everyExprChild, everyStmtChild } from "../../ir/traverse.js"; +import type { IrExpr, IrFunction, IrStmt } from "../../ir/ir.js"; +import { everyExpr, everyExprChild, everyStmtChild } from "../../ir/traverse.js"; +import { InternalCompilerError } from "../../errors.js"; import { f64Lit } from "./common.js"; import type { LlValue, LlvmEmitterContext } from "./expr-context.js"; import { exactInteger, widenInteger } from "./integer-values.js"; @@ -30,6 +31,295 @@ export interface SpeculativeLoopPlan { * than through its condition, `break` or a failed check: no calls, no * throwing operations, no labeled jumps and no `return`. */ observerFree: boolean; + /** The writes that speculate a rewritten value (rewriteSpeculable). */ + values: Map; +} + +/** A write's value as speculation evaluates it: calls of small pure number + * functions expanded in place, and other number leaves (`leaves`) replaced + * by references to fresh locals. The fast copy evaluates each leaf once, + * stores it to its local and checks it as an integer; a failed check + * resumes at a copy of the write that reads the locals, so a leaf is never + * evaluated twice. */ +export interface SpeculativeValue { + value: IrExpr; + leaves: { localId: string; expr: IrExpr; ref: IrExpr }[]; +} + +/** What the planner may assume about the code around a loop. */ +export interface SpeculationContext { + /** The parameters and the returned expression of `callee` when it is a + * small pure number function (inlinableNumberFunction). */ + inlinable?: (callee: string) => InlinableFunction | undefined; + /** A local that no call can write: not captured, boxed or global. */ + unaliased?: (localId: string) => boolean; + /** A fresh local id for a hoisted leaf. */ + leafId?: () => string; +} + +export interface InlinableFunction { + params: readonly string[]; + body: IrExpr; +} + +/** `s` (a speculated write) writing `value` instead. */ +export function withSpeculativeValue(s: IrStmt, value: IrExpr): IrStmt { + switch (s.kind) { + case "assign": + return { ...s, value }; + case "varDecl": + return { ...s, init: value }; + case "exprStmt": + if (s.expr.kind === "assignExpr") return { ...s, expr: { ...s.expr, value } }; + break; + default: + break; + } + throw new InternalCompilerError(`llvm emitter bug: speculated ${s.kind} with a rewritten value`); +} + +/** Calls nest at most this deep in one expansion. */ +const MAX_INLINE_DEPTH = 4; +/** An expansion larger than this stays a call. */ +const MAX_INLINE_NODES = 64; + +/** A function that only returns arithmetic over its number parameters: + * `function trim(v, n) { return v - Math.floor(v / Math.pow(2, n)) * Math.pow(2, n); }`. + * Its call evaluates nothing but its arguments, so a speculated write may + * expand it in place. A call in its body must be of another such function + * (`resolve`): any other call could change what the expansion reads. */ +export function inlinableNumberFunction( + fn: IrFunction, + resolve?: (callee: string) => InlinableFunction | undefined, +): InlinableFunction | undefined { + if ( + fn.async || + fn.generator || + (fn.captures?.length ?? 0) > 0 || + (fn.classCaptures?.length ?? 0) > 0 || + fn.returnType.kind !== "f64" || + fn.params.some((p) => p.type.kind !== "f64" && p.type.kind !== "union") || + fn.body.length !== 1 || + fn.body[0]!.kind !== "return" || + fn.body[0]!.value === undefined + ) + return undefined; + const params = fn.params.map((p) => p.localId); + const known = new Set(params); + let nodes = 0; + // A parameter of a union type (`number | undefined`) is read only + // through its tag test and narrowing, which a wrapped argument folds. + const union = (e: IrExpr): boolean => e.kind === "varRef" && known.has(e.localId); + const pure = (e: IrExpr): boolean => { + if (++nodes > MAX_INLINE_NODES) return false; + if (e.kind === "boolLit") return true; + if (e.kind === "unionIsTag") return union(e.value); + if (e.kind === "unionNarrow") return e.type.kind === "f64" && union(e.value); + if (e.type.kind !== "f64") return false; + switch (e.kind) { + case "numLit": + return true; + case "varRef": + return known.has(e.localId); + case "ternary": + return pure(e.cond) && pure(e.then) && pure(e.else_); + case "bin": + return (INTEGER_OPS.has(e.op) || e.op === "/") && pure(e.left) && pure(e.right); + case "unary": + return e.op === "-" && pure(e.operand); + case "libCall": + return (e.fn === "math.floor" || e.fn === "math.pow") && e.args.every(pure); + case "call": + return resolve?.(e.callee) !== undefined && e.args.every(pure); + default: + return false; + } + }; + const body = fn.body[0]!.value!; + return pure(body) ? { params, body } : undefined; +} + +/** `e` with each parameter reference replaced by its argument. A tag test + * or narrowing of a wrapped argument folds, and so does a ternary over a + * folded test. */ +function substitute(e: IrExpr, args: ReadonlyMap): IrExpr { + switch (e.kind) { + case "varRef": + return args.get(e.localId) ?? e; + case "unionIsTag": { + const value = substitute(e.value, args); + if (value.kind !== "unionWrap" || value.unionId !== e.unionId) return { ...e, value }; + return { + kind: "boolLit", + value: (value.tag === e.tag) !== (e.negated === true), + type: e.type, + loc: e.loc, + }; + } + case "unionNarrow": { + const value = substitute(e.value, args); + return value.kind === "unionWrap" && value.unionId === e.unionId && value.tag === e.tag + ? value.value + : { ...e, value }; + } + case "ternary": { + const cond = substitute(e.cond, args); + if (cond.kind === "boolLit") return substitute(cond.value ? e.then : e.else_, args); + return { ...e, cond, then: substitute(e.then, args), else_: substitute(e.else_, args) }; + } + case "bin": + return { ...e, left: substitute(e.left, args), right: substitute(e.right, args) }; + case "unary": + return { ...e, operand: substitute(e.operand, args) }; + case "libCall": + case "call": + return { ...e, args: e.args.map((a) => substitute(a, args)) }; + default: + return e; + } +} + +/** k when `value` is exactly 2^k with k in [0, 52], else null. + * Math.log2 rounds: 2^52 + 1 gives 52, so the power is compared back. */ +function exactPowerOfTwo(value: number): number | null { + if (!(value >= 1)) return null; + const k = Math.round(Math.log2(value)); + return k >= 0 && k <= 52 && 2 ** k === value ? k : null; +} + +/** Whether `e` is a power of two speculation can shift by: + * `Math.pow(2, n)` or a literal 2^k with k in [0, 52]. */ +function powerOfTwo(e: IrExpr): boolean { + if (e.kind === "numLit") return exactPowerOfTwo(e.value) !== null; + return ( + e.kind === "libCall" && + e.fn === "math.pow" && + e.args.length === 2 && + e.args[0]!.kind === "numLit" && + e.args[0]!.value === 2 + ); +} + +/** Typed-array element kinds whose elements are integers. */ +const INTEGER_ELEMENTS = new Set(["u8", "u8c", "i8", "u16", "i16", "u32", "i32"]); + +/** A typed-array element read of an integer kind with no effect: an + * out-of-range read gives NaN, which fails its check. Float elements would + * only bail. */ +function isSpeculativeByteRead(e: IrExpr): boolean { + return ( + e.kind === "bytesIntrinsic" && + e.method === "get" && + e.type.kind === "f64" && + e.receiver.type.kind === "bytes" && + INTEGER_ELEMENTS.has(e.receiver.type.elem) && + e.args.length === 1 && + isEffectFree(e.receiver) && + isEffectFree(e.args[0]!) + ); +} + +/** Whether evaluating `e` assigns no binding: no `++`/`--`, no assignment + * expression, and no statement inside it. */ +function writesNoLocal(e: IrExpr): boolean { + return everyExpr(e, { + expr: (node) => node.kind !== "incDec" && node.kind !== "assignExpr", + stmt: () => false, + }); +} + +/** Rewrite a write's value that isSpeculableValue refuses into one it can + * speculate (SpeculativeValue), or null. Expansion and the new operations + * (`Math.pow(2, n)`, `Math.floor(a / 2^n)`, byte reads) evaluate nothing + * but numbers, so the fast copy may evaluate them in any order. Hoisted + * leaves run in the write's evaluation order before the arithmetic; that + * reorders nothing only when the write reads no heap and every local it + * reads is one no call can write. */ +export function rewriteSpeculable( + value: IrExpr, + context: SpeculationContext, +): SpeculativeValue | null { + const leaves: SpeculativeValue["leaves"] = []; + let reads = false; + const locals: string[] = []; + const rewrite = (e: IrExpr, depth: number): IrExpr | null => { + if (e.type.kind !== "f64") return null; + switch (e.kind) { + case "numLit": + return Number.isSafeInteger(e.value) && !Object.is(e.value, -0) ? e : null; + case "varRef": + locals.push(e.localId); + return e; + case "bin": { + if (!ARITHMETIC_OPS.has(e.op)) break; + const left = rewrite(e.left, depth); + const right = left && rewrite(e.right, depth); + return left && right ? { ...e, left, right } : null; + } + case "unary": { + if (e.op !== "-") break; + const operand = rewrite(e.operand, depth); + return operand ? { ...e, operand } : null; + } + case "libCall": { + if (powerOfTwo(e) && e.kind === "libCall") { + const exponent = rewrite(e.args[1]!, depth); + return exponent ? { ...e, args: [e.args[0]!, exponent] } : null; + } + const quotient = e.args[0]; + if ( + e.fn === "math.floor" && + e.args.length === 1 && + quotient?.kind === "bin" && + quotient.op === "/" && + powerOfTwo(quotient.right) + ) { + const left = rewrite(quotient.left, depth); + const right = left && rewrite(quotient.right, depth); + return left && right ? { ...e, args: [{ ...quotient, left, right }] } : null; + } + break; + } + case "call": { + const fn = depth < MAX_INLINE_DEPTH ? context.inlinable?.(e.callee) : undefined; + // Arguments are evaluated before the body: an expansion may read + // them in another order, or twice, only when they are pure. + const plain = (a: IrExpr): boolean => + isSpeculableValue(a, false) || + (a.kind === "unionWrap" && isSpeculableValue(a.value, false)); + if (fn !== undefined && fn.params.length === e.args.length && e.args.every(plain)) { + const args = new Map(fn.params.map((id, i) => [id, e.args[i]!])); + return rewrite(substitute(fn.body, args), depth + 1); + } + break; + } + default: + if (isSpeculativeRead(e) || isSpeculativeByteRead(e)) { + reads = true; + return e; + } + break; + } + // Only calls of program functions are hoisted: other leaves (a string's + // length, a Map read) are cheap where they stand, and would turn every + // loop that adds one into a second copy. + if (context.leafId === undefined || e.kind !== "call") return null; + // Hoisting runs the arguments before the rest of the write: one that + // assigns a local would move that write ahead of the write's earlier + // reads of it. + if (!e.args.every(writesNoLocal)) return null; + const ref: IrExpr = { kind: "varRef", localId: context.leafId(), type: e.type, loc: e.loc }; + leaves.push({ localId: (ref as { localId: string }).localId, expr: e, ref }); + return ref; + }; + const rewritten = rewrite(value, 0); + if (rewritten === null) return null; + if (leaves.length > 0) { + const leafIds = new Set(leaves.map((leaf) => leaf.localId)); + if (reads || !locals.every((id) => leafIds.has(id) || context.unaliased?.(id) === true)) + return null; + } + return { value: rewritten, leaves }; } type SpeculativeLoop = Extract; @@ -170,13 +460,16 @@ function isEffectFree(e: IrExpr): boolean { } } -/** The speculative reads of a speculable value in evaluation order. */ +/** The speculative reads of a speculable value in evaluation order, those + * of a rewritten value (rewriteSpeculable) included. */ export function speculativeReads(e: IrExpr, out: IrExpr[] = []): IrExpr[] { - if (e.kind === "bin" && ARITHMETIC_OPS.has(e.op)) { + if (e.kind === "bin" && (ARITHMETIC_OPS.has(e.op) || e.op === "/")) { speculativeReads(e.left, out); speculativeReads(e.right, out); } else if (e.kind === "unary" && e.op === "-") speculativeReads(e.operand, out); - else if (isSpeculativeRead(e)) out.push(e); + else if (e.kind === "libCall" && (e.fn === "math.floor" || e.fn === "math.pow")) + for (const arg of e.args) speculativeReads(arg, out); + else if (isSpeculativeRead(e) || isSpeculativeByteRead(e)) out.push(e); return out; } @@ -195,6 +488,8 @@ function hasProduct(e: IrExpr): boolean { if (e.kind === "bin") return e.op === "*" || e.op === "%" || hasProduct(e.left) || hasProduct(e.right); if (e.kind === "unary") return hasProduct(e.operand); + // Math.floor(a / 2^n) and Math.pow(2, n) replace a division and a call. + if (e.kind === "libCall") return e.fn === "math.floor" || e.fn === "math.pow"; return false; } @@ -202,6 +497,7 @@ function hasArithmetic(e: IrExpr): boolean { if (e.kind === "bin") return ARITHMETIC_OPS.has(e.op) || hasArithmetic(e.left) || hasArithmetic(e.right); if (e.kind === "unary") return e.op === "-" || hasArithmetic(e.operand); + if (e.kind === "libCall") return e.fn === "math.floor" || e.fn === "math.pow"; return false; } @@ -236,8 +532,11 @@ export function planSpeculativeLoop( integral: (localId: string) => boolean = () => false, proven: (e: IrExpr) => boolean = () => false, invariantEligible: (localId: string) => boolean = eligible, + context: SpeculationContext = {}, ): SpeculativeLoopPlan | null { const writes = new Map(); + const values = new Map(); + const leafIds = new Set(); const unspeculable = new Set(); const reads = new Set(); // Locals with at least one read the static analysis does not prove. @@ -263,11 +562,9 @@ export function planSpeculativeLoop( // make the body's work integer. let inCondition = false; - const expr = (e: IrExpr): boolean => { - if (rejected) return false; - nodes++; - if (!OBSERVER_FREE_EXPRESSIONS.has(e.kind) || (e.kind === "bin" && e.left.type.kind !== "f64")) - observerFree = false; + // What the integer writes read and consume. A rewritten value is noted + // here alone: its original, which `expr` walks, has its effects. + const note = (e: IrExpr): void => { if (e.kind === "varRef") { reads.add(e.localId); if (!proven(e)) unproven.add(e.localId); @@ -281,6 +578,20 @@ export function planSpeculativeLoop( } else if (e.kind === "arrayGet" || e.kind === "arrayHas" || e.kind === "arrayState") consume(e.index); else if (e.kind === "bytesIntrinsic" || e.kind === "strIntrinsic") e.args.forEach(consume); + }; + // An expanded call is emitted in the fast copy: it counts toward the + // duplicated code. + const noteAll = (e: IrExpr): boolean => { + nodes++; + note(e); + return everyExprChild(e, noteAll, () => true); + }; + const expr = (e: IrExpr): boolean => { + if (rejected) return false; + nodes++; + if (!OBSERVER_FREE_EXPRESSIONS.has(e.kind) || (e.kind === "bin" && e.left.type.kind !== "f64")) + observerFree = false; + note(e); return everyExprChild(e, expr, nested); }; // Statements inside expressions (sequence expressions) run mid-expression, @@ -307,7 +618,11 @@ export function planSpeculativeLoop( if (s.kind === "arraySet" || s.kind === "bytesSet") consume(s.index); const target = speculativeWriteTarget(s); if (target) { - if (target.value === null || isSpeculableValue(target.value)) { + const rewritten = + target.value === null || isSpeculableValue(target.value) + ? null + : rewriteSpeculable(target.value, context); + if (target.value === null || isSpeculableValue(target.value) || rewritten !== null) { const list = writes.get(target.localId) ?? []; list.push(s); writes.set(target.localId, list); @@ -315,6 +630,11 @@ export function planSpeculativeLoop( reads.add(target.localId); if (s.kind !== "exprStmt" || !proven(s.expr)) unproven.add(target.localId); } else expr(target.value); + if (rewritten !== null) { + values.set(s, rewritten); + for (const leaf of rewritten.leaves) leafIds.add(leaf.localId); + noteAll(rewritten.value); + } return !rejected; } unspeculable.add(target.localId); @@ -341,25 +661,28 @@ export function planSpeculativeLoop( // write runs: speculated itself, already an integer, or loop-invariant // (checked on entry). A value computed from doubles the loop writes // would only fail its check on every entry. + const valueOf = (s: IrStmt): IrExpr | null => + values.get(s)?.value ?? speculativeWriteTarget(s)!.value; const grounded = (e: IrExpr): boolean => { if (e.kind === "varRef") { const id = e.localId; return ( + leafIds.has(id) || proven(e) || locals.has(id) || integral(id) || (!writes.has(id) && !unspeculable.has(id) && invariantEligible(id)) ); } - if (isSpeculativeRead(e)) return true; + if (isSpeculativeRead(e) || isSpeculativeByteRead(e)) return true; return everyExprChild(e, grounded, () => false); }; for (let changed = true; changed;) { changed = false; for (const id of [...locals.keys()]) { for (const s of writes.get(id)!) { - const target = speculativeWriteTarget(s)!; - if (target.value !== null && !grounded(target.value)) { + const value = valueOf(s); + if (value !== null && !grounded(value)) { locals.delete(id); changed = true; break; @@ -376,6 +699,7 @@ export function planSpeculativeLoop( if (e.kind === "varRef") { const id = e.localId; if ( + !leafIds.has(id) && !proven(e) && !locals.has(id) && !writes.has(id) && @@ -388,11 +712,11 @@ export function planSpeculativeLoop( return everyExprChild(e, invariant, () => true); }; for (const s of speculated) { - const target = speculativeWriteTarget(s)!; - if (target.value === null) arithmetic = true; + const value = valueOf(s); + if (value === null) arithmetic = true; else { - if (hasArithmetic(target.value)) arithmetic = true; - invariant(target.value); + if (hasArithmetic(value)) arithmetic = true; + invariant(value); } } if (!arithmetic) return null; @@ -401,13 +725,12 @@ export function planSpeculativeLoop( [...locals.keys()].some((id) => multiplied.has(id)) || [...speculated].some((s) => { const target = speculativeWriteTarget(s)!; - return ( - target.value !== null && - (hasProduct(target.value) || isVariableStep(target.localId, target.value)) - ); + const value = valueOf(s); + return value !== null && (hasProduct(value) || isVariableStep(target.localId, value)); }); if (!heavy) return null; - return { locals, writes: speculated, observerFree }; + for (const s of values.keys()) if (!speculated.has(s)) values.delete(s); + return { locals, writes: speculated, observerFree, values }; } /** A proven integer loop. The bindings integer-ranges.ts bounds over a for @@ -682,6 +1005,19 @@ export function emitSpeculativeInteger( bail(negativeZero); return result; } + case "libCall": { + // A rewritten value's Math.pow(2, n) and Math.floor(a / 2^n) + // (rewriteSpeculable): with n an integer in [0, 52], 2^n is exact and + // so is a / 2^n for a safe integer a, whose floor the arithmetic shift + // computes, also for negative a. No result is -0. + if (e.fn === "math.pow") return shiftOne(host, exponentOf(host, e, bail, reads, integral)); + if (e.fn !== "math.floor" || e.args[0]?.kind !== "bin") break; + const a = emitSpeculativeInteger(host, e.args[0].left, bail, reads, integral); + const n = exponentOf(host, e.args[0].right, bail, reads, integral); + const result = B.tmp(); + B.line(`${result} = ashr i64 ${a}, ${n}`); + return result; + } case "unary": { if (e.op !== "-") break; const a = emitSpeculativeInteger(host, e.operand, bail, reads, integral); @@ -699,6 +1035,35 @@ export function emitSpeculativeInteger( return speculativeLeaf(host, reads.get(e) ?? host.emitExpr(e), e, bail); } +/** The exponent of a power of two (Math.pow(2, n) or a literal 2^k) as an + * i64 in [0, 52]; other exponents bail. */ +function exponentOf( + host: LlvmEmitterContext, + e: IrExpr, + bail: (condition: string) => void, + reads: ReadonlyMap, + integral: boolean, +): string { + if (e.kind === "numLit") { + const k = exactPowerOfTwo(e.value); + if (k === null) throw new InternalCompilerError("llvm emitter bug: speculative power of two"); + return String(k); + } + if (e.kind !== "libCall") + throw new InternalCompilerError("llvm emitter bug: speculative power of two"); + const n = emitSpeculativeInteger(host, e.args[1]!, bail, reads, integral); + const outside = host.B.tmp(); + host.B.line(`${outside} = icmp ugt i64 ${n}, 52`); + bail(outside); + return n; +} + +function shiftOne(host: LlvmEmitterContext, n: string): string { + const result = host.B.tmp(); + host.B.line(`${result} = shl i64 1, ${n}`); + return result; +} + /** The safe i64 of an evaluated number: its proven integer, or a checked * conversion. */ function speculativeLeaf( diff --git a/packages/compiler/test/ts7/baselines/order-parity/tests/corpus/7757-speculative-key-chains.ts.json b/packages/compiler/test/ts7/baselines/order-parity/tests/corpus/7757-speculative-key-chains.ts.json new file mode 100644 index 000000000..274151b18 --- /dev/null +++ b/packages/compiler/test/ts7/baselines/order-parity/tests/corpus/7757-speculative-key-chains.ts.json @@ -0,0 +1,6 @@ +{ + "order": [ + "/tests/corpus/7757-speculative-key-chains.ts" + ], + "diags": [] +} diff --git a/tests/corpus/7757-speculative-key-chains.ts b/tests/corpus/7757-speculative-key-chains.ts new file mode 100644 index 000000000..1ef08a942 --- /dev/null +++ b/tests/corpus/7757-speculative-key-chains.ts @@ -0,0 +1,124 @@ +// Speculative integer loops over writes that call small pure number +// functions, use Math.pow(2, n) and Math.floor(a / 2^n), read bytes, or +// call other functions: every fallback gives Node's numbers. +function trim(value: number, length: number): number { + return value - Math.floor(value / Math.pow(2, length)) * Math.pow(2, length); +} +function maybeTrim(value: number | undefined, length: number): number { + return (value ?? 0) - Math.floor((value ?? 0) / Math.pow(2, length)) * Math.pow(2, length); +} + +function chain(codes: Uint8Array, length: number): number { + const mask = length * 2 - 2; + let key = 0; + let sum = 0; + for (let i = 0; i < codes.length; ++i) { + sum = (sum + key) % 1000000007; + key = trim(key, mask) * 4 + codes[i]!; + } + return sum * 31 + key; +} + +function optionalChain(codes: Uint8Array, length: number): number { + let key = 0; + let sum = 0; + for (let i = 0; i < codes.length; ++i) { + sum = (sum * 3 + key) % 1000003; + key = maybeTrim(key, length * 2) * 4 + codes[i]!; + } + return sum + key; +} + +// Exponents past 52, negative and fractional values, and NaN reads leave +// the integer copy and finish in the double one. +function shifts(start: number, exponent: number, steps: number): number[] { + const out: number[] = []; + let v = start; + for (let i = 0; i < steps; i++) { + v = Math.floor(v / Math.pow(2, exponent)) * 3 + i * 7; + out.push(v); + } + return out; +} + +let calls = 0; +function noisy(x: number): number { + calls++; + return x % 5 === 0 ? x / 2 : x * 2; +} +// A call that the write cannot expand runs once per iteration, also on +// the iteration whose check fails. +function hoisted(n: number): number { + let acc = 1; + for (let i = 0; i < n; i++) acc = (acc * 7 + noisy(i)) % 1000003; + return acc; +} + +// A check that fails before a write with a hoisted call: that write still +// makes its call on the iteration that bails. +function tally(x: number): number { + calls++; + return x * 2 + 1; +} +function chained(n: number): number { + let a = 1; + let b = 1; + for (let i = 0; i < n; i++) { + a = a * 3 + i; + b = (b * 7 + tally(i) + (a % 11)) % 1000003; + } + return b; +} + +// A divisor next to a power of two is not one: 2^52 + 1 divides as itself. +function nearPowers(start: number, steps: number): number[] { + const out: number[] = []; + let v = start; + for (let i = 0; i < steps; i++) { + v = Math.floor(v / 4503599627370497) * 4 + Math.floor(v / 4503599627370496) + i; + out.push(v); + } + return out; +} + +// A helper that calls a function with effects is not expanded: its +// argument is read before the call runs. +let level = 0; +function raise(): number { + level += 100; + return 1; +} +function lifted(x: number): number { + return raise() + x; +} +function effectfulHelper(n: number): number { + let acc = 0; + for (let i = 0; i < n; i++) acc = (acc * 3 + lifted(level)) % 1000003; + return acc; +} + +// A call whose argument assigns a local the write also reads stays where it +// is: `i` is read before `i++` runs. +function stepped(n: number): number { + let acc = 1; + let i = 0; + while (i < n) acc = (acc * 3 + i + tally(i++)) % 1000003; + return acc; +} + +const codes = new Uint8Array(5000); +let seed = 7; +for (let i = 0; i < codes.length; i++) { + seed = (seed * 1103515245 + 12345) % 2147483648; + codes[i] = seed % 4; +} +for (const length of [1, 2, 3, 6, 12, 18, 27]) console.log(length, chain(codes, length), optionalChain(codes, length)); +console.log(shifts(1000, 3, 6), shifts(-1000, 2, 6), shifts(5, 53, 3), shifts(7.5, 1, 3), shifts(2 ** 52, 1, 3)); +console.log(hoisted(20), calls); +const short = codes.subarray(0, 3); +console.log(chain(short, 2), chain(new Uint8Array(0), 2)); +calls = 0; +console.log(chained(34), chained(60), calls); +console.log(nearPowers(4503599627370496, 3), nearPowers(9007199254740991, 2)); +console.log(effectfulHelper(10), level); +console.log(stepped(40)); diff --git a/tests/harness/corpus-costs/7757-speculative-key-chains.ts.json b/tests/harness/corpus-costs/7757-speculative-key-chains.ts.json new file mode 100644 index 000000000..be69de5f7 --- /dev/null +++ b/tests/harness/corpus-costs/7757-speculative-key-chains.ts.json @@ -0,0 +1 @@ +{ "plain": 2200, "san": 2900 }