fix: Support for PrimitiveOffsetInfo in render pipeline draw...Indirect methods

packages/typegpu-testing-utility/src/extendedIt.ts

@@ -23,6 +23,8 @@ export const it = base
       },
       draw: vi.fn(),
       drawIndexed: vi.fn(),
+      drawIndirect: vi.fn(),
+      drawIndexedIndirect: vi.fn(),
       end: vi.fn(),
       setBindGroup: vi.fn(),
       setPipeline: vi.fn(),

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

@@ -1,7 +1,6 @@
 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 { applyBindGroups } from './applyPipelineState.ts';
@@ -19,15 +18,16 @@ import {
 import { isGPUCommandEncoder, isGPUComputePassEncoder } from './typeGuards.ts';
 import { logDataFromGPU } from '../../tgsl/consoleLog/deserializers.ts';
 import type { LogResources } from '../../tgsl/consoleLog/types.ts';
-import type { ResolutionCtx, SelfResolvable } from '../../types.ts';
+import { isGPUBuffer, type ResolutionCtx, type SelfResolvable } from '../../types.ts';
 import { wgslExtensions, wgslExtensionToFeatureName } from '../../wgslExtensions.ts';
 import type { IORecord } from '../function/fnTypes.ts';
 import type { TgpuComputeFn } from '../function/tgpuComputeFn.ts';
 import { namespace } from '../resolve/namespace.ts';
 import type { ExperimentalTgpuRoot } from '../root/rootTypes.ts';
 import type { TgpuSlot } from '../slot/slotTypes.ts';
 
-import { memoryLayoutOf, type PrimitiveOffsetInfo } from '../../data/offsetUtils.ts';
+import type { PrimitiveOffsetInfo } from '../../data/offsetUtils.ts';
+import { resolveIndirectOffset } from './pipelineUtils.ts';
 import {
   createWithPerformanceCallback,
   createWithTimestampWrites,
@@ -72,11 +72,11 @@ export interface TgpuComputePipeline extends TgpuNamable, SelfResolvable, Timeab
    * The buffer must contain 3 consecutive u32 values (x, y, z workgroup counts).
    * To get the correct offset within complex data structures, use `d.memoryLayoutOf(...)`.
    *
-   * @param indirectBuffer - Buffer marked with 'indirect' usage containing dispatch parameters
+   * @param indirectBuffer - Buffer marked with 'indirect' usage containing dispatch parameters or raw GPUBuffer
    * @param start - PrimitiveOffsetInfo pointing to the first dispatch parameter. If not provided, starts at offset 0. To obtain safe offsets, use `d.memoryLayoutOf(...)`.
    */
   dispatchWorkgroupsIndirect<T extends AnyWgslData>(
-    indirectBuffer: TgpuBuffer<T> & IndirectFlag,
+    indirectBuffer: (TgpuBuffer<T> & IndirectFlag) | GPUBuffer,
     start?: PrimitiveOffsetInfo | number,
   ): void;
 }
@@ -113,25 +113,6 @@ 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}`);
-  }
-}
-
 const _lastAppliedCompute = new WeakMap<GPUComputePassEncoder, TgpuComputePipelineImpl>();
 
 class TgpuComputePipelineImpl implements TgpuComputePipeline {
@@ -252,46 +233,20 @@ class TgpuComputePipelineImpl implements TgpuComputePipeline {
   }
 
   dispatchWorkgroupsIndirect<T extends AnyWgslData>(
-    indirectBuffer: TgpuBuffer<T> & IndirectFlag,
+    indirectBuffer: (TgpuBuffer<T> & IndirectFlag) | GPUBuffer,
     start?: PrimitiveOffsetInfo | number,
   ): void {
     const DISPATCH_SIZE = 12; // 3 x u32 (x, y, z)
 
-    let offsetInfo = start ?? memoryLayoutOf(indirectBuffer.dataType);
-
-    if (typeof offsetInfo === 'number') {
-      if (offsetInfo === 0) {
-        offsetInfo = memoryLayoutOf(indirectBuffer.dataType);
-      } else {
-        console.warn(
-          `dispatchWorkgroupsIndirect: Provided start offset ${offsetInfo} as a raw number. Use d.memoryLayoutOf(...) 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,
+    const rawBuffer = isGPUBuffer(indirectBuffer) ? indirectBuffer : indirectBuffer.buffer;
+    const offset = resolveIndirectOffset(
+      indirectBuffer,
+      start,
       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),
-    );
+    this._executeComputePass((pass) => pass.dispatchWorkgroupsIndirect(rawBuffer, offset));
   }
 
   private _applyComputeState(pass: GPUComputePassEncoder): void {

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

@@ -0,0 +1,60 @@
+import type { IndirectFlag, TgpuBuffer } from '../buffer/buffer.ts';
+import { memoryLayoutOf, type PrimitiveOffsetInfo } from '../../data/offsetUtils.ts';
+import { sizeOf } from '../../data/sizeOf.ts';
+import type { BaseData } from '../../data/wgslTypes.ts';
+import { isGPUBuffer } from '../../types.ts';
+
+type IndirectOperation = 'dispatchWorkgroupsIndirect' | 'drawIndirect' | 'drawIndexedIndirect';
+const IndirectOperationToRequiredData = {
+  dispatchWorkgroupsIndirect: '3 x u32',
+  drawIndirect: '4 x u32',
+  drawIndexedIndirect: '3 x u32, i32, u32',
+} as const satisfies Record<IndirectOperation, string>;
+
+function validateIndirectBufferSize(
+  bufferSize: number,
+  offset: number,
+  requiredBytes: number,
+  operation: IndirectOperation,
+): 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}`);
+  }
+}
+
+export function resolveIndirectOffset(
+  indirectBuffer: (TgpuBuffer<BaseData> & IndirectFlag) | GPUBuffer,
+  start: PrimitiveOffsetInfo | number | undefined,
+  requiredSize: number,
+  operation: IndirectOperation,
+): number {
+  if (isGPUBuffer(indirectBuffer)) {
+    const offset = typeof start === 'number' ? start : (start?.offset ?? 0);
+    validateIndirectBufferSize(indirectBuffer.size, offset, requiredSize, operation);
+    return offset;
+  }
+
+  const offsetInfo = start
+    ? typeof start === 'number'
+      ? { offset: start, contiguous: requiredSize }
+      : start
+    : memoryLayoutOf(indirectBuffer.dataType);
+
+  const { offset, contiguous } = offsetInfo;
+
+  validateIndirectBufferSize(sizeOf(indirectBuffer.dataType), offset, requiredSize, operation);
+
+  if (contiguous < requiredSize) {
+    console.warn(
+      `${operation}: Starting at offset ${offset}, only ${contiguous} contiguous bytes are available before padding. '${operation}' requires ${requiredSize} bytes (${IndirectOperationToRequiredData[operation]}). Reading across padding may result in undefined behavior.`,
+    );
+  }
+
+  return offset;
+}

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

@@ -1,5 +1,5 @@
 import type { AnyBuiltin, OmitBuiltins } from '../../builtin.ts';
-import type { IndexFlag, TgpuBuffer, VertexFlag } from '../../core/buffer/buffer.ts';
+import type { IndexFlag, IndirectFlag, TgpuBuffer, VertexFlag } from '../../core/buffer/buffer.ts';
 import type { TgpuQuerySet } from '../../core/querySet/querySet.ts';
 import { isBuiltin } from '../../data/attributes.ts';
 import { type Disarray, getCustomLocation, type UndecorateRecord } from '../../data/dataTypes.ts';
@@ -81,6 +81,11 @@ import {
   type TimestampWritesPriors,
   triggerPerformanceCallback,
 } from './timeable.ts';
+import { type PrimitiveOffsetInfo } from '../../data/offsetUtils.ts';
+import { resolveIndirectOffset } from './pipelineUtils.ts';
+
+const DRAW_INDIRECT_SIZE = 16; // 4 x 4
+const DRAW_INDEXED_INDIRECT_SIZE = 20; // 5 x 4
 
 interface RenderPipelineInternals {
   readonly core: RenderPipelineCore;
@@ -185,11 +190,30 @@ export interface TgpuRenderPipeline<in Targets = never>
     firstInstance?: number,
   ): void;
 
-  drawIndirect(indirectBuffer: TgpuBuffer<BaseData> | GPUBuffer, indirectOffset?: GPUSize64): void;
+  /**
+   * Draws primitives using parameters read from a buffer.
+   * The buffer must contain 4 consecutive u32 values (vertexCount, instanceCount, firstVertex, firstInstance).
+   * To get the correct offset within complex data structures, use `d.memoryLayoutOf(...)`.
+   *
+   * @param indirectBuffer - Buffer marked with 'indirect' usage containing draw parameters or raw GPUBuffer
+   * @param indirectOffset - PrimitiveOffsetInfo pointing to the first draw parameter. If not provided, starts at offset 0. To obtain safe offsets, use `d.memoryLayoutOf(...)`.
+   */
+  drawIndirect(
+    indirectBuffer: (TgpuBuffer<BaseData> & IndirectFlag) | GPUBuffer,
+    indirectOffset?: PrimitiveOffsetInfo | number,
+  ): void;
 
+  /**
+   * Draws indexed primitives using parameters read from a buffer.
+   * The buffer must contain 5 consecutive 32-bit integer values (indexCount u32, instanceCount u32, firstIndex u32, baseVertex i32, firstInstance u32).
+   * To get the correct offset within complex data structures, use `d.memoryLayoutOf(...)`.
+   *
+   * @param indirectBuffer - Buffer marked with 'indirect' usage containing draw parameters or raw GPUBuffer
+   * @param indirectOffset - PrimitiveOffsetInfo pointing to the first draw parameter. If not provided, starts at offset 0. To obtain safe offsets, use `d.memoryLayoutOf(...)`.
+   */
   drawIndexedIndirect(
-    indirectBuffer: TgpuBuffer<BaseData> | GPUBuffer,
-    indirectOffset?: GPUSize64,
+    indirectBuffer: (TgpuBuffer<BaseData> & IndirectFlag) | GPUBuffer,
+    indirectOffset?: PrimitiveOffsetInfo | number,
   ): void;
 }
 
@@ -869,26 +893,32 @@ class TgpuRenderPipelineImpl implements TgpuRenderPipeline {
   }
 
   drawIndirect(
-    indirectBuffer: TgpuBuffer<BaseData> | GPUBuffer,
-    indirectOffset: GPUSize64 = 0,
+    indirectBuffer: (TgpuBuffer<BaseData> & IndirectFlag) | GPUBuffer,
+    indirectOffset?: PrimitiveOffsetInfo | number,
   ): void {
     const internals = this[$internal];
     const { root } = internals.core.options;
-    const rawBuffer = isGPUBuffer(indirectBuffer) ? indirectBuffer : root.unwrap(indirectBuffer);
+    const rawBuffer = isGPUBuffer(indirectBuffer) ? indirectBuffer : indirectBuffer.buffer;
+    const offset = resolveIndirectOffset(
+      indirectBuffer,
+      indirectOffset,
+      DRAW_INDIRECT_SIZE,
+      'drawIndirect',
+    );
 
     if (internals.priors.externalRenderEncoder) {
       if (_lastAppliedRender.get(internals.priors.externalRenderEncoder) !== this) {
         this._applyRenderState(internals.priors.externalRenderEncoder);
         _lastAppliedRender.set(internals.priors.externalRenderEncoder, this);
       }
-      internals.priors.externalRenderEncoder.drawIndirect(rawBuffer, indirectOffset);
+      internals.priors.externalRenderEncoder.drawIndirect(rawBuffer, offset);
       return;
     }
 
     if (internals.priors.externalEncoder) {
       const pass = this._createRenderPass(internals.priors.externalEncoder);
       this._applyRenderState(pass);
-      pass.drawIndirect(rawBuffer, indirectOffset);
+      pass.drawIndirect(rawBuffer, offset);
       pass.end();
       return;
     }
@@ -898,7 +928,7 @@ class TgpuRenderPipelineImpl implements TgpuRenderPipeline {
     const commandEncoder = root.device.createCommandEncoder();
     const pass = this._createRenderPass(commandEncoder);
     this._applyRenderState(pass);
-    pass.drawIndirect(rawBuffer, indirectOffset);
+    pass.drawIndirect(rawBuffer, offset);
     pass.end();
     root.device.queue.submit([commandEncoder.finish()]);
 
@@ -912,28 +942,34 @@ class TgpuRenderPipelineImpl implements TgpuRenderPipeline {
   }
 
   drawIndexedIndirect(
-    indirectBuffer: TgpuBuffer<BaseData> | GPUBuffer,
-    indirectOffset: GPUSize64 = 0,
+    indirectBuffer: (TgpuBuffer<BaseData> & IndirectFlag) | GPUBuffer,
+    indirectOffset?: PrimitiveOffsetInfo | number,
   ): void {
     const internals = this[$internal];
     const { root } = internals.core.options;
     const rawBuffer = isGPUBuffer(indirectBuffer) ? indirectBuffer : root.unwrap(indirectBuffer);
+    const offset = resolveIndirectOffset(
+      indirectBuffer,
+      indirectOffset,
+      DRAW_INDEXED_INDIRECT_SIZE,
+      'drawIndexedIndirect',
+    );
 
     if (internals.priors.externalRenderEncoder) {
       if (_lastAppliedRender.get(internals.priors.externalRenderEncoder) !== this) {
         this._applyRenderState(internals.priors.externalRenderEncoder);
         this._setIndexBuffer(internals.priors.externalRenderEncoder);
         _lastAppliedRender.set(internals.priors.externalRenderEncoder, this);
       }
-      internals.priors.externalRenderEncoder.drawIndexedIndirect(rawBuffer, indirectOffset);
+      internals.priors.externalRenderEncoder.drawIndexedIndirect(rawBuffer, offset);
       return;
     }
 
     if (internals.priors.externalEncoder) {
       const pass = this._createRenderPass(internals.priors.externalEncoder);
       this._applyRenderState(pass);
       this._setIndexBuffer(pass);
-      pass.drawIndexedIndirect(rawBuffer, indirectOffset);
+      pass.drawIndexedIndirect(rawBuffer, offset);
       pass.end();
       return;
     }
@@ -944,7 +980,7 @@ class TgpuRenderPipelineImpl implements TgpuRenderPipeline {
     const pass = this._createRenderPass(commandEncoder);
     this._applyRenderState(pass);
     this._setIndexBuffer(pass);
-    pass.drawIndexedIndirect(rawBuffer, indirectOffset);
+    pass.drawIndexedIndirect(rawBuffer, offset);
     pass.end();
     root.device.queue.submit([commandEncoder.finish()]);