From 604b1c3ec00a6a195c3f0e0e459002a5be24788b Mon Sep 17 00:00:00 2001 From: Akaash Parthasarathy Date: Sun, 16 Aug 2026 07:45:15 -0400 Subject: [PATCH] [Web] Avoid copies when uploading Wasm memory to WebGPU queue.writeBuffer copies its source before returning, so aligned uploads can read directly from Wasm linear memory instead of first copying into another Uint8Array. Use a checked view for aligned transfers. Keep the padded buffer for unaligned lengths, and align GPU allocations and readback copies to WebGPU copy requirements. --- web/src/memory.ts | 28 +++++++++++ web/src/webgpu.ts | 46 +++++++++++++----- web/tests/node/test_memory.js | 63 +++++++++++++++++++++++++ web/tests/node/test_webgpu.js | 89 +++++++++++++++++++++++++++++++++-- 4 files changed, 209 insertions(+), 17 deletions(-) diff --git a/web/src/memory.ts b/web/src/memory.ts index 00dc17d4126a..5da8b596409f 100644 --- a/web/src/memory.ts +++ b/web/src/memory.ts @@ -137,6 +137,34 @@ export class Memory { result.set(this.viewU8.subarray(ptr, ptr + numBytes)); return result; } + /** + * Return a borrowed view of raw bytes in Wasm memory. + * + * The returned view aliases the current WebAssembly.Memory buffer and must + * not be retained across a call that can grow the memory. + * + * @param ptr The head address. + * @param numBytes The number of bytes. + */ + viewRawBytes(ptr: Pointer, numBytes: number): Uint8Array { + if (this.buffer != this.memory.buffer) { + this.updateViews(); + } + if (!Number.isSafeInteger(ptr) || ptr < 0) { + throw new Error(`Invalid Wasm memory pointer: ${ptr}`); + } + if (!Number.isSafeInteger(numBytes) || numBytes < 0) { + throw new Error(`Invalid Wasm memory byte length: ${numBytes}`); + } + const end = ptr + numBytes; + if (!Number.isSafeInteger(end) || end > this.viewU8.byteLength) { + throw new Error( + `Wasm memory range [${ptr}, ${end}) exceeds memory size ` + + `${this.viewU8.byteLength}`, + ); + } + return this.viewU8.subarray(ptr, end); + } /** * Load null-terminated C-string from ptr. * @param ptr The head address diff --git a/web/src/webgpu.ts b/web/src/webgpu.ts index 60b553bdac77..038f9accca94 100644 --- a/web/src/webgpu.ts +++ b/web/src/webgpu.ts @@ -30,6 +30,23 @@ export interface GPUDeviceDetectOutput { device: GPUDevice; } +function roundUpToFourBytes(nbytes: number): number { + if (!Number.isSafeInteger(nbytes) || nbytes < 0) { + throw new Error(`Invalid WebGPU buffer size: ${nbytes}`); + } + const aligned = Math.ceil(nbytes / 4) * 4; + if (!Number.isSafeInteger(aligned)) { + throw new Error(`WebGPU buffer size is too large to align: ${nbytes}`); + } + return aligned; +} + +function validateWebGPUCopyOffset(offset: number, name: string): void { + if (!Number.isSafeInteger(offset) || offset < 0 || offset % 4 != 0) { + throw new Error(`${name} must be a nonnegative multiple of four: ${offset}`); + } +} + /** * DetectGPU device in the environment. */ @@ -914,16 +931,14 @@ export class WebGPUContext { // DeviceAPI private deviceAllocDataSpace(nbytes: number): GPUPointer { - // allocate 0 bytes buffer as 1 bytes buffer. - if (nbytes == 0) { - nbytes = 1; - } + // WebGPU buffer copies and queue writes operate in four-byte units. + const allocationBytes = Math.max(4, roundUpToFourBytes(nbytes)); const buffer = tryCreateBuffer(this.device, { - size: nbytes, + size: allocationBytes, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST, }); - this.currAllocatedBytes += nbytes; - this.allAllocatedBytes += nbytes; + this.currAllocatedBytes += buffer.size; + this.allAllocatedBytes += buffer.size; if (this.currAllocatedBytes > this.peakAllocatedBytes) { this.peakAllocatedBytes = this.currAllocatedBytes; } @@ -951,11 +966,12 @@ export class WebGPUContext { toOffset: number, nbytes: number ): void { + validateWebGPUCopyOffset(toOffset, "WebGPU destination offset"); // Flush batched compute passes before writing to a GPU buffer, // otherwise the write may be reordered before pending dispatches // that read from the same buffer. this.flushCommands(); - let rawBytes = this.memory.loadRawBytes(from, nbytes); + let rawBytes = this.memory.viewRawBytes(from, nbytes); if (rawBytes.length % 4 !== 0) { // writeBuffer requires length to be multiples of 4, so we pad here const toPad = 4 - rawBytes.length % 4; @@ -1016,9 +1032,15 @@ export class WebGPUContext { to: Pointer, nbytes: number ): void { + validateWebGPUCopyOffset(fromOffset, "WebGPU source offset"); // Flush batched compute passes before the readback copy. this.flushCommands(); - const gpuTemp = this.getOrCreateReadStagingBuffer(nbytes); + if (nbytes == 0) { + this.memory.storeRawBytes(to, new Uint8Array(0)); + return; + } + const copyBytes = roundUpToFourBytes(nbytes); + const gpuTemp = this.getOrCreateReadStagingBuffer(copyBytes); const copyEncoder = this.device.createCommandEncoder(); copyEncoder.copyBufferToBuffer( @@ -1026,14 +1048,14 @@ export class WebGPUContext { fromOffset, gpuTemp, 0, - nbytes + copyBytes ); const copyCommands = copyEncoder.finish(); this.device.queue.submit([copyCommands]); const readPromise = gpuTemp.mapAsync(GPUMapMode.READ).then(() => { - const data = gpuTemp.getMappedRange(0, nbytes); - this.memory.storeRawBytes(to, new Uint8Array(data)); + const data = gpuTemp.getMappedRange(0, copyBytes); + this.memory.storeRawBytes(to, new Uint8Array(data).subarray(0, nbytes)); this.recycleReadStagingBuffer(gpuTemp); }); // Chain with any existing pending read so sync() awaits all of them. diff --git a/web/tests/node/test_memory.js b/web/tests/node/test_memory.js index 39de27835c2d..2daa9ff9b4bb 100644 --- a/web/tests/node/test_memory.js +++ b/web/tests/node/test_memory.js @@ -46,6 +46,69 @@ test("loadRawBytes preserves the requested length at the end of memory", () => { expect(Array.from(result)).toEqual([5, 6, 0, 0]); }); +test("viewRawBytes returns a borrowed Wasm memory view", () => { + const wasmMemory = new WebAssembly.Memory({ initial: 1 }); + const memory = new Memory(wasmMemory); + const source = new Uint8Array(wasmMemory.buffer, 8, 4); + source.set([1, 2, 3, 4]); + + const result = memory.viewRawBytes(8, 4); + + expect(Array.from(result)).toEqual([1, 2, 3, 4]); + expect(result.buffer).toBe(wasmMemory.buffer); + source[0] = 10; + result[1] = 20; + expect(Array.from(result)).toEqual([10, 20, 3, 4]); + expect(Array.from(source)).toEqual([10, 20, 3, 4]); +}); + +test("viewRawBytes refreshes its backing view after memory growth", () => { + const wasmMemory = new WebAssembly.Memory({ initial: 1, maximum: 2 }); + const memory = new Memory(wasmMemory); + const oldBuffer = wasmMemory.buffer; + + wasmMemory.grow(1); + const source = new Uint8Array(wasmMemory.buffer, 65536, 4); + source.set([5, 6, 7, 8]); + const result = memory.viewRawBytes(65536, 4); + + expect(result.buffer).toBe(wasmMemory.buffer); + expect(result.buffer).not.toBe(oldBuffer); + expect(Array.from(result)).toEqual([5, 6, 7, 8]); +}); + +test("viewRawBytes supports shared Wasm memory", () => { + const wasmMemory = new WebAssembly.Memory({ + initial: 1, + maximum: 2, + shared: true, + }); + const memory = new Memory(wasmMemory); + const source = new Uint8Array(wasmMemory.buffer, 16, 4); + source.set([1, 2, 3, 4]); + + const result = memory.viewRawBytes(16, 4); + + expect(result.buffer).toBe(wasmMemory.buffer); + expect(result.buffer).toBeInstanceOf(SharedArrayBuffer); + expect(Array.from(result)).toEqual([1, 2, 3, 4]); +}); + +test.each([ + [-1, 1, "pointer"], + [0.5, 1, "pointer"], + [Number.MAX_SAFE_INTEGER + 1, 1, "pointer"], + [0, -1, "byte length"], + [0, 0.5, "byte length"], + [0, Number.MAX_SAFE_INTEGER + 1, "byte length"], + [65536, 1, "exceeds memory size"], + [65535, 2, "exceeds memory size"], + [Number.MAX_SAFE_INTEGER, 1, "exceeds memory size"], +])("viewRawBytes rejects invalid range (%p, %p)", (ptr, nbytes, message) => { + const memory = new Memory(new WebAssembly.Memory({ initial: 1 })); + expect(() => memory.viewRawBytes(ptr, nbytes)).toThrow(message); +}); + test("CachedCallStack commits a view of its cached bytes", () => { const memory = { wasm32: true, diff --git a/web/tests/node/test_webgpu.js b/web/tests/node/test_webgpu.js index dcd0be2e09c1..30f011d7e18e 100644 --- a/web/tests/node/test_webgpu.js +++ b/web/tests/node/test_webgpu.js @@ -111,6 +111,7 @@ function createContext(deviceOptions) { const gpu = createMockDevice(deviceOptions); const memory = { loadRawBytes: jest.fn(), + viewRawBytes: jest.fn(), storeRawBytes: jest.fn(), }; const context = new WebGPUContext(memory, gpu.device); @@ -179,12 +180,16 @@ test("a host write flushes pending GPU copies before writeBuffer", () => { test("an aligned CPU to GPU copy writes the requested bytes", () => { const { context, device, queue, memory, destination } = createContext(); const copyToGPU = context.getDeviceAPI("deviceCopyToGPU"); - const rawBytes = new Uint8Array([1, 2, 3, 4, 5, 6, 7, 8]); - memory.loadRawBytes.mockReturnValue(rawBytes); + const wasmMemory = new WebAssembly.Memory({ initial: 1 }); + const rawBytes = new Uint8Array(wasmMemory.buffer, 128, 8); + rawBytes.set([1, 2, 3, 4, 5, 6, 7, 8]); + memory.viewRawBytes.mockReturnValue(rawBytes); copyToGPU(128, destination, 12, rawBytes.length); - expect(memory.loadRawBytes).toHaveBeenCalledWith(128, rawBytes.length); + expect(memory.viewRawBytes).toHaveBeenCalledWith(128, rawBytes.length); + expect(memory.loadRawBytes).not.toHaveBeenCalled(); + expect(queue.writeBuffer.mock.calls[0][2].buffer).toBe(wasmMemory.buffer); expect(queue.writeBuffer).toHaveBeenCalledWith( device.createBuffer.mock.results[1].value, 12, @@ -198,11 +203,12 @@ test("an unaligned CPU to GPU copy pads the write to four bytes", () => { const { context, device, queue, memory, destination } = createContext(); const copyToGPU = context.getDeviceAPI("deviceCopyToGPU"); const rawBytes = new Uint8Array([1, 2, 3]); - memory.loadRawBytes.mockReturnValue(rawBytes); + memory.viewRawBytes.mockReturnValue(rawBytes); copyToGPU(256, destination, 4, rawBytes.length); - expect(memory.loadRawBytes).toHaveBeenCalledWith(256, rawBytes.length); + expect(memory.viewRawBytes).toHaveBeenCalledWith(256, rawBytes.length); + expect(memory.loadRawBytes).not.toHaveBeenCalled(); expect(queue.writeBuffer).toHaveBeenCalledTimes(1); const [buffer, toOffset, data, dataOffset, nbytes] = queue.writeBuffer.mock.calls[0]; @@ -213,6 +219,79 @@ test("an unaligned CPU to GPU copy pads the write to four bytes", () => { expect(nbytes).toBe(4); }); +test("a non-four-byte GPU allocation is rounded up for padded writes", () => { + const { context, device } = createContext(); + const allocate = context.getDeviceAPI("deviceAllocDataSpace"); + + allocate(3); + + expect(device.createBuffer).toHaveBeenLastCalledWith({ + size: 4, + usage: GPUBufferUsage.STORAGE | + GPUBufferUsage.COPY_SRC | + GPUBufferUsage.COPY_DST, + }); + expect(context.currAllocatedBytes).toBe(64 + 64 + 4); +}); + +test.each([ + [-1, "destination offset"], + [0.5, "destination offset"], + [2, "destination offset"], +])("a CPU to GPU copy rejects invalid offset %p", (offset, message) => { + const { context, memory, destination } = createContext(); + const copyToGPU = context.getDeviceAPI("deviceCopyToGPU"); + + expect(() => copyToGPU(128, destination, offset, 4)).toThrow(message); + expect(memory.viewRawBytes).not.toHaveBeenCalled(); +}); + +test.each([ + [-1, "source offset"], + [0.5, "source offset"], + [2, "source offset"], +])("a GPU readback rejects invalid offset %p", (offset, message) => { + const { context, source } = createContext(); + const copyFromGPU = context.getDeviceAPI("deviceCopyFromGPU"); + + expect(() => copyFromGPU(source, offset, 128, 4)).toThrow(message); +}); + +test("an unaligned GPU readback copies four bytes and stores the logical bytes", async () => { + const { + context, + device, + memory, + source, + } = createContext(); + const copyFromGPU = context.getDeviceAPI("deviceCopyFromGPU"); + device.createBuffer.mockImplementationOnce((descriptor) => { + const mappedData = new Uint8Array([1, 2, 3, 99]).buffer; + return { + size: descriptor.size, + destroy: jest.fn(), + mapAsync: jest.fn(() => Promise.resolve()), + getMappedRange: jest.fn(() => mappedData), + unmap: jest.fn(), + }; + }); + + copyFromGPU(source, 0, 128, 3); + await context.sync(); + + const copyEncoder = device.createCommandEncoder.mock.results[0].value; + expect(copyEncoder.copyBufferToBuffer).toHaveBeenCalledWith( + device.createBuffer.mock.results[0].value, + 0, + device.createBuffer.mock.results[2].value, + 0, + 4, + ); + expect(memory.storeRawBytes).toHaveBeenCalledTimes(1); + expect(memory.storeRawBytes.mock.calls[0][0]).toBe(128); + expect(Array.from(memory.storeRawBytes.mock.calls[0][1])).toEqual([1, 2, 3]); +}); + test("a GPU readback flushes pending copies before its own submission", async () => { const { context,