From dd48fc1b041b22319f13966c2d6a9a49a02a0cb9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Fri, 24 Jul 2026 13:29:09 +0200 Subject: [PATCH 01/39] fix: prevent out-of-bounds writes to pfm_pairs --- src_rbd_shaders/broad_phase/narrow_phase.rs | 34 +++++++++++++++------ 1 file changed, 25 insertions(+), 9 deletions(-) diff --git a/src_rbd_shaders/broad_phase/narrow_phase.rs b/src_rbd_shaders/broad_phase/narrow_phase.rs index 20a4ead..4d4f20d 100644 --- a/src_rbd_shaders/broad_phase/narrow_phase.rs +++ b/src_rbd_shaders/broad_phase/narrow_phase.rs @@ -171,8 +171,7 @@ pub fn gpu_narrow_phase_shape_shape( let target_contact_index = atomic_add_u32(contacts_len, 1) as usize; // NOTE: if we exceed the contacts allocation size, just skip - // the contact. It’s up to the caller to resize the buffer - // and re-run the narrow-phase. + // the contact. if target_contact_index < contacts_batch_capacity { let mat1 = collider_materials[pair.colliders.x as usize]; let mat2 = collider_materials[pair.colliders.y as usize]; @@ -283,7 +282,10 @@ pub fn gpu_narrow_phase_shape_shape_deferred( colliders: pair.colliders, }; let pfm_index = atomic_add_u32(pfm_pairs_len, 1); - pfm_pairs.write(pfm_index as usize, pfm_pair); + // NOTE: if we exceed capacity, just skip the pair. + if (pfm_index as usize) < contacts_batch_capacity { + pfm_pairs.write(pfm_index as usize, pfm_pair); + } // The actual calculations are deferred to another kernel. continue; @@ -302,6 +304,7 @@ pub fn gpu_narrow_phase_shape_shape_deferred( pair.colliders, &mut pfm_pairs, pfm_pairs_len, + contacts_batch_capacity, vertices, indices, ); @@ -319,6 +322,7 @@ pub fn gpu_narrow_phase_shape_shape_deferred( UVec2::new(pair.colliders.y, pair.colliders.x), &mut pfm_pairs, pfm_pairs_len, + contacts_batch_capacity, vertices, indices, ); @@ -337,6 +341,7 @@ pub fn gpu_narrow_phase_shape_shape_deferred( pair.colliders, &mut pfm_pairs, pfm_pairs_len, + contacts_batch_capacity, vertices, indices, ); @@ -354,6 +359,7 @@ pub fn gpu_narrow_phase_shape_shape_deferred( UVec2::new(pair.colliders.y, pair.colliders.x), &mut pfm_pairs, pfm_pairs_len, + contacts_batch_capacity, vertices, indices, ); @@ -370,6 +376,7 @@ fn trimesh_convex( colliders: UVec2, pfm_pairs: &mut SliceMut, pfm_pairs_len: &mut u32, + pfm_pairs_capacity: usize, vertices: &[PaddedVector], indices: &[u32], ) { @@ -413,7 +420,10 @@ fn trimesh_convex( colliders, }; let pfm_index = atomic_add_u32(pfm_pairs_len, 1); - pfm_pairs.write(pfm_index as usize, pfm_pair); + // Skip (don’t write) on overflow; the caller resizes and re-runs. + if (pfm_index as usize) < pfm_pairs_capacity { + pfm_pairs.write(pfm_index as usize, pfm_pair); + } // Continue traversal. curr = idx.exit_index; @@ -436,6 +446,7 @@ fn polyline_convex( colliders: UVec2, pfm_pairs: &mut SliceMut, pfm_pairs_len: &mut u32, + pfm_pairs_capacity: usize, vertices: &[PaddedVector], indices: &[u32], ) { @@ -482,7 +493,10 @@ fn polyline_convex( colliders, }; let pfm_index = atomic_add_u32(pfm_pairs_len, 1); - pfm_pairs.write(pfm_index as usize, pfm_pair); + // Skip (don’t write) on overflow; the caller resizes and re-runs. + if (pfm_index as usize) < pfm_pairs_capacity { + pfm_pairs.write(pfm_index as usize, pfm_pair); + } // Continue traversal. curr = idx.exit_index; @@ -564,7 +578,11 @@ pub fn gpu_narrow_phase_pfm_pfm( let collider_materials = batch_ids.coll_batch(batch_id, collider_materials); let pfm_pairs = batch_ids.contact_batch(batch_id, pfm_pairs); let contacts_len = contacts_len.at_mut(batch_id as usize); - let pfm_pairs_len = pfm_pairs_len.read(batch_id as usize); + // The producer counter can exceed the allocation on overflow (writes are + // skipped past capacity); clamp so we never read uninitialized slots. + let pfm_pairs_len = pfm_pairs_len + .read(batch_id as usize) + .min(contacts_batch_capacity as u32); for i in StepRng::new(invocation_id.x..pfm_pairs_len, num_threads) { let pair = pfm_pairs[i as usize]; @@ -592,9 +610,7 @@ pub fn gpu_narrow_phase_pfm_pfm( if manifold.len > 0 && manifold.points_a.at(0).dist < PREDICTION { let target_contact_index = atomic_add_u32(contacts_len, 1) as usize; - // NOTE: if we exceed the contacts allocation size, just skip - // the contact. It’s up to the caller to resize the buffer - // and re-run the narrow-phase. + // NOTE: if we exceed capacity, just skip the pair. if target_contact_index < contacts_batch_capacity { let mat1 = collider_materials[pair.colliders.x as usize]; let mat2 = collider_materials[pair.colliders.y as usize]; From caa9395dbd8aea0042736da6f342c977a2bb0463 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Fri, 24 Jul 2026 13:45:53 +0200 Subject: [PATCH 02/39] fix: convert LBVH pair-traversal while-loop to bounded for --- src_rbd_shaders/broad_phase/lbvh.rs | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/src_rbd_shaders/broad_phase/lbvh.rs b/src_rbd_shaders/broad_phase/lbvh.rs index 350a49d..fa9ff93 100644 --- a/src_rbd_shaders/broad_phase/lbvh.rs +++ b/src_rbd_shaders/broad_phase/lbvh.rs @@ -518,7 +518,13 @@ pub fn gpu_lbvh_find_collision_pairs( let mut stack_len = 1u32; stack.write(0, 0); - while stack_len != 0 { + // NOTE: we use a fixed-size for loop to avoid miscompilation issues of + // while loops on MacOs. Each tree node is pushed at most once per + // traversal, so `2 * num_bodies` (≥ node count) bounds the loop. + for _ in 0..2 * num_bodies { + if stack_len == 0 { + break; + } stack_len -= 1; let curr_id = stack.read(stack_len as usize); let node = tree.at(curr_id as usize); From 0e5edcd06aa1dca2843c13ffb8d2fd03652dc6cf Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Fri, 24 Jul 2026 14:19:28 +0200 Subject: [PATCH 03/39] perf: replace GPU color cursors with pre-built per-color uniforms --- src_rbd/dynamics/joint.rs | 44 ++++++++++--------- .../dynamics/multibody/loop_closing_joints.rs | 1 - .../multibody/multibody_from_rapier.rs | 2 - src_rbd/dynamics/multibody/multibody_set.rs | 12 +++-- .../dynamics/multibody/multibody_solver.rs | 40 +++++++---------- src_rbd/dynamics/solver.rs | 36 +++++++-------- src_rbd/pipeline/insertion_removal.rs | 1 + src_rbd/pipeline/rbd_state.rs | 11 +++++ src_rbd/pipeline/rbd_state_from_rapier.rs | 1 + src_rbd/pipeline/rbd_step.rs | 18 +++++++- src_rbd_shaders/dynamics/joint_constraint.rs | 33 -------------- src_rbd_shaders/dynamics/solver.rs | 33 -------------- 12 files changed, 93 insertions(+), 139 deletions(-) diff --git a/src_rbd/dynamics/joint.rs b/src_rbd/dynamics/joint.rs index 08f29d9..2a9efbe 100644 --- a/src_rbd/dynamics/joint.rs +++ b/src_rbd/dynamics/joint.rs @@ -5,7 +5,7 @@ use crate::math::Pose; use crate::shaders::dynamics::{ - GpuIncJointColor, GpuInitJointConstraints, GpuRemoveJointBias, GpuResetJointColor, + GpuInitJointConstraints, GpuRemoveJointBias, GpuSolveJointConstraints, GpuUpdateJointConstraints, ImpulseJoint, JointConstraint, JointConstraintBuilder, LocalMassProperties, RbdSimParams, Velocity, WorldMassProperties, }; @@ -83,8 +83,9 @@ pub struct GpuImpulseJointSet { /// Identical across batches by the equal-topology invariant. num_active_joints: u32, num_colors: u32, - max_color_group_len: u32, - curr_color: Tensor, + /// Host copy of the per-color prefix sums in `color_groups` (identical + /// across batches), used to size each per-color solve dispatch exactly. + color_groups_cpu: Vec, color_groups: Tensor, joints: Tensor, builders: Tensor, @@ -170,7 +171,6 @@ impl GpuImpulseJointSet { let max_joints = filtered_lens.iter().copied().max().unwrap_or(0); let mut global_num_colors = 0u32; - let mut global_max_color_group_len = 0u32; // Per-environment sorted joints and color groups. let mut per_env_sorted_joints: Vec> = Vec::new(); @@ -255,8 +255,6 @@ impl GpuImpulseJointSet { color_groups[*color as usize] += 1; } - let env_max_color_group_len = color_groups.iter().copied().max().unwrap_or_default(); - // Prefix sum. for i in 0..color_groups.len().saturating_sub(1) { color_groups[i + 1] += color_groups[i]; @@ -273,7 +271,6 @@ impl GpuImpulseJointSet { } global_num_colors = global_num_colors.max(env_num_colors); - global_max_color_group_len = global_max_color_group_len.max(env_max_color_group_len); per_env_sorted_joints.push(sorted_gpu_joints); per_env_color_groups.push(color_groups); @@ -329,8 +326,7 @@ impl GpuImpulseJointSet { // invariant, so the per-batch active count is any env's count. num_active_joints: filtered_lens.first().copied().unwrap_or(0), num_colors: global_num_colors, - max_color_group_len: global_max_color_group_len, - curr_color: Tensor::scalar(backend, 0u32, usage | BufferUsages::UNIFORM).unwrap(), + color_groups_cpu: all_color_groups.clone(), color_groups: Tensor::vector(backend, &all_color_groups, usage).unwrap(), joints: Tensor::vector(backend, &all_joints, usage).unwrap(), builders: Tensor::matrix_uninit(backend, num_batches, max_joints, usage).unwrap(), @@ -372,8 +368,6 @@ pub struct GpuJointSolver { init_joint_constraints: GpuInitJointConstraints, /// Updates joint constraints each substep. update_joint_constraints: GpuUpdateJointConstraints, - reset_joint_color: GpuResetJointColor, - inc_joint_color: GpuIncJointColor, /// Solves joint constraints. solve_joint_constraints: GpuSolveJointConstraints, /// Removes bias from joint constraints. @@ -394,6 +388,9 @@ pub struct JointSolverArgs<'a> { pub local_mprops: &'a Tensor, /// Shared per-batch capacity / section-offset uniform. pub batch_indices: &'a Tensor, + /// Per-color-index uniform tensors (`color_uniforms[c] == c`), + /// shared with the contact solver. + pub color_uniforms: &'a [Tensor], } impl GpuJointSolver { @@ -464,23 +461,28 @@ impl GpuJointSolver { )?; } - self.reset_joint_color - .call(pass, 1u32, &mut args.joints.curr_color)?; - - for _ in 0..args.joints.num_colors { - // TODO PERF: figure out a way to dispatch a number of threads that fits - // more tightly the size of the current color. + // One dispatch per color, sized exactly to that color's group (the + // prefix sums are known on the host). The color index is bound as a + // tiny pre-built uniform instead of a GPU-incremented cursor. + for c in 0..args.joints.num_colors as usize { + let start = if c > 0 { + args.joints.color_groups_cpu[c - 1] + } else { + 0 + }; + let group_len = args.joints.color_groups_cpu[c] - start; + if group_len == 0 { + continue; + } self.solve_joint_constraints.call( pass, - [args.joints.max_color_group_len, args.num_batches, 1], + [group_len, args.num_batches, 1], &mut args.joints.constraints, solver_vels, &args.joints.color_groups, - &args.joints.curr_color, + &args.color_uniforms[c], args.batch_indices, )?; - self.inc_joint_color - .call(pass, 1u32, &mut args.joints.curr_color)?; } Ok(()) diff --git a/src_rbd/dynamics/multibody/loop_closing_joints.rs b/src_rbd/dynamics/multibody/loop_closing_joints.rs index e9ba6c8..92d012a 100644 --- a/src_rbd/dynamics/multibody/loop_closing_joints.rs +++ b/src_rbd/dynamics/multibody/loop_closing_joints.rs @@ -303,7 +303,6 @@ impl GpuMultibodySet { } self.mb_imp_joint_color_groups = Tensor::vector(backend, &all_color_groups, storage).unwrap(); - self.mb_imp_joint_curr_color = Tensor::scalar(backend, 0u32, usage_u).unwrap(); self.mb_imp_joint_num_colors = global_num_colors; self.mb_imp_joint_max_color_group_len = global_max_color_group_len; } diff --git a/src_rbd/dynamics/multibody/multibody_from_rapier.rs b/src_rbd/dynamics/multibody/multibody_from_rapier.rs index cb390b3..7398244 100644 --- a/src_rbd/dynamics/multibody/multibody_from_rapier.rs +++ b/src_rbd/dynamics/multibody/multibody_from_rapier.rs @@ -366,7 +366,6 @@ impl GpuMultibodySet { } let storage = BufferUsages::STORAGE | BufferUsages::COPY_DST; - let usage_u = storage | BufferUsages::UNIFORM; Self { num_batches, @@ -506,7 +505,6 @@ impl GpuMultibodySet { storage, ) .unwrap(), - mb_imp_joint_curr_color: Tensor::scalar(backend, 0u32, usage_u).unwrap(), mb_imp_joint_num_colors: 0, mb_imp_joint_max_color_group_len: 0, joint_constraints_per_batch: cons_cap, diff --git a/src_rbd/dynamics/multibody/multibody_set.rs b/src_rbd/dynamics/multibody/multibody_set.rs index 89734fe..c6bfd40 100644 --- a/src_rbd/dynamics/multibody/multibody_set.rs +++ b/src_rbd/dynamics/multibody/multibody_set.rs @@ -112,10 +112,9 @@ pub struct GpuMultibodySet { /// Per-batch prefix-sum over the color-sorted `mb_imp_joint_builders`. /// Built at init time by `set_impulse_joints` (greedy graph coloring). pub(super) mb_imp_joint_color_groups: Tensor, - /// Scalar color cursor incremented by the host color loop. - pub(super) mb_imp_joint_curr_color: Tensor, - /// Number of colors (host color-loop trip count). CPU mirror. - pub(super) mb_imp_joint_num_colors: u32, + /// Number of colors (per-batch stride of `mb_imp_joint_color_groups`, + /// and the host color-loop trip count). CPU mirror. + pub(crate) mb_imp_joint_num_colors: u32, /// Largest color group across batches — the per-color dispatch width. pub(super) mb_imp_joint_max_color_group_len: u32, /// Per-batch capacities of the joint / contact constraint slabs (CPU-side @@ -153,6 +152,11 @@ impl GpuMultibodySet { self.multibodies_per_batch == 0 || self.links_per_batch == 0 } + /// Number of colors used by the colored multibody impulse-joint sweeps. + pub fn mb_imp_joint_num_colors(&self) -> u32 { + self.mb_imp_joint_num_colors + } + /// GPU buffer holding generalized velocities followed by per-DOF damping. /// The velocity section is `[0, dof_batch_capacity * num_batches)`; the /// damping section follows. Callers reading velocities should use only the diff --git a/src_rbd/dynamics/multibody/multibody_solver.rs b/src_rbd/dynamics/multibody/multibody_solver.rs index 20ef22f..a255a0a 100644 --- a/src_rbd/dynamics/multibody/multibody_solver.rs +++ b/src_rbd/dynamics/multibody/multibody_solver.rs @@ -4,14 +4,15 @@ use super::multibody_set::*; use crate::math::Pose; use crate::queries::GpuIndexedContact; use crate::shaders::dynamics::{ - GpuIncJointColor, GpuMbComputeDynamicsPre, GpuMbComputeDynamicsWithoutCoriolisPre, - GpuMbFinalizeContactConstraints, GpuMbFinalizeImpulseJointConstraints, GpuMbGravityAndLu, - GpuMbInitContactConstraints, GpuMbInitJointConstraints, GpuMbIntegrate, - GpuMbIntegrateVelocities, GpuMbRemoveContactConstraintBias, - GpuMbRemoveImpulseJointConstraintBias, GpuMbRemoveSolveJointNoBias, GpuMbResetContactWarmstart, - GpuMbSolveContactConstraints, GpuMbSolveImpulseJointConstraints, GpuMbSolveJointConstraints, - GpuMbUpdateImpulseJointConstraints, GpuMbWarmstartContactConstraints, GpuResetJointColor, - Velocity, WorldMassProperties, + GpuMbComputeDynamicsPre, + GpuMbComputeDynamicsWithoutCoriolisPre, + GpuMbFinalizeContactConstraints, GpuMbGravityAndLu, GpuMbInitContactConstraints, + GpuMbInitJointConstraints, GpuMbIntegrate, GpuMbIntegrateVelocities, + GpuMbRemoveContactConstraintBias, GpuMbRemoveImpulseJointConstraintBias, + GpuMbResetContactWarmstart, GpuMbWarmstartContactConstraints, + GpuMbRemoveSolveJointNoBias, GpuMbSolveContactConstraints, GpuMbSolveImpulseJointConstraints, + GpuMbFinalizeImpulseJointConstraints, GpuMbSolveJointConstraints, + GpuMbUpdateImpulseJointConstraints, Velocity, WorldMassProperties, }; use crate::shaders::utils::BatchIndices; use khal::Shader; @@ -41,10 +42,6 @@ pub struct GpuMultibodySolver { /// split out so the build pass fits 8 storage buffers. finalize_impulse_joint_constraints: GpuMbFinalizeImpulseJointConstraints, solve_impulse_joint_constraints: GpuMbSolveImpulseJointConstraints, - /// Color cursor reset / increment for the colored impulse-joint solve - /// loop. Reuses the free-body joint color kernels (generic `&mut u32`). - reset_imp_joint_color: GpuResetJointColor, - inc_imp_joint_color: GpuIncJointColor, remove_impulse_joint_constraint_bias: GpuMbRemoveImpulseJointConstraintBias, integrate_velocities: GpuMbIntegrateVelocities, integrate: GpuMbIntegrate, @@ -71,6 +68,9 @@ pub struct MultibodySolverArgs<'a> { /// Shared `BatchIndices` uniform — per-batch caps and packed-section /// offsets read by every multibody kernel. Owned by `RbdState`. pub batch_indices: &'a Tensor, + /// Per-color-index uniform tensors (`color_uniforms[c]` holds `c`), + /// shared with the contact/joint solvers. + pub color_uniforms: &'a [Tensor], } impl GpuMultibodySolver { @@ -357,9 +357,7 @@ impl GpuMultibodySolver { // Colored PGS sweep WITH bias: one dispatch per color, each // color's joints solved race-free in parallel (graph coloring // done at init in `set_impulse_joints`). - self.reset_imp_joint_color - .call(pass, 1u32, &mut mb.mb_imp_joint_curr_color)?; - for _ in 0..mb.mb_imp_joint_num_colors { + for c in 0..mb.mb_imp_joint_num_colors as usize { self.solve_impulse_joint_constraints.call( pass, // One workgroup (MB_LU_LANES threads) per joint; thread @@ -374,13 +372,11 @@ impl GpuMultibodySolver { &mb.mb_imp_joint_jacobians, &mb.mb_imp_joint_color_groups, args.batch_indices, - &mb.mb_imp_joint_curr_color, + &args.color_uniforms[c], &mb.multibody_info, &mut mb.dof_state, args.solver_vels, )?; - self.inc_imp_joint_color - .call(pass, 1u32, &mut mb.mb_imp_joint_curr_color)?; } } @@ -488,9 +484,7 @@ impl GpuMultibodySolver { if mb.mb_imp_joints_per_batch > 0 { // Final stabilization sweep WITHOUT bias — colored, one // dispatch per color (see the with-bias sweep above). - self.reset_imp_joint_color - .call(pass, 1u32, &mut mb.mb_imp_joint_curr_color)?; - for _ in 0..mb.mb_imp_joint_num_colors { + for c in 0..mb.mb_imp_joint_num_colors as usize { self.solve_impulse_joint_constraints.call( pass, // One workgroup (MB_LU_LANES threads) per joint; thread @@ -505,13 +499,11 @@ impl GpuMultibodySolver { &mb.mb_imp_joint_jacobians, &mb.mb_imp_joint_color_groups, args.batch_indices, - &mb.mb_imp_joint_curr_color, + &args.color_uniforms[c], &mb.multibody_info, &mut mb.dof_state, args.solver_vels, )?; - self.inc_imp_joint_color - .call(pass, 1u32, &mut mb.mb_imp_joint_curr_color)?; } } diff --git a/src_rbd/dynamics/solver.rs b/src_rbd/dynamics/solver.rs index 3b89840..81f3208 100644 --- a/src_rbd/dynamics/solver.rs +++ b/src_rbd/dynamics/solver.rs @@ -12,9 +12,10 @@ use crate::queries::GpuIndexedContact; use crate::shaders::dynamics::{ GpuApplySolverVelsInc, GpuInitSolverBodies, GpuInitSolverVelsInc, GpuIntegrateLinearized, GpuRemoveCfmAndBiasKernel, GpuSolverCleanup, GpuSolverCountConstraints, GpuSolverFinalize, - GpuSolverIncColor, GpuSolverInitConstraints, GpuSolverResetColor, GpuSolverSortConstraints, - GpuSolverUpdateConstraints, GpuStepGaussSeidel, GpuWarmstart, LocalMassProperties, - RbdSimParams, TwoBodyConstraint, TwoBodyConstraintBuilder, Velocity, WorldMassProperties, + GpuSolverInitConstraints, GpuSolverSortConstraints, + GpuSolverUpdateConstraints, GpuStepGaussSeidel, GpuWarmstart, GpuWarmstartWithoutColors, + LocalMassProperties, RbdSimParams, TwoBodyConstraint, TwoBodyConstraintBuilder, Velocity, + WorldMassProperties, }; use crate::utils::{GpuPrefixSum, PrefixSumWorkspace}; use khal::Shader; @@ -25,8 +26,6 @@ use vortx::tensor::Tensor; #[derive(Shader)] pub struct GpuSolver { sort_constraints: GpuSolverSortConstraints, - reset_color: GpuSolverResetColor, - inc_color: GpuSolverIncColor, /// Initializes constraints from contact manifolds. init_constraints: GpuSolverInitConstraints, /// Companion counting pass to `init_constraints` (split out to keep each @@ -115,8 +114,11 @@ pub struct SolverArgs<'a> { pub body_constraint_ids: &'a mut Tensor, /// Color assigned to each constraint by graph coloring. pub constraints_colors: &'a Tensor, - /// Current color being processed. - pub curr_color: &'a mut Tensor, + /// Per-color-index uniform tensors: `color_uniforms[c]` holds the constant + /// `c`. Bound (instead of a GPU-incremented cursor) by each color-sweep + /// dispatch, which removes the 1-thread `reset_color`/`inc_color` + /// dispatches (and their barriers) from every color loop. + pub color_uniforms: &'a [Tensor], /// Prefix sum shader for building constraint ranges. pub prefix_sum: &'a GpuPrefixSum, /// Number of solver iterations (max across all environments). @@ -277,6 +279,7 @@ impl GpuSolver { contacts_len: args.contacts_len, solver_vels: &mut *args.solver_vels, batch_indices: args.batch_indices, + color_uniforms: args.color_uniforms, }; solver.$method(pass, state, &mut mb_args $(, $extra)*)?; } @@ -316,8 +319,8 @@ impl GpuSolver { 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 { + // NOTE: contact colors start at 1 (0 = unassigned). + for c in 1..=args.num_colors { self.warmstart.call( pass, args.contacts_len_indirect, @@ -325,10 +328,9 @@ impl GpuSolver { args.solver_vels, args.constraints_colors, args.contacts_len, - args.curr_color, + &args.color_uniforms[c as usize], args.batch_indices, )?; - self.inc_color.call(pass, 1u32, args.curr_color)? } /* @@ -336,8 +338,7 @@ impl GpuSolver { */ 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 { + for c in 1..=args.num_colors { self.step_gauss_seidel.call( pass, args.contacts_len_indirect, @@ -345,10 +346,9 @@ impl GpuSolver { args.solver_vels, args.constraints_colors, args.contacts_len, - args.curr_color, + &args.color_uniforms[c as usize], args.batch_indices, )?; - self.inc_color.call(pass, 1u32, args.curr_color)? } /* @@ -376,8 +376,7 @@ impl GpuSolver { args.contacts_len, args.batch_indices, )?; - self.reset_color.call(pass, 1u32, args.curr_color)?; - for _ in 0..args.num_colors { + for c in 1..=args.num_colors { self.step_gauss_seidel.call( pass, args.contacts_len_indirect, @@ -385,10 +384,9 @@ impl GpuSolver { args.solver_vels, args.constraints_colors, args.contacts_len, - args.curr_color, + &args.color_uniforms[c as usize], args.batch_indices, )?; - self.inc_color.call(pass, 1u32, args.curr_color)? } } diff --git a/src_rbd/pipeline/insertion_removal.rs b/src_rbd/pipeline/insertion_removal.rs index b5e1332..fe5ec94 100644 --- a/src_rbd/pipeline/insertion_removal.rs +++ b/src_rbd/pipeline/insertion_removal.rs @@ -294,6 +294,7 @@ impl RbdState { .unwrap(), old_body_constraint_ids, new_body_constraint_ids, + color_uniforms: Vec::new(), prefix_sum_workspace: PrefixSumWorkspace::default(), lbvh: LbvhState::with_usages(backend, lbvh_usages), max_colors: capacities.solver_colors, diff --git a/src_rbd/pipeline/rbd_state.rs b/src_rbd/pipeline/rbd_state.rs index 73e0c92..2df3e9b 100644 --- a/src_rbd/pipeline/rbd_state.rs +++ b/src_rbd/pipeline/rbd_state.rs @@ -199,6 +199,9 @@ pub struct RbdState { pub(super) colored: Tensor, pub(super) constraints_rands: Tensor, pub(super) curr_color: Tensor, + /// Pre-built per-color-index uniforms: `color_uniforms[c] == c`. + /// [`Self::ensure_color_uniforms`]. + pub(super) color_uniforms: Vec>, pub(super) uncolored: Tensor, pub(super) uncolored_staging: Tensor, pub(super) lbvh: LbvhState, @@ -266,6 +269,14 @@ impl RbdState { self.max_colors = max_colors.max(1); } + /// Grows [`Self::color_uniforms`] so indices `0..n` are available. + pub(super) fn ensure_color_uniforms(&mut self, backend: &GpuBackend, n: u32) { + for c in self.color_uniforms.len() as u32..n { + self.color_uniforms + .push(Tensor::scalar(backend, c, BufferUsages::UNIFORM).unwrap()); + } + } + /// Returns the configured max color count. pub fn max_colors(&self) -> u32 { self.max_colors diff --git a/src_rbd/pipeline/rbd_state_from_rapier.rs b/src_rbd/pipeline/rbd_state_from_rapier.rs index 92bdd1e..20364c5 100644 --- a/src_rbd/pipeline/rbd_state_from_rapier.rs +++ b/src_rbd/pipeline/rbd_state_from_rapier.rs @@ -771,6 +771,7 @@ impl RbdState { | BufferUsages::COPY_SRC, ) .unwrap(), + color_uniforms: Vec::new(), uncolored: Tensor::scalar( backend, 0, diff --git a/src_rbd/pipeline/rbd_step.rs b/src_rbd/pipeline/rbd_step.rs index 64817d2..ff78c78 100644 --- a/src_rbd/pipeline/rbd_step.rs +++ b/src_rbd/pipeline/rbd_step.rs @@ -68,6 +68,18 @@ impl RbdPipeline { ) -> Result { let mut stats = RunStats::default(); + // Make sure the color index uniforms are up-to-date. + // This is the maximum over the colors needed for contacts, joints, and multibodies. + { + let mut needed = state.max_colors + 2; + needed = needed.max(state.joints.num_colors() + 1); + #[cfg(feature = "dim3")] + { + needed = needed.max(state.multibodies.mb_imp_joint_num_colors() + 1); + } + state.ensure_color_uniforms(backend, needed); + } + // Phase 0: Multibody once-per-visible-step setup (3D only for now). #[cfg(feature = "dim3")] { @@ -83,6 +95,7 @@ impl RbdPipeline { contacts_len: &state.contacts_len, solver_vels: &mut state.solver_vels, batch_indices: &state.batch_indices, + color_uniforms: &state.color_uniforms, }; self.multibody_solver .init_step(&mut pass, &mut state.multibodies, &mut args)?; @@ -231,7 +244,7 @@ impl RbdPipeline { body_constraint_counts: &mut state.new_constraints_counts, body_constraint_ids: &mut state.new_body_constraint_ids, constraints_colors: &state.constraints_colors, - curr_color: &mut state.curr_color, + color_uniforms: &state.color_uniforms, prefix_sum: &self.prefix_sum, num_colors: 0, num_batches: state.num_batches, @@ -312,7 +325,7 @@ impl RbdPipeline { body_constraint_counts: &mut state.new_constraints_counts, body_constraint_ids: &mut state.new_body_constraint_ids, constraints_colors: &state.constraints_colors, - curr_color: &mut state.curr_color, + color_uniforms: &state.color_uniforms, prefix_sum: &self.prefix_sum, num_colors, num_batches: state.num_batches, @@ -330,6 +343,7 @@ impl RbdPipeline { local_mprops: &state.local_mprops, joints: &mut state.joints, batch_indices: &state.batch_indices, + color_uniforms: &state.color_uniforms, }; { diff --git a/src_rbd_shaders/dynamics/joint_constraint.rs b/src_rbd_shaders/dynamics/joint_constraint.rs index 63774ff..b00f355 100644 --- a/src_rbd_shaders/dynamics/joint_constraint.rs +++ b/src_rbd_shaders/dynamics/joint_constraint.rs @@ -160,39 +160,6 @@ pub struct JointConstraintElement { pub impulse_bounds: Vec2, } -/// Resets the joint color to 0. -#[spirv_bindgen] -#[spirv(compute(threads(1)))] -pub fn gpu_reset_joint_color( - #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] curr_color: &mut u32, -) { - // NOTE: this `for` loop is silly. It doesn’t do anything - // more than a `*curr_color = 0` in a convoluted - // way because otherwise rustgpu apparently does not generate - // the spirv for this kernel (seems to happen if the kernel is - // too trivial. - for k in 0..1 { - // NOTE: for joints, our first colors start at 0. - *curr_color = k; - } -} - -/// Increments the joint color. -#[spirv_bindgen] -#[spirv(compute(threads(1)))] -pub fn gpu_inc_joint_color( - #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] curr_color: &mut u32, -) { - // NOTE: this `for` loop is silly. It doesn’t do anything - // more than a `*curr_color += 1` in a convoluted - // way because otherwise rustgpu apparently does not generate - // the spirv for this kernel (seems to happen if the kernel is - // too trivial. - for k in 0..1 { - *curr_color += 1 + k; - } -} - /// Initializes joint constraint builders and constraints. #[spirv_bindgen] #[spirv(compute(threads(64)))] diff --git a/src_rbd_shaders/dynamics/solver.rs b/src_rbd_shaders/dynamics/solver.rs index 6dfbe62..a94cc78 100644 --- a/src_rbd_shaders/dynamics/solver.rs +++ b/src_rbd_shaders/dynamics/solver.rs @@ -18,39 +18,6 @@ use crate::utils::{BatchIndices, Slice, SliceMut}; const WORKGROUP_SIZE: u32 = 64; -/// Resets the current color to 1 (for graph coloring). -#[spirv_bindgen] -#[spirv(compute(threads(1)))] -pub fn gpu_solver_reset_color( - #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] curr_color: &mut u32, -) { - // NOTE: this `for` loop is silly. It doesn't do anything - // more than a `*curr_color = 1` in a convoluted - // way because otherwise rustgpu apparently does not generate - // the spirv for this kernel (seems to happen if the kernel is - // too trivial. - for k in 0..1 { - // NOTE: our first colors start at 1 instead of 0. - *curr_color = 1 + k; - } -} - -/// Increments the current color. -#[spirv_bindgen] -#[spirv(compute(threads(1)))] -pub fn gpu_solver_inc_color( - #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] curr_color: &mut u32, -) { - // NOTE: this `for` loop is silly. It doesn't do anything - // more than a `*curr_color += 1` in a convoluted - // way because otherwise rustgpu apparently does not generate - // the spirv for this kernel (seems to happen if the kernel is - // too trivial. - for k in 0..1 { - *curr_color += 1 + k; - } -} - /// Initializes constraints from contact manifolds. /// /// Split into two passes to stay within WebGPU's 8-storage-buffer per-stage From 8eabde53315257640e9bba918a088c6a4f7507f3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Fri, 24 Jul 2026 14:49:45 +0200 Subject: [PATCH 04/39] perf: bucket-sort contact constraints by color so each color only iterates its own constraints --- src_rbd/dynamics/coloring.rs | 69 +++++++++++++ src_rbd/dynamics/mod.rs | 2 +- src_rbd/dynamics/solver.rs | 24 ++--- src_rbd/pipeline/insertion_removal.rs | 14 +++ src_rbd/pipeline/rbd_state.rs | 11 ++ src_rbd/pipeline/rbd_state_from_rapier.rs | 18 ++++ src_rbd/pipeline/rbd_step.rs | 40 +++++++- src_rbd_shaders/dynamics/color_buckets.rs | 120 ++++++++++++++++++++++ src_rbd_shaders/dynamics/mod.rs | 2 + src_rbd_shaders/dynamics/solver.rs | 67 ++++++------ src_rbd_shaders/utils/indices.rs | 4 + 11 files changed, 325 insertions(+), 46 deletions(-) create mode 100644 src_rbd_shaders/dynamics/color_buckets.rs diff --git a/src_rbd/dynamics/coloring.rs b/src_rbd/dynamics/coloring.rs index 84b2cb5..88f2e33 100644 --- a/src_rbd/dynamics/coloring.rs +++ b/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, + /// Number of contacts per batch. + pub contacts_len: &'a Tensor, + /// Per-batch per-color counts (stride `solver_color_buckets_stride`). + pub color_bucket_counts: &'a mut Tensor, + /// Per-batch per-color exclusive prefix sums. + pub color_bucket_starts: &'a mut Tensor, + /// Scatter cursors (seeded from the starts). + pub color_bucket_cursors: &'a mut Tensor, + /// Constraint ids bucket-sorted by color (contacts layout). + pub color_sorted_ids: &'a mut Tensor, + /// Shared per-batch capacity / section-offset uniform. + pub batch_indices: &'a Tensor, } /// 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, diff --git a/src_rbd/dynamics/mod.rs b/src_rbd/dynamics/mod.rs index b534ea5..023e512 100644 --- a/src_rbd/dynamics/mod.rs +++ b/src_rbd/dynamics/mod.rs @@ -1,7 +1,7 @@ //! Rigid-body dynamics: forces, velocities, constraints, and solvers. pub use crate::shaders::dynamics::RbdSimParams; -pub use coloring::{ColoringArgs, GpuColoring}; +pub use coloring::{ColorBucketsArgs, ColoringArgs, GpuColoring}; pub use joint::{GpuImpulseJointSet, GpuJointSolver, JointSolverArgs}; pub use mprops_update::{GpuMpropsUpdate, GpuSyncColliderPosesShader}; #[cfg(feature = "dim3")] diff --git a/src_rbd/dynamics/solver.rs b/src_rbd/dynamics/solver.rs index 81f3208..cfdcd06 100644 --- a/src_rbd/dynamics/solver.rs +++ b/src_rbd/dynamics/solver.rs @@ -112,12 +112,12 @@ pub struct SolverArgs<'a> { /// All constraints of all the bodies part of the same multibody are in the same list associated /// to the multibody’s root. pub body_constraint_ids: &'a mut Tensor, - /// Color assigned to each constraint by graph coloring. - pub constraints_colors: &'a Tensor, - /// Per-color-index uniform tensors: `color_uniforms[c]` holds the constant - /// `c`. Bound (instead of a GPU-incremented cursor) by each color-sweep - /// dispatch, which removes the 1-thread `reset_color`/`inc_color` - /// dispatches (and their barriers) from every color loop. + /// Per-batch per-color exclusive prefix sums over the color-bucketed + /// constraint ids (stride `BatchIndices::solver_color_buckets_stride`). + pub color_bucket_starts: &'a Tensor, + /// Constraint ids bucket-sorted by color (contacts layout). + pub color_sorted_ids: &'a Tensor, + /// Per-color-index uniform tensors: `color_uniforms[c] == c`. pub color_uniforms: &'a [Tensor], /// Prefix sum shader for building constraint ranges. pub prefix_sum: &'a GpuPrefixSum, @@ -326,8 +326,8 @@ impl GpuSolver { args.contacts_len_indirect, args.constraints, args.solver_vels, - args.constraints_colors, - args.contacts_len, + args.color_bucket_starts, + args.color_sorted_ids, &args.color_uniforms[c as usize], args.batch_indices, )?; @@ -344,8 +344,8 @@ impl GpuSolver { args.contacts_len_indirect, args.constraints, args.solver_vels, - args.constraints_colors, - args.contacts_len, + args.color_bucket_starts, + args.color_sorted_ids, &args.color_uniforms[c as usize], args.batch_indices, )?; @@ -382,8 +382,8 @@ impl GpuSolver { args.contacts_len_indirect, args.constraints, args.solver_vels, - args.constraints_colors, - args.contacts_len, + args.color_bucket_starts, + args.color_sorted_ids, &args.color_uniforms[c as usize], args.batch_indices, )?; diff --git a/src_rbd/pipeline/insertion_removal.rs b/src_rbd/pipeline/insertion_removal.rs index fe5ec94..0cc8438 100644 --- a/src_rbd/pipeline/insertion_removal.rs +++ b/src_rbd/pipeline/insertion_removal.rs @@ -182,6 +182,15 @@ impl RbdState { Tensor::vector_uninit(backend, collisions_capacity * num_batches, storage).unwrap(); let constraints_rands = Tensor::vector_uninit(backend, collisions_capacity * num_batches, storage).unwrap(); + let color_buckets_stride = capacities.solver_colors + 3; + let color_bucket_counts = + Tensor::vector_uninit(backend, color_buckets_stride * num_batches, storage).unwrap(); + let color_bucket_starts = + Tensor::vector_uninit(backend, color_buckets_stride * num_batches, storage).unwrap(); + let color_bucket_cursors = + Tensor::vector_uninit(backend, color_buckets_stride * num_batches, storage).unwrap(); + let color_sorted_ids = + Tensor::vector_uninit(backend, collisions_capacity * num_batches, storage).unwrap(); let old_constraints_counts = Tensor::vector_uninit(backend, num_colliders_per_batch * num_batches, storage).unwrap(); let new_constraints_counts = @@ -209,6 +218,7 @@ impl RbdState { contacts_batch_capacity: contacts_per_batch_cpu, impulse_joints_batch_capacity: joints.joints_per_batch(), impulse_joints_len: joints.num_active_joints(), + solver_color_buckets_stride: color_buckets_stride, ..Default::default() }; #[cfg(feature = "dim3")] @@ -271,6 +281,10 @@ impl RbdState { constraints_colors, colored, constraints_rands, + color_bucket_counts, + color_bucket_starts, + color_bucket_cursors, + color_sorted_ids, curr_color: Tensor::scalar( backend, 0u32, diff --git a/src_rbd/pipeline/rbd_state.rs b/src_rbd/pipeline/rbd_state.rs index 2df3e9b..920f008 100644 --- a/src_rbd/pipeline/rbd_state.rs +++ b/src_rbd/pipeline/rbd_state.rs @@ -198,6 +198,16 @@ pub struct RbdState { pub(super) constraints_colors: Tensor, pub(super) colored: Tensor, pub(super) constraints_rands: Tensor, + /// Per-batch per-color constraint counts (stride `max_colors + 3`), see + /// the `gpu_color_buckets_*` kernels. + pub(super) color_bucket_counts: Tensor, + /// Per-batch per-color exclusive prefix sums over the counts: color `c` + /// owns `color_sorted_ids[starts[c]..starts[c + 1]]`. + pub(super) color_bucket_starts: Tensor, + /// Scatter cursors (seeded from the starts each step). + pub(super) color_bucket_cursors: Tensor, + /// Constraint indices bucket-sorted by color (contacts layout). + pub(super) color_sorted_ids: Tensor, pub(super) curr_color: Tensor, /// Pre-built per-color-index uniforms: `color_uniforms[c] == c`. /// [`Self::ensure_color_uniforms`]. @@ -245,6 +255,7 @@ impl RbdState { contacts_batch_capacity: self.contacts_per_batch_cpu, impulse_joints_batch_capacity: self.joints.joints_per_batch(), impulse_joints_len: self.joints.num_active_joints(), + solver_color_buckets_stride: self.max_colors + 3, ..Default::default() }; #[cfg(feature = "dim3")] diff --git a/src_rbd/pipeline/rbd_state_from_rapier.rs b/src_rbd/pipeline/rbd_state_from_rapier.rs index 20364c5..64aa081 100644 --- a/src_rbd/pipeline/rbd_state_from_rapier.rs +++ b/src_rbd/pipeline/rbd_state_from_rapier.rs @@ -637,6 +637,19 @@ impl RbdState { storage, ) .unwrap(); + let color_buckets_stride = capacities.solver_colors + 3; + let color_bucket_counts = + Tensor::vector_uninit(backend, color_buckets_stride * num_batches, storage).unwrap(); + let color_bucket_starts = + Tensor::vector_uninit(backend, color_buckets_stride * num_batches, storage).unwrap(); + let color_bucket_cursors = + Tensor::vector_uninit(backend, color_buckets_stride * num_batches, storage).unwrap(); + let color_sorted_ids = Tensor::vector_uninit( + backend, + capacities.collisions_capacity * num_batches, + storage, + ) + .unwrap(); let old_constraints_counts = Tensor::vector_uninit( backend, num_colliders_per_batch as u32 * num_batches, @@ -679,6 +692,7 @@ impl RbdState { contacts_batch_capacity: contacts_per_batch_cpu, impulse_joints_batch_capacity: joints.joints_per_batch(), impulse_joints_len: joints.num_active_joints(), + solver_color_buckets_stride: color_buckets_stride, ..Default::default() }; #[cfg(feature = "dim3")] @@ -762,6 +776,10 @@ impl RbdState { constraints_colors, colored, constraints_rands, + color_bucket_counts, + color_bucket_starts, + color_bucket_cursors, + color_sorted_ids, curr_color: Tensor::scalar( backend, 0u32, diff --git a/src_rbd/pipeline/rbd_step.rs b/src_rbd/pipeline/rbd_step.rs index ff78c78..a3d3c1a 100644 --- a/src_rbd/pipeline/rbd_step.rs +++ b/src_rbd/pipeline/rbd_step.rs @@ -243,7 +243,8 @@ impl RbdPipeline { local_mprops: &state.local_mprops, body_constraint_counts: &mut state.new_constraints_counts, body_constraint_ids: &mut state.new_body_constraint_ids, - constraints_colors: &state.constraints_colors, + color_bucket_starts: &state.color_bucket_starts, + color_sorted_ids: &state.color_sorted_ids, color_uniforms: &state.color_uniforms, prefix_sum: &self.prefix_sum, num_colors: 0, @@ -294,6 +295,25 @@ impl RbdPipeline { self.coloring .dispatch_topo_gc_bounded(&mut pass, coloring_args, state.max_colors)?; + // Bucket-sort the constraint ids by color so each colored solver + // sweep only touches its own constraints. + let bucket_args = crate::dynamics::ColorBucketsArgs { + contacts_len_indirect: &state.contacts_indirect, + constraints_colors: &state.constraints_colors, + contacts_len: &state.contacts_len, + color_bucket_counts: &mut state.color_bucket_counts, + color_bucket_starts: &mut state.color_bucket_starts, + color_bucket_cursors: &mut state.color_bucket_cursors, + color_sorted_ids: &mut state.color_sorted_ids, + batch_indices: &state.batch_indices, + }; + self.coloring.dispatch_build_color_buckets( + &mut pass, + bucket_args, + state.max_colors + 3, + state.num_batches, + )?; + // `+1` because solver iterates 1..=max_colors (color 0 is unassigned). let num_colors = state.max_colors + 1; stats.num_colors = num_colors; @@ -324,7 +344,8 @@ impl RbdPipeline { local_mprops: &state.local_mprops, body_constraint_counts: &mut state.new_constraints_counts, body_constraint_ids: &mut state.new_body_constraint_ids, - constraints_colors: &state.constraints_colors, + color_bucket_starts: &state.color_bucket_starts, + color_sorted_ids: &state.color_sorted_ids, color_uniforms: &state.color_uniforms, prefix_sum: &self.prefix_sum, num_colors, @@ -421,6 +442,19 @@ impl RbdPipeline { && coloring_converged == 0 { state.max_colors += 5; + + // The color-bucket buffers are strided by `max_colors + 3`: + // regrow them and update the stride in `BatchIndices`. + let storage: BufferUsages = BufferUsages::STORAGE | BufferUsages::COPY_SRC; + let stride = state.max_colors + 3; + let nb = state.num_batches; + state.color_bucket_counts = + Tensor::vector_uninit(backend, stride * nb, storage)?; + state.color_bucket_starts = + Tensor::vector_uninit(backend, stride * nb, storage)?; + state.color_bucket_cursors = + Tensor::vector_uninit(backend, stride * nb, storage)?; + state.rebuild_batch_indices(backend); } // Lazy resize based on the *previous* frame's max pair count. @@ -466,6 +500,8 @@ impl RbdPipeline { state.colored = Tensor::vector_uninit(backend, new_capacity * nb, storage)?; state.constraints_rands = Tensor::vector_uninit(backend, new_capacity * nb, storage)?; + state.color_sorted_ids = + Tensor::vector_uninit(backend, new_capacity * nb, storage)?; state.collision_pairs_per_batch_cpu = new_capacity; state.contacts_per_batch_cpu = new_capacity; diff --git a/src_rbd_shaders/dynamics/color_buckets.rs b/src_rbd_shaders/dynamics/color_buckets.rs new file mode 100644 index 0000000..0ee2fe2 --- /dev/null +++ b/src_rbd_shaders/dynamics/color_buckets.rs @@ -0,0 +1,120 @@ +//! Bucket-sort of contact constraints by graph-coloring color. +//! +//! After the per-step coloring converges, the constraint indices are +//! bucket-sorted by color (`color_sorted_ids`, contacts layout) with +//! per-batch per-color exclusive prefix sums (`color_starts`), so each +//! colored solver sweep iterates only its own bucket instead of scanning the +//! whole constraint buffer. The count/start/cursor buffers are flat +//! `[num_batches × stride]` arrays with `stride = +//! BatchIndices::solver_color_buckets_stride` (= `max_colors + 3`, keeping +//! `starts[c + 1]` in bounds for every swept color). + +use khal_std::glamx::UVec3; +use khal_std::macros::{spirv, spirv_bindgen}; +use khal_std::{index::MaybeIndexUnchecked, iter::StepRng, sync::atomic_add_u32}; + +use crate::utils::BatchIndices; + +const WORKGROUP_SIZE: u32 = 64; + +/// Zeroes the per-batch per-color constraint counts. +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn gpu_color_buckets_reset( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] color_counts: &mut [u32], + #[spirv(uniform, descriptor_set = 0, binding = 1)] batch_ids: &BatchIndices, +) { + let stride = batch_ids.solver_color_buckets_stride; + let batch_id = invocation_id.y; + let i = invocation_id.x; + + if i < stride { + color_counts.write((batch_id * stride + i) as usize, 0); + } +} + +/// Counts, per batch, how many constraints hold each color. +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn gpu_color_buckets_count( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[spirv(num_workgroups)] num_workgroups: UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] constraints_colors: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] contacts_len: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] color_counts: &mut [u32], + #[spirv(uniform, descriptor_set = 0, binding = 3)] batch_ids: &BatchIndices, +) { + let num_threads = num_workgroups.x * WORKGROUP_SIZE; + let batch_id = invocation_id.y; + let stride = batch_ids.solver_color_buckets_stride; + + let constraints_colors = batch_ids.contact_batch(batch_id, constraints_colors); + let len = contacts_len + .read(batch_id as usize) + .min(batch_ids.contacts_batch_capacity); + + for i in StepRng::new(invocation_id.x..len, num_threads) { + let color = constraints_colors[i as usize]; + // Colors past the swept range (can happen if the bounded coloring + // didn't converge) are dropped; they were never solved before either. + if color < stride - 1 { + atomic_add_u32(color_counts.at_mut((batch_id * stride + color) as usize), 1); + } + } +} + +/// Per-batch serial exclusive prefix sum over the (hopefully very small) per-color counts, +/// producing bucket start offsets. +#[spirv_bindgen] +#[spirv(compute(threads(1)))] +pub fn gpu_color_buckets_scan( + #[spirv(workgroup_id)] workgroup_id: UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] color_counts: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] color_starts: &mut [u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] color_cursors: &mut [u32], + #[spirv(uniform, descriptor_set = 0, binding = 3)] batch_ids: &BatchIndices, +) { + let stride = batch_ids.solver_color_buckets_stride; + let batch_id = workgroup_id.y; + let base = (batch_id * stride) as usize; + + let mut acc = 0u32; + for c in 0..stride as usize { + color_starts.write(base + c, acc); + color_cursors.write(base + c, acc); + acc += color_counts.read(base + c); + } +} + +/// Scatters each constraint index into its color's bucket. +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn gpu_color_buckets_scatter( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[spirv(num_workgroups)] num_workgroups: UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] constraints_colors: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] contacts_len: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] color_cursors: &mut [u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] color_sorted_ids: &mut [u32], + #[spirv(uniform, descriptor_set = 0, binding = 4)] batch_ids: &BatchIndices, +) { + let num_threads = num_workgroups.x * WORKGROUP_SIZE; + let batch_id = invocation_id.y; + let stride = batch_ids.solver_color_buckets_stride; + + let constraints_colors = batch_ids.contact_batch(batch_id, constraints_colors); + let mut color_sorted_ids = batch_ids.contact_batch_mut(batch_id, color_sorted_ids); + let len = contacts_len + .read(batch_id as usize) + .min(batch_ids.contacts_batch_capacity); + + for i in StepRng::new(invocation_id.x..len, num_threads) { + let color = constraints_colors[i as usize]; + if color < stride - 1 { + let dst = + atomic_add_u32(color_cursors.at_mut((batch_id * stride + color) as usize), 1); + color_sorted_ids[dst as usize] = i; + } + } +} diff --git a/src_rbd_shaders/dynamics/mod.rs b/src_rbd_shaders/dynamics/mod.rs index eca83ab..d1303a9 100644 --- a/src_rbd_shaders/dynamics/mod.rs +++ b/src_rbd_shaders/dynamics/mod.rs @@ -19,6 +19,7 @@ mod solver_utils; mod warmstart; // GPU compute shader kernels +mod color_buckets; mod coloring; mod mprops_update; mod prep_render; @@ -35,6 +36,7 @@ pub use joint_constraint_builder::{JointConstraintBuilder, JointConstraintHelper pub use multibody::*; pub use sim_params::*; // Re-export solver items; update_constraint comes from joint_constraint_builder for joints +pub use color_buckets::*; pub use coloring::*; pub use mprops_update::*; pub use prep_render::*; diff --git a/src_rbd_shaders/dynamics/solver.rs b/src_rbd_shaders/dynamics/solver.rs index a94cc78..507d839 100644 --- a/src_rbd_shaders/dynamics/solver.rs +++ b/src_rbd_shaders/dynamics/solver.rs @@ -335,34 +335,37 @@ pub fn gpu_warmstart( #[spirv(num_workgroups)] num_workgroups: UVec3, #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] constraints: &[TwoBodyConstraint], #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] solver_vels: &mut [Velocity], - #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] constraints_colors: &[u32], - #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] contacts_len: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] color_starts: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] color_sorted_ids: &[u32], #[spirv(uniform, descriptor_set = 0, binding = 4)] curr_color: &u32, #[spirv(uniform, descriptor_set = 0, binding = 5)] batch_ids: &BatchIndices, ) { let num_threads = num_workgroups.x * WORKGROUP_SIZE; let batch_id = invocation_id.y; + let stride = batch_ids.solver_color_buckets_stride; let constraints = batch_ids.contact_batch(batch_id, constraints); - let constraints_colors = batch_ids.contact_batch(batch_id, constraints_colors); + let color_sorted_ids = batch_ids.contact_batch(batch_id, color_sorted_ids); let mut solver_vels = batch_ids.coll_batch_mut(batch_id, solver_vels); - let len = contacts_len.read(batch_id as usize); let color = *curr_color; - for i in StepRng::new(invocation_id.x..len, num_threads) { - if constraints_colors[i as usize] == color { - let constraint = &constraints[i as usize]; - let solver_id1 = constraint.solver_body_a as usize; - let solver_id2 = constraint.solver_body_b as usize; + let bucket = (batch_id * stride + color) as usize; + let start = color_starts.read(bucket); + let end = color_starts.read(bucket + 1); - let mut solver_vel1 = solver_vels[solver_id1]; - let mut solver_vel2 = solver_vels[solver_id2]; + for k in StepRng::new(start + invocation_id.x..end, num_threads) { + let i = color_sorted_ids[k as usize]; + let constraint = &constraints[i as usize]; + let solver_id1 = constraint.solver_body_a as usize; + let solver_id2 = constraint.solver_body_b as usize; - constraint.warmstart_constraint(&mut solver_vel1, &mut solver_vel2); + let mut solver_vel1 = solver_vels[solver_id1]; + let mut solver_vel2 = solver_vels[solver_id2]; - solver_vels[solver_id1] = solver_vel1; - solver_vels[solver_id2] = solver_vel2; - } + constraint.warmstart_constraint(&mut solver_vel1, &mut solver_vel2); + + solver_vels[solver_id1] = solver_vel1; + solver_vels[solver_id2] = solver_vel2; } } @@ -375,35 +378,37 @@ pub fn gpu_step_gauss_seidel( #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] constraints: &mut [TwoBodyConstraint], #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] solver_vels: &mut [Velocity], - #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] constraints_colors: &[u32], - #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] contacts_len: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] color_starts: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] color_sorted_ids: &[u32], #[spirv(uniform, descriptor_set = 0, binding = 4)] curr_color: &u32, #[spirv(uniform, descriptor_set = 0, binding = 5)] batch_ids: &BatchIndices, ) { let num_threads = num_workgroups.x * WORKGROUP_SIZE; let batch_id = invocation_id.y; + let stride = batch_ids.solver_color_buckets_stride; let mut constraints = batch_ids.contact_batch_mut(batch_id, constraints); - let constraints_colors = batch_ids.contact_batch(batch_id, constraints_colors); + let color_sorted_ids = batch_ids.contact_batch(batch_id, color_sorted_ids); let mut solver_vels = batch_ids.coll_batch_mut(batch_id, solver_vels); - let len = contacts_len.read(batch_id as usize); let color = *curr_color; - for i in StepRng::new(invocation_id.x..len, num_threads) { - // Only process constraints of the current color (for parallelization) - if constraints_colors[i as usize] == color { - let solver_id1 = constraints[i as usize].solver_body_a as usize; - let solver_id2 = constraints[i as usize].solver_body_b as usize; + let bucket = (batch_id * stride + color) as usize; + let start = color_starts.read(bucket); + let end = color_starts.read(bucket + 1); - let mut solver_vel1 = solver_vels[solver_id1]; - let mut solver_vel2 = solver_vels[solver_id2]; + for k in StepRng::new(start + invocation_id.x..end, num_threads) { + let i = color_sorted_ids[k as usize]; + let solver_id1 = constraints[i as usize].solver_body_a as usize; + let solver_id2 = constraints[i as usize].solver_body_b as usize; - constraints[i as usize] - .solve_constraint_gauss_seidel(&mut solver_vel1, &mut solver_vel2); + let mut solver_vel1 = solver_vels[solver_id1]; + let mut solver_vel2 = solver_vels[solver_id2]; - solver_vels[solver_id1] = solver_vel1; - solver_vels[solver_id2] = solver_vel2; - } + constraints[i as usize] + .solve_constraint_gauss_seidel(&mut solver_vel1, &mut solver_vel2); + + solver_vels[solver_id1] = solver_vel1; + solver_vels[solver_id2] = solver_vel2; } } diff --git a/src_rbd_shaders/utils/indices.rs b/src_rbd_shaders/utils/indices.rs index df22253..19623ef 100644 --- a/src_rbd_shaders/utils/indices.rs +++ b/src_rbd_shaders/utils/indices.rs @@ -53,6 +53,10 @@ pub struct BatchIndices { /// contrast, are stored single-batch (identical coloring across batches) /// and read at offset 0. pub mb_imp_joint_color_groups_batch_capacity: u32, + /// Per-batch stride of the contact-solver color-bucket buffers + /// (`color_counts` / `color_starts` / `color_cursors`), = `max_colors + 3` + /// so that `starts[c + 1]` is in bounds for every swept color. + pub solver_color_buckets_stride: u32, /* * Intra-batch offsets for multi-purpose buffers. From 3cbd8ac9449feb08a4183a15ea50adae2dd8c5bc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Fri, 24 Jul 2026 15:27:58 +0200 Subject: [PATCH 05/39] perf: fold bias removal into the contact/joint solve kernels --- src_rbd/dynamics/joint.rs | 16 ++++------ src_rbd/dynamics/solver.rs | 15 ++++------ src_rbd_shaders/dynamics/joint_constraint.rs | 27 ++--------------- .../dynamics/joint_constraint_builder.rs | 9 ++++-- src_rbd_shaders/dynamics/solver.rs | 29 +++++-------------- src_rbd_shaders/dynamics/solver_utils.rs | 29 ++++--------------- 6 files changed, 33 insertions(+), 92 deletions(-) diff --git a/src_rbd/dynamics/joint.rs b/src_rbd/dynamics/joint.rs index 2a9efbe..8e35a58 100644 --- a/src_rbd/dynamics/joint.rs +++ b/src_rbd/dynamics/joint.rs @@ -5,7 +5,7 @@ use crate::math::Pose; use crate::shaders::dynamics::{ - GpuInitJointConstraints, GpuRemoveJointBias, + GpuInitJointConstraints, GpuSolveJointConstraints, GpuUpdateJointConstraints, ImpulseJoint, JointConstraint, JointConstraintBuilder, LocalMassProperties, RbdSimParams, Velocity, WorldMassProperties, }; @@ -370,8 +370,6 @@ pub struct GpuJointSolver { update_joint_constraints: GpuUpdateJointConstraints, /// Solves joint constraints. solve_joint_constraints: GpuSolveJointConstraints, - /// Removes bias from joint constraints. - remove_joint_bias: GpuRemoveJointBias, } /// Arguments given to the joint solver. @@ -452,14 +450,9 @@ impl GpuJointSolver { return Ok(()); } - if !use_bias { - self.remove_joint_bias.call( - pass, - [args.joints.len, args.num_batches, 1], - &mut args.joints.constraints, - args.batch_indices, - )?; - } + // Convert `use_bias` to its uniform. + // We can use `color_uniforms` since `color_uniforms[k] == k`. + let use_bias_uniform = &args.color_uniforms[use_bias as usize]; // One dispatch per color, sized exactly to that color's group (the // prefix sums are known on the host). The color index is bound as a @@ -482,6 +475,7 @@ impl GpuJointSolver { &args.joints.color_groups, &args.color_uniforms[c], args.batch_indices, + use_bias_uniform, )?; } diff --git a/src_rbd/dynamics/solver.rs b/src_rbd/dynamics/solver.rs index cfdcd06..7c06322 100644 --- a/src_rbd/dynamics/solver.rs +++ b/src_rbd/dynamics/solver.rs @@ -11,7 +11,7 @@ use crate::math::Pose; use crate::queries::GpuIndexedContact; use crate::shaders::dynamics::{ GpuApplySolverVelsInc, GpuInitSolverBodies, GpuInitSolverVelsInc, GpuIntegrateLinearized, - GpuRemoveCfmAndBiasKernel, GpuSolverCleanup, GpuSolverCountConstraints, GpuSolverFinalize, + GpuSolverCleanup, GpuSolverCountConstraints, GpuSolverFinalize, GpuSolverInitConstraints, GpuSolverSortConstraints, GpuSolverUpdateConstraints, GpuStepGaussSeidel, GpuWarmstart, GpuWarmstartWithoutColors, LocalMassProperties, RbdSimParams, TwoBodyConstraint, TwoBodyConstraintBuilder, Velocity, @@ -51,8 +51,6 @@ pub struct GpuSolver { /// Writes solver velocities and converts the COM-centered solver poses /// back to body-origin poses. finalize: GpuSolverFinalize, - /// Removes CFM and bias terms for velocity-only solving. - remove_cfm_and_bias_kernel: GpuRemoveCfmAndBiasKernel, } /// Arguments for constraint solver dispatch, used by [`GpuSolver::prepare`] and @@ -348,6 +346,8 @@ impl GpuSolver { args.color_sorted_ids, &args.color_uniforms[c as usize], args.batch_indices, + // use_bias = 1 (`color_uniforms[c]` holds the constant `c`). + &args.color_uniforms[1], )?; } @@ -369,13 +369,6 @@ 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, - )?; for c in 1..=args.num_colors { self.step_gauss_seidel.call( pass, @@ -386,6 +379,8 @@ impl GpuSolver { args.color_sorted_ids, &args.color_uniforms[c as usize], args.batch_indices, + // use_bias = 0 (`color_uniforms[c]` holds the constant `c`). + &args.color_uniforms[0], )?; } } diff --git a/src_rbd_shaders/dynamics/joint_constraint.rs b/src_rbd_shaders/dynamics/joint_constraint.rs index b00f355..04ba1e2 100644 --- a/src_rbd_shaders/dynamics/joint_constraint.rs +++ b/src_rbd_shaders/dynamics/joint_constraint.rs @@ -243,29 +243,6 @@ pub fn gpu_update_joint_constraints( } } -/// Removes bias from joint constraints for the final substep. -#[spirv_bindgen] -#[spirv(compute(threads(64)))] -pub fn gpu_remove_joint_bias( - #[spirv(global_invocation_id)] invocation_id: UVec3, - #[spirv(num_workgroups)] num_workgroups: UVec3, - #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] constraints: &mut [JointConstraint], - #[spirv(uniform, descriptor_set = 0, binding = 1)] batch_ids: &BatchIndices, -) { - let num_threads = num_workgroups.x * WORKGROUP_SIZE; - let batch_id = invocation_id.y; - let mut constraints = batch_ids.impulse_joints_batch_mut(batch_id, constraints); - - let len = batch_ids.impulse_joints_len; - - for i in StepRng::new(invocation_id.x..len, num_threads) { - let idx = i as usize; - for j in 0..(constraints[idx].len as usize) { - constraints[idx].elements.at_mut(j).rhs = constraints[idx].elements.at(j).rhs_wo_bias; - } - } -} - /// Solves joint constraints. #[spirv_bindgen] #[spirv(compute(threads(64)))] @@ -277,12 +254,14 @@ pub fn gpu_solve_joint_constraints( #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] all_color_groups: &[u32], #[spirv(uniform, descriptor_set = 0, binding = 3)] curr_color: &u32, #[spirv(uniform, descriptor_set = 0, binding = 4)] batch_ids: &BatchIndices, + #[spirv(uniform, descriptor_set = 0, binding = 5)] use_bias: &u32, ) { let num_threads = num_workgroups.x * WORKGROUP_SIZE; let batch_id = invocation_id.y; let mut constraints = batch_ids.impulse_joints_batch_mut(batch_id, constraints); let mut solver_vels = batch_ids.coll_batch_mut(batch_id, solver_vels); + let use_bias = *use_bias != 0; let color = *curr_color as usize; // Coloring is identical across batches (enforced on the host), so the @@ -297,6 +276,6 @@ pub fn gpu_solve_joint_constraints( let end = color_groups[color]; for i in StepRng::new(start + invocation_id.x..end, num_threads) { - constraints[i as usize].solve_joint_constraint(&mut solver_vels); + constraints[i as usize].solve_joint_constraint(&mut solver_vels, use_bias); } } diff --git a/src_rbd_shaders/dynamics/joint_constraint_builder.rs b/src_rbd_shaders/dynamics/joint_constraint_builder.rs index 7ce3cdb..a1e5707 100644 --- a/src_rbd_shaders/dynamics/joint_constraint_builder.rs +++ b/src_rbd_shaders/dynamics/joint_constraint_builder.rs @@ -330,7 +330,7 @@ impl JointConstraintHelper { impl JointConstraint { /// Solves a joint constraint. - pub fn solve_joint_constraint(&mut self, solver_vels: &mut SliceMut) { + pub fn solve_joint_constraint(&mut self, solver_vels: &mut SliceMut, use_bias: bool) { let mut solver_vel1 = solver_vels[self.solver_vel_a as usize]; let mut solver_vel2 = solver_vels[self.solver_vel_b as usize]; @@ -340,7 +340,12 @@ impl JointConstraint { let dangvel = gdot(element.ang_jac_b, solver_vel2.angular) - gdot(element.ang_jac_a, solver_vel1.angular); - let rhs = dlinvel + dangvel + element.rhs; + let el_rhs = if use_bias { + element.rhs + } else { + element.rhs_wo_bias + }; + let rhs = dlinvel + dangvel + el_rhs; let total_impulse = (element.impulse + element.inv_lhs * (rhs - element.cfm_gain * element.impulse)) .clamp(element.impulse_bounds.x, element.impulse_bounds.y); diff --git a/src_rbd_shaders/dynamics/solver.rs b/src_rbd_shaders/dynamics/solver.rs index 507d839..98185b3 100644 --- a/src_rbd_shaders/dynamics/solver.rs +++ b/src_rbd_shaders/dynamics/solver.rs @@ -382,6 +382,7 @@ pub fn gpu_step_gauss_seidel( #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] color_sorted_ids: &[u32], #[spirv(uniform, descriptor_set = 0, binding = 4)] curr_color: &u32, #[spirv(uniform, descriptor_set = 0, binding = 5)] batch_ids: &BatchIndices, + #[spirv(uniform, descriptor_set = 0, binding = 6)] use_bias: &u32, ) { let num_threads = num_workgroups.x * WORKGROUP_SIZE; let batch_id = invocation_id.y; @@ -391,6 +392,7 @@ pub fn gpu_step_gauss_seidel( let color_sorted_ids = batch_ids.contact_batch(batch_id, color_sorted_ids); let mut solver_vels = batch_ids.coll_batch_mut(batch_id, solver_vels); let color = *curr_color; + let use_bias = *use_bias != 0; let bucket = (batch_id * stride + color) as usize; let start = color_starts.read(bucket); @@ -404,8 +406,11 @@ pub fn gpu_step_gauss_seidel( let mut solver_vel1 = solver_vels[solver_id1]; let mut solver_vel2 = solver_vels[solver_id2]; - constraints[i as usize] - .solve_constraint_gauss_seidel(&mut solver_vel1, &mut solver_vel2); + constraints[i as usize].solve_constraint_gauss_seidel( + &mut solver_vel1, + &mut solver_vel2, + use_bias, + ); solver_vels[solver_id1] = solver_vel1; solver_vels[solver_id2] = solver_vel2; @@ -501,23 +506,3 @@ pub fn gpu_solver_finalize( } } -/// Removes CFM and bias from constraints for the final substep iteration. -#[spirv_bindgen] -#[spirv(compute(threads(64)))] -pub fn gpu_remove_cfm_and_bias_kernel( - #[spirv(global_invocation_id)] invocation_id: UVec3, - #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] - constraints: &mut [TwoBodyConstraint], - #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] contacts_len: &[u32], - #[spirv(uniform, descriptor_set = 0, binding = 2)] batch_ids: &BatchIndices, -) { - let batch_id = invocation_id.y; - let i = invocation_id.x; - - let mut constraints = batch_ids.contact_batch_mut(batch_id, constraints); - let len = contacts_len.read(batch_id as usize); - - if i < len { - constraints[i as usize].remove_cfm_and_bias(); - } -} diff --git a/src_rbd_shaders/dynamics/solver_utils.rs b/src_rbd_shaders/dynamics/solver_utils.rs index 2f65380..4f48234 100644 --- a/src_rbd_shaders/dynamics/solver_utils.rs +++ b/src_rbd_shaders/dynamics/solver_utils.rs @@ -538,12 +538,15 @@ impl TwoBodyConstraint { &mut self, solver_vel1: &mut Velocity, solver_vel2: &mut Velocity, + use_bias: bool, ) { let dir_a = self.dir_a; let friction_coeff = self.limit; let im_a = self.im_a; let im_b = self.im_b; - let cfm_factor = self.cfm_factor; + // The stabilization (no-bias) sweep uses the bias-free rhs and no CFM. + // This replaces the former `remove_cfm_and_bias` full-buffer pass. + let cfm_factor = if use_bias { self.cfm_factor } else { 1.0 }; #[cfg(feature = "dim2")] let tangent_a = Vec2::new(-dir_a.y, dir_a.x); @@ -558,11 +561,12 @@ impl TwoBodyConstraint { let ii_torque_dir_a = c.ii_torque_dir_a; let ii_torque_dir_b = c.ii_torque_dir_b; + let rhs = if use_bias { c.rhs } else { c.rhs_wo_bias }; let dvel = dir_a.dot(solver_vel1.linear) + gdot(c.torque_dir_a, solver_vel1.angular) - dir_a.dot(solver_vel2.linear) + gdot(c.torque_dir_b, solver_vel2.angular) - + c.rhs; + + rhs; let new_impulse = cfm_factor * (c.impulse - c.r * dvel).max(0.0); let delta_impulse = new_impulse - c.impulse; @@ -651,24 +655,3 @@ impl TwoBodyConstraint { } } -impl TwoBodyConstraint { - /// Removes CFM and bias from constraints for the final substep iteration. - #[cfg(feature = "dim2")] - #[inline(always)] - pub fn remove_cfm_and_bias(&mut self) { - self.elements.at_mut(0).normal_part.rhs = self.elements.at(0).normal_part.rhs_wo_bias; - self.elements.at_mut(1).normal_part.rhs = self.elements.at(1).normal_part.rhs_wo_bias; - self.cfm_factor = 1.0; - } - - /// Removes CFM and bias from constraints for the final substep iteration. - #[cfg(feature = "dim3")] - #[inline(always)] - pub fn remove_cfm_and_bias(&mut self) { - self.elements.at_mut(0).normal_part.rhs = self.elements.at(0).normal_part.rhs_wo_bias; - self.elements.at_mut(1).normal_part.rhs = self.elements.at(1).normal_part.rhs_wo_bias; - self.elements.at_mut(2).normal_part.rhs = self.elements.at(2).normal_part.rhs_wo_bias; - self.elements.at_mut(3).normal_part.rhs = self.elements.at(3).normal_part.rhs_wo_bias; - self.cfm_factor = 1.0; - } -} From c912154df9aa69e352073a5b1db8b10e7651e0e4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Fri, 24 Jul 2026 15:39:03 +0200 Subject: [PATCH 06/39] perf: cache radix-sort pass uniforms and n_sort_flat across calls --- src_rbd/utils/radix_sort/mod.rs | 90 +++++++++++++++++++-------------- 1 file changed, 52 insertions(+), 38 deletions(-) diff --git a/src_rbd/utils/radix_sort/mod.rs b/src_rbd/utils/radix_sort/mod.rs index 6111942..667520e 100644 --- a/src_rbd/utils/radix_sort/mod.rs +++ b/src_rbd/utils/radix_sort/mod.rs @@ -38,6 +38,9 @@ pub struct RadixSort { /// intermediate buffers as needed. pub struct RadixSortWorkspace { pass_uniforms: Vec>, + /// Configuration `(mode, sorting_bits, max_keys, num_batches)` the cached + /// `pass_uniforms` / `n_sort_flat` were built for. + uniforms_key: (u32, u32, u32, u32), reduced_buf: Tensor, // Tensor of size BLOCK_SIZE count_buf: Tensor, num_wgs: Tensor<[u32; 3]>, @@ -57,6 +60,7 @@ impl RadixSortWorkspace { let zeros = vec![0u32; BLOCK_SIZE as usize]; Self { pass_uniforms: vec![], + uniforms_key: (u32::MAX, 0, 0, 0), reduced_buf: Tensor::vector(backend, &zeros, BufferUsages::STORAGE).unwrap(), count_buf: Tensor::vector_uninit(backend, 0, BufferUsages::STORAGE).unwrap(), num_wgs: Tensor::scalar( @@ -254,18 +258,22 @@ impl RadixSort { let num_passes = sorting_bits.div_ceil(4); - // Create uniforms (has_aux=0 for single batch). - workspace.pass_uniforms.clear(); - for pass_id in 0..num_passes { - workspace.pass_uniforms.push(Tensor::scalar( - backend, - SortUniforms { - shift: pass_id * 4, - max_keys_per_batch: per_batch_max, - has_aux: 0, - }, - BufferUsages::STORAGE | BufferUsages::UNIFORM, - )?); + // Create uniforms (has_aux=0 for single batch), cached across calls. + let uniforms_key = (0, sorting_bits, per_batch_max, 1); + if workspace.uniforms_key != uniforms_key { + workspace.pass_uniforms.clear(); + for pass_id in 0..num_passes { + workspace.pass_uniforms.push(Tensor::scalar( + backend, + SortUniforms { + shift: pass_id * 4, + max_keys_per_batch: per_batch_max, + has_aux: 0, + }, + BufferUsages::STORAGE | BufferUsages::UNIFORM, + )?); + } + workspace.uniforms_key = uniforms_key; } let mut output_keys = output_keys; @@ -402,41 +410,47 @@ impl RadixSort { Tensor::vector_uninit(backend, total_n, BufferUsages::STORAGE)?; } - // n_sort_flat = [total_n] for the flattened single-batch view. - workspace.n_sort_flat = Tensor::scalar(backend, total_n, BufferUsages::STORAGE)?; + // The pass uniforms and the flattened count only depend on + // `(sorting_bits, total_n, num_batches)`; cache them across calls. + let init_uniform_idx = total_passes as usize; + let uniforms_key = (1, sorting_bits, total_n, num_batches); + if workspace.uniforms_key != uniforms_key { + // n_sort_flat = [total_n] for the flattened single-batch view. + workspace.n_sort_flat = Tensor::scalar(backend, total_n, BufferUsages::STORAGE)?; + + // Create uniforms for all passes. + workspace.pass_uniforms.clear(); + for pass_id in 0..total_passes { + let shift = if pass_id < key_passes { + pass_id * 4 + } else { + (pass_id - key_passes) * 4 + }; + workspace.pass_uniforms.push(Tensor::scalar( + backend, + SortUniforms { + shift, + max_keys_per_batch: total_n, + has_aux: 1, + }, + BufferUsages::STORAGE | BufferUsages::UNIFORM, + )?); + } - // Create uniforms for all passes. - workspace.pass_uniforms.clear(); - for pass_id in 0..total_passes { - let shift = if pass_id < key_passes { - pass_id * 4 - } else { - (pass_id - key_passes) * 4 - }; + // Extra uniform for init_batched (max_keys_per_batch = per_batch, not total_n). + // shift is repurposed to carry num_batches for this kernel. workspace.pass_uniforms.push(Tensor::scalar( backend, SortUniforms { - shift, - max_keys_per_batch: total_n, - has_aux: 1, + shift: num_batches, + max_keys_per_batch: per_batch, + has_aux: 0, }, BufferUsages::STORAGE | BufferUsages::UNIFORM, )?); + workspace.uniforms_key = uniforms_key; } - // Extra uniform for init_batched (max_keys_per_batch = per_batch, not total_n). - // shift is repurposed to carry num_batches for this kernel. - let init_uniform_idx = total_passes as usize; - workspace.pass_uniforms.push(Tensor::scalar( - backend, - SortUniforms { - shift: num_batches, - max_keys_per_batch: per_batch, - has_aux: 0, - }, - BufferUsages::STORAGE | BufferUsages::UNIFORM, - )?); - // Init writes to output buffers. After even total_passes, data stays in output. // NOTE: call() takes a thread count, not workgroup count (khal resolves internally). self.init_batched.call( From 3811f97c0bcd47cadc78f1d16377e591b5afbe4c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Fri, 24 Jul 2026 16:10:59 +0200 Subject: [PATCH 07/39] perf: enable cuboid-cuboid SAT early-exits --- src_rbd_shaders/queries/contact.rs | 23 +++++++++++------------ src_rbd_shaders/queries/sat.rs | 3 +++ 2 files changed, 14 insertions(+), 12 deletions(-) diff --git a/src_rbd_shaders/queries/contact.rs b/src_rbd_shaders/queries/contact.rs index aa4aa0f..26b8678 100644 --- a/src_rbd_shaders/queries/contact.rs +++ b/src_rbd_shaders/queries/contact.rs @@ -226,17 +226,17 @@ pub fn cuboid_cuboid( */ let sep1 = sat::cuboid_cuboid_find_local_separating_normal_oneway(cuboid1, cuboid2, pose12); - // TODO PERF: support the prediction early-exit. - // if sep1.separation > prediction { - // return ContactManifold::default(); - // } + // Early-exit: any contact point's distance is >= the separation along a + // separating axis, so the caller would drop the manifold anyway. + if sep1.separation > prediction { + return ContactManifold::default(); + } let sep2 = sat::cuboid_cuboid_find_local_separating_normal_oneway(cuboid2, cuboid1, pose21); - // TODO PERF: support the prediction early-exit. - // if sep2.separation > prediction { - // return ContactManifold::default(); - // } + if sep2.separation > prediction { + return ContactManifold::default(); + } /* * @@ -248,10 +248,9 @@ pub fn cuboid_cuboid( #[cfg(feature = "dim3")] let sep3 = sat::cuboid_cuboid_find_local_separating_edge_twoway(cuboid1, cuboid2, pose12); - // TODO PERF: support the prediction early-exit. - // if sep3.separation > prediction { - // return ContactManifold::default(); - // } + if sep3.separation > prediction { + return ContactManifold::default(); + } /* * diff --git a/src_rbd_shaders/queries/sat.rs b/src_rbd_shaders/queries/sat.rs index bcbbea6..c018ab7 100644 --- a/src_rbd_shaders/queries/sat.rs +++ b/src_rbd_shaders/queries/sat.rs @@ -40,6 +40,7 @@ pub const EPSILON: f32 = 1.1920929E-7; #[cfg(feature = "dim3")] /// Computes the separation of two cuboids along `axis1`. +#[inline(always)] pub fn cuboid_cuboid_compute_separation_wrt_local_line( cuboid1: &Cuboid, cuboid2: &Cuboid, @@ -65,6 +66,7 @@ pub fn cuboid_cuboid_compute_separation_wrt_local_line( /// /// All combinations of edges from both cuboids are taken into /// account. +#[inline(always)] pub fn cuboid_cuboid_find_local_separating_edge_twoway( cuboid1: &Cuboid, cuboid2: &Cuboid, @@ -116,6 +118,7 @@ pub fn cuboid_cuboid_find_local_separating_edge_twoway( /// Finds the best separating normal between two cuboids. /// /// Only the normals from `cuboid1` are tested. +#[inline(always)] pub fn cuboid_cuboid_find_local_separating_normal_oneway( cuboid1: &Cuboid, cuboid2: &Cuboid, From 9ca0dad36a867ee57f3192b39ec534045d027ca3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Fri, 24 Jul 2026 16:32:48 +0200 Subject: [PATCH 08/39] perf: check the analytic-pair predicate before loading poses in the deferred narrow phase --- src_rbd_shaders/broad_phase/narrow_phase.rs | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/src_rbd_shaders/broad_phase/narrow_phase.rs b/src_rbd_shaders/broad_phase/narrow_phase.rs index 4d4f20d..3805d09 100644 --- a/src_rbd_shaders/broad_phase/narrow_phase.rs +++ b/src_rbd_shaders/broad_phase/narrow_phase.rs @@ -233,13 +233,10 @@ pub fn gpu_narrow_phase_shape_shape_deferred( // contact is written. for i in StepRng::new(invocation_id.x..len, num_threads) { let pair = collision_pairs[i as usize]; - let pose1 = poses[pair.colliders.x as usize]; - let pose2 = poses[pair.colliders.y as usize]; let shape1 = &shapes[pair.colliders.x as usize]; let shape2 = &shapes[pair.colliders.y as usize]; let shape_ty1 = shape1.shape_type(); let shape_ty2 = shape2.shape_type(); - let pose12 = pose1.inverse() * pose2; // Mirror pass 1's analytic-pair predicate (ball/cuboid) so those pairs // are skipped here — they were already turned into contacts. Only the @@ -266,6 +263,13 @@ pub fn gpu_narrow_phase_shape_shape_deferred( if !checked && shape_ty1 == SHAPE_TYPE_CUBOID && shape_ty2 == SHAPE_TYPE_CUBOID { checked = true; } + if checked { + continue; + } + + let pose1 = poses[pair.colliders.x as usize]; + let pose2 = poses[pair.colliders.y as usize]; + let pose12 = pose1.inverse() * pose2; // PFM - PFM (generic convex shapes via GJK/EPA) if !checked { From b63faad1cc60a04f8aa3bb8bd9fba7eee3d069fb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Fri, 24 Jul 2026 18:50:37 +0200 Subject: [PATCH 09/39] perf: bound constraint init/count loops by the indirect grid, not capacity --- src_rbd_shaders/dynamics/solver.rs | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/src_rbd_shaders/dynamics/solver.rs b/src_rbd_shaders/dynamics/solver.rs index 98185b3..c751a89 100644 --- a/src_rbd_shaders/dynamics/solver.rs +++ b/src_rbd_shaders/dynamics/solver.rs @@ -51,11 +51,7 @@ pub fn gpu_solver_init_constraints( let solver_body_poses = batch_ids.coll_batch(batch_id, solver_body_poses); let vels = batch_ids.coll_batch(batch_id, vels); let mprops = batch_ids.coll_batch(batch_id, mprops); - // Iterating to `cap` (instead of `contacts_len[batch]`) lets us drop the - // `contacts_len` storage binding. Empty / unused contact slots have - // `contact.len == 0` and are skipped — narrow-phase zero-initialises the - // buffer so the sentinel is reliable. - let cap = batch_ids.contacts_batch_capacity; + let cap = batch_ids.contacts_batch_capacity.min(num_threads); for i in StepRng::new(invocation_id.x..cap, num_threads) { let im = &contacts[i as usize]; @@ -94,7 +90,9 @@ pub fn gpu_solver_count_constraints( let mut body_constraint_counts = batch_ids.coll_batch_mut(batch_id, body_constraint_counts); let body_group = batch_ids.coll_batch(batch_id, body_group); let mprops = batch_ids.coll_batch(batch_id, mprops); - let cap = batch_ids.contacts_batch_capacity; + // See `gpu_solver_init_constraints` — the indirect grid bounds the active + // range much tighter than the capacity. + let cap = batch_ids.contacts_batch_capacity.min(num_threads); for i in StepRng::new(invocation_id.x..cap, num_threads) { let im = &contacts[i as usize]; From 9f6a3e9825de4e9e6d20f10a0990d1664a374606 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Fri, 24 Jul 2026 19:30:23 +0200 Subject: [PATCH 10/39] perf: seed the topo-gc coloring from the previous frame's colors --- src_rbd/dynamics/coloring.rs | 22 ++++++ src_rbd/dynamics/warmstart.rs | 52 +++++++++++++- src_rbd/pipeline/insertion_removal.rs | 7 ++ src_rbd/pipeline/rbd_state.rs | 1 + src_rbd/pipeline/rbd_state_from_rapier.rs | 7 ++ src_rbd/pipeline/rbd_step.rs | 44 +++++++++++- src_rbd_shaders/dynamics/warmstart.rs | 82 +++++++++++++++++++++++ 7 files changed, 213 insertions(+), 2 deletions(-) diff --git a/src_rbd/dynamics/coloring.rs b/src_rbd/dynamics/coloring.rs index 88f2e33..f187bfd 100644 --- a/src_rbd/dynamics/coloring.rs +++ b/src_rbd/dynamics/coloring.rs @@ -299,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)?; diff --git a/src_rbd/dynamics/warmstart.rs b/src_rbd/dynamics/warmstart.rs index 00c1e2e..b0b07f3 100644 --- a/src_rbd/dynamics/warmstart.rs +++ b/src_rbd/dynamics/warmstart.rs @@ -1,7 +1,8 @@ //! Warmstarting: reuses previous-frame impulses for faster solver convergence. use crate::shaders::dynamics::{ - GpuTransferWarmstartImpulses, TwoBodyConstraint, TwoBodyConstraintBuilder, + GpuSeedColorsFromWarmstart, GpuTransferWarmstartImpulses, TwoBodyConstraint, + TwoBodyConstraintBuilder, }; use crate::shaders::utils::BatchIndices; use khal::Shader; @@ -16,6 +17,9 @@ use vortx::tensor::Tensor; pub struct GpuWarmstart { /// Compute pipeline that matches contacts and transfers impulses. transfer_warmstart_impulses_kernel: GpuTransferWarmstartImpulses, + /// Seeds the topo-gc coloring from the previous frame's colors (same + /// old/new body-pair matching as the impulse transfer). + seed_colors_kernel: GpuSeedColorsFromWarmstart, } /// Arguments for warmstart dispatch. @@ -42,6 +46,30 @@ pub struct WarmstartArgs<'a> { pub batch_indices: &'a Tensor, } +/// Arguments for the coloring seed dispatch. +pub struct SeedColorsArgs<'a> { + /// Number of contacts in current frame. + pub contacts_len: &'a Tensor, + /// Constraint counts per body from previous frame. + pub old_body_constraint_counts: &'a Tensor, + /// Constraint IDs per body from previous frame. + pub old_body_constraint_ids: &'a Tensor, + /// Solver constraints from previous frame. + pub old_constraints: &'a Tensor, + /// Solver constraints for current frame. + pub new_constraints: &'a Tensor, + /// Colors assigned to the previous frame's constraints. + pub old_constraints_colors: &'a Tensor, + /// Output: colors for the current frame's constraints (seeded slots only). + pub constraints_colors: &'a mut Tensor, + /// Output: per-constraint colored flag consumed by topo-gc. + pub colored: &'a mut Tensor, + /// Indirect dispatch arguments based on contact count. + pub contacts_len_indirect: &'a Tensor<[u32; 3]>, + /// Shared per-batch index uniform. + pub batch_indices: &'a Tensor, +} + impl GpuWarmstart { /// Transfers warmstart impulses from old constraints to new constraints. pub fn transfer_warmstart_impulses<'a>( @@ -62,4 +90,26 @@ impl GpuWarmstart { args.batch_indices, ) } + + /// Seeds the topo-gc coloring from the previous frame's colors. Must run + /// after the topo-gc reset and before its iterations. + pub fn seed_colors_from_warmstart( + &self, + pass: &mut GpuPass, + args: SeedColorsArgs<'_>, + ) -> Result<(), GpuBackendError> { + self.seed_colors_kernel.call( + pass, + args.contacts_len_indirect, + args.old_body_constraint_counts, + args.old_body_constraint_ids, + args.old_constraints, + args.new_constraints, + args.old_constraints_colors, + args.constraints_colors, + args.colored, + args.contacts_len, + args.batch_indices, + ) + } } diff --git a/src_rbd/pipeline/insertion_removal.rs b/src_rbd/pipeline/insertion_removal.rs index 0cc8438..8933ea8 100644 --- a/src_rbd/pipeline/insertion_removal.rs +++ b/src_rbd/pipeline/insertion_removal.rs @@ -178,6 +178,12 @@ impl RbdState { Tensor::vector_uninit(backend, collisions_capacity * num_batches, storage).unwrap(); let constraints_colors = Tensor::vector_uninit(backend, collisions_capacity * num_batches, storage).unwrap(); + let old_constraints_colors = Tensor::vector( + backend, + &vec![0u32; (collisions_capacity * num_batches) as usize], + storage, + ) + .unwrap(); let colored = Tensor::vector_uninit(backend, collisions_capacity * num_batches, storage).unwrap(); let constraints_rands = @@ -279,6 +285,7 @@ impl RbdState { new_constraint_builders, new_constraints_counts, constraints_colors, + old_constraints_colors, colored, constraints_rands, color_bucket_counts, diff --git a/src_rbd/pipeline/rbd_state.rs b/src_rbd/pipeline/rbd_state.rs index 920f008..60b7ae2 100644 --- a/src_rbd/pipeline/rbd_state.rs +++ b/src_rbd/pipeline/rbd_state.rs @@ -196,6 +196,7 @@ pub struct RbdState { pub(super) old_constraints_counts: Tensor, pub(super) old_body_constraint_ids: Tensor, pub(super) constraints_colors: Tensor, + pub(super) old_constraints_colors: Tensor, pub(super) colored: Tensor, pub(super) constraints_rands: Tensor, /// Per-batch per-color constraint counts (stride `max_colors + 3`), see diff --git a/src_rbd/pipeline/rbd_state_from_rapier.rs b/src_rbd/pipeline/rbd_state_from_rapier.rs index 64aa081..24fbaf8 100644 --- a/src_rbd/pipeline/rbd_state_from_rapier.rs +++ b/src_rbd/pipeline/rbd_state_from_rapier.rs @@ -625,6 +625,12 @@ impl RbdState { storage, ) .unwrap(); + let old_constraints_colors = Tensor::vector( + backend, + &vec![0u32; (capacities.collisions_capacity * num_batches) as usize], + storage, + ) + .unwrap(); let colored = Tensor::vector_uninit( backend, capacities.collisions_capacity * num_batches, @@ -774,6 +780,7 @@ impl RbdState { new_constraint_builders, new_constraints_counts, constraints_colors, + old_constraints_colors, colored, constraints_rands, color_bucket_counts, diff --git a/src_rbd/pipeline/rbd_step.rs b/src_rbd/pipeline/rbd_step.rs index a3d3c1a..f01a705 100644 --- a/src_rbd/pipeline/rbd_step.rs +++ b/src_rbd/pipeline/rbd_step.rs @@ -277,6 +277,40 @@ impl RbdPipeline { self.warmstart .transfer_warmstart_impulses(&mut pass, warmstart_args)?; + let coloring_args = ColoringArgs { + contacts_len_indirect: &state.contacts_indirect, + body_constraint_counts: &state.new_constraints_counts, + body_constraint_ids: &state.new_body_constraint_ids, + constraints: &state.new_constraints, + constraints_colors: &mut state.constraints_colors, + constraints_rands: &mut state.constraints_rands, + curr_color: &mut state.curr_color, + uncolored: &mut state.uncolored, + uncolored_staging: &state.uncolored_staging, + contacts_len: &state.contacts_len, + colored: &mut state.colored, + batch_indices: &state.batch_indices, + body_group: &state.body_group, + }; + self.coloring.dispatch_topo_gc_reset(&mut pass, coloring_args)?; + + // Seed the coloring from the previous frame's colors (contacts + // persist, so most constraints can reuse their old color and the + // topo-gc iterations converge in 1-2 rounds instead of ~num_colors). + let seed_args = crate::dynamics::warmstart::SeedColorsArgs { + contacts_len: &state.contacts_len, + old_body_constraint_counts: &state.old_constraints_counts, + old_body_constraint_ids: &state.old_body_constraint_ids, + old_constraints: &state.old_constraints, + new_constraints: &state.new_constraints, + old_constraints_colors: &state.old_constraints_colors, + constraints_colors: &mut state.constraints_colors, + colored: &mut state.colored, + contacts_len_indirect: &state.contacts_indirect, + batch_indices: &state.batch_indices, + }; + self.warmstart.seed_colors_from_warmstart(&mut pass, seed_args)?; + let coloring_args = ColoringArgs { contacts_len_indirect: &state.contacts_indirect, body_constraint_counts: &state.new_constraints_counts, @@ -293,7 +327,7 @@ impl RbdPipeline { body_group: &state.body_group, }; self.coloring - .dispatch_topo_gc_bounded(&mut pass, coloring_args, state.max_colors)?; + .dispatch_topo_gc_iterations(&mut pass, coloring_args, state.max_colors)?; // Bucket-sort the constraint ids by color so each colored solver // sweep only touches its own constraints. @@ -407,6 +441,10 @@ impl RbdPipeline { &mut state.old_constraints_counts, &mut state.new_constraints_counts, ); + std::mem::swap( + &mut state.old_constraints_colors, + &mut state.constraints_colors, + ); Ok(stats) } @@ -497,6 +535,10 @@ impl RbdPipeline { Tensor::vector_uninit(backend, new_capacity * 2 * nb, storage)?; state.constraints_colors = Tensor::vector_uninit(backend, new_capacity * nb, storage)?; + // Zeroed (not uninit): 0 = "uncolored" disables color seeding + // for the frame right after the resize. + state.old_constraints_colors = + Tensor::vector(backend, &vec![0u32; (new_capacity * nb) as usize], storage)?; state.colored = Tensor::vector_uninit(backend, new_capacity * nb, storage)?; state.constraints_rands = Tensor::vector_uninit(backend, new_capacity * nb, storage)?; diff --git a/src_rbd_shaders/dynamics/warmstart.rs b/src_rbd_shaders/dynamics/warmstart.rs index 6e68659..ff0b210 100644 --- a/src_rbd_shaders/dynamics/warmstart.rs +++ b/src_rbd_shaders/dynamics/warmstart.rs @@ -56,6 +56,88 @@ pub fn gpu_transfer_warmstart_impulses( } } +/// Seeds the topo-gc coloring from the previous frame's colors (same +/// body-pair matching as the warmstart impulse transfer), between the +/// topo-gc reset and its iterations. +/// +/// Contacts persist across frames, so most constraints can reuse last +/// frame's color and the iterations only color the genuinely new ones; the +/// fix-conflicts pass still validates every seed, so a stale seed is simply +/// uncolored and recomputed. +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn gpu_seed_colors_from_warmstart( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] old_body_constraint_counts: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] old_body_constraint_ids: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] + old_constraints: &[TwoBodyConstraint], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] new_constraints: &[TwoBodyConstraint], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] old_constraints_colors: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] constraints_colors: &mut [u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 6)] colored: &mut [u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 7)] contacts_len: &[u32], + #[spirv(uniform, descriptor_set = 0, binding = 8)] batch_ids: &BatchIndices, +) { + let batch_id = invocation_id.y; + let contacts_start = batch_ids.contacts_start(batch_id); + let colliders_start = batch_ids.coll_start(batch_id); + let bci_start = batch_id as usize * 2 * batch_ids.contacts_batch_capacity as usize; + + let old_body_constraint_counts = Slice(old_body_constraint_counts, colliders_start); + let old_body_constraint_ids = Slice(old_body_constraint_ids, bci_start); + let old_constraints = Slice(old_constraints, contacts_start); + let old_constraints_colors = Slice(old_constraints_colors, contacts_start); + let mut constraints_colors = SliceMut(constraints_colors, contacts_start); + let mut colored = SliceMut(colored, contacts_start); + let new_constraints = Slice(new_constraints, contacts_start); + + let len = contacts_len.read(batch_id as usize); + let i = invocation_id.x as usize; + + if (i as u32) < len { + let body_a = new_constraints[i].solver_body_a; + let body_b = new_constraints[i].solver_body_b; + + let first_a = if body_a != 0 { + old_body_constraint_counts[body_a as usize - 1] as usize + } else { + 0 + }; + let last_a = old_body_constraint_counts[body_a as usize] as usize; + let first_b = if body_b != 0 { + old_body_constraint_counts[body_b as usize - 1] as usize + } else { + 0 + }; + let last_b = old_body_constraint_counts[body_b as usize] as usize; + + let len_a = last_a - first_a; + let len_b = last_b - first_b; + let (first_ref, last_ref) = if len_a != 0 && len_a < len_b { + (first_a, last_a) + } else { + (first_b, last_b) + }; + + for j in first_ref..last_ref { + let cid_old = old_body_constraint_ids[j] as usize; + if old_constraints[cid_old].solver_body_a == body_a + && old_constraints[cid_old].solver_body_b == body_b + { + let old_color = old_constraints_colors[cid_old]; + // Colors 1..64 are the valid topo-gc range; anything else + // (e.g. stale data after a buffer resize) stays uncolored. + if old_color > 0 && old_color < 64 { + constraints_colors[i] = old_color; + colored[i] = 1; + } + break; + } + } + } +} + /// Transfers warmstart impulses from previous frame to current frame. /// /// NOTE: this assumes that the solver body ids in the constraints match the index of the body itself. From 6d2ce60168c009c977653d18ca4b1e635777ec39 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Sat, 25 Jul 2026 10:09:32 +0200 Subject: [PATCH 11/39] perf: bound LBVH internal refit by the active collider count --- src_rbd_shaders/broad_phase/lbvh.rs | 13 +++++-------- 1 file changed, 5 insertions(+), 8 deletions(-) diff --git a/src_rbd_shaders/broad_phase/lbvh.rs b/src_rbd_shaders/broad_phase/lbvh.rs index fa9ff93..1d9455a 100644 --- a/src_rbd_shaders/broad_phase/lbvh.rs +++ b/src_rbd_shaders/broad_phase/lbvh.rs @@ -342,12 +342,7 @@ pub fn gpu_lbvh_refit_internal( let first_leaf_id = num_bodies - 1; let mut tree = SliceMut(tree, root_id(colliders_start) as usize); - - // All threads must execute the same number of outer loop iterations for uniform control flow. - // NOTE: we calculate the interation count based on `colliders_batch_capacity` instead of - // `num_bodies` since the latter is non-uniform because it originates from a storage - // buffer. - let num_iterations = batch_ids.colliders_batch_capacity.div_ceil(num_threads); + let num_iterations = num_bodies.div_ceil(num_threads); // NOTE: using unchecked indexing (via MaybeIndexUnchecked) because otherwise the bounds // checking inserted by rustgpu breaks the shader when targeting some NVidia graphics @@ -449,8 +444,10 @@ pub fn gpu_lbvh_refit( // Propagate to ancestors. let mut curr_id = tree.at(curr_leaf_id as usize).parent; - loop { - let refit_count = atomic_add_u32(&mut tree.at_mut(curr_id as usize).refit_count, 1); + // NOTE: bounded `for` (tree depth <= 32 in practice) instead of `loop` + // to avoid the MacOS miscompilation bug. + for _ in 0..32u32 { + let refit_count = atomic_add_u32(&mut tree.at_mut(curr_id as usize).refit_count_or_max_subtree_index, 1); if refit_count == 0 { // If `refit_count` was 0 then the other thread hasn't reached this node From b411384316421e820976a9e71a2bc2694bb8d940 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Sat, 25 Jul 2026 11:25:45 +0200 Subject: [PATCH 12/39] =?UTF-8?q?perf:=20multibody=20=E2=80=94=20uniform?= =?UTF-8?q?=20loop=20bounds=20+=20lane-parallel=20contact=20finalize?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../dynamics/multibody/multibody_from_rapier.rs | 7 +++++++ src_rbd/dynamics/multibody/multibody_set.rs | 8 ++++++++ src_rbd/dynamics/multibody/multibody_solver.rs | 6 +++++- .../dynamics/multibody/compute_dynamics_pre.rs | 15 +++++++-------- .../dynamics/multibody/contact_constraints.rs | 14 +++++++++----- .../dynamics/multibody/gravity_and_lu.rs | 9 +++++---- src_rbd_shaders/dynamics/multibody/lu.rs | 14 +++++--------- src_rbd_shaders/utils/indices.rs | 5 +++++ 8 files changed, 51 insertions(+), 27 deletions(-) diff --git a/src_rbd/dynamics/multibody/multibody_from_rapier.rs b/src_rbd/dynamics/multibody/multibody_from_rapier.rs index 7398244..69e7379 100644 --- a/src_rbd/dynamics/multibody/multibody_from_rapier.rs +++ b/src_rbd/dynamics/multibody/multibody_from_rapier.rs @@ -48,6 +48,9 @@ impl GpuMultibodySet { let mut global_max_mb = 0u32; let mut global_max_links = 0u32; + // Per-multibody maxima (not per-env sums) for the uniform loop bounds. + let mut max_mb_ndofs = 0u32; + let mut max_mb_links = 0u32; let mut global_max_dofs = 0u32; let mut global_max_jac = 0u32; let mut global_max_mm = 0u32; @@ -99,6 +102,8 @@ impl GpuMultibodySet { }; let ndofs = mb.ndofs() as u32 - root_ndof_adjust; let num_links = mb.num_links() as u32; + max_mb_ndofs = max_mb_ndofs.max(ndofs); + max_mb_links = max_mb_links.max(num_links); // Count maximum constraint slots this multibody could need: for // each non-root non-kinematic joint, every free axis with a limit @@ -507,6 +512,8 @@ impl GpuMultibodySet { .unwrap(), mb_imp_joint_num_colors: 0, mb_imp_joint_max_color_group_len: 0, + max_ndofs: max_mb_ndofs, + max_links: max_mb_links, joint_constraints_per_batch: cons_cap, joint_constraint_columns_per_batch: cons_col_cap, contact_constraints_per_batch: contact_cons_cap, diff --git a/src_rbd/dynamics/multibody/multibody_set.rs b/src_rbd/dynamics/multibody/multibody_set.rs index c6bfd40..2a24546 100644 --- a/src_rbd/dynamics/multibody/multibody_set.rs +++ b/src_rbd/dynamics/multibody/multibody_set.rs @@ -115,6 +115,12 @@ pub struct GpuMultibodySet { /// Number of colors (per-batch stride of `mb_imp_joint_color_groups`, /// and the host color-loop trip count). CPU mirror. pub(crate) mb_imp_joint_num_colors: u32, + /// Max `ndofs` across every multibody in every batch (CPU mirror of + /// `BatchIndices::mb_max_ndofs`). + pub(super) max_ndofs: u32, + /// Max link count across every multibody in every batch (CPU mirror of + /// `BatchIndices::mb_max_links`). + pub(super) max_links: u32, /// Largest color group across batches — the per-color dispatch width. pub(super) mb_imp_joint_max_color_group_len: u32, /// Per-batch capacities of the joint / contact constraint slabs (CPU-side @@ -296,6 +302,8 @@ impl GpuMultibodySet { dst.mb_imp_joint_constraints_batch_capacity = self.mb_imp_joint_constraints_per_batch; dst.mb_imp_joint_jacobians_batch_capacity = self.mb_imp_joint_jacobians_per_batch; dst.mb_imp_joint_color_groups_batch_capacity = self.mb_imp_joint_num_colors.max(1); + dst.mb_max_ndofs = self.max_ndofs; + dst.mb_max_links = self.max_links; dst.coriolis_w_section_offset = self.coriolis_entries_per_batch * self.num_batches; dst.i_coriolis_dt_section_offset = 2 * self.coriolis_entries_per_batch * self.num_batches; dst.dof_damping_section_offset = self.dofs_per_batch * self.num_batches; diff --git a/src_rbd/dynamics/multibody/multibody_solver.rs b/src_rbd/dynamics/multibody/multibody_solver.rs index a255a0a..b3a9940 100644 --- a/src_rbd/dynamics/multibody/multibody_solver.rs +++ b/src_rbd/dynamics/multibody/multibody_solver.rs @@ -253,9 +253,13 @@ impl GpuMultibodySolver { args.contacts, )?; + // One 64-lane workgroup per multibody: the per-constraint LU + // back-solves are independent, so they run one-per-lane instead of + // sequentially on a single thread. + let finalize_dispatch = [mb.multibodies_per_batch * MB_LU_LANES, mb.num_batches, 1]; self.finalize_contact_constraints.call( pass, - dispatch, + finalize_dispatch, &mb.multibody_info, &mb.mass_matrices, &mb.lu_pivots, diff --git a/src_rbd_shaders/dynamics/multibody/compute_dynamics_pre.rs b/src_rbd_shaders/dynamics/multibody/compute_dynamics_pre.rs index a15512f..b3f43de 100644 --- a/src_rbd_shaders/dynamics/multibody/compute_dynamics_pre.rs +++ b/src_rbd_shaders/dynamics/multibody/compute_dynamics_pre.rs @@ -107,6 +107,7 @@ pub fn gpu_mb_compute_dynamics_pre( mb_jac_base, ndofs, num_links, + batch_ids.mb_max_links, &stat_slice, &ws_slice.as_ref(), body_jacobians, @@ -128,9 +129,7 @@ pub fn gpu_mb_compute_dynamics_pre( workgroup_memory_barrier_with_group_sync(); - // NOTE: fixed number of iterations for uniform control flow. - // TODO(PERF): on non-web platforms we could just use `mb.num_links` as the upper bound. - for k in 0..MAX_MB_DOFS as u32 { + for k in 0..batch_ids.mb_max_links { let loop_is_active = k < num_links; let mut inv_mass_x = 0.0; let mut mass = 0.0; @@ -545,6 +544,7 @@ pub fn gpu_mb_compute_dynamics_without_coriolis_pre( mb_jac_base, ndofs, num_links, + batch_ids.mb_max_links, &stat_slice, &ws_slice.as_ref(), body_jacobians, @@ -561,9 +561,8 @@ pub fn gpu_mb_compute_dynamics_without_coriolis_pre( fill_par(mass_matrices, acc_augmented_mass, 0.0, lane, LANES); workgroup_memory_barrier_with_group_sync(); - // NOTE: fixed number of iterations for uniform control flow. - // TODO(PERF): on non-web platforms we could just use `num_links` as the upper bound. - for k in 0..MAX_MB_DOFS as u32 { + // NOTE: uniform trip count (from the `BatchIndices` uniform). + for k in 0..batch_ids.mb_max_links { let mut active = k < num_links; if active { let lmp = stat_slice[k as usize].local_mprops; @@ -721,16 +720,16 @@ fn update_body_jacobians( mb_jac_base: usize, ndofs: u32, num_links: u32, + max_links: u32, stat_slice: &Slice, ws_slice: &Slice, body_jacobians: &mut [f32], ) { - // TODO(PERF): on non-web platforms we could just use `mb.num_links` as the upper bound. // TODO(PERF): instead of copying the body jacobian over and over for each body, we should // precompute a bit set that indicates which dofs are part of the kinematic tree // of each node. For a max number of DOFs set to 32, this means a single addition 32-bits // value per node. - for k in 0..MAX_MB_DOFS as u32 { + for k in 0..max_links { let mut parent_to_world = Pose::default(); let link_j = MatSlice::dense( mb_jac_base + (k as usize) * SPATIAL_DIM * (ndofs as usize), diff --git a/src_rbd_shaders/dynamics/multibody/contact_constraints.rs b/src_rbd_shaders/dynamics/multibody/contact_constraints.rs index 922015a..6b5ddee 100644 --- a/src_rbd_shaders/dynamics/multibody/contact_constraints.rs +++ b/src_rbd_shaders/dynamics/multibody/contact_constraints.rs @@ -15,6 +15,7 @@ use khal_std::glamx::UVec3; use khal_std::index::MaybeIndexUnchecked; +use khal_std::iter::StepRng; use khal_std::macros::{spirv, spirv_bindgen}; use crate::dynamics::ConstraintSoftness; @@ -644,9 +645,10 @@ pub fn gpu_mb_warmstart_contact_constraints( /// (the row produced by the init kernel) and set `inv_lhs = 1 / (Jᵀ · /// column + free_body_inv_r)`. #[spirv_bindgen] -#[spirv(compute(threads(1)))] +#[spirv(compute(threads(64)))] pub fn gpu_mb_finalize_contact_constraints( - #[spirv(global_invocation_id)] invocation_id: UVec3, + #[spirv(workgroup_id)] workgroup_id: UVec3, + #[spirv(local_invocation_id)] local_id: UVec3, #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] multibody_info: &[MultibodyInfo], #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] mass_matrices: &[f32], #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] lu_pivots: &[u32], @@ -657,8 +659,10 @@ pub fn gpu_mb_finalize_contact_constraints( contact_constraint_columns: &mut [f32], #[spirv(uniform, descriptor_set = 0, binding = 6)] batch_ids: &BatchIndices, ) { - let batch_id = invocation_id.y; - let mb_idx = invocation_id.x; + const LANES: u32 = 64; + let batch_id = workgroup_id.y; + let mb_idx = workgroup_id.x; + let lane = local_id.x; let num_mb = batch_ids.multibodies_len; if mb_idx >= num_mb { return; @@ -683,7 +687,7 @@ pub fn gpu_mb_finalize_contact_constraints( let m = MatSlice::dense(mb_mm_base, ndofs, ndofs); let count = mb.contact_constraint_count; - for s in 0..count { + for s in StepRng::new(lane..count, LANES) { let col_offset = col_base + (s as usize) * dofs_stride; // 1) Copy J^T row into the column buffer (it'll be overwritten by the // LU solve with the M⁻¹·Jᵀ result). diff --git a/src_rbd_shaders/dynamics/multibody/gravity_and_lu.rs b/src_rbd_shaders/dynamics/multibody/gravity_and_lu.rs index b2eeb78..52c0e7a 100644 --- a/src_rbd_shaders/dynamics/multibody/gravity_and_lu.rs +++ b/src_rbd_shaders/dynamics/multibody/gravity_and_lu.rs @@ -57,6 +57,8 @@ pub fn gpu_mb_gravity_and_lu( let batch_id = wg_id.y; let mb_idx = wg_id.x; let lane = lid.x; + let max_ndofs = batch_ids.mb_max_ndofs; + let max_links = batch_ids.mb_max_links; let mb = batch_ids .mb_batch(batch_id, multibody_info) @@ -98,9 +100,7 @@ pub fn gpu_mb_gravity_and_lu( let g = Vec2::new(gravity.x, gravity.y); // ---- Phase 2: per-link gravity / Coriolis-force assembly. ---- - // NOTE: fixed number of iterations for uniform control flow. - // TODO(PERF): on non-web platforms we could just use `num_links` as the upper bound. - for k in 0..MAX_MB_DOFS as u32 { + for k in 0..max_links { let active = k < num_links; let mut acc_lin = Vector::ZERO; #[cfg(feature = "dim3")] @@ -239,6 +239,7 @@ pub fn gpu_mb_gravity_and_lu( lu_factor_in_shared( ndofs, + max_ndofs, lane, mat, lu_pivots, @@ -257,7 +258,7 @@ pub fn gpu_mb_gravity_and_lu( // ---- Phase 4: solve M·x = τ for the gravity rhs. ---- lu_apply_pivots(ndofs, lane, lu_pivots, piv_offset, x); - lu_triangular_solve_in_place(ndofs, lane, mat, x, partial); + lu_triangular_solve_in_place(ndofs, max_ndofs, lane, mat, x, partial); if lane < ndofs { gen_forces.write(rhs_offset + lane as usize, x.read(lane as usize)); diff --git a/src_rbd_shaders/dynamics/multibody/lu.rs b/src_rbd_shaders/dynamics/multibody/lu.rs index af33947..d6c9909 100644 --- a/src_rbd_shaders/dynamics/multibody/lu.rs +++ b/src_rbd_shaders/dynamics/multibody/lu.rs @@ -27,6 +27,7 @@ pub(super) fn sm_idx(r: u32, c: u32) -> usize { #[inline] pub(super) fn lu_factor_in_shared( n: u32, + max_n: u32, lane: u32, mat: &mut [f32; MAX_MB_DOFS * MAX_MB_DOFS], pivots_dst: &mut [u32], @@ -34,9 +35,7 @@ pub(super) fn lu_factor_in_shared( pivot_row_shared: &mut u32, inv_akk_shared: &mut f32, ) { - // NOTE: fixed number of iterations for uniform control flow. - // TODO(PERF): on non-web platforms we could just use `n` as the upper bound. - for k in 0..MAX_DOFS_U32 { + for k in 0..max_n { let active = k < n; if active && lane == 0 { let mut pivot_row = k; @@ -104,14 +103,13 @@ pub(super) fn lu_factor_in_shared( #[inline] pub(super) fn lu_triangular_solve_in_place( n: u32, + max_n: u32, lane: u32, mat: &[f32; MAX_MB_DOFS * MAX_MB_DOFS], x: &mut [f32; MAX_MB_DOFS], partial: &mut [f32; LANES as usize], ) { - // NOTE: fixed number of iterations for uniform control flow. - // TODO(PERF): on non-web platforms we could just use `n` as the upper bound. - for i in 0..MAX_DOFS_U32 { + for i in 0..max_n { let active = i < n; let s = if active && lane < i { mat.read(sm_idx(i, lane)) * x.read(lane as usize) @@ -135,9 +133,7 @@ pub(super) fn lu_triangular_solve_in_place( workgroup_memory_barrier_with_group_sync(); } - // NOTE: fixed number of iterations for uniform control flow. - // TODO(PERF): on non-web platforms we could just use `n` as the upper bound. - for step in 0..MAX_DOFS_U32 { + for step in 0..max_n { let active = step < n; // For dummy iterations (step >= n), `i` is not meaningful — guard // every use of it behind `active`. diff --git a/src_rbd_shaders/utils/indices.rs b/src_rbd_shaders/utils/indices.rs index 19623ef..ecac949 100644 --- a/src_rbd_shaders/utils/indices.rs +++ b/src_rbd_shaders/utils/indices.rs @@ -53,6 +53,11 @@ pub struct BatchIndices { /// contrast, are stored single-batch (identical coloring across batches) /// and read at offset 0. pub mb_imp_joint_color_groups_batch_capacity: u32, + /// Actual max `ndofs` across every multibody in every batch (often smaller + /// than the fixed `MAX_MB_DOFS` limit). + pub mb_max_ndofs: u32, + /// Actual max link count across every multibody in every batch. + pub mb_max_links: u32, /// Per-batch stride of the contact-solver color-bucket buffers /// (`color_counts` / `color_starts` / `color_cursors`), = `max_colors + 3` /// so that `starts[c + 1]` is in bounds for every swept color. From 5bc7b9cbfe9c6e60bb36effeac19675089e31a1e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Sat, 25 Jul 2026 11:43:58 +0200 Subject: [PATCH 13/39] perf: misc rbd pipeline cleanup --- src_rbd/dynamics/multibody/multibody_set.rs | 6 +- src_rbd/dynamics/solver.rs | 2 - src_rbd/pipeline/insertion_removal.rs | 98 ++++++++++++--------- src_rbd/pipeline/rbd_state.rs | 10 +-- src_rbd/pipeline/rbd_state_from_rapier.rs | 1 - src_rbd/pipeline/rbd_step.rs | 2 - 6 files changed, 66 insertions(+), 53 deletions(-) diff --git a/src_rbd/dynamics/multibody/multibody_set.rs b/src_rbd/dynamics/multibody/multibody_set.rs index 2a24546..b559ba7 100644 --- a/src_rbd/dynamics/multibody/multibody_set.rs +++ b/src_rbd/dynamics/multibody/multibody_set.rs @@ -154,8 +154,12 @@ impl GpuMultibodySet { } /// True if the set contains no multibodies in any batch. + /// + /// Uses the *active* count: the per-batch capacity is padded to >= 1 to + /// avoid zero-sized buffers, so testing it would run the whole multibody + /// kernel chain every step for scenes without any multibody. pub fn is_empty(&self) -> bool { - self.multibodies_per_batch == 0 || self.links_per_batch == 0 + self.num_active_multibodies == 0 || self.links_per_batch == 0 } /// Number of colors used by the colored multibody impulse-joint sweeps. diff --git a/src_rbd/dynamics/solver.rs b/src_rbd/dynamics/solver.rs index 7c06322..fa1c2d3 100644 --- a/src_rbd/dynamics/solver.rs +++ b/src_rbd/dynamics/solver.rs @@ -92,8 +92,6 @@ pub struct SolverArgs<'a> { pub vels: &'a mut Tensor, /// Solver working velocities. pub solver_vels: &'a mut Tensor, - /// Solver output velocities (currently unused). - pub solver_vels_out: &'a Tensor, /// Accumulated velocity increments during substeps. pub solver_vels_inc: &'a mut Tensor, /// World-space mass properties. diff --git a/src_rbd/pipeline/insertion_removal.rs b/src_rbd/pipeline/insertion_removal.rs index 8933ea8..7520d96 100644 --- a/src_rbd/pipeline/insertion_removal.rs +++ b/src_rbd/pipeline/insertion_removal.rs @@ -244,7 +244,6 @@ impl RbdState { sim_params: Tensor::vector(backend, &all_sim_params, BufferUsages::STORAGE).unwrap(), vels: Tensor::vector(backend, &all_vels, rw).unwrap(), solver_vels: Tensor::vector(backend, &all_vels, storage).unwrap(), - solver_vels_out: Tensor::vector(backend, &all_vels, storage).unwrap(), solver_vels_inc: Tensor::vector(backend, &all_vels, storage).unwrap(), joints, #[cfg(feature = "dim3")] @@ -482,6 +481,23 @@ impl RbdState { crate::rapier::geometry::InteractionTestMode::And, ); + let staging_usages = + BufferUsages::STORAGE | BufferUsages::COPY_SRC | BufferUsages::COPY_DST; + let mut enc = backend.begin_encoding(); + let mut any_copy = false; + let mut staging_pose = backend.uninit_buffer::(1, staging_usages)?; + let mut staging_local_mprops = + backend.uninit_buffer::(1, staging_usages)?; + let mut staging_mprops = backend.uninit_buffer::(1, staging_usages)?; + let mut staging_vels = backend.uninit_buffer::(1, staging_usages)?; + let mut staging_shapes = backend.uninit_buffer::(1, staging_usages)?; + let mut staging_groups = backend + .uninit_buffer::(1, staging_usages)?; + let mut staging_materials = + backend.uninit_buffer::(1, staging_usages)?; + // Deferred `(buffer-kind, slot)` neutralisation writes. + let mut neutralize: Vec = Vec::new(); + for local in locals { let active = self.num_active_colliders as usize; if active == 0 || local >= active { @@ -497,61 +513,37 @@ impl RbdState { // Relocate the last active body into the freed slot. A staging // buffer is used to avoid same-buffer overlapping copies. macro_rules! relocate { - ($t:expr) => {{ - let mut staging = backend.uninit_buffer( - 1, - BufferUsages::STORAGE - | BufferUsages::COPY_SRC - | BufferUsages::COPY_DST, - )?; - let mut enc = backend.begin_encoding(); + ($t:expr, $staging:expr) => {{ enc.copy_buffer_to_buffer( $t.buffer(), last_global, - &mut staging, + &mut $staging, 0, 1, )?; enc.copy_buffer_to_buffer( - &staging, + &$staging, 0, $t.buffer_mut(), hole_global, 1, )?; - backend.submit(enc)?; + any_copy = true; }}; } - relocate!(self.body_poses); - relocate!(self.solver_body_poses); - relocate!(self.collider_world_poses); - relocate!(self.collider_local_poses); - relocate!(self.local_mprops); - relocate!(self.mprops); - relocate!(self.vels); - relocate!(self.shapes); - relocate!(self.collision_groups); - relocate!(self.collider_materials); + relocate!(self.body_poses, staging_pose); + relocate!(self.solver_body_poses, staging_pose); + relocate!(self.collider_world_poses, staging_pose); + relocate!(self.collider_local_poses, staging_pose); + relocate!(self.local_mprops, staging_local_mprops); + relocate!(self.mprops, staging_mprops); + relocate!(self.vels, staging_vels); + relocate!(self.shapes, staging_shapes); + relocate!(self.collision_groups, staging_groups); + relocate!(self.collider_materials, staging_materials); } - // The now-topmost slot becomes inactive padding: neutralize it so - // it never participates in collisions even if a kernel scans up - // to the per-batch capacity. - backend.write_buffer( - self.collision_groups.buffer_mut(), - last_global as u64, - &[none_groups], - )?; - backend.write_buffer( - self.local_mprops.buffer_mut(), - last_global as u64, - &[GpuLocalMassProperties::default()], - )?; - backend.write_buffer( - self.mprops.buffer_mut(), - last_global as u64, - &[GpuWorldMassProperties::default()], - )?; + neutralize.push(last_global); } if local != last { @@ -560,6 +552,32 @@ impl RbdState { self.num_active_colliders = (active - 1) as u32; } + if any_copy { + backend.submit(enc)?; + } + + // The now-topmost slots become inactive padding: neutralize them so + // they never participate in collisions even if a kernel scans up to + // the per-batch capacity. Done after the relocation submit so the + // copies read the pre-neutralisation data. + for last_global in neutralize { + backend.write_buffer( + self.collision_groups.buffer_mut(), + last_global as u64, + &[none_groups], + )?; + backend.write_buffer( + self.local_mprops.buffer_mut(), + last_global as u64, + &[GpuLocalMassProperties::default()], + )?; + backend.write_buffer( + self.mprops.buffer_mut(), + last_global as u64, + &[GpuWorldMassProperties::default()], + )?; + } + // `collider_parent` is the identity mapping on the incremental (one // collider per body) path and stays identity under swap-remove, so it // needs no relocation; only the active body count tracks the colliders. diff --git a/src_rbd/pipeline/rbd_state.rs b/src_rbd/pipeline/rbd_state.rs index 60b7ae2..b464963 100644 --- a/src_rbd/pipeline/rbd_state.rs +++ b/src_rbd/pipeline/rbd_state.rs @@ -137,7 +137,6 @@ pub struct RbdState { pub(super) mprops: Tensor, pub(super) vels: Tensor, pub(super) solver_vels: Tensor, - pub(super) solver_vels_out: Tensor, pub(super) solver_vels_inc: Tensor, pub(super) vertex_buffers: Tensor, pub(super) index_buffers: Tensor, @@ -261,12 +260,9 @@ impl RbdState { }; #[cfg(feature = "dim3")] self.multibodies.fill_batch_indices(&mut bi); - self.batch_indices = Tensor::scalar( - backend, - bi, - BufferUsages::STORAGE | BufferUsages::UNIFORM | BufferUsages::COPY_DST, - ) - .unwrap(); + backend + .write_buffer(self.batch_indices.buffer_mut(), 0, &[bi]) + .unwrap(); } /// Shared per-batch index uniform — see `Self::rebuild_batch_indices`. diff --git a/src_rbd/pipeline/rbd_state_from_rapier.rs b/src_rbd/pipeline/rbd_state_from_rapier.rs index 24fbaf8..962f3c7 100644 --- a/src_rbd/pipeline/rbd_state_from_rapier.rs +++ b/src_rbd/pipeline/rbd_state_from_rapier.rs @@ -718,7 +718,6 @@ impl RbdState { sim_params: Tensor::vector(backend, &all_sim_params, BufferUsages::STORAGE).unwrap(), vels: Tensor::vector(backend, &all_vels, storage).unwrap(), solver_vels: Tensor::vector(backend, &all_vels, storage).unwrap(), - solver_vels_out: Tensor::vector(backend, &all_vels, storage).unwrap(), solver_vels_inc: Tensor::vector(backend, &all_vels, storage).unwrap(), joints, #[cfg(feature = "dim3")] diff --git a/src_rbd/pipeline/rbd_step.rs b/src_rbd/pipeline/rbd_step.rs index f01a705..2bc2938 100644 --- a/src_rbd/pipeline/rbd_step.rs +++ b/src_rbd/pipeline/rbd_step.rs @@ -237,7 +237,6 @@ impl RbdPipeline { collider_world_poses: &state.collider_world_poses, vels: &mut state.vels, solver_vels: &mut state.solver_vels, - solver_vels_out: &state.solver_vels_out, solver_vels_inc: &mut state.solver_vels_inc, mprops: &state.mprops, local_mprops: &state.local_mprops, @@ -372,7 +371,6 @@ impl RbdPipeline { collider_world_poses: &state.collider_world_poses, vels: &mut state.vels, solver_vels: &mut state.solver_vels, - solver_vels_out: &state.solver_vels_out, solver_vels_inc: &mut state.solver_vels_inc, mprops: &state.mprops, local_mprops: &state.local_mprops, From d6d0b0447d1c52407115ee50b4cf6df163270a7d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Sat, 25 Jul 2026 13:16:56 +0200 Subject: [PATCH 14/39] perf: use the gather (colorless) warmstart when no multibody is present --- src_rbd/dynamics/solver.rs | 40 ++++++++++++++++++++++++++++-------- src_rbd/pipeline/rbd_step.rs | 7 +++++++ 2 files changed, 39 insertions(+), 8 deletions(-) diff --git a/src_rbd/dynamics/solver.rs b/src_rbd/dynamics/solver.rs index fa1c2d3..26d1ff6 100644 --- a/src_rbd/dynamics/solver.rs +++ b/src_rbd/dynamics/solver.rs @@ -36,7 +36,10 @@ pub struct GpuSolver { /// Clears solver velocities and constraint counts. cleanup: GpuSolverCleanup, /// Applies warmstart impulses from previous frame. + #[allow(dead_code)] warmstart: GpuWarmstart, + /// Applies warmstart impulses from previous frame, without relying on graph coloring. + warmstart_without_colors: GpuWarmstartWithoutColors, /// Gauss-Seidel iteration step (sequential per color). step_gauss_seidel: GpuStepGaussSeidel, /// Initializes solver velocity increments. @@ -121,6 +124,12 @@ pub struct SolverArgs<'a> { pub num_solver_iterations: u32, /// Per-body graph-coloring group id (multibody-aware). pub body_group: &'a Tensor, + /// When `true` (no multibody in the scene), warmstart uses the single + /// gather-per-body dispatch instead of one scatter dispatch per color. + /// The gather variant looks bodies up by their own id, which is only + /// correct when `body_group` is the identity (multibody constraints are + /// counted on their root's slot with link-id constraint sides). + pub colorless_warmstart: bool, /// Shared per-batch capacity / section-offset uniform — see /// [`crate::shaders::utils::BatchIndices`]. Consumed by the (refactored) /// multibody kernels via `MultibodySolverArgs::batch_indices`; the RBD @@ -315,22 +324,37 @@ impl GpuSolver { args.batch_indices, )?; joint_solver.update(pass, &mut joint_args, args.solver_body_poses)?; - // NOTE: contact colors start at 1 (0 = unassigned). - for c in 1..=args.num_colors { - self.warmstart.call( + if args.colorless_warmstart { + // One gather dispatch over bodies instead of `num_colors` + // scatter dispatches (each constraint is visited once per + // body side, but the dispatch count drops by ~num_colors). + self.warmstart_without_colors.call( pass, - args.contacts_len_indirect, + [args.num_colliders, args.num_batches, 1], + args.body_constraint_counts, + args.body_constraint_ids, args.constraints, args.solver_vels, - args.color_bucket_starts, - args.color_sorted_ids, - &args.color_uniforms[c as usize], args.batch_indices, )?; + } else { + // NOTE: contact colors start at 1 (0 = unassigned). + for c in 1..=args.num_colors { + self.warmstart.call( + pass, + args.contacts_len_indirect, + args.constraints, + args.solver_vels, + args.color_bucket_starts, + args.color_sorted_ids, + &args.color_uniforms[c as usize], + args.batch_indices, + )?; + } } /* - * P3/F3 — solve ALL joints + contacts WITH bias. + * Solve all joints + contacts with bias. */ mb_phase!(substep_solve_with_bias); joint_solver.solve(pass, &mut joint_args, args.solver_vels, true)?; diff --git a/src_rbd/pipeline/rbd_step.rs b/src_rbd/pipeline/rbd_step.rs index 2bc2938..91f35ce 100644 --- a/src_rbd/pipeline/rbd_step.rs +++ b/src_rbd/pipeline/rbd_step.rs @@ -252,6 +252,7 @@ impl RbdPipeline { num_solver_iterations: state.num_solver_iterations, body_group: &state.body_group, batch_indices: &state.batch_indices, + colorless_warmstart: false, }; self.solver.prepare( backend, @@ -386,6 +387,12 @@ impl RbdPipeline { num_solver_iterations: state.num_solver_iterations, body_group: &state.body_group, batch_indices: &state.batch_indices, + // The gather warmstart is only valid without multibody grouping — + // see `SolverArgs::colorless_warmstart`. + #[cfg(feature = "dim3")] + colorless_warmstart: state.multibodies.is_empty(), + #[cfg(not(feature = "dim3"))] + colorless_warmstart: true, }; // Phase 3: Solve constraints From 23d6d7d8eb57315b5ce578cd9a07e19fecfa3ac9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Sat, 25 Jul 2026 13:57:00 +0200 Subject: [PATCH 15/39] perf: prune whole LBVH subtrees in the pair traversal --- src_rbd_shaders/broad_phase/lbvh.rs | 35 ++++++++++++++++++++--------- 1 file changed, 25 insertions(+), 10 deletions(-) diff --git a/src_rbd_shaders/broad_phase/lbvh.rs b/src_rbd_shaders/broad_phase/lbvh.rs index 1d9455a..6794f86 100644 --- a/src_rbd_shaders/broad_phase/lbvh.rs +++ b/src_rbd_shaders/broad_phase/lbvh.rs @@ -41,8 +41,12 @@ pub struct LbvhNode { pub right: u32, /// Parent node index. pub parent: u32, - /// Counter for bottom-up refitting (0, 1, or 2). - pub refit_count: u32, + /// During refit: bottom-up arrival counter; each thread atomically + /// increments it and only the second one continues upward (both children + /// ready). After refit, `refit_leaves` sets it to the maximum sorted leaf + /// index in this node’s subtree, which the pair traversal uses to prune + /// subtrees that can only produce duplicate pairs. + pub refit_count_or_max_subtree_index: u32, } /// Resets the collision pairs counter. @@ -319,6 +323,9 @@ pub fn gpu_lbvh_refit_leaves( tree.at_mut(curr_leaf_id as usize).aabb = leaf_shape.compute_aabb(leaf_pose, vertices); tree.at_mut(curr_leaf_id as usize).left = leaf_collider; + // For leaves, we can set their index here. They don’t use the `refit_count` + // mechanism which is for internal nodes only. + tree.at_mut(curr_leaf_id as usize).refit_count_or_max_subtree_index = i; } } @@ -380,6 +387,11 @@ pub fn gpu_lbvh_refit_internal( let right = tree.at(right_idx as usize).aabb; tree.at_mut(curr_id as usize).aabb = left.merged(&right); + // Set `refit_count_or_max_subtree_index` to the max subtree leaf index. + let max_l = tree.at(left_idx as usize).refit_count_or_max_subtree_index; + let max_r = tree.at(right_idx as usize).refit_count_or_max_subtree_index; + tree.at_mut(curr_id as usize).refit_count_or_max_subtree_index = max_l.max(max_r); + if curr_id == 0 { // We reached the root, can't go higher. thread_is_active = false; @@ -539,8 +551,11 @@ pub fn gpu_lbvh_find_collision_pairs( continue; } - // NOTE: we don't have to compare i < j to avoid duplicates since that comparison already happened - // alongside the AABB check. + // Duplicates were already pruned during the descent (sorted + // leaf-index comparison). Emit the pair in ascending collider + // order so the narrow phase / warmstart see a stable ordering + // regardless of the traversal's dedup basis. + let (ci, cj) = if i < j { (i, j) } else { (j, i) }; let target_pair_index = atomic_add_u32(collision_pairs_len.at_mut(batch_id as usize), 1); @@ -555,16 +570,17 @@ pub fn gpu_lbvh_find_collision_pairs( // intermediate pfm-pair buffer) narrow, and keeping // `collider_parent` out of the broad phase entirely. collision_pairs[target_pair_index as usize] = CollisionPair { - colliders: UVec2::new(i, j), + colliders: UVec2::new(ci, cj), }; } } else { let left = node.left; let right = node.right; - // Go on the child only if the AABB intersects and either the child isn't a leaf, or it is a leaf with associated collider - // smaller than `i` (to avoid duplicate pairs). - if (left < first_leaf_id || i < tree.at(left as usize).left) + // Descend only if the subtree contains leaf id smaller than + // `leaf_i`. That way we prune the part of the tree that’s + // on the "left" of `leaf_i`, avoiding duplicate pairs/traversals. + if leaf_i < tree.at(left as usize).refit_count_or_max_subtree_index && aabb1.intersects(&tree.at(left as usize).aabb) && stack_len < 64 { @@ -572,8 +588,7 @@ pub fn gpu_lbvh_find_collision_pairs( stack_len += 1; } - // NOTE: on leaves (including tree[right]), the collider id is stored as the left child index. - if (right < first_leaf_id || i < tree.at(right as usize).left) + if leaf_i < tree.at(right as usize).refit_count_or_max_subtree_index && aabb1.intersects(&tree.at(right as usize).aabb) && stack_len < 64 { From 793b47df033ae128d50915a021621e92a1327b06 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Sun, 26 Jul 2026 10:10:22 +0200 Subject: [PATCH 16/39] chore: drop unused MAX_MB_DOFS import --- src_rbd_shaders/dynamics/multibody/compute_dynamics_pre.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src_rbd_shaders/dynamics/multibody/compute_dynamics_pre.rs b/src_rbd_shaders/dynamics/multibody/compute_dynamics_pre.rs index b3f43de..07233bc 100644 --- a/src_rbd_shaders/dynamics/multibody/compute_dynamics_pre.rs +++ b/src_rbd_shaders/dynamics/multibody/compute_dynamics_pre.rs @@ -19,7 +19,7 @@ use crate::dynamics::joint::SPATIAL_DIM; #[cfg(feature = "dim3")] use crate::utils::linalg::gemm_skew_lhs_cross_buf_par; use crate::utils::linalg::{ - MAX_MB_DOFS, MatSlice, copy_from_par, fill_par, gemm_inertia_lhs_par, + MatSlice, copy_from_par, fill_par, gemm_inertia_lhs_par, gemm_omega_skew_tr_cross_buf_par, gemm_skew_tr_lhs_cross_buf_par, gemm_skew_tr_lhs_par, gemm_tr_par, quadform_spatial_par, }; From a8bff4993fecacc4078786de8891001d7656ef34 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Sun, 26 Jul 2026 10:50:34 +0200 Subject: [PATCH 17/39] fix: count constraints over contacts_len, not the padded capacity --- src_rbd/dynamics/solver.rs | 1 + src_rbd_shaders/dynamics/solver.rs | 13 ++++--------- 2 files changed, 5 insertions(+), 9 deletions(-) diff --git a/src_rbd/dynamics/solver.rs b/src_rbd/dynamics/solver.rs index 26d1ff6..67c8b9c 100644 --- a/src_rbd/dynamics/solver.rs +++ b/src_rbd/dynamics/solver.rs @@ -194,6 +194,7 @@ impl GpuSolver { args.body_constraint_counts, args.body_group, args.mprops, + args.contacts_len, args.batch_indices, )?; diff --git a/src_rbd_shaders/dynamics/solver.rs b/src_rbd_shaders/dynamics/solver.rs index c751a89..e45e192 100644 --- a/src_rbd_shaders/dynamics/solver.rs +++ b/src_rbd_shaders/dynamics/solver.rs @@ -81,7 +81,8 @@ pub fn gpu_solver_count_constraints( #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] body_constraint_counts: &mut [u32], #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] body_group: &[u32], #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] mprops: &[WorldMassProperties], - #[spirv(uniform, descriptor_set = 0, binding = 4)] batch_ids: &BatchIndices, + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] contacts_len: &[u32], + #[spirv(uniform, descriptor_set = 0, binding = 5)] batch_ids: &BatchIndices, ) { let num_threads = num_workgroups.x * WORKGROUP_SIZE; let batch_id = invocation_id.y; @@ -90,16 +91,10 @@ pub fn gpu_solver_count_constraints( let mut body_constraint_counts = batch_ids.coll_batch_mut(batch_id, body_constraint_counts); let body_group = batch_ids.coll_batch(batch_id, body_group); let mprops = batch_ids.coll_batch(batch_id, mprops); - // See `gpu_solver_init_constraints` — the indirect grid bounds the active - // range much tighter than the capacity. - let cap = batch_ids.contacts_batch_capacity.min(num_threads); + let len = contacts_len.read(batch_id as usize); - for i in StepRng::new(invocation_id.x..cap, num_threads) { + for i in StepRng::new(invocation_id.x..len, num_threads) { let im = &contacts[i as usize]; - if im.contact.len == 0 { - continue; - } - let body1 = im.bodies.x; let body2 = im.bodies.y; let group1 = body_group[body1 as usize]; From e6e0fce3da04b1362a35ebe64a4501fd84b6e9f5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Sun, 26 Jul 2026 11:18:57 +0200 Subject: [PATCH 18/39] perf: pack per-multibody kernels into full workgroups --- src_rbd/dynamics/multibody/multibody_set.rs | 29 +- .../dynamics/multibody/multibody_solver.rs | 147 ++++---- src_rbd/pipeline/insertion_removal.rs | 1 + src_rbd/pipeline/rbd_state.rs | 1 + src_rbd/pipeline/rbd_state_from_rapier.rs | 1 + .../multibody/compute_dynamics_pre.rs | 124 ++++--- .../dynamics/multibody/contact_constraints.rs | 50 +-- .../dynamics/multibody/gravity_and_lu.rs | 318 +++++++++++++++++- .../dynamics/multibody/integrate.rs | 16 +- .../dynamics/multibody/joint_constraints.rs | 29 +- src_rbd_shaders/dynamics/multibody/lu.rs | 192 +++++++++++ src_rbd_shaders/utils/indices.rs | 6 + 12 files changed, 734 insertions(+), 180 deletions(-) diff --git a/src_rbd/dynamics/multibody/multibody_set.rs b/src_rbd/dynamics/multibody/multibody_set.rs index b559ba7..b7ad7c1 100644 --- a/src_rbd/dynamics/multibody/multibody_set.rs +++ b/src_rbd/dynamics/multibody/multibody_set.rs @@ -153,11 +153,29 @@ impl GpuMultibodySet { self.multibodies_per_batch } - /// True if the set contains no multibodies in any batch. - /// - /// Uses the *active* count: the per-batch capacity is padded to >= 1 to - /// avoid zero-sized buffers, so testing it would run the whole multibody - /// kernel chain every step for scenes without any multibody. + /// Thread-count grid for the per-multibody kernels, with `(multibody, + /// batch)` flattened into X. The kernels decode + /// `batch_id = x / multibodies_len`, `mb_idx = x % multibodies_len`. + pub(crate) fn flat_mb_dispatch(&self) -> [u32; 3] { + [self.num_active_multibodies * self.num_batches, 1, 1] + } + + /// Lanes per multibody for the packed per-multibody workgroup kernels — + /// mirrored into `BatchIndices::mb_pack_lanes`. + pub(crate) fn pack_lanes(&self) -> u32 { + self.max_ndofs.next_power_of_two().clamp(8, MB_LU_LANES) + } + + /// Thread-count grid for the packed per-multibody workgroup kernels + /// (`compute_dynamics_pre`, `gravity_and_lu`): `64 / pack_lanes` + /// multibodies per 64-lane workgroup, flattened `(multibody, batch)`. + pub(crate) fn packed_wg_dispatch(&self) -> [u32; 3] { + let slots = MB_LU_LANES / self.pack_lanes(); + let total = self.num_active_multibodies * self.num_batches; + [total.div_ceil(slots) * MB_LU_LANES, 1, 1] + } + + /// True if the set contains no active multibodies in any batch. pub fn is_empty(&self) -> bool { self.num_active_multibodies == 0 || self.links_per_batch == 0 } @@ -308,6 +326,7 @@ impl GpuMultibodySet { dst.mb_imp_joint_color_groups_batch_capacity = self.mb_imp_joint_num_colors.max(1); dst.mb_max_ndofs = self.max_ndofs; dst.mb_max_links = self.max_links; + dst.mb_pack_lanes = self.pack_lanes(); dst.coriolis_w_section_offset = self.coriolis_entries_per_batch * self.num_batches; dst.i_coriolis_dt_section_offset = 2 * self.coriolis_entries_per_batch * self.num_batches; dst.dof_damping_section_offset = self.dofs_per_batch * self.num_batches; diff --git a/src_rbd/dynamics/multibody/multibody_solver.rs b/src_rbd/dynamics/multibody/multibody_solver.rs index b3a9940..ba3f0c6 100644 --- a/src_rbd/dynamics/multibody/multibody_solver.rs +++ b/src_rbd/dynamics/multibody/multibody_solver.rs @@ -6,7 +6,8 @@ use crate::queries::GpuIndexedContact; use crate::shaders::dynamics::{ GpuMbComputeDynamicsPre, GpuMbComputeDynamicsWithoutCoriolisPre, - GpuMbFinalizeContactConstraints, GpuMbGravityAndLu, GpuMbInitContactConstraints, + GpuMbFinalizeContactConstraints, GpuMbGravityAndLu, GpuMbGravityAndLuT8, + GpuMbGravityAndLuT16, GpuMbGravityAndLuT32, GpuMbInitContactConstraints, GpuMbInitJointConstraints, GpuMbIntegrate, GpuMbIntegrateVelocities, GpuMbRemoveContactConstraintBias, GpuMbRemoveImpulseJointConstraintBias, GpuMbResetContactWarmstart, GpuMbWarmstartContactConstraints, @@ -23,6 +24,13 @@ use vortx::tensor::Tensor; #[derive(Shader)] pub struct GpuMultibodySolver { gravity_and_lu: GpuMbGravityAndLu, + /// Packed tiers of `gravity_and_lu` — `64/T` multibodies per workgroup + /// with a `T×T` shared tile each, selected by `max_ndofs`. The fallback + /// `gravity_and_lu` (one multibody per workgroup, 64×64 tile) only runs + /// for `max_ndofs > 32`. + gravity_and_lu_t8: GpuMbGravityAndLuT8, + gravity_and_lu_t16: GpuMbGravityAndLuT16, + gravity_and_lu_t32: GpuMbGravityAndLuT32, compute_dynamics_pre: GpuMbComputeDynamicsPre, compute_dynamics_without_coriolis_pre: GpuMbComputeDynamicsWithoutCoriolisPre, solve_joint_with_bias: GpuMbSolveJointConstraints, @@ -83,67 +91,11 @@ impl GpuMultibodySolver { mb: &mut GpuMultibodySet, args: MultibodySolverArgs<'_>, ) -> Result<(), GpuBackendError> { + let mut args = args; if mb.is_empty() { return Ok(()); } - // Fused FK + body-jacobians + velocity propagation + CRBA-with-Coriolis - // mass-matrix assembly (4 dispatches → 1) — see - // `gpu_mb_compute_dynamics_pre`. Only the implicit-Coriolis path is - // wired through the fused kernel; the explicit-Coriolis fallback keeps - // the legacy split path. - let pre_dispatch = [mb.multibodies_per_batch * MB_LU_LANES, mb.num_batches, 1]; - if mb.implicit_coriolis { - self.compute_dynamics_pre.call( - pass, - pre_dispatch, - &mb.multibody_info, - &mb.links_static, - &mut mb.links_workspace, - args.poses, - &mut mb.body_jacobians, - &mut mb.mass_matrices, - &mut mb.coriolis_packed, - &mb.dof_state, - &mb.dt, - args.batch_indices, - )?; - } else { - self.compute_dynamics_without_coriolis_pre.call( - pass, - pre_dispatch, - &mb.multibody_info, - &mb.links_static, - &mut mb.links_workspace, - args.poses, - &mut mb.body_jacobians, - &mut mb.mass_matrices, - &mb.dof_state, - &mb.dt, - args.batch_indices, - )?; - } - - // Fused: gravity / Coriolis force assembly + LU factor + LU solve in - // a single dispatch. Replaces the previous 2-dispatch chain - // (apply_gravity_with_coriolis → lu_factor_and_solve) — drops one - // WebGPU dispatch per `compute_dynamics` call. - let grav_lu_dispatch = [mb.multibodies_per_batch * MB_LU_LANES, mb.num_batches, 1]; - self.gravity_and_lu.call( - pass, - grav_lu_dispatch, - &mb.multibody_info, - &mb.links_static, - &mut mb.links_workspace, - &mb.body_jacobians, - &mut mb.gen_forces, - &mut mb.mass_matrices, - &mut mb.lu_pivots, - &mb.dof_state, - &mb.gravity, - args.batch_indices, - )?; - - Ok(()) + self.compute_dynamics(pass, mb, &mut args) } /// Once-per-visible-step setup. After this call, `gen_forces` holds the @@ -163,7 +115,7 @@ impl GpuMultibodySolver { // starts cold (within a frame they are then preserved across substeps). self.reset_contact_warmstart.call( pass, - [mb.multibodies_per_batch, mb.num_batches, 1], + mb.flat_mb_dispatch(), &mb.multibody_info, &mut mb.contact_constraints, args.batch_indices, @@ -186,7 +138,7 @@ impl GpuMultibodySolver { if mb.is_empty() { return Ok(()); } - let dispatch = [mb.multibodies_per_batch, mb.num_batches, 1]; + let dispatch = mb.flat_mb_dispatch(); self.integrate_velocities.call( pass, dispatch, @@ -209,7 +161,7 @@ impl GpuMultibodySolver { if mb.is_empty() { return Ok(()); } - let dispatch = [mb.multibodies_per_batch, mb.num_batches, 1]; + let dispatch = mb.flat_mb_dispatch(); if mb.has_joint_constraints { // TODO(PERF): joints init could parallelized. We either need to rework @@ -299,7 +251,7 @@ impl GpuMultibodySolver { if mb.is_empty() { return Ok(()); } - let dispatch = [mb.multibodies_per_batch, mb.num_batches, 1]; + let dispatch = mb.flat_mb_dispatch(); if mb.has_joint_constraints { self.solve_joint_with_bias.call( @@ -400,7 +352,7 @@ impl GpuMultibodySolver { if mb.is_empty() { return Ok(()); } - let dispatch = [mb.multibodies_per_batch, mb.num_batches, 1]; + let dispatch = mb.flat_mb_dispatch(); self.integrate.call( pass, @@ -441,7 +393,7 @@ impl GpuMultibodySolver { if mb.is_empty() { return Ok(()); } - let dispatch = [mb.multibodies_per_batch, mb.num_batches, 1]; + let dispatch = mb.flat_mb_dispatch(); if mb.has_joint_constraints { self.remove_solve_joint_no_bias.call( @@ -523,8 +475,10 @@ impl GpuMultibodySolver { mb: &mut GpuMultibodySet, args: &mut MultibodySolverArgs<'_>, ) -> Result<(), GpuBackendError> { - // Fused FK + body-jacobians + velocity propagation + Mass-matrix assembly - let pre_dispatch = [mb.multibodies_per_batch * MB_LU_LANES, mb.num_batches, 1]; + // Fused FK + body-jacobians + velocity propagation + Mass-matrix + // assembly. Packed: `64 / mb_pack_lanes` multibodies per workgroup, + // flattened (multibody, batch) grid. + let pre_dispatch = mb.packed_wg_dispatch(); if mb.implicit_coriolis { self.compute_dynamics_pre.call( pass, @@ -556,22 +510,49 @@ impl GpuMultibodySolver { )?; } - // Fused gravity + LU factor + LU solve. - let grav_lu_dispatch = [mb.multibodies_per_batch * MB_LU_LANES, mb.num_batches, 1]; - self.gravity_and_lu.call( - pass, - grav_lu_dispatch, - &mb.multibody_info, - &mb.links_static, - &mut mb.links_workspace, - &mb.body_jacobians, - &mut mb.gen_forces, - &mut mb.mass_matrices, - &mut mb.lu_pivots, - &mb.dof_state, - &mb.gravity, - args.batch_indices, - )?; + // Fused gravity + LU factor + LU solve. Select an implementation based on how many + // environments we can pack on the same workgroup (depending on its dofs). + macro_rules! grav_lu { + ($kernel:ident) => { + self.$kernel.call( + pass, + mb.packed_wg_dispatch(), + &mb.multibody_info, + &mb.links_static, + &mut mb.links_workspace, + &mb.body_jacobians, + &mut mb.gen_forces, + &mut mb.mass_matrices, + &mut mb.lu_pivots, + &mb.dof_state, + &mb.gravity, + args.batch_indices, + )? + }; + } + match mb.pack_lanes() { + 8 => grav_lu!(gravity_and_lu_t8), + 16 => grav_lu!(gravity_and_lu_t16), + 32 => grav_lu!(gravity_and_lu_t32), + _ => { + let grav_lu_dispatch = + [mb.multibodies_per_batch * MB_LU_LANES, mb.num_batches, 1]; + self.gravity_and_lu.call( + pass, + grav_lu_dispatch, + &mb.multibody_info, + &mb.links_static, + &mut mb.links_workspace, + &mb.body_jacobians, + &mut mb.gen_forces, + &mut mb.mass_matrices, + &mut mb.lu_pivots, + &mb.dof_state, + &mb.gravity, + args.batch_indices, + )?; + } + } Ok(()) } diff --git a/src_rbd/pipeline/insertion_removal.rs b/src_rbd/pipeline/insertion_removal.rs index 7520d96..a079521 100644 --- a/src_rbd/pipeline/insertion_removal.rs +++ b/src_rbd/pipeline/insertion_removal.rs @@ -216,6 +216,7 @@ impl RbdState { let collision_pairs_per_batch_cpu = collisions_capacity; #[allow(unused_mut)] // Only mutated with the dim3 (multibody) feature. let mut bi = BatchIndices { + num_batches, colliders_batch_capacity: num_colliders_per_batch, // No body is active initially; bodies are added later via `append_bodies`. colliders_len: 0, diff --git a/src_rbd/pipeline/rbd_state.rs b/src_rbd/pipeline/rbd_state.rs index b464963..b8a87fc 100644 --- a/src_rbd/pipeline/rbd_state.rs +++ b/src_rbd/pipeline/rbd_state.rs @@ -248,6 +248,7 @@ impl RbdState { pub(super) fn rebuild_batch_indices(&mut self, backend: &GpuBackend) { #[allow(unused_mut)] // Only mutated with the dim3 (multibody) feature. let mut bi = BatchIndices { + num_batches: self.num_batches, colliders_batch_capacity: self.num_colliders_per_batch, colliders_len: self.num_active_colliders, bodies_len: self.num_active_bodies, diff --git a/src_rbd/pipeline/rbd_state_from_rapier.rs b/src_rbd/pipeline/rbd_state_from_rapier.rs index 962f3c7..a8aed63 100644 --- a/src_rbd/pipeline/rbd_state_from_rapier.rs +++ b/src_rbd/pipeline/rbd_state_from_rapier.rs @@ -691,6 +691,7 @@ impl RbdState { let collision_pairs_per_batch_cpu = capacities.collisions_capacity; #[allow(unused_mut)] // Only mutated with the dim3 (multibody) feature. let mut bi = BatchIndices { + num_batches, colliders_batch_capacity: num_colliders_per_batch as u32, colliders_len: num_colliders as u32, bodies_len: num_bodies as u32, diff --git a/src_rbd_shaders/dynamics/multibody/compute_dynamics_pre.rs b/src_rbd_shaders/dynamics/multibody/compute_dynamics_pre.rs index 07233bc..c3152a2 100644 --- a/src_rbd_shaders/dynamics/multibody/compute_dynamics_pre.rs +++ b/src_rbd_shaders/dynamics/multibody/compute_dynamics_pre.rs @@ -27,7 +27,28 @@ use crate::utils::{BatchIndices, Slice, SliceMut}; use crate::{ANG_DIM, AngVector, DIM, Pose, Vector, gcross_av}; use parry::math::VectorExt; -const LANES: u32 = 64; +/// Packed slot decode shared by the two `pre` kernels: `64 / mb_pack_lanes` +/// multibodies per 64-lane workgroup, `(multibody, batch)` flattened into the +/// workgroup X dimension. Returns `(t, lane, batch_id, mb_idx, active_slot)`; +/// inactive slots get clamped indices (their loops all no-op on the zeroed +/// dummy `MultibodyInfo` the caller substitutes). `mb_pack_lanes` is +/// uniform-sourced so the decode keeps uniform control flow for barriers. +#[inline(always)] +fn packed_decode(wg_id: UVec3, lid: UVec3, batch_ids: &BatchIndices) -> (u32, u32, u32, u32, bool) { + let t = batch_ids.mb_pack_lanes; + let slot = lid.x / t; + let lane = lid.x % t; + let slots = 64 / t; + + let num_mb = batch_ids.multibodies_len; + let total_mb = num_mb * batch_ids.num_batches; + let global_mb = wg_id.x * slots + slot; + let active_slot = global_mb < total_mb; + let clamped_mb = if active_slot { global_mb } else { total_mb - 1 }; + let batch_id = clamped_mb / num_mb; + let mb_idx = clamped_mb % num_mb; + (t, lane, batch_id, mb_idx, active_slot) +} // TODO: refactor into multiple functions (but single kernel) to share between the coriolis and non-coriolis versions. /// Fused FK + body-jacobians + velocity propagation + CRBA-with-Coriolis. @@ -49,21 +70,17 @@ pub fn gpu_mb_compute_dynamics_pre( #[spirv(uniform, descriptor_set = 0, binding = 8)] dt_uniform: &f32, #[spirv(uniform, descriptor_set = 0, binding = 9)] batch_ids: &BatchIndices, ) { - let batch_id = wg_id.y; - let mb_idx = wg_id.x; - let lane = lid.x; - // Padding multibody slots have `num_links == 0` and `ndofs == 0` so all - // per-link / per-DOF loops below iterate zero times. No early-return — - // WGSL's naga frontend can't prove a storage-loaded comparison is - // uniform across the workgroup, so any subsequent `workgroupBarrier()` - // would be flagged "called from non-uniform control flow". See - // `gpu_mb_lu_decompose` for the rationale. + let (t, lane, batch_id, mb_idx, active_slot) = packed_decode(wg_id, lid, batch_ids); let dt = *dt_uniform; - let mb = batch_ids - .mb_batch(batch_id, multibody_info) - .read(mb_idx as usize); + let mb = if active_slot { + batch_ids + .mb_batch(batch_id, multibody_info) + .read(mb_idx as usize) + } else { + MultibodyInfo::default() + }; let num_links = mb.num_links; let ndofs = mb.ndofs; let mb_jac_base = batch_ids.jac_start(batch_id) + mb.jacobian_offset as usize; @@ -96,7 +113,7 @@ pub fn gpu_mb_compute_dynamics_pre( let vel_slice = Slice(dof_state, vel_base); // 1) Forward Kinematics (single-threaded) - if lane == 0 { + if active_slot && num_links > 0 && lane == 0 { forward_kinematics(&mb, &stat_slice, &mut poses_slice, &mut ws_slice, num_links); } workgroup_memory_barrier_with_group_sync(); @@ -104,6 +121,7 @@ pub fn gpu_mb_compute_dynamics_pre( // 2) Update body jacobians update_body_jacobians( lane, + t, mb_jac_base, ndofs, num_links, @@ -114,14 +132,14 @@ pub fn gpu_mb_compute_dynamics_pre( ); // 3) Propagate velocities (single-threaded) - if lane == 0 { + if active_slot && num_links > 0 && lane == 0 { propagate_velocities(num_links, &stat_slice, &vel_slice, &mut ws_slice); } workgroup_memory_barrier_with_group_sync(); // 3) Mass matrix (with semi-implicit coriolis handling). let acc_augmented_mass = MatSlice::dense(mb_mm_base, ndofs, ndofs); - fill_par(mass_matrices, acc_augmented_mass, 0.0, lane, LANES); + fill_par(mass_matrices, acc_augmented_mass, 0.0, lane, t); let i_coriolis_dt_view = MatSlice::dense(mb_icdt_base, SPATIAL_DIM as u32, ndofs); let i_coriolis_dt_v = i_coriolis_dt_view.fixed_rows(0, DIM); @@ -145,7 +163,7 @@ pub fn gpu_mb_compute_dynamics_pre( DIM, ndofs, ); - fill_par(coriolis_packed, coriolis_block, 0.0, lane, LANES); + fill_par(coriolis_packed, coriolis_block, 0.0, lane, t); fill_par( coriolis_packed, MatSlice::dense( @@ -155,7 +173,7 @@ pub fn gpu_mb_compute_dynamics_pre( ), 0.0, lane, - LANES, + t, ); } } @@ -210,7 +228,7 @@ pub fn gpu_mb_compute_dynamics_pre( body_jacobian, 1.0, lane, - LANES, + t, ); if k != 0 { @@ -240,14 +258,14 @@ pub fn gpu_mb_compute_dynamics_pre( coriolis_v_i, parent_coriolis_v, lane, - LANES, + t, ); copy_from_par( coriolis_packed, coriolis_w_i, parent_coriolis_w, lane, - LANES, + t, ); gemm_skew_tr_lhs_par( @@ -258,7 +276,7 @@ pub fn gpu_mb_compute_dynamics_pre( parent_coriolis_w, 1.0, lane, - LANES, + t, ); let dvel = crate::gcross_av(ws.rb_vels.angular, ws.shift02) @@ -272,7 +290,7 @@ pub fn gpu_mb_compute_dynamics_pre( parent_j_w, 1.0, lane, - LANES, + t, ); gemm_skew_tr_lhs_cross_buf_par( @@ -284,7 +302,7 @@ pub fn gpu_mb_compute_dynamics_pre( parent_j_w, 1.0, lane, - LANES, + t, ); gemm_omega_skew_tr_cross_buf_par( @@ -297,7 +315,7 @@ pub fn gpu_mb_compute_dynamics_pre( parent_j_w, 1.0, lane, - LANES, + t, ); #[cfg(feature = "dim3")] @@ -311,7 +329,7 @@ pub fn gpu_mb_compute_dynamics_pre( parent_j_w, 1.0, lane, - LANES, + t, ); } } @@ -370,8 +388,8 @@ pub fn gpu_mb_compute_dynamics_pre( } } } else { - fill_par(coriolis_packed, coriolis_v_i, 0.0, lane, LANES); - fill_par(coriolis_packed, coriolis_w_i, 0.0, lane, LANES); + fill_par(coriolis_packed, coriolis_v_i, 0.0, lane, t); + fill_par(coriolis_packed, coriolis_w_i, 0.0, lane, t); } } @@ -387,7 +405,7 @@ pub fn gpu_mb_compute_dynamics_pre( coriolis_w_i, 1.0, lane, - LANES, + t, ); let dvel_23 = crate::gcross_av(ws.rb_vels.angular, ws.shift23); @@ -400,7 +418,7 @@ pub fn gpu_mb_compute_dynamics_pre( rb_j_w, 1.0, lane, - LANES, + t, ); gemm_omega_skew_tr_cross_buf_par( @@ -413,7 +431,7 @@ pub fn gpu_mb_compute_dynamics_pre( rb_j_w, 1.0, lane, - LANES, + t, ); } @@ -439,7 +457,7 @@ pub fn gpu_mb_compute_dynamics_pre( coriolis_w_i, 0.0, lane, - LANES, + t, ); } @@ -456,7 +474,7 @@ pub fn gpu_mb_compute_dynamics_pre( i_coriolis_dt_view, 1.0, lane, - LANES, + t, ); } @@ -494,18 +512,19 @@ pub fn gpu_mb_compute_dynamics_without_coriolis_pre( #[spirv(uniform, descriptor_set = 0, binding = 7)] dt_uniform: &f32, #[spirv(uniform, descriptor_set = 0, binding = 8)] batch_ids: &BatchIndices, ) { - let batch_id = wg_id.y; - let mb_idx = wg_id.x; - let lane = lid.x; - // No early-return on out-of-range `mb_idx` — see `gpu_mb_lu_decompose` - // for the WGSL uniformity rationale. Dummy multibody slots have zero - // links / DOFs, so all per-link loops below iterate zero times. + // Packed layout — see `packed_decode` and the uniformity note on + // `gpu_mb_compute_dynamics_pre`. + let (t, lane, batch_id, mb_idx, active_slot) = packed_decode(wg_id, lid, batch_ids); let dt = *dt_uniform; - let mb = batch_ids - .mb_batch(batch_id, multibody_info) - .read(mb_idx as usize); + let mb = if active_slot { + batch_ids + .mb_batch(batch_id, multibody_info) + .read(mb_idx as usize) + } else { + MultibodyInfo::default() + }; let num_links = mb.num_links; let ndofs = mb.ndofs; let mb_jac_base = batch_ids.jac_start(batch_id) + mb.jacobian_offset as usize; @@ -533,7 +552,7 @@ pub fn gpu_mb_compute_dynamics_without_coriolis_pre( let vel_slice = Slice(dof_state, vel_base); // 1) Forward Kinematics (single-threaded) - if lane == 0 { + if active_slot && num_links > 0 && lane == 0 { forward_kinematics(&mb, &stat_slice, &mut poses_slice, &mut ws_slice, num_links); } workgroup_memory_barrier_with_group_sync(); @@ -541,6 +560,7 @@ pub fn gpu_mb_compute_dynamics_without_coriolis_pre( // 2) Update body jacobians update_body_jacobians( lane, + t, mb_jac_base, ndofs, num_links, @@ -551,14 +571,14 @@ pub fn gpu_mb_compute_dynamics_without_coriolis_pre( ); // 3) Velocities propagation (single-threaded) - if lane == 0 { + if active_slot && num_links > 0 && lane == 0 { propagate_velocities(num_links, &stat_slice, &vel_slice, &mut ws_slice); } workgroup_memory_barrier_with_group_sync(); // 4) Mass matrix (without coriolis). let acc_augmented_mass = MatSlice::dense(mb_mm_base, ndofs, ndofs); - fill_par(mass_matrices, acc_augmented_mass, 0.0, lane, LANES); + fill_par(mass_matrices, acc_augmented_mass, 0.0, lane, t); workgroup_memory_barrier_with_group_sync(); // NOTE: uniform trip count (from the `BatchIndices` uniform). @@ -593,7 +613,7 @@ pub fn gpu_mb_compute_dynamics_without_coriolis_pre( body_jacobian, 1.0, lane, - LANES, + t, ); } @@ -717,6 +737,8 @@ fn forward_kinematics( fn update_body_jacobians( lane: u32, + // Lanes owned by this multibody's slot (`BatchIndices::mb_pack_lanes`). + lanes: u32, mb_jac_base: usize, ndofs: u32, num_links: u32, @@ -751,7 +773,7 @@ fn update_body_jacobians( let parent_link = &ws_slice[link_infos.parent_link_id as usize]; parent_to_world = parent_link.local_to_world; - copy_from_par(body_jacobians, link_j, parent_j, lane, LANES); + copy_from_par(body_jacobians, link_j, parent_j, lane, lanes); let link_j_v = link_j.fixed_rows(0, DIM); let parent_j_w = parent_j.fixed_rows(DIM, ANG_DIM); gemm_skew_tr_lhs_par( @@ -762,10 +784,10 @@ fn update_body_jacobians( parent_j_w, 1.0, lane, - LANES, + lanes, ); } else { - fill_par(body_jacobians, link_j, 0.0, lane, LANES); + fill_par(body_jacobians, link_j, 0.0, lane, lanes); } } @@ -779,7 +801,7 @@ fn update_body_jacobians( body_jacobians, link_j_part, lane, - LANES, + lanes, ); } @@ -796,7 +818,7 @@ fn update_body_jacobians( link_j_w, 1.0, lane, - LANES, + lanes, ); } diff --git a/src_rbd_shaders/dynamics/multibody/contact_constraints.rs b/src_rbd_shaders/dynamics/multibody/contact_constraints.rs index 6b5ddee..b17db0e 100644 --- a/src_rbd_shaders/dynamics/multibody/contact_constraints.rs +++ b/src_rbd_shaders/dynamics/multibody/contact_constraints.rs @@ -119,7 +119,7 @@ fn fill_contact_jac_row( /// `contact_constraint_jacs`. Multibody-multibody contacts (each side a /// different multibody) are not handled — such contacts are skipped. #[spirv_bindgen] -#[spirv(compute(threads(1)))] +#[spirv(compute(threads(64)))] pub fn gpu_mb_init_contact_constraints( #[spirv(global_invocation_id)] invocation_id: UVec3, #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] @@ -135,8 +135,16 @@ pub fn gpu_mb_init_contact_constraints( #[spirv(storage_buffer, descriptor_set = 1, binding = 2)] contacts: &[IndexedManifold], #[spirv(uniform, descriptor_set = 0, binding = 6)] batch_ids: &BatchIndices, ) { - let batch_id = invocation_id.y; - let mb_idx = invocation_id.x; + // Flattened (multibody, batch) grid — see `BatchIndices::num_batches`. + // Only active multibody slots are visited now; the `ndofs == 0` sentinel + // below is kept for all-locked (zero-dof) multibodies. Padding slots past + // `multibodies_len` are never read (every consumer guards on it). + let num_mb = batch_ids.multibodies_len; + if invocation_id.x >= num_mb * batch_ids.num_batches { + return; + } + let batch_id = invocation_id.x / num_mb; + let mb_idx = invocation_id.x % num_mb; // Soft-constraint coefficients (rapier TGS-soft), precomputed on the host. // The old path used a rigid `erp = 1/dt` with zero CFM, which overshoots // penetration recovery (~14× too stiff for the defaults) and jitters. @@ -546,7 +554,7 @@ pub fn gpu_mb_init_contact_constraints( /// `gpu_mb_init_contact_constraints` preserves the impulse across substeps and /// `gpu_mb_warmstart_contact_constraints` re-applies it each substep. #[spirv_bindgen] -#[spirv(compute(threads(1)))] +#[spirv(compute(threads(64)))] pub fn gpu_mb_reset_contact_warmstart( #[spirv(global_invocation_id)] invocation_id: UVec3, #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] multibody_info: &[MultibodyInfo], @@ -554,12 +562,13 @@ pub fn gpu_mb_reset_contact_warmstart( contact_constraints: &mut [MultibodyContactConstraint], #[spirv(uniform, descriptor_set = 0, binding = 2)] batch_ids: &BatchIndices, ) { - let batch_id = invocation_id.y; - let mb_idx = invocation_id.x; + // Flattened (multibody, batch) grid — see `BatchIndices::num_batches`. let num_mb = batch_ids.multibodies_len; - if mb_idx >= num_mb { + if invocation_id.x >= num_mb * batch_ids.num_batches { return; } + let batch_id = invocation_id.x / num_mb; + let mb_idx = invocation_id.x % num_mb; let mb_start = batch_ids.mb_start(batch_id); let cons_start = batch_ids.mb_contact_constraints_start(batch_id); let mb = multibody_info.read(mb_start + mb_idx as usize); @@ -581,7 +590,7 @@ pub fn gpu_mb_reset_contact_warmstart( /// solver velocities. Applies the FULL accumulated impulse (no `rhs` term, no /// clamping). #[spirv_bindgen] -#[spirv(compute(threads(1)))] +#[spirv(compute(threads(64)))] pub fn gpu_mb_warmstart_contact_constraints( #[spirv(global_invocation_id)] invocation_id: UVec3, #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] multibody_info: &[MultibodyInfo], @@ -592,12 +601,13 @@ pub fn gpu_mb_warmstart_contact_constraints( #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] solver_vels: &mut [Velocity], #[spirv(uniform, descriptor_set = 0, binding = 5)] batch_ids: &BatchIndices, ) { - let batch_id = invocation_id.y; - let mb_idx = invocation_id.x; + // Flattened (multibody, batch) grid — see `BatchIndices::num_batches`. let num_mb = batch_ids.multibodies_len; - if mb_idx >= num_mb { + if invocation_id.x >= num_mb * batch_ids.num_batches { return; } + let batch_id = invocation_id.x / num_mb; + let mb_idx = invocation_id.x % num_mb; let mb_start = batch_ids.mb_start(batch_id); let cons_start = batch_ids.mb_contact_constraints_start(batch_id); @@ -730,7 +740,7 @@ pub fn gpu_mb_finalize_contact_constraints( /// One PGS sweep over the multibody's active contact constraints. Updates /// the multibody's `dof_velocities` and the free body's `solver_vels`. #[spirv_bindgen] -#[spirv(compute(threads(1)))] +#[spirv(compute(threads(64)))] pub fn gpu_mb_solve_contact_constraints( #[spirv(global_invocation_id)] invocation_id: UVec3, #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] multibody_info: &[MultibodyInfo], @@ -742,12 +752,13 @@ pub fn gpu_mb_solve_contact_constraints( #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] solver_vels: &mut [Velocity], #[spirv(uniform, descriptor_set = 0, binding = 6)] batch_ids: &BatchIndices, ) { - let batch_id = invocation_id.y; - let mb_idx = invocation_id.x; + // Flattened (multibody, batch) grid — see `BatchIndices::num_batches`. let num_mb = batch_ids.multibodies_len; - if mb_idx >= num_mb { + if invocation_id.x >= num_mb * batch_ids.num_batches { return; } + let batch_id = invocation_id.x / num_mb; + let mb_idx = invocation_id.x % num_mb; let mb_start = batch_ids.mb_start(batch_id); let cons_start = batch_ids.mb_contact_constraints_start(batch_id); @@ -839,7 +850,7 @@ pub fn gpu_mb_solve_contact_constraints( /// Strip the positional bias from each active contact constraint's `rhs`, /// matching `gpu_mb_remove_joint_constraint_bias`. #[spirv_bindgen] -#[spirv(compute(threads(1)))] +#[spirv(compute(threads(64)))] pub fn gpu_mb_remove_contact_constraint_bias( #[spirv(global_invocation_id)] invocation_id: UVec3, #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] @@ -847,12 +858,13 @@ pub fn gpu_mb_remove_contact_constraint_bias( #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] multibody_info: &[MultibodyInfo], #[spirv(uniform, descriptor_set = 0, binding = 2)] batch_ids: &BatchIndices, ) { - let batch_id = invocation_id.y; - let mb_idx = invocation_id.x; + // Flattened (multibody, batch) grid — see `BatchIndices::num_batches`. let num_mb = batch_ids.multibodies_len; - if mb_idx >= num_mb { + if invocation_id.x >= num_mb * batch_ids.num_batches { return; } + let batch_id = invocation_id.x / num_mb; + let mb_idx = invocation_id.x % num_mb; let mb_start = batch_ids.mb_start(batch_id); let cons_start = batch_ids.mb_contact_constraints_start(batch_id); diff --git a/src_rbd_shaders/dynamics/multibody/gravity_and_lu.rs b/src_rbd_shaders/dynamics/multibody/gravity_and_lu.rs index 52c0e7a..4582c42 100644 --- a/src_rbd_shaders/dynamics/multibody/gravity_and_lu.rs +++ b/src_rbd_shaders/dynamics/multibody/gravity_and_lu.rs @@ -22,7 +22,9 @@ use crate::utils::{BatchIndices, Slice}; use crate::{AngVector, Vector, gcross_av}; use super::lu::{ - LANES, lu_apply_pivots, lu_factor_in_shared, lu_triangular_solve_in_place, sm_idx, + LANES, lu_apply_pivots, lu_apply_pivots_packed, lu_factor_in_shared, + lu_factor_in_shared_packed, lu_triangular_solve_in_place, + lu_triangular_solve_in_place_packed, sm_idx, sm_idx_packed, }; use super::types::{MultibodyInfo, MultibodyLinkStatic, MultibodyLinkWorkspace}; @@ -264,3 +266,317 @@ pub fn gpu_mb_gravity_and_lu( gen_forces.write(rhs_offset + lane as usize, x.read(lane as usize)); } } + +/// Packed version of [`gpu_mb_gravity_and_lu`]: `SLOTS = 64 / T` multibodies per +/// 64-lane workgroup, each owning `T` lanes and a `T×T` shared tile. +#[inline(always)] +#[allow(clippy::too_many_arguments)] +fn gravity_and_lu_packed_impl( + wg_id: UVec3, + lid: UVec3, + multibody_info: &[MultibodyInfo], + links_static: &[MultibodyLinkStatic], + links_workspace: &mut [MultibodyLinkWorkspace], + body_jacobians: &[f32], + gen_forces: &mut [f32], + mass_matrices: &mut [f32], + lu_pivots: &mut [u32], + dof_state: &[f32], + gravity: &Vec4, + batch_ids: &BatchIndices, + mat: &mut [f32; MATN], + x: &mut [f32; 64], + partial: &mut [f32; 64], + pivot_row_shared: &mut [u32; SLOTS], + inv_akk_shared: &mut [f32; SLOTS], +) { + let slot = lid.x / T; + let lane = lid.x % T; + let seg = (slot * T) as usize; + + let num_mb = batch_ids.multibodies_len; + let total_mb = num_mb * batch_ids.num_batches; + let global_mb = wg_id.x * SLOTS as u32 + slot; + let active_slot = global_mb < total_mb; + // Clamped so index math stays in-bounds for inactive slots; their loops + // all no-op (dummy `mb`) and every store is guarded. + let clamped_mb = if active_slot { global_mb } else { total_mb - 1 }; + let batch_id = clamped_mb / num_mb; + let mb_idx = clamped_mb % num_mb; + let max_ndofs = batch_ids.mb_max_ndofs; + let max_links = batch_ids.mb_max_links; + + let mb = if active_slot { + batch_ids + .mb_batch(batch_id, multibody_info) + .read(mb_idx as usize) + } else { + MultibodyInfo::default() + }; + let num_links = mb.num_links; + let ndofs = mb.ndofs; + let mb_jac_base = batch_ids.jac_start(batch_id) + mb.jacobian_offset as usize; + let gen_base = batch_ids.dof_start(batch_id) + mb.first_dof as usize; + let mb_mm_base = batch_ids.mm_start(batch_id) + mb.mass_matrix_offset as usize; + let piv_offset = gen_base; + let rhs_offset = gen_base; + + let stat_slice = batch_ids + .mb_links_batch(batch_id, links_static) + .offset(mb.first_link as usize); + let mut ws_slice = batch_ids + .mb_links_batch_mut(batch_id, links_workspace) + .offset(mb.first_link as usize); + let vel_slice = Slice(dof_state, gen_base); + let damping_slice = Slice( + dof_state, + batch_ids.dof_damping_section_offset as usize + gen_base, + ); + + // ---- Phase 1: zero the generalized-force vector (parallel across DOFs). ---- + let accelerations = MatSlice::dense(gen_base, ndofs, 1); + fill_par(gen_forces, accelerations, 0.0, lane, T); + workgroup_memory_barrier_with_group_sync(); + + #[cfg(feature = "dim3")] + let g = Vec3::new(gravity.x, gravity.y, gravity.z); + #[cfg(feature = "dim2")] + let g = Vec2::new(gravity.x, gravity.y); + + // ---- Phase 2: per-link gravity / Coriolis-force assembly. ---- + for k in 0..max_links { + let active = k < num_links; + let mut acc_lin = Vector::ZERO; + #[cfg(feature = "dim3")] + let mut acc_ang: AngVector = AngVector::ZERO; + #[cfg(feature = "dim2")] + let mut acc_ang: AngVector = 0.0; + + if active { + let ( + self_joint_vel_lin, + self_joint_vel_ang, + self_shift02, + self_shift23, + _self_local_to_world, + self_rb_ang, + ) = { + let ws = &ws_slice[k as usize]; + ( + ws.joint_velocity.linear, + ws.joint_velocity.angular, + ws.shift02, + ws.shift23, + ws.local_to_world, + ws.rb_vels.angular, + ) + }; + + if k != 0 { + let stat = stat_slice[k as usize]; + let parent_ws = &ws_slice[stat.parent_link_id as usize]; + let parent_acc_lin = parent_ws.kinematic_acc.linear; + let parent_acc_ang = parent_ws.kinematic_acc.angular; + let parent_ang = parent_ws.rb_vels.angular; + + acc_lin = parent_acc_lin; + acc_ang = parent_acc_ang; + + acc_lin += gcross_av(parent_ang, self_joint_vel_lin) * 2.0; + #[cfg(feature = "dim3")] + { + acc_ang += parent_ang.cross(self_joint_vel_ang); + } + #[cfg(feature = "dim2")] + { + let _ = self_joint_vel_ang; + } + acc_lin += gcross_av(parent_ang, gcross_av(parent_ang, self_shift02)); + acc_lin += gcross_av(parent_acc_ang, self_shift02); + } else { + let _ = self_joint_vel_ang; + let _ = self_shift02; + } + let rb_ang = self_rb_ang; + acc_lin += gcross_av(rb_ang, gcross_av(rb_ang, self_shift23)); + acc_lin += gcross_av(acc_ang, self_shift23); + + if lane == 0 { + ws_slice[k as usize].kinematic_acc = Velocity::new(acc_lin, acc_ang); + } + } + + // Top-level barrier: reached uniformly by every lane on every outer + // iteration so children see the just-published `kinematic_acc`. + workgroup_memory_barrier_with_group_sync(); + + if active { + #[cfg(feature = "dim3")] + let rb_ang = ws_slice[k as usize].rb_vels.angular; + let lmp = stat_slice[k as usize].local_mprops; + let inv_mass_x = lmp.inv_mass.x; + if inv_mass_x != 0.0 { + let mass = 1.0 / inv_mass_x; + let rb_inertia = ws_slice[k as usize].link_world_inertia(&lmp); + + #[cfg(feature = "dim3")] + let gyroscopic = { + let i_omega = rb_inertia * rb_ang; + rb_ang.cross(i_omega) + }; + #[cfg(feature = "dim2")] + let gyroscopic: AngVector = 0.0; + + let i_acc_ang = rb_inertia * acc_ang; + + let f_lin = (g - acc_lin) * mass; + let f_ang = -gyroscopic - i_acc_ang; + + let body_jacobian = MatSlice::dense( + mb_jac_base + (k as usize) * SPATIAL_DIM * (ndofs as usize), + SPATIAL_DIM as u32, + ndofs, + ); + + gemv_tr_spatial_split_par( + gen_forces, + gen_base, + 1.0, + body_jacobians, + body_jacobian, + f_lin, + f_ang, + 1.0, + lane, + T, + ); + } + } + } + + // Apply damping. + workgroup_memory_barrier_with_group_sync(); + let i = lane; + if i < ndofs { + let idx = gen_base + i as usize; + let cur = gen_forces.read(idx); + gen_forces.write(idx, cur - damping_slice[i as usize] * vel_slice[i as usize]); + } + workgroup_memory_barrier_with_group_sync(); + + // ---- Phase 3: load M into this slot's shared tile, factor in place. ---- + let m_view = MatSlice::dense(mb_mm_base, ndofs, ndofs); + if lane < ndofs { + for r in 0..ndofs { + mat.write( + sm_idx_packed::(slot, r, lane), + mass_matrices.read(m_view.idx(r, lane)), + ); + } + x.write(seg + lane as usize, gen_forces.read(rhs_offset + lane as usize)); + } + workgroup_memory_barrier_with_group_sync(); + + lu_factor_in_shared_packed::( + ndofs, + max_ndofs, + slot, + lane, + active_slot, + mat, + lu_pivots, + piv_offset, + pivot_row_shared, + inv_akk_shared, + ); + + // Persist LU factors to global memory (joint / contact constraint init + // reuses them for unit-RHS solves). + if lane < ndofs { + for r in 0..ndofs { + mass_matrices.write(m_view.idx(r, lane), mat.read(sm_idx_packed::(slot, r, lane))); + } + } + + // ---- Phase 4: solve M·x = τ for the gravity rhs. ---- + lu_apply_pivots_packed::(ndofs, slot, lane, active_slot, lu_pivots, piv_offset, x); + lu_triangular_solve_in_place_packed::( + ndofs, + max_ndofs, + slot, + lane, + active_slot, + mat, + x, + partial, + ); + + if lane < ndofs { + gen_forces.write(rhs_offset + lane as usize, x.read(seg + lane as usize)); + } +} + +/// Stamps one packed-tier entry point of the fused gravity + LU kernel. +/// `MATN = 64·T`, `SLOTS = 64/T`. +macro_rules! gravity_and_lu_packed_entry { + ($(#[$doc:meta])* $name:ident, $t:literal, $matn:literal, $slots:literal) => { + $(#[$doc])* + #[spirv_bindgen] + #[spirv(compute(threads(64, 1, 1)))] + pub fn $name( + #[spirv(workgroup_id)] wg_id: UVec3, + #[spirv(local_invocation_id)] lid: UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + multibody_info: &[MultibodyInfo], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] + links_static: &[MultibodyLinkStatic], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] + links_workspace: &mut [MultibodyLinkWorkspace], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] body_jacobians: &[f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] gen_forces: &mut [f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] mass_matrices: &mut [f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 6)] lu_pivots: &mut [u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 7)] dof_state: &[f32], + #[spirv(uniform, descriptor_set = 0, binding = 8)] gravity: &Vec4, + #[spirv(uniform, descriptor_set = 0, binding = 9)] batch_ids: &BatchIndices, + #[spirv(workgroup)] mat: &mut [f32; $matn], + #[spirv(workgroup)] x: &mut [f32; 64], + #[spirv(workgroup)] partial: &mut [f32; 64], + #[spirv(workgroup)] pivot_row_shared: &mut [u32; $slots], + #[spirv(workgroup)] inv_akk_shared: &mut [f32; $slots], + ) { + gravity_and_lu_packed_impl::<$t, $matn, $slots>( + wg_id, + lid, + multibody_info, + links_static, + links_workspace, + body_jacobians, + gen_forces, + mass_matrices, + lu_pivots, + dof_state, + gravity, + batch_ids, + mat, + x, + partial, + pivot_row_shared, + inv_akk_shared, + ); + } + }; +} + +gravity_and_lu_packed_entry!( + /// We have `max_ndofs ≤ 8` => 8 multibodies per workgroup. + gpu_mb_gravity_and_lu_t8, 8u32, 512, 8 +); +gravity_and_lu_packed_entry!( + /// We have `max_ndofs ≤ 16` => 4 multibodies per workgroup. + gpu_mb_gravity_and_lu_t16, 16u32, 1024, 4 +); +gravity_and_lu_packed_entry!( + /// We have `max_ndofs ≤ 32` => 2 multibodies per workgroup. + gpu_mb_gravity_and_lu_t32, 32u32, 2048, 2 +); diff --git a/src_rbd_shaders/dynamics/multibody/integrate.rs b/src_rbd_shaders/dynamics/multibody/integrate.rs index 5c1b5b6..cbc1d1b 100644 --- a/src_rbd_shaders/dynamics/multibody/integrate.rs +++ b/src_rbd_shaders/dynamics/multibody/integrate.rs @@ -25,7 +25,7 @@ use super::types::{MultibodyInfo, MultibodyLinkStatic, MultibodyLinkWorkspace}; /// constraints can run in between (rapier's order: velocity update → constraint /// solver → position update). #[spirv_bindgen] -#[spirv(compute(threads(1)))] +#[spirv(compute(threads(64)))] pub fn gpu_mb_integrate_velocities( #[spirv(global_invocation_id)] invocation_id: UVec3, #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] multibody_info: &[MultibodyInfo], @@ -34,12 +34,12 @@ pub fn gpu_mb_integrate_velocities( #[spirv(uniform, descriptor_set = 0, binding = 3)] dt_uniform: &f32, #[spirv(uniform, descriptor_set = 0, binding = 4)] batch_ids: &BatchIndices, ) { - let batch_id = invocation_id.y; - let mb_idx = invocation_id.x; let num_mb = batch_ids.multibodies_len; - if mb_idx >= num_mb { + if invocation_id.x >= num_mb * batch_ids.num_batches { return; } + let batch_id = invocation_id.x / num_mb; + let mb_idx = invocation_id.x % num_mb; let dt = *dt_uniform; let mb = batch_ids @@ -57,7 +57,7 @@ pub fn gpu_mb_integrate_velocities( } #[spirv_bindgen] -#[spirv(compute(threads(1)))] +#[spirv(compute(threads(64)))] pub fn gpu_mb_integrate( #[spirv(global_invocation_id)] invocation_id: UVec3, #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] multibody_info: &[MultibodyInfo], @@ -70,12 +70,12 @@ pub fn gpu_mb_integrate( #[spirv(uniform, descriptor_set = 0, binding = 5)] dt_uniform: &f32, #[spirv(uniform, descriptor_set = 0, binding = 6)] batch_ids: &BatchIndices, ) { - let batch_id = invocation_id.y; - let mb_idx = invocation_id.x; let num_mb = batch_ids.multibodies_len; - if mb_idx >= num_mb { + if invocation_id.x >= num_mb * batch_ids.num_batches { return; } + let batch_id = invocation_id.x / num_mb; + let mb_idx = invocation_id.x % num_mb; let dt = *dt_uniform; let mb = batch_ids diff --git a/src_rbd_shaders/dynamics/multibody/joint_constraints.rs b/src_rbd_shaders/dynamics/multibody/joint_constraints.rs index 29a4511..cb60d12 100644 --- a/src_rbd_shaders/dynamics/multibody/joint_constraints.rs +++ b/src_rbd_shaders/dynamics/multibody/joint_constraints.rs @@ -515,7 +515,7 @@ fn emit_motor_constraint( /// Must run after `gpu_mb_lu_decompose` — the LU factors of `M` are used to compute /// the per-constraint M⁻¹ column and effective inverse mass. #[spirv_bindgen] -#[spirv(compute(threads(1)))] +#[spirv(compute(threads(64)))] pub fn gpu_mb_init_joint_constraints( #[spirv(global_invocation_id)] invocation_id: UVec3, #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] multibody_info: &[MultibodyInfo], @@ -532,12 +532,13 @@ pub fn gpu_mb_init_joint_constraints( #[spirv(uniform, descriptor_set = 0, binding = 7)] softness: &ConstraintSoftness, #[spirv(uniform, descriptor_set = 0, binding = 8)] batch_ids: &BatchIndices, ) { - let batch_id = invocation_id.y; - let mb_idx = invocation_id.x; + // Flattened (multibody, batch) grid — see `BatchIndices::num_batches`. let num_mb = batch_ids.multibodies_len; - if mb_idx >= num_mb { + if invocation_id.x >= num_mb * batch_ids.num_batches { return; } + let batch_id = invocation_id.x / num_mb; + let mb_idx = invocation_id.x % num_mb; init_joint_constraints_body( multibody_info, links_static, @@ -559,7 +560,7 @@ pub fn gpu_mb_init_joint_constraints( /// updates `dof_velocities` in place. Mirrors rapier's `JointConstraint::solve_generic` /// for a 1-DOF jacobian. #[spirv_bindgen] -#[spirv(compute(threads(1)))] +#[spirv(compute(threads(64)))] pub fn gpu_mb_solve_joint_constraints( #[spirv(global_invocation_id)] invocation_id: UVec3, #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] multibody_info: &[MultibodyInfo], @@ -570,12 +571,13 @@ pub fn gpu_mb_solve_joint_constraints( #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] dof_state: &mut [f32], #[spirv(uniform, descriptor_set = 0, binding = 4)] batch_ids: &BatchIndices, ) { - let batch_id = invocation_id.y; - let mb_idx = invocation_id.x; + // Flattened (multibody, batch) grid — see `BatchIndices::num_batches`. let num_mb = batch_ids.multibodies_len; - if mb_idx >= num_mb { + if invocation_id.x >= num_mb * batch_ids.num_batches { return; } + let batch_id = invocation_id.x / num_mb; + let mb_idx = invocation_id.x % num_mb; solve_joint_constraints_body( multibody_info, joint_constraints, @@ -589,9 +591,9 @@ pub fn gpu_mb_solve_joint_constraints( /// Fused `remove_bias + solve_without_bias` for joint constraints — runs once /// per substep at the end of the substep, after position integration. Drops -/// one threads(1) dispatch per substep. +/// one per-multibody dispatch per substep. #[spirv_bindgen] -#[spirv(compute(threads(1)))] +#[spirv(compute(threads(64)))] pub fn gpu_mb_remove_solve_joint_no_bias( #[spirv(global_invocation_id)] invocation_id: UVec3, #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] multibody_info: &[MultibodyInfo], @@ -601,12 +603,13 @@ pub fn gpu_mb_remove_solve_joint_no_bias( #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] dof_state: &mut [f32], #[spirv(uniform, descriptor_set = 0, binding = 4)] batch_ids: &BatchIndices, ) { - let batch_id = invocation_id.y; - let mb_idx = invocation_id.x; + // Flattened (multibody, batch) grid — see `BatchIndices::num_batches`. let num_mb = batch_ids.multibodies_len; - if mb_idx >= num_mb { + if invocation_id.x >= num_mb * batch_ids.num_batches { return; } + let batch_id = invocation_id.x / num_mb; + let mb_idx = invocation_id.x % num_mb; let mb = batch_ids .mb_batch(batch_id, multibody_info) diff --git a/src_rbd_shaders/dynamics/multibody/lu.rs b/src_rbd_shaders/dynamics/multibody/lu.rs index d6c9909..da41437 100644 --- a/src_rbd_shaders/dynamics/multibody/lu.rs +++ b/src_rbd_shaders/dynamics/multibody/lu.rs @@ -162,6 +162,198 @@ pub(super) fn lu_triangular_solve_in_place( } } +/* + * Packed variants: `SLOTS = 64 / T` multibodies per 64-lane workgroup, each + * owning `T` lanes and a `T×T` shared tile. Shrinking the tile from + * `MAX_MB_DOFS² = 16 KB` to `64·T` floats lifts the shared-memory occupancy + * cap that made the one-multibody-per-workgroup layout latency-bound when + * every environment holds one small robot. All barriers stay at uniform + * points: every slot executes the same `max_n`-bounded loops in lock-step, + * inactive slots simply skip their stores. + */ + +/// Index into the packed shared tile: slot-`slot`'s `T×T` column-major tile. +#[inline] +pub(super) fn sm_idx_packed(slot: u32, r: u32, c: u32) -> usize { + (slot * T * T + c * T + r) as usize +} + +/// Packed [`lu_factor_in_shared`]: factor each slot's `T×T` tile in place. +/// `lane` is slot-relative (`0..T`). +#[inline] +#[allow(clippy::too_many_arguments)] +pub(super) fn lu_factor_in_shared_packed( + n: u32, + max_n: u32, + slot: u32, + lane: u32, + active_slot: bool, + mat: &mut [f32; MATN], + pivots_dst: &mut [u32], + pivots_offset: usize, + pivot_row_shared: &mut [u32; SLOTS], + inv_akk_shared: &mut [f32; SLOTS], +) { + for k in 0..max_n { + let active = active_slot && k < n; + if active && lane == 0 { + let mut pivot_row = k; + let mut pivot_val = { + let v = mat.read(sm_idx_packed::(slot, k, k)); + if v >= 0.0 { v } else { -v } + }; + for i in (k + 1)..n { + let v = mat.read(sm_idx_packed::(slot, i, k)); + let av = if v >= 0.0 { v } else { -v }; + if av > pivot_val { + pivot_val = av; + pivot_row = i; + } + } + pivot_row_shared.write(slot as usize, pivot_row); + pivots_dst.write(pivots_offset + k as usize, pivot_row); + } + workgroup_memory_barrier_with_group_sync(); + let pivot_row = pivot_row_shared.read(slot as usize); + + if active && pivot_row != k && lane < n { + let c = lane; + let a = mat.read(sm_idx_packed::(slot, k, c)); + let b = mat.read(sm_idx_packed::(slot, pivot_row, c)); + mat.write(sm_idx_packed::(slot, k, c), b); + mat.write(sm_idx_packed::(slot, pivot_row, c), a); + } + workgroup_memory_barrier_with_group_sync(); + + if active && lane == 0 { + let akk = mat.read(sm_idx_packed::(slot, k, k)); + inv_akk_shared.write(slot as usize, if akk != 0.0 { 1.0 / akk } else { 0.0 }); + } + workgroup_memory_barrier_with_group_sync(); + let inv_akk = inv_akk_shared.read(slot as usize); + + if active { + let r = k + 1 + lane; + if r < n { + let v = mat.read(sm_idx_packed::(slot, r, k)) * inv_akk; + mat.write(sm_idx_packed::(slot, r, k), v); + } + } + workgroup_memory_barrier_with_group_sync(); + + if active { + let j = k + 1 + lane; + if j < n { + let akj = mat.read(sm_idx_packed::(slot, k, j)); + for r in (k + 1)..n { + let lik = mat.read(sm_idx_packed::(slot, r, k)); + let v = mat.read(sm_idx_packed::(slot, r, j)) - lik * akj; + mat.write(sm_idx_packed::(slot, r, j), v); + } + } + } + workgroup_memory_barrier_with_group_sync(); + } +} + +/// Packed [`lu_triangular_solve_in_place`]: per-slot `x`/`partial` segments +/// live at `slot·T`; the tree reduction runs `log2(T)` levels within each +/// slot's segment, all slots in lock-step. +#[inline] +pub(super) fn lu_triangular_solve_in_place_packed( + n: u32, + max_n: u32, + slot: u32, + lane: u32, + active_slot: bool, + mat: &[f32; MATN], + x: &mut [f32; 64], + partial: &mut [f32; 64], +) { + let seg = (slot * T) as usize; + let log2_t = T.trailing_zeros(); + + for i in 0..max_n { + let active = active_slot && i < n; + let s = if active && lane < i { + mat.read(sm_idx_packed::(slot, i, lane)) * x.read(seg + lane as usize) + } else { + 0.0f32 + }; + partial.write(seg + lane as usize, s); + workgroup_memory_barrier_with_group_sync(); + for step in 0..log2_t { + let stride = T >> (step + 1); + if lane < stride { + let v = partial.read(seg + lane as usize) + + partial.read(seg + (lane + stride) as usize); + partial.write(seg + lane as usize, v); + } + workgroup_memory_barrier_with_group_sync(); + } + if active && lane == 0 { + let cur = x.read(seg + i as usize); + x.write(seg + i as usize, cur - partial.read(seg)); + } + workgroup_memory_barrier_with_group_sync(); + } + + for step in 0..max_n { + let active = active_slot && step < n; + // For dummy iterations (step >= n), `i` is not meaningful; guard + // every use of it behind `active`. + let i = if active { n - 1 - step } else { 0 }; + let s = if active && lane > i && lane < n { + mat.read(sm_idx_packed::(slot, i, lane)) * x.read(seg + lane as usize) + } else { + 0.0f32 + }; + partial.write(seg + lane as usize, s); + workgroup_memory_barrier_with_group_sync(); + for r in 0..log2_t { + let stride = T >> (r + 1); + if lane < stride { + let v = partial.read(seg + lane as usize) + + partial.read(seg + (lane + stride) as usize); + partial.write(seg + lane as usize, v); + } + workgroup_memory_barrier_with_group_sync(); + } + if active && lane == 0 { + let u = mat.read(sm_idx_packed::(slot, i, i)); + let cur = x.read(seg + i as usize) - partial.read(seg); + x.write(seg + i as usize, if u != 0.0 { cur / u } else { 0.0 }); + } + workgroup_memory_barrier_with_group_sync(); + } +} + +/// Packed [`lu_apply_pivots`]: sequential on each slot's lane 0. +#[inline] +pub(super) fn lu_apply_pivots_packed( + n: u32, + slot: u32, + lane: u32, + active_slot: bool, + buf_pivots: &[u32], + pivots_offset: usize, + x: &mut [f32; 64], +) { + let seg = (slot * T) as usize; + if active_slot && lane == 0 { + for k in 0..n { + let p = buf_pivots.read(pivots_offset + k as usize); + if p != k { + let a = x.read(seg + k as usize); + let b = x.read(seg + p as usize); + x.write(seg + k as usize, b); + x.write(seg + p as usize, a); + } + } + } + workgroup_memory_barrier_with_group_sync(); +} + /// Apply the recorded pivots (sequential — lane 0 only). `n` is small so the /// extra parallelism wouldn't pay for the barrier. #[inline] diff --git a/src_rbd_shaders/utils/indices.rs b/src_rbd_shaders/utils/indices.rs index ecac949..24e7996 100644 --- a/src_rbd_shaders/utils/indices.rs +++ b/src_rbd_shaders/utils/indices.rs @@ -9,6 +9,9 @@ use crate::utils::{Slice, SliceMut}; #[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] #[repr(C)] pub struct BatchIndices { + /// Total number of simulation batches (environments). + pub num_batches: u32, + /* * RBD / collision-detection capacities. */ @@ -58,6 +61,9 @@ pub struct BatchIndices { pub mb_max_ndofs: u32, /// Actual max link count across every multibody in every batch. pub mb_max_links: u32, + /// Lanes per multibody for the packed per-multibody workgroup kernels: + /// `next_power_of_two(mb_max_ndofs).clamp(8, 64)`. + pub mb_pack_lanes: u32, /// Per-batch stride of the contact-solver color-bucket buffers /// (`color_counts` / `color_starts` / `color_cursors`), = `max_colors + 3` /// so that `starts[c + 1]` is in bounds for every swept color. From c86471be1dfecc3dc791f2cdfccabcf236301972 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Sun, 26 Jul 2026 11:39:42 +0200 Subject: [PATCH 19/39] perf: filter same-body and disabled self-contact pairs in the broad phase --- src_rbd/broad_phase/lbvh.rs | 2 ++ src_rbd/pipeline/insertion_removal.rs | 9 ++++++++ src_rbd/pipeline/rbd_state.rs | 5 ++++ src_rbd/pipeline/rbd_state_from_rapier.rs | 28 +++++++++++++++++++++++ src_rbd/pipeline/rbd_step.rs | 1 + src_rbd_shaders/broad_phase/lbvh.rs | 17 +++++++------- 6 files changed, 54 insertions(+), 8 deletions(-) diff --git a/src_rbd/broad_phase/lbvh.rs b/src_rbd/broad_phase/lbvh.rs index ead559c..4a86eb1 100644 --- a/src_rbd/broad_phase/lbvh.rs +++ b/src_rbd/broad_phase/lbvh.rs @@ -269,6 +269,7 @@ impl Lbvh { collision_pairs_len: &mut Tensor, collision_pairs_indirect: &mut Tensor<[u32; 3]>, collision_groups: &Tensor, + 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; @@ -286,6 +287,7 @@ impl Lbvh { collision_pairs_len, collision_groups, batch_indices, + pair_filter, )?; self.shaders.lbvh_init_indirect_args.call( pass, diff --git a/src_rbd/pipeline/insertion_removal.rs b/src_rbd/pipeline/insertion_removal.rs index a079521..d4f9f70 100644 --- a/src_rbd/pipeline/insertion_removal.rs +++ b/src_rbd/pipeline/insertion_removal.rs @@ -77,9 +77,12 @@ impl RbdState { // collider_parent: identity within each batch initially (no body is // active yet). `append_bodies` overwrites the active prefix. let mut all_collider_parent = Vec::with_capacity(num_bodies_total); + let mut all_pair_filter = Vec::with_capacity(num_bodies_total); for _ in 0..num_batches { for c in 0..capacity_per_batch { all_collider_parent.push(c); + // Identity parent, no multibody key. + all_pair_filter.push([c, 0u32]); } } @@ -123,6 +126,7 @@ impl RbdState { let collider_local_poses = Tensor::vector(backend, &all_collider_local_poses, rw).unwrap(); let collider_parent = Tensor::vector(backend, &all_collider_parent, rw).unwrap(); let collision_groups = Tensor::vector(backend, &all_collision_groups, rw).unwrap(); + let pair_filter = Tensor::vector(backend, &all_pair_filter, rw).unwrap(); // Padding colliders are inert; default material keeps the buffer sized. let all_collider_materials = vec![GpuColliderMaterial::default(); num_bodies_total]; let collider_materials = Tensor::vector(backend, &all_collider_materials, rw).unwrap(); @@ -261,6 +265,7 @@ impl RbdState { collider_local_poses, collider_parent, collision_groups, + pair_filter, collider_materials, collision_pairs, collision_pairs_len, @@ -426,6 +431,8 @@ impl RbdState { // body's collider slot equals its body slot: `collider_parent` is the // identity over the appended (env-local) range. let parents: Vec = (active as u32..(active + bodies.len()) as u32).collect(); + // NOTE: appended bodies are free bodies (never multibody links). + let pair_filters: Vec<[u32; 2]> = parents.iter().map(|&p| [p, 0u32]).collect(); // Write the same body data into every batch's slot range so all // environments share the same topology. @@ -440,6 +447,7 @@ impl RbdState { &collider_local_poses, )?; backend.write_buffer(self.collider_parent.buffer_mut(), base, &parents)?; + backend.write_buffer(self.pair_filter.buffer_mut(), base, &pair_filters)?; backend.write_buffer(self.local_mprops.buffer_mut(), base, &local_mprops)?; backend.write_buffer(self.mprops.buffer_mut(), base, &mprops)?; backend.write_buffer(self.shapes.buffer_mut(), base, &shapes)?; @@ -582,6 +590,7 @@ impl RbdState { // `collider_parent` is the identity mapping on the incremental (one // collider per body) path and stays identity under swap-remove, so it // needs no relocation; only the active body count tracks the colliders. + // The same holds for `pair_filter` (`[identity, 0]` on this path). self.num_active_bodies = self.num_active_colliders; self.rebuild_batch_indices(backend); Ok(remaps) diff --git a/src_rbd/pipeline/rbd_state.rs b/src_rbd/pipeline/rbd_state.rs index b8a87fc..328e88e 100644 --- a/src_rbd/pipeline/rbd_state.rs +++ b/src_rbd/pipeline/rbd_state.rs @@ -151,6 +151,11 @@ pub struct RbdState { pub(super) collider_world_poses: Tensor, /// Per-collider [`crate::rapier::geometry::InteractionGroups`]. pub(super) collision_groups: Tensor, + /// Per-collider broad-phase pair-filter key: + /// - `[0]`: to prevent colliders of the same body from colliding. + /// - `[1]`: to prevent colliders of adjacent links from a multibody from coliding. + /// Nonzero keys that are equal never collide. + pub(super) pair_filter: Tensor<[u32; 2]>, /// Per-collider friction / restitution coefficients (+ combine rules), pub(super) collider_materials: Tensor, pub(super) collision_pairs: Tensor, diff --git a/src_rbd/pipeline/rbd_state_from_rapier.rs b/src_rbd/pipeline/rbd_state_from_rapier.rs index a8aed63..bd4de5d 100644 --- a/src_rbd/pipeline/rbd_state_from_rapier.rs +++ b/src_rbd/pipeline/rbd_state_from_rapier.rs @@ -141,6 +141,7 @@ impl RbdState { // Bodies and colliders share the `max_colliders` per-batch stride but // form distinct index spaces — a body may own several colliders. let mut all_collider_parent: Vec = Vec::new(); + let mut all_pair_filter: Vec<[u32; 2]> = Vec::new(); // Per-environment count of *active* rigid bodies (distinct collider // parents + one synthetic body per parentless collider). Used to assert // the equal-topology invariant and to set `BatchIndices::bodies_len`. @@ -229,6 +230,20 @@ impl RbdState { } }; + // Handle bodies whose multibody disables self-contacts. + #[cfg(feature = "dim3")] + let no_self_collide: HashMap = { + let mut map = HashMap::new(); + for (mb_ord, mb) in multibody_joints.multibodies().enumerate() { + if !mb.self_contacts_enabled() { + for link in mb.links() { + map.insert(link.rigid_body_handle(), mb_ord as u32 + 1); + } + } + } + map + }; + for (_, co) in colliders.iter() { // Resolve (allocating if needed) the parent body's env-local slot. // @@ -281,6 +296,16 @@ impl RbdState { all_collider_materials.push(collider_material_from_rapier(co)); // Env-local body slot; the kernels apply the per-batch stride. all_collider_parent.push(body_local); + + // Broad-phase pair-filter key. + #[cfg(feature = "dim3")] + let mb_key = co + .parent() + .and_then(|h| no_self_collide.get(&h).copied()) + .unwrap_or(0); + #[cfg(feature = "dim2")] + let mb_key = 0u32; + all_pair_filter.push([body_local, mb_key]); } // Give every multibody link a body slot too, even collider-less ones @@ -327,6 +352,7 @@ impl RbdState { all_collider_materials.push(GpuColliderMaterial::default()); // Padding colliders are inert; parent body 0 is a safe placeholder. all_collider_parent.push(0); + all_pair_filter.push([u32::MAX, 0]); } // Pad bodies to the shared per-batch stride (`max_colliders` = @@ -539,6 +565,7 @@ impl RbdState { Tensor::vector(backend, &all_collider_local_poses, storage).unwrap(); let collider_parent = Tensor::vector(backend, &all_collider_parent, storage).unwrap(); let collision_groups = Tensor::vector(backend, &all_collision_groups, storage).unwrap(); + let pair_filter = Tensor::vector(backend, &all_pair_filter, storage).unwrap(); let collider_materials = Tensor::vector(backend, &all_collider_materials, storage).unwrap(); let collision_pairs = Tensor::vector_uninit( @@ -756,6 +783,7 @@ impl RbdState { collider_local_poses, collider_parent, collision_groups, + pair_filter, collider_materials, collision_pairs, collision_pairs_len, diff --git a/src_rbd/pipeline/rbd_step.rs b/src_rbd/pipeline/rbd_step.rs index 91f35ce..c0b9a0a 100644 --- a/src_rbd/pipeline/rbd_step.rs +++ b/src_rbd/pipeline/rbd_step.rs @@ -178,6 +178,7 @@ impl RbdPipeline { &mut state.collision_pairs_len, &mut state.collision_pairs_indirect, &state.collision_groups, + &state.pair_filter, )?; drop(pass); diff --git a/src_rbd_shaders/broad_phase/lbvh.rs b/src_rbd_shaders/broad_phase/lbvh.rs index 6794f86..0d5b600 100644 --- a/src_rbd_shaders/broad_phase/lbvh.rs +++ b/src_rbd_shaders/broad_phase/lbvh.rs @@ -502,6 +502,7 @@ pub fn gpu_lbvh_find_collision_pairs( #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] collision_groups: &[InteractionGroups], #[spirv(uniform, descriptor_set = 0, binding = 4)] batch_ids: &BatchIndices, + #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] pair_filter: &[[u32; 2]], ) { let num_threads = num_workgroups.x * WORKGROUP_SIZE; let batch_id = invocation_id.y; @@ -512,10 +513,12 @@ pub fn gpu_lbvh_find_collision_pairs( let mut collision_pairs = batch_ids.collision_pairs_batch_mut(batch_id, collision_pairs); let tree = Slice(tree, root_id(colliders_start) as usize); let collision_groups = batch_ids.coll_batch(batch_id, collision_groups); + let pair_filter = batch_ids.coll_batch(batch_id, pair_filter); for leaf_i in StepRng::new(invocation_id.x..num_bodies, num_threads) { let i = tree.at((first_leaf_id + leaf_i) as usize).left; let groups_i = collision_groups[i as usize]; + let filter_i = pair_filter[i as usize]; let mut aabb1 = tree.at((first_leaf_id + leaf_i) as usize).aabb; let prediction = 2.0e-3; // TODO: should be configurable. let dilation = Vector::splat(prediction); @@ -544,18 +547,16 @@ pub fn gpu_lbvh_find_collision_pairs( let groups_j = collision_groups[j as usize]; // Skip pairs whose collision groups don't authorize an interaction. - // NOTE: same-body collider pairs are *not* filtered here — that - // skip is deferred to the narrow-phase so the broad phase - // never has to touch `collider_parent`. if !groups_i.test(groups_j) { continue; } - // Duplicates were already pruned during the descent (sorted - // leaf-index comparison). Emit the pair in ascending collider - // order so the narrow phase / warmstart see a stable ordering - // regardless of the traversal's dedup basis. - let (ci, cj) = if i < j { (i, j) } else { (j, i) }; + // Apply computed filters (same-body and self-contact). + let filter_j = pair_filter[j as usize]; + if filter_i[0] == filter_j[0] || (filter_i[1] != 0 && filter_i[1] == filter_j[1]) { + continue; + } + let target_pair_index = atomic_add_u32(collision_pairs_len.at_mut(batch_id as usize), 1); From 8ddb97cf398173610aab32f1d01f895ba7eb07de Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Sun, 26 Jul 2026 13:12:53 +0200 Subject: [PATCH 20/39] perf: parallelize the per-batch max reductions feeding indirect dispatches --- src_rbd/broad_phase/lbvh.rs | 4 +- src_rbd/broad_phase/narrow_phase.rs | 7 +- src_rbd_shaders/broad_phase/lbvh.rs | 85 +++++++++++++-------- src_rbd_shaders/broad_phase/narrow_phase.rs | 67 +++++----------- 4 files changed, 81 insertions(+), 82 deletions(-) diff --git a/src_rbd/broad_phase/lbvh.rs b/src_rbd/broad_phase/lbvh.rs index 4a86eb1..d5b27d4 100644 --- a/src_rbd/broad_phase/lbvh.rs +++ b/src_rbd/broad_phase/lbvh.rs @@ -276,7 +276,7 @@ impl Lbvh { self.shaders.reset_collision_pairs.call( pass, - [1u32, num_batches, 1], + [num_batches, 1, 1], collision_pairs_len, )?; self.shaders.find_collision_pairs.call( @@ -291,7 +291,7 @@ impl Lbvh { )?; self.shaders.lbvh_init_indirect_args.call( pass, - 1u32, + 256u32, collision_pairs_len, collision_pairs_indirect, )?; diff --git a/src_rbd/broad_phase/narrow_phase.rs b/src_rbd/broad_phase/narrow_phase.rs index b37cfb5..61a59a6 100644 --- a/src_rbd/broad_phase/narrow_phase.rs +++ b/src_rbd/broad_phase/narrow_phase.rs @@ -51,7 +51,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 +84,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 +98,9 @@ impl GpuNarrowPhase { collider_parent, collider_materials, )?; + // Single 256-lane workgroup: parallel max over the per-batch counts. self.init_contacts_indirect_args - .call(pass, 1u32, contacts_len, contacts_indirect)?; + .call(pass, 256u32, contacts_len, contacts_indirect)?; Ok(()) } diff --git a/src_rbd_shaders/broad_phase/lbvh.rs b/src_rbd_shaders/broad_phase/lbvh.rs index 0d5b600..6e2c2e6 100644 --- a/src_rbd_shaders/broad_phase/lbvh.rs +++ b/src_rbd_shaders/broad_phase/lbvh.rs @@ -49,49 +49,74 @@ pub struct LbvhNode { pub refit_count_or_max_subtree_index: u32, } -/// Resets the collision pairs counter. +/// Resets the collision pairs counter. One thread per batch. #[spirv_bindgen] -#[spirv(compute(threads(1)))] +#[spirv(compute(threads(64)))] pub fn gpu_lbvh_reset_collision_pairs( - #[spirv(workgroup_id)] workgroup_id: UVec3, + #[spirv(global_invocation_id)] invocation_id: UVec3, #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] collision_pairs_len: &mut [u32], ) { - let batch_id = workgroup_id.y as usize; - - // NOTE: this `for` loop is silly. It doesn’t do anything - // more than a `*collision_pairs_len = 0` in a convoluted - // way because otherwise rustgpu apparently does not generate - // the spirv for this kernel (seems to happen if the kernel is - // too trivial. - for k in 0..1 { - collision_pairs_len.write(batch_id, k); + let batch_id = invocation_id.x as usize; + if batch_id < collision_pairs_len.len() { + collision_pairs_len.write(batch_id, 0); + } +} + +/// Number of lanes used by the per-batch-count max reductions below. Their +/// host dispatch is a single workgroup: `.call(pass, MAX_REDUCE_LANES, ...)`. +pub const MAX_REDUCE_LANES: u32 = 256; + +/// Workgroup-parallel `max` over the per-batch counts, then writes the +/// `[ceil(max/64), num_batches, 1]` indirect grid. +/// +/// NOTE: `lens` is mutable even though we don't modify it: the loads must be +/// atomic or they occasionally read stale data (breaks Windows+Nvidia+wgpu, see +/// https://github.com/gfx-rs/wgpu/issues/9221). +#[inline(always)] +pub(crate) fn max_len_indirect_args( + lane: u32, + lens: &mut [u32], + indirect_args: &mut [u32; 3], + partial: &mut [u32; MAX_REDUCE_LANES as usize], +) { + let num_batches = lens.len(); + + let mut m = 0u32; + for i in StepRng::new(lane..num_batches as u32, MAX_REDUCE_LANES) { + m = m.max(atomic_load_u32(lens.at_mut(i as usize))); + } + partial.write(lane as usize, m); + workgroup_memory_barrier_with_group_sync(); + + // Tree reduction over the 256 lanes (8 halving steps). + for step in 0..8u32 { + let stride = MAX_REDUCE_LANES >> (step + 1); + if lane < stride { + let v = partial + .read(lane as usize) + .max(partial.read((lane + stride) as usize)); + partial.write(lane as usize, v); + } + workgroup_memory_barrier_with_group_sync(); + } + + if lane == 0 { + *indirect_args.at_mut(0) = partial.read(0).div_ceil(WORKGROUP_SIZE); + *indirect_args.at_mut(1) = num_batches as u32; + *indirect_args.at_mut(2) = 1; } } /// Initializes indirect dispatch arguments for narrow phase. #[spirv_bindgen] -#[spirv(compute(threads(1)))] +#[spirv(compute(threads(256)))] pub fn gpu_lbvh_init_dispatch( - // TODO: take the batch dimension as argument (instead of relying on the len of `collision_pairs_len`)? - // NOTE: the `collision_pairs_len` is mutable here even though we don’t modify it. That’s - // because we access it with an atomic load otherwise it would occasionally read - // stale data (on Windows+Nvidia+wgpu backend). This might be caused by: - // https://github.com/gfx-rs/wgpu/issues/9221 + #[spirv(local_invocation_id)] lid: UVec3, #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] collision_pairs_len: &mut [u32], #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] indirect_args: &mut [u32; 3], + #[spirv(workgroup)] partial: &mut [u32; MAX_REDUCE_LANES as usize], ) { - // For indirect dispatch, get the largest length along all batch dimensions. - let num_batches = collision_pairs_len.len(); - let mut highest_pairs_len = 0; - for batch_id in 0..num_batches { - // NOTE: atomic_load is needed for correctness on some platforms (see comment above `collision_pairs_len`). - highest_pairs_len = - highest_pairs_len.max(atomic_load_u32(collision_pairs_len.at_mut(batch_id))); - } - - *indirect_args.at_mut(0) = highest_pairs_len.div_ceil(WORKGROUP_SIZE); - *indirect_args.at_mut(1) = num_batches as u32; - *indirect_args.at_mut(2) = 1; + max_len_indirect_args(lid.x, collision_pairs_len, indirect_args, partial); } /// Runs a reduction to compute the AABB of the collider positions. diff --git a/src_rbd_shaders/broad_phase/narrow_phase.rs b/src_rbd_shaders/broad_phase/narrow_phase.rs index 3805d09..0bdd84f 100644 --- a/src_rbd_shaders/broad_phase/narrow_phase.rs +++ b/src_rbd_shaders/broad_phase/narrow_phase.rs @@ -19,55 +19,39 @@ use khal_std::{ sync::{atomic_add_u32, atomic_load_u32}, }; +use super::lbvh::{MAX_REDUCE_LANES, max_len_indirect_args}; use crate::broad_phase::CollisionPair; use crate::utils::{BatchIndices, SliceMut}; use glamx::UVec2; const WORKGROUP_SIZE: u32 = 64; -/// Resets the contacts counter. +/// Resets the contacts counter. One thread per batch. #[spirv_bindgen] -#[spirv(compute(threads(1)))] +#[spirv(compute(threads(64)))] pub fn gpu_reset_narrow_phase( - #[spirv(workgroup_id)] workgroup_id: UVec3, + #[spirv(global_invocation_id)] invocation_id: UVec3, #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] contacts_len: &mut [u32], #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] pfm_pairs_len: &mut [u32], ) { - let batch_id = workgroup_id.y as usize; - - // NOTE: this `for` loop is silly. It doesn’t do anything - // more than a `*contacts_len = 0` in a convoluted - // way because otherwise rustgpu apparently does not generate - // the spirv for this kernel (seems to happen if the kernel is - // too trivial. - for k in 0..1 { - contacts_len.write(batch_id, k); - pfm_pairs_len.write(batch_id, k); + let batch_id = invocation_id.x as usize; + if batch_id < contacts_len.len() { + contacts_len.write(batch_id, 0); + pfm_pairs_len.write(batch_id, 0); } } -/// Initializes indirect dispatch arguments for constraint solver. +/// Initializes indirect dispatch arguments for constraint solver. Dispatch one +/// [`MAX_REDUCE_LANES`]-thread workgroup. #[spirv_bindgen] -#[spirv(compute(threads(1)))] +#[spirv(compute(threads(256)))] pub fn gpu_narrow_phase_init_contacts_dispatch( - // NOTE: the `contacts_len` is mutable here even though we don’t modify it. That’s - // because we access it with an atomic load otherwise it would occasionally read - // stale data (on Windows+Nvidia+wgpu backend). This might be caused by: - // https://github.com/gfx-rs/wgpu/issues/9221 + #[spirv(local_invocation_id)] lid: UVec3, #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] contacts_len: &mut [u32], #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] indirect_args: &mut [u32; 3], + #[spirv(workgroup)] partial: &mut [u32; MAX_REDUCE_LANES as usize], ) { - // For indirect dispatch, get the largest length along all batch dimensions. - let num_batches = contacts_len.len(); - let mut highest_contacts_len = 0; - for i in 0..num_batches { - // NOTE: atomic_load is needed for correctness on some platforms (see comment above `contacts_len`). - highest_contacts_len = highest_contacts_len.max(atomic_load_u32(contacts_len.at_mut(i))); - } - - *indirect_args.at_mut(0) = highest_contacts_len.div_ceil(WORKGROUP_SIZE); - *indirect_args.at_mut(1) = num_batches as u32; - *indirect_args.at_mut(2) = 1; + max_len_indirect_args(lid.x, contacts_len, indirect_args, partial); } const PREDICTION: f32 = 2.0e-3; // TODO: make the prediction configurable. @@ -527,28 +511,17 @@ pub struct NarrowPhasePfmPair { colliders: UVec2, } -/// Initializes PFM-PFM dispatch arguments for constraint solver. +/// Initializes PFM-PFM dispatch arguments for constraint solver. Dispatch one +/// [`MAX_REDUCE_LANES`]-thread workgroup. #[spirv_bindgen] -#[spirv(compute(threads(1)))] +#[spirv(compute(threads(256)))] pub fn gpu_init_pfm_pfm_dispatch( - // NOTE: the `pfm_pairs_len` is mutable here even though we don’t modify it. That’s - // because we access it with an atomic load otherwise it would occasionally read - // stale data (on Windows+Nvidia+wgpu backend). This might be caused by: - // https://github.com/gfx-rs/wgpu/issues/9221 + #[spirv(local_invocation_id)] lid: UVec3, #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] pfm_pairs_len: &mut [u32], #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] indirect_args: &mut [u32; 3], + #[spirv(workgroup)] partial: &mut [u32; MAX_REDUCE_LANES as usize], ) { - let num_batches = pfm_pairs_len.len(); - let mut highest_pfm_pairs_len = 0; - for batch_id in 0..num_batches { - // NOTE: atomic_load is needed for correctness on some platforms (see comment above `pfm_pairs_len`). - highest_pfm_pairs_len = - highest_pfm_pairs_len.max(atomic_load_u32(pfm_pairs_len.at_mut(batch_id))); - } - // TODO PERF: pfm_pfm is very divergent. Use a smaller workgroup size? - *indirect_args.at_mut(0) = highest_pfm_pairs_len.div_ceil(WORKGROUP_SIZE); - *indirect_args.at_mut(1) = num_batches as u32; - *indirect_args.at_mut(2) = 1; + max_len_indirect_args(lid.x, pfm_pairs_len, indirect_args, partial); } #[spirv_bindgen] From 761a89c8118089ae318a824c7423e9ab225afff3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Sun, 26 Jul 2026 13:46:30 +0200 Subject: [PATCH 21/39] =?UTF-8?q?perf:=20brute-force=20O(n=C2=B2)=20broad?= =?UTF-8?q?=20phase=20for=20tiny=20environments?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src_rbd/broad_phase/lbvh.rs | 88 ++++++++++++++- src_rbd/pipeline/rbd_step.rs | 116 ++++++++++++-------- src_rbd_shaders/broad_phase/brute_force.rs | 108 ++++++++++++++++++ src_rbd_shaders/broad_phase/mod.rs | 4 +- src_rbd_shaders/broad_phase/narrow_phase.rs | 4 +- 5 files changed, 270 insertions(+), 50 deletions(-) create mode 100644 src_rbd_shaders/broad_phase/brute_force.rs diff --git a/src_rbd/broad_phase/lbvh.rs b/src_rbd/broad_phase/lbvh.rs index d5b27d4..067a43c 100644 --- a/src_rbd/broad_phase/lbvh.rs +++ b/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, tree: Tensor, 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, } /// 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 { @@ -297,4 +315,68 @@ impl Lbvh { )?; 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, + vertex_buffers: &Tensor, + shapes: &Tensor, + batch_indices: &Tensor, + collision_pairs: &mut Tensor, + collision_pairs_len: &mut Tensor, + collision_pairs_indirect: &mut Tensor<[u32; 3]>, + collision_groups: &Tensor, + 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, + 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; diff --git a/src_rbd/pipeline/rbd_step.rs b/src_rbd/pipeline/rbd_step.rs index c0b9a0a..6b728df 100644 --- a/src_rbd/pipeline/rbd_step.rs +++ b/src_rbd/pipeline/rbd_step.rs @@ -1,6 +1,6 @@ //! The [`RbdPipeline`] running one full simulation step on the GPU. -use crate::broad_phase::{GpuNarrowPhase, Lbvh}; +use crate::broad_phase::{BRUTE_FORCE_MAX_COLLIDERS, GpuNarrowPhase, Lbvh}; #[cfg(feature = "dim3")] use crate::dynamics::GpuMultibodySolver; use crate::dynamics::{ @@ -135,54 +135,82 @@ impl RbdPipeline { drop(pass); - // Build LBVH and find collision pairs. - self.lbvh.update_tree( - backend, - &mut encoder, - &mut state.lbvh, - state.collider_local_poses.len() as u32, - state.num_active_colliders, - state.num_batches, - &state.collider_world_poses, - &state.vertex_buffers, - &state.shapes, - &state.batch_indices, - timestamps.as_deref_mut(), - )?; - - // Debug: validate LBVH topology after tree construction - if crate::VALIDATE_LBVH_TOPOLOGY { + let use_bf = state.num_active_colliders <= BRUTE_FORCE_MAX_COLLIDERS + && std::env::var("NEXUS_DISABLE_BF").is_err(); + if use_bf { + let mut pass = + encoder.begin_pass("[RBD] bf-find-pairs", timestamps.as_deref_mut()); + self.lbvh.brute_force_pairs( + backend, + &mut pass, + &mut state.lbvh, + state.collider_local_poses.len() as u32, + state.num_active_colliders, + state.num_batches, + &state.collider_world_poses, + &state.vertex_buffers, + &state.shapes, + &state.batch_indices, + &mut state.collision_pairs, + &mut state.collision_pairs_len, + &mut state.collision_pairs_indirect, + &state.collision_groups, + &state.pair_filter, + )?; + drop(pass); backend.submit(encoder)?; - - let num_colliders = state.collider_world_poses.len() as u32; - let tree: Vec = - futures::executor::block_on(backend.slow_read_vec(state.lbvh.tree().buffer()))?; - let sorted_colliders: Vec = futures::executor::block_on( - backend.slow_read_vec(state.lbvh.sorted_colliders().buffer()), + } else { + // Build LBVH and find collision pairs. + self.lbvh.update_tree( + backend, + &mut encoder, + &mut state.lbvh, + state.collider_local_poses.len() as u32, + state.num_active_colliders, + state.num_batches, + &state.collider_world_poses, + &state.vertex_buffers, + &state.shapes, + &state.batch_indices, + timestamps.as_deref_mut(), )?; - validate_lbvh_topology(&tree, &sorted_colliders, num_colliders); - encoder = backend.begin_encoding(); - let _pass = - encoder.begin_pass("[RBD] broad-phase-find-pairs", timestamps.as_deref_mut()); - } + // Debug: validate LBVH topology after tree construction + if crate::VALIDATE_LBVH_TOPOLOGY { + backend.submit(encoder)?; + + let num_colliders = state.collider_world_poses.len() as u32; + let tree: Vec = futures::executor::block_on( + backend.slow_read_vec(state.lbvh.tree().buffer()), + )?; + let sorted_colliders: Vec = futures::executor::block_on( + backend.slow_read_vec(state.lbvh.sorted_colliders().buffer()), + )?; + validate_lbvh_topology(&tree, &sorted_colliders, num_colliders); + + encoder = backend.begin_encoding(); + let _pass = encoder + .begin_pass("[RBD] broad-phase-find-pairs", timestamps.as_deref_mut()); + } - let mut pass = encoder.begin_pass("[RBD] lbvh-find-pairs", timestamps.as_deref_mut()); - self.lbvh.find_pairs( - &mut pass, - &mut state.lbvh, - state.num_active_colliders, - state.num_batches, - &state.batch_indices, - &mut state.collision_pairs, - &mut state.collision_pairs_len, - &mut state.collision_pairs_indirect, - &state.collision_groups, - &state.pair_filter, - )?; + let mut pass = + encoder.begin_pass("[RBD] lbvh-find-pairs", timestamps.as_deref_mut()); + self.lbvh.find_pairs( + &mut pass, + &mut state.lbvh, + state.num_active_colliders, + state.num_batches, + &state.batch_indices, + &mut state.collision_pairs, + &mut state.collision_pairs_len, + &mut state.collision_pairs_indirect, + &state.collision_groups, + &state.pair_filter, + )?; - drop(pass); - backend.submit(encoder)?; + drop(pass); + backend.submit(encoder)?; + } } // Phase 2a: Narrow phase. Split out from solver-prep + coloring diff --git a/src_rbd_shaders/broad_phase/brute_force.rs b/src_rbd_shaders/broad_phase/brute_force.rs new file mode 100644 index 0000000..fada07e --- /dev/null +++ b/src_rbd_shaders/broad_phase/brute_force.rs @@ -0,0 +1,108 @@ +//! Brute-force O(n²) broad phase for tiny environments. +//! +//! When each batch holds only a handful of colliders (which happens often +//! in single-robot many-batches simulations), the BVH broad-phase is significantly +//! slower than a naive brute-force approach. +use khal_std::glamx::UVec3; +use khal_std::index::MaybeIndexUnchecked; +use khal_std::macros::{spirv, spirv_bindgen}; +use khal_std::sync::atomic_add_u32; + +use crate::bounding_volumes::Aabb; +use crate::broad_phase::CollisionPair; +use crate::shapes::Shape; +use crate::utils::BatchIndices; +use crate::{PaddedVector, Pose, Vector}; +use glamx::UVec2; +use rapier::geometry::InteractionGroups; + +use super::narrow_phase::PREDICTION; + +/// Computes every active collider's world AABB. +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn gpu_bf_compute_aabbs( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] poses: &[Pose], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] shapes: &[Shape], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] aabbs: &mut [Aabb], + #[spirv(uniform, descriptor_set = 0, binding = 3)] batch_ids: &BatchIndices, + #[spirv(storage_buffer, descriptor_set = 1, binding = 0)] vertices: &[PaddedVector], +) { + let n = batch_ids.colliders_len; + if invocation_id.x >= n * batch_ids.num_batches { + return; + } + let batch_id = invocation_id.x / n; + let i = invocation_id.x % n; + + let poses = batch_ids.coll_batch(batch_id, poses); + let shapes = batch_ids.coll_batch(batch_id, shapes); + let out = batch_ids.coll_start(batch_id) + i as usize; + aabbs.write(out, shapes[i as usize].compute_aabb(poses[i as usize], vertices)); +} + +/// Tests every collider pair of every batch and appends the intersecting, +/// unfiltered ones to `collision_pairs`. +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn gpu_bf_find_pairs( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] aabbs: &[Aabb], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] + collision_pairs: &mut [CollisionPair], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] collision_pairs_len: &mut [u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] + collision_groups: &[InteractionGroups], + #[spirv(uniform, descriptor_set = 0, binding = 4)] batch_ids: &BatchIndices, + #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] pair_filter: &[[u32; 2]], +) { + let n = batch_ids.colliders_len; + let nn = n * n; + if invocation_id.x >= nn * batch_ids.num_batches { + return; + } + let batch_id = invocation_id.x / nn; + let r = invocation_id.x % nn; + let i = r / n; + let j = r % n; + if i >= j { + return; + } + + let collision_groups = batch_ids.coll_batch(batch_id, collision_groups); + let pair_filter = batch_ids.coll_batch(batch_id, pair_filter); + + // Skip pairs whose collision groups don't authorize an interaction. + if !collision_groups[i as usize].test(collision_groups[j as usize]) { + return; + } + + // Built-in filters (same-body and adjacent-multibody-links). + let filter_i = pair_filter[i as usize]; + let filter_j = pair_filter[j as usize]; + if filter_i[0] == filter_j[0] || (filter_i[1] != 0 && filter_i[1] == filter_j[1]) { + return; + } + + let coll_start = batch_ids.coll_start(batch_id); + // Dilate one side by the contact prediction distance. + let mut aabb_i = aabbs.read(coll_start + i as usize); + let dilation = Vector::splat(PREDICTION); + aabb_i.mins -= dilation; + aabb_i.maxs += dilation; + let aabb_j = aabbs.read(coll_start + j as usize); + if !aabb_i.intersects(&aabb_j) { + return; + } + + let target_pair_index = atomic_add_u32(collision_pairs_len.at_mut(batch_id as usize), 1); + + // If we exceed capacity, keep counting the pairs but don’t store any more to avoid overflow. + if target_pair_index < batch_ids.collision_pairs_batch_capacity { + let mut collision_pairs = batch_ids.collision_pairs_batch_mut(batch_id, collision_pairs); + collision_pairs[target_pair_index as usize] = CollisionPair { + colliders: UVec2::new(i, j), + }; + } +} diff --git a/src_rbd_shaders/broad_phase/mod.rs b/src_rbd_shaders/broad_phase/mod.rs index 36d279f..43e94b7 100644 --- a/src_rbd_shaders/broad_phase/mod.rs +++ b/src_rbd_shaders/broad_phase/mod.rs @@ -5,6 +5,7 @@ //! - LBVH (Linear Bounding Volume Hierarchy, for large scenes) // Data structures and algorithms +mod brute_force; mod lbvh; // GPU compute shader kernels @@ -14,12 +15,13 @@ use glamx::UVec2; // Re-export non-spirv items explicitly to avoid ambiguous glob re-exports. // The div_ceil functions have different signatures (u32 vs i32) so we pick one. // Spirv-only items (functions and generated structs) are re-exported via glob. +pub use brute_force::*; pub use lbvh::*; #[cfg(feature = "dim2")] pub use lbvh::{expand_bits_2d, morton_2d}; pub use narrow_phase::*; -#[derive(Copy, Clone, PartialEq, Eq)] +#[derive(Copy, Clone, PartialEq, Eq, Default)] #[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] #[repr(C)] pub struct CollisionPair { diff --git a/src_rbd_shaders/broad_phase/narrow_phase.rs b/src_rbd_shaders/broad_phase/narrow_phase.rs index 0bdd84f..ae9364b 100644 --- a/src_rbd_shaders/broad_phase/narrow_phase.rs +++ b/src_rbd_shaders/broad_phase/narrow_phase.rs @@ -16,7 +16,7 @@ use khal_std::index::MaybeIndexUnchecked; use khal_std::macros::{spirv, spirv_bindgen}; use khal_std::{ iter::StepRng, - sync::{atomic_add_u32, atomic_load_u32}, + sync::atomic_add_u32, }; use super::lbvh::{MAX_REDUCE_LANES, max_len_indirect_args}; @@ -54,7 +54,7 @@ pub fn gpu_narrow_phase_init_contacts_dispatch( max_len_indirect_args(lid.x, contacts_len, indirect_args, partial); } -const PREDICTION: f32 = 2.0e-3; // TODO: make the prediction configurable. +pub(crate) const PREDICTION: f32 = 2.0e-3; // TODO: make the prediction configurable. /// Narrow phase, pass 1 of 2: analytic shape-shape contacts for ball / cuboid /// pairs, written straight into the `contacts` buffer. From 3755e54525c1c25158421cd8418361788f9a8b09 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Sun, 26 Jul 2026 14:13:49 +0200 Subject: [PATCH 22/39] perf: bound the multibody contact scan by the real contact count --- .../multibody/multibody_from_rapier.rs | 1 + .../dynamics/multibody/multibody_solver.rs | 127 ++++++--- src_rbd/dynamics/solver.rs | 264 +++++++++++------- src_rbd/pipeline/rbd_step.rs | 5 +- .../dynamics/multibody/contact_constraints.rs | 30 +- src_rbd_shaders/dynamics/multibody/types.rs | 3 + 6 files changed, 281 insertions(+), 149 deletions(-) diff --git a/src_rbd/dynamics/multibody/multibody_from_rapier.rs b/src_rbd/dynamics/multibody/multibody_from_rapier.rs index 69e7379..7547774 100644 --- a/src_rbd/dynamics/multibody/multibody_from_rapier.rs +++ b/src_rbd/dynamics/multibody/multibody_from_rapier.rs @@ -152,6 +152,7 @@ impl GpuMultibodySet { max_constraints, self_contacts_enabled: if mb.self_contacts_enabled() { 1 } else { 0 }, contact_constraint_count: 0, + batch_contacts_len: 0, }); // `assembly_id` is not exposed publicly on `MultibodyLink`, so we diff --git a/src_rbd/dynamics/multibody/multibody_solver.rs b/src_rbd/dynamics/multibody/multibody_solver.rs index ba3f0c6..464723a 100644 --- a/src_rbd/dynamics/multibody/multibody_solver.rs +++ b/src_rbd/dynamics/multibody/multibody_solver.rs @@ -10,7 +10,7 @@ use crate::shaders::dynamics::{ GpuMbGravityAndLuT16, GpuMbGravityAndLuT32, GpuMbInitContactConstraints, GpuMbInitJointConstraints, GpuMbIntegrate, GpuMbIntegrateVelocities, GpuMbRemoveContactConstraintBias, GpuMbRemoveImpulseJointConstraintBias, - GpuMbResetContactWarmstart, GpuMbWarmstartContactConstraints, + GpuMbResetContactWarmstart, GpuMbStashContactsLen, GpuMbWarmstartContactConstraints, GpuMbRemoveSolveJointNoBias, GpuMbSolveContactConstraints, GpuMbSolveImpulseJointConstraints, GpuMbFinalizeImpulseJointConstraints, GpuMbSolveJointConstraints, GpuMbUpdateImpulseJointConstraints, Velocity, WorldMassProperties, @@ -42,6 +42,10 @@ pub struct GpuMultibodySolver { solve_contact_constraints: GpuMbSolveContactConstraints, /// Zero the accumulated contact impulses once per frame (warmstart reset). reset_contact_warmstart: GpuMbResetContactWarmstart, + /// Copy `contacts_len[batch]` into each `MultibodyInfo` once per step so + /// `init_contact_constraints` (at the 8-storage-buffer limit) can bound + /// its manifold scan by the actual count instead of the capacity. + stash_contacts_len: GpuMbStashContactsLen, /// Re-apply the accumulated contact impulse each substep (warmstart). warmstart_contact_constraints: GpuMbWarmstartContactConstraints, remove_contact_constraint_bias: GpuMbRemoveContactConstraintBias, @@ -123,6 +127,29 @@ impl GpuMultibodySolver { self.compute_dynamics(pass, mb, args) } + /// Copy `contacts_len[batch]` into each `MultibodyInfo`. + /// + /// This is a workaround for kernels that are already at the 8-storage-binding + /// web limit and could therefore not bind `contacts_len`. + pub fn stash_contacts_len( + &self, + pass: &mut GpuPass, + mb: &mut GpuMultibodySet, + args: &mut MultibodySolverArgs<'_>, + ) -> Result<(), GpuBackendError> { + if mb.is_empty() { + return Ok(()); + } + self.stash_contacts_len.call( + pass, + mb.flat_mb_dispatch(), + &mut mb.multibody_info, + args.contacts_len, + args.batch_indices, + )?; + Ok(()) + } + // Per-substep work is split into five phases so the pipeline can interleave // them with the rigid-body substep: `substep_integrate_velocities` (P1), // `substep_build_constraints` (P2), `substep_solve_with_bias` (P3), @@ -154,25 +181,21 @@ impl GpuMultibodySolver { /// constraints, then warmstart the contacts. pub fn substep_build_constraints( &self, - pass: &mut GpuPass, + encoder: &mut khal::backend::GpuEncoder, + mut timestamps: Option<&mut khal::backend::GpuTimestamps>, mb: &mut GpuMultibodySet, args: &mut MultibodySolverArgs<'_>, ) -> Result<(), GpuBackendError> { + use khal::backend::Encoder; if mb.is_empty() { return Ok(()); } let dispatch = mb.flat_mb_dispatch(); if mb.has_joint_constraints { - // TODO(PERF): joints init could parallelized. We either need to rework - // the flow of the kernel to that the LU parts are not in - // potentially diverging code paths, or we need to have - // each link have its limits/motors generated by a separate - // threadgroups (which might actually be better for lower - // divergence and allow us to theadgroup-parallelize the LU - // solve). + let mut pass = encoder.begin_pass("[RBD] mbb/init-joint", timestamps.as_deref_mut()); self.init_joint_with_bias.call( - pass, + &mut pass, dispatch, &mb.multibody_info, &mb.links_static, @@ -190,52 +213,64 @@ impl GpuMultibodySolver { // multibody pairs only). `init` PRESERVES the accumulated impulse across // substeps (zeroed once per frame by `reset_contact_warmstart` in // `init_step`); `finalize` recomputes `inv_lhs` and the M⁻¹Jᵀ columns. - self.init_contact_constraints.call( - pass, - dispatch, - &mut mb.multibody_info, - &mb.body_jacobians, - &mb.body_to_link, - &mut mb.contact_constraints, - &mut mb.contact_constraint_jacs, - &mb.constraint_softness, - args.batch_indices, - args.mprops, - args.collider_world_poses, - args.contacts, - )?; + { + let mut pass = + encoder.begin_pass("[RBD] mbb/init-contact", timestamps.as_deref_mut()); + self.init_contact_constraints.call( + &mut pass, + dispatch, + &mut mb.multibody_info, + &mb.body_jacobians, + &mb.body_to_link, + &mut mb.contact_constraints, + &mut mb.contact_constraint_jacs, + &mb.constraint_softness, + args.batch_indices, + args.mprops, + args.collider_world_poses, + args.contacts, + )?; + } // One 64-lane workgroup per multibody: the per-constraint LU // back-solves are independent, so they run one-per-lane instead of // sequentially on a single thread. - let finalize_dispatch = [mb.multibodies_per_batch * MB_LU_LANES, mb.num_batches, 1]; - self.finalize_contact_constraints.call( - pass, - finalize_dispatch, - &mb.multibody_info, - &mb.mass_matrices, - &mb.lu_pivots, - &mut mb.contact_constraints, - &mb.contact_constraint_jacs, - &mut mb.contact_constraint_columns, - args.batch_indices, - )?; + { + let mut pass = + encoder.begin_pass("[RBD] mbb/finalize-contact", timestamps.as_deref_mut()); + let finalize_dispatch = [mb.multibodies_per_batch * MB_LU_LANES, mb.num_batches, 1]; + self.finalize_contact_constraints.call( + &mut pass, + finalize_dispatch, + &mb.multibody_info, + &mb.mass_matrices, + &mb.lu_pivots, + &mut mb.contact_constraints, + &mb.contact_constraint_jacs, + &mut mb.contact_constraint_columns, + args.batch_indices, + )?; + } // Warmstart: re-apply the accumulated contact impulse to dof_state (and // the free-body solver velocities) so the contact starts "warm" each // substep — mirrors rapier's per-substep `contact_constraints.warmstart` // and matches what the rigid-body solver does for free contacts. On the // first substep the impulse was just reset to 0, so this is a no-op. - self.warmstart_contact_constraints.call( - pass, - dispatch, - &mb.multibody_info, - &mb.contact_constraints, - &mb.contact_constraint_columns, - &mut mb.dof_state, - args.solver_vels, - args.batch_indices, - )?; + { + let mut pass = + encoder.begin_pass("[RBD] mbb/warmstart-contact", timestamps.as_deref_mut()); + self.warmstart_contact_constraints.call( + &mut pass, + dispatch, + &mb.multibody_info, + &mb.contact_constraints, + &mb.contact_constraint_columns, + &mut mb.dof_state, + args.solver_vels, + args.batch_indices, + )?; + } Ok(()) } diff --git a/src_rbd/dynamics/solver.rs b/src_rbd/dynamics/solver.rs index 67c8b9c..8e22384 100644 --- a/src_rbd/dynamics/solver.rs +++ b/src_rbd/dynamics/solver.rs @@ -19,7 +19,7 @@ use crate::shaders::dynamics::{ }; use crate::utils::{GpuPrefixSum, PrefixSumWorkspace}; use khal::Shader; -use khal::backend::{GpuBackend, GpuBackendError, GpuPass}; +use khal::backend::{Encoder, GpuBackend, GpuBackendError, GpuEncoder, GpuPass, GpuTimestamps}; use vortx::tensor::Tensor; /// GPU shader bundle for the constraint solver. @@ -229,7 +229,8 @@ impl GpuSolver { /// `velocity_solver::solve_constraints`. pub fn solve_tgs<'a>( &self, - pass: &mut GpuPass, + encoder: &mut GpuEncoder, + mut timestamps: Option<&mut GpuTimestamps>, joint_solver: &GpuJointSolver, args: SolverArgs<'a>, mut joint_args: JointSolverArgs<'a>, @@ -245,16 +246,36 @@ impl GpuSolver { /* * Init solver vel increments. */ - self.init_solver_vels_inc.call( - pass, - [args.num_colliders, args.num_batches, 1], - args.solver_vels_inc, - args.mprops, - args.sim_params, - args.batch_indices, - )?; + { + let mut pass = encoder.begin_pass("[RBD] slv/init", timestamps.as_deref_mut()); + self.init_solver_vels_inc.call( + &mut pass, + [args.num_colliders, args.num_batches, 1], + args.solver_vels_inc, + args.mprops, + args.sim_params, + args.batch_indices, + )?; + + joint_solver.init(&mut pass, &mut joint_args)?; - joint_solver.init(pass, &mut joint_args)?; + // Bound for the per-substep contact scans: runs after the narrow + // phase wrote `contacts_len`, before the first substep build. + #[cfg(feature = "dim3")] + if let (Some(solver), Some(state)) = (mb_solver, mb_state.as_deref_mut()) { + let mut mb_args = MultibodySolverArgs { + poses: &mut *args.solver_body_poses, + collider_world_poses: args.collider_world_poses, + mprops: args.mprops, + contacts: args.contacts, + contacts_len: args.contacts_len, + solver_vels: &mut *args.solver_vels, + batch_indices: args.batch_indices, + color_uniforms: args.color_uniforms, + }; + solver.stash_contacts_len(&mut pass, state, &mut mb_args)?; + } + } // Per substep, the multibody work is split into five phases that are // INTERLEAVED with the matching rigid-body phases, mirroring rapier's @@ -274,9 +295,10 @@ impl GpuSolver { // of `args.solver_vels` / `args.solver_body_poses` must be released // before the interleaved rigid-body call that touches the same buffers. macro_rules! mb_phase { - ($method:ident $(, $extra:expr)*) => {{ + ($label:expr, $method:ident $(, $extra:expr)*) => {{ #[cfg(feature = "dim3")] if let (Some(solver), Some(state)) = (mb_solver, mb_state.as_deref_mut()) { + let mut pass = encoder.begin_pass($label, timestamps.as_deref_mut()); let mut mb_args = MultibodySolverArgs { poses: &mut *args.solver_body_poses, collider_world_poses: args.collider_world_poses, @@ -287,7 +309,7 @@ impl GpuSolver { batch_indices: args.batch_indices, color_uniforms: args.color_uniforms, }; - solver.$method(pass, state, &mut mb_args $(, $extra)*)?; + solver.$method(&mut pass, state, &mut mb_args $(, $extra)*)?; } }}; } @@ -301,47 +323,99 @@ 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, - )?; + mb_phase!("[RBD] slv/mb-integrate-vels", substep_integrate_velocities); + { + let mut pass = + encoder.begin_pass("[RBD] slv/rb-apply-inc", timestamps.as_deref_mut()); + self.apply_solver_vels_inc.call( + &mut 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)?; - if args.colorless_warmstart { - // One gather dispatch over bodies instead of `num_colors` - // scatter dispatches (each constraint is visited once per - // body side, but the dispatch count drops by ~num_colors). - self.warmstart_without_colors.call( + { + #[cfg(feature = "dim3")] + if let (Some(solver), Some(state)) = (mb_solver, mb_state.as_deref_mut()) { + let mut mb_args = MultibodySolverArgs { + poses: &mut *args.solver_body_poses, + collider_world_poses: args.collider_world_poses, + mprops: args.mprops, + contacts: args.contacts, + contacts_len: args.contacts_len, + solver_vels: &mut *args.solver_vels, + batch_indices: args.batch_indices, + color_uniforms: args.color_uniforms, + }; + solver.substep_build_constraints( + encoder, + timestamps.as_deref_mut(), + state, + &mut mb_args, + )?; + } + } + { + let mut pass = + encoder.begin_pass("[RBD] slv/rb-build-warmstart", timestamps.as_deref_mut()); + let pass = &mut pass; + self.update_constraints.call( pass, - [args.num_colliders, args.num_batches, 1], - args.body_constraint_counts, - args.body_constraint_ids, + args.contacts_len_indirect, args.constraints, - args.solver_vels, + args.constraint_builders, + args.contacts_len, + args.solver_body_poses, + args.sim_params, args.batch_indices, )?; - } else { - // NOTE: contact colors start at 1 (0 = unassigned). + joint_solver.update(pass, &mut joint_args, args.solver_body_poses)?; + if args.colorless_warmstart { + // One gather dispatch over bodies instead of `num_colors` + // scatter dispatches (each constraint is visited once per + // body side, but the dispatch count drops by ~num_colors). + self.warmstart_without_colors.call( + pass, + [args.num_colliders, args.num_batches, 1], + args.body_constraint_counts, + args.body_constraint_ids, + args.constraints, + args.solver_vels, + args.batch_indices, + )?; + } else { + // NOTE: contact colors start at 1 (0 = unassigned). + for c in 1..=args.num_colors { + self.warmstart.call( + pass, + args.contacts_len_indirect, + args.constraints, + args.solver_vels, + args.color_bucket_starts, + args.color_sorted_ids, + &args.color_uniforms[c as usize], + args.batch_indices, + )?; + } + } + } + + /* + * Solve all joints + contacts with bias. + */ + mb_phase!("[RBD] slv/mb-solve-bias", substep_solve_with_bias); + { + let mut pass = + encoder.begin_pass("[RBD] slv/rb-solve-bias", timestamps.as_deref_mut()); + let pass = &mut pass; + joint_solver.solve(pass, &mut joint_args, args.solver_vels, true)?; for c in 1..=args.num_colors { - self.warmstart.call( + self.step_gauss_seidel.call( pass, args.contacts_len_indirect, args.constraints, @@ -350,61 +424,56 @@ impl GpuSolver { args.color_sorted_ids, &args.color_uniforms[c as usize], args.batch_indices, + // use_bias = 1 (the `color_uniform[1]` contains the value 1) + &args.color_uniforms[1], )?; } } /* - * Solve all joints + contacts with bias. + * Integrate all positions once. */ - mb_phase!(substep_solve_with_bias); - joint_solver.solve(pass, &mut joint_args, args.solver_vels, true)?; - for c in 1..=args.num_colors { - self.step_gauss_seidel.call( - pass, - args.contacts_len_indirect, - args.constraints, + mb_phase!( + "[RBD] slv/mb-integrate-pos", + substep_integrate_positions, + is_last_substep + ); + { + let mut pass = + encoder.begin_pass("[RBD] slv/rb-integrate", timestamps.as_deref_mut()); + self.integrate_linearized.call( + &mut pass, + [args.num_colliders, args.num_batches, 1], + args.solver_body_poses, args.solver_vels, - args.color_bucket_starts, - args.color_sorted_ids, - &args.color_uniforms[c as usize], + args.sim_params, args.batch_indices, - // use_bias = 1 (`color_uniforms[c]` holds the constant `c`). - &args.color_uniforms[1], )?; } - /* - * P4/F4 — integrate ALL positions once. - */ - mb_phase!(substep_integrate_positions, is_last_substep); - self.integrate_linearized.call( - pass, - [args.num_colliders, args.num_batches, 1], - args.solver_body_poses, - args.solver_vels, - args.sim_params, - args.batch_indices, - )?; - /* * P5/F5 — solve ALL joints + contacts WITHOUT bias (stabilization). */ - mb_phase!(substep_solve_no_bias); - joint_solver.solve(pass, &mut joint_args, args.solver_vels, false)?; - for c in 1..=args.num_colors { - self.step_gauss_seidel.call( - pass, - args.contacts_len_indirect, - args.constraints, - args.solver_vels, - args.color_bucket_starts, - args.color_sorted_ids, - &args.color_uniforms[c as usize], - args.batch_indices, - // use_bias = 0 (`color_uniforms[c]` holds the constant `c`). - &args.color_uniforms[0], - )?; + mb_phase!("[RBD] slv/mb-solve-nobias", substep_solve_no_bias); + { + let mut pass = + encoder.begin_pass("[RBD] slv/rb-solve-nobias", timestamps.as_deref_mut()); + let pass = &mut pass; + joint_solver.solve(pass, &mut joint_args, args.solver_vels, false)?; + for c in 1..=args.num_colors { + self.step_gauss_seidel.call( + pass, + args.contacts_len_indirect, + args.constraints, + args.solver_vels, + args.color_bucket_starts, + args.color_sorted_ids, + &args.color_uniforms[c as usize], + args.batch_indices, + // use_bias = 0 (the `color_uniform[0]` contains the value 0) + &args.color_uniforms[0], + )?; + } } } @@ -412,16 +481,19 @@ impl GpuSolver { * Writeback body velocities and convert COM-centered solver poses * back to body-origin poses. */ - self.finalize.call( - pass, - [args.num_colliders, args.num_batches, 1], - args.vels, - args.solver_vels, - args.body_poses, - args.solver_body_poses, - args.local_mprops, - args.batch_indices, - )?; + { + let mut pass = encoder.begin_pass("[RBD] slv/finalize", timestamps.as_deref_mut()); + self.finalize.call( + &mut pass, + [args.num_colliders, args.num_batches, 1], + args.vels, + args.solver_vels, + args.body_poses, + args.solver_body_poses, + args.local_mprops, + args.batch_indices, + )?; + } Ok(()) } diff --git a/src_rbd/pipeline/rbd_step.rs b/src_rbd/pipeline/rbd_step.rs index 6b728df..9f0f64e 100644 --- a/src_rbd/pipeline/rbd_step.rs +++ b/src_rbd/pipeline/rbd_step.rs @@ -437,7 +437,6 @@ impl RbdPipeline { { let mut encoder = backend.begin_encoding(); - let mut pass = encoder.begin_pass("[RBD] solver", timestamps.as_deref_mut()); #[cfg(feature = "dim3")] let mb = if state.multibodies.is_empty() { None @@ -445,14 +444,14 @@ impl RbdPipeline { Some((&self.multibody_solver, &mut state.multibodies)) }; self.solver.solve_tgs( - &mut pass, + &mut encoder, + timestamps.as_deref_mut(), &self.joint_solver, solver_args, joint_solver_args, #[cfg(feature = "dim3")] mb, )?; - drop(pass); // Resolve all accumulated timestamps before the final submit. if let Some(ts) = ×tamps { diff --git a/src_rbd_shaders/dynamics/multibody/contact_constraints.rs b/src_rbd_shaders/dynamics/multibody/contact_constraints.rs index b17db0e..00a35f4 100644 --- a/src_rbd_shaders/dynamics/multibody/contact_constraints.rs +++ b/src_rbd_shaders/dynamics/multibody/contact_constraints.rs @@ -181,10 +181,7 @@ pub fn gpu_mb_init_contact_constraints( col_start + (mb_idx as usize) * (MAX_MB_CONTACT_CONSTRAINTS_PER_MB as usize) * dofs_stride; let contacts_slice = batch_ids.contact_batch(batch_id, contacts); - // Iterate to capacity (instead of reading a `contacts_len` storage buffer): - // empty slots have `contact.len == 0` and are skipped. Drops one binding to - // fit the 8-storage-buffer WebGPU limit; matches `gpu_solver_init_constraints`. - let n_contacts = batch_ids.contacts_batch_capacity; + let n_contacts = mb.batch_contacts_len.min(batch_ids.contacts_batch_capacity); let mut count = 0u32; for ci in 0..n_contacts { @@ -548,6 +545,31 @@ pub fn gpu_mb_init_contact_constraints( multibody_info.write(mb_start + mb_idx as usize, mb); } +/// HACK: stash `contacts_len[batch]` into each multibody's `batch_contacts_len`. +/// +/// This exists only to work around the web 8-storage-bindings limit for kernels +/// that bind multibodies but don’t have any room left to bind `contacts_len`. +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn gpu_mb_stash_contacts_len( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + multibody_info: &mut [MultibodyInfo], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] contacts_len: &[u32], + #[spirv(uniform, descriptor_set = 0, binding = 2)] batch_ids: &BatchIndices, +) { + let num_mb = batch_ids.multibodies_len; + if invocation_id.x >= num_mb * batch_ids.num_batches { + return; + } + let batch_id = invocation_id.x / num_mb; + let mb_idx = invocation_id.x % num_mb; + let mb_start = batch_ids.mb_start(batch_id); + let mut mb = multibody_info.read(mb_start + mb_idx as usize); + mb.batch_contacts_len = contacts_len.read(batch_id as usize); + multibody_info.write(mb_start + mb_idx as usize, mb); +} + /// Zero the accumulated impulse of every contact-constraint slot for each /// multibody. Called ONCE per visible frame (from `init_step`, before the /// substep loop) so the first substep's warmstart starts cold; within a frame diff --git a/src_rbd_shaders/dynamics/multibody/types.rs b/src_rbd_shaders/dynamics/multibody/types.rs index f06e2ee..235eaea 100644 --- a/src_rbd_shaders/dynamics/multibody/types.rs +++ b/src_rbd_shaders/dynamics/multibody/types.rs @@ -351,4 +351,7 @@ pub struct MultibodyInfo { /// multibody. Written by `gpu_mb_init_contact_constraints`, read by the /// warmstart / finalize / solve / remove-bias contact kernels. pub contact_constraint_count: u32, + /// Per-step copy of `contacts_len[batch]` (to work around the web 8 storage + /// bindings count limit). + pub batch_contacts_len: u32, } From ab4c07736891b7a3d6975c2ee78650702fcfae44 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Sun, 26 Jul 2026 14:25:33 +0200 Subject: [PATCH 23/39] perf: lane-parallelize the multibody joint-constraint LU back-solves --- .../dynamics/multibody/multibody_solver.rs | 6 +- .../dynamics/multibody/joint_constraints.rs | 269 ++++++++---------- 2 files changed, 123 insertions(+), 152 deletions(-) diff --git a/src_rbd/dynamics/multibody/multibody_solver.rs b/src_rbd/dynamics/multibody/multibody_solver.rs index 464723a..0266d2d 100644 --- a/src_rbd/dynamics/multibody/multibody_solver.rs +++ b/src_rbd/dynamics/multibody/multibody_solver.rs @@ -194,9 +194,13 @@ impl GpuMultibodySolver { if mb.has_joint_constraints { let mut pass = encoder.begin_pass("[RBD] mbb/init-joint", timestamps.as_deref_mut()); + // One 64-lane workgroup per multibody: lane 0 emits the constraint + // metadata serially (cheap), then the per-constraint M⁻¹-column LU + // back-solves run one-per-lane instead of sequentially. + let init_joint_dispatch = [mb.multibodies_per_batch * MB_LU_LANES, mb.num_batches, 1]; self.init_joint_with_bias.call( &mut pass, - dispatch, + init_joint_dispatch, &mb.multibody_info, &mb.links_static, &mb.links_workspace, diff --git a/src_rbd_shaders/dynamics/multibody/joint_constraints.rs b/src_rbd_shaders/dynamics/multibody/joint_constraints.rs index cb60d12..a248502 100644 --- a/src_rbd_shaders/dynamics/multibody/joint_constraints.rs +++ b/src_rbd_shaders/dynamics/multibody/joint_constraints.rs @@ -6,7 +6,9 @@ use khal_std::glamx::UVec3; use khal_std::index::MaybeIndexUnchecked; +use khal_std::iter::StepRng; use khal_std::macros::{spirv, spirv_bindgen}; +use khal_std::sync::control_barrier; use crate::dynamics::ConstraintSoftness; use crate::dynamics::joint::SPATIAL_DIM; @@ -109,41 +111,24 @@ fn solve_joint_constraints_body( } } +/// Serial (lane-0) emission walk: writes the metadata of every active +/// limit/motor constraint slot. The expensive M⁻¹-column back-solves happen +/// afterwards, lane-parallel, in `gpu_mb_init_joint_constraints`' finalize +/// stage. Slot zeroing also happens there (lane-parallel, before this walk). #[inline] -fn init_joint_constraints_body( - multibody_info: &[MultibodyInfo], +fn emit_joint_constraints( links_static: &[MultibodyLinkStatic], links_workspace: &[MultibodyLinkWorkspace], - mass_matrices: &[f32], - lu_pivots: &[u32], joint_constraints: &mut [MultibodyJointConstraint], - joint_constraint_columns: &mut [f32], + mb: &MultibodyInfo, + cons_base: usize, batch_id: u32, - mb_idx: u32, dt: f32, joint_erp_inv_dt: f32, joint_cfm_coeff: f32, batch_ids: &BatchIndices, ) { - let mb = batch_ids - .mb_batch(batch_id, multibody_info) - .read(mb_idx as usize); let num_links = mb.num_links; - let ndofs = mb.ndofs; - if ndofs == 0 { - return; - } - let mb_mm_base = batch_ids.mm_start(batch_id) + mb.mass_matrix_offset as usize; - let piv_offset = batch_ids.dof_start(batch_id) + mb.first_dof as usize; - let cons_base = batch_ids.mb_joint_constraints_start(batch_id) + mb.first_constraint as usize; - // One column of M⁻¹ per constraint slot — `dof_batch_capacity` floats - // per slot (only the first `ndofs` of each are meaningful, but we use - // the batch-wide max as the stride to match the host allocation - // `cons_col_cap = cons_cap * dofs_cap` and to avoid two multibodies - // with different ndofs stomping on each other's columns). - let dofs_stride = batch_ids.dof_batch_capacity as usize; - let col_base = batch_ids.mb_joint_constraint_columns_start(batch_id) - + (mb.first_constraint as usize) * dofs_stride; let stat_slice = batch_ids .mb_links_batch(batch_id, links_static) @@ -151,14 +136,6 @@ fn init_joint_constraints_body( let ws_slice = batch_ids .mb_links_batch(batch_id, links_workspace) .offset(mb.first_link as usize); - let m = MatSlice::dense(mb_mm_base, ndofs, ndofs); - - for s in 0..mb.max_constraints { - let mut cz: MultibodyJointConstraint = joint_constraints.read(cons_base + s as usize); - cz.kind = 0; - cz.impulse = 0.0; - joint_constraints.write(cons_base + s as usize, cz); - } let inv_dt = if dt != 0.0 { 1.0 / dt } else { 0.0 }; @@ -196,13 +173,9 @@ fn init_joint_constraints_body( let limit_max = stat.data.limits[axis as usize].max; emit_motor_constraint( joint_constraints, - joint_constraint_columns, cons_base, - col_base, - dofs_stride, slot, abs_dof, - ndofs, curr_pos, inv_dt, dt, @@ -210,23 +183,15 @@ fn init_joint_constraints_body( has_limits, limit_min, limit_max, - mass_matrices, - m, - lu_pivots, - piv_offset, ); slot += 1; } if (limit_axes & (1 << axis)) != 0 { emit_limit_constraint( joint_constraints, - joint_constraint_columns, cons_base, - col_base, - dofs_stride, slot, abs_dof, - ndofs, curr_pos, [ stat.data.limits[axis as usize].min, @@ -234,10 +199,6 @@ fn init_joint_constraints_body( ], joint_erp_inv_dt, joint_cfm_coeff, - mass_matrices, - m, - lu_pivots, - piv_offset, ); slot += 1; } @@ -255,13 +216,9 @@ fn init_joint_constraints_body( if (limit_axes & (1 << axis)) != 0 { emit_limit_constraint( joint_constraints, - joint_constraint_columns, cons_base, - col_base, - dofs_stride, slot, abs_dof, - ndofs, curr_pos, [ stat.data.limits[axis as usize].min, @@ -269,10 +226,6 @@ fn init_joint_constraints_body( ], joint_erp_inv_dt, joint_cfm_coeff, - mass_matrices, - m, - lu_pivots, - piv_offset, ); slot += 1; } @@ -282,13 +235,9 @@ fn init_joint_constraints_body( let limit_max = stat.data.limits[axis as usize].max; emit_motor_constraint( joint_constraints, - joint_constraint_columns, cons_base, - col_base, - dofs_stride, slot, abs_dof, - ndofs, curr_pos, inv_dt, dt, @@ -296,10 +245,6 @@ fn init_joint_constraints_body( has_limits, limit_min, limit_max, - mass_matrices, - m, - lu_pivots, - piv_offset, ); slot += 1; } @@ -345,24 +290,21 @@ fn inv(x: f32) -> f32 { /// Initialize a single limit constraint slot. Mirrors rapier's /// `unit_joint_limit_constraint`. +/// +/// Emits METADATA ONLY: `inv_lhs` is left 0 and `cfm_gain` holds the +/// pre-fold gain (0 for limits); the lane-parallel finalize stage of +/// `gpu_mb_init_joint_constraints` back-solves the M⁻¹ column and applies +/// rapier's `finalize_generic_constraints` fold. #[inline] fn emit_limit_constraint( joint_constraints: &mut [MultibodyJointConstraint], - joint_constraint_columns: &mut [f32], cons_base: usize, - col_base: usize, - dofs_stride: usize, slot: u32, dof_id: u32, - ndofs: u32, curr_pos: f32, limits: [f32; 2], erp_inv_dt: f32, cfm_coeff: f32, - mass_matrices: &[f32], - m: MatSlice, - lu_pivots: &[u32], - piv_offset: usize, ) { // rapier (`limit_*` builder): erp_inv_dt = joint.softness.erp_inv_dt(dt), // cfm_coeff = joint.softness.cfm_coeff(dt), cfm_gain = 0 — configurable via @@ -379,27 +321,6 @@ fn emit_limit_constraint( let rhs_bias = (hi_excess - lo_excess) * erp_inv_dt; let rhs_wo_bias = 0.0f32; - let lhs = compute_constraint_column( - joint_constraint_columns, - col_base, - slot, - dofs_stride, - ndofs, - dof_id, - mass_matrices, - m, - lu_pivots, - piv_offset, - ); - // rapier `finalize_generic_constraints` (the multibody-internal constraints - // ARE finalized, after `unit_joint_limit_constraint` sets the preliminary - // 1/lhs): cfm_gain = lhs·cfm_coeff + cfm_gain_init; inv_lhs = 1/(lhs + - // cfm_gain). The PGS sweep then applies `cfm_gain` directly. For limits - // cfm_gain_init = 0 and (near-rigid) cfm_coeff ≈ 2e-6, so this is ~1/lhs — - // but we replicate the formula exactly so it's correct for any softness. - let cfm_gain = lhs * cfm_coeff; - let inv_lhs = inv(lhs + cfm_gain); - let max_neg_impulse = if min_enabled { -MAX_FLT } else { 0.0 }; let max_pos_impulse = if max_enabled { MAX_FLT } else { 0.0 }; @@ -410,29 +331,24 @@ fn emit_limit_constraint( _pad0: 0, rhs: rhs_wo_bias + rhs_bias, rhs_wo_bias, - inv_lhs, + inv_lhs: 0.0, impulse: 0.0, impulse_lo: max_neg_impulse, impulse_hi: max_pos_impulse, cfm_coeff, - cfm_gain, + // This will be calculated in the finalize (orthogonalization) step. + cfm_gain: 0.0, }; joint_constraints.write(cons_base + slot as usize, cons); } -/// Initialize a single motor constraint slot. Mirrors rapier's -/// `unit_joint_motor_constraint`. `has_limits` + `(limit_min, limit_max)` flatten -/// rapier's `Option<[Real; 2]>` parameter (rust-gpu can't represent enums). +/// Initialize a single motor constraint slot.. #[inline] fn emit_motor_constraint( joint_constraints: &mut [MultibodyJointConstraint], - joint_constraint_columns: &mut [f32], cons_base: usize, - col_base: usize, - dofs_stride: usize, slot: u32, dof_id: u32, - ndofs: u32, curr_pos: f32, inv_dt: f32, dt: f32, @@ -440,10 +356,6 @@ fn emit_motor_constraint( has_limits: bool, limit_min: f32, limit_max: f32, - mass_matrices: &[f32], - m: MatSlice, - lu_pivots: &[u32], - piv_offset: usize, ) { let (erp_inv_dt, cfm_coeff, cfm_gain, _, max_impulse) = motor_params(motor, dt); @@ -465,30 +377,6 @@ fn emit_motor_constraint( } rhs_wo_bias += -target_vel; - let lhs = compute_constraint_column( - joint_constraint_columns, - col_base, - slot, - dofs_stride, - ndofs, - dof_id, - mass_matrices, - m, - lu_pivots, - piv_offset, - ); - // rapier `finalize_generic_constraints` (run after `unit_joint_motor_ - // constraint` sets the preliminary 1/lhs): cfm_gain = lhs·cfm_coeff + - // cfm_gain_init; inv_lhs = 1/(lhs + cfm_gain). The PGS sweep then applies - // `cfm_gain` (`impulse + inv_lhs·(rhs - cfm_gain·impulse)`). This is the - // ACCELERATION-based motor's compliance: for an acceleration-based position - // servo (``), `cfm_coeff` is large (∝ 1/(dt²·stiffness)), so the - // fold dominates and makes the effective gain inertia-independent — without - // it the servo is far too weak and the robot sags. `cfm_gain` here is - // `motor_params.cfm_gain` (nonzero only for force-based motors). - let cfm_gain = lhs * cfm_coeff + cfm_gain; - let inv_lhs = inv(lhs + cfm_gain); - let cons = MultibodyJointConstraint { dof_id, kind: 2, @@ -496,7 +384,7 @@ fn emit_motor_constraint( _pad0: 0, rhs: rhs_wo_bias, rhs_wo_bias, - inv_lhs, + inv_lhs: 0.0, impulse: 0.0, impulse_lo: -max_impulse, impulse_hi: max_impulse, @@ -514,10 +402,17 @@ fn emit_motor_constraint( /// /// Must run after `gpu_mb_lu_decompose` — the LU factors of `M` are used to compute /// the per-constraint M⁻¹ column and effective inverse mass. +/// +/// One 64-lane workgroup per (multibody, batch), in three stages: +/// 1. lane-parallel: zero all constraint slots; +/// 2. lane 0: the serial link walk emitting constraint metadata (cheap); +/// 3. lane-parallel: one M⁻¹-column LU back-solve per emitted slot plus +/// rapier's `finalize_generic_constraints`. #[spirv_bindgen] #[spirv(compute(threads(64)))] pub fn gpu_mb_init_joint_constraints( - #[spirv(global_invocation_id)] invocation_id: UVec3, + #[spirv(workgroup_id)] workgroup_id: UVec3, + #[spirv(local_invocation_id)] local_id: UVec3, #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] multibody_info: &[MultibodyInfo], #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] links_static: &[MultibodyLinkStatic], @@ -532,28 +427,100 @@ pub fn gpu_mb_init_joint_constraints( #[spirv(uniform, descriptor_set = 0, binding = 7)] softness: &ConstraintSoftness, #[spirv(uniform, descriptor_set = 0, binding = 8)] batch_ids: &BatchIndices, ) { - // Flattened (multibody, batch) grid — see `BatchIndices::num_batches`. + const LANES: u32 = 64; + + // One workgroup per (multibody, batch): grid `[mbs · LANES, batches, 1]`. + let batch_id = workgroup_id.y; + let mb_idx = workgroup_id.x; + let lane = local_id.x; let num_mb = batch_ids.multibodies_len; - if invocation_id.x >= num_mb * batch_ids.num_batches { + if mb_idx >= num_mb { return; } - let batch_id = invocation_id.x / num_mb; - let mb_idx = invocation_id.x % num_mb; - init_joint_constraints_body( - multibody_info, - links_static, - links_workspace, - mass_matrices, - lu_pivots, - joint_constraints, - joint_constraint_columns, - batch_id, - mb_idx, - softness.dt, - softness.joint_erp_inv_dt, - softness.joint_cfm_coeff, - batch_ids, - ); + + let mb = batch_ids + .mb_batch(batch_id, multibody_info) + .read(mb_idx as usize); + let ndofs = mb.ndofs; + // Uniform per workgroup: every lane of this group returns together. + if ndofs == 0 { + return; + } + + let mb_mm_base = batch_ids.mm_start(batch_id) + mb.mass_matrix_offset as usize; + let piv_offset = batch_ids.dof_start(batch_id) + mb.first_dof as usize; + let cons_base = batch_ids.mb_joint_constraints_start(batch_id) + mb.first_constraint as usize; + // One column of M⁻¹ per constraint slot . + let dofs_stride = batch_ids.dof_batch_capacity as usize; + let col_base = batch_ids.mb_joint_constraint_columns_start(batch_id) + + (mb.first_constraint as usize) * dofs_stride; + let m = MatSlice::dense(mb_mm_base, ndofs, ndofs); + + // Stage 1: lane-parallel slot reset. + for s in StepRng::new(lane..mb.max_constraints, LANES) { + let mut cz: MultibodyJointConstraint = joint_constraints.read(cons_base + s as usize); + cz.kind = 0; + cz.impulse = 0.0; + joint_constraints.write(cons_base + s as usize, cz); + } + + control_barrier::< + { khal_std::memory::Scope::Workgroup as u32 }, + { khal_std::memory::Scope::QueueFamily as u32 }, + { + khal_std::memory::Semantics::UNIFORM_MEMORY.bits() + | khal_std::memory::Semantics::ACQUIRE_RELEASE.bits() + }, + >(); + + // Stage 2: serial metadata emission on lane 0. + if lane == 0 { + emit_joint_constraints( + links_static, + links_workspace, + joint_constraints, + &mb, + cons_base, + batch_id, + softness.dt, + softness.joint_erp_inv_dt, + softness.joint_cfm_coeff, + batch_ids, + ); + } + + control_barrier::< + { khal_std::memory::Scope::Workgroup as u32 }, + { khal_std::memory::Scope::QueueFamily as u32 }, + { + khal_std::memory::Semantics::UNIFORM_MEMORY.bits() + | khal_std::memory::Semantics::ACQUIRE_RELEASE.bits() + }, + >(); + + // Stage 3: lane-parallel finalize. + for s in StepRng::new(lane..mb.max_constraints, LANES) { + let mut cons = joint_constraints.read(cons_base + s as usize); + if cons.kind == 0 { + continue; + } + let lhs = compute_constraint_column( + joint_constraint_columns, + col_base, + s, + dofs_stride, + ndofs, + cons.dof_id, + mass_matrices, + m, + lu_pivots, + piv_offset, + ); + let cfm_gain = lhs * cons.cfm_coeff + cons.cfm_gain; + cons.cfm_gain = cfm_gain; + cons.inv_lhs = inv(lhs + cfm_gain); + joint_constraints.write(cons_base + s as usize, cons); + } } /// One PGS sweep: iterates the multibody's active limit/motor constraints and From d27f079396a11817399f59f5f05533aa00891787 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Sun, 26 Jul 2026 14:50:55 +0200 Subject: [PATCH 24/39] perf: fuse the colored solver sweeps into per-batch single-workgroup dispatches --- src_rbd/dynamics/solver.rs | 74 +++++++++++++-- src_rbd/pipeline/rbd_step.rs | 19 ++++ src_rbd_shaders/dynamics/solver.rs | 141 ++++++++++++++++++++++++++++- 3 files changed, 224 insertions(+), 10 deletions(-) diff --git a/src_rbd/dynamics/solver.rs b/src_rbd/dynamics/solver.rs index 8e22384..eec6817 100644 --- a/src_rbd/dynamics/solver.rs +++ b/src_rbd/dynamics/solver.rs @@ -13,7 +13,8 @@ use crate::shaders::dynamics::{ GpuApplySolverVelsInc, GpuInitSolverBodies, GpuInitSolverVelsInc, GpuIntegrateLinearized, GpuSolverCleanup, GpuSolverCountConstraints, GpuSolverFinalize, GpuSolverInitConstraints, GpuSolverSortConstraints, - GpuSolverUpdateConstraints, GpuStepGaussSeidel, GpuWarmstart, GpuWarmstartWithoutColors, + GpuSolverUpdateConstraints, GpuStepGaussSeidel, GpuStepGaussSeidelFused, GpuWarmstart, + GpuWarmstartFused, GpuWarmstartWithoutColors, LocalMassProperties, RbdSimParams, TwoBodyConstraint, TwoBodyConstraintBuilder, Velocity, WorldMassProperties, }; @@ -42,6 +43,12 @@ pub struct GpuSolver { warmstart_without_colors: GpuWarmstartWithoutColors, /// Gauss-Seidel iteration step (sequential per color). step_gauss_seidel: GpuStepGaussSeidel, + /// Fused variant of the colored warmstart sweep: one workgroup per batch + /// loops every color internally (barrier between colors). Used when + /// per-batch constraint counts are small. + warmstart_fused: GpuWarmstartFused, + /// Fused variant of the colored Gauss-Seidel sweep (same rationale). + step_gauss_seidel_fused: GpuStepGaussSeidelFused, /// Initializes solver velocity increments. init_solver_vels_inc: GpuInitSolverVelsInc, /// Seeds the COM-centered solver poses from the body world poses @@ -130,6 +137,11 @@ pub struct SolverArgs<'a> { /// correct when `body_group` is the identity (multibody constraints are /// counted on their root's slot with link-id constraint sides). pub colorless_warmstart: bool, + /// Whether the fused colored kernels are used. + /// + /// This is generally used when the number of constraints is small wrt. + /// the number of environments. + pub fused_color_sweeps: bool, /// Shared per-batch capacity / section-offset uniform — see /// [`crate::shaders::utils::BatchIndices`]. Consumed by the (refactored) /// multibody kernels via `MultibodySolverArgs::batch_indices`; the RBD @@ -388,6 +400,20 @@ impl GpuSolver { args.solver_vels, args.batch_indices, )?; + } else if args.fused_color_sweeps { + // One dispatch, one workgroup per batch, colors looped + // internally. `color_uniforms[num_colors]` holds the + // constant `num_colors`. + self.warmstart_fused.call( + pass, + [64, args.num_batches, 1], + args.constraints, + args.solver_vels, + args.color_bucket_starts, + args.color_sorted_ids, + &args.color_uniforms[args.num_colors as usize], + args.batch_indices, + )?; } else { // NOTE: contact colors start at 1 (0 = unassigned). for c in 1..=args.num_colors { @@ -414,19 +440,34 @@ impl GpuSolver { encoder.begin_pass("[RBD] slv/rb-solve-bias", timestamps.as_deref_mut()); let pass = &mut pass; joint_solver.solve(pass, &mut joint_args, args.solver_vels, true)?; - for c in 1..=args.num_colors { - self.step_gauss_seidel.call( + if args.fused_color_sweeps { + self.step_gauss_seidel_fused.call( pass, - args.contacts_len_indirect, + [64, args.num_batches, 1], args.constraints, args.solver_vels, args.color_bucket_starts, args.color_sorted_ids, - &args.color_uniforms[c as usize], + &args.color_uniforms[args.num_colors as usize], args.batch_indices, // use_bias = 1 (the `color_uniform[1]` contains the value 1) &args.color_uniforms[1], )?; + } else { + for c in 1..=args.num_colors { + self.step_gauss_seidel.call( + pass, + args.contacts_len_indirect, + args.constraints, + args.solver_vels, + args.color_bucket_starts, + args.color_sorted_ids, + &args.color_uniforms[c as usize], + args.batch_indices, + // use_bias = 1 (the `color_uniform[1]` contains the value 1) + &args.color_uniforms[1], + )?; + } } } @@ -460,19 +501,34 @@ impl GpuSolver { encoder.begin_pass("[RBD] slv/rb-solve-nobias", timestamps.as_deref_mut()); let pass = &mut pass; joint_solver.solve(pass, &mut joint_args, args.solver_vels, false)?; - for c in 1..=args.num_colors { - self.step_gauss_seidel.call( + if args.fused_color_sweeps { + self.step_gauss_seidel_fused.call( pass, - args.contacts_len_indirect, + [64, args.num_batches, 1], args.constraints, args.solver_vels, args.color_bucket_starts, args.color_sorted_ids, - &args.color_uniforms[c as usize], + &args.color_uniforms[args.num_colors as usize], args.batch_indices, // use_bias = 0 (the `color_uniform[0]` contains the value 0) &args.color_uniforms[0], )?; + } else { + for c in 1..=args.num_colors { + self.step_gauss_seidel.call( + pass, + args.contacts_len_indirect, + args.constraints, + args.solver_vels, + args.color_bucket_starts, + args.color_sorted_ids, + &args.color_uniforms[c as usize], + args.batch_indices, + // use_bias = 0 (the `color_uniform[0]` contains the value 0) + &args.color_uniforms[0], + )?; + } } } } diff --git a/src_rbd/pipeline/rbd_step.rs b/src_rbd/pipeline/rbd_step.rs index 9f0f64e..0f0c8bf 100644 --- a/src_rbd/pipeline/rbd_step.rs +++ b/src_rbd/pipeline/rbd_step.rs @@ -245,6 +245,23 @@ impl RbdPipeline { backend.submit(encoder)?; } + // Colored-sweep strategy: with few constraints per batch, the + // `num_colors` dispatches per sweep (and their empty buckets) dominate + // — run each sweep as one dispatch with one workgroup per batch + // looping the colors internally. The gate is perf-only (the fused + // kernel is correct for any size, just serialized past ~64 lanes): + // use the lagging pair-count readback when auto-resize keeps it fresh, + // else the fixed capacity. + let readback_enabled = state.capacities.solver_colors_resize_policy + != RbdResizePolicy::Fixed + || state.capacities.collisions_resize_policy != RbdResizePolicy::Fixed; + let est_pairs = if readback_enabled { + state.collision_pairs_len_cpu + } else { + state.collision_pairs_per_batch_cpu + }; + let fused_color_sweeps = est_pairs <= 128; + // Phase 2b: solver-prep + warmstart + bounded coloring. Separate // submit from narrow-phase to enable CPU/GPU overlap with the // upcoming Phase 3 solver substep loop. @@ -282,6 +299,7 @@ impl RbdPipeline { body_group: &state.body_group, batch_indices: &state.batch_indices, colorless_warmstart: false, + fused_color_sweeps, }; self.solver.prepare( backend, @@ -422,6 +440,7 @@ impl RbdPipeline { colorless_warmstart: state.multibodies.is_empty(), #[cfg(not(feature = "dim3"))] colorless_warmstart: true, + fused_color_sweeps, }; // Phase 3: Solve constraints diff --git a/src_rbd_shaders/dynamics/solver.rs b/src_rbd_shaders/dynamics/solver.rs index e45e192..cabad45 100644 --- a/src_rbd_shaders/dynamics/solver.rs +++ b/src_rbd_shaders/dynamics/solver.rs @@ -6,7 +6,11 @@ use khal_std::glamx::UVec3; use khal_std::macros::{spirv, spirv_bindgen}; use crate::{AngVector, Pose, Vector}; -use khal_std::{index::MaybeIndexUnchecked, iter::StepRng, sync::atomic_add_u32}; +use khal_std::{ + index::MaybeIndexUnchecked, + iter::StepRng, + sync::{atomic_add_u32, control_barrier}, +}; use super::body::{LocalMassProperties, Velocity, WorldMassProperties}; use super::constraint::{TwoBodyConstraint, TwoBodyConstraintBuilder}; @@ -410,6 +414,141 @@ pub fn gpu_step_gauss_seidel( } } +/// Fused colored warmstart: only one 64-lane workgroup per batch walks every color +/// bucket. +/// +/// Used for small scenes where the contact count is small wrt. the environment count. +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn gpu_warmstart_fused( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] constraints: &[TwoBodyConstraint], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] solver_vels: &mut [Velocity], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] color_starts: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] color_sorted_ids: &[u32], + #[spirv(uniform, descriptor_set = 0, binding = 4)] num_colors: &u32, + #[spirv(uniform, descriptor_set = 0, binding = 5)] batch_ids: &BatchIndices, +) { + let lane = invocation_id.x; + let batch_id = invocation_id.y; + let stride = batch_ids.solver_color_buckets_stride; + + let constraints = batch_ids.contact_batch(batch_id, constraints); + let color_sorted_ids = batch_ids.contact_batch(batch_id, color_sorted_ids); + let mut solver_vels = batch_ids.coll_batch_mut(batch_id, solver_vels); + let num_colors = *num_colors; + + let base = (batch_id * stride) as usize; + if color_starts.read(base + 1) == color_starts.read(base + num_colors as usize + 1) { + // Every color bucket is empty. + return; + } + + for color in 1..=num_colors { + let bucket = base + color as usize; + let start = color_starts.read(bucket); + let end = color_starts.read(bucket + 1); + if start == end { + // Empty color. + continue; + } + + for k in StepRng::new(start + lane..end, WORKGROUP_SIZE) { + let i = color_sorted_ids[k as usize]; + let constraint = &constraints[i as usize]; + let solver_id1 = constraint.solver_body_a as usize; + let solver_id2 = constraint.solver_body_b as usize; + + let mut solver_vel1 = solver_vels[solver_id1]; + let mut solver_vel2 = solver_vels[solver_id2]; + + constraint.warmstart_constraint(&mut solver_vel1, &mut solver_vel2); + + solver_vels[solver_id1] = solver_vel1; + solver_vels[solver_id2] = solver_vel2; + } + + control_barrier::< + { khal_std::memory::Scope::Workgroup as u32 }, + { khal_std::memory::Scope::QueueFamily as u32 }, + { + khal_std::memory::Semantics::UNIFORM_MEMORY.bits() + | khal_std::memory::Semantics::ACQUIRE_RELEASE.bits() + }, + >(); + } +} + +/// Fused colored Gauss-Seidel sweep: only one 64-lane workgroup per batch walks +/// every color bucket with a storage barrier between colors. +/// +/// Used for small scenes where the contact count is small wrt. the environment count. +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn gpu_step_gauss_seidel_fused( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + constraints: &mut [TwoBodyConstraint], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] solver_vels: &mut [Velocity], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] color_starts: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] color_sorted_ids: &[u32], + #[spirv(uniform, descriptor_set = 0, binding = 4)] num_colors: &u32, + #[spirv(uniform, descriptor_set = 0, binding = 5)] batch_ids: &BatchIndices, + #[spirv(uniform, descriptor_set = 0, binding = 6)] use_bias: &u32, +) { + let lane = invocation_id.x; + let batch_id = invocation_id.y; + let stride = batch_ids.solver_color_buckets_stride; + + let mut constraints = batch_ids.contact_batch_mut(batch_id, constraints); + let color_sorted_ids = batch_ids.contact_batch(batch_id, color_sorted_ids); + let mut solver_vels = batch_ids.coll_batch_mut(batch_id, solver_vels); + let num_colors = *num_colors; + let use_bias = *use_bias != 0; + + // Early-out / empty-color skip: see `gpu_warmstart_fused`. + let base = (batch_id * stride) as usize; + if color_starts.read(base + 1) == color_starts.read(base + num_colors as usize + 1) { + return; + } + + for color in 1..=num_colors { + let bucket = base + color as usize; + let start = color_starts.read(bucket); + let end = color_starts.read(bucket + 1); + if start == end { + continue; + } + + for k in StepRng::new(start + lane..end, WORKGROUP_SIZE) { + let i = color_sorted_ids[k as usize]; + let solver_id1 = constraints[i as usize].solver_body_a as usize; + let solver_id2 = constraints[i as usize].solver_body_b as usize; + + let mut solver_vel1 = solver_vels[solver_id1]; + let mut solver_vel2 = solver_vels[solver_id2]; + + constraints[i as usize].solve_constraint_gauss_seidel( + &mut solver_vel1, + &mut solver_vel2, + use_bias, + ); + + solver_vels[solver_id1] = solver_vel1; + solver_vels[solver_id2] = solver_vel2; + } + + control_barrier::< + { khal_std::memory::Scope::Workgroup as u32 }, + { khal_std::memory::Scope::QueueFamily as u32 }, + { + khal_std::memory::Semantics::UNIFORM_MEMORY.bits() + | khal_std::memory::Semantics::ACQUIRE_RELEASE.bits() + }, + >(); + } +} + /// Integrates velocity to update poses. #[spirv_bindgen] #[spirv(compute(threads(64)))] From 37196d5aa79d35e8b12868b79648fcb785a8a2b4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Sun, 26 Jul 2026 15:04:10 +0200 Subject: [PATCH 25/39] perf: fused shared-memory multibody PGS sweep + parallel warmstart/reset --- .../dynamics/multibody/multibody_solver.rs | 94 +++----- .../dynamics/multibody/contact_constraints.rs | 199 +++------------- .../dynamics/multibody/joint_constraints.rs | 142 +----------- src_rbd_shaders/dynamics/multibody/mod.rs | 2 + .../dynamics/multibody/solve_constraints.rs | 217 ++++++++++++++++++ 5 files changed, 295 insertions(+), 359 deletions(-) create mode 100644 src_rbd_shaders/dynamics/multibody/solve_constraints.rs diff --git a/src_rbd/dynamics/multibody/multibody_solver.rs b/src_rbd/dynamics/multibody/multibody_solver.rs index 0266d2d..6425db5 100644 --- a/src_rbd/dynamics/multibody/multibody_solver.rs +++ b/src_rbd/dynamics/multibody/multibody_solver.rs @@ -9,10 +9,10 @@ use crate::shaders::dynamics::{ GpuMbFinalizeContactConstraints, GpuMbGravityAndLu, GpuMbGravityAndLuT8, GpuMbGravityAndLuT16, GpuMbGravityAndLuT32, GpuMbInitContactConstraints, GpuMbInitJointConstraints, GpuMbIntegrate, GpuMbIntegrateVelocities, - GpuMbRemoveContactConstraintBias, GpuMbRemoveImpulseJointConstraintBias, + GpuMbRemoveImpulseJointConstraintBias, GpuMbResetContactWarmstart, GpuMbStashContactsLen, GpuMbWarmstartContactConstraints, - GpuMbRemoveSolveJointNoBias, GpuMbSolveContactConstraints, GpuMbSolveImpulseJointConstraints, - GpuMbFinalizeImpulseJointConstraints, GpuMbSolveJointConstraints, + GpuMbSolveConstraints, GpuMbSolveImpulseJointConstraints, + GpuMbFinalizeImpulseJointConstraints, GpuMbUpdateImpulseJointConstraints, Velocity, WorldMassProperties, }; use crate::shaders::utils::BatchIndices; @@ -33,13 +33,12 @@ pub struct GpuMultibodySolver { gravity_and_lu_t32: GpuMbGravityAndLuT32, compute_dynamics_pre: GpuMbComputeDynamicsPre, compute_dynamics_without_coriolis_pre: GpuMbComputeDynamicsWithoutCoriolisPre, - solve_joint_with_bias: GpuMbSolveJointConstraints, init_joint_with_bias: GpuMbInitJointConstraints, - /// Fused remove-bias + solve-without-bias for the stabilization sweep. - remove_solve_joint_no_bias: GpuMbRemoveSolveJointNoBias, init_contact_constraints: GpuMbInitContactConstraints, finalize_contact_constraints: GpuMbFinalizeContactConstraints, - solve_contact_constraints: GpuMbSolveContactConstraints, + /// Fused joint+contact PGS sweep (one workgroup per multibody, shared- + /// memory dof velocities). + solve_constraints: GpuMbSolveConstraints, /// Zero the accumulated contact impulses once per frame (warmstart reset). reset_contact_warmstart: GpuMbResetContactWarmstart, /// Copy `contacts_len[batch]` into each `MultibodyInfo` once per step so @@ -48,7 +47,6 @@ pub struct GpuMultibodySolver { stash_contacts_len: GpuMbStashContactsLen, /// Re-apply the accumulated contact impulse each substep (warmstart). warmstart_contact_constraints: GpuMbWarmstartContactConstraints, - remove_contact_constraint_bias: GpuMbRemoveContactConstraintBias, update_impulse_joint_constraints: GpuMbUpdateImpulseJointConstraints, /// Finalize pass for the impulse-joint build (LU back-solve + `inv_lhs`), /// split out so the build pass fits 8 storage buffers. @@ -117,9 +115,10 @@ impl GpuMultibodySolver { } // Zero the accumulated contact impulses so the first substep's warmstart // starts cold (within a frame they are then preserved across substeps). + // One 64-lane workgroup per multibody (lanes stride the slots). self.reset_contact_warmstart.call( pass, - mb.flat_mb_dispatch(), + [mb.multibodies_per_batch * MB_LU_LANES, mb.num_batches, 1], &mb.multibody_info, &mut mb.contact_constraints, args.batch_indices, @@ -261,12 +260,15 @@ impl GpuMultibodySolver { // substep — mirrors rapier's per-substep `contact_constraints.warmstart` // and matches what the rigid-body solver does for free contacts. On the // first substep the impulse was just reset to 0, so this is a no-op. + // One 64-lane workgroup per multibody (one DOF per lane). { let mut pass = encoder.begin_pass("[RBD] mbb/warmstart-contact", timestamps.as_deref_mut()); + let warmstart_dispatch = + [mb.multibodies_per_batch * MB_LU_LANES, mb.num_batches, 1]; self.warmstart_contact_constraints.call( &mut pass, - dispatch, + warmstart_dispatch, &mb.multibody_info, &mb.contact_constraints, &mb.contact_constraint_columns, @@ -290,30 +292,24 @@ impl GpuMultibodySolver { if mb.is_empty() { return Ok(()); } - let dispatch = mb.flat_mb_dispatch(); - - if mb.has_joint_constraints { - self.solve_joint_with_bias.call( - pass, - dispatch, - &mb.multibody_info, - &mut mb.joint_constraints, - &mut mb.joint_constraint_columns, - &mut mb.dof_state, - args.batch_indices, - )?; - } - self.solve_contact_constraints.call( + // Fused joint+contact sweep: one 64-lane workgroup per multibody with + // the generalized velocities held in workgroup memory + // (`color_uniforms[1]` holds the constant 1 = use_bias). + let solve_dispatch = [mb.multibodies_per_batch * MB_LU_LANES, mb.num_batches, 1]; + self.solve_constraints.call( pass, - dispatch, + solve_dispatch, &mb.multibody_info, + &mut mb.joint_constraints, + &mb.joint_constraint_columns, &mut mb.contact_constraints, &mb.contact_constraint_jacs, &mb.contact_constraint_columns, + &args.color_uniforms[1], + args.batch_indices, &mut mb.dof_state, args.solver_vels, - args.batch_indices, )?; // Multibody-touching impulse joints — generic (rb-mb / mb-mb) @@ -432,25 +428,26 @@ impl GpuMultibodySolver { if mb.is_empty() { return Ok(()); } - let dispatch = mb.flat_mb_dispatch(); - if mb.has_joint_constraints { - self.remove_solve_joint_no_bias.call( - pass, - dispatch, - &mb.multibody_info, - &mut mb.joint_constraints, - &mb.joint_constraint_columns, - &mut mb.dof_state, - args.batch_indices, - )?; - } - self.remove_contact_constraint_bias.call( + // Fused joint+contact stabilization sweep: `use_bias = 0` + // (`color_uniforms[0]`) makes the kernel read `rhs_wo_bias` directly, + // which replaces the former remove-bias read-modify-write dispatches + // (every constraint is re-initialized next substep, so the persistent + // `rhs` rewrite was never needed). + let solve_dispatch = [mb.multibodies_per_batch * MB_LU_LANES, mb.num_batches, 1]; + self.solve_constraints.call( pass, - dispatch, - &mut mb.contact_constraints, + solve_dispatch, &mb.multibody_info, + &mut mb.joint_constraints, + &mb.joint_constraint_columns, + &mut mb.contact_constraints, + &mb.contact_constraint_jacs, + &mb.contact_constraint_columns, + &args.color_uniforms[0], args.batch_indices, + &mut mb.dof_state, + args.solver_vels, )?; if mb.mb_imp_joints_per_batch > 0 { let imp_dispatch = [mb.mb_imp_joints_per_batch, mb.num_batches, 1]; @@ -462,21 +459,6 @@ impl GpuMultibodySolver { &mb.mb_imp_joint_count, args.batch_indices, )?; - } - - // (joint sweep WITHOUT bias was fused into `remove_solve_joint_no_bias` above.) - self.solve_contact_constraints.call( - pass, - dispatch, - &mb.multibody_info, - &mut mb.contact_constraints, - &mb.contact_constraint_jacs, - &mb.contact_constraint_columns, - &mut mb.dof_state, - args.solver_vels, - args.batch_indices, - )?; - if mb.mb_imp_joints_per_batch > 0 { // Final stabilization sweep WITHOUT bias — colored, one // dispatch per color (see the with-bias sweep above). for c in 0..mb.mb_imp_joint_num_colors as usize { diff --git a/src_rbd_shaders/dynamics/multibody/contact_constraints.rs b/src_rbd_shaders/dynamics/multibody/contact_constraints.rs index 00a35f4..eac48dc 100644 --- a/src_rbd_shaders/dynamics/multibody/contact_constraints.rs +++ b/src_rbd_shaders/dynamics/multibody/contact_constraints.rs @@ -575,22 +575,25 @@ pub fn gpu_mb_stash_contacts_len( /// substep loop) so the first substep's warmstart starts cold; within a frame /// `gpu_mb_init_contact_constraints` preserves the impulse across substeps and /// `gpu_mb_warmstart_contact_constraints` re-applies it each substep. +/// One 64-lane workgroup per (multibody, batch). #[spirv_bindgen] #[spirv(compute(threads(64)))] pub fn gpu_mb_reset_contact_warmstart( - #[spirv(global_invocation_id)] invocation_id: UVec3, + #[spirv(workgroup_id)] workgroup_id: UVec3, + #[spirv(local_invocation_id)] local_id: UVec3, #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] multibody_info: &[MultibodyInfo], #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] contact_constraints: &mut [MultibodyContactConstraint], #[spirv(uniform, descriptor_set = 0, binding = 2)] batch_ids: &BatchIndices, ) { - // Flattened (multibody, batch) grid — see `BatchIndices::num_batches`. + const LANES: u32 = 64; + let batch_id = workgroup_id.y; + let mb_idx = workgroup_id.x; + let lane = local_id.x; let num_mb = batch_ids.multibodies_len; - if invocation_id.x >= num_mb * batch_ids.num_batches { + if mb_idx >= num_mb { return; } - let batch_id = invocation_id.x / num_mb; - let mb_idx = invocation_id.x % num_mb; let mb_start = batch_ids.mb_start(batch_id); let cons_start = batch_ids.mb_contact_constraints_start(batch_id); let mb = multibody_info.read(mb_start + mb_idx as usize); @@ -600,7 +603,7 @@ pub fn gpu_mb_reset_contact_warmstart( let cons_base = cons_start + (mb_idx as usize) * (MAX_MB_CONTACT_CONSTRAINTS_PER_MB as usize); // Zero to capacity (the per-frame contact count isn't known here, and last // frame's count may be smaller than this frame's). - for s in 0..MAX_MB_CONTACT_CONSTRAINTS_PER_MB { + for s in StepRng::new(lane..MAX_MB_CONTACT_CONSTRAINTS_PER_MB, LANES) { let mut cons = contact_constraints.read(cons_base + s as usize); cons.impulse = 0.0; contact_constraints.write(cons_base + s as usize, cons); @@ -611,10 +614,13 @@ pub fn gpu_mb_reset_contact_warmstart( /// to the multibody generalized velocities (`dof_state`) and the free-body /// solver velocities. Applies the FULL accumulated impulse (no `rhs` term, no /// clamping). +/// +/// One 64-lane workgroup per (multibody, batch). #[spirv_bindgen] #[spirv(compute(threads(64)))] pub fn gpu_mb_warmstart_contact_constraints( - #[spirv(global_invocation_id)] invocation_id: UVec3, + #[spirv(workgroup_id)] workgroup_id: UVec3, + #[spirv(local_invocation_id)] local_id: UVec3, #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] multibody_info: &[MultibodyInfo], #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] contact_constraints: &[MultibodyContactConstraint], @@ -623,13 +629,13 @@ pub fn gpu_mb_warmstart_contact_constraints( #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] solver_vels: &mut [Velocity], #[spirv(uniform, descriptor_set = 0, binding = 5)] batch_ids: &BatchIndices, ) { - // Flattened (multibody, batch) grid — see `BatchIndices::num_batches`. + let batch_id = workgroup_id.y; + let mb_idx = workgroup_id.x; + let lane = local_id.x; let num_mb = batch_ids.multibodies_len; - if invocation_id.x >= num_mb * batch_ids.num_batches { + if mb_idx >= num_mb { return; } - let batch_id = invocation_id.x / num_mb; - let mb_idx = invocation_id.x % num_mb; let mb_start = batch_ids.mb_start(batch_id); let cons_start = batch_ids.mb_contact_constraints_start(batch_id); @@ -648,21 +654,31 @@ pub fn gpu_mb_warmstart_contact_constraints( col_start + (mb_idx as usize) * (MAX_MB_CONTACT_CONSTRAINTS_PER_MB as usize) * dofs_stride; let count = mb.contact_constraint_count; + // No accumulated impulses to re-apply: skip the dof round-trip. + if count == 0 { + return; + } + + // This lane's DOF velocity, accumulated in a register across every + // constraint. + let mut v_lane = if lane < ndofs { + dof_state.read(v_base + lane as usize) + } else { + 0.0 + }; for s in 0..count { let cons = contact_constraints.read(cons_base + s as usize); let imp = cons.impulse; if imp != 0.0 { let col_offset = col_base + (s as usize) * dofs_stride; // Multibody side: v += impulse · column (column = M⁻¹ Jᵀ). - for i in 0..ndofs { - let v_idx = v_base + i as usize; - let cur = dof_state.read(v_idx); - let col = contact_constraint_columns.read(col_offset + i as usize); - dof_state.write(v_idx, cur + imp * col); + if lane < ndofs { + let col = contact_constraint_columns.read(col_offset + lane as usize); + v_lane += imp * col; } // Free body side (skipped for self-contacts). let is_self = cons.free_body_id == u32::MAX; - if !is_self { + if lane == 0 && !is_self { let free = solver_vels.read(colliders_start + cons.free_body_id as usize); let mut new_free = free; new_free.linear += cons.lin_jac * (cons.free_body_im * imp); @@ -671,6 +687,10 @@ pub fn gpu_mb_warmstart_contact_constraints( } } } + + if lane < ndofs { + dof_state.write(v_base + lane as usize, v_lane); + } } /// Pass 2: for each emitted constraint, LU back-solve `M · column = Jᵀ` @@ -758,148 +778,3 @@ pub fn gpu_mb_finalize_contact_constraints( contact_constraints.write(cons_base + s as usize, cons); } } - -/// One PGS sweep over the multibody's active contact constraints. Updates -/// the multibody's `dof_velocities` and the free body's `solver_vels`. -#[spirv_bindgen] -#[spirv(compute(threads(64)))] -pub fn gpu_mb_solve_contact_constraints( - #[spirv(global_invocation_id)] invocation_id: UVec3, - #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] multibody_info: &[MultibodyInfo], - #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] - contact_constraints: &mut [MultibodyContactConstraint], - #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] contact_constraint_jacs: &[f32], - #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] contact_constraint_columns: &[f32], - #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] dof_state: &mut [f32], - #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] solver_vels: &mut [Velocity], - #[spirv(uniform, descriptor_set = 0, binding = 6)] batch_ids: &BatchIndices, -) { - // Flattened (multibody, batch) grid — see `BatchIndices::num_batches`. - let num_mb = batch_ids.multibodies_len; - if invocation_id.x >= num_mb * batch_ids.num_batches { - return; - } - let batch_id = invocation_id.x / num_mb; - let mb_idx = invocation_id.x % num_mb; - - let mb_start = batch_ids.mb_start(batch_id); - let cons_start = batch_ids.mb_contact_constraints_start(batch_id); - let col_start = batch_ids.mb_contact_constraint_columns_start(batch_id); - let colliders_start = batch_ids.coll_start(batch_id); - - let mb = multibody_info.read(mb_start + mb_idx as usize); - let ndofs = mb.ndofs; - if ndofs == 0 { - return; - } - let v_base = batch_ids.dof_start(batch_id) + mb.first_dof as usize; - let cons_base = cons_start + (mb_idx as usize) * (MAX_MB_CONTACT_CONSTRAINTS_PER_MB as usize); - let dofs_stride = batch_ids.dof_batch_capacity as usize; - let col_base = - col_start + (mb_idx as usize) * (MAX_MB_CONTACT_CONSTRAINTS_PER_MB as usize) * dofs_stride; - - let count = mb.contact_constraint_count; - for s in 0..count { - let mut cons = contact_constraints.read(cons_base + s as usize); - let col_offset = col_base + (s as usize) * dofs_stride; - - // J · u = J_mb · v_mb_dofs + J_free · v_free. - let is_self = cons.free_body_id == u32::MAX; - let mut j_dot_v = 0.0f32; - for i in 0..ndofs { - let j = contact_constraint_jacs.read(col_offset + i as usize); - let v = dof_state.read(v_base + i as usize); - j_dot_v += j * v; - } - let free = if is_self { - Velocity::default() - } else { - solver_vels.read(colliders_start + cons.free_body_id as usize) - }; - if !is_self { - j_dot_v += cons.lin_jac.dot(free.linear) + gdot(cons.ang_jac, free.angular); - } - - let rhs_total = j_dot_v + cons.rhs; - // CFM-factor form (rapier's `*ContactConstraintNormalPart::generic_solve`): - // `new = cfm_factor · (impulse − r · Δvel)`. `cfm_factor < 1` provides the - // soft-constraint compliance that keeps resting contacts from jittering. - let raw_imp = cons.cfm_factor * (cons.impulse - cons.inv_lhs * rhs_total); - - // Normal: clamp to ≥ 0 (no separation impulse). Friction tangent: - // clamp to `±μ · normal_impulse` — looks up the paired normal slot - // for the current accumulated impulse. Mirrors rapier's - // `ContactConstraintNormalPart::generic_solve` / - // `ContactConstraintTangentPart::generic_solve` (independent - // per-tangent clamp, i.e. box friction; rapier's circular-cone - // joint clamp is a future refinement). - let new_imp = if cons.kind == MB_CONTACT_KIND_TANGENT { - let normal = contact_constraints.read(cons_base + cons.normal_constraint_slot as usize); - let limit = cons.friction_coeff * normal.impulse; - if raw_imp > limit { - limit - } else if raw_imp < -limit { - -limit - } else { - raw_imp - } - } else if raw_imp < 0.0 { - 0.0 - } else { - raw_imp - }; - let delta = new_imp - cons.impulse; - cons.impulse = new_imp; - contact_constraints.write(cons_base + s as usize, cons); - - if delta != 0.0 { - for i in 0..ndofs { - let v_idx = v_base + i as usize; - let cur = dof_state.read(v_idx); - let col = contact_constraint_columns.read(col_offset + i as usize); - dof_state.write(v_idx, cur + delta * col); - } - if !is_self { - let mut new_free = free; - new_free.linear += cons.lin_jac * (cons.free_body_im * delta); - new_free.angular += cons.ii_ang_jac * delta; - solver_vels.write(colliders_start + cons.free_body_id as usize, new_free); - } - } - } -} - -/// Strip the positional bias from each active contact constraint's `rhs`, -/// matching `gpu_mb_remove_joint_constraint_bias`. -#[spirv_bindgen] -#[spirv(compute(threads(64)))] -pub fn gpu_mb_remove_contact_constraint_bias( - #[spirv(global_invocation_id)] invocation_id: UVec3, - #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] - contact_constraints: &mut [MultibodyContactConstraint], - #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] multibody_info: &[MultibodyInfo], - #[spirv(uniform, descriptor_set = 0, binding = 2)] batch_ids: &BatchIndices, -) { - // Flattened (multibody, batch) grid — see `BatchIndices::num_batches`. - let num_mb = batch_ids.multibodies_len; - if invocation_id.x >= num_mb * batch_ids.num_batches { - return; - } - let batch_id = invocation_id.x / num_mb; - let mb_idx = invocation_id.x % num_mb; - - let mb_start = batch_ids.mb_start(batch_id); - let cons_start = batch_ids.mb_contact_constraints_start(batch_id); - let cons_base = cons_start + (mb_idx as usize) * (MAX_MB_CONTACT_CONSTRAINTS_PER_MB as usize); - let mb = multibody_info.read(mb_start + mb_idx as usize); - let count = mb.contact_constraint_count; - - for s in 0..count { - let mut cons = contact_constraints.read(cons_base + s as usize); - if cons.kind == 0 { - continue; - } - cons.rhs = cons.rhs_wo_bias; - contact_constraints.write(cons_base + s as usize, cons); - } -} diff --git a/src_rbd_shaders/dynamics/multibody/joint_constraints.rs b/src_rbd_shaders/dynamics/multibody/joint_constraints.rs index a248502..4a69c94 100644 --- a/src_rbd_shaders/dynamics/multibody/joint_constraints.rs +++ b/src_rbd_shaders/dynamics/multibody/joint_constraints.rs @@ -55,66 +55,7 @@ fn lu_solve_unit( lu_solve_in_place(buf_m, m, buf_pivots, pivots_offset, dst, dst_offset); } -/// PGS sweep body — shared between `gpu_mb_solve_joint_constraints` and -/// the fused init+solve / remove-bias+solve kernels. Writes back `cons` -/// before subtracting `delta · column` from `v`. -#[inline] -fn solve_joint_constraints_body( - multibody_info: &[MultibodyInfo], - joint_constraints: &mut [MultibodyJointConstraint], - joint_constraint_columns: &[f32], - dof_state: &mut [f32], - batch_id: u32, - mb_idx: u32, - batch_ids: &BatchIndices, -) { - let mb = batch_ids - .mb_batch(batch_id, multibody_info) - .read(mb_idx as usize); - let ndofs = mb.ndofs; - if ndofs == 0 || mb.max_constraints == 0 { - return; - } - let v_base = batch_ids.dof_start(batch_id) + mb.first_dof as usize; - let cons_base = batch_ids.mb_joint_constraints_start(batch_id) + mb.first_constraint as usize; - let dofs_stride = batch_ids.dof_batch_capacity as usize; - let col_base = batch_ids.mb_joint_constraint_columns_start(batch_id) - + (mb.first_constraint as usize) * dofs_stride; - - for s in 0..mb.max_constraints { - let mut cons = joint_constraints.read(cons_base + s as usize); - if cons.kind == 0 { - continue; - } - - let v_d = dof_state.read(v_base + cons.dof_id as usize); - let rhs_total = v_d + cons.rhs; - let raw_imp = cons.impulse + cons.inv_lhs * (rhs_total - cons.cfm_gain * cons.impulse); - let mut new_imp = raw_imp; - if new_imp < cons.impulse_lo { - new_imp = cons.impulse_lo; - } - if new_imp > cons.impulse_hi { - new_imp = cons.impulse_hi; - } - let delta = new_imp - cons.impulse; - cons.impulse = new_imp; - joint_constraints.write(cons_base + s as usize, cons); - - for i in 0..ndofs { - let v_idx = v_base + i as usize; - let cur = dof_state.read(v_idx); - let col = - joint_constraint_columns.read(col_base + (s as usize) * dofs_stride + i as usize); - dof_state.write(v_idx, cur - delta * col); - } - } -} - -/// Serial (lane-0) emission walk: writes the metadata of every active -/// limit/motor constraint slot. The expensive M⁻¹-column back-solves happen -/// afterwards, lane-parallel, in `gpu_mb_init_joint_constraints`' finalize -/// stage. Slot zeroing also happens there (lane-parallel, before this walk). +/// Serially writes the metadata of every active limit/motor constraint slot. #[inline] fn emit_joint_constraints( links_static: &[MultibodyLinkStatic], @@ -522,84 +463,3 @@ pub fn gpu_mb_init_joint_constraints( joint_constraints.write(cons_base + s as usize, cons); } } - -/// One PGS sweep: iterates the multibody's active limit/motor constraints and -/// updates `dof_velocities` in place. Mirrors rapier's `JointConstraint::solve_generic` -/// for a 1-DOF jacobian. -#[spirv_bindgen] -#[spirv(compute(threads(64)))] -pub fn gpu_mb_solve_joint_constraints( - #[spirv(global_invocation_id)] invocation_id: UVec3, - #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] multibody_info: &[MultibodyInfo], - #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] - joint_constraints: &mut [MultibodyJointConstraint], - #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] - joint_constraint_columns: &mut [f32], - #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] dof_state: &mut [f32], - #[spirv(uniform, descriptor_set = 0, binding = 4)] batch_ids: &BatchIndices, -) { - // Flattened (multibody, batch) grid — see `BatchIndices::num_batches`. - let num_mb = batch_ids.multibodies_len; - if invocation_id.x >= num_mb * batch_ids.num_batches { - return; - } - let batch_id = invocation_id.x / num_mb; - let mb_idx = invocation_id.x % num_mb; - solve_joint_constraints_body( - multibody_info, - joint_constraints, - joint_constraint_columns, - dof_state, - batch_id, - mb_idx, - batch_ids, - ); -} - -/// Fused `remove_bias + solve_without_bias` for joint constraints — runs once -/// per substep at the end of the substep, after position integration. Drops -/// one per-multibody dispatch per substep. -#[spirv_bindgen] -#[spirv(compute(threads(64)))] -pub fn gpu_mb_remove_solve_joint_no_bias( - #[spirv(global_invocation_id)] invocation_id: UVec3, - #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] multibody_info: &[MultibodyInfo], - #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] - joint_constraints: &mut [MultibodyJointConstraint], - #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] joint_constraint_columns: &[f32], - #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] dof_state: &mut [f32], - #[spirv(uniform, descriptor_set = 0, binding = 4)] batch_ids: &BatchIndices, -) { - // Flattened (multibody, batch) grid — see `BatchIndices::num_batches`. - let num_mb = batch_ids.multibodies_len; - if invocation_id.x >= num_mb * batch_ids.num_batches { - return; - } - let batch_id = invocation_id.x / num_mb; - let mb_idx = invocation_id.x % num_mb; - - let mb = batch_ids - .mb_batch(batch_id, multibody_info) - .read(mb_idx as usize); - let cons_base = batch_ids.mb_joint_constraints_start(batch_id) + mb.first_constraint as usize; - - // Inlined `remove_bias`: replace `rhs` with `rhs_wo_bias` for active slots. - for s in 0..mb.max_constraints { - let mut cons = joint_constraints.read(cons_base + s as usize); - if cons.kind == 0 { - continue; - } - cons.rhs = cons.rhs_wo_bias; - joint_constraints.write(cons_base + s as usize, cons); - } - - solve_joint_constraints_body( - multibody_info, - joint_constraints, - joint_constraint_columns, - dof_state, - batch_id, - mb_idx, - batch_ids, - ); -} diff --git a/src_rbd_shaders/dynamics/multibody/mod.rs b/src_rbd_shaders/dynamics/multibody/mod.rs index 8b64d1a..5cde22b 100644 --- a/src_rbd_shaders/dynamics/multibody/mod.rs +++ b/src_rbd_shaders/dynamics/multibody/mod.rs @@ -22,6 +22,7 @@ mod jacobian; mod joint_constraints; mod lu; mod mass_matrix; +mod solve_constraints; mod types; mod utils; @@ -31,5 +32,6 @@ pub use gravity_and_lu::*; pub use impulse_joint_constraints::*; pub use integrate::*; pub use joint_constraints::*; +pub use solve_constraints::*; pub use types::*; pub use utils::*; diff --git a/src_rbd_shaders/dynamics/multibody/solve_constraints.rs b/src_rbd_shaders/dynamics/multibody/solve_constraints.rs new file mode 100644 index 0000000..ecdbd36 --- /dev/null +++ b/src_rbd_shaders/dynamics/multibody/solve_constraints.rs @@ -0,0 +1,217 @@ +//! Fused multibody PGS sweep: joint limit/motor constraints followed by +//! contact constraints, in one dispatch per substep phase. + +use khal_std::glamx::UVec3; +use khal_std::index::MaybeIndexUnchecked; +use khal_std::iter::StepRng; +use khal_std::macros::{spirv, spirv_bindgen}; +use khal_std::sync::workgroup_memory_barrier_with_group_sync; + +use crate::dynamics::body::Velocity; +use crate::gdot; +use crate::utils::BatchIndices; +use crate::utils::linalg::MAX_MB_DOFS; + +use super::types::{ + MAX_MB_CONTACT_CONSTRAINTS_PER_MB, MB_CONTACT_KIND_TANGENT, MultibodyContactConstraint, + MultibodyInfo, MultibodyJointConstraint, +}; + +const LANES: u32 = 64; + +/// One PGS sweep over a multibody's joint (limit/motor) constraints followed +/// by its contact constraints. +/// +/// Dispatch: one 64-lane workgroup per (multibody, batch). +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn gpu_mb_solve_constraints( + #[spirv(workgroup_id)] workgroup_id: UVec3, + #[spirv(local_invocation_id)] local_id: UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] multibody_info: &[MultibodyInfo], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] + joint_constraints: &mut [MultibodyJointConstraint], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] joint_constraint_columns: &[f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] + contact_constraints: &mut [MultibodyContactConstraint], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] contact_constraint_jacs: &[f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] contact_constraint_columns: &[f32], + #[spirv(uniform, descriptor_set = 0, binding = 6)] use_bias: &u32, + #[spirv(uniform, descriptor_set = 0, binding = 7)] batch_ids: &BatchIndices, + #[spirv(storage_buffer, descriptor_set = 1, binding = 0)] dof_state: &mut [f32], + #[spirv(storage_buffer, descriptor_set = 1, binding = 1)] solver_vels: &mut [Velocity], + #[spirv(workgroup)] dof_v: &mut [f32; MAX_MB_DOFS as usize], + #[spirv(workgroup)] scratch: &mut [f32; LANES as usize], + #[spirv(workgroup)] imp_shared: &mut [f32; MAX_MB_CONTACT_CONSTRAINTS_PER_MB as usize], +) { + let batch_id = workgroup_id.y; + let mb_idx = workgroup_id.x; + let lane = local_id.x; + let num_mb = batch_ids.multibodies_len; + if mb_idx >= num_mb { + return; + } + + let mb_start = batch_ids.mb_start(batch_id); + let mb = multibody_info.read(mb_start + mb_idx as usize); + let ndofs = mb.ndofs; + // Uniform per workgroup: every lane of this group returns together. + if ndofs == 0 { + return; + } + let use_bias = *use_bias != 0; + + let v_base = batch_ids.dof_start(batch_id) + mb.first_dof as usize; + let dofs_stride = batch_ids.dof_batch_capacity as usize; + let colliders_start = batch_ids.coll_start(batch_id); + + let jcons_base = + batch_ids.mb_joint_constraints_start(batch_id) + mb.first_constraint as usize; + let jcol_base = batch_ids.mb_joint_constraint_columns_start(batch_id) + + (mb.first_constraint as usize) * dofs_stride; + + let ccons_base = batch_ids.mb_contact_constraints_start(batch_id) + + (mb_idx as usize) * (MAX_MB_CONTACT_CONSTRAINTS_PER_MB as usize); + let ccol_base = batch_ids.mb_contact_constraint_columns_start(batch_id) + + (mb_idx as usize) * (MAX_MB_CONTACT_CONSTRAINTS_PER_MB as usize) * dofs_stride; + + let contact_count = mb.contact_constraint_count; + if mb.max_constraints == 0 && contact_count == 0 { + // Nothing to solve. + return; + } + + // Load the generalized velocities and accumulated contact impulses into + // workgroup memory. + if lane < ndofs { + dof_v[lane as usize] = dof_state.read(v_base + lane as usize); + } + for s in StepRng::new(lane..contact_count, LANES) { + imp_shared[s as usize] = contact_constraints.read(ccons_base + s as usize).impulse; + } + workgroup_memory_barrier_with_group_sync(); + + + // Joint limits/motors + for s in 0..mb.max_constraints { + let cons = joint_constraints.read(jcons_base + s as usize); + if cons.kind == 0 { + // Uniform skip: all lanes take it together (barrier-safe). + continue; + } + + let rhs = if use_bias { cons.rhs } else { cons.rhs_wo_bias }; + let v_d = dof_v[cons.dof_id as usize]; + let rhs_total = v_d + rhs; + let raw_imp = cons.impulse + cons.inv_lhs * (rhs_total - cons.cfm_gain * cons.impulse); + let mut new_imp = raw_imp; + if new_imp < cons.impulse_lo { + new_imp = cons.impulse_lo; + } + if new_imp > cons.impulse_hi { + new_imp = cons.impulse_hi; + } + let delta = new_imp - cons.impulse; + + if lane == 0 { + let mut cons = cons; + cons.impulse = new_imp; + joint_constraints.write(jcons_base + s as usize, cons); + } + + // All lanes read `dof_v[dof_id]` above; sync before overwriting it. + workgroup_memory_barrier_with_group_sync(); + if lane < ndofs { + let col = joint_constraint_columns + .read(jcol_base + (s as usize) * dofs_stride + lane as usize); + dof_v[lane as usize] -= delta * col; + } + workgroup_memory_barrier_with_group_sync(); + } + + + // Contacts. + for s in 0..contact_count { + let cons = contact_constraints.read(ccons_base + s as usize); + let col_offset = ccol_base + (s as usize) * dofs_stride; + let is_self = cons.free_body_id == u32::MAX; + + // Multibody side of J · u, one product per lane; lane 0 sums them in + // DOF order. + scratch[lane as usize] = if lane < ndofs { + contact_constraint_jacs.read(col_offset + lane as usize) * dof_v[lane as usize] + } else { + 0.0 + }; + workgroup_memory_barrier_with_group_sync(); + + if lane == 0 { + let mut j_dot_v = 0.0f32; + for i in 0..ndofs { + j_dot_v += scratch[i as usize]; + } + // Free-body side stays lane-0-local. + let free = if is_self { + Velocity::default() + } else { + solver_vels.read(colliders_start + cons.free_body_id as usize) + }; + if !is_self { + j_dot_v += cons.lin_jac.dot(free.linear) + gdot(cons.ang_jac, free.angular); + } + + let rhs = if use_bias { cons.rhs } else { cons.rhs_wo_bias }; + let impulse = imp_shared[s as usize]; + let rhs_total = j_dot_v + rhs; + // CFM-factor form (rapier's `*ContactConstraintNormalPart::generic_solve`). + let raw_imp = cons.cfm_factor * (impulse - cons.inv_lhs * rhs_total); + + // Normal: clamp to ≥ 0. Friction tangent: clamp to + // `±μ · normal_impulse` (box friction), reading the paired normal + // slot's current impulse from shared memory. + let new_imp = if cons.kind == MB_CONTACT_KIND_TANGENT { + let limit = + cons.friction_coeff * imp_shared[cons.normal_constraint_slot as usize]; + if raw_imp > limit { + limit + } else if raw_imp < -limit { + -limit + } else { + raw_imp + } + } else if raw_imp < 0.0 { + 0.0 + } else { + raw_imp + }; + let delta = new_imp - impulse; + imp_shared[s as usize] = new_imp; + scratch[0] = delta; + + if delta != 0.0 && !is_self { + let mut new_free = free; + new_free.linear += cons.lin_jac * (cons.free_body_im * delta); + new_free.angular += cons.ii_ang_jac * delta; + solver_vels.write(colliders_start + cons.free_body_id as usize, new_free); + } + } + workgroup_memory_barrier_with_group_sync(); + + let delta = scratch[0]; + if delta != 0.0 && lane < ndofs { + let col = contact_constraint_columns.read(col_offset + lane as usize); + dof_v[lane as usize] += delta * col; + } + workgroup_memory_barrier_with_group_sync(); + } + + // Writeback + if lane < ndofs { + dof_state.write(v_base + lane as usize, dof_v[lane as usize]); + } + for s in StepRng::new(lane..contact_count, LANES) { + let mut cons = contact_constraints.read(ccons_base + s as usize); + cons.impulse = imp_shared[s as usize]; + contact_constraints.write(ccons_base + s as usize, cons); + } +} From 02d0e56b8e33994e724c2ed0850351da6f9e36b1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Fri, 31 Jul 2026 11:44:28 +0200 Subject: [PATCH 26/39] perf: MuJoCo-style explicit coriolis in the menagerie example --- crates/examples3d/mujoco_menagerie3.rs | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/crates/examples3d/mujoco_menagerie3.rs b/crates/examples3d/mujoco_menagerie3.rs index bde6198..5da2377 100644 --- a/crates/examples3d/mujoco_menagerie3.rs +++ b/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) } From ac0d099dbc3e78397f9e6a995c26b7ce2db1dce8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Fri, 31 Jul 2026 12:06:17 +0200 Subject: [PATCH 27/39] perf: lane-parallelize the multibody contact-constraint jacobian fills --- .../dynamics/multibody/multibody_solver.rs | 7 +- .../dynamics/multibody/contact_constraints.rs | 69 ++++++++++++------- 2 files changed, 51 insertions(+), 25 deletions(-) diff --git a/src_rbd/dynamics/multibody/multibody_solver.rs b/src_rbd/dynamics/multibody/multibody_solver.rs index 6425db5..efd2174 100644 --- a/src_rbd/dynamics/multibody/multibody_solver.rs +++ b/src_rbd/dynamics/multibody/multibody_solver.rs @@ -189,7 +189,6 @@ impl GpuMultibodySolver { if mb.is_empty() { return Ok(()); } - let dispatch = mb.flat_mb_dispatch(); if mb.has_joint_constraints { let mut pass = encoder.begin_pass("[RBD] mbb/init-joint", timestamps.as_deref_mut()); @@ -216,12 +215,16 @@ impl GpuMultibodySolver { // multibody pairs only). `init` PRESERVES the accumulated impulse across // substeps (zeroed once per frame by `reset_contact_warmstart` in // `init_step`); `finalize` recomputes `inv_lhs` and the M⁻¹Jᵀ columns. + // One 64-lane workgroup per multibody: the uniform emission walk runs + // redundantly on every lane, the per-DOF `Jᵀ`-row fills one-per-lane. { let mut pass = encoder.begin_pass("[RBD] mbb/init-contact", timestamps.as_deref_mut()); + let init_contact_dispatch = + [mb.multibodies_per_batch * MB_LU_LANES, mb.num_batches, 1]; self.init_contact_constraints.call( &mut pass, - dispatch, + init_contact_dispatch, &mut mb.multibody_info, &mb.body_jacobians, &mb.body_to_link, diff --git a/src_rbd_shaders/dynamics/multibody/contact_constraints.rs b/src_rbd_shaders/dynamics/multibody/contact_constraints.rs index eac48dc..9d65018 100644 --- a/src_rbd_shaders/dynamics/multibody/contact_constraints.rs +++ b/src_rbd_shaders/dynamics/multibody/contact_constraints.rs @@ -60,6 +60,8 @@ fn orthonormal_vector(v: Vec2) -> Vec2 { /// **adding** the resulting `Jᵀ` row to `out_jacs[col_offset ..]` (so two /// calls accumulate — used by self-collisions, which combine the two /// touched links into a single net `Jᵀ` row). +/// +/// One DOF per lane. #[inline] fn fill_contact_jac_row( body_jacobians: &[f32], @@ -71,13 +73,15 @@ fn fill_contact_jac_row( out_jacs: &mut [f32], col_offset: usize, accumulate: bool, + lane: u32, ) { // Per-link SPATIAL_DIM × ndofs jacobian (rows 0..DIM = J_v, rows // DIM..SPATIAL_DIM = J_w). let link_jac_base = mb_jac_base + (link_id as usize) * SPATIAL_DIM * (ndofs as usize); let link_j = MatSlice::dense(link_jac_base, SPATIAL_DIM as u32, ndofs); let (link_j_v, link_j_w) = link_j.rows_range_pair(0, DIM, DIM, ANG_DIM); - for j in 0..ndofs { + let j = lane; + if j < ndofs { // Linear contribution: `unit_force · J_v[:, j]`. let dot; #[cfg(feature = "dim3")] @@ -118,10 +122,12 @@ fn fill_contact_jac_row( /// `MultibodyContactConstraint` and writes the multibody-side `Jᵀ` row into /// `contact_constraint_jacs`. Multibody-multibody contacts (each side a /// different multibody) are not handled — such contacts are skipped. +/// One 64-lane workgroup per (multibody, batch). #[spirv_bindgen] #[spirv(compute(threads(64)))] pub fn gpu_mb_init_contact_constraints( - #[spirv(global_invocation_id)] invocation_id: UVec3, + #[spirv(workgroup_id)] workgroup_id: UVec3, + #[spirv(local_invocation_id)] local_id: UVec3, #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] multibody_info: &mut [MultibodyInfo], #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] body_jacobians: &[f32], @@ -135,16 +141,16 @@ pub fn gpu_mb_init_contact_constraints( #[spirv(storage_buffer, descriptor_set = 1, binding = 2)] contacts: &[IndexedManifold], #[spirv(uniform, descriptor_set = 0, binding = 6)] batch_ids: &BatchIndices, ) { - // Flattened (multibody, batch) grid — see `BatchIndices::num_batches`. // Only active multibody slots are visited now; the `ndofs == 0` sentinel // below is kept for all-locked (zero-dof) multibodies. Padding slots past // `multibodies_len` are never read (every consumer guards on it). let num_mb = batch_ids.multibodies_len; - if invocation_id.x >= num_mb * batch_ids.num_batches { + let batch_id = workgroup_id.y; + let mb_idx = workgroup_id.x; + let lane = local_id.x; + if mb_idx >= num_mb { return; } - let batch_id = invocation_id.x / num_mb; - let mb_idx = invocation_id.x % num_mb; // Soft-constraint coefficients (rapier TGS-soft), precomputed on the host. // The old path used a rigid `erp = 1/dt` with zero CFM, which overshoots // penetration recovery (~14× too stiff for the defaults) and jitters. @@ -167,8 +173,11 @@ pub fn gpu_mb_init_contact_constraints( let mut mb = multibody_info.read(mb_start + mb_idx as usize); let ndofs = mb.ndofs; if ndofs == 0 { - mb.contact_constraint_count = 0; - multibody_info.write(mb_start + mb_idx as usize, mb); + // Uniform per workgroup: every lane returns together. + if lane == 0 { + mb.contact_constraint_count = 0; + multibody_info.write(mb_start + mb_idx as usize, mb); + } return; } let mb_jac_base = batch_ids.jac_start(batch_id) + mb.jacobian_offset as usize; @@ -299,13 +308,13 @@ pub fn gpu_mb_init_contact_constraints( let normal_slot = count; let normal_col_offset = col_base + (normal_slot as usize) * dofs_stride; - // Warmstart: preserve the accumulated impulse from the previous - // substep (same contact slot — within a frame the manifolds are - // fixed). `gpu_mb_reset_contact_warmstart` zeroes these once per - // frame so the first substep starts cold. - let warmstart_normal_impulse = contact_constraints - .read(cons_base + normal_slot as usize) - .impulse; + let warmstart_normal_impulse = if lane == 0 { + contact_constraints + .read(cons_base + normal_slot as usize) + .impulse + } else { + 0.0 + }; fill_contact_jac_row( body_jacobians, @@ -317,6 +326,7 @@ pub fn gpu_mb_init_contact_constraints( contact_constraint_jacs, normal_col_offset, false, + lane, ); // B-side fold-in for self-contacts, free body for the rest. The @@ -341,6 +351,7 @@ pub fn gpu_mb_init_contact_constraints( contact_constraint_jacs, normal_col_offset, true, + lane, ); #[cfg(feature = "dim3")] { @@ -403,7 +414,9 @@ pub fn gpu_mb_init_contact_constraints( cfm_factor, _unused_cfm: 0.0, }; - contact_constraints.write(cons_base + normal_slot as usize, normal_cons); + if lane == 0 { + contact_constraints.write(cons_base + normal_slot as usize, normal_cons); + } count += 1; // Friction tangent constraints — same contact point, tangent @@ -441,10 +454,14 @@ pub fn gpu_mb_init_contact_constraints( let tang_slot = count; let tang_col_offset = col_base + (tang_slot as usize) * dofs_stride; // Warmstart: preserve the accumulated tangent impulse (see the - // normal slot above). - let warmstart_tang_impulse = contact_constraints - .read(cons_base + tang_slot as usize) - .impulse; + // normal slot above; lane 0 only). + let warmstart_tang_impulse = if lane == 0 { + contact_constraints + .read(cons_base + tang_slot as usize) + .impulse + } else { + 0.0 + }; fill_contact_jac_row( body_jacobians, @@ -456,6 +473,7 @@ pub fn gpu_mb_init_contact_constraints( contact_constraint_jacs, tang_col_offset, false, + lane, ); let (ang_jac_tang, ii_ang_jac_tang) = if is_self { @@ -472,6 +490,7 @@ pub fn gpu_mb_init_contact_constraints( contact_constraint_jacs, tang_col_offset, true, + lane, ); #[cfg(feature = "dim3")] { @@ -533,7 +552,9 @@ pub fn gpu_mb_init_contact_constraints( cfm_factor, _unused_cfm: 0.0, }; - contact_constraints.write(cons_base + tang_slot as usize, tang_cons); + if lane == 0 { + contact_constraints.write(cons_base + tang_slot as usize, tang_cons); + } count += 1; } } @@ -541,8 +562,10 @@ pub fn gpu_mb_init_contact_constraints( // The solve / finalize / remove-bias kernels iterate `0..count` so we // don't need to mark surplus slots inactive — they're never read. - mb.contact_constraint_count = count; - multibody_info.write(mb_start + mb_idx as usize, mb); + if lane == 0 { + mb.contact_constraint_count = count; + multibody_info.write(mb_start + mb_idx as usize, mb); + } } /// HACK: stash `contacts_len[batch]` into each multibody's `batch_contacts_len`. From 72a6d37e7ee641e41b8a1e9d45ae4621b12358d6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Fri, 31 Jul 2026 12:10:42 +0200 Subject: [PATCH 28/39] perf: trim a barrier from the mb contact sweep + split the init-step passes --- .../dynamics/multibody/multibody_solver.rs | 24 ++++++++++++------- src_rbd/pipeline/rbd_step.rs | 13 +++++----- .../dynamics/multibody/solve_constraints.rs | 7 +++--- 3 files changed, 26 insertions(+), 18 deletions(-) diff --git a/src_rbd/dynamics/multibody/multibody_solver.rs b/src_rbd/dynamics/multibody/multibody_solver.rs index efd2174..5e0b7ab 100644 --- a/src_rbd/dynamics/multibody/multibody_solver.rs +++ b/src_rbd/dynamics/multibody/multibody_solver.rs @@ -106,24 +106,30 @@ impl GpuMultibodySolver { /// the last call carrying `is_last_substep = true`. pub fn init_step( &self, - pass: &mut GpuPass, + encoder: &mut khal::backend::GpuEncoder, + mut timestamps: Option<&mut khal::backend::GpuTimestamps>, mb: &mut GpuMultibodySet, args: &mut MultibodySolverArgs<'_>, ) -> Result<(), GpuBackendError> { + use khal::backend::Encoder; if mb.is_empty() { return Ok(()); } // Zero the accumulated contact impulses so the first substep's warmstart // starts cold (within a frame they are then preserved across substeps). // One 64-lane workgroup per multibody (lanes stride the slots). - self.reset_contact_warmstart.call( - pass, - [mb.multibodies_per_batch * MB_LU_LANES, mb.num_batches, 1], - &mb.multibody_info, - &mut mb.contact_constraints, - args.batch_indices, - )?; - self.compute_dynamics(pass, mb, args) + { + let mut pass = encoder.begin_pass("[RBD] mbi/reset", timestamps.as_deref_mut()); + self.reset_contact_warmstart.call( + &mut pass, + [mb.multibodies_per_batch * MB_LU_LANES, mb.num_batches, 1], + &mb.multibody_info, + &mut mb.contact_constraints, + args.batch_indices, + )?; + } + let mut pass = encoder.begin_pass("[RBD] mbi/dynamics", timestamps.as_deref_mut()); + self.compute_dynamics(&mut pass, mb, args) } /// Copy `contacts_len[batch]` into each `MultibodyInfo`. diff --git a/src_rbd/pipeline/rbd_step.rs b/src_rbd/pipeline/rbd_step.rs index 0f0c8bf..66b4e3b 100644 --- a/src_rbd/pipeline/rbd_step.rs +++ b/src_rbd/pipeline/rbd_step.rs @@ -85,8 +85,6 @@ impl RbdPipeline { { if !state.multibodies.is_empty() { let mut encoder = backend.begin_encoding(); - let mut pass = - encoder.begin_pass("[RBD] multibody-init-step", timestamps.as_deref_mut()); let mut args = crate::dynamics::MultibodySolverArgs { poses: &mut state.body_poses, collider_world_poses: &state.collider_world_poses, @@ -95,11 +93,14 @@ impl RbdPipeline { contacts_len: &state.contacts_len, solver_vels: &mut state.solver_vels, batch_indices: &state.batch_indices, - color_uniforms: &state.color_uniforms, + color_uniforms: &state.color_uniforms, }; - self.multibody_solver - .init_step(&mut pass, &mut state.multibodies, &mut args)?; - drop(pass); + self.multibody_solver.init_step( + &mut encoder, + timestamps.as_deref_mut(), + &mut state.multibodies, + &mut args, + )?; backend.submit(encoder)?; } } diff --git a/src_rbd_shaders/dynamics/multibody/solve_constraints.rs b/src_rbd_shaders/dynamics/multibody/solve_constraints.rs index ecdbd36..6160eaf 100644 --- a/src_rbd_shaders/dynamics/multibody/solve_constraints.rs +++ b/src_rbd_shaders/dynamics/multibody/solve_constraints.rs @@ -43,6 +43,7 @@ pub fn gpu_mb_solve_constraints( #[spirv(workgroup)] dof_v: &mut [f32; MAX_MB_DOFS as usize], #[spirv(workgroup)] scratch: &mut [f32; LANES as usize], #[spirv(workgroup)] imp_shared: &mut [f32; MAX_MB_CONTACT_CONSTRAINTS_PER_MB as usize], + #[spirv(workgroup)] delta_shared: &mut f32, ) { let batch_id = workgroup_id.y; let mb_idx = workgroup_id.x; @@ -186,7 +187,7 @@ pub fn gpu_mb_solve_constraints( }; let delta = new_imp - impulse; imp_shared[s as usize] = new_imp; - scratch[0] = delta; + *delta_shared = delta; if delta != 0.0 && !is_self { let mut new_free = free; @@ -197,12 +198,12 @@ pub fn gpu_mb_solve_constraints( } workgroup_memory_barrier_with_group_sync(); - let delta = scratch[0]; + // Per-lane `dof_v[lane]` update. + let delta = *delta_shared; if delta != 0.0 && lane < ndofs { let col = contact_constraint_columns.read(col_offset + lane as usize); dof_v[lane as usize] += delta * col; } - workgroup_memory_barrier_with_group_sync(); } // Writeback From c826829768138ee03987f8606bd6923a8a87d661 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Fri, 31 Jul 2026 12:31:12 +0200 Subject: [PATCH 29/39] perf: build multibody contact constraints once per step in explicit-coriolis mode --- .../dynamics/multibody/multibody_solver.rs | 85 +++++++++++++------ src_rbd/dynamics/solver.rs | 28 ++++++ 2 files changed, 86 insertions(+), 27 deletions(-) diff --git a/src_rbd/dynamics/multibody/multibody_solver.rs b/src_rbd/dynamics/multibody/multibody_solver.rs index 5e0b7ab..f58439e 100644 --- a/src_rbd/dynamics/multibody/multibody_solver.rs +++ b/src_rbd/dynamics/multibody/multibody_solver.rs @@ -217,10 +217,64 @@ impl GpuMultibodySolver { )?; } - // Build + finalize contact constraints (normal-only, free body × - // multibody pairs only). `init` PRESERVES the accumulated impulse across - // substeps (zeroed once per frame by `reset_contact_warmstart` in - // `init_step`); `finalize` recomputes `inv_lhs` and the M⁻¹Jᵀ columns. + // With implicit coriolis, the mass matrix / LU / body jacobians are + // recomputed every substep, so the contact constraints (whose M⁻¹Jᵀ + // columns depend on them) must be rebuilt every substep too. In the + // explicit mode every input is a per-step constant, so the pipeline + // builds them ONCE per step instead (see `build_contact_constraints`). + if mb.implicit_coriolis { + self.build_contact_constraints(encoder, timestamps.as_deref_mut(), mb, args)?; + } + + // Warmstart: re-apply the accumulated contact impulse to dof_state (and + // the free-body solver velocities) so the contact starts "warm" each + // substep — mirrors rapier's per-substep `contact_constraints.warmstart` + // and matches what the rigid-body solver does for free contacts. On the + // first substep the impulse was just reset to 0, so this is a no-op. + // One 64-lane workgroup per multibody (one DOF per lane). + { + let mut pass = + encoder.begin_pass("[RBD] mbb/warmstart-contact", timestamps.as_deref_mut()); + let warmstart_dispatch = + [mb.multibodies_per_batch * MB_LU_LANES, mb.num_batches, 1]; + self.warmstart_contact_constraints.call( + &mut pass, + warmstart_dispatch, + &mb.multibody_info, + &mb.contact_constraints, + &mb.contact_constraint_columns, + &mut mb.dof_state, + args.solver_vels, + args.batch_indices, + )?; + } + + Ok(()) + } + + /// Build + finalize the contact constraints (normal + friction slots, + /// free-body × multibody and self-contact pairs). `init` PRESERVES the + /// accumulated impulse across substeps (zeroed once per frame by + /// `reset_contact_warmstart` in `init_step`); `finalize` computes + /// `inv_lhs` and the M⁻¹Jᵀ columns. + /// + /// Inputs are the narrow-phase manifolds, the collider world poses, the + /// body jacobians and the mass-matrix LU. The first two only change once + /// per step; the last two change per substep ONLY with implicit coriolis. + /// So this runs once per step (from `solve_tgs`'s init pass, after the + /// narrow phase) in the explicit mode, and once per substep otherwise. + pub fn build_contact_constraints( + &self, + encoder: &mut khal::backend::GpuEncoder, + mut timestamps: Option<&mut khal::backend::GpuTimestamps>, + mb: &mut GpuMultibodySet, + args: &mut MultibodySolverArgs<'_>, + ) -> Result<(), GpuBackendError> { + use khal::backend::Encoder; + if mb.is_empty() { + return Ok(()); + } + // One 64-lane workgroup per multibody: the uniform emission walk runs // redundantly on every lane, the per-DOF `Jᵀ`-row fills one-per-lane. { @@ -264,29 +318,6 @@ impl GpuMultibodySolver { )?; } - // Warmstart: re-apply the accumulated contact impulse to dof_state (and - // the free-body solver velocities) so the contact starts "warm" each - // substep — mirrors rapier's per-substep `contact_constraints.warmstart` - // and matches what the rigid-body solver does for free contacts. On the - // first substep the impulse was just reset to 0, so this is a no-op. - // One 64-lane workgroup per multibody (one DOF per lane). - { - let mut pass = - encoder.begin_pass("[RBD] mbb/warmstart-contact", timestamps.as_deref_mut()); - let warmstart_dispatch = - [mb.multibodies_per_batch * MB_LU_LANES, mb.num_batches, 1]; - self.warmstart_contact_constraints.call( - &mut pass, - warmstart_dispatch, - &mb.multibody_info, - &mb.contact_constraints, - &mb.contact_constraint_columns, - &mut mb.dof_state, - args.solver_vels, - args.batch_indices, - )?; - } - Ok(()) } diff --git a/src_rbd/dynamics/solver.rs b/src_rbd/dynamics/solver.rs index eec6817..3a2afc0 100644 --- a/src_rbd/dynamics/solver.rs +++ b/src_rbd/dynamics/solver.rs @@ -289,6 +289,34 @@ impl GpuSolver { } } + // Explicit-coriolis mode: the multibody contact constraints only + // depend on per-step constants (manifolds, collider world poses, the + // once-per-step body jacobians and mass-matrix LU), so build them ONCE + // here instead of once per substep. The accumulated impulses still + // persist in the constraint slots across substeps, exactly as with + // the per-substep rebuild (which preserved them explicitly). + #[cfg(feature = "dim3")] + if let (Some(solver), Some(state)) = (mb_solver, mb_state.as_deref_mut()) { + if !state.implicit_coriolis() { + let mut mb_args = MultibodySolverArgs { + poses: &mut *args.solver_body_poses, + collider_world_poses: args.collider_world_poses, + mprops: args.mprops, + contacts: args.contacts, + contacts_len: args.contacts_len, + solver_vels: &mut *args.solver_vels, + batch_indices: args.batch_indices, + color_uniforms: args.color_uniforms, + }; + solver.build_contact_constraints( + encoder, + timestamps.as_deref_mut(), + state, + &mut mb_args, + )?; + } + } + // Per substep, the multibody work is split into five phases that are // INTERLEAVED with the matching rigid-body phases, mirroring rapier's // `velocity_solver::solve_constraints` order: From da69b58be06415c483fe21a5e652b682e73133ad Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Fri, 31 Jul 2026 13:39:04 +0200 Subject: [PATCH 30/39] perf: skip the rigid-body contact pipeline when it is provably inert --- src_rbd/dynamics/solver.rs | 80 +++++++++++++---------- src_rbd/pipeline/insertion_removal.rs | 7 ++ src_rbd/pipeline/rbd_state.rs | 8 +++ src_rbd/pipeline/rbd_state_from_rapier.rs | 11 ++++ src_rbd/pipeline/rbd_step.rs | 10 +++ 5 files changed, 83 insertions(+), 33 deletions(-) diff --git a/src_rbd/dynamics/solver.rs b/src_rbd/dynamics/solver.rs index 3a2afc0..05f5d2b 100644 --- a/src_rbd/dynamics/solver.rs +++ b/src_rbd/dynamics/solver.rs @@ -142,10 +142,9 @@ pub struct SolverArgs<'a> { /// This is generally used when the number of constraints is small wrt. /// the number of environments. pub fused_color_sweeps: bool, - /// Shared per-batch capacity / section-offset uniform — see - /// [`crate::shaders::utils::BatchIndices`]. Consumed by the (refactored) - /// multibody kernels via `MultibodySolverArgs::batch_indices`; the RBD - /// constraint-solver kernels will migrate to it next. + /// `true` when every rigid-body contact constraint is provably a no-op. + pub rb_contacts_inert: bool, + /// Shared per-batch indices. pub batch_indices: &'a Tensor, } @@ -183,6 +182,10 @@ impl GpuSolver { args.batch_indices, )?; + if args.rb_contacts_inert { + return Ok(()); + } + self.init_constraints.call( pass, args.contacts_len_indirect, @@ -255,19 +258,24 @@ impl GpuSolver { None => (None, None), }; + let skip_rb = args.rb_contacts_inert; + let joints_empty = joint_args.joints.is_empty(); + /* * Init solver vel increments. */ { let mut pass = encoder.begin_pass("[RBD] slv/init", timestamps.as_deref_mut()); - self.init_solver_vels_inc.call( - &mut pass, - [args.num_colliders, args.num_batches, 1], - args.solver_vels_inc, - args.mprops, - args.sim_params, - args.batch_indices, - )?; + if !skip_rb { + self.init_solver_vels_inc.call( + &mut pass, + [args.num_colliders, args.num_batches, 1], + args.solver_vels_inc, + args.mprops, + args.sim_params, + args.batch_indices, + )?; + } joint_solver.init(&mut pass, &mut joint_args)?; @@ -364,7 +372,7 @@ impl GpuSolver { * P1/F1 — integrate velocities (apply `a · dt'` / gravity increment). */ mb_phase!("[RBD] slv/mb-integrate-vels", substep_integrate_velocities); - { + if !skip_rb { let mut pass = encoder.begin_pass("[RBD] slv/rb-apply-inc", timestamps.as_deref_mut()); self.apply_solver_vels_inc.call( @@ -400,25 +408,27 @@ impl GpuSolver { )?; } } - { + if !skip_rb || !joints_empty { let mut pass = encoder.begin_pass("[RBD] slv/rb-build-warmstart", timestamps.as_deref_mut()); let pass = &mut pass; - 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, - )?; + if !skip_rb { + 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.colorless_warmstart { - // One gather dispatch over bodies instead of `num_colors` - // scatter dispatches (each constraint is visited once per - // body side, but the dispatch count drops by ~num_colors). + if skip_rb { + // Contact warmstart skipped: no rigid-body contact + // constraint can carry an impulse here. + } else if args.colorless_warmstart { self.warmstart_without_colors.call( pass, [args.num_colliders, args.num_batches, 1], @@ -463,12 +473,14 @@ impl GpuSolver { * Solve all joints + contacts with bias. */ mb_phase!("[RBD] slv/mb-solve-bias", substep_solve_with_bias); - { + if !skip_rb || !joints_empty { let mut pass = encoder.begin_pass("[RBD] slv/rb-solve-bias", timestamps.as_deref_mut()); let pass = &mut pass; joint_solver.solve(pass, &mut joint_args, args.solver_vels, true)?; - if args.fused_color_sweeps { + if skip_rb { + // Contact sweeps skipped (inert constraints). + } else if args.fused_color_sweeps { self.step_gauss_seidel_fused.call( pass, [64, args.num_batches, 1], @@ -507,7 +519,7 @@ impl GpuSolver { substep_integrate_positions, is_last_substep ); - { + if !skip_rb { let mut pass = encoder.begin_pass("[RBD] slv/rb-integrate", timestamps.as_deref_mut()); self.integrate_linearized.call( @@ -524,12 +536,14 @@ impl GpuSolver { * P5/F5 — solve ALL joints + contacts WITHOUT bias (stabilization). */ mb_phase!("[RBD] slv/mb-solve-nobias", substep_solve_no_bias); - { + if !skip_rb || !joints_empty { let mut pass = encoder.begin_pass("[RBD] slv/rb-solve-nobias", timestamps.as_deref_mut()); let pass = &mut pass; joint_solver.solve(pass, &mut joint_args, args.solver_vels, false)?; - if args.fused_color_sweeps { + if skip_rb { + // Contact sweeps skipped (inert constraints). + } else if args.fused_color_sweeps { self.step_gauss_seidel_fused.call( pass, [64, args.num_batches, 1], diff --git a/src_rbd/pipeline/insertion_removal.rs b/src_rbd/pipeline/insertion_removal.rs index d4f9f70..c20e42d 100644 --- a/src_rbd/pipeline/insertion_removal.rs +++ b/src_rbd/pipeline/insertion_removal.rs @@ -324,6 +324,7 @@ impl RbdState { prefix_sum_workspace: PrefixSumWorkspace::default(), lbvh: LbvhState::with_usages(backend, lbvh_usages), max_colors: capacities.solver_colors, + rb_contacts_inert: false, num_active_colliders: 0, num_active_bodies: 0, } @@ -369,6 +370,12 @@ impl RbdState { let body_pose = *rb.position(); let collider_local_pose = co.position_wrt_parent().copied().unwrap_or(Pose::IDENTITY); let is_dynamic = rb.is_dynamic(); + if is_dynamic { + // A free dynamic body: its contacts need the rigid-body + // constraint pipeline. Never set back on removal; a stale + // `false` only costs performance. + self.rb_contacts_inert = false; + } let (local, world) = if is_dynamic { // A standalone rigid-body carries no collider mass: rapier only // folds a collider's mass into the body once the collider is diff --git a/src_rbd/pipeline/rbd_state.rs b/src_rbd/pipeline/rbd_state.rs index 328e88e..6591c8c 100644 --- a/src_rbd/pipeline/rbd_state.rs +++ b/src_rbd/pipeline/rbd_state.rs @@ -232,6 +232,9 @@ pub struct RbdState { pub(super) prefix_sum_workspace: PrefixSumWorkspace, /// Maximum number of constraint colors the solver will iterate. pub(super) max_colors: u32, + /// `true` when every body is either non-dynamic or multibody-controlled + /// (its rb-side `inv_mass` is zero),i.e., we can skip the contact pipelines. + pub(super) rb_contacts_inert: bool, /// CPU-side mirror of the number of *active* colliders per batch. Identical /// across all batches by the equal-topology invariant; slots in /// `[num_active_colliders .. num_colliders_per_batch)` are reserved padding. @@ -295,6 +298,11 @@ impl RbdState { pub fn max_colors(&self) -> u32 { self.max_colors } + + /// `true` when every rigid-body contact constraint is provably a no-op. + pub fn rb_contacts_inert(&self) -> bool { + self.rb_contacts_inert + } } impl RbdState { diff --git a/src_rbd/pipeline/rbd_state_from_rapier.rs b/src_rbd/pipeline/rbd_state_from_rapier.rs index bd4de5d..8d427b4 100644 --- a/src_rbd/pipeline/rbd_state_from_rapier.rs +++ b/src_rbd/pipeline/rbd_state_from_rapier.rs @@ -515,6 +515,16 @@ impl RbdState { } } + let rb_contacts_inert = all_env_body_counts + .iter() + .enumerate() + .all(|(batch, &count)| { + let start = batch * max_colliders; + all_local_mprops[start..start + count] + .iter() + .all(|m| m.inv_mass == Vector::ZERO) + }); + // Build the per-body "graph group" lookup. Free bodies map to themselves // (one body = one graph node). Bodies belonging to a multibody all map // to a single shared group id (= the body id of the multibody's root @@ -842,6 +852,7 @@ impl RbdState { prefix_sum_workspace: PrefixSumWorkspace::default(), lbvh: LbvhState::with_usages(backend, lbvh_usages), max_colors: capacities.solver_colors, + rb_contacts_inert, num_active_colliders: num_colliders as u32, num_active_bodies: num_bodies as u32, } diff --git a/src_rbd/pipeline/rbd_step.rs b/src_rbd/pipeline/rbd_step.rs index 66b4e3b..a7ad103 100644 --- a/src_rbd/pipeline/rbd_step.rs +++ b/src_rbd/pipeline/rbd_step.rs @@ -301,6 +301,7 @@ impl RbdPipeline { batch_indices: &state.batch_indices, colorless_warmstart: false, fused_color_sweeps, + rb_contacts_inert: state.rb_contacts_inert, }; self.solver.prepare( backend, @@ -309,6 +310,12 @@ impl RbdPipeline { &mut state.prefix_sum_workspace, )?; + if state.rb_contacts_inert { + stats.num_colors = state.max_colors + 1; + drop(pass); + backend.submit(encoder)?; + } else { + // Warmstart let warmstart_args = WarmstartArgs { contacts_len: &state.contacts_len, @@ -402,6 +409,7 @@ impl RbdPipeline { drop(pass); backend.submit(encoder)?; + } } let num_colors = stats.num_colors; @@ -442,6 +450,7 @@ impl RbdPipeline { #[cfg(not(feature = "dim3"))] colorless_warmstart: true, fused_color_sweeps, + rb_contacts_inert: state.rb_contacts_inert, }; // Phase 3: Solve constraints @@ -531,6 +540,7 @@ impl RbdPipeline { // count earlier, before it gets a chance to fail. if state.capacities.solver_colors_resize_policy != RbdResizePolicy::Fixed && coloring_converged == 0 + && !state.rb_contacts_inert { state.max_colors += 5; From 48865f669c367518e0cb0cd31775fddbc6ebf539 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Fri, 31 Jul 2026 14:09:01 +0200 Subject: [PATCH 31/39] perf: constraint-space (Delassus) multibody contact sweep --- .../multibody/multibody_from_rapier.rs | 24 ++ src_rbd/dynamics/multibody/multibody_set.rs | 10 + .../dynamics/multibody/multibody_solver.rs | 135 ++++++-- .../dynamics/multibody/solve_constraints.rs | 327 ++++++++++++++++++ 4 files changed, 458 insertions(+), 38 deletions(-) diff --git a/src_rbd/dynamics/multibody/multibody_from_rapier.rs b/src_rbd/dynamics/multibody/multibody_from_rapier.rs index 7547774..9b5cad2 100644 --- a/src_rbd/dynamics/multibody/multibody_from_rapier.rs +++ b/src_rbd/dynamics/multibody/multibody_from_rapier.rs @@ -471,6 +471,30 @@ impl GpuMultibodySet { storage, ) .unwrap(), + // Per-multibody Delassus blocks for the constraint-space contact + // sweep: MAX_MB_CONTACT_CONSTRAINTS_PER_MB² floats each (147 KB + // in 3D), so only small total multibody counts get them; larger + // batched scenes keep the dof-space solve. + contact_delassus: { + // Sized by the capacity stride (the kernels index blocks by + // `batch · multibodies_batch_capacity + mb_idx`). + let total_mbs = mb_cap * num_batches; + if global_max_mb > 0 && total_mbs <= MAX_DELASSUS_MULTIBODIES { + let block = (MAX_MB_CONTACT_CONSTRAINTS_PER_MB + * MAX_MB_CONTACT_CONSTRAINTS_PER_MB) + as usize; + Some( + Tensor::vector_uninit( + backend, + (total_mbs as usize * block) as u32, + storage, + ) + .unwrap(), + ) + } else { + None + } + }, // Impulse-joint buffers are sized for "no MB-touching joints" by // default — `set_impulse_joints` resizes them at pipeline build diff --git a/src_rbd/dynamics/multibody/multibody_set.rs b/src_rbd/dynamics/multibody/multibody_set.rs index b7ad7c1..82a3aad 100644 --- a/src_rbd/dynamics/multibody/multibody_set.rs +++ b/src_rbd/dynamics/multibody/multibody_set.rs @@ -19,6 +19,12 @@ use vortx::tensor::Tensor; /// `gpu_mb_lu_solve`. pub(super) const MB_LU_LANES: u32 = 64; +/// Maximum total multibody count (capacity × batches) for which the +/// constraint-space (Delassus) contact solve is enabled: each multibody's +/// Delassus block costs `MAX_MB_CONTACT_CONSTRAINTS_PER_MB²` floats (~147 KB +/// in 3D), so huge batched scenes would run out of memory. +pub(super) const MAX_DELASSUS_MULTIBODIES: u32 = 128; + use crate::shaders::dynamics::{GenericJoint, JointLimits, JointMotor}; /// GPU-resident articulated multibody set, packed across simulation batches. @@ -91,6 +97,10 @@ pub struct GpuMultibodySet { pub(super) contact_constraint_jacs: Tensor, /// Per-constraint M⁻¹·Jᵀ column (length `ndofs`). pub(super) contact_constraint_columns: Tensor, + /// Per-multibody Delassus blocks (`MAX_MB_CONTACT_CONSTRAINTS_PER_MB²` + /// floats each) only allocated when the total multibody count is at most + /// [`MAX_DELASSUS_MULTIBODIES`]. + pub(super) contact_delassus: Option>, /// Per-batch number of multibody-touching impulse joints (body1 OR body2 /// part of any multibody). diff --git a/src_rbd/dynamics/multibody/multibody_solver.rs b/src_rbd/dynamics/multibody/multibody_solver.rs index f58439e..24cafb8 100644 --- a/src_rbd/dynamics/multibody/multibody_solver.rs +++ b/src_rbd/dynamics/multibody/multibody_solver.rs @@ -4,14 +4,15 @@ use super::multibody_set::*; use crate::math::Pose; use crate::queries::GpuIndexedContact; use crate::shaders::dynamics::{ - GpuMbComputeDynamicsPre, + GpuMbBuildContactDelassus, GpuMbComputeDynamicsPre, GpuMbComputeDynamicsWithoutCoriolisPre, GpuMbFinalizeContactConstraints, GpuMbGravityAndLu, GpuMbGravityAndLuT8, GpuMbGravityAndLuT16, GpuMbGravityAndLuT32, GpuMbInitContactConstraints, GpuMbInitJointConstraints, GpuMbIntegrate, GpuMbIntegrateVelocities, GpuMbRemoveImpulseJointConstraintBias, GpuMbResetContactWarmstart, GpuMbStashContactsLen, GpuMbWarmstartContactConstraints, - GpuMbSolveConstraints, GpuMbSolveImpulseJointConstraints, + GpuMbSolveConstraints, GpuMbSolveContactsDelassus, GpuMbSolveImpulseJointConstraints, + GpuMbSolveJoints, GpuMbFinalizeImpulseJointConstraints, GpuMbUpdateImpulseJointConstraints, Velocity, WorldMassProperties, }; @@ -39,6 +40,17 @@ pub struct GpuMultibodySolver { /// Fused joint+contact PGS sweep (one workgroup per multibody, shared- /// memory dof velocities). solve_constraints: GpuMbSolveConstraints, + /// Joint-only half of the sweep, used with the Delassus contact path + /// (one kernel binding both joint and Delassus buffers would exceed the + /// 8-storage-buffer budget). + solve_joints: GpuMbSolveJoints, + /// Fills the per-multibody Delassus blocks (`D = J M⁻¹ Jᵀ` + free-body + /// coupling) right after the contact columns are finalized. + build_contact_delassus: GpuMbBuildContactDelassus, + /// Constraint-space contact sweep: `a = J·u` tracked incrementally in + /// shared memory via the Delassus rows, breaking the per-iteration + /// dof-space latency chain. + solve_contacts_delassus: GpuMbSolveContactsDelassus, /// Zero the accumulated contact impulses once per frame (warmstart reset). reset_contact_warmstart: GpuMbResetContactWarmstart, /// Copy `contacts_len[batch]` into each `MultibodyInfo` once per step so @@ -318,6 +330,82 @@ impl GpuMultibodySolver { )?; } + // Delassus blocks for the constraint-space contact sweep (consumes + // the columns finalized just above). + if let Some(delassus) = &mut mb.contact_delassus { + let mut pass = + encoder.begin_pass("[RBD] mbb/build-delassus", timestamps.as_deref_mut()); + let dispatch = [mb.multibodies_per_batch * MB_LU_LANES, mb.num_batches, 1]; + self.build_contact_delassus.call( + &mut pass, + dispatch, + &mb.multibody_info, + &mb.contact_constraints, + &mb.contact_constraint_jacs, + &mb.contact_constraint_columns, + delassus, + args.batch_indices, + )?; + } + + Ok(()) + } + + /// One joint+contact PGS sweep: the dof-space fused kernel, or (when the + /// Delassus blocks are allocated) the joint-only kernel followed by the + /// constraint-space contact kernel. `use_bias_idx` indexes + /// `color_uniforms` (0 or 1, holding those constants). + fn dispatch_solve( + &self, + pass: &mut GpuPass, + mb: &mut GpuMultibodySet, + args: &mut MultibodySolverArgs<'_>, + solve_dispatch: [u32; 3], + use_bias_idx: usize, + ) -> Result<(), GpuBackendError> { + let use_bias = &args.color_uniforms[use_bias_idx]; + if let Some(delassus) = &mb.contact_delassus { + if mb.has_joint_constraints { + self.solve_joints.call( + pass, + solve_dispatch, + &mb.multibody_info, + &mut mb.joint_constraints, + &mb.joint_constraint_columns, + &mut mb.dof_state, + use_bias, + args.batch_indices, + )?; + } + self.solve_contacts_delassus.call( + pass, + solve_dispatch, + &mb.multibody_info, + &mut mb.contact_constraints, + &mb.contact_constraint_jacs, + &mb.contact_constraint_columns, + delassus, + use_bias, + args.batch_indices, + &mut mb.dof_state, + args.solver_vels, + )?; + } else { + self.solve_constraints.call( + pass, + solve_dispatch, + &mb.multibody_info, + &mut mb.joint_constraints, + &mb.joint_constraint_columns, + &mut mb.contact_constraints, + &mb.contact_constraint_jacs, + &mb.contact_constraint_columns, + use_bias, + args.batch_indices, + &mut mb.dof_state, + args.solver_vels, + )?; + } Ok(()) } @@ -333,24 +421,12 @@ impl GpuMultibodySolver { return Ok(()); } - // Fused joint+contact sweep: one 64-lane workgroup per multibody with - // the generalized velocities held in workgroup memory - // (`color_uniforms[1]` holds the constant 1 = use_bias). + // One 64-lane workgroup per multibody with the generalized velocities + // held in workgroup memory (`color_uniforms[1]` holds the constant + // 1 = use_bias). With the Delassus blocks allocated, the contact half + // runs in constraint space instead (joints first, same order). let solve_dispatch = [mb.multibodies_per_batch * MB_LU_LANES, mb.num_batches, 1]; - self.solve_constraints.call( - pass, - solve_dispatch, - &mb.multibody_info, - &mut mb.joint_constraints, - &mb.joint_constraint_columns, - &mut mb.contact_constraints, - &mb.contact_constraint_jacs, - &mb.contact_constraint_columns, - &args.color_uniforms[1], - args.batch_indices, - &mut mb.dof_state, - args.solver_vels, - )?; + self.dispatch_solve(pass, mb, args, solve_dispatch, 1)?; // Multibody-touching impulse joints — generic (rb-mb / mb-mb) // constraints. Mirrors rapier's `JointGenericExternalConstraintBuilder::update` @@ -469,26 +545,9 @@ impl GpuMultibodySolver { return Ok(()); } - // Fused joint+contact stabilization sweep: `use_bias = 0` - // (`color_uniforms[0]`) makes the kernel read `rhs_wo_bias` directly, - // which replaces the former remove-bias read-modify-write dispatches - // (every constraint is re-initialized next substep, so the persistent - // `rhs` rewrite was never needed). + // Stabilization sweep: `use_bias = 0` (`color_uniforms[0] == 0`). let solve_dispatch = [mb.multibodies_per_batch * MB_LU_LANES, mb.num_batches, 1]; - self.solve_constraints.call( - pass, - solve_dispatch, - &mb.multibody_info, - &mut mb.joint_constraints, - &mb.joint_constraint_columns, - &mut mb.contact_constraints, - &mb.contact_constraint_jacs, - &mb.contact_constraint_columns, - &args.color_uniforms[0], - args.batch_indices, - &mut mb.dof_state, - args.solver_vels, - )?; + self.dispatch_solve(pass, mb, args, solve_dispatch, 0)?; if mb.mb_imp_joints_per_batch > 0 { let imp_dispatch = [mb.mb_imp_joints_per_batch, mb.num_batches, 1]; self.remove_impulse_joint_constraint_bias.call( diff --git a/src_rbd_shaders/dynamics/multibody/solve_constraints.rs b/src_rbd_shaders/dynamics/multibody/solve_constraints.rs index 6160eaf..4f8af6a 100644 --- a/src_rbd_shaders/dynamics/multibody/solve_constraints.rs +++ b/src_rbd_shaders/dynamics/multibody/solve_constraints.rs @@ -216,3 +216,330 @@ pub fn gpu_mb_solve_constraints( contact_constraints.write(ccons_base + s as usize, cons); } } + +/// Joint-only PGS sweep (the joint half of [`gpu_mb_solve_constraints`]), +/// used by the Delassus path where the contact half runs in constraint space +/// as a separate dispatch (to avoid exceeding the 8-storage-buffer budget). +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn gpu_mb_solve_joints( + #[spirv(workgroup_id)] workgroup_id: UVec3, + #[spirv(local_invocation_id)] local_id: UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] multibody_info: &[MultibodyInfo], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] + joint_constraints: &mut [MultibodyJointConstraint], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] joint_constraint_columns: &[f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] dof_state: &mut [f32], + #[spirv(uniform, descriptor_set = 0, binding = 4)] use_bias: &u32, + #[spirv(uniform, descriptor_set = 0, binding = 5)] batch_ids: &BatchIndices, + #[spirv(workgroup)] dof_v: &mut [f32; MAX_MB_DOFS as usize], +) { + let batch_id = workgroup_id.y; + let mb_idx = workgroup_id.x; + let lane = local_id.x; + let num_mb = batch_ids.multibodies_len; + if mb_idx >= num_mb { + return; + } + + let mb_start = batch_ids.mb_start(batch_id); + let mb = multibody_info.read(mb_start + mb_idx as usize); + let ndofs = mb.ndofs; + // Uniform per workgroup: every lane of this group returns together. + if ndofs == 0 || mb.max_constraints == 0 { + return; + } + let use_bias = *use_bias != 0; + + let v_base = batch_ids.dof_start(batch_id) + mb.first_dof as usize; + let dofs_stride = batch_ids.dof_batch_capacity as usize; + let jcons_base = + batch_ids.mb_joint_constraints_start(batch_id) + mb.first_constraint as usize; + let jcol_base = batch_ids.mb_joint_constraint_columns_start(batch_id) + + (mb.first_constraint as usize) * dofs_stride; + + if lane < ndofs { + dof_v[lane as usize] = dof_state.read(v_base + lane as usize); + } + workgroup_memory_barrier_with_group_sync(); + + for s in 0..mb.max_constraints { + let cons = joint_constraints.read(jcons_base + s as usize); + if cons.kind == 0 { + // Uniform skip: all lanes take it together (barrier-safe). + continue; + } + + let rhs = if use_bias { cons.rhs } else { cons.rhs_wo_bias }; + let v_d = dof_v[cons.dof_id as usize]; + let rhs_total = v_d + rhs; + let raw_imp = cons.impulse + cons.inv_lhs * (rhs_total - cons.cfm_gain * cons.impulse); + let mut new_imp = raw_imp; + if new_imp < cons.impulse_lo { + new_imp = cons.impulse_lo; + } + if new_imp > cons.impulse_hi { + new_imp = cons.impulse_hi; + } + let delta = new_imp - cons.impulse; + + if lane == 0 { + let mut cons = cons; + cons.impulse = new_imp; + joint_constraints.write(jcons_base + s as usize, cons); + } + + workgroup_memory_barrier_with_group_sync(); + if lane < ndofs { + let col = joint_constraint_columns + .read(jcol_base + (s as usize) * dofs_stride + lane as usize); + dof_v[lane as usize] -= delta * col; + } + workgroup_memory_barrier_with_group_sync(); + } + + if lane < ndofs { + dof_state.write(v_base + lane as usize, dof_v[lane as usize]); + } +} + +/// Fills the per-multibody Delassus block `D[s][j] = ∂a[j]/∂impulse[s]` (row +/// `s` = the effect of constraint `s`, laid out row-contiguously so the solve +/// kernel's per-iteration row update reads coalesced). Runs right after +/// `gpu_mb_finalize_contact_constraints` (it consumes the M⁻¹Jᵀ columns). +/// +/// One 64-lane workgroup per (multibody, batch). +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn gpu_mb_build_contact_delassus( + #[spirv(workgroup_id)] workgroup_id: UVec3, + #[spirv(local_invocation_id)] local_id: UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] multibody_info: &[MultibodyInfo], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] + contact_constraints: &[MultibodyContactConstraint], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] contact_constraint_jacs: &[f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] contact_constraint_columns: &[f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] delassus: &mut [f32], + #[spirv(uniform, descriptor_set = 0, binding = 5)] batch_ids: &BatchIndices, +) { + const MAXC: u32 = MAX_MB_CONTACT_CONSTRAINTS_PER_MB; + let batch_id = workgroup_id.y; + let mb_idx = workgroup_id.x; + let lane = local_id.x; + let num_mb = batch_ids.multibodies_len; + if mb_idx >= num_mb { + return; + } + + let mb_start = batch_ids.mb_start(batch_id); + let mb = multibody_info.read(mb_start + mb_idx as usize); + let ndofs = mb.ndofs; + let count = mb.contact_constraint_count; + if ndofs == 0 || count == 0 { + return; + } + + let cons_base = batch_ids.mb_contact_constraints_start(batch_id) + + (mb_idx as usize) * (MAXC as usize); + let dofs_stride = batch_ids.dof_batch_capacity as usize; + let col_base = batch_ids.mb_contact_constraint_columns_start(batch_id) + + (mb_idx as usize) * (MAXC as usize) * dofs_stride; + let d_base = ((batch_id * batch_ids.multibodies_batch_capacity + mb_idx) as usize) + * (MAXC as usize) + * (MAXC as usize); + + // Pair `p = s · count + j`: consecutive lanes share the source row `s` + // and vary the target `j`, so the column reads of `s` broadcast and the + // `D` writes coalesce. + let num_pairs = count * count; + for p in StepRng::new(lane..num_pairs, LANES) { + let s = p / count; + let j = p % count; + + // Multibody coupling: jac_j · (M⁻¹ jac_sᵀ). + let jac_j_off = col_base + (j as usize) * dofs_stride; + let col_s_off = col_base + (s as usize) * dofs_stride; + let mut v = 0.0f32; + for i in 0..ndofs { + let jj = contact_constraint_jacs.read(jac_j_off + i as usize); + let cs = contact_constraint_columns.read(col_s_off + i as usize); + v += jj * cs; + } + + // Free-body coupling (impulse at `s` moves the shared free body, + // which feeds `a[j]`'s free-side term). Zero for self-contacts and + // static free bodies. + let cons_s = contact_constraints.read(cons_base + s as usize); + let cons_j = contact_constraints.read(cons_base + j as usize); + if cons_s.free_body_id != u32::MAX && cons_s.free_body_id == cons_j.free_body_id { + v += cons_s.free_body_im * cons_j.lin_jac.dot(cons_s.lin_jac) + + gdot(cons_j.ang_jac, cons_s.ii_ang_jac); + } + + delassus.write(d_base + (s * MAXC + j) as usize, v); + } +} + +/// Constraint-space contact sweep: tracks `a[s] = J_s · u` incrementally in +/// workgroup memory using the precomputed Delassus rows, so each PGS +/// iteration is a couple of shared-memory scalars plus one lane-parallel row +/// update. +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn gpu_mb_solve_contacts_delassus( + #[spirv(workgroup_id)] workgroup_id: UVec3, + #[spirv(local_invocation_id)] local_id: UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] multibody_info: &[MultibodyInfo], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] + contact_constraints: &mut [MultibodyContactConstraint], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] contact_constraint_jacs: &[f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] contact_constraint_columns: &[f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] delassus: &[f32], + #[spirv(uniform, descriptor_set = 0, binding = 5)] use_bias: &u32, + #[spirv(uniform, descriptor_set = 0, binding = 6)] batch_ids: &BatchIndices, + #[spirv(storage_buffer, descriptor_set = 1, binding = 0)] dof_state: &mut [f32], + #[spirv(storage_buffer, descriptor_set = 1, binding = 1)] solver_vels: &mut [Velocity], + #[spirv(workgroup)] dof_v: &mut [f32; MAX_MB_DOFS as usize], + #[spirv(workgroup)] a_shared: &mut [f32; MAX_MB_CONTACT_CONSTRAINTS_PER_MB as usize], + #[spirv(workgroup)] imp_shared: &mut [f32; MAX_MB_CONTACT_CONSTRAINTS_PER_MB as usize], + #[spirv(workgroup)] rhs_shared: &mut [f32; MAX_MB_CONTACT_CONSTRAINTS_PER_MB as usize], + #[spirv(workgroup)] inv_lhs_shared: &mut [f32; MAX_MB_CONTACT_CONSTRAINTS_PER_MB as usize], + #[spirv(workgroup)] cfm_shared: &mut [f32; MAX_MB_CONTACT_CONSTRAINTS_PER_MB as usize], + #[spirv(workgroup)] friction_shared: &mut [f32; MAX_MB_CONTACT_CONSTRAINTS_PER_MB as usize], + #[spirv(workgroup)] meta_shared: &mut [u32; MAX_MB_CONTACT_CONSTRAINTS_PER_MB as usize], +) { + const MAXC: u32 = MAX_MB_CONTACT_CONSTRAINTS_PER_MB; + let batch_id = workgroup_id.y; + let mb_idx = workgroup_id.x; + let lane = local_id.x; + let num_mb = batch_ids.multibodies_len; + if mb_idx >= num_mb { + return; + } + + let mb_start = batch_ids.mb_start(batch_id); + let mb = multibody_info.read(mb_start + mb_idx as usize); + let ndofs = mb.ndofs; + let count = mb.contact_constraint_count; + // Uniform per workgroup: every lane of this group returns together. + if ndofs == 0 || count == 0 { + return; + } + let use_bias = *use_bias != 0; + + let v_base = batch_ids.dof_start(batch_id) + mb.first_dof as usize; + let colliders_start = batch_ids.coll_start(batch_id); + let cons_base = batch_ids.mb_contact_constraints_start(batch_id) + + (mb_idx as usize) * (MAXC as usize); + let dofs_stride = batch_ids.dof_batch_capacity as usize; + let col_base = batch_ids.mb_contact_constraint_columns_start(batch_id) + + (mb_idx as usize) * (MAXC as usize) * dofs_stride; + let d_base = ((batch_id * batch_ids.multibodies_batch_capacity + mb_idx) as usize) + * (MAXC as usize) + * (MAXC as usize); + + if lane < ndofs { + dof_v[lane as usize] = dof_state.read(v_base + lane as usize); + } + + // Preload the per-constraint solve scalars into shared SoA arrays so the + // serial recurrence below never touches storage on its critical path. + // `meta` packs the kind, the paired normal slot, and whether the + // free-body side needs the fire-and-forget storage velocity update. + for s in StepRng::new(lane..count, LANES) { + let cons = contact_constraints.read(cons_base + s as usize); + imp_shared[s as usize] = cons.impulse; + rhs_shared[s as usize] = if use_bias { cons.rhs } else { cons.rhs_wo_bias }; + inv_lhs_shared[s as usize] = cons.inv_lhs; + cfm_shared[s as usize] = cons.cfm_factor; + friction_shared[s as usize] = cons.friction_coeff; + let is_self = cons.free_body_id == u32::MAX; + let free_active = !is_self + && (cons.free_body_im != 0.0 || gdot(cons.ii_ang_jac, cons.ii_ang_jac) != 0.0); + meta_shared[s as usize] = (cons.kind & 0xff) + | ((cons.normal_constraint_slot & 0xffff) << 8) + | (if free_active { 1 << 24 } else { 0 }); + } + workgroup_memory_barrier_with_group_sync(); + + // Fresh `a[s] = J_s · u` under the current (post-joint-sweep, post- + // warmstart) velocities. + for s in StepRng::new(lane..count, LANES) { + let jac_off = col_base + (s as usize) * dofs_stride; + let mut dot = 0.0f32; + for i in 0..ndofs { + dot += contact_constraint_jacs.read(jac_off + i as usize) * dof_v[i as usize]; + } + let cons = contact_constraints.read(cons_base + s as usize); + if cons.free_body_id != u32::MAX { + let free = solver_vels.read(colliders_start + cons.free_body_id as usize); + dot += cons.lin_jac.dot(free.linear) + gdot(cons.ang_jac, free.angular); + } + a_shared[s as usize] = dot; + } + workgroup_memory_barrier_with_group_sync(); + + for s in 0..count { + let meta = meta_shared[s as usize]; + let kind = meta & 0xff; + let normal_slot = (meta >> 8) & 0xffff; + let free_active = (meta >> 24) != 0; + + let impulse = imp_shared[s as usize]; + let rhs_total = a_shared[s as usize] + rhs_shared[s as usize]; + let raw_imp = cfm_shared[s as usize] * (impulse - inv_lhs_shared[s as usize] * rhs_total); + + let new_imp = if kind == MB_CONTACT_KIND_TANGENT { + let limit = friction_shared[s as usize] * imp_shared[normal_slot as usize]; + if raw_imp > limit { + limit + } else if raw_imp < -limit { + -limit + } else { + raw_imp + } + } else if raw_imp < 0.0 { + 0.0 + } else { + raw_imp + }; + let delta = new_imp - impulse; + + if delta != 0.0 { + if lane == 0 { + imp_shared[s as usize] = new_imp; + + if free_active { + let cons = contact_constraints.read(cons_base + s as usize); + let mut free = + solver_vels.read(colliders_start + cons.free_body_id as usize); + free.linear += cons.lin_jac * (cons.free_body_im * delta); + free.angular += cons.ii_ang_jac * delta; + solver_vels.write(colliders_start + cons.free_body_id as usize, free); + } + } + // Lane-parallel Delassus row update (row `s` is contiguous), plus + // the off-path dof update (each lane owns its DOF). + let d_row = d_base + (s * MAXC) as usize; + for j in StepRng::new(lane..count, LANES) { + a_shared[j as usize] += delta * delassus.read(d_row + j as usize); + } + if lane < ndofs { + let col = contact_constraint_columns + .read(col_base + (s as usize) * dofs_stride + lane as usize); + dof_v[lane as usize] += delta * col; + } + workgroup_memory_barrier_with_group_sync(); + } + } + + // Writeback. + if lane < ndofs { + dof_state.write(v_base + lane as usize, dof_v[lane as usize]); + } + for s in StepRng::new(lane..count, LANES) { + let mut cons = contact_constraints.read(cons_base + s as usize); + cons.impulse = imp_shared[s as usize]; + contact_constraints.write(cons_base + s as usize, cons); + } +} From 9e186b4c3295ffc40e84db2e3dfc77c1eec6b344 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Fri, 31 Jul 2026 14:48:55 +0200 Subject: [PATCH 32/39] perf: per-substep joint-constraint refresh in explicit-coriolis mode --- .../dynamics/multibody/multibody_solver.rs | 67 +++--- src_rbd/dynamics/solver.rs | 1 + .../dynamics/multibody/joint_constraints.rs | 195 +++++++++++++----- .../dynamics/multibody/solve_constraints.rs | 14 +- src_rbd_shaders/dynamics/multibody/types.rs | 12 ++ 5 files changed, 212 insertions(+), 77 deletions(-) diff --git a/src_rbd/dynamics/multibody/multibody_solver.rs b/src_rbd/dynamics/multibody/multibody_solver.rs index 24cafb8..8525e48 100644 --- a/src_rbd/dynamics/multibody/multibody_solver.rs +++ b/src_rbd/dynamics/multibody/multibody_solver.rs @@ -9,7 +9,7 @@ use crate::shaders::dynamics::{ GpuMbFinalizeContactConstraints, GpuMbGravityAndLu, GpuMbGravityAndLuT8, GpuMbGravityAndLuT16, GpuMbGravityAndLuT32, GpuMbInitContactConstraints, GpuMbInitJointConstraints, GpuMbIntegrate, GpuMbIntegrateVelocities, - GpuMbRemoveImpulseJointConstraintBias, + GpuMbRefreshJointConstraints, GpuMbRemoveImpulseJointConstraintBias, GpuMbResetContactWarmstart, GpuMbStashContactsLen, GpuMbWarmstartContactConstraints, GpuMbSolveConstraints, GpuMbSolveContactsDelassus, GpuMbSolveImpulseJointConstraints, GpuMbSolveJoints, @@ -35,6 +35,10 @@ pub struct GpuMultibodySolver { compute_dynamics_pre: GpuMbComputeDynamicsPre, compute_dynamics_without_coriolis_pre: GpuMbComputeDynamicsWithoutCoriolisPre, init_joint_with_bias: GpuMbInitJointConstraints, + /// Explicit-coriolis fast path: per-substep refresh of the joint rhs / + /// limit activity (the columns and `inv_lhs` are per-step constants + /// there, so the full build + back-solves run once per step). + refresh_joint_constraints: GpuMbRefreshJointConstraints, init_contact_constraints: GpuMbInitContactConstraints, finalize_contact_constraints: GpuMbFinalizeContactConstraints, /// Fused joint+contact PGS sweep (one workgroup per multibody, shared- @@ -202,42 +206,40 @@ impl GpuMultibodySolver { mut timestamps: Option<&mut khal::backend::GpuTimestamps>, mb: &mut GpuMultibodySet, args: &mut MultibodySolverArgs<'_>, + first_substep: bool, ) -> Result<(), GpuBackendError> { use khal::backend::Encoder; if mb.is_empty() { return Ok(()); } - if mb.has_joint_constraints { - let mut pass = encoder.begin_pass("[RBD] mbb/init-joint", timestamps.as_deref_mut()); - // One 64-lane workgroup per multibody: lane 0 emits the constraint - // metadata serially (cheap), then the per-constraint M⁻¹-column LU - // back-solves run one-per-lane instead of sequentially. - let init_joint_dispatch = [mb.multibodies_per_batch * MB_LU_LANES, mb.num_batches, 1]; - self.init_joint_with_bias.call( + // With implicit coriolis, the mass matrix / LU / body jacobians are + // recomputed every substep, so the joint + contact constraints (whose + // M⁻¹Jᵀ columns depend on them) must be rebuilt every substep too. In + // the explicit mode every column-derived quantity is a per-step + // constant: the full build runs ONCE per step (see + // `build_contact_constraints`) and each substep only refreshes the + // joint rhs / limit activity / accumulated impulse from the + // integrated joint positions (a no-op on the first substep — the + // once-per-step build just wrote those exact values). + if mb.implicit_coriolis { + self.build_contact_constraints(encoder, timestamps.as_deref_mut(), mb, args)?; + } else if mb.has_joint_constraints && !first_substep { + let mut pass = + encoder.begin_pass("[RBD] mbb/refresh-joint", timestamps.as_deref_mut()); + let dispatch = [mb.multibodies_per_batch * MB_LU_LANES, mb.num_batches, 1]; + self.refresh_joint_constraints.call( &mut pass, - init_joint_dispatch, + dispatch, &mb.multibody_info, &mb.links_static, &mb.links_workspace, - &mb.mass_matrices, - &mb.lu_pivots, &mut mb.joint_constraints, - &mut mb.joint_constraint_columns, &mb.constraint_softness, args.batch_indices, )?; } - // With implicit coriolis, the mass matrix / LU / body jacobians are - // recomputed every substep, so the contact constraints (whose M⁻¹Jᵀ - // columns depend on them) must be rebuilt every substep too. In the - // explicit mode every input is a per-step constant, so the pipeline - // builds them ONCE per step instead (see `build_contact_constraints`). - if mb.implicit_coriolis { - self.build_contact_constraints(encoder, timestamps.as_deref_mut(), mb, args)?; - } - // Warmstart: re-apply the accumulated contact impulse to dof_state (and // the free-body solver velocities) so the contact starts "warm" each // substep — mirrors rapier's per-substep `contact_constraints.warmstart` @@ -287,8 +289,27 @@ impl GpuMultibodySolver { return Ok(()); } - // One 64-lane workgroup per multibody: the uniform emission walk runs - // redundantly on every lane, the per-DOF `Jᵀ`-row fills one-per-lane. + // Joint limit/motor constraints: one 64-lane workgroup per multibody + // (lane 0 emits the metadata serially). + if mb.has_joint_constraints { + let mut pass = encoder.begin_pass("[RBD] mbb/init-joint", timestamps.as_deref_mut()); + let init_joint_dispatch = [mb.multibodies_per_batch * MB_LU_LANES, mb.num_batches, 1]; + self.init_joint_with_bias.call( + &mut pass, + init_joint_dispatch, + &mb.multibody_info, + &mb.links_static, + &mb.links_workspace, + &mb.mass_matrices, + &mb.lu_pivots, + &mut mb.joint_constraints, + &mut mb.joint_constraint_columns, + &mb.constraint_softness, + args.batch_indices, + )?; + } + + // One 64-lane workgroup per multibody. { let mut pass = encoder.begin_pass("[RBD] mbb/init-contact", timestamps.as_deref_mut()); diff --git a/src_rbd/dynamics/solver.rs b/src_rbd/dynamics/solver.rs index 05f5d2b..d2b1e1b 100644 --- a/src_rbd/dynamics/solver.rs +++ b/src_rbd/dynamics/solver.rs @@ -405,6 +405,7 @@ impl GpuSolver { timestamps.as_deref_mut(), state, &mut mb_args, + substep_id == 0, )?; } } diff --git a/src_rbd_shaders/dynamics/multibody/joint_constraints.rs b/src_rbd_shaders/dynamics/multibody/joint_constraints.rs index 4a69c94..41938a0 100644 --- a/src_rbd_shaders/dynamics/multibody/joint_constraints.rs +++ b/src_rbd_shaders/dynamics/multibody/joint_constraints.rs @@ -17,7 +17,9 @@ use crate::utils::linalg::{MatSlice, lu_solve_in_place}; use crate::{DIM, MAX_FLT}; use super::types::{ - MultibodyInfo, MultibodyJointConstraint, MultibodyLinkStatic, MultibodyLinkWorkspace, + MB_JOINT_KIND_INACTIVE, MB_JOINT_KIND_LIMIT, MB_JOINT_KIND_LIMIT_INACTIVE, + MB_JOINT_KIND_MOTOR, MultibodyInfo, MultibodyJointConstraint, MultibodyLinkStatic, + MultibodyLinkWorkspace, }; /// Compute joint motor parameters mirroring rapier's `JointMotor::motor_params`. @@ -112,11 +114,10 @@ fn emit_joint_constraints( let has_limits = (limit_axes & (1 << axis)) != 0; let limit_min = stat.data.limits[axis as usize].min; let limit_max = stat.data.limits[axis as usize].max; - emit_motor_constraint( - joint_constraints, - cons_base, - slot, + let cons = build_motor_constraint( abs_dof, + k, + axis, curr_pos, inv_dt, dt, @@ -125,14 +126,14 @@ fn emit_joint_constraints( limit_min, limit_max, ); + joint_constraints.write(cons_base + slot as usize, cons); slot += 1; } if (limit_axes & (1 << axis)) != 0 { - emit_limit_constraint( - joint_constraints, - cons_base, - slot, + let cons = build_limit_constraint( abs_dof, + k, + axis, curr_pos, [ stat.data.limits[axis as usize].min, @@ -141,6 +142,7 @@ fn emit_joint_constraints( joint_erp_inv_dt, joint_cfm_coeff, ); + joint_constraints.write(cons_base + slot as usize, cons); slot += 1; } curr_free_dof += 1; @@ -155,11 +157,10 @@ fn emit_joint_constraints( let curr_pos = ws.coords.read(axis as usize); if (limit_axes & (1 << axis)) != 0 { - emit_limit_constraint( - joint_constraints, - cons_base, - slot, + let cons = build_limit_constraint( abs_dof, + k, + axis, curr_pos, [ stat.data.limits[axis as usize].min, @@ -168,17 +169,17 @@ fn emit_joint_constraints( joint_erp_inv_dt, joint_cfm_coeff, ); + joint_constraints.write(cons_base + slot as usize, cons); slot += 1; } if (motor_axes & (1 << axis)) != 0 { let has_limits = (limit_axes & (1 << axis)) != 0; let limit_min = stat.data.limits[axis as usize].min; let limit_max = stat.data.limits[axis as usize].max; - emit_motor_constraint( - joint_constraints, - cons_base, - slot, + let cons = build_motor_constraint( abs_dof, + k, + axis, curr_pos, inv_dt, dt, @@ -187,6 +188,7 @@ fn emit_joint_constraints( limit_min, limit_max, ); + joint_constraints.write(cons_base + slot as usize, cons); slot += 1; } curr_free_dof += 1; @@ -194,6 +196,107 @@ fn emit_joint_constraints( } } +/// Per-substep joint-constraint refresh — the explicit-coriolis fast path. +/// +/// With explicit coriolis the mass-matrix LU (and therefore every slot's M⁻¹ +/// column, `inv_lhs` and folded `cfm_gain`) is a per-step constant: only the +/// rhs (from the integrated joint positions), the limit activity and the +/// accumulated impulse change per substep. This kernel recomputes exactly +/// those from the slot's stashed (link, axis) — the full emission walk and +/// the back-solves run once per step instead of once per substep. +/// +/// One 64-lane workgroup per (multibody, batch); lanes stride the slots. +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn gpu_mb_refresh_joint_constraints( + #[spirv(workgroup_id)] workgroup_id: UVec3, + #[spirv(local_invocation_id)] local_id: UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] multibody_info: &[MultibodyInfo], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] + links_static: &[MultibodyLinkStatic], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] + links_workspace: &[MultibodyLinkWorkspace], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] + joint_constraints: &mut [MultibodyJointConstraint], + #[spirv(uniform, descriptor_set = 0, binding = 4)] softness: &ConstraintSoftness, + #[spirv(uniform, descriptor_set = 0, binding = 5)] batch_ids: &BatchIndices, +) { + const LANES: u32 = 64; + let batch_id = workgroup_id.y; + let mb_idx = workgroup_id.x; + let lane = local_id.x; + let num_mb = batch_ids.multibodies_len; + if mb_idx >= num_mb { + return; + } + + let mb = batch_ids + .mb_batch(batch_id, multibody_info) + .read(mb_idx as usize); + if mb.ndofs == 0 || mb.max_constraints == 0 { + return; + } + let cons_base = batch_ids.mb_joint_constraints_start(batch_id) + mb.first_constraint as usize; + + let stat_slice = batch_ids + .mb_links_batch(batch_id, links_static) + .offset(mb.first_link as usize); + let ws_slice = batch_ids + .mb_links_batch(batch_id, links_workspace) + .offset(mb.first_link as usize); + + let dt = softness.dt; + let inv_dt = if dt != 0.0 { 1.0 / dt } else { 0.0 }; + + for s in StepRng::new(lane..mb.max_constraints, LANES) { + let old = joint_constraints.read(cons_base + s as usize); + if old.kind == MB_JOINT_KIND_INACTIVE { + continue; + } + let link_id = old._kind_extra & 0xffff; + let axis = old._kind_extra >> 16; + let stat = &stat_slice[link_id as usize]; + let ws = &ws_slice[link_id as usize]; + let curr_pos = ws.coords.read(axis as usize); + + // Rebuild the per-substep fields with the SAME formulas as the full + // emission, then graft the per-step constants (column-derived + // `inv_lhs` and folded `cfm_gain`) from the existing slot. + let mut fresh = if old.kind == MB_JOINT_KIND_MOTOR { + let locked = stat.data.locked_axes; + let has_limits = (stat.data.limit_axes & !locked & (1 << axis)) != 0; + build_motor_constraint( + old.dof_id, + link_id, + axis, + curr_pos, + inv_dt, + dt, + &stat.data.motors[axis as usize], + has_limits, + stat.data.limits[axis as usize].min, + stat.data.limits[axis as usize].max, + ) + } else { + build_limit_constraint( + old.dof_id, + link_id, + axis, + curr_pos, + [ + stat.data.limits[axis as usize].min, + stat.data.limits[axis as usize].max, + ], + softness.joint_erp_inv_dt, + softness.joint_cfm_coeff, + ) + }; + fresh.inv_lhs = old.inv_lhs; + fresh.cfm_gain = old.cfm_gain; + joint_constraints.write(cons_base + s as usize, fresh); + } +} + /// Solve `M · column = e_{dof_id}` (writes the M⁻¹ column) and return the raw /// `lhs = column[dof_id]` for J = e_{dof_id}. #[inline] @@ -229,34 +332,24 @@ fn inv(x: f32) -> f32 { if x != 0.0 { 1.0 / x } else { 0.0 } } -/// Initialize a single limit constraint slot. Mirrors rapier's -/// `unit_joint_limit_constraint`. -/// -/// Emits METADATA ONLY: `inv_lhs` is left 0 and `cfm_gain` holds the -/// pre-fold gain (0 for limits); the lane-parallel finalize stage of -/// `gpu_mb_init_joint_constraints` back-solves the M⁻¹ column and applies -/// rapier's `finalize_generic_constraints` fold. +/// Initialize a single limit constraint slot. #[inline] -fn emit_limit_constraint( - joint_constraints: &mut [MultibodyJointConstraint], - cons_base: usize, - slot: u32, +#[allow(clippy::too_many_arguments)] +fn build_limit_constraint( dof_id: u32, + link_id: u32, + axis: u32, curr_pos: f32, limits: [f32; 2], erp_inv_dt: f32, cfm_coeff: f32, -) { +) -> MultibodyJointConstraint { // rapier (`limit_*` builder): erp_inv_dt = joint.softness.erp_inv_dt(dt), // cfm_coeff = joint.softness.cfm_coeff(dt), cfm_gain = 0 — configurable via // `joint_natural_frequency` / `joint_damping_ratio` (defaults make this // near-rigid, matching the old hardcoded `1/dt`). let min_enabled = curr_pos < limits[0]; let max_enabled = limits[1] < curr_pos; - // No limit is active, skip the constraint for the current substep. - if !min_enabled && !max_enabled { - return; - } let lo_excess = (limits[0] - curr_pos).max(0.0); let hi_excess = (curr_pos - limits[1]).max(0.0); let rhs_bias = (hi_excess - lo_excess) * erp_inv_dt; @@ -265,10 +358,18 @@ fn emit_limit_constraint( let max_neg_impulse = if min_enabled { -MAX_FLT } else { 0.0 }; let max_pos_impulse = if max_enabled { MAX_FLT } else { 0.0 }; - let cons = MultibodyJointConstraint { + let kind = if min_enabled || max_enabled { + MB_JOINT_KIND_LIMIT + } else { + // Inactive this substep: the solve skips it, the finalize stage still + // back-solves its column for later refreshes. + MB_JOINT_KIND_LIMIT_INACTIVE + }; + + MultibodyJointConstraint { dof_id, - kind: 1, - _kind_extra: 0, + kind, + _kind_extra: link_id | (axis << 16), _pad0: 0, rhs: rhs_wo_bias + rhs_bias, rhs_wo_bias, @@ -279,17 +380,16 @@ fn emit_limit_constraint( cfm_coeff, // This will be calculated in the finalize (orthogonalization) step. cfm_gain: 0.0, - }; - joint_constraints.write(cons_base + slot as usize, cons); + } } /// Initialize a single motor constraint slot.. #[inline] -fn emit_motor_constraint( - joint_constraints: &mut [MultibodyJointConstraint], - cons_base: usize, - slot: u32, +#[allow(clippy::too_many_arguments)] +fn build_motor_constraint( dof_id: u32, + link_id: u32, + axis: u32, curr_pos: f32, inv_dt: f32, dt: f32, @@ -297,7 +397,7 @@ fn emit_motor_constraint( has_limits: bool, limit_min: f32, limit_max: f32, -) { +) -> MultibodyJointConstraint { let (erp_inv_dt, cfm_coeff, cfm_gain, _, max_impulse) = motor_params(motor, dt); let mut rhs_wo_bias = 0.0f32; @@ -318,10 +418,10 @@ fn emit_motor_constraint( } rhs_wo_bias += -target_vel; - let cons = MultibodyJointConstraint { + MultibodyJointConstraint { dof_id, - kind: 2, - _kind_extra: 0, + kind: MB_JOINT_KIND_MOTOR, + _kind_extra: link_id | (axis << 16), _pad0: 0, rhs: rhs_wo_bias, rhs_wo_bias, @@ -331,8 +431,7 @@ fn emit_motor_constraint( impulse_hi: max_impulse, cfm_coeff, cfm_gain, - }; - joint_constraints.write(cons_base + slot as usize, cons); + } } /// Initialize the multibody's joint-limit / joint-motor unit constraints. diff --git a/src_rbd_shaders/dynamics/multibody/solve_constraints.rs b/src_rbd_shaders/dynamics/multibody/solve_constraints.rs index 4f8af6a..d47ac96 100644 --- a/src_rbd_shaders/dynamics/multibody/solve_constraints.rs +++ b/src_rbd_shaders/dynamics/multibody/solve_constraints.rs @@ -13,8 +13,8 @@ use crate::utils::BatchIndices; use crate::utils::linalg::MAX_MB_DOFS; use super::types::{ - MAX_MB_CONTACT_CONSTRAINTS_PER_MB, MB_CONTACT_KIND_TANGENT, MultibodyContactConstraint, - MultibodyInfo, MultibodyJointConstraint, + MAX_MB_CONTACT_CONSTRAINTS_PER_MB, MB_CONTACT_KIND_TANGENT, MB_JOINT_KIND_LIMIT, + MB_JOINT_KIND_MOTOR, MultibodyContactConstraint, MultibodyInfo, MultibodyJointConstraint, }; const LANES: u32 = 64; @@ -96,8 +96,9 @@ pub fn gpu_mb_solve_constraints( // Joint limits/motors for s in 0..mb.max_constraints { let cons = joint_constraints.read(jcons_base + s as usize); - if cons.kind == 0 { - // Uniform skip: all lanes take it together (barrier-safe). + if cons.kind != MB_JOINT_KIND_LIMIT && cons.kind != MB_JOINT_KIND_MOTOR { + // Unused slot or inactive limit. Uniform skip: all lanes take it + // together (barrier-safe). continue; } @@ -265,8 +266,9 @@ pub fn gpu_mb_solve_joints( for s in 0..mb.max_constraints { let cons = joint_constraints.read(jcons_base + s as usize); - if cons.kind == 0 { - // Uniform skip: all lanes take it together (barrier-safe). + if cons.kind != MB_JOINT_KIND_LIMIT && cons.kind != MB_JOINT_KIND_MOTOR { + // Unused slot or inactive limit. Uniform skip: all lanes take it + // together (barrier-safe). continue; } diff --git a/src_rbd_shaders/dynamics/multibody/types.rs b/src_rbd_shaders/dynamics/multibody/types.rs index 235eaea..65149c6 100644 --- a/src_rbd_shaders/dynamics/multibody/types.rs +++ b/src_rbd_shaders/dynamics/multibody/types.rs @@ -47,6 +47,18 @@ pub const MB_CONTACT_KIND_NORMAL: u32 = 1; /// `normal_constraint_slot` (relative to the multibody's `cons_base`). pub const MB_CONTACT_KIND_TANGENT: u32 = 2; +/// Joint-constraint `kind`: unused slot. +pub const MB_JOINT_KIND_INACTIVE: u32 = 0; +/// Joint-constraint `kind`: active limit. +pub const MB_JOINT_KIND_LIMIT: u32 = 1; +/// Joint-constraint `kind`: active motor. +pub const MB_JOINT_KIND_MOTOR: u32 = 2; +/// Joint-constraint `kind`: limit slot that is INACTIVE this substep. The +/// solve skips it, but the slot keeps a valid M⁻¹ column / `inv_lhs` / +/// `cfm_gain` so the per-substep refresh can flip it active without a +/// back-solve (see `gpu_mb_refresh_joint_constraints`). +pub const MB_JOINT_KIND_LIMIT_INACTIVE: u32 = 3; + /// Sentinel marking a link with no parent (the root). pub const MULTIBODY_ROOT: u32 = u32::MAX; From 0c5fa845515aec8d617294c3836bd93f6d11073d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Sat, 1 Aug 2026 10:57:53 +0200 Subject: [PATCH 33/39] perf: merge small-scene submits + skip the provably-empty first warmstart --- .../dynamics/multibody/multibody_solver.rs | 6 +-- src_rbd/pipeline/rbd_step.rs | 54 ++++++++++--------- 2 files changed, 30 insertions(+), 30 deletions(-) diff --git a/src_rbd/dynamics/multibody/multibody_solver.rs b/src_rbd/dynamics/multibody/multibody_solver.rs index 8525e48..98d2ead 100644 --- a/src_rbd/dynamics/multibody/multibody_solver.rs +++ b/src_rbd/dynamics/multibody/multibody_solver.rs @@ -242,11 +242,9 @@ impl GpuMultibodySolver { // Warmstart: re-apply the accumulated contact impulse to dof_state (and // the free-body solver velocities) so the contact starts "warm" each - // substep — mirrors rapier's per-substep `contact_constraints.warmstart` - // and matches what the rigid-body solver does for free contacts. On the - // first substep the impulse was just reset to 0, so this is a no-op. + // substep. // One 64-lane workgroup per multibody (one DOF per lane). - { + if !first_substep { let mut pass = encoder.begin_pass("[RBD] mbb/warmstart-contact", timestamps.as_deref_mut()); let warmstart_dispatch = diff --git a/src_rbd/pipeline/rbd_step.rs b/src_rbd/pipeline/rbd_step.rs index a7ad103..a4a20a0 100644 --- a/src_rbd/pipeline/rbd_step.rs +++ b/src_rbd/pipeline/rbd_step.rs @@ -80,11 +80,12 @@ impl RbdPipeline { state.ensure_color_uniforms(backend, needed); } + let mut encoder = backend.begin_encoding(); + // Phase 0: Multibody once-per-visible-step setup (3D only for now). #[cfg(feature = "dim3")] { if !state.multibodies.is_empty() { - let mut encoder = backend.begin_encoding(); let mut args = crate::dynamics::MultibodySolverArgs { poses: &mut state.body_poses, collider_world_poses: &state.collider_world_poses, @@ -101,13 +102,11 @@ impl RbdPipeline { &mut state.multibodies, &mut args, )?; - backend.submit(encoder)?; } } // Phase 1: Update mass properties, build LBVH, and find collision pairs. { - let mut encoder = backend.begin_encoding(); let mut pass = encoder.begin_pass("[RBD] update-mprops", timestamps.as_deref_mut()); // Update mass properties — uses body world poses to compute the @@ -214,11 +213,28 @@ impl RbdPipeline { } } + let readback_enabled = state.capacities.solver_colors_resize_policy + != RbdResizePolicy::Fixed + || state.capacities.collisions_resize_policy != RbdResizePolicy::Fixed; + let est_pairs = if readback_enabled { + state.collision_pairs_len_cpu + } else { + state.collision_pairs_per_batch_cpu + }; + + // Choose the kernel depending on the expected pairs count. + // Small pairs with many environment benefit from the fused kernels. + let fused_color_sweeps = est_pairs <= 128; + + // In small scenes, submit less frequently. In big scenes submit more + // to overlap compute and encoding. + let merge_submits = fused_color_sweeps && state.num_batches <= 64; + // Phase 2a: Narrow phase. Split out from solver-prep + coloring // so its CPU encoding overlaps with Phase 1's GPU work and its // own GPU work overlaps with Phase 2b's CPU encoding. + let mut encoder = backend.begin_encoding(); { - let mut encoder = backend.begin_encoding(); let mut pass = encoder.begin_pass("[RBD] narrow-phase", timestamps.as_deref_mut()); self.narrow_phase.dispatch( @@ -243,31 +259,16 @@ impl RbdPipeline { )?; drop(pass); - backend.submit(encoder)?; + if !merge_submits { + backend.submit(encoder)?; + encoder = backend.begin_encoding(); + } } - // Colored-sweep strategy: with few constraints per batch, the - // `num_colors` dispatches per sweep (and their empty buckets) dominate - // — run each sweep as one dispatch with one workgroup per batch - // looping the colors internally. The gate is perf-only (the fused - // kernel is correct for any size, just serialized past ~64 lanes): - // use the lagging pair-count readback when auto-resize keeps it fresh, - // else the fixed capacity. - let readback_enabled = state.capacities.solver_colors_resize_policy - != RbdResizePolicy::Fixed - || state.capacities.collisions_resize_policy != RbdResizePolicy::Fixed; - let est_pairs = if readback_enabled { - state.collision_pairs_len_cpu - } else { - state.collision_pairs_per_batch_cpu - }; - let fused_color_sweeps = est_pairs <= 128; - // Phase 2b: solver-prep + warmstart + bounded coloring. Separate // submit from narrow-phase to enable CPU/GPU overlap with the // upcoming Phase 3 solver substep loop. { - let mut encoder = backend.begin_encoding(); let mut pass = encoder.begin_pass("[RBD] solver-prep", timestamps.as_deref_mut()); // Solver preparation - create args here to avoid borrow conflicts @@ -313,7 +314,6 @@ impl RbdPipeline { if state.rb_contacts_inert { stats.num_colors = state.max_colors + 1; drop(pass); - backend.submit(encoder)?; } else { // Warmstart @@ -408,7 +408,10 @@ impl RbdPipeline { stats.num_colors = num_colors; drop(pass); - backend.submit(encoder)?; + } + if !merge_submits { + backend.submit(encoder)?; + encoder = backend.begin_encoding(); } } @@ -465,7 +468,6 @@ impl RbdPipeline { }; { - let mut encoder = backend.begin_encoding(); #[cfg(feature = "dim3")] let mb = if state.multibodies.is_empty() { None From 226d3dde414310286e74e419472e2219b21ace94 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Sat, 1 Aug 2026 12:33:40 +0200 Subject: [PATCH 34/39] perf: serial dynamics tier for large batches of small robots --- src_rbd/dynamics/multibody/multibody_set.rs | 22 ++- .../dynamics/multibody/multibody_solver.rs | 8 +- .../multibody/compute_dynamics_pre.rs | 46 ++--- .../dynamics/multibody/gravity_and_lu.rs | 179 +++++++++++++++++- 4 files changed, 224 insertions(+), 31 deletions(-) diff --git a/src_rbd/dynamics/multibody/multibody_set.rs b/src_rbd/dynamics/multibody/multibody_set.rs index 82a3aad..c14fe06 100644 --- a/src_rbd/dynamics/multibody/multibody_set.rs +++ b/src_rbd/dynamics/multibody/multibody_set.rs @@ -170,10 +170,26 @@ impl GpuMultibodySet { [self.num_active_multibodies * self.num_batches, 1, 1] } - /// Lanes per multibody for the packed per-multibody workgroup kernels — - /// mirrored into `BatchIndices::mb_pack_lanes`. + /// Lanes per multibody for the packed per-multibody dynamics kernels + /// (`compute_dynamics_pre`, `gravity_and_lu`) — mirrored into + /// `BatchIndices::mb_pack_lanes`. + /// + /// `1` selects the SERIAL tier: one thread runs its + /// multibody's whole FK/CRBA/LU chain with no barriers at all, 64 + /// multibodies per workgroup with every lane busy. For small robots this + /// beats lane-parallelism — whose ~60-barrier dependency chain caps how + /// fast one multibody can finish — but ONLY once there are enough + /// multibodies for the thread count to hide the long serial chain's + /// latency (measured crossover between 1024 and 4096 on Apple M-series; + /// below that, spreading each robot across 8 lanes wins despite the + /// barriers). pub(crate) fn pack_lanes(&self) -> u32 { - self.max_ndofs.next_power_of_two().clamp(8, MB_LU_LANES) + let total_mb = self.num_active_multibodies * self.num_batches; + if self.max_ndofs <= 8 && total_mb >= 2048 { + 1 + } else { + self.max_ndofs.next_power_of_two().clamp(8, MB_LU_LANES) + } } /// Thread-count grid for the packed per-multibody workgroup kernels diff --git a/src_rbd/dynamics/multibody/multibody_solver.rs b/src_rbd/dynamics/multibody/multibody_solver.rs index 98d2ead..648abae 100644 --- a/src_rbd/dynamics/multibody/multibody_solver.rs +++ b/src_rbd/dynamics/multibody/multibody_solver.rs @@ -6,7 +6,7 @@ use crate::queries::GpuIndexedContact; use crate::shaders::dynamics::{ GpuMbBuildContactDelassus, GpuMbComputeDynamicsPre, GpuMbComputeDynamicsWithoutCoriolisPre, - GpuMbFinalizeContactConstraints, GpuMbGravityAndLu, GpuMbGravityAndLuT8, + GpuMbFinalizeContactConstraints, GpuMbGravityAndLu, GpuMbGravityAndLuT1, GpuMbGravityAndLuT8, GpuMbGravityAndLuT16, GpuMbGravityAndLuT32, GpuMbInitContactConstraints, GpuMbInitJointConstraints, GpuMbIntegrate, GpuMbIntegrateVelocities, GpuMbRefreshJointConstraints, GpuMbRemoveImpulseJointConstraintBias, @@ -25,10 +25,7 @@ use vortx::tensor::Tensor; #[derive(Shader)] pub struct GpuMultibodySolver { gravity_and_lu: GpuMbGravityAndLu, - /// Packed tiers of `gravity_and_lu` — `64/T` multibodies per workgroup - /// with a `T×T` shared tile each, selected by `max_ndofs`. The fallback - /// `gravity_and_lu` (one multibody per workgroup, 64×64 tile) only runs - /// for `max_ndofs > 32`. + gravity_and_lu_t1: GpuMbGravityAndLuT1, gravity_and_lu_t8: GpuMbGravityAndLuT8, gravity_and_lu_t16: GpuMbGravityAndLuT16, gravity_and_lu_t32: GpuMbGravityAndLuT32, @@ -670,6 +667,7 @@ impl GpuMultibodySolver { }; } match mb.pack_lanes() { + 1 => grav_lu!(gravity_and_lu_t1), 8 => grav_lu!(gravity_and_lu_t8), 16 => grav_lu!(gravity_and_lu_t16), 32 => grav_lu!(gravity_and_lu_t32), diff --git a/src_rbd_shaders/dynamics/multibody/compute_dynamics_pre.rs b/src_rbd_shaders/dynamics/multibody/compute_dynamics_pre.rs index c3152a2..3bb24c8 100644 --- a/src_rbd_shaders/dynamics/multibody/compute_dynamics_pre.rs +++ b/src_rbd_shaders/dynamics/multibody/compute_dynamics_pre.rs @@ -27,12 +27,14 @@ use crate::utils::{BatchIndices, Slice, SliceMut}; use crate::{ANG_DIM, AngVector, DIM, Pose, Vector, gcross_av}; use parry::math::VectorExt; -/// Packed slot decode shared by the two `pre` kernels: `64 / mb_pack_lanes` -/// multibodies per 64-lane workgroup, `(multibody, batch)` flattened into the -/// workgroup X dimension. Returns `(t, lane, batch_id, mb_idx, active_slot)`; -/// inactive slots get clamped indices (their loops all no-op on the zeroed -/// dummy `MultibodyInfo` the caller substitutes). `mb_pack_lanes` is -/// uniform-sourced so the decode keeps uniform control flow for barriers. +#[inline(always)] +fn sync_slots(t: u32) { + if t > 1 { + workgroup_memory_barrier_with_group_sync(); + } +} + +/// Returns `(t, lane, batch_id, mb_idx, active_slot)`. #[inline(always)] fn packed_decode(wg_id: UVec3, lid: UVec3, batch_ids: &BatchIndices) -> (u32, u32, u32, u32, bool) { let t = batch_ids.mb_pack_lanes; @@ -116,7 +118,7 @@ pub fn gpu_mb_compute_dynamics_pre( if active_slot && num_links > 0 && lane == 0 { forward_kinematics(&mb, &stat_slice, &mut poses_slice, &mut ws_slice, num_links); } - workgroup_memory_barrier_with_group_sync(); + sync_slots(t); // 2) Update body jacobians update_body_jacobians( @@ -135,7 +137,7 @@ pub fn gpu_mb_compute_dynamics_pre( if active_slot && num_links > 0 && lane == 0 { propagate_velocities(num_links, &stat_slice, &vel_slice, &mut ws_slice); } - workgroup_memory_barrier_with_group_sync(); + sync_slots(t); // 3) Mass matrix (with semi-implicit coriolis handling). let acc_augmented_mass = MatSlice::dense(mb_mm_base, ndofs, ndofs); @@ -145,7 +147,7 @@ pub fn gpu_mb_compute_dynamics_pre( let i_coriolis_dt_v = i_coriolis_dt_view.fixed_rows(0, DIM); let i_coriolis_dt_w = i_coriolis_dt_view.fixed_rows(DIM, ANG_DIM); - workgroup_memory_barrier_with_group_sync(); + sync_slots(t); for k in 0..batch_ids.mb_max_links { let loop_is_active = k < num_links; @@ -179,7 +181,7 @@ pub fn gpu_mb_compute_dynamics_pre( } // Uniform barrier so subsequent parent-coriolis reads see consistent // state — WebGPU forbids a barrier inside divergent control flow. - workgroup_memory_barrier_with_group_sync(); + sync_slots(t); let loop_is_active = k < num_links && inv_mass_x != 0.0; let coriolis_v_i = MatSlice::dense( @@ -335,7 +337,7 @@ pub fn gpu_mb_compute_dynamics_pre( } } - workgroup_memory_barrier_with_group_sync(); + sync_slots(t); if loop_is_active { if k != 0 { @@ -393,7 +395,7 @@ pub fn gpu_mb_compute_dynamics_pre( } } - workgroup_memory_barrier_with_group_sync(); + sync_slots(t); if loop_is_active { let ws = &ws_slice[k as usize]; @@ -435,7 +437,7 @@ pub fn gpu_mb_compute_dynamics_pre( ); } - workgroup_memory_barrier_with_group_sync(); + sync_slots(t); if loop_is_active { // i_coriolis_dt assembly: dt · (mass·coriolis_v, I·coriolis_w). @@ -461,7 +463,7 @@ pub fn gpu_mb_compute_dynamics_pre( ); } - workgroup_memory_barrier_with_group_sync(); + sync_slots(t); if loop_is_active { gemm_tr_par( @@ -478,7 +480,7 @@ pub fn gpu_mb_compute_dynamics_pre( ); } - workgroup_memory_barrier_with_group_sync(); + sync_slots(t); } // Diagonal: M[i, i] += damping[i] * dt + armature[i] — parallel. @@ -555,7 +557,7 @@ pub fn gpu_mb_compute_dynamics_without_coriolis_pre( if active_slot && num_links > 0 && lane == 0 { forward_kinematics(&mb, &stat_slice, &mut poses_slice, &mut ws_slice, num_links); } - workgroup_memory_barrier_with_group_sync(); + sync_slots(t); // 2) Update body jacobians update_body_jacobians( @@ -574,12 +576,12 @@ pub fn gpu_mb_compute_dynamics_without_coriolis_pre( if active_slot && num_links > 0 && lane == 0 { propagate_velocities(num_links, &stat_slice, &vel_slice, &mut ws_slice); } - workgroup_memory_barrier_with_group_sync(); + sync_slots(t); // 4) Mass matrix (without coriolis). let acc_augmented_mass = MatSlice::dense(mb_mm_base, ndofs, ndofs); fill_par(mass_matrices, acc_augmented_mass, 0.0, lane, t); - workgroup_memory_barrier_with_group_sync(); + sync_slots(t); // NOTE: uniform trip count (from the `BatchIndices` uniform). for k in 0..batch_ids.mb_max_links { @@ -617,7 +619,7 @@ pub fn gpu_mb_compute_dynamics_without_coriolis_pre( ); } - workgroup_memory_barrier_with_group_sync(); + sync_slots(t); } // Diagonal: M[i, i] += damping[i] * dt + armature[i] — parallel. @@ -791,7 +793,7 @@ fn update_body_jacobians( } } - workgroup_memory_barrier_with_group_sync(); + sync_slots(lanes); if k < num_links { let link_infos = &stat_slice[k as usize]; @@ -805,7 +807,7 @@ fn update_body_jacobians( ); } - workgroup_memory_barrier_with_group_sync(); + sync_slots(lanes); if k < num_links { let link = &ws_slice[k as usize]; @@ -822,7 +824,7 @@ fn update_body_jacobians( ); } - workgroup_memory_barrier_with_group_sync(); + sync_slots(lanes); } } diff --git a/src_rbd_shaders/dynamics/multibody/gravity_and_lu.rs b/src_rbd_shaders/dynamics/multibody/gravity_and_lu.rs index 4582c42..2b86861 100644 --- a/src_rbd_shaders/dynamics/multibody/gravity_and_lu.rs +++ b/src_rbd_shaders/dynamics/multibody/gravity_and_lu.rs @@ -17,7 +17,9 @@ use glamx::Vec4; use crate::dynamics::body::Velocity; use crate::dynamics::joint::SPATIAL_DIM; -use crate::utils::linalg::{MAX_MB_DOFS, MatSlice, fill_par, gemv_tr_spatial_split_par}; +use crate::utils::linalg::{ + MAX_MB_DOFS, MatSlice, fill_par, gemv_tr_spatial_split_par, lu_decompose, lu_solve_in_place, +}; use crate::utils::{BatchIndices, Slice}; use crate::{AngVector, Vector, gcross_av}; @@ -516,6 +518,181 @@ fn gravity_and_lu_packed_impl= num_mb * batch_ids.num_batches { + return; + } + let batch_id = invocation_id.x / num_mb; + let mb_idx = invocation_id.x % num_mb; + + let mb = batch_ids + .mb_batch(batch_id, multibody_info) + .read(mb_idx as usize); + let num_links = mb.num_links; + let ndofs = mb.ndofs; + if ndofs == 0 { + return; + } + let mb_jac_base = batch_ids.jac_start(batch_id) + mb.jacobian_offset as usize; + let gen_base = batch_ids.dof_start(batch_id) + mb.first_dof as usize; + let mb_mm_base = batch_ids.mm_start(batch_id) + mb.mass_matrix_offset as usize; + let piv_offset = gen_base; + let rhs_offset = gen_base; + + let stat_slice = batch_ids + .mb_links_batch(batch_id, links_static) + .offset(mb.first_link as usize); + let mut ws_slice = batch_ids + .mb_links_batch_mut(batch_id, links_workspace) + .offset(mb.first_link as usize); + let vel_slice = Slice(dof_state, gen_base); + let damping_slice = Slice( + dof_state, + batch_ids.dof_damping_section_offset as usize + gen_base, + ); + + // ---- Phase 1: zero the generalized-force vector. ---- + for d in 0..ndofs { + gen_forces.write(gen_base + d as usize, 0.0); + } + + #[cfg(feature = "dim3")] + let g = Vec3::new(gravity.x, gravity.y, gravity.z); + #[cfg(feature = "dim2")] + let g = Vec2::new(gravity.x, gravity.y); + + // ---- Phase 2: per-link gravity / Coriolis-force assembly (serial: + // parents precede children in link order, so `kinematic_acc` reads see + // the parent's write in program order). ---- + for k in 0..num_links { + let mut acc_lin = Vector::ZERO; + #[cfg(feature = "dim3")] + let mut acc_ang: AngVector = AngVector::ZERO; + #[cfg(feature = "dim2")] + let mut acc_ang: AngVector = 0.0; + + let (self_joint_vel_lin, self_joint_vel_ang, self_shift02, self_shift23, self_rb_ang) = { + let ws = &ws_slice[k as usize]; + ( + ws.joint_velocity.linear, + ws.joint_velocity.angular, + ws.shift02, + ws.shift23, + ws.rb_vels.angular, + ) + }; + + if k != 0 { + let stat = stat_slice[k as usize]; + let parent_ws = &ws_slice[stat.parent_link_id as usize]; + let parent_acc_lin = parent_ws.kinematic_acc.linear; + let parent_acc_ang = parent_ws.kinematic_acc.angular; + let parent_ang = parent_ws.rb_vels.angular; + + acc_lin = parent_acc_lin; + acc_ang = parent_acc_ang; + + acc_lin += gcross_av(parent_ang, self_joint_vel_lin) * 2.0; + #[cfg(feature = "dim3")] + { + acc_ang += parent_ang.cross(self_joint_vel_ang); + } + #[cfg(feature = "dim2")] + { + let _ = self_joint_vel_ang; + } + acc_lin += gcross_av(parent_ang, gcross_av(parent_ang, self_shift02)); + acc_lin += gcross_av(parent_acc_ang, self_shift02); + } else { + let _ = self_joint_vel_ang; + let _ = self_shift02; + } + let rb_ang = self_rb_ang; + acc_lin += gcross_av(rb_ang, gcross_av(rb_ang, self_shift23)); + acc_lin += gcross_av(acc_ang, self_shift23); + + ws_slice[k as usize].kinematic_acc = Velocity::new(acc_lin, acc_ang); + + let lmp = stat_slice[k as usize].local_mprops; + let inv_mass_x = lmp.inv_mass.x; + if inv_mass_x != 0.0 { + let mass = 1.0 / inv_mass_x; + let rb_inertia = ws_slice[k as usize].link_world_inertia(&lmp); + + #[cfg(feature = "dim3")] + let gyroscopic = { + let i_omega = rb_inertia * rb_ang; + rb_ang.cross(i_omega) + }; + #[cfg(feature = "dim2")] + let gyroscopic: AngVector = 0.0; + + let i_acc_ang = rb_inertia * acc_ang; + + let f_lin = (g - acc_lin) * mass; + let f_ang = -gyroscopic - i_acc_ang; + + let body_jacobian = MatSlice::dense( + mb_jac_base + (k as usize) * SPATIAL_DIM * (ndofs as usize), + SPATIAL_DIM as u32, + ndofs, + ); + + // Single lane owns the whole gemv (lane = 0, lanes = 1). + gemv_tr_spatial_split_par( + gen_forces, + gen_base, + 1.0, + body_jacobians, + body_jacobian, + f_lin, + f_ang, + 1.0, + 0, + 1, + ); + } + } + + // Damping subtraction. + for i in 0..ndofs { + let idx = gen_base + i as usize; + let cur = gen_forces.read(idx); + gen_forces.write(idx, cur - damping_slice[i as usize] * vel_slice[i as usize]); + } + + // ---- Phase 3 + 4: factor M in place in global memory, then solve + // M·x = τ in place on the gravity rhs. ---- + let m_view = MatSlice::dense(mb_mm_base, ndofs, ndofs); + lu_decompose(mass_matrices, m_view, lu_pivots, piv_offset); + lu_solve_in_place( + mass_matrices, + m_view, + lu_pivots, + piv_offset, + gen_forces, + rhs_offset, + ); +} + /// Stamps one packed-tier entry point of the fused gravity + LU kernel. /// `MATN = 64·T`, `SLOTS = 64/T`. macro_rules! gravity_and_lu_packed_entry { From 15087ad656038733237b4e602cb6491f3aeebe07 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Sat, 1 Aug 2026 13:08:05 +0200 Subject: [PATCH 35/39] perf: zero-workgroup indirect gating for the multibody contact pipeline --- src_rbd/broad_phase/narrow_phase.rs | 12 +++-- .../dynamics/multibody/multibody_solver.rs | 51 +++++++++++++------ src_rbd/dynamics/solver.rs | 7 +++ src_rbd/pipeline/insertion_removal.rs | 3 ++ src_rbd/pipeline/rbd_state.rs | 3 ++ src_rbd/pipeline/rbd_state_from_rapier.rs | 3 ++ src_rbd/pipeline/rbd_step.rs | 6 ++- src_rbd_shaders/broad_phase/narrow_phase.rs | 20 +++++++- .../dynamics/multibody/contact_constraints.rs | 37 +++++--------- 9 files changed, 98 insertions(+), 44 deletions(-) diff --git a/src_rbd/broad_phase/narrow_phase.rs b/src_rbd/broad_phase/narrow_phase.rs index 61a59a6..82d6ed6 100644 --- a/src_rbd/broad_phase/narrow_phase.rs +++ b/src_rbd/broad_phase/narrow_phase.rs @@ -42,6 +42,7 @@ impl GpuNarrowPhase { contacts: &mut Tensor, contacts_len: &mut Tensor, contacts_indirect: &mut Tensor<[u32; 3]>, + mb_sweep_indirect: &mut Tensor<[u32; 3]>, pfm_pairs: &mut Tensor, pfm_pairs_len: &mut Tensor, pfm_pairs_indirect: &mut Tensor<[u32; 3]>, @@ -98,9 +99,14 @@ impl GpuNarrowPhase { collider_parent, collider_materials, )?; - // Single 256-lane workgroup: parallel max over the per-batch counts. - self.init_contacts_indirect_args - .call(pass, 256u32, contacts_len, contacts_indirect)?; + self.init_contacts_indirect_args.call( + pass, + 256u32, + contacts_len, + contacts_indirect, + mb_sweep_indirect, + batch_indices, + )?; Ok(()) } diff --git a/src_rbd/dynamics/multibody/multibody_solver.rs b/src_rbd/dynamics/multibody/multibody_solver.rs index 648abae..b79e4d2 100644 --- a/src_rbd/dynamics/multibody/multibody_solver.rs +++ b/src_rbd/dynamics/multibody/multibody_solver.rs @@ -94,6 +94,9 @@ pub struct MultibodySolverArgs<'a> { /// Per-color-index uniform tensors (`color_uniforms[c]` holds `c`), /// shared with the contact/joint solvers. pub color_uniforms: &'a [Tensor], + /// GPU-written workgroup grid for the per-multibody contact-constraint + /// dispatches: `[multibodies_batch_capacity, num_batches, 1]`. + pub mb_sweep_indirect: &'a Tensor<[u32; 3]>, } impl GpuMultibodySolver { @@ -130,13 +133,15 @@ impl GpuMultibodySolver { } // Zero the accumulated contact impulses so the first substep's warmstart // starts cold (within a frame they are then preserved across substeps). - // One 64-lane workgroup per multibody (lanes stride the slots). + // Flat (slot, multibody, batch) grid (impulse-field-only stores). { let mut pass = encoder.begin_pass("[RBD] mbi/reset", timestamps.as_deref_mut()); + let total_slots = mb.num_active_multibodies + * mb.num_batches + * crate::shaders::dynamics::MAX_MB_CONTACT_CONSTRAINTS_PER_MB; self.reset_contact_warmstart.call( &mut pass, - [mb.multibodies_per_batch * MB_LU_LANES, mb.num_batches, 1], - &mb.multibody_info, + [total_slots, 1, 1], &mut mb.contact_constraints, args.batch_indices, )?; @@ -244,11 +249,11 @@ impl GpuMultibodySolver { if !first_substep { let mut pass = encoder.begin_pass("[RBD] mbb/warmstart-contact", timestamps.as_deref_mut()); - let warmstart_dispatch = - [mb.multibodies_per_batch * MB_LU_LANES, mb.num_batches, 1]; + // Contact-only work: indirect grid collapses to zero workgroups + // when no batch has any contact this step. self.warmstart_contact_constraints.call( &mut pass, - warmstart_dispatch, + args.mb_sweep_indirect, &mb.multibody_info, &mb.contact_constraints, &mb.contact_constraint_columns, @@ -326,16 +331,13 @@ impl GpuMultibodySolver { )?; } - // One 64-lane workgroup per multibody: the per-constraint LU - // back-solves are independent, so they run one-per-lane instead of - // sequentially on a single thread. + // One 64-lane workgroup per multibody. { let mut pass = encoder.begin_pass("[RBD] mbb/finalize-contact", timestamps.as_deref_mut()); - let finalize_dispatch = [mb.multibodies_per_batch * MB_LU_LANES, mb.num_batches, 1]; self.finalize_contact_constraints.call( &mut pass, - finalize_dispatch, + args.mb_sweep_indirect, &mb.multibody_info, &mb.mass_matrices, &mb.lu_pivots, @@ -351,10 +353,9 @@ impl GpuMultibodySolver { if let Some(delassus) = &mut mb.contact_delassus { let mut pass = encoder.begin_pass("[RBD] mbb/build-delassus", timestamps.as_deref_mut()); - let dispatch = [mb.multibodies_per_batch * MB_LU_LANES, mb.num_batches, 1]; self.build_contact_delassus.call( &mut pass, - dispatch, + args.mb_sweep_indirect, &mb.multibody_info, &mb.contact_constraints, &mb.contact_constraint_jacs, @@ -393,9 +394,11 @@ impl GpuMultibodySolver { args.batch_indices, )?; } + // Contact-only work: indirect grid collapses to zero workgroups + // on contact-free steps. self.solve_contacts_delassus.call( pass, - solve_dispatch, + args.mb_sweep_indirect, &mb.multibody_info, &mut mb.contact_constraints, &mb.contact_constraint_jacs, @@ -406,7 +409,7 @@ impl GpuMultibodySolver { &mut mb.dof_state, args.solver_vels, )?; - } else { + } else if mb.has_joint_constraints { self.solve_constraints.call( pass, solve_dispatch, @@ -421,6 +424,24 @@ impl GpuMultibodySolver { &mut mb.dof_state, args.solver_vels, )?; + } else { + // No joint limits/motors anywhere: the fused sweep is contact-only + // work, so the indirect grid (zero workgroups on contact-free + // steps) replaces the full per-(multibody, batch) launch. + self.solve_constraints.call( + pass, + args.mb_sweep_indirect, + &mb.multibody_info, + &mut mb.joint_constraints, + &mb.joint_constraint_columns, + &mut mb.contact_constraints, + &mb.contact_constraint_jacs, + &mb.contact_constraint_columns, + use_bias, + args.batch_indices, + &mut mb.dof_state, + args.solver_vels, + )?; } Ok(()) } diff --git a/src_rbd/dynamics/solver.rs b/src_rbd/dynamics/solver.rs index d2b1e1b..b235f73 100644 --- a/src_rbd/dynamics/solver.rs +++ b/src_rbd/dynamics/solver.rs @@ -146,6 +146,9 @@ pub struct SolverArgs<'a> { pub rb_contacts_inert: bool, /// Shared per-batch indices. pub batch_indices: &'a Tensor, + /// GPU-written workgroup grid for the per-multibody contact-constraint + /// dispatches (zero workgroups on contact-free steps). + pub mb_sweep_indirect: &'a Tensor<[u32; 3]>, } impl GpuSolver { @@ -292,6 +295,7 @@ impl GpuSolver { solver_vels: &mut *args.solver_vels, batch_indices: args.batch_indices, color_uniforms: args.color_uniforms, + mb_sweep_indirect: args.mb_sweep_indirect, }; solver.stash_contacts_len(&mut pass, state, &mut mb_args)?; } @@ -315,6 +319,7 @@ impl GpuSolver { solver_vels: &mut *args.solver_vels, batch_indices: args.batch_indices, color_uniforms: args.color_uniforms, + mb_sweep_indirect: args.mb_sweep_indirect, }; solver.build_contact_constraints( encoder, @@ -356,6 +361,7 @@ impl GpuSolver { solver_vels: &mut *args.solver_vels, batch_indices: args.batch_indices, color_uniforms: args.color_uniforms, + mb_sweep_indirect: args.mb_sweep_indirect, }; solver.$method(&mut pass, state, &mut mb_args $(, $extra)*)?; } @@ -399,6 +405,7 @@ impl GpuSolver { solver_vels: &mut *args.solver_vels, batch_indices: args.batch_indices, color_uniforms: args.color_uniforms, + mb_sweep_indirect: args.mb_sweep_indirect, }; solver.substep_build_constraints( encoder, diff --git a/src_rbd/pipeline/insertion_removal.rs b/src_rbd/pipeline/insertion_removal.rs index c20e42d..2c9870b 100644 --- a/src_rbd/pipeline/insertion_removal.rs +++ b/src_rbd/pipeline/insertion_removal.rs @@ -162,6 +162,8 @@ impl RbdState { .unwrap(); let contacts_indirect = Tensor::scalar_uninit(backend, BufferUsages::STORAGE | BufferUsages::INDIRECT).unwrap(); + let mb_sweep_indirect = + Tensor::scalar_uninit(backend, BufferUsages::STORAGE | BufferUsages::INDIRECT).unwrap(); let pfm_pairs_indirect = Tensor::scalar_uninit(backend, BufferUsages::STORAGE | BufferUsages::INDIRECT).unwrap(); let pfm_pairs = @@ -280,6 +282,7 @@ impl RbdState { contacts, contacts_len, contacts_indirect, + mb_sweep_indirect, pfm_pairs, pfm_pairs_len, pfm_pairs_indirect, diff --git a/src_rbd/pipeline/rbd_state.rs b/src_rbd/pipeline/rbd_state.rs index 6591c8c..4a7efff 100644 --- a/src_rbd/pipeline/rbd_state.rs +++ b/src_rbd/pipeline/rbd_state.rs @@ -191,6 +191,9 @@ pub struct RbdState { pub(super) contacts: Tensor, pub(super) contacts_len: Tensor, pub(super) contacts_indirect: Tensor<[u32; 3]>, + /// Workgroup grid for the per-multibody contact-constraint dispatches: + /// `[multibodies_batch_capacity, num_batches, 1]`. + pub(super) mb_sweep_indirect: Tensor<[u32; 3]>, pub(super) new_constraints: Tensor, pub(super) new_constraint_builders: Tensor, pub(super) new_constraints_counts: Tensor, diff --git a/src_rbd/pipeline/rbd_state_from_rapier.rs b/src_rbd/pipeline/rbd_state_from_rapier.rs index 8d427b4..b84a3e7 100644 --- a/src_rbd/pipeline/rbd_state_from_rapier.rs +++ b/src_rbd/pipeline/rbd_state_from_rapier.rs @@ -618,6 +618,8 @@ impl RbdState { .unwrap(); let contacts_indirect = Tensor::scalar_uninit(backend, BufferUsages::STORAGE | BufferUsages::INDIRECT).unwrap(); + let mb_sweep_indirect = + Tensor::scalar_uninit(backend, BufferUsages::STORAGE | BufferUsages::INDIRECT).unwrap(); let pfm_pairs_indirect = Tensor::scalar_uninit(backend, BufferUsages::STORAGE | BufferUsages::INDIRECT).unwrap(); let pfm_pairs = Tensor::vector_uninit( @@ -808,6 +810,7 @@ impl RbdState { contacts, contacts_len, contacts_indirect, + mb_sweep_indirect, pfm_pairs, pfm_pairs_len, pfm_pairs_indirect, diff --git a/src_rbd/pipeline/rbd_step.rs b/src_rbd/pipeline/rbd_step.rs index a4a20a0..f2a83cf 100644 --- a/src_rbd/pipeline/rbd_step.rs +++ b/src_rbd/pipeline/rbd_step.rs @@ -95,6 +95,7 @@ impl RbdPipeline { solver_vels: &mut state.solver_vels, batch_indices: &state.batch_indices, color_uniforms: &state.color_uniforms, + mb_sweep_indirect: &state.mb_sweep_indirect, }; self.multibody_solver.init_step( &mut encoder, @@ -250,6 +251,7 @@ impl RbdPipeline { &mut state.contacts, &mut state.contacts_len, &mut state.contacts_indirect, + &mut state.mb_sweep_indirect, &mut state.pfm_pairs, &mut state.pfm_pairs_len, &mut state.pfm_pairs_indirect, @@ -300,6 +302,7 @@ impl RbdPipeline { num_solver_iterations: state.num_solver_iterations, body_group: &state.body_group, batch_indices: &state.batch_indices, + mb_sweep_indirect: &state.mb_sweep_indirect, colorless_warmstart: false, fused_color_sweeps, rb_contacts_inert: state.rb_contacts_inert, @@ -446,7 +449,8 @@ impl RbdPipeline { num_solver_iterations: state.num_solver_iterations, body_group: &state.body_group, batch_indices: &state.batch_indices, - // The gather warmstart is only valid without multibody grouping — + mb_sweep_indirect: &state.mb_sweep_indirect, + // The gather warmstart is only valid without multibody grouping; // see `SolverArgs::colorless_warmstart`. #[cfg(feature = "dim3")] colorless_warmstart: state.multibodies.is_empty(), diff --git a/src_rbd_shaders/broad_phase/narrow_phase.rs b/src_rbd_shaders/broad_phase/narrow_phase.rs index ae9364b..ab919d8 100644 --- a/src_rbd_shaders/broad_phase/narrow_phase.rs +++ b/src_rbd_shaders/broad_phase/narrow_phase.rs @@ -41,17 +41,33 @@ pub fn gpu_reset_narrow_phase( } } -/// Initializes indirect dispatch arguments for constraint solver. Dispatch one -/// [`MAX_REDUCE_LANES`]-thread workgroup. +/// Initializes indirect dispatch arguments for constraint solver. +/// +/// Also inits `mb_sweep_indirect`, the workgroup grid for the per-multibody +/// contact-constraint dispatches (`[multibodies_batch_capacity, num_batches, +/// 1]`). #[spirv_bindgen] #[spirv(compute(threads(256)))] pub fn gpu_narrow_phase_init_contacts_dispatch( #[spirv(local_invocation_id)] lid: UVec3, #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] contacts_len: &mut [u32], #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] indirect_args: &mut [u32; 3], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] mb_sweep_indirect: &mut [u32; 3], + #[spirv(uniform, descriptor_set = 0, binding = 3)] batch_ids: &BatchIndices, #[spirv(workgroup)] partial: &mut [u32; MAX_REDUCE_LANES as usize], ) { max_len_indirect_args(lid.x, contacts_len, indirect_args, partial); + // `partial[0]` holds the max after the reduction (all lanes synced). + if lid.x == 0 { + let any_contacts = partial.read(0) > 0; + *mb_sweep_indirect.at_mut(0) = if any_contacts { + batch_ids.multibodies_batch_capacity + } else { + 0 + }; + *mb_sweep_indirect.at_mut(1) = batch_ids.num_batches; + *mb_sweep_indirect.at_mut(2) = 1; + } } pub(crate) const PREDICTION: f32 = 2.0e-3; // TODO: make the prediction configurable. diff --git a/src_rbd_shaders/dynamics/multibody/contact_constraints.rs b/src_rbd_shaders/dynamics/multibody/contact_constraints.rs index 9d65018..8197f66 100644 --- a/src_rbd_shaders/dynamics/multibody/contact_constraints.rs +++ b/src_rbd_shaders/dynamics/multibody/contact_constraints.rs @@ -602,35 +602,26 @@ pub fn gpu_mb_stash_contacts_len( #[spirv_bindgen] #[spirv(compute(threads(64)))] pub fn gpu_mb_reset_contact_warmstart( - #[spirv(workgroup_id)] workgroup_id: UVec3, - #[spirv(local_invocation_id)] local_id: UVec3, - #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] multibody_info: &[MultibodyInfo], - #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] contact_constraints: &mut [MultibodyContactConstraint], - #[spirv(uniform, descriptor_set = 0, binding = 2)] batch_ids: &BatchIndices, + #[spirv(uniform, descriptor_set = 0, binding = 1)] batch_ids: &BatchIndices, ) { - const LANES: u32 = 64; - let batch_id = workgroup_id.y; - let mb_idx = workgroup_id.x; - let lane = local_id.x; + // One thread per (slot, multibody, batch), flattened. + const MAXC: u32 = MAX_MB_CONTACT_CONSTRAINTS_PER_MB; let num_mb = batch_ids.multibodies_len; - if mb_idx >= num_mb { + let per_batch = num_mb * MAXC; + if invocation_id.x >= per_batch * batch_ids.num_batches { return; } - let mb_start = batch_ids.mb_start(batch_id); + let batch_id = invocation_id.x / per_batch; + let r = invocation_id.x % per_batch; + let mb_idx = r / MAXC; + let s = r % MAXC; + let cons_start = batch_ids.mb_contact_constraints_start(batch_id); - let mb = multibody_info.read(mb_start + mb_idx as usize); - if mb.ndofs == 0 { - return; - } - let cons_base = cons_start + (mb_idx as usize) * (MAX_MB_CONTACT_CONSTRAINTS_PER_MB as usize); - // Zero to capacity (the per-frame contact count isn't known here, and last - // frame's count may be smaller than this frame's). - for s in StepRng::new(lane..MAX_MB_CONTACT_CONSTRAINTS_PER_MB, LANES) { - let mut cons = contact_constraints.read(cons_base + s as usize); - cons.impulse = 0.0; - contact_constraints.write(cons_base + s as usize, cons); - } + let idx = cons_start + (mb_idx * MAXC + s) as usize; + contact_constraints.at_mut(idx).impulse = 0.0; } /// Warmstart: re-apply each active contact constraint's accumulated `impulse` From d8da3d10abdfbe4df31b3b50db71497ebb84c200 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Sat, 1 Aug 2026 13:34:51 +0200 Subject: [PATCH 36/39] perf: batch-interleaved layout for multibody dynamics buffers --- .../multibody/multibody_from_rapier.rs | 21 +++ src_rbd/dynamics/multibody/multibody_set.rs | 5 +- .../multibody/compute_dynamics_pre.rs | 117 +++++++-------- .../dynamics/multibody/contact_constraints.rs | 52 ++++--- .../dynamics/multibody/gravity_and_lu.rs | 133 +++++++++--------- .../impulse_joint_constraints/helper.rs | 33 ++--- .../impulse_joint_constraints/jacobians.rs | 25 ++-- .../impulse_joint_constraints/kernels.rs | 40 +++--- .../impulse_joint_constraints/update.rs | 50 ++++--- .../dynamics/multibody/integrate.rs | 30 ++-- .../dynamics/multibody/joint_constraints.rs | 30 ++-- src_rbd_shaders/dynamics/multibody/lu.rs | 18 +-- .../dynamics/multibody/solve_constraints.rs | 39 ++--- src_rbd_shaders/utils/indices.rs | 91 +++++------- src_rbd_shaders/utils/linalg.rs | 123 ++++++++++++---- src_rbd_shaders/utils/mod.rs | 2 +- src_rbd_shaders/utils/slice.rs | 112 +++++++++++++++ 17 files changed, 565 insertions(+), 356 deletions(-) diff --git a/src_rbd/dynamics/multibody/multibody_from_rapier.rs b/src_rbd/dynamics/multibody/multibody_from_rapier.rs index 9b5cad2..61cc4d9 100644 --- a/src_rbd/dynamics/multibody/multibody_from_rapier.rs +++ b/src_rbd/dynamics/multibody/multibody_from_rapier.rs @@ -373,6 +373,27 @@ impl GpuMultibodySet { let storage = BufferUsages::STORAGE | BufferUsages::COPY_DST; + // Batch-interleaved (batch-minor) layout for the + // dynamics buffers: element `k` of batch `b` lives at `k · nb + b`. + fn interleave(data: &[T], cap: u32, nb: usize) -> Vec { + let cap = cap as usize; + let mut out = Vec::with_capacity(data.len()); + for k in 0..cap { + for b in 0..nb { + out.push(data[b * cap + k]); + } + } + out + } + let nb = num_batches as usize; + let all_infos = interleave(&all_infos, mb_cap, nb); + let all_statics = interleave(&all_statics, links_cap, nb); + let all_ws = interleave(&all_ws, links_cap, nb); + let all_dof_vals = interleave(&all_dof_vals, dofs_cap, nb); + let all_dof_vels = interleave(&all_dof_vels, dofs_cap, nb); + let all_dof_damping = interleave(&all_dof_damping, dofs_cap, nb); + let all_dof_armature = interleave(&all_dof_armature, dofs_cap, nb); + Self { num_batches, multibodies_per_batch: mb_cap, diff --git a/src_rbd/dynamics/multibody/multibody_set.rs b/src_rbd/dynamics/multibody/multibody_set.rs index c14fe06..e041016 100644 --- a/src_rbd/dynamics/multibody/multibody_set.rs +++ b/src_rbd/dynamics/multibody/multibody_set.rs @@ -296,8 +296,9 @@ impl GpuMultibodySet { axis: JointAxis, target_vel: f32, ) -> Result<(), GpuBackendError> { - let stride = self.links_per_batch; - let global_idx = (batch * stride + link_id) as usize; + // Batch-interleaved links layout: element `link_id` of batch `batch` + // lives at `link_id · num_batches + batch` (mirror included). + let global_idx = (link_id * self.num_batches + batch) as usize; let axis_id = axis as usize; let entry = match self.links_static_mirror.get_mut(global_idx) { Some(e) => e, diff --git a/src_rbd_shaders/dynamics/multibody/compute_dynamics_pre.rs b/src_rbd_shaders/dynamics/multibody/compute_dynamics_pre.rs index 3bb24c8..89a234c 100644 --- a/src_rbd_shaders/dynamics/multibody/compute_dynamics_pre.rs +++ b/src_rbd_shaders/dynamics/multibody/compute_dynamics_pre.rs @@ -23,7 +23,7 @@ use crate::utils::linalg::{ gemm_omega_skew_tr_cross_buf_par, gemm_skew_tr_lhs_cross_buf_par, gemm_skew_tr_lhs_par, gemm_tr_par, quadform_spatial_par, }; -use crate::utils::{BatchIndices, Slice, SliceMut}; +use crate::utils::{BatchIndices, ISlice, ISliceMut, SliceMut}; use crate::{ANG_DIM, AngVector, DIM, Pose, Vector, gcross_av}; use parry::math::VectorExt; @@ -78,41 +78,38 @@ pub fn gpu_mb_compute_dynamics_pre( let mb = if active_slot { batch_ids - .mb_batch(batch_id, multibody_info) + .ib(batch_id, multibody_info) .read(mb_idx as usize) } else { MultibodyInfo::default() }; let num_links = mb.num_links; let ndofs = mb.ndofs; - let mb_jac_base = batch_ids.jac_start(batch_id) + mb.jacobian_offset as usize; - let mb_mm_base = batch_ids.mm_start(batch_id) + mb.mass_matrix_offset as usize; - let mb_cor_base = batch_ids.cor_start(batch_id) + mb.coriolis_offset as usize; - let mb_cor_w_base = batch_ids.coriolis_w_section_offset as usize + mb_cor_base; - let mb_icdt_base = batch_ids.i_coriolis_dt_section_offset as usize - + batch_ids.icdt_start(batch_id) - + mb.i_coriolis_dt_offset as usize; - let vel_base = batch_ids.dof_start(batch_id) + mb.first_dof as usize; + let mb_jac_base = mb.jacobian_offset as usize; + let mb_mm_base = mb.mass_matrix_offset as usize; + let mb_cor_base = mb.coriolis_offset as usize; + let mb_cor_w_base = batch_ids.coriolis_batch_capacity as usize + mb_cor_base; + let mb_icdt_base = + 2 * batch_ids.coriolis_batch_capacity as usize + mb.i_coriolis_dt_offset as usize; + let vel_base = mb.first_dof as usize; let stat_slice = batch_ids - .mb_links_batch(batch_id, links_static) + .ib(batch_id, links_static) .offset(mb.first_link as usize); let mut ws_slice = batch_ids - .mb_links_batch_mut(batch_id, links_workspace) + .ib_mut(batch_id, links_workspace) .offset(mb.first_link as usize); let mut poses_slice = batch_ids.coll_batch_mut(batch_id, poses); - let damping_slice = Slice( - dof_state, - vel_base + batch_ids.dof_damping_section_offset as usize, - ); + let damping_slice = batch_ids + .ib(batch_id, dof_state) + .offset(batch_ids.dof_batch_capacity as usize + vel_base); // Armature (reflected rotor inertia) section sits right after damping, at // `2 · dof_damping_section_offset` (= 2·N). Added to the mass-matrix diagonal // alongside `damping·dt`, matching rapier's `update_mass_matrix`. - let armature_slice = Slice( - dof_state, - vel_base + 2 * batch_ids.dof_damping_section_offset as usize, - ); - let vel_slice = Slice(dof_state, vel_base); + let armature_slice = batch_ids + .ib(batch_id, dof_state) + .offset(2 * batch_ids.dof_batch_capacity as usize + vel_base); + let vel_slice = batch_ids.ib(batch_id, dof_state).offset(vel_base); // 1) Forward Kinematics (single-threaded) if active_slot && num_links > 0 && lane == 0 { @@ -131,6 +128,8 @@ pub fn gpu_mb_compute_dynamics_pre( &stat_slice, &ws_slice.as_ref(), body_jacobians, + batch_ids, + batch_id, ); // 3) Propagate velocities (single-threaded) @@ -140,10 +139,10 @@ pub fn gpu_mb_compute_dynamics_pre( sync_slots(t); // 3) Mass matrix (with semi-implicit coriolis handling). - let acc_augmented_mass = MatSlice::dense(mb_mm_base, ndofs, ndofs); + let acc_augmented_mass = batch_ids.imat(batch_id, mb_mm_base, ndofs, ndofs); fill_par(mass_matrices, acc_augmented_mass, 0.0, lane, t); - let i_coriolis_dt_view = MatSlice::dense(mb_icdt_base, SPATIAL_DIM as u32, ndofs); + let i_coriolis_dt_view = batch_ids.imat(batch_id, mb_icdt_base, SPATIAL_DIM as u32, ndofs); let i_coriolis_dt_v = i_coriolis_dt_view.fixed_rows(0, DIM); let i_coriolis_dt_w = i_coriolis_dt_view.fixed_rows(DIM, ANG_DIM); @@ -160,7 +159,7 @@ pub fn gpu_mb_compute_dynamics_pre( inv_mass_x = lmp.inv_mass.x; if inv_mass_x == 0.0 { - let coriolis_block = MatSlice::dense( + let coriolis_block = batch_ids.imat(batch_id, mb_cor_base + (k as usize) * (DIM as usize) * (ndofs as usize), DIM, ndofs, @@ -168,7 +167,7 @@ pub fn gpu_mb_compute_dynamics_pre( fill_par(coriolis_packed, coriolis_block, 0.0, lane, t); fill_par( coriolis_packed, - MatSlice::dense( + batch_ids.imat(batch_id, mb_cor_w_base + (k as usize) * (DIM as usize) * (ndofs as usize), DIM, ndofs, @@ -184,17 +183,17 @@ pub fn gpu_mb_compute_dynamics_pre( sync_slots(t); let loop_is_active = k < num_links && inv_mass_x != 0.0; - let coriolis_v_i = MatSlice::dense( + let coriolis_v_i = batch_ids.imat(batch_id, mb_cor_base + (k as usize) * (DIM as usize) * (ndofs as usize), DIM, ndofs, ); - let coriolis_w_i = MatSlice::dense( + let coriolis_w_i = batch_ids.imat(batch_id, mb_cor_w_base + (k as usize) * (DIM as usize) * (ndofs as usize), ANG_DIM, ndofs, ); - let body_jacobian = MatSlice::dense( + let body_jacobian = batch_ids.imat(batch_id, mb_jac_base + (k as usize) * SPATIAL_DIM * (ndofs as usize), SPATIAL_DIM as u32, ndofs, @@ -237,18 +236,18 @@ pub fn gpu_mb_compute_dynamics_pre( let stat = stat_slice[k as usize]; let parent_id = stat.parent_link_id; let parent_link = &ws_slice[parent_id as usize]; - let parent_j = MatSlice::dense( + let parent_j = batch_ids.imat(batch_id, mb_jac_base + (parent_id as usize) * SPATIAL_DIM * (ndofs as usize), SPATIAL_DIM as u32, ndofs, ); let parent_j_w = parent_j.fixed_rows(DIM, ANG_DIM); - let parent_coriolis_v = MatSlice::dense( + let parent_coriolis_v = batch_ids.imat(batch_id, mb_cor_base + (parent_id as usize) * (DIM as usize) * (ndofs as usize), DIM, ndofs, ); - let parent_coriolis_w = MatSlice::dense( + let parent_coriolis_w = batch_ids.imat(batch_id, mb_cor_w_base + (parent_id as usize) * (DIM as usize) * (ndofs as usize), ANG_DIM, ndofs, @@ -522,36 +521,34 @@ pub fn gpu_mb_compute_dynamics_without_coriolis_pre( let mb = if active_slot { batch_ids - .mb_batch(batch_id, multibody_info) + .ib(batch_id, multibody_info) .read(mb_idx as usize) } else { MultibodyInfo::default() }; let num_links = mb.num_links; let ndofs = mb.ndofs; - let mb_jac_base = batch_ids.jac_start(batch_id) + mb.jacobian_offset as usize; - let mb_mm_base = batch_ids.mm_start(batch_id) + mb.mass_matrix_offset as usize; - let vel_base = batch_ids.dof_start(batch_id) + mb.first_dof as usize; + let mb_jac_base = mb.jacobian_offset as usize; + let mb_mm_base = mb.mass_matrix_offset as usize; + let vel_base = mb.first_dof as usize; let stat_slice = batch_ids - .mb_links_batch(batch_id, links_static) + .ib(batch_id, links_static) .offset(mb.first_link as usize); let mut ws_slice = batch_ids - .mb_links_batch_mut(batch_id, links_workspace) + .ib_mut(batch_id, links_workspace) .offset(mb.first_link as usize); let mut poses_slice = batch_ids.coll_batch_mut(batch_id, poses); - let damping_slice = Slice( - dof_state, - vel_base + batch_ids.dof_damping_section_offset as usize, - ); + let damping_slice = batch_ids + .ib(batch_id, dof_state) + .offset(batch_ids.dof_batch_capacity as usize + vel_base); // Armature (reflected rotor inertia) section sits right after damping, at // `2 · dof_damping_section_offset` (= 2·N). Added to the mass-matrix diagonal // alongside `damping·dt`, matching rapier's `update_mass_matrix`. - let armature_slice = Slice( - dof_state, - vel_base + 2 * batch_ids.dof_damping_section_offset as usize, - ); - let vel_slice = Slice(dof_state, vel_base); + let armature_slice = batch_ids + .ib(batch_id, dof_state) + .offset(2 * batch_ids.dof_batch_capacity as usize + vel_base); + let vel_slice = batch_ids.ib(batch_id, dof_state).offset(vel_base); // 1) Forward Kinematics (single-threaded) if active_slot && num_links > 0 && lane == 0 { @@ -570,6 +567,8 @@ pub fn gpu_mb_compute_dynamics_without_coriolis_pre( &stat_slice, &ws_slice.as_ref(), body_jacobians, + batch_ids, + batch_id, ); // 3) Velocities propagation (single-threaded) @@ -579,7 +578,7 @@ pub fn gpu_mb_compute_dynamics_without_coriolis_pre( sync_slots(t); // 4) Mass matrix (without coriolis). - let acc_augmented_mass = MatSlice::dense(mb_mm_base, ndofs, ndofs); + let acc_augmented_mass = batch_ids.imat(batch_id, mb_mm_base, ndofs, ndofs); fill_par(mass_matrices, acc_augmented_mass, 0.0, lane, t); sync_slots(t); @@ -599,7 +598,7 @@ pub fn gpu_mb_compute_dynamics_without_coriolis_pre( let mass = 1.0 / lmp.inv_mass.x; let inertia = ws.link_world_inertia(&lmp); - let body_jacobian = MatSlice::dense( + let body_jacobian = batch_ids.imat(batch_id, mb_jac_base + (k as usize) * SPATIAL_DIM * (ndofs as usize), SPATIAL_DIM as u32, ndofs, @@ -641,7 +640,7 @@ pub fn gpu_mb_compute_dynamics_without_coriolis_pre( fn jacobian_mul_coordinates( locked_axes: u32, assembly_id: u32, - vel_slice: &Slice, + vel_slice: &ISlice, ) -> (Vector, AngVector) { let mut lin = Vector::ZERO; #[cfg(feature = "dim3")] @@ -688,9 +687,9 @@ fn jacobian_mul_coordinates( // sequentially on a single thread. fn forward_kinematics( mb: &MultibodyInfo, - stat_slice: &Slice, + stat_slice: &ISlice, poses_slice: &mut SliceMut, - ws_slice: &mut SliceMut, + ws_slice: &mut ISliceMut, num_links: u32, ) { // Root pose. @@ -745,9 +744,11 @@ fn update_body_jacobians( ndofs: u32, num_links: u32, max_links: u32, - stat_slice: &Slice, - ws_slice: &Slice, + stat_slice: &ISlice, + ws_slice: &ISlice, body_jacobians: &mut [f32], + batch_ids: &BatchIndices, + batch_id: u32, ) { // TODO(PERF): instead of copying the body jacobian over and over for each body, we should // precompute a bit set that indicates which dofs are part of the kinematic tree @@ -755,7 +756,7 @@ fn update_body_jacobians( // value per node. for k in 0..max_links { let mut parent_to_world = Pose::default(); - let link_j = MatSlice::dense( + let link_j = batch_ids.imat(batch_id, mb_jac_base + (k as usize) * SPATIAL_DIM * (ndofs as usize), SPATIAL_DIM as u32, ndofs, @@ -766,7 +767,7 @@ fn update_body_jacobians( let link = &ws_slice[k as usize]; if k != 0 { - let parent_j = MatSlice::dense( + let parent_j = batch_ids.imat(batch_id, mb_jac_base + (link_infos.parent_link_id as usize) * SPATIAL_DIM * (ndofs as usize), SPATIAL_DIM as u32, @@ -830,9 +831,9 @@ fn update_body_jacobians( fn propagate_velocities( num_links: u32, - stat_slice: &Slice, - vel_slice: &Slice, - ws_slice: &mut SliceMut, + stat_slice: &ISlice, + vel_slice: &ISlice, + ws_slice: &mut ISliceMut, ) { for k in 0..num_links { let k_usize = k as usize; diff --git a/src_rbd_shaders/dynamics/multibody/contact_constraints.rs b/src_rbd_shaders/dynamics/multibody/contact_constraints.rs index 8197f66..e4aa8e7 100644 --- a/src_rbd_shaders/dynamics/multibody/contact_constraints.rs +++ b/src_rbd_shaders/dynamics/multibody/contact_constraints.rs @@ -23,7 +23,7 @@ use crate::dynamics::body::{Velocity, WorldMassProperties}; use crate::dynamics::joint::SPATIAL_DIM; use crate::queries::IndexedManifold; use crate::utils::BatchIndices; -use crate::utils::linalg::{MatSlice, lu_solve_in_place}; +use crate::utils::linalg::{MatSlice, VSlice, lu_solve_in_place}; use crate::{ANG_DIM, AngVector, DIM, Pose, Vector, gcross, gdot}; use super::types::{ @@ -66,6 +66,9 @@ fn orthonormal_vector(v: Vec2) -> Vec2 { fn fill_contact_jac_row( body_jacobians: &[f32], mb_jac_base: usize, + // Interleave parameters of `body_jacobians` (`num_batches`, `batch_id`). + jac_stride: u32, + jac_shift: u32, ndofs: u32, link_id: u32, unit_force: Vector, @@ -78,7 +81,8 @@ fn fill_contact_jac_row( // Per-link SPATIAL_DIM × ndofs jacobian (rows 0..DIM = J_v, rows // DIM..SPATIAL_DIM = J_w). let link_jac_base = mb_jac_base + (link_id as usize) * SPATIAL_DIM * (ndofs as usize); - let link_j = MatSlice::dense(link_jac_base, SPATIAL_DIM as u32, ndofs); + let link_j = + MatSlice::interleaved(link_jac_base, SPATIAL_DIM as u32, ndofs, jac_stride, jac_shift); let (link_j_v, link_j_w) = link_j.rows_range_pair(0, DIM, DIM, ANG_DIM); let j = lane; if j < ndofs { @@ -160,7 +164,6 @@ pub fn gpu_mb_init_contact_constraints( let max_corr_velocity = softness.max_corr_velocity; let cfm_factor = softness.cfm_factor; - let mb_start = batch_ids.mb_start(batch_id); let cons_start = batch_ids.mb_contact_constraints_start(batch_id); let col_start = batch_ids.mb_contact_constraint_columns_start(batch_id); let colliders_start = batch_ids.coll_start(batch_id); @@ -170,17 +173,17 @@ pub fn gpu_mb_init_contact_constraints( // Per-multibody early-out: padding multibody slots have `ndofs == 0`, // which we use here as the sentinel (replaces the `num_multibodies` // storage binding the kernel used to read). - let mut mb = multibody_info.read(mb_start + mb_idx as usize); + let mut mb = multibody_info.read(batch_ids.mbi(batch_id, mb_idx as usize)); let ndofs = mb.ndofs; if ndofs == 0 { // Uniform per workgroup: every lane returns together. if lane == 0 { mb.contact_constraint_count = 0; - multibody_info.write(mb_start + mb_idx as usize, mb); + multibody_info.write(batch_ids.mbi(batch_id, mb_idx as usize), mb); } return; } - let mb_jac_base = batch_ids.jac_start(batch_id) + mb.jacobian_offset as usize; + let mb_jac_base = mb.jacobian_offset as usize; let cons_base = cons_start + (mb_idx as usize) * (MAX_MB_CONTACT_CONSTRAINTS_PER_MB as usize); // Each constraint slot reserves `dof_batch_capacity` floats in the // column buffer (matches the allocation in `from_rapier` and avoids any @@ -319,6 +322,8 @@ pub fn gpu_mb_init_contact_constraints( fill_contact_jac_row( body_jacobians, mb_jac_base, + batch_ids.num_batches, + batch_id, ndofs, mb_link_id_a, mb_normal, @@ -344,6 +349,8 @@ pub fn gpu_mb_init_contact_constraints( fill_contact_jac_row( body_jacobians, mb_jac_base, + batch_ids.num_batches, + batch_id, ndofs, mb_link_id_b, lin_jac, @@ -466,6 +473,8 @@ pub fn gpu_mb_init_contact_constraints( fill_contact_jac_row( body_jacobians, mb_jac_base, + batch_ids.num_batches, + batch_id, ndofs, mb_link_id_a, mb_tangent, @@ -483,6 +492,8 @@ pub fn gpu_mb_init_contact_constraints( fill_contact_jac_row( body_jacobians, mb_jac_base, + batch_ids.num_batches, + batch_id, ndofs, mb_link_id_b, free_tangent, @@ -564,7 +575,7 @@ pub fn gpu_mb_init_contact_constraints( // don't need to mark surplus slots inactive — they're never read. if lane == 0 { mb.contact_constraint_count = count; - multibody_info.write(mb_start + mb_idx as usize, mb); + multibody_info.write(batch_ids.mbi(batch_id, mb_idx as usize), mb); } } @@ -587,10 +598,9 @@ pub fn gpu_mb_stash_contacts_len( } let batch_id = invocation_id.x / num_mb; let mb_idx = invocation_id.x % num_mb; - let mb_start = batch_ids.mb_start(batch_id); - let mut mb = multibody_info.read(mb_start + mb_idx as usize); + let mut mb = multibody_info.read(batch_ids.mbi(batch_id, mb_idx as usize)); mb.batch_contacts_len = contacts_len.read(batch_id as usize); - multibody_info.write(mb_start + mb_idx as usize, mb); + multibody_info.write(batch_ids.mbi(batch_id, mb_idx as usize), mb); } /// Zero the accumulated impulse of every contact-constraint slot for each @@ -651,17 +661,16 @@ pub fn gpu_mb_warmstart_contact_constraints( return; } - let mb_start = batch_ids.mb_start(batch_id); let cons_start = batch_ids.mb_contact_constraints_start(batch_id); let col_start = batch_ids.mb_contact_constraint_columns_start(batch_id); let colliders_start = batch_ids.coll_start(batch_id); - let mb = multibody_info.read(mb_start + mb_idx as usize); + let mb = multibody_info.read(batch_ids.mbi(batch_id, mb_idx as usize)); let ndofs = mb.ndofs; if ndofs == 0 { return; } - let v_base = batch_ids.dof_start(batch_id) + mb.first_dof as usize; + let v_base = mb.first_dof as usize; let cons_base = cons_start + (mb_idx as usize) * (MAX_MB_CONTACT_CONSTRAINTS_PER_MB as usize); let dofs_stride = batch_ids.dof_batch_capacity as usize; let col_base = @@ -676,7 +685,7 @@ pub fn gpu_mb_warmstart_contact_constraints( // This lane's DOF velocity, accumulated in a register across every // constraint. let mut v_lane = if lane < ndofs { - dof_state.read(v_base + lane as usize) + dof_state.read(batch_ids.mbi(batch_id, v_base + lane as usize)) } else { 0.0 }; @@ -703,7 +712,7 @@ pub fn gpu_mb_warmstart_contact_constraints( } if lane < ndofs { - dof_state.write(v_base + lane as usize, v_lane); + dof_state.write(batch_ids.mbi(batch_id, v_base + lane as usize), v_lane); } } @@ -734,23 +743,22 @@ pub fn gpu_mb_finalize_contact_constraints( return; } - let mb_start = batch_ids.mb_start(batch_id); let cons_start = batch_ids.mb_contact_constraints_start(batch_id); let col_start = batch_ids.mb_contact_constraint_columns_start(batch_id); - let mb = multibody_info.read(mb_start + mb_idx as usize); + let mb = multibody_info.read(batch_ids.mbi(batch_id, mb_idx as usize)); let ndofs = mb.ndofs; if ndofs == 0 { return; } - let mb_mm_base = batch_ids.mm_start(batch_id) + mb.mass_matrix_offset as usize; - let piv_offset = batch_ids.dof_start(batch_id) + mb.first_dof as usize; + let mb_mm_base = mb.mass_matrix_offset as usize; + let piv = batch_ids.ivec(batch_id, mb.first_dof as usize); let cons_base = cons_start + (mb_idx as usize) * (MAX_MB_CONTACT_CONSTRAINTS_PER_MB as usize); let dofs_stride = batch_ids.dof_batch_capacity as usize; let col_base = col_start + (mb_idx as usize) * (MAX_MB_CONTACT_CONSTRAINTS_PER_MB as usize) * dofs_stride; - let m = MatSlice::dense(mb_mm_base, ndofs, ndofs); + let m = batch_ids.imat(batch_id, mb_mm_base, ndofs, ndofs); let count = mb.contact_constraint_count; for s in StepRng::new(lane..count, LANES) { @@ -766,9 +774,9 @@ pub fn gpu_mb_finalize_contact_constraints( mass_matrices, m, lu_pivots, - piv_offset, + piv, contact_constraint_columns, - col_offset, + VSlice::dense(col_offset), ); // 3) inv_r_mb = J · column. let mut inv_r_mb = 0.0f32; diff --git a/src_rbd_shaders/dynamics/multibody/gravity_and_lu.rs b/src_rbd_shaders/dynamics/multibody/gravity_and_lu.rs index 2b86861..e2a8cca 100644 --- a/src_rbd_shaders/dynamics/multibody/gravity_and_lu.rs +++ b/src_rbd_shaders/dynamics/multibody/gravity_and_lu.rs @@ -65,30 +65,28 @@ pub fn gpu_mb_gravity_and_lu( let max_links = batch_ids.mb_max_links; let mb = batch_ids - .mb_batch(batch_id, multibody_info) + .ib(batch_id, multibody_info) .read(mb_idx as usize); let num_links = mb.num_links; let ndofs = mb.ndofs; - let mb_jac_base = batch_ids.jac_start(batch_id) + mb.jacobian_offset as usize; - let gen_base = batch_ids.dof_start(batch_id) + mb.first_dof as usize; - let mb_mm_base = batch_ids.mm_start(batch_id) + mb.mass_matrix_offset as usize; - let piv_offset = gen_base; - let rhs_offset = gen_base; + let mb_jac_base = mb.jacobian_offset as usize; + let gen_base = mb.first_dof as usize; + let mb_mm_base = mb.mass_matrix_offset as usize; + let piv = batch_ids.ivec(batch_id, gen_base); let stat_slice = batch_ids - .mb_links_batch(batch_id, links_static) + .ib(batch_id, links_static) .offset(mb.first_link as usize); let mut ws_slice = batch_ids - .mb_links_batch_mut(batch_id, links_workspace) + .ib_mut(batch_id, links_workspace) .offset(mb.first_link as usize); - let vel_slice = Slice(dof_state, gen_base); - let damping_slice = Slice( - dof_state, - batch_ids.dof_damping_section_offset as usize + gen_base, - ); + let vel_slice = batch_ids.ib(batch_id, dof_state).offset(gen_base); + let damping_slice = batch_ids + .ib(batch_id, dof_state) + .offset(batch_ids.dof_batch_capacity as usize + gen_base); // ---- Phase 1: zero the generalized-force vector (parallel across DOFs). ---- - let accelerations = MatSlice::dense(gen_base, ndofs, 1); + let accelerations = batch_ids.imat(batch_id, gen_base, ndofs, 1); // TODO(perf): up to a certain number of degrees of freedom, we could actually run all the // calculations in shared memory and only write the result in the end. // Currently, the max number of dofs is 32 but we still accumulate forces/accelerations @@ -195,7 +193,7 @@ pub fn gpu_mb_gravity_and_lu( let f_lin = (g - acc_lin) * mass; let f_ang = -gyroscopic - i_acc_ang; - let body_jacobian = MatSlice::dense( + let body_jacobian = batch_ids.imat(batch_id, mb_jac_base + (k as usize) * SPATIAL_DIM * (ndofs as usize), SPATIAL_DIM as u32, ndofs, @@ -203,7 +201,7 @@ pub fn gpu_mb_gravity_and_lu( gemv_tr_spatial_split_par( gen_forces, - gen_base, + batch_ids.ivec(batch_id, gen_base), 1.0, body_jacobians, body_jacobian, @@ -225,19 +223,22 @@ pub fn gpu_mb_gravity_and_lu( workgroup_memory_barrier_with_group_sync(); let i = lane; if i < ndofs { - let idx = gen_base + i as usize; + let idx = batch_ids.mbi(batch_id, gen_base + i as usize); let cur = gen_forces.read(idx); gen_forces.write(idx, cur - damping_slice[i as usize] * vel_slice[i as usize]); } workgroup_memory_barrier_with_group_sync(); // ---- Phase 3: load M into shared memory, factor in place. ---- - let m_view = MatSlice::dense(mb_mm_base, ndofs, ndofs); + let m_view = batch_ids.imat(batch_id, mb_mm_base, ndofs, ndofs); if lane < ndofs { for r in 0..ndofs { mat.write(sm_idx(r, lane), mass_matrices.read(m_view.idx(r, lane))); } - x.write(lane as usize, gen_forces.read(rhs_offset + lane as usize)); + x.write( + lane as usize, + gen_forces.read(batch_ids.mbi(batch_id, gen_base + lane as usize)), + ); } workgroup_memory_barrier_with_group_sync(); @@ -247,7 +248,7 @@ pub fn gpu_mb_gravity_and_lu( lane, mat, lu_pivots, - piv_offset, + piv, pivot_row_shared, inv_akk_shared, ); @@ -261,11 +262,11 @@ pub fn gpu_mb_gravity_and_lu( } // ---- Phase 4: solve M·x = τ for the gravity rhs. ---- - lu_apply_pivots(ndofs, lane, lu_pivots, piv_offset, x); + lu_apply_pivots(ndofs, lane, lu_pivots, piv, x); lu_triangular_solve_in_place(ndofs, max_ndofs, lane, mat, x, partial); if lane < ndofs { - gen_forces.write(rhs_offset + lane as usize, x.read(lane as usize)); + gen_forces.write(batch_ids.mbi(batch_id, gen_base + lane as usize), x.read(lane as usize)); } } @@ -310,33 +311,31 @@ fn gravity_and_lu_packed_impl(ndofs, slot, lane, active_slot, lu_pivots, piv_offset, x); + lu_apply_pivots_packed::(ndofs, slot, lane, active_slot, lu_pivots, piv, x); lu_triangular_solve_in_place_packed::( ndofs, max_ndofs, @@ -514,7 +516,10 @@ fn gravity_and_lu_packed_impl 0 { - let mb = multibody_info.read(mb_start + c.side_a_id as usize); + let mb = multibody_info.read(il.atz(c.side_a_id as usize)); solve_mb_wj( jacobians, c.j_id_a, c.ndofs_a, &mb, mass_matrices, - mm_start, lu_pivots, - dof_start, + il, ); } if c.side_b_kind == SIDE_KIND_MB && c.ndofs_b > 0 { - let mb = multibody_info.read(mb_start + c.side_b_id as usize); + let mb = multibody_info.read(il.atz(c.side_b_id as usize)); solve_mb_wj( jacobians, c.j_id_b, c.ndofs_b, &mb, mass_matrices, - mm_start, lu_pivots, - dof_start, + il, ); } c.finalize_generic_constraint(jacobians); @@ -210,8 +205,7 @@ pub fn gpu_mb_solve_impulse_joint_constraints( let joints_start = batch_ids.mb_imp_joints_start(batch_id); let cons_start = batch_ids.mb_imp_joint_constraints_start(batch_id); - let mb_start = batch_ids.mb_start(batch_id); - let dof_start = batch_ids.dof_start(batch_id); + let il = VSlice::interleaved(0, batch_ids.num_batches, batch_id); let colliders_start = batch_ids.coll_start(batch_id); // `color_groups` is a per-batch prefix-sum over the color-sorted @@ -240,16 +234,16 @@ pub fn gpu_mb_solve_impulse_joint_constraints( // Per-multibody dof base: same for every axis constraint of this joint. let dof_base_a = if builder.side_a_kind == SIDE_KIND_MB { - let mb = multibody_info.at(mb_start + builder.side_a_id as usize); - dof_start + mb.first_dof as usize + let mb = multibody_info.at(il.atz(builder.side_a_id as usize)); + VSlice::interleaved(mb.first_dof as usize, il.stride, il.shift) } else { - 0 + VSlice::dense(0) }; let dof_base_b = if builder.side_b_kind == SIDE_KIND_MB { - let mb = multibody_info.at(mb_start + builder.side_b_id as usize); - dof_start + mb.first_dof as usize + let mb = multibody_info.at(il.atz(builder.side_b_id as usize)); + VSlice::interleaved(mb.first_dof as usize, il.stride, il.shift) } else { - 0 + VSlice::dense(0) }; // TODO(PERF): load jacobians into shared memory and keep the velocity deltat on shared diff --git a/src_rbd_shaders/dynamics/multibody/impulse_joint_constraints/update.rs b/src_rbd_shaders/dynamics/multibody/impulse_joint_constraints/update.rs index 5e3dabd..d38b360 100644 --- a/src_rbd_shaders/dynamics/multibody/impulse_joint_constraints/update.rs +++ b/src_rbd_shaders/dynamics/multibody/impulse_joint_constraints/update.rs @@ -5,7 +5,7 @@ use khal_std::index::MaybeIndexUnchecked; use crate::dynamics::body::WorldMassProperties; use crate::dynamics::joint::{ANG_AXES_MASK, LIN_AXES_MASK, SPATIAL_DIM}; -use crate::utils::linalg::{MatSlice, lu_solve_in_place}; +use crate::utils::linalg::{MatSlice, lu_solve_in_place, VSlice}; use crate::{DIM, Pose}; use super::super::types::{MultibodyInfo, MultibodyLinkWorkspace}; @@ -24,9 +24,10 @@ pub(super) fn solve_mb_wj( ndofs: u32, mb: &MultibodyInfo, mass_matrices: &[f32], - mm_start: usize, lu_pivots: &[u32], - dof_start: usize, + // Interleaved dynamics-buffer view (`stride = num_batches`, `shift = + // batch_id`). + il: VSlice, ) { // Copy J into the W·J slot, then LU back-solve in place (matches the old // fused path: `wj = M⁻¹·j`). @@ -35,10 +36,15 @@ pub(super) fn solve_mb_wj( let v = jacobians.read(j_id as usize + k as usize); jacobians.write(wj_base + k as usize, v); } - let mb_mm_base = mm_start + mb.mass_matrix_offset as usize; - let m = MatSlice::dense(mb_mm_base, ndofs, ndofs); - let piv_offset = dof_start + mb.first_dof as usize; - lu_solve_in_place(mass_matrices, m, lu_pivots, piv_offset, jacobians, wj_base); + let m = MatSlice::interleaved( + mb.mass_matrix_offset as usize, + ndofs, + ndofs, + il.stride, + il.shift, + ); + let piv = VSlice::interleaved(mb.first_dof as usize, il.stride, il.shift); + lu_solve_in_place(mass_matrices, m, lu_pivots, piv, jacobians, VSlice::dense(wj_base)); } impl MbImpulseJointBuilder { @@ -50,11 +56,11 @@ impl MbImpulseJointBuilder { jacobians: &mut [f32], jac_buf_start: usize, multibody_info: &[MultibodyInfo], - mb_start: usize, links_workspace: &[MultibodyLinkWorkspace], - links_start: usize, body_jacobians: &[f32], - body_jac_start: usize, + // Interleaved dynamics-buffer view (`stride = num_batches`, `shift = + // batch_id`). + il: VSlice, poses: &[Pose], colliders_start: usize, mprops: &[WorldMassProperties], @@ -78,12 +84,12 @@ impl MbImpulseJointBuilder { // SPIR-V's "pointer to arbitrary element" restriction). Free / fixed // sides ignore the read. let mb_a = if self.side_a_kind == SIDE_KIND_MB { - multibody_info.read(mb_start + self.side_a_id as usize) + multibody_info.read(il.atz(self.side_a_id as usize)) } else { MultibodyInfo::default() }; let mb_b = if self.side_b_kind == SIDE_KIND_MB { - multibody_info.read(mb_start + self.side_b_id as usize) + multibody_info.read(il.atz(self.side_b_id as usize)) } else { MultibodyInfo::default() }; @@ -94,7 +100,7 @@ impl MbImpulseJointBuilder { self.side_a_link, &mb_a, links_workspace, - links_start, + il, poses, colliders_start, ); @@ -104,7 +110,7 @@ impl MbImpulseJointBuilder { self.side_b_link, &mb_b, links_workspace, - links_start, + il, poses, colliders_start, ); @@ -187,7 +193,7 @@ impl MbImpulseJointBuilder { j_id_a, j_id_b, body_jacobians, - body_jac_start, + il, mprops, colliders_start, ); @@ -219,7 +225,7 @@ impl MbImpulseJointBuilder { j_id_a, j_id_b, body_jacobians, - body_jac_start, + il, mprops, colliders_start, ); @@ -251,7 +257,7 @@ impl MbImpulseJointBuilder { j_id_a, j_id_b, body_jacobians, - body_jac_start, + il, mprops, colliders_start, ); @@ -283,7 +289,7 @@ impl MbImpulseJointBuilder { j_id_a, j_id_b, body_jacobians, - body_jac_start, + il, mprops, colliders_start, ); @@ -317,7 +323,7 @@ impl MbImpulseJointBuilder { j_id_a, j_id_b, body_jacobians, - body_jac_start, + il, mprops, colliders_start, ); @@ -351,7 +357,7 @@ impl MbImpulseJointBuilder { j_id_a, j_id_b, body_jacobians, - body_jac_start, + il, mprops, colliders_start, ); @@ -381,7 +387,7 @@ pub(super) fn side_world_pose( side_link: u32, mb: &MultibodyInfo, links_workspace: &[MultibodyLinkWorkspace], - links_start: usize, + il: VSlice, poses: &[Pose], colliders_start: usize, ) -> Pose { @@ -391,6 +397,6 @@ pub(super) fn side_world_pose( if side_kind == SIDE_KIND_BODY { return poses.read(colliders_start + side_id as usize); } - let link_global = links_start + mb.first_link as usize + side_link as usize; + let link_global = il.atz(mb.first_link as usize + side_link as usize); links_workspace.read(link_global).local_to_world } diff --git a/src_rbd_shaders/dynamics/multibody/integrate.rs b/src_rbd_shaders/dynamics/multibody/integrate.rs index cbc1d1b..fc2339d 100644 --- a/src_rbd_shaders/dynamics/multibody/integrate.rs +++ b/src_rbd_shaders/dynamics/multibody/integrate.rs @@ -10,7 +10,7 @@ use khal_std::macros::{spirv, spirv_bindgen}; #[cfg(feature = "dim2")] use crate::rotation_from_angle; -use crate::utils::{BatchIndices, Slice, SliceMut}; +use crate::utils::BatchIndices; use crate::{ANG_DIM, DIM}; #[cfg(feature = "dim3")] use crate::{Vector, rotation_from_scaled_axis, rotation_renormalize_fast}; @@ -43,12 +43,15 @@ pub fn gpu_mb_integrate_velocities( let dt = *dt_uniform; let mb = batch_ids - .mb_batch(batch_id, multibody_info) + .ib(batch_id, multibody_info) .read(mb_idx as usize); - let gen_base = batch_ids.dof_start(batch_id) + mb.first_dof as usize; - let mut dof_vel = SliceMut(dof_state, gen_base); - let acc = Slice(gen_accelerations, gen_base); + let mut dof_vel = batch_ids + .ib_mut(batch_id, dof_state) + .offset(mb.first_dof as usize); + let acc = batch_ids + .ib(batch_id, gen_accelerations) + .offset(mb.first_dof as usize); for d in 0..mb.ndofs { let di = d as usize; @@ -79,19 +82,22 @@ pub fn gpu_mb_integrate( let dt = *dt_uniform; let mb = batch_ids - .mb_batch(batch_id, multibody_info) + .ib(batch_id, multibody_info) .read(mb_idx as usize); let num_links = mb.num_links; - let gen_base = batch_ids.dof_start(batch_id) + mb.first_dof as usize; let stat_slice = batch_ids - .mb_links_batch(batch_id, links_static) + .ib(batch_id, links_static) .offset(mb.first_link as usize); let mut ws_slice = batch_ids - .mb_links_batch_mut(batch_id, links_workspace) + .ib_mut(batch_id, links_workspace) .offset(mb.first_link as usize); - let dof_val = SliceMut(dof_values, gen_base); - let dof_vel = Slice(dof_state, gen_base); + let dof_val = batch_ids + .ib_mut(batch_id, dof_values) + .offset(mb.first_dof as usize); + let dof_vel = batch_ids + .ib(batch_id, dof_state) + .offset(mb.first_dof as usize); // Per-link coord / joint_rot update (uses the already-corrected `dof_velocities`). // @@ -154,5 +160,5 @@ pub fn gpu_mb_integrate( // Silence dof_val unused warning — it will be used once we also support // setting coords directly (e.g. user-controlled kinematic DOFs). - let _ = dof_val.0; + let _ = dof_val.buf; } diff --git a/src_rbd_shaders/dynamics/multibody/joint_constraints.rs b/src_rbd_shaders/dynamics/multibody/joint_constraints.rs index 41938a0..0651e86 100644 --- a/src_rbd_shaders/dynamics/multibody/joint_constraints.rs +++ b/src_rbd_shaders/dynamics/multibody/joint_constraints.rs @@ -13,7 +13,7 @@ use khal_std::sync::control_barrier; use crate::dynamics::ConstraintSoftness; use crate::dynamics::joint::SPATIAL_DIM; use crate::utils::BatchIndices; -use crate::utils::linalg::{MatSlice, lu_solve_in_place}; +use crate::utils::linalg::{MatSlice, VSlice, lu_solve_in_place}; use crate::{DIM, MAX_FLT}; use super::types::{ @@ -44,7 +44,7 @@ fn lu_solve_unit( buf_m: &[f32], m: MatSlice, buf_pivots: &[u32], - pivots_offset: usize, + piv: VSlice, dst: &mut [f32], dst_offset: usize, dof_id: u32, @@ -54,7 +54,7 @@ fn lu_solve_unit( for i in 0..n { dst[dst_offset + i as usize] = if i == dof_id { 1.0 } else { 0.0 }; } - lu_solve_in_place(buf_m, m, buf_pivots, pivots_offset, dst, dst_offset); + lu_solve_in_place(buf_m, m, buf_pivots, piv, dst, VSlice::dense(dst_offset)); } /// Serially writes the metadata of every active limit/motor constraint slot. @@ -74,10 +74,10 @@ fn emit_joint_constraints( let num_links = mb.num_links; let stat_slice = batch_ids - .mb_links_batch(batch_id, links_static) + .ib(batch_id, links_static) .offset(mb.first_link as usize); let ws_slice = batch_ids - .mb_links_batch(batch_id, links_workspace) + .ib(batch_id, links_workspace) .offset(mb.first_link as usize); let inv_dt = if dt != 0.0 { 1.0 / dt } else { 0.0 }; @@ -231,7 +231,7 @@ pub fn gpu_mb_refresh_joint_constraints( } let mb = batch_ids - .mb_batch(batch_id, multibody_info) + .ib(batch_id, multibody_info) .read(mb_idx as usize); if mb.ndofs == 0 || mb.max_constraints == 0 { return; @@ -239,10 +239,10 @@ pub fn gpu_mb_refresh_joint_constraints( let cons_base = batch_ids.mb_joint_constraints_start(batch_id) + mb.first_constraint as usize; let stat_slice = batch_ids - .mb_links_batch(batch_id, links_static) + .ib(batch_id, links_static) .offset(mb.first_link as usize); let ws_slice = batch_ids - .mb_links_batch(batch_id, links_workspace) + .ib(batch_id, links_workspace) .offset(mb.first_link as usize); let dt = softness.dt; @@ -310,7 +310,7 @@ fn compute_constraint_column( mass_matrices: &[f32], m: MatSlice, lu_pivots: &[u32], - piv_offset: usize, + piv: VSlice, ) -> f32 { let _ = ndofs; let col_offset = col_base + (slot as usize) * dofs_stride; @@ -318,7 +318,7 @@ fn compute_constraint_column( mass_matrices, m, lu_pivots, - piv_offset, + piv, joint_constraint_columns, col_offset, dof_id, @@ -479,7 +479,7 @@ pub fn gpu_mb_init_joint_constraints( } let mb = batch_ids - .mb_batch(batch_id, multibody_info) + .ib(batch_id, multibody_info) .read(mb_idx as usize); let ndofs = mb.ndofs; // Uniform per workgroup: every lane of this group returns together. @@ -487,14 +487,14 @@ pub fn gpu_mb_init_joint_constraints( return; } - let mb_mm_base = batch_ids.mm_start(batch_id) + mb.mass_matrix_offset as usize; - let piv_offset = batch_ids.dof_start(batch_id) + mb.first_dof as usize; + let mb_mm_base = mb.mass_matrix_offset as usize; + let piv = batch_ids.ivec(batch_id, mb.first_dof as usize); let cons_base = batch_ids.mb_joint_constraints_start(batch_id) + mb.first_constraint as usize; // One column of M⁻¹ per constraint slot . let dofs_stride = batch_ids.dof_batch_capacity as usize; let col_base = batch_ids.mb_joint_constraint_columns_start(batch_id) + (mb.first_constraint as usize) * dofs_stride; - let m = MatSlice::dense(mb_mm_base, ndofs, ndofs); + let m = batch_ids.imat(batch_id, mb_mm_base, ndofs, ndofs); // Stage 1: lane-parallel slot reset. for s in StepRng::new(lane..mb.max_constraints, LANES) { @@ -554,7 +554,7 @@ pub fn gpu_mb_init_joint_constraints( mass_matrices, m, lu_pivots, - piv_offset, + piv, ); let cfm_gain = lhs * cons.cfm_coeff + cons.cfm_gain; cons.cfm_gain = cfm_gain; diff --git a/src_rbd_shaders/dynamics/multibody/lu.rs b/src_rbd_shaders/dynamics/multibody/lu.rs index da41437..f6e01a8 100644 --- a/src_rbd_shaders/dynamics/multibody/lu.rs +++ b/src_rbd_shaders/dynamics/multibody/lu.rs @@ -7,7 +7,7 @@ use khal_std::index::MaybeIndexUnchecked; use khal_std::sync::workgroup_memory_barrier_with_group_sync; -use crate::utils::linalg::MAX_MB_DOFS; +use crate::utils::linalg::{MAX_MB_DOFS, VSlice}; /// Workgroup width for the parallelised LU kernels. Must match the /// `threads(N, 1, 1)` attribute and `MB_LU_LANES` on the host side. @@ -31,7 +31,7 @@ pub(super) fn lu_factor_in_shared( lane: u32, mat: &mut [f32; MAX_MB_DOFS * MAX_MB_DOFS], pivots_dst: &mut [u32], - pivots_offset: usize, + piv: VSlice, pivot_row_shared: &mut u32, inv_akk_shared: &mut f32, ) { @@ -52,7 +52,7 @@ pub(super) fn lu_factor_in_shared( } } *pivot_row_shared = pivot_row; - pivots_dst.write(pivots_offset + k as usize, pivot_row); + pivots_dst.write(piv.at(k), pivot_row); } workgroup_memory_barrier_with_group_sync(); let pivot_row = *pivot_row_shared; @@ -190,7 +190,7 @@ pub(super) fn lu_factor_in_shared_packed( lane: u32, active_slot: bool, buf_pivots: &[u32], - pivots_offset: usize, + piv: VSlice, x: &mut [f32; 64], ) { let seg = (slot * T) as usize; if active_slot && lane == 0 { for k in 0..n { - let p = buf_pivots.read(pivots_offset + k as usize); + let p = buf_pivots.read(piv.at(k)); if p != k { let a = x.read(seg + k as usize); let b = x.read(seg + p as usize); @@ -361,12 +361,12 @@ pub(super) fn lu_apply_pivots( n: u32, lane: u32, buf_pivots: &[u32], - pivots_offset: usize, + piv: VSlice, x: &mut [f32; MAX_MB_DOFS], ) { if lane == 0 { for k in 0..n { - let p = buf_pivots.read(pivots_offset + k as usize); + let p = buf_pivots.read(piv.at(k)); if p != k { let a = x.read(k as usize); let b = x.read(p as usize); diff --git a/src_rbd_shaders/dynamics/multibody/solve_constraints.rs b/src_rbd_shaders/dynamics/multibody/solve_constraints.rs index d47ac96..49f26ff 100644 --- a/src_rbd_shaders/dynamics/multibody/solve_constraints.rs +++ b/src_rbd_shaders/dynamics/multibody/solve_constraints.rs @@ -53,8 +53,7 @@ pub fn gpu_mb_solve_constraints( return; } - let mb_start = batch_ids.mb_start(batch_id); - let mb = multibody_info.read(mb_start + mb_idx as usize); + let mb = multibody_info.read(batch_ids.mbi(batch_id, mb_idx as usize)); let ndofs = mb.ndofs; // Uniform per workgroup: every lane of this group returns together. if ndofs == 0 { @@ -62,7 +61,7 @@ pub fn gpu_mb_solve_constraints( } let use_bias = *use_bias != 0; - let v_base = batch_ids.dof_start(batch_id) + mb.first_dof as usize; + let v_base = mb.first_dof as usize; let dofs_stride = batch_ids.dof_batch_capacity as usize; let colliders_start = batch_ids.coll_start(batch_id); @@ -85,7 +84,7 @@ pub fn gpu_mb_solve_constraints( // Load the generalized velocities and accumulated contact impulses into // workgroup memory. if lane < ndofs { - dof_v[lane as usize] = dof_state.read(v_base + lane as usize); + dof_v[lane as usize] = dof_state.read(batch_ids.mbi(batch_id, v_base + lane as usize)); } for s in StepRng::new(lane..contact_count, LANES) { imp_shared[s as usize] = contact_constraints.read(ccons_base + s as usize).impulse; @@ -209,7 +208,10 @@ pub fn gpu_mb_solve_constraints( // Writeback if lane < ndofs { - dof_state.write(v_base + lane as usize, dof_v[lane as usize]); + dof_state.write( + batch_ids.mbi(batch_id, v_base + lane as usize), + dof_v[lane as usize], + ); } for s in StepRng::new(lane..contact_count, LANES) { let mut cons = contact_constraints.read(ccons_base + s as usize); @@ -243,8 +245,7 @@ pub fn gpu_mb_solve_joints( return; } - let mb_start = batch_ids.mb_start(batch_id); - let mb = multibody_info.read(mb_start + mb_idx as usize); + let mb = multibody_info.read(batch_ids.mbi(batch_id, mb_idx as usize)); let ndofs = mb.ndofs; // Uniform per workgroup: every lane of this group returns together. if ndofs == 0 || mb.max_constraints == 0 { @@ -252,7 +253,7 @@ pub fn gpu_mb_solve_joints( } let use_bias = *use_bias != 0; - let v_base = batch_ids.dof_start(batch_id) + mb.first_dof as usize; + let v_base = mb.first_dof as usize; let dofs_stride = batch_ids.dof_batch_capacity as usize; let jcons_base = batch_ids.mb_joint_constraints_start(batch_id) + mb.first_constraint as usize; @@ -260,7 +261,7 @@ pub fn gpu_mb_solve_joints( + (mb.first_constraint as usize) * dofs_stride; if lane < ndofs { - dof_v[lane as usize] = dof_state.read(v_base + lane as usize); + dof_v[lane as usize] = dof_state.read(batch_ids.mbi(batch_id, v_base + lane as usize)); } workgroup_memory_barrier_with_group_sync(); @@ -301,7 +302,10 @@ pub fn gpu_mb_solve_joints( } if lane < ndofs { - dof_state.write(v_base + lane as usize, dof_v[lane as usize]); + dof_state.write( + batch_ids.mbi(batch_id, v_base + lane as usize), + dof_v[lane as usize], + ); } } @@ -333,8 +337,7 @@ pub fn gpu_mb_build_contact_delassus( return; } - let mb_start = batch_ids.mb_start(batch_id); - let mb = multibody_info.read(mb_start + mb_idx as usize); + let mb = multibody_info.read(batch_ids.mbi(batch_id, mb_idx as usize)); let ndofs = mb.ndofs; let count = mb.contact_constraint_count; if ndofs == 0 || count == 0 { @@ -419,8 +422,7 @@ pub fn gpu_mb_solve_contacts_delassus( return; } - let mb_start = batch_ids.mb_start(batch_id); - let mb = multibody_info.read(mb_start + mb_idx as usize); + let mb = multibody_info.read(batch_ids.mbi(batch_id, mb_idx as usize)); let ndofs = mb.ndofs; let count = mb.contact_constraint_count; // Uniform per workgroup: every lane of this group returns together. @@ -429,7 +431,7 @@ pub fn gpu_mb_solve_contacts_delassus( } let use_bias = *use_bias != 0; - let v_base = batch_ids.dof_start(batch_id) + mb.first_dof as usize; + let v_base = mb.first_dof as usize; let colliders_start = batch_ids.coll_start(batch_id); let cons_base = batch_ids.mb_contact_constraints_start(batch_id) + (mb_idx as usize) * (MAXC as usize); @@ -441,7 +443,7 @@ pub fn gpu_mb_solve_contacts_delassus( * (MAXC as usize); if lane < ndofs { - dof_v[lane as usize] = dof_state.read(v_base + lane as usize); + dof_v[lane as usize] = dof_state.read(batch_ids.mbi(batch_id, v_base + lane as usize)); } // Preload the per-constraint solve scalars into shared SoA arrays so the @@ -537,7 +539,10 @@ pub fn gpu_mb_solve_contacts_delassus( // Writeback. if lane < ndofs { - dof_state.write(v_base + lane as usize, dof_v[lane as usize]); + dof_state.write( + batch_ids.mbi(batch_id, v_base + lane as usize), + dof_v[lane as usize], + ); } for s in StepRng::new(lane..count, LANES) { let mut cons = contact_constraints.read(cons_base + s as usize); diff --git a/src_rbd_shaders/utils/indices.rs b/src_rbd_shaders/utils/indices.rs index 24e7996..f182e39 100644 --- a/src_rbd_shaders/utils/indices.rs +++ b/src_rbd_shaders/utils/indices.rs @@ -1,4 +1,5 @@ -use crate::utils::{Slice, SliceMut}; +use crate::utils::linalg::{MatSlice, VSlice}; +use crate::utils::{ISlice, ISliceMut, Slice, SliceMut}; /// Per-batch capacities and packed-buffer section offsets, shared by every /// kernel that needs to slice a flat tensor into its batch's slot. @@ -91,54 +92,60 @@ impl BatchIndices { batch_id as usize * self.colliders_batch_capacity as usize } + /// Interleaved flat index for the multibody dynamics buffers. #[inline] - pub fn collision_pairs_start(&self, batch_id: u32) -> usize { - batch_id as usize * self.collision_pairs_batch_capacity as usize - } - - #[inline] - pub fn contacts_start(&self, batch_id: u32) -> usize { - batch_id as usize * self.contacts_batch_capacity as usize - } - - #[inline] - pub fn impulse_joints_start(&self, batch_id: u32) -> usize { - batch_id as usize * self.impulse_joints_batch_capacity as usize + pub fn mbi(&self, batch_id: u32, intra: usize) -> usize { + intra * self.num_batches as usize + batch_id as usize } + /// Interleaved view of a multibody dynamics buffer for batch `batch_id` + /// (use `.offset(...)` for the intra-batch element offset). #[inline] - pub fn mb_start(&self, batch_id: u32) -> usize { - batch_id as usize * self.multibodies_batch_capacity as usize + pub fn ib<'s, T>(&self, batch_id: u32, slice: &'s [T]) -> ISlice<'s, T> { + ISlice { + buf: slice, + base: 0, + stride: self.num_batches, + shift: batch_id, + } } + /// Mutable interleaved view. #[inline] - pub fn links_start(&self, batch_id: u32) -> usize { - batch_id as usize * self.links_batch_capacity as usize + pub fn ib_mut<'s, T>(&self, batch_id: u32, slice: &'s mut [T]) -> ISliceMut<'s, T> { + ISliceMut { + buf: slice, + base: 0, + stride: self.num_batches, + shift: batch_id, + } } + /// Interleaved dense matrix view at intra-batch element offset `offset`. #[inline] - pub fn jac_start(&self, batch_id: u32) -> usize { - batch_id as usize * self.jacobians_batch_capacity as usize + pub fn imat(&self, batch_id: u32, offset: usize, rows: u32, cols: u32) -> MatSlice { + MatSlice::interleaved(offset, rows, cols, self.num_batches, batch_id) } + /// Interleaved vector view at intra-batch element offset `offset`. #[inline] - pub fn mm_start(&self, batch_id: u32) -> usize { - batch_id as usize * self.mass_matrix_batch_capacity as usize + pub fn ivec(&self, batch_id: u32, offset: usize) -> VSlice { + VSlice::interleaved(offset, self.num_batches, batch_id) } #[inline] - pub fn cor_start(&self, batch_id: u32) -> usize { - batch_id as usize * self.coriolis_batch_capacity as usize + pub fn collision_pairs_start(&self, batch_id: u32) -> usize { + batch_id as usize * self.collision_pairs_batch_capacity as usize } #[inline] - pub fn icdt_start(&self, batch_id: u32) -> usize { - batch_id as usize * self.i_coriolis_dt_batch_capacity as usize + pub fn contacts_start(&self, batch_id: u32) -> usize { + batch_id as usize * self.contacts_batch_capacity as usize } #[inline] - pub fn dof_start(&self, batch_id: u32) -> usize { - batch_id as usize * self.dof_batch_capacity as usize + pub fn impulse_joints_start(&self, batch_id: u32) -> usize { + batch_id as usize * self.impulse_joints_batch_capacity as usize } #[inline] @@ -233,36 +240,6 @@ impl BatchIndices { SliceMut(slice, self.impulse_joints_start(batch_id)) } - #[inline] - pub fn mb_batch<'s, T>(&self, batch_id: u32, slice: &'s [T]) -> Slice<'s, T> { - Slice(slice, self.mb_start(batch_id)) - } - - #[inline] - pub fn mb_batch_mut<'s, T>(&self, batch_id: u32, slice: &'s mut [T]) -> SliceMut<'s, T> { - SliceMut(slice, self.mb_start(batch_id)) - } - - #[inline] - pub fn mb_links_batch<'s, T>(&self, batch_id: u32, slice: &'s [T]) -> Slice<'s, T> { - Slice(slice, self.links_start(batch_id)) - } - - #[inline] - pub fn mb_links_batch_mut<'s, T>(&self, batch_id: u32, slice: &'s mut [T]) -> SliceMut<'s, T> { - SliceMut(slice, self.links_start(batch_id)) - } - - #[inline] - pub fn dof_batch<'s, T>(&self, batch_id: u32, slice: &'s [T]) -> Slice<'s, T> { - Slice(slice, self.dof_start(batch_id)) - } - - #[inline] - pub fn dof_batch_mut<'s, T>(&self, batch_id: u32, slice: &'s mut [T]) -> SliceMut<'s, T> { - SliceMut(slice, self.dof_start(batch_id)) - } - #[inline] pub fn mb_joint_constraints_batch<'s, T>(&self, batch_id: u32, slice: &'s [T]) -> Slice<'s, T> { Slice(slice, self.mb_joint_constraints_start(batch_id)) diff --git a/src_rbd_shaders/utils/linalg.rs b/src_rbd_shaders/utils/linalg.rs index 7918699..d206885 100644 --- a/src_rbd_shaders/utils/linalg.rs +++ b/src_rbd_shaders/utils/linalg.rs @@ -22,14 +22,22 @@ pub const MAX_MB_DOFS: usize = 64; /// A column-major matrix view into a flat f32 buffer. #[derive(Copy, Clone)] pub struct MatSlice { - /// Offset (in f32 entries) of the (0, 0) element inside the backing buffer. + /// Offset (in element entries, multiplied by `stride`) of the (0, 0) + /// element inside the backing buffer. pub offset: usize, /// Number of rows. pub rows: u32, /// Number of columns. pub cols: u32, - /// Leading dimension — distance between columns, in f32 entries. + /// Leading dimension: distance between columns, in element entries. pub lead: u32, + /// Element stride: flat index = `(offset + c·lead + r) · stride + shift`. + /// `1` for plain dense storage; `num_batches` for the batch-interleaved + /// multibody dynamics buffers (with `shift = batch_id`). + pub stride: u32, + /// Additive shift applied after the stride (the batch id for + /// batch-interleaved buffers). + pub shift: u32, } impl MatSlice { @@ -41,16 +49,34 @@ impl MatSlice { rows, cols, lead: rows, + stride: 1, + shift: 0, + } + } + + /// Batch-interleaved dense view: element `k` (intra-batch offset + /// `offset + k`) of batch `shift` lives at `(offset + k) · stride + + /// shift`. + #[inline] + pub fn interleaved(offset: usize, rows: u32, cols: u32, stride: u32, shift: u32) -> Self { + Self { + offset, + rows, + cols, + lead: rows, + stride, + shift, } } /// Flat index of element `(r, c)`. #[inline] pub fn idx(&self, r: u32, c: u32) -> usize { - self.offset + (c * self.lead + r) as usize + (self.offset + (c * self.lead + r) as usize) * self.stride as usize + self.shift as usize } - /// Sub-view starting at `(r0, c0)` with shape `(nr × nc)`. Inherits `lead`. + /// Sub-view starting at `(r0, c0)` with shape `(nr × nc)`. Inherits + /// `lead`, `stride` and `shift`. #[inline] pub fn view(&self, r0: u32, c0: u32, nr: u32, nc: u32) -> Self { Self { @@ -58,6 +84,8 @@ impl MatSlice { rows: nr, cols: nc, lead: self.lead, + stride: self.stride, + shift: self.shift, } } @@ -80,6 +108,51 @@ impl MatSlice { } } +/// Strided vector view: element `i` lives at `(offset + i) · stride + shift`. +/// The vector counterpart of [`MatSlice`], used for the LU pivot / rhs +/// buffers which may be batch-interleaved (`stride = num_batches`, `shift = +/// batch_id`) or plain dense (`stride = 1, shift = 0`). +#[derive(Copy, Clone)] +pub struct VSlice { + pub offset: usize, + pub stride: u32, + pub shift: u32, +} + +impl VSlice { + /// Plain dense view at `offset`. + #[inline] + pub fn dense(offset: usize) -> Self { + Self { + offset, + stride: 1, + shift: 0, + } + } + + /// Batch-interleaved view. + #[inline] + pub fn interleaved(offset: usize, stride: u32, shift: u32) -> Self { + Self { + offset, + stride, + shift, + } + } + + /// Flat index of element `i`. + #[inline] + pub fn at(&self, i: u32) -> usize { + (self.offset + i as usize) * self.stride as usize + self.shift as usize + } + + /// Flat index of element `i` (usize form). + #[inline] + pub fn atz(&self, i: usize) -> usize { + (self.offset + i) * self.stride as usize + self.shift as usize + } +} + /// `m := val` (element-wise). #[inline] pub fn fill(buf: &mut [f32], m: MatSlice, val: f32) { @@ -266,7 +339,7 @@ pub fn gemv_tr_spatial( #[inline] pub fn gemv_tr_spatial_split( buf_y: &mut [f32], - y_offset: usize, + y: VSlice, alpha: f32, buf_a: &[f32], a: MatSlice, @@ -281,7 +354,7 @@ pub fn gemv_tr_spatial_split( + buf_a.read(a.idx(3, c)) * x_ang.x + buf_a.read(a.idx(4, c)) * x_ang.y + buf_a.read(a.idx(5, c)) * x_ang.z; - let idx = y_offset + c as usize; + let idx = y.at(c); let cur = buf_y.read(idx); buf_y.write(idx, beta * cur + alpha * s); } @@ -291,7 +364,7 @@ pub fn gemv_tr_spatial_split( #[inline] pub fn gemv_tr_spatial_split( buf_y: &mut [f32], - y_offset: usize, + y: VSlice, alpha: f32, buf_a: &[f32], a: MatSlice, @@ -303,7 +376,7 @@ pub fn gemv_tr_spatial_split( let s = buf_a.read(a.idx(0, c)) * x_lin.x + buf_a.read(a.idx(1, c)) * x_lin.y + buf_a.read(a.idx(2, c)) * x_ang; - let idx = y_offset + c as usize; + let idx = y.at(c); let cur = buf_y.read(idx); buf_y.write(idx, beta * cur + alpha * s); } @@ -622,7 +695,7 @@ pub fn gemm_tr( /// `pivots[k]` is the row that was swapped with row `k` during elimination step /// `k`. `pivots_offset` is where this multibody's pivot slot starts in `buf_pivots`. #[inline] -pub fn lu_decompose(buf_m: &mut [f32], m: MatSlice, buf_pivots: &mut [u32], pivots_offset: usize) { +pub fn lu_decompose(buf_m: &mut [f32], m: MatSlice, buf_pivots: &mut [u32], piv: VSlice) { let n = m.rows; for k in 0..n { // Partial pivot: find max |M[i, k]| for i in k..n. @@ -639,7 +712,7 @@ pub fn lu_decompose(buf_m: &mut [f32], m: MatSlice, buf_pivots: &mut [u32], pivo pivot_row = i; } } - buf_pivots.write(pivots_offset + k as usize, pivot_row); + buf_pivots.write(piv.at(k), pivot_row); // Row swap k ↔ pivot_row (full row since we haven't computed past col k). if pivot_row != k { @@ -681,18 +754,18 @@ pub fn lu_solve_in_place( buf_m: &[f32], m: MatSlice, buf_pivots: &[u32], - pivots_offset: usize, + piv: VSlice, buf_rhs: &mut [f32], - rhs_offset: usize, + rhs: VSlice, ) { let n = m.rows; // Permute rhs in place according to the recorded pivots. for k in 0..n { - let p = buf_pivots.read(pivots_offset + k as usize); + let p = buf_pivots.read(piv.at(k)); if p != k { - let ki = rhs_offset + k as usize; - let pi = rhs_offset + p as usize; + let ki = rhs.at(k); + let pi = rhs.at(p); let a = buf_rhs.read(ki); let b = buf_rhs.read(pi); buf_rhs.write(ki, b); @@ -702,22 +775,22 @@ pub fn lu_solve_in_place( // Forward substitution: L · y = P · rhs (L is unit-lower — implicit diag = 1). for i in 0..n { - let mut s = buf_rhs.read(rhs_offset + i as usize); + let mut s = buf_rhs.read(rhs.at(i)); for j in 0..i { - s -= buf_m.read(m.idx(i, j)) * buf_rhs.read(rhs_offset + j as usize); + s -= buf_m.read(m.idx(i, j)) * buf_rhs.read(rhs.at(j)); } - buf_rhs.write(rhs_offset + i as usize, s); + buf_rhs.write(rhs.at(i), s); } // Back substitution: U · x = y (reverse iteration — equivalent to `for ii in (0..n).rev()`). for step in 0..n { let ii = n - 1 - step; - let mut s = buf_rhs.read(rhs_offset + ii as usize); + let mut s = buf_rhs.read(rhs.at(ii)); for j in (ii + 1)..n { - s -= buf_m.read(m.idx(ii, j)) * buf_rhs.read(rhs_offset + j as usize); + s -= buf_m.read(m.idx(ii, j)) * buf_rhs.read(rhs.at(j)); } let u = buf_m.read(m.idx(ii, ii)); - buf_rhs.write(rhs_offset + ii as usize, if u != 0.0 { s / u } else { 0.0 }); + buf_rhs.write(rhs.at(ii), if u != 0.0 { s / u } else { 0.0 }); } } @@ -1189,7 +1262,7 @@ pub fn gemm_omega_skew_tr_cross_buf_par( #[inline] pub fn gemv_tr_spatial_split_par( buf_y: &mut [f32], - y_offset: usize, + y: VSlice, alpha: f32, buf_a: &[f32], a: MatSlice, @@ -1207,7 +1280,7 @@ pub fn gemv_tr_spatial_split_par( + buf_a.read(a.idx(3, c)) * x_ang.x + buf_a.read(a.idx(4, c)) * x_ang.y + buf_a.read(a.idx(5, c)) * x_ang.z; - let idx = y_offset + c as usize; + let idx = y.at(c); let cur = buf_y.read(idx); buf_y.write(idx, beta * cur + alpha * s); } @@ -1217,7 +1290,7 @@ pub fn gemv_tr_spatial_split_par( #[inline] pub fn gemv_tr_spatial_split_par( buf_y: &mut [f32], - y_offset: usize, + y: VSlice, alpha: f32, buf_a: &[f32], a: MatSlice, @@ -1232,7 +1305,7 @@ pub fn gemv_tr_spatial_split_par( let s = buf_a.read(a.idx(0, c)) * x_lin.x + buf_a.read(a.idx(1, c)) * x_lin.y + buf_a.read(a.idx(2, c)) * x_ang; - let idx = y_offset + c as usize; + let idx = y.at(c); let cur = buf_y.read(idx); buf_y.write(idx, beta * cur + alpha * s); } diff --git a/src_rbd_shaders/utils/mod.rs b/src_rbd_shaders/utils/mod.rs index e28d2f8..a9916b5 100644 --- a/src_rbd_shaders/utils/mod.rs +++ b/src_rbd_shaders/utils/mod.rs @@ -9,7 +9,7 @@ mod slice; pub use basis::orthonormal_basis3; pub use indices::BatchIndices; -pub use slice::{Slice, SliceMut}; +pub use slice::{ISlice, ISliceMut, Slice, SliceMut}; /// Division with ceiling (signed). pub fn div_ceil(x: i32, y: i32) -> i32 { diff --git a/src_rbd_shaders/utils/slice.rs b/src_rbd_shaders/utils/slice.rs index 4c0b2cd..121e226 100644 --- a/src_rbd_shaders/utils/slice.rs +++ b/src_rbd_shaders/utils/slice.rs @@ -80,3 +80,115 @@ impl IndexMut for SliceMut<'_, T> { self.0.at_mut(self.1 + i) } } + +// Batch-interleaved slice: element `i` of the view lives at +// `(base + i) · stride + shift` in the backing buffer. +pub struct ISlice<'a, T> { + pub buf: &'a [T], + pub base: usize, + pub stride: u32, + pub shift: u32, +} + +impl<'a, T: Copy> ISlice<'a, T> { + #[inline(always)] + fn flat(&self, i: usize) -> usize { + (self.base + i) * self.stride as usize + self.shift as usize + } + + #[inline] + pub fn at(&self, i: usize) -> &'a T { + self.buf.at(self.flat(i)) + } + + #[inline] + pub fn read(&self, i: usize) -> T { + self.buf.read(self.flat(i)) + } + + #[inline] + pub fn offset(self, offset: usize) -> Self { + ISlice { + base: self.base + offset, + ..self + } + } +} + +pub struct ISliceMut<'a, T> { + pub buf: &'a mut [T], + pub base: usize, + pub stride: u32, + pub shift: u32, +} + +impl<'a, T: Copy> ISliceMut<'a, T> { + #[inline(always)] + fn flat(&self, i: usize) -> usize { + (self.base + i) * self.stride as usize + self.shift as usize + } + + #[inline] + pub fn as_ref(&self) -> ISlice<'_, T> { + ISlice { + buf: &*self.buf, + base: self.base, + stride: self.stride, + shift: self.shift, + } + } + + #[inline] + pub fn at(&self, i: usize) -> &T { + self.buf.at(self.flat(i)) + } + + #[inline] + pub fn read(&self, i: usize) -> T { + self.buf.read(self.flat(i)) + } + + #[inline] + pub fn at_mut(&mut self, i: usize) -> &mut T { + let idx = self.flat(i); + self.buf.at_mut(idx) + } + + #[inline] + pub fn write(&mut self, i: usize, value: T) { + let idx = self.flat(i); + self.buf.write(idx, value) + } + + #[inline] + pub fn offset(self, offset: usize) -> Self { + ISliceMut { + base: self.base + offset, + ..self + } + } +} + +impl Index for ISlice<'_, T> { + type Output = T; + #[inline(always)] + fn index(&self, i: usize) -> &T { + self.buf.at(self.flat(i)) + } +} + +impl Index for ISliceMut<'_, T> { + type Output = T; + #[inline(always)] + fn index(&self, i: usize) -> &T { + self.buf.at(self.flat(i)) + } +} + +impl IndexMut for ISliceMut<'_, T> { + #[inline(always)] + fn index_mut(&mut self, i: usize) -> &mut T { + let idx = self.flat(i); + self.buf.at_mut(idx) + } +} From 1b953f381559ce4afef28a290241ad13c6a0a95c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Sat, 1 Aug 2026 15:17:30 +0200 Subject: [PATCH 37/39] perf: SoA (vec4-quad) layout for the multibody link workspace --- .../multibody/multibody_from_rapier.rs | 8 +- src_rbd/dynamics/multibody/multibody_set.rs | 19 +- .../multibody/compute_dynamics_pre.rs | 161 ++++---- .../dynamics/multibody/gravity_and_lu.rs | 113 +++--- .../impulse_joint_constraints/kernels.rs | 5 +- .../impulse_joint_constraints/update.rs | 12 +- .../dynamics/multibody/integrate.rs | 54 ++- .../dynamics/multibody/joint_constraints.rs | 25 +- src_rbd_shaders/dynamics/multibody/mod.rs | 2 + src_rbd_shaders/dynamics/multibody/ws_soa.rs | 373 ++++++++++++++++++ 10 files changed, 575 insertions(+), 197 deletions(-) create mode 100644 src_rbd_shaders/dynamics/multibody/ws_soa.rs diff --git a/src_rbd/dynamics/multibody/multibody_from_rapier.rs b/src_rbd/dynamics/multibody/multibody_from_rapier.rs index 61cc4d9..3a955a7 100644 --- a/src_rbd/dynamics/multibody/multibody_from_rapier.rs +++ b/src_rbd/dynamics/multibody/multibody_from_rapier.rs @@ -388,7 +388,6 @@ impl GpuMultibodySet { let nb = num_batches as usize; let all_infos = interleave(&all_infos, mb_cap, nb); let all_statics = interleave(&all_statics, links_cap, nb); - let all_ws = interleave(&all_ws, links_cap, nb); let all_dof_vals = interleave(&all_dof_vals, dofs_cap, nb); let all_dof_vels = interleave(&all_dof_vels, dofs_cap, nb); let all_dof_damping = interleave(&all_dof_damping, dofs_cap, nb); @@ -411,7 +410,12 @@ impl GpuMultibodySet { links_static: Tensor::vector(backend, &all_statics, storage | BufferUsages::COPY_DST) .unwrap(), links_static_mirror: all_statics.clone(), - links_workspace: Tensor::vector(backend, &all_ws, storage).unwrap(), + links_workspace: Tensor::vector( + backend, + &crate::shaders::dynamics::ws_soa_from_structs(&all_ws, links_cap, num_batches), + storage, + ) + .unwrap(), dof_values: Tensor::vector(backend, &all_dof_vals, storage).unwrap(), dof_state: { // Pack [velocities (N), damping (N), armature (N)] back-to-back diff --git a/src_rbd/dynamics/multibody/multibody_set.rs b/src_rbd/dynamics/multibody/multibody_set.rs index e041016..7023a03 100644 --- a/src_rbd/dynamics/multibody/multibody_set.rs +++ b/src_rbd/dynamics/multibody/multibody_set.rs @@ -60,8 +60,8 @@ pub struct GpuMultibodySet { /// CPU-side mirror of [`Self::links_static`] used to support runtime /// mutations like motor changes without round-tripping through a GPU read. pub(super) links_static_mirror: Vec, - /// Per-batch per-step link workspace. - pub(super) links_workspace: Tensor, + /// Per-batch per-step link workspace, SoA quad layout. + pub(super) links_workspace: Tensor, /// Generalized coordinates (flat). pub(super) dof_values: Tensor, /// Packed buffer holding generalized velocities (offset 0) and per-DOF @@ -171,21 +171,10 @@ impl GpuMultibodySet { } /// Lanes per multibody for the packed per-multibody dynamics kernels - /// (`compute_dynamics_pre`, `gravity_and_lu`) — mirrored into - /// `BatchIndices::mb_pack_lanes`. - /// - /// `1` selects the SERIAL tier: one thread runs its - /// multibody's whole FK/CRBA/LU chain with no barriers at all, 64 - /// multibodies per workgroup with every lane busy. For small robots this - /// beats lane-parallelism — whose ~60-barrier dependency chain caps how - /// fast one multibody can finish — but ONLY once there are enough - /// multibodies for the thread count to hide the long serial chain's - /// latency (measured crossover between 1024 and 4096 on Apple M-series; - /// below that, spreading each robot across 8 lanes wins despite the - /// barriers). + /// (`compute_dynamics_pre`, `gravity_and_lu`). pub(crate) fn pack_lanes(&self) -> u32 { let total_mb = self.num_active_multibodies * self.num_batches; - if self.max_ndofs <= 8 && total_mb >= 2048 { + if self.max_ndofs <= 8 && total_mb >= 1024 { 1 } else { self.max_ndofs.next_power_of_two().clamp(8, MB_LU_LANES) diff --git a/src_rbd_shaders/dynamics/multibody/compute_dynamics_pre.rs b/src_rbd_shaders/dynamics/multibody/compute_dynamics_pre.rs index 89a234c..3534409 100644 --- a/src_rbd_shaders/dynamics/multibody/compute_dynamics_pre.rs +++ b/src_rbd_shaders/dynamics/multibody/compute_dynamics_pre.rs @@ -9,21 +9,27 @@ //! (gravity rhs + LU factor + LU solve). use khal_std::glamx::UVec3; +use glamx::Vec4; use khal_std::index::MaybeIndexUnchecked; use khal_std::macros::{spirv, spirv_bindgen}; use khal_std::sync::workgroup_memory_barrier_with_group_sync; -use super::types::{MultibodyInfo, MultibodyLinkStatic, MultibodyLinkWorkspace}; +use super::types::{MultibodyInfo, MultibodyLinkStatic}; +use super::ws_soa::{ + WS_JOINT_ROT, WS_JOINT_VEL, WS_LTP, WS_LTW, WS_RB_VELS, WS_SHIFT02, WS_SHIFT23, WsAddr, + ws_coords, ws_pose, ws_rot, ws_set_pose, ws_set_vec, ws_set_vel, ws_vec, ws_vel, ws_vel_ang, + ws_world_inertia, +}; use crate::dynamics::body::Velocity; use crate::dynamics::joint::SPATIAL_DIM; #[cfg(feature = "dim3")] use crate::utils::linalg::gemm_skew_lhs_cross_buf_par; use crate::utils::linalg::{ - MatSlice, copy_from_par, fill_par, gemm_inertia_lhs_par, + copy_from_par, fill_par, gemm_inertia_lhs_par, gemm_omega_skew_tr_cross_buf_par, gemm_skew_tr_lhs_cross_buf_par, gemm_skew_tr_lhs_par, gemm_tr_par, quadform_spatial_par, }; -use crate::utils::{BatchIndices, ISlice, ISliceMut, SliceMut}; +use crate::utils::{BatchIndices, ISlice, SliceMut}; use crate::{ANG_DIM, AngVector, DIM, Pose, Vector, gcross_av}; use parry::math::VectorExt; @@ -63,7 +69,7 @@ pub fn gpu_mb_compute_dynamics_pre( #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] links_static: &[MultibodyLinkStatic], #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] - links_workspace: &mut [MultibodyLinkWorkspace], + links_workspace: &mut [Vec4], #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] poses: &mut [Pose], #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] body_jacobians: &mut [f32], #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] mass_matrices: &mut [f32], @@ -96,9 +102,7 @@ pub fn gpu_mb_compute_dynamics_pre( let stat_slice = batch_ids .ib(batch_id, links_static) .offset(mb.first_link as usize); - let mut ws_slice = batch_ids - .ib_mut(batch_id, links_workspace) - .offset(mb.first_link as usize); + let wa = WsAddr::new(mb.first_link as usize, batch_ids.num_batches, batch_id); let mut poses_slice = batch_ids.coll_batch_mut(batch_id, poses); let damping_slice = batch_ids .ib(batch_id, dof_state) @@ -113,7 +117,7 @@ pub fn gpu_mb_compute_dynamics_pre( // 1) Forward Kinematics (single-threaded) if active_slot && num_links > 0 && lane == 0 { - forward_kinematics(&mb, &stat_slice, &mut poses_slice, &mut ws_slice, num_links); + forward_kinematics(&mb, &stat_slice, &mut poses_slice, links_workspace, wa, num_links); } sync_slots(t); @@ -126,7 +130,8 @@ pub fn gpu_mb_compute_dynamics_pre( num_links, batch_ids.mb_max_links, &stat_slice, - &ws_slice.as_ref(), + links_workspace, + wa, body_jacobians, batch_ids, batch_id, @@ -134,7 +139,7 @@ pub fn gpu_mb_compute_dynamics_pre( // 3) Propagate velocities (single-threaded) if active_slot && num_links > 0 && lane == 0 { - propagate_velocities(num_links, &stat_slice, &vel_slice, &mut ws_slice); + propagate_velocities(num_links, &stat_slice, &vel_slice, links_workspace, wa); } sync_slots(t); @@ -202,14 +207,13 @@ pub fn gpu_mb_compute_dynamics_pre( let mut rb_inertia = Default::default(); if loop_is_active { - let ws = &ws_slice[k as usize]; let lmp = stat_slice[k as usize].local_mprops; mass = 1.0 / inv_mass_x; - rb_inertia = ws.link_world_inertia(&lmp); + rb_inertia = ws_world_inertia(links_workspace, wa, k, &lmp); #[cfg(feature = "dim3")] let augmented_inertia = { - let angvel = ws.rb_vels.angular; + let angvel = ws_vel_ang(links_workspace, wa, k, WS_RB_VELS); let w_skew = crate::utils::linalg::skew(angvel); let i_omega = rb_inertia * angvel; let i_omega_skew = crate::utils::linalg::skew(i_omega); @@ -235,7 +239,6 @@ pub fn gpu_mb_compute_dynamics_pre( if k != 0 { let stat = stat_slice[k as usize]; let parent_id = stat.parent_link_id; - let parent_link = &ws_slice[parent_id as usize]; let parent_j = batch_ids.imat(batch_id, mb_jac_base + (parent_id as usize) * SPATIAL_DIM * (ndofs as usize), SPATIAL_DIM as u32, @@ -252,7 +255,9 @@ pub fn gpu_mb_compute_dynamics_pre( ANG_DIM, ndofs, ); - let parent_w = parent_link.rb_vels.angular; + let parent_w = ws_vel_ang(links_workspace, wa, parent_id, WS_RB_VELS); + let ws_shift02 = ws_vec(links_workspace, wa, k, WS_SHIFT02); + let ws_joint_vel = ws_vel(links_workspace, wa, k, WS_JOINT_VEL); copy_from_par( coriolis_packed, @@ -273,15 +278,15 @@ pub fn gpu_mb_compute_dynamics_pre( coriolis_packed, coriolis_v_i, 1.0, - ws.shift02, + ws_shift02, parent_coriolis_w, 1.0, lane, t, ); - let dvel = crate::gcross_av(ws.rb_vels.angular, ws.shift02) - + ws.joint_velocity.linear * 2.0; + let ws_rb_ang = ws_vel_ang(links_workspace, wa, k, WS_RB_VELS); + let dvel = crate::gcross_av(ws_rb_ang, ws_shift02) + ws_joint_vel.linear * 2.0; gemm_skew_tr_lhs_cross_buf_par( coriolis_packed, coriolis_v_i, @@ -298,7 +303,7 @@ pub fn gpu_mb_compute_dynamics_pre( coriolis_packed, coriolis_v_i, 1.0, - ws.joint_velocity.linear, + ws_joint_vel.linear, body_jacobians, parent_j_w, 1.0, @@ -311,7 +316,7 @@ pub fn gpu_mb_compute_dynamics_pre( coriolis_v_i, 1.0, parent_w, - ws.shift02, + ws_shift02, body_jacobians, parent_j_w, 1.0, @@ -325,7 +330,7 @@ pub fn gpu_mb_compute_dynamics_pre( coriolis_packed, coriolis_w_i, -1.0, - ws.joint_velocity.angular, + ws_joint_vel.angular, body_jacobians, parent_j_w, 1.0, @@ -342,17 +347,18 @@ pub fn gpu_mb_compute_dynamics_pre( if k != 0 { let stat = stat_slice[k as usize]; let parent_id = stat.parent_link_id; - let parent_link = &ws_slice[parent_id as usize]; if stat.kinematic == 0 { - let transform_rot = - parent_link.local_to_world.rotation * stat.data.local_frame_a.rotation; + let transform_rot = ws_rot(links_workspace, wa, parent_id, WS_LTW) + * stat.data.local_frame_a.rotation; let coriolis_v_part = coriolis_v_i.columns(stat.assembly_id, stat.ndofs); let coriolis_w_part = coriolis_w_i.columns(stat.assembly_id, stat.ndofs); #[cfg(feature = "dim3")] { - let parent_w_skew = crate::utils::linalg::skew(parent_link.rb_vels.angular); + let parent_w_skew = crate::utils::linalg::skew( + ws_vel_ang(links_workspace, wa, parent_id, WS_RB_VELS), + ); let c = lane; if c < stat.ndofs { let (jv, jw) = stat.joint_jacobian_column(transform_rot, c); @@ -374,7 +380,7 @@ pub fn gpu_mb_compute_dynamics_pre( } #[cfg(feature = "dim2")] { - let parent_w = parent_link.rb_vels.angular; + let parent_w = ws_vel_ang(links_workspace, wa, parent_id, WS_RB_VELS); let c = lane; if c < stat.ndofs { let (jv, _) = stat.joint_jacobian_column(transform_rot, c); @@ -397,19 +403,20 @@ pub fn gpu_mb_compute_dynamics_pre( sync_slots(t); if loop_is_active { - let ws = &ws_slice[k as usize]; + let ws_shift23 = ws_vec(links_workspace, wa, k, WS_SHIFT23); + let ws_rb_ang = ws_vel_ang(links_workspace, wa, k, WS_RB_VELS); gemm_skew_tr_lhs_par( coriolis_packed, coriolis_v_i, 1.0, - ws.shift23, + ws_shift23, coriolis_w_i, 1.0, lane, t, ); - let dvel_23 = crate::gcross_av(ws.rb_vels.angular, ws.shift23); + let dvel_23 = crate::gcross_av(ws_rb_ang, ws_shift23); gemm_skew_tr_lhs_cross_buf_par( coriolis_packed, coriolis_v_i, @@ -426,8 +433,8 @@ pub fn gpu_mb_compute_dynamics_pre( coriolis_packed, coriolis_v_i, 1.0, - ws.rb_vels.angular, - ws.shift23, + ws_rb_ang, + ws_shift23, body_jacobians, rb_j_w, 1.0, @@ -505,7 +512,7 @@ pub fn gpu_mb_compute_dynamics_without_coriolis_pre( #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] links_static: &[MultibodyLinkStatic], #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] - links_workspace: &mut [MultibodyLinkWorkspace], + links_workspace: &mut [Vec4], #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] poses: &mut [Pose], #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] body_jacobians: &mut [f32], #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] mass_matrices: &mut [f32], @@ -535,9 +542,7 @@ pub fn gpu_mb_compute_dynamics_without_coriolis_pre( let stat_slice = batch_ids .ib(batch_id, links_static) .offset(mb.first_link as usize); - let mut ws_slice = batch_ids - .ib_mut(batch_id, links_workspace) - .offset(mb.first_link as usize); + let wa = WsAddr::new(mb.first_link as usize, batch_ids.num_batches, batch_id); let mut poses_slice = batch_ids.coll_batch_mut(batch_id, poses); let damping_slice = batch_ids .ib(batch_id, dof_state) @@ -552,7 +557,7 @@ pub fn gpu_mb_compute_dynamics_without_coriolis_pre( // 1) Forward Kinematics (single-threaded) if active_slot && num_links > 0 && lane == 0 { - forward_kinematics(&mb, &stat_slice, &mut poses_slice, &mut ws_slice, num_links); + forward_kinematics(&mb, &stat_slice, &mut poses_slice, links_workspace, wa, num_links); } sync_slots(t); @@ -565,7 +570,8 @@ pub fn gpu_mb_compute_dynamics_without_coriolis_pre( num_links, batch_ids.mb_max_links, &stat_slice, - &ws_slice.as_ref(), + links_workspace, + wa, body_jacobians, batch_ids, batch_id, @@ -573,7 +579,7 @@ pub fn gpu_mb_compute_dynamics_without_coriolis_pre( // 3) Velocities propagation (single-threaded) if active_slot && num_links > 0 && lane == 0 { - propagate_velocities(num_links, &stat_slice, &vel_slice, &mut ws_slice); + propagate_velocities(num_links, &stat_slice, &vel_slice, links_workspace, wa); } sync_slots(t); @@ -593,10 +599,9 @@ pub fn gpu_mb_compute_dynamics_without_coriolis_pre( } if active { - let ws = &ws_slice[k as usize]; let lmp = stat_slice[k as usize].local_mprops; let mass = 1.0 / lmp.inv_mass.x; - let inertia = ws.link_world_inertia(&lmp); + let inertia = ws_world_inertia(links_workspace, wa, k, &lmp); let body_jacobian = batch_ids.imat(batch_id, mb_jac_base + (k as usize) * SPATIAL_DIM * (ndofs as usize), @@ -689,7 +694,8 @@ fn forward_kinematics( mb: &MultibodyInfo, stat_slice: &ISlice, poses_slice: &mut SliceMut, - ws_slice: &mut ISliceMut, + ws: &mut [Vec4], + wa: WsAddr, num_links: u32, ) { // Root pose. @@ -697,26 +703,22 @@ fn forward_kinematics( let root_pose = if mb.root_is_dynamic == 0 { poses_slice[root_config.rb_id as usize] } else { - let ws_ref = &ws_slice[0]; - let pose = root_config.body_to_parent(ws_ref.joint_rot, &ws_ref.coords); + let jr = ws_rot(ws, wa, 0, WS_JOINT_ROT); + let coords = ws_coords(ws, wa, 0); + let pose = root_config.body_to_parent(jr, &coords); poses_slice[root_config.rb_id as usize] = pose; pose }; - let link0 = &mut ws_slice[0]; - link0.local_to_parent = root_pose; - link0.local_to_world = root_pose; + ws_set_pose(ws, wa, 0, WS_LTP, root_pose); + ws_set_pose(ws, wa, 0, WS_LTW, root_pose); for k in 1..num_links { let k_usize = k as usize; let stat = &stat_slice[k_usize]; - let local_to_parent; - let parent_to_world; - { - let ws_ref = &ws_slice[k_usize]; - let parent_ref = &ws_slice[stat.parent_link_id as usize]; - parent_to_world = parent_ref.local_to_world; - local_to_parent = stat.body_to_parent(ws_ref.joint_rot, &ws_ref.coords); - } + let parent_to_world = ws_pose(ws, wa, stat.parent_link_id, WS_LTW); + let jr = ws_rot(ws, wa, k, WS_JOINT_ROT); + let coords = ws_coords(ws, wa, k); + let local_to_parent = stat.body_to_parent(jr, &coords); let local_to_world = parent_to_world * local_to_parent; let parent_lmp = stat_slice[stat.parent_link_id as usize].local_mprops; @@ -727,11 +729,10 @@ fn forward_kinematics( let shift02 = child_anchor_world - parent_com_world; let shift23 = world_com - child_anchor_world; - let link_mut = &mut ws_slice[k_usize]; - link_mut.local_to_parent = local_to_parent; - link_mut.local_to_world = local_to_world; - link_mut.shift02 = shift02; - link_mut.shift23 = shift23; + ws_set_pose(ws, wa, k, WS_LTP, local_to_parent); + ws_set_pose(ws, wa, k, WS_LTW, local_to_world); + ws_set_vec(ws, wa, k, WS_SHIFT02, shift02); + ws_set_vec(ws, wa, k, WS_SHIFT23, shift23); poses_slice[stat.rb_id as usize] = local_to_world; } } @@ -745,7 +746,8 @@ fn update_body_jacobians( num_links: u32, max_links: u32, stat_slice: &ISlice, - ws_slice: &ISlice, + ws: &[Vec4], + wa: WsAddr, body_jacobians: &mut [f32], batch_ids: &BatchIndices, batch_id: u32, @@ -764,7 +766,6 @@ fn update_body_jacobians( if k < num_links { let link_infos = &stat_slice[k as usize]; - let link = &ws_slice[k as usize]; if k != 0 { let parent_j = batch_ids.imat(batch_id, @@ -773,8 +774,7 @@ fn update_body_jacobians( SPATIAL_DIM as u32, ndofs, ); - let parent_link = &ws_slice[link_infos.parent_link_id as usize]; - parent_to_world = parent_link.local_to_world; + parent_to_world = ws_pose(ws, wa, link_infos.parent_link_id, WS_LTW); copy_from_par(body_jacobians, link_j, parent_j, lane, lanes); let link_j_v = link_j.fixed_rows(0, DIM); @@ -783,7 +783,7 @@ fn update_body_jacobians( body_jacobians, link_j_v, 1.0, - link.shift02, + ws_vec(ws, wa, k, WS_SHIFT02), parent_j_w, 1.0, lane, @@ -811,13 +811,12 @@ fn update_body_jacobians( sync_slots(lanes); if k < num_links { - let link = &ws_slice[k as usize]; let (link_j_v, link_j_w) = link_j.rows_range_pair(0, DIM, DIM, ANG_DIM); gemm_skew_tr_lhs_par( body_jacobians, link_j_v, 1.0, - link.shift23, + ws_vec(ws, wa, k, WS_SHIFT23), link_j_w, 1.0, lane, @@ -833,7 +832,8 @@ fn propagate_velocities( num_links: u32, stat_slice: &ISlice, vel_slice: &ISlice, - ws_slice: &mut ISliceMut, + ws: &mut [Vec4], + wa: WsAddr, ) { for k in 0..num_links { let k_usize = k as usize; @@ -846,14 +846,14 @@ fn propagate_velocities( let jv = Velocity::new(jv_local_lin, jv_local_ang); (jv, jv) } else { - let parent_id = stat.parent_link_id as usize; - let parent_ws = &ws_slice[parent_id]; - let parent_to_world_rot = parent_ws.local_to_world.rotation; - let parent_world_com_pose = parent_ws.local_to_world; - let parent_rb_lin = parent_ws.rb_vels.linear; - let parent_rb_ang = parent_ws.rb_vels.angular; - - let parent_lmp = stat_slice[parent_id].local_mprops; + let parent_id = stat.parent_link_id; + let parent_world_com_pose = ws_pose(ws, wa, parent_id, WS_LTW); + let parent_to_world_rot = parent_world_com_pose.rotation; + let parent_rb = ws_vel(ws, wa, parent_id, WS_RB_VELS); + let parent_rb_lin = parent_rb.linear; + let parent_rb_ang = parent_rb.angular; + + let parent_lmp = stat_slice[parent_id as usize].local_mprops; let transform_rot = parent_to_world_rot * stat.data.local_frame_a.rotation; #[cfg(feature = "dim3")] @@ -862,10 +862,8 @@ fn propagate_velocities( #[cfg(feature = "dim2")] let joint_velocity = Velocity::new(transform_rot * jv_local_lin, jv_local_ang); - let (self_local_to_world, self_shift23) = { - let ws_ref = &ws_slice[k_usize]; - (ws_ref.local_to_world, ws_ref.shift23) - }; + let self_local_to_world = ws_pose(ws, wa, k, WS_LTW); + let self_shift23 = ws_vec(ws, wa, k, WS_SHIFT23); let lmp = stat.local_mprops; let world_com = self_local_to_world * lmp.com; @@ -880,8 +878,7 @@ fn propagate_velocities( (joint_velocity, Velocity::new(new_lin, new_ang)) }; - let link_mut = &mut ws_slice[k_usize]; - link_mut.joint_velocity = joint_velocity; - link_mut.rb_vels = rb_vels; + ws_set_vel(ws, wa, k, WS_JOINT_VEL, joint_velocity); + ws_set_vel(ws, wa, k, WS_RB_VELS, rb_vels); } } diff --git a/src_rbd_shaders/dynamics/multibody/gravity_and_lu.rs b/src_rbd_shaders/dynamics/multibody/gravity_and_lu.rs index e2a8cca..cd981e7 100644 --- a/src_rbd_shaders/dynamics/multibody/gravity_and_lu.rs +++ b/src_rbd_shaders/dynamics/multibody/gravity_and_lu.rs @@ -18,9 +18,9 @@ use glamx::Vec4; use crate::dynamics::body::Velocity; use crate::dynamics::joint::SPATIAL_DIM; use crate::utils::linalg::{ - MAX_MB_DOFS, MatSlice, fill_par, gemv_tr_spatial_split_par, lu_decompose, lu_solve_in_place, + MAX_MB_DOFS, fill_par, gemv_tr_spatial_split_par, lu_decompose, lu_solve_in_place, }; -use crate::utils::{BatchIndices, Slice}; +use crate::utils::BatchIndices; use crate::{AngVector, Vector, gcross_av}; use super::lu::{ @@ -28,7 +28,11 @@ use super::lu::{ lu_factor_in_shared_packed, lu_triangular_solve_in_place, lu_triangular_solve_in_place_packed, sm_idx, sm_idx_packed, }; -use super::types::{MultibodyInfo, MultibodyLinkStatic, MultibodyLinkWorkspace}; +use super::types::{MultibodyInfo, MultibodyLinkStatic}; +use super::ws_soa::{ + WS_JOINT_VEL, WS_KIN_ACC, WS_LTW, WS_RB_VELS, WS_SHIFT02, WS_SHIFT23, WsAddr, ws_pose, + ws_set_vel, ws_vec, ws_vel, ws_vel_ang, ws_world_inertia, +}; /// Fused gravity / Coriolis-force assembly + LU factor + LU solve. #[spirv_bindgen] @@ -40,7 +44,7 @@ pub fn gpu_mb_gravity_and_lu( #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] links_static: &[MultibodyLinkStatic], #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] - links_workspace: &mut [MultibodyLinkWorkspace], + links_workspace: &mut [Vec4], #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] body_jacobians: &[f32], #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] gen_forces: &mut [f32], #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] mass_matrices: &mut [f32], @@ -77,9 +81,7 @@ pub fn gpu_mb_gravity_and_lu( let stat_slice = batch_ids .ib(batch_id, links_static) .offset(mb.first_link as usize); - let mut ws_slice = batch_ids - .ib_mut(batch_id, links_workspace) - .offset(mb.first_link as usize); + let wa = WsAddr::new(mb.first_link as usize, batch_ids.num_batches, batch_id); let vel_slice = batch_ids.ib(batch_id, dof_state).offset(gen_base); let damping_slice = batch_ids .ib(batch_id, dof_state) @@ -119,23 +121,24 @@ pub fn gpu_mb_gravity_and_lu( _self_local_to_world, self_rb_ang, ) = { - let ws = &ws_slice[k as usize]; + let jv = ws_vel(links_workspace, wa, k, WS_JOINT_VEL); ( - ws.joint_velocity.linear, - ws.joint_velocity.angular, - ws.shift02, - ws.shift23, - ws.local_to_world, - ws.rb_vels.angular, + jv.linear, + jv.angular, + ws_vec(links_workspace, wa, k, WS_SHIFT02), + ws_vec(links_workspace, wa, k, WS_SHIFT23), + ws_pose(links_workspace, wa, k, WS_LTW), + ws_vel_ang(links_workspace, wa, k, WS_RB_VELS), ) }; if k != 0 { let stat = stat_slice[k as usize]; - let parent_ws = &ws_slice[stat.parent_link_id as usize]; - let parent_acc_lin = parent_ws.kinematic_acc.linear; - let parent_acc_ang = parent_ws.kinematic_acc.angular; - let parent_ang = parent_ws.rb_vels.angular; + let pid = stat.parent_link_id; + let parent_acc = ws_vel(links_workspace, wa, pid, WS_KIN_ACC); + let parent_acc_lin = parent_acc.linear; + let parent_acc_ang = parent_acc.angular; + let parent_ang = ws_vel_ang(links_workspace, wa, pid, WS_RB_VELS); acc_lin = parent_acc_lin; acc_ang = parent_acc_ang; @@ -160,7 +163,7 @@ pub fn gpu_mb_gravity_and_lu( acc_lin += gcross_av(acc_ang, self_shift23); if lane == 0 { - ws_slice[k as usize].kinematic_acc = Velocity::new(acc_lin, acc_ang); + ws_set_vel(links_workspace, wa, k, WS_KIN_ACC, Velocity::new(acc_lin, acc_ang)); } } @@ -170,12 +173,12 @@ pub fn gpu_mb_gravity_and_lu( if active { #[cfg(feature = "dim3")] - let rb_ang = ws_slice[k as usize].rb_vels.angular; + let rb_ang = ws_vel_ang(links_workspace, wa, k, WS_RB_VELS); let lmp = stat_slice[k as usize].local_mprops; let inv_mass_x = lmp.inv_mass.x; if inv_mass_x != 0.0 { let mass = 1.0 / inv_mass_x; - let rb_inertia = ws_slice[k as usize].link_world_inertia(&lmp); + let rb_inertia = ws_world_inertia(links_workspace, wa, k, &lmp); #[cfg(feature = "dim3")] let gyroscopic = { @@ -279,7 +282,7 @@ fn gravity_and_lu_packed_impl> 16; let stat = &stat_slice[link_id as usize]; - let ws = &ws_slice[link_id as usize]; - let curr_pos = ws.coords.read(axis as usize); + let curr_pos = ws_coord(links_workspace, wa, link_id, axis); // Rebuild the per-substep fields with the SAME formulas as the full // emission, then graft the per-step constants (column-derived @@ -457,7 +452,7 @@ pub fn gpu_mb_init_joint_constraints( #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] links_static: &[MultibodyLinkStatic], #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] - links_workspace: &[MultibodyLinkWorkspace], + links_workspace: &[Vec4], #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] mass_matrices: &[f32], #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] lu_pivots: &[u32], #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] diff --git a/src_rbd_shaders/dynamics/multibody/mod.rs b/src_rbd_shaders/dynamics/multibody/mod.rs index 5cde22b..45ed4e2 100644 --- a/src_rbd_shaders/dynamics/multibody/mod.rs +++ b/src_rbd_shaders/dynamics/multibody/mod.rs @@ -25,6 +25,7 @@ mod mass_matrix; mod solve_constraints; mod types; mod utils; +mod ws_soa; pub use compute_dynamics_pre::*; pub use contact_constraints::*; @@ -35,3 +36,4 @@ pub use joint_constraints::*; pub use solve_constraints::*; pub use types::*; pub use utils::*; +pub use ws_soa::*; diff --git a/src_rbd_shaders/dynamics/multibody/ws_soa.rs b/src_rbd_shaders/dynamics/multibody/ws_soa.rs new file mode 100644 index 0000000..1449ed4 --- /dev/null +++ b/src_rbd_shaders/dynamics/multibody/ws_soa.rs @@ -0,0 +1,373 @@ +//! SoA layout for the per-link multibody workspace. + +use glamx::Vec4; +use khal_std::index::MaybeIndexUnchecked; + +#[cfg(feature = "dim2")] +use glamx::Vec2; +#[cfg(feature = "dim3")] +use glamx::Vec3; + +use super::types::MultibodyLinkWorkspace; +use crate::dynamics::body::Velocity; +use crate::{Pose, Rotation, Vector}; + +/* + * Per-link Vec4 offsets of each field. + */ +#[cfg(feature = "dim3")] +mod layout { + /// Joint rotation quat (xyzw). + pub const WS_JOINT_ROT: u32 = 0; + /// Generalized coordinates c0..c3 | c4, c5, pad, pad. + pub const WS_COORDS: u32 = 1; + /// Local-to-parent: rot quat | trans xyz, pad. + pub const WS_LTP: u32 = 3; + /// Local-to-world: rot quat | trans xyz, pad. + pub const WS_LTW: u32 = 5; + /// shift02 xyz, pad. + pub const WS_SHIFT02: u32 = 7; + /// shift23 xyz, pad. + pub const WS_SHIFT23: u32 = 8; + /// Joint velocity: lin xyz, pad | ang xyz, pad. + pub const WS_JOINT_VEL: u32 = 9; + /// Rigid-body velocity: lin | ang. + pub const WS_RB_VELS: u32 = 11; + /// Kinematic acceleration: lin | ang. + pub const WS_KIN_ACC: u32 = 13; + /// Total quads per link (per-link stride, in quad units). + pub const WS_QUADS: u32 = 15; +} + +/* + * Per-link QUAD offsets of each field (dim2): 9 quads / 144 B per link. + */ +#[cfg(feature = "dim2")] +mod layout { + /// Joint rotation (re, im, pad, pad). + pub const WS_JOINT_ROT: u32 = 0; + /// Generalized coordinates c0..c2, pad. + pub const WS_COORDS: u32 = 1; + /// Local-to-parent: (rot.re, rot.im, trans.x, trans.y), one quad. + pub const WS_LTP: u32 = 2; + /// Local-to-world: (rot.re, rot.im, trans.x, trans.y), one quad. + pub const WS_LTW: u32 = 3; + /// shift02 x, y, pad, pad. + pub const WS_SHIFT02: u32 = 4; + /// shift23 x, y, pad, pad. + pub const WS_SHIFT23: u32 = 5; + /// Joint velocity: (lin.x, lin.y, ang, pad). + pub const WS_JOINT_VEL: u32 = 6; + /// Rigid-body velocity: (lin.x, lin.y, ang, pad). + pub const WS_RB_VELS: u32 = 7; + /// Kinematic acceleration: (lin.x, lin.y, ang, pad). + pub const WS_KIN_ACC: u32 = 8; + /// Total quads per link (per-link stride, in quad units). + pub const WS_QUADS: u32 = 9; +} + +pub use layout::*; + +/// Addressing view over the SoA workspace buffer for one multibody of one +/// batch: `base` is the multibody's first link (intra-batch), `stride` / +/// `shift` the batch interleave (`num_batches` / `batch_id`). +#[derive(Copy, Clone)] +pub struct WsAddr { + pub base: usize, + pub stride: u32, + pub shift: u32, +} + +impl WsAddr { + /// View of batch `shift`'s workspace, based at link `base`. + #[inline] + pub fn new(base: usize, stride: u32, shift: u32) -> Self { + Self { + base, + stride, + shift, + } + } + + /// Re-based view (like `Slice::offset`). + #[inline] + pub fn offset(self, links: usize) -> Self { + Self { + base: self.base + links, + ..self + } + } + + /// Flat index of quad `quad` (a `WS_*` field offset + quad index) of + /// link `k` (relative to `base`). + #[inline] + pub fn at(&self, k: u32, quad: u32) -> usize { + ((self.base + k as usize) * WS_QUADS as usize + quad as usize) * self.stride as usize + + self.shift as usize + } +} + +#[cfg(feature = "dim3")] +#[inline] +pub fn ws_rot(buf: &[Vec4], a: WsAddr, k: u32, f: u32) -> Rotation { + let q = buf.read(a.at(k, f)); + Rotation::from_xyzw(q.x, q.y, q.z, q.w) +} + +#[cfg(feature = "dim2")] +#[inline] +pub fn ws_rot(buf: &[Vec4], a: WsAddr, k: u32, f: u32) -> Rotation { + let q = buf.read(a.at(k, f)); + Rotation::from_cos_sin_unchecked(q.x, q.y) +} + +#[cfg(feature = "dim3")] +#[inline] +pub fn ws_set_rot(buf: &mut [Vec4], a: WsAddr, k: u32, f: u32, r: Rotation) { + buf.write(a.at(k, f), Vec4::new(r.x, r.y, r.z, r.w)); +} + +#[cfg(feature = "dim2")] +#[inline] +pub fn ws_set_rot(buf: &mut [Vec4], a: WsAddr, k: u32, f: u32, r: Rotation) { + buf.write(a.at(k, f), Vec4::new(r.re, r.im, 0.0, 0.0)); +} + +#[cfg(feature = "dim3")] +#[inline] +pub fn ws_vec(buf: &[Vec4], a: WsAddr, k: u32, f: u32) -> Vector { + let q = buf.read(a.at(k, f)); + Vec3::new(q.x, q.y, q.z) +} + +#[cfg(feature = "dim2")] +#[inline] +pub fn ws_vec(buf: &[Vec4], a: WsAddr, k: u32, f: u32) -> Vector { + let q = buf.read(a.at(k, f)); + Vec2::new(q.x, q.y) +} + +#[cfg(feature = "dim3")] +#[inline] +pub fn ws_set_vec(buf: &mut [Vec4], a: WsAddr, k: u32, f: u32, v: Vector) { + buf.write(a.at(k, f), Vec4::new(v.x, v.y, v.z, 0.0)); +} + +#[cfg(feature = "dim2")] +#[inline] +pub fn ws_set_vec(buf: &mut [Vec4], a: WsAddr, k: u32, f: u32, v: Vector) { + buf.write(a.at(k, f), Vec4::new(v.x, v.y, 0.0, 0.0)); +} + +/// Pose accessors. 3D: rotation quad + translation quad. 2D: one quad +/// `(rot.re, rot.im, trans.x, trans.y)`. +#[cfg(feature = "dim3")] +#[inline] +pub fn ws_pose(buf: &[Vec4], a: WsAddr, k: u32, f: u32) -> Pose { + let rot = ws_rot(buf, a, k, f); + let tr = ws_vec(buf, a, k, f + 1); + Pose::from_parts(tr, rot) +} + +#[cfg(feature = "dim2")] +#[inline] +pub fn ws_pose(buf: &[Vec4], a: WsAddr, k: u32, f: u32) -> Pose { + let q = buf.read(a.at(k, f)); + Pose::from_parts( + Vec2::new(q.z, q.w), + Rotation::from_cos_sin_unchecked(q.x, q.y), + ) +} + +#[cfg(feature = "dim3")] +#[inline] +pub fn ws_set_pose(buf: &mut [Vec4], a: WsAddr, k: u32, f: u32, p: Pose) { + ws_set_rot(buf, a, k, f, p.rotation); + ws_set_vec(buf, a, k, f + 1, p.translation); +} + +#[cfg(feature = "dim2")] +#[inline] +pub fn ws_set_pose(buf: &mut [Vec4], a: WsAddr, k: u32, f: u32, p: Pose) { + buf.write( + a.at(k, f), + Vec4::new( + p.rotation.re, + p.rotation.im, + p.translation.x, + p.translation.y, + ), + ); +} + +/// Velocity accessors. 3D: linear quad + angular quad. 2D: one quad +/// `(lin.x, lin.y, ang, pad)`. +#[cfg(feature = "dim3")] +#[inline] +pub fn ws_vel(buf: &[Vec4], a: WsAddr, k: u32, f: u32) -> Velocity { + Velocity::new(ws_vec(buf, a, k, f), ws_vec(buf, a, k, f + 1)) +} + +#[cfg(feature = "dim2")] +#[inline] +pub fn ws_vel(buf: &[Vec4], a: WsAddr, k: u32, f: u32) -> Velocity { + let q = buf.read(a.at(k, f)); + Velocity::new(Vec2::new(q.x, q.y), q.z) +} + +/// Angular part of a velocity field (3D loads only the angular quad instead +/// of the whole velocity). +#[cfg(feature = "dim3")] +#[inline] +pub fn ws_vel_ang(buf: &[Vec4], a: WsAddr, k: u32, f: u32) -> crate::AngVector { + ws_vec(buf, a, k, f + 1) +} + +#[cfg(feature = "dim2")] +#[inline] +pub fn ws_vel_ang(buf: &[Vec4], a: WsAddr, k: u32, f: u32) -> crate::AngVector { + buf.read(a.at(k, f)).z +} + +#[cfg(feature = "dim3")] +#[inline] +pub fn ws_set_vel(buf: &mut [Vec4], a: WsAddr, k: u32, f: u32, v: Velocity) { + ws_set_vec(buf, a, k, f, v.linear); + ws_set_vec(buf, a, k, f + 1, v.angular); +} + +#[cfg(feature = "dim2")] +#[inline] +pub fn ws_set_vel(buf: &mut [Vec4], a: WsAddr, k: u32, f: u32, v: Velocity) { + buf.write(a.at(k, f), Vec4::new(v.linear.x, v.linear.y, v.angular, 0.0)); +} + +/// Extract component `i` (0..4) of a `Vec4` by value (no reference indexing, +/// which would create SPIR-V pointer phis). +#[inline] +fn vec4_get(v: Vec4, i: u32) -> f32 { + if i == 0 { + v.x + } else if i == 1 { + v.y + } else if i == 2 { + v.z + } else { + v.w + } +} + +#[inline] +fn vec4_set(v: Vec4, i: u32, val: f32) -> Vec4 { + let mut out = v; + if i == 0 { + out.x = val; + } else if i == 1 { + out.y = val; + } else if i == 2 { + out.z = val; + } else { + out.w = val; + } + out +} + +/// Single generalized-coordinate accessors (`i < MAX_JOINT_DOFS`). +#[inline] +pub fn ws_coord(buf: &[Vec4], a: WsAddr, k: u32, i: u32) -> f32 { + vec4_get(buf.read(a.at(k, WS_COORDS + i / 4)), i % 4) +} + +/// Read-modify-write of one coordinate's quad. +#[inline] +pub fn ws_set_coord(buf: &mut [Vec4], a: WsAddr, k: u32, i: u32, v: f32) { + let idx = a.at(k, WS_COORDS + i / 4); + let q = buf.read(idx); + buf.write(idx, vec4_set(q, i % 4, v)); +} + +/// Load the whole coords array (for `body_to_parent`). +#[cfg(feature = "dim3")] +#[inline] +pub fn ws_coords(buf: &[Vec4], a: WsAddr, k: u32) -> [f32; 6] { + let q0 = buf.read(a.at(k, WS_COORDS)); + let q1 = buf.read(a.at(k, WS_COORDS + 1)); + [q0.x, q0.y, q0.z, q0.w, q1.x, q1.y] +} + +#[cfg(feature = "dim2")] +#[inline] +pub fn ws_coords(buf: &[Vec4], a: WsAddr, k: u32) -> [f32; 3] { + let q0 = buf.read(a.at(k, WS_COORDS)); + [q0.x, q0.y, q0.z] +} + +/// World-space inertia of link `k`, the SoA counterpart of the former +/// `MultibodyLinkWorkspace::link_world_inertia` (reads only the +/// local-to-world rotation). +#[cfg(feature = "dim3")] +#[inline] +pub fn ws_world_inertia( + buf: &[Vec4], + a: WsAddr, + k: u32, + lmp: &crate::dynamics::body::LocalMassProperties, +) -> glamx::Mat3 { + use crate::rotation_to_matrix; + let ipi = lmp.inv_principal_inertia; + let px = if ipi.x != 0.0 { 1.0 / ipi.x } else { 0.0 }; + let py = if ipi.y != 0.0 { 1.0 / ipi.y } else { 0.0 }; + let pz = if ipi.z != 0.0 { 1.0 / ipi.z } else { 0.0 }; + let rot = ws_rot(buf, a, k, WS_LTW); + let r = rotation_to_matrix(rot * lmp.inertia_ref_frame); + // M = r · diag(px, py, pz) (column-scale); I = M · rᵀ. + let m = glamx::Mat3::from_cols(r.x_axis * px, r.y_axis * py, r.z_axis * pz); + m * r.transpose() +} + +#[cfg(feature = "dim2")] +#[inline] +pub fn ws_world_inertia( + _buf: &[Vec4], + _a: WsAddr, + _k: u32, + lmp: &crate::dynamics::body::LocalMassProperties, +) -> f32 { + if lmp.inv_inertia != 0.0 { + 1.0 / lmp.inv_inertia + } else { + 0.0 + } +} + +/* + * Host-side conversion of the AoS structs into the SoA buffer. `data` is + * batch-major (`batch · links_cap + link`). + */ +#[cfg(not(target_arch_is_gpu))] +pub fn ws_soa_from_structs( + data: &[MultibodyLinkWorkspace], + links_cap: u32, + num_batches: u32, +) -> std::vec::Vec { + let mut out = + std::vec![Vec4::ZERO; links_cap as usize * WS_QUADS as usize * num_batches as usize]; + for b in 0..num_batches { + let a = WsAddr::new(0, num_batches, b); + for k in 0..links_cap { + let ws = &data[(b * links_cap + k) as usize]; + ws_set_rot(&mut out, a, k, WS_JOINT_ROT, ws.joint_rot); + for (i, &c) in ws.coords.iter().enumerate() { + ws_set_coord(&mut out, a, k, i as u32, c); + } + ws_set_pose(&mut out, a, k, WS_LTP, ws.local_to_parent); + ws_set_pose(&mut out, a, k, WS_LTW, ws.local_to_world); + ws_set_vec(&mut out, a, k, WS_SHIFT02, ws.shift02); + ws_set_vec(&mut out, a, k, WS_SHIFT23, ws.shift23); + ws_set_vel(&mut out, a, k, WS_JOINT_VEL, ws.joint_velocity); + ws_set_vel(&mut out, a, k, WS_RB_VELS, ws.rb_vels); + ws_set_vel(&mut out, a, k, WS_KIN_ACC, ws.kinematic_acc); + } + } + out +} From 2d6a57f76c5b147884d7dfc0ab329306a92970f5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Sun, 2 Aug 2026 09:30:35 +0200 Subject: [PATCH 38/39] feat: improve multibody/mjcf stability --- .run/all_examples3.run.xml | 2 +- .../multibody/multibody_from_rapier.rs | 117 ++++++- src_rbd/dynamics/multibody/multibody_set.rs | 34 +- .../dynamics/multibody/multibody_solver.rs | 105 ++---- src_rbd/dynamics/solver.rs | 48 +-- src_rbd_shaders/broad_phase/narrow_phase.rs | 2 +- src_rbd_shaders/dynamics/constraint.rs | 8 +- .../multibody/compute_dynamics_pre.rs | 216 ++++--------- .../dynamics/multibody/contact_constraints.rs | 21 +- .../dynamics/multibody/gravity_and_lu.rs | 282 +++++++++++++++-- .../dynamics/multibody/joint_constraints.rs | 205 ++++++------ .../dynamics/multibody/solve_constraints.rs | 48 ++- src_rbd_shaders/dynamics/multibody/types.rs | 62 +++- src_rbd_shaders/dynamics/sim_params.rs | 156 ++++++--- src_rbd_shaders/dynamics/solver.rs | 61 +++- src_rbd_shaders/dynamics/solver_utils.rs | 298 +++++++++--------- src_rbd_shaders/dynamics/warmstart.rs | 47 ++- src_rbd_shaders/utils/indices.rs | 13 + src_rbd_shaders/utils/linalg.rs | 25 +- 19 files changed, 1085 insertions(+), 665 deletions(-) diff --git a/.run/all_examples3.run.xml b/.run/all_examples3.run.xml index c53f5f1..a1d673c 100644 --- a/.run/all_examples3.run.xml +++ b/.run/all_examples3.run.xml @@ -1,7 +1,7 @@