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.
164    pub own_group: u32,
165    /// Additional bind group layouts, each tagged with the `@group(N)` it occupies.
166    /// Every index from 0 up to the highest one used (including `own_group`) must be
167    /// covered exactly once, or `build_compute` panics — this makes group assignment
168    /// explicit instead of inferred from field order.
169    pub extra_layouts: Vec<super::layout::OwnedGroupLayout>,
170}
171
172impl<'a> Default for ComputeDescriptor<'a> {
173    fn default() -> Self {
174        Self {
175            label: None,
176            shader_source: "",
177            entry_point: Some("cs_main"),
178            entries: Vec::new(),
179            own_group: 0,
180            extra_layouts: Vec::new(),
181        }
182    }
183}
184
185pub fn build_bind_group_layout(
186    device: &wgpu::Device,
187    label: Option<&str>,
188    entries: &[ComputeBindingEntry],
189) -> wgpu::BindGroupLayout {
190    let layout_entries: Vec<_> = entries.iter().map(|e| e.kind.layout_entry(e.binding)).collect();
191
192    let mut seen = std::collections::HashSet::new();
193    for e in entries {
194        if !seen.insert(e.binding) {
195            panic!(
196                "binding {} assigned more than once building bind group layout{} (entry '{}')",
197                e.binding,
198                label.map(|l| format!(" '{l}'")).unwrap_or_default(),
199                e.name
200            );
201        }
202    }
203
204    device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
205        label,
206        entries: &layout_entries,
207    })
208}
209
210pub fn build_compute(
211    device: &wgpu::Device,
212    desc: &ComputeDescriptor,
213) -> (wgpu::ComputePipeline, wgpu::BindGroupLayout) {
214    let layout = build_bind_group_layout(device, desc.label, &desc.entries);
215
216    let module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
217        label: desc.label,
218        source: wgpu::ShaderSource::Wgsl(desc.shader_source.into()),
219    });
220
221    let mut slots: Vec<super::layout::GroupLayout> = desc
222        .extra_layouts
223        .iter()
224        .map(|g| super::layout::GroupLayout { group: g.group, layout: &g.layout })
225        .collect();
226    slots.push(super::layout::GroupLayout { group: desc.own_group, layout: &layout });
227    let bind_group_layouts = super::layout::assemble_bind_group_layouts(desc.label, slots);
228
229    let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
230        label: desc.label,
231        bind_group_layouts: &bind_group_layouts,
232        immediate_size: 0,
233    });
234
235    let pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
236        label: desc.label,
237        layout: Some(&pipeline_layout),
238        module: &module,
239        entry_point: desc.entry_point,
240        compilation_options: Default::default(),
241        cache: None,
242    });
243
244    (pipeline, layout)
245}
246
247pub struct GPUCompute {
248    pub pipeline: wgpu::ComputePipeline,
249    pub layout: wgpu::BindGroupLayout,
250    pub entries: Vec<ComputeBindingEntry>,
251}
252
253impl Asset<WGPUBackend> for GPUCompute {
254    type Source = ComputeDescriptor<'static>;
255    type Deps<'a> = ();
256
257    fn upload<'a>(source: &ComputeDescriptor, backend: &WGPUBackend, _deps: &()) -> Option<Self> {
258        let (pipeline, layout) = build_compute(&backend.device, source);
259
260        Some(Self {
261            pipeline,
262            layout,
263            entries: source.entries.to_vec(),
264        })
265    }
266}
267
268#[derive(Default)]
269pub struct ComputePlugin;
270impl ComputePlugin {
271    pub fn new() -> Self {
272        Self
273    }
274}
275
276impl Plugin for ComputePlugin {
277    fn build(&self, app: &mut App) {
278        app.add_plugin(AssetPlugin::<super::backend::WGPUBackend, GPUCompute>::new());
279    }
280}
pebble/wgpu/compute.rs

pebble/wgpu/
compute.rs
