perf(rbd): fuse the solver color loops — one workgroup per env (stacked on #21)

src_rbd/broad_phase/lbvh.rs

@@ -232,7 +232,8 @@ impl Lbvh {
         let mut pass = encoder.begin_pass("[RBD] lbvh-refit_leaves", timestamps.as_deref_mut());
         self.shaders.refit_leaves.call(
             &mut pass,
-            [colliders_per_batch, num_batches, 1],
+            // Flat over all batches' leaves (kernel recovers batch by division).
+            [colliders_per_batch * num_batches, 1, 1],
             poses,
             shapes,
             &state.sorted_colliders,

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/dynamics/solver.rs

@@ -13,7 +13,8 @@ use crate::shaders::dynamics::{
     GpuApplySolverVelsInc, GpuInitSolverBodies, GpuInitSolverVelsInc, GpuIntegrateLinearized,
     GpuRemoveCfmAndBiasKernel, GpuSolverCleanup, GpuSolverCountConstraints, GpuSolverFinalize,
     GpuSolverIncColor, GpuSolverInitConstraints, GpuSolverResetColor, GpuSolverSortConstraints,
-    GpuSolverUpdateConstraints, GpuStepGaussSeidel, GpuWarmstart, LocalMassProperties,
+    GpuSolverUpdateConstraints, GpuStepGaussSeidel, GpuStepGaussSeidelFused,
+    GpuStepGaussSeidelFusedNoBias, GpuWarmstart, GpuWarmstartFused, LocalMassProperties,
     RbdSimParams, TwoBodyConstraint, TwoBodyConstraintBuilder, Velocity, WorldMassProperties,
 };
 use crate::utils::{GpuPrefixSum, PrefixSumWorkspace};
@@ -40,6 +41,13 @@ pub struct GpuSolver {
     warmstart: GpuWarmstart,
     /// Gauss-Seidel iteration step (sequential per color).
     step_gauss_seidel: GpuStepGaussSeidel,
+    /// One-workgroup-per-env fused variants of the color loops (velocities
+    /// staged in shared memory, colors ordered by workgroup barriers). Used
+    /// when every batch's body count fits the shared stage; the per-color
+    /// dispatch chain above is the fallback.
+    warmstart_fused: GpuWarmstartFused,
+    step_gauss_seidel_fused: GpuStepGaussSeidelFused,
+    step_gauss_seidel_fused_no_bias: GpuStepGaussSeidelFusedNoBias,
     /// Initializes solver velocity increments.
     init_solver_vels_inc: GpuInitSolverVelsInc,
     /// Seeds the COM-centered solver poses from the body world poses
@@ -71,6 +79,13 @@ pub struct SolverArgs<'a> {
     pub contacts_len: &'a Tensor<u32>,
     /// Indirect dispatch arguments based on contact count.
     pub contacts_len_indirect: &'a Tensor<[u32; 3]>,
+    /// `num_colors` as a single-element uniform, consumed by the fused color
+    /// kernels (must equal [`Self::num_colors`]).
+    pub num_colors_uniform: &'a Tensor<u32>,
+    /// Run the fused one-workgroup-per-env color kernels instead of the
+    /// per-color dispatch chain. Only valid when every batch's body count is
+    /// <= `FUSED_SOLVE_MAX_BODIES`.
+    pub fuse_color_loops: bool,
     /// Solver constraints (output from constraint initialization).
     pub constraints: &'a mut Tensor<TwoBodyConstraint>,
     /// Builder data for initializing constraints.
@@ -293,62 +308,100 @@ impl GpuSolver {
              * P1/F1 — integrate velocities (apply `a · dt'` / gravity increment).
              */
             mb_phase!(substep_integrate_velocities);
-            self.apply_solver_vels_inc.call(
-                pass,
-                [args.num_colliders, args.num_batches, 1],
-                args.solver_vels,
-                args.solver_vels_inc,
-                args.batch_indices,
-            )?;
+            if !args.fuse_color_loops {
+                // Fused path folds the increment into gpu_warmstart_fused's
+                // shared-memory staging.
+                self.apply_solver_vels_inc.call(
+                    pass,
+                    [args.num_colliders, args.num_batches, 1],
+                    args.solver_vels,
+                    args.solver_vels_inc,
+                    args.batch_indices,
+                )?;
+            }
 
             /*
              * P2/F2 — build + warmstart constraints.
              */
             mb_phase!(substep_build_constraints);
-            self.update_constraints.call(
-                pass,
-                args.contacts_len_indirect,
-                args.constraints,
-                args.constraint_builders,
-                args.contacts_len,
-                args.solver_body_poses,
-                args.sim_params,
-                args.batch_indices,
-            )?;
-            joint_solver.update(pass, &mut joint_args, args.solver_body_poses)?;
-            self.reset_color.call(pass, 1u32, args.curr_color)?;
-            for _ in 0..args.num_colors {
-                self.warmstart.call(
+            if !args.fuse_color_loops {
+                // Fused path folds the constraint refresh into
+                // gpu_warmstart_fused's prologue.
+                self.update_constraints.call(
                     pass,
                     args.contacts_len_indirect,
                     args.constraints,
+                    args.constraint_builders,
+                    args.contacts_len,
+                    args.solver_body_poses,
+                    args.sim_params,
+                    args.batch_indices,
+                )?;
+            }
+            joint_solver.update(pass, &mut joint_args, args.solver_body_poses)?;
+            if args.fuse_color_loops {
+                self.warmstart_fused.call(
+                    pass,
+                    [64u32, args.num_batches, 1],
+                    args.constraints,
                     args.solver_vels,
                     args.constraints_colors,
                     args.contacts_len,
-                    args.curr_color,
+                    args.num_colors_uniform,
                     args.batch_indices,
+                    args.solver_vels_inc,
+                    args.constraint_builders,
+                    args.solver_body_poses,
+                    args.sim_params,
                 )?;
-                self.inc_color.call(pass, 1u32, args.curr_color)?
+            } else {
+                self.reset_color.call(pass, 1u32, args.curr_color)?;
+                for _ in 0..args.num_colors {
+                    self.warmstart.call(
+                        pass,
+                        args.contacts_len_indirect,
+                        args.constraints,
+                        args.solver_vels,
+                        args.constraints_colors,
+                        args.contacts_len,
+                        args.curr_color,
+                        args.batch_indices,
+                    )?;
+                    self.inc_color.call(pass, 1u32, args.curr_color)?
+                }
             }
 
             /*
              * P3/F3 — solve ALL joints + contacts WITH bias.
              */
             mb_phase!(substep_solve_with_bias);
             joint_solver.solve(pass, &mut joint_args, args.solver_vels, true)?;
-            self.reset_color.call(pass, 1u32, args.curr_color)?;
-            for _ in 0..args.num_colors {
-                self.step_gauss_seidel.call(
+            if args.fuse_color_loops {
+                self.step_gauss_seidel_fused.call(
                     pass,
-                    args.contacts_len_indirect,
+                    [64u32, args.num_batches, 1],
                     args.constraints,
                     args.solver_vels,
                     args.constraints_colors,
                     args.contacts_len,
-                    args.curr_color,
+                    args.num_colors_uniform,
                     args.batch_indices,
                 )?;
-                self.inc_color.call(pass, 1u32, args.curr_color)?
+            } else {
+                self.reset_color.call(pass, 1u32, args.curr_color)?;
+                for _ in 0..args.num_colors {
+                    self.step_gauss_seidel.call(
+                        pass,
+                        args.contacts_len_indirect,
+                        args.constraints,
+                        args.solver_vels,
+                        args.constraints_colors,
+                        args.contacts_len,
+                        args.curr_color,
+                        args.batch_indices,
+                    )?;
+                    self.inc_color.call(pass, 1u32, args.curr_color)?
+                }
             }
 
             /*
@@ -369,26 +422,43 @@ impl GpuSolver {
              */
             mb_phase!(substep_solve_no_bias);
             joint_solver.solve(pass, &mut joint_args, args.solver_vels, false)?;
-            self.remove_cfm_and_bias_kernel.call(
-                pass,
-                args.contacts_len_indirect,
-                args.constraints,
-                args.contacts_len,
-                args.batch_indices,
-            )?;
-            self.reset_color.call(pass, 1u32, args.curr_color)?;
-            for _ in 0..args.num_colors {
-                self.step_gauss_seidel.call(
+            if !args.fuse_color_loops {
+                // Fused path folds the CFM/bias strip into the no-bias kernel's
+                // prologue.
+                self.remove_cfm_and_bias_kernel.call(
                     pass,
                     args.contacts_len_indirect,
                     args.constraints,
+                    args.contacts_len,
+                    args.batch_indices,
+                )?;
+            }
+            if args.fuse_color_loops {
+                self.step_gauss_seidel_fused_no_bias.call(
+                    pass,
+                    [64u32, args.num_batches, 1],
+                    args.constraints,
                     args.solver_vels,
                     args.constraints_colors,
                     args.contacts_len,
-                    args.curr_color,
+                    args.num_colors_uniform,
                     args.batch_indices,
                 )?;
-                self.inc_color.call(pass, 1u32, args.curr_color)?
+            } else {
+                self.reset_color.call(pass, 1u32, args.curr_color)?;
+                for _ in 0..args.num_colors {
+                    self.step_gauss_seidel.call(
+                        pass,
+                        args.contacts_len_indirect,
+                        args.constraints,
+                        args.solver_vels,
+                        args.constraints_colors,
+                        args.contacts_len,
+                        args.curr_color,
+                        args.batch_indices,
+                    )?;
+                    self.inc_color.call(pass, 1u32, args.curr_color)?
+                }
             }
         }
 

src_rbd/pipeline/insertion_removal.rs

@@ -168,6 +168,16 @@ 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 num_colors_uniform = Tensor::scalar(
+            backend,
+            capacities.solver_colors + 1,
+            BufferUsages::UNIFORM | BufferUsages::COPY_DST,
+        )
+        .unwrap();
         let old_constraints =
             Tensor::vector_uninit(backend, collisions_capacity * num_batches, storage).unwrap();
         let old_constraint_builders =
@@ -262,6 +272,10 @@ impl RbdState {
             pfm_pairs,
             pfm_pairs_len,
             pfm_pairs_indirect,
+            pairs_flat_offsets,
+            pfm_flat_offsets,
+            num_colors_uniform,
+            num_colors_uniform_cpu: capacities.solver_colors + 1,
             old_constraints,
             old_constraint_builders,
             old_constraints_counts,

src_rbd/pipeline/rbd_state.rs

@@ -184,6 +184,16 @@ 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>,
+    /// `max_colors + 1` as a uniform for the fused color-loop kernels, plus a
+    /// CPU mirror so the tensor is only recreated when the count grows.
+    pub(super) num_colors_uniform: Tensor<u32>,
+    pub(super) num_colors_uniform_cpu: u32,
     pub(super) contacts: Tensor<GpuIndexedContact>,
     pub(super) contacts_len: Tensor<u32>,
     pub(super) contacts_indirect: Tensor<[u32; 3]>,