From c8398abebf3adf04dfdd0448af7850da7ac213db Mon Sep 17 00:00:00 2001 From: pepijn kooijmans Date: Fri, 17 Jul 2026 14:31:00 +0200 Subject: [PATCH 1/4] =?UTF-8?q?perf(rbd):=20flat=201-D=20narrow-phase=20di?= =?UTF-8?q?spatch=20=E2=80=94=20pack=20warps=20across=20batches?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The narrow-phase kernels dispatched [max_len/64, num_batches, 1]: one 64-lane workgroup per batch rounded up from that batch's live pair count. A batched robot env has ~7 pairs, so at 2048+ envs the GPU ran thousands of workgroups at ~11% lane occupancy and narrow-phase scaled ~linearly with env count (0.02 ms at 1 env -> 18.6 ms at 4096). gpu_flatten_batches_dispatch (one thread, same style as the existing max-scan init kernels) now builds exclusive prefix offsets over the per-batch work-lists plus a flat [total/64, 1, 1] grid; the classify, deferred and PFM kernels walk 0..total and recover (batch, item) with a binary search over the offsets. Warps fill with real pairs from consecutive batches. Buffer layout is unchanged — only the dispatch shape and index math moved. RTX 5090, 12-DOF biped batch, dt=5ms, steady-state contacts: narrow-phase 2048 envs: 10.31 -> 5.09 ms 4096 envs: 18.57 -> 6.94 ms whole step 2048 envs: 22.45 -> 16.50 ms/step (91k -> 124k env-steps/s) Physics bit-identical (robot trajectory + cube settle unchanged to printed precision). Co-Authored-By: Claude Fable 5 --- src_rbd/broad_phase/narrow_phase.rs | 50 +++++-- src_rbd/pipeline/insertion_removal.rs | 6 + src_rbd/pipeline/rbd_state.rs | 6 + src_rbd/pipeline/rbd_state_from_rapier.rs | 6 + src_rbd/pipeline/rbd_step.rs | 6 +- src_rbd_shaders/broad_phase/narrow_phase.rs | 141 +++++++++++++++----- 6 files changed, 165 insertions(+), 50 deletions(-) diff --git a/src_rbd/broad_phase/narrow_phase.rs b/src_rbd/broad_phase/narrow_phase.rs index b37cfb5..53e0e8c 100644 --- a/src_rbd/broad_phase/narrow_phase.rs +++ b/src_rbd/broad_phase/narrow_phase.rs @@ -4,9 +4,9 @@ use crate::math::Pose; use crate::queries::GpuIndexedContact; use crate::shaders::PaddedVector; use crate::shaders::broad_phase::{ - CollisionPair, GpuInitPfmPfmDispatch, GpuNarrowPhaseInitContactsDispatch, GpuNarrowPhasePfmPfm, - GpuNarrowPhaseShapeShape, GpuNarrowPhaseShapeShapeDeferred, GpuResetNarrowPhase, - NarrowPhasePfmPair, + CollisionPair, GpuFlattenBatchesDispatch, GpuNarrowPhaseInitContactsDispatch, + GpuNarrowPhasePfmPfm, GpuNarrowPhaseShapeShape, GpuNarrowPhaseShapeShapeDeferred, + GpuResetNarrowPhase, NarrowPhasePfmPair, }; use crate::shaders::shapes::Shape; use khal::Shader; @@ -22,7 +22,11 @@ pub struct GpuNarrowPhase { /// `pfm_pairs` work-list. Split from `narrow_phase` to fit 8 storage buffers. narrow_phase_deferred: GpuNarrowPhaseShapeShapeDeferred, narrow_phase_pfm_pfm: GpuNarrowPhasePfmPfm, - init_pfm_pfm_indirect_args: GpuInitPfmPfmDispatch, + /// Builds the flat 1-D dispatch grid + prefix offsets for a per-batch + /// work-list (used for both the collision pairs and the PFM pairs), so the + /// kernels pack items from many batches into full warps instead of one + /// mostly-idle workgroup per batch. + flatten_batches: GpuFlattenBatchesDispatch, init_contacts_indirect_args: GpuNarrowPhaseInitContactsDispatch, } @@ -37,8 +41,8 @@ impl GpuNarrowPhase { vertices: &Tensor, indices: &Tensor, collision_pairs: &Tensor, - collision_pairs_len: &Tensor, - collision_pairs_indirect: &Tensor<[u32; 3]>, + collision_pairs_len: &mut Tensor, + collision_pairs_indirect: &mut Tensor<[u32; 3]>, contacts: &mut Tensor, contacts_len: &mut Tensor, contacts_indirect: &mut Tensor<[u32; 3]>, @@ -48,16 +52,30 @@ impl GpuNarrowPhase { batch_indices: &Tensor, collider_parent: &Tensor, collider_materials: &Tensor, + pairs_offsets: &mut Tensor, + pfm_offsets: &mut Tensor, ) -> Result<(), GpuBackendError> { let num_batches = contacts_len.len() as u32; self.reset_narrow_phase .call(pass, [1u32, num_batches, 1], contacts_len, pfm_pairs_len)?; - self.narrow_phase.call( + // The broad phase wrote a `[max/64, num_batches, 1]` grid into + // `collision_pairs_indirect`; rewrite it (and derive the offsets) for + // the flat layout. Nothing else consumes the batched form. + self.flatten_batches.call( pass, + 1u32, + collision_pairs_len, + pairs_offsets, collision_pairs_indirect, + batch_indices, + )?; + + self.narrow_phase.call( + pass, + &*collision_pairs_indirect, collision_pairs, - collision_pairs_len, + pairs_offsets, poses, shapes, contacts, @@ -71,9 +89,9 @@ impl GpuNarrowPhase { // separate dispatch so each pass fits 8 storage buffers). self.narrow_phase_deferred.call( pass, - collision_pairs_indirect, + &*collision_pairs_indirect, collision_pairs, - collision_pairs_len, + pairs_offsets, poses, shapes, pfm_pairs, @@ -83,15 +101,21 @@ impl GpuNarrowPhase { indices, )?; - self.init_pfm_pfm_indirect_args - .call(pass, 1u32, pfm_pairs_len, pfm_pairs_indirect)?; + self.flatten_batches.call( + pass, + 1u32, + pfm_pairs_len, + pfm_offsets, + pfm_pairs_indirect, + batch_indices, + )?; self.narrow_phase_pfm_pfm.call( pass, &*pfm_pairs_indirect, contacts, contacts_len, pfm_pairs, - pfm_pairs_len, + pfm_offsets, batch_indices, vertices, indices, diff --git a/src_rbd/pipeline/insertion_removal.rs b/src_rbd/pipeline/insertion_removal.rs index b5e1332..051891b 100644 --- a/src_rbd/pipeline/insertion_removal.rs +++ b/src_rbd/pipeline/insertion_removal.rs @@ -168,6 +168,10 @@ impl RbdState { BufferUsages::STORAGE | BufferUsages::COPY_SRC, ) .unwrap(); + let pairs_flat_offsets = + Tensor::vector_uninit(backend, num_batches + 1, BufferUsages::STORAGE).unwrap(); + let pfm_flat_offsets = + Tensor::vector_uninit(backend, num_batches + 1, BufferUsages::STORAGE).unwrap(); let old_constraints = Tensor::vector_uninit(backend, collisions_capacity * num_batches, storage).unwrap(); let old_constraint_builders = @@ -262,6 +266,8 @@ impl RbdState { pfm_pairs, pfm_pairs_len, pfm_pairs_indirect, + pairs_flat_offsets, + pfm_flat_offsets, old_constraints, old_constraint_builders, old_constraints_counts, diff --git a/src_rbd/pipeline/rbd_state.rs b/src_rbd/pipeline/rbd_state.rs index 73e0c92..0deea59 100644 --- a/src_rbd/pipeline/rbd_state.rs +++ b/src_rbd/pipeline/rbd_state.rs @@ -184,6 +184,12 @@ pub struct RbdState { pub(super) pfm_pairs: Tensor, pub(super) pfm_pairs_len: Tensor, pub(super) pfm_pairs_indirect: Tensor<[u32; 3]>, + /// Flat-dispatch prefix offsets (`num_batches + 1`) over the per-batch + /// collision-pair / PFM work-lists, rebuilt on the GPU each step by + /// `gpu_flatten_batches_dispatch` so the narrow-phase kernels can pack + /// items from many batches into full warps. + pub(super) pairs_flat_offsets: Tensor, + pub(super) pfm_flat_offsets: Tensor, pub(super) contacts: Tensor, pub(super) contacts_len: Tensor, pub(super) contacts_indirect: Tensor<[u32; 3]>, diff --git a/src_rbd/pipeline/rbd_state_from_rapier.rs b/src_rbd/pipeline/rbd_state_from_rapier.rs index 92bdd1e..84c3b50 100644 --- a/src_rbd/pipeline/rbd_state_from_rapier.rs +++ b/src_rbd/pipeline/rbd_state_from_rapier.rs @@ -595,6 +595,10 @@ impl RbdState { BufferUsages::STORAGE | BufferUsages::COPY_SRC, ) .unwrap(); + let pairs_flat_offsets = + Tensor::vector_uninit(backend, num_batches + 1, BufferUsages::STORAGE).unwrap(); + let pfm_flat_offsets = + Tensor::vector_uninit(backend, num_batches + 1, BufferUsages::STORAGE).unwrap(); let old_constraints = Tensor::vector_uninit( backend, capacities.collisions_capacity * num_batches, @@ -753,6 +757,8 @@ impl RbdState { pfm_pairs, pfm_pairs_len, pfm_pairs_indirect, + pairs_flat_offsets, + pfm_flat_offsets, old_constraints, old_constraint_builders, old_constraints_counts, diff --git a/src_rbd/pipeline/rbd_step.rs b/src_rbd/pipeline/rbd_step.rs index 64817d2..59352af 100644 --- a/src_rbd/pipeline/rbd_step.rs +++ b/src_rbd/pipeline/rbd_step.rs @@ -186,8 +186,8 @@ impl RbdPipeline { &state.vertex_buffers, &state.index_buffers, &state.collision_pairs, - &state.collision_pairs_len, - &state.collision_pairs_indirect, + &mut state.collision_pairs_len, + &mut state.collision_pairs_indirect, &mut state.contacts, &mut state.contacts_len, &mut state.contacts_indirect, @@ -197,6 +197,8 @@ impl RbdPipeline { &state.batch_indices, &state.collider_parent, &state.collider_materials, + &mut state.pairs_flat_offsets, + &mut state.pfm_flat_offsets, )?; drop(pass); diff --git a/src_rbd_shaders/broad_phase/narrow_phase.rs b/src_rbd_shaders/broad_phase/narrow_phase.rs index 20a4ead..9234176 100644 --- a/src_rbd_shaders/broad_phase/narrow_phase.rs +++ b/src_rbd_shaders/broad_phase/narrow_phase.rs @@ -70,6 +70,65 @@ pub fn gpu_narrow_phase_init_contacts_dispatch( *indirect_args.at_mut(2) = 1; } +/// Builds the flat-dispatch layout for a per-batch work-list: exclusive prefix +/// offsets (so item `t` of the flat range maps back to a batch via +/// [`find_batch`]) and the matching 1-D indirect grid. +/// +/// This replaces the max-over-batches indirect grids for the narrow-phase +/// kernels: with `[max/64, num_batches, 1]` every batch rounds its handful of +/// pairs up to a full 64-lane workgroup (a robot env has ~7 pairs → ~11% lane +/// occupancy, thousands of near-empty workgroups). The flat grid packs items +/// from consecutive batches into the same warps: `[total/64, 1, 1]`. +/// +/// Serial over batches in one thread — same pattern (and cost) as the existing +/// `gpu_narrow_phase_init_contacts_dispatch` max-scan. +#[spirv_bindgen] +#[spirv(compute(threads(1)))] +pub fn gpu_flatten_batches_dispatch( + // NOTE: `lens` is mutable only for `atomic_load_u32` (see the note on + // `gpu_narrow_phase_init_contacts_dispatch`). + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] lens: &mut [u32], + // `num_batches + 1` entries; `offsets[num_batches]` is the total. + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] offsets: &mut [u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] indirect_args: &mut [u32; 3], + #[spirv(uniform, descriptor_set = 0, binding = 3)] batch_ids: &BatchIndices, +) { + let num_batches = lens.len(); + // Same clamp as the consuming kernels: a batch's list may overflow its + // capacity slot; the overflowing tail was never written and must not be + // walked. + let capacity = batch_ids.contacts_batch_capacity; + let mut total = 0u32; + for i in 0..num_batches { + offsets.write(i, total); + total += atomic_load_u32(lens.at_mut(i)).min(capacity); + } + offsets.write(num_batches, total); + *indirect_args.at_mut(0) = total.div_ceil(WORKGROUP_SIZE); + *indirect_args.at_mut(1) = 1; + *indirect_args.at_mut(2) = 1; +} + +/// Largest `b` with `offsets[b] <= t` — the batch owning flat item `t`. +/// Invariant: `offsets[0] == 0 <= t < offsets[num_batches]`. +fn find_batch(offsets: &[u32], num_batches: u32, t: u32) -> u32 { + let mut lo = 0u32; + let mut hi = num_batches; + // Bounded loop instead of `while` (see the trimesh BVH walk for why). + for _ in 0..32 { + if lo + 1 >= hi { + break; + } + let mid = (lo + hi) / 2; + if offsets.read(mid as usize) <= t { + lo = mid; + } else { + hi = mid; + } + } + lo +} + const PREDICTION: f32 = 2.0e-3; // TODO: make the prediction configurable. /// Narrow phase, pass 1 of 2: analytic shape-shape contacts for ball / cuboid @@ -83,7 +142,10 @@ pub fn gpu_narrow_phase_shape_shape( #[spirv(global_invocation_id)] invocation_id: UVec3, #[spirv(num_workgroups)] num_workgroups: UVec3, #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] collision_pairs: &[CollisionPair], - #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] collision_pairs_len: &[u32], + // Flat-dispatch prefix offsets from `gpu_flatten_batches_dispatch` + // (`num_batches + 1` entries; replaces the per-batch `collision_pairs_len`, + // which it already folds in, clamped to capacity). + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] pairs_offsets: &[u32], #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] poses: &[Pose], #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] shapes: &[Shape], #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] contacts: &mut [IndexedManifold], @@ -97,23 +159,25 @@ pub fn gpu_narrow_phase_shape_shape( collider_materials: &[ColliderMaterial], ) { let num_threads = num_workgroups.x * WORKGROUP_SIZE; - let batch_id = invocation_id.y; let contacts_batch_capacity = batch_ids.contacts_batch_capacity as usize; - let collision_pairs = batch_ids.contact_batch(batch_id, collision_pairs); - let poses = batch_ids.coll_batch(batch_id, poses); - let shapes = batch_ids.coll_batch(batch_id, shapes); - let collider_materials = batch_ids.coll_batch(batch_id, collider_materials); - let mut contacts = batch_ids.contact_batch_mut(batch_id, contacts); - let contacts_len = contacts_len.at_mut(batch_id as usize); + // Flat over all batches' pairs: consecutive lanes take consecutive pairs + // regardless of which batch owns them, so warps stay packed even when each + // batch only has a handful. + let num_batches = pairs_offsets.len() - 1; + let total = pairs_offsets.read(num_batches); + + for t in StepRng::new(invocation_id.x..total, num_threads) { + let batch_id = find_batch(pairs_offsets, num_batches as u32, t); + let i = t - pairs_offsets.read(batch_id as usize); - // 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); + let collision_pairs = batch_ids.contact_batch(batch_id, collision_pairs); + let poses = batch_ids.coll_batch(batch_id, poses); + let shapes = batch_ids.coll_batch(batch_id, shapes); + let collider_materials = batch_ids.coll_batch(batch_id, collider_materials); + let mut contacts = SliceMut(&mut *contacts, batch_ids.contacts_start(batch_id)); + let contacts_len = contacts_len.at_mut(batch_id as usize); - for i in StepRng::new(invocation_id.x..len, num_threads) { let pair = collision_pairs[i as usize]; // Resolve the parent rigid-bodies here (the broad phase no longer does) // and skip pairs whose colliders share the same body. @@ -200,7 +264,8 @@ pub fn gpu_narrow_phase_shape_shape_deferred( #[spirv(global_invocation_id)] invocation_id: UVec3, #[spirv(num_workgroups)] num_workgroups: UVec3, #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] collision_pairs: &[CollisionPair], - #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] collision_pairs_len: &[u32], + // Flat-dispatch prefix offsets (see `gpu_narrow_phase_shape_shape`). + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] pairs_offsets: &[u32], #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] poses: &[Pose], #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] shapes: &[Shape], #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] @@ -214,25 +279,25 @@ pub fn gpu_narrow_phase_shape_shape_deferred( #[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 contacts_batch_capacity = batch_ids.contacts_batch_capacity as usize; - - let collision_pairs = batch_ids.contact_batch(batch_id, collision_pairs); - let poses = batch_ids.coll_batch(batch_id, poses); - let shapes = batch_ids.coll_batch(batch_id, shapes); - let mut pfm_pairs = batch_ids.contact_batch_mut(batch_id, pfm_pairs); - let pfm_pairs_len = pfm_pairs_len.at_mut(batch_id as usize); - let len = collision_pairs_len - .read(batch_id as usize) - .min(contacts_batch_capacity as u32); + let num_batches = pairs_offsets.len() - 1; + let total = pairs_offsets.read(num_batches); // NOTE: same-body collider pairs are *not* filtered in this pass — it is // already at the 8-storage-buffer WebGPU limit and can't take the // `collider_parent` binding. The complex pairs it emits are filtered // downstream in `gpu_narrow_phase_pfm_pfm` (which has room) before any // contact is written. - for i in StepRng::new(invocation_id.x..len, num_threads) { + for t in StepRng::new(invocation_id.x..total, num_threads) { + let batch_id = find_batch(pairs_offsets, num_batches as u32, t); + let i = t - pairs_offsets.read(batch_id as usize); + + let collision_pairs = batch_ids.contact_batch(batch_id, collision_pairs); + let poses = batch_ids.coll_batch(batch_id, poses); + let shapes = batch_ids.coll_batch(batch_id, shapes); + let mut pfm_pairs = SliceMut(&mut *pfm_pairs, batch_ids.contacts_start(batch_id)); + let pfm_pairs_len = pfm_pairs_len.at_mut(batch_id as usize); + let pair = collision_pairs[i as usize]; let pose1 = poses[pair.colliders.x as usize]; let pose2 = poses[pair.colliders.y as usize]; @@ -541,7 +606,9 @@ pub fn gpu_narrow_phase_pfm_pfm( #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] contacts: &mut [IndexedManifold], #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] contacts_len: &mut [u32], #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] pfm_pairs: &[NarrowPhasePfmPair], - #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] pfm_pairs_len: &[u32], + // Flat-dispatch prefix offsets over the per-batch PFM work-lists (see + // `gpu_narrow_phase_shape_shape`; replaces the per-batch `pfm_pairs_len`). + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] pfm_offsets: &[u32], // NOTE: we assume that max_pfm_pairs == contacts_batch_capacity // And we assume all batch dimensions are given the same buffer allocation sizes // (i.e. the same `contacts_batch_capacity`). @@ -557,16 +624,20 @@ pub fn gpu_narrow_phase_pfm_pfm( collider_materials: &[ColliderMaterial], ) { 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); - let pfm_pairs_len = pfm_pairs_len.read(batch_id as usize); + let num_batches = pfm_offsets.len() - 1; + let total = pfm_offsets.read(num_batches); + + for t in StepRng::new(invocation_id.x..total, num_threads) { + let batch_id = find_batch(pfm_offsets, num_batches as u32, t); + let i = t - pfm_offsets.read(batch_id as usize); + + let mut contacts = SliceMut(&mut *contacts, batch_ids.contacts_start(batch_id)); + let collider_materials = batch_ids.coll_batch(batch_id, collider_materials); + let pfm_pairs = batch_ids.contact_batch(batch_id, pfm_pairs); + let contacts_len = contacts_len.at_mut(batch_id as usize); - for i in StepRng::new(invocation_id.x..pfm_pairs_len, num_threads) { let pair = pfm_pairs[i as usize]; // Resolve the parent rigid-bodies and skip same-body collider pairs. This // is where the deferred (PFM / trimesh / polyline) pairs get the same-body From 4e83d28168877f2535ed77db1c6b9080fe0081bf Mon Sep 17 00:00:00 2001 From: pepijn kooijmans Date: Fri, 17 Jul 2026 15:47:48 +0200 Subject: [PATCH 2/4] =?UTF-8?q?perf(rbd):=20fuse=20the=20solver=20color=20?= =?UTF-8?q?loops=20=E2=80=94=20one=20workgroup=20per=20env,=20colors=20via?= =?UTF-8?q?=20workgroup=20barriers?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The contact solver ran reset_color + num_colors x (solve + inc_color) as sequential device-wide dispatches, three times per substep — ~230 dependent launches per step whose cost was launch/barrier latency, not lane math (the color kernels touch ~1-2 contacts per env; flattening their lanes was measured neutral). But the ordering the colors enforce is only ever *within* an env — bodies are never shared across batches — so the barrier only needs workgroup scope. gpu_warmstart_fused / gpu_step_gauss_seidel_fused run one 64-lane workgroup per batch: stage the batch's solver_vels in shared memory (which the workgroup barrier fences), walk colors 1..=num_colors with a barrier between, write back once. Each replaces a whole reset+N x (solve+inc) chain with a single dispatch; the per-color path remains as fallback for batches with more than FUSED_SOLVE_MAX_BODIES (64) bodies. num_colors reaches the kernels via a tiny uniform refreshed only when the Grow policy raises it. RTX 5090, 12-DOF biped batch, dt=5ms, steady-state, wgpu — on top of the flat narrow-phase dispatch: 2048 envs: 16.50 -> 8.68 ms/step (124k -> 236k env-steps/s) 4096 envs: 27.27 -> 13.72 ms/step (150k -> 298k env-steps/s) Cumulative vs pre-flatten baseline: 91k -> 298k env-steps/s at 4096 (3.3x). Physics bit-identical (robot trajectory + cube settle unchanged to printed precision) — same-color contacts touch disjoint bodies, so intra-color order cannot matter, and inter-color order is preserved. Co-Authored-By: Claude Fable 5 --- src_rbd/dynamics/solver.rs | 91 +++++++++++---- src_rbd/pipeline/insertion_removal.rs | 8 ++ src_rbd/pipeline/rbd_state.rs | 4 + src_rbd/pipeline/rbd_state_from_rapier.rs | 8 ++ src_rbd/pipeline/rbd_step.rs | 22 ++++ src_rbd_shaders/dynamics/solver.rs | 129 +++++++++++++++++++++- 6 files changed, 242 insertions(+), 20 deletions(-) diff --git a/src_rbd/dynamics/solver.rs b/src_rbd/dynamics/solver.rs index 3b89840..93838e5 100644 --- a/src_rbd/dynamics/solver.rs +++ b/src_rbd/dynamics/solver.rs @@ -13,7 +13,8 @@ use crate::shaders::dynamics::{ GpuApplySolverVelsInc, GpuInitSolverBodies, GpuInitSolverVelsInc, GpuIntegrateLinearized, GpuRemoveCfmAndBiasKernel, GpuSolverCleanup, GpuSolverCountConstraints, GpuSolverFinalize, GpuSolverIncColor, GpuSolverInitConstraints, GpuSolverResetColor, GpuSolverSortConstraints, - GpuSolverUpdateConstraints, GpuStepGaussSeidel, GpuWarmstart, LocalMassProperties, + GpuSolverUpdateConstraints, GpuStepGaussSeidel, GpuStepGaussSeidelFused, GpuWarmstart, + GpuWarmstartFused, LocalMassProperties, RbdSimParams, TwoBodyConstraint, TwoBodyConstraintBuilder, Velocity, WorldMassProperties, }; use crate::utils::{GpuPrefixSum, PrefixSumWorkspace}; @@ -40,6 +41,12 @@ pub struct GpuSolver { warmstart: GpuWarmstart, /// Gauss-Seidel iteration step (sequential per color). step_gauss_seidel: GpuStepGaussSeidel, + /// One-workgroup-per-env fused variants of the color loops (velocities + /// staged in shared memory, colors ordered by workgroup barriers). Used + /// when every batch's body count fits the shared stage; the per-color + /// dispatch chain above is the fallback. + warmstart_fused: GpuWarmstartFused, + step_gauss_seidel_fused: GpuStepGaussSeidelFused, /// Initializes solver velocity increments. init_solver_vels_inc: GpuInitSolverVelsInc, /// Seeds the COM-centered solver poses from the body world poses @@ -71,6 +78,13 @@ pub struct SolverArgs<'a> { pub contacts_len: &'a Tensor, /// Indirect dispatch arguments based on contact count. pub contacts_len_indirect: &'a Tensor<[u32; 3]>, + /// `num_colors` as a single-element uniform, consumed by the fused color + /// kernels (must equal [`Self::num_colors`]). + pub num_colors_uniform: &'a Tensor, + /// Run the fused one-workgroup-per-env color kernels instead of the + /// per-color dispatch chain. Only valid when every batch's body count is + /// <= `FUSED_SOLVE_MAX_BODIES`. + pub fuse_color_loops: bool, /// Solver constraints (output from constraint initialization). pub constraints: &'a mut Tensor, /// Builder data for initializing constraints. @@ -316,19 +330,32 @@ impl GpuSolver { args.batch_indices, )?; joint_solver.update(pass, &mut joint_args, args.solver_body_poses)?; - self.reset_color.call(pass, 1u32, args.curr_color)?; - for _ in 0..args.num_colors { - self.warmstart.call( + if args.fuse_color_loops { + self.warmstart_fused.call( pass, - args.contacts_len_indirect, + [64u32, args.num_batches, 1], args.constraints, args.solver_vels, args.constraints_colors, args.contacts_len, - args.curr_color, + args.num_colors_uniform, args.batch_indices, )?; - self.inc_color.call(pass, 1u32, args.curr_color)? + } else { + self.reset_color.call(pass, 1u32, args.curr_color)?; + for _ in 0..args.num_colors { + self.warmstart.call( + pass, + args.contacts_len_indirect, + args.constraints, + args.solver_vels, + args.constraints_colors, + args.contacts_len, + args.curr_color, + args.batch_indices, + )?; + self.inc_color.call(pass, 1u32, args.curr_color)? + } } /* @@ -336,19 +363,32 @@ impl GpuSolver { */ mb_phase!(substep_solve_with_bias); joint_solver.solve(pass, &mut joint_args, args.solver_vels, true)?; - self.reset_color.call(pass, 1u32, args.curr_color)?; - for _ in 0..args.num_colors { - self.step_gauss_seidel.call( + if args.fuse_color_loops { + self.step_gauss_seidel_fused.call( pass, - args.contacts_len_indirect, + [64u32, args.num_batches, 1], args.constraints, args.solver_vels, args.constraints_colors, args.contacts_len, - args.curr_color, + args.num_colors_uniform, args.batch_indices, )?; - self.inc_color.call(pass, 1u32, args.curr_color)? + } else { + self.reset_color.call(pass, 1u32, args.curr_color)?; + for _ in 0..args.num_colors { + self.step_gauss_seidel.call( + pass, + args.contacts_len_indirect, + args.constraints, + args.solver_vels, + args.constraints_colors, + args.contacts_len, + args.curr_color, + args.batch_indices, + )?; + self.inc_color.call(pass, 1u32, args.curr_color)? + } } /* @@ -376,19 +416,32 @@ impl GpuSolver { args.contacts_len, args.batch_indices, )?; - self.reset_color.call(pass, 1u32, args.curr_color)?; - for _ in 0..args.num_colors { - self.step_gauss_seidel.call( + if args.fuse_color_loops { + self.step_gauss_seidel_fused.call( pass, - args.contacts_len_indirect, + [64u32, args.num_batches, 1], args.constraints, args.solver_vels, args.constraints_colors, args.contacts_len, - args.curr_color, + args.num_colors_uniform, args.batch_indices, )?; - self.inc_color.call(pass, 1u32, args.curr_color)? + } else { + self.reset_color.call(pass, 1u32, args.curr_color)?; + for _ in 0..args.num_colors { + self.step_gauss_seidel.call( + pass, + args.contacts_len_indirect, + args.constraints, + args.solver_vels, + args.constraints_colors, + args.contacts_len, + args.curr_color, + args.batch_indices, + )?; + self.inc_color.call(pass, 1u32, args.curr_color)? + } } } diff --git a/src_rbd/pipeline/insertion_removal.rs b/src_rbd/pipeline/insertion_removal.rs index 051891b..5026dc0 100644 --- a/src_rbd/pipeline/insertion_removal.rs +++ b/src_rbd/pipeline/insertion_removal.rs @@ -172,6 +172,12 @@ impl RbdState { Tensor::vector_uninit(backend, num_batches + 1, BufferUsages::STORAGE).unwrap(); let pfm_flat_offsets = Tensor::vector_uninit(backend, num_batches + 1, BufferUsages::STORAGE).unwrap(); + let num_colors_uniform = Tensor::scalar( + backend, + capacities.solver_colors + 1, + BufferUsages::UNIFORM | BufferUsages::COPY_DST, + ) + .unwrap(); let old_constraints = Tensor::vector_uninit(backend, collisions_capacity * num_batches, storage).unwrap(); let old_constraint_builders = @@ -268,6 +274,8 @@ impl RbdState { pfm_pairs_indirect, pairs_flat_offsets, pfm_flat_offsets, + num_colors_uniform, + num_colors_uniform_cpu: capacities.solver_colors + 1, old_constraints, old_constraint_builders, old_constraints_counts, diff --git a/src_rbd/pipeline/rbd_state.rs b/src_rbd/pipeline/rbd_state.rs index 0deea59..6bedead 100644 --- a/src_rbd/pipeline/rbd_state.rs +++ b/src_rbd/pipeline/rbd_state.rs @@ -190,6 +190,10 @@ pub struct RbdState { /// items from many batches into full warps. pub(super) pairs_flat_offsets: Tensor, pub(super) pfm_flat_offsets: Tensor, + /// `max_colors + 1` as a uniform for the fused color-loop kernels, plus a + /// CPU mirror so the tensor is only recreated when the count grows. + pub(super) num_colors_uniform: Tensor, + pub(super) num_colors_uniform_cpu: u32, pub(super) contacts: Tensor, pub(super) contacts_len: Tensor, pub(super) contacts_indirect: Tensor<[u32; 3]>, diff --git a/src_rbd/pipeline/rbd_state_from_rapier.rs b/src_rbd/pipeline/rbd_state_from_rapier.rs index 84c3b50..1bae60d 100644 --- a/src_rbd/pipeline/rbd_state_from_rapier.rs +++ b/src_rbd/pipeline/rbd_state_from_rapier.rs @@ -599,6 +599,12 @@ impl RbdState { Tensor::vector_uninit(backend, num_batches + 1, BufferUsages::STORAGE).unwrap(); let pfm_flat_offsets = Tensor::vector_uninit(backend, num_batches + 1, BufferUsages::STORAGE).unwrap(); + let num_colors_uniform = Tensor::scalar( + backend, + capacities.solver_colors + 1, + BufferUsages::UNIFORM | BufferUsages::COPY_DST, + ) + .unwrap(); let old_constraints = Tensor::vector_uninit( backend, capacities.collisions_capacity * num_batches, @@ -759,6 +765,8 @@ impl RbdState { pfm_pairs_indirect, pairs_flat_offsets, pfm_flat_offsets, + num_colors_uniform, + num_colors_uniform_cpu: capacities.solver_colors + 1, old_constraints, old_constraint_builders, old_constraints_counts, diff --git a/src_rbd/pipeline/rbd_step.rs b/src_rbd/pipeline/rbd_step.rs index 59352af..4e34295 100644 --- a/src_rbd/pipeline/rbd_step.rs +++ b/src_rbd/pipeline/rbd_step.rs @@ -217,6 +217,9 @@ impl RbdPipeline { contacts: &state.contacts, contacts_len: &state.contacts_len, contacts_len_indirect: &state.contacts_indirect, + num_colors_uniform: &state.num_colors_uniform, + // Prep never runs the color loops, so no fusion here. + fuse_color_loops: false, constraints: &mut state.new_constraints, constraint_builders: &mut state.new_constraint_builders, sim_params: &state.sim_params, @@ -293,11 +296,30 @@ impl RbdPipeline { let num_colors = stats.num_colors; + // Keep the fused-kernel color-count uniform in sync (the count only + // moves when the Grow policy raises `max_colors`, so this is rare). + if state.num_colors_uniform_cpu != num_colors { + state.num_colors_uniform = vortx::tensor::Tensor::scalar( + backend, + num_colors, + BufferUsages::UNIFORM | BufferUsages::COPY_DST, + ) + .unwrap(); + state.num_colors_uniform_cpu = num_colors; + } + // One 64-lane workgroup per env stages that env's velocities in shared + // memory; batches with more bodies than the stage fall back to the + // per-color dispatch chain. + let fuse_color_loops = state.num_colliders_per_batch + <= crate::shaders::dynamics::FUSED_SOLVE_MAX_BODIES as u32; + // Create solver_args for solve phase (after coloring is complete) let solver_args = SolverArgs { contacts: &state.contacts, contacts_len: &state.contacts_len, contacts_len_indirect: &state.contacts_indirect, + num_colors_uniform: &state.num_colors_uniform, + fuse_color_loops, constraints: &mut state.new_constraints, constraint_builders: &mut state.new_constraint_builders, sim_params: &state.sim_params, diff --git a/src_rbd_shaders/dynamics/solver.rs b/src_rbd_shaders/dynamics/solver.rs index 6dfbe62..28a15a0 100644 --- a/src_rbd_shaders/dynamics/solver.rs +++ b/src_rbd_shaders/dynamics/solver.rs @@ -6,7 +6,11 @@ use khal_std::glamx::UVec3; use khal_std::macros::{spirv, spirv_bindgen}; use crate::{AngVector, Pose, Vector}; -use khal_std::{index::MaybeIndexUnchecked, iter::StepRng, sync::atomic_add_u32}; +use khal_std::{ + index::MaybeIndexUnchecked, + iter::StepRng, + sync::{atomic_add_u32, workgroup_memory_barrier_with_group_sync}, +}; use super::body::{LocalMassProperties, Velocity, WorldMassProperties}; use super::constraint::{TwoBodyConstraint, TwoBodyConstraintBuilder}; @@ -440,6 +444,129 @@ pub fn gpu_step_gauss_seidel( } } + +/// Max rigid bodies per batch supported by the fused (one-workgroup-per-env) +/// contact solvers below — bounds the shared-memory velocity stage. The host +/// falls back to the per-color dispatch loop when a batch exceeds it. +pub const FUSED_SOLVE_MAX_BODIES: usize = 64; + +/// Fused warmstart: the whole per-substep color loop in ONE dispatch, one +/// 64-lane workgroup per batch (env). +/// +/// The per-color dispatch chain exists to order contacts that share a body — +/// but bodies are only ever shared *within* a batch, so the ordering barrier +/// only needs workgroup scope, not a device-wide dispatch boundary. Each +/// workgroup stages its batch's `solver_vels` in shared memory (which is what +/// the workgroup barrier fences), walks colors `1..=num_colors` with a barrier +/// between them, and writes velocities back once. This replaces +/// `reset_color + num_colors x (warmstart + inc_color)` — dozens of dependent +/// dispatches whose cost was launch/barrier latency, not lane math. +/// +/// Constraint data needs no fence: each contact belongs to exactly one color +/// and is touched by exactly one lane in the whole kernel. +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn gpu_warmstart_fused( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] constraints: &[TwoBodyConstraint], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] solver_vels: &mut [Velocity], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] constraints_colors: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] contacts_len: &[u32], + #[spirv(uniform, descriptor_set = 0, binding = 4)] num_colors: &u32, + #[spirv(uniform, descriptor_set = 0, binding = 5)] batch_ids: &BatchIndices, + #[spirv(workgroup)] vels_smem: &mut [Velocity; FUSED_SOLVE_MAX_BODIES], +) { + // Grid is [1, num_batches, 1] workgroups, so global x == lane in workgroup. + let lane = invocation_id.x; + let batch_id = invocation_id.y; + let num_bodies = batch_ids.colliders_len; + let cbase = batch_ids.contacts_start(batch_id); + let vbase = batch_ids.coll_start(batch_id); + let len = contacts_len + .read(batch_id as usize) + .min(batch_ids.contacts_batch_capacity); + let nc = *num_colors; + + for i in StepRng::new(lane..num_bodies, 64) { + vels_smem[i as usize] = *solver_vels.at(vbase + i as usize); + } + workgroup_memory_barrier_with_group_sync(); + + // NOTE: bounded loop over a uniform value; every lane sees the same bound + // so the barriers stay in uniform control flow. + for color in 1..=nc { + for i in StepRng::new(lane..len, 64) { + if constraints_colors.read(cbase + i as usize) == color { + let constraint = constraints.at(cbase + i as usize); + let a = constraint.solver_body_a as usize; + let b = constraint.solver_body_b as usize; + let mut va = vels_smem[a]; + let mut vb = vels_smem[b]; + constraint.warmstart_constraint(&mut va, &mut vb); + vels_smem[a] = va; + vels_smem[b] = vb; + } + } + workgroup_memory_barrier_with_group_sync(); + } + + for i in StepRng::new(lane..num_bodies, 64) { + solver_vels.write(vbase + i as usize, vels_smem[i as usize]); + } +} + +/// Fused Gauss-Seidel: same one-workgroup-per-env structure as +/// [`gpu_warmstart_fused`] (see there for the why), replacing the per-color +/// `step_gauss_seidel` dispatch chain in both the with-bias and no-bias phases. +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn gpu_step_gauss_seidel_fused( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + constraints: &mut [TwoBodyConstraint], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] solver_vels: &mut [Velocity], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] constraints_colors: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] contacts_len: &[u32], + #[spirv(uniform, descriptor_set = 0, binding = 4)] num_colors: &u32, + #[spirv(uniform, descriptor_set = 0, binding = 5)] batch_ids: &BatchIndices, + #[spirv(workgroup)] vels_smem: &mut [Velocity; FUSED_SOLVE_MAX_BODIES], +) { + let lane = invocation_id.x; + let batch_id = invocation_id.y; + let num_bodies = batch_ids.colliders_len; + let cbase = batch_ids.contacts_start(batch_id); + let vbase = batch_ids.coll_start(batch_id); + let len = contacts_len + .read(batch_id as usize) + .min(batch_ids.contacts_batch_capacity); + let nc = *num_colors; + + for i in StepRng::new(lane..num_bodies, 64) { + vels_smem[i as usize] = *solver_vels.at(vbase + i as usize); + } + workgroup_memory_barrier_with_group_sync(); + + for color in 1..=nc { + for i in StepRng::new(lane..len, 64) { + if constraints_colors.read(cbase + i as usize) == color { + let constraint = constraints.at_mut(cbase + i as usize); + let a = constraint.solver_body_a as usize; + let b = constraint.solver_body_b as usize; + let mut va = vels_smem[a]; + let mut vb = vels_smem[b]; + constraint.solve_constraint_gauss_seidel(&mut va, &mut vb); + vels_smem[a] = va; + vels_smem[b] = vb; + } + } + workgroup_memory_barrier_with_group_sync(); + } + + for i in StepRng::new(lane..num_bodies, 64) { + solver_vels.write(vbase + i as usize, vels_smem[i as usize]); + } +} + /// Integrates velocity to update poses. #[spirv_bindgen] #[spirv(compute(threads(64)))] From a69d6a8e9332acf6605a6d34c5486f3773fc405f Mon Sep 17 00:00:00 2001 From: pepijn kooijmans Date: Fri, 17 Jul 2026 16:00:30 +0200 Subject: [PATCH 3/4] perf(rbd): absorb per-substep contact dispatches into the fused kernels + flat refit_leaves MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Round 2 of the one-workgroup-per-env fusion: - gpu_warmstart_fused absorbs gpu_apply_solver_vels_inc (increment added while staging velocities into shared memory, same bodies_len bound) and gpu_solver_update_constraints (each lane refreshes its own contacts' constraints before the color walk — a constraint is only read by the lane that updated it, so no extra barrier). Exactly fills the 8-storage-buffer budget. - gpu_step_gauss_seidel_fused_no_bias replaces the final color loop, with gpu_remove_cfm_and_bias_kernel absorbed as a prologue (same lane-locality argument). - gpu_lbvh_refit_leaves goes flat: the per-batch collider count is uniform, so batch/index recover by division. A 14-collider env used 14 of 64 lanes per workgroup, and per-leaf work is mesh-AABB computation — this pass was 17% of the profiled frame. Removes another ~12 dispatches per step in the fused path; the standalone kernels remain for the >FUSED_SOLVE_MAX_BODIES fallback. RTX 5090, 12-DOF biped batch, dt=5ms, wgpu, on top of the round-1 fusion: 2048 envs: 8.68 -> 6.52 ms/step (236k -> 314k env-steps/s) Physics bit-identical (robot trajectory endpoint + cube settle to printed precision). Co-Authored-By: Claude Fable 5 --- src_rbd/broad_phase/lbvh.rs | 3 +- src_rbd/dynamics/solver.rs | 71 +++++++++++++--------- src_rbd_shaders/broad_phase/lbvh.rs | 23 ++++--- src_rbd_shaders/dynamics/solver.rs | 94 ++++++++++++++++++++++++++++- 4 files changed, 153 insertions(+), 38 deletions(-) diff --git a/src_rbd/broad_phase/lbvh.rs b/src_rbd/broad_phase/lbvh.rs index ead559c..c9412b4 100644 --- a/src_rbd/broad_phase/lbvh.rs +++ b/src_rbd/broad_phase/lbvh.rs @@ -232,7 +232,8 @@ impl Lbvh { let mut pass = encoder.begin_pass("[RBD] lbvh-refit_leaves", timestamps.as_deref_mut()); self.shaders.refit_leaves.call( &mut pass, - [colliders_per_batch, num_batches, 1], + // Flat over all batches' leaves (kernel recovers batch by division). + [colliders_per_batch * num_batches, 1, 1], poses, shapes, &state.sorted_colliders, diff --git a/src_rbd/dynamics/solver.rs b/src_rbd/dynamics/solver.rs index 93838e5..4c48a3f 100644 --- a/src_rbd/dynamics/solver.rs +++ b/src_rbd/dynamics/solver.rs @@ -13,8 +13,8 @@ use crate::shaders::dynamics::{ GpuApplySolverVelsInc, GpuInitSolverBodies, GpuInitSolverVelsInc, GpuIntegrateLinearized, GpuRemoveCfmAndBiasKernel, GpuSolverCleanup, GpuSolverCountConstraints, GpuSolverFinalize, GpuSolverIncColor, GpuSolverInitConstraints, GpuSolverResetColor, GpuSolverSortConstraints, - GpuSolverUpdateConstraints, GpuStepGaussSeidel, GpuStepGaussSeidelFused, GpuWarmstart, - GpuWarmstartFused, LocalMassProperties, + GpuSolverUpdateConstraints, GpuStepGaussSeidel, GpuStepGaussSeidelFused, + GpuStepGaussSeidelFusedNoBias, GpuWarmstart, GpuWarmstartFused, LocalMassProperties, RbdSimParams, TwoBodyConstraint, TwoBodyConstraintBuilder, Velocity, WorldMassProperties, }; use crate::utils::{GpuPrefixSum, PrefixSumWorkspace}; @@ -47,6 +47,7 @@ pub struct GpuSolver { /// dispatch chain above is the fallback. warmstart_fused: GpuWarmstartFused, step_gauss_seidel_fused: GpuStepGaussSeidelFused, + step_gauss_seidel_fused_no_bias: GpuStepGaussSeidelFusedNoBias, /// Initializes solver velocity increments. init_solver_vels_inc: GpuInitSolverVelsInc, /// Seeds the COM-centered solver poses from the body world poses @@ -307,28 +308,36 @@ impl GpuSolver { * P1/F1 — integrate velocities (apply `a · dt'` / gravity increment). */ mb_phase!(substep_integrate_velocities); - self.apply_solver_vels_inc.call( - pass, - [args.num_colliders, args.num_batches, 1], - args.solver_vels, - args.solver_vels_inc, - args.batch_indices, - )?; + if !args.fuse_color_loops { + // Fused path folds the increment into gpu_warmstart_fused's + // shared-memory staging. + self.apply_solver_vels_inc.call( + pass, + [args.num_colliders, args.num_batches, 1], + args.solver_vels, + args.solver_vels_inc, + args.batch_indices, + )?; + } /* * P2/F2 — build + warmstart constraints. */ mb_phase!(substep_build_constraints); - self.update_constraints.call( - pass, - args.contacts_len_indirect, - args.constraints, - args.constraint_builders, - args.contacts_len, - args.solver_body_poses, - args.sim_params, - args.batch_indices, - )?; + if !args.fuse_color_loops { + // Fused path folds the constraint refresh into + // gpu_warmstart_fused's prologue. + self.update_constraints.call( + pass, + args.contacts_len_indirect, + args.constraints, + args.constraint_builders, + args.contacts_len, + args.solver_body_poses, + args.sim_params, + args.batch_indices, + )?; + } joint_solver.update(pass, &mut joint_args, args.solver_body_poses)?; if args.fuse_color_loops { self.warmstart_fused.call( @@ -340,6 +349,10 @@ impl GpuSolver { args.contacts_len, args.num_colors_uniform, args.batch_indices, + args.solver_vels_inc, + args.constraint_builders, + args.solver_body_poses, + args.sim_params, )?; } else { self.reset_color.call(pass, 1u32, args.curr_color)?; @@ -409,15 +422,19 @@ impl GpuSolver { */ mb_phase!(substep_solve_no_bias); joint_solver.solve(pass, &mut joint_args, args.solver_vels, false)?; - self.remove_cfm_and_bias_kernel.call( - pass, - args.contacts_len_indirect, - args.constraints, - args.contacts_len, - args.batch_indices, - )?; + if !args.fuse_color_loops { + // Fused path folds the CFM/bias strip into the no-bias kernel's + // prologue. + self.remove_cfm_and_bias_kernel.call( + pass, + args.contacts_len_indirect, + args.constraints, + args.contacts_len, + args.batch_indices, + )?; + } if args.fuse_color_loops { - self.step_gauss_seidel_fused.call( + self.step_gauss_seidel_fused_no_bias.call( pass, [64u32, args.num_batches, 1], args.constraints, diff --git a/src_rbd_shaders/broad_phase/lbvh.rs b/src_rbd_shaders/broad_phase/lbvh.rs index 350a49d..835b49b 100644 --- a/src_rbd_shaders/broad_phase/lbvh.rs +++ b/src_rbd_shaders/broad_phase/lbvh.rs @@ -301,17 +301,24 @@ pub fn gpu_lbvh_refit_leaves( // workgroup. // 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 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 mut tree = SliceMut(tree, root_id(colliders_start) as usize); - - for i in StepRng::new(invocation_id.x..num_colliders, num_threads) { + // Flat over all batches' leaves: the per-batch collider count is uniform, + // so batch/index recover by division — no per-batch workgroup rounding + // (a 14-collider robot env used 14 of 64 lanes per workgroup, and the + // per-leaf work is mesh-AABB computation, which is genuinely expensive). + let num_batches = (sorted_colliders.len() / batch_ids.colliders_batch_capacity as usize) as u32; + let total = num_colliders * num_batches; + + for t in StepRng::new(invocation_id.x..total, num_threads) { + let batch_id = t / num_colliders; + let i = t % num_colliders; + let colliders_start = batch_ids.coll_start(batch_id) as u32; + 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 mut tree = SliceMut(&mut *tree, root_id(colliders_start) as usize); let curr_leaf_id = first_leaf_id + i; let leaf_collider = sorted_colliders[i as usize]; let leaf_pose = poses[leaf_collider as usize]; diff --git a/src_rbd_shaders/dynamics/solver.rs b/src_rbd_shaders/dynamics/solver.rs index 28a15a0..6518ffc 100644 --- a/src_rbd_shaders/dynamics/solver.rs +++ b/src_rbd_shaders/dynamics/solver.rs @@ -468,27 +468,61 @@ pub const FUSED_SOLVE_MAX_BODIES: usize = 64; #[spirv(compute(threads(64)))] pub fn gpu_warmstart_fused( #[spirv(global_invocation_id)] invocation_id: UVec3, - #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] constraints: &[TwoBodyConstraint], + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + constraints: &mut [TwoBodyConstraint], #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] solver_vels: &mut [Velocity], #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] constraints_colors: &[u32], #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] contacts_len: &[u32], #[spirv(uniform, descriptor_set = 0, binding = 4)] num_colors: &u32, #[spirv(uniform, descriptor_set = 0, binding = 5)] batch_ids: &BatchIndices, + // Absorbed `gpu_apply_solver_vels_inc`: the increment is added while + // staging velocities into shared memory. + #[spirv(storage_buffer, descriptor_set = 0, binding = 6)] solver_vels_inc: &[Velocity], + // Absorbed `gpu_solver_update_constraints`: each lane refreshes its own + // contacts' constraints before the color walk. No barrier needed — a + // constraint is only ever read by the same lane that updated it. + #[spirv(storage_buffer, descriptor_set = 0, binding = 7)] + constraint_builders: &[TwoBodyConstraintBuilder], + #[spirv(storage_buffer, descriptor_set = 0, binding = 8)] solver_body_poses: &[Pose], + #[spirv(storage_buffer, descriptor_set = 0, binding = 9)] all_params: &[RbdSimParams], #[spirv(workgroup)] vels_smem: &mut [Velocity; FUSED_SOLVE_MAX_BODIES], ) { // Grid is [1, num_batches, 1] workgroups, so global x == lane in workgroup. let lane = invocation_id.x; let batch_id = invocation_id.y; let num_bodies = batch_ids.colliders_len; + let num_dyn_bodies = batch_ids.bodies_len; let cbase = batch_ids.contacts_start(batch_id); let vbase = batch_ids.coll_start(batch_id); let len = contacts_len .read(batch_id as usize) .min(batch_ids.contacts_batch_capacity); let nc = *num_colors; + let params = all_params.at(batch_id as usize); for i in StepRng::new(lane..num_bodies, 64) { - vels_smem[i as usize] = *solver_vels.at(vbase + i as usize); + let g = vbase + i as usize; + let mut v = *solver_vels.at(g); + // Same bound as the absorbed kernel: increments only exist for the + // first `bodies_len` slots. + if i < num_dyn_bodies { + let inc = solver_vels_inc.at(g); + v.linear += inc.linear; + v.angular += inc.angular; + } + vels_smem[i as usize] = v; + } + + { + let solver_body_poses = batch_ids.coll_batch(batch_id, solver_body_poses); + for i in StepRng::new(lane..len, 64) { + let ci = cbase + i as usize; + constraints.at_mut(ci).update_constraint( + constraint_builders.at(ci), + &solver_body_poses, + params, + ); + } } workgroup_memory_barrier_with_group_sync(); @@ -567,6 +601,62 @@ pub fn gpu_step_gauss_seidel_fused( } } +/// The no-bias variant of [`gpu_step_gauss_seidel_fused`]: identical color +/// walk, plus the absorbed `gpu_remove_cfm_and_bias_kernel` as a prologue +/// (each lane strips CFM/bias from its own contacts before solving them, so +/// no barrier is needed between the two). +#[spirv_bindgen] +#[spirv(compute(threads(64)))] +pub fn gpu_step_gauss_seidel_fused_no_bias( + #[spirv(global_invocation_id)] invocation_id: UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] + constraints: &mut [TwoBodyConstraint], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] solver_vels: &mut [Velocity], + #[spirv(storage_buffer, descriptor_set = 0, binding = 2)] constraints_colors: &[u32], + #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] contacts_len: &[u32], + #[spirv(uniform, descriptor_set = 0, binding = 4)] num_colors: &u32, + #[spirv(uniform, descriptor_set = 0, binding = 5)] batch_ids: &BatchIndices, + #[spirv(workgroup)] vels_smem: &mut [Velocity; FUSED_SOLVE_MAX_BODIES], +) { + let lane = invocation_id.x; + let batch_id = invocation_id.y; + let num_bodies = batch_ids.colliders_len; + let cbase = batch_ids.contacts_start(batch_id); + let vbase = batch_ids.coll_start(batch_id); + let len = contacts_len + .read(batch_id as usize) + .min(batch_ids.contacts_batch_capacity); + let nc = *num_colors; + + for i in StepRng::new(lane..num_bodies, 64) { + vels_smem[i as usize] = *solver_vels.at(vbase + i as usize); + } + for i in StepRng::new(lane..len, 64) { + constraints.at_mut(cbase + i as usize).remove_cfm_and_bias(); + } + workgroup_memory_barrier_with_group_sync(); + + for color in 1..=nc { + for i in StepRng::new(lane..len, 64) { + if constraints_colors.read(cbase + i as usize) == color { + let constraint = constraints.at_mut(cbase + i as usize); + let a = constraint.solver_body_a as usize; + let b = constraint.solver_body_b as usize; + let mut va = vels_smem[a]; + let mut vb = vels_smem[b]; + constraint.solve_constraint_gauss_seidel(&mut va, &mut vb); + vels_smem[a] = va; + vels_smem[b] = vb; + } + } + workgroup_memory_barrier_with_group_sync(); + } + + for i in StepRng::new(lane..num_bodies, 64) { + solver_vels.write(vbase + i as usize, vels_smem[i as usize]); + } +} + /// Integrates velocity to update poses. #[spirv_bindgen] #[spirv(compute(threads(64)))] From f1dbab85cbc48512ba6b3e36e8fd1b51c82b857d Mon Sep 17 00:00:00 2001 From: pepijn kooijmans Date: Fri, 17 Jul 2026 18:05:41 +0200 Subject: [PATCH 4/4] fix(rbd): apply the per-batch stride to collider_parent reads in the narrow phase MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Collision pairs carry env-local collider ids, and collider_parent is batch-strided like every other per-collider buffer (its construction comment says so: 'Env-local body slot; the kernels apply the per-batch stride') — but the classify and pfm_pfm kernels read it unsliced, resolving every batch's parents through batch 0's table. With identical environments the tables coincide and nothing observable goes wrong — which is why every bit-exactness check passed. With heterogeneous environments (the point of per-env MJCF insertion) contacts are silently mis-parented: a pair that is same-body in batch 0 gets skipped in batches where it isn't, and solved impulses can target the wrong bodies. Repro (now a stacking test): env0 = body with two glued boxes + a single box; env1 = single box dropped onto a two-glued-box body (equal body and collider counts, different parent tables). Pre-fix the falling box never rests on the stack; post-fix it settles on top (z=0.75). Identical-env scenes are bit-identical before/after, as the aliasing argument predicts. Co-Authored-By: Claude Fable 5 --- src_rbd_shaders/broad_phase/narrow_phase.rs | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/src_rbd_shaders/broad_phase/narrow_phase.rs b/src_rbd_shaders/broad_phase/narrow_phase.rs index 9234176..7610a0a 100644 --- a/src_rbd_shaders/broad_phase/narrow_phase.rs +++ b/src_rbd_shaders/broad_phase/narrow_phase.rs @@ -180,9 +180,14 @@ pub fn gpu_narrow_phase_shape_shape( let pair = collision_pairs[i as usize]; // Resolve the parent rigid-bodies here (the broad phase no longer does) - // and skip pairs whose colliders share the same body. - let body1 = collider_parent.read(pair.colliders.x as usize); - let body2 = collider_parent.read(pair.colliders.y as usize); + // and skip pairs whose colliders share the same body. Pair ids are + // env-local, and `collider_parent` is batch-strided like the other + // per-collider buffers — an unsliced read here silently returned + // batch 0's parents for every batch (masked whenever all envs are + // identical, wrong the moment they aren't). + let coll_base = batch_ids.coll_start(batch_id); + let body1 = collider_parent.read(coll_base + pair.colliders.x as usize); + let body2 = collider_parent.read(coll_base + pair.colliders.y as usize); if body1 == body2 { continue; } @@ -643,8 +648,9 @@ pub fn gpu_narrow_phase_pfm_pfm( // is where the deferred (PFM / trimesh / polyline) pairs get the same-body // filtering that the analytic pass does inline — the broad phase no longer // does it, and the deferred pass has no spare storage binding for it. - let body1 = collider_parent.read(pair.colliders.x as usize); - let body2 = collider_parent.read(pair.colliders.y as usize); + let coll_base = batch_ids.coll_start(batch_id); + let body1 = collider_parent.read(coll_base + pair.colliders.x as usize); + let body2 = collider_parent.read(coll_base + pair.colliders.y as usize); if body1 == body2 { continue; }