dynamis_gpu/
pipeline.rs

1use std::sync::{Arc, OnceLock};
2use wgpu::{
3    BindGroup, BindGroupEntry, BindGroupLayout, BindGroupLayoutDescriptor, BindGroupLayoutEntry,
4    BindingType, BufferBindingType, ComputePipeline as WgpuComputePipeline,
5    ComputePipelineDescriptor, Device, PipelineLayout, PipelineLayoutDescriptor,
6    ShaderModuleDescriptor, ShaderSource, ShaderStages,
7};
8
9#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
10pub enum BindingKind {
11    Uniform,
12    ReadOnlyStorage,
13    ReadWriteStorage,
14}
15
16impl BindingKind {
17    fn layout_entry(self, binding: u32) -> BindGroupLayoutEntry {
18        BindGroupLayoutEntry {
19            binding,
20            visibility: ShaderStages::COMPUTE,
21            ty: match self {
22                Self::Uniform => BindingType::Buffer {
23                    ty: BufferBindingType::Uniform,
24                    has_dynamic_offset: false,
25                    min_binding_size: None,
26                },
27                Self::ReadOnlyStorage => BindingType::Buffer {
28                    ty: BufferBindingType::Storage { read_only: true },
29                    has_dynamic_offset: false,
30                    min_binding_size: None,
31                },
32                Self::ReadWriteStorage => BindingType::Buffer {
33                    ty: BufferBindingType::Storage { read_only: false },
34                    has_dynamic_offset: false,
35                    min_binding_size: None,
36                },
37            },
38            count: None,
39        }
40    }
41}
42
43#[derive(Clone, Copy, PartialEq, Eq, Hash)]
44pub struct BindingSpec {
45    pub binding: u32,
46    pub kind: BindingKind,
47}
48
49#[derive(PartialEq, Eq, Hash)]
50pub struct ComputeProgram {
51    label: String,
52    shader: Arc<str>,
53    entry: String,
54    groups: Vec<Vec<BindingSpec>>,
55}
56
57impl ComputeProgram {
58    pub fn new(
59        label: &str,
60        shader: impl Into<Arc<str>>,
61        entry: &str,
62        groups: &[&[BindingSpec]],
63    ) -> Self {
64        Self {
65            label: label.to_owned(),
66            shader: shader.into(),
67            entry: entry.to_owned(),
68            groups: groups.iter().map(|group| group.to_vec()).collect(),
69        }
70    }
71
72    pub fn label(&self) -> &str {
73        &self.label
74    }
75}
76
77pub struct ComputeLayout {
78    pipeline: PipelineLayout,
79    groups: Vec<BindGroupLayout>,
80}
81
82impl ComputeLayout {
83    pub(crate) fn new(device: &Device, program: &ComputeProgram) -> Self {
84        let groups = program
85            .groups
86            .iter()
87            .map(|bindings| {
88                let entries = bindings
89                    .iter()
90                    .map(|spec| spec.kind.layout_entry(spec.binding))
91                    .collect::<Vec<_>>();
92                device.create_bind_group_layout(&BindGroupLayoutDescriptor {
93                    label: Some(program.label()),
94                    entries: &entries,
95                })
96            })
97            .collect::<Vec<_>>();
98        let layouts = groups.iter().map(Some).collect::<Vec<_>>();
99        let pipeline = device.create_pipeline_layout(&PipelineLayoutDescriptor {
100            label: Some(program.label()),
101            bind_group_layouts: &layouts,
102            immediate_size: 0,
103        });
104        Self { pipeline, groups }
105    }
106
107    pub fn create_bind_group(
108        &self,
109        device: &Device,
110        group: usize,
111        entries: &[BindGroupEntry<'_>],
112    ) -> BindGroup {
113        device.create_bind_group(&wgpu::BindGroupDescriptor {
114            label: None,
115            layout: &self.groups[group],
116            entries,
117        })
118    }
119}
120
121pub struct ComputePipeline {
122    pipeline: WgpuComputePipeline,
123}
124
125impl ComputePipeline {
126    fn compile(device: &Device, program: &ComputeProgram, layout: &ComputeLayout) -> Self {
127        let module = device.create_shader_module(ShaderModuleDescriptor {
128            label: Some(program.label()),
129            source: ShaderSource::Wgsl(program.shader.as_ref().into()),
130        });
131        let pipeline = device.create_compute_pipeline(&ComputePipelineDescriptor {
132            label: Some(program.label()),
133            layout: Some(&layout.pipeline),
134            module: &module,
135            entry_point: Some(&program.entry),
136            compilation_options: Default::default(),
137            cache: None,
138        });
139        Self { pipeline }
140    }
141
142    pub(crate) fn wgpu(&self) -> &WgpuComputePipeline {
143        &self.pipeline
144    }
145}
146
147struct PipelineSlot {
148    program: Arc<ComputeProgram>,
149    layout: ComputeLayout,
150    compiled: OnceLock<ComputePipeline>,
151}
152
153#[derive(Clone)]
154pub struct PipelineHandle {
155    slot: Arc<PipelineSlot>,
156}
157
158impl PipelineHandle {
159    pub(crate) fn new(program: Arc<ComputeProgram>, layout: ComputeLayout) -> Self {
160        Self {
161            slot: Arc::new(PipelineSlot {
162                program,
163                layout,
164                compiled: OnceLock::new(),
165            }),
166        }
167    }
168
169    pub fn label(&self) -> &str {
170        self.slot.program.label()
171    }
172
173    pub fn is_warmed(&self) -> bool {
174        self.slot.compiled.get().is_some()
175    }
176
177    pub fn pipeline(&self) -> &ComputePipeline {
178        self.slot.compiled.get().unwrap_or_else(|| {
179            panic!(
180                "pipeline {:?} must be warmed before it is recorded",
181                self.slot.program.label()
182            )
183        })
184    }
185
186    pub fn create_bind_group(
187        &self,
188        device: &Device,
189        group: usize,
190        entries: &[BindGroupEntry<'_>],
191    ) -> BindGroup {
192        self.slot.layout.create_bind_group(device, group, entries)
193    }
194
195    pub(crate) fn compile(&self, device: &Device) {
196        let pipeline = ComputePipeline::compile(device, &self.slot.program, &self.slot.layout);
197        assert!(
198            self.slot.compiled.set(pipeline).is_ok(),
199            "pipeline {:?} compiled twice",
200            self.slot.program.label()
201        );
202    }
203}
dynamis_gpu/pipeline.rs

dynamis_gpu/
pipeline.rs
