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,