feat: optimizations of the rigid-body solver

.run/all_examples3.run.xml

@@ -1,7 +1,7 @@
 <component name="ProjectRunConfigurationManager">
   <configuration default="false" name="all_examples3" type="CargoCommandRunConfiguration" factoryName="Cargo Command">
     <option name="buildProfileId" value="release" />
-    <option name="command" value="run --package nexus_examples_3d --bin all_examples3 --features cpu" />
+    <option name="command" value="run --package nexus_examples_3d --bin all_examples3 --features cpu,metal" />
     <option name="workingDirectory" value="file://$PROJECT_DIR$" />
     <envs />
     <option name="emulateTerminal" value="false" />

crates/examples3d/mujoco_menagerie3.rs

@@ -320,6 +320,12 @@ async fn load_scene(
     // MJCF is Z-up: gravity points along -Z (set after `finalize`, which builds
     // the rigid-body state with the default -Y gravity).
     state.set_rbd_gravity(viewer.backend(), [0.0, 0.0, -9.81]);
+    // MuJoCo-style explicit coriolis: the mass matrix / LU / gravity solve
+    // runs once per step instead of once per substep. This matches how
+    // MuJoCo integrates these models and saves ~25% of the step time.
+    if let Some(rbd) = state.rbd.as_mut() {
+        rbd.multibodies_mut().set_implicit_coriolis(false);
+    }
     Ok(state)
 }
 

src_rbd/broad_phase/lbvh.rs

@@ -7,9 +7,9 @@ use crate::math::Pose;
 use crate::shaders::PaddedVector;
 use crate::shaders::bounding_volumes::Aabb;
 use crate::shaders::broad_phase::{
-    CollisionPair, GpuLbvhBuild, GpuLbvhComputeDomain, GpuLbvhComputeMorton,
-    GpuLbvhFindCollisionPairs, GpuLbvhInitDispatch, GpuLbvhRefitInternal, GpuLbvhRefitLeaves,
-    GpuLbvhResetCollisionPairs, LbvhNode,
+    CollisionPair, GpuBfComputeAabbs, GpuBfFindPairs, GpuLbvhBuild, GpuLbvhComputeDomain,
+    GpuLbvhComputeMorton, GpuLbvhFindCollisionPairs, GpuLbvhInitDispatch, GpuLbvhRefitInternal,
+    GpuLbvhRefitLeaves, GpuLbvhResetCollisionPairs, LbvhNode,
 };
 use crate::shaders::shapes::Shape;
 use crate::utils::{RadixSort, RadixSortWorkspace};
@@ -33,6 +33,10 @@ pub struct GpuLbvh {
     reset_collision_pairs: GpuLbvhResetCollisionPairs,
     find_collision_pairs: GpuLbvhFindCollisionPairs,
     lbvh_init_indirect_args: GpuLbvhInitDispatch,
+    // Kernels for brute-force broad-phase for small scenes
+    // (typically, small scenes but many batches).
+    bf_compute_aabbs: GpuBfComputeAabbs,
+    bf_find_pairs: GpuBfFindPairs,
 }
 
 /// GPU-resident state for LBVH construction and queries.
@@ -53,6 +57,10 @@ pub struct LbvhState {
     sorted_colliders: Tensor<u32>,
     tree: Tensor<LbvhNode>,
     sort_workspace: RadixSortWorkspace,
+    /// Per-collider world AABBs, only used by the brute-force tiny-batch path
+    /// (strided by the per-batch collider capacity, like the other
+    /// per-collider buffers).
+    aabbs: Tensor<Aabb>,
 }
 
 /// High-level LBVH broad-phase interface.
@@ -79,6 +87,7 @@ impl LbvhState {
             sorted_colliders: Tensor::vector_uninit(backend, 0, usages).unwrap(),
             tree: Tensor::vector_uninit(backend, 0, usages).unwrap(),
             sort_workspace: RadixSortWorkspace::new(backend),
+            aabbs: Tensor::vector_uninit(backend, 0, usages).unwrap(),
             buffer_usages: usages,
         }
     }
@@ -135,6 +144,15 @@ impl LbvhState {
             self.n_sort_active = None;
         }
     }
+
+    /// Sizes the brute-force path's AABB buffer (per-collider, capacity
+    /// stride). Kept separate from [`Self::resize_buffers`] so the tree
+    /// buffers aren't allocated when only the brute-force path runs.
+    fn resize_bf_buffers(&mut self, backend: &GpuBackend, colliders_len: u32) {
+        if (self.aabbs.len() as u32) < colliders_len {
+            self.aabbs = Tensor::vector_uninit(backend, colliders_len, self.buffer_usages).unwrap();
+        }
+    }
 }
 
 impl Lbvh {
@@ -269,13 +287,14 @@ impl Lbvh {
         collision_pairs_len: &mut Tensor<u32>,
         collision_pairs_indirect: &mut Tensor<[u32; 3]>,
         collision_groups: &Tensor<crate::rapier::geometry::InteractionGroups>,
+        pair_filter: &Tensor<[u32; 2]>,
     ) -> Result<(), GpuBackendError> {
         // One thread per live collider (leaf); padding slots aren't in the tree.
         let colliders_per_batch = active_per_batch;
 
         self.shaders.reset_collision_pairs.call(
             pass,
-            [1u32, num_batches, 1],
+            [num_batches, 1, 1],
             collision_pairs_len,
         )?;
         self.shaders.find_collision_pairs.call(
@@ -286,13 +305,78 @@ impl Lbvh {
             collision_pairs_len,
             collision_groups,
             batch_indices,
+            pair_filter,
+        )?;
+        self.shaders.lbvh_init_indirect_args.call(
+            pass,
+            256u32,
+            collision_pairs_len,
+            collision_pairs_indirect,
+        )?;
+        Ok(())
+    }
+
+    /// Brute-force O(n²) replacement for [`Self::update_tree`] +
+    /// [`Self::find_pairs`], used when each batch holds at most
+    /// [`Self::BRUTE_FORCE_MAX_COLLIDERS`] colliders. One AABB pass and one
+    /// all-pairs pass emit the same pair set as the whole tree pipeline.
+    #[allow(clippy::too_many_arguments)]
+    pub fn brute_force_pairs(
+        &self,
+        backend: &GpuBackend,
+        pass: &mut GpuPass,
+        state: &mut LbvhState,
+        colliders_len: u32,
+        active_per_batch: u32,
+        num_batches: u32,
+        poses: &Tensor<Pose>,
+        vertex_buffers: &Tensor<PaddedVector>,
+        shapes: &Tensor<Shape>,
+        batch_indices: &Tensor<crate::shaders::utils::BatchIndices>,
+        collision_pairs: &mut Tensor<CollisionPair>,
+        collision_pairs_len: &mut Tensor<u32>,
+        collision_pairs_indirect: &mut Tensor<[u32; 3]>,
+        collision_groups: &Tensor<crate::rapier::geometry::InteractionGroups>,
+        pair_filter: &Tensor<[u32; 2]>,
+    ) -> Result<(), GpuBackendError> {
+        state.resize_bf_buffers(backend, colliders_len);
+
+        self.shaders.bf_compute_aabbs.call(
+            pass,
+            [active_per_batch * num_batches, 1, 1],
+            poses,
+            shapes,
+            &mut state.aabbs,
+            batch_indices,
+            vertex_buffers,
+        )?;
+        self.shaders.reset_collision_pairs.call(
+            pass,
+            [num_batches, 1, 1],
+            collision_pairs_len,
+        )?;
+        self.shaders.bf_find_pairs.call(
+            pass,
+            [active_per_batch * active_per_batch * num_batches, 1, 1],
+            &state.aabbs,
+            collision_pairs,
+            collision_pairs_len,
+            collision_groups,
+            batch_indices,
+            pair_filter,
         )?;
+        // Single 256-lane workgroup: parallel max over the per-batch counts.
         self.shaders.lbvh_init_indirect_args.call(
             pass,
-            1u32,
+            256u32,
             collision_pairs_len,
             collision_pairs_indirect,
         )?;
         Ok(())
     }
 }
+
+/// Per-batch collider count at or below which the pipeline uses
+/// [`Lbvh::brute_force_pairs`] instead of building a tree.
+/// `NEXUS_DISABLE_BF=1` forces the LBVH path (A/B debugging).
+pub const BRUTE_FORCE_MAX_COLLIDERS: u32 = 64;

src_rbd/broad_phase/narrow_phase.rs

@@ -42,6 +42,7 @@ impl GpuNarrowPhase {
         contacts: &mut Tensor<GpuIndexedContact>,
         contacts_len: &mut Tensor<u32>,
         contacts_indirect: &mut Tensor<[u32; 3]>,
+        mb_sweep_indirect: &mut Tensor<[u32; 3]>,
         pfm_pairs: &mut Tensor<NarrowPhasePfmPair>,
         pfm_pairs_len: &mut Tensor<u32>,
         pfm_pairs_indirect: &mut Tensor<[u32; 3]>,
@@ -51,7 +52,7 @@ impl GpuNarrowPhase {
     ) -> 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)?;
+            .call(pass, [num_batches, 1, 1], contacts_len, pfm_pairs_len)?;
 
         self.narrow_phase.call(
             pass,
@@ -84,7 +85,7 @@ impl GpuNarrowPhase {
         )?;
 
         self.init_pfm_pfm_indirect_args
-            .call(pass, 1u32, pfm_pairs_len, pfm_pairs_indirect)?;
+            .call(pass, 256u32, pfm_pairs_len, pfm_pairs_indirect)?;
         self.narrow_phase_pfm_pfm.call(
             pass,
             &*pfm_pairs_indirect,
@@ -98,8 +99,14 @@ impl GpuNarrowPhase {
             collider_parent,
             collider_materials,
         )?;
-        self.init_contacts_indirect_args
-            .call(pass, 1u32, contacts_len, contacts_indirect)?;
+        self.init_contacts_indirect_args.call(
+            pass,
+            256u32,
+            contacts_len,
+            contacts_indirect,
+            mb_sweep_indirect,
+            batch_indices,
+        )?;
 
         Ok(())
     }

src_rbd/dynamics/coloring.rs

@@ -10,6 +10,7 @@
 use crate::pipeline::RunStats;
 use crate::shaders::dynamics::TwoBodyConstraint;
 use crate::shaders::dynamics::{
+    GpuColorBucketsCount, GpuColorBucketsReset, GpuColorBucketsScan, GpuColorBucketsScatter,
     GpuFixConflictsTopoGc, GpuResetCompletionFlagTopoGc, GpuResetLuby, GpuResetTopoGc,
     GpuStepGraphColoringLuby, GpuStepGraphColoringTopoGc,
 };
@@ -33,6 +34,32 @@ pub struct GpuColoring {
     /// Detects and fixes conflicts in TOPO-GC coloring.
     fix_conflicts_topo_gc_kernel: GpuFixConflictsTopoGc,
     reset_completion_flag_topo_gc: GpuResetCompletionFlagTopoGc,
+    // Workspace for bucket-sorting constraint ids by color so each color iteration
+    // only touches their own constraint.
+    color_buckets_reset: GpuColorBucketsReset,
+    color_buckets_count: GpuColorBucketsCount,
+    color_buckets_scan: GpuColorBucketsScan,
+    color_buckets_scatter: GpuColorBucketsScatter,
+}
+
+/// Buffers for the per-color constraint bucket sort.
+pub struct ColorBucketsArgs<'a> {
+    /// Indirect dispatch arguments based on contact count.
+    pub contacts_len_indirect: &'a Tensor<[u32; 3]>,
+    /// Color assigned to each constraint by graph coloring.
+    pub constraints_colors: &'a Tensor<u32>,
+    /// Number of contacts per batch.
+    pub contacts_len: &'a Tensor<u32>,
+    /// Per-batch per-color counts (stride `solver_color_buckets_stride`).
+    pub color_bucket_counts: &'a mut Tensor<u32>,
+    /// Per-batch per-color exclusive prefix sums.
+    pub color_bucket_starts: &'a mut Tensor<u32>,
+    /// Scatter cursors (seeded from the starts).
+    pub color_bucket_cursors: &'a mut Tensor<u32>,
+    /// Constraint ids bucket-sorted by color (contacts layout).
+    pub color_sorted_ids: &'a mut Tensor<u32>,
+    /// Shared per-batch capacity / section-offset uniform.
+    pub batch_indices: &'a Tensor<crate::shaders::utils::BatchIndices>,
 }
 
 /// Arguments for graph coloring dispatch.
@@ -221,6 +248,48 @@ impl GpuColoring {
         num_colors
     }
 
+    /// Bucket-sorts the constraint ids by their color.
+    pub fn dispatch_build_color_buckets(
+        &self,
+        pass: &mut GpuPass,
+        args: ColorBucketsArgs<'_>,
+        color_buckets_stride: u32,
+        num_batches: u32,
+    ) -> Result<(), GpuBackendError> {
+        self.color_buckets_reset.call(
+            pass,
+            [color_buckets_stride, num_batches, 1],
+            args.color_bucket_counts,
+            args.batch_indices,
+        )?;
+        self.color_buckets_count.call(
+            pass,
+            args.contacts_len_indirect,
+            args.constraints_colors,
+            args.contacts_len,
+            args.color_bucket_counts,
+            args.batch_indices,
+        )?;
+        self.color_buckets_scan.call(
+            pass,
+            [1, num_batches, 1],
+            args.color_bucket_counts,
+            args.color_bucket_starts,
+            args.color_bucket_cursors,
+            args.batch_indices,
+        )?;
+        self.color_buckets_scatter.call(
+            pass,
+            args.contacts_len_indirect,
+            args.constraints_colors,
+            args.contacts_len,
+            args.color_bucket_cursors,
+            args.color_sorted_ids,
+            args.batch_indices,
+        )?;
+        Ok(())
+    }
+
     /// Runs a fixed number of iterations of the topo-gc coloring.
     pub fn dispatch_topo_gc_bounded<'a>(
         &self,
@@ -230,6 +299,28 @@ impl GpuColoring {
     ) -> Result<(), GpuBackendError> {
         // Reset coloring state.
         self.dispatch_reset_topo_gc(pass, &mut args)?;
+        self.dispatch_topo_gc_iterations(pass, args, max_colors)
+    }
+
+    /// Resets the topo-gc coloring state (all constraints uncolored). Public
+    /// so a seeding pass (e.g. warmstart color transfer) can run between the
+    /// reset and [`Self::dispatch_topo_gc_iterations`].
+    pub fn dispatch_topo_gc_reset<'a>(
+        &self,
+        pass: &mut GpuPass,
+        mut args: ColoringArgs<'a>,
+    ) -> Result<(), GpuBackendError> {
+        self.dispatch_reset_topo_gc(pass, &mut args)
+    }
+
+    /// Runs the bounded topo-gc step/fix-conflicts iterations, assuming the
+    /// coloring state was already reset (and possibly seeded).
+    pub fn dispatch_topo_gc_iterations<'a>(
+        &self,
+        pass: &mut GpuPass,
+        mut args: ColoringArgs<'a>,
+        max_colors: u32,
+    ) -> Result<(), GpuBackendError> {
         for _ in 0..max_colors {
             self.reset_completion_flag_topo_gc
                 .call(pass, 1u32, args.uncolored)?;