dynamis_gpu/
compute.rs

1use wgpu::{
2    BindGroup, BindGroupEntry, BindGroupLayout, BindGroupLayoutDescriptor, BindGroupLayoutEntry,
3    BindingType, BufferBindingType, CommandEncoder, ComputePassDescriptor,
4    ComputePassTimestampWrites, ComputePipeline as WgpuComputePipeline, ComputePipelineDescriptor,
5    Device, PipelineLayoutDescriptor, ShaderModuleDescriptor, ShaderSource, ShaderStages,
6};
7
8pub struct ComputeRecorder<'a> {
9    pass: wgpu::ComputePass<'a>,
10}
11
12impl<'a> ComputeRecorder<'a> {
13    pub fn begin(encoder: &'a mut CommandEncoder, label: &'a str) -> Self {
14        Self::begin_timed(encoder, label, None)
15    }
16
17    pub fn begin_timed(
18        encoder: &'a mut CommandEncoder,
19        label: &'a str,
20        timing: Option<ComputePassTimestampWrites<'a>>,
21    ) -> Self {
22        let pass = encoder.begin_compute_pass(&ComputePassDescriptor {
23            label: Some(label),
24            timestamp_writes: timing,
25        });
26        Self { pass }
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.pipeline());
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.dispatch_workgroups(count, 1, 1);
38    }
39}
40
41#[derive(Clone, Copy, PartialEq, Eq, Hash)]
42pub enum BindingKind {
43    Uniform,
44    ReadOnlyStorage,
45    ReadWriteStorage,
46}
47
48pub struct BindingSpec {
49    pub binding: u32,
50    pub kind: BindingKind,
51}
52
53#[derive(Clone)]
54pub struct ComputePipeline {
55    pipeline: WgpuComputePipeline,
56    bind_group_layouts: Vec<BindGroupLayout>,
57}
58
59impl ComputePipeline {
60    pub fn new(
61        device: &Device,
62        label: &str,
63        shader: &str,
64        entry: &str,
65        groups: &[&[BindingSpec]],
66    ) -> Self {
67        let module = device.create_shader_module(ShaderModuleDescriptor {
68            label: Some(label),
69            source: ShaderSource::Wgsl(shader.into()),
70        });
71        let bind_group_layouts = groups
72            .iter()
73            .map(|bindings| {
74                let entries: Vec<BindGroupLayoutEntry> = bindings
75                    .iter()
76                    .map(|spec| BindGroupLayoutEntry {
77                        binding: spec.binding,
78                        visibility: ShaderStages::COMPUTE,
79                        ty: match spec.kind {
80                            BindingKind::Uniform => BindingType::Buffer {
81                                ty: BufferBindingType::Uniform,
82                                has_dynamic_offset: false,
83                                min_binding_size: None,
84                            },
85                            BindingKind::ReadOnlyStorage => BindingType::Buffer {
86                                ty: BufferBindingType::Storage { read_only: true },
87                                has_dynamic_offset: false,
88                                min_binding_size: None,
89                            },
90                            BindingKind::ReadWriteStorage => BindingType::Buffer {
91                                ty: BufferBindingType::Storage { read_only: false },
92                                has_dynamic_offset: false,
93                                min_binding_size: None,
94                            },
95                        },
96                        count: None,
97                    })
98                    .collect();
99                device.create_bind_group_layout(&BindGroupLayoutDescriptor {
100                    label: Some(label),
101                    entries: &entries,
102                })
103            })
104            .collect::<Vec<_>>();
105        let layouts = bind_group_layouts.iter().map(Some).collect::<Vec<_>>();
106        let pipeline_layout = device.create_pipeline_layout(&PipelineLayoutDescriptor {
107            label: Some(label),
108            bind_group_layouts: &layouts,
109            immediate_size: 0,
110        });
111        let pipeline = device.create_compute_pipeline(&ComputePipelineDescriptor {
112            label: Some(label),
113            layout: Some(&pipeline_layout),
114            module: &module,
115            entry_point: Some(entry),
116            compilation_options: Default::default(),
117            cache: None,
118        });
119        Self {
120            pipeline,
121            bind_group_layouts,
122        }
123    }
124
125    pub fn pipeline(&self) -> &WgpuComputePipeline {
126        &self.pipeline
127    }
128
129    pub fn create_bind_group(
130        &self,
131        device: &Device,
132        group: usize,
133        entries: &[BindGroupEntry<'_>],
134    ) -> BindGroup {
135        device.create_bind_group(&wgpu::BindGroupDescriptor {
136            label: None,
137            layout: &self.bind_group_layouts[group],
138            entries,
139        })
140    }
141}
dynamis_gpu/compute.rs

dynamis_gpu/
compute.rs
