From 991dd354cb916119f410147502824bfa2f7af2b9 Mon Sep 17 00:00:00 2001 From: Preetinder Kalsi <60548673+facelessdevhack@users.noreply.github.com> Date: Fri, 16 Oct 2026 00:32:55 +0530 Subject: [PATCH] solver: share Cholesky and LDLT triangular solves --- crates/solver/src/cholesky.rs | 171 ++++++++++++++++++++++++--------- crates/solver/src/ldlt.rs | 175 ++++++++++++++++------------------ 2 files changed, 212 insertions(+), 134 deletions(-) diff --git a/crates/solver/src/cholesky.rs b/crates/solver/src/cholesky.rs index 11c0703..a753e66 100644 --- a/crates/solver/src/cholesky.rs +++ b/crates/solver/src/cholesky.rs @@ -531,17 +531,144 @@ pub(crate) trait Supernode: Send + Sync { fn values(&self) -> &[f64]; /// The pivots of D (symmetric indefinite factors), `None` for Cholesky. fn pivots(&self) -> Option<&[Pivot]>; + /// Local pivot position to symbolic column within this supernode; identity for Cholesky. + #[inline] + fn order(&self) -> Option<&[usize]> { + None + } } impl Supernode for Vec { + #[inline] fn values(&self) -> &[f64] { self } + #[inline] fn pivots(&self) -> Option<&[Pivot]> { None } } +/// Shared triangular solve, with Cholesky's diagonal or LDLᵀ's unit L and block diagonal D. +pub(crate) fn solve_supernodes(n: usize, perm: &[usize], s: &Structure, blocks: &[T], b: &mut [f64]) -> Result<()> { + if b.len() != n { + return Err(Error::Internal("right-hand side length".into())); + } + let mut x: Vec = perm.iter().map(|&j| b.get(j).copied().unwrap_or(0.0)).collect(); + let get = |v: &[f64], i: usize| v.get(i).copied().unwrap_or(0.0); + let mut scratch = Vec::new(); + // L z = P b, then D⁻¹, by supernodes; the supernode's own unknowns are left in pivot + // order for the backward pass. + for (k, blk) in blocks.iter().enumerate() { + let rows = s.rows(k); + let (m, f, w) = (rows.len(), s.cols(k).start, s.cols(k).len()); + let t = Panel { m, data: blk.values() }; + let (own, tail) = x.get_mut(f..).and_then(|v| v.split_at_mut_checked(w)).ok_or_else(|| Error::Internal("solve supernode range".into()))?; + let pivots = blk.pivots(); + let order = blk.order(); + scratch.clear(); + let y = if let Some(order) = order { + scratch.extend(order.iter().map(|&c| get(own, c))); + scratch.as_mut_slice() + } else { + &mut *own + }; + for q in 0..w { + let end = q + if pivots.and_then(|p| p.get(q)) == Some(&Pivot::First) { 2 } else { 1 }; + let mut yq = get(y, q); + if pivots.is_none() { + yq /= t.data.get(offset(m, q)).copied().unwrap_or(1.0); + if let Some(v) = y.get_mut(q) { + *v = yq; + } + } + for (v, l) in y.get_mut(end..w).unwrap_or_default().iter_mut().zip(t.column(q, end, w - end)) { + *v -= l * yq; + } + for (&r, l) in rows.get(w..).unwrap_or_default().iter().zip(t.column(q, w, m - w)) { + if let Some(v) = tail.get_mut(r.wrapping_sub(f + w)) { + *v -= l * yq; + } + } + } + if pivots.is_some() { + let mut c = 0; + while c < w { + if pivots.and_then(|p| p.get(c)) == Some(&Pivot::First) { + let (d11, d21, d22) = (t.at(c, c), t.at(c + 1, c), t.at(c + 1, c + 1)); + let det = d11 * d22 - d21 * d21; + let (p, q) = (get(y, c), get(y, c + 1)); + if let Some(v) = y.get_mut(c) { + *v = (d22 * p - d21 * q) / det; + } + if let Some(v) = y.get_mut(c + 1) { + *v = (d11 * q - d21 * p) / det; + } + c += 2; + } else { + if let Some(v) = y.get_mut(c) { + *v /= t.at(c, c); + } + c += 1; + } + } + } + if order.is_some() { + own.copy_from_slice(&scratch); + } + } + // Lᵀ x = z, supernodes in reverse; each scatters its unknowns back from pivot order. + for (k, blk) in blocks.iter().enumerate().rev() { + let rows = s.rows(k); + let (m, f, w) = (rows.len(), s.cols(k).start, s.cols(k).len()); + let t = Panel { m, data: blk.values() }; + let (own, tail) = x.get_mut(f..).and_then(|v| v.split_at_mut_checked(w)).ok_or_else(|| Error::Internal("solve supernode range".into()))?; + let pivots = blk.pivots(); + let order = blk.order(); + scratch.clear(); + let y = if order.is_some() { + scratch.extend_from_slice(own); + scratch.as_mut_slice() + } else { + &mut *own + }; + for c in (0..w).rev() { + // Reverse pivot groups, keeping ascending columns within each 2 × 2 group. + let (q, end) = match pivots.and_then(|p| p.get(c)) { + Some(Pivot::Second) if c > 0 => (c - 1, c + 1), + Some(Pivot::First) => (c + 1, c + 2), + _ => (c, c + 1), + }; + let mut sum = get(y, q); + for (v, l) in y.get(end..w).unwrap_or_default().iter().zip(t.column(q, end, w - end)) { + sum -= l * v; + } + for (&r, l) in rows.get(w..).unwrap_or_default().iter().zip(t.column(q, w, m - w)) { + sum -= l * get(tail, r.wrapping_sub(f + w)); + } + if pivots.is_none() { + sum /= t.data.get(offset(m, q)).copied().unwrap_or(1.0); + } + if let Some(v) = y.get_mut(q) { + *v = sum; + } + } + if let Some(order) = order { + for (pos, &c) in order.iter().enumerate() { + if let Some(v) = own.get_mut(c) { + *v = get(&scratch, pos); + } + } + } + } + for (&j, v) in perm.iter().zip(&x) { + if let Some(t) = b.get_mut(j) { + *t = *v; + } + } + Ok(()) +} + fn internal(what: &str) -> FactorError { FactorError::Internal(format!("supernodal factorization: {what}")) } @@ -981,49 +1108,7 @@ impl Factor { /// Solve A x = b in place. pub fn solve(&self, b: &mut [f64]) -> Result<()> { - if b.len() != self.n { - return Err(Error::Internal("right-hand side length".into())); - } - let mut x: Vec = self.perm.iter().map(|&j| b.get(j).copied().unwrap_or(0.0)).collect(); - let s = &self.structure; - // L y = b, by columns. - for (k, l) in self.values.iter().enumerate() { - let rows = s.rows(k); - let m = rows.len(); - for (c, j) in s.cols(k).enumerate() { - let start = offset(m, c); - let yj = x.get(j).copied().unwrap_or(0.0) / l.get(start).copied().unwrap_or(1.0); - if let Some(v) = x.get_mut(j) { - *v = yj; - } - for (r, v) in rows.get(c + 1..).unwrap_or_default().iter().zip(l.get(start + 1..).unwrap_or_default()) { - if let Some(t) = x.get_mut(*r) { - *t -= v * yj; - } - } - } - } - // Lᵀ x = y. - for (k, l) in self.values.iter().enumerate().rev() { - let rows = s.rows(k); - let m = rows.len(); - for (c, j) in s.cols(k).enumerate().rev() { - let start = offset(m, c); - let mut sum = x.get(j).copied().unwrap_or(0.0); - for (r, v) in rows.get(c + 1..).unwrap_or_default().iter().zip(l.get(start + 1..).unwrap_or_default()) { - sum -= v * x.get(*r).copied().unwrap_or(0.0); - } - if let Some(t) = x.get_mut(j) { - *t = sum / l.get(start).copied().unwrap_or(1.0); - } - } - } - for (&j, v) in self.perm.iter().zip(&x) { - if let Some(t) = b.get_mut(j) { - *t = *v; - } - } - Ok(()) + solve_supernodes(self.n, &self.perm, &self.structure, &self.values, b) } } diff --git a/crates/solver/src/ldlt.rs b/crates/solver/src/ldlt.rs index 1c99f4b..0d2f72c 100644 --- a/crates/solver/src/ldlt.rs +++ b/crates/solver/src/ldlt.rs @@ -60,12 +60,12 @@ use serde::Serialize; +use crate::Result; use crate::cholesky::{ - Context, Factor, FactorError, PANEL, PIVOT_TOLERANCE, Panel, Pivot, Structure, Supernode, Symbolic, UpperColumns, offset, products, subtract, - times_d, + Context, Factor, FactorError, PANEL, PIVOT_TOLERANCE, Panel, Pivot, Structure, Supernode, Symbolic, UpperColumns, offset, products, + solve_supernodes, subtract, times_d, }; use crate::sparse::Csr; -use crate::{Error, Result}; /// (1 + √17)/8: the growth bound of Bunch and Kaufman. const ALPHA: f64 = 0.640_388_203_202_207_6; @@ -94,12 +94,18 @@ struct Block { } impl Supernode for Block { + #[inline] fn values(&self) -> &[f64] { &self.values } + #[inline] fn pivots(&self) -> Option<&[Pivot]> { Some(&self.pivots) } + #[inline] + fn order(&self) -> Option<&[usize]> { + Some(&self.order) + } } /// The pivot a search chose, at supernode-local positions not yet moved to k. @@ -410,94 +416,7 @@ impl LdltFactor { /// Solve A x = b in place. pub fn solve(&self, b: &mut [f64]) -> Result<()> { - if b.len() != self.n { - return Err(Error::Internal("right-hand side length".into())); - } - let mut x: Vec = self.perm.iter().map(|&j| b.get(j).copied().unwrap_or(0.0)).collect(); - let s = &self.structure; - let get = |v: &[f64], i: usize| v.get(i).copied().unwrap_or(0.0); - let mut y = Vec::new(); - // L z = P b, then D⁻¹, by supernodes; the supernode's own unknowns are left in pivot - // order for the backward pass. - for (k, blk) in self.blocks.iter().enumerate() { - let rows = s.rows(k); - let (m, f, w) = (rows.len(), s.cols(k).start, s.cols(k).len()); - let t = Panel { m, data: &blk.values }; - y.clear(); - y.extend(blk.order.iter().map(|c| get(&x, f + c))); - let mut c = 0; - while c < w { - let width = if blk.pivots.get(c) == Some(&Pivot::First) { 2 } else { 1 }; - for q in c..c + width { - let yq = get(&y, q); - for (r, l) in (c + width..m).zip(t.column(q, c + width, m - c - width)) { - let target = if r < w { y.get_mut(r) } else { rows.get(r).and_then(|g| x.get_mut(*g)) }; - if let Some(v) = target { - *v -= l * yq; - } - } - } - c += width; - } - let mut c = 0; - while c < w { - if blk.pivots.get(c) == Some(&Pivot::First) { - let (d11, d21, d22) = (t.at(c, c), t.at(c + 1, c), t.at(c + 1, c + 1)); - let det = d11 * d22 - d21 * d21; - let (p, q) = (get(&y, c), get(&y, c + 1)); - if let Some(v) = y.get_mut(c) { - *v = (d22 * p - d21 * q) / det; - } - if let Some(v) = y.get_mut(c + 1) { - *v = (d11 * q - d21 * p) / det; - } - c += 2; - } else { - if let Some(v) = y.get_mut(c) { - *v /= t.at(c, c); - } - c += 1; - } - } - if let Some(dst) = x.get_mut(f..f + w) { - dst.copy_from_slice(&y); - } - } - // Lᵀ x = z, supernodes in reverse; each scatters its unknowns back from pivot order. - for (k, blk) in self.blocks.iter().enumerate().rev() { - let rows = s.rows(k); - let (m, f, w) = (rows.len(), s.cols(k).start, s.cols(k).len()); - let t = Panel { m, data: &blk.values }; - y.clear(); - y.extend(x.get(f..f + w).unwrap_or_default()); - let mut c = w; - while c > 0 { - let width = if c >= 2 && blk.pivots.get(c - 1) == Some(&Pivot::Second) { 2 } else { 1 }; - let first = c - width; - for q in first..c { - let mut sum = get(&y, q); - for (r, l) in (c..m).zip(t.column(q, c, m - c)) { - let v = if r < w { get(&y, r) } else { rows.get(r).map_or(0.0, |g| get(&x, *g)) }; - sum -= l * v; - } - if let Some(v) = y.get_mut(q) { - *v = sum; - } - } - c = first; - } - for (pos, c) in blk.order.iter().enumerate() { - if let Some(v) = x.get_mut(f + c) { - *v = get(&y, pos); - } - } - } - for (&j, v) in self.perm.iter().zip(&x) { - if let Some(t) = b.get_mut(j) { - *t = *v; - } - } - Ok(()) + solve_supernodes(self.n, &self.perm, &self.structure, &self.blocks, b) } } @@ -606,6 +525,7 @@ impl Symmetric { #[cfg(test)] mod tests { use super::*; + use crate::Error; use fieldcraft_parallel as parallel; fn csr(rows: &[Vec<(usize, f64)>]) -> Csr { @@ -790,6 +710,79 @@ mod tests { assert!(err < 1e-9, "{err}"); } + #[test] + fn solve_across_pivoted_supernodes_and_rhs_errors() { + // Two dense blocks sharing a separator; interleaving their columns forces a global + // postorder, while the first block needs a local interchange and a 2 × 2 pivot. + let mut rows = vec![Vec::new(); 7]; + for (i, row) in [[0.0, 0.1, 3.0], [0.1, 5.0, 0.2], [3.0, 0.2, 0.5]].iter().enumerate() { + rows[i].extend(row.iter().enumerate().map(|(j, v)| (j, *v))); + } + for i in 3..6 { + rows[i].extend((3..6).map(|j| (j, if i == j { 4.0 } else { 0.25 }))); + } + for i in 0..6 { + rows[i].push((6, 0.15)); + rows[6].push((i, 0.15)); + } + rows[6].push((6, 8.0)); + let order = [3, 0, 4, 1, 5, 2, 6]; + let inverse: Vec<_> = (0..7).map(|i| order.iter().position(|&j| j == i).unwrap()).collect(); + let mixed = csr(&order.iter().map(|&i| rows[i].iter().map(|&(j, v)| (inverse[j], v)).collect()).collect::>()); + let f = LdltFactor::new(&UpperColumns::from_symmetric(&mixed), usize::MAX).unwrap(); + assert!(f.perm.iter().enumerate().any(|(i, &j)| i != j)); + assert!(f.blocks.len() > 1); + assert!(f.blocks.iter().enumerate().any(|(k, b)| { + b.pivots.contains(&Pivot::One) + && b.pivots.contains(&Pivot::First) + && b.order.iter().enumerate().any(|(i, &j)| i != j) + && f.structure.rows(k).len() > b.order.len() + })); + let cases = [ + csr(&[]), + csr(&[vec![(0, 4.0)]]), + csr(&[vec![(0, -4.0)]]), + mixed, + shifted_grid([4, 3, 3], -0.5).0, + shifted_grid([4, 3, 3], 2.7).0, + shifted_grid([4, 3, 3], 7.1).0, + ]; + for (case, a) in cases.iter().enumerate() { + let upper = UpperColumns::from_symmetric(a); + let ldlt = LdltFactor::new(&upper, usize::MAX).unwrap(); + let chol = Factor::new(&upper, usize::MAX).ok(); + for seed in [0.0, 0.5, 3.25] { + let expected: Vec<_> = (0..a.n).map(|i| ((i as f64 + seed) * 0.37).sin()).collect(); + let mut rhs = vec![0.0; a.n]; + a.mul(&expected, &mut rhs); + let mut actual = rhs.clone(); + ldlt.solve(&mut actual).unwrap(); + assert!(actual.iter().zip(&expected).all(|(x, y)| (x - y).abs() < 1e-10), "case {case}"); + let mut repeated = rhs.clone(); + ldlt.solve(&mut repeated).unwrap(); + assert!(actual.iter().zip(&repeated).all(|(x, y)| x.to_bits() == y.to_bits())); + if let Some(chol) = &chol { + let mut actual = rhs; + chol.solve(&mut actual).unwrap(); + assert!(actual.iter().zip(&expected).all(|(x, y)| (x - y).abs() < 1e-10), "case {case}"); + } + } + for len in [a.n + 1, a.n.saturating_sub(1)] { + if len == a.n { + continue; + } + let before = vec![0.25; len]; + let mut rhs = before.clone(); + assert!(matches!(ldlt.solve(&mut rhs), Err(Error::Internal(_)))); + assert_eq!(rhs, before); + if let Some(chol) = &chol { + assert!(matches!(chol.solve(&mut rhs), Err(Error::Internal(_)))); + assert_eq!(rhs, before); + } + } + } + } + #[test] fn auto_selection_and_refinement() { let (spd, _) = shifted_grid([4, 4, 4], -1.0);