pebble/wgpu/
compute.rs

1use crate::{
2    assets::{handle::Handle, storage::Assets, upload::Asset},
3    wgpu::{
4        backend::WGPUBackend,
5        binding::{BindGroupLayout, BindGroupLayoutBuilder, BindingEntry},
6        flags::ShaderStages,
7    },
8};
9
10/// A `wgpu::ComputePipeline`, opaque — built only via [`build_compute`]/
11/// [`GPUCompute`]'s `Asset::upload`. Bind it against a
12/// [`ComputePass`](super::compute_pass::ComputePass) via
13/// [`ComputePass::set_pipeline`](super::compute_pass::ComputePass::set_pipeline);
14/// there's no way to reach the underlying `wgpu::ComputePipeline` from
15/// outside this crate.
16pub struct ComputePipeline(wgpu::ComputePipeline);
17
18impl ComputePipeline {
19    pub(crate) fn raw(&self) -> &wgpu::ComputePipeline {
20        &self.0
21    }
22}
23
24/// Describes a compute pipeline + its own bind group, the source type
25/// [`GPUCompute`] is built from via [`build_compute`]. Fields are private —
26/// start from [`Compute::new`] and chain the setters below rather than
27/// constructing one as a struct literal.
28pub struct Compute {
29    /// Debug label, threaded through to the shader module, pipeline, and
30    /// bind group layout.
31    label: Option<&'static str>,
32    /// WGSL source for the compute stage.
33    shader_source: &'static str,
34    /// Compute stage entry point. Defaults to `"cs_main"`.
35    entry_point: Option<&'static str>,
36    /// This compute pass's own bind group entries. See
37    /// [`BindingKind`](super::binding::BindingKind) for what a
38    /// compute-appropriate entry looks like — [`build_compute`] panics if
39    /// any entry here isn't exactly `COMPUTE`-visible.
40    entries: Vec<BindingEntry>,
41    /// Which `@group(N)` the layout built from `entries` occupies in the pipeline, or
42    /// `None` if this compute pass has no entries of its own (e.g. it only uses `extra_layouts`).
43    own_group: Option<u32>,
44    /// Additional bind group layouts, each tagged with the `@group(N)` it occupies.
45    /// Every index from 0 up to the highest one used (including `own_group`, if set) must
46    /// be covered exactly once, or `build_compute` panics — this makes group assignment
47    /// explicit instead of inferred from field order.
48    extra_layouts: Vec<super::layout::OwnedGroupLayout>,
49}
50
51impl Default for Compute {
52    fn default() -> Self {
53        Self {
54            label: None,
55            shader_source: "",
56            entry_point: Some("cs_main"),
57            entries: Vec::new(),
58            own_group: Some(0),
59            extra_layouts: Vec::new(),
60        }
61    }
62}
63
64impl Compute {
65    /// Start building a compute pass with the given WGSL shader source.
66    /// All other fields are set to their defaults (see [`Default`]).
67    pub fn new(shader_source: &'static str) -> Self {
68        Self { shader_source, ..Self::default() }
69    }
70
71    pub fn label(mut self, label: &'static str) -> Self {
72        self.label = Some(label);
73        self
74    }
75
76    pub fn entry_point(mut self, entry: &'static str) -> Self {
77        self.entry_point = Some(entry);
78        self
79    }
80
81    pub fn entries(mut self, entries: Vec<BindingEntry>) -> Self {
82        self.entries = entries;
83        self
84    }
85
86    pub fn own_group(mut self, group: u32) -> Self {
87        self.own_group = Some(group);
88        self
89    }
90
91    /// Clear `own_group` — this compute pass has no bind group entries of
92    /// its own (only [`extra_layouts`](Self::extra_layouts)). The
93    /// counterpart to [`own_group`](Self::own_group), which can only set it
94    /// to `Some`.
95    pub fn no_own_group(mut self) -> Self {
96        self.own_group = None;
97        self
98    }
99
100    pub fn extra_layouts(mut self, layouts: Vec<super::layout::OwnedGroupLayout>) -> Self {
101        self.extra_layouts = layouts;
102        self
103    }
104
105    /// Logs a WARN if this pass has no bind group entries at all (neither
106    /// its own nor `extra_layouts`) — not fatal, since a shader could
107    /// legitimately need no bindings, but a compute pass with nothing to
108    /// read or write is unusual enough to flag.
109    fn validate(&self) {
110        if self.entries.is_empty() && self.extra_layouts.is_empty() {
111            tracing::warn!(
112                "Compute{}: no bind group entries at all — this pass can't read or write \
113                 anything; consider calling .entries(...) or .extra_layouts(...)",
114                self.label.map(|l| format!(" '{l}'")).unwrap_or_default(),
115            );
116        }
117    }
118
119    /// Consume the builder and return the finished [`Compute`] value.
120    pub fn build(self) -> Self {
121        self.validate();
122        self
123    }
124
125    /// Consume the builder, insert into `assets` under `name`, and return
126    /// the resulting [`Handle<Compute>`].
127    pub fn build_asset(self, name: &str, assets: &mut Assets<Self>) -> Handle<Self> {
128        self.validate();
129        assets.insert(name, self)
130    }
131}
132
133/// Builds a compute pipeline and its own bind group layout from `desc`.
134///
135/// Panics if any of `desc.entries` isn't visible to exactly the compute
136/// stage — [`BindingKind`](super::binding::BindingKind) is shared with
137/// [`Material`](super::material::Material), and this is
138/// the check that catches a material entry (`FRAGMENT`/`VERTEX_FRAGMENT`)
139/// accidentally reused in a compute pass instead of letting it fail deep
140/// inside wgpu with a less specific error. The bind group layout itself
141/// comes from [`binding::BindGroupLayoutBuilder`](super::binding::BindGroupLayoutBuilder).
142/// The pipeline layout is assembled from `desc.own_group` (this pass's own
143/// layout) plus `desc.extra_layouts`, keyed by explicit `@group(N)` —
144/// panics on a gap or a collision across `0..=max`, turning a mismatched
145/// `@group(N)` in the shader into an immediate, specific error instead of
146/// an opaque wgpu validation failure at draw time.
147pub fn build_compute(backend: &WGPUBackend, desc: &Compute) -> (ComputePipeline, BindGroupLayout) {
148    build_compute_raw(&backend.device, desc)
149}
150
151/// Internal primitive behind [`build_compute`] — used directly only by
152/// tests, which have a raw `wgpu::Device` but no full [`WGPUBackend`].
153pub(crate) fn build_compute_raw(
154    device: &wgpu::Device,
155    desc: &Compute,
156) -> (ComputePipeline, BindGroupLayout) {
157    for entry in &desc.entries {
158        if entry.kind.visibility() != ShaderStages::COMPUTE {
159            panic!(
160                "compute pass{}: entry '{}' is not visible to exactly the compute stage — \
161                 compute bind group entries must be visible to exactly COMPUTE",
162                desc.label.map(|l| format!(" '{l}'")).unwrap_or_default(),
163                entry.name,
164            );
165        }
166    }
167
168    let layout = BindGroupLayoutBuilder::new()
169        .label(desc.label)
170        .entries(desc.entries.iter().cloned())
171        .build_raw(device);
172
173    let module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
174        label: desc.label,
175        source: wgpu::ShaderSource::Wgsl(desc.shader_source.into()),
176    });
177
178    let mut slots: Vec<super::layout::GroupLayout> = desc
179        .extra_layouts
180        .iter()
181        .map(|g| super::layout::GroupLayout { group: g.group, layout: &g.layout })
182        .collect();
183    if let Some(own_group) = desc.own_group {
184        slots.push(super::layout::GroupLayout { group: own_group, layout: &layout });
185    }
186    let bind_group_layouts = super::layout::assemble_bind_group_layouts(desc.label, slots);
187
188    let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
189        label: desc.label,
190        bind_group_layouts: &bind_group_layouts,
191        immediate_size: 0,
192    });
193
194    let pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
195        label: desc.label,
196        layout: Some(&pipeline_layout),
197        module: &module,
198        entry_point: desc.entry_point,
199        compilation_options: Default::default(),
200        cache: None,
201    });
202
203    (ComputePipeline(pipeline), layout)
204}
205
206/// A compute pass uploaded to the GPU: a compute pipeline plus the bind
207/// group layout entries it expects.
208pub struct GPUCompute {
209    pub pipeline: ComputePipeline,
210    layout: BindGroupLayout,
211    entries: Vec<BindingEntry>,
212}
213
214impl super::binding::BindGroupTarget for GPUCompute {
215    fn bind_group_layout(&self) -> &BindGroupLayout {
216        &self.layout
217    }
218    fn binding_entries(&self) -> &[BindingEntry] {
219        &self.entries
220    }
221}
222
223impl Asset<WGPUBackend> for GPUCompute {
224    type Source = Compute;
225    type Deps<'a> = ();
226
227    fn upload<'a>(source: &Compute, backend: &WGPUBackend, _deps: &()) -> Option<Self> {
228        let (pipeline, layout) = build_compute(backend, source);
229
230        Some(Self {
231            pipeline,
232            layout,
233            entries: source.entries.to_vec(),
234        })
235    }
236}
237
238crate::wgpu::plugin_macros::asset_plugin! {
239    /// Registers the [`GPUCompute`] asset pipeline (`Assets<Compute>`
240    /// → `ProcessedAssets<GPUCompute>`). Included by
241    /// [`WGPUPlugin`](super::backend::WGPUPlugin); add directly only if you're
242    /// assembling the `wgpu` module's plugins by hand.
243    ComputePlugin, GPUCompute
244}
245
246#[cfg(test)]
247mod tests {
248    use super::*;
249    use crate::wgpu::binding::{BindingEntry, BindingKind};
250    use crate::wgpu::test_util::with_device;
251
252    const MINIMAL_COMPUTE_SHADER: &str = r#"
253        @compute @workgroup_size(1)
254        fn cs_main() {}
255    "#;
256
257    #[test]
258    fn a_fragment_visible_entry_panics_before_touching_the_device() {
259        with_device!(device, _queue, {
260            let desc = Compute {
261                shader_source: MINIMAL_COMPUTE_SHADER,
262                entries: vec![BindingEntry {
263                    name: "bad",
264                    binding: 0,
265                    kind: BindingKind::sampler(ShaderStages::FRAGMENT),
266                }],
267                ..Default::default()
268            };
269            let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
270                build_compute_raw(&device, &desc);
271            }));
272            assert!(result.is_err(), "expected a panic for a non-COMPUTE-visible compute entry");
273        });
274    }
275
276    #[test]
277    fn a_vertex_fragment_visible_entry_also_panics() {
278        // Not just "wrong stage" but "wrong stage in addition to COMPUTE" —
279        // build_compute requires visibility == exactly COMPUTE, so a
280        // COMPUTE | FRAGMENT entry (reused from a material by mistake, say)
281        // must panic too, not just entries missing COMPUTE entirely.
282        with_device!(device, _queue, {
283            let desc = Compute {
284                shader_source: MINIMAL_COMPUTE_SHADER,
285                entries: vec![BindingEntry {
286                    name: "bad",
287                    binding: 0,
288                    kind: BindingKind::storage_buffer_read_write(
289                        ShaderStages::COMPUTE | ShaderStages::FRAGMENT,
290                    ),
291                }],
292                ..Default::default()
293            };
294            let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
295                build_compute_raw(&device, &desc);
296            }));
297            assert!(result.is_err(), "expected a panic for a COMPUTE | FRAGMENT compute entry");
298        });
299    }
300
301    #[test]
302    fn a_compute_only_entry_builds_without_panicking() {
303        with_device!(device, _queue, {
304            let desc = Compute {
305                shader_source: MINIMAL_COMPUTE_SHADER,
306                entries: vec![],
307                own_group: None,
308                ..Default::default()
309            };
310            build_compute_raw(&device, &desc);
311        });
312    }
313}
pebble/wgpu/compute.rs

pebble/wgpu/
compute.rs
