Version two-counter loops over typed arrays (#2021)

packages/compiler/src/backend/llvm/bytes-loop-versions.test.ts

@@ -11,7 +11,7 @@ import {
   type IrType,
 } from "../../ir/ir.js";
 import { analyzeIntegerRanges } from "../../ir/integer-ranges.js";
-import { matchIntegerCountedForLoop } from "../../ir/integer-loops.js";
+import { matchIntegerCountedForLoop, matchIntegerPairedForLoop } from "../../ir/integer-loops.js";
 import { findInitializedByteLoopBindings } from "./byte-windows.js";
 import { matchBytesLoopVersion, type BytesLoopVersion } from "./bytes-loop-versions.js";
 import { emitLlvmModule } from "./emitter.js";
@@ -623,3 +623,78 @@ test("while loop counters below a string's length version its reads", () => {
   expect(version.bounded.map((c) => c.slack)).toEqual([0]);
   expect(version.fixed).toEqual([]);
 });
+
+test("a two-counter loop versions the typed arrays it indexes at either counter", () => {
+  // for (let i = 0, j = k; i < j; i++, j--) { const t = out[i]; out[i] = out[j]; out[j] = t; }
+  const loop: IrStmt & { kind: "for" } = {
+    kind: "for",
+    init: {
+      kind: "block",
+      body: [
+        { kind: "varDecl", localId: "i", init: num(0), loc },
+        { kind: "varDecl", localId: "j", init: ref("k"), loc },
+      ],
+      loc,
+    },
+    cond: { kind: "bin", op: "<", left: ref("i"), right: ref("j"), type: BOOL, loc },
+    update: {
+      kind: "block",
+      body: [
+        { kind: "assign", localId: "i", value: bin("+", ref("i"), num(1)), loc },
+        { kind: "assign", localId: "j", value: bin("-", ref("j"), num(1)), loc },
+      ],
+      loc,
+    },
+    body: [
+      { kind: "varDecl", localId: "t", init: get(ref("out", f64s), ref("i")), loc },
+      set(ref("out", f64s), ref("i"), get(ref("out", f64s), ref("j"))),
+      set(ref("out", f64s), ref("j"), ref("t")),
+      // Another index stays checked.
+      set(ref("src", f64s), bin("+", ref("i"), num(1)), num(0)),
+    ],
+    loc,
+  };
+  const kernel = fn([loop], [local("t", F64)]);
+  const ranges = analyzeIntegerRanges(kernel);
+  const locals = new Map(kernel.locals.map((l) => [l.id, l]));
+  const paired = matchIntegerPairedForLoop(loop, locals, ranges)!;
+  expect(paired).not.toBeNull();
+  const version = matchBytesLoopVersion(
+    loop,
+    null,
+    {
+      locals,
+      invariantGlobals: new Set(),
+      captured: new Set(),
+      ranges,
+      initialized: findInitializedByteLoopBindings(kernel).get(loop) ?? new Set(),
+      proven: new Set(),
+      workers: false,
+      owned: new Set(),
+      scalarStrings: new Set(),
+    },
+    paired,
+  )!;
+  expect(version.paired?.loop).toBe(paired);
+  expect(version.paired?.receivers.map((r) => version.receivers[r]!.ref)).toEqual([
+    ref("out", f64s),
+  ]);
+  expect(version.bounds.size).toBe(4);
+  const mod: IrModule = {
+    irVersion: 17,
+    sourceFile: loc.file,
+    entry: "__main",
+    functions: [
+      { name: "__main", params: [], returnType: VOID, locals: [], body: [], loc },
+      kernel,
+    ],
+  };
+  const body = /^define internal [^\n]*@sc_bf_kernel[^]*?^}/m.exec(emitLlvmModule(mod))![0];
+  const guard = body.slice(0, body.search(/^bytes\.version\.(fast|slow)\d+:$/m));
+  expect(guard).toMatch(/icmp sge i64 %\S+, 0/);
+  expect(guard).toMatch(/icmp slt i64 %\S+, %\S+/);
+  const fast = /^bytes\.version\.fast\d+:$[^]*?^bytes\.version\.done\d+:$/m.exec(body)![0];
+  // The fast copy keeps only the `src[i + 1]` check, and no second guard.
+  expect(fast.match(/bytes\.index\.invalid\d+:/g)).toHaveLength(1);
+  expect(fast).not.toContain("counted.fast");
+});

packages/compiler/src/backend/llvm/bytes-loop-versions.ts

@@ -1,11 +1,15 @@
 import { F64, type IrExpr, type IrLocal, type IrStmt } from "../../ir/ir.js";
 import type { IntegerRange, IntegerRanges } from "../../ir/integer-ranges.js";
-import type { IntegerCountedForLoop } from "../../ir/integer-loops.js";
+import type { IntegerCountedForLoop, IntegerPairedForLoop } from "../../ir/integer-loops.js";
 import type { LlvmEmitterContext } from "./expr-context.js";
 import { byteNumberAccess } from "../../ir/byte-numbers.js";
 import { everyExprChild, everyStmtChild } from "../../ir/traverse.js";
 import { exactInteger, widenInteger } from "./integer-values.js";
-import { emitCountedLoopLimit, emitCountedLoopSafety } from "./counted-loops.js";
+import {
+  emitCountedLoopLimit,
+  emitCountedLoopSafety,
+  emitExactIntegerTest,
+} from "./counted-loops.js";
 import { emitBytesHeaderLength, emitBytesSharedStorage } from "./expr-bytes.js";
 import { f64Lit } from "./common.js";
 import { emitStringAsciiFlag, emitStringByteLength } from "./string-reads.js";
@@ -79,6 +83,10 @@ export interface BytesLoopVersion {
   guarded: ReadonlySet<string>;
   /** The counted loop whose iteration space the affine checks cover. */
   counted: IntegerCountedForLoop | null;
+  /** A two-counter loop (`for (let i = lo, j = hi; i < j; i++, j--)`) and
+   * the receivers it indexes at either counter. Both counters stay between
+   * the two starts, so the guard requires `lo >= 0` and `hi < length`. */
+  paired: { loop: IntegerPairedForLoop; receivers: number[] } | null;
   /** The guard also requires every counter value within [-bound, bound]. */
   counterBound: number | null;
   /** Integer ranges of affine index operands that hold in the unchecked
@@ -153,6 +161,7 @@ export function matchBytesLoopVersion(
   loop: IrStmt & { kind: "for" | "while" },
   counted: IntegerCountedForLoop | null,
   facts: BytesLoopFacts,
+  paired: IntegerPairedForLoop | null = null,
 ): BytesLoopVersion | null {
   const written = new Set<string>();
   const temps = new Map<string, IrExpr>();
@@ -438,6 +447,9 @@ export function matchBytesLoopVersion(
     string,
     { check: BoundedCheck; nodes: { node: Node; access: WhileAccess }[] }
   >();
+  const pairedCounters =
+    paired !== null && counted === null ? new Set([paired.low.localId, paired.high.localId]) : null;
+  const pairedNodes = new Map<number, Node[]>();
   const access = (node: Node, array: IrExpr, index: IrExpr, width: number): void => {
     const ref = root(array);
     if (!ref) return;
@@ -448,6 +460,20 @@ export function matchBytesLoopVersion(
       else owned.set(receiver, [node]);
     }
     if (facts.proven.has(node)) return;
+    // An element at either counter of a two-counter loop.
+    if (
+      pairedCounters !== null &&
+      width === 1 &&
+      ref.type.kind === "bytes" &&
+      index.kind === "varRef" &&
+      pairedCounters.has(index.localId)
+    ) {
+      const receiver = receiverOf(ref);
+      const nodes = pairedNodes.get(receiver);
+      if (nodes) nodes.push(node);
+      else pairedNodes.set(receiver, [node]);
+      return;
+    }
     const shape = counter !== null ? affine(index) : null;
     if (shape && shape.counter !== 0) {
       const operands = new Set<IrExpr>();
@@ -543,11 +569,16 @@ export function matchBytesLoopVersion(
     Math.max(0, MAX_CHECKS - fixedKept.length - affineKept.length),
   );
   const ownedKept = [...owned.entries()].slice(0, MAX_CHECKS);
+  const pairedKept = [...pairedNodes.entries()].slice(
+    0,
+    Math.max(0, MAX_CHECKS - fixedKept.length - affineKept.length - boundedKept.length),
+  );
   if (
     fixedKept.length === 0 &&
     affineKept.length === 0 &&
     boundedKept.length === 0 &&
-    ownedKept.length === 0
+    ownedKept.length === 0 &&
+    pairedKept.length === 0
   )
     return null;
   const bounds = new Set<Node>();
@@ -616,6 +647,10 @@ export function matchBytesLoopVersion(
       if (access.index !== access.ref) narrow(access.index, { min: 0, max: SAFE });
     }
   }
+  for (const [receiver, nodes] of pairedKept) {
+    used.add(receiver);
+    for (const node of nodes) bounds.add(node);
+  }
   const counterChecks: CounterCheck[] = [];
   for (const [id, min] of counterMins) {
     const info = whileAccesses!.counters.get(id)!;
@@ -658,6 +693,10 @@ export function matchBytesLoopVersion(
     counters: counterChecks,
     guarded,
     counted: affineKept.length > 0 ? counted : null,
+    paired:
+      paired !== null && pairedKept.length > 0
+        ? { loop: paired, receivers: pairedKept.map(([receiver]) => renumber.get(receiver)!) }
+        : null,
     counterBound: needsBound ? bound : null,
     ranges,
   };
@@ -1143,6 +1182,25 @@ export function emitBytesLoopVersionGuard(
       conditions.push(both);
     }
   }
+  if (version.paired !== null) {
+    // Both counters stay within [lo, hi] (exact integers below 2^53): every
+    // access at either one is in bounds when `lo >= 0` and `hi < length`.
+    const { loop } = version.paired;
+    const [low, high] = [loop.low.start, loop.high.start].map((expr) => {
+      const tested = emitExactIntegerTest(host, host.emitExpr(expr).name, -SAFE, SAFE);
+      conditions.push(tested.ok);
+      return tested.integer;
+    });
+    const nonNegative = B.tmp();
+    B.line(`${nonNegative} = icmp sge i64 ${low}, 0`);
+    conditions.push(nonNegative);
+    for (const receiver of version.paired.receivers) {
+      // Lengths are at most 2^53 elements: a signed comparison is exact.
+      const below = B.tmp();
+      B.line(`${below} = icmp slt i64 ${high}, ${lengths[receiver]}`);
+      conditions.push(below);
+    }
+  }
   let valid = "true";
   for (const cond of conditions) {
     if (valid === "true") valid = cond;
@@ -1156,39 +1214,6 @@ export function emitBytesLoopVersionGuard(
   else B.condBr(valid, fast, slow);
 }
 
-/** Tests that a double is an exact integer within [min, max] (not -0; NaN
- * fails the ordered comparisons) and gives its i64. The integer is frozen,
- * so it is a fixed value even when the test fails. */
-function emitExactIntegerTest(
-  host: LlvmEmitterContext,
-  value: string,
-  min: number,
-  max: number,
-): { ok: string; integer: string } {
-  const B = host.B;
-  const low = B.tmp(),
-    high = B.tmp(),
-    inRange = B.tmp(),
-    raw = B.tmp(),
-    integer = B.tmp(),
-    back = B.tmp(),
-    bits = B.tmp(),
-    backBits = B.tmp(),
-    exact = B.tmp(),
-    ok = B.tmp();
-  B.line(`${low} = fcmp oge double ${value}, ${f64Lit(min)}`);
-  B.line(`${high} = fcmp ole double ${value}, ${f64Lit(max)}`);
-  B.line(`${inRange} = and i1 ${low}, ${high}`);
-  B.line(`${raw} = fptosi double ${value} to i64`);
-  B.line(`${integer} = freeze i64 ${raw}`);
-  B.line(`${back} = sitofp i64 ${integer} to double`);
-  B.line(`${bits} = bitcast double ${value} 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}`);
-  return { ok, integer };
-}
-
 function scaleCounter(host: LlvmEmitterContext, value: string, scale: number): string {
   if (scale === 1) return value;
   const scaled = host.B.tmp();

packages/compiler/src/backend/llvm/counted-loops.ts

@@ -151,3 +151,36 @@ function emitExactStart(host: LlvmEmitterContext, expr: IrExpr): string {
   B.line(`${ok} = and i1 ${inRange}, ${exact}`);
   return ok;
 }
+
+/** Tests that a double is an exact integer within [min, max] (not -0; NaN
+ * fails the ordered comparisons) and gives its i64. The integer is frozen,
+ * so it is a fixed value even when the test fails. */
+export function emitExactIntegerTest(
+  host: LlvmEmitterContext,
+  value: string,
+  min: number,
+  max: number,
+): { ok: string; integer: string } {
+  const B = host.B;
+  const low = B.tmp(),
+    high = B.tmp(),
+    inRange = B.tmp(),
+    raw = B.tmp(),
+    integer = B.tmp(),
+    back = B.tmp(),
+    bits = B.tmp(),
+    backBits = B.tmp(),
+    exact = B.tmp(),
+    ok = B.tmp();
+  B.line(`${low} = fcmp oge double ${value}, ${f64Lit(min)}`);
+  B.line(`${high} = fcmp ole double ${value}, ${f64Lit(max)}`);
+  B.line(`${inRange} = and i1 ${low}, ${high}`);
+  B.line(`${raw} = fptosi double ${value} to i64`);
+  B.line(`${integer} = freeze i64 ${raw}`);
+  B.line(`${back} = sitofp i64 ${integer} to double`);
+  B.line(`${bits} = bitcast double ${value} 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}`);
+  return { ok, integer };
+}

packages/compiler/src/backend/llvm/dense-array-loops.ts

@@ -36,8 +36,7 @@ import { everyExprChild, everyStmtChild } from "../../ir/traverse.js";
 import type { IntegerCountedForLoop, IntegerPairedForLoop } from "../../ir/integer-loops.js";
 import type { LlvmEmitterContext } from "./expr-context.js";
 import { loadArrCount, SCR_ARR_FIELD } from "./common.js";
-import { emitCountedLoopLimit } from "./counted-loops.js";
-import { f64Lit } from "./common.js";
+import { emitCountedLoopLimit, emitExactIntegerTest } from "./counted-loops.js";
 
 export const INT32_MAX = 2147483647;
 
@@ -453,28 +452,10 @@ function emitPairedDenseGuard(
   header: (array: string, field: number) => string,
 ): string {
   const B = host.B;
-  const integer = (expr: IrExpr, min: number, max: number): { value: string; ok: string } => {
-    const start = host.emitExpr(expr);
-    const low = B.tmp(),
-      high = B.tmp(),
-      raw = B.tmp(),
-      value = B.tmp(),
-      back = B.tmp(),
-      exact = B.tmp(),
-      inRange = B.tmp(),
-      ok = B.tmp();
-    B.line(`${low} = fcmp oge double ${start.name}, ${f64Lit(min)}`);
-    B.line(`${high} = fcmp ole double ${start.name}, ${f64Lit(max)}`);
-    B.line(`${inRange} = and i1 ${low}, ${high}`);
-    B.line(`${raw} = fptosi double ${start.name} to i64`);
-    B.line(`${value} = freeze i64 ${raw}`);
-    B.line(`${back} = sitofp i64 ${value} to double`);
-    B.line(`${exact} = fcmp oeq double ${back}, ${start.name}`);
-    B.line(`${ok} = and i1 ${inRange}, ${exact}`);
-    return { value, ok };
-  };
-  const low = integer(pair.lowStart, 0, INT32_MAX - 1);
-  const high = integer(pair.highStart, -1, INT32_MAX - 1);
+  const start = (expr: IrExpr, min: number, max: number) =>
+    emitExactIntegerTest(host, host.emitExpr(expr).name, min, max);
+  const low = start(pair.lowStart, 0, INT32_MAX - 1);
+  const high = start(pair.highStart, -1, INT32_MAX - 1);
   let valid = B.tmp();
   B.line(`${valid} = and i1 ${low.ok}, ${high.ok}`);
   for (const ref of loop.arrays.values()) {
@@ -486,7 +467,7 @@ function emitPairedDenseGuard(
       both = B.tmp();
     B.line(`${filled} = load i32, ptr ${filledPtr}${host.fieldAliasAttachment(filledPtr)}`);
     B.line(`${wide} = zext i32 ${filled} to i64`);
-    B.line(`${covered} = icmp slt i64 ${high.value}, ${wide}`);
+    B.line(`${covered} = icmp slt i64 ${high.integer}, ${wide}`);
     B.line(`${both} = and i1 ${valid}, ${covered}`);
     valid = both;
   }

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

@@ -1051,15 +1051,17 @@ export function emitStmtBody(emitter: LlEmitter, s: IrStmt): void {
               LOOPS.of(emitter).unaliasedArrays,
             )
           : countedLoop;
-      const version = versionable ? emitter.matchBytesLoopVersion(s, versionCounted) : null;
+      const version = versionable
+        ? emitter.matchBytesLoopVersion(s, versionCounted, versionCounted ? null : pairedLoop)
+        : null;
       if (version) {
         emitVersionedLoop(emitter, s, version);
         break;
       }
       let countedJoin: string | null = null;
       if (
-        (countedLoop?.guarded && !LOOPS.of(emitter).provenCountedGuards.has(s)) ||
-        pairedLoop?.guarded
+        (countedLoop?.guarded || pairedLoop?.guarded) &&
+        !LOOPS.of(emitter).provenCountedGuards.has(s)
       ) {
         const safe = countedLoop
           ? emitCountedLoopSafety(emitter, countedLoop)
@@ -2094,7 +2096,8 @@ function emitVersionedLoop(
     if (version.owned.size > 0) emitter.bytesOwned = new Set([...originalOwned, ...version.owned]);
     if (version.ranges.size > 0)
       emitter.integerRanges = new Map([...originalRanges, ...version.ranges]);
-    if (version.counted?.guarded) LOOPS.of(emitter).provenCountedGuards.add(s);
+    if (version.counted?.guarded || version.paired?.loop.guarded)
+      LOOPS.of(emitter).provenCountedGuards.add(s);
     emitter.emitStmt(s);
     B.br(done);
   } finally {

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

@@ -60,7 +60,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 { IntegerCountedForLoop } from "../../ir/integer-loops.js";
+import type { IntegerCountedForLoop, IntegerPairedForLoop } from "../../ir/integer-loops.js";
 import {
   materializeSplitPiece,
   type SplitSpan,
@@ -4263,6 +4263,7 @@ export class LlEmitter {
   matchBytesLoopVersion(
     s: IrStmt & { kind: "for" | "while" },
     counted: IntegerCountedForLoop | null,
+    paired: IntegerPairedForLoop | null = null,
   ): ReturnType<typeof matchBytesLoopVersion> {
     const fn = LOOPS.of(this).bytesLoopFunction;
     if (!fn) return null;
@@ -4272,17 +4273,22 @@ export class LlEmitter {
     const invariantGlobals = new Set(this.stableGlobals);
     for (const global of this.initializerBindings.get(fn.name) ?? [])
       invariantGlobals.add(global.id);
-    return matchBytesLoopVersion(s, counted, {
-      locals: this.currentLocals,
-      invariantGlobals,
-      captured: this.captureIds,
-      ranges: this.integerRanges,
-      initialized,
-      proven: this.bytesBounds,
-      workers: this.mod.workers === true,
-      owned: this.bytesOwned,
-      scalarStrings: new Set(this.stringSlices.keys()),
-    });
+    return matchBytesLoopVersion(
+      s,
+      counted,
+      {
+        locals: this.currentLocals,
+        invariantGlobals,
+        captured: this.captureIds,
+        ranges: this.integerRanges,
+        initialized,
+        proven: this.bytesBounds,
+        workers: this.mod.workers === true,
+        owned: this.bytesOwned,
+        scalarStrings: new Set(this.stringSlices.keys()),
+      },
+      paired,
+    );
   }
 
   /** The alias tag of a whole-slot load of a binding (module globals only). */

packages/compiler/test/ts7/baselines/order-parity/tests/corpus/7981-typed-array-paired-loops.ts.json

@@ -0,0 +1,6 @@
+{
+ "order": [
+  "<repo>/tests/corpus/7981-typed-array-paired-loops.ts"
+ ],
+ "diags": []
+}

tests/corpus/7981-typed-array-paired-loops.ts

@@ -0,0 +1,57 @@
+// Two-counter loops over typed arrays (`for (let i = lo, j = hi; i < j;
+// i++, j--)`), with starts inside the arrays and outside them: past the
+// end, negative, fractional and NaN. Views over one buffer, other element
+// types, inclusive conditions and steps other than one.
+
+function reverse(a: Int32Array, lo: number, hi: number): string {
+  for (let i = lo, j = hi; i < j; i++, j--) {
+    const t = a[i]!;
+    a[i] = a[j]!;
+    a[j] = t;
+  }
+  return Array.from(a).join();
+}
+
+function mirror(a: Float64Array, b: Float64Array, hi: number): number {
+  let sum = 0;
+  for (let i = 0, j = hi; i <= j; i++, j--) {
+    b[i] = a[j]! * 2;
+    sum += a[i]! - b[j]!;
+  }
+  return sum;
+}
+
+function skipping(a: Uint8ClampedArray, hi: number): string {
+  for (let i = 0, j = hi; i < j; i += 2, j -= 3) {
+    a[i] = a[j]! + 200;
+    a[j] = i - 10;
+  }
+  return Array.from(a).join();
+}
+
+function palindrome(text: Uint16Array, n: number): boolean {
+  for (let i = 0, j = n - 1; i < j; i++, j--) if (text[i] !== text[j]) return false;
+  return true;
+}
+
+console.log(reverse(new Int32Array([1, 2, 3, 4, 5, 6]), 0, 5));
+console.log(reverse(new Int32Array([1, 2, 3, 4, 5, 6]), 1, 3));
+console.log(reverse(new Int32Array([1, 2, 3, 4]), 0, 6));
+console.log(reverse(new Int32Array([1, 2, 3, 4]), -2, 2));
+console.log(reverse(new Int32Array([1, 2, 3, 4]), 0.5, 3.5));
+console.log(reverse(new Int32Array([1, 2, 3, 4]), 0, NaN));
+console.log(reverse(new Int32Array([1, 2, 3, 4]), 3, 0));
+
+const buffer = new ArrayBuffer(48);
+const whole = new Float64Array(buffer);
+for (let k = 0; k < whole.length; k++) whole[k] = k + 0.5;
+const tail = new Float64Array(buffer, 16, 3);
+console.log(mirror(whole, tail, 4), Array.from(whole).join());
+console.log(mirror(tail, tail, 2), Array.from(whole).join());
+console.log(mirror(tail, whole, 7), Array.from(whole).join());
+
+console.log(skipping(new Uint8ClampedArray([10, 20, 30, 40, 50, 60, 70, 80]), 7));
+console.log(skipping(new Uint8ClampedArray([10, 20, 30]), 9));
+
+const word = new Uint16Array([114, 97, 99, 101, 99, 97, 114]);
+console.log(palindrome(word, 7), palindrome(word, 6), palindrome(word, 9), palindrome(word, 0));

tests/harness/corpus-costs/7981-typed-array-paired-loops.ts.json

@@ -0,0 +1 @@
+{ "plain": 3500, "san": 4500 }