From 9207f3679e3b7ef2c7c8a61185bf998441027162 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Fri, 7 Aug 2026 09:48:05 +0200 Subject: [PATCH 1/6] chore: build against rapier 0.35 --- Cargo.toml | 19 +++++++++++++------ 1 file changed, 13 insertions(+), 6 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index b997730..c204371 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -38,12 +38,12 @@ bvh = "0.7" # Physics / geometry. default-features off for the no_std shader crates; host # crates opt back in with `features = ["default"]`. -rapier2d = { version = "0.34", default-features = false } -rapier3d = { version = "0.34", default-features = false } -rapier3d-urdf = "0.34" -rapier3d-mjcf = { version = "0.34", features = ["stl", "wavefront", "msh"] } -parry2d = { version = "0.29", default-features = false } -parry3d = { version = "0.29", default-features = false } +rapier2d = { version = "0.35", default-features = false } +rapier3d = { version = "0.35", default-features = false } +rapier3d-urdf = "0.35" +rapier3d-mjcf = { version = "0.35", features = ["stl", "wavefront", "msh"] } +parry2d = { version = "0.30", default-features = false } +parry3d = { version = "0.30", default-features = false } # Viewer / examples deps kiss3d = "0.45.1" @@ -85,6 +85,13 @@ rust.unexpected_cfgs = { level = "warn", check-cfg = [ ] } [patch.crates-io] +# Compare against the rapier checkout the reference example runs, not the +# crates.io release (their solver defaults differ). +#rapier2d = { path = "../rapier/crates/rapier2d" } +#rapier3d = { path = "../rapier/crates/rapier3d" } +#rapier3d-mjcf = { path = "../rapier/crates/rapier3d-mjcf" } +#rapier3d-urdf = { path = "../rapier/crates/rapier3d-urdf" } +#rapier3d-meshloader = { path = "../rapier/crates/rapier3d-meshloader" } ## Local glam clone with SPIR-V vector-arithmetic intrinsics (Vec3 add/sub/mul/scale). #glam = { path = "../glam-rs" } # 30% faster for loop in P2G From 19411dded7de0926e09b1a091ed999a7f9852a7b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Fri, 7 Aug 2026 10:02:32 +0200 Subject: [PATCH 2/6] feat: match rapier's multibody dynamics, joints and contact solving --- src/state.rs | 11 +- .../dynamics/multibody/loop_closing_joints.rs | 46 +- .../multibody/multibody_from_rapier.rs | 39 +- src_rbd/dynamics/multibody/multibody_set.rs | 100 ++- .../dynamics/multibody/multibody_solver.rs | 124 +++- src_rbd/dynamics/solver.rs | 8 + src_rbd/pipeline/insertion_removal.rs | 8 +- src_rbd/pipeline/rbd_state.rs | 37 +- src_rbd/pipeline/rbd_state_from_rapier.rs | 11 +- src_rbd/pipeline/rbd_step.rs | 3 + .../dynamics/joint_constraint_builder.rs | 14 + src_rbd_shaders/dynamics/mod.rs | 1 + .../multibody/compute_dynamics_pre.rs | 32 +- .../dynamics/multibody/contact_constraints.rs | 667 +++++++++++++----- .../dynamics/multibody/gravity_and_lu.rs | 26 +- .../impulse_joint_constraints/helper.rs | 118 +++- .../impulse_joint_constraints/kernels.rs | 10 +- .../impulse_joint_constraints/mod.rs | 2 +- .../impulse_joint_constraints/update.rs | 66 +- .../dynamics/multibody/solve_constraints.rs | 232 ++++-- src_rbd_shaders/dynamics/multibody/types.rs | 98 ++- src_rbd_shaders/dynamics/multibody/utils.rs | 21 + src_rbd_shaders/dynamics/multibody/ws_soa.rs | 85 ++- src_rbd_shaders/dynamics/sim_params.rs | 7 +- src_rbd_shaders/dynamics/solver.rs | 11 +- src_rbd_shaders/utils/indices.rs | 6 +- 26 files changed, 1370 insertions(+), 413 deletions(-) diff --git a/src/state.rs b/src/state.rs index e4dd821..20b788b 100644 --- a/src/state.rs +++ b/src/state.rs @@ -141,11 +141,12 @@ impl NexusState { // TODO: resize the GPU buffers too. } - /// Sets the rigid-body multibody gravity vector, e.g. `[0.0, 0.0, -9.81]` - /// for a Z-up scene. No-op until the rigid-body state is built, so call it - /// after [`Self::finalize`]. Free (non-multibody) bodies keep the solver's - /// fixed gravity. - #[cfg(all(feature = "dim3", feature = "rbd"))] + /// Sets the rigid-body gravity vector, e.g. `[0.0, 0.0, -9.81]` for a Z-up + /// scene. Every solver path reads the same uniform, so this applies to free + /// rigid-bodies and multibody links alike (in 2D the third component is + /// ignored). No-op until the rigid-body state is built, so call it after + /// [`Self::finalize`]. + #[cfg(feature = "rbd")] pub fn set_rbd_gravity(&mut self, backend: &GpuBackend, gravity: [f32; 3]) { if let Some(rbd) = self.rbd.as_mut() { rbd.set_gravity(backend, gravity); diff --git a/src_rbd/dynamics/multibody/loop_closing_joints.rs b/src_rbd/dynamics/multibody/loop_closing_joints.rs index 92d012a..45d06d0 100644 --- a/src_rbd/dynamics/multibody/loop_closing_joints.rs +++ b/src_rbd/dynamics/multibody/loop_closing_joints.rs @@ -113,31 +113,31 @@ impl GpuMultibodySet { continue; // Both sides static — no constraint to solve. } - // Mirror rapier `GenericJoint::transform_to_solver_body_space`: - // shift the anchor frame's translation into COM space — but ONLY - // for FREE-BODY sides, whose solver pose IS the center of mass. - // A multibody-link side is positioned by its `local_to_world`, - // which is the link ORIGIN frame (not the COM), so the shift must - // NOT be applied there — the anchor stays origin-relative and the - // lever arm is taken against the COM separately (see - // `world_com` in `update_one_joint`). Applying the shift to MB - // links offsets the anchor by `local_com` (≈0.25 m for Cassie's - // rods), producing a huge spurious loop-closure violation. This - // matches rapier's `generic_joint_constraint_builder` (the shift - // is applied to `LinkOrBody::Body` sides only). Fixed-side fold - // is still a TODO mirroring rapier's `is_fixed` branch. + // Move each anchor frame into the space the solver resolves it + // against. A free-body side is positioned by its COM-centered + // solver pose, so its anchor shifts by `-local_com`. A fixed + // side has no solver pose at all (the kernel resolves it against + // the identity), so its world transform is folded in here. A + // multibody-link side is positioned by its `local_to_world`, an + // origin frame, so the anchor stays as-is — shifting it there + // would offset it by `local_com` and fabricate a large loop + // violation. let mut joint_data = convert_generic_joint(joint.data); - if side_a_kind == SIDE_KIND_BODY - && let Some(rb) = rb1 - { - let com = rb.mass_properties().local_mprops.local_com; - joint_data.local_frame_a.translation -= com; + if let Some(rb) = rb1 { + if side_a_kind == SIDE_KIND_FIXED { + joint_data.local_frame_a = *rb.position() * joint_data.local_frame_a; + } else if side_a_kind == SIDE_KIND_BODY { + joint_data.local_frame_a.translation -= + rb.mass_properties().local_mprops.local_com; + } } - if side_b_kind == SIDE_KIND_BODY - && let Some(rb) = rb2 - { - let com = rb.mass_properties().local_mprops.local_com; - joint_data.local_frame_b.translation -= com; + if let Some(rb) = rb2 { + if side_b_kind == SIDE_KIND_FIXED { + joint_data.local_frame_b = *rb.position() * joint_data.local_frame_b; + } else if side_b_kind == SIDE_KIND_BODY { + joint_data.local_frame_b.translation -= + rb.mass_properties().local_mprops.local_com; + } } // Per-axis stride = 2 * (ndofs_a + ndofs_b); reserve diff --git a/src_rbd/dynamics/multibody/multibody_from_rapier.rs b/src_rbd/dynamics/multibody/multibody_from_rapier.rs index 75698c0..cd6bf74 100644 --- a/src_rbd/dynamics/multibody/multibody_from_rapier.rs +++ b/src_rbd/dynamics/multibody/multibody_from_rapier.rs @@ -8,7 +8,6 @@ use crate::shaders::dynamics::{ RbdSimParams, }; use crate::shaders::utils::linalg::MAX_MB_DOFS; -use glamx::Vec4; use khal::BufferUsages; use khal::backend::GpuBackend; use vortx::tensor::Tensor; @@ -31,7 +30,6 @@ impl GpuMultibodySet { &HashMap, &RigidBodySet, )], - gravity: [f32; 3], colliders_per_batch: u32, ) -> Self { let num_batches = environments.len() as u32; @@ -245,6 +243,11 @@ impl GpuMultibodySet { let mut ws = make_workspace_init(); ws.coords = link.joint.coords(); ws.joint_rot = link.joint.joint_rot(); + if let Some(rb) = bodies.get(link.rigid_body_handle()) { + ws.gravity_scale = rb.gravity_scale(); + ws.external_force = rb.user_force(); + ws.external_torque = rb.user_torque(); + } // For free joints at the root, copy the rigid-body pose directly. if link.joint.data.locked_axes.is_empty() @@ -467,6 +470,7 @@ impl GpuMultibodySet { out } let nb = num_batches as usize; + let info_mirror = all_infos.clone(); let all_infos = interleave(&all_infos, mb_cap, nb); let all_statics = interleave(&all_statics, links_cap, nb); let all_dof_vals = interleave(&all_dof_vals, dofs_cap, nb); @@ -487,18 +491,19 @@ impl GpuMultibodySet { mass_matrix_entries_per_batch: mm_cap, coriolis_entries_per_batch: cor_cap, i_coriolis_dt_entries_per_batch: icdt_cap, - // Rapier's scheme (the default): the coriolis-melded matrix only - // drives the free-acceleration solve, constraints see the plain - // matrix, and everything is built once per step. The implicit - // legacy mode (single melded matrix, per-substep rebuilds) stays - // available via `set_implicit_coriolis(true)`. - implicit_coriolis: false, + // Default: implicit coriolis. The acceleration solve uses a mass + // matrix augmented with the coriolis/gyroscopic derivatives, while + // constraints keep the plain one. `set_implicit_coriolis(false)` + // falls back to a single plain matrix with explicit-only + // coriolis/gyroscopic forces (cheaper, but less stable). + implicit_coriolis: true, has_joint_constraints: all_infos.iter().any(|info| info.max_constraints > 0), multibody_info: Tensor::vector(backend, &all_infos, storage).unwrap(), links_static: Tensor::vector(backend, &all_statics, storage | BufferUsages::COPY_DST) .unwrap(), links_static_mirror: all_statics.clone(), + info_mirror, links_workspace: Tensor::vector( backend, &crate::shaders::dynamics::ws_soa_from_structs(&all_ws, links_cap, num_batches), @@ -581,6 +586,15 @@ impl GpuMultibodySet { storage, ) .unwrap(), + old_contact_constraints: Tensor::vector( + backend, + vec![ + MultibodyContactConstraint::default(); + (contact_cons_cap * num_batches) as usize + ], + storage, + ) + .unwrap(), contact_constraint_jacs: Tensor::vector( backend, vec![0.0f32; (contact_cons_col_cap * num_batches) as usize], @@ -667,14 +681,8 @@ impl GpuMultibodySet { contact_constraint_columns_per_batch: contact_cons_col_cap, num_solver_iterations: 4, + num_internal_pgs_iterations: 1, - // FIXME: should be read from the simulation settings. - gravity: Tensor::scalar( - backend, - Vec4::new(gravity[0], gravity[1], gravity[2], 0.0), - BufferUsages::STORAGE | BufferUsages::UNIFORM | BufferUsages::COPY_DST, - ) - .unwrap(), dt: Tensor::scalar( backend, 1.0f32 / 60.0, @@ -689,6 +697,7 @@ impl GpuMultibodySet { BufferUsages::STORAGE | BufferUsages::UNIFORM | BufferUsages::COPY_DST, ) .unwrap(), + warmstart_coefficient: RbdSimParams::default().warmstart_coefficient, } } } diff --git a/src_rbd/dynamics/multibody/multibody_set.rs b/src_rbd/dynamics/multibody/multibody_set.rs index 9693f5f..ab2989b 100644 --- a/src_rbd/dynamics/multibody/multibody_set.rs +++ b/src_rbd/dynamics/multibody/multibody_set.rs @@ -1,5 +1,5 @@ //! The [`GpuMultibodySet`] buffers: struct definition, accessors and -//! runtime-mutation entry points (motors, gravity, dt, softness). +//! runtime-mutation entry points (motors, dt, softness). use crate::math::Pose; use crate::shaders::dynamics::{ @@ -8,7 +8,6 @@ use crate::shaders::dynamics::{ MultibodyJointConstraint, MultibodyLinkStatic, MultibodyLinkWorkspace, RbdSimParams, }; use crate::shaders::utils::BatchIndices; -use glamx::Vec4; use khal::BufferUsages; use khal::backend::{Backend, GpuBackend, GpuBackendError}; use rapier3d::prelude::JointAxis; @@ -23,7 +22,7 @@ pub(super) const MB_LU_LANES: u32 = 64; /// 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; +pub(super) const MAX_DELASSUS_MULTIBODIES: u32 = 0; // 128; use crate::shaders::dynamics::{GenericJoint, JointLimits, JointMotor}; @@ -57,6 +56,9 @@ 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, + /// Host copy of the per-multibody descriptors, batch-major (before the + /// batch interleave), indexed `batch * multibodies_per_batch + mb_idx`. + pub(super) info_mirror: Vec, /// Per-batch per-step link workspace, SoA quad layout. pub(super) links_workspace: Tensor, /// Generalized coordinates (flat). @@ -95,6 +97,9 @@ pub struct GpuMultibodySet { /// Per-multibody bank of contact constraints (1 normal + 2 friction per /// touched contact point). pub(super) contact_constraints: Tensor, + /// Snapshot of `contact_constraints` taken at the start of the step; the + /// warmstart transfer matches this frame's slots against it. + pub(super) old_contact_constraints: Tensor, /// Per-constraint `Jᵀ` row (length `ndofs`) — the multibody side's /// contribution to the constraint Jacobian. pub(super) contact_constraint_jacs: Tensor, @@ -145,14 +150,20 @@ pub struct GpuMultibodySet { /// Number of solver iterations to run on `joint_constraints` per `step()`. pub(super) num_solver_iterations: u32, + /// PGS iterations over the joint + contact constraints per substep, in the + /// biased pass. One is enough for simple articulations; servo-driven robots + /// resting on contacts need several to stop the motor and contact rows + /// fighting each other. + pub(super) num_internal_pgs_iterations: u32, - /// Gravity vector (only the first 3 components are read by the shaders). - pub(super) gravity: Tensor, /// Current integration timestep. pub(super) dt: Tensor, /// Precomputed soft-constraint coefficients (contact + joint, rapier /// TGS-soft). pub(super) constraint_softness: Tensor, + /// CPU mirror of `ConstraintSoftness::warmstart_coefficient`, so the solver + /// can skip the warmstart passes entirely when it is zero. + pub(super) warmstart_coefficient: f32, } impl GpuMultibodySet { @@ -223,7 +234,7 @@ impl GpuMultibodySet { &self.gen_forces } - /// Indicates if the implicit treatment of coriolis forces is enabled. + /// Enables or disables the implicit treatment of coriolis forces. pub fn set_implicit_coriolis(&mut self, enabled: bool) { self.implicit_coriolis = enabled; } @@ -247,6 +258,17 @@ impl GpuMultibodySet { self.num_solver_iterations = n; } + /// Sets how many PGS iterations the biased pass runs per substep (default 1). + pub fn set_num_internal_pgs_iterations(&mut self, n: u32) { + self.num_internal_pgs_iterations = n.max(1); + } + + /// PGS iterations per substep in the biased pass. + pub fn num_internal_pgs_iterations(&self) -> u32 { + self.num_internal_pgs_iterations + } + + /// Upload the visible-frame `dt`. Internally divides by `num_solver_iterations` /// and stores the *substep* dt (which is what the GPU kernels read). pub fn set_visible_dt(&mut self, backend: &GpuBackend, visible_dt: f32) { @@ -263,6 +285,7 @@ impl GpuMultibodySet { /// (substep) sim params. Must be called whenever the contact softness / /// timestep changes. pub fn set_constraint_softness(&mut self, backend: &GpuBackend, params: &RbdSimParams) { + self.warmstart_coefficient = params.warmstart_coefficient; self.constraint_softness = Tensor::scalar( backend, ConstraintSoftness::from_params(params), @@ -305,16 +328,6 @@ impl GpuMultibodySet { ) } - /// Upload a new gravity vector. - pub fn set_gravity(&mut self, backend: &GpuBackend, g: [f32; 3]) { - self.gravity = Tensor::scalar( - backend, - Vec4::new(g[0], g[1], g[2], 0.0), - BufferUsages::STORAGE | BufferUsages::UNIFORM | BufferUsages::COPY_DST, - ) - .unwrap(); - } - /// Number of multibody-touching impulse joints in any batch. pub fn mb_impulse_joints_per_batch(&self) -> u32 { self.mb_imp_joints_per_batch @@ -348,14 +361,64 @@ impl GpuMultibodySet { 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; + // Implicit coriolis needs two matrices: the coriolis-augmented one (acc + // section) for the acceleration solve, the plain one for constraints. + // With the flag off, a single plain matrix serves both. dst.mass_matrix_acc_section_offset = if self.implicit_coriolis { - 0 - } else { self.mass_matrix_entries_per_batch + } else { + 0 }; dst.mb_dof_couplings_batch_capacity = self.couplings_per_batch; } + /// Sets the world-space force and torque applied to `link_id` of multibody + /// `mb_idx` in `batch_id`, plus that link's multiplier on the global + /// gravity. They stay applied until overwritten. + pub fn set_link_external_wrench( + &mut self, + backend: &GpuBackend, + batch_id: u32, + mb_idx: u32, + link_id: u32, + force: crate::math::Vector, + torque: crate::math::AngVector, + gravity_scale: f32, + ) -> Result<(), khal::backend::GpuBackendError> { + use crate::shaders::dynamics::{WS_EXT_FORCE, WS_EXT_TORQUE, WsAddr}; + + let info = self.info_mirror[(batch_id * self.multibodies_per_batch + mb_idx) as usize]; + let k = info.first_link + link_id; + let a = WsAddr::new(0, self.num_batches, batch_id); + + #[cfg(feature = "dim3")] + { + let f = glamx::Vec4::new(force.x, force.y, force.z, gravity_scale); + let t = glamx::Vec4::new(torque.x, torque.y, torque.z, 0.0); + backend.write_buffer( + self.links_workspace.buffer_mut(), + a.at(k, WS_EXT_FORCE) as u64, + &[f], + )?; + backend.write_buffer( + self.links_workspace.buffer_mut(), + a.at(k, WS_EXT_TORQUE) as u64, + &[t], + )?; + } + #[cfg(feature = "dim2")] + { + let _ = WS_EXT_TORQUE; + let f = glamx::Vec4::new(force.x, force.y, torque, gravity_scale); + backend.write_buffer( + self.links_workspace.buffer_mut(), + a.at(k, WS_EXT_FORCE) as u64, + &[f], + )?; + } + Ok(()) + } + /// Upload a new integration timestep. pub fn set_dt(&mut self, backend: &GpuBackend, dt: f32) { self.dt = Tensor::scalar( @@ -412,6 +475,7 @@ pub(super) fn convert_generic_joint(j: crate::rapier::dynamics::GenericJoint) -> pub(super) fn make_workspace_init() -> MultibodyLinkWorkspace { let mut w: MultibodyLinkWorkspace = bytemuck::Zeroable::zeroed(); w.joint_rot = glamx::Quat::IDENTITY; + w.gravity_scale = 1.0; w.local_to_parent = Pose::default(); w.local_to_world = Pose::default(); w diff --git a/src_rbd/dynamics/multibody/multibody_solver.rs b/src_rbd/dynamics/multibody/multibody_solver.rs index 0c9a408..f7126c5 100644 --- a/src_rbd/dynamics/multibody/multibody_solver.rs +++ b/src_rbd/dynamics/multibody/multibody_solver.rs @@ -9,7 +9,8 @@ use crate::shaders::dynamics::{ GpuMbGravityAndLuT16, GpuMbGravityAndLuT32, GpuMbInitContactConstraints, GpuMbInitJointConstraints, GpuMbIntegrate, GpuMbIntegrateVelocities, GpuMbRemoveImpulseJointConstraintBias, - GpuMbResetContactWarmstart, GpuMbStashContactsLen, GpuMbWarmstartContactConstraints, + GpuMbApplyContactRestitution, GpuMbSeedContactRestitution, GpuMbSnapshotContactWarmstart, + GpuMbStashContactsLen, GpuMbTransferContactWarmstart, GpuMbWarmstartContactConstraints, GpuMbSolveConstraints, GpuMbSolveContactsDelassus, GpuMbSolveImpulseJointConstraints, GpuMbSolveJoints, GpuMbFinalizeImpulseJointConstraints, @@ -35,7 +36,7 @@ pub struct GpuMultibodySolver { init_joint_with_bias: GpuMbInitJointConstraints, init_contact_constraints: GpuMbInitContactConstraints, finalize_contact_constraints: GpuMbFinalizeContactConstraints, - /// Fused joint+contact PGS sweep (one workgroup per multibody, shared- + /// Fused joint+contact PGS iteration (one workgroup per multibody, shared- /// memory dof velocities). solve_constraints: GpuMbSolveConstraints, /// Joint-only half of the sweep, used with the Delassus contact path @@ -49,14 +50,20 @@ pub struct GpuMultibodySolver { /// 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, + /// Snapshot the contact impulses once per frame, for the cross-frame match. + snapshot_contact_warmstart: GpuMbSnapshotContactWarmstart, + /// Carry the snapshotted impulses over to this frame's matching contacts. + transfer_contact_warmstart: GpuMbTransferContactWarmstart, /// 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, + /// Capture each bouncy contact's approach velocity at the start of the step. + seed_contact_restitution: GpuMbSeedContactRestitution, + /// Restore that approach velocity once every substep is done. + apply_contact_restitution: GpuMbApplyContactRestitution, 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. @@ -88,6 +95,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, + /// The one gravity uniform every rigid-body and multibody kernel reads. + /// Owned by `RbdState`. + pub gravity: &'a Tensor, /// Per-color-index uniform tensors (`color_uniforms[c]` holds `c`), /// shared with the contact/joint solvers. pub color_uniforms: &'a [Tensor], @@ -128,18 +138,19 @@ impl GpuMultibodySolver { 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). - // Flat (slot, multibody, batch) grid (impulse-field-only stores). + // Snapshot the contact slab this frame's build will overwrite, so the + // warmstart transfer can still match against it. + // Flat (slot, multibody, batch) grid. { - let mut pass = encoder.begin_pass("[RBD] mbi/reset", timestamps.as_deref_mut()); + let mut pass = encoder.begin_pass("[RBD] mbi/snapshot", 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( + self.snapshot_contact_warmstart.call( &mut pass, [total_slots, 1, 1], - &mut mb.contact_constraints, + &mb.contact_constraints, + &mut mb.old_contact_constraints, args.batch_indices, )?; } @@ -213,13 +224,50 @@ impl GpuMultibodySolver { } // Full rebuild of the joint + contact constraints every substep. - self.build_contact_constraints(encoder, timestamps.as_deref_mut(), mb, args)?; + self.build_contact_constraints( + encoder, + timestamps.as_deref_mut(), + mb, + args, + first_substep, + )?; + + // Carry the previous frame's impulses over before anything reads them. + if first_substep && mb.warmstart_coefficient != 0.0 { + let mut pass = + encoder.begin_pass("[RBD] mbb/transfer-warmstart", timestamps.as_deref_mut()); + self.transfer_contact_warmstart.call( + &mut pass, + args.mb_sweep_indirect, + &mb.multibody_info, + &mut mb.contact_constraints, + &mb.old_contact_constraints, + args.batch_indices, + &mb.constraint_softness, + )?; + } + + // Restitution is measured once, from the velocities the step starts with. + if first_substep { + let mut pass = + encoder.begin_pass("[RBD] mbb/seed-restitution", timestamps.as_deref_mut()); + self.seed_contact_restitution.call( + &mut pass, + args.mb_sweep_indirect, + &mb.multibody_info, + &mut mb.contact_constraints, + &mb.contact_constraint_jacs, + &mb.dof_state, + args.solver_vels, + 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. + // substep, including the first, which carries the previous frame's. // One 64-lane workgroup per multibody (one DOF per lane). - if !first_substep { + if mb.warmstart_coefficient != 0.0 { let mut pass = encoder.begin_pass("[RBD] mbb/warmstart-contact", timestamps.as_deref_mut()); // Contact-only work: indirect grid collapses to zero workgroups @@ -247,6 +295,7 @@ 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() { @@ -284,15 +333,16 @@ impl GpuMultibodySolver { &mut pass, init_contact_dispatch, &mut mb.multibody_info, - &mb.body_jacobians, + &mb.links_workspace, &mb.body_to_link, &mut mb.contact_constraints, - &mut mb.contact_constraint_jacs, &mb.constraint_softness, args.batch_indices, + &args.color_uniforms[first_substep as usize], args.mprops, args.collider_world_poses, args.contacts, + args.poses, )?; } @@ -307,8 +357,10 @@ impl GpuMultibodySolver { &mb.mass_matrices, &mb.lu_pivots, &mut mb.contact_constraints, - &mb.contact_constraint_jacs, + &mut mb.contact_constraint_jacs, &mut mb.contact_constraint_columns, + &mb.links_static, + &mb.body_jacobians, args.batch_indices, )?; } @@ -333,7 +385,7 @@ impl GpuMultibodySolver { Ok(()) } - /// One joint+contact PGS sweep: the dof-space fused kernel, or (when the + /// One joint+contact PGS iteration: 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). @@ -411,7 +463,7 @@ impl GpuMultibodySolver { Ok(()) } - /// P3: one PGS sweep with bias over the joint, contact, and multibody- + /// P3: one PGS iteration with bias over the joint, contact, and multibody- /// touching impulse-joint constraints. pub fn substep_solve_with_bias( &self, @@ -428,7 +480,9 @@ impl GpuMultibodySolver { // 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.dispatch_solve(pass, mb, args, solve_dispatch, 1)?; + for _ in 0..mb.num_internal_pgs_iterations() { + self.dispatch_solve(pass, mb, args, solve_dispatch, 1)?; + } // Multibody-touching impulse joints — generic (rb-mb / mb-mb) // constraints. @@ -461,8 +515,9 @@ impl GpuMultibodySolver { &mb.multibody_info, &mb.mass_matrices, &mb.lu_pivots, + &mb.links_static, )?; - // Colored PGS sweep WITH bias: one dispatch per color, each + // Colored PGS iteration WITH bias: one dispatch per color, each // color's joints solved race-free in parallel (graph coloring // done at init in `set_impulse_joints`). for c in 0..mb.mb_imp_joint_num_colors as usize { @@ -527,7 +582,7 @@ impl GpuMultibodySolver { Ok(()) } - /// P5: stabilization — fused remove-bias + final PGS sweep WITHOUT bias for + /// P5: stabilization — fused remove-bias + final PGS iteration WITHOUT bias for /// joint limits/motors, contacts, and multibody-touching impulse joints. /// Settles velocity along constrained DOFs to zero (no rebound from the /// positional bias). @@ -582,6 +637,29 @@ impl GpuMultibodySolver { Ok(()) } + /// End-of-step restitution pass, run once after the last substep. + pub fn apply_restitution( + &self, + pass: &mut GpuPass, + mb: &mut GpuMultibodySet, + args: &mut MultibodySolverArgs<'_>, + ) -> Result<(), GpuBackendError> { + if mb.is_empty() { + return Ok(()); + } + self.apply_contact_restitution.call( + pass, + args.mb_sweep_indirect, + &mb.multibody_info, + &mut mb.contact_constraints, + &mb.contact_constraint_jacs, + &mb.contact_constraint_columns, + args.batch_indices, + &mut mb.dof_state, + args.solver_vels, + ) + } + /// Recompute the dynamics (mass matrix, LU factors, generalized /// acceleration). After this call, `gen_forces` holds the generalized /// acceleration `a` for the *next* substep's velocity update. @@ -625,7 +703,7 @@ impl GpuMultibodySolver { &mut mb.mass_matrices, &mut mb.lu_pivots, &mb.dof_state, - &mb.gravity, + args.gravity, args.batch_indices, &mb.dt, )? @@ -650,7 +728,7 @@ impl GpuMultibodySolver { &mut mb.mass_matrices, &mut mb.lu_pivots, &mb.dof_state, - &mb.gravity, + args.gravity, args.batch_indices, &mb.dt, )?; diff --git a/src_rbd/dynamics/solver.rs b/src_rbd/dynamics/solver.rs index 05cac55..2a8dee6 100644 --- a/src_rbd/dynamics/solver.rs +++ b/src_rbd/dynamics/solver.rs @@ -149,6 +149,8 @@ pub struct SolverArgs<'a> { pub rb_contacts_inert: bool, /// Shared per-batch indices. pub batch_indices: &'a Tensor, + /// The one gravity uniform every rigid-body and multibody kernel reads. + pub gravity: &'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]>, @@ -278,6 +280,7 @@ impl GpuSolver { args.mprops, args.sim_params, args.batch_indices, + args.gravity, )?; } @@ -295,6 +298,7 @@ impl GpuSolver { contacts_len: args.contacts_len, solver_vels: &mut *args.solver_vels, batch_indices: args.batch_indices, + gravity: args.gravity, color_uniforms: args.color_uniforms, mb_sweep_indirect: args.mb_sweep_indirect, }; @@ -319,6 +323,7 @@ impl GpuSolver { contacts_len: args.contacts_len, solver_vels: &mut *args.solver_vels, batch_indices: args.batch_indices, + gravity: args.gravity, color_uniforms: args.color_uniforms, mb_sweep_indirect: args.mb_sweep_indirect, }; @@ -363,6 +368,7 @@ impl GpuSolver { contacts_len: args.contacts_len, solver_vels: &mut *args.solver_vels, batch_indices: args.batch_indices, + gravity: args.gravity, color_uniforms: args.color_uniforms, mb_sweep_indirect: args.mb_sweep_indirect, }; @@ -554,6 +560,8 @@ impl GpuSolver { } } + mb_phase!("[RBD] slv/mb-restitution", apply_restitution); + /* * Writeback body velocities and convert COM-centered solver poses * back to body-origin poses. diff --git a/src_rbd/pipeline/insertion_removal.rs b/src_rbd/pipeline/insertion_removal.rs index 2c9870b..2f7315f 100644 --- a/src_rbd/pipeline/insertion_removal.rs +++ b/src_rbd/pipeline/insertion_removal.rs @@ -101,12 +101,7 @@ impl RbdState { let mb_refs: Vec<_> = (0..num_batches as usize) .map(|_| (&empty_mb, &empty_body_ids, &empty_bodies)) .collect(); - let mut mb = GpuMultibodySet::from_rapier( - backend, - &mb_refs, - [0.0, -9.81, 0.0], - capacity_per_batch, - ); + let mut mb = GpuMultibodySet::from_rapier(backend, &mb_refs, capacity_per_batch); mb.set_constraint_softness(backend, &all_sim_params[0]); mb }; @@ -255,6 +250,7 @@ impl RbdState { joints, #[cfg(feature = "dim3")] multibodies, + gravity: Self::gravity_tensor(backend, [0.0, -9.81, 0.0]), body_group, local_mprops: Tensor::vector(backend, &all_local_mprops, rw).unwrap(), mprops: Tensor::vector(backend, &all_mprops, rw).unwrap(), diff --git a/src_rbd/pipeline/rbd_state.rs b/src_rbd/pipeline/rbd_state.rs index 001659d..3b5db97 100644 --- a/src_rbd/pipeline/rbd_state.rs +++ b/src_rbd/pipeline/rbd_state.rs @@ -221,6 +221,8 @@ pub struct RbdState { pub(super) joints: GpuImpulseJointSet, #[cfg(feature = "dim3")] pub(super) multibodies: GpuMultibodySet, + /// The one gravity uniform every rigid-body and multibody kernel reads. + pub(super) gravity: Tensor, /// Per-body "graph group" id, used by graph coloring to treat all bodies of /// the same multibody as a single node. For free bodies, `body_group[i] = i`; /// bodies of a multibody all share the group id of the root link, so two @@ -331,13 +333,29 @@ impl RbdState { self.collision_pairs.capacity() as u32 } - /// Uploads a new gravity vector for the multibody solver, e.g. - /// `[0.0, 0.0, -9.81]` for a Z-up scene. Affects multibody links; free - /// (non-multibody) rigid-bodies use a fixed gravity baked into the solver - /// shader. - #[cfg(feature = "dim3")] + /// Uploads a new gravity vector, e.g. `[0.0, 0.0, -9.81]` for a Z-up scene. + /// Every solver path reads this one uniform, so it applies to free + /// rigid-bodies and multibody links alike. In 2D the third component is + /// ignored. pub fn set_gravity(&mut self, backend: &GpuBackend, gravity: [f32; 3]) { - self.multibodies.set_gravity(backend, gravity); + self.gravity = Self::gravity_tensor(backend, gravity); + } + + /// The gravity uniform shared by every solver kernel. + pub fn gravity(&self) -> &Tensor { + &self.gravity + } + + pub(super) fn gravity_tensor( + backend: &GpuBackend, + gravity: [f32; 3], + ) -> Tensor { + Tensor::scalar( + backend, + glamx::Vec4::new(gravity[0], gravity[1], gravity[2], 0.0), + BufferUsages::STORAGE | BufferUsages::UNIFORM | BufferUsages::COPY_DST, + ) + .unwrap() } /// Per-collider world pose. @@ -363,6 +381,13 @@ impl RbdState { &self.multibodies } + /// Enables or disables the implicit treatment of multibody coriolis forces. + #[cfg(feature = "dim3")] + pub fn set_implicit_coriolis(&mut self, backend: &GpuBackend, enabled: bool) { + self.multibodies.set_implicit_coriolis(enabled); + self.rebuild_batch_indices(backend); + } + /// Returns a reference to the GPU buffer containing collision shapes. /// /// Each shape corresponds to one rigid body in the simulation. diff --git a/src_rbd/pipeline/rbd_state_from_rapier.rs b/src_rbd/pipeline/rbd_state_from_rapier.rs index 6dd8527..9bf22ca 100644 --- a/src_rbd/pipeline/rbd_state_from_rapier.rs +++ b/src_rbd/pipeline/rbd_state_from_rapier.rs @@ -454,12 +454,10 @@ impl RbdState { .iter() .map(|(mb, ids, bodies)| (*mb, ids, *bodies)) .collect(); - let mut mb = GpuMultibodySet::from_rapier( - backend, - &mb_refs, - [0.0, -9.81, 0.0], - max_colliders as u32, - ); + let mut mb = GpuMultibodySet::from_rapier(backend, &mb_refs, max_colliders as u32); + // `set_visible_dt` divides by the substep count, so that has to be + // in place first or the multibody integrates at the wrong rate. + mb.set_num_solver_iterations(num_solver_iterations); mb.set_visible_dt(backend, multibody_dt); // Soft contact coefficients (rapier TGS-soft) from the substep sim // params, so multibody-vs-floor contacts use the same soft ERP + CFM @@ -759,6 +757,7 @@ impl RbdState { joints, #[cfg(feature = "dim3")] multibodies, + gravity: RbdState::gravity_tensor(backend, [0.0, -9.81, 0.0]), body_group, local_mprops: Tensor::vector(backend, &all_local_mprops, storage).unwrap(), mprops: Tensor::vector(backend, &all_mprops, storage).unwrap(), diff --git a/src_rbd/pipeline/rbd_step.rs b/src_rbd/pipeline/rbd_step.rs index f2a83cf..b28397a 100644 --- a/src_rbd/pipeline/rbd_step.rs +++ b/src_rbd/pipeline/rbd_step.rs @@ -96,6 +96,7 @@ impl RbdPipeline { batch_indices: &state.batch_indices, color_uniforms: &state.color_uniforms, mb_sweep_indirect: &state.mb_sweep_indirect, + gravity: &state.gravity, }; self.multibody_solver.init_step( &mut encoder, @@ -306,6 +307,7 @@ impl RbdPipeline { colorless_warmstart: false, fused_color_sweeps, rb_contacts_inert: state.rb_contacts_inert, + gravity: &state.gravity, }; self.solver.prepare( backend, @@ -458,6 +460,7 @@ impl RbdPipeline { colorless_warmstart: true, fused_color_sweeps, rb_contacts_inert: state.rb_contacts_inert, + gravity: &state.gravity, }; // Phase 3: Solve constraints diff --git a/src_rbd_shaders/dynamics/joint_constraint_builder.rs b/src_rbd_shaders/dynamics/joint_constraint_builder.rs index a1e5707..4630a63 100644 --- a/src_rbd_shaders/dynamics/joint_constraint_builder.rs +++ b/src_rbd_shaders/dynamics/joint_constraint_builder.rs @@ -86,6 +86,20 @@ fn gcross_matrix(r: Vec3) -> Mat3 { ) } +/// Computes the smallest absolute difference between two half-angle sines, +/// i.e. the sine-space analog of [`smallest_abs_diff_between_angles`] where a +/// full turn maps to a span of 2. +pub(crate) fn smallest_abs_diff_between_sin_angles(a: f32, b: f32) -> f32 { + let s_err = a - b; + let sgn = if s_err < 0.0 { -1.0 } else { 1.0 }; + let s_err_complement = s_err - sgn * 2.0; + if s_err.abs() < s_err_complement.abs() { + s_err + } else { + s_err_complement + } +} + /// Computes the smallest absolute difference between two angles. fn smallest_abs_diff_between_angles(a: f32, b: f32) -> f32 { // Select the smallest path among the two angles to reach the target. diff --git a/src_rbd_shaders/dynamics/mod.rs b/src_rbd_shaders/dynamics/mod.rs index d1303a9..432533a 100644 --- a/src_rbd_shaders/dynamics/mod.rs +++ b/src_rbd_shaders/dynamics/mod.rs @@ -33,6 +33,7 @@ pub use joint::{ }; pub use joint_constraint::*; pub use joint_constraint_builder::{JointConstraintBuilder, JointConstraintHelper, new_helper}; +pub(crate) use joint_constraint_builder::smallest_abs_diff_between_sin_angles; pub use multibody::*; pub use sim_params::*; // Re-export solver items; update_constraint comes from joint_constraint_builder for joints diff --git a/src_rbd_shaders/dynamics/multibody/compute_dynamics_pre.rs b/src_rbd_shaders/dynamics/multibody/compute_dynamics_pre.rs index 880ee02..ac93c7a 100644 --- a/src_rbd_shaders/dynamics/multibody/compute_dynamics_pre.rs +++ b/src_rbd_shaders/dynamics/multibody/compute_dynamics_pre.rs @@ -16,9 +16,9 @@ use khal_std::sync::workgroup_memory_barrier_with_group_sync; 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, + WS_JOINT_ROT, WS_JOINT_VEL, WS_LTP, WS_LTW, WS_RB_VELS, WS_SHIFT02, WS_SHIFT23, WS_WORLD_COM, + 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; @@ -58,8 +58,9 @@ fn packed_decode(wg_id: UVec3, lid: UVec3, batch_ids: &BatchIndices) -> (u32, u3 (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. +/// Fused FK + body-jacobians + velocity propagation + CRBA mass matrix. +/// Coriolis blocks run only when implicit coriolis is enabled (non-zero +/// `mass_matrix_acc_section_offset`). #[spirv_bindgen(force_cpu_coroutines)] #[spirv(compute(threads(64, 1, 1)))] pub fn gpu_mb_compute_dynamics_pre( @@ -150,7 +151,10 @@ pub fn gpu_mb_compute_dynamics_pre( } sync_slots(t); - // 3) Mass matrices. + // 3) Mass matrices. `split` (implicit coriolis on) builds two matrices: + // the plain one for constraints and the coriolis-augmented acc section for + // the acceleration solve. Otherwise only the plain matrix is built and the + // coriolis blocks are skipped entirely (forces stay explicit). let acc_section = batch_ids.mass_matrix_acc_section_offset as usize; let split = acc_section != 0; let acc_augmented_mass = if split { @@ -180,8 +184,8 @@ pub fn gpu_mb_compute_dynamics_pre( inv_mass_x = lmp.inv_mass.x; - if inv_mass_x == 0.0 { - let coriolis_block = batch_ids.imat(batch_id, + if split && inv_mass_x == 0.0 { + let coriolis_block = batch_ids.imat(batch_id, mb_cor_base + (k as usize) * (DIM as usize) * (ndofs as usize), DIM, ndofs, @@ -242,7 +246,7 @@ pub fn gpu_mb_compute_dynamics_pre( t, ); - if k != 0 { + if split && k != 0 { let stat = stat_slice[k as usize]; let parent_id = stat.parent_link_id; let parent_j = batch_ids.imat(batch_id, @@ -349,7 +353,7 @@ pub fn gpu_mb_compute_dynamics_pre( sync_slots(t); - if loop_is_active { + if loop_is_active && split { if k != 0 { let stat = stat_slice[k as usize]; let parent_id = stat.parent_link_id; @@ -408,7 +412,7 @@ pub fn gpu_mb_compute_dynamics_pre( sync_slots(t); - if loop_is_active { + if loop_is_active && split { 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( @@ -451,7 +455,7 @@ pub fn gpu_mb_compute_dynamics_pre( sync_slots(t); - if loop_is_active { + if loop_is_active && split { // i_coriolis_dt assembly: dt · (mass·coriolis_v, I·coriolis_w). { let scale = mass * dt; @@ -496,7 +500,7 @@ pub fn gpu_mb_compute_dynamics_pre( sync_slots(t); - if loop_is_active { + if loop_is_active && split { gemm_tr_par( mass_matrices, acc_augmented_mass, @@ -617,6 +621,7 @@ fn forward_kinematics( }; ws_set_pose(ws, wa, 0, WS_LTP, root_pose); ws_set_pose(ws, wa, 0, WS_LTW, root_pose); + ws_set_vec(ws, wa, 0, WS_WORLD_COM, root_pose * root_config.local_mprops.com); for k in 1..num_links { let k_usize = k as usize; @@ -639,6 +644,7 @@ fn forward_kinematics( 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); + ws_set_vec(ws, wa, k, WS_WORLD_COM, world_com); poses_slice[stat.rb_id as usize] = local_to_world; } } diff --git a/src_rbd_shaders/dynamics/multibody/contact_constraints.rs b/src_rbd_shaders/dynamics/multibody/contact_constraints.rs index b23d6b2..aeb991b 100644 --- a/src_rbd_shaders/dynamics/multibody/contact_constraints.rs +++ b/src_rbd_shaders/dynamics/multibody/contact_constraints.rs @@ -13,23 +13,28 @@ //! 3. `gpu_mb_solve_contact_constraints` //! 4. `gpu_mb_remove_contact_constraint_bias` +use glamx::Vec4; 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::ConstraintSoftness; use crate::dynamics::body::{Velocity, WorldMassProperties}; use crate::dynamics::joint::SPATIAL_DIM; -use crate::queries::IndexedManifold; +use crate::queries::{IndexedManifold, MAX_MANIFOLD_POINTS}; use crate::utils::BatchIndices; -use crate::utils::linalg::{MatSlice, VSlice, lu_solve_in_place}; +use crate::utils::linalg::{MAX_MB_DOFS, MatSlice, VSlice, lu_solve_in_place}; use crate::{ANG_DIM, AngVector, DIM, Pose, Vector, gcross, gdot}; use super::types::{ - CONTACT_CONSTRAINTS_PER_POINT, MAX_MB_CONTACT_CONSTRAINTS_PER_MB, MAX_MB_CONTACTS_PER_MB, - MB_CONTACT_KIND_NORMAL, MB_CONTACT_KIND_TANGENT, MultibodyContactConstraint, MultibodyInfo, + CONTACT_CONSTRAINTS_PER_POINT, MAX_MB_CONTACT_CONSTRAINTS_PER_MB, + MB_CONTACT_KIND_INACTIVE, MB_CONTACT_KIND_NORMAL, MB_CONTACT_KIND_TANGENT, + MultibodyContactConstraint, MultibodyInfo, MultibodyLinkStatic, }; +use super::utils::zero_kinematic_dofs; +use super::ws_soa::{WS_LTW, WS_WORLD_COM, WsAddr, ws_pose, ws_vec}; #[cfg(feature = "dim2")] use glamx::Vec2; @@ -60,8 +65,7 @@ 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. +#[allow(clippy::too_many_arguments)] #[inline] fn fill_contact_jac_row( body_jacobians: &[f32], @@ -76,7 +80,6 @@ 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). @@ -84,8 +87,7 @@ fn fill_contact_jac_row( 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 { + for j in 0..ndofs { // Linear contribution: `unit_force · J_v[:, j]`. let dot; #[cfg(feature = "dim3")] @@ -123,9 +125,10 @@ fn fill_contact_jac_row( /// /// Pass 1: scans every contact in `contacts[batch]` and, for each contact /// point touching a link of this multibody, emits a normal-direction -/// `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. +/// `MultibodyContactConstraint` plus its friction slots. The multibody-side +/// `Jᵀ` rows are assembled later, by the finalize pass. 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)))] @@ -134,16 +137,17 @@ pub fn gpu_mb_init_contact_constraints( #[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], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] links_workspace: &[Vec4], #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] body_to_link: &[[u32; 2]], #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] contact_constraints: &mut [MultibodyContactConstraint], - #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] contact_constraint_jacs: &mut [f32], - #[spirv(uniform, descriptor_set = 0, binding = 5)] softness: &ConstraintSoftness, + #[spirv(uniform, descriptor_set = 0, binding = 4)] softness: &ConstraintSoftness, #[spirv(storage_buffer, descriptor_set = 1, binding = 0)] mprops: &[WorldMassProperties], #[spirv(storage_buffer, descriptor_set = 1, binding = 1)] poses: &[Pose], #[spirv(storage_buffer, descriptor_set = 1, binding = 2)] contacts: &[IndexedManifold], - #[spirv(uniform, descriptor_set = 0, binding = 6)] batch_ids: &BatchIndices, + #[spirv(storage_buffer, descriptor_set = 1, binding = 3)] solver_body_poses: &[Pose], + #[spirv(uniform, descriptor_set = 0, binding = 5)] batch_ids: &BatchIndices, + #[spirv(uniform, descriptor_set = 0, binding = 6)] first_substep: &u32, ) { // Only active multibody slots are visited now; the `ndofs == 0` sentinel // below is kept for all-locked (zero-dof) multibodies. Padding slots past @@ -158,9 +162,12 @@ pub fn gpu_mb_init_contact_constraints( let inv_dt = softness.inv_dt; let max_corr_velocity = softness.max_corr_velocity; + let warmstart_coeff = softness.warmstart_coefficient; + // On the first substep of a step the anchors are (re)frozen from the fresh + // manifold, afterwards they are updated normally. + let freeze_anchors = *first_substep != 0; 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); // `body_to_link` is laid out with stride = colliders_batch_capacity. let b2l_start = colliders_start; @@ -178,21 +185,18 @@ pub fn gpu_mb_init_contact_constraints( } return; } - 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 - // overlap between multibodies of differing `ndofs`). - 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 wa = WsAddr::new(mb.first_link as usize, batch_ids.num_batches, batch_id); let contacts_slice = batch_ids.contact_batch(batch_id, contacts); let n_contacts = mb.batch_contacts_len.min(batch_ids.contacts_batch_capacity); + let prev_count = mb.contact_constraint_count; let mut count = 0u32; for ci in 0..n_contacts { - if count >= MAX_MB_CONTACTS_PER_MB { + if count + (MAX_MANIFOLD_POINTS as u32) * CONTACT_CONSTRAINTS_PER_POINT + > MAX_MB_CONTACT_CONSTRAINTS_PER_MB + { break; } let im = contacts_slice[ci as usize]; @@ -200,7 +204,6 @@ pub fn gpu_mb_init_contact_constraints( continue; } let id1 = im.colliders.x; - let id2 = im.colliders.y; let b1 = im.bodies.x; let b2 = im.bodies.y; @@ -240,7 +243,6 @@ pub fn gpu_mb_init_contact_constraints( } let pose1 = poses.read(colliders_start + id1 as usize); - let pose2 = poses.read(colliders_start + id2 as usize); let world_normal = pose1.rotation * im.contact.normal_a; let lin_jac = if is_self || mb_on_1 { world_normal @@ -268,23 +270,40 @@ pub fn gpu_mb_init_contact_constraints( softness.cfm_factor }; - // Multibody-link origins come from the collider poses buffer instead - // of `links_workspace`. - let link_origin_a = if is_self || mb_on_1 { - pose1.translation + // The body jacobian of a link measures its velocity at the link's + // center of mass, so every torque arm is taken from there. + let com_a = ws_vec(links_workspace, wa, mb_link_id_a, WS_WORLD_COM); + let pose_a = ws_pose(links_workspace, wa, mb_link_id_a, WS_LTW); + // Frame the other side's anchor lives in: the second link for a + // self-contact, the free body's center-of-mass solver pose otherwise. + let pose_b = if is_self { + ws_pose(links_workspace, wa, mb_link_id_b, WS_LTW) } else { - pose2.translation + solver_body_poses.read(colliders_start + free_body_id as usize) }; - let link_origin_b_default = link_origin_a; for k in 0..im.contact.len { // One contact point produces 1 normal + (DIM-1) friction slots. if count + CONTACT_CONSTRAINTS_PER_POINT > MAX_MB_CONTACT_CONSTRAINTS_PER_MB { break; } - let pt_local = im.contact.points_a.read(k as usize).pt; - let dist = im.contact.points_a.read(k as usize).dist; - let pt_world = pose1 * (pt_local + im.contact.normal_a * (dist * 0.5)); + let normal_slot = count; + let prev = contact_constraints.read(cons_base + normal_slot as usize); + + // Re-resolve both anchors through the current poses, then track the + // separation as their drift along the contact normal. + let (local_p1, local_p2, base_dist) = if freeze_anchors { + let pt_local = im.contact.points_a.read(k as usize).pt; + let d = im.contact.points_a.read(k as usize).dist; + let pt = pose1 * (pt_local + im.contact.normal_a * (d * 0.5)); + (pose_a.inverse() * pt, pose_b.inverse() * pt, d) + } else { + (prev.local_p1, prev.local_p2, prev.base_dist) + }; + let p1 = pose_a * local_p1; + let p2 = pose_b * local_p2; + let pt_world = p1; + let dist = base_dist + (p1 - p2).dot(mb_normal); // Tangent basis — matches rapier's fallback path // (`OrthonormalBasis::orthonormal_vector(force_dir1)` then @@ -297,7 +316,7 @@ pub fn gpu_mb_init_contact_constraints( // A-side (link `mb_link_id_a`, rapier's body 1): impulse along // `force_dir1 = -world_normal_a = mb_normal`. - let shift_a = pt_world - link_origin_a; + let shift_a = pt_world - com_a; let torque_a_normal = gcross(shift_a, mb_normal); let torque_a_t0 = gcross(shift_a, mb_tangent0); #[cfg(feature = "dim3")] @@ -306,71 +325,38 @@ pub fn gpu_mb_init_contact_constraints( let rhs_bias = (erp_inv_dt * dist).clamp(-max_corr_velocity, 0.0); let rhs_wo_bias = if dist > 0.0 { dist * inv_dt } else { 0.0 }; - let normal_slot = count; - let normal_col_offset = col_base + (normal_slot as usize) * dofs_stride; - let warmstart_normal_impulse = if lane == 0 { - contact_constraints - .read(cons_base + normal_slot as usize) - .impulse - } else { + let warmstart_normal_impulse = if freeze_anchors { 0.0 + } else { + prev.impulse * warmstart_coeff }; - fill_contact_jac_row( - body_jacobians, - mb_jac_base, - batch_ids.num_batches, - batch_id, - ndofs, - mb_link_id_a, - mb_normal, - torque_a_normal, - contact_constraint_jacs, - normal_col_offset, - false, - lane, - ); - // B-side fold-in for self-contacts, free body for the rest. The // ang_jac fields below describe the FREE body side; for self // contacts they collapse to zero because both sides are folded // into `J_mb` already. - let (ang_jac_normal, ii_ang_jac_normal) = if is_self { - // Self-contact: B-side link is the collider at `id2`, so its - // world pose is `pose2` (already loaded). Avoids a - // `links_workspace` binding. - let link_origin_b = pose2.translation; - let _ = link_origin_b_default; - let shift_b = pt_world - link_origin_b; - let torque_b_normal = gcross(shift_b, lin_jac); - fill_contact_jac_row( - body_jacobians, - mb_jac_base, - batch_ids.num_batches, - batch_id, - ndofs, - mb_link_id_b, - lin_jac, - torque_b_normal, - contact_constraint_jacs, - normal_col_offset, - true, - lane, - ); + let (torque_b_normal, ang_jac_normal, ii_ang_jac_normal) = if is_self { + let shift_b = p2 - ws_vec(links_workspace, wa, mb_link_id_b, WS_WORLD_COM); #[cfg(feature = "dim3")] { - (AngVector::ZERO, AngVector::ZERO) + (gcross(shift_b, lin_jac), AngVector::ZERO, AngVector::ZERO) } #[cfg(feature = "dim2")] { - (0.0f32, 0.0f32) + (gcross(shift_b, lin_jac), 0.0f32, 0.0f32) } } else { - let _ = link_origin_b_default; - let free_shift = pt_world - free_mp.com; + let free_shift = p2 - pose_b.translation; let aj = gcross(free_shift, lin_jac); let iiaj = free_mp.inv_inertia_mul(aj); - (aj, iiaj) + #[cfg(feature = "dim3")] + { + (AngVector::ZERO, aj, iiaj) + } + #[cfg(feature = "dim2")] + { + (0.0f32, aj, iiaj) + } }; // Normal constraint slot. @@ -383,7 +369,7 @@ pub fn gpu_mb_init_contact_constraints( free_body_im: free_im, friction_coeff: im.friction, normal_constraint_slot: normal_slot, - _pad0: 0, + link_id_b: mb_link_id_b, lin_jac, _pad1: 0, ang_jac: ang_jac_normal, @@ -395,8 +381,17 @@ pub fn gpu_mb_init_contact_constraints( rhs_wo_bias, impulse: warmstart_normal_impulse, cfm_factor, - _unused_cfm: 0.0, - _pad4: [0; 2], + restitution_seed: prev.restitution_seed, + restitution: im.restitution, + _pad4: 0, + torque_a: torque_a_normal, + _pad5: 0, + torque_b: torque_b_normal, + _pad6: 0, + local_p1, + base_dist, + local_p2, + _pad7: 0, }; #[cfg(feature = "dim2")] let normal_cons = MultibodyContactConstraint { @@ -409,14 +404,21 @@ pub fn gpu_mb_init_contact_constraints( ii_ang_jac: ii_ang_jac_normal, friction_coeff: im.friction, normal_constraint_slot: normal_slot, - _pad0: [0; 1], + link_id_b: mb_link_id_b, lin_jac, inv_lhs: 0.0, rhs: rhs_wo_bias + rhs_bias, rhs_wo_bias, impulse: warmstart_normal_impulse, cfm_factor, - _unused_cfm: 0.0, + restitution_seed: prev.restitution_seed, + restitution: im.restitution, + torque_a: torque_a_normal, + torque_b: torque_b_normal, + local_p1, + local_p2, + base_dist, + _pad1: [0.0; 2], }; if lane == 0 { contact_constraints.write(cons_base + normal_slot as usize, normal_cons); @@ -456,66 +458,48 @@ pub fn gpu_mb_init_contact_constraints( }; let free_tangent = -mb_tangent; 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; lane 0 only). - let warmstart_tang_impulse = if lane == 0 { - contact_constraints - .read(cons_base + tang_slot as usize) - .impulse - } else { + let tang_prev = contact_constraints.read(cons_base + tang_slot as usize); + let warmstart_tang_impulse = if freeze_anchors { 0.0 + } else { + tang_prev.impulse * warmstart_coeff }; - fill_contact_jac_row( - body_jacobians, - mb_jac_base, - batch_ids.num_batches, - batch_id, - ndofs, - mb_link_id_a, - mb_tangent, - torque_a_tang, - contact_constraint_jacs, - tang_col_offset, - false, - lane, - ); - - let (ang_jac_tang, ii_ang_jac_tang) = if is_self { - let link_origin_b = pose2.translation; - let shift_b = pt_world - link_origin_b; - let torque_b_tang = gcross(shift_b, free_tangent); - fill_contact_jac_row( - body_jacobians, - mb_jac_base, - batch_ids.num_batches, - batch_id, - ndofs, - mb_link_id_b, - free_tangent, - torque_b_tang, - contact_constraint_jacs, - tang_col_offset, - true, - lane, - ); + let (torque_b_tang, ang_jac_tang, ii_ang_jac_tang) = if is_self { + let shift_b = + p2 - ws_vec(links_workspace, wa, mb_link_id_b, WS_WORLD_COM); #[cfg(feature = "dim3")] { - (AngVector::ZERO, AngVector::ZERO) + ( + gcross(shift_b, free_tangent), + AngVector::ZERO, + AngVector::ZERO, + ) } #[cfg(feature = "dim2")] { - (0.0f32, 0.0f32) + (gcross(shift_b, free_tangent), 0.0f32, 0.0f32) } } else { - let free_shift = pt_world - free_mp.com; + let free_shift = p2 - pose_b.translation; let aj = gcross(free_shift, free_tangent); let iiaj = free_mp.inv_inertia_mul(aj); - (aj, iiaj) + #[cfg(feature = "dim3")] + { + (AngVector::ZERO, aj, iiaj) + } + #[cfg(feature = "dim2")] + { + (0.0f32, aj, iiaj) + } }; - // No surface velocity (TODO: conveyor belts) → rhs = 0. + // Positional bias along the tangent: pull the two anchors back + // together so friction sticks instead of drifting. No surface + // velocity yet (TODO: conveyor belts), so `rhs_wo_bias` is 0. + let tang_bias = (p1 - p2).dot(mb_tangent) * inv_dt; #[cfg(feature = "dim3")] let tang_cons = MultibodyContactConstraint { multibody_id: mb_idx, @@ -525,7 +509,7 @@ pub fn gpu_mb_init_contact_constraints( free_body_im: free_im, friction_coeff: im.friction, normal_constraint_slot: normal_slot, - _pad0: 0, + link_id_b: mb_link_id_b, lin_jac: free_tangent, _pad1: 0, ang_jac: ang_jac_tang, @@ -533,12 +517,21 @@ pub fn gpu_mb_init_contact_constraints( ii_ang_jac: ii_ang_jac_tang, _pad3: 0, inv_lhs: 0.0, - rhs: 0.0, + rhs: tang_bias, rhs_wo_bias: 0.0, impulse: warmstart_tang_impulse, cfm_factor, - _unused_cfm: 0.0, - _pad4: [0; 2], + restitution_seed: tang_prev.restitution_seed, + restitution: im.restitution, + _pad4: 0, + torque_a: torque_a_tang, + _pad5: 0, + torque_b: torque_b_tang, + _pad6: 0, + local_p1, + base_dist, + local_p2, + _pad7: 0, }; #[cfg(feature = "dim2")] let tang_cons = MultibodyContactConstraint { @@ -551,14 +544,21 @@ pub fn gpu_mb_init_contact_constraints( ii_ang_jac: ii_ang_jac_tang, friction_coeff: im.friction, normal_constraint_slot: normal_slot, - _pad0: [0; 1], + link_id_b: mb_link_id_b, lin_jac: free_tangent, inv_lhs: 0.0, - rhs: 0.0, + rhs: tang_bias, rhs_wo_bias: 0.0, impulse: warmstart_tang_impulse, cfm_factor, - _unused_cfm: 0.0, + restitution_seed: tang_prev.restitution_seed, + restitution: im.restitution, + torque_a: torque_a_tang, + torque_b: torque_b_tang, + local_p1, + local_p2, + base_dist, + _pad1: [0.0; 2], }; if lane == 0 { contact_constraints.write(cons_base + tang_slot as usize, tang_cons); @@ -568,9 +568,15 @@ 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. + // The solve kernels only iterate `0..count`, but next frame's warmstart + // match scans the whole slab, so the leftovers of the previous build have + // to be marked inactive. if lane == 0 { + for s in count..prev_count.min(MAX_MB_CONTACT_CONSTRAINTS_PER_MB) { + let mut stale = contact_constraints.read(cons_base + s as usize); + stale.kind = MB_CONTACT_KIND_INACTIVE; + contact_constraints.write(cons_base + s as usize, stale); + } mb.contact_constraint_count = count; multibody_info.write(batch_ids.mbi(batch_id, mb_idx as usize), mb); } @@ -600,19 +606,20 @@ pub fn gpu_mb_stash_contacts_len( multibody_info.write(batch_ids.mbi(batch_id, 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 -/// `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). +/// Snapshot every contact-constraint slot into the "previous frame" slab that +/// `gpu_mb_transfer_contact_warmstart` matches against. Called once per visible +/// frame from `init_step`, before the substep loop rebuilds the live slab. +/// +/// One thread per (slot, multibody, batch). #[spirv_bindgen] #[spirv(compute(threads(64)))] -pub fn gpu_mb_reset_contact_warmstart( +pub fn gpu_mb_snapshot_contact_warmstart( #[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 = 1)] batch_ids: &BatchIndices, + contact_constraints: &[MultibodyContactConstraint], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] + old_contact_constraints: &mut [MultibodyContactConstraint], + #[spirv(uniform, descriptor_set = 0, binding = 2)] batch_ids: &BatchIndices, ) { // One thread per (slot, multibody, batch), flattened. const MAXC: u32 = MAX_MB_CONTACT_CONSTRAINTS_PER_MB; @@ -628,7 +635,7 @@ pub fn gpu_mb_reset_contact_warmstart( let cons_start = batch_ids.mb_contact_constraints_start(batch_id); let idx = cons_start + (mb_idx * MAXC + s) as usize; - contact_constraints.at_mut(idx).impulse = 0.0; + old_contact_constraints.write(idx, contact_constraints.read(idx)); } /// Warmstart: re-apply each active contact constraint's accumulated `impulse` @@ -713,9 +720,9 @@ pub fn gpu_mb_warmstart_contact_constraints( } } -/// Pass 2: for each emitted constraint, LU back-solve `M · column = Jᵀ` -/// (the row produced by the init kernel) and set `inv_lhs = 1 / (Jᵀ · -/// column + free_body_inv_r)`. +/// Pass 2: for each emitted constraint, build its multibody-side `Jᵀ` row, +/// LU back-solve `M · column = Jᵀ` and set `inv_lhs = 1 / (Jᵀ · column + +/// free_body_inv_r)`. #[spirv_bindgen] #[spirv(compute(threads(64)))] pub fn gpu_mb_finalize_contact_constraints( @@ -726,10 +733,13 @@ pub fn gpu_mb_finalize_contact_constraints( #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] lu_pivots: &[u32], #[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 = 4)] contact_constraint_jacs: &mut [f32], #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] contact_constraint_columns: &mut [f32], - #[spirv(uniform, descriptor_set = 0, binding = 6)] batch_ids: &BatchIndices, + #[spirv(storage_buffer, descriptor_set = 0, binding = 6)] + links_static: &[MultibodyLinkStatic], + #[spirv(storage_buffer, descriptor_set = 0, binding = 7)] body_jacobians: &[f32], + #[spirv(uniform, descriptor_set = 0, binding = 8)] batch_ids: &BatchIndices, ) { const LANES: u32 = 64; let batch_id = workgroup_id.y; @@ -749,6 +759,7 @@ pub fn gpu_mb_finalize_contact_constraints( return; } let mb_mm_base = mb.mass_matrix_offset as usize; + let mb_jac_base = mb.jacobian_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; @@ -757,16 +768,53 @@ pub fn gpu_mb_finalize_contact_constraints( let m = batch_ids.imat(batch_id, mb_mm_base, ndofs, ndofs); let count = mb.contact_constraint_count; + let stat_slice = batch_ids + .ib(batch_id, links_static) + .offset(mb.first_link as usize); 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 + let mut cons = contact_constraints.read(cons_base + s as usize); + let is_self = cons.free_body_id == u32::MAX; + + // 1) Build the multibody-side Jᵀ row from the stored contact wrench, + // folding both touched links in for a self-contact. + fill_contact_jac_row( + body_jacobians, + mb_jac_base, + batch_ids.num_batches, + batch_id, + ndofs, + cons.link_id, + -cons.lin_jac, + cons.torque_a, + contact_constraint_jacs, + col_offset, + false, + ); + if is_self { + fill_contact_jac_row( + body_jacobians, + mb_jac_base, + batch_ids.num_batches, + batch_id, + ndofs, + cons.link_id_b, + cons.lin_jac, + cons.torque_b, + contact_constraint_jacs, + col_offset, + true, + ); + } + + // 2) Copy J^T row into the column buffer (it'll be overwritten by the // LU solve with the M⁻¹·Jᵀ result). for i in 0..ndofs { let v = contact_constraint_jacs.read(col_offset + i as usize); contact_constraint_columns.write(col_offset + i as usize, v); } - // 2) Solve M · column = J^T (in place). + // 3) Solve M · column = J^T (in place). lu_solve_in_place( mass_matrices, m, @@ -775,18 +823,23 @@ pub fn gpu_mb_finalize_contact_constraints( contact_constraint_columns, VSlice::dense(col_offset), ); - // 3) inv_r_mb = J · column. + // 3b) Kinematic dofs are user-driven: the impulse must not move them. + zero_kinematic_dofs( + contact_constraint_columns, + col_offset, + &stat_slice, + mb.num_links, + ); + // 4) inv_r_mb = J · column. let mut inv_r_mb = 0.0f32; for i in 0..ndofs { let j = contact_constraint_jacs.read(col_offset + i as usize); let c = contact_constraint_columns.read(col_offset + i as usize); inv_r_mb += j * c; } - // 4) Add free body's contribution: im (since lin_jac is unit) + + // 5) Add free body's contribution: im (since lin_jac is unit) + // ang_jac · ii_ang_jac. For self-contacts the B-side is folded into // `J_mb`, so there's no free-body term. - let mut cons = contact_constraints.read(cons_base + s as usize); - let is_self = cons.free_body_id == u32::MAX; let inv_r_free = if is_self { 0.0 } else { @@ -797,3 +850,287 @@ pub fn gpu_mb_finalize_contact_constraints( contact_constraints.write(cons_base + s as usize, cons); } } + + +/// Carries the accumulated contact impulses of the previous frame over to this +/// frame's freshly built slots. A point is matched by the pair of links (or +/// link and free body) it touches plus the proximity of both frozen local +/// anchors; friction is re-projected through world space onto the new tangent +/// basis. Runs once per frame, right after the first build. +/// +/// One 64-lane workgroup per (multibody, batch). +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn gpu_mb_transfer_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)] + contact_constraints: &mut [MultibodyContactConstraint], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] + old_contact_constraints: &[MultibodyContactConstraint], + #[spirv(uniform, descriptor_set = 0, binding = 3)] batch_ids: &BatchIndices, + #[spirv(uniform, descriptor_set = 0, binding = 4)] softness: &ConstraintSoftness, +) { + const LANES: u32 = 64; + // Anchors this far apart (in each side's own frame) are taken to be the + // same contact point from one frame to the next. + const MATCH_DIST: f32 = 1.0e-1; + + let batch_id = workgroup_id.y; + let mb_idx = workgroup_id.x; + let lane = local_id.x; + if mb_idx >= batch_ids.multibodies_len { + return; + } + + let mb = multibody_info.read(batch_ids.mbi(batch_id, mb_idx as usize)); + let count = mb.contact_constraint_count; + if mb.ndofs == 0 || count == 0 { + return; + } + + let cons_base = batch_ids.mb_contact_constraints_start(batch_id) + + (mb_idx as usize) * (MAX_MB_CONTACT_CONSTRAINTS_PER_MB as usize); + let sq_threshold = MATCH_DIST * MATCH_DIST; + let warmstart_coeff = softness.warmstart_coefficient; + + // Each lane owns whole contact points (their normal slot), so the tangent + // writes below stay disjoint. + for s in StepRng::new(lane..count, LANES) { + let mut cons = contact_constraints.read(cons_base + s as usize); + if cons.kind != MB_CONTACT_KIND_NORMAL { + continue; + } + + for j in 0..MAX_MB_CONTACT_CONSTRAINTS_PER_MB { + let old = old_contact_constraints.read(cons_base + j as usize); + if old.kind != MB_CONTACT_KIND_NORMAL + || old.link_id != cons.link_id + || old.link_id_b != cons.link_id_b + || old.free_body_id != cons.free_body_id + { + continue; + } + let d1 = old.local_p1 - cons.local_p1; + let d2 = old.local_p2 - cons.local_p2; + if d1.dot(d1) >= sq_threshold || d2.dot(d2) >= sq_threshold { + continue; + } + + cons.impulse = old.impulse * warmstart_coeff; + contact_constraints.write(cons_base + s as usize, cons); + + // Friction rows follow their normal row contiguously. + #[cfg(feature = "dim3")] + { + let old_t0 = old_contact_constraints.read(cons_base + (j + 1) as usize); + let old_t1 = old_contact_constraints.read(cons_base + (j + 2) as usize); + let world = (-old_t0.lin_jac * old_t0.impulse - old_t1.lin_jac * old_t1.impulse) + * warmstart_coeff; + + let mut new_t0 = contact_constraints.read(cons_base + (s + 1) as usize); + let mut new_t1 = contact_constraints.read(cons_base + (s + 2) as usize); + new_t0.impulse = world.dot(-new_t0.lin_jac); + new_t1.impulse = world.dot(-new_t1.lin_jac); + contact_constraints.write(cons_base + (s + 1) as usize, new_t0); + contact_constraints.write(cons_base + (s + 2) as usize, new_t1); + } + #[cfg(feature = "dim2")] + { + let old_t0 = old_contact_constraints.read(cons_base + (j + 1) as usize); + let mut new_t0 = contact_constraints.read(cons_base + (s + 1) as usize); + new_t0.impulse = old_t0.impulse * warmstart_coeff; + contact_constraints.write(cons_base + (s + 1) as usize, new_t0); + } + break; + } + } +} + +/// Captures each bouncy contact point's approaching normal velocity at the +/// start of the step, so [`gpu_mb_apply_contact_restitution`] can drive the +/// point back to it once every substep is done. Dispatched once per step, +/// right after the first `finalize`. +/// +/// One 64-lane workgroup per (multibody, batch). +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn gpu_mb_seed_contact_restitution( + #[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)] dof_state: &[f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] solver_vels: &[Velocity], + #[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; + if mb_idx >= batch_ids.multibodies_len { + return; + } + + 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 { + return; + } + + let colliders_start = batch_ids.coll_start(batch_id); + let v_base = mb.first_dof as usize; + let cons_base = batch_ids.mb_contact_constraints_start(batch_id) + + (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 = batch_ids.mb_contact_constraint_columns_start(batch_id) + + (mb_idx as usize) * (MAX_MB_CONTACT_CONSTRAINTS_PER_MB as usize) * dofs_stride; + + for s in StepRng::new(lane..count, LANES) { + let mut cons = contact_constraints.read(cons_base + s as usize); + if cons.kind != MB_CONTACT_KIND_NORMAL { + continue; + } + let jac_off = col_base + (s as usize) * dofs_stride; + let mut j_dot_v = 0.0f32; + for i in 0..ndofs { + j_dot_v += contact_constraint_jacs.read(jac_off + i as usize) + * dof_state.read(batch_ids.mbi(batch_id, v_base + i as usize)); + } + if cons.free_body_id != u32::MAX { + let free = solver_vels.read(colliders_start + cons.free_body_id as usize); + j_dot_v += cons.lin_jac.dot(free.linear) + gdot(cons.ang_jac, free.angular); + } + + // A fresh contact bounces whenever restitution is non-zero; one that + // survived the previous step is resting unless restitution is maximal. + let is_new = cons.impulse == 0.0; + let bouncy = if is_new { + cons.restitution > 0.0 + } else { + cons.restitution >= 1.0 + }; + cons.restitution_seed = if bouncy { cons.restitution * j_dot_v } else { 0.0 }; + contact_constraints.write(cons_base + s as usize, cons); + } +} + +/// End-of-step restitution pass: drives every bouncy point that carried an +/// impulse back to its seeded approach velocity, by re-running the normal +/// solve with `rhs = restitution_seed` and no CFM. +/// +/// One 64-lane workgroup per (multibody, batch). +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn gpu_mb_apply_contact_restitution( + #[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(uniform, descriptor_set = 0, binding = 4)] 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; 64], + #[spirv(workgroup)] delta_shared: &mut f32, +) { + let batch_id = workgroup_id.y; + let mb_idx = workgroup_id.x; + let lane = local_id.x; + if mb_idx >= batch_ids.multibodies_len { + return; + } + + 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 { + return; + } + + let colliders_start = batch_ids.coll_start(batch_id); + let v_base = mb.first_dof as usize; + let cons_base = batch_ids.mb_contact_constraints_start(batch_id) + + (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 = batch_ids.mb_contact_constraint_columns_start(batch_id) + + (mb_idx as usize) * (MAX_MB_CONTACT_CONSTRAINTS_PER_MB as usize) * dofs_stride; + + if lane < ndofs { + dof_v[lane as usize] = dof_state.read(batch_ids.mbi(batch_id, v_base + lane as usize)); + } + workgroup_memory_barrier_with_group_sync(); + + for s in 0..count { + let cons = contact_constraints.read(cons_base + s as usize); + // Only approaching, load-bearing points bounce. + if cons.kind != MB_CONTACT_KIND_NORMAL + || cons.restitution_seed >= 0.0 + || cons.impulse <= 0.0 + { + continue; + } + let col_offset = col_base + (s as usize) * dofs_stride; + let is_self = cons.free_body_id == u32::MAX; + + 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]; + } + 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 raw = cons.impulse - cons.inv_lhs * (j_dot_v + cons.restitution_seed); + let new_imp = if raw < 0.0 { 0.0 } else { raw }; + let delta = new_imp - cons.impulse; + *delta_shared = delta; + + let mut updated = cons; + updated.impulse = new_imp; + contact_constraints.write(cons_base + s as usize, updated); + + 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 = *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(); + } + + if lane < ndofs { + dof_state.write( + batch_ids.mbi(batch_id, v_base + lane as usize), + dof_v[lane as usize], + ); + } +} diff --git a/src_rbd_shaders/dynamics/multibody/gravity_and_lu.rs b/src_rbd_shaders/dynamics/multibody/gravity_and_lu.rs index 62aa99d..2ce9953 100644 --- a/src_rbd_shaders/dynamics/multibody/gravity_and_lu.rs +++ b/src_rbd_shaders/dynamics/multibody/gravity_and_lu.rs @@ -31,7 +31,7 @@ use super::lu::{ 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_coord, - ws_pose, ws_set_vel, ws_vec, ws_vel, ws_vel_ang, ws_world_inertia, + ws_ext_wrench, ws_pose, ws_set_vel, ws_vec, ws_vel, ws_vel_ang, ws_world_inertia, }; /// Adds the per-DoF joint-spring generalized forces (rapier's @@ -242,12 +242,11 @@ pub fn gpu_mb_gravity_and_lu( let gyroscopic: AngVector = 0.0; let i_acc_ang = rb_inertia * acc_ang; + let (ext_force, ext_torque, gravity_scale) = + ws_ext_wrench(links_workspace, wa, k); - #[cfg(feature = "dim3")] - let f_lin = (g - acc_lin) * mass; - #[cfg(feature = "dim2")] - let f_lin = (g - acc_lin) * mass; - let f_ang = -gyroscopic - i_acc_ang; + let f_lin = g * (mass * gravity_scale) + ext_force - acc_lin * mass; + let f_ang = ext_torque - gyroscopic - i_acc_ang; let body_jacobian = batch_ids.imat(batch_id, mb_jac_base + (k as usize) * SPATIAL_DIM * (ndofs as usize), @@ -343,8 +342,6 @@ pub fn gpu_mb_gravity_and_lu( inv_akk_shared, ); - // Legacy mode only: this factor doubles as the constraints' LU; persist - // it (joint / contact constraint init reuses it for unit-RHS solves). if !split && lane < ndofs { for r in 0..ndofs { mass_matrices.write(m_view.idx(r, lane), mat.read(sm_idx(r, lane))); @@ -555,9 +552,11 @@ fn gravity_and_lu_packed_impl(slot, r, lane))); @@ -849,9 +846,10 @@ pub fn gpu_mb_gravity_and_lu_t1( let gyroscopic: AngVector = 0.0; let i_acc_ang = rb_inertia * acc_ang; + let (ext_force, ext_torque, gravity_scale) = ws_ext_wrench(links_workspace, wa, k); - let f_lin = (g - acc_lin) * mass; - let f_ang = -gyroscopic - i_acc_ang; + let f_lin = g * (mass * gravity_scale) + ext_force - acc_lin * mass; + let f_ang = ext_torque - gyroscopic - i_acc_ang; let body_jacobian = batch_ids.imat(batch_id, mb_jac_base + (k as usize) * SPATIAL_DIM * (ndofs as usize), diff --git a/src_rbd_shaders/dynamics/multibody/impulse_joint_constraints/helper.rs b/src_rbd_shaders/dynamics/multibody/impulse_joint_constraints/helper.rs index 49ecc03..477d6d9 100644 --- a/src_rbd_shaders/dynamics/multibody/impulse_joint_constraints/helper.rs +++ b/src_rbd_shaders/dynamics/multibody/impulse_joint_constraints/helper.rs @@ -8,6 +8,7 @@ use khal_std::index::MaybeIndexUnchecked; use crate::dynamics::body::WorldMassProperties; use crate::dynamics::joint::JointMotor; +use crate::dynamics::smallest_abs_diff_between_sin_angles; use crate::{AngVector, MAX_FLT, Pose, Vector, rotation_to_matrix}; use super::super::types::MultibodyInfo; @@ -37,6 +38,10 @@ pub(super) struct JointConstraintHelper { ang_err: crate::Rotation, #[cfg(feature = "dim3")] ang_err: [f32; 3], + /// Real (scalar) part of the angular error quaternion, sign-matched to + /// `ang_err`. Needed to recover the re-centered angle of a limit row. + #[cfg(feature = "dim3")] + ang_err_w: f32, } #[cfg(feature = "dim3")] @@ -137,8 +142,44 @@ pub(super) fn new_helper( ang_basis, lin_err, ang_err, + ang_err_w: quat_err.w * sgn, + } + } +} + +/// Row parameters of an angular limit allowing `[min, max]` (radians): the row +/// measures the angle from the middle of the range against `±half_range`. +pub(super) struct AngularLimitParams { + /// Middle of the allowed range. + center: f32, + /// Half the allowed range. Larger than π means the row is unconstrained. + half_range: f32, +} + +impl AngularLimitParams { + pub(super) fn new(min: f32, max: f32) -> Self { + let half_range = (max - min) * 0.5; + // A range of a full turn or more is indistinguishable from "no limit" + // for an angle read off a relative rotation. `!(x < PI)` also catches + // NaN bounds before they poison the row. + if !(half_range < core::f32::consts::PI) { + return Self { + center: 0.0, + half_range: 10.0, + }; + } + + Self { + center: (min + max) * 0.5, + half_range, } } + + /// Symmetric bounds the re-centered angle is tested against. + #[inline] + pub(super) fn bounds(&self) -> [f32; 2] { + [-self.half_range, self.half_range] + } } #[inline] @@ -343,16 +384,52 @@ impl JointConstraintHelper { #[inline] #[cfg(feature = "dim2")] - pub(super) fn motor_ang_jac(&self, _axis: usize) -> AngVector { + pub(super) fn axis_ang_jac(&self, _axis: usize) -> AngVector { 1.0 } #[inline] #[cfg(feature = "dim3")] - pub(super) fn motor_ang_jac(&self, axis: usize) -> AngVector { + pub(super) fn axis_ang_jac(&self, axis: usize) -> AngVector { self.basis.col(axis) } + /// Angle of the relative rotation about `axis`, measured from the middle of + /// the limit's range and wrapped to `(-π, π]`. Its gradient is the plain + /// joint axis, so it pairs with [`Self::axis_ang_jac`]. + #[inline] + #[cfg(feature = "dim3")] + pub(super) fn recentered_angle(&self, axis: usize, limit: &AngularLimitParams) -> f32 { + let c_cos = crate::cos(limit.center * 0.5); + let c_sin = crate::sin(limit.center * 0.5); + let x = self.ang_err.read(axis); + let w = self.ang_err_w; + let sin_half = c_cos * x - c_sin * w; + let cos_half = c_cos * w + c_sin * x; + // The re-centered HALF angle; doubling it must wrap back to `(-π, π]`, + // which is a ±π shift whenever the half angle leaves `(-π/2, π/2]`. + let half = crate::atan2(sin_half, cos_half); + let shift = if half >= 0.0 { + core::f32::consts::PI + } else { + -core::f32::consts::PI + }; + let wrapped_half = if crate::abs(half) > core::f32::consts::FRAC_PI_2 { + half - shift + } else { + half + }; + wrapped_half * 2.0 + } + + #[inline] + #[cfg(feature = "dim2")] + pub(super) fn recentered_angle(&self, _axis: usize, limit: &AngularLimitParams) -> f32 { + let d = crate::rotation_angle(self.ang_err) - limit.center; + // `atan2(sin, cos)` wraps the result to `(-π, π]`. + crate::atan2(crate::sin(d), crate::cos(d)) + } + #[inline] #[cfg(feature = "dim2")] pub(super) fn ang_err_axis(&self, _axis: usize) -> f32 { @@ -476,6 +553,7 @@ pub(super) fn limit_linear_generic( limits: [f32; 2], erp_inv_dt_val: f32, cfm_coeff: f32, + max_corr_velocity: f32, jacobians: &mut [f32], j_id_a: u32, j_id_b: u32, @@ -517,7 +595,11 @@ pub(super) fn limit_linear_generic( let dist = helper.lin_err.dot(lin_jac); let min_enabled = dist <= limits[0]; let max_enabled = limits[1] <= dist; - let rhs_bias = ((dist - limits[1]).max(0.0) - (limits[0] - dist).max(0.0)) * erp_inv_dt_val; + // Cap the bias so a deep violation recovers over a few steps instead of + // catapulting the bodies. + let rhs_bias = (((dist - limits[1]).max(0.0) - (limits[0] - dist).max(0.0)) * erp_inv_dt_val) + .max(-max_corr_velocity) + .min(max_corr_velocity); out.rhs_wo_bias = 0.0; out.rhs = rhs_bias; out.impulse_lo = if min_enabled { -MAX_FLT } else { 0.0 }; @@ -536,6 +618,7 @@ pub(super) fn limit_angular_generic( limits: [f32; 2], erp_inv_dt_val: f32, cfm_coeff: f32, + max_corr_velocity: f32, jacobians: &mut [f32], j_id_a: u32, j_id_b: u32, @@ -544,7 +627,10 @@ pub(super) fn limit_angular_generic( mprops: &[WorldMassProperties], colliders_start: usize, ) { - let ang_jac = helper.ang_jac_for_axis(limited_axis); + // The row measures the wrapped angle from the middle of the allowed range; + // its gradient is the plain joint axis, like the angular motor row. + let limit = AngularLimitParams::new(limits[0], limits[1]); + let ang_jac = helper.axis_ang_jac(limited_axis); lock_jacobians_generic( out, jacobians, @@ -565,12 +651,16 @@ pub(super) fn limit_angular_generic( out.writeback_axis = (DIM_USIZE + limited_axis) as u32; out.cfm_coeff = cfm_coeff; - let s_limits = [crate::sin(limits[0] * 0.5), crate::sin(limits[1] * 0.5)]; - let s_ang = helper.ang_err_axis(limited_axis); - let min_enabled = s_ang <= s_limits[0]; - let max_enabled = s_limits[1] <= s_ang; - let rhs_bias = - ((s_ang - s_limits[1]).max(0.0) - (s_limits[0] - s_ang).max(0.0)) * erp_inv_dt_val; + let ang_limits = limit.bounds(); + let ang = helper.recentered_angle(limited_axis, &limit); + let min_enabled = ang <= ang_limits[0]; + let max_enabled = ang_limits[1] <= ang; + // Cap the bias so a deep violation recovers over a few steps instead of + // catapulting the bodies. + let rhs_bias = (((ang - ang_limits[1]).max(0.0) - (ang_limits[0] - ang).max(0.0)) + * erp_inv_dt_val) + .max(-max_corr_velocity) + .min(max_corr_velocity); out.rhs_wo_bias = 0.0; out.rhs = rhs_bias; out.impulse_lo = if min_enabled { -MAX_FLT } else { 0.0 }; @@ -661,7 +751,7 @@ pub(super) fn motor_angular_generic( colliders_start: usize, ) { let mp = motor.motor_params(dt); - let ang_jac = helper.motor_ang_jac(motor_axis); + let ang_jac = helper.axis_ang_jac(motor_axis); lock_jacobians_generic( out, jacobians, @@ -688,10 +778,8 @@ pub(super) fn motor_angular_generic( #[cfg(feature = "dim2")] let s_ang_dist = crate::sin(crate::rotation_angle(helper.ang_err) * 0.5); let s_target_ang = crate::sin(mp.target_pos * 0.5); - // smallest_abs_diff_between_sin_angles — using the simpler form - // (dist - target) since the two-pi wrap concerns the rotation part - // and we operate on sin-half-angles already. - rhs_wo_bias += (s_ang_dist - s_target_ang) * mp.erp_inv_dt; + rhs_wo_bias += + smallest_abs_diff_between_sin_angles(s_ang_dist, s_target_ang) * mp.erp_inv_dt; } rhs_wo_bias += -mp.target_vel; diff --git a/src_rbd_shaders/dynamics/multibody/impulse_joint_constraints/kernels.rs b/src_rbd_shaders/dynamics/multibody/impulse_joint_constraints/kernels.rs index 25cd281..bbaeb0a 100644 --- a/src_rbd_shaders/dynamics/multibody/impulse_joint_constraints/kernels.rs +++ b/src_rbd_shaders/dynamics/multibody/impulse_joint_constraints/kernels.rs @@ -14,7 +14,7 @@ use crate::utils::BatchIndices; use crate::utils::linalg::VSlice; use super::super::lu::LANES; -use super::super::types::MultibodyInfo; +use super::super::types::{MultibodyInfo, MultibodyLinkStatic}; use super::jacobians::*; use super::types::*; @@ -55,6 +55,7 @@ pub fn gpu_mb_update_impulse_joint_constraints( // hardcoded `0.8/dt` + `cfm = 0`. let lock_erp_inv_dt = softness.joint_erp_inv_dt; let lock_cfm_coeff = softness.joint_cfm_coeff; + let max_corr_velocity = softness.max_corr_velocity; let joints_start = batch_ids.mb_imp_joints_start(batch_id); let cons_start = batch_ids.mb_imp_joint_constraints_start(batch_id); @@ -91,6 +92,7 @@ pub fn gpu_mb_update_impulse_joint_constraints( dt, lock_erp_inv_dt, lock_cfm_coeff, + max_corr_velocity, ); } i += num_threads; @@ -113,6 +115,8 @@ pub fn gpu_mb_finalize_impulse_joint_constraints( #[spirv(storage_buffer, descriptor_set = 1, binding = 0)] multibody_info: &[MultibodyInfo], #[spirv(storage_buffer, descriptor_set = 1, binding = 1)] mass_matrices: &[f32], #[spirv(storage_buffer, descriptor_set = 1, binding = 2)] lu_pivots: &[u32], + #[spirv(storage_buffer, descriptor_set = 1, binding = 3)] + links_static: &[MultibodyLinkStatic], #[spirv(uniform, descriptor_set = 0, binding = 3)] batch_ids: &BatchIndices, ) { let num_threads = num_workgroups.x * 64; @@ -150,6 +154,7 @@ pub fn gpu_mb_finalize_impulse_joint_constraints( &mb, mass_matrices, lu_pivots, + links_static, il, ); } @@ -162,6 +167,7 @@ pub fn gpu_mb_finalize_impulse_joint_constraints( &mb, mass_matrices, lu_pivots, + links_static, il, ); } @@ -174,7 +180,7 @@ pub fn gpu_mb_finalize_impulse_joint_constraints( } } -/// One PGS sweep over the multibody-touching impulse-joint axis constraints +/// One PGS iteration over the multibody-touching impulse-joint axis constraints /// of a single color — **one workgroup per joint**, the lanes cooperating on /// that joint's per-axis `J·v` reductions and `W·J` applies. /// diff --git a/src_rbd_shaders/dynamics/multibody/impulse_joint_constraints/mod.rs b/src_rbd_shaders/dynamics/multibody/impulse_joint_constraints/mod.rs index a6d1ecc..b92043f 100644 --- a/src_rbd_shaders/dynamics/multibody/impulse_joint_constraints/mod.rs +++ b/src_rbd_shaders/dynamics/multibody/impulse_joint_constraints/mod.rs @@ -8,7 +8,7 @@ //! //! 1. `gpu_mb_init_impulse_joint_constraints` — once per step, after FK / LU. //! 2. `gpu_mb_update_impulse_joint_constraints` — once per substep. -//! 3. `gpu_mb_solve_impulse_joint_constraints` — one PGS sweep, updates both +//! 3. `gpu_mb_solve_impulse_joint_constraints` — one PGS iteration, updates both //! sides' velocities. //! 4. `gpu_mb_remove_impulse_joint_constraint_bias` — strips the positional //! bias from `rhs` before the stabilization sweep. 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 8143248..b0632cb 100644 --- a/src_rbd_shaders/dynamics/multibody/impulse_joint_constraints/update.rs +++ b/src_rbd_shaders/dynamics/multibody/impulse_joint_constraints/update.rs @@ -6,11 +6,13 @@ 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::ISlice; use crate::utils::linalg::{MatSlice, lu_solve_in_place, VSlice}; use crate::{DIM, Pose}; -use super::super::types::MultibodyInfo; -use super::super::ws_soa::{WS_LTW, WsAddr, ws_pose}; +use super::super::types::{MultibodyInfo, MultibodyLinkStatic}; +use super::super::utils::zero_kinematic_dofs; +use super::super::ws_soa::{WS_LTW, WS_WORLD_COM, WsAddr, ws_pose, ws_vec}; use super::helper::*; use super::jacobians::*; use super::types::*; @@ -27,6 +29,7 @@ pub(super) fn solve_mb_wj( mb: &MultibodyInfo, mass_matrices: &[f32], lu_pivots: &[u32], + links_static: &[MultibodyLinkStatic], // Interleaved dynamics-buffer view (`stride = num_batches`, `shift = // batch_id`). il: VSlice, @@ -47,6 +50,15 @@ pub(super) fn solve_mb_wj( ); 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)); + + // Kinematic dofs are user-driven: the impulse must not move them. + let stat_slice = ISlice { + buf: links_static, + base: mb.first_link as usize, + stride: il.stride, + shift: il.shift, + }; + zero_kinematic_dofs(jacobians, wj_base, &stat_slice, mb.num_links); } impl MbImpulseJointBuilder { @@ -69,6 +81,7 @@ impl MbImpulseJointBuilder { dt: f32, lock_erp_inv_dt: f32, lock_cfm_coeff: f32, + max_corr_velocity: f32, ) { let cons_base = cons_start + self.constraint_id as usize; // Mark all axis-constraint slots inactive up-front; the active branches @@ -119,8 +132,22 @@ impl MbImpulseJointBuilder { let frame1 = pose_a * self.joint.local_frame_a; let frame2 = pose_b * self.joint.local_frame_b; - let world_com1 = pose_a.translation; - let world_com2 = pose_b.translation; + let world_com1 = side_world_com( + self.side_a_kind, + self.side_a_link, + &mb_a, + links_workspace, + il, + pose_a, + ); + let world_com2 = side_world_com( + self.side_b_kind, + self.side_b_link, + &mb_b, + links_workspace, + il, + pose_b, + ); let helper = new_helper( frame1, @@ -321,6 +348,7 @@ impl MbImpulseJointBuilder { [lim.min, lim.max], lock_erp_inv_dt, lock_cfm_coeff, + max_corr_velocity, jacobians, j_id_a, j_id_b, @@ -355,6 +383,7 @@ impl MbImpulseJointBuilder { [lim.min, lim.max], lock_erp_inv_dt, lock_cfm_coeff, + max_corr_velocity, jacobians, j_id_a, j_id_b, @@ -375,13 +404,32 @@ impl MbImpulseJointBuilder { } } +/// Center of mass the side's lever arm is measured against. A free-body side +/// is positioned by its COM-centered solver pose, so that pose's origin is it; +/// a multibody link is positioned by its body-origin frame, so its center of +/// mass comes from the workspace. +#[inline] +pub(super) fn side_world_com( + side_kind: u32, + side_link: u32, + mb: &MultibodyInfo, + links_workspace: &[Vec4], + il: VSlice, + side_pose: Pose, +) -> crate::Vector { + if side_kind == SIDE_KIND_MB { + let wa = WsAddr::new(mb.first_link as usize, il.stride, il.shift); + ws_vec(links_workspace, wa, side_link, WS_WORLD_COM) + } else { + side_pose.translation + } +} + /// Look up the world-space pose of a side. Free-body sides read from the /// shared `poses` buffer (COM-centered solver pose); multibody sides take -/// their link's `local_to_world` from the multibody workspace (which also -/// stores body-origin = COM-centered, since multibody links have a zeroed -/// `local_com`, as set up by the host pipeline). The `mb` argument is read -/// by value to keep SPIR-V happy and is only meaningful when `side_kind == -/// SIDE_KIND_MB`. +/// their link's body-origin `local_to_world` from the multibody workspace. +/// The `mb` argument is read by value to keep SPIR-V happy and is only +/// meaningful when `side_kind == SIDE_KIND_MB`. #[inline] pub(super) fn side_world_pose( side_kind: u32, diff --git a/src_rbd_shaders/dynamics/multibody/solve_constraints.rs b/src_rbd_shaders/dynamics/multibody/solve_constraints.rs index cdba815..cd6ff3e 100644 --- a/src_rbd_shaders/dynamics/multibody/solve_constraints.rs +++ b/src_rbd_shaders/dynamics/multibody/solve_constraints.rs @@ -1,4 +1,4 @@ -//! Fused multibody PGS sweep: joint limit/motor constraints followed by +//! Fused multibody PGS iteration: joint limit/motor constraints followed by //! contact constraints, in one dispatch per substep phase. use khal_std::glamx::UVec3; @@ -20,7 +20,21 @@ use super::types::{ const LANES: u32 = 64; -/// One PGS sweep over a multibody's joint (limit/motor) constraints followed +/// Caps a friction impulse to the circular cone of radius `limit`. In 3D both +/// tangent rows of a contact point are capped jointly; in 2D there is a single +/// row and this degenerates to a scalar clamp. +#[inline] +fn cap_friction(t0: f32, t1: f32, limit: f32) -> (f32, f32) { + let norm_sq = t0 * t0 + t1 * t1; + if norm_sq > limit * limit && norm_sq > 0.0 { + let scale = limit / crate::sqrt(norm_sq); + (t0 * scale, t1 * scale) + } else { + (t0, t1) + } +} + +/// One PGS iteration over a multibody's joint (limit/motor) constraints followed /// by its contact constraints. /// /// Dispatch: one 64-lane workgroup per (multibody, batch). @@ -45,6 +59,7 @@ pub fn gpu_mb_solve_constraints( #[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, + #[spirv(workgroup)] delta2_shared: &mut f32, ) { let batch_id = workgroup_id.y; let mb_idx = workgroup_id.x; @@ -138,14 +153,27 @@ pub fn gpu_mb_solve_constraints( } - // Contacts. + // Contacts. In 3D the two friction rows of a contact point are solved + // together so their impulse can be capped to the friction cone; the second + // row is handled by its sibling and skipped here. for s in 0..contact_count { let cons = contact_constraints.read(ccons_base + s as usize); + let is_tangent = cons.kind == MB_CONTACT_KIND_TANGENT; // Friction is only solved during the relaxation phase. - if use_bias && cons.kind == MB_CONTACT_KIND_TANGENT { + if use_bias && is_tangent { continue; } + #[cfg(feature = "dim3")] + if is_tangent && s != cons.normal_constraint_slot + 1 { + continue; + } + #[cfg(feature = "dim3")] + let has_pair = is_tangent; + #[cfg(feature = "dim2")] + let has_pair = false; + let col_offset = ccol_base + (s as usize) * dofs_stride; + let col_offset2 = col_offset + 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 @@ -157,10 +185,30 @@ pub fn gpu_mb_solve_constraints( }; workgroup_memory_barrier_with_group_sync(); + let mut j_dot_v0 = 0.0f32; if lane == 0 { - let mut j_dot_v = 0.0f32; for i in 0..ndofs { - j_dot_v += scratch[i as usize]; + j_dot_v0 += scratch[i as usize]; + } + } + workgroup_memory_barrier_with_group_sync(); + + if has_pair { + scratch[lane as usize] = if lane < ndofs { + contact_constraint_jacs.read(col_offset2 + lane as usize) * dof_v[lane as usize] + } else { + 0.0 + }; + } + workgroup_memory_barrier_with_group_sync(); + + if lane == 0 { + let cons2 = contact_constraints.read(ccons_base + (s + 1) as usize); + let mut j_dot_v1 = 0.0f32; + if has_pair { + for i in 0..ndofs { + j_dot_v1 += scratch[i as usize]; + } } // Free-body side stays lane-0-local. let free = if is_self { @@ -169,52 +217,76 @@ pub fn gpu_mb_solve_constraints( 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); + j_dot_v0 += cons.lin_jac.dot(free.linear) + gdot(cons.ang_jac, free.angular); + if has_pair { + j_dot_v1 += cons2.lin_jac.dot(free.linear) + gdot(cons2.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; let cfm_factor = if use_bias { cons.cfm_factor } else { 1.0 }; - let raw_imp = 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 { + let impulse0 = imp_shared[s as usize]; + let rhs0 = if use_bias { cons.rhs } else { cons.rhs_wo_bias }; + let raw0 = cfm_factor * (impulse0 - cons.inv_lhs * (j_dot_v0 + rhs0)); + + let impulse1 = if has_pair { + imp_shared[(s + 1) as usize] + } else { 0.0 + }; + let raw1 = if has_pair { + let rhs1 = if use_bias { cons2.rhs } else { cons2.rhs_wo_bias }; + cfm_factor * (impulse1 - cons2.inv_lhs * (j_dot_v1 + rhs1)) } else { - raw_imp + 0.0 }; - let delta = new_imp - impulse; - imp_shared[s as usize] = new_imp; - *delta_shared = delta; - if delta != 0.0 && !is_self { + // Normal: clamp to ≥ 0. Friction: cap the tangent pair to the + // circular cone `μ · normal_impulse`. + let (new0, new1) = if is_tangent { + let limit = cons.friction_coeff * imp_shared[cons.normal_constraint_slot as usize]; + cap_friction(raw0, raw1, limit) + } else if raw0 < 0.0 { + (0.0, 0.0) + } else { + (raw0, 0.0) + }; + + let delta0 = new0 - impulse0; + let delta1 = if has_pair { new1 - impulse1 } else { 0.0 }; + imp_shared[s as usize] = new0; + if has_pair { + imp_shared[(s + 1) as usize] = new1; + } + *delta_shared = delta0; + *delta2_shared = delta1; + + if !is_self && (delta0 != 0.0 || delta1 != 0.0) { let mut new_free = free; - new_free.linear += cons.lin_jac * (cons.free_body_im * delta); - new_free.angular += cons.ii_ang_jac * delta; + new_free.linear += cons.lin_jac * (cons.free_body_im * delta0); + new_free.angular += cons.ii_ang_jac * delta0; + if has_pair { + new_free.linear += cons2.lin_jac * (cons2.free_body_im * delta1); + new_free.angular += cons2.ii_ang_jac * delta1; + } solver_vels.write(colliders_start + cons.free_body_id as usize, new_free); } } workgroup_memory_barrier_with_group_sync(); // 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; + let delta0 = *delta_shared; + let delta1 = *delta2_shared; + if lane < ndofs { + if delta0 != 0.0 { + let col = contact_constraint_columns.read(col_offset + lane as usize); + dof_v[lane as usize] += delta0 * col; + } + if has_pair && delta1 != 0.0 { + let col = contact_constraint_columns.read(col_offset2 + lane as usize); + dof_v[lane as usize] += delta1 * col; + } } + workgroup_memory_barrier_with_group_sync(); } // Writeback @@ -231,7 +303,7 @@ pub fn gpu_mb_solve_constraints( } } -/// Joint-only PGS sweep (the joint half of [`gpu_mb_solve_constraints`]), +/// Joint-only PGS iteration (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] @@ -500,60 +572,98 @@ pub fn gpu_mb_solve_contacts_delassus( } workgroup_memory_barrier_with_group_sync(); + // In 3D the two friction rows of a contact point are solved together so + // their impulse can be capped to the friction cone; the second row is + // handled by its sibling and skipped here. 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 is_tangent = kind == MB_CONTACT_KIND_TANGENT; - if use_bias && kind == MB_CONTACT_KIND_TANGENT { + if use_bias && is_tangent { // Friction is only solved during the stabilization sweep. continue; } + #[cfg(feature = "dim3")] + if is_tangent && s != normal_slot + 1 { + continue; + } + #[cfg(feature = "dim3")] + let has_pair = is_tangent; + #[cfg(feature = "dim2")] + let has_pair = false; + + let impulse0 = imp_shared[s as usize]; + let raw0 = cfm_shared[s as usize] + * (impulse0 - inv_lhs_shared[s as usize] * (a_shared[s as usize] + rhs_shared[s as usize])); + let impulse1 = if has_pair { + imp_shared[(s + 1) as usize] + } else { + 0.0 + }; + let raw1 = if has_pair { + cfm_shared[(s + 1) as usize] + * (impulse1 + - inv_lhs_shared[(s + 1) as usize] + * (a_shared[(s + 1) as usize] + rhs_shared[(s + 1) as usize])) + } else { + 0.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 (new0, new1) = if is_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 + cap_friction(raw0, raw1, limit) + } else if raw0 < 0.0 { + (0.0, 0.0) } else { - raw_imp + (raw0, 0.0) }; - let delta = new_imp - impulse; + let delta0 = new0 - impulse0; + let delta1 = if has_pair { new1 - impulse1 } else { 0.0 }; - if delta != 0.0 { + if delta0 != 0.0 || delta1 != 0.0 { if lane == 0 { - imp_shared[s as usize] = new_imp; + imp_shared[s as usize] = new0; + if has_pair { + imp_shared[(s + 1) as usize] = new1; + } 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; + free.linear += cons.lin_jac * (cons.free_body_im * delta0); + free.angular += cons.ii_ang_jac * delta0; + if has_pair { + let cons2 = contact_constraints.read(cons_base + (s + 1) as usize); + free.linear += cons2.lin_jac * (cons2.free_body_im * delta1); + free.angular += cons2.ii_ang_jac * delta1; + } 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; + let d_row2 = d_base + ((s + 1) * MAXC) as usize; for j in StepRng::new(lane..count, LANES) { - a_shared[j as usize] += delta * delassus.read(d_row + j as usize); + let mut acc = delta0 * delassus.read(d_row + j as usize); + if has_pair { + acc += delta1 * delassus.read(d_row2 + j as usize); + } + a_shared[j as usize] += acc; } 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; + dof_v[lane as usize] += delta0 * col; + if has_pair { + let col2 = contact_constraint_columns + .read(col_base + ((s + 1) as usize) * dofs_stride + lane as usize); + dof_v[lane as usize] += delta1 * col2; + } } workgroup_memory_barrier_with_group_sync(); } diff --git a/src_rbd_shaders/dynamics/multibody/types.rs b/src_rbd_shaders/dynamics/multibody/types.rs index b5c9974..2253530 100644 --- a/src_rbd_shaders/dynamics/multibody/types.rs +++ b/src_rbd_shaders/dynamics/multibody/types.rs @@ -16,11 +16,10 @@ use crate::dynamics::joint::{GenericJoint, SPATIAL_DIM}; /// Equivalent to `SPATIAL_DIM`. pub const MAX_JOINT_DOFS: usize = SPATIAL_DIM; -/// Maximum number of simultaneously-active multibody contact **points** per -/// multibody. Sized for typical use (a single multibody touching the -/// environment with up to ~32 contact points × 2 manifold sides). Per-multibody -/// banks of this size are pre-allocated; surplus slots are left inactive. -pub const MAX_MB_CONTACTS_PER_MB: u32 = 64; +/// Maximum number of simultaneously-active multibody contact points per +/// multibody. +/// TODO: make this configurable/auto-resizeable +pub const MAX_MB_CONTACTS_PER_MB: u32 = 128; /// Number of constraint slots reserved per contact point — one normal + /// `DIM-1` friction tangents (Coulomb friction). Mirrors rapier's @@ -132,6 +131,13 @@ pub struct MultibodyLinkWorkspace { /// Per-link kinematic acceleration (rapier's `workspace.accs[i]`). /// Populated by the Coriolis variant of `apply_gravity`. pub kinematic_acc: Velocity, + /// User-applied force on this link, in world space. + pub external_force: Vec3, + /// Per-link multiplier on the global gravity. + pub gravity_scale: f32, + /// User-applied torque on this link, in world space. + pub external_torque: Vec3, + pub _pad3: u32, } /// Per-link workspace updated every step. @@ -163,6 +169,12 @@ pub struct MultibodyLinkWorkspace { pub rb_vels: Velocity, /// Per-link kinematic acceleration (rapier's `workspace.accs[i]`). pub kinematic_acc: Velocity, + /// User-applied force on this link, in world space. + pub external_force: Vec2, + /// User-applied torque on this link. + pub external_torque: f32, + /// Per-link multiplier on the global gravity. + pub gravity_scale: f32, } /// One unit (1-DOF) constraint generated from a multibody joint's limit or @@ -264,7 +276,8 @@ pub struct MultibodyContactConstraint { /// `cons[normal_constraint_slot].impulse` to compute their clamp limit /// `±μ · normal_impulse`. For normal slots this is just self. pub normal_constraint_slot: u32, - pub _pad0: u32, + /// Second touched link for a self-contact, `u32::MAX` otherwise. + pub link_id_b: u32, /// Free-body linear jacobian: `+jac_dir` on body B's side or /// `-jac_dir` on body A's side, depending on which side of the contact @@ -277,7 +290,7 @@ pub struct MultibodyContactConstraint { pub _pad2: u32, /// Same as `ang_jac` but pre-multiplied by the free body's /// `effective_world_inv_inertia`. Used to update `solver_vels.angular` - /// without re-multiplying every PGS sweep. + /// without re-multiplying every PGS iteration. pub ii_ang_jac: Vec3, pub _pad3: u32, @@ -292,11 +305,34 @@ pub struct MultibodyContactConstraint { pub impulse: f32, /// Contact CFM factor `1/(1+cfm_coeff)` — rapier's generic-contact form: - /// multiplies the impulse each PGS sweep for compliance (replaces the old + /// multiplies the impulse each PGS iteration for compliance (replaces the old /// rigid `cfm_coeff`/`cfm_gain` generic-joint form, which was always 0). pub cfm_factor: f32, - pub _unused_cfm: f32, - pub _pad4: [u32; 2], + /// Approaching normal velocity captured at the start of the step, scaled by + /// the restitution coefficient. Zero on non-bouncy points. + pub restitution_seed: f32, + /// Combined restitution coefficient of the two colliders. + pub restitution: f32, + pub _pad4: u32, + + /// Torque arm of the `link_id` side about that link's center of mass, + /// crossed with the multibody-side direction (`-lin_jac`). + pub torque_a: Vec3, + pub _pad5: u32, + /// Same for the `link_id_b` side of a self-contact, crossed with `lin_jac`. + pub torque_b: Vec3, + pub _pad6: u32, + + /// Contact anchor frozen in the `link_id` side's body frame. + pub local_p1: Vec3, + /// Separation at the substep the anchors were frozen; the live separation + /// is this plus the drift of the two anchors along the contact normal. + pub base_dist: f32, + /// Contact anchor frozen in the other side's frame: the second link's body + /// frame for a self-contact, the free body's solver (center-of-mass) frame + /// otherwise. + pub local_p2: Vec3, + pub _pad7: u32, } /// 2D variant of [`MultibodyContactConstraint`] — angular jacobian collapses @@ -311,6 +347,22 @@ pub struct MultibodyContactConstraint { pub kind: u32, pub free_body_id: u32, + /// Slot index (relative to `cons_base`) of the associated normal + /// constraint. Tangents read `cons[normal_constraint_slot].impulse` to + /// compute their clamp limit `±μ · normal_impulse`. + pub normal_constraint_slot: u32, + /// Second touched link for a self-contact, `u32::MAX` otherwise. + pub link_id_b: u32, + /// Free-body linear jacobian. + pub lin_jac: Vec2, + + /// Contact anchor frozen in the `link_id` side's body frame. + pub local_p1: Vec2, + /// Contact anchor frozen in the other side's frame: the second link's body + /// frame for a self-contact, the free body's solver (center-of-mass) frame + /// otherwise. + pub local_p2: Vec2, + pub free_body_im: f32, /// Free-body angular jacobian (`r_free × jac_dir`) — scalar in 2D. pub ang_jac: f32, @@ -319,14 +371,6 @@ pub struct MultibodyContactConstraint { /// Coulomb friction coefficient `μ`. pub friction_coeff: f32, - /// Slot index (relative to `cons_base`) of the associated normal - /// constraint. Tangents read `cons[normal_constraint_slot].impulse` to - /// compute their clamp limit `±μ · normal_impulse`. - pub normal_constraint_slot: u32, - pub _pad0: [u32; 1], - /// Free-body linear jacobian. - pub lin_jac: Vec2, - pub inv_lhs: f32, pub rhs: f32, pub rhs_wo_bias: f32, @@ -334,7 +378,23 @@ pub struct MultibodyContactConstraint { /// Contact CFM factor `1/(1+cfm_coeff)` (rapier's generic-contact form). pub cfm_factor: f32, - pub _unused_cfm: f32, + /// Approaching normal velocity captured at the start of the step, scaled by + /// the restitution coefficient. Zero on non-bouncy points. + pub restitution_seed: f32, + /// Combined restitution coefficient of the two colliders. + pub restitution: f32, + /// Separation at the substep the anchors were frozen; the live separation + /// is this plus the drift of the two anchors along the contact normal. + pub base_dist: f32, + + /// Torque arm of the `link_id` side about that link's center of mass, + /// crossed with the multibody-side direction (`-lin_jac`). + pub torque_a: f32, + /// Same for the `link_id_b` side of a self-contact, crossed with `lin_jac`. + pub torque_b: f32, + /// Pads the struct to a multiple of 16 bytes, which std430 requires of the + /// array stride given the vector members. + pub _pad1: [f32; 2], } /// Descriptor for one multibody: where its links live, how many DOFs it has, and diff --git a/src_rbd_shaders/dynamics/multibody/utils.rs b/src_rbd_shaders/dynamics/multibody/utils.rs index fa92418..ec19a49 100644 --- a/src_rbd_shaders/dynamics/multibody/utils.rs +++ b/src_rbd_shaders/dynamics/multibody/utils.rs @@ -1,12 +1,33 @@ //! Small math / coordinate helpers shared across multibody kernels. use crate::dynamics::joint::{ANG_AXES_MASK, LIN_AXES_MASK}; +use crate::utils::ISlice; use crate::{DIM, Pose, Rotation, Vector}; use khal_std::index::MaybeIndexUnchecked; use parry::math::VectorExt; use super::types::{MAX_JOINT_DOFS, MultibodyLinkStatic}; +/// Zeroes the entries of a dense dof-space vector (`dst[base .. base + ndofs]`) +/// that belong to kinematic joints. Those velocities are user-driven, so a +/// constraint's `M⁻¹Jᵀ` response must never touch them. +#[inline] +pub fn zero_kinematic_dofs( + dst: &mut [f32], + base: usize, + stat_slice: &ISlice, + num_links: u32, +) { + for k in 0..num_links { + let stat = stat_slice[k as usize]; + if stat.kinematic != 0 { + for d in 0..stat.ndofs { + dst.write(base + (stat.assembly_id + d) as usize, 0.0); + } + } + } +} + /// Number of free DOFs implied by a `locked_axes` bitmask. #[inline] pub fn count_free_dofs(locked: u32) -> u32 { diff --git a/src_rbd_shaders/dynamics/multibody/ws_soa.rs b/src_rbd_shaders/dynamics/multibody/ws_soa.rs index 1449ed4..796470c 100644 --- a/src_rbd_shaders/dynamics/multibody/ws_soa.rs +++ b/src_rbd_shaders/dynamics/multibody/ws_soa.rs @@ -35,12 +35,18 @@ mod layout { pub const WS_RB_VELS: u32 = 11; /// Kinematic acceleration: lin | ang. pub const WS_KIN_ACC: u32 = 13; + /// World-space center of mass of the link: xyz, pad. + pub const WS_WORLD_COM: u32 = 15; + /// User-applied force: xyz | gravity scale in `w`. + pub const WS_EXT_FORCE: u32 = 16; + /// User-applied torque: xyz, pad. + pub const WS_EXT_TORQUE: u32 = 17; /// Total quads per link (per-link stride, in quad units). - pub const WS_QUADS: u32 = 15; + pub const WS_QUADS: u32 = 18; } /* - * Per-link QUAD offsets of each field (dim2): 9 quads / 144 B per link. + * Per-link QUAD offsets of each field (dim2): 11 quads / 176 B per link. */ #[cfg(feature = "dim2")] mod layout { @@ -62,8 +68,14 @@ mod layout { pub const WS_RB_VELS: u32 = 7; /// Kinematic acceleration: (lin.x, lin.y, ang, pad). pub const WS_KIN_ACC: u32 = 8; + /// World-space center of mass of the link: x, y, pad, pad. + pub const WS_WORLD_COM: u32 = 9; + /// User-applied wrench: force x, y | torque | gravity scale. + pub const WS_EXT_FORCE: u32 = 10; + /// Alias of [`WS_EXT_FORCE`]: in 2D the whole wrench fits one quad. + pub const WS_EXT_TORQUE: u32 = 10; /// Total quads per link (per-link stride, in quad units). - pub const WS_QUADS: u32 = 9; + pub const WS_QUADS: u32 = 11; } pub use layout::*; @@ -159,6 +171,65 @@ 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)); } +/// User-applied external wrench and per-link gravity scale, packed into +/// [`WS_EXT_FORCE`] (and [`WS_EXT_TORQUE`] in 3D). Written by the host, read +/// by the generalized-force assembly; nothing in the step loop clears them. +#[cfg(feature = "dim3")] +#[inline] +pub fn ws_ext_wrench(buf: &[Vec4], a: WsAddr, k: u32) -> (Vector, crate::AngVector, f32) { + let f = buf.read(a.at(k, WS_EXT_FORCE)); + let t = buf.read(a.at(k, WS_EXT_TORQUE)); + ( + Vec3::new(f.x, f.y, f.z), + Vec3::new(t.x, t.y, t.z), + f.w, + ) +} + +#[cfg(feature = "dim2")] +#[inline] +pub fn ws_ext_wrench(buf: &[Vec4], a: WsAddr, k: u32) -> (Vector, crate::AngVector, f32) { + let q = buf.read(a.at(k, WS_EXT_FORCE)); + (Vec2::new(q.x, q.y), q.z, q.w) +} + +/// Writes the pair of quads [`ws_ext_wrench`] reads back. +#[cfg(feature = "dim3")] +#[inline] +pub fn ws_set_ext_wrench( + buf: &mut [Vec4], + a: WsAddr, + k: u32, + force: Vector, + torque: crate::AngVector, + gravity_scale: f32, +) { + buf.write( + a.at(k, WS_EXT_FORCE), + Vec4::new(force.x, force.y, force.z, gravity_scale), + ); + buf.write( + a.at(k, WS_EXT_TORQUE), + Vec4::new(torque.x, torque.y, torque.z, 0.0), + ); +} + +#[cfg(feature = "dim2")] +#[inline] +pub fn ws_set_ext_wrench( + buf: &mut [Vec4], + a: WsAddr, + k: u32, + force: Vector, + torque: crate::AngVector, + gravity_scale: f32, +) { + buf.write( + a.at(k, WS_EXT_FORCE), + Vec4::new(force.x, force.y, torque, gravity_scale), + ); +} + /// Pose accessors. 3D: rotation quad + translation quad. 2D: one quad /// `(rot.re, rot.im, trans.x, trans.y)`. #[cfg(feature = "dim3")] @@ -367,6 +438,14 @@ pub fn ws_soa_from_structs( 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); + ws_set_ext_wrench( + &mut out, + a, + k, + ws.external_force, + ws.external_torque, + ws.gravity_scale, + ); } } out diff --git a/src_rbd_shaders/dynamics/sim_params.rs b/src_rbd_shaders/dynamics/sim_params.rs index e5c7762..545c6df 100644 --- a/src_rbd_shaders/dynamics/sim_params.rs +++ b/src_rbd_shaders/dynamics/sim_params.rs @@ -40,8 +40,9 @@ pub struct ConstraintSoftness { pub static_erp_inv_dt: f32, /// [`Self::cfm_factor`] for contacts touching a fixed body. pub static_cfm_factor: f32, - /// Unused; keeps the uniform a 16-byte multiple. - pub _padding0: f32, + /// Coefficient in `[0, 1]` applied to a contact impulse before it is + /// re-used as the next substep's (or next frame's) initial guess. + pub warmstart_coefficient: f32, /// Unused; keeps the uniform a 16-byte multiple. pub _padding1: f32, } @@ -62,7 +63,7 @@ impl ConstraintSoftness { dt: params.dt, static_erp_inv_dt: params.static_contact_erp_inv_dt(), static_cfm_factor: params.static_contact_cfm_factor(), - _padding0: 0.0, + warmstart_coefficient: params.warmstart_coefficient, _padding1: 0.0, } } diff --git a/src_rbd_shaders/dynamics/solver.rs b/src_rbd_shaders/dynamics/solver.rs index 3b2ddf3..0862dd5 100644 --- a/src_rbd_shaders/dynamics/solver.rs +++ b/src_rbd_shaders/dynamics/solver.rs @@ -276,6 +276,7 @@ pub fn gpu_init_solver_vels_inc( #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] mprops: &[WorldMassProperties], #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] all_params: &[RbdSimParams], #[spirv(uniform, descriptor_set = 0, binding = 3)] batch_ids: &BatchIndices, + #[spirv(uniform, descriptor_set = 0, binding = 4)] gravity: &glamx::Vec4, ) { let batch_id = invocation_id.y; let params = all_params.at(batch_id as usize); @@ -292,10 +293,12 @@ pub fn gpu_init_solver_vels_inc( // TODO: this isn't a very pretty way of detecting static bodies. if mprops[idx].inv_mass != Vector::ZERO { - // TODO: this currently only handles gravity. - // TODO: make the gravity configurable - let gravity = Vector::Y * -9.81; - solver_vels_inc[idx].linear = gravity * params.dt; + // TODO: this currently only handles gravity (no user forces yet). + #[cfg(feature = "dim3")] + let g = Vector::new(gravity.x, gravity.y, gravity.z); + #[cfg(feature = "dim2")] + let g = Vector::new(gravity.x, gravity.y); + solver_vels_inc[idx].linear = g * params.dt; } } } diff --git a/src_rbd_shaders/utils/indices.rs b/src_rbd_shaders/utils/indices.rs index bfd81a6..806a48d 100644 --- a/src_rbd_shaders/utils/indices.rs +++ b/src_rbd_shaders/utils/indices.rs @@ -81,8 +81,10 @@ pub struct BatchIndices { /// Offset (in f32 entries, within a batch's `mass_matrices` view) of the /// section holding the coriolis-aware "acceleration" mass matrix /// (rapier's `acc_augmented_mass`). - /// Non-zero = split-matrix. - /// Zero = single melded matrix for both (not recommended). + /// Non-zero = implicit coriolis on: split matrices, the acc section drives + /// the acceleration solve while the plain matrix drives constraints. + /// Zero = implicit coriolis off: a single plain (coriolis-free) matrix + /// serves both; coriolis/gyroscopic forces are applied explicitly only. pub mass_matrix_acc_section_offset: u32, /// Per-batch stride (capacity) of the multibody DoF-coupling buffer. pub mb_dof_couplings_batch_capacity: u32, From 5fb1a1e008b90afbe4697c2f40e33776f613b0ac Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Fri, 7 Aug 2026 10:17:12 +0200 Subject: [PATCH 3/6] fix: match rapier's manifold reduction and its degenerate-selection guards --- src_rbd/pipeline/rbd_state.rs | 6 ++ src_rbd_shaders/queries/polygonal_feature.rs | 94 +++++++++++--------- 2 files changed, 59 insertions(+), 41 deletions(-) diff --git a/src_rbd/pipeline/rbd_state.rs b/src_rbd/pipeline/rbd_state.rs index 3b5db97..5f69ce7 100644 --- a/src_rbd/pipeline/rbd_state.rs +++ b/src_rbd/pipeline/rbd_state.rs @@ -401,6 +401,12 @@ impl RbdState { } /// The contact manifold buffer (post narrow-phase + body resolution). + /// Per-batch number of manifolds currently in [`Self::contacts`]. + pub fn contacts_len(&self) -> &Tensor { + &self.contacts_len + } + + /// GPU buffer holding the contact manifolds. pub fn contacts(&self) -> &Tensor { &self.contacts } diff --git a/src_rbd_shaders/queries/polygonal_feature.rs b/src_rbd_shaders/queries/polygonal_feature.rs index 29c8ddd..68bcd11 100644 --- a/src_rbd_shaders/queries/polygonal_feature.rs +++ b/src_rbd_shaders/queries/polygonal_feature.rs @@ -579,63 +579,79 @@ mod dim3 { result } + /// Reduces the candidate set to at most `MAX_MANIFOLD_POINTS` solver + /// contacts. Mirrors `reduce_manifold_naive`: pick the deepest point, then + /// the one furthest from it, then the two extremes along the tangent of + /// that segment, considering only points within `prediction`. pub fn manifold_reduction( candidates: &[ContactPoint; MAX_CANDIDATE_POINTS], num_candidates: u32, normal: Vector, + prediction: f32, ) -> ContactManifold { let mut result = ContactManifold::default(); let num = num_candidates as usize; if num <= MAX_MANIFOLD_POINTS { - result.points_a.write(0, candidates.read(0)); - result.points_a.write(1, candidates.read(1)); - result.points_a.write(2, candidates.read(2)); - result.points_a.write(3, candidates.read(3)); + for i in 0..num { + result.points_a.write(i, candidates.read(i)); + } result.len = num_candidates; return result; } - // Run contact reduction so we only have up to four solver contacts. - // 1. Find the deepest contact. - let mut deepest_dist = candidates.at(0).dist; - let mut selected = [ - 0usize, - MAX_CANDIDATE_POINTS, - MAX_CANDIDATE_POINTS, - MAX_CANDIDATE_POINTS, - ]; + const NONE: usize = MAX_CANDIDATE_POINTS; - for i in 1..num { + // 1. Find the deepest contact. + let mut selected = [NONE, NONE, NONE, NONE]; + let mut deepest_dist = MAX_FLT; + for i in 0..num { if candidates.at(i).dist < deepest_dist { deepest_dist = candidates.at(i).dist; selected.write(0, i); } } + if selected.read(0) == NONE { + return result; + } + // 2. Find the point that is the furthest from the deepest one. let selected_a = candidates.at(selected.read(0)).pt; let mut furthest_dist = -MAX_FLT; - for i in 0..num { - let pt_sel = selected_a - candidates.at(i).pt; - let dist = pt_sel.dot(pt_sel); - if i != selected.read(0) && dist > furthest_dist { + let d = candidates.at(i).pt - selected_a; + let dist = d.dot(d); + if i != selected.read(0) + && candidates.at(i).dist <= prediction + && dist > furthest_dist + { furthest_dist = dist; selected.write(1, i); } } + result.points_a.write(0, candidates.read(selected.read(0))); + result.len = 1; + if selected.read(1) == NONE { + return result; + } + // 3. Now find the two points furthest from the segment we built so far. + // A zero-length segment has no tangent, so it stays a single contact. let selected_b = candidates.at(selected.read(1)).pt; - let selected_ab = selected_b - selected_a; - let tangent = selected_ab.cross(normal); + if selected_a == selected_b { + return result; + } + let tangent = (selected_b - selected_a).cross(normal); let mut min_dot = MAX_FLT; let mut max_dot = -MAX_FLT; - for i in 0..num { - if i == selected.read(0) || i == selected.read(1) { + if i == selected.read(0) + || i == selected.read(1) + || candidates.at(i).dist > prediction + { continue; } @@ -644,32 +660,28 @@ mod dim3 { min_dot = d; selected.write(2, i); } - if d > max_dot { max_dot = d; selected.write(3, i); } } - if selected.read(2) == MAX_CANDIDATE_POINTS { - selected.write(2, selected.read(3)); - selected.write(3, MAX_CANDIDATE_POINTS); - } - - result.points_a.write(0, candidates.read(selected.read(0))); result.points_a.write(1, candidates.read(selected.read(1))); result.len = 2; + if selected.read(2) == NONE { + return result; + } - if selected.read(2) != MAX_CANDIDATE_POINTS { - result.points_a.write(2, candidates.read(selected.read(2))); - result.len = 3; - - if selected.read(3) != MAX_CANDIDATE_POINTS { - result.points_a.write(3, candidates.read(selected.read(3))); - result.len = 4; - } + result.points_a.write(2, candidates.read(selected.read(2))); + result.len = 3; + // The min and max extremes come from one pass, so a single remaining + // candidate is picked for both; keeping it once leaves three contacts. + if selected.read(2) == selected.read(3) { + return result; } + result.points_a.write(3, candidates.read(selected.read(3))); + result.len = 4; result } @@ -783,7 +795,7 @@ mod dim3 { } if !any_point_is_outside { - return manifold_reduction(&candidates, num_candidates, sep_axis1); + return manifold_reduction(&candidates, num_candidates, sep_axis1, prediction); } } } @@ -841,7 +853,7 @@ mod dim3 { } if !any_point_is_outside { - return manifold_reduction(&candidates, num_candidates, sep_axis1); + return manifold_reduction(&candidates, num_candidates, sep_axis1, prediction); } } } @@ -877,13 +889,13 @@ mod dim3 { } if num_candidates as usize == MAX_CANDIDATE_POINTS { - return manifold_reduction(&candidates, num_candidates, sep_axis1); + return manifold_reduction(&candidates, num_candidates, sep_axis1, prediction); } } } } - manifold_reduction(&candidates, num_candidates, sep_axis1) + manifold_reduction(&candidates, num_candidates, sep_axis1, prediction) } } From be15c9f0d91b9dca43d99ea5cea8722f2faadebf Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Fri, 7 Aug 2026 10:35:33 +0200 Subject: [PATCH 4/6] feat: expose multibody motor retargeting and solver introspection --- src/state.rs | 14 ++ src_rbd/dynamics/joint.rs | 5 +- src_rbd/dynamics/mod.rs | 2 +- src_rbd/dynamics/multibody/multibody_set.rs | 172 ++++++++++++++++++++ 4 files changed, 191 insertions(+), 2 deletions(-) diff --git a/src/state.rs b/src/state.rs index 20b788b..e23f83b 100644 --- a/src/state.rs +++ b/src/state.rs @@ -234,6 +234,20 @@ impl NexusState { &mut self.rbd_envs[env] } + /// Mutable access to environment `env`'s rapier world that does **not** mark + /// the rbd state dirty, for use after [`Self::finalize`]. + /// + /// Nothing written here reaches the GPU on its own: the rapier sets are the + /// build-time source the GPU buffers were baked from, and marking them dirty + /// would rebuild those buffers and snap the simulation back to the authored + /// state. Use this to run rapier-side helpers whose output you then push + /// through a runtime setter — e.g. driving an MJCF actuator model and + /// forwarding the resulting motors with + /// [`GpuMultibodySet::set_motors`](crate::rbd::dynamics::GpuMultibodySet::set_motors). + pub fn rbd_world_mut_untracked(&mut self, env: usize) -> &mut PhysicsWorld { + &mut self.rbd_envs[env] + } + pub fn insert_rigid_body(&mut self, body: RigidBody, collider: Collider) -> RigidBodyHandle { self.insert_rigid_body_in(0, body, collider) } diff --git a/src_rbd/dynamics/joint.rs b/src_rbd/dynamics/joint.rs index 8e35a58..4f743c1 100644 --- a/src_rbd/dynamics/joint.rs +++ b/src_rbd/dynamics/joint.rs @@ -33,7 +33,10 @@ fn convert_joint_limits(limits: RapierJointLimits) -> JointLimits { } } -fn convert_joint_motor(motor: RapierJointMotor) -> JointMotor { +/// Converts a rapier joint motor into the GPU representation. Public so callers +/// driving actuators at runtime can push a freshly configured motor through +/// [`GpuMultibodySet::set_motor`](crate::dynamics::GpuMultibodySet::set_motor). +pub fn convert_joint_motor(motor: RapierJointMotor) -> JointMotor { JointMotor { target_vel: motor.target_vel, target_pos: motor.target_pos, diff --git a/src_rbd/dynamics/mod.rs b/src_rbd/dynamics/mod.rs index 023e512..e67b91e 100644 --- a/src_rbd/dynamics/mod.rs +++ b/src_rbd/dynamics/mod.rs @@ -2,7 +2,7 @@ pub use crate::shaders::dynamics::RbdSimParams; pub use coloring::{ColorBucketsArgs, ColoringArgs, GpuColoring}; -pub use joint::{GpuImpulseJointSet, GpuJointSolver, JointSolverArgs}; +pub use joint::{GpuImpulseJointSet, GpuJointSolver, JointSolverArgs, convert_joint_motor}; pub use mprops_update::{GpuMpropsUpdate, GpuSyncColliderPosesShader}; #[cfg(feature = "dim3")] pub use multibody::{GpuMultibodySet, GpuMultibodySolver, MultibodySolverArgs}; diff --git a/src_rbd/dynamics/multibody/multibody_set.rs b/src_rbd/dynamics/multibody/multibody_set.rs index ab2989b..223805c 100644 --- a/src_rbd/dynamics/multibody/multibody_set.rs +++ b/src_rbd/dynamics/multibody/multibody_set.rs @@ -224,6 +224,12 @@ impl GpuMultibodySet { } /// GPU buffer for generalized coordinates. + /// Per-batch stride of the DoF buffers (the length of each section of + /// [`Self::dof_state`]). + pub fn dofs_per_batch(&self) -> u32 { + self.dofs_per_batch + } + pub fn dof_values(&self) -> &Tensor { &self.dof_values } @@ -294,6 +300,99 @@ impl GpuMultibodySet { .unwrap(); } + /// Overwrites one joint motor of a multibody link and uploads the changed + /// link to the GPU, enabling the axis so the solver drives it. + /// + /// `link_id` is the global link id within the batch (it matches the body + /// index given to [`from_rapier`](Self::from_rapier)) and `axis` indexes the + /// 6-DoF spatial layout (`0..DIM` linear, `DIM..` angular). Motors are baked + /// into the GPU state at finalization, so per-step actuation has to come + /// through here. + pub fn set_motor( + &mut self, + backend: &GpuBackend, + batch: u32, + link_id: u32, + axis: usize, + motor: JointMotor, + ) -> Result<(), GpuBackendError> { + if axis >= 6 { + return Ok(()); + } + let global_idx = (link_id * self.num_batches + batch) as usize; + let entry = match self.links_static_mirror.get_mut(global_idx) { + Some(e) => e, + None => return Ok(()), + }; + // `impulse` is solver state, not configuration: keep the accumulated + // value so retargeting a servo does not drop its warmstart. + let impulse = entry.data.motors[axis].impulse; + entry.data.motors[axis] = motor; + entry.data.motors[axis].impulse = impulse; + entry.data.motor_axes |= 1u32 << axis; + let snapshot = *entry; + backend.write_buffer( + self.links_static.buffer_mut(), + global_idx as u64, + std::slice::from_ref(&snapshot), + ) + } + + /// The motor currently configured on `axis` of multibody link `link_id`, + /// as last uploaded. Use it to adjust one field of a live motor without + /// rebuilding the rest. + pub fn motor(&self, batch: u32, link_id: u32, axis: usize) -> Option { + if axis >= 6 { + return None; + } + let global_idx = (link_id * self.num_batches + batch) as usize; + self.links_static_mirror + .get(global_idx) + .map(|e| e.data.motors[axis]) + } + + /// Batched [`Self::set_motor`]: applies every `(link_id, axis, motor)` and + /// uploads each touched link once, rather than once per axis. + /// + /// This is the entry point for per-step actuation (a position-servo robot + /// re-targets several axes of many links every frame), so it is worth + /// keeping the upload count down to one per link. + pub fn set_motors( + &mut self, + backend: &GpuBackend, + batch: u32, + updates: &[(u32, usize, JointMotor)], + ) -> Result<(), GpuBackendError> { + let mut touched: Vec = Vec::with_capacity(updates.len()); + for &(link_id, axis, motor) in updates { + if axis >= 6 { + continue; + } + let global_idx = (link_id * self.num_batches + batch) as usize; + let Some(entry) = self.links_static_mirror.get_mut(global_idx) else { + continue; + }; + // `impulse` is solver state, not configuration: keep the accumulated + // value so retargeting a servo does not drop its warmstart. + let impulse = entry.data.motors[axis].impulse; + entry.data.motors[axis] = motor; + entry.data.motors[axis].impulse = impulse; + entry.data.motor_axes |= 1u32 << axis; + touched.push(global_idx); + } + touched.sort_unstable(); + touched.dedup(); + for global_idx in touched { + let snapshot = self.links_static_mirror[global_idx]; + backend.write_buffer( + self.links_static.buffer_mut(), + global_idx as u64, + std::slice::from_ref(&snapshot), + )?; + } + Ok(()) + } + /// Sets a motor's target velocity on a multibody joint and uploads the /// updated link to the GPU. `link_id` is the global link id within the /// batch (matches the body / collider index that was given to @@ -419,6 +518,79 @@ impl GpuMultibodySet { Ok(()) } + /// Per-multibody descriptors (contact counts, dof offsets, ...). + pub fn multibody_info(&self) -> &Tensor { + &self.multibody_info + } + + /// The per-multibody contact-constraint slabs. + pub fn contact_constraints(&self) -> &Tensor { + &self.contact_constraints + } + + /// Coefficient applied to a contact impulse before it is re-used as the + /// next substep's or next frame's initial guess. Zero disables warmstarting. + pub fn warmstart_coefficient(&self) -> f32 { + self.warmstart_coefficient + } + + /// Per-batch stride of [`Self::contact_constraints`]. + pub fn contact_constraints_per_batch(&self) -> u32 { + self.contact_constraints_per_batch + } + + /// Per-constraint `Jᵀ` rows of the contact constraints (`ndofs` floats each, + /// laid out like [`Self::contact_constraints`]). + pub fn contact_constraint_jacs(&self) -> &Tensor { + &self.contact_constraint_jacs + } + + /// Per-constraint `M⁻¹·Jᵀ` columns of the contact constraints, laid out + /// like [`Self::contact_constraint_jacs`]. + pub fn contact_constraint_columns(&self) -> &Tensor { + &self.contact_constraint_columns + } + + /// Per-link `SPATIAL_DIM × ndofs` column-major body jacobians, indexed from + /// each multibody's [`MultibodyInfo::jacobian_offset`]. + pub fn body_jacobians(&self) -> &Tensor { + &self.body_jacobians + } + + /// Per-multibody `ndofs × ndofs` mass matrices, indexed from each + /// multibody's [`MultibodyInfo::mass_matrix_offset`]. Doubles as the LU work + /// buffer, so after a step this holds the factorization, not `M` itself. + pub fn mass_matrices(&self) -> &Tensor { + &self.mass_matrices + } + + /// Reads back the generalized coordinate of every DoF of batch `batch_id`, + /// in assembly order (the same order as [`Self::dof_state`]'s velocity + /// section). The coordinates live in the link workspace, so this unpacks the + /// SoA layout for callers. + pub async fn read_dof_coords( + &self, + backend: &GpuBackend, + batch_id: u32, + ) -> Result, khal::backend::GpuBackendError> { + use crate::shaders::dynamics::{WsAddr, ws_coord}; + + let ws: Vec = backend.slow_read_vec(self.links_workspace.buffer()).await?; + let a = WsAddr::new(0, self.num_batches, batch_id); + let mut out = Vec::new(); + for k in 0..self.links_per_batch { + let stat = &self.links_static_mirror + [(batch_id * self.links_per_batch + k) as usize]; + let locked = stat.data.locked_axes; + for axis in 0..6u32 { + if locked & (1 << axis) == 0 { + out.push(ws_coord(&ws, a, k, axis)); + } + } + } + Ok(out) + } + /// Upload a new integration timestep. pub fn set_dt(&mut self, backend: &GpuBackend, dt: f32) { self.dt = Tensor::scalar( From 61a8c9159646225b8d47160680a7f38d1eb1f7cd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Fri, 7 Aug 2026 10:57:24 +0200 Subject: [PATCH 5/6] chore: drive the mjcf menagerie demo with the reference controls --- crates/examples3d/mujoco_menagerie3.rs | 736 ++++++++++++++++++------- 1 file changed, 529 insertions(+), 207 deletions(-) diff --git a/crates/examples3d/mujoco_menagerie3.rs b/crates/examples3d/mujoco_menagerie3.rs index 5da2377..a2ead81 100644 --- a/crates/examples3d/mujoco_menagerie3.rs +++ b/crates/examples3d/mujoco_menagerie3.rs @@ -3,7 +3,9 @@ use kiss3d::egui; use nexus_viewer3d::{NexusViewer, RenderMaterial}; use nexus3d::prelude::{NexusPipeline, NexusState}; use rapier3d::prelude::*; -use rapier3d_mjcf::{MjcfLoaderOptions, MjcfMultibodyOptions, MjcfRobot}; +use nexus3d::rbd::dynamics::convert_joint_motor; +use nexus3d::rbd::shaders::dynamics::JointMotor; +use rapier3d_mjcf::{MjcfLoaderOptions, MjcfMultibodyOptions, MjcfRobot, MjcfRobotHandles}; use std::fs; use std::path::{Path, PathBuf}; @@ -118,123 +120,304 @@ struct VisualMeshReg { texture: Option, material: Option, } +/// Panel state, mirroring the Example Settings of rapier's `mujoco_menagerie3`. +#[derive(Clone, Copy, PartialEq)] +struct Settings { + use_multibody: bool, + render_colliders: bool, + render_visual_meshes: bool, + render_visual_primitives: bool, + disable_collisions: bool, + enable_controls: bool, + enable_springs: bool, + actuator_strength: f32, + /// Index into the keyframe picker: 0 is "(none)", `i + 1` is keyframe `i`. + keyframe: usize, +} + +impl Default for Settings { + fn default() -> Self { + Self { + use_multibody: true, + render_colliders: false, + render_visual_meshes: true, + render_visual_primitives: false, + disable_collisions: true, + enable_controls: true, + enable_springs: true, + actuator_strength: 1.0, + keyframe: 0, + } + } +} + +impl Settings { + /// With the actuators driving the model, switching keyframe retargets the + /// servos live; without them nothing tracks the target, so the pose has to + /// be applied by reloading instead. + fn keyframe_is_live(&self) -> bool { + self.use_multibody && self.enable_controls + } -/// Loads a single MuJoCo Menagerie MJCF model into a fresh [`NexusState`], -/// registers its render shapes and a floor with `viewer`, frames the camera on -/// it, and finalizes the state ready for simulation. + /// Whether moving from `self` to `next` requires rebuilding the scene. + /// Actuator strength is read live every step, and so is the keyframe while + /// the servos are driving. + fn needs_reload(&self, next: &Self) -> bool { + self.use_multibody != next.use_multibody + || self.render_colliders != next.render_colliders + || self.render_visual_meshes != next.render_visual_meshes + || self.render_visual_primitives != next.render_visual_primitives + || self.disable_collisions != next.disable_collisions + || self.enable_controls != next.enable_controls + || self.enable_springs != next.enable_springs + || (self.keyframe != next.keyframe && !next.keyframe_is_live()) + } +} + +/// Everything the per-step actuator drive needs. Multibody path only. +struct Controls { + handles: MjcfRobotHandles>, + /// One control vector per keyframe, precomputed at load. + per_keyframe_ctrl: Vec>, + /// Control vector for "(none)": hold the neutral pose. + neutral: Vec, +} + +/// A loaded model plus the picker state that depends on it. +struct Loaded { + state: NexusState, + controls: Option, + /// "(none)" followed by one entry per declared keyframe. + keyframe_names: Vec, +} + +/// Merge the keyframes from a sibling `keyframes.xml` (next to the scene file) +/// into `robot`, skipping any whose name is already present. /// -/// The model is kept in its native Z-up frame (no rotation): the viewer is -/// configured Z-up by the caller and the rigid-body gravity is set to -Z below, -/// so MJCF data is consumed as-authored. -async fn load_scene( - viewer: &mut NexusViewer, - scene: &Path, - render_colliders: bool, -) -> anyhow::Result { - let mut state = NexusState::default(); +/// Menagerie models often keep their keyframes in a standalone file meant to be +/// ``d, which the scene itself does not reference; without this they +/// would never reach the picker. +fn merge_sibling_keyframes(robot: &mut MjcfRobot, scene_path: &Path) { + let Some(kf_path) = scene_path.parent().map(|d| d.join("keyframes.xml")) else { + return; + }; + if !kf_path.exists() { + return; + } + match MjcfRobot::from_file(&kf_path, loader_options()) { + Ok((kf_robot, _)) => { + let existing: std::collections::HashSet = robot + .keyframes + .iter() + .filter_map(|k| k.name.clone()) + .collect(); + for k in kf_robot.keyframes { + if k.name.as_ref().is_none_or(|n| !existing.contains(n)) { + robot.keyframes.push(k); + } + } + } + Err(e) => eprintln!( + "Failed to load sibling keyframes `{}`: {e}.", + kf_path.display() + ), + } +} - // `` collision shapes get density 0 — the physical mass comes from the - // model's `` tags. Roots stay dynamic so free-based robots fall - // and land on the floor; set `make_roots_fixed: true` to anchor them. - let options = MjcfLoaderOptions { +/// Picker entries for a model's keyframes, prefixed with "(none)". +fn keyframe_names(robot: &MjcfRobot) -> Vec { + let mut names = vec!["(none)".to_string()]; + for (i, k) in robot.keyframes.iter().enumerate() { + names.push(k.name.clone().unwrap_or_else(|| format!("key {i}"))); + } + names +} + +/// The keyframe a freshly picked model starts on: `home` if it declares one, +/// else its first, else "(none)". +fn default_keyframe(names: &[String]) -> usize { + names + .iter() + .position(|n| n == "home") + .unwrap_or(if names.len() > 1 { 1 } else { 0 }) +} + +/// The MJCF loader options shared by the pre-flight DoF check and the real load. +fn loader_options() -> MjcfLoaderOptions { + MjcfLoaderOptions { skip_plane_geoms: true, make_roots_fixed: false, // Surface visual-only geoms as `MjcfBody::visual_meshes` (forwarded to // the viewer below) instead of turning them into colliders. create_colliders_from_visual_shapes: false, + // Density 0: the physical mass comes from the model's `` tags. collider_blueprint: ColliderBuilder::default().density(0.0), - // No `shift`: the model stays in its native MJCF Z-up frame. The viewer - // is set Z-up and gravity points -Z, so nothing needs rotating. ..MjcfLoaderOptions::default() - }; + } +} + +/// Loads a single MuJoCo Menagerie MJCF model into a fresh [`NexusState`] under +/// `settings`, registers its render shapes and a floor with `viewer`, frames the +/// camera on it, and finalizes the state ready for simulation. +/// +/// The model is kept in its native Z-up frame (no rotation): the viewer is +/// configured Z-up by the caller and gravity is set to -Z below, so MJCF data is +/// consumed as-authored. +async fn load_scene( + viewer: &mut NexusViewer, + scene: &Path, + settings: &Settings, +) -> anyhow::Result { + let mut state = NexusState::default(); - // Collected during loading, registered once the world borrow ends. Both the - // visual meshes (textured/PBR) and every collision collider are gathered so - // the render mode can be chosen at registration time. Each collider entry is - // tagged with whether its body has visual meshes, so the visual-mesh mode can - // still fall back to colliders for links that have none. + // Collected during loading, registered once the world borrow ends. let mut visual_meshes: Vec = Vec::new(); let mut collider_shapes: Vec<(RigidBodyHandle, SharedShape, Pose, bool)> = Vec::new(); - // Fixed cuboid floor, sized from the loaded model's bounding box. let mut floor: Option<(Vec3, Vec3)> = None; - // Camera framing for the loaded model: `(eye, target)`. let mut camera: Option<(Vec3, Vec3)> = None; + let mut controls = None; + let mut names = vec!["(none)".to_string()]; + let mut gravity = -9.81; println!("Loading MJCF scene `{}`.", scene.display()); - match MjcfRobot::from_file(scene, options) { - Ok((robot, _model)) => { - // Every `` collider is kept as-is. Collider-less links need no - // placeholder: the GPU pipeline now gives every body its own slot. - let world = state.rbd_world_mut(0); - // `insert_using_multibody_joints` consumes the robot, so clone - // it and keep the original around for its visual meshes. - let handles = robot.clone().insert_using_multibody_joints( - &mut world.bodies, - &mut world.colliders, - &mut world.multibody_joints, - &mut world.impulse_joints, - MjcfMultibodyOptions::DISABLE_SELF_CONTACTS, - ); + match MjcfRobot::from_file(scene, loader_options()) { + Ok((mut robot, model)) => { + merge_sibling_keyframes(&mut robot, scene); + names = keyframe_names(&robot); + let keyframe = settings + .keyframe + .checked_sub(1) + .and_then(|i| robot.keyframes.get(i)) + .cloned(); + + // MJCF gives gravity as a 3-vector, normally (0, 0, -9.81) since the + // format is Z-up. Keep only the magnitude and lock it to -Z so + // physics and rendering stay aligned whatever the model declares. + let g = model.option.gravity; + let mag = ((g[0] * g[0] + g[1] * g[1] + g[2] * g[2]) as f32).sqrt(); + gravity = -mag; + + if settings.disable_collisions { + for link in &mut robot.bodies { + for collider in &mut link.colliders { + collider.set_collision_groups(InteractionGroups::new( + Group::GROUP_1, + Group::GROUP_2, + Default::default(), + )); + } + } + } + + let mut mb_options = if settings.disable_collisions { + MjcfMultibodyOptions::DISABLE_SELF_CONTACTS + } else { + MjcfMultibodyOptions::default() + }; + // `` passive springs are integrated implicitly by + // default; unchecking strips them (e.g. cassie's leg springs). + if !settings.enable_springs { + mb_options |= MjcfMultibodyOptions::SKIP_JOINT_SPRINGS; + } - // Activate the model's actuators so position-servo robots hold their - // pose instead of folding under gravity (e.g. anymal_c's - // `` joint servos). A zero control vector targets the - // neutral/rest pose — the same as the rapier testbed's "Enable joint - // controls" (which applies `ctrl = 0` every frame). The actuator motor - // config (target + stiffness) is static for a constant `ctrl`, so we - // configure it once here, before `finalize` bakes the multibody into - // the GPU state. Without this, ``-actuated models collapse. - let ctrl = vec![0.0; handles.actuators.len()]; - handles.apply_controls_multibody(&mut world.bodies, &mut world.multibody_joints, &ctrl); - - // Forward each body's visual meshes to the viewer; for bodies - // without visual meshes, render their collision shapes instead. - for (i, body_handle) in handles.bodies.iter().enumerate() { - let Some(body_handle) = body_handle else { - continue; - }; + let world = state.rbd_world_mut(0); + // `insert_using_*` consumes the robot, so clone it and keep the + // original around for its visual meshes and keyframes. + let body_handles: Vec> = if settings.use_multibody { + let handles = robot.clone().insert_using_multibody_joints( + &mut world.bodies, + &mut world.colliders, + &mut world.multibody_joints, + &mut world.impulse_joints, + mb_options, + ); + if let Some(key) = &keyframe { + handles.apply_keyframe( + &mut world.bodies, + &mut world.multibody_joints, + &robot, + key, + ); + } + let bodies = handles + .bodies + .iter() + .map(|b| b.as_ref().map(|h| h.body)) + .collect(); + if settings.enable_controls { + let per_keyframe_ctrl = robot + .keyframes + .iter() + .map(|k| robot.keyframe_controls(k)) + .collect(); + let neutral = vec![0.0; handles.actuators.len()]; + controls = Some(Controls { + handles, + per_keyframe_ctrl, + neutral, + }); + } + bodies + } else { + let handles = robot.clone().insert_using_impulse_joints( + &mut world.bodies, + &mut world.colliders, + &mut world.impulse_joints, + ); + if let Some(key) = &keyframe { + handles.apply_keyframe(&mut world.bodies, &robot, key); + } + handles.bodies.iter().map(|b| b.as_ref().map(|h| h.body)).collect() + }; + + // Gather each body's render geometry. Visual meshes carry the + // authored color / texture / UVs; colliders are the fallback for + // links that declare none. + for (i, body) in body_handles.iter().enumerate() { + let Some(body) = *body else { continue }; let mjcf_body = &robot.bodies[i]; - let has_visual = !mjcf_body.visual_meshes.is_empty(); - // Collect every collision collider at its body-local pose (a body - // can own several now), tagged with whether the body has visuals. - for collider in &body_handle.colliders { - let c = &world.colliders[collider.handle]; + let visuals: Vec<_> = mjcf_body + .visual_meshes + .iter() + // "Render visual primitives" keeps the capsules and boxes some + // models declare in their visual channel; by default only the + // .obj-derived meshes are drawn. + .filter(|vm| settings.render_visual_primitives || vm.shape.as_trimesh().is_some()) + .collect(); + let has_visual = !visuals.is_empty(); + for (handle, _) in world + .colliders + .iter() + .filter(|(_, c)| c.parent() == Some(body)) + .map(|(h, c)| (h, c)) + .collect::>() + { + let c = &world.colliders[handle]; let local_pose = c.position_wrt_parent().copied().unwrap_or(Pose::IDENTITY); - collider_shapes.push(( - body_handle.body, - c.shared_shape().clone(), - local_pose, - has_visual, - )); + collider_shapes.push((body, c.shared_shape().clone(), local_pose, has_visual)); } - if has_visual { - for vm in &mjcf_body.visual_meshes { - // Resolve the base color: geom/material rgba if set, - // white behind a texture (so it shows in native colors), - // else a neutral grey. Mirrors rapier's testbed. - let textured = vm.texture.is_some(); - let color = vm.rgba.unwrap_or(if textured { - [1.0, 1.0, 1.0, 1.0] - } else { - [0.7, 0.7, 0.75, 1.0] - }); - let material = vm.material.map(|m| RenderMaterial { + for vm in visuals { + let textured = vm.texture.is_some(); + let color = vm.rgba.unwrap_or(if textured { + [1.0, 1.0, 1.0, 1.0] + } else { + [0.7, 0.7, 0.75, 1.0] + }); + visual_meshes.push(VisualMeshReg { + body, + shape: vm.shape.clone(), + local_pose: vm.local_pose, + color, + uvs: vm.uvs.clone(), + normals: vm.normals.clone(), + texture: vm.texture.clone(), + material: vm.material.map(|m| RenderMaterial { metallic: m.metallic, roughness: m.roughness, reflectance: m.reflectance, emissive: m.emissive, - }); - visual_meshes.push(VisualMeshReg { - body: body_handle.body, - shape: vm.shape.clone(), - local_pose: vm.local_pose, - color, - uvs: vm.uvs.clone(), - normals: vm.normals.clone(), - texture: vm.texture.clone(), - material, - }); - } + }), + }); } } @@ -248,28 +431,19 @@ async fn load_scene( let center = aabb.center(); let he = aabb.half_extents(); let footprint = he.x.max(he.y).max(0.5); - - // A wide, thin floor just below the model (Z is up, so it's thin - // on Z and sits at the model's lowest Z). let floor_thick = 0.1; - let floor_he = Vec3::new(footprint * 6.0, footprint * 6.0, floor_thick); - let floor_center = Vec3::new(center.x, center.y, center.z - he.z - floor_thick); - floor = Some((floor_center, floor_he)); - - // Frame the model from a 3/4 view (Z up, so the elevation is +Z). + floor = Some(( + Vec3::new(center.x, center.y, center.z - he.z - floor_thick), + Vec3::new(footprint * 6.0, footprint * 6.0, floor_thick), + )); let radius = (he.x * he.x + he.y * he.y + he.z * he.z).sqrt().max(0.5); let target = Vec3::new(center.x, center.y, center.z); - let eye = target + Vec3::new(radius * 2.2, -radius * 2.2, radius * 1.6); - camera = Some((eye, target)); + camera = Some((target + Vec3::new(radius * 2.2, -radius * 2.2, radius * 1.6), target)); } } - Err(e) => { - eprintln!("Failed to load MJCF scene `{}`: {e}.", scene.display()); - } + Err(e) => eprintln!("Failed to load MJCF scene `{}`: {e}.", scene.display()), } - // Floor (inserted through `NexusState` so it participates in the GPU sim and - // gets a render shape registered). if let Some((center, he)) = floor { let body = RigidBodyBuilder::fixed().translation(center).build(); let collider = ColliderBuilder::cuboid(he.x, he.y, he.z).build(); @@ -278,16 +452,11 @@ async fn load_scene( viewer.insert_shape(handle, &shape, Pose::IDENTITY); } - if render_colliders { - // Collider view: every collision shape (instanced, colored by shape - // type), at its body-local pose. No visual meshes. + if settings.render_colliders || !settings.render_visual_meshes { for (body, shape, local_pose, _) in &collider_shapes { viewer.insert_visual_shape(0, *body, shape, *local_pose); } } else { - // Visual-mesh view: the authored color/texture/UVs/normals/PBR meshes — - // rendered the way MuJoCo's own viewer shows them — plus colliders only - // for links that have no visual mesh (so nothing is invisible). for vm in &visual_meshes { viewer.insert_visual_mesh( 0, @@ -301,6 +470,7 @@ async fn load_scene( vm.material, ); } + // Links with no visual mesh would otherwise be invisible. for (body, shape, local_pose, has_visual) in &collider_shapes { if !has_visual { viewer.insert_visual_shape(0, *body, shape, *local_pose); @@ -311,35 +481,98 @@ async fn load_scene( if let Some((eye, target)) = camera { viewer.set_camera(eye, target); } - viewer .scene3d_mut() .add_directional_light(glamx::Vec3::new(-1.0, 1.0, -1.0)); + // The impulse-joint path needs a much finer step to stay stable; the + // multibody path instead raises the PGS iterations per substep. Mirrors the + // reference example. + let mut sim_params = nexus3d::rbd::shaders::dynamics::RbdSimParams::default(); + if !settings.use_multibody { + sim_params.dt = 1.0 / 240.0; + sim_params.num_solver_iterations = 12; + } + state.set_rbd_sim_params(0, sim_params); + state.finalize(viewer.backend()).await?; - // 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. + state.set_rbd_gravity(viewer.backend(), [0.0, 0.0, gravity]); if let Some(rbd) = state.rbd.as_mut() { - rbd.multibodies_mut().set_implicit_coriolis(false); + if settings.use_multibody { + rbd.multibodies_mut().set_num_internal_pgs_iterations(4); + } + // MuJoCo-style explicit coriolis: a single plain mass matrix, with + // coriolis / gyroscopic forces applied explicitly on the rhs. + rbd.set_implicit_coriolis(viewer.backend(), false); + } + Ok(Loaded { + state, + controls, + keyframe_names: names, + }) +} + +/// Drives the model's actuators toward `ctrl`, scaled by `gain`. +/// +/// The motor configuration is baked into the GPU state at finalization, so this +/// runs the MJCF actuator model on the CPU-side joints and then pushes each +/// touched motor across. +fn apply_controls( + state: &mut NexusState, + backend: &khal::backend::GpuBackend, + controls: &Controls, + ctrl: &[Real], + gain: Real, +) { + let mut updates: Vec<(u32, usize, JointMotor)> = Vec::new(); + { + // Untracked: the rapier sets are only the scratch the MJCF actuator + // model writes into. Marking them dirty would rebuild the GPU buffers + // from the authored poses and reset the model every step. + let world = state.rbd_world_mut_untracked(0); + controls.handles.apply_controls_multibody_scaled( + &mut world.bodies, + &mut world.multibody_joints, + ctrl, + gain, + ); + for ah in &controls.handles.actuators { + let Some(Some(handle)) = ah.joint else { continue }; + let Some((mb, link_id)) = world.multibody_joints.get(handle) else { + continue; + }; + let Some(link) = mb.links().nth(link_id) else { continue }; + // The GPU link id is the body index (see `GpuMultibodySet::set_motor`). + let body_idx = link.rigid_body_handle().into_raw_parts().0; + let axes = link.joint().data.motor_axes.bits(); + for axis in 0..6 { + if axes & (1 << axis) != 0 { + updates.push(( + body_idx, + axis, + convert_joint_motor(link.joint().data.motors[axis]), + )); + } + } + } + } + if let Some(rbd) = state.rbd.as_mut() { + let _ = rbd.multibodies_mut().set_motors(backend, 0, &updates); } - Ok(state) } /// Picks a scene: first runs the cheap DoF pre-check (no mesh I/O); if the model /// is within the GPU solver's DoF cap it tears down the current scene and loads -/// it, returning `Ok(state)`. If it exceeds the cap, nothing is loaded and an -/// `Err(message)` is returned for display in the picker. +/// it. If it exceeds the cap, nothing is loaded and an `Err(message)` is +/// returned for display in the picker. async fn select_scene( viewer: &mut NexusViewer, scene: &Path, - render_colliders: bool, -) -> anyhow::Result> { + settings: &Settings, +) -> anyhow::Result> { if let Some(dofs) = scene_max_dofs(scene) && dofs > MAX_MB_DOFS + && settings.use_multibody { return Ok(Err(format!( "{} needs {dofs} DoFs (max {MAX_MB_DOFS}) — not supported by the GPU solver.", @@ -347,14 +580,13 @@ async fn select_scene( ))); } viewer.clear_scene(); - let state = load_scene(viewer, scene, render_colliders).await?; - Ok(Ok(state)) + Ok(Ok(load_scene(viewer, scene, settings).await?)) } /// Loads MuJoCo Menagerie MJCF models and simulates them on the GPU rigid-body -/// pipeline, with a floating egui window to switch between the discovered models -/// at runtime (mirroring the scene picker in rapier's `mujoco_menagerie3` -/// example). +/// pipeline, with a floating egui window carrying the same controls as rapier's +/// `mujoco_menagerie3` example: model picker, render modes, collision / spring / +/// actuator toggles, a keyframe picker and a live actuator-strength slider. /// /// Scenes are discovered under `MUJOCO_MENAGERIE_DIR` (default: /// `../mujoco_menagerie` next to the workspace). The initial model is the one @@ -379,8 +611,6 @@ pub async fn run( println!("Discovered {} MuJoCo Menagerie scene(s).", scenes.len()); } - // Pick the initial scene (substring match), defaulting to unitree_a1, and - // falling back to the first discovered scene otherwise. let wanted = std::env::var("MUJOCO_MENAGERIE_SCENE").unwrap_or_else(|_| "unitree_a1".into()); let mut selected = scenes .iter() @@ -388,76 +618,133 @@ pub async fn run( .unwrap_or(0); // MJCF models are Z-up: orient the viewer's camera accordingly so the model - // stands upright without rotating its data. Done once; preserved across the - // per-model `set_camera` calls in `load_scene`. + // stands upright without rotating its data. viewer.set_up_axis(Vec3::Z); let mut timestamps = GpuTimestamps::new(viewer.backend(), 2048); - - // Red message shown in the picker when the highlighted model can't be loaded - // (currently: too many DoFs for the GPU solver). let mut error: Option = None; - // Render mode: false = textured visual meshes (default), true = the collision - // shapes. Toggled via the picker checkbox; a change reloads the scene. - let mut render_colliders = false; + let mut settings = Settings::default(); + let mut controls = None; + let mut keyframe_names = vec!["(none)".to_string()]; let mut state = match scenes.get(selected) { - Some(scene) => match select_scene(viewer, scene, render_colliders).await? { - Ok(state) => state, - Err(msg) => { - eprintln!("{msg}"); - error = Some(msg); - let mut state = NexusState::default(); - state.finalize(viewer.backend()).await?; - state + Some(scene) => { + // A freshly picked model starts on its default keyframe, which is + // only known once it is loaded: probe the names, then load for real. + match select_scene(viewer, scene, &settings).await? { + Ok(loaded) => { + settings.keyframe = default_keyframe(&loaded.keyframe_names); + keyframe_names = loaded.keyframe_names; + if settings.keyframe != 0 { + let reloaded = select_scene(viewer, scene, &settings).await?; + match reloaded { + Ok(l) => { + controls = l.controls; + l.state + } + Err(msg) => { + error = Some(msg); + let mut s = NexusState::default(); + s.finalize(viewer.backend()).await?; + s + } + } + } else { + controls = loaded.controls; + loaded.state + } + } + Err(msg) => { + eprintln!("{msg}"); + error = Some(msg); + let mut s = NexusState::default(); + s.finalize(viewer.backend()).await?; + s + } } - }, + } None => { - let mut state = NexusState::default(); - state.finalize(viewer.backend()).await?; - state + let mut s = NexusState::default(); + s.finalize(viewer.backend()).await?; + s } }; - // Model selection requested through the picker this frame, applied after the - // UI pass so we don't rebuild the scene mid-borrow. - let mut pending: Option = None; - // Render-mode change requested through the picker this frame. - let mut pending_mode: Option = None; + // Requested through the picker this frame, applied after the UI pass so we + // never rebuild the scene mid-borrow. + let mut pending_scene: Option = None; + let mut pending_settings: Option = None; while viewer.render_frame().await { - // Floating model-picker window (in addition to the viewer's main panel). - if !labels.is_empty() { + { let current = selected; let labels = &labels; - let pending = &mut pending; + let names = &keyframe_names; + let now = settings; + let pending_scene = &mut pending_scene; + let pending_settings = &mut pending_settings; let error = error.as_deref(); let count = labels.len(); - let render_colliders_now = render_colliders; - let pending_mode = &mut pending_mode; viewer.draw_custom_ui(move |ctx| { egui::Window::new("MuJoCo Menagerie") .default_pos([24.0, 220.0]) .resizable(true) .show(ctx, |ui| { - // Render-mode toggle: visual meshes (default) vs colliders. - let mut rc = render_colliders_now; - if ui.checkbox(&mut rc, "Render colliders").changed() { - *pending_mode = Some(rc); - } - ui.separator(); - // Previous / next buttons cycle through the models, - // wrapping around at either end. + let mut next = now; + ui.checkbox(&mut next.use_multibody, "Use multibody joints"); + ui.checkbox(&mut next.render_colliders, "Render colliders"); + ui.checkbox(&mut next.render_visual_meshes, "Render visual meshes"); + ui.checkbox( + &mut next.render_visual_primitives, + "Render visual primitives", + ); + ui.checkbox(&mut next.disable_collisions, "Disable collisions"); + ui.checkbox(&mut next.enable_controls, "Enable joint controls"); + ui.checkbox(&mut next.enable_springs, "Enable joint springs"); + ui.add( + egui::Slider::new(&mut next.actuator_strength, 0.02..=2.0) + .text("Actuator strength"), + ); ui.horizontal(|ui| { + ui.label("Keyframe"); + // Prev / next step through the model's keyframes, + // wrapping at either end, like the scene picker. + let n = names.len().max(1); if ui.button("<").clicked() { - *pending = Some((current + count - 1) % count); + next.keyframe = (next.keyframe + n - 1) % n; } if ui.button(">").clicked() { - *pending = Some((current + 1) % count); + next.keyframe = (next.keyframe + 1) % n; } - ui.label(format!("{}/{}", current + 1, count)); + egui::ComboBox::from_id_salt("keyframe") + .selected_text( + names + .get(next.keyframe) + .cloned() + .unwrap_or_else(|| "(none)".into()), + ) + .show_ui(ui, |ui| { + for (i, name) in names.iter().enumerate() { + ui.selectable_value(&mut next.keyframe, i, name); + } + }); }); - // Red error for an unsupported (e.g. too-many-DoF) model. + if next != now { + *pending_settings = Some(next); + } + + ui.separator(); + if count > 0 { + ui.horizontal(|ui| { + if ui.button("<").clicked() { + *pending_scene = Some((current + count - 1) % count); + } + if ui.button(">").clicked() { + *pending_scene = Some((current + 1) % count); + } + ui.label(format!("{}/{}", current + 1, count)); + }); + } if let Some(msg) = error { ui.colored_label(egui::Color32::RED, msg); } @@ -467,7 +754,7 @@ pub async fn run( .show(ui, |ui| { for (i, label) in labels.iter().enumerate() { if ui.selectable_label(current == i, label).clicked() { - *pending = Some(i); + *pending_scene = Some(i); } } }); @@ -475,18 +762,58 @@ pub async fn run( }); } - // Apply a model selection: the highlight moves immediately (so prev/next - // can step past an unsupported model), but the scene is only rebuilt when - // the model is within the GPU solver's DoF cap — otherwise the current - // scene stays and the picker shows a red error. - if let Some(i) = pending.take() - && i != selected - { - selected = i; - match select_scene(viewer, &scenes[selected], render_colliders).await? { - Ok(new_state) => { - state = new_state; + // A settings change either retargets the running sim (actuator strength, + // and the keyframe while the servos are driving) or rebuilds the scene. + let mut reload = false; + if let Some(next) = pending_settings.take() { + reload = settings.needs_reload(&next); + settings = next; + } + // A model change always rebuilds, and resets the keyframe to the new + // model's default. + let scene_changed = if let Some(i) = pending_scene.take() { + if i != selected { + selected = i; + settings.keyframe = 0; + reload = true; + true + } else { + false + } + } else { + false + }; + + if reload && let Some(scene) = scenes.get(selected) { + match select_scene(viewer, scene, &settings).await? { + Ok(loaded) => { + keyframe_names = loaded.keyframe_names; error = None; + if scene_changed { + // Only known now that the model is loaded; re-load once + // so it actually starts in that pose. + let def = default_keyframe(&keyframe_names); + if def != settings.keyframe { + settings.keyframe = def; + match select_scene(viewer, scene, &settings).await? { + Ok(l) => { + keyframe_names = l.keyframe_names; + controls = l.controls; + state = l.state; + } + Err(msg) => { + eprintln!("{msg}"); + error = Some(msg); + } + } + } else { + controls = loaded.controls; + state = loaded.state; + } + } else { + controls = loaded.controls; + state = loaded.state; + } } Err(msg) => { eprintln!("{msg}"); @@ -495,26 +822,21 @@ pub async fn run( } } - // Apply a render-mode toggle: reload the current model with the new mode. - if let Some(new_mode) = pending_mode.take() - && new_mode != render_colliders - { - render_colliders = new_mode; - if let Some(scene) = scenes.get(selected) { - match select_scene(viewer, scene, render_colliders).await? { - Ok(new_state) => { - state = new_state; - error = None; - } - Err(msg) => { - eprintln!("{msg}"); - error = Some(msg); - } - } - } - } - if viewer.simulating() { + if let Some(controls) = controls.as_ref() { + let ctrl = settings + .keyframe + .checked_sub(1) + .and_then(|i| controls.per_keyframe_ctrl.get(i)) + .unwrap_or(&controls.neutral); + apply_controls( + &mut state, + viewer.backend(), + controls, + ctrl, + settings.actuator_strength, + ); + } pipeline .simulate(viewer.backend(), &mut state, Some(&mut timestamps)) .await?; From 7370cd5823d016adc59e28c6d5e922a04449c0fe Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Fri, 14 Aug 2026 18:40:25 +0200 Subject: [PATCH 6/6] chore: fix the CI checks --- crates/examples3d/mujoco_menagerie3.rs | 26 ++- src/state.rs | 2 +- src_rbd/broad_phase/lbvh.rs | 18 +- src_rbd/dynamics/joint.rs | 8 +- .../multibody/multibody_from_rapier.rs | 15 +- src_rbd/dynamics/multibody/multibody_set.rs | 10 +- .../dynamics/multibody/multibody_solver.rs | 37 ++-- src_rbd/dynamics/solver.rs | 13 +- src_rbd/pipeline/insertion_removal.rs | 7 +- src_rbd/pipeline/rbd_state.rs | 6 +- src_rbd/pipeline/rbd_state_from_rapier.rs | 7 +- src_rbd/pipeline/rbd_step.rs | 194 +++++++++--------- src_rbd_shaders/broad_phase/brute_force.rs | 5 +- src_rbd_shaders/broad_phase/lbvh.rs | 25 ++- src_rbd_shaders/broad_phase/narrow_phase.rs | 5 +- src_rbd_shaders/dynamics/color_buckets.rs | 6 +- src_rbd_shaders/dynamics/mod.rs | 2 +- .../multibody/compute_dynamics_pre.rs | 102 +++++---- .../dynamics/multibody/contact_constraints.rs | 27 ++- .../dynamics/multibody/gravity_and_lu.rs | 77 ++++--- .../impulse_joint_constraints/helper.rs | 8 +- .../impulse_joint_constraints/jacobians.rs | 12 +- .../impulse_joint_constraints/kernels.rs | 5 +- .../impulse_joint_constraints/update.rs | 11 +- .../dynamics/multibody/integrate.rs | 23 ++- .../dynamics/multibody/joint_constraints.rs | 17 +- .../dynamics/multibody/mass_matrix.rs | 38 ---- src_rbd_shaders/dynamics/multibody/mod.rs | 1 - .../dynamics/multibody/solve_constraints.rs | 44 ++-- src_rbd_shaders/dynamics/multibody/ws_soa.rs | 11 +- src_rbd_shaders/dynamics/solver.rs | 11 +- src_rbd_shaders/dynamics/solver_utils.rs | 4 +- src_rbd_shaders/dynamics/warmstart.rs | 6 +- src_rbd_shaders/queries/polygonal_feature.rs | 15 +- src_rbd_shaders/tests/linalg.rs | 68 ++++-- src_rbd_shaders/utils/linalg.rs | 6 +- 36 files changed, 474 insertions(+), 398 deletions(-) delete mode 100644 src_rbd_shaders/dynamics/multibody/mass_matrix.rs diff --git a/crates/examples3d/mujoco_menagerie3.rs b/crates/examples3d/mujoco_menagerie3.rs index a2ead81..9c27130 100644 --- a/crates/examples3d/mujoco_menagerie3.rs +++ b/crates/examples3d/mujoco_menagerie3.rs @@ -2,9 +2,9 @@ use khal::backend::GpuTimestamps; use kiss3d::egui; use nexus_viewer3d::{NexusViewer, RenderMaterial}; use nexus3d::prelude::{NexusPipeline, NexusState}; -use rapier3d::prelude::*; use nexus3d::rbd::dynamics::convert_joint_motor; use nexus3d::rbd::shaders::dynamics::JointMotor; +use rapier3d::prelude::*; use rapier3d_mjcf::{MjcfLoaderOptions, MjcfMultibodyOptions, MjcfRobot, MjcfRobotHandles}; use std::fs; use std::path::{Path, PathBuf}; @@ -367,7 +367,11 @@ async fn load_scene( if let Some(key) = &keyframe { handles.apply_keyframe(&mut world.bodies, &robot, key); } - handles.bodies.iter().map(|b| b.as_ref().map(|h| h.body)).collect() + handles + .bodies + .iter() + .map(|b| b.as_ref().map(|h| h.body)) + .collect() }; // Gather each body's render geometry. Visual meshes carry the @@ -382,14 +386,15 @@ async fn load_scene( // "Render visual primitives" keeps the capsules and boxes some // models declare in their visual channel; by default only the // .obj-derived meshes are drawn. - .filter(|vm| settings.render_visual_primitives || vm.shape.as_trimesh().is_some()) + .filter(|vm| { + settings.render_visual_primitives || vm.shape.as_trimesh().is_some() + }) .collect(); let has_visual = !visuals.is_empty(); for (handle, _) in world .colliders .iter() .filter(|(_, c)| c.parent() == Some(body)) - .map(|(h, c)| (h, c)) .collect::>() { let c = &world.colliders[handle]; @@ -438,7 +443,10 @@ async fn load_scene( )); let radius = (he.x * he.x + he.y * he.y + he.z * he.z).sqrt().max(0.5); let target = Vec3::new(center.x, center.y, center.z); - camera = Some((target + Vec3::new(radius * 2.2, -radius * 2.2, radius * 1.6), target)); + camera = Some(( + target + Vec3::new(radius * 2.2, -radius * 2.2, radius * 1.6), + target, + )); } } Err(e) => eprintln!("Failed to load MJCF scene `{}`: {e}.", scene.display()), @@ -537,11 +545,15 @@ fn apply_controls( gain, ); for ah in &controls.handles.actuators { - let Some(Some(handle)) = ah.joint else { continue }; + let Some(Some(handle)) = ah.joint else { + continue; + }; let Some((mb, link_id)) = world.multibody_joints.get(handle) else { continue; }; - let Some(link) = mb.links().nth(link_id) else { continue }; + let Some(link) = mb.links().nth(link_id) else { + continue; + }; // The GPU link id is the body index (see `GpuMultibodySet::set_motor`). let body_idx = link.rigid_body_handle().into_raw_parts().0; let axes = link.joint().data.motor_axes.bits(); diff --git a/src/state.rs b/src/state.rs index e23f83b..3e24bcc 100644 --- a/src/state.rs +++ b/src/state.rs @@ -243,7 +243,7 @@ impl NexusState { /// state. Use this to run rapier-side helpers whose output you then push /// through a runtime setter — e.g. driving an MJCF actuator model and /// forwarding the resulting motors with - /// [`GpuMultibodySet::set_motors`](crate::rbd::dynamics::GpuMultibodySet::set_motors). + /// `GpuMultibodySet::set_motors` (3D only, hence no intra-doc link here). pub fn rbd_world_mut_untracked(&mut self, env: usize) -> &mut PhysicsWorld { &mut self.rbd_envs[env] } diff --git a/src_rbd/broad_phase/lbvh.rs b/src_rbd/broad_phase/lbvh.rs index 067a43c..8d9f012 100644 --- a/src_rbd/broad_phase/lbvh.rs +++ b/src_rbd/broad_phase/lbvh.rs @@ -292,11 +292,9 @@ impl Lbvh { // One thread per live collider (leaf); padding slots aren't in the tree. let colliders_per_batch = active_per_batch; - self.shaders.reset_collision_pairs.call( - pass, - [num_batches, 1, 1], - collision_pairs_len, - )?; + self.shaders + .reset_collision_pairs + .call(pass, [num_batches, 1, 1], collision_pairs_len)?; self.shaders.find_collision_pairs.call( pass, [colliders_per_batch, num_batches, 1], @@ -318,7 +316,7 @@ impl Lbvh { /// 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 + /// [`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( @@ -350,11 +348,9 @@ impl Lbvh { batch_indices, vertex_buffers, )?; - self.shaders.reset_collision_pairs.call( - pass, - [num_batches, 1, 1], - collision_pairs_len, - )?; + 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], diff --git a/src_rbd/dynamics/joint.rs b/src_rbd/dynamics/joint.rs index 4f743c1..9c89214 100644 --- a/src_rbd/dynamics/joint.rs +++ b/src_rbd/dynamics/joint.rs @@ -5,9 +5,9 @@ use crate::math::Pose; use crate::shaders::dynamics::{ - GpuInitJointConstraints, - GpuSolveJointConstraints, GpuUpdateJointConstraints, ImpulseJoint, JointConstraint, - JointConstraintBuilder, LocalMassProperties, RbdSimParams, Velocity, WorldMassProperties, + GpuInitJointConstraints, GpuSolveJointConstraints, GpuUpdateJointConstraints, ImpulseJoint, + JointConstraint, JointConstraintBuilder, LocalMassProperties, RbdSimParams, Velocity, + WorldMassProperties, }; use bytemuck::Zeroable; use khal::Shader; @@ -35,7 +35,7 @@ fn convert_joint_limits(limits: RapierJointLimits) -> JointLimits { /// Converts a rapier joint motor into the GPU representation. Public so callers /// driving actuators at runtime can push a freshly configured motor through -/// [`GpuMultibodySet::set_motor`](crate::dynamics::GpuMultibodySet::set_motor). +/// `GpuMultibodySet::set_motor`. pub fn convert_joint_motor(motor: RapierJointMotor) -> JointMotor { JointMotor { target_vel: motor.target_vel, diff --git a/src_rbd/dynamics/multibody/multibody_from_rapier.rs b/src_rbd/dynamics/multibody/multibody_from_rapier.rs index cd6bf74..3784237 100644 --- a/src_rbd/dynamics/multibody/multibody_from_rapier.rs +++ b/src_rbd/dynamics/multibody/multibody_from_rapier.rs @@ -2,10 +2,9 @@ use super::multibody_set::*; use crate::shaders::dynamics::{ - ConstraintSoftness, MAX_AXIS_CONSTRAINTS, MAX_MB_CONTACT_CONSTRAINTS_PER_MB, - MbDofCoupling, MbImpulseJointBuilder, MbImpulseJointConstraint, MultibodyContactConstraint, - MultibodyInfo, MultibodyJointConstraint, MultibodyLinkStatic, MultibodyLinkWorkspace, - RbdSimParams, + ConstraintSoftness, MAX_AXIS_CONSTRAINTS, MAX_MB_CONTACT_CONSTRAINTS_PER_MB, MbDofCoupling, + MbImpulseJointBuilder, MbImpulseJointConstraint, MultibodyContactConstraint, MultibodyInfo, + MultibodyJointConstraint, MultibodyLinkStatic, MultibodyLinkWorkspace, RbdSimParams, }; use crate::shaders::utils::linalg::MAX_MB_DOFS; use khal::BufferUsages; @@ -506,7 +505,7 @@ impl GpuMultibodySet { info_mirror, links_workspace: Tensor::vector( backend, - &crate::shaders::dynamics::ws_soa_from_structs(&all_ws, links_cap, num_batches), + crate::shaders::dynamics::ws_soa_from_structs(&all_ws, links_cap, num_batches), storage, ) .unwrap(), @@ -615,7 +614,11 @@ impl GpuMultibodySet { // 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 { + // `MAX_DELASSUS_MULTIBODIES` is currently 0, which disables the + // path; the bound is kept so raising the constant re-enables it. + #[allow(clippy::absurd_extreme_comparisons)] + let use_delassus = global_max_mb > 0 && total_mbs <= MAX_DELASSUS_MULTIBODIES; + if use_delassus { let block = (MAX_MB_CONTACT_CONSTRAINTS_PER_MB * MAX_MB_CONTACT_CONSTRAINTS_PER_MB) as usize; diff --git a/src_rbd/dynamics/multibody/multibody_set.rs b/src_rbd/dynamics/multibody/multibody_set.rs index 223805c..5f7c61d 100644 --- a/src_rbd/dynamics/multibody/multibody_set.rs +++ b/src_rbd/dynamics/multibody/multibody_set.rs @@ -4,8 +4,8 @@ use crate::math::Pose; use crate::shaders::dynamics::{ ConstraintSoftness, LocalMassProperties, MbDofCoupling, MbImpulseJointBuilder, - MbImpulseJointConstraint, MultibodyContactConstraint, MultibodyInfo, - MultibodyJointConstraint, MultibodyLinkStatic, MultibodyLinkWorkspace, RbdSimParams, + MbImpulseJointConstraint, MultibodyContactConstraint, MultibodyInfo, MultibodyJointConstraint, + MultibodyLinkStatic, MultibodyLinkWorkspace, RbdSimParams, }; use crate::shaders::utils::BatchIndices; use khal::BufferUsages; @@ -223,13 +223,13 @@ impl GpuMultibodySet { &self.dof_state } - /// GPU buffer for generalized coordinates. /// Per-batch stride of the DoF buffers (the length of each section of /// [`Self::dof_state`]). pub fn dofs_per_batch(&self) -> u32 { self.dofs_per_batch } + /// GPU buffer for generalized coordinates. pub fn dof_values(&self) -> &Tensor { &self.dof_values } @@ -274,7 +274,6 @@ impl GpuMultibodySet { self.num_internal_pgs_iterations } - /// Upload the visible-frame `dt`. Internally divides by `num_solver_iterations` /// and stores the *substep* dt (which is what the GPU kernels read). pub fn set_visible_dt(&mut self, backend: &GpuBackend, visible_dt: f32) { @@ -579,8 +578,7 @@ impl GpuMultibodySet { let a = WsAddr::new(0, self.num_batches, batch_id); let mut out = Vec::new(); for k in 0..self.links_per_batch { - let stat = &self.links_static_mirror - [(batch_id * self.links_per_batch + k) as usize]; + let stat = &self.links_static_mirror[(batch_id * self.links_per_batch + k) as usize]; let locked = stat.data.locked_axes; for axis in 0..6u32 { if locked & (1 << axis) == 0 { diff --git a/src_rbd/dynamics/multibody/multibody_solver.rs b/src_rbd/dynamics/multibody/multibody_solver.rs index f7126c5..11f2d1a 100644 --- a/src_rbd/dynamics/multibody/multibody_solver.rs +++ b/src_rbd/dynamics/multibody/multibody_solver.rs @@ -4,17 +4,15 @@ use super::multibody_set::*; use crate::math::Pose; use crate::queries::GpuIndexedContact; use crate::shaders::dynamics::{ - GpuMbBuildContactDelassus, GpuMbComputeDynamicsPre, - GpuMbFinalizeContactConstraints, GpuMbGravityAndLu, GpuMbGravityAndLuT1, GpuMbGravityAndLuT8, - GpuMbGravityAndLuT16, GpuMbGravityAndLuT32, GpuMbInitContactConstraints, - GpuMbInitJointConstraints, GpuMbIntegrate, GpuMbIntegrateVelocities, - GpuMbRemoveImpulseJointConstraintBias, - GpuMbApplyContactRestitution, GpuMbSeedContactRestitution, GpuMbSnapshotContactWarmstart, - GpuMbStashContactsLen, GpuMbTransferContactWarmstart, GpuMbWarmstartContactConstraints, - GpuMbSolveConstraints, GpuMbSolveContactsDelassus, GpuMbSolveImpulseJointConstraints, - GpuMbSolveJoints, - GpuMbFinalizeImpulseJointConstraints, - GpuMbUpdateImpulseJointConstraints, Velocity, WorldMassProperties, + GpuMbApplyContactRestitution, GpuMbBuildContactDelassus, GpuMbComputeDynamicsPre, + GpuMbFinalizeContactConstraints, GpuMbFinalizeImpulseJointConstraints, GpuMbGravityAndLu, + GpuMbGravityAndLuT1, GpuMbGravityAndLuT8, GpuMbGravityAndLuT16, GpuMbGravityAndLuT32, + GpuMbInitContactConstraints, GpuMbInitJointConstraints, GpuMbIntegrate, + GpuMbIntegrateVelocities, GpuMbRemoveImpulseJointConstraintBias, GpuMbSeedContactRestitution, + GpuMbSnapshotContactWarmstart, GpuMbSolveConstraints, GpuMbSolveContactsDelassus, + GpuMbSolveImpulseJointConstraints, GpuMbSolveJoints, GpuMbStashContactsLen, + GpuMbTransferContactWarmstart, GpuMbUpdateImpulseJointConstraints, + GpuMbWarmstartContactConstraints, Velocity, WorldMassProperties, }; use crate::shaders::utils::BatchIndices; use khal::Shader; @@ -154,7 +152,7 @@ impl GpuMultibodySolver { args.batch_indices, )?; } - let mut pass = encoder.begin_pass("[RBD] mbi/dynamics", timestamps.as_deref_mut()); + let mut pass = encoder.begin_pass("[RBD] mbi/dynamics", timestamps); self.compute_dynamics(&mut pass, mb, args) } @@ -268,8 +266,7 @@ impl GpuMultibodySolver { // substep, including the first, which carries the previous frame's. // One 64-lane workgroup per multibody (one DOF per lane). if mb.warmstart_coefficient != 0.0 { - let mut pass = - encoder.begin_pass("[RBD] mbb/warmstart-contact", timestamps.as_deref_mut()); + let mut pass = encoder.begin_pass("[RBD] mbb/warmstart-contact", timestamps); // Contact-only work: indirect grid collapses to zero workgroups // when no batch has any contact this step. self.warmstart_contact_constraints.call( @@ -325,10 +322,8 @@ impl GpuMultibodySolver { // One 64-lane workgroup per multibody. { - 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]; + 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, init_contact_dispatch, @@ -368,8 +363,7 @@ 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 mut pass = encoder.begin_pass("[RBD] mbb/build-delassus", timestamps); self.build_contact_delassus.call( &mut pass, args.mb_sweep_indirect, @@ -715,8 +709,7 @@ impl GpuMultibodySolver { 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]; + let grav_lu_dispatch = [mb.multibodies_per_batch * MB_LU_LANES, mb.num_batches, 1]; self.gravity_and_lu.call( pass, grav_lu_dispatch, diff --git a/src_rbd/dynamics/solver.rs b/src_rbd/dynamics/solver.rs index 2a8dee6..f1acaa6 100644 --- a/src_rbd/dynamics/solver.rs +++ b/src_rbd/dynamics/solver.rs @@ -11,12 +11,11 @@ use crate::math::Pose; use crate::queries::GpuIndexedContact; use crate::shaders::dynamics::{ GpuApplySolverVelsInc, GpuInitSolverBodies, GpuInitSolverVelsInc, GpuIntegrateLinearized, - GpuSolverCleanup, GpuSolverCountConstraints, GpuSolverFinalize, - GpuSolverInitConstraints, GpuSolverRefreshRhsWoBias, GpuSolverSortConstraints, - GpuSolverUpdateConstraints, GpuStepGaussSeidel, GpuStepGaussSeidelFused, GpuWarmstart, - GpuWarmstartFused, GpuWarmstartWithoutColors, - LocalMassProperties, RbdSimParams, TwoBodyConstraint, TwoBodyConstraintBuilder, Velocity, - WorldMassProperties, + GpuSolverCleanup, GpuSolverCountConstraints, GpuSolverFinalize, GpuSolverInitConstraints, + GpuSolverRefreshRhsWoBias, GpuSolverSortConstraints, GpuSolverUpdateConstraints, + GpuStepGaussSeidel, GpuStepGaussSeidelFused, GpuWarmstart, GpuWarmstartFused, + GpuWarmstartWithoutColors, LocalMassProperties, RbdSimParams, TwoBodyConstraint, + TwoBodyConstraintBuilder, Velocity, WorldMassProperties, }; use crate::utils::{GpuPrefixSum, PrefixSumWorkspace}; use khal::Shader; @@ -567,7 +566,7 @@ impl GpuSolver { * back to body-origin poses. */ { - let mut pass = encoder.begin_pass("[RBD] slv/finalize", timestamps.as_deref_mut()); + let mut pass = encoder.begin_pass("[RBD] slv/finalize", timestamps); self.finalize.call( &mut pass, [args.num_colliders, args.num_batches, 1], diff --git a/src_rbd/pipeline/insertion_removal.rs b/src_rbd/pipeline/insertion_removal.rs index 2f7315f..a19e042 100644 --- a/src_rbd/pipeline/insertion_removal.rs +++ b/src_rbd/pipeline/insertion_removal.rs @@ -181,7 +181,7 @@ impl RbdState { 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], + vec![0u32; (collisions_capacity * num_batches) as usize], storage, ) .unwrap(); @@ -503,7 +503,8 @@ impl RbdState { 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_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 @@ -543,9 +544,9 @@ impl RbdState { hole_global, 1, )?; - any_copy = true; }}; } + any_copy = true; relocate!(self.body_poses, staging_pose); relocate!(self.solver_body_poses, staging_pose); relocate!(self.collider_world_poses, staging_pose); diff --git a/src_rbd/pipeline/rbd_state.rs b/src_rbd/pipeline/rbd_state.rs index 5f69ce7..2016e00 100644 --- a/src_rbd/pipeline/rbd_state.rs +++ b/src_rbd/pipeline/rbd_state.rs @@ -149,6 +149,7 @@ pub struct RbdState { /// 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), @@ -346,10 +347,7 @@ impl RbdState { &self.gravity } - pub(super) fn gravity_tensor( - backend: &GpuBackend, - gravity: [f32; 3], - ) -> Tensor { + pub(super) fn gravity_tensor(backend: &GpuBackend, gravity: [f32; 3]) -> Tensor { Tensor::scalar( backend, glamx::Vec4::new(gravity[0], gravity[1], gravity[2], 0.0), diff --git a/src_rbd/pipeline/rbd_state_from_rapier.rs b/src_rbd/pipeline/rbd_state_from_rapier.rs index 9bf22ca..042de1c 100644 --- a/src_rbd/pipeline/rbd_state_from_rapier.rs +++ b/src_rbd/pipeline/rbd_state_from_rapier.rs @@ -232,7 +232,10 @@ impl RbdState { // Handle bodies whose multibody disables self-contacts. #[cfg(feature = "dim3")] - let no_self_collide: HashMap = { + let no_self_collide: HashMap< + crate::rapier::dynamics::RigidBodyHandle, + u32, + > = { let mut map = HashMap::new(); for (mb_ord, mb) in multibody_joints.multibodies().enumerate() { if !mb.self_contacts_enabled() { @@ -661,7 +664,7 @@ impl RbdState { .unwrap(); let old_constraints_colors = Tensor::vector( backend, - &vec![0u32; (capacities.collisions_capacity * num_batches) as usize], + vec![0u32; (capacities.collisions_capacity * num_batches) as usize], storage, ) .unwrap(); diff --git a/src_rbd/pipeline/rbd_step.rs b/src_rbd/pipeline/rbd_step.rs index b28397a..d208f32 100644 --- a/src_rbd/pipeline/rbd_step.rs +++ b/src_rbd/pipeline/rbd_step.rs @@ -140,8 +140,7 @@ impl RbdPipeline { 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()); + let mut pass = encoder.begin_pass("[RBD] bf-find-pairs", timestamps.as_deref_mut()); self.lbvh.brute_force_pairs( backend, &mut pass, @@ -320,99 +319,103 @@ impl RbdPipeline { stats.num_colors = state.max_colors + 1; drop(pass); } else { + // Warmstart + let warmstart_args = WarmstartArgs { + contacts_len: &state.contacts_len, + old_body_constraint_counts: &state.old_constraints_counts, + old_constraint_builders: &state.old_constraint_builders, + old_body_constraint_ids: &state.old_body_constraint_ids, + old_constraints: &state.old_constraints, + new_constraints: &mut state.new_constraints, + new_constraint_builders: &state.new_constraint_builders, + contacts_len_indirect: &state.contacts_indirect, + batch_indices: &state.batch_indices, + }; - // Warmstart - let warmstart_args = WarmstartArgs { - contacts_len: &state.contacts_len, - old_body_constraint_counts: &state.old_constraints_counts, - old_constraint_builders: &state.old_constraint_builders, - old_body_constraint_ids: &state.old_body_constraint_ids, - old_constraints: &state.old_constraints, - new_constraints: &mut state.new_constraints, - new_constraint_builders: &state.new_constraint_builders, - contacts_len_indirect: &state.contacts_indirect, - batch_indices: &state.batch_indices, - }; - - 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)?; + 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)?; - 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_iterations(&mut pass, coloring_args, state.max_colors)?; + // 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, + 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_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. - 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, - )?; + // 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; + // `+1` because solver iterates 1..=max_colors (color 0 is unassigned). + let num_colors = state.max_colors + 1; + stats.num_colors = num_colors; - drop(pass); + drop(pass); } if !merge_submits { backend.submit(encoder)?; @@ -558,12 +561,9 @@ impl RbdPipeline { 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.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); } @@ -610,7 +610,7 @@ impl RbdPipeline { // 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)?; + 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/broad_phase/brute_force.rs b/src_rbd_shaders/broad_phase/brute_force.rs index fada07e..94601b1 100644 --- a/src_rbd_shaders/broad_phase/brute_force.rs +++ b/src_rbd_shaders/broad_phase/brute_force.rs @@ -39,7 +39,10 @@ pub fn gpu_bf_compute_aabbs( 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)); + 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, diff --git a/src_rbd_shaders/broad_phase/lbvh.rs b/src_rbd_shaders/broad_phase/lbvh.rs index 368e8d7..31051b4 100644 --- a/src_rbd_shaders/broad_phase/lbvh.rs +++ b/src_rbd_shaders/broad_phase/lbvh.rs @@ -71,7 +71,7 @@ pub const MAX_REDUCE_LANES: u32 = 256; /// /// 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, @@ -307,7 +307,8 @@ pub fn gpu_lbvh_build( tree.at_mut(node_id as usize).left = left as u32; tree.at_mut(node_id as usize).right = right as u32; - tree.at_mut(node_id as usize).refit_count_or_max_subtree_index = 0; // Might as well reset the refit count here. + tree.at_mut(node_id as usize) + .refit_count_or_max_subtree_index = 0; // Might as well reset the refit count here. tree.at_mut(left as usize).parent = node_id; tree.at_mut(right as usize).parent = node_id; } @@ -350,7 +351,8 @@ pub fn gpu_lbvh_refit_leaves( 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; + tree.at_mut(curr_leaf_id as usize) + .refit_count_or_max_subtree_index = i; } } @@ -393,7 +395,12 @@ pub fn gpu_lbvh_refit_internal( // Maximum tree depth is log2(num_colliders), but we use 32 as a safe upper bound. for _level in 0..32u32 { if thread_is_active { - let refit_count = atomic_add_u32(&mut tree.at_mut(curr_id as usize).refit_count_or_max_subtree_index, 1); + 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 @@ -415,7 +422,8 @@ pub fn gpu_lbvh_refit_internal( // 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); + 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. @@ -484,7 +492,12 @@ pub fn gpu_lbvh_refit( // 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); + 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 diff --git a/src_rbd_shaders/broad_phase/narrow_phase.rs b/src_rbd_shaders/broad_phase/narrow_phase.rs index d910bb2..86d0e1e 100644 --- a/src_rbd_shaders/broad_phase/narrow_phase.rs +++ b/src_rbd_shaders/broad_phase/narrow_phase.rs @@ -14,10 +14,7 @@ use crate::{PaddedVector, Pose, Vector}; use khal_std::glamx::UVec3; use khal_std::index::MaybeIndexUnchecked; use khal_std::macros::{spirv, spirv_bindgen}; -use khal_std::{ - iter::StepRng, - sync::atomic_add_u32, -}; +use khal_std::{iter::StepRng, sync::atomic_add_u32}; use super::lbvh::{MAX_REDUCE_LANES, max_len_indirect_args}; use crate::broad_phase::CollisionPair; diff --git a/src_rbd_shaders/dynamics/color_buckets.rs b/src_rbd_shaders/dynamics/color_buckets.rs index 0ee2fe2..b0478ab 100644 --- a/src_rbd_shaders/dynamics/color_buckets.rs +++ b/src_rbd_shaders/dynamics/color_buckets.rs @@ -112,8 +112,10 @@ pub fn gpu_color_buckets_scatter( 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); + 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 432533a..5a8ec75 100644 --- a/src_rbd_shaders/dynamics/mod.rs +++ b/src_rbd_shaders/dynamics/mod.rs @@ -32,8 +32,8 @@ pub use joint::{ JointMotor, LIN_AXES_MASK, MotorParameters, SPATIAL_DIM, }; pub use joint_constraint::*; -pub use joint_constraint_builder::{JointConstraintBuilder, JointConstraintHelper, new_helper}; pub(crate) use joint_constraint_builder::smallest_abs_diff_between_sin_angles; +pub use joint_constraint_builder::{JointConstraintBuilder, JointConstraintHelper, new_helper}; pub use multibody::*; pub use sim_params::*; // Re-export solver items; update_constraint comes from joint_constraint_builder for joints diff --git a/src_rbd_shaders/dynamics/multibody/compute_dynamics_pre.rs b/src_rbd_shaders/dynamics/multibody/compute_dynamics_pre.rs index ac93c7a..8dac2a8 100644 --- a/src_rbd_shaders/dynamics/multibody/compute_dynamics_pre.rs +++ b/src_rbd_shaders/dynamics/multibody/compute_dynamics_pre.rs @@ -8,8 +8,8 @@ //! The follow-up `gpu_mb_gravity_and_lu` kernel finishes the dynamics pipeline //! (gravity rhs + LU factor + LU solve). -use khal_std::glamx::UVec3; use glamx::Vec4; +use khal_std::glamx::UVec3; use khal_std::index::MaybeIndexUnchecked; use khal_std::macros::{spirv, spirv_bindgen}; use khal_std::sync::workgroup_memory_barrier_with_group_sync; @@ -22,13 +22,12 @@ use super::ws_soa::{ }; use crate::dynamics::body::Velocity; use crate::dynamics::joint::SPATIAL_DIM; -#[cfg(feature = "dim3")] -use crate::utils::linalg::{gemm_inertia_lhs_cross_buf_par, gemm_skew_lhs_cross_buf_par}; use crate::utils::linalg::{ - axpy_mat_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, + axpy_mat_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, }; +#[cfg(feature = "dim3")] +use crate::utils::linalg::{gemm_inertia_lhs_cross_buf_par, gemm_skew_lhs_cross_buf_par}; use crate::utils::{BatchIndices, ISlice, SliceMut}; use crate::{ANG_DIM, AngVector, DIM, Pose, Vector, gcross_av}; use parry::math::VectorExt; @@ -69,8 +68,7 @@ pub fn gpu_mb_compute_dynamics_pre( #[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 [Vec4], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] 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], @@ -84,9 +82,7 @@ pub fn gpu_mb_compute_dynamics_pre( let dt = *dt_uniform; let mb = if active_slot { - batch_ids - .ib(batch_id, multibody_info) - .read(mb_idx as usize) + batch_ids.ib(batch_id, multibody_info).read(mb_idx as usize) } else { MultibodyInfo::default() }; @@ -125,7 +121,14 @@ 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, links_workspace, wa, num_links); + forward_kinematics( + &mb, + &stat_slice, + &mut poses_slice, + links_workspace, + wa, + num_links, + ); } sync_slots(t); @@ -185,7 +188,8 @@ pub fn gpu_mb_compute_dynamics_pre( inv_mass_x = lmp.inv_mass.x; if split && inv_mass_x == 0.0 { - let coriolis_block = batch_ids.imat(batch_id, + let coriolis_block = batch_ids.imat( + batch_id, mb_cor_base + (k as usize) * (DIM as usize) * (ndofs as usize), DIM, ndofs, @@ -193,7 +197,8 @@ pub fn gpu_mb_compute_dynamics_pre( fill_par(coriolis_packed, coriolis_block, 0.0, lane, t); fill_par( coriolis_packed, - batch_ids.imat(batch_id, + batch_ids.imat( + batch_id, mb_cor_w_base + (k as usize) * (DIM as usize) * (ndofs as usize), DIM, ndofs, @@ -209,17 +214,20 @@ 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 = batch_ids.imat(batch_id, + 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 = batch_ids.imat(batch_id, + 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 = batch_ids.imat(batch_id, + let body_jacobian = batch_ids.imat( + batch_id, mb_jac_base + (k as usize) * SPATIAL_DIM * (ndofs as usize), SPATIAL_DIM as u32, ndofs, @@ -232,7 +240,11 @@ pub fn gpu_mb_compute_dynamics_pre( mass = 1.0 / inv_mass_x; rb_inertia = ws_world_inertia(links_workspace, wa, k, &lmp); - let quad_target = if split { plain_mass } else { acc_augmented_mass }; + let quad_target = if split { + plain_mass + } else { + acc_augmented_mass + }; quadform_spatial_par( mass_matrices, quad_target, @@ -249,18 +261,21 @@ pub fn gpu_mb_compute_dynamics_pre( if split && k != 0 { let stat = stat_slice[k as usize]; let parent_id = stat.parent_link_id; - let parent_j = batch_ids.imat(batch_id, + 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 = batch_ids.imat(batch_id, + 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 = batch_ids.imat(batch_id, + 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, @@ -269,20 +284,8 @@ pub fn gpu_mb_compute_dynamics_pre( 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, - coriolis_v_i, - parent_coriolis_v, - lane, - t, - ); - copy_from_par( - coriolis_packed, - coriolis_w_i, - parent_coriolis_w, - lane, - t, - ); + copy_from_par(coriolis_packed, coriolis_v_i, parent_coriolis_v, lane, t); + copy_from_par(coriolis_packed, coriolis_w_i, parent_coriolis_w, lane, t); gemm_skew_tr_lhs_par( coriolis_packed, @@ -366,9 +369,12 @@ pub fn gpu_mb_compute_dynamics_pre( #[cfg(feature = "dim3")] { - let parent_w_skew = crate::utils::linalg::skew( - ws_vel_ang(links_workspace, wa, parent_id, WS_RB_VELS), - ); + 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); @@ -524,7 +530,11 @@ pub fn gpu_mb_compute_dynamics_pre( let diag = damping_slice[d as usize] * dt + armature_slice[d as usize] + stiffness_slice[d as usize] * dt * dt; - let diag_target = if split { plain_mass } else { acc_augmented_mass }; + let diag_target = if split { + plain_mass + } else { + acc_augmented_mass + }; let diag_idx = diag_target.idx(d, d); let cur = mass_matrices.read(diag_idx); mass_matrices.write(diag_idx, cur + diag); @@ -621,7 +631,13 @@ fn forward_kinematics( }; ws_set_pose(ws, wa, 0, WS_LTP, root_pose); ws_set_pose(ws, wa, 0, WS_LTW, root_pose); - ws_set_vec(ws, wa, 0, WS_WORLD_COM, root_pose * root_config.local_mprops.com); + ws_set_vec( + ws, + wa, + 0, + WS_WORLD_COM, + root_pose * root_config.local_mprops.com, + ); for k in 1..num_links { let k_usize = k as usize; @@ -670,7 +686,8 @@ fn update_body_jacobians( // value per node. for k in 0..max_links { let mut parent_to_world = Pose::default(); - let link_j = batch_ids.imat(batch_id, + let link_j = batch_ids.imat( + batch_id, mb_jac_base + (k as usize) * SPATIAL_DIM * (ndofs as usize), SPATIAL_DIM as u32, ndofs, @@ -680,7 +697,8 @@ fn update_body_jacobians( let link_infos = &stat_slice[k as usize]; if k != 0 { - let parent_j = batch_ids.imat(batch_id, + 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, diff --git a/src_rbd_shaders/dynamics/multibody/contact_constraints.rs b/src_rbd_shaders/dynamics/multibody/contact_constraints.rs index aeb991b..075e07f 100644 --- a/src_rbd_shaders/dynamics/multibody/contact_constraints.rs +++ b/src_rbd_shaders/dynamics/multibody/contact_constraints.rs @@ -29,9 +29,9 @@ use crate::utils::linalg::{MAX_MB_DOFS, MatSlice, VSlice, lu_solve_in_place}; use crate::{ANG_DIM, AngVector, DIM, Pose, Vector, gcross, gdot}; use super::types::{ - CONTACT_CONSTRAINTS_PER_POINT, MAX_MB_CONTACT_CONSTRAINTS_PER_MB, - MB_CONTACT_KIND_INACTIVE, MB_CONTACT_KIND_NORMAL, MB_CONTACT_KIND_TANGENT, - MultibodyContactConstraint, MultibodyInfo, MultibodyLinkStatic, + CONTACT_CONSTRAINTS_PER_POINT, MAX_MB_CONTACT_CONSTRAINTS_PER_MB, MB_CONTACT_KIND_INACTIVE, + MB_CONTACT_KIND_NORMAL, MB_CONTACT_KIND_TANGENT, MultibodyContactConstraint, MultibodyInfo, + MultibodyLinkStatic, }; use super::utils::zero_kinematic_dofs; use super::ws_soa::{WS_LTW, WS_WORLD_COM, WsAddr, ws_pose, ws_vec}; @@ -84,8 +84,13 @@ 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::interleaved(link_jac_base, SPATIAL_DIM as u32, ndofs, jac_stride, jac_shift); + 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); for j in 0..ndofs { // Linear contribution: `unit_force · J_v[:, j]`. @@ -468,8 +473,7 @@ pub fn gpu_mb_init_contact_constraints( }; let (torque_b_tang, ang_jac_tang, ii_ang_jac_tang) = if is_self { - let shift_b = - p2 - ws_vec(links_workspace, wa, mb_link_id_b, WS_WORLD_COM); + let shift_b = p2 - ws_vec(links_workspace, wa, mb_link_id_b, WS_WORLD_COM); #[cfg(feature = "dim3")] { ( @@ -851,7 +855,6 @@ pub fn gpu_mb_finalize_contact_constraints( } } - /// Carries the accumulated contact impulses of the previous frame over to this /// frame's freshly built slots. A point is matched by the pair of links (or /// link and free body) it touches plus the proximity of both frozen local @@ -1014,7 +1017,11 @@ pub fn gpu_mb_seed_contact_restitution( } else { cons.restitution >= 1.0 }; - cons.restitution_seed = if bouncy { cons.restitution * j_dot_v } else { 0.0 }; + cons.restitution_seed = if bouncy { + cons.restitution * j_dot_v + } else { + 0.0 + }; contact_constraints.write(cons_base + s as usize, cons); } } @@ -1037,7 +1044,7 @@ pub fn gpu_mb_apply_contact_restitution( #[spirv(uniform, descriptor_set = 0, binding = 4)] 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)] dof_v: &mut [f32; MAX_MB_DOFS], #[spirv(workgroup)] scratch: &mut [f32; 64], #[spirv(workgroup)] delta_shared: &mut f32, ) { diff --git a/src_rbd_shaders/dynamics/multibody/gravity_and_lu.rs b/src_rbd_shaders/dynamics/multibody/gravity_and_lu.rs index 2ce9953..f328c9c 100644 --- a/src_rbd_shaders/dynamics/multibody/gravity_and_lu.rs +++ b/src_rbd_shaders/dynamics/multibody/gravity_and_lu.rs @@ -25,8 +25,8 @@ use crate::{AngVector, Vector, gcross_av}; use super::lu::{ 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, + 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}; use super::ws_soa::{ @@ -88,8 +88,7 @@ pub fn gpu_mb_gravity_and_lu( #[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 [Vec4], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] 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], @@ -114,9 +113,7 @@ pub fn gpu_mb_gravity_and_lu( let max_ndofs = batch_ids.mb_max_ndofs; let max_links = batch_ids.mb_max_links; - let mb = batch_ids - .ib(batch_id, multibody_info) - .read(mb_idx as usize); + let mb = batch_ids.ib(batch_id, multibody_info).read(mb_idx as usize); let num_links = mb.num_links; let ndofs = mb.ndofs; let mb_jac_base = mb.jacobian_offset as usize; @@ -216,7 +213,13 @@ pub fn gpu_mb_gravity_and_lu( acc_lin += gcross_av(acc_ang, self_shift23); if lane == 0 { - ws_set_vel(links_workspace, wa, k, WS_KIN_ACC, Velocity::new(acc_lin, acc_ang)); + ws_set_vel( + links_workspace, + wa, + k, + WS_KIN_ACC, + Velocity::new(acc_lin, acc_ang), + ); } } @@ -242,13 +245,13 @@ pub fn gpu_mb_gravity_and_lu( let gyroscopic: AngVector = 0.0; let i_acc_ang = rb_inertia * acc_ang; - let (ext_force, ext_torque, gravity_scale) = - ws_ext_wrench(links_workspace, wa, k); + let (ext_force, ext_torque, gravity_scale) = ws_ext_wrench(links_workspace, wa, k); let f_lin = g * (mass * gravity_scale) + ext_force - acc_lin * mass; let f_ang = ext_torque - gyroscopic - i_acc_ang; - let body_jacobian = batch_ids.imat(batch_id, + let body_jacobian = batch_ids.imat( + batch_id, mb_jac_base + (k as usize) * SPATIAL_DIM * (ndofs as usize), SPATIAL_DIM as u32, ndofs, @@ -353,7 +356,10 @@ pub fn gpu_mb_gravity_and_lu( lu_triangular_solve_in_place(ndofs, max_ndofs, lane, mat, x, partial); if lane < ndofs { - gen_forces.write(batch_ids.mbi(batch_id, gen_base + 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), + ); } // ---- Phase 5 (split mode only): factor the plain matrix and persist its @@ -425,9 +431,7 @@ fn gravity_and_lu_packed_impl(slot, r, lane))); + mass_matrices.write( + m_view.idx(r, lane), + mat.read(sm_idx_packed::(slot, r, lane)), + ); } } @@ -704,8 +717,10 @@ fn gravity_and_lu_packed_impl(slot, r, lane))); + mass_matrices.write( + m_view.idx(r, lane), + mat.read(sm_idx_packed::(slot, r, lane)), + ); } } } @@ -719,8 +734,7 @@ pub fn gpu_mb_gravity_and_lu_t1( #[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 [Vec4], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] 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], @@ -737,9 +751,7 @@ pub fn gpu_mb_gravity_and_lu_t1( let batch_id = invocation_id.x / num_mb; let mb_idx = invocation_id.x % num_mb; - let mb = batch_ids - .ib(batch_id, multibody_info) - .read(mb_idx as usize); + let mb = batch_ids.ib(batch_id, multibody_info).read(mb_idx as usize); let num_links = mb.num_links; let ndofs = mb.ndofs; if ndofs == 0 { @@ -829,7 +841,13 @@ pub fn gpu_mb_gravity_and_lu_t1( acc_lin += gcross_av(rb_ang, gcross_av(rb_ang, self_shift23)); acc_lin += gcross_av(acc_ang, self_shift23); - ws_set_vel(links_workspace, wa, k, WS_KIN_ACC, Velocity::new(acc_lin, acc_ang)); + ws_set_vel( + links_workspace, + wa, + k, + WS_KIN_ACC, + Velocity::new(acc_lin, acc_ang), + ); let lmp = stat_slice[k as usize].local_mprops; let inv_mass_x = lmp.inv_mass.x; @@ -851,7 +869,8 @@ pub fn gpu_mb_gravity_and_lu_t1( let f_lin = g * (mass * gravity_scale) + ext_force - acc_lin * mass; let f_ang = ext_torque - gyroscopic - i_acc_ang; - let body_jacobian = batch_ids.imat(batch_id, + let body_jacobian = batch_ids.imat( + batch_id, mb_jac_base + (k as usize) * SPATIAL_DIM * (ndofs as usize), SPATIAL_DIM as u32, ndofs, diff --git a/src_rbd_shaders/dynamics/multibody/impulse_joint_constraints/helper.rs b/src_rbd_shaders/dynamics/multibody/impulse_joint_constraints/helper.rs index 477d6d9..7a875a8 100644 --- a/src_rbd_shaders/dynamics/multibody/impulse_joint_constraints/helper.rs +++ b/src_rbd_shaders/dynamics/multibody/impulse_joint_constraints/helper.rs @@ -13,8 +13,8 @@ use crate::{AngVector, MAX_FLT, Pose, Vector, rotation_to_matrix}; use super::super::types::MultibodyInfo; use super::jacobians::*; -use crate::utils::linalg::VSlice; use super::types::*; +use crate::utils::linalg::VSlice; /// `JointConstraintHelper`-equivalent: precomputed per-joint quantities used /// by `lock_*`, `limit_*`, `motor_*`. Mirrors the homonymous rapier struct. @@ -160,9 +160,9 @@ impl AngularLimitParams { pub(super) fn new(min: f32, max: f32) -> Self { let half_range = (max - min) * 0.5; // A range of a full turn or more is indistinguishable from "no limit" - // for an angle read off a relative rotation. `!(x < PI)` also catches - // NaN bounds before they poison the row. - if !(half_range < core::f32::consts::PI) { + // for an angle read off a relative rotation. NaN bounds are rejected + // here too, before they poison the row. + if half_range.is_nan() || half_range >= core::f32::consts::PI { return Self { center: 0.0, half_range: 10.0, diff --git a/src_rbd_shaders/dynamics/multibody/impulse_joint_constraints/jacobians.rs b/src_rbd_shaders/dynamics/multibody/impulse_joint_constraints/jacobians.rs index c7a589d..2098eec 100644 --- a/src_rbd_shaders/dynamics/multibody/impulse_joint_constraints/jacobians.rs +++ b/src_rbd_shaders/dynamics/multibody/impulse_joint_constraints/jacobians.rs @@ -92,8 +92,7 @@ pub(super) fn side_dot_vel_par( } else { // SIDE_KIND_MB if lane < ndofs { - jacobians.read(j_id as usize + lane as usize) - * dof_vels.read(dof_base_for_mb.at(lane)) + jacobians.read(j_id as usize + lane as usize) * dof_vels.read(dof_base_for_mb.at(lane)) } else { 0.0f32 } @@ -250,8 +249,13 @@ pub(super) fn fill_mb_jacobians( let ndofs = mb.ndofs; let mb_jac_base = mb.jacobian_offset as usize; let link_jac_base = mb_jac_base + (link_id as usize) * SPATIAL_DIM * (ndofs as usize); - let link_j = - MatSlice::interleaved(link_jac_base, SPATIAL_DIM as u32, ndofs, il.stride, il.shift); + let link_j = MatSlice::interleaved( + link_jac_base, + SPATIAL_DIM as u32, + ndofs, + il.stride, + il.shift, + ); let (link_j_v, link_j_w) = link_j.rows_range_pair(0, DIM, DIM, ANG_DIM); // 1) j = link_J^T · (unit_force, unit_torque). Same kernel used by diff --git a/src_rbd_shaders/dynamics/multibody/impulse_joint_constraints/kernels.rs b/src_rbd_shaders/dynamics/multibody/impulse_joint_constraints/kernels.rs index bbaeb0a..66643ee 100644 --- a/src_rbd_shaders/dynamics/multibody/impulse_joint_constraints/kernels.rs +++ b/src_rbd_shaders/dynamics/multibody/impulse_joint_constraints/kernels.rs @@ -1,8 +1,8 @@ //! The four compute entry points of the multibody impulse-joint pipeline //! (update / finalize / solve / remove-bias). -use khal_std::glamx::UVec3; use glamx::Vec4; +use khal_std::glamx::UVec3; use khal_std::index::MaybeIndexUnchecked; use khal_std::macros::{spirv, spirv_bindgen}; use khal_std::sync::workgroup_memory_barrier_with_group_sync; @@ -36,8 +36,7 @@ pub fn gpu_mb_update_impulse_joint_constraints( #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] jacobians: &mut [f32], #[spirv(uniform, descriptor_set = 0, binding = 3)] softness: &ConstraintSoftness, #[spirv(storage_buffer, descriptor_set = 1, binding = 0)] multibody_info: &[MultibodyInfo], - #[spirv(storage_buffer, descriptor_set = 1, binding = 1)] - links_workspace: &[Vec4], + #[spirv(storage_buffer, descriptor_set = 1, binding = 1)] links_workspace: &[Vec4], #[spirv(storage_buffer, descriptor_set = 1, binding = 2)] body_jacobians: &[f32], #[spirv(storage_buffer, descriptor_set = 1, binding = 3)] poses: &[Pose], #[spirv(storage_buffer, descriptor_set = 1, binding = 4)] mprops: &[WorldMassProperties], 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 b0632cb..db284db 100644 --- a/src_rbd_shaders/dynamics/multibody/impulse_joint_constraints/update.rs +++ b/src_rbd_shaders/dynamics/multibody/impulse_joint_constraints/update.rs @@ -7,7 +7,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::ISlice; -use crate::utils::linalg::{MatSlice, lu_solve_in_place, VSlice}; +use crate::utils::linalg::{MatSlice, VSlice, lu_solve_in_place}; use crate::{DIM, Pose}; use super::super::types::{MultibodyInfo, MultibodyLinkStatic}; @@ -49,7 +49,14 @@ pub(super) fn solve_mb_wj( 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)); + lu_solve_in_place( + mass_matrices, + m, + lu_pivots, + piv, + jacobians, + VSlice::dense(wj_base), + ); // Kinematic dofs are user-driven: the impulse must not move them. let stat_slice = ISlice { diff --git a/src_rbd_shaders/dynamics/multibody/integrate.rs b/src_rbd_shaders/dynamics/multibody/integrate.rs index 9414a3b..4d4e03f 100644 --- a/src_rbd_shaders/dynamics/multibody/integrate.rs +++ b/src_rbd_shaders/dynamics/multibody/integrate.rs @@ -4,9 +4,8 @@ //! After this pass, callers are expected to re-run forward kinematics to //! refresh link poses. -use khal_std::glamx::UVec3; use glamx::Vec4; -use khal_std::index::MaybeIndexUnchecked; +use khal_std::glamx::UVec3; use khal_std::macros::{spirv, spirv_bindgen}; #[cfg(feature = "dim2")] @@ -19,7 +18,9 @@ use crate::{Vector, rotation_from_scaled_axis, rotation_renormalize_fast}; use parry::math::VectorExt; use super::types::{MultibodyInfo, MultibodyLinkStatic}; -use super::ws_soa::{WS_JOINT_ROT, WsAddr, ws_coord, ws_rot, ws_set_coord, ws_set_rot}; +#[cfg(feature = "dim3")] +use super::ws_soa::ws_rot; +use super::ws_soa::{WS_JOINT_ROT, WsAddr, ws_coord, ws_set_coord, ws_set_rot}; /// Update generalized velocities: `v += a · dt`. /// @@ -44,9 +45,7 @@ pub fn gpu_mb_integrate_velocities( let mb_idx = invocation_id.x % num_mb; let dt = *dt_uniform; - let mb = batch_ids - .ib(batch_id, multibody_info) - .read(mb_idx as usize); + let mb = batch_ids.ib(batch_id, multibody_info).read(mb_idx as usize); let mut dof_vel = batch_ids .ib_mut(batch_id, dof_state) @@ -82,9 +81,7 @@ pub fn gpu_mb_integrate( let mb_idx = invocation_id.x % num_mb; let dt = *dt_uniform; - let mb = batch_ids - .ib(batch_id, multibody_info) - .read(mb_idx as usize); + let mb = batch_ids.ib(batch_id, multibody_info).read(mb_idx as usize); let num_links = mb.num_links; let stat_slice = batch_ids @@ -143,7 +140,13 @@ pub fn gpu_mb_integrate( let v = dof_vel[aid + curr_free as usize]; let new = ws_coord(links_workspace, wa, k, DIM) + v * dt; ws_set_coord(links_workspace, wa, k, DIM, new); - ws_set_rot(links_workspace, wa, k, WS_JOINT_ROT, rotation_from_angle(new)); + ws_set_rot( + links_workspace, + wa, + k, + WS_JOINT_ROT, + rotation_from_angle(new), + ); } } else if num_ang == 3 { #[cfg(feature = "dim3")] diff --git a/src_rbd_shaders/dynamics/multibody/joint_constraints.rs b/src_rbd_shaders/dynamics/multibody/joint_constraints.rs index cc2a292..19f669d 100644 --- a/src_rbd_shaders/dynamics/multibody/joint_constraints.rs +++ b/src_rbd_shaders/dynamics/multibody/joint_constraints.rs @@ -4,8 +4,8 @@ //! sweeps. Per-multibody, all constraint slots are scanned (`kind == 0` ones //! are skipped). -use khal_std::glamx::UVec3; use glamx::Vec4; +use khal_std::glamx::UVec3; use khal_std::index::MaybeIndexUnchecked; use khal_std::iter::StepRng; use khal_std::macros::{spirv, spirv_bindgen}; @@ -18,9 +18,8 @@ use crate::utils::linalg::{MatSlice, VSlice, lu_solve_in_place}; use crate::{DIM, MAX_FLT}; use super::types::{ - MB_JOINT_KIND_COUPLING, MB_JOINT_KIND_LIMIT, MB_JOINT_KIND_LIMIT_INACTIVE, - MB_JOINT_KIND_MOTOR, MbDofCoupling, MultibodyInfo, MultibodyJointConstraint, - MultibodyLinkStatic, + MB_JOINT_KIND_COUPLING, MB_JOINT_KIND_LIMIT, MB_JOINT_KIND_LIMIT_INACTIVE, MB_JOINT_KIND_MOTOR, + MbDofCoupling, MultibodyInfo, MultibodyJointConstraint, MultibodyLinkStatic, }; use super::ws_soa::{WsAddr, ws_coord}; @@ -435,16 +434,14 @@ pub fn gpu_mb_init_joint_constraints( #[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: &[Vec4], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] 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)] joint_constraints: &mut [MultibodyJointConstraint], #[spirv(storage_buffer, descriptor_set = 0, binding = 6)] joint_constraint_columns: &mut [f32], - #[spirv(storage_buffer, descriptor_set = 0, binding = 7)] - dof_couplings: &[MbDofCoupling], + #[spirv(storage_buffer, descriptor_set = 0, binding = 7)] dof_couplings: &[MbDofCoupling], #[spirv(uniform, descriptor_set = 0, binding = 8)] softness: &ConstraintSoftness, #[spirv(uniform, descriptor_set = 0, binding = 9)] batch_ids: &BatchIndices, ) { @@ -459,9 +456,7 @@ pub fn gpu_mb_init_joint_constraints( return; } - let mb = batch_ids - .ib(batch_id, multibody_info) - .read(mb_idx as usize); + let mb = batch_ids.ib(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 { diff --git a/src_rbd_shaders/dynamics/multibody/mass_matrix.rs b/src_rbd_shaders/dynamics/multibody/mass_matrix.rs deleted file mode 100644 index a9a2f08..0000000 --- a/src_rbd_shaders/dynamics/multibody/mass_matrix.rs +++ /dev/null @@ -1,38 +0,0 @@ -#[cfg(feature = "dim3")] -use glamx::Mat3; - -use crate::dynamics::body::LocalMassProperties; -#[cfg(feature = "dim3")] -use crate::rotation_to_matrix; - -use super::types::MultibodyLinkWorkspace; - -impl MultibodyLinkWorkspace { - /// World-space inertia for this link. - /// - /// In 3D returns a `Mat3` (`I_world = R · diag(principal_inertia) · Rᵀ`). In 2D - /// returns the scalar moment of inertia (already in world frame because there - /// is only one rotational DOF). - #[cfg(feature = "dim3")] - #[inline] - pub(super) fn link_world_inertia(&self, lmp: &LocalMassProperties) -> Mat3 { - 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 r = rotation_to_matrix(self.local_to_world.rotation * lmp.inertia_ref_frame); - // M = r · diag(px, py, pz) (column-scale); I = M · rᵀ. - let m = Mat3::from_cols(r.x_axis * px, r.y_axis * py, r.z_axis * pz); - m * r.transpose() - } - - #[cfg(feature = "dim2")] - #[inline] - pub(super) fn link_world_inertia(&self, lmp: &LocalMassProperties) -> f32 { - if lmp.inv_inertia != 0.0 { - 1.0 / lmp.inv_inertia - } else { - 0.0 - } - } -} diff --git a/src_rbd_shaders/dynamics/multibody/mod.rs b/src_rbd_shaders/dynamics/multibody/mod.rs index 45ed4e2..2b66446 100644 --- a/src_rbd_shaders/dynamics/multibody/mod.rs +++ b/src_rbd_shaders/dynamics/multibody/mod.rs @@ -21,7 +21,6 @@ mod integrate; mod jacobian; mod joint_constraints; mod lu; -mod mass_matrix; mod solve_constraints; mod types; mod utils; diff --git a/src_rbd_shaders/dynamics/multibody/solve_constraints.rs b/src_rbd_shaders/dynamics/multibody/solve_constraints.rs index cd6ff3e..0c89879 100644 --- a/src_rbd_shaders/dynamics/multibody/solve_constraints.rs +++ b/src_rbd_shaders/dynamics/multibody/solve_constraints.rs @@ -55,7 +55,7 @@ pub fn gpu_mb_solve_constraints( #[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)] dof_v: &mut [f32; MAX_MB_DOFS], #[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, @@ -81,8 +81,7 @@ pub fn gpu_mb_solve_constraints( 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 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; @@ -107,7 +106,6 @@ pub fn gpu_mb_solve_constraints( } 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); @@ -123,8 +121,7 @@ pub fn gpu_mb_solve_constraints( // Generalized `J·v` for `J = e_{dof_id} - coupling_coeff*e_{dof2_id}` // (coupling rows); collapses to `v[dof_id]` for limit / motor rows // (their `coupling_coeff` is 0). - let v_d = dof_v[cons.dof_id as usize] - - cons.coupling_coeff * dof_v[cons.dof2_id as usize]; + let v_d = dof_v[cons.dof_id as usize] - cons.coupling_coeff * dof_v[cons.dof2_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; @@ -152,7 +149,6 @@ pub fn gpu_mb_solve_constraints( workgroup_memory_barrier_with_group_sync(); } - // Contacts. In 3D the two friction rows of a contact point are solved // together so their impulse can be capped to the friction cone; the second // row is handled by its sibling and skipped here. @@ -234,7 +230,11 @@ pub fn gpu_mb_solve_constraints( 0.0 }; let raw1 = if has_pair { - let rhs1 = if use_bias { cons2.rhs } else { cons2.rhs_wo_bias }; + let rhs1 = if use_bias { + cons2.rhs + } else { + cons2.rhs_wo_bias + }; cfm_factor * (impulse1 - cons2.inv_lhs * (j_dot_v1 + rhs1)) } else { 0.0 @@ -318,7 +318,7 @@ pub fn gpu_mb_solve_joints( #[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], + #[spirv(workgroup)] dof_v: &mut [f32; MAX_MB_DOFS], ) { let batch_id = workgroup_id.y; let mb_idx = workgroup_id.x; @@ -338,8 +338,7 @@ pub fn gpu_mb_solve_joints( 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; + 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; @@ -362,8 +361,7 @@ pub fn gpu_mb_solve_joints( // Generalized `J·v` for `J = e_{dof_id} - coupling_coeff*e_{dof2_id}` // (coupling rows); collapses to `v[dof_id]` for limit / motor rows // (their `coupling_coeff` is 0). - let v_d = dof_v[cons.dof_id as usize] - - cons.coupling_coeff * dof_v[cons.dof2_id as usize]; + let v_d = dof_v[cons.dof_id as usize] - cons.coupling_coeff * dof_v[cons.dof2_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; @@ -433,8 +431,8 @@ pub fn gpu_mb_build_contact_delassus( return; } - let cons_base = batch_ids.mb_contact_constraints_start(batch_id) - + (mb_idx as usize) * (MAXC as usize); + 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; @@ -493,7 +491,7 @@ pub fn gpu_mb_solve_contacts_delassus( #[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)] dof_v: &mut [f32; MAX_MB_DOFS], #[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], @@ -522,8 +520,8 @@ pub fn gpu_mb_solve_contacts_delassus( 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); + 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; @@ -547,8 +545,8 @@ pub fn gpu_mb_solve_contacts_delassus( cfm_shared[s as usize] = if use_bias { cons.cfm_factor } else { 1.0 }; 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); + 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 }); @@ -597,7 +595,8 @@ pub fn gpu_mb_solve_contacts_delassus( let impulse0 = imp_shared[s as usize]; let raw0 = cfm_shared[s as usize] - * (impulse0 - inv_lhs_shared[s as usize] * (a_shared[s as usize] + rhs_shared[s as usize])); + * (impulse0 + - inv_lhs_shared[s as usize] * (a_shared[s as usize] + rhs_shared[s as usize])); let impulse1 = if has_pair { imp_shared[(s + 1) as usize] } else { @@ -632,8 +631,7 @@ pub fn gpu_mb_solve_contacts_delassus( 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); + let mut free = solver_vels.read(colliders_start + cons.free_body_id as usize); free.linear += cons.lin_jac * (cons.free_body_im * delta0); free.angular += cons.ii_ang_jac * delta0; if has_pair { diff --git a/src_rbd_shaders/dynamics/multibody/ws_soa.rs b/src_rbd_shaders/dynamics/multibody/ws_soa.rs index 796470c..3e05bb5 100644 --- a/src_rbd_shaders/dynamics/multibody/ws_soa.rs +++ b/src_rbd_shaders/dynamics/multibody/ws_soa.rs @@ -179,11 +179,7 @@ pub fn ws_set_vec(buf: &mut [Vec4], a: WsAddr, k: u32, f: u32, v: Vector) { pub fn ws_ext_wrench(buf: &[Vec4], a: WsAddr, k: u32) -> (Vector, crate::AngVector, f32) { let f = buf.read(a.at(k, WS_EXT_FORCE)); let t = buf.read(a.at(k, WS_EXT_TORQUE)); - ( - Vec3::new(f.x, f.y, f.z), - Vec3::new(t.x, t.y, t.z), - f.w, - ) + (Vec3::new(f.x, f.y, f.z), Vec3::new(t.x, t.y, t.z), f.w) } #[cfg(feature = "dim2")] @@ -310,7 +306,10 @@ pub fn ws_set_vel(buf: &mut [Vec4], a: WsAddr, k: u32, f: u32, v: Velocity) { #[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)); + 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, diff --git a/src_rbd_shaders/dynamics/solver.rs b/src_rbd_shaders/dynamics/solver.rs index 0862dd5..327ddaf 100644 --- a/src_rbd_shaders/dynamics/solver.rs +++ b/src_rbd_shaders/dynamics/solver.rs @@ -617,7 +617,15 @@ pub fn gpu_integrate_linearized( let max_ang = params.max_angular_velocity(); #[cfg(feature = "dim2")] if vels.angular.abs() > max_ang { - vels.angular = vels.angular.signum() * max_ang; + // Explicit sign select rather than `signum`: `f32::signum` compiles to + // a comparison against a NaN constant, and naga rejects a NaN literal + // outright, so the whole module fails to translate at pipeline + // creation. The guard above rules out zero, so the two cases suffice. + vels.angular = if vels.angular > 0.0 { + max_ang + } else { + -max_ang + }; } #[cfg(feature = "dim3")] { @@ -695,4 +703,3 @@ pub fn gpu_solver_finalize( body_poses[idx] = solver_body_poses[idx].prepend_translation(-local_mprops[idx].com); } } - diff --git a/src_rbd_shaders/dynamics/solver_utils.rs b/src_rbd_shaders/dynamics/solver_utils.rs index 974379f..869a6b7 100644 --- a/src_rbd_shaders/dynamics/solver_utils.rs +++ b/src_rbd_shaders/dynamics/solver_utils.rs @@ -560,8 +560,7 @@ impl TwoBodyConstraint { } else { (c.rhs_wo_bias, 1.0) }; - let dvel = dir_a.dot(solver_vel1.linear) - + gdot(c.torque_dir_a, solver_vel1.angular) + 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) + rhs; @@ -660,4 +659,3 @@ impl TwoBodyConstraint { } } } - diff --git a/src_rbd_shaders/dynamics/warmstart.rs b/src_rbd_shaders/dynamics/warmstart.rs index 512eaf4..b5be104 100644 --- a/src_rbd_shaders/dynamics/warmstart.rs +++ b/src_rbd_shaders/dynamics/warmstart.rs @@ -72,7 +72,8 @@ pub fn gpu_seed_colors_from_warmstart( #[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 = 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], @@ -256,8 +257,7 @@ pub fn transfer_warmstart_impulses( let old_c = &old_constraints[cid_old]; let old_t0 = old_c.tangent_a; let old_t1 = old_c.dir_a.cross(old_t0); - let old_impulse = - old_c.elements.at(k_old).tangent_part.impulse; + let old_impulse = old_c.elements.at(k_old).tangent_part.impulse; let world = old_t0 * old_impulse.x + old_t1 * old_impulse.y; let new_t0 = new_constraints[i].tangent_a; diff --git a/src_rbd_shaders/queries/polygonal_feature.rs b/src_rbd_shaders/queries/polygonal_feature.rs index 68bcd11..ddec14e 100644 --- a/src_rbd_shaders/queries/polygonal_feature.rs +++ b/src_rbd_shaders/queries/polygonal_feature.rs @@ -622,9 +622,7 @@ mod dim3 { for i in 0..num { let d = candidates.at(i).pt - selected_a; let dist = d.dot(d); - if i != selected.read(0) - && candidates.at(i).dist <= prediction - && dist > furthest_dist + if i != selected.read(0) && candidates.at(i).dist <= prediction && dist > furthest_dist { furthest_dist = dist; selected.write(1, i); @@ -648,9 +646,7 @@ mod dim3 { let mut min_dot = MAX_FLT; let mut max_dot = -MAX_FLT; for i in 0..num { - if i == selected.read(0) - || i == selected.read(1) - || candidates.at(i).dist > prediction + if i == selected.read(0) || i == selected.read(1) || candidates.at(i).dist > prediction { continue; } @@ -889,7 +885,12 @@ mod dim3 { } if num_candidates as usize == MAX_CANDIDATE_POINTS { - return manifold_reduction(&candidates, num_candidates, sep_axis1, prediction); + return manifold_reduction( + &candidates, + num_candidates, + sep_axis1, + prediction, + ); } } } diff --git a/src_rbd_shaders/tests/linalg.rs b/src_rbd_shaders/tests/linalg.rs index b9921f0..b16d199 100644 --- a/src_rbd_shaders/tests/linalg.rs +++ b/src_rbd_shaders/tests/linalg.rs @@ -12,7 +12,7 @@ //! with multiple right-hand sides. use crate::utils::linalg::{ - MatSlice, axpy_mat, copy_from, fill, gemm, gemm_mat3_lhs, gemm_tr, gemv_tr_spatial, + MatSlice, VSlice, axpy_mat, copy_from, fill, gemm, gemm_mat3_lhs, gemm_tr, gemv_tr_spatial, lu_decompose, lu_solve_in_place, quadform_spatial, skew, skew_tr, }; use glamx::{Mat3, Vec3}; @@ -494,11 +494,18 @@ fn lu_solve_identity_recovers_rhs() { buf_m[m.idx(i, i)] = 1.0; } let mut pivots = vec![0u32; n as usize]; - lu_decompose(&mut buf_m, m, &mut pivots, 0); + lu_decompose(&mut buf_m, m, &mut pivots, VSlice::dense(0)); let mut rhs = vec![7.0, -3.0, 2.5, 1.25]; let want = rhs.clone(); - lu_solve_in_place(&buf_m, m, &pivots, 0, &mut rhs, 0); + lu_solve_in_place( + &buf_m, + m, + &pivots, + VSlice::dense(0), + &mut rhs, + VSlice::dense(0), + ); assert_slice_eq(&rhs, &want); } @@ -516,11 +523,18 @@ fn lu_solve_spd_matrix() { // Decompose. let mut pivots = vec![0u32; n as usize]; - lu_decompose(&mut buf, m, &mut pivots, 0); + lu_decompose(&mut buf, m, &mut pivots, VSlice::dense(0)); // Solve M · x = b with b = [1, 2, 3]. let mut rhs = vec![1.0f32, 2.0, 3.0]; - lu_solve_in_place(&buf, m, &pivots, 0, &mut rhs, 0); + lu_solve_in_place( + &buf, + m, + &pivots, + VSlice::dense(0), + &mut rhs, + VSlice::dense(0), + ); // Verify by multiplying back: M · x ≈ b. let mx = matvec(rows, &rhs); @@ -541,14 +555,21 @@ fn lu_solve_requires_pivoting() { pack_matrix(&mut buf, m, rows); let mut pivots = vec![0u32; n as usize]; - lu_decompose(&mut buf, m, &mut pivots, 0); + lu_decompose(&mut buf, m, &mut pivots, VSlice::dense(0)); // Pick a known solution, compute b = M·x, then confirm the solve recovers x. let x_true = [1.5f32, -0.5, 2.0]; let b = matvec(rows, &x_true); let mut rhs = b.clone(); - lu_solve_in_place(&buf, m, &pivots, 0, &mut rhs, 0); + lu_solve_in_place( + &buf, + m, + &pivots, + VSlice::dense(0), + &mut rhs, + VSlice::dense(0), + ); assert_slice_eq(&rhs, &x_true); } @@ -569,7 +590,7 @@ fn lu_factor_reused_across_multiple_rhs() { // Decompose once. let mut pivots = vec![0u32; n as usize]; - lu_decompose(&mut buf, m, &mut pivots, 0); + lu_decompose(&mut buf, m, &mut pivots, VSlice::dense(0)); // Solve three different RHSes with the same factorization. let rhss = [ @@ -579,7 +600,14 @@ fn lu_factor_reused_across_multiple_rhs() { ]; for b in &rhss { let mut rhs = b.to_vec(); - lu_solve_in_place(&buf, m, &pivots, 0, &mut rhs, 0); + lu_solve_in_place( + &buf, + m, + &pivots, + VSlice::dense(0), + &mut rhs, + VSlice::dense(0), + ); // Verify: M · rhs ≈ b. let mx = matvec(rows, &rhs); assert_slice_eq(&mx, b); @@ -602,8 +630,8 @@ fn lu_solve_respects_offsets() { pack_matrix(&mut buf_m, m_b, rows_b); let mut pivots = vec![0u32; 4]; - lu_decompose(&mut buf_m, m_a, &mut pivots, 0); - lu_decompose(&mut buf_m, m_b, &mut pivots, 2); + lu_decompose(&mut buf_m, m_a, &mut pivots, VSlice::dense(0)); + lu_decompose(&mut buf_m, m_b, &mut pivots, VSlice::dense(2)); let x_a_true = [1.0f32, -1.0]; let x_b_true = [2.0f32, 0.5]; @@ -611,8 +639,22 @@ fn lu_solve_respects_offsets() { let b_b = matvec(rows_b, &x_b_true); let mut buf_rhs = vec![b_a[0], b_a[1], b_b[0], b_b[1]]; - lu_solve_in_place(&buf_m, m_a, &pivots, 0, &mut buf_rhs, 0); - lu_solve_in_place(&buf_m, m_b, &pivots, 2, &mut buf_rhs, 2); + lu_solve_in_place( + &buf_m, + m_a, + &pivots, + VSlice::dense(0), + &mut buf_rhs, + VSlice::dense(0), + ); + lu_solve_in_place( + &buf_m, + m_b, + &pivots, + VSlice::dense(2), + &mut buf_rhs, + VSlice::dense(2), + ); assert_slice_eq(&buf_rhs[0..2], &x_a_true); assert_slice_eq(&buf_rhs[2..4], &x_b_true); diff --git a/src_rbd_shaders/utils/linalg.rs b/src_rbd_shaders/utils/linalg.rs index ef8133f..209674f 100644 --- a/src_rbd_shaders/utils/linalg.rs +++ b/src_rbd_shaders/utils/linalg.rs @@ -716,7 +716,7 @@ pub fn gemm_tr( /// strictly-below-diagonal entries hold `L` (with implicit unit diagonal), the /// diagonal and above hold `U`. Row pivots are written to `pivots[0..n]` — /// `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`. +/// `k`. `piv` is the view of this multibody's pivot slot inside `buf_pivots`. #[inline] pub fn lu_decompose(buf_m: &mut [f32], m: MatSlice, buf_pivots: &mut [u32], piv: VSlice) { let n = m.rows; @@ -770,8 +770,8 @@ pub fn lu_decompose(buf_m: &mut [f32], m: MatSlice, buf_pivots: &mut [u32], piv: /// Solve `M · x = rhs` in-place, using LU factors produced by /// [`lu_decompose`] (and its pivot array). The result overwrites `rhs`. /// -/// `m` and `pivots` must be the exact outputs of a previous `lu_decompose` call. -/// `rhs` is an `n`-element column vector; `rhs_offset` is where it starts in `buf_rhs`. +/// `m` and `piv` must designate the exact outputs of a previous `lu_decompose` +/// call. `rhs` is the view of an `n`-element column vector inside `buf_rhs`. #[inline] pub fn lu_solve_in_place( buf_m: &[f32],