bevy_shader_plugin/compute/
mod.rs

1mod node;
2mod pipeline;
3mod reflect;
4
5use std::{
6    any::{TypeId, type_name},
7    collections::HashSet,
8};
9
10use bevy::{
11    core_pipeline::schedule::camera_driver,
12    prelude::*,
13    render::{
14        Extract, ExtractSchedule, Render, RenderApp, RenderStartup, RenderSystems,
15        renderer::{RenderGraph, RenderGraphSystems},
16    },
17    shader::Shader,
18    window::PrimaryWindow,
19};
20
21use crate::{ShaderPassSystems, ShaderSystems, add_core_plugin, gpu_buffer::GpuBufferLengths};
22
23use reflect::{ComputeReflection, WorkgroupSize};
24
25/// WebGPU's guaranteed maximum workgroup count per dispatch dimension.
26const MAX_WORKGROUPS_PER_DIMENSION: u32 = 65535;
27
28/// How many workgroups a compute pass dispatches.
29///
30/// Every variant except [`exact`](Self::exact) is expressed as a number of *threads*;
31/// the plugin divides it by the entry point's `@workgroup_size`, read from the WGSL
32/// source, so the size is never repeated on the Rust side.
33#[derive(Clone, Debug)]
34pub struct Workgroups(WorkgroupsKind);
35
36#[derive(Clone, Debug)]
37enum WorkgroupsKind {
38    Window,
39    Buffer { tag: TypeId, tag_name: &'static str },
40    Threads(UVec3),
41    Exact(UVec3),
42}
43
44impl Workgroups {
45    /// One thread per physical pixel of the primary window, on x and y.
46    /// Use `global_invocation_id.xy` as the pixel coordinate.
47    pub fn window() -> Self {
48        Self(WorkgroupsKind::Window)
49    }
50
51    /// One thread per element of the GPU buffer registered with tag `Tag` via
52    /// [`register_gpu_buffer`](crate::ShaderAppExt::register_gpu_buffer).
53    /// Dispatches along x only; use `global_invocation_id.x` as the element index.
54    pub fn over<Tag: 'static>() -> Self {
55        Self(WorkgroupsKind::Buffer {
56            tag: TypeId::of::<Tag>(),
57            tag_name: type_name::<Tag>(),
58        })
59    }
60
61    /// An explicit total number of threads on each axis.
62    pub fn threads(threads: UVec3) -> Self {
63        Self(WorkgroupsKind::Threads(threads))
64    }
65
66    /// Raw workgroup counts, passed straight to `dispatch_workgroups`. No reflection is
67    /// needed, so this also works for shaders whose workgroup size can't be reflected.
68    pub fn exact(workgroups: UVec3) -> Self {
69        Self(WorkgroupsKind::Exact(workgroups))
70    }
71}
72
73/// One configured compute pass: an entry point and how to size its dispatch.
74#[derive(Clone, Debug)]
75pub(crate) struct PassConfig {
76    pub entry: &'static str,
77    pub workgroups: Workgroups,
78    /// Init passes run once (and again on [`ComputePasses::rerun`]) instead of every frame.
79    pub init: bool,
80}
81
82/// Bevy plugin that runs compute shader passes once per frame, before any camera renders.
83///
84/// Every buffer registered through [`ShaderAppExt`](crate::ShaderAppExt) is visible to
85/// both the compute passes and the [`FullscreenFragmentPlugin`](crate::FullscreenFragmentPlugin)
86/// shader, at the same `@group`/`@binding`. Each binding must be declared with the same
87/// access mode (`read` or `read_write`) in every shader that uses it.
88///
89/// Only one `ComputeShaderPlugin` is supported per app. Add it after `DefaultPlugins`.
90///
91/// # Example
92///
93/// ```rust,ignore
94/// struct Particles;
95///
96/// App::new()
97///     .add_plugins(DefaultPlugins)
98///     .add_plugins(
99///         ComputeShaderPlugin::new("shaders/particles.wgsl")
100///             .init_pass("init", Workgroups::over::<Particles>())
101///             .pass("update", Workgroups::over::<Particles>()),
102///     )
103///     .register_gpu_buffer::<Particles, Particle>(0, 0, 1024)
104///     .run();
105/// ```
106pub struct ComputeShaderPlugin {
107    /// Path to the compute shader asset (e.g. `"shaders/simulation.wgsl"`).
108    pub shader_path: &'static str,
109    passes: Vec<PassConfig>,
110}
111
112impl ComputeShaderPlugin {
113    /// Creates a plugin for the compute shader at the given asset path. Add passes with
114    /// [`pass`](Self::pass) and [`init_pass`](Self::init_pass).
115    pub fn new(shader_path: &'static str) -> Self {
116        Self {
117            shader_path,
118            passes: Vec::new(),
119        }
120    }
121
122    /// Run the `@compute` entry point `entry` every frame. Passes run in the order they
123    /// are added, and each sees the writes of the passes before it. The same entry point
124    /// may be added more than once (e.g. for simulation substeps).
125    pub fn pass(mut self, entry: &'static str, workgroups: Workgroups) -> Self {
126        self.passes.push(PassConfig {
127            entry,
128            workgroups,
129            init: false,
130        });
131        self
132    }
133
134    /// Run the `@compute` entry point `entry` once, on the first frame everything is
135    /// ready. Every-frame passes don't start until all init passes have run. Use
136    /// [`ComputePasses::rerun`] to run it again later.
137    pub fn init_pass(mut self, entry: &'static str, workgroups: Workgroups) -> Self {
138        self.passes.push(PassConfig {
139            entry,
140            workgroups,
141            init: true,
142        });
143        self
144    }
145}
146
147/// Main-world control over the compute passes.
148///
149/// ```rust,ignore
150/// fn reset(keys: Res<ButtonInput<KeyCode>>, mut passes: ResMut<ComputePasses>) {
151///     if keys.just_pressed(KeyCode::KeyR) {
152///         passes.rerun("init");
153///     }
154/// }
155/// ```
156#[derive(Resource, Default)]
157pub struct ComputePasses {
158    reruns: Vec<String>,
159}
160
161impl ComputePasses {
162    /// Run the init pass(es) for entry point `entry` again on the next frame.
163    pub fn rerun(&mut self, entry: &str) {
164        self.reruns.push(entry.to_owned());
165    }
166}
167
168/// Main-world configuration, inserted by [`ComputeShaderPlugin`].
169#[derive(Resource)]
170pub(crate) struct ComputeConfig {
171    pub shader_path: &'static str,
172    pub shader: Handle<Shader>,
173    pub passes: Vec<PassConfig>,
174}
175
176/// This frame's dispatch plan, built in the main world and extracted to the render world.
177#[derive(Resource, Default)]
178pub(crate) struct ComputeFrame {
179    /// Workgroup count per pass (same order as the configured passes), or `None` if the
180    /// pass can't be dispatched this frame.
181    pub counts: Vec<Option<UVec3>>,
182    /// Indices of init passes to run again.
183    pub reruns: Vec<usize>,
184}
185
186impl Plugin for ComputeShaderPlugin {
187    fn build(&self, app: &mut App) {
188        add_core_plugin(app);
189
190        let shader = app.world().resource::<AssetServer>().load(self.shader_path);
191
192        app.init_resource::<GpuBufferLengths>()
193            .init_resource::<ComputePasses>()
194            .init_resource::<ComputeReflection>()
195            .init_resource::<ComputeFrame>()
196            .insert_resource(ComputeConfig {
197                shader_path: self.shader_path,
198                shader,
199                passes: self.passes.clone(),
200            })
201            .add_systems(
202                PostUpdate,
203                (reflect::reflect_workgroup_sizes, plan_dispatches).chain(),
204            );
205
206        let Some(render_app) = app.get_sub_app_mut(RenderApp) else {
207            return;
208        };
209
210        render_app
211            .insert_resource(pipeline::ComputeRenderConfig {
212                shader_path: self.shader_path,
213                passes: self.passes.clone(),
214            })
215            .insert_resource(pipeline::ComputeRunState {
216                init_done: vec![false; self.passes.len()],
217            })
218            .init_resource::<pipeline::ComputeRunList>()
219            .init_resource::<ComputeFrame>()
220            .add_systems(ExtractSchedule, extract_compute_frame)
221            .add_systems(
222                RenderStartup,
223                pipeline::init_compute_pipelines.after(ShaderSystems::CompileLayouts),
224            )
225            .add_systems(
226                Render,
227                pipeline::prepare_compute_run
228                    .in_set(RenderSystems::PrepareBindGroups)
229                    .after(ShaderSystems::FinalizeBindGroups),
230            );
231
232        // Before `camera_driver`, so compute work is submitted ahead of every camera.
233        render_app.add_systems(
234            RenderGraph,
235            node::compute_pass
236                .in_set(ShaderPassSystems::Compute)
237                .in_set(RenderGraphSystems::Render)
238                .before(camera_driver),
239        );
240    }
241}
242
243/// Converts a thread count to a workgroup count for the given workgroup size.
244/// Returns `Ok(None)` when there is nothing to dispatch (e.g. a minimised window).
245fn workgroup_count(threads: UVec3, size: UVec3) -> Result<Option<UVec3>, String> {
246    if size.cmpeq(UVec3::ZERO).any() {
247        return Err(format!("invalid workgroup size {size}"));
248    }
249    let count = UVec3::new(
250        threads.x.div_ceil(size.x),
251        threads.y.div_ceil(size.y),
252        threads.z.div_ceil(size.z),
253    );
254    check_count(count)
255}
256
257fn check_count(count: UVec3) -> Result<Option<UVec3>, String> {
258    if count.max_element() > MAX_WORKGROUPS_PER_DIMENSION {
259        return Err(format!(
260            "dispatch of {count} workgroups exceeds the limit of {MAX_WORKGROUPS_PER_DIMENSION} per dimension"
261        ));
262    }
263    if count.cmpeq(UVec3::ZERO).any() {
264        return Ok(None);
265    }
266    Ok(Some(count))
267}
268
269/// Works out this frame's workgroup counts for every pass. Problems are logged once per
270/// pass and message, and that pass is skipped until they're resolved.
271fn plan_dispatches(
272    config: Res<ComputeConfig>,
273    reflection: Res<ComputeReflection>,
274    lengths: Res<GpuBufferLengths>,
275    windows: Query<&Window, With<PrimaryWindow>>,
276    mut passes: ResMut<ComputePasses>,
277    mut frame: ResMut<ComputeFrame>,
278    mut reported: Local<HashSet<(usize, String)>>,
279) {
280    let window_size = windows
281        .single()
282        .ok()
283        .map(|window| UVec3::new(window.physical_width(), window.physical_height(), 1));
284
285    let mut report = |index: usize, message: String| {
286        if reported.insert((index, message.clone())) {
287            warn!(
288                "compute pass `{}` skipped: {message}",
289                config.passes[index].entry
290            );
291        }
292    };
293
294    frame.counts.clear();
295    for (index, pass) in config.passes.iter().enumerate() {
296        let reflected_size = || -> Result<Option<UVec3>, String> {
297            let Some(sizes) = &reflection.sizes else {
298                // Shader not loaded yet (or failed to parse, which is logged separately).
299                return Ok(None);
300            };
301            match sizes.get(pass.entry) {
302                Some(WorkgroupSize::Fixed(size)) => Ok(Some(*size)),
303                Some(WorkgroupSize::Overridden) => Err(
304                    "its @workgroup_size uses override constants; use `Workgroups::exact`".into(),
305                ),
306                None => Err("no compute entry point with this name".into()),
307            }
308        };
309
310        let count = match &pass.workgroups.0 {
311            WorkgroupsKind::Exact(count) => check_count(*count),
312            kind => reflected_size().and_then(|size| {
313                let Some(size) = size else { return Ok(None) };
314                let threads = match kind {
315                    WorkgroupsKind::Window => window_size.unwrap_or(UVec3::ZERO),
316                    WorkgroupsKind::Buffer { tag, tag_name } => {
317                        let Some(&len) = lengths.0.get(tag) else {
318                            return Err(format!(
319                                "no GPU buffer is registered with tag `{tag_name}`"
320                            ));
321                        };
322                        if size.y != 1 || size.z != 1 {
323                            return Err(format!(
324                                "`Workgroups::over` dispatches along x only, but the workgroup size is {size}; use @workgroup_size(N)"
325                            ));
326                        }
327                        UVec3::new(len, 1, 1)
328                    }
329                    WorkgroupsKind::Threads(threads) => *threads,
330                    WorkgroupsKind::Exact(_) => unreachable!(),
331                };
332                workgroup_count(threads, size)
333            }),
334        };
335
336        frame.counts.push(count.unwrap_or_else(|message| {
337            report(index, message);
338            None
339        }));
340    }
341
342    frame.reruns.clear();
343    for entry in passes.reruns.drain(..) {
344        let matching: Vec<usize> = config
345            .passes
346            .iter()
347            .enumerate()
348            .filter(|(_, pass)| pass.init && pass.entry == entry)
349            .map(|(index, _)| index)
350            .collect();
351        if matching.is_empty() {
352            warn!("ComputePasses::rerun: `{entry}` is not an init pass");
353        }
354        frame.reruns.extend(matching);
355    }
356}
357
358/// Copies the plan into the render world's `ComputeFrame`, reusing its memory.
359fn extract_compute_frame(
360    frame: Extract<Res<ComputeFrame>>,
361    mut render_frame: ResMut<ComputeFrame>,
362) {
363    render_frame.counts.clone_from(&frame.counts);
364    render_frame.reruns.clone_from(&frame.reruns);
365}
366
367#[cfg(test)]
368mod tests {
369    use super::*;
370
371    #[test]
372    fn rounds_thread_counts_up() {
373        assert_eq!(
374            workgroup_count(UVec3::new(500, 1, 1), UVec3::new(64, 1, 1)),
375            Ok(Some(UVec3::new(8, 1, 1)))
376        );
377        assert_eq!(
378            workgroup_count(UVec3::new(1920, 1080, 1), UVec3::new(8, 8, 1)),
379            Ok(Some(UVec3::new(240, 135, 1)))
380        );
381    }
382
383    #[test]
384    fn skips_empty_dispatches() {
385        assert_eq!(
386            workgroup_count(UVec3::new(0, 0, 1), UVec3::new(8, 8, 1)),
387            Ok(None)
388        );
389    }
390
391    #[test]
392    fn rejects_oversized_dispatches() {
393        assert!(workgroup_count(UVec3::new(65536 * 64 + 1, 1, 1), UVec3::new(64, 1, 1)).is_err());
394        assert!(check_count(UVec3::new(65536, 1, 1)).is_err());
395    }
396
397    #[test]
398    fn rejects_zero_workgroup_size() {
399        assert!(workgroup_count(UVec3::ONE, UVec3::new(0, 1, 1)).is_err());
400    }
401}
bevy_shader_plugin/compute/mod.rs

bevy_shader_plugin/compute/
mod.rs
