pebble/wgpu/
compute.rs

1use crate::{
2    app::App,
3    assets::{plugin::AssetPlugin, upload::Asset},
4    ecs::plugin::Plugin,
5    wgpu::backend::WGPUBackend,
6};
7
8#[derive(Copy, Clone, PartialEq, Eq, Hash)]
9pub enum ComputeBindingKind {
10    StorageBufferReadOnly {
11        has_dynamic_offset: bool,
12        min_binding_size: Option<wgpu::BufferSize>,
13    },
14    StorageBufferReadWrite {
15        has_dynamic_offset: bool,
16        min_binding_size: Option<wgpu::BufferSize>,
17    },
18    UniformBuffer {
19        has_dynamic_offset: bool,
20        min_binding_size: Option<wgpu::BufferSize>,
21    },
22    Texture {
23        sample_type: wgpu::TextureSampleType,
24        view_dimension: wgpu::TextureViewDimension,
25        multisampled: bool,
26    },
27    StorageTexture {
28        format: wgpu::TextureFormat,
29        access: wgpu::StorageTextureAccess,
30        view_dimension: wgpu::TextureViewDimension,
31    },
32    Sampler,
33    ComparisonSampler,
34}
35
36impl ComputeBindingKind {
37    pub fn texture_2d() -> Self {
38        Self::Texture {
39            sample_type: wgpu::TextureSampleType::Float { filterable: true },
40            view_dimension: wgpu::TextureViewDimension::D2,
41            multisampled: false,
42        }
43    }
44
45    pub fn storage_buffer_read_only() -> Self {
46        Self::StorageBufferReadOnly { has_dynamic_offset: false, min_binding_size: None }
47    }
48
49    pub fn storage_buffer_read_write() -> Self {
50        Self::StorageBufferReadWrite { has_dynamic_offset: false, min_binding_size: None }
51    }
52
53    pub fn uniform_buffer() -> Self {
54        Self::UniformBuffer { has_dynamic_offset: false, min_binding_size: None }
55    }
56
57    /// A uniform buffer bound with a per-dispatch dynamic offset, e.g. one large buffer
58    /// holding many elements' data, rebound at a different offset via
59    /// `ComputePass::set_bind_group`'s dynamic offsets slice instead of a bind group per
60    /// dispatch. `element_size` is the size in bytes of a single element (before alignment
61    /// padding). Use [`crate::wgpu::buffers::build_dynamic_uniform_buffer`] to allocate the
62    /// backing buffer and [`crate::wgpu::buffers::dynamic_buffer_binding`] (not
63    /// `buffer.as_entire_binding()`) to build the bind group entry for it — the entry must
64    /// be scoped to one element's size, not the whole buffer, or dynamic offsets will fail
65    /// validation.
66    pub fn dynamic_uniform_buffer(element_size: u64) -> Self {
67        Self::UniformBuffer { has_dynamic_offset: true, min_binding_size: wgpu::BufferSize::new(element_size) }
68    }
69
70    /// A storage buffer bound with a per-dispatch dynamic offset. See [`Self::dynamic_uniform_buffer`].
71    pub fn dynamic_storage_buffer(element_size: u64, read_only: bool) -> Self {
72        let has_dynamic_offset = true;
73        let min_binding_size = wgpu::BufferSize::new(element_size);
74        if read_only {
75            Self::StorageBufferReadOnly { has_dynamic_offset, min_binding_size }
76        } else {
77            Self::StorageBufferReadWrite { has_dynamic_offset, min_binding_size }
78        }
79    }
80
81    pub fn layout_entry(&self, binding: u32) -> wgpu::BindGroupLayoutEntry {
82        match self {
83            ComputeBindingKind::StorageBufferReadOnly { has_dynamic_offset, min_binding_size } => wgpu::BindGroupLayoutEntry {
84                binding,
85                visibility: wgpu::ShaderStages::COMPUTE,
86                ty: wgpu::BindingType::Buffer {
87                    ty: wgpu::BufferBindingType::Storage { read_only: true },
88                    has_dynamic_offset: *has_dynamic_offset,
89                    min_binding_size: *min_binding_size,
90                },
91                count: None,
92            },
93            ComputeBindingKind::StorageBufferReadWrite { has_dynamic_offset, min_binding_size } => wgpu::BindGroupLayoutEntry {
94                binding,
95                visibility: wgpu::ShaderStages::COMPUTE,
96                ty: wgpu::BindingType::Buffer {
97                    ty: wgpu::BufferBindingType::Storage { read_only: false },
98                    has_dynamic_offset: *has_dynamic_offset,
99                    min_binding_size: *min_binding_size,
100                },
101                count: None,
102            },
103            ComputeBindingKind::UniformBuffer { has_dynamic_offset, min_binding_size } => wgpu::BindGroupLayoutEntry {
104                binding,
105                visibility: wgpu::ShaderStages::COMPUTE,
106                ty: wgpu::BindingType::Buffer {
107                    ty: wgpu::BufferBindingType::Uniform,
108                    has_dynamic_offset: *has_dynamic_offset,
109                    min_binding_size: *min_binding_size,
110                },
111                count: None,
112            },
113            ComputeBindingKind::Texture { sample_type, view_dimension, multisampled } => wgpu::BindGroupLayoutEntry {
114                binding,
115                visibility: wgpu::ShaderStages::COMPUTE,
116                ty: wgpu::BindingType::Texture {
117                    sample_type: *sample_type,
118                    view_dimension: *view_dimension,
119                    multisampled: *multisampled,
120                },
121                count: None,
122            },
123            ComputeBindingKind::StorageTexture { format, access, view_dimension } => wgpu::BindGroupLayoutEntry {
124                binding,
125                visibility: wgpu::ShaderStages::COMPUTE,
126                ty: wgpu::BindingType::StorageTexture {
127                    access: *access,
128                    format: *format,
129                    view_dimension: *view_dimension,
130                },
131                count: None,
132            },
133            ComputeBindingKind::Sampler => wgpu::BindGroupLayoutEntry {
134                binding,
135                visibility: wgpu::ShaderStages::COMPUTE,
136                ty: wgpu::BindingType::Sampler(wgpu::SamplerBindingType::Filtering),
137                count: None,
138            },
139            ComputeBindingKind::ComparisonSampler => wgpu::BindGroupLayoutEntry {
140                binding,
141                visibility: wgpu::ShaderStages::COMPUTE,
142                ty: wgpu::BindingType::Sampler(wgpu::SamplerBindingType::Comparison),
143                count: None,
144            },
145        }
146    }
147}
148
149#[derive(Clone)]
150pub struct ComputeBindingEntry {
151    pub name: &'static str,
152    /// The `@binding(N)` this entry occupies within its bind group. Explicit rather than
153    /// inferred from position in `entries`, so it matches the shader unambiguously.
154    pub binding: u32,
155    pub kind: ComputeBindingKind,
156}
157
158pub struct ComputeDescriptor<'a> {
159    pub label: Option<&'a str>,
160    pub shader_source: &'a str,
161    pub entry_point: Option<&'a str>,
162    pub entries: Vec<ComputeBindingEntry>,
163    /// Which `@group(N)` the layout built from `entries` occupies in the pipeline, or
164    /// `None` if this compute pass has no entries of its own (e.g. it only uses `extra_layouts`).
165    pub own_group: Option<u32>,
166    /// Additional bind group layouts, each tagged with the `@group(N)` it occupies.
167    /// Every index from 0 up to the highest one used (including `own_group`, if set) must
168    /// be covered exactly once, or `build_compute` panics — this makes group assignment
169    /// explicit instead of inferred from field order.
170    pub extra_layouts: Vec<super::layout::OwnedGroupLayout>,
171}
172
173impl<'a> Default for ComputeDescriptor<'a> {
174    fn default() -> Self {
175        Self {
176            label: None,
177            shader_source: "",
178            entry_point: Some("cs_main"),
179            entries: Vec::new(),
180            own_group: Some(0),
181            extra_layouts: Vec::new(),
182        }
183    }
184}
185
186pub fn build_bind_group_layout(
187    device: &wgpu::Device,
188    label: Option<&str>,
189    entries: &[ComputeBindingEntry],
190) -> wgpu::BindGroupLayout {
191    let layout_entries: Vec<_> = entries.iter().map(|e| e.kind.layout_entry(e.binding)).collect();
192
193    let mut seen = std::collections::HashSet::new();
194    for e in entries {
195        if !seen.insert(e.binding) {
196            panic!(
197                "binding {} assigned more than once building bind group layout{} (entry '{}')",
198                e.binding,
199                label.map(|l| format!(" '{l}'")).unwrap_or_default(),
200                e.name
201            );
202        }
203    }
204
205    device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
206        label,
207        entries: &layout_entries,
208    })
209}
210
211pub fn build_compute(
212    device: &wgpu::Device,
213    desc: &ComputeDescriptor,
214) -> (wgpu::ComputePipeline, wgpu::BindGroupLayout) {
215    let layout = build_bind_group_layout(device, desc.label, &desc.entries);
216
217    let module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
218        label: desc.label,
219        source: wgpu::ShaderSource::Wgsl(desc.shader_source.into()),
220    });
221
222    let mut slots: Vec<super::layout::GroupLayout> = desc
223        .extra_layouts
224        .iter()
225        .map(|g| super::layout::GroupLayout { group: g.group, layout: &g.layout })
226        .collect();
227    if let Some(own_group) = desc.own_group {
228        slots.push(super::layout::GroupLayout { group: own_group, layout: &layout });
229    }
230    let bind_group_layouts = super::layout::assemble_bind_group_layouts(desc.label, slots);
231
232    let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
233        label: desc.label,
234        bind_group_layouts: &bind_group_layouts,
235        immediate_size: 0,
236    });
237
238    let pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
239        label: desc.label,
240        layout: Some(&pipeline_layout),
241        module: &module,
242        entry_point: desc.entry_point,
243        compilation_options: Default::default(),
244        cache: None,
245    });
246
247    (pipeline, layout)
248}
249
250pub struct GPUCompute {
251    pub pipeline: wgpu::ComputePipeline,
252    pub layout: wgpu::BindGroupLayout,
253    pub entries: Vec<ComputeBindingEntry>,
254}
255
256impl Asset<WGPUBackend> for GPUCompute {
257    type Source = ComputeDescriptor<'static>;
258    type Deps<'a> = ();
259
260    fn upload<'a>(source: &ComputeDescriptor, backend: &WGPUBackend, _deps: &()) -> Option<Self> {
261        let (pipeline, layout) = build_compute(&backend.device, source);
262
263        Some(Self {
264            pipeline,
265            layout,
266            entries: source.entries.to_vec(),
267        })
268    }
269}
270
271#[derive(Default)]
272pub struct ComputePlugin;
273impl ComputePlugin {
274    pub fn new() -> Self {
275        Self
276    }
277}
278
279impl Plugin for ComputePlugin {
280    fn build(&self, app: &mut App) {
281        app.add_plugin(AssetPlugin::<super::backend::WGPUBackend, GPUCompute>::new());
282    }
283}
pebble/wgpu/compute.rs

pebble/wgpu/
compute.rs
