Version typed-array loops that contain nested loops

An outer loop's statically ranged accesses (the sieve's flags[i]) are
proved once per entry even when its body holds an inner loop. Versions
nest at most three deep and never inside their own copies, and fixed
checks only cover lengths a buffer can have, so an outer loop never pays
for a guard on an inner loop's unbounded counter.

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

@@ -64,6 +64,7 @@ export interface BytesLoopFacts {
 }
 
 const MAX_COST = 400;
+const MAX_FIXED_END = 2 ** 32;
 const MAX_CHECKS = 8;
 const MAX_OFFSET = 2 ** 31;
 const SAFE = Number.MAX_SAFE_INTEGER;
@@ -81,9 +82,10 @@ const INVARIANT_OPS = new Set(["+", "-", "*", "&", "|", "^", "<<", ">>", ">>>"])
  *   and last values in bounds, such as `out[base + i]` for `i < n` or
  *   `view.getUint32(i * 4)`.
  *
- * Only innermost synchronous loops are versioned, so the body is emitted
- * twice at most. The checked loop runs whenever the guard fails, and keeps
- * every out-of-bounds read, ignored write and NaN-index behavior. */
+ * Only small synchronous loops are versioned (nested loops included), and
+ * the emitter bounds how deeply versions nest. The checked loop runs
+ * whenever the guard fails, and keeps every out-of-bounds read, ignored
+ * write and NaN-index behavior. */
 export function matchBytesLoopVersion(
   loop: IrStmt & { kind: "for" },
   counted: IntegerCountedForLoop | null,
@@ -92,7 +94,6 @@ export function matchBytesLoopVersion(
   const written = new Set<string>();
   const temps = new Map<string, IrExpr>();
   let cost = 0;
-  let innermost = true;
   const scanExpr = (e: IrExpr): boolean => {
     if (++cost > MAX_COST) return false;
     if (e.kind === "assignExpr" || e.kind === "incDec") written.add(e.localId);
@@ -101,12 +102,9 @@ export function matchBytesLoopVersion(
   const scanStmt = (s: IrStmt): boolean => {
     if (++cost > MAX_COST) return false;
     switch (s.kind) {
-      case "for":
       case "forOf":
-      case "while":
-      case "doWhile":
-        innermost = false;
-        return false;
+        written.add(s.localId);
+        break;
       case "varDecl": {
         if (temps.has(s.localId) || written.has(s.localId)) temps.delete(s.localId);
         else if (s.init) temps.set(s.localId, s.init);
@@ -126,7 +124,6 @@ export function matchBytesLoopVersion(
   };
   if (
     !loop.body.every(scanStmt) ||
-    !innermost ||
     (loop.cond !== null && !scanExpr(loop.cond)) ||
     (loop.update !== null && !scanStmt(loop.update))
   )
@@ -339,7 +336,9 @@ export function matchBytesLoopVersion(
     const range = facts.ranges.get(index);
     if (!range) return;
     const end = range.max + width - 1;
-    if (range.min < 0 || !Number.isSafeInteger(end)) return;
+    // Only a length a buffer can have: an outer loop must not pay for a
+    // guard that bounds an inner loop's unbounded counter.
+    if (range.min < 0 || end >= MAX_FIXED_END) return;
     const receiver = receiverOf(ref);
     const group = fixed.get(receiver);
     if (group) {

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

@@ -620,6 +620,9 @@ function mentionsLocal(e: IrExpr, localId: string): boolean {
   return found;
 }
 
+/** Typed-array loop versions nest at most this deep (each doubles a body). */
+const MAX_BYTES_LOOP_VERSION_DEPTH = 3;
+
 export class LlEmitter {
   localArrayReads = new Map<string, LocalArrayRead>();
   private loopArrayBorrows: ReadonlySet<IrStmt> = new Set();
@@ -1071,6 +1074,9 @@ export class LlEmitter {
   /** Typed-array loop versioning (bytes-loop-versions.ts) for the current
    * function: off in debug and suspending functions and inside a copy. */
   private bytesLoopVersionsEnabled = false;
+  /** Loops being versioned around the statement being emitted: a loop is
+   * never versioned inside its own copies, and versions nest boundedly. */
+  private readonly bytesLoopVersioning = new Set<IrStmt>();
   private bytesLoopEntries: ReturnType<typeof findInitializedByteLoopBindings> | null = null;
   private bytesLoopFunction: IrFunction | null = null;
   /** Counted loops whose integer-form safety test an enclosing version
@@ -8712,9 +8718,12 @@ export class LlEmitter {
           B.startBlock(done);
           break;
         }
-        const version = this.bytesLoopVersionsEnabled
-          ? this.matchBytesLoopVersion(s, countedLoop)
-          : null;
+        const version =
+          this.bytesLoopVersionsEnabled &&
+          this.bytesLoopVersioning.size < MAX_BYTES_LOOP_VERSION_DEPTH &&
+          !this.bytesLoopVersioning.has(s)
+            ? this.matchBytesLoopVersion(s, countedLoop)
+            : null;
         if (version) {
           const fast = B.newLabel("bytes.version.fast"),
             slow = B.newLabel("bytes.version.slow"),
@@ -8722,7 +8731,7 @@ export class LlEmitter {
           emitBytesLoopVersionGuard(this, version, fast, slow);
           const originalBounds = this.bytesBounds,
             originalRanges = this.integerRanges;
-          this.bytesLoopVersionsEnabled = false;
+          this.bytesLoopVersioning.add(s);
           try {
             B.startBlock(slow);
             this.emitStmt(s);
@@ -8738,7 +8747,7 @@ export class LlEmitter {
             this.bytesBounds = originalBounds;
             this.integerRanges = originalRanges;
             this.provenCountedGuards.delete(s);
-            this.bytesLoopVersionsEnabled = true;
+            this.bytesLoopVersioning.delete(s);
           }
           B.startBlock(done);
           break;

tests/corpus/5467-typed-array-loop-versions.ts

@@ -216,6 +216,21 @@ try {
   console.log((e as Error).message, Array.from(u32).join(","));
 }
 
+// An outer loop proved in bounds around an inner loop that is not.
+function nested(grid: Int32Array, w: number, h: number): number {
+  let sum = 0;
+  for (let y = 0; y < h; y++) {
+    sum += grid[y & 7];
+    for (let x = 0; x < w; x++) {
+      grid[y * w + x] = grid[y * w + x] + 1;
+      sum += grid[y * w + x] | 0;
+    }
+  }
+  return sum;
+}
+const grid = new Int32Array(16);
+console.log("nested", nested(grid, 4, 4), nested(grid, 5, 4), nested(grid, 4, 5), grid.join(","));
+
 // A matrix product over module-level constant arrays.
 const N = 6;
 const ma: Float64Array = new Float64Array(N * N);