diff --git a/src_rbd_shaders/dynamics/multibody/contact_constraints.rs b/src_rbd_shaders/dynamics/multibody/contact_constraints.rs index 922015a..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::{ @@ -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_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 { - 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 diff --git a/src_rbd_shaders/utils/linalg.rs b/src_rbd_shaders/utils/linalg.rs index 7918699..ac7af15 100644 --- a/src_rbd_shaders/utils/linalg.rs +++ b/src_rbd_shaders/utils/linalg.rs @@ -721,6 +721,57 @@ 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 { + // 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; + } + } + } + + // 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.