refactor(engine): share computation dispatch setup (#2833)

Author: ibgreen-openai

modules/engine/src/compute/computation.ts

@@ -225,18 +225,7 @@ export class Computation {
   dispatch(computePass: ComputePass, x: number, y?: number, z?: number): void {
     try {
       this._logDrawCallStart();
-
-      // Check if the pipeline is invalidated
-      // TODO - this is likely the worst place to do this from performance perspective. Perhaps add a predraw()?
-      this.pipeline = this._updatePipeline();
-
-      // Set pipeline state, we may be sharing a pipeline so we need to set all state on every draw
-      // Any caching needs to be done inside the pipeline functions
-      this.pipeline.setBindings(this.bindings);
-      computePass.setPipeline(this.pipeline);
-      // @ts-expect-error
-      computePass.setBindings({});
-
+      this._setPipeline(computePass);
       computePass.dispatch(x, y, z);
     } finally {
       this._logDrawCallEnd();
@@ -247,19 +236,26 @@ export class Computation {
   dispatchIndirect(computePass: ComputePass, indirectBuffer: Buffer, indirectOffset = 0): void {
     try {
       this._logDrawCallStart();
-
-      this.pipeline = this._updatePipeline();
-      this.pipeline.setBindings(this.bindings);
-      computePass.setPipeline(this.pipeline);
-      // @ts-expect-error ComputePass implementations expose binding application internally.
-      computePass.setBindings({});
-
+      this._setPipeline(computePass);
       computePass.dispatchIndirect(indirectBuffer, indirectOffset);
     } finally {
       this._logDrawCallEnd();
     }
   }
 
+  private _setPipeline(computePass: ComputePass): void {
+    // Check if the pipeline is invalidated
+    // TODO - this is likely the worst place to do this from performance perspective. Perhaps add a predraw()?
+    this.pipeline = this._updatePipeline();
+
+    // Set pipeline state, we may be sharing a pipeline so we need to set all state on every draw
+    // Any caching needs to be done inside the pipeline functions
+    this.pipeline.setBindings(this.bindings);
+    computePass.setPipeline(this.pipeline);
+    // @ts-expect-error ComputePass implementations expose binding application internally.
+    computePass.setBindings({});
+  }
+
   // Update fixed fields (can trigger pipeline rebuild)
 
   // Update dynamic fields

modules/engine/test/compute/computation.spec.ts

@@ -93,7 +93,7 @@ test('Computation#dispatchIndirect', async t => {
       usage: Buffer.STORAGE | Buffer.COPY_SRC
     });
     const dispatchBuffer = webgpuDevice.createBuffer({
-      data: new Uint32Array([4, 1, 1]),
+      data: new Uint32Array([4, 1, 1, 1, 1, 1]),
       usage: Buffer.INDIRECT
     });
     computation.setBindings({data: workBuffer});
@@ -103,11 +103,43 @@ test('Computation#dispatchIndirect', async t => {
     computePass.end();
     webgpuDevice.submit();
 
-    const computedData = new Int32Array(await workBuffer.readAsync());
+    const computedBytes = await workBuffer.readAsync();
+    const computedData = new Int32Array(
+      computedBytes.buffer,
+      computedBytes.byteOffset,
+      computedBytes.byteLength / Int32Array.BYTES_PER_ELEMENT
+    );
     t.deepEqual(Array.from(computedData), [2, 4, 6, 8], 'GPU-written dimensions drive dispatch');
 
+    const offsetWorkBuffer = webgpuDevice.createBuffer({
+      data: new Int32Array([1, 2, 3, 4]),
+      usage: Buffer.STORAGE | Buffer.COPY_SRC
+    });
+    computation.setBindings({data: offsetWorkBuffer});
+    const offsetComputePass = webgpuDevice.beginComputePass({});
+    computation.dispatchIndirect(
+      offsetComputePass,
+      dispatchBuffer,
+      3 * Uint32Array.BYTES_PER_ELEMENT
+    );
+    offsetComputePass.end();
+    webgpuDevice.submit();
+
+    const offsetComputedBytes = await offsetWorkBuffer.readAsync();
+    const offsetComputedData = new Int32Array(
+      offsetComputedBytes.buffer,
+      offsetComputedBytes.byteOffset,
+      offsetComputedBytes.byteLength / Int32Array.BYTES_PER_ELEMENT
+    );
+    t.deepEqual(
+      Array.from(offsetComputedData),
+      [2, 2, 3, 4],
+      'nonzero byte offset selects the requested dispatch record'
+    );
+
     computation.destroy();
     workBuffer.destroy();
+    offsetWorkBuffer.destroy();
     dispatchBuffer.destroy();
   }
   t.end();