dynamis_gpu/
count_args.rs

1use crate::buffer::GpuBuffer;
2use crate::{BindingKind, BindingSpec, ComputePipeline, ComputeRecorder, GpuContext};
3use wgpu::{BindGroup, BindGroupEntry, Device};
4
5/// Converts a GPU-written element count into `dispatch_workgroups_indirect`
6/// arguments for both the 64-thread and 256-thread workgroup sizes.
7///
8/// Every simulation stage launches workgroups from a runtime count rather
9/// than from its pre-allocated capacity; without this conversion each
10/// indirect dispatch would schedule one workgroup per element, idling
11/// all but the first lane of every workgroup.
12pub struct GpuCountArgs {
13    pipeline: ComputePipeline,
14    bindings: std::sync::Mutex<Option<(u64, u64, BindGroup)>>,
15}
16
17impl GpuCountArgs {
18    pub fn new(context: &GpuContext, label: &str) -> Self {
19        let shader = "@group(0) @binding(0) var<storage, read> count: array<u32>;\n\
20                      @group(0) @binding(1) var<storage, read_write> args: array<u32>;\n\n\
21                      @compute @workgroup_size(1u)\n\
22                      fn main() {\n\
23                          let elements = count[0];\n\
24                          args[0] = (elements + 63u) / 64u;\n\
25                          args[1] = 1u;\n\
26                          args[2] = 1u;\n\
27                          args[3] = 0u;\n\
28                          args[4] = (elements + 255u) / 256u;\n\
29                          args[5] = 1u;\n\
30                          args[6] = 1u;\n\
31                          args[7] = 0u;\n\
32                      }\n";
33        let specs = [
34            BindingSpec {
35                binding: 0,
36                kind: BindingKind::ReadOnlyStorage,
37            },
38            BindingSpec {
39                binding: 1,
40                kind: BindingKind::ReadWriteStorage,
41            },
42        ];
43        Self {
44            pipeline: context.compute_pipeline(label, shader, "main", &[&specs[..]], 1),
45            bindings: std::sync::Mutex::new(None),
46        }
47    }
48
49    fn bindings(&self, device: &Device, count: &GpuBuffer, args: &GpuBuffer) -> BindGroup {
50        let mut guard = self.bindings.lock().unwrap();
51        let key = (count.token(), args.token());
52        if guard
53            .as_ref()
54            .is_none_or(|(count_key, args_key, _)| (*count_key, *args_key) != key)
55        {
56            let group = self.pipeline.create_bind_group(
57                device,
58                0,
59                &[
60                    BindGroupEntry {
61                        binding: 0,
62                        resource: count.as_binding(),
63                    },
64                    BindGroupEntry {
65                        binding: 1,
66                        resource: args.as_binding(),
67                    },
68                ],
69            );
70            *guard = Some((count.token(), args.token(), group));
71        }
72        guard
73            .as_ref()
74            .expect("bindings ensured just above")
75            .2
76            .clone()
77    }
78
79    pub fn encode(
80        &self,
81        device: &Device,
82        recorder: &mut ComputeRecorder,
83        count: &GpuBuffer,
84        args: &GpuBuffer,
85    ) {
86        let group = self.bindings(device, count, args);
87        recorder.record(&self.pipeline, &[&group], 1);
88    }
89}
dynamis_gpu/count_args.rs

dynamis_gpu/
count_args.rs
