Skip to content

Commit 8837df1

Browse files
committed
perf(ldlt): improve factorization kernel
- Preserve the tiny-dimension update shape for D2-D5 to avoid regressing the core fixed-size path - Fuse multiplier computation with trailing updates for larger dimensions to reduce extra column walks - Rely on the LDLT factorization proof instead of a redundant final finite-storage scan Closes #146
1 parent 9a869fc commit 8837df1

2 files changed

Lines changed: 85 additions & 19 deletions

File tree

src/ldlt.rs

Lines changed: 84 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -104,32 +104,53 @@ impl<const D: usize> Ldlt<D> {
104104
return Err(LaError::Singular { pivot_col: j });
105105
}
106106

107-
// Compute L multipliers below the diagonal in column j.
108-
for i in (j + 1)..D {
109-
let l = rows[i][j] / d;
110-
if !l.is_finite() {
111-
cold_path();
112-
return Err(LaError::non_finite_cell(i, j));
107+
if D <= 5 {
108+
// Tiny matrices benchmark better when column normalization stays
109+
// separate from the trailing update.
110+
for i in (j + 1)..D {
111+
let l = rows[i][j] / d;
112+
if !l.is_finite() {
113+
cold_path();
114+
return Err(LaError::non_finite_cell(i, j));
115+
}
116+
rows[i][j] = l;
113117
}
114-
rows[i][j] = l;
115-
}
116118

117-
// Update the trailing submatrix (lower triangle): A := A - (L_col * d) * L_col^T.
118-
for i in (j + 1)..D {
119-
let l_i = rows[i][j];
120-
let l_i_d = l_i * d;
119+
for i in (j + 1)..D {
120+
let l_i = rows[i][j];
121+
let l_i_d = l_i * d;
122+
123+
for k in (j + 1)..=i {
124+
let l_k = rows[k][j];
125+
let new_val = (-l_i_d).mul_add(l_k, rows[i][k]);
126+
rows[i][k] = new_val;
127+
}
128+
}
129+
} else {
130+
// Larger fixed dimensions avoid an extra column walk by updating
131+
// each lower-triangular row prefix as soon as its multiplier is finite.
132+
for i in (j + 1)..D {
133+
let l_i = rows[i][j] / d;
134+
if !l_i.is_finite() {
135+
cold_path();
136+
return Err(LaError::non_finite_cell(i, j));
137+
}
138+
rows[i][j] = l_i;
139+
140+
let l_i_d = l_i * d;
121141

122-
for k in (j + 1)..=i {
123-
let l_k = rows[k][j];
124-
let new_val = (-l_i_d).mul_add(l_k, rows[i][k]);
125-
rows[i][k] = new_val;
142+
for k in (j + 1)..=i {
143+
let l_k = rows[k][j];
144+
let new_val = (-l_i_d).mul_add(l_k, rows[i][k]);
145+
rows[i][k] = new_val;
146+
}
126147
}
127148
}
128149
}
129150
}
130151

131-
let f = f.validate_finite()?;
132-
152+
// Every computed lower-triangular entry is checked when it becomes a
153+
// pivot or multiplier; the untouched upper triangle remains finite input.
133154
Ok(Self {
134155
factors: LdltFactors::new_unchecked(f),
135156
})
@@ -473,6 +494,29 @@ mod tests {
473494
);
474495
}
475496

497+
#[test]
498+
fn nonfinite_l_multiplier_overflow_fused_branch_6d() {
499+
// D > 5 uses the fused LDLT update path. Keep the same overflow shape
500+
// as the 2D test while forcing that branch.
501+
let mut rows = [[0.0; 6]; 6];
502+
for (i, row) in rows.iter_mut().enumerate() {
503+
row[i] = 1.0;
504+
}
505+
rows[0][0] = 1e-11;
506+
rows[0][5] = 1e300;
507+
rows[5][0] = 1e300;
508+
509+
let a = Matrix::<6>::try_from_rows(rows).unwrap();
510+
let err = a.ldlt(DEFAULT_SINGULAR_TOL).unwrap_err();
511+
assert_eq!(
512+
err,
513+
LaError::NonFinite {
514+
row: Some(5),
515+
col: 0
516+
}
517+
);
518+
}
519+
476520
#[test]
477521
fn nonfinite_trailing_submatrix_overflow() {
478522
// L multiplier is finite (1e200), but the rank-1 update
@@ -488,6 +532,28 @@ mod tests {
488532
);
489533
}
490534

535+
#[test]
536+
fn nonfinite_trailing_submatrix_overflow_fused_branch_6d() {
537+
// D > 5 uses the fused LDLT update path. The overflowing trailing
538+
// diagonal is detected when it later becomes a pivot.
539+
let mut rows = [[0.0; 6]; 6];
540+
for (i, row) in rows.iter_mut().enumerate() {
541+
row[i] = 1.0;
542+
}
543+
rows[0][5] = 1e200;
544+
rows[5][0] = 1e200;
545+
546+
let a = Matrix::<6>::try_from_rows(rows).unwrap();
547+
let err = a.ldlt(DEFAULT_SINGULAR_TOL).unwrap_err();
548+
assert_eq!(
549+
err,
550+
LaError::NonFinite {
551+
row: Some(5),
552+
col: 5
553+
}
554+
);
555+
}
556+
491557
#[test]
492558
fn nonfinite_solve_forward_substitution_overflow() {
493559
// SPD matrix with large L multiplier: L[1,0] = 1e153.

src/matrix.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -138,7 +138,7 @@ impl<const D: usize> Matrix<D> {
138138
/// Mutably borrow raw row-major storage without preserving the finite invariant.
139139
///
140140
/// This is reserved for internal factorization temporaries whose results are
141-
/// validated before becoming observable API values.
141+
/// validated or otherwise proven finite before becoming observable API values.
142142
#[inline]
143143
pub(crate) const fn rows_mut_unchecked(&mut self) -> &mut [[f64; D]; D] {
144144
&mut self.rows

0 commit comments

Comments
 (0)