From ef6f8838d0ebda35515d4651165e55527ef0286b Mon Sep 17 00:00:00 2001 From: Chris Tate Date: Sun, 25 Oct 2026 23:53:54 -0500 Subject: [PATCH] Address review feedback --- .../src/backend/llvm/emit-statements.ts | 2 +- packages/compiler/src/backend/llvm/emitter.ts | 12 ++++++---- .../compiler/src/backend/llvm/expr-bytes.ts | 12 +++++----- .../test/counted-loop-emission.test.ts | 22 ++++++++++++++++++- 4 files changed, 37 insertions(+), 11 deletions(-) diff --git a/packages/compiler/src/backend/llvm/emit-statements.ts b/packages/compiler/src/backend/llvm/emit-statements.ts index 310c5abb4f..66bbf34d0b 100644 --- a/packages/compiler/src/backend/llvm/emit-statements.ts +++ b/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 { diff --git a/packages/compiler/src/backend/llvm/emitter.ts b/packages/compiler/src/backend/llvm/emitter.ts index 095515e314..6a058299e5 100644 --- a/packages/compiler/src/backend/llvm/emitter.ts +++ b/packages/compiler/src/backend/llvm/emitter.ts @@ -1824,11 +1824,15 @@ export class LlEmitter { * guard proved thread-owned. */ bytesOwned: ReadonlySet = 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). */ diff --git a/packages/compiler/src/backend/llvm/expr-bytes.ts b/packages/compiler/src/backend/llvm/expr-bytes.ts index 10363c19c5..2493a343d1 100644 --- a/packages/compiler/src/backend/llvm/expr-bytes.ts +++ b/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]; diff --git a/packages/compiler/test/counted-loop-emission.test.ts b/packages/compiler/test/counted-loop-emission.test.ts index fe48deeeb5..e9a79a71ab 100644 --- a/packages/compiler/test/counted-loop-emission.test.ts +++ b/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 }); }