This is an automated email from the ASF dual-hosted git repository.
akaashrp pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new abe2526b3e [Web] Batch GPU-to-GPU copies, fix WebGPU synchronization,
and add tests for command batching (#20059)
abe2526b3e is described below
commit abe2526b3e54bd47eaa880c0e832ae7edf75a918
Author: Akaash Parthasarathy <[email protected]>
AuthorDate: Fri Jul 31 01:08:14 2026 -0700
[Web] Batch GPU-to-GPU copies, fix WebGPU synchronization, and add tests
for command batching (#20059)
Batch GPU-to-GPU buffer copies with pending compute dispatches in the
same `GPUCommandEncoder` instead of finishing and submitting after every
copy. Since `copyBufferToBuffer` records an encoder command,
dependencies remain ordered while copies can be submitted at the next
required synchronization boundary (see
https://www.w3.org/TR/webgpu/#buffer-copies).
Preserve pending GPU-to-CPU readbacks separately from their queue-tail
status so `sync()` waits for both host readback completion and any later
queue work when necessary. Also flush pending commands before canvas
drawing to prevent rendering stale buffer contents.
Add WebGPU device mocks and tests covering:
- Mixed compute and GPU-copy batching
- Host-write, readback, deallocation, and canvas-draw flush boundaries
- `deviceCopyFromGPU` followed by a batched GPU copy and `sync()`
- Readback error propagation
Regarding importing from `../../src/webgpu` instead of `../../dist`:
`WebGPUContext` is not re-exported from `index.ts`, so it is unavailable
through the built package like for the other tests. Exporting it solely
for this test would unnecessarily expand the public API. Importing
`../../src/webgpu` is a relatively simple alternative.
---
web/jest.config.js | 5 +
web/src/webgpu.ts | 63 ++--
web/tests/node/test_webgpu.js | 330 +++++++++++++++++++++
.../typescript_transformer.cjs} | 18 +-
4 files changed, 386 insertions(+), 30 deletions(-)
diff --git a/web/jest.config.js b/web/jest.config.js
index 01ea70e9d9..6cd117872d 100644
--- a/web/jest.config.js
+++ b/web/jest.config.js
@@ -19,6 +19,11 @@
module.exports = {
testEnvironment: "node",
+ transform: {
+ "^.+\\.ts$": "<rootDir>/tests/typescript_transformer.cjs",
+ "^.+\\.js$": "babel-jest",
+ },
+
testMatch: [
"**/tests/node/*.js"
],
diff --git a/web/src/webgpu.ts b/web/src/webgpu.ts
index 5a25833d1a..60b553bdac 100644
--- a/web/src/webgpu.ts
+++ b/web/src/webgpu.ts
@@ -419,14 +419,12 @@ export class WebGPUContext {
// Pool of MAP_READ staging buffers to avoid per-copy create/destroy overhead
private readStagingBufferPool: Array<{ buffer: GPUBuffer; size: number }> =
[];
private maxReadStagingBuffers = 4;
- // Pending mapAsync promise from the last GPU→CPU copy.
- // Used in sync() as a fast path: if the last queue operation was a
- // GPU→CPU copy, awaiting its mapAsync is sufficient (no need for
- // the heavier onSubmittedWorkDone). Reset to null after any non-copy
- // queue submission so we fall back to onSubmittedWorkDone.
+ // Pending GPU→CPU copies, including storing the mapped data in WASM memory.
private pendingGPUToCPUCopy: Promise<void> | null = null;
- // Batched command encoding: accumulate compute passes in a single encoder,
- // submit only on flush to reduce JS-native transition overhead.
+ // Whether a pending GPU→CPU copy is still the last queue operation.
+ private pendingGPUToCPUCopyIsQueueTail = false;
+ // Batched command encoding: accumulate compute passes and GPU copies in a
+ // single encoder, and submit only on flush to reduce JS-native transition
overhead.
private pendingEncoder: GPUCommandEncoder | null = null;
// Pool of uniform buffers reused across flushes. Each dispatch in a batch
// gets its own buffer (indexed by pendingDispatchCount). The pool grows
@@ -455,14 +453,14 @@ export class WebGPUContext {
}
/**
- * Flush all pending compute passes by finishing and submitting the
+ * Flush all pending GPU commands by finishing and submitting the
* accumulated command encoder.
*
* Must be called before:
* - GPU→CPU readback (deviceCopyFromGPU)
* - CPU→GPU writes (deviceCopyToGPU, copyRawBytesToBuffer)
- * - GPU↔GPU copies (deviceCopyWithinGPU)
* - Buffer deallocation (deviceFreeDataSpace)
+ * - Canvas drawing (drawImageFromBuffer)
* - Queue sync (sync)
*/
flushCommands(): void {
@@ -470,9 +468,7 @@ export class WebGPUContext {
this.device.queue.submit([this.pendingEncoder.finish()]);
this.pendingEncoder = null;
this.pendingDispatchCount = 0;
- // A compute submission is now the last queue operation, so the
- // GPU→CPU copy fast path in sync() is no longer valid.
- this.pendingGPUToCPUCopy = null;
+ this.pendingGPUToCPUCopyIsQueueTail = false;
}
}
@@ -502,14 +498,22 @@ export class WebGPUContext {
* Wait for all pending GPU tasks to complete
*/
async sync(): Promise<void> {
- // Flush any batched compute passes before waiting on the queue.
this.flushCommands();
- if (this.pendingGPUToCPUCopy) {
- const p = this.pendingGPUToCPUCopy;
- this.pendingGPUToCPUCopy = null;
- await p;
+
+ const pendingRead = this.pendingGPUToCPUCopy;
+ const pendingReadIsQueueTail = this.pendingGPUToCPUCopyIsQueueTail;
+ this.pendingGPUToCPUCopy = null;
+ this.pendingGPUToCPUCopyIsQueueTail = false;
+
+ if (pendingRead && pendingReadIsQueueTail) {
+ await pendingRead;
} else {
- await this.device.queue.onSubmittedWorkDone();
+ const queueDone = this.device.queue.onSubmittedWorkDone();
+ if (pendingRead) {
+ await Promise.all([pendingRead, queueDone]);
+ } else {
+ await queueDone;
+ }
}
}
@@ -533,7 +537,9 @@ export class WebGPUContext {
if (this.canvasRenderManager == undefined) {
throw Error("Do not have a canvas context, call bindCanvas first");
}
+ this.flushCommands();
this.canvasRenderManager.draw(this.gpuBufferFromPtr(ptr), height, width);
+ this.pendingGPUToCPUCopyIsQueueTail = false;
}
/**
@@ -559,12 +565,16 @@ export class WebGPUContext {
0,
nbytes
);
+ this.pendingGPUToCPUCopyIsQueueTail = false;
}
/**
* Clear canvas
*/
clearCanvas() {
- this.canvasRenderManager?.clear();
+ if (this.canvasRenderManager) {
+ this.canvasRenderManager.clear();
+ this.pendingGPUToCPUCopyIsQueueTail = false;
+ }
}
/**
@@ -961,6 +971,7 @@ export class WebGPUContext {
0,
nbytes
);
+ this.pendingGPUToCPUCopyIsQueueTail = false;
}
/**
@@ -1029,6 +1040,7 @@ export class WebGPUContext {
this.pendingGPUToCPUCopy = this.pendingGPUToCPUCopy
? this.pendingGPUToCPUCopy.then(() => readPromise)
: readPromise;
+ this.pendingGPUToCPUCopyIsQueueTail = true;
}
private deviceCopyWithinGPU(
@@ -1038,18 +1050,19 @@ export class WebGPUContext {
toOffset: number,
nbytes: number
): void {
- // Flush batched compute passes before the GPU-to-GPU copy.
- this.flushCommands();
- const copyEncoder = this.device.createCommandEncoder();
- copyEncoder.copyBufferToBuffer(
+ // Keep copies in the same command encoder as compute dispatches. Command
+ // ordering within the encoder preserves dependencies, while a later
+ // readback, CPU write, deallocation, or sync provides the flush point.
+ if (!this.pendingEncoder) {
+ this.pendingEncoder = this.device.createCommandEncoder();
+ }
+ this.pendingEncoder.copyBufferToBuffer(
this.gpuBufferFromPtr(from),
fromOffset,
this.gpuBufferFromPtr(to),
toOffset,
nbytes
);
- const copyCommands = copyEncoder.finish();
- this.device.queue.submit([copyCommands]);
}
private gpuBufferFromPtr(ptr: GPUPointer): GPUBuffer {
diff --git a/web/tests/node/test_webgpu.js b/web/tests/node/test_webgpu.js
new file mode 100644
index 0000000000..30c6c8c581
--- /dev/null
+++ b/web/tests/node/test_webgpu.js
@@ -0,0 +1,330 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+const { WebGPUContext } = require("../../src/webgpu");
+
+global.GPUBufferUsage = {
+ MAP_READ: 1 << 0,
+ COPY_DST: 1 << 1,
+ COPY_SRC: 1 << 2,
+ STORAGE: 1 << 3,
+ UNIFORM: 1 << 4,
+};
+global.GPUMapMode = {
+ READ: 1,
+};
+global.GPUShaderStage = {
+ COMPUTE: 1,
+};
+
+function createDeferred() {
+ let resolve;
+ let reject;
+ const promise = new Promise((resolvePromise, rejectPromise) => {
+ resolve = resolvePromise;
+ reject = rejectPromise;
+ });
+ return { promise, resolve, reject };
+}
+
+function createMockDevice({
+ mapAsync = () => Promise.resolve(),
+ onSubmittedWorkDone = () => Promise.resolve(),
+} = {}) {
+ const events = [];
+ const encoders = [];
+
+ const queue = {
+ submit: jest.fn(() => events.push("submit")),
+ writeBuffer: jest.fn(() => events.push("writeBuffer")),
+ onSubmittedWorkDone: jest.fn(onSubmittedWorkDone),
+ };
+
+ const device = {
+ queue,
+ createCommandEncoder: jest.fn(() => {
+ const commands = [];
+ const encoderId = encoders.length;
+ const encoder = {
+ commands,
+ beginComputePass: jest.fn(() => ({
+ setPipeline: jest.fn(),
+ setBindGroup: jest.fn(),
+ dispatchWorkgroups: jest.fn(() => {
+ commands.push("compute");
+ events.push("compute");
+ }),
+ end: jest.fn(),
+ })),
+ copyBufferToBuffer: jest.fn(() => {
+ commands.push("copy");
+ events.push("copy");
+ }),
+ finish: jest.fn(() => {
+ const commandBuffer = { encoderId, commands: commands.slice() };
+ events.push("finish");
+ return commandBuffer;
+ }),
+ };
+ encoders.push(encoder);
+ return encoder;
+ }),
+ createBuffer: jest.fn((descriptor) => {
+ const mappedData = new ArrayBuffer(descriptor.size);
+ return {
+ size: descriptor.size,
+ destroy: jest.fn(() => events.push("destroy")),
+ mapAsync: jest.fn(mapAsync),
+ getMappedRange: jest.fn(() => mappedData),
+ unmap: jest.fn(),
+ };
+ }),
+ createBindGroupLayout: jest.fn(() => ({})),
+ createPipelineLayout: jest.fn(() => ({})),
+ createShaderModule: jest.fn(() => ({})),
+ createComputePipeline: jest.fn(() => ({})),
+ createBindGroup: jest.fn(() => ({})),
+ pushErrorScope: jest.fn(),
+ popErrorScope: jest.fn(() => Promise.resolve(null)),
+ destroy: jest.fn(),
+ };
+
+ return { device, queue, events, encoders };
+}
+
+function createContext(deviceOptions) {
+ const gpu = createMockDevice(deviceOptions);
+ const memory = {
+ storeRawBytes: jest.fn(),
+ };
+ const context = new WebGPUContext(memory, gpu.device);
+ const allocate = context.getDeviceAPI("deviceAllocDataSpace");
+
+ return {
+ ...gpu,
+ context,
+ memory,
+ source: allocate(64),
+ destination: allocate(64),
+ };
+}
+
+test("compute dispatches and GPU copies share one submission", async () => {
+ const { context, device, queue, encoders, source, destination } =
createContext();
+ const copyWithinGPU = context.getDeviceAPI("deviceCopyWithinGPU");
+ const shader = context.createShader(
+ {
+ name: "main",
+ arg_types: [],
+ launch_param_tags: [],
+ },
+ "@compute @workgroup_size(1) fn main() {}"
+ );
+
+ shader();
+ copyWithinGPU(source, 0, destination, 0, 16);
+ copyWithinGPU(destination, 16, source, 32, 16);
+
+ expect(device.createCommandEncoder).toHaveBeenCalledTimes(1);
+ expect(queue.submit).not.toHaveBeenCalled();
+ expect(encoders[0].commands).toEqual(["compute", "copy", "copy"]);
+
+ await context.sync();
+
+ expect(encoders[0].finish).toHaveBeenCalledTimes(1);
+ expect(queue.submit).toHaveBeenCalledTimes(1);
+ expect(queue.submit.mock.calls[0][0]).toEqual([
+ {
+ encoderId: 0,
+ commands: encoders[0].commands,
+ },
+ ]);
+ expect(queue.onSubmittedWorkDone).toHaveBeenCalledTimes(1);
+
+ await context.sync();
+ expect(queue.submit).toHaveBeenCalledTimes(1);
+});
+
+test("a host write flushes pending GPU copies before writeBuffer", () => {
+ const { context, queue, events, source, destination } = createContext();
+ const copyWithinGPU = context.getDeviceAPI("deviceCopyWithinGPU");
+ const rawBytes = new Uint8Array([1, 2, 3, 4]);
+
+ copyWithinGPU(source, 0, destination, 0, rawBytes.length);
+ expect(queue.submit).not.toHaveBeenCalled();
+
+ context.copyRawBytesToBuffer(rawBytes, destination, 4, rawBytes.length);
+
+ expect(queue.submit).toHaveBeenCalledTimes(1);
+ expect(queue.writeBuffer).toHaveBeenCalledTimes(1);
+ expect(events).toEqual(["copy", "finish", "submit", "writeBuffer"]);
+});
+
+test("a GPU readback flushes pending copies before its own submission", async
() => {
+ const {
+ context,
+ queue,
+ events,
+ encoders,
+ memory,
+ source,
+ destination,
+ } = createContext();
+ const copyWithinGPU = context.getDeviceAPI("deviceCopyWithinGPU");
+ const copyFromGPU = context.getDeviceAPI("deviceCopyFromGPU");
+
+ copyWithinGPU(source, 0, destination, 0, 16);
+ copyFromGPU(destination, 0, 128, 16);
+
+ expect(queue.submit).toHaveBeenCalledTimes(2);
+ expect(encoders).toHaveLength(2);
+ expect(encoders[0].commands).toEqual(["copy"]);
+ expect(encoders[1].commands).toEqual(["copy"]);
+ expect(events).toEqual(["copy", "finish", "submit", "copy", "finish",
"submit"]);
+
+ await context.sync();
+
+ expect(memory.storeRawBytes).toHaveBeenCalledTimes(1);
+ expect(memory.storeRawBytes.mock.calls[0][0]).toBe(128);
+ expect(memory.storeRawBytes.mock.calls[0][1]).toHaveLength(16);
+ expect(queue.onSubmittedWorkDone).not.toHaveBeenCalled();
+});
+
+test("buffer deallocation flushes pending copies before destroy", () => {
+ const { context, queue, events, source, destination } = createContext();
+ const copyWithinGPU = context.getDeviceAPI("deviceCopyWithinGPU");
+ const free = context.getDeviceAPI("deviceFreeDataSpace");
+
+ copyWithinGPU(source, 0, destination, 0, 16);
+ free(source);
+
+ expect(queue.submit).toHaveBeenCalledTimes(1);
+ expect(events).toEqual(["copy", "finish", "submit", "destroy"]);
+});
+
+test("drawing flushes pending copies first", () => {
+ const { context, queue, events, source, destination } = createContext();
+ const copyWithinGPU = context.getDeviceAPI("deviceCopyWithinGPU");
+ const canvasRenderManager = {
+ draw: jest.fn(() => events.push("draw")),
+ };
+ context.canvasRenderManager = canvasRenderManager;
+
+ copyWithinGPU(source, 0, destination, 0, 16);
+ context.drawImageFromBuffer(destination, 2, 2);
+
+ expect(queue.submit).toHaveBeenCalledTimes(1);
+ expect(canvasRenderManager.draw).toHaveBeenCalledTimes(1);
+ expect(events).toEqual(["copy", "finish", "submit", "draw"]);
+});
+
+test("sync awaits a readback and a later batched GPU copy", async () => {
+ const readback = createDeferred();
+ const queueDone = createDeferred();
+ const {
+ context,
+ queue,
+ memory,
+ source,
+ destination,
+ } = createContext({
+ mapAsync: () => readback.promise,
+ onSubmittedWorkDone: () => queueDone.promise,
+ });
+ const copyFromGPU = context.getDeviceAPI("deviceCopyFromGPU");
+ const copyWithinGPU = context.getDeviceAPI("deviceCopyWithinGPU");
+
+ copyFromGPU(source, 0, 128, 16);
+ copyWithinGPU(source, 0, destination, 0, 16);
+
+ let syncResolved = false;
+ const syncPromise = context.sync().then(() => {
+ syncResolved = true;
+ });
+
+ expect(queue.submit).toHaveBeenCalledTimes(2);
+ expect(queue.onSubmittedWorkDone).toHaveBeenCalledTimes(1);
+
+ queueDone.resolve();
+ await Promise.resolve();
+ expect(syncResolved).toBe(false);
+ expect(memory.storeRawBytes).not.toHaveBeenCalled();
+
+ readback.resolve();
+ await syncPromise;
+
+ expect(memory.storeRawBytes).toHaveBeenCalledTimes(1);
+ expect(memory.storeRawBytes.mock.calls[0][0]).toBe(128);
+ expect(memory.storeRawBytes.mock.calls[0][1]).toHaveLength(16);
+});
+
+test("a host write after a readback makes sync wait for the queue", async ()
=> {
+ const queueDone = createDeferred();
+ const {
+ context,
+ queue,
+ memory,
+ source,
+ destination,
+ } = createContext({
+ onSubmittedWorkDone: () => queueDone.promise,
+ });
+ const copyFromGPU = context.getDeviceAPI("deviceCopyFromGPU");
+
+ copyFromGPU(source, 0, 128, 16);
+ context.copyRawBytesToBuffer(
+ new Uint8Array([1, 2, 3, 4]),
+ destination,
+ 0,
+ 4
+ );
+
+ let syncResolved = false;
+ const syncPromise = context.sync().then(() => {
+ syncResolved = true;
+ });
+ expect(queue.onSubmittedWorkDone).toHaveBeenCalledTimes(1);
+
+ await Promise.resolve();
+ expect(syncResolved).toBe(false);
+
+ queueDone.resolve();
+ await syncPromise;
+
+ expect(memory.storeRawBytes).toHaveBeenCalledTimes(1);
+});
+
+test("sync propagates a pending readback failure", async () => {
+ const readError = new Error("mapAsync failed");
+ const {
+ context,
+ queue,
+ source,
+ destination,
+ } = createContext({
+ mapAsync: () => Promise.reject(readError),
+ });
+ const copyFromGPU = context.getDeviceAPI("deviceCopyFromGPU");
+ const copyWithinGPU = context.getDeviceAPI("deviceCopyWithinGPU");
+
+ copyFromGPU(source, 0, 128, 16);
+ copyWithinGPU(source, 0, destination, 0, 16);
+
+ await expect(context.sync()).rejects.toBe(readError);
+ expect(queue.onSubmittedWorkDone).toHaveBeenCalledTimes(1);
+});
diff --git a/web/jest.config.js b/web/tests/typescript_transformer.cjs
similarity index 72%
copy from web/jest.config.js
copy to web/tests/typescript_transformer.cjs
index 01ea70e9d9..f001dff94f 100644
--- a/web/jest.config.js
+++ b/web/tests/typescript_transformer.cjs
@@ -16,10 +16,18 @@
* specific language governing permissions and limitations
* under the License.
*/
-module.exports = {
- testEnvironment: "node",
+const ts = require("typescript");
- testMatch: [
- "**/tests/node/*.js"
- ],
+module.exports = {
+ process(sourceText, sourcePath) {
+ return {
+ code: ts.transpileModule(sourceText, {
+ compilerOptions: {
+ module: ts.ModuleKind.CommonJS,
+ target: ts.ScriptTarget.ES2018,
+ },
+ fileName: sourcePath,
+ }).outputText,
+ };
+ },
};
