[Web] Batch GPU-to-GPU copies, fix WebGPU synchronization, and add tests for command batching

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"
   ],

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,24 +453,22 @@ 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 {
     if (this.pendingEncoder) {
       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 {