dynamis_gpu/
compute.rs

1use crate::buffer::GpuBuffer;
2use wgpu::{
3    BindGroup, BindGroupEntry, BindGroupLayout, BindGroupLayoutDescriptor, BindGroupLayoutEntry,
4    BindingType, BufferBindingType, CommandEncoder, ComputePassDescriptor,
5    ComputePipeline as WgpuComputePipeline, ComputePipelineDescriptor, Device,
6    PipelineLayoutDescriptor, ShaderModuleDescriptor, ShaderSource, ShaderStages,
7};
8
9/// Records every dispatch of a physics frame into a single compute pass.
10///
11/// The runtime tracks buffer usage per dispatch and inserts the required
12/// barriers between dispatches inside the pass, so splitting a frame into
13/// dozens of one-dispatch passes buys nothing but submission overhead.
14/// A recorder owns the pass until dropped, at which point the encoder is
15/// released for copies and buffer operations.
16pub struct ComputeRecorder<'a> {
17    pass: wgpu::ComputePass<'a>,
18}
19
20impl<'a> ComputeRecorder<'a> {
21    pub fn begin(encoder: &'a mut CommandEncoder, label: &str) -> Self {
22        let pass = encoder.begin_compute_pass(&ComputePassDescriptor {
23            label: Some(label),
24            timestamp_writes: None,
25        });
26        Self { pass }
27    }
28
29    pub fn record(&mut self, pipeline: &ComputePipeline, bind_groups: &[&BindGroup], count: u32) {
30        self.pass.set_pipeline(pipeline.pipeline());
31        for (group, bind_group) in bind_groups.iter().enumerate() {
32            self.pass.set_bind_group(group as u32, *bind_group, &[]);
33        }
34        self.pass.dispatch_workgroups(count, 1, 1);
35    }
36
37    pub fn record_indirect(
38        &mut self,
39        pipeline: &ComputePipeline,
40        bind_groups: &[&BindGroup],
41        args: &GpuBuffer,
42        offset: u64,
43    ) {
44        self.pass.set_pipeline(pipeline.pipeline());
45        for (group, bind_group) in bind_groups.iter().enumerate() {
46            self.pass.set_bind_group(group as u32, *bind_group, &[]);
47        }
48        self.pass
49            .dispatch_workgroups_indirect(args.as_indirect_args(), offset);
50    }
51}
52
53#[derive(Clone, Copy, PartialEq, Eq, Hash)]
54pub enum BindingKind {
55    Uniform,
56    ReadOnlyStorage,
57    ReadWriteStorage,
58}
59
60pub struct BindingSpec {
61    pub binding: u32,
62    pub kind: BindingKind,
63}
64
65#[derive(Clone)]
66pub struct ComputePipeline {
67    pipeline: WgpuComputePipeline,
68    bind_group_layouts: Vec<BindGroupLayout>,
69    workgroup_size: u32,
70}
71
72impl ComputePipeline {
73    pub fn new(
74        device: &Device,
75        label: &str,
76        shader: &str,
77        entry: &str,
78        groups: &[&[BindingSpec]],
79        workgroup_size: u32,
80    ) -> Self {
81        let module = device.create_shader_module(ShaderModuleDescriptor {
82            label: Some(label),
83            source: ShaderSource::Wgsl(shader.into()),
84        });
85        let bind_group_layouts = groups
86            .iter()
87            .map(|bindings| {
88                let entries: Vec<BindGroupLayoutEntry> = bindings
89                    .iter()
90                    .map(|spec| BindGroupLayoutEntry {
91                        binding: spec.binding,
92                        visibility: ShaderStages::COMPUTE,
93                        ty: match spec.kind {
94                            BindingKind::Uniform => BindingType::Buffer {
95                                ty: BufferBindingType::Uniform,
96                                has_dynamic_offset: false,
97                                min_binding_size: None,
98                            },
99                            BindingKind::ReadOnlyStorage => BindingType::Buffer {
100                                ty: BufferBindingType::Storage { read_only: true },
101                                has_dynamic_offset: false,
102                                min_binding_size: None,
103                            },
104                            BindingKind::ReadWriteStorage => BindingType::Buffer {
105                                ty: BufferBindingType::Storage { read_only: false },
106                                has_dynamic_offset: false,
107                                min_binding_size: None,
108                            },
109                        },
110                        count: None,
111                    })
112                    .collect();
113                device.create_bind_group_layout(&BindGroupLayoutDescriptor {
114                    label: Some(label),
115                    entries: &entries,
116                })
117            })
118            .collect::<Vec<_>>();
119        let layouts = bind_group_layouts.iter().map(Some).collect::<Vec<_>>();
120        let pipeline_layout = device.create_pipeline_layout(&PipelineLayoutDescriptor {
121            label: Some(label),
122            bind_group_layouts: &layouts,
123            immediate_size: 0,
124        });
125        let pipeline = device.create_compute_pipeline(&ComputePipelineDescriptor {
126            label: Some(label),
127            layout: Some(&pipeline_layout),
128            module: &module,
129            entry_point: Some(entry),
130            compilation_options: Default::default(),
131            cache: None,
132        });
133        Self {
134            pipeline,
135            bind_group_layouts,
136            workgroup_size,
137        }
138    }
139
140    pub fn bind_group_layout(&self, group: usize) -> &BindGroupLayout {
141        &self.bind_group_layouts[group]
142    }
143
144    pub fn pipeline(&self) -> &WgpuComputePipeline {
145        &self.pipeline
146    }
147
148    pub fn workgroup_size(&self) -> u32 {
149        self.workgroup_size
150    }
151
152    pub fn create_bind_group(
153        &self,
154        device: &Device,
155        group: usize,
156        entries: &[BindGroupEntry<'_>],
157    ) -> BindGroup {
158        device.create_bind_group(&wgpu::BindGroupDescriptor {
159            label: None,
160            layout: &self.bind_group_layouts[group],
161            entries,
162        })
163    }
164
165    pub fn workgroup_count(&self, elements: u32) -> u32 {
166        elements.div_ceil(self.workgroup_size)
167    }
168}
dynamis_gpu/compute.rs

dynamis_gpu/
compute.rs
