Store a table entry for if/else chains that store literals to one element

An if/else-if chain that compares one number local with integer literals
and stores a number literal to the same typed-array element in every arm,
including a final else, such as

    if (x === 97) a[i] = 0;
    else if (x === 116) a[i] = 1;
    else if (x === 99) a[i] = 2;
    else a[i] = 3;

is now emitted as one store of an entry of a constant table indexed by
the local's value (constant-choices.ts). On input whose cases do not
repeat in a pattern, such as the bases of a DNA sequence, the branches
mispredict; LLVM keeps them, because cases this far apart are too sparse
for its own switch tables.

The chain needs at least three distinct values within 256 of each other.
The table covers the cases' span, or the local's proven integer range
when that range has at most 256 values, in which case no test is
emitted. Otherwise a value that is not an exact integer in the table's
span (NaN, a fraction, or a value outside it) stores the final else's
literal, chosen with a select. The first arm wins for a repeated
value, and -0 selects the arm for 0, as strict equality does. The store
is the arms' own (the same receiver and index), converted to the
element's type as each arm's would be. Debug builds keep the arms.

k-nucleotide's transform of the input sequence uses this form: its
median wall time drops from 2.71 s to 2.21 s.

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();
+});

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<IrStmt, { kind: "bytesSet" }>;
+
+/** 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<number>();
+  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] ?? "") : "",
+  );
+}

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");

packages/compiler/test/ts7/baselines/order-parity/tests/corpus/constant-choice-chains.ts.json

@@ -0,0 +1,6 @@
+{
+ "order": [
+  "<repo>/tests/corpus/constant-choice-chains.ts"
+ ],
+ "diags": []
+}

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);

tests/harness/corpus-costs/constant-choice-chains.ts.json

@@ -0,0 +1 @@
+{ "plain": 3921, "san": 5592 }