feat: Add Indirect buffer usage and support .dispatchWorkgroupsIndirect API (#2105)

Author: reczkok

apps/typegpu-docs/src/examples/algorithms/mnist-inference/index.ts

@@ -19,7 +19,7 @@ const hasTimestampQuery = root.enabledFeatures.has('timestamp-query');
 const hasSubgroups = root.enabledFeatures.has('subgroups');
 let useSubgroups = hasSubgroups;
 
-const canvasData = Array.from({ length: (SIZE ** 2) }, () => 0);
+const canvasData = Array.from({ length: SIZE ** 2 }, () => 0);
 
 // Shaders
 
@@ -314,7 +314,7 @@ function centerImage(data: number[]) {
   const offsetX = Math.round(SIZE / 2 - x);
   const offsetY = Math.round(SIZE / 2 - y);
 
-  const newData = Array.from({ length: (SIZE * SIZE) }, () => 0);
+  const newData = Array.from({ length: SIZE * SIZE }, () => 0);
   for (let i = 0; i < SIZE; i++) {
     for (let j = 0; j < SIZE; j++) {
       const index = i * SIZE + j;

packages/typegpu-three/src/typegpu-node.ts

@@ -397,7 +397,7 @@ export const fromTSL = tgpu.comptime(
     let nodeType: string | null = null;
     try { // sometimes it needs information (overrideNodes) from compilation context which is not present
       nodeType = node.getNodeType(sharedBuilder);
-    } catch (e) {
+    } catch {
       console.warn(
         `fromTSL: failed to infer node type via getNodeType; skipping type comparison.`,
       );

packages/typegpu/src/core/buffer/buffer.ts

@@ -56,17 +56,24 @@ export interface IndexFlag {
   usableAsIndex: true;
 }
 
+export interface IndirectFlag {
+  usableAsIndirect: true;
+}
+
 /**
  * @deprecated Use VertexFlag instead.
  */
 export type Vertex = VertexFlag;
 
-type LiteralToUsageType<T extends 'uniform' | 'storage' | 'vertex' | 'index'> =
-  T extends 'uniform' ? UniformFlag
-    : T extends 'storage' ? StorageFlag
-    : T extends 'vertex' ? VertexFlag
-    : T extends 'index' ? IndexFlag
-    : never;
+type UsageLiteral = 'uniform' | 'storage' | 'vertex' | 'index' | 'indirect';
+
+type LiteralToUsageType<T extends UsageLiteral> = T extends 'uniform'
+  ? UniformFlag
+  : T extends 'storage' ? StorageFlag
+  : T extends 'vertex' ? VertexFlag
+  : T extends 'index' ? IndexFlag
+  : T extends 'indirect' ? IndirectFlag
+  : never;
 
 type ViewUsages<TBuffer extends TgpuBuffer<BaseData>> =
   | (boolean extends TBuffer['usableAsUniform'] ? never : 'uniform')
@@ -89,7 +96,9 @@ type InnerValidUsagesFor<T> = {
     | (IsValidStorageSchema<T> extends true ? 'storage' : never)
     | (IsValidUniformSchema<T> extends true ? 'uniform' : never)
     | (IsValidVertexSchema<T> extends true ? 'vertex' : never)
-    | (IsValidIndexSchema<T> extends true ? 'index' : never);
+    | (IsValidIndexSchema<T> extends true ? 'index' : never)
+    // there is no way to check at the type level if a buffer can be used as indirect (size >= 12 bytes)
+    | 'indirect';
 };
 
 export type ValidUsagesFor<T> = InnerValidUsagesFor<T>['usage'];
@@ -107,6 +116,7 @@ export interface TgpuBuffer<TData extends BaseData> extends TgpuNamable {
   usableAsStorage: boolean;
   usableAsVertex: boolean;
   usableAsIndex: boolean;
+  usableAsIndirect: boolean;
 
   $usage<
     T extends [
@@ -184,13 +194,13 @@ class TgpuBufferImpl<TData extends BaseData> implements TgpuBuffer<TData> {
   usableAsStorage = false;
   usableAsVertex = false;
   usableAsIndex = false;
+  usableAsIndirect = false;
 
   constructor(
     root: ExperimentalTgpuRoot,
     public readonly dataType: TData,
     public readonly initialOrBuffer?: Infer<TData> | GPUBuffer,
-    private readonly _disallowedUsages?:
-      ('uniform' | 'storage' | 'vertex' | 'index')[],
+    private readonly _disallowedUsages?: UsageLiteral[],
   ) {
     this.#device = root.device;
     if (isGPUBuffer(initialOrBuffer)) {
@@ -236,7 +246,7 @@ class TgpuBufferImpl<TData extends BaseData> implements TgpuBuffer<TData> {
     return this;
   }
 
-  $usage<T extends ('uniform' | 'storage' | 'vertex' | 'index')[]>(
+  $usage<T extends UsageLiteral[]>(
     ...usages: T
   ): this & UnionToIntersection<LiteralToUsageType<T[number]>> {
     for (const usage of usages) {
@@ -250,10 +260,12 @@ class TgpuBufferImpl<TData extends BaseData> implements TgpuBuffer<TData> {
       this.flags |= usage === 'storage' ? GPUBufferUsage.STORAGE : 0;
       this.flags |= usage === 'vertex' ? GPUBufferUsage.VERTEX : 0;
       this.flags |= usage === 'index' ? GPUBufferUsage.INDEX : 0;
+      this.flags |= usage === 'indirect' ? GPUBufferUsage.INDIRECT : 0;
       this.usableAsUniform = this.usableAsUniform || usage === 'uniform';
       this.usableAsStorage = this.usableAsStorage || usage === 'storage';
       this.usableAsVertex = this.usableAsVertex || usage === 'vertex';
       this.usableAsIndex = this.usableAsIndex || usage === 'index';
+      this.usableAsIndirect = this.usableAsIndirect || usage === 'indirect';
     }
     return this as this & UnionToIntersection<LiteralToUsageType<T[number]>>;
   }

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

@@ -1,11 +1,14 @@
 import type { AnyComputeBuiltin } from '../../builtin.ts';
 import type { TgpuQuerySet } from '../../core/querySet/querySet.ts';
 import { type ResolvedSnippet, snip } from '../../data/snippet.ts';
+import { sizeOf } from '../../data/sizeOf.ts';
+import type { AnyWgslData } from '../../data/wgslTypes.ts';
 import { Void } from '../../data/wgslTypes.ts';
 import { MissingBindGroupsError } from '../../errors.ts';
 import { type ResolutionResult, resolve } from '../../resolutionCtx.ts';
 import type { TgpuNamable } from '../../shared/meta.ts';
 import { getName, PERF, setName } from '../../shared/meta.ts';
+
 import { $getNameForward, $internal, $resolve } from '../../shared/symbols.ts';
 import {
   isBindGroup,
@@ -26,6 +29,10 @@ import { namespace } from '../resolve/namespace.ts';
 import type { ExperimentalTgpuRoot } from '../root/rootTypes.ts';
 import type { TgpuSlot } from '../slot/slotTypes.ts';
 import { warnIfOverflow } from './limitsOverflow.ts';
+import {
+  getOffsetInfoAt,
+  type PrimitiveOffsetInfo,
+} from '../../data/offsetUtils.ts';
 import {
   createWithPerformanceCallback,
   createWithTimestampWrites,
@@ -34,6 +41,7 @@ import {
   type TimestampWritesPriors,
   triggerPerformanceCallback,
 } from './timeable.ts';
+import type { IndirectFlag, TgpuBuffer } from '../buffer/buffer.ts';
 
 interface ComputePipelineInternals {
   readonly rawPipeline: GPUComputePipeline;
@@ -65,6 +73,19 @@ export interface TgpuComputePipeline
     y?: number,
     z?: number,
   ): void;
+
+  /**
+   * Dispatches compute workgroups using parameters read from a buffer.
+   * The buffer must contain 3 consecutive u32 values (x, y, z workgroup counts).
+   * To get the correct offset within complex data structures, use `d.getOffsetInfoAt(...)`.
+   *
+   * @param indirectBuffer - Buffer marked with 'indirect' usage containing dispatch parameters
+   * @param start - PrimitiveOffsetInfo pointing to the first dispatch parameter. If not provided, starts at offset 0. To obtain safe offsets, use `getOffsetInfoAt(...)`.
+   */
+  dispatchWorkgroupsIndirect<T extends AnyWgslData>(
+    indirectBuffer: TgpuBuffer<T> & IndirectFlag,
+    start?: PrimitiveOffsetInfo | number,
+  ): void;
 }
 
 export declare namespace TgpuComputePipeline {
@@ -101,6 +122,27 @@ type Memo = {
   logResources: LogResources | undefined;
 };
 
+function validateIndirectBufferSize(
+  bufferSize: number,
+  offset: number,
+  requiredBytes: number,
+  operation: string,
+): void {
+  if (offset + requiredBytes > bufferSize) {
+    throw new Error(
+      `Buffer too small for ${operation}. ` +
+        `Required: ${requiredBytes} bytes at offset ${offset}, ` +
+        `but buffer is only ${bufferSize} bytes.`,
+    );
+  }
+
+  if (offset % 4 !== 0) {
+    throw new Error(
+      `Indirect buffer offset must be a multiple of 4. Got: ${offset}`,
+    );
+  }
+}
+
 class TgpuComputePipelineImpl implements TgpuComputePipeline {
   public readonly [$internal]: ComputePipelineInternals;
   public readonly resourceType = 'compute-pipeline';
@@ -192,6 +234,55 @@ class TgpuComputePipelineImpl implements TgpuComputePipeline {
     x: number,
     y?: number,
     z?: number,
+  ): void {
+    this._executeComputePass((pass) => pass.dispatchWorkgroups(x, y, z));
+  }
+
+  dispatchWorkgroupsIndirect<T extends AnyWgslData>(
+    indirectBuffer: TgpuBuffer<T> & IndirectFlag,
+    start?: PrimitiveOffsetInfo | number,
+  ): void {
+    const DISPATCH_SIZE = 12; // 3 x u32 (x, y, z)
+
+    let offsetInfo = start ?? getOffsetInfoAt(indirectBuffer.dataType);
+
+    if (typeof offsetInfo === 'number') {
+      if (offsetInfo === 0) {
+        offsetInfo = getOffsetInfoAt(indirectBuffer.dataType);
+      } else {
+        console.warn(
+          `dispatchWorkgroupsIndirect: Provided start offset ${offsetInfo} as a raw number. Use getOffsetInfoAt(...) to include contiguous padding info for safer validation.`,
+        );
+        // When only an offset is provided, assume we have at least 12 bytes contiguous.
+        offsetInfo = {
+          offset: offsetInfo,
+          contiguous: DISPATCH_SIZE,
+        };
+      }
+    }
+
+    const { offset, contiguous } = offsetInfo;
+
+    validateIndirectBufferSize(
+      sizeOf(indirectBuffer.dataType),
+      offset,
+      DISPATCH_SIZE,
+      'dispatchWorkgroupsIndirect',
+    );
+
+    if (contiguous < DISPATCH_SIZE) {
+      console.warn(
+        `dispatchWorkgroupsIndirect: Starting at offset ${offset}, only ${contiguous} contiguous bytes are available before padding. Dispatch requires ${DISPATCH_SIZE} bytes (3 x u32). Reading across padding may result in undefined behavior.`,
+      );
+    }
+
+    this._executeComputePass((pass) =>
+      pass.dispatchWorkgroupsIndirect(indirectBuffer.buffer, offset)
+    );
+  }
+
+  private _executeComputePass(
+    dispatch: (pass: GPUComputePassEncoder) => void,
   ): void {
     const memo = this._core.unwrap();
     const { root } = this._core;
@@ -231,7 +322,7 @@ class TgpuComputePipelineImpl implements TgpuComputePipeline {
       throw new MissingBindGroupsError(missingBindGroups);
     }
 
-    pass.dispatchWorkgroups(x, y, z);
+    dispatch(pass);
     pass.end();
     root.device.queue.submit([commandEncoder.finish()]);
 

packages/typegpu/src/data/alignmentOf.ts

@@ -25,17 +25,17 @@ const knownAlignmentMap: Record<string, number> = {
   vec2h: 4,
   vec2i: 8,
   vec2u: 8,
-  vec2b: 8,
+  'vec2<bool>': 8,
   vec3f: 16,
   vec3h: 8,
   vec3i: 16,
   vec3u: 16,
-  vec3b: 16,
+  'vec3<bool>': 16,
   vec4f: 16,
   vec4h: 8,
   vec4i: 16,
   vec4u: 16,
-  vec4b: 16,
+  'vec4<bool>': 16,
   mat2x2f: 8,
   mat3x3f: 16,
   mat4x4f: 16,

packages/typegpu/src/data/getLongestContiguousPrefix.ts

@@ -0,0 +1,14 @@
+import type { AnyData } from './dataTypes.ts';
+import type { BaseData } from './wgslTypes.ts';
+import { getLayoutInfo } from './schemaMemoryLayout.ts';
+
+export function getLongestContiguousPrefix(schema: BaseData): number {
+  return getLayoutInfo(schema, 'longestContiguousPrefix');
+}
+
+/**
+ * Returns the size (in bytes) of the longest contiguous memory prefix of data represented by the `schema`.
+ */
+export function PUBLIC_getLongestContiguousPrefix(schema: AnyData): number {
+  return getLongestContiguousPrefix(schema);
+}

packages/typegpu/src/data/index.ts

@@ -211,6 +211,11 @@ export {
   isUnstruct,
 } from './dataTypes.ts';
 export { PUBLIC_sizeOf as sizeOf } from './sizeOf.ts';
+export { PUBLIC_isContiguous as isContiguous } from './isContiguous.ts';
+export {
+  PUBLIC_getLongestContiguousPrefix as getLongestContiguousPrefix,
+} from './getLongestContiguousPrefix.ts';
+export { getOffsetInfoAt } from './offsetUtils.ts';
 export { PUBLIC_alignmentOf as alignmentOf } from './alignmentOf.ts';
 export { builtin } from '../builtin.ts';
 export { deepEqual } from './deepEqual.ts';

packages/typegpu/src/data/isContiguous.ts

@@ -0,0 +1,14 @@
+import type { AnyData } from './dataTypes.ts';
+import type { BaseData } from './wgslTypes.ts';
+import { getLayoutInfo } from './schemaMemoryLayout.ts';
+
+export function isContiguous(schema: BaseData): boolean {
+  return getLayoutInfo(schema, 'isContiguous');
+}
+
+/**
+ * Returns `true` if data represented by the `schema` doesn't have padding.
+ */
+export function PUBLIC_isContiguous(schema: AnyData): boolean {
+  return isContiguous(schema);
+}