pebble/wgpu/
compute.rs

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

pebble/wgpu/
compute.rs
