From 93d2cf143285279e4c939a8afb50e520e2ec0af4 Mon Sep 17 00:00:00 2001 From: Haixuan Xavier Tao Date: Mon, 20 Jul 2026 15:25:29 +0200 Subject: [PATCH 1/3] perf(mb): keep the contact-constraint back-solve column in registers MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `gpu_mb_finalize_contact_constraints` solved `M·column = Jᵀ` in place in the GLOBAL `contact_constraint_columns` buffer. The dense LU back-solve (permute + forward + backward substitution) read-modify-writes that column O(n²) times, and each dependent read stalls on an L2 round-trip: GPU L1 is write-evict for global stores, so a write isn't visible to the next read until it reaches L2. Hold the column in a per-thread local `[f32; 64]` instead (registers / L1-backed local memory) and write it out once at the end. `lu_solve_in_place` is `#[inline]`, so it lowers onto the local array with direct indexing — no cross-function storage traffic. Same arithmetic and iteration order → bit-identical `inv_lhs` and columns. Measured on a downstream fork (a tree-sparse LᵀDL variant of this kernel), the identical change cut the finalize kernel 2.36× (1449→615 µs/call, G1 29-DOF, 8192 envs). Upstream's dense LU does even more global RMW per solve, so the effect should be at least as large. I did not have an upstream-stock GPU setup to re-measure on — flagging that honestly. --- .../dynamics/multibody/contact_constraints.rs | 34 +++++++++---------- 1 file changed, 17 insertions(+), 17 deletions(-) diff --git a/src_rbd_shaders/dynamics/multibody/contact_constraints.rs b/src_rbd_shaders/dynamics/multibody/contact_constraints.rs index 922015a..7eaa57e 100644 --- a/src_rbd_shaders/dynamics/multibody/contact_constraints.rs +++ b/src_rbd_shaders/dynamics/multibody/contact_constraints.rs @@ -685,27 +685,27 @@ pub fn gpu_mb_finalize_contact_constraints( for s in 0..count { let col_offset = col_base + (s as usize) * dofs_stride; - // 1) Copy J^T row into the column buffer (it'll be overwritten by the - // LU solve with the M⁻¹·Jᵀ result). + // The LU back-solve read-modify-writes the column O(n²) times (permute + // + forward + backward substitution). Doing that in the GLOBAL + // `contact_constraint_columns` buffer costs an L2 round-trip on every + // step of the dependence chain — GPU L1 is write-evict for global + // stores, so each write is not seen by the next dependent read until it + // reaches L2. Hold the column in a per-thread LOCAL array instead + // (registers / L1-backed local memory), and write it out once at the + // end. Same arithmetic and iteration order → bit-identical result. + let mut col = [0.0f32; 64]; + // 1) J^T row into the local column. 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); + col[i as usize] = contact_constraint_jacs.read(col_offset + i as usize); } - // 2) Solve M · column = J^T (in place). - lu_solve_in_place( - mass_matrices, - m, - lu_pivots, - piv_offset, - contact_constraint_columns, - col_offset, - ); - // 3) inv_r_mb = J · column. + // 2) Solve M · column = J^T in place in the local vector. + lu_solve_in_place(mass_matrices, m, lu_pivots, piv_offset, &mut col, 0); + // 3) Write the finished column out + inv_r_mb = J · column in one pass. let mut inv_r_mb = 0.0f32; for i in 0..ndofs { - let j = contact_constraint_jacs.read(col_offset + i as usize); - let c = contact_constraint_columns.read(col_offset + i as usize); - inv_r_mb += j * c; + let c = col[i as usize]; + contact_constraint_columns.write(col_offset + i as usize, c); + inv_r_mb += contact_constraint_jacs.read(col_offset + i as usize) * c; } // 4) Add free body's contribution: im (since lin_jac is unit) + // ang_jac · ii_ang_jac. For self-contacts the B-side is folded into From c6ca275ef036c0c7db66f366453542f8358ce0f1 Mon Sep 17 00:00:00 2001 From: haixuantao Date: Fri, 7 Aug 2026 12:09:34 +0200 Subject: [PATCH 2/3] fix(mb): fixed-size-array LU back-solve variant for the local column MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit rust-gpu cannot unsize a thread-local [f32; 64] into the &mut [f32] argument of lu_solve_in_place, so the finalize kernel failed to compile to SPIR-V. Add lu_solve_in_place_local taking &mut [f32; N] directly — same arithmetic and iteration order. Co-Authored-By: Claude Fable 5 --- .../dynamics/multibody/contact_constraints.rs | 4 +- src_rbd_shaders/utils/linalg.rs | 45 +++++++++++++++++++ 2 files changed, 47 insertions(+), 2 deletions(-) diff --git a/src_rbd_shaders/dynamics/multibody/contact_constraints.rs b/src_rbd_shaders/dynamics/multibody/contact_constraints.rs index 7eaa57e..d62e720 100644 --- a/src_rbd_shaders/dynamics/multibody/contact_constraints.rs +++ b/src_rbd_shaders/dynamics/multibody/contact_constraints.rs @@ -22,7 +22,7 @@ use crate::dynamics::body::{Velocity, WorldMassProperties}; use crate::dynamics::joint::SPATIAL_DIM; use crate::queries::IndexedManifold; use crate::utils::BatchIndices; -use crate::utils::linalg::{MatSlice, lu_solve_in_place}; +use crate::utils::linalg::{MatSlice, lu_solve_in_place_local}; use crate::{ANG_DIM, AngVector, DIM, Pose, Vector, gcross, gdot}; use super::types::{ @@ -699,7 +699,7 @@ pub fn gpu_mb_finalize_contact_constraints( col[i as usize] = contact_constraint_jacs.read(col_offset + i as usize); } // 2) Solve M · column = J^T in place in the local vector. - lu_solve_in_place(mass_matrices, m, lu_pivots, piv_offset, &mut col, 0); + lu_solve_in_place_local(mass_matrices, m, lu_pivots, piv_offset, &mut col); // 3) Write the finished column out + inv_r_mb = J · column in one pass. let mut inv_r_mb = 0.0f32; for i in 0..ndofs { diff --git a/src_rbd_shaders/utils/linalg.rs b/src_rbd_shaders/utils/linalg.rs index 7918699..676a14b 100644 --- a/src_rbd_shaders/utils/linalg.rs +++ b/src_rbd_shaders/utils/linalg.rs @@ -721,6 +721,51 @@ pub fn lu_solve_in_place( } } +/// Variant of [`lu_solve_in_place`] whose right-hand side lives in a +/// thread-local fixed-size array instead of a storage buffer: rust-gpu cannot +/// unsize a local `[f32; N]` into a `&mut [f32]` argument, so the array is +/// taken by reference directly. Same arithmetic and iteration order. +#[inline] +pub fn lu_solve_in_place_local( + buf_m: &[f32], + m: MatSlice, + buf_pivots: &[u32], + pivots_offset: usize, + rhs: &mut [f32; N], +) { + let n = m.rows; + + // Permute rhs in place according to the recorded pivots. + for k in 0..n { + let p = buf_pivots.read(pivots_offset + k as usize); + if p != k { + let a = rhs[k as usize]; + rhs[k as usize] = rhs[p as usize]; + rhs[p as usize] = a; + } + } + + // Forward substitution: L · y = P · rhs (L is unit-lower — implicit diag = 1). + for i in 0..n { + let mut s = rhs[i as usize]; + for j in 0..i { + s -= buf_m.read(m.idx(i, j)) * rhs[j as usize]; + } + rhs[i as usize] = s; + } + + // Back substitution: U · x = y (reverse iteration — equivalent to `for ii in (0..n).rev()`). + for step in 0..n { + let ii = n - 1 - step; + let mut s = rhs[ii as usize]; + for j in (ii + 1)..n { + s -= buf_m.read(m.idx(ii, j)) * rhs[j as usize]; + } + let u = buf_m.read(m.idx(ii, ii)); + rhs[ii as usize] = if u != 0.0 { s / u } else { 0.0 }; + } +} + // // Workgroup-parallel variants. Mirror the sequential primitives above but // partition each iteration's work across `lanes` lanes of a SIMT workgroup. From 9adbedc733e2f070b401ff8de8a9e1b1b1987113 Mon Sep 17 00:00:00 2001 From: haixuantao Date: Fri, 7 Aug 2026 12:27:14 +0200 Subject: [PATCH 3/3] fix ci: allow clippy::manual_swap in lu_solve_in_place_local MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit rhs.swap() autoderefs the local array to a slice — the same unsizing cast rust-gpu rejects (verified: the SPIR-V build fails with it). Co-Authored-By: Claude Fable 5 --- src_rbd_shaders/utils/linalg.rs | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/src_rbd_shaders/utils/linalg.rs b/src_rbd_shaders/utils/linalg.rs index 676a14b..ac7af15 100644 --- a/src_rbd_shaders/utils/linalg.rs +++ b/src_rbd_shaders/utils/linalg.rs @@ -739,9 +739,15 @@ pub fn lu_solve_in_place_local( for k in 0..n { let p = buf_pivots.read(pivots_offset + k as usize); if p != k { - let a = rhs[k as usize]; - rhs[k as usize] = rhs[p as usize]; - rhs[p as usize] = a; + // NOTE: not `rhs.swap(..)` — the method autoderefs the local + // array into a slice, which is the exact `*[f32; N]` → `*[f32]` + // cast rust-gpu rejects. + #[allow(clippy::manual_swap)] + { + let a = rhs[k as usize]; + rhs[k as usize] = rhs[p as usize]; + rhs[p as usize] = a; + } } }