Read typed-array header fields once at a speculative loop's entry (#2321)

packages/compiler/src/backend/llvm/bytes-headers.ts

@@ -0,0 +1,77 @@
+/* Typed-array header fields read once at a speculative loop's entry.
+ *
+ * A view's length, data pointer and shared storage are set at construction
+ * and never change, and a view held by a const local or a parameter that
+ * the loop does not rebind is the same view in every iteration. The entry
+ * reads those fields once, and every element access in the loop's copies
+ * (the ordinary copy, the fast copy and the copy a failed check resumes in)
+ * uses the values. No header read is then left inside the copies for LLVM's
+ * load PRE to move ahead of the fast copy's exits: such a move keeps the
+ * fast loop's first iteration in the loop, with the checks that one
+ * iteration of it makes redundant. */
+import type { IrExpr, IrStmt } from "../../ir/ir.js";
+import { everyStmt } from "../../ir/traverse.js";
+import { emitBytesData, emitBytesHeaderLength, emitBytesSharedStorage } from "./expr-bytes.js";
+import type { LlEmitter } from "./emitter.js";
+
+export interface BytesHeader {
+  length: string;
+  data: string;
+  /** The view's shared storage (worker programs only). */
+  shared: string | null;
+}
+
+export interface BytesHeaderScope {
+  /** The function the scope's values belong to. */
+  fn: string;
+  locals: ReadonlyMap<string, BytesHeader>;
+  /** The header of each receiver value read from one of `locals`. */
+  values: Map<string, BytesHeader>;
+}
+
+/** Reads the header fields of the views `loop` accesses through a const
+ * local or a parameter it does not rebind, at the current block, which must
+ * dominate every copy of the loop. Null when there are none. */
+export function emitLoopBytesHeaders(
+  emitter: LlEmitter,
+  loop: IrStmt,
+  parameters: ReadonlySet<string>,
+): BytesHeaderScope | null {
+  const receivers = new Set<string>();
+  const rebound = new Set<string>();
+  const receiver = (e: IrExpr): void => {
+    if (e.kind === "varRef") receivers.add(e.localId);
+  };
+  everyStmt(loop, {
+    expr: (e) => {
+      if (e.kind === "bytesIntrinsic") receiver(e.receiver);
+      else if (e.kind === "assignExpr" || e.kind === "incDec") rebound.add(e.localId);
+      return true;
+    },
+    stmt: (s) => {
+      if (s.kind === "bytesSet") receiver(s.arr);
+      else if (s.kind === "assign" || s.kind === "forOf" || s.kind === "varDecl")
+        rebound.add(s.localId);
+      else if (s.kind === "tryCatch" && s.catchLocalId !== null) rebound.add(s.catchLocalId);
+      return true;
+    },
+  });
+  const locals = new Map<string, BytesHeader>();
+  for (const id of receivers) {
+    if (rebound.has(id)) continue;
+    const local = emitter.currentLocals.get(id);
+    if (local === undefined || local.boxed === true || local.type.kind !== "bytes") continue;
+    if (local.mutable && !parameters.has(id)) continue;
+    if (emitter.integerLoopBindings.has(id) || emitter.int32Locals.has(id)) continue;
+    const B = emitter.B;
+    const view = B.tmp();
+    B.line(`${view} = load ptr, ptr ${emitter.binding(id).slot}${emitter.slotAlias(id)}`);
+    locals.set(id, {
+      length: emitBytesHeaderLength(emitter, view),
+      data: emitBytesData(emitter, view),
+      shared: emitter.mod.workers ? emitBytesSharedStorage(emitter, view) : null,
+    });
+  }
+  if (locals.size === 0) return null;
+  return { fn: emitter.currentFunctionName, locals, values: new Map() };
+}

packages/compiler/src/backend/llvm/emit-function.ts

@@ -83,6 +83,7 @@ export function emitFunction(emitter: LlEmitter, fn: IrFunction): string {
   emitter.unwindCleanups.clear();
   emitter.jumpTargets = [];
   emitter.currentLocals = new Map(fn.locals.map((l) => [l.id, l]));
+  emitter.currentParameters = new Set(fn.params.map((p) => p.localId));
   emitter.lazyCaptureBoxes.clear();
   emitter.lazyStackBoxes.clear();
   // An unchanged captured parameter keeps its value in a plain slot; the

packages/compiler/src/backend/llvm/emit-statements.ts

@@ -1,4 +1,6 @@
 /* LlEmitter's statement emission: the statement switch, try/catch, switch, loop copies, speculative writes and string self-concatenation. */
+import { emitLoopBytesHeaders } from "./bytes-headers.js";
+import { noteBytesReceiver } from "./expr-bytes.js";
 import {
   type LlEmitter,
   type LlScopeEntry,
@@ -81,6 +83,7 @@ import {
   planProvenIntegerLoop,
   planSpeculativeLoop,
   type ProvenLoopPlan,
+  type SpeculativeLoopPlan,
   speculativeReads,
   speculativeWriteTarget,
 } from "./speculative-integers.js";
@@ -577,6 +580,7 @@ export function emitStmtBody(emitter: LlEmitter, s: IrStmt): void {
       // IrBytesElem is static, so never rediscover it through the
       // generic runtime switch in a hot loop.
       const arr = emitter.emitStableReceiver(s.arr, [s.index, s.value]);
+      noteBytesReceiver(emitter, s.arr, arr.name);
       const idx = emitter.emitExpr(s.index);
       const v = emitter.emitExpr(s.value);
       if (s.arr.type.kind !== "bytes")
@@ -2381,6 +2385,22 @@ export function emitLoopCopies(
     else ordinary(exit);
     return;
   }
+  const outer = emitter.bytesHeaders;
+  emitter.bytesHeaders = emitLoopBytesHeaders(emitter, loop, emitter.currentParameters) ?? outer;
+  try {
+    emitSpeculativeLoop(emitter, plan, exit, emitCopy);
+  } finally {
+    emitter.bytesHeaders = outer;
+  }
+}
+
+/** The copies of a speculative loop (emitLoopCopies). */
+function emitSpeculativeLoop(
+  emitter: LlEmitter,
+  plan: SpeculativeLoopPlan,
+  exit: string,
+  emitCopy: (exit: string) => void,
+): void {
   const B = emitter.B;
   const entries = new Map<string, string>();
   let ok = "true";

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

@@ -64,6 +64,7 @@ import { captureFreeFunctions, observesStacks } from "../stack-observations.js";
 import { findInitializedByteLoopBindings } from "./byte-windows.js";
 import { matchBytesLoopVersion } from "./bytes-loop-versions.js";
 import { globalSlotAlias } from "./bytes-alias.js";
+import type { BytesHeaderScope } from "./bytes-headers.js";
 import type { IntegerCountedForLoop, IntegerPairedForLoop } from "../../ir/integer-loops.js";
 import {
   materializeSplitPiece,
@@ -1841,6 +1842,11 @@ export class LlEmitter {
    * one the element access being emitted carries (bytesElementAlias). */
   bytesAliasScopes: ReadonlyMap<IrExpr | IrStmt, string> = new Map();
   bytesAccessScope = "";
+  /** Header fields a speculative loop's entry read for its views
+   * (bytes-headers.ts), while its copies are emitted. */
+  bytesHeaders: BytesHeaderScope | null = null;
+  /** The parameters of the function being emitted. */
+  currentParameters: ReadonlySet<string> = new Set();
   /** Loop versions with alias scopes in the current function; it names
    * their scopes. */
   aliasScopedVersions = 0;

packages/compiler/src/backend/llvm/expr-bytes.ts

@@ -1,6 +1,7 @@
 /* Focused LLVM expression emission extracted from emitter.ts. */
 import { InternalCompilerError } from "../../errors.js";
 import { emitBorrowedInput } from "./borrowed-inputs.js";
+import type { BytesHeader } from "./bytes-headers.js";
 import { BYTES_ELEMENT_SIZE, F64, type IrBytesElem, type IrExpr } from "../../ir/ir.js";
 import type { LlvmEmitterContext, LlValue } from "./expr-context.js";
 import { F64_INF, exactIndexLines, f64Lit, sizeToDouble } from "./common.js";
@@ -160,6 +161,22 @@ export function emitBytesIndex(
   return idx;
 }
 
+/** Records that `value` is the view `receiver` reads, so the header reads of
+ * an access through it use the values a loop entry read (bytes-headers.ts). */
+export function noteBytesReceiver(host: LlvmEmitterContext, receiver: IrExpr, value: string): void {
+  const scope = host.bytesHeaders;
+  if (scope === null || receiver.kind !== "varRef" || scope.fn !== host.currentFunctionName) return;
+  const header = scope.locals.get(receiver.localId);
+  if (header !== undefined) scope.values.set(value, header);
+}
+
+/** The header fields a loop entry read for the view `receiver`, if any. */
+function cachedHeader(host: LlvmEmitterContext, receiver: string): BytesHeader | undefined {
+  const scope = host.bytesHeaders;
+  if (scope === null || scope.fn !== host.currentFunctionName) return undefined;
+  return scope.values.get(receiver);
+}
+
 /** The element count of a typed array, in the target's size type. `value`:
  * the program reads the count itself (bytesHeaderValueAlias), rather than
  * only bounding an access. */
@@ -168,6 +185,8 @@ export function emitBytesHeaderLength(
   receiver: string,
   value = false,
 ): string {
+  const cached = cachedHeader(host, receiver);
+  if (cached !== undefined) return cached.length;
   const p = host.B.tmp();
   const len = host.B.tmp();
   const alias = value ? bytesHeaderValueAlias(host) : bytesHeaderAlias(host);
@@ -177,6 +196,8 @@ export function emitBytesHeaderLength(
 }
 
 export function emitBytesData(host: LlvmEmitterContext, receiver: string): string {
+  const cached = cachedHeader(host, receiver);
+  if (cached !== undefined) return cached.data;
   const p = host.B.tmp();
   const data = host.B.tmp();
   host.B.line(`${p} = getelementptr inbounds %ScrBytes, ptr ${receiver}, i64 0, i32 3`);
@@ -203,6 +224,8 @@ export function emitBytesLength(
 /** A view's shared backing (null for a buffer this thread owns). It is
  * set at construction, before script code can see the view. */
 export function emitBytesSharedStorage(host: LlvmEmitterContext, receiver: string): string {
+  const cached = cachedHeader(host, receiver)?.shared;
+  if (cached !== undefined && cached !== null) return cached;
   const B = host.B;
   const slot = B.tmp(),
     storage = B.tmp();
@@ -813,6 +836,7 @@ export function emitBytesIntrinsic(
   const r = directElementAccess
     ? host.emitStableReceiver(e.receiver, e.args)
     : emitBorrowedInput(host, e.receiver);
+  if (directElementAccess) noteBytesReceiver(host, e.receiver, r.name);
   // Byte operations borrow every input; views and identity results return
   // independent owners, while writes retain no operand beyond the call.
   const args = e.args.map((a) => emitBorrowedInput(host, a));

packages/compiler/test/speculative-loop-emission.test.ts

@@ -183,3 +183,47 @@ console.log(key(new Uint8Array([1, 2, 3, 0]), 6));
   expect(fast).not.toMatch(/ashr i64 %/);
   expect(fast).not.toContain("@llvm.smul.with.overflow.i64");
 });
+
+test("a speculative loop reads its views' header fields once, before its copies", async () => {
+  // A view held by a const local or by a parameter is the same view in
+  // every iteration. Its data pointer is read before the branch to the fast
+  // copy, and no copy (ordinary, fast, or the one a failed check resumes
+  // in) reads it again.
+  const llvm = await emitLlvm(`
+let last = 42;
+function fill(array: ArrayBuffer, count: number): void {
+  const buf = new Float64Array(array);
+  let l = last;
+  for (let i = 0; i < count; i++) {
+    l = (l * 3877 + 29573) % 139968;
+    buf[i] = l;
+  }
+  last = l;
+}
+function fillInto(buf: Float64Array, count: number): void {
+  let l = last;
+  for (let i = 0; i < count; i++) {
+    l = (l * 3877 + 29573) % 139968;
+    buf[i] = l;
+  }
+  last = l;
+}
+const array = new ArrayBuffer(8 * 16);
+fill(array, 16);
+const view = new Float64Array(8);
+fillInto(view, 8);
+console.log(new Float64Array(array)[15], view[7], last);
+`);
+  const dataReads = /getelementptr inbounds %ScrBytes, ptr %\w+, i64 0, i32 3$/gm;
+  for (const name of ["fill", "fillInto"]) {
+    const fn =
+      new RegExp(`^define internal void @sc_bf_${name}(?:\\.i32)?\\([^]*?^}`, "m").exec(
+        llvm,
+      )?.[0] ?? "";
+    const branch = /^ {2}br i1 %\w+, label %spec\.fast\d+, label %spec\.slow\d+$/m.exec(fn);
+    expect(branch, name).not.toBeNull();
+    const copies = fn.slice(branch!.index);
+    expect(copies.match(dataReads), name).toBeNull();
+    expect(fn.slice(0, branch!.index).match(dataReads)?.length, name).toBeGreaterThan(0);
+  }
+});

packages/compiler/test/ts7/baselines/order-parity/tests/corpus/speculative-loop-view-headers.ts.json

@@ -0,0 +1,6 @@
+{
+ "order": [
+  "<repo>/tests/corpus/speculative-loop-view-headers.ts"
+ ],
+ "diags": []
+}

tests/corpus/speculative-loop-view-headers.ts

@@ -0,0 +1,36 @@
+// Speculative integer loops over typed-array views held by a const local
+// and by a parameter. The fast copy, the copy a failed check resumes in and
+// the ordinary copy all read and write the same views: a product that
+// leaves the safe integers fails the fast copy's check partway through,
+// and writes past the end are ignored.
+function fill(buffer: ArrayBuffer, count: number, seed: number): number {
+  const view = new Float64Array(buffer);
+  let l = seed;
+  for (let i = 0; i < count; i++) {
+    l = (l * 3877 + 29573) % 139968;
+    view[i] = l;
+  }
+  return l;
+}
+
+function grow(out: Int32Array, count: number, start: number): number {
+  let x = start;
+  for (let i = 0; i < count; i++) {
+    x = x * 3 + 1;
+    out[i] = x % 1000;
+  }
+  return x;
+}
+
+function sum(view: Float64Array | Int32Array): number {
+  let s = 0;
+  for (const v of view) s += v;
+  return s;
+}
+
+const buffer = new ArrayBuffer(8 * 64);
+console.log(fill(buffer, 64, 42), sum(new Float64Array(buffer)));
+console.log(fill(buffer, 100, 7), new Float64Array(buffer)[63]);
+const out = new Int32Array(40);
+console.log(grow(out, 60, 1), sum(out), out[0], out[39]);
+console.log(grow(out, 10, 5), out.join(","));

tests/harness/corpus-costs/speculative-loop-view-headers.ts.json

@@ -0,0 +1 @@
+{ "plain": 2400 }