Version two-counter loops over typed arrays

A loop with two counters closing in on each other, such as
`for (let i = lo, j = hi; i < j; i++, j--)`, checked the bounds of every
typed-array element it read or wrote at either counter. Within the body
the low counter only rises and stays below the high one, which only
falls, so every such index lies between the two starts.

Typed-array loop versions now cover these loops. The guard tests that
both starts are exact integers, that the low start is at least 0, and
that the high start is below the length of each view the loop indexes at
a counter. The fast copy then reads and writes those elements without
checks, and the checked copy runs whenever the guard fails. In the fast
copy the guard also stands for the loop's own start test.

The exact-integer test of a start is now one function, shared by the
typed-array version guards and the dense array two-counter guard. The
dense guard, which compared the round trip with `fcmp oeq`, now also
leaves a -0 start to the checked loop.

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 }