From 3e7473cec7b954b12b744357d68ab3aa0c0af274 Mon Sep 17 00:00:00 2001
From: John Owens <jowens@ece.ucdavis.edu>
Date: Tue, 8 Oct 2024 15:30:40 -0700
Subject: [PATCH] from Brandon, checking in so I can source from github

---
 timestamp-helper.ts | 175 ++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 175 insertions(+)
 create mode 100644 timestamp-helper.ts

diff --git a/timestamp-helper.ts b/timestamp-helper.ts
new file mode 100644
index 0000000..b34631a
--- /dev/null
+++ b/timestamp-helper.ts
@@ -0,0 +1,175 @@
+/** Number of timing samples to collect before computing the average */
+const AVG_SAMPLE_COUNT = 30;
+
+interface TimestampWrites {
+  querySet: GPUQuerySet,
+  beginningOfPassWriteIndex: number,
+  endOfPassWriteIndex: number,
+};
+
+interface TimestampResults {
+  [Key: string]: number;
+}
+
+/**
+ * A utility that makes it a bit easier to gather timings from WebGPU compute and render passes.
+ */
+export class TimestampHelper {
+  device: GPUDevice;
+
+  #timestampsSupported = false;
+  #timestampQuerySet: GPUQuerySet;
+  #timestampResolveBuffer: GPUBuffer;
+  #timestampReadbackBuffers: GPUBuffer[] = [];
+  #readbackBufferCount = 0;
+  #currentReadbackBuffer?: GPUBuffer = null;
+
+  #passTimings = new Map();
+  #maxPassCount = 0;
+  #nextQueryIndex = 0;
+  #queriesUsed = new Map();
+
+  #averages: TimestampResults = {};
+
+  constructor(device: GPUDevice, maxPassCount = 16) {
+    this.device = device;
+    this.#maxPassCount = maxPassCount;
+    this.#timestampsSupported = this.device.features.has('timestamp-query');
+
+    if (this.#timestampsSupported) {
+      this.#timestampQuerySet = this.device.createQuerySet({
+          label: 'Timestamp Helper',
+          type: 'timestamp',
+          count: this.#maxPassCount,
+      });
+
+      this.#timestampResolveBuffer = this.device.createBuffer({
+          size: BigUint64Array.BYTES_PER_ELEMENT * this.#maxPassCount,
+          usage: GPUBufferUsage.QUERY_RESOLVE | GPUBufferUsage.COPY_SRC,
+      });
+    }
+  }
+
+  get timestampsSupported() {
+    return this.#timestampsSupported;
+  }
+
+  timestampWrites(name: string): TimestampWrites {
+    if (!this.#timestampsSupported) { return undefined; }
+
+    if (this.#currentReadbackBuffer) {
+      throw new Error('Must read back the previous resolve before new timestamps can be added.');
+    }
+
+    if (this.#nextQueryIndex >= this.#maxPassCount * 2) {
+      throw new Error('Exceeded the number of passes that can be queried in a single resolve.');
+    }
+
+    const timestampWrites = {
+      querySet: this.#timestampQuerySet,
+      beginningOfPassWriteIndex: this.#nextQueryIndex,
+      endOfPassWriteIndex: this.#nextQueryIndex + 1
+    };
+
+    this.#queriesUsed.set(name, {
+      begin: timestampWrites.beginningOfPassWriteIndex,
+      end: timestampWrites.endOfPassWriteIndex,
+    });
+    this.#nextQueryIndex += 2;
+    return timestampWrites;
+  }
+
+  resolve(commandEncoder: GPUCommandEncoder) {
+    if (!this.#timestampsSupported) { return; }
+
+    if (this.#currentReadbackBuffer) {
+      throw new Error('Must read back the previous resolve before resolve can be called again.');
+    }
+
+    if (this.#timestampReadbackBuffers.length > 0) {
+      this.#currentReadbackBuffer = this.#timestampReadbackBuffers.pop();
+    } else {
+      this.#currentReadbackBuffer = this.device.createBuffer({
+        label: `Timestamp Readback ${this.#readbackBufferCount}`,
+        size: this.#timestampResolveBuffer.size,
+        usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
+      });
+      this.#readbackBufferCount++;
+    }
+
+    commandEncoder.resolveQuerySet(this.#timestampQuerySet, 0, this.#nextQueryIndex, this.#timestampResolveBuffer, 0);
+    commandEncoder.copyBufferToBuffer(this.#timestampResolveBuffer, 0, this.#currentReadbackBuffer, 0, this.#timestampResolveBuffer.size);
+  }
+
+  async read(): Promise<TimestampResults> {
+    if (!this.#currentReadbackBuffer) { return; }
+
+    let readbackBuffer = this.#currentReadbackBuffer;
+    let queries = new Map(this.#queriesUsed);
+
+    this.#currentReadbackBuffer = null;
+    this.#queriesUsed.clear();
+    this.#nextQueryIndex = 0;
+
+    await readbackBuffer.mapAsync(GPUMapMode.READ);
+    const mappedArray = new BigUint64Array(readbackBuffer.getMappedRange());
+
+    const results: TimestampResults = {};
+
+    for (const [ name, query ] of queries.entries()) {
+      const passTime = Number(mappedArray[query.end] - mappedArray[query.begin]);
+      // Discard negative times
+      if (passTime >= 0) {
+        let passTimings = this.#passTimings.get(name);
+        if (!passTimings) {
+          passTimings = {
+            values: new Array(AVG_SAMPLE_COUNT),
+            index: 0,
+          };
+          this.#passTimings.set(name, passTimings);
+        }
+        // Storing pass timings in µs
+        let passTimeMicro = passTime / 1000;
+        passTimings.values[passTimings.index++ % AVG_SAMPLE_COUNT] = passTimeMicro;
+        if (passTimings.index % AVG_SAMPLE_COUNT == 0) {
+          // Update the average
+          let avg = 0;
+          for (const value of passTimings.values) {
+            avg += value;
+          }
+          this.#averages[name] = avg / AVG_SAMPLE_COUNT;
+        }
+        results[name] = passTimeMicro;
+      }
+    }
+
+    // Update any existing keys not in the queries to 0. This way if one of them
+    // isn't present in the readback it will zero out instead of appearing to
+    // stay steady at the last reading.
+    for (const name of this.#passTimings.keys()) {
+      if (!queries.has(name)) {
+        let passTimings = this.#passTimings.get(name);
+        passTimings.values[passTimings.index++ % AVG_SAMPLE_COUNT] = 0;
+        if (passTimings.index % AVG_SAMPLE_COUNT == 0) {
+          // Update the average
+          let avg = 0;
+          for (const value of passTimings.values) {
+            avg += value;
+          }
+          this.#averages[name] = avg / AVG_SAMPLE_COUNT;
+        }
+      }
+    }
+
+    readbackBuffer.unmap();
+    this.#timestampReadbackBuffers.push(readbackBuffer);
+
+    return results;
+  }
+
+  hasPendingResults(): boolean { return this.#nextQueryIndex != 0; }
+
+  get averages(): TimestampResults {
+    return this.#averages;
+  }
+}
\ No newline at end of file