perf(mb): keep the contact-constraint back-solve column in registers

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

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<const N: usize>(
+    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.