diff --git a/Cargo.toml b/Cargo.toml index d2ced59a..ba3757b0 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -42,10 +42,10 @@ 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.35.2", default-features = false } -rapier3d = { version = "0.35.2", default-features = false } -rapier3d-urdf = "0.35" -rapier3d-mjcf = { version = "0.35", features = ["stl", "wavefront", "msh"] } +rapier2d = { version = "0.35.3", default-features = false } +rapier3d = { version = "0.35.3", default-features = false } +rapier3d-urdf = "0.35.3" +rapier3d-mjcf = { version = "0.35.3", features = ["stl", "wavefront", "msh"] } parry2d = { version = "0.30", default-features = false } parry3d = { version = "0.30", default-features = false } @@ -93,8 +93,6 @@ 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" } diff --git a/crates/examples3d/rbd_joint_ball3.rs b/crates/examples3d/rbd_joint_ball3.rs index ee088dcd..3c849249 100644 --- a/crates/examples3d/rbd_joint_ball3.rs +++ b/crates/examples3d/rbd_joint_ball3.rs @@ -42,7 +42,7 @@ pub async fn run( let collider = if status == RigidBodyType::Fixed { ColliderBuilder::cuboid(rad, rad, rad).build() } else { - ColliderBuilder::ball(rad).density(10.0).build() + ColliderBuilder::ball(rad).density(40.0).build() }; let shape = collider.shared_shape().clone(); let child_handle = state.insert_rigid_body(rigid_body, collider, no_coupling); diff --git a/crates/nexus_python3d/src/loaders.rs b/crates/nexus_python3d/src/loaders.rs index 58447299..b19602d9 100644 --- a/crates/nexus_python3d/src/loaders.rs +++ b/crates/nexus_python3d/src/loaders.rs @@ -115,12 +115,19 @@ struct VisualMeshReg { /// floor, camera and light with the viewer. Mirrors the Rust `mujoco_menagerie3` /// example's `load_scene` (minus the runtime model picker). Gravity is left to /// the caller (set after `finalize`). +/// Handles of the robot loaded by [`insert_mjcf`], kept by the Python +/// `NexusState` so per-step actuator control (`apply_actuator_controls`) can +/// reuse `rapier3d-mjcf`'s MJCF actuator semantics. +pub type MjcfHandles = + rapier3d_mjcf::MjcfRobotHandles>; + pub fn insert_mjcf( state: &mut nexus3d::prelude::NexusState, mut viewer: PyRefMut, scene_path: &std::path::Path, render_colliders: bool, -) -> PyResult { + env: usize, +) -> PyResult<(MjcfSceneInfo, Option)> { use nexus3d::prelude::RbdCoupling; use pyo3::exceptions::PyRuntimeError; use rapier3d::parry::bounding_volume::BoundingVolume; // for `Aabb::merge` @@ -141,9 +148,9 @@ pub fn insert_mjcf( let mut floor: Option<(glamx::Vec3, glamx::Vec3)> = None; let mut camera: Option<(glamx::Vec3, glamx::Vec3)> = None; - match MjcfRobot::from_file(scene_path, options) { + let robot_handles: Option = match MjcfRobot::from_file(scene_path, options) { Ok((robot, _model)) => { - let world = state.rbd_world_mut(0); + let world = state.rbd_world_mut(env); let handles = robot.clone().insert_using_multibody_joints( &mut world.bodies, &mut world.colliders, @@ -225,6 +232,7 @@ pub fn insert_mjcf( let eye = target + glamx::Vec3::new(radius * 2.2, -radius * 2.2, radius * 1.6); camera = Some((eye, target)); } + Some(handles) } Err(e) => { return Err(PyRuntimeError::new_err(format!( @@ -232,7 +240,7 @@ pub fn insert_mjcf( scene_path.display() ))); } - } + }; let loaded = camera.is_some(); let v = viewer.rust_mut(); @@ -242,11 +250,15 @@ pub fn insert_mjcf( let body = rp::RigidBodyBuilder::fixed().translation(center).build(); let collider = rp::ColliderBuilder::cuboid(he.x, he.y, he.z).build(); let shape = collider.shared_shape().clone(); - let handle = state.insert_rigid_body(body, collider, RbdCoupling::None); - v.insert_shape(handle, &shape, rp::Pose::IDENTITY); + let handle = state.insert_rigid_body_in(env, body, collider, RbdCoupling::None); + if env == 0 { + v.insert_shape(handle, &shape, rp::Pose::IDENTITY); + } } - if render_colliders { + if env != 0 { + // Batch environments are physics-only: the viewer draws environment 0. + } else if render_colliders { for (body, shape, local_pose, _) in &collider_shapes { v.insert_visual_shape(0, *body, shape, *local_pose); } @@ -271,11 +283,13 @@ pub fn insert_mjcf( } } - if let Some((eye, target)) = camera { - v.set_camera(eye, target); + if env == 0 { + if let Some((eye, target)) = camera { + v.set_camera(eye, target); + } + v.scene3d_mut() + .add_directional_light(glamx::Vec3::new(-1.0, 1.0, -1.0)); } - v.scene3d_mut() - .add_directional_light(glamx::Vec3::new(-1.0, 1.0, -1.0)); - Ok(MjcfSceneInfo { z_up: true, loaded }) + Ok((MjcfSceneInfo { z_up: true, loaded }, robot_handles)) } diff --git a/crates/nexus_python3d/src/nexus.rs b/crates/nexus_python3d/src/nexus.rs index ef9648d3..b0bb9e60 100644 --- a/crates/nexus_python3d/src/nexus.rs +++ b/crates/nexus_python3d/src/nexus.rs @@ -15,6 +15,7 @@ use nexus3d::prelude::{ NexusPipeline as RNexusPipeline, NexusPipelineMask, NexusState as RNexusState, RbdCoupling as RRbdCoupling, }; +use numpy::PyArray2; use pyo3::exceptions::PyRuntimeError; use pyo3::prelude::*; use rapier3d::prelude as rp; @@ -95,15 +96,17 @@ impl GpuTimestamps { } /// The GPU-resident state of a multiphysics simulation -/// (`nexus3d::prelude::NexusState`). +/// (`nexus3d::prelude::NexusState`). The second field keeps the +/// `rapier3d-mjcf` robot handles of the last `insert_mjcf`, so +/// `apply_actuator_controls` can drive the robot's actuators per step. #[pyclass(name = "NexusState", unsendable)] -pub struct NexusState(pub RNexusState); +pub struct NexusState(pub RNexusState, pub Option); #[pymethods] impl NexusState { #[new] fn new() -> Self { - NexusState(RNexusState::default()) + NexusState(RNexusState::default(), None) } // --- rigid bodies ----------------------------------------------------- @@ -322,17 +325,145 @@ impl NexusState { }) } - /// Loads a MuJoCo MJCF scene into environment 0 as multibodies, registering - /// its render shapes (and a sized floor) with `viewer`. Returns scene info - /// (suggested camera + whether the scene is Z-up). Call `finalize` after. - #[pyo3(signature = (viewer, scene_path, render_colliders=false))] + /// Per-environment collision-pair capacity (default 4096). Lower it before + /// `finalize` when batching many small environments: pair-keyed GPU + /// workspaces scale with `capacity x num_envs`. + fn set_rbd_collisions_capacity(&mut self, capacity: u32) { + self.0.set_rbd_collisions_capacity(capacity); + } + + /// Loads a MuJoCo MJCF scene into environment `env` as multibodies, + /// registering its render shapes (and a sized floor) with `viewer`. Returns + /// scene info (suggested camera + whether the scene is Z-up). Call + /// `finalize` after. + #[pyo3(signature = (viewer, scene_path, render_colliders=false, env=0))] fn insert_mjcf( &mut self, viewer: PyRefMut, scene_path: std::path::PathBuf, render_colliders: bool, + env: usize, ) -> PyResult { - crate::loaders::insert_mjcf(&mut self.0, viewer, &scene_path, render_colliders) + let (info, handles) = + crate::loaders::insert_mjcf(&mut self.0, viewer, &scene_path, render_colliders, env)?; + self.1 = handles; + Ok(info) + } + + // --- MJCF actuation ----------------------------------------------------- + + /// Names of the MJCF ``s of the robot loaded by `insert_mjcf`, in + /// actuator (control-vector) order. Unnamed actuators fall back to the name + /// of the joint they drive. Empty before `insert_mjcf`. + fn actuator_names(&self) -> Vec { + self.1 + .as_ref() + .map(|h| { + h.actuators + .iter() + .map(|a| { + a.actuator + .name + .clone() + .or_else(|| a.actuator.joint.clone()) + .unwrap_or_default() + }) + .collect() + }) + .unwrap_or_default() + } + + /// Applies one MJCF control vector (one entry per actuator, in + /// `actuator_names` order) to every environment's copy of the robot loaded + /// by `insert_mjcf`, with full MJCF actuator semantics (`` servos + /// with kp/kv, `` force/gear, force limits), and pushes the resulting + /// joint-motor state to the GPU. + /// + /// Call once per control step, after `finalize`; the next + /// `NexusPipeline.simulate` steps the solver against the new targets. + #[pyo3(signature = (viewer, ctrl))] + fn apply_actuator_controls( + &mut self, + viewer: PyRef, + ctrl: Vec, + ) -> PyResult<()> { + let Some(handles) = self.1.as_ref() else { + return Err(PyRuntimeError::new_err( + "no MJCF robot loaded (call insert_mjcf first)", + )); + }; + if ctrl.len() != handles.actuators.len() { + return Err(PyRuntimeError::new_err(format!( + "ctrl has {} entries but the robot has {} actuators", + ctrl.len(), + handles.actuators.len() + ))); + } + let handles = handles.clone(); + self.0 + .control_multibody_motors(viewer.backend(), |_, world| { + handles.apply_controls_multibody( + &mut world.bodies, + &mut world.multibody_joints, + &ctrl, + ); + }) + .map_err(gpu_err) + } + + /// Reads every environment's multibody link states back from the GPU in one + /// transfer. Returns five float32 numpy arrays with + /// `num_environments * multibody_links_per_env` rows, environment-major; + /// links follow the GPU build's traversal order (multibodies, then links, + /// parent before child), the same order `apply_actuator_controls` drives: + /// + /// - `coords (n, 6)`: generalized joint coordinates (only the joint's DOF + /// count is meaningful; a revolute joint's angle is `coords[5]`), + /// - `positions (n, 3)` / `quats (n, 4)`: link world pose (`w, x, y, z`), + /// - `linvels (n, 3)` / `angvels (n, 3)`: world-space velocities, valid + /// after the first simulated step. + /// + /// Use `multibody_links_per_env()` to slice a single environment out. + #[allow(clippy::type_complexity)] + fn read_multibody_links<'py>( + &self, + py: Python<'py>, + viewer: PyRef, + ) -> ( + Bound<'py, PyArray2>, + Bound<'py, PyArray2>, + Bound<'py, PyArray2>, + Bound<'py, PyArray2>, + Bound<'py, PyArray2>, + ) { + let links = pollster::block_on(self.0.read_multibody_links(viewer.backend())); + let mut coords = Vec::with_capacity(links.len()); + let mut positions = Vec::with_capacity(links.len()); + let mut quats = Vec::with_capacity(links.len()); + let mut linvels = Vec::with_capacity(links.len()); + let mut angvels = Vec::with_capacity(links.len()); + for ws in &links { + coords.push(ws.coords.to_vec()); + let (t, q) = (ws.local_to_world.translation, ws.local_to_world.rotation); + positions.push(vec![t.x, t.y, t.z]); + quats.push(vec![q.w, q.x, q.y, q.z]); + let (l, a) = (ws.rb_vels.linear, ws.rb_vels.angular); + linvels.push(vec![l.x, l.y, l.z]); + angvels.push(vec![a.x, a.y, a.z]); + } + ( + PyArray2::from_vec2(py, &coords).unwrap(), + PyArray2::from_vec2(py, &positions).unwrap(), + PyArray2::from_vec2(py, &quats).unwrap(), + PyArray2::from_vec2(py, &linvels).unwrap(), + PyArray2::from_vec2(py, &angvels).unwrap(), + ) + } + + /// Number of link slots per environment, the stride of + /// `read_multibody_links`. + fn multibody_links_per_env(&self) -> u32 { + self.0.multibody_links_per_env() } // --- rbd config ------------------------------------------------------- diff --git a/src/pipeline.rs b/src/pipeline.rs index a330e721..e79b5424 100644 --- a/src/pipeline.rs +++ b/src/pipeline.rs @@ -1,3 +1,4 @@ +#[cfg(feature = "mpm")] use crate::mpm::pipeline::MpmPipeline; use crate::rbd::pipeline::RbdPipeline; use crate::state::NexusState; @@ -15,6 +16,7 @@ bitflags::bitflags! { #[derive(Default)] pub struct NexusPipeline { pub rbd_pipeline: Option, + #[cfg(feature = "mpm")] pub mpm_pipeline: Option, } @@ -27,6 +29,7 @@ impl NexusPipeline { if pipelines.contains(NexusPipelineMask::RBD) && self.rbd_pipeline.is_none() { self.rbd_pipeline = Some(RbdPipeline::new(backend)?); } + #[cfg(feature = "mpm")] if pipelines.contains(NexusPipelineMask::MPM) && self.mpm_pipeline.is_none() { self.mpm_pipeline = Some(MpmPipeline::new(backend)?); } @@ -73,6 +76,7 @@ impl NexusPipeline { } // MPM pipeline + #[cfg(feature = "mpm")] if let Some(mpm) = state.mpm.as_mut() { self.preload_pipelines(backend, NexusPipelineMask::MPM)?; let pipeline = self.mpm_pipeline.as_mut().unwrap_or_else(|| unreachable!()); @@ -91,6 +95,7 @@ impl NexusPipeline { // bodies as static. Push that copy back so rendering and the next // step's broad phase see a boundary that actually moved. // FIXME: the RBD pipeline should remain in charge of moving the bodies. + #[cfg(feature = "mpm")] if let (Some(rbd), Some(mpm)) = (state.rbd.as_mut(), state.mpm.as_ref()) { let pipeline = self.mpm_pipeline.as_ref().unwrap_or_else(|| unreachable!()); pipeline.writeback_body_poses(backend, mpm, rbd.body_poses_mut())?; diff --git a/src/state.rs b/src/state.rs index 76ab81db..94d2b474 100644 --- a/src/state.rs +++ b/src/state.rs @@ -1,4 +1,6 @@ +#[cfg(feature = "mpm")] use crate::mpm::pipeline::{MpmCapacities, MpmState}; +#[cfg(feature = "mpm")] use crate::mpm::solver::{BoundaryCondition, Particle, SimulationParams}; use crate::rapier::data::{Arena, Coarena, Index}; use crate::rapier::prelude::{ @@ -10,6 +12,8 @@ use crate::rbd::dynamics::{ body::{BodyCoupling, RapierBodyCouplingEntry}, }; use crate::rbd::pipeline::{RbdCapacities, RbdResizePolicy, RbdState, RunStats}; +#[cfg(feature = "dim3")] +use khal::backend::Backend; use khal::backend::{GpuBackend, GpuBackendError}; /// Handle referencing a rigid-body managed by a [`NexusState`]. @@ -28,7 +32,9 @@ pub struct NexusParticleChunk(Index); #[derive(Copy, Clone, PartialEq, Debug)] pub enum RbdCoupling { None, + #[cfg(feature = "mpm")] MpmOneWay(BoundaryCondition), + #[cfg(feature = "mpm")] MpmTwoWay(BoundaryCondition), } @@ -38,6 +44,7 @@ pub struct NexusCapacities { /// Rigid-body solver capacities. pub rbd: RbdCapacities, /// MPM solver capacities. + #[cfg(feature = "mpm")] pub mpm: MpmCapacities, } @@ -57,6 +64,7 @@ impl NexusCapacities { self } + #[cfg(feature = "mpm")] pub fn mpm_grid_size(mut self, num_chunks: u32) -> Self { self.mpm.grid_size = num_chunks; self @@ -67,6 +75,7 @@ impl NexusCapacities { self } + #[cfg(feature = "mpm")] pub fn mpm_particles(mut self, capacity: u32) -> Self { self.mpm.particles_capacity = capacity; self @@ -114,6 +123,7 @@ pub struct NexusState { pub rbd: Option, /// MPM sub-state, allocated on the first [`Self::add_particles`] (or the /// first coupled rigid-body insertion). + #[cfg(feature = "mpm")] pub mpm: Option, pub run_stats: RunStats, @@ -124,6 +134,7 @@ pub struct NexusState { /// Live particle count per MPM chunk (the arena key is the public /// [`NexusParticleChunk`] handle). + #[cfg(feature = "mpm")] mpm_chunks: Arena, /// Owning chunk for each GPU particle slot, kept in sync under the /// swap-removal performed by [`Self::remove_chunk`] / @@ -131,16 +142,21 @@ pub struct NexusState { slot2chunk: Vec, /// MPM simulation params / grid cell width requested before the MPM /// sub-state is lazily created. + #[cfg(feature = "mpm")] mpm_params: Option, + #[cfg(feature = "mpm")] mpm_cell_width: f32, /// Number of MPM substeps run per [`NexusPipeline::simulate`](crate::pipeline::NexusPipeline::simulate) call. + #[cfg(feature = "mpm")] pub mpm_substeps: u32, /// Desired CPIC rigid-coupling flag, kept here so it survives until the MPM /// sub-state is lazily created (and is what [`Self::mpm_use_cpic`] reports /// meanwhile). + #[cfg(feature = "mpm")] mpm_use_cpic: bool, /// Set when particles or MPM-coupled bodies change; consumed by /// [`Self::finalize`] to rebuild the MPM↔rapier coupling. + #[cfg(feature = "mpm")] mpm_dirty: bool, // Initial capacities used to allocate the states lazily. @@ -179,6 +195,7 @@ impl NexusState { pub fn new(capacities: NexusCapacities) -> Self { Self { rbd: None, + #[cfg(feature = "mpm")] mpm: None, run_stats: RunStats::default(), rbd_envs: vec![PhysicsWorld::default()], @@ -187,12 +204,18 @@ impl NexusState { rbd_steps_per_frame: 1, rbd_reserve_per_env: 0, rbd2gpu: vec![Coarena::new()], + #[cfg(feature = "mpm")] mpm_chunks: Arena::new(), slot2chunk: Vec::new(), + #[cfg(feature = "mpm")] mpm_params: None, + #[cfg(feature = "mpm")] mpm_cell_width: 1.0, + #[cfg(feature = "mpm")] mpm_substeps: 20, + #[cfg(feature = "mpm")] mpm_use_cpic: true, + #[cfg(feature = "mpm")] mpm_dirty: false, capacities, } @@ -211,6 +234,7 @@ impl NexusState { /// width. Call before the first [`Self::add_particles`]; the values are /// applied when the MPM sub-state is created. If MPM already exists they are /// applied immediately (the grid is reset, so prefer calling this first). + #[cfg(feature = "mpm")] pub fn set_mpm_params( &mut self, backend: &GpuBackend, @@ -228,11 +252,13 @@ impl NexusState { /// Sets the number of MPM substeps run per [`NexusPipeline::simulate`](crate::pipeline::NexusPipeline::simulate) call (default /// 20). More substeps → smaller timestep → more stable but slower. + #[cfg(feature = "mpm")] pub fn set_mpm_substeps(&mut self, substeps: u32) { self.mpm_substeps = substeps.max(1); } /// Number of MPM substeps run per [`NexusPipeline::simulate`](crate::pipeline::NexusPipeline::simulate) call. + #[cfg(feature = "mpm")] pub fn mpm_substeps(&self) -> u32 { self.mpm_substeps } @@ -240,6 +266,7 @@ impl NexusState { /// Enables/disables CPIC (compatible particle-in-cell) rigid coupling. The /// preference is stored so it survives until MPM is lazily allocated. Not /// overwritten by [`Self::finalize`] unless the coupling set changes. + #[cfg(feature = "mpm")] pub fn set_mpm_use_cpic(&mut self, enabled: bool) { self.mpm_use_cpic = enabled; if let Some(mpm) = self.mpm.as_mut() { @@ -249,6 +276,7 @@ impl NexusState { /// Whether CPIC rigid coupling is enabled. Falls back to the stored /// preference before MPM is lazily allocated. + #[cfg(feature = "mpm")] pub fn mpm_use_cpic(&self) -> bool { self.mpm .as_ref() @@ -260,12 +288,14 @@ impl NexusState { /// via [`Self::set_mpm_params`], even before the sub-state is lazily /// allocated on the first [`Self::add_particles`], so a particle emitter /// that starts empty still reports its MPM usage. + #[cfg(feature = "mpm")] pub fn has_mpm(&self) -> bool { self.mpm.is_some() || self.mpm_params.is_some() } /// Sets the MPM gravity vector. Applied on the next [`NexusPipeline::simulate`](crate::pipeline::NexusPipeline::simulate) (the /// per-substep params are re-uploaded each frame), so this is cheap. + #[cfg(feature = "mpm")] pub fn set_mpm_gravity(&mut self, gravity: crate::rbd::math::Vector) { // Keep the stored params authoritative so the gravity survives until MPM // is lazily allocated (and is what `mpm_gravity` reports meanwhile). @@ -280,6 +310,7 @@ impl NexusState { /// Current MPM gravity vector. Falls back to the gravity configured via /// [`Self::set_mpm_params`] before MPM is lazily allocated, and only to zero /// if no params were ever set. + #[cfg(feature = "mpm")] pub fn mpm_gravity(&self) -> crate::rbd::math::Vector { self.mpm .as_ref() @@ -302,6 +333,15 @@ impl NexusState { // ── Rigid-body runtime settings ───────────────────────────────────── + /// Overrides the per-environment collision-pair capacity used when the GPU + /// rigid-body state is (re)allocated at `finalize`. The default (4096) is + /// sized for one busy scene, not thousands of small batched envs: the + /// pair-keyed workspaces scale as `capacity x num_envs x sizeof(manifold)`, + /// which binds ~9 GiB at 2048 envs unless lowered. + pub fn set_rbd_collisions_capacity(&mut self, capacity: u32) { + self.capacities.rbd.collisions_capacity = capacity.max(1); + } + /// Sets the number of rigid-body solver steps advanced per /// [`NexusPipeline::simulate`](crate::pipeline::NexusPipeline::simulate) call (default 1). Acts as a simulation-speed control. pub fn set_rbd_steps_per_frame(&mut self, steps: u32) { @@ -334,6 +374,7 @@ impl NexusState { c.collision_pairs = rbd.collision_pairs_len() as usize; c.collision_pairs_capacity = rbd.collision_pairs_capacity() as usize; } + #[cfg(feature = "mpm")] if let Some(mpm) = self.mpm.as_ref() { c.particles = mpm.particles.len(); } @@ -343,6 +384,7 @@ impl NexusState { /// Returns a mutable reference to the MPM sub-state, allocating an empty one /// (sized from the stored capacities, configured from [`Self::set_mpm_params`]) /// if it doesn’t exist yet. + #[cfg(feature = "mpm")] fn mpm_or_insert(&mut self, backend: &GpuBackend) -> Result<&mut MpmState, GpuBackendError> { if self.mpm.is_none() { let grid_capacity = self.capacities.mpm.grid_size; @@ -401,6 +443,89 @@ impl NexusState { &mut self.rbd_envs[env] } + /// Runtime actuation entry point: mutates every environment's rapier + /// multibody joints through `f`, then pushes the refreshed joint data + /// (motor targets and gains, limits) to the GPU. + /// + /// `f` receives each environment in turn, so one control vector can drive + /// the whole batch or `f` can pick per-environment targets. Unlike + /// [`Self::rbd_world_mut`] this does not mark the world dirty: motor + /// updates are per-step control, not a topology change, so they trigger no + /// GPU rebuild. Call after [`Self::finalize`]; a no-op before it. + #[cfg(feature = "dim3")] + pub fn control_multibody_motors( + &mut self, + backend: &GpuBackend, + mut f: F, + ) -> Result<(), GpuBackendError> + where + F: FnMut(usize, &mut PhysicsWorld), + { + for (env, world) in self.rbd_envs.iter_mut().enumerate() { + f(env, world); + if let Some(rbd) = self.rbd.as_mut() { + rbd.multibodies_mut().sync_joint_data_from_rapier( + backend, + env as u32, + &world.multibody_joints, + &world.bodies, + )?; + } + } + Ok(()) + } + + /// Reads every environment's multibody link workspace back from the GPU in + /// one transfer: per link, the generalized joint coordinates, accumulated + /// joint rotation, world pose and world-space velocity. + /// + /// The result is `num_environments() * multibody_links_per_env()` entries, + /// environment-major. Links follow the GPU build's traversal order + /// (multibodies, then links, parent before child), the same order + /// [`Self::control_multibody_motors`] targets. Empty when there is no + /// multibody state. + /// + /// Velocities only become meaningful after the first simulated step; the + /// coordinates and poses are valid from `finalize` on. + #[cfg(feature = "dim3")] + pub async fn read_multibody_links( + &self, + backend: &GpuBackend, + ) -> Vec { + let Some(rbd) = self.rbd.as_ref() else { + return Vec::new(); + }; + let mbs = rbd.multibodies(); + if mbs.links_per_batch() == 0 { + return Vec::new(); + } + // The workspace is batch-interleaved SoA quads, so read the raw buffer + // and decode it back into one struct per link. + let mut raw = vec![crate::rbd::glamx::Vec4::ZERO; mbs.links_workspace().len() as usize]; + if backend + .slow_read_buffer(mbs.links_workspace().buffer(), &mut raw) + .await + .is_err() + { + return Vec::new(); + } + crate::rbd::shaders::dynamics::ws_soa_to_structs( + &raw, + mbs.links_per_batch(), + mbs.num_batches(), + ) + } + + /// Number of link slots per environment, the stride of + /// [`Self::read_multibody_links`]. + #[cfg(feature = "dim3")] + pub fn multibody_links_per_env(&self) -> u32 { + self.rbd + .as_ref() + .map(|rbd| rbd.multibodies().links_per_batch()) + .unwrap_or(0) + } + /// Mutable access to environment `env`'s rapier world that does **not** mark /// the rbd state dirty, for use after [`Self::finalize`]. /// @@ -443,6 +568,7 @@ impl NexusState { self.rbd_dirty = true; // MPM-coupled boundary colliders live only in environment 0 and feed the // MPM coupling rebuild in `finalize`. + #[cfg(feature = "mpm")] if env == 0 && coupling != RbdCoupling::None { self.mpm_dirty = true; } @@ -543,6 +669,7 @@ impl NexusState { } self.rbd_dirty = true; } + #[cfg(feature = "mpm")] if couplings.iter().any(|c| *c != RbdCoupling::None) { self.mpm_dirty = true; } @@ -662,6 +789,7 @@ impl NexusState { } /// Appends a new chunk of MPM particles (`O(added)`) and returns its handle. + #[cfg(feature = "mpm")] pub fn add_particles( &mut self, backend: &GpuBackend, @@ -679,6 +807,7 @@ impl NexusState { } /// Appends more particles to an existing chunk (`O(added)`). + #[cfg(feature = "mpm")] pub fn extend_chunk( &mut self, backend: &GpuBackend, @@ -699,11 +828,13 @@ impl NexusState { } /// MPM background-grid cell width. + #[cfg(feature = "mpm")] pub fn mpm_cell_width(&self) -> f32 { self.mpm_cell_width } /// Removes every particle of a chunk (`O(removed)`) and drops the handle. + #[cfg(feature = "mpm")] pub fn remove_chunk( &mut self, backend: &GpuBackend, @@ -723,6 +854,7 @@ impl NexusState { /// Removes up to `count` particles from a chunk (`O(removed)`), returning the /// number actually removed. The chunk itself is kept (even if emptied). + #[cfg(feature = "mpm")] pub fn remove_particles_from_chunk( &mut self, backend: &GpuBackend, @@ -749,6 +881,7 @@ impl NexusState { /// Swap-removes the given GPU particle slots and patches `slot2chunk` to /// follow the relocations the GPU performed. + #[cfg(feature = "mpm")] fn swap_remove_particle_slots( &mut self, backend: &GpuBackend, @@ -919,6 +1052,7 @@ impl NexusState { // as rigid bodies tagged `RbdCoupling::Mpm*`; rebuild the coupling // (sampled rigid particles, uploaded body set) whenever those bodies or // the particle set changed. + #[cfg(feature = "mpm")] if (rbd_was_dirty || self.mpm_dirty) && self.mpm.is_some() { let world = &self.rbd_envs[0]; let mut coupling = Vec::new(); diff --git a/src_rbd/broad_phase/lbvh.rs b/src_rbd/broad_phase/lbvh.rs index 8d9f0120..a081fab2 100644 --- a/src_rbd/broad_phase/lbvh.rs +++ b/src_rbd/broad_phase/lbvh.rs @@ -336,6 +336,7 @@ impl Lbvh { collision_pairs_indirect: &mut Tensor<[u32; 3]>, collision_groups: &Tensor, pair_filter: &Tensor<[u32; 2]>, + sim_params: &Tensor, ) -> Result<(), GpuBackendError> { state.resize_bf_buffers(backend, colliders_len); @@ -360,6 +361,7 @@ impl Lbvh { collision_groups, batch_indices, pair_filter, + sim_params, )?; // Single 256-lane workgroup: parallel max over the per-batch counts. self.shaders.lbvh_init_indirect_args.call( diff --git a/src_rbd/broad_phase/narrow_phase.rs b/src_rbd/broad_phase/narrow_phase.rs index 82d6ed60..94cbdecc 100644 --- a/src_rbd/broad_phase/narrow_phase.rs +++ b/src_rbd/broad_phase/narrow_phase.rs @@ -3,6 +3,8 @@ use crate::math::Pose; use crate::queries::GpuIndexedContact; use crate::shaders::PaddedVector; +#[cfg(feature = "dim3")] +use crate::shaders::broad_phase::GpuReduceContacts; use crate::shaders::broad_phase::{ CollisionPair, GpuInitPfmPfmDispatch, GpuNarrowPhaseInitContactsDispatch, GpuNarrowPhasePfmPfm, GpuNarrowPhaseShapeShape, GpuNarrowPhaseShapeShapeDeferred, GpuResetNarrowPhase, @@ -22,6 +24,8 @@ pub struct GpuNarrowPhase { /// `pfm_pairs` work-list. Split from `narrow_phase` to fit 8 storage buffers. narrow_phase_deferred: GpuNarrowPhaseShapeShapeDeferred, narrow_phase_pfm_pfm: GpuNarrowPhasePfmPfm, + #[cfg(feature = "dim3")] + reduce_contacts: GpuReduceContacts, init_pfm_pfm_indirect_args: GpuInitPfmPfmDispatch, init_contacts_indirect_args: GpuNarrowPhaseInitContactsDispatch, } @@ -49,6 +53,10 @@ impl GpuNarrowPhase { batch_indices: &Tensor, collider_parent: &Tensor, collider_materials: &Tensor, + sim_params: &Tensor, + // Optional: merge each collider pair's manifolds into one before the + // solvers see them. `false` skips the kernel entirely. + reduce_contacts: bool, ) -> Result<(), GpuBackendError> { let num_batches = contacts_len.len() as u32; self.reset_narrow_phase @@ -66,6 +74,7 @@ impl GpuNarrowPhase { batch_indices, collider_parent, collider_materials, + sim_params, )?; // Pass 2: defer the complex shape pairs into `pfm_pairs` (kept as a @@ -80,6 +89,7 @@ impl GpuNarrowPhase { pfm_pairs, pfm_pairs_len, batch_indices, + sim_params, vertices, indices, )?; @@ -98,7 +108,23 @@ impl GpuNarrowPhase { indices, collider_parent, collider_materials, + sim_params, )?; + // Reduction rewrites `contacts_len`, so it has to run before the + // indirect args are derived from it. + #[cfg(feature = "dim3")] + if reduce_contacts { + self.reduce_contacts.call( + pass, + [1u32, num_batches, 1], + contacts, + contacts_len, + batch_indices, + sim_params, + )?; + } + #[cfg(not(feature = "dim3"))] + let _ = reduce_contacts; self.init_contacts_indirect_args.call( pass, 256u32, diff --git a/src_rbd/dynamics/mod.rs b/src_rbd/dynamics/mod.rs index c20e1f41..743a91fd 100644 --- a/src_rbd/dynamics/mod.rs +++ b/src_rbd/dynamics/mod.rs @@ -6,7 +6,9 @@ pub use coloring::{ColorBucketsArgs, ColoringArgs, GpuColoring}; pub use joint::{GpuImpulseJointSet, GpuJointSolver, JointSolverArgs, convert_joint_motor}; pub use mprops_update::{GpuMpropsUpdate, GpuSyncColliderPosesShader}; #[cfg(feature = "dim3")] -pub use multibody::{GpuMultibodySet, GpuMultibodySolver, MultibodySolverArgs}; +pub use multibody::{ + GpuMultibodySet, GpuMultibodySnapshot, GpuMultibodySolver, MultibodySolverArgs, +}; pub use prep_render::{RbdInstanceDesc, WgRbdPrepRender}; pub use solver::{GpuSolver, SolverArgs}; pub use warmstart::{GpuWarmstart, WarmstartArgs}; diff --git a/src_rbd/dynamics/multibody/env_reset.rs b/src_rbd/dynamics/multibody/env_reset.rs new file mode 100644 index 00000000..8d8008c5 --- /dev/null +++ b/src_rbd/dynamics/multibody/env_reset.rs @@ -0,0 +1,412 @@ +//! Per-environment reset primitives for batched RL envs. +//! +//! Two paths sit on top of the `gpu_mb_env_reset*` kernels: +//! +//! - [`GpuMultibodySnapshot`] plus [`GpuMultibodySet::reset_env_from_snapshot`]: +//! one staging upload and one dispatch per reset, no GPU to CPU readback. +//! - [`GpuMultibodySet::publish_reset_templates`] plus +//! [`GpuMultibodySet::encode_reset_envs_batch`]: the templates live on the +//! GPU permanently and N resets ride a single dispatch, with the teleport +//! offset applied in-kernel. +//! +//! Reset loops should prefer the second: it is what keeps a rollout free of +//! per-step host writes, and therefore capturable into a CUDA graph. + +use super::multibody_set::GpuMultibodySet; +use crate::math::Vector; +use crate::shaders::dynamics::{ + GpuMbEnvReset, GpuMbEnvResetBatch, GpuMbEnvResetBatchDofs, MULTIBODY_ROOT, MultibodyLinkStatic, + MultibodyLinkWorkspace, WS_QUADS, ws_soa_from_structs, ws_soa_to_structs, +}; +use glamx::{UVec4, Vec4}; +use khal::BufferUsages; +use khal::Shader; +use khal::backend::{Backend, GpuBackend}; +use vortx::tensor::Tensor; + +/// CPU snapshot of one (single-batch) multibody template: the AoS per-link +/// workspace, the static link descriptors, and the generalized coordinates and +/// velocities of batch 0. +#[derive(Clone)] +pub struct GpuMultibodySnapshot { + /// AoS per-link workspace of batch 0, `links_per_batch` entries including + /// the padding slots. Converted to the SoA quad layout on upload. + pub(super) links_workspace: Vec, + pub(super) links_static: Vec, + /// Generalized coordinates of batch 0 (`dofs_per_batch`). + pub(super) dof_values: Vec, + /// Generalized velocities of batch 0 (`dofs_per_batch`): the velocity + /// section of `dof_state`, the sections after it being static config. + pub(super) dof_vels: Vec, +} + +impl GpuMultibodySnapshot { + /// True for entries describing a real link. The buffers are padded to + /// `links_per_batch` with zeroed slots (rb_id 0, parent 0, ndofs 0), a + /// combination no real link can have: a chain's body-0 link is its root, + /// and roots carry `parent_link_id == MULTIBODY_ROOT`. + pub(super) fn link_is_valid(ls: &MultibodyLinkStatic) -> bool { + ls.parent_link_id == MULTIBODY_ROOT || ls.ndofs > 0 || ls.rb_id != 0 + } + + /// Whether multibody `multibody_id` has an unlocked (floating) root, i.e. + /// whether an offset reset may move it. + pub(super) fn mb_root_is_free(&self, multibody_id: u32) -> bool { + self.links_static.iter().any(|ls| { + Self::link_is_valid(ls) + && ls.multibody_id == multibody_id + && ls.parent_link_id == MULTIBODY_ROOT + && ls.data.locked_axes == 0 + }) + } + + /// Calls `f(rb_id)` for every rigid body backing a link of a free-rooted + /// (floating-base) multibody: the set of bodies an offset reset moves. + pub(crate) fn for_each_link_rb_id(&self, mut f: impl FnMut(u32)) { + for ls in &self.links_static { + if Self::link_is_valid(ls) && self.mb_root_is_free(ls.multibody_id) { + f(ls.rb_id); + } + } + } + + /// A copy with every floating-base multibody translated by `offset` (world + /// frame). Rotations, joint coordinates past the free linear DoFs, + /// velocities and `dof_values` are translation-invariant; the free root's + /// world position lives in `coords[0..3]` and `local_to_parent` (a root's + /// parent frame is the world), and each link's `local_to_world` carries its + /// body pose. Fixed-base multibodies are untouched. `body_poses`, owned by + /// the caller, must be translated for the same rb ids. + pub(crate) fn translated(&self, offset: Vector) -> GpuMultibodySnapshot { + let mut out = self.clone(); + for (ws, ls) in out.links_workspace.iter_mut().zip(&self.links_static) { + if !Self::link_is_valid(ls) || !self.mb_root_is_free(ls.multibody_id) { + continue; + } + ws.local_to_world.translation += offset; + if ls.parent_link_id == MULTIBODY_ROOT { + ws.local_to_parent.translation += offset; + ws.coords[0] += offset.x; + ws.coords[1] += offset.y; + ws.coords[2] += offset.z; + } + } + out + } +} + +/// `#[derive(Shader)]` supplies `from_backend`, loading the embedded entry. +#[derive(Shader)] +struct EnvResetShader { + kernel: GpuMbEnvReset, +} + +/// `#[derive(Shader)]` supplies `from_backend` for the batched entry. +#[derive(Shader)] +struct EnvResetBatchShader { + kernel: GpuMbEnvResetBatch, + /// Static-link and DoF half, split off so each pass fits 8 storage buffers. + dofs: GpuMbEnvResetBatchDofs, +} + +/// Shader bundle plus persistent staging buffers for the per-env reset +/// scatter. Created on first reset, so the allocations stay outside any +/// captured region. +pub(super) struct EnvResetBundle { + shader: EnvResetShader, + staging_ws: Tensor, + staging_links: Tensor, + staging_dofs: Tensor, + params: Tensor, +} + +impl EnvResetBundle { + fn new(backend: &GpuBackend, lpb: u32, dpb: u32) -> Self { + let storage = BufferUsages::STORAGE | BufferUsages::COPY_DST; + let uniform = BufferUsages::STORAGE | BufferUsages::UNIFORM | BufferUsages::COPY_DST; + Self { + shader: EnvResetShader::from_backend(backend).unwrap(), + staging_ws: Tensor::vector( + backend, + vec![Vec4::ZERO; (lpb * WS_QUADS).max(1) as usize], + storage, + ) + .unwrap(), + staging_links: Tensor::vector( + backend, + vec![::zeroed(); lpb.max(1) as usize], + storage, + ) + .unwrap(), + staging_dofs: Tensor::vector(backend, vec![0.0f32; (2 * dpb).max(1) as usize], storage) + .unwrap(), + params: Tensor::scalar(backend, UVec4::new(0, 0, lpb, dpb), uniform).unwrap(), + } + } +} + +/// GPU-resident reset templates plus the batch-reset shader. +pub(super) struct ResetTemplatesMb { + ws: Tensor, + links: Tensor, + dofs: Tensor, + flags: Tensor, + shader: EnvResetBatchShader, + /// Host copies, used to keep the `links_static` mirror in step. + mirror_links: Vec>, +} + +impl GpuMultibodySet { + /// Reads this set's batch-0 state off the GPU into a CPU snapshot. Call it + /// once per template at setup (typically on a single-env set) and pass the + /// result to [`Self::reset_env_from_snapshot`] for readback-free resets. + pub async fn snapshot(&self, backend: &GpuBackend) -> GpuMultibodySnapshot { + let nb = self.num_batches; + let lpb = self.links_per_batch as usize; + let dpb = self.dofs_per_batch as usize; + + let mut ws_soa: Vec = bytemuck::zeroed_vec(self.links_workspace.len() as usize); + backend + .slow_read_buffer(self.links_workspace.buffer(), &mut ws_soa) + .await + .unwrap(); + let mut ls_all: Vec = + bytemuck::zeroed_vec(self.links_static.len() as usize); + backend + .slow_read_buffer(self.links_static.buffer(), &mut ls_all) + .await + .unwrap(); + let mut dv_all: Vec = bytemuck::zeroed_vec(self.dof_values.len() as usize); + backend + .slow_read_buffer(self.dof_values.buffer(), &mut dv_all) + .await + .unwrap(); + let mut ds_all: Vec = bytemuck::zeroed_vec(self.dof_state.len() as usize); + backend + .slow_read_buffer(self.dof_state.buffer(), &mut ds_all) + .await + .unwrap(); + + // Gather batch 0 out of the interleave; the workspace is de-SoA'd + // through the shared layout accessors, so it stays one source of truth + // with the kernels. `ws_soa_to_structs` lays batch `b` out at + // `b * links_cap`, so batch 0 is the leading `lpb` entries. + let mut links_workspace = ws_soa_to_structs(&ws_soa, lpb as u32, nb); + links_workspace.truncate(lpb); + GpuMultibodySnapshot { + links_workspace, + links_static: (0..lpb).map(|k| ls_all[k * nb as usize]).collect(), + dof_values: (0..dpb).map(|d| dv_all[d * nb as usize]).collect(), + dof_vels: (0..dpb).map(|d| ds_all[d * nb as usize]).collect(), + } + } + + /// Resets env `dst_env` from a CPU snapshot. + pub fn reset_env_from_snapshot( + &mut self, + backend: &GpuBackend, + dst_env: u32, + snap: &GpuMultibodySnapshot, + ) { + if self.is_empty() { + return; + } + let nb = self.num_batches; + let lpb = self.links_per_batch; + let dpb = self.dofs_per_batch; + debug_assert_eq!(snap.links_static.len(), lpb as usize); + debug_assert_eq!(snap.dof_values.len(), dpb as usize); + + // Keep the host mirror in lockstep: the motor setters read-modify-write + // it. + for k in 0..lpb as usize { + self.links_static_mirror[k * nb as usize + dst_env as usize] = snap.links_static[k]; + } + + // Take the bundle out so the live buffers below can be borrowed + // mutably at the same time. + let mut bundle = match self.env_reset.take() { + Some(b) => b, + None => EnvResetBundle::new(backend, lpb, dpb), + }; + + let ws = ws_soa_from_structs(&snap.links_workspace, lpb, 1); + backend + .write_buffer(bundle.staging_ws.buffer_mut(), 0, &ws) + .unwrap(); + backend + .write_buffer(bundle.staging_links.buffer_mut(), 0, &snap.links_static) + .unwrap(); + let mut dofs = snap.dof_values.clone(); + dofs.extend_from_slice(&snap.dof_vels); + if !dofs.is_empty() { + backend + .write_buffer(bundle.staging_dofs.buffer_mut(), 0, &dofs) + .unwrap(); + } + bundle.params = Tensor::scalar( + backend, + UVec4::new(dst_env, nb, lpb, dpb), + BufferUsages::STORAGE | BufferUsages::UNIFORM | BufferUsages::COPY_DST, + ) + .unwrap(); + + let mut encoder = backend.begin_encoding(); + { + use khal::backend::Encoder as _; + let mut pass = encoder.begin_pass("[RBD] mb-env-reset", None); + bundle + .shader + .kernel + .call( + &mut pass, + lpb * WS_QUADS, + &bundle.staging_ws, + &bundle.staging_links, + &bundle.staging_dofs, + &mut self.links_workspace, + &mut self.links_static, + &mut self.dof_values, + &mut self.dof_state, + &bundle.params, + ) + .unwrap(); + } + backend.submit(encoder).unwrap(); + self.env_reset = Some(bundle); + } + + /// Uploads the reset templates once as GPU-resident blobs (SoA workspace, + /// links, coords and velocities) plus the per-link translate flags the + /// batch kernel needs, enabling [`Self::encode_reset_envs_batch`]. A host + /// copy of each template's `links_static` is kept so the batch reset can + /// maintain the CPU mirror. + pub fn publish_reset_templates( + &mut self, + backend: &GpuBackend, + snaps: &[&GpuMultibodySnapshot], + ) { + if self.is_empty() || snaps.is_empty() { + return; + } + let lpb = self.links_per_batch as usize; + let dpb = self.dofs_per_batch as usize; + let storage = BufferUsages::STORAGE | BufferUsages::COPY_DST; + + let mut ws = Vec::with_capacity(snaps.len() * lpb * WS_QUADS as usize); + let mut links = Vec::with_capacity(snaps.len() * lpb); + let mut dofs = Vec::with_capacity(snaps.len() * 2 * dpb); + let mut mirror_links = Vec::with_capacity(snaps.len()); + for snap in snaps { + debug_assert_eq!(snap.links_static.len(), lpb); + debug_assert_eq!(snap.dof_values.len(), dpb); + ws.extend_from_slice(&ws_soa_from_structs(&snap.links_workspace, lpb as u32, 1)); + links.extend_from_slice(&snap.links_static); + dofs.extend_from_slice(&snap.dof_values); + dofs.extend_from_slice(&snap.dof_vels); + mirror_links.push(snap.links_static.clone()); + } + // Per-link translate flags, constant per robot and identical across + // templates: bit 0 = valid link of a free-root multibody, bit 1 = the + // root link itself. Matches `GpuMultibodySnapshot::translated`. + let flags: Vec = snaps[0] + .links_static + .iter() + .map(|ls| { + let movable = GpuMultibodySnapshot::link_is_valid(ls) + && snaps[0].mb_root_is_free(ls.multibody_id); + (movable as u32) | (((movable && ls.parent_link_id == MULTIBODY_ROOT) as u32) << 1) + }) + .collect(); + + self.reset_templates = Some(ResetTemplatesMb { + ws: Tensor::vector(backend, &ws, storage).unwrap(), + links: Tensor::vector(backend, &links, storage).unwrap(), + dofs: Tensor::vector(backend, &dofs, storage).unwrap(), + flags: Tensor::vector(backend, &flags, storage).unwrap(), + shader: EnvResetBatchShader::from_backend(backend).unwrap(), + mirror_links, + }); + } + + /// Encodes one dispatch resetting every `(dst_env, template)` in `resets` + /// from the resident templates, translating each by its `offsets` entry and + /// writing its `dof_vels` slice (`dofs_per_batch` floats per reset) into + /// the velocity section. Only the compact reset list is uploaded. The host + /// `links_static` mirror is refreshed for the reset envs. + /// + /// [`Self::publish_reset_templates`] must have run first. + pub fn encode_reset_envs_batch( + &mut self, + backend: &GpuBackend, + enc: &mut ::Encoder, + resets: &[UVec4], + offsets: &[Vec4], + dof_vels: &[f32], + ) { + use khal::backend::Encoder as _; + let n = resets.len() as u32; + if n == 0 || self.is_empty() { + return; + } + let nb = self.num_batches; + let lpb = self.links_per_batch; + let dpb = self.dofs_per_batch; + debug_assert_eq!(dof_vels.len(), (n * dpb) as usize); + let tpl = self + .reset_templates + .take() + .expect("publish_reset_templates must run first"); + + // Host mirror lockstep: the motor setters read-modify-write it. + for meta in resets { + let (env, t) = (meta.x as usize, meta.y as usize); + for (k, ls) in tpl.mirror_links[t].iter().enumerate() { + self.links_static_mirror[k * nb as usize + env] = *ls; + } + } + + let storage = BufferUsages::STORAGE | BufferUsages::COPY_DST; + let t_resets = Tensor::vector(backend, resets, storage).unwrap(); + let t_offs = Tensor::vector(backend, offsets, storage).unwrap(); + let t_vels = Tensor::vector(backend, dof_vels, storage).unwrap(); + let params = Tensor::scalar( + backend, + UVec4::new(nb, lpb, dpb, n), + BufferUsages::STORAGE | BufferUsages::UNIFORM | BufferUsages::COPY_DST, + ) + .unwrap(); + { + let mut pass = enc.begin_pass("[RBD] mb-env-reset-batch", None); + tpl.shader + .kernel + .call( + &mut pass, + [lpb * WS_QUADS, n, 1], + &tpl.ws, + &tpl.flags, + &t_resets, + &t_offs, + &mut self.links_workspace, + ¶ms, + ) + .unwrap(); + tpl.shader + .dofs + .call( + &mut pass, + [lpb.max(dpb), n, 1], + &tpl.links, + &tpl.dofs, + &t_resets, + &t_vels, + &mut self.links_static, + &mut self.dof_values, + &mut self.dof_state, + ¶ms, + ) + .unwrap(); + } + self.reset_templates = Some(tpl); + } +} diff --git a/src_rbd/dynamics/multibody/mod.rs b/src_rbd/dynamics/multibody/mod.rs index 19bca7fa..b0bfea14 100644 --- a/src_rbd/dynamics/multibody/mod.rs +++ b/src_rbd/dynamics/multibody/mod.rs @@ -8,10 +8,12 @@ #![cfg(feature = "dim3")] +mod env_reset; mod loop_closing_joints; mod multibody_from_rapier; mod multibody_set; mod multibody_solver; +pub use env_reset::GpuMultibodySnapshot; pub use multibody_set::GpuMultibodySet; pub use multibody_solver::{GpuMultibodySolver, MultibodySolverArgs}; diff --git a/src_rbd/dynamics/multibody/multibody_from_rapier.rs b/src_rbd/dynamics/multibody/multibody_from_rapier.rs index ce742aa9..ae9a4141 100644 --- a/src_rbd/dynamics/multibody/multibody_from_rapier.rs +++ b/src_rbd/dynamics/multibody/multibody_from_rapier.rs @@ -43,6 +43,7 @@ impl GpuMultibodySet { let mut per_env_dof_vels: Vec> = Vec::with_capacity(num_batches as usize); let mut per_env_dof_damping: Vec> = Vec::with_capacity(num_batches as usize); let mut per_env_dof_armature: Vec> = Vec::with_capacity(num_batches as usize); + let mut per_env_dof_friction: Vec> = Vec::with_capacity(num_batches as usize); let mut per_env_dof_stiffness: Vec> = Vec::with_capacity(num_batches as usize); let mut per_env_dof_spring_ref: Vec> = Vec::with_capacity(num_batches as usize); let mut per_env_dof_kinematic: Vec> = Vec::with_capacity(num_batches as usize); @@ -63,6 +64,12 @@ impl GpuMultibodySet { let mut global_max_cons = 0u32; let mut global_max_couplings = 0u32; + // Whether any multibody anywhere declares dry joint friction. + let scene_has_joint_friction = environments.iter().any(|(set, _, _)| { + set.multibodies() + .any(|mb| mb.frictions().iter().any(|f| *f > 0.0)) + }); + for (set, body_ids, bodies) in environments { let mut infos = Vec::new(); let mut statics = Vec::new(); @@ -71,6 +78,7 @@ impl GpuMultibodySet { let mut dof_vels = Vec::new(); let mut dof_damping = Vec::new(); let mut dof_armature = Vec::new(); + let mut dof_friction = Vec::new(); let mut dof_stiffness = Vec::new(); let mut dof_spring_ref = Vec::new(); let mut dof_kinematic = Vec::new(); @@ -151,7 +159,10 @@ impl GpuMultibodySet { // One extra constraint slot per DoF coupling (they are solved // among the joint constraints). let num_couplings = mb.couplings().len() as u32; - let max_constraints = max_constraints + num_couplings; + // One dry-friction row per DoF, emitted only for DoFs whose + // friction is non-zero (see `gpu_mb_init_joint_constraints`). + let friction_slots = if scene_has_joint_friction { ndofs } else { 0 }; + let max_constraints = max_constraints + num_couplings + friction_slots; max_mb_joint_constraints = max_mb_joint_constraints.max(max_constraints); infos.push(MultibodyInfo { @@ -185,6 +196,8 @@ impl GpuMultibodySet { // `assembly_counter` (which drops the fixed root's DoFs). let mb_damping = mb.damping(); let mb_armature = mb.armature(); + // Per-DoF dry joint friction. + let mb_friction = mb.frictions(); let mut rapier_assembly = 0usize; let mb_statics_start = statics.len(); for (link_idx, link) in mb.links().enumerate() { @@ -292,6 +305,7 @@ impl GpuMultibodySet { dof_vels.push(0.0); dof_damping.push(mb_damping[rapier_assembly + d]); dof_armature.push(mb_armature[rapier_assembly + d]); + dof_friction.push(mb_friction[rapier_assembly + d]); let (k_s, rest) = link.joint().spring(free_axis_of_dof[d]); dof_stiffness.push(k_s); dof_spring_ref.push(rest); @@ -345,6 +359,7 @@ impl GpuMultibodySet { per_env_dof_vels.push(dof_vels); per_env_dof_damping.push(dof_damping); per_env_dof_armature.push(dof_armature); + per_env_dof_friction.push(dof_friction); per_env_dof_stiffness.push(dof_stiffness); per_env_dof_spring_ref.push(dof_spring_ref); per_env_dof_kinematic.push(dof_kinematic); @@ -403,6 +418,7 @@ impl GpuMultibodySet { let mut all_dof_vels: Vec = Vec::with_capacity((dofs_cap * num_batches) as usize); let mut all_dof_damping: Vec = Vec::with_capacity((dofs_cap * num_batches) as usize); let mut all_dof_armature: Vec = Vec::with_capacity((dofs_cap * num_batches) as usize); + let mut all_dof_friction: Vec = Vec::with_capacity((dofs_cap * num_batches) as usize); let mut all_dof_stiffness: Vec = Vec::with_capacity((dofs_cap * num_batches) as usize); let mut all_dof_spring_ref: Vec = Vec::with_capacity((dofs_cap * num_batches) as usize); @@ -442,6 +458,9 @@ impl GpuMultibodySet { all_dof_armature.extend_from_slice(&per_env_dof_armature[i]); let pad = (dofs_cap as usize).saturating_sub(per_env_dof_armature[i].len()); all_dof_armature.resize(all_dof_armature.len() + pad, 0.0); + all_dof_friction.extend_from_slice(&per_env_dof_friction[i]); + let pad = (dofs_cap as usize).saturating_sub(per_env_dof_friction[i].len()); + all_dof_friction.resize(all_dof_friction.len() + pad, 0.0); all_dof_stiffness.extend_from_slice(&per_env_dof_stiffness[i]); let pad = (dofs_cap as usize).saturating_sub(per_env_dof_stiffness[i].len()); all_dof_stiffness.resize(all_dof_stiffness.len() + pad, 0.0); @@ -478,6 +497,7 @@ impl GpuMultibodySet { let all_dof_vels = interleave(&all_dof_vels, dofs_cap, nb); let all_dof_damping = interleave(&all_dof_damping, dofs_cap, nb); let all_dof_armature = interleave(&all_dof_armature, dofs_cap, nb); + let all_dof_friction = interleave(&all_dof_friction, dofs_cap, nb); let all_dof_stiffness = interleave(&all_dof_stiffness, dofs_cap, nb); let all_dof_spring_ref = interleave(&all_dof_spring_ref, dofs_cap, nb); let all_dof_kinematic = interleave(&all_dof_kinematic, dofs_cap, nb); @@ -498,7 +518,10 @@ impl GpuMultibodySet { // falls back to a single plain matrix with explicit-only // coriolis/gyroscopic forces (cheaper, but less stable). implicit_coriolis: true, + coriolis_in_uniform: true, has_joint_constraints: all_infos.iter().any(|info| info.max_constraints > 0), + frictionloss_slots_reserved: scene_has_joint_friction, + constraint_caps_dirty: false, multibody_info: Tensor::vector(backend, &all_infos, storage).unwrap(), max_contact_constraints: Tensor::scalar( @@ -511,27 +534,35 @@ impl GpuMultibodySet { .unwrap(), links_static_mirror: all_statics.clone(), info_mirror, + // COPY_SRC so hosts can read joint/link state back (observation + // pipelines); see `GpuMultibodySet::links_workspace`. links_workspace: Tensor::vector( backend, crate::shaders::dynamics::ws_soa_from_structs(&all_ws, links_cap, num_batches), - storage, + storage | BufferUsages::COPY_SRC, ) .unwrap(), dof_values: Tensor::vector(backend, &all_dof_vals, storage).unwrap(), dof_state: { // Pack [velocities, damping, armature, spring stiffness, - // spring rest, kinematic mask] back-to-back, each section - // N = dofs_cap * num_batches long. The shaders address - // section `s` at intra-batch offset `s · dof_batch_capacity`. + // spring rest, kinematic mask, frictionloss] back-to-back, + // each section N = dofs_cap * num_batches long. The shaders + // address section `s` at intra-batch offset + // `s · dof_batch_capacity`. The friction section comes from + // rapier's `Multibody::frictions` (MJCF ``); `RbdState::set_dof_frictionloss` overrides + // it afterwards. let n = (dofs_cap * num_batches) as usize; - let mut buf = Vec::with_capacity(6 * n); + let mut buf = Vec::with_capacity(7 * n); buf.extend_from_slice(&all_dof_vels); buf.extend_from_slice(&all_dof_damping); buf.extend_from_slice(&all_dof_armature); buf.extend_from_slice(&all_dof_stiffness); buf.extend_from_slice(&all_dof_spring_ref); buf.extend_from_slice(&all_dof_kinematic); - debug_assert_eq!(buf.len(), 6 * n); + buf.extend_from_slice(&all_dof_friction); + buf.resize(7 * n, 0.0); + debug_assert_eq!(buf.len(), 7 * n); Tensor::vector(backend, &buf, storage).unwrap() }, gen_forces: Tensor::vector( @@ -584,6 +615,42 @@ impl GpuMultibodySet { dof_couplings: Tensor::vector(backend, &all_couplings, storage).unwrap(), couplings_per_batch: couplings_cap, body_to_link: Tensor::vector(backend, &all_body_to_link, storage).unwrap(), + body_to_link_host: all_body_to_link, + body_to_link_cap, + motor_delay_state: Tensor::vector( + backend, + vec![0.0f32; ((2 + links_cap) * num_batches) as usize], + storage | BufferUsages::COPY_DST, + ) + .unwrap(), + motor_delay_params: Tensor::scalar( + backend, + glamx::UVec4::new(num_batches, 2 + links_cap, 0, 0), + BufferUsages::STORAGE | BufferUsages::UNIFORM, + ) + .unwrap(), + delay_update_cache: None, + contact_sensor_links: Tensor::vector( + backend, + [u32::MAX; crate::shaders::dynamics::MAX_CONTACT_SENSORS as usize], + storage, + ) + .unwrap(), + contact_sensor_out: Tensor::vector( + backend, + vec![ + 0.0f32; + (mb_cap * num_batches * crate::shaders::dynamics::MAX_CONTACT_SENSORS) as usize + ], + storage | BufferUsages::COPY_SRC, + ) + .unwrap(), + num_contact_sensors: 0, + substep_refresh: true, + substep_refresh_light: false, + env_reset: None, + reset_templates: None, + scatter_caches: Vec::new(), contact_constraints: Tensor::vector( backend, vec![ diff --git a/src_rbd/dynamics/multibody/multibody_set.rs b/src_rbd/dynamics/multibody/multibody_set.rs index 5a017157..cd90f43f 100644 --- a/src_rbd/dynamics/multibody/multibody_set.rs +++ b/src_rbd/dynamics/multibody/multibody_set.rs @@ -9,6 +9,7 @@ use crate::shaders::dynamics::{ }; use crate::shaders::utils::BatchIndices; use khal::BufferUsages; +use khal::Shader; use khal::backend::{Backend, GpuBackend, GpuBackendError}; use rapier3d::prelude::JointAxis; use vortx::tensor::Tensor; @@ -26,6 +27,42 @@ pub(super) const MAX_DELASSUS_MULTIBODIES: u32 = 0; // 128; use crate::shaders::dynamics::{GenericJoint, JointLimits, JointMotor}; +/// The motor-target scatter entry point, loaded once per backend. +#[derive(Shader)] +pub(super) struct MotorScatterBundle { + scatter: crate::shaders::dynamics::GpuScatterMotorTargets, +} + +/// Shader plus the constant tensors of one motor-target scatter configuration. +/// The link ids, counts and axis do not change between steps, so a per-call +/// `from_backend` and four allocations would cost more than the dispatch. +pub(super) struct MotorScatterCache { + shader: MotorScatterBundle, + t_links: Tensor, + u_na: Tensor, + u_ne: Tensor, + u_ax: Tensor, + num_actuated: u32, + axis: u32, + link_ids: Vec, +} + +/// The on-device delay-state refresh entry point. +#[derive(Shader)] +pub(super) struct DelayUpdateBundle { + kernel: crate::shaders::dynamics::GpuMbDelayStateUpdate, +} + +/// Shader plus the constant tensors of the delay-state refresh. The link ids +/// and counts do not change between steps, so a per-call `from_backend` plus +/// uniform allocations would cost more than the upload this path removes. +pub(super) struct DelayUpdateCache { + shader: DelayUpdateBundle, + t_links: Tensor, + params: Tensor, + num_actuated: u32, +} + /// GPU-resident articulated multibody set, packed across simulation batches. /// /// Every buffer is a flat tensor with per-batch capacity (`*_batch_capacity`) and @@ -45,9 +82,31 @@ pub struct GpuMultibodySet { pub(super) coriolis_entries_per_batch: u32, pub(super) i_coriolis_dt_entries_per_batch: u32, pub(super) implicit_coriolis: bool, + /// What [`Self::implicit_coriolis`] was the last time the `batch_indices` + /// uniform was built. The kernels read the flag from that uniform, so the + /// two drifting apart silently changes which dynamics path runs; the + /// dispatch guard compares them. + pub(super) coriolis_in_uniform: bool, + /// Rebuild the joint and contact constraints from scratch every substep. + /// On (the default) this matches a per-substep constraint refresh; off, the + /// full build runs once per step and each later substep only refreshes the + /// joint rhs / limit activity, which is cheaper and closer to how MuJoCo + /// and Genesis step. + pub(super) substep_refresh: bool, + /// Split cadence: refresh the constraints every substep but keep the mass + /// matrix and its LU factors per step. Ignored when `substep_refresh` is on. + pub(super) substep_refresh_light: bool, /// When `false` (no joint limits / motors anywhere), the joint constraint /// kernel chain is skipped on the host side. pub(super) has_joint_constraints: bool, + /// `true` once the joint-constraint bank has been grown to hold a + /// dry-friction row per DoF; see `reserve_frictionloss_slots`. + pub(super) frictionloss_slots_reserved: bool, + /// Set when a capacity edit here has invalidated the shared `BatchIndices` + /// uniform. The next `RbdPipeline` step re-uploads it and clears this; + /// without that the kernels would index the resized buffers with stale + /// per-batch capacities. + pub(crate) constraint_caps_dirty: bool, /// Per-batch multibody descriptors. pub(super) multibody_info: Tensor, @@ -96,6 +155,34 @@ pub struct GpuMultibodySet { /// Per-body lookup `[multibody_idx, link_idx]` (`u32::MAX` sentinel for /// free / non-multibody bodies). Indexed by the per-batch local body id. pub(super) body_to_link: Tensor<[u32; 2]>, + /// Actuator-delay state, per batch `[tick, k, prev_target x + /// links_per_batch]`. All zeros (the default) means no delay. + pub(super) motor_delay_state: Tensor, + /// `(num_batches, stride, 0, 0)` uniform for the delay tick dispatch. + pub(super) motor_delay_params: Tensor, + /// Cached shader and constants for the on-device delay-state refresh. + pub(super) delay_update_cache: Option, + /// The sensed multibody link ids, `MAX_CONTACT_SENSORS` slots padded with + /// `u32::MAX`. The same set is sensed on every multibody in every batch. + pub(super) contact_sensor_links: Tensor, + /// Per-(multibody, slot) summed normal-contact impulse, written once per + /// step by `gpu_mb_sense_contact_impulses`. + pub(super) contact_sensor_out: Tensor, + /// Number of configured contact sensors; 0 skips the readout dispatch. + pub(super) num_contact_sensors: u32, + /// Shader plus staging buffers for the single-env reset scatter, created + /// on first use. + pub(super) env_reset: Option, + /// GPU-resident reset templates, published by `publish_reset_templates`. + pub(super) reset_templates: Option, + /// One entry per (axis, actuated link set) the caller has scattered motor + /// targets for. Grown on demand by [`Self::encode_scatter_motor_targets`]. + pub(super) scatter_caches: Vec, + /// CPU mirror of [`Self::body_to_link`], batch-major with a + /// `body_to_link_cap` stride. Backs [`Self::link_of_body`]. + pub(super) body_to_link_host: Vec<[u32; 2]>, + /// Per-batch stride of [`Self::body_to_link_host`] (colliders per batch). + pub(super) body_to_link_cap: u32, /// Per-multibody bank of contact constraints (1 normal + 2 friction per /// touched contact point). @@ -223,12 +310,38 @@ impl GpuMultibodySet { /// GPU buffer holding six back-to-back per-DOF sections of /// `dof_batch_capacity * num_batches` floats each: generalized /// velocities, damping, armature, spring stiffness, spring rest position, - /// and the kinematic-DOF mask. Callers reading velocities should use only - /// the first section. + /// the kinematic-DOF mask, and Coulomb joint friction. Callers reading + /// velocities should use only the first section. pub fn dof_state(&self) -> &Tensor { &self.dof_state } + /// Sets whether the joint and contact constraints are rebuilt from scratch + /// every substep. See [`Self::substep_refresh`]; on by default. + pub fn set_substep_refresh(&mut self, enabled: bool) { + self.substep_refresh = enabled; + } + + /// Whether the per-substep constraint rebuild is enabled. + pub fn substep_refresh(&self) -> bool { + self.substep_refresh + } + + /// Sets the split cadence: constraints per substep, mass matrix and LU per + /// step. Ignored while [`Self::substep_refresh`] is on. + pub fn set_substep_refresh_light(&mut self, enabled: bool) { + self.substep_refresh_light = enabled; + } + + /// Mutable view of [`Self::dof_state`], for callers that push generalized + /// velocities straight into the buffer (e.g. an external RL env resetting + /// one environment). Section offsets are the ones documented on + /// [`Self::dof_state`]; writing past the velocity section overwrites the + /// damping, armature and spring parameters. + pub fn dof_state_mut(&mut self) -> &mut Tensor { + &mut self.dof_state + } + /// Per-batch stride of the DoF buffers (the length of each section of /// [`Self::dof_state`]). pub fn dofs_per_batch(&self) -> u32 { @@ -432,6 +545,148 @@ impl GpuMultibodySet { ) } + /// Scatters per-(actuated joint, env) motor target positions into + /// `links_static` on the GPU, reading the targets from a GPU buffer so a + /// GPU-resident policy can drive the motors with no host round-trip. + /// + /// `targets` is row-major `[num_actuated x num_batches]` (element + /// `(j, env)` at `j · num_batches + env`) and `actuated_link_ids[j]` is the + /// link index of actuated joint `j`. Sets `motors[axis].target_pos` and the + /// matching `motor_axes` bit. + /// + /// This bypasses `links_static_mirror`: the scattered targets live only on + /// the GPU, so do not interleave it with [`Self::set_motor`] / + /// [`Self::set_motors`] on the same axis. + pub fn scatter_motor_targets( + &mut self, + backend: &GpuBackend, + targets: &Tensor, + actuated_link_ids: &[u32], + axis: u32, + ) -> Result<(), GpuBackendError> { + let mut enc = backend.begin_encoding(); + self.encode_scatter_motor_targets(backend, &mut enc, targets, actuated_link_ids, axis)?; + backend.submit(enc) + } + + /// [`Self::scatter_motor_targets`], recorded into an existing encoder so + /// the control step shares one submit with the caller's other work. + pub fn encode_scatter_motor_targets( + &mut self, + backend: &GpuBackend, + enc: &mut ::Encoder, + targets: &Tensor, + actuated_link_ids: &[u32], + axis: u32, + ) -> Result<(), GpuBackendError> { + use khal::backend::Encoder as _; + + // Take the matching entry out so `self.links_static` can be borrowed + // mutably for the dispatch; it goes back at the end. + let hit = self + .scatter_caches + .iter() + .position(|c| c.axis == axis && c.link_ids == actuated_link_ids); + let cache = match hit { + Some(i) => self.scatter_caches.swap_remove(i), + None => { + let num_actuated = actuated_link_ids.len() as u32; + let uu = BufferUsages::STORAGE | BufferUsages::UNIFORM; + MotorScatterCache { + shader: MotorScatterBundle::from_backend(backend)?, + t_links: Tensor::vector(backend, actuated_link_ids, BufferUsages::STORAGE)?, + u_na: Tensor::scalar(backend, num_actuated, uu)?, + u_ne: Tensor::scalar(backend, self.num_batches, uu)?, + u_ax: Tensor::scalar(backend, axis, uu)?, + num_actuated, + axis, + link_ids: actuated_link_ids.to_vec(), + } + } + }; + { + let mut pass = enc.begin_pass("[RBD] mb/scatter-motor-targets", None); + cache.shader.scatter.call( + &mut pass, + [cache.num_actuated, self.num_batches, 1], + targets, + &mut self.links_static, + &cache.t_links, + &cache.u_na, + &cache.u_ne, + &cache.u_ax, + )?; + } + self.scatter_caches.push(cache); + Ok(()) + } + + /// Per-batch per-step link workspace (generalized coordinates, joint + /// rotations, world-space link velocities), in the batch-interleaved SoA + /// quad layout the kernels index. Read it back with `slow_read_buffer` for + /// joint/base state observation and decode it with `ws_soa_to_structs`, + /// which yields one struct per link laid out `env * links_per_batch + link` + /// in [`from_rapier`](Self::from_rapier)'s link traversal order. + pub fn links_workspace(&self) -> &Tensor { + &self.links_workspace + } + + /// Per-batch static link data (joint definitions, motors, limits, mass + /// properties), batch-interleaved like [`Self::links_workspace`]. Exposed + /// for diagnostics; use [`Self::set_motor`] / [`Self::set_motors`] to + /// mutate motors so the CPU mirror stays in sync. + pub fn links_static(&self) -> &Tensor { + &self.links_static + } + + /// Number of link slots per environment (the stride of + /// [`Self::links_workspace`] and `links_static`). + pub fn links_per_batch(&self) -> u32 { + self.links_per_batch + } + + /// Refreshes every link's joint parameters (motor targets/gains, limits) of + /// environment `env` from a rapier multibody set laid out identically to the + /// one this GPU set was built from (same multibody/link traversal order as + /// [`from_rapier`](Self::from_rapier)), then uploads the `links_static` + /// buffer in one write. + /// + /// This is the per-step control path for actuated robots: mutate the motors + /// on the CPU rapier joints, then call this to push them to the GPU. Only + /// joint data is refreshed (coordinates, velocities and mass properties are + /// untouched), so this cannot teleport links. + pub fn sync_joint_data_from_rapier( + &mut self, + backend: &GpuBackend, + env: u32, + set: &crate::rapier::dynamics::MultibodyJointSet, + bodies: &crate::rapier::dynamics::RigidBodySet, + ) -> Result<(), GpuBackendError> { + let base = (env * self.links_per_batch) as usize; + let mut offset = 0usize; + for mb in set.multibodies() { + // Mirror `from_rapier`'s fixed-root handling: a non-dynamic root has + // all 6 DOFs locked on the GPU even though rapier models it as free. + let root_is_dynamic = mb + .link(0) + .and_then(|r| bodies.get(r.rigid_body_handle())) + .map(|rb| rb.is_dynamic()) + .unwrap_or(false); + for (link_idx, link) in mb.links().enumerate() { + let Some(entry) = self.links_static_mirror.get_mut(base + offset) else { + return Ok(()); + }; + let mut data = convert_generic_joint(link.joint().data); + if link_idx == 0 && !root_is_dynamic { + data.locked_axes = 0x3f; + } + entry.data = data; + offset += 1; + } + } + backend.write_buffer(self.links_static.buffer_mut(), 0, &self.links_static_mirror) + } + /// Number of multibody-touching impulse joints in any batch. pub fn mb_impulse_joints_per_batch(&self) -> u32 { self.mb_imp_joints_per_batch @@ -441,7 +696,8 @@ impl GpuMultibodySet { /// RBD-side fields (`colliders_batch_capacity`, `contacts_batch_capacity`, /// `collision_pairs_batch_capacity`, `impulse_joints_batch_capacity`, /// `color_groups_batch_capacity`) untouched — the caller fills those. - pub(crate) fn fill_batch_indices(&self, dst: &mut BatchIndices) { + pub(crate) fn fill_batch_indices(&mut self, dst: &mut BatchIndices) { + self.coriolis_in_uniform = self.implicit_coriolis; dst.multibodies_batch_capacity = self.multibodies_per_batch; dst.multibodies_len = self.num_active_multibodies; dst.links_batch_capacity = self.links_per_batch; @@ -545,6 +801,271 @@ impl GpuMultibodySet { self.contact_constraints_per_batch } + /// The per-multibody bank of unit (1-DoF) joint limit / motor constraints. + pub fn joint_constraints(&self) -> &Tensor { + &self.joint_constraints + } + + /// Per-batch stride of [`Self::joint_constraints`]. + pub fn joint_constraints_per_batch(&self) -> u32 { + self.joint_constraints_per_batch + } + + /// Per-batch stride of the joint-constraint `M⁻¹` column buffer. + pub fn joint_constraint_columns_per_batch(&self) -> u32 { + self.joint_constraint_columns_per_batch + } + + /// Overwrites the per-DoF armature (reflected rotor inertia) section of + /// [`Self::dof_state`]. `values` is `dofs_per_batch * num_batches` in + /// env-major order (env outer, DoF inner). + pub fn set_dof_armature(&mut self, backend: &GpuBackend, values: &[f32]) { + self.write_dof_section(backend, 2, values, "armature"); + } + + /// Overwrites the per-DoF dry joint friction coefficients. + pub fn set_dof_frictionloss(&mut self, backend: &GpuBackend, values: &[f32]) { + if values.iter().any(|v| *v > 0.0) { + self.reserve_frictionloss_slots(backend); + } + self.write_dof_section(backend, 6, values, "frictionloss"); + } + + /// Grows the joint-constraint bank by one slot per DoF of every multibody, + /// the worst case for `gpu_mb_init_joint_constraints`' dry-friction rows + /// (emitted only for DoFs whose `frictionloss` is non-zero). Idempotent, + /// and never called unless some frictionloss is actually set, so scenes + /// without joint friction keep the tighter limit/motor-only capacity. + fn reserve_frictionloss_slots(&mut self, backend: &GpuBackend) { + if self.frictionloss_slots_reserved { + return; + } + self.frictionloss_slots_reserved = true; + + let mb_cap = self.multibodies_per_batch as usize; + let nb = self.num_batches as usize; + let mut cons_cap = 0u32; + let mut max_constraints = 0u32; + for b in 0..nb { + let mut cons_off = 0u32; + for i in 0..mb_cap { + let info = &mut self.info_mirror[b * mb_cap + i]; + // Padding slots stay untouched: they hold no links, so the + // kernels bail out before reading their offsets. + if info.num_links == 0 { + continue; + } + info.first_constraint = cons_off; + info.max_constraints += info.ndofs; + cons_off += info.max_constraints; + max_constraints = max_constraints.max(info.max_constraints); + } + cons_cap = cons_cap.max(cons_off); + } + let cons_cap = cons_cap.max(1); + let cons_col_cap = cons_cap.saturating_mul(self.dofs_per_batch).max(1); + + // Batch-interleaved (batch-minor) upload, matching the build path. + let mut interleaved = Vec::with_capacity(mb_cap * nb); + for k in 0..mb_cap { + for b in 0..nb { + interleaved.push(self.info_mirror[b * mb_cap + k]); + } + } + let storage = BufferUsages::STORAGE | BufferUsages::COPY_DST; + self.multibody_info = Tensor::vector(backend, &interleaved, storage).unwrap(); + self.joint_constraints = Tensor::vector( + backend, + vec![MultibodyJointConstraint::default(); cons_cap as usize * nb], + storage, + ) + .unwrap(); + self.joint_constraint_columns = + Tensor::vector(backend, vec![0.0f32; cons_col_cap as usize * nb], storage).unwrap(); + self.joint_constraints_per_batch = cons_cap; + self.joint_constraint_columns_per_batch = cons_col_cap; + self.max_joint_constraints = max_constraints; + self.has_joint_constraints = max_constraints > 0; + // `BatchIndices` now disagrees with these capacities. `RbdState:: + // set_dof_frictionloss` rebuilds it immediately; for callers who came + // in through `GpuMultibodySet` directly, the next step does. + self.constraint_caps_dirty = true; + } + + /// Transposes an env-major `dofs_per_batch * num_batches` block into the + /// batch-interleaved layout and writes it over section `section` of + /// [`Self::dof_state`]. + fn write_dof_section( + &mut self, + backend: &GpuBackend, + section: u64, + values: &[f32], + what: &str, + ) { + let cap = self.dofs_per_batch as usize; + let nb = self.num_batches as usize; + assert_eq!( + values.len(), + cap * nb, + "{what}: expected dofs_per_batch * num_batches values" + ); + let mut interleaved = vec![0.0f32; cap * nb]; + for b in 0..nb { + for k in 0..cap { + interleaved[k * nb + b] = values[b * cap + k]; + } + } + backend + .write_buffer( + self.dof_state.buffer_mut(), + section * (cap * nb) as u64, + &interleaved, + ) + .unwrap(); + } + + /// Per-batch stride of the actuator-delay state buffer: + /// `[tick, k, prev_target x links_per_batch]`. + pub fn motor_delay_stride(&self) -> u32 { + 2 + self.links_per_batch + } + + /// Uploads the actuator-delay state for every batch. `data.len()` must be + /// `motor_delay_stride() * num_batches`; all zeros disables the delay. + /// + /// While a control step's substep counter `tick` is below that batch's `k`, + /// every motor tracks `prev_target[link]` instead of its current target, so + /// latency costs no mid-step host writes. Call this before the step's + /// kernels are queued: an upload issued between queued substeps stalls the + /// stream, which is exactly what the GPU-side delay exists to avoid. + pub fn write_motor_delay_state( + &mut self, + backend: &GpuBackend, + data: &[f32], + ) -> Result<(), GpuBackendError> { + assert_eq!( + data.len(), + (self.motor_delay_stride() * self.num_batches) as usize, + "motor delay state: expected motor_delay_stride() * num_batches values" + ); + backend.write_buffer(self.motor_delay_state.buffer_mut(), 0, data) + } + + /// Per-step actuator-delay refresh on device (see + /// `gpu_mb_delay_state_update`): `tick <- 0`, `k <- k_eff`, and the + /// actuated links' prev-target lanes copied from `prev_targets`, the motor + /// target tensor as it stood *before* this step's scatter. + /// + /// This replaces the full `stride * num_batches` host rebuild and upload + /// that [`Self::write_motor_delay_state`] performs. + pub fn update_motor_delay_state_gpu( + &mut self, + backend: &GpuBackend, + prev_targets: &Tensor, + k_eff: &Tensor, + actuated_link_ids: &[u32], + ) -> Result<(), GpuBackendError> { + let mut enc = backend.begin_encoding(); + self.encode_update_motor_delay_state( + backend, + &mut enc, + prev_targets, + k_eff, + actuated_link_ids, + )?; + backend.submit(enc) + } + + /// [`Self::update_motor_delay_state_gpu`], recorded into an existing + /// encoder so the delay refresh and the target scatter share one submit. + pub fn encode_update_motor_delay_state( + &mut self, + backend: &GpuBackend, + enc: &mut ::Encoder, + prev_targets: &Tensor, + k_eff: &Tensor, + actuated_link_ids: &[u32], + ) -> Result<(), GpuBackendError> { + use khal::backend::Encoder as _; + let stride = self.motor_delay_stride(); + let cache = match self.delay_update_cache.take() { + Some(c) => c, + None => { + let num_actuated = actuated_link_ids.len() as u32; + DelayUpdateCache { + shader: DelayUpdateBundle::from_backend(backend)?, + t_links: Tensor::vector(backend, actuated_link_ids, BufferUsages::STORAGE)?, + params: Tensor::scalar( + backend, + glamx::UVec4::new(num_actuated, self.num_batches, stride, 0), + BufferUsages::STORAGE | BufferUsages::UNIFORM, + )?, + num_actuated, + } + } + }; + { + let mut pass = enc.begin_pass("[RBD] mb/delay-state-update", None); + cache.shader.kernel.call( + &mut pass, + [cache.num_actuated, self.num_batches, 1], + prev_targets, + k_eff, + &cache.t_links, + &mut self.motor_delay_state, + &cache.params, + )?; + } + self.delay_update_cache = Some(cache); + Ok(()) + } + + /// Configures the contact force sensor: senses the summed normal-contact + /// impulse on these multibody links (at most + /// [`MAX_CONTACT_SENSORS`](crate::shaders::dynamics::MAX_CONTACT_SENSORS); + /// the same links are sensed on every multibody in every batch). Translate + /// a local body / collider id with [`Self::link_of_body`] first. An empty + /// slice disables the readout. + pub fn set_contact_sensor_links(&mut self, backend: &GpuBackend, links: &[u32]) { + const MAX: usize = crate::shaders::dynamics::MAX_CONTACT_SENSORS as usize; + assert!( + links.len() <= MAX, + "at most {MAX} contact sensors supported (got {})", + links.len() + ); + let mut padded = [u32::MAX; MAX]; + padded[..links.len()].copy_from_slice(links); + backend + .write_buffer(self.contact_sensor_links.buffer_mut(), 0, &padded) + .unwrap(); + self.num_contact_sensors = links.len() as u32; + } + + /// The contact force-sensor readout, interleaved like the other per-mb + /// buffers: slot `s` of multibody `m` in batch `b` at + /// `(m · num_batches + b) · MAX_CONTACT_SENSORS + s`. Read it after a step; + /// the values are accumulated normal impulses, so divide by the step `dt` + /// for an average force. + pub fn contact_sensor_out(&self) -> &Tensor { + &self.contact_sensor_out + } + + /// Number of configured contact sensors (0 means sensing is disabled). + pub fn num_contact_sensors(&self) -> u32 { + self.num_contact_sensors + } + + /// `[multibody_idx, link_idx]` of the local body / collider id + /// `local_body_id` within `batch`, or `[u32::MAX; 2]` when that body is not + /// a multibody link. Resolved on the CPU mirror, so it costs no readback. + pub fn link_of_body(&self, batch: u32, local_body_id: u32) -> [u32; 2] { + let idx = batch as usize * self.body_to_link_cap as usize + local_body_id as usize; + self.body_to_link_host + .get(idx) + .copied() + .unwrap_or([u32::MAX; 2]) + } + /// Per-constraint `Jᵀ` rows of the contact constraints (`ndofs` floats each, /// laid out like [`Self::contact_constraints`]). pub fn contact_constraint_jacs(&self) -> &Tensor { diff --git a/src_rbd/dynamics/multibody/multibody_solver.rs b/src_rbd/dynamics/multibody/multibody_solver.rs index f52a414a..4270ba68 100644 --- a/src_rbd/dynamics/multibody/multibody_solver.rs +++ b/src_rbd/dynamics/multibody/multibody_solver.rs @@ -5,14 +5,16 @@ use crate::math::Pose; use crate::queries::GpuIndexedContact; use crate::shaders::dynamics::{ GpuMbApplyContactRestitution, GpuMbBuildContactDelassus, GpuMbComputeDynamicsPre, - GpuMbComputeSolveBounds, GpuMbFinalizeContactConstraints, GpuMbFinalizeImpulseJointConstraints, - GpuMbGravityAndLu, GpuMbGravityAndLuT1, GpuMbGravityAndLuT8, GpuMbGravityAndLuT16, - GpuMbGravityAndLuT32, GpuMbInitContactConstraints, GpuMbInitJointConstraints, GpuMbIntegrate, - GpuMbIntegrateVelocities, GpuMbRemoveImpulseJointConstraintBias, GpuMbSeedContactRestitution, - GpuMbSnapshotContactWarmstart, GpuMbSolveConstraints, GpuMbSolveContactsDelassus, - GpuMbSolveImpulseJointConstraints, GpuMbSolveJoints, GpuMbStashContactsLen, - GpuMbTransferContactWarmstart, GpuMbUpdateImpulseJointConstraints, - GpuMbWarmstartContactConstraints, Velocity, WorldMassProperties, + GpuMbComputeSolveBounds, GpuMbDelayTick, GpuMbFinalizeContactConstraints, + GpuMbFinalizeImpulseJointConstraints, GpuMbFinalizeJointConstraints, GpuMbGravityAndLu, + GpuMbGravityAndLuT1, GpuMbGravityAndLuT8, GpuMbGravityAndLuT16, GpuMbGravityAndLuT32, + GpuMbInitContactConstraints, GpuMbInitJointConstraints, GpuMbIntegrate, + GpuMbIntegrateVelocities, GpuMbRefreshJointConstraints, GpuMbRemoveImpulseJointConstraintBias, + GpuMbSeedContactRestitution, GpuMbSenseContactImpulses, GpuMbSnapshotContactWarmstart, + GpuMbSolveConstraints, GpuMbSolveContactsDelassus, GpuMbSolveImpulseJointConstraints, + GpuMbSolveJoints, GpuMbStashContactsLen, GpuMbTransferContactWarmstart, + GpuMbUpdateImpulseJointConstraints, GpuMbWarmstartContactConstraints, Velocity, + WorldMassProperties, }; use crate::shaders::utils::BatchIndices; use khal::Shader; @@ -32,7 +34,18 @@ pub struct GpuMultibodySolver { gravity_and_lu_t32: GpuMbGravityAndLuT32, compute_dynamics_pre: GpuMbComputeDynamicsPre, init_joint_with_bias: GpuMbInitJointConstraints, + /// M⁻¹-column back-solve for the joint constraints, split from the build + /// pass so each fits 8 storage buffers. + finalize_joint_constraints: GpuMbFinalizeJointConstraints, init_contact_constraints: GpuMbInitContactConstraints, + /// Cheap per-substep refresh of the joint rhs / limit activity, used when + /// the full build runs only once per step. + refresh_joint_constraints: GpuMbRefreshJointConstraints, + /// Advances the actuator-delay step counter, once per physics step. + delay_tick: GpuMbDelayTick, + /// Contact force-sensor readout, dispatched once per step after the last + /// substep's stabilization sweep and only when sensors are configured. + sense_contact_impulses: GpuMbSenseContactImpulses, finalize_contact_constraints: GpuMbFinalizeContactConstraints, /// Fused joint+contact PGS iteration (one workgroup per multibody, shared- /// memory dof velocities). @@ -224,14 +237,33 @@ impl GpuMultibodySolver { return Ok(()); } - // Full rebuild of the joint + contact constraints every substep. - self.build_contact_constraints( - encoder, - timestamps.as_deref_mut(), - mb, - args, - first_substep, - )?; + // Full rebuild of the joint + contact constraints. With the refresh + // cadences off, every column-derived quantity is a per-step constant, + // so the full build runs only on the first substep and each later one + // just refreshes the joint rhs / limit activity / accumulated impulse + // from the integrated joint positions. + if mb.implicit_coriolis || mb.substep_refresh || mb.substep_refresh_light || first_substep { + self.build_contact_constraints( + encoder, + timestamps.as_deref_mut(), + mb, + args, + first_substep, + )?; + } else if mb.has_joint_constraints { + let mut pass = encoder.begin_pass("[RBD] mbb/refresh-joint", timestamps.as_deref_mut()); + self.refresh_joint_constraints.call( + &mut pass, + [mb.multibodies_per_batch * MB_LU_LANES, mb.num_batches, 1], + &mb.multibody_info, + &mb.links_static, + &mb.links_workspace, + &mut mb.joint_constraints, + &mb.motor_delay_state, + &mb.constraint_softness, + args.batch_indices, + )?; + } // Carry the previous frame's impulses over before anything reads them. if first_substep && mb.warmstart_coefficient != 0.0 { @@ -302,6 +334,18 @@ impl GpuMultibodySolver { return Ok(()); } + // Actuator delay: advance the step counter before anything reads it, + // so every substep of this step sees the same tick. + if first_substep { + let mut pass = encoder.begin_pass("[RBD] mbb/delay-tick", timestamps.as_deref_mut()); + self.delay_tick.call( + &mut pass, + mb.num_batches, + &mut mb.motor_delay_state, + &mb.motor_delay_params, + )?; + } + // Joint limit/motor constraints: one 64-lane workgroup per multibody // (lane 0 emits the metadata serially). if mb.has_joint_constraints { @@ -313,14 +357,30 @@ impl GpuMultibodySolver { &mb.multibody_info, &mb.links_static, &mb.links_workspace, - &mb.mass_matrices, - &mb.lu_pivots, &mut mb.joint_constraints, - &mut mb.joint_constraint_columns, &mb.dof_couplings, + &mb.motor_delay_state, + &mb.dof_state, &mb.constraint_softness, args.batch_indices, )?; + drop(pass); + + // The M⁻¹-column back-solve is a separate dispatch: one kernel + // binding both the emission inputs and the LU factors would exceed + // the 8-storage-buffer budget. + let mut pass = + encoder.begin_pass("[RBD] mbb/finalize-joint", timestamps.as_deref_mut()); + self.finalize_joint_constraints.call( + &mut pass, + init_joint_dispatch, + &mb.multibody_info, + &mut mb.joint_constraints, + &mut mb.joint_constraint_columns, + &mb.mass_matrices, + &mb.lu_pivots, + args.batch_indices, + )?; } // One 64-lane workgroup per multibody. @@ -585,8 +645,10 @@ impl GpuMultibodySolver { )?; // Recompute the dynamics (FK, mass matrices, LU, accelerations) for - // the next substep. - if !is_last_substep { + // the next substep. With implicit Coriolis off and the refresh cadence + // relaxed, the mass matrix is refreshed once per step instead, which is + // the main win of disabling implicit Coriolis in the first place. + if !is_last_substep && (mb.implicit_coriolis || mb.substep_refresh) { self.compute_dynamics(pass, mb, args)?; } @@ -648,6 +710,31 @@ impl GpuMultibodySolver { Ok(()) } + /// Contact force-sensor readout: folds each sensed link's normal-contact + /// impulses into `contact_sensor_out`. Run it once per step, after the last + /// substep's stabilization sweep and before [`Self::apply_restitution`], so + /// the value is the accumulated contact impulse rather than a + /// restitution-adjusted one. A no-op when no sensors are configured. + pub fn sense_contact_impulses( + &self, + pass: &mut GpuPass, + mb: &mut GpuMultibodySet, + args: &mut MultibodySolverArgs<'_>, + ) -> Result<(), GpuBackendError> { + if mb.is_empty() || mb.num_contact_sensors == 0 { + return Ok(()); + } + self.sense_contact_impulses.call( + pass, + [mb.multibodies_per_batch, mb.num_batches, 1], + &mb.multibody_info, + &mb.contact_constraints, + &mb.contact_sensor_links, + &mut mb.contact_sensor_out, + args.batch_indices, + ) + } + /// End-of-step restitution pass, run once after the last substep. pub fn apply_restitution( &self, @@ -684,6 +771,15 @@ impl GpuMultibodySolver { // Fused FK + body-jacobians + velocity propagation + Mass-matrix // assembly. Packed: `64 / mb_pack_lanes` multibodies per workgroup, // flattened (multibody, batch) grid. + // The kernels read implicit-coriolis out of `batch_indices`, not off + // this struct, so a flag set without rebuilding that uniform silently + // keeps running the (much more expensive) Coriolis path. + // `RbdState::set_implicit_coriolis` keeps the two in step. + debug_assert_eq!( + mb.implicit_coriolis, mb.coriolis_in_uniform, + "implicit-coriolis changed without rebuilding batch_indices: use \ + RbdState::set_implicit_coriolis, not GpuMultibodySet's" + ); let pre_dispatch = mb.packed_wg_dispatch(); self.compute_dynamics_pre.call( pass, diff --git a/src_rbd/dynamics/solver.rs b/src_rbd/dynamics/solver.rs index f1acaa6f..28c08271 100644 --- a/src_rbd/dynamics/solver.rs +++ b/src_rbd/dynamics/solver.rs @@ -559,6 +559,7 @@ impl GpuSolver { } } + mb_phase!("[RBD] slv/mb-sense-contacts", sense_contact_impulses); mb_phase!("[RBD] slv/mb-restitution", apply_restitution); /* diff --git a/src_rbd/pipeline/bench_narrow_phase.rs b/src_rbd/pipeline/bench_narrow_phase.rs new file mode 100644 index 00000000..5dda6284 --- /dev/null +++ b/src_rbd/pipeline/bench_narrow_phase.rs @@ -0,0 +1,100 @@ +//! Headless timing harness for the batched narrow phase. +//! +//! Builds `num_envs` copies of a small scene (the many-small-environments shape +//! the flat dispatch targets) and reports wall-clock ms per step. Run with +//! `cargo test -p nexus_rbd3d --features metal bench_narrow_phase -- --nocapture --ignored`. + +use crate::pipeline::{RbdCapacities, RbdPipeline, RbdState}; +use crate::rapier::prelude::*; +use crate::shaders::dynamics::RbdSimParams; +use khal::backend::{Backend, GpuBackend}; + +/// Metal when it is available, else the default WebGPU device. Note that the +/// WebGPU path currently trips the 8-storage-buffer limit in +/// `gpu_mb_init_joint_constraints` when the pipeline is created. +async fn bench_backend() -> GpuBackend { + #[cfg(feature = "metal")] + { + GpuBackend::Metal(khal::backend::metal::Metal::new().unwrap()) + } + #[cfg(not(feature = "metal"))] + { + GpuBackend::WebGpu(khal::backend::WebGpu::default().await.unwrap()) + } +} + +/// One environment: a ground cuboid plus `num_boxes` stacked dynamic cuboids. +fn build_env(num_boxes: usize) -> (RigidBodySet, ColliderSet) { + let mut bodies = RigidBodySet::new(); + let mut colliders = ColliderSet::new(); + + let ground = bodies.insert(RigidBodyBuilder::fixed().translation(Vec3::new(0.0, -0.5, 0.0))); + colliders.insert_with_parent(ColliderBuilder::cuboid(5.0, 0.5, 5.0), ground, &mut bodies); + + for i in 0..num_boxes { + let y = 0.6 + i as f32 * 0.45; + let handle = bodies.insert(RigidBodyBuilder::dynamic().translation(Vec3::new(0.0, y, 0.0))); + colliders.insert_with_parent(ColliderBuilder::cuboid(0.2, 0.2, 0.2), handle, &mut bodies); + } + + (bodies, colliders) +} + +async fn run_bench(num_envs: u32, num_boxes: usize, num_steps: u32) { + let backend = bench_backend().await; + + let envs: Vec<_> = (0..num_envs).map(|_| build_env(num_boxes)).collect(); + let joints = ImpulseJointSet::new(); + let mb_joints = MultibodyJointSet::new(); + let params = RbdSimParams::tgs_soft(); + let refs: Vec<_> = envs + .iter() + .map(|(b, c)| (b, c, &joints, &mb_joints, ¶ms)) + .collect(); + + let capacities = RbdCapacities { + batches: num_envs, + // ~5 colliders per env: keep the pair-keyed buffers small enough + // to stay under wgpu's 256 MiB max-buffer-size at 4096 envs. + collisions_capacity: 32, + ..Default::default() + }; + let mut state = RbdState::from_rapier(&backend, &refs, capacities); + let pipeline = RbdPipeline::new(&backend).unwrap(); + + // Warm up: buffer growth and coloring settle over the first few steps. + for _ in 0..20 { + pipeline.step(&backend, &mut state, None).unwrap(); + } + backend.synchronize().unwrap(); + + let start = web_time::Instant::now(); + for _ in 0..num_steps { + pipeline.step(&backend, &mut state, None).unwrap(); + } + backend.synchronize().unwrap(); + let elapsed = start.elapsed(); + + let per_step = elapsed.as_secs_f64() * 1000.0 / num_steps as f64; + println!( + "envs={num_envs:5} boxes/env={num_boxes} -> {per_step:8.3} ms/step \ + ({:.0} env-steps/s)", + num_envs as f64 / (per_step / 1000.0) + ); +} + +#[futures_test::test] +#[serial_test::serial] +#[ignore] +async fn bench_narrow_phase_sweep() { + // The 4096-env rung needs a buffer past wgpu's default 256 MiB cap, so it + // only runs on backends with a higher limit. + #[cfg(feature = "metal")] + let sweep: &[u32] = &[1, 64, 256, 1024, 4096]; + #[cfg(not(feature = "metal"))] + let sweep: &[u32] = &[1, 64, 256, 1024]; + + for &envs in sweep { + run_bench(envs, 4, 200).await; + } +} diff --git a/src_rbd/pipeline/insertion_removal.rs b/src_rbd/pipeline/insertion_removal.rs index a19e0425..f8bf0991 100644 --- a/src_rbd/pipeline/insertion_removal.rs +++ b/src_rbd/pipeline/insertion_removal.rs @@ -48,7 +48,6 @@ impl RbdState { let num_solver_iterations = 4u32; let mut base_sim_params = RbdSimParams::default(); base_sim_params.dt /= num_solver_iterations as f32; - let all_sim_params = vec![base_sim_params; num_batches as usize]; // Inactive (padding) slots use empty collision groups so the broad-phase // never matches them with anything. @@ -95,14 +94,14 @@ impl RbdState { let joints = GpuImpulseJointSet::from_rapier_filtered(backend, &joint_env_refs, &[], &[]); #[cfg(feature = "dim3")] - let multibodies = { + let mut multibodies = { let empty_mb = MultibodyJointSet::new(); let empty_bodies = RigidBodySet::new(); 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, capacity_per_batch); - mb.set_constraint_softness(backend, &all_sim_params[0]); + mb.set_constraint_softness(backend, &base_sim_params); mb }; @@ -243,8 +242,16 @@ impl RbdState { num_batches, num_colliders_per_batch, num_solver_iterations, - sim_params: Tensor::vector(backend, &all_sim_params, BufferUsages::STORAGE).unwrap(), + sim_params: Tensor::scalar( + backend, + base_sim_params, + BufferUsages::UNIFORM | BufferUsages::COPY_DST, + ) + .unwrap(), + sim_params_cpu: base_sim_params, vels: Tensor::vector(backend, &all_vels, rw).unwrap(), + #[cfg(feature = "dim3")] + reset_templates_bodies: None, solver_vels: Tensor::vector(backend, &all_vels, storage).unwrap(), solver_vels_inc: Tensor::vector(backend, &all_vels, storage).unwrap(), joints, diff --git a/src_rbd/pipeline/mod.rs b/src_rbd/pipeline/mod.rs index 4875c540..2b80368d 100644 --- a/src_rbd/pipeline/mod.rs +++ b/src_rbd/pipeline/mod.rs @@ -4,11 +4,15 @@ //! simulation step on the GPU. The pipeline manages collision detection, contact generation, //! constraint solving, and integration. +#[cfg(all(test, feature = "dim3"))] +mod bench_narrow_phase; mod insertion_removal; mod lbvh_validation; mod rbd_state; mod rbd_state_from_rapier; mod rbd_step; +#[cfg(feature = "dim3")] +pub use rbd_state::RbdSnapshot; pub use rbd_state::{RbdCapacities, RbdResizePolicy, RbdState, RunStats}; -pub use rbd_step::RbdPipeline; +pub use rbd_step::{FORCE_FUSED_SWEEPS, RbdPipeline}; diff --git a/src_rbd/pipeline/rbd_state.rs b/src_rbd/pipeline/rbd_state.rs index ac691fac..0ee03857 100644 --- a/src_rbd/pipeline/rbd_state.rs +++ b/src_rbd/pipeline/rbd_state.rs @@ -123,6 +123,9 @@ pub struct RbdState { pub(super) num_colliders_per_batch: u32, pub(super) num_solver_iterations: u32, pub(super) sim_params: Tensor, + /// CPU mirror of `sim_params`, so runtime setters can patch one field and + /// re-upload without rebuilding the whole state. + pub(super) sim_params_cpu: RbdSimParams, /// Per-body world-origin pose (matches rapier's `RigidBody::position`). pub(super) body_poses: Tensor, /// Per-body COM-centered pose (rapier's `SolverPose`) used temporarily by @@ -131,6 +134,10 @@ pub struct RbdState { pub(super) local_mprops: Tensor, pub(super) mprops: Tensor, pub(super) vels: Tensor, + /// GPU-resident rigid-body reset templates, published by + /// `publish_reset_templates` and consumed by `reset_envs_from_templates`. + #[cfg(feature = "dim3")] + pub(super) reset_templates_bodies: Option, pub(super) solver_vels: Tensor, pub(super) solver_vels_inc: Tensor, pub(super) vertex_buffers: Tensor, @@ -160,6 +167,8 @@ pub struct RbdState { /// Single-element scratch holding the max of `collision_pairs_len` across all /// batches, computed on the GPU (only used when `num_batches > 1`). pub(super) collision_pairs_len_max: Tensor, + /// Cosine of the maximum angle between two contact normals for their + /// manifolds to be clustered together (see `gpu_reduce_contacts`). /// `num_batches` as a uniform, the scan length for the max reduction. pub(super) num_batches_uniform: Tensor, /// Non-blocking readback of `[max collision_pairs_len, uncolored]` used by @@ -338,6 +347,18 @@ impl RbdState { self.collision_pairs_len_cpu } + /// GPU buffer of the broad-phase collision pairs found this step. + pub fn collision_pairs(&self) -> &Tensor { + &self.collision_pairs + } + + /// GPU buffer holding the per-batch collision-pair counts. Unlike + /// [`Self::collision_pairs_len`], which returns the CPU mirror from the + /// last readback, this is the value the current step wrote. + pub fn collision_pairs_len_gpu(&self) -> &Tensor { + &self.collision_pairs_len + } + /// The max number a collision pairs the state can currently store. pub fn collision_pairs_capacity(&self) -> u32 { self.collision_pairs.capacity() as u32 @@ -351,6 +372,20 @@ impl RbdState { self.gravity = Self::gravity_tensor(backend, gravity); } + /// Sets how nearly parallel two contact normals must be for their + /// manifolds to be clustered by `gpu_reduce_contacts`, as a cosine. + /// + /// Defaults to [`COS_MERGE_ANGLE`](crate::shaders::broad_phase::COS_MERGE_ANGLE) + /// (~5.1 degrees), matching rapier. Pass `-1.0` to merge every manifold of + /// a collider pair regardless of normal: cheaper, but a single averaged + /// normal then stands in for a ridge or a step edge. + pub fn set_contact_merge_cos(&mut self, backend: &GpuBackend, cos: f32) { + let mut params = self.sim_params_cpu; + params.contact_merge_cos = cos; + let _ = backend.write_buffer(self.sim_params.buffer_mut(), 0, &[params]); + self.sim_params_cpu = params; + } + /// The gravity uniform shared by every solver kernel. pub fn gravity(&self) -> &Tensor { &self.gravity @@ -395,6 +430,14 @@ impl RbdState { self.rebuild_batch_indices(backend); } + /// Sets the per-DoF dry joint friction (N·m). + #[cfg(feature = "dim3")] + pub fn set_dof_frictionloss(&mut self, backend: &GpuBackend, values: &[f32]) { + self.multibodies.set_dof_frictionloss(backend, values); + self.rebuild_batch_indices(backend); + self.multibodies.constraint_caps_dirty = false; + } + /// Returns a reference to the GPU buffer containing collision shapes. /// /// Each shape corresponds to one rigid body in the simulation. @@ -561,3 +604,230 @@ pub(super) fn world_mprops_from_local( } } } + +/// GPU-resident rigid-body reset templates (see +/// [`RbdState::publish_reset_templates`]). +#[cfg(feature = "dim3")] +pub(super) struct ResetTemplatesBodies { + poses: Tensor, + vels: Tensor, + mask: Tensor, + kernel: crate::shaders::dynamics::GpuEnvResetBodies, +} + +/// CPU snapshot of one (single-batch) physics template: body poses, velocities +/// and the multibody joint-space state, read off the GPU once so per-env resets +/// need no readback. See [`RbdState::snapshot`]. +#[cfg(feature = "dim3")] +#[derive(Clone)] +pub struct RbdSnapshot { + body_poses: Vec, + vels: Vec, + mb: crate::dynamics::GpuMultibodySnapshot, +} + +#[cfg(feature = "dim3")] +impl RbdSnapshot { + /// A copy with every floating-base multibody translated by `offset`: the + /// affected links' `body_poses` plus the multibody workspace (root + /// free-joint coords, local-to-parent, per-link local-to-world). Fixed + /// bodies (ground, terrain) and velocities are untouched. + pub fn translated(&self, offset: crate::math::Vector) -> RbdSnapshot { + let mut out = self.clone(); + out.mb = self.mb.translated(offset); + self.mb.for_each_link_rb_id(|rb_id| { + if let Some(p) = out.body_poses.get_mut(rb_id as usize) { + p.translation += offset; + } + }); + out + } +} + +#[cfg(feature = "dim3")] +impl RbdState { + /// Reads this (template) physics state off the GPU into a CPU snapshot. + /// Call it once per template at setup and pass the result to + /// [`Self::reset_env_from_snapshot`] for readback-free per-env resets. + pub async fn snapshot(&self, backend: &GpuBackend) -> RbdSnapshot { + let mut body_poses = bytemuck::zeroed_vec(self.body_poses.len() as usize); + backend + .slow_read_buffer(self.body_poses.buffer(), &mut body_poses) + .await + .unwrap(); + let mut vels = bytemuck::zeroed_vec(self.vels.len() as usize); + backend + .slow_read_buffer(self.vels.buffer(), &mut vels) + .await + .unwrap(); + let mb = self.multibodies.snapshot(backend).await; + RbdSnapshot { + body_poses, + vels, + mb, + } + } + + /// Resets env `dst_env` from a CPU snapshot using `write_buffer` only. + pub fn reset_env_from_snapshot( + &mut self, + backend: &GpuBackend, + dst_env: u32, + snap: &RbdSnapshot, + ) { + let nb = self.num_batches as u64; + let bps = (self.body_poses.len() / nb) as usize; + backend + .write_buffer( + self.body_poses.buffer_mut(), + dst_env as u64 * bps as u64, + &snap.body_poses[..bps], + ) + .unwrap(); + let vs = (self.vels.len() / nb) as usize; + backend + .write_buffer( + self.vels.buffer_mut(), + dst_env as u64 * vs as u64, + &snap.vels[..vs], + ) + .unwrap(); + self.multibodies + .reset_env_from_snapshot(backend, dst_env, &snap.mb); + } + + /// [`Self::reset_env_from_snapshot`] with the robot rigidly translated by + /// `offset` (world frame): the teleport primitive for terrain-curriculum + /// spawn placement. Only floating-base multibody links move; fixed bodies + /// keep their snapshot poses. Costs one single-env-sized snapshot clone per + /// call, so prefer [`Self::reset_envs_from_templates`] in reset loops. + pub fn reset_env_from_snapshot_offset( + &mut self, + backend: &GpuBackend, + dst_env: u32, + snap: &RbdSnapshot, + offset: crate::math::Vector, + ) { + let moved = snap.translated(offset); + self.reset_env_from_snapshot(backend, dst_env, &moved); + } + + /// Uploads the reset templates once (rigid-body poses and velocities here, + /// the multibody blobs via + /// [`GpuMultibodySet::publish_reset_templates`][mb]), enabling the batched + /// [`Self::reset_envs_from_templates`]. + /// + /// [mb]: crate::dynamics::GpuMultibodySet::publish_reset_templates + pub fn publish_reset_templates(&mut self, backend: &GpuBackend, snaps: &[&RbdSnapshot]) { + use crate::shaders::dynamics::GpuEnvResetBodies; + use khal::Shader as _; + if snaps.is_empty() { + return; + } + let nb = self.num_batches as usize; + let bps = self.body_poses.len() as usize / nb; + let vs = self.vels.len() as usize / nb; + let storage = BufferUsages::STORAGE | BufferUsages::COPY_DST; + + let mut poses = Vec::with_capacity(snaps.len() * bps); + let mut vels = Vec::with_capacity(snaps.len() * vs); + for snap in snaps { + poses.extend_from_slice(&snap.body_poses[..bps]); + vels.extend_from_slice(&snap.vels[..vs]); + } + // The bodies a teleport offset applies to: free-multibody links, per + // `RbdSnapshot::translated`. Ground and terrain stay put. + let mut mask = vec![0u32; bps]; + snaps[0].mb.for_each_link_rb_id(|rb_id| { + if let Some(m) = mask.get_mut(rb_id as usize) { + *m = 1; + } + }); + + /// `#[derive(Shader)]` supplies `from_backend` for the embedded entry. + #[derive(khal::Shader)] + struct EnvResetBodiesShader { + kernel: GpuEnvResetBodies, + } + let shader = EnvResetBodiesShader::from_backend(backend).unwrap(); + self.reset_templates_bodies = Some(ResetTemplatesBodies { + poses: Tensor::vector(backend, &poses, storage).unwrap(), + vels: Tensor::vector(backend, &vels, storage).unwrap(), + mask: Tensor::vector(backend, &mask, storage).unwrap(), + kernel: shader.kernel, + }); + let mb_snaps: Vec<&crate::dynamics::GpuMultibodySnapshot> = + snaps.iter().map(|s| &s.mb).collect(); + self.multibodies.publish_reset_templates(backend, &mb_snaps); + } + + /// Batched reset: restores every `(dst_env, template)` in `resets` from the + /// GPU-resident templates, translated by the matching `offsets` entry, with + /// `dof_vels` (`dofs_per_batch` floats per reset, a randomized reset draw or + /// zeros) written into the generalized-velocity section. + /// + /// One compact upload, two dispatches and one submit for the whole batch, + /// replacing the per-env snapshot clone, staging uploads and strided + /// velocity writes. [`Self::publish_reset_templates`] must have run first. + pub fn reset_envs_from_templates( + &mut self, + backend: &GpuBackend, + resets: &[(u32, u32)], + offsets: &[crate::math::Vector], + dof_vels: &[f32], + ) { + use glamx::{UVec4, Vec4}; + use khal::backend::Encoder as _; + let n = resets.len() as u32; + if n == 0 { + return; + } + let nb = self.num_batches; + let bps = self.body_poses.len() as u32 / nb; + let vs = self.vels.len() as u32 / nb; + let meta: Vec = resets + .iter() + .map(|&(env, t)| UVec4::new(env, t, 0, 0)) + .collect(); + let offs: Vec = offsets + .iter() + .map(|o| Vec4::new(o.x, o.y, o.z, 0.0)) + .collect(); + let storage = BufferUsages::STORAGE | BufferUsages::COPY_DST; + let t_meta = Tensor::vector(backend, &meta, storage).unwrap(); + let t_offs = Tensor::vector(backend, &offs, storage).unwrap(); + let params = Tensor::scalar( + backend, + UVec4::new(bps, vs, n, 0), + BufferUsages::STORAGE | BufferUsages::UNIFORM | BufferUsages::COPY_DST, + ) + .unwrap(); + + let tpl = self + .reset_templates_bodies + .take() + .expect("publish_reset_templates must run first"); + let mut enc = backend.begin_encoding(); + { + let mut pass = enc.begin_pass("[RBD] env-reset-bodies", None); + tpl.kernel + .call( + &mut pass, + [bps.max(vs), n, 1], + &tpl.poses, + &tpl.vels, + &tpl.mask, + &t_meta, + &t_offs, + &mut self.body_poses, + &mut self.vels, + ¶ms, + ) + .unwrap(); + } + self.multibodies + .encode_reset_envs_batch(backend, &mut enc, &meta, &offs, dof_vels); + backend.submit(enc).unwrap(); + self.reset_templates_bodies = Some(tpl); + } +} diff --git a/src_rbd/pipeline/rbd_state_from_rapier.rs b/src_rbd/pipeline/rbd_state_from_rapier.rs index 042de1c3..7c51e93f 100644 --- a/src_rbd/pipeline/rbd_state_from_rapier.rs +++ b/src_rbd/pipeline/rbd_state_from_rapier.rs @@ -87,6 +87,11 @@ impl RbdState { (env 0 has {}, env {env} has {})", sp0.num_solver_iterations, sp.num_solver_iterations ); + assert!( + sp == sp0, + "batched rbd requires identical simulation parameters in every \ + environment (env {env} differs from env 0)" + ); } } @@ -147,25 +152,33 @@ impl RbdState { // the equal-topology invariant and to set `BatchIndices::bodies_len`. let mut all_env_body_counts: Vec = Vec::new(); let mut shape_buffers = ShapeBuffers::default(); + // A `SharedShape` cloned across envs (e.g. shared terrain) is + // serialized into `shape_buffers` once, keyed by its parry data + // pointer, and every clone reuses the resulting `Shape` descriptor. + let mut trimesh_cache: HashMap = HashMap::new(); let mut joint_envs: Vec<( &ImpulseJointSet, HashMap, )> = Vec::new(); - // Collect per-batch sim params, adjusting dt for substeps. - let num_solver_iterations = environments - .iter() - .map(|(_, _, _, _, sp)| sp.num_solver_iterations) - .max() - .unwrap_or(4); - let all_sim_params: Vec = environments - .iter() + // Simulation parameters are global: every batch shares one struct, + // bound as a uniform. Identical across environments by the invariant + // asserted above, so the first environment's is authoritative. + let sim_params = environments + .first() .map(|(_, _, _, _, sp)| { let mut sp = **sp; sp.dt /= sp.num_solver_iterations as f32; sp }) - .collect(); + .unwrap_or_else(|| { + let mut sp = RbdSimParams::default(); + sp.dt /= sp.num_solver_iterations as f32; + sp + }); + // Unchanged by the dt division above, so this is every environment's + // solver-iteration count (and 4, the default, when there are none). + let num_solver_iterations = sim_params.num_solver_iterations; // Pick representative dt (outer dt, not the per-substep one) from any batch. #[cfg(feature = "dim3")] let multibody_dt = environments @@ -291,9 +304,24 @@ impl RbdState { } }; - all_shapes.push( - shape_from_parry(co.shape(), &mut shape_buffers).expect("Unsupported shape"), - ); + let gpu_shape = match co.shape().as_typed_shape() { + crate::parry::shape::TypedShape::TriMesh(tm) => { + let key = tm as *const _ as *const u8 as usize; + match trimesh_cache.get(&key) { + Some(&s) => s, + None => { + let s = shape_from_parry(co.shape(), &mut shape_buffers) + .expect("Unsupported shape"); + trimesh_cache.insert(key, s); + s + } + } + } + _ => { + shape_from_parry(co.shape(), &mut shape_buffers).expect("Unsupported shape") + } + }; + all_shapes.push(gpu_shape); all_collider_local_poses.push(collider_local_pose); all_collision_groups.push(co.collision_groups()); all_collider_materials.push(collider_material_from_rapier(co)); @@ -448,7 +476,7 @@ impl RbdState { // Convert multibodies (3D only). #[cfg(feature = "dim3")] - let multibodies = { + let mut multibodies = { let mb_refs: Vec<( &MultibodyJointSet, &HashMap, @@ -465,7 +493,7 @@ impl RbdState { // Soft contact coefficients (rapier TGS-soft) from the substep sim // params, so multibody-vs-floor contacts use the same soft ERP + CFM // as the free-body path (and as rapier) instead of a rigid `1/dt`. - mb.set_constraint_softness(backend, &all_sim_params[0]); + mb.set_constraint_softness(backend, &sim_params); // Route MB-touching impulse joints (those skipped by the // regular `GpuImpulseJointSet`) to the multibody generic @@ -753,8 +781,16 @@ impl RbdState { num_batches, num_colliders_per_batch: num_colliders_per_batch as u32, num_solver_iterations, - sim_params: Tensor::vector(backend, &all_sim_params, BufferUsages::STORAGE).unwrap(), - vels: Tensor::vector(backend, &all_vels, storage).unwrap(), + sim_params: Tensor::scalar( + backend, + sim_params, + BufferUsages::UNIFORM | BufferUsages::COPY_DST, + ) + .unwrap(), + sim_params_cpu: sim_params, + vels: Tensor::vector(backend, &all_vels, storage | BufferUsages::COPY_DST).unwrap(), + #[cfg(feature = "dim3")] + reset_templates_bodies: None, solver_vels: Tensor::vector(backend, &all_vels, storage).unwrap(), solver_vels_inc: Tensor::vector(backend, &all_vels, storage).unwrap(), joints, @@ -767,7 +803,7 @@ impl RbdState { body_poses: Tensor::vector( backend, &all_poses, - BufferUsages::STORAGE | BufferUsages::COPY_SRC, + BufferUsages::STORAGE | BufferUsages::COPY_SRC | BufferUsages::COPY_DST, ) .unwrap(), // Sized like `body_poses`. Will be (re-)seeded each step before diff --git a/src_rbd/pipeline/rbd_step.rs b/src_rbd/pipeline/rbd_step.rs index d208f328..39afd1c1 100644 --- a/src_rbd/pipeline/rbd_step.rs +++ b/src_rbd/pipeline/rbd_step.rs @@ -18,6 +18,12 @@ use khal::backend::{Backend, Encoder, GpuBackend, GpuBackendError, GpuTimestamps use vortx::Reduce; use vortx::tensor::Tensor; +/// Forces the fused colored-sweep kernels regardless of the estimated pair +/// count: the programmatic twin of `NEXUS_FUSED_SWEEPS=1`, for targets without +/// environment variables (wasm). +pub static FORCE_FUSED_SWEEPS: core::sync::atomic::AtomicBool = + core::sync::atomic::AtomicBool::new(false); + /// The main GPU physics pipeline coordinating all simulation stages. pub struct RbdPipeline { mprops_update: GpuMpropsUpdate, @@ -32,6 +38,9 @@ pub struct RbdPipeline { coloring: GpuColoring, warmstart: GpuWarmstart, reduce: Reduce, + /// Optional (default `false`): merge each collider pair's manifolds + /// (e.g. per-triangle trimesh contacts) into one before the solvers. + pub contact_reduction: bool, } impl RbdPipeline { @@ -54,6 +63,7 @@ impl RbdPipeline { coloring: GpuColoring::from_backend(backend)?, warmstart: GpuWarmstart::from_backend(backend)?, reduce: Reduce::from_backend(backend)?, + contact_reduction: false, }) } @@ -61,13 +71,65 @@ impl RbdPipeline { /// /// Automatically resizes buffers (next power of two) if collision pair count exceeds capacity. pub fn step( + &self, + backend: &GpuBackend, + state: &mut RbdState, + timestamps: Option<&mut GpuTimestamps>, + ) -> Result { + let mut encoder = backend.begin_encoding(); + let stats = self.step_impl(backend, state, timestamps, &mut encoder, true)?; + backend.submit(encoder)?; + Ok(stats) + } + + /// Records one timestep into a caller-owned `encoder` instead of managing + /// (and submitting) its own. Nothing is submitted: the caller submits, so + /// its own dispatches can share the command buffer with the physics step. + /// + /// The intra-step submits `step` uses to overlap CPU encoding with GPU work + /// are skipped here; WebGPU guarantees dispatch-order visibility within a + /// single encoder, so the step stays correct without them. + pub fn step_encoded( + &self, + backend: &GpuBackend, + state: &mut RbdState, + timestamps: Option<&mut GpuTimestamps>, + encoder: &mut ::Encoder, + ) -> Result { + self.step_impl(backend, state, timestamps, encoder, false) + } + + fn step_impl( &self, backend: &GpuBackend, state: &mut RbdState, mut timestamps: Option<&mut GpuTimestamps>, + encoder: &mut ::Encoder, + allow_splits: bool, ) -> Result { + // Submit what is recorded so far and start a fresh encoder, so CPU + // encoding overlaps GPU work. A no-op in encoded mode, where everything + // stays in the caller's encoder. + let split = |enc: &mut ::Encoder| -> Result<(), GpuBackendError> { + if allow_splits { + let done = std::mem::replace(enc, backend.begin_encoding()); + backend.submit(done)?; + } + Ok(()) + }; let mut stats = RunStats::default(); + // A multibody capacity edit (e.g. reserving the dry-friction + // constraint slots on the first `set_dof_frictionloss`) leaves the + // shared `BatchIndices` uniform describing the old buffer sizes, so + // the kernels would index the resized buffers with stale per-batch + // capacities. Re-upload it before anything reads it. + #[cfg(feature = "dim3")] + if state.multibodies.constraint_caps_dirty { + state.rebuild_batch_indices(backend); + state.multibodies.constraint_caps_dirty = false; + } + // Make sure the color index uniforms are up-to-date. // This is the maximum over the colors needed for contacts, joints, and multibodies. { @@ -80,8 +142,6 @@ impl RbdPipeline { state.ensure_color_uniforms(backend, needed); } - let mut encoder = backend.begin_encoding(); - // Phase 0: Multibody once-per-visible-step setup (3D only for now). #[cfg(feature = "dim3")] { @@ -99,7 +159,7 @@ impl RbdPipeline { gravity: &state.gravity, }; self.multibody_solver.init_step( - &mut encoder, + &mut *encoder, timestamps.as_deref_mut(), &mut state.multibodies, &mut args, @@ -157,14 +217,15 @@ impl RbdPipeline { &mut state.collision_pairs_indirect, &state.collision_groups, &state.pair_filter, + &state.sim_params, )?; drop(pass); - backend.submit(encoder)?; + split(&mut *encoder)?; } else { // Build LBVH and find collision pairs. self.lbvh.update_tree( backend, - &mut encoder, + &mut *encoder, &mut state.lbvh, state.collider_local_poses.len() as u32, state.num_active_colliders, @@ -177,8 +238,8 @@ impl RbdPipeline { )?; // Debug: validate LBVH topology after tree construction - if crate::VALIDATE_LBVH_TOPOLOGY { - backend.submit(encoder)?; + if crate::VALIDATE_LBVH_TOPOLOGY && allow_splits { + split(&mut *encoder)?; let num_colliders = state.collider_world_poses.len() as u32; let tree: Vec = futures::executor::block_on( @@ -189,7 +250,6 @@ impl RbdPipeline { )?; validate_lbvh_topology(&tree, &sorted_colliders, num_colliders); - encoder = backend.begin_encoding(); let _pass = encoder .begin_pass("[RBD] broad-phase-find-pairs", timestamps.as_deref_mut()); } @@ -210,7 +270,7 @@ impl RbdPipeline { )?; drop(pass); - backend.submit(encoder)?; + split(&mut *encoder)?; } } @@ -223,9 +283,16 @@ impl RbdPipeline { state.collision_pairs_per_batch_cpu }; - // Choose the kernel depending on the expected pairs count. - // Small pairs with many environment benefit from the fused kernels. - let fused_color_sweeps = est_pairs <= 128; + // Choose the kernel depending on the expected pairs count: small pair + // counts with many environments benefit from the fused kernels. The + // fused path can also be forced regardless of size (an A/B knob: an env + // var natively, [`FORCE_FUSED_SWEEPS`] on wasm where env vars do not + // exist). It is correct at any size, just serialized past ~64 lanes, + // which may still win where per-dispatch latency rules, i.e. small + // batch counts in the browser. + let fused_color_sweeps = est_pairs <= 128 + || FORCE_FUSED_SWEEPS.load(core::sync::atomic::Ordering::Relaxed) + || std::env::var("NEXUS_FUSED_SWEEPS").as_deref() == Ok("1"); // In small scenes, submit less frequently. In big scenes submit more // to overlap compute and encoding. @@ -234,7 +301,6 @@ impl RbdPipeline { // Phase 2a: Narrow phase. Split out from solver-prep + coloring // so its CPU encoding overlaps with Phase 1's GPU work and its // own GPU work overlaps with Phase 2b's CPU encoding. - let mut encoder = backend.begin_encoding(); { let mut pass = encoder.begin_pass("[RBD] narrow-phase", timestamps.as_deref_mut()); @@ -258,12 +324,13 @@ impl RbdPipeline { &state.batch_indices, &state.collider_parent, &state.collider_materials, + &state.sim_params, + self.contact_reduction, )?; drop(pass); if !merge_submits { - backend.submit(encoder)?; - encoder = backend.begin_encoding(); + split(&mut *encoder)?; } } @@ -418,8 +485,7 @@ impl RbdPipeline { drop(pass); } if !merge_submits { - backend.submit(encoder)?; - encoder = backend.begin_encoding(); + split(&mut *encoder)?; } } @@ -485,7 +551,7 @@ impl RbdPipeline { Some((&self.multibody_solver, &mut state.multibodies)) }; self.solver.solve_tgs( - &mut encoder, + &mut *encoder, timestamps.as_deref_mut(), &self.joint_solver, solver_args, @@ -496,9 +562,9 @@ impl RbdPipeline { // Resolve all accumulated timestamps before the final submit. if let Some(ts) = ×tamps { - ts.resolve(&mut encoder); + ts.resolve(&mut *encoder); } - backend.submit(encoder)?; + split(&mut *encoder)?; } // Swap buffers for warm-starting next frame diff --git a/src_rbd_shaders/broad_phase/brute_force.rs b/src_rbd_shaders/broad_phase/brute_force.rs index 94601b1a..5b095e9f 100644 --- a/src_rbd_shaders/broad_phase/brute_force.rs +++ b/src_rbd_shaders/broad_phase/brute_force.rs @@ -10,14 +10,13 @@ use khal_std::sync::atomic_add_u32; use crate::bounding_volumes::Aabb; use crate::broad_phase::CollisionPair; +use crate::dynamics::RbdSimParams; use crate::shapes::Shape; use crate::utils::BatchIndices; use crate::{PaddedVector, Pose, Vector}; use glamx::UVec2; use rapier::geometry::InteractionGroups; -use super::narrow_phase::PREDICTION; - /// Computes every active collider's world AABB. #[spirv_bindgen] #[spirv(compute(threads(64)))] @@ -59,6 +58,7 @@ pub fn gpu_bf_find_pairs( collision_groups: &[InteractionGroups], #[spirv(uniform, descriptor_set = 0, binding = 4)] batch_ids: &BatchIndices, #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] pair_filter: &[[u32; 2]], + #[spirv(uniform, descriptor_set = 0, binding = 6)] params: &RbdSimParams, ) { let n = batch_ids.colliders_len; let nn = n * n; @@ -91,7 +91,7 @@ pub fn gpu_bf_find_pairs( let coll_start = batch_ids.coll_start(batch_id); // Dilate one side by the contact prediction distance. let mut aabb_i = aabbs.read(coll_start + i as usize); - let dilation = Vector::splat(PREDICTION); + let dilation = Vector::splat(params.prediction_distance()); aabb_i.mins -= dilation; aabb_i.maxs += dilation; let aabb_j = aabbs.read(coll_start + j as usize); diff --git a/src_rbd_shaders/broad_phase/narrow_phase.rs b/src_rbd_shaders/broad_phase/narrow_phase.rs index 86d0e1e0..8213d85b 100644 --- a/src_rbd_shaders/broad_phase/narrow_phase.rs +++ b/src_rbd_shaders/broad_phase/narrow_phase.rs @@ -2,10 +2,13 @@ //! //! Computes contact manifolds from collision pairs detected by the broad phase. +use crate::dynamics::RbdSimParams; use crate::queries::{ ColliderMaterial, ContactManifold, IndexedManifold, ball_ball, ball_convex, convex_ball, cuboid_cuboid, pfm_pfm, }; +#[cfg(feature = "dim3")] +use crate::queries::{ContactPoint, MAX_MANIFOLD_POINTS, manifold_reduction}; use crate::shapes::{ Capsule, Polyline, SHAPE_TYPE_BALL, SHAPE_TYPE_CAPSULE, SHAPE_TYPE_CONE, SHAPE_TYPE_CUBOID, SHAPE_TYPE_CYLINDER, SHAPE_TYPE_POLYLINE, SHAPE_TYPE_TRIMESH, Shape, TriMesh, @@ -67,7 +70,154 @@ pub fn gpu_narrow_phase_init_contacts_dispatch( } } -pub(crate) const PREDICTION: f32 = 2.0e-2; // TODO: make the prediction configurable. +/// Default cluster threshold: normals must agree within ~5.1 degrees, +/// matching rapier's `contact_clustering::COS_MERGE_ANGLE`. Passed as a +/// uniform so it can be loosened (`-1` merges every manifold of a pair, +/// whatever its normal) to trade contact fidelity for solver cost. +pub const COS_MERGE_ANGLE: f32 = 0.996; + +/// Pools `pt` into `cand`, deduplicating against points already there. +/// +/// Composite shapes emit near-coincident points on both sides of a shared +/// triangle edge; rapier's clustering collapses those within a quarter of the +/// prediction distance, keeping the deeper one. Same rule here. +#[cfg(feature = "dim3")] +#[inline] +fn pool_dedup(cand: &mut [ContactPoint; 8], num: &mut usize, pt: ContactPoint, dedup_eps_sq: f32) { + let mut hit = false; + for k in 0..*num { + let d = cand.read(k).pt - pt.pt; + if !hit && d.dot(d) < dedup_eps_sq { + if pt.dist < cand.read(k).dist { + cand.write(k, pt); + } + hit = true; + } + } + if !hit && *num < 8 { + cand.write(*num, pt); + *num += 1; + } +} + +/// Optional contact reduction: compacts each batch's contacts in place by +/// merging manifolds that share both a collider pair and a (nearly) parallel +/// normal into a single `MAX_MANIFOLD_POINTS` manifold. This mirrors rapier's +/// `cluster_manifolds_for_solver` + `reduce_manifold_naive`: cluster by +/// normal, deduplicate near-coincident points, then keep the deepest point, +/// the point furthest from it, and the two tangent extremes. +/// +/// Per-triangle trimesh contacts share one `colliders` key and one collider-A +/// local frame, so a flat patch collapses to one manifold while a ridge keeps +/// one cluster per face. The first record of a cluster is kept verbatim, so +/// single-manifold pairs are bit-identical to the unreduced path. +/// +/// Two deliberate divergences from rapier. Clusters are reduced incrementally +/// at each merge against an 8-point pool, where rapier accumulates every point +/// (up to 255) and reduces once, so the selection here depends on manifold +/// emission order. And the cluster's normal comes from the deepest point +/// rather than from the manifold that opened it: identical in effect at the +/// default threshold, where every member is within ~5.1 degrees, but it keeps +/// the choice sane when `merge_cos` is loosened. +/// +/// Grid `[1, num_batches, 1]`, serial per batch. +#[cfg(feature = "dim3")] +#[spirv_bindgen] +#[spirv(compute(threads(1)))] +pub fn gpu_reduce_contacts( + #[spirv(workgroup_id)] workgroup_id: UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] contacts: &mut [IndexedManifold], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] contacts_len: &mut [u32], + #[spirv(uniform, descriptor_set = 0, binding = 2)] batch_ids: &BatchIndices, + #[spirv(uniform, descriptor_set = 0, binding = 3)] params: &RbdSimParams, +) { + let prediction = params.prediction_distance(); + let merge_cos = params.contact_merge_cos; + let batch_id = workgroup_id.y; + let capacity = batch_ids.contacts_batch_capacity as usize; + let mut contacts = batch_ids.contact_batch_mut(batch_id, contacts); + let n = (contacts_len.read(batch_id as usize) as usize).min(capacity); + + // Write cursor: always <= the read cursor, so compacting in place is safe. + let mut w = 0usize; + for i in 0..n { + let im = contacts[i]; + let mut merged = false; + for j in 0..w { + let out = contacts[j]; + if out.colliders.x == im.colliders.x + && out.colliders.y == im.colliders.y + && out.contact.normal_a.dot(im.contact.normal_a) >= merge_cos + { + // Pool the two manifolds' points (same collider-A local frame), + // dropping near-duplicates as rapier's clustering does. + let na = (out.contact.len as usize).min(MAX_MANIFOLD_POINTS); + let nb = (im.contact.len as usize).min(MAX_MANIFOLD_POINTS); + let dedup_eps = prediction * 0.25; + let dedup_eps_sq = dedup_eps * dedup_eps; + let mut cand = [ContactPoint::default(); 8]; + let mut num = 0usize; + for k in 0..na { + pool_dedup( + &mut cand, + &mut num, + out.contact.points_a.read(k), + dedup_eps_sq, + ); + } + for k in 0..nb { + pool_dedup( + &mut cand, + &mut num, + im.contact.points_a.read(k), + dedup_eps_sq, + ); + } + // Normal of whichever manifold holds the deepest point. rapier + // keeps the opener's normal instead, which it can afford + // because its ~5.1 degree cone makes every member equivalent; + // this degrades gracefully when `merge_cos` is loosened, and + // agrees with rapier's choice when it is not. + let mut deep_out = out.contact.points_a.at(0).dist; + for k in 1..na { + let d = out.contact.points_a.at(k).dist; + if d < deep_out { + deep_out = d; + } + } + let mut deep_in = im.contact.points_a.at(0).dist; + for k in 1..nb { + let d = im.contact.points_a.at(k).dist; + if d < deep_in { + deep_in = d; + } + } + let normal = if deep_in < deep_out { + im.contact.normal_a + } else { + out.contact.normal_a + }; + let mut reduced = manifold_reduction(&cand, num as u32, normal, prediction); + // `manifold_reduction` fills points/len only. + reduced.normal_a = normal; + let mut kept = out; + kept.contact = reduced; + contacts[j] = kept; + merged = true; + break; + } + } + if !merged { + contacts[w] = im; + w += 1; + } + } + // Compacted count; plain store, single writer per batch. (Loop shell per + // the `gpu_reset_narrow_phase` rustgpu-triviality workaround.) + for _ in 0..1 { + contacts_len.write(batch_id as usize, w as u32); + } +} /// Narrow phase, pass 1 of 2: analytic shape-shape contacts for ball / cuboid /// pairs, written straight into the `contacts` buffer. @@ -92,7 +242,9 @@ pub fn gpu_narrow_phase_shape_shape( #[spirv(storage_buffer, descriptor_set = 0, binding = 7)] collider_parent: &[u32], #[spirv(storage_buffer, descriptor_set = 0, binding = 8)] collider_materials: &[ColliderMaterial], + #[spirv(uniform, descriptor_set = 0, binding = 9)] params: &RbdSimParams, ) { + let prediction = params.prediction_distance(); let num_threads = num_workgroups.x * WORKGROUP_SIZE; let batch_id = invocation_id.y; let contacts_batch_capacity = batch_ids.contacts_batch_capacity as usize; @@ -113,9 +265,12 @@ pub fn gpu_narrow_phase_shape_shape( for i in StepRng::new(invocation_id.x..len, num_threads) { let pair = collision_pairs[i as usize]; // Resolve the parent rigid-bodies here (the broad phase no longer does) - // and skip pairs whose colliders share the same body. - let body1 = collider_parent.read(pair.colliders.x as usize); - let body2 = collider_parent.read(pair.colliders.y as usize); + // and skip pairs whose colliders share the same body. Pair ids are + // env-local and `collider_parent` is batch-strided, so the stride is + // required: without it every batch reads batch 0's parents. + let coll_base = batch_ids.coll_start(batch_id); + let body1 = collider_parent.read(coll_base + pair.colliders.x as usize); + let body2 = collider_parent.read(coll_base + pair.colliders.y as usize); if body1 == body2 { continue; } @@ -159,12 +314,12 @@ pub fn gpu_narrow_phase_shape_shape( if shape_ty1 == SHAPE_TYPE_CUBOID && shape_ty2 == SHAPE_TYPE_CUBOID { let cuboid1 = shape1.to_cuboid(); let cuboid2 = shape2.to_cuboid(); - manifold = cuboid_cuboid(pose12, &cuboid1, &cuboid2, PREDICTION); + manifold = cuboid_cuboid(pose12, &cuboid1, &cuboid2, prediction); } // Everything else (PFM / trimesh / polyline) is handled by the deferred // pass; `manifold.len` stays 0 here so nothing is written. - if manifold.len > 0 && manifold.points_a.at(0).dist < PREDICTION { + if manifold.len > 0 && manifold.points_a.at(0).dist < prediction { let target_contact_index = atomic_add_u32(contacts_len, 1) as usize; // NOTE: if we exceed the contacts allocation size, just skip @@ -208,7 +363,9 @@ pub fn gpu_narrow_phase_shape_shape_deferred( // And we assume all batch dimensions are given the same buffer allocation sizes // (i.e. the same `contacts_batch_capacity`). #[spirv(uniform, descriptor_set = 0, binding = 6)] batch_ids: &BatchIndices, + #[spirv(uniform, descriptor_set = 0, binding = 7)] params: &RbdSimParams, ) { + let prediction = params.prediction_distance(); let num_threads = num_workgroups.x * WORKGROUP_SIZE; let batch_id = invocation_id.y; let contacts_batch_capacity = batch_ids.contacts_batch_capacity as usize; @@ -299,6 +456,7 @@ pub fn gpu_narrow_phase_shape_shape_deferred( let mesh = shape1.to_trimesh(); let convex = shape2; trimesh_convex( + prediction, pose12, &mesh, convex, @@ -317,6 +475,7 @@ pub fn gpu_narrow_phase_shape_shape_deferred( let mesh = shape2.to_trimesh(); // NOTE: pair indices are flipped. trimesh_convex( + prediction, pose12.inverse(), &mesh, convex, @@ -336,6 +495,7 @@ pub fn gpu_narrow_phase_shape_shape_deferred( let pline = shape1.to_polyline(); let convex = shape2; polyline_convex( + prediction, pose12, &pline, convex, @@ -354,6 +514,7 @@ pub fn gpu_narrow_phase_shape_shape_deferred( let pline = shape2.to_polyline(); // NOTE: pair indices are flipped. polyline_convex( + prediction, pose12.inverse(), &pline, convex, @@ -371,6 +532,7 @@ pub fn gpu_narrow_phase_shape_shape_deferred( /// Collision detection between a triangle mesh and a convex shape. fn trimesh_convex( + prediction: f32, pose12: Pose, mesh: &TriMesh, convex: &Shape, @@ -387,10 +549,10 @@ fn trimesh_convex( return; } - // Get the convex shape's AABB in the trimesh's local space, and enlarge with the PREDICTION. + // Get the convex shape's AABB in the trimesh's local space, and enlarge with the prediction distance. let mut test_aabb = convex.compute_aabb(pose12, vertices); - test_aabb.mins -= Vector::splat(PREDICTION); - test_aabb.maxs += Vector::splat(PREDICTION); + test_aabb.mins -= Vector::splat(prediction); + test_aabb.maxs += Vector::splat(prediction); if !test_aabb.intersects(&mesh.root_aabb) { // No collision possible. @@ -441,6 +603,7 @@ fn trimesh_convex( /// Collision detection between a polyline and a convex shape. fn polyline_convex( + prediction: f32, pose12: Pose, mesh: &Polyline, convex: &Shape, @@ -457,11 +620,11 @@ fn polyline_convex( return; } - // Get the convex shape's AABB in the polyline's local space, and enlarge with the PREDICTION. + // Get the convex shape's AABB in the polyline's local space, and enlarge with the prediction distance. let thickness = 0.4; // TODO: make thickness configurable or part of the polyline struct let mut test_aabb = convex.compute_aabb(pose12, vertices); - test_aabb.mins -= Vector::splat(PREDICTION + thickness); - test_aabb.maxs += Vector::splat(PREDICTION + thickness); + test_aabb.mins -= Vector::splat(prediction + thickness); + test_aabb.maxs += Vector::splat(prediction + thickness); if !test_aabb.intersects(&mesh.root_aabb) { // No collision possible. @@ -559,7 +722,9 @@ pub fn gpu_narrow_phase_pfm_pfm( #[spirv(storage_buffer, descriptor_set = 0, binding = 7)] collider_parent: &[u32], #[spirv(storage_buffer, descriptor_set = 0, binding = 8)] collider_materials: &[ColliderMaterial], + #[spirv(uniform, descriptor_set = 0, binding = 9)] params: &RbdSimParams, ) { + let prediction = params.prediction_distance(); let num_threads = num_workgroups.x * WORKGROUP_SIZE; let batch_id = invocation_id.y; let contacts_batch_capacity = batch_ids.contacts_batch_capacity as usize; @@ -580,8 +745,9 @@ pub fn gpu_narrow_phase_pfm_pfm( // is where the deferred (PFM / trimesh / polyline) pairs get the same-body // filtering that the analytic pass does inline — the broad phase no longer // does it, and the deferred pass has no spare storage binding for it. - let body1 = collider_parent.read(pair.colliders.x as usize); - let body2 = collider_parent.read(pair.colliders.y as usize); + let coll_base = batch_ids.coll_start(batch_id); + let body1 = collider_parent.read(coll_base + pair.colliders.x as usize); + let body2 = collider_parent.read(coll_base + pair.colliders.y as usize); if body1 == body2 { continue; } @@ -591,13 +757,13 @@ pub fn gpu_narrow_phase_pfm_pfm( pair.thickness1, &pair.shape2, pair.thickness2, - PREDICTION, + prediction, vertices, #[cfg(feature = "dim3")] indices, ); - if manifold.len > 0 && manifold.points_a.at(0).dist < PREDICTION { + if manifold.len > 0 && manifold.points_a.at(0).dist < prediction { let target_contact_index = atomic_add_u32(contacts_len, 1) as usize; // NOTE: if we exceed capacity, just skip the pair. diff --git a/src_rbd_shaders/dynamics/joint_constraint.rs b/src_rbd_shaders/dynamics/joint_constraint.rs index 04ba1e2f..e32a4d8e 100644 --- a/src_rbd_shaders/dynamics/joint_constraint.rs +++ b/src_rbd_shaders/dynamics/joint_constraint.rs @@ -17,7 +17,6 @@ use khal_std::macros::{spirv, spirv_bindgen}; use crate::Pose; use crate::utils::{BatchIndices, Slice}; -use khal_std::index::MaybeIndexUnchecked; use super::body::{LocalMassProperties, Velocity, WorldMassProperties}; use super::joint::ImpulseJoint; @@ -223,12 +222,11 @@ pub fn gpu_update_joint_constraints( #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] constraints: &mut [JointConstraint], #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] poses: &[Pose], #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] mprops: &[WorldMassProperties], - #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] all_params: &[RbdSimParams], + #[spirv(uniform, descriptor_set = 0, binding = 4)] params: &RbdSimParams, #[spirv(uniform, descriptor_set = 0, binding = 5)] batch_ids: &BatchIndices, ) { let num_threads = num_workgroups.x * WORKGROUP_SIZE; let batch_id = invocation_id.y; - let params = all_params.at(batch_id as usize); let builders = batch_ids.impulse_joints_batch(batch_id, builders); let mut constraints = batch_ids.impulse_joints_batch_mut(batch_id, constraints); diff --git a/src_rbd_shaders/dynamics/multibody/contact_sensor.rs b/src_rbd_shaders/dynamics/multibody/contact_sensor.rs new file mode 100644 index 00000000..03adc551 --- /dev/null +++ b/src_rbd_shaders/dynamics/multibody/contact_sensor.rs @@ -0,0 +1,68 @@ +//! Contact "force sensor" readout for RL observations. + +use super::types::{ + MAX_MB_CONTACT_CONSTRAINTS_PER_MB, MB_CONTACT_KIND_NORMAL, MultibodyContactConstraint, + MultibodyInfo, +}; +use crate::utils::BatchIndices; +use khal_std::glamx::UVec3; +use khal_std::index::MaybeIndexUnchecked; +use khal_std::macros::{spirv, spirv_bindgen}; + +/// Maximum sensed links per multibody for the contact force-sensor readout. +pub const MAX_CONTACT_SENSORS: u32 = 4; + +/// Per sensed link, sums the accumulated normal-constraint impulses. Dispatch +/// it once per step, after the last substep's stabilization sweep: the value is +/// then the step's total accumulated normal impulse (divide by the step `dt` +/// for an average force) when the constraints are built once per step, or the +/// last substep's impulse when they are rebuilt per substep. +/// +/// Slots whose sensed link has no active normal rows read exactly 0.0: the +/// kernel zeroes its slots before accumulating, so no host-side clear pass is +/// needed and the dispatch is graph-capture safe. +/// +/// `contact_sensor_links` holds `MAX_CONTACT_SENSORS` multibody link ids +/// (`u32::MAX` marks an unused slot); the same set is sensed for every +/// multibody in every batch. The output is interleaved like the other per-mb +/// buffers: `contact_sensor_out[batch_ids.mbi(batch, mb_idx) * +/// MAX_CONTACT_SENSORS + slot]`. +#[spirv_bindgen] +#[spirv(compute(threads(1)))] +pub fn gpu_mb_sense_contact_impulses( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] multibody_info: &[MultibodyInfo], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] + contact_constraints: &[MultibodyContactConstraint], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] contact_sensor_links: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] contact_sensor_out: &mut [f32], + #[spirv(uniform, descriptor_set = 0, binding = 4)] batch_ids: &BatchIndices, +) { + let batch_id = invocation_id.y; + let mb_idx = invocation_id.x; + let out_base = batch_ids.mbi(batch_id, mb_idx as usize) * (MAX_CONTACT_SENSORS as usize); + for s in 0..MAX_CONTACT_SENSORS { + contact_sensor_out.write(out_base + s as usize, 0.0); + } + if mb_idx >= batch_ids.multibodies_len { + return; + } + + let mb = multibody_info.read(batch_ids.mbi(batch_id, mb_idx as usize)); + let cons_start = batch_ids.mb_contact_constraints_start(batch_id); + let cons_base = cons_start + (mb_idx as usize) * (MAX_MB_CONTACT_CONSTRAINTS_PER_MB as usize); + let count = mb.contact_constraint_count; + + for c in 0..count { + let cons = contact_constraints.read(cons_base + c as usize); + if cons.kind != MB_CONTACT_KIND_NORMAL { + continue; + } + for s in 0..MAX_CONTACT_SENSORS { + if contact_sensor_links.read(s as usize) == cons.link_id { + let cur = contact_sensor_out.read(out_base + s as usize); + contact_sensor_out.write(out_base + s as usize, cur + cons.impulse); + } + } + } +} diff --git a/src_rbd_shaders/dynamics/multibody/env_reset.rs b/src_rbd_shaders/dynamics/multibody/env_reset.rs new file mode 100644 index 00000000..712df56f --- /dev/null +++ b/src_rbd_shaders/dynamics/multibody/env_reset.rs @@ -0,0 +1,227 @@ +//! Per-environment reset scatter for RL teleport / reset primitives. +//! +//! Copies one environment's carry-over multibody state (SoA link workspace, +//! static link descriptors, generalized coordinates and velocities) from a +//! compact contiguous staging blob into the batch-interleaved live buffers. + +use glamx::{UVec4, Vec4}; +use khal_std::glamx::UVec3; +use khal_std::index::MaybeIndexUnchecked; +use khal_std::macros::{spirv, spirv_bindgen}; + +use super::types::MultibodyLinkStatic; +use super::ws_soa::{WS_COORDS, WS_LTP, WS_LTW, WS_QUADS}; + +/// Scatters one staged env state into the interleaved buffers. Dispatch +/// `[links_per_batch · WS_QUADS, 1, 1]` threads, the largest of the three +/// per-element loops (`links_per_batch · WS_QUADS >= links_per_batch`, and +/// `dofs_per_batch <= links_per_batch · WS_QUADS` for any real multibody). +/// +/// `staging_dofs` holds `dofs_per_batch` generalized coordinates followed by +/// `dofs_per_batch` generalized velocities. Only the velocity section of +/// `dof_state` is written; the sections after it are static configuration +/// (damping, armature, springs), not per-episode state. +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn gpu_mb_env_reset( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] staging_ws: &[Vec4], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] + staging_links: &[MultibodyLinkStatic], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] staging_dofs: &[f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] links_workspace: &mut [Vec4], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] + links_static: &mut [MultibodyLinkStatic], + #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] dof_values: &mut [f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 6)] dof_state: &mut [f32], + // x = dst_env, y = num_batches, z = links_per_batch, w = dofs_per_batch. + #[spirv(uniform, descriptor_set = 0, binding = 7)] params: &UVec4, +) { + let i = invocation_id.x; + let env = params.x; + let nb = params.y; + let lpb = params.z; + let dpb = params.w; + + if i < lpb * WS_QUADS { + links_workspace.write((i * nb + env) as usize, staging_ws.read(i as usize)); + } + if i < lpb { + links_static.write((i * nb + env) as usize, staging_links.read(i as usize)); + } + if i < dpb { + dof_values.write((i * nb + env) as usize, staging_dofs.read(i as usize)); + dof_state.write( + (i * nb + env) as usize, + staging_dofs.read((dpb + i) as usize), + ); + } +} + +/// Batched, template-resident variant of [`gpu_mb_env_reset`]: N resets in one +/// dispatch, reading from template blobs that live on the GPU permanently +/// (uploaded once at build) instead of a per-reset staging upload. A terrain +/// teleport offset is applied in-kernel to the free root's world position +/// (local-to-world / local-to-parent translations plus coords c0..c2), so the +/// host never clones and translates a snapshot per reset. +/// +/// This pass writes the link workspace only; [`gpu_mb_env_reset_batch_dofs`] +/// writes the static links and the DoF sections. They are split so each fits +/// the 8-storage-buffer WebGPU limit. +/// +/// Dispatch `[lpb · WS_QUADS, num_resets, 1]` threads. +/// +/// `link_flags` is constant per robot: bit 0 = a valid link of a free-root +/// multibody (translate `WS_LTW`), bit 1 = that link is the root (translate +/// `WS_LTP` and coords c0..c2). +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn gpu_mb_env_reset_batch( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] templates_ws: &[Vec4], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] link_flags: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] resets: &[UVec4], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] offsets: &[Vec4], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] links_workspace: &mut [Vec4], + // x = num_batches, y = links_per_batch, z = dofs_per_batch, w = num_resets. + #[spirv(uniform, descriptor_set = 0, binding = 5)] params: &UVec4, +) { + let i = invocation_id.x; + let r = invocation_id.y; + let nb = params.x; + let lpb = params.y; + if r >= params.w { + return; + } + let meta = resets.read(r as usize); + let env = meta.x; + let t = meta.y; + let off = offsets.read(r as usize); + + if i < lpb * WS_QUADS { + let mut v = templates_ws.read((t * lpb * WS_QUADS + i) as usize); + let link = i / WS_QUADS; + let q = i % WS_QUADS; + let f = link_flags.read(link as usize); + // `WS_LTW` and `WS_LTP` are rot|trans quad pairs, so `+1` is the + // translation quad. Coords c0..c3 share quad `WS_COORDS` (c3 is a + // rotational DoF, never offset, so its `.w` lane stays untouched). + if f & 1 != 0 && q == WS_LTW + 1 { + v.x += off.x; + v.y += off.y; + v.z += off.z; + } + if f & 2 != 0 && (q == WS_LTP + 1 || q == WS_COORDS) { + v.x += off.x; + v.y += off.y; + v.z += off.z; + } + links_workspace.write((i * nb + env) as usize, v); + } +} + +/// Static-link and DoF half of the batched reset, split from +/// [`gpu_mb_env_reset_batch`] so each pass fits 8 storage buffers. +/// +/// The teleport offset is not needed here: static links carry no world +/// position, and generalized coords are translation-invariant (the free root's +/// world position lives in the workspace coords quad the other pass handles). +/// +/// Dispatch `[max(lpb, dpb), num_resets, 1]` threads. +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn gpu_mb_env_reset_batch_dofs( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + templates_links: &[MultibodyLinkStatic], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] templates_dofs: &[f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] resets: &[UVec4], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] dof_vels: &[f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] + links_static: &mut [MultibodyLinkStatic], + #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] dof_values: &mut [f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 6)] dof_state: &mut [f32], + // x = num_batches, y = links_per_batch, z = dofs_per_batch, w = num_resets. + #[spirv(uniform, descriptor_set = 0, binding = 7)] params: &UVec4, +) { + let i = invocation_id.x; + let r = invocation_id.y; + let nb = params.x; + let lpb = params.y; + let dpb = params.z; + if r >= params.w { + return; + } + let meta = resets.read(r as usize); + let env = meta.x; + let t = meta.y; + + if i < lpb { + links_static.write( + (i * nb + env) as usize, + templates_links.read((t * lpb + i) as usize), + ); + } + if i < dpb { + dof_values.write( + (i * nb + env) as usize, + templates_dofs.read((t * 2 * dpb + i) as usize), + ); + dof_state.write( + (i * nb + env) as usize, + dof_vels.read((r * dpb + i) as usize), + ); + } +} + +/// Rigid-body half of the batched reset: copies each reset env's `body_poses` +/// and `vels` slices (env-major, unlike the interleaved multibody buffers) +/// from the resident templates, adding the teleport offset to the poses of the +/// bodies flagged in `body_mask` (free-multibody links; ground and terrain stay +/// put). +/// +/// Dispatch `[max(bodies_per_env, vels_per_env), num_resets, 1]` threads. +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn gpu_env_reset_bodies( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] templates_poses: &[crate::Pose], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] + templates_vels: &[crate::dynamics::body::Velocity], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] body_mask: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] resets: &[UVec4], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] offsets: &[Vec4], + #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] body_poses: &mut [crate::Pose], + #[spirv(storage_buffer, descriptor_set = 0, binding = 6)] + vels: &mut [crate::dynamics::body::Velocity], + // x = bodies_per_env, y = vels_per_env, z = num_resets. + #[spirv(uniform, descriptor_set = 0, binding = 7)] params: &UVec4, +) { + let i = invocation_id.x; + let r = invocation_id.y; + let bps = params.x; + let vs = params.y; + if r >= params.z { + return; + } + let meta = resets.read(r as usize); + let env = meta.x; + let t = meta.y; + let off = offsets.read(r as usize); + + if i < bps { + let mut p = templates_poses.read((t * bps + i) as usize); + if body_mask.read(i as usize) != 0 { + p.translation.x += off.x; + p.translation.y += off.y; + p.translation.z += off.z; + } + body_poses.write((env * bps + i) as usize, p); + } + if i < vs { + vels.write( + (env * vs + i) as usize, + templates_vels.read((t * vs + i) 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 f328c9cf..fc005f5f 100644 --- a/src_rbd_shaders/dynamics/multibody/gravity_and_lu.rs +++ b/src_rbd_shaders/dynamics/multibody/gravity_and_lu.rs @@ -283,7 +283,8 @@ pub fn gpu_mb_gravity_and_lu( if i < ndofs { let idx = batch_ids.mbi(batch_id, gen_base + i as usize); let cur = gen_forces.read(idx); - gen_forces.write(idx, cur - damping_slice[i as usize] * vel_slice[i as usize]); + let v = vel_slice[i as usize]; + gen_forces.write(idx, cur - damping_slice[i as usize] * v); } workgroup_memory_barrier_with_group_sync(); @@ -596,7 +597,8 @@ fn gravity_and_lu_packed_impl usize { + 2 + batch_ids.links_batch_capacity as usize +} + +/// The motor position target to track this substep. +/// +/// `tick` counts the physics steps begun since the host last refreshed the +/// delay state, incremented once per step by `gpu_mb_delay_tick`; it is stable +/// for a whole step, so every substep of a step agrees. While it is at most the +/// batch's delay `k`, the motor tracks the previous control step's target +/// instead of the current one, modelling actuator latency with no mid-step host +/// writes. An all-zero delay buffer (the default) leaves `target` untouched: +/// the first step's `tick` is 1, already past `k = 0`. +#[inline] +fn delayed_motor_target( + motor_delay_state: &[f32], + batch_ids: &BatchIndices, + batch_id: u32, + link_id: u32, + target: f32, +) -> f32 { + let base = batch_id as usize * motor_delay_stride(batch_ids); + let tick = motor_delay_state.read(base); + let delay_k = motor_delay_state.read(base + 1); + if tick <= delay_k { + motor_delay_state.read(base + 2 + link_id as usize) + } else { + target + } +} + /// Compute joint motor parameters mirroring rapier's `JointMotor::motor_params`. #[inline] fn motor_params(motor: &crate::dynamics::joint::JointMotor, dt: f32) -> (f32, f32, f32, f32, f32) { @@ -75,12 +110,14 @@ fn emit_joint_constraints( links_workspace: &[Vec4], dof_couplings: &[MbDofCoupling], joint_constraints: &mut [MultibodyJointConstraint], + dof_state: &[f32], mb: &MultibodyInfo, cons_base: usize, batch_id: u32, dt: f32, joint_erp_inv_dt: f32, joint_cfm_coeff: f32, + motor_delay_state: &[f32], batch_ids: &BatchIndices, ) { let num_links = mb.num_links; @@ -131,6 +168,13 @@ fn emit_joint_constraints( inv_dt, dt, stat.data.motors.at(axis as usize), + delayed_motor_target( + motor_delay_state, + batch_ids, + batch_id, + mb.first_link + k, + stat.data.motors.read(axis as usize).target_pos, + ), has_limits, limit_min, limit_max, @@ -193,6 +237,13 @@ fn emit_joint_constraints( inv_dt, dt, stat.data.motors.at(axis as usize), + delayed_motor_target( + motor_delay_state, + batch_ids, + batch_id, + mb.first_link + k, + stat.data.motors.read(axis as usize).target_pos, + ), has_limits, limit_min, limit_max, @@ -225,6 +276,29 @@ fn emit_joint_constraints( joint_constraints.write(cons_base + slot as usize, cons); slot += 1; } + + // Joint dry friction (MJCF `frictionloss`): one box-bounded row per DoF + // that has a non-zero loss. Purely DoF-indexed, so no link walk is needed. + // The frictionloss section is all-zero unless the host reserved the extra + // slots through `RbdState::set_dof_frictionloss`, so this emits nothing + // (and cannot overflow `max_constraints`) by default. + let dof_cap = batch_ids.dof_batch_capacity as usize; + let frictionloss_slice = batch_ids + .ib(batch_id, dof_state) + .offset(6 * dof_cap + mb.first_dof as usize); + let kin_mask_slice = batch_ids + .ib(batch_id, dof_state) + .offset(5 * dof_cap + mb.first_dof as usize); + for d in 0..mb.ndofs { + let fl = frictionloss_slice.read(d as usize); + // Kinematic DoFs have a prescribed velocity; a friction row would + // fight it. + if fl > 0.0 && kin_mask_slice.read(d as usize) == 0.0 { + let cons = build_friction_constraint(d, fl, dt, joint_cfm_coeff); + joint_constraints.write(cons_base + slot as usize, cons); + slot += 1; + } + } } /// Solve `M · column = J` (writes the `M⁻¹·J` column) and return the raw @@ -357,6 +431,43 @@ fn build_coupling_constraint( } } +/// Initialize a single dry-friction constraint slot (MJCF `frictionloss`). +/// +/// The row has no position residual: it simply drives the DoF velocity to zero +/// with an impulse clamped to `±frictionloss·dt`, which is MuJoCo's friction +/// loss (a load-independent force bound, unlike Coulomb friction). +/// +/// `cfm_coeff` is the shared joint softness (rapier's `joint.softness.cfm_coeff(dt)`). +#[inline] +fn build_friction_constraint( + dof_id: u32, + frictionloss: f32, + dt: f32, + cfm_coeff: f32, +) -> MultibodyJointConstraint { + let max_impulse = frictionloss * dt; + + MultibodyJointConstraint { + dof_id, + kind: MB_JOINT_KIND_FRICTION, + _kind_extra: 0, + dof2_id: 0, + rhs: 0.0, + rhs_wo_bias: 0.0, + inv_lhs: 0.0, + impulse: 0.0, + impulse_lo: -max_impulse, + impulse_hi: max_impulse, + cfm_coeff, + // Folded with the row's `lhs` by the finalize stage. + cfm_gain: 0.0, + coupling_coeff: 0.0, + coupling_offset: 0.0, + _kind_extra2: 0, + _pad1: 0, + } +} + /// Initialize a single motor constraint slot.. #[inline] #[allow(clippy::too_many_arguments)] @@ -368,6 +479,9 @@ fn build_motor_constraint( inv_dt: f32, dt: f32, motor: &crate::dynamics::joint::JointMotor, + // The position target to track, which actuator delay may pull from a + // previous control step (see `delayed_motor_target`). + target_pos: f32, has_limits: bool, limit_min: f32, limit_max: f32, @@ -376,7 +490,7 @@ fn build_motor_constraint( let mut rhs_wo_bias = 0.0f32; if erp_inv_dt != 0.0 { - rhs_wo_bias += (curr_pos - motor.target_pos) * erp_inv_dt; + rhs_wo_bias += (curr_pos - target_pos) * erp_inv_dt; } let mut target_vel = motor.target_vel; @@ -421,11 +535,12 @@ fn build_motor_constraint( /// Must run after `gpu_mb_lu_decompose` — the LU factors of `M` are used to compute /// the per-constraint M⁻¹ column and effective inverse mass. /// -/// One 64-lane workgroup per (multibody, batch), in three stages: +/// One 64-lane workgroup per (multibody, batch), in two stages: /// 1. lane-parallel: zero all constraint slots; -/// 2. lane 0: the serial link walk emitting constraint metadata (cheap); -/// 3. lane-parallel: one M⁻¹-column LU back-solve per emitted slot plus -/// rapier's `finalize_generic_constraints`. +/// 2. lane 0: the serial link walk emitting constraint metadata (cheap). +/// +/// The M⁻¹-column back-solve is a separate dispatch +/// ([`gpu_mb_finalize_joint_constraints`]), so each pass fits 8 storage buffers. #[spirv_bindgen] #[spirv(compute(threads(64)))] pub fn gpu_mb_init_joint_constraints( @@ -435,15 +550,17 @@ pub fn gpu_mb_init_joint_constraints( #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] links_static: &[MultibodyLinkStatic], #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] links_workspace: &[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)] + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] 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(uniform, descriptor_set = 0, binding = 8)] softness: &ConstraintSoftness, - #[spirv(uniform, descriptor_set = 0, binding = 9)] batch_ids: &BatchIndices, + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] dof_couplings: &[MbDofCoupling], + // Actuator-delay state, per-batch `[tick, k, prev_target x links]`. Zeroed + // (the default) means no delay; see `delayed_motor_target`. + #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] motor_delay_state: &[f32], + // Packed per-DoF sections; only the kinematic mask (5) and frictionloss + // (6) ones are read here. + #[spirv(storage_buffer, descriptor_set = 0, binding = 6)] dof_state: &[f32], + #[spirv(uniform, descriptor_set = 0, binding = 7)] softness: &ConstraintSoftness, + #[spirv(uniform, descriptor_set = 0, binding = 8)] batch_ids: &BatchIndices, ) { const LANES: u32 = 64; @@ -468,14 +585,7 @@ pub fn gpu_mb_init_joint_constraints( } let active = in_range && ndofs != 0; - let mb_mm_base = mb.mass_matrix_offset as usize; - let piv = batch_ids.ivec(batch_id, mb.first_dof as usize); let cons_base = batch_ids.mb_joint_constraints_start(batch_id) + mb.first_constraint as usize; - // One column of M⁻¹ per constraint slot . - let dofs_stride = batch_ids.dof_batch_capacity as usize; - let col_base = batch_ids.mb_joint_constraint_columns_start(batch_id) - + (mb.first_constraint as usize) * dofs_stride; - let m = batch_ids.imat(batch_id, mb_mm_base, ndofs, ndofs); // Stage 1: lane-parallel slot reset. if active { @@ -503,26 +613,70 @@ pub fn gpu_mb_init_joint_constraints( links_workspace, dof_couplings, joint_constraints, + dof_state, &mb, cons_base, batch_id, softness.dt, softness.joint_erp_inv_dt, softness.joint_cfm_coeff, + motor_delay_state, batch_ids, ); } +} - control_barrier::< - { khal_std::memory::Scope::Workgroup as u32 }, - { khal_std::memory::Scope::QueueFamily as u32 }, - { - khal_std::memory::Semantics::UNIFORM_MEMORY.bits() - | khal_std::memory::Semantics::ACQUIRE_RELEASE.bits() - }, - >(); +/// Back-solves one M⁻¹ column per emitted joint-constraint slot and applies +/// rapier's `finalize_generic_constraints`. +/// +/// Split from [`gpu_mb_init_joint_constraints`] so that each pass stays within +/// the 8-storage-buffer WebGPU limit. Must run after it, and after +/// `gpu_mb_lu_decompose`, whose LU factors of `M` it consumes. +/// +/// One 64-lane workgroup per (multibody, batch); lanes stride the slots. +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn gpu_mb_finalize_joint_constraints( + #[spirv(workgroup_id)] workgroup_id: UVec3, + #[spirv(local_invocation_id)] local_id: UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] multibody_info: &[MultibodyInfo], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] + joint_constraints: &mut [MultibodyJointConstraint], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] + joint_constraint_columns: &mut [f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] mass_matrices: &[f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] lu_pivots: &[u32], + #[spirv(uniform, descriptor_set = 0, binding = 5)] batch_ids: &BatchIndices, +) { + const LANES: u32 = 64; + + let batch_id = workgroup_id.y; + let mb_idx = workgroup_id.x; + let lane = local_id.x; + let num_mb = batch_ids.multibodies_len; + let in_range = mb_idx < num_mb; + #[cfg(not(feature = "web-compat"))] + if !in_range { + return; + } + let slot = if in_range { mb_idx } else { 0 }; + + let mb = batch_ids.ib(batch_id, multibody_info).read(slot as usize); + let ndofs = mb.ndofs; + #[cfg(not(feature = "web-compat"))] + if ndofs == 0 { + return; + } + let active = in_range && ndofs != 0; + + let mb_mm_base = mb.mass_matrix_offset as usize; + let piv = batch_ids.ivec(batch_id, mb.first_dof as usize); + let cons_base = batch_ids.mb_joint_constraints_start(batch_id) + mb.first_constraint as usize; + let dofs_stride = batch_ids.dof_batch_capacity as usize; + let col_base = batch_ids.mb_joint_constraint_columns_start(batch_id) + + (mb.first_constraint as usize) * dofs_stride; + let m = batch_ids.imat(batch_id, mb_mm_base, ndofs, ndofs); - // Stage 3: lane-parallel finalize. if active { for s in StepRng::new(lane..mb.max_constraints, LANES) { let mut cons = joint_constraints.read(cons_base + s as usize); @@ -550,3 +704,111 @@ pub fn gpu_mb_init_joint_constraints( } } } + +/// Per-substep refresh of the joint limit / motor slots, the cheap alternative +/// to a full rebuild. +/// +/// One 64-lane workgroup per (multibody, batch); lanes stride the slots. +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn gpu_mb_refresh_joint_constraints( + #[spirv(workgroup_id)] workgroup_id: UVec3, + #[spirv(local_invocation_id)] local_id: UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] multibody_info: &[MultibodyInfo], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] + links_static: &[MultibodyLinkStatic], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] links_workspace: &[Vec4], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] + joint_constraints: &mut [MultibodyJointConstraint], + // Actuator-delay state; see `gpu_mb_init_joint_constraints`. + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] motor_delay_state: &[f32], + #[spirv(uniform, descriptor_set = 0, binding = 5)] softness: &ConstraintSoftness, + #[spirv(uniform, descriptor_set = 0, binding = 6)] 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 = batch_ids.ib(batch_id, multibody_info).read(mb_idx as usize); + if mb.ndofs == 0 || mb.max_constraints == 0 { + return; + } + let cons_base = batch_ids.mb_joint_constraints_start(batch_id) + mb.first_constraint as usize; + let stat_slice = batch_ids + .ib(batch_id, links_static) + .offset(mb.first_link as usize); + let wa = WsAddr::new(mb.first_link as usize, batch_ids.num_batches, batch_id); + + let dt = softness.dt; + let inv_dt = if dt != 0.0 { 1.0 / dt } else { 0.0 }; + + for s in StepRng::new(lane..mb.max_constraints, LANES) { + let old = joint_constraints.read(cons_base + s as usize); + // Friction rows are per-step constants except for the accumulated + // impulse, which must restart from zero so the `±frictionloss·dt` + // bound applies per substep rather than per step. + if old.kind == MB_JOINT_KIND_FRICTION { + let mut fresh = old; + fresh.impulse = 0.0; + joint_constraints.write(cons_base + s as usize, fresh); + continue; + } + // Coupling rows are per-step constants; inactive slots stay inactive. + if old.kind != MB_JOINT_KIND_MOTOR + && old.kind != MB_JOINT_KIND_LIMIT + && old.kind != MB_JOINT_KIND_LIMIT_INACTIVE + { + continue; + } + let link_id = old._kind_extra & 0xffff; + let axis = old._kind_extra >> 16; + let stat = &stat_slice[link_id as usize]; + let curr_pos = ws_coord(links_workspace, wa, link_id, axis); + let limit_min = stat.data.limits.read(axis as usize).min; + let limit_max = stat.data.limits.read(axis as usize).max; + + // Rebuild the per-substep fields with the same formulas the full + // emission uses, then graft back the per-step constants (the + // column-derived `inv_lhs` and the folded `cfm_gain`). + let mut fresh = if old.kind == MB_JOINT_KIND_MOTOR { + let locked = stat.data.locked_axes; + let has_limits = (stat.data.limit_axes & !locked & (1 << axis)) != 0; + build_motor_constraint( + old.dof_id, + link_id, + axis, + curr_pos, + inv_dt, + dt, + stat.data.motors.at(axis as usize), + delayed_motor_target( + motor_delay_state, + batch_ids, + batch_id, + mb.first_link + link_id, + stat.data.motors.read(axis as usize).target_pos, + ), + has_limits, + limit_min, + limit_max, + ) + } else { + build_limit_constraint( + old.dof_id, + link_id, + axis, + curr_pos, + [limit_min, limit_max], + softness.joint_erp_inv_dt, + softness.joint_cfm_coeff, + ) + }; + fresh.inv_lhs = old.inv_lhs; + fresh.cfm_gain = old.cfm_gain; + joint_constraints.write(cons_base + s as usize, fresh); + } +} diff --git a/src_rbd_shaders/dynamics/multibody/mod.rs b/src_rbd_shaders/dynamics/multibody/mod.rs index 2b664467..4f63ca57 100644 --- a/src_rbd_shaders/dynamics/multibody/mod.rs +++ b/src_rbd_shaders/dynamics/multibody/mod.rs @@ -15,12 +15,16 @@ mod compute_dynamics_pre; mod contact_constraints; +mod contact_sensor; +#[cfg(feature = "dim3")] +mod env_reset; mod gravity_and_lu; mod impulse_joint_constraints; mod integrate; mod jacobian; mod joint_constraints; mod lu; +mod scatter_motor; mod solve_constraints; mod types; mod utils; @@ -28,10 +32,14 @@ mod ws_soa; pub use compute_dynamics_pre::*; pub use contact_constraints::*; +pub use contact_sensor::*; +#[cfg(feature = "dim3")] +pub use env_reset::*; pub use gravity_and_lu::*; pub use impulse_joint_constraints::*; pub use integrate::*; pub use joint_constraints::*; +pub use scatter_motor::*; pub use solve_constraints::*; pub use types::*; pub use utils::*; diff --git a/src_rbd_shaders/dynamics/multibody/scatter_motor.rs b/src_rbd_shaders/dynamics/multibody/scatter_motor.rs new file mode 100644 index 00000000..c5b9a83d --- /dev/null +++ b/src_rbd_shaders/dynamics/multibody/scatter_motor.rs @@ -0,0 +1,113 @@ +//! GPU motor-target scatter: writes per-(env, actuated-joint) target positions +//! straight into `links_static` on the GPU, replacing a host-side +//! `set_motors` + whole-mirror upload every step. This is what lets an RL +//! policy drive the motors without a host round-trip, and therefore what makes +//! a rollout capturable into a CUDA graph (no per-step host writes). +//! +//! `links_static` is batch-interleaved: link `l` of env `e` lives at +//! `l · num_envs + e`. Targets are row-major `[num_actuated x num_envs]`, +//! element `(j, env)` at `j · num_envs + env`, matching the policy action +//! buffer layout. + +use khal_std::glamx::{UVec3, UVec4}; +use khal_std::index::MaybeIndexUnchecked; +use khal_std::macros::{spirv, spirv_bindgen}; + +use super::types::MultibodyLinkStatic; + +/// One thread per (actuated joint `x`, env `y`). Writes `target_pos` into the +/// matching motor and sets its `motor_axes` bit, like `set_motor` does on the +/// host, but without touching the CPU mirror. +#[spirv_bindgen] +#[spirv(compute(threads(1)))] +pub fn gpu_scatter_motor_targets( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] motor_targets: &[f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] + links_static: &mut [MultibodyLinkStatic], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] actuated_link_ids: &[u32], + #[spirv(uniform, descriptor_set = 0, binding = 3)] num_actuated: &u32, + #[spirv(uniform, descriptor_set = 0, binding = 4)] num_envs: &u32, + #[spirv(uniform, descriptor_set = 0, binding = 5)] axis_id: &u32, +) { + let j = invocation_id.x; + let env = invocation_id.y; + if j >= *num_actuated || env >= *num_envs { + return; + } + let link_id = actuated_link_ids[j as usize]; + // Batch-interleaved links layout. + let global_idx = (link_id * *num_envs + env) as usize; + let target = motor_targets[(j * *num_envs + env) as usize]; + + // The single-iteration loop matches `gpu_lbvh_reset_collision_pairs`: + // rust-gpu sometimes prunes the SPIR-V for kernels it deems trivial, and + // the loop shell keeps the entry point emitted. + for _ in 0..1 { + let link = &mut links_static[global_idx]; + link.data.motors[*axis_id as usize].target_pos = target; + link.data.motor_axes |= 1u32 << *axis_id; + } +} + +/// Per-step actuator-delay state refresh, on device: `tick <- 0`, +/// `k <- k_eff[env]`, and the actuated links' `prev_target` lanes copied from +/// the previous step's motor-target tensor (row-major `[num_actuated x n]`, the +/// same buffer the target scatter consumed last step, read before this step's +/// scatter overwrites it). +/// +/// This replaces a full `stride * n` host rebuild and upload every step with +/// one `[n]` upload (`k_eff`) plus this dispatch. Non-actuated `prev` lanes keep +/// their existing value, so held joints stay wherever the host last put them. +/// +/// Dispatch `[num_actuated, num_envs, 1]` threads; lane `j == 0` also writes the +/// two scalar lanes. +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn gpu_mb_delay_state_update( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] prev_targets: &[f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] k_eff: &[f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] actuated_link_ids: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] delay_state: &mut [f32], + // x = num_actuated, y = num_envs, z = stride (2 + links_per_batch). + #[spirv(uniform, descriptor_set = 0, binding = 4)] params: &UVec4, +) { + let j = invocation_id.x; + let env = invocation_id.y; + if j >= params.x || env >= params.y { + return; + } + let base = (env * params.z) as usize; + if j == 0 { + delay_state.write(base, 0.0); + delay_state.write(base + 1, k_eff.read(env as usize)); + } + let link = actuated_link_ids.read(j as usize); + delay_state.write( + base + 2 + link as usize, + prev_targets.read((j * params.y + env) as usize), + ); +} + +/// Advances the actuator-delay step counter by one, for every batch. +/// +/// Dispatched once per physics step, before the joint constraints are built, so +/// the tick is stable across all of that step's substeps regardless of the +/// constraint-refresh cadence. Dispatch `[num_envs, 1, 1]` threads. +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn gpu_mb_delay_tick( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] delay_state: &mut [f32], + // x = num_envs, y = stride (2 + links_per_batch). + #[spirv(uniform, descriptor_set = 0, binding = 1)] params: &UVec4, +) { + let env = invocation_id.x; + if env >= params.x { + return; + } + let base = (env * params.y) as usize; + let tick = delay_state.read(base); + delay_state.write(base, tick + 1.0); +} diff --git a/src_rbd_shaders/dynamics/multibody/solve_constraints.rs b/src_rbd_shaders/dynamics/multibody/solve_constraints.rs index 3110926f..f12109a4 100644 --- a/src_rbd_shaders/dynamics/multibody/solve_constraints.rs +++ b/src_rbd_shaders/dynamics/multibody/solve_constraints.rs @@ -14,8 +14,8 @@ use crate::utils::linalg::MAX_MB_DOFS; use super::types::{ MAX_MB_CONTACT_CONSTRAINTS_PER_MB, MB_CONTACT_KIND_TANGENT, MB_JOINT_KIND_COUPLING, - MB_JOINT_KIND_LIMIT, MB_JOINT_KIND_MOTOR, MultibodyContactConstraint, MultibodyInfo, - MultibodyJointConstraint, + MB_JOINT_KIND_FRICTION, MB_JOINT_KIND_LIMIT, MB_JOINT_KIND_MOTOR, MultibodyContactConstraint, + MultibodyInfo, MultibodyJointConstraint, }; const LANES: u32 = 64; @@ -177,7 +177,8 @@ pub fn gpu_mb_solve_constraints( let solve = slot_active && (cons.kind == MB_JOINT_KIND_LIMIT || cons.kind == MB_JOINT_KIND_MOTOR - || cons.kind == MB_JOINT_KIND_COUPLING); + || cons.kind == MB_JOINT_KIND_COUPLING + || cons.kind == MB_JOINT_KIND_FRICTION); #[cfg(not(feature = "web-compat"))] if !solve { // Unused slot or inactive limit. @@ -473,7 +474,8 @@ pub fn gpu_mb_solve_joints( let solve = slot_active && (cons.kind == MB_JOINT_KIND_LIMIT || cons.kind == MB_JOINT_KIND_MOTOR - || cons.kind == MB_JOINT_KIND_COUPLING); + || cons.kind == MB_JOINT_KIND_COUPLING + || cons.kind == MB_JOINT_KIND_FRICTION); #[cfg(not(feature = "web-compat"))] if !solve { // Unused slot or inactive limit. diff --git a/src_rbd_shaders/dynamics/multibody/types.rs b/src_rbd_shaders/dynamics/multibody/types.rs index 22535308..4dc0ff1e 100644 --- a/src_rbd_shaders/dynamics/multibody/types.rs +++ b/src_rbd_shaders/dynamics/multibody/types.rs @@ -62,6 +62,8 @@ pub const MB_JOINT_KIND_LIMIT_INACTIVE: u32 = 3; /// `J = e_{dof_id} − coeff·e_{dof2_id}` and a bias pulling the position drift /// back to zero. pub const MB_JOINT_KIND_COUPLING: u32 = 4; +/// Joint-constraint `kind`: dry joint friction. +pub const MB_JOINT_KIND_FRICTION: u32 = 5; /// Sentinel marking a link with no parent (the root). pub const MULTIBODY_ROOT: u32 = u32::MAX; diff --git a/src_rbd_shaders/dynamics/multibody/ws_soa.rs b/src_rbd_shaders/dynamics/multibody/ws_soa.rs index 3e05bb5c..9fab1029 100644 --- a/src_rbd_shaders/dynamics/multibody/ws_soa.rs +++ b/src_rbd_shaders/dynamics/multibody/ws_soa.rs @@ -449,3 +449,38 @@ pub fn ws_soa_from_structs( } out } + +/* + * Host-side conversion of the SoA buffer back into the AoS structs, the inverse + * of `ws_soa_from_structs`. Used by observation readbacks, which want one struct + * per link rather than the interleaved quads the kernels index. + */ +#[cfg(not(target_arch_is_gpu))] +pub fn ws_soa_to_structs( + buf: &[Vec4], + links_cap: u32, + num_batches: u32, +) -> std::vec::Vec { + let mut out: std::vec::Vec = + bytemuck::zeroed_vec(links_cap as usize * num_batches as usize); + for b in 0..num_batches { + let a = WsAddr::new(0, num_batches, b); + for k in 0..links_cap { + let ws = &mut out[(b * links_cap + k) as usize]; + ws.joint_rot = ws_rot(buf, a, k, WS_JOINT_ROT); + ws.coords = ws_coords(buf, a, k); + ws.local_to_parent = ws_pose(buf, a, k, WS_LTP); + ws.local_to_world = ws_pose(buf, a, k, WS_LTW); + ws.shift02 = ws_vec(buf, a, k, WS_SHIFT02); + ws.shift23 = ws_vec(buf, a, k, WS_SHIFT23); + ws.joint_velocity = ws_vel(buf, a, k, WS_JOINT_VEL); + ws.rb_vels = ws_vel(buf, a, k, WS_RB_VELS); + ws.kinematic_acc = ws_vel(buf, a, k, WS_KIN_ACC); + let (force, torque, gravity_scale) = ws_ext_wrench(buf, a, k); + ws.external_force = force; + ws.external_torque = torque; + ws.gravity_scale = gravity_scale; + } + } + out +} diff --git a/src_rbd_shaders/dynamics/sim_params.rs b/src_rbd_shaders/dynamics/sim_params.rs index 545c6dff..fbf8d662 100644 --- a/src_rbd_shaders/dynamics/sim_params.rs +++ b/src_rbd_shaders/dynamics/sim_params.rs @@ -70,7 +70,7 @@ impl ConstraintSoftness { } /// Parameters for a time-step of the physics engine. -#[derive(Clone, Copy)] +#[derive(Clone, Copy, PartialEq)] #[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] #[repr(C)] pub struct RbdSimParams { @@ -160,6 +160,13 @@ pub struct RbdSimParams { /// The number of solver iterations run by the constraints solver for calculating forces (default: `4`). pub num_solver_iterations: u32, + + /// Minimum cosine between two manifold normals for the contact-reduction + /// pass to cluster them (default: `0.996`, ~5.1 degrees, matching rapier). + /// + /// `-1.0` merges every manifold of a collider pair regardless of normal: + /// cheaper, but one averaged normal then stands in for a ridge or a step. + pub contact_merge_cos: f32, } impl RbdSimParams { @@ -181,6 +188,7 @@ impl RbdSimParams { normalized_allowed_linear_error: 0.005, normalized_max_corrective_velocity: 3.0, normalized_prediction_distance: 0.02, + contact_merge_cos: crate::broad_phase::COS_MERGE_ANGLE, normalized_max_linear_velocity: 400.0, length_unit: 1.0, } diff --git a/src_rbd_shaders/dynamics/solver.rs b/src_rbd_shaders/dynamics/solver.rs index 596a8fa5..3c596dfa 100644 --- a/src_rbd_shaders/dynamics/solver.rs +++ b/src_rbd_shaders/dynamics/solver.rs @@ -41,12 +41,11 @@ pub fn gpu_solver_init_constraints( #[spirv(storage_buffer, descriptor_set = 1, binding = 1)] solver_body_poses: &[Pose], #[spirv(storage_buffer, descriptor_set = 1, binding = 2)] vels: &[Velocity], #[spirv(storage_buffer, descriptor_set = 1, binding = 3)] mprops: &[WorldMassProperties], - #[spirv(storage_buffer, descriptor_set = 1, binding = 4)] all_params: &[RbdSimParams], + #[spirv(uniform, descriptor_set = 1, binding = 4)] params: &RbdSimParams, #[spirv(uniform, descriptor_set = 1, binding = 5)] batch_ids: &BatchIndices, ) { let num_threads = num_workgroups.x * WORKGROUP_SIZE; let batch_id = invocation_id.y; - let params = all_params.at(batch_id as usize); let contacts = batch_ids.contact_batch(batch_id, contacts); let mut constraints = batch_ids.contact_batch_mut(batch_id, constraints); @@ -132,12 +131,11 @@ pub fn gpu_solver_update_constraints( constraint_builders: &[TwoBodyConstraintBuilder], #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] contacts_len: &[u32], #[spirv(storage_buffer, descriptor_set = 1, binding = 0)] solver_body_poses: &[Pose], - #[spirv(storage_buffer, descriptor_set = 1, binding = 1)] all_params: &[RbdSimParams], + #[spirv(uniform, descriptor_set = 1, binding = 1)] params: &RbdSimParams, #[spirv(uniform, descriptor_set = 1, binding = 2)] batch_ids: &BatchIndices, ) { let num_threads = num_workgroups.x * WORKGROUP_SIZE; let batch_id = invocation_id.y; - let params = all_params.at(batch_id as usize); let mut constraints = batch_ids.contact_batch_mut(batch_id, constraints); let constraint_builders = batch_ids.contact_batch(batch_id, constraint_builders); @@ -166,12 +164,11 @@ pub fn gpu_solver_refresh_rhs_wo_bias( constraint_builders: &[TwoBodyConstraintBuilder], #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] contacts_len: &[u32], #[spirv(storage_buffer, descriptor_set = 1, binding = 0)] solver_body_poses: &[Pose], - #[spirv(storage_buffer, descriptor_set = 1, binding = 1)] all_params: &[RbdSimParams], + #[spirv(uniform, descriptor_set = 1, binding = 1)] params: &RbdSimParams, #[spirv(uniform, descriptor_set = 1, binding = 2)] batch_ids: &BatchIndices, ) { let num_threads = num_workgroups.x * WORKGROUP_SIZE; let batch_id = invocation_id.y; - let params = all_params.at(batch_id as usize); let mut constraints = batch_ids.contact_batch_mut(batch_id, constraints); let constraint_builders = batch_ids.contact_batch(batch_id, constraint_builders); @@ -274,12 +271,11 @@ pub fn gpu_init_solver_vels_inc( #[spirv(global_invocation_id)] invocation_id: UVec3, #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] solver_vels_inc: &mut [Velocity], #[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 = 2)] 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); let i = invocation_id.x; let num_bodies = batch_ids.bodies_len; @@ -603,11 +599,10 @@ pub fn gpu_integrate_linearized( #[spirv(global_invocation_id)] invocation_id: UVec3, #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] poses: &mut [Pose], #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] solver_vels: &mut [Velocity], - #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] all_params: &[RbdSimParams], + #[spirv(uniform, descriptor_set = 0, binding = 2)] params: &RbdSimParams, #[spirv(uniform, descriptor_set = 0, binding = 3)] batch_ids: &BatchIndices, ) { let batch_id = invocation_id.y; - let params = all_params.at(batch_id as usize); let i = invocation_id.x; let num_bodies = batch_ids.bodies_len; diff --git a/src_rbd_shaders/queries/polygonal_feature.rs b/src_rbd_shaders/queries/polygonal_feature.rs index ddec14e5..c1f23dbe 100644 --- a/src_rbd_shaders/queries/polygonal_feature.rs +++ b/src_rbd_shaders/queries/polygonal_feature.rs @@ -288,6 +288,9 @@ mod dim2 { // 3D Implementation // ==================== +#[cfg(feature = "dim3")] +pub use dim3::manifold_reduction; + #[cfg(feature = "dim3")] mod dim3 { use super::*;