pebble/graphics/pipeline/
compute.rs

1use crate::{
2    assets::{handle::Handle, storage::Assets, upload::{Asset, AssetSource}},
3    ecs::resources::Read,
4    graphics::{
5        pipeline::{
6            binding::{BindGroupLayout, BindGroupLayoutBuilder, BindingKind},
7            buffers::{BindGroup, Buffer, DynamicBuffer},
8            cubemap::Cubemap,
9            layout::{
10                ComputePipelineCache, ComputePipelineKey, GlobalLayoutPool, GroupEntry, OwnEntriesBuilder, PipelineKind,
11                assemble_group_layouts, find_own_entries,
12            },
13            params::{BindGroupParams, BindingValue, build_bind_group},
14            samplers::{GlobalSamplers, SamplerKind},
15            texture_array::TextureArray,
16            texture_view::TextureView,
17            textures::Texture,
18        },
19        render::Backend,
20        types::flags::ShaderStages,
21    },
22};
23
24pub use pebble_derive::ComputeParams;
25
26/// A compiled GPU compute pipeline, wrapping `wgpu::ComputePipeline`. Cheap
27/// to `Clone` — `wgpu::ComputePipeline` is itself an `Arc`-backed handle —
28/// which is what lets [`ComputePipelineCache`] hand out a cache hit without
29/// recompiling.
30#[derive(Clone)]
31pub struct ComputePipeline(wgpu::ComputePipeline);
32
33impl ComputePipeline {
34    pub(crate) fn raw(&self) -> &wgpu::ComputePipeline {
35        &self.0
36    }
37}
38
39/// A compute pipeline asset, plus the bind group values (buffers/textures)
40/// it dispatches with — WGSL shader source and its bind group layout
41/// compile into a `wgpu::ComputePipeline`; many `Compute`s sharing the same
42/// shader automatically share one compiled pipeline (see
43/// [`ComputePipelineCache`]). Dispatch it via
44/// [`Backend::dispatch_compute`](crate::graphics::render::Backend::dispatch_compute).
45pub struct Compute {
46    label: Option<&'static str>,
47    shader_source: &'static str,
48    entry_point: Option<&'static str>,
49    own_entries: OwnEntriesBuilder,
50    extra_groups: Vec<GroupEntry>,
51    params: BindGroupParams,
52}
53
54impl Default for Compute {
55    fn default() -> Self {
56        Self {
57            label: None,
58            shader_source: "",
59            entry_point: Some("cs_main"),
60            own_entries: OwnEntriesBuilder::new(),
61            extra_groups: Vec::new(),
62            params: BindGroupParams::new(),
63        }
64    }
65}
66
67impl Compute {
68    pub fn new(shader_source: &'static str) -> Self {
69        Self { shader_source, ..Self::default() }
70    }
71
72    pub fn with_label(mut self, label: &'static str) -> Self {
73        self.label = Some(label);
74        self
75    }
76
77    pub fn with_entry_point(mut self, entry: &'static str) -> Self {
78        self.entry_point = Some(entry);
79        self
80    }
81
82    /// Declares one of this pass's own (group 0) bind group entries, at the
83    /// next auto-assigned binding index — the low-level counterpart to the
84    /// streamlined `.texture(...)`/`.buffer(...)`/etc. calls, for a `kind`
85    /// one of those doesn't produce (a dynamic-offset buffer, a non-default
86    /// sample type). Pair it with the matching value-only `.with_texture(...)`/etc.
87    pub fn with_entry(mut self, name: &'static str, kind: BindingKind) -> Self {
88        self.own_entries = self.own_entries.with_entry(name, kind);
89        self
90    }
91
92    /// Same as [`with_entry`](Self::with_entry), pinning an explicit
93    /// binding index instead of auto-assigning the next one.
94    pub fn with_entry_at(mut self, name: &'static str, binding: u32, kind: BindingKind) -> Self {
95        self.own_entries = self.own_entries.with_entry_at(name, binding, kind);
96        self
97    }
98
99    /// Appends a bind group beyond this pass's own (group 0) — typically
100    /// [`GroupEntry::Global`], a layout shared with other materials/computes
101    /// via [`GlobalLayoutPool`]. Groups append in call order, starting at
102    /// group 1.
103    pub fn with_extra_group(mut self, group: GroupEntry) -> Self {
104        self.extra_groups.push(group);
105        self
106    }
107
108    /// Value-only counterpart to `.texture(...)` — see [`texture`](Self::texture).
109    pub fn with_texture(mut self, name: &'static str, handle: Handle<Texture>) -> Self {
110        self.params = self.params.with_texture(name, handle);
111        self
112    }
113
114    /// Value-only counterpart to `.texture_array(...)` — see [`texture`](Self::texture).
115    pub fn with_texture_array(mut self, name: &'static str, handle: Handle<TextureArray>) -> Self {
116        self.params = self.params.with_texture_array(name, handle);
117        self
118    }
119
120    /// Value-only counterpart to `.cubemap(...)` — see [`texture`](Self::texture).
121    pub fn with_cubemap(mut self, name: &'static str, handle: Handle<Cubemap>) -> Self {
122        self.params = self.params.with_cubemap(name, handle);
123        self
124    }
125
126    /// Binds an already-built [`TextureView`] directly — e.g. one mip level
127    /// from [`GPUTexture::get_view`](super::textures::GPUTexture::get_view),
128    /// or a standalone render target from
129    /// [`Texture::empty`](super::textures::Texture::empty). Unlike
130    /// `.with_texture`/`.with_texture_array`/`.with_cubemap`, no `Handle`
131    /// lookup happens at upload time — `view` must already exist.
132    pub fn with_texture_view(mut self, name: &'static str, view: TextureView) -> Self {
133        self.params = self.params.with_texture_view(name, view);
134        self
135    }
136
137    /// Value-only counterpart to `.sampler(...)` — see [`texture`](Self::texture).
138    pub fn with_sampler(mut self, name: &'static str, kind: SamplerKind) -> Self {
139        self.params = self.params.with_sampler(name, kind);
140        self
141    }
142
143    /// Value-only counterpart to `.uniform(...)` — see [`texture`](Self::texture).
144    pub fn with_uniform(mut self, name: &'static str, data: Vec<u8>) -> Self {
145        self.params = self.params.with_uniform(name, data);
146        self
147    }
148
149    /// Value-only counterpart to `.storage(...)` — see [`texture`](Self::texture).
150    pub fn with_storage(mut self, name: &'static str, data: Vec<u8>) -> Self {
151        self.params = self.params.with_storage(name, data);
152        self
153    }
154
155    /// Same as [`with_uniform`](Self::with_uniform), but takes a typed
156    /// value instead of pre-packed bytes — uses `encase` to lay it out with
157    /// correct WGSL `uniform` (std140) alignment. Value-only counterpart to
158    /// `.uniform_value(...)`.
159    pub fn with_uniform_value<T>(mut self, name: &'static str, value: &T) -> Self
160    where
161        T: encase::ShaderType + encase::internal::WriteInto,
162    {
163        self.params = self.params.with_uniform_value(name, value);
164        self
165    }
166
167    /// Same as [`with_storage`](Self::with_storage), but takes a typed
168    /// value instead of pre-packed bytes — uses `encase` to lay it out with
169    /// correct WGSL `storage` (std430) alignment. Value-only counterpart to
170    /// `.storage_value(...)`.
171    pub fn with_storage_value<T>(mut self, name: &'static str, value: &T) -> Self
172    where
173        T: encase::ShaderType + encase::internal::WriteInto,
174    {
175        self.params = self.params.with_storage_value(name, value);
176        self
177    }
178
179    /// Declares a compute-visible `texture_2d<f32>` entry at the next
180    /// auto-assigned binding index *and* binds `handle` to it — the
181    /// streamlined one-call form of `.with_entry(name, BindingKind::texture_2d(COMPUTE))`
182    /// followed by `.with_texture(name, handle)`. Reach for those two
183    /// directly for a non-default sample type or an explicit binding index
184    /// — visibility is always `COMPUTE` for a compute pass, so there's no
185    /// visibility to override here.
186    pub fn texture(mut self, name: &'static str, handle: Handle<Texture>) -> Self {
187        self.own_entries = self.own_entries.with_entry(name, BindingKind::texture_2d(ShaderStages::COMPUTE));
188        self.with_texture(name, handle)
189    }
190
191    /// Streamlined form of `.texture_array(...)` — see [`texture`](Self::texture).
192    pub fn texture_array(mut self, name: &'static str, handle: Handle<TextureArray>) -> Self {
193        self.own_entries = self.own_entries.with_entry(name, BindingKind::texture_2d_array(ShaderStages::COMPUTE));
194        self.with_texture_array(name, handle)
195    }
196
197    /// Streamlined form of `.cubemap(...)` — see [`texture`](Self::texture).
198    pub fn cubemap(mut self, name: &'static str, handle: Handle<Cubemap>) -> Self {
199        self.own_entries = self.own_entries.with_entry(name, BindingKind::texture_cubemap(ShaderStages::COMPUTE));
200        self.with_cubemap(name, handle)
201    }
202
203    /// Streamlined form of `.sampler(...)` — see [`texture`](Self::texture).
204    pub fn sampler(mut self, name: &'static str, kind: SamplerKind) -> Self {
205        self.own_entries = self.own_entries.with_entry(name, BindingKind::sampler(ShaderStages::COMPUTE));
206        self.with_sampler(name, kind)
207    }
208
209    /// Streamlined form of `.uniform(...)` — see [`texture`](Self::texture).
210    pub fn uniform(mut self, name: &'static str, data: Vec<u8>) -> Self {
211        self.own_entries = self.own_entries.with_entry(name, BindingKind::uniform_buffer(ShaderStages::COMPUTE));
212        self.with_uniform(name, data)
213    }
214
215    /// Streamlined form of `.storage(...)` (read-write — a compute pass
216    /// binding a storage buffer usually means to write it) — see
217    /// [`texture`](Self::texture). Use `.with_entry(...)` +
218    /// `.with_storage(...)` directly for a read-only one.
219    pub fn storage(mut self, name: &'static str, data: Vec<u8>) -> Self {
220        self.own_entries = self.own_entries.with_entry(name, BindingKind::storage_buffer_read_write(ShaderStages::COMPUTE));
221        self.with_storage(name, data)
222    }
223
224    /// Streamlined, typed form of `.uniform(...)` — declares the entry and
225    /// binds an `encase`-laid-out value in one call. See [`texture`](Self::texture).
226    pub fn uniform_value<T>(mut self, name: &'static str, value: &T) -> Self
227    where
228        T: encase::ShaderType + encase::internal::WriteInto,
229    {
230        self.own_entries = self.own_entries.with_entry(name, BindingKind::uniform_buffer(ShaderStages::COMPUTE));
231        self.with_uniform_value(name, value)
232    }
233
234    /// Streamlined, typed form of `.storage(...)` (read-write) — see
235    /// [`storage`](Self::storage).
236    pub fn storage_value<T>(mut self, name: &'static str, value: &T) -> Self
237    where
238        T: encase::ShaderType + encase::internal::WriteInto,
239    {
240        self.own_entries = self.own_entries.with_entry(name, BindingKind::storage_buffer_read_write(ShaderStages::COMPUTE));
241        self.with_storage_value(name, value)
242    }
243
244    /// Binds an existing [`Buffer`] instead of uploading raw bytes — for a
245    /// buffer another pass already wrote to. Unlike `.with_uniform`/
246    /// `.with_storage`, no buffer is created here; `buffer` must already
247    /// carry the usage flags this binding needs.
248    pub fn with_buffer(mut self, name: &'static str, buffer: Buffer) -> Self {
249        self.params = self.params.with_buffer(name, buffer);
250        self
251    }
252
253    /// Binds an existing [`DynamicBuffer`] — the dynamic-offset counterpart
254    /// to `.with_buffer`.
255    pub fn with_dynamic_buffer(mut self, name: &'static str, buffer: DynamicBuffer) -> Self {
256        self.params = self.params.with_dynamic_buffer(name, buffer);
257        self
258    }
259
260    pub fn with_param(mut self, name: &'static str, entry: BindingValue) -> Self {
261        self.params = self.params.with_param(name, entry);
262        self
263    }
264
265    /// This pass's full bind group list — its own entries (group 0, from
266    /// `.texture(...)`/`.with_entry(...)`/etc.) followed by whatever
267    /// `.with_extra_group(...)` appended (group 1 and up).
268    fn groups(&self) -> Vec<GroupEntry> {
269        std::iter::once(GroupEntry::Own(self.own_entries.entries().to_vec()))
270            .chain(self.extra_groups.iter().cloned())
271            .collect()
272    }
273
274    fn validate(&self) {
275        if self.own_entries.entries().is_empty() && self.extra_groups.is_empty() {
276            tracing::warn!(
277                "Compute{}: no bind groups at all — this pass can't read or write \
278                 anything; consider calling .texture(...)/.buffer(...)/etc.",
279                self.label.map(|l| format!(" '{l}'")).unwrap_or_default(),
280            );
281        }
282        if self.params.is_empty() {
283            tracing::warn!(
284                "Compute{}: no bind group params — this pass won't bind anything against \
285                 its own entries; did you forget to chain .with_texture(...)/.with_buffer(...)/etc.?",
286                self.label.map(|l| format!(" '{l}'")).unwrap_or_default(),
287            );
288        }
289    }
290
291    pub fn build_asset(self, name: &str, assets: &mut Assets<Compute>) -> Handle<Compute> {
292        self.validate();
293        assets.insert(name, self)
294    }
295}
296
297/// Compiles a [`Compute`] into a raw pipeline + bind group layout. Used
298/// internally by the asset upload path (behind [`ComputePipelineCache`] —
299/// this always compiles, never checks the cache); exposed for callers
300/// assembling pipelines outside the usual [`Assets`] flow.
301pub fn build_compute(backend: &Backend, desc: &Compute, pool: &GlobalLayoutPool) -> Option<(ComputePipeline, BindGroupLayout)> {
302    let groups = desc.groups();
303    let own_entries = find_own_entries(desc.label, PipelineKind::Compute, &groups);
304    for entry in own_entries {
305        if entry.kind.visibility() != ShaderStages::COMPUTE {
306            panic!(
307                "compute pass{}: entry '{}' is not visible to exactly the compute stage — \
308                 compute bind group entries must be visible to exactly COMPUTE",
309                desc.label.map(|l| format!(" '{l}'")).unwrap_or_default(),
310                entry.name,
311            );
312        }
313    }
314
315    let layout = BindGroupLayoutBuilder::new()
316        .with_label(desc.label)
317        .with_entries(own_entries.iter().cloned())
318        .build(backend);
319
320    let device = &backend.device;
321    let module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
322        label: desc.label,
323        source: wgpu::ShaderSource::Wgsl(desc.shader_source.into()),
324    });
325
326    let bind_group_layouts = assemble_group_layouts(
327        desc.label,
328        &groups,
329        &layout,
330        pool,
331        device.limits().max_bind_groups,
332    )?;
333
334    let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
335        label: desc.label,
336        bind_group_layouts: &bind_group_layouts,
337        immediate_size: 0,
338    });
339
340    let pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
341        label: desc.label,
342        layout: Some(&pipeline_layout),
343        module: &module,
344        entry_point: desc.entry_point,
345        compilation_options: Default::default(),
346        cache: None,
347    });
348
349    Some((ComputePipeline(pipeline), layout))
350}
351
352/// The GPU-resident form an uploaded [`Compute`] produces — its compiled
353/// pipeline (possibly shared with other `Compute`s, see
354/// [`ComputePipelineCache`]) plus its own bind group.
355pub struct GPUCompute {
356    pub pipeline: ComputePipeline,
357    pub bind_group: BindGroup,
358    buffers: Vec<(&'static str, Buffer)>,
359    dynamic_buffers: Vec<(&'static str, DynamicBuffer)>,
360}
361
362impl GPUCompute {
363    /// Overwrites a named uniform/storage buffer's contents in place —
364    /// avoids rebuilding the whole bind group for a per-frame update.
365    pub fn update(&self, name: &str, data: &[u8]) {
366        match self.buffer(name) {
367            Some(buf) => buf.write(data),
368            None => tracing::warn!(
369                "GPUCompute::update: no bound buffer named '{name}' — check for a typo \
370                 against this pass's own .with_uniform(...)/.with_storage(...) entries"
371            ),
372        }
373    }
374
375    /// Same as [`update`](Self::update), but takes a typed value instead of
376    /// raw bytes — same `encase` layout `Compute::with_uniform_value`/
377    /// `with_storage_value` use.
378    pub fn update_value<T>(&self, name: &str, value: &T)
379    where
380        T: encase::ShaderType + encase::internal::WriteInto,
381    {
382        let mut buffer = encase::UniformBuffer::new(Vec::new());
383        buffer
384            .write(value)
385            .expect("encase: failed to write value — this shouldn't happen for a #[derive(ShaderType)] struct");
386        self.update(name, &buffer.into_inner());
387    }
388
389    pub fn buffer(&self, name: &str) -> Option<&Buffer> {
390        self.buffers.iter().find(|(n, _)| *n == name).map(|(_, buf)| buf)
391    }
392
393    /// Same as [`buffer`](Self::buffer), for a binding made via
394    /// `.with_dynamic_buffer` — use `DynamicBuffer::write_element` on the
395    /// result to update one element in place.
396    pub fn dynamic_buffer(&self, name: &str) -> Option<&DynamicBuffer> {
397        self.dynamic_buffers.iter().find(|(n, _)| *n == name).map(|(_, buf)| buf)
398    }
399}
400
401impl AssetSource for Compute {
402    type Processed = GPUCompute;
403}
404
405impl Asset<Backend> for Compute {
406    type Deps<'a> = (
407        Read<'a, GlobalLayoutPool>,
408        Read<'a, ComputePipelineCache>,
409        Read<'a, Assets<Texture>>,
410        Read<'a, Assets<TextureArray>>,
411        Read<'a, Assets<Cubemap>>,
412        Read<'a, GlobalSamplers>,
413    );
414
415    fn upload<'a>(&self, backend: &Backend, deps: &Self::Deps<'a>) -> Option<GPUCompute> {
416        let (layout_pool, pipeline_cache, textures, texture_arrays, cubemaps, samplers) = deps;
417
418        let groups = self.groups();
419        let key = ComputePipelineKey::new(self.shader_source, self.entry_point, &groups);
420        let (pipeline, layout) = pipeline_cache.get_or_compile(key, || build_compute(backend, self, layout_pool))?;
421        let entries = find_own_entries(self.label, PipelineKind::Compute, &groups);
422
423        let built = build_bind_group(backend, &self.params, &layout, entries, textures, texture_arrays, cubemaps, samplers)?;
424
425        Some(GPUCompute {
426            pipeline,
427            bind_group: built.bind_group,
428            buffers: built.buffers,
429            dynamic_buffers: built.dynamic_buffers,
430        })
431    }
432}
433
434#[cfg(test)]
435mod tests {
436    use super::*;
437
438    #[derive(ComputeParams)]
439    struct TestParams {
440        #[storage(0)]
441        data: f32,
442        #[texture(1)]
443        tex: Handle<Texture>,
444    }
445
446    #[test]
447    fn compute_params_derive_uses_compute_visibility_and_read_write_storage() {
448        let params = TestParams { data: 1.0, tex: Handle::default() };
449        let compute = params.into_compute(Compute::new("shader"));
450
451        assert!(!compute.params.is_empty());
452
453        let groups = compute.groups();
454        assert_eq!(groups.len(), 1);
455        let GroupEntry::Own(entries) = &groups[0] else { panic!("group 0 should be Own") };
456        assert_eq!(entries.len(), 2);
457        assert_eq!(entries[0].binding, 0);
458        assert!(entries[0].kind.visibility() == ShaderStages::COMPUTE);
459        assert_eq!(entries[1].binding, 1);
460        assert!(entries[1].kind.visibility() == ShaderStages::COMPUTE);
461    }
462}
pebble/graphics/pipeline/compute.rs

pebble/graphics/pipeline/
compute.rs
