1use crate::ComputePipeline;
2use wgpu::{
3 BindGroup, CommandEncoder, ComputePass, ComputePassDescriptor, ComputePassTimestampWrites,
4};
5
6pub struct ComputeRecorder<'a> {
7 pass: ComputePass<'a>,
8 per_row: u32,
9}
10
11impl<'a> ComputeRecorder<'a> {
12 pub fn begin(encoder: &'a mut CommandEncoder, label: &'a str, per_row: u32) -> Self {
13 Self::begin_timed(encoder, label, None, per_row)
14 }
15
16 pub fn begin_timed(
17 encoder: &'a mut CommandEncoder,
18 label: &'a str,
19 timing: Option<ComputePassTimestampWrites<'a>>,
20 per_row: u32,
21 ) -> Self {
22 let pass = encoder.begin_compute_pass(&ComputePassDescriptor {
23 label: Some(label),
24 timestamp_writes: timing,
25 });
26 Self { pass, per_row }
27 }
28
29 pub fn record(&mut self, pipeline: &ComputePipeline, bind_groups: &[&BindGroup], count: u32) {
30 if count == 0 {
31 return;
32 }
33 self.pass.set_pipeline(pipeline.wgpu());
34 for (group, bind_group) in bind_groups.iter().enumerate() {
35 self.pass.set_bind_group(group as u32, *bind_group, &[]);
36 }
37 self.pass
38 .dispatch_workgroups(count.min(self.per_row), count.div_ceil(self.per_row), 1);
39 }
40
41 pub fn record_indirect(
42 &mut self,
43 pipeline: &ComputePipeline,
44 bind_groups: &[&BindGroup],
45 table: &crate::DispatchTable,
46 slot: u32,
47 ) {
48 self.pass.set_pipeline(pipeline.wgpu());
49 for (group, bind_group) in bind_groups.iter().enumerate() {
50 self.pass.set_bind_group(group as u32, *bind_group, &[]);
51 }
52 self.pass
53 .dispatch_workgroups_indirect(table.buffer().buffer(), table.offset(slot));
54 }
55}