Pack multiple effects per particle slab (#508)

Sub-allocate each particle slab with the content of multiple effects.
- Introduce a `base_particle` value per effect instance, which is the
equivalent of the `base_vertex` for rendering, and corresponds to the
index of the first particle in the sub-allocated slice for that effect,
inside the overall slab buffer.
- Store that `base_particle` in the `Spawner` and give access to all
shaders which need it (most of them).
- Restore the default 64k particle count per slab, which allows packing
multiple effects per buffer/slab.

When `debug_assertions` is active (in Debug build), fill the first value
of each particle to a `NaN` (0xFFFFFFFF) to make it easier to see in
RenderDoc or any other GPU debugger that the particle is unused.

Author: djeedai

src/lib.rs

@@ -1248,7 +1248,7 @@ fn append_spawn_events_{0}(base_child_index: u32, particle_index: u32, count: u3
         {
             writeln!(
                 &mut writeback_code,
-                "    particle_buffer.particles[particle_index].{0} = particle.{0};",
+                "    particle_buffer.particles[base_particle + particle_index].{0} = particle.{0};",
                 attribute.name()
             )
             .unwrap();

src/render/effect_cache.rs

@@ -56,8 +56,9 @@ impl PartialOrd for EffectSlice {
 /// A reference to a slice allocated inside an [`ParticleSlab`].
 #[derive(Debug, Default, Clone, PartialEq, Eq)]
 pub struct SlabSliceRef {
-    /// Range into an [`ParticleSlab`], in item count.
+    /// Range into a [`ParticleSlab`], in item count.
     range: Range<u32>,
+    /// Particle layout for the effect stored in that slice.
     pub(crate) particle_layout: ParticleLayout,
 }
 
@@ -223,17 +224,17 @@ pub struct ParticleSlab {
 
 impl ParticleSlab {
     /// Minimum buffer capacity to allocate, in number of particles.
-    // FIXME - Batching is broken due to binding a single GpuSpawnerParam instead of
-    // N, and inability for a particle index to tell which Spawner it should
-    // use. Setting this to 1 effectively ensures that all new buffers just fit
-    // the effect, so batching never occurs.
-    pub const MIN_CAPACITY: u32 = 1; // 65536; // at least 64k particles
+    pub const MIN_CAPACITY: u32 = 65536; // at least 64k particles
 
     /// Create a new slab and the GPU resources to back it up.
     ///
     /// The slab cannot contain less than [`MIN_CAPACITY`] particles. If the
     /// input `capacity` is smaller, it's rounded up to [`MIN_CAPACITY`].
     ///
+    /// # Panics
+    ///
+    /// This panics if the `capacity` is zero.
+    ///
     /// [`MIN_CAPACITY`]: Self::MIN_CAPACITY
     pub fn new(
         slab_id: SlabId,
@@ -252,31 +253,47 @@ impl ParticleSlab {
 
         // Calculate the clamped capacity of the group, in number of particles.
         let capacity = capacity.max(Self::MIN_CAPACITY);
-        debug_assert!(
+        assert!(
             capacity > 0,
             "Attempted to create a zero-sized effect buffer."
         );
 
         // Allocate the particle buffer itself, containing the attributes of each
         // particle.
+        #[cfg(debug_assertions)]
+        let mapped_at_creation = true;
+        #[cfg(not(debug_assertions))]
+        let mapped_at_creation = false;
         let particle_capacity_bytes: BufferAddress =
             capacity as u64 * particle_layout.min_binding_size().get();
         let particle_label = format!("hanabi:buffer:slab{}:particle", slab_id.0);
         let particle_buffer = render_device.create_buffer(&BufferDescriptor {
             label: Some(&particle_label),
             size: particle_capacity_bytes,
             usage: BufferUsages::COPY_DST | BufferUsages::STORAGE,
-            mapped_at_creation: false,
+            mapped_at_creation,
         });
+        // Set content
+        #[cfg(debug_assertions)]
+        {
+            // Scope get_mapped_range_mut() to force a drop before unmap()
+            {
+                let slice: &mut [u8] = &mut particle_buffer
+                    .slice(..particle_capacity_bytes)
+                    .get_mapped_range_mut();
+                let slice: &mut [u32] = cast_slice_mut(slice);
+                slice.fill(0xFFFFFFFF);
+            }
+            particle_buffer.unmap();
+        }
 
         // Each indirect buffer stores 3 arrays of u32, of length the number of
         // particles.
-        let capacity_bytes: BufferAddress = capacity as u64 * 4 * 3;
-
+        let indirect_capacity_bytes: BufferAddress = capacity as u64 * 4 * 3;
         let indirect_label = format!("hanabi:buffer:slab{}:indirect", slab_id.0);
         let indirect_index_buffer = render_device.create_buffer(&BufferDescriptor {
             label: Some(&indirect_label),
-            size: capacity_bytes,
+            size: indirect_capacity_bytes,
             usage: BufferUsages::COPY_DST | BufferUsages::STORAGE,
             mapped_at_creation: true,
         });
@@ -285,11 +302,11 @@ impl ParticleSlab {
             // Scope get_mapped_range_mut() to force a drop before unmap()
             {
                 let slice: &mut [u8] = &mut indirect_index_buffer
-                    .slice(..capacity_bytes)
+                    .slice(..indirect_capacity_bytes)
                     .get_mapped_range_mut();
                 let slice: &mut [u32] = cast_slice_mut(slice);
                 for index in 0..capacity {
-                    slice[3 * index as usize + 2] = capacity - 1 - index;
+                    slice[3 * index as usize + 2] = index;
                 }
             }
             indirect_index_buffer.unmap();
@@ -337,11 +354,11 @@ impl ParticleSlab {
             },
         ];
         let label = format!(
-            "hanabi:bind_group_layout:render:particles@1:vfx{}",
+            "hanabi:bind_group_layout:render:particles@1:slab{}",
             slab_id.0
         );
         trace!(
-            "Creating render layout '{}' with {} entries",
+            "Creating particles@1 layout '{}' for render pass with {} entries",
             label,
             entries.len(),
         );
@@ -378,16 +395,6 @@ impl ParticleSlab {
         &self.indirect_index_buffer
     }
 
-    #[inline]
-    pub fn particle_offset(&self, row: u32) -> u32 {
-        self.particle_layout.min_binding_size().get() as u32 * row
-    }
-
-    #[inline]
-    pub fn indirect_index_offset(&self, row: u32) -> u32 {
-        row * 12
-    }
-
     /// Return a binding for the entire particle buffer.
     pub fn as_entire_binding_particle(&self) -> BindingResource<'_> {
         let capacity_bytes = self.capacity as u64 * self.particle_layout.min_binding_size().get();

src/render/event.rs

@@ -166,6 +166,8 @@ pub(crate) struct CachedParentInfo {
 pub(crate) struct CachedChildInfo {
     /// ID of the slab storing the parent effect.
     pub parent_slab_id: SlabId,
+    /// Offset into the slab of the parent's particles.
+    pub parent_slab_offset: u32,
     /// Parent's particle layout.
     pub parent_particle_layout: ParticleLayout,
     /// Parent's buffer.
@@ -186,6 +188,7 @@ pub(crate) struct CachedChildInfo {
 impl CachedChildInfo {
     pub fn is_locally_equal(&self, other: &CachedChildInfo) -> bool {
         self.parent_slab_id == other.parent_slab_id
+        && self.parent_slab_offset == other.parent_slab_offset
         && self.parent_particle_layout == other.parent_particle_layout
         && self.parent_buffer_binding_source == other.parent_buffer_binding_source
         && self.local_child_index == other.local_child_index

src/render/mod.rs

@@ -361,6 +361,13 @@ pub(crate) struct GpuSpawnerParams {
     /// Index of the [`GpuDrawIndirect`] or [`GpuDrawIndexedIndirect`] for this
     /// effect.
     draw_indirect_index: u32,
+    /// Start offset of the particles and indirect indices into the effect's
+    /// slab, in number of particles (row index).
+    slab_offset: u32,
+    /// Start offset of the particles and indirect indices into the parent effect's
+    /// slab (if the effect has a parent effect), in number of particles (row index).
+    /// This is ignored if the effect has no parent.
+    parent_slab_offset: u32,
 }
 
 /// GPU representation of an indirect compute dispatch input.
@@ -2933,6 +2940,8 @@ impl EffectsMeta {
         global_transform: &GlobalTransform,
         spawn_count: u32,
         prng_seed: u32,
+        slab_offset: u32,
+        parent_slab_offset: Option<u32>,
         effect_metadata_buffer_table_id: BufferTableId,
         maybe_cached_draw_indirect_args: Option<&CachedDrawIndirectArgs>,
     ) -> u32 {
@@ -2953,6 +2962,8 @@ impl EffectsMeta {
             draw_indirect_index: maybe_cached_draw_indirect_args
                 .map(|cdia| cdia.get_row().0)
                 .unwrap_or_default(),
+            slab_offset,
+            parent_slab_offset: parent_slab_offset.unwrap_or(u32::MAX),
             ..default()
         };
         trace!("spawner params = {:?}", spawner_params);
@@ -3539,6 +3550,7 @@ pub fn allocate_parent_child_infos(
 
         let new_cached_child_info = CachedChildInfo {
             parent_slab_id: parent_cached_effect.slab_id,
+            parent_slab_offset: parent_cached_effect.slice.range().start,
             parent_particle_layout: parent_cached_effect.slice.particle_layout.clone(),
             parent_buffer_binding_source,
             local_child_index,
@@ -4422,11 +4434,15 @@ pub(crate) fn prepare_batch_inputs(
         trace!("layout_flags = {:?}", extracted_effect.layout_flags);
         trace!("particle_layout = {:?}", effect_slice.particle_layout);
 
+        let parent_slab_offset = maybe_cached_child_info.map(|cci| cci.parent_slab_offset);
+
         assert!(cached_effect_metadata.table_id.is_valid());
         let spawner_index = effects_meta.allocate_spawner(
             &extracted_spawner.transform,
             extracted_spawner.spawn_count,
             extracted_spawner.prng_seed,
+            cached_effect.slice.range().start,
+            parent_slab_offset,
             cached_effect_metadata.table_id,
             maybe_cached_draw_indirect_args,
         );
@@ -6410,6 +6426,7 @@ pub(crate) fn prepare_bind_groups(
                 particle_buffer,
                 indirect_index_buffer,
                 effect_metadata_buffer,
+                &spawner_buffer,
             ) {
                 error!(
                     "Failed to create sort-fill bind group @0 for ribbon effect: {:?}",
@@ -6420,9 +6437,11 @@ pub(crate) fn prepare_bind_groups(
 
             // Bind group @0 of sort-copy pass
             let indirect_index_buffer = effect_buffer.indirect_index_buffer();
-            if let Err(err) = sort_bind_groups
-                .ensure_sort_copy_bind_group(indirect_index_buffer, effect_metadata_buffer)
-            {
+            if let Err(err) = sort_bind_groups.ensure_sort_copy_bind_group(
+                indirect_index_buffer,
+                effect_metadata_buffer,
+                &spawner_buffer,
+            ) {
                 error!(
                     "Failed to create sort-copy bind group @0 for ribbon effect: {:?}",
                     err
@@ -7343,6 +7362,11 @@ impl Node for VfxSimulateNode {
                         return Ok(());
                     }
 
+                    let spawner_base = effect_batch.spawner_base;
+                    let spawner_aligned_size = effects_meta.spawner_buffer.aligned_size();
+                    assert!(spawner_aligned_size >= GpuSpawnerParams::min_size().get() as usize);
+                    let spawner_offset = spawner_base * spawner_aligned_size as u32;
+
                     // Bind group sort_fill@0
                     let particle_buffer = effect_buffer.particle_buffer();
                     let indirect_index_buffer = effect_buffer.indirect_index_buffer();
@@ -7355,21 +7379,14 @@ impl Node for VfxSimulateNode {
                         compute_pass.insert_debug_marker("ERROR:MissingSortFillBindGroup");
                         continue;
                     };
-                    let particle_offset = effect_buffer.particle_offset(effect_batch.slice.start);
-                    let indirect_index_offset =
-                        effect_buffer.indirect_index_offset(effect_batch.slice.start);
                     let effect_metadata_offset = effects_meta
                         .gpu_limits
                         .effect_metadata_offset(effect_batch.metadata_table_id.0)
                         as u32;
                     compute_pass.set_bind_group(
                         0,
                         bind_group,
-                        &[
-                            particle_offset,
-                            indirect_index_offset,
-                            effect_metadata_offset,
-                        ],
+                        &[effect_metadata_offset, spawner_offset],
                     );
 
                     compute_pass
@@ -7418,6 +7435,11 @@ impl Node for VfxSimulateNode {
                         return Ok(());
                     }
 
+                    let spawner_base = effect_batch.spawner_base;
+                    let spawner_aligned_size = effects_meta.spawner_buffer.aligned_size();
+                    assert!(spawner_aligned_size >= GpuSpawnerParams::min_size().get() as usize);
+                    let spawner_offset = spawner_base * spawner_aligned_size as u32;
+
                     // Bind group sort_copy@0
                     let indirect_index_buffer = effect_buffer.indirect_index_buffer();
                     let Some(bind_group) = sort_bind_groups.sort_copy_bind_group(
@@ -7428,14 +7450,13 @@ impl Node for VfxSimulateNode {
                         compute_pass.insert_debug_marker("ERROR:MissingSortCopyBindGroup");
                         continue;
                     };
-                    let indirect_index_offset = effect_batch.slice.start;
                     let effect_metadata_offset = effects_meta
                         .effect_metadata_buffer
                         .dynamic_offset(effect_batch.metadata_table_id);
                     compute_pass.set_bind_group(
                         0,
                         bind_group,
-                        &[indirect_index_offset, effect_metadata_offset],
+                        &[effect_metadata_offset, spawner_offset],
                     );
 
                     compute_pass