From 11d1278b65693f4e23b56ce864a3bd7f86cfc736 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Sat, 8 Aug 2026 10:12:37 +0200 Subject: [PATCH 01/41] fix(rbd): apply the per-batch stride to collider_parent reads in the narrow phase Replaces #21 Co-Authored-By: Haixuan Xavier Tao --- src_rbd_shaders/broad_phase/narrow_phase.rs | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/src_rbd_shaders/broad_phase/narrow_phase.rs b/src_rbd_shaders/broad_phase/narrow_phase.rs index 86d0e1e0..d73f0f1b 100644 --- a/src_rbd_shaders/broad_phase/narrow_phase.rs +++ b/src_rbd_shaders/broad_phase/narrow_phase.rs @@ -113,9 +113,14 @@ 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 like the other + // per-collider buffers — an unsliced read here silently returned + // batch 0's parents for every batch (masked whenever all envs are + // identical, wrong the moment they aren't). + 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; } @@ -580,8 +585,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; } From b4f7e55fa9137146a3bb8f6c902d80cab7c462b0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Sat, 8 Aug 2026 11:47:05 +0200 Subject: [PATCH 02/41] feat: per-environment collision-pair capacity override Replaces #24 Co-Authored-By: Haixuan Xavier Tao --- crates/nexus_python3d/src/nexus.rs | 7 +++++++ src/state.rs | 9 +++++++++ 2 files changed, 16 insertions(+) diff --git a/crates/nexus_python3d/src/nexus.rs b/crates/nexus_python3d/src/nexus.rs index ef9648d3..16e00d16 100644 --- a/crates/nexus_python3d/src/nexus.rs +++ b/crates/nexus_python3d/src/nexus.rs @@ -341,6 +341,13 @@ impl NexusState { self.0.set_rbd_steps_per_frame(steps); } + /// Overrides the per-environment collision-pair capacity (default 4096) + /// used when the GPU state is allocated at `finalize`. Lower it for many + /// small batched envs; pair-keyed buffers scale as capacity × envs. + fn set_rbd_collisions_capacity(&mut self, capacity: u32) { + self.0.set_rbd_collisions_capacity(capacity); + } + fn set_rbd_gravity(&mut self, viewer: PyRef, gravity: Vec3) { self.0 .set_rbd_gravity(viewer.backend(), [gravity.0.x, gravity.0.y, gravity.0.z]); diff --git a/src/state.rs b/src/state.rs index 76ab81db..42860afd 100644 --- a/src/state.rs +++ b/src/state.rs @@ -308,6 +308,15 @@ impl NexusState { self.rbd_steps_per_frame = steps.max(1); } + /// 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 — + /// pair-keyed workspaces scale as `capacity x num_envs x sizeof(manifold)`, + /// which at 2048 envs binds ~9 GiB unless this is lowered. + pub fn set_rbd_collisions_capacity(&mut self, capacity: u32) { + self.capacities.rbd.collisions_capacity = capacity.max(1); + } + /// Number of rigid-body solver steps per [`NexusPipeline::simulate`](crate::pipeline::NexusPipeline::simulate) call. pub fn rbd_steps_per_frame(&self) -> u32 { self.rbd_steps_per_frame From 9fb04fddf5ee852335b56270b409d967c63e80bf Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Sat, 8 Aug 2026 15:23:51 +0200 Subject: [PATCH 03/41] feat(rbd): make the narrow-phase contact prediction distance configurable Replaces #28 Co-Authored-By: Haixuan Xavier Tao --- src_rbd/broad_phase/narrow_phase.rs | 4 +++ src_rbd/pipeline/insertion_removal.rs | 7 +++++ src_rbd/pipeline/rbd_state.rs | 3 ++ src_rbd/pipeline/rbd_state_from_rapier.rs | 7 +++++ src_rbd/pipeline/rbd_step.rs | 1 + src_rbd_shaders/broad_phase/narrow_phase.rs | 31 +++++++++++++-------- 6 files changed, 42 insertions(+), 11 deletions(-) diff --git a/src_rbd/broad_phase/narrow_phase.rs b/src_rbd/broad_phase/narrow_phase.rs index 82d6ed60..664b0bc0 100644 --- a/src_rbd/broad_phase/narrow_phase.rs +++ b/src_rbd/broad_phase/narrow_phase.rs @@ -49,6 +49,7 @@ impl GpuNarrowPhase { batch_indices: &Tensor, collider_parent: &Tensor, collider_materials: &Tensor, + prediction: &Tensor, ) -> Result<(), GpuBackendError> { let num_batches = contacts_len.len() as u32; self.reset_narrow_phase @@ -66,6 +67,7 @@ impl GpuNarrowPhase { batch_indices, collider_parent, collider_materials, + prediction, )?; // Pass 2: defer the complex shape pairs into `pfm_pairs` (kept as a @@ -80,6 +82,7 @@ impl GpuNarrowPhase { pfm_pairs, pfm_pairs_len, batch_indices, + prediction, vertices, indices, )?; @@ -98,6 +101,7 @@ impl GpuNarrowPhase { indices, collider_parent, collider_materials, + prediction, )?; self.init_contacts_indirect_args.call( pass, diff --git a/src_rbd/pipeline/insertion_removal.rs b/src_rbd/pipeline/insertion_removal.rs index a19e0425..4cdbec54 100644 --- a/src_rbd/pipeline/insertion_removal.rs +++ b/src_rbd/pipeline/insertion_removal.rs @@ -143,6 +143,12 @@ impl RbdState { BufferUsages::STORAGE | BufferUsages::UNIFORM, ) .unwrap(); + let prediction = Tensor::scalar( + backend, + all_sim_params[0].prediction_distance(), + BufferUsages::UNIFORM | BufferUsages::COPY_DST, + ) + .unwrap(); // Two-element readback: the (max) collision-pair count and the uncolored count. let resize_readback = GpuReadback::new(backend, 2).unwrap(); let collision_pairs_indirect = @@ -269,6 +275,7 @@ impl RbdState { collision_pairs_len, collision_pairs_len_max, num_batches_uniform, + prediction, resize_readback, collision_pairs_indirect, contacts_per_batch_cpu, diff --git a/src_rbd/pipeline/rbd_state.rs b/src_rbd/pipeline/rbd_state.rs index ac691fac..f9f2b80f 100644 --- a/src_rbd/pipeline/rbd_state.rs +++ b/src_rbd/pipeline/rbd_state.rs @@ -160,6 +160,9 @@ 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, + /// Contact prediction distance (`RbdSimParams::prediction_distance`), + /// consumed by the narrow-phase kernels. + pub(super) prediction: Tensor, /// `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 diff --git a/src_rbd/pipeline/rbd_state_from_rapier.rs b/src_rbd/pipeline/rbd_state_from_rapier.rs index 042de1c3..361dac90 100644 --- a/src_rbd/pipeline/rbd_state_from_rapier.rs +++ b/src_rbd/pipeline/rbd_state_from_rapier.rs @@ -597,6 +597,12 @@ impl RbdState { BufferUsages::STORAGE | BufferUsages::UNIFORM, ) .unwrap(); + let prediction = Tensor::scalar( + backend, + all_sim_params[0].prediction_distance(), + BufferUsages::UNIFORM | BufferUsages::COPY_DST, + ) + .unwrap(); // Two-element readback: the (max) collision-pair count and the uncolored count. let resize_readback = GpuReadback::new(backend, 2).unwrap(); let collision_pairs_indirect = @@ -800,6 +806,7 @@ impl RbdState { collision_pairs_len, collision_pairs_len_max, num_batches_uniform, + prediction, resize_readback, collision_pairs_indirect, contacts_per_batch_cpu, diff --git a/src_rbd/pipeline/rbd_step.rs b/src_rbd/pipeline/rbd_step.rs index d208f328..500aae5f 100644 --- a/src_rbd/pipeline/rbd_step.rs +++ b/src_rbd/pipeline/rbd_step.rs @@ -258,6 +258,7 @@ impl RbdPipeline { &state.batch_indices, &state.collider_parent, &state.collider_materials, + &state.prediction, )?; drop(pass); diff --git a/src_rbd_shaders/broad_phase/narrow_phase.rs b/src_rbd_shaders/broad_phase/narrow_phase.rs index d73f0f1b..169f0055 100644 --- a/src_rbd_shaders/broad_phase/narrow_phase.rs +++ b/src_rbd_shaders/broad_phase/narrow_phase.rs @@ -67,7 +67,6 @@ pub fn gpu_narrow_phase_init_contacts_dispatch( } } -pub(crate) const PREDICTION: f32 = 2.0e-2; // TODO: make the prediction configurable. /// Narrow phase, pass 1 of 2: analytic shape-shape contacts for ball / cuboid /// pairs, written straight into the `contacts` buffer. @@ -92,6 +91,8 @@ 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], + // Contact prediction distance (`RbdSimParams::prediction_distance`). + #[spirv(uniform, descriptor_set = 0, binding = 9)] prediction: &f32, ) { let num_threads = num_workgroups.x * WORKGROUP_SIZE; let batch_id = invocation_id.y; @@ -164,12 +165,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 @@ -213,6 +214,7 @@ 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)] prediction: &f32, ) { let num_threads = num_workgroups.x * WORKGROUP_SIZE; let batch_id = invocation_id.y; @@ -304,6 +306,7 @@ pub fn gpu_narrow_phase_shape_shape_deferred( let mesh = shape1.to_trimesh(); let convex = shape2; trimesh_convex( + *prediction, pose12, &mesh, convex, @@ -322,6 +325,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, @@ -341,6 +345,7 @@ pub fn gpu_narrow_phase_shape_shape_deferred( let pline = shape1.to_polyline(); let convex = shape2; polyline_convex( + *prediction, pose12, &pline, convex, @@ -359,6 +364,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, @@ -376,6 +382,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, @@ -392,10 +399,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. @@ -446,6 +453,7 @@ fn trimesh_convex( /// Collision detection between a polyline and a convex shape. fn polyline_convex( + prediction: f32, pose12: Pose, mesh: &Polyline, convex: &Shape, @@ -462,11 +470,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. @@ -564,6 +572,7 @@ 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)] prediction: &f32, ) { let num_threads = num_workgroups.x * WORKGROUP_SIZE; let batch_id = invocation_id.y; @@ -597,13 +606,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. From c64d717dee44f786db282243470e3f9f2305f7ff Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Sun, 9 Aug 2026 09:38:14 +0200 Subject: [PATCH 04/41] fix(rbd): thread the configurable prediction distance through the brute-force broad phase Completes #28 --- src_rbd/broad_phase/lbvh.rs | 2 ++ src_rbd/pipeline/rbd_step.rs | 1 + src_rbd_shaders/broad_phase/brute_force.rs | 6 +++--- src_rbd_shaders/broad_phase/narrow_phase.rs | 5 ++--- 4 files changed, 8 insertions(+), 6 deletions(-) diff --git a/src_rbd/broad_phase/lbvh.rs b/src_rbd/broad_phase/lbvh.rs index 8d9f0120..e97e3e96 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]>, + prediction: &Tensor, ) -> Result<(), GpuBackendError> { state.resize_bf_buffers(backend, colliders_len); @@ -360,6 +361,7 @@ impl Lbvh { collision_groups, batch_indices, pair_filter, + prediction, )?; // Single 256-lane workgroup: parallel max over the per-batch counts. self.shaders.lbvh_init_indirect_args.call( diff --git a/src_rbd/pipeline/rbd_step.rs b/src_rbd/pipeline/rbd_step.rs index 500aae5f..10a9a4e1 100644 --- a/src_rbd/pipeline/rbd_step.rs +++ b/src_rbd/pipeline/rbd_step.rs @@ -157,6 +157,7 @@ impl RbdPipeline { &mut state.collision_pairs_indirect, &state.collision_groups, &state.pair_filter, + &state.prediction, )?; drop(pass); backend.submit(encoder)?; diff --git a/src_rbd_shaders/broad_phase/brute_force.rs b/src_rbd_shaders/broad_phase/brute_force.rs index 94601b1a..b16b2522 100644 --- a/src_rbd_shaders/broad_phase/brute_force.rs +++ b/src_rbd_shaders/broad_phase/brute_force.rs @@ -16,8 +16,6 @@ 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 +57,8 @@ 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]], + // Contact prediction distance (`RbdSimParams::prediction_distance`). + #[spirv(uniform, descriptor_set = 0, binding = 6)] prediction: &f32, ) { 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(*prediction); 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 169f0055..f1348e6a 100644 --- a/src_rbd_shaders/broad_phase/narrow_phase.rs +++ b/src_rbd_shaders/broad_phase/narrow_phase.rs @@ -67,7 +67,6 @@ pub fn gpu_narrow_phase_init_contacts_dispatch( } } - /// Narrow phase, pass 1 of 2: analytic shape-shape contacts for ball / cuboid /// pairs, written straight into the `contacts` buffer. /// @@ -165,7 +164,7 @@ 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 @@ -606,7 +605,7 @@ pub fn gpu_narrow_phase_pfm_pfm( pair.thickness1, &pair.shape2, pair.thickness2, - (*prediction), + *prediction, vertices, #[cfg(feature = "dim3")] indices, From 6c60a18fb63dd25cd4aedb96e2ec1e548c1ba06c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Sun, 9 Aug 2026 14:05:42 +0200 Subject: [PATCH 05/41] feat(python): per-environment MJCF insertion Replaces #16 Co-Authored-By: Haixuan Xavier Tao --- crates/nexus_python3d/src/loaders.rs | 23 +++++++++++++++-------- crates/nexus_python3d/src/nexus.rs | 19 ++++++++++++++----- src/state.rs | 9 +++++++++ 3 files changed, 38 insertions(+), 13 deletions(-) diff --git a/crates/nexus_python3d/src/loaders.rs b/crates/nexus_python3d/src/loaders.rs index 58447299..e0c85316 100644 --- a/crates/nexus_python3d/src/loaders.rs +++ b/crates/nexus_python3d/src/loaders.rs @@ -120,6 +120,7 @@ pub fn insert_mjcf( mut viewer: PyRefMut, scene_path: &std::path::Path, render_colliders: bool, + env: usize, ) -> PyResult { use nexus3d::prelude::RbdCoupling; use pyo3::exceptions::PyRuntimeError; @@ -143,7 +144,7 @@ pub fn insert_mjcf( 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, @@ -242,11 +243,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); + 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 +276,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 }) } diff --git a/crates/nexus_python3d/src/nexus.rs b/crates/nexus_python3d/src/nexus.rs index 16e00d16..873870ef 100644 --- a/crates/nexus_python3d/src/nexus.rs +++ b/crates/nexus_python3d/src/nexus.rs @@ -322,17 +322,26 @@ 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 this + /// 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) + crate::loaders::insert_mjcf(&mut self.0, viewer, &scene_path, render_colliders, env) } // --- rbd config ------------------------------------------------------- diff --git a/src/state.rs b/src/state.rs index 42860afd..b1e0ab17 100644 --- a/src/state.rs +++ b/src/state.rs @@ -302,6 +302,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 — + /// pair-keyed workspaces scale as `capacity x num_envs x sizeof(manifold)`, + /// which at 2048 envs binds ~9 GiB unless this is 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) { From f8adbd952087f89f97c3fe0bf6275afe4e2e32c7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Sun, 9 Aug 2026 20:41:09 +0200 Subject: [PATCH 06/41] fix(python): drop the duplicated collisions-capacity setter and pass the RbdCoupling to insert_rigid_body_in Completes #16 --- crates/nexus_python3d/src/loaders.rs | 2 +- crates/nexus_python3d/src/nexus.rs | 7 ------- src/state.rs | 9 --------- 3 files changed, 1 insertion(+), 17 deletions(-) diff --git a/crates/nexus_python3d/src/loaders.rs b/crates/nexus_python3d/src/loaders.rs index e0c85316..18f1d8d6 100644 --- a/crates/nexus_python3d/src/loaders.rs +++ b/crates/nexus_python3d/src/loaders.rs @@ -243,7 +243,7 @@ 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_in(env, body, collider); + let handle = state.insert_rigid_body_in(env, body, collider, RbdCoupling::None); if env == 0 { v.insert_shape(handle, &shape, rp::Pose::IDENTITY); } diff --git a/crates/nexus_python3d/src/nexus.rs b/crates/nexus_python3d/src/nexus.rs index 873870ef..3d19ffc6 100644 --- a/crates/nexus_python3d/src/nexus.rs +++ b/crates/nexus_python3d/src/nexus.rs @@ -350,13 +350,6 @@ impl NexusState { self.0.set_rbd_steps_per_frame(steps); } - /// Overrides the per-environment collision-pair capacity (default 4096) - /// used when the GPU state is allocated at `finalize`. Lower it for many - /// small batched envs; pair-keyed buffers scale as capacity × envs. - fn set_rbd_collisions_capacity(&mut self, capacity: u32) { - self.0.set_rbd_collisions_capacity(capacity); - } - fn set_rbd_gravity(&mut self, viewer: PyRef, gravity: Vec3) { self.0 .set_rbd_gravity(viewer.backend(), [gravity.0.x, gravity.0.y, gravity.0.z]); diff --git a/src/state.rs b/src/state.rs index b1e0ab17..d021ddcc 100644 --- a/src/state.rs +++ b/src/state.rs @@ -317,15 +317,6 @@ impl NexusState { self.rbd_steps_per_frame = steps.max(1); } - /// 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 — - /// pair-keyed workspaces scale as `capacity x num_envs x sizeof(manifold)`, - /// which at 2048 envs binds ~9 GiB unless this is lowered. - pub fn set_rbd_collisions_capacity(&mut self, capacity: u32) { - self.capacities.rbd.collisions_capacity = capacity.max(1); - } - /// Number of rigid-body solver steps per [`NexusPipeline::simulate`](crate::pipeline::NexusPipeline::simulate) call. pub fn rbd_steps_per_frame(&self) -> u32 { self.rbd_steps_per_frame From 3deaad3698042d5a2b8b64b864f2962e5684a8c7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Fri, 14 Aug 2026 19:26:33 +0200 Subject: [PATCH 07/41] feat(python): per-step MJCF actuator control + multibody state readback Replaces #12 Co-Authored-By: Haixuan Xavier Tao --- crates/nexus_python3d/src/loaders.rs | 12 ++- crates/nexus_python3d/src/nexus.rs | 76 ++++++++++++++++++- crates/nexus_python3d/src/viewer.rs | 48 ++++++++++++ src/state.rs | 31 ++++++++ .../multibody/multibody_from_rapier.rs | 4 +- src_rbd/dynamics/multibody/multibody_set.rs | 73 ++++++++++++++++++ src_viewer/viewer.rs | 37 +++++++++ 7 files changed, 274 insertions(+), 7 deletions(-) diff --git a/crates/nexus_python3d/src/loaders.rs b/crates/nexus_python3d/src/loaders.rs index 18f1d8d6..d550232c 100644 --- a/crates/nexus_python3d/src/loaders.rs +++ b/crates/nexus_python3d/src/loaders.rs @@ -115,13 +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, env: usize, -) -> PyResult { +) -> PyResult<(MjcfSceneInfo, Option)> { use nexus3d::prelude::RbdCoupling; use pyo3::exceptions::PyRuntimeError; use rapier3d::parry::bounding_volume::BoundingVolume; // for `Aabb::merge` @@ -142,6 +148,7 @@ pub fn insert_mjcf( let mut floor: Option<(glamx::Vec3, glamx::Vec3)> = None; let mut camera: Option<(glamx::Vec3, glamx::Vec3)> = None; + let mut robot_handles: Option = None; match MjcfRobot::from_file(scene_path, options) { Ok((robot, _model)) => { let world = state.rbd_world_mut(env); @@ -226,6 +233,7 @@ pub fn insert_mjcf( let eye = target + glamx::Vec3::new(radius * 2.2, -radius * 2.2, radius * 1.6); camera = Some((eye, target)); } + robot_handles = Some(handles); } Err(e) => { return Err(PyRuntimeError::new_err(format!( @@ -284,5 +292,5 @@ pub fn insert_mjcf( .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 3d19ffc6..5e9462dd 100644 --- a/crates/nexus_python3d/src/nexus.rs +++ b/crates/nexus_python3d/src/nexus.rs @@ -95,15 +95,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 ----------------------------------------------------- @@ -341,7 +343,73 @@ impl NexusState { render_colliders: bool, env: usize, ) -> PyResult { - crate::loaders::insert_mjcf(&mut self.0, viewer, &scene_path, render_colliders, env) + 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 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 in one buffer write. + /// + /// Call once per control step, after `finalize`; the next + /// `NexusPipeline.simulate` steps the solver against the new targets. This + /// is the GPU counterpart of stepping rapier natively with actuators. + #[pyo3(signature = (viewer, ctrl, env=0))] + fn apply_actuator_controls( + &mut self, + viewer: PyRef, + ctrl: Vec, + env: usize, + ) -> 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(), env, |world| { + handles.apply_controls_multibody( + &mut world.bodies, + &mut world.multibody_joints, + &ctrl, + ); + }) + .map_err(gpu_err) } // --- rbd config ------------------------------------------------------- diff --git a/crates/nexus_python3d/src/viewer.rs b/crates/nexus_python3d/src/viewer.rs index d5e6983e..c19e0276 100644 --- a/crates/nexus_python3d/src/viewer.rs +++ b/crates/nexus_python3d/src/viewer.rs @@ -282,6 +282,54 @@ impl NexusViewer { .map_err(|e| PyRuntimeError::new_err(format!("{e:?}"))) } + /// Reads back environment `env`'s multibody link states from the GPU in one + /// readback. Returns five float32 numpy arrays, one row per link (in the + /// GPU build's traversal order — multibodies, then links, parent before + /// child; the same order `NexusState.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). + #[pyo3(signature = (state, env=0))] + #[allow(clippy::type_complexity)] + fn read_multibody_links<'py>( + &mut self, + py: Python<'py>, + state: PyRef, + env: u32, + ) -> ( + Bound<'py, PyArray2>, + Bound<'py, PyArray2>, + Bound<'py, PyArray2>, + Bound<'py, PyArray2>, + Bound<'py, PyArray2>, + ) { + let links = pollster::block_on(self.inner_mut().read_multibody_links(&state.0, env)); + 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(), + ) + } + // --- misc ------------------------------------------------------------- fn clear_scene(&mut self) { diff --git a/src/state.rs b/src/state.rs index d021ddcc..a8a845eb 100644 --- a/src/state.rs +++ b/src/state.rs @@ -410,6 +410,37 @@ impl NexusState { &mut self.rbd_envs[env] } + /// Runtime actuation entry point: mutates environment `env`'s rapier + /// multibody joints through `f` (e.g. `rapier3d-mjcf`'s + /// `apply_controls_multibody`, which implements MJCF actuator semantics), + /// then pushes the refreshed joint data — motor targets/gains, limits — to + /// the GPU multibody links in one buffer write. + /// + /// Unlike [`Self::rbd_world_mut`] this does NOT mark the world dirty: motor + /// updates are per-step control, not a topology change, so no GPU rebuild + /// is triggered. Call after [`Self::finalize`]; a no-op before it. + pub fn control_multibody_motors( + &mut self, + backend: &GpuBackend, + env: usize, + f: F, + ) -> Result<(), GpuBackendError> + where + F: FnOnce(&mut PhysicsWorld), + { + let world = &mut self.rbd_envs[env]; + f(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(()) + } + /// Mutable access to environment `env`'s rapier world that does **not** mark /// the rbd state dirty, for use after [`Self::finalize`]. /// diff --git a/src_rbd/dynamics/multibody/multibody_from_rapier.rs b/src_rbd/dynamics/multibody/multibody_from_rapier.rs index ce742aa9..b9025f2c 100644 --- a/src_rbd/dynamics/multibody/multibody_from_rapier.rs +++ b/src_rbd/dynamics/multibody/multibody_from_rapier.rs @@ -511,10 +511,12 @@ 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(), diff --git a/src_rbd/dynamics/multibody/multibody_set.rs b/src_rbd/dynamics/multibody/multibody_set.rs index 5a017157..84ef47d7 100644 --- a/src_rbd/dynamics/multibody/multibody_set.rs +++ b/src_rbd/dynamics/multibody/multibody_set.rs @@ -432,6 +432,79 @@ impl GpuMultibodySet { ) } + /// Per-batch per-step link workspace (generalized coordinates, joint + /// rotations, world-space link velocities). Read it back with + /// `slow_read_buffer` for joint/base state observation; entries are 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 + } + + /// 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 (e.g. via `rapier3d-mjcf`'s + /// `apply_controls_multibody`, which implements the MJCF actuator + /// semantics), then call this to push the new motor state to the GPU. Only + /// joint data is refreshed — coordinates, velocities and mass properties are + /// untouched, so this cannot be used to 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, + ) + } + + /// Upload a new gravity vector. + pub fn set_gravity(&mut self, backend: &GpuBackend, g: [f32; 3]) { + self.gravity = Tensor::scalar( + backend, + Vec4::new(g[0], g[1], g[2], 0.0), + BufferUsages::STORAGE | BufferUsages::UNIFORM | BufferUsages::COPY_DST, + ) + .unwrap(); + } + /// Number of multibody-touching impulse joints in any batch. pub fn mb_impulse_joints_per_batch(&self) -> u32 { self.mb_imp_joints_per_batch diff --git a/src_viewer/viewer.rs b/src_viewer/viewer.rs index c07b8d54..23a76773 100644 --- a/src_viewer/viewer.rs +++ b/src_viewer/viewer.rs @@ -1217,6 +1217,43 @@ impl NexusViewer { .set_denoise(enabled); } + /// Reads back environment `env`'s multibody link workspaces from the GPU in + /// one readback: per link, the generalized joint coordinates, accumulated + /// joint rotation, world pose, and world-space velocity. Links are in the + /// GPU build's traversal order (multibodies, then links, parent before + /// child) — the same order `NexusState::control_multibody_motors` targets. + /// Empty when no multibody state exists. + /// + /// Velocities are only meaningful after the first simulated step (the + /// forward-kinematics pass fills them); coordinates and poses are valid + /// from `finalize`. + pub async fn read_multibody_links( + &mut self, + state: &NexusState, + env: u32, + ) -> Vec { + let Some(rbd) = state.rbd.as_ref() else { + return Vec::new(); + }; + let mbs = rbd.multibodies(); + let stride = mbs.links_per_batch() as usize; + if stride == 0 { + return Vec::new(); + } + let mut all = bytemuck::zeroed_vec(mbs.links_workspace().len() as usize); + if self + .backend() + .slow_read_buffer(mbs.links_workspace().buffer(), &mut all) + .await + .is_err() + { + return Vec::new(); + } + let start = (env as usize * stride).min(all.len()); + let end = (start + stride).min(all.len()); + all[start..end].to_vec() + } + /// Draws example-specific egui widgets into the current frame's UI pass. /// /// Call this once per frame, after [`Self::render_frame`], to overlay a From 88221d342f17e4f6cefbafbff88ac6d28b25db90 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Fri, 14 Aug 2026 22:03:57 +0200 Subject: [PATCH 08/41] fix(python): gate multibody control/readback on dim3, add the missing PyArray2 import Replaces #12 Co-Authored-By: Haixuan Xavier Tao --- crates/nexus_python3d/src/loaders.rs | 7 +++---- crates/nexus_python3d/src/viewer.rs | 2 +- src/state.rs | 1 + src_rbd/dynamics/multibody/multibody_set.rs | 6 +----- src_viewer/viewer.rs | 1 + 5 files changed, 7 insertions(+), 10 deletions(-) diff --git a/crates/nexus_python3d/src/loaders.rs b/crates/nexus_python3d/src/loaders.rs index d550232c..b19602d9 100644 --- a/crates/nexus_python3d/src/loaders.rs +++ b/crates/nexus_python3d/src/loaders.rs @@ -148,8 +148,7 @@ pub fn insert_mjcf( let mut floor: Option<(glamx::Vec3, glamx::Vec3)> = None; let mut camera: Option<(glamx::Vec3, glamx::Vec3)> = None; - let mut robot_handles: Option = 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(env); let handles = robot.clone().insert_using_multibody_joints( @@ -233,7 +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)); } - robot_handles = Some(handles); + Some(handles) } Err(e) => { return Err(PyRuntimeError::new_err(format!( @@ -241,7 +240,7 @@ pub fn insert_mjcf( scene_path.display() ))); } - } + }; let loaded = camera.is_some(); let v = viewer.rust_mut(); diff --git a/crates/nexus_python3d/src/viewer.rs b/crates/nexus_python3d/src/viewer.rs index c19e0276..25ad62d7 100644 --- a/crates/nexus_python3d/src/viewer.rs +++ b/crates/nexus_python3d/src/viewer.rs @@ -10,7 +10,7 @@ use crate::nexus::{GpuTimestamps, NexusState}; use crate::rbd::{RigidBodyHandle, SharedShape}; use khal::backend::GpuBackend; use nexus_viewer3d::NexusViewer as RViewer; -use numpy::{IntoPyArray, PyArray3, PyArrayMethods}; +use numpy::{IntoPyArray, PyArray2, PyArray3, PyArrayMethods}; use pyo3::exceptions::PyRuntimeError; use pyo3::prelude::*; diff --git a/src/state.rs b/src/state.rs index a8a845eb..de2c0810 100644 --- a/src/state.rs +++ b/src/state.rs @@ -419,6 +419,7 @@ impl NexusState { /// Unlike [`Self::rbd_world_mut`] this does NOT mark the world dirty: motor /// updates are per-step control, not a topology change, so no GPU rebuild /// is triggered. Call after [`Self::finalize`]; a no-op before it. + #[cfg(all(feature = "dim3", feature = "rbd"))] pub fn control_multibody_motors( &mut self, backend: &GpuBackend, diff --git a/src_rbd/dynamics/multibody/multibody_set.rs b/src_rbd/dynamics/multibody/multibody_set.rs index 84ef47d7..3e407594 100644 --- a/src_rbd/dynamics/multibody/multibody_set.rs +++ b/src_rbd/dynamics/multibody/multibody_set.rs @@ -488,11 +488,7 @@ impl GpuMultibodySet { offset += 1; } } - backend.write_buffer( - self.links_static.buffer_mut(), - 0, - &self.links_static_mirror, - ) + backend.write_buffer(self.links_static.buffer_mut(), 0, &self.links_static_mirror) } /// Upload a new gravity vector. diff --git a/src_viewer/viewer.rs b/src_viewer/viewer.rs index 23a76773..83dab43c 100644 --- a/src_viewer/viewer.rs +++ b/src_viewer/viewer.rs @@ -1227,6 +1227,7 @@ impl NexusViewer { /// Velocities are only meaningful after the first simulated step (the /// forward-kinematics pass fills them); coordinates and poses are valid /// from `finalize`. + #[cfg(feature = "dim3")] pub async fn read_multibody_links( &mut self, state: &NexusState, From e22c94e81729b56c31400a41d29c672b806d0580 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Sat, 15 Aug 2026 10:55:21 +0200 Subject: [PATCH 09/41] fix(rbd): decode the SoA link workspace for the multibody readback and drop the stale set_gravity copy Completes #12 --- src_rbd/dynamics/multibody/multibody_set.rs | 21 ++++-------- src_rbd_shaders/dynamics/multibody/ws_soa.rs | 35 ++++++++++++++++++++ src_viewer/viewer.rs | 11 ++++-- 3 files changed, 50 insertions(+), 17 deletions(-) diff --git a/src_rbd/dynamics/multibody/multibody_set.rs b/src_rbd/dynamics/multibody/multibody_set.rs index 3e407594..3d9c301a 100644 --- a/src_rbd/dynamics/multibody/multibody_set.rs +++ b/src_rbd/dynamics/multibody/multibody_set.rs @@ -433,11 +433,12 @@ impl GpuMultibodySet { } /// Per-batch per-step link workspace (generalized coordinates, joint - /// rotations, world-space link velocities). Read it back with - /// `slow_read_buffer` for joint/base state observation; entries are laid out - /// `env * links_per_batch + link`, in [`from_rapier`](Self::from_rapier)'s - /// link traversal order. - pub fn links_workspace(&self) -> &Tensor { + /// 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 } @@ -491,16 +492,6 @@ impl GpuMultibodySet { backend.write_buffer(self.links_static.buffer_mut(), 0, &self.links_static_mirror) } - /// Upload a new gravity vector. - pub fn set_gravity(&mut self, backend: &GpuBackend, g: [f32; 3]) { - self.gravity = Tensor::scalar( - backend, - Vec4::new(g[0], g[1], g[2], 0.0), - BufferUsages::STORAGE | BufferUsages::UNIFORM | BufferUsages::COPY_DST, - ) - .unwrap(); - } - /// Number of multibody-touching impulse joints in any batch. pub fn mb_impulse_joints_per_batch(&self) -> u32 { self.mb_imp_joints_per_batch 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_viewer/viewer.rs b/src_viewer/viewer.rs index 83dab43c..29edddc1 100644 --- a/src_viewer/viewer.rs +++ b/src_viewer/viewer.rs @@ -1241,15 +1241,22 @@ impl NexusViewer { if stride == 0 { return Vec::new(); } - let mut all = bytemuck::zeroed_vec(mbs.links_workspace().len() as usize); + // The workspace is stored batch-interleaved SoA (quads), so read the + // raw buffer and decode it back into one struct per link. + let mut raw: Vec = bytemuck::zeroed_vec(mbs.links_workspace().len() as usize); if self .backend() - .slow_read_buffer(mbs.links_workspace().buffer(), &mut all) + .slow_read_buffer(mbs.links_workspace().buffer(), &mut raw) .await .is_err() { return Vec::new(); } + let all = nexus::rbd::shaders::dynamics::ws_soa_to_structs( + &raw, + mbs.links_per_batch(), + mbs.num_batches(), + ); let start = (env as usize * stride).min(all.len()); let end = (start + stride).min(all.len()); all[start..end].to_vec() From 0748a8af0283b69c4ca178ed6dd2d59e44a6c5c7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Sat, 15 Aug 2026 13:12:48 +0200 Subject: [PATCH 10/41] perf(rbd): dedupe shared TriMesh uploads in from_rapier Replaces #19 Co-Authored-By: Haixuan Xavier Tao --- src_rbd/pipeline/rbd_state_from_rapier.rs | 25 ++++++++++++++++++++--- 1 file changed, 22 insertions(+), 3 deletions(-) diff --git a/src_rbd/pipeline/rbd_state_from_rapier.rs b/src_rbd/pipeline/rbd_state_from_rapier.rs index 361dac90..3f7acaba 100644 --- a/src_rbd/pipeline/rbd_state_from_rapier.rs +++ b/src_rbd/pipeline/rbd_state_from_rapier.rs @@ -147,6 +147,10 @@ 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(); + // TriMesh dedupe: a `SharedShape` cloned across envs (e.g. shared + // terrain) is serialized into `shape_buffers` once and its `Shape` + // descriptor reused — keyed by the parry shape data pointer. + let mut trimesh_cache: HashMap = HashMap::new(); let mut joint_envs: Vec<( &ImpulseJointSet, HashMap, @@ -291,9 +295,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)); From 9c796e136e06123995b523c0c8f89a1119c0b642 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Sat, 15 Aug 2026 17:48:02 +0200 Subject: [PATCH 11/41] perf(rbd): optional GPU contact reduction, merging per-pair manifolds to <=4 points Replaces #17 Co-Authored-By: Haixuan Xavier Tao --- src_rbd/broad_phase/narrow_phase.rs | 21 +++++ src_rbd/pipeline/rbd_step.rs | 5 ++ src_rbd_shaders/broad_phase/narrow_phase.rs | 85 ++++++++++++++++++++ src_rbd_shaders/queries/polygonal_feature.rs | 5 ++ 4 files changed, 116 insertions(+) diff --git a/src_rbd/broad_phase/narrow_phase.rs b/src_rbd/broad_phase/narrow_phase.rs index 664b0bc0..f0433425 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, } @@ -50,6 +54,9 @@ impl GpuNarrowPhase { collider_parent: &Tensor, collider_materials: &Tensor, prediction: &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 @@ -103,6 +110,20 @@ impl GpuNarrowPhase { collider_materials, prediction, )?; + // 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, + )?; + } + #[cfg(not(feature = "dim3"))] + let _ = reduce_contacts; self.init_contacts_indirect_args.call( pass, 256u32, diff --git a/src_rbd/pipeline/rbd_step.rs b/src_rbd/pipeline/rbd_step.rs index 10a9a4e1..b98c6fb6 100644 --- a/src_rbd/pipeline/rbd_step.rs +++ b/src_rbd/pipeline/rbd_step.rs @@ -32,6 +32,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 +57,7 @@ impl RbdPipeline { coloring: GpuColoring::from_backend(backend)?, warmstart: GpuWarmstart::from_backend(backend)?, reduce: Reduce::from_backend(backend)?, + contact_reduction: false, }) } @@ -260,6 +264,7 @@ impl RbdPipeline { &state.collider_parent, &state.collider_materials, &state.prediction, + self.contact_reduction, )?; drop(pass); diff --git a/src_rbd_shaders/broad_phase/narrow_phase.rs b/src_rbd_shaders/broad_phase/narrow_phase.rs index f1348e6a..4512ec16 100644 --- a/src_rbd_shaders/broad_phase/narrow_phase.rs +++ b/src_rbd_shaders/broad_phase/narrow_phase.rs @@ -6,6 +6,8 @@ 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,6 +69,89 @@ pub fn gpu_narrow_phase_init_contacts_dispatch( } } +/// Optional contact reduction: compacts each batch's contacts in place by +/// merging all manifolds of a collider pair (e.g. per-triangle trimesh +/// contacts, which share one `colliders` key and one collider-A local frame) +/// into a single `MAX_MANIFOLD_POINTS` manifold via `manifold_reduction`, +/// keeping the deeper manifold's normal. The first record of a pair is kept +/// verbatim, so single-manifold pairs are bit-identical to the unreduced +/// path. Approximations: one normal per merged manifold, greedy merging in +/// emission order. 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, +) { + 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); + + let mut w = 0usize; // write cursor — always ≤ read cursor, in-place safe + 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 { + // Pool the two manifolds' points (same collider-A local frame). + let na = (out.contact.len as usize).min(MAX_MANIFOLD_POINTS); + let nb = (im.contact.len as usize).min(MAX_MANIFOLD_POINTS); + let mut cand = [ContactPoint::default(); 8]; + for k in 0..na { + cand.write(k, out.contact.points_a.read(k)); + } + for k in 0..nb { + cand.write(na + k, im.contact.points_a.read(k)); + } + // Normal of whichever manifold holds the deepest point. + 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, (na + nb) as u32, normal); + // `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); + } +} + +const PREDICTION: f32 = 2.0e-3; // TODO: make the prediction configurable. + /// Narrow phase, pass 1 of 2: analytic shape-shape contacts for ball / cuboid /// pairs, written straight into the `contacts` buffer. /// diff --git a/src_rbd_shaders/queries/polygonal_feature.rs b/src_rbd_shaders/queries/polygonal_feature.rs index ddec14e5..4d41fefd 100644 --- a/src_rbd_shaders/queries/polygonal_feature.rs +++ b/src_rbd_shaders/queries/polygonal_feature.rs @@ -288,6 +288,11 @@ mod dim2 { // 3D Implementation // ==================== +/// Re-export of the ≤8 → ≤4 keep-deepest-then-spread contact selector for the +/// optional contact-reduction pass (see `gpu_reduce_contacts`). +#[cfg(feature = "dim3")] +pub use dim3::manifold_reduction; + #[cfg(feature = "dim3")] mod dim3 { use super::*; From 4897931853067137d9848c32c8fa5806f0741aaa Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Sun, 16 Aug 2026 09:14:26 +0200 Subject: [PATCH 12/41] fix(rbd): pass the prediction distance to manifold_reduction in the contact-reduction kernel Completes #17 --- src_rbd/broad_phase/narrow_phase.rs | 1 + src_rbd_shaders/broad_phase/narrow_phase.rs | 5 ++++- 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/src_rbd/broad_phase/narrow_phase.rs b/src_rbd/broad_phase/narrow_phase.rs index f0433425..017908ab 100644 --- a/src_rbd/broad_phase/narrow_phase.rs +++ b/src_rbd/broad_phase/narrow_phase.rs @@ -120,6 +120,7 @@ impl GpuNarrowPhase { contacts, contacts_len, batch_indices, + prediction, )?; } #[cfg(not(feature = "dim3"))] diff --git a/src_rbd_shaders/broad_phase/narrow_phase.rs b/src_rbd_shaders/broad_phase/narrow_phase.rs index 4512ec16..78e03edf 100644 --- a/src_rbd_shaders/broad_phase/narrow_phase.rs +++ b/src_rbd_shaders/broad_phase/narrow_phase.rs @@ -85,6 +85,9 @@ pub fn gpu_reduce_contacts( #[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, + // Contact prediction distance: `manifold_reduction` only keeps candidates + // within it, exactly as the narrow-phase passes that produced them. + #[spirv(uniform, descriptor_set = 0, binding = 3)] prediction: &f32, ) { let batch_id = workgroup_id.y; let capacity = batch_ids.contacts_batch_capacity as usize; @@ -128,7 +131,7 @@ pub fn gpu_reduce_contacts( } else { out.contact.normal_a }; - let mut reduced = manifold_reduction(&cand, (na + nb) as u32, normal); + let mut reduced = manifold_reduction(&cand, (na + nb) as u32, normal, *prediction); // `manifold_reduction` fills points/len only. reduced.normal_a = normal; let mut kept = out; From 5c090a0e49cddabfc70caf0665eabe766aebb929 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Sun, 16 Aug 2026 11:30:44 +0200 Subject: [PATCH 13/41] perf(rbd): flat 1-D narrow-phase dispatch, packing warps across batches Replaces #21 Co-Authored-By: Haixuan Xavier Tao --- src_rbd/broad_phase/narrow_phase.rs | 50 +++++-- src_rbd/pipeline/insertion_removal.rs | 6 + src_rbd/pipeline/rbd_state.rs | 6 + src_rbd/pipeline/rbd_state_from_rapier.rs | 6 + src_rbd/pipeline/rbd_step.rs | 6 +- src_rbd_shaders/broad_phase/narrow_phase.rs | 147 ++++++++++++++------ 6 files changed, 166 insertions(+), 55 deletions(-) diff --git a/src_rbd/broad_phase/narrow_phase.rs b/src_rbd/broad_phase/narrow_phase.rs index 017908ab..2c523c61 100644 --- a/src_rbd/broad_phase/narrow_phase.rs +++ b/src_rbd/broad_phase/narrow_phase.rs @@ -6,9 +6,9 @@ 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, - NarrowPhasePfmPair, + CollisionPair, GpuFlattenBatchesDispatch, GpuNarrowPhaseInitContactsDispatch, + GpuNarrowPhasePfmPfm, GpuNarrowPhaseShapeShape, GpuNarrowPhaseShapeShapeDeferred, + GpuResetNarrowPhase, NarrowPhasePfmPair, }; use crate::shaders::shapes::Shape; use khal::Shader; @@ -26,7 +26,11 @@ pub struct GpuNarrowPhase { narrow_phase_pfm_pfm: GpuNarrowPhasePfmPfm, #[cfg(feature = "dim3")] reduce_contacts: GpuReduceContacts, - init_pfm_pfm_indirect_args: GpuInitPfmPfmDispatch, + /// Builds the flat 1-D dispatch grid + prefix offsets for a per-batch + /// work-list (used for both the collision pairs and the PFM pairs), so the + /// kernels pack items from many batches into full warps instead of one + /// mostly-idle workgroup per batch. + flatten_batches: GpuFlattenBatchesDispatch, init_contacts_indirect_args: GpuNarrowPhaseInitContactsDispatch, } @@ -41,8 +45,8 @@ impl GpuNarrowPhase { vertices: &Tensor, indices: &Tensor, collision_pairs: &Tensor, - collision_pairs_len: &Tensor, - collision_pairs_indirect: &Tensor<[u32; 3]>, + collision_pairs_len: &mut Tensor, + collision_pairs_indirect: &mut Tensor<[u32; 3]>, contacts: &mut Tensor, contacts_len: &mut Tensor, contacts_indirect: &mut Tensor<[u32; 3]>, @@ -57,16 +61,30 @@ impl GpuNarrowPhase { // Optional: merge each collider pair's manifolds into one before the // solvers see them. `false` skips the kernel entirely. reduce_contacts: bool, + pairs_offsets: &mut Tensor, + pfm_offsets: &mut Tensor, ) -> Result<(), GpuBackendError> { let num_batches = contacts_len.len() as u32; self.reset_narrow_phase .call(pass, [num_batches, 1, 1], contacts_len, pfm_pairs_len)?; - self.narrow_phase.call( + // The broad phase wrote a `[max/64, num_batches, 1]` grid into + // `collision_pairs_indirect`; rewrite it (and derive the offsets) for + // the flat layout. Nothing else consumes the batched form. + self.flatten_batches.call( pass, + 1u32, + collision_pairs_len, + pairs_offsets, collision_pairs_indirect, + batch_indices, + )?; + + self.narrow_phase.call( + pass, + &*collision_pairs_indirect, collision_pairs, - collision_pairs_len, + pairs_offsets, poses, shapes, contacts, @@ -81,9 +99,9 @@ impl GpuNarrowPhase { // separate dispatch so each pass fits 8 storage buffers). self.narrow_phase_deferred.call( pass, - collision_pairs_indirect, + &*collision_pairs_indirect, collision_pairs, - collision_pairs_len, + pairs_offsets, poses, shapes, pfm_pairs, @@ -94,15 +112,21 @@ impl GpuNarrowPhase { indices, )?; - self.init_pfm_pfm_indirect_args - .call(pass, 256u32, pfm_pairs_len, pfm_pairs_indirect)?; + self.flatten_batches.call( + pass, + 1u32, + pfm_pairs_len, + pfm_offsets, + pfm_pairs_indirect, + batch_indices, + )?; self.narrow_phase_pfm_pfm.call( pass, &*pfm_pairs_indirect, contacts, contacts_len, pfm_pairs, - pfm_pairs_len, + pfm_offsets, batch_indices, vertices, indices, diff --git a/src_rbd/pipeline/insertion_removal.rs b/src_rbd/pipeline/insertion_removal.rs index 4cdbec54..6d03941e 100644 --- a/src_rbd/pipeline/insertion_removal.rs +++ b/src_rbd/pipeline/insertion_removal.rs @@ -175,6 +175,10 @@ impl RbdState { BufferUsages::STORAGE | BufferUsages::COPY_SRC, ) .unwrap(); + let pairs_flat_offsets = + Tensor::vector_uninit(backend, num_batches + 1, BufferUsages::STORAGE).unwrap(); + let pfm_flat_offsets = + Tensor::vector_uninit(backend, num_batches + 1, BufferUsages::STORAGE).unwrap(); let old_constraints = Tensor::vector_uninit(backend, collisions_capacity * num_batches, storage).unwrap(); let old_constraint_builders = @@ -289,6 +293,8 @@ impl RbdState { pfm_pairs, pfm_pairs_len, pfm_pairs_indirect, + pairs_flat_offsets, + pfm_flat_offsets, old_constraints, old_constraint_builders, old_constraints_counts, diff --git a/src_rbd/pipeline/rbd_state.rs b/src_rbd/pipeline/rbd_state.rs index f9f2b80f..05fed0cf 100644 --- a/src_rbd/pipeline/rbd_state.rs +++ b/src_rbd/pipeline/rbd_state.rs @@ -187,6 +187,12 @@ pub struct RbdState { pub(super) pfm_pairs: Tensor, pub(super) pfm_pairs_len: Tensor, pub(super) pfm_pairs_indirect: Tensor<[u32; 3]>, + /// Flat-dispatch prefix offsets (`num_batches + 1`) over the per-batch + /// collision-pair / PFM work-lists, rebuilt on the GPU each step by + /// `gpu_flatten_batches_dispatch` so the narrow-phase kernels can pack + /// items from many batches into full warps. + pub(super) pairs_flat_offsets: Tensor, + pub(super) pfm_flat_offsets: Tensor, pub(super) contacts: Tensor, pub(super) contacts_len: Tensor, pub(super) contacts_indirect: Tensor<[u32; 3]>, diff --git a/src_rbd/pipeline/rbd_state_from_rapier.rs b/src_rbd/pipeline/rbd_state_from_rapier.rs index 3f7acaba..9cad90d4 100644 --- a/src_rbd/pipeline/rbd_state_from_rapier.rs +++ b/src_rbd/pipeline/rbd_state_from_rapier.rs @@ -657,6 +657,10 @@ impl RbdState { BufferUsages::STORAGE | BufferUsages::COPY_SRC, ) .unwrap(); + let pairs_flat_offsets = + Tensor::vector_uninit(backend, num_batches + 1, BufferUsages::STORAGE).unwrap(); + let pfm_flat_offsets = + Tensor::vector_uninit(backend, num_batches + 1, BufferUsages::STORAGE).unwrap(); let old_constraints = Tensor::vector_uninit( backend, capacities.collisions_capacity * num_batches, @@ -839,6 +843,8 @@ impl RbdState { pfm_pairs, pfm_pairs_len, pfm_pairs_indirect, + pairs_flat_offsets, + pfm_flat_offsets, old_constraints, old_constraint_builders, old_constraints_counts, diff --git a/src_rbd/pipeline/rbd_step.rs b/src_rbd/pipeline/rbd_step.rs index b98c6fb6..9d4e4cba 100644 --- a/src_rbd/pipeline/rbd_step.rs +++ b/src_rbd/pipeline/rbd_step.rs @@ -251,8 +251,8 @@ impl RbdPipeline { &state.vertex_buffers, &state.index_buffers, &state.collision_pairs, - &state.collision_pairs_len, - &state.collision_pairs_indirect, + &mut state.collision_pairs_len, + &mut state.collision_pairs_indirect, &mut state.contacts, &mut state.contacts_len, &mut state.contacts_indirect, @@ -265,6 +265,8 @@ impl RbdPipeline { &state.collider_materials, &state.prediction, self.contact_reduction, + &mut state.pairs_flat_offsets, + &mut state.pfm_flat_offsets, )?; drop(pass); diff --git a/src_rbd_shaders/broad_phase/narrow_phase.rs b/src_rbd_shaders/broad_phase/narrow_phase.rs index 78e03edf..f8948580 100644 --- a/src_rbd_shaders/broad_phase/narrow_phase.rs +++ b/src_rbd_shaders/broad_phase/narrow_phase.rs @@ -153,6 +153,65 @@ pub fn gpu_reduce_contacts( } } +/// Builds the flat-dispatch layout for a per-batch work-list: exclusive prefix +/// offsets (so item `t` of the flat range maps back to a batch via +/// [`find_batch`]) and the matching 1-D indirect grid. +/// +/// This replaces the max-over-batches indirect grids for the narrow-phase +/// kernels: with `[max/64, num_batches, 1]` every batch rounds its handful of +/// pairs up to a full 64-lane workgroup (a robot env has ~7 pairs → ~11% lane +/// occupancy, thousands of near-empty workgroups). The flat grid packs items +/// from consecutive batches into the same warps: `[total/64, 1, 1]`. +/// +/// Serial over batches in one thread — same pattern (and cost) as the existing +/// `gpu_narrow_phase_init_contacts_dispatch` max-scan. +#[spirv_bindgen] +#[spirv(compute(threads(1)))] +pub fn gpu_flatten_batches_dispatch( + // NOTE: `lens` is mutable only for `atomic_load_u32` (see the note on + // `gpu_narrow_phase_init_contacts_dispatch`). + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] lens: &mut [u32], + // `num_batches + 1` entries; `offsets[num_batches]` is the total. + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] offsets: &mut [u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] indirect_args: &mut [u32; 3], + #[spirv(uniform, descriptor_set = 0, binding = 3)] batch_ids: &BatchIndices, +) { + let num_batches = lens.len(); + // Same clamp as the consuming kernels: a batch's list may overflow its + // capacity slot; the overflowing tail was never written and must not be + // walked. + let capacity = batch_ids.contacts_batch_capacity; + let mut total = 0u32; + for i in 0..num_batches { + offsets.write(i, total); + total += atomic_load_u32(lens.at_mut(i)).min(capacity); + } + offsets.write(num_batches, total); + *indirect_args.at_mut(0) = total.div_ceil(WORKGROUP_SIZE); + *indirect_args.at_mut(1) = 1; + *indirect_args.at_mut(2) = 1; +} + +/// Largest `b` with `offsets[b] <= t` — the batch owning flat item `t`. +/// Invariant: `offsets[0] == 0 <= t < offsets[num_batches]`. +fn find_batch(offsets: &[u32], num_batches: u32, t: u32) -> u32 { + let mut lo = 0u32; + let mut hi = num_batches; + // Bounded loop instead of `while` (see the trimesh BVH walk for why). + for _ in 0..32 { + if lo + 1 >= hi { + break; + } + let mid = (lo + hi) / 2; + if offsets.read(mid as usize) <= t { + lo = mid; + } else { + hi = mid; + } + } + lo +} + const PREDICTION: f32 = 2.0e-3; // TODO: make the prediction configurable. /// Narrow phase, pass 1 of 2: analytic shape-shape contacts for ball / cuboid @@ -166,7 +225,10 @@ pub fn gpu_narrow_phase_shape_shape( #[spirv(global_invocation_id)] invocation_id: UVec3, #[spirv(num_workgroups)] num_workgroups: UVec3, #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] collision_pairs: &[CollisionPair], - #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] collision_pairs_len: &[u32], + // Flat-dispatch prefix offsets from `gpu_flatten_batches_dispatch` + // (`num_batches + 1` entries; replaces the per-batch `collision_pairs_len`, + // which it already folds in, clamped to capacity). + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] pairs_offsets: &[u32], #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] poses: &[Pose], #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] shapes: &[Shape], #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] contacts: &mut [IndexedManifold], @@ -182,23 +244,25 @@ pub fn gpu_narrow_phase_shape_shape( #[spirv(uniform, descriptor_set = 0, binding = 9)] prediction: &f32, ) { 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; - let collision_pairs = batch_ids.contact_batch(batch_id, collision_pairs); - let poses = batch_ids.coll_batch(batch_id, poses); - let shapes = batch_ids.coll_batch(batch_id, shapes); - let collider_materials = batch_ids.coll_batch(batch_id, collider_materials); - let mut contacts = batch_ids.contact_batch_mut(batch_id, contacts); - let contacts_len = contacts_len.at_mut(batch_id as usize); + // Flat over all batches' pairs: consecutive lanes take consecutive pairs + // regardless of which batch owns them, so warps stay packed even when each + // batch only has a handful. + let num_batches = pairs_offsets.len() - 1; + let total = pairs_offsets.read(num_batches); - // NOTE: `collision_pairs_len` might be greater than `contacts_batch_apacity` if the - // narrow-phase found more pairs than the buffer can contain. - let len = collision_pairs_len - .read(batch_id as usize) - .min(contacts_batch_capacity as u32); + for t in StepRng::new(invocation_id.x..total, num_threads) { + let batch_id = find_batch(pairs_offsets, num_batches as u32, t); + let i = t - pairs_offsets.read(batch_id as usize); + + let collision_pairs = batch_ids.contact_batch(batch_id, collision_pairs); + let poses = batch_ids.coll_batch(batch_id, poses); + let shapes = batch_ids.coll_batch(batch_id, shapes); + let collider_materials = batch_ids.coll_batch(batch_id, collider_materials); + let mut contacts = SliceMut(&mut *contacts, batch_ids.contacts_start(batch_id)); + let contacts_len = contacts_len.at_mut(batch_id as usize); - 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. Pair ids are @@ -289,7 +353,8 @@ pub fn gpu_narrow_phase_shape_shape_deferred( #[spirv(global_invocation_id)] invocation_id: UVec3, #[spirv(num_workgroups)] num_workgroups: UVec3, #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] collision_pairs: &[CollisionPair], - #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] collision_pairs_len: &[u32], + // Flat-dispatch prefix offsets (see `gpu_narrow_phase_shape_shape`). + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] pairs_offsets: &[u32], #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] poses: &[Pose], #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] shapes: &[Shape], #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] @@ -304,25 +369,25 @@ pub fn gpu_narrow_phase_shape_shape_deferred( #[spirv(uniform, descriptor_set = 0, binding = 7)] prediction: &f32, ) { 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; - - let collision_pairs = batch_ids.contact_batch(batch_id, collision_pairs); - let poses = batch_ids.coll_batch(batch_id, poses); - let shapes = batch_ids.coll_batch(batch_id, shapes); - let mut pfm_pairs = batch_ids.contact_batch_mut(batch_id, pfm_pairs); - let pfm_pairs_len = pfm_pairs_len.at_mut(batch_id as usize); - let len = collision_pairs_len - .read(batch_id as usize) - .min(contacts_batch_capacity as u32); + let num_batches = pairs_offsets.len() - 1; + let total = pairs_offsets.read(num_batches); // NOTE: same-body collider pairs are *not* filtered in this pass — it is // already at the 8-storage-buffer WebGPU limit and can't take the // `collider_parent` binding. The complex pairs it emits are filtered // downstream in `gpu_narrow_phase_pfm_pfm` (which has room) before any // contact is written. - for i in StepRng::new(invocation_id.x..len, num_threads) { + for t in StepRng::new(invocation_id.x..total, num_threads) { + let batch_id = find_batch(pairs_offsets, num_batches as u32, t); + let i = t - pairs_offsets.read(batch_id as usize); + + let collision_pairs = batch_ids.contact_batch(batch_id, collision_pairs); + let poses = batch_ids.coll_batch(batch_id, poses); + let shapes = batch_ids.coll_batch(batch_id, shapes); + let mut pfm_pairs = SliceMut(&mut *pfm_pairs, batch_ids.contacts_start(batch_id)); + let pfm_pairs_len = pfm_pairs_len.at_mut(batch_id as usize); + let pair = collision_pairs[i as usize]; let shape1 = &shapes[pair.colliders.x as usize]; let shape2 = &shapes[pair.colliders.y as usize]; @@ -645,7 +710,9 @@ pub fn gpu_narrow_phase_pfm_pfm( #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] contacts: &mut [IndexedManifold], #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] contacts_len: &mut [u32], #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] pfm_pairs: &[NarrowPhasePfmPair], - #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] pfm_pairs_len: &[u32], + // Flat-dispatch prefix offsets over the per-batch PFM work-lists (see + // `gpu_narrow_phase_shape_shape`; replaces the per-batch `pfm_pairs_len`). + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] pfm_offsets: &[u32], // NOTE: we assume that max_pfm_pairs == contacts_batch_capacity // And we assume all batch dimensions are given the same buffer allocation sizes // (i.e. the same `contacts_batch_capacity`). @@ -662,20 +729,20 @@ pub fn gpu_narrow_phase_pfm_pfm( #[spirv(uniform, descriptor_set = 0, binding = 9)] prediction: &f32, ) { 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; - let mut contacts = batch_ids.contact_batch_mut(batch_id, contacts); - let collider_materials = batch_ids.coll_batch(batch_id, collider_materials); - let pfm_pairs = batch_ids.contact_batch(batch_id, pfm_pairs); - let contacts_len = contacts_len.at_mut(batch_id as usize); - // The producer counter can exceed the allocation on overflow (writes are - // skipped past capacity); clamp so we never read uninitialized slots. - let pfm_pairs_len = pfm_pairs_len - .read(batch_id as usize) - .min(contacts_batch_capacity as u32); - - for i in StepRng::new(invocation_id.x..pfm_pairs_len, num_threads) { + let num_batches = pfm_offsets.len() - 1; + let total = pfm_offsets.read(num_batches); + + for t in StepRng::new(invocation_id.x..total, num_threads) { + let batch_id = find_batch(pfm_offsets, num_batches as u32, t); + let i = t - pfm_offsets.read(batch_id as usize); + + let mut contacts = SliceMut(&mut *contacts, batch_ids.contacts_start(batch_id)); + let collider_materials = batch_ids.coll_batch(batch_id, collider_materials); + let pfm_pairs = batch_ids.contact_batch(batch_id, pfm_pairs); + let contacts_len = contacts_len.at_mut(batch_id as usize); + let pair = pfm_pairs[i as usize]; // Resolve the parent rigid-bodies and skip same-body collider pairs. This // is where the deferred (PFM / trimesh / polyline) pairs get the same-body From 673d80497de2c5e34e6533bc3a938140b60dd47b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Sun, 16 Aug 2026 16:22:11 +0200 Subject: [PATCH 14/41] fix(rbd): restore the contacts capacity binding and import atomic_load_u32 for the flat dispatch Completes #21 --- src_rbd_shaders/broad_phase/narrow_phase.rs | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/src_rbd_shaders/broad_phase/narrow_phase.rs b/src_rbd_shaders/broad_phase/narrow_phase.rs index f8948580..b0b17eb0 100644 --- a/src_rbd_shaders/broad_phase/narrow_phase.rs +++ b/src_rbd_shaders/broad_phase/narrow_phase.rs @@ -16,7 +16,10 @@ use crate::{PaddedVector, Pose, Vector}; use khal_std::glamx::UVec3; use khal_std::index::MaybeIndexUnchecked; use khal_std::macros::{spirv, spirv_bindgen}; -use khal_std::{iter::StepRng, sync::atomic_add_u32}; +use khal_std::{ + iter::StepRng, + sync::{atomic_add_u32, atomic_load_u32}, +}; use super::lbvh::{MAX_REDUCE_LANES, max_len_indirect_args}; use crate::broad_phase::CollisionPair; @@ -155,7 +158,7 @@ pub fn gpu_reduce_contacts( /// Builds the flat-dispatch layout for a per-batch work-list: exclusive prefix /// offsets (so item `t` of the flat range maps back to a batch via -/// [`find_batch`]) and the matching 1-D indirect grid. +/// `find_batch`) and the matching 1-D indirect grid. /// /// This replaces the max-over-batches indirect grids for the narrow-phase /// kernels: with `[max/64, num_batches, 1]` every batch rounds its handful of @@ -369,6 +372,8 @@ pub fn gpu_narrow_phase_shape_shape_deferred( #[spirv(uniform, descriptor_set = 0, binding = 7)] prediction: &f32, ) { let num_threads = num_workgroups.x * WORKGROUP_SIZE; + // Every batch is allocated the same capacity, so this is batch-independent. + let contacts_batch_capacity = batch_ids.contacts_batch_capacity as usize; let num_batches = pairs_offsets.len() - 1; let total = pairs_offsets.read(num_batches); From e01c065c4fb7a50075553698fc4a72281b53c1a5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Sun, 16 Aug 2026 21:07:39 +0200 Subject: [PATCH 15/41] fix(rbd): drop the stale 2mm PREDICTION constant reintroduced by the flat-dispatch port Completes #21 --- src_rbd_shaders/broad_phase/narrow_phase.rs | 2 -- 1 file changed, 2 deletions(-) diff --git a/src_rbd_shaders/broad_phase/narrow_phase.rs b/src_rbd_shaders/broad_phase/narrow_phase.rs index b0b17eb0..11f28253 100644 --- a/src_rbd_shaders/broad_phase/narrow_phase.rs +++ b/src_rbd_shaders/broad_phase/narrow_phase.rs @@ -215,8 +215,6 @@ fn find_batch(offsets: &[u32], num_batches: u32, t: u32) -> u32 { lo } -const PREDICTION: f32 = 2.0e-3; // TODO: make the prediction configurable. - /// Narrow phase, pass 1 of 2: analytic shape-shape contacts for ball / cuboid /// pairs, written straight into the `contacts` buffer. /// From 7fe25f10a7d52213de2038cabc72a48e8c87ccd2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Fri, 21 Aug 2026 18:33:52 +0200 Subject: [PATCH 16/41] fix mpm feature-gating --- src/pipeline.rs | 5 +++++ src/state.rs | 40 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 45 insertions(+) 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 de2c0810..39f1a752 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::{ @@ -28,7 +30,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 +42,7 @@ pub struct NexusCapacities { /// Rigid-body solver capacities. pub rbd: RbdCapacities, /// MPM solver capacities. + #[cfg(feature = "mpm")] pub mpm: MpmCapacities, } @@ -57,6 +62,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 +73,7 @@ impl NexusCapacities { self } + #[cfg(feature = "mpm")] pub fn mpm_particles(mut self, capacity: u32) -> Self { self.mpm.particles_capacity = capacity; self @@ -114,6 +121,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 +132,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 +140,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 +193,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 +202,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 +232,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 +250,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 +264,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 +274,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 +286,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 +308,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() @@ -343,6 +372,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(); } @@ -352,6 +382,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; @@ -484,6 +515,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; } @@ -584,6 +616,7 @@ impl NexusState { } self.rbd_dirty = true; } + #[cfg(feature = "mpm")] if couplings.iter().any(|c| *c != RbdCoupling::None) { self.mpm_dirty = true; } @@ -703,6 +736,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, @@ -720,6 +754,7 @@ impl NexusState { } /// Appends more particles to an existing chunk (`O(added)`). + #[cfg(feature = "mpm")] pub fn extend_chunk( &mut self, backend: &GpuBackend, @@ -740,11 +775,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, @@ -764,6 +801,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, @@ -790,6 +828,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, @@ -960,6 +999,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(); From c2a18dc11bcd5ffec9c9402b04af3ac653aca717 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Fri, 21 Aug 2026 20:19:07 +0200 Subject: [PATCH 17/41] feat(rbd): expose dof_state_mut, links_static, joint_constraints and link_of_body --- .../multibody/multibody_from_rapier.rs | 2 + src_rbd/dynamics/multibody/multibody_set.rs | 43 +++++++++++++++++++ 2 files changed, 45 insertions(+) diff --git a/src_rbd/dynamics/multibody/multibody_from_rapier.rs b/src_rbd/dynamics/multibody/multibody_from_rapier.rs index b9025f2c..827fb34a 100644 --- a/src_rbd/dynamics/multibody/multibody_from_rapier.rs +++ b/src_rbd/dynamics/multibody/multibody_from_rapier.rs @@ -586,6 +586,8 @@ 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, 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 3d9c301a..5132044b 100644 --- a/src_rbd/dynamics/multibody/multibody_set.rs +++ b/src_rbd/dynamics/multibody/multibody_set.rs @@ -96,6 +96,11 @@ 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]>, + /// 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). @@ -229,6 +234,15 @@ impl GpuMultibodySet { &self.dof_state } + /// 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 { @@ -442,6 +456,14 @@ impl GpuMultibodySet { &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 { @@ -605,6 +627,27 @@ 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 + } + + /// `[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 { From 0929b65ab9918d8fcbb022ca4616aa33341db255 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Fri, 21 Aug 2026 23:41:28 +0200 Subject: [PATCH 18/41] feat(rbd): env-reset primitives, GPU motor scatter, contact sensors, actuator delay, encoded step, substep-refresh cadence and per-DoF armature/frictionloss --- src_rbd/dynamics/mod.rs | 4 +- src_rbd/dynamics/multibody/env_reset.rs | 402 ++++++++++++++++++ src_rbd/dynamics/multibody/mod.rs | 2 + .../multibody/multibody_from_rapier.rs | 49 ++- src_rbd/dynamics/multibody/multibody_set.rs | 349 ++++++++++++++- .../dynamics/multibody/multibody_solver.rs | 96 ++++- src_rbd/dynamics/solver.rs | 1 + src_rbd/pipeline/insertion_removal.rs | 2 + src_rbd/pipeline/mod.rs | 4 +- src_rbd/pipeline/rbd_state.rs | 245 +++++++++++ src_rbd/pipeline/rbd_state_from_rapier.rs | 7 +- src_rbd/pipeline/rbd_step.rs | 88 +++- .../dynamics/multibody/contact_sensor.rs | 68 +++ .../dynamics/multibody/env_reset.rs | 198 +++++++++ .../dynamics/multibody/gravity_and_lu.rs | 51 ++- .../dynamics/multibody/joint_constraints.rs | 170 +++++++- src_rbd_shaders/dynamics/multibody/mod.rs | 10 + .../dynamics/multibody/scatter_motor.rs | 113 +++++ 18 files changed, 1810 insertions(+), 49 deletions(-) create mode 100644 src_rbd/dynamics/multibody/env_reset.rs create mode 100644 src_rbd_shaders/dynamics/multibody/contact_sensor.rs create mode 100644 src_rbd_shaders/dynamics/multibody/env_reset.rs create mode 100644 src_rbd_shaders/dynamics/multibody/scatter_motor.rs 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..da4bf468 --- /dev/null +++ b/src_rbd/dynamics/multibody/env_reset.rs @@ -0,0 +1,402 @@ +//! 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, 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, +} + +/// 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: one staging upload and one + /// scatter dispatch, with no readback and no per-element strided writes. + 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.links, + &tpl.dofs, + &tpl.flags, + &t_resets, + &t_offs, + &t_vels, + &mut self.links_workspace, + &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 827fb34a..e149abd3 100644 --- a/src_rbd/dynamics/multibody/multibody_from_rapier.rs +++ b/src_rbd/dynamics/multibody/multibody_from_rapier.rs @@ -522,18 +522,22 @@ impl GpuMultibodySet { 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`. Frictionloss has no rapier + // counterpart, so it starts at zero (off) and is filled in by + // `GpuMultibodySet::set_dof_frictionloss`. 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.resize(7 * n, 0.0); + debug_assert_eq!(buf.len(), 7 * n); Tensor::vector(backend, &buf, storage).unwrap() }, gen_forces: Tensor::vector( @@ -588,6 +592,41 @@ impl GpuMultibodySet { 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 5132044b..3ccb04ce 100644 --- a/src_rbd/dynamics/multibody/multibody_set.rs +++ b/src_rbd/dynamics/multibody/multibody_set.rs @@ -8,6 +8,7 @@ use crate::shaders::dynamics::{ MultibodyLinkStatic, MultibodyLinkWorkspace, RbdSimParams, }; use crate::shaders::utils::BatchIndices; +use khal::Shader; use khal::BufferUsages; use khal::backend::{Backend, GpuBackend, GpuBackendError}; use rapier3d::prelude::JointAxis; @@ -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,6 +82,15 @@ pub struct GpuMultibodySet { pub(super) coriolis_entries_per_batch: u32, pub(super) i_coriolis_dt_entries_per_batch: u32, pub(super) implicit_coriolis: 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, @@ -96,6 +142,29 @@ 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]>, @@ -228,12 +297,29 @@ 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 @@ -446,6 +532,82 @@ 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 @@ -637,6 +799,189 @@ impl GpuMultibodySet { self.joint_constraints_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) and is transposed here into the + /// batch-interleaved layout the kernels index. + /// + /// A post-build override for callers whose rapier scenes carry no armature, + /// or that randomize it per environment; the build path already seeds this + /// section from `mb.armature()`. + pub fn set_dof_armature(&mut self, backend: &GpuBackend, values: &[f32]) { + self.write_dof_section(backend, 2, values, "armature"); + } + + /// Overwrites the per-DoF Coulomb joint friction section of + /// [`Self::dof_state`] (MJCF `frictionloss`, N·m), applied by the gravity + /// kernels as `-fl·sign(q̇)`. Zero, the default, disables it. `values` uses + /// the same env-major layout as [`Self::set_dof_armature`]. + pub fn set_dof_frictionloss(&mut self, backend: &GpuBackend, values: &[f32]) { + self.write_dof_section(backend, 6, values, "frictionloss"); + } + + /// 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. diff --git a/src_rbd/dynamics/multibody/multibody_solver.rs b/src_rbd/dynamics/multibody/multibody_solver.rs index f52a414a..b976bdba 100644 --- a/src_rbd/dynamics/multibody/multibody_solver.rs +++ b/src_rbd/dynamics/multibody/multibody_solver.rs @@ -8,8 +8,8 @@ use crate::shaders::dynamics::{ GpuMbComputeSolveBounds, GpuMbFinalizeContactConstraints, GpuMbFinalizeImpulseJointConstraints, GpuMbGravityAndLu, GpuMbGravityAndLuT1, GpuMbGravityAndLuT8, GpuMbGravityAndLuT16, GpuMbGravityAndLuT32, GpuMbInitContactConstraints, GpuMbInitJointConstraints, GpuMbIntegrate, - GpuMbIntegrateVelocities, GpuMbRemoveImpulseJointConstraintBias, GpuMbSeedContactRestitution, - GpuMbSnapshotContactWarmstart, GpuMbSolveConstraints, GpuMbSolveContactsDelassus, + GpuMbIntegrateVelocities, GpuMbRemoveImpulseJointConstraintBias, GpuMbDelayTick, GpuMbRefreshJointConstraints, GpuMbSeedContactRestitution, + GpuMbSenseContactImpulses, GpuMbSnapshotContactWarmstart, GpuMbSolveConstraints, GpuMbSolveContactsDelassus, GpuMbSolveImpulseJointConstraints, GpuMbSolveJoints, GpuMbStashContactsLen, GpuMbTransferContactWarmstart, GpuMbUpdateImpulseJointConstraints, GpuMbWarmstartContactConstraints, Velocity, WorldMassProperties, @@ -33,6 +33,14 @@ pub struct GpuMultibodySolver { compute_dynamics_pre: GpuMbComputeDynamicsPre, init_joint_with_bias: GpuMbInitJointConstraints, 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 +232,38 @@ 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, + &mut 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 { @@ -318,6 +362,7 @@ impl GpuMultibodySolver { &mut mb.joint_constraints, &mut mb.joint_constraint_columns, &mb.dof_couplings, + &mut mb.motor_delay_state, &mb.constraint_softness, args.batch_indices, )?; @@ -585,8 +630,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 +695,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, 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/insertion_removal.rs b/src_rbd/pipeline/insertion_removal.rs index 6d03941e..aba87576 100644 --- a/src_rbd/pipeline/insertion_removal.rs +++ b/src_rbd/pipeline/insertion_removal.rs @@ -255,6 +255,8 @@ impl RbdState { num_solver_iterations, sim_params: Tensor::vector(backend, &all_sim_params, BufferUsages::STORAGE).unwrap(), 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..7b24c0ee 100644 --- a/src_rbd/pipeline/mod.rs +++ b/src_rbd/pipeline/mod.rs @@ -11,4 +11,6 @@ mod rbd_state_from_rapier; mod rbd_step; pub use rbd_state::{RbdCapacities, RbdResizePolicy, RbdState, RunStats}; -pub use rbd_step::RbdPipeline; +#[cfg(feature = "dim3")] +pub use rbd_state::RbdSnapshot; +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 05fed0cf..e7a49fca 100644 --- a/src_rbd/pipeline/rbd_state.rs +++ b/src_rbd/pipeline/rbd_state.rs @@ -131,6 +131,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, @@ -347,6 +351,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 @@ -570,3 +586,232 @@ 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, with + /// no GPU to CPU readback. This is what removes the handful of per-reset + /// sync stalls that otherwise dominate reset cost on the WebGPU backend. + 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 9cad90d4..12debe65 100644 --- a/src_rbd/pipeline/rbd_state_from_rapier.rs +++ b/src_rbd/pipeline/rbd_state_from_rapier.rs @@ -783,7 +783,10 @@ impl RbdState { 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(), + 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, @@ -796,7 +799,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 9d4e4cba..f0b04622 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, @@ -65,11 +71,52 @@ 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(); // Make sure the color index uniforms are up-to-date. @@ -84,8 +131,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")] { @@ -103,7 +148,7 @@ impl RbdPipeline { gravity: &state.gravity, }; self.multibody_solver.init_step( - &mut encoder, + &mut *encoder, timestamps.as_deref_mut(), &mut state.multibodies, &mut args, @@ -164,12 +209,12 @@ impl RbdPipeline { &state.prediction, )?; 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, @@ -182,8 +227,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( @@ -194,7 +239,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()); } @@ -215,7 +259,7 @@ impl RbdPipeline { )?; drop(pass); - backend.submit(encoder)?; + split(&mut *encoder)?; } } @@ -228,9 +272,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. @@ -239,7 +290,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()); @@ -271,8 +321,7 @@ impl RbdPipeline { drop(pass); if !merge_submits { - backend.submit(encoder)?; - encoder = backend.begin_encoding(); + split(&mut *encoder)?; } } @@ -427,8 +476,7 @@ impl RbdPipeline { drop(pass); } if !merge_submits { - backend.submit(encoder)?; - encoder = backend.begin_encoding(); + split(&mut *encoder)?; } } @@ -494,7 +542,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, @@ -505,9 +553,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/dynamics/multibody/contact_sensor.rs b/src_rbd_shaders/dynamics/multibody/contact_sensor.rs new file mode 100644 index 00000000..ee24be98 --- /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::{ + MB_CONTACT_KIND_NORMAL, MAX_MB_CONTACT_CONSTRAINTS_PER_MB, 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..6dd0dea0 --- /dev/null +++ b/src_rbd_shaders/dynamics/multibody/env_reset.rs @@ -0,0 +1,198 @@ +//! 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. +//! That is one staging upload plus one dispatch per reset, instead of the +//! hundreds of strided `write_buffer`s the interleaved layout would otherwise +//! force: an env's data is strided across the whole buffer (element `intra` of +//! batch `b` lives at `intra · num_batches + b`, workspace quads at +//! `(link · WS_QUADS + q) · num_batches + b`). The staging blob is exactly the +//! `num_batches = 1` interleaving, so its source indices are the flat `0..len`. + +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. The DoF velocity +/// section comes from `dof_vels` (host-randomized reset velocities, or zeros), +/// replacing a per-DoF strided `write_buffer` loop. +/// +/// 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)] + templates_links: &[MultibodyLinkStatic], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] templates_dofs: &[f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] link_flags: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] resets: &[UVec4], + #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] offsets: &[Vec4], + #[spirv(storage_buffer, descriptor_set = 0, binding = 6)] dof_vels: &[f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 7)] links_workspace: &mut [Vec4], + #[spirv(storage_buffer, descriptor_set = 0, binding = 8)] + links_static: &mut [MultibodyLinkStatic], + #[spirv(storage_buffer, descriptor_set = 0, binding = 9)] dof_values: &mut [f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 10)] dof_state: &mut [f32], + // x = num_batches, y = links_per_batch, z = dofs_per_batch, w = num_resets. + #[spirv(uniform, descriptor_set = 0, binding = 11)] 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; + 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); + } + if i < lpb { + links_static.write( + (i * nb + env) as usize, + templates_links.read((t * lpb + i) as usize), + ); + } + if i < dpb { + // Generalized coords are translation-invariant: the free root's world + // position lives in the workspace coords quad handled above. + 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..9ebb63f5 100644 --- a/src_rbd_shaders/dynamics/multibody/gravity_and_lu.rs +++ b/src_rbd_shaders/dynamics/multibody/gravity_and_lu.rs @@ -129,6 +129,10 @@ pub fn gpu_mb_gravity_and_lu( let damping_slice = batch_ids .ib(batch_id, dof_state) .offset(batch_ids.dof_batch_capacity as usize + gen_base); + // Coulomb joint friction (MJCF `frictionloss`) sits in the 7th section. + let frictionloss_slice = batch_ids + .ib(batch_id, dof_state) + .offset(6 * batch_ids.dof_batch_capacity as usize + gen_base); let stiffness_slice = batch_ids .ib(batch_id, dof_state) .offset(3 * batch_ids.dof_batch_capacity as usize + gen_base); @@ -283,7 +287,18 @@ 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]; + // Coulomb joint friction: -fl·sign(v), MuJoCo `frictionloss` semantics. + // Zero (the default) leaves the DoF untouched. + let fl = frictionloss_slice[i as usize]; + let fric = if v > 0.0 { + -fl + } else if v < 0.0 { + fl + } else { + 0.0 + }; + gen_forces.write(idx, cur - damping_slice[i as usize] * v + fric); } workgroup_memory_barrier_with_group_sync(); @@ -450,6 +465,10 @@ fn gravity_and_lu_packed_impl 0.0 { + -fl + } else if v < 0.0 { + fl + } else { + 0.0 + }; + gen_forces.write(idx, cur - damping_slice[i as usize] * v + fric); } workgroup_memory_barrier_with_group_sync(); @@ -770,6 +800,10 @@ pub fn gpu_mb_gravity_and_lu_t1( let damping_slice = batch_ids .ib(batch_id, dof_state) .offset(batch_ids.dof_batch_capacity as usize + gen_base); + // Coulomb joint friction (MJCF `frictionloss`) sits in the 7th section. + let frictionloss_slice = batch_ids + .ib(batch_id, dof_state) + .offset(6 * batch_ids.dof_batch_capacity as usize + gen_base); let stiffness_slice = batch_ids .ib(batch_id, dof_state) .offset(3 * batch_ids.dof_batch_capacity as usize + gen_base); @@ -896,7 +930,18 @@ pub fn gpu_mb_gravity_and_lu_t1( for i in 0..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]; + // Coulomb joint friction: -fl·sign(v), MuJoCo `frictionloss` semantics. + // Zero (the default) leaves the DoF untouched. + let fl = frictionloss_slice[i as usize]; + let fric = if v > 0.0 { + -fl + } else if v < 0.0 { + fl + } else { + 0.0 + }; + gen_forces.write(idx, cur - damping_slice[i as usize] * v + fric); } // Per-DoF joint springs. diff --git a/src_rbd_shaders/dynamics/multibody/joint_constraints.rs b/src_rbd_shaders/dynamics/multibody/joint_constraints.rs index b0dd32a5..35bf823f 100644 --- a/src_rbd_shaders/dynamics/multibody/joint_constraints.rs +++ b/src_rbd_shaders/dynamics/multibody/joint_constraints.rs @@ -23,6 +23,40 @@ use super::types::{ }; use super::ws_soa::{WsAddr, ws_coord}; +/// Per-batch stride of the actuator-delay state: `[tick, k, prev_target x +/// links_batch_capacity]`. +#[inline] +fn motor_delay_stride(batch_ids: &BatchIndices) -> 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) { @@ -81,6 +115,7 @@ fn emit_joint_constraints( 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 +166,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 +235,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, @@ -368,6 +417,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 +428,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; @@ -442,8 +494,12 @@ pub fn gpu_mb_init_joint_constraints( #[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, + // 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 = 8)] + motor_delay_state: &[f32], + #[spirv(uniform, descriptor_set = 0, binding = 9)] softness: &ConstraintSoftness, + #[spirv(uniform, descriptor_set = 0, binding = 10)] batch_ids: &BatchIndices, ) { const LANES: u32 = 64; @@ -509,6 +565,7 @@ pub fn gpu_mb_init_joint_constraints( softness.dt, softness.joint_erp_inv_dt, softness.joint_cfm_coeff, + motor_delay_state, batch_ids, ); } @@ -550,3 +607,110 @@ pub fn gpu_mb_init_joint_constraints( } } } + +/// Per-substep refresh of the joint limit / motor slots, the cheap alternative +/// to a full rebuild. +/// +/// When the constraint columns and `inv_lhs` are per-step constants (no +/// implicit Coriolis, no per-substep mass-matrix refresh), the only things that +/// change between substeps are the rhs, the limit activity and the accumulated +/// impulse. This recomputes exactly those from the slot's stashed `(link, +/// axis)`, so the full emission walk and the LU back-solves run once per step +/// instead of once per substep. +/// +/// One 64-lane workgroup per (multibody, batch); lanes stride the slots. +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn gpu_mb_refresh_joint_constraints( + #[spirv(workgroup_id)] workgroup_id: UVec3, + #[spirv(local_invocation_id)] local_id: UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] multibody_info: &[MultibodyInfo], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] + links_static: &[MultibodyLinkStatic], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] links_workspace: &[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); + // 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..35690962 100644 --- a/src_rbd_shaders/dynamics/multibody/mod.rs +++ b/src_rbd_shaders/dynamics/multibody/mod.rs @@ -15,12 +15,18 @@ mod compute_dynamics_pre; mod contact_constraints; +mod contact_sensor; +// The RL env-reset primitives (terrain teleport offsets, template blobs) +// are 3D-only, like the RbdSnapshot host API on top of them. +#[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 +34,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); +} From 851f67242f781aaff224d254f3790f3e863439d6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Sat, 22 Aug 2026 11:08:16 +0200 Subject: [PATCH 19/41] fix(rbd): guard against implicit-coriolis drifting from the batch_indices uniform --- src_rbd/dynamics/multibody/multibody_from_rapier.rs | 1 + src_rbd/dynamics/multibody/multibody_set.rs | 8 +++++++- src_rbd/dynamics/multibody/multibody_solver.rs | 9 +++++++++ src_rbd/pipeline/insertion_removal.rs | 2 +- src_rbd/pipeline/rbd_state_from_rapier.rs | 2 +- 5 files changed, 19 insertions(+), 3 deletions(-) diff --git a/src_rbd/dynamics/multibody/multibody_from_rapier.rs b/src_rbd/dynamics/multibody/multibody_from_rapier.rs index e149abd3..c18bc0d8 100644 --- a/src_rbd/dynamics/multibody/multibody_from_rapier.rs +++ b/src_rbd/dynamics/multibody/multibody_from_rapier.rs @@ -498,6 +498,7 @@ 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), multibody_info: Tensor::vector(backend, &all_infos, storage).unwrap(), diff --git a/src_rbd/dynamics/multibody/multibody_set.rs b/src_rbd/dynamics/multibody/multibody_set.rs index 3ccb04ce..e22c0425 100644 --- a/src_rbd/dynamics/multibody/multibody_set.rs +++ b/src_rbd/dynamics/multibody/multibody_set.rs @@ -82,6 +82,11 @@ 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 @@ -685,7 +690,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; diff --git a/src_rbd/dynamics/multibody/multibody_solver.rs b/src_rbd/dynamics/multibody/multibody_solver.rs index b976bdba..dbf5eeb8 100644 --- a/src_rbd/dynamics/multibody/multibody_solver.rs +++ b/src_rbd/dynamics/multibody/multibody_solver.rs @@ -756,6 +756,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/pipeline/insertion_removal.rs b/src_rbd/pipeline/insertion_removal.rs index aba87576..6c5657ac 100644 --- a/src_rbd/pipeline/insertion_removal.rs +++ b/src_rbd/pipeline/insertion_removal.rs @@ -95,7 +95,7 @@ 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) diff --git a/src_rbd/pipeline/rbd_state_from_rapier.rs b/src_rbd/pipeline/rbd_state_from_rapier.rs index 12debe65..4a519ca9 100644 --- a/src_rbd/pipeline/rbd_state_from_rapier.rs +++ b/src_rbd/pipeline/rbd_state_from_rapier.rs @@ -467,7 +467,7 @@ impl RbdState { // Convert multibodies (3D only). #[cfg(feature = "dim3")] - let multibodies = { + let mut multibodies = { let mb_refs: Vec<( &MultibodyJointSet, &HashMap, From 71517831ed8af8d39ac51a1e711680cf7558b4e1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Sat, 22 Aug 2026 16:44:33 +0200 Subject: [PATCH 20/41] feat(rbd): cluster contact manifolds by normal, matching rapier, with a tunable threshold --- src_rbd/broad_phase/narrow_phase.rs | 4 +- src_rbd/pipeline/insertion_removal.rs | 7 ++ src_rbd/pipeline/rbd_state.rs | 19 +++++ src_rbd/pipeline/rbd_state_from_rapier.rs | 7 ++ src_rbd/pipeline/rbd_step.rs | 1 + src_rbd_shaders/broad_phase/narrow_phase.rs | 88 ++++++++++++++++++--- 6 files changed, 112 insertions(+), 14 deletions(-) diff --git a/src_rbd/broad_phase/narrow_phase.rs b/src_rbd/broad_phase/narrow_phase.rs index 2c523c61..b4ca4933 100644 --- a/src_rbd/broad_phase/narrow_phase.rs +++ b/src_rbd/broad_phase/narrow_phase.rs @@ -61,6 +61,7 @@ impl GpuNarrowPhase { // Optional: merge each collider pair's manifolds into one before the // solvers see them. `false` skips the kernel entirely. reduce_contacts: bool, + merge_cos: &Tensor, pairs_offsets: &mut Tensor, pfm_offsets: &mut Tensor, ) -> Result<(), GpuBackendError> { @@ -145,10 +146,11 @@ impl GpuNarrowPhase { contacts_len, batch_indices, prediction, + merge_cos, )?; } #[cfg(not(feature = "dim3"))] - let _ = reduce_contacts; + let _ = (reduce_contacts, merge_cos); self.init_contacts_indirect_args.call( pass, 256u32, diff --git a/src_rbd/pipeline/insertion_removal.rs b/src_rbd/pipeline/insertion_removal.rs index 6c5657ac..cb8f533d 100644 --- a/src_rbd/pipeline/insertion_removal.rs +++ b/src_rbd/pipeline/insertion_removal.rs @@ -143,6 +143,12 @@ impl RbdState { BufferUsages::STORAGE | BufferUsages::UNIFORM, ) .unwrap(); + let contact_merge_cos = Tensor::scalar( + backend, + crate::shaders::broad_phase::COS_MERGE_ANGLE, + BufferUsages::STORAGE | BufferUsages::UNIFORM | BufferUsages::COPY_DST, + ) + .unwrap(); let prediction = Tensor::scalar( backend, all_sim_params[0].prediction_distance(), @@ -282,6 +288,7 @@ impl RbdState { collision_pairs_len_max, num_batches_uniform, prediction, + contact_merge_cos, resize_readback, collision_pairs_indirect, contacts_per_batch_cpu, diff --git a/src_rbd/pipeline/rbd_state.rs b/src_rbd/pipeline/rbd_state.rs index e7a49fca..c5b240ab 100644 --- a/src_rbd/pipeline/rbd_state.rs +++ b/src_rbd/pipeline/rbd_state.rs @@ -167,6 +167,9 @@ pub struct RbdState { /// Contact prediction distance (`RbdSimParams::prediction_distance`), /// consumed by the narrow-phase kernels. pub(super) prediction: Tensor, + /// Cosine of the maximum angle between two contact normals for their + /// manifolds to be clustered together (see `gpu_reduce_contacts`). + pub(super) contact_merge_cos: Tensor, /// `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 @@ -376,6 +379,22 @@ 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) { + self.contact_merge_cos = Tensor::scalar( + backend, + cos, + BufferUsages::STORAGE | BufferUsages::UNIFORM | BufferUsages::COPY_DST, + ) + .unwrap(); + } + /// The gravity uniform shared by every solver kernel. pub fn gravity(&self) -> &Tensor { &self.gravity diff --git a/src_rbd/pipeline/rbd_state_from_rapier.rs b/src_rbd/pipeline/rbd_state_from_rapier.rs index 4a519ca9..a507383d 100644 --- a/src_rbd/pipeline/rbd_state_from_rapier.rs +++ b/src_rbd/pipeline/rbd_state_from_rapier.rs @@ -616,6 +616,12 @@ impl RbdState { BufferUsages::STORAGE | BufferUsages::UNIFORM, ) .unwrap(); + let contact_merge_cos = Tensor::scalar( + backend, + crate::shaders::broad_phase::COS_MERGE_ANGLE, + BufferUsages::STORAGE | BufferUsages::UNIFORM | BufferUsages::COPY_DST, + ) + .unwrap(); let prediction = Tensor::scalar( backend, all_sim_params[0].prediction_distance(), @@ -833,6 +839,7 @@ impl RbdState { collision_pairs_len_max, num_batches_uniform, prediction, + contact_merge_cos, resize_readback, collision_pairs_indirect, contacts_per_batch_cpu, diff --git a/src_rbd/pipeline/rbd_step.rs b/src_rbd/pipeline/rbd_step.rs index f0b04622..81432b7d 100644 --- a/src_rbd/pipeline/rbd_step.rs +++ b/src_rbd/pipeline/rbd_step.rs @@ -315,6 +315,7 @@ impl RbdPipeline { &state.collider_materials, &state.prediction, self.contact_reduction, + &state.contact_merge_cos, &mut state.pairs_flat_offsets, &mut state.pfm_flat_offsets, )?; diff --git a/src_rbd_shaders/broad_phase/narrow_phase.rs b/src_rbd_shaders/broad_phase/narrow_phase.rs index 11f28253..28086891 100644 --- a/src_rbd_shaders/broad_phase/narrow_phase.rs +++ b/src_rbd_shaders/broad_phase/narrow_phase.rs @@ -72,14 +72,62 @@ pub fn gpu_narrow_phase_init_contacts_dispatch( } } +/// 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 all manifolds of a collider pair (e.g. per-triangle trimesh -/// contacts, which share one `colliders` key and one collider-A local frame) -/// into a single `MAX_MANIFOLD_POINTS` manifold via `manifold_reduction`, -/// keeping the deeper manifold's normal. The first record of a pair is kept -/// verbatim, so single-manifold pairs are bit-identical to the unreduced -/// path. Approximations: one normal per merged manifold, greedy merging in -/// emission order. Grid `[1, num_batches, 1]`, serial per batch. +/// merging manifolds that share 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)))] @@ -91,6 +139,9 @@ pub fn gpu_reduce_contacts( // Contact prediction distance: `manifold_reduction` only keeps candidates // within it, exactly as the narrow-phase passes that produced them. #[spirv(uniform, descriptor_set = 0, binding = 3)] prediction: &f32, + // Cosine of the maximum angle between two manifolds' normals for them to + // share a cluster. See [`COS_MERGE_ANGLE`]. + #[spirv(uniform, descriptor_set = 0, binding = 4)] merge_cos: &f32, ) { let batch_id = workgroup_id.y; let capacity = batch_ids.contacts_batch_capacity as usize; @@ -103,18 +154,29 @@ pub fn gpu_reduce_contacts( 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 { - // Pool the two manifolds' points (same collider-A local frame). + 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 { - cand.write(k, out.contact.points_a.read(k)); + pool_dedup(&mut cand, &mut num, out.contact.points_a.read(k), dedup_eps_sq); } for k in 0..nb { - cand.write(na + k, im.contact.points_a.read(k)); + pool_dedup(&mut cand, &mut num, im.contact.points_a.read(k), dedup_eps_sq); } - // Normal of whichever manifold holds the deepest point. + // 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; @@ -134,7 +196,7 @@ pub fn gpu_reduce_contacts( } else { out.contact.normal_a }; - let mut reduced = manifold_reduction(&cand, (na + nb) as u32, normal, *prediction); + 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; From 31694bf461b147a7eda497b4dd3e7eda2eb85d83 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Sun, 23 Aug 2026 10:26:05 +0200 Subject: [PATCH 21/41] feat(rbd): model multibody joint frictionloss as a constraint instead of a force --- crates/examples3d/Cargo.toml | 20 ++ crates/examples3d/test_frictionloss.rs | 183 +++++++++++++++ crates/examples3d/test_frictionloss_caps.rs | 137 ++++++++++++ crates/examples3d/test_frictionloss_stress.rs | 208 ++++++++++++++++++ crates/examples3d/test_refresh_nan.rs | 193 ++++++++++++++++ .../multibody/multibody_from_rapier.rs | 7 +- src_rbd/dynamics/multibody/multibody_set.rs | 97 +++++++- .../dynamics/multibody/multibody_solver.rs | 28 ++- src_rbd/pipeline/rbd_state.rs | 13 ++ src_rbd/pipeline/rbd_step.rs | 11 + .../dynamics/multibody/gravity_and_lu.rs | 48 +--- .../dynamics/multibody/joint_constraints.rs | 95 +++++++- .../dynamics/multibody/solve_constraints.rs | 10 +- src_rbd_shaders/dynamics/multibody/types.rs | 8 + 14 files changed, 978 insertions(+), 80 deletions(-) create mode 100644 crates/examples3d/test_frictionloss.rs create mode 100644 crates/examples3d/test_frictionloss_caps.rs create mode 100644 crates/examples3d/test_frictionloss_stress.rs create mode 100644 crates/examples3d/test_refresh_nan.rs diff --git a/crates/examples3d/Cargo.toml b/crates/examples3d/Cargo.toml index 9d60f77a..0ce62bc3 100644 --- a/crates/examples3d/Cargo.toml +++ b/crates/examples3d/Cargo.toml @@ -54,3 +54,23 @@ path = "all_examples3.rs" #[[bin]] #name = "bench_joints3" #path = "bench_joints3.rs" + +[[bin]] +name = "test_frictionloss" +path = "test_frictionloss.rs" +required-features = ["metal"] + +[[bin]] +name = "test_frictionloss_stress" +path = "test_frictionloss_stress.rs" +required-features = ["metal"] + +[[bin]] +name = "test_frictionloss_caps" +path = "test_frictionloss_caps.rs" +required-features = ["metal"] + +[[bin]] +name = "test_refresh_nan" +path = "test_refresh_nan.rs" +required-features = ["metal"] diff --git a/crates/examples3d/test_frictionloss.rs b/crates/examples3d/test_frictionloss.rs new file mode 100644 index 00000000..af93f044 --- /dev/null +++ b/crates/examples3d/test_frictionloss.rs @@ -0,0 +1,183 @@ +//! Headless probe for MJCF-style joint dry friction (`frictionloss`). +//! +//! A single-link pendulum hinged at the origin, its rod lying along +X so +//! gravity applies a torque `m·g·l` about the hinge. MuJoCo models friction +//! loss as a constraint (a bound on the force friction may generate), not as a +//! `-f·sign(q̇)` force, so the expected behaviour is: +//! +//! * `frictionloss = 0`: the link falls freely. +//! * `frictionloss` above the gravity torque: the link sticks, exactly. A +//! force-based implementation cannot do this; it chatters around `q̇ = 0`. +//! * `frictionloss` below the gravity torque: the link falls, but slower. +//! +//! Run with `cargo run --release --bin test_frictionloss --features metal`. + +use khal::backend::{Backend, GpuBackend}; +use nexus3d::prelude::{NexusPipeline, NexusState, RbdCoupling}; +use nexus3d::rbd::dynamics::RbdSimParams; +use rapier3d::prelude::*; + +const LINK_LEN: f32 = 1.0; +const RAD: f32 = 0.05; +const STEPS: usize = 60; + +/// Builds the one-link pendulum and returns its state. `motor` adds a velocity +/// motor and a limit on the hinge, so the friction rows have to share the +/// constraint bank with them. `joint_frequency`, when given, softens the shared +/// joint constraint softness the friction rows draw their CFM from. +fn make_state(motor: bool, joint_frequency: Option) -> NexusState { + let mut state = NexusState::default(); + if let Some(hz) = joint_frequency { + let mut params = RbdSimParams::default(); + params.joint_natural_frequency = hz; + state.set_rbd_sim_params(0, params); + } + let no_coupling = RbdCoupling::None; + + let root = RigidBodyBuilder::fixed().build(); + let root_collider = ColliderBuilder::cuboid(RAD, RAD, RAD) + .collision_groups(InteractionGroups::none()) + .build(); + let root_handle = state.insert_rigid_body(root, root_collider, no_coupling); + + let link = RigidBodyBuilder::dynamic() + .translation(Vec3::new(LINK_LEN, 0.0, 0.0)) + .build(); + let collider = ColliderBuilder::cuboid(LINK_LEN * 0.5, RAD, RAD) + .collision_groups(InteractionGroups::none()) + .build(); + let link_handle = state.insert_rigid_body(link, collider, no_coupling); + + // Hinge about Z at the origin, so gravity (-Y) torques the joint. + let mut builder = RevoluteJointBuilder::new(Vec3::Z) + .local_anchor1(Vec3::ZERO) + .local_anchor2(Vec3::new(-LINK_LEN, 0.0, 0.0)); + if motor { + builder = builder + .limits([-2.0, 2.0]) + .motor_velocity(2.0, 0.0) + .motor_max_force(100.0); + } + let joint = builder.build(); + state.insert_multibody_joint(root_handle, link_handle, joint); + + state +} + +/// Steps the pendulum for `STEPS` frames and returns `(drop, |q̇|)`: how far +/// the link's centre of mass fell, and the joint velocity at the end. +async fn run_case( + backend: &GpuBackend, + frictionloss: f32, + motor: bool, + joint_frequency: Option, +) -> Result<(f32, f32), khal::backend::GpuBackendError> { + let mut state = make_state(motor, joint_frequency); + let mut pipeline = NexusPipeline::default(); + + // The frictionloss slots are reserved on the first non-zero write, which + // needs the GPU state to exist: finalize before setting it. + state.finalize(backend).await?; + let rbd = state.rbd.as_mut().expect("no rbd state"); + let ndofs = rbd.multibodies().dofs_per_batch() as usize; + rbd.set_dof_frictionloss(backend, &vec![frictionloss; ndofs]); + + for _ in 0..STEPS { + pipeline.simulate(backend, &mut state, None).await?; + } + + let rbd = state.rbd.as_ref().expect("no rbd state"); + let poses: Vec = + backend.slow_read_vec(rbd.body_poses().buffer()).await?; + let dof_state: Vec = backend + .slow_read_vec(rbd.multibodies().dof_state().buffer()) + .await?; + // Body 1 is the link; it starts level with the hinge, so any fall shows up + // as a negative y. + Ok((-poses[1].translation.y, dof_state[0].abs())) +} + +fn main() -> anyhow::Result<()> { + pollster::block_on(run()) +} + +async fn run() -> anyhow::Result<()> { + let backend = + GpuBackend::Metal(khal::backend::Metal::new().map_err(|e| anyhow::anyhow!("{e:?}"))?); + + // Gravity torque about the hinge for a rod of half-length `LINK_LEN/2`. + // The collider density is rapier's default (1.0). + let mass = 2.0 * (LINK_LEN * 0.5) * (2.0 * RAD) * (2.0 * RAD); + let gravity_torque = mass * 9.81 * LINK_LEN; + println!("mass ≈ {mass:.5} kg, gravity torque ≈ {gravity_torque:.5} N·m\n"); + + let cases = [ + ("free (fl = 0)", 0.0, false, None), + ("weak (fl = 0.25 · τ_g)", 0.25 * gravity_torque, false, None), + ("locked (fl = 4 · τ_g)", 4.0 * gravity_torque, false, None), + // A force-based `-fl·sign(q̇)` blows up here; a bounded constraint row + // simply never exceeds what it takes to stop the DoF. + ( + "extreme (fl = 1000 · τ_g)", + 1000.0 * gravity_torque, + false, + None, + ), + // Friction rows sharing the constraint bank with a limit and a motor. + ("motor + fl = 0.1 · τ_g", 0.1 * gravity_torque, true, None), + // Same locked case, but with a compliant joint softness: the friction + // row picks up CFM and the DoF creeps under load instead of sticking. + ( + "locked, soft joints (2 Hz)", + 4.0 * gravity_torque, + false, + Some(2.0), + ), + ]; + + let mut results = Vec::new(); + for (name, fl, motor, hz) in cases { + let (drop, qd) = run_case(&backend, fl, motor, hz).await?; + println!("{name:<26} drop = {drop:.6} m |q̇| = {qd:.6} rad/s"); + results.push((name, drop, qd)); + } + + let (_, free_drop, _) = results[0]; + let (_, weak_drop, _) = results[1]; + let (_, locked_drop, locked_qd) = results[2]; + let (_, extreme_drop, extreme_qd) = results[3]; + let (_, _, motor_qd) = results[4]; + let (_, soft_drop, _) = results[5]; + + println!(); + assert!( + free_drop > 0.1, + "frictionless pendulum should fall: drop = {free_drop}" + ); + assert!( + weak_drop < free_drop * 0.9, + "sub-gravity friction should slow the fall: {weak_drop} vs {free_drop}" + ); + assert!( + locked_drop.abs() < 1.0e-4 && locked_qd < 1.0e-4, + "friction above the gravity torque should hold the joint at rest: \ + drop = {locked_drop}, |q̇| = {locked_qd}" + ); + assert!( + extreme_drop.abs() < 1.0e-4 && extreme_qd < 1.0e-4, + "an oversized friction bound must stay inert, not chatter: \ + drop = {extreme_drop}, |q̇| = {extreme_qd}" + ); + assert!( + (motor_qd - 2.0).abs() < 0.1, + "the motor should still reach its 2 rad/s target through weak \ + friction: |q̇| = {motor_qd}" + ); + assert!( + soft_drop > 1.0e-3 && soft_drop < free_drop, + "a compliant joint softness should let the held DoF creep, without \ + letting it fall freely: drop = {soft_drop}" + ); + println!("OK"); + Ok(()) +} diff --git a/crates/examples3d/test_frictionloss_caps.rs b/crates/examples3d/test_frictionloss_caps.rs new file mode 100644 index 00000000..f45e390a --- /dev/null +++ b/crates/examples3d/test_frictionloss_caps.rs @@ -0,0 +1,137 @@ +//! Checks that the `BatchIndices` uniform agrees with the joint-constraint +//! buffer sizes after `set_dof_frictionloss` grows them. +//! +//! Reserving the dry-friction slots reallocates `joint_constraints` / +//! `joint_constraint_columns` and changes their per-batch capacities. Those +//! capacities are mirrored into the shared `BatchIndices` uniform, which the +//! kernels use to locate each multibody's slab. If the uniform is not +//! re-uploaded, every kernel indexes the resized buffers with stale strides. +//! +//! This asserts on the uniform directly rather than watching for NaN, which +//! only shows up once the resulting garbage happens to be large. +//! +//! Run with `cargo run --release --bin test_frictionloss_caps --features metal`. + +use khal::backend::{Backend, GpuBackend}; +use nexus3d::prelude::{NexusPipeline, NexusState, RbdCoupling}; +use nexus3d::rbd::shaders::utils::BatchIndices; +use rapier3d::prelude::*; + +const RAD: f32 = 0.05; +const LINK_LEN: f32 = 0.5; + +fn build_env(state: &mut NexusState, env: usize) { + for (c, num_links) in [2usize, 4, 3].iter().enumerate() { + let z = c as f32 * 4.0; + let root = RigidBodyBuilder::fixed() + .translation(Vec3::new(0.0, 0.0, z)) + .build(); + let rc = ColliderBuilder::cuboid(RAD, RAD, RAD) + .collision_groups(InteractionGroups::none()) + .build(); + let mut parent = state.insert_rigid_body_in(env, root, rc, RbdCoupling::None); + for i in 0..*num_links { + let body = RigidBodyBuilder::dynamic() + .translation(Vec3::new((i as f32 + 1.0) * LINK_LEN * 2.0, 0.0, z)) + .build(); + let collider = ColliderBuilder::cuboid(LINK_LEN, RAD, RAD) + .collision_groups(InteractionGroups::none()) + .build(); + let handle = state.insert_rigid_body_in(env, body, collider, RbdCoupling::None); + let anchor = if i == 0 { + Vec3::ZERO + } else { + Vec3::new(LINK_LEN, 0.0, 0.0) + }; + let joint = RevoluteJointBuilder::new(Vec3::Z) + .local_anchor1(anchor) + .local_anchor2(Vec3::new(-LINK_LEN, 0.0, 0.0)) + .limits([-1.5, 1.5]) + .motor_velocity(0.5, 0.1) + .motor_max_force(50.0); + state.insert_multibody_joint_in(env, parent, handle, joint.build()); + parent = handle; + } + } +} + +/// `(uniform capacity, actual capacity)` for the joint-constraint bank and its +/// column buffer, after `steps` pipeline steps. +async fn caps( + backend: &GpuBackend, + via_multibody_set: bool, + steps: usize, +) -> Result<((u32, u32), (u32, u32)), khal::backend::GpuBackendError> { + let mut state = NexusState::default(); + build_env(&mut state, 0); + for _ in 1..2 { + let env = state.add_environment(); + build_env(&mut state, env); + } + let mut pipeline = NexusPipeline::default(); + state.finalize(backend).await?; + + let rbd = state.rbd.as_mut().expect("no rbd state"); + let n = rbd.multibodies().dofs_per_batch() as usize * rbd.multibodies().num_batches() as usize; + let values = vec![0.1f32; n]; + if via_multibody_set { + rbd.multibodies_mut().set_dof_frictionloss(backend, &values); + } else { + rbd.set_dof_frictionloss(backend, &values); + } + + for _ in 0..steps { + pipeline.simulate(backend, &mut state, None).await?; + } + + let rbd = state.rbd.as_ref().expect("no rbd state"); + let bi: Vec = backend.slow_read_vec(rbd.batch_indices().buffer()).await?; + let bi = bi[0]; + Ok(( + ( + bi.mb_joint_constraints_batch_capacity, + rbd.multibodies().joint_constraints_per_batch(), + ), + ( + bi.mb_joint_constraint_columns_batch_capacity, + rbd.multibodies().joint_constraint_columns_per_batch(), + ), + )) +} + +fn main() -> anyhow::Result<()> { + pollster::block_on(run()) +} + +async fn run() -> anyhow::Result<()> { + let backend = + GpuBackend::Metal(khal::backend::Metal::new().map_err(|e| anyhow::anyhow!("{e:?}"))?); + + let mut bad = 0; + for via_set in [false, true] { + for steps in [0usize, 1] { + let (rows, cols) = caps(&backend, via_set, steps).await?; + let api = if via_set { "mb_set" } else { "rbd " }; + let ok = rows.0 == rows.1 && cols.0 == cols.1; + println!( + "{api}, after {steps} step(s): rows uniform/actual = {}/{}, cols = {}/{} {}", + rows.0, + rows.1, + cols.0, + cols.1, + if ok { "ok" } else { "MISMATCH" } + ); + // Before any step has run, the `mb_set` entry point is expected to + // carry a stale uniform: the step is what re-uploads it. + if !ok && steps > 0 { + bad += 1; + } + } + } + + if bad > 0 { + anyhow::bail!("{bad} case(s) still stale after a step"); + } + println!("\nOK"); + Ok(()) +} diff --git a/crates/examples3d/test_frictionloss_stress.rs b/crates/examples3d/test_frictionloss_stress.rs new file mode 100644 index 00000000..823b2ab5 --- /dev/null +++ b/crates/examples3d/test_frictionloss_stress.rs @@ -0,0 +1,208 @@ +//! Stress probe for the `frictionloss` constraint slot reservation. +//! +//! `reserve_frictionloss_slots` recomputes every multibody's `first_constraint` +//! offset and grows the joint-constraint bank. The single-pendulum probe in +//! `test_frictionloss.rs` never exercises that: it has one multibody in one +//! batch, so every offset is zero. This one builds several multibodies of +//! differing DoF counts across several batches, with limits and motors mixed +//! in, and checks that nothing goes non-finite. +//! +//! Run with `cargo run --release --bin test_frictionloss_stress --features metal`. + +use khal::backend::{Backend, GpuBackend}; +use nexus3d::prelude::{NexusPipeline, NexusState, RbdCoupling}; +use rapier3d::prelude::*; + +const RAD: f32 = 0.05; +const LINK_LEN: f32 = 0.5; +const STEPS: usize = 90; + +/// Chain lengths, in links. Differing DoF counts are the point: they make each +/// multibody's `first_constraint` offset distinct. +const CHAINS: [usize; 3] = [2, 4, 3]; + +/// Builds `CHAINS.len()` pendulum chains in environment `env`, offset along Z +/// so they don't overlap. Chain `c` gets limits and motors on its first joint +/// only, so limit/motor rows and friction rows share the bank unevenly. +fn build_env(state: &mut NexusState, env: usize, contacts: bool) { + if contacts { + let ground = RigidBodyBuilder::fixed() + .translation(Vec3::new(0.0, -3.0, 0.0)) + .build(); + let ground_collider = ColliderBuilder::cuboid(40.0, 0.5, 40.0).build(); + state.insert_rigid_body_in(env, ground, ground_collider, RbdCoupling::None); + } + let groups = if contacts { + InteractionGroups::all() + } else { + InteractionGroups::none() + }; + for (c, num_links) in CHAINS.iter().enumerate() { + let z = c as f32 * 4.0; + let root = RigidBodyBuilder::fixed() + .translation(Vec3::new(0.0, 0.0, z)) + .build(); + let root_collider = ColliderBuilder::cuboid(RAD, RAD, RAD) + .collision_groups(groups) + .build(); + let mut parent = state.insert_rigid_body_in(env, root, root_collider, RbdCoupling::None); + + for i in 0..*num_links { + let x = (i as f32 + 1.0) * LINK_LEN * 2.0; + let body = RigidBodyBuilder::dynamic() + .translation(Vec3::new(x, 0.0, z)) + .build(); + let collider = ColliderBuilder::cuboid(LINK_LEN, RAD, RAD) + .collision_groups(groups) + .build(); + let handle = state.insert_rigid_body_in(env, body, collider, RbdCoupling::None); + + let parent_anchor = if i == 0 { + Vec3::ZERO + } else { + Vec3::new(LINK_LEN, 0.0, 0.0) + }; + let mut builder = RevoluteJointBuilder::new(Vec3::Z) + .local_anchor1(parent_anchor) + .local_anchor2(Vec3::new(-LINK_LEN, 0.0, 0.0)); + if i == 0 { + builder = builder + .limits([-1.5, 1.5]) + .motor_velocity(0.5, 0.1) + .motor_max_force(50.0); + } + state.insert_multibody_joint_in(env, parent, handle, builder.build()); + parent = handle; + } + } +} + +#[allow(clippy::type_complexity)] +async fn run_case( + backend: &GpuBackend, + num_envs: usize, + frictionloss: f32, + // `false` turns off the implicit-Coriolis path, which is what makes the + // solver take the per-substep `gpu_mb_refresh_joint_constraints` branch + // instead of a full rebuild each substep. The friction rows are only + // touched by that kernel here. + implicit_coriolis: bool, + contacts: bool, + // `true` reproduces the zealot call pattern: the frictionloss is set + // through `GpuMultibodySet` directly rather than through `RbdState`, so + // nothing rebuilds `BatchIndices` at the call site. + via_multibody_set: bool, +) -> Result<(usize, f32), khal::backend::GpuBackendError> { + let mut state = NexusState::default(); + build_env(&mut state, 0, contacts); + for _ in 1..num_envs { + let env = state.add_environment(); + build_env(&mut state, env, contacts); + } + + let mut pipeline = NexusPipeline::default(); + state.finalize(backend).await?; + + let rbd = state.rbd.as_mut().expect("no rbd state"); + let per_batch = rbd.multibodies().dofs_per_batch() as usize; + let batches = rbd.multibodies().num_batches() as usize; + if frictionloss > 0.0 { + let values = vec![frictionloss; per_batch * batches]; + if via_multibody_set { + rbd.multibodies_mut().set_dof_frictionloss(backend, &values); + } else { + rbd.set_dof_frictionloss(backend, &values); + } + } + if !implicit_coriolis { + rbd.multibodies_mut().set_substep_refresh(false); + rbd.multibodies_mut().set_substep_refresh_light(false); + rbd.set_implicit_coriolis(backend, false); + } + + for _ in 0..STEPS { + pipeline.simulate(backend, &mut state, None).await?; + } + + let rbd = state.rbd.as_ref().expect("no rbd state"); + let poses: Vec = + backend.slow_read_vec(rbd.body_poses().buffer()).await?; + let dof_state: Vec = backend + .slow_read_vec(rbd.multibodies().dof_state().buffer()) + .await?; + + let bad_poses = poses + .iter() + .filter(|p| { + !p.translation.x.is_finite() + || !p.translation.y.is_finite() + || !p.translation.z.is_finite() + }) + .count(); + // Only the velocity section is integrated state; the rest are parameters. + let bad_vels = dof_state[..per_batch * batches] + .iter() + .filter(|v| !v.is_finite()) + .count(); + let max_speed = dof_state[..per_batch * batches] + .iter() + .filter(|v| v.is_finite()) + .fold(0.0f32, |a, v| a.max(v.abs())); + + Ok((bad_poses + bad_vels, max_speed)) +} + +fn main() -> anyhow::Result<()> { + pollster::block_on(run()) +} + +async fn run() -> anyhow::Result<()> { + let backend = + GpuBackend::Metal(khal::backend::Metal::new().map_err(|e| anyhow::anyhow!("{e:?}"))?); + + let mut failures = 0; + let mut known = 0; + // `implicit = false` switches the multibody to explicit Coriolis forces, + // which diverges on free-swinging chains of four or more links: the + // velocities grow smoothly and exponentially until they overflow. It does + // so with *or without* frictionloss, and identically on the pre-frictionloss + // tree, so it is reported but not counted. See `test_refresh_nan` for the + // isolated repro; it is the documented stability cost of that path, not a + // regression here. + for via_set in [false, true] { + for contacts in [false, true] { + for implicit in [true, false] { + for num_envs in [1usize, 2, 4] { + // The last two values are deliberately absurd: the impulse bound is + // `frictionloss · dt`, so a large enough loss can overflow the + // accumulated impulse. + for fl in [0.0f32, 0.05, 5.0, 1.0e12, 1.0e30] { + let (bad, max_speed) = + run_case(&backend, num_envs, fl, implicit, contacts, via_set).await?; + let c = if contacts { "contacts" } else { "free " }; + let path = if implicit { "rebuild" } else { "refresh" }; + let api = if via_set { "mb_set " } else { "rbd " }; + let tag = format!("{api}, {path}, {c}, envs = {num_envs}, fl = {fl:e}"); + if bad > 0 && !implicit { + println!( + "{tag:<43} KNOWN (explicit-Coriolis divergence): {bad} non-finite" + ); + known += 1; + } else if bad > 0 { + println!("{tag:<43} FAIL: {bad} non-finite values"); + failures += 1; + } else { + println!("{tag:<43} ok, max |q̇| = {max_speed:.4}"); + } + } + } + } + } + } + + if failures > 0 { + anyhow::bail!("{failures} configuration(s) produced non-finite state"); + } + println!("\nOK ({known} known explicit-Coriolis divergences ignored)"); + Ok(()) +} diff --git a/crates/examples3d/test_refresh_nan.rs b/crates/examples3d/test_refresh_nan.rs new file mode 100644 index 00000000..906bd672 --- /dev/null +++ b/crates/examples3d/test_refresh_nan.rs @@ -0,0 +1,193 @@ +//! Repro for the non-finite state seen with `implicit_coriolis = false`. +//! +//! This is *numerical divergence*, not corruption, and not a constraint bug: +//! it reproduces with no limits and no motors anywhere (`rows = none`), where +//! `gpu_mb_refresh_joint_constraints` is never even dispatched, and the +//! velocity trace grows smoothly and exponentially until it overflows +//! (`TRACE=1` prints it). It matches the tradeoff the solver already +//! documents: the explicit-Coriolis path is cheaper but less stable. +//! +//! Frictionloss is not involved: this reproduces at `frictionloss = 0`, and +//! identically on the tree from before joint friction became a constraint. +//! +//! The sweep isolates the variable: the same scenes are run with implicit +//! Coriolis on and off, reporting the first step at which any DoF velocity or +//! body pose stops being finite. Chains of four or more links diverge with it +//! off and are stable with it on. +//! +//! Run with `cargo run --release --bin test_refresh_nan --features metal`, +//! and `TRACE=1 ...` to see the per-step velocity growth. + +use khal::backend::{Backend, GpuBackend}; +use nexus3d::prelude::{NexusPipeline, NexusState, RbdCoupling}; +use rapier3d::prelude::*; + +const RAD: f32 = 0.05; +const LINK_LEN: f32 = 0.5; +const MAX_STEPS: usize = 100; + +#[derive(Clone, Copy)] +struct Case { + num_chains: usize, + num_links: usize, + limits: bool, + motor: bool, + /// Put the limit / motor rows on the chain's first joint only, leaving the + /// rest of its DoFs with no constraint slot at all. + first_joint_only: bool, +} + +fn build(state: &mut NexusState, env: usize, case: Case) { + for c in 0..case.num_chains { + let z = c as f32 * 4.0; + let root = RigidBodyBuilder::fixed() + .translation(Vec3::new(0.0, 0.0, z)) + .build(); + let rc = ColliderBuilder::cuboid(RAD, RAD, RAD) + .collision_groups(InteractionGroups::none()) + .build(); + let mut parent = state.insert_rigid_body_in(env, root, rc, RbdCoupling::None); + + // Chain `c` gets one extra link, so the multibodies have distinct DoF + // counts and therefore distinct constraint-slab offsets. + for i in 0..(case.num_links + c) { + let body = RigidBodyBuilder::dynamic() + .translation(Vec3::new((i as f32 + 1.0) * LINK_LEN * 2.0, 0.0, z)) + .build(); + let collider = ColliderBuilder::cuboid(LINK_LEN, RAD, RAD) + .collision_groups(InteractionGroups::none()) + .build(); + let handle = state.insert_rigid_body_in(env, body, collider, RbdCoupling::None); + + let anchor = if i == 0 { + Vec3::ZERO + } else { + Vec3::new(LINK_LEN, 0.0, 0.0) + }; + let mut j = RevoluteJointBuilder::new(Vec3::Z) + .local_anchor1(anchor) + .local_anchor2(Vec3::new(-LINK_LEN, 0.0, 0.0)); + let rows_here = !case.first_joint_only || i == 0; + if case.limits && rows_here { + j = j.limits([-1.5, 1.5]); + } + if case.motor && rows_here { + j = j.motor_velocity(0.5, 0.1).motor_max_force(50.0); + } + state.insert_multibody_joint_in(env, parent, handle, j.build()); + parent = handle; + } + } +} + +/// Runs `case` on the refresh path and returns the first step index at which +/// state goes non-finite, or `None` if it stays finite for `MAX_STEPS`. +async fn first_bad_step( + backend: &GpuBackend, + case: Case, + implicit_coriolis: bool, +) -> Result, khal::backend::GpuBackendError> { + let mut state = NexusState::default(); + build(&mut state, 0, case); + let mut pipeline = NexusPipeline::default(); + state.finalize(backend).await?; + + if !implicit_coriolis { + let rbd = state.rbd.as_mut().expect("no rbd state"); + rbd.multibodies_mut().set_substep_refresh(false); + rbd.multibodies_mut().set_substep_refresh_light(false); + rbd.set_implicit_coriolis(backend, false); + } + + for step in 0..MAX_STEPS { + pipeline.simulate(backend, &mut state, None).await?; + let rbd = state.rbd.as_ref().expect("no rbd state"); + let n = + rbd.multibodies().dofs_per_batch() as usize * rbd.multibodies().num_batches() as usize; + let dof: Vec = backend + .slow_read_vec(rbd.multibodies().dof_state().buffer()) + .await?; + let poses: Vec = + backend.slow_read_vec(rbd.body_poses().buffer()).await?; + let bad = dof[..n].iter().any(|v| !v.is_finite()) + || poses.iter().any(|p| { + !p.translation.x.is_finite() + || !p.translation.y.is_finite() + || !p.translation.z.is_finite() + }); + if std::env::var("TRACE").is_ok() { + let m = dof[..n] + .iter() + .filter(|v| v.is_finite()) + .fold(0.0f32, |a, v| a.max(v.abs())); + if step % 5 == 0 || bad { + println!(" step {step:>3}: max |q̇| = {m:e}"); + } + } + if bad { + return Ok(Some(step)); + } + } + Ok(None) +} + +fn main() -> anyhow::Result<()> { + pollster::block_on(run()) +} + +async fn run() -> anyhow::Result<()> { + let backend = + GpuBackend::Metal(khal::backend::Metal::new().map_err(|e| anyhow::anyhow!("{e:?}"))?); + + let mut any = false; + for implicit in [false, true] { + for first_joint_only in [true] { + for num_chains in [1usize] { + for num_links in [2usize, 4, 6] { + for (limits, motor) in + [(false, false), (true, false), (false, true), (true, true)] + { + let case = Case { + num_chains, + num_links, + limits, + motor, + first_joint_only, + }; + let rows = match (limits, motor) { + (false, false) => "none ", + (true, false) => "limit ", + (false, true) => "motor ", + (true, true) => "limit+motor ", + }; + let ic = if implicit { + "coriolis=implicit" + } else { + "coriolis=explicit" + }; + match first_bad_step(&backend, case, implicit).await? { + Some(step) => { + println!( + "{ic}, links = {num_links}, rows = {rows} DIVERGED at step {step}" + ); + any = true; + } + None => println!( + "{ic}, links = {num_links}, rows = {rows} finite for {MAX_STEPS} steps" + ), + } + } + } + } + } + } + + if any { + println!( + "\nReproduced (explicit-Coriolis divergence; run with TRACE=1 to see the growth)." + ); + } else { + println!("\nNo failure in this sweep."); + } + Ok(()) +} diff --git a/src_rbd/dynamics/multibody/multibody_from_rapier.rs b/src_rbd/dynamics/multibody/multibody_from_rapier.rs index c18bc0d8..4bd4d143 100644 --- a/src_rbd/dynamics/multibody/multibody_from_rapier.rs +++ b/src_rbd/dynamics/multibody/multibody_from_rapier.rs @@ -500,6 +500,8 @@ impl GpuMultibodySet { implicit_coriolis: true, coriolis_in_uniform: true, has_joint_constraints: all_infos.iter().any(|info| info.max_constraints > 0), + frictionloss_slots_reserved: false, + constraint_caps_dirty: false, multibody_info: Tensor::vector(backend, &all_infos, storage).unwrap(), max_contact_constraints: Tensor::scalar( @@ -528,7 +530,7 @@ impl GpuMultibodySet { // address section `s` at intra-batch offset // `s · dof_batch_capacity`. Frictionloss has no rapier // counterpart, so it starts at zero (off) and is filled in by - // `GpuMultibodySet::set_dof_frictionloss`. + // `RbdState::set_dof_frictionloss`. let n = (dofs_cap * num_batches) as usize; let mut buf = Vec::with_capacity(7 * n); buf.extend_from_slice(&all_dof_vels); @@ -616,8 +618,7 @@ impl GpuMultibodySet { backend, vec![ 0.0f32; - (mb_cap * num_batches * crate::shaders::dynamics::MAX_CONTACT_SENSORS) - as usize + (mb_cap * num_batches * crate::shaders::dynamics::MAX_CONTACT_SENSORS) as usize ], storage | BufferUsages::COPY_SRC, ) diff --git a/src_rbd/dynamics/multibody/multibody_set.rs b/src_rbd/dynamics/multibody/multibody_set.rs index e22c0425..d61c6f22 100644 --- a/src_rbd/dynamics/multibody/multibody_set.rs +++ b/src_rbd/dynamics/multibody/multibody_set.rs @@ -8,8 +8,8 @@ use crate::shaders::dynamics::{ MultibodyLinkStatic, MultibodyLinkWorkspace, RbdSimParams, }; use crate::shaders::utils::BatchIndices; -use khal::Shader; use khal::BufferUsages; +use khal::Shader; use khal::backend::{Backend, GpuBackend, GpuBackendError}; use rapier3d::prelude::JointAxis; use vortx::tensor::Tensor; @@ -99,6 +99,14 @@ pub struct GpuMultibodySet { /// 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, @@ -805,6 +813,11 @@ impl GpuMultibodySet { 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) and is transposed here into the @@ -817,14 +830,88 @@ impl GpuMultibodySet { self.write_dof_section(backend, 2, values, "armature"); } - /// Overwrites the per-DoF Coulomb joint friction section of - /// [`Self::dof_state`] (MJCF `frictionloss`, N·m), applied by the gravity - /// kernels as `-fl·sign(q̇)`. Zero, the default, disables it. `values` uses - /// the same env-major layout as [`Self::set_dof_armature`]. + /// Overwrites the per-DoF dry joint friction section of + /// [`Self::dof_state`] (MJCF `frictionloss`, N·m). Zero, the default, + /// disables it. `values` uses the same env-major layout as + /// [`Self::set_dof_armature`]. + /// + /// Friction loss is a constraint, not a force: each DoF with a non-zero + /// loss gets a solver row driving its velocity to zero, with the impulse + /// bounded by `frictionloss · dt` (a load-independent bound, which is why + /// MuJoCo distinguishes it from Coulomb friction). The first non-zero call + /// reserves the extra constraint slots, which changes per-batch capacities + /// and so invalidates the shared `BatchIndices` uniform; the next + /// `RbdPipeline` step re-uploads it. + /// [`RbdState::set_dof_frictionloss`](crate::pipeline::RbdState::set_dof_frictionloss) + /// does it up front instead, if you would rather not carry a dirty uniform. 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`]. diff --git a/src_rbd/dynamics/multibody/multibody_solver.rs b/src_rbd/dynamics/multibody/multibody_solver.rs index dbf5eeb8..894295af 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, GpuMbDelayTick, GpuMbRefreshJointConstraints, GpuMbSeedContactRestitution, - GpuMbSenseContactImpulses, GpuMbSnapshotContactWarmstart, GpuMbSolveConstraints, GpuMbSolveContactsDelassus, - GpuMbSolveImpulseJointConstraints, GpuMbSolveJoints, GpuMbStashContactsLen, - GpuMbTransferContactWarmstart, GpuMbUpdateImpulseJointConstraints, - GpuMbWarmstartContactConstraints, Velocity, WorldMassProperties, + GpuMbComputeSolveBounds, GpuMbDelayTick, GpuMbFinalizeContactConstraints, + GpuMbFinalizeImpulseJointConstraints, 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; @@ -237,11 +239,7 @@ impl GpuMultibodySolver { // 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 - { + if mb.implicit_coriolis || mb.substep_refresh || mb.substep_refresh_light || first_substep { self.build_contact_constraints( encoder, timestamps.as_deref_mut(), @@ -250,8 +248,7 @@ impl GpuMultibodySolver { first_substep, )?; } else if mb.has_joint_constraints { - let mut pass = - encoder.begin_pass("[RBD] mbb/refresh-joint", timestamps.as_deref_mut()); + 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], @@ -363,6 +360,7 @@ impl GpuMultibodySolver { &mut mb.joint_constraint_columns, &mb.dof_couplings, &mut mb.motor_delay_state, + &mb.dof_state, &mb.constraint_softness, args.batch_indices, )?; diff --git a/src_rbd/pipeline/rbd_state.rs b/src_rbd/pipeline/rbd_state.rs index c5b240ab..f4e76d02 100644 --- a/src_rbd/pipeline/rbd_state.rs +++ b/src_rbd/pipeline/rbd_state.rs @@ -439,6 +439,19 @@ impl RbdState { self.rebuild_batch_indices(backend); } + /// Sets the per-DoF dry joint friction (MJCF `frictionloss`, N·m), in the + /// env-major `dofs_per_batch * num_batches` layout described on + /// [`GpuMultibodySet::set_dof_frictionloss`](crate::dynamics::GpuMultibodySet::set_dof_frictionloss). + /// + /// The first non-zero call grows the joint-constraint bank, so the shared + /// `BatchIndices` uniform is rebuilt here. + #[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. diff --git a/src_rbd/pipeline/rbd_step.rs b/src_rbd/pipeline/rbd_step.rs index 81432b7d..cfe9ce49 100644 --- a/src_rbd/pipeline/rbd_step.rs +++ b/src_rbd/pipeline/rbd_step.rs @@ -119,6 +119,17 @@ impl RbdPipeline { }; 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. { diff --git a/src_rbd_shaders/dynamics/multibody/gravity_and_lu.rs b/src_rbd_shaders/dynamics/multibody/gravity_and_lu.rs index 9ebb63f5..fc005f5f 100644 --- a/src_rbd_shaders/dynamics/multibody/gravity_and_lu.rs +++ b/src_rbd_shaders/dynamics/multibody/gravity_and_lu.rs @@ -129,10 +129,6 @@ pub fn gpu_mb_gravity_and_lu( let damping_slice = batch_ids .ib(batch_id, dof_state) .offset(batch_ids.dof_batch_capacity as usize + gen_base); - // Coulomb joint friction (MJCF `frictionloss`) sits in the 7th section. - let frictionloss_slice = batch_ids - .ib(batch_id, dof_state) - .offset(6 * batch_ids.dof_batch_capacity as usize + gen_base); let stiffness_slice = batch_ids .ib(batch_id, dof_state) .offset(3 * batch_ids.dof_batch_capacity as usize + gen_base); @@ -288,17 +284,7 @@ pub fn gpu_mb_gravity_and_lu( let idx = batch_ids.mbi(batch_id, gen_base + i as usize); let cur = gen_forces.read(idx); let v = vel_slice[i as usize]; - // Coulomb joint friction: -fl·sign(v), MuJoCo `frictionloss` semantics. - // Zero (the default) leaves the DoF untouched. - let fl = frictionloss_slice[i as usize]; - let fric = if v > 0.0 { - -fl - } else if v < 0.0 { - fl - } else { - 0.0 - }; - gen_forces.write(idx, cur - damping_slice[i as usize] * v + fric); + gen_forces.write(idx, cur - damping_slice[i as usize] * v); } workgroup_memory_barrier_with_group_sync(); @@ -465,10 +451,6 @@ fn gravity_and_lu_packed_impl 0.0 { - -fl - } else if v < 0.0 { - fl - } else { - 0.0 - }; - gen_forces.write(idx, cur - damping_slice[i as usize] * v + fric); + gen_forces.write(idx, cur - damping_slice[i as usize] * v); } workgroup_memory_barrier_with_group_sync(); @@ -800,10 +772,6 @@ pub fn gpu_mb_gravity_and_lu_t1( let damping_slice = batch_ids .ib(batch_id, dof_state) .offset(batch_ids.dof_batch_capacity as usize + gen_base); - // Coulomb joint friction (MJCF `frictionloss`) sits in the 7th section. - let frictionloss_slice = batch_ids - .ib(batch_id, dof_state) - .offset(6 * batch_ids.dof_batch_capacity as usize + gen_base); let stiffness_slice = batch_ids .ib(batch_id, dof_state) .offset(3 * batch_ids.dof_batch_capacity as usize + gen_base); @@ -931,17 +899,7 @@ pub fn gpu_mb_gravity_and_lu_t1( let idx = batch_ids.mbi(batch_id, gen_base + i as usize); let cur = gen_forces.read(idx); let v = vel_slice[i as usize]; - // Coulomb joint friction: -fl·sign(v), MuJoCo `frictionloss` semantics. - // Zero (the default) leaves the DoF untouched. - let fl = frictionloss_slice[i as usize]; - let fric = if v > 0.0 { - -fl - } else if v < 0.0 { - fl - } else { - 0.0 - }; - gen_forces.write(idx, cur - damping_slice[i as usize] * v + fric); + gen_forces.write(idx, cur - damping_slice[i as usize] * v); } // Per-DoF joint springs. diff --git a/src_rbd_shaders/dynamics/multibody/joint_constraints.rs b/src_rbd_shaders/dynamics/multibody/joint_constraints.rs index 35bf823f..45611b67 100644 --- a/src_rbd_shaders/dynamics/multibody/joint_constraints.rs +++ b/src_rbd_shaders/dynamics/multibody/joint_constraints.rs @@ -18,8 +18,9 @@ use crate::utils::linalg::{MatSlice, VSlice, lu_solve_in_place}; use crate::{DIM, MAX_FLT}; use super::types::{ - MB_JOINT_KIND_COUPLING, MB_JOINT_KIND_LIMIT, MB_JOINT_KIND_LIMIT_INACTIVE, MB_JOINT_KIND_MOTOR, - MbDofCoupling, MultibodyInfo, MultibodyJointConstraint, MultibodyLinkStatic, + MB_JOINT_KIND_COUPLING, MB_JOINT_KIND_FRICTION, MB_JOINT_KIND_LIMIT, + MB_JOINT_KIND_LIMIT_INACTIVE, MB_JOINT_KIND_MOTOR, MbDofCoupling, MultibodyInfo, + MultibodyJointConstraint, MultibodyLinkStatic, }; use super::ws_soa::{WsAddr, ws_coord}; @@ -109,6 +110,7 @@ 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, @@ -274,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 @@ -406,6 +431,49 @@ 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)`), the same compliance the limit rows use; +/// it is MuJoCo's `solreffriction` knob. Only CFM applies here, never ERP: +/// there is no position error for a bias to chase. With the default +/// near-rigid joint softness the coefficient is ~0 and a DoF whose driving +/// force stays under the bound sticks exactly; softening it lets the DoF +/// creep under load instead. +#[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)] @@ -496,10 +564,12 @@ pub fn gpu_mb_init_joint_constraints( #[spirv(storage_buffer, descriptor_set = 0, binding = 7)] 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 = 8)] - motor_delay_state: &[f32], - #[spirv(uniform, descriptor_set = 0, binding = 9)] softness: &ConstraintSoftness, - #[spirv(uniform, descriptor_set = 0, binding = 10)] batch_ids: &BatchIndices, + #[spirv(storage_buffer, descriptor_set = 0, binding = 8)] 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 = 9)] dof_state: &[f32], + #[spirv(uniform, descriptor_set = 0, binding = 10)] softness: &ConstraintSoftness, + #[spirv(uniform, descriptor_set = 0, binding = 11)] batch_ids: &BatchIndices, ) { const LANES: u32 = 64; @@ -559,6 +629,7 @@ pub fn gpu_mb_init_joint_constraints( links_workspace, dof_couplings, joint_constraints, + dof_state, &mb, cons_base, batch_id, @@ -631,8 +702,7 @@ pub fn gpu_mb_refresh_joint_constraints( #[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(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, ) { @@ -659,6 +729,15 @@ pub fn gpu_mb_refresh_joint_constraints( 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 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..8bf23fb6 100644 --- a/src_rbd_shaders/dynamics/multibody/types.rs +++ b/src_rbd_shaders/dynamics/multibody/types.rs @@ -62,6 +62,14 @@ 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 (MJCF `frictionloss`). A row +/// with jacobian `J = e_{dof_id}`, zero target velocity and no position +/// residual, whose impulse is clamped to `±frictionloss·dt` and which carries +/// the shared joint CFM softness. MuJoCo models friction loss this way rather +/// than as a `-f·sign(q̇)` force: the bound is load-independent (not Coulomb +/// friction), and only a constraint can hold a DoF at rest instead of +/// chattering around zero velocity. +pub const MB_JOINT_KIND_FRICTION: u32 = 5; /// Sentinel marking a link with no parent (the root). pub const MULTIBODY_ROOT: u32 = u32::MAX; From d17da57aa9347b2a8e06add97ada171b3767e634 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Sun, 23 Aug 2026 12:51:47 +0200 Subject: [PATCH 22/41] feat(rbd): seed per-DoF joint friction from rapier's Multibody::frictions --- Cargo.toml | 14 +- crates/examples3d/Cargo.toml | 5 + crates/examples3d/test_frictionloss_mjcf.rs | 134 ++++++++++++++++++ .../multibody/multibody_from_rapier.rs | 36 ++++- 4 files changed, 179 insertions(+), 10 deletions(-) create mode 100644 crates/examples3d/test_frictionloss_mjcf.rs diff --git a/Cargo.toml b/Cargo.toml index d2ced59a..ebf90a18 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -95,11 +95,15 @@ rust.unexpected_cfgs = { level = "warn", check-cfg = [ [patch.crates-io] # Compare against the rapier checkout the reference example runs, not the # crates.io release (their solver defaults differ). -#rapier2d = { path = "../rapier/crates/rapier2d" } -#rapier3d = { path = "../rapier/crates/rapier3d" } -#rapier3d-mjcf = { path = "../rapier/crates/rapier3d-mjcf" } -#rapier3d-urdf = { path = "../rapier/crates/rapier3d-urdf" } -#rapier3d-meshloader = { path = "../rapier/crates/rapier3d-meshloader" } +# +# Also carries `Multibody::frictions` (MJCF ``), which the +# multibody build path reads into the GPU per-DoF friction section. All five +# must be patched together so the rapier types unify. +rapier2d = { path = "../rapier/crates/rapier2d" } +rapier3d = { path = "../rapier/crates/rapier3d" } +rapier3d-mjcf = { path = "../rapier/crates/rapier3d-mjcf" } +rapier3d-urdf = { path = "../rapier/crates/rapier3d-urdf" } +rapier3d-meshloader = { path = "../rapier/crates/rapier3d-meshloader" } ## Local glam clone with SPIR-V vector-arithmetic intrinsics (Vec3 add/sub/mul/scale). #glam = { path = "../glam-rs" } # 30% faster for loop in P2G diff --git a/crates/examples3d/Cargo.toml b/crates/examples3d/Cargo.toml index 0ce62bc3..daacc1c7 100644 --- a/crates/examples3d/Cargo.toml +++ b/crates/examples3d/Cargo.toml @@ -74,3 +74,8 @@ required-features = ["metal"] name = "test_refresh_nan" path = "test_refresh_nan.rs" required-features = ["metal"] + +[[bin]] +name = "test_frictionloss_mjcf" +path = "test_frictionloss_mjcf.rs" +required-features = ["metal"] diff --git a/crates/examples3d/test_frictionloss_mjcf.rs b/crates/examples3d/test_frictionloss_mjcf.rs new file mode 100644 index 00000000..ba389d31 --- /dev/null +++ b/crates/examples3d/test_frictionloss_mjcf.rs @@ -0,0 +1,134 @@ +//! End-to-end check that MJCF `` reaches the GPU solver. +//! +//! The path is: `` → rapier's `Multibody::frictions` (via +//! `rapier3d-mjcf`'s `add_frictionloss_to_multibody`) → the per-DoF friction +//! section of `dof_state` at build time → one `MB_JOINT_KIND_FRICTION` row per +//! non-zero DoF in `gpu_mb_init_joint_constraints`. +//! +//! Nothing calls `set_dof_frictionloss` here: the whole point is that loading +//! a model is enough. +//! +//! Run with `cargo run --release --bin test_frictionloss_mjcf --features metal`. + +use khal::backend::{Backend, GpuBackend}; +use nexus3d::prelude::{NexusPipeline, NexusState, RbdCoupling}; +use rapier3d::prelude::*; +use rapier3d_mjcf::{MjcfLoaderOptions, MjcfMultibodyOptions, MjcfRobot}; + +const STEPS: usize = 15; + +/// A hinge at the origin with a 1 m arm along +X, so gravity torques it. +/// `frictionloss` is substituted in. +fn model(frictionloss: f32) -> String { + format!( + r#" + + +"# + ) +} + +/// Loads the model into a `NexusState`, steps it, and returns how far the arm +/// fell plus the joint speed at the end. +async fn run_case( + backend: &GpuBackend, + frictionloss: f32, +) -> Result<(f32, f32), khal::backend::GpuBackendError> { + let xml = model(frictionloss); + let (robot, _) = MjcfRobot::from_str(&xml, MjcfLoaderOptions::default(), ".").unwrap(); + + let mut state = NexusState::default(); + { + let world = state.rbd_world_mut(0); + robot.insert_using_multibody_joints( + &mut world.bodies, + &mut world.colliders, + &mut world.multibody_joints, + &mut world.impulse_joints, + MjcfMultibodyOptions::empty(), + ); + } + + let mut pipeline = NexusPipeline::default(); + state.finalize(backend).await?; + + // Confirm the value survived the loader before trusting the simulation. + let loaded = state + .rbd_world(0) + .multibody_joints + .multibodies() + .flat_map(|mb| mb.frictions().iter().copied().collect::>()) + .fold(0.0f32, f32::max); + assert!( + (loaded - frictionloss).abs() < 1.0e-6, + "rapier's Multibody::frictions should carry the MJCF value: \ + got {loaded}, expected {frictionloss}" + ); + + for _ in 0..STEPS { + pipeline.simulate(backend, &mut state, None).await?; + } + + let rbd = state.rbd.as_ref().expect("no rbd state"); + let poses: Vec = + backend.slow_read_vec(rbd.body_poses().buffer()).await?; + let dof: Vec = backend + .slow_read_vec(rbd.multibodies().dof_state().buffer()) + .await?; + let drop = poses + .iter() + .map(|p| -p.translation.y) + .fold(0.0f32, f32::max); + Ok((drop, dof[0].abs())) +} + +fn main() -> anyhow::Result<()> { + pollster::block_on(run()) +} + +async fn run() -> anyhow::Result<()> { + let backend = + GpuBackend::Metal(khal::backend::Metal::new().map_err(|e| anyhow::anyhow!("{e:?}"))?); + + // Gravity torque about the hinge for a 1 kg arm with its centre 1 m out. + let gravity_torque = 1.0 * 9.81 * 1.0; + let cases = [ + ("frictionloss = 0", 0.0), + ("frictionloss = 0.25 · τ_g", 0.25 * gravity_torque), + ("frictionloss = 4 · τ_g", 4.0 * gravity_torque), + ]; + + let mut results = Vec::new(); + for (name, fl) in cases { + let (drop, speed) = run_case(&backend, fl).await?; + println!("{name:<28} drop = {drop:.6} m |q̇| = {speed:.6} rad/s"); + results.push((drop, speed)); + } + + let (free, _) = results[0]; + let (weak, _) = results[1]; + let (locked, locked_speed) = results[2]; + + println!(); + assert!(free > 0.02, "the arm should fall with no friction: {free}"); + assert!( + weak < free * 0.95, + "sub-gravity friction should slow the fall: {weak} vs {free}" + ); + assert!( + locked.abs() < 1.0e-4 && locked_speed < 1.0e-4, + "friction above the gravity torque should hold the joint: \ + drop = {locked}, |q̇| = {locked_speed}" + ); + println!("OK"); + Ok(()) +} diff --git a/src_rbd/dynamics/multibody/multibody_from_rapier.rs b/src_rbd/dynamics/multibody/multibody_from_rapier.rs index 4bd4d143..c43c027d 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,16 @@ impl GpuMultibodySet { let mut global_max_cons = 0u32; let mut global_max_couplings = 0u32; + // Whether any multibody anywhere declares dry joint friction. The + // constraint-slot reservation below is all-or-nothing, matching + // `GpuMultibodySet::reserve_frictionloss_slots`: reserving for only the + // multibodies that currently have friction would leave a later + // `set_dof_frictionloss` on a different one with nowhere to emit. + 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 +82,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 +163,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 +200,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 (MJCF ``). + 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 +309,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 +363,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 +422,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 +462,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 +501,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); @@ -500,7 +524,7 @@ impl GpuMultibodySet { implicit_coriolis: true, coriolis_in_uniform: true, has_joint_constraints: all_infos.iter().any(|info| info.max_constraints > 0), - frictionloss_slots_reserved: false, + frictionloss_slots_reserved: scene_has_joint_friction, constraint_caps_dirty: false, multibody_info: Tensor::vector(backend, &all_infos, storage).unwrap(), @@ -528,9 +552,10 @@ impl GpuMultibodySet { // 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`. Frictionloss has no rapier - // counterpart, so it starts at zero (off) and is filled in by - // `RbdState::set_dof_frictionloss`. + // `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(7 * n); buf.extend_from_slice(&all_dof_vels); @@ -539,6 +564,7 @@ impl GpuMultibodySet { buf.extend_from_slice(&all_dof_stiffness); buf.extend_from_slice(&all_dof_spring_ref); buf.extend_from_slice(&all_dof_kinematic); + 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() From 54d3f2385a63496eb216c0223c8cdd795a284384 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Sun, 23 Aug 2026 15:37:22 +0200 Subject: [PATCH 23/41] refactor(rbd): read the contact prediction distance from RbdSimParams instead of a dedicated uniform --- src_rbd/broad_phase/lbvh.rs | 4 +- src_rbd/broad_phase/narrow_phase.rs | 10 ++-- src_rbd/pipeline/insertion_removal.rs | 14 +++-- src_rbd/pipeline/rbd_state.rs | 3 -- src_rbd/pipeline/rbd_state_from_rapier.rs | 17 +++--- src_rbd/pipeline/rbd_step.rs | 4 +- src_rbd_shaders/broad_phase/brute_force.rs | 6 +-- src_rbd_shaders/broad_phase/narrow_phase.rs | 57 ++++++++++++--------- 8 files changed, 57 insertions(+), 58 deletions(-) diff --git a/src_rbd/broad_phase/lbvh.rs b/src_rbd/broad_phase/lbvh.rs index e97e3e96..a081fab2 100644 --- a/src_rbd/broad_phase/lbvh.rs +++ b/src_rbd/broad_phase/lbvh.rs @@ -336,7 +336,7 @@ impl Lbvh { collision_pairs_indirect: &mut Tensor<[u32; 3]>, collision_groups: &Tensor, pair_filter: &Tensor<[u32; 2]>, - prediction: &Tensor, + sim_params: &Tensor, ) -> Result<(), GpuBackendError> { state.resize_bf_buffers(backend, colliders_len); @@ -361,7 +361,7 @@ impl Lbvh { collision_groups, batch_indices, pair_filter, - prediction, + 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 b4ca4933..e7c7b0b5 100644 --- a/src_rbd/broad_phase/narrow_phase.rs +++ b/src_rbd/broad_phase/narrow_phase.rs @@ -57,7 +57,7 @@ impl GpuNarrowPhase { batch_indices: &Tensor, collider_parent: &Tensor, collider_materials: &Tensor, - prediction: &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, @@ -93,7 +93,7 @@ impl GpuNarrowPhase { batch_indices, collider_parent, collider_materials, - prediction, + sim_params, )?; // Pass 2: defer the complex shape pairs into `pfm_pairs` (kept as a @@ -108,7 +108,7 @@ impl GpuNarrowPhase { pfm_pairs, pfm_pairs_len, batch_indices, - prediction, + sim_params, vertices, indices, )?; @@ -133,7 +133,7 @@ impl GpuNarrowPhase { indices, collider_parent, collider_materials, - prediction, + sim_params, )?; // Reduction rewrites `contacts_len`, so it has to run before the // indirect args are derived from it. @@ -145,7 +145,7 @@ impl GpuNarrowPhase { contacts, contacts_len, batch_indices, - prediction, + sim_params, merge_cos, )?; } diff --git a/src_rbd/pipeline/insertion_removal.rs b/src_rbd/pipeline/insertion_removal.rs index cb8f533d..804e310b 100644 --- a/src_rbd/pipeline/insertion_removal.rs +++ b/src_rbd/pipeline/insertion_removal.rs @@ -149,12 +149,6 @@ impl RbdState { BufferUsages::STORAGE | BufferUsages::UNIFORM | BufferUsages::COPY_DST, ) .unwrap(); - let prediction = Tensor::scalar( - backend, - all_sim_params[0].prediction_distance(), - BufferUsages::UNIFORM | BufferUsages::COPY_DST, - ) - .unwrap(); // Two-element readback: the (max) collision-pair count and the uncolored count. let resize_readback = GpuReadback::new(backend, 2).unwrap(); let collision_pairs_indirect = @@ -259,7 +253,12 @@ 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::vector( + backend, + &all_sim_params, + BufferUsages::STORAGE | BufferUsages::UNIFORM, + ) + .unwrap(), vels: Tensor::vector(backend, &all_vels, rw).unwrap(), #[cfg(feature = "dim3")] reset_templates_bodies: None, @@ -287,7 +286,6 @@ impl RbdState { collision_pairs_len, collision_pairs_len_max, num_batches_uniform, - prediction, contact_merge_cos, resize_readback, collision_pairs_indirect, diff --git a/src_rbd/pipeline/rbd_state.rs b/src_rbd/pipeline/rbd_state.rs index f4e76d02..6f33232d 100644 --- a/src_rbd/pipeline/rbd_state.rs +++ b/src_rbd/pipeline/rbd_state.rs @@ -164,9 +164,6 @@ 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, - /// Contact prediction distance (`RbdSimParams::prediction_distance`), - /// consumed by the narrow-phase kernels. - pub(super) prediction: Tensor, /// Cosine of the maximum angle between two contact normals for their /// manifolds to be clustered together (see `gpu_reduce_contacts`). pub(super) contact_merge_cos: Tensor, diff --git a/src_rbd/pipeline/rbd_state_from_rapier.rs b/src_rbd/pipeline/rbd_state_from_rapier.rs index a507383d..7e458c91 100644 --- a/src_rbd/pipeline/rbd_state_from_rapier.rs +++ b/src_rbd/pipeline/rbd_state_from_rapier.rs @@ -622,12 +622,6 @@ impl RbdState { BufferUsages::STORAGE | BufferUsages::UNIFORM | BufferUsages::COPY_DST, ) .unwrap(); - let prediction = Tensor::scalar( - backend, - all_sim_params[0].prediction_distance(), - BufferUsages::UNIFORM | BufferUsages::COPY_DST, - ) - .unwrap(); // Two-element readback: the (max) collision-pair count and the uncolored count. let resize_readback = GpuReadback::new(backend, 2).unwrap(); let collision_pairs_indirect = @@ -788,9 +782,13 @@ 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 | BufferUsages::COPY_DST) - .unwrap(), + sim_params: Tensor::vector( + backend, + &all_sim_params, + BufferUsages::STORAGE | BufferUsages::UNIFORM, + ) + .unwrap(), + 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(), @@ -838,7 +836,6 @@ impl RbdState { collision_pairs_len, collision_pairs_len_max, num_batches_uniform, - prediction, contact_merge_cos, resize_readback, collision_pairs_indirect, diff --git a/src_rbd/pipeline/rbd_step.rs b/src_rbd/pipeline/rbd_step.rs index cfe9ce49..c834c2e9 100644 --- a/src_rbd/pipeline/rbd_step.rs +++ b/src_rbd/pipeline/rbd_step.rs @@ -217,7 +217,7 @@ impl RbdPipeline { &mut state.collision_pairs_indirect, &state.collision_groups, &state.pair_filter, - &state.prediction, + &state.sim_params, )?; drop(pass); split(&mut *encoder)?; @@ -324,7 +324,7 @@ impl RbdPipeline { &state.batch_indices, &state.collider_parent, &state.collider_materials, - &state.prediction, + &state.sim_params, self.contact_reduction, &state.contact_merge_cos, &mut state.pairs_flat_offsets, diff --git a/src_rbd_shaders/broad_phase/brute_force.rs b/src_rbd_shaders/broad_phase/brute_force.rs index b16b2522..5b095e9f 100644 --- a/src_rbd_shaders/broad_phase/brute_force.rs +++ b/src_rbd_shaders/broad_phase/brute_force.rs @@ -10,6 +10,7 @@ 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}; @@ -57,8 +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]], - // Contact prediction distance (`RbdSimParams::prediction_distance`). - #[spirv(uniform, descriptor_set = 0, binding = 6)] prediction: &f32, + #[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 28086891..8e8c7cfe 100644 --- a/src_rbd_shaders/broad_phase/narrow_phase.rs +++ b/src_rbd_shaders/broad_phase/narrow_phase.rs @@ -2,6 +2,7 @@ //! //! 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, @@ -85,12 +86,7 @@ pub const COS_MERGE_ANGLE: f32 = 0.996; /// 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, -) { +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; @@ -136,13 +132,12 @@ pub fn gpu_reduce_contacts( #[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, - // Contact prediction distance: `manifold_reduction` only keeps candidates - // within it, exactly as the narrow-phase passes that produced them. - #[spirv(uniform, descriptor_set = 0, binding = 3)] prediction: &f32, + #[spirv(uniform, descriptor_set = 0, binding = 3)] params: &RbdSimParams, // Cosine of the maximum angle between two manifolds' normals for them to // share a cluster. See [`COS_MERGE_ANGLE`]. #[spirv(uniform, descriptor_set = 0, binding = 4)] merge_cos: &f32, ) { + let prediction = params.prediction_distance(); 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); @@ -162,15 +157,25 @@ pub fn gpu_reduce_contacts( // 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 = 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); + 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); + 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 @@ -196,7 +201,7 @@ pub fn gpu_reduce_contacts( } else { out.contact.normal_a }; - let mut reduced = manifold_reduction(&cand, num as u32, normal, *prediction); + 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; @@ -303,9 +308,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], - // Contact prediction distance (`RbdSimParams::prediction_distance`). - #[spirv(uniform, descriptor_set = 0, binding = 9)] prediction: &f32, + #[spirv(uniform, descriptor_set = 0, binding = 9)] params: &RbdSimParams, ) { + let prediction = params.prediction_distance(); let num_threads = num_workgroups.x * WORKGROUP_SIZE; let contacts_batch_capacity = batch_ids.contacts_batch_capacity as usize; @@ -379,12 +384,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 @@ -429,8 +434,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)] prediction: &f32, + #[spirv(uniform, descriptor_set = 0, binding = 7)] params: &RbdSimParams, ) { + let prediction = params.prediction_distance(); let num_threads = num_workgroups.x * WORKGROUP_SIZE; // Every batch is allocated the same capacity, so this is batch-independent. let contacts_batch_capacity = batch_ids.contacts_batch_capacity as usize; @@ -523,7 +529,7 @@ pub fn gpu_narrow_phase_shape_shape_deferred( let mesh = shape1.to_trimesh(); let convex = shape2; trimesh_convex( - *prediction, + prediction, pose12, &mesh, convex, @@ -542,7 +548,7 @@ pub fn gpu_narrow_phase_shape_shape_deferred( let mesh = shape2.to_trimesh(); // NOTE: pair indices are flipped. trimesh_convex( - *prediction, + prediction, pose12.inverse(), &mesh, convex, @@ -562,7 +568,7 @@ pub fn gpu_narrow_phase_shape_shape_deferred( let pline = shape1.to_polyline(); let convex = shape2; polyline_convex( - *prediction, + prediction, pose12, &pline, convex, @@ -581,7 +587,7 @@ pub fn gpu_narrow_phase_shape_shape_deferred( let pline = shape2.to_polyline(); // NOTE: pair indices are flipped. polyline_convex( - *prediction, + prediction, pose12.inverse(), &pline, convex, @@ -791,8 +797,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)] prediction: &f32, + #[spirv(uniform, descriptor_set = 0, binding = 9)] params: &RbdSimParams, ) { + let prediction = params.prediction_distance(); let num_threads = num_workgroups.x * WORKGROUP_SIZE; let contacts_batch_capacity = batch_ids.contacts_batch_capacity as usize; @@ -825,13 +832,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. From 3c7dc2a0e060ba69fb3f4ad6aa8361b8e49cb4a6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Sun, 23 Aug 2026 19:12:58 +0200 Subject: [PATCH 24/41] chore: cargo fmt --- src_rbd/dynamics/multibody/env_reset.rs | 12 ++++++++---- src_rbd/pipeline/mod.rs | 2 +- src_rbd_shaders/dynamics/multibody/contact_sensor.rs | 2 +- src_rbd_shaders/dynamics/multibody/env_reset.rs | 10 ++++++++-- 4 files changed, 18 insertions(+), 8 deletions(-) diff --git a/src_rbd/dynamics/multibody/env_reset.rs b/src_rbd/dynamics/multibody/env_reset.rs index da4bf468..d9863de9 100644 --- a/src_rbd/dynamics/multibody/env_reset.rs +++ b/src_rbd/dynamics/multibody/env_reset.rs @@ -15,8 +15,8 @@ use super::multibody_set::GpuMultibodySet; use crate::math::Vector; use crate::shaders::dynamics::{ - GpuMbEnvReset, GpuMbEnvResetBatch, MULTIBODY_ROOT, MultibodyLinkStatic, - MultibodyLinkWorkspace, WS_QUADS, ws_soa_from_structs, ws_soa_to_structs, + GpuMbEnvReset, GpuMbEnvResetBatch, MULTIBODY_ROOT, MultibodyLinkStatic, MultibodyLinkWorkspace, + WS_QUADS, ws_soa_from_structs, ws_soa_to_structs, }; use glamx::{UVec4, Vec4}; use khal::BufferUsages; @@ -136,8 +136,12 @@ impl EnvResetBundle { storage, ) .unwrap(), - staging_dofs: Tensor::vector(backend, &vec![0.0f32; (2 * dpb).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(), } } diff --git a/src_rbd/pipeline/mod.rs b/src_rbd/pipeline/mod.rs index 7b24c0ee..65a40df4 100644 --- a/src_rbd/pipeline/mod.rs +++ b/src_rbd/pipeline/mod.rs @@ -10,7 +10,7 @@ mod rbd_state; mod rbd_state_from_rapier; mod rbd_step; -pub use rbd_state::{RbdCapacities, RbdResizePolicy, RbdState, RunStats}; #[cfg(feature = "dim3")] pub use rbd_state::RbdSnapshot; +pub use rbd_state::{RbdCapacities, RbdResizePolicy, RbdState, RunStats}; pub use rbd_step::{FORCE_FUSED_SWEEPS, RbdPipeline}; diff --git a/src_rbd_shaders/dynamics/multibody/contact_sensor.rs b/src_rbd_shaders/dynamics/multibody/contact_sensor.rs index ee24be98..03adc551 100644 --- a/src_rbd_shaders/dynamics/multibody/contact_sensor.rs +++ b/src_rbd_shaders/dynamics/multibody/contact_sensor.rs @@ -1,7 +1,7 @@ //! Contact "force sensor" readout for RL observations. use super::types::{ - MB_CONTACT_KIND_NORMAL, MAX_MB_CONTACT_CONSTRAINTS_PER_MB, MultibodyContactConstraint, + MAX_MB_CONTACT_CONSTRAINTS_PER_MB, MB_CONTACT_KIND_NORMAL, MultibodyContactConstraint, MultibodyInfo, }; use crate::utils::BatchIndices; diff --git a/src_rbd_shaders/dynamics/multibody/env_reset.rs b/src_rbd_shaders/dynamics/multibody/env_reset.rs index 6dd0dea0..36532767 100644 --- a/src_rbd_shaders/dynamics/multibody/env_reset.rs +++ b/src_rbd_shaders/dynamics/multibody/env_reset.rs @@ -57,7 +57,10 @@ pub fn gpu_mb_env_reset( } 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)); + dof_state.write( + (i * nb + env) as usize, + staging_dofs.read((dpb + i) as usize), + ); } } @@ -141,7 +144,10 @@ pub fn gpu_mb_env_reset_batch( (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)); + dof_state.write( + (i * nb + env) as usize, + dof_vels.read((r * dpb + i) as usize), + ); } } From 05b535133c70d27f0db0b4e17eb0849fe362cdcb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Fri, 28 Aug 2026 08:14:31 +0200 Subject: [PATCH 25/41] refactor(rbd): move the contact merge cosine into RbdSimParams --- src_rbd/broad_phase/narrow_phase.rs | 4 +--- src_rbd/pipeline/insertion_removal.rs | 10 ++-------- src_rbd/pipeline/rbd_state.rs | 16 +++++++++------- src_rbd/pipeline/rbd_state_from_rapier.rs | 10 ++-------- src_rbd/pipeline/rbd_step.rs | 1 - src_rbd_shaders/broad_phase/narrow_phase.rs | 6 ++---- src_rbd_shaders/dynamics/sim_params.rs | 8 ++++++++ 7 files changed, 24 insertions(+), 31 deletions(-) diff --git a/src_rbd/broad_phase/narrow_phase.rs b/src_rbd/broad_phase/narrow_phase.rs index e7c7b0b5..ef3995c1 100644 --- a/src_rbd/broad_phase/narrow_phase.rs +++ b/src_rbd/broad_phase/narrow_phase.rs @@ -61,7 +61,6 @@ impl GpuNarrowPhase { // Optional: merge each collider pair's manifolds into one before the // solvers see them. `false` skips the kernel entirely. reduce_contacts: bool, - merge_cos: &Tensor, pairs_offsets: &mut Tensor, pfm_offsets: &mut Tensor, ) -> Result<(), GpuBackendError> { @@ -146,11 +145,10 @@ impl GpuNarrowPhase { contacts_len, batch_indices, sim_params, - merge_cos, )?; } #[cfg(not(feature = "dim3"))] - let _ = (reduce_contacts, merge_cos); + let _ = reduce_contacts; self.init_contacts_indirect_args.call( pass, 256u32, diff --git a/src_rbd/pipeline/insertion_removal.rs b/src_rbd/pipeline/insertion_removal.rs index 804e310b..98a8effb 100644 --- a/src_rbd/pipeline/insertion_removal.rs +++ b/src_rbd/pipeline/insertion_removal.rs @@ -143,12 +143,6 @@ impl RbdState { BufferUsages::STORAGE | BufferUsages::UNIFORM, ) .unwrap(); - let contact_merge_cos = Tensor::scalar( - backend, - crate::shaders::broad_phase::COS_MERGE_ANGLE, - BufferUsages::STORAGE | BufferUsages::UNIFORM | BufferUsages::COPY_DST, - ) - .unwrap(); // Two-element readback: the (max) collision-pair count and the uncolored count. let resize_readback = GpuReadback::new(backend, 2).unwrap(); let collision_pairs_indirect = @@ -256,9 +250,10 @@ impl RbdState { sim_params: Tensor::vector( backend, &all_sim_params, - BufferUsages::STORAGE | BufferUsages::UNIFORM, + BufferUsages::STORAGE | BufferUsages::UNIFORM | BufferUsages::COPY_DST, ) .unwrap(), + all_sim_params, vels: Tensor::vector(backend, &all_vels, rw).unwrap(), #[cfg(feature = "dim3")] reset_templates_bodies: None, @@ -286,7 +281,6 @@ impl RbdState { collision_pairs_len, collision_pairs_len_max, num_batches_uniform, - contact_merge_cos, resize_readback, collision_pairs_indirect, contacts_per_batch_cpu, diff --git a/src_rbd/pipeline/rbd_state.rs b/src_rbd/pipeline/rbd_state.rs index 6f33232d..fcf327fa 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) all_sim_params: Vec, /// 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 @@ -166,7 +169,6 @@ pub struct RbdState { 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`). - pub(super) contact_merge_cos: Tensor, /// `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 @@ -384,12 +386,12 @@ impl RbdState { /// 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) { - self.contact_merge_cos = Tensor::scalar( - backend, - cos, - BufferUsages::STORAGE | BufferUsages::UNIFORM | BufferUsages::COPY_DST, - ) - .unwrap(); + let mut params = self.all_sim_params.clone(); + for p in &mut params { + p.contact_merge_cos = cos; + } + backend.write_buffer(self.sim_params.buffer_mut(), 0, ¶ms); + self.all_sim_params = params; } /// The gravity uniform shared by every solver kernel. diff --git a/src_rbd/pipeline/rbd_state_from_rapier.rs b/src_rbd/pipeline/rbd_state_from_rapier.rs index 7e458c91..cb0a2c43 100644 --- a/src_rbd/pipeline/rbd_state_from_rapier.rs +++ b/src_rbd/pipeline/rbd_state_from_rapier.rs @@ -616,12 +616,6 @@ impl RbdState { BufferUsages::STORAGE | BufferUsages::UNIFORM, ) .unwrap(); - let contact_merge_cos = Tensor::scalar( - backend, - crate::shaders::broad_phase::COS_MERGE_ANGLE, - BufferUsages::STORAGE | BufferUsages::UNIFORM | BufferUsages::COPY_DST, - ) - .unwrap(); // Two-element readback: the (max) collision-pair count and the uncolored count. let resize_readback = GpuReadback::new(backend, 2).unwrap(); let collision_pairs_indirect = @@ -785,9 +779,10 @@ impl RbdState { sim_params: Tensor::vector( backend, &all_sim_params, - BufferUsages::STORAGE | BufferUsages::UNIFORM, + BufferUsages::STORAGE | BufferUsages::UNIFORM | BufferUsages::COPY_DST, ) .unwrap(), + all_sim_params, vels: Tensor::vector(backend, &all_vels, storage | BufferUsages::COPY_DST).unwrap(), #[cfg(feature = "dim3")] reset_templates_bodies: None, @@ -836,7 +831,6 @@ impl RbdState { collision_pairs_len, collision_pairs_len_max, num_batches_uniform, - contact_merge_cos, resize_readback, collision_pairs_indirect, contacts_per_batch_cpu, diff --git a/src_rbd/pipeline/rbd_step.rs b/src_rbd/pipeline/rbd_step.rs index c834c2e9..7342189c 100644 --- a/src_rbd/pipeline/rbd_step.rs +++ b/src_rbd/pipeline/rbd_step.rs @@ -326,7 +326,6 @@ impl RbdPipeline { &state.collider_materials, &state.sim_params, self.contact_reduction, - &state.contact_merge_cos, &mut state.pairs_flat_offsets, &mut state.pfm_flat_offsets, )?; diff --git a/src_rbd_shaders/broad_phase/narrow_phase.rs b/src_rbd_shaders/broad_phase/narrow_phase.rs index 8e8c7cfe..1f776509 100644 --- a/src_rbd_shaders/broad_phase/narrow_phase.rs +++ b/src_rbd_shaders/broad_phase/narrow_phase.rs @@ -133,11 +133,9 @@ pub fn gpu_reduce_contacts( #[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, - // Cosine of the maximum angle between two manifolds' normals for them to - // share a cluster. See [`COS_MERGE_ANGLE`]. - #[spirv(uniform, descriptor_set = 0, binding = 4)] merge_cos: &f32, ) { 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); @@ -151,7 +149,7 @@ pub fn gpu_reduce_contacts( 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 + && 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. diff --git a/src_rbd_shaders/dynamics/sim_params.rs b/src_rbd_shaders/dynamics/sim_params.rs index 545c6dff..4746b68c 100644 --- a/src_rbd_shaders/dynamics/sim_params.rs +++ b/src_rbd_shaders/dynamics/sim_params.rs @@ -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, } From 73df1fca2f3867150c8d502dbaec16f77be15865 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Fri, 28 Aug 2026 09:02:19 +0200 Subject: [PATCH 26/41] refactor: move read_multibody_links onto NexusState and drive every env from control_multibody_motors --- crates/nexus_python3d/src/nexus.rs | 72 +++++++++++++++++++--- crates/nexus_python3d/src/viewer.rs | 50 +-------------- src/state.rs | 94 ++++++++++++++++++++++------- src_rbd/pipeline/rbd_state.rs | 2 +- src_viewer/viewer.rs | 45 -------------- 5 files changed, 137 insertions(+), 126 deletions(-) diff --git a/crates/nexus_python3d/src/nexus.rs b/crates/nexus_python3d/src/nexus.rs index 5e9462dd..e1c9c3ff 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; @@ -373,20 +374,18 @@ impl NexusState { } /// Applies one MJCF control vector (one entry per actuator, in - /// `actuator_names` order) to 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 in one buffer write. + /// `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. This - /// is the GPU counterpart of stepping rapier natively with actuators. - #[pyo3(signature = (viewer, ctrl, env=0))] + /// `NexusPipeline.simulate` steps the solver against the new targets. + #[pyo3(signature = (viewer, ctrl))] fn apply_actuator_controls( &mut self, viewer: PyRef, ctrl: Vec, - env: usize, ) -> PyResult<()> { let Some(handles) = self.1.as_ref() else { return Err(PyRuntimeError::new_err( @@ -402,7 +401,7 @@ impl NexusState { } let handles = handles.clone(); self.0 - .control_multibody_motors(viewer.backend(), env, |world| { + .control_multibody_motors(viewer.backend(), |_, world| { handles.apply_controls_multibody( &mut world.bodies, &mut world.multibody_joints, @@ -412,6 +411,61 @@ impl NexusState { .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 ------------------------------------------------------- fn set_rbd_steps_per_frame(&mut self, steps: u32) { diff --git a/crates/nexus_python3d/src/viewer.rs b/crates/nexus_python3d/src/viewer.rs index 25ad62d7..d5e6983e 100644 --- a/crates/nexus_python3d/src/viewer.rs +++ b/crates/nexus_python3d/src/viewer.rs @@ -10,7 +10,7 @@ use crate::nexus::{GpuTimestamps, NexusState}; use crate::rbd::{RigidBodyHandle, SharedShape}; use khal::backend::GpuBackend; use nexus_viewer3d::NexusViewer as RViewer; -use numpy::{IntoPyArray, PyArray2, PyArray3, PyArrayMethods}; +use numpy::{IntoPyArray, PyArray3, PyArrayMethods}; use pyo3::exceptions::PyRuntimeError; use pyo3::prelude::*; @@ -282,54 +282,6 @@ impl NexusViewer { .map_err(|e| PyRuntimeError::new_err(format!("{e:?}"))) } - /// Reads back environment `env`'s multibody link states from the GPU in one - /// readback. Returns five float32 numpy arrays, one row per link (in the - /// GPU build's traversal order — multibodies, then links, parent before - /// child; the same order `NexusState.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). - #[pyo3(signature = (state, env=0))] - #[allow(clippy::type_complexity)] - fn read_multibody_links<'py>( - &mut self, - py: Python<'py>, - state: PyRef, - env: u32, - ) -> ( - Bound<'py, PyArray2>, - Bound<'py, PyArray2>, - Bound<'py, PyArray2>, - Bound<'py, PyArray2>, - Bound<'py, PyArray2>, - ) { - let links = pollster::block_on(self.inner_mut().read_multibody_links(&state.0, env)); - 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(), - ) - } - // --- misc ------------------------------------------------------------- fn clear_scene(&mut self) { diff --git a/src/state.rs b/src/state.rs index 39f1a752..0c54e160 100644 --- a/src/state.rs +++ b/src/state.rs @@ -12,7 +12,7 @@ use crate::rbd::dynamics::{ body::{BodyCoupling, RapierBodyCouplingEntry}, }; use crate::rbd::pipeline::{RbdCapacities, RbdResizePolicy, RbdState, RunStats}; -use khal::backend::{GpuBackend, GpuBackendError}; +use khal::backend::{Backend, GpuBackend, GpuBackendError}; /// Handle referencing a rigid-body managed by a [`NexusState`]. #[derive(Copy, Clone, PartialEq, Eq, Debug, Hash)] @@ -441,38 +441,88 @@ impl NexusState { &mut self.rbd_envs[env] } - /// Runtime actuation entry point: mutates environment `env`'s rapier - /// multibody joints through `f` (e.g. `rapier3d-mjcf`'s - /// `apply_controls_multibody`, which implements MJCF actuator semantics), - /// then pushes the refreshed joint data — motor targets/gains, limits — to - /// the GPU multibody links in one buffer write. + /// 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. /// - /// Unlike [`Self::rbd_world_mut`] this does NOT mark the world dirty: motor - /// updates are per-step control, not a topology change, so no GPU rebuild - /// is triggered. Call after [`Self::finalize`]; a no-op before it. - #[cfg(all(feature = "dim3", feature = "rbd"))] + /// `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. pub fn control_multibody_motors( &mut self, backend: &GpuBackend, - env: usize, - f: F, + mut f: F, ) -> Result<(), GpuBackendError> where - F: FnOnce(&mut PhysicsWorld), + F: FnMut(usize, &mut PhysicsWorld), { - let world = &mut self.rbd_envs[env]; - f(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, - )?; + 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`]. /// diff --git a/src_rbd/pipeline/rbd_state.rs b/src_rbd/pipeline/rbd_state.rs index fcf327fa..be519756 100644 --- a/src_rbd/pipeline/rbd_state.rs +++ b/src_rbd/pipeline/rbd_state.rs @@ -390,7 +390,7 @@ impl RbdState { for p in &mut params { p.contact_merge_cos = cos; } - backend.write_buffer(self.sim_params.buffer_mut(), 0, ¶ms); + let _ = backend.write_buffer(self.sim_params.buffer_mut(), 0, ¶ms); self.all_sim_params = params; } diff --git a/src_viewer/viewer.rs b/src_viewer/viewer.rs index 29edddc1..c07b8d54 100644 --- a/src_viewer/viewer.rs +++ b/src_viewer/viewer.rs @@ -1217,51 +1217,6 @@ impl NexusViewer { .set_denoise(enabled); } - /// Reads back environment `env`'s multibody link workspaces from the GPU in - /// one readback: per link, the generalized joint coordinates, accumulated - /// joint rotation, world pose, and world-space velocity. Links are in the - /// GPU build's traversal order (multibodies, then links, parent before - /// child) — the same order `NexusState::control_multibody_motors` targets. - /// Empty when no multibody state exists. - /// - /// Velocities are only meaningful after the first simulated step (the - /// forward-kinematics pass fills them); coordinates and poses are valid - /// from `finalize`. - #[cfg(feature = "dim3")] - pub async fn read_multibody_links( - &mut self, - state: &NexusState, - env: u32, - ) -> Vec { - let Some(rbd) = state.rbd.as_ref() else { - return Vec::new(); - }; - let mbs = rbd.multibodies(); - let stride = mbs.links_per_batch() as usize; - if stride == 0 { - return Vec::new(); - } - // The workspace is stored batch-interleaved SoA (quads), so read the - // raw buffer and decode it back into one struct per link. - let mut raw: Vec = bytemuck::zeroed_vec(mbs.links_workspace().len() as usize); - if self - .backend() - .slow_read_buffer(mbs.links_workspace().buffer(), &mut raw) - .await - .is_err() - { - return Vec::new(); - } - let all = nexus::rbd::shaders::dynamics::ws_soa_to_structs( - &raw, - mbs.links_per_batch(), - mbs.num_batches(), - ); - let start = (env as usize * stride).min(all.len()); - let end = (start + stride).min(all.len()); - all[start..end].to_vec() - } - /// Draws example-specific egui widgets into the current frame's UI pass. /// /// Call this once per frame, after [`Self::render_frame`], to overlay a From 1835b6f49fc1f320310f6ab97969a54b86bd9138 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Fri, 28 Aug 2026 09:47:44 +0200 Subject: [PATCH 27/41] test(rbd): add a headless many-small-environments step-timing harness --- src_rbd/pipeline/bench_narrow_phase.rs | 79 ++++++++++++++++++++++++++ src_rbd/pipeline/mod.rs | 2 + 2 files changed, 81 insertions(+) create mode 100644 src_rbd/pipeline/bench_narrow_phase.rs diff --git a/src_rbd/pipeline/bench_narrow_phase.rs b/src_rbd/pipeline/bench_narrow_phase.rs new file mode 100644 index 00000000..3765d5cd --- /dev/null +++ b/src_rbd/pipeline/bench_narrow_phase.rs @@ -0,0 +1,79 @@ +//! 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}; + +/// 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) { + // Metal, not WebGPU: `gpu_mb_init_joint_constraints` currently binds 10 + // storage buffers, over WebGPU's per-stage limit of 8. + let backend = GpuBackend::Metal(khal::backend::metal::Metal::new().unwrap()); + + 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, + collisions_capacity: 256, + ..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() { + for envs in [1u32, 64, 256, 1024, 4096] { + run_bench(envs, 4, 200).await; + } +} diff --git a/src_rbd/pipeline/mod.rs b/src_rbd/pipeline/mod.rs index 65a40df4..352f617c 100644 --- a/src_rbd/pipeline/mod.rs +++ b/src_rbd/pipeline/mod.rs @@ -4,6 +4,8 @@ //! simulation step on the GPU. The pipeline manages collision detection, contact generation, //! constraint solving, and integration. +#[cfg(test)] +mod bench_narrow_phase; mod insertion_removal; mod lbvh_validation; mod rbd_state; From 5901f0b14bf624a251ed08bb7483b19cddd883c4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Fri, 28 Aug 2026 10:23:06 +0200 Subject: [PATCH 28/41] revert(rbd): drop the flat 1-D narrow-phase dispatch Measured 7-21% slower on Metal. --- src_rbd/broad_phase/narrow_phase.rs | 50 ++----- src_rbd/pipeline/insertion_removal.rs | 6 - src_rbd/pipeline/rbd_state.rs | 6 - src_rbd/pipeline/rbd_state_from_rapier.rs | 6 - src_rbd/pipeline/rbd_step.rs | 6 +- src_rbd_shaders/broad_phase/narrow_phase.rs | 152 ++++++-------------- 6 files changed, 56 insertions(+), 170 deletions(-) diff --git a/src_rbd/broad_phase/narrow_phase.rs b/src_rbd/broad_phase/narrow_phase.rs index ef3995c1..94cbdecc 100644 --- a/src_rbd/broad_phase/narrow_phase.rs +++ b/src_rbd/broad_phase/narrow_phase.rs @@ -6,9 +6,9 @@ use crate::shaders::PaddedVector; #[cfg(feature = "dim3")] use crate::shaders::broad_phase::GpuReduceContacts; use crate::shaders::broad_phase::{ - CollisionPair, GpuFlattenBatchesDispatch, GpuNarrowPhaseInitContactsDispatch, - GpuNarrowPhasePfmPfm, GpuNarrowPhaseShapeShape, GpuNarrowPhaseShapeShapeDeferred, - GpuResetNarrowPhase, NarrowPhasePfmPair, + CollisionPair, GpuInitPfmPfmDispatch, GpuNarrowPhaseInitContactsDispatch, GpuNarrowPhasePfmPfm, + GpuNarrowPhaseShapeShape, GpuNarrowPhaseShapeShapeDeferred, GpuResetNarrowPhase, + NarrowPhasePfmPair, }; use crate::shaders::shapes::Shape; use khal::Shader; @@ -26,11 +26,7 @@ pub struct GpuNarrowPhase { narrow_phase_pfm_pfm: GpuNarrowPhasePfmPfm, #[cfg(feature = "dim3")] reduce_contacts: GpuReduceContacts, - /// Builds the flat 1-D dispatch grid + prefix offsets for a per-batch - /// work-list (used for both the collision pairs and the PFM pairs), so the - /// kernels pack items from many batches into full warps instead of one - /// mostly-idle workgroup per batch. - flatten_batches: GpuFlattenBatchesDispatch, + init_pfm_pfm_indirect_args: GpuInitPfmPfmDispatch, init_contacts_indirect_args: GpuNarrowPhaseInitContactsDispatch, } @@ -45,8 +41,8 @@ impl GpuNarrowPhase { vertices: &Tensor, indices: &Tensor, collision_pairs: &Tensor, - collision_pairs_len: &mut Tensor, - collision_pairs_indirect: &mut Tensor<[u32; 3]>, + collision_pairs_len: &Tensor, + collision_pairs_indirect: &Tensor<[u32; 3]>, contacts: &mut Tensor, contacts_len: &mut Tensor, contacts_indirect: &mut Tensor<[u32; 3]>, @@ -61,30 +57,16 @@ impl GpuNarrowPhase { // Optional: merge each collider pair's manifolds into one before the // solvers see them. `false` skips the kernel entirely. reduce_contacts: bool, - pairs_offsets: &mut Tensor, - pfm_offsets: &mut Tensor, ) -> Result<(), GpuBackendError> { let num_batches = contacts_len.len() as u32; self.reset_narrow_phase .call(pass, [num_batches, 1, 1], contacts_len, pfm_pairs_len)?; - // The broad phase wrote a `[max/64, num_batches, 1]` grid into - // `collision_pairs_indirect`; rewrite it (and derive the offsets) for - // the flat layout. Nothing else consumes the batched form. - self.flatten_batches.call( - pass, - 1u32, - collision_pairs_len, - pairs_offsets, - collision_pairs_indirect, - batch_indices, - )?; - self.narrow_phase.call( pass, - &*collision_pairs_indirect, + collision_pairs_indirect, collision_pairs, - pairs_offsets, + collision_pairs_len, poses, shapes, contacts, @@ -99,9 +81,9 @@ impl GpuNarrowPhase { // separate dispatch so each pass fits 8 storage buffers). self.narrow_phase_deferred.call( pass, - &*collision_pairs_indirect, + collision_pairs_indirect, collision_pairs, - pairs_offsets, + collision_pairs_len, poses, shapes, pfm_pairs, @@ -112,21 +94,15 @@ impl GpuNarrowPhase { indices, )?; - self.flatten_batches.call( - pass, - 1u32, - pfm_pairs_len, - pfm_offsets, - pfm_pairs_indirect, - batch_indices, - )?; + self.init_pfm_pfm_indirect_args + .call(pass, 256u32, pfm_pairs_len, pfm_pairs_indirect)?; self.narrow_phase_pfm_pfm.call( pass, &*pfm_pairs_indirect, contacts, contacts_len, pfm_pairs, - pfm_offsets, + pfm_pairs_len, batch_indices, vertices, indices, diff --git a/src_rbd/pipeline/insertion_removal.rs b/src_rbd/pipeline/insertion_removal.rs index 98a8effb..45411b80 100644 --- a/src_rbd/pipeline/insertion_removal.rs +++ b/src_rbd/pipeline/insertion_removal.rs @@ -169,10 +169,6 @@ impl RbdState { BufferUsages::STORAGE | BufferUsages::COPY_SRC, ) .unwrap(); - let pairs_flat_offsets = - Tensor::vector_uninit(backend, num_batches + 1, BufferUsages::STORAGE).unwrap(); - let pfm_flat_offsets = - Tensor::vector_uninit(backend, num_batches + 1, BufferUsages::STORAGE).unwrap(); let old_constraints = Tensor::vector_uninit(backend, collisions_capacity * num_batches, storage).unwrap(); let old_constraint_builders = @@ -294,8 +290,6 @@ impl RbdState { pfm_pairs, pfm_pairs_len, pfm_pairs_indirect, - pairs_flat_offsets, - pfm_flat_offsets, old_constraints, old_constraint_builders, old_constraints_counts, diff --git a/src_rbd/pipeline/rbd_state.rs b/src_rbd/pipeline/rbd_state.rs index be519756..e214d761 100644 --- a/src_rbd/pipeline/rbd_state.rs +++ b/src_rbd/pipeline/rbd_state.rs @@ -193,12 +193,6 @@ pub struct RbdState { pub(super) pfm_pairs: Tensor, pub(super) pfm_pairs_len: Tensor, pub(super) pfm_pairs_indirect: Tensor<[u32; 3]>, - /// Flat-dispatch prefix offsets (`num_batches + 1`) over the per-batch - /// collision-pair / PFM work-lists, rebuilt on the GPU each step by - /// `gpu_flatten_batches_dispatch` so the narrow-phase kernels can pack - /// items from many batches into full warps. - pub(super) pairs_flat_offsets: Tensor, - pub(super) pfm_flat_offsets: Tensor, pub(super) contacts: Tensor, pub(super) contacts_len: Tensor, pub(super) contacts_indirect: Tensor<[u32; 3]>, diff --git a/src_rbd/pipeline/rbd_state_from_rapier.rs b/src_rbd/pipeline/rbd_state_from_rapier.rs index cb0a2c43..b6530c3d 100644 --- a/src_rbd/pipeline/rbd_state_from_rapier.rs +++ b/src_rbd/pipeline/rbd_state_from_rapier.rs @@ -651,10 +651,6 @@ impl RbdState { BufferUsages::STORAGE | BufferUsages::COPY_SRC, ) .unwrap(); - let pairs_flat_offsets = - Tensor::vector_uninit(backend, num_batches + 1, BufferUsages::STORAGE).unwrap(); - let pfm_flat_offsets = - Tensor::vector_uninit(backend, num_batches + 1, BufferUsages::STORAGE).unwrap(); let old_constraints = Tensor::vector_uninit( backend, capacities.collisions_capacity * num_batches, @@ -844,8 +840,6 @@ impl RbdState { pfm_pairs, pfm_pairs_len, pfm_pairs_indirect, - pairs_flat_offsets, - pfm_flat_offsets, old_constraints, old_constraint_builders, old_constraints_counts, diff --git a/src_rbd/pipeline/rbd_step.rs b/src_rbd/pipeline/rbd_step.rs index 7342189c..39afd1c1 100644 --- a/src_rbd/pipeline/rbd_step.rs +++ b/src_rbd/pipeline/rbd_step.rs @@ -312,8 +312,8 @@ impl RbdPipeline { &state.vertex_buffers, &state.index_buffers, &state.collision_pairs, - &mut state.collision_pairs_len, - &mut state.collision_pairs_indirect, + &state.collision_pairs_len, + &state.collision_pairs_indirect, &mut state.contacts, &mut state.contacts_len, &mut state.contacts_indirect, @@ -326,8 +326,6 @@ impl RbdPipeline { &state.collider_materials, &state.sim_params, self.contact_reduction, - &mut state.pairs_flat_offsets, - &mut state.pfm_flat_offsets, )?; drop(pass); diff --git a/src_rbd_shaders/broad_phase/narrow_phase.rs b/src_rbd_shaders/broad_phase/narrow_phase.rs index 1f776509..34e26746 100644 --- a/src_rbd_shaders/broad_phase/narrow_phase.rs +++ b/src_rbd_shaders/broad_phase/narrow_phase.rs @@ -17,10 +17,7 @@ use crate::{PaddedVector, Pose, Vector}; use khal_std::glamx::UVec3; use khal_std::index::MaybeIndexUnchecked; use khal_std::macros::{spirv, spirv_bindgen}; -use khal_std::{ - iter::StepRng, - sync::{atomic_add_u32, atomic_load_u32}, -}; +use khal_std::{iter::StepRng, sync::atomic_add_u32}; use super::lbvh::{MAX_REDUCE_LANES, max_len_indirect_args}; use crate::broad_phase::CollisionPair; @@ -221,64 +218,7 @@ pub fn gpu_reduce_contacts( } } -/// Builds the flat-dispatch layout for a per-batch work-list: exclusive prefix -/// offsets (so item `t` of the flat range maps back to a batch via -/// `find_batch`) and the matching 1-D indirect grid. -/// -/// This replaces the max-over-batches indirect grids for the narrow-phase -/// kernels: with `[max/64, num_batches, 1]` every batch rounds its handful of -/// pairs up to a full 64-lane workgroup (a robot env has ~7 pairs → ~11% lane -/// occupancy, thousands of near-empty workgroups). The flat grid packs items -/// from consecutive batches into the same warps: `[total/64, 1, 1]`. -/// -/// Serial over batches in one thread — same pattern (and cost) as the existing -/// `gpu_narrow_phase_init_contacts_dispatch` max-scan. -#[spirv_bindgen] -#[spirv(compute(threads(1)))] -pub fn gpu_flatten_batches_dispatch( - // NOTE: `lens` is mutable only for `atomic_load_u32` (see the note on - // `gpu_narrow_phase_init_contacts_dispatch`). - #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] lens: &mut [u32], - // `num_batches + 1` entries; `offsets[num_batches]` is the total. - #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] offsets: &mut [u32], - #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] indirect_args: &mut [u32; 3], - #[spirv(uniform, descriptor_set = 0, binding = 3)] batch_ids: &BatchIndices, -) { - let num_batches = lens.len(); - // Same clamp as the consuming kernels: a batch's list may overflow its - // capacity slot; the overflowing tail was never written and must not be - // walked. - let capacity = batch_ids.contacts_batch_capacity; - let mut total = 0u32; - for i in 0..num_batches { - offsets.write(i, total); - total += atomic_load_u32(lens.at_mut(i)).min(capacity); - } - offsets.write(num_batches, total); - *indirect_args.at_mut(0) = total.div_ceil(WORKGROUP_SIZE); - *indirect_args.at_mut(1) = 1; - *indirect_args.at_mut(2) = 1; -} - -/// Largest `b` with `offsets[b] <= t` — the batch owning flat item `t`. -/// Invariant: `offsets[0] == 0 <= t < offsets[num_batches]`. -fn find_batch(offsets: &[u32], num_batches: u32, t: u32) -> u32 { - let mut lo = 0u32; - let mut hi = num_batches; - // Bounded loop instead of `while` (see the trimesh BVH walk for why). - for _ in 0..32 { - if lo + 1 >= hi { - break; - } - let mid = (lo + hi) / 2; - if offsets.read(mid as usize) <= t { - lo = mid; - } else { - hi = mid; - } - } - lo -} +const PREDICTION: f32 = 2.0e-3; // TODO: make the prediction configurable. /// Narrow phase, pass 1 of 2: analytic shape-shape contacts for ball / cuboid /// pairs, written straight into the `contacts` buffer. @@ -291,10 +231,7 @@ pub fn gpu_narrow_phase_shape_shape( #[spirv(global_invocation_id)] invocation_id: UVec3, #[spirv(num_workgroups)] num_workgroups: UVec3, #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] collision_pairs: &[CollisionPair], - // Flat-dispatch prefix offsets from `gpu_flatten_batches_dispatch` - // (`num_batches + 1` entries; replaces the per-batch `collision_pairs_len`, - // which it already folds in, clamped to capacity). - #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] pairs_offsets: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] collision_pairs_len: &[u32], #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] poses: &[Pose], #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] shapes: &[Shape], #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] contacts: &mut [IndexedManifold], @@ -310,25 +247,23 @@ pub fn gpu_narrow_phase_shape_shape( ) { 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; - // Flat over all batches' pairs: consecutive lanes take consecutive pairs - // regardless of which batch owns them, so warps stay packed even when each - // batch only has a handful. - let num_batches = pairs_offsets.len() - 1; - let total = pairs_offsets.read(num_batches); - - for t in StepRng::new(invocation_id.x..total, num_threads) { - let batch_id = find_batch(pairs_offsets, num_batches as u32, t); - let i = t - pairs_offsets.read(batch_id as usize); + let collision_pairs = batch_ids.contact_batch(batch_id, collision_pairs); + let poses = batch_ids.coll_batch(batch_id, poses); + let shapes = batch_ids.coll_batch(batch_id, shapes); + let collider_materials = batch_ids.coll_batch(batch_id, collider_materials); + let mut contacts = batch_ids.contact_batch_mut(batch_id, contacts); + let contacts_len = contacts_len.at_mut(batch_id as usize); - let collision_pairs = batch_ids.contact_batch(batch_id, collision_pairs); - let poses = batch_ids.coll_batch(batch_id, poses); - let shapes = batch_ids.coll_batch(batch_id, shapes); - let collider_materials = batch_ids.coll_batch(batch_id, collider_materials); - let mut contacts = SliceMut(&mut *contacts, batch_ids.contacts_start(batch_id)); - let contacts_len = contacts_len.at_mut(batch_id as usize); + // NOTE: `collision_pairs_len` might be greater than `contacts_batch_apacity` if the + // narrow-phase found more pairs than the buffer can contain. + let len = collision_pairs_len + .read(batch_id as usize) + .min(contacts_batch_capacity as u32); + 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. Pair ids are @@ -419,8 +354,7 @@ pub fn gpu_narrow_phase_shape_shape_deferred( #[spirv(global_invocation_id)] invocation_id: UVec3, #[spirv(num_workgroups)] num_workgroups: UVec3, #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] collision_pairs: &[CollisionPair], - // Flat-dispatch prefix offsets (see `gpu_narrow_phase_shape_shape`). - #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] pairs_offsets: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] collision_pairs_len: &[u32], #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] poses: &[Pose], #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] shapes: &[Shape], #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] @@ -436,27 +370,25 @@ pub fn gpu_narrow_phase_shape_shape_deferred( ) { let prediction = params.prediction_distance(); let num_threads = num_workgroups.x * WORKGROUP_SIZE; - // Every batch is allocated the same capacity, so this is batch-independent. + let batch_id = invocation_id.y; let contacts_batch_capacity = batch_ids.contacts_batch_capacity as usize; - let num_batches = pairs_offsets.len() - 1; - let total = pairs_offsets.read(num_batches); + let collision_pairs = batch_ids.contact_batch(batch_id, collision_pairs); + let poses = batch_ids.coll_batch(batch_id, poses); + let shapes = batch_ids.coll_batch(batch_id, shapes); + let mut pfm_pairs = batch_ids.contact_batch_mut(batch_id, pfm_pairs); + let pfm_pairs_len = pfm_pairs_len.at_mut(batch_id as usize); + + let len = collision_pairs_len + .read(batch_id as usize) + .min(contacts_batch_capacity as u32); // NOTE: same-body collider pairs are *not* filtered in this pass — it is // already at the 8-storage-buffer WebGPU limit and can't take the // `collider_parent` binding. The complex pairs it emits are filtered // downstream in `gpu_narrow_phase_pfm_pfm` (which has room) before any // contact is written. - for t in StepRng::new(invocation_id.x..total, num_threads) { - let batch_id = find_batch(pairs_offsets, num_batches as u32, t); - let i = t - pairs_offsets.read(batch_id as usize); - - let collision_pairs = batch_ids.contact_batch(batch_id, collision_pairs); - let poses = batch_ids.coll_batch(batch_id, poses); - let shapes = batch_ids.coll_batch(batch_id, shapes); - let mut pfm_pairs = SliceMut(&mut *pfm_pairs, batch_ids.contacts_start(batch_id)); - let pfm_pairs_len = pfm_pairs_len.at_mut(batch_id as usize); - + for i in StepRng::new(invocation_id.x..len, num_threads) { let pair = collision_pairs[i as usize]; let shape1 = &shapes[pair.colliders.x as usize]; let shape2 = &shapes[pair.colliders.y as usize]; @@ -779,9 +711,7 @@ pub fn gpu_narrow_phase_pfm_pfm( #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] contacts: &mut [IndexedManifold], #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] contacts_len: &mut [u32], #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] pfm_pairs: &[NarrowPhasePfmPair], - // Flat-dispatch prefix offsets over the per-batch PFM work-lists (see - // `gpu_narrow_phase_shape_shape`; replaces the per-batch `pfm_pairs_len`). - #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] pfm_offsets: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] pfm_pairs_len: &[u32], // NOTE: we assume that max_pfm_pairs == contacts_batch_capacity // And we assume all batch dimensions are given the same buffer allocation sizes // (i.e. the same `contacts_batch_capacity`). @@ -799,20 +729,20 @@ pub fn gpu_narrow_phase_pfm_pfm( ) { 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; - let num_batches = pfm_offsets.len() - 1; - let total = pfm_offsets.read(num_batches); - - for t in StepRng::new(invocation_id.x..total, num_threads) { - let batch_id = find_batch(pfm_offsets, num_batches as u32, t); - let i = t - pfm_offsets.read(batch_id as usize); - - let mut contacts = SliceMut(&mut *contacts, batch_ids.contacts_start(batch_id)); - let collider_materials = batch_ids.coll_batch(batch_id, collider_materials); - let pfm_pairs = batch_ids.contact_batch(batch_id, pfm_pairs); - let contacts_len = contacts_len.at_mut(batch_id as usize); - + let mut contacts = batch_ids.contact_batch_mut(batch_id, contacts); + let collider_materials = batch_ids.coll_batch(batch_id, collider_materials); + let pfm_pairs = batch_ids.contact_batch(batch_id, pfm_pairs); + let contacts_len = contacts_len.at_mut(batch_id as usize); + // The producer counter can exceed the allocation on overflow (writes are + // skipped past capacity); clamp so we never read uninitialized slots. + let pfm_pairs_len = pfm_pairs_len + .read(batch_id as usize) + .min(contacts_batch_capacity as u32); + + for i in StepRng::new(invocation_id.x..pfm_pairs_len, num_threads) { let pair = pfm_pairs[i as usize]; // Resolve the parent rigid-bodies and skip same-body collider pairs. This // is where the deferred (PFM / trimesh / polyline) pairs get the same-body From cc44854a7fba776a47b53bcfd7d4d5f52328d6b7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Fri, 28 Aug 2026 10:58:52 +0200 Subject: [PATCH 29/41] =?UTF-8?q?chore:=E2=80=AFcleanup=20comments?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- crates/nexus_python3d/src/nexus.rs | 6 +++--- src/state.rs | 8 ++++---- src_rbd/dynamics/multibody/multibody_set.rs | 8 +++----- src_rbd/pipeline/rbd_state_from_rapier.rs | 6 +++--- src_rbd_shaders/broad_phase/narrow_phase.rs | 11 +++++------ 5 files changed, 18 insertions(+), 21 deletions(-) diff --git a/crates/nexus_python3d/src/nexus.rs b/crates/nexus_python3d/src/nexus.rs index e1c9c3ff..b0bb9e60 100644 --- a/crates/nexus_python3d/src/nexus.rs +++ b/crates/nexus_python3d/src/nexus.rs @@ -325,9 +325,9 @@ impl NexusState { }) } - /// Per-environment collision-pair capacity (default 4096). Lower this - /// before `finalize` when batching many small environments — pair-keyed - /// GPU workspaces scale with `capacity x num_envs`. + /// 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); } diff --git a/src/state.rs b/src/state.rs index 0c54e160..b204664a 100644 --- a/src/state.rs +++ b/src/state.rs @@ -331,11 +331,11 @@ 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 — + /// 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 at 2048 envs binds ~9 GiB unless this is lowered. + /// 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); } diff --git a/src_rbd/dynamics/multibody/multibody_set.rs b/src_rbd/dynamics/multibody/multibody_set.rs index d61c6f22..5db79745 100644 --- a/src_rbd/dynamics/multibody/multibody_set.rs +++ b/src_rbd/dynamics/multibody/multibody_set.rs @@ -652,11 +652,9 @@ impl GpuMultibodySet { /// buffer in one write. /// /// This is the per-step control path for actuated robots: mutate the motors - /// on the CPU rapier joints (e.g. via `rapier3d-mjcf`'s - /// `apply_controls_multibody`, which implements the MJCF actuator - /// semantics), then call this to push the new motor state to the GPU. Only - /// joint data is refreshed — coordinates, velocities and mass properties are - /// untouched, so this cannot be used to teleport links. + /// 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, diff --git a/src_rbd/pipeline/rbd_state_from_rapier.rs b/src_rbd/pipeline/rbd_state_from_rapier.rs index b6530c3d..6a4b80d1 100644 --- a/src_rbd/pipeline/rbd_state_from_rapier.rs +++ b/src_rbd/pipeline/rbd_state_from_rapier.rs @@ -147,9 +147,9 @@ 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(); - // TriMesh dedupe: a `SharedShape` cloned across envs (e.g. shared - // terrain) is serialized into `shape_buffers` once and its `Shape` - // descriptor reused — keyed by the parry shape data pointer. + // 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, diff --git a/src_rbd_shaders/broad_phase/narrow_phase.rs b/src_rbd_shaders/broad_phase/narrow_phase.rs index 34e26746..1a14216b 100644 --- a/src_rbd_shaders/broad_phase/narrow_phase.rs +++ b/src_rbd_shaders/broad_phase/narrow_phase.rs @@ -101,7 +101,7 @@ fn pool_dedup(cand: &mut [ContactPoint; 8], num: &mut usize, pt: ContactPoint, d } /// Optional contact reduction: compacts each batch's contacts in place by -/// merging manifolds that share a collider pair AND a (nearly) parallel +/// 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, @@ -138,7 +138,8 @@ pub fn gpu_reduce_contacts( let mut contacts = batch_ids.contact_batch_mut(batch_id, contacts); let n = (contacts_len.read(batch_id as usize) as usize).min(capacity); - let mut w = 0usize; // write cursor — always ≤ read cursor, in-place safe + // 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; @@ -267,10 +268,8 @@ pub fn gpu_narrow_phase_shape_shape( 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. Pair ids are - // env-local, and `collider_parent` is batch-strided like the other - // per-collider buffers — an unsliced read here silently returned - // batch 0's parents for every batch (masked whenever all envs are - // identical, wrong the moment they aren't). + // 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); From 234ddb6480271861c406332d44da2530fce9b29b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Fri, 28 Aug 2026 11:31:15 +0200 Subject: [PATCH 30/41] fix: gate control_multibody_motors on dim3 so the 2D build still compiles --- src/state.rs | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/state.rs b/src/state.rs index b204664a..94d2b474 100644 --- a/src/state.rs +++ b/src/state.rs @@ -12,7 +12,9 @@ use crate::rbd::dynamics::{ body::{BodyCoupling, RapierBodyCouplingEntry}, }; use crate::rbd::pipeline::{RbdCapacities, RbdResizePolicy, RbdState, RunStats}; -use khal::backend::{Backend, GpuBackend, GpuBackendError}; +#[cfg(feature = "dim3")] +use khal::backend::Backend; +use khal::backend::{GpuBackend, GpuBackendError}; /// Handle referencing a rigid-body managed by a [`NexusState`]. #[derive(Copy, Clone, PartialEq, Eq, Debug, Hash)] @@ -450,6 +452,7 @@ impl NexusState { /// [`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, From 345ae818f3cbd6286b65010db86ee66d1400f9e3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Fri, 28 Aug 2026 13:53:00 +0200 Subject: [PATCH 31/41] fix instability in joint-ball3 demo --- crates/examples3d/rbd_joint_ball3.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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); From afdceef243195a9f1568c3a3e69e104f848da19c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Fri, 28 Aug 2026 14:07:02 +0200 Subject: [PATCH 32/41] chore: remove debug test files --- crates/examples3d/Cargo.toml | 25 --- crates/examples3d/test_frictionloss.rs | 183 --------------- crates/examples3d/test_frictionloss_caps.rs | 137 ------------ crates/examples3d/test_frictionloss_mjcf.rs | 134 ----------- crates/examples3d/test_frictionloss_stress.rs | 208 ------------------ crates/examples3d/test_refresh_nan.rs | 193 ---------------- 6 files changed, 880 deletions(-) delete mode 100644 crates/examples3d/test_frictionloss.rs delete mode 100644 crates/examples3d/test_frictionloss_caps.rs delete mode 100644 crates/examples3d/test_frictionloss_mjcf.rs delete mode 100644 crates/examples3d/test_frictionloss_stress.rs delete mode 100644 crates/examples3d/test_refresh_nan.rs diff --git a/crates/examples3d/Cargo.toml b/crates/examples3d/Cargo.toml index daacc1c7..9d60f77a 100644 --- a/crates/examples3d/Cargo.toml +++ b/crates/examples3d/Cargo.toml @@ -54,28 +54,3 @@ path = "all_examples3.rs" #[[bin]] #name = "bench_joints3" #path = "bench_joints3.rs" - -[[bin]] -name = "test_frictionloss" -path = "test_frictionloss.rs" -required-features = ["metal"] - -[[bin]] -name = "test_frictionloss_stress" -path = "test_frictionloss_stress.rs" -required-features = ["metal"] - -[[bin]] -name = "test_frictionloss_caps" -path = "test_frictionloss_caps.rs" -required-features = ["metal"] - -[[bin]] -name = "test_refresh_nan" -path = "test_refresh_nan.rs" -required-features = ["metal"] - -[[bin]] -name = "test_frictionloss_mjcf" -path = "test_frictionloss_mjcf.rs" -required-features = ["metal"] diff --git a/crates/examples3d/test_frictionloss.rs b/crates/examples3d/test_frictionloss.rs deleted file mode 100644 index af93f044..00000000 --- a/crates/examples3d/test_frictionloss.rs +++ /dev/null @@ -1,183 +0,0 @@ -//! Headless probe for MJCF-style joint dry friction (`frictionloss`). -//! -//! A single-link pendulum hinged at the origin, its rod lying along +X so -//! gravity applies a torque `m·g·l` about the hinge. MuJoCo models friction -//! loss as a constraint (a bound on the force friction may generate), not as a -//! `-f·sign(q̇)` force, so the expected behaviour is: -//! -//! * `frictionloss = 0`: the link falls freely. -//! * `frictionloss` above the gravity torque: the link sticks, exactly. A -//! force-based implementation cannot do this; it chatters around `q̇ = 0`. -//! * `frictionloss` below the gravity torque: the link falls, but slower. -//! -//! Run with `cargo run --release --bin test_frictionloss --features metal`. - -use khal::backend::{Backend, GpuBackend}; -use nexus3d::prelude::{NexusPipeline, NexusState, RbdCoupling}; -use nexus3d::rbd::dynamics::RbdSimParams; -use rapier3d::prelude::*; - -const LINK_LEN: f32 = 1.0; -const RAD: f32 = 0.05; -const STEPS: usize = 60; - -/// Builds the one-link pendulum and returns its state. `motor` adds a velocity -/// motor and a limit on the hinge, so the friction rows have to share the -/// constraint bank with them. `joint_frequency`, when given, softens the shared -/// joint constraint softness the friction rows draw their CFM from. -fn make_state(motor: bool, joint_frequency: Option) -> NexusState { - let mut state = NexusState::default(); - if let Some(hz) = joint_frequency { - let mut params = RbdSimParams::default(); - params.joint_natural_frequency = hz; - state.set_rbd_sim_params(0, params); - } - let no_coupling = RbdCoupling::None; - - let root = RigidBodyBuilder::fixed().build(); - let root_collider = ColliderBuilder::cuboid(RAD, RAD, RAD) - .collision_groups(InteractionGroups::none()) - .build(); - let root_handle = state.insert_rigid_body(root, root_collider, no_coupling); - - let link = RigidBodyBuilder::dynamic() - .translation(Vec3::new(LINK_LEN, 0.0, 0.0)) - .build(); - let collider = ColliderBuilder::cuboid(LINK_LEN * 0.5, RAD, RAD) - .collision_groups(InteractionGroups::none()) - .build(); - let link_handle = state.insert_rigid_body(link, collider, no_coupling); - - // Hinge about Z at the origin, so gravity (-Y) torques the joint. - let mut builder = RevoluteJointBuilder::new(Vec3::Z) - .local_anchor1(Vec3::ZERO) - .local_anchor2(Vec3::new(-LINK_LEN, 0.0, 0.0)); - if motor { - builder = builder - .limits([-2.0, 2.0]) - .motor_velocity(2.0, 0.0) - .motor_max_force(100.0); - } - let joint = builder.build(); - state.insert_multibody_joint(root_handle, link_handle, joint); - - state -} - -/// Steps the pendulum for `STEPS` frames and returns `(drop, |q̇|)`: how far -/// the link's centre of mass fell, and the joint velocity at the end. -async fn run_case( - backend: &GpuBackend, - frictionloss: f32, - motor: bool, - joint_frequency: Option, -) -> Result<(f32, f32), khal::backend::GpuBackendError> { - let mut state = make_state(motor, joint_frequency); - let mut pipeline = NexusPipeline::default(); - - // The frictionloss slots are reserved on the first non-zero write, which - // needs the GPU state to exist: finalize before setting it. - state.finalize(backend).await?; - let rbd = state.rbd.as_mut().expect("no rbd state"); - let ndofs = rbd.multibodies().dofs_per_batch() as usize; - rbd.set_dof_frictionloss(backend, &vec![frictionloss; ndofs]); - - for _ in 0..STEPS { - pipeline.simulate(backend, &mut state, None).await?; - } - - let rbd = state.rbd.as_ref().expect("no rbd state"); - let poses: Vec = - backend.slow_read_vec(rbd.body_poses().buffer()).await?; - let dof_state: Vec = backend - .slow_read_vec(rbd.multibodies().dof_state().buffer()) - .await?; - // Body 1 is the link; it starts level with the hinge, so any fall shows up - // as a negative y. - Ok((-poses[1].translation.y, dof_state[0].abs())) -} - -fn main() -> anyhow::Result<()> { - pollster::block_on(run()) -} - -async fn run() -> anyhow::Result<()> { - let backend = - GpuBackend::Metal(khal::backend::Metal::new().map_err(|e| anyhow::anyhow!("{e:?}"))?); - - // Gravity torque about the hinge for a rod of half-length `LINK_LEN/2`. - // The collider density is rapier's default (1.0). - let mass = 2.0 * (LINK_LEN * 0.5) * (2.0 * RAD) * (2.0 * RAD); - let gravity_torque = mass * 9.81 * LINK_LEN; - println!("mass ≈ {mass:.5} kg, gravity torque ≈ {gravity_torque:.5} N·m\n"); - - let cases = [ - ("free (fl = 0)", 0.0, false, None), - ("weak (fl = 0.25 · τ_g)", 0.25 * gravity_torque, false, None), - ("locked (fl = 4 · τ_g)", 4.0 * gravity_torque, false, None), - // A force-based `-fl·sign(q̇)` blows up here; a bounded constraint row - // simply never exceeds what it takes to stop the DoF. - ( - "extreme (fl = 1000 · τ_g)", - 1000.0 * gravity_torque, - false, - None, - ), - // Friction rows sharing the constraint bank with a limit and a motor. - ("motor + fl = 0.1 · τ_g", 0.1 * gravity_torque, true, None), - // Same locked case, but with a compliant joint softness: the friction - // row picks up CFM and the DoF creeps under load instead of sticking. - ( - "locked, soft joints (2 Hz)", - 4.0 * gravity_torque, - false, - Some(2.0), - ), - ]; - - let mut results = Vec::new(); - for (name, fl, motor, hz) in cases { - let (drop, qd) = run_case(&backend, fl, motor, hz).await?; - println!("{name:<26} drop = {drop:.6} m |q̇| = {qd:.6} rad/s"); - results.push((name, drop, qd)); - } - - let (_, free_drop, _) = results[0]; - let (_, weak_drop, _) = results[1]; - let (_, locked_drop, locked_qd) = results[2]; - let (_, extreme_drop, extreme_qd) = results[3]; - let (_, _, motor_qd) = results[4]; - let (_, soft_drop, _) = results[5]; - - println!(); - assert!( - free_drop > 0.1, - "frictionless pendulum should fall: drop = {free_drop}" - ); - assert!( - weak_drop < free_drop * 0.9, - "sub-gravity friction should slow the fall: {weak_drop} vs {free_drop}" - ); - assert!( - locked_drop.abs() < 1.0e-4 && locked_qd < 1.0e-4, - "friction above the gravity torque should hold the joint at rest: \ - drop = {locked_drop}, |q̇| = {locked_qd}" - ); - assert!( - extreme_drop.abs() < 1.0e-4 && extreme_qd < 1.0e-4, - "an oversized friction bound must stay inert, not chatter: \ - drop = {extreme_drop}, |q̇| = {extreme_qd}" - ); - assert!( - (motor_qd - 2.0).abs() < 0.1, - "the motor should still reach its 2 rad/s target through weak \ - friction: |q̇| = {motor_qd}" - ); - assert!( - soft_drop > 1.0e-3 && soft_drop < free_drop, - "a compliant joint softness should let the held DoF creep, without \ - letting it fall freely: drop = {soft_drop}" - ); - println!("OK"); - Ok(()) -} diff --git a/crates/examples3d/test_frictionloss_caps.rs b/crates/examples3d/test_frictionloss_caps.rs deleted file mode 100644 index f45e390a..00000000 --- a/crates/examples3d/test_frictionloss_caps.rs +++ /dev/null @@ -1,137 +0,0 @@ -//! Checks that the `BatchIndices` uniform agrees with the joint-constraint -//! buffer sizes after `set_dof_frictionloss` grows them. -//! -//! Reserving the dry-friction slots reallocates `joint_constraints` / -//! `joint_constraint_columns` and changes their per-batch capacities. Those -//! capacities are mirrored into the shared `BatchIndices` uniform, which the -//! kernels use to locate each multibody's slab. If the uniform is not -//! re-uploaded, every kernel indexes the resized buffers with stale strides. -//! -//! This asserts on the uniform directly rather than watching for NaN, which -//! only shows up once the resulting garbage happens to be large. -//! -//! Run with `cargo run --release --bin test_frictionloss_caps --features metal`. - -use khal::backend::{Backend, GpuBackend}; -use nexus3d::prelude::{NexusPipeline, NexusState, RbdCoupling}; -use nexus3d::rbd::shaders::utils::BatchIndices; -use rapier3d::prelude::*; - -const RAD: f32 = 0.05; -const LINK_LEN: f32 = 0.5; - -fn build_env(state: &mut NexusState, env: usize) { - for (c, num_links) in [2usize, 4, 3].iter().enumerate() { - let z = c as f32 * 4.0; - let root = RigidBodyBuilder::fixed() - .translation(Vec3::new(0.0, 0.0, z)) - .build(); - let rc = ColliderBuilder::cuboid(RAD, RAD, RAD) - .collision_groups(InteractionGroups::none()) - .build(); - let mut parent = state.insert_rigid_body_in(env, root, rc, RbdCoupling::None); - for i in 0..*num_links { - let body = RigidBodyBuilder::dynamic() - .translation(Vec3::new((i as f32 + 1.0) * LINK_LEN * 2.0, 0.0, z)) - .build(); - let collider = ColliderBuilder::cuboid(LINK_LEN, RAD, RAD) - .collision_groups(InteractionGroups::none()) - .build(); - let handle = state.insert_rigid_body_in(env, body, collider, RbdCoupling::None); - let anchor = if i == 0 { - Vec3::ZERO - } else { - Vec3::new(LINK_LEN, 0.0, 0.0) - }; - let joint = RevoluteJointBuilder::new(Vec3::Z) - .local_anchor1(anchor) - .local_anchor2(Vec3::new(-LINK_LEN, 0.0, 0.0)) - .limits([-1.5, 1.5]) - .motor_velocity(0.5, 0.1) - .motor_max_force(50.0); - state.insert_multibody_joint_in(env, parent, handle, joint.build()); - parent = handle; - } - } -} - -/// `(uniform capacity, actual capacity)` for the joint-constraint bank and its -/// column buffer, after `steps` pipeline steps. -async fn caps( - backend: &GpuBackend, - via_multibody_set: bool, - steps: usize, -) -> Result<((u32, u32), (u32, u32)), khal::backend::GpuBackendError> { - let mut state = NexusState::default(); - build_env(&mut state, 0); - for _ in 1..2 { - let env = state.add_environment(); - build_env(&mut state, env); - } - let mut pipeline = NexusPipeline::default(); - state.finalize(backend).await?; - - let rbd = state.rbd.as_mut().expect("no rbd state"); - let n = rbd.multibodies().dofs_per_batch() as usize * rbd.multibodies().num_batches() as usize; - let values = vec![0.1f32; n]; - if via_multibody_set { - rbd.multibodies_mut().set_dof_frictionloss(backend, &values); - } else { - rbd.set_dof_frictionloss(backend, &values); - } - - for _ in 0..steps { - pipeline.simulate(backend, &mut state, None).await?; - } - - let rbd = state.rbd.as_ref().expect("no rbd state"); - let bi: Vec = backend.slow_read_vec(rbd.batch_indices().buffer()).await?; - let bi = bi[0]; - Ok(( - ( - bi.mb_joint_constraints_batch_capacity, - rbd.multibodies().joint_constraints_per_batch(), - ), - ( - bi.mb_joint_constraint_columns_batch_capacity, - rbd.multibodies().joint_constraint_columns_per_batch(), - ), - )) -} - -fn main() -> anyhow::Result<()> { - pollster::block_on(run()) -} - -async fn run() -> anyhow::Result<()> { - let backend = - GpuBackend::Metal(khal::backend::Metal::new().map_err(|e| anyhow::anyhow!("{e:?}"))?); - - let mut bad = 0; - for via_set in [false, true] { - for steps in [0usize, 1] { - let (rows, cols) = caps(&backend, via_set, steps).await?; - let api = if via_set { "mb_set" } else { "rbd " }; - let ok = rows.0 == rows.1 && cols.0 == cols.1; - println!( - "{api}, after {steps} step(s): rows uniform/actual = {}/{}, cols = {}/{} {}", - rows.0, - rows.1, - cols.0, - cols.1, - if ok { "ok" } else { "MISMATCH" } - ); - // Before any step has run, the `mb_set` entry point is expected to - // carry a stale uniform: the step is what re-uploads it. - if !ok && steps > 0 { - bad += 1; - } - } - } - - if bad > 0 { - anyhow::bail!("{bad} case(s) still stale after a step"); - } - println!("\nOK"); - Ok(()) -} diff --git a/crates/examples3d/test_frictionloss_mjcf.rs b/crates/examples3d/test_frictionloss_mjcf.rs deleted file mode 100644 index ba389d31..00000000 --- a/crates/examples3d/test_frictionloss_mjcf.rs +++ /dev/null @@ -1,134 +0,0 @@ -//! End-to-end check that MJCF `` reaches the GPU solver. -//! -//! The path is: `` → rapier's `Multibody::frictions` (via -//! `rapier3d-mjcf`'s `add_frictionloss_to_multibody`) → the per-DoF friction -//! section of `dof_state` at build time → one `MB_JOINT_KIND_FRICTION` row per -//! non-zero DoF in `gpu_mb_init_joint_constraints`. -//! -//! Nothing calls `set_dof_frictionloss` here: the whole point is that loading -//! a model is enough. -//! -//! Run with `cargo run --release --bin test_frictionloss_mjcf --features metal`. - -use khal::backend::{Backend, GpuBackend}; -use nexus3d::prelude::{NexusPipeline, NexusState, RbdCoupling}; -use rapier3d::prelude::*; -use rapier3d_mjcf::{MjcfLoaderOptions, MjcfMultibodyOptions, MjcfRobot}; - -const STEPS: usize = 15; - -/// A hinge at the origin with a 1 m arm along +X, so gravity torques it. -/// `frictionloss` is substituted in. -fn model(frictionloss: f32) -> String { - format!( - r#" - - -"# - ) -} - -/// Loads the model into a `NexusState`, steps it, and returns how far the arm -/// fell plus the joint speed at the end. -async fn run_case( - backend: &GpuBackend, - frictionloss: f32, -) -> Result<(f32, f32), khal::backend::GpuBackendError> { - let xml = model(frictionloss); - let (robot, _) = MjcfRobot::from_str(&xml, MjcfLoaderOptions::default(), ".").unwrap(); - - let mut state = NexusState::default(); - { - let world = state.rbd_world_mut(0); - robot.insert_using_multibody_joints( - &mut world.bodies, - &mut world.colliders, - &mut world.multibody_joints, - &mut world.impulse_joints, - MjcfMultibodyOptions::empty(), - ); - } - - let mut pipeline = NexusPipeline::default(); - state.finalize(backend).await?; - - // Confirm the value survived the loader before trusting the simulation. - let loaded = state - .rbd_world(0) - .multibody_joints - .multibodies() - .flat_map(|mb| mb.frictions().iter().copied().collect::>()) - .fold(0.0f32, f32::max); - assert!( - (loaded - frictionloss).abs() < 1.0e-6, - "rapier's Multibody::frictions should carry the MJCF value: \ - got {loaded}, expected {frictionloss}" - ); - - for _ in 0..STEPS { - pipeline.simulate(backend, &mut state, None).await?; - } - - let rbd = state.rbd.as_ref().expect("no rbd state"); - let poses: Vec = - backend.slow_read_vec(rbd.body_poses().buffer()).await?; - let dof: Vec = backend - .slow_read_vec(rbd.multibodies().dof_state().buffer()) - .await?; - let drop = poses - .iter() - .map(|p| -p.translation.y) - .fold(0.0f32, f32::max); - Ok((drop, dof[0].abs())) -} - -fn main() -> anyhow::Result<()> { - pollster::block_on(run()) -} - -async fn run() -> anyhow::Result<()> { - let backend = - GpuBackend::Metal(khal::backend::Metal::new().map_err(|e| anyhow::anyhow!("{e:?}"))?); - - // Gravity torque about the hinge for a 1 kg arm with its centre 1 m out. - let gravity_torque = 1.0 * 9.81 * 1.0; - let cases = [ - ("frictionloss = 0", 0.0), - ("frictionloss = 0.25 · τ_g", 0.25 * gravity_torque), - ("frictionloss = 4 · τ_g", 4.0 * gravity_torque), - ]; - - let mut results = Vec::new(); - for (name, fl) in cases { - let (drop, speed) = run_case(&backend, fl).await?; - println!("{name:<28} drop = {drop:.6} m |q̇| = {speed:.6} rad/s"); - results.push((drop, speed)); - } - - let (free, _) = results[0]; - let (weak, _) = results[1]; - let (locked, locked_speed) = results[2]; - - println!(); - assert!(free > 0.02, "the arm should fall with no friction: {free}"); - assert!( - weak < free * 0.95, - "sub-gravity friction should slow the fall: {weak} vs {free}" - ); - assert!( - locked.abs() < 1.0e-4 && locked_speed < 1.0e-4, - "friction above the gravity torque should hold the joint: \ - drop = {locked}, |q̇| = {locked_speed}" - ); - println!("OK"); - Ok(()) -} diff --git a/crates/examples3d/test_frictionloss_stress.rs b/crates/examples3d/test_frictionloss_stress.rs deleted file mode 100644 index 823b2ab5..00000000 --- a/crates/examples3d/test_frictionloss_stress.rs +++ /dev/null @@ -1,208 +0,0 @@ -//! Stress probe for the `frictionloss` constraint slot reservation. -//! -//! `reserve_frictionloss_slots` recomputes every multibody's `first_constraint` -//! offset and grows the joint-constraint bank. The single-pendulum probe in -//! `test_frictionloss.rs` never exercises that: it has one multibody in one -//! batch, so every offset is zero. This one builds several multibodies of -//! differing DoF counts across several batches, with limits and motors mixed -//! in, and checks that nothing goes non-finite. -//! -//! Run with `cargo run --release --bin test_frictionloss_stress --features metal`. - -use khal::backend::{Backend, GpuBackend}; -use nexus3d::prelude::{NexusPipeline, NexusState, RbdCoupling}; -use rapier3d::prelude::*; - -const RAD: f32 = 0.05; -const LINK_LEN: f32 = 0.5; -const STEPS: usize = 90; - -/// Chain lengths, in links. Differing DoF counts are the point: they make each -/// multibody's `first_constraint` offset distinct. -const CHAINS: [usize; 3] = [2, 4, 3]; - -/// Builds `CHAINS.len()` pendulum chains in environment `env`, offset along Z -/// so they don't overlap. Chain `c` gets limits and motors on its first joint -/// only, so limit/motor rows and friction rows share the bank unevenly. -fn build_env(state: &mut NexusState, env: usize, contacts: bool) { - if contacts { - let ground = RigidBodyBuilder::fixed() - .translation(Vec3::new(0.0, -3.0, 0.0)) - .build(); - let ground_collider = ColliderBuilder::cuboid(40.0, 0.5, 40.0).build(); - state.insert_rigid_body_in(env, ground, ground_collider, RbdCoupling::None); - } - let groups = if contacts { - InteractionGroups::all() - } else { - InteractionGroups::none() - }; - for (c, num_links) in CHAINS.iter().enumerate() { - let z = c as f32 * 4.0; - let root = RigidBodyBuilder::fixed() - .translation(Vec3::new(0.0, 0.0, z)) - .build(); - let root_collider = ColliderBuilder::cuboid(RAD, RAD, RAD) - .collision_groups(groups) - .build(); - let mut parent = state.insert_rigid_body_in(env, root, root_collider, RbdCoupling::None); - - for i in 0..*num_links { - let x = (i as f32 + 1.0) * LINK_LEN * 2.0; - let body = RigidBodyBuilder::dynamic() - .translation(Vec3::new(x, 0.0, z)) - .build(); - let collider = ColliderBuilder::cuboid(LINK_LEN, RAD, RAD) - .collision_groups(groups) - .build(); - let handle = state.insert_rigid_body_in(env, body, collider, RbdCoupling::None); - - let parent_anchor = if i == 0 { - Vec3::ZERO - } else { - Vec3::new(LINK_LEN, 0.0, 0.0) - }; - let mut builder = RevoluteJointBuilder::new(Vec3::Z) - .local_anchor1(parent_anchor) - .local_anchor2(Vec3::new(-LINK_LEN, 0.0, 0.0)); - if i == 0 { - builder = builder - .limits([-1.5, 1.5]) - .motor_velocity(0.5, 0.1) - .motor_max_force(50.0); - } - state.insert_multibody_joint_in(env, parent, handle, builder.build()); - parent = handle; - } - } -} - -#[allow(clippy::type_complexity)] -async fn run_case( - backend: &GpuBackend, - num_envs: usize, - frictionloss: f32, - // `false` turns off the implicit-Coriolis path, which is what makes the - // solver take the per-substep `gpu_mb_refresh_joint_constraints` branch - // instead of a full rebuild each substep. The friction rows are only - // touched by that kernel here. - implicit_coriolis: bool, - contacts: bool, - // `true` reproduces the zealot call pattern: the frictionloss is set - // through `GpuMultibodySet` directly rather than through `RbdState`, so - // nothing rebuilds `BatchIndices` at the call site. - via_multibody_set: bool, -) -> Result<(usize, f32), khal::backend::GpuBackendError> { - let mut state = NexusState::default(); - build_env(&mut state, 0, contacts); - for _ in 1..num_envs { - let env = state.add_environment(); - build_env(&mut state, env, contacts); - } - - let mut pipeline = NexusPipeline::default(); - state.finalize(backend).await?; - - let rbd = state.rbd.as_mut().expect("no rbd state"); - let per_batch = rbd.multibodies().dofs_per_batch() as usize; - let batches = rbd.multibodies().num_batches() as usize; - if frictionloss > 0.0 { - let values = vec![frictionloss; per_batch * batches]; - if via_multibody_set { - rbd.multibodies_mut().set_dof_frictionloss(backend, &values); - } else { - rbd.set_dof_frictionloss(backend, &values); - } - } - if !implicit_coriolis { - rbd.multibodies_mut().set_substep_refresh(false); - rbd.multibodies_mut().set_substep_refresh_light(false); - rbd.set_implicit_coriolis(backend, false); - } - - for _ in 0..STEPS { - pipeline.simulate(backend, &mut state, None).await?; - } - - let rbd = state.rbd.as_ref().expect("no rbd state"); - let poses: Vec = - backend.slow_read_vec(rbd.body_poses().buffer()).await?; - let dof_state: Vec = backend - .slow_read_vec(rbd.multibodies().dof_state().buffer()) - .await?; - - let bad_poses = poses - .iter() - .filter(|p| { - !p.translation.x.is_finite() - || !p.translation.y.is_finite() - || !p.translation.z.is_finite() - }) - .count(); - // Only the velocity section is integrated state; the rest are parameters. - let bad_vels = dof_state[..per_batch * batches] - .iter() - .filter(|v| !v.is_finite()) - .count(); - let max_speed = dof_state[..per_batch * batches] - .iter() - .filter(|v| v.is_finite()) - .fold(0.0f32, |a, v| a.max(v.abs())); - - Ok((bad_poses + bad_vels, max_speed)) -} - -fn main() -> anyhow::Result<()> { - pollster::block_on(run()) -} - -async fn run() -> anyhow::Result<()> { - let backend = - GpuBackend::Metal(khal::backend::Metal::new().map_err(|e| anyhow::anyhow!("{e:?}"))?); - - let mut failures = 0; - let mut known = 0; - // `implicit = false` switches the multibody to explicit Coriolis forces, - // which diverges on free-swinging chains of four or more links: the - // velocities grow smoothly and exponentially until they overflow. It does - // so with *or without* frictionloss, and identically on the pre-frictionloss - // tree, so it is reported but not counted. See `test_refresh_nan` for the - // isolated repro; it is the documented stability cost of that path, not a - // regression here. - for via_set in [false, true] { - for contacts in [false, true] { - for implicit in [true, false] { - for num_envs in [1usize, 2, 4] { - // The last two values are deliberately absurd: the impulse bound is - // `frictionloss · dt`, so a large enough loss can overflow the - // accumulated impulse. - for fl in [0.0f32, 0.05, 5.0, 1.0e12, 1.0e30] { - let (bad, max_speed) = - run_case(&backend, num_envs, fl, implicit, contacts, via_set).await?; - let c = if contacts { "contacts" } else { "free " }; - let path = if implicit { "rebuild" } else { "refresh" }; - let api = if via_set { "mb_set " } else { "rbd " }; - let tag = format!("{api}, {path}, {c}, envs = {num_envs}, fl = {fl:e}"); - if bad > 0 && !implicit { - println!( - "{tag:<43} KNOWN (explicit-Coriolis divergence): {bad} non-finite" - ); - known += 1; - } else if bad > 0 { - println!("{tag:<43} FAIL: {bad} non-finite values"); - failures += 1; - } else { - println!("{tag:<43} ok, max |q̇| = {max_speed:.4}"); - } - } - } - } - } - } - - if failures > 0 { - anyhow::bail!("{failures} configuration(s) produced non-finite state"); - } - println!("\nOK ({known} known explicit-Coriolis divergences ignored)"); - Ok(()) -} diff --git a/crates/examples3d/test_refresh_nan.rs b/crates/examples3d/test_refresh_nan.rs deleted file mode 100644 index 906bd672..00000000 --- a/crates/examples3d/test_refresh_nan.rs +++ /dev/null @@ -1,193 +0,0 @@ -//! Repro for the non-finite state seen with `implicit_coriolis = false`. -//! -//! This is *numerical divergence*, not corruption, and not a constraint bug: -//! it reproduces with no limits and no motors anywhere (`rows = none`), where -//! `gpu_mb_refresh_joint_constraints` is never even dispatched, and the -//! velocity trace grows smoothly and exponentially until it overflows -//! (`TRACE=1` prints it). It matches the tradeoff the solver already -//! documents: the explicit-Coriolis path is cheaper but less stable. -//! -//! Frictionloss is not involved: this reproduces at `frictionloss = 0`, and -//! identically on the tree from before joint friction became a constraint. -//! -//! The sweep isolates the variable: the same scenes are run with implicit -//! Coriolis on and off, reporting the first step at which any DoF velocity or -//! body pose stops being finite. Chains of four or more links diverge with it -//! off and are stable with it on. -//! -//! Run with `cargo run --release --bin test_refresh_nan --features metal`, -//! and `TRACE=1 ...` to see the per-step velocity growth. - -use khal::backend::{Backend, GpuBackend}; -use nexus3d::prelude::{NexusPipeline, NexusState, RbdCoupling}; -use rapier3d::prelude::*; - -const RAD: f32 = 0.05; -const LINK_LEN: f32 = 0.5; -const MAX_STEPS: usize = 100; - -#[derive(Clone, Copy)] -struct Case { - num_chains: usize, - num_links: usize, - limits: bool, - motor: bool, - /// Put the limit / motor rows on the chain's first joint only, leaving the - /// rest of its DoFs with no constraint slot at all. - first_joint_only: bool, -} - -fn build(state: &mut NexusState, env: usize, case: Case) { - for c in 0..case.num_chains { - let z = c as f32 * 4.0; - let root = RigidBodyBuilder::fixed() - .translation(Vec3::new(0.0, 0.0, z)) - .build(); - let rc = ColliderBuilder::cuboid(RAD, RAD, RAD) - .collision_groups(InteractionGroups::none()) - .build(); - let mut parent = state.insert_rigid_body_in(env, root, rc, RbdCoupling::None); - - // Chain `c` gets one extra link, so the multibodies have distinct DoF - // counts and therefore distinct constraint-slab offsets. - for i in 0..(case.num_links + c) { - let body = RigidBodyBuilder::dynamic() - .translation(Vec3::new((i as f32 + 1.0) * LINK_LEN * 2.0, 0.0, z)) - .build(); - let collider = ColliderBuilder::cuboid(LINK_LEN, RAD, RAD) - .collision_groups(InteractionGroups::none()) - .build(); - let handle = state.insert_rigid_body_in(env, body, collider, RbdCoupling::None); - - let anchor = if i == 0 { - Vec3::ZERO - } else { - Vec3::new(LINK_LEN, 0.0, 0.0) - }; - let mut j = RevoluteJointBuilder::new(Vec3::Z) - .local_anchor1(anchor) - .local_anchor2(Vec3::new(-LINK_LEN, 0.0, 0.0)); - let rows_here = !case.first_joint_only || i == 0; - if case.limits && rows_here { - j = j.limits([-1.5, 1.5]); - } - if case.motor && rows_here { - j = j.motor_velocity(0.5, 0.1).motor_max_force(50.0); - } - state.insert_multibody_joint_in(env, parent, handle, j.build()); - parent = handle; - } - } -} - -/// Runs `case` on the refresh path and returns the first step index at which -/// state goes non-finite, or `None` if it stays finite for `MAX_STEPS`. -async fn first_bad_step( - backend: &GpuBackend, - case: Case, - implicit_coriolis: bool, -) -> Result, khal::backend::GpuBackendError> { - let mut state = NexusState::default(); - build(&mut state, 0, case); - let mut pipeline = NexusPipeline::default(); - state.finalize(backend).await?; - - if !implicit_coriolis { - let rbd = state.rbd.as_mut().expect("no rbd state"); - rbd.multibodies_mut().set_substep_refresh(false); - rbd.multibodies_mut().set_substep_refresh_light(false); - rbd.set_implicit_coriolis(backend, false); - } - - for step in 0..MAX_STEPS { - pipeline.simulate(backend, &mut state, None).await?; - let rbd = state.rbd.as_ref().expect("no rbd state"); - let n = - rbd.multibodies().dofs_per_batch() as usize * rbd.multibodies().num_batches() as usize; - let dof: Vec = backend - .slow_read_vec(rbd.multibodies().dof_state().buffer()) - .await?; - let poses: Vec = - backend.slow_read_vec(rbd.body_poses().buffer()).await?; - let bad = dof[..n].iter().any(|v| !v.is_finite()) - || poses.iter().any(|p| { - !p.translation.x.is_finite() - || !p.translation.y.is_finite() - || !p.translation.z.is_finite() - }); - if std::env::var("TRACE").is_ok() { - let m = dof[..n] - .iter() - .filter(|v| v.is_finite()) - .fold(0.0f32, |a, v| a.max(v.abs())); - if step % 5 == 0 || bad { - println!(" step {step:>3}: max |q̇| = {m:e}"); - } - } - if bad { - return Ok(Some(step)); - } - } - Ok(None) -} - -fn main() -> anyhow::Result<()> { - pollster::block_on(run()) -} - -async fn run() -> anyhow::Result<()> { - let backend = - GpuBackend::Metal(khal::backend::Metal::new().map_err(|e| anyhow::anyhow!("{e:?}"))?); - - let mut any = false; - for implicit in [false, true] { - for first_joint_only in [true] { - for num_chains in [1usize] { - for num_links in [2usize, 4, 6] { - for (limits, motor) in - [(false, false), (true, false), (false, true), (true, true)] - { - let case = Case { - num_chains, - num_links, - limits, - motor, - first_joint_only, - }; - let rows = match (limits, motor) { - (false, false) => "none ", - (true, false) => "limit ", - (false, true) => "motor ", - (true, true) => "limit+motor ", - }; - let ic = if implicit { - "coriolis=implicit" - } else { - "coriolis=explicit" - }; - match first_bad_step(&backend, case, implicit).await? { - Some(step) => { - println!( - "{ic}, links = {num_links}, rows = {rows} DIVERGED at step {step}" - ); - any = true; - } - None => println!( - "{ic}, links = {num_links}, rows = {rows} finite for {MAX_STEPS} steps" - ), - } - } - } - } - } - } - - if any { - println!( - "\nReproduced (explicit-Coriolis divergence; run with TRACE=1 to see the growth)." - ); - } else { - println!("\nNo failure in this sweep."); - } - Ok(()) -} From 96c626676c895bf62001e332abaa0256414c35cf Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Fri, 28 Aug 2026 15:42:49 +0200 Subject: [PATCH 33/41] chore: cleanups --- src_rbd/dynamics/multibody/env_reset.rs | 3 +-- .../multibody/multibody_from_rapier.rs | 8 ++----- src_rbd/dynamics/multibody/multibody_set.rs | 22 ++----------------- src_rbd/pipeline/rbd_state.rs | 11 ++-------- src_rbd_shaders/broad_phase/narrow_phase.rs | 2 -- .../dynamics/multibody/env_reset.rs | 6 ----- .../dynamics/multibody/joint_constraints.rs | 15 +------------ src_rbd_shaders/dynamics/multibody/mod.rs | 2 -- src_rbd_shaders/dynamics/multibody/types.rs | 8 +------ src_rbd_shaders/queries/polygonal_feature.rs | 2 -- 10 files changed, 9 insertions(+), 70 deletions(-) diff --git a/src_rbd/dynamics/multibody/env_reset.rs b/src_rbd/dynamics/multibody/env_reset.rs index d9863de9..55d78c77 100644 --- a/src_rbd/dynamics/multibody/env_reset.rs +++ b/src_rbd/dynamics/multibody/env_reset.rs @@ -203,8 +203,7 @@ impl GpuMultibodySet { } } - /// Resets env `dst_env` from a CPU snapshot: one staging upload and one - /// scatter dispatch, with no readback and no per-element strided writes. + /// Resets env `dst_env` from a CPU snapshot. pub fn reset_env_from_snapshot( &mut self, backend: &GpuBackend, diff --git a/src_rbd/dynamics/multibody/multibody_from_rapier.rs b/src_rbd/dynamics/multibody/multibody_from_rapier.rs index c43c027d..f80ab405 100644 --- a/src_rbd/dynamics/multibody/multibody_from_rapier.rs +++ b/src_rbd/dynamics/multibody/multibody_from_rapier.rs @@ -64,11 +64,7 @@ impl GpuMultibodySet { let mut global_max_cons = 0u32; let mut global_max_couplings = 0u32; - // Whether any multibody anywhere declares dry joint friction. The - // constraint-slot reservation below is all-or-nothing, matching - // `GpuMultibodySet::reserve_frictionloss_slots`: reserving for only the - // multibodies that currently have friction would leave a later - // `set_dof_frictionloss` on a different one with nowhere to emit. + // 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)) @@ -200,7 +196,7 @@ 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 (MJCF ``). + // Per-DoF dry joint friction. let mb_friction = mb.frictions(); let mut rapier_assembly = 0usize; let mb_statics_start = statics.len(); diff --git a/src_rbd/dynamics/multibody/multibody_set.rs b/src_rbd/dynamics/multibody/multibody_set.rs index 5db79745..cd90f43f 100644 --- a/src_rbd/dynamics/multibody/multibody_set.rs +++ b/src_rbd/dynamics/multibody/multibody_set.rs @@ -818,30 +818,12 @@ impl GpuMultibodySet { /// 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) and is transposed here into the - /// batch-interleaved layout the kernels index. - /// - /// A post-build override for callers whose rapier scenes carry no armature, - /// or that randomize it per environment; the build path already seeds this - /// section from `mb.armature()`. + /// 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 section of - /// [`Self::dof_state`] (MJCF `frictionloss`, N·m). Zero, the default, - /// disables it. `values` uses the same env-major layout as - /// [`Self::set_dof_armature`]. - /// - /// Friction loss is a constraint, not a force: each DoF with a non-zero - /// loss gets a solver row driving its velocity to zero, with the impulse - /// bounded by `frictionloss · dt` (a load-independent bound, which is why - /// MuJoCo distinguishes it from Coulomb friction). The first non-zero call - /// reserves the extra constraint slots, which changes per-batch capacities - /// and so invalidates the shared `BatchIndices` uniform; the next - /// `RbdPipeline` step re-uploads it. - /// [`RbdState::set_dof_frictionloss`](crate::pipeline::RbdState::set_dof_frictionloss) - /// does it up front instead, if you would rather not carry a dirty uniform. + /// 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); diff --git a/src_rbd/pipeline/rbd_state.rs b/src_rbd/pipeline/rbd_state.rs index e214d761..dfb2819b 100644 --- a/src_rbd/pipeline/rbd_state.rs +++ b/src_rbd/pipeline/rbd_state.rs @@ -432,12 +432,7 @@ impl RbdState { self.rebuild_batch_indices(backend); } - /// Sets the per-DoF dry joint friction (MJCF `frictionloss`, N·m), in the - /// env-major `dofs_per_batch * num_batches` layout described on - /// [`GpuMultibodySet::set_dof_frictionloss`](crate::dynamics::GpuMultibodySet::set_dof_frictionloss). - /// - /// The first non-zero call grows the joint-constraint bank, so the shared - /// `BatchIndices` uniform is rebuilt here. + /// 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); @@ -675,9 +670,7 @@ impl RbdState { } } - /// Resets env `dst_env` from a CPU snapshot using `write_buffer` only, with - /// no GPU to CPU readback. This is what removes the handful of per-reset - /// sync stalls that otherwise dominate reset cost on the WebGPU backend. + /// Resets env `dst_env` from a CPU snapshot using `write_buffer` only. pub fn reset_env_from_snapshot( &mut self, backend: &GpuBackend, diff --git a/src_rbd_shaders/broad_phase/narrow_phase.rs b/src_rbd_shaders/broad_phase/narrow_phase.rs index 1a14216b..8213d85b 100644 --- a/src_rbd_shaders/broad_phase/narrow_phase.rs +++ b/src_rbd_shaders/broad_phase/narrow_phase.rs @@ -219,8 +219,6 @@ pub fn gpu_reduce_contacts( } } -const PREDICTION: f32 = 2.0e-3; // TODO: make the prediction configurable. - /// Narrow phase, pass 1 of 2: analytic shape-shape contacts for ball / cuboid /// pairs, written straight into the `contacts` buffer. /// diff --git a/src_rbd_shaders/dynamics/multibody/env_reset.rs b/src_rbd_shaders/dynamics/multibody/env_reset.rs index 36532767..d6b7246f 100644 --- a/src_rbd_shaders/dynamics/multibody/env_reset.rs +++ b/src_rbd_shaders/dynamics/multibody/env_reset.rs @@ -3,12 +3,6 @@ //! 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. -//! That is one staging upload plus one dispatch per reset, instead of the -//! hundreds of strided `write_buffer`s the interleaved layout would otherwise -//! force: an env's data is strided across the whole buffer (element `intra` of -//! batch `b` lives at `intra · num_batches + b`, workspace quads at -//! `(link · WS_QUADS + q) · num_batches + b`). The staging blob is exactly the -//! `num_batches = 1` interleaving, so its source indices are the flat `0..len`. use glamx::{UVec4, Vec4}; use khal_std::glamx::UVec3; diff --git a/src_rbd_shaders/dynamics/multibody/joint_constraints.rs b/src_rbd_shaders/dynamics/multibody/joint_constraints.rs index 45611b67..fd9d1e59 100644 --- a/src_rbd_shaders/dynamics/multibody/joint_constraints.rs +++ b/src_rbd_shaders/dynamics/multibody/joint_constraints.rs @@ -437,13 +437,7 @@ fn build_coupling_constraint( /// 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)`), the same compliance the limit rows use; -/// it is MuJoCo's `solreffriction` knob. Only CFM applies here, never ERP: -/// there is no position error for a bias to chase. With the default -/// near-rigid joint softness the coefficient is ~0 and a DoF whose driving -/// force stays under the bound sticks exactly; softening it lets the DoF -/// creep under load instead. +/// `cfm_coeff` is the shared joint softness (rapier's `joint.softness.cfm_coeff(dt)`). #[inline] fn build_friction_constraint( dof_id: u32, @@ -682,13 +676,6 @@ pub fn gpu_mb_init_joint_constraints( /// Per-substep refresh of the joint limit / motor slots, the cheap alternative /// to a full rebuild. /// -/// When the constraint columns and `inv_lhs` are per-step constants (no -/// implicit Coriolis, no per-substep mass-matrix refresh), the only things that -/// change between substeps are the rhs, the limit activity and the accumulated -/// impulse. This recomputes exactly those from the slot's stashed `(link, -/// axis)`, so the full emission walk and the LU back-solves run once per step -/// instead of once per substep. -/// /// One 64-lane workgroup per (multibody, batch); lanes stride the slots. #[spirv_bindgen] #[spirv(compute(threads(64)))] diff --git a/src_rbd_shaders/dynamics/multibody/mod.rs b/src_rbd_shaders/dynamics/multibody/mod.rs index 35690962..4f63ca57 100644 --- a/src_rbd_shaders/dynamics/multibody/mod.rs +++ b/src_rbd_shaders/dynamics/multibody/mod.rs @@ -16,8 +16,6 @@ mod compute_dynamics_pre; mod contact_constraints; mod contact_sensor; -// The RL env-reset primitives (terrain teleport offsets, template blobs) -// are 3D-only, like the RbdSnapshot host API on top of them. #[cfg(feature = "dim3")] mod env_reset; mod gravity_and_lu; diff --git a/src_rbd_shaders/dynamics/multibody/types.rs b/src_rbd_shaders/dynamics/multibody/types.rs index 8bf23fb6..4dc0ff1e 100644 --- a/src_rbd_shaders/dynamics/multibody/types.rs +++ b/src_rbd_shaders/dynamics/multibody/types.rs @@ -62,13 +62,7 @@ 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 (MJCF `frictionloss`). A row -/// with jacobian `J = e_{dof_id}`, zero target velocity and no position -/// residual, whose impulse is clamped to `±frictionloss·dt` and which carries -/// the shared joint CFM softness. MuJoCo models friction loss this way rather -/// than as a `-f·sign(q̇)` force: the bound is load-independent (not Coulomb -/// friction), and only a constraint can hold a DoF at rest instead of -/// chattering around zero velocity. +/// Joint-constraint `kind`: dry joint friction. pub const MB_JOINT_KIND_FRICTION: u32 = 5; /// Sentinel marking a link with no parent (the root). diff --git a/src_rbd_shaders/queries/polygonal_feature.rs b/src_rbd_shaders/queries/polygonal_feature.rs index 4d41fefd..c1f23dbe 100644 --- a/src_rbd_shaders/queries/polygonal_feature.rs +++ b/src_rbd_shaders/queries/polygonal_feature.rs @@ -288,8 +288,6 @@ mod dim2 { // 3D Implementation // ==================== -/// Re-export of the ≤8 → ≤4 keep-deepest-then-spread contact selector for the -/// optional contact-reduction pass (see `gpu_reduce_contacts`). #[cfg(feature = "dim3")] pub use dim3::manifold_reduction; From 81b17623b0ccaa2d0a428bc59c2694f8f8d5ce8a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Fri, 28 Aug 2026 15:45:12 +0200 Subject: [PATCH 34/41] fix(rbd): build the bench harness without the metal feature and only in 3D --- src_rbd/pipeline/bench_narrow_phase.rs | 18 +++++++++++++++--- src_rbd/pipeline/mod.rs | 2 +- 2 files changed, 16 insertions(+), 4 deletions(-) diff --git a/src_rbd/pipeline/bench_narrow_phase.rs b/src_rbd/pipeline/bench_narrow_phase.rs index 3765d5cd..8b58f2a8 100644 --- a/src_rbd/pipeline/bench_narrow_phase.rs +++ b/src_rbd/pipeline/bench_narrow_phase.rs @@ -9,6 +9,20 @@ 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(); @@ -27,9 +41,7 @@ fn build_env(num_boxes: usize) -> (RigidBodySet, ColliderSet) { } async fn run_bench(num_envs: u32, num_boxes: usize, num_steps: u32) { - // Metal, not WebGPU: `gpu_mb_init_joint_constraints` currently binds 10 - // storage buffers, over WebGPU's per-stage limit of 8. - let backend = GpuBackend::Metal(khal::backend::metal::Metal::new().unwrap()); + let backend = bench_backend().await; let envs: Vec<_> = (0..num_envs).map(|_| build_env(num_boxes)).collect(); let joints = ImpulseJointSet::new(); diff --git a/src_rbd/pipeline/mod.rs b/src_rbd/pipeline/mod.rs index 352f617c..2b80368d 100644 --- a/src_rbd/pipeline/mod.rs +++ b/src_rbd/pipeline/mod.rs @@ -4,7 +4,7 @@ //! simulation step on the GPU. The pipeline manages collision detection, contact generation, //! constraint solving, and integration. -#[cfg(test)] +#[cfg(all(test, feature = "dim3"))] mod bench_narrow_phase; mod insertion_removal; mod lbvh_validation; From 497d367fba7079e2dd650f8806264e3ba0f31027 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Fri, 28 Aug 2026 15:47:36 +0200 Subject: [PATCH 35/41] fix(rbd): silence the clippy needless-borrow and unnecessary-mut lints --- src_rbd/dynamics/multibody/env_reset.rs | 12 ++++-------- src_rbd/dynamics/multibody/multibody_from_rapier.rs | 2 +- src_rbd/dynamics/multibody/multibody_solver.rs | 4 ++-- 3 files changed, 7 insertions(+), 11 deletions(-) diff --git a/src_rbd/dynamics/multibody/env_reset.rs b/src_rbd/dynamics/multibody/env_reset.rs index 55d78c77..2d0eaa79 100644 --- a/src_rbd/dynamics/multibody/env_reset.rs +++ b/src_rbd/dynamics/multibody/env_reset.rs @@ -126,22 +126,18 @@ impl EnvResetBundle { shader: EnvResetShader::from_backend(backend).unwrap(), staging_ws: Tensor::vector( backend, - &vec![Vec4::ZERO; (lpb * WS_QUADS).max(1) as usize], + 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], + 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(), } } diff --git a/src_rbd/dynamics/multibody/multibody_from_rapier.rs b/src_rbd/dynamics/multibody/multibody_from_rapier.rs index f80ab405..ae9a4141 100644 --- a/src_rbd/dynamics/multibody/multibody_from_rapier.rs +++ b/src_rbd/dynamics/multibody/multibody_from_rapier.rs @@ -632,7 +632,7 @@ impl GpuMultibodySet { delay_update_cache: None, contact_sensor_links: Tensor::vector( backend, - &[u32::MAX; crate::shaders::dynamics::MAX_CONTACT_SENSORS as usize], + [u32::MAX; crate::shaders::dynamics::MAX_CONTACT_SENSORS as usize], storage, ) .unwrap(), diff --git a/src_rbd/dynamics/multibody/multibody_solver.rs b/src_rbd/dynamics/multibody/multibody_solver.rs index 894295af..87fcbd76 100644 --- a/src_rbd/dynamics/multibody/multibody_solver.rs +++ b/src_rbd/dynamics/multibody/multibody_solver.rs @@ -256,7 +256,7 @@ impl GpuMultibodySolver { &mb.links_static, &mb.links_workspace, &mut mb.joint_constraints, - &mut mb.motor_delay_state, + &mb.motor_delay_state, &mb.constraint_softness, args.batch_indices, )?; @@ -359,7 +359,7 @@ impl GpuMultibodySolver { &mut mb.joint_constraints, &mut mb.joint_constraint_columns, &mb.dof_couplings, - &mut mb.motor_delay_state, + &mb.motor_delay_state, &mb.dof_state, &mb.constraint_softness, args.batch_indices, From 1656dfa2a0203b5e818e7d92e1f670b88916a8a2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Fri, 28 Aug 2026 16:01:47 +0200 Subject: [PATCH 36/41] fix(rbd): split the joint-constraint back-solve into its own dispatch to fit 8 storage buffers --- .../dynamics/multibody/multibody_solver.rs | 31 +++++-- .../dynamics/multibody/joint_constraints.rs | 92 +++++++++++++------ 2 files changed, 86 insertions(+), 37 deletions(-) diff --git a/src_rbd/dynamics/multibody/multibody_solver.rs b/src_rbd/dynamics/multibody/multibody_solver.rs index 87fcbd76..4270ba68 100644 --- a/src_rbd/dynamics/multibody/multibody_solver.rs +++ b/src_rbd/dynamics/multibody/multibody_solver.rs @@ -6,10 +6,10 @@ use crate::queries::GpuIndexedContact; use crate::shaders::dynamics::{ GpuMbApplyContactRestitution, GpuMbBuildContactDelassus, GpuMbComputeDynamicsPre, GpuMbComputeSolveBounds, GpuMbDelayTick, GpuMbFinalizeContactConstraints, - GpuMbFinalizeImpulseJointConstraints, GpuMbGravityAndLu, GpuMbGravityAndLuT1, - GpuMbGravityAndLuT8, GpuMbGravityAndLuT16, GpuMbGravityAndLuT32, GpuMbInitContactConstraints, - GpuMbInitJointConstraints, GpuMbIntegrate, GpuMbIntegrateVelocities, - GpuMbRefreshJointConstraints, GpuMbRemoveImpulseJointConstraintBias, + GpuMbFinalizeImpulseJointConstraints, GpuMbFinalizeJointConstraints, GpuMbGravityAndLu, + GpuMbGravityAndLuT1, GpuMbGravityAndLuT8, GpuMbGravityAndLuT16, GpuMbGravityAndLuT32, + GpuMbInitContactConstraints, GpuMbInitJointConstraints, GpuMbIntegrate, + GpuMbIntegrateVelocities, GpuMbRefreshJointConstraints, GpuMbRemoveImpulseJointConstraintBias, GpuMbSeedContactRestitution, GpuMbSenseContactImpulses, GpuMbSnapshotContactWarmstart, GpuMbSolveConstraints, GpuMbSolveContactsDelassus, GpuMbSolveImpulseJointConstraints, GpuMbSolveJoints, GpuMbStashContactsLen, GpuMbTransferContactWarmstart, @@ -34,6 +34,9 @@ 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. @@ -354,16 +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. diff --git a/src_rbd_shaders/dynamics/multibody/joint_constraints.rs b/src_rbd_shaders/dynamics/multibody/joint_constraints.rs index fd9d1e59..136a5f21 100644 --- a/src_rbd_shaders/dynamics/multibody/joint_constraints.rs +++ b/src_rbd_shaders/dynamics/multibody/joint_constraints.rs @@ -535,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( @@ -549,21 +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(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 = 8)] motor_delay_state: &[f32], + #[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 = 9)] dof_state: &[f32], - #[spirv(uniform, descriptor_set = 0, binding = 10)] softness: &ConstraintSoftness, - #[spirv(uniform, descriptor_set = 0, binding = 11)] batch_ids: &BatchIndices, + #[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; @@ -588,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 { @@ -634,17 +624,59 @@ pub fn gpu_mb_init_joint_constraints( 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); From b3b4a73c8c6e718e81e6dea6141924c768a6fab6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Fri, 28 Aug 2026 16:04:23 +0200 Subject: [PATCH 37/41] test(rbd): keep the bench harness under wgpu's default buffer-size limit --- src_rbd/pipeline/bench_narrow_phase.rs | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/src_rbd/pipeline/bench_narrow_phase.rs b/src_rbd/pipeline/bench_narrow_phase.rs index 8b58f2a8..5dda6284 100644 --- a/src_rbd/pipeline/bench_narrow_phase.rs +++ b/src_rbd/pipeline/bench_narrow_phase.rs @@ -54,7 +54,9 @@ async fn run_bench(num_envs: u32, num_boxes: usize, num_steps: u32) { let capacities = RbdCapacities { batches: num_envs, - collisions_capacity: 256, + // ~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); @@ -85,7 +87,14 @@ async fn run_bench(num_envs: u32, num_boxes: usize, num_steps: u32) { #[serial_test::serial] #[ignore] async fn bench_narrow_phase_sweep() { - for envs in [1u32, 64, 256, 1024, 4096] { + // 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; } } From 8f2bac6dd6109437382166c98902ca43b75de13e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Fri, 28 Aug 2026 16:09:34 +0200 Subject: [PATCH 38/41] fix(rbd): split the batched env reset into pose and DoF passes to fit 8 storage buffers --- src_rbd/dynamics/multibody/env_reset.rs | 31 +++++++-- .../dynamics/multibody/env_reset.rs | 67 +++++++++++++------ 2 files changed, 74 insertions(+), 24 deletions(-) diff --git a/src_rbd/dynamics/multibody/env_reset.rs b/src_rbd/dynamics/multibody/env_reset.rs index 2d0eaa79..47644724 100644 --- a/src_rbd/dynamics/multibody/env_reset.rs +++ b/src_rbd/dynamics/multibody/env_reset.rs @@ -15,8 +15,8 @@ use super::multibody_set::GpuMultibodySet; use crate::math::Vector; use crate::shaders::dynamics::{ - GpuMbEnvReset, GpuMbEnvResetBatch, MULTIBODY_ROOT, MultibodyLinkStatic, MultibodyLinkWorkspace, - WS_QUADS, ws_soa_from_structs, ws_soa_to_structs, + GpuMbEnvReset, GpuMbEnvResetBatch, GpuMbEnvResetBatchDofs, MULTIBODY_ROOT, MultibodyLinkStatic, + MultibodyLinkWorkspace, WS_QUADS, ws_soa_from_structs, ws_soa_to_structs, }; use glamx::{UVec4, Vec4}; use khal::BufferUsages; @@ -105,6 +105,8 @@ struct EnvResetShader { #[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 @@ -382,13 +384,22 @@ impl GpuMultibodySet { &mut pass, [lpb * WS_QUADS, n, 1], &tpl.ws, - &tpl.links, - &tpl.dofs, &tpl.flags, &t_resets, &t_offs, - &t_vels, &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, @@ -399,3 +410,13 @@ impl GpuMultibodySet { self.reset_templates = Some(tpl); } } + +/// Builds both env-reset shader bundles, so the pipeline-limit smoke test can +/// reach them without making the private bundle types public. +#[cfg(test)] +pub(crate) fn env_reset_shaders_for_test(backend: &GpuBackend) { + use khal::Shader; + EnvResetShader::from_backend(backend).expect("env-reset shader exceeds the WebGPU limits"); + EnvResetBatchShader::from_backend(backend) + .expect("batched env-reset shader exceeds the WebGPU limits"); +} diff --git a/src_rbd_shaders/dynamics/multibody/env_reset.rs b/src_rbd_shaders/dynamics/multibody/env_reset.rs index d6b7246f..712df56f 100644 --- a/src_rbd_shaders/dynamics/multibody/env_reset.rs +++ b/src_rbd_shaders/dynamics/multibody/env_reset.rs @@ -63,9 +63,11 @@ pub fn gpu_mb_env_reset( /// (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. The DoF velocity -/// section comes from `dof_vels` (host-randomized reset velocities, or zeros), -/// replacing a per-DoF strided `write_buffer` loop. +/// 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. /// @@ -77,26 +79,17 @@ pub fn gpu_mb_env_reset( 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)] - templates_links: &[MultibodyLinkStatic], - #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] templates_dofs: &[f32], - #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] link_flags: &[u32], - #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] resets: &[UVec4], - #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] offsets: &[Vec4], - #[spirv(storage_buffer, descriptor_set = 0, binding = 6)] dof_vels: &[f32], - #[spirv(storage_buffer, descriptor_set = 0, binding = 7)] links_workspace: &mut [Vec4], - #[spirv(storage_buffer, descriptor_set = 0, binding = 8)] - links_static: &mut [MultibodyLinkStatic], - #[spirv(storage_buffer, descriptor_set = 0, binding = 9)] dof_values: &mut [f32], - #[spirv(storage_buffer, descriptor_set = 0, binding = 10)] dof_state: &mut [f32], + #[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 = 11)] params: &UVec4, + #[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; - let dpb = params.z; if r >= params.w { return; } @@ -125,6 +118,44 @@ pub fn gpu_mb_env_reset_batch( } 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, @@ -132,8 +163,6 @@ pub fn gpu_mb_env_reset_batch( ); } if i < dpb { - // Generalized coords are translation-invariant: the free root's world - // position lives in the workspace coords quad handled above. dof_values.write( (i * nb + env) as usize, templates_dofs.read((t * 2 * dpb + i) as usize), From f7dcd9c044a58e0869302d4d89a8d8dd171c422a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Fri, 28 Aug 2026 18:03:44 +0200 Subject: [PATCH 39/41] chore: switch to the published rapier version --- Cargo.toml | 24 +++++++++--------------- 1 file changed, 9 insertions(+), 15 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index ebf90a18..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,17 +93,11 @@ 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). -# -# Also carries `Multibody::frictions` (MJCF ``), which the -# multibody build path reads into the GPU per-DoF friction section. All five -# must be patched together so the rapier types unify. -rapier2d = { path = "../rapier/crates/rapier2d" } -rapier3d = { path = "../rapier/crates/rapier3d" } -rapier3d-mjcf = { path = "../rapier/crates/rapier3d-mjcf" } -rapier3d-urdf = { path = "../rapier/crates/rapier3d-urdf" } -rapier3d-meshloader = { path = "../rapier/crates/rapier3d-meshloader" } +#rapier2d = { path = "../rapier/crates/rapier2d" } +#rapier3d = { path = "../rapier/crates/rapier3d" } +#rapier3d-mjcf = { path = "../rapier/crates/rapier3d-mjcf" } +#rapier3d-urdf = { path = "../rapier/crates/rapier3d-urdf" } +#rapier3d-meshloader = { path = "../rapier/crates/rapier3d-meshloader" } ## Local glam clone with SPIR-V vector-arithmetic intrinsics (Vec3 add/sub/mul/scale). #glam = { path = "../glam-rs" } # 30% faster for loop in P2G From 82093c61265c60d9117c8d41580e85a74068b196 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Fri, 28 Aug 2026 18:56:40 +0200 Subject: [PATCH 40/41] fix: make all envs share the same RbdSimParams --- src_rbd/pipeline/insertion_removal.rs | 11 +++--- src_rbd/pipeline/rbd_state.rs | 12 +++---- src_rbd/pipeline/rbd_state_from_rapier.rs | 37 ++++++++++++-------- src_rbd_shaders/dynamics/joint_constraint.rs | 4 +-- src_rbd_shaders/dynamics/sim_params.rs | 2 +- src_rbd_shaders/dynamics/solver.rs | 15 +++----- 6 files changed, 40 insertions(+), 41 deletions(-) diff --git a/src_rbd/pipeline/insertion_removal.rs b/src_rbd/pipeline/insertion_removal.rs index 45411b80..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. @@ -102,7 +101,7 @@ impl RbdState { .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,13 +242,13 @@ impl RbdState { num_batches, num_colliders_per_batch, num_solver_iterations, - sim_params: Tensor::vector( + sim_params: Tensor::scalar( backend, - &all_sim_params, - BufferUsages::STORAGE | BufferUsages::UNIFORM | BufferUsages::COPY_DST, + base_sim_params, + BufferUsages::UNIFORM | BufferUsages::COPY_DST, ) .unwrap(), - all_sim_params, + sim_params_cpu: base_sim_params, vels: Tensor::vector(backend, &all_vels, rw).unwrap(), #[cfg(feature = "dim3")] reset_templates_bodies: None, diff --git a/src_rbd/pipeline/rbd_state.rs b/src_rbd/pipeline/rbd_state.rs index dfb2819b..0ee03857 100644 --- a/src_rbd/pipeline/rbd_state.rs +++ b/src_rbd/pipeline/rbd_state.rs @@ -125,7 +125,7 @@ pub struct RbdState { 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) all_sim_params: Vec, + 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 @@ -380,12 +380,10 @@ impl RbdState { /// 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.all_sim_params.clone(); - for p in &mut params { - p.contact_merge_cos = cos; - } - let _ = backend.write_buffer(self.sim_params.buffer_mut(), 0, ¶ms); - self.all_sim_params = params; + 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. diff --git a/src_rbd/pipeline/rbd_state_from_rapier.rs b/src_rbd/pipeline/rbd_state_from_rapier.rs index 6a4b80d1..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)" + ); } } @@ -156,20 +161,24 @@ impl RbdState { 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 @@ -484,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 @@ -772,13 +781,13 @@ impl RbdState { num_batches, num_colliders_per_batch: num_colliders_per_batch as u32, num_solver_iterations, - sim_params: Tensor::vector( + sim_params: Tensor::scalar( backend, - &all_sim_params, - BufferUsages::STORAGE | BufferUsages::UNIFORM | BufferUsages::COPY_DST, + sim_params, + BufferUsages::UNIFORM | BufferUsages::COPY_DST, ) .unwrap(), - all_sim_params, + sim_params_cpu: sim_params, vels: Tensor::vector(backend, &all_vels, storage | BufferUsages::COPY_DST).unwrap(), #[cfg(feature = "dim3")] reset_templates_bodies: None, 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/sim_params.rs b/src_rbd_shaders/dynamics/sim_params.rs index 4746b68c..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 { 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; From e70e5f789d87f0f2eeee531a8825f4f43c1c927c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Fri, 28 Aug 2026 19:12:11 +0200 Subject: [PATCH 41/41] chore: clippy fixes --- src_rbd/dynamics/multibody/env_reset.rs | 10 ---------- 1 file changed, 10 deletions(-) diff --git a/src_rbd/dynamics/multibody/env_reset.rs b/src_rbd/dynamics/multibody/env_reset.rs index 47644724..8d8008c5 100644 --- a/src_rbd/dynamics/multibody/env_reset.rs +++ b/src_rbd/dynamics/multibody/env_reset.rs @@ -410,13 +410,3 @@ impl GpuMultibodySet { self.reset_templates = Some(tpl); } } - -/// Builds both env-reset shader bundles, so the pipeline-limit smoke test can -/// reach them without making the private bundle types public. -#[cfg(test)] -pub(crate) fn env_reset_shaders_for_test(backend: &GpuBackend) { - use khal::Shader; - EnvResetShader::from_backend(backend).expect("env-reset shader exceeds the WebGPU limits"); - EnvResetBatchShader::from_backend(backend) - .expect("batched env-reset shader exceeds the WebGPU limits"); -}