From 93d2cf143285279e4c939a8afb50e520e2ec0af4 Mon Sep 17 00:00:00 2001 From: Haixuan Xavier Tao Date: Mon, 20 Jul 2026 15:25:29 +0200 Subject: [PATCH] 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