trending/scriptc
1
0
Fork 0
mirror of https://github.com/vercel-labs/scriptc.git synced 2026-11-02 07:43:42 +08:00
Code Issues Packages Projects Releases Wiki Activity

Read and write shared typed-array elements at a counted loop's counter with plain accesses (#2309)

Browse Source
This commit is contained in:
Chris Tate
2026-10-26 00:23:15 -05:00
committed by GitHub
parent b090950ffe ef6f8838d0
commit e7b60a5302
9 changed files with 301 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/compiler/src/backend/llvm/bytes-loop-versions.test.ts
+5 -4
View File
@@ -320,16 +320,17 @@ test("worker programs test each receiver's shared storage once in the guard", ()
const fast = /^bytes\.version\.fast\d+:$[^]*?^bytes\.version\.fast\.shared\d+:$/m.exec(body)![0];
expect(fast).not.toContain("bytes.shared");
expect(fast).not.toMatch(/(load|store) atomic double/);
// The checked copy keeps the shared-memory path for each access, and so
// does the copy for shared views whose bounds hold.
// The checked copy keeps the shared-memory path for each access. The copy
// for shared views whose bounds hold is a counted loop that polls between
// blocks, so its accesses at the counter are plain.
const slow = body.slice(body.search(/^bytes\.version\.slow\d+:$/m), body.indexOf(fast));
expect(slow).toContain("load atomic double");
expect(slow).toContain("store atomic double");
const shared = /^bytes\.version\.fast\.shared\d+:$[^]*?^bytes\.version\.done\d+:$/m.exec(
body,
)![0];
expect(shared).toContain("load atomic double");
expect(shared).toContain("store atomic double");
expect(shared).toMatch(/^loop\.poll\d+:$/m);
expect(shared).not.toMatch(/(load|store) atomic double/);
});
test("index temps of a version's counter are integers in the fast copy", () => {
packages/compiler/src/backend/llvm/emit-statements.ts
+20 -6
View File
@@ -1265,6 +1265,7 @@ export function emitStmtBody(emitter: LlEmitter, s: IrStmt): void {
if (stripMined) {
const counter: StripMinedCounter = countedLoop
? {
localId: countedLoop.localId,
slot: integerSlot!,
type: "i64",
signed: true,
@@ -1276,6 +1277,7 @@ export function emitStmtBody(emitter: LlEmitter, s: IrStmt): void {
limit: () => countedLimit!,
}
: {
localId: integerLoop!.localId,
slot: integerSlot!,
type: emitter.sizeType,
signed: false,
@@ -2705,10 +2707,11 @@ function boundedLoopRuns(emitter: LlEmitter, loop: IrStmt): number {
return trips * (1 + nested);
}
/** The counter of a strip-mined loop: its slot, the LLVM integer type and
* signedness of its values, its step, the counters stepped beside it, and
* the limit, emitted in the outer header before each block. */
/** The counter of a strip-mined loop: its binding, its slot, the LLVM
* integer type and signedness of its values, its step, the counters stepped
* beside it, and the limit, emitted in the outer header before each block. */
interface StripMinedCounter {
localId: string;
slot: string;
type: string;
signed: boolean;
@@ -2800,7 +2803,7 @@ function emitStripMinedCountedLoop(
const after = runSlot === null ? poll : B.newLabel("loop.end");
B.condBr(inBlock, lb, after);
B.startBlock(lb);
emitter.jumpTargets.push({
const target: LlEmitter["jumpTargets"][number] = {
kind: "loop",
brkLabel: exit,
contLabel: lu,
@@ -2808,8 +2811,19 @@ function emitStripMinedCountedLoop(
frameDepth: emitter.frames.length,
scopeDepth: emitter.scopes.length,
finallyDepth: emitter.finallyStack.length,
});
emitter.emitBlock(s.body);
};
emitter.jumpTargets.push(target);
// The counter steps by a nonzero constant and the body does not assign
// it, so an element at the counter is a different element on every
// 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, depth: emitter.loopDepth };
try {
emitter.emitBlock(s.body);
} finally {
emitter.plainSharedLoop = outer;
}
emitter.jumpTargets.pop();
B.br(lu);
B.startBlock(lu);
packages/compiler/src/backend/llvm/emitter.ts
+10
View File
@@ -1823,6 +1823,16 @@ export class LlEmitter {
/** Worker programs: typed-array accesses whose views an enclosing version
* 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, 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];
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
+27 -2
View File
@@ -317,7 +317,7 @@ function emitBytesLoad(
owned: boolean,
): LlValue {
// `owned`: an enclosing version guard proved the view thread-owned.
if (host.mod.workers && !owned) {
if (host.mod.workers && !owned && !plainSharedElement(host, elem, expr)) {
const idx = emitBytesIndex(host, receiver, index, expr, invalid, inBounds);
return emitSharedOrLocalLoad(host, elem, receiver, idx);
}
@@ -328,6 +328,31 @@ function emitBytesLoad(
return emitBytesLoadAt(host, elem, data, idx);
}
/** 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 (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(
host: LlvmEmitterContext,
elem: IrBytesElem,
index: IrExpr | undefined,
): boolean {
const loop = host.plainSharedLoop;
if (loop === null || elem === "f16" || index?.kind !== "varRef") 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];
return innermost === loop.target;
}
/** Worker programs: the element at a checked index of a view that may be
* shared. An integer element joins as its 32-bit value and converts to a
* number after the join, as the local load does, so compares and
@@ -666,7 +691,7 @@ function emitScopedBytesSet(
ty = "i32";
operand = stored!;
}
if (host.mod.workers && !owned) {
if (host.mod.workers && !owned && !plainSharedElement(host, elem, indexExpr)) {
// A shared view's element takes a relaxed atomic store of its width.
const local = emitSharedBytesSplit(host, receiver);
const p = B.tmp();
packages/compiler/test/counted-loop-emission.test.ts
+99
View File
@@ -244,3 +244,102 @@ test("worker programs poll no loop whose proven ranges bound its runs inside ano
await rm(dir, { recursive: true, force: true });
}
});
test("worker programs read and write shared elements at a strip-mined loop's counter with plain accesses", async () => {
const dir = await mkdtemp(join(tmpdir(), "scriptc-counted-shared-"));
try {
const entry = join(dir, "shared.ts");
await writeFile(
entry,
[
'import { isMainThread, Worker } from "node:worker_threads";',
"export function scale(a: Float64Array, b: Float64Array, n: number): void {",
" for (let i = 0; i < n; i++) b[i] = a[i]! * 2;",
"}",
"export function codes(a: Uint8Array): void {",
" for (let i = 0; i < a.length; ++i) a[i] = a[i] === 97 ? 0 : 3;",
"}",
"export function shift(a: Uint8Array): void {",
" for (let i = 0; i < a.length - 1; i++) a[i] = a[i + 1]!;",
"}",
"export function rows(a: Int32Array, n: number): number {",
" let t = 0;",
" 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++;",
" return n;",
"}",
"const a = new Float64Array(new SharedArrayBuffer(64)).fill(1);",
"const b = new Float64Array(8);",
"scale(a, b, 8);",
"const c = new Uint8Array(new SharedArrayBuffer(8)).fill(97);",
"codes(c);",
"shift(c);",
"const d = new Int32Array(new SharedArrayBuffer(16)).fill(1);",
"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"),
);
const outPath = join(dir, "shared.ir.json");
const result = await compile(entry, { outPath, outDir: dir, outputKind: "ir" });
if (!result.ok) throw new Error(result.diagnostics.map((d) => d.message).join("\n"));
const ll = emitLlvmModule(deserializeModule(await readFile(outPath, "utf8")));
const bodies = (name: string): string =>
[
...ll.matchAll(
new RegExp(
`define internal [^\\n]+ @sc_b?f_${name}(?:\\.i32)?\\([^\\n]*\\n([\\s\\S]*?)\\n\\}`,
"g",
),
),
]
.map((match) => match[1]!)
.join("\n");
// Element loads (the poll's flag read is atomic too, from a global).
const elementLoads = (body: string): number =>
body.match(/load atomic \w+, ptr %/g)?.length ?? 0;
// Elements at the counter of a strip-mined loop: plain loads and stores.
for (const name of ["scale", "codes"]) {
const body = bodies(name);
expect(body).toMatch(/^loop\.poll\d+:$/m);
expect(elementLoads(body)).toBe(0);
expect(body).not.toContain("store atomic");
}
// `a[i + 1]` is not at the counter: its read stays atomic, the store
// at the counter is plain.
const shift = bodies("shift");
expect(shift).toMatch(/load atomic i8, ptr %/);
expect(shift).not.toContain("store atomic");
// The outer counter read inside the inner loop reads one element on every
// inner iteration: it stays atomic; the inner counter's read is plain.
const rows = bodies("rows");
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 });
}
});
packages/compiler/test/ts7/baselines/order-parity/tests/corpus/8332-worker-shared-counted-loops.ts.json
+6
View File
@@ -0,0 +1,6 @@
{
"order": [
"<repo>/tests/corpus/8332-worker-shared-counted-loops.ts"
],
"diags": []
}
packages/compiler/test/worker-shared-loop-version-emission.test.ts
+5 -3
View File
@@ -56,11 +56,13 @@ if (isMainThread) {
// The guard tests the bounds first, then whether the view is shared.
expect(fn).toMatch(/^bytes\.version\.owned\d*:/m);
expect(fn).toMatch(/^bytes\.version\.fast\.shared\d*:/m);
// The shared copy reads with no bounds test, and shared-safe: a view
// with shared storage loads atomically.
// The shared copy reads with no bounds test. Each copy is a counted loop
// that polls between blocks, so its read at the counter is a plain load
// (plainSharedElement in expr-bytes.ts).
const shared = blocks(fn, "bytes.version.fast.shared");
expect(shared).not.toMatch(/icmp (ult|uge|ugt) i64 %\w+, %\w+/);
expect(fn).toMatch(/load atomic double, ptr %\w+ monotonic/);
expect(fn).toMatch(/load double, ptr %\w+/);
expect(fn).not.toMatch(/load atomic double/);
});
test("a program without workers keeps two copies of the loop", async () => {
tests/corpus/8332-worker-shared-counted-loops.ts
+128
View File
@@ -0,0 +1,128 @@
// Counted loops over typed arrays backed by a SharedArrayBuffer in a program
// with workers: loops to a length and to a bound, ascending, descending and
// by two, reads and writes at the counter and next to it, nested loops, and
// each element type. The worker and the main thread run the same loops over
// the same shared memory, one after the other.
import { isMainThread, parentPort, Worker, workerData } from "node:worker_threads";
interface Shared {
bytes: Uint8Array;
shorts: Int16Array;
words: Uint32Array;
ints: Int32Array;
floats: Float32Array;
doubles: Float64Array;
}
function codes(a: Uint8Array): void {
for (let i = 0; i < a.length; ++i) {
const x = a[i];
if (x === 97) a[i] = 0;
else if (x === 116) a[i] = 1;
else if (x === 99) a[i] = 2;
else a[i] = 3;
}
}
function checksum(a: Uint8Array): number {
let t = 0;
for (let i = 0; i < a.length; i++) t = (t * 31 + a[i]!) % 1000003;
return t;
}
function shiftLeft(a: Uint8Array): void {
for (let i = 0; i < a.length - 1; i++) a[i] = a[i + 1]!;
}
function descend(a: Int16Array, n: number): number {
let t = 0;
for (let i = n - 1; i >= 0; i--) {
a[i] = a[i]! - i * 3;
t += a[i]!;
}
return t;
}
function everyOther(a: Uint32Array): number {
let t = 0;
for (let i = 0; i < a.length; i += 2) {
a[i] = a[i]! * 2 + 4000000000;
t += a[i]! % 1000;
}
return t;
}
function grid(a: Int32Array, n: number): number {
let t = 0;
for (let i = 0; i < n; i++) for (let j = 0; j < n; j++) t += a[i]! * a[j]! - j;
return t;
}
function mix(f: Float32Array, d: Float64Array, n: number): number {
for (let i = 0; i < n; i++) {
f[i] = f[i]! * 0.5 + i;
d[i] = d[i]! + f[i]! / 3;
}
let t = 0;
for (let i = 0; i < n; i++) t += d[i]!;
return Math.round(t * 1000) / 1000;
}
function fill(s: Shared): void {
const letters = [97, 116, 99, 103];
let seed = 7;
for (let i = 0; i < s.bytes.length; i++) {
seed = (seed * 1103515245 + 12345) % 2147483648;
s.bytes[i] = letters[(seed >> 16) % 4]!;
}
for (let i = 0; i < s.shorts.length; i++) s.shorts[i] = (i * 37) % 2000 - 1000;
for (let i = 0; i < s.words.length; i++) s.words[i] = i * 2654435761;
for (let i = 0; i < s.ints.length; i++) s.ints[i] = (i % 7) - 3;
for (let i = 0; i < s.floats.length; i++) s.floats[i] = i / 4;
for (let i = 0; i < s.doubles.length; i++) s.doubles[i] = i * 1.25;
}
function run(where: string, s: Shared): string {
codes(s.bytes);
const coded = checksum(s.bytes);
shiftLeft(s.bytes);
const shifted = checksum(s.bytes);
const down = descend(s.shorts, s.shorts.length);
const other = everyOther(s.words);
const cells = grid(s.ints, 40);
const mixed = mix(s.floats, s.doubles, s.floats.length);
return `${where}: ${coded} ${shifted} ${down} ${other} ${cells} ${mixed} ${s.bytes[0]} ${s.words[2]}`;
}
function views(buffers: SharedArrayBuffer[]): Shared {
return {
bytes: new Uint8Array(buffers[0]!),
shorts: new Int16Array(buffers[1]!),
words: new Uint32Array(buffers[2]!),
ints: new Int32Array(buffers[3]!),
floats: new Float32Array(buffers[4]!),
doubles: new Float64Array(buffers[5]!),
};
}
if (isMainThread) {
const buffers = [
new SharedArrayBuffer(10007),
new SharedArrayBuffer(2 * 515),
new SharedArrayBuffer(4 * 333),
new SharedArrayBuffer(4 * 40),
new SharedArrayBuffer(4 * 129),
new SharedArrayBuffer(8 * 129),
];
const shared = views(buffers);
fill(shared);
const worker = new Worker(new URL(import.meta.url), { workerData: buffers });
worker.on("message", (line: string) => {
console.log(line);
console.log(run("main", shared));
fill(shared);
console.log(run("main again", shared));
});
} else {
parentPort!.postMessage(run("worker", views(workerData as SharedArrayBuffer[])));
}
tests/harness/corpus-costs/8332-worker-shared-counted-loops.ts.json
+1
View File
@@ -0,0 +1 @@
{ "plain": 4011, "san": 2208 }