impr: Better withPerformanceCallback API (#2309)

Author: reczkok

packages/typegpu/src/core/pipeline/computePipeline.ts

@@ -220,7 +220,21 @@ class TgpuComputePipelineImpl implements TgpuComputePipeline {
   }
 
   withPerformanceCallback(callback: (start: bigint, end: bigint) => void | Promise<void>): this {
-    const newPriors = createWithPerformanceCallback(this.#priors, callback, this.#core.root);
+    if (this.#priors.timestampWrites) {
+      return new TgpuComputePipelineImpl(this.#core, {
+        ...this.#priors,
+        performanceCallback: callback,
+      }) as this;
+    }
+
+    const querySet = this.#core.performanceCallbackQuerySet;
+    if (!querySet) {
+      console.warn(
+        'Performance callback cannot be used because the timestamp-query feature is not enabled on the root.',
+      );
+      return this;
+    }
+    const newPriors = createWithPerformanceCallback(this.#priors, callback, querySet);
     return new TgpuComputePipelineImpl(this.#core, newPriors) as this;
   }
 
@@ -353,6 +367,7 @@ class ComputePipelineCore implements SelfResolvable {
 
   #slotBindings: [TgpuSlot<unknown>, unknown][];
   #descriptor: TgpuComputePipeline.Descriptor;
+  #performanceCallbackQuerySet: TgpuQuerySet<'timestamp'> | undefined;
 
   constructor(
     root: ExperimentalTgpuRoot,
@@ -375,6 +390,13 @@ class ComputePipelineCore implements SelfResolvable {
     return 'computePipelineCore';
   }
 
+  get performanceCallbackQuerySet() {
+    if (!this.root.enabledFeatures.has('timestamp-query')) {
+      return undefined;
+    }
+    return (this.#performanceCallbackQuerySet ??= this.root.createQuerySet('timestamp', 2));
+  }
+
   public unwrap(): Memo {
     if (this._memo === undefined) {
       const device = this.root.device;

packages/typegpu/src/core/pipeline/renderPipeline.ts

@@ -550,11 +550,22 @@ class TgpuRenderPipelineImpl implements TgpuRenderPipeline {
 
   withPerformanceCallback(callback: (start: bigint, end: bigint) => void | Promise<void>): this {
     const internals = this[$internal];
-    const newPriors = createWithPerformanceCallback(
-      internals.priors,
-      callback,
-      internals.core.options.root,
-    );
+
+    if (internals.priors.timestampWrites) {
+      return new TgpuRenderPipelineImpl(internals.core, {
+        ...internals.priors,
+        performanceCallback: callback,
+      }) as this;
+    }
+
+    const querySet = internals.core.performanceCallbackQuerySet;
+    if (!querySet) {
+      console.warn(
+        'Performance callback cannot be used because the timestamp-query feature is not enabled on the root.',
+      );
+      return this;
+    }
+    const newPriors = createWithPerformanceCallback(internals.priors, callback, querySet);
     return new TgpuRenderPipelineImpl(internals.core, newPriors) as this;
   }
 
@@ -955,6 +966,7 @@ class RenderPipelineCore implements SelfResolvable {
 
   #latestAutoVertexIn: TgpuVertexFn.In | undefined;
   #latestAutoFragmentOut: BaseData | undefined;
+  #performanceCallbackQuerySet: TgpuQuerySet<'timestamp'> | undefined;
 
   constructor(options: RenderPipelineCoreOptions) {
     this.options = options;
@@ -1011,6 +1023,13 @@ class RenderPipelineCore implements SelfResolvable {
     return 'renderPipelineCore';
   }
 
+  get performanceCallbackQuerySet() {
+    if (!this.options.root.enabledFeatures.has('timestamp-query')) {
+      return undefined;
+    }
+    return (this.#performanceCallbackQuerySet ??= this.options.root.createQuerySet('timestamp', 2));
+  }
+
   public unwrap(): Memo {
     if (this._memo !== undefined) {
       return this._memo;

packages/typegpu/src/core/pipeline/timeable.ts

@@ -19,27 +19,19 @@ export type TimestampWritesPriors = {
     endOfPassWriteIndex?: number;
   };
   readonly performanceCallback?: (start: bigint, end: bigint) => void | Promise<void>;
-  readonly hasAutoQuerySet?: boolean;
 };
 
 export function createWithPerformanceCallback<T extends TimestampWritesPriors>(
   currentPriors: T,
   callback: (start: bigint, end: bigint) => void | Promise<void>,
-  root: ExperimentalTgpuRoot,
+  querySet: TgpuQuerySet<'timestamp'>,
 ): T {
-  if (!root.enabledFeatures.has('timestamp-query')) {
-    throw new Error(
-      'Performance callback requires the "timestamp-query" feature to be enabled on GPU device.',
-    );
-  }
-
   if (!currentPriors.timestampWrites) {
     return {
       ...currentPriors,
       performanceCallback: callback,
-      hasAutoQuerySet: true,
       timestampWrites: {
-        querySet: root.createQuerySet('timestamp', 2),
+        querySet,
         beginningOfPassWriteIndex: 0,
         endOfPassWriteIndex: 1,
       },
@@ -67,10 +59,6 @@ export function createWithTimestampWrites<T extends TimestampWritesPriors>(
     );
   }
 
-  if (currentPriors.hasAutoQuerySet && currentPriors.timestampWrites) {
-    currentPriors.timestampWrites.querySet.destroy();
-  }
-
   const timestampWrites: TimestampWritesPriors['timestampWrites'] = {
     querySet: options.querySet,
   };
@@ -84,7 +72,6 @@ export function createWithTimestampWrites<T extends TimestampWritesPriors>(
 
   return {
     ...currentPriors,
-    hasAutoQuerySet: false,
     timestampWrites,
   };
 }

packages/typegpu/tests/computePipeline.test.ts

@@ -141,23 +141,24 @@ describe('TgpuComputePipeline', () => {
       expect(pipeline[$internal].priors.performanceCallback).not.toBe(callback1);
     });
 
-    it('should throw error if timestamp-query feature is not enabled', ({ root, device }) => {
-      const originalFeatures = device.features;
-      //@ts-expect-error
+    it('should warn if timestamp-query feature is not enabled', ({ root, device }) => {
+      // @ts-expect-error
       device.features = new Set();
+      using consoleWarnSpy = vi.spyOn(console, 'warn').mockImplementation(() => {});
 
       const entryFn = tgpu.computeFn({ workgroupSize: [1] })(() => {});
 
       const callback = vi.fn();
 
       expect(() => {
-        root.createComputePipeline({ compute: entryFn }).withPerformanceCallback(callback);
-      }).toThrow(
-        'Performance callback requires the "timestamp-query" feature to be enabled on GPU device.',
+        const before = root.createComputePipeline({ compute: entryFn });
+        const after = before.withPerformanceCallback(callback);
+        // no-op
+        expect(after).toBe(before);
+      }).not.toThrow();
+      expect(consoleWarnSpy).toHaveBeenCalledWith(
+        'Performance callback cannot be used because the timestamp-query feature is not enabled on the root.',
       );
-
-      //@ts-expect-error
-      device.features = originalFeatures;
     });
   });
 
@@ -331,6 +332,7 @@ describe('TgpuComputePipeline', () => {
   describe('Combined Performance callback and Timestamp Writes', () => {
     it('should work with both performance callback and custom timestamp writes', ({
       root,
+      device,
       commandEncoder,
     }) => {
       const entryFn = tgpu.computeFn({ workgroupSize: [1] })(() => {});
@@ -371,6 +373,12 @@ describe('TgpuComputePipeline', () => {
         querySet[$internal].resolveBuffer,
         0,
       );
+
+      expect(device.mock.createQuerySet).toHaveBeenCalledTimes(1);
+      expect(device.mock.createQuerySet).toHaveBeenCalledWith({
+        type: 'timestamp',
+        count: 10,
+      });
     });
 
     it('should prioritize custom timestamp writes over automatic ones', ({
@@ -394,8 +402,6 @@ describe('TgpuComputePipeline', () => {
         endOfPassWriteIndex: 5,
       });
 
-      expect((autoQuerySet as TgpuQuerySet<'timestamp'>).destroyed).toBe(true);
-
       const priors = pipeline[$internal].priors;
       expect(priors.performanceCallback).toBe(callback);
       expect(priors.timestampWrites?.querySet).toBe(querySet);

packages/typegpu/tests/renderPipeline.test.ts

@@ -501,10 +501,10 @@ describe('root.withVertex(...).withFragment(...)', () => {
       expect(pipeline[$internal].priors.performanceCallback).not.toBe(callback1);
     });
 
-    it('should throw error if timestamp-query feature is not enabled', ({ root, device }) => {
-      const originalFeatures = device.features;
+    it('should warn if timestamp-query feature is not enabled', ({ root, device }) => {
       //@ts-expect-error
       device.features = new Set();
+      using consoleWarnSpy = vi.spyOn(console, 'warn').mockImplementation(() => {});
 
       const vertexFn = tgpu.vertexFn({
         out: { pos: d.builtin.position },
@@ -517,17 +517,17 @@ describe('root.withVertex(...).withFragment(...)', () => {
       const callback = vi.fn();
 
       expect(() => {
-        root
-          .withVertex(vertexFn, {})
-          .withFragment(fragmentFn, { color: { format: 'rgba8unorm' } })
-          .createPipeline()
-          .withPerformanceCallback(callback);
-      }).toThrow(
-        'Performance callback requires the "timestamp-query" feature to be enabled on GPU device.',
+        const before = root.createRenderPipeline({
+          vertex: vertexFn,
+          fragment: fragmentFn,
+        });
+        const after = before.withPerformanceCallback(callback);
+        // no-op
+        expect(after).toBe(before);
+      }).not.toThrow();
+      expect(consoleWarnSpy).toHaveBeenCalledWith(
+        'Performance callback cannot be used because the timestamp-query feature is not enabled on the root.',
       );
-
-      //@ts-expect-error
-      device.features = originalFeatures;
     });
 
     it("should not throw 'A color target was not provided to the shader'", ({ root, device }) => {