perf(rbd): flat 1-D narrow-phase dispatch — pack warps across batches

src_rbd/broad_phase/narrow_phase.rs

@@ -4,9 +4,9 @@ use crate::math::Pose;
 use crate::queries::GpuIndexedContact;
 use crate::shaders::PaddedVector;
 use crate::shaders::broad_phase::{
-    CollisionPair, GpuInitPfmPfmDispatch, GpuNarrowPhaseInitContactsDispatch, GpuNarrowPhasePfmPfm,
-    GpuNarrowPhaseShapeShape, GpuNarrowPhaseShapeShapeDeferred, GpuResetNarrowPhase,
-    NarrowPhasePfmPair,
+    CollisionPair, GpuFlattenBatchesDispatch, GpuNarrowPhaseInitContactsDispatch,
+    GpuNarrowPhasePfmPfm, GpuNarrowPhaseShapeShape, GpuNarrowPhaseShapeShapeDeferred,
+    GpuResetNarrowPhase, NarrowPhasePfmPair,
 };
 use crate::shaders::shapes::Shape;
 use khal::Shader;
@@ -22,7 +22,11 @@ pub struct GpuNarrowPhase {
     /// `pfm_pairs` work-list. Split from `narrow_phase` to fit 8 storage buffers.
     narrow_phase_deferred: GpuNarrowPhaseShapeShapeDeferred,
     narrow_phase_pfm_pfm: GpuNarrowPhasePfmPfm,
-    init_pfm_pfm_indirect_args: GpuInitPfmPfmDispatch,
+    /// Builds the flat 1-D dispatch grid + prefix offsets for a per-batch
+    /// work-list (used for both the collision pairs and the PFM pairs), so the
+    /// kernels pack items from many batches into full warps instead of one
+    /// mostly-idle workgroup per batch.
+    flatten_batches: GpuFlattenBatchesDispatch,
     init_contacts_indirect_args: GpuNarrowPhaseInitContactsDispatch,
 }
 
@@ -37,8 +41,8 @@ impl GpuNarrowPhase {
         vertices: &Tensor<PaddedVector>,
         indices: &Tensor<u32>,
         collision_pairs: &Tensor<CollisionPair>,
-        collision_pairs_len: &Tensor<u32>,
-        collision_pairs_indirect: &Tensor<[u32; 3]>,
+        collision_pairs_len: &mut Tensor<u32>,
+        collision_pairs_indirect: &mut Tensor<[u32; 3]>,
         contacts: &mut Tensor<GpuIndexedContact>,
         contacts_len: &mut Tensor<u32>,
         contacts_indirect: &mut Tensor<[u32; 3]>,
@@ -48,16 +52,30 @@ impl GpuNarrowPhase {
         batch_indices: &Tensor<crate::shaders::utils::BatchIndices>,
         collider_parent: &Tensor<u32>,
         collider_materials: &Tensor<crate::shaders::queries::ColliderMaterial>,
+        pairs_offsets: &mut Tensor<u32>,
+        pfm_offsets: &mut Tensor<u32>,
     ) -> Result<(), GpuBackendError> {
         let num_batches = contacts_len.len() as u32;
         self.reset_narrow_phase
             .call(pass, [1u32, num_batches, 1], contacts_len, pfm_pairs_len)?;
 
-        self.narrow_phase.call(
+        // The broad phase wrote a `[max/64, num_batches, 1]` grid into
+        // `collision_pairs_indirect`; rewrite it (and derive the offsets) for
+        // the flat layout. Nothing else consumes the batched form.
+        self.flatten_batches.call(
             pass,
+            1u32,
+            collision_pairs_len,
+            pairs_offsets,
             collision_pairs_indirect,
+            batch_indices,
+        )?;
+
+        self.narrow_phase.call(
+            pass,
+            &*collision_pairs_indirect,
             collision_pairs,
-            collision_pairs_len,
+            pairs_offsets,
             poses,
             shapes,
             contacts,
@@ -71,9 +89,9 @@ impl GpuNarrowPhase {
         // separate dispatch so each pass fits 8 storage buffers).
         self.narrow_phase_deferred.call(
             pass,
-            collision_pairs_indirect,
+            &*collision_pairs_indirect,
             collision_pairs,
-            collision_pairs_len,
+            pairs_offsets,
             poses,
             shapes,
             pfm_pairs,
@@ -83,15 +101,21 @@ impl GpuNarrowPhase {
             indices,
         )?;
 
-        self.init_pfm_pfm_indirect_args
-            .call(pass, 1u32, pfm_pairs_len, pfm_pairs_indirect)?;
+        self.flatten_batches.call(
+            pass,
+            1u32,
+            pfm_pairs_len,
+            pfm_offsets,
+            pfm_pairs_indirect,
+            batch_indices,
+        )?;
         self.narrow_phase_pfm_pfm.call(
             pass,
             &*pfm_pairs_indirect,
             contacts,
             contacts_len,
             pfm_pairs,
-            pfm_pairs_len,
+            pfm_offsets,
             batch_indices,
             vertices,
             indices,

src_rbd/pipeline/insertion_removal.rs

@@ -168,6 +168,10 @@ impl RbdState {
             BufferUsages::STORAGE | BufferUsages::COPY_SRC,
         )
         .unwrap();
+        let pairs_flat_offsets =
+            Tensor::vector_uninit(backend, num_batches + 1, BufferUsages::STORAGE).unwrap();
+        let pfm_flat_offsets =
+            Tensor::vector_uninit(backend, num_batches + 1, BufferUsages::STORAGE).unwrap();
         let old_constraints =
             Tensor::vector_uninit(backend, collisions_capacity * num_batches, storage).unwrap();
         let old_constraint_builders =
@@ -262,6 +266,8 @@ impl RbdState {
             pfm_pairs,
             pfm_pairs_len,
             pfm_pairs_indirect,
+            pairs_flat_offsets,
+            pfm_flat_offsets,
             old_constraints,
             old_constraint_builders,
             old_constraints_counts,

src_rbd/pipeline/rbd_state.rs

@@ -184,6 +184,12 @@ pub struct RbdState {
     pub(super) pfm_pairs: Tensor<NarrowPhasePfmPair>,
     pub(super) pfm_pairs_len: Tensor<u32>,
     pub(super) pfm_pairs_indirect: Tensor<[u32; 3]>,
+    /// Flat-dispatch prefix offsets (`num_batches + 1`) over the per-batch
+    /// collision-pair / PFM work-lists, rebuilt on the GPU each step by
+    /// `gpu_flatten_batches_dispatch` so the narrow-phase kernels can pack
+    /// items from many batches into full warps.
+    pub(super) pairs_flat_offsets: Tensor<u32>,
+    pub(super) pfm_flat_offsets: Tensor<u32>,
     pub(super) contacts: Tensor<GpuIndexedContact>,
     pub(super) contacts_len: Tensor<u32>,
     pub(super) contacts_indirect: Tensor<[u32; 3]>,

src_rbd/pipeline/rbd_state_from_rapier.rs

@@ -595,6 +595,10 @@ impl RbdState {
             BufferUsages::STORAGE | BufferUsages::COPY_SRC,
         )
         .unwrap();
+        let pairs_flat_offsets =
+            Tensor::vector_uninit(backend, num_batches + 1, BufferUsages::STORAGE).unwrap();
+        let pfm_flat_offsets =
+            Tensor::vector_uninit(backend, num_batches + 1, BufferUsages::STORAGE).unwrap();
         let old_constraints = Tensor::vector_uninit(
             backend,
             capacities.collisions_capacity * num_batches,
@@ -753,6 +757,8 @@ impl RbdState {
             pfm_pairs,
             pfm_pairs_len,
             pfm_pairs_indirect,
+            pairs_flat_offsets,
+            pfm_flat_offsets,
             old_constraints,
             old_constraint_builders,
             old_constraints_counts,

src_rbd/pipeline/rbd_step.rs

@@ -186,8 +186,8 @@ impl RbdPipeline {
                 &state.vertex_buffers,
                 &state.index_buffers,
                 &state.collision_pairs,
-                &state.collision_pairs_len,
-                &state.collision_pairs_indirect,
+                &mut state.collision_pairs_len,
+                &mut state.collision_pairs_indirect,
                 &mut state.contacts,
                 &mut state.contacts_len,
                 &mut state.contacts_indirect,
@@ -197,6 +197,8 @@ impl RbdPipeline {
                 &state.batch_indices,
                 &state.collider_parent,
                 &state.collider_materials,
+                &mut state.pairs_flat_offsets,
+                &mut state.pfm_flat_offsets,
             )?;
 
             drop(pass);