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::dynamic_uniform_offset_stride`] or
62    /// [`crate::wgpu::buffers::build_dynamic_uniform_buffer`] to compute the actual stride.
63    pub fn dynamic_uniform_buffer(element_size: u64) -> Self {
64        Self::UniformBuffer { has_dynamic_offset: true, min_binding_size: wgpu::BufferSize::new(element_size) }
65    }
66
67    /// A storage buffer bound with a per-dispatch dynamic offset. See [`Self::dynamic_uniform_buffer`].
68    pub fn dynamic_storage_buffer(element_size: u64, read_only: bool) -> Self {
69        let has_dynamic_offset = true;
70        let min_binding_size = wgpu::BufferSize::new(element_size);
71        if read_only {
72            Self::StorageBufferReadOnly { has_dynamic_offset, min_binding_size }
73        } else {
74            Self::StorageBufferReadWrite { has_dynamic_offset, min_binding_size }
75        }
76    }
77
78    pub fn layout_entry(&self, binding: u32) -> wgpu::BindGroupLayoutEntry {
79        match self {
80            ComputeBindingKind::StorageBufferReadOnly { has_dynamic_offset, min_binding_size } => wgpu::BindGroupLayoutEntry {
81                binding,
82                visibility: wgpu::ShaderStages::COMPUTE,
83                ty: wgpu::BindingType::Buffer {
84                    ty: wgpu::BufferBindingType::Storage { read_only: true },
85                    has_dynamic_offset: *has_dynamic_offset,
86                    min_binding_size: *min_binding_size,
87                },
88                count: None,
89            },
90            ComputeBindingKind::StorageBufferReadWrite { has_dynamic_offset, min_binding_size } => wgpu::BindGroupLayoutEntry {
91                binding,
92                visibility: wgpu::ShaderStages::COMPUTE,
93                ty: wgpu::BindingType::Buffer {
94                    ty: wgpu::BufferBindingType::Storage { read_only: false },
95                    has_dynamic_offset: *has_dynamic_offset,
96                    min_binding_size: *min_binding_size,
97                },
98                count: None,
99            },
100            ComputeBindingKind::UniformBuffer { has_dynamic_offset, min_binding_size } => wgpu::BindGroupLayoutEntry {
101                binding,
102                visibility: wgpu::ShaderStages::COMPUTE,
103                ty: wgpu::BindingType::Buffer {
104                    ty: wgpu::BufferBindingType::Uniform,
105                    has_dynamic_offset: *has_dynamic_offset,
106                    min_binding_size: *min_binding_size,
107                },
108                count: None,
109            },
110            ComputeBindingKind::Texture { sample_type, view_dimension, multisampled } => wgpu::BindGroupLayoutEntry {
111                binding,
112                visibility: wgpu::ShaderStages::COMPUTE,
113                ty: wgpu::BindingType::Texture {
114                    sample_type: *sample_type,
115                    view_dimension: *view_dimension,
116                    multisampled: *multisampled,
117                },
118                count: None,
119            },
120            ComputeBindingKind::StorageTexture { format, access, view_dimension } => wgpu::BindGroupLayoutEntry {
121                binding,
122                visibility: wgpu::ShaderStages::COMPUTE,
123                ty: wgpu::BindingType::StorageTexture {
124                    access: *access,
125                    format: *format,
126                    view_dimension: *view_dimension,
127                },
128                count: None,
129            },
130            ComputeBindingKind::Sampler => wgpu::BindGroupLayoutEntry {
131                binding,
132                visibility: wgpu::ShaderStages::COMPUTE,
133                ty: wgpu::BindingType::Sampler(wgpu::SamplerBindingType::Filtering),
134                count: None,
135            },
136            ComputeBindingKind::ComparisonSampler => wgpu::BindGroupLayoutEntry {
137                binding,
138                visibility: wgpu::ShaderStages::COMPUTE,
139                ty: wgpu::BindingType::Sampler(wgpu::SamplerBindingType::Comparison),
140                count: None,
141            },
142        }
143    }
144}
145
146#[derive(Clone)]
147pub struct ComputeBindingEntry {
148    pub name: &'static str,
149    pub kind: ComputeBindingKind,
150}
151
152pub struct ComputeDescriptor<'a> {
153    pub label: Option<&'a str>,
154    pub shader_source: &'a str,
155    pub entry_point: Option<&'a str>,
156    pub entries: Vec<ComputeBindingEntry>,
157    pub extra_layouts: Vec<wgpu::BindGroupLayout>,
158}
159
160impl<'a> Default for ComputeDescriptor<'a> {
161    fn default() -> Self {
162        Self {
163            label: None,
164            shader_source: "",
165            entry_point: Some("cs_main"),
166            entries: Vec::new(),
167            extra_layouts: Vec::new(),
168        }
169    }
170}
171
172pub fn build_bind_group_layout(
173    device: &wgpu::Device,
174    label: Option<&str>,
175    entries: &[ComputeBindingEntry],
176) -> wgpu::BindGroupLayout {
177    let layout_entries: Vec<_> = entries
178        .iter()
179        .enumerate()
180        .map(|(i, e)| e.kind.layout_entry(i as u32))
181        .collect();
182
183    device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
184        label,
185        entries: &layout_entries,
186    })
187}
188
189pub fn build_compute(
190    device: &wgpu::Device,
191    desc: &ComputeDescriptor,
192) -> (wgpu::ComputePipeline, wgpu::BindGroupLayout) {
193    let layout = build_bind_group_layout(device, desc.label, &desc.entries);
194
195    let module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
196        label: desc.label,
197        source: wgpu::ShaderSource::Wgsl(desc.shader_source.into()),
198    });
199
200    let mut bind_group_layouts: Vec<&wgpu::BindGroupLayout> = desc.extra_layouts.iter().collect();
201    bind_group_layouts.push(&layout);
202    let bind_group_layouts: Vec<Option<&wgpu::BindGroupLayout>> =
203        bind_group_layouts.into_iter().map(Some).collect();
204
205    let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
206        label: desc.label,
207        bind_group_layouts: &bind_group_layouts,
208        immediate_size: 0,
209    });
210
211    let pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
212        label: desc.label,
213        layout: Some(&pipeline_layout),
214        module: &module,
215        entry_point: desc.entry_point,
216        compilation_options: Default::default(),
217        cache: None,
218    });
219
220    (pipeline, layout)
221}
222
223pub struct GPUCompute {
224    pub pipeline: wgpu::ComputePipeline,
225    pub layout: wgpu::BindGroupLayout,
226    pub entries: Vec<ComputeBindingEntry>,
227}
228
229impl Asset<WGPUBackend> for GPUCompute {
230    type Source = ComputeDescriptor<'static>;
231    type Deps<'a> = ();
232
233    fn upload<'a>(source: &ComputeDescriptor, backend: &WGPUBackend, _deps: &()) -> Option<Self> {
234        let (pipeline, layout) = build_compute(&backend.device, source);
235
236        Some(Self {
237            pipeline,
238            layout,
239            entries: source.entries.to_vec(),
240        })
241    }
242}
243
244#[derive(Default)]
245pub struct ComputePlugin;
246impl ComputePlugin {
247    pub fn new() -> Self {
248        Self
249    }
250}
251
252impl Plugin for ComputePlugin {
253    fn build(&self, app: &mut App) {
254        app.add_plugin(AssetPlugin::<super::backend::WGPUBackend, GPUCompute>::new());
255    }
256}
pebble/wgpu/compute.rs

pebble/wgpu/
compute.rs
