diff --git a/packages/compiler/src/backend/llvm/constant-choices.test.ts b/packages/compiler/src/backend/llvm/constant-choices.test.ts new file mode 100644 index 000000000..f19b5b19d --- /dev/null +++ b/packages/compiler/src/backend/llvm/constant-choices.test.ts @@ -0,0 +1,109 @@ +import { expect, test } from "vitest"; +import { BOOL, F64, type IrExpr, type IrStmt } from "../../ir/ir.js"; +import { matchConstantChoice } from "./constant-choices.js"; + +const loc = { file: "choices.ts", start: 0, end: 0 }; +const bytes = { kind: "bytes", elem: "u8" } as const; +const num = (value: number): IrExpr => ({ kind: "numLit", value, type: F64, loc }); +const ref = (localId: string, type: IrExpr["type"] = F64): IrExpr => ({ + kind: "varRef", + localId, + type, + loc, +}); +const is = (left: IrExpr, right: IrExpr): IrExpr => ({ + kind: "bin", + op: "===", + left, + right, + type: BOOL, + loc, +}); +const store = (value: IrExpr, index: IrExpr = ref("i"), arr = "a"): IrStmt => ({ + kind: "bytesSet", + arr: ref(arr, bytes), + index, + value, + loc, +}); +const branch = (cond: IrExpr, then: IrStmt[], else_: IrStmt[] | null): IrStmt => ({ + kind: "if", + cond, + then, + else_, + loc, +}); +/** `if (x === c0) a[i] = r0; else if ... else a[i] = fallback`. */ +function chain(cases: [number, number][], fallback: IrStmt[] | null, subject = "x"): IrStmt { + let rest = fallback; + for (let k = cases.length - 1; k >= 0; k--) { + const [value, result] = cases[k]!; + rest = [branch(is(ref(subject), num(value)), [store(num(result))], rest)]; + } + return rest![0]!; +} + +test("a chain of literal stores to one element is a constant choice", () => { + const choice = matchConstantChoice( + chain( + [ + [97, 0], + [116, 1], + [99, 2], + ], + [store(num(3))], + ), + ); + expect(choice?.cases).toEqual([ + { value: 97, result: 0 }, + { value: 116, result: 1 }, + { value: 99, result: 2 }, + ]); + expect(choice?.fallback).toBe(3); + expect(choice?.stores).toHaveLength(4); +}); + +test("a literal on the left, -0 and a repeated value keep strict equality's first arm", () => { + const first: IrStmt = branch( + is(num(-0), ref("x")), + [store(num(5))], + [ + chain( + [ + [1, 6], + [0, 7], + [2, 8], + ], + [store(num(9))], + ), + ], + ); + expect(matchConstantChoice(first)?.cases).toEqual([ + { value: 0, result: 5 }, + { value: 1, result: 6 }, + { value: 2, result: 8 }, + ]); +}); + +test("chains that are not one choice of literals keep their arms", () => { + const three: [number, number][] = [ + [1, 10], + [2, 20], + [3, 30], + ]; + // Two values, no final else, a fraction, or cases spanning more than 256 values. + expect(matchConstantChoice(chain(three.slice(0, 2), [store(num(0))]))).toBeNull(); + expect(matchConstantChoice(chain(three, null))).toBeNull(); + expect(matchConstantChoice(chain([...three, [1.5, 40]], [store(num(0))]))).toBeNull(); + expect(matchConstantChoice(chain([...three, [256, 40]], [store(num(0))]))).not.toBeNull(); + expect(matchConstantChoice(chain([...three, [257, 40]], [store(num(0))]))).toBeNull(); + // A final else that is not a literal, or that stores elsewhere. + expect(matchConstantChoice(chain(three, [store(ref("y"))]))).toBeNull(); + expect(matchConstantChoice(chain(three, [store(num(0), ref("j"))]))).toBeNull(); + expect(matchConstantChoice(chain(three, [store(num(0), ref("i"), "b")]))).toBeNull(); + // An arm with a second statement. + expect(matchConstantChoice(chain(three, [store(num(0)), store(num(1))]))).toBeNull(); + // A second subject. + const mixed = branch(is(ref("y"), num(0)), [store(num(1))], [chain(three, [store(num(0))])]); + expect(matchConstantChoice(mixed)).toBeNull(); +}); diff --git a/packages/compiler/src/backend/llvm/constant-choices.ts b/packages/compiler/src/backend/llvm/constant-choices.ts new file mode 100644 index 000000000..208fac36a --- /dev/null +++ b/packages/compiler/src/backend/llvm/constant-choices.ts @@ -0,0 +1,236 @@ +import { createHash } from "node:crypto"; +import { InternalCompilerError } from "../../errors.js"; +import type { IrExpr, IrStmt } from "../../ir/ir.js"; +import { exactInteger, widenInteger } from "./integer-values.js"; +import type { LlValue, LlvmEmitterContext } from "./expr-context.js"; +import { f64Lit } from "./common.js"; + +type BytesSet = Extract; + +/** An if/else-if chain that stores one of several numeric literals to the + * same typed-array element, choosing it by comparing one number local with + * integer literals: + * + * if (x === 97) a[i] = 0; else if (x === 116) a[i] = 1; ... else a[i] = 3; + * + * The chain is emitted as one store of a value read from a constant table, + * with no branch on the comparisons. Branches on unpredictable input (the + * bases of a DNA sequence, the bytes of a text) mispredict about half the + * time; LLVM keeps them, because sparse cases (97, 99, 116) are too far + * apart for its own switch tables. */ +export interface ConstantChoice { + subject: IrExpr; + /** The first arm for each distinct value, in chain order. */ + cases: readonly { value: number; result: number }[]; + fallback: number; + /** Every arm's store, the final else's last. */ + stores: readonly BytesSet[]; +} + +/** At least this many distinct values: shorter chains predict as well as a + * table load. */ +const MIN_CASES = 3; +/** The table covers the subject's values from the smallest case (or the + * subject's proven minimum) up to this many entries. */ +const MAX_TABLE = 256; + +function integerLiteral(expr: IrExpr): number | null { + if (expr.kind !== "numLit" || !Number.isSafeInteger(expr.value)) return null; + // `x === -0` and `x === 0` select the same values. + return expr.value === 0 ? 0 : expr.value; +} + +function sameVarRef(a: IrExpr, b: IrExpr): boolean { + return a.kind === "varRef" && b.kind === "varRef" && a.localId === b.localId; +} + +/** A local read or a number literal: reading it has no effect. */ +function sameOperand(a: IrExpr, b: IrExpr): boolean { + if (a.kind === "numLit" && b.kind === "numLit") return Object.is(a.value, b.value); + return sameVarRef(a, b); +} + +/** The store of a single-statement arm, when it stores a number literal to + * an element named by a local and a local or a literal index (so evaluating + * them once, after the comparisons, reads what every arm would). */ +function literalStore(body: readonly IrStmt[] | null): BytesSet | null { + if (body === null || body.length !== 1) return null; + const s = body[0]!; + if (s.kind !== "bytesSet" || s.value.kind !== "numLit") return null; + if (s.arr.kind !== "varRef" || (s.index.kind !== "varRef" && s.index.kind !== "numLit")) + return null; + if (s.arr.type.kind !== "bytes" || s.index.type.kind !== "f64") return null; + return s; +} + +export function matchConstantChoice(s: IrStmt): ConstantChoice | null { + if (s.kind !== "if") return null; + let subject: IrExpr | null = null; + const cases: { value: number; result: number }[] = []; + const seen = new Set(); + const stores: BytesSet[] = []; + let current: IrStmt = s; + for (;;) { + if (current.kind !== "if") return null; + const cond = current.cond; + if (cond.kind !== "bin" || cond.op !== "===") return null; + const [local, literal] = + cond.left.kind === "varRef" ? [cond.left, cond.right] : [cond.right, cond.left]; + if (local.kind !== "varRef" || local.type.kind !== "f64") return null; + if (subject === null) subject = local; + else if (!sameVarRef(subject, local)) return null; + const value = integerLiteral(literal); + if (value === null) return null; + const store = literalStore(current.then); + if (store === null) return null; + stores.push(store); + // A repeated value never reaches its later arm. + if (!seen.has(value)) { + seen.add(value); + cases.push({ value, result: (store.value as IrExpr & { kind: "numLit" }).value }); + } + const rest: readonly IrStmt[] | null = current.else_; + if (rest === null || rest.length !== 1) return null; + if (rest[0]!.kind === "if") { + current = rest[0]!; + continue; + } + const last = literalStore(rest); + if (last === null) return null; + stores.push(last); + break; + } + const first = stores[0]!; + for (const store of stores) { + if (!sameVarRef(store.arr, first.arr) || !sameOperand(store.index, first.index)) return null; + if (store.arr.type.kind !== "bytes" || first.arr.type.kind !== "bytes") return null; + if (store.arr.type.elem !== first.arr.type.elem) return null; + } + if (cases.length < MIN_CASES) return null; + const values = cases.map((c) => c.value); + if (Math.max(...values) - Math.min(...values) >= MAX_TABLE) return null; + const fallback = (stores[stores.length - 1]!.value as IrExpr & { kind: "numLit" }).value; + return { subject: subject!, cases, fallback, stores }; +} + +function int32(value: number): boolean { + return ( + Number.isInteger(value) && value >= -2147483648 && value <= 2147483647 && !Object.is(value, -0) + ); +} + +/** The chosen number: the table entry for the subject's value, or the final + * else's literal for a value no case names (including NaN, a fraction and + * every value outside the table). */ +function emitChoice(host: LlvmEmitterContext, choice: ConstantChoice): LlValue { + const B = host.B; + const subject = host.emitExpr(choice.subject); + const results = [...choice.cases.map((c) => c.result), choice.fallback]; + const integral = results.every(int32); + const element = integral ? "i32" : "double"; + const literal = (value: number): string => (integral ? String(value) : f64Lit(value)); + + const proven = exactInteger(host, subject, choice.subject); + const values = choice.cases.map((c) => c.value); + let low = Math.min(...values); + let high = Math.max(...values); + // A proven range that fits a table needs no test: every value it allows + // has an entry. + const covered = + proven !== null && proven.range.max - proven.range.min < MAX_TABLE ? proven : null; + if (covered !== null) { + low = covered.range.min; + high = covered.range.max; + } + const size = high - low + 1; + const entries: string[] = []; + for (let value = low; value <= high; value++) { + const hit = choice.cases.find((c) => c.value === value); + entries.push(`${element} ${literal(hit === undefined ? choice.fallback : hit.result)}`); + } + const table = `@sc_choice_${createHash("sha256") + .update(`${element}:${entries.join(",")}`) + .digest("hex") + .slice(0, 16)}`; + host.declare( + `${table} = private unnamed_addr constant [${size} x ${element}] [${entries.join(", ")}]`, + ); + + let index: string; + let ok: string | null = null; + if (covered !== null) { + const wide = widenInteger(host, covered); + index = B.tmp(); + B.line(`${index} = sub i64 ${wide}, ${low}`); + } else { + // freeze keeps NaN and out-of-range conversions from producing poison; + // the round trip then rejects every value that is not an exact integer + // (-0 converts to 0, which `=== 0` selects too). + const raw = B.tmp(); + const selector = B.tmp(); + const restored = B.tmp(); + const exact = B.tmp(); + const offset = B.tmp(); + const inside = B.tmp(); + ok = B.tmp(); + index = B.tmp(); + B.line(`${raw} = fptosi double ${subject.name} to i64`); + B.line(`${selector} = freeze i64 ${raw}`); + B.line(`${restored} = sitofp i64 ${selector} to double`); + B.line(`${exact} = fcmp oeq double ${subject.name}, ${restored}`); + B.line(`${offset} = sub i64 ${selector}, ${low}`); + B.line(`${inside} = icmp ult i64 ${offset}, ${size}`); + B.line(`${ok} = and i1 ${exact}, ${inside}`); + B.line(`${index} = select i1 ${ok}, i64 ${offset}, i64 0`); + } + const slot = B.tmp(); + const entry = B.tmp(); + B.line( + `${slot} = getelementptr inbounds [${size} x ${element}], ptr ${table}, i64 0, i64 ${index}`, + ); + B.line(`${entry} = load ${element}, ptr ${slot}`); + let chosen = entry; + if (ok !== null) { + chosen = B.tmp(); + B.line( + `${chosen} = select i1 ${ok}, ${element} ${entry}, ${element} ${literal(choice.fallback)}`, + ); + } + if (!integral) return { name: chosen, type: choice.stores[0]!.value.type }; + const number = B.tmp(); + B.line(`${number} = sitofp i32 ${chosen} to double`); + return { + name: number, + type: choice.stores[0]!.value.type, + uint32: chosen, + integer: { + name: chosen, + type: "i32", + signed: Math.min(...results) < 0, + range: { min: Math.min(...results), max: Math.max(...results) }, + }, + }; +} + +/** Emits a matched chain as one store of the chosen literal. The element + * write is the arms' own: every arm names the same receiver and index. */ +export function emitConstantChoice(host: LlvmEmitterContext, choice: ConstantChoice): void { + const first = choice.stores[0]!; + if (first.arr.type.kind !== "bytes") + throw new InternalCompilerError("llvm emitter bug: constant choice on non-bytes"); + const arr = host.emitStableReceiver(first.arr, [first.index]); + const index = host.emitExpr(first.index); + const value = emitChoice(host, choice); + const scopes = choice.stores.map((store) => host.bytesAliasScopes.get(store)); + host.emitBytesSet( + first.arr.type.elem, + arr.name, + index, + value, + first.index, + undefined, + choice.stores.every((store) => host.bytesBounds.has(store)), + choice.stores.every((store) => host.bytesOwned.has(store)), + scopes.every((scope) => scope === scopes[0]) ? (scopes[0] ?? "") : "", + ); +} diff --git a/packages/compiler/src/backend/llvm/emit-statements.ts b/packages/compiler/src/backend/llvm/emit-statements.ts index 66bbf34d0..8bba1e7cf 100644 --- a/packages/compiler/src/backend/llvm/emit-statements.ts +++ b/packages/compiler/src/backend/llvm/emit-statements.ts @@ -93,6 +93,7 @@ import { import { emitStringInputs } from "./string-lifetimes.js"; import { LOOPS } from "./loop-state.js"; import { loopAccessesTypedArrays } from "../loop-polls.js"; +import { emitConstantChoice, matchConstantChoice } from "./constant-choices.js"; /** A callback loop version copies the loop; keep the copies small. */ const MAX_CALLBACK_VERSION_COST = 400; @@ -837,6 +838,13 @@ export function emitStmtBody(emitter: LlEmitter, s: IrStmt): void { break; } case "if": { + // A chain storing one of several literals to one element stores a + // table entry instead (debug builds keep the arms to step through). + const choice = emitter.debug === null ? matchConstantChoice(s) : null; + if (choice !== null) { + emitConstantChoice(emitter, choice); + break; + } const cond = emitter.emitCondition(s.cond); const trueLabel = B.newLabel("if.t"); const joinLabel = B.newLabel("if.j"); diff --git a/packages/compiler/test/ts7/baselines/order-parity/tests/corpus/constant-choice-chains.ts.json b/packages/compiler/test/ts7/baselines/order-parity/tests/corpus/constant-choice-chains.ts.json new file mode 100644 index 000000000..059a21156 --- /dev/null +++ b/packages/compiler/test/ts7/baselines/order-parity/tests/corpus/constant-choice-chains.ts.json @@ -0,0 +1,6 @@ +{ + "order": [ + "/tests/corpus/constant-choice-chains.ts" + ], + "diags": [] +} diff --git a/tests/corpus/constant-choice-chains.ts b/tests/corpus/constant-choice-chains.ts new file mode 100644 index 000000000..8d749daee --- /dev/null +++ b/tests/corpus/constant-choice-chains.ts @@ -0,0 +1,99 @@ +// if/else-if chains that store one of several literals to one typed-array +// element, chosen by comparing one number with integer literals: the +// subject's values that no case names, repeated and negative cases, +// results that do not fit the element, and chains that keep their arms. + +function bases(arr: Uint8Array): void { + for (let i = 0; i < arr.length; ++i) { + const x = arr[i]; + if (x === 97) arr[i] = 0; + else if (x === 116) arr[i] = 1; + else if (x === 99) arr[i] = 2; + else arr[i] = 3; + } +} + +const sequence = new Uint8Array(64); +for (let i = 0; i < sequence.length; i++) sequence[i] = "actgACTGnx"[i % 10]!.charCodeAt(0); +bases(sequence); +console.log("bases", sequence.join("")); + +// A number subject: NaN, -0, fractions, infinities and values far outside the cases. +const subjects = [0, -0, 1, 2, 3, 4, -1, 0.5, 2.0000001, NaN, Infinity, -Infinity, 2 ** 40, -(2 ** 40), 1e300]; +function classify(out: Int16Array, at: number, v: number): void { + if (v === 0) out[at] = 10; + else if (v === 1) out[at] = -20; + else if (v === 2) out[at] = 30; + else if (v === 1) out[at] = 99; // never reached + else if (v === -1) out[at] = 40; + else out[at] = -50; +} +const classes = new Int16Array(subjects.length); +for (let i = 0; i < subjects.length; i++) classify(classes, i, subjects[i]!); +console.log("classes", classes.join(" ")); + +// The literal on the left, and results the element converts: 300 wraps in a +// Uint8Array, -1 is 255, 1.5 truncates, and a clamped array rounds and clamps. +function convert(u8: Uint8Array, c: Uint8ClampedArray, f: Float64Array, i: Int8Array, v: number): void { + if (5 === v) u8[0] = 300; + else if (6 === v) u8[0] = -1; + else if (7 === v) u8[0] = 1.5; + else u8[0] = 8589934599; // 2 ** 33 + 7 + if (v === 5) c[0] = 300; + else if (v === 6) c[0] = -1; + else if (v === 7) c[0] = 2.5; + else c[0] = 3.5; + if (v === 5) f[0] = 0.25; + else if (v === 6) f[0] = -0; + else if (v === 7) f[0] = 1e300; + else f[0] = NaN; + if (v === 5) i[0] = 128; + else if (v === 6) i[0] = -129; + else if (v === 7) i[0] = -128; + else i[0] = 127; +} +for (const v of [5, 6, 7, 8, NaN]) { + const u8 = new Uint8Array(1); + const c = new Uint8ClampedArray(1); + const f = new Float64Array(1); + const i = new Int8Array(1); + convert(u8, c, f, i, v); + console.log("convert", v, u8[0], c[0], Object.is(f[0], -0) ? "-0" : f[0], i[0]); +} + +// An index past the end ignores the store, whatever the chain chooses. +const short = new Uint8Array(2); +for (let k = 0; k < 4; k++) { + const v = k + 1; + if (v === 1) short[k] = 11; + else if (v === 2) short[k] = 22; + else if (v === 3) short[k] = 33; + else short[k] = 44; +} +console.log("short", short.join(",")); + +// Chains that keep their arms: an arm with a second statement, no final +// else, a different index, and a subject that is not a local. +let calls = 0; +function count(): number { + calls++; + return 9; +} +const kept = new Uint8Array(8); +for (let k = 0; k < 8; k++) { + const v = k % 5; + if (v === 0) kept[k] = 1; + else if (v === 1) { + kept[k] = 2; + calls += 100; + } else if (v === 2) kept[k] = 3; + else kept[k] = 4; + if (v === 0) kept[0] += 1; + else if (v === 1) kept[0] += 2; + else if (v === 2) kept[0] += 3; + if (count() === v) kept[k] = 7; + else if (v === 3) kept[(k + 1) % 8] = 5; + else if (v === 4) kept[k] = 6; + else kept[k] = kept[k]!; +} +console.log("kept", kept.join(","), calls); diff --git a/tests/harness/corpus-costs/constant-choice-chains.ts.json b/tests/harness/corpus-costs/constant-choice-chains.ts.json new file mode 100644 index 000000000..6562c9db3 --- /dev/null +++ b/tests/harness/corpus-costs/constant-choice-chains.ts.json @@ -0,0 +1 @@ +{ "plain": 3921, "san": 5592 }