Merge remote-tracking branch 'origin/perf/fold-constant-globals' into HEAD

packages/compiler/src/backend/llvm/emitter.ts

@@ -100,6 +100,7 @@ import { everyExprChild, everyStmtChild, everyStmtList } from "../../ir/traverse
 import { analyzeIntegerRanges, INT32_RANGE, type IntegerRanges } from "../../ir/integer-ranges.js";
 import { analyzeInt32Slots, type Int32Slots } from "../../ir/int32-slots.js";
 import { analyzeFinalFields, constructorClass, type FinalFields } from "../../ir/final-fields.js";
+import { foldConstantGlobals } from "../../ir/constant-globals.js";
 import { findIntegerViews } from "./integer-views.js";
 import { findInitializerBindings, withInitializerBindings } from "../../ir/initializer-bindings.js";
 import { findConstantNumericTables, type ConstantNumericTable } from "../../ir/constant-tables.js";
@@ -405,7 +406,7 @@ export function emitLlvmModule(mod: IrModule, options: LlvmTargetOptions = {}):
   // Keep source storage intact for debugger inspection in dev builds.
   return new LlEmitter(
     options.debugSources === undefined
-      ? scalarizeNumericRecords(specializeNumericCalls(mod, mangleFunction))
+      ? scalarizeNumericRecords(specializeNumericCalls(foldConstantGlobals(mod), mangleFunction))
       : mod,
     options,
   ).emit();
@@ -418,7 +419,7 @@ export function emitLlvmModuleSource(
 ): string | readonly string[] {
   const parts = new LlEmitter(
     options.debugSources === undefined
-      ? scalarizeNumericRecords(specializeNumericCalls(mod, mangleFunction))
+      ? scalarizeNumericRecords(specializeNumericCalls(foldConstantGlobals(mod), mangleFunction))
       : mod,
     options,
   ).emitParts();
@@ -6057,14 +6058,49 @@ export class LlEmitter {
         }
         let countedJoin: string | null = null;
         if (countedLoop?.guarded) {
-          const bound = this.emitExpr(countedLoop.limit);
-          const safe = B.tmp();
+          let safe = "true";
+          if (countedLoop.limitGuarded) {
+            const bound = this.emitExpr(countedLoop.limit);
+            safe = B.tmp();
+            B.line(
+              `${safe} = fcmp ${countedLoop.step > 0 ? "ole" : "oge"} double ${bound.name}, ${f64Lit(countedLoop.guardLimit)}`,
+            );
+          }
+          if (countedLoop.startGuarded) {
+            // An exact integer of at most 2^53 - 1 in magnitude, and not -0
+            // (the bit patterns of the value and its round trip agree).
+            // NaN and infinities fail the ordered range tests.
+            const start = this.emitExpr(countedLoop.start);
+            const low = B.tmp(),
+              high = B.tmp(),
+              raw = B.tmp(),
+              wide = B.tmp(),
+              back = B.tmp(),
+              bits = B.tmp(),
+              backBits = B.tmp(),
+              exact = B.tmp(),
+              inRange = B.tmp(),
+              ok = B.tmp();
+            B.line(`${low} = fcmp oge double ${start.name}, ${f64Lit(-Number.MAX_SAFE_INTEGER)}`);
+            B.line(`${high} = fcmp ole double ${start.name}, ${f64Lit(Number.MAX_SAFE_INTEGER)}`);
+            B.line(`${inRange} = and i1 ${low}, ${high}`);
+            B.line(`${raw} = fptosi double ${start.name} to i64`);
+            B.line(`${wide} = freeze i64 ${raw}`);
+            B.line(`${back} = sitofp i64 ${wide} to double`);
+            B.line(`${bits} = bitcast double ${start.name} to i64`);
+            B.line(`${backBits} = bitcast double ${back} to i64`);
+            B.line(`${exact} = icmp eq i64 ${bits}, ${backBits}`);
+            B.line(`${ok} = and i1 ${inRange}, ${exact}`);
+            if (safe === "true") safe = ok;
+            else {
+              const both = B.tmp();
+              B.line(`${both} = and i1 ${safe}, ${ok}`);
+              safe = both;
+            }
+          }
           const fast = B.newLabel("counted.fast");
           const slow = B.newLabel("counted.slow");
           countedJoin = B.newLabel("counted.done");
-          B.line(
-            `${safe} = fcmp ${countedLoop.step > 0 ? "ole" : "oge"} double ${bound.name}, ${f64Lit(countedLoop.guardLimit)}`,
-          );
           B.condBr(safe, fast, slow);
           B.startBlock(slow);
           this.countedLoopsEnabled = false;

packages/compiler/src/ir/constant-globals.test.ts

@@ -0,0 +1,95 @@
+import { expect, test } from "vitest";
+import { foldConstantGlobals } from "./constant-globals.js";
+import { F64, VOID, type IrExpr, type IrFunction, type IrModule, type IrStmt } from "./ir.js";
+
+const loc = { file: "constant-globals.ts", start: 0, end: 0 };
+const num = (value: number): IrExpr => ({ loc, kind: "numLit", value, type: F64 });
+const ref = (localId: string): IrExpr => ({ loc, kind: "varRef", localId, type: F64 });
+const assign = (localId: string, value: IrExpr): IrStmt => ({
+  loc,
+  kind: "assign",
+  localId,
+  value,
+});
+const ret = (value: IrExpr): IrStmt => ({ loc, kind: "return", value });
+const fn = (name: string, body: IrStmt[]): IrFunction => ({
+  loc,
+  name,
+  params: [],
+  locals: [],
+  returnType: name === "%init" ? VOID : F64,
+  body,
+});
+const global = (id: string, extra: object = {}) => ({
+  id,
+  name: id,
+  type: F64,
+  mutable: false,
+  ...extra,
+});
+
+test("immutable globals stored once with a literal read as that literal", () => {
+  const mod: IrModule = {
+    irVersion: 15,
+    sourceFile: "constant-globals.ts",
+    entry: "%init",
+    functions: [
+      fn("%init", [
+        assign("%g.DOT", num(46)),
+        assign("%g.LATE", num(7)),
+        assign("%g.TWICE", num(1)),
+        assign("%g.TWICE", num(2)),
+        assign("%g.MUT", num(3)),
+      ]),
+      fn("read", [
+        ret({
+          loc,
+          kind: "bin",
+          op: "+",
+          left: ref("%g.DOT"),
+          right: {
+            loc,
+            kind: "bin",
+            op: "+",
+            left: ref("%g.LATE"),
+            right: {
+              loc,
+              kind: "bin",
+              op: "+",
+              left: ref("%g.TWICE"),
+              right: ref("%g.MUT"),
+              type: F64,
+            },
+            type: F64,
+          },
+          type: F64,
+        }),
+      ]),
+    ],
+    classes: [],
+    records: [],
+    unions: [],
+    globals: [
+      global("%g.DOT"),
+      global("%g.LATE", { initFlag: "%g.LATE.initialized" }),
+      global("%g.TWICE"),
+      global("%g.MUT", { mutable: true }),
+    ],
+  };
+  const folded = foldConstantGlobals(mod);
+  const reads: string[] = [];
+  const walk = (e: IrExpr): void => {
+    if (e.kind === "varRef") reads.push(e.localId);
+    else if (e.kind === "numLit") reads.push(String(e.value));
+    else if (e.kind === "bin") {
+      walk(e.left);
+      walk(e.right);
+    }
+  };
+  const body = folded.functions.find((f) => f.name === "read")!.body[0]!;
+  if (body.kind === "return" && body.value) walk(body.value);
+  expect(reads).toEqual(["46", "%g.LATE", "%g.TWICE", "%g.MUT"]);
+  // The initializer keeps its store; the input module is not modified.
+  expect(folded.functions[0]).toBe(mod.functions[0]);
+  expect(mod.functions[1]!.body[0]).not.toBe(body);
+});

packages/compiler/src/ir/constant-globals.ts

@@ -0,0 +1,76 @@
+import type { IrExpr, IrFunction, IrModule, IrStmt } from "./ir.js";
+import { everyStmtList, mapExprChildren, mapStmtChildren } from "./traverse.js";
+
+type Literal = (IrExpr & { kind: "numLit" }) | (IrExpr & { kind: "boolLit" });
+
+/** Module constants such as `const DOT = 0x2e` read as their literal.
+ *
+ * A module-level binding is a global that the module initializer stores
+ * once; every other function loads it (a thread-local load in worker
+ * programs), which also hides the value from integer ranges, switch
+ * dispatch and LLVM's constant folding. When the global is immutable, has
+ * exactly one write in the program and that write stores a number or
+ * boolean literal, every read sees that literal: reads that could run before
+ * the store carry the frontend's initialization flag (`initFlag`) or a TDZ
+ * sentinel, and those globals are left alone. The store stays, so module
+ * namespace views and debuggers still find the value in the global. Library
+ * builds keep their globals: a host may observe them across the boundary. */
+export function foldConstantGlobals(mod: IrModule): IrModule {
+  if (mod.lib || !mod.globals?.length) return mod;
+  const candidates = new Set<string>();
+  for (const g of mod.globals)
+    if (
+      !g.mutable &&
+      g.tdz !== true &&
+      g.initFlag === undefined &&
+      (g.type.kind === "f64" || g.type.kind === "bool")
+    )
+      candidates.add(g.id);
+  if (candidates.size === 0) return mod;
+
+  const writes = new Map<string, number>();
+  const values = new Map<string, Literal>();
+  const write = (id: string, value: IrExpr | null): void => {
+    if (!candidates.has(id)) return;
+    writes.set(id, (writes.get(id) ?? 0) + 1);
+    if (value && (value.kind === "numLit" || value.kind === "boolLit")) values.set(id, value);
+    else values.delete(id);
+  };
+  for (const fn of mod.functions)
+    everyStmtList(fn.body, {
+      // Every node other than a read that names the binding is a write
+      // (assignments, loop and catch bindings, compound updates); only a
+      // plain assignment or declaration contributes a value.
+      stmt: (s) => {
+        if (s.kind === "assign") write(s.localId, s.value);
+        else if (s.kind === "varDecl") write(s.localId, s.init);
+        else if ("localId" in s && typeof s.localId === "string") write(s.localId, null);
+        return true;
+      },
+      expr: (e) => {
+        if (e.kind !== "varRef" && "localId" in e && typeof e.localId === "string")
+          write(e.localId, null);
+        return true;
+      },
+    });
+  const constants = new Map<string, Literal>();
+  for (const [id, value] of values) if (writes.get(id) === 1) constants.set(id, value);
+  if (constants.size === 0) return mod;
+
+  const reads = (e: IrExpr): boolean => !(e.kind === "varRef" && constants.has(e.localId));
+  const rewriteExpr = (e: IrExpr): IrExpr => {
+    if (e.kind === "varRef") {
+      const value = constants.get(e.localId);
+      return value ? { ...value, type: e.type, loc: e.loc } : e;
+    }
+    return mapExprChildren(e, rewriteExpr, rewriteStmt);
+  };
+  const rewriteStmt = (s: IrStmt): IrStmt => mapStmtChildren(s, rewriteExpr, rewriteStmt);
+  let changed = false;
+  const functions = mod.functions.map((fn): IrFunction => {
+    if (everyStmtList(fn.body, { stmt: () => true, expr: reads })) return fn;
+    changed = true;
+    return { ...fn, body: fn.body.map(rewriteStmt) };
+  });
+  return changed ? { ...mod, functions } : mod;
+}

packages/compiler/src/ir/integer-loops.test.ts

@@ -164,3 +164,21 @@ test("versioning is bounded and requires an integer consumer", () => {
   expect(matchIntegerCountedForLoop(nested, withLimit)).toBeNull();
   expect(matchIntegerCountedForLoop(counted(number(4)), locals)?.guarded).toBe(false);
 });
+
+test("an unproven start versions the loop on an exact integer start", () => {
+  const withStart = new Map(locals);
+  withStart.set("start", { id: "start", name: "start", type: F64, mutable: false });
+  const s = counted(number(8));
+  s.init = { kind: "varDecl", localId: "i", init: ref("start"), loc };
+  expect(matchIntegerCountedForLoop(s, withStart)).toMatchObject({
+    guarded: true,
+    startGuarded: true,
+    limitGuarded: false,
+    range: { min: -Number.MAX_SAFE_INTEGER, max: 7 },
+  });
+  // The versioned body is duplicated, so the usual size bound applies.
+  const large = counted(number(8));
+  large.init = { kind: "varDecl", localId: "i", init: ref("start"), loc };
+  large.body = Array.from({ length: 81 }, () => large.body[0]!);
+  expect(matchIntegerCountedForLoop(large, withStart)).toBeNull();
+});

packages/compiler/src/ir/integer-loops.ts

@@ -20,6 +20,11 @@ export interface IntegerCountedForLoop {
   step: number;
   inclusive: boolean;
   guarded: boolean;
+  /** The limit needs the run-time `guardLimit` test (part of `guarded`). */
+  limitGuarded: boolean;
+  /** The start has no proven integer range: the versioned loop also tests
+   * that it is an exact integer (not -0) of at most 2^53 in magnitude. */
+  startGuarded: boolean;
   guardLimit: number;
   range: IntegerRange;
 }
@@ -104,10 +109,13 @@ export function matchIntegerCountedForLoop(
   const local = locals.get(init.localId);
   if (local?.type.kind !== "f64" || !local.mutable || local.boxed || local.tdz) return null;
   const start = init.init;
-  const startRange =
+  let startRange =
     start.kind === "numLit" && Number.isSafeInteger(start.value) && !Object.is(start.value, -0)
       ? { min: start.value, max: start.value }
       : ranges.get(start);
+  // An unproven start is tested at run time, on a versioned loop.
+  const startGuarded = !startRange && start.kind !== "numLit" && start.type.kind === "f64";
+  if (startGuarded) startRange = { min: -Number.MAX_SAFE_INTEGER, max: Number.MAX_SAFE_INTEGER };
   if (
     !startRange ||
     startRange.min < -Number.MAX_SAFE_INTEGER ||
@@ -132,9 +140,10 @@ export function matchIntegerCountedForLoop(
   const guardLimit = step > 0 ? 2 ** 53 - (inclusive ? 1 : 0) : -(2 ** 53) + (inclusive ? 1 : 0);
   const boundRange =
     limit.kind === "numLit" ? { min: limit.value, max: limit.value } : ranges.get(limit);
-  const guarded =
+  const limitGuarded =
     !boundRange || (step > 0 ? !(boundRange.max <= guardLimit) : !(boundRange.min >= guardLimit));
-  if (guarded && limit.kind === "numLit") return null;
+  if (limitGuarded && limit.kind === "numLit") return null;
+  const guarded = limitGuarded || startGuarded;
   let useful = false;
   let cost = 0;
   const counter = (e: IrExpr): boolean =>
@@ -204,7 +213,18 @@ export function matchIntegerCountedForLoop(
           ),
           max: startRange.max,
         };
-  return { localId: local.id, start, limit, step, inclusive, guarded, guardLimit, range };
+  return {
+    localId: local.id,
+    start,
+    limit,
+    step,
+    inclusive,
+    guarded,
+    limitGuarded,
+    startGuarded,
+    guardLimit,
+    range,
+  };
 }
 
 /** True when a lowered subtree writes `localId`. Local ids are unique per

tests/corpus/constant-module-globals/consts.ts

@@ -0,0 +1,10 @@
+export const SLASH = 0x2f;
+export const DOT = 0x2e;
+export const NEGATIVE = -1;
+export const HALF = 0.5;
+export const ENABLED = true;
+export const MINUS_ZERO = -0;
+export let counter = 0;
+export function bump(): number {
+  return ++counter;
+}

tests/corpus/constant-module-globals/main.ts

@@ -0,0 +1,51 @@
+// Module constants read as their literal values; bindings that change or
+// can be read before their declaration keep their global.
+import * as ns from "./consts.ts";
+import { DOT, ENABLED, HALF, MINUS_ZERO, NEGATIVE, SLASH, bump, counter } from "./consts.ts";
+
+const COLON = 0x3a;
+const LIMIT = 3;
+
+function classify(code: number): string {
+  switch (code) {
+    case SLASH:
+      return "slash";
+    case DOT:
+      return "dot";
+    case COLON:
+      return "colon";
+    case NEGATIVE:
+      return "negative";
+    default:
+      return "other";
+  }
+}
+
+const path = "a/b.c:d";
+const kinds: string[] = [];
+for (let i = 0; i < path.length; i++) kinds.push(classify(path.charCodeAt(i)));
+console.log(kinds.join(","), classify(-1), classify(NaN));
+console.log(HALF * 3, ENABLED && LIMIT > 2, Object.is(MINUS_ZERO, -0), 1 / MINUS_ZERO);
+console.log(ns.DOT, ns.SLASH, Object.keys(ns).sort().join(" "));
+
+// A mutable export keeps observing writes.
+bump();
+bump();
+console.log(counter, ns.counter);
+
+// A function hoisted above the declaration it reads still throws in the TDZ.
+try {
+  console.log(early());
+} catch (error) {
+  console.log((error as Error).name, (error as Error).message);
+}
+const LATE = 7;
+function early(): number {
+  return LATE * 2;
+}
+console.log(early());
+
+// A loop binding at module scope is written once per iteration.
+const totals: number[] = [];
+for (const step of [1, 2, 3]) totals.push(step * LIMIT);
+console.log(totals.join(" "));

tests/corpus/exact-index-round-trip.ts

@@ -96,3 +96,17 @@ console.log(
   rangeHash(cursor.units, 0.5, 4),
   rangeHash(cursor.units, opaque(odd, 1), opaque(odd, 3)),
 );
+
+// Counted loops from an unproven start run as integer loops only when the
+// start is an exact integer; -0, fractions, NaN and huge starts keep the
+// double loop, and the counter's own value stays observable.
+function observe(chars: Uint16Array, start: number, end: number): string {
+  const seen: string[] = [];
+  for (let i = start; i < end; i++) {
+    seen.push(`${Object.is(i, -0) ? "-0" : i}:${chars[i] ^ 1}:${(i * 2) | 0}`);
+    if (seen.length > 6) break;
+  }
+  return seen.join(" ");
+}
+for (const start of [0, -0, 1, 2.5, -2, NaN, 2 ** 53 - 2, -(2 ** 53), 1e300, -Infinity])
+  console.log(String(start), observe(cursor.units, opaque([start], 0), opaque([start + 4, 9], start === -Infinity ? 1 : 0)));