Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add compute app function for running compute shaders #976

Merged
merged 22 commits into from
Sep 16, 2024

Conversation

tychedelia
Copy link
Collaborator

@tychedelia tychedelia commented Sep 7, 2024

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.
game_of_life.mp4
  • particle_mouse sets up a simple particle system that attracts particles to the mouse.
particles.mp4
  • particle_sdf uses a sdf in the compute shader to energize particles.
indirect_particles.mp4

@tombh
Copy link

tombh commented Sep 13, 2024

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
Copy link
Collaborator Author

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
Copy link

tombh commented Sep 14, 2024

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 🤓

@tychedelia tychedelia merged commit 83a59b6 into nannou-org:bevy-refactor Sep 16, 2024
14 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants