diff --git a/src_rbd_shaders/dynamics/multibody/gravity_and_lu.rs b/src_rbd_shaders/dynamics/multibody/gravity_and_lu.rs index b2eeb78d..4340ffff 100644 --- a/src_rbd_shaders/dynamics/multibody/gravity_and_lu.rs +++ b/src_rbd_shaders/dynamics/multibody/gravity_and_lu.rs @@ -21,9 +21,7 @@ use crate::utils::linalg::{MAX_MB_DOFS, MatSlice, fill_par, gemv_tr_spatial_spli use crate::utils::{BatchIndices, Slice}; use crate::{AngVector, Vector, gcross_av}; -use super::lu::{ - LANES, lu_apply_pivots, lu_factor_in_shared, lu_triangular_solve_in_place, sm_idx, -}; +use super::lu::{LANES, NO_PARENT, ltdl_factor_in_shared, ltdl_solve_in_shared, sm_idx}; use super::types::{MultibodyInfo, MultibodyLinkStatic, MultibodyLinkWorkspace}; /// Fused gravity / Coriolis-force assembly + LU factor + LU solve. @@ -235,20 +233,49 @@ pub fn gpu_mb_gravity_and_lu( } x.write(lane as usize, gen_forces.read(rhs_offset + lane as usize)); } + + // Per-DOF parent array (tree metadata for the sparse LᵀDL factor/solves), + // written into the old pivots buffer so downstream solves read the tree + // from the binding they already have. A joint's DOFs chain internally; + // its first DOF hangs off the last DOF of the nearest ancestor link that + // has any (welded links contribute none). + if lane == 0 { + let mut k = 0u32; + while k < num_links { + let stat = stat_slice[k as usize]; + if stat.ndofs > 0 { + let mut p = NO_PARENT; + if k != 0 { + let mut a = stat.parent_link_id; + for _ in 0..MAX_MB_DOFS as u32 { + let s = stat_slice[a as usize]; + if s.ndofs > 0 { + p = s.assembly_id + s.ndofs - 1; + break; + } + if a == 0 { + break; + } + a = s.parent_link_id; + } + } + lu_pivots.write(piv_offset + stat.assembly_id as usize, p); + for t in 1..stat.ndofs { + lu_pivots.write( + piv_offset + (stat.assembly_id + t) as usize, + stat.assembly_id + t - 1, + ); + } + } + k += 1; + } + } workgroup_memory_barrier_with_group_sync(); - lu_factor_in_shared( - ndofs, - lane, - mat, - lu_pivots, - piv_offset, - pivot_row_shared, - inv_akk_shared, - ); + ltdl_factor_in_shared(ndofs, lane, mat, lu_pivots, piv_offset); - // Persist LU factors to global memory (joint / contact constraint init - // reuses them for unit-RHS solves). + // Persist the LᵀDL factors to global memory (joint / contact constraint + // init reuses them for unit-RHS solves). if lane < ndofs { for r in 0..ndofs { mass_matrices.write(m_view.idx(r, lane), mat.read(sm_idx(r, lane))); @@ -256,9 +283,11 @@ pub fn gpu_mb_gravity_and_lu( } // ---- Phase 4: solve M·x = τ for the gravity rhs. ---- - lu_apply_pivots(ndofs, lane, lu_pivots, piv_offset, x); - lu_triangular_solve_in_place(ndofs, lane, mat, x, partial); + ltdl_solve_in_shared(ndofs, lane, mat, lu_pivots, piv_offset, x); + let _ = partial; + let _ = pivot_row_shared; + let _ = inv_akk_shared; if lane < ndofs { gen_forces.write(rhs_offset + lane as usize, x.read(lane as usize)); } diff --git a/src_rbd_shaders/dynamics/multibody/lu.rs b/src_rbd_shaders/dynamics/multibody/lu.rs index af339475..2f6e3797 100644 --- a/src_rbd_shaders/dynamics/multibody/lu.rs +++ b/src_rbd_shaders/dynamics/multibody/lu.rs @@ -1,18 +1,38 @@ -//! LU decomposition + solve. +//! Tree-sparse LᵀDL factorization + solve for the multibody mass matrix. //! -//! Split into two kernels so the factorization can be reused across multiple -//! right-hand sides within a frame (e.g. gravity τ, contact impulses, …). -//! Dispatched one workgroup per `(multibody, batch)` pair. +//! The mass matrix of a kinematic tree has branch-induced sparsity: +//! `M[i, j] != 0` only when DOFs `i` and `j` lie on the same root-to-leaf +//! path. Eliminating leaves-to-root (descending DOF order, since parents are +//! assembled before children) factors `M = Lᵀ·D·L` with ZERO fill-in — +//! `L[k, i]` is nonzero only for `i` on `k`'s ancestor chain (Featherstone, +//! RBDA §6; MuJoCo's `mj_factorM`). Factor cost drops from O(n³) to +//! O(n·depth²) and every solve from O(n²) to O(n·depth). +//! +//! The per-DOF parent index (`u32::MAX` at roots) is stored in the buffer +//! that used to hold the LU pivots, so every downstream solve reads the tree +//! from a binding it already had. Factors live in the dense `mass_matrices` +//! storage: strict-lower ancestor-chain entries hold `L` (unit diagonal +//! implied), the diagonal holds `D`, and the upper triangle keeps stale +//! symmetric copies of M that no solve reads. +//! +//! The factor/solve loops are data-dependent but contain NO barriers, so a +//! single lane runs them serially — legal under WGSL uniformity rules (the +//! old dense kernels needed fixed 64-iteration loops only because they +//! barriered inside). At n ≤ 64, depth ≈ 7, the serial sparse path does a +//! few hundred FMAs versus the dense path's ~192 workgroup barriers. use khal_std::index::MaybeIndexUnchecked; use khal_std::sync::workgroup_memory_barrier_with_group_sync; use crate::utils::linalg::MAX_MB_DOFS; -/// Workgroup width for the parallelised LU kernels. Must match the +/// Workgroup width for the multibody solver kernels. Must match the /// `threads(N, 1, 1)` attribute and `MB_LU_LANES` on the host side. pub(super) const LANES: u32 = 64; +/// Parent sentinel for root DOFs in the per-DOF parent array. +pub(super) const NO_PARENT: u32 = u32::MAX; + /// Side length of the workgroup-shared matrix tile. Must equal both the lane /// count and the maximum supported `ndofs`. const MAX_DOFS_U32: u32 = MAX_MB_DOFS as u32; @@ -23,168 +43,96 @@ pub(super) fn sm_idx(r: u32, c: u32) -> usize { (c * MAX_DOFS_U32 + r) as usize } -/// Workgroup-parallel LU factorization on the shared `mat` tile in place. +/// Tree-sparse LᵀDL factorization of the shared `mat` tile in place, run +/// serially by lane 0 (no barriers inside — the loops are data-dependent). +/// `parents` holds the per-DOF parent indices (`NO_PARENT` at roots), read +/// from the pivots buffer at `parents_offset`. +/// +/// Ends with a workgroup barrier so every lane sees the factors. #[inline] -pub(super) fn lu_factor_in_shared( +pub(super) fn ltdl_factor_in_shared( n: u32, lane: u32, - mat: &mut [f32; MAX_MB_DOFS * MAX_MB_DOFS], - pivots_dst: &mut [u32], - pivots_offset: usize, - pivot_row_shared: &mut u32, - inv_akk_shared: &mut f32, + mat: &mut impl MaybeIndexUnchecked, + parents: &[u32], + parents_offset: usize, ) { - // NOTE: fixed number of iterations for uniform control flow. - // TODO(PERF): on non-web platforms we could just use `n` as the upper bound. - for k in 0..MAX_DOFS_U32 { - let active = k < n; - if active && lane == 0 { - let mut pivot_row = k; - let mut pivot_val = { - let v = mat.read(sm_idx(k, k)); - if v >= 0.0 { v } else { -v } - }; - for i in (k + 1)..n { - let v = mat.read(sm_idx(i, k)); - let av = if v >= 0.0 { v } else { -v }; - if av > pivot_val { - pivot_val = av; - pivot_row = i; + if lane == 0 { + // Eliminate leaves-to-root: descending k, updating only the + // ancestor-chain entries of rows above. + for step in 0..n { + let k = n - 1 - step; + let d = mat.read(sm_idx(k, k)); + let inv_d = if d != 0.0 { 1.0 / d } else { 0.0 }; + let mut i = parents.read(parents_offset + k as usize); + // Bounded loops (parents are strictly decreasing) so a corrupt + // parent array can't hang the GPU. + for _ in 0..MAX_DOFS_U32 { + if i == NO_PARENT { + break; } - } - *pivot_row_shared = pivot_row; - pivots_dst.write(pivots_offset + k as usize, pivot_row); - } - workgroup_memory_barrier_with_group_sync(); - let pivot_row = *pivot_row_shared; - - if active && pivot_row != k && lane < n { - let c = lane; - let a = mat.read(sm_idx(k, c)); - let b = mat.read(sm_idx(pivot_row, c)); - mat.write(sm_idx(k, c), b); - mat.write(sm_idx(pivot_row, c), a); - } - workgroup_memory_barrier_with_group_sync(); - - if active && lane == 0 { - let akk = mat.read(sm_idx(k, k)); - *inv_akk_shared = if akk != 0.0 { 1.0 / akk } else { 0.0 }; - } - workgroup_memory_barrier_with_group_sync(); - let inv_akk = *inv_akk_shared; - - if active { - let r = k + 1 + lane; - if r < n { - let v = mat.read(sm_idx(r, k)) * inv_akk; - mat.write(sm_idx(r, k), v); - } - } - workgroup_memory_barrier_with_group_sync(); - - if active { - let j = k + 1 + lane; - if j < n { - let akj = mat.read(sm_idx(k, j)); - for r in (k + 1)..n { - let lik = mat.read(sm_idx(r, k)); - let v = mat.read(sm_idx(r, j)) - lik * akj; - mat.write(sm_idx(r, j), v); + let a = mat.read(sm_idx(k, i)) * inv_d; + let mut j = i; + for _ in 0..MAX_DOFS_U32 { + if j == NO_PARENT { + break; + } + let v = mat.read(sm_idx(i, j)) - a * mat.read(sm_idx(k, j)); + mat.write(sm_idx(i, j), v); + j = parents.read(parents_offset + j as usize); } + mat.write(sm_idx(k, i), a); + i = parents.read(parents_offset + i as usize); } } - workgroup_memory_barrier_with_group_sync(); } + workgroup_memory_barrier_with_group_sync(); } -/// Inner of the workgroup-parallel triangular solves used by the LU -/// solve/factor-and-solve kernels. Preconditions: `mat` -/// already holds the LU factors and `x` already holds the permuted rhs. +/// Solve `M·x = b` in place on the shared tile holding LᵀDL factors, run +/// serially by lane 0. Ends with a workgroup barrier. #[inline] -pub(super) fn lu_triangular_solve_in_place( +pub(super) fn ltdl_solve_in_shared( n: u32, lane: u32, - mat: &[f32; MAX_MB_DOFS * MAX_MB_DOFS], - x: &mut [f32; MAX_MB_DOFS], - partial: &mut [f32; LANES as usize], + mat: &impl MaybeIndexUnchecked, + parents: &[u32], + parents_offset: usize, + x: &mut impl MaybeIndexUnchecked, ) { - // NOTE: fixed number of iterations for uniform control flow. - // TODO(PERF): on non-web platforms we could just use `n` as the upper bound. - for i in 0..MAX_DOFS_U32 { - let active = i < n; - let s = if active && lane < i { - mat.read(sm_idx(i, lane)) * x.read(lane as usize) - } else { - 0.0f32 - }; - partial.write(lane as usize, s); - workgroup_memory_barrier_with_group_sync(); - for step in 0..6u32 { - let stride = 1u32 << (5 - step); - if lane < stride { - let v = partial.read(lane as usize) + partial.read((lane + stride) as usize); - partial.write(lane as usize, v); - } - workgroup_memory_barrier_with_group_sync(); - } - if active && lane == 0 { - let cur = x.read(i as usize); - x.write(i as usize, cur - partial.read(0)); - } - workgroup_memory_barrier_with_group_sync(); - } - - // NOTE: fixed number of iterations for uniform control flow. - // TODO(PERF): on non-web platforms we could just use `n` as the upper bound. - for step in 0..MAX_DOFS_U32 { - let active = step < n; - // For dummy iterations (step >= n), `i` is not meaningful — guard - // every use of it behind `active`. - let i = if active { n - 1 - step } else { 0 }; - let s = if active && lane > i && lane < n { - mat.read(sm_idx(i, lane)) * x.read(lane as usize) - } else { - 0.0f32 - }; - partial.write(lane as usize, s); - workgroup_memory_barrier_with_group_sync(); - for r in 0..6u32 { - let stride = 1u32 << (5 - r); - if lane < stride { - let v = partial.read(lane as usize) + partial.read((lane + stride) as usize); - partial.write(lane as usize, v); + if lane == 0 { + // Solve Lᵀ·z = b: scatter descending (descendants before ancestors). + for step in 0..n { + let i = n - 1 - step; + let xi = x.read(i as usize); + let mut j = parents.read(parents_offset + i as usize); + for _ in 0..MAX_DOFS_U32 { + if j == NO_PARENT { + break; + } + let v = x.read(j as usize) - mat.read(sm_idx(i, j)) * xi; + x.write(j as usize, v); + j = parents.read(parents_offset + j as usize); } - workgroup_memory_barrier_with_group_sync(); } - if active && lane == 0 { - let u = mat.read(sm_idx(i, i)); - let cur = x.read(i as usize) - partial.read(0); - x.write(i as usize, if u != 0.0 { cur / u } else { 0.0 }); + // z = D⁻¹·z. + for i in 0..n { + let d = mat.read(sm_idx(i, i)); + let v = x.read(i as usize); + x.write(i as usize, if d != 0.0 { v / d } else { 0.0 }); } - workgroup_memory_barrier_with_group_sync(); - } -} - -/// Apply the recorded pivots (sequential — lane 0 only). `n` is small so the -/// extra parallelism wouldn't pay for the barrier. -#[inline] -pub(super) fn lu_apply_pivots( - n: u32, - lane: u32, - buf_pivots: &[u32], - pivots_offset: usize, - x: &mut [f32; MAX_MB_DOFS], -) { - if lane == 0 { - for k in 0..n { - let p = buf_pivots.read(pivots_offset + k as usize); - if p != k { - let a = x.read(k as usize); - let b = x.read(p as usize); - x.write(k as usize, b); - x.write(p as usize, a); + // Solve L·x = z: gather ascending (ancestors before descendants). + for i in 0..n { + let mut s = x.read(i as usize); + let mut j = parents.read(parents_offset + i as usize); + for _ in 0..MAX_DOFS_U32 { + if j == NO_PARENT { + break; + } + s -= mat.read(sm_idx(i, j)) * x.read(j as usize); + j = parents.read(parents_offset + j as usize); } + x.write(i as usize, s); } } workgroup_memory_barrier_with_group_sync(); diff --git a/src_rbd_shaders/tests/linalg.rs b/src_rbd_shaders/tests/linalg.rs index b9921f01..6929a217 100644 --- a/src_rbd_shaders/tests/linalg.rs +++ b/src_rbd_shaders/tests/linalg.rs @@ -8,12 +8,12 @@ //! - `quadform_spatial` (block-diagonal spatial-mass CRBA step), //! - `gemv_tr_spatial`, //! - `skew_tr`, -//! - `lu_decompose` + `lu_solve_in_place`, including reuse of the same LU factor +//! - `ltdl_decompose` + `lu_solve_in_place`, including reuse of the same factor //! with multiple right-hand sides. use crate::utils::linalg::{ MatSlice, axpy_mat, copy_from, fill, gemm, gemm_mat3_lhs, gemm_tr, gemv_tr_spatial, - lu_decompose, lu_solve_in_place, quadform_spatial, skew, skew_tr, + ltdl_decompose, lu_solve_in_place, quadform_spatial, skew, skew_tr, }; use glamx::{Mat3, Vec3}; @@ -486,25 +486,27 @@ fn matvec(m_rows: &[&[f32]], x: &[f32]) -> Vec { } #[test] -fn lu_solve_identity_recovers_rhs() { +fn ltdl_solve_identity_recovers_rhs() { + // Identity matrix = a forest of root DOFs (no parents anywhere). let n = 4u32; let mut buf_m = vec![0.0f32; (n * n) as usize]; let m = MatSlice::dense(0, n, n); for i in 0..n { buf_m[m.idx(i, i)] = 1.0; } - let mut pivots = vec![0u32; n as usize]; - lu_decompose(&mut buf_m, m, &mut pivots, 0); + let parents = vec![u32::MAX; n as usize]; + ltdl_decompose(&mut buf_m, m, &parents, 0); let mut rhs = vec![7.0, -3.0, 2.5, 1.25]; let want = rhs.clone(); - lu_solve_in_place(&buf_m, m, &pivots, 0, &mut rhs, 0); + lu_solve_in_place(&buf_m, m, &parents, 0, &mut rhs, 0); assert_slice_eq(&rhs, &want); } #[test] -fn lu_solve_spd_matrix() { - // A symmetric positive-definite 3×3 (mimics a small mass matrix). +fn ltdl_solve_spd_matrix() { + // A symmetric positive-definite 3×3 (mimics a small mass matrix). A dense + // SPD matrix is the chain-tree special case: parents[k] = k - 1. // M = [[4, 1, 2], // [1, 5, 3], // [2, 3, 6]] @@ -514,13 +516,12 @@ fn lu_solve_spd_matrix() { let m = MatSlice::dense(0, n, n); pack_matrix(&mut buf, m, rows); - // Decompose. - let mut pivots = vec![0u32; n as usize]; - lu_decompose(&mut buf, m, &mut pivots, 0); + let parents = vec![u32::MAX, 0, 1]; + ltdl_decompose(&mut buf, m, &parents, 0); // Solve M · x = b with b = [1, 2, 3]. let mut rhs = vec![1.0f32, 2.0, 3.0]; - lu_solve_in_place(&buf, m, &pivots, 0, &mut rhs, 0); + lu_solve_in_place(&buf, m, &parents, 0, &mut rhs, 0); // Verify by multiplying back: M · x ≈ b. let mx = matvec(rows, &rhs); @@ -528,34 +529,48 @@ fn lu_solve_spd_matrix() { } #[test] -fn lu_solve_requires_pivoting() { - // This matrix has a zero top-left entry, so Doolittle without pivoting would fail. - // Partial pivoting should swap rows and produce a correct factor + solve. - // M = [[0, 2, 1], - // [1, 0, 3], - // [4, 1, 0]] - let rows: &[&[f32]] = &[&[0.0, 2.0, 1.0], &[1.0, 0.0, 3.0], &[4.0, 1.0, 0.0]]; - let n = 3u32; +fn ltdl_solve_branched_tree() { + // Branch-induced sparsity: DOFs 1 and 2 hang off root 0 on DIFFERENT + // branches (M[1,2] = M[2,1] = 0), DOF 3 continues branch 1's chain. + // 0 + // / \ + // 1 2 + // | + // 3 + // M = [[6, 1, 2, 1], + // [1, 5, 0, 2], + // [2, 0, 4, 0], + // [1, 2, 0, 3]] + let rows: &[&[f32]] = &[ + &[6.0, 1.0, 2.0, 1.0], + &[1.0, 5.0, 0.0, 2.0], + &[2.0, 0.0, 4.0, 0.0], + &[1.0, 2.0, 0.0, 3.0], + ]; + let n = 4u32; let mut buf = vec![0.0f32; (n * n) as usize]; let m = MatSlice::dense(0, n, n); pack_matrix(&mut buf, m, rows); - let mut pivots = vec![0u32; n as usize]; - lu_decompose(&mut buf, m, &mut pivots, 0); + let parents = vec![u32::MAX, 0, 0, 1]; + ltdl_decompose(&mut buf, m, &parents, 0); + + // The factorization must not have filled in the cross-branch zeros. + approx_eq(buf[m.idx(2, 1)], 0.0); // Pick a known solution, compute b = M·x, then confirm the solve recovers x. - let x_true = [1.5f32, -0.5, 2.0]; + let x_true = [1.5f32, -0.5, 2.0, 0.75]; let b = matvec(rows, &x_true); let mut rhs = b.clone(); - lu_solve_in_place(&buf, m, &pivots, 0, &mut rhs, 0); + lu_solve_in_place(&buf, m, &parents, 0, &mut rhs, 0); assert_slice_eq(&rhs, &x_true); } #[test] -fn lu_factor_reused_across_multiple_rhs() { - // The point of splitting `lu_decompose` and `lu_solve` is that we should be able - // to solve multiple right-hand sides without re-factoring. Verify that here. +fn ltdl_factor_reused_across_multiple_rhs() { + // The point of splitting `ltdl_decompose` and the solve is that we should + // be able to solve multiple right-hand sides without re-factoring. let rows: &[&[f32]] = &[ &[6.0, 2.0, 1.0, 0.0], &[2.0, 5.0, 2.0, 1.0], @@ -567,9 +582,9 @@ fn lu_factor_reused_across_multiple_rhs() { let m = MatSlice::dense(0, n, n); pack_matrix(&mut buf, m, rows); - // Decompose once. - let mut pivots = vec![0u32; n as usize]; - lu_decompose(&mut buf, m, &mut pivots, 0); + // Decompose once (dense SPD -> chain tree). + let parents = vec![u32::MAX, 0, 1, 2]; + ltdl_decompose(&mut buf, m, &parents, 0); // Solve three different RHSes with the same factorization. let rhss = [ @@ -579,7 +594,7 @@ fn lu_factor_reused_across_multiple_rhs() { ]; for b in &rhss { let mut rhs = b.to_vec(); - lu_solve_in_place(&buf, m, &pivots, 0, &mut rhs, 0); + lu_solve_in_place(&buf, m, &parents, 0, &mut rhs, 0); // Verify: M · rhs ≈ b. let mx = matvec(rows, &rhs); assert_slice_eq(&mx, b); @@ -587,23 +602,23 @@ fn lu_factor_reused_across_multiple_rhs() { } #[test] -fn lu_solve_respects_offsets() { +fn ltdl_solve_respects_offsets() { // Run two independent systems in one set of buffers using nonzero offsets — // this is how the multibody kernels share a flat buffer across multibodies. let rows_a: &[&[f32]] = &[&[3.0, 1.0], &[1.0, 2.0]]; let rows_b: &[&[f32]] = &[&[5.0, 2.0], &[2.0, 4.0]]; let n = 2u32; - // Layout: [A (4) | B (4)] for matrices, pivots similarly, RHS similarly. + // Layout: [A (4) | B (4)] for matrices, parents similarly, RHS similarly. let mut buf_m = vec![0.0f32; 8]; let m_a = MatSlice::dense(0, n, n); let m_b = MatSlice::dense(4, n, n); pack_matrix(&mut buf_m, m_a, rows_a); pack_matrix(&mut buf_m, m_b, rows_b); - let mut pivots = vec![0u32; 4]; - lu_decompose(&mut buf_m, m_a, &mut pivots, 0); - lu_decompose(&mut buf_m, m_b, &mut pivots, 2); + let parents = vec![u32::MAX, 0, u32::MAX, 0]; + ltdl_decompose(&mut buf_m, m_a, &parents, 0); + ltdl_decompose(&mut buf_m, m_b, &parents, 2); let x_a_true = [1.0f32, -1.0]; let x_b_true = [2.0f32, 0.5]; @@ -611,8 +626,8 @@ fn lu_solve_respects_offsets() { let b_b = matvec(rows_b, &x_b_true); let mut buf_rhs = vec![b_a[0], b_a[1], b_b[0], b_b[1]]; - lu_solve_in_place(&buf_m, m_a, &pivots, 0, &mut buf_rhs, 0); - lu_solve_in_place(&buf_m, m_b, &pivots, 2, &mut buf_rhs, 2); + lu_solve_in_place(&buf_m, m_a, &parents, 0, &mut buf_rhs, 0); + lu_solve_in_place(&buf_m, m_b, &parents, 2, &mut buf_rhs, 2); assert_slice_eq(&buf_rhs[0..2], &x_a_true); assert_slice_eq(&buf_rhs[2..4], &x_b_true); diff --git a/src_rbd_shaders/utils/linalg.rs b/src_rbd_shaders/utils/linalg.rs index 7918699a..8e764996 100644 --- a/src_rbd_shaders/utils/linalg.rs +++ b/src_rbd_shaders/utils/linalg.rs @@ -610,72 +610,57 @@ pub fn gemm_tr( } // -// LU decomposition + solve, split so the factorization can be reused with -// multiple right-hand sides (mirrors `nalgebra::LU` / `LU::solve_mut`). +// Tree-sparse LᵀDL decomposition + solve, split so the factorization can be +// reused with multiple right-hand sides. // -/// In-place LU factorization with partial pivoting (Doolittle form). +/// In-place tree-sparse LᵀDL factorization (`M = Lᵀ·D·L`) of an SPD matrix +/// with branch-induced sparsity, eliminating leaves-to-root (Featherstone, +/// RBDA ch. 6). /// -/// Overwrites `m` (a square `n × n` dense column-major view) with the LU factors: -/// strictly-below-diagonal entries hold `L` (with implicit unit diagonal), the -/// diagonal and above hold `U`. Row pivots are written to `pivots[0..n]` — -/// `pivots[k]` is the row that was swapped with row `k` during elimination step -/// `k`. `pivots_offset` is where this multibody's pivot slot starts in `buf_pivots`. +/// `buf_parents[parents_offset + k]` is DOF `k`'s parent index (`u32::MAX` +/// at roots); `M[i, j]` must be zero whenever `i` and `j` do not lie on the +/// same root-to-leaf chain. Overwrites the strict-lower ancestor-chain +/// entries of `m` with `L` (implicit unit diagonal) and the diagonal with +/// `D`; the upper triangle is left untouched (stale symmetric copies). +/// A dense SPD matrix is the chain-tree special case (`parents[k] = k - 1`). #[inline] -pub fn lu_decompose(buf_m: &mut [f32], m: MatSlice, buf_pivots: &mut [u32], pivots_offset: usize) { +pub fn ltdl_decompose(buf_m: &mut [f32], m: MatSlice, buf_parents: &[u32], parents_offset: usize) { let n = m.rows; - for k in 0..n { - // Partial pivot: find max |M[i, k]| for i in k..n. - let mut pivot_row = k; - let mut pivot_val = { - let v = buf_m.read(m.idx(k, k)); - if v >= 0.0 { v } else { -v } - }; - for i in (k + 1)..n { - let v = buf_m.read(m.idx(i, k)); - let av = if v >= 0.0 { v } else { -v }; - if av > pivot_val { - pivot_val = av; - pivot_row = i; + const NO_PARENT: u32 = u32::MAX; + for step in 0..n { + let k = n - 1 - step; + let d = buf_m.read(m.idx(k, k)); + let inv_d = if d != 0.0 { 1.0 / d } else { 0.0 }; + let mut i = buf_parents.read(parents_offset + k as usize); + // Bounded (parents strictly decrease) so corrupt data can't hang. + for _ in 0..n { + if i == NO_PARENT { + break; } - } - buf_pivots.write(pivots_offset + k as usize, pivot_row); - - // Row swap k ↔ pivot_row (full row since we haven't computed past col k). - if pivot_row != k { - for c in 0..n { - let a = buf_m.read(m.idx(k, c)); - let b = buf_m.read(m.idx(pivot_row, c)); - buf_m.write(m.idx(k, c), b); - buf_m.write(m.idx(pivot_row, c), a); - } - } - - // Scale sub-column below the pivot: M[i, k] /= M[k, k]. - let akk = buf_m.read(m.idx(k, k)); - let inv_akk = if akk != 0.0 { 1.0 / akk } else { 0.0 }; - for r in (k + 1)..n { - let v = buf_m.read(m.idx(r, k)) * inv_akk; - buf_m.write(m.idx(r, k), v); - } - - // Trailing sub-matrix update: M[i, j] -= M[i, k] * M[k, j]. - for j in (k + 1)..n { - let akj = buf_m.read(m.idx(k, j)); - for i2 in (k + 1)..n { - let lik = buf_m.read(m.idx(i2, k)); - let v = buf_m.read(m.idx(i2, j)) - lik * akj; - buf_m.write(m.idx(i2, j), v); + let a = buf_m.read(m.idx(k, i)) * inv_d; + let mut j = i; + for _ in 0..n { + if j == NO_PARENT { + break; + } + let v = buf_m.read(m.idx(i, j)) - a * buf_m.read(m.idx(k, j)); + buf_m.write(m.idx(i, j), v); + j = buf_parents.read(parents_offset + j as usize); } + buf_m.write(m.idx(k, i), a); + i = buf_parents.read(parents_offset + i as usize); } } } -/// Solve `M · x = rhs` in-place, using LU factors produced by -/// [`lu_decompose`] (and its pivot array). The result overwrites `rhs`. +/// Solve `M · x = rhs` in-place, using LᵀDL factors produced by +/// [`ltdl_decompose`] (or the workgroup variant in `dynamics::multibody::lu`) +/// and the same per-DOF parent array. The result overwrites `rhs`. /// -/// `m` and `pivots` must be the exact outputs of a previous `lu_decompose` call. -/// `rhs` is an `n`-element column vector; `rhs_offset` is where it starts in `buf_rhs`. +/// `m` must hold the exact factor output and `pivots` the parent array the +/// factorization used. `rhs` is an `n`-element column vector; `rhs_offset` is +/// where it starts in `buf_rhs`. #[inline] pub fn lu_solve_in_place( buf_m: &[f32], @@ -686,39 +671,51 @@ pub fn lu_solve_in_place( rhs_offset: usize, ) { let n = m.rows; + const NO_PARENT: u32 = u32::MAX; - // 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 ki = rhs_offset + k as usize; - let pi = rhs_offset + p as usize; - let a = buf_rhs.read(ki); - let b = buf_rhs.read(pi); - buf_rhs.write(ki, b); - buf_rhs.write(pi, a); + // `m` holds tree-sparse LᵀDL factors (see `dynamics::multibody::lu`): + // strict-lower ancestor-chain entries are L (unit diagonal), the diagonal + // is D, and `buf_pivots` holds each DOF's parent (NO_PARENT at roots). + // Every walk is O(depth), making the whole solve O(n·depth). + + // Solve Lᵀ·z = b: scatter descending (descendants before ancestors). + for step in 0..n { + let i = n - 1 - step; + let xi = buf_rhs.read(rhs_offset + i as usize); + let mut j = buf_pivots.read(pivots_offset + i as usize); + // Bounded (parents strictly decrease) so corrupt data can't hang. + for _ in 0..n { + if j == NO_PARENT { + break; + } + let idx = rhs_offset + j as usize; + let v = buf_rhs.read(idx) - buf_m.read(m.idx(i, j)) * xi; + buf_rhs.write(idx, v); + j = buf_pivots.read(pivots_offset + j as usize); } } - // Forward substitution: L · y = P · rhs (L is unit-lower — implicit diag = 1). + // z = D⁻¹·z. + for i in 0..n { + let d = buf_m.read(m.idx(i, i)); + let idx = rhs_offset + i as usize; + let v = buf_rhs.read(idx); + buf_rhs.write(idx, if d != 0.0 { v / d } else { 0.0 }); + } + + // Solve L·x = z: gather ascending (ancestors before descendants). for i in 0..n { let mut s = buf_rhs.read(rhs_offset + i as usize); - for j in 0..i { + let mut j = buf_pivots.read(pivots_offset + i as usize); + for _ in 0..n { + if j == NO_PARENT { + break; + } s -= buf_m.read(m.idx(i, j)) * buf_rhs.read(rhs_offset + j as usize); + j = buf_pivots.read(pivots_offset + j as usize); } buf_rhs.write(rhs_offset + 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 = buf_rhs.read(rhs_offset + ii as usize); - for j in (ii + 1)..n { - s -= buf_m.read(m.idx(ii, j)) * buf_rhs.read(rhs_offset + j as usize); - } - let u = buf_m.read(m.idx(ii, ii)); - buf_rhs.write(rhs_offset + ii as usize, if u != 0.0 { s / u } else { 0.0 }); - } } //