mirror of
https://github.com/vercel-labs/scriptc.git
synced 2026-10-27 06:43:19 +08:00
Read and write shared typed-array elements inline
This commit is contained in:
Chris Tate
@@ -176,10 +176,10 @@ export function emitBytesLength(
|
||||
return { name: out, type: F64 };
|
||||
}
|
||||
|
||||
/** Worker programs: a view of a SharedArrayBuffer (`shared` is fixed before
|
||||
* script code can see the view) is read and written by the runtime under
|
||||
* its shared-memory lock; every other buffer belongs to this thread and
|
||||
* takes the ordinary inline path. The test follows the index check: a
|
||||
/** Worker programs: an element of a view of a SharedArrayBuffer (`shared` is
|
||||
* fixed before script code can see the view) is read and written with one
|
||||
* relaxed atomic access of its width, as the runtime's shared accessors do;
|
||||
* every other buffer belongs to this thread and takes the ordinary path. The test follows the index check: a
|
||||
* view's length never changes, so both paths use the one checked integer
|
||||
* index. Branches to the returned label for a thread-owned buffer and
|
||||
* continues in the shared block. */
|
||||
@@ -276,8 +276,9 @@ function emitBytesLoad(
|
||||
* 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
|
||||
* ToInt32/ToUint32 of the result stay integer operations. The shared read
|
||||
* is an ordinary call: another thread's write must stay visible to a loop
|
||||
* that reads the same element again. */
|
||||
* is a monotonic atomic load, which LLVM neither hoists nor merges: another
|
||||
* thread's write stays visible to a loop that reads the same element again.
|
||||
* Float16 elements read through the runtime. */
|
||||
function emitSharedOrLocalLoad(
|
||||
host: LlvmEmitterContext,
|
||||
elem: IrBytesElem,
|
||||
@@ -286,12 +287,18 @@ function emitSharedOrLocalLoad(
|
||||
): LlValue {
|
||||
const B = host.B;
|
||||
const floating = elem === "f16" || elem === "f32" || elem === "f64";
|
||||
const fn = floating ? "scr_bytes_shared_get_f64" : "scr_bytes_shared_get_int";
|
||||
const ret = floating ? "double" : "i32";
|
||||
host.declare(`declare ${ret} @${fn}(ptr, ${host.sizeType})`);
|
||||
const local = emitSharedBytesSplit(host, receiver);
|
||||
const shared = B.tmp();
|
||||
B.line(`${shared} = call ${ret} @${fn}(ptr ${receiver}, ${host.sizeType} ${idx}) cold`);
|
||||
let shared: string;
|
||||
if (elem === "f16") {
|
||||
host.declare(`declare double @scr_bytes_shared_get_f64(ptr, ${host.sizeType})`);
|
||||
shared = B.tmp();
|
||||
B.line(
|
||||
`${shared} = call double @scr_bytes_shared_get_f64(ptr ${receiver}, ${host.sizeType} ${idx}) cold`,
|
||||
);
|
||||
} else {
|
||||
const loaded = emitSharedLoadAt(host, elem, host.emitBytesData(receiver), idx);
|
||||
shared = loaded.uint32 ?? loaded.name;
|
||||
}
|
||||
const sharedDone = B.newLabel("bytes.shared.done");
|
||||
B.br(sharedDone);
|
||||
B.startBlock(local);
|
||||
@@ -316,6 +323,43 @@ function emitSharedOrLocalLoad(
|
||||
return { name: out, type: F64, uint32: u32 };
|
||||
}
|
||||
|
||||
/** The element at a checked index of a shared view, given its data pointer:
|
||||
* a relaxed atomic load of the element's width, aligned as shared views
|
||||
* are. Not for Float16 elements. */
|
||||
function emitSharedLoadAt(
|
||||
host: LlvmEmitterContext,
|
||||
elem: IrBytesElem,
|
||||
data: string,
|
||||
idx: string,
|
||||
): LlValue {
|
||||
const B = host.B;
|
||||
const size = BYTES_ELEMENT_SIZE[elem];
|
||||
const p = B.tmp();
|
||||
if (elem === "f32" || elem === "f64") {
|
||||
const ty = elem === "f32" ? "float" : "double";
|
||||
const raw = B.tmp();
|
||||
B.line(`${p} = getelementptr inbounds ${ty}, ptr ${data}, ${host.sizeType} ${idx}`);
|
||||
B.line(`${raw} = load atomic ${ty}, ptr ${p} monotonic, align ${size}`);
|
||||
if (elem === "f64") return { name: raw, type: F64 };
|
||||
const out = B.tmp();
|
||||
B.line(`${out} = fpext float ${raw} to double`);
|
||||
return { name: out, type: F64 };
|
||||
}
|
||||
const bits = size * 8;
|
||||
const raw = B.tmp();
|
||||
B.line(`${p} = getelementptr inbounds i${bits}, ptr ${data}, ${host.sizeType} ${idx}`);
|
||||
B.line(`${raw} = load atomic i${bits}, ptr ${p} monotonic, align ${size}`);
|
||||
const signed = elem === "i8" || elem === "i16" || elem === "i32";
|
||||
let wide = raw;
|
||||
if (bits < 32) {
|
||||
wide = B.tmp();
|
||||
B.line(`${wide} = ${signed ? "sext" : "zext"} i${bits} ${raw} to i32`);
|
||||
}
|
||||
const out = B.tmp();
|
||||
B.line(`${out} = ${signed ? "sitofp" : "uitofp"} i32 ${wide} to double`);
|
||||
return { name: out, type: F64, uint32: wide };
|
||||
}
|
||||
|
||||
/** The element at a checked index of a thread-owned view, given its data
|
||||
* pointer. */
|
||||
function emitBytesLoadAt(
|
||||
@@ -477,7 +521,7 @@ export function emitBytesSet(
|
||||
const done = B.newLabel("bytes.store.done");
|
||||
const data = host.emitBytesData(receiver);
|
||||
const idx = emitBytesIndex(host, receiver, index, indexExpr, done, inBounds);
|
||||
if (host.mod.workers) {
|
||||
if (host.mod.workers && elem === "f16") {
|
||||
host.declare(`declare void @scr_bytes_shared_set(ptr, ${host.sizeType}, double)`);
|
||||
const local = emitSharedBytesSplit(host, receiver);
|
||||
B.line(
|
||||
@@ -498,32 +542,48 @@ export function emitBytesSet(
|
||||
elem === "f16" || elem === "f32" || elem === "f64"
|
||||
? null
|
||||
: host.emitToUint32(value.name, expr, value.uint32);
|
||||
const p = B.tmp();
|
||||
if (elem === "f16") {
|
||||
// One round-to-nearest-even step from the double, in the runtime.
|
||||
host.declare(`declare zeroext i16 @scr_f16_from_f64(double)`);
|
||||
const narrowed = B.tmp();
|
||||
const p = B.tmp();
|
||||
B.line(`${narrowed} = call zeroext i16 @scr_f16_from_f64(double ${value.name})`);
|
||||
B.line(`${p} = getelementptr inbounds i16, ptr ${data}, ${host.sizeType} ${idx}`);
|
||||
B.line(`store i16 ${narrowed}, ptr ${p}, align 1${bytesElementAlias(host)}`);
|
||||
} else if (BYTES_ELEMENT_SIZE[elem] < 4) {
|
||||
const bits = BYTES_ELEMENT_SIZE[elem] * 8;
|
||||
const byte = B.tmp();
|
||||
B.line(`${byte} = trunc i32 ${stored!} to i${bits}`);
|
||||
B.line(`${p} = getelementptr inbounds i${bits}, ptr ${data}, ${host.sizeType} ${idx}`);
|
||||
B.line(`store i${bits} ${byte}, ptr ${p}, align 1${bytesElementAlias(host)}`);
|
||||
} else if (elem === "f32") {
|
||||
const narrowed = B.tmp();
|
||||
B.line(`${narrowed} = fptrunc double ${value.name} to float`);
|
||||
B.line(`${p} = getelementptr inbounds float, ptr ${data}, ${host.sizeType} ${idx}`);
|
||||
B.line(`store float ${narrowed}, ptr ${p}, align 1${bytesElementAlias(host)}`);
|
||||
} else if (elem === "f64") {
|
||||
B.line(`${p} = getelementptr inbounds double, ptr ${data}, ${host.sizeType} ${idx}`);
|
||||
B.line(`store double ${value.name}, ptr ${p}, align 1${bytesElementAlias(host)}`);
|
||||
} else {
|
||||
B.line(`${p} = getelementptr inbounds i32, ptr ${data}, ${host.sizeType} ${idx}`);
|
||||
B.line(`store i32 ${stored!}, ptr ${p}, align 1${bytesElementAlias(host)}`);
|
||||
B.br(done);
|
||||
B.startBlock(done);
|
||||
return;
|
||||
}
|
||||
// The element's bits as an LLVM value of the element's width.
|
||||
let ty: string;
|
||||
let operand: string;
|
||||
if (BYTES_ELEMENT_SIZE[elem] < 4) {
|
||||
ty = `i${BYTES_ELEMENT_SIZE[elem] * 8}`;
|
||||
operand = B.tmp();
|
||||
B.line(`${operand} = trunc i32 ${stored!} to ${ty}`);
|
||||
} else if (elem === "f32") {
|
||||
ty = "float";
|
||||
operand = B.tmp();
|
||||
B.line(`${operand} = fptrunc double ${value.name} to float`);
|
||||
} else if (elem === "f64") {
|
||||
ty = "double";
|
||||
operand = value.name;
|
||||
} else {
|
||||
ty = "i32";
|
||||
operand = stored!;
|
||||
}
|
||||
if (host.mod.workers) {
|
||||
// A shared view's element takes a relaxed atomic store of its width.
|
||||
const local = emitSharedBytesSplit(host, receiver);
|
||||
const p = B.tmp();
|
||||
B.line(`${p} = getelementptr inbounds ${ty}, ptr ${data}, ${host.sizeType} ${idx}`);
|
||||
B.line(`store atomic ${ty} ${operand}, ptr ${p} monotonic, align ${BYTES_ELEMENT_SIZE[elem]}`);
|
||||
B.br(done);
|
||||
B.startBlock(local);
|
||||
}
|
||||
const p = B.tmp();
|
||||
B.line(`${p} = getelementptr inbounds ${ty}, ptr ${data}, ${host.sizeType} ${idx}`);
|
||||
B.line(`store ${ty} ${operand}, ptr ${p}, align 1${bytesElementAlias(host)}`);
|
||||
B.br(done);
|
||||
B.startBlock(done);
|
||||
}
|
||||
|
||||
@@ -335,11 +335,17 @@ test("a body mutation keeps the byte loop on the general f64 path", () => {
|
||||
|
||||
test("worker programs read and write shared views at the checked integer index", () => {
|
||||
const ll = emitLlvmModule({ ...fixture(), workers: true });
|
||||
// The shared path receives the index the inline path checked, as an integer,
|
||||
// and integer elements come back as 32-bit values.
|
||||
expect(ll).toMatch(/call i32 @scr_bytes_shared_get_int\(ptr %\w+, i64 %?\w+\) cold/);
|
||||
expect(ll).toMatch(/call double @scr_bytes_shared_get_f64\(ptr %\w+, i64 %?\w+\) cold/);
|
||||
expect(ll).toMatch(/call void @scr_bytes_shared_set\(ptr %\w+, i64 %?\w+, double %?[\w.]+\)/);
|
||||
// The shared path uses the index the inline path checked, as an integer:
|
||||
// each element is one relaxed atomic access of its width, aligned to it.
|
||||
expect(ll).toMatch(/load atomic i8, ptr %\w+ monotonic, align 1/);
|
||||
expect(ll).toMatch(/load atomic i32, ptr %\w+ monotonic, align 4/);
|
||||
expect(ll).toMatch(/load atomic float, ptr %\w+ monotonic, align 4/);
|
||||
expect(ll).toMatch(/store atomic i8 [%\w.]+, ptr %\w+ monotonic, align 1/);
|
||||
expect(ll).toMatch(/store atomic i32 [%\w.]+, ptr %\w+ monotonic, align 4/);
|
||||
expect(ll).toMatch(/store atomic float [%\w.]+, ptr %\w+ monotonic, align 4/);
|
||||
expect(ll).not.toContain("@scr_bytes_shared_get_int");
|
||||
expect(ll).not.toContain("@scr_bytes_shared_get_f64");
|
||||
expect(ll).not.toContain("@scr_bytes_shared_set");
|
||||
expect(ll).not.toContain("@scr_bytes_get_or_nan");
|
||||
expect(ll).not.toContain("@scr_bytes_set(");
|
||||
// Both paths join as one integer, which converts to a number after the join.
|
||||
|
||||
+6
@@ -0,0 +1,6 @@
|
||||
{
|
||||
"order": [
|
||||
"<repo>/tests/corpus/6864-shared-typed-array-elements.ts"
|
||||
],
|
||||
"diags": []
|
||||
}
|
||||
@@ -0,0 +1,120 @@
|
||||
// Element reads and writes of typed arrays over a SharedArrayBuffer, for every
|
||||
// element kind, from the main thread and from a worker: each kind's
|
||||
// conversion on write (wrapping, clamping, rounding to float32 and float16),
|
||||
// NaN, -0 and the infinities, and values one thread writes and the other
|
||||
// reads after a message. A worker loop that reads a flag without Atomics
|
||||
// sees the main thread's write.
|
||||
import { isMainThread, parentPort, Worker, workerData } from "node:worker_threads";
|
||||
|
||||
const KINDS = 10;
|
||||
|
||||
const VALUES = [0, -0, 1, -1, 1.5, 2.5, -2.5, 127, 128, 255, 256, -129, 65535, 65536, 2 ** 31, -(2 ** 31) - 1, 2 ** 32 + 7, 0.1, 1e40, -1e-50, NaN, Infinity, -Infinity];
|
||||
const LENGTH = VALUES.length;
|
||||
|
||||
interface Views {
|
||||
i8: Int8Array;
|
||||
u8: Uint8Array;
|
||||
u8c: Uint8ClampedArray;
|
||||
i16: Int16Array;
|
||||
u16: Uint16Array;
|
||||
i32: Int32Array;
|
||||
u32: Uint32Array;
|
||||
f16: Float16Array;
|
||||
f32: Float32Array;
|
||||
f64: Float64Array;
|
||||
}
|
||||
|
||||
/** One region of LENGTH elements per kind, each aligned to 8 bytes. */
|
||||
function viewsOf(sab: SharedArrayBuffer): Views {
|
||||
const region = LENGTH * 8;
|
||||
return {
|
||||
i8: new Int8Array(sab, 0 * region, LENGTH),
|
||||
u8: new Uint8Array(sab, 1 * region, LENGTH),
|
||||
u8c: new Uint8ClampedArray(sab, 2 * region, LENGTH),
|
||||
i16: new Int16Array(sab, 3 * region, LENGTH),
|
||||
u16: new Uint16Array(sab, 4 * region, LENGTH),
|
||||
i32: new Int32Array(sab, 5 * region, LENGTH),
|
||||
u32: new Uint32Array(sab, 6 * region, LENGTH),
|
||||
f16: new Float16Array(sab, 7 * region, LENGTH),
|
||||
f32: new Float32Array(sab, 8 * region, LENGTH),
|
||||
f64: new Float64Array(sab, 9 * region, LENGTH),
|
||||
};
|
||||
}
|
||||
|
||||
function show(x: number): string {
|
||||
return Object.is(x, -0) ? "-0" : String(x);
|
||||
}
|
||||
|
||||
/** Writes VALUES (scaled by `scale`) through every view. */
|
||||
function writeAll(v: Views, scale: number): void {
|
||||
for (let i = 0; i < LENGTH; i++) {
|
||||
const x = VALUES[i]! * scale;
|
||||
v.i8[i] = x;
|
||||
v.u8[i] = x;
|
||||
v.u8c[i] = x;
|
||||
v.i16[i] = x;
|
||||
v.u16[i] = x;
|
||||
v.i32[i] = x;
|
||||
v.u32[i] = x;
|
||||
v.f16[i] = x;
|
||||
v.f32[i] = x;
|
||||
v.f64[i] = x;
|
||||
}
|
||||
}
|
||||
|
||||
/** Every element of every view, read one at a time, and a sum per view. */
|
||||
function readAll(v: Views): string[] {
|
||||
const lines: string[] = [];
|
||||
const row = (kind: string, read: (i: number) => number): void => {
|
||||
const cells: string[] = [];
|
||||
for (let i = 0; i < LENGTH; i++) cells.push(show(read(i)));
|
||||
lines.push(`${kind}: ${cells.join(" ")}`);
|
||||
};
|
||||
row("Int8Array", (i) => v.i8[i]!);
|
||||
row("Uint8Array", (i) => v.u8[i]!);
|
||||
row("Uint8ClampedArray", (i) => v.u8c[i]!);
|
||||
row("Int16Array", (i) => v.i16[i]!);
|
||||
row("Uint16Array", (i) => v.u16[i]!);
|
||||
row("Int32Array", (i) => v.i32[i]!);
|
||||
row("Uint32Array", (i) => v.u32[i]!);
|
||||
row("Float16Array", (i) => v.f16[i]!);
|
||||
row("Float32Array", (i) => v.f32[i]!);
|
||||
row("Float64Array", (i) => v.f64[i]!);
|
||||
let total = 0;
|
||||
for (let i = 0; i < LENGTH; i++) total += v.i32[i]! + v.u16[i]! + (v.u8c[i]! | 0);
|
||||
lines.push(`integer total ${total}`);
|
||||
return lines;
|
||||
}
|
||||
|
||||
if (isMainThread) {
|
||||
const sab = new SharedArrayBuffer(KINDS * LENGTH * 8 + 16);
|
||||
const views = viewsOf(sab);
|
||||
const flag = new Int32Array(sab, KINDS * LENGTH * 8, 2);
|
||||
writeAll(views, 1);
|
||||
console.log("main wrote, main reads");
|
||||
for (const line of readAll(views)) console.log(line);
|
||||
|
||||
const worker = new Worker(new URL(import.meta.url), { workerData: sab });
|
||||
worker.on("message", (lines: string[]) => {
|
||||
if (lines.length === 1) {
|
||||
console.log(lines[0]);
|
||||
return;
|
||||
}
|
||||
console.log("main wrote, worker reads; worker wrote, main reads");
|
||||
for (const line of lines) console.log(line);
|
||||
for (const line of readAll(views)) console.log(line);
|
||||
flag[1] = 7;
|
||||
flag[0] = 1;
|
||||
});
|
||||
} else {
|
||||
const sab = workerData as SharedArrayBuffer;
|
||||
const views = viewsOf(sab);
|
||||
const flag = new Int32Array(sab, KINDS * LENGTH * 8, 2);
|
||||
const seen = readAll(views);
|
||||
writeAll(views, -3);
|
||||
parentPort!.postMessage(seen);
|
||||
while (flag[0] === 0) {
|
||||
// Spin until the main thread's write is visible.
|
||||
}
|
||||
parentPort!.postMessage([`worker saw the flag, then ${flag[1]}`]);
|
||||
}
|
||||
@@ -0,0 +1 @@
|
||||
{ "plain": 2290, "san": 5070 }
|
||||
Reference in New Issue
Block a user