Add `compute` app function for running compute shaders

[tychedelia]

compute

Adds a new lifecycle function that can be used to dispatch a compute pass before view is called

There's a bit more boilerplate than usual to implement:

fn main() {
    nannou::app(model)
        .compute(compute)
        .run();
}

// State defines the different stages of our compute shader
#[derive(Default, Debug, Eq, PartialEq, Hash, Clone)]
enum State {
    #[default]
    Init,
    Update,
}

// Our compute model will be passed as bindings to the shader
#[derive(AsBindGroup, Clone)]
struct ComputeModel {
    #[uniform(0)]
    radius: f32,
}

impl Compute for ComputeModel {
    type State = State;

    // The location of the compute shader in our assets folder
    fn shader() -> ShaderRef {
        "foo.wgsl"
    }

    // How to derive the entrypoint for our shader from the state
    fn entry(state: &Self::State) -> &'static str {
        match state {
            State::Init => "init",
            State::Update => "update",
        }
    }

    // The dispatch size for our shader pass
    fn dispatch_size(_state: &Self::State) -> (u32, u32, u32) {
        (1, 1, 1)
    }
}

// The `compute` fn accepts takes the old state as an argument and is expected to produce
// the next state, as well as the compute model (typically derived from the app model) to use
// for this pass
fn compute(app: &App, model: &Model, state: State, view: Entity) -> (State, ComputeModel) {
    match state {
        State::Init => (State::Update, ComputeModel { radius: app.time() }),
        State::Update => (State::Update, ComputeModel { radius: app.time() }),
    }
}

The idea here is that the Compute trait helps define a state machine driven by Compute::State. This of course can only be a single state, but is helpful for the typical init -> update lifecycle seen in things like particle systems.

We also introduce a new concept ComputeModel which is used to populate the uniform passed to the compute shader. This must implement AsBindGroup and the associated Bevy traits necessary to generate a bind group.

Draw commands

Additionally, two new draw commands have been added that are very useful for writing compute shaders:

• Instanced causes the supplied primitive to be rendered over the provided instance range allowing access to instance_index in the vertex shader:

draw.instanced()
    .primitive(draw.rect().w_h(5.0, 5.0))
    .range(0..NUM_PARTICLES);

Typically, for performance reasons we render all primitives that share a material into a single mesh. However, this has the downside of resulting in a single draw call with one instance per mesh. Using explicit instancing helps avoid this.

• Indirect allows the user to supply a Ssbo handle that can be supplied with draw arguments from a compute shader:

  draw.indirect()
  .primitive(draw.rect().w_h(2.0, 2.0))
  .buffer(model.indirect_params.clone());

  This is an advanced technique, but allows eliminating almost all CPU overhead, uploading only the vertex and index buffer for a single mesh each frame. In the future we can support building the vertex buffer in the GPU or techniques like vertex pulling as well.

See the examples for more details:

• game_of_life demonstrates writing to texture storage in a ping pong manner to run a game of life simulation.

https://github.com/user-attachments/assets/9e4b9bfa-b691-4471-a2b9-72f6ddd96413

• particle_mouse sets up a simple particle system that attracts particles to the mouse.

https://github.com/user-attachments/assets/8c474c0c-09f1-45b6-b203-bb59ec123661

• particle_sdf uses a sdf in the compute shader to energize particles.

https://github.com/user-attachments/assets/c7413190-54fc-4e67-99b4-35e738416277

[tombh]

Wow, I just got put onto this from This Week In Bevy. The API is so good 🥹

It can't be used independently as a plugin in Bevy right? The most similar project I know of is https://github.com/AnthonyTornetta/bevy_easy_compute, but it has a couple of hacks at the moment. Did you get inspiration from somewhere or is this all completely made from scratch?

[tychedelia]

Thanks for checking it out!

It can't be used independently as a plugin in Bevy right?

Although this PR is built as a Bevy plugin, it's not currently exposed in a way that would be convenient, but it is definitely a goal! Perhaps this is a prod that I need to refactor it a bit. :) I'm also personally interested in continuing to improve compute infrastructure in Bevy itself. Although our refactoring work still isn't totally mature, our goal is to be interoperable with Bevy.

Did you get inspiration from somewhere or is this all completely made from scratch?

This is all from scratch!

[tombh]

Sure, no pressure! I'm getting by fine with that bevy_easy_compute plugin, I just think your approach is more idiomatic and up to date. So I'm going to study it nonetheless 🤓