guan404ming commented on code in PR #20059:
URL: https://github.com/apache/tvm/pull/20059#discussion_r3669449614
##########
web/src/webgpu.ts:
##########
@@ -470,7 +469,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
+ // A command submission is now the last queue operation, so the
// GPU→CPU copy fast path in sync() is no longer valid.
this.pendingGPUToCPUCopy = null;
Review Comment:
Nice cleanup! One catch here: clearing `pendingGPUToCPUCopy` drops the
outstanding `mapAsync` promise, so `sync()` falls back to
`onSubmittedWorkDone()` and no longer awaits the readback.
Repro: `deviceCopyFromGPU` -> `deviceCopyWithinGPU` -> `sync()`. Before this
PR the copy was submitted directly and the fast path stayed valid, so
`storeRawBytes` had run by the time `sync()` resolved. Now the flush nulls it
out and `sync()` can return before host memory is written. I checked this
against a mock device where `mapAsync` resolves on a later task: it passes on
`main` and fails on this branch.
Could we await both instead of discarding one? Something like:
```ts
async sync(): Promise<void> {
const pendingRead = this.pendingGPUToCPUCopy;
this.pendingGPUToCPUCopy = null;
const flushed = this.pendingEncoder !== null;
this.flushCommands();
if (pendingRead) await pendingRead;
if (flushed || !pendingRead) await this.device.queue.onSubmittedWorkDone();
}
```
(The same promise drop already exists on the compute-dispatch path, so it
may be worth fixing both together.)
##########
web/tests/node/test_webgpu.js:
##########
@@ -0,0 +1,205 @@
+/*
+ * 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 createMockDevice() {
+ const events = [];
+ const encoders = [];
+
+ const queue = {
+ submit: jest.fn(() => events.push("submit")),
+ writeBuffer: jest.fn(() => events.push("writeBuffer")),
+ onSubmittedWorkDone: jest.fn(() => Promise.resolve()),
+ };
+
+ 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(() => Promise.resolve()),
+ 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() {
+ const gpu = createMockDevice();
+ 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"]);
+});
Review Comment:
Would you mind adding one more case that interleaves `deviceCopyFromGPU`,
`deviceCopyWithinGPU`, then `sync()`? That is the sequence where the
`pendingGPUToCPUCopy` fast path interacts with the newly batched copies, and it
is currently uncovered.
##########
web/src/webgpu.ts:
##########
@@ -455,13 +455,12 @@ 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)
* - Queue sync (sync)
Review Comment:
Minor: worth adding `drawImageFromBuffer` to this list too, since it submits
its own encoder and needs a flush point as well.
##########
web/jest.config.js:
##########
@@ -19,6 +19,10 @@
module.exports = {
testEnvironment: "node",
+ transform: {
Review Comment:
Heads-up: `transform` replaces Jest's default map rather than merging into
it, so `.js` files no longer go through `babel-jest`. The existing tests are
plain CommonJS so nothing breaks today, but adding an explicit `"^.+\\.js$":
"babel-jest"` entry would keep this from biting later.
##########
web/tests/node/test_webgpu.js:
##########
@@ -0,0 +1,205 @@
+/*
+ * 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");
Review Comment:
Small note: the other `tests/node/*.js` files import the built bundle from
`../../dist`. Importing `src` directly makes sense given `WebGPUContext` is not
exported from `index.ts`, but it might be worth calling that out in the PR
description so the new transformer does not look surprising.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
