diff --git a/apps/typegpu-docs/src/content/docs/apis/pipelines.mdx b/apps/typegpu-docs/src/content/docs/apis/pipelines.mdx index 00a1926688..997b5242e7 100644 --- a/apps/typegpu-docs/src/content/docs/apis/pipelines.mdx +++ b/apps/typegpu-docs/src/content/docs/apis/pipelines.mdx @@ -242,6 +242,64 @@ tgpu.resolve([innerPipeline]); ``` ::: +## Initialization + +Pipeline initialization involves resolving the pipeline code, creating the shader module, and creating the underlying WebGPU pipeline. +This happens automatically the first time the pipeline is executed via `.draw`, `.dispatchWorkgroups`, or a similar method. +The `initSync` method lets you start the initialization early. +To wait until initialization actually finishes on the device (fully avoiding a stall on first execution), use `initAsync` instead. + +```ts twoslash +import tgpu from 'typegpu'; +const root = await tgpu.init(); +const mainCompute = tgpu.computeFn({ workgroupSize: [1] })(() => {}); +const pipeline = root.createComputePipeline({ + compute: mainCompute, +}); +// ---cut--- +// Automatically calls `pipeline.initSync();`, +// then enqueues a dispatch. +pipeline.dispatchWorkgroups(1); +``` + +```ts twoslash +import tgpu from 'typegpu'; +const root = await tgpu.init(); +const mainCompute = tgpu.computeFn({ workgroupSize: [1] })(() => {}); +const pipeline = root.createComputePipeline({ + compute: mainCompute, +}); +// ---cut--- +// If not already initialized, runs JS initialization, +// and issues pipeline initialization steps on the device. +pipeline.initSync(); + +// Enqueues a dispatch. +// Stall partially avoided because the pipeline is already resolved. +pipeline.dispatchWorkgroups(1); +``` + +```ts twoslash +import tgpu from 'typegpu'; +const root = await tgpu.init(); +const mainCompute = tgpu.computeFn({ workgroupSize: [1] })(() => {}); +const pipeline = root.createComputePipeline({ + compute: mainCompute, +}); +// ---cut--- +// If not already initialized, runs JS initialization, +// and issues pipeline initialization steps on the device. +const initPromise = pipeline.initAsync(); + +// Awaits until the initialization finishes. +await initPromise; + +// Enqueues a dispatch. +// Stall avoided, the pipeline is already resolved and initialized on the device. +pipeline.dispatchWorkgroups(1); +``` + + ## Execution ```ts diff --git a/apps/typegpu-docs/src/examples/algorithms/matrix-next/index.ts b/apps/typegpu-docs/src/examples/algorithms/matrix-next/index.ts index afc3e10130..b9405a0714 100644 --- a/apps/typegpu-docs/src/examples/algorithms/matrix-next/index.ts +++ b/apps/typegpu-docs/src/examples/algorithms/matrix-next/index.ts @@ -35,6 +35,10 @@ const buffers = { let bindGroup = createBindGroup(); const pipelines = createPipelines(); +const pipelinePromises = [ + pipelines['gpu-optimized'].initAsync(), + pipelines['gpu-simple'].initAsync(), +]; function createBindGroup() { return root.createBindGroup(computeLayout, { @@ -107,6 +111,7 @@ function updateInputDisplays() { let isComputing = false; +await Promise.all(pipelinePromises); async function compute() { if (isComputing) { console.warn('Computation already in progress'); diff --git a/apps/typegpu-docs/src/examples/algorithms/mnist-inference/index.ts b/apps/typegpu-docs/src/examples/algorithms/mnist-inference/index.ts index 4fe7ed0d89..3fd99f61ff 100644 --- a/apps/typegpu-docs/src/examples/algorithms/mnist-inference/index.ts +++ b/apps/typegpu-docs/src/examples/algorithms/mnist-inference/index.ts @@ -85,6 +85,7 @@ const pipelines = { ? root.createComputePipeline({ compute: subgroupCompute }) : null, }; +const pipelinePromises = [pipelines.default.initAsync(), pipelines.subgroup?.initAsync()]; // Definitions for the network @@ -181,6 +182,7 @@ function createNetwork(layers: [LayerData, LayerData][]): Network { } const network = createNetwork(await downloadLayers(root)); +await Promise.all(pipelinePromises); // #region Example controls and cleanup diff --git a/apps/typegpu-docs/src/examples/image-processing/background-segmentation/index.ts b/apps/typegpu-docs/src/examples/image-processing/background-segmentation/index.ts index 0c08feb8f6..6b86aa54bd 100644 --- a/apps/typegpu-docs/src/examples/image-processing/background-segmentation/index.ts +++ b/apps/typegpu-docs/src/examples/image-processing/background-segmentation/index.ts @@ -113,6 +113,7 @@ let blurBindGroups: TgpuBindGroup[]; const prepareModelInputPipeline = root .with(paramsAccess, paramsUniform) .createGuardedComputePipeline(prepareModelInput); +prepareModelInputPipeline.initSync(); let currentModelIndex = 0; let session = await prepareSession( @@ -138,6 +139,7 @@ async function switchModel(modelIndex: number) { } const generateMaskFromOutputPipeline = root.createGuardedComputePipeline(generateMaskFromOutput); +generateMaskFromOutputPipeline.initSync(); const blurPipelines = [false, true].map((flip) => root.with(flipAccess, flip).createComputePipeline({ compute: computeFn }), diff --git a/apps/typegpu-docs/src/examples/rendering/function-visualizer/index.ts b/apps/typegpu-docs/src/examples/rendering/function-visualizer/index.ts index b426e67e53..61b11c223d 100644 --- a/apps/typegpu-docs/src/examples/rendering/function-visualizer/index.ts +++ b/apps/typegpu-docs/src/examples/rendering/function-visualizer/index.ts @@ -90,6 +90,7 @@ const createComputePipeline = (exprCode: string) => { const computePipelines: Array = initialFunctions.map( (functionData, _) => createComputePipeline(functionData.code), ); +const pipelinePromises = computePipelines.map((guarded) => guarded.initAsync()); // Render background shader @@ -224,6 +225,7 @@ function draw() { requestAnimationFrame(draw); } +await Promise.all(pipelinePromises); requestAnimationFrame(draw); function runComputePass(functionNumber: number) { diff --git a/apps/typegpu-docs/src/examples/simulation/gravity/index.ts b/apps/typegpu-docs/src/examples/simulation/gravity/index.ts index 39601f5c9e..daed85a102 100644 --- a/apps/typegpu-docs/src/examples/simulation/gravity/index.ts +++ b/apps/typegpu-docs/src/examples/simulation/gravity/index.ts @@ -113,6 +113,11 @@ const renderPipeline = root }, }); +const pipelinePromises = [ + computeCollisionsPipeline.initAsync(), + computeGravityPipeline.initAsync(), +]; + let depthTexture = root.device.createTexture({ size: [canvas.width, canvas.height, 1], format: 'depth24plus', @@ -167,6 +172,7 @@ function frame(timestamp: DOMHighResTimeStamp) { render(); requestAnimationFrame(frame); } +await Promise.all(pipelinePromises); requestAnimationFrame(frame); async function loadPreset(preset: Preset): Promise { diff --git a/apps/typegpu-docs/src/examples/simulation/stable-fluid/index.ts b/apps/typegpu-docs/src/examples/simulation/stable-fluid/index.ts index 420c81d294..a53ea37370 100644 --- a/apps/typegpu-docs/src/examples/simulation/stable-fluid/index.ts +++ b/apps/typegpu-docs/src/examples/simulation/stable-fluid/index.ts @@ -126,6 +126,19 @@ const projectPipeline = createComputePipeline(c.projectFn); const advectInkPipeline = createComputePipeline(c.advectInkFn); const addInkPipeline = createComputePipeline(c.addInkFn); +// Eagerly initialize all pipelines +await Promise.all([ + brushPipeline.initAsync(), + addForcePipeline.initAsync(), + advectPipeline.initAsync(), + diffusionPipeline.initAsync(), + divergencePipeline.initAsync(), + pressurePipeline.initAsync(), + projectPipeline.initAsync(), + advectInkPipeline.initAsync(), + addInkPipeline.initAsync(), +]); + // Create render pipelines function createRenderPipeline(fragmentFn: TgpuFragmentFn<{ uv: d.Vec2f }, d.Vec4f>) { return root.createRenderPipeline({ diff --git a/apps/typegpu-docs/tests/individual-example-tests/matrix-next.test.ts b/apps/typegpu-docs/tests/individual-example-tests/matrix-next.test.ts index fd5ba40cfe..6dac9c8e87 100644 --- a/apps/typegpu-docs/tests/individual-example-tests/matrix-next.test.ts +++ b/apps/typegpu-docs/tests/individual-example-tests/matrix-next.test.ts @@ -15,7 +15,7 @@ describe('matrix(next) example', () => { category: 'algorithms', name: 'matrix-next', controlTriggers: ['Compute'], - expectedCalls: 1, + expectedCalls: 2, }, device, ); @@ -84,6 +84,39 @@ describe('matrix(next) example', () => { let outputIndex = getIndex(globalRow, globalCol, (*dimensions).secondColumnCount); resultMatrix[outputIndex] = accumulatedResult; } + } + + struct MatrixInfo { + firstRowCount: u32, + firstColumnCount: u32, + secondColumnCount: u32, + } + + @group(0) @binding(3) var dimensions: MatrixInfo; + + fn getIndex(row: u32, col: u32, columns: u32) -> u32 { + return (col + (row * columns)); + } + + @group(0) @binding(0) var firstMatrix: array; + + @group(0) @binding(1) var secondMatrix: array; + + @group(0) @binding(2) var resultMatrix: array; + + @compute @workgroup_size(16, 16) fn computeSimple(@builtin(global_invocation_id) gid: vec3u) { + let row = gid.x; + let col = gid.y; + if (((row >= dimensions.firstRowCount) || (col >= dimensions.secondColumnCount))) { + return; + } + var result = 0; + for (var k = 0u; (k < dimensions.firstColumnCount); k++) { + let aValue = firstMatrix[getIndex(row, k, dimensions.firstColumnCount)]; + let bValue = secondMatrix[getIndex(k, col, dimensions.secondColumnCount)]; + result += (aValue * bValue); + } + resultMatrix[getIndex(row, col, dimensions.secondColumnCount)] = result; }" `); }); diff --git a/apps/typegpu-docs/tests/individual-example-tests/mnist-inference.test.ts b/apps/typegpu-docs/tests/individual-example-tests/mnist-inference.test.ts index 2174bedee4..1160c266ea 100644 --- a/apps/typegpu-docs/tests/individual-example-tests/mnist-inference.test.ts +++ b/apps/typegpu-docs/tests/individual-example-tests/mnist-inference.test.ts @@ -23,7 +23,31 @@ describe('mnist inference example', () => { ); expect(shaderCodes).toMatchInlineSnapshot(` - "enable subgroups; + "@group(0) @binding(0) var input: array; + + @group(1) @binding(0) var weights: array; + + @group(1) @binding(1) var biases: array; + + @group(0) @binding(1) var output: array; + + fn relu(x: f32) -> f32 { + return max(0f, x); + } + + @compute @workgroup_size(1) fn defaultCompute(@builtin(global_invocation_id) gid: vec3u) { + let inputSize = arrayLength(&input); + let i = gid.x; + let weightsOffset = (i * inputSize); + var sum = 0f; + for (var j = 0u; (j < inputSize); j++) { + sum = fma(input[j], weights[(weightsOffset + j)], sum); + } + let total = (sum + biases[i]); + output[i] = relu(total); + } + + enable subgroups; @group(0) @binding(0) var input: array; diff --git a/apps/typegpu-docs/tests/individual-example-tests/stable-fluid.test.ts b/apps/typegpu-docs/tests/individual-example-tests/stable-fluid.test.ts index 27477efbed..8d5d6fa4d3 100644 --- a/apps/typegpu-docs/tests/individual-example-tests/stable-fluid.test.ts +++ b/apps/typegpu-docs/tests/individual-example-tests/stable-fluid.test.ts @@ -19,13 +19,67 @@ describe('stable-fluid example', () => { mockImageLoading(); mockCreateImageBitmap(); }, - expectedCalls: 7, + expectedCalls: 9, }, device, ); expect(shaderCodes).toMatchInlineSnapshot(` - "@group(0) @binding(0) var src: texture_2d; + "struct BrushParams { + pos: vec2i, + delta: vec2f, + radius: f32, + forceScale: f32, + inkAmount: f32, + } + + @group(0) @binding(0) var brushParams: BrushParams; + + @group(0) @binding(1) var forceDst: texture_storage_2d; + + @group(0) @binding(2) var inkDst: texture_storage_2d; + + @compute @workgroup_size(16, 16) fn brushFn(@builtin(global_invocation_id) gid: vec3u) { + let pixelPos = gid.xy; + let brushSettings = (&brushParams); + var forceVec = vec2f(); + var inkAmount = 0f; + let deltaX = (f32(pixelPos.x) - f32((*brushSettings).pos.x)); + let deltaY = (f32(pixelPos.y) - f32((*brushSettings).pos.y)); + let distSquared = ((deltaX * deltaX) + (deltaY * deltaY)); + let radiusSquared = ((*brushSettings).radius * (*brushSettings).radius); + if ((distSquared < radiusSquared)) { + let brushWeight = exp((-(distSquared) / radiusSquared)); + forceVec = (((*brushSettings).forceScale * brushWeight) * (*brushSettings).delta); + inkAmount = ((*brushSettings).inkAmount * brushWeight); + } + textureStore(forceDst, pixelPos, vec4f(forceVec, 0f, 1f)); + textureStore(inkDst, pixelPos, vec4f(inkAmount, 0f, 0f, 1f)); + } + + @group(0) @binding(0) var src: texture_2d; + + @group(0) @binding(2) var force: texture_2d; + + struct ShaderParams { + dt: f32, + viscosity: f32, + } + + @group(0) @binding(3) var simParams: ShaderParams; + + @group(0) @binding(1) var dst: texture_storage_2d; + + @compute @workgroup_size(16, 16) fn addForcesFn(@builtin(global_invocation_id) gid: vec3u) { + let pixelPos = gid.xy; + let currentVel = textureLoad(src, pixelPos, 0).xy; + let forceVec = textureLoad(force, pixelPos, 0).xy; + let timeStep = simParams.dt; + let newVel = (currentVel + (timeStep * forceVec)); + textureStore(dst, pixelPos, vec4f(newVel, 0f, 1f)); + } + + @group(0) @binding(0) var src: texture_2d; @group(0) @binding(1) var dst: texture_storage_2d; @@ -248,6 +302,19 @@ describe('stable-fluid example', () => { textureStore(dst, pixelPos, inkVal); } + @group(0) @binding(2) var add: texture_2d; + + @group(0) @binding(0) var src: texture_2d; + + @group(0) @binding(1) var dst: texture_storage_2d; + + @compute @workgroup_size(16, 16) fn addInkFn(@builtin(global_invocation_id) gid: vec3u) { + let pixelPos = gid.xy; + let addVal = textureLoad(add, pixelPos, 0).x; + let srcVal = textureLoad(src, pixelPos, 0).x; + textureStore(dst, pixelPos, vec4f((addVal + srcVal), 0f, 0f, 1f)); + } + struct renderFn_Output { @builtin(position) pos: vec4f, @location(0) uv: vec2f, diff --git a/packages/typegpu-testing-utility/src/extendedIt.ts b/packages/typegpu-testing-utility/src/extendedIt.ts index 00ba17e02c..ea3b4a774b 100644 --- a/packages/typegpu-testing-utility/src/extendedIt.ts +++ b/packages/typegpu-testing-utility/src/extendedIt.ts @@ -114,6 +114,7 @@ export const it = base return commandEncoder; }), createComputePipeline: vi.fn(() => mockComputePipeline), + createComputePipelineAsync: vi.fn(async () => mockComputePipeline), createPipelineLayout: vi.fn(() => 'mockPipelineLayout'), createQuerySet: vi.fn(({ type, count }: GPUQuerySetDescriptor) => { const querySet = Object.create(mockQuerySet); diff --git a/packages/typegpu/src/core/pipeline/computePipeline.ts b/packages/typegpu/src/core/pipeline/computePipeline.ts index f5e8dd7590..c1de8e902f 100644 --- a/packages/typegpu/src/core/pipeline/computePipeline.ts +++ b/packages/typegpu/src/core/pipeline/computePipeline.ts @@ -67,6 +67,22 @@ export interface TgpuComputePipeline extends TgpuNamable, SelfResolvable, Timeab dispatchWorkgroups(x: number, y?: number, z?: number): void; + /** + * Immediately resolves the pipeline, then calls `device.createComputePipelineAsync()`. + * + * NOTE: while it is not necessary to initialize pipeline manually, + * it is generally preferable to use this method whenever possible, + * as it prevents blocking of GPU operation execution on pipeline compilation. + */ + initAsync(): Promise; + + /** + * Immediately resolves the pipeline and creates WebGPU resources. + * + * NOTE: it is not necessary to initialize pipeline manually. + */ + initSync(): void; + /** * Dispatches compute workgroups using parameters read from a buffer. * The buffer must contain 3 consecutive u32 values (x, y, z workgroup counts). @@ -249,6 +265,14 @@ class TgpuComputePipelineImpl implements TgpuComputePipeline { this._executeComputePass((pass) => pass.dispatchWorkgroupsIndirect(rawBuffer, offset)); } + initAsync(): Promise { + return this.#core.initAsync(); + } + + initSync() { + this.#core.initSync(); + } + private _applyComputeState(pass: GPUComputePassEncoder): void { const memo = this.#core.unwrap(); const { root } = this.#core; @@ -318,7 +342,10 @@ class TgpuComputePipelineImpl implements TgpuComputePipeline { class ComputePipelineCore implements SelfResolvable { readonly [$internal] = true; readonly root: ExperimentalTgpuRoot; - private _memo: Memo | undefined; + #performanceTracker: PerformanceTracker; + + #initAsyncPromise: Promise | undefined; + #memo: Memo | undefined; #slotBindings: [TgpuSlot, unknown][]; #descriptor: TgpuComputePipeline.Descriptor; @@ -332,6 +359,9 @@ class ComputePipelineCore implements SelfResolvable { this.root = root; this.#slotBindings = slotBindings; this.#descriptor = descriptor; + this.#performanceTracker = PERF?.enabled + ? new PerformanceTrackerImpl() + : new NullPerformanceTracker(); } [$resolve](ctx: ResolutionCtx) { @@ -352,82 +382,152 @@ class ComputePipelineCore implements SelfResolvable { return (this.#performanceCallbackQuerySet ??= this.root.createQuerySet('timestamp', 2)); } - public unwrap(): Memo { - if (this._memo === undefined) { - const device = this.root.device; - const enableExtensions = wgslEnableExtensions.filter((extension) => - this.root.enabledFeatures.has(wgslEnableExtensionToFeatureName[extension]), - ); - - // Resolving code - let resolutionResult: ResolutionResult; - - let resolveMeasure: PerformanceMeasure | undefined; - const ns = namespace({ names: this.root.nameRegistrySetting }); - if (PERF?.enabled) { - const resolveStart = performance.mark('typegpu:resolution:start'); - resolutionResult = resolve(this, { - namespace: ns, - enableExtensions, - shaderGenerator: this.root.shaderGenerator, - root: this.root, - }); - resolveMeasure = performance.measure('typegpu:resolution', { - start: resolveStart.name, - }); - } else { - resolutionResult = resolve(this, { - namespace: ns, - enableExtensions, - shaderGenerator: this.root.shaderGenerator, - root: this.root, - }); - } - - const { code, usedBindGroupLayouts, catchall, logResources } = resolutionResult; - - if (catchall !== undefined) { - usedBindGroupLayouts[catchall[0]]?.$name( - `${getName(this) ?? ''} - Automatic Bind Group & Layout`, - ); - } + /** + * @privateRemarks + * This function cannot be a regular async function + * because when called multiple times before the promise finishes, + * we want it to return the same promise each time. + */ + initAsync(): Promise { + if (this.#memo !== undefined) { + // the pipeline was already resolved & compiled + return Promise.resolve(); + } - const module = device.createShaderModule({ - label: `${getName(this) ?? ''} - Shader`, - code, - }); + if (this.#initAsyncPromise === undefined) { + // the pipeline did not start resolution & compilation + const device = this.root.device; + const { resolutionResult, module } = this.resolveAndCreateShaderModule(); + const { usedBindGroupLayouts, catchall, logResources } = resolutionResult; - this._memo = { - pipeline: device.createComputePipeline({ + this.#initAsyncPromise = device + .createComputePipelineAsync({ label: getName(this) ?? '', layout: device.createPipelineLayout({ label: `${getName(this) ?? ''} - Pipeline Layout`, bindGroupLayouts: usedBindGroupLayouts.map((l) => this.root.unwrap(l)), }), compute: { module }, + }) + .then((pipeline) => { + this.#memo = { pipeline, usedBindGroupLayouts, catchall, logResources }; + this.#performanceTracker.measureCompile(device); + }) + .finally(() => { + this.#initAsyncPromise = undefined; + }); + } + return this.#initAsyncPromise; + } + + initSync() { + if (this.#memo !== undefined) { + return; + } + + if (this.#initAsyncPromise !== undefined) { + throw new Error("'pipeline.initAsync()' was called and is not yet resolved."); + } + + const device = this.root.device; + const { resolutionResult, module } = this.resolveAndCreateShaderModule(); + const { usedBindGroupLayouts, catchall, logResources } = resolutionResult; + + this.#memo = { + pipeline: device.createComputePipeline({ + label: getName(this) ?? '', + layout: device.createPipelineLayout({ + label: `${getName(this) ?? ''} - Pipeline Layout`, + bindGroupLayouts: usedBindGroupLayouts.map((l) => this.root.unwrap(l)), }), - usedBindGroupLayouts, - catchall, - logResources, - }; + compute: { module }, + }), + usedBindGroupLayouts, + catchall, + logResources, + }; - if (PERF?.enabled) { - void (async () => { - const start = performance.mark('typegpu:compile-start'); - await device.queue.onSubmittedWorkDone(); - const compileMeasure = performance.measure('typegpu:compiled', { - start: start.name, - }); - - PERF?.record('resolution', { - resolveDuration: resolveMeasure?.duration, - compileDuration: compileMeasure.duration, - wgslSize: code.length, - }); - })(); - } + this.#performanceTracker.measureCompile(device); + } + + public unwrap(): Memo { + this.initSync(); + return this.#memo as Memo; + } + + private resolveAndCreateShaderModule() { + const device = this.root.device; + const enableExtensions = wgslEnableExtensions.filter((extension) => + this.root.enabledFeatures.has(wgslEnableExtensionToFeatureName[extension]), + ); + + // Resolving code + const ns = namespace({ names: this.root.nameRegistrySetting }); + const resolutionResult = this.#performanceTracker.measureResolve(() => + resolve(this, { + namespace: ns, + enableExtensions, + shaderGenerator: this.root.shaderGenerator, + root: this.root, + }), + ); + const { code, usedBindGroupLayouts, catchall } = resolutionResult; + + if (catchall !== undefined) { + usedBindGroupLayouts[catchall[0]]?.$name( + `${getName(this) ?? ''} - Automatic Bind Group & Layout`, + ); } - return this._memo; + const module = device.createShaderModule({ + label: `${getName(this) ?? ''} - Shader`, + code, + }); + + return { resolutionResult, module }; } } + +interface PerformanceTracker { + measureResolve(callback: () => ResolutionResult): ResolutionResult; + measureCompile(device: GPUDevice): void; +} + +class PerformanceTrackerImpl implements PerformanceTracker { + #resolveMeasure: PerformanceMeasure | undefined; + #wgslSize: number | undefined; + + measureResolve(callback: () => ResolutionResult): ResolutionResult { + const resolveStart = performance.mark('typegpu:resolution:start'); + const result = callback(); + this.#resolveMeasure = performance.measure('typegpu:resolution', { + start: resolveStart.name, + }); + this.#wgslSize = result.code.length; + return result; + } + + measureCompile(device: GPUDevice): void { + void (async () => { + const start = performance.mark('typegpu:compile-start'); + await device.queue.onSubmittedWorkDone(); + const compileMeasure = performance.measure('typegpu:compiled', { + start: start.name, + }); + + PERF?.record('resolution', { + resolveDuration: this.#resolveMeasure?.duration, + compileDuration: compileMeasure.duration, + wgslSize: this.#wgslSize, + }); + })(); + } +} + +class NullPerformanceTracker implements PerformanceTracker { + measureResolve(callback: () => ResolutionResult): ResolutionResult { + return callback(); + } + + measureCompile(): void {} +} diff --git a/packages/typegpu/src/core/root/init.ts b/packages/typegpu/src/core/root/init.ts index 075ac5e36a..1dfa1ed9ef 100644 --- a/packages/typegpu/src/core/root/init.ts +++ b/packages/typegpu/src/core/root/init.ts @@ -184,6 +184,14 @@ export class TgpuGuardedComputePipelineImpl< this.#pipeline.dispatchWorkgroups(workgroupCount.x, workgroupCount.y, workgroupCount.z); } + initAsync() { + return this.pipeline.initAsync(); + } + + initSync() { + return this.pipeline.initSync(); + } + get pipeline() { return this.#pipeline; } diff --git a/packages/typegpu/src/core/root/rootTypes.ts b/packages/typegpu/src/core/root/rootTypes.ts index 1a52e07654..7952e52d76 100644 --- a/packages/typegpu/src/core/root/rootTypes.ts +++ b/packages/typegpu/src/core/root/rootTypes.ts @@ -108,6 +108,22 @@ export interface TgpuGuardedComputePipeline e */ dispatchThreads(...args: TArgs): void; + /** + * Immediately resolves the pipeline, then calls `device.createComputePipelineAsync()`. + * + * NOTE: while it is not necessary to initialize pipeline manually, + * it is generally preferable to use this method whenever possible, + * as it prevents blocking of GPU operation execution on pipeline compilation. + */ + initAsync(): Promise; + + /** + * Immediately resolves the pipeline and creates WebGPU resources. + * + * NOTE: it is not necessary to initialize pipeline manually. + */ + initSync(): void; + /** * The underlying pipeline used during `dispatchThreads`. */ diff --git a/packages/typegpu/tests/pipelineInit.test.ts b/packages/typegpu/tests/pipelineInit.test.ts new file mode 100644 index 0000000000..91c1880879 --- /dev/null +++ b/packages/typegpu/tests/pipelineInit.test.ts @@ -0,0 +1,70 @@ +import { describe, expect } from 'vitest'; +import tgpu from '../src/index.js'; +import { it } from 'typegpu-testing-utility'; + +describe('pipeline initialization', () => { + describe('compute pipeline', () => { + const computeFn = tgpu.computeFn({ workgroupSize: [1, 1, 1] })(() => {}); + describe('initSync', () => { + it('resolves and creates a pipeline', ({ root, device }) => { + const pipeline = root.createComputePipeline({ compute: computeFn }); + + pipeline.initSync(); + + expect(device.mock.createComputePipeline).toHaveBeenCalled(); + expect(tgpu.resolve([pipeline])).toMatchInlineSnapshot(` + "@compute @workgroup_size(1, 1, 1) fn computeFn() { + + }" + `); + }); + }); + + describe('initAsync', () => { + it('resolves and creates a pipeline', async ({ root, device }) => { + const pipeline = root.createComputePipeline({ compute: computeFn }); + + await pipeline.initAsync(); + + expect(device.mock.createComputePipelineAsync).toHaveBeenCalled(); + expect(() => root.unwrap(pipeline)).not.toThrow(); // this means that memo already exists + expect(tgpu.resolve([pipeline])).toMatchInlineSnapshot(` + "@compute @workgroup_size(1, 1, 1) fn computeFn() { + + }" + `); + expect(device.mock.createComputePipeline).not.toHaveBeenCalled(); + }); + + it('throws when attempting to unwrap when not resolved', async ({ root, device }) => { + const pipeline = root.createComputePipeline({ compute: computeFn }); + + // oxlint-disable-next-line typescript-eslint/no-floating-promises -- it's a test + pipeline.initAsync(); + + expect(() => root.unwrap(pipeline)).toThrowErrorMatchingInlineSnapshot( + `[Error: 'pipeline.initAsync()' was called and is not yet resolved.]`, + ); + }); + + it('returns the promise again if not resolved', async ({ root, device }) => { + const pipeline = root.createComputePipeline({ compute: computeFn }); + + const p1 = pipeline.initAsync(); + const p2 = pipeline.initAsync(); + + expect(p1).toBe(p2); + }); + + it('returns a resolved promise if pipeline was already created', async ({ root, device }) => { + const pipeline = root.createComputePipeline({ compute: computeFn }); + + root.unwrap(pipeline); // resolves & compiles the pipeline + await pipeline.initAsync(); // should not enqueue device operations + + expect(device.mock.createComputePipeline).toHaveBeenCalled(); + expect(device.mock.createComputePipelineAsync).not.toHaveBeenCalled(); + }); + }); + }); +});