[WIP] feat: Asynchronous pipeline creation

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

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');

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
 

apps/typegpu-docs/src/examples/image-processing/background-segmentation/index.ts

@@ -113,6 +113,7 @@ let blurBindGroups: TgpuBindGroup<typeof blurLayout.entries>[];
 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 }),

apps/typegpu-docs/src/examples/rendering/function-visualizer/index.ts

@@ -90,6 +90,7 @@ const createComputePipeline = (exprCode: string) => {
 const computePipelines: Array<TgpuGuardedComputePipeline> = 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) {

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<DynamicResources> {

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({

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<uniform> dimensions: MatrixInfo;
+
+      fn getIndex(row: u32, col: u32, columns: u32) -> u32 {
+        return (col + (row * columns));
+      }
+
+      @group(0) @binding(0) var<storage, read> firstMatrix: array<i32>;
+
+      @group(0) @binding(1) var<storage, read> secondMatrix: array<i32>;
+
+      @group(0) @binding(2) var<storage, read_write> resultMatrix: array<i32>;
+
+      @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;
       }"
     `);
   });

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<storage, read> input: array<f32>;
+
+      @group(1) @binding(0) var<storage, read> weights: array<f32>;
+
+      @group(1) @binding(1) var<storage, read> biases: array<f32>;
+
+      @group(0) @binding(1) var<storage, read_write> output: array<f32>;
+
+      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<storage, read> input: array<f32>;
 

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<f32>;
+      "struct BrushParams {
+        pos: vec2i,
+        delta: vec2f,
+        radius: f32,
+        forceScale: f32,
+        inkAmount: f32,
+      }
+
+      @group(0) @binding(0) var<uniform> brushParams: BrushParams;
+
+      @group(0) @binding(1) var forceDst: texture_storage_2d<rgba16float, write>;
+
+      @group(0) @binding(2) var inkDst: texture_storage_2d<rgba16float, write>;
+
+      @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<f32>;
+
+      @group(0) @binding(2) var force: texture_2d<f32>;
+
+      struct ShaderParams {
+        dt: f32,
+        viscosity: f32,
+      }
+
+      @group(0) @binding(3) var<uniform> simParams: ShaderParams;
+
+      @group(0) @binding(1) var dst: texture_storage_2d<rgba16float, write>;
+
+      @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<f32>;
 
       @group(0) @binding(1) var dst: texture_storage_2d<rgba16float, write>;
 
@@ -248,6 +302,19 @@ describe('stable-fluid example', () => {
         textureStore(dst, pixelPos, inkVal);
       }
 
+      @group(0) @binding(2) var add: texture_2d<f32>;
+
+      @group(0) @binding(0) var src: texture_2d<f32>;
+
+      @group(0) @binding(1) var dst: texture_storage_2d<rgba16float, write>;
+
+      @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,

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);