Finish sdf example.

examples/assets/shaders/particle_sdf_compute.wgsl

@@ -6,11 +6,20 @@ struct Particle {
     color: vec4<f32>,
 };
 
+struct DrawIndirectArgs {
+    index_count: u32,
+    instance_count: atomic<u32>,
+    first_index: u32,
+    base_vertex: i32,
+    first_instance: u32,
+}
+
 @group(0) @binding(0) var<storage, read_write> particles: array<Particle>;
 @group(0) @binding(1) var<uniform> sphere_center: vec4<f32>;
 @group(0) @binding(2) var<uniform> sphere_radius: f32;
-@group(0) @binding(3) var<uniform> particle_count: u32;
+@group(0) @binding(3) var<uniform> scaling_factor: u32;
 @group(0) @binding(4) var<uniform> resolution: vec2<u32>;
+@group(0) @binding(5) var<storage, read_write> indirect_args: array<DrawIndirectArgs>;
 
 fn sdf_sphere(point: vec3<f32>, center: vec3<f32>, radius: f32) -> f32 {
     return length(point - center) - radius;
@@ -48,6 +57,14 @@ fn closest_point_on_sphere(point: vec2<f32>) -> vec3<f32> {
 
 @compute @workgroup_size(64)
 fn update(@builtin(global_invocation_id) global_id: vec3<u32>) {
+    let pixels = resolution.x * resolution.y;
+    let particle_count = pixels / scaling_factor;
+
+    /// Update our indirect params if we are thread 0
+    if (global_id.x == 0u && global_id.y == 0u && global_id.z == 0u) {
+        atomicStore(&indirect_args[0].instance_count, particle_count);
+    }
+
     let index = global_id.x;
     if (index >= particle_count) {
         return;
@@ -159,6 +176,14 @@ fn update(@builtin(global_invocation_id) global_id: vec3<u32>) {
 
 @compute @workgroup_size(64)
 fn init(@builtin(global_invocation_id) global_id: vec3<u32>) {
+    let pixels = resolution.x * resolution.y;
+    let particle_count = pixels / scaling_factor;
+
+    /// Update our indirect params if we are thread 0
+    if (global_id.x == 0u && global_id.y == 0u && global_id.z == 0u) {
+        atomicStore(&indirect_args[0].instance_count, particle_count);
+    }
+
     let index = global_id.x;
     if (index >= particle_count) {
         return;
@@ -177,7 +202,7 @@ fn init(@builtin(global_invocation_id) global_id: vec3<u32>) {
     let grid_x = f32(index % u32(grid_width)) / (grid_width - 1.0);
     let grid_y = floor(f32(index) / grid_width) / (grid_height - 1.0);
 
-    // Map grid position to screen space (-1 to 1 for both axes)
+    // Map grid position to screen space
     let screen_x = grid_x * 2.0 - 1.0;
     let screen_y = grid_y * 2.0 - 1.0;
 

examples/compute/particle_sdf.rs

@@ -3,7 +3,6 @@ use nannou::prelude::bevy_render::storage::ShaderStorageBuffer;
 use nannou::prelude::*;
 use std::sync::Arc;
 
-const NUM_PARTICLES: u32 = 100000;
 const WORKGROUP_SIZE: u32 = 64;
 
 fn main() {
@@ -24,6 +23,9 @@ struct Model {
     particles: Handle<ShaderStorageBuffer>,
     indirect_params: Handle<ShaderStorageBuffer>,
     circle_radius: f32,
+    size: u32,
+    buffer_size: u32,
+    scaling_factor: u32,
 }
 
 impl Model {
@@ -34,38 +36,47 @@ impl Model {
     }
 }
 
+// This struct isn't used on the CPU side, but it's necessary to define the layout of the data
+// to correctly size the buffer on the GPU side.
 #[repr(C)]
-#[derive(ShaderType, Copy, Clone, Debug, Default, bytemuck::Pod, bytemuck::Zeroable)]
+#[derive(ShaderType)]
 struct Particle {
     position: Vec2,
+    original_position: Vec2,
     velocity: Vec2,
+    energy: f32,
+    _pad: f32,
     color: Vec4,
 }
 
+// The draw indirect args struct is used to pass the number of instances to draw to the GPU
+// The compute shader will write the number of particles to draw to `instance_count`
 #[repr(C)]
-#[derive(ShaderType, Copy, Clone, Debug, Default, bytemuck::Pod, bytemuck::Zeroable)]
+#[derive(ShaderType)]
 struct DrawIndirectArgs {
-    /// The number of indices to draw.
     pub index_count: u32,
-    /// The number of instances to draw.
     pub instance_count: u32,
-    /// The first index within the index buffer.
     pub first_index: u32,
-    /// The value added to the vertex index before indexing into the vertex buffer.
     pub base_vertex: i32,
-    /// The instance ID of the first instance to draw.
-    ///
-    /// Has to be 0, unless [`Features::INDIRECT_FIRST_INSTANCE`](crate::Features::INDIRECT_FIRST_INSTANCE) is enabled.
     pub first_instance: u32,
 }
 
-#[derive(Default, Debug, Eq, PartialEq, Hash, Clone)]
+// Our compute shader has two states, `Init` and `Update`
+// The `Init` state is used to initialize the particles
+// The `Update` state is used to run the simulation
+#[derive(Debug, Eq, PartialEq, Hash, Clone)]
 enum State {
-    #[default]
-    Init,
-    Update,
+    Init(u32),
+    Update(u32),
 }
 
+impl Default for State {
+    fn default() -> Self {
+        State::Init(1)
+    }
+}
+
+// Compute model, defining uniforms and storage buffers that will be passed to the compute shader
 #[derive(AsBindGroup, Clone)]
 struct ComputeModel {
     #[storage(0, visibility(compute))]
@@ -75,9 +86,11 @@ struct ComputeModel {
     #[uniform(2)]
     circle_radius: f32,
     #[uniform(3)]
-    particle_count: u32,
+    scaling_factor: u32,
     #[uniform(4)]
     resolution: UVec2,
+    #[storage(5, visibility(compute))]
+    indirect_params: Handle<ShaderStorageBuffer>,
 }
 
 impl Compute for ComputeModel {
@@ -89,16 +102,22 @@ impl Compute for ComputeModel {
 
     fn entry(state: &Self::State) -> &'static str {
         match state {
-            State::Init => "init",
-            State::Update => "update",
+            State::Init(_) => "init",
+            State::Update(_) => "update",
         }
     }
 
-    fn dispatch_size(_state: &Self::State) -> (u32, u32, u32) {
-        ((NUM_PARTICLES + WORKGROUP_SIZE - 1) / WORKGROUP_SIZE, 1, 1)
+    fn dispatch_size(state: &Self::State) -> (u32, u32, u32) {
+        let size = match state {
+            State::Init(size) => size,
+            State::Update(size) => size,
+        };
+
+        ((size + WORKGROUP_SIZE - 1) / WORKGROUP_SIZE, 1, 1)
     }
 }
 
+// Our shader model that will be used to render the particles
 #[shader_model(
     fragment = "shaders/particle_sdf_material.wgsl",
     vertex = "shaders/particle_sdf_material.wgsl"
@@ -117,44 +136,59 @@ fn model(app: &App) -> Model {
         .build();
 
     // Create a buffer to store the particles.
-    let particle_size = Particle::min_size().get() as usize;
-    let mut particles = ShaderStorageBuffer::with_size(
-        NUM_PARTICLES as usize * particle_size * 2,
-        RenderAssetUsages::RENDER_WORLD,
-    );
-    particles.buffer_description.label = Some("particles");
-    particles.buffer_description.usage |= BufferUsages::STORAGE | BufferUsages::VERTEX;
-    let particles = app.assets_mut().add(particles);
+    let size = 1; // This will be updated in the `update` function
+    let particles = create_particle_buffer(app, size);
 
     // Create a buffer store our indirect draw params
-    let mut indirect_params = ShaderStorageBuffer::from(DrawIndirectArgs {
-        index_count: 21,
-        instance_count: NUM_PARTICLES,
-        first_index: 0,
-        base_vertex: 0,
-        first_instance: 0,
-    });
-    indirect_params.buffer_description.label = Some("indirect_params");
-    indirect_params.buffer_description.usage |= BufferUsages::STORAGE | BufferUsages::INDIRECT;
-    let indirect_params = app.assets_mut().add(indirect_params);
+    let indirect_params = create_indirect_params_buffer(app, size);
 
     Model {
         particles,
         indirect_params,
         circle_radius: 0.5,
+        size,
+        buffer_size: size,
+        scaling_factor: 50,
     }
 }
 
+// Controls:
+// - Arrow keys to change the circle radius
+// - Left and right arrow keys to change the scaling factor (i.e. density)
 fn update(app: &App, model: &mut Model) {
     if app.keys().pressed(KeyCode::ArrowUp) {
         model.circle_radius += 0.01;
     }
     if app.keys().pressed(KeyCode::ArrowDown) {
         model.circle_radius -= 0.01;
     }
+    if app.keys().pressed(KeyCode::ArrowRight) {
+        model.scaling_factor += 1;
+    }
+    if app.keys().pressed(KeyCode::ArrowLeft) {
+        // Don't let the scaling factor go below 20
+        model.scaling_factor = model.scaling_factor.saturating_sub(1).max(20);
+    }
+
+
+    let pixels = app.main_window().size_pixels().element_product();
+    let new_size = pixels / model.scaling_factor.clamp(20, 1000);
+
+    // Resize the buffer if necessary
+    if new_size > model.buffer_size {
+        model.particles = create_particle_buffer(app, new_size);
+        model.buffer_size = new_size;
+    }
+
+    model.size = new_size;
 }
 
-fn compute(app: &App, model: &Model, state: State, view: Entity) -> (State, ComputeModel) {
+fn compute(
+    app: &App,
+    model: &Model,
+    previous_state: State,
+    _view: Entity,
+) -> (State, ComputeModel) {
     let window = app.main_window();
     let window_rect = window.rect();
 
@@ -168,13 +202,30 @@ fn compute(app: &App, model: &Model, state: State, view: Entity) -> (State, Comp
         particles: model.particles.clone(),
         circle_center: mouse_norm.extend(1.0).extend(0.0),
         circle_radius: model.circle_radius,
-        particle_count: NUM_PARTICLES,
+        scaling_factor: model.scaling_factor,
         resolution: window.size_pixels(),
+        indirect_params: model.indirect_params.clone(),
     };
 
-    match state {
-        State::Init => (State::Update, compute_model),
-        State::Update => (State::Update, compute_model),
+    // If the size has changed, we need to re-initialize the particles
+    // Otherwise, we can just update them
+    match previous_state {
+        State::Init(size) => {
+            // Even if we are in the `Init` state, we still need to update the size of the buffer
+            // in case the window has been resized
+            if size != model.size {
+                (State::Init(model.size), compute_model)
+            } else {
+                (State::Update(model.size), compute_model)
+            }
+        }
+        State::Update(size) => {
+            if size != model.size {
+                (State::Init(model.size), compute_model)
+            } else {
+                (State::Update(model.size), compute_model)
+            }
+        }
     }
 }
 
@@ -184,6 +235,32 @@ fn view(app: &App, model: &Model) {
 
     let draw = draw.material(model.material());
     draw.indirect()
-        .primitive(draw.ellipse().w_h(5.0, 5.0))
+        .primitive(draw.rect().w_h(2.0, 2.0))
         .buffer(model.indirect_params.clone());
 }
+
+fn create_particle_buffer(app: &App, size: u32) -> Handle<ShaderStorageBuffer> {
+    let particle_size = Particle::min_size().get() as usize;
+    let mut particles = ShaderStorageBuffer::with_size(
+        size as usize * particle_size * 2,
+        RenderAssetUsages::RENDER_WORLD,
+    );
+    particles.buffer_description.label = Some("particles");
+    particles.buffer_description.usage |= BufferUsages::STORAGE | BufferUsages::VERTEX;
+    let particles = app.assets_mut().add(particles);
+    particles
+}
+
+fn create_indirect_params_buffer(app: &App, size: u32) -> Handle<ShaderStorageBuffer> {
+    let mut indirect_params = ShaderStorageBuffer::from(DrawIndirectArgs {
+        index_count: 6, // Hardcoded for now, 2 triangles
+        instance_count: size,
+        first_index: 0,
+        base_vertex: 0,
+        first_instance: 0,
+    });
+    indirect_params.buffer_description.label = Some("indirect_params");
+    indirect_params.buffer_description.usage |= BufferUsages::STORAGE | BufferUsages::INDIRECT;
+    let indirect_params = app.assets_mut().add(indirect_params);
+    indirect_params
+}