pebble/wgpu/
compute.rs

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

pebble/wgpu/
compute.rs
