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, Clone)]
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            .add_systems(ExtractSchedule, extract_compute_frame)
220            .add_systems(
221                RenderStartup,
222                pipeline::init_compute_pipelines.after(ShaderSystems::CompileLayouts),
223            )
224            .add_systems(
225                Render,
226                pipeline::prepare_compute_run
227                    .in_set(RenderSystems::PrepareBindGroups)
228                    .after(ShaderSystems::FinalizeBindGroups),
229            );
230
231        // Before `camera_driver`, so compute work is submitted ahead of every camera.
232        render_app.add_systems(
233            RenderGraph,
234            node::compute_pass
235                .in_set(ShaderPassSystems::Compute)
236                .in_set(RenderGraphSystems::Render)
237                .before(camera_driver),
238        );
239    }
240}
241
242/// Converts a thread count to a workgroup count for the given workgroup size.
243/// Returns `Ok(None)` when there is nothing to dispatch (e.g. a minimised window).
244fn workgroup_count(threads: UVec3, size: UVec3) -> Result<Option<UVec3>, String> {
245    if size.cmpeq(UVec3::ZERO).any() {
246        return Err(format!("invalid workgroup size {size}"));
247    }
248    let count = UVec3::new(
249        threads.x.div_ceil(size.x),
250        threads.y.div_ceil(size.y),
251        threads.z.div_ceil(size.z),
252    );
253    check_count(count)
254}
255
256fn check_count(count: UVec3) -> Result<Option<UVec3>, String> {
257    if count.max_element() > MAX_WORKGROUPS_PER_DIMENSION {
258        return Err(format!(
259            "dispatch of {count} workgroups exceeds the limit of {MAX_WORKGROUPS_PER_DIMENSION} per dimension"
260        ));
261    }
262    if count.cmpeq(UVec3::ZERO).any() {
263        return Ok(None);
264    }
265    Ok(Some(count))
266}
267
268/// Works out this frame's workgroup counts for every pass. Problems are logged once per
269/// pass and message, and that pass is skipped until they're resolved.
270fn plan_dispatches(
271    config: Res<ComputeConfig>,
272    reflection: Res<ComputeReflection>,
273    lengths: Res<GpuBufferLengths>,
274    windows: Query<&Window, With<PrimaryWindow>>,
275    mut passes: ResMut<ComputePasses>,
276    mut frame: ResMut<ComputeFrame>,
277    mut reported: Local<HashSet<(usize, String)>>,
278) {
279    let window_size = windows
280        .single()
281        .ok()
282        .map(|window| UVec3::new(window.physical_width(), window.physical_height(), 1));
283
284    let mut report = |index: usize, message: String| {
285        if reported.insert((index, message.clone())) {
286            warn!(
287                "compute pass `{}` skipped: {message}",
288                config.passes[index].entry
289            );
290        }
291    };
292
293    frame.counts.clear();
294    for (index, pass) in config.passes.iter().enumerate() {
295        let reflected_size = || -> Result<Option<UVec3>, String> {
296            let Some(sizes) = &reflection.sizes else {
297                // Shader not loaded yet (or failed to parse, which is logged separately).
298                return Ok(None);
299            };
300            match sizes.get(pass.entry) {
301                Some(WorkgroupSize::Fixed(size)) => Ok(Some(*size)),
302                Some(WorkgroupSize::Overridden) => Err(
303                    "its @workgroup_size uses override constants; use `Workgroups::exact`".into(),
304                ),
305                None => Err("no compute entry point with this name".into()),
306            }
307        };
308
309        let count = match &pass.workgroups.0 {
310            WorkgroupsKind::Exact(count) => check_count(*count),
311            kind => reflected_size().and_then(|size| {
312                let Some(size) = size else { return Ok(None) };
313                let threads = match kind {
314                    WorkgroupsKind::Window => window_size.unwrap_or(UVec3::ZERO),
315                    WorkgroupsKind::Buffer { tag, tag_name } => {
316                        let Some(&len) = lengths.0.get(tag) else {
317                            return Err(format!(
318                                "no GPU buffer is registered with tag `{tag_name}`"
319                            ));
320                        };
321                        if size.y != 1 || size.z != 1 {
322                            return Err(format!(
323                                "`Workgroups::over` dispatches along x only, but the workgroup size is {size}; use @workgroup_size(N)"
324                            ));
325                        }
326                        UVec3::new(len, 1, 1)
327                    }
328                    WorkgroupsKind::Threads(threads) => *threads,
329                    WorkgroupsKind::Exact(_) => unreachable!(),
330                };
331                workgroup_count(threads, size)
332            }),
333        };
334
335        frame.counts.push(count.unwrap_or_else(|message| {
336            report(index, message);
337            None
338        }));
339    }
340
341    frame.reruns.clear();
342    for entry in passes.reruns.drain(..) {
343        let matching: Vec<usize> = config
344            .passes
345            .iter()
346            .enumerate()
347            .filter(|(_, pass)| pass.init && pass.entry == entry)
348            .map(|(index, _)| index)
349            .collect();
350        if matching.is_empty() {
351            warn!("ComputePasses::rerun: `{entry}` is not an init pass");
352        }
353        frame.reruns.extend(matching);
354    }
355}
356
357fn extract_compute_frame(mut commands: Commands, frame: Extract<Res<ComputeFrame>>) {
358    commands.insert_resource(frame.clone());
359}
360
361#[cfg(test)]
362mod tests {
363    use super::*;
364
365    #[test]
366    fn rounds_thread_counts_up() {
367        assert_eq!(
368            workgroup_count(UVec3::new(500, 1, 1), UVec3::new(64, 1, 1)),
369            Ok(Some(UVec3::new(8, 1, 1)))
370        );
371        assert_eq!(
372            workgroup_count(UVec3::new(1920, 1080, 1), UVec3::new(8, 8, 1)),
373            Ok(Some(UVec3::new(240, 135, 1)))
374        );
375    }
376
377    #[test]
378    fn skips_empty_dispatches() {
379        assert_eq!(
380            workgroup_count(UVec3::new(0, 0, 1), UVec3::new(8, 8, 1)),
381            Ok(None)
382        );
383    }
384
385    #[test]
386    fn rejects_oversized_dispatches() {
387        assert!(workgroup_count(UVec3::new(65536 * 64 + 1, 1, 1), UVec3::new(64, 1, 1)).is_err());
388        assert!(check_count(UVec3::new(65536, 1, 1)).is_err());
389    }
390
391    #[test]
392    fn rejects_zero_workgroup_size() {
393        assert!(workgroup_count(UVec3::ONE, UVec3::new(0, 1, 1)).is_err());
394    }
395}
bevy_shader_plugin/compute/mod.rs

bevy_shader_plugin/compute/
mod.rs
