diff --git a/packages/typegpu/src/core/root/init.ts b/packages/typegpu/src/core/root/init.ts index bc84ad408b..0ef85a27c9 100644 --- a/packages/typegpu/src/core/root/init.ts +++ b/packages/typegpu/src/core/root/init.ts @@ -152,6 +152,30 @@ export class TgpuGuardedComputePipelineImpl< ); } + withPerformanceCallback( + callback: (start: bigint, end: bigint) => void | Promise, + ): TgpuGuardedComputePipeline { + return new TgpuGuardedComputePipelineImpl( + this.#root, + this.#pipeline.withPerformanceCallback(callback), + this.#sizeUniform, + this.#workgroupSize, + ); + } + + withTimestampWrites(options: { + querySet: TgpuQuerySet<'timestamp'> | GPUQuerySet; + beginningOfPassWriteIndex?: number; + endOfPassWriteIndex?: number; + }): TgpuGuardedComputePipeline { + return new TgpuGuardedComputePipelineImpl( + this.#root, + this.#pipeline.withTimestampWrites(options), + this.#sizeUniform, + this.#workgroupSize, + ); + } + dispatchThreads(...threads: TArgs): void { const sanitizedSize = toVec3(threads); const workgroupCount = ceil(vec3f(sanitizedSize).div(vec3f(this.#workgroupSize))); diff --git a/packages/typegpu/src/core/root/rootTypes.ts b/packages/typegpu/src/core/root/rootTypes.ts index e6c5c8bf66..afa6cb08df 100644 --- a/packages/typegpu/src/core/root/rootTypes.ts +++ b/packages/typegpu/src/core/root/rootTypes.ts @@ -1,6 +1,7 @@ import type { AnyComputeBuiltin, AnyFragmentInputBuiltin, OmitBuiltins } from '../../builtin.ts'; import type { TgpuQuerySet } from '../../core/querySet/querySet.ts'; import type { AnyData, Disarray, UndecorateRecord } from '../../data/dataTypes.ts'; +import type { InstanceToSchema } from '../../data/instanceToSchema.ts'; import type { WgslComparisonSamplerProps, WgslSamplerProps } from '../../data/sampler.ts'; import type { AnyWgslData, @@ -12,6 +13,7 @@ import type { Void, WgslArray, } from '../../data/wgslTypes.ts'; +import type { TgpuNamable } from '../../shared/meta.ts'; import type { ExtractInvalidSchemaError, InferGPURecord, @@ -33,7 +35,13 @@ import type { ShaderGenerator } from '../../tgsl/shaderGenerator.ts'; import type { Unwrapper } from '../../unwrapper.ts'; import type { TgpuBuffer, VertexFlag } from '../buffer/buffer.ts'; import type { TgpuMutable, TgpuReadonly, TgpuUniform } from '../buffer/bufferShorthand.ts'; -import type { TgpuFixedComparisonSampler, TgpuFixedSampler } from '../sampler/sampler.ts'; +import type { + AnyAutoCustoms, + AutoFragmentIn, + AutoFragmentOut, + AutoVertexIn, + AutoVertexOut, +} from '../function/autoIO.ts'; import type { IORecord } from '../function/fnTypes.ts'; import type { FragmentInConstrained, @@ -44,6 +52,7 @@ import type { import type { TgpuVertexFn } from '../function/tgpuVertexFn.ts'; import type { TgpuComputePipeline } from '../pipeline/computePipeline.ts'; import type { FragmentOutToTargets, TgpuRenderPipeline } from '../pipeline/renderPipeline.ts'; +import type { TgpuFixedComparisonSampler, TgpuFixedSampler } from '../sampler/sampler.ts'; import type { Eventual, TgpuAccessor, TgpuMutableAccessor, TgpuSlot } from '../slot/slotTypes.ts'; import type { TgpuTexture } from '../texture/texture.ts'; import type { @@ -52,15 +61,6 @@ import type { } from '../vertexLayout/vertexAttribute.ts'; import type { TgpuVertexLayout } from '../vertexLayout/vertexLayout.ts'; import type { TgpuComputeFn } from './../function/tgpuComputeFn.ts'; -import type { TgpuNamable } from '../../shared/meta.ts'; -import type { - AnyAutoCustoms, - AutoFragmentIn, - AutoFragmentOut, - AutoVertexIn, - AutoVertexOut, -} from '../function/autoIO.ts'; -import type { InstanceToSchema } from '../../data/instanceToSchema.ts'; // ---------- // Public API @@ -80,6 +80,24 @@ export interface TgpuGuardedComputePipeline e */ with(encoder: GPUCommandEncoder): TgpuGuardedComputePipeline; + /** + * Returns a pipeline wrapper with the given performance callback attached. + * Analogous to `TgpuComputePipeline.withPerformanceCallback(callback)`. + */ + withPerformanceCallback( + callback: (start: bigint, end: bigint) => void | Promise, + ): TgpuGuardedComputePipeline; + + /** + * Returns a pipeline wrapper with the given timestamp writes configuration. + * Analogous to `TgpuComputePipeline.withTimestampWrites(options)`. + */ + withTimestampWrites(options: { + querySet: TgpuQuerySet<'timestamp'> | GPUQuerySet; + beginningOfPassWriteIndex?: number; + endOfPassWriteIndex?: number; + }): TgpuGuardedComputePipeline; + /** * Dispatches the pipeline. * Unlike `TgpuComputePipeline.dispatchWorkgroups()`, this method takes in the @@ -378,7 +396,7 @@ export interface WithBinding extends Withable { /** * Creates a compute pipeline that executes the given callback in an exact number of threads. - * This is different from `withCompute(...).createPipeline()` in that it does a bounds check on the + * This is different from `createComputePipeline()` in that it does a bounds check on the * thread id, where as regular pipelines do not and work in units of workgroups. * * @param callback A function converted to WGSL and executed on the GPU. diff --git a/packages/typegpu/tests/guardedComputePipeline.test.ts b/packages/typegpu/tests/guardedComputePipeline.test.ts index 065101fc82..58484d5e2e 100644 --- a/packages/typegpu/tests/guardedComputePipeline.test.ts +++ b/packages/typegpu/tests/guardedComputePipeline.test.ts @@ -1,4 +1,4 @@ -import { describe, expect } from 'vitest'; +import { describe, expect, vi } from 'vitest'; import { it } from 'typegpu-testing-utility'; import { getName } from '../src/shared/meta.ts'; import { bindGroupLayout } from '../src/tgpuBindGroupLayout.ts'; @@ -31,4 +31,34 @@ describe('TgpuGuardedComputePipeline', () => { expect(getName(pipeline)).toBe('myPipeline'); expect(getName(pipeline.pipeline)).toBe('myPipeline'); }); + + it('delegates `withPerformanceCallback` to the underlying pipeline', ({ root }) => { + const callback = vi.fn(); + const guarded = root.createGuardedComputePipeline(() => { + 'use gpu'; + }); + + const spy = vi.spyOn(guarded.pipeline, 'withPerformanceCallback'); + guarded.withPerformanceCallback(callback); + + expect(spy).toHaveBeenCalledWith(callback); + }); + + it('delegates `withTimestampWrites` to the underlying pipeline', ({ root }) => { + const querySet = root.createQuerySet('timestamp', 2); + const guarded = root.createGuardedComputePipeline(() => { + 'use gpu'; + }); + + const options = { + querySet, + beginningOfPassWriteIndex: 0, + endOfPassWriteIndex: 1, + }; + + const spy = vi.spyOn(guarded.pipeline, 'withTimestampWrites'); + guarded.withTimestampWrites(options); + + expect(spy).toHaveBeenCalledWith(options); + }); });