Skip to content

Commit 85f5ff0

Browse files
cieplypolarcieplypolar
committed
Batch implemented with tests
1 parent b39c1b3 commit 85f5ff0

File tree

6 files changed

+333
-12
lines changed

6 files changed

+333
-12
lines changed

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -621,6 +621,8 @@ class TgpuRenderPipelineImpl implements TgpuRenderPipeline {
621621
root: branch,
622622
priors: internals.priors,
623623
})
624+
: branch[$internal].ongoingBatch
625+
? ({})
624626
: branch[$internal].flush();
625627
}
626628
}

‎packages/typegpu/src/core/root/init.ts‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -326,15 +326,14 @@ class TgpuRootImpl extends WithBindingImpl
326326
}(device);
327327
}
328328

329-
batch(callback: () => void): Promise<undefined> {
329+
batch(callback: () => void) {
330330
this[$internal].ongoingBatch = true;
331331
try {
332332
callback();
333333
} finally {
334334
this[$internal].ongoingBatch = false;
335335
this[$internal].flush();
336336
}
337-
return this.device.queue.onSubmittedWorkDone();
338337
}
339338

340339
get enabledFeatures() {

‎packages/typegpu/src/core/root/rootTypes.ts‎

Lines changed: 6 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -626,16 +626,17 @@ export interface TgpuRootInternals {
626626
*/
627627
ongoingBatch: boolean;
628628
/**
629-
* The current command encoder. This property will
630-
* hold the same value throughout the entire `batch()` invocation.
631-
* In case of single `draw()` or `dispatchWorkgroups()` call, getter will be used
629+
* The current command encoder. This property
630+
* holds the same value throughout the entire `batch()` invocation,
631+
* unless you use pipeline with performance callback.
632+
* In case of single `draw()` or `drawIndexed()` or `dispatchWorkgroups()` call, getter will be used
632633
* to create a single-use command encoder.
633634
*/
634635
readonly commandEncoder: GPUCommandEncoder;
635636
/**
636637
* Causes all commands enqueued by pipelines to be
637638
* submitted to the GPU.
638-
* If a single `draw()` or `dispatchWorkgroups()` is called, `flush()` is executed after each command.
639+
* If a single `draw()` or `drawIndexed()` or `dispatchWorkgroups()` is called, `flush()` is executed after each command.
639640
*/
640641
flush(): void;
641642
}
@@ -694,8 +695,6 @@ export interface ExperimentalTgpuRoot extends TgpuRoot, WithBinding {
694695
* performance callbacks/timestamp writes are flushed immediately.
695696
*
696697
* @param callback A function with GPU computations to be batched.
697-
*
698-
* Returns a Promise that resolves when the batch is completed.
699698
*/
700-
batch(callback: () => void): Promise<undefined>;
699+
batch(callback: () => void): void;
701700
}

‎packages/typegpu/tests/batch.test.ts‎

Lines changed: 241 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,241 @@
1+
import { describe, expect, vi } from 'vitest';
2+
import * as d from '../src/data/index.ts';
3+
import tgpu from '../src/index.ts';
4+
import { $internal } from '../src/shared/symbols.ts';
5+
import { it } from './utils/extendedIt.ts';
6+
7+
describe('Batch', () => {
8+
const entryFn = tgpu['~unstable'].computeFn({ workgroupSize: [7] })(() => {});
9+
const vertexFn = tgpu['~unstable'].vertexFn({
10+
out: { pos: d.builtin.position },
11+
})(() => {
12+
return {
13+
pos: d.vec4f(),
14+
};
15+
});
16+
const fragmentFn = tgpu['~unstable'].fragmentFn({
17+
out: d.vec4f,
18+
})(() => d.vec4f());
19+
20+
it('flushes only once when used without performance callback', ({ root }) => {
21+
const renderPipeline1 = root
22+
.withVertex(vertexFn, {})
23+
.withFragment(fragmentFn, { format: 'rgba8unorm' })
24+
.createPipeline()
25+
.withColorAttachment({
26+
view: {} as unknown as GPUTextureView,
27+
loadOp: 'clear',
28+
storeOp: 'store',
29+
});
30+
31+
const renderPipeline2 = root
32+
.withVertex(vertexFn, {})
33+
.withFragment(fragmentFn, { format: 'rgba8unorm' })
34+
.createPipeline()
35+
.withColorAttachment({
36+
view: {} as unknown as GPUTextureView,
37+
loadOp: 'clear',
38+
storeOp: 'store',
39+
});
40+
41+
const computePipeline = root
42+
.withCompute(entryFn)
43+
.createPipeline();
44+
45+
vi.spyOn(root[$internal], 'flush');
46+
47+
root.batch(() => {
48+
renderPipeline1.draw(7);
49+
computePipeline.dispatchWorkgroups(7);
50+
renderPipeline2.draw(7);
51+
expect(root[$internal].flush).toBeCalledTimes(0);
52+
});
53+
54+
expect(root[$internal].flush).toBeCalledTimes(1);
55+
});
56+
57+
it('flushes immediatelly when used with performance callback', ({ root }) => {
58+
const querySet = root.createQuerySet('timestamp', 2);
59+
const callback = vi.fn();
60+
61+
const renderPipelineWithPerformance = root
62+
.withVertex(vertexFn, {})
63+
.withFragment(fragmentFn, { format: 'rgba8unorm' })
64+
.createPipeline()
65+
.withColorAttachment({
66+
view: {} as unknown as GPUTextureView,
67+
loadOp: 'clear',
68+
storeOp: 'store',
69+
})
70+
.withPerformanceCallback(callback);
71+
72+
const renderPipelineWithTimestampWrites = root
73+
.withVertex(vertexFn, {})
74+
.withFragment(fragmentFn, { format: 'rgba8unorm' })
75+
.createPipeline()
76+
.withColorAttachment({
77+
view: {} as unknown as GPUTextureView,
78+
loadOp: 'clear',
79+
storeOp: 'store',
80+
})
81+
.withTimestampWrites({
82+
querySet,
83+
beginningOfPassWriteIndex: 0,
84+
endOfPassWriteIndex: 1,
85+
});
86+
87+
const computePipeline = root
88+
.withCompute(entryFn)
89+
.createPipeline();
90+
91+
vi.spyOn(root[$internal], 'flush');
92+
93+
// trying different permutations
94+
root.batch(() => {
95+
computePipeline.dispatchWorkgroups(7);
96+
expect(root[$internal].flush).toBeCalledTimes(0);
97+
renderPipelineWithPerformance.draw(7);
98+
expect(root[$internal].flush).toBeCalledTimes(1);
99+
renderPipelineWithTimestampWrites.draw(7);
100+
expect(root[$internal].flush).toBeCalledTimes(1);
101+
});
102+
expect(root[$internal].flush).toBeCalledTimes(2);
103+
104+
vi.spyOn(root[$internal], 'flush');
105+
106+
root.batch(() => {
107+
renderPipelineWithPerformance.draw(7);
108+
expect(root[$internal].flush).toBeCalledTimes(1);
109+
computePipeline.dispatchWorkgroups(7);
110+
expect(root[$internal].flush).toBeCalledTimes(1);
111+
renderPipelineWithTimestampWrites.draw(7);
112+
expect(root[$internal].flush).toBeCalledTimes(1);
113+
});
114+
expect(root[$internal].flush).toBeCalledTimes(2);
115+
116+
vi.spyOn(root[$internal], 'flush');
117+
118+
root.batch(() => {
119+
renderPipelineWithTimestampWrites.draw(7);
120+
expect(root[$internal].flush).toBeCalledTimes(0);
121+
computePipeline.dispatchWorkgroups(7);
122+
expect(root[$internal].flush).toBeCalledTimes(0);
123+
renderPipelineWithPerformance.draw(7);
124+
expect(root[$internal].flush).toBeCalledTimes(1);
125+
});
126+
expect(root[$internal].flush).toBeCalledTimes(2);
127+
});
128+
129+
it('flushes properly with drawIndexed', ({ root }) => {
130+
const callback = vi.fn();
131+
const querySet = root.createQuerySet('timestamp', 2);
132+
const indexBuffer = root.createBuffer(d.arrayOf(d.u16, 2)).$usage('index');
133+
134+
const renderPipeline1 = root
135+
.withVertex(vertexFn, {})
136+
.withFragment(fragmentFn, { format: 'rgba8unorm' })
137+
.createPipeline()
138+
.withColorAttachment({
139+
view: {} as unknown as GPUTextureView,
140+
loadOp: 'clear',
141+
storeOp: 'store',
142+
})
143+
.withIndexBuffer(indexBuffer);
144+
145+
const renderPipeline2 = root
146+
.withVertex(vertexFn, {})
147+
.withFragment(fragmentFn, { format: 'rgba8unorm' })
148+
.createPipeline()
149+
.withColorAttachment({
150+
view: {} as unknown as GPUTextureView,
151+
loadOp: 'clear',
152+
storeOp: 'store',
153+
})
154+
.withPerformanceCallback(callback)
155+
.withIndexBuffer(indexBuffer);
156+
157+
const renderPipeline3 = root
158+
.withVertex(vertexFn, {})
159+
.withFragment(fragmentFn, { format: 'rgba8unorm' })
160+
.createPipeline()
161+
.withColorAttachment({
162+
view: {} as unknown as GPUTextureView,
163+
loadOp: 'clear',
164+
storeOp: 'store',
165+
})
166+
.withTimestampWrites({
167+
querySet,
168+
beginningOfPassWriteIndex: 0,
169+
endOfPassWriteIndex: 1,
170+
})
171+
.withIndexBuffer(indexBuffer);
172+
173+
vi.spyOn(root[$internal], 'flush');
174+
175+
root.batch(() => {
176+
renderPipeline1.drawIndexed(7);
177+
expect(root[$internal].flush).toBeCalledTimes(0);
178+
renderPipeline2.drawIndexed(7);
179+
expect(root[$internal].flush).toBeCalledTimes(1);
180+
renderPipeline3.drawIndexed(7);
181+
expect(root[$internal].flush).toBeCalledTimes(1);
182+
});
183+
184+
expect(root[$internal].flush).toBeCalledTimes(2);
185+
});
186+
187+
it('flushes properly when used with beginRenderPass', ({ root }) => {
188+
const renderPipeline1 = root
189+
.withVertex(vertexFn, {})
190+
.withFragment(fragmentFn, { format: 'rgba8unorm' })
191+
.createPipeline();
192+
193+
const bindGroupLayout = tgpu.bindGroupLayout({});
194+
const bindGroup = root.createBindGroup(bindGroupLayout, {});
195+
196+
vi.spyOn(root[$internal], 'flush');
197+
198+
root['~unstable'].beginRenderPass(
199+
{ colorAttachments: [] },
200+
(pass) => {
201+
pass.setPipeline(renderPipeline1);
202+
pass.setBindGroup(bindGroupLayout, bindGroup);
203+
pass.draw(7);
204+
},
205+
);
206+
expect(root[$internal].flush).toBeCalledTimes(1);
207+
root['~unstable'].beginRenderPass(
208+
{ colorAttachments: [] },
209+
(pass) => {
210+
pass.setPipeline(renderPipeline1);
211+
pass.setBindGroup(bindGroupLayout, bindGroup);
212+
pass.draw(7);
213+
},
214+
);
215+
expect(root[$internal].flush).toBeCalledTimes(2);
216+
217+
vi.spyOn(root[$internal], 'flush');
218+
219+
root['~unstable'].batch(() => {
220+
root['~unstable'].beginRenderPass(
221+
{ colorAttachments: [] },
222+
(pass) => {
223+
pass.setPipeline(renderPipeline1);
224+
pass.setBindGroup(bindGroupLayout, bindGroup);
225+
pass.draw(7);
226+
},
227+
);
228+
expect(root[$internal].flush).toBeCalledTimes(0);
229+
root['~unstable'].beginRenderPass(
230+
{ colorAttachments: [] },
231+
(pass) => {
232+
pass.setPipeline(renderPipeline1);
233+
pass.setBindGroup(bindGroupLayout, bindGroup);
234+
pass.draw(7);
235+
},
236+
);
237+
expect(root[$internal].flush).toBeCalledTimes(0);
238+
});
239+
expect(root[$internal].flush).toBeCalledTimes(1);
240+
});
241+
});

‎packages/typegpu/tests/computePipeline.test.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -563,7 +563,7 @@ describe('TgpuComputePipeline', () => {
563563
expect(root[$internal].flush).toBeCalledTimes(1);
564564
});
565565
it('flushes after draw with timestamp writes', ({ root }) => {
566-
const querySet = root.createQuerySet('timestamp', 4);
566+
const querySet = root.createQuerySet('timestamp', 2);
567567

568568
const pipeline = root
569569
.withCompute(entryFn)

0 commit comments

Comments
 (0)