From ac33d1f3da0bcb55e5789d88fc3ad29a5e289b16 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Sat, 29 Aug 2026 21:07:41 +0200 Subject: [PATCH 1/7] perf(rbd): batch-interleaved flat buffer layout with demand-sized constraint arenas --- src/state.rs | 29 +- src_rbd/broad_phase/lbvh.rs | 21 +- src_rbd/broad_phase/narrow_phase.rs | 96 +++-- src_rbd/dynamics/coloring.rs | 68 ++-- src_rbd/dynamics/joint.rs | 17 +- src_rbd/dynamics/mprops_update.rs | 4 +- .../dynamics/multibody/loop_closing_joints.rs | 45 ++- .../multibody/multibody_from_rapier.rs | 94 +++-- src_rbd/dynamics/multibody/multibody_set.rs | 114 ++++-- .../dynamics/multibody/multibody_solver.rs | 69 ++-- src_rbd/dynamics/solver.rs | 47 +-- src_rbd/dynamics/warmstart.rs | 10 +- src_rbd/pipeline/insertion_removal.rs | 270 +++++++------ src_rbd/pipeline/rbd_state.rs | 97 +++-- src_rbd/pipeline/rbd_state_from_rapier.rs | 107 +++-- src_rbd/pipeline/rbd_step.rs | 229 ++++++----- src_rbd_shaders/broad_phase/brute_force.rs | 29 +- src_rbd_shaders/broad_phase/lbvh.rs | 111 +++-- src_rbd_shaders/broad_phase/narrow_phase.rs | 333 ++++++++------- src_rbd_shaders/dynamics/color_buckets.rs | 99 ++--- src_rbd_shaders/dynamics/coloring.rs | 107 +++-- src_rbd_shaders/dynamics/joint_constraint.rs | 103 +++-- .../dynamics/joint_constraint_builder.rs | 8 +- src_rbd_shaders/dynamics/mprops_update.rs | 22 +- .../multibody/compute_dynamics_pre.rs | 112 +++--- .../dynamics/multibody/contact_constraints.rs | 378 ++++++++++-------- .../dynamics/multibody/contact_sensor.rs | 5 +- .../dynamics/multibody/env_reset.rs | 14 +- .../dynamics/multibody/gravity_and_lu.rs | 157 ++++---- .../impulse_joint_constraints/helper.rs | 31 +- .../impulse_joint_constraints/jacobians.rs | 34 +- .../impulse_joint_constraints/kernels.rs | 115 +++--- .../impulse_joint_constraints/update.rs | 70 ++-- .../dynamics/multibody/integrate.rs | 7 +- .../dynamics/multibody/joint_constraints.rs | 41 +- .../dynamics/multibody/solve_constraints.rs | 155 +++---- src_rbd_shaders/dynamics/multibody/types.rs | 7 +- src_rbd_shaders/dynamics/multibody/ws_soa.rs | 7 +- src_rbd_shaders/dynamics/solver.rs | 269 ++++++------- src_rbd_shaders/dynamics/warmstart.rs | 54 ++- src_rbd_shaders/utils/indices.rs | 222 ++-------- src_rbd_shaders/utils/mod.rs | 2 +- src_viewer/ui.rs | 2 + 43 files changed, 1952 insertions(+), 1859 deletions(-) diff --git a/src/state.rs b/src/state.rs index 94d2b474..4840b6ee 100644 --- a/src/state.rs +++ b/src/state.rs @@ -64,6 +64,11 @@ impl NexusCapacities { self } + pub fn rbd_mb_contact_constraints(mut self, capacity: u32) -> Self { + self.rbd.mb_contact_constraints_capacity = capacity; + self + } + #[cfg(feature = "mpm")] pub fn mpm_grid_size(mut self, num_chunks: u32) -> Self { self.mpm.grid_size = num_chunks; @@ -109,6 +114,8 @@ pub struct NexusCounts { pub multibody_dofs: usize, pub collision_pairs: usize, pub collision_pairs_capacity: usize, + pub mb_contact_constraints: usize, + pub mb_contact_constraints_capacity: usize, pub particles: usize, } @@ -342,6 +349,10 @@ impl NexusState { self.capacities.rbd.collisions_capacity = capacity.max(1); } + pub fn set_rbd_mb_contact_constraints_capacity(&mut self, capacity: u32) { + self.capacities.rbd.mb_contact_constraints_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) { @@ -373,6 +384,12 @@ impl NexusState { if let Some(rbd) = self.rbd.as_ref() { c.collision_pairs = rbd.collision_pairs_len() as usize; c.collision_pairs_capacity = rbd.collision_pairs_capacity() as usize; + #[cfg(feature = "dim3")] + { + c.mb_contact_constraints = rbd.mb_contact_constraints_len() as usize; + c.mb_contact_constraints_capacity = + rbd.mb_contact_constraints_capacity() as usize; + } } #[cfg(feature = "mpm")] if let Some(mpm) = self.mpm.as_ref() { @@ -640,13 +657,13 @@ impl NexusState { <= rbd.num_colliders_per_batch() as usize => { let range = rbd.append_bodies(backend, &gpu_pairs)?; - // Single environment: the per-batch local slot is the gpu_id. + let nb = rbd.num_batches(); for (i, (&handle, &coupling)) in handles.iter().zip(&couplings).enumerate() { self.rbd2gpu[0].insert( handle.0, GpuRigidBodyRef { coupling, - gpu_id: range.start + i as u32, + gpu_id: (range.start + i as u32) * nb, }, ); } @@ -981,9 +998,7 @@ impl NexusState { // `gpu_id` is its *body* slot, not a collider slot, since a body may // own several colliders. Body slots are assigned in the order // `from_rapier` uses (the first time each parent body is seen while - // iterating colliders) and are laid out env-major with stride - // `num_colliders_per_batch`. - let stride = rbd_state.num_colliders_per_batch(); + let nb = rbd_state.num_batches(); for (env_idx, world) in self.rbd_envs.iter().enumerate() { let mut body_slot: std::collections::HashMap<_, u32> = std::collections::HashMap::new(); @@ -1012,7 +1027,7 @@ impl NexusState { body_handle.0, GpuRigidBodyRef { coupling, - gpu_id: env_idx as u32 * stride + slot, + gpu_id: slot * nb + env_idx as u32, }, ); } @@ -1038,7 +1053,7 @@ impl NexusState { body_handle.0, GpuRigidBodyRef { coupling, - gpu_id: env_idx as u32 * stride + slot, + gpu_id: slot * nb + env_idx as u32, }, ); } diff --git a/src_rbd/broad_phase/lbvh.rs b/src_rbd/broad_phase/lbvh.rs index a081fab2..347a8804 100644 --- a/src_rbd/broad_phase/lbvh.rs +++ b/src_rbd/broad_phase/lbvh.rs @@ -7,8 +7,8 @@ use crate::math::Pose; use crate::shaders::PaddedVector; use crate::shaders::bounding_volumes::Aabb; use crate::shaders::broad_phase::{ - CollisionPair, GpuBfComputeAabbs, GpuBfFindPairs, GpuLbvhBuild, GpuLbvhComputeDomain, - GpuLbvhComputeMorton, GpuLbvhFindCollisionPairs, GpuLbvhInitDispatch, GpuLbvhRefitInternal, + CollisionPair, GpuBfComputeAabbs, GpuBfFindPairs, GpuFlatListDispatch, GpuLbvhBuild, + GpuLbvhComputeDomain, GpuLbvhComputeMorton, GpuLbvhFindCollisionPairs, GpuLbvhRefitInternal, GpuLbvhRefitLeaves, GpuLbvhResetCollisionPairs, LbvhNode, }; use crate::shaders::shapes::Shape; @@ -32,7 +32,7 @@ pub struct GpuLbvh { refit_internal: GpuLbvhRefitInternal, reset_collision_pairs: GpuLbvhResetCollisionPairs, find_collision_pairs: GpuLbvhFindCollisionPairs, - lbvh_init_indirect_args: GpuLbvhInitDispatch, + flat_list_dispatch: GpuFlatListDispatch, // Kernels for brute-force broad-phase for small scenes // (typically, small scenes but many batches). bf_compute_aabbs: GpuBfComputeAabbs, @@ -294,7 +294,7 @@ impl Lbvh { self.shaders .reset_collision_pairs - .call(pass, [num_batches, 1, 1], collision_pairs_len)?; + .call(pass, [1u32, 1, 1], collision_pairs_len)?; self.shaders.find_collision_pairs.call( pass, [colliders_per_batch, num_batches, 1], @@ -305,11 +305,12 @@ impl Lbvh { batch_indices, pair_filter, )?; - self.shaders.lbvh_init_indirect_args.call( + self.shaders.flat_list_dispatch.call( pass, - 256u32, + 1u32, collision_pairs_len, collision_pairs_indirect, + batch_indices, )?; Ok(()) } @@ -351,7 +352,7 @@ impl Lbvh { )?; self.shaders .reset_collision_pairs - .call(pass, [num_batches, 1, 1], collision_pairs_len)?; + .call(pass, [1u32, 1, 1], collision_pairs_len)?; self.shaders.bf_find_pairs.call( pass, [active_per_batch * active_per_batch * num_batches, 1, 1], @@ -363,12 +364,12 @@ impl Lbvh { pair_filter, sim_params, )?; - // Single 256-lane workgroup: parallel max over the per-batch counts. - self.shaders.lbvh_init_indirect_args.call( + self.shaders.flat_list_dispatch.call( pass, - 256u32, + 1u32, collision_pairs_len, collision_pairs_indirect, + batch_indices, )?; Ok(()) } diff --git a/src_rbd/broad_phase/narrow_phase.rs b/src_rbd/broad_phase/narrow_phase.rs index 94cbdecc..94a47c1f 100644 --- a/src_rbd/broad_phase/narrow_phase.rs +++ b/src_rbd/broad_phase/narrow_phase.rs @@ -6,9 +6,10 @@ use crate::shaders::PaddedVector; #[cfg(feature = "dim3")] use crate::shaders::broad_phase::GpuReduceContacts; use crate::shaders::broad_phase::{ - CollisionPair, GpuInitPfmPfmDispatch, GpuNarrowPhaseInitContactsDispatch, GpuNarrowPhasePfmPfm, + CollisionPair, GpuContactOffsetsScan, GpuCountPairsPerBatch, GpuCountPfmPerBatch, + GpuFlatListDispatch, GpuNarrowPhaseInitContactsDispatch, GpuNarrowPhasePfmPfm, GpuNarrowPhaseShapeShape, GpuNarrowPhaseShapeShapeDeferred, GpuResetNarrowPhase, - NarrowPhasePfmPair, + GpuZeroContactLens, NarrowPhasePfmPair, }; use crate::shaders::shapes::Shape; use khal::Shader; @@ -26,7 +27,11 @@ pub struct GpuNarrowPhase { narrow_phase_pfm_pfm: GpuNarrowPhasePfmPfm, #[cfg(feature = "dim3")] reduce_contacts: GpuReduceContacts, - init_pfm_pfm_indirect_args: GpuInitPfmPfmDispatch, + count_pairs_per_batch: GpuCountPairsPerBatch, + count_pfm_per_batch: GpuCountPfmPerBatch, + contact_offsets_scan: GpuContactOffsetsScan, + zero_contact_lens: GpuZeroContactLens, + flat_list_dispatch: GpuFlatListDispatch, init_contacts_indirect_args: GpuNarrowPhaseInitContactsDispatch, } @@ -41,11 +46,13 @@ impl GpuNarrowPhase { vertices: &Tensor, indices: &Tensor, collision_pairs: &Tensor, - collision_pairs_len: &Tensor, - collision_pairs_indirect: &Tensor<[u32; 3]>, + collision_pairs_len: &mut Tensor, contacts: &mut Tensor, contacts_len: &mut Tensor, contacts_indirect: &mut Tensor<[u32; 3]>, + contact_offsets: &mut Tensor, + pair_batch_counts: &mut Tensor, + pfm_batch_counts: &mut Tensor, mb_sweep_indirect: &mut Tensor<[u32; 3]>, pfm_pairs: &mut Tensor, pfm_pairs_len: &mut Tensor, @@ -57,24 +64,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, + collision_pairs_indirect: &Tensor<[u32; 3]>, ) -> 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( + self.reset_narrow_phase.call( pass, - collision_pairs_indirect, - collision_pairs, - collision_pairs_len, - poses, - shapes, - contacts, + [num_batches, 1, 1], contacts_len, - batch_indices, - collider_parent, - collider_materials, - sim_params, + pfm_pairs_len, + pair_batch_counts, + pfm_batch_counts, )?; // Pass 2: defer the complex shape pairs into `pfm_pairs` (kept as a @@ -94,15 +93,64 @@ impl GpuNarrowPhase { indices, )?; - self.init_pfm_pfm_indirect_args - .call(pass, 256u32, pfm_pairs_len, pfm_pairs_indirect)?; + self.count_pairs_per_batch.call( + pass, + collision_pairs_indirect, + collision_pairs, + collision_pairs_len, + pair_batch_counts, + batch_indices, + )?; + self.flat_list_dispatch + .call(pass, 1u32, pfm_pairs_len, pfm_pairs_indirect, batch_indices)?; + self.count_pfm_per_batch.call( + pass, + &*pfm_pairs_indirect, + pfm_pairs, + pfm_pairs_len, + pfm_batch_counts, + batch_indices, + )?; + self.contact_offsets_scan.call( + pass, + 1u32, + pair_batch_counts, + pfm_batch_counts, + collision_pairs_len, + pfm_pairs_len, + contact_offsets, + contacts_indirect, + batch_indices, + )?; + self.zero_contact_lens.call( + pass, + &*contacts_indirect, + contacts, + contact_offsets, + batch_indices, + )?; + + self.narrow_phase.call( + pass, + collision_pairs_indirect, + collision_pairs, + contact_offsets, + poses, + shapes, + contacts, + contacts_len, + batch_indices, + collider_parent, + collider_materials, + sim_params, + )?; self.narrow_phase_pfm_pfm.call( pass, &*pfm_pairs_indirect, contacts, contacts_len, pfm_pairs, - pfm_pairs_len, + contact_offsets, batch_indices, vertices, indices, @@ -110,8 +158,6 @@ impl GpuNarrowPhase { collider_materials, sim_params, )?; - // Reduction rewrites `contacts_len`, so it has to run before the - // indirect args are derived from it. #[cfg(feature = "dim3")] if reduce_contacts { self.reduce_contacts.call( @@ -119,6 +165,7 @@ impl GpuNarrowPhase { [1u32, num_batches, 1], contacts, contacts_len, + contact_offsets, batch_indices, sim_params, )?; @@ -129,7 +176,6 @@ impl GpuNarrowPhase { pass, 256u32, contacts_len, - contacts_indirect, mb_sweep_indirect, batch_indices, )?; diff --git a/src_rbd/dynamics/coloring.rs b/src_rbd/dynamics/coloring.rs index f187bfd3..2c21a64e 100644 --- a/src_rbd/dynamics/coloring.rs +++ b/src_rbd/dynamics/coloring.rs @@ -10,10 +10,11 @@ use crate::pipeline::RunStats; use crate::shaders::dynamics::TwoBodyConstraint; use crate::shaders::dynamics::{ - GpuColorBucketsCount, GpuColorBucketsReset, GpuColorBucketsScan, GpuColorBucketsScatter, - GpuFixConflictsTopoGc, GpuResetCompletionFlagTopoGc, GpuResetLuby, GpuResetTopoGc, - GpuStepGraphColoringLuby, GpuStepGraphColoringTopoGc, + GpuColorBucketsCount, GpuColorBucketsReset, GpuColorBucketsScatter, GpuFixConflictsTopoGc, + GpuResetCompletionFlagTopoGc, GpuResetLuby, GpuResetTopoGc, GpuStepGraphColoringLuby, + GpuStepGraphColoringTopoGc, }; +use crate::utils::{GpuPrefixSum, PrefixSumWorkspace}; use khal::Shader; use khal::backend::{Backend, Encoder, GpuBackend, GpuBackendError, GpuPass, GpuTimestamps}; use vortx::tensor::Tensor; @@ -38,7 +39,6 @@ pub struct GpuColoring { // only touches their own constraint. color_buckets_reset: GpuColorBucketsReset, color_buckets_count: GpuColorBucketsCount, - color_buckets_scan: GpuColorBucketsScan, color_buckets_scatter: GpuColorBucketsScatter, } @@ -48,15 +48,9 @@ pub struct ColorBucketsArgs<'a> { pub contacts_len_indirect: &'a Tensor<[u32; 3]>, /// Color assigned to each constraint by graph coloring. pub constraints_colors: &'a Tensor, - /// Number of contacts per batch. - pub contacts_len: &'a Tensor, - /// Per-batch per-color counts (stride `solver_color_buckets_stride`). - pub color_bucket_counts: &'a mut Tensor, - /// Per-batch per-color exclusive prefix sums. - pub color_bucket_starts: &'a mut Tensor, - /// Scatter cursors (seeded from the starts). - pub color_bucket_cursors: &'a mut Tensor, - /// Constraint ids bucket-sorted by color (contacts layout). + pub constraints: &'a Tensor, + pub contact_offsets: &'a Tensor, + pub color_buckets: &'a mut Tensor, pub color_sorted_ids: &'a mut Tensor, /// Shared per-batch capacity / section-offset uniform. pub batch_indices: &'a Tensor, @@ -84,8 +78,7 @@ pub struct ColoringArgs<'a> { pub uncolored: &'a mut Tensor, /// Staging buffer for reading uncolored count on CPU. pub uncolored_staging: &'a Tensor, - /// Total number of contacts. - pub contacts_len: &'a Tensor, + pub contact_offsets: &'a Tensor, /// Buffer tracking which constraints are colored. pub colored: &'a mut Tensor, /// Shared per-batch capacity / section-offset uniform. @@ -108,7 +101,8 @@ impl GpuColoring { args.contacts_len_indirect, args.constraints_colors, args.constraints_rands, - args.contacts_len, + args.constraints, + args.contact_offsets, args.batch_indices, )?; Ok(()) @@ -131,7 +125,7 @@ impl GpuColoring { args.uncolored, args.body_group, args.curr_color, - args.contacts_len, + args.contact_offsets, args.batch_indices, )?; Ok(()) @@ -148,7 +142,8 @@ impl GpuColoring { args.contacts_len_indirect, args.constraints_colors, args.colored, - args.contacts_len, + args.constraints, + args.contact_offsets, args.batch_indices, )?; Ok(()) @@ -169,7 +164,7 @@ impl GpuColoring { args.constraints_colors, args.colored, args.uncolored, - args.contacts_len, + args.contact_offsets, args.body_group, args.batch_indices, )?; @@ -191,7 +186,7 @@ impl GpuColoring { args.constraints_colors, args.colored, args.uncolored, - args.contacts_len, + args.contact_offsets, args.body_group, args.batch_indices, )?; @@ -251,39 +246,32 @@ impl GpuColoring { /// Bucket-sorts the constraint ids by their color. pub fn dispatch_build_color_buckets( &self, + backend: &GpuBackend, pass: &mut GpuPass, args: ColorBucketsArgs<'_>, - color_buckets_stride: u32, - num_batches: u32, + prefix_sum: &GpuPrefixSum, + prefix_workspace: &mut PrefixSumWorkspace, ) -> Result<(), GpuBackendError> { - self.color_buckets_reset.call( - pass, - [color_buckets_stride, num_batches, 1], - args.color_bucket_counts, - args.batch_indices, - )?; + let num_buckets = args.color_buckets.len() as u32; + self.color_buckets_reset + .call(pass, [num_buckets, 1, 1], args.color_buckets)?; self.color_buckets_count.call( pass, args.contacts_len_indirect, args.constraints_colors, - args.contacts_len, - args.color_bucket_counts, - args.batch_indices, - )?; - self.color_buckets_scan.call( - pass, - [1, num_batches, 1], - args.color_bucket_counts, - args.color_bucket_starts, - args.color_bucket_cursors, + args.constraints, + args.contact_offsets, + args.color_buckets, args.batch_indices, )?; + prefix_sum.launch(backend, pass, prefix_workspace, args.color_buckets, 1)?; self.color_buckets_scatter.call( pass, args.contacts_len_indirect, args.constraints_colors, - args.contacts_len, - args.color_bucket_cursors, + args.constraints, + args.contact_offsets, + args.color_buckets, args.color_sorted_ids, args.batch_indices, )?; diff --git a/src_rbd/dynamics/joint.rs b/src_rbd/dynamics/joint.rs index 9c89214e..f4d4eb70 100644 --- a/src_rbd/dynamics/joint.rs +++ b/src_rbd/dynamics/joint.rs @@ -309,12 +309,11 @@ impl GpuImpulseJointSet { // Build flat joint buffer [num_batches * max_joints], padded with zeroed joints. let dummy_joint = ImpulseJoint::zeroed(); - let mut all_joints = Vec::with_capacity(num_batches as usize * max_joints as usize); - for sorted_joints in &per_env_sorted_joints { - all_joints.extend_from_slice(sorted_joints); - // Pad to max_joints. - for _ in sorted_joints.len()..max_joints as usize { - all_joints.push(dummy_joint); + let mut all_joints = + vec![dummy_joint; num_batches as usize * max_joints as usize]; + for (batch, sorted_joints) in per_env_sorted_joints.iter().enumerate() { + for (j, joint) in sorted_joints.iter().enumerate() { + all_joints[j * num_batches as usize + batch] = *joint; } } @@ -407,7 +406,7 @@ impl GpuJointSolver { self.init_joint_constraints.call( pass, - [args.joints.len, args.num_batches, 1], + args.joints.len * args.num_batches, &args.joints.joints, &mut args.joints.builders, &mut args.joints.constraints, @@ -430,7 +429,7 @@ impl GpuJointSolver { self.update_joint_constraints.call( pass, - [args.joints.len, args.num_batches, 1], + args.joints.len * args.num_batches, &args.joints.builders, &mut args.joints.constraints, poses, @@ -472,7 +471,7 @@ impl GpuJointSolver { } self.solve_joint_constraints.call( pass, - [group_len, args.num_batches, 1], + group_len * args.num_batches, &mut args.joints.constraints, solver_vels, &args.joints.color_groups, diff --git a/src_rbd/dynamics/mprops_update.rs b/src_rbd/dynamics/mprops_update.rs index 05968c51..cb62506f 100644 --- a/src_rbd/dynamics/mprops_update.rs +++ b/src_rbd/dynamics/mprops_update.rs @@ -29,7 +29,7 @@ impl GpuMpropsUpdate { ) -> Result<(), GpuBackendError> { self.update_mprops_kernel.call( pass, - [num_bodies, num_batches, 1], + num_bodies * num_batches, mprops, local_mprops, body_poses, @@ -67,7 +67,7 @@ impl GpuSyncColliderPosesShader { ) -> Result<(), GpuBackendError> { self.sync_kernel.call( pass, - [num_colliders, num_batches, 1], + num_colliders * num_batches, body_poses, collider_local_poses, collider_world_poses, diff --git a/src_rbd/dynamics/multibody/loop_closing_joints.rs b/src_rbd/dynamics/multibody/loop_closing_joints.rs index 45d06d05..9db35330 100644 --- a/src_rbd/dynamics/multibody/loop_closing_joints.rs +++ b/src_rbd/dynamics/multibody/loop_closing_joints.rs @@ -250,27 +250,37 @@ impl GpuMultibodySet { let joints_cap = max_joints.max(1); let cons_cap = (joints_cap * MAX_AXIS_CONSTRAINTS).max(1); let jac_cap = max_jac_floats.max(1); + let nb = self.num_batches as usize; - let mut all_builders: Vec = - Vec::with_capacity((joints_cap * self.num_batches) as usize); - let mut all_counts: Vec = Vec::with_capacity(self.num_batches as usize); - // Padding builder: both sides marked FIXED so the GPU kernel can - // skip them by sentinel check (replaces the per-batch `num_joints` - // storage binding the kernel used to read for early-out). + for (b, env) in per_env_builders.iter().enumerate() { + assert_eq!( + env.len(), + per_env_builders[0].len(), + "batch {b}: multibody-touching impulse-joint count differs from batch 0 \ + (batched envs must have identical topology)" + ); + for (i, (builder, b0)) in env.iter().zip(per_env_builders[0].iter()).enumerate() { + assert!( + builder.jacobian_offset == b0.jacobian_offset + && builder.jacobian_capacity == b0.jacobian_capacity + && builder.constraint_id == b0.constraint_id, + "batch {b} joint slot {i}: jacobian/constraint layout differs from batch 0 \ + (batched envs must have identical topology)" + ); + } + } let mut dummy: MbImpulseJointBuilder = bytemuck::Zeroable::zeroed(); dummy.side_a_kind = SIDE_KIND_FIXED; dummy.side_b_kind = SIDE_KIND_FIXED; - for env in &per_env_builders { - all_counts.push(env.len() as u32); - all_builders.extend_from_slice(env); - for _ in env.len()..joints_cap as usize { - all_builders.push(dummy); + let mut all_builders: Vec = + vec![dummy; joints_cap as usize * nb]; + for (b, env) in per_env_builders.iter().enumerate() { + for (i, builder) in env.iter().enumerate() { + all_builders[i * nb + b] = *builder; } } let storage = BufferUsages::STORAGE | BufferUsages::COPY_DST; - let usage_u = storage | BufferUsages::UNIFORM; - self.mb_imp_joint_count = Tensor::vector(backend, &all_counts, usage_u).unwrap(); self.mb_imp_joint_builders = Tensor::vector(backend, &all_builders, storage).unwrap(); self.mb_imp_joint_constraints = Tensor::vector( backend, @@ -293,12 +303,11 @@ impl GpuMultibodySet { // colors are no-ops (start == end). `cols` is clamped to ≥1 so the // buffer is always a valid non-empty binding even with no joints. let cols = global_num_colors.max(1); - let mut all_color_groups = Vec::with_capacity((cols * self.num_batches) as usize); - for env_cg in &per_env_color_groups { + let mut all_color_groups = vec![0u32; cols as usize * nb]; + for (b, env_cg) in per_env_color_groups.iter().enumerate() { let last = env_cg.last().copied().unwrap_or(0); - all_color_groups.extend_from_slice(env_cg); - for _ in env_cg.len()..cols as usize { - all_color_groups.push(last); + for c in 0..cols as usize { + all_color_groups[c * nb + b] = env_cg.get(c).copied().unwrap_or(last); } } self.mb_imp_joint_color_groups = diff --git a/src_rbd/dynamics/multibody/multibody_from_rapier.rs b/src_rbd/dynamics/multibody/multibody_from_rapier.rs index ae9a4141..a9ccaf68 100644 --- a/src_rbd/dynamics/multibody/multibody_from_rapier.rs +++ b/src_rbd/dynamics/multibody/multibody_from_rapier.rs @@ -30,6 +30,7 @@ impl GpuMultibodySet { &RigidBodySet, )], colliders_per_batch: u32, + contact_constraint_slots: u32, ) -> Self { let num_batches = environments.len() as u32; @@ -141,14 +142,12 @@ impl GpuMultibodySet { let j = link.joint().data; let locked = j.locked_axes.bits() as u32; let limit_axes = j.limit_axes.bits() as u32 & !locked; - let motor_axes = j.motor_axes.bits() as u32 & !locked; - // 1 per active limit + 1 per active motor (axis-wise). let mut n = 0u32; for ax in 0u32..6 { if (limit_axes >> ax) & 1 != 0 { n += 1; } - if (motor_axes >> ax) & 1 != 0 { + if (locked >> ax) & 1 == 0 { n += 1; } } @@ -179,7 +178,11 @@ impl GpuMultibodySet { max_constraints, self_contacts_enabled: if mb.self_contacts_enabled() { 1 } else { 0 }, contact_constraint_count: 0, + contact_constraint_start: 0, + old_contact_constraint_start: 0, + old_contact_constraint_count: 0, batch_contacts_len: 0, + batch_contacts_start: 0, first_coupling: coupling_off, num_couplings, }); @@ -379,12 +382,9 @@ impl GpuMultibodySet { // One length-`dofs_cap` column of `M⁻¹` per constraint slot. let cons_col_cap = cons_cap.saturating_mul(dofs_cap).max(1); - // Per-multibody contact-constraint banks: every multibody owns a - // fixed-size slab of `MAX_MB_CONTACT_CONSTRAINTS_PER_MB` slots — - // each contact point produces 1 normal + (DIM-1) friction tangent - // constraint slots. The init kernel marks unused slots as `kind = 0`. - let contact_cons_cap = mb_cap - .saturating_mul(MAX_MB_CONTACT_CONSTRAINTS_PER_MB) + let contact_cons_cap = contact_constraint_slots + .saturating_mul(num_batches) + .max((mb_cap * num_batches).saturating_mul(crate::shaders::dynamics::MB_CONS_SLOT_RESERVE)) .max(1); let contact_cons_col_cap = contact_cons_cap.saturating_mul(dofs_cap).max(1); let body_to_link_cap = colliders_per_batch.max(1); @@ -429,6 +429,28 @@ impl GpuMultibodySet { let dummy_info = MultibodyInfo::default(); let dummy_stat: MultibodyLinkStatic = bytemuck::Zeroable::zeroed(); let dummy_ws = make_workspace_init(); + for (b, env) in per_env_infos.iter().enumerate() { + assert_eq!( + env.len(), + per_env_infos[0].len(), + "batch {b}: multibody count differs from batch 0 \ + (batched envs must have identical topology)" + ); + for (i, (info, i0)) in env.iter().zip(per_env_infos[0].iter()).enumerate() { + assert!( + info.first_link == i0.first_link + && info.num_links == i0.num_links + && info.first_dof == i0.first_dof + && info.ndofs == i0.ndofs + && info.jacobian_offset == i0.jacobian_offset + && info.mass_matrix_offset == i0.mass_matrix_offset + && info.coriolis_offset == i0.coriolis_offset + && info.i_coriolis_dt_offset == i0.i_coriolis_dt_offset, + "batch {b} multibody {i}: dynamics-arena layout differs from batch 0 \ + (batched envs must have identical topology)" + ); + } + } for i in 0..num_batches as usize { all_infos.extend_from_slice(&per_env_infos[i]); @@ -508,10 +530,8 @@ impl GpuMultibodySet { num_active_multibodies: global_max_mb, links_per_batch: links_cap, dofs_per_batch: dofs_cap, - jacobian_entries_per_batch: jac_cap, mass_matrix_entries_per_batch: mm_cap, coriolis_entries_per_batch: cor_cap, - i_coriolis_dt_entries_per_batch: icdt_cap, // Default: implicit coriolis. The acceleration solve uses a mass // matrix augmented with the coriolis/gyroscopic derivatives, while // constraints keep the plain one. `set_implicit_coriolis(false)` @@ -523,7 +543,8 @@ impl GpuMultibodySet { frictionloss_slots_reserved: scene_has_joint_friction, constraint_caps_dirty: false, - multibody_info: Tensor::vector(backend, &all_infos, storage).unwrap(), + multibody_info: Tensor::vector(backend, &all_infos, storage | BufferUsages::COPY_SRC) + .unwrap(), max_contact_constraints: Tensor::scalar( backend, 0u32, @@ -614,7 +635,17 @@ impl GpuMultibodySet { .unwrap(), 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: { + let cap = body_to_link_cap as usize; + let nb = num_batches as usize; + let mut interleaved = vec![[u32::MAX, u32::MAX]; all_body_to_link.len()]; + for local in 0..cap { + for b in 0..nb { + interleaved[local * nb + b] = all_body_to_link[b * cap + local]; + } + } + Tensor::vector(backend, &interleaved, storage).unwrap() + }, body_to_link_host: all_body_to_link, body_to_link_cap, motor_delay_state: Tensor::vector( @@ -653,31 +684,19 @@ impl GpuMultibodySet { scatter_caches: Vec::new(), contact_constraints: Tensor::vector( backend, - vec![ - MultibodyContactConstraint::default(); - (contact_cons_cap * num_batches) as usize - ], - storage, + vec![MultibodyContactConstraint::default(); contact_cons_cap as usize], + storage | BufferUsages::COPY_SRC, ) .unwrap(), old_contact_constraints: Tensor::vector( backend, - vec![ - MultibodyContactConstraint::default(); - (contact_cons_cap * num_batches) as usize - ], - storage, - ) - .unwrap(), - contact_constraint_jacs: Tensor::vector( - backend, - vec![0.0f32; (contact_cons_col_cap * num_batches) as usize], + vec![MultibodyContactConstraint::default(); contact_cons_cap as usize], storage, ) .unwrap(), - contact_constraint_columns: Tensor::vector( + contact_jac_cols: Tensor::vector( backend, - vec![0.0f32; (contact_cons_col_cap * num_batches) as usize], + vec![0.0f32; 2 * contact_cons_col_cap as usize], storage, ) .unwrap(), @@ -713,12 +732,6 @@ impl GpuMultibodySet { // Impulse-joint buffers are sized for "no MB-touching joints" by // default — `set_impulse_joints` resizes them at pipeline build // time when the host has actually counted the joints. - mb_imp_joint_count: Tensor::vector( - backend, - vec![0u32; num_batches as usize], - storage | BufferUsages::UNIFORM, - ) - .unwrap(), mb_imp_joint_builders: Tensor::vector( backend, vec![::zeroed(); num_batches as usize], @@ -756,8 +769,13 @@ impl GpuMultibodySet { max_joint_constraints: max_mb_joint_constraints, joint_constraints_per_batch: cons_cap, joint_constraint_columns_per_batch: cons_col_cap, - contact_constraints_per_batch: contact_cons_cap, - contact_constraint_columns_per_batch: contact_cons_col_cap, + contact_constraints_capacity: contact_cons_cap, + mb_cons_demand: Tensor::vector( + backend, + &[0u32], + storage | BufferUsages::COPY_SRC, + ) + .unwrap(), num_solver_iterations: 4, num_internal_pgs_iterations: 1, diff --git a/src_rbd/dynamics/multibody/multibody_set.rs b/src_rbd/dynamics/multibody/multibody_set.rs index cd90f43f..9e8bdca4 100644 --- a/src_rbd/dynamics/multibody/multibody_set.rs +++ b/src_rbd/dynamics/multibody/multibody_set.rs @@ -77,10 +77,8 @@ pub struct GpuMultibodySet { pub(super) num_active_multibodies: u32, pub(super) links_per_batch: u32, pub(super) dofs_per_batch: u32, - pub(super) jacobian_entries_per_batch: u32, pub(super) mass_matrix_entries_per_batch: u32, 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 @@ -190,20 +188,12 @@ pub struct GpuMultibodySet { /// Snapshot of `contact_constraints` taken at the start of the step; the /// warmstart transfer matches this frame's slots against it. pub(super) old_contact_constraints: Tensor, - /// Per-constraint `Jᵀ` row (length `ndofs`) — the multibody side's - /// contribution to the constraint Jacobian. - pub(super) contact_constraint_jacs: Tensor, - /// Per-constraint M⁻¹·Jᵀ column (length `ndofs`). - pub(super) contact_constraint_columns: Tensor, + pub(super) contact_jac_cols: Tensor, /// Per-multibody Delassus blocks (`MAX_MB_CONTACT_CONSTRAINTS_PER_MB²` /// floats each) only allocated when the total multibody count is at most /// [`MAX_DELASSUS_MULTIBODIES`]. pub(super) contact_delassus: Option>, - /// Per-batch number of multibody-touching impulse joints (body1 OR body2 - /// part of any multibody). - pub(super) mb_imp_joint_count: Tensor, - /// Per-batch slab of impulse-joint builder descriptors. pub(super) mb_imp_joint_builders: Tensor, /// Per-batch slab of axis constraints. pub(super) mb_imp_joint_constraints: Tensor, @@ -238,8 +228,8 @@ pub struct GpuMultibodySet { /// mirror). Stored so `RbdState` can rebuild its `BatchIndices` when caps change. pub(super) joint_constraints_per_batch: u32, pub(super) joint_constraint_columns_per_batch: u32, - pub(super) contact_constraints_per_batch: u32, - pub(super) contact_constraint_columns_per_batch: u32, + pub(super) contact_constraints_capacity: u32, + pub(super) mb_cons_demand: Tensor, /// Number of solver iterations to run on `joint_constraints` per `step()`. pub(super) num_solver_iterations: u32, @@ -662,7 +652,7 @@ impl GpuMultibodySet { set: &crate::rapier::dynamics::MultibodyJointSet, bodies: &crate::rapier::dynamics::RigidBodySet, ) -> Result<(), GpuBackendError> { - let base = (env * self.links_per_batch) as usize; + let nb = self.num_batches as usize; let mut offset = 0usize; for mb in set.multibodies() { // Mirror `from_rapier`'s fixed-root handling: a non-dynamic root has @@ -673,13 +663,17 @@ impl GpuMultibodySet { .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 { + let Some(entry) = self.links_static_mirror.get_mut(offset * nb + env as usize) + else { return Ok(()); }; let mut data = convert_generic_joint(link.joint().data); if link_idx == 0 && !root_is_dynamic { data.locked_axes = 0x3f; } + for axis in 0..6 { + data.motors[axis].impulse = entry.data.motors[axis].impulse; + } entry.data = data; offset += 1; } @@ -701,27 +695,16 @@ impl GpuMultibodySet { dst.multibodies_batch_capacity = self.multibodies_per_batch; dst.multibodies_len = self.num_active_multibodies; dst.links_batch_capacity = self.links_per_batch; - dst.jacobians_batch_capacity = self.jacobian_entries_per_batch; - dst.mass_matrix_batch_capacity = self.mass_matrix_entries_per_batch; dst.coriolis_batch_capacity = self.coriolis_entries_per_batch; - dst.i_coriolis_dt_batch_capacity = self.i_coriolis_dt_entries_per_batch; dst.dof_batch_capacity = self.dofs_per_batch; dst.mb_joint_constraints_batch_capacity = self.joint_constraints_per_batch; dst.mb_joint_constraint_columns_batch_capacity = self.joint_constraint_columns_per_batch; - dst.mb_contact_constraints_batch_capacity = self.contact_constraints_per_batch; - dst.mb_contact_constraint_columns_batch_capacity = - self.contact_constraint_columns_per_batch; + dst.mb_contact_constraints_capacity = self.contact_constraints_capacity; dst.mb_imp_joints_batch_capacity = self.mb_imp_joints_per_batch.max(1); - dst.mb_imp_joint_constraints_batch_capacity = self.mb_imp_joint_constraints_per_batch; - dst.mb_imp_joint_jacobians_batch_capacity = self.mb_imp_joint_jacobians_per_batch; - dst.mb_imp_joint_color_groups_batch_capacity = self.mb_imp_joint_num_colors.max(1); dst.mb_max_ndofs = self.max_ndofs; dst.mb_max_links = self.max_links; dst.mb_max_joint_constraints = self.max_joint_constraints; dst.mb_pack_lanes = self.pack_lanes(); - dst.coriolis_w_section_offset = self.coriolis_entries_per_batch * self.num_batches; - dst.i_coriolis_dt_section_offset = 2 * self.coriolis_entries_per_batch * self.num_batches; - dst.dof_damping_section_offset = self.dofs_per_batch * self.num_batches; // Implicit coriolis needs two matrices: the coriolis-augmented one (acc // section) for the acceleration solve, the plain one for constraints. // With the flag off, a single plain matrix serves both. @@ -796,9 +779,68 @@ impl GpuMultibodySet { self.warmstart_coefficient } - /// Per-batch stride of [`Self::contact_constraints`]. - pub fn contact_constraints_per_batch(&self) -> u32 { - self.contact_constraints_per_batch + pub fn contact_constraints_capacity(&self) -> u32 { + self.contact_constraints_capacity + } + pub fn body_to_link(&self) -> &Tensor<[u32; 2]> { + &self.body_to_link + } + + pub fn mb_cons_demand(&self) -> &Tensor { + &self.mb_cons_demand + } + pub(crate) fn min_contact_slab_capacity(&self) -> u32 { + (self.num_active_multibodies * self.num_batches) + .saturating_mul(crate::shaders::dynamics::MB_CONS_SLOT_RESERVE) + .max(64) + } + pub async fn debug_cons_layout( + &self, + backend: &GpuBackend, + ) -> (Vec<(u32, u32, u32, u32)>, u32) { + let infos: Vec = backend + .slow_read_vec(self.multibody_info.buffer()) + .await + .unwrap_or_default(); + let demand: Vec = backend + .slow_read_vec(self.mb_cons_demand.buffer()) + .await + .unwrap_or_default(); + let n = (self.num_active_multibodies * self.num_batches) as usize; + let out = infos + .iter() + .take(n) + .map(|i| { + ( + i.contact_constraint_start, + i.contact_constraint_count, + i.batch_contacts_start, + i.batch_contacts_len, + ) + }) + .collect(); + (out, demand.first().copied().unwrap_or(0)) + } + pub(crate) fn resize_contact_slabs(&mut self, backend: &GpuBackend, new_capacity: u32) { + use khal::BufferUsages; + let storage = BufferUsages::STORAGE | BufferUsages::COPY_DST; + let n = new_capacity.max(1); + let cols = n.saturating_mul(self.dofs_per_batch.max(1)); + self.contact_constraints = Tensor::vector( + backend, + vec![crate::shaders::dynamics::MultibodyContactConstraint::default(); n as usize], + storage | BufferUsages::COPY_SRC, + ) + .unwrap(); + self.old_contact_constraints = Tensor::vector( + backend, + vec![crate::shaders::dynamics::MultibodyContactConstraint::default(); n as usize], + storage, + ) + .unwrap(); + self.contact_jac_cols = + Tensor::vector(backend, vec![0.0f32; 2 * cols as usize], storage).unwrap(); + self.contact_constraints_capacity = n; } /// The per-multibody bank of unit (1-DoF) joint limit / motor constraints. @@ -1066,16 +1108,8 @@ impl GpuMultibodySet { .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 { - &self.contact_constraint_jacs - } - - /// Per-constraint `M⁻¹·Jᵀ` columns of the contact constraints, laid out - /// like [`Self::contact_constraint_jacs`]. - pub fn contact_constraint_columns(&self) -> &Tensor { - &self.contact_constraint_columns + pub fn contact_jac_cols(&self) -> &Tensor { + &self.contact_jac_cols } /// Per-link `SPATIAL_DIM × ndofs` column-major body jacobians, indexed from diff --git a/src_rbd/dynamics/multibody/multibody_solver.rs b/src_rbd/dynamics/multibody/multibody_solver.rs index 4270ba68..ff39f1ef 100644 --- a/src_rbd/dynamics/multibody/multibody_solver.rs +++ b/src_rbd/dynamics/multibody/multibody_solver.rs @@ -12,7 +12,8 @@ use crate::shaders::dynamics::{ GpuMbIntegrateVelocities, GpuMbRefreshJointConstraints, GpuMbRemoveImpulseJointConstraintBias, GpuMbSeedContactRestitution, GpuMbSenseContactImpulses, GpuMbSnapshotContactWarmstart, GpuMbSolveConstraints, GpuMbSolveContactsDelassus, GpuMbSolveImpulseJointConstraints, - GpuMbSolveJoints, GpuMbStashContactsLen, GpuMbTransferContactWarmstart, + GpuMbSolveJoints, GpuMbCountContactConstraints, GpuMbConsOffsetsScan, GpuMbSavePrevConsBounds, + GpuMbStashContactsLen, GpuMbTransferContactWarmstart, GpuMbUpdateImpulseJointConstraints, GpuMbWarmstartContactConstraints, Velocity, WorldMassProperties, }; @@ -66,6 +67,9 @@ pub struct GpuMultibodySolver { solve_contacts_delassus: GpuMbSolveContactsDelassus, /// Snapshot the contact impulses once per frame, for the cross-frame match. snapshot_contact_warmstart: GpuMbSnapshotContactWarmstart, + save_prev_cons_bounds: GpuMbSavePrevConsBounds, + count_contact_constraints: GpuMbCountContactConstraints, + cons_offsets_scan: GpuMbConsOffsetsScan, /// Carry the snapshotted impulses over to this frame's matching contacts. transfer_contact_warmstart: GpuMbTransferContactWarmstart, /// Copy `contacts_len[batch]` into each `MultibodyInfo` once per step so @@ -104,6 +108,7 @@ pub struct MultibodySolverArgs<'a> { pub contacts: &'a Tensor, /// Per-batch contact count (parallel to `contacts`). pub contacts_len: &'a Tensor, + pub contact_offsets: &'a Tensor, /// Free-body solver velocities (updated in place by `solve_contact_constraints`). pub solver_vels: &'a mut Tensor, /// Shared `BatchIndices` uniform — per-batch caps and packed-section @@ -157,12 +162,15 @@ impl GpuMultibodySolver { // Flat (slot, multibody, batch) grid. { let mut pass = encoder.begin_pass("[RBD] mbi/snapshot", timestamps.as_deref_mut()); - let total_slots = mb.num_active_multibodies - * mb.num_batches - * crate::shaders::dynamics::MAX_MB_CONTACT_CONSTRAINTS_PER_MB; + self.save_prev_cons_bounds.call( + &mut pass, + mb.flat_mb_dispatch(), + &mut mb.multibody_info, + args.batch_indices, + )?; self.snapshot_contact_warmstart.call( &mut pass, - [total_slots, 1, 1], + mb.contact_constraints_capacity, &mb.contact_constraints, &mut mb.old_contact_constraints, args.batch_indices, @@ -190,6 +198,22 @@ impl GpuMultibodySolver { mb.flat_mb_dispatch(), &mut mb.multibody_info, args.contacts_len, + args.contact_offsets, + args.batch_indices, + )?; + self.count_contact_constraints.call( + pass, + mb.flat_mb_dispatch(), + &mut mb.multibody_info, + args.contacts, + &mb.body_to_link, + args.batch_indices, + )?; + self.cons_offsets_scan.call( + pass, + 1u32, + &mut mb.multibody_info, + &mut mb.mb_cons_demand, args.batch_indices, )?; Ok(()) @@ -289,7 +313,7 @@ impl GpuMultibodySolver { args.mb_sweep_indirect, &mb.multibody_info, &mut mb.contact_constraints, - &mb.contact_constraint_jacs, + &mb.contact_jac_cols, &mb.dof_state, args.solver_vels, args.batch_indices, @@ -309,7 +333,7 @@ impl GpuMultibodySolver { args.mb_sweep_indirect, &mb.multibody_info, &mb.contact_constraints, - &mb.contact_constraint_columns, + &mb.contact_jac_cols, &mut mb.dof_state, args.solver_vels, args.batch_indices, @@ -415,8 +439,7 @@ impl GpuMultibodySolver { &mb.mass_matrices, &mb.lu_pivots, &mut mb.contact_constraints, - &mut mb.contact_constraint_jacs, - &mut mb.contact_constraint_columns, + &mut mb.contact_jac_cols, &mb.links_static, &mb.body_jacobians, args.batch_indices, @@ -443,8 +466,7 @@ impl GpuMultibodySolver { args.mb_sweep_indirect, &mb.multibody_info, &mb.contact_constraints, - &mb.contact_constraint_jacs, - &mb.contact_constraint_columns, + &mb.contact_jac_cols, delassus, args.batch_indices, )?; @@ -486,8 +508,7 @@ impl GpuMultibodySolver { args.mb_sweep_indirect, &mb.multibody_info, &mut mb.contact_constraints, - &mb.contact_constraint_jacs, - &mb.contact_constraint_columns, + &mb.contact_jac_cols, delassus, use_bias, args.batch_indices, @@ -503,8 +524,7 @@ impl GpuMultibodySolver { &mut mb.joint_constraints, &mb.joint_constraint_columns, &mut mb.contact_constraints, - &mb.contact_constraint_jacs, - &mb.contact_constraint_columns, + &mb.contact_jac_cols, use_bias, args.batch_indices, &mb.max_contact_constraints, @@ -522,8 +542,7 @@ impl GpuMultibodySolver { &mut mb.joint_constraints, &mb.joint_constraint_columns, &mut mb.contact_constraints, - &mb.contact_constraint_jacs, - &mb.contact_constraint_columns, + &mb.contact_jac_cols, use_bias, args.batch_indices, &mb.max_contact_constraints, @@ -558,7 +577,7 @@ impl GpuMultibodySolver { // Multibody-touching impulse joints — generic (rb-mb / mb-mb) // constraints. if mb.mb_imp_joints_per_batch > 0 { - let imp_dispatch = [mb.mb_imp_joints_per_batch, mb.num_batches, 1]; + let imp_dispatch = [mb.mb_imp_joints_per_batch * mb.num_batches, 1, 1]; self.update_impulse_joint_constraints.call( pass, imp_dispatch, @@ -597,8 +616,8 @@ impl GpuMultibodySolver { // One workgroup (MB_LU_LANES threads) per joint; thread // count = joints-in-largest-color × workgroup size. [ - mb.mb_imp_joint_max_color_group_len * MB_LU_LANES, - mb.num_batches, + mb.mb_imp_joint_max_color_group_len * mb.num_batches * MB_LU_LANES, + 1, 1, ], &mb.mb_imp_joint_builders, @@ -673,13 +692,12 @@ impl GpuMultibodySolver { let solve_dispatch = [mb.multibodies_per_batch * MB_LU_LANES, mb.num_batches, 1]; self.dispatch_solve(pass, mb, args, solve_dispatch, 0)?; if mb.mb_imp_joints_per_batch > 0 { - let imp_dispatch = [mb.mb_imp_joints_per_batch, mb.num_batches, 1]; + let imp_dispatch = [mb.mb_imp_joints_per_batch * mb.num_batches, 1, 1]; self.remove_impulse_joint_constraint_bias.call( pass, imp_dispatch, &mb.mb_imp_joint_builders, &mut mb.mb_imp_joint_constraints, - &mb.mb_imp_joint_count, args.batch_indices, )?; // Final stabilization sweep WITHOUT bias — colored, one @@ -690,8 +708,8 @@ impl GpuMultibodySolver { // One workgroup (MB_LU_LANES threads) per joint; thread // count = joints-in-largest-color × workgroup size. [ - mb.mb_imp_joint_max_color_group_len * MB_LU_LANES, - mb.num_batches, + mb.mb_imp_joint_max_color_group_len * mb.num_batches * MB_LU_LANES, + 1, 1, ], &mb.mb_imp_joint_builders, @@ -750,8 +768,7 @@ impl GpuMultibodySolver { args.mb_sweep_indirect, &mb.multibody_info, &mut mb.contact_constraints, - &mb.contact_constraint_jacs, - &mb.contact_constraint_columns, + &mb.contact_jac_cols, args.batch_indices, &mb.max_contact_constraints, &mut mb.dof_state, diff --git a/src_rbd/dynamics/solver.rs b/src_rbd/dynamics/solver.rs index 28c08271..752ac401 100644 --- a/src_rbd/dynamics/solver.rs +++ b/src_rbd/dynamics/solver.rs @@ -78,7 +78,7 @@ pub struct SolverArgs<'a> { pub contacts: &'a Tensor, /// Number of contacts (per batch). pub contacts_len: &'a Tensor, - /// Indirect dispatch arguments based on contact count. + pub contact_offsets: &'a Tensor, pub contacts_len_indirect: &'a Tensor<[u32; 3]>, /// Solver constraints (output from constraint initialization). pub constraints: &'a mut Tensor, @@ -120,10 +120,7 @@ pub struct SolverArgs<'a> { /// All constraints of all the bodies part of the same multibody are in the same list associated /// to the multibody’s root. pub body_constraint_ids: &'a mut Tensor, - /// Per-batch per-color exclusive prefix sums over the color-bucketed - /// constraint ids (stride `BatchIndices::solver_color_buckets_stride`). - pub color_bucket_starts: &'a Tensor, - /// Constraint ids bucket-sorted by color (contacts layout). + pub color_buckets: &'a Tensor, pub color_sorted_ids: &'a Tensor, /// Per-color-index uniform tensors: `color_uniforms[c] == c`. pub color_uniforms: &'a [Tensor], @@ -167,7 +164,7 @@ impl GpuSolver { // Cleanup zeroes body_constraint_counts, solver_vels, vels, mprops. self.cleanup.call( pass, - [args.num_colliders, args.num_batches, 1], + args.num_colliders * args.num_batches, args.body_constraint_counts, args.solver_vels, args.vels, @@ -182,7 +179,7 @@ impl GpuSolver { // below reads `solver_body_poses` so we barrier before it. self.init_solver_bodies.call( pass, - [args.num_colliders, args.num_batches, 1], + args.num_colliders * args.num_batches, args.body_poses, args.local_mprops, args.solver_body_poses, @@ -199,6 +196,7 @@ impl GpuSolver { args.contacts, args.constraints, args.constraint_builders, + args.contact_offsets, args.collider_world_poses, args.solver_body_poses, args.vels, @@ -216,7 +214,7 @@ impl GpuSolver { args.body_constraint_counts, args.body_group, args.mprops, - args.contacts_len, + args.contact_offsets, args.batch_indices, )?; @@ -225,7 +223,7 @@ impl GpuSolver { pass, prefix_sum_workspace, args.body_constraint_counts, - args.num_batches, + 1, )?; self.sort_constraints.call( @@ -234,7 +232,7 @@ impl GpuSolver { args.body_constraint_counts, args.mprops, args.contacts, - args.contacts_len, + args.contact_offsets, args.body_constraint_ids, args.body_group, args.batch_indices, @@ -274,7 +272,7 @@ impl GpuSolver { if !skip_rb { self.init_solver_vels_inc.call( &mut pass, - [args.num_colliders, args.num_batches, 1], + args.num_colliders * args.num_batches, args.solver_vels_inc, args.mprops, args.sim_params, @@ -295,6 +293,7 @@ impl GpuSolver { mprops: args.mprops, contacts: args.contacts, contacts_len: args.contacts_len, + contact_offsets: args.contact_offsets, solver_vels: &mut *args.solver_vels, batch_indices: args.batch_indices, gravity: args.gravity, @@ -320,6 +319,7 @@ impl GpuSolver { mprops: args.mprops, contacts: args.contacts, contacts_len: args.contacts_len, + contact_offsets: args.contact_offsets, solver_vels: &mut *args.solver_vels, batch_indices: args.batch_indices, gravity: args.gravity, @@ -346,7 +346,7 @@ impl GpuSolver { encoder.begin_pass("[RBD] slv/rb-apply-inc", timestamps.as_deref_mut()); self.apply_solver_vels_inc.call( &mut pass, - [args.num_colliders, args.num_batches, 1], + args.num_colliders * args.num_batches, args.solver_vels, args.solver_vels_inc, args.batch_indices, @@ -365,6 +365,7 @@ impl GpuSolver { mprops: args.mprops, contacts: args.contacts, contacts_len: args.contacts_len, + contact_offsets: args.contact_offsets, solver_vels: &mut *args.solver_vels, batch_indices: args.batch_indices, gravity: args.gravity, @@ -390,7 +391,7 @@ impl GpuSolver { args.contacts_len_indirect, args.constraints, args.constraint_builders, - args.contacts_len, + args.contact_offsets, args.solver_body_poses, args.sim_params, args.batch_indices, @@ -403,7 +404,7 @@ impl GpuSolver { } else if args.colorless_warmstart { self.warmstart_without_colors.call( pass, - [args.num_colliders, args.num_batches, 1], + args.num_colliders * args.num_batches, args.body_constraint_counts, args.body_constraint_ids, args.constraints, @@ -419,7 +420,7 @@ impl GpuSolver { [64, args.num_batches, 1], args.constraints, args.solver_vels, - args.color_bucket_starts, + args.color_buckets, args.color_sorted_ids, &args.color_uniforms[args.num_colors as usize], args.batch_indices, @@ -432,7 +433,7 @@ impl GpuSolver { args.contacts_len_indirect, args.constraints, args.solver_vels, - args.color_bucket_starts, + args.color_buckets, args.color_sorted_ids, &args.color_uniforms[c as usize], args.batch_indices, @@ -458,7 +459,7 @@ impl GpuSolver { [64, args.num_batches, 1], args.constraints, args.solver_vels, - args.color_bucket_starts, + args.color_buckets, args.color_sorted_ids, &args.color_uniforms[args.num_colors as usize], args.batch_indices, @@ -472,7 +473,7 @@ impl GpuSolver { args.contacts_len_indirect, args.constraints, args.solver_vels, - args.color_bucket_starts, + args.color_buckets, args.color_sorted_ids, &args.color_uniforms[c as usize], args.batch_indices, @@ -496,7 +497,7 @@ impl GpuSolver { encoder.begin_pass("[RBD] slv/rb-integrate", timestamps.as_deref_mut()); self.integrate_linearized.call( &mut pass, - [args.num_colliders, args.num_batches, 1], + args.num_colliders * args.num_batches, args.solver_body_poses, args.solver_vels, args.sim_params, @@ -518,7 +519,7 @@ impl GpuSolver { args.contacts_len_indirect, args.constraints, args.constraint_builders, - args.contacts_len, + args.contact_offsets, args.solver_body_poses, args.sim_params, args.batch_indices, @@ -533,7 +534,7 @@ impl GpuSolver { [64, args.num_batches, 1], args.constraints, args.solver_vels, - args.color_bucket_starts, + args.color_buckets, args.color_sorted_ids, &args.color_uniforms[args.num_colors as usize], args.batch_indices, @@ -547,7 +548,7 @@ impl GpuSolver { args.contacts_len_indirect, args.constraints, args.solver_vels, - args.color_bucket_starts, + args.color_buckets, args.color_sorted_ids, &args.color_uniforms[c as usize], args.batch_indices, @@ -570,7 +571,7 @@ impl GpuSolver { let mut pass = encoder.begin_pass("[RBD] slv/finalize", timestamps); self.finalize.call( &mut pass, - [args.num_colliders, args.num_batches, 1], + args.num_colliders * args.num_batches, args.vels, args.solver_vels, args.body_poses, diff --git a/src_rbd/dynamics/warmstart.rs b/src_rbd/dynamics/warmstart.rs index b0b07f3d..5b4e271c 100644 --- a/src_rbd/dynamics/warmstart.rs +++ b/src_rbd/dynamics/warmstart.rs @@ -26,8 +26,7 @@ pub struct GpuWarmstart { /// /// Contains buffers for both old (previous frame) and new (current frame) constraint data. pub struct WarmstartArgs<'a> { - /// Number of contacts in current frame. - pub contacts_len: &'a Tensor, + pub contact_offsets: &'a Tensor, /// Constraint counts per body from previous frame. pub old_body_constraint_counts: &'a Tensor, /// Constraint IDs per body from previous frame. @@ -48,8 +47,7 @@ pub struct WarmstartArgs<'a> { /// Arguments for the coloring seed dispatch. pub struct SeedColorsArgs<'a> { - /// Number of contacts in current frame. - pub contacts_len: &'a Tensor, + pub contact_offsets: &'a Tensor, /// Constraint counts per body from previous frame. pub old_body_constraint_counts: &'a Tensor, /// Constraint IDs per body from previous frame. @@ -86,7 +84,7 @@ impl GpuWarmstart { args.old_constraint_builders, args.new_constraints, args.new_constraint_builders, - args.contacts_len, + args.contact_offsets, args.batch_indices, ) } @@ -108,7 +106,7 @@ impl GpuWarmstart { args.old_constraints_colors, args.constraints_colors, args.colored, - args.contacts_len, + args.contact_offsets, args.batch_indices, ) } diff --git a/src_rbd/pipeline/insertion_removal.rs b/src_rbd/pipeline/insertion_removal.rs index f8bf0991..dcef8bc9 100644 --- a/src_rbd/pipeline/insertion_removal.rs +++ b/src_rbd/pipeline/insertion_removal.rs @@ -65,25 +65,12 @@ impl RbdState { let all_collision_groups = vec![none_groups; num_bodies_total]; let all_vels = vec![GpuVelocity::default(); num_bodies_total]; - // body_group: per-batch local indices (free bodies map to themselves). - let mut all_body_group = Vec::with_capacity(num_bodies_total); - for _ in 0..num_batches { - for b in 0..capacity_per_batch { - all_body_group.push(b); - } - } + let all_body_group: Vec = (0..num_bodies_total as u32).collect(); - // collider_parent: identity within each batch initially (no body is - // active yet). `append_bodies` overwrites the active prefix. - let mut all_collider_parent = Vec::with_capacity(num_bodies_total); - let mut all_pair_filter = Vec::with_capacity(num_bodies_total); - for _ in 0..num_batches { - for c in 0..capacity_per_batch { - all_collider_parent.push(c); - // Identity parent, no multibody key. - all_pair_filter.push([c, 0u32]); - } - } + let all_collider_parent: Vec = (0..num_bodies_total as u32).collect(); + let all_pair_filter: Vec<[u32; 2]> = (0..num_bodies_total as u32) + .map(|i| [i / num_batches, 0u32]) + .collect(); // Empty joints / multibodies, one (empty) environment per batch. let empty_joints = ImpulseJointSet::new(); @@ -100,7 +87,12 @@ impl RbdState { let mb_refs: Vec<_> = (0..num_batches as usize) .map(|_| (&empty_mb, &empty_body_ids, &empty_bodies)) .collect(); - let mut mb = GpuMultibodySet::from_rapier(backend, &mb_refs, capacity_per_batch); + let mut mb = GpuMultibodySet::from_rapier( + backend, + &mb_refs, + capacity_per_batch, + capacities.mb_contact_constraints_capacity, + ); mb.set_constraint_softness(backend, &base_sim_params); mb }; @@ -127,23 +119,16 @@ impl RbdState { let collision_pairs = Tensor::vector_uninit(backend, collisions_capacity * num_batches, storage).unwrap(); - let collision_pairs_len = Tensor::vector_uninit( + let collision_pairs_len = Tensor::vector( backend, - num_batches, + &[0u32], BufferUsages::STORAGE | BufferUsages::COPY_SRC, ) .unwrap(); - let collision_pairs_len_max = - Tensor::vector_uninit(backend, 1, BufferUsages::STORAGE | BufferUsages::COPY_SRC) - .unwrap(); - let num_batches_uniform = Tensor::scalar( - backend, - collision_pairs_len.layout().into(), - BufferUsages::STORAGE | BufferUsages::UNIFORM, - ) - .unwrap(); - // Two-element readback: the (max) collision-pair count and the uncolored count. - let resize_readback = GpuReadback::new(backend, 2).unwrap(); + #[cfg(feature = "dim3")] + let resize_readback = GpuReadback::new(backend, 4).unwrap(); + #[cfg(not(feature = "dim3"))] + let resize_readback = GpuReadback::new(backend, 3).unwrap(); let collision_pairs_indirect = Tensor::scalar_uninit(backend, BufferUsages::STORAGE | BufferUsages::INDIRECT).unwrap(); let contacts = @@ -162,9 +147,9 @@ impl RbdState { Tensor::scalar_uninit(backend, BufferUsages::STORAGE | BufferUsages::INDIRECT).unwrap(); let pfm_pairs = Tensor::vector_uninit(backend, collisions_capacity * num_batches, storage).unwrap(); - let pfm_pairs_len = Tensor::vector_uninit( + let pfm_pairs_len = Tensor::vector( backend, - num_batches, + &[0u32], BufferUsages::STORAGE | BufferUsages::COPY_SRC, ) .unwrap(); @@ -189,16 +174,26 @@ impl RbdState { let constraints_rands = Tensor::vector_uninit(backend, collisions_capacity * num_batches, storage).unwrap(); let color_buckets_stride = capacities.solver_colors + 3; - let color_bucket_counts = - Tensor::vector_uninit(backend, color_buckets_stride * num_batches, storage).unwrap(); - let color_bucket_starts = - Tensor::vector_uninit(backend, color_buckets_stride * num_batches, storage).unwrap(); - let color_bucket_cursors = + let color_buckets = Tensor::vector_uninit(backend, color_buckets_stride * num_batches, storage).unwrap(); + let contact_offsets = Tensor::vector( + backend, + vec![0u32; num_batches as usize + 3], + storage, + ) + .unwrap(); + let pair_batch_counts = + Tensor::vector(backend, vec![0u32; num_batches as usize], storage).unwrap(); + let pfm_batch_counts = + Tensor::vector(backend, vec![0u32; num_batches as usize], storage).unwrap(); let color_sorted_ids = Tensor::vector_uninit(backend, collisions_capacity * num_batches, storage).unwrap(); - let old_constraints_counts = - Tensor::vector_uninit(backend, num_colliders_per_batch * num_batches, storage).unwrap(); + let old_constraints_counts = Tensor::vector( + backend, + vec![0u32; (num_colliders_per_batch * num_batches) as usize], + storage, + ) + .unwrap(); let new_constraints_counts = Tensor::vector_uninit(backend, num_colliders_per_batch * num_batches, storage).unwrap(); let old_body_constraint_ids = @@ -212,8 +207,8 @@ impl RbdState { BufferUsages::STORAGE }; - let contacts_per_batch_cpu = collisions_capacity; - let collision_pairs_per_batch_cpu = collisions_capacity; + let contacts_capacity_cpu = collisions_capacity * num_batches; + let collision_pairs_capacity_cpu = collisions_capacity * num_batches; #[allow(unused_mut)] // Only mutated with the dim3 (multibody) feature. let mut bi = BatchIndices { num_batches, @@ -221,9 +216,8 @@ impl RbdState { // No body is active initially; bodies are added later via `append_bodies`. colliders_len: 0, bodies_len: 0, - collision_pairs_batch_capacity: collision_pairs_per_batch_cpu, - contacts_batch_capacity: contacts_per_batch_cpu, - impulse_joints_batch_capacity: joints.joints_per_batch(), + collision_pairs_capacity: collision_pairs_capacity_cpu, + contacts_capacity: contacts_capacity_cpu, impulse_joints_len: joints.num_active_joints(), solver_color_buckets_stride: color_buckets_stride, ..Default::default() @@ -274,17 +268,20 @@ impl RbdState { collider_materials, collision_pairs, collision_pairs_len, - collision_pairs_len_max, - num_batches_uniform, resize_readback, collision_pairs_indirect, - contacts_per_batch_cpu, - collision_pairs_per_batch_cpu, + contacts_capacity_cpu, + collision_pairs_capacity_cpu, collision_pairs_len_cpu: 0, + #[cfg(feature = "dim3")] + mb_cons_demand_cpu: 0, batch_indices, contacts, contacts_len, contacts_indirect, + contact_offsets, + pair_batch_counts, + pfm_batch_counts, mb_sweep_indirect, pfm_pairs, pfm_pairs_len, @@ -299,9 +296,7 @@ impl RbdState { old_constraints_colors, colored, constraints_rands, - color_bucket_counts, - color_bucket_starts, - color_bucket_cursors, + color_buckets, color_sorted_ids, curr_color: Tensor::scalar( backend, @@ -328,6 +323,7 @@ impl RbdState { new_body_constraint_ids, color_uniforms: Vec::new(), prefix_sum_workspace: PrefixSumWorkspace::default(), + bucket_prefix_workspace: PrefixSumWorkspace::default(), lbvh: LbvhState::with_usages(backend, lbvh_usages), max_colors: capacities.solver_colors, rb_contacts_inert: false, @@ -442,32 +438,59 @@ impl RbdState { // The incremental path attaches exactly one collider per body, so a // body's collider slot equals its body slot: `collider_parent` is the - // identity over the appended (env-local) range. - let parents: Vec = (active as u32..(active + bodies.len()) as u32).collect(); - // NOTE: appended bodies are free bodies (never multibody links). - let pair_filters: Vec<[u32; 2]> = parents.iter().map(|&p| [p, 0u32]).collect(); - - // Write the same body data into every batch's slot range so all - // environments share the same topology. - for batch_id in 0..self.num_batches as usize { - let base = (batch_id * cap + active) as u64; - backend.write_buffer(self.body_poses.buffer_mut(), base, &poses)?; - backend.write_buffer(self.solver_body_poses.buffer_mut(), base, &poses)?; - backend.write_buffer(self.collider_world_poses.buffer_mut(), base, &poses)?; - backend.write_buffer( - self.collider_local_poses.buffer_mut(), - base, - &collider_local_poses, - )?; - backend.write_buffer(self.collider_parent.buffer_mut(), base, &parents)?; - backend.write_buffer(self.pair_filter.buffer_mut(), base, &pair_filters)?; - backend.write_buffer(self.local_mprops.buffer_mut(), base, &local_mprops)?; - backend.write_buffer(self.mprops.buffer_mut(), base, &mprops)?; - backend.write_buffer(self.shapes.buffer_mut(), base, &shapes)?; - backend.write_buffer(self.collision_groups.buffer_mut(), base, &collision_groups)?; - backend.write_buffer(self.collider_materials.buffer_mut(), base, &materials)?; - backend.write_buffer(self.vels.buffer_mut(), base, &vels)?; + let nb = self.num_batches as usize; + let parents: Vec = ((active * nb) as u32..((active + bodies.len()) * nb) as u32) + .collect(); + let pair_filters: Vec<[u32; 2]> = (0..bodies.len() * nb) + .map(|i| [(active + i / nb) as u32, 0u32]) + .collect(); + + fn replicate(v: &[T], nb: usize) -> Vec { + let mut out = Vec::with_capacity(v.len() * nb); + for &x in v { + for _ in 0..nb { + out.push(x); + } + } + out } + let base = (active * nb) as u64; + backend.write_buffer(self.body_poses.buffer_mut(), base, &replicate(&poses, nb))?; + backend.write_buffer( + self.solver_body_poses.buffer_mut(), + base, + &replicate(&poses, nb), + )?; + backend.write_buffer( + self.collider_world_poses.buffer_mut(), + base, + &replicate(&poses, nb), + )?; + backend.write_buffer( + self.collider_local_poses.buffer_mut(), + base, + &replicate(&collider_local_poses, nb), + )?; + backend.write_buffer(self.collider_parent.buffer_mut(), base, &parents)?; + backend.write_buffer(self.pair_filter.buffer_mut(), base, &pair_filters)?; + backend.write_buffer( + self.local_mprops.buffer_mut(), + base, + &replicate(&local_mprops, nb), + )?; + backend.write_buffer(self.mprops.buffer_mut(), base, &replicate(&mprops, nb))?; + backend.write_buffer(self.shapes.buffer_mut(), base, &replicate(&shapes, nb))?; + backend.write_buffer( + self.collision_groups.buffer_mut(), + base, + &replicate(&collision_groups, nb), + )?; + backend.write_buffer( + self.collider_materials.buffer_mut(), + base, + &replicate(&materials, nb), + )?; + backend.write_buffer(self.vels.buffer_mut(), base, &replicate(&vels, nb))?; let new_active = (active + bodies.len()) as u32; self.num_active_colliders = new_active; @@ -488,7 +511,6 @@ impl RbdState { backend: &GpuBackend, local_indices: &[u32], ) -> Result, GpuBackendError> { - let cap = self.num_colliders_per_batch as usize; let mut remaps = Vec::new(); // Process local slots in descending order so removing one doesn't @@ -503,22 +525,22 @@ impl RbdState { crate::rapier::geometry::InteractionTestMode::And, ); + let nb = self.num_batches as usize; let staging_usages = BufferUsages::STORAGE | BufferUsages::COPY_SRC | BufferUsages::COPY_DST; let mut enc = backend.begin_encoding(); let mut any_copy = false; - let mut staging_pose = backend.uninit_buffer::(1, staging_usages)?; + let mut staging_pose = backend.uninit_buffer::(nb, staging_usages)?; let mut staging_local_mprops = - backend.uninit_buffer::(1, staging_usages)?; + backend.uninit_buffer::(nb, staging_usages)?; let mut staging_mprops = - backend.uninit_buffer::(1, staging_usages)?; - let mut staging_vels = backend.uninit_buffer::(1, staging_usages)?; - let mut staging_shapes = backend.uninit_buffer::(1, staging_usages)?; + backend.uninit_buffer::(nb, staging_usages)?; + let mut staging_vels = backend.uninit_buffer::(nb, staging_usages)?; + let mut staging_shapes = backend.uninit_buffer::(nb, staging_usages)?; let mut staging_groups = backend - .uninit_buffer::(1, staging_usages)?; + .uninit_buffer::(nb, staging_usages)?; let mut staging_materials = - backend.uninit_buffer::(1, staging_usages)?; - // Deferred `(buffer-kind, slot)` neutralisation writes. + backend.uninit_buffer::(nb, staging_usages)?; let mut neutralize: Vec = Vec::new(); for local in locals { @@ -528,47 +550,30 @@ impl RbdState { } let last = active - 1; - for batch in 0..self.num_batches as usize { - let hole_global = batch * cap + local; - let last_global = batch * cap + last; - - if local != last { - // Relocate the last active body into the freed slot. A staging - // buffer is used to avoid same-buffer overlapping copies. - macro_rules! relocate { - ($t:expr, $staging:expr) => {{ - enc.copy_buffer_to_buffer( - $t.buffer(), - last_global, - &mut $staging, - 0, - 1, - )?; - enc.copy_buffer_to_buffer( - &$staging, - 0, - $t.buffer_mut(), - hole_global, - 1, - )?; - }}; - } - any_copy = true; - relocate!(self.body_poses, staging_pose); - relocate!(self.solver_body_poses, staging_pose); - relocate!(self.collider_world_poses, staging_pose); - relocate!(self.collider_local_poses, staging_pose); - relocate!(self.local_mprops, staging_local_mprops); - relocate!(self.mprops, staging_mprops); - relocate!(self.vels, staging_vels); - relocate!(self.shapes, staging_shapes); - relocate!(self.collision_groups, staging_groups); - relocate!(self.collider_materials, staging_materials); + if local != last { + let hole_global = local * nb; + let last_global = last * nb; + macro_rules! relocate { + ($t:expr, $staging:expr) => {{ + enc.copy_buffer_to_buffer($t.buffer(), last_global, &mut $staging, 0, nb)?; + enc.copy_buffer_to_buffer(&$staging, 0, $t.buffer_mut(), hole_global, nb)?; + }}; } - - neutralize.push(last_global); + any_copy = true; + relocate!(self.body_poses, staging_pose); + relocate!(self.solver_body_poses, staging_pose); + relocate!(self.collider_world_poses, staging_pose); + relocate!(self.collider_local_poses, staging_pose); + relocate!(self.local_mprops, staging_local_mprops); + relocate!(self.mprops, staging_mprops); + relocate!(self.vels, staging_vels); + relocate!(self.shapes, staging_shapes); + relocate!(self.collision_groups, staging_groups); + relocate!(self.collider_materials, staging_materials); } + neutralize.push(last); + if local != last { remaps.push((last as u32, local as u32)); } @@ -583,21 +588,22 @@ impl RbdState { // they never participate in collisions even if a kernel scans up to // the per-batch capacity. Done after the relocation submit so the // copies read the pre-neutralisation data. - for last_global in neutralize { + for slot in neutralize { + let base = (slot * nb) as u64; backend.write_buffer( self.collision_groups.buffer_mut(), - last_global as u64, - &[none_groups], + base, + &vec![none_groups; nb], )?; backend.write_buffer( self.local_mprops.buffer_mut(), - last_global as u64, - &[GpuLocalMassProperties::default()], + base, + &vec![GpuLocalMassProperties::default(); nb], )?; backend.write_buffer( self.mprops.buffer_mut(), - last_global as u64, - &[GpuWorldMassProperties::default()], + base, + &vec![GpuWorldMassProperties::default(); nb], )?; } diff --git a/src_rbd/pipeline/rbd_state.rs b/src_rbd/pipeline/rbd_state.rs index 0ee03857..8a9a1252 100644 --- a/src_rbd/pipeline/rbd_state.rs +++ b/src_rbd/pipeline/rbd_state.rs @@ -20,7 +20,6 @@ use crate::utils::PrefixSumWorkspace; use khal::BufferUsages; use khal::backend::{Backend, GpuBackend, GpuReadback}; use std::time::Duration; -use vortx::shaders::linalg::Shape as TensorShape; use vortx::tensor::Tensor; /// Performance statistics collected during a physics simulation step. @@ -59,6 +58,7 @@ pub struct RbdCapacities { /// /// This may or may not be automatically resized depending on [`Self::collisions_resize_policy`]. pub collisions_capacity: u32, + pub mb_contact_constraints_capacity: u32, /// How internal collision buffers gets automatically resized (or not). /// /// Note that setting both [`Self::collisions_resize_policy`] and @@ -89,6 +89,7 @@ impl Default for RbdCapacities { batches: 1, body_capacity: 65536, collisions_capacity: 4096, + mb_contact_constraints_capacity: 256, collisions_resize_policy: RbdResizePolicy::Grow, solver_colors: 8, solver_colors_resize_policy: RbdResizePolicy::Grow, @@ -164,14 +165,6 @@ pub struct RbdState { pub(super) collision_pairs: Tensor, /// Per-batch live collision-pair counts (length `num_batches`). pub(super) collision_pairs_len: Tensor, - /// Single-element scratch holding the max of `collision_pairs_len` across all - /// batches, computed on the GPU (only used when `num_batches > 1`). - pub(super) collision_pairs_len_max: Tensor, - /// Cosine of the maximum angle between two contact normals for their - /// manifolds to be clustered together (see `gpu_reduce_contacts`). - /// `num_batches` as a uniform, the scan length for the max reduction. - pub(super) num_batches_uniform: Tensor, - /// Non-blocking readback of `[max collision_pairs_len, uncolored]` used by /// [`RbdPipeline::auto_resize_buffers`](crate::pipeline::RbdPipeline::auto_resize_buffers) /// to grow buffers without stalling. pub(super) resize_readback: GpuReadback, @@ -179,12 +172,12 @@ pub struct RbdState { /// CPU-side mirrors of the dynamic batch capacities. The capacity values /// live in the [`BatchIndices`] uniform; these mirrors let /// [`Self::rebuild_batch_indices`] re-emit it whenever a buffer grows. - pub(super) contacts_per_batch_cpu: u32, - pub(super) collision_pairs_per_batch_cpu: u32, - /// Most recently read live collision-pair count — the max across all batches, - /// harvested by the non-blocking readback in [`RbdPipeline::auto_resize_buffers`](crate::pipeline::RbdPipeline::auto_resize_buffers). + pub(super) contacts_capacity_cpu: u32, + pub(super) collision_pairs_capacity_cpu: u32, /// Surfaced in the viewer UI; lags the GPU by a frame or two like the resize. pub(super) collision_pairs_len_cpu: u32, + #[cfg(feature = "dim3")] + pub(super) mb_cons_demand_cpu: u32, /// Single uniform aggregating every per-batch capacity and packed-buffer /// section offset consumed by the compute kernels (multibody and RBD /// sides). Rebuilt by [`Self::rebuild_batch_indices`] whenever any of its @@ -196,6 +189,9 @@ pub struct RbdState { pub(super) contacts: Tensor, pub(super) contacts_len: Tensor, pub(super) contacts_indirect: Tensor<[u32; 3]>, + pub(super) contact_offsets: Tensor, + pub(super) pair_batch_counts: Tensor, + pub(super) pfm_batch_counts: Tensor, /// Workgroup grid for the per-multibody contact-constraint dispatches: /// `[multibodies_batch_capacity, num_batches, 1]`. pub(super) mb_sweep_indirect: Tensor<[u32; 3]>, @@ -211,15 +207,7 @@ pub struct RbdState { pub(super) old_constraints_colors: Tensor, pub(super) colored: Tensor, pub(super) constraints_rands: Tensor, - /// Per-batch per-color constraint counts (stride `max_colors + 3`), see - /// the `gpu_color_buckets_*` kernels. - pub(super) color_bucket_counts: Tensor, - /// Per-batch per-color exclusive prefix sums over the counts: color `c` - /// owns `color_sorted_ids[starts[c]..starts[c + 1]]`. - pub(super) color_bucket_starts: Tensor, - /// Scatter cursors (seeded from the starts each step). - pub(super) color_bucket_cursors: Tensor, - /// Constraint indices bucket-sorted by color (contacts layout). + pub(super) color_buckets: Tensor, pub(super) color_sorted_ids: Tensor, pub(super) curr_color: Tensor, /// Pre-built per-color-index uniforms: `color_uniforms[c] == c`. @@ -240,6 +228,7 @@ pub struct RbdState { /// assigned the same color. pub(super) body_group: Tensor, pub(super) prefix_sum_workspace: PrefixSumWorkspace, + pub(super) bucket_prefix_workspace: PrefixSumWorkspace, /// Maximum number of constraint colors the solver will iterate. pub(super) max_colors: u32, /// `true` when every body is either non-dynamic or multibody-controlled @@ -270,9 +259,8 @@ impl RbdState { colliders_batch_capacity: self.num_colliders_per_batch, colliders_len: self.num_active_colliders, bodies_len: self.num_active_bodies, - collision_pairs_batch_capacity: self.collision_pairs_per_batch_cpu, - contacts_batch_capacity: self.contacts_per_batch_cpu, - impulse_joints_batch_capacity: self.joints.joints_per_batch(), + collision_pairs_capacity: self.collision_pairs_capacity_cpu, + contacts_capacity: self.contacts_capacity_cpu, impulse_joints_len: self.joints.num_active_joints(), solver_color_buckets_stride: self.max_colors + 3, ..Default::default() @@ -416,6 +404,14 @@ impl RbdState { pub fn multibodies_mut(&mut self) -> &mut crate::dynamics::GpuMultibodySet { &mut self.multibodies } + #[cfg(feature = "dim3")] + pub fn mb_contact_constraints_len(&self) -> u32 { + self.mb_cons_demand_cpu + } + #[cfg(feature = "dim3")] + pub fn mb_contact_constraints_capacity(&self) -> u32 { + self.multibodies.contact_constraints_capacity() + } /// Immutable access to the multibody set (e.g. to read back `dof_state`). #[cfg(feature = "dim3")] @@ -628,6 +624,10 @@ pub struct RbdSnapshot { #[cfg(feature = "dim3")] impl RbdSnapshot { + pub fn debug_body_pose(&self, body_id: usize) -> Pose { + self.body_poses[body_id] + } + /// 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 @@ -650,16 +650,21 @@ impl RbdState { /// 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); + let nb = self.num_batches as usize; + let mut all_poses: Vec = bytemuck::zeroed_vec(self.body_poses.len() as usize); backend - .slow_read_buffer(self.body_poses.buffer(), &mut body_poses) + .slow_read_buffer(self.body_poses.buffer(), &mut all_poses) .await .unwrap(); - let mut vels = bytemuck::zeroed_vec(self.vels.len() as usize); + let mut all_vels: Vec = bytemuck::zeroed_vec(self.vels.len() as usize); backend - .slow_read_buffer(self.vels.buffer(), &mut vels) + .slow_read_buffer(self.vels.buffer(), &mut all_vels) .await .unwrap(); + let bps = all_poses.len() / nb; + let body_poses = (0..bps).map(|i| all_poses[i * nb]).collect(); + let vs = all_vels.len() / nb; + let vels = (0..vs).map(|i| all_vels[i * nb]).collect(); let mb = self.multibodies.snapshot(backend).await; RbdSnapshot { body_poses, @@ -677,21 +682,25 @@ impl RbdState { ) { 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(); + for (i, pose) in snap.body_poses[..bps].iter().enumerate() { + backend + .write_buffer( + self.body_poses.buffer_mut(), + i as u64 * nb + dst_env as u64, + core::slice::from_ref(pose), + ) + .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(); + for (i, vel) in snap.vels[..vs].iter().enumerate() { + backend + .write_buffer( + self.vels.buffer_mut(), + i as u64 * nb + dst_env as u64, + core::slice::from_ref(vel), + ) + .unwrap(); + } self.multibodies .reset_env_from_snapshot(backend, dst_env, &snap.mb); } @@ -798,7 +807,7 @@ impl RbdState { let t_offs = Tensor::vector(backend, &offs, storage).unwrap(); let params = Tensor::scalar( backend, - UVec4::new(bps, vs, n, 0), + UVec4::new(bps, vs, n, nb), BufferUsages::STORAGE | BufferUsages::UNIFORM | BufferUsages::COPY_DST, ) .unwrap(); diff --git a/src_rbd/pipeline/rbd_state_from_rapier.rs b/src_rbd/pipeline/rbd_state_from_rapier.rs index 7c51e93f..06cb42d2 100644 --- a/src_rbd/pipeline/rbd_state_from_rapier.rs +++ b/src_rbd/pipeline/rbd_state_from_rapier.rs @@ -485,7 +485,12 @@ impl RbdState { .iter() .map(|(mb, ids, bodies)| (*mb, ids, *bodies)) .collect(); - let mut mb = GpuMultibodySet::from_rapier(backend, &mb_refs, max_colliders as u32); + let mut mb = GpuMultibodySet::from_rapier( + backend, + &mb_refs, + max_colliders as u32, + capacities.mb_contact_constraints_capacity, + ); // `set_visible_dt` divides by the substep count, so that has to be // in place first or the multibody integrates at the wrong rate. mb.set_num_solver_iterations(num_solver_iterations); @@ -585,10 +590,40 @@ impl RbdState { } } } - let body_group = Tensor::vector(backend, &all_body_group, BufferUsages::STORAGE).unwrap(); let num_colliders_per_batch = max_colliders; let num_bodies_total = num_colliders_per_batch * num_batches as usize; + let nb = num_batches as usize; + fn interleave_batches(v: &[T], nb: usize) -> Vec { + let per_batch = v.len() / nb.max(1); + let mut out = Vec::with_capacity(v.len()); + for local in 0..per_batch { + for b in 0..nb { + out.push(v[b * per_batch + local]); + } + } + out + } + for (idx, v) in all_collider_parent.iter_mut().enumerate() { + let batch = idx / max_colliders; + *v = *v * nb as u32 + batch as u32; + } + for (idx, v) in all_body_group.iter_mut().enumerate() { + let batch = idx / max_colliders; + *v = *v * nb as u32 + batch as u32; + } + let all_poses = interleave_batches(&all_poses, nb); + let all_vels = interleave_batches(&all_vels, nb); + let all_local_mprops = interleave_batches(&all_local_mprops, nb); + let all_mprops = interleave_batches(&all_mprops, nb); + let all_shapes = interleave_batches(&all_shapes, nb); + let all_collider_local_poses = interleave_batches(&all_collider_local_poses, nb); + let all_collider_parent = interleave_batches(&all_collider_parent, nb); + let all_collision_groups = interleave_batches(&all_collision_groups, nb); + let all_pair_filter = interleave_batches(&all_pair_filter, nb); + let all_collider_materials = interleave_batches(&all_collider_materials, nb); + let all_body_group = interleave_batches(&all_body_group, nb); + let body_group = Tensor::vector(backend, &all_body_group, BufferUsages::STORAGE).unwrap(); // Initial body velocities were accumulated in body-slot order alongside // `all_poses`; zero-filling here would silently drop each body's initial @@ -610,23 +645,16 @@ impl RbdState { storage, ) .unwrap(); - let collision_pairs_len = Tensor::vector_uninit( + let collision_pairs_len = Tensor::vector( backend, - num_batches, + &[0u32], BufferUsages::STORAGE | BufferUsages::COPY_SRC, ) .unwrap(); - let collision_pairs_len_max = - Tensor::vector_uninit(backend, 1, BufferUsages::STORAGE | BufferUsages::COPY_SRC) - .unwrap(); - let num_batches_uniform = Tensor::scalar( - backend, - collision_pairs_len.layout().into(), - BufferUsages::STORAGE | BufferUsages::UNIFORM, - ) - .unwrap(); - // Two-element readback: the (max) collision-pair count and the uncolored count. - let resize_readback = GpuReadback::new(backend, 2).unwrap(); + #[cfg(feature = "dim3")] + let resize_readback = GpuReadback::new(backend, 4).unwrap(); + #[cfg(not(feature = "dim3"))] + let resize_readback = GpuReadback::new(backend, 3).unwrap(); let collision_pairs_indirect = Tensor::scalar_uninit(backend, BufferUsages::STORAGE | BufferUsages::INDIRECT).unwrap(); @@ -654,9 +682,9 @@ impl RbdState { storage, ) .unwrap(); - let pfm_pairs_len = Tensor::vector_uninit( + let pfm_pairs_len = Tensor::vector( backend, - num_batches, + &[0u32], BufferUsages::STORAGE | BufferUsages::COPY_SRC, ) .unwrap(); @@ -709,21 +737,27 @@ impl RbdState { ) .unwrap(); let color_buckets_stride = capacities.solver_colors + 3; - let color_bucket_counts = - Tensor::vector_uninit(backend, color_buckets_stride * num_batches, storage).unwrap(); - let color_bucket_starts = - Tensor::vector_uninit(backend, color_buckets_stride * num_batches, storage).unwrap(); - let color_bucket_cursors = + let color_buckets = Tensor::vector_uninit(backend, color_buckets_stride * num_batches, storage).unwrap(); + let contact_offsets = Tensor::vector( + backend, + vec![0u32; num_batches as usize + 3], + storage, + ) + .unwrap(); + let pair_batch_counts = + Tensor::vector(backend, vec![0u32; num_batches as usize], storage).unwrap(); + let pfm_batch_counts = + Tensor::vector(backend, vec![0u32; num_batches as usize], storage).unwrap(); let color_sorted_ids = Tensor::vector_uninit( backend, capacities.collisions_capacity * num_batches, storage, ) .unwrap(); - let old_constraints_counts = Tensor::vector_uninit( + let old_constraints_counts = Tensor::vector( backend, - num_colliders_per_batch as u32 * num_batches, + vec![0u32; (num_colliders_per_batch as u32 * num_batches) as usize], storage, ) .unwrap(); @@ -752,17 +786,16 @@ impl RbdState { BufferUsages::STORAGE }; - let contacts_per_batch_cpu = capacities.collisions_capacity; - let collision_pairs_per_batch_cpu = capacities.collisions_capacity; + let contacts_capacity_cpu = capacities.collisions_capacity * num_batches; + let collision_pairs_capacity_cpu = capacities.collisions_capacity * num_batches; #[allow(unused_mut)] // Only mutated with the dim3 (multibody) feature. let mut bi = BatchIndices { num_batches, colliders_batch_capacity: num_colliders_per_batch as u32, colliders_len: num_colliders as u32, bodies_len: num_bodies as u32, - collision_pairs_batch_capacity: collision_pairs_per_batch_cpu, - contacts_batch_capacity: contacts_per_batch_cpu, - impulse_joints_batch_capacity: joints.joints_per_batch(), + collision_pairs_capacity: collision_pairs_capacity_cpu, + contacts_capacity: contacts_capacity_cpu, impulse_joints_len: joints.num_active_joints(), solver_color_buckets_stride: color_buckets_stride, ..Default::default() @@ -834,17 +867,20 @@ impl RbdState { collider_materials, collision_pairs, collision_pairs_len, - collision_pairs_len_max, - num_batches_uniform, resize_readback, collision_pairs_indirect, - contacts_per_batch_cpu, - collision_pairs_per_batch_cpu, + contacts_capacity_cpu, + collision_pairs_capacity_cpu, collision_pairs_len_cpu: 0, + #[cfg(feature = "dim3")] + mb_cons_demand_cpu: 0, batch_indices, contacts, contacts_len, contacts_indirect, + contact_offsets, + pair_batch_counts, + pfm_batch_counts, mb_sweep_indirect, pfm_pairs, pfm_pairs_len, @@ -859,9 +895,7 @@ impl RbdState { old_constraints_colors, colored, constraints_rands, - color_bucket_counts, - color_bucket_starts, - color_bucket_cursors, + color_buckets, color_sorted_ids, curr_color: Tensor::scalar( backend, @@ -888,6 +922,7 @@ impl RbdState { old_body_constraint_ids, new_body_constraint_ids, prefix_sum_workspace: PrefixSumWorkspace::default(), + bucket_prefix_workspace: PrefixSumWorkspace::default(), lbvh: LbvhState::with_usages(backend, lbvh_usages), max_colors: capacities.solver_colors, rb_contacts_inert, diff --git a/src_rbd/pipeline/rbd_step.rs b/src_rbd/pipeline/rbd_step.rs index 39afd1c1..702b8fe8 100644 --- a/src_rbd/pipeline/rbd_step.rs +++ b/src_rbd/pipeline/rbd_step.rs @@ -15,7 +15,6 @@ use super::lbvh_validation::validate_lbvh_topology; use super::rbd_state::*; use khal::BufferUsages; 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 @@ -37,7 +36,6 @@ pub struct RbdPipeline { lbvh: Lbvh, 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, @@ -62,7 +60,6 @@ impl RbdPipeline { lbvh: Lbvh::from_backend(backend), coloring: GpuColoring::from_backend(backend)?, warmstart: GpuWarmstart::from_backend(backend)?, - reduce: Reduce::from_backend(backend)?, contact_reduction: false, }) } @@ -152,6 +149,7 @@ impl RbdPipeline { mprops: &state.mprops, contacts: &state.contacts, contacts_len: &state.contacts_len, + contact_offsets: &state.contact_offsets, solver_vels: &mut state.solver_vels, batch_indices: &state.batch_indices, color_uniforms: &state.color_uniforms, @@ -278,9 +276,11 @@ impl RbdPipeline { != RbdResizePolicy::Fixed || state.capacities.collisions_resize_policy != RbdResizePolicy::Fixed; let est_pairs = if readback_enabled { - state.collision_pairs_len_cpu + state.collision_pairs_len_cpu.div_ceil(state.num_batches) } else { - state.collision_pairs_per_batch_cpu + state + .collision_pairs_capacity_cpu + .div_ceil(state.num_batches) }; // Choose the kernel depending on the expected pairs count: small pair @@ -312,11 +312,13 @@ 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.contacts, &mut state.contacts_len, &mut state.contacts_indirect, + &mut state.contact_offsets, + &mut state.pair_batch_counts, + &mut state.pfm_batch_counts, &mut state.mb_sweep_indirect, &mut state.pfm_pairs, &mut state.pfm_pairs_len, @@ -326,6 +328,7 @@ impl RbdPipeline { &state.collider_materials, &state.sim_params, self.contact_reduction, + &state.collision_pairs_indirect, )?; drop(pass); @@ -344,6 +347,7 @@ impl RbdPipeline { let prepare_args = SolverArgs { contacts: &state.contacts, contacts_len: &state.contacts_len, + contact_offsets: &state.contact_offsets, contacts_len_indirect: &state.contacts_indirect, constraints: &mut state.new_constraints, constraint_builders: &mut state.new_constraint_builders, @@ -359,7 +363,7 @@ impl RbdPipeline { local_mprops: &state.local_mprops, body_constraint_counts: &mut state.new_constraints_counts, body_constraint_ids: &mut state.new_body_constraint_ids, - color_bucket_starts: &state.color_bucket_starts, + color_buckets: &state.color_buckets, color_sorted_ids: &state.color_sorted_ids, color_uniforms: &state.color_uniforms, prefix_sum: &self.prefix_sum, @@ -388,7 +392,7 @@ impl RbdPipeline { } else { // Warmstart let warmstart_args = WarmstartArgs { - contacts_len: &state.contacts_len, + contact_offsets: &state.contact_offsets, old_body_constraint_counts: &state.old_constraints_counts, old_constraint_builders: &state.old_constraint_builders, old_body_constraint_ids: &state.old_body_constraint_ids, @@ -412,7 +416,7 @@ impl RbdPipeline { curr_color: &mut state.curr_color, uncolored: &mut state.uncolored, uncolored_staging: &state.uncolored_staging, - contacts_len: &state.contacts_len, + contact_offsets: &state.contact_offsets, colored: &mut state.colored, batch_indices: &state.batch_indices, body_group: &state.body_group, @@ -424,7 +428,7 @@ impl RbdPipeline { // persist, so most constraints can reuse their old color and the // topo-gc iterations converge in 1-2 rounds instead of ~num_colors). let seed_args = crate::dynamics::warmstart::SeedColorsArgs { - contacts_len: &state.contacts_len, + contact_offsets: &state.contact_offsets, old_body_constraint_counts: &state.old_constraints_counts, old_body_constraint_ids: &state.old_body_constraint_ids, old_constraints: &state.old_constraints, @@ -448,7 +452,7 @@ impl RbdPipeline { curr_color: &mut state.curr_color, uncolored: &mut state.uncolored, uncolored_staging: &state.uncolored_staging, - contacts_len: &state.contacts_len, + contact_offsets: &state.contact_offsets, colored: &mut state.colored, batch_indices: &state.batch_indices, body_group: &state.body_group, @@ -464,18 +468,18 @@ impl RbdPipeline { let bucket_args = crate::dynamics::ColorBucketsArgs { contacts_len_indirect: &state.contacts_indirect, constraints_colors: &state.constraints_colors, - contacts_len: &state.contacts_len, - color_bucket_counts: &mut state.color_bucket_counts, - color_bucket_starts: &mut state.color_bucket_starts, - color_bucket_cursors: &mut state.color_bucket_cursors, + constraints: &state.new_constraints, + contact_offsets: &state.contact_offsets, + color_buckets: &mut state.color_buckets, color_sorted_ids: &mut state.color_sorted_ids, batch_indices: &state.batch_indices, }; self.coloring.dispatch_build_color_buckets( + backend, &mut pass, bucket_args, - state.max_colors + 3, - state.num_batches, + &self.prefix_sum, + &mut state.bucket_prefix_workspace, )?; // `+1` because solver iterates 1..=max_colors (color 0 is unassigned). @@ -495,6 +499,7 @@ impl RbdPipeline { let solver_args = SolverArgs { contacts: &state.contacts, contacts_len: &state.contacts_len, + contact_offsets: &state.contact_offsets, contacts_len_indirect: &state.contacts_indirect, constraints: &mut state.new_constraints, constraint_builders: &mut state.new_constraint_builders, @@ -510,7 +515,7 @@ impl RbdPipeline { local_mprops: &state.local_mprops, body_constraint_counts: &mut state.new_constraints_counts, body_constraint_ids: &mut state.new_body_constraint_ids, - color_bucket_starts: &state.color_bucket_starts, + color_buckets: &state.color_buckets, color_sorted_ids: &state.color_sorted_ids, color_uniforms: &state.color_uniforms, prefix_sum: &self.prefix_sum, @@ -604,22 +609,79 @@ impl RbdPipeline { != RbdResizePolicy::Fixed || state.capacities.collisions_resize_policy != RbdResizePolicy::Fixed; - // The readback holds `[max collision-pair count across batches, uncolored - // count]`. The max is computed on the GPU. - let mut counts = [0u32; 2]; + #[cfg(feature = "dim3")] + let mut counts = [0u32; 4]; + #[cfg(not(feature = "dim3"))] + let mut counts = [0u32; 3]; if state.resize_readback.try_take(backend, &mut counts) { // TODO: make the coloring update optional (and pre-configurable) too? - let collision_pairs_len = counts[0]; - let coloring_converged = counts[1]; - state.collision_pairs_len_cpu = collision_pairs_len; + let pairs_len = counts[0].max(counts[1]); + let coloring_converged = counts[2]; + state.collision_pairs_len_cpu = counts[0]; + let nb = state.num_batches; // TODO: Fit will act like Grow. To be able to auto-shrink the max color count, we need // to readback the actual color count. This would also allow us to grow the color // count earlier, before it gets a chance to fail. - if state.capacities.solver_colors_resize_policy != RbdResizePolicy::Fixed + let grow_colors = state.capacities.solver_colors_resize_policy + != RbdResizePolicy::Fixed && coloring_converged == 0 - && !state.rb_contacts_inert - { + && !state.rb_contacts_inert; + let total_capacity = state.collision_pairs_capacity_cpu; + let safe_total = pairs_len.saturating_add(pairs_len / 4); + let new_total = pairs_len + .saturating_add(pairs_len / 2) + .max(state.capacities.collisions_capacity.saturating_mul(nb)); + let resize_pairs = match state.capacities.collisions_resize_policy { + RbdResizePolicy::Fixed => false, + RbdResizePolicy::Grow => safe_total >= total_capacity, + RbdResizePolicy::Fit => safe_total >= total_capacity || total_capacity >= new_total, + }; + + let contact_demand = counts[0].saturating_add(counts[1]); + let contacts_capacity = state.contacts_capacity_cpu; + let safe_contacts = contact_demand.saturating_add(contact_demand / 4); + let new_contacts = contact_demand + .saturating_add(contact_demand / 2) + .max(state.capacities.collisions_capacity.saturating_mul(nb)); + let resize_contacts = match state.capacities.collisions_resize_policy { + RbdResizePolicy::Fixed => false, + RbdResizePolicy::Grow => safe_contacts >= contacts_capacity, + RbdResizePolicy::Fit => { + safe_contacts >= contacts_capacity || contacts_capacity >= new_contacts + } + }; + + #[cfg(feature = "dim3")] + let (resize_mb, new_mb) = { + let mb_demand = counts[3]; + state.mb_cons_demand_cpu = mb_demand; + let mb_capacity = state.multibodies.contact_constraints_capacity(); + let safe_mb = mb_demand.saturating_add(mb_demand / 4); + let new_mb = mb_demand + .saturating_add(mb_demand / 2) + .max(state.multibodies.min_contact_slab_capacity()) + .max( + state + .capacities + .mb_contact_constraints_capacity + .saturating_mul(nb), + ); + let resize_mb = match state.capacities.collisions_resize_policy { + RbdResizePolicy::Fixed => false, + RbdResizePolicy::Grow => safe_mb >= mb_capacity, + RbdResizePolicy::Fit => safe_mb >= mb_capacity || mb_capacity >= new_mb * 2, + }; + (resize_mb, new_mb) + }; + #[cfg(not(feature = "dim3"))] + let resize_mb = false; + + if grow_colors || resize_pairs || resize_contacts || resize_mb { + backend.synchronize()?; + } + + if grow_colors { state.max_colors += 5; // The color-bucket buffers are strided by `max_colors + 3`: @@ -627,89 +689,72 @@ impl RbdPipeline { let storage: BufferUsages = BufferUsages::STORAGE | BufferUsages::COPY_SRC; let stride = state.max_colors + 3; let nb = state.num_batches; - state.color_bucket_counts = Tensor::vector_uninit(backend, stride * nb, storage)?; - state.color_bucket_starts = Tensor::vector_uninit(backend, stride * nb, storage)?; - state.color_bucket_cursors = Tensor::vector_uninit(backend, stride * nb, storage)?; + state.color_buckets = Tensor::vector_uninit(backend, stride * nb, storage)?; state.rebuild_batch_indices(backend); } - // Lazy resize based on the *previous* frame's max pair count. - let per_batch_capacity = - (state.collision_pairs.len() as u32).div_ceil(state.num_batches); + let storage: BufferUsages = BufferUsages::STORAGE | BufferUsages::COPY_SRC; - // Since the auto-resize always lags a bit behind, consider resizing if we have less than 25% - // padding available, reducing the risks of missing contacts. - let safe_capacity = collision_pairs_len.saturating_add(collision_pairs_len / 4); - // Add a 50% extra so we don’t need to reallocate immediately if the - // collision count grows further. Can never be smaller than the `RbdCapacities::collisions_capacity`. - let new_capacity = collision_pairs_len - .saturating_add(collision_pairs_len / 2) - .max(state.capacities.collisions_capacity); - - let resize = match state.capacities.collisions_resize_policy { - RbdResizePolicy::Fixed => false, - RbdResizePolicy::Grow => safe_capacity >= per_batch_capacity, - RbdResizePolicy::Fit => { - safe_capacity >= per_batch_capacity || per_batch_capacity >= new_capacity - } - }; - - if resize { - let storage: BufferUsages = BufferUsages::STORAGE | BufferUsages::COPY_SRC; - let nb = state.num_batches; + if resize_pairs { + state.collision_pairs = Tensor::vector_uninit(backend, new_total, storage)?; + state.pfm_pairs = Tensor::vector_uninit(backend, new_total, storage)?; + state.collision_pairs_capacity_cpu = new_total; + } - state.collision_pairs = Tensor::vector_uninit(backend, new_capacity * nb, storage)?; - state.contacts = Tensor::vector_uninit(backend, new_capacity * nb, storage)?; - state.pfm_pairs = Tensor::vector_uninit(backend, new_capacity * nb, storage)?; - state.old_constraints = Tensor::vector_uninit(backend, new_capacity * nb, storage)?; + if resize_contacts { + state.contacts = Tensor::vector_uninit(backend, new_contacts, storage)?; + state.old_constraints = Tensor::vector_uninit(backend, new_contacts, storage)?; state.old_constraint_builders = - Tensor::vector_uninit(backend, new_capacity * nb, storage)?; + Tensor::vector_uninit(backend, new_contacts, storage)?; state.old_body_constraint_ids = - Tensor::vector_uninit(backend, new_capacity * 2 * nb, storage)?; - state.new_constraints = Tensor::vector_uninit(backend, new_capacity * nb, storage)?; + Tensor::vector_uninit(backend, new_contacts * 2, storage)?; + state.new_constraints = Tensor::vector_uninit(backend, new_contacts, storage)?; state.new_constraint_builders = - Tensor::vector_uninit(backend, new_capacity * nb, storage)?; + Tensor::vector_uninit(backend, new_contacts, storage)?; state.new_body_constraint_ids = - Tensor::vector_uninit(backend, new_capacity * 2 * nb, storage)?; - state.constraints_colors = - Tensor::vector_uninit(backend, new_capacity * nb, storage)?; + Tensor::vector_uninit(backend, new_contacts * 2, storage)?; + state.constraints_colors = Tensor::vector_uninit(backend, new_contacts, storage)?; // Zeroed (not uninit): 0 = "uncolored" disables color seeding // for the frame right after the resize. state.old_constraints_colors = - Tensor::vector(backend, vec![0u32; (new_capacity * nb) as usize], storage)?; - state.colored = Tensor::vector_uninit(backend, new_capacity * nb, storage)?; - state.constraints_rands = - Tensor::vector_uninit(backend, new_capacity * nb, storage)?; - state.color_sorted_ids = - Tensor::vector_uninit(backend, new_capacity * nb, storage)?; - - state.collision_pairs_per_batch_cpu = new_capacity; - state.contacts_per_batch_cpu = new_capacity; + Tensor::vector(backend, vec![0u32; new_contacts as usize], storage)?; + state.colored = Tensor::vector_uninit(backend, new_contacts, storage)?; + state.constraints_rands = Tensor::vector_uninit(backend, new_contacts, storage)?; + state.color_sorted_ids = Tensor::vector_uninit(backend, new_contacts, storage)?; + let counts_len = state.old_constraints_counts.len() as usize; + state.old_constraints_counts = + Tensor::vector(backend, vec![0u32; counts_len], storage)?; + + state.contacts_capacity_cpu = new_contacts; + } + #[cfg(feature = "dim3")] + if resize_mb { + state.multibodies.resize_contact_slabs(backend, new_mb); + } + if resize_pairs || resize_contacts || resize_mb { state.rebuild_batch_indices(backend); } } if readback_enabled && state.resize_readback.is_idle() { - let pairs_source = if state.num_batches > 1 { - let mut encoder = backend.begin_encoding(); - let mut pass = encoder.begin_pass("[RBD] calc-max-coll-len", None); - self.reduce.reduce_max_u32.call( - &mut pass, - 1, - &state.num_batches_uniform, - &state.collision_pairs_len, - &mut state.collision_pairs_len_max, - )?; - drop(pass); - backend.submit(encoder)?; - state.collision_pairs_len_max.buffer() - } else { - state.collision_pairs_len.buffer() - }; - + #[cfg(feature = "dim3")] + state.resize_readback.request( + backend, + &[ + (state.collision_pairs_len.buffer(), 0, 1), + (state.pfm_pairs_len.buffer(), 0, 1), + (state.uncolored.buffer(), 0, 1), + (state.multibodies.mb_cons_demand().buffer(), 0, 1), + ], + )?; + #[cfg(not(feature = "dim3"))] state.resize_readback.request( backend, - &[(pairs_source, 0, 1), (state.uncolored.buffer(), 0, 1)], + &[ + (state.collision_pairs_len.buffer(), 0, 1), + (state.pfm_pairs_len.buffer(), 0, 1), + (state.uncolored.buffer(), 0, 1), + ], )?; } diff --git a/src_rbd_shaders/broad_phase/brute_force.rs b/src_rbd_shaders/broad_phase/brute_force.rs index 5b095e9f..50f97281 100644 --- a/src_rbd_shaders/broad_phase/brute_force.rs +++ b/src_rbd_shaders/broad_phase/brute_force.rs @@ -35,9 +35,9 @@ pub fn gpu_bf_compute_aabbs( let batch_id = invocation_id.x / n; let i = invocation_id.x % n; - let poses = batch_ids.coll_batch(batch_id, poses); - let shapes = batch_ids.coll_batch(batch_id, shapes); - let out = batch_ids.coll_start(batch_id) + i as usize; + let poses = batch_ids.ib(batch_id, poses); + let shapes = batch_ids.ib(batch_id, shapes); + let out = (crate::broad_phase::scratch_start(batch_ids, batch_id) + i) as usize; aabbs.write( out, shapes[i as usize].compute_aabb(poses[i as usize], vertices), @@ -73,8 +73,8 @@ pub fn gpu_bf_find_pairs( return; } - let collision_groups = batch_ids.coll_batch(batch_id, collision_groups); - let pair_filter = batch_ids.coll_batch(batch_id, pair_filter); + let collision_groups = batch_ids.ib(batch_id, collision_groups); + let pair_filter = batch_ids.ib(batch_id, pair_filter); // Skip pairs whose collision groups don't authorize an interaction. if !collision_groups[i as usize].test(collision_groups[j as usize]) { @@ -88,7 +88,7 @@ pub fn gpu_bf_find_pairs( return; } - let coll_start = batch_ids.coll_start(batch_id); + let coll_start = crate::broad_phase::scratch_start(batch_ids, batch_id) as usize; // Dilate one side by the contact prediction distance. let mut aabb_i = aabbs.read(coll_start + i as usize); let dilation = Vector::splat(params.prediction_distance()); @@ -99,13 +99,18 @@ pub fn gpu_bf_find_pairs( return; } - let target_pair_index = atomic_add_u32(collision_pairs_len.at_mut(batch_id as usize), 1); + let target_pair_index = atomic_add_u32(collision_pairs_len.at_mut(0), 1); // If we exceed capacity, keep counting the pairs but don’t store any more to avoid overflow. - if target_pair_index < batch_ids.collision_pairs_batch_capacity { - let mut collision_pairs = batch_ids.collision_pairs_batch_mut(batch_id, collision_pairs); - collision_pairs[target_pair_index as usize] = CollisionPair { - colliders: UVec2::new(i, j), - }; + if target_pair_index < batch_ids.collision_pairs_capacity { + collision_pairs.write( + target_pair_index as usize, + CollisionPair { + colliders: UVec2::new( + batch_ids.body_global(batch_id, i) as u32, + batch_ids.body_global(batch_id, j) as u32, + ), + }, + ); } } diff --git a/src_rbd_shaders/broad_phase/lbvh.rs b/src_rbd_shaders/broad_phase/lbvh.rs index 31051b47..9647282b 100644 --- a/src_rbd_shaders/broad_phase/lbvh.rs +++ b/src_rbd_shaders/broad_phase/lbvh.rs @@ -56,9 +56,23 @@ pub fn gpu_lbvh_reset_collision_pairs( #[spirv(global_invocation_id)] invocation_id: UVec3, #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] collision_pairs_len: &mut [u32], ) { - let batch_id = invocation_id.x as usize; - if batch_id < collision_pairs_len.len() { - collision_pairs_len.write(batch_id, 0); + let i = invocation_id.x as usize; + if i < collision_pairs_len.len() { + collision_pairs_len.write(i, 0); + } +} +#[spirv_bindgen] +#[spirv(compute(threads(1)))] +pub fn gpu_flat_list_dispatch( + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] len: &mut [u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] indirect_args: &mut [u32; 3], + #[spirv(uniform, descriptor_set = 0, binding = 2)] batch_ids: &BatchIndices, +) { + for _ in 0..1 { + let total = atomic_load_u32(len.at_mut(0)).min(batch_ids.collision_pairs_capacity); + *indirect_args.at_mut(0) = total.div_ceil(WORKGROUP_SIZE); + *indirect_args.at_mut(1) = 1; + *indirect_args.at_mut(2) = 1; } } @@ -73,10 +87,9 @@ pub const MAX_REDUCE_LANES: u32 = 256; /// atomic or they occasionally read stale data (breaks Windows+Nvidia+wgpu, see /// ). #[inline(always)] -pub(crate) fn max_len_indirect_args( +pub(crate) fn reduce_max_lens( lane: u32, lens: &mut [u32], - indirect_args: &mut [u32; 3], partial: &mut [u32; MAX_REDUCE_LANES as usize], ) { let num_batches = lens.len(); @@ -99,24 +112,6 @@ pub(crate) fn max_len_indirect_args( } workgroup_memory_barrier_with_group_sync(); } - - if lane == 0 { - *indirect_args.at_mut(0) = partial.read(0).div_ceil(WORKGROUP_SIZE); - *indirect_args.at_mut(1) = num_batches as u32; - *indirect_args.at_mut(2) = 1; - } -} - -/// Initializes indirect dispatch arguments for narrow phase. -#[spirv_bindgen] -#[spirv(compute(threads(256)))] -pub fn gpu_lbvh_init_dispatch( - #[spirv(local_invocation_id)] lid: UVec3, - #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] collision_pairs_len: &mut [u32], - #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] indirect_args: &mut [u32; 3], - #[spirv(workgroup)] partial: &mut [u32; MAX_REDUCE_LANES as usize], -) { - max_len_indirect_args(lid.x, collision_pairs_len, indirect_args, partial); } /// Runs a reduction to compute the AABB of the collider positions. @@ -135,14 +130,12 @@ pub fn gpu_lbvh_compute_domain( let thread_id = global_id.x; *workspace_mins.at_mut(thread_id as usize) = Vector::splat(MAX_FLT); *workspace_maxs.at_mut(thread_id as usize) = Vector::splat(-MAX_FLT); - let colliders_start = batch_ids.coll_start(batch_id) as u32; - let colliders_end = colliders_start + batch_ids.colliders_len; for i in StepRng::new( - colliders_start + thread_id..colliders_end, + thread_id..batch_ids.colliders_len, REDUCTION_WORKGROUP_SIZE, ) { - let val_i = poses.at(i as usize).translation; + let val_i = poses.at(batch_ids.body_global(batch_id, i)).translation; *workspace_mins.at_mut(thread_id as usize) = workspace_mins.at(thread_id as usize).min(val_i); *workspace_maxs.at_mut(thread_id as usize) = @@ -194,17 +187,13 @@ pub fn gpu_lbvh_compute_morton( let num_threads = num_workgroups.x * WORKGROUP_SIZE; let batch_id = invocation_id.y; let domain_aabb = domain_aabb.read(batch_id as usize); - let colliders_start = batch_ids.coll_start(batch_id) as u32; - let colliders_end = colliders_start + batch_ids.colliders_len; + let scratch = scratch_start(batch_ids, batch_id); - for i in StepRng::new( - colliders_start + invocation_id.x..colliders_end, - num_threads, - ) { - let center = poses.at(i as usize).translation; + for i in StepRng::new(invocation_id.x..batch_ids.colliders_len, num_threads) { + let center = poses.at(batch_ids.body_global(batch_id, i)).translation; let normalized = (center - domain_aabb.mins) / (domain_aabb.maxs - domain_aabb.mins); let morton_key = morton(normalized); - morton_keys.write(i as usize, morton_key); + morton_keys.write((scratch + i) as usize, morton_key); } } @@ -223,7 +212,7 @@ pub fn gpu_lbvh_build( ) { let num_threads = num_workgroups.x * WORKGROUP_SIZE; let batch_id = invocation_id.y; - let colliders_start = batch_ids.coll_start(batch_id) as u32; + let colliders_start = scratch_start(batch_ids, batch_id); let num_bodies = batch_ids.colliders_len; let num_internal_nodes = num_bodies - 1; let first_leaf_id = num_internal_nodes; @@ -332,13 +321,13 @@ pub fn gpu_lbvh_refit_leaves( // Bottom-up refit. Leaf index starts at `num_colliders`. let num_threads = num_workgroups.x * WORKGROUP_SIZE; let batch_id = invocation_id.y; - let colliders_start = batch_ids.coll_start(batch_id) as u32; + let colliders_start = scratch_start(batch_ids, batch_id); let num_colliders = batch_ids.colliders_len; let first_leaf_id = num_colliders - 1; - let poses = batch_ids.coll_batch(batch_id, poses); - let shapes = batch_ids.coll_batch(batch_id, shapes); - let sorted_colliders = batch_ids.coll_batch(batch_id, sorted_colliders); + let poses = batch_ids.ib(batch_id, poses); + let shapes = batch_ids.ib(batch_id, shapes); + let sorted_colliders = Slice(sorted_colliders, colliders_start as usize); let mut tree = SliceMut(tree, root_id(colliders_start) as usize); for i in StepRng::new(invocation_id.x..num_colliders, num_threads) { @@ -371,7 +360,7 @@ pub fn gpu_lbvh_refit_internal( // Bottom-up refit. Leaf index starts at `num_colliders`. let num_threads = 256u32; let batch_id = workgroup_id.y; - let colliders_start = batch_ids.coll_start(batch_id) as u32; + let colliders_start = scratch_start(batch_ids, batch_id); let num_bodies = batch_ids.colliders_len; let first_leaf_id = num_bodies - 1; @@ -468,13 +457,13 @@ pub fn gpu_lbvh_refit( // Bottom-up refit. Leaf index starts at `num_colliders`. let batch_id = invocation_id.y; let num_threads = num_workgroups.x * WORKGROUP_SIZE; - let colliders_start = batch_ids.coll_start(batch_id) as u32; + let colliders_start = scratch_start(batch_ids, batch_id); let num_bodies = batch_ids.colliders_len; let first_leaf_id = num_bodies - 1; - let poses = batch_ids.coll_batch(batch_id, poses); - let shapes = batch_ids.coll_batch(batch_id, shapes); - let sorted_colliders = batch_ids.coll_batch(batch_id, sorted_colliders); + let poses = batch_ids.ib(batch_id, poses); + let shapes = batch_ids.ib(batch_id, shapes); + let sorted_colliders = Slice(sorted_colliders, colliders_start as usize); let mut tree = SliceMut(tree, root_id(colliders_start) as usize); for i in StepRng::new(invocation_id.x..num_bodies, num_threads) { @@ -544,14 +533,13 @@ pub fn gpu_lbvh_find_collision_pairs( ) { let num_threads = num_workgroups.x * WORKGROUP_SIZE; let batch_id = invocation_id.y; - let colliders_start = batch_ids.coll_start(batch_id) as u32; + let colliders_start = scratch_start(batch_ids, batch_id); let num_bodies = batch_ids.colliders_len; let first_leaf_id = num_bodies - 1; - let mut collision_pairs = batch_ids.collision_pairs_batch_mut(batch_id, collision_pairs); let tree = Slice(tree, root_id(colliders_start) as usize); - let collision_groups = batch_ids.coll_batch(batch_id, collision_groups); - let pair_filter = batch_ids.coll_batch(batch_id, pair_filter); + let collision_groups = batch_ids.ib(batch_id, collision_groups); + let pair_filter = batch_ids.ib(batch_id, pair_filter); for leaf_i in StepRng::new(invocation_id.x..num_bodies, num_threads) { let i = tree.at((first_leaf_id + leaf_i) as usize).left; @@ -595,23 +583,25 @@ pub fn gpu_lbvh_find_collision_pairs( continue; } - let target_pair_index = - atomic_add_u32(collision_pairs_len.at_mut(batch_id as usize), 1); + let target_pair_index = atomic_add_u32(collision_pairs_len.at_mut(0), 1); // NOTE: if the index is out-of-bounds (meaning the `collision_pairs` isn't // big enough), don't write. But keep traversing so we get the exact count we need // for reallocating the buffers. - if target_pair_index < batch_ids.collision_pairs_batch_capacity { - // NOTE: we only store the collider pair here. The parent body - // ids are resolved lazily, at the very last moment, when - // the narrow-phase writes the `IndexedManifold` consumed + if target_pair_index < batch_ids.collision_pairs_capacity { // by the solver — keeping this hot buffer (and the // intermediate pfm-pair buffer) narrow, and keeping // `collider_parent` out of the broad phase entirely. let (ci, cj) = if i < j { (i, j) } else { (j, i) }; - collision_pairs[target_pair_index as usize] = CollisionPair { - colliders: UVec2::new(ci, cj), - }; + collision_pairs.write( + target_pair_index as usize, + CollisionPair { + colliders: UVec2::new( + batch_ids.body_global(batch_id, ci) as u32, + batch_ids.body_global(batch_id, cj) as u32, + ), + }, + ); } } else { let left = node.left; @@ -718,6 +708,11 @@ pub fn prefix_len( } } +#[inline] +pub fn scratch_start(batch_ids: &BatchIndices, batch_id: u32) -> u32 { + batch_id * batch_ids.colliders_batch_capacity +} + fn root_id(collider_start_id: u32) -> u32 { // Every LBVH tree contains `n - 1` internal nodes and `n` leaves, where // `n` is its number of colliders. This is a total of `2n - 1`, but to diff --git a/src_rbd_shaders/broad_phase/narrow_phase.rs b/src_rbd_shaders/broad_phase/narrow_phase.rs index 8213d85b..c274df65 100644 --- a/src_rbd_shaders/broad_phase/narrow_phase.rs +++ b/src_rbd_shaders/broad_phase/narrow_phase.rs @@ -17,9 +17,12 @@ 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 super::lbvh::{MAX_REDUCE_LANES, reduce_max_lens}; use crate::broad_phase::CollisionPair; use crate::utils::{BatchIndices, SliceMut}; use glamx::UVec2; @@ -33,17 +36,110 @@ pub fn gpu_reset_narrow_phase( #[spirv(global_invocation_id)] invocation_id: UVec3, #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] contacts_len: &mut [u32], #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] pfm_pairs_len: &mut [u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] pair_batch_counts: &mut [u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] pfm_batch_counts: &mut [u32], ) { - let batch_id = invocation_id.x as usize; - if batch_id < contacts_len.len() { - contacts_len.write(batch_id, 0); - pfm_pairs_len.write(batch_id, 0); + let i = invocation_id.x as usize; + if i < contacts_len.len() { + contacts_len.write(i, 0); + } + if i < pfm_pairs_len.len() { + pfm_pairs_len.write(i, 0); + } + if i < pair_batch_counts.len() { + pair_batch_counts.write(i, 0); + } + if i < pfm_batch_counts.len() { + pfm_batch_counts.write(i, 0); } } -/// Initializes indirect dispatch arguments for constraint solver. +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn gpu_count_pairs_per_batch( + #[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: &mut [u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] pair_batch_counts: &mut [u32], + #[spirv(uniform, descriptor_set = 0, binding = 3)] batch_ids: &BatchIndices, +) { + let num_threads = num_workgroups.x * WORKGROUP_SIZE; + let total = + atomic_load_u32(collision_pairs_len.at_mut(0)).min(batch_ids.collision_pairs_capacity); + for t in StepRng::new(invocation_id.x..total, num_threads) { + let pair = collision_pairs.read(t as usize); + let batch_id = batch_ids.collider_batch(pair.colliders.x); + atomic_add_u32(pair_batch_counts.at_mut(batch_id as usize), 1); + } +} +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn gpu_count_pfm_per_batch( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[spirv(num_workgroups)] num_workgroups: UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] pfm_pairs: &[NarrowPhasePfmPair], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] pfm_pairs_len: &mut [u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] pfm_batch_counts: &mut [u32], + #[spirv(uniform, descriptor_set = 0, binding = 3)] batch_ids: &BatchIndices, +) { + let num_threads = num_workgroups.x * WORKGROUP_SIZE; + let total = atomic_load_u32(pfm_pairs_len.at_mut(0)).min(batch_ids.collision_pairs_capacity); + for t in StepRng::new(invocation_id.x..total, num_threads) { + let pair = pfm_pairs.read(t as usize); + let batch_id = batch_ids.collider_batch(pair.colliders.x); + atomic_add_u32(pfm_batch_counts.at_mut(batch_id as usize), 1); + } +} /// -/// Also inits `mb_sweep_indirect`, the workgroup grid for the per-multibody +#[spirv_bindgen] +#[spirv(compute(threads(1)))] +pub fn gpu_contact_offsets_scan( + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] pair_batch_counts: &mut [u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] pfm_batch_counts: &mut [u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] collision_pairs_len: &mut [u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] pfm_pairs_len: &mut [u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] contact_offsets: &mut [u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] indirect_args: &mut [u32; 3], + #[spirv(uniform, descriptor_set = 0, binding = 6)] batch_ids: &BatchIndices, +) { + let num_batches = batch_ids.num_batches as usize; + let capacity = batch_ids.contacts_capacity; + let mut total = 0u32; + for b in 0..num_batches { + contact_offsets.write(b, total); + let bound = atomic_load_u32(pair_batch_counts.at_mut(b)) + + atomic_load_u32(pfm_batch_counts.at_mut(b)); + total = (total + bound).min(capacity); + } + contact_offsets.write(num_batches, total); + contact_offsets.write( + num_batches + 1, + atomic_load_u32(collision_pairs_len.at_mut(0)).min(batch_ids.collision_pairs_capacity), + ); + contact_offsets.write( + num_batches + 2, + atomic_load_u32(pfm_pairs_len.at_mut(0)).min(batch_ids.collision_pairs_capacity), + ); + *indirect_args.at_mut(0) = total.div_ceil(WORKGROUP_SIZE); + *indirect_args.at_mut(1) = 1; + *indirect_args.at_mut(2) = 1; +} +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn gpu_zero_contact_lens( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[spirv(num_workgroups)] num_workgroups: UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] contacts: &mut [IndexedManifold], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] contact_offsets: &[u32], + #[spirv(uniform, descriptor_set = 0, binding = 2)] batch_ids: &BatchIndices, +) { + let num_threads = num_workgroups.x * WORKGROUP_SIZE; + let total = contact_offsets.read(batch_ids.num_batches as usize); + for t in StepRng::new(invocation_id.x..total, num_threads) { + contacts.at_mut(t as usize).contact.len = 0; + } +} /// contact-constraint dispatches (`[multibodies_batch_capacity, num_batches, /// 1]`). #[spirv_bindgen] @@ -51,12 +147,11 @@ pub fn gpu_reset_narrow_phase( pub fn gpu_narrow_phase_init_contacts_dispatch( #[spirv(local_invocation_id)] lid: UVec3, #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] contacts_len: &mut [u32], - #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] indirect_args: &mut [u32; 3], - #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] mb_sweep_indirect: &mut [u32; 3], - #[spirv(uniform, descriptor_set = 0, binding = 3)] batch_ids: &BatchIndices, + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] mb_sweep_indirect: &mut [u32; 3], + #[spirv(uniform, descriptor_set = 0, binding = 2)] batch_ids: &BatchIndices, #[spirv(workgroup)] partial: &mut [u32; MAX_REDUCE_LANES as usize], ) { - max_len_indirect_args(lid.x, contacts_len, indirect_args, partial); + reduce_max_lens(lid.x, contacts_len, partial); // `partial[0]` holds the max after the reduction (all lanes synced). if lid.x == 0 { let any_contacts = partial.read(0) > 0; @@ -128,15 +223,19 @@ pub fn gpu_reduce_contacts( #[spirv(workgroup_id)] workgroup_id: UVec3, #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] contacts: &mut [IndexedManifold], #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] contacts_len: &mut [u32], - #[spirv(uniform, descriptor_set = 0, binding = 2)] batch_ids: &BatchIndices, - #[spirv(uniform, descriptor_set = 0, binding = 3)] params: &RbdSimParams, + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] contact_offsets: &[u32], + #[allow(unused_variables)] + #[spirv(uniform, descriptor_set = 0, binding = 3)] + batch_ids: &BatchIndices, + #[spirv(uniform, descriptor_set = 0, binding = 4)] params: &RbdSimParams, ) { let prediction = params.prediction_distance(); let merge_cos = params.contact_merge_cos; let batch_id = workgroup_id.y; - let capacity = batch_ids.contacts_batch_capacity as usize; - let mut contacts = batch_ids.contact_batch_mut(batch_id, contacts); - let n = (contacts_len.read(batch_id as usize) as usize).min(capacity); + let seg_start = contact_offsets.read(batch_id as usize) as usize; + let seg_end = contact_offsets.read(batch_id as usize + 1) as usize; + let mut contacts = SliceMut(contacts, seg_start); + let n = (contacts_len.read(batch_id as usize) as usize).min(seg_end - seg_start); // Write cursor: always <= the read cursor, so compacting in place is safe. let mut w = 0usize; @@ -212,11 +311,10 @@ pub fn gpu_reduce_contacts( 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); + for i in w..n { + contacts[i].contact.len = 0; } + contacts_len.write(batch_id as usize, w as u32); } /// Narrow phase, pass 1 of 2: analytic shape-shape contacts for ball / cuboid @@ -230,7 +328,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], - #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] collision_pairs_len: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] contact_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], @@ -246,38 +344,27 @@ 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; - - 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); - - // 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]; + + let total = contact_offsets.read(batch_ids.num_batches as usize + 1); + + for t in StepRng::new(invocation_id.x..total, num_threads) { + let pair = collision_pairs.read(t as usize); + let batch_id = batch_ids.collider_batch(pair.colliders.x); + let seg_start = contact_offsets.read(batch_id as usize); + let seg_end = contact_offsets.read(batch_id as usize + 1); + let contacts_len = contacts_len.at_mut(batch_id 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, 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); + let body1 = collider_parent.read(pair.colliders.x as usize); + let body2 = collider_parent.read(pair.colliders.y as usize); if body1 == body2 { continue; } - let pose1 = poses[pair.colliders.x as usize]; - let pose2 = poses[pair.colliders.y as usize]; - let shape1 = &shapes[pair.colliders.x as usize]; - let shape2 = &shapes[pair.colliders.y as usize]; + let pose1 = poses.read(pair.colliders.x as usize); + let pose2 = poses.read(pair.colliders.y as usize); + let shape1 = shapes.at(pair.colliders.x as usize); + let shape2 = shapes.at(pair.colliders.y as usize); let shape_ty1 = shape1.shape_type(); let shape_ty2 = shape2.shape_type(); let mut manifold = ContactManifold::default(); @@ -320,21 +407,22 @@ pub fn gpu_narrow_phase_shape_shape( // 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 { - let target_contact_index = atomic_add_u32(contacts_len, 1) as usize; - - // NOTE: if we exceed the contacts allocation size, just skip - // the contact. - if target_contact_index < contacts_batch_capacity { - let mat1 = collider_materials[pair.colliders.x as usize]; - let mat2 = collider_materials[pair.colliders.y as usize]; - contacts[target_contact_index] = IndexedManifold { - contact: manifold, - colliders: pair.colliders, - bodies: UVec2::new(body1, body2), - friction: mat1.combined_friction(&mat2), - restitution: mat1.combined_restitution(&mat2), - _padding: [0.0; 2], - }; + let idx = seg_start + atomic_add_u32(contacts_len, 1); + + if idx < seg_end { + let mat1 = collider_materials.read(pair.colliders.x as usize); + let mat2 = collider_materials.read(pair.colliders.y as usize); + contacts.write( + idx as usize, + IndexedManifold { + contact: manifold, + colliders: pair.colliders, + bodies: UVec2::new(body1, body2), + friction: mat1.combined_friction(&mat2), + restitution: mat1.combined_restitution(&mat2), + _padding: [0.0; 2], + }, + ); } } } @@ -351,7 +439,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], - #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] collision_pairs_len: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] collision_pairs_len: &mut [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)] @@ -367,28 +455,21 @@ pub fn gpu_narrow_phase_shape_shape_deferred( ) { 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 pfm_capacity = batch_ids.collision_pairs_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 total = atomic_load_u32(collision_pairs_len.at_mut(0)).min(batch_ids.collision_pairs_capacity); // 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) { - let pair = collision_pairs[i as usize]; - let shape1 = &shapes[pair.colliders.x as usize]; - let shape2 = &shapes[pair.colliders.y as usize]; + for t in StepRng::new(invocation_id.x..total, num_threads) { + let mut pfm_pairs = SliceMut(&mut *pfm_pairs, 0); + let pfm_pairs_len = pfm_pairs_len.at_mut(0); + let pair = collision_pairs.read(t as usize); + let shape1 = shapes.at(pair.colliders.x as usize); + let shape2 = shapes.at(pair.colliders.y as usize); let shape_ty1 = shape1.shape_type(); let shape_ty2 = shape2.shape_type(); @@ -421,8 +502,8 @@ pub fn gpu_narrow_phase_shape_shape_deferred( continue; } - let pose1 = poses[pair.colliders.x as usize]; - let pose2 = poses[pair.colliders.y as usize]; + let pose1 = poses.read(pair.colliders.x as usize); + let pose2 = poses.read(pair.colliders.y as usize); let pose12 = pose1.inverse() * pose2; // PFM - PFM (generic convex shapes via GJK/EPA) @@ -441,7 +522,7 @@ pub fn gpu_narrow_phase_shape_shape_deferred( }; let pfm_index = atomic_add_u32(pfm_pairs_len, 1); // NOTE: if we exceed capacity, just skip the pair. - if (pfm_index as usize) < contacts_batch_capacity { + if (pfm_index as usize) < pfm_capacity { pfm_pairs.write(pfm_index as usize, pfm_pair); } @@ -463,7 +544,7 @@ pub fn gpu_narrow_phase_shape_shape_deferred( pair.colliders, &mut pfm_pairs, pfm_pairs_len, - contacts_batch_capacity, + pfm_capacity, vertices, indices, ); @@ -482,7 +563,7 @@ pub fn gpu_narrow_phase_shape_shape_deferred( UVec2::new(pair.colliders.y, pair.colliders.x), &mut pfm_pairs, pfm_pairs_len, - contacts_batch_capacity, + pfm_capacity, vertices, indices, ); @@ -502,7 +583,7 @@ pub fn gpu_narrow_phase_shape_shape_deferred( pair.colliders, &mut pfm_pairs, pfm_pairs_len, - contacts_batch_capacity, + pfm_capacity, vertices, indices, ); @@ -521,7 +602,7 @@ pub fn gpu_narrow_phase_shape_shape_deferred( UVec2::new(pair.colliders.y, pair.colliders.x), &mut pfm_pairs, pfm_pairs_len, - contacts_batch_capacity, + pfm_capacity, vertices, indices, ); @@ -687,19 +768,6 @@ pub struct NarrowPhasePfmPair { colliders: UVec2, } -/// Initializes PFM-PFM dispatch arguments for constraint solver. Dispatch one -/// [`MAX_REDUCE_LANES`]-thread workgroup. -#[spirv_bindgen] -#[spirv(compute(threads(256)))] -pub fn gpu_init_pfm_pfm_dispatch( - #[spirv(local_invocation_id)] lid: UVec3, - #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] pfm_pairs_len: &mut [u32], - #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] indirect_args: &mut [u32; 3], - #[spirv(workgroup)] partial: &mut [u32; MAX_REDUCE_LANES as usize], -) { - max_len_indirect_args(lid.x, pfm_pairs_len, indirect_args, partial); -} - #[spirv_bindgen] #[spirv(compute(threads(64)))] // TODO PERF: pfm_pfm is very divergent. Use a smaller workgroup size? pub fn gpu_narrow_phase_pfm_pfm( @@ -708,10 +776,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], - #[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`). + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] contact_offsets: &[u32], #[spirv(uniform, descriptor_set = 0, binding = 4)] batch_ids: &BatchIndices, #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] vertices: &[PaddedVector], #[allow(unused_variables)] @@ -726,28 +791,22 @@ 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 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]; + + let total = contact_offsets.read(batch_ids.num_batches as usize + 2); + + for t in StepRng::new(invocation_id.x..total, num_threads) { + let pair = pfm_pairs.read(t as usize); + let batch_id = batch_ids.collider_batch(pair.colliders.x); + let seg_start = contact_offsets.read(batch_id as usize); + let seg_end = contact_offsets.read(batch_id as usize + 1); + let contacts_len = contacts_len.at_mut(batch_id 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 // 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 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); + let body1 = collider_parent.read(pair.colliders.x as usize); + let body2 = collider_parent.read(pair.colliders.y as usize); if body1 == body2 { continue; } @@ -764,20 +823,22 @@ pub fn gpu_narrow_phase_pfm_pfm( ); if manifold.len > 0 && manifold.points_a.at(0).dist < prediction { - let target_contact_index = atomic_add_u32(contacts_len, 1) as usize; - - // NOTE: if we exceed capacity, just skip the pair. - if target_contact_index < contacts_batch_capacity { - let mat1 = collider_materials[pair.colliders.x as usize]; - let mat2 = collider_materials[pair.colliders.y as usize]; - contacts[target_contact_index] = IndexedManifold { - contact: manifold, - colliders: pair.colliders, - bodies: UVec2::new(body1, body2), - friction: mat1.combined_friction(&mat2), - restitution: mat1.combined_restitution(&mat2), - _padding: [0.0; 2], - }; + let idx = seg_start + atomic_add_u32(contacts_len, 1); + + if idx < seg_end { + let mat1 = collider_materials.read(pair.colliders.x as usize); + let mat2 = collider_materials.read(pair.colliders.y as usize); + contacts.write( + idx as usize, + IndexedManifold { + contact: manifold, + colliders: pair.colliders, + bodies: UVec2::new(body1, body2), + friction: mat1.combined_friction(&mat2), + restitution: mat1.combined_restitution(&mat2), + _padding: [0.0; 2], + }, + ); } } } diff --git a/src_rbd_shaders/dynamics/color_buckets.rs b/src_rbd_shaders/dynamics/color_buckets.rs index b0478abc..82a9c969 100644 --- a/src_rbd_shaders/dynamics/color_buckets.rs +++ b/src_rbd_shaders/dynamics/color_buckets.rs @@ -13,7 +13,8 @@ use khal_std::glamx::UVec3; use khal_std::macros::{spirv, spirv_bindgen}; use khal_std::{index::MaybeIndexUnchecked, iter::StepRng, sync::atomic_add_u32}; -use crate::utils::BatchIndices; +use super::constraint::TwoBodyConstraint; +use crate::utils::{BatchIndices, Slice}; const WORKGROUP_SIZE: u32 = 64; @@ -22,15 +23,11 @@ const WORKGROUP_SIZE: u32 = 64; #[spirv(compute(threads(64)))] pub fn gpu_color_buckets_reset( #[spirv(global_invocation_id)] invocation_id: UVec3, - #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] color_counts: &mut [u32], - #[spirv(uniform, descriptor_set = 0, binding = 1)] batch_ids: &BatchIndices, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] color_buckets: &mut [u32], ) { - let stride = batch_ids.solver_color_buckets_stride; - let batch_id = invocation_id.y; - let i = invocation_id.x; - - if i < stride { - color_counts.write((batch_id * stride + i) as usize, 0); + let i = invocation_id.x as usize; + if i < color_buckets.len() { + color_buckets.write(i, 0); } } @@ -41,82 +38,50 @@ pub fn gpu_color_buckets_count( #[spirv(global_invocation_id)] invocation_id: UVec3, #[spirv(num_workgroups)] num_workgroups: UVec3, #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] constraints_colors: &[u32], - #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] contacts_len: &[u32], - #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] color_counts: &mut [u32], - #[spirv(uniform, descriptor_set = 0, binding = 3)] batch_ids: &BatchIndices, + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] constraints: &[TwoBodyConstraint], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] contact_offsets: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] color_buckets: &mut [u32], + #[spirv(uniform, descriptor_set = 0, binding = 4)] batch_ids: &BatchIndices, ) { let num_threads = num_workgroups.x * WORKGROUP_SIZE; - let batch_id = invocation_id.y; + let nb = batch_ids.num_batches; let stride = batch_ids.solver_color_buckets_stride; + let total = contact_offsets.read(batch_ids.num_batches as usize); + let constraints = Slice(constraints, 0); - let constraints_colors = batch_ids.contact_batch(batch_id, constraints_colors); - let len = contacts_len - .read(batch_id as usize) - .min(batch_ids.contacts_batch_capacity); - - for i in StepRng::new(invocation_id.x..len, num_threads) { - let color = constraints_colors[i as usize]; - // Colors past the swept range (can happen if the bounded coloring - // didn't converge) are dropped; they were never solved before either. - if color < stride - 1 { - atomic_add_u32(color_counts.at_mut((batch_id * stride + color) as usize), 1); + for i in StepRng::new(invocation_id.x..total, num_threads) { + let color = constraints_colors.read(i as usize); + if color != 0 && color < stride - 1 { + let batch = batch_ids.collider_batch(constraints[i as usize].solver_body_a); + atomic_add_u32(color_buckets.at_mut((color * nb + batch) as usize), 1); } } } -/// Per-batch serial exclusive prefix sum over the (hopefully very small) per-color counts, -/// producing bucket start offsets. -#[spirv_bindgen] -#[spirv(compute(threads(1)))] -pub fn gpu_color_buckets_scan( - #[spirv(workgroup_id)] workgroup_id: UVec3, - #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] color_counts: &[u32], - #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] color_starts: &mut [u32], - #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] color_cursors: &mut [u32], - #[spirv(uniform, descriptor_set = 0, binding = 3)] batch_ids: &BatchIndices, -) { - let stride = batch_ids.solver_color_buckets_stride; - let batch_id = workgroup_id.y; - let base = (batch_id * stride) as usize; - - let mut acc = 0u32; - for c in 0..stride as usize { - color_starts.write(base + c, acc); - color_cursors.write(base + c, acc); - acc += color_counts.read(base + c); - } -} - -/// Scatters each constraint index into its color's bucket. #[spirv_bindgen] #[spirv(compute(threads(64)))] pub fn gpu_color_buckets_scatter( #[spirv(global_invocation_id)] invocation_id: UVec3, #[spirv(num_workgroups)] num_workgroups: UVec3, #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] constraints_colors: &[u32], - #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] contacts_len: &[u32], - #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] color_cursors: &mut [u32], - #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] color_sorted_ids: &mut [u32], - #[spirv(uniform, descriptor_set = 0, binding = 4)] batch_ids: &BatchIndices, + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] constraints: &[TwoBodyConstraint], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] contact_offsets: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] color_buckets: &mut [u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] color_sorted_ids: &mut [u32], + #[spirv(uniform, descriptor_set = 0, binding = 5)] batch_ids: &BatchIndices, ) { let num_threads = num_workgroups.x * WORKGROUP_SIZE; - let batch_id = invocation_id.y; + let nb = batch_ids.num_batches; let stride = batch_ids.solver_color_buckets_stride; + let total = contact_offsets.read(batch_ids.num_batches as usize); + let constraints = Slice(constraints, 0); - let constraints_colors = batch_ids.contact_batch(batch_id, constraints_colors); - let mut color_sorted_ids = batch_ids.contact_batch_mut(batch_id, color_sorted_ids); - let len = contacts_len - .read(batch_id as usize) - .min(batch_ids.contacts_batch_capacity); - - for i in StepRng::new(invocation_id.x..len, num_threads) { - let color = constraints_colors[i as usize]; - if color < stride - 1 { - let dst = atomic_add_u32( - color_cursors.at_mut((batch_id * stride + color) as usize), - 1, - ); - color_sorted_ids[dst as usize] = i; + for i in StepRng::new(invocation_id.x..total, num_threads) { + let color = constraints_colors.read(i as usize); + if color != 0 && color < stride - 1 { + let batch = batch_ids.collider_batch(constraints[i as usize].solver_body_a); + let dst = atomic_add_u32(color_buckets.at_mut((color * nb + batch) as usize), 1); + color_sorted_ids.write(dst as usize, i); } } } diff --git a/src_rbd_shaders/dynamics/coloring.rs b/src_rbd_shaders/dynamics/coloring.rs index 38d5fed3..9ab942ed 100644 --- a/src_rbd_shaders/dynamics/coloring.rs +++ b/src_rbd_shaders/dynamics/coloring.rs @@ -10,7 +10,7 @@ use khal_std::{ sync::{atomic_add_u32, atomic_max_u32}, }; -use crate::utils::{BatchIndices, Slice}; +use crate::utils::{BatchIndices, Slice, SliceMut}; use khal_std::index::MaybeIndexUnchecked; use super::constraint::TwoBodyConstraint; @@ -47,22 +47,22 @@ pub fn gpu_reset_luby( #[spirv(global_invocation_id)] invocation_id: UVec3, #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] constraints_colors: &mut [u32], #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] constraints_rands: &mut [u32], - #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] contacts_len: &[u32], - #[spirv(uniform, descriptor_set = 0, binding = 3)] batch_ids: &BatchIndices, + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] constraints: &[TwoBodyConstraint], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] contact_offsets: &[u32], + #[spirv(uniform, descriptor_set = 0, binding = 4)] batch_ids: &BatchIndices, ) { - let batch_id = invocation_id.y; - let mut constraints_colors = batch_ids.contact_batch_mut(batch_id, constraints_colors); - let mut constraints_rands = batch_ids.contact_batch_mut(batch_id, constraints_rands); - let len = contacts_len.read(batch_id as usize); - + let total = contact_offsets.read(batch_ids.num_batches as usize); let i = invocation_id.x; - if i < len { + if i < total { let idx = i as usize; - // Mark as uncolored - constraints_colors[idx] = MAX_U32; + if constraints.at(idx).len == 0 { + constraints_colors.write(idx, 0); + } else { + constraints_colors.write(idx, MAX_U32); + } // Assign random weight - constraints_rands[idx] = hash(i); + constraints_rands.write(idx, hash(i)); } } @@ -80,24 +80,22 @@ pub fn gpu_step_graph_coloring_luby( #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] uncolored: &mut u32, #[spirv(storage_buffer, descriptor_set = 0, binding = 6)] body_group: &[u32], #[spirv(uniform, descriptor_set = 0, binding = 7)] curr_color: &u32, - #[spirv(storage_buffer, descriptor_set = 0, binding = 8)] contacts_len: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 8)] contact_offsets: &[u32], #[spirv(uniform, descriptor_set = 0, binding = 9)] batch_ids: &BatchIndices, ) { let num_threads = num_workgroups.x * WORKGROUP_SIZE; - let batch_id = invocation_id.y; - let bci_start = batch_id as usize * 2 * batch_ids.contacts_batch_capacity as usize; - let body_constraint_counts = batch_ids.coll_batch(batch_id, body_constraint_counts); - let body_constraint_ids = Slice(body_constraint_ids, bci_start); - let body_group = batch_ids.coll_batch(batch_id, body_group); - let constraints = batch_ids.contact_batch(batch_id, constraints); - let mut constraints_colors = batch_ids.contact_batch_mut(batch_id, constraints_colors); - let constraints_rands = batch_ids.contact_batch(batch_id, constraints_rands); + let total = contact_offsets.read(batch_ids.num_batches as usize); + let body_constraint_counts = Slice(body_constraint_counts, 0); + let body_constraint_ids = Slice(body_constraint_ids, 0); + let body_group = Slice(body_group, 0); + let constraints = Slice(constraints, 0); + let mut constraints_colors = SliceMut(constraints_colors, 0); + let constraints_rands = Slice(constraints_rands, 0); - let len = contacts_len.read(batch_id as usize); let color = *curr_color; - for constraint_i in StepRng::new(invocation_id.x..len, num_threads) { + for constraint_i in StepRng::new(invocation_id.x..total, num_threads) { let i = constraint_i as usize; if constraints_colors[i] == MAX_U32 { @@ -185,21 +183,19 @@ pub fn gpu_reset_topo_gc( #[spirv(global_invocation_id)] invocation_id: UVec3, #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] constraints_colors: &mut [u32], #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] colored: &mut [u32], - #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] contacts_len: &[u32], - #[spirv(uniform, descriptor_set = 0, binding = 3)] batch_ids: &BatchIndices, + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] constraints: &[TwoBodyConstraint], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] contact_offsets: &[u32], + #[spirv(uniform, descriptor_set = 0, binding = 4)] batch_ids: &BatchIndices, ) { - let batch_id = invocation_id.y; - let mut constraints_colors = batch_ids.contact_batch_mut(batch_id, constraints_colors); - let mut colored = batch_ids.contact_batch_mut(batch_id, colored); - let len = contacts_len.read(batch_id as usize); - + let total = contact_offsets.read(batch_ids.num_batches as usize); let i = invocation_id.x; - if i < len { + if i < total { let idx = i as usize; // Color 0 is reserved for "uncolored" state - constraints_colors[idx] = 0; - colored[idx] = 0; + constraints_colors.write(idx, 0); + let inert = if constraints.at(idx).len == 0 { 1 } else { 0 }; + colored.write(idx, inert); } } @@ -237,24 +233,21 @@ pub fn gpu_step_graph_coloring_topo_gc( #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] constraints_colors: &mut [u32], #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] colored: &mut [u32], #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] num_colors: &mut u32, - #[spirv(storage_buffer, descriptor_set = 0, binding = 6)] contacts_len: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 6)] contact_offsets: &[u32], #[spirv(storage_buffer, descriptor_set = 0, binding = 7)] body_group: &[u32], #[spirv(uniform, descriptor_set = 0, binding = 8)] batch_ids: &BatchIndices, ) { let num_threads = num_workgroups.x * WORKGROUP_SIZE; - let batch_id = invocation_id.y; - let bci_start = batch_id as usize * 2 * batch_ids.contacts_batch_capacity as usize; - - let body_constraint_counts = batch_ids.coll_batch(batch_id, body_constraint_counts); - let body_constraint_ids = Slice(body_constraint_ids, bci_start); - let body_group = batch_ids.coll_batch(batch_id, body_group); - let constraints = batch_ids.contact_batch(batch_id, constraints); - let mut constraints_colors = batch_ids.contact_batch_mut(batch_id, constraints_colors); - let mut colored = batch_ids.contact_batch_mut(batch_id, colored); - let len = contacts_len.read(batch_id as usize); + let total = contact_offsets.read(batch_ids.num_batches as usize); + let body_constraint_counts = Slice(body_constraint_counts, 0); + let body_constraint_ids = Slice(body_constraint_ids, 0); + let body_group = Slice(body_group, 0); + let constraints = Slice(constraints, 0); + let mut constraints_colors = SliceMut(constraints_colors, 0); + let mut colored = SliceMut(colored, 0); - for constraint_i in StepRng::new(invocation_id.x..len, num_threads) { + for constraint_i in StepRng::new(invocation_id.x..total, num_threads) { let i = constraint_i as usize; if colored[i] == 0 { @@ -330,25 +323,25 @@ pub fn gpu_fix_conflicts_topo_gc( #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] constraints_colors: &[u32], #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] colored: &mut [u32], #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] num_colors: &mut u32, - #[spirv(storage_buffer, descriptor_set = 0, binding = 6)] contacts_len: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 6)] contact_offsets: &[u32], #[spirv(storage_buffer, descriptor_set = 0, binding = 7)] body_group: &[u32], #[spirv(uniform, descriptor_set = 0, binding = 8)] batch_ids: &BatchIndices, ) { let num_threads = num_workgroups.x * WORKGROUP_SIZE; - let batch_id = invocation_id.y; - let bci_start = batch_id as usize * 2 * batch_ids.contacts_batch_capacity as usize; - let body_constraint_counts = batch_ids.coll_batch(batch_id, body_constraint_counts); - let body_constraint_ids = Slice(body_constraint_ids, bci_start); - let body_group = batch_ids.coll_batch(batch_id, body_group); - let constraints = batch_ids.contact_batch(batch_id, constraints); - let constraints_colors = batch_ids.contact_batch(batch_id, constraints_colors); - let mut colored = batch_ids.contact_batch_mut(batch_id, colored); + let total = contact_offsets.read(batch_ids.num_batches as usize); + let body_constraint_counts = Slice(body_constraint_counts, 0); + let body_constraint_ids = Slice(body_constraint_ids, 0); + let body_group = Slice(body_group, 0); + let constraints = Slice(constraints, 0); + let constraints_colors = Slice(constraints_colors, 0); + let mut colored = SliceMut(colored, 0); - let len = contacts_len.read(batch_id as usize); - - for constraint_i in StepRng::new(invocation_id.x..len, num_threads) { + for constraint_i in StepRng::new(invocation_id.x..total, num_threads) { let i = constraint_i as usize; + if constraints[i].len == 0 { + continue; + } let color_i = constraints_colors[i]; // NOTE: this `num_colors` read doesn't need to be atomic. Any non-zero value is indicative of a finished diff --git a/src_rbd_shaders/dynamics/joint_constraint.rs b/src_rbd_shaders/dynamics/joint_constraint.rs index e32a4d8e..be7409f7 100644 --- a/src_rbd_shaders/dynamics/joint_constraint.rs +++ b/src_rbd_shaders/dynamics/joint_constraint.rs @@ -12,11 +12,12 @@ use glamx::Vec2; use glamx::{Mat4, Vec2}; use khal_std::glamx::UVec3; +use khal_std::index::MaybeIndexUnchecked; use khal_std::iter::StepRng; use khal_std::macros::{spirv, spirv_bindgen}; use crate::Pose; -use crate::utils::{BatchIndices, Slice}; +use crate::utils::{BatchIndices, ISlice, ISliceMut, Slice}; use super::body::{LocalMassProperties, Velocity, WorldMassProperties}; use super::joint::ImpulseJoint; @@ -174,18 +175,17 @@ pub fn gpu_init_joint_constraints( #[spirv(uniform, descriptor_set = 0, binding = 4)] batch_ids: &BatchIndices, ) { let num_threads = num_workgroups.x * WORKGROUP_SIZE; - let batch_id = invocation_id.y; + let nb = batch_ids.num_batches; - let joints = batch_ids.impulse_joints_batch(batch_id, joints); - let mut builders = batch_ids.impulse_joints_batch_mut(batch_id, builders); - let mut constraints = batch_ids.impulse_joints_batch_mut(batch_id, constraints); - let local_mprops = batch_ids.coll_batch(batch_id, local_mprops); + let total = batch_ids.impulse_joints_len * nb; - let len = batch_ids.impulse_joints_len; - - for i in StepRng::new(invocation_id.x..len, num_threads) { - let idx = i as usize; - let joint = &joints[idx]; + for t in StepRng::new(invocation_id.x..total, num_threads) { + let idx = t as usize; + let batch_id = t % nb; + let bix = batch_ids.body_ix(batch_id); + let joint = joints.at(idx); + let body_a = bix.at(joint.body_a) as u32; + let body_b = bix.at(joint.body_b) as u32; // Mirror rapier `GenericJoint::transform_to_solver_body_space`: the // joint's local anchor frames are expressed relative to the body's @@ -193,22 +193,26 @@ pub fn gpu_init_joint_constraints( // body local center of mass from each anchor's translation. // TODO: handle the rapier "is_fixed" branch (`local_frame = body_pose * local_frame`). let mut joint_data = joint.data; - joint_data.local_frame_a.translation -= local_mprops[joint.body_a as usize].com; - joint_data.local_frame_b.translation -= local_mprops[joint.body_b as usize].com; - - builders[idx] = JointConstraintBuilder { - body1: joint.body_a, - body2: joint.body_b, - joint_id: i, - joint: joint_data, - constraint_id: i, - }; - - constraints[idx].solver_vel_a = joint.body_a; - constraints[idx].solver_vel_b = joint.body_b; - constraints[idx].im_a = local_mprops[joint.body_a as usize].inv_mass; - constraints[idx].im_b = local_mprops[joint.body_b as usize].inv_mass; - constraints[idx].len = 0; // Constraint elements will be filled later. + joint_data.local_frame_a.translation -= local_mprops.at(body_a as usize).com; + joint_data.local_frame_b.translation -= local_mprops.at(body_b as usize).com; + + builders.write( + idx, + JointConstraintBuilder { + body1: body_a, + body2: body_b, + joint_id: t, + joint: joint_data, + constraint_id: t, + }, + ); + + let cons = constraints.at_mut(idx); + cons.solver_vel_a = body_a; + cons.solver_vel_b = body_b; + cons.im_a = local_mprops.at(body_a as usize).inv_mass; + cons.im_b = local_mprops.at(body_b as usize).inv_mass; + cons.len = 0; } } @@ -226,18 +230,26 @@ pub fn gpu_update_joint_constraints( #[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 builders = batch_ids.impulse_joints_batch(batch_id, builders); - let mut constraints = batch_ids.impulse_joints_batch_mut(batch_id, constraints); - let poses = batch_ids.coll_batch(batch_id, poses); - let mprops = batch_ids.coll_batch(batch_id, mprops); - let len = batch_ids.impulse_joints_len; + let total = batch_ids.impulse_joints_len * batch_ids.num_batches; + let poses = ISlice { + buf: poses, + base: 0, + stride: 1, + shift: 0, + }; + let mprops = ISlice { + buf: mprops, + base: 0, + stride: 1, + shift: 0, + }; - for i in StepRng::new(invocation_id.x..len, num_threads) { - let idx = i as usize; - builders[idx].update_constraint(&mut constraints[idx], &poses, &mprops, params); + for t in StepRng::new(invocation_id.x..total, num_threads) { + let idx = t as usize; + builders + .at(idx) + .update_constraint(constraints.at_mut(idx), &poses, &mprops, params); } } @@ -255,10 +267,13 @@ pub fn gpu_solve_joint_constraints( #[spirv(uniform, descriptor_set = 0, binding = 5)] use_bias: &u32, ) { let num_threads = num_workgroups.x * WORKGROUP_SIZE; - let batch_id = invocation_id.y; - - let mut constraints = batch_ids.impulse_joints_batch_mut(batch_id, constraints); - let mut solver_vels = batch_ids.coll_batch_mut(batch_id, solver_vels); + let nb = batch_ids.num_batches; + let mut solver_vels = ISliceMut { + buf: solver_vels, + base: 0, + stride: 1, + shift: 0, + }; let use_bias = *use_bias != 0; let color = *curr_color as usize; @@ -273,7 +288,9 @@ pub fn gpu_solve_joint_constraints( }; let end = color_groups[color]; - for i in StepRng::new(start + invocation_id.x..end, num_threads) { - constraints[i as usize].solve_joint_constraint(&mut solver_vels, use_bias); + for t in StepRng::new(start * nb + invocation_id.x..end * nb, num_threads) { + constraints + .at_mut(t as usize) + .solve_joint_constraint(&mut solver_vels, use_bias); } } diff --git a/src_rbd_shaders/dynamics/joint_constraint_builder.rs b/src_rbd_shaders/dynamics/joint_constraint_builder.rs index 4630a632..e34959f9 100644 --- a/src_rbd_shaders/dynamics/joint_constraint_builder.rs +++ b/src_rbd_shaders/dynamics/joint_constraint_builder.rs @@ -11,7 +11,7 @@ use super::sim_params::{RbdSimParams, TWO_PI}; use crate::Rotation; #[cfg(feature = "dim2")] use crate::rotation_angle; -use crate::utils::{Slice, SliceMut}; +use crate::utils::{ISlice, ISliceMut}; use crate::{AngVector, ColumnIndex, MAX_FLT, Pose, Vector, gdot, rotation_to_matrix}; use khal_std::index::MaybeIndexUnchecked; @@ -344,7 +344,7 @@ impl JointConstraintHelper { impl JointConstraint { /// Solves a joint constraint. - pub fn solve_joint_constraint(&mut self, solver_vels: &mut SliceMut, use_bias: bool) { + pub fn solve_joint_constraint(&mut self, solver_vels: &mut ISliceMut, use_bias: bool) { let mut solver_vel1 = solver_vels[self.solver_vel_a as usize]; let mut solver_vel2 = solver_vels[self.solver_vel_b as usize]; @@ -710,8 +710,8 @@ impl JointConstraintBuilder { pub fn update_constraint( &self, constraint: &mut JointConstraint, - poses: &Slice, - mprops: &Slice, + poses: &ISlice, + mprops: &ISlice, params: &RbdSimParams, ) { // NOTE: right now, the "update", is basically reconstructing all the diff --git a/src_rbd_shaders/dynamics/mprops_update.rs b/src_rbd_shaders/dynamics/mprops_update.rs index 7d621ee8..78a0dc39 100644 --- a/src_rbd_shaders/dynamics/mprops_update.rs +++ b/src_rbd_shaders/dynamics/mprops_update.rs @@ -3,6 +3,7 @@ //! This module contains the actual GPU compute shader entry points for mass properties update. use khal_std::glamx::UVec3; +use khal_std::index::MaybeIndexUnchecked; use khal_std::iter::StepRng; use khal_std::macros::{spirv, spirv_bindgen}; @@ -26,17 +27,13 @@ pub fn gpu_update_mprops( #[spirv(uniform, descriptor_set = 0, binding = 3)] batch_ids: &BatchIndices, ) { let num_threads = num_workgroups.x * WORKGROUP_SIZE; - let batch_id = invocation_id.y; - let num_bodies = batch_ids.bodies_len; - let mut mprops = batch_ids.coll_batch_mut(batch_id, mprops); - let local_mprops = batch_ids.coll_batch(batch_id, local_mprops); - let poses = batch_ids.coll_batch(batch_id, poses); + let num_bodies = batch_ids.bodies_len * batch_ids.num_batches; for i in StepRng::new(invocation_id.x..num_bodies, num_threads) { let idx = i as usize; - let new_mprops = local_mprops[idx].to_world(&poses[idx]); - mprops[idx] = new_mprops; + let new_mprops = local_mprops.at(idx).to_world(poses.at(idx)); + mprops.write(idx, new_mprops); } } @@ -59,17 +56,12 @@ pub fn gpu_sync_collider_poses( #[spirv(uniform, descriptor_set = 0, binding = 4)] batch_ids: &BatchIndices, ) { let num_threads = num_workgroups.x * WORKGROUP_SIZE; - let batch_id = invocation_id.y; - let num_colliders = batch_ids.colliders_len; - let body_poses = batch_ids.coll_batch(batch_id, body_poses); - let collider_local_poses = batch_ids.coll_batch(batch_id, collider_local_poses); - let mut collider_world_poses = batch_ids.coll_batch_mut(batch_id, collider_world_poses); - let collider_parent = batch_ids.coll_batch(batch_id, collider_parent); + let num_colliders = batch_ids.colliders_len * batch_ids.num_batches; for i in StepRng::new(invocation_id.x..num_colliders, num_threads) { let idx = i as usize; - let body = collider_parent[idx] as usize; - collider_world_poses[idx] = body_poses[body] * collider_local_poses[idx]; + let body = collider_parent.read(idx) as usize; + collider_world_poses.write(idx, *body_poses.at(body) * *collider_local_poses.at(idx)); } } diff --git a/src_rbd_shaders/dynamics/multibody/compute_dynamics_pre.rs b/src_rbd_shaders/dynamics/multibody/compute_dynamics_pre.rs index 8dac2a8b..5e48b898 100644 --- a/src_rbd_shaders/dynamics/multibody/compute_dynamics_pre.rs +++ b/src_rbd_shaders/dynamics/multibody/compute_dynamics_pre.rs @@ -23,12 +23,13 @@ use super::ws_soa::{ use crate::dynamics::body::Velocity; use crate::dynamics::joint::SPATIAL_DIM; use crate::utils::linalg::{ - axpy_mat_par, copy_from_par, fill_par, gemm_inertia_lhs_par, gemm_omega_skew_tr_cross_buf_par, - gemm_skew_tr_lhs_cross_buf_par, gemm_skew_tr_lhs_par, gemm_tr_par, quadform_spatial_par, + MatSlice, axpy_mat_par, copy_from_par, fill_par, gemm_inertia_lhs_par, + gemm_omega_skew_tr_cross_buf_par, gemm_skew_tr_lhs_cross_buf_par, gemm_skew_tr_lhs_par, + gemm_tr_par, quadform_spatial_par, }; #[cfg(feature = "dim3")] use crate::utils::linalg::{gemm_inertia_lhs_cross_buf_par, gemm_skew_lhs_cross_buf_par}; -use crate::utils::{BatchIndices, ISlice, SliceMut}; +use crate::utils::{BatchIndices, ISlice, ISliceMut}; use crate::{ANG_DIM, AngVector, DIM, Pose, Vector, gcross_av}; use parry::math::VectorExt; @@ -88,19 +89,30 @@ pub fn gpu_mb_compute_dynamics_pre( }; let num_links = mb.num_links; let ndofs = mb.ndofs; - let mb_jac_base = mb.jacobian_offset as usize; - let mb_mm_base = mb.mass_matrix_offset as usize; - let mb_cor_base = mb.coriolis_offset as usize; - let mb_cor_w_base = batch_ids.coriolis_batch_capacity as usize + mb_cor_base; - let mb_icdt_base = - 2 * batch_ids.coriolis_batch_capacity as usize + mb.i_coriolis_dt_offset as usize; + let jac0 = batch_ids.mb_region( + batch_id, + mb.jacobian_offset, + num_links * SPATIAL_DIM as u32 * ndofs, + ); + let cor_len = num_links * DIM * ndofs; + let cor_v0 = batch_ids.mb_region(batch_id, mb.coriolis_offset, cor_len); + let cor_w0 = batch_ids.mb_region( + batch_id, + batch_ids.coriolis_batch_capacity + mb.coriolis_offset, + cor_len, + ); + let icdt0 = batch_ids.mb_region( + batch_id, + 2 * batch_ids.coriolis_batch_capacity + mb.i_coriolis_dt_offset, + SPATIAL_DIM as u32 * ndofs, + ); let vel_base = mb.first_dof 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 mut poses_slice = batch_ids.coll_batch_mut(batch_id, poses); + let mut poses_slice = batch_ids.ib_mut(batch_id, poses); let damping_slice = batch_ids .ib(batch_id, dof_state) .offset(batch_ids.dof_batch_capacity as usize + vel_base); @@ -136,7 +148,7 @@ pub fn gpu_mb_compute_dynamics_pre( update_body_jacobians( lane, t, - mb_jac_base, + jac0, ndofs, num_links, batch_ids.mb_max_links, @@ -144,8 +156,6 @@ pub fn gpu_mb_compute_dynamics_pre( links_workspace, wa, body_jacobians, - batch_ids, - batch_id, ); // 3) Propagate velocities (single-threaded) @@ -158,20 +168,29 @@ pub fn gpu_mb_compute_dynamics_pre( // the plain one for constraints and the coriolis-augmented acc section for // the acceleration solve. Otherwise only the plain matrix is built and the // coriolis blocks are skipped entirely (forces stay explicit). - let acc_section = batch_ids.mass_matrix_acc_section_offset as usize; + let acc_section = batch_ids.mass_matrix_acc_section_offset; let split = acc_section != 0; + let mm_len = ndofs * ndofs; + let plain_mass = MatSlice::dense( + batch_ids.mb_region(batch_id, mb.mass_matrix_offset, mm_len), + ndofs, + ndofs, + ); let acc_augmented_mass = if split { - batch_ids.imat(batch_id, acc_section + mb_mm_base, ndofs, ndofs) + MatSlice::dense( + batch_ids.mb_region(batch_id, acc_section + mb.mass_matrix_offset, mm_len), + ndofs, + ndofs, + ) } else { - batch_ids.imat(batch_id, mb_mm_base, ndofs, ndofs) + plain_mass }; - let plain_mass = batch_ids.imat(batch_id, mb_mm_base, ndofs, ndofs); fill_par(mass_matrices, acc_augmented_mass, 0.0, lane, t); if split { fill_par(mass_matrices, plain_mass, 0.0, lane, t); } - let i_coriolis_dt_view = batch_ids.imat(batch_id, mb_icdt_base, SPATIAL_DIM as u32, ndofs); + let i_coriolis_dt_view = MatSlice::dense(icdt0, SPATIAL_DIM as u32, ndofs); let i_coriolis_dt_v = i_coriolis_dt_view.fixed_rows(0, DIM); let i_coriolis_dt_w = i_coriolis_dt_view.fixed_rows(DIM, ANG_DIM); @@ -188,18 +207,16 @@ pub fn gpu_mb_compute_dynamics_pre( inv_mass_x = lmp.inv_mass.x; if split && inv_mass_x == 0.0 { - let coriolis_block = batch_ids.imat( - batch_id, - mb_cor_base + (k as usize) * (DIM as usize) * (ndofs as usize), + let coriolis_block = MatSlice::dense( + cor_v0 + (k as usize) * (DIM as usize) * (ndofs as usize), DIM, ndofs, ); fill_par(coriolis_packed, coriolis_block, 0.0, lane, t); fill_par( coriolis_packed, - batch_ids.imat( - batch_id, - mb_cor_w_base + (k as usize) * (DIM as usize) * (ndofs as usize), + MatSlice::dense( + cor_w0 + (k as usize) * (DIM as usize) * (ndofs as usize), DIM, ndofs, ), @@ -214,21 +231,18 @@ pub fn gpu_mb_compute_dynamics_pre( sync_slots(t); let loop_is_active = k < num_links && inv_mass_x != 0.0; - let coriolis_v_i = batch_ids.imat( - batch_id, - mb_cor_base + (k as usize) * (DIM as usize) * (ndofs as usize), + let coriolis_v_i = MatSlice::dense( + cor_v0 + (k as usize) * (DIM as usize) * (ndofs as usize), DIM, ndofs, ); - let coriolis_w_i = batch_ids.imat( - batch_id, - mb_cor_w_base + (k as usize) * (DIM as usize) * (ndofs as usize), + let coriolis_w_i = MatSlice::dense( + cor_w0 + (k as usize) * (DIM as usize) * (ndofs as usize), ANG_DIM, ndofs, ); - let body_jacobian = batch_ids.imat( - batch_id, - mb_jac_base + (k as usize) * SPATIAL_DIM * (ndofs as usize), + let body_jacobian = MatSlice::dense( + jac0 + (k as usize) * SPATIAL_DIM * (ndofs as usize), SPATIAL_DIM as u32, ndofs, ); @@ -261,22 +275,19 @@ pub fn gpu_mb_compute_dynamics_pre( if split && k != 0 { let stat = stat_slice[k as usize]; let parent_id = stat.parent_link_id; - let parent_j = batch_ids.imat( - batch_id, - mb_jac_base + (parent_id as usize) * SPATIAL_DIM * (ndofs as usize), + let parent_j = MatSlice::dense( + jac0 + (parent_id as usize) * SPATIAL_DIM * (ndofs as usize), SPATIAL_DIM as u32, ndofs, ); let parent_j_w = parent_j.fixed_rows(DIM, ANG_DIM); - let parent_coriolis_v = batch_ids.imat( - batch_id, - mb_cor_base + (parent_id as usize) * (DIM as usize) * (ndofs as usize), + let parent_coriolis_v = MatSlice::dense( + cor_v0 + (parent_id as usize) * (DIM as usize) * (ndofs as usize), DIM, ndofs, ); - let parent_coriolis_w = batch_ids.imat( - batch_id, - mb_cor_w_base + (parent_id as usize) * (DIM as usize) * (ndofs as usize), + let parent_coriolis_w = MatSlice::dense( + cor_w0 + (parent_id as usize) * (DIM as usize) * (ndofs as usize), ANG_DIM, ndofs, ); @@ -613,7 +624,7 @@ fn jacobian_mul_coordinates( fn forward_kinematics( mb: &MultibodyInfo, stat_slice: &ISlice, - poses_slice: &mut SliceMut, + poses_slice: &mut ISliceMut, ws: &mut [Vec4], wa: WsAddr, num_links: u32, @@ -669,7 +680,7 @@ fn update_body_jacobians( lane: u32, // Lanes owned by this multibody's slot (`BatchIndices::mb_pack_lanes`). lanes: u32, - mb_jac_base: usize, + jac0: usize, ndofs: u32, num_links: u32, max_links: u32, @@ -677,8 +688,6 @@ fn update_body_jacobians( ws: &[Vec4], wa: WsAddr, body_jacobians: &mut [f32], - batch_ids: &BatchIndices, - batch_id: u32, ) { // TODO(PERF): instead of copying the body jacobian over and over for each body, we should // precompute a bit set that indicates which dofs are part of the kinematic tree @@ -686,9 +695,8 @@ fn update_body_jacobians( // value per node. for k in 0..max_links { let mut parent_to_world = Pose::default(); - let link_j = batch_ids.imat( - batch_id, - mb_jac_base + (k as usize) * SPATIAL_DIM * (ndofs as usize), + let link_j = MatSlice::dense( + jac0 + (k as usize) * SPATIAL_DIM * (ndofs as usize), SPATIAL_DIM as u32, ndofs, ); @@ -697,10 +705,8 @@ fn update_body_jacobians( let link_infos = &stat_slice[k as usize]; if k != 0 { - let parent_j = batch_ids.imat( - batch_id, - mb_jac_base - + (link_infos.parent_link_id as usize) * SPATIAL_DIM * (ndofs as usize), + let parent_j = MatSlice::dense( + jac0 + (link_infos.parent_link_id as usize) * SPATIAL_DIM * (ndofs as usize), SPATIAL_DIM as u32, ndofs, ); diff --git a/src_rbd_shaders/dynamics/multibody/contact_constraints.rs b/src_rbd_shaders/dynamics/multibody/contact_constraints.rs index 46211483..6557e4c4 100644 --- a/src_rbd_shaders/dynamics/multibody/contact_constraints.rs +++ b/src_rbd_shaders/dynamics/multibody/contact_constraints.rs @@ -23,15 +23,14 @@ use khal_std::sync::workgroup_memory_barrier_with_group_sync; use crate::dynamics::ConstraintSoftness; use crate::dynamics::body::{Velocity, WorldMassProperties}; use crate::dynamics::joint::SPATIAL_DIM; -use crate::queries::{IndexedManifold, MAX_MANIFOLD_POINTS}; -use crate::utils::BatchIndices; +use crate::queries::IndexedManifold; use crate::utils::linalg::{MAX_MB_DOFS, MatSlice, VSlice, lu_solve_in_place}; +use crate::utils::{BatchIndices, Slice}; use crate::{ANG_DIM, AngVector, DIM, Pose, Vector, gcross, gdot}; use super::types::{ - CONTACT_CONSTRAINTS_PER_POINT, MAX_MB_CONTACT_CONSTRAINTS_PER_MB, MB_CONTACT_KIND_INACTIVE, - MB_CONTACT_KIND_NORMAL, MB_CONTACT_KIND_TANGENT, MultibodyContactConstraint, MultibodyInfo, - MultibodyLinkStatic, + CONTACT_CONSTRAINTS_PER_POINT, MB_CONS_SLOT_RESERVE, MB_CONTACT_KIND_NORMAL, + MB_CONTACT_KIND_TANGENT, MultibodyContactConstraint, MultibodyInfo, MultibodyLinkStatic, }; use super::utils::zero_kinematic_dofs; use super::ws_soa::{WS_LTW, WS_WORLD_COM, WsAddr, ws_pose, ws_vec}; @@ -69,10 +68,7 @@ fn orthonormal_vector(v: Vec2) -> Vec2 { #[inline] fn fill_contact_jac_row( body_jacobians: &[f32], - mb_jac_base: usize, - // Interleave parameters of `body_jacobians` (`num_batches`, `batch_id`). - jac_stride: u32, - jac_shift: u32, + jac0: usize, ndofs: u32, link_id: u32, unit_force: Vector, @@ -83,13 +79,10 @@ fn fill_contact_jac_row( ) { // Per-link SPATIAL_DIM × ndofs jacobian (rows 0..DIM = J_v, rows // DIM..SPATIAL_DIM = J_w). - let link_jac_base = mb_jac_base + (link_id as usize) * SPATIAL_DIM * (ndofs as usize); - let link_j = MatSlice::interleaved( - link_jac_base, + let link_j = MatSlice::dense( + jac0 + (link_id as usize) * SPATIAL_DIM * (ndofs as usize), SPATIAL_DIM as u32, ndofs, - jac_stride, - jac_shift, ); let (link_j_v, link_j_w) = link_j.rows_range_pair(0, DIM, DIM, ANG_DIM); for j in 0..ndofs { @@ -125,6 +118,125 @@ fn fill_contact_jac_row( out_jacs.write(col_offset + j as usize, prev + dot); } } +#[inline(always)] +fn mb_contact_demand( + im: &IndexedManifold, + mb_idx: u32, + self_contacts_enabled: u32, + body_to_link: &[[u32; 2]], +) -> u32 { + if im.contact.len == 0 { + return 0; + } + let l1 = body_to_link.read(im.bodies.x as usize); + let l2 = body_to_link.read(im.bodies.y as usize); + let mb_on_1 = l1[0] == mb_idx; + let mb_on_2 = l2[0] == mb_idx; + if !mb_on_1 && !mb_on_2 { + return 0; + } + if l1[0] != u32::MAX && l2[0] != u32::MAX && l1[0] != l2[0] { + return 0; + } + let is_self = mb_on_1 && mb_on_2; + if is_self && self_contacts_enabled == 0 { + return 0; + } + if is_self && l1[1] == l2[1] { + return 0; + } + im.contact.len * CONTACT_CONSTRAINTS_PER_POINT +} +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn gpu_mb_count_contact_constraints( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + multibody_info: &mut [MultibodyInfo], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] contacts: &[IndexedManifold], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] body_to_link: &[[u32; 2]], + #[spirv(uniform, descriptor_set = 0, binding = 3)] batch_ids: &BatchIndices, +) { + let num_mb = batch_ids.multibodies_len; + if invocation_id.x >= num_mb * batch_ids.num_batches { + return; + } + let batch_id = invocation_id.x / num_mb; + let mb_idx = invocation_id.x % num_mb; + let mut mb = multibody_info.read(batch_ids.mbi(batch_id, mb_idx as usize)); + let mut count = 0u32; + if mb.ndofs != 0 { + let contacts_slice = Slice(contacts, mb.batch_contacts_start as usize); + for ci in 0..mb.batch_contacts_len { + count += mb_contact_demand( + contacts_slice.at(ci as usize), + mb_idx, + mb.self_contacts_enabled, + body_to_link, + ); + } + } + mb.contact_constraint_count = count; + multibody_info.write(batch_ids.mbi(batch_id, mb_idx as usize), mb); +} +#[spirv_bindgen] +#[spirv(compute(threads(1)))] +pub fn gpu_mb_cons_offsets_scan( + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + multibody_info: &mut [MultibodyInfo], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] mb_cons_demand: &mut [u32], + #[spirv(uniform, descriptor_set = 0, binding = 2)] batch_ids: &BatchIndices, +) { + let total_infos = batch_ids.multibodies_len * batch_ids.num_batches; + let capacity = batch_ids.mb_contact_constraints_capacity; + let mut demand = 0u32; + let mut reserved_total = 0u32; + for i in 0..total_infos { + let count = multibody_info.read(i as usize).contact_constraint_count; + demand += count; + reserved_total += count.min(MB_CONS_SLOT_RESERVE); + } + let mut extra_budget = if capacity > reserved_total { + capacity - reserved_total + } else { + 0 + }; + + let mut acc = 0u32; + for i in 0..total_infos { + let mut mb = multibody_info.read(i as usize); + let count = mb.contact_constraint_count; + let reserve = count.min(MB_CONS_SLOT_RESERVE); + let extra = (count - reserve).min(extra_budget); + extra_budget -= extra; + let start = acc.min(capacity); + let avail = (reserve + extra).min(capacity - start); + mb.contact_constraint_start = start; + mb.contact_constraint_count = avail; + multibody_info.write(i as usize, mb); + acc = start + avail; + } + mb_cons_demand.write(0, demand); +} +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn gpu_mb_save_prev_cons_bounds( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + multibody_info: &mut [MultibodyInfo], + #[spirv(uniform, descriptor_set = 0, binding = 1)] batch_ids: &BatchIndices, +) { + let num_mb = batch_ids.multibodies_len; + if invocation_id.x >= num_mb * batch_ids.num_batches { + return; + } + let batch_id = invocation_id.x / num_mb; + let mb_idx = invocation_id.x % num_mb; + let mut mb = multibody_info.read(batch_ids.mbi(batch_id, mb_idx as usize)); + mb.old_contact_constraint_start = mb.contact_constraint_start; + mb.old_contact_constraint_count = mb.contact_constraint_count; + multibody_info.write(batch_ids.mbi(batch_id, mb_idx as usize), mb); +} /// Pack the per-link world-space contact point into the constraint. /// @@ -172,11 +284,6 @@ pub fn gpu_mb_init_contact_constraints( // manifold, afterwards they are updated normally. let freeze_anchors = *first_substep != 0; - let cons_start = batch_ids.mb_contact_constraints_start(batch_id); - let colliders_start = batch_ids.coll_start(batch_id); - // `body_to_link` is laid out with stride = colliders_batch_capacity. - let b2l_start = colliders_start; - // Per-multibody early-out: padding multibody slots have `ndofs == 0`, // which we use here as the sentinel (replaces the `num_multibodies` // storage binding the kernel used to read). @@ -190,50 +297,31 @@ pub fn gpu_mb_init_contact_constraints( } return; } - let cons_base = cons_start + (mb_idx as usize) * (MAX_MB_CONTACT_CONSTRAINTS_PER_MB as usize); + let cons_base = mb.contact_constraint_start as usize; + let avail = mb.contact_constraint_count; let wa = WsAddr::new(mb.first_link as usize, batch_ids.num_batches, batch_id); - let contacts_slice = batch_ids.contact_batch(batch_id, contacts); - let n_contacts = mb.batch_contacts_len.min(batch_ids.contacts_batch_capacity); - let prev_count = mb.contact_constraint_count; + let contacts_slice = Slice(contacts, mb.batch_contacts_start as usize); + let n_contacts = mb.batch_contacts_len; let mut count = 0u32; for ci in 0..n_contacts { - if count + (MAX_MANIFOLD_POINTS as u32) * CONTACT_CONSTRAINTS_PER_POINT - > MAX_MB_CONTACT_CONSTRAINTS_PER_MB - { - break; - } let im = contacts_slice[ci as usize]; - if im.contact.len == 0 { + let demand = mb_contact_demand(&im, mb_idx, mb.self_contacts_enabled, body_to_link); + if demand == 0 { continue; } + if count + demand > avail { + break; + } let id1 = im.colliders.x; let b1 = im.bodies.x; let b2 = im.bodies.y; - let l1 = body_to_link.read(b2l_start + b1 as usize); - let l2 = body_to_link.read(b2l_start + b2 as usize); + let l1 = body_to_link.read(b1 as usize); + let l2 = body_to_link.read(b2 as usize); let mb_on_1 = l1[0] == mb_idx; - let mb_on_2 = l2[0] == mb_idx; - - if !mb_on_1 && !mb_on_2 { - continue; - } - // Inter-multibody contacts (each side is a DIFFERENT multibody) are - // not yet handled — skip them. Self-collisions (both sides on this - // SAME multibody) are handled below. - if l1[0] != u32::MAX && l2[0] != u32::MAX && l1[0] != l2[0] { - continue; - } - - let is_self = mb_on_1 && mb_on_2; - // Honor rapier's `Multibody::self_contacts_enabled` (MJCF - // `DISABLE_SELF_CONTACTS`): skip contacts between two links of the same - // multibody when self-contacts are disabled. - if is_self && mb.self_contacts_enabled == 0 { - continue; - } + let is_self = mb_on_1 && l2[0] == mb_idx; let (mb_link_id_a, mb_link_id_b, free_body_id) = if is_self { (l1[1], l2[1], u32::MAX) } else if mb_on_1 { @@ -242,12 +330,7 @@ pub fn gpu_mb_init_contact_constraints( (l2[1], u32::MAX, b1) }; - // Skip degenerate self-contacts on the same link. - if is_self && mb_link_id_a == mb_link_id_b { - continue; - } - - let pose1 = poses.read(colliders_start + id1 as usize); + let pose1 = poses.read(id1 as usize); let world_normal = pose1.rotation * im.contact.normal_a; let lin_jac = if is_self || mb_on_1 { world_normal @@ -259,7 +342,7 @@ pub fn gpu_mb_init_contact_constraints( let free_mp = if is_self { WorldMassProperties::default() } else { - mprops.read(colliders_start + free_body_id as usize) + mprops.read(free_body_id as usize) }; let free_im = if is_self { 0.0 } else { free_mp.inv_mass.x }; @@ -284,12 +367,12 @@ pub fn gpu_mb_init_contact_constraints( let pose_b = if is_self { ws_pose(links_workspace, wa, mb_link_id_b, WS_LTW) } else { - solver_body_poses.read(colliders_start + free_body_id as usize) + solver_body_poses.read(free_body_id as usize) }; for k in 0..im.contact.len { // One contact point produces 1 normal + (DIM-1) friction slots. - if count + CONTACT_CONSTRAINTS_PER_POINT > MAX_MB_CONTACT_CONSTRAINTS_PER_MB { + if count + CONTACT_CONSTRAINTS_PER_POINT > avail { break; } let normal_slot = count; @@ -576,11 +659,6 @@ pub fn gpu_mb_init_contact_constraints( // match scans the whole slab, so the leftovers of the previous build have // to be marked inactive. if lane == 0 { - for s in count..prev_count.min(MAX_MB_CONTACT_CONSTRAINTS_PER_MB) { - let mut stale = contact_constraints.read(cons_base + s as usize); - stale.kind = MB_CONTACT_KIND_INACTIVE; - contact_constraints.write(cons_base + s as usize, stale); - } mb.contact_constraint_count = count; multibody_info.write(batch_ids.mbi(batch_id, mb_idx as usize), mb); } @@ -597,7 +675,8 @@ pub fn gpu_mb_stash_contacts_len( #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] multibody_info: &mut [MultibodyInfo], #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] contacts_len: &[u32], - #[spirv(uniform, descriptor_set = 0, binding = 2)] batch_ids: &BatchIndices, + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] contact_offsets: &[u32], + #[spirv(uniform, descriptor_set = 0, binding = 3)] batch_ids: &BatchIndices, ) { let num_mb = batch_ids.multibodies_len; if invocation_id.x >= num_mb * batch_ids.num_batches { @@ -605,8 +684,11 @@ pub fn gpu_mb_stash_contacts_len( } let batch_id = invocation_id.x / num_mb; let mb_idx = invocation_id.x % num_mb; + let seg_start = contact_offsets.read(batch_id as usize); + let seg_end = contact_offsets.read(batch_id as usize + 1); let mut mb = multibody_info.read(batch_ids.mbi(batch_id, mb_idx as usize)); - mb.batch_contacts_len = contacts_len.read(batch_id as usize); + mb.batch_contacts_len = contacts_len.read(batch_id as usize).min(seg_end - seg_start); + mb.batch_contacts_start = seg_start; multibody_info.write(batch_ids.mbi(batch_id, mb_idx as usize), mb); } @@ -625,21 +707,10 @@ pub fn gpu_mb_snapshot_contact_warmstart( old_contact_constraints: &mut [MultibodyContactConstraint], #[spirv(uniform, descriptor_set = 0, binding = 2)] batch_ids: &BatchIndices, ) { - // One thread per (slot, multibody, batch), flattened. - const MAXC: u32 = MAX_MB_CONTACT_CONSTRAINTS_PER_MB; - let num_mb = batch_ids.multibodies_len; - let per_batch = num_mb * MAXC; - if invocation_id.x >= per_batch * batch_ids.num_batches { - return; + let i = invocation_id.x; + if i < batch_ids.mb_contact_constraints_capacity { + old_contact_constraints.write(i as usize, contact_constraints.read(i as usize)); } - let batch_id = invocation_id.x / per_batch; - let r = invocation_id.x % per_batch; - let mb_idx = r / MAXC; - let s = r % MAXC; - - let cons_start = batch_ids.mb_contact_constraints_start(batch_id); - let idx = cons_start + (mb_idx * MAXC + s) as usize; - old_contact_constraints.write(idx, contact_constraints.read(idx)); } /// Warmstart: re-apply each active contact constraint's accumulated `impulse` @@ -656,7 +727,7 @@ pub fn gpu_mb_warmstart_contact_constraints( #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] multibody_info: &[MultibodyInfo], #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] contact_constraints: &[MultibodyContactConstraint], - #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] contact_constraint_columns: &[f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] contact_jac_cols: &[f32], #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] dof_state: &mut [f32], #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] solver_vels: &mut [Velocity], #[spirv(uniform, descriptor_set = 0, binding = 5)] batch_ids: &BatchIndices, @@ -669,20 +740,15 @@ pub fn gpu_mb_warmstart_contact_constraints( return; } - let cons_start = batch_ids.mb_contact_constraints_start(batch_id); - let col_start = batch_ids.mb_contact_constraint_columns_start(batch_id); - let colliders_start = batch_ids.coll_start(batch_id); - let mb = multibody_info.read(batch_ids.mbi(batch_id, mb_idx as usize)); let ndofs = mb.ndofs; if ndofs == 0 { return; } let v_base = mb.first_dof as usize; - let cons_base = cons_start + (mb_idx as usize) * (MAX_MB_CONTACT_CONSTRAINTS_PER_MB as usize); + let cons_base = mb.contact_constraint_start as usize; let dofs_stride = batch_ids.dof_batch_capacity as usize; - let col_base = - col_start + (mb_idx as usize) * (MAX_MB_CONTACT_CONSTRAINTS_PER_MB as usize) * dofs_stride; + let jc_base = cons_base * 2 * dofs_stride; let count = mb.contact_constraint_count; // No accumulated impulses to re-apply: skip the dof round-trip. @@ -701,20 +767,20 @@ pub fn gpu_mb_warmstart_contact_constraints( let cons = contact_constraints.read(cons_base + s as usize); let imp = cons.impulse; if imp != 0.0 { - let col_offset = col_base + (s as usize) * dofs_stride; + let col_offset = jc_base + (s as usize) * 2 * dofs_stride + dofs_stride; // Multibody side: v += impulse · column (column = M⁻¹ Jᵀ). if lane < ndofs { - let col = contact_constraint_columns.read(col_offset + lane as usize); + let col = contact_jac_cols.read(col_offset + lane as usize); v_lane += imp * col; } // Free body side (skipped for self-contacts). let is_self = cons.free_body_id == u32::MAX; if lane == 0 && !is_self { - let free = solver_vels.read(colliders_start + cons.free_body_id as usize); + let free = solver_vels.read(cons.free_body_id as usize); let mut new_free = free; new_free.linear += cons.lin_jac * (cons.free_body_im * imp); new_free.angular += cons.ii_ang_jac * imp; - solver_vels.write(colliders_start + cons.free_body_id as usize, new_free); + solver_vels.write(cons.free_body_id as usize, new_free); } } } @@ -737,13 +803,11 @@ pub fn gpu_mb_finalize_contact_constraints( #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] lu_pivots: &[u32], #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] contact_constraints: &mut [MultibodyContactConstraint], - #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] contact_constraint_jacs: &mut [f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] contact_jac_cols: &mut [f32], #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] - contact_constraint_columns: &mut [f32], - #[spirv(storage_buffer, descriptor_set = 0, binding = 6)] links_static: &[MultibodyLinkStatic], - #[spirv(storage_buffer, descriptor_set = 0, binding = 7)] body_jacobians: &[f32], - #[spirv(uniform, descriptor_set = 0, binding = 8)] batch_ids: &BatchIndices, + #[spirv(storage_buffer, descriptor_set = 0, binding = 6)] body_jacobians: &[f32], + #[spirv(uniform, descriptor_set = 0, binding = 7)] batch_ids: &BatchIndices, ) { const LANES: u32 = 64; let batch_id = workgroup_id.y; @@ -754,30 +818,34 @@ pub fn gpu_mb_finalize_contact_constraints( return; } - let cons_start = batch_ids.mb_contact_constraints_start(batch_id); - let col_start = batch_ids.mb_contact_constraint_columns_start(batch_id); - let mb = multibody_info.read(batch_ids.mbi(batch_id, mb_idx as usize)); let ndofs = mb.ndofs; if ndofs == 0 { return; } - let mb_mm_base = mb.mass_matrix_offset as usize; - let mb_jac_base = mb.jacobian_offset as usize; - let piv = batch_ids.ivec(batch_id, mb.first_dof as usize); - let cons_base = cons_start + (mb_idx as usize) * (MAX_MB_CONTACT_CONSTRAINTS_PER_MB as usize); + let jac0 = batch_ids.mb_region( + batch_id, + mb.jacobian_offset, + mb.num_links * SPATIAL_DIM as u32 * ndofs, + ); + let piv = VSlice::dense(batch_ids.mb_region(batch_id, mb.first_dof, ndofs)); + let cons_base = mb.contact_constraint_start as usize; let dofs_stride = batch_ids.dof_batch_capacity as usize; - let col_base = - col_start + (mb_idx as usize) * (MAX_MB_CONTACT_CONSTRAINTS_PER_MB as usize) * dofs_stride; + let jc_base = cons_base * 2 * dofs_stride; - let m = batch_ids.imat(batch_id, mb_mm_base, ndofs, ndofs); + let m = MatSlice::dense( + batch_ids.mb_region(batch_id, mb.mass_matrix_offset, ndofs * ndofs), + ndofs, + ndofs, + ); let count = mb.contact_constraint_count; let stat_slice = batch_ids .ib(batch_id, links_static) .offset(mb.first_link as usize); for s in StepRng::new(lane..count, LANES) { - let col_offset = col_base + (s as usize) * dofs_stride; + let jac_offset = jc_base + (s as usize) * 2 * dofs_stride; + let col_offset = jac_offset + dofs_stride; let mut cons = contact_constraints.read(cons_base + s as usize); let is_self = cons.free_body_id == u32::MAX; @@ -785,29 +853,25 @@ pub fn gpu_mb_finalize_contact_constraints( // folding both touched links in for a self-contact. fill_contact_jac_row( body_jacobians, - mb_jac_base, - batch_ids.num_batches, - batch_id, + jac0, ndofs, cons.link_id, -cons.lin_jac, cons.torque_a, - contact_constraint_jacs, - col_offset, + contact_jac_cols, + jac_offset, false, ); if is_self { fill_contact_jac_row( body_jacobians, - mb_jac_base, - batch_ids.num_batches, - batch_id, + jac0, ndofs, cons.link_id_b, cons.lin_jac, cons.torque_b, - contact_constraint_jacs, - col_offset, + contact_jac_cols, + jac_offset, true, ); } @@ -815,8 +879,8 @@ pub fn gpu_mb_finalize_contact_constraints( // 2) Copy J^T row into the column buffer (it'll be overwritten by the // LU solve with the M⁻¹·Jᵀ result). for i in 0..ndofs { - let v = contact_constraint_jacs.read(col_offset + i as usize); - contact_constraint_columns.write(col_offset + i as usize, v); + let v = contact_jac_cols.read(jac_offset + i as usize); + contact_jac_cols.write(col_offset + i as usize, v); } // 3) Solve M · column = J^T (in place). lu_solve_in_place( @@ -824,21 +888,16 @@ pub fn gpu_mb_finalize_contact_constraints( m, lu_pivots, piv, - contact_constraint_columns, + contact_jac_cols, VSlice::dense(col_offset), ); // 3b) Kinematic dofs are user-driven: the impulse must not move them. - zero_kinematic_dofs( - contact_constraint_columns, - col_offset, - &stat_slice, - mb.num_links, - ); + zero_kinematic_dofs(contact_jac_cols, col_offset, &stat_slice, mb.num_links); // 4) inv_r_mb = J · column. let mut inv_r_mb = 0.0f32; for i in 0..ndofs { - let j = contact_constraint_jacs.read(col_offset + i as usize); - let c = contact_constraint_columns.read(col_offset + i as usize); + let j = contact_jac_cols.read(jac_offset + i as usize); + let c = contact_jac_cols.read(col_offset + i as usize); inv_r_mb += j * c; } // 5) Add free body's contribution: im (since lin_jac is unit) + @@ -893,8 +952,15 @@ pub fn gpu_mb_transfer_contact_warmstart( return; } - let cons_base = batch_ids.mb_contact_constraints_start(batch_id) - + (mb_idx as usize) * (MAX_MB_CONTACT_CONSTRAINTS_PER_MB as usize); + let cons_base = mb.contact_constraint_start as usize; + let old_base = mb.old_contact_constraint_start as usize; + let old_count = if mb.old_contact_constraint_start + mb.old_contact_constraint_count + <= batch_ids.mb_contact_constraints_capacity + { + mb.old_contact_constraint_count + } else { + 0 + }; let sq_threshold = MATCH_DIST * MATCH_DIST; let warmstart_coeff = softness.warmstart_coefficient; @@ -906,8 +972,8 @@ pub fn gpu_mb_transfer_contact_warmstart( continue; } - for j in 0..MAX_MB_CONTACT_CONSTRAINTS_PER_MB { - let old = old_contact_constraints.read(cons_base + j as usize); + for j in 0..old_count { + let old = old_contact_constraints.read(old_base + j as usize); if old.kind != MB_CONTACT_KIND_NORMAL || old.link_id != cons.link_id || old.link_id_b != cons.link_id_b @@ -927,8 +993,8 @@ pub fn gpu_mb_transfer_contact_warmstart( // Friction rows follow their normal row contiguously. #[cfg(feature = "dim3")] { - let old_t0 = old_contact_constraints.read(cons_base + (j + 1) as usize); - let old_t1 = old_contact_constraints.read(cons_base + (j + 2) as usize); + let old_t0 = old_contact_constraints.read(old_base + (j + 1) as usize); + let old_t1 = old_contact_constraints.read(old_base + (j + 2) as usize); let world = (-old_t0.lin_jac * old_t0.impulse - old_t1.lin_jac * old_t1.impulse) * warmstart_coeff; @@ -941,7 +1007,7 @@ pub fn gpu_mb_transfer_contact_warmstart( } #[cfg(feature = "dim2")] { - let old_t0 = old_contact_constraints.read(cons_base + (j + 1) as usize); + let old_t0 = old_contact_constraints.read(old_base + (j + 1) as usize); let mut new_t0 = contact_constraints.read(cons_base + (s + 1) as usize); new_t0.impulse = old_t0.impulse * warmstart_coeff; contact_constraints.write(cons_base + (s + 1) as usize, new_t0); @@ -965,7 +1031,7 @@ pub fn gpu_mb_seed_contact_restitution( #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] multibody_info: &[MultibodyInfo], #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] contact_constraints: &mut [MultibodyContactConstraint], - #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] contact_constraint_jacs: &[f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] contact_jac_cols: &[f32], #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] dof_state: &[f32], #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] solver_vels: &[Velocity], #[spirv(uniform, descriptor_set = 0, binding = 5)] batch_ids: &BatchIndices, @@ -985,27 +1051,24 @@ pub fn gpu_mb_seed_contact_restitution( return; } - let colliders_start = batch_ids.coll_start(batch_id); let v_base = mb.first_dof as usize; - let cons_base = batch_ids.mb_contact_constraints_start(batch_id) - + (mb_idx as usize) * (MAX_MB_CONTACT_CONSTRAINTS_PER_MB as usize); + let cons_base = mb.contact_constraint_start as usize; let dofs_stride = batch_ids.dof_batch_capacity as usize; - let col_base = batch_ids.mb_contact_constraint_columns_start(batch_id) - + (mb_idx as usize) * (MAX_MB_CONTACT_CONSTRAINTS_PER_MB as usize) * dofs_stride; + let jc_base = cons_base * 2 * dofs_stride; for s in StepRng::new(lane..count, LANES) { let mut cons = contact_constraints.read(cons_base + s as usize); if cons.kind != MB_CONTACT_KIND_NORMAL { continue; } - let jac_off = col_base + (s as usize) * dofs_stride; + let jac_off = jc_base + (s as usize) * 2 * dofs_stride; let mut j_dot_v = 0.0f32; for i in 0..ndofs { - j_dot_v += contact_constraint_jacs.read(jac_off + i as usize) + j_dot_v += contact_jac_cols.read(jac_off + i as usize) * dof_state.read(batch_ids.mbi(batch_id, v_base + i as usize)); } if cons.free_body_id != u32::MAX { - let free = solver_vels.read(colliders_start + cons.free_body_id as usize); + let free = solver_vels.read(cons.free_body_id as usize); j_dot_v += cons.lin_jac.dot(free.linear) + gdot(cons.ang_jac, free.angular); } @@ -1039,10 +1102,9 @@ pub fn gpu_mb_apply_contact_restitution( #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] multibody_info: &[MultibodyInfo], #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] contact_constraints: &mut [MultibodyContactConstraint], - #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] contact_constraint_jacs: &[f32], - #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] contact_constraint_columns: &[f32], - #[spirv(uniform, descriptor_set = 0, binding = 4)] batch_ids: &BatchIndices, - #[spirv(uniform, descriptor_set = 0, binding = 5)] max_contact_constraints: &u32, + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] contact_jac_cols: &[f32], + #[spirv(uniform, descriptor_set = 0, binding = 3)] batch_ids: &BatchIndices, + #[spirv(uniform, descriptor_set = 0, binding = 4)] max_contact_constraints: &u32, #[spirv(storage_buffer, descriptor_set = 1, binding = 0)] dof_state: &mut [f32], #[spirv(storage_buffer, descriptor_set = 1, binding = 1)] solver_vels: &mut [Velocity], #[spirv(workgroup)] dof_v: &mut [f32; MAX_MB_DOFS], @@ -1068,13 +1130,10 @@ pub fn gpu_mb_apply_contact_restitution( } let active = in_range && ndofs != 0 && count != 0; - let colliders_start = batch_ids.coll_start(batch_id); let v_base = mb.first_dof as usize; - let cons_base = batch_ids.mb_contact_constraints_start(batch_id) - + (mb_idx as usize) * (MAX_MB_CONTACT_CONSTRAINTS_PER_MB as usize); + let cons_base = mb.contact_constraint_start as usize; let dofs_stride = batch_ids.dof_batch_capacity as usize; - let col_base = batch_ids.mb_contact_constraint_columns_start(batch_id) - + (mb_idx as usize) * (MAX_MB_CONTACT_CONSTRAINTS_PER_MB as usize) * dofs_stride; + let jc_base = cons_base * 2 * dofs_stride; if active && lane < ndofs { dof_v.write( @@ -1108,15 +1167,14 @@ pub fn gpu_mb_apply_contact_restitution( if !solve { continue; } - let col_offset = col_base + (s as usize) * dofs_stride; + let jac_offset = jc_base + (s as usize) * 2 * dofs_stride; let is_self = cons.free_body_id == u32::MAX; if solve { scratch.write( lane as usize, if lane < ndofs { - contact_constraint_jacs.read(col_offset + lane as usize) - * dof_v.read(lane as usize) + contact_jac_cols.read(jac_offset + lane as usize) * dof_v.read(lane as usize) } else { 0.0 }, @@ -1132,7 +1190,7 @@ pub fn gpu_mb_apply_contact_restitution( let free = if is_self { Velocity::default() } else { - solver_vels.read(colliders_start + cons.free_body_id as usize) + solver_vels.read(cons.free_body_id as usize) }; if !is_self { j_dot_v += cons.lin_jac.dot(free.linear) + gdot(cons.ang_jac, free.angular); @@ -1151,14 +1209,14 @@ pub fn gpu_mb_apply_contact_restitution( let mut new_free = free; new_free.linear += cons.lin_jac * (cons.free_body_im * delta); new_free.angular += cons.ii_ang_jac * delta; - solver_vels.write(colliders_start + cons.free_body_id as usize, new_free); + solver_vels.write(cons.free_body_id as usize, new_free); } } workgroup_memory_barrier_with_group_sync(); let delta = *delta_shared; if solve && delta != 0.0 && lane < ndofs { - let col = contact_constraint_columns.read(col_offset + lane as usize); + let col = contact_jac_cols.read(jac_offset + dofs_stride + lane as usize); dof_v.write(lane as usize, dof_v.read(lane as usize) + delta * col); } workgroup_memory_barrier_with_group_sync(); diff --git a/src_rbd_shaders/dynamics/multibody/contact_sensor.rs b/src_rbd_shaders/dynamics/multibody/contact_sensor.rs index 03adc551..e12a9874 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::{ - MAX_MB_CONTACT_CONSTRAINTS_PER_MB, MB_CONTACT_KIND_NORMAL, MultibodyContactConstraint, + MB_CONTACT_KIND_NORMAL, MultibodyContactConstraint, MultibodyInfo, }; use crate::utils::BatchIndices; @@ -49,8 +49,7 @@ pub fn gpu_mb_sense_contact_impulses( } 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 cons_base = mb.contact_constraint_start as usize; let count = mb.contact_constraint_count; for c in 0..count { diff --git a/src_rbd_shaders/dynamics/multibody/env_reset.rs b/src_rbd_shaders/dynamics/multibody/env_reset.rs index 712df56f..1ed94c57 100644 --- a/src_rbd_shaders/dynamics/multibody/env_reset.rs +++ b/src_rbd_shaders/dynamics/multibody/env_reset.rs @@ -44,7 +44,12 @@ pub fn gpu_mb_env_reset( let dpb = params.w; if i < lpb * WS_QUADS { - links_workspace.write((i * nb + env) as usize, staging_ws.read(i as usize)); + let link = i / WS_QUADS; + let q = i % WS_QUADS; + links_workspace.write( + ((link * nb + env) * WS_QUADS + q) 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)); @@ -116,7 +121,7 @@ pub fn gpu_mb_env_reset_batch( v.y += off.y; v.z += off.z; } - links_workspace.write((i * nb + env) as usize, v); + links_workspace.write(((link * nb + env) * WS_QUADS + q) as usize, v); } } @@ -201,6 +206,7 @@ pub fn gpu_env_reset_bodies( let r = invocation_id.y; let bps = params.x; let vs = params.y; + let nb = params.w; if r >= params.z { return; } @@ -216,11 +222,11 @@ pub fn gpu_env_reset_bodies( p.translation.y += off.y; p.translation.z += off.z; } - body_poses.write((env * bps + i) as usize, p); + body_poses.write((i * nb + env) as usize, p); } if i < vs { vels.write( - (env * vs + i) as usize, + (i * nb + env) 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 fc005f5f..bdcc4f7f 100644 --- a/src_rbd_shaders/dynamics/multibody/gravity_and_lu.rs +++ b/src_rbd_shaders/dynamics/multibody/gravity_and_lu.rs @@ -18,7 +18,8 @@ use glamx::Vec4; use crate::dynamics::body::Velocity; use crate::dynamics::joint::SPATIAL_DIM; use crate::utils::linalg::{ - MAX_MB_DOFS, fill_par, gemv_tr_spatial_split_par, lu_decompose, lu_solve_in_place, + MAX_MB_DOFS, MatSlice, VSlice, fill_par, gemv_tr_spatial_split_par, lu_decompose, + lu_solve_in_place, }; use crate::utils::{BatchIndices, ISlice}; use crate::{AngVector, Vector, gcross_av}; @@ -43,9 +44,7 @@ use super::ws_soa::{ #[inline] fn apply_spring_forces( gen_forces: &mut [f32], - batch_ids: &BatchIndices, - batch_id: u32, - gen_base: usize, + gen0: usize, stat_slice: &ISlice, links_workspace: &[Vec4], wa: WsAddr, @@ -70,7 +69,7 @@ fn apply_spring_forces( if k_s != 0.0 { let q = ws_coord(links_workspace, wa, k, axis); let rest = spring_ref_slice[dof]; - let idx = batch_ids.mbi(batch_id, gen_base + dof); + let idx = gen0 + dof; let cur = gen_forces.read(idx); gen_forces.write(idx, cur - k_s * (q - rest) - k_s * dt * vel_slice[dof]); } @@ -116,10 +115,14 @@ pub fn gpu_mb_gravity_and_lu( let mb = batch_ids.ib(batch_id, multibody_info).read(mb_idx as usize); let num_links = mb.num_links; let ndofs = mb.ndofs; - let mb_jac_base = mb.jacobian_offset as usize; + let jac0 = batch_ids.mb_region( + batch_id, + mb.jacobian_offset, + num_links * SPATIAL_DIM as u32 * ndofs, + ); + let gen0 = batch_ids.mb_region(batch_id, mb.first_dof, ndofs); let gen_base = mb.first_dof as usize; - let mb_mm_base = mb.mass_matrix_offset as usize; - let piv = batch_ids.ivec(batch_id, gen_base); + let piv = VSlice::dense(gen0); let stat_slice = batch_ids .ib(batch_id, links_static) @@ -140,7 +143,7 @@ pub fn gpu_mb_gravity_and_lu( .offset(5 * batch_ids.dof_batch_capacity as usize + gen_base); // ---- Phase 1: zero the generalized-force vector (parallel across DOFs). ---- - let accelerations = batch_ids.imat(batch_id, gen_base, ndofs, 1); + let accelerations = MatSlice::dense(gen0, ndofs, 1); // TODO(perf): up to a certain number of degrees of freedom, we could actually run all the // calculations in shared memory and only write the result in the end. // Currently, the max number of dofs is 32 but we still accumulate forces/accelerations @@ -250,16 +253,15 @@ pub fn gpu_mb_gravity_and_lu( let f_lin = g * (mass * gravity_scale) + ext_force - acc_lin * mass; let f_ang = ext_torque - gyroscopic - i_acc_ang; - let body_jacobian = batch_ids.imat( - batch_id, - mb_jac_base + (k as usize) * SPATIAL_DIM * (ndofs as usize), + let body_jacobian = MatSlice::dense( + jac0 + (k as usize) * SPATIAL_DIM * (ndofs as usize), SPATIAL_DIM as u32, ndofs, ); gemv_tr_spatial_split_par( gen_forces, - batch_ids.ivec(batch_id, gen_base), + VSlice::dense(gen0), 1.0, body_jacobians, body_jacobian, @@ -281,7 +283,7 @@ pub fn gpu_mb_gravity_and_lu( workgroup_memory_barrier_with_group_sync(); let i = lane; if i < ndofs { - let idx = batch_ids.mbi(batch_id, gen_base + i as usize); + let idx = gen0 + i as usize; let cur = gen_forces.read(idx); let v = vel_slice[i as usize]; gen_forces.write(idx, cur - damping_slice[i as usize] * v); @@ -292,9 +294,7 @@ pub fn gpu_mb_gravity_and_lu( if lane == 0 { apply_spring_forces( gen_forces, - batch_ids, - batch_id, - gen_base, + gen0, &stat_slice, links_workspace, wa, @@ -311,16 +311,25 @@ pub fn gpu_mb_gravity_and_lu( // identity (see the `pre` kernel), so the solve passes the rhs through: // zeroing it here pins their velocities to the user-driven values. if i < ndofs && kin_mask_slice[i as usize] != 0.0 { - gen_forces.write(batch_ids.mbi(batch_id, gen_base + i as usize), 0.0); + gen_forces.write(gen0 + i as usize, 0.0); } workgroup_memory_barrier_with_group_sync(); // ---- Phase 3: factor the acceleration matrix, solve M·x = τ. ---- - let m_view = batch_ids.imat(batch_id, mb_mm_base, ndofs, ndofs); - let acc_section = batch_ids.mass_matrix_acc_section_offset as usize; + let mm_len = ndofs * ndofs; + let m_view = MatSlice::dense( + batch_ids.mb_region(batch_id, mb.mass_matrix_offset, mm_len), + ndofs, + ndofs, + ); + let acc_section = batch_ids.mass_matrix_acc_section_offset; let split = acc_section != 0; let m_acc_view = if split { - batch_ids.imat(batch_id, acc_section + mb_mm_base, ndofs, ndofs) + MatSlice::dense( + batch_ids.mb_region(batch_id, acc_section + mb.mass_matrix_offset, mm_len), + ndofs, + ndofs, + ) } else { m_view }; @@ -328,10 +337,7 @@ pub fn gpu_mb_gravity_and_lu( for r in 0..ndofs { mat.write(sm_idx(r, lane), mass_matrices.read(m_acc_view.idx(r, lane))); } - x.write( - lane as usize, - gen_forces.read(batch_ids.mbi(batch_id, gen_base + lane as usize)), - ); + x.write(lane as usize, gen_forces.read(gen0 + lane as usize)); } workgroup_memory_barrier_with_group_sync(); @@ -357,10 +363,7 @@ pub fn gpu_mb_gravity_and_lu( lu_triangular_solve_in_place(ndofs, max_ndofs, lane, mat, x, partial); if lane < ndofs { - gen_forces.write( - batch_ids.mbi(batch_id, gen_base + lane as usize), - x.read(lane as usize), - ); + gen_forces.write(gen0 + lane as usize, x.read(lane as usize)); } // ---- Phase 5 (split mode only): factor the plain matrix and persist its @@ -438,10 +441,14 @@ fn gravity_and_lu_packed_impl f32 { @@ -84,7 +85,7 @@ pub(super) fn side_dot_vel_par( 0.0f32 } else if kind == SIDE_KIND_BODY { if lane < SPATIAL_DIM as u32 { - let v = solver_vels.read(colliders_start + body_id as usize); + let v = solver_vels.read(bix.at(body_id)); jacobians.read(j_id as usize + lane as usize) * spatial_component(v, lane) } else { 0.0f32 @@ -132,7 +133,7 @@ pub(super) fn side_apply_impulse_par( dof_vels: &mut [f32], dof_base_for_mb: VSlice, solver_vels: &mut [Velocity], - colliders_start: usize, + bix: BodyIx, lane: u32, ) { // All operands are workgroup-uniform, so this early-out is uniform. @@ -143,7 +144,7 @@ pub(super) fn side_apply_impulse_par( let scaled = sign * delta; if kind == SIDE_KIND_BODY { if lane == 0 { - let coll_idx = colliders_start + body_id as usize; + let coll_idx = bix.at(body_id); let mut v = solver_vels.read(coll_idx); #[cfg(feature = "dim3")] { @@ -190,9 +191,9 @@ pub(super) fn fill_body_jacobians( unit_force: Vector, unit_torque: AngVector, mprops: &[WorldMassProperties], - colliders_start: usize, + bix: BodyIx, ) { - let mp = mprops.read(colliders_start + body_id as usize); + let mp = mprops.read(bix.at(body_id)); let im = mp.inv_mass; let base = j_id as usize; @@ -247,14 +248,12 @@ pub(super) fn fill_mb_jacobians( il: VSlice, ) { let ndofs = mb.ndofs; - let mb_jac_base = mb.jacobian_offset as usize; - let link_jac_base = mb_jac_base + (link_id as usize) * SPATIAL_DIM * (ndofs as usize); - let link_j = MatSlice::interleaved( - link_jac_base, + let jac0 = mb.jacobian_offset as usize * il.stride as usize + + il.shift as usize * (mb.num_links * SPATIAL_DIM as u32 * ndofs) as usize; + let link_j = MatSlice::dense( + jac0 + (link_id as usize) * SPATIAL_DIM * (ndofs as usize), SPATIAL_DIM as u32, ndofs, - il.stride, - il.shift, ); let (link_j_v, link_j_w) = link_j.rows_range_pair(0, DIM, DIM, ANG_DIM); @@ -319,11 +318,12 @@ pub(super) fn fill_relative_mb_jacobians( il: VSlice, ) { let ndofs = mb.ndofs; - let mb_jac_base = mb.jacobian_offset as usize; - let la_base = mb_jac_base + (link_a as usize) * SPATIAL_DIM * (ndofs as usize); - let lb_base = mb_jac_base + (link_b as usize) * SPATIAL_DIM * (ndofs as usize); - let la = MatSlice::interleaved(la_base, SPATIAL_DIM as u32, ndofs, il.stride, il.shift); - let lb = MatSlice::interleaved(lb_base, SPATIAL_DIM as u32, ndofs, il.stride, il.shift); + let jac0 = mb.jacobian_offset as usize * il.stride as usize + + il.shift as usize * (mb.num_links * SPATIAL_DIM as u32 * ndofs) as usize; + let la_base = jac0 + (link_a as usize) * SPATIAL_DIM * (ndofs as usize); + let lb_base = jac0 + (link_b as usize) * SPATIAL_DIM * (ndofs as usize); + let la = MatSlice::dense(la_base, SPATIAL_DIM as u32, ndofs); + let lb = MatSlice::dense(lb_base, SPATIAL_DIM as u32, ndofs); let (la_v, la_w) = la.rows_range_pair(0, DIM, DIM, ANG_DIM); let (lb_v, lb_w) = lb.rows_range_pair(0, DIM, DIM, ANG_DIM); diff --git a/src_rbd_shaders/dynamics/multibody/impulse_joint_constraints/kernels.rs b/src_rbd_shaders/dynamics/multibody/impulse_joint_constraints/kernels.rs index 66643eeb..3592a9bc 100644 --- a/src_rbd_shaders/dynamics/multibody/impulse_joint_constraints/kernels.rs +++ b/src_rbd_shaders/dynamics/multibody/impulse_joint_constraints/kernels.rs @@ -5,6 +5,7 @@ use glamx::Vec4; use khal_std::glamx::UVec3; use khal_std::index::MaybeIndexUnchecked; use khal_std::macros::{spirv, spirv_bindgen}; +use khal_std::iter::StepRng; use khal_std::sync::workgroup_memory_barrier_with_group_sync; use crate::Pose; @@ -43,11 +44,6 @@ pub fn gpu_mb_update_impulse_joint_constraints( #[spirv(uniform, descriptor_set = 0, binding = 4)] batch_ids: &BatchIndices, ) { let num_threads = num_workgroups.x * 64; - let batch_id = invocation_id.y; - let cap = batch_ids.mb_imp_joints_batch_capacity; - if invocation_id.x >= cap { - return; - } let dt = softness.dt; // Joint lock/limit softness — configurable via `joint_natural_frequency` / // `joint_damping_ratio` (rapier's `joint.softness`), replacing the old @@ -56,37 +52,25 @@ pub fn gpu_mb_update_impulse_joint_constraints( let lock_cfm_coeff = softness.joint_cfm_coeff; let max_corr_velocity = softness.max_corr_velocity; - let joints_start = batch_ids.mb_imp_joints_start(batch_id); - let cons_start = batch_ids.mb_imp_joint_constraints_start(batch_id); - let jac_buf_start = batch_ids.mb_imp_joint_jacobians_start(batch_id); - // Interleaved dynamics-buffer view (multibody_info / links_workspace / - // body_jacobians / mass_matrices / lu_pivots / dof_state). - let il = VSlice::interleaved(0, batch_ids.num_batches, batch_id); - let colliders_start = batch_ids.coll_start(batch_id); - - // Loop chunked across `num_threads` so a single workgroup row processes - // all joints in a batch (matches `gpu_init_joint_constraints` style). - // Iterating to `cap` instead of `len` lets us drop the per-batch - // `num_joints` storage binding — the host pads unused builder slots - // with `side_a_kind == SIDE_KIND_FIXED && side_b_kind == SIDE_KIND_FIXED` - // which we use as the inactive-slot sentinel below. - let mut i = invocation_id.x; - while i < cap { - let builder = builders.read(joints_start + i as usize); + let nb = batch_ids.num_batches; + let total = batch_ids.mb_imp_joints_batch_capacity * nb; + for t in StepRng::new(invocation_id.x..total, num_threads) { + let builder = builders.read(t as usize); let is_dummy = builder.side_a_kind == SIDE_KIND_FIXED && builder.side_b_kind == SIDE_KIND_FIXED; if !is_dummy { + let batch_id = t % nb; + let il = VSlice::interleaved(0, nb, batch_id); + let bix = batch_ids.body_ix(batch_id); builder.update_one_joint( constraints, - cons_start, jacobians, - jac_buf_start, multibody_info, links_workspace, body_jacobians, il, poses, - colliders_start, + bix, mprops, dt, lock_erp_inv_dt, @@ -94,7 +78,6 @@ pub fn gpu_mb_update_impulse_joint_constraints( max_corr_velocity, ); } - i += num_threads; } } @@ -119,25 +102,19 @@ pub fn gpu_mb_finalize_impulse_joint_constraints( #[spirv(uniform, descriptor_set = 0, binding = 3)] batch_ids: &BatchIndices, ) { let num_threads = num_workgroups.x * 64; - let batch_id = invocation_id.y; - let cap = batch_ids.mb_imp_joints_batch_capacity; - if invocation_id.x >= cap { - return; - } - - let joints_start = batch_ids.mb_imp_joints_start(batch_id); - let cons_start = batch_ids.mb_imp_joint_constraints_start(batch_id); - let il = VSlice::interleaved(0, batch_ids.num_batches, batch_id); + let nb = batch_ids.num_batches; + let total = batch_ids.mb_imp_joints_batch_capacity * nb; - let mut i = invocation_id.x; - while i < cap { - let builder = builders.read(joints_start + i as usize); + for t in StepRng::new(invocation_id.x..total, num_threads) { + let builder = builders.read(t as usize); let is_dummy = builder.side_a_kind == SIDE_KIND_FIXED && builder.side_b_kind == SIDE_KIND_FIXED; if !is_dummy { - let cons_base = cons_start + builder.constraint_id as usize; + let batch_id = t % nb; + let il = VSlice::interleaved(0, nb, batch_id); for s in 0..MAX_AXIS_CONSTRAINTS { - let mut c = constraints.read(cons_base + s as usize); + let cons_idx = il.atz(builder.constraint_id as usize + s as usize); + let mut c = constraints.read(cons_idx); if c.kind != 0 { // Multibody side(s): LU back-solve `M⁻¹·Jᵀ`. Free-body sides // already have their `W·J` (= M⁻¹·Jᵀ) filled by the build @@ -171,11 +148,10 @@ pub fn gpu_mb_finalize_impulse_joint_constraints( ); } c.finalize_generic_constraint(jacobians); - constraints.write(cons_base + s as usize, c); + constraints.write(cons_idx, c); } } } - i += num_threads; } } @@ -206,27 +182,26 @@ pub fn gpu_mb_solve_impulse_joint_constraints( // Per-lane scratch for the J·v tree reductions. #[spirv(workgroup)] partial: &mut [f32; LANES as usize], ) { - let batch_id = wg_id.y; + let nb = batch_ids.num_batches; + let batch_id = wg_id.x % nb; + let k = wg_id.x / nb; let lane = lid.x; - let joints_start = batch_ids.mb_imp_joints_start(batch_id); - let cons_start = batch_ids.mb_imp_joint_constraints_start(batch_id); - let il = VSlice::interleaved(0, batch_ids.num_batches, batch_id); - let colliders_start = batch_ids.coll_start(batch_id); + let il = VSlice::interleaved(0, nb, batch_id); + let bix = batch_ids.body_ix(batch_id); // `color_groups` is a per-batch prefix-sum over the color-sorted // builders: color `c` owns the sorted-builder range // `[color_groups[c-1], color_groups[c])` (start `0` for color `0`). let color = *curr_color as usize; - let color_groups = batch_ids.mb_imp_joint_color_groups_batch(batch_id, all_color_groups); let start = if color > 0 { - color_groups[color - 1] + all_color_groups.read(il.atz(color - 1)) } else { 0 }; - let end = color_groups[color]; + let end = all_color_groups.read(il.atz(color)); - let mut j = start + wg_id.x; + let mut j = start + k; let workgroup_is_active = j < end; if !workgroup_is_active { // Technically, if we enter here, we should return. However, on the web, a return would @@ -235,8 +210,7 @@ pub fn gpu_mb_solve_impulse_joint_constraints( j = start; // Any valid index will do. } - let builder = builders.at(joints_start + j as usize); - let cons_base = cons_start + builder.constraint_id as usize; + let builder = builders.at(il.atz(j as usize)); // Per-multibody dof base: same for every axis constraint of this joint. let dof_base_a = if builder.side_a_kind == SIDE_KIND_MB { @@ -255,7 +229,7 @@ pub fn gpu_mb_solve_impulse_joint_constraints( // TODO(PERF): load jacobians into shared memory and keep the velocity deltat on shared // memory and only writeback after all the axis constraints are solved. for s in 0..MAX_AXIS_CONSTRAINTS { - let c = constraints.at_mut(cons_base + s as usize); + let c = constraints.at_mut(il.atz(builder.constraint_id as usize + s as usize)); let active = workgroup_is_active && c.kind != 0; // dvel = J_b · v_b - J_a · v_a (rapier's `vel2 - vel1`). @@ -269,7 +243,7 @@ pub fn gpu_mb_solve_impulse_joint_constraints( dof_state, dof_base_a, solver_vels, - colliders_start, + bix, lane, partial, ); @@ -283,7 +257,7 @@ pub fn gpu_mb_solve_impulse_joint_constraints( dof_state, dof_base_b, solver_vels, - colliders_start, + bix, lane, partial, ); @@ -318,7 +292,7 @@ pub fn gpu_mb_solve_impulse_joint_constraints( dof_state, dof_base_a, solver_vels, - colliders_start, + bix, lane, ); side_apply_impulse_par( @@ -333,7 +307,7 @@ pub fn gpu_mb_solve_impulse_joint_constraints( dof_state, dof_base_b, solver_vels, - colliders_start, + bix, lane, ); @@ -351,24 +325,25 @@ pub fn gpu_mb_remove_impulse_joint_constraint_bias( #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] builders: &[MbImpulseJointBuilder], #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] constraints: &mut [MbImpulseJointConstraint], - #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] num_joints: &[u32], - #[spirv(uniform, descriptor_set = 0, binding = 3)] batch_ids: &BatchIndices, + #[spirv(uniform, descriptor_set = 0, binding = 2)] batch_ids: &BatchIndices, ) { let num_threads = num_workgroups.x * 64; - let batch_id = invocation_id.y; - let len = num_joints.read(batch_id as usize); - let joints_start = batch_ids.mb_imp_joints_start(batch_id); - let cons_start = batch_ids.mb_imp_joint_constraints_start(batch_id); - let mut i = invocation_id.x; - while i < len { - let builder = builders.at(joints_start + i as usize); - let cons_base = cons_start + builder.constraint_id as usize; + let nb = batch_ids.num_batches; + let total = batch_ids.mb_imp_joints_batch_capacity * nb; + for t in StepRng::new(invocation_id.x..total, num_threads) { + let builder = builders.at(t as usize); + let is_dummy = + builder.side_a_kind == SIDE_KIND_FIXED && builder.side_b_kind == SIDE_KIND_FIXED; + if is_dummy { + continue; + } + let batch_id = t % nb; + let il = VSlice::interleaved(0, nb, batch_id); for s in 0..MAX_AXIS_CONSTRAINTS { - let c = constraints.at_mut(cons_base + s as usize); + let c = constraints.at_mut(il.atz(builder.constraint_id as usize + s as usize)); if c.kind != 0 { c.rhs = c.rhs_wo_bias; } } - i += num_threads; } } diff --git a/src_rbd_shaders/dynamics/multibody/impulse_joint_constraints/update.rs b/src_rbd_shaders/dynamics/multibody/impulse_joint_constraints/update.rs index db284dbc..64089cc9 100644 --- a/src_rbd_shaders/dynamics/multibody/impulse_joint_constraints/update.rs +++ b/src_rbd_shaders/dynamics/multibody/impulse_joint_constraints/update.rs @@ -6,7 +6,7 @@ use khal_std::index::MaybeIndexUnchecked; use crate::dynamics::body::WorldMassProperties; use crate::dynamics::joint::{ANG_AXES_MASK, LIN_AXES_MASK, SPATIAL_DIM}; -use crate::utils::ISlice; +use crate::utils::{ISlice, BodyIx}; use crate::utils::linalg::{MatSlice, VSlice, lu_solve_in_place}; use crate::{DIM, Pose}; @@ -41,14 +41,15 @@ pub(super) fn solve_mb_wj( let v = jacobians.read(j_id as usize + k as usize); jacobians.write(wj_base + k as usize, v); } - let m = MatSlice::interleaved( - mb.mass_matrix_offset as usize, + let m = MatSlice::dense( + mb.mass_matrix_offset as usize * il.stride as usize + + il.shift as usize * (ndofs * ndofs) as usize, ndofs, ndofs, - il.stride, - il.shift, ); - let piv = VSlice::interleaved(mb.first_dof as usize, il.stride, il.shift); + let piv = VSlice::dense( + mb.first_dof as usize * il.stride as usize + il.shift as usize * ndofs as usize, + ); lu_solve_in_place( mass_matrices, m, @@ -73,9 +74,7 @@ impl MbImpulseJointBuilder { pub(super) fn update_one_joint( &self, constraints: &mut [MbImpulseJointConstraint], - cons_start: usize, jacobians: &mut [f32], - jac_buf_start: usize, multibody_info: &[MultibodyInfo], links_workspace: &[Vec4], body_jacobians: &[f32], @@ -83,23 +82,23 @@ impl MbImpulseJointBuilder { // batch_id`). il: VSlice, poses: &[Pose], - colliders_start: usize, + bix: BodyIx, mprops: &[WorldMassProperties], dt: f32, lock_erp_inv_dt: f32, lock_cfm_coeff: f32, max_corr_velocity: f32, ) { - let cons_base = cons_start + self.constraint_id as usize; + let cons_base = self.constraint_id as usize; // Mark all axis-constraint slots inactive up-front; the active branches // below overwrite the live ones (rapier rebuilds the entire // `out[start..len]` slab each `update` call, so unfilled slots are // guaranteed inactive). for s in 0..MAX_AXIS_CONSTRAINTS { - let mut cz = constraints.read(cons_base + s as usize); + let mut cz = constraints.read(il.atz(cons_base + s as usize)); cz.kind = 0; cz.impulse = 0.0; - constraints.write(cons_base + s as usize, cz); + constraints.write(il.atz(cons_base + s as usize), cz); } // Resolve per-side multibody descriptors (read by value to avoid @@ -124,7 +123,7 @@ impl MbImpulseJointBuilder { links_workspace, il, poses, - colliders_start, + bix, ); let pose_b = side_world_pose( self.side_b_kind, @@ -134,7 +133,7 @@ impl MbImpulseJointBuilder { links_workspace, il, poses, - colliders_start, + bix, ); let frame1 = pose_a * self.joint.local_frame_a; @@ -193,7 +192,8 @@ impl MbImpulseJointBuilder { mb: mb_b, }; let stride = axis_stride(ndofs_a, ndofs_b); - let j_base = jac_buf_start + self.jacobian_offset as usize; + let j_base = self.jacobian_offset as usize * il.stride as usize + + il.shift as usize * self.jacobian_capacity as usize; // `lock_erp_inv_dt` / `lock_cfm_coeff` are passed in from the configurable // joint softness (rapier's `joint.softness.{erp_inv_dt,cfm_coeff}(dt)`). @@ -213,7 +213,7 @@ impl MbImpulseJointBuilder { if len >= MAX_AXIS_CONSTRAINTS { break; } - let mut c = constraints.read(cons_base + len as usize); + let mut c = constraints.read(il.atz(cons_base + len as usize)); let j_id_a = j_off; let j_id_b = j_off + 2 * ndofs_a; motor_angular_generic( @@ -231,9 +231,9 @@ impl MbImpulseJointBuilder { body_jacobians, il, mprops, - colliders_start, + bix, ); - constraints.write(cons_base + len as usize, c); + constraints.write(il.atz(cons_base + len as usize), c); len += 1; j_off += stride; } @@ -245,7 +245,7 @@ impl MbImpulseJointBuilder { if len >= MAX_AXIS_CONSTRAINTS { break; } - let mut c = constraints.read(cons_base + len as usize); + let mut c = constraints.read(il.atz(cons_base + len as usize)); let j_id_a = j_off; let j_id_b = j_off + 2 * ndofs_a; motor_linear_generic( @@ -263,9 +263,9 @@ impl MbImpulseJointBuilder { body_jacobians, il, mprops, - colliders_start, + bix, ); - constraints.write(cons_base + len as usize, c); + constraints.write(il.atz(cons_base + len as usize), c); len += 1; j_off += stride; } @@ -277,7 +277,7 @@ impl MbImpulseJointBuilder { if len >= MAX_AXIS_CONSTRAINTS { break; } - let mut c = constraints.read(cons_base + len as usize); + let mut c = constraints.read(il.atz(cons_base + len as usize)); let j_id_a = j_off; let j_id_b = j_off + 2 * ndofs_a; lock_angular_generic( @@ -295,9 +295,9 @@ impl MbImpulseJointBuilder { body_jacobians, il, mprops, - colliders_start, + bix, ); - constraints.write(cons_base + len as usize, c); + constraints.write(il.atz(cons_base + len as usize), c); len += 1; j_off += stride; } @@ -309,7 +309,7 @@ impl MbImpulseJointBuilder { if len >= MAX_AXIS_CONSTRAINTS { break; } - let mut c = constraints.read(cons_base + len as usize); + let mut c = constraints.read(il.atz(cons_base + len as usize)); let j_id_a = j_off; let j_id_b = j_off + 2 * ndofs_a; lock_linear_generic( @@ -327,9 +327,9 @@ impl MbImpulseJointBuilder { body_jacobians, il, mprops, - colliders_start, + bix, ); - constraints.write(cons_base + len as usize, c); + constraints.write(il.atz(cons_base + len as usize), c); len += 1; j_off += stride; } @@ -341,7 +341,7 @@ impl MbImpulseJointBuilder { if len >= MAX_AXIS_CONSTRAINTS { break; } - let mut c = constraints.read(cons_base + len as usize); + let mut c = constraints.read(il.atz(cons_base + len as usize)); let j_id_a = j_off; let j_id_b = j_off + 2 * ndofs_a; let lim = self.joint.limits.at(i as usize); @@ -362,9 +362,9 @@ impl MbImpulseJointBuilder { body_jacobians, il, mprops, - colliders_start, + bix, ); - constraints.write(cons_base + len as usize, c); + constraints.write(il.atz(cons_base + len as usize), c); len += 1; j_off += stride; } @@ -376,7 +376,7 @@ impl MbImpulseJointBuilder { if len >= MAX_AXIS_CONSTRAINTS { break; } - let mut c = constraints.read(cons_base + len as usize); + let mut c = constraints.read(il.atz(cons_base + len as usize)); let j_id_a = j_off; let j_id_b = j_off + 2 * ndofs_a; let lim = self.joint.limits.at(i as usize); @@ -397,9 +397,9 @@ impl MbImpulseJointBuilder { body_jacobians, il, mprops, - colliders_start, + bix, ); - constraints.write(cons_base + len as usize, c); + constraints.write(il.atz(cons_base + len as usize), c); len += 1; j_off += stride; } @@ -446,13 +446,13 @@ pub(super) fn side_world_pose( links_workspace: &[Vec4], il: VSlice, poses: &[Pose], - colliders_start: usize, + bix: BodyIx, ) -> Pose { if side_kind == SIDE_KIND_FIXED { return Pose::IDENTITY; } if side_kind == SIDE_KIND_BODY { - return poses.read(colliders_start + side_id as usize); + return poses.read(bix.at(side_id)); } let wa = WsAddr::new(mb.first_link as usize, il.stride, il.shift); ws_pose(links_workspace, wa, side_link, WS_LTW) diff --git a/src_rbd_shaders/dynamics/multibody/integrate.rs b/src_rbd_shaders/dynamics/multibody/integrate.rs index 4d4e03f3..f75dd47c 100644 --- a/src_rbd_shaders/dynamics/multibody/integrate.rs +++ b/src_rbd_shaders/dynamics/multibody/integrate.rs @@ -6,6 +6,7 @@ use glamx::Vec4; use khal_std::glamx::UVec3; +use khal_std::index::MaybeIndexUnchecked; use khal_std::macros::{spirv, spirv_bindgen}; #[cfg(feature = "dim2")] @@ -50,13 +51,11 @@ pub fn gpu_mb_integrate_velocities( let mut dof_vel = batch_ids .ib_mut(batch_id, dof_state) .offset(mb.first_dof as usize); - let acc = batch_ids - .ib(batch_id, gen_accelerations) - .offset(mb.first_dof as usize); + let acc0 = batch_ids.mb_region(batch_id, mb.first_dof, mb.ndofs); for d in 0..mb.ndofs { let di = d as usize; - dof_vel[di] += acc[di] * dt; + dof_vel[di] += gen_accelerations.read(acc0 + di) * dt; } } diff --git a/src_rbd_shaders/dynamics/multibody/joint_constraints.rs b/src_rbd_shaders/dynamics/multibody/joint_constraints.rs index 136a5f21..fe361342 100644 --- a/src_rbd_shaders/dynamics/multibody/joint_constraints.rs +++ b/src_rbd_shaders/dynamics/multibody/joint_constraints.rs @@ -179,8 +179,10 @@ fn emit_joint_constraints( limit_min, limit_max, ); - joint_constraints.write(cons_base + slot as usize, cons); - slot += 1; + if slot < mb.max_constraints { + joint_constraints.write(cons_base + slot as usize, cons); + slot += 1; + } } if (limit_axes & (1 << axis)) != 0 { let cons = build_limit_constraint( @@ -195,8 +197,10 @@ fn emit_joint_constraints( joint_erp_inv_dt, joint_cfm_coeff, ); - joint_constraints.write(cons_base + slot as usize, cons); - slot += 1; + if slot < mb.max_constraints { + joint_constraints.write(cons_base + slot as usize, cons); + slot += 1; + } } curr_free_dof += 1; } @@ -222,8 +226,10 @@ fn emit_joint_constraints( joint_erp_inv_dt, joint_cfm_coeff, ); - joint_constraints.write(cons_base + slot as usize, cons); - slot += 1; + if slot < mb.max_constraints { + joint_constraints.write(cons_base + slot as usize, cons); + slot += 1; + } } if (motor_axes & (1 << axis)) != 0 { let has_limits = (limit_axes & (1 << axis)) != 0; @@ -248,8 +254,10 @@ fn emit_joint_constraints( limit_min, limit_max, ); - joint_constraints.write(cons_base + slot as usize, cons); - slot += 1; + if slot < mb.max_constraints { + joint_constraints.write(cons_base + slot as usize, cons); + slot += 1; + } } curr_free_dof += 1; } @@ -273,8 +281,10 @@ fn emit_joint_constraints( coupling.link_axis2 >> 16, ); let cons = build_coupling_constraint(&coupling, q1, q2, joint_erp_inv_dt); - joint_constraints.write(cons_base + slot as usize, cons); - slot += 1; + if slot < mb.max_constraints { + joint_constraints.write(cons_base + slot as usize, cons); + slot += 1; + } } // Joint dry friction (MJCF `frictionloss`): one box-bounded row per DoF @@ -293,7 +303,7 @@ fn emit_joint_constraints( 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 { + if fl > 0.0 && kin_mask_slice.read(d as usize) == 0.0 && slot < mb.max_constraints { let cons = build_friction_constraint(d, fl, dt, joint_cfm_coeff); joint_constraints.write(cons_base + slot as usize, cons); slot += 1; @@ -669,13 +679,16 @@ pub fn gpu_mb_finalize_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 piv = VSlice::dense(batch_ids.mb_region(batch_id, mb.first_dof, ndofs)); 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); + let m = MatSlice::dense( + batch_ids.mb_region(batch_id, mb.mass_matrix_offset, ndofs * ndofs), + ndofs, + ndofs, + ); if active { for s in StepRng::new(lane..mb.max_constraints, LANES) { diff --git a/src_rbd_shaders/dynamics/multibody/solve_constraints.rs b/src_rbd_shaders/dynamics/multibody/solve_constraints.rs index f12109a4..90d56b5a 100644 --- a/src_rbd_shaders/dynamics/multibody/solve_constraints.rs +++ b/src_rbd_shaders/dynamics/multibody/solve_constraints.rs @@ -88,16 +88,14 @@ pub fn gpu_mb_solve_constraints( #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] joint_constraint_columns: &[f32], #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] contact_constraints: &mut [MultibodyContactConstraint], - #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] contact_constraint_jacs: &[f32], - #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] contact_constraint_columns: &[f32], - #[spirv(uniform, descriptor_set = 0, binding = 6)] use_bias: &u32, - #[spirv(uniform, descriptor_set = 0, binding = 7)] batch_ids: &BatchIndices, - #[spirv(uniform, descriptor_set = 0, binding = 8)] max_contact_constraints: &u32, + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] contact_jac_cols: &[f32], + #[spirv(uniform, descriptor_set = 0, binding = 5)] use_bias: &u32, + #[spirv(uniform, descriptor_set = 0, binding = 6)] batch_ids: &BatchIndices, + #[spirv(uniform, descriptor_set = 0, binding = 7)] max_contact_constraints: &u32, #[spirv(storage_buffer, descriptor_set = 1, binding = 0)] dof_state: &mut [f32], #[spirv(storage_buffer, descriptor_set = 1, binding = 1)] solver_vels: &mut [Velocity], #[spirv(workgroup)] dof_v: &mut [f32; MAX_MB_DOFS], #[spirv(workgroup)] scratch: &mut [f32; LANES as usize], - #[spirv(workgroup)] imp_shared: &mut [f32; MAX_MB_CONTACT_CONSTRAINTS_PER_MB as usize], #[spirv(workgroup)] delta_shared: &mut f32, #[spirv(workgroup)] delta2_shared: &mut f32, ) { @@ -123,16 +121,13 @@ pub fn gpu_mb_solve_constraints( let v_base = mb.first_dof as usize; let dofs_stride = batch_ids.dof_batch_capacity as usize; - let colliders_start = batch_ids.coll_start(batch_id); let jcons_base = batch_ids.mb_joint_constraints_start(batch_id) + mb.first_constraint as usize; let jcol_base = batch_ids.mb_joint_constraint_columns_start(batch_id) + (mb.first_constraint as usize) * dofs_stride; - let ccons_base = batch_ids.mb_contact_constraints_start(batch_id) - + (mb_idx as usize) * (MAX_MB_CONTACT_CONSTRAINTS_PER_MB as usize); - let ccol_base = batch_ids.mb_contact_constraint_columns_start(batch_id) - + (mb_idx as usize) * (MAX_MB_CONTACT_CONSTRAINTS_PER_MB as usize) * dofs_stride; + let ccons_base = mb.contact_constraint_start as usize; + let cjc_base = ccons_base * 2 * dofs_stride; let contact_count = mb.contact_constraint_count; #[cfg(not(feature = "web-compat"))] @@ -142,21 +137,11 @@ pub fn gpu_mb_solve_constraints( } let active = in_range && ndofs != 0 && (mb.max_constraints != 0 || contact_count != 0); - // Load the generalized velocities and accumulated contact impulses into - // workgroup memory. - if active { - if lane < ndofs { - dof_v.write( - lane as usize, - dof_state.read(batch_ids.mbi(batch_id, v_base + lane as usize)), - ); - } - for s in StepRng::new(lane..contact_count, LANES) { - imp_shared.write( - s as usize, - contact_constraints.read(ccons_base + s as usize).impulse, - ); - } + if active && lane < ndofs { + dof_v.write( + lane as usize, + dof_state.read(batch_ids.mbi(batch_id, v_base + lane as usize)), + ); } workgroup_memory_barrier_with_group_sync(); @@ -258,8 +243,8 @@ pub fn gpu_mb_solve_constraints( #[cfg(feature = "dim2")] let has_pair = false; - let col_offset = ccol_base + (s as usize) * dofs_stride; - let col_offset2 = col_offset + dofs_stride; + let jac_offset = cjc_base + (s as usize) * 2 * dofs_stride; + let jac_offset2 = jac_offset + 2 * dofs_stride; let is_self = cons.free_body_id == u32::MAX; // Multibody side of J · u, one product per lane; lane 0 sums them in @@ -268,8 +253,7 @@ pub fn gpu_mb_solve_constraints( scratch.write( lane as usize, if lane < ndofs { - contact_constraint_jacs.read(col_offset + lane as usize) - * dof_v.read(lane as usize) + contact_jac_cols.read(jac_offset + lane as usize) * dof_v.read(lane as usize) } else { 0.0 }, @@ -289,8 +273,7 @@ pub fn gpu_mb_solve_constraints( scratch.write( lane as usize, if lane < ndofs { - contact_constraint_jacs.read(col_offset2 + lane as usize) - * dof_v.read(lane as usize) + contact_jac_cols.read(jac_offset2 + lane as usize) * dof_v.read(lane as usize) } else { 0.0 }, @@ -310,7 +293,7 @@ pub fn gpu_mb_solve_constraints( let free = if is_self { Velocity::default() } else { - solver_vels.read(colliders_start + cons.free_body_id as usize) + solver_vels.read(cons.free_body_id as usize) }; if !is_self { j_dot_v0 += cons.lin_jac.dot(free.linear) + gdot(cons.ang_jac, free.angular); @@ -320,15 +303,11 @@ pub fn gpu_mb_solve_constraints( } let cfm_factor = if use_bias { cons.cfm_factor } else { 1.0 }; - let impulse0 = imp_shared.read(s as usize); + let impulse0 = cons.impulse; let rhs0 = if use_bias { cons.rhs } else { cons.rhs_wo_bias }; let raw0 = cfm_factor * (impulse0 - cons.inv_lhs * (j_dot_v0 + rhs0)); - let impulse1 = if has_pair { - imp_shared.read((s + 1) as usize) - } else { - 0.0 - }; + let impulse1 = if has_pair { cons2.impulse } else { 0.0 }; let raw1 = if has_pair { let rhs1 = if use_bias { cons2.rhs @@ -343,8 +322,10 @@ pub fn gpu_mb_solve_constraints( // Normal: clamp to ≥ 0. Friction: cap the tangent pair to the // circular cone `μ · normal_impulse`. let (new0, new1) = if is_tangent { - let limit = - cons.friction_coeff * imp_shared.read(cons.normal_constraint_slot as usize); + let limit = cons.friction_coeff + * contact_constraints + .at(ccons_base + cons.normal_constraint_slot as usize) + .impulse; cap_friction(raw0, raw1, limit) } else if raw0 < 0.0 { (0.0, 0.0) @@ -354,9 +335,9 @@ pub fn gpu_mb_solve_constraints( let delta0 = new0 - impulse0; let delta1 = if has_pair { new1 - impulse1 } else { 0.0 }; - imp_shared.write(s as usize, new0); + contact_constraints.at_mut(cons_idx).impulse = new0; if has_pair { - imp_shared.write((s + 1) as usize, new1); + contact_constraints.at_mut(cons_idx + 1).impulse = new1; } *delta_shared = delta0; *delta2_shared = delta1; @@ -369,7 +350,7 @@ pub fn gpu_mb_solve_constraints( new_free.linear += cons2.lin_jac * (cons2.free_body_im * delta1); new_free.angular += cons2.ii_ang_jac * delta1; } - solver_vels.write(colliders_start + cons.free_body_id as usize, new_free); + solver_vels.write(cons.free_body_id as usize, new_free); } } workgroup_memory_barrier_with_group_sync(); @@ -379,30 +360,22 @@ pub fn gpu_mb_solve_constraints( let delta1 = *delta2_shared; if solve && lane < ndofs { if delta0 != 0.0 { - let col = contact_constraint_columns.read(col_offset + lane as usize); + let col = contact_jac_cols.read(jac_offset + dofs_stride + lane as usize); dof_v.write(lane as usize, dof_v.read(lane as usize) + delta0 * col); } if has_pair && delta1 != 0.0 { - let col = contact_constraint_columns.read(col_offset2 + lane as usize); + let col = contact_jac_cols.read(jac_offset2 + dofs_stride + lane as usize); dof_v.write(lane as usize, dof_v.read(lane as usize) + delta1 * col); } } workgroup_memory_barrier_with_group_sync(); } - // Writeback - if active { - if lane < ndofs { - dof_state.write( - batch_ids.mbi(batch_id, v_base + lane as usize), - dof_v.read(lane as usize), - ); - } - for s in StepRng::new(lane..contact_count, LANES) { - let mut cons = contact_constraints.read(ccons_base + s as usize); - cons.impulse = imp_shared.read(s as usize); - contact_constraints.write(ccons_base + s as usize, cons); - } + if active && lane < ndofs { + dof_state.write( + batch_ids.mbi(batch_id, v_base + lane as usize), + dof_v.read(lane as usize), + ); } } @@ -539,10 +512,9 @@ pub fn gpu_mb_build_contact_delassus( #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] multibody_info: &[MultibodyInfo], #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] contact_constraints: &[MultibodyContactConstraint], - #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] contact_constraint_jacs: &[f32], - #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] contact_constraint_columns: &[f32], - #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] delassus: &mut [f32], - #[spirv(uniform, descriptor_set = 0, binding = 5)] batch_ids: &BatchIndices, + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] contact_jac_cols: &[f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] delassus: &mut [f32], + #[spirv(uniform, descriptor_set = 0, binding = 4)] batch_ids: &BatchIndices, ) { const MAXC: u32 = MAX_MB_CONTACT_CONSTRAINTS_PER_MB; let batch_id = workgroup_id.y; @@ -555,16 +527,14 @@ pub fn gpu_mb_build_contact_delassus( let mb = multibody_info.read(batch_ids.mbi(batch_id, mb_idx as usize)); let ndofs = mb.ndofs; - let count = mb.contact_constraint_count; + let count = mb.contact_constraint_count.min(MAXC); if ndofs == 0 || count == 0 { return; } - let cons_base = - batch_ids.mb_contact_constraints_start(batch_id) + (mb_idx as usize) * (MAXC as usize); + let cons_base = mb.contact_constraint_start as usize; let dofs_stride = batch_ids.dof_batch_capacity as usize; - let col_base = batch_ids.mb_contact_constraint_columns_start(batch_id) - + (mb_idx as usize) * (MAXC as usize) * dofs_stride; + let jc_base = cons_base * 2 * dofs_stride; let d_base = ((batch_id * batch_ids.multibodies_batch_capacity + mb_idx) as usize) * (MAXC as usize) * (MAXC as usize); @@ -578,12 +548,12 @@ pub fn gpu_mb_build_contact_delassus( let j = p % count; // Multibody coupling: jac_j · (M⁻¹ jac_sᵀ). - let jac_j_off = col_base + (j as usize) * dofs_stride; - let col_s_off = col_base + (s as usize) * dofs_stride; + let jac_j_off = jc_base + (j as usize) * 2 * dofs_stride; + let col_s_off = jc_base + (s as usize) * 2 * dofs_stride + dofs_stride; let mut v = 0.0f32; for i in 0..ndofs { - let jj = contact_constraint_jacs.read(jac_j_off + i as usize); - let cs = contact_constraint_columns.read(col_s_off + i as usize); + let jj = contact_jac_cols.read(jac_j_off + i as usize); + let cs = contact_jac_cols.read(col_s_off + i as usize); v += jj * cs; } @@ -613,12 +583,11 @@ pub fn gpu_mb_solve_contacts_delassus( #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] multibody_info: &[MultibodyInfo], #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] contact_constraints: &mut [MultibodyContactConstraint], - #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] contact_constraint_jacs: &[f32], - #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] contact_constraint_columns: &[f32], - #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] delassus: &[f32], - #[spirv(uniform, descriptor_set = 0, binding = 5)] use_bias: &u32, - #[spirv(uniform, descriptor_set = 0, binding = 6)] batch_ids: &BatchIndices, - #[spirv(uniform, descriptor_set = 0, binding = 7)] max_contact_constraints: &u32, + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] contact_jac_cols: &[f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] delassus: &[f32], + #[spirv(uniform, descriptor_set = 0, binding = 4)] use_bias: &u32, + #[spirv(uniform, descriptor_set = 0, binding = 5)] batch_ids: &BatchIndices, + #[spirv(uniform, descriptor_set = 0, binding = 6)] max_contact_constraints: &u32, #[spirv(storage_buffer, descriptor_set = 1, binding = 0)] dof_state: &mut [f32], #[spirv(storage_buffer, descriptor_set = 1, binding = 1)] solver_vels: &mut [Velocity], #[spirv(workgroup)] dof_v: &mut [f32; MAX_MB_DOFS], @@ -644,7 +613,7 @@ pub fn gpu_mb_solve_contacts_delassus( let mb = multibody_info.read(batch_ids.mbi(batch_id, slot as usize)); let ndofs = mb.ndofs; - let count = mb.contact_constraint_count; + let count = mb.contact_constraint_count.min(MAXC); // Uniform per workgroup: every lane of this group returns together. #[cfg(not(feature = "web-compat"))] if ndofs == 0 || count == 0 { @@ -654,12 +623,9 @@ pub fn gpu_mb_solve_contacts_delassus( let use_bias = *use_bias != 0; let v_base = mb.first_dof as usize; - let colliders_start = batch_ids.coll_start(batch_id); - let cons_base = - batch_ids.mb_contact_constraints_start(batch_id) + (mb_idx as usize) * (MAXC as usize); + let cons_base = mb.contact_constraint_start as usize; let dofs_stride = batch_ids.dof_batch_capacity as usize; - let col_base = batch_ids.mb_contact_constraint_columns_start(batch_id) - + (mb_idx as usize) * (MAXC as usize) * dofs_stride; + let jc_base = cons_base * 2 * dofs_stride; let d_base = ((batch_id * batch_ids.multibodies_batch_capacity + mb_idx) as usize) * (MAXC as usize) * (MAXC as usize); @@ -703,14 +669,14 @@ pub fn gpu_mb_solve_contacts_delassus( // warmstart) velocities. if active { for s in StepRng::new(lane..count, LANES) { - let jac_off = col_base + (s as usize) * dofs_stride; + let jac_off = jc_base + (s as usize) * 2 * dofs_stride; let mut dot = 0.0f32; for i in 0..ndofs { - dot += contact_constraint_jacs.read(jac_off + i as usize) * dof_v.read(i as usize); + dot += contact_jac_cols.read(jac_off + i as usize) * dof_v.read(i as usize); } let cons = contact_constraints.read(cons_base + s as usize); if cons.free_body_id != u32::MAX { - let free = solver_vels.read(colliders_start + cons.free_body_id as usize); + let free = solver_vels.read(cons.free_body_id as usize); dot += cons.lin_jac.dot(free.linear) + gdot(cons.ang_jac, free.angular); } a_shared.write(s as usize, dot); @@ -799,7 +765,7 @@ pub fn gpu_mb_solve_contacts_delassus( if free_active { let cons = contact_constraints.read(cons_base + s as usize); let mut free = - solver_vels.read(colliders_start + cons.free_body_id as usize); + solver_vels.read(cons.free_body_id as usize); free.linear += cons.lin_jac * (cons.free_body_im * delta0); free.angular += cons.ii_ang_jac * delta0; if has_pair { @@ -807,7 +773,7 @@ pub fn gpu_mb_solve_contacts_delassus( free.linear += cons2.lin_jac * (cons2.free_body_im * delta1); free.angular += cons2.ii_ang_jac * delta1; } - solver_vels.write(colliders_start + cons.free_body_id as usize, free); + solver_vels.write(cons.free_body_id as usize, free); } } // Lane-parallel Delassus row update (row `s` is contiguous), plus @@ -822,12 +788,15 @@ pub fn gpu_mb_solve_contacts_delassus( a_shared.write(j as usize, a_shared.read(j as usize) + acc); } if lane < ndofs { - let col = contact_constraint_columns - .read(col_base + (s as usize) * dofs_stride + lane as usize); + let col = contact_jac_cols + .read(jc_base + (s as usize) * 2 * dofs_stride + dofs_stride + lane as usize); dof_v.write(lane as usize, dof_v.read(lane as usize) + delta0 * col); if has_pair { - let col2 = contact_constraint_columns - .read(col_base + ((s + 1) as usize) * dofs_stride + lane as usize); + let col2 = contact_jac_cols.read( + jc_base + ((s + 1) as usize) * 2 * dofs_stride + + dofs_stride + + lane as usize, + ); dof_v.write(lane as usize, dof_v.read(lane as usize) + delta1 * col2); } } diff --git a/src_rbd_shaders/dynamics/multibody/types.rs b/src_rbd_shaders/dynamics/multibody/types.rs index 4dc0ff1e..3070d5fe 100644 --- a/src_rbd_shaders/dynamics/multibody/types.rs +++ b/src_rbd_shaders/dynamics/multibody/types.rs @@ -35,6 +35,7 @@ pub const CONTACT_CONSTRAINTS_PER_POINT: u32 = 3; /// Total constraint slots reserved per multibody (= contact points × DIM). pub const MAX_MB_CONTACT_CONSTRAINTS_PER_MB: u32 = MAX_MB_CONTACTS_PER_MB * CONTACT_CONSTRAINTS_PER_POINT; +pub const MB_CONS_SLOT_RESERVE: u32 = 8 * CONTACT_CONSTRAINTS_PER_POINT; /// `kind` value: inactive / unused slot. pub const MB_CONTACT_KIND_INACTIVE: u32 = 0; @@ -446,9 +447,11 @@ pub struct MultibodyInfo { /// multibody. Written by `gpu_mb_init_contact_constraints`, read by the /// warmstart / finalize / solve / remove-bias contact kernels. pub contact_constraint_count: u32, - /// Per-step copy of `contacts_len[batch]` (to work around the web 8 storage - /// bindings count limit). + pub contact_constraint_start: u32, + pub old_contact_constraint_start: u32, + pub old_contact_constraint_count: u32, pub batch_contacts_len: u32, + pub batch_contacts_start: u32, /// First entry of this multibody's DoF couplings in the `dof_couplings` /// buffer (relative to the batch's coupling slice). pub first_coupling: u32, diff --git a/src_rbd_shaders/dynamics/multibody/ws_soa.rs b/src_rbd_shaders/dynamics/multibody/ws_soa.rs index 9fab1029..f3c666d2 100644 --- a/src_rbd_shaders/dynamics/multibody/ws_soa.rs +++ b/src_rbd_shaders/dynamics/multibody/ws_soa.rs @@ -114,8 +114,9 @@ impl WsAddr { /// link `k` (relative to `base`). #[inline] pub fn at(&self, k: u32, quad: u32) -> usize { - ((self.base + k as usize) * WS_QUADS as usize + quad as usize) * self.stride as usize - + self.shift as usize + ((self.base + k as usize) * self.stride as usize + self.shift as usize) + * WS_QUADS as usize + + quad as usize } } @@ -453,7 +454,7 @@ pub fn ws_soa_from_structs( /* * 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. + * per link rather than the record-interleaved quads the kernels index. */ #[cfg(not(target_arch_is_gpu))] pub fn ws_soa_to_structs( diff --git a/src_rbd_shaders/dynamics/solver.rs b/src_rbd_shaders/dynamics/solver.rs index 3c596dfa..2d3d6e90 100644 --- a/src_rbd_shaders/dynamics/solver.rs +++ b/src_rbd_shaders/dynamics/solver.rs @@ -37,6 +37,7 @@ pub fn gpu_solver_init_constraints( constraints: &mut [TwoBodyConstraint], #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] constraint_builders: &mut [TwoBodyConstraintBuilder], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] contact_offsets: &[u32], #[spirv(storage_buffer, descriptor_set = 1, binding = 0)] collider_world_poses: &[Pose], #[spirv(storage_buffer, descriptor_set = 1, binding = 1)] solver_body_poses: &[Pose], #[spirv(storage_buffer, descriptor_set = 1, binding = 2)] vels: &[Velocity], @@ -45,20 +46,18 @@ pub fn gpu_solver_init_constraints( #[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 contacts = batch_ids.contact_batch(batch_id, contacts); - let mut constraints = batch_ids.contact_batch_mut(batch_id, constraints); - let mut constraint_builders = batch_ids.contact_batch_mut(batch_id, constraint_builders); - let collider_world_poses = batch_ids.coll_batch(batch_id, collider_world_poses); - let solver_body_poses = batch_ids.coll_batch(batch_id, solver_body_poses); - let vels = batch_ids.coll_batch(batch_id, vels); - let mprops = batch_ids.coll_batch(batch_id, mprops); - let cap = batch_ids.contacts_batch_capacity.min(num_threads); + let total = contact_offsets.read(batch_ids.num_batches as usize); + let collider_world_poses = Slice(collider_world_poses, 0); + let solver_body_poses = Slice(solver_body_poses, 0); + let vels = Slice(vels, 0); + let mprops = Slice(mprops, 0); - for i in StepRng::new(invocation_id.x..cap, num_threads) { - let im = &contacts[i as usize]; + for i in StepRng::new(invocation_id.x..total, num_threads) { + let i = i as usize; + let im = contacts.at(i); if im.contact.len == 0 { + constraints.at_mut(i).len = 0; continue; } im.contact_to_constraint( @@ -67,8 +66,8 @@ pub fn gpu_solver_init_constraints( &solver_body_poses, &vels, params, - &mut constraints[i as usize], - &mut constraint_builders[i as usize], + constraints.at_mut(i), + constraint_builders.at_mut(i), ); } } @@ -84,20 +83,22 @@ pub fn gpu_solver_count_constraints( #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] body_constraint_counts: &mut [u32], #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] body_group: &[u32], #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] mprops: &[WorldMassProperties], - #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] contacts_len: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] contact_offsets: &[u32], #[spirv(uniform, descriptor_set = 0, binding = 5)] batch_ids: &BatchIndices, ) { let num_threads = num_workgroups.x * WORKGROUP_SIZE; - let batch_id = invocation_id.y; - let contacts = batch_ids.contact_batch(batch_id, contacts); - let mut body_constraint_counts = batch_ids.coll_batch_mut(batch_id, body_constraint_counts); - let body_group = batch_ids.coll_batch(batch_id, body_group); - let mprops = batch_ids.coll_batch(batch_id, mprops); - let len = contacts_len.read(batch_id as usize); + let total = contact_offsets.read(batch_ids.num_batches as usize); + let contacts = Slice(contacts, 0); + let mut body_constraint_counts = SliceMut(body_constraint_counts, 0); + let body_group = Slice(body_group, 0); + let mprops = Slice(mprops, 0); - for i in StepRng::new(invocation_id.x..len, num_threads) { + for i in StepRng::new(invocation_id.x..total, num_threads) { let im = &contacts[i as usize]; + if im.contact.len == 0 { + continue; + } let body1 = im.bodies.x; let body2 = im.bodies.y; let group1 = body_group[body1 as usize]; @@ -129,20 +130,22 @@ pub fn gpu_solver_update_constraints( constraints: &mut [TwoBodyConstraint], #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] constraint_builders: &[TwoBodyConstraintBuilder], - #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] contacts_len: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] contact_offsets: &[u32], #[spirv(storage_buffer, descriptor_set = 1, binding = 0)] solver_body_poses: &[Pose], #[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 mut constraints = batch_ids.contact_batch_mut(batch_id, constraints); - let constraint_builders = batch_ids.contact_batch(batch_id, constraint_builders); - let solver_body_poses = batch_ids.coll_batch(batch_id, solver_body_poses); - let len = contacts_len.read(batch_id as usize); + let total = contact_offsets.read(batch_ids.num_batches as usize); + let mut constraints = SliceMut(constraints, 0); + let constraint_builders = Slice(constraint_builders, 0); + let solver_body_poses = Slice(solver_body_poses, 0); - for i in StepRng::new(invocation_id.x..len, num_threads) { + for i in StepRng::new(invocation_id.x..total, num_threads) { + if constraints[i as usize].len == 0 { + continue; + } constraints[i as usize].update_constraint( &constraint_builders[i as usize], &solver_body_poses, @@ -162,20 +165,22 @@ pub fn gpu_solver_refresh_rhs_wo_bias( constraints: &mut [TwoBodyConstraint], #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] constraint_builders: &[TwoBodyConstraintBuilder], - #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] contacts_len: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] contact_offsets: &[u32], #[spirv(storage_buffer, descriptor_set = 1, binding = 0)] solver_body_poses: &[Pose], #[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 mut constraints = batch_ids.contact_batch_mut(batch_id, constraints); - let constraint_builders = batch_ids.contact_batch(batch_id, constraint_builders); - let solver_body_poses = batch_ids.coll_batch(batch_id, solver_body_poses); - let len = contacts_len.read(batch_id as usize); + let total = contact_offsets.read(batch_ids.num_batches as usize); + let mut constraints = SliceMut(constraints, 0); + let constraint_builders = Slice(constraint_builders, 0); + let solver_body_poses = Slice(solver_body_poses, 0); - for i in StepRng::new(invocation_id.x..len, num_threads) { + for i in StepRng::new(invocation_id.x..total, num_threads) { + if constraints[i as usize].len == 0 { + continue; + } constraints[i as usize].refresh_rhs_wo_bias( &constraint_builders[i as usize], &solver_body_poses, @@ -192,23 +197,24 @@ pub fn gpu_solver_sort_constraints( #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] body_constraint_counts: &mut [u32], #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] mprops: &[WorldMassProperties], #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] contacts: &[IndexedManifold], - #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] contacts_len: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] contact_offsets: &[u32], #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] body_constraint_ids: &mut [u32], #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] body_group: &[u32], #[spirv(uniform, descriptor_set = 0, binding = 6)] batch_ids: &BatchIndices, ) { let num_threads = num_workgroups.x * WORKGROUP_SIZE; - let batch_id = invocation_id.y; - let bci_start = batch_id as usize * 2 * batch_ids.contacts_batch_capacity as usize; - let contacts = batch_ids.contact_batch(batch_id, contacts); - let mut body_constraint_counts = batch_ids.coll_batch_mut(batch_id, body_constraint_counts); - let body_group = batch_ids.coll_batch(batch_id, body_group); - let mprops = batch_ids.coll_batch(batch_id, mprops); - let mut body_constraint_ids = SliceMut(body_constraint_ids, bci_start); - let len = contacts_len.read(batch_id as usize); + let total = contact_offsets.read(batch_ids.num_batches as usize); + let contacts = Slice(contacts, 0); + let mut body_constraint_counts = SliceMut(body_constraint_counts, 0); + let body_group = Slice(body_group, 0); + let mprops = Slice(mprops, 0); + let mut body_constraint_ids = SliceMut(body_constraint_ids, 0); - for i in StepRng::new(invocation_id.x..len, num_threads) { + for i in StepRng::new(invocation_id.x..total, num_threads) { + if contacts[i as usize].contact.len == 0 { + continue; + } let body1 = contacts[i as usize].bodies.x as usize; let body2 = contacts[i as usize].bodies.y as usize; let group1 = body_group[body1] as usize; @@ -241,25 +247,19 @@ pub fn gpu_solver_cleanup( #[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 num_bodies = batch_ids.colliders_len; + let num_slots = batch_ids.colliders_batch_capacity * batch_ids.num_batches; - let mut body_constraint_counts = batch_ids.coll_batch_mut(batch_id, body_constraint_counts); - let mut solver_vels = batch_ids.coll_batch_mut(batch_id, solver_vels); - let vels = batch_ids.coll_batch(batch_id, vels); - let mprops = batch_ids.coll_batch(batch_id, mprops); - - for i in StepRng::new(invocation_id.x..num_bodies, num_threads) { + for i in StepRng::new(invocation_id.x..num_slots, num_threads) { let idx = i as usize; - body_constraint_counts[idx] = 0; + body_constraint_counts.write(idx, 0); // HACK: to handle static bodies. - if mprops[idx].inv_mass != Vector::ZERO { - solver_vels[idx].linear = vels[idx].linear; - solver_vels[idx].angular = vels[idx].angular; + if mprops.at(idx).inv_mass != Vector::ZERO { + solver_vels.at_mut(idx).linear = vels.at(idx).linear; + solver_vels.at_mut(idx).angular = vels.at(idx).angular; } else { - solver_vels[idx].linear = Vector::ZERO; - solver_vels[idx].angular = AngVector::default(); + solver_vels.at_mut(idx).linear = Vector::ZERO; + solver_vels.at_mut(idx).angular = AngVector::default(); } } } @@ -275,26 +275,23 @@ pub fn gpu_init_solver_vels_inc( #[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 i = invocation_id.x; - let num_bodies = batch_ids.bodies_len; - let mut solver_vels_inc = batch_ids.coll_batch_mut(batch_id, solver_vels_inc); - let mprops = batch_ids.coll_batch(batch_id, mprops); + let num_bodies = batch_ids.bodies_len * batch_ids.num_batches; if i < num_bodies { let idx = i as usize; - solver_vels_inc[idx].linear = Vector::ZERO; - solver_vels_inc[idx].angular = AngVector::default(); + solver_vels_inc.at_mut(idx).linear = Vector::ZERO; + solver_vels_inc.at_mut(idx).angular = AngVector::default(); // TODO: this isn't a very pretty way of detecting static bodies. - if mprops[idx].inv_mass != Vector::ZERO { + if mprops.at(idx).inv_mass != Vector::ZERO { // TODO: this currently only handles gravity (no user forces yet). #[cfg(feature = "dim3")] let g = Vector::new(gravity.x, gravity.y, gravity.z); #[cfg(feature = "dim2")] let g = Vector::new(gravity.x, gravity.y); - solver_vels_inc[idx].linear = g * params.dt; + solver_vels_inc.at_mut(idx).linear = g * params.dt; } } } @@ -308,17 +305,14 @@ pub fn gpu_apply_solver_vels_inc( #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] solver_vels_inc: &[Velocity], #[spirv(uniform, descriptor_set = 0, binding = 2)] batch_ids: &BatchIndices, ) { - let batch_id = invocation_id.y; let i = invocation_id.x; - let num_bodies = batch_ids.bodies_len; - let mut solver_vels = batch_ids.coll_batch_mut(batch_id, solver_vels); - let solver_vels_inc = batch_ids.coll_batch(batch_id, solver_vels_inc); + let num_bodies = batch_ids.bodies_len * batch_ids.num_batches; if i < num_bodies { let idx = i as usize; - solver_vels[idx].linear += solver_vels_inc[idx].linear; - solver_vels[idx].angular += solver_vels_inc[idx].angular; + solver_vels.at_mut(idx).linear += solver_vels_inc.at(idx).linear; + solver_vels.at_mut(idx).angular += solver_vels_inc.at(idx).angular; } } @@ -335,14 +329,12 @@ pub fn gpu_warmstart_without_colors( #[spirv(uniform, descriptor_set = 0, binding = 4)] batch_ids: &BatchIndices, ) { let num_threads = num_workgroups.x * WORKGROUP_SIZE; - let batch_id = invocation_id.y; - let bci_start = batch_id as usize * 2 * batch_ids.contacts_batch_capacity as usize; - let num_bodies = batch_ids.bodies_len; + let num_bodies = batch_ids.bodies_len * batch_ids.num_batches; - let body_constraint_counts = batch_ids.coll_batch(batch_id, body_constraint_counts); - let body_constraint_ids = Slice(body_constraint_ids, bci_start); - let constraints = batch_ids.contact_batch(batch_id, constraints); - let mut solver_vels = batch_ids.coll_batch_mut(batch_id, solver_vels); + let body_constraint_counts = Slice(body_constraint_counts, 0); + let body_constraint_ids = Slice(body_constraint_ids, 0); + let constraints = Slice(constraints, 0); + let mut solver_vels = SliceMut(solver_vels, 0); for body_id in StepRng::new(invocation_id.x..num_bodies, num_threads) { let mut solver_vel = solver_vels[body_id as usize]; @@ -371,17 +363,15 @@ pub fn gpu_warmstart( #[spirv(uniform, descriptor_set = 0, binding = 5)] batch_ids: &BatchIndices, ) { let num_threads = num_workgroups.x * WORKGROUP_SIZE; - let batch_id = invocation_id.y; - let stride = batch_ids.solver_color_buckets_stride; + let nb = batch_ids.num_batches; - let constraints = batch_ids.contact_batch(batch_id, constraints); - let color_sorted_ids = batch_ids.contact_batch(batch_id, color_sorted_ids); - let mut solver_vels = batch_ids.coll_batch_mut(batch_id, solver_vels); + let constraints = Slice(constraints, 0); + let color_sorted_ids = Slice(color_sorted_ids, 0); + let mut solver_vels = SliceMut(solver_vels, 0); let color = *curr_color; - let bucket = (batch_id * stride + color) as usize; - let start = color_starts.read(bucket); - let end = color_starts.read(bucket + 1); + let start = color_starts.read((color * nb - 1) as usize); + let end = color_starts.read(((color + 1) * nb - 1) as usize); for k in StepRng::new(start + invocation_id.x..end, num_threads) { let i = color_sorted_ids[k as usize]; @@ -415,18 +405,16 @@ pub fn gpu_step_gauss_seidel( #[spirv(uniform, descriptor_set = 0, binding = 6)] use_bias: &u32, ) { let num_threads = num_workgroups.x * WORKGROUP_SIZE; - let batch_id = invocation_id.y; - let stride = batch_ids.solver_color_buckets_stride; + let nb = batch_ids.num_batches; - let mut constraints = batch_ids.contact_batch_mut(batch_id, constraints); - let color_sorted_ids = batch_ids.contact_batch(batch_id, color_sorted_ids); - let mut solver_vels = batch_ids.coll_batch_mut(batch_id, solver_vels); + let mut constraints = SliceMut(constraints, 0); + let color_sorted_ids = Slice(color_sorted_ids, 0); + let mut solver_vels = SliceMut(solver_vels, 0); let color = *curr_color; let use_bias = *use_bias != 0; - let bucket = (batch_id * stride + color) as usize; - let start = color_starts.read(bucket); - let end = color_starts.read(bucket + 1); + let start = color_starts.read((color * nb - 1) as usize); + let end = color_starts.read(((color + 1) * nb - 1) as usize); for k in StepRng::new(start + invocation_id.x..end, num_threads) { let i = color_sorted_ids[k as usize]; @@ -464,32 +452,24 @@ pub fn gpu_warmstart_fused( ) { let lane = invocation_id.x; let batch_id = invocation_id.y; - let stride = batch_ids.solver_color_buckets_stride; + let nb = batch_ids.num_batches; - let constraints = batch_ids.contact_batch(batch_id, constraints); - let color_sorted_ids = batch_ids.contact_batch(batch_id, color_sorted_ids); - let mut solver_vels = batch_ids.coll_batch_mut(batch_id, solver_vels); + let constraints = Slice(constraints, 0); + let color_sorted_ids = Slice(color_sorted_ids, 0); + let mut solver_vels = SliceMut(solver_vels, 0); let num_colors = *num_colors; - let base = (batch_id * stride) as usize; - let any_work = color_starts.read(base + 1) != color_starts.read(base + num_colors as usize + 1); - #[cfg(not(feature = "web-compat"))] - if !any_work { - // Every color bucket is empty. - return; - } - for color in 1..=num_colors { - let bucket = base + color as usize; - let start = color_starts.read(bucket); - let end = color_starts.read(bucket + 1); + let bucket = (color * nb + batch_id) as usize; + let start = color_starts.read(bucket - 1); + let end = color_starts.read(bucket); #[cfg(not(feature = "web-compat"))] if start == end { // Empty color. continue; } - if any_work && start != end { + if start != end { for k in StepRng::new(start + lane..end, WORKGROUP_SIZE) { let i = color_sorted_ids[k as usize]; let constraint = &constraints[i as usize]; @@ -536,32 +516,24 @@ pub fn gpu_step_gauss_seidel_fused( ) { let lane = invocation_id.x; let batch_id = invocation_id.y; - let stride = batch_ids.solver_color_buckets_stride; + let nb = batch_ids.num_batches; - let mut constraints = batch_ids.contact_batch_mut(batch_id, constraints); - let color_sorted_ids = batch_ids.contact_batch(batch_id, color_sorted_ids); - let mut solver_vels = batch_ids.coll_batch_mut(batch_id, solver_vels); + let mut constraints = SliceMut(constraints, 0); + let color_sorted_ids = Slice(color_sorted_ids, 0); + let mut solver_vels = SliceMut(solver_vels, 0); let num_colors = *num_colors; let use_bias = *use_bias != 0; - // Early-out / empty-color skip: see `gpu_warmstart_fused`. - let base = (batch_id * stride) as usize; - let any_work = color_starts.read(base + 1) != color_starts.read(base + num_colors as usize + 1); - #[cfg(not(feature = "web-compat"))] - if !any_work { - return; - } - for color in 1..=num_colors { - let bucket = base + color as usize; - let start = color_starts.read(bucket); - let end = color_starts.read(bucket + 1); + let bucket = (color * nb + batch_id) as usize; + let start = color_starts.read(bucket - 1); + let end = color_starts.read(bucket); #[cfg(not(feature = "web-compat"))] if start == end { continue; } - if any_work && start != end { + if start != end { for k in StepRng::new(start + lane..end, WORKGROUP_SIZE) { let i = color_sorted_ids[k as usize]; let solver_id1 = constraints[i as usize].solver_body_a as usize; @@ -602,16 +574,13 @@ pub fn gpu_integrate_linearized( #[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 i = invocation_id.x; - let num_bodies = batch_ids.bodies_len; - let mut poses = batch_ids.coll_batch_mut(batch_id, poses); - let mut solver_vels = batch_ids.coll_batch_mut(batch_id, solver_vels); + let num_bodies = batch_ids.bodies_len * batch_ids.num_batches; if i < num_bodies { let idx = i as usize; - let mut vels = solver_vels[idx]; + let mut vels = solver_vels.read(idx); let max_lin = params.max_linear_velocity(); let lin_norm = vels.linear.length(); @@ -640,8 +609,8 @@ pub fn gpu_integrate_linearized( } } - solver_vels[idx] = vels; - let pose = &mut poses[idx]; + solver_vels.write(idx, vels); + let pose = poses.at_mut(idx); vels.integrate_linearized(params.dt, &mut pose.translation, &mut pose.rotation); } } @@ -660,17 +629,16 @@ pub fn gpu_init_solver_bodies( #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] solver_body_poses: &mut [Pose], #[spirv(uniform, descriptor_set = 0, binding = 3)] batch_ids: &BatchIndices, ) { - let batch_id = invocation_id.y; let i = invocation_id.x; - let num_bodies = batch_ids.bodies_len; - let body_poses = batch_ids.coll_batch(batch_id, body_poses); - let local_mprops = batch_ids.coll_batch(batch_id, local_mprops); - let mut solver_body_poses = batch_ids.coll_batch_mut(batch_id, solver_body_poses); + let num_bodies = batch_ids.bodies_len * batch_ids.num_batches; if i < num_bodies { let idx = i as usize; - solver_body_poses[idx] = body_poses[idx].prepend_translation(local_mprops[idx].com); + solver_body_poses.write( + idx, + body_poses.read(idx).prepend_translation(local_mprops.at(idx).com), + ); } } @@ -691,20 +659,19 @@ pub fn gpu_solver_finalize( local_mprops: &[LocalMassProperties], #[spirv(uniform, descriptor_set = 0, binding = 5)] batch_ids: &BatchIndices, ) { - let batch_id = invocation_id.y; let i = invocation_id.x; - let num_bodies = batch_ids.bodies_len; - let mut vels = batch_ids.coll_batch_mut(batch_id, vels); - let solver_vels = batch_ids.coll_batch(batch_id, solver_vels); - let mut body_poses = batch_ids.coll_batch_mut(batch_id, body_poses); - let solver_body_poses = batch_ids.coll_batch(batch_id, solver_body_poses); - let local_mprops = batch_ids.coll_batch(batch_id, local_mprops); + let num_bodies = batch_ids.bodies_len * batch_ids.num_batches; if i < num_bodies { let idx = i as usize; - vels[idx].linear = solver_vels[idx].linear; - vels[idx].angular = solver_vels[idx].angular; - body_poses[idx] = solver_body_poses[idx].prepend_translation(-local_mprops[idx].com); + vels.at_mut(idx).linear = solver_vels.at(idx).linear; + vels.at_mut(idx).angular = solver_vels.at(idx).angular; + body_poses.write( + idx, + solver_body_poses + .read(idx) + .prepend_translation(-local_mprops.at(idx).com), + ); } } diff --git a/src_rbd_shaders/dynamics/warmstart.rs b/src_rbd_shaders/dynamics/warmstart.rs index b5be104a..4edf8034 100644 --- a/src_rbd_shaders/dynamics/warmstart.rs +++ b/src_rbd_shaders/dynamics/warmstart.rs @@ -25,25 +25,23 @@ pub fn gpu_transfer_warmstart_impulses( new_constraints: &mut [TwoBodyConstraint], #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] new_constraint_builders: &[TwoBodyConstraintBuilder], - #[spirv(storage_buffer, descriptor_set = 0, binding = 6)] contacts_len: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 6)] contact_offsets: &[u32], #[spirv(uniform, descriptor_set = 0, binding = 7)] batch_ids: &BatchIndices, ) { - let batch_id = invocation_id.y; - let contacts_start = batch_ids.contacts_start(batch_id); - let colliders_start = batch_ids.coll_start(batch_id); - let bci_start = batch_id as usize * 2 * batch_ids.contacts_batch_capacity as usize; + let total = contact_offsets.read(batch_ids.num_batches as usize); + let old_body_constraint_counts = Slice(old_body_constraint_counts, 0); + let old_body_constraint_ids = Slice(old_body_constraint_ids, 0); + let old_constraints = Slice(old_constraints, 0); + let old_constraint_builders = Slice(old_constraint_builders, 0); + let mut new_constraints = SliceMut(new_constraints, 0); + let new_constraint_builders = Slice(new_constraint_builders, 0); - let old_body_constraint_counts = Slice(old_body_constraint_counts, colliders_start); - let old_body_constraint_ids = Slice(old_body_constraint_ids, bci_start); - let old_constraints = Slice(old_constraints, contacts_start); - let old_constraint_builders = Slice(old_constraint_builders, contacts_start); - let mut new_constraints = SliceMut(new_constraints, contacts_start); - let new_constraint_builders = Slice(new_constraint_builders, contacts_start); - - let len = contacts_len.read(batch_id as usize); let cid_new = invocation_id.x; - if cid_new < len { + if cid_new < total { + if new_constraints[cid_new as usize].len == 0 { + return; + } transfer_warmstart_impulses( cid_new, &old_body_constraint_counts, @@ -77,26 +75,24 @@ pub fn gpu_seed_colors_from_warmstart( #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] old_constraints_colors: &[u32], #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] constraints_colors: &mut [u32], #[spirv(storage_buffer, descriptor_set = 0, binding = 6)] colored: &mut [u32], - #[spirv(storage_buffer, descriptor_set = 0, binding = 7)] contacts_len: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 7)] contact_offsets: &[u32], #[spirv(uniform, descriptor_set = 0, binding = 8)] batch_ids: &BatchIndices, ) { - let batch_id = invocation_id.y; - let contacts_start = batch_ids.contacts_start(batch_id); - let colliders_start = batch_ids.coll_start(batch_id); - let bci_start = batch_id as usize * 2 * batch_ids.contacts_batch_capacity as usize; + let total = contact_offsets.read(batch_ids.num_batches as usize); + let old_body_constraint_counts = Slice(old_body_constraint_counts, 0); + let old_body_constraint_ids = Slice(old_body_constraint_ids, 0); + let old_constraints = Slice(old_constraints, 0); + let old_constraints_colors = Slice(old_constraints_colors, 0); + let mut constraints_colors = SliceMut(constraints_colors, 0); + let mut colored = SliceMut(colored, 0); + let new_constraints = Slice(new_constraints, 0); - let old_body_constraint_counts = Slice(old_body_constraint_counts, colliders_start); - let old_body_constraint_ids = Slice(old_body_constraint_ids, bci_start); - let old_constraints = Slice(old_constraints, contacts_start); - let old_constraints_colors = Slice(old_constraints_colors, contacts_start); - let mut constraints_colors = SliceMut(constraints_colors, contacts_start); - let mut colored = SliceMut(colored, contacts_start); - let new_constraints = Slice(new_constraints, contacts_start); - - let len = contacts_len.read(batch_id as usize); let i = invocation_id.x as usize; - if (i as u32) < len { + if (i as u32) < total { + if new_constraints[i].len == 0 { + return; + } let body_a = new_constraints[i].solver_body_a; let body_b = new_constraints[i].solver_body_b; diff --git a/src_rbd_shaders/utils/indices.rs b/src_rbd_shaders/utils/indices.rs index 8c1627c4..8ee15ec2 100644 --- a/src_rbd_shaders/utils/indices.rs +++ b/src_rbd_shaders/utils/indices.rs @@ -1,5 +1,4 @@ -use crate::utils::linalg::{MatSlice, VSlice}; -use crate::utils::{ISlice, ISliceMut, Slice, SliceMut}; +use crate::utils::{ISlice, ISliceMut}; /// Per-batch capacities and packed-buffer section offsets, shared by every /// kernel that needs to slice a flat tensor into its batch's slot. @@ -21,11 +20,8 @@ pub struct BatchIndices { pub colliders_len: u32, /// Number of *active* rigid bodies per batch. pub bodies_len: u32, - pub collision_pairs_batch_capacity: u32, - pub contacts_batch_capacity: u32, - /// Free-body impulse joints — buffer stride (capacity) per batch. - pub impulse_joints_batch_capacity: u32, - /// Number of *active* free-body impulse joints per batch (the loop bound). + pub collision_pairs_capacity: u32, + pub contacts_capacity: u32, pub impulse_joints_len: u32, /* @@ -36,10 +32,7 @@ pub struct BatchIndices { /// per-multibody kernels). pub multibodies_len: u32, pub links_batch_capacity: u32, - pub jacobians_batch_capacity: u32, - pub mass_matrix_batch_capacity: u32, pub coriolis_batch_capacity: u32, - pub i_coriolis_dt_batch_capacity: u32, pub dof_batch_capacity: u32, /* @@ -47,16 +40,8 @@ pub struct BatchIndices { */ pub mb_joint_constraints_batch_capacity: u32, pub mb_joint_constraint_columns_batch_capacity: u32, - pub mb_contact_constraints_batch_capacity: u32, - pub mb_contact_constraint_columns_batch_capacity: u32, + pub mb_contact_constraints_capacity: u32, pub mb_imp_joints_batch_capacity: u32, - pub mb_imp_joint_constraints_batch_capacity: u32, - pub mb_imp_joint_jacobians_batch_capacity: u32, - /// Multibody-touching impulse-joint color-group slab (per-batch stride - /// = number of colors). The free-body impulse-joint color groups, by - /// contrast, are stored single-batch (identical coloring across batches) - /// and read at offset 0. - pub mb_imp_joint_color_groups_batch_capacity: u32, /// Actual max `ndofs` across every multibody in every batch (often smaller /// than the fixed `MAX_MB_DOFS` limit). pub mb_max_ndofs: u32, @@ -77,9 +62,6 @@ pub struct BatchIndices { * These are buffers that were combined into a single storage * buffer to comply with the 10 storage buffers limit on the web. */ - pub coriolis_w_section_offset: u32, - pub i_coriolis_dt_section_offset: u32, - pub dof_damping_section_offset: u32, /// Offset (in f32 entries, within a batch's `mass_matrices` view) of the /// section holding the coriolis-aware "acceleration" mass matrix /// (rapier's `acc_augmented_mass`). @@ -100,8 +82,15 @@ impl BatchIndices { * `MatSlice::dense(base, ...)`). */ #[inline] - pub fn coll_start(&self, batch_id: u32) -> usize { - batch_id as usize * self.colliders_batch_capacity as usize + pub fn body_global(&self, batch_id: u32, local: u32) -> usize { + local as usize * self.num_batches as usize + batch_id as usize + } + #[inline] + pub fn body_ix(&self, batch_id: u32) -> BodyIx { + BodyIx { + stride: self.num_batches, + shift: batch_id, + } } /// Interleaved flat index for the multibody dynamics buffers. @@ -135,29 +124,14 @@ impl BatchIndices { /// Interleaved dense matrix view at intra-batch element offset `offset`. #[inline] - pub fn imat(&self, batch_id: u32, offset: usize, rows: u32, cols: u32) -> MatSlice { - MatSlice::interleaved(offset, rows, cols, self.num_batches, batch_id) + pub fn mb_region(&self, batch_id: u32, offset: u32, len: u32) -> usize { + offset as usize * self.num_batches as usize + batch_id as usize * len as usize } /// Interleaved vector view at intra-batch element offset `offset`. #[inline] - pub fn ivec(&self, batch_id: u32, offset: usize) -> VSlice { - VSlice::interleaved(offset, self.num_batches, batch_id) - } - - #[inline] - pub fn collision_pairs_start(&self, batch_id: u32) -> usize { - batch_id as usize * self.collision_pairs_batch_capacity as usize - } - - #[inline] - pub fn contacts_start(&self, batch_id: u32) -> usize { - batch_id as usize * self.contacts_batch_capacity as usize - } - - #[inline] - pub fn impulse_joints_start(&self, batch_id: u32) -> usize { - batch_id as usize * self.impulse_joints_batch_capacity as usize + pub fn collider_batch(&self, collider_id: u32) -> u32 { + collider_id % self.num_batches } #[inline] @@ -170,163 +144,21 @@ impl BatchIndices { batch_id as usize * self.mb_joint_constraint_columns_batch_capacity as usize } - #[inline] - pub fn mb_contact_constraints_start(&self, batch_id: u32) -> usize { - batch_id as usize * self.mb_contact_constraints_batch_capacity as usize - } - - #[inline] - pub fn mb_contact_constraint_columns_start(&self, batch_id: u32) -> usize { - batch_id as usize * self.mb_contact_constraint_columns_batch_capacity as usize - } - #[inline] pub fn mb_dof_couplings_start(&self, batch_id: u32) -> usize { batch_id as usize * self.mb_dof_couplings_batch_capacity as usize } +} - #[inline] - pub fn mb_imp_joints_start(&self, batch_id: u32) -> usize { - batch_id as usize * self.mb_imp_joints_batch_capacity as usize - } - - #[inline] - pub fn mb_imp_joint_constraints_start(&self, batch_id: u32) -> usize { - batch_id as usize * self.mb_imp_joint_constraints_batch_capacity as usize - } - - #[inline] - pub fn mb_imp_joint_jacobians_start(&self, batch_id: u32) -> usize { - batch_id as usize * self.mb_imp_joint_jacobians_batch_capacity as usize - } - - #[inline] - pub fn mb_imp_joint_color_groups_start(&self, batch_id: u32) -> usize { - batch_id as usize * self.mb_imp_joint_color_groups_batch_capacity as usize - } - - /* - * Typed batch slices — for buffers consumed via `Slice` / `SliceMut` - * wrappers rather than as raw f32 arrays. - */ - #[inline] - pub fn coll_batch<'s, T>(&self, batch_id: u32, slice: &'s [T]) -> Slice<'s, T> { - Slice(slice, self.coll_start(batch_id)) - } - - #[inline] - pub fn coll_batch_mut<'s, T>(&self, batch_id: u32, slice: &'s mut [T]) -> SliceMut<'s, T> { - SliceMut(slice, self.coll_start(batch_id)) - } - - #[inline] - pub fn collision_pairs_batch<'s, T>(&self, batch_id: u32, slice: &'s [T]) -> Slice<'s, T> { - Slice(slice, self.collision_pairs_start(batch_id)) - } - - #[inline] - pub fn collision_pairs_batch_mut<'s, T>( - &self, - batch_id: u32, - slice: &'s mut [T], - ) -> SliceMut<'s, T> { - SliceMut(slice, self.collision_pairs_start(batch_id)) - } - - #[inline] - pub fn contact_batch<'s, T>(&self, batch_id: u32, slice: &'s [T]) -> Slice<'s, T> { - Slice(slice, self.contacts_start(batch_id)) - } - - #[inline] - pub fn contact_batch_mut<'s, T>(&self, batch_id: u32, slice: &'s mut [T]) -> SliceMut<'s, T> { - SliceMut(slice, self.contacts_start(batch_id)) - } - - #[inline] - pub fn impulse_joints_batch<'s, T>(&self, batch_id: u32, slice: &'s [T]) -> Slice<'s, T> { - Slice(slice, self.impulse_joints_start(batch_id)) - } - - #[inline] - pub fn impulse_joints_batch_mut<'s, T>( - &self, - batch_id: u32, - slice: &'s mut [T], - ) -> SliceMut<'s, T> { - SliceMut(slice, self.impulse_joints_start(batch_id)) - } - - #[inline] - pub fn mb_joint_constraints_batch<'s, T>(&self, batch_id: u32, slice: &'s [T]) -> Slice<'s, T> { - Slice(slice, self.mb_joint_constraints_start(batch_id)) - } - - #[inline] - pub fn mb_joint_constraints_batch_mut<'s, T>( - &self, - batch_id: u32, - slice: &'s mut [T], - ) -> SliceMut<'s, T> { - SliceMut(slice, self.mb_joint_constraints_start(batch_id)) - } - - #[inline] - pub fn mb_contact_constraints_batch<'s, T>( - &self, - batch_id: u32, - slice: &'s [T], - ) -> Slice<'s, T> { - Slice(slice, self.mb_contact_constraints_start(batch_id)) - } - - #[inline] - pub fn mb_contact_constraints_batch_mut<'s, T>( - &self, - batch_id: u32, - slice: &'s mut [T], - ) -> SliceMut<'s, T> { - SliceMut(slice, self.mb_contact_constraints_start(batch_id)) - } - - #[inline] - pub fn mb_imp_joints_batch<'s, T>(&self, batch_id: u32, slice: &'s [T]) -> Slice<'s, T> { - Slice(slice, self.mb_imp_joints_start(batch_id)) - } - - #[inline] - pub fn mb_imp_joints_batch_mut<'s, T>( - &self, - batch_id: u32, - slice: &'s mut [T], - ) -> SliceMut<'s, T> { - SliceMut(slice, self.mb_imp_joints_start(batch_id)) - } - - #[inline] - pub fn mb_imp_joint_constraints_batch<'s, T>( - &self, - batch_id: u32, - slice: &'s [T], - ) -> Slice<'s, T> { - Slice(slice, self.mb_imp_joint_constraints_start(batch_id)) - } - - #[inline] - pub fn mb_imp_joint_constraints_batch_mut<'s, T>( - &self, - batch_id: u32, - slice: &'s mut [T], - ) -> SliceMut<'s, T> { - SliceMut(slice, self.mb_imp_joint_constraints_start(batch_id)) - } +#[derive(Copy, Clone)] +pub struct BodyIx { + pub stride: u32, + pub shift: u32, +} - #[inline] - pub fn mb_imp_joint_color_groups_batch<'s, T>( - &self, - batch_id: u32, - slice: &'s [T], - ) -> Slice<'s, T> { - Slice(slice, self.mb_imp_joint_color_groups_start(batch_id)) +impl BodyIx { + #[inline(always)] + pub fn at(self, id: u32) -> usize { + id as usize * self.stride as usize + self.shift as usize } } diff --git a/src_rbd_shaders/utils/mod.rs b/src_rbd_shaders/utils/mod.rs index a9916b52..cce3e5df 100644 --- a/src_rbd_shaders/utils/mod.rs +++ b/src_rbd_shaders/utils/mod.rs @@ -8,7 +8,7 @@ pub mod radix_sort; mod slice; pub use basis::orthonormal_basis3; -pub use indices::BatchIndices; +pub use indices::{BatchIndices, BodyIx}; pub use slice::{ISlice, ISliceMut, Slice, SliceMut}; /// Division with ceiling (signed). diff --git a/src_viewer/ui.rs b/src_viewer/ui.rs index 3f8886d1..e79e58a4 100644 --- a/src_viewer/ui.rs +++ b/src_viewer/ui.rs @@ -266,6 +266,8 @@ fn performance_ui( row("Impulse joints:", counts.impulse_joints); row("Multibodies:", counts.multibodies); row("Multibody DOFs:", counts.multibody_dofs); + row("MB contact slots:", counts.mb_contact_constraints); + row("MB contact capacity:", counts.mb_contact_constraints_capacity); } if counts.particles > 0 { From 7c5a6f2a89d725b8bf6210767c1e3c726cf475f5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Sat, 29 Aug 2026 21:07:41 +0200 Subject: [PATCH 2/7] test(rbd): add a batched-stacks parity probe --- src_rbd/pipeline/mod.rs | 2 + src_rbd/pipeline/test_batched_stacks.rs | 572 ++++++++++++++++++++++++ 2 files changed, 574 insertions(+) create mode 100644 src_rbd/pipeline/test_batched_stacks.rs diff --git a/src_rbd/pipeline/mod.rs b/src_rbd/pipeline/mod.rs index 2b80368d..1347ff60 100644 --- a/src_rbd/pipeline/mod.rs +++ b/src_rbd/pipeline/mod.rs @@ -6,6 +6,8 @@ #[cfg(all(test, feature = "dim3"))] mod bench_narrow_phase; +#[cfg(all(test, feature = "dim3"))] +mod test_batched_stacks; mod insertion_removal; mod lbvh_validation; mod rbd_state; diff --git a/src_rbd/pipeline/test_batched_stacks.rs b/src_rbd/pipeline/test_batched_stacks.rs new file mode 100644 index 00000000..f71d1c23 --- /dev/null +++ b/src_rbd/pipeline/test_batched_stacks.rs @@ -0,0 +1,572 @@ +use crate::math::Pose; +use crate::pipeline::{RbdCapacities, RbdPipeline, RbdResizePolicy, RbdState}; +use crate::rapier::prelude::*; +use crate::shaders::dynamics::RbdSimParams; +use khal::backend::{Backend, GpuBackend}; + +async fn test_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()) + } +} +fn build_env(num_stacks: usize, stack_height: 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(50.0, 0.5, 50.0), ground, &mut bodies); + + for s in 0..num_stacks { + let x = (s % 8) as f32 * 2.0; + let z = (s / 8) as f32 * 2.0; + for i in 0..stack_height { + let y = 0.25 + i as f32 * 0.45; + let handle = + bodies.insert(RigidBodyBuilder::dynamic().translation(Vec3::new(x, y, z))); + colliders.insert_with_parent(ColliderBuilder::cuboid(0.2, 0.2, 0.2), handle, &mut bodies); + } + } + + (bodies, colliders) +} +async fn run_case(num_envs: u32, num_stacks: usize, stack_height: usize, collisions_capacity: u32) { + let backend = test_backend().await; + + let envs: Vec<_> = (0..num_envs) + .map(|_| build_env(num_stacks, stack_height)) + .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, + ..Default::default() + }; + let mut state = RbdState::from_rapier(&backend, &refs, capacities); + let pipeline = RbdPipeline::new(&backend).unwrap(); + + for _ in 0..250 { + pipeline.step(&backend, &mut state, None).unwrap(); + pipeline.auto_resize_buffers(&backend, &mut state).unwrap(); + } + backend.synchronize().unwrap(); + + let poses: Vec = backend + .slow_read_vec(state.body_poses().buffer()) + .await + .unwrap(); + let nb = num_envs as usize; + let boxes_per_env = num_stacks * stack_height; + for env in 0..num_envs as usize { + for b in 0..boxes_per_env { + let pose = poses[(1 + b) * nb + env]; + assert!( + pose.translation.is_finite(), + "env {env} box {b}: non-finite pose {:?}", + pose.translation + ); + let level = b % stack_height; + let expected_y = 0.2 + level as f32 * 0.4; + let y = pose.translation.y; + assert!( + (y - expected_y).abs() < 0.1, + "env {env} box {b}: y = {y}, expected ~{expected_y}" + ); + if env > 0 { + let ref_pose = poses[(1 + b) * nb]; + let d = (pose.translation - ref_pose.translation).length(); + assert!(d < 5.0e-2, "env {env} box {b}: diverged from env 0 by {d}"); + } + } + } + println!( + "OK: envs={num_envs} stacks={num_stacks} height={stack_height} cap={collisions_capacity}" + ); +} +#[futures_test::test] +#[serial_test::serial] +#[ignore] +async fn test_stacks_1_tiny() { + run_case(1, 1, 4, 64).await; +} +#[futures_test::test] +#[serial_test::serial] +#[ignore] +async fn test_stacks_2_batched() { + run_case(64, 1, 4, 64).await; +} +#[futures_test::test] +#[serial_test::serial] +#[ignore] +async fn test_stacks_3_lbvh() { + run_case(1, 32, 4, 1024).await; +} +#[futures_test::test] +#[serial_test::serial] +#[ignore] +async fn test_stacks_4_lbvh_batched() { + run_case(4, 32, 4, 1024).await; +} +#[futures_test::test] +#[serial_test::serial] +#[ignore] +async fn test_stacks_5_overflow_resize() { + run_case(64, 1, 4, 8).await; +} +#[futures_test::test] +#[serial_test::serial] +#[ignore] +async fn test_stacks_8_impulse_joint() { + let backend = test_backend().await; + + let num_envs = 4u32; + let build = || { + let mut bodies = RigidBodySet::new(); + let mut colliders = ColliderSet::new(); + let mut joints = ImpulseJointSet::new(); + let anchor = bodies.insert(RigidBodyBuilder::fixed().translation(Vec3::new(0.0, 2.0, 0.0))); + colliders.insert_with_parent(ColliderBuilder::ball(0.1), anchor, &mut bodies); + let bob = + bodies.insert(RigidBodyBuilder::dynamic().translation(Vec3::new(1.0, 2.0, 0.0))); + colliders.insert_with_parent(ColliderBuilder::ball(0.1), bob, &mut bodies); + let joint = RevoluteJointBuilder::new(Vec3::Z) + .local_anchor1(Vec3::ZERO) + .local_anchor2(Vec3::new(-1.0, 0.0, 0.0)); + joints.insert(anchor, bob, joint, true); + (bodies, colliders, joints) + }; + let envs: Vec<_> = (0..num_envs).map(|_| build()).collect(); + let mb_joints = MultibodyJointSet::new(); + let params = RbdSimParams::tgs_soft(); + let refs: Vec<_> = envs + .iter() + .map(|(b, c, j)| (b, c, j, &mb_joints, ¶ms)) + .collect(); + + let capacities = RbdCapacities { + batches: num_envs, + collisions_capacity: 16, + ..Default::default() + }; + let mut state = RbdState::from_rapier(&backend, &refs, capacities); + let pipeline = RbdPipeline::new(&backend).unwrap(); + for _ in 0..400 { + pipeline.step(&backend, &mut state, None).unwrap(); + } + backend.synchronize().unwrap(); + + let poses: Vec = backend + .slow_read_vec(state.body_poses().buffer()) + .await + .unwrap(); + let nb = num_envs as usize; + for env in 0..nb { + let p = poses[nb + env].translation; + assert!(p.is_finite(), "env {env}: non-finite bob pose {p:?}"); + let d = (p - Vec3::new(0.0, 2.0, 0.0)).length(); + assert!( + (0.9..1.1).contains(&d), + "env {env}: bob at distance {d} from anchor, expected ~1" + ); + } + println!("OK: impulse-joint pendulum holds, envs={num_envs}"); +} +#[futures_test::test] +#[serial_test::serial] +#[ignore] +async fn test_stacks_9_env_reset() { + let backend = test_backend().await; + + let num_envs = 4u32; + let envs: Vec<_> = (0..num_envs).map(|_| build_env(1, 4)).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: 32, + ..Default::default() + }; + let mut state = RbdState::from_rapier(&backend, &refs, capacities); + let pipeline = RbdPipeline::new(&backend).unwrap(); + + for _ in 0..50 { + pipeline.step(&backend, &mut state, None).unwrap(); + } + backend.synchronize().unwrap(); + let snap = state.snapshot(&backend).await; + state.publish_reset_templates(&backend, &[&snap]); + + for _ in 0..50 { + pipeline.step(&backend, &mut state, None).unwrap(); + } + state.reset_envs_from_templates(&backend, &[(1, 0)], &[crate::math::Vector::ZERO], &[]); + backend.synchronize().unwrap(); + + let poses: Vec = backend + .slow_read_vec(state.body_poses().buffer()) + .await + .unwrap(); + let nb = num_envs as usize; + for b in 0..5 { + let restored = poses[b * nb + 1].translation; + let template = snap_pose(&snap, b); + let d = (restored - template).length(); + assert!( + d < 1.0e-6, + "body {b}: env 1 not restored to the template (off by {d})" + ); + } + println!("OK: env reset restores the template, envs={num_envs}"); +} +#[cfg(feature = "dim3")] +fn snap_pose(snap: &crate::pipeline::RbdSnapshot, b: usize) -> Vec3 { + snap.debug_body_pose(b).translation +} +#[futures_test::test] +#[serial_test::serial] +#[ignore] +async fn test_stacks_7_pfm_shapes() { + let backend = test_backend().await; + + let num_envs = 4u32; + let build = || { + 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(20.0, 0.5, 20.0), ground, &mut bodies); + for i in 0..3 { + let b = bodies.insert( + RigidBodyBuilder::dynamic().translation(Vec3::new(i as f32 * 1.5, 0.6, 0.0)), + ); + colliders.insert_with_parent(ColliderBuilder::capsule_y(0.15, 0.1), b, &mut bodies); + } + for i in 0..2 { + let b = bodies.insert( + RigidBodyBuilder::dynamic().translation(Vec3::new(i as f32 * 1.5 + 0.5, 0.6, 1.5)), + ); + colliders.insert_with_parent(ColliderBuilder::ball(0.2), b, &mut bodies); + } + (bodies, colliders) + }; + let envs: Vec<_> = (0..num_envs).map(|_| build()).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: 32, + ..Default::default() + }; + let mut state = RbdState::from_rapier(&backend, &refs, capacities); + let pipeline = RbdPipeline::new(&backend).unwrap(); + for _ in 0..250 { + pipeline.step(&backend, &mut state, None).unwrap(); + pipeline.auto_resize_buffers(&backend, &mut state).unwrap(); + } + backend.synchronize().unwrap(); + + let poses: Vec = backend + .slow_read_vec(state.body_poses().buffer()) + .await + .unwrap(); + let nb = num_envs as usize; + for env in 0..num_envs as usize { + for b in 0..5 { + let pose = poses[(1 + b) * nb + env]; + assert!(pose.translation.is_finite()); + let y = pose.translation.y; + assert!( + (0.05..0.4).contains(&y), + "env {env} body {b}: y = {y}, expected resting on the ground" + ); + } + } + println!("OK: pfm shapes rest, envs={num_envs}"); +} +fn build_mb_env(ball_colliders: bool) -> (RigidBodySet, ColliderSet, MultibodyJointSet) { + let mut bodies = RigidBodySet::new(); + let mut colliders = ColliderSet::new(); + let mut mb_joints = MultibodyJointSet::new(); + + let ground = bodies.insert(RigidBodyBuilder::fixed().translation(Vec3::new(0.0, -0.5, 0.0))); + colliders.insert_with_parent(ColliderBuilder::cuboid(20.0, 0.5, 20.0), ground, &mut bodies); + let mut prev = None; + for i in 0..3 { + let x = i as f32 * 0.5; + let link = bodies.insert(RigidBodyBuilder::dynamic().translation(Vec3::new(x, 1.0, 0.0))); + let shape = if ball_colliders { + ColliderBuilder::ball(0.1) + } else { + ColliderBuilder::cuboid(0.2, 0.1, 0.1) + }; + colliders.insert_with_parent(shape, link, &mut bodies); + if let Some(prev) = prev { + let joint = RevoluteJointBuilder::new(Vec3::Z) + .local_anchor1(Vec3::new(0.25, 0.0, 0.0)) + .local_anchor2(Vec3::new(-0.25, 0.0, 0.0)); + mb_joints.insert(prev, link, joint, true); + } + prev = Some(link); + } + + (bodies, colliders, mb_joints) +} +#[futures_test::test] +#[serial_test::serial] +#[ignore] +async fn test_stacks_6_multibody() { + let backend = test_backend().await; + + let num_envs = 4u32; + let envs: Vec<_> = (0..num_envs).map(|_| build_mb_env(false)).collect(); + let joints = ImpulseJointSet::new(); + let params = RbdSimParams::tgs_soft(); + let refs: Vec<_> = envs + .iter() + .map(|(b, c, mb)| (b, c, &joints, mb, ¶ms)) + .collect(); + + let mut capacities = RbdCapacities { + batches: num_envs, + collisions_capacity: 32, + ..Default::default() + }; + if std::env::var("NEXUS_TEST_FIXED").is_ok() { + capacities.collisions_resize_policy = RbdResizePolicy::Fixed; + capacities.solver_colors_resize_policy = RbdResizePolicy::Fixed; + } + let mut state = RbdState::from_rapier(&backend, &refs, capacities); + let pipeline = RbdPipeline::new(&backend).unwrap(); + + let debug = std::env::var("NEXUS_TEST_DEBUG").is_ok(); + for step in 0..300 { + pipeline.step(&backend, &mut state, None).unwrap(); + pipeline.auto_resize_buffers(&backend, &mut state).unwrap(); + if debug && (20..40).contains(&step) { + backend.synchronize().unwrap(); + let (layout, demand) = state.multibodies().debug_cons_layout(&backend).await; + let poses: Vec = backend + .slow_read_vec(state.body_poses().buffer()) + .await + .unwrap(); + let finite = poses.iter().all(|p| p.translation.is_finite()); + println!("step {step}: cap={} demand={demand} layout={layout:?} finite={finite}", state.multibodies().contact_constraints_capacity()); + if demand > 0 { + let cons: Vec = backend + .slow_read_vec(state.multibodies().contact_constraints().buffer()) + .await + .unwrap_or_default(); + for env in [0usize, 2] { + let (st, ct, _, _) = layout[env]; + for k in 0..ct.min(2) { + let c = &cons[(st + k) as usize]; + println!( + " env{env} slot{k}: kind={} imp={} inv_lhs={} rhs={}", + c.kind, c.impulse, c.inv_lhs, c.rhs + ); + } + } + } + if !finite { + for (i, p) in poses.iter().enumerate() { + if !p.translation.is_finite() { + println!(" non-finite body slot {i}"); + } + } + break; + } + } + } + backend.synchronize().unwrap(); + + let poses: Vec = backend + .slow_read_vec(state.body_poses().buffer()) + .await + .unwrap(); + let nb = num_envs as usize; + for env in 0..num_envs as usize { + for l in 0..3 { + let pose = poses[(1 + l) * nb + env]; + assert!( + pose.translation.is_finite(), + "env {env} link {l}: non-finite pose {:?}", + pose.translation + ); + let y = pose.translation.y; + assert!( + (0.0..0.5).contains(&y), + "env {env} link {l}: y = {y}, expected resting near 0.1" + ); + } + } + println!("OK: multibody chain rests, envs={num_envs}"); +} +#[futures_test::test] +#[serial_test::serial] +#[ignore] +async fn test_stacks_11_mb_point_contacts() { + let backend = test_backend().await; + + let num_envs = 4u32; + let envs: Vec<_> = (0..num_envs).map(|_| build_mb_env(true)).collect(); + let joints = ImpulseJointSet::new(); + let params = RbdSimParams::tgs_soft(); + let refs: Vec<_> = envs + .iter() + .map(|(b, c, mb)| (b, c, &joints, mb, ¶ms)) + .collect(); + + let capacities = RbdCapacities { + batches: num_envs, + collisions_capacity: 32, + ..Default::default() + }; + let mut state = RbdState::from_rapier(&backend, &refs, capacities); + let pipeline = RbdPipeline::new(&backend).unwrap(); + for _ in 0..300 { + pipeline.step(&backend, &mut state, None).unwrap(); + pipeline.auto_resize_buffers(&backend, &mut state).unwrap(); + } + backend.synchronize().unwrap(); + + let poses: Vec = backend + .slow_read_vec(state.body_poses().buffer()) + .await + .unwrap(); + let nb = num_envs as usize; + for env in 0..nb { + for l in 0..3 { + let pose = poses[(1 + l) * nb + env]; + assert!( + pose.translation.is_finite(), + "env {env} link {l}: non-finite pose {:?}", + pose.translation + ); + let y = pose.translation.y; + assert!( + (0.05..0.3).contains(&y), + "env {env} link {l}: y = {y}, expected resting near 0.1" + ); + } + } + println!("OK: 1-point-manifold multibody contacts rest, envs={num_envs}"); +} +#[futures_test::test] +#[serial_test::serial] +#[ignore] +async fn test_stacks_10_mb_impulse_joint() { + let backend = test_backend().await; + + let num_envs = 4u32; + let anchor_y = |e: u32| 2.0 + 0.25 * e as f32; + let build = |e: u32| { + let mut bodies = RigidBodySet::new(); + let mut colliders = ColliderSet::new(); + let mut joints = ImpulseJointSet::new(); + let mut mb_joints = MultibodyJointSet::new(); + let y = anchor_y(e); + let anchor = bodies.insert(RigidBodyBuilder::fixed().translation(Vec3::new(0.0, y, 0.0))); + colliders.insert_with_parent(ColliderBuilder::ball(0.05), anchor, &mut bodies); + let mut links = Vec::new(); + let mut prev = anchor; + for i in 0..2 { + let x = 0.5 + i as f32 * 0.5; + let link = + bodies.insert(RigidBodyBuilder::dynamic().translation(Vec3::new(x, y, 0.0))); + colliders.insert_with_parent(ColliderBuilder::ball(0.05), link, &mut bodies); + let joint = RevoluteJointBuilder::new(Vec3::Z) + .local_anchor1(Vec3::new(if i == 0 { 0.0 } else { 0.25 }, 0.0, 0.0)) + .local_anchor2(Vec3::new(if i == 0 { -0.5 } else { -0.25 }, 0.0, 0.0)); + mb_joints.insert(prev, link, joint, true); + links.push(link); + prev = link; + } + for (i, &link) in links.iter().enumerate() { + let x = 0.5 + i as f32 * 0.5; + let bob = + bodies.insert(RigidBodyBuilder::dynamic().translation(Vec3::new(x, y - 0.4, 0.0))); + colliders.insert_with_parent(ColliderBuilder::ball(0.05), bob, &mut bodies); + let joint = RevoluteJointBuilder::new(Vec3::Z) + .local_anchor1(Vec3::ZERO) + .local_anchor2(Vec3::new(0.0, 0.4, 0.0)); + joints.insert(link, bob, joint, true); + } + + (bodies, colliders, joints, mb_joints) + }; + let envs: Vec<_> = (0..num_envs).map(build).collect(); + let params = RbdSimParams::tgs_soft(); + let refs: Vec<_> = envs + .iter() + .map(|(b, c, j, mb)| (b, c, j, mb, ¶ms)) + .collect(); + + let capacities = RbdCapacities { + batches: num_envs, + collisions_capacity: 16, + ..Default::default() + }; + let mut state = RbdState::from_rapier(&backend, &refs, capacities); + assert!( + state.multibodies().mb_imp_joint_num_colors() >= 2, + "expected >= 2 impulse-joint colors (both joints touch the same multibody)" + ); + let pipeline = RbdPipeline::new(&backend).unwrap(); + for _ in 0..400 { + pipeline.step(&backend, &mut state, None).unwrap(); + } + backend.synchronize().unwrap(); + + let poses: Vec = backend + .slow_read_vec(state.body_poses().buffer()) + .await + .unwrap(); + let nb = num_envs as usize; + for env in 0..nb { + let at = |slot: usize| poses[slot * nb + env]; + let ay = at(0).translation.y; + assert!( + (ay - anchor_y(env as u32)).abs() < 1.0e-4, + "env {env}: anchor y = {ay}, expected {}", + anchor_y(env as u32) + ); + for slot in 0..5 { + let p = at(slot).translation; + assert!(p.is_finite(), "env {env} slot {slot}: non-finite pose {p:?}"); + } + for (link_slot, bob_slot) in [(1usize, 3usize), (2, 4)] { + let link = at(link_slot); + let bob = at(bob_slot); + let bob_anchor = bob.translation + bob.rotation * Vec3::new(0.0, 0.4, 0.0); + let err = (bob_anchor - link.translation).length(); + assert!( + err < 0.05, + "env {env}: bob {bob_slot} anchor drifts {err} from link {link_slot}" + ); + } + } + println!("OK: multibody impulse joints hold, envs={num_envs}"); +} From f8ac510f41deac1e9b4458e790d0adbe18da5093 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Sat, 29 Aug 2026 21:07:41 +0200 Subject: [PATCH 3/7] refactor(rbd): remove the never-maintained dof_values mirror --- src_rbd/dynamics/multibody/env_reset.rs | 28 +++---------------- .../multibody/multibody_from_rapier.rs | 12 +------- src_rbd/dynamics/multibody/multibody_set.rs | 7 ----- .../dynamics/multibody/multibody_solver.rs | 1 - .../dynamics/multibody/env_reset.rs | 27 ++++++------------ .../dynamics/multibody/integrate.rs | 13 ++------- 6 files changed, 16 insertions(+), 72 deletions(-) diff --git a/src_rbd/dynamics/multibody/env_reset.rs b/src_rbd/dynamics/multibody/env_reset.rs index 8d8008c5..2f782eea 100644 --- a/src_rbd/dynamics/multibody/env_reset.rs +++ b/src_rbd/dynamics/multibody/env_reset.rs @@ -33,8 +33,6 @@ pub struct GpuMultibodySnapshot { /// 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, @@ -138,7 +136,7 @@ impl EnvResetBundle { storage, ) .unwrap(), - staging_dofs: Tensor::vector(backend, vec![0.0f32; (2 * dpb).max(1) as usize], storage) + staging_dofs: Tensor::vector(backend, vec![0.0f32; dpb.max(1) as usize], storage) .unwrap(), params: Tensor::scalar(backend, UVec4::new(0, 0, lpb, dpb), uniform).unwrap(), } @@ -149,7 +147,6 @@ impl EnvResetBundle { 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. @@ -176,11 +173,6 @@ impl GpuMultibodySet { .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) @@ -196,7 +188,6 @@ impl GpuMultibodySet { 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(), } } @@ -215,7 +206,7 @@ impl GpuMultibodySet { 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); + debug_assert_eq!(snap.dof_vels.len(), dpb as usize); // Keep the host mirror in lockstep: the motor setters read-modify-write // it. @@ -237,11 +228,9 @@ impl GpuMultibodySet { 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() { + if !snap.dof_vels.is_empty() { backend - .write_buffer(bundle.staging_dofs.buffer_mut(), 0, &dofs) + .write_buffer(bundle.staging_dofs.buffer_mut(), 0, &snap.dof_vels) .unwrap(); } bundle.params = Tensor::scalar( @@ -266,7 +255,6 @@ impl GpuMultibodySet { &bundle.staging_dofs, &mut self.links_workspace, &mut self.links_static, - &mut self.dof_values, &mut self.dof_state, &bundle.params, ) @@ -290,20 +278,15 @@ impl GpuMultibodySet { 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 @@ -322,7 +305,6 @@ impl GpuMultibodySet { 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, @@ -397,11 +379,9 @@ impl GpuMultibodySet { &mut pass, [lpb.max(dpb), n, 1], &tpl.links, - &tpl.dofs, &t_resets, &t_vels, &mut self.links_static, - &mut self.dof_values, &mut self.dof_state, ¶ms, ) diff --git a/src_rbd/dynamics/multibody/multibody_from_rapier.rs b/src_rbd/dynamics/multibody/multibody_from_rapier.rs index a9ccaf68..618dc781 100644 --- a/src_rbd/dynamics/multibody/multibody_from_rapier.rs +++ b/src_rbd/dynamics/multibody/multibody_from_rapier.rs @@ -40,7 +40,6 @@ impl GpuMultibodySet { Vec::with_capacity(num_batches as usize); let mut per_env_links_workspace: Vec> = Vec::with_capacity(num_batches as usize); - let mut per_env_dof_values: Vec> = Vec::with_capacity(num_batches as usize); 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); @@ -75,7 +74,6 @@ impl GpuMultibodySet { let mut infos = Vec::new(); let mut statics = Vec::new(); let mut workspaces = Vec::new(); - let mut dof_vals = Vec::new(); let mut dof_vels = Vec::new(); let mut dof_damping = Vec::new(); let mut dof_armature = Vec::new(); @@ -304,7 +302,6 @@ impl GpuMultibodySet { } } for d in 0..link_ndofs as usize { - dof_vals.push(0.0); dof_vels.push(0.0); dof_damping.push(mb_damping[rapier_assembly + d]); dof_armature.push(mb_armature[rapier_assembly + d]); @@ -347,7 +344,7 @@ impl GpuMultibodySet { global_max_mb = global_max_mb.max(infos.len() as u32); global_max_links = global_max_links.max(statics.len() as u32); - global_max_dofs = global_max_dofs.max(dof_vals.len() as u32); + global_max_dofs = global_max_dofs.max(dof_vels.len() as u32); global_max_jac = global_max_jac.max(jac_off); global_max_mm = global_max_mm.max(mm_off); global_max_cor = global_max_cor.max(cor_off); @@ -358,7 +355,6 @@ impl GpuMultibodySet { per_env_infos.push(infos); per_env_links_static.push(statics); per_env_links_workspace.push(workspaces); - per_env_dof_values.push(dof_vals); per_env_dof_vels.push(dof_vels); per_env_dof_damping.push(dof_damping); per_env_dof_armature.push(dof_armature); @@ -414,7 +410,6 @@ impl GpuMultibodySet { Vec::with_capacity((links_cap * num_batches) as usize); let mut all_ws: Vec = Vec::with_capacity((links_cap * num_batches) as usize); - let mut all_dof_vals: Vec = Vec::with_capacity((dofs_cap * num_batches) as usize); 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); @@ -468,9 +463,6 @@ impl GpuMultibodySet { all_ws.push(dummy_ws); } - all_dof_vals.extend_from_slice(&per_env_dof_values[i]); - let pad = (dofs_cap as usize).saturating_sub(per_env_dof_values[i].len()); - all_dof_vals.resize(all_dof_vals.len() + pad, 0.0); all_dof_vels.extend_from_slice(&per_env_dof_vels[i]); let pad = (dofs_cap as usize).saturating_sub(per_env_dof_vels[i].len()); all_dof_vels.resize(all_dof_vels.len() + pad, 0.0); @@ -515,7 +507,6 @@ impl GpuMultibodySet { let info_mirror = all_infos.clone(); let all_infos = interleave(&all_infos, mb_cap, nb); let all_statics = interleave(&all_statics, links_cap, nb); - let all_dof_vals = interleave(&all_dof_vals, dofs_cap, nb); 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); @@ -563,7 +554,6 @@ impl GpuMultibodySet { storage | BufferUsages::COPY_SRC, ) .unwrap(), - dof_values: Tensor::vector(backend, &all_dof_vals, storage).unwrap(), dof_state: { // Pack [velocities, damping, armature, spring stiffness, // spring rest, kinematic mask, frictionloss] back-to-back, diff --git a/src_rbd/dynamics/multibody/multibody_set.rs b/src_rbd/dynamics/multibody/multibody_set.rs index 9e8bdca4..a0764728 100644 --- a/src_rbd/dynamics/multibody/multibody_set.rs +++ b/src_rbd/dynamics/multibody/multibody_set.rs @@ -121,8 +121,6 @@ pub struct GpuMultibodySet { pub(super) info_mirror: Vec, /// Per-batch per-step link workspace, SoA quad layout. pub(super) links_workspace: Tensor, - /// Generalized coordinates (flat). - pub(super) dof_values: Tensor, /// Packed buffer holding generalized velocities (offset 0) and per-DOF /// damping coefficients (offset `damping_section_offset`). Callers reading /// velocities should use only the velocity section. @@ -338,11 +336,6 @@ impl GpuMultibodySet { self.dofs_per_batch } - /// GPU buffer for generalized coordinates. - pub fn dof_values(&self) -> &Tensor { - &self.dof_values - } - /// GPU buffer for the last-computed generalized accelerations (populated by /// `GpuMultibodySolver::solve_gravity`). pub fn gen_accelerations(&self) -> &Tensor { diff --git a/src_rbd/dynamics/multibody/multibody_solver.rs b/src_rbd/dynamics/multibody/multibody_solver.rs index ff39f1ef..b05efa4a 100644 --- a/src_rbd/dynamics/multibody/multibody_solver.rs +++ b/src_rbd/dynamics/multibody/multibody_solver.rs @@ -657,7 +657,6 @@ impl GpuMultibodySolver { &mb.multibody_info, &mb.links_static, &mut mb.links_workspace, - &mut mb.dof_values, &mb.dof_state, &mb.dt, args.batch_indices, diff --git a/src_rbd_shaders/dynamics/multibody/env_reset.rs b/src_rbd_shaders/dynamics/multibody/env_reset.rs index 1ed94c57..cdeaa6c8 100644 --- a/src_rbd_shaders/dynamics/multibody/env_reset.rs +++ b/src_rbd_shaders/dynamics/multibody/env_reset.rs @@ -32,10 +32,9 @@ pub fn gpu_mb_env_reset( #[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], + #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] 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, + #[spirv(uniform, descriptor_set = 0, binding = 6)] params: &UVec4, ) { let i = invocation_id.x; let env = params.x; @@ -55,11 +54,7 @@ pub fn gpu_mb_env_reset( 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), - ); + dof_state.write((i * nb + env) as usize, staging_dofs.read(i as usize)); } } @@ -139,15 +134,13 @@ 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)] + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] resets: &[UVec4], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] dof_vels: &[f32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] 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], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] 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, + #[spirv(uniform, descriptor_set = 0, binding = 5)] params: &UVec4, ) { let i = invocation_id.x; let r = invocation_id.y; @@ -168,10 +161,6 @@ pub fn gpu_mb_env_reset_batch_dofs( ); } if i < dpb { - dof_values.write( - (i * nb + env) as usize, - templates_dofs.read((t * 2 * dpb + i) as usize), - ); dof_state.write( (i * nb + env) as usize, dof_vels.read((r * dpb + i) as usize), diff --git a/src_rbd_shaders/dynamics/multibody/integrate.rs b/src_rbd_shaders/dynamics/multibody/integrate.rs index f75dd47c..12e9c5fd 100644 --- a/src_rbd_shaders/dynamics/multibody/integrate.rs +++ b/src_rbd_shaders/dynamics/multibody/integrate.rs @@ -67,10 +67,9 @@ pub fn gpu_mb_integrate( #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] links_static: &[MultibodyLinkStatic], #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] links_workspace: &mut [Vec4], - #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] dof_values: &mut [f32], - #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] dof_state: &[f32], - #[spirv(uniform, descriptor_set = 0, binding = 5)] dt_uniform: &f32, - #[spirv(uniform, descriptor_set = 0, binding = 6)] batch_ids: &BatchIndices, + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] dof_state: &[f32], + #[spirv(uniform, descriptor_set = 0, binding = 4)] dt_uniform: &f32, + #[spirv(uniform, descriptor_set = 0, binding = 5)] batch_ids: &BatchIndices, ) { let num_mb = batch_ids.multibodies_len; if invocation_id.x >= num_mb * batch_ids.num_batches { @@ -87,9 +86,6 @@ pub fn gpu_mb_integrate( .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 dof_val = batch_ids - .ib_mut(batch_id, dof_values) - .offset(mb.first_dof as usize); let dof_vel = batch_ids .ib(batch_id, dof_state) .offset(mb.first_dof as usize); @@ -174,7 +170,4 @@ pub fn gpu_mb_integrate( // num_ang == 0: no-op. } - // Silence dof_val unused warning — it will be used once we also support - // setting coords directly (e.g. user-controlled kinematic DOFs). - let _ = dof_val.buf; } From 7d74c54d3ab4b1515b83f2d4e17bf3cbf1805ff4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Sun, 30 Aug 2026 10:48:22 +0200 Subject: [PATCH 4/7] perf(rbd): positional contact slots, per-pair manifold reduction, contact-to-multibody index --- src_rbd/broad_phase/narrow_phase.rs | 203 ++++--- src_rbd/dynamics/coloring.rs | 25 +- .../multibody/multibody_from_rapier.rs | 23 +- src_rbd/dynamics/multibody/multibody_set.rs | 7 +- .../dynamics/multibody/multibody_solver.rs | 71 +-- src_rbd/dynamics/solver.rs | 41 +- src_rbd/dynamics/warmstart.rs | 16 +- src_rbd/pipeline/insertion_removal.rs | 85 ++- src_rbd/pipeline/rbd_state.rs | 20 +- src_rbd/pipeline/rbd_state_from_rapier.rs | 129 ++-- src_rbd/pipeline/rbd_step.rs | 64 +- src_rbd/pipeline/test_batched_stacks.rs | 155 ++++- src_rbd/utils/radix_sort/mod.rs | 2 +- src_rbd_shaders/broad_phase/lbvh.rs | 43 +- src_rbd_shaders/broad_phase/narrow_phase.rs | 561 ++++++++---------- src_rbd_shaders/dynamics/color_buckets.rs | 17 +- src_rbd_shaders/dynamics/coloring.rs | 28 +- .../dynamics/multibody/contact_constraints.rs | 235 +++++--- src_rbd_shaders/dynamics/multibody/types.rs | 17 +- src_rbd_shaders/dynamics/solver.rs | 30 +- src_rbd_shaders/dynamics/warmstart.rs | 13 +- 21 files changed, 890 insertions(+), 895 deletions(-) diff --git a/src_rbd/broad_phase/narrow_phase.rs b/src_rbd/broad_phase/narrow_phase.rs index 94a47c1f..55c3fb7f 100644 --- a/src_rbd/broad_phase/narrow_phase.rs +++ b/src_rbd/broad_phase/narrow_phase.rs @@ -6,41 +6,82 @@ use crate::shaders::PaddedVector; #[cfg(feature = "dim3")] use crate::shaders::broad_phase::GpuReduceContacts; use crate::shaders::broad_phase::{ - CollisionPair, GpuContactOffsetsScan, GpuCountPairsPerBatch, GpuCountPfmPerBatch, - GpuFlatListDispatch, GpuNarrowPhaseInitContactsDispatch, GpuNarrowPhasePfmPfm, - GpuNarrowPhaseShapeShape, GpuNarrowPhaseShapeShapeDeferred, GpuResetNarrowPhase, - GpuZeroContactLens, NarrowPhasePfmPair, + CollisionPair, ContactPlan, GpuContactPlan, GpuNarrowPhasePfmPfm, GpuNarrowPhaseShapeShape, + GpuNarrowPhaseShapeShapeDeferred, GpuPfmSortKeys, GpuResetNarrowPhase, NarrowPhasePfmPair, }; use crate::shaders::shapes::Shape; -use khal::Shader; -use khal::backend::{GpuBackendError, GpuPass}; +use crate::utils::{RadixSort, RadixSortWorkspace}; +use khal::backend::{GpuBackend, GpuBackendError, GpuPass}; +use khal::{BufferUsages, Shader}; use vortx::tensor::Tensor; -/// GPU shader for narrow-phase collision detection. #[derive(Shader)] -pub struct GpuNarrowPhase { +struct GpuNarrowPhaseShaders { reset_narrow_phase: GpuResetNarrowPhase, narrow_phase: GpuNarrowPhaseShapeShape, - /// Pass 2: defers complex shape pairs (PFM / trimesh / polyline) into the /// `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, - count_pairs_per_batch: GpuCountPairsPerBatch, - count_pfm_per_batch: GpuCountPfmPerBatch, - contact_offsets_scan: GpuContactOffsetsScan, - zero_contact_lens: GpuZeroContactLens, - flat_list_dispatch: GpuFlatListDispatch, - init_contacts_indirect_args: GpuNarrowPhaseInitContactsDispatch, + contact_plan: GpuContactPlan, + pfm_sort_keys: GpuPfmSortKeys, +} + +pub struct GpuNarrowPhase { + shaders: GpuNarrowPhaseShaders, + sort: RadixSort, +} + +impl GpuNarrowPhase { + pub fn from_backend(backend: &GpuBackend) -> Result { + Ok(Self { + shaders: GpuNarrowPhaseShaders::from_backend(backend)?, + sort: RadixSort::from_backend(backend)?, + }) + } +} + +pub struct PfmSortState { + keys: Tensor, + identity: Tensor, + sorted_keys: Tensor, + sorted_values: Tensor, + sort_len: Tensor, + workspace: RadixSortWorkspace, +} + +impl PfmSortState { + pub fn new(backend: &GpuBackend, capacity: u32) -> Self { + let storage = BufferUsages::STORAGE; + let identity: Vec = (0..capacity).collect(); + Self { + keys: Tensor::vector_uninit(backend, capacity.max(1), storage).unwrap(), + identity: Tensor::vector(backend, &identity, storage).unwrap(), + sorted_keys: Tensor::vector_uninit(backend, capacity.max(1), storage).unwrap(), + sorted_values: Tensor::vector_uninit(backend, capacity.max(1), storage).unwrap(), + sort_len: Tensor::vector(backend, &[0u32], storage).unwrap(), + workspace: RadixSortWorkspace::new(backend), + } + } + + pub fn resize(&mut self, backend: &GpuBackend, capacity: u32) { + *self = Self::new(backend, capacity); + } + + #[cfg(feature = "dim3")] + pub fn sorted_keys(&self) -> &Tensor { + &self.sorted_keys + } } impl GpuNarrowPhase { /// Dispatches the narrow-phase collision detection pipeline. + #[allow(clippy::too_many_arguments)] pub fn dispatch( &self, + backend: &GpuBackend, pass: &mut GpuPass, - _num_colliders: u32, poses: &Tensor, shapes: &Tensor, vertices: &Tensor, @@ -48,37 +89,29 @@ impl GpuNarrowPhase { collision_pairs: &Tensor, collision_pairs_len: &mut Tensor, contacts: &mut Tensor, - contacts_len: &mut Tensor, contacts_indirect: &mut Tensor<[u32; 3]>, - contact_offsets: &mut Tensor, - pair_batch_counts: &mut Tensor, - pfm_batch_counts: &mut Tensor, + contact_plan: &mut Tensor, mb_sweep_indirect: &mut Tensor<[u32; 3]>, pfm_pairs: &mut Tensor, pfm_pairs_len: &mut Tensor, pfm_pairs_indirect: &mut Tensor<[u32; 3]>, + pfm_sort: &mut PfmSortState, batch_indices: &Tensor, collider_parent: &Tensor, collider_materials: &Tensor, sim_params: &Tensor, // Optional: merge each collider pair's manifolds into one before the - // solvers see them. `false` skips the kernel entirely. reduce_contacts: bool, collision_pairs_indirect: &Tensor<[u32; 3]>, + collision_pairs_capacity: u32, ) -> 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, - pair_batch_counts, - pfm_batch_counts, - )?; + let reduce_contacts = reduce_contacts && cfg!(feature = "dim3"); + + self.shaders + .reset_narrow_phase + .call(pass, 1u32, pfm_pairs_len)?; - // Pass 2: defer the complex shape pairs into `pfm_pairs` (kept as a - // separate dispatch so each pass fits 8 storage buffers). - self.narrow_phase_deferred.call( + self.shaders.narrow_phase_deferred.call( pass, collision_pairs_indirect, collision_pairs, @@ -93,92 +126,94 @@ impl GpuNarrowPhase { indices, )?; - self.count_pairs_per_batch.call( - pass, - collision_pairs_indirect, - collision_pairs, - collision_pairs_len, - pair_batch_counts, - batch_indices, - )?; - self.flat_list_dispatch - .call(pass, 1u32, pfm_pairs_len, pfm_pairs_indirect, batch_indices)?; - self.count_pfm_per_batch.call( - pass, - &*pfm_pairs_indirect, - pfm_pairs, - pfm_pairs_len, - pfm_batch_counts, - batch_indices, - )?; - self.contact_offsets_scan.call( + self.shaders.contact_plan.call( pass, 1u32, - pair_batch_counts, - pfm_batch_counts, collision_pairs_len, pfm_pairs_len, - contact_offsets, + contact_plan, + &mut pfm_sort.sort_len, contacts_indirect, - batch_indices, - )?; - self.zero_contact_lens.call( - pass, - &*contacts_indirect, - contacts, - contact_offsets, + pfm_pairs_indirect, + mb_sweep_indirect, batch_indices, )?; - self.narrow_phase.call( + self.shaders.narrow_phase.call( pass, collision_pairs_indirect, collision_pairs, - contact_offsets, + &*contact_plan, poses, shapes, contacts, - contacts_len, - batch_indices, collider_parent, collider_materials, sim_params, )?; - self.narrow_phase_pfm_pfm.call( + + if reduce_contacts { + self.shaders.pfm_sort_keys.call( + pass, + &*pfm_pairs_indirect, + pfm_pairs, + &*contact_plan, + &mut pfm_sort.keys, + )?; + let sorting_bits = + (32 - collision_pairs_capacity.saturating_sub(1).leading_zeros()).max(1); + let PfmSortState { + keys, + identity, + sorted_keys, + sorted_values, + sort_len, + workspace, + } = pfm_sort; + self.sort.dispatch( + backend, + pass, + workspace, + keys, + identity, + sort_len, + sorting_bits, + 1, + sorted_keys, + sorted_values, + )?; + } + + let pfm_order = if reduce_contacts { + &pfm_sort.sorted_values + } else { + &pfm_sort.identity + }; + self.shaders.narrow_phase_pfm_pfm.call( pass, &*pfm_pairs_indirect, contacts, - contacts_len, pfm_pairs, - contact_offsets, - batch_indices, + pfm_order, + &*contact_plan, vertices, indices, collider_parent, collider_materials, sim_params, )?; + #[cfg(feature = "dim3")] if reduce_contacts { - self.reduce_contacts.call( + self.shaders.reduce_contacts.call( pass, - [1u32, num_batches, 1], + &*pfm_pairs_indirect, contacts, - contacts_len, - contact_offsets, - batch_indices, + &pfm_sort.sorted_keys, + &*contact_plan, sim_params, )?; } - #[cfg(not(feature = "dim3"))] - let _ = reduce_contacts; - self.init_contacts_indirect_args.call( - pass, - 256u32, - contacts_len, - mb_sweep_indirect, - batch_indices, - )?; Ok(()) } diff --git a/src_rbd/dynamics/coloring.rs b/src_rbd/dynamics/coloring.rs index 2c21a64e..985690b2 100644 --- a/src_rbd/dynamics/coloring.rs +++ b/src_rbd/dynamics/coloring.rs @@ -8,6 +8,7 @@ //! and handles arbitrary constraint graphs. use crate::pipeline::RunStats; +use crate::shaders::broad_phase::ContactPlan; use crate::shaders::dynamics::TwoBodyConstraint; use crate::shaders::dynamics::{ GpuColorBucketsCount, GpuColorBucketsReset, GpuColorBucketsScatter, GpuFixConflictsTopoGc, @@ -44,12 +45,11 @@ pub struct GpuColoring { /// Buffers for the per-color constraint bucket sort. pub struct ColorBucketsArgs<'a> { - /// Indirect dispatch arguments based on contact count. pub contacts_len_indirect: &'a Tensor<[u32; 3]>, /// Color assigned to each constraint by graph coloring. pub constraints_colors: &'a Tensor, pub constraints: &'a Tensor, - pub contact_offsets: &'a Tensor, + pub contact_plan: &'a Tensor, pub color_buckets: &'a mut Tensor, pub color_sorted_ids: &'a mut Tensor, /// Shared per-batch capacity / section-offset uniform. @@ -78,7 +78,7 @@ pub struct ColoringArgs<'a> { pub uncolored: &'a mut Tensor, /// Staging buffer for reading uncolored count on CPU. pub uncolored_staging: &'a Tensor, - pub contact_offsets: &'a Tensor, + pub contact_plan: &'a Tensor, /// Buffer tracking which constraints are colored. pub colored: &'a mut Tensor, /// Shared per-batch capacity / section-offset uniform. @@ -102,8 +102,7 @@ impl GpuColoring { args.constraints_colors, args.constraints_rands, args.constraints, - args.contact_offsets, - args.batch_indices, + args.contact_plan, )?; Ok(()) } @@ -125,8 +124,7 @@ impl GpuColoring { args.uncolored, args.body_group, args.curr_color, - args.contact_offsets, - args.batch_indices, + args.contact_plan, )?; Ok(()) } @@ -143,8 +141,7 @@ impl GpuColoring { args.constraints_colors, args.colored, args.constraints, - args.contact_offsets, - args.batch_indices, + args.contact_plan, )?; Ok(()) } @@ -164,9 +161,8 @@ impl GpuColoring { args.constraints_colors, args.colored, args.uncolored, - args.contact_offsets, + args.contact_plan, args.body_group, - args.batch_indices, )?; Ok(()) } @@ -186,9 +182,8 @@ impl GpuColoring { args.constraints_colors, args.colored, args.uncolored, - args.contact_offsets, + args.contact_plan, args.body_group, - args.batch_indices, )?; Ok(()) } @@ -260,7 +255,7 @@ impl GpuColoring { args.contacts_len_indirect, args.constraints_colors, args.constraints, - args.contact_offsets, + args.contact_plan, args.color_buckets, args.batch_indices, )?; @@ -270,7 +265,7 @@ impl GpuColoring { args.contacts_len_indirect, args.constraints_colors, args.constraints, - args.contact_offsets, + args.contact_plan, args.color_buckets, args.color_sorted_ids, args.batch_indices, diff --git a/src_rbd/dynamics/multibody/multibody_from_rapier.rs b/src_rbd/dynamics/multibody/multibody_from_rapier.rs index 618dc781..627eb9ce 100644 --- a/src_rbd/dynamics/multibody/multibody_from_rapier.rs +++ b/src_rbd/dynamics/multibody/multibody_from_rapier.rs @@ -179,8 +179,8 @@ impl GpuMultibodySet { contact_constraint_start: 0, old_contact_constraint_start: 0, old_contact_constraint_count: 0, - batch_contacts_len: 0, - batch_contacts_start: 0, + contact_index_len: 0, + contact_index_start: 0, first_coupling: coupling_off, num_couplings, }); @@ -380,7 +380,10 @@ impl GpuMultibodySet { let contact_cons_cap = contact_constraint_slots .saturating_mul(num_batches) - .max((mb_cap * num_batches).saturating_mul(crate::shaders::dynamics::MB_CONS_SLOT_RESERVE)) + .max( + (mb_cap * num_batches) + .saturating_mul(crate::shaders::dynamics::MB_CONS_SLOT_RESERVE), + ) .max(1); let contact_cons_col_cap = contact_cons_cap.saturating_mul(dofs_cap).max(1); let body_to_link_cap = colliders_per_batch.max(1); @@ -760,10 +763,18 @@ impl GpuMultibodySet { joint_constraints_per_batch: cons_cap, joint_constraint_columns_per_batch: cons_col_cap, contact_constraints_capacity: contact_cons_cap, - mb_cons_demand: Tensor::vector( + mb_cons_demand: Tensor::vector(backend, &[0u32], storage | BufferUsages::COPY_SRC) + .unwrap(), + mb_cons_counts: Tensor::vector( backend, - &[0u32], - storage | BufferUsages::COPY_SRC, + vec![0u32; (mb_cap * num_batches) as usize], + storage, + ) + .unwrap(), + mb_index_counts: Tensor::vector( + backend, + vec![0u32; (mb_cap * num_batches) as usize], + storage, ) .unwrap(), diff --git a/src_rbd/dynamics/multibody/multibody_set.rs b/src_rbd/dynamics/multibody/multibody_set.rs index a0764728..1ef62a20 100644 --- a/src_rbd/dynamics/multibody/multibody_set.rs +++ b/src_rbd/dynamics/multibody/multibody_set.rs @@ -228,6 +228,8 @@ pub struct GpuMultibodySet { pub(super) joint_constraint_columns_per_batch: u32, pub(super) contact_constraints_capacity: u32, pub(super) mb_cons_demand: Tensor, + pub(super) mb_cons_counts: Tensor, + pub(super) mb_index_counts: Tensor, /// Number of solver iterations to run on `joint_constraints` per `step()`. pub(super) num_solver_iterations: u32, @@ -807,13 +809,14 @@ impl GpuMultibodySet { ( i.contact_constraint_start, i.contact_constraint_count, - i.batch_contacts_start, - i.batch_contacts_len, + i.contact_index_start, + i.contact_index_len, ) }) .collect(); (out, demand.first().copied().unwrap_or(0)) } + pub(crate) fn resize_contact_slabs(&mut self, backend: &GpuBackend, new_capacity: u32) { use khal::BufferUsages; let storage = BufferUsages::STORAGE | BufferUsages::COPY_DST; diff --git a/src_rbd/dynamics/multibody/multibody_solver.rs b/src_rbd/dynamics/multibody/multibody_solver.rs index b05efa4a..aa2b9220 100644 --- a/src_rbd/dynamics/multibody/multibody_solver.rs +++ b/src_rbd/dynamics/multibody/multibody_solver.rs @@ -3,19 +3,20 @@ use super::multibody_set::*; use crate::math::Pose; use crate::queries::GpuIndexedContact; +use crate::shaders::broad_phase::ContactPlan; use crate::shaders::dynamics::{ GpuMbApplyContactRestitution, GpuMbBuildContactDelassus, GpuMbComputeDynamicsPre, - GpuMbComputeSolveBounds, GpuMbDelayTick, GpuMbFinalizeContactConstraints, - GpuMbFinalizeImpulseJointConstraints, GpuMbFinalizeJointConstraints, GpuMbGravityAndLu, - GpuMbGravityAndLuT1, GpuMbGravityAndLuT8, GpuMbGravityAndLuT16, GpuMbGravityAndLuT32, - GpuMbInitContactConstraints, GpuMbInitJointConstraints, GpuMbIntegrate, - GpuMbIntegrateVelocities, GpuMbRefreshJointConstraints, GpuMbRemoveImpulseJointConstraintBias, - GpuMbSeedContactRestitution, GpuMbSenseContactImpulses, GpuMbSnapshotContactWarmstart, - GpuMbSolveConstraints, GpuMbSolveContactsDelassus, GpuMbSolveImpulseJointConstraints, - GpuMbSolveJoints, GpuMbCountContactConstraints, GpuMbConsOffsetsScan, GpuMbSavePrevConsBounds, - GpuMbStashContactsLen, GpuMbTransferContactWarmstart, - GpuMbUpdateImpulseJointConstraints, GpuMbWarmstartContactConstraints, Velocity, - WorldMassProperties, + GpuMbComputeSolveBounds, GpuMbConsOffsetsScan, GpuMbCountContactConstraints, GpuMbDelayTick, + GpuMbFinalizeContactConstraints, GpuMbFinalizeImpulseJointConstraints, + GpuMbFinalizeJointConstraints, GpuMbGravityAndLu, GpuMbGravityAndLuT1, GpuMbGravityAndLuT8, + GpuMbGravityAndLuT16, GpuMbGravityAndLuT32, GpuMbInitContactConstraints, + GpuMbInitJointConstraints, GpuMbIntegrate, GpuMbIntegrateVelocities, + GpuMbRefreshJointConstraints, GpuMbRemoveImpulseJointConstraintBias, GpuMbSavePrevConsBounds, + GpuMbScatterContactIndex, GpuMbSeedContactRestitution, GpuMbSenseContactImpulses, + GpuMbSnapshotContactWarmstart, GpuMbSolveConstraints, GpuMbSolveContactsDelassus, + GpuMbSolveImpulseJointConstraints, GpuMbSolveJoints, GpuMbTransferContactWarmstart, + GpuMbUpdateImpulseJointConstraints, GpuMbWarmstartContactConstraints, MbContactIndexEntry, + Velocity, WorldMassProperties, }; use crate::shaders::utils::BatchIndices; use khal::Shader; @@ -70,12 +71,9 @@ pub struct GpuMultibodySolver { save_prev_cons_bounds: GpuMbSavePrevConsBounds, count_contact_constraints: GpuMbCountContactConstraints, cons_offsets_scan: GpuMbConsOffsetsScan, + scatter_contact_index: GpuMbScatterContactIndex, /// Carry the snapshotted impulses over to this frame's matching contacts. transfer_contact_warmstart: GpuMbTransferContactWarmstart, - /// Copy `contacts_len[batch]` into each `MultibodyInfo` once per step so - /// `init_contact_constraints` (at the 8-storage-buffer limit) can bound - /// its manifold scan by the actual count instead of the capacity. - stash_contacts_len: GpuMbStashContactsLen, /// Re-apply the accumulated contact impulse each substep (warmstart). warmstart_contact_constraints: GpuMbWarmstartContactConstraints, /// Capture each bouncy contact's approach velocity at the start of the step. @@ -104,11 +102,10 @@ pub struct MultibodySolverArgs<'a> { pub collider_world_poses: &'a Tensor, /// Free-body world mass properties (read by `init_contact_constraints`). pub mprops: &'a Tensor, - /// Per-batch contact manifold list (filled by narrow-phase). pub contacts: &'a Tensor, - /// Per-batch contact count (parallel to `contacts`). - pub contacts_len: &'a Tensor, - pub contact_offsets: &'a Tensor, + pub contact_plan: &'a Tensor, + pub contacts_indirect: &'a Tensor<[u32; 3]>, + pub mb_contact_index: &'a mut Tensor, /// Free-body solver velocities (updated in place by `solve_contact_constraints`). pub solver_vels: &'a mut Tensor, /// Shared `BatchIndices` uniform — per-batch caps and packed-section @@ -180,11 +177,7 @@ impl GpuMultibodySolver { self.compute_dynamics(&mut pass, mb, args) } - /// Copy `contacts_len[batch]` into each `MultibodyInfo`. - /// - /// This is a workaround for kernels that are already at the 8-storage-binding - /// web limit and could therefore not bind `contacts_len`. - pub fn stash_contacts_len( + pub fn layout_contact_constraints( &self, pass: &mut GpuPass, mb: &mut GpuMultibodySet, @@ -193,29 +186,37 @@ impl GpuMultibodySolver { if mb.is_empty() { return Ok(()); } - self.stash_contacts_len.call( - pass, - mb.flat_mb_dispatch(), - &mut mb.multibody_info, - args.contacts_len, - args.contact_offsets, - args.batch_indices, - )?; self.count_contact_constraints.call( pass, - mb.flat_mb_dispatch(), - &mut mb.multibody_info, + args.contacts_indirect, + &mb.multibody_info, args.contacts, &mb.body_to_link, + &mut mb.mb_cons_counts, + &mut mb.mb_index_counts, + args.contact_plan, args.batch_indices, )?; self.cons_offsets_scan.call( pass, 1u32, &mut mb.multibody_info, + &mut mb.mb_cons_counts, + &mut mb.mb_index_counts, &mut mb.mb_cons_demand, args.batch_indices, )?; + self.scatter_contact_index.call( + pass, + args.contacts_indirect, + &mb.multibody_info, + args.contacts, + &mb.body_to_link, + &mut mb.mb_index_counts, + args.mb_contact_index, + args.contact_plan, + args.batch_indices, + )?; Ok(()) } @@ -416,7 +417,7 @@ impl GpuMultibodySolver { init_contact_dispatch, &mut mb.multibody_info, &mb.links_workspace, - &mb.body_to_link, + &*args.mb_contact_index, &mut mb.contact_constraints, &mb.constraint_softness, args.batch_indices, diff --git a/src_rbd/dynamics/solver.rs b/src_rbd/dynamics/solver.rs index 752ac401..26a24a4f 100644 --- a/src_rbd/dynamics/solver.rs +++ b/src_rbd/dynamics/solver.rs @@ -9,6 +9,9 @@ use crate::dynamics::joint::{GpuJointSolver, JointSolverArgs}; use crate::dynamics::multibody::{GpuMultibodySet, GpuMultibodySolver, MultibodySolverArgs}; use crate::math::Pose; use crate::queries::GpuIndexedContact; +use crate::shaders::broad_phase::ContactPlan; +#[cfg(feature = "dim3")] +use crate::shaders::dynamics::MbContactIndexEntry; use crate::shaders::dynamics::{ GpuApplySolverVelsInc, GpuInitSolverBodies, GpuInitSolverVelsInc, GpuIntegrateLinearized, GpuSolverCleanup, GpuSolverCountConstraints, GpuSolverFinalize, GpuSolverInitConstraints, @@ -76,9 +79,9 @@ pub struct SolverArgs<'a> { pub num_colliders: u32, /// Contact manifolds generated by narrow-phase. pub contacts: &'a Tensor, - /// Number of contacts (per batch). - pub contacts_len: &'a Tensor, - pub contact_offsets: &'a Tensor, + pub contact_plan: &'a Tensor, + #[cfg(feature = "dim3")] + pub mb_contact_index: &'a mut Tensor, pub contacts_len_indirect: &'a Tensor<[u32; 3]>, /// Solver constraints (output from constraint initialization). pub constraints: &'a mut Tensor, @@ -196,13 +199,12 @@ impl GpuSolver { args.contacts, args.constraints, args.constraint_builders, - args.contact_offsets, + args.contact_plan, args.collider_world_poses, args.solver_body_poses, args.vels, args.mprops, args.sim_params, - args.batch_indices, )?; // Counting runs as a separate dispatch (same indirect grid) so the @@ -214,8 +216,7 @@ impl GpuSolver { args.body_constraint_counts, args.body_group, args.mprops, - args.contact_offsets, - args.batch_indices, + args.contact_plan, )?; args.prefix_sum.launch( @@ -232,10 +233,9 @@ impl GpuSolver { args.body_constraint_counts, args.mprops, args.contacts, - args.contact_offsets, + args.contact_plan, args.body_constraint_ids, args.body_group, - args.batch_indices, )?; Ok(()) @@ -292,15 +292,16 @@ impl GpuSolver { collider_world_poses: args.collider_world_poses, mprops: args.mprops, contacts: args.contacts, - contacts_len: args.contacts_len, - contact_offsets: args.contact_offsets, + contact_plan: args.contact_plan, + contacts_indirect: args.contacts_len_indirect, + mb_contact_index: &mut *args.mb_contact_index, solver_vels: &mut *args.solver_vels, batch_indices: args.batch_indices, gravity: args.gravity, color_uniforms: args.color_uniforms, mb_sweep_indirect: args.mb_sweep_indirect, }; - solver.stash_contacts_len(&mut pass, state, &mut mb_args)?; + solver.layout_contact_constraints(&mut pass, state, &mut mb_args)?; } } @@ -318,8 +319,9 @@ impl GpuSolver { collider_world_poses: args.collider_world_poses, mprops: args.mprops, contacts: args.contacts, - contacts_len: args.contacts_len, - contact_offsets: args.contact_offsets, + contact_plan: args.contact_plan, + contacts_indirect: args.contacts_len_indirect, + mb_contact_index: &mut *args.mb_contact_index, solver_vels: &mut *args.solver_vels, batch_indices: args.batch_indices, gravity: args.gravity, @@ -364,8 +366,9 @@ impl GpuSolver { collider_world_poses: args.collider_world_poses, mprops: args.mprops, contacts: args.contacts, - contacts_len: args.contacts_len, - contact_offsets: args.contact_offsets, + contact_plan: args.contact_plan, + contacts_indirect: args.contacts_len_indirect, + mb_contact_index: &mut *args.mb_contact_index, solver_vels: &mut *args.solver_vels, batch_indices: args.batch_indices, gravity: args.gravity, @@ -391,10 +394,9 @@ impl GpuSolver { args.contacts_len_indirect, args.constraints, args.constraint_builders, - args.contact_offsets, + args.contact_plan, args.solver_body_poses, args.sim_params, - args.batch_indices, )?; } joint_solver.update(pass, &mut joint_args, args.solver_body_poses)?; @@ -519,10 +521,9 @@ impl GpuSolver { args.contacts_len_indirect, args.constraints, args.constraint_builders, - args.contact_offsets, + args.contact_plan, args.solver_body_poses, args.sim_params, - args.batch_indices, )?; } joint_solver.solve(pass, &mut joint_args, args.solver_vels, false)?; diff --git a/src_rbd/dynamics/warmstart.rs b/src_rbd/dynamics/warmstart.rs index 5b4e271c..b7d67959 100644 --- a/src_rbd/dynamics/warmstart.rs +++ b/src_rbd/dynamics/warmstart.rs @@ -1,10 +1,10 @@ //! Warmstarting: reuses previous-frame impulses for faster solver convergence. +use crate::shaders::broad_phase::ContactPlan; use crate::shaders::dynamics::{ GpuSeedColorsFromWarmstart, GpuTransferWarmstartImpulses, TwoBodyConstraint, TwoBodyConstraintBuilder, }; -use crate::shaders::utils::BatchIndices; use khal::Shader; use khal::backend::{GpuBackendError, GpuPass}; use vortx::tensor::Tensor; @@ -26,7 +26,7 @@ pub struct GpuWarmstart { /// /// Contains buffers for both old (previous frame) and new (current frame) constraint data. pub struct WarmstartArgs<'a> { - pub contact_offsets: &'a Tensor, + pub contact_plan: &'a Tensor, /// Constraint counts per body from previous frame. pub old_body_constraint_counts: &'a Tensor, /// Constraint IDs per body from previous frame. @@ -41,13 +41,11 @@ pub struct WarmstartArgs<'a> { pub new_constraint_builders: &'a Tensor, /// Indirect dispatch arguments based on contact count. pub contacts_len_indirect: &'a Tensor<[u32; 3]>, - /// Shared per-batch index uniform. - pub batch_indices: &'a Tensor, } /// Arguments for the coloring seed dispatch. pub struct SeedColorsArgs<'a> { - pub contact_offsets: &'a Tensor, + pub contact_plan: &'a Tensor, /// Constraint counts per body from previous frame. pub old_body_constraint_counts: &'a Tensor, /// Constraint IDs per body from previous frame. @@ -64,8 +62,6 @@ pub struct SeedColorsArgs<'a> { pub colored: &'a mut Tensor, /// Indirect dispatch arguments based on contact count. pub contacts_len_indirect: &'a Tensor<[u32; 3]>, - /// Shared per-batch index uniform. - pub batch_indices: &'a Tensor, } impl GpuWarmstart { @@ -84,8 +80,7 @@ impl GpuWarmstart { args.old_constraint_builders, args.new_constraints, args.new_constraint_builders, - args.contact_offsets, - args.batch_indices, + args.contact_plan, ) } @@ -106,8 +101,7 @@ impl GpuWarmstart { args.old_constraints_colors, args.constraints_colors, args.colored, - args.contact_offsets, - args.batch_indices, + args.contact_plan, ) } } diff --git a/src_rbd/pipeline/insertion_removal.rs b/src_rbd/pipeline/insertion_removal.rs index dcef8bc9..1aaef983 100644 --- a/src_rbd/pipeline/insertion_removal.rs +++ b/src_rbd/pipeline/insertion_removal.rs @@ -1,11 +1,12 @@ //! Incremental construction of [`RbdState`]: empty allocation, append and removal of bodies. -use crate::broad_phase::LbvhState; +use crate::broad_phase::{LbvhState, PfmSortState}; use crate::dynamics::GpuImpulseJointSet; #[cfg(feature = "dim3")] use crate::dynamics::GpuMultibodySet; use crate::math::{Pose, Vector}; use crate::queries::GpuColliderMaterial; +use crate::shaders::broad_phase::ContactPlan; use crate::shaders::dynamics::{ LocalMassProperties as GpuLocalMassProperties, RbdSimParams, Velocity as GpuVelocity, WorldMassProperties as GpuWorldMassProperties, @@ -117,8 +118,9 @@ impl RbdState { let all_collider_materials = vec![GpuColliderMaterial::default(); num_bodies_total]; let collider_materials = Tensor::vector(backend, &all_collider_materials, rw).unwrap(); - let collision_pairs = - Tensor::vector_uninit(backend, collisions_capacity * num_batches, storage).unwrap(); + let pairs_capacity = collisions_capacity * num_batches; + let contacts_capacity = pairs_capacity * 2; + let collision_pairs = Tensor::vector_uninit(backend, pairs_capacity, storage).unwrap(); let collision_pairs_len = Tensor::vector( backend, &[0u32], @@ -131,63 +133,42 @@ impl RbdState { let resize_readback = GpuReadback::new(backend, 3).unwrap(); let collision_pairs_indirect = Tensor::scalar_uninit(backend, BufferUsages::STORAGE | BufferUsages::INDIRECT).unwrap(); - let contacts = - Tensor::vector_uninit(backend, collisions_capacity * num_batches, storage).unwrap(); - let contacts_len = Tensor::vector_uninit( - backend, - num_batches, - BufferUsages::STORAGE | BufferUsages::COPY_SRC, - ) - .unwrap(); + let contacts = Tensor::vector_uninit(backend, contacts_capacity, storage).unwrap(); + #[cfg(feature = "dim3")] + let mb_contact_index = Tensor::vector_uninit(backend, contacts_capacity, storage).unwrap(); let contacts_indirect = Tensor::scalar_uninit(backend, BufferUsages::STORAGE | BufferUsages::INDIRECT).unwrap(); let mb_sweep_indirect = Tensor::scalar_uninit(backend, BufferUsages::STORAGE | BufferUsages::INDIRECT).unwrap(); let pfm_pairs_indirect = Tensor::scalar_uninit(backend, BufferUsages::STORAGE | BufferUsages::INDIRECT).unwrap(); - let pfm_pairs = - Tensor::vector_uninit(backend, collisions_capacity * num_batches, storage).unwrap(); + let pfm_pairs = Tensor::vector_uninit(backend, pairs_capacity, storage).unwrap(); + let pfm_sort = PfmSortState::new(backend, pairs_capacity); let pfm_pairs_len = Tensor::vector( backend, &[0u32], BufferUsages::STORAGE | BufferUsages::COPY_SRC, ) .unwrap(); - let old_constraints = - Tensor::vector_uninit(backend, collisions_capacity * num_batches, storage).unwrap(); + let old_constraints = Tensor::vector_uninit(backend, contacts_capacity, storage).unwrap(); let old_constraint_builders = - Tensor::vector_uninit(backend, collisions_capacity * num_batches, storage).unwrap(); - let new_constraints = - Tensor::vector_uninit(backend, collisions_capacity * num_batches, storage).unwrap(); + Tensor::vector_uninit(backend, contacts_capacity, storage).unwrap(); + let new_constraints = Tensor::vector_uninit(backend, contacts_capacity, storage).unwrap(); let new_constraint_builders = - Tensor::vector_uninit(backend, collisions_capacity * num_batches, storage).unwrap(); + Tensor::vector_uninit(backend, contacts_capacity, storage).unwrap(); let constraints_colors = - Tensor::vector_uninit(backend, collisions_capacity * num_batches, storage).unwrap(); - let old_constraints_colors = Tensor::vector( - backend, - vec![0u32; (collisions_capacity * num_batches) as usize], - storage, - ) - .unwrap(); - let colored = - Tensor::vector_uninit(backend, collisions_capacity * num_batches, storage).unwrap(); - let constraints_rands = - Tensor::vector_uninit(backend, collisions_capacity * num_batches, storage).unwrap(); + Tensor::vector_uninit(backend, contacts_capacity, storage).unwrap(); + let old_constraints_colors = + Tensor::vector(backend, vec![0u32; contacts_capacity as usize], storage).unwrap(); + let colored = Tensor::vector_uninit(backend, contacts_capacity, storage).unwrap(); + let constraints_rands = Tensor::vector_uninit(backend, contacts_capacity, storage).unwrap(); let color_buckets_stride = capacities.solver_colors + 3; let color_buckets = Tensor::vector_uninit(backend, color_buckets_stride * num_batches, storage).unwrap(); - let contact_offsets = Tensor::vector( - backend, - vec![0u32; num_batches as usize + 3], - storage, - ) - .unwrap(); - let pair_batch_counts = - Tensor::vector(backend, vec![0u32; num_batches as usize], storage).unwrap(); - let pfm_batch_counts = - Tensor::vector(backend, vec![0u32; num_batches as usize], storage).unwrap(); - let color_sorted_ids = - Tensor::vector_uninit(backend, collisions_capacity * num_batches, storage).unwrap(); + let contact_plan = + Tensor::scalar(backend, ContactPlan::default(), storage | BufferUsages::UNIFORM) + .unwrap(); + let color_sorted_ids = Tensor::vector_uninit(backend, contacts_capacity, storage).unwrap(); let old_constraints_counts = Tensor::vector( backend, vec![0u32; (num_colliders_per_batch * num_batches) as usize], @@ -197,9 +178,9 @@ impl RbdState { let new_constraints_counts = Tensor::vector_uninit(backend, num_colliders_per_batch * num_batches, storage).unwrap(); let old_body_constraint_ids = - Tensor::vector_uninit(backend, collisions_capacity * 2 * num_batches, storage).unwrap(); + Tensor::vector_uninit(backend, contacts_capacity * 2, storage).unwrap(); let new_body_constraint_ids = - Tensor::vector_uninit(backend, collisions_capacity * 2 * num_batches, storage).unwrap(); + Tensor::vector_uninit(backend, contacts_capacity * 2, storage).unwrap(); let lbvh_usages = if crate::VALIDATE_LBVH_TOPOLOGY { BufferUsages::STORAGE | BufferUsages::COPY_SRC @@ -207,8 +188,8 @@ impl RbdState { BufferUsages::STORAGE }; - let contacts_capacity_cpu = collisions_capacity * num_batches; - let collision_pairs_capacity_cpu = collisions_capacity * num_batches; + let contacts_capacity_cpu = contacts_capacity; + let collision_pairs_capacity_cpu = pairs_capacity; #[allow(unused_mut)] // Only mutated with the dim3 (multibody) feature. let mut bi = BatchIndices { num_batches, @@ -277,11 +258,11 @@ impl RbdState { mb_cons_demand_cpu: 0, batch_indices, contacts, - contacts_len, contacts_indirect, - contact_offsets, - pair_batch_counts, - pfm_batch_counts, + contact_plan, + pfm_sort, + #[cfg(feature = "dim3")] + mb_contact_index, mb_sweep_indirect, pfm_pairs, pfm_pairs_len, @@ -439,8 +420,8 @@ impl RbdState { // The incremental path attaches exactly one collider per body, so a // body's collider slot equals its body slot: `collider_parent` is the let nb = self.num_batches as usize; - let parents: Vec = ((active * nb) as u32..((active + bodies.len()) * nb) as u32) - .collect(); + let parents: Vec = + ((active * nb) as u32..((active + bodies.len()) * nb) as u32).collect(); let pair_filters: Vec<[u32; 2]> = (0..bodies.len() * nb) .map(|i| [(active + i / nb) as u32, 0u32]) .collect(); diff --git a/src_rbd/pipeline/rbd_state.rs b/src_rbd/pipeline/rbd_state.rs index 8a9a1252..e435a4e8 100644 --- a/src_rbd/pipeline/rbd_state.rs +++ b/src_rbd/pipeline/rbd_state.rs @@ -1,13 +1,15 @@ //! GPU-resident rigid-body state ([`RbdState`]): buffer definitions, accessors, //! run statistics and capacity/resize policies. -use crate::broad_phase::LbvhState; +use crate::broad_phase::{LbvhState, PfmSortState}; use crate::dynamics::GpuImpulseJointSet; #[cfg(feature = "dim3")] use crate::dynamics::GpuMultibodySet; use crate::math::Pose; use crate::queries::{GpuColliderMaterial, GpuIndexedContact}; use crate::shaders::PaddedVector; -use crate::shaders::broad_phase::{CollisionPair, NarrowPhasePfmPair}; +use crate::shaders::broad_phase::{CollisionPair, ContactPlan, NarrowPhasePfmPair}; +#[cfg(feature = "dim3")] +use crate::shaders::dynamics::MbContactIndexEntry; use crate::shaders::dynamics::{ LocalMassProperties as GpuLocalMassProperties, RbdSimParams, TwoBodyConstraint, TwoBodyConstraintBuilder, Velocity as GpuVelocity, @@ -187,11 +189,11 @@ pub struct RbdState { pub(super) pfm_pairs_len: Tensor, pub(super) pfm_pairs_indirect: Tensor<[u32; 3]>, pub(super) contacts: Tensor, - pub(super) contacts_len: Tensor, pub(super) contacts_indirect: Tensor<[u32; 3]>, - pub(super) contact_offsets: Tensor, - pub(super) pair_batch_counts: Tensor, - pub(super) pfm_batch_counts: Tensor, + pub(super) contact_plan: Tensor, + pub(super) pfm_sort: PfmSortState, + #[cfg(feature = "dim3")] + pub(super) mb_contact_index: Tensor, /// Workgroup grid for the per-multibody contact-constraint dispatches: /// `[multibodies_batch_capacity, num_batches, 1]`. pub(super) mb_sweep_indirect: Tensor<[u32; 3]>, @@ -446,12 +448,6 @@ impl RbdState { &self.collider_parent } - /// The contact manifold buffer (post narrow-phase + body resolution). - /// Per-batch number of manifolds currently in [`Self::contacts`]. - pub fn contacts_len(&self) -> &Tensor { - &self.contacts_len - } - /// GPU buffer holding the contact manifolds. pub fn contacts(&self) -> &Tensor { &self.contacts diff --git a/src_rbd/pipeline/rbd_state_from_rapier.rs b/src_rbd/pipeline/rbd_state_from_rapier.rs index 06cb42d2..59ab3d9a 100644 --- a/src_rbd/pipeline/rbd_state_from_rapier.rs +++ b/src_rbd/pipeline/rbd_state_from_rapier.rs @@ -1,11 +1,12 @@ //! Initialization of [`RbdState`] from CPU-side Rapier data structures. -use crate::broad_phase::LbvhState; +use crate::broad_phase::{LbvhState, PfmSortState}; use crate::dynamics::GpuImpulseJointSet; #[cfg(feature = "dim3")] use crate::dynamics::GpuMultibodySet; use crate::math::{Pose, Vector}; use crate::queries::GpuColliderMaterial; +use crate::shaders::broad_phase::ContactPlan; use crate::shaders::dynamics::{ LocalMassProperties as GpuLocalMassProperties, RbdSimParams, Velocity as GpuVelocity, WorldMassProperties as GpuWorldMassProperties, @@ -658,18 +659,11 @@ impl RbdState { let collision_pairs_indirect = Tensor::scalar_uninit(backend, BufferUsages::STORAGE | BufferUsages::INDIRECT).unwrap(); - let contacts = Tensor::vector_uninit( - backend, - capacities.collisions_capacity * num_batches, - storage, - ) - .unwrap(); - let contacts_len = Tensor::vector_uninit( - backend, - num_batches, - BufferUsages::STORAGE | BufferUsages::COPY_SRC, - ) - .unwrap(); + let pairs_capacity = capacities.collisions_capacity * num_batches; + let contacts_capacity = pairs_capacity * 2; + let contacts = Tensor::vector_uninit(backend, contacts_capacity, storage).unwrap(); + #[cfg(feature = "dim3")] + let mb_contact_index = Tensor::vector_uninit(backend, contacts_capacity, storage).unwrap(); let contacts_indirect = Tensor::scalar_uninit(backend, BufferUsages::STORAGE | BufferUsages::INDIRECT).unwrap(); let mb_sweep_indirect = @@ -688,73 +682,26 @@ impl RbdState { BufferUsages::STORAGE | BufferUsages::COPY_SRC, ) .unwrap(); - let old_constraints = Tensor::vector_uninit( - backend, - capacities.collisions_capacity * num_batches, - storage, - ) - .unwrap(); - let old_constraint_builders = Tensor::vector_uninit( - backend, - capacities.collisions_capacity * num_batches, - storage, - ) - .unwrap(); - let new_constraints = Tensor::vector_uninit( - backend, - capacities.collisions_capacity * num_batches, - storage, - ) - .unwrap(); - let new_constraint_builders = Tensor::vector_uninit( - backend, - capacities.collisions_capacity * num_batches, - storage, - ) - .unwrap(); - let constraints_colors = Tensor::vector_uninit( - backend, - capacities.collisions_capacity * num_batches, - storage, - ) - .unwrap(); - let old_constraints_colors = Tensor::vector( - backend, - vec![0u32; (capacities.collisions_capacity * num_batches) as usize], - storage, - ) - .unwrap(); - let colored = Tensor::vector_uninit( - backend, - capacities.collisions_capacity * num_batches, - storage, - ) - .unwrap(); - let constraints_rands = Tensor::vector_uninit( - backend, - capacities.collisions_capacity * num_batches, - storage, - ) - .unwrap(); + let old_constraints = Tensor::vector_uninit(backend, contacts_capacity, storage).unwrap(); + let old_constraint_builders = + Tensor::vector_uninit(backend, contacts_capacity, storage).unwrap(); + let new_constraints = Tensor::vector_uninit(backend, contacts_capacity, storage).unwrap(); + let new_constraint_builders = + Tensor::vector_uninit(backend, contacts_capacity, storage).unwrap(); + let constraints_colors = + Tensor::vector_uninit(backend, contacts_capacity, storage).unwrap(); + let old_constraints_colors = + Tensor::vector(backend, vec![0u32; contacts_capacity as usize], storage).unwrap(); + let colored = Tensor::vector_uninit(backend, contacts_capacity, storage).unwrap(); + let constraints_rands = Tensor::vector_uninit(backend, contacts_capacity, storage).unwrap(); let color_buckets_stride = capacities.solver_colors + 3; let color_buckets = Tensor::vector_uninit(backend, color_buckets_stride * num_batches, storage).unwrap(); - let contact_offsets = Tensor::vector( - backend, - vec![0u32; num_batches as usize + 3], - storage, - ) - .unwrap(); - let pair_batch_counts = - Tensor::vector(backend, vec![0u32; num_batches as usize], storage).unwrap(); - let pfm_batch_counts = - Tensor::vector(backend, vec![0u32; num_batches as usize], storage).unwrap(); - let color_sorted_ids = Tensor::vector_uninit( - backend, - capacities.collisions_capacity * num_batches, - storage, - ) - .unwrap(); + let contact_plan = + Tensor::scalar(backend, ContactPlan::default(), storage | BufferUsages::UNIFORM) + .unwrap(); + let pfm_sort = PfmSortState::new(backend, pairs_capacity); + let color_sorted_ids = Tensor::vector_uninit(backend, contacts_capacity, storage).unwrap(); let old_constraints_counts = Tensor::vector( backend, vec![0u32; (num_colliders_per_batch as u32 * num_batches) as usize], @@ -767,18 +714,10 @@ impl RbdState { storage, ) .unwrap(); - let old_body_constraint_ids = Tensor::vector_uninit( - backend, - capacities.collisions_capacity * 2 * num_batches, - storage, - ) - .unwrap(); - let new_body_constraint_ids = Tensor::vector_uninit( - backend, - capacities.collisions_capacity * 2 * num_batches, - storage, - ) - .unwrap(); + let old_body_constraint_ids = + Tensor::vector_uninit(backend, contacts_capacity * 2, storage).unwrap(); + let new_body_constraint_ids = + Tensor::vector_uninit(backend, contacts_capacity * 2, storage).unwrap(); let lbvh_usages = if crate::VALIDATE_LBVH_TOPOLOGY { BufferUsages::STORAGE | BufferUsages::COPY_SRC @@ -786,8 +725,8 @@ impl RbdState { BufferUsages::STORAGE }; - let contacts_capacity_cpu = capacities.collisions_capacity * num_batches; - let collision_pairs_capacity_cpu = capacities.collisions_capacity * num_batches; + let contacts_capacity_cpu = contacts_capacity; + let collision_pairs_capacity_cpu = pairs_capacity; #[allow(unused_mut)] // Only mutated with the dim3 (multibody) feature. let mut bi = BatchIndices { num_batches, @@ -876,11 +815,11 @@ impl RbdState { mb_cons_demand_cpu: 0, batch_indices, contacts, - contacts_len, contacts_indirect, - contact_offsets, - pair_batch_counts, - pfm_batch_counts, + contact_plan, + pfm_sort, + #[cfg(feature = "dim3")] + mb_contact_index, mb_sweep_indirect, pfm_pairs, pfm_pairs_len, diff --git a/src_rbd/pipeline/rbd_step.rs b/src_rbd/pipeline/rbd_step.rs index 702b8fe8..30203418 100644 --- a/src_rbd/pipeline/rbd_step.rs +++ b/src_rbd/pipeline/rbd_step.rs @@ -148,8 +148,9 @@ impl RbdPipeline { collider_world_poses: &state.collider_world_poses, mprops: &state.mprops, contacts: &state.contacts, - contacts_len: &state.contacts_len, - contact_offsets: &state.contact_offsets, + contact_plan: &state.contact_plan, + contacts_indirect: &state.contacts_indirect, + mb_contact_index: &mut state.mb_contact_index, solver_vels: &mut state.solver_vels, batch_indices: &state.batch_indices, color_uniforms: &state.color_uniforms, @@ -305,8 +306,8 @@ impl RbdPipeline { let mut pass = encoder.begin_pass("[RBD] narrow-phase", timestamps.as_deref_mut()); self.narrow_phase.dispatch( + backend, &mut pass, - state.body_poses.len() as u32, &state.collider_world_poses, &state.shapes, &state.vertex_buffers, @@ -314,21 +315,20 @@ impl RbdPipeline { &state.collision_pairs, &mut state.collision_pairs_len, &mut state.contacts, - &mut state.contacts_len, &mut state.contacts_indirect, - &mut state.contact_offsets, - &mut state.pair_batch_counts, - &mut state.pfm_batch_counts, + &mut state.contact_plan, &mut state.mb_sweep_indirect, &mut state.pfm_pairs, &mut state.pfm_pairs_len, &mut state.pfm_pairs_indirect, + &mut state.pfm_sort, &state.batch_indices, &state.collider_parent, &state.collider_materials, &state.sim_params, self.contact_reduction, &state.collision_pairs_indirect, + state.collision_pairs_capacity_cpu, )?; drop(pass); @@ -346,8 +346,9 @@ impl RbdPipeline { // Solver preparation - create args here to avoid borrow conflicts let prepare_args = SolverArgs { contacts: &state.contacts, - contacts_len: &state.contacts_len, - contact_offsets: &state.contact_offsets, + contact_plan: &state.contact_plan, + #[cfg(feature = "dim3")] + mb_contact_index: &mut state.mb_contact_index, contacts_len_indirect: &state.contacts_indirect, constraints: &mut state.new_constraints, constraint_builders: &mut state.new_constraint_builders, @@ -392,7 +393,7 @@ impl RbdPipeline { } else { // Warmstart let warmstart_args = WarmstartArgs { - contact_offsets: &state.contact_offsets, + contact_plan: &state.contact_plan, old_body_constraint_counts: &state.old_constraints_counts, old_constraint_builders: &state.old_constraint_builders, old_body_constraint_ids: &state.old_body_constraint_ids, @@ -400,7 +401,6 @@ impl RbdPipeline { new_constraints: &mut state.new_constraints, new_constraint_builders: &state.new_constraint_builders, contacts_len_indirect: &state.contacts_indirect, - batch_indices: &state.batch_indices, }; self.warmstart @@ -416,7 +416,7 @@ impl RbdPipeline { curr_color: &mut state.curr_color, uncolored: &mut state.uncolored, uncolored_staging: &state.uncolored_staging, - contact_offsets: &state.contact_offsets, + contact_plan: &state.contact_plan, colored: &mut state.colored, batch_indices: &state.batch_indices, body_group: &state.body_group, @@ -428,7 +428,7 @@ impl RbdPipeline { // persist, so most constraints can reuse their old color and the // topo-gc iterations converge in 1-2 rounds instead of ~num_colors). let seed_args = crate::dynamics::warmstart::SeedColorsArgs { - contact_offsets: &state.contact_offsets, + contact_plan: &state.contact_plan, old_body_constraint_counts: &state.old_constraints_counts, old_body_constraint_ids: &state.old_body_constraint_ids, old_constraints: &state.old_constraints, @@ -437,7 +437,6 @@ impl RbdPipeline { constraints_colors: &mut state.constraints_colors, colored: &mut state.colored, contacts_len_indirect: &state.contacts_indirect, - batch_indices: &state.batch_indices, }; self.warmstart .seed_colors_from_warmstart(&mut pass, seed_args)?; @@ -452,7 +451,7 @@ impl RbdPipeline { curr_color: &mut state.curr_color, uncolored: &mut state.uncolored, uncolored_staging: &state.uncolored_staging, - contact_offsets: &state.contact_offsets, + contact_plan: &state.contact_plan, colored: &mut state.colored, batch_indices: &state.batch_indices, body_group: &state.body_group, @@ -469,7 +468,7 @@ impl RbdPipeline { contacts_len_indirect: &state.contacts_indirect, constraints_colors: &state.constraints_colors, constraints: &state.new_constraints, - contact_offsets: &state.contact_offsets, + contact_plan: &state.contact_plan, color_buckets: &mut state.color_buckets, color_sorted_ids: &mut state.color_sorted_ids, batch_indices: &state.batch_indices, @@ -498,8 +497,9 @@ impl RbdPipeline { // Create solver_args for solve phase (after coloring is complete) let solver_args = SolverArgs { contacts: &state.contacts, - contacts_len: &state.contacts_len, - contact_offsets: &state.contact_offsets, + contact_plan: &state.contact_plan, + #[cfg(feature = "dim3")] + mb_contact_index: &mut state.mb_contact_index, contacts_len_indirect: &state.contacts_indirect, constraints: &mut state.new_constraints, constraint_builders: &mut state.new_constraint_builders, @@ -638,20 +638,6 @@ impl RbdPipeline { RbdResizePolicy::Fit => safe_total >= total_capacity || total_capacity >= new_total, }; - let contact_demand = counts[0].saturating_add(counts[1]); - let contacts_capacity = state.contacts_capacity_cpu; - let safe_contacts = contact_demand.saturating_add(contact_demand / 4); - let new_contacts = contact_demand - .saturating_add(contact_demand / 2) - .max(state.capacities.collisions_capacity.saturating_mul(nb)); - let resize_contacts = match state.capacities.collisions_resize_policy { - RbdResizePolicy::Fixed => false, - RbdResizePolicy::Grow => safe_contacts >= contacts_capacity, - RbdResizePolicy::Fit => { - safe_contacts >= contacts_capacity || contacts_capacity >= new_contacts - } - }; - #[cfg(feature = "dim3")] let (resize_mb, new_mb) = { let mb_demand = counts[3]; @@ -677,15 +663,13 @@ impl RbdPipeline { #[cfg(not(feature = "dim3"))] let resize_mb = false; - if grow_colors || resize_pairs || resize_contacts || resize_mb { + if grow_colors || resize_pairs || resize_mb { backend.synchronize()?; } if grow_colors { state.max_colors += 5; - // The color-bucket buffers are strided by `max_colors + 3`: - // regrow them and update the stride in `BatchIndices`. let storage: BufferUsages = BufferUsages::STORAGE | BufferUsages::COPY_SRC; let stride = state.max_colors + 3; let nb = state.num_batches; @@ -698,11 +682,15 @@ impl RbdPipeline { if resize_pairs { state.collision_pairs = Tensor::vector_uninit(backend, new_total, storage)?; state.pfm_pairs = Tensor::vector_uninit(backend, new_total, storage)?; + state.pfm_sort.resize(backend, new_total); state.collision_pairs_capacity_cpu = new_total; - } - if resize_contacts { + let new_contacts = new_total * 2; state.contacts = Tensor::vector_uninit(backend, new_contacts, storage)?; + #[cfg(feature = "dim3")] + { + state.mb_contact_index = Tensor::vector_uninit(backend, new_contacts, storage)?; + } state.old_constraints = Tensor::vector_uninit(backend, new_contacts, storage)?; state.old_constraint_builders = Tensor::vector_uninit(backend, new_contacts, storage)?; @@ -731,7 +719,7 @@ impl RbdPipeline { if resize_mb { state.multibodies.resize_contact_slabs(backend, new_mb); } - if resize_pairs || resize_contacts || resize_mb { + if resize_pairs || resize_mb { state.rebuild_batch_indices(backend); } } diff --git a/src_rbd/pipeline/test_batched_stacks.rs b/src_rbd/pipeline/test_batched_stacks.rs index f71d1c23..b118e70b 100644 --- a/src_rbd/pipeline/test_batched_stacks.rs +++ b/src_rbd/pipeline/test_batched_stacks.rs @@ -14,26 +14,35 @@ async fn test_backend() -> GpuBackend { GpuBackend::WebGpu(khal::backend::WebGpu::default().await.unwrap()) } } + fn build_env(num_stacks: usize, stack_height: 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(50.0, 0.5, 50.0), ground, &mut bodies); + colliders.insert_with_parent( + ColliderBuilder::cuboid(50.0, 0.5, 50.0), + ground, + &mut bodies, + ); for s in 0..num_stacks { let x = (s % 8) as f32 * 2.0; let z = (s / 8) as f32 * 2.0; for i in 0..stack_height { let y = 0.25 + i as f32 * 0.45; - let handle = - bodies.insert(RigidBodyBuilder::dynamic().translation(Vec3::new(x, y, z))); - colliders.insert_with_parent(ColliderBuilder::cuboid(0.2, 0.2, 0.2), handle, &mut bodies); + let handle = bodies.insert(RigidBodyBuilder::dynamic().translation(Vec3::new(x, y, z))); + colliders.insert_with_parent( + ColliderBuilder::cuboid(0.2, 0.2, 0.2), + handle, + &mut bodies, + ); } } (bodies, colliders) } + async fn run_case(num_envs: u32, num_stacks: usize, stack_height: usize, collisions_capacity: u32) { let backend = test_backend().await; @@ -137,8 +146,7 @@ async fn test_stacks_8_impulse_joint() { let mut joints = ImpulseJointSet::new(); let anchor = bodies.insert(RigidBodyBuilder::fixed().translation(Vec3::new(0.0, 2.0, 0.0))); colliders.insert_with_parent(ColliderBuilder::ball(0.1), anchor, &mut bodies); - let bob = - bodies.insert(RigidBodyBuilder::dynamic().translation(Vec3::new(1.0, 2.0, 0.0))); + let bob = bodies.insert(RigidBodyBuilder::dynamic().translation(Vec3::new(1.0, 2.0, 0.0))); colliders.insert_with_parent(ColliderBuilder::ball(0.1), bob, &mut bodies); let joint = RevoluteJointBuilder::new(Vec3::Z) .local_anchor1(Vec3::ZERO) @@ -251,17 +259,25 @@ async fn test_stacks_7_pfm_shapes() { 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(20.0, 0.5, 20.0), ground, &mut bodies); + colliders.insert_with_parent( + ColliderBuilder::cuboid(20.0, 0.5, 20.0), + ground, + &mut bodies, + ); for i in 0..3 { - let b = bodies.insert( - RigidBodyBuilder::dynamic().translation(Vec3::new(i as f32 * 1.5, 0.6, 0.0)), - ); + let b = bodies.insert(RigidBodyBuilder::dynamic().translation(Vec3::new( + i as f32 * 1.5, + 0.6, + 0.0, + ))); colliders.insert_with_parent(ColliderBuilder::capsule_y(0.15, 0.1), b, &mut bodies); } for i in 0..2 { - let b = bodies.insert( - RigidBodyBuilder::dynamic().translation(Vec3::new(i as f32 * 1.5 + 0.5, 0.6, 1.5)), - ); + let b = bodies.insert(RigidBodyBuilder::dynamic().translation(Vec3::new( + i as f32 * 1.5 + 0.5, + 0.6, + 1.5, + ))); colliders.insert_with_parent(ColliderBuilder::ball(0.2), b, &mut bodies); } (bodies, colliders) @@ -306,13 +322,19 @@ async fn test_stacks_7_pfm_shapes() { } println!("OK: pfm shapes rest, envs={num_envs}"); } + fn build_mb_env(ball_colliders: bool) -> (RigidBodySet, ColliderSet, MultibodyJointSet) { let mut bodies = RigidBodySet::new(); let mut colliders = ColliderSet::new(); let mut mb_joints = MultibodyJointSet::new(); let ground = bodies.insert(RigidBodyBuilder::fixed().translation(Vec3::new(0.0, -0.5, 0.0))); - colliders.insert_with_parent(ColliderBuilder::cuboid(20.0, 0.5, 20.0), ground, &mut bodies); + colliders.insert_with_parent( + ColliderBuilder::cuboid(20.0, 0.5, 20.0), + ground, + &mut bodies, + ); + let mut prev = None; for i in 0..3 { let x = i as f32 * 0.5; @@ -373,7 +395,10 @@ async fn test_stacks_6_multibody() { .await .unwrap(); let finite = poses.iter().all(|p| p.translation.is_finite()); - println!("step {step}: cap={} demand={demand} layout={layout:?} finite={finite}", state.multibodies().contact_constraints_capacity()); + println!( + "step {step}: cap={} demand={demand} layout={layout:?} finite={finite}", + state.multibodies().contact_constraints_capacity() + ); if demand > 0 { let cons: Vec = backend .slow_read_vec(state.multibodies().contact_constraints().buffer()) @@ -488,14 +513,14 @@ async fn test_stacks_10_mb_impulse_joint() { let mut joints = ImpulseJointSet::new(); let mut mb_joints = MultibodyJointSet::new(); let y = anchor_y(e); + let anchor = bodies.insert(RigidBodyBuilder::fixed().translation(Vec3::new(0.0, y, 0.0))); colliders.insert_with_parent(ColliderBuilder::ball(0.05), anchor, &mut bodies); let mut links = Vec::new(); let mut prev = anchor; for i in 0..2 { let x = 0.5 + i as f32 * 0.5; - let link = - bodies.insert(RigidBodyBuilder::dynamic().translation(Vec3::new(x, y, 0.0))); + let link = bodies.insert(RigidBodyBuilder::dynamic().translation(Vec3::new(x, y, 0.0))); colliders.insert_with_parent(ColliderBuilder::ball(0.05), link, &mut bodies); let joint = RevoluteJointBuilder::new(Vec3::Z) .local_anchor1(Vec3::new(if i == 0 { 0.0 } else { 0.25 }, 0.0, 0.0)) @@ -555,7 +580,10 @@ async fn test_stacks_10_mb_impulse_joint() { ); for slot in 0..5 { let p = at(slot).translation; - assert!(p.is_finite(), "env {env} slot {slot}: non-finite pose {p:?}"); + assert!( + p.is_finite(), + "env {env} slot {slot}: non-finite pose {p:?}" + ); } for (link_slot, bob_slot) in [(1usize, 3usize), (2, 4)] { let link = at(link_slot); @@ -570,3 +598,94 @@ async fn test_stacks_10_mb_impulse_joint() { } println!("OK: multibody impulse joints hold, envs={num_envs}"); } + +#[futures_test::test] +#[serial_test::serial] +#[ignore] +async fn test_stacks_12_trimesh_reduction() { + let backend = test_backend().await; + + let num_envs = 4u32; + let build = || { + let mut bodies = RigidBodySet::new(); + let mut colliders = ColliderSet::new(); + + let ground = bodies.insert(RigidBodyBuilder::fixed()); + let nsubdivs = 8; + let heights = Array2::from_fn(nsubdivs + 1, nsubdivs + 1, |_, _| 0.0f32); + let (vertices, indices) = HeightField::new(heights, Vec3::new(8.0, 1.0, 8.0)).to_trimesh(); + colliders.insert_with_parent( + ColliderBuilder::trimesh_with_flags( + vertices, + indices, + TriMeshFlags::MERGE_DUPLICATE_VERTICES, + ) + .unwrap(), + ground, + &mut bodies, + ); + + let wide = bodies.insert(RigidBodyBuilder::dynamic().translation(Vec3::new(0.0, 0.3, 0.0))); + colliders.insert_with_parent(ColliderBuilder::cuboid(1.5, 0.2, 1.5), wide, &mut bodies); + + for i in 0..3 { + let b = bodies.insert(RigidBodyBuilder::dynamic().translation(Vec3::new( + 0.0, + 0.7 + i as f32 * 0.45, + 0.0, + ))); + colliders.insert_with_parent(ColliderBuilder::cuboid(0.2, 0.2, 0.2), b, &mut bodies); + } + (bodies, colliders) + }; + let envs: Vec<_> = (0..num_envs).map(|_| build()).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: 64, + ..Default::default() + }; + let mut state = RbdState::from_rapier(&backend, &refs, capacities); + let mut pipeline = RbdPipeline::new(&backend).unwrap(); + pipeline.contact_reduction = true; + for _ in 0..250 { + pipeline.step(&backend, &mut state, None).unwrap(); + pipeline.auto_resize_buffers(&backend, &mut state).unwrap(); + } + backend.synchronize().unwrap(); + + let poses: Vec = backend + .slow_read_vec(state.body_poses().buffer()) + .await + .unwrap(); + let nb = num_envs as usize; + for env in 0..num_envs as usize { + let expected = [0.2f32, 0.6, 1.0, 1.4]; + for (b, expected_y) in expected.iter().enumerate() { + let pose = poses[(1 + b) * nb + env]; + assert!( + pose.translation.is_finite(), + "env {env} box {b}: non-finite pose {:?}", + pose.translation + ); + let y = pose.translation.y; + assert!( + (y - expected_y).abs() < 0.1, + "env {env} box {b}: y = {y}, expected ~{expected_y}" + ); + if env > 0 { + let ref_pose = poses[(1 + b) * nb]; + let d = (pose.translation - ref_pose.translation).length(); + assert!(d < 5.0e-2, "env {env} box {b}: diverged from env 0 by {d}"); + } + } + } + println!("OK: trimesh floor + contact reduction rests, envs={num_envs}"); +} diff --git a/src_rbd/utils/radix_sort/mod.rs b/src_rbd/utils/radix_sort/mod.rs index 667520ed..c39cae12 100644 --- a/src_rbd/utils/radix_sort/mod.rs +++ b/src_rbd/utils/radix_sort/mod.rs @@ -186,7 +186,7 @@ impl RadixSort { // Sort by permutation: sort indices by their corresponding key. indices.clear(); indices.extend(0..n as u32); - indices.sort_unstable_by_key(|&i| keys[i as usize]); + indices.sort_by_key(|&i| keys[i as usize]); // Scatter using the sorted permutation. for (dst, &src) in indices.iter().enumerate() { diff --git a/src_rbd_shaders/broad_phase/lbvh.rs b/src_rbd_shaders/broad_phase/lbvh.rs index 9647282b..9ac3af5d 100644 --- a/src_rbd_shaders/broad_phase/lbvh.rs +++ b/src_rbd_shaders/broad_phase/lbvh.rs @@ -76,44 +76,6 @@ pub fn gpu_flat_list_dispatch( } } -/// Number of lanes used by the per-batch-count max reductions below. Their -/// host dispatch is a single workgroup: `.call(pass, MAX_REDUCE_LANES, ...)`. -pub const MAX_REDUCE_LANES: u32 = 256; - -/// Workgroup-parallel `max` over the per-batch counts, then writes the -/// `[ceil(max/64), num_batches, 1]` indirect grid. -/// -/// NOTE: `lens` is mutable even though we don't modify it: the loads must be -/// atomic or they occasionally read stale data (breaks Windows+Nvidia+wgpu, see -/// ). -#[inline(always)] -pub(crate) fn reduce_max_lens( - lane: u32, - lens: &mut [u32], - partial: &mut [u32; MAX_REDUCE_LANES as usize], -) { - let num_batches = lens.len(); - - let mut m = 0u32; - for i in StepRng::new(lane..num_batches as u32, MAX_REDUCE_LANES) { - m = m.max(atomic_load_u32(lens.at_mut(i as usize))); - } - partial.write(lane as usize, m); - workgroup_memory_barrier_with_group_sync(); - - // Tree reduction over the 256 lanes (8 halving steps). - for step in 0..8u32 { - let stride = MAX_REDUCE_LANES >> (step + 1); - if lane < stride { - let v = partial - .read(lane as usize) - .max(partial.read((lane + stride) as usize)); - partial.write(lane as usize, v); - } - workgroup_memory_barrier_with_group_sync(); - } -} - /// Runs a reduction to compute the AABB of the collider positions. /// Needs to be called with a single workgroup. #[spirv_bindgen] @@ -131,10 +93,7 @@ pub fn gpu_lbvh_compute_domain( *workspace_mins.at_mut(thread_id as usize) = Vector::splat(MAX_FLT); *workspace_maxs.at_mut(thread_id as usize) = Vector::splat(-MAX_FLT); - for i in StepRng::new( - thread_id..batch_ids.colliders_len, - REDUCTION_WORKGROUP_SIZE, - ) { + for i in StepRng::new(thread_id..batch_ids.colliders_len, REDUCTION_WORKGROUP_SIZE) { let val_i = poses.at(batch_ids.body_global(batch_id, i)).translation; *workspace_mins.at_mut(thread_id as usize) = workspace_mins.at(thread_id as usize).min(val_i); diff --git a/src_rbd_shaders/broad_phase/narrow_phase.rs b/src_rbd_shaders/broad_phase/narrow_phase.rs index c274df65..28def5f6 100644 --- a/src_rbd_shaders/broad_phase/narrow_phase.rs +++ b/src_rbd_shaders/broad_phase/narrow_phase.rs @@ -22,146 +22,84 @@ use khal_std::{ sync::{atomic_add_u32, atomic_load_u32}, }; -use super::lbvh::{MAX_REDUCE_LANES, reduce_max_lens}; use crate::broad_phase::CollisionPair; use crate::utils::{BatchIndices, SliceMut}; use glamx::UVec2; const WORKGROUP_SIZE: u32 = 64; -/// Resets the contacts counter. One thread per batch. +#[derive(Copy, Clone, Default)] +#[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] +#[repr(C)] +pub struct ContactPlan { + pub bound: u32, + pub pfm_base: u32, + pub pfm_len: u32, +} + #[spirv_bindgen] #[spirv(compute(threads(64)))] pub fn gpu_reset_narrow_phase( #[spirv(global_invocation_id)] invocation_id: UVec3, - #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] contacts_len: &mut [u32], - #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] pfm_pairs_len: &mut [u32], - #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] pair_batch_counts: &mut [u32], - #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] pfm_batch_counts: &mut [u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] pfm_pairs_len: &mut [u32], ) { let i = invocation_id.x as usize; - if i < contacts_len.len() { - contacts_len.write(i, 0); - } if i < pfm_pairs_len.len() { pfm_pairs_len.write(i, 0); } - if i < pair_batch_counts.len() { - pair_batch_counts.write(i, 0); - } - if i < pfm_batch_counts.len() { - pfm_batch_counts.write(i, 0); - } } #[spirv_bindgen] -#[spirv(compute(threads(64)))] -pub fn gpu_count_pairs_per_batch( - #[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: &mut [u32], - #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] pair_batch_counts: &mut [u32], - #[spirv(uniform, descriptor_set = 0, binding = 3)] batch_ids: &BatchIndices, +#[spirv(compute(threads(1)))] +pub fn gpu_contact_plan( + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] collision_pairs_len: &mut [u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] pfm_pairs_len: &mut [u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] contact_plan: &mut ContactPlan, + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] pfm_sort_len: &mut [u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] contacts_indirect: &mut [u32; 3], + #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] pfm_indirect: &mut [u32; 3], + #[spirv(storage_buffer, descriptor_set = 0, binding = 6)] mb_sweep_indirect: &mut [u32; 3], + #[spirv(uniform, descriptor_set = 0, binding = 7)] batch_ids: &BatchIndices, ) { - let num_threads = num_workgroups.x * WORKGROUP_SIZE; - let total = + let pairs = atomic_load_u32(collision_pairs_len.at_mut(0)).min(batch_ids.collision_pairs_capacity); - for t in StepRng::new(invocation_id.x..total, num_threads) { - let pair = collision_pairs.read(t as usize); - let batch_id = batch_ids.collider_batch(pair.colliders.x); - atomic_add_u32(pair_batch_counts.at_mut(batch_id as usize), 1); - } + let pfm = atomic_load_u32(pfm_pairs_len.at_mut(0)).min(batch_ids.collision_pairs_capacity); + let bound = pairs + pfm; + + contact_plan.bound = bound; + contact_plan.pfm_base = pairs; + contact_plan.pfm_len = pfm; + pfm_sort_len.write(0, pfm); + + *contacts_indirect.at_mut(0) = bound.div_ceil(WORKGROUP_SIZE); + *contacts_indirect.at_mut(1) = 1; + *contacts_indirect.at_mut(2) = 1; + *pfm_indirect.at_mut(0) = pfm.div_ceil(WORKGROUP_SIZE); + *pfm_indirect.at_mut(1) = 1; + *pfm_indirect.at_mut(2) = 1; + + *mb_sweep_indirect.at_mut(0) = if bound > 0 { + batch_ids.multibodies_batch_capacity + } else { + 0 + }; + *mb_sweep_indirect.at_mut(1) = batch_ids.num_batches; + *mb_sweep_indirect.at_mut(2) = 1; } + #[spirv_bindgen] #[spirv(compute(threads(64)))] -pub fn gpu_count_pfm_per_batch( +pub fn gpu_pfm_sort_keys( #[spirv(global_invocation_id)] invocation_id: UVec3, #[spirv(num_workgroups)] num_workgroups: UVec3, #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] pfm_pairs: &[NarrowPhasePfmPair], - #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] pfm_pairs_len: &mut [u32], - #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] pfm_batch_counts: &mut [u32], - #[spirv(uniform, descriptor_set = 0, binding = 3)] batch_ids: &BatchIndices, -) { - let num_threads = num_workgroups.x * WORKGROUP_SIZE; - let total = atomic_load_u32(pfm_pairs_len.at_mut(0)).min(batch_ids.collision_pairs_capacity); - for t in StepRng::new(invocation_id.x..total, num_threads) { - let pair = pfm_pairs.read(t as usize); - let batch_id = batch_ids.collider_batch(pair.colliders.x); - atomic_add_u32(pfm_batch_counts.at_mut(batch_id as usize), 1); - } -} -/// -#[spirv_bindgen] -#[spirv(compute(threads(1)))] -pub fn gpu_contact_offsets_scan( - #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] pair_batch_counts: &mut [u32], - #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] pfm_batch_counts: &mut [u32], - #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] collision_pairs_len: &mut [u32], - #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] pfm_pairs_len: &mut [u32], - #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] contact_offsets: &mut [u32], - #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] indirect_args: &mut [u32; 3], - #[spirv(uniform, descriptor_set = 0, binding = 6)] batch_ids: &BatchIndices, -) { - let num_batches = batch_ids.num_batches as usize; - let capacity = batch_ids.contacts_capacity; - let mut total = 0u32; - for b in 0..num_batches { - contact_offsets.write(b, total); - let bound = atomic_load_u32(pair_batch_counts.at_mut(b)) - + atomic_load_u32(pfm_batch_counts.at_mut(b)); - total = (total + bound).min(capacity); - } - contact_offsets.write(num_batches, total); - contact_offsets.write( - num_batches + 1, - atomic_load_u32(collision_pairs_len.at_mut(0)).min(batch_ids.collision_pairs_capacity), - ); - contact_offsets.write( - num_batches + 2, - atomic_load_u32(pfm_pairs_len.at_mut(0)).min(batch_ids.collision_pairs_capacity), - ); - *indirect_args.at_mut(0) = total.div_ceil(WORKGROUP_SIZE); - *indirect_args.at_mut(1) = 1; - *indirect_args.at_mut(2) = 1; -} -#[spirv_bindgen] -#[spirv(compute(threads(64)))] -pub fn gpu_zero_contact_lens( - #[spirv(global_invocation_id)] invocation_id: UVec3, - #[spirv(num_workgroups)] num_workgroups: UVec3, - #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] contacts: &mut [IndexedManifold], - #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] contact_offsets: &[u32], - #[spirv(uniform, descriptor_set = 0, binding = 2)] batch_ids: &BatchIndices, + #[spirv(uniform, descriptor_set = 0, binding = 1)] contact_plan: &ContactPlan, + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] sort_keys: &mut [u32], ) { let num_threads = num_workgroups.x * WORKGROUP_SIZE; - let total = contact_offsets.read(batch_ids.num_batches as usize); + let total = contact_plan.pfm_len; for t in StepRng::new(invocation_id.x..total, num_threads) { - contacts.at_mut(t as usize).contact.len = 0; - } -} -/// contact-constraint dispatches (`[multibodies_batch_capacity, num_batches, -/// 1]`). -#[spirv_bindgen] -#[spirv(compute(threads(256)))] -pub fn gpu_narrow_phase_init_contacts_dispatch( - #[spirv(local_invocation_id)] lid: UVec3, - #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] contacts_len: &mut [u32], - #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] mb_sweep_indirect: &mut [u32; 3], - #[spirv(uniform, descriptor_set = 0, binding = 2)] batch_ids: &BatchIndices, - #[spirv(workgroup)] partial: &mut [u32; MAX_REDUCE_LANES as usize], -) { - reduce_max_lens(lid.x, contacts_len, partial); - // `partial[0]` holds the max after the reduction (all lanes synced). - if lid.x == 0 { - let any_contacts = partial.read(0) > 0; - *mb_sweep_indirect.at_mut(0) = if any_contacts { - batch_ids.multibodies_batch_capacity - } else { - 0 - }; - *mb_sweep_indirect.at_mut(1) = batch_ids.num_batches; - *mb_sweep_indirect.at_mut(2) = 1; + sort_keys.write(t as usize, pfm_pairs.read(t as usize).pair_index); } } @@ -215,106 +153,110 @@ fn pool_dedup(cand: &mut [ContactPoint; 8], num: &mut usize, pt: ContactPoint, d /// 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)))] +#[spirv(compute(threads(64)))] pub fn gpu_reduce_contacts( - #[spirv(workgroup_id)] workgroup_id: UVec3, + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[spirv(num_workgroups)] num_workgroups: 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(storage_buffer, descriptor_set = 0, binding = 2)] contact_offsets: &[u32], - #[allow(unused_variables)] - #[spirv(uniform, descriptor_set = 0, binding = 3)] - batch_ids: &BatchIndices, - #[spirv(uniform, descriptor_set = 0, binding = 4)] params: &RbdSimParams, + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] sorted_pfm_keys: &[u32], + #[spirv(uniform, descriptor_set = 0, binding = 2)] contact_plan: &ContactPlan, + #[spirv(uniform, descriptor_set = 0, binding = 3)] params: &RbdSimParams, ) { let prediction = params.prediction_distance(); let merge_cos = params.contact_merge_cos; - let batch_id = workgroup_id.y; - let seg_start = contact_offsets.read(batch_id as usize) as usize; - let seg_end = contact_offsets.read(batch_id as usize + 1) as usize; - let mut contacts = SliceMut(contacts, seg_start); - let n = (contacts_len.read(batch_id as usize) as usize).min(seg_end - seg_start); - - // Write cursor: always <= the read cursor, so compacting in place is safe. - let mut w = 0usize; - for i in 0..n { - let im = contacts[i]; - let mut merged = false; - for j in 0..w { - let out = contacts[j]; - if out.colliders.x == im.colliders.x - && out.colliders.y == im.colliders.y - && out.contact.normal_a.dot(im.contact.normal_a) >= merge_cos - { - // Pool the two manifolds' points (same collider-A local frame), - // dropping near-duplicates as rapier's clustering does. - let na = (out.contact.len as usize).min(MAX_MANIFOLD_POINTS); - let nb = (im.contact.len as usize).min(MAX_MANIFOLD_POINTS); - let dedup_eps = prediction * 0.25; - let dedup_eps_sq = dedup_eps * dedup_eps; - let mut cand = [ContactPoint::default(); 8]; - let mut num = 0usize; - for k in 0..na { - pool_dedup( - &mut cand, - &mut num, - out.contact.points_a.read(k), - dedup_eps_sq, - ); - } - for k in 0..nb { - pool_dedup( - &mut cand, - &mut num, - im.contact.points_a.read(k), - dedup_eps_sq, - ); - } - // Normal of whichever manifold holds the deepest point. rapier - // keeps the opener's normal instead, which it can afford - // because its ~5.1 degree cone makes every member equivalent; - // this degrades gracefully when `merge_cos` is loosened, and - // agrees with rapier's choice when it is not. - let mut deep_out = out.contact.points_a.at(0).dist; - for k in 1..na { - let d = out.contact.points_a.at(k).dist; - if d < deep_out { - deep_out = d; + let num_threads = num_workgroups.x * WORKGROUP_SIZE; + let total = contact_plan.pfm_len as usize; + let base = contact_plan.pfm_base as usize; + + for t in StepRng::new(invocation_id.x..total as u32, num_threads) { + let i = t as usize; + let key = sorted_pfm_keys.read(i); + if i > 0 && sorted_pfm_keys.read(i - 1) == key { + continue; + } + let mut n = 1usize; + for j in (i + 1)..total { + if sorted_pfm_keys.read(j) != key { + break; + } + n += 1; + } + if n <= 1 { + continue; + } + let mut contacts = SliceMut(contacts, base + i); + + let mut w = 0usize; + for i in 0..n { + let im = contacts[i]; + if im.contact.len == 0 { + continue; + } + let mut merged = false; + for j in 0..w { + let out = contacts[j]; + if out.contact.normal_a.dot(im.contact.normal_a) >= merge_cos { + let na = (out.contact.len as usize).min(MAX_MANIFOLD_POINTS); + let nb = (im.contact.len as usize).min(MAX_MANIFOLD_POINTS); + let dedup_eps = prediction * 0.25; + let dedup_eps_sq = dedup_eps * dedup_eps; + let mut cand = [ContactPoint::default(); 8]; + let mut num = 0usize; + for k in 0..na { + pool_dedup( + &mut cand, + &mut num, + out.contact.points_a.read(k), + dedup_eps_sq, + ); } - } - 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; + for k in 0..nb { + pool_dedup( + &mut cand, + &mut num, + im.contact.points_a.read(k), + dedup_eps_sq, + ); + } + let mut deep_out = out.contact.points_a.at(0).dist; + for k in 1..na { + let d = out.contact.points_a.at(k).dist; + if d < deep_out { + deep_out = d; + } } + let mut deep_in = im.contact.points_a.at(0).dist; + for k in 1..nb { + let d = im.contact.points_a.at(k).dist; + if d < deep_in { + deep_in = d; + } + } + let normal = if deep_in < deep_out { + im.contact.normal_a + } else { + out.contact.normal_a + }; + let mut reduced = manifold_reduction(&cand, num as u32, normal, prediction); + reduced.normal_a = normal; + let mut kept = out; + kept.contact = reduced; + contacts[j] = kept; + merged = true; + break; } - let normal = if deep_in < deep_out { - im.contact.normal_a - } else { - out.contact.normal_a - }; - let mut reduced = manifold_reduction(&cand, num as u32, normal, prediction); - // `manifold_reduction` fills points/len only. - reduced.normal_a = normal; - let mut kept = out; - kept.contact = reduced; - contacts[j] = kept; - merged = true; - break; + } + if !merged { + contacts[w] = im; + w += 1; } } - if !merged { - contacts[w] = im; - w += 1; + for i in w..n { + contacts[i].contact.len = 0; } } - for i in w..n { - contacts[i].contact.len = 0; - } - contacts_len.write(batch_id as usize, w as u32); } /// Narrow phase, pass 1 of 2: analytic shape-shape contacts for ball / cuboid @@ -328,102 +270,89 @@ 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)] contact_offsets: &[u32], + #[spirv(uniform, descriptor_set = 0, binding = 1)] contact_plan: &ContactPlan, #[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], - #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] contacts_len: &mut [u32], - #[spirv(uniform, descriptor_set = 0, binding = 6)] batch_ids: &BatchIndices, // Per-collider parent body id, used to resolve `IndexedManifold::bodies` here, // at the last moment before the solver consumes it (instead of carrying the // body ids all the way through the broad-phase collision-pair buffer). - #[spirv(storage_buffer, descriptor_set = 0, binding = 7)] collider_parent: &[u32], - #[spirv(storage_buffer, descriptor_set = 0, binding = 8)] + #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] collider_parent: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 6)] collider_materials: &[ColliderMaterial], - #[spirv(uniform, descriptor_set = 0, binding = 9)] params: &RbdSimParams, + #[spirv(uniform, descriptor_set = 0, binding = 7)] params: &RbdSimParams, ) { let prediction = params.prediction_distance(); let num_threads = num_workgroups.x * WORKGROUP_SIZE; - let total = contact_offsets.read(batch_ids.num_batches as usize + 1); + let total = contact_plan.pfm_base; for t in StepRng::new(invocation_id.x..total, num_threads) { let pair = collision_pairs.read(t as usize); - let batch_id = batch_ids.collider_batch(pair.colliders.x); - let seg_start = contact_offsets.read(batch_id as usize); - let seg_end = contact_offsets.read(batch_id as usize + 1); - let contacts_len = contacts_len.at_mut(batch_id 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 let body1 = collider_parent.read(pair.colliders.x as usize); let body2 = collider_parent.read(pair.colliders.y as usize); - if body1 == body2 { - continue; - } - let pose1 = poses.read(pair.colliders.x as usize); - let pose2 = poses.read(pair.colliders.y as usize); - let shape1 = shapes.at(pair.colliders.x as usize); - let shape2 = shapes.at(pair.colliders.y as usize); - let shape_ty1 = shape1.shape_type(); - let shape_ty2 = shape2.shape_type(); let mut manifold = ContactManifold::default(); - let pose12 = pose1.inverse() * pose2; + if body1 != body2 { + let pose1 = poses.read(pair.colliders.x as usize); + let pose2 = poses.read(pair.colliders.y as usize); + let shape1 = shapes.at(pair.colliders.x as usize); + let shape2 = shapes.at(pair.colliders.y as usize); + let shape_ty1 = shape1.shape_type(); + let shape_ty2 = shape2.shape_type(); + let pose12 = pose1.inverse() * pose2; + + if shape_ty1 == SHAPE_TYPE_BALL { + if shape_ty2 == SHAPE_TYPE_BALL { + let ball1 = shape1.to_ball(); + let ball2 = shape2.to_ball(); + manifold = ball_ball(pose12, &ball1, &ball2); + } else if shape_ty2 == SHAPE_TYPE_CUBOID + || shape_ty2 == SHAPE_TYPE_CAPSULE + || shape_ty2 == SHAPE_TYPE_CONE + || shape_ty2 == SHAPE_TYPE_CYLINDER + { + let ball1 = shape1.to_ball(); + manifold = ball_convex(pose12, &ball1, shape2); + } + } - // Ball - Convex - if shape_ty1 == SHAPE_TYPE_BALL { - if shape_ty2 == SHAPE_TYPE_BALL { - let ball1 = shape1.to_ball(); - let ball2 = shape2.to_ball(); - manifold = ball_ball(pose12, &ball1, &ball2); - } else if shape_ty2 == SHAPE_TYPE_CUBOID - || shape_ty2 == SHAPE_TYPE_CAPSULE - || shape_ty2 == SHAPE_TYPE_CONE - || shape_ty2 == SHAPE_TYPE_CYLINDER + if shape_ty2 == SHAPE_TYPE_BALL + && (shape_ty1 == SHAPE_TYPE_CUBOID + || shape_ty1 == SHAPE_TYPE_CAPSULE + || shape_ty1 == SHAPE_TYPE_CONE + || shape_ty1 == SHAPE_TYPE_CYLINDER) { - let ball1 = shape1.to_ball(); - manifold = ball_convex(pose12, &ball1, shape2); + let ball2 = shape2.to_ball(); + manifold = convex_ball(pose12, shape1, &ball2); } - } - - // Convex - Ball - if shape_ty2 == SHAPE_TYPE_BALL - && (shape_ty1 == SHAPE_TYPE_CUBOID - || shape_ty1 == SHAPE_TYPE_CAPSULE - || shape_ty1 == SHAPE_TYPE_CONE - || shape_ty1 == SHAPE_TYPE_CYLINDER) - { - let ball2 = shape2.to_ball(); - manifold = convex_ball(pose12, shape1, &ball2); - } - // Cuboid - Cuboid - 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); + 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); + } } // 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 { - let idx = seg_start + atomic_add_u32(contacts_len, 1); - - if idx < seg_end { - let mat1 = collider_materials.read(pair.colliders.x as usize); - let mat2 = collider_materials.read(pair.colliders.y as usize); - contacts.write( - idx as usize, - IndexedManifold { - contact: manifold, - colliders: pair.colliders, - bodies: UVec2::new(body1, body2), - friction: mat1.combined_friction(&mat2), - restitution: mat1.combined_restitution(&mat2), - _padding: [0.0; 2], - }, - ); - } + let mat1 = collider_materials.read(pair.colliders.x as usize); + let mat2 = collider_materials.read(pair.colliders.y as usize); + contacts.write( + t as usize, + IndexedManifold { + contact: manifold, + colliders: pair.colliders, + bodies: UVec2::new(body1, body2), + friction: mat1.combined_friction(&mat2), + restitution: mat1.combined_restitution(&mat2), + _padding: [0.0; 2], + }, + ); + } else { + contacts.at_mut(t as usize).contact.len = 0; } } } @@ -457,7 +386,8 @@ pub fn gpu_narrow_phase_shape_shape_deferred( let num_threads = num_workgroups.x * WORKGROUP_SIZE; let pfm_capacity = batch_ids.collision_pairs_capacity as usize; - let total = atomic_load_u32(collision_pairs_len.at_mut(0)).min(batch_ids.collision_pairs_capacity); + let total = + atomic_load_u32(collision_pairs_len.at_mut(0)).min(batch_ids.collision_pairs_capacity); // 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 @@ -519,6 +449,8 @@ pub fn gpu_narrow_phase_shape_shape_deferred( thickness1: sub1.thickness, thickness2: sub2.thickness, colliders: pair.colliders, + pair_index: t, + _padding: [0; 3], }; let pfm_index = atomic_add_u32(pfm_pairs_len, 1); // NOTE: if we exceed capacity, just skip the pair. @@ -542,6 +474,7 @@ pub fn gpu_narrow_phase_shape_shape_deferred( &mesh, convex, pair.colliders, + t, &mut pfm_pairs, pfm_pairs_len, pfm_capacity, @@ -561,6 +494,7 @@ pub fn gpu_narrow_phase_shape_shape_deferred( &mesh, convex, UVec2::new(pair.colliders.y, pair.colliders.x), + t, &mut pfm_pairs, pfm_pairs_len, pfm_capacity, @@ -581,6 +515,7 @@ pub fn gpu_narrow_phase_shape_shape_deferred( &pline, convex, pair.colliders, + t, &mut pfm_pairs, pfm_pairs_len, pfm_capacity, @@ -600,6 +535,7 @@ pub fn gpu_narrow_phase_shape_shape_deferred( &pline, convex, UVec2::new(pair.colliders.y, pair.colliders.x), + t, &mut pfm_pairs, pfm_pairs_len, pfm_capacity, @@ -618,6 +554,7 @@ fn trimesh_convex( mesh: &TriMesh, convex: &Shape, colliders: UVec2, + pair_index: u32, pfm_pairs: &mut SliceMut, pfm_pairs_len: &mut u32, pfm_pairs_capacity: usize, @@ -662,6 +599,8 @@ fn trimesh_convex( thickness1: sub1.thickness, thickness2: sub2.thickness, colliders, + pair_index, + _padding: [0; 3], }; let pfm_index = atomic_add_u32(pfm_pairs_len, 1); // Skip (don’t write) on overflow; the caller resizes and re-runs. @@ -689,6 +628,7 @@ fn polyline_convex( mesh: &Polyline, convex: &Shape, colliders: UVec2, + pair_index: u32, pfm_pairs: &mut SliceMut, pfm_pairs_len: &mut u32, pfm_pairs_capacity: usize, @@ -736,6 +676,8 @@ fn polyline_convex( thickness1: sub1.thickness, thickness2: sub2.thickness, colliders, + pair_index, + _padding: [0; 3], }; let pfm_index = atomic_add_u32(pfm_pairs_len, 1); // Skip (don’t write) on overflow; the caller resizes and re-runs. @@ -766,6 +708,8 @@ pub struct NarrowPhasePfmPair { thickness1: f32, thickness2: f32, colliders: UVec2, + pair_index: u32, + _padding: [u32; 3], } #[spirv_bindgen] @@ -774,32 +718,29 @@ pub fn gpu_narrow_phase_pfm_pfm( #[spirv(global_invocation_id)] invocation_id: UVec3, #[spirv(num_workgroups)] num_workgroups: 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(storage_buffer, descriptor_set = 0, binding = 2)] pfm_pairs: &[NarrowPhasePfmPair], - #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] contact_offsets: &[u32], - #[spirv(uniform, descriptor_set = 0, binding = 4)] batch_ids: &BatchIndices, - #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] vertices: &[PaddedVector], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] pfm_pairs: &[NarrowPhasePfmPair], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] pfm_order: &[u32], + #[spirv(uniform, descriptor_set = 0, binding = 3)] contact_plan: &ContactPlan, + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] vertices: &[PaddedVector], #[allow(unused_variables)] - #[spirv(storage_buffer, descriptor_set = 0, binding = 6)] + #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] indices: &[u32], // Per-collider parent body id, used to resolve `IndexedManifold::bodies` here // (see the note on `gpu_narrow_phase_shape_shape`). - #[spirv(storage_buffer, descriptor_set = 0, binding = 7)] collider_parent: &[u32], - #[spirv(storage_buffer, descriptor_set = 0, binding = 8)] + #[spirv(storage_buffer, descriptor_set = 0, binding = 6)] collider_parent: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 7)] collider_materials: &[ColliderMaterial], - #[spirv(uniform, descriptor_set = 0, binding = 9)] params: &RbdSimParams, + #[spirv(uniform, descriptor_set = 0, binding = 8)] params: &RbdSimParams, ) { let prediction = params.prediction_distance(); let num_threads = num_workgroups.x * WORKGROUP_SIZE; - let total = contact_offsets.read(batch_ids.num_batches as usize + 2); + let total = contact_plan.pfm_len; + let base = contact_plan.pfm_base; for t in StepRng::new(invocation_id.x..total, num_threads) { - let pair = pfm_pairs.read(t as usize); - let batch_id = batch_ids.collider_batch(pair.colliders.x); - let seg_start = contact_offsets.read(batch_id as usize); - let seg_end = contact_offsets.read(batch_id as usize + 1); - let contacts_len = contacts_len.at_mut(batch_id as usize); + let pair = pfm_pairs.read(pfm_order.read(t as usize) as usize); + let slot = (base + t) 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 @@ -807,39 +748,37 @@ pub fn gpu_narrow_phase_pfm_pfm( // 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); - if body1 == body2 { - continue; + let mut manifold = ContactManifold::default(); + if body1 != body2 { + manifold = pfm_pfm( + pair.pose12, + &pair.shape1, + pair.thickness1, + &pair.shape2, + pair.thickness2, + prediction, + vertices, + #[cfg(feature = "dim3")] + indices, + ); } - let manifold = pfm_pfm( - pair.pose12, - &pair.shape1, - pair.thickness1, - &pair.shape2, - pair.thickness2, - prediction, - vertices, - #[cfg(feature = "dim3")] - indices, - ); if manifold.len > 0 && manifold.points_a.at(0).dist < prediction { - let idx = seg_start + atomic_add_u32(contacts_len, 1); - - if idx < seg_end { - let mat1 = collider_materials.read(pair.colliders.x as usize); - let mat2 = collider_materials.read(pair.colliders.y as usize); - contacts.write( - idx as usize, - IndexedManifold { - contact: manifold, - colliders: pair.colliders, - bodies: UVec2::new(body1, body2), - friction: mat1.combined_friction(&mat2), - restitution: mat1.combined_restitution(&mat2), - _padding: [0.0; 2], - }, - ); - } + let mat1 = collider_materials.read(pair.colliders.x as usize); + let mat2 = collider_materials.read(pair.colliders.y as usize); + contacts.write( + slot, + IndexedManifold { + contact: manifold, + colliders: pair.colliders, + bodies: UVec2::new(body1, body2), + friction: mat1.combined_friction(&mat2), + restitution: mat1.combined_restitution(&mat2), + _padding: [0.0; 2], + }, + ); + } else { + contacts.at_mut(slot).contact.len = 0; } } } diff --git a/src_rbd_shaders/dynamics/color_buckets.rs b/src_rbd_shaders/dynamics/color_buckets.rs index 82a9c969..3c22fba5 100644 --- a/src_rbd_shaders/dynamics/color_buckets.rs +++ b/src_rbd_shaders/dynamics/color_buckets.rs @@ -1,14 +1,7 @@ //! Bucket-sort of contact constraints by graph-coloring color. //! -//! After the per-step coloring converges, the constraint indices are -//! bucket-sorted by color (`color_sorted_ids`, contacts layout) with -//! per-batch per-color exclusive prefix sums (`color_starts`), so each -//! colored solver sweep iterates only its own bucket instead of scanning the -//! whole constraint buffer. The count/start/cursor buffers are flat -//! `[num_batches × stride]` arrays with `stride = -//! BatchIndices::solver_color_buckets_stride` (= `max_colors + 3`, keeping -//! `starts[c + 1]` in bounds for every swept color). +use crate::broad_phase::ContactPlan; use khal_std::glamx::UVec3; use khal_std::macros::{spirv, spirv_bindgen}; use khal_std::{index::MaybeIndexUnchecked, iter::StepRng, sync::atomic_add_u32}; @@ -39,14 +32,14 @@ pub fn gpu_color_buckets_count( #[spirv(num_workgroups)] num_workgroups: UVec3, #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] constraints_colors: &[u32], #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] constraints: &[TwoBodyConstraint], - #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] contact_offsets: &[u32], + #[spirv(uniform, descriptor_set = 0, binding = 2)] contact_plan: &ContactPlan, #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] color_buckets: &mut [u32], #[spirv(uniform, descriptor_set = 0, binding = 4)] batch_ids: &BatchIndices, ) { let num_threads = num_workgroups.x * WORKGROUP_SIZE; let nb = batch_ids.num_batches; let stride = batch_ids.solver_color_buckets_stride; - let total = contact_offsets.read(batch_ids.num_batches as usize); + let total = contact_plan.bound; let constraints = Slice(constraints, 0); for i in StepRng::new(invocation_id.x..total, num_threads) { @@ -65,7 +58,7 @@ pub fn gpu_color_buckets_scatter( #[spirv(num_workgroups)] num_workgroups: UVec3, #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] constraints_colors: &[u32], #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] constraints: &[TwoBodyConstraint], - #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] contact_offsets: &[u32], + #[spirv(uniform, descriptor_set = 0, binding = 2)] contact_plan: &ContactPlan, #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] color_buckets: &mut [u32], #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] color_sorted_ids: &mut [u32], #[spirv(uniform, descriptor_set = 0, binding = 5)] batch_ids: &BatchIndices, @@ -73,7 +66,7 @@ pub fn gpu_color_buckets_scatter( let num_threads = num_workgroups.x * WORKGROUP_SIZE; let nb = batch_ids.num_batches; let stride = batch_ids.solver_color_buckets_stride; - let total = contact_offsets.read(batch_ids.num_batches as usize); + let total = contact_plan.bound; let constraints = Slice(constraints, 0); for i in StepRng::new(invocation_id.x..total, num_threads) { diff --git a/src_rbd_shaders/dynamics/coloring.rs b/src_rbd_shaders/dynamics/coloring.rs index 9ab942ed..7ac713dc 100644 --- a/src_rbd_shaders/dynamics/coloring.rs +++ b/src_rbd_shaders/dynamics/coloring.rs @@ -3,6 +3,7 @@ //! Assigns colors to constraints so that no two constraints sharing a body get the //! same color. Implements Jones-Plassmann-Luby and Topo-GC algorithms. +use crate::broad_phase::ContactPlan; use khal_std::glamx::UVec3; use khal_std::macros::{spirv, spirv_bindgen}; use khal_std::{ @@ -10,7 +11,7 @@ use khal_std::{ sync::{atomic_add_u32, atomic_max_u32}, }; -use crate::utils::{BatchIndices, Slice, SliceMut}; +use crate::utils::{Slice, SliceMut}; use khal_std::index::MaybeIndexUnchecked; use super::constraint::TwoBodyConstraint; @@ -48,10 +49,9 @@ pub fn gpu_reset_luby( #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] constraints_colors: &mut [u32], #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] constraints_rands: &mut [u32], #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] constraints: &[TwoBodyConstraint], - #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] contact_offsets: &[u32], - #[spirv(uniform, descriptor_set = 0, binding = 4)] batch_ids: &BatchIndices, + #[spirv(uniform, descriptor_set = 0, binding = 3)] contact_plan: &ContactPlan, ) { - let total = contact_offsets.read(batch_ids.num_batches as usize); + let total = contact_plan.bound; let i = invocation_id.x; if i < total { @@ -80,12 +80,11 @@ pub fn gpu_step_graph_coloring_luby( #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] uncolored: &mut u32, #[spirv(storage_buffer, descriptor_set = 0, binding = 6)] body_group: &[u32], #[spirv(uniform, descriptor_set = 0, binding = 7)] curr_color: &u32, - #[spirv(storage_buffer, descriptor_set = 0, binding = 8)] contact_offsets: &[u32], - #[spirv(uniform, descriptor_set = 0, binding = 9)] batch_ids: &BatchIndices, + #[spirv(uniform, descriptor_set = 0, binding = 8)] contact_plan: &ContactPlan, ) { let num_threads = num_workgroups.x * WORKGROUP_SIZE; - let total = contact_offsets.read(batch_ids.num_batches as usize); + let total = contact_plan.bound; let body_constraint_counts = Slice(body_constraint_counts, 0); let body_constraint_ids = Slice(body_constraint_ids, 0); let body_group = Slice(body_group, 0); @@ -184,10 +183,9 @@ pub fn gpu_reset_topo_gc( #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] constraints_colors: &mut [u32], #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] colored: &mut [u32], #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] constraints: &[TwoBodyConstraint], - #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] contact_offsets: &[u32], - #[spirv(uniform, descriptor_set = 0, binding = 4)] batch_ids: &BatchIndices, + #[spirv(uniform, descriptor_set = 0, binding = 3)] contact_plan: &ContactPlan, ) { - let total = contact_offsets.read(batch_ids.num_batches as usize); + let total = contact_plan.bound; let i = invocation_id.x; if i < total { @@ -233,13 +231,12 @@ pub fn gpu_step_graph_coloring_topo_gc( #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] constraints_colors: &mut [u32], #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] colored: &mut [u32], #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] num_colors: &mut u32, - #[spirv(storage_buffer, descriptor_set = 0, binding = 6)] contact_offsets: &[u32], + #[spirv(uniform, descriptor_set = 0, binding = 6)] contact_plan: &ContactPlan, #[spirv(storage_buffer, descriptor_set = 0, binding = 7)] body_group: &[u32], - #[spirv(uniform, descriptor_set = 0, binding = 8)] batch_ids: &BatchIndices, ) { let num_threads = num_workgroups.x * WORKGROUP_SIZE; - let total = contact_offsets.read(batch_ids.num_batches as usize); + let total = contact_plan.bound; let body_constraint_counts = Slice(body_constraint_counts, 0); let body_constraint_ids = Slice(body_constraint_ids, 0); let body_group = Slice(body_group, 0); @@ -323,13 +320,12 @@ pub fn gpu_fix_conflicts_topo_gc( #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] constraints_colors: &[u32], #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] colored: &mut [u32], #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] num_colors: &mut u32, - #[spirv(storage_buffer, descriptor_set = 0, binding = 6)] contact_offsets: &[u32], + #[spirv(uniform, descriptor_set = 0, binding = 6)] contact_plan: &ContactPlan, #[spirv(storage_buffer, descriptor_set = 0, binding = 7)] body_group: &[u32], - #[spirv(uniform, descriptor_set = 0, binding = 8)] batch_ids: &BatchIndices, ) { let num_threads = num_workgroups.x * WORKGROUP_SIZE; - let total = contact_offsets.read(batch_ids.num_batches as usize); + let total = contact_plan.bound; let body_constraint_counts = Slice(body_constraint_counts, 0); let body_constraint_ids = Slice(body_constraint_ids, 0); let body_group = Slice(body_group, 0); diff --git a/src_rbd_shaders/dynamics/multibody/contact_constraints.rs b/src_rbd_shaders/dynamics/multibody/contact_constraints.rs index 6557e4c4..9535f97b 100644 --- a/src_rbd_shaders/dynamics/multibody/contact_constraints.rs +++ b/src_rbd_shaders/dynamics/multibody/contact_constraints.rs @@ -18,20 +18,24 @@ use khal_std::glamx::UVec3; use khal_std::index::MaybeIndexUnchecked; use khal_std::iter::StepRng; use khal_std::macros::{spirv, spirv_bindgen}; -use khal_std::sync::workgroup_memory_barrier_with_group_sync; +use khal_std::sync::{atomic_add_u32, atomic_load_u32, workgroup_memory_barrier_with_group_sync}; +use crate::broad_phase::ContactPlan; use crate::dynamics::ConstraintSoftness; use crate::dynamics::body::{Velocity, WorldMassProperties}; use crate::dynamics::joint::SPATIAL_DIM; use crate::queries::IndexedManifold; +use crate::utils::BatchIndices; use crate::utils::linalg::{MAX_MB_DOFS, MatSlice, VSlice, lu_solve_in_place}; -use crate::utils::{BatchIndices, Slice}; use crate::{ANG_DIM, AngVector, DIM, Pose, Vector, gcross, gdot}; use super::types::{ CONTACT_CONSTRAINTS_PER_POINT, MB_CONS_SLOT_RESERVE, MB_CONTACT_KIND_NORMAL, - MB_CONTACT_KIND_TANGENT, MultibodyContactConstraint, MultibodyInfo, MultibodyLinkStatic, + MB_CONTACT_KIND_TANGENT, MbContactIndexEntry, MultibodyContactConstraint, MultibodyInfo, + MultibodyLinkStatic, }; + +const MB_SWEEP_WG: u32 = 64; use super::utils::zero_kinematic_dofs; use super::ws_soa::{WS_LTW, WS_WORLD_COM, WsAddr, ws_pose, ws_vec}; @@ -118,84 +122,113 @@ fn fill_contact_jac_row( out_jacs.write(col_offset + j as usize, prev + dot); } } + +struct MbContactOwner { + mb: u32, + batch: u32, + link_a: u32, + link_b: u32, + mb_on_first: u32, +} + +const MB_OWNER_SKIP: MbContactOwner = MbContactOwner { + mb: u32::MAX, + batch: 0, + link_a: u32::MAX, + link_b: u32::MAX, + mb_on_first: 0, +}; + #[inline(always)] -fn mb_contact_demand( +fn mb_contact_owner( im: &IndexedManifold, - mb_idx: u32, - self_contacts_enabled: u32, body_to_link: &[[u32; 2]], -) -> u32 { + multibody_info: &[MultibodyInfo], + batch_ids: &BatchIndices, +) -> MbContactOwner { if im.contact.len == 0 { - return 0; + return MB_OWNER_SKIP; } let l1 = body_to_link.read(im.bodies.x as usize); let l2 = body_to_link.read(im.bodies.y as usize); - let mb_on_1 = l1[0] == mb_idx; - let mb_on_2 = l2[0] == mb_idx; - if !mb_on_1 && !mb_on_2 { - return 0; + if l1[0] == u32::MAX && l2[0] == u32::MAX { + return MB_OWNER_SKIP; } if l1[0] != u32::MAX && l2[0] != u32::MAX && l1[0] != l2[0] { - return 0; - } - let is_self = mb_on_1 && mb_on_2; - if is_self && self_contacts_enabled == 0 { - return 0; + return MB_OWNER_SKIP; } + let is_self = l1[0] != u32::MAX && l2[0] != u32::MAX; if is_self && l1[1] == l2[1] { - return 0; + return MB_OWNER_SKIP; + } + let mb_on_first = l1[0] != u32::MAX; + let owner = if mb_on_first { l1[0] } else { l2[0] }; + let batch = batch_ids.collider_batch(im.bodies.x); + let mb = multibody_info.read(batch_ids.mbi(batch, owner as usize)); + if mb.ndofs == 0 { + return MB_OWNER_SKIP; + } + if is_self && mb.self_contacts_enabled == 0 { + return MB_OWNER_SKIP; + } + MbContactOwner { + mb: owner, + batch, + link_a: if mb_on_first { l1[1] } else { l2[1] }, + link_b: if is_self { l2[1] } else { u32::MAX }, + mb_on_first: mb_on_first as u32, } - im.contact.len * CONTACT_CONSTRAINTS_PER_POINT } + #[spirv_bindgen] #[spirv(compute(threads(64)))] pub fn gpu_mb_count_contact_constraints( #[spirv(global_invocation_id)] invocation_id: UVec3, - #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] - multibody_info: &mut [MultibodyInfo], + #[spirv(num_workgroups)] num_workgroups: UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] multibody_info: &[MultibodyInfo], #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] contacts: &[IndexedManifold], #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] body_to_link: &[[u32; 2]], - #[spirv(uniform, descriptor_set = 0, binding = 3)] batch_ids: &BatchIndices, + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] mb_cons_counts: &mut [u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] mb_index_counts: &mut [u32], + #[spirv(uniform, descriptor_set = 0, binding = 5)] contact_plan: &ContactPlan, + #[spirv(uniform, descriptor_set = 0, binding = 6)] batch_ids: &BatchIndices, ) { - let num_mb = batch_ids.multibodies_len; - if invocation_id.x >= num_mb * batch_ids.num_batches { - return; - } - let batch_id = invocation_id.x / num_mb; - let mb_idx = invocation_id.x % num_mb; - let mut mb = multibody_info.read(batch_ids.mbi(batch_id, mb_idx as usize)); - let mut count = 0u32; - if mb.ndofs != 0 { - let contacts_slice = Slice(contacts, mb.batch_contacts_start as usize); - for ci in 0..mb.batch_contacts_len { - count += mb_contact_demand( - contacts_slice.at(ci as usize), - mb_idx, - mb.self_contacts_enabled, - body_to_link, - ); + let num_threads = num_workgroups.x * MB_SWEEP_WG; + let total = contact_plan.bound; + for t in StepRng::new(invocation_id.x..total, num_threads) { + let im = contacts.read(t as usize); + let owner = mb_contact_owner(&im, body_to_link, multibody_info, batch_ids); + if owner.mb == u32::MAX { + continue; } + let slot = batch_ids.mbi(owner.batch, owner.mb as usize); + let demand = im.contact.len * CONTACT_CONSTRAINTS_PER_POINT; + atomic_add_u32(mb_cons_counts.at_mut(slot), demand); + atomic_add_u32(mb_index_counts.at_mut(slot), 1); } - mb.contact_constraint_count = count; - multibody_info.write(batch_ids.mbi(batch_id, mb_idx as usize), mb); } + #[spirv_bindgen] #[spirv(compute(threads(1)))] pub fn gpu_mb_cons_offsets_scan( #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] multibody_info: &mut [MultibodyInfo], - #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] mb_cons_demand: &mut [u32], - #[spirv(uniform, descriptor_set = 0, binding = 2)] batch_ids: &BatchIndices, + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] mb_cons_counts: &mut [u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] mb_index_counts: &mut [u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] mb_cons_demand: &mut [u32], + #[spirv(uniform, descriptor_set = 0, binding = 4)] batch_ids: &BatchIndices, ) { let total_infos = batch_ids.multibodies_len * batch_ids.num_batches; let capacity = batch_ids.mb_contact_constraints_capacity; + let mut demand = 0u32; let mut reserved_total = 0u32; for i in 0..total_infos { - let count = multibody_info.read(i as usize).contact_constraint_count; + let count = atomic_load_u32(mb_cons_counts.at_mut(i as usize)); demand += count; reserved_total += count.min(MB_CONS_SLOT_RESERVE); } + #[allow(clippy::implicit_saturating_sub)] let mut extra_budget = if capacity > reserved_total { capacity - reserved_total } else { @@ -203,9 +236,10 @@ pub fn gpu_mb_cons_offsets_scan( }; let mut acc = 0u32; + let mut index_acc = 0u32; for i in 0..total_infos { let mut mb = multibody_info.read(i as usize); - let count = mb.contact_constraint_count; + let count = atomic_load_u32(mb_cons_counts.at_mut(i as usize)); let reserve = count.min(MB_CONS_SLOT_RESERVE); let extra = (count - reserve).min(extra_budget); extra_budget -= extra; @@ -213,11 +247,56 @@ pub fn gpu_mb_cons_offsets_scan( let avail = (reserve + extra).min(capacity - start); mb.contact_constraint_start = start; mb.contact_constraint_count = avail; - multibody_info.write(i as usize, mb); acc = start + avail; + + mb.contact_index_start = index_acc; + mb.contact_index_len = atomic_load_u32(mb_index_counts.at_mut(i as usize)); + index_acc += mb.contact_index_len; + + multibody_info.write(i as usize, mb); + mb_cons_counts.write(i as usize, 0); + mb_index_counts.write(i as usize, 0); } mb_cons_demand.write(0, demand); } + +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn gpu_mb_scatter_contact_index( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[spirv(num_workgroups)] num_workgroups: UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] multibody_info: &[MultibodyInfo], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] contacts: &[IndexedManifold], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] body_to_link: &[[u32; 2]], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] mb_index_counts: &mut [u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] + mb_contact_index: &mut [MbContactIndexEntry], + #[spirv(uniform, descriptor_set = 0, binding = 5)] contact_plan: &ContactPlan, + #[spirv(uniform, descriptor_set = 0, binding = 6)] batch_ids: &BatchIndices, +) { + let num_threads = num_workgroups.x * MB_SWEEP_WG; + let total = contact_plan.bound; + for t in StepRng::new(invocation_id.x..total, num_threads) { + let im = contacts.read(t as usize); + let owner = mb_contact_owner(&im, body_to_link, multibody_info, batch_ids); + if owner.mb == u32::MAX { + continue; + } + let slot = batch_ids.mbi(owner.batch, owner.mb as usize); + let mb = multibody_info.read(slot); + let pos = atomic_add_u32(mb_index_counts.at_mut(slot), 1); + mb_contact_index.write( + (mb.contact_index_start + pos) as usize, + MbContactIndexEntry { + contact_slot: t, + link_a: owner.link_a, + link_b: owner.link_b, + mb_on_first: owner.mb_on_first, + }, + ); + } +} + #[spirv_bindgen] #[spirv(compute(threads(64)))] pub fn gpu_mb_save_prev_cons_bounds( @@ -255,7 +334,8 @@ pub fn gpu_mb_init_contact_constraints( #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] multibody_info: &mut [MultibodyInfo], #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] links_workspace: &[Vec4], - #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] body_to_link: &[[u32; 2]], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] + mb_contact_index: &[MbContactIndexEntry], #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] contact_constraints: &mut [MultibodyContactConstraint], #[spirv(uniform, descriptor_set = 0, binding = 4)] softness: &ConstraintSoftness, @@ -301,16 +381,14 @@ pub fn gpu_mb_init_contact_constraints( let avail = mb.contact_constraint_count; let wa = WsAddr::new(mb.first_link as usize, batch_ids.num_batches, batch_id); - let contacts_slice = Slice(contacts, mb.batch_contacts_start as usize); - let n_contacts = mb.batch_contacts_len; + let idx_base = mb.contact_index_start as usize; + let n_entries = mb.contact_index_len; let mut count = 0u32; - for ci in 0..n_contacts { - let im = contacts_slice[ci as usize]; - let demand = mb_contact_demand(&im, mb_idx, mb.self_contacts_enabled, body_to_link); - if demand == 0 { - continue; - } + for ci in 0..n_entries { + let entry = mb_contact_index.read(idx_base + ci as usize); + let im = contacts.read(entry.contact_slot as usize); + let demand = im.contact.len * CONTACT_CONSTRAINTS_PER_POINT; if count + demand > avail { break; } @@ -318,16 +396,14 @@ pub fn gpu_mb_init_contact_constraints( let b1 = im.bodies.x; let b2 = im.bodies.y; - let l1 = body_to_link.read(b1 as usize); - let l2 = body_to_link.read(b2 as usize); - let mb_on_1 = l1[0] == mb_idx; - let is_self = mb_on_1 && l2[0] == mb_idx; + let mb_on_1 = entry.mb_on_first != 0; + let is_self = entry.link_b != u32::MAX; let (mb_link_id_a, mb_link_id_b, free_body_id) = if is_self { - (l1[1], l2[1], u32::MAX) + (entry.link_a, entry.link_b, u32::MAX) } else if mb_on_1 { - (l1[1], u32::MAX, b2) + (entry.link_a, u32::MAX, b2) } else { - (l2[1], u32::MAX, b1) + (entry.link_a, u32::MAX, b1) }; let pose1 = poses.read(id1 as usize); @@ -655,48 +731,13 @@ pub fn gpu_mb_init_contact_constraints( } } - // The solve kernels only iterate `0..count`, but next frame's warmstart - // match scans the whole slab, so the leftovers of the previous build have - // to be marked inactive. if lane == 0 { mb.contact_constraint_count = count; multibody_info.write(batch_ids.mbi(batch_id, mb_idx as usize), mb); } } -/// HACK: stash `contacts_len[batch]` into each multibody's `batch_contacts_len`. -/// -/// This exists only to work around the web 8-storage-bindings limit for kernels -/// that bind multibodies but don’t have any room left to bind `contacts_len`. -#[spirv_bindgen] -#[spirv(compute(threads(64)))] -pub fn gpu_mb_stash_contacts_len( - #[spirv(global_invocation_id)] invocation_id: UVec3, - #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] - multibody_info: &mut [MultibodyInfo], - #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] contacts_len: &[u32], - #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] contact_offsets: &[u32], - #[spirv(uniform, descriptor_set = 0, binding = 3)] batch_ids: &BatchIndices, -) { - let num_mb = batch_ids.multibodies_len; - if invocation_id.x >= num_mb * batch_ids.num_batches { - return; - } - let batch_id = invocation_id.x / num_mb; - let mb_idx = invocation_id.x % num_mb; - let seg_start = contact_offsets.read(batch_id as usize); - let seg_end = contact_offsets.read(batch_id as usize + 1); - let mut mb = multibody_info.read(batch_ids.mbi(batch_id, mb_idx as usize)); - mb.batch_contacts_len = contacts_len.read(batch_id as usize).min(seg_end - seg_start); - mb.batch_contacts_start = seg_start; - multibody_info.write(batch_ids.mbi(batch_id, mb_idx as usize), mb); -} - -/// Snapshot every contact-constraint slot into the "previous frame" slab that -/// `gpu_mb_transfer_contact_warmstart` matches against. Called once per visible -/// frame from `init_step`, before the substep loop rebuilds the live slab. /// -/// One thread per (slot, multibody, batch). #[spirv_bindgen] #[spirv(compute(threads(64)))] pub fn gpu_mb_snapshot_contact_warmstart( diff --git a/src_rbd_shaders/dynamics/multibody/types.rs b/src_rbd_shaders/dynamics/multibody/types.rs index 3070d5fe..75fea041 100644 --- a/src_rbd_shaders/dynamics/multibody/types.rs +++ b/src_rbd_shaders/dynamics/multibody/types.rs @@ -443,15 +443,12 @@ pub struct MultibodyInfo { /// `DISABLE_SELF_CONTACTS`). The contact-constraint kernel skips self /// contacts when this is `0`. pub self_contacts_enabled: u32, - /// Per-frame count of active multibody contact constraints emitted for this - /// multibody. Written by `gpu_mb_init_contact_constraints`, read by the - /// warmstart / finalize / solve / remove-bias contact kernels. pub contact_constraint_count: u32, pub contact_constraint_start: u32, pub old_contact_constraint_start: u32, pub old_contact_constraint_count: u32, - pub batch_contacts_len: u32, - pub batch_contacts_start: u32, + pub contact_index_len: u32, + pub contact_index_start: u32, /// First entry of this multibody's DoF couplings in the `dof_couplings` /// buffer (relative to the batch's coupling slice). pub first_coupling: u32, @@ -459,6 +456,16 @@ pub struct MultibodyInfo { pub num_couplings: u32, } +#[derive(Copy, Clone, Default)] +#[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] +#[repr(C)] +pub struct MbContactIndexEntry { + pub contact_slot: u32, + pub link_a: u32, + pub link_b: u32, + pub mb_on_first: u32, +} + /// One holonomic coupling `q2 = coeff·q1 + offset` between two generalized /// coordinates of the same multibody (rapier's `MultibodyDofCoupling`), /// converted to the GPU's re-numbered assembly ids at build time. diff --git a/src_rbd_shaders/dynamics/solver.rs b/src_rbd_shaders/dynamics/solver.rs index 2d3d6e90..4a23544e 100644 --- a/src_rbd_shaders/dynamics/solver.rs +++ b/src_rbd_shaders/dynamics/solver.rs @@ -2,6 +2,7 @@ //! //! This module contains the actual GPU compute shader entry points for the physics solver. +use crate::broad_phase::ContactPlan; use khal_std::glamx::UVec3; use khal_std::macros::{spirv, spirv_bindgen}; @@ -37,17 +38,16 @@ pub fn gpu_solver_init_constraints( constraints: &mut [TwoBodyConstraint], #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] constraint_builders: &mut [TwoBodyConstraintBuilder], - #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] contact_offsets: &[u32], + #[spirv(uniform, descriptor_set = 0, binding = 3)] contact_plan: &ContactPlan, #[spirv(storage_buffer, descriptor_set = 1, binding = 0)] collider_world_poses: &[Pose], #[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(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 total = contact_offsets.read(batch_ids.num_batches as usize); + let total = contact_plan.bound; let collider_world_poses = Slice(collider_world_poses, 0); let solver_body_poses = Slice(solver_body_poses, 0); let vels = Slice(vels, 0); @@ -83,12 +83,11 @@ pub fn gpu_solver_count_constraints( #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] body_constraint_counts: &mut [u32], #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] body_group: &[u32], #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] mprops: &[WorldMassProperties], - #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] contact_offsets: &[u32], - #[spirv(uniform, descriptor_set = 0, binding = 5)] batch_ids: &BatchIndices, + #[spirv(uniform, descriptor_set = 0, binding = 4)] contact_plan: &ContactPlan, ) { let num_threads = num_workgroups.x * WORKGROUP_SIZE; - let total = contact_offsets.read(batch_ids.num_batches as usize); + let total = contact_plan.bound; let contacts = Slice(contacts, 0); let mut body_constraint_counts = SliceMut(body_constraint_counts, 0); let body_group = Slice(body_group, 0); @@ -130,14 +129,13 @@ pub fn gpu_solver_update_constraints( constraints: &mut [TwoBodyConstraint], #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] constraint_builders: &[TwoBodyConstraintBuilder], - #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] contact_offsets: &[u32], + #[spirv(uniform, descriptor_set = 0, binding = 2)] contact_plan: &ContactPlan, #[spirv(storage_buffer, descriptor_set = 1, binding = 0)] solver_body_poses: &[Pose], #[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 total = contact_offsets.read(batch_ids.num_batches as usize); + let total = contact_plan.bound; let mut constraints = SliceMut(constraints, 0); let constraint_builders = Slice(constraint_builders, 0); let solver_body_poses = Slice(solver_body_poses, 0); @@ -165,14 +163,13 @@ pub fn gpu_solver_refresh_rhs_wo_bias( constraints: &mut [TwoBodyConstraint], #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] constraint_builders: &[TwoBodyConstraintBuilder], - #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] contact_offsets: &[u32], + #[spirv(uniform, descriptor_set = 0, binding = 2)] contact_plan: &ContactPlan, #[spirv(storage_buffer, descriptor_set = 1, binding = 0)] solver_body_poses: &[Pose], #[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 total = contact_offsets.read(batch_ids.num_batches as usize); + let total = contact_plan.bound; let mut constraints = SliceMut(constraints, 0); let constraint_builders = Slice(constraint_builders, 0); let solver_body_poses = Slice(solver_body_poses, 0); @@ -197,14 +194,13 @@ pub fn gpu_solver_sort_constraints( #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] body_constraint_counts: &mut [u32], #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] mprops: &[WorldMassProperties], #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] contacts: &[IndexedManifold], - #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] contact_offsets: &[u32], + #[spirv(uniform, descriptor_set = 0, binding = 3)] contact_plan: &ContactPlan, #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] body_constraint_ids: &mut [u32], #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] body_group: &[u32], - #[spirv(uniform, descriptor_set = 0, binding = 6)] batch_ids: &BatchIndices, ) { let num_threads = num_workgroups.x * WORKGROUP_SIZE; - let total = contact_offsets.read(batch_ids.num_batches as usize); + let total = contact_plan.bound; let contacts = Slice(contacts, 0); let mut body_constraint_counts = SliceMut(body_constraint_counts, 0); let body_group = Slice(body_group, 0); @@ -637,7 +633,9 @@ pub fn gpu_init_solver_bodies( let idx = i as usize; solver_body_poses.write( idx, - body_poses.read(idx).prepend_translation(local_mprops.at(idx).com), + body_poses + .read(idx) + .prepend_translation(local_mprops.at(idx).com), ); } } diff --git a/src_rbd_shaders/dynamics/warmstart.rs b/src_rbd_shaders/dynamics/warmstart.rs index 4edf8034..5861b5f1 100644 --- a/src_rbd_shaders/dynamics/warmstart.rs +++ b/src_rbd_shaders/dynamics/warmstart.rs @@ -3,11 +3,12 @@ //! Transfers impulses from frame `n-1` to frame `n` as initial guesses for the solver. //! Contacts are matched by proximity in local coordinates (threshold: 10cm). +use crate::broad_phase::ContactPlan; use khal_std::glamx::UVec3; use khal_std::macros::{spirv, spirv_bindgen}; use super::constraint::{TwoBodyConstraint, TwoBodyConstraintBuilder}; -use crate::utils::{BatchIndices, Slice, SliceMut}; +use crate::utils::{Slice, SliceMut}; use khal_std::index::MaybeIndexUnchecked; /// Transfers warmstart impulses from previous frame to current frame. @@ -25,10 +26,9 @@ pub fn gpu_transfer_warmstart_impulses( new_constraints: &mut [TwoBodyConstraint], #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] new_constraint_builders: &[TwoBodyConstraintBuilder], - #[spirv(storage_buffer, descriptor_set = 0, binding = 6)] contact_offsets: &[u32], - #[spirv(uniform, descriptor_set = 0, binding = 7)] batch_ids: &BatchIndices, + #[spirv(uniform, descriptor_set = 0, binding = 6)] contact_plan: &ContactPlan, ) { - let total = contact_offsets.read(batch_ids.num_batches as usize); + let total = contact_plan.bound; let old_body_constraint_counts = Slice(old_body_constraint_counts, 0); let old_body_constraint_ids = Slice(old_body_constraint_ids, 0); let old_constraints = Slice(old_constraints, 0); @@ -75,10 +75,9 @@ pub fn gpu_seed_colors_from_warmstart( #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] old_constraints_colors: &[u32], #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] constraints_colors: &mut [u32], #[spirv(storage_buffer, descriptor_set = 0, binding = 6)] colored: &mut [u32], - #[spirv(storage_buffer, descriptor_set = 0, binding = 7)] contact_offsets: &[u32], - #[spirv(uniform, descriptor_set = 0, binding = 8)] batch_ids: &BatchIndices, + #[spirv(uniform, descriptor_set = 0, binding = 7)] contact_plan: &ContactPlan, ) { - let total = contact_offsets.read(batch_ids.num_batches as usize); + let total = contact_plan.bound; let old_body_constraint_counts = Slice(old_body_constraint_counts, 0); let old_body_constraint_ids = Slice(old_body_constraint_ids, 0); let old_constraints = Slice(old_constraints, 0); From a3e20be92677de33f54110c27dca6aabf144e1eb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Sun, 30 Aug 2026 11:43:56 +0200 Subject: [PATCH 5/7] chore(rbd): update docs --- src/state.rs | 5 + src_rbd/broad_phase/lbvh.rs | 2 + src_rbd/broad_phase/narrow_phase.rs | 42 ++++++++ src_rbd/dynamics/coloring.rs | 12 ++- src_rbd/dynamics/joint.rs | 5 +- src_rbd/dynamics/multibody/env_reset.rs | 12 +-- .../dynamics/multibody/loop_closing_joints.rs | 25 ++++- .../multibody/multibody_from_rapier.rs | 29 +++++- src_rbd/dynamics/multibody/multibody_set.rs | 58 +++++++++-- .../dynamics/multibody/multibody_solver.rs | 39 ++++++-- src_rbd/dynamics/solver.rs | 15 ++- src_rbd/dynamics/warmstart.rs | 2 + src_rbd/pipeline/insertion_removal.rs | 33 +++++++ src_rbd/pipeline/rbd_state.rs | 56 ++++++++++- src_rbd/pipeline/rbd_state_from_rapier.rs | 32 ++++++- src_rbd/pipeline/rbd_step.rs | 31 ++++++ src_rbd/pipeline/test_batched_stacks.rs | 96 +++++++++++++++++++ src_rbd/utils/radix_sort/mod.rs | 2 + src_rbd_shaders/broad_phase/brute_force.rs | 2 + src_rbd_shaders/broad_phase/lbvh.rs | 22 ++++- src_rbd_shaders/broad_phase/narrow_phase.rs | 94 +++++++++++++++++- src_rbd_shaders/dynamics/color_buckets.rs | 18 +++- src_rbd_shaders/dynamics/coloring.rs | 5 + src_rbd_shaders/dynamics/joint_constraint.rs | 2 +- .../multibody/compute_dynamics_pre.rs | 5 +- .../dynamics/multibody/contact_constraints.rs | 91 +++++++++++++++++- .../dynamics/multibody/env_reset.rs | 13 ++- .../dynamics/multibody/gravity_and_lu.rs | 1 + .../impulse_joint_constraints/jacobians.rs | 7 +- .../impulse_joint_constraints/kernels.rs | 12 ++- .../impulse_joint_constraints/update.rs | 9 ++ .../dynamics/multibody/integrate.rs | 2 + .../dynamics/multibody/solve_constraints.rs | 19 +++- src_rbd_shaders/dynamics/multibody/types.rs | 35 +++++++ src_rbd_shaders/dynamics/multibody/ws_soa.rs | 4 +- src_rbd_shaders/dynamics/solver.rs | 12 +++ src_rbd_shaders/dynamics/warmstart.rs | 2 + src_rbd_shaders/utils/indices.rs | 34 ++++++- 38 files changed, 818 insertions(+), 67 deletions(-) diff --git a/src/state.rs b/src/state.rs index 4840b6ee..fa720d82 100644 --- a/src/state.rs +++ b/src/state.rs @@ -349,6 +349,9 @@ impl NexusState { self.capacities.rbd.collisions_capacity = capacity.max(1); } + /// Sets the per-batch multibody contact-constraint slot budget (see + /// [`RbdCapacities::mb_contact_constraints_capacity`]). Takes effect on the + /// next state (re)build. pub fn set_rbd_mb_contact_constraints_capacity(&mut self, capacity: u32) { self.capacities.rbd.mb_contact_constraints_capacity = capacity.max(1); } @@ -998,6 +1001,8 @@ impl NexusState { // `gpu_id` is its *body* slot, not a collider slot, since a body may // own several colliders. Body slots are assigned in the order // `from_rapier` uses (the first time each parent body is seen while + // iterating colliders); the per-body buffers are batch-interleaved, + // so `gpu_id = local_slot * num_batches + env`. let nb = rbd_state.num_batches(); for (env_idx, world) in self.rbd_envs.iter().enumerate() { let mut body_slot: std::collections::HashMap<_, u32> = diff --git a/src_rbd/broad_phase/lbvh.rs b/src_rbd/broad_phase/lbvh.rs index 347a8804..80fb69d9 100644 --- a/src_rbd/broad_phase/lbvh.rs +++ b/src_rbd/broad_phase/lbvh.rs @@ -32,6 +32,8 @@ pub struct GpuLbvh { refit_internal: GpuLbvhRefitInternal, reset_collision_pairs: GpuLbvhResetCollisionPairs, find_collision_pairs: GpuLbvhFindCollisionPairs, + /// Writes the `[total/64, 1, 1]` indirect grid from the single global pair + /// counter. flat_list_dispatch: GpuFlatListDispatch, // Kernels for brute-force broad-phase for small scenes // (typically, small scenes but many batches). diff --git a/src_rbd/broad_phase/narrow_phase.rs b/src_rbd/broad_phase/narrow_phase.rs index 55c3fb7f..c902d9aa 100644 --- a/src_rbd/broad_phase/narrow_phase.rs +++ b/src_rbd/broad_phase/narrow_phase.rs @@ -15,25 +15,33 @@ use khal::backend::{GpuBackend, GpuBackendError, GpuPass}; use khal::{BufferUsages, Shader}; use vortx::tensor::Tensor; +/// Narrow-phase kernel bundle (see [`GpuNarrowPhase`]). #[derive(Shader)] struct GpuNarrowPhaseShaders { reset_narrow_phase: GpuResetNarrowPhase, narrow_phase: GpuNarrowPhaseShapeShape, + /// Defers complex shape pairs (PFM / trimesh / polyline) into the /// `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, + /// Publishes the clamped list totals + every derived dispatch grid. contact_plan: GpuContactPlan, + /// Extracts each PFM entry's pair index into the radix-sort key buffer. pfm_sort_keys: GpuPfmSortKeys, } +/// GPU shader for narrow-phase collision detection. pub struct GpuNarrowPhase { shaders: GpuNarrowPhaseShaders, + /// Groups the PFM entries per originating pair (contact-reduction only). sort: RadixSort, } impl GpuNarrowPhase { + /// Loads the narrow-phase kernels (and the radix sort they use) from the + /// given backend. pub fn from_backend(backend: &GpuBackend) -> Result { Ok(Self { shaders: GpuNarrowPhaseShaders::from_backend(backend)?, @@ -42,16 +50,26 @@ impl GpuNarrowPhase { } } +/// Buffers backing the per-pair PFM sort of the contact-reduction path. All +/// sized by the collision-pair capacity (the PFM list shares it); `resize` +/// must be called whenever that capacity changes. pub struct PfmSortState { + /// Per-entry pair-index keys, written by `gpu_pfm_sort_keys`. keys: Tensor, + /// The identity permutation `0..capacity`: the sort's value input, and + /// the `pfm_order` indirection of the unsorted path. identity: Tensor, + /// Sort outputs (pair-grouped keys + entry permutation). sorted_keys: Tensor, sorted_values: Tensor, + /// Clamped PFM entry count (the sort's GPU-side `n_sort`), written by + /// `gpu_contact_plan`. sort_len: Tensor, workspace: RadixSortWorkspace, } impl PfmSortState { + /// Allocates the sort buffers for a PFM list of `capacity` entries. pub fn new(backend: &GpuBackend, capacity: u32) -> Self { let storage = BufferUsages::STORAGE; let identity: Vec = (0..capacity).collect(); @@ -65,10 +83,12 @@ impl PfmSortState { } } + /// Regrow after a collision-pair capacity change. pub fn resize(&mut self, backend: &GpuBackend, capacity: u32) { *self = Self::new(backend, capacity); } + /// The sorted-keys tensor consumed by `gpu_reduce_contacts`. #[cfg(feature = "dim3")] pub fn sorted_keys(&self) -> &Tensor { &self.sorted_keys @@ -101,16 +121,25 @@ impl GpuNarrowPhase { collider_materials: &Tensor, sim_params: &Tensor, // Optional: merge each collider pair's manifolds into one before the + // solvers see them. Enables the per-pair PFM sort (the manifolds of a + // pair must land in contiguous slots for the per-run reduction). reduce_contacts: bool, + // The `[total/64, 1, 1]` grid written by the broad phase from the + // single global pair counter. collision_pairs_indirect: &Tensor<[u32; 3]>, + // The flat pair/PFM buffer capacity, bounding the sort-key width. collision_pairs_capacity: u32, ) -> Result<(), GpuBackendError> { + // The per-run reduction kernel is 3D-only; without it the sort would + // group entries nobody consumes. let reduce_contacts = reduce_contacts && cfg!(feature = "dim3"); self.shaders .reset_narrow_phase .call(pass, 1u32, pfm_pairs_len)?; + // Defer the complex shape pairs into `pfm_pairs` FIRST: both list + // counters must be final before the plan is published. self.shaders.narrow_phase_deferred.call( pass, collision_pairs_indirect, @@ -126,6 +155,8 @@ impl GpuNarrowPhase { indices, )?; + // Clamped totals + every derived grid (contacts sweep, PFM sweep, + // multibody contact sweep) in one serial thread. self.shaders.contact_plan.call( pass, 1u32, @@ -139,6 +170,7 @@ impl GpuNarrowPhase { batch_indices, )?; + // Analytic pairs: pair `t` writes contact slot `t` (inert on a miss). self.shaders.narrow_phase.call( pass, collision_pairs_indirect, @@ -152,6 +184,11 @@ impl GpuNarrowPhase { sim_params, )?; + // Contact reduction needs each pair's manifolds contiguous: group the + // PFM entries per originating pair with a stable radix sort (which + // also makes the PFM contact slots deterministic). Without reduction + // the entries are consumed in emission order through the identity + // permutation. if reduce_contacts { self.shaders.pfm_sort_keys.call( pass, @@ -160,8 +197,11 @@ impl GpuNarrowPhase { &*contact_plan, &mut pfm_sort.keys, )?; + // Keys are flat pair indices: bounded by the pair capacity. let sorting_bits = (32 - collision_pairs_capacity.saturating_sub(1).leading_zeros()).max(1); + // Split borrows: the sort reads `keys`/`identity` and writes the + // `sorted_*` pair. let PfmSortState { keys, identity, @@ -184,6 +224,8 @@ impl GpuNarrowPhase { )?; } + // PFM entry `i` (in sorted order when the sort ran) writes contact + // slot `pairs_total + i` (inert on a miss). let pfm_order = if reduce_contacts { &pfm_sort.sorted_values } else { diff --git a/src_rbd/dynamics/coloring.rs b/src_rbd/dynamics/coloring.rs index 985690b2..fc377717 100644 --- a/src_rbd/dynamics/coloring.rs +++ b/src_rbd/dynamics/coloring.rs @@ -45,12 +45,19 @@ pub struct GpuColoring { /// Buffers for the per-color constraint bucket sort. pub struct ColorBucketsArgs<'a> { + /// Flat dispatch grid over the whole contacts range. pub contacts_len_indirect: &'a Tensor<[u32; 3]>, /// Color assigned to each constraint by graph coloring. pub constraints_colors: &'a Tensor, + /// The colored constraints (batch recovered from their global body ids). pub constraints: &'a Tensor, + /// Clamped per-frame list totals (the flat contact sweep bound). pub contact_plan: &'a Tensor, + /// The single `(color, batch)` bucket buffer, color-major + /// (`solver_color_buckets_stride * num_batches` entries): counts, then + /// scanned exclusive starts, then post-scatter exclusive ends. pub color_buckets: &'a mut Tensor, + /// Constraint ids bucket-sorted by `(color, batch)`. pub color_sorted_ids: &'a mut Tensor, /// Shared per-batch capacity / section-offset uniform. pub batch_indices: &'a Tensor, @@ -78,6 +85,7 @@ pub struct ColoringArgs<'a> { pub uncolored: &'a mut Tensor, /// Staging buffer for reading uncolored count on CPU. pub uncolored_staging: &'a Tensor, + /// Clamped per-frame list totals (the flat contact sweep bound). pub contact_plan: &'a Tensor, /// Buffer tracking which constraints are colored. pub colored: &'a mut Tensor, @@ -238,7 +246,9 @@ impl GpuColoring { num_colors } - /// Bucket-sorts the constraint ids by their color. + /// Bucket-sorts the constraint ids by `(color, batch)`: zero the buckets, + /// count, exclusive-prefix-scan them in place, then scatter (which turns + /// the starts into exclusive ends, the form the sweeps read). pub fn dispatch_build_color_buckets( &self, backend: &GpuBackend, diff --git a/src_rbd/dynamics/joint.rs b/src_rbd/dynamics/joint.rs index f4d4eb70..6cda04b5 100644 --- a/src_rbd/dynamics/joint.rs +++ b/src_rbd/dynamics/joint.rs @@ -307,7 +307,10 @@ impl GpuImpulseJointSet { } } - // Build flat joint buffer [num_batches * max_joints], padded with zeroed joints. + // Build the flat joint buffer [num_batches * max_joints], padded with + // zeroed joints and batch-interleaved (`joint * num_batches + batch`) + // so the solver's flat sweeps put the same joint of consecutive + // batches on adjacent lanes. let dummy_joint = ImpulseJoint::zeroed(); let mut all_joints = vec![dummy_joint; num_batches as usize * max_joints as usize]; diff --git a/src_rbd/dynamics/multibody/env_reset.rs b/src_rbd/dynamics/multibody/env_reset.rs index 2f782eea..94fcc75e 100644 --- a/src_rbd/dynamics/multibody/env_reset.rs +++ b/src_rbd/dynamics/multibody/env_reset.rs @@ -25,8 +25,8 @@ 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. +/// workspace (which carries the generalized coordinates), the static link +/// descriptors, and the generalized velocities of batch 0. #[derive(Clone)] pub struct GpuMultibodySnapshot { /// AoS per-link workspace of batch 0, `links_per_batch` entries including @@ -69,8 +69,8 @@ impl GpuMultibodySnapshot { } /// 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 + /// frame). Rotations, joint coordinates past the free linear DoFs and + /// velocities 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 @@ -264,8 +264,8 @@ impl GpuMultibodySet { 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 + /// Uploads the reset templates once as GPU-resident blobs (SoA workspace + /// and links) 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. diff --git a/src_rbd/dynamics/multibody/loop_closing_joints.rs b/src_rbd/dynamics/multibody/loop_closing_joints.rs index 9db35330..314ec27e 100644 --- a/src_rbd/dynamics/multibody/loop_closing_joints.rs +++ b/src_rbd/dynamics/multibody/loop_closing_joints.rs @@ -246,12 +246,22 @@ impl GpuMultibodySet { per_env_builders.push(sorted_builders); } - // Stage 2 — flatten with per-batch padding to `max_joints`. + // Stage 2 — batch-interleaved upload: joint slot `i` of batch `b` + // lands at `i * num_batches + b`, and constraint slots follow the same + // per-slot interleave (their ids inside the builders are batch-local). + // The jacobians buffer is interleaved at per-joint-region granularity + // instead: joint `j`'s region for batch `b` starts at + // `jacobian_offset_j * nb + b * jacobian_capacity_j` and is dense + // inside, so the solver's lane-parallel row reads stay contiguous. let joints_cap = max_joints.max(1); let cons_cap = (joints_cap * MAX_AXIS_CONSTRAINTS).max(1); let jac_cap = max_jac_floats.max(1); let nb = self.num_batches as usize; + // The region tiling above only works if every batch agrees on each + // joint's jacobian layout (offsets are a prefix sum of capacities, so + // the regions tile the buffer exactly). The equal-topology invariant + // guarantees this; make it explicit instead of padding around it. for (b, env) in per_env_builders.iter().enumerate() { assert_eq!( env.len(), @@ -269,6 +279,10 @@ impl GpuMultibodySet { ); } } + + // Padding builder: both sides marked FIXED so the GPU kernels can + // skip unused slots by sentinel check (no per-batch `num_joints` + // binding needed). let mut dummy: MbImpulseJointBuilder = bytemuck::Zeroable::zeroed(); dummy.side_a_kind = SIDE_KIND_FIXED; dummy.side_b_kind = SIDE_KIND_FIXED; @@ -298,10 +312,11 @@ impl GpuMultibodySet { self.mb_imp_joint_constraints_per_batch = cons_cap; self.mb_imp_joint_jacobians_per_batch = jac_cap; - // Flat color-groups buffer [num_batches * cols]. Envs with fewer - // colors are padded with their last prefix value so the extra - // colors are no-ops (start == end). `cols` is clamped to ≥1 so the - // buffer is always a valid non-empty binding even with no joints. + // Color-groups buffer, color-major interleaved (`color * nb + batch`). + // Envs with fewer colors are padded with their last prefix value so + // the extra colors are no-ops (start == end). `cols` is clamped to ≥1 + // so the buffer is always a valid non-empty binding even with no + // joints. let cols = global_num_colors.max(1); let mut all_color_groups = vec![0u32; cols as usize * nb]; for (b, env_cg) in per_env_color_groups.iter().enumerate() { diff --git a/src_rbd/dynamics/multibody/multibody_from_rapier.rs b/src_rbd/dynamics/multibody/multibody_from_rapier.rs index 627eb9ce..38ab4bf4 100644 --- a/src_rbd/dynamics/multibody/multibody_from_rapier.rs +++ b/src_rbd/dynamics/multibody/multibody_from_rapier.rs @@ -30,6 +30,8 @@ impl GpuMultibodySet { &RigidBodySet, )], colliders_per_batch: u32, + // Per-batch contact-constraint slot budget + // (`RbdCapacities::mb_contact_constraints_capacity`). contact_constraint_slots: u32, ) -> Self { let num_batches = environments.len() as u32; @@ -123,10 +125,12 @@ impl GpuMultibodySet { max_mb_links = max_mb_links.max(num_links); // Count maximum constraint slots this multibody could need: for - // each non-root non-kinematic joint, every free axis with a limit - // OR a motor enabled produces one constraint slot, plus an - // additional one if BOTH limit and motor are enabled on the same - // axis (rapier emits them as separate constraints). + // each non-root non-kinematic joint, every free axis with a + // limit produces one constraint slot, plus one MOTOR slot per + // free axis whether or not a motor is enabled at build time — + // `set_motors` can enable motors at runtime (the MJCF actuator + // path does exactly that every frame), and the emission kernel + // must never outgrow this baked-in segment. let max_constraints = mb .links() .enumerate() @@ -146,6 +150,7 @@ impl GpuMultibodySet { n += 1; } if (locked >> ax) & 1 == 0 { + // Reserved motor row (runtime-enableable). n += 1; } } @@ -378,6 +383,12 @@ impl GpuMultibodySet { // One length-`dofs_cap` column of `M⁻¹` per constraint slot. let cons_col_cap = cons_cap.saturating_mul(dofs_cap).max(1); + // Flat contact-constraint buffer: per-multibody segments within it are + // demand-sized every frame (count pass + offsets scan). The initial + // capacity is the configured per-batch slot budget, floored by the + // per-multibody overflow reservation; the auto-resize grows it from + // the readback of the actual demand. Each contact point produces 1 + // normal + (DIM-1) friction tangent constraint slots. let contact_cons_cap = contact_constraint_slots .saturating_mul(num_batches) .max( @@ -427,6 +438,12 @@ impl GpuMultibodySet { let dummy_info = MultibodyInfo::default(); let dummy_stat: MultibodyLinkStatic = bytemuck::Zeroable::zeroed(); let dummy_ws = make_workspace_init(); + + // The dynamics arenas (mass matrices, LU pivots, body jacobians, + // coriolis, generalized forces) are interleaved at per-multibody-region + // granularity (`BatchIndices::mb_region`): the tiling requires every + // batch to agree on each multibody's layout offsets. The equal-topology + // invariant guarantees it; make it explicit. for (b, env) in per_env_infos.iter().enumerate() { assert_eq!( env.len(), @@ -628,6 +645,8 @@ impl GpuMultibodySet { .unwrap(), dof_couplings: Tensor::vector(backend, &all_couplings, storage).unwrap(), couplings_per_batch: couplings_cap, + // The GPU buffer is indexed by global (batch-interleaved) body id; + // the host mirror stays batch-major for `link_of_body`. body_to_link: { let cap = body_to_link_cap as usize; let nb = num_batches as usize; @@ -765,6 +784,8 @@ impl GpuMultibodySet { contact_constraints_capacity: contact_cons_cap, mb_cons_demand: Tensor::vector(backend, &[0u32], storage | BufferUsages::COPY_SRC) .unwrap(), + // Zeroed: the count pass atomically accumulates into these and the + // offsets scan re-zeroes them after consuming them. mb_cons_counts: Tensor::vector( backend, vec![0u32; (mb_cap * num_batches) as usize], diff --git a/src_rbd/dynamics/multibody/multibody_set.rs b/src_rbd/dynamics/multibody/multibody_set.rs index 1ef62a20..7e59df1a 100644 --- a/src_rbd/dynamics/multibody/multibody_set.rs +++ b/src_rbd/dynamics/multibody/multibody_set.rs @@ -186,27 +186,38 @@ pub struct GpuMultibodySet { /// Snapshot of `contact_constraints` taken at the start of the step; the /// warmstart transfer matches this frame's slots against it. pub(super) old_contact_constraints: Tensor, + /// Paired per-slot jac/column arena of the contact constraints: slot `s` + /// owns `2 * dofs_per_batch` dense floats, its `Jᵀ` row followed by its + /// `M⁻¹·Jᵀ` column (same pairing as the impulse-joint jacobians arena). pub(super) contact_jac_cols: Tensor, /// Per-multibody Delassus blocks (`MAX_MB_CONTACT_CONSTRAINTS_PER_MB²` /// floats each) only allocated when the total multibody count is at most /// [`MAX_DELASSUS_MULTIBODIES`]. pub(super) contact_delassus: Option>, + /// Batch-interleaved impulse-joint builder descriptors (joint slot `i` of + /// batch `b` at `i * num_batches + b`; unused slots padded with a + /// FIXED/FIXED sentinel). pub(super) mb_imp_joint_builders: Tensor, - /// Per-batch slab of axis constraints. + /// Batch-interleaved axis constraints (batch-local slot ids live in the + /// builders). pub(super) mb_imp_joint_constraints: Tensor, - /// Per-batch flat jacobians buffer — stores `J / W·J` for both sides - /// of every axis constraint of every joint. + /// Jacobians buffer — stores `J / W·J` for both sides of every axis + /// constraint of every joint. Interleaved at per-joint-region granularity + /// (joint `j`, batch `b` at `jacobian_offset_j * nb + b * + /// jacobian_capacity_j`, dense inside); constraints hold absolute ids. pub(super) mb_imp_joint_jacobians: Tensor, - /// Capacities (per-batch strides) for the impulse-joint slabs above. - /// Mirrored into `BatchIndices` via [`Self::fill_batch_indices`]. + /// Per-batch slot counts for the interleaved impulse-joint buffers above + /// (`mb_imp_joints_per_batch` is mirrored into `BatchIndices` as the flat + /// sweeps' loop bound; the others only size the buffers). pub(super) mb_imp_joints_per_batch: u32, pub(super) mb_imp_joint_constraints_per_batch: u32, pub(super) mb_imp_joint_jacobians_per_batch: u32, - /// Per-batch prefix-sum over the color-sorted `mb_imp_joint_builders`. - /// Built at init time by `set_impulse_joints` (greedy graph coloring). + /// Per-batch prefix-sums over the color-sorted `mb_imp_joint_builders`, + /// color-major interleaved (`color * num_batches + batch`). Built at init + /// time by `set_impulse_joints` (greedy graph coloring). pub(super) mb_imp_joint_color_groups: Tensor, /// Number of colors (per-batch stride of `mb_imp_joint_color_groups`, /// and the host color-loop trip count). CPU mirror. @@ -226,9 +237,18 @@ pub struct GpuMultibodySet { /// mirror). Stored so `RbdState` can rebuild its `BatchIndices` when caps change. pub(super) joint_constraints_per_batch: u32, pub(super) joint_constraint_columns_per_batch: u32, + /// Total capacity (in slots) of the flat contact-constraint buffer; the + /// jacobian/column arenas hold `dofs_per_batch` floats per slot. pub(super) contact_constraints_capacity: u32, + /// Total contact-constraint slot demand of the last stepped frame, + /// written by `gpu_mb_cons_offsets_scan` and read back by the auto-resize. pub(super) mb_cons_demand: Tensor, + /// Per-(multibody, batch) contact-constraint slot demand, accumulated by + /// the flat count pass and consumed (then re-zeroed) by the offsets scan. + /// Indexed like `multibody_info` (interleaved). pub(super) mb_cons_counts: Tensor, + /// Per-(multibody, batch) contact-index entry counts; after the scan they + /// double as the scatter pass's write cursors. pub(super) mb_index_counts: Tensor, /// Number of solver iterations to run on `joint_constraints` per `step()`. @@ -647,6 +667,9 @@ impl GpuMultibodySet { set: &crate::rapier::dynamics::MultibodyJointSet, bodies: &crate::rapier::dynamics::RigidBodySet, ) -> Result<(), GpuBackendError> { + // The mirror shares the GPU buffer's batch-interleaved layout: link + // slot `i` of env `e` lives at `i * num_batches + e` (same convention + // as `set_motors`). let nb = self.num_batches as usize; let mut offset = 0usize; for mb in set.multibodies() { @@ -666,6 +689,9 @@ impl GpuMultibodySet { if link_idx == 0 && !root_is_dynamic { data.locked_axes = 0x3f; } + // Motor `impulse` is solver state, not configuration: keep the + // accumulated values so retargeting does not drop the servos' + // warmstart (mirrors `set_motors`). for axis in 0..6 { data.motors[axis].impulse = entry.data.motors[axis].impulse; } @@ -774,9 +800,14 @@ impl GpuMultibodySet { self.warmstart_coefficient } + /// Total slot capacity of [`Self::contact_constraints`]. pub fn contact_constraints_capacity(&self) -> u32 { self.contact_constraints_capacity } + + /// Total contact-constraint slot demand of the last stepped frame (GPU + /// buffer, read back by the auto-resize). + /// Debug: the body-id → (multibody, link) lookup buffer. pub fn body_to_link(&self) -> &Tensor<[u32; 2]> { &self.body_to_link } @@ -784,11 +815,17 @@ impl GpuMultibodySet { pub fn mb_cons_demand(&self) -> &Tensor { &self.mb_cons_demand } + + /// Smallest flat contact-constraint capacity that honors every + /// multibody's overflow reservation (see `MB_CONS_SLOT_RESERVE`). pub(crate) fn min_contact_slab_capacity(&self) -> u32 { (self.num_active_multibodies * self.num_batches) .saturating_mul(crate::shaders::dynamics::MB_CONS_SLOT_RESERVE) .max(64) } + + /// Debug: per (multibody, batch) `(start, count, old_start, old_count)` of + /// the dynamic contact-constraint segments, plus the last total demand. pub async fn debug_cons_layout( &self, backend: &GpuBackend, @@ -817,6 +854,10 @@ impl GpuMultibodySet { (out, demand.first().copied().unwrap_or(0)) } + /// Reallocates the flat contact-constraint buffer (and its parallel + /// jacobian / column arenas and previous-frame copy) at `new_capacity` + /// slots. The previous-frame copy is zeroed, so the warmstart transfer + /// skips one frame after a resize. pub(crate) fn resize_contact_slabs(&mut self, backend: &GpuBackend, new_capacity: u32) { use khal::BufferUsages; let storage = BufferUsages::STORAGE | BufferUsages::COPY_DST; @@ -1104,6 +1145,9 @@ impl GpuMultibodySet { .unwrap_or([u32::MAX; 2]) } + /// Paired per-slot jac/column arena of the contact constraints: slot `s` + /// (laid out like [`Self::contact_constraints`]) owns `2 * dofs_per_batch` + /// dense floats — its `Jᵀ` row followed by its `M⁻¹·Jᵀ` column. pub fn contact_jac_cols(&self) -> &Tensor { &self.contact_jac_cols } diff --git a/src_rbd/dynamics/multibody/multibody_solver.rs b/src_rbd/dynamics/multibody/multibody_solver.rs index aa2b9220..90fee1af 100644 --- a/src_rbd/dynamics/multibody/multibody_solver.rs +++ b/src_rbd/dynamics/multibody/multibody_solver.rs @@ -68,9 +68,17 @@ pub struct GpuMultibodySolver { solve_contacts_delassus: GpuMbSolveContactsDelassus, /// Snapshot the contact impulses once per frame, for the cross-frame match. snapshot_contact_warmstart: GpuMbSnapshotContactWarmstart, + /// Saves the previous frame's constraint-segment bounds before the new + /// layout is computed (the warmstart transfer matches against them). save_prev_cons_bounds: GpuMbSavePrevConsBounds, + /// Predicts the per-(multibody, batch) contact-constraint slot and index + /// counts (one flat sweep over the contacts). count_contact_constraints: GpuMbCountContactConstraints, + /// Turns the predictions into dynamic segments of the flat constraint + /// buffer and of the contact→multibody index (and publishes the total + /// demand for the auto-resize). cons_offsets_scan: GpuMbConsOffsetsScan, + /// Fills the contact→multibody index segments laid out by the scan. scatter_contact_index: GpuMbScatterContactIndex, /// Carry the snapshotted impulses over to this frame's matching contacts. transfer_contact_warmstart: GpuMbTransferContactWarmstart, @@ -102,9 +110,16 @@ pub struct MultibodySolverArgs<'a> { pub collider_world_poses: &'a Tensor, /// Free-body world mass properties (read by `init_contact_constraints`). pub mprops: &'a Tensor, + /// Flat contact manifold list (filled by narrow-phase; positional slots). pub contacts: &'a Tensor, + /// Clamped per-frame list totals (see `gpu_contact_plan`); `[PLAN_BOUND]` + /// bounds the flat contact sweeps. pub contact_plan: &'a Tensor, + /// Flat dispatch grid over the whole contacts range (written by + /// `gpu_contact_plan`). pub contacts_indirect: &'a Tensor<[u32; 3]>, + /// Contact→multibody index (per-multibody segments of contact slots), + /// rebuilt once per step by [`GpuMultibodySolver::layout_contact_constraints`]. pub mb_contact_index: &'a mut Tensor, /// Free-body solver velocities (updated in place by `solve_contact_constraints`). pub solver_vels: &'a mut Tensor, @@ -154,9 +169,9 @@ impl GpuMultibodySolver { if mb.is_empty() { return Ok(()); } - // Snapshot the contact slab this frame's build will overwrite, so the - // warmstart transfer can still match against it. - // Flat (slot, multibody, batch) grid. + // Snapshot the contact-constraint buffer (and its segment bounds) that + // this frame's build will overwrite, so the warmstart transfer can + // still match against it. { let mut pass = encoder.begin_pass("[RBD] mbi/snapshot", timestamps.as_deref_mut()); self.save_prev_cons_bounds.call( @@ -177,6 +192,12 @@ impl GpuMultibodySolver { self.compute_dynamics(&mut pass, mb, args) } + /// Lay out this frame's dynamic contact-constraint segments and build the + /// contact→multibody index, once per step after the narrow phase: + /// predict the per-(multibody, batch) slot / index counts (one flat sweep + /// over the contacts), prefix-scan them into segment starts (also + /// publishing the total demand for the auto-resize readback), then + /// scatter each contact's slot into its owner's index segment. pub fn layout_contact_constraints( &self, pass: &mut GpuPass, @@ -578,6 +599,7 @@ impl GpuMultibodySolver { // Multibody-touching impulse joints — generic (rb-mb / mb-mb) // constraints. if mb.mb_imp_joints_per_batch > 0 { + // Flat 1-D sweep over the interleaved joint slots. let imp_dispatch = [mb.mb_imp_joints_per_batch * mb.num_batches, 1, 1]; self.update_impulse_joint_constraints.call( pass, @@ -614,8 +636,9 @@ impl GpuMultibodySolver { for c in 0..mb.mb_imp_joint_num_colors as usize { self.solve_impulse_joint_constraints.call( pass, - // One workgroup (MB_LU_LANES threads) per joint; thread - // count = joints-in-largest-color × workgroup size. + // One workgroup (MB_LU_LANES threads) per (joint, batch), + // batch-fastest; thread count = joints-in-largest-color × + // batches × workgroup size. [ mb.mb_imp_joint_max_color_group_len * mb.num_batches * MB_LU_LANES, 1, @@ -692,6 +715,7 @@ impl GpuMultibodySolver { let solve_dispatch = [mb.multibodies_per_batch * MB_LU_LANES, mb.num_batches, 1]; self.dispatch_solve(pass, mb, args, solve_dispatch, 0)?; if mb.mb_imp_joints_per_batch > 0 { + // Flat 1-D sweep over the interleaved joint slots. let imp_dispatch = [mb.mb_imp_joints_per_batch * mb.num_batches, 1, 1]; self.remove_impulse_joint_constraint_bias.call( pass, @@ -705,8 +729,9 @@ impl GpuMultibodySolver { for c in 0..mb.mb_imp_joint_num_colors as usize { self.solve_impulse_joint_constraints.call( pass, - // One workgroup (MB_LU_LANES threads) per joint; thread - // count = joints-in-largest-color × workgroup size. + // One workgroup (MB_LU_LANES threads) per (joint, batch), + // batch-fastest; thread count = joints-in-largest-color × + // batches × workgroup size. [ mb.mb_imp_joint_max_color_group_len * mb.num_batches * MB_LU_LANES, 1, diff --git a/src_rbd/dynamics/solver.rs b/src_rbd/dynamics/solver.rs index 26a24a4f..06076903 100644 --- a/src_rbd/dynamics/solver.rs +++ b/src_rbd/dynamics/solver.rs @@ -79,9 +79,14 @@ pub struct SolverArgs<'a> { pub num_colliders: u32, /// Contact manifolds generated by narrow-phase. pub contacts: &'a Tensor, + /// Clamped per-frame list totals (see `gpu_contact_plan`); `[PLAN_BOUND]` + /// bounds the flat contact sweeps. pub contact_plan: &'a Tensor, + /// Contact→multibody index storage, rebuilt each step by + /// `GpuMultibodySolver::layout_contact_constraints`. #[cfg(feature = "dim3")] pub mb_contact_index: &'a mut Tensor, + /// Flat dispatch grid over the whole contacts range. pub contacts_len_indirect: &'a Tensor<[u32; 3]>, /// Solver constraints (output from constraint initialization). pub constraints: &'a mut Tensor, @@ -123,7 +128,10 @@ pub struct SolverArgs<'a> { /// All constraints of all the bodies part of the same multibody are in the same list associated /// to the multibody’s root. pub body_constraint_ids: &'a mut Tensor, + /// The `(color, batch)` bucket buffer (color-major, post-scatter exclusive + /// ends): bucket `k` of `color_sorted_ids` spans `[buckets[k-1], buckets[k])`. pub color_buckets: &'a Tensor, + /// Constraint ids bucket-sorted by `(color, batch)`. pub color_sorted_ids: &'a Tensor, /// Per-color-index uniform tensors: `color_uniforms[c] == c`. pub color_uniforms: &'a [Tensor], @@ -219,6 +227,8 @@ impl GpuSolver { args.contact_plan, )?; + // One global cumulative scan (bodies of every batch): the constraint + // ranges live in one flat `body_constraint_ids` list. args.prefix_sum.launch( backend, pass, @@ -283,8 +293,9 @@ impl GpuSolver { joint_solver.init(&mut pass, &mut joint_args)?; - // Bound for the per-substep contact scans: runs after the narrow - // phase wrote `contacts_len`, before the first substep build. + // Lay out the multibody contact-constraint segments and the + // contact→multibody index, after the narrow phase and before the + // first substep build. #[cfg(feature = "dim3")] if let (Some(solver), Some(state)) = (mb_solver, mb_state.as_deref_mut()) { let mut mb_args = MultibodySolverArgs { diff --git a/src_rbd/dynamics/warmstart.rs b/src_rbd/dynamics/warmstart.rs index b7d67959..4d17016a 100644 --- a/src_rbd/dynamics/warmstart.rs +++ b/src_rbd/dynamics/warmstart.rs @@ -26,6 +26,7 @@ pub struct GpuWarmstart { /// /// Contains buffers for both old (previous frame) and new (current frame) constraint data. pub struct WarmstartArgs<'a> { + /// Clamped per-frame list totals (the flat contact sweep bound). pub contact_plan: &'a Tensor, /// Constraint counts per body from previous frame. pub old_body_constraint_counts: &'a Tensor, @@ -45,6 +46,7 @@ pub struct WarmstartArgs<'a> { /// Arguments for the coloring seed dispatch. pub struct SeedColorsArgs<'a> { + /// Clamped per-frame list totals (the flat contact sweep bound). pub contact_plan: &'a Tensor, /// Constraint counts per body from previous frame. pub old_body_constraint_counts: &'a Tensor, diff --git a/src_rbd/pipeline/insertion_removal.rs b/src_rbd/pipeline/insertion_removal.rs index 1aaef983..75623179 100644 --- a/src_rbd/pipeline/insertion_removal.rs +++ b/src_rbd/pipeline/insertion_removal.rs @@ -66,8 +66,15 @@ impl RbdState { let all_collision_groups = vec![none_groups; num_bodies_total]; let all_vels = vec![GpuVelocity::default(); num_bodies_total]; + // body_group: global body ids (free bodies map to themselves). The + // per-body buffers are batch-interleaved, so the identity mapping is + // simply each slot's own flat index. let all_body_group: Vec = (0..num_bodies_total as u32).collect(); + // collider_parent: global parent body ids, identity initially (no body + // is active yet). `append_bodies` overwrites the active prefix. The + // pair-filter key stays env-local (compared within a batch only); the + // interleaved slot `i` belongs to local collider `i / num_batches`. let all_collider_parent: Vec = (0..num_bodies_total as u32).collect(); let all_pair_filter: Vec<[u32; 2]> = (0..num_bodies_total as u32) .map(|i| [i / num_batches, 0u32]) @@ -118,15 +125,23 @@ impl RbdState { let all_collider_materials = vec![GpuColliderMaterial::default(); num_bodies_total]; let collider_materials = Tensor::vector(backend, &all_collider_materials, rw).unwrap(); + // The flat pair buffer is shared by every batch: `collisions_capacity` + // stays a per-batch sizing hint, so the initial total is `× num_batches`. let pairs_capacity = collisions_capacity * num_batches; + // Positional contact slots: pair `t` owns slot `t`, PFM entry `i` owns + // slot `pairs_total + i`, so the contacts buffer (and every + // contacts-keyed buffer) is sized `pairs + pfm = 2 ×` the pair capacity. let contacts_capacity = pairs_capacity * 2; let collision_pairs = Tensor::vector_uninit(backend, pairs_capacity, storage).unwrap(); + // Single global pair counter. let collision_pairs_len = Tensor::vector( backend, &[0u32], BufferUsages::STORAGE | BufferUsages::COPY_SRC, ) .unwrap(); + // Readback: pair count, PFM count, uncolored count (+ the multibody + // contact-constraint demand on dim3). #[cfg(feature = "dim3")] let resize_readback = GpuReadback::new(backend, 4).unwrap(); #[cfg(not(feature = "dim3"))] @@ -144,6 +159,7 @@ impl RbdState { Tensor::scalar_uninit(backend, BufferUsages::STORAGE | BufferUsages::INDIRECT).unwrap(); let pfm_pairs = Tensor::vector_uninit(backend, pairs_capacity, storage).unwrap(); let pfm_sort = PfmSortState::new(backend, pairs_capacity); + // Single global PFM work-list counter. let pfm_pairs_len = Tensor::vector( backend, &[0u32], @@ -165,10 +181,13 @@ impl RbdState { let color_buckets_stride = capacities.solver_colors + 3; let color_buckets = Tensor::vector_uninit(backend, color_buckets_stride * num_batches, storage).unwrap(); + // Written as a storage buffer by `gpu_contact_plan`, read as a uniform + // by every consumer. let contact_plan = Tensor::scalar(backend, ContactPlan::default(), storage | BufferUsages::UNIFORM) .unwrap(); let color_sorted_ids = Tensor::vector_uninit(backend, contacts_capacity, storage).unwrap(); + // Zeroed: see `RbdState::from_rapier`. let old_constraints_counts = Tensor::vector( backend, vec![0u32; (num_colliders_per_batch * num_batches) as usize], @@ -419,13 +438,20 @@ impl RbdState { // The incremental path attaches exactly one collider per body, so a // body's collider slot equals its body slot: `collider_parent` is the + // identity (each interleaved slot's own flat index). let nb = self.num_batches as usize; let parents: Vec = ((active * nb) as u32..((active + bodies.len()) * nb) as u32).collect(); + // NOTE: appended bodies are free bodies (never multibody links). The + // pair-filter key stays env-local. let pair_filters: Vec<[u32; 2]> = (0..bodies.len() * nb) .map(|i| [(active + i / nb) as u32, 0u32]) .collect(); + // The per-body buffers are batch-interleaved: entity slots `[active, + // active + n)` of EVERY batch form one contiguous range `[active * nb, + // (active + n) * nb)`, so one replicated write per buffer covers all + // environments (identical topology). fn replicate(v: &[T], nb: usize) -> Vec { let mut out = Vec::with_capacity(v.len() * nb); for &x in v { @@ -511,6 +537,9 @@ impl RbdState { BufferUsages::STORAGE | BufferUsages::COPY_SRC | BufferUsages::COPY_DST; let mut enc = backend.begin_encoding(); let mut any_copy = false; + // The per-body buffers are batch-interleaved: entity slot `k` of every + // batch is one contiguous range `[k * nb, (k + 1) * nb)`, so one + // nb-element copy relocates the slot in ALL environments. let mut staging_pose = backend.uninit_buffer::(nb, staging_usages)?; let mut staging_local_mprops = backend.uninit_buffer::(nb, staging_usages)?; @@ -522,6 +551,7 @@ impl RbdState { .uninit_buffer::(nb, staging_usages)?; let mut staging_materials = backend.uninit_buffer::(nb, staging_usages)?; + // Deferred entity-slot neutralisation writes. let mut neutralize: Vec = Vec::new(); for local in locals { @@ -534,6 +564,9 @@ impl RbdState { if local != last { let hole_global = local * nb; let last_global = last * nb; + // Relocate the last active body into the freed slot (in every + // batch at once). A staging buffer avoids same-buffer + // overlapping copies. macro_rules! relocate { ($t:expr, $staging:expr) => {{ enc.copy_buffer_to_buffer($t.buffer(), last_global, &mut $staging, 0, nb)?; diff --git a/src_rbd/pipeline/rbd_state.rs b/src_rbd/pipeline/rbd_state.rs index e435a4e8..e05a4438 100644 --- a/src_rbd/pipeline/rbd_state.rs +++ b/src_rbd/pipeline/rbd_state.rs @@ -60,6 +60,11 @@ pub struct RbdCapacities { /// /// This may or may not be automatically resized depending on [`Self::collisions_resize_policy`]. pub collisions_capacity: u32, + /// Maximum number of multibody contact-constraint slots reserved per batch + /// (each contact point involving a multibody costs one normal slot plus + /// `DIM - 1` friction slots). Resized according to + /// [`Self::collisions_resize_policy`], like the collision buffers, and never + /// shrunk below the per-multibody overflow reservation. pub mb_contact_constraints_capacity: u32, /// How internal collision buffers gets automatically resized (or not). /// @@ -165,19 +170,29 @@ pub struct RbdState { /// Per-collider friction / restitution coefficients (+ combine rules), pub(super) collider_materials: Tensor, pub(super) collision_pairs: Tensor, - /// Per-batch live collision-pair counts (length `num_batches`). + /// Single global live collision-pair count (length 1): every batch appends + /// to the same flat pair buffer. pub(super) collision_pairs_len: Tensor, + /// Non-blocking readback of `[collision_pairs_len, pfm_pairs_len, + /// uncolored]` used by /// [`RbdPipeline::auto_resize_buffers`](crate::pipeline::RbdPipeline::auto_resize_buffers) /// to grow buffers without stalling. pub(super) resize_readback: GpuReadback, pub(super) collision_pairs_indirect: Tensor<[u32; 3]>, - /// CPU-side mirrors of the dynamic batch capacities. The capacity values + /// CPU-side mirrors of the dynamic capacities. The capacity values /// live in the [`BatchIndices`] uniform; these mirrors let /// [`Self::rebuild_batch_indices`] re-emit it whenever a buffer grows. + /// Total capacity of the flat contacts buffer (and of every contacts-keyed + /// buffer). pub(super) contacts_capacity_cpu: u32, + /// Total capacity of the flat collision-pair (and PFM) buffer. pub(super) collision_pairs_capacity_cpu: u32, + /// Most recently read live collision-pair count — the total across all + /// batches, harvested by the non-blocking readback in [`RbdPipeline::auto_resize_buffers`](crate::pipeline::RbdPipeline::auto_resize_buffers). /// Surfaced in the viewer UI; lags the GPU by a frame or two like the resize. pub(super) collision_pairs_len_cpu: u32, + /// CPU mirror of the multibody contact-constraint slot demand, refreshed + /// by the same (asynchronous) readback as `collision_pairs_len_cpu`. #[cfg(feature = "dim3")] pub(super) mb_cons_demand_cpu: u32, /// Single uniform aggregating every per-batch capacity and packed-buffer @@ -186,12 +201,21 @@ pub struct RbdState { /// constituent caps changes (e.g. when the contacts buffer grows). pub(super) batch_indices: Tensor, pub(super) pfm_pairs: Tensor, + /// Single global live PFM work-list count (length 1). pub(super) pfm_pairs_len: Tensor, pub(super) pfm_pairs_indirect: Tensor<[u32; 3]>, pub(super) contacts: Tensor, + /// Flat dispatch grid over the whole contacts range, written by + /// `gpu_contact_plan`. pub(super) contacts_indirect: Tensor<[u32; 3]>, + /// Clamped per-frame list totals (see `gpu_contact_plan`): the sweep + /// bound and the positional-slot bases of the flat contacts buffer. pub(super) contact_plan: Tensor, + /// Buffers backing the per-pair PFM sort of the contact-reduction path. pub(super) pfm_sort: PfmSortState, + /// Contact→multibody index: per-(multibody, batch) segments of contact + /// slots (one entry per contact touching a multibody link), rebuilt each + /// step. Sized like `contacts` (each contact owns at most one entry). #[cfg(feature = "dim3")] pub(super) mb_contact_index: Tensor, /// Workgroup grid for the per-multibody contact-constraint dispatches: @@ -209,7 +233,11 @@ pub struct RbdState { pub(super) old_constraints_colors: Tensor, pub(super) colored: Tensor, pub(super) constraints_rands: Tensor, + /// The single `(color, batch)` bucket buffer, color-major, of length + /// `(max_colors + 3) * num_batches`: counts, then scanned exclusive + /// starts, then post-scatter exclusive ends (what the sweeps read). pub(super) color_buckets: Tensor, + /// Constraint indices bucket-sorted by `(color, batch)`. pub(super) color_sorted_ids: Tensor, pub(super) curr_color: Tensor, /// Pre-built per-color-index uniforms: `color_uniforms[c] == c`. @@ -230,6 +258,9 @@ pub struct RbdState { /// assigned the same color. pub(super) body_group: Tensor, pub(super) prefix_sum_workspace: PrefixSumWorkspace, + /// Separate workspace for the color-bucket prefix scan (different length + /// than the body-count scan, so sharing one workspace would thrash its + /// cached sizing). pub(super) bucket_prefix_workspace: PrefixSumWorkspace, /// Maximum number of constraint colors the solver will iterate. pub(super) max_colors: u32, @@ -330,8 +361,8 @@ impl RbdState { &mut self.body_poses } - /// Live collision-pair count (batch 0) most recently harvested by the - /// non-blocking readback in [`RbdPipeline::auto_resize_buffers`](crate::pipeline::RbdPipeline::auto_resize_buffers). Lags the GPU by a + /// Live collision-pair count (total across all batches) most recently + /// harvested by the non-blocking readback in [`RbdPipeline::auto_resize_buffers`](crate::pipeline::RbdPipeline::auto_resize_buffers). Lags the GPU by a /// frame or two; `0` until the first readback completes. pub fn collision_pairs_len(&self) -> u32 { self.collision_pairs_len_cpu @@ -342,7 +373,7 @@ impl RbdState { &self.collision_pairs } - /// GPU buffer holding the per-batch collision-pair counts. Unlike + /// GPU buffer holding the single global collision-pair count. 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 { @@ -406,10 +437,18 @@ impl RbdState { pub fn multibodies_mut(&mut self) -> &mut crate::dynamics::GpuMultibodySet { &mut self.multibodies } + + /// Last known multibody contact-constraint slot demand (each contact point + /// costs one normal + `DIM - 1` friction slots). Refreshed by the same + /// asynchronous readback as [`Self::collision_pairs_len`], so it lags a + /// frame or two behind the GPU. #[cfg(feature = "dim3")] pub fn mb_contact_constraints_len(&self) -> u32 { self.mb_cons_demand_cpu } + + /// Current capacity (in slots) of the flat multibody contact-constraint + /// buffer (see [`RbdCapacities::mb_contact_constraints_capacity`]). #[cfg(feature = "dim3")] pub fn mb_contact_constraints_capacity(&self) -> u32 { self.multibodies.contact_constraints_capacity() @@ -620,6 +659,7 @@ pub struct RbdSnapshot { #[cfg(feature = "dim3")] impl RbdSnapshot { + /// Debug/test accessor: body `body_id`'s snapshotted (dense per-env) pose. pub fn debug_body_pose(&self, body_id: usize) -> Pose { self.body_poses[body_id] } @@ -657,6 +697,8 @@ impl RbdState { .slow_read_buffer(self.vels.buffer(), &mut all_vels) .await .unwrap(); + // The live buffers are batch-interleaved; a snapshot holds environment + // 0's state as a dense per-env array. let bps = all_poses.len() / nb; let body_poses = (0..bps).map(|i| all_poses[i * nb]).collect(); let vs = all_vels.len() / nb; @@ -670,6 +712,10 @@ impl RbdState { } /// Resets env `dst_env` from a CPU snapshot using `write_buffer` only. + /// + /// The per-body buffers are batch-interleaved, so this issues one strided + /// write per body: fine for the documented slow path, but prefer + /// [`Self::reset_envs_from_templates`] in reset loops. pub fn reset_env_from_snapshot( &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 59ab3d9a..5febe240 100644 --- a/src_rbd/pipeline/rbd_state_from_rapier.rs +++ b/src_rbd/pipeline/rbd_state_from_rapier.rs @@ -564,9 +564,10 @@ impl RbdState { // when computing constraint adjacency, so contacts touching different // bodies of the same multibody correctly conflict and never share a // color. - // `body_group` stores PER-BATCH local indices (so a kernel can use the - // same `Slice(buf, colliders_start)` pattern as for `body_constraint_*` - // and just index by `group_local`). + // `body_group` stores global group ids (indexed by global body id): + // free bodies map to themselves, every link of a multibody maps to the + // multibody root's global body id. Built with LOCAL values batch-major + // here; globalized and interleaved with the other per-body vecs below. let mut all_body_group: Vec = Vec::with_capacity(max_colliders * num_batches as usize); for _batch_idx in 0..num_batches as usize { for b in 0..max_colliders { @@ -594,6 +595,13 @@ impl RbdState { let num_colliders_per_batch = max_colliders; let num_bodies_total = num_colliders_per_batch * num_batches as usize; + + /* + * Re-layout: the per-body/per-collider buffers are batch-interleaved + * (global id = local * num_batches + batch), so the batch-major vecs + * built above are transposed before upload, and buffers holding body + * ids switch from env-local to global values. + */ let nb = num_batches as usize; fn interleave_batches(v: &[T], nb: usize) -> Vec { let per_batch = v.len() / nb.max(1); @@ -605,6 +613,8 @@ impl RbdState { } out } + // Env-local body ids → global (interleaved) ids, still batch-major + // positioned; the position transpose follows. for (idx, v) in all_collider_parent.iter_mut().enumerate() { let batch = idx / max_colliders; *v = *v * nb as u32 + batch as u32; @@ -640,18 +650,23 @@ impl RbdState { let pair_filter = Tensor::vector(backend, &all_pair_filter, storage).unwrap(); let collider_materials = Tensor::vector(backend, &all_collider_materials, storage).unwrap(); + // The flat pair buffer is shared by every batch: `collisions_capacity` + // stays a per-batch sizing hint, so the initial total is `× num_batches`. let collision_pairs = Tensor::vector_uninit( backend, capacities.collisions_capacity * num_batches, storage, ) .unwrap(); + // Single global pair counter. let collision_pairs_len = Tensor::vector( backend, &[0u32], BufferUsages::STORAGE | BufferUsages::COPY_SRC, ) .unwrap(); + // Readback: pair count, PFM count, uncolored count (+ the multibody + // contact-constraint demand on dim3). #[cfg(feature = "dim3")] let resize_readback = GpuReadback::new(backend, 4).unwrap(); #[cfg(not(feature = "dim3"))] @@ -659,6 +674,9 @@ impl RbdState { let collision_pairs_indirect = Tensor::scalar_uninit(backend, BufferUsages::STORAGE | BufferUsages::INDIRECT).unwrap(); + // Positional contact slots: pair `t` owns slot `t`, PFM entry `i` owns + // slot `pairs_total + i`, so the contacts buffer (and every + // contacts-keyed buffer) is sized `pairs + pfm = 2 ×` the pair capacity. let pairs_capacity = capacities.collisions_capacity * num_batches; let contacts_capacity = pairs_capacity * 2; let contacts = Tensor::vector_uninit(backend, contacts_capacity, storage).unwrap(); @@ -676,6 +694,7 @@ impl RbdState { storage, ) .unwrap(); + // Single global PFM work-list counter. let pfm_pairs_len = Tensor::vector( backend, &[0u32], @@ -697,11 +716,18 @@ impl RbdState { let color_buckets_stride = capacities.solver_colors + 3; let color_buckets = Tensor::vector_uninit(backend, color_buckets_stride * num_batches, storage).unwrap(); + // Clamped per-frame list totals (see `gpu_contact_plan`). + // Zero-initialized: a resize readback can race the first step. + // Written as a storage buffer by `gpu_contact_plan`, read as a uniform + // by every consumer. let contact_plan = Tensor::scalar(backend, ContactPlan::default(), storage | BufferUsages::UNIFORM) .unwrap(); let pfm_sort = PfmSortState::new(backend, pairs_capacity); let color_sorted_ids = Tensor::vector_uninit(backend, contacts_capacity, storage).unwrap(); + // Zeroed (not uninit): the first frame's warmstart transfer walks the + // "old" counts before any step has written them; zero counts = empty + // ranges. let old_constraints_counts = Tensor::vector( backend, vec![0u32; (num_colliders_per_batch as u32 * num_batches) as usize], diff --git a/src_rbd/pipeline/rbd_step.rs b/src_rbd/pipeline/rbd_step.rs index 30203418..666623ae 100644 --- a/src_rbd/pipeline/rbd_step.rs +++ b/src_rbd/pipeline/rbd_step.rs @@ -276,6 +276,8 @@ impl RbdPipeline { let readback_enabled = state.capacities.solver_colors_resize_policy != RbdResizePolicy::Fixed || state.capacities.collisions_resize_policy != RbdResizePolicy::Fixed; + // Estimated pairs per batch: the counter (and the capacity fallback) + // are totals over the whole flat pair buffer. let est_pairs = if readback_enabled { state.collision_pairs_len_cpu.div_ceil(state.num_batches) } else { @@ -609,12 +611,17 @@ impl RbdPipeline { != RbdResizePolicy::Fixed || state.capacities.collisions_resize_policy != RbdResizePolicy::Fixed; + // The readback holds `[total collision-pair count, total PFM-pair + // count, uncolored count]` (+ the multibody contact-constraint demand + // on dim3). #[cfg(feature = "dim3")] let mut counts = [0u32; 4]; #[cfg(not(feature = "dim3"))] let mut counts = [0u32; 3]; if state.resize_readback.try_take(backend, &mut counts) { // TODO: make the coloring update optional (and pre-configurable) too? + // The flat pair and PFM work-lists share one buffer capacity, so + // whichever is larger drives that resize. let pairs_len = counts[0].max(counts[1]); let coloring_converged = counts[2]; state.collision_pairs_len_cpu = counts[0]; @@ -627,6 +634,13 @@ impl RbdPipeline { != RbdResizePolicy::Fixed && coloring_converged == 0 && !state.rb_contacts_inert; + + // Decide every resize up front, then drain the GPU once before + // applying any of them: `rebuild_batch_indices` rewrites the shared + // uniform, and on Metal a `write_buffer` is observed by every + // still-queued submission — an in-flight step would read the new + // capacities while bound to the old (smaller) buffers and write out + // of bounds. Resizes are rare, so the stall is negligible. let total_capacity = state.collision_pairs_capacity_cpu; let safe_total = pairs_len.saturating_add(pairs_len / 4); let new_total = pairs_len @@ -670,6 +684,8 @@ impl RbdPipeline { if grow_colors { state.max_colors += 5; + // The color-bucket buffer is strided by `max_colors + 3`: + // regrow it and update the stride in `BatchIndices`. let storage: BufferUsages = BufferUsages::STORAGE | BufferUsages::COPY_SRC; let stride = state.max_colors + 3; let nb = state.num_batches; @@ -679,12 +695,23 @@ impl RbdPipeline { let storage: BufferUsages = BufferUsages::STORAGE | BufferUsages::COPY_SRC; + // Flat pair / PFM buffers: sized by the TOTAL demand across all + // batches (the whole point of the flat layout: one hot batch no + // longer multiplies every batch's capacity). + // + // Since the auto-resize always lags a bit behind, resizes trigger + // with less than 25% padding available and grow with 50% slack; + // never below the configured per-batch floor × num_batches. if resize_pairs { state.collision_pairs = Tensor::vector_uninit(backend, new_total, storage)?; state.pfm_pairs = Tensor::vector_uninit(backend, new_total, storage)?; state.pfm_sort.resize(backend, new_total); state.collision_pairs_capacity_cpu = new_total; + // Contacts-keyed buffers follow the pair capacity: contact + // slots are positional (pair `t` owns slot `t`, PFM entry `i` + // owns slot `pairs_total + i`), so their capacity is always + // `2 ×` the pair/PFM capacity. let new_contacts = new_total * 2; state.contacts = Tensor::vector_uninit(backend, new_contacts, storage)?; #[cfg(feature = "dim3")] @@ -709,12 +736,16 @@ impl RbdPipeline { state.colored = Tensor::vector_uninit(backend, new_contacts, storage)?; state.constraints_rands = Tensor::vector_uninit(backend, new_contacts, storage)?; state.color_sorted_ids = Tensor::vector_uninit(backend, new_contacts, storage)?; + // The old counts index the old (now discarded) constraint list: + // zero them so the next warmstart transfer sees empty ranges + // instead of stale offsets into the fresh buffers. let counts_len = state.old_constraints_counts.len() as usize; state.old_constraints_counts = Tensor::vector(backend, vec![0u32; counts_len], storage)?; state.contacts_capacity_cpu = new_contacts; } + // Multibody contact-constraint slots: flat and demand-sized. #[cfg(feature = "dim3")] if resize_mb { state.multibodies.resize_contact_slabs(backend, new_mb); diff --git a/src_rbd/pipeline/test_batched_stacks.rs b/src_rbd/pipeline/test_batched_stacks.rs index b118e70b..26e995a2 100644 --- a/src_rbd/pipeline/test_batched_stacks.rs +++ b/src_rbd/pipeline/test_batched_stacks.rs @@ -1,3 +1,11 @@ +//! Headless correctness probe for the batched contact pipeline. +//! +//! Steps small box stacks across many identical environments and asserts that +//! every environment settles to the same resting configuration. Covers both +//! broad-phase paths (brute-force for tiny envs, LBVH for the large scene) and +//! the pair-buffer auto-resize. Run with +//! `cargo test -p nexus_rbd3d --features metal test_batched_stacks -- --nocapture --ignored`. + use crate::math::Pose; use crate::pipeline::{RbdCapacities, RbdPipeline, RbdResizePolicy, RbdState}; use crate::rapier::prelude::*; @@ -15,6 +23,8 @@ async fn test_backend() -> GpuBackend { } } +/// One environment: a ground cuboid plus `num_stacks` stacks of `stack_height` +/// dynamic cuboids (half-extent 0.2). fn build_env(num_stacks: usize, stack_height: usize) -> (RigidBodySet, ColliderSet) { let mut bodies = RigidBodySet::new(); let mut colliders = ColliderSet::new(); @@ -43,6 +53,8 @@ fn build_env(num_stacks: usize, stack_height: usize) -> (RigidBodySet, ColliderS (bodies, colliders) } +/// Steps `num_envs` copies of the scene and checks every dynamic box settled +/// at a plausible stack height, identically across environments. async fn run_case(num_envs: u32, num_stacks: usize, stack_height: usize, collisions_capacity: u32) { let backend = test_backend().await; @@ -77,6 +89,10 @@ async fn run_case(num_envs: u32, num_stacks: usize, stack_height: usize, collisi .unwrap(); let nb = num_envs as usize; let boxes_per_env = num_stacks * stack_height; + + // Body slot 0 is the ground; slots 1..=boxes_per_env are the boxes in + // insertion order (stack by stack, bottom to top). Per-body buffers are + // batch-interleaved: `global = local * num_envs + env`. for env in 0..num_envs as usize { for b in 0..boxes_per_env { let pose = poses[(1 + b) * nb + env]; @@ -92,6 +108,9 @@ async fn run_case(num_envs: u32, num_stacks: usize, stack_height: usize, collisi (y - expected_y).abs() < 0.1, "env {env} box {b}: y = {y}, expected ~{expected_y}" ); + + // Identical topology + identical inputs: every environment must + // settle to (approximately) the same configuration as env 0. if env > 0 { let ref_pose = poses[(1 + b) * nb]; let d = (pose.translation - ref_pose.translation).length(); @@ -103,36 +122,54 @@ async fn run_case(num_envs: u32, num_stacks: usize, stack_height: usize, collisi "OK: envs={num_envs} stacks={num_stacks} height={stack_height} cap={collisions_capacity}" ); } + +// Split into separately runnable cases (smallest first) so a regression can +// be bisected one GPU workload at a time. + +/// One tiny env (brute-force broad-phase path), ample capacity. #[futures_test::test] #[serial_test::serial] #[ignore] async fn test_stacks_1_tiny() { run_case(1, 1, 4, 64).await; } + +/// Tiny envs, many batches, no overflow. #[futures_test::test] #[serial_test::serial] #[ignore] async fn test_stacks_2_batched() { run_case(64, 1, 4, 64).await; } + +/// Large single env: LBVH broad-phase path (> 64 colliders). #[futures_test::test] #[serial_test::serial] #[ignore] async fn test_stacks_3_lbvh() { run_case(1, 32, 4, 1024).await; } + +/// A few large envs: LBVH path with batching. #[futures_test::test] #[serial_test::serial] #[ignore] async fn test_stacks_4_lbvh_batched() { run_case(4, 32, 4, 1024).await; } + +/// Overflow stress: the tiny pair capacity forces counter overflow and the +/// flat-buffer auto-resize (the clamped kernels must never walk past capacity +/// while the resize catches up). #[futures_test::test] #[serial_test::serial] #[ignore] async fn test_stacks_5_overflow_resize() { run_case(64, 1, 4, 8).await; } + +/// A revolute-joint pendulum across a few envs: covers the free-body +/// impulse-joint solver path. #[futures_test::test] #[serial_test::serial] #[ignore] @@ -146,6 +183,7 @@ async fn test_stacks_8_impulse_joint() { let mut joints = ImpulseJointSet::new(); let anchor = bodies.insert(RigidBodyBuilder::fixed().translation(Vec3::new(0.0, 2.0, 0.0))); colliders.insert_with_parent(ColliderBuilder::ball(0.1), anchor, &mut bodies); + // Bob starts horizontal; swings down to hang 1m below the anchor. let bob = bodies.insert(RigidBodyBuilder::dynamic().translation(Vec3::new(1.0, 2.0, 0.0))); colliders.insert_with_parent(ColliderBuilder::ball(0.1), bob, &mut bodies); let joint = RevoluteJointBuilder::new(Vec3::Z) @@ -180,9 +218,11 @@ async fn test_stacks_8_impulse_joint() { .unwrap(); let nb = num_envs as usize; for env in 0..nb { + // Body slot 1 is the bob (slot 0 = anchor). let p = poses[nb + env].translation; assert!(p.is_finite(), "env {env}: non-finite bob pose {p:?}"); let d = (p - Vec3::new(0.0, 2.0, 0.0)).length(); + // The joint must hold the bob at ~1m from the anchor while it swings. assert!( (0.9..1.1).contains(&d), "env {env}: bob at distance {d} from anchor, expected ~1" @@ -190,6 +230,10 @@ async fn test_stacks_8_impulse_joint() { } println!("OK: impulse-joint pendulum holds, envs={num_envs}"); } + +/// Template-based environment reset: steps, publishes env 0 as a template, +/// keeps stepping, then restores one env from the template (covers the +/// GPU-resident env-reset path). #[futures_test::test] #[serial_test::serial] #[ignore] @@ -224,6 +268,7 @@ async fn test_stacks_9_env_reset() { for _ in 0..50 { pipeline.step(&backend, &mut state, None).unwrap(); } + // Restore env 1 from the template (no teleport offset, zero DoF vels). state.reset_envs_from_templates(&backend, &[(1, 0)], &[crate::math::Vector::ZERO], &[]); backend.synchronize().unwrap(); @@ -243,10 +288,15 @@ async fn test_stacks_9_env_reset() { } println!("OK: env reset restores the template, envs={num_envs}"); } + +/// Test-only accessor: body `b`'s translation in a snapshot (dense per-env). #[cfg(feature = "dim3")] fn snap_pose(snap: &crate::pipeline::RbdSnapshot, b: usize) -> Vec3 { snap.debug_body_pose(b).translation } + +/// Balls and capsules dropped on the ground across a few envs: covers the +/// analytic ball paths and the deferred PFM (GJK/EPA) narrow-phase path. #[futures_test::test] #[serial_test::serial] #[ignore] @@ -314,6 +364,7 @@ async fn test_stacks_7_pfm_shapes() { let pose = poses[(1 + b) * nb + env]; assert!(pose.translation.is_finite()); let y = pose.translation.y; + // Capsules rest between 0.1 (lying) and 0.25 (upright); balls at 0.2. assert!( (0.05..0.4).contains(&y), "env {env} body {b}: y = {y}, expected resting on the ground" @@ -323,6 +374,8 @@ async fn test_stacks_7_pfm_shapes() { println!("OK: pfm shapes rest, envs={num_envs}"); } +/// One environment: ground plus a free-floating 3-link multibody chain that +/// falls onto it (covers the multibody contact-constraint path). fn build_mb_env(ball_colliders: bool) -> (RigidBodySet, ColliderSet, MultibodyJointSet) { let mut bodies = RigidBodySet::new(); let mut colliders = ColliderSet::new(); @@ -335,6 +388,10 @@ fn build_mb_env(ball_colliders: bool) -> (RigidBodySet, ColliderSet, MultibodyJo &mut bodies, ); + // Horizontal chain 1m above the ground; links 0.4 long, 0.1 thick. + // `ball_colliders` swaps the link cuboids for balls: their 1-point + // manifolds cover the demand-exact constraint segments (a 4-point cuboid + // manifold makes any conservative max-points bound exact by accident). let mut prev = None; for i in 0..3 { let x = i as f32 * 0.5; @@ -356,6 +413,9 @@ fn build_mb_env(ball_colliders: bool) -> (RigidBodySet, ColliderSet, MultibodyJo (bodies, colliders, mb_joints) } + +/// Multibody chain dropped on the ground across a few identical envs: links +/// must come to rest on (not through, not far above) the ground plane. #[futures_test::test] #[serial_test::serial] #[ignore] @@ -432,6 +492,8 @@ async fn test_stacks_6_multibody() { .await .unwrap(); let nb = num_envs as usize; + + // Body slot 0 is the ground; slots 1..=3 the chain links. for env in 0..num_envs as usize { for l in 0..3 { let pose = poses[(1 + l) * nb + env]; @@ -441,6 +503,8 @@ async fn test_stacks_6_multibody() { pose.translation ); let y = pose.translation.y; + // Half-thickness 0.1: resting links sit at ~0.1, with slack for + // chain articulation. Above 0.5 = floating, below 0 = tunnelled. assert!( (0.0..0.5).contains(&y), "env {env} link {l}: y = {y}, expected resting near 0.1" @@ -449,6 +513,13 @@ async fn test_stacks_6_multibody() { } println!("OK: multibody chain rests, envs={num_envs}"); } + +/// Same as `test_stacks_6_multibody` but with BALL link colliders: every +/// foot-ground manifold has a single point, so each multibody's demand-sized +/// constraint segment is much smaller than a max-points-per-manifold bound +/// would predict. Regression test for the emission pass truncating trailing +/// contacts of demand-exact segments (menagerie robots with ball/capsule feet +/// fell through the floor). #[futures_test::test] #[serial_test::serial] #[ignore] @@ -482,6 +553,9 @@ async fn test_stacks_11_mb_point_contacts() { .await .unwrap(); let nb = num_envs as usize; + // Body slot 0 is the ground; slots 1..=3 the chain links. Every link must + // rest ON the ground (ball radius 0.1): below 0 = its contact constraints + // were dropped and it tunnelled, far above = floating/jittering. for env in 0..nb { for l in 0..3 { let pose = poses[(1 + l) * nb + env]; @@ -499,6 +573,12 @@ async fn test_stacks_11_mb_point_contacts() { } println!("OK: 1-point-manifold multibody contacts rest, envs={num_envs}"); } + +/// Multibody-touching impulse joints (the `MbImpulseJointConstraint` path): +/// a fixed-rooted 2-link multibody chain with a free bob hung off each link +/// by a revolute impulse joint. Two joints share the multibody, so the greedy +/// coloring is forced to at least 2 colors. Envs differ (per-env anchor +/// height) so a cross-batch indexing mixup shows up as a wrong-env pose. #[futures_test::test] #[serial_test::serial] #[ignore] @@ -514,6 +594,7 @@ async fn test_stacks_10_mb_impulse_joint() { let mut mb_joints = MultibodyJointSet::new(); let y = anchor_y(e); + // Slot 0: fixed anchor. Slots 1-2: the chain links (one multibody). let anchor = bodies.insert(RigidBodyBuilder::fixed().translation(Vec3::new(0.0, y, 0.0))); colliders.insert_with_parent(ColliderBuilder::ball(0.05), anchor, &mut bodies); let mut links = Vec::new(); @@ -529,6 +610,9 @@ async fn test_stacks_10_mb_impulse_joint() { links.push(link); prev = link; } + + // Slots 3-4: free bobs, each hung 0.4 below its link by a revolute + // IMPULSE joint (side A = multibody link, side B = free body). for (i, &link) in links.iter().enumerate() { let x = 0.5 + i as f32 * 0.5; let bob = @@ -572,6 +656,8 @@ async fn test_stacks_10_mb_impulse_joint() { let nb = num_envs as usize; for env in 0..nb { let at = |slot: usize| poses[slot * nb + env]; + // The fixed anchor must still sit at ITS env's height (a cross-batch + // mixup would show another env's value here). let ay = at(0).translation.y; assert!( (ay - anchor_y(env as u32)).abs() < 1.0e-4, @@ -585,6 +671,8 @@ async fn test_stacks_10_mb_impulse_joint() { "env {env} slot {slot}: non-finite pose {p:?}" ); } + // Each bob's joint anchor (0.4 above the bob) must coincide with its + // link's center: the impulse joint holds under gravity. for (link_slot, bob_slot) in [(1usize, 3usize), (2, 4)] { let link = at(link_slot); let bob = at(bob_slot); @@ -599,6 +687,10 @@ async fn test_stacks_10_mb_impulse_joint() { println!("OK: multibody impulse joints hold, envs={num_envs}"); } +/// A flat trimesh floor plus a wide box spanning many floor triangles (its +/// per-triangle manifolds form one flat patch) carrying a small stack, with +/// contact reduction enabled: covers the per-pair PFM sort and the per-run +/// manifold reduction. #[futures_test::test] #[serial_test::serial] #[ignore] @@ -610,6 +702,7 @@ async fn test_stacks_12_trimesh_reduction() { let mut bodies = RigidBodySet::new(); let mut colliders = ColliderSet::new(); + // Flat 8x8-cell trimesh floor at y = 0 spanning [-4, 4]^2. let ground = bodies.insert(RigidBodyBuilder::fixed()); let nsubdivs = 8; let heights = Array2::from_fn(nsubdivs + 1, nsubdivs + 1, |_, _| 0.0f32); @@ -625,9 +718,11 @@ async fn test_stacks_12_trimesh_reduction() { &mut bodies, ); + // Wide box: touches several 1x1 floor cells at once. let wide = bodies.insert(RigidBodyBuilder::dynamic().translation(Vec3::new(0.0, 0.3, 0.0))); colliders.insert_with_parent(ColliderBuilder::cuboid(1.5, 0.2, 1.5), wide, &mut bodies); + // Small stack on top of the wide box. for i in 0..3 { let b = bodies.insert(RigidBodyBuilder::dynamic().translation(Vec3::new( 0.0, @@ -667,6 +762,7 @@ async fn test_stacks_12_trimesh_reduction() { .unwrap(); let nb = num_envs as usize; for env in 0..num_envs as usize { + // Slot 0 is the floor; slot 1 the wide box; slots 2..=4 the stack. let expected = [0.2f32, 0.6, 1.0, 1.4]; for (b, expected_y) in expected.iter().enumerate() { let pose = poses[(1 + b) * nb + env]; diff --git a/src_rbd/utils/radix_sort/mod.rs b/src_rbd/utils/radix_sort/mod.rs index c39cae12..f06d9ae5 100644 --- a/src_rbd/utils/radix_sort/mod.rs +++ b/src_rbd/utils/radix_sort/mod.rs @@ -184,6 +184,8 @@ impl RadixSort { let keys = &in_keys[offset..offset + n]; // Sort by permutation: sort indices by their corresponding key. + // Stable, like the GPU radix sort (equal keys keep their order; + // the per-pair PFM grouping relies on it). indices.clear(); indices.extend(0..n as u32); indices.sort_by_key(|&i| keys[i as usize]); diff --git a/src_rbd_shaders/broad_phase/brute_force.rs b/src_rbd_shaders/broad_phase/brute_force.rs index 50f97281..75e224c2 100644 --- a/src_rbd_shaders/broad_phase/brute_force.rs +++ b/src_rbd_shaders/broad_phase/brute_force.rs @@ -37,6 +37,7 @@ pub fn gpu_bf_compute_aabbs( let poses = batch_ids.ib(batch_id, poses); let shapes = batch_ids.ib(batch_id, shapes); + // The AABB scratch stays batch-major (broad-phase internal layout). let out = (crate::broad_phase::scratch_start(batch_ids, batch_id) + i) as usize; aabbs.write( out, @@ -99,6 +100,7 @@ pub fn gpu_bf_find_pairs( return; } + // Single global counter: every batch appends to the same flat buffer. let target_pair_index = atomic_add_u32(collision_pairs_len.at_mut(0), 1); // If we exceed capacity, keep counting the pairs but don’t store any more to avoid overflow. diff --git a/src_rbd_shaders/broad_phase/lbvh.rs b/src_rbd_shaders/broad_phase/lbvh.rs index 9ac3af5d..8633b11c 100644 --- a/src_rbd_shaders/broad_phase/lbvh.rs +++ b/src_rbd_shaders/broad_phase/lbvh.rs @@ -49,7 +49,7 @@ pub struct LbvhNode { pub refit_count_or_max_subtree_index: u32, } -/// Resets the collision pairs counter. One thread per batch. +/// Resets the (single, global) collision pairs counter. #[spirv_bindgen] #[spirv(compute(threads(64)))] pub fn gpu_lbvh_reset_collision_pairs( @@ -61,6 +61,12 @@ pub fn gpu_lbvh_reset_collision_pairs( collision_pairs_len.write(i, 0); } } + +/// Writes the flat 1-D indirect grid for a kernel iterating a flat work-list +/// (collision pairs or PFM pairs): `[ceil(min(len, capacity) / 64), 1, 1]`. +/// +/// NOTE: the load must be atomic or it occasionally reads stale data (breaks +/// Windows+Nvidia+wgpu, see ). #[spirv_bindgen] #[spirv(compute(threads(1)))] pub fn gpu_flat_list_dispatch( @@ -68,6 +74,8 @@ pub fn gpu_flat_list_dispatch( #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] indirect_args: &mut [u32; 3], #[spirv(uniform, descriptor_set = 0, binding = 2)] batch_ids: &BatchIndices, ) { + // Loop shell per the `gpu_reset_completion_flag_topo_gc` rustgpu-triviality + // workaround (too-trivial kernels don't get their spirv generated). for _ in 0..1 { let total = atomic_load_u32(len.at_mut(0)).min(batch_ids.collision_pairs_capacity); *indirect_args.at_mut(0) = total.div_ceil(WORKGROUP_SIZE); @@ -542,12 +550,19 @@ pub fn gpu_lbvh_find_collision_pairs( continue; } + // Single global counter: every batch appends to the same flat + // buffer (pairs from different batches interleave freely). let target_pair_index = atomic_add_u32(collision_pairs_len.at_mut(0), 1); // NOTE: if the index is out-of-bounds (meaning the `collision_pairs` isn't // big enough), don't write. But keep traversing so we get the exact count we need // for reallocating the buffers. if target_pair_index < batch_ids.collision_pairs_capacity { + // NOTE: we only store the collider pair here, with global + // collider ids (the owning batch is recovered from the + // id downstream). The parent body ids are resolved + // lazily, at the very last moment, when the + // narrow-phase writes the `IndexedManifold` consumed // by the solver — keeping this hot buffer (and the // intermediate pfm-pair buffer) narrow, and keeping // `collider_parent` out of the broad phase entirely. @@ -667,6 +682,11 @@ pub fn prefix_len( } } +/// Start of `batch_id`'s segment in the broad-phase internal scratch buffers +/// (morton keys, sorted collider ids, tree, brute-force AABBs). These stay +/// batch-major: the radix sort works on contiguous per-batch key segments. +/// Only the shared per-collider state (poses, shapes, groups, filters) is +/// batch-interleaved. #[inline] pub fn scratch_start(batch_ids: &BatchIndices, batch_id: u32) -> u32 { batch_id * batch_ids.colliders_batch_capacity diff --git a/src_rbd_shaders/broad_phase/narrow_phase.rs b/src_rbd_shaders/broad_phase/narrow_phase.rs index 28def5f6..2cbb43c9 100644 --- a/src_rbd_shaders/broad_phase/narrow_phase.rs +++ b/src_rbd_shaders/broad_phase/narrow_phase.rs @@ -28,15 +28,29 @@ use glamx::UVec2; const WORKGROUP_SIZE: u32 = 64; +/// The clamped per-frame list totals every contacts-keyed kernel reads, +/// written once per frame by [`gpu_contact_plan`]. +/// +/// Contact slots are positional: pair `t` owns contact slot `t`, and the +/// `i`-th PFM entry (in sorted order when the sort runs) owns slot +/// `pfm_base + i`. Every slot in `[0, bound)` is (re)written each frame by +/// exactly one producer (`len = 0` when the pair yields no manifold), so no +/// zeroing pass is needed and the flat consumers can sweep the whole bound. #[derive(Copy, Clone, Default)] #[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] #[repr(C)] pub struct ContactPlan { + /// Total contact-slot bound (`pfm_base + pfm_len`): the sweep range of + /// every flat contacts-keyed kernel. pub bound: u32, + /// Clamped flat collision-pair total; also the base contact slot of the + /// PFM entries. pub pfm_base: u32, + /// Clamped flat PFM work-list total. pub pfm_len: u32, } +/// Resets the (single, global) PFM work-list counter. #[spirv_bindgen] #[spirv(compute(threads(64)))] pub fn gpu_reset_narrow_phase( @@ -49,6 +63,18 @@ pub fn gpu_reset_narrow_phase( } } +/// Publishes this frame's clamped list totals (the `contact_plan`) and every +/// grid derived from them: the flat contacts sweep grid, the PFM sweep grid, +/// the (clamped) PFM sort count, and the per-multibody contact sweep grid +/// (`[multibodies_batch_capacity, num_batches, 1]`, zero workgroups when the +/// frame cannot produce any contact). +/// +/// Runs after the deferred pass (both list counters final) and before +/// everything that consumes contact slots. Serial in one thread. +/// +/// NOTE: the counter loads must be atomic or they occasionally read stale +/// data (breaks Windows+Nvidia+wgpu, see +/// ). #[spirv_bindgen] #[spirv(compute(threads(1)))] pub fn gpu_contact_plan( @@ -64,6 +90,8 @@ pub fn gpu_contact_plan( let pairs = atomic_load_u32(collision_pairs_len.at_mut(0)).min(batch_ids.collision_pairs_capacity); let pfm = atomic_load_u32(pfm_pairs_len.at_mut(0)).min(batch_ids.collision_pairs_capacity); + // `contacts_capacity = 2 * collision_pairs_capacity` (host invariant), so + // the positional slots `[0, pairs)` and `[pairs, pairs + pfm)` always fit. let bound = pairs + pfm; contact_plan.bound = bound; @@ -87,6 +115,9 @@ pub fn gpu_contact_plan( *mb_sweep_indirect.at_mut(2) = 1; } +/// Copies each PFM entry's originating pair index into the flat sort-key +/// buffer consumed by the radix sort that groups the entries per pair (only +/// dispatched when contact reduction is enabled). #[spirv_bindgen] #[spirv(compute(threads(64)))] pub fn gpu_pfm_sort_keys( @@ -153,6 +184,12 @@ fn pool_dedup(cand: &mut [ContactPoint; 8], num: &mut usize, pt: ContactPoint, d /// default threshold, where every member is within ~5.1 degrees, but it keeps /// the choice sane when `merge_cos` is loosened. /// +/// One thread per sorted PFM entry; the entries were radix-sorted by their +/// originating pair (a prerequisite of reduction), so each pair's manifolds +/// sit in one contiguous run of contact slots and the run leader (the first +/// entry of its run) compacts the run in place. Analytic pairs emit exactly +/// one manifold and never enter the PFM list, so they need no reduction. +/// Grid: `pfm_indirect`. #[cfg(feature = "dim3")] #[spirv_bindgen] #[spirv(compute(threads(64)))] @@ -173,9 +210,11 @@ pub fn gpu_reduce_contacts( for t in StepRng::new(invocation_id.x..total as u32, num_threads) { let i = t as usize; let key = sorted_pfm_keys.read(i); + // Only the leader (first entry) of each same-pair run does the work. if i > 0 && sorted_pfm_keys.read(i - 1) == key { continue; } + // Run length: entries sharing this pair key (clamped by `total`). let mut n = 1usize; for j in (i + 1)..total { if sorted_pfm_keys.read(j) != key { @@ -184,20 +223,28 @@ pub fn gpu_reduce_contacts( n += 1; } if n <= 1 { + // Single-manifold pairs are bit-identical to the unreduced path. continue; } let mut contacts = SliceMut(contacts, base + i); + // 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]; + // PFM misses left an inert slot; nothing to merge or keep. if im.contact.len == 0 { continue; } let mut merged = false; for j in 0..w { let out = contacts[j]; + // Every entry of the run shares one collider pair (and one + // collider-A local frame): cluster on the normal alone. if 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; @@ -220,6 +267,11 @@ pub fn gpu_reduce_contacts( dedup_eps_sq, ); } + // Normal of whichever manifold holds the deepest point. rapier + // keeps the opener's normal instead, which it can afford + // because its ~5.1 degree cone makes every member equivalent; + // this degrades gracefully when `merge_cos` is loosened, and + // agrees with rapier's choice when it is not. let mut deep_out = out.contact.points_a.at(0).dist; for k in 1..na { let d = out.contact.points_a.at(k).dist; @@ -240,6 +292,7 @@ pub fn gpu_reduce_contacts( out.contact.normal_a }; let mut reduced = manifold_reduction(&cand, num as u32, normal, prediction); + // `manifold_reduction` fills points/len only. reduced.normal_a = normal; let mut kept = out; kept.contact = reduced; @@ -253,6 +306,8 @@ pub fn gpu_reduce_contacts( w += 1; } } + // The compaction leaves stale duplicates in `[w, n)`; mark them inert + // so the flat consumers (which walk the whole bound) skip them. for i in w..n { contacts[i].contact.len = 0; } @@ -262,6 +317,11 @@ pub fn gpu_reduce_contacts( /// Narrow phase, pass 1 of 2: analytic shape-shape contacts for ball / cuboid /// pairs, written straight into the `contacts` buffer. /// +/// Contact slots are positional: pair `t` owns contact slot `t` and this pass +/// writes every slot in `[0, pairs_total)` exactly once (`len = 0` when the +/// pair yields no manifold here: separated, same-body, or deferred), so the +/// flat consumers can sweep the whole bound without a zeroing pass. +/// /// The complex cases (generic convex via PFM, trimesh, polyline) are deferred /// to `gpu_narrow_phase_shape_shape_deferred`. #[spirv_bindgen] @@ -285,6 +345,9 @@ pub fn gpu_narrow_phase_shape_shape( let prediction = params.prediction_distance(); let num_threads = num_workgroups.x * WORKGROUP_SIZE; + // Flat over the single mixed-batch pair list: consecutive lanes take + // consecutive pairs regardless of which batch owns them, so warps stay + // packed even when each batch only has a handful. let total = contact_plan.pfm_base; for t in StepRng::new(invocation_id.x..total, num_threads) { @@ -292,6 +355,7 @@ pub fn gpu_narrow_phase_shape_shape( // Resolve the parent rigid-bodies here (the broad phase no longer does) // and skip pairs whose colliders share the same body. Pair ids are + // global, and `collider_parent` maps them to global body ids. let body1 = collider_parent.read(pair.colliders.x as usize); let body2 = collider_parent.read(pair.colliders.y as usize); let mut manifold = ContactManifold::default(); @@ -304,6 +368,7 @@ pub fn gpu_narrow_phase_shape_shape( let shape_ty2 = shape2.shape_type(); let pose12 = pose1.inverse() * pose2; + // Ball - Convex if shape_ty1 == SHAPE_TYPE_BALL { if shape_ty2 == SHAPE_TYPE_BALL { let ball1 = shape1.to_ball(); @@ -319,6 +384,7 @@ pub fn gpu_narrow_phase_shape_shape( } } + // Convex - Ball if shape_ty2 == SHAPE_TYPE_BALL && (shape_ty1 == SHAPE_TYPE_CUBOID || shape_ty1 == SHAPE_TYPE_CAPSULE @@ -329,6 +395,7 @@ pub fn gpu_narrow_phase_shape_shape( manifold = convex_ball(pose12, shape1, &ball2); } + // Cuboid - Cuboid if shape_ty1 == SHAPE_TYPE_CUBOID && shape_ty2 == SHAPE_TYPE_CUBOID { let cuboid1 = shape1.to_cuboid(); let cuboid2 = shape2.to_cuboid(); @@ -337,9 +404,11 @@ pub fn gpu_narrow_phase_shape_shape( } // Everything else (PFM / trimesh / polyline) is handled by the deferred + // pass; `manifold.len` stays 0 here so the pair's slot reads as inert. if manifold.len > 0 && manifold.points_a.at(0).dist < prediction { let mat1 = collider_materials.read(pair.colliders.x as usize); let mat2 = collider_materials.read(pair.colliders.y as usize); + // Contacts carry global collider/body ids. contacts.write( t as usize, IndexedManifold { @@ -352,6 +421,8 @@ pub fn gpu_narrow_phase_shape_shape( }, ); } else { + // The slot is owned by this pair either way; only its `len` gates + // every consumer, so a field write avoids the full-struct store. contacts.at_mut(t as usize).contact.len = 0; } } @@ -368,17 +439,19 @@ 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], + // Single global pair count (see `gpu_narrow_phase_shape_shape`). #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] collision_pairs_len: &mut [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)] pfm_pairs: &mut [NarrowPhasePfmPair], + // Single global PFM work-list count. #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] pfm_pairs_len: &mut [u32], #[spirv(storage_buffer, descriptor_set = 1, binding = 0)] vertices: &[PaddedVector], #[spirv(storage_buffer, descriptor_set = 1, binding = 1)] indices: &[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`). + // NOTE: the flat PFM work-list shares the collision-pair buffer's capacity + // (`collision_pairs_capacity`); both buffers are allocated the same + // total size. #[spirv(uniform, descriptor_set = 0, binding = 6)] batch_ids: &BatchIndices, #[spirv(uniform, descriptor_set = 0, binding = 7)] params: &RbdSimParams, ) { @@ -397,6 +470,8 @@ pub fn gpu_narrow_phase_shape_shape_deferred( for t in StepRng::new(invocation_id.x..total, num_threads) { let mut pfm_pairs = SliceMut(&mut *pfm_pairs, 0); let pfm_pairs_len = pfm_pairs_len.at_mut(0); + + // Pair collider ids are global; the emitted PFM pairs keep them global. let pair = collision_pairs.read(t as usize); let shape1 = shapes.at(pair.colliders.x as usize); let shape2 = shapes.at(pair.colliders.y as usize); @@ -708,10 +783,22 @@ pub struct NarrowPhasePfmPair { thickness1: f32, thickness2: f32, colliders: UVec2, + /// Index of the originating pair in the flat collision-pair list; the + /// per-pair sort key of the contact-reduction path. pair_index: u32, _padding: [u32; 3], } +/// PFM (GJK/EPA) manifold computation for the deferred work-list entries. +/// +/// Contact slots are positional: the `i`-th PFM entry owns contact slot +/// `plan.pfm_base + i` and every slot in that range is written exactly +/// once (`len = 0` on a miss or a same-body pair). +/// +/// The `pfm_order` indirection selects which entry lane `i` processes: the +/// identity permutation normally, or the pair-sorted permutation when contact +/// reduction is enabled (so each pair's manifolds land in one contiguous run +/// of slots for `gpu_reduce_contacts`). #[spirv_bindgen] #[spirv(compute(threads(64)))] // TODO PERF: pfm_pfm is very divergent. Use a smaller workgroup size? pub fn gpu_narrow_phase_pfm_pfm( @@ -766,6 +853,7 @@ pub fn gpu_narrow_phase_pfm_pfm( if manifold.len > 0 && manifold.points_a.at(0).dist < prediction { let mat1 = collider_materials.read(pair.colliders.x as usize); let mat2 = collider_materials.read(pair.colliders.y as usize); + // Contacts carry global collider/body ids. contacts.write( slot, IndexedManifold { diff --git a/src_rbd_shaders/dynamics/color_buckets.rs b/src_rbd_shaders/dynamics/color_buckets.rs index 3c22fba5..b8cff93f 100644 --- a/src_rbd_shaders/dynamics/color_buckets.rs +++ b/src_rbd_shaders/dynamics/color_buckets.rs @@ -1,5 +1,11 @@ //! Bucket-sort of contact constraints by graph-coloring color. //! +//! After the per-step (global) coloring converges, the constraint indices are +//! bucket-sorted by `(color, batch)` into `color_sorted_ids`. Buckets are laid +//! out color-major (`bucket = color * num_batches + batch`, buffer length +//! `solver_color_buckets_stride * num_batches`), so one color's constraints +//! are contiguous across every batch (per-color solver sweeps) while each +//! `(color, batch)` cell stays contiguous too (fused per-batch sweeps). use crate::broad_phase::ContactPlan; use khal_std::glamx::UVec3; @@ -11,7 +17,8 @@ use crate::utils::{BatchIndices, Slice}; const WORKGROUP_SIZE: u32 = 64; -/// Zeroes the per-batch per-color constraint counts. +/// Zeroes the `(color, batch)` bucket counts (flat 1-D grid over the whole +/// bucket buffer). #[spirv_bindgen] #[spirv(compute(threads(64)))] pub fn gpu_color_buckets_reset( @@ -24,7 +31,7 @@ pub fn gpu_color_buckets_reset( } } -/// Counts, per batch, how many constraints hold each color. +/// Counts how many constraints fall in each `(color, batch)` bucket. #[spirv_bindgen] #[spirv(compute(threads(64)))] pub fn gpu_color_buckets_count( @@ -44,6 +51,10 @@ pub fn gpu_color_buckets_count( for i in StepRng::new(invocation_id.x..total, num_threads) { let color = constraints_colors.read(i as usize); + // Color 0 (uncolored / gap slots) is never swept; colors past the + // swept range (bounded coloring didn't converge) are dropped. They + // were never solved before either. Skipping color 0 also keeps the + // stale body ids of gap slots from being dereferenced. if color != 0 && color < stride - 1 { let batch = batch_ids.collider_batch(constraints[i as usize].solver_body_a); atomic_add_u32(color_buckets.at_mut((color * nb + batch) as usize), 1); @@ -51,6 +62,9 @@ pub fn gpu_color_buckets_count( } } +/// Scatters each constraint index into its `(color, batch)` bucket. The bucket +/// buffer holds the scanned exclusive starts, used as cursors; after this pass +/// every entry is its bucket's exclusive end. #[spirv_bindgen] #[spirv(compute(threads(64)))] pub fn gpu_color_buckets_scatter( diff --git a/src_rbd_shaders/dynamics/coloring.rs b/src_rbd_shaders/dynamics/coloring.rs index 7ac713dc..be112356 100644 --- a/src_rbd_shaders/dynamics/coloring.rs +++ b/src_rbd_shaders/dynamics/coloring.rs @@ -57,8 +57,10 @@ pub fn gpu_reset_luby( if i < total { let idx = i as usize; if constraints.at(idx).len == 0 { + // Gap / inert slot: pre-colored with 0 so the Luby steps skip it. constraints_colors.write(idx, 0); } else { + // Mark as uncolored constraints_colors.write(idx, MAX_U32); } // Assign random weight @@ -192,6 +194,8 @@ pub fn gpu_reset_topo_gc( let idx = i as usize; // Color 0 is reserved for "uncolored" state constraints_colors.write(idx, 0); + // Gaps / inert slots are pre-marked colored so the topo-gc iterations + // skip them (and converge). let inert = if constraints.at(idx).len == 0 { 1 } else { 0 }; colored.write(idx, inert); } @@ -335,6 +339,7 @@ pub fn gpu_fix_conflicts_topo_gc( for constraint_i in StepRng::new(invocation_id.x..total, num_threads) { let i = constraint_i as usize; + // Gap / inert slot: its stale body ids must not be dereferenced. if constraints[i].len == 0 { continue; } diff --git a/src_rbd_shaders/dynamics/joint_constraint.rs b/src_rbd_shaders/dynamics/joint_constraint.rs index be7409f7..21f9c983 100644 --- a/src_rbd_shaders/dynamics/joint_constraint.rs +++ b/src_rbd_shaders/dynamics/joint_constraint.rs @@ -212,7 +212,7 @@ pub fn gpu_init_joint_constraints( cons.solver_vel_b = body_b; cons.im_a = local_mprops.at(body_a as usize).inv_mass; cons.im_b = local_mprops.at(body_b as usize).inv_mass; - cons.len = 0; + cons.len = 0; // Constraint elements will be filled later. } } diff --git a/src_rbd_shaders/dynamics/multibody/compute_dynamics_pre.rs b/src_rbd_shaders/dynamics/multibody/compute_dynamics_pre.rs index 5e48b898..4eec9a2b 100644 --- a/src_rbd_shaders/dynamics/multibody/compute_dynamics_pre.rs +++ b/src_rbd_shaders/dynamics/multibody/compute_dynamics_pre.rs @@ -117,8 +117,8 @@ pub fn gpu_mb_compute_dynamics_pre( .ib(batch_id, dof_state) .offset(batch_ids.dof_batch_capacity as usize + vel_base); // Armature (reflected rotor inertia) section sits right after damping, at - // `2 · dof_damping_section_offset` (= 2·N). Added to the mass-matrix diagonal - // alongside `damping·dt`, matching rapier's `update_mass_matrix`. + // `2 · dof_batch_capacity`. Added to the mass-matrix diagonal alongside + // `damping·dt`, matching rapier's `update_mass_matrix`. let armature_slice = batch_ids .ib(batch_id, dof_state) .offset(2 * batch_ids.dof_batch_capacity as usize + vel_base); @@ -680,6 +680,7 @@ fn update_body_jacobians( lane: u32, // Lanes owned by this multibody's slot (`BatchIndices::mb_pack_lanes`). lanes: u32, + // Dense base of this multibody's body-jacobians region. jac0: usize, ndofs: u32, num_links: u32, diff --git a/src_rbd_shaders/dynamics/multibody/contact_constraints.rs b/src_rbd_shaders/dynamics/multibody/contact_constraints.rs index 9535f97b..fde27b7d 100644 --- a/src_rbd_shaders/dynamics/multibody/contact_constraints.rs +++ b/src_rbd_shaders/dynamics/multibody/contact_constraints.rs @@ -35,6 +35,8 @@ use super::types::{ MultibodyLinkStatic, }; +/// Workgroup width of the flat contact sweeps below (must match the +/// `contacts_indirect` grid written by `gpu_contact_plan`). const MB_SWEEP_WG: u32 = 64; use super::utils::zero_kinematic_dofs; use super::ws_soa::{WS_LTW, WS_WORLD_COM, WsAddr, ws_pose, ws_vec}; @@ -72,6 +74,7 @@ fn orthonormal_vector(v: Vec2) -> Vec2 { #[inline] fn fill_contact_jac_row( body_jacobians: &[f32], + // Dense base of the multibody's body-jacobians region. jac0: usize, ndofs: u32, link_id: u32, @@ -123,11 +126,26 @@ fn fill_contact_jac_row( } } +/// Resolves which multibody (if any) owns contact `im`'s constraint slots, +/// from the contact's global body ids. This is the single source of truth for +/// the ownership decision: `gpu_mb_count_contact_constraints` sizes the +/// dynamic segments with it and `gpu_mb_scatter_contact_index` builds the +/// contact→multibody index with it, so the prediction always matches the +/// index the emission consumes. +/// +/// Skipped contacts (empty manifold, free-free, inter-multibody which is not +/// yet handled, disabled or same-link self contact) return `mb == u32::MAX`. +/// The self-contacts-enabled bit comes from the owner's `MultibodyInfo`. struct MbContactOwner { + /// Env-local index of the owning multibody, `u32::MAX` when skipped. mb: u32, + /// Owning batch (recovered from the interleaved global body ids). batch: u32, + /// Touched link of the owning multibody (the constraint's A side). link_a: u32, + /// Second touched link for self-contacts, `u32::MAX` otherwise. link_b: u32, + /// 1 when the multibody sits on the contact's `bodies.x` side. mb_on_first: u32, } @@ -154,17 +172,23 @@ fn mb_contact_owner( if l1[0] == u32::MAX && l2[0] == u32::MAX { return MB_OWNER_SKIP; } + // Inter-multibody contacts (each side a different multibody) are not yet + // handled. if l1[0] != u32::MAX && l2[0] != u32::MAX && l1[0] != l2[0] { return MB_OWNER_SKIP; } let is_self = l1[0] != u32::MAX && l2[0] != u32::MAX; + // Degenerate self-contact on the same link. if is_self && l1[1] == l2[1] { return MB_OWNER_SKIP; } let mb_on_first = l1[0] != u32::MAX; let owner = if mb_on_first { l1[0] } else { l2[0] }; + // Bodies are batch-interleaved, so the contact's global body ids carry + // the owning batch. let batch = batch_ids.collider_batch(im.bodies.x); let mb = multibody_info.read(batch_ids.mbi(batch, owner as usize)); + // All-locked (zero-dof) multibodies take no contact constraints. if mb.ndofs == 0 { return MB_OWNER_SKIP; } @@ -180,6 +204,16 @@ fn mb_contact_owner( } } +/// Predicts, per (multibody, batch), how many contact-constraint slots the +/// emission pass needs, and how many contacts touch each multibody. One flat +/// sweep over the contacts: each contact resolves its owner and atomically +/// bumps that owner's counters, so the pass costs one visit per contact +/// instead of one full-list scan per (multibody, batch). +/// +/// `mb_cons_counts` / `mb_index_counts` are indexed like `multibody_info` +/// (interleaved `mb * num_batches + batch`) and must be zero on entry (the +/// offsets scan re-zeroes them after consuming them). Grid: +/// `contacts_indirect`. #[spirv_bindgen] #[spirv(compute(threads(64)))] pub fn gpu_mb_count_contact_constraints( @@ -208,6 +242,16 @@ pub fn gpu_mb_count_contact_constraints( } } +/// Turns the per-(multibody, batch) predictions into dynamic segments: +/// contact-constraint segments (`MultibodyInfo::contact_constraint_start/ +/// count`, exclusive prefix cumulatively clamped to the buffer capacity so +/// the emission never overruns) and contact-index segments +/// (`contact_index_start/len`, unclamped: the index buffer is sized like the +/// contacts buffer and each contact owns at most one entry). Also publishes +/// the total slot demand for the host's auto-resize readback, and re-zeroes +/// `mb_cons_counts` (for the next frame) and `mb_index_counts` (which the +/// scatter pass reuses as its write cursors). Serial in one thread (the +/// multibody count per scene is small). #[spirv_bindgen] #[spirv(compute(threads(1)))] pub fn gpu_mb_cons_offsets_scan( @@ -218,6 +262,8 @@ pub fn gpu_mb_cons_offsets_scan( #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] mb_cons_demand: &mut [u32], #[spirv(uniform, descriptor_set = 0, binding = 4)] batch_ids: &BatchIndices, ) { + // The interleaved `multibody_info` layout makes the active multibodies of + // every batch the contiguous prefix `[0, multibodies_len * num_batches)`. let total_infos = batch_ids.multibodies_len * batch_ids.num_batches; let capacity = batch_ids.mb_contact_constraints_capacity; @@ -228,6 +274,7 @@ pub fn gpu_mb_cons_offsets_scan( demand += count; reserved_total += count.min(MB_CONS_SLOT_RESERVE); } + // NOTE: not `saturating_sub`, which rust-gpu fails to compile. #[allow(clippy::implicit_saturating_sub)] let mut extra_budget = if capacity > reserved_total { capacity - reserved_total @@ -254,12 +301,20 @@ pub fn gpu_mb_cons_offsets_scan( index_acc += mb.contact_index_len; multibody_info.write(i as usize, mb); + // Zeroed for the next frame's count pass / for the scatter cursors. mb_cons_counts.write(i as usize, 0); mb_index_counts.write(i as usize, 0); } mb_cons_demand.write(0, demand); } +/// Builds the contact→multibody index: one flat sweep over the contacts, +/// each contact appending its entry to its owner's segment (laid out by the +/// offsets scan; `mb_index_counts` was re-zeroed there and serves as the +/// per-multibody write cursors). Entry order within a segment follows the +/// atomic race; the emission's warmstart matching is key-based, so the order +/// only affects the (already nondeterministic) impulse iteration order. +/// Grid: `contacts_indirect`. #[spirv_bindgen] #[spirv(compute(threads(64)))] pub fn gpu_mb_scatter_contact_index( @@ -297,6 +352,10 @@ pub fn gpu_mb_scatter_contact_index( } } +/// Saves the current (about to be superseded) constraint-segment bounds as the +/// "previous frame" bounds the warmstart transfer matches against. Runs at +/// step start, before the new layout is computed. One thread per (multibody, +/// batch). #[spirv_bindgen] #[spirv(compute(threads(64)))] pub fn gpu_mb_save_prev_cons_bounds( @@ -319,12 +378,13 @@ pub fn gpu_mb_save_prev_cons_bounds( /// Pack the per-link world-space contact point into the constraint. /// -/// Pass 1: scans every contact in `contacts[batch]` and, for each contact -/// point touching a link of this multibody, emits a normal-direction +/// Pass 1: walks this multibody's segment of the contact→multibody index +/// (built by `gpu_mb_scatter_contact_index`) and, for each contact point +/// touching one of its links, emits a normal-direction /// `MultibodyContactConstraint` plus its friction slots. The multibody-side /// `Jᵀ` rows are assembled later, by the finalize pass. Multibody-multibody /// contacts (each side a different multibody) are not handled — such contacts -/// are skipped. +/// never enter the index. /// One 64-lane workgroup per (multibody, batch). #[spirv_bindgen] #[spirv(compute(threads(64)))] @@ -377,10 +437,17 @@ pub fn gpu_mb_init_contact_constraints( } return; } + // This multibody's dynamic segment of the flat constraint buffer, sized by + // the count pass and the offsets scan; `avail` is the (possibly clamped) + // slot budget the emission below must stay within. let cons_base = mb.contact_constraint_start as usize; let avail = mb.contact_constraint_count; let wa = WsAddr::new(mb.first_link as usize, batch_ids.num_batches, batch_id); + // This multibody's segment of the contact→multibody index. Entries were + // pre-filtered by `mb_contact_owner` (shared with the count pass, so the + // segment prediction always matches the emission below). Contact + // collider/body ids are global. let idx_base = mb.contact_index_start as usize; let n_entries = mb.contact_index_len; let mut count = 0u32; @@ -731,13 +798,20 @@ pub fn gpu_mb_init_contact_constraints( } } + // Next frame's warmstart match only scans `[old_start, old_start + + // old_count)`, so no stale-slot invalidation is needed. if lane == 0 { mb.contact_constraint_count = count; multibody_info.write(batch_ids.mbi(batch_id, mb_idx as usize), mb); } } +/// Snapshot the (flat) contact-constraint buffer into the "previous frame" +/// copy that `gpu_mb_transfer_contact_warmstart` matches against (the segment +/// bounds are saved separately by `gpu_mb_save_prev_cons_bounds`). Called once +/// per visible frame from `init_step`, before the new layout is computed. /// +/// One thread per slot. #[spirv_bindgen] #[spirv(compute(threads(64)))] pub fn gpu_mb_snapshot_contact_warmstart( @@ -788,6 +862,8 @@ pub fn gpu_mb_warmstart_contact_constraints( } let v_base = mb.first_dof as usize; let cons_base = mb.contact_constraint_start as usize; + // Paired jac/column arena: slot `g` owns `2 * dofs_stride` dense floats, + // the `J` row first, then its `M^-1*J^T` column. let dofs_stride = batch_ids.dof_batch_capacity as usize; let jc_base = cons_base * 2 * dofs_stride; @@ -817,6 +893,7 @@ pub fn gpu_mb_warmstart_contact_constraints( // Free body side (skipped for self-contacts). let is_self = cons.free_body_id == u32::MAX; if lane == 0 && !is_self { + // `free_body_id` is a global body id. let free = solver_vels.read(cons.free_body_id as usize); let mut new_free = free; new_free.linear += cons.lin_jac * (cons.free_body_im * imp); @@ -917,8 +994,8 @@ pub fn gpu_mb_finalize_contact_constraints( ); } - // 2) Copy J^T row into the column buffer (it'll be overwritten by the - // LU solve with the M⁻¹·Jᵀ result). + // 2) Copy the J^T row into the paired column half (it'll be + // overwritten by the LU solve with the M⁻¹·Jᵀ result). for i in 0..ndofs { let v = contact_jac_cols.read(jac_offset + i as usize); contact_jac_cols.write(col_offset + i as usize, v); @@ -994,6 +1071,9 @@ pub fn gpu_mb_transfer_contact_warmstart( } let cons_base = mb.contact_constraint_start as usize; + // Previous frame's segment (bounds saved at step start). The bounds guard + // covers the one frame right after a buffer resize, where they still + // describe the discarded (differently sized) buffer. let old_base = mb.old_contact_constraint_start as usize; let old_count = if mb.old_contact_constraint_start + mb.old_contact_constraint_count <= batch_ids.mb_contact_constraints_capacity @@ -1109,6 +1189,7 @@ pub fn gpu_mb_seed_contact_restitution( * dof_state.read(batch_ids.mbi(batch_id, v_base + i as usize)); } if cons.free_body_id != u32::MAX { + // `free_body_id` is a global body id. let free = solver_vels.read(cons.free_body_id as usize); j_dot_v += cons.lin_jac.dot(free.linear) + gdot(cons.ang_jac, free.angular); } diff --git a/src_rbd_shaders/dynamics/multibody/env_reset.rs b/src_rbd_shaders/dynamics/multibody/env_reset.rs index cdeaa6c8..b5dab844 100644 --- a/src_rbd_shaders/dynamics/multibody/env_reset.rs +++ b/src_rbd_shaders/dynamics/multibody/env_reset.rs @@ -17,10 +17,10 @@ use super::ws_soa::{WS_COORDS, WS_LTP, WS_LTW, WS_QUADS}; /// 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. +/// `staging_dofs` holds `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. +/// The generalized coordinates live in the workspace quads copied above. #[spirv_bindgen] #[spirv(compute(threads(64)))] pub fn gpu_mb_env_reset( @@ -43,6 +43,8 @@ pub fn gpu_mb_env_reset( let dpb = params.w; if i < lpb * WS_QUADS { + // Staging is per-env dense; the live workspace is interleaved at + // per-link-record granularity (record = `WS_QUADS` dense quads). let link = i / WS_QUADS; let q = i % WS_QUADS; links_workspace.write( @@ -116,6 +118,7 @@ pub fn gpu_mb_env_reset_batch( v.y += off.y; v.z += off.z; } + // Live workspace is interleaved at per-link-record granularity. links_workspace.write(((link * nb + env) * WS_QUADS + q) as usize, v); } } @@ -204,6 +207,8 @@ pub fn gpu_env_reset_bodies( let t = meta.y; let off = offsets.read(r as usize); + // Templates are dense per-env arrays; the live per-body buffers are + // batch-interleaved (`local * num_batches + env`). if i < bps { let mut p = templates_poses.read((t * bps + i) as usize); if body_mask.read(i as usize) != 0 { diff --git a/src_rbd_shaders/dynamics/multibody/gravity_and_lu.rs b/src_rbd_shaders/dynamics/multibody/gravity_and_lu.rs index bdcc4f7f..82caa9c9 100644 --- a/src_rbd_shaders/dynamics/multibody/gravity_and_lu.rs +++ b/src_rbd_shaders/dynamics/multibody/gravity_and_lu.rs @@ -44,6 +44,7 @@ use super::ws_soa::{ #[inline] fn apply_spring_forces( gen_forces: &mut [f32], + // Dense base of this multibody's generalized-force region. gen0: usize, stat_slice: &ISlice, links_workspace: &[Vec4], diff --git a/src_rbd_shaders/dynamics/multibody/impulse_joint_constraints/jacobians.rs b/src_rbd_shaders/dynamics/multibody/impulse_joint_constraints/jacobians.rs index 90c2f4da..89f7120b 100644 --- a/src_rbd_shaders/dynamics/multibody/impulse_joint_constraints/jacobians.rs +++ b/src_rbd_shaders/dynamics/multibody/impulse_joint_constraints/jacobians.rs @@ -14,9 +14,10 @@ use super::super::lu::LANES; use super::super::types::MultibodyInfo; use super::types::*; -// Returns the multibody side jacobian's offset / size in the per-batch -// jacobians buffer. `wj_id` is the start of the corresponding `M⁻¹·J` -// block (= `j_id + ndofs`). +// Returns the multibody side jacobian's offset in the jacobians buffer. +// Jacobian ids are absolute dense indices into the joint's own region (the +// buffer is interleaved at per-joint-region granularity, dense inside). +// `wj_id` is the start of the corresponding `M⁻¹·J` block (= `j_id + ndofs`). #[inline] pub(super) fn wj_id(j_id: u32, ndofs: u32) -> usize { (j_id + ndofs) as usize diff --git a/src_rbd_shaders/dynamics/multibody/impulse_joint_constraints/kernels.rs b/src_rbd_shaders/dynamics/multibody/impulse_joint_constraints/kernels.rs index 3592a9bc..b55c8cbd 100644 --- a/src_rbd_shaders/dynamics/multibody/impulse_joint_constraints/kernels.rs +++ b/src_rbd_shaders/dynamics/multibody/impulse_joint_constraints/kernels.rs @@ -52,6 +52,11 @@ pub fn gpu_mb_update_impulse_joint_constraints( let lock_cfm_coeff = softness.joint_cfm_coeff; let max_corr_velocity = softness.max_corr_velocity; + // Flat sweep over the interleaved builder slots (slot `t` = joint `t / nb` + // of batch `t % nb`). Iterating to `cap * nb` instead of the per-batch + // lengths lets us drop the `num_joints` storage binding — the host pads + // unused slots with `side_a_kind == SIDE_KIND_FIXED && side_b_kind == + // SIDE_KIND_FIXED` which we use as the inactive-slot sentinel below. let nb = batch_ids.num_batches; let total = batch_ids.mb_imp_joints_batch_capacity * nb; for t in StepRng::new(invocation_id.x..total, num_threads) { @@ -60,6 +65,10 @@ pub fn gpu_mb_update_impulse_joint_constraints( builder.side_a_kind == SIDE_KIND_FIXED && builder.side_b_kind == SIDE_KIND_FIXED; if !is_dummy { let batch_id = t % nb; + // Interleaved view shared by the dynamics buffers (multibody_info / + // links_workspace / body_jacobians / dof_state) and the constraint + // slots. The jacobians buffer is interleaved at per-joint-region + // granularity instead (dense inside; see `update_one_joint`). let il = VSlice::interleaved(0, nb, batch_id); let bix = batch_ids.body_ix(batch_id); builder.update_one_joint( @@ -191,7 +200,8 @@ pub fn gpu_mb_solve_impulse_joint_constraints( let bix = batch_ids.body_ix(batch_id); // `color_groups` is a per-batch prefix-sum over the color-sorted - // builders: color `c` owns the sorted-builder range + // builders, stored color-major interleaved (`color * nb + batch`): + // color `c` owns the sorted-builder range // `[color_groups[c-1], color_groups[c])` (start `0` for color `0`). let color = *curr_color as usize; let start = if color > 0 { diff --git a/src_rbd_shaders/dynamics/multibody/impulse_joint_constraints/update.rs b/src_rbd_shaders/dynamics/multibody/impulse_joint_constraints/update.rs index 64089cc9..90822997 100644 --- a/src_rbd_shaders/dynamics/multibody/impulse_joint_constraints/update.rs +++ b/src_rbd_shaders/dynamics/multibody/impulse_joint_constraints/update.rs @@ -41,6 +41,8 @@ pub(super) fn solve_mb_wj( let v = jacobians.read(j_id as usize + k as usize); jacobians.write(wj_base + k as usize, v); } + // Dense per-multibody-region views (`il` carries `num_batches` / + // `batch_id` as stride / shift). let m = MatSlice::dense( mb.mass_matrix_offset as usize * il.stride as usize + il.shift as usize * (ndofs * ndofs) as usize, @@ -89,6 +91,8 @@ impl MbImpulseJointBuilder { lock_cfm_coeff: f32, max_corr_velocity: f32, ) { + // Constraint / jacobian slots are batch-local ids into the + // batch-interleaved buffers (resolved through `il`). let cons_base = self.constraint_id as usize; // Mark all axis-constraint slots inactive up-front; the active branches // below overwrite the live ones (rapier rebuilds the entire @@ -192,6 +196,11 @@ impl MbImpulseJointBuilder { mb: mb_b, }; let stride = axis_stride(ndofs_a, ndofs_b); + // The jacobians buffer is interleaved at per-joint-region granularity: + // this joint's region for the current batch starts at + // `jacobian_offset * num_batches + batch * jacobian_capacity` and is + // dense inside. All `j_id`s below (stored in the constraints) are + // absolute dense indices into that region. let j_base = self.jacobian_offset as usize * il.stride as usize + il.shift as usize * self.jacobian_capacity as usize; diff --git a/src_rbd_shaders/dynamics/multibody/integrate.rs b/src_rbd_shaders/dynamics/multibody/integrate.rs index 12e9c5fd..76fce2f7 100644 --- a/src_rbd_shaders/dynamics/multibody/integrate.rs +++ b/src_rbd_shaders/dynamics/multibody/integrate.rs @@ -51,6 +51,8 @@ pub fn gpu_mb_integrate_velocities( let mut dof_vel = batch_ids .ib_mut(batch_id, dof_state) .offset(mb.first_dof as usize); + // The accelerations live in the per-multibody-region `gen_forces` buffer + // (dense region per (multibody, batch); `dof_state` stays interleaved). let acc0 = batch_ids.mb_region(batch_id, mb.first_dof, mb.ndofs); for d in 0..mb.ndofs { diff --git a/src_rbd_shaders/dynamics/multibody/solve_constraints.rs b/src_rbd_shaders/dynamics/multibody/solve_constraints.rs index 90d56b5a..b2c23a7d 100644 --- a/src_rbd_shaders/dynamics/multibody/solve_constraints.rs +++ b/src_rbd_shaders/dynamics/multibody/solve_constraints.rs @@ -126,6 +126,9 @@ pub fn gpu_mb_solve_constraints( let jcol_base = batch_ids.mb_joint_constraint_columns_start(batch_id) + (mb.first_constraint as usize) * dofs_stride; + // This multibody's dynamic segment in the flat constraint buffer, and the + // paired jac/column arena: slot `g` owns `2 * dofs_stride` dense floats, + // the `J` row first, then its `M^-1*J^T` column. let ccons_base = mb.contact_constraint_start as usize; let cjc_base = ccons_base * 2 * dofs_stride; @@ -137,6 +140,10 @@ pub fn gpu_mb_solve_constraints( } let active = in_range && ndofs != 0 && (mb.max_constraints != 0 || contact_count != 0); + // Load the generalized velocities into workgroup memory. The accumulated + // contact impulses stay in storage: every impulse access below is lane-0 + // only, so same-invocation ordering makes storage reads-after-writes safe + // and no compile-time per-multibody bound is needed. if active && lane < ndofs { dof_v.write( lane as usize, @@ -289,7 +296,8 @@ pub fn gpu_mb_solve_constraints( j_dot_v1 += scratch.read(i as usize); } } - // Free-body side stays lane-0-local. + // Free-body side stays lane-0-local (`free_body_id` is a global + // body id). let free = if is_self { Velocity::default() } else { @@ -322,6 +330,8 @@ pub fn gpu_mb_solve_constraints( // Normal: clamp to ≥ 0. Friction: cap the tangent pair to the // circular cone `μ · normal_impulse`. let (new0, new1) = if is_tangent { + // The paired normal was updated earlier in this sweep by this + // same lane, so the storage read observes the fresh value. let limit = cons.friction_coeff * contact_constraints .at(ccons_base + cons.normal_constraint_slot as usize) @@ -371,6 +381,8 @@ pub fn gpu_mb_solve_constraints( workgroup_memory_barrier_with_group_sync(); } + // Writeback (the contact impulses were updated in storage as they were + // solved). if active && lane < ndofs { dof_state.write( batch_ids.mbi(batch_id, v_base + lane as usize), @@ -527,12 +539,15 @@ pub fn gpu_mb_build_contact_delassus( let mb = multibody_info.read(batch_ids.mbi(batch_id, mb_idx as usize)); let ndofs = mb.ndofs; + // The Delassus path's per-multibody blocks and shared SoA arrays are + // compile-time sized: clamp the (otherwise unbounded) dynamic count. let count = mb.contact_constraint_count.min(MAXC); if ndofs == 0 || count == 0 { return; } let cons_base = mb.contact_constraint_start as usize; + // Paired jac/column arena (see `gpu_mb_solve_constraints`). let dofs_stride = batch_ids.dof_batch_capacity as usize; let jc_base = cons_base * 2 * dofs_stride; let d_base = ((batch_id * batch_ids.multibodies_batch_capacity + mb_idx) as usize) @@ -613,6 +628,7 @@ pub fn gpu_mb_solve_contacts_delassus( let mb = multibody_info.read(batch_ids.mbi(batch_id, slot as usize)); let ndofs = mb.ndofs; + // Clamped: see `gpu_mb_build_contact_delassus`. let count = mb.contact_constraint_count.min(MAXC); // Uniform per workgroup: every lane of this group returns together. #[cfg(not(feature = "web-compat"))] @@ -624,6 +640,7 @@ pub fn gpu_mb_solve_contacts_delassus( let v_base = mb.first_dof as usize; let cons_base = mb.contact_constraint_start as usize; + // Paired jac/column arena (see `gpu_mb_solve_constraints`). let dofs_stride = batch_ids.dof_batch_capacity as usize; let jc_base = cons_base * 2 * dofs_stride; let d_base = ((batch_id * batch_ids.multibodies_batch_capacity + mb_idx) as usize) diff --git a/src_rbd_shaders/dynamics/multibody/types.rs b/src_rbd_shaders/dynamics/multibody/types.rs index 75fea041..64884878 100644 --- a/src_rbd_shaders/dynamics/multibody/types.rs +++ b/src_rbd_shaders/dynamics/multibody/types.rs @@ -33,8 +33,20 @@ pub const CONTACT_CONSTRAINTS_PER_POINT: u32 = 2; pub const CONTACT_CONSTRAINTS_PER_POINT: u32 = 3; /// Total constraint slots reserved per multibody (= contact points × DIM). +/// Per-multibody constraint bound used only by the (currently disabled) +/// Delassus constraint-space solve path, whose per-multibody blocks and shared +/// SoA arrays need a compile-time size. The live dof-space pipeline has no +/// per-multibody cap: its flat constraint buffer is demand-sized, so +/// re-enabling the Delassus path requires clamping the per-multibody counts to +/// this bound again. pub const MAX_MB_CONTACT_CONSTRAINTS_PER_MB: u32 = MAX_MB_CONTACTS_PER_MB * CONTACT_CONSTRAINTS_PER_POINT; + +/// Per-(multibody, batch) slot reservation honored by the segment scan when +/// the flat constraint buffer overflows: every multibody is guaranteed up to +/// this many slots (8 contact points) before the leftover capacity is handed +/// out in order, so an overflow degrades every environment a little instead of +/// starving the last ones entirely while the auto-resize catches up. pub const MB_CONS_SLOT_RESERVE: u32 = 8 * CONTACT_CONSTRAINTS_PER_POINT; /// `kind` value: inactive / unused slot. @@ -443,11 +455,24 @@ pub struct MultibodyInfo { /// `DISABLE_SELF_CONTACTS`). The contact-constraint kernel skips self /// contacts when this is `0`. pub self_contacts_enabled: u32, + /// Per-frame count of active multibody contact constraints for this + /// multibody: predicted by `gpu_mb_count_contact_constraints`, clamped by + /// the segment scan, finalized by `gpu_mb_init_contact_constraints`; read + /// by the warmstart / finalize / solve / remove-bias contact kernels. pub contact_constraint_count: u32, + /// Per-frame start of this multibody's dynamic segment in the flat + /// contact-constraint buffer (written by `gpu_mb_cons_offsets_scan`). pub contact_constraint_start: u32, + /// Previous frame's segment bounds (saved at step start, before the new + /// layout is computed), matched against by the contact warmstart transfer. pub old_contact_constraint_start: u32, pub old_contact_constraint_count: u32, + /// Per-step count of this multibody's entries in the contact→multibody + /// index (one entry per contact touching one of its links), written by + /// `gpu_mb_cons_offsets_scan`. pub contact_index_len: u32, + /// Per-step start of this multibody's segment in the contact→multibody + /// index buffer (same scan). pub contact_index_start: u32, /// First entry of this multibody's DoF couplings in the `dof_couplings` /// buffer (relative to the batch's coupling slice). @@ -456,13 +481,23 @@ pub struct MultibodyInfo { pub num_couplings: u32, } +/// One entry of the per-step contact→multibody index: contact +/// `contact_slot` touches multibody link `link_a` (and `link_b` when it is a +/// self-contact). Built by `gpu_mb_scatter_contact_index` into per-multibody +/// segments (`MultibodyInfo::contact_index_start/len`) so the emission pass +/// visits only its own contacts instead of scanning the whole flat list. #[derive(Copy, Clone, Default)] #[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] #[repr(C)] pub struct MbContactIndexEntry { + /// Slot of the contact in the flat contacts buffer. pub contact_slot: u32, + /// Touched link of the owning multibody (the constraint's A side). pub link_a: u32, + /// Second touched link for self-contacts, `u32::MAX` otherwise. pub link_b: u32, + /// 1 when the multibody is on the contact's `bodies.x` side (fixes the + /// jacobian sign and which side is the free body). pub mb_on_first: u32, } diff --git a/src_rbd_shaders/dynamics/multibody/ws_soa.rs b/src_rbd_shaders/dynamics/multibody/ws_soa.rs index f3c666d2..2c331438 100644 --- a/src_rbd_shaders/dynamics/multibody/ws_soa.rs +++ b/src_rbd_shaders/dynamics/multibody/ws_soa.rs @@ -111,7 +111,9 @@ impl WsAddr { } /// Flat index of quad `quad` (a `WS_*` field offset + quad index) of - /// link `k` (relative to `base`). + /// link `k` (relative to `base`). The workspace is interleaved at + /// per-link-record granularity: link `L` of batch `b` owns the dense + /// `WS_QUADS` quads at `(L * num_batches + b) * WS_QUADS`. #[inline] pub fn at(&self, k: u32, quad: u32) -> usize { ((self.base + k as usize) * self.stride as usize + self.shift as usize) diff --git a/src_rbd_shaders/dynamics/solver.rs b/src_rbd_shaders/dynamics/solver.rs index 4a23544e..cebfbd89 100644 --- a/src_rbd_shaders/dynamics/solver.rs +++ b/src_rbd_shaders/dynamics/solver.rs @@ -57,6 +57,8 @@ pub fn gpu_solver_init_constraints( let i = i as usize; let im = contacts.at(i); if im.contact.len == 0 { + // Gap or inert slot: clear the (stale) constraint so every flat + // consumer skips it. constraints.at_mut(i).len = 0; continue; } @@ -87,6 +89,7 @@ pub fn gpu_solver_count_constraints( ) { let num_threads = num_workgroups.x * WORKGROUP_SIZE; + // Flat over all contacts of all batches: ids are global, counts cumulative. let total = contact_plan.bound; let contacts = Slice(contacts, 0); let mut body_constraint_counts = SliceMut(body_constraint_counts, 0); @@ -142,6 +145,7 @@ pub fn gpu_solver_update_constraints( for i in StepRng::new(invocation_id.x..total, num_threads) { if constraints[i as usize].len == 0 { + // Gap / inert slot. continue; } constraints[i as usize].update_constraint( @@ -176,6 +180,7 @@ pub fn gpu_solver_refresh_rhs_wo_bias( for i in StepRng::new(invocation_id.x..total, num_threads) { if constraints[i as usize].len == 0 { + // Gap / inert slot. continue; } constraints[i as usize].refresh_rhs_wo_bias( @@ -366,6 +371,9 @@ pub fn gpu_warmstart( let mut solver_vels = SliceMut(solver_vels, 0); let color = *curr_color; + // Buckets are color-major (`color * num_batches + batch`) and the buffer + // holds post-scatter exclusive ENDS, so color `c` (over every batch) spans + // `[ends[c*nb - 1], ends[(c+1)*nb - 1])`. `c >= 1` keeps the index valid. let start = color_starts.read((color * nb - 1) as usize); let end = color_starts.read(((color + 1) * nb - 1) as usize); @@ -409,6 +417,7 @@ pub fn gpu_step_gauss_seidel( let color = *curr_color; let use_bias = *use_bias != 0; + // Color-major bucket ends; see `gpu_warmstart`. let start = color_starts.read((color * nb - 1) as usize); let end = color_starts.read(((color + 1) * nb - 1) as usize); @@ -456,6 +465,8 @@ pub fn gpu_warmstart_fused( let num_colors = *num_colors; for color in 1..=num_colors { + // This batch's bucket for `color` (color-major layout, post-scatter + // exclusive ends; the index is >= num_batches >= 1 for color >= 1). let bucket = (color * nb + batch_id) as usize; let start = color_starts.read(bucket - 1); let end = color_starts.read(bucket); @@ -521,6 +532,7 @@ pub fn gpu_step_gauss_seidel_fused( let use_bias = *use_bias != 0; for color in 1..=num_colors { + // Empty-color skip: see `gpu_warmstart_fused`. let bucket = (color * nb + batch_id) as usize; let start = color_starts.read(bucket - 1); let end = color_starts.read(bucket); diff --git a/src_rbd_shaders/dynamics/warmstart.rs b/src_rbd_shaders/dynamics/warmstart.rs index 5861b5f1..d590a026 100644 --- a/src_rbd_shaders/dynamics/warmstart.rs +++ b/src_rbd_shaders/dynamics/warmstart.rs @@ -39,6 +39,7 @@ pub fn gpu_transfer_warmstart_impulses( let cid_new = invocation_id.x; if cid_new < total { + // Gap / inert slot: stale body ids must not be dereferenced. if new_constraints[cid_new as usize].len == 0 { return; } @@ -89,6 +90,7 @@ pub fn gpu_seed_colors_from_warmstart( let i = invocation_id.x as usize; if (i as u32) < total { + // Gap / inert slot: stale body ids must not be dereferenced. if new_constraints[i].len == 0 { return; } diff --git a/src_rbd_shaders/utils/indices.rs b/src_rbd_shaders/utils/indices.rs index 8ee15ec2..a5602d58 100644 --- a/src_rbd_shaders/utils/indices.rs +++ b/src_rbd_shaders/utils/indices.rs @@ -20,8 +20,16 @@ pub struct BatchIndices { pub colliders_len: u32, /// Number of *active* rigid bodies per batch. pub bodies_len: u32, + /// Total capacity of the flat collision-pair buffer (all batches share it; + /// also the capacity of the flat PFM work-list, which is sized identically). pub collision_pairs_capacity: u32, + /// Total capacity of the flat contacts buffer (and of every contacts-keyed + /// buffer: constraints, builders, colors, sorted ids). Contact slots are + /// positional (pair `t` owns slot `t`, PFM entry `i` owns slot + /// `pairs_total + i`), so this is always `2 * collision_pairs_capacity`. pub contacts_capacity: u32, + /// Number of *active* free-body impulse joints per batch (the loop bound; + /// the joint-keyed buffers are batch-interleaved). pub impulse_joints_len: u32, /* @@ -40,7 +48,14 @@ pub struct BatchIndices { */ pub mb_joint_constraints_batch_capacity: u32, pub mb_joint_constraint_columns_batch_capacity: u32, + /// Total capacity (in slots) of the flat multibody contact-constraint + /// buffer; per-multibody segments within it are dynamic (see + /// `MultibodyInfo::contact_constraint_start`). The paired jac/column arena + /// holds `2 * dof_batch_capacity` floats per slot (`Jᵀ` row, then its + /// `M⁻¹·Jᵀ` column). pub mb_contact_constraints_capacity: u32, + /// Per-batch multibody-touching impulse-joint slot count (loop bound for + /// the flat sweeps). pub mb_imp_joints_batch_capacity: u32, /// Actual max `ndofs` across every multibody in every batch (often smaller /// than the fixed `MAX_MB_DOFS` limit). @@ -81,10 +96,16 @@ impl BatchIndices { * compute base indices into flat f32 buffers (e.g. when constructing a * `MatSlice::dense(base, ...)`). */ + /// Global id of body/collider `local` of `batch_id`: bodies are stored + /// batch-interleaved (`local * num_batches + batch`), so the same + /// topological entity across environments sits on adjacent lanes. #[inline] pub fn body_global(&self, batch_id: u32, local: u32) -> usize { local as usize * self.num_batches as usize + batch_id as usize } + + /// Strided body indexer for `batch_id` (see [`BodyIx`]), for helpers that + /// resolve env-local body ids without carrying a slice view. #[inline] pub fn body_ix(&self, batch_id: u32) -> BodyIx { BodyIx { @@ -122,13 +143,19 @@ impl BatchIndices { } } - /// Interleaved dense matrix view at intra-batch element offset `offset`. + /// Base of batch `batch_id`'s dense region in a per-multibody-region + /// buffer (mass matrices, LU pivots, body jacobians, coriolis blocks, + /// generalized forces, impulse-joint jacobians): regions tile as + /// `offset * num_batches + batch * len`, dense inside. #[inline] pub fn mb_region(&self, batch_id: u32, offset: u32, len: u32) -> usize { offset as usize * self.num_batches as usize + batch_id as usize * len as usize } - /// Interleaved vector view at intra-batch element offset `offset`. + /// Batch owning the collider with global id `collider_id` (the flat + /// collision-pair and contact buffers store global collider/body ids; + /// bodies are batch-interleaved so the batch is the id modulo the batch + /// count). #[inline] pub fn collider_batch(&self, collider_id: u32) -> u32 { collider_id % self.num_batches @@ -150,6 +177,9 @@ impl BatchIndices { } } +/// Strided indexer mapping an env-local body/collider id to its global slot in +/// the batch-interleaved per-body buffers: `global = id * stride + shift` +/// (`stride = num_batches`, `shift = batch_id`). #[derive(Copy, Clone)] pub struct BodyIx { pub stride: u32, From 47c922702e5da925a0da2e65235dfc3e21d4ce3b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Sun, 30 Aug 2026 18:32:21 +0200 Subject: [PATCH 6/7] feat(rbd): use stronger typing in shaders when possible --- src_rbd/dynamics/multibody/env_reset.rs | 39 ++++-- .../multibody/multibody_from_rapier.rs | 12 +- src_rbd/dynamics/multibody/multibody_set.rs | 16 ++- src_rbd/pipeline/rbd_state.rs | 14 +- .../dynamics/multibody/env_reset.rs | 121 +++++++++++++----- .../dynamics/multibody/scatter_motor.rs | 43 +++++-- 6 files changed, 179 insertions(+), 66 deletions(-) diff --git a/src_rbd/dynamics/multibody/env_reset.rs b/src_rbd/dynamics/multibody/env_reset.rs index 94fcc75e..a2abef99 100644 --- a/src_rbd/dynamics/multibody/env_reset.rs +++ b/src_rbd/dynamics/multibody/env_reset.rs @@ -15,10 +15,11 @@ use super::multibody_set::GpuMultibodySet; use crate::math::Vector; use crate::shaders::dynamics::{ - GpuMbEnvReset, GpuMbEnvResetBatch, GpuMbEnvResetBatchDofs, MULTIBODY_ROOT, MultibodyLinkStatic, - MultibodyLinkWorkspace, WS_QUADS, ws_soa_from_structs, ws_soa_to_structs, + EnvResetRecord, GpuMbEnvReset, GpuMbEnvResetBatch, GpuMbEnvResetBatchDofs, MULTIBODY_ROOT, + MbEnvResetBatchParams, MbEnvResetParams, MultibodyLinkStatic, MultibodyLinkWorkspace, WS_QUADS, + ws_soa_from_structs, ws_soa_to_structs, }; -use glamx::{UVec4, Vec4}; +use glamx::Vec4; use khal::BufferUsages; use khal::Shader; use khal::backend::{Backend, GpuBackend}; @@ -115,7 +116,7 @@ pub(super) struct EnvResetBundle { staging_ws: Tensor, staging_links: Tensor, staging_dofs: Tensor, - params: Tensor, + params: Tensor, } impl EnvResetBundle { @@ -138,7 +139,17 @@ impl EnvResetBundle { .unwrap(), staging_dofs: Tensor::vector(backend, vec![0.0f32; dpb.max(1) as usize], storage) .unwrap(), - params: Tensor::scalar(backend, UVec4::new(0, 0, lpb, dpb), uniform).unwrap(), + params: Tensor::scalar( + backend, + MbEnvResetParams { + dst_env: 0, + num_batches: 0, + links_per_batch: lpb, + dofs_per_batch: dpb, + }, + uniform, + ) + .unwrap(), } } } @@ -235,7 +246,12 @@ impl GpuMultibodySet { } bundle.params = Tensor::scalar( backend, - UVec4::new(dst_env, nb, lpb, dpb), + MbEnvResetParams { + dst_env, + num_batches: nb, + links_per_batch: lpb, + dofs_per_batch: dpb, + }, BufferUsages::STORAGE | BufferUsages::UNIFORM | BufferUsages::COPY_DST, ) .unwrap(); @@ -322,7 +338,7 @@ impl GpuMultibodySet { &mut self, backend: &GpuBackend, enc: &mut ::Encoder, - resets: &[UVec4], + resets: &[EnvResetRecord], offsets: &[Vec4], dof_vels: &[f32], ) { @@ -342,7 +358,7 @@ impl GpuMultibodySet { // 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); + let (env, t) = (meta.env as usize, meta.template as usize); for (k, ls) in tpl.mirror_links[t].iter().enumerate() { self.links_static_mirror[k * nb as usize + env] = *ls; } @@ -354,7 +370,12 @@ impl GpuMultibodySet { let t_vels = Tensor::vector(backend, dof_vels, storage).unwrap(); let params = Tensor::scalar( backend, - UVec4::new(nb, lpb, dpb, n), + MbEnvResetBatchParams { + num_batches: nb, + links_per_batch: lpb, + dofs_per_batch: dpb, + num_resets: n, + }, BufferUsages::STORAGE | BufferUsages::UNIFORM | BufferUsages::COPY_DST, ) .unwrap(); diff --git a/src_rbd/dynamics/multibody/multibody_from_rapier.rs b/src_rbd/dynamics/multibody/multibody_from_rapier.rs index 38ab4bf4..68c0b213 100644 --- a/src_rbd/dynamics/multibody/multibody_from_rapier.rs +++ b/src_rbd/dynamics/multibody/multibody_from_rapier.rs @@ -2,9 +2,10 @@ use super::multibody_set::*; use crate::shaders::dynamics::{ - ConstraintSoftness, MAX_AXIS_CONSTRAINTS, MAX_MB_CONTACT_CONSTRAINTS_PER_MB, MbDofCoupling, - MbImpulseJointBuilder, MbImpulseJointConstraint, MultibodyContactConstraint, MultibodyInfo, - MultibodyJointConstraint, MultibodyLinkStatic, MultibodyLinkWorkspace, RbdSimParams, + ConstraintSoftness, MAX_AXIS_CONSTRAINTS, MAX_MB_CONTACT_CONSTRAINTS_PER_MB, MbDelayTickParams, + MbDofCoupling, MbImpulseJointBuilder, MbImpulseJointConstraint, MultibodyContactConstraint, + MultibodyInfo, MultibodyJointConstraint, MultibodyLinkStatic, MultibodyLinkWorkspace, + RbdSimParams, }; use crate::shaders::utils::linalg::MAX_MB_DOFS; use khal::BufferUsages; @@ -668,7 +669,10 @@ impl GpuMultibodySet { .unwrap(), motor_delay_params: Tensor::scalar( backend, - glamx::UVec4::new(num_batches, 2 + links_cap, 0, 0), + MbDelayTickParams { + num_envs: num_batches, + stride: 2 + links_cap, + }, BufferUsages::STORAGE | BufferUsages::UNIFORM, ) .unwrap(), diff --git a/src_rbd/dynamics/multibody/multibody_set.rs b/src_rbd/dynamics/multibody/multibody_set.rs index 7e59df1a..6da098fa 100644 --- a/src_rbd/dynamics/multibody/multibody_set.rs +++ b/src_rbd/dynamics/multibody/multibody_set.rs @@ -3,9 +3,9 @@ use crate::math::Pose; use crate::shaders::dynamics::{ - ConstraintSoftness, LocalMassProperties, MbDofCoupling, MbImpulseJointBuilder, - MbImpulseJointConstraint, MultibodyContactConstraint, MultibodyInfo, MultibodyJointConstraint, - MultibodyLinkStatic, MultibodyLinkWorkspace, RbdSimParams, + ConstraintSoftness, LocalMassProperties, MbDelayStateParams, MbDelayTickParams, MbDofCoupling, + MbImpulseJointBuilder, MbImpulseJointConstraint, MultibodyContactConstraint, MultibodyInfo, + MultibodyJointConstraint, MultibodyLinkStatic, MultibodyLinkWorkspace, RbdSimParams, }; use crate::shaders::utils::BatchIndices; use khal::BufferUsages; @@ -59,7 +59,7 @@ pub(super) struct DelayUpdateBundle { pub(super) struct DelayUpdateCache { shader: DelayUpdateBundle, t_links: Tensor, - params: Tensor, + params: Tensor, num_actuated: u32, } @@ -155,7 +155,7 @@ pub struct GpuMultibodySet { /// 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, + 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 @@ -1076,7 +1076,11 @@ impl GpuMultibodySet { t_links: Tensor::vector(backend, actuated_link_ids, BufferUsages::STORAGE)?, params: Tensor::scalar( backend, - glamx::UVec4::new(num_actuated, self.num_batches, stride, 0), + MbDelayStateParams { + num_actuated, + num_envs: self.num_batches, + stride, + }, BufferUsages::STORAGE | BufferUsages::UNIFORM, )?, num_actuated, diff --git a/src_rbd/pipeline/rbd_state.rs b/src_rbd/pipeline/rbd_state.rs index e05a4438..e4bf3ab5 100644 --- a/src_rbd/pipeline/rbd_state.rs +++ b/src_rbd/pipeline/rbd_state.rs @@ -827,7 +827,8 @@ impl RbdState { offsets: &[crate::math::Vector], dof_vels: &[f32], ) { - use glamx::{UVec4, Vec4}; + use crate::shaders::dynamics::{EnvResetBodiesParams, EnvResetRecord}; + use glamx::Vec4; use khal::backend::Encoder as _; let n = resets.len() as u32; if n == 0 { @@ -836,9 +837,9 @@ impl RbdState { 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 + let meta: Vec = resets .iter() - .map(|&(env, t)| UVec4::new(env, t, 0, 0)) + .map(|&(env, template)| EnvResetRecord { env, template }) .collect(); let offs: Vec = offsets .iter() @@ -849,7 +850,12 @@ impl RbdState { let t_offs = Tensor::vector(backend, &offs, storage).unwrap(); let params = Tensor::scalar( backend, - UVec4::new(bps, vs, n, nb), + EnvResetBodiesParams { + bodies_per_env: bps, + vels_per_env: vs, + num_resets: n, + num_batches: nb, + }, BufferUsages::STORAGE | BufferUsages::UNIFORM | BufferUsages::COPY_DST, ) .unwrap(); diff --git a/src_rbd_shaders/dynamics/multibody/env_reset.rs b/src_rbd_shaders/dynamics/multibody/env_reset.rs index b5dab844..2184acea 100644 --- a/src_rbd_shaders/dynamics/multibody/env_reset.rs +++ b/src_rbd_shaders/dynamics/multibody/env_reset.rs @@ -4,7 +4,7 @@ //! static link descriptors, generalized coordinates and velocities) from a //! compact contiguous staging blob into the batch-interleaved live buffers. -use glamx::{UVec4, Vec4}; +use glamx::Vec4; use khal_std::glamx::UVec3; use khal_std::index::MaybeIndexUnchecked; use khal_std::macros::{spirv, spirv_bindgen}; @@ -12,6 +12,65 @@ use khal_std::macros::{spirv, spirv_bindgen}; use super::types::MultibodyLinkStatic; use super::ws_soa::{WS_COORDS, WS_LTP, WS_LTW, WS_QUADS}; +/// One entry of the batched-reset list: restore environment `env` from the +/// GPU-resident template `template`. +#[derive(Copy, Clone, Default)] +#[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] +#[repr(C)] +pub struct EnvResetRecord { + /// Destination environment (batch) index. + pub env: u32, + /// Index of the resident template to restore from. + pub template: u32, +} + +/// Parameters of the single-env staging reset ([`gpu_mb_env_reset`]). +#[derive(Copy, Clone, Default)] +#[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] +#[repr(C)] +pub struct MbEnvResetParams { + /// Destination environment (batch) index. + pub dst_env: u32, + /// Total number of simulation batches (the interleave stride). + pub num_batches: u32, + /// Links per batch (with padding slots). + pub links_per_batch: u32, + /// Generalized-velocity entries per batch. + pub dofs_per_batch: u32, +} + +/// Parameters shared by the two batched multibody reset passes +/// ([`gpu_mb_env_reset_batch`] and [`gpu_mb_env_reset_batch_dofs`]). +#[derive(Copy, Clone, Default)] +#[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] +#[repr(C)] +pub struct MbEnvResetBatchParams { + /// Total number of simulation batches (the interleave stride). + pub num_batches: u32, + /// Links per batch (with padding slots). + pub links_per_batch: u32, + /// Generalized-velocity entries per batch. + pub dofs_per_batch: u32, + /// Number of entries in the reset list. + pub num_resets: u32, +} + +/// Parameters of the rigid-body half of the batched reset +/// ([`gpu_env_reset_bodies`]). +#[derive(Copy, Clone, Default)] +#[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] +#[repr(C)] +pub struct EnvResetBodiesParams { + /// Body slots per environment (template stride). + pub bodies_per_env: u32, + /// Velocity slots per environment (template stride). + pub vels_per_env: u32, + /// Number of entries in the reset list. + pub num_resets: u32, + /// Total number of simulation batches (the interleave stride). + pub num_batches: u32, +} + /// 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 @@ -33,14 +92,13 @@ pub fn gpu_mb_env_reset( #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] links_static: &mut [MultibodyLinkStatic], #[spirv(storage_buffer, descriptor_set = 0, binding = 5)] dof_state: &mut [f32], - // x = dst_env, y = num_batches, z = links_per_batch, w = dofs_per_batch. - #[spirv(uniform, descriptor_set = 0, binding = 6)] params: &UVec4, + #[spirv(uniform, descriptor_set = 0, binding = 6)] params: &MbEnvResetParams, ) { let i = invocation_id.x; - let env = params.x; - let nb = params.y; - let lpb = params.z; - let dpb = params.w; + let env = params.dst_env; + let nb = params.num_batches; + let lpb = params.links_per_batch; + let dpb = params.dofs_per_batch; if i < lpb * WS_QUADS { // Staging is per-env dense; the live workspace is interleaved at @@ -82,22 +140,21 @@ pub fn gpu_mb_env_reset_batch( #[spirv(global_invocation_id)] invocation_id: UVec3, #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] templates_ws: &[Vec4], #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] link_flags: &[u32], - #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] resets: &[UVec4], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] resets: &[EnvResetRecord], #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] offsets: &[Vec4], #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] links_workspace: &mut [Vec4], - // x = num_batches, y = links_per_batch, z = dofs_per_batch, w = num_resets. - #[spirv(uniform, descriptor_set = 0, binding = 5)] params: &UVec4, + #[spirv(uniform, descriptor_set = 0, binding = 5)] params: &MbEnvResetBatchParams, ) { let i = invocation_id.x; let r = invocation_id.y; - let nb = params.x; - let lpb = params.y; - if r >= params.w { + let nb = params.num_batches; + let lpb = params.links_per_batch; + if r >= params.num_resets { return; } let meta = resets.read(r as usize); - let env = meta.x; - let t = meta.y; + let env = meta.env; + let t = meta.template; let off = offsets.read(r as usize); if i < lpb * WS_QUADS { @@ -137,25 +194,24 @@ 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)] resets: &[UVec4], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] resets: &[EnvResetRecord], #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] dof_vels: &[f32], #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] links_static: &mut [MultibodyLinkStatic], #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] dof_state: &mut [f32], - // x = num_batches, y = links_per_batch, z = dofs_per_batch, w = num_resets. - #[spirv(uniform, descriptor_set = 0, binding = 5)] params: &UVec4, + #[spirv(uniform, descriptor_set = 0, binding = 5)] params: &MbEnvResetBatchParams, ) { 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 { + let nb = params.num_batches; + let lpb = params.links_per_batch; + let dpb = params.dofs_per_batch; + if r >= params.num_resets { return; } let meta = resets.read(r as usize); - let env = meta.x; - let t = meta.y; + let env = meta.env; + let t = meta.template; if i < lpb { links_static.write( @@ -186,25 +242,24 @@ pub fn gpu_env_reset_bodies( #[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 = 3)] resets: &[EnvResetRecord], #[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, + #[spirv(uniform, descriptor_set = 0, binding = 7)] params: &EnvResetBodiesParams, ) { let i = invocation_id.x; let r = invocation_id.y; - let bps = params.x; - let vs = params.y; - let nb = params.w; - if r >= params.z { + let bps = params.bodies_per_env; + let vs = params.vels_per_env; + let nb = params.num_batches; + if r >= params.num_resets { return; } let meta = resets.read(r as usize); - let env = meta.x; - let t = meta.y; + let env = meta.env; + let t = meta.template; let off = offsets.read(r as usize); // Templates are dense per-env arrays; the live per-body buffers are diff --git a/src_rbd_shaders/dynamics/multibody/scatter_motor.rs b/src_rbd_shaders/dynamics/multibody/scatter_motor.rs index c5b9a83d..16ef0da6 100644 --- a/src_rbd_shaders/dynamics/multibody/scatter_motor.rs +++ b/src_rbd_shaders/dynamics/multibody/scatter_motor.rs @@ -9,12 +9,37 @@ //! element `(j, env)` at `j · num_envs + env`, matching the policy action //! buffer layout. -use khal_std::glamx::{UVec3, UVec4}; +use khal_std::glamx::UVec3; use khal_std::index::MaybeIndexUnchecked; use khal_std::macros::{spirv, spirv_bindgen}; use super::types::MultibodyLinkStatic; +/// Parameters of the on-device delay-state refresh +/// ([`gpu_mb_delay_state_update`]). +#[derive(Copy, Clone, Default)] +#[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] +#[repr(C)] +pub struct MbDelayStateParams { + /// Number of actuated joints (rows of the target tensor). + pub num_actuated: u32, + /// Number of environments (columns of the target tensor). + pub num_envs: u32, + /// Per-env stride of the delay-state buffer (`2 + links_per_batch`). + pub stride: u32, +} + +/// Parameters of the per-step delay tick ([`gpu_mb_delay_tick`]). +#[derive(Copy, Clone, Default)] +#[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] +#[repr(C)] +pub struct MbDelayTickParams { + /// Number of environments. + pub num_envs: u32, + /// Per-env stride of the delay-state buffer (`2 + links_per_batch`). + pub stride: u32, +} + /// 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. @@ -70,15 +95,14 @@ pub fn gpu_mb_delay_state_update( #[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, + #[spirv(uniform, descriptor_set = 0, binding = 4)] params: &MbDelayStateParams, ) { let j = invocation_id.x; let env = invocation_id.y; - if j >= params.x || env >= params.y { + if j >= params.num_actuated || env >= params.num_envs { return; } - let base = (env * params.z) as usize; + let base = (env * params.stride) as usize; if j == 0 { delay_state.write(base, 0.0); delay_state.write(base + 1, k_eff.read(env as usize)); @@ -86,7 +110,7 @@ pub fn gpu_mb_delay_state_update( 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), + prev_targets.read((j * params.num_envs + env) as usize), ); } @@ -100,14 +124,13 @@ pub fn gpu_mb_delay_state_update( 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, + #[spirv(uniform, descriptor_set = 0, binding = 1)] params: &MbDelayTickParams, ) { let env = invocation_id.x; - if env >= params.x { + if env >= params.num_envs { return; } - let base = (env * params.y) as usize; + let base = (env * params.stride) as usize; let tick = delay_state.read(base); delay_state.write(base, tick + 1.0); } From 2d9aa586ce1f87ee0360655a8e784607dda05d84 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Crozet?= Date: Sun, 30 Aug 2026 18:38:15 +0200 Subject: [PATCH 7/7] =?UTF-8?q?chore:=20CI=E2=80=AFfixes?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/state.rs | 3 +-- src_rbd/broad_phase/narrow_phase.rs | 2 +- src_rbd/dynamics/joint.rs | 3 +-- .../dynamics/multibody/loop_closing_joints.rs | 3 +-- .../multibody/multibody_from_rapier.rs | 2 +- src_rbd/dynamics/multibody/multibody_set.rs | 4 ++-- src_rbd/pipeline/insertion_removal.rs | 13 +++++++----- src_rbd/pipeline/mod.rs | 4 ++-- src_rbd/pipeline/rbd_state_from_rapier.rs | 13 +++++++----- .../dynamics/joint_constraint_builder.rs | 6 +++++- .../dynamics/multibody/contact_sensor.rs | 5 +---- .../impulse_joint_constraints/helper.rs | 20 ++----------------- .../impulse_joint_constraints/kernels.rs | 2 +- .../impulse_joint_constraints/update.rs | 2 +- .../dynamics/multibody/integrate.rs | 1 - .../dynamics/multibody/solve_constraints.rs | 11 +++++----- src_rbd_shaders/dynamics/multibody/ws_soa.rs | 3 +-- src_viewer/ui.rs | 5 ++++- 18 files changed, 46 insertions(+), 56 deletions(-) diff --git a/src/state.rs b/src/state.rs index fa720d82..9a0976d2 100644 --- a/src/state.rs +++ b/src/state.rs @@ -390,8 +390,7 @@ impl NexusState { #[cfg(feature = "dim3")] { c.mb_contact_constraints = rbd.mb_contact_constraints_len() as usize; - c.mb_contact_constraints_capacity = - rbd.mb_contact_constraints_capacity() as usize; + c.mb_contact_constraints_capacity = rbd.mb_contact_constraints_capacity() as usize; } } #[cfg(feature = "mpm")] diff --git a/src_rbd/broad_phase/narrow_phase.rs b/src_rbd/broad_phase/narrow_phase.rs index c902d9aa..f23847f4 100644 --- a/src_rbd/broad_phase/narrow_phase.rs +++ b/src_rbd/broad_phase/narrow_phase.rs @@ -78,7 +78,7 @@ impl PfmSortState { identity: Tensor::vector(backend, &identity, storage).unwrap(), sorted_keys: Tensor::vector_uninit(backend, capacity.max(1), storage).unwrap(), sorted_values: Tensor::vector_uninit(backend, capacity.max(1), storage).unwrap(), - sort_len: Tensor::vector(backend, &[0u32], storage).unwrap(), + sort_len: Tensor::vector(backend, [0u32], storage).unwrap(), workspace: RadixSortWorkspace::new(backend), } } diff --git a/src_rbd/dynamics/joint.rs b/src_rbd/dynamics/joint.rs index 6cda04b5..39fe7111 100644 --- a/src_rbd/dynamics/joint.rs +++ b/src_rbd/dynamics/joint.rs @@ -312,8 +312,7 @@ impl GpuImpulseJointSet { // so the solver's flat sweeps put the same joint of consecutive // batches on adjacent lanes. let dummy_joint = ImpulseJoint::zeroed(); - let mut all_joints = - vec![dummy_joint; num_batches as usize * max_joints as usize]; + let mut all_joints = vec![dummy_joint; num_batches as usize * max_joints as usize]; for (batch, sorted_joints) in per_env_sorted_joints.iter().enumerate() { for (j, joint) in sorted_joints.iter().enumerate() { all_joints[j * num_batches as usize + batch] = *joint; diff --git a/src_rbd/dynamics/multibody/loop_closing_joints.rs b/src_rbd/dynamics/multibody/loop_closing_joints.rs index 314ec27e..a301767b 100644 --- a/src_rbd/dynamics/multibody/loop_closing_joints.rs +++ b/src_rbd/dynamics/multibody/loop_closing_joints.rs @@ -286,8 +286,7 @@ impl GpuMultibodySet { let mut dummy: MbImpulseJointBuilder = bytemuck::Zeroable::zeroed(); dummy.side_a_kind = SIDE_KIND_FIXED; dummy.side_b_kind = SIDE_KIND_FIXED; - let mut all_builders: Vec = - vec![dummy; joints_cap as usize * nb]; + let mut all_builders: Vec = vec![dummy; joints_cap as usize * nb]; for (b, env) in per_env_builders.iter().enumerate() { for (i, builder) in env.iter().enumerate() { all_builders[i * nb + b] = *builder; diff --git a/src_rbd/dynamics/multibody/multibody_from_rapier.rs b/src_rbd/dynamics/multibody/multibody_from_rapier.rs index 68c0b213..cd4ef8a4 100644 --- a/src_rbd/dynamics/multibody/multibody_from_rapier.rs +++ b/src_rbd/dynamics/multibody/multibody_from_rapier.rs @@ -786,7 +786,7 @@ impl GpuMultibodySet { joint_constraints_per_batch: cons_cap, joint_constraint_columns_per_batch: cons_col_cap, contact_constraints_capacity: contact_cons_cap, - mb_cons_demand: Tensor::vector(backend, &[0u32], storage | BufferUsages::COPY_SRC) + mb_cons_demand: Tensor::vector(backend, [0u32], storage | BufferUsages::COPY_SRC) .unwrap(), // Zeroed: the count pass atomically accumulates into these and the // offsets scan re-zeroes them after consuming them. diff --git a/src_rbd/dynamics/multibody/multibody_set.rs b/src_rbd/dynamics/multibody/multibody_set.rs index 6da098fa..6bb9b173 100644 --- a/src_rbd/dynamics/multibody/multibody_set.rs +++ b/src_rbd/dynamics/multibody/multibody_set.rs @@ -805,13 +805,13 @@ impl GpuMultibodySet { self.contact_constraints_capacity } - /// Total contact-constraint slot demand of the last stepped frame (GPU - /// buffer, read back by the auto-resize). /// Debug: the body-id → (multibody, link) lookup buffer. pub fn body_to_link(&self) -> &Tensor<[u32; 2]> { &self.body_to_link } + /// Total contact-constraint slot demand of the last stepped frame (GPU + /// buffer, read back by the auto-resize). pub fn mb_cons_demand(&self) -> &Tensor { &self.mb_cons_demand } diff --git a/src_rbd/pipeline/insertion_removal.rs b/src_rbd/pipeline/insertion_removal.rs index 75623179..098d6545 100644 --- a/src_rbd/pipeline/insertion_removal.rs +++ b/src_rbd/pipeline/insertion_removal.rs @@ -136,7 +136,7 @@ impl RbdState { // Single global pair counter. let collision_pairs_len = Tensor::vector( backend, - &[0u32], + [0u32], BufferUsages::STORAGE | BufferUsages::COPY_SRC, ) .unwrap(); @@ -162,7 +162,7 @@ impl RbdState { // Single global PFM work-list counter. let pfm_pairs_len = Tensor::vector( backend, - &[0u32], + [0u32], BufferUsages::STORAGE | BufferUsages::COPY_SRC, ) .unwrap(); @@ -183,9 +183,12 @@ impl RbdState { Tensor::vector_uninit(backend, color_buckets_stride * num_batches, storage).unwrap(); // Written as a storage buffer by `gpu_contact_plan`, read as a uniform // by every consumer. - let contact_plan = - Tensor::scalar(backend, ContactPlan::default(), storage | BufferUsages::UNIFORM) - .unwrap(); + let contact_plan = Tensor::scalar( + backend, + ContactPlan::default(), + storage | BufferUsages::UNIFORM, + ) + .unwrap(); let color_sorted_ids = Tensor::vector_uninit(backend, contacts_capacity, storage).unwrap(); // Zeroed: see `RbdState::from_rapier`. let old_constraints_counts = Tensor::vector( diff --git a/src_rbd/pipeline/mod.rs b/src_rbd/pipeline/mod.rs index 1347ff60..cf08ce98 100644 --- a/src_rbd/pipeline/mod.rs +++ b/src_rbd/pipeline/mod.rs @@ -6,13 +6,13 @@ #[cfg(all(test, feature = "dim3"))] mod bench_narrow_phase; -#[cfg(all(test, feature = "dim3"))] -mod test_batched_stacks; mod insertion_removal; mod lbvh_validation; mod rbd_state; mod rbd_state_from_rapier; mod rbd_step; +#[cfg(all(test, feature = "dim3"))] +mod test_batched_stacks; #[cfg(feature = "dim3")] pub use rbd_state::RbdSnapshot; diff --git a/src_rbd/pipeline/rbd_state_from_rapier.rs b/src_rbd/pipeline/rbd_state_from_rapier.rs index 5febe240..a758c053 100644 --- a/src_rbd/pipeline/rbd_state_from_rapier.rs +++ b/src_rbd/pipeline/rbd_state_from_rapier.rs @@ -661,7 +661,7 @@ impl RbdState { // Single global pair counter. let collision_pairs_len = Tensor::vector( backend, - &[0u32], + [0u32], BufferUsages::STORAGE | BufferUsages::COPY_SRC, ) .unwrap(); @@ -697,7 +697,7 @@ impl RbdState { // Single global PFM work-list counter. let pfm_pairs_len = Tensor::vector( backend, - &[0u32], + [0u32], BufferUsages::STORAGE | BufferUsages::COPY_SRC, ) .unwrap(); @@ -720,9 +720,12 @@ impl RbdState { // Zero-initialized: a resize readback can race the first step. // Written as a storage buffer by `gpu_contact_plan`, read as a uniform // by every consumer. - let contact_plan = - Tensor::scalar(backend, ContactPlan::default(), storage | BufferUsages::UNIFORM) - .unwrap(); + let contact_plan = Tensor::scalar( + backend, + ContactPlan::default(), + storage | BufferUsages::UNIFORM, + ) + .unwrap(); let pfm_sort = PfmSortState::new(backend, pairs_capacity); let color_sorted_ids = Tensor::vector_uninit(backend, contacts_capacity, storage).unwrap(); // Zeroed (not uninit): the first frame's warmstart transfer walks the diff --git a/src_rbd_shaders/dynamics/joint_constraint_builder.rs b/src_rbd_shaders/dynamics/joint_constraint_builder.rs index e34959f9..9abfa033 100644 --- a/src_rbd_shaders/dynamics/joint_constraint_builder.rs +++ b/src_rbd_shaders/dynamics/joint_constraint_builder.rs @@ -344,7 +344,11 @@ impl JointConstraintHelper { impl JointConstraint { /// Solves a joint constraint. - pub fn solve_joint_constraint(&mut self, solver_vels: &mut ISliceMut, use_bias: bool) { + pub fn solve_joint_constraint( + &mut self, + solver_vels: &mut ISliceMut, + use_bias: bool, + ) { let mut solver_vel1 = solver_vels[self.solver_vel_a as usize]; let mut solver_vel2 = solver_vels[self.solver_vel_b as usize]; diff --git a/src_rbd_shaders/dynamics/multibody/contact_sensor.rs b/src_rbd_shaders/dynamics/multibody/contact_sensor.rs index e12a9874..89589861 100644 --- a/src_rbd_shaders/dynamics/multibody/contact_sensor.rs +++ b/src_rbd_shaders/dynamics/multibody/contact_sensor.rs @@ -1,9 +1,6 @@ //! Contact "force sensor" readout for RL observations. -use super::types::{ - MB_CONTACT_KIND_NORMAL, MultibodyContactConstraint, - MultibodyInfo, -}; +use super::types::{MB_CONTACT_KIND_NORMAL, MultibodyContactConstraint, MultibodyInfo}; use crate::utils::BatchIndices; use khal_std::glamx::UVec3; use khal_std::index::MaybeIndexUnchecked; diff --git a/src_rbd_shaders/dynamics/multibody/impulse_joint_constraints/helper.rs b/src_rbd_shaders/dynamics/multibody/impulse_joint_constraints/helper.rs index b4975a99..921c2745 100644 --- a/src_rbd_shaders/dynamics/multibody/impulse_joint_constraints/helper.rs +++ b/src_rbd_shaders/dynamics/multibody/impulse_joint_constraints/helper.rs @@ -278,15 +278,7 @@ pub(super) fn lock_jacobians_generic( out.ndofs_a = a.ndofs; out.j_id_a = j_id_a; if a.side_kind == SIDE_KIND_BODY { - fill_body_jacobians( - jacobians, - j_id_a, - a.side_id, - lin_jac, - ang_jac1, - mprops, - bix, - ); + fill_body_jacobians(jacobians, j_id_a, a.side_id, lin_jac, ang_jac1, mprops, bix); } else if a.side_kind == SIDE_KIND_MB { fill_mb_jacobians( jacobians, @@ -306,15 +298,7 @@ pub(super) fn lock_jacobians_generic( out.ndofs_b = b.ndofs; out.j_id_b = j_id_b; if b.side_kind == SIDE_KIND_BODY { - fill_body_jacobians( - jacobians, - j_id_b, - b.side_id, - lin_jac, - ang_jac2, - mprops, - bix, - ); + fill_body_jacobians(jacobians, j_id_b, b.side_id, lin_jac, ang_jac2, mprops, bix); } else if b.side_kind == SIDE_KIND_MB { fill_mb_jacobians( jacobians, diff --git a/src_rbd_shaders/dynamics/multibody/impulse_joint_constraints/kernels.rs b/src_rbd_shaders/dynamics/multibody/impulse_joint_constraints/kernels.rs index b55c8cbd..dd74d16c 100644 --- a/src_rbd_shaders/dynamics/multibody/impulse_joint_constraints/kernels.rs +++ b/src_rbd_shaders/dynamics/multibody/impulse_joint_constraints/kernels.rs @@ -4,8 +4,8 @@ use glamx::Vec4; use khal_std::glamx::UVec3; use khal_std::index::MaybeIndexUnchecked; -use khal_std::macros::{spirv, spirv_bindgen}; use khal_std::iter::StepRng; +use khal_std::macros::{spirv, spirv_bindgen}; use khal_std::sync::workgroup_memory_barrier_with_group_sync; use crate::Pose; diff --git a/src_rbd_shaders/dynamics/multibody/impulse_joint_constraints/update.rs b/src_rbd_shaders/dynamics/multibody/impulse_joint_constraints/update.rs index 90822997..08ac15d4 100644 --- a/src_rbd_shaders/dynamics/multibody/impulse_joint_constraints/update.rs +++ b/src_rbd_shaders/dynamics/multibody/impulse_joint_constraints/update.rs @@ -6,8 +6,8 @@ use khal_std::index::MaybeIndexUnchecked; use crate::dynamics::body::WorldMassProperties; use crate::dynamics::joint::{ANG_AXES_MASK, LIN_AXES_MASK, SPATIAL_DIM}; -use crate::utils::{ISlice, BodyIx}; use crate::utils::linalg::{MatSlice, VSlice, lu_solve_in_place}; +use crate::utils::{BodyIx, ISlice}; use crate::{DIM, Pose}; use super::super::types::{MultibodyInfo, MultibodyLinkStatic}; diff --git a/src_rbd_shaders/dynamics/multibody/integrate.rs b/src_rbd_shaders/dynamics/multibody/integrate.rs index 76fce2f7..6593b4d7 100644 --- a/src_rbd_shaders/dynamics/multibody/integrate.rs +++ b/src_rbd_shaders/dynamics/multibody/integrate.rs @@ -171,5 +171,4 @@ pub fn gpu_mb_integrate( } // num_ang == 0: no-op. } - } diff --git a/src_rbd_shaders/dynamics/multibody/solve_constraints.rs b/src_rbd_shaders/dynamics/multibody/solve_constraints.rs index b2c23a7d..7f22a835 100644 --- a/src_rbd_shaders/dynamics/multibody/solve_constraints.rs +++ b/src_rbd_shaders/dynamics/multibody/solve_constraints.rs @@ -781,8 +781,7 @@ pub fn gpu_mb_solve_contacts_delassus( if free_active { let cons = contact_constraints.read(cons_base + s as usize); - let mut free = - solver_vels.read(cons.free_body_id as usize); + let mut free = solver_vels.read(cons.free_body_id as usize); free.linear += cons.lin_jac * (cons.free_body_im * delta0); free.angular += cons.ii_ang_jac * delta0; if has_pair { @@ -805,12 +804,14 @@ pub fn gpu_mb_solve_contacts_delassus( a_shared.write(j as usize, a_shared.read(j as usize) + acc); } if lane < ndofs { - let col = contact_jac_cols - .read(jc_base + (s as usize) * 2 * dofs_stride + dofs_stride + lane as usize); + let col = contact_jac_cols.read( + jc_base + (s as usize) * 2 * dofs_stride + dofs_stride + lane as usize, + ); dof_v.write(lane as usize, dof_v.read(lane as usize) + delta0 * col); if has_pair { let col2 = contact_jac_cols.read( - jc_base + ((s + 1) as usize) * 2 * dofs_stride + jc_base + + ((s + 1) as usize) * 2 * dofs_stride + dofs_stride + lane as usize, ); diff --git a/src_rbd_shaders/dynamics/multibody/ws_soa.rs b/src_rbd_shaders/dynamics/multibody/ws_soa.rs index 2c331438..0eba6c21 100644 --- a/src_rbd_shaders/dynamics/multibody/ws_soa.rs +++ b/src_rbd_shaders/dynamics/multibody/ws_soa.rs @@ -116,8 +116,7 @@ impl WsAddr { /// `WS_QUADS` quads at `(L * num_batches + b) * WS_QUADS`. #[inline] pub fn at(&self, k: u32, quad: u32) -> usize { - ((self.base + k as usize) * self.stride as usize + self.shift as usize) - * WS_QUADS as usize + ((self.base + k as usize) * self.stride as usize + self.shift as usize) * WS_QUADS as usize + quad as usize } } diff --git a/src_viewer/ui.rs b/src_viewer/ui.rs index e79e58a4..9731ecce 100644 --- a/src_viewer/ui.rs +++ b/src_viewer/ui.rs @@ -267,7 +267,10 @@ fn performance_ui( row("Multibodies:", counts.multibodies); row("Multibody DOFs:", counts.multibody_dofs); row("MB contact slots:", counts.mb_contact_constraints); - row("MB contact capacity:", counts.mb_contact_constraints_capacity); + row( + "MB contact capacity:", + counts.mb_contact_constraints_capacity, + ); } if counts.particles > 0 {