feat: linarize most bached simulation kernels by switching to batch-interleaved formats (#37)

* perf(rbd): batch-interleaved flat buffer layout with demand-sized constraint arenas

* test(rbd): add a batched-stacks parity probe

* refactor(rbd): remove the never-maintained dof_values mirror

* perf(rbd): positional contact slots, per-pair manifold reduction, contact-to-multibody index

* chore(rbd): update docs

* feat(rbd): use stronger typing in shaders when possible

* chore: CI fixes

Author: sebcrozet

src/state.rs

@@ -64,6 +64,11 @@ impl NexusCapacities {
         self
     }
 
+    pub fn rbd_mb_contact_constraints(mut self, capacity: u32) -> Self {
+        self.rbd.mb_contact_constraints_capacity = capacity;
+        self
+    }
+
     #[cfg(feature = "mpm")]
     pub fn mpm_grid_size(mut self, num_chunks: u32) -> Self {
         self.mpm.grid_size = num_chunks;
@@ -109,6 +114,8 @@ pub struct NexusCounts {
     pub multibody_dofs: usize,
     pub collision_pairs: usize,
     pub collision_pairs_capacity: usize,
+    pub mb_contact_constraints: usize,
+    pub mb_contact_constraints_capacity: usize,
     pub particles: usize,
 }
 
@@ -342,6 +349,13 @@ impl NexusState {
         self.capacities.rbd.collisions_capacity = capacity.max(1);
     }
 
+    /// Sets the per-batch multibody contact-constraint slot budget (see
+    /// [`RbdCapacities::mb_contact_constraints_capacity`]). Takes effect on the
+    /// next state (re)build.
+    pub fn set_rbd_mb_contact_constraints_capacity(&mut self, capacity: u32) {
+        self.capacities.rbd.mb_contact_constraints_capacity = capacity.max(1);
+    }
+
     /// Sets the number of rigid-body solver steps advanced per
     /// [`NexusPipeline::simulate`](crate::pipeline::NexusPipeline::simulate) call (default 1). Acts as a simulation-speed control.
     pub fn set_rbd_steps_per_frame(&mut self, steps: u32) {
@@ -373,6 +387,11 @@ impl NexusState {
         if let Some(rbd) = self.rbd.as_ref() {
             c.collision_pairs = rbd.collision_pairs_len() as usize;
             c.collision_pairs_capacity = rbd.collision_pairs_capacity() as usize;
+            #[cfg(feature = "dim3")]
+            {
+                c.mb_contact_constraints = rbd.mb_contact_constraints_len() as usize;
+                c.mb_contact_constraints_capacity = rbd.mb_contact_constraints_capacity() as usize;
+            }
         }
         #[cfg(feature = "mpm")]
         if let Some(mpm) = self.mpm.as_ref() {
@@ -640,13 +659,13 @@ impl NexusState {
                     <= rbd.num_colliders_per_batch() as usize =>
             {
                 let range = rbd.append_bodies(backend, &gpu_pairs)?;
-                // Single environment: the per-batch local slot is the gpu_id.
+                let nb = rbd.num_batches();
                 for (i, (&handle, &coupling)) in handles.iter().zip(&couplings).enumerate() {
                     self.rbd2gpu[0].insert(
                         handle.0,
                         GpuRigidBodyRef {
                             coupling,
-                            gpu_id: range.start + i as u32,
+                            gpu_id: (range.start + i as u32) * nb,
                         },
                     );
                 }
@@ -981,9 +1000,9 @@ impl NexusState {
             // `gpu_id` is its *body* slot, not a collider slot, since a body may
             // own several colliders. Body slots are assigned in the order
             // `from_rapier` uses (the first time each parent body is seen while
-            // iterating colliders) and are laid out env-major with stride
-            // `num_colliders_per_batch`.
-            let stride = rbd_state.num_colliders_per_batch();
+            // iterating colliders); the per-body buffers are batch-interleaved,
+            // so `gpu_id = local_slot * num_batches + env`.
+            let nb = rbd_state.num_batches();
             for (env_idx, world) in self.rbd_envs.iter().enumerate() {
                 let mut body_slot: std::collections::HashMap<_, u32> =
                     std::collections::HashMap::new();
@@ -1012,7 +1031,7 @@ impl NexusState {
                         body_handle.0,
                         GpuRigidBodyRef {
                             coupling,
-                            gpu_id: env_idx as u32 * stride + slot,
+                            gpu_id: slot * nb + env_idx as u32,
                         },
                     );
                 }
@@ -1038,7 +1057,7 @@ impl NexusState {
                             body_handle.0,
                             GpuRigidBodyRef {
                                 coupling,
-                                gpu_id: env_idx as u32 * stride + slot,
+                                gpu_id: slot * nb + env_idx as u32,
                             },
                         );
                     }

src_rbd/broad_phase/lbvh.rs

@@ -7,8 +7,8 @@ use crate::math::Pose;
 use crate::shaders::PaddedVector;
 use crate::shaders::bounding_volumes::Aabb;
 use crate::shaders::broad_phase::{
-    CollisionPair, GpuBfComputeAabbs, GpuBfFindPairs, GpuLbvhBuild, GpuLbvhComputeDomain,
-    GpuLbvhComputeMorton, GpuLbvhFindCollisionPairs, GpuLbvhInitDispatch, GpuLbvhRefitInternal,
+    CollisionPair, GpuBfComputeAabbs, GpuBfFindPairs, GpuFlatListDispatch, GpuLbvhBuild,
+    GpuLbvhComputeDomain, GpuLbvhComputeMorton, GpuLbvhFindCollisionPairs, GpuLbvhRefitInternal,
     GpuLbvhRefitLeaves, GpuLbvhResetCollisionPairs, LbvhNode,
 };
 use crate::shaders::shapes::Shape;
@@ -32,7 +32,9 @@ pub struct GpuLbvh {
     refit_internal: GpuLbvhRefitInternal,
     reset_collision_pairs: GpuLbvhResetCollisionPairs,
     find_collision_pairs: GpuLbvhFindCollisionPairs,
-    lbvh_init_indirect_args: GpuLbvhInitDispatch,
+    /// Writes the `[total/64, 1, 1]` indirect grid from the single global pair
+    /// counter.
+    flat_list_dispatch: GpuFlatListDispatch,
     // Kernels for brute-force broad-phase for small scenes
     // (typically, small scenes but many batches).
     bf_compute_aabbs: GpuBfComputeAabbs,
@@ -294,7 +296,7 @@ impl Lbvh {
 
         self.shaders
             .reset_collision_pairs
-            .call(pass, [num_batches, 1, 1], collision_pairs_len)?;
+            .call(pass, [1u32, 1, 1], collision_pairs_len)?;
         self.shaders.find_collision_pairs.call(
             pass,
             [colliders_per_batch, num_batches, 1],
@@ -305,11 +307,12 @@ impl Lbvh {
             batch_indices,
             pair_filter,
         )?;
-        self.shaders.lbvh_init_indirect_args.call(
+        self.shaders.flat_list_dispatch.call(
             pass,
-            256u32,
+            1u32,
             collision_pairs_len,
             collision_pairs_indirect,
+            batch_indices,
         )?;
         Ok(())
     }
@@ -351,7 +354,7 @@ impl Lbvh {
         )?;
         self.shaders
             .reset_collision_pairs
-            .call(pass, [num_batches, 1, 1], collision_pairs_len)?;
+            .call(pass, [1u32, 1, 1], collision_pairs_len)?;
         self.shaders.bf_find_pairs.call(
             pass,
             [active_per_batch * active_per_batch * num_batches, 1, 1],
@@ -363,12 +366,12 @@ impl Lbvh {
             pair_filter,
             sim_params,
         )?;
-        // Single 256-lane workgroup: parallel max over the per-batch counts.
-        self.shaders.lbvh_init_indirect_args.call(
+        self.shaders.flat_list_dispatch.call(
             pass,
-            256u32,
+            1u32,
             collision_pairs_len,
             collision_pairs_indirect,
+            batch_indices,
         )?;
         Ok(())
     }