Skip to content

Commit 1f5024f

Browse files
cieplypolarcieplypolar
authored
fix: Support for PrimitiveOffsetInfo in render pipeline draw...Indirect methods (#2337)
1 parent 7950868 commit 1f5024f

6 files changed

Lines changed: 350 additions & 74 deletions

File tree

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,8 @@ export const it = base
2323
},
2424
draw: vi.fn(),
2525
drawIndexed: vi.fn(),
26+
drawIndirect: vi.fn(),
27+
drawIndexedIndirect: vi.fn(),
2628
end: vi.fn(),
2729
setBindGroup: vi.fn(),
2830
setPipeline: vi.fn(),

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

Lines changed: 11 additions & 56 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
import type { AnyComputeBuiltin } from '../../builtin.ts';
22
import type { TgpuQuerySet } from '../../core/querySet/querySet.ts';
33
import { type ResolvedSnippet, snip } from '../../data/snippet.ts';
4-
import { sizeOf } from '../../data/sizeOf.ts';
54
import type { AnyWgslData } from '../../data/wgslTypes.ts';
65
import { Void } from '../../data/wgslTypes.ts';
76
import { applyBindGroups } from './applyPipelineState.ts';
@@ -19,15 +18,16 @@ import {
1918
import { isGPUCommandEncoder, isGPUComputePassEncoder } from './typeGuards.ts';
2019
import { logDataFromGPU } from '../../tgsl/consoleLog/deserializers.ts';
2120
import type { LogResources } from '../../tgsl/consoleLog/types.ts';
22-
import type { ResolutionCtx, SelfResolvable } from '../../types.ts';
21+
import { isGPUBuffer, type ResolutionCtx, type SelfResolvable } from '../../types.ts';
2322
import { wgslExtensions, wgslExtensionToFeatureName } from '../../wgslExtensions.ts';
2423
import type { IORecord } from '../function/fnTypes.ts';
2524
import type { TgpuComputeFn } from '../function/tgpuComputeFn.ts';
2625
import { namespace } from '../resolve/namespace.ts';
2726
import type { ExperimentalTgpuRoot } from '../root/rootTypes.ts';
2827
import type { TgpuSlot } from '../slot/slotTypes.ts';
2928

30-
import { memoryLayoutOf, type PrimitiveOffsetInfo } from '../../data/offsetUtils.ts';
29+
import type { PrimitiveOffsetInfo } from '../../data/offsetUtils.ts';
30+
import { resolveIndirectOffset } from './pipelineUtils.ts';
3131
import {
3232
createWithPerformanceCallback,
3333
createWithTimestampWrites,
@@ -72,11 +72,11 @@ export interface TgpuComputePipeline extends TgpuNamable, SelfResolvable, Timeab
7272
* The buffer must contain 3 consecutive u32 values (x, y, z workgroup counts).
7373
* To get the correct offset within complex data structures, use `d.memoryLayoutOf(...)`.
7474
*
75-
* @param indirectBuffer - Buffer marked with 'indirect' usage containing dispatch parameters
75+
* @param indirectBuffer - Buffer marked with 'indirect' usage containing dispatch parameters or raw GPUBuffer
7676
* @param start - PrimitiveOffsetInfo pointing to the first dispatch parameter. If not provided, starts at offset 0. To obtain safe offsets, use `d.memoryLayoutOf(...)`.
7777
*/
7878
dispatchWorkgroupsIndirect<T extends AnyWgslData>(
79-
indirectBuffer: TgpuBuffer<T> & IndirectFlag,
79+
indirectBuffer: (TgpuBuffer<T> & IndirectFlag) | GPUBuffer,
8080
start?: PrimitiveOffsetInfo | number,
8181
): void;
8282
}
@@ -113,25 +113,6 @@ type Memo = {
113113
logResources: LogResources | undefined;
114114
};
115115

116-
function validateIndirectBufferSize(
117-
bufferSize: number,
118-
offset: number,
119-
requiredBytes: number,
120-
operation: string,
121-
): void {
122-
if (offset + requiredBytes > bufferSize) {
123-
throw new Error(
124-
`Buffer too small for ${operation}. ` +
125-
`Required: ${requiredBytes} bytes at offset ${offset}, ` +
126-
`but buffer is only ${bufferSize} bytes.`,
127-
);
128-
}
129-
130-
if (offset % 4 !== 0) {
131-
throw new Error(`Indirect buffer offset must be a multiple of 4. Got: ${offset}`);
132-
}
133-
}
134-
135116
const _lastAppliedCompute = new WeakMap<GPUComputePassEncoder, TgpuComputePipelineImpl>();
136117

137118
class TgpuComputePipelineImpl implements TgpuComputePipeline {
@@ -252,46 +233,20 @@ class TgpuComputePipelineImpl implements TgpuComputePipeline {
252233
}
253234

254235
dispatchWorkgroupsIndirect<T extends AnyWgslData>(
255-
indirectBuffer: TgpuBuffer<T> & IndirectFlag,
236+
indirectBuffer: (TgpuBuffer<T> & IndirectFlag) | GPUBuffer,
256237
start?: PrimitiveOffsetInfo | number,
257238
): void {
258239
const DISPATCH_SIZE = 12; // 3 x u32 (x, y, z)
259240

260-
let offsetInfo = start ?? memoryLayoutOf(indirectBuffer.dataType);
261-
262-
if (typeof offsetInfo === 'number') {
263-
if (offsetInfo === 0) {
264-
offsetInfo = memoryLayoutOf(indirectBuffer.dataType);
265-
} else {
266-
console.warn(
267-
`dispatchWorkgroupsIndirect: Provided start offset ${offsetInfo} as a raw number. Use d.memoryLayoutOf(...) to include contiguous padding info for safer validation.`,
268-
);
269-
// When only an offset is provided, assume we have at least 12 bytes contiguous.
270-
offsetInfo = {
271-
offset: offsetInfo,
272-
contiguous: DISPATCH_SIZE,
273-
};
274-
}
275-
}
276-
277-
const { offset, contiguous } = offsetInfo;
278-
279-
validateIndirectBufferSize(
280-
sizeOf(indirectBuffer.dataType),
281-
offset,
241+
const rawBuffer = isGPUBuffer(indirectBuffer) ? indirectBuffer : indirectBuffer.buffer;
242+
const offset = resolveIndirectOffset(
243+
indirectBuffer,
244+
start,
282245
DISPATCH_SIZE,
283246
'dispatchWorkgroupsIndirect',
284247
);
285248

286-
if (contiguous < DISPATCH_SIZE) {
287-
console.warn(
288-
`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.`,
289-
);
290-
}
291-
292-
this._executeComputePass((pass) =>
293-
pass.dispatchWorkgroupsIndirect(indirectBuffer.buffer, offset),
294-
);
249+
this._executeComputePass((pass) => pass.dispatchWorkgroupsIndirect(rawBuffer, offset));
295250
}
296251

297252
private _applyComputeState(pass: GPUComputePassEncoder): void {

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

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
1+
import type { IndirectFlag, TgpuBuffer } from '../buffer/buffer.ts';
2+
import { memoryLayoutOf, type PrimitiveOffsetInfo } from '../../data/offsetUtils.ts';
3+
import { sizeOf } from '../../data/sizeOf.ts';
4+
import type { BaseData } from '../../data/wgslTypes.ts';
5+
import { isGPUBuffer } from '../../types.ts';
6+
7+
type IndirectOperation = 'dispatchWorkgroupsIndirect' | 'drawIndirect' | 'drawIndexedIndirect';
8+
const IndirectOperationToRequiredData = {
9+
dispatchWorkgroupsIndirect: '3 x u32',
10+
drawIndirect: '4 x u32',
11+
drawIndexedIndirect: '3 x u32, i32, u32',
12+
} as const satisfies Record<IndirectOperation, string>;
13+
14+
function validateIndirectBufferSize(
15+
bufferSize: number,
16+
offset: number,
17+
requiredBytes: number,
18+
operation: IndirectOperation,
19+
): void {
20+
if (offset + requiredBytes > bufferSize) {
21+
throw new Error(
22+
`Buffer too small for ${operation}. Required: ${requiredBytes} bytes at offset ${offset}, but buffer is only ${bufferSize} bytes.`,
23+
);
24+
}
25+
26+
if (offset % 4 !== 0) {
27+
throw new Error(`Indirect buffer offset must be a multiple of 4. Got: ${offset}`);
28+
}
29+
}
30+
31+
export function resolveIndirectOffset(
32+
indirectBuffer: (TgpuBuffer<BaseData> & IndirectFlag) | GPUBuffer,
33+
start: PrimitiveOffsetInfo | number | undefined,
34+
requiredSize: number,
35+
operation: IndirectOperation,
36+
): number {
37+
if (isGPUBuffer(indirectBuffer)) {
38+
const offset = typeof start === 'number' ? start : (start?.offset ?? 0);
39+
validateIndirectBufferSize(indirectBuffer.size, offset, requiredSize, operation);
40+
return offset;
41+
}
42+
43+
const offsetInfo = start
44+
? typeof start === 'number'
45+
? { offset: start, contiguous: requiredSize }
46+
: start
47+
: memoryLayoutOf(indirectBuffer.dataType);
48+
49+
const { offset, contiguous } = offsetInfo;
50+
51+
validateIndirectBufferSize(sizeOf(indirectBuffer.dataType), offset, requiredSize, operation);
52+
53+
if (contiguous < requiredSize) {
54+
console.warn(
55+
`${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.`,
56+
);
57+
}
58+
59+
return offset;
60+
}

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

Lines changed: 51 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
import type { AnyBuiltin, OmitBuiltins } from '../../builtin.ts';
2-
import type { IndexFlag, TgpuBuffer, VertexFlag } from '../../core/buffer/buffer.ts';
2+
import type { IndexFlag, IndirectFlag, TgpuBuffer, VertexFlag } from '../../core/buffer/buffer.ts';
33
import type { TgpuQuerySet } from '../../core/querySet/querySet.ts';
44
import { isBuiltin } from '../../data/attributes.ts';
55
import { type Disarray, getCustomLocation, type UndecorateRecord } from '../../data/dataTypes.ts';
@@ -81,6 +81,11 @@ import {
8181
type TimestampWritesPriors,
8282
triggerPerformanceCallback,
8383
} from './timeable.ts';
84+
import { type PrimitiveOffsetInfo } from '../../data/offsetUtils.ts';
85+
import { resolveIndirectOffset } from './pipelineUtils.ts';
86+
87+
const DRAW_INDIRECT_SIZE = 16; // 4 x 4
88+
const DRAW_INDEXED_INDIRECT_SIZE = 20; // 5 x 4
8489

8590
interface RenderPipelineInternals {
8691
readonly core: RenderPipelineCore;
@@ -185,11 +190,30 @@ export interface TgpuRenderPipeline<in Targets = never>
185190
firstInstance?: number,
186191
): void;
187192

188-
drawIndirect(indirectBuffer: TgpuBuffer<BaseData> | GPUBuffer, indirectOffset?: GPUSize64): void;
193+
/**
194+
* Draws primitives using parameters read from a buffer.
195+
* The buffer must contain 4 consecutive u32 values (vertexCount, instanceCount, firstVertex, firstInstance).
196+
* To get the correct offset within complex data structures, use `d.memoryLayoutOf(...)`.
197+
*
198+
* @param indirectBuffer - Buffer marked with 'indirect' usage containing draw parameters or raw GPUBuffer
199+
* @param indirectOffset - PrimitiveOffsetInfo pointing to the first draw parameter. If not provided, starts at offset 0. To obtain safe offsets, use `d.memoryLayoutOf(...)`.
200+
*/
201+
drawIndirect(
202+
indirectBuffer: (TgpuBuffer<BaseData> & IndirectFlag) | GPUBuffer,
203+
indirectOffset?: PrimitiveOffsetInfo | number,
204+
): void;
189205

206+
/**
207+
* Draws indexed primitives using parameters read from a buffer.
208+
* The buffer must contain 5 consecutive 32-bit integer values (indexCount u32, instanceCount u32, firstIndex u32, baseVertex i32, firstInstance u32).
209+
* To get the correct offset within complex data structures, use `d.memoryLayoutOf(...)`.
210+
*
211+
* @param indirectBuffer - Buffer marked with 'indirect' usage containing draw parameters or raw GPUBuffer
212+
* @param indirectOffset - PrimitiveOffsetInfo pointing to the first draw parameter. If not provided, starts at offset 0. To obtain safe offsets, use `d.memoryLayoutOf(...)`.
213+
*/
190214
drawIndexedIndirect(
191-
indirectBuffer: TgpuBuffer<BaseData> | GPUBuffer,
192-
indirectOffset?: GPUSize64,
215+
indirectBuffer: (TgpuBuffer<BaseData> & IndirectFlag) | GPUBuffer,
216+
indirectOffset?: PrimitiveOffsetInfo | number,
193217
): void;
194218
}
195219

@@ -869,26 +893,32 @@ class TgpuRenderPipelineImpl implements TgpuRenderPipeline {
869893
}
870894

871895
drawIndirect(
872-
indirectBuffer: TgpuBuffer<BaseData> | GPUBuffer,
873-
indirectOffset: GPUSize64 = 0,
896+
indirectBuffer: (TgpuBuffer<BaseData> & IndirectFlag) | GPUBuffer,
897+
indirectOffset?: PrimitiveOffsetInfo | number,
874898
): void {
875899
const internals = this[$internal];
876900
const { root } = internals.core.options;
877-
const rawBuffer = isGPUBuffer(indirectBuffer) ? indirectBuffer : root.unwrap(indirectBuffer);
901+
const rawBuffer = isGPUBuffer(indirectBuffer) ? indirectBuffer : indirectBuffer.buffer;
902+
const offset = resolveIndirectOffset(
903+
indirectBuffer,
904+
indirectOffset,
905+
DRAW_INDIRECT_SIZE,
906+
'drawIndirect',
907+
);
878908

879909
if (internals.priors.externalRenderEncoder) {
880910
if (_lastAppliedRender.get(internals.priors.externalRenderEncoder) !== this) {
881911
this._applyRenderState(internals.priors.externalRenderEncoder);
882912
_lastAppliedRender.set(internals.priors.externalRenderEncoder, this);
883913
}
884-
internals.priors.externalRenderEncoder.drawIndirect(rawBuffer, indirectOffset);
914+
internals.priors.externalRenderEncoder.drawIndirect(rawBuffer, offset);
885915
return;
886916
}
887917

888918
if (internals.priors.externalEncoder) {
889919
const pass = this._createRenderPass(internals.priors.externalEncoder);
890920
this._applyRenderState(pass);
891-
pass.drawIndirect(rawBuffer, indirectOffset);
921+
pass.drawIndirect(rawBuffer, offset);
892922
pass.end();
893923
return;
894924
}
@@ -898,7 +928,7 @@ class TgpuRenderPipelineImpl implements TgpuRenderPipeline {
898928
const commandEncoder = root.device.createCommandEncoder();
899929
const pass = this._createRenderPass(commandEncoder);
900930
this._applyRenderState(pass);
901-
pass.drawIndirect(rawBuffer, indirectOffset);
931+
pass.drawIndirect(rawBuffer, offset);
902932
pass.end();
903933
root.device.queue.submit([commandEncoder.finish()]);
904934

@@ -912,28 +942,34 @@ class TgpuRenderPipelineImpl implements TgpuRenderPipeline {
912942
}
913943

914944
drawIndexedIndirect(
915-
indirectBuffer: TgpuBuffer<BaseData> | GPUBuffer,
916-
indirectOffset: GPUSize64 = 0,
945+
indirectBuffer: (TgpuBuffer<BaseData> & IndirectFlag) | GPUBuffer,
946+
indirectOffset?: PrimitiveOffsetInfo | number,
917947
): void {
918948
const internals = this[$internal];
919949
const { root } = internals.core.options;
920950
const rawBuffer = isGPUBuffer(indirectBuffer) ? indirectBuffer : root.unwrap(indirectBuffer);
951+
const offset = resolveIndirectOffset(
952+
indirectBuffer,
953+
indirectOffset,
954+
DRAW_INDEXED_INDIRECT_SIZE,
955+
'drawIndexedIndirect',
956+
);
921957

922958
if (internals.priors.externalRenderEncoder) {
923959
if (_lastAppliedRender.get(internals.priors.externalRenderEncoder) !== this) {
924960
this._applyRenderState(internals.priors.externalRenderEncoder);
925961
this._setIndexBuffer(internals.priors.externalRenderEncoder);
926962
_lastAppliedRender.set(internals.priors.externalRenderEncoder, this);
927963
}
928-
internals.priors.externalRenderEncoder.drawIndexedIndirect(rawBuffer, indirectOffset);
964+
internals.priors.externalRenderEncoder.drawIndexedIndirect(rawBuffer, offset);
929965
return;
930966
}
931967

932968
if (internals.priors.externalEncoder) {
933969
const pass = this._createRenderPass(internals.priors.externalEncoder);
934970
this._applyRenderState(pass);
935971
this._setIndexBuffer(pass);
936-
pass.drawIndexedIndirect(rawBuffer, indirectOffset);
972+
pass.drawIndexedIndirect(rawBuffer, offset);
937973
pass.end();
938974
return;
939975
}
@@ -944,7 +980,7 @@ class TgpuRenderPipelineImpl implements TgpuRenderPipeline {
944980
const pass = this._createRenderPass(commandEncoder);
945981
this._applyRenderState(pass);
946982
this._setIndexBuffer(pass);
947-
pass.drawIndexedIndirect(rawBuffer, indirectOffset);
983+
pass.drawIndexedIndirect(rawBuffer, offset);
948984
pass.end();
949985
root.device.queue.submit([commandEncoder.finish()]);
950986

0 commit comments

Comments
 (0)