Add timestamp query support to compute boids (#323)

* Add timestamp query support to compute boids

* Create buffers each frame

* Fix lint

* Address kai's feedback

* 100-frame average in microseconds; add comments

* lint

---------

Co-authored-by: Kai Ninomiya <[email protected]>

Author: beaufortfrancois

package-lock.json

@@ -39,7 +39,7 @@
     "@types/stats.js": "^0.17.0",
     "@typescript-eslint/eslint-plugin": "^5.41.0",
     "@typescript-eslint/parser": "^5.41.0",
-    "@webgpu/types": "^0.1.21",
+    "@webgpu/types": "^0.1.38",
     "eslint": "^8.26.0",
     "eslint-config-prettier": "^8.5.0",
     "eslint-plugin-prettier": "^4.2.1",

package.json

@@ -6,7 +6,27 @@ import updateSpritesWGSL from './updateSprites.wgsl';
 const init: SampleInit = async ({ canvas, pageState, gui }) => {
   const adapter = await navigator.gpu.requestAdapter();
   assert(adapter, 'requestAdapter returned null');
-  const device = await adapter.requestDevice();
+
+  const hasTimestampQuery = adapter.features.has('timestamp-query');
+  const device = await adapter.requestDevice({
+    requiredFeatures: hasTimestampQuery ? ['timestamp-query'] : [],
+  });
+
+  const perfDisplayContainer = document.createElement('div');
+  perfDisplayContainer.style.color = 'white';
+  perfDisplayContainer.style.background = 'black';
+  perfDisplayContainer.style.position = 'absolute';
+  perfDisplayContainer.style.top = '10px';
+  perfDisplayContainer.style.left = '10px';
+  perfDisplayContainer.style.textAlign = 'left';
+
+  const perfDisplay = document.createElement('pre');
+  perfDisplayContainer.appendChild(perfDisplay);
+  if (canvas.parentNode) {
+    canvas.parentNode.appendChild(perfDisplayContainer);
+  } else {
+    console.error('canvas.parentNode is null');
+  }
 
   if (!pageState.active) return;
   const context = canvas.getContext('webgpu') as GPUCanvasContext;
@@ -86,7 +106,7 @@ const init: SampleInit = async ({ canvas, pageState, gui }) => {
     },
   });
 
-  const renderPassDescriptor = {
+  const renderPassDescriptor: GPURenderPassDescriptor = {
     colorAttachments: [
       {
         view: undefined as GPUTextureView, // Assigned later
@@ -97,6 +117,38 @@ const init: SampleInit = async ({ canvas, pageState, gui }) => {
     ],
   };
 
+  const computePassDescriptor: GPUComputePassDescriptor = {};
+
+  /** Storage for timestamp query results */
+  let querySet: GPUQuerySet | undefined = undefined;
+  /** Timestamps are resolved into this buffer */
+  let resolveBuffer: GPUBuffer | undefined = undefined;
+  /** Pool of spare buffers for MAP_READing the timestamps back to CPU. A buffer
+   * is taken from the pool (if available) when a readback is needed, and placed
+   * back into the pool once the readback is done and it's unmapped. */
+  const spareResultBuffers = [];
+
+  if (hasTimestampQuery) {
+    querySet = device.createQuerySet({
+      type: 'timestamp',
+      count: 4,
+    });
+    resolveBuffer = device.createBuffer({
+      size: 4 * BigInt64Array.BYTES_PER_ELEMENT,
+      usage: GPUBufferUsage.QUERY_RESOLVE | GPUBufferUsage.COPY_SRC,
+    });
+    computePassDescriptor.timestampWrites = {
+      querySet,
+      beginningOfPassWriteIndex: 0,
+      endOfPassWriteIndex: 1,
+    };
+    renderPassDescriptor.timestampWrites = {
+      querySet,
+      beginningOfPassWriteIndex: 2,
+      endOfPassWriteIndex: 3,
+    };
+  }
+
   // prettier-ignore
   const vertexBufferData = new Float32Array([
     -0.01, -0.02, 0.01,
@@ -207,6 +259,8 @@ const init: SampleInit = async ({ canvas, pageState, gui }) => {
   }
 
   let t = 0;
+  let computePassDurationSum = 0;
+  let renderPassDurationSum = 0;
   function frame() {
     // Sample is no longer the active page.
     if (!pageState.active) return;
@@ -217,7 +271,9 @@ const init: SampleInit = async ({ canvas, pageState, gui }) => {
 
     const commandEncoder = device.createCommandEncoder();
     {
-      const passEncoder = commandEncoder.beginComputePass();
+      const passEncoder = commandEncoder.beginComputePass(
+        computePassDescriptor
+      );
       passEncoder.setPipeline(computePipeline);
       passEncoder.setBindGroup(0, particleBindGroups[t % 2]);
       passEncoder.dispatchWorkgroups(Math.ceil(numParticles / 64));
@@ -231,8 +287,54 @@ const init: SampleInit = async ({ canvas, pageState, gui }) => {
       passEncoder.draw(3, numParticles, 0, 0);
       passEncoder.end();
     }
+
+    let resultBuffer: GPUBuffer | undefined = undefined;
+    if (hasTimestampQuery) {
+      resultBuffer =
+        spareResultBuffers.pop() ||
+        device.createBuffer({
+          size: 4 * BigInt64Array.BYTES_PER_ELEMENT,
+          usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ,
+        });
+      commandEncoder.resolveQuerySet(querySet, 0, 4, resolveBuffer, 0);
+      commandEncoder.copyBufferToBuffer(
+        resolveBuffer,
+        0,
+        resultBuffer,
+        0,
+        resultBuffer.size
+      );
+    }
+
     device.queue.submit([commandEncoder.finish()]);
 
+    if (hasTimestampQuery) {
+      resultBuffer.mapAsync(GPUMapMode.READ).then(() => {
+        const times = new BigInt64Array(resultBuffer.getMappedRange());
+        computePassDurationSum += Number(times[1] - times[0]);
+        renderPassDurationSum += Number(times[3] - times[2]);
+        resultBuffer.unmap();
+
+        // Periodically update the text for the timer stats
+        const kNumTimerSamples = 100;
+        if (t % kNumTimerSamples === 0) {
+          const avgComputeMicroseconds = Math.round(
+            computePassDurationSum / kNumTimerSamples / 1000
+          );
+          const avgRenderMicroseconds = Math.round(
+            renderPassDurationSum / kNumTimerSamples / 1000
+          );
+          perfDisplay.textContent = `\
+avg compute pass duration: ${avgComputeMicroseconds}µs
+avg render pass duration: ${avgRenderMicroseconds}µs
+spare readback buffers: ${spareResultBuffers.length}`;
+          computePassDurationSum = 0;
+          renderPassDurationSum = 0;
+        }
+        spareResultBuffers.push(resultBuffer);
+      });
+    }
+
     ++t;
     requestAnimationFrame(frame);
   }