Address review feedback

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

@@ -2818,7 +2818,7 @@ function emitStripMinedCountedLoop(
   // iteration: its plain load or store is not merged with another
   // iteration's or moved out of the loop.
   const outer = emitter.plainSharedLoop;
-  emitter.plainSharedLoop = { counter: counter.localId, target };
+  emitter.plainSharedLoop = { counter: counter.localId, target, depth: emitter.loopDepth };
   try {
     emitter.emitBlock(s.body);
   } finally {

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

@@ -1824,11 +1824,15 @@ export class LlEmitter {
    * guard proved thread-owned. */
   bytesOwned: ReadonlySet<IrExpr | IrStmt> = new Set();
   /** Worker programs: the strip-mined counted loop whose body is being
-   * emitted (emitStripMinedCountedLoop), with its counter and its jump
-   * target. An element access of a view that may be shared, at the counter
-   * and directly in that loop's body, is a plain load or store
+   * emitted (emitStripMinedCountedLoop), with its counter, its jump target
+   * and its loopDepth. An element access of a view that may be shared, at
+   * the counter and directly in that loop's body, is a plain load or store
    * (plainSharedElement in expr-bytes.ts). */
-  plainSharedLoop: { counter: string; target: LlEmitter["jumpTargets"][number] } | null = null;
+  plainSharedLoop: {
+    counter: string;
+    target: LlEmitter["jumpTargets"][number];
+    depth: number;
+  } | null = null;
   /** The alias-scope attachment of each typed-array element read and write
    * in the current loop version's fast copy (emitVersionedLoop), and the
    * one the element access being emitted carries (bytesElementAlias). */

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

@@ -331,10 +331,11 @@ function emitBytesLoad(
 /** Worker programs: whether an element access at `index` of a view that
  * may be shared is a plain load or store, like a thread-owned view's. It is
  * when the index is the counter of the strip-mined counted loop whose body
- * holds the access, outside any loop nested in that body: every iteration
- * then reaches a different element, so LLVM cannot merge an element's reads
- * across iterations or move them out of the loop, and the loop can be
- * vectorized. A racing write from another thread can still be seen or
+ * holds the access, outside any loop nested in that body (its condition and
+ * update included: a nested loop that waits on the element reads it on
+ * every iteration, so the read stays atomic): every iteration then reaches
+ * a different element, so LLVM cannot merge an element's reads across
+ * iterations or move them out of the loop, and the loop can be vectorized. A racing write from another thread can still be seen or
  * missed within one iteration, the documented gap of non-atomic shared
  * accesses. Float16 elements keep their runtime path. */
 function plainSharedElement(
@@ -344,7 +345,8 @@ function plainSharedElement(
 ): boolean {
   const loop = host.plainSharedLoop;
   if (loop === null || elem === "f16" || index?.kind !== "varRef") return false;
-  if (index.localId !== loop.counter) return false;
+  // A nested loop raises loopDepth before its condition is emitted.
+  if (index.localId !== loop.counter || host.loopDepth !== loop.depth) return false;
   let innermost: (typeof host.jumpTargets)[number] | undefined;
   for (let i = host.jumpTargets.length - 1; i >= 0 && innermost === undefined; i--)
     if (host.jumpTargets[i]!.kind === "loop") innermost = host.jumpTargets[i];

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

@@ -267,6 +267,17 @@ test("worker programs read and write shared elements at a strip-mined loop's cou
         "  for (let i = 0; i < n; i++) for (let j = 0; j < n; j++) t += a[i]! + a[j]!;",
         "  return t;",
         "}",
+        "export function waits(a: Int32Array, out: Uint8Array, n: number): number {",
+        "  let spins = 0;",
+        "  for (let i = 0; i < n; i++) {",
+        "    out[i] = 1;",
+        "    while (a[i] === 0) spins++;",
+        "    do spins++; while (a[i] === 1);",
+        "    for (let k = 0; a[i]! > k; k++) spins++;",
+        "    for (let k = 0; k < 3; k += a[i]!) spins++;",
+        "  }",
+        "  return spins;",
+        "}",
         "export function spin(a: Int32Array): number {",
         "  let n = 0;",
         "  while (a[0] === 0 && n < 10) n++;",
@@ -279,7 +290,7 @@ test("worker programs read and write shared elements at a strip-mined loop's cou
         "codes(c);",
         "shift(c);",
         "const d = new Int32Array(new SharedArrayBuffer(16)).fill(1);",
-        "console.log(b[7], c[0], rows(d, 4), spin(d));",
+        "console.log(b[7], c[0], rows(d, 4), spin(d), waits(new Int32Array(new SharedArrayBuffer(16)).fill(2), new Uint8Array(new SharedArrayBuffer(4)), 4));",
         "if (!isMainThread) new Worker(new URL(import.meta.url));",
       ].join("\n"),
     );
@@ -319,6 +330,15 @@ test("worker programs read and write shared elements at a strip-mined loop's cou
     expect(rows).toMatch(/load atomic i32, ptr %/);
     // A loop that is not counted keeps atomic reads.
     expect(bodies("spin")).toMatch(/load atomic i32, ptr %/);
+    // A nested loop that tests the outer counter's element in its condition
+    // or steps by it reads it again on every iteration: each read stays
+    // atomic, so another thread's write ends the wait.
+    const waits = bodies("waits");
+    expect(waits).toMatch(/^loop\.poll\d+:$/m);
+    expect(waits.match(/load atomic i32, ptr %/g)?.length).toBeGreaterThanOrEqual(4);
+    // The outer loop's write at its counter stays plain (speculation counters
+    // are atomic stores to globals).
+    expect(waits).not.toMatch(/store atomic i8 \S+, ptr %/);
   } finally {
     await rm(dir, { recursive: true, force: true });
   }