Diff that causes Naga validation error

shaders/physics/src/cell.rs

@@ -5,7 +5,9 @@ use spirv_std::{arch::IndexUnchecked, glam::Vec2, num_traits::Euclid};
 
 use wrach_cpu_gpu_shared::{self as shared, WorldSettings};
 
-use crate::{particle::Particle, particles::Particles, PREFIX_SUM_OFFSET_HACK};
+use crate::{
+    particle::Particle, particles::Particles, PREFIX_SUM_OFFSET_HACK, WORKGROUP_MEMORY_SIZE,
+};
 
 /// The amount of extra space for over-packed cells. If the `MIN_DISTANCE` is right then this
 /// should not generally be needed. I think it's most useful for the very beginning of a simulation
@@ -32,6 +34,9 @@ pub struct Cell<'world> {
     pub current: usize,
     /// Is this the final spatial bin cell?
     pub is_last: bool,
+
+    pub first_particle_index: usize,
+
     /// Config, like viewport position etc.
     pub settings: &'world WorldSettings,
 
@@ -49,9 +54,9 @@ pub struct Cell<'world> {
     /// Indices of aux particles that surround main cells.
     pub indices_aux: &'world mut [u32],
     /// Aux cell particle positions for reading.
-    pub positions_aux: &'world [Vec2],
+    pub positions_workgroup: &'world [Vec2; WORKGROUP_MEMORY_SIZE],
     /// Aux cell particle positions for writing.
-    pub velocities_aux: &'world [Vec2],
+    pub velocities_workgroup: &'world [Vec2; WORKGROUP_MEMORY_SIZE],
 }
 
 impl Cell<'_> {
@@ -70,10 +75,11 @@ impl Cell<'_> {
             self.indices_aux,
             self.current,
             self.settings.grid_dimensions.x,
+            self.first_particle_index,
             self.positions_in,
             self.velocities_in,
-            self.positions_aux,
-            self.velocities_aux,
+            self.positions_workgroup,
+            self.velocities_workgroup,
         );
         particles.pairs_in_cell();
         particles.auxiliares_around_cell();
@@ -157,48 +163,49 @@ impl Cell<'_> {
     }
 }
 
-#[cfg(test)]
-mod tests {
-
-    use bevy::math::{UVec2, Vec2};
-
-    use super::*;
-
-    #[test]
-    fn border_checks() {
-        let cell = Cell {
-            current: 0,
-            is_last: false,
-            settings: &WorldSettings {
-                view_dimensions: Vec2::ZERO,
-                view_anchor: Vec2::ZERO,
-                grid_dimensions: UVec2::new(10, 10),
-                cell_size: 2,
-                particles_in_frame_count: 0,
-            },
-
-            indices_main: &mut [],
-            positions_in: &[],
-            velocities_in: &[],
-            positions_out: &mut [],
-            velocities_out: &mut [],
-
-            indices_aux: &mut [],
-            positions_aux: &mut [],
-            velocities_aux: &[],
-        };
-
-        assert!(!cell.is_cell_at_right_edge(0));
-        assert!(!cell.is_cell_at_right_edge(1));
-        assert!(cell.is_cell_at_right_edge(9));
-        assert!(cell.is_cell_at_right_edge(59));
-        assert!(cell.is_cell_at_right_edge(99));
-        assert!(!cell.is_cell_at_right_edge(100));
-
-        assert!(!cell.is_cell_at_top_edge(0));
-        assert!(!cell.is_cell_at_top_edge(55));
-        assert!(cell.is_cell_at_top_edge(90));
-        assert!(cell.is_cell_at_top_edge(99));
-        assert!(!cell.is_cell_at_top_edge(100));
-    }
-}
+// #[cfg(test)]
+// mod tests {
+//
+//     use bevy::math::{UVec2, Vec2};
+//
+//     use super::*;
+//
+//     #[test]
+//     fn border_checks() {
+//         let cell = Cell {
+//             current: 0,
+//             is_last: false,
+//             first_particle_index: 0,
+//             settings: &WorldSettings {
+//                 view_dimensions: Vec2::ZERO,
+//                 view_anchor: Vec2::ZERO,
+//                 grid_dimensions: UVec2::new(10, 10),
+//                 cell_size: 2,
+//                 particles_in_frame_count: 0,
+//             },
+//
+//             indices_main: &mut [],
+//             positions_in: &[],
+//             velocities_in: &[],
+//             positions_out: &mut [],
+//             velocities_out: &mut [],
+//
+//             indices_aux: &mut [],
+//             positions_workgroup: &mut [],
+//             velocities_workgroup: &[],
+//         };
+//
+//         assert!(!cell.is_cell_at_right_edge(0));
+//         assert!(!cell.is_cell_at_right_edge(1));
+//         assert!(cell.is_cell_at_right_edge(9));
+//         assert!(cell.is_cell_at_right_edge(59));
+//         assert!(cell.is_cell_at_right_edge(99));
+//         assert!(!cell.is_cell_at_right_edge(100));
+//
+//         assert!(!cell.is_cell_at_top_edge(0));
+//         assert!(!cell.is_cell_at_top_edge(55));
+//         assert!(cell.is_cell_at_top_edge(90));
+//         assert!(cell.is_cell_at_top_edge(99));
+//         assert!(!cell.is_cell_at_top_edge(100));
+//     }
+// }

shaders/physics/src/indices.rs

@@ -8,6 +8,7 @@ use crate::cell::MAX_PARTICLES_IN_CELL;
 #[allow(clippy::missing_docs_in_private_items)]
 #[derive(Debug)]
 pub struct Indices {
+    pub workgroup_offset: usize,
     pub centre: Storage,
     pub bottom_left: Storage,
     pub bottom_right: Storage,
@@ -48,37 +49,49 @@ impl Location {
 
 impl Indices {
     /// Instantiate
-    pub fn new(index_main: &[u32], index_aux: &[u32], cell_index: usize, grid_width: u32) -> Self {
+    pub fn new(
+        index_main: &[u32],
+        index_aux: &[u32],
+        cell_index: usize,
+        grid_width: u32,
+        workgroup_offset: usize,
+    ) -> Self {
         Self {
+            workgroup_offset,
             centre: Self::get_start_end_indices_of_particles_in_cell(
                 index_main,
                 Location::CENTRE,
                 cell_index,
                 grid_width,
+                0,
             ),
             bottom_left: Self::get_start_end_indices_of_particles_in_cell(
                 index_aux,
                 Location::BOTTOM_LEFT,
                 cell_index,
                 grid_width,
+                workgroup_offset,
             ),
             bottom_right: Self::get_start_end_indices_of_particles_in_cell(
                 index_aux,
                 Location::BOTTOM_RIGHT,
                 cell_index,
                 grid_width,
+                workgroup_offset,
             ),
             top_left: Self::get_start_end_indices_of_particles_in_cell(
                 index_aux,
                 Location::TOP_LEFT,
                 cell_index,
                 grid_width,
+                workgroup_offset,
             ),
             top_right: Self::get_start_end_indices_of_particles_in_cell(
                 index_aux,
                 Location::TOP_RIGHT,
                 cell_index,
                 grid_width,
+                workgroup_offset,
             ),
         }
     }
@@ -90,6 +103,7 @@ impl Indices {
         location: u32,
         centre_cell_index: usize,
         grid_width: u32,
+        workgroup_offset: usize,
     ) -> Storage {
         let aux_grid_width = grid_width as usize + 2;
 
@@ -126,8 +140,8 @@ impl Indices {
             end_by: particles_end_by as usize,
         };
         let local = Range {
-            from: local_index,
-            end_by: local_index + particles_count as usize,
+            from: local_index - workgroup_offset,
+            end_by: local_index + particles_count as usize - workgroup_offset,
         };
 
         Storage { global, local }
@@ -150,7 +164,7 @@ mod tests {
                 0, 0, 0,
                 0, 0, 1
         ];
-        Indices::new(&indices_main, &indices_aux, cell_index, 2)
+        Indices::new(&indices_main, &indices_aux, cell_index, 2, 0)
     }
 
     #[test]

shaders/physics/src/lib.rs

@@ -9,8 +9,9 @@
 #![allow(clippy::explicit_counter_loop)]
 #![allow(clippy::needless_range_loop)]
 
-use cell::Cell;
+use cell::{Cell, MAX_PARTICLES_IN_CELL};
 use spirv_std::{
+    arch::workgroup_memory_barrier_with_group_sync,
     glam::{UVec3, Vec2},
     spirv,
 };
@@ -21,10 +22,15 @@ mod indices;
 mod particle;
 mod particles;
 
+const THREADS_PER_WORKGROUP: u32 = 64;
+pub const WORKGROUP_MEMORY_SIZE: usize = THREADS_PER_WORKGROUP as usize * MAX_PARTICLES_IN_CELL * 4;
+
 /// Physics entrypoint
 #[spirv(compute(threads(64)))]
 pub fn main(
-    #[spirv(global_invocation_id)] id: UVec3,
+    // #[spirv(global_invocation_id)] global_invocation: UVec3,
+    #[spirv(workgroup_id)] workgroup_invocation: UVec3,
+    #[spirv(local_invocation_id)] local_invocation: UVec3,
     #[spirv(uniform, descriptor_set = 0, binding = 0)] settings: &WorldSettings,
 
     #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] indices_main: &mut [u32],
@@ -36,12 +42,54 @@ pub fn main(
     #[spirv(storage_buffer, descriptor_set = 0, binding = 6)] indices_aux: &mut [u32],
     #[spirv(storage_buffer, descriptor_set = 0, binding = 7)] positions_aux: &[Vec2],
     #[spirv(storage_buffer, descriptor_set = 0, binding = 8)] velocities_aux: &[Vec2],
+
+    #[spirv(workgroup)] positions_workgroup: &mut [Vec2; WORKGROUP_MEMORY_SIZE],
+    #[spirv(workgroup)] velocities_workgroup: &mut [Vec2; WORKGROUP_MEMORY_SIZE],
 ) {
-    let current_cell = (id.x + PREFIX_SUM_OFFSET_HACK) as usize;
+    // let current_cell = (global_invocation.x + PREFIX_SUM_OFFSET_HACK) as usize;
+    let first_cell_in_workgroup = workgroup_invocation.x * THREADS_PER_WORKGROUP;
+    let last_cell_in_workgroup = first_cell_in_workgroup + THREADS_PER_WORKGROUP - 1;
+
+    let current_cell =
+        (first_cell_in_workgroup + local_invocation.x + PREFIX_SUM_OFFSET_HACK) as usize;
+
+    let mut first_particle_index: usize = 0;
+    if local_invocation.x == 1 {
+        let grid_width = settings.grid_dimensions.x;
+        let aux_steps = [0, 1, grid_width, grid_width + 1];
+        let mut workgroup_memory_index = 0;
+        for cell_index in first_cell_in_workgroup..=last_cell_in_workgroup {
+            for s in 0..4 {
+                let step = aux_steps[s];
+                let aux_index = (cell_index + step) as usize;
+                let (particles_start_at, marker) =
+                    (indices_aux[aux_index], indices_aux[aux_index + 1]);
+                if first_particle_index == 0 {
+                    first_particle_index = particles_start_at as usize;
+                }
+                let particles_count = marker - particles_start_at;
+                let particles_end_by = particles_start_at + particles_count;
+                for particle_index in particles_start_at..particles_end_by {
+                    positions_workgroup[workgroup_memory_index] =
+                        positions_aux[particle_index as usize];
+                    velocities_workgroup[workgroup_memory_index] =
+                        velocities_aux[particle_index as usize];
+
+                    workgroup_memory_index += 1;
+                }
+            }
+        }
+    }
+
+    // Safety: I think `unsafe` is required because this can be difficult to manage correctly.
+    unsafe {
+        workgroup_memory_barrier_with_group_sync();
+    };
 
     let mut world = Cell {
         current: current_cell,
         is_last: false,
+        first_particle_index,
         settings,
 
         indices_main,
@@ -51,8 +99,8 @@ pub fn main(
         velocities_out,
 
         indices_aux,
-        positions_aux,
-        velocities_aux,
+        positions_workgroup,
+        velocities_workgroup,
     };
 
     world.physics_for_cell();

shaders/physics/src/particle.rs

@@ -6,6 +6,8 @@ use spirv_std::{
 };
 use wrach_cpu_gpu_shared::WorldSettings;
 
+use crate::WORKGROUP_MEMORY_SIZE;
+
 /// Convenient representation of a particle
 #[derive(Default, Copy, Clone)]
 pub struct Particle {
@@ -37,6 +39,24 @@ impl<'particle> Particle {
         }
     }
 
+    pub fn new_aux(
+        index: usize,
+        positions_input: &'particle [Vec2; WORKGROUP_MEMORY_SIZE],
+        velocities_input: &'particle [Vec2; WORKGROUP_MEMORY_SIZE],
+    ) -> Self {
+        // SAFETY:
+        //   Getting data with bounds checks is obviously undefined behaviour. We rely on the
+        //   rest of the pipeline to ensure that indices are always within limits.
+        #[allow(clippy::multiple_unsafe_ops_per_block)]
+        unsafe {
+            Self {
+                index,
+                position: *positions_input.index_unchecked(index),
+                velocity: *velocities_input.index_unchecked(index),
+            }
+        }
+    }
+
     /// Enforce particle limits like bouundaries and speed
     pub fn enforce_limits(&mut self, world_config: &WorldSettings) {
         self.enforce_boundaries(world_config);