pebble/wgpu/
compute.rs

1use crate::{
2    assets::upload::Asset,
3    wgpu::{backend::WGPUBackend, binding::{BindGroupLayout, BindGroupLayoutBuilder, BindingEntry}},
4};
5
6/// A `wgpu::ComputePipeline`, opaque — built only via [`build_compute`]/
7/// [`GPUCompute`]'s `Asset::upload`. Bind it against a
8/// [`ComputePass`](super::compute_pass::ComputePass) via
9/// [`ComputePass::set_pipeline`](super::compute_pass::ComputePass::set_pipeline);
10/// there's no way to reach the underlying `wgpu::ComputePipeline` from
11/// outside this crate.
12pub struct ComputePipeline(wgpu::ComputePipeline);
13
14impl ComputePipeline {
15    pub(crate) fn raw(&self) -> &wgpu::ComputePipeline {
16        &self.0
17    }
18}
19
20/// Describes a compute pipeline + its own bind group, the source type
21/// [`GPUCompute`] is built from via [`build_compute`].
22pub struct ComputeDescriptor<'a> {
23    /// Debug label, threaded through to the shader module, pipeline, and
24    /// bind group layout.
25    pub label: Option<&'a str>,
26    /// WGSL source for the compute stage.
27    pub shader_source: &'a str,
28    /// Compute stage entry point. Defaults to `"cs_main"`.
29    pub entry_point: Option<&'a str>,
30    /// This compute pass's own bind group entries. See
31    /// [`BindingKind`](super::binding::BindingKind) for what a
32    /// compute-appropriate entry looks like — [`build_compute`] panics if
33    /// any entry here isn't exactly `COMPUTE`-visible.
34    pub entries: Vec<BindingEntry>,
35    /// Which `@group(N)` the layout built from `entries` occupies in the pipeline, or
36    /// `None` if this compute pass has no entries of its own (e.g. it only uses `extra_layouts`).
37    pub own_group: Option<u32>,
38    /// Additional bind group layouts, each tagged with the `@group(N)` it occupies.
39    /// Every index from 0 up to the highest one used (including `own_group`, if set) must
40    /// be covered exactly once, or `build_compute` panics — this makes group assignment
41    /// explicit instead of inferred from field order.
42    pub extra_layouts: Vec<super::layout::OwnedGroupLayout>,
43}
44
45impl<'a> Default for ComputeDescriptor<'a> {
46    fn default() -> Self {
47        Self {
48            label: None,
49            shader_source: "",
50            entry_point: Some("cs_main"),
51            entries: Vec::new(),
52            own_group: Some(0),
53            extra_layouts: Vec::new(),
54        }
55    }
56}
57
58/// Builds a compute pipeline and its own bind group layout from `desc`.
59///
60/// Panics if any of `desc.entries` isn't visible to exactly the compute
61/// stage — [`BindingKind`](super::binding::BindingKind) is shared with
62/// [`MaterialDescriptor`](super::material::MaterialDescriptor), and this is
63/// the check that catches a material entry (`FRAGMENT`/`VERTEX_FRAGMENT`)
64/// accidentally reused in a compute pass instead of letting it fail deep
65/// inside wgpu with a less specific error. The bind group layout itself
66/// comes from [`binding::BindGroupLayoutBuilder`](super::binding::BindGroupLayoutBuilder).
67/// The pipeline layout is assembled from `desc.own_group` (this pass's own
68/// layout) plus `desc.extra_layouts`, keyed by explicit `@group(N)` —
69/// panics on a gap or a collision across `0..=max`, turning a mismatched
70/// `@group(N)` in the shader into an immediate, specific error instead of
71/// an opaque wgpu validation failure at draw time.
72pub fn build_compute(
73    device: &wgpu::Device,
74    desc: &ComputeDescriptor,
75) -> (ComputePipeline, BindGroupLayout) {
76    for entry in &desc.entries {
77        if entry.kind.visibility() != wgpu::ShaderStages::COMPUTE {
78            panic!(
79                "compute pass{}: entry '{}' has visibility {:?} — compute bind group entries \
80                 must be visible to exactly the compute stage",
81                desc.label.map(|l| format!(" '{l}'")).unwrap_or_default(),
82                entry.name,
83                entry.kind.visibility()
84            );
85        }
86    }
87
88    let layout = BindGroupLayoutBuilder::new()
89        .label(desc.label)
90        .entries(desc.entries.iter().cloned())
91        .build(device);
92
93    let module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
94        label: desc.label,
95        source: wgpu::ShaderSource::Wgsl(desc.shader_source.into()),
96    });
97
98    let mut slots: Vec<super::layout::GroupLayout> = desc
99        .extra_layouts
100        .iter()
101        .map(|g| super::layout::GroupLayout { group: g.group, layout: &g.layout })
102        .collect();
103    if let Some(own_group) = desc.own_group {
104        slots.push(super::layout::GroupLayout { group: own_group, layout: &layout });
105    }
106    let bind_group_layouts = super::layout::assemble_bind_group_layouts(desc.label, slots);
107
108    let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
109        label: desc.label,
110        bind_group_layouts: &bind_group_layouts,
111        immediate_size: 0,
112    });
113
114    let pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
115        label: desc.label,
116        layout: Some(&pipeline_layout),
117        module: &module,
118        entry_point: desc.entry_point,
119        compilation_options: Default::default(),
120        cache: None,
121    });
122
123    (ComputePipeline(pipeline), layout)
124}
125
126/// A compute pass uploaded to the GPU: a compute pipeline plus the bind
127/// group layout entries it expects.
128pub struct GPUCompute {
129    pub pipeline: ComputePipeline,
130    layout: BindGroupLayout,
131    entries: Vec<BindingEntry>,
132}
133
134impl super::binding::BindGroupTarget for GPUCompute {
135    fn bind_group_layout(&self) -> &BindGroupLayout {
136        &self.layout
137    }
138    fn binding_entries(&self) -> &[BindingEntry] {
139        &self.entries
140    }
141}
142
143impl Asset<WGPUBackend> for GPUCompute {
144    type Source = ComputeDescriptor<'static>;
145    type Deps<'a> = ();
146
147    fn upload<'a>(source: &ComputeDescriptor, backend: &WGPUBackend, _deps: &()) -> Option<Self> {
148        let (pipeline, layout) = build_compute(&backend.device, source);
149
150        Some(Self {
151            pipeline,
152            layout,
153            entries: source.entries.to_vec(),
154        })
155    }
156}
157
158crate::wgpu::plugin_macros::asset_plugin! {
159    /// Registers the [`GPUCompute`] asset pipeline (`Assets<ComputeDescriptor>`
160    /// → `ProcessedAssets<GPUCompute>`). Included by
161    /// [`WGPUPlugin`](super::backend::WGPUPlugin); add directly only if you're
162    /// assembling the `wgpu` module's plugins by hand.
163    ComputePlugin, GPUCompute
164}
165
166#[cfg(test)]
167mod tests {
168    use super::*;
169    use crate::wgpu::binding::{BindingEntry, BindingKind};
170    use crate::wgpu::test_util::with_device;
171
172    const MINIMAL_COMPUTE_SHADER: &str = r#"
173        @compute @workgroup_size(1)
174        fn cs_main() {}
175    "#;
176
177    #[test]
178    fn a_fragment_visible_entry_panics_before_touching_the_device() {
179        with_device!(device, _queue, {
180            let desc = ComputeDescriptor {
181                shader_source: MINIMAL_COMPUTE_SHADER,
182                entries: vec![BindingEntry {
183                    name: "bad",
184                    binding: 0,
185                    kind: BindingKind::sampler(wgpu::ShaderStages::FRAGMENT),
186                }],
187                ..Default::default()
188            };
189            let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
190                build_compute(&device, &desc);
191            }));
192            assert!(result.is_err(), "expected a panic for a non-COMPUTE-visible compute entry");
193        });
194    }
195
196    #[test]
197    fn a_vertex_fragment_visible_entry_also_panics() {
198        // Not just "wrong stage" but "wrong stage in addition to COMPUTE" —
199        // build_compute requires visibility == exactly COMPUTE, so a
200        // COMPUTE | FRAGMENT entry (reused from a material by mistake, say)
201        // must panic too, not just entries missing COMPUTE entirely.
202        with_device!(device, _queue, {
203            let desc = ComputeDescriptor {
204                shader_source: MINIMAL_COMPUTE_SHADER,
205                entries: vec![BindingEntry {
206                    name: "bad",
207                    binding: 0,
208                    kind: BindingKind::storage_buffer_read_write(
209                        wgpu::ShaderStages::COMPUTE | wgpu::ShaderStages::FRAGMENT,
210                    ),
211                }],
212                ..Default::default()
213            };
214            let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
215                build_compute(&device, &desc);
216            }));
217            assert!(result.is_err(), "expected a panic for a COMPUTE | FRAGMENT compute entry");
218        });
219    }
220
221    #[test]
222    fn a_compute_only_entry_builds_without_panicking() {
223        with_device!(device, _queue, {
224            let desc = ComputeDescriptor {
225                shader_source: MINIMAL_COMPUTE_SHADER,
226                entries: vec![],
227                own_group: None,
228                ..Default::default()
229            };
230            build_compute(&device, &desc);
231        });
232    }
233}
pebble/wgpu/compute.rs

pebble/wgpu/
compute.rs
