pebble/wgpu/
compute.rs

1use crate::{
2    assets::upload::Asset,
3    wgpu::{backend::WGPUBackend, binding::{BindGroupLayoutBuilder, 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::BindGroupLayoutBuilder`](super::binding::BindGroupLayoutBuilder).
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 = BindGroupLayoutBuilder::new()
75        .label(desc.label)
76        .entries(desc.entries.iter().cloned())
77        .build(device);
78
79    let module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
80        label: desc.label,
81        source: wgpu::ShaderSource::Wgsl(desc.shader_source.into()),
82    });
83
84    let mut slots: Vec<super::layout::GroupLayout> = desc
85        .extra_layouts
86        .iter()
87        .map(|g| super::layout::GroupLayout { group: g.group, layout: &g.layout })
88        .collect();
89    if let Some(own_group) = desc.own_group {
90        slots.push(super::layout::GroupLayout { group: own_group, layout: &layout });
91    }
92    let bind_group_layouts = super::layout::assemble_bind_group_layouts(desc.label, slots);
93
94    let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
95        label: desc.label,
96        bind_group_layouts: &bind_group_layouts,
97        immediate_size: 0,
98    });
99
100    let pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
101        label: desc.label,
102        layout: Some(&pipeline_layout),
103        module: &module,
104        entry_point: desc.entry_point,
105        compilation_options: Default::default(),
106        cache: None,
107    });
108
109    (pipeline, layout)
110}
111
112/// A compute pass uploaded to the GPU: a compute pipeline plus the bind
113/// group layout entries it expects.
114pub struct GPUCompute {
115    pub pipeline: wgpu::ComputePipeline,
116    pub layout: wgpu::BindGroupLayout,
117    pub entries: Vec<BindingEntry>,
118}
119
120impl super::binding::BindGroupTarget for GPUCompute {
121    fn bind_group_layout(&self) -> &wgpu::BindGroupLayout {
122        &self.layout
123    }
124    fn binding_entries(&self) -> &[BindingEntry] {
125        &self.entries
126    }
127}
128
129impl Asset<WGPUBackend> for GPUCompute {
130    type Source = ComputeDescriptor<'static>;
131    type Deps<'a> = ();
132
133    fn upload<'a>(source: &ComputeDescriptor, backend: &WGPUBackend, _deps: &()) -> Option<Self> {
134        let (pipeline, layout) = build_compute(&backend.device, source);
135
136        Some(Self {
137            pipeline,
138            layout,
139            entries: source.entries.to_vec(),
140        })
141    }
142}
143
144crate::wgpu::plugin_macros::asset_plugin! {
145    /// Registers the [`GPUCompute`] asset pipeline (`Assets<ComputeDescriptor>`
146    /// → `ProcessedAssets<GPUCompute>`). Included by
147    /// [`WGPUPlugin`](super::backend::WGPUPlugin); add directly only if you're
148    /// assembling the `wgpu` module's plugins by hand.
149    ComputePlugin, GPUCompute
150}
151
152#[cfg(test)]
153mod tests {
154    use super::*;
155    use crate::wgpu::binding::{BindingEntry, BindingKind};
156    use crate::wgpu::test_util::with_device;
157
158    const MINIMAL_COMPUTE_SHADER: &str = r#"
159        @compute @workgroup_size(1)
160        fn cs_main() {}
161    "#;
162
163    #[test]
164    fn a_fragment_visible_entry_panics_before_touching_the_device() {
165        with_device!(device, _queue, {
166            let desc = ComputeDescriptor {
167                shader_source: MINIMAL_COMPUTE_SHADER,
168                entries: vec![BindingEntry {
169                    name: "bad",
170                    binding: 0,
171                    kind: BindingKind::sampler(wgpu::ShaderStages::FRAGMENT),
172                }],
173                ..Default::default()
174            };
175            let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
176                build_compute(&device, &desc);
177            }));
178            assert!(result.is_err(), "expected a panic for a non-COMPUTE-visible compute entry");
179        });
180    }
181
182    #[test]
183    fn a_vertex_fragment_visible_entry_also_panics() {
184        // Not just "wrong stage" but "wrong stage in addition to COMPUTE" —
185        // build_compute requires visibility == exactly COMPUTE, so a
186        // COMPUTE | FRAGMENT entry (reused from a material by mistake, say)
187        // must panic too, not just entries missing COMPUTE entirely.
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::storage_buffer_read_write(
195                        wgpu::ShaderStages::COMPUTE | wgpu::ShaderStages::FRAGMENT,
196                    ),
197                }],
198                ..Default::default()
199            };
200            let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
201                build_compute(&device, &desc);
202            }));
203            assert!(result.is_err(), "expected a panic for a COMPUTE | FRAGMENT compute entry");
204        });
205    }
206
207    #[test]
208    fn a_compute_only_entry_builds_without_panicking() {
209        with_device!(device, _queue, {
210            let desc = ComputeDescriptor {
211                shader_source: MINIMAL_COMPUTE_SHADER,
212                entries: vec![],
213                own_group: None,
214                ..Default::default()
215            };
216            build_compute(&device, &desc);
217        });
218    }
219}
pebble/wgpu/compute.rs

pebble/wgpu/
compute.rs
