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    pub kind: ComputeBindingKind,
153}
154
155pub struct ComputeDescriptor<'a> {
156    pub label: Option<&'a str>,
157    pub shader_source: &'a str,
158    pub entry_point: Option<&'a str>,
159    pub entries: Vec<ComputeBindingEntry>,
160    pub extra_layouts: Vec<wgpu::BindGroupLayout>,
161}
162
163impl<'a> Default for ComputeDescriptor<'a> {
164    fn default() -> Self {
165        Self {
166            label: None,
167            shader_source: "",
168            entry_point: Some("cs_main"),
169            entries: Vec::new(),
170            extra_layouts: Vec::new(),
171        }
172    }
173}
174
175pub fn build_bind_group_layout(
176    device: &wgpu::Device,
177    label: Option<&str>,
178    entries: &[ComputeBindingEntry],
179) -> wgpu::BindGroupLayout {
180    let layout_entries: Vec<_> = entries
181        .iter()
182        .enumerate()
183        .map(|(i, e)| e.kind.layout_entry(i as u32))
184        .collect();
185
186    device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
187        label,
188        entries: &layout_entries,
189    })
190}
191
192pub fn build_compute(
193    device: &wgpu::Device,
194    desc: &ComputeDescriptor,
195) -> (wgpu::ComputePipeline, wgpu::BindGroupLayout) {
196    let layout = build_bind_group_layout(device, desc.label, &desc.entries);
197
198    let module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
199        label: desc.label,
200        source: wgpu::ShaderSource::Wgsl(desc.shader_source.into()),
201    });
202
203    let mut bind_group_layouts: Vec<&wgpu::BindGroupLayout> = desc.extra_layouts.iter().collect();
204    bind_group_layouts.push(&layout);
205    let bind_group_layouts: Vec<Option<&wgpu::BindGroupLayout>> =
206        bind_group_layouts.into_iter().map(Some).collect();
207
208    let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
209        label: desc.label,
210        bind_group_layouts: &bind_group_layouts,
211        immediate_size: 0,
212    });
213
214    let pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
215        label: desc.label,
216        layout: Some(&pipeline_layout),
217        module: &module,
218        entry_point: desc.entry_point,
219        compilation_options: Default::default(),
220        cache: None,
221    });
222
223    (pipeline, layout)
224}
225
226pub struct GPUCompute {
227    pub pipeline: wgpu::ComputePipeline,
228    pub layout: wgpu::BindGroupLayout,
229    pub entries: Vec<ComputeBindingEntry>,
230}
231
232impl Asset<WGPUBackend> for GPUCompute {
233    type Source = ComputeDescriptor<'static>;
234    type Deps<'a> = ();
235
236    fn upload<'a>(source: &ComputeDescriptor, backend: &WGPUBackend, _deps: &()) -> Option<Self> {
237        let (pipeline, layout) = build_compute(&backend.device, source);
238
239        Some(Self {
240            pipeline,
241            layout,
242            entries: source.entries.to_vec(),
243        })
244    }
245}
246
247#[derive(Default)]
248pub struct ComputePlugin;
249impl ComputePlugin {
250    pub fn new() -> Self {
251        Self
252    }
253}
254
255impl Plugin for ComputePlugin {
256    fn build(&self, app: &mut App) {
257        app.add_plugin(AssetPlugin::<super::backend::WGPUBackend, GPUCompute>::new());
258    }
259}
pebble/wgpu/compute.rs

pebble/wgpu/
compute.rs
