1use dyn_stack::{MemBuffer, MemStack};
2use faer::diag::{Diag, DiagRef};
3use faer::linalg::solvers::{self, Solve};
4pub use faer::linalg::solvers::{
5 Lblt as FaerLblt, Ldlt as FaerLdlt, Llt as FaerLlt, Solve as FaerSolve,
6};
7use faer::linalg::svd::{self, ComputeSvdVectors};
8use faer::prelude::ReborrowMut;
9use faer::{Conj, Mat, MatMut, MatRef, Par, Side, Unbind, get_global_parallelism};
10use ndarray::{Array1, Array2, ArrayBase, ArrayViewMut1, Data, Ix1, Ix2};
11use std::marker::PhantomData;
12use std::panic::{AssertUnwindSafe, catch_unwind};
13use thiserror::Error;
14
15const RRQR_RANK_ALPHA: f64 = 100.0;
16
17thread_local! {
18 static NESTED_PARALLEL_DEPTH: std::cell::Cell<usize> = const { std::cell::Cell::new(0) };
19}
20
21struct NestedParallelGuard;
22
23impl NestedParallelGuard {
24 #[inline]
25 fn enter() -> Self {
26 NESTED_PARALLEL_DEPTH.with(|depth| depth.set(depth.get().saturating_add(1)));
27 Self
28 }
29}
30
31impl Drop for NestedParallelGuard {
32 #[inline]
33 fn drop(&mut self) {
34 NESTED_PARALLEL_DEPTH.with(|depth| depth.set(depth.get().saturating_sub(1)));
35 }
36}
37
38#[inline]
48pub fn with_nested_parallel<T>(body: impl FnOnce() -> T) -> T {
49 let guard = NestedParallelGuard::enter();
50 let out = body();
51 drop(guard);
52 out
53}
54
55#[inline]
58pub fn in_nested_parallel_region() -> bool {
59 NESTED_PARALLEL_DEPTH.with(|depth| depth.get() > 0)
60}
61
62#[inline]
70pub fn effective_global_parallelism() -> Par {
71 if in_nested_parallel_region() {
72 Par::Seq
73 } else {
74 get_global_parallelism()
75 }
76}
77
78static FAER_SEQ_STATE: std::sync::Mutex<FaerSeqState> = std::sync::Mutex::new(FaerSeqState {
100 depth: 0,
101 saved: None,
102});
103
104struct FaerSeqState {
105 depth: usize,
106 saved: Option<Par>,
107}
108
109#[must_use = "the sequential scope only holds while the guard is alive"]
118pub struct FaerSequentialScope {
119 _private: (),
120}
121
122impl FaerSequentialScope {
123 pub fn enter() -> Self {
125 let mut state = FAER_SEQ_STATE
126 .lock()
127 .unwrap_or_else(std::sync::PoisonError::into_inner);
128 if state.depth == 0 {
129 state.saved = Some(get_global_parallelism());
130 faer::set_global_parallelism(Par::Seq);
131 }
132 state.depth += 1;
133 Self { _private: () }
134 }
135}
136
137impl Drop for FaerSequentialScope {
138 fn drop(&mut self) {
139 let mut state = FAER_SEQ_STATE
140 .lock()
141 .unwrap_or_else(std::sync::PoisonError::into_inner);
142 state.depth -= 1;
143 if state.depth == 0 {
144 if let Some(par) = state.saved.take() {
145 faer::set_global_parallelism(par);
146 }
147 }
148 }
149}
150
151#[inline]
156pub fn with_faer_sequential<T>(body: impl FnOnce() -> T) -> T {
157 let faer_seq_guard = FaerSequentialScope::enter();
158 let out = body();
159 drop(faer_seq_guard);
160 out
161}
162
163#[derive(Debug, Error)]
164pub enum FaerLinalgError {
165 #[error("Factorization failed in {context}")]
166 FactorizationFailed { context: &'static str },
167 #[error("SVD failed to converge in {context}")]
168 SvdNoConvergence { context: &'static str },
169 #[error("Self-adjoint eigendecomposition input contains non-finite values in {context}")]
170 SelfAdjointEigenNonFiniteInput { context: &'static str },
171 #[error("Strict self-adjoint eigendecomposition rejected its input: {reason}")]
172 StrictSelfAdjointEigenInvalidInput { reason: String },
173 #[error("Self-adjoint eigendecomposition failed: {0:?}")]
174 SelfAdjointEigen(solvers::EvdError),
175 #[error("Cholesky factorization failed: {0:?}")]
176 Cholesky(solvers::LltError),
177 #[error("LDLT factorization failed: {0:?}")]
178 Ldlt(solvers::LdltError),
179}
180
181pub enum FaerSymmetricFactor {
182 Llt(FaerLlt<f64>),
183 Ldlt(FaerLdlt<f64>),
184 Lblt(FaerLblt<f64>),
185}
186
187#[inline]
188pub fn cholesky_factor_logdet(factor: MatRef<'_, f64>) -> f64 {
189 2.0 * diagonal_log_sum(factor.diagonal())
190}
191
192#[inline]
193fn diagonal_log_sum(diagonal: DiagRef<'_, f64>) -> f64 {
194 diagonal
195 .column_vector()
196 .iter()
197 .map(|&x| x.ln())
198 .sum::<f64>()
199}
200
201impl FaerSymmetricFactor {
202 #[inline]
204 pub fn n(&self) -> usize {
205 use faer::linalg::solvers::ShapeCore;
206 match self {
207 FaerSymmetricFactor::Llt(f) => f.nrows(),
208 FaerSymmetricFactor::Ldlt(f) => f.nrows(),
209 FaerSymmetricFactor::Lblt(f) => f.nrows(),
210 }
211 }
212
213 #[inline]
214 pub fn solve(&self, rhs: MatRef<'_, f64>) -> Mat<f64> {
215 match self {
216 FaerSymmetricFactor::Llt(f) => f.solve(rhs),
217 FaerSymmetricFactor::Ldlt(f) => f.solve(rhs),
218 FaerSymmetricFactor::Lblt(f) => f.solve(rhs),
219 }
220 }
221
222 #[inline]
223 pub fn solve_in_place(&self, rhs: MatMut<'_, f64>) {
224 match self {
225 FaerSymmetricFactor::Llt(f) => f.solve_in_place(rhs),
226 FaerSymmetricFactor::Ldlt(f) => f.solve_in_place(rhs),
227 FaerSymmetricFactor::Lblt(f) => f.solve_in_place(rhs),
228 }
229 }
230}
231
232impl crate::matrix::FactorizedSystem for FaerSymmetricFactor {
233 fn solve(&self, rhs: &Array1<f64>) -> Result<Array1<f64>, String> {
234 let mut out = rhs.clone();
235 let mut out_mat = array1_to_col_matmut(&mut out);
236 self.solve_in_place(out_mat.as_mut());
237 if !out.iter().all(|v| v.is_finite()) {
238 return Err("symmetric factor solve produced non-finite values".to_string());
239 }
240 Ok(out)
241 }
242
243 fn solvemulti(&self, rhs: &Array2<f64>) -> Result<Array2<f64>, String> {
244 let mut out = Array2::<f64>::zeros(rhs.raw_dim());
245 for j in 0..rhs.ncols() {
246 for i in 0..rhs.nrows() {
247 out[[i, j]] = rhs[[i, j]];
248 }
249 }
250 let mut out_mat = array2_to_matmut(&mut out);
251 self.solve_in_place(out_mat.as_mut());
252 if !out.iter().all(|v| v.is_finite()) {
253 return Err("symmetric factor multi-solve produced non-finite values".to_string());
254 }
255 Ok(out)
256 }
257
258 fn logdet(&self) -> f64 {
259 match self {
260 FaerSymmetricFactor::Llt(f) => cholesky_factor_logdet(f.L()),
261 FaerSymmetricFactor::Ldlt(f) => diagonal_log_sum(f.D()),
262 FaerSymmetricFactor::Lblt(..) => {
263 f64::NAN
267 }
268 }
269 }
270}
271
272#[inline]
274pub fn factorize_symmetricwith_fallback(
275 matrix: MatRef<'_, f64>,
276 side: Side,
277) -> Result<FaerSymmetricFactor, FaerLinalgError> {
278 if let Ok(llt) = FaerLlt::new(matrix, side) {
279 return Ok(FaerSymmetricFactor::Llt(llt));
280 }
281 let ldlt_err = match FaerLdlt::new(matrix, side) {
282 Ok(ldlt) => return Ok(FaerSymmetricFactor::Ldlt(ldlt)),
283 Err(err) => err,
284 };
285 let lblt = catch_unwind(AssertUnwindSafe(|| FaerLblt::new(matrix, side)))
286 .map_err(|_| FaerLinalgError::Ldlt(ldlt_err))?;
287 Ok(FaerSymmetricFactor::Lblt(lblt))
288}
289
290#[inline]
291const fn should_use_faer_matmul(m: usize, n: usize, k: usize) -> bool {
292 const MIN_DIM: usize = 32;
296 const MIN_FLOP_SCALE: usize = 64 * 64;
297 (m >= MIN_DIM || n >= MIN_DIM || k >= MIN_DIM)
298 && m.saturating_mul(n).saturating_mul(k) >= MIN_FLOP_SCALE
299}
300
301#[inline]
302pub fn matmul_parallelism(m: usize, n: usize, k: usize) -> Par {
303 const PAR_MIN_FLOP_SCALE: usize = 2_000_000;
307 const PAR_MIN_LONG_DIM: usize = 256;
308 let flop_scale = m.saturating_mul(n).saturating_mul(k);
309 let long_dim = m.max(n).max(k);
310 if flop_scale >= PAR_MIN_FLOP_SCALE && long_dim >= PAR_MIN_LONG_DIM {
311 effective_global_parallelism()
315 } else {
316 Par::Seq
317 }
318}
319
320#[inline]
321pub fn array2_to_matmut(array: &mut Array2<f64>) -> MatMut<'_, f64> {
322 let (rows, cols) = array.dim();
323 let strides = array.strides();
324
325 let s0 = strides[0];
332 let s1 = strides[1];
333
334 unsafe { MatMut::from_raw_parts_mut(array.as_mut_ptr(), rows, cols, s0, s1) }
338}
339
340pub fn array2_to_nested_vec(array: &Array2<f64>) -> Vec<Vec<f64>> {
343 array.rows().into_iter().map(|row| row.to_vec()).collect()
344}
345
346#[inline]
347pub fn array1_to_col_matmut(array: &mut Array1<f64>) -> MatMut<'_, f64> {
348 let len = array.len();
349 let stride = array.strides()[0];
350 unsafe {
354 MatMut::from_raw_parts_mut(
355 array.as_mut_ptr(),
356 len,
357 1,
358 stride,
359 0, )
361 }
362}
363
364#[inline]
371pub fn fast_ata<S: Data<Elem = f64>>(a: &ArrayBase<S, Ix2>) -> Array2<f64> {
372 let p = a.ncols();
373 let mut out = Array2::<f64>::zeros((p, p));
374 fast_ata_into(a, &mut out);
375 out
376}
377
378#[inline]
381pub fn fast_ata_into<S: Data<Elem = f64>>(a: &ArrayBase<S, Ix2>, out: &mut Array2<f64>) {
382 use faer::Accum;
383 use faer::linalg::matmul::triangular::{BlockStructure, matmul as tri_matmul};
384
385 let (n, p) = a.dim();
386 assert_eq!(out.nrows(), p, "output rows must match p");
387 assert_eq!(out.ncols(), p, "output cols must match p");
388
389 if !should_use_faer_matmul(p, p, n) {
390 out.assign(&a.t().dot(a));
391 return;
392 }
393
394 let mut outview = array2_to_matmut(out);
395
396 let aview = FaerArrayView::new(a);
397 let a_ref = aview.as_ref();
398 let a_t = a_ref.transpose();
399 let par = matmul_parallelism(p, p, n);
400 tri_matmul(
401 outview.as_mut(),
402 BlockStructure::TriangularLower,
403 Accum::Replace,
404 a_t,
405 BlockStructure::Rectangular,
406 a_ref,
407 BlockStructure::Rectangular,
408 1.0,
409 par,
410 );
411 for i in 0..p {
413 for j in (i + 1)..p {
414 out[[i, j]] = out[[j, i]];
415 }
416 }
417}
418
419#[inline]
423pub fn fast_atb<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
424 a: &ArrayBase<S1, Ix2>,
425 b: &ArrayBase<S2, Ix2>,
426) -> Array2<f64> {
427 if let Some(out) =
428 crate::gpu_hook::gpu_dispatch().and_then(|d| d.try_fast_atb(a.view(), b.view()))
429 {
430 return out;
431 }
432 let (n_a, p) = a.dim();
433 let q = b.ncols();
434 fast_atb_with_parallelism(a, b, matmul_parallelism(p, q, n_a))
435}
436
437#[inline]
440pub fn fast_atb_with_parallelism<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
441 a: &ArrayBase<S1, Ix2>,
442 b: &ArrayBase<S2, Ix2>,
443 par: Par,
444) -> Array2<f64> {
445 use faer::linalg::matmul::matmul;
446 use faer::{Accum, Mat};
447
448 let (n_a, p) = a.dim();
449 let (n_b, q) = b.dim();
450 assert_eq!(n_a, n_b, "A and B must have same number of rows");
451
452 if !should_use_faer_matmul(p, q, n_a) {
454 return a.t().dot(b);
455 }
456
457 let mut result = Mat::<f64>::zeros(p, q);
458
459 let aview = FaerArrayView::new(a);
460 let bview = FaerArrayView::new(b);
461 let a_ref = aview.as_ref();
462 let b_ref = bview.as_ref();
463
464 matmul(
466 result.as_mut(),
467 Accum::Replace,
468 a_ref.transpose(),
469 b_ref,
470 1.0,
471 par,
472 );
473
474 mat_to_array(result.as_ref())
475}
476
477#[inline]
480pub fn fast_abt<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
481 a: &ArrayBase<S1, Ix2>,
482 b: &ArrayBase<S2, Ix2>,
483) -> Array2<f64> {
484 use faer::linalg::matmul::matmul;
485 use faer::{Accum, Mat};
486
487 let (m, k_a) = a.dim();
488 let (n, k_b) = b.dim();
489 assert_eq!(
490 k_a, k_b,
491 "A and B must have same number of columns for A·Bᵀ"
492 );
493
494 if !should_use_faer_matmul(m, n, k_a) {
495 return a.dot(&b.t());
496 }
497
498 let mut result = Mat::<f64>::zeros(m, n);
499 let aview = FaerArrayView::new(a);
500 let bview = FaerArrayView::new(b);
501 let par = matmul_parallelism(m, n, k_a);
502 matmul(
503 result.as_mut(),
504 Accum::Replace,
505 aview.as_ref(),
506 bview.as_ref().transpose(),
507 1.0,
508 par,
509 );
510 mat_to_array(result.as_ref())
511}
512
513#[inline]
517pub fn fast_ab<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
518 a: &ArrayBase<S1, Ix2>,
519 b: &ArrayBase<S2, Ix2>,
520) -> Array2<f64> {
521 if let Some(out) =
522 crate::gpu_hook::gpu_dispatch().and_then(|d| d.try_fast_ab(a.view(), b.view()))
523 {
524 return out;
525 }
526 let n = a.nrows();
527 let q = b.ncols();
528 let mut out = Array2::<f64>::zeros((n, q));
529 fast_ab_into(a, b, &mut out);
530 out
531}
532
533const FMA_LANES: usize = 8;
558
559const KERNEL_PAR_MIN_FLOP: usize = 1 << 18; const AV_PAR_CHUNK_ROWS: usize = 1024;
567
568const ATV_BLOCK_ROWS: usize = 512;
573
574#[inline]
575fn kernel_should_parallelize(n: usize, p: usize) -> bool {
576 !in_nested_parallel_region()
577 && n.saturating_mul(p) >= KERNEL_PAR_MIN_FLOP
578 && rayon::current_num_threads() > 1
579}
580
581#[inline(always)]
596fn fma_dot(a: &[f64], b: &[f64]) -> f64 {
597 assert_eq!(a.len(), b.len(), "fma_dot: operand length mismatch");
598 let mut sum = [0.0f64; FMA_LANES];
599 let mut comp = [0.0f64; FMA_LANES];
600 let mut ca = a.chunks_exact(FMA_LANES);
601 let mut cb = b.chunks_exact(FMA_LANES);
602 for (xa, xb) in ca.by_ref().zip(cb.by_ref()) {
603 for l in 0..FMA_LANES {
604 let x = xa[l];
605 let y = xb[l];
606 let p = x * y;
608 let ep = x.mul_add(y, -p);
609 let s = sum[l] + p;
611 let bb = s - sum[l];
612 let es = (sum[l] - (s - bb)) + (p - bb);
613 sum[l] = s;
614 comp[l] += ep + es;
615 }
616 }
617 let mut sr = 0.0f64;
619 let mut cr = 0.0f64;
620 for (&x, &y) in ca.remainder().iter().zip(cb.remainder().iter()) {
621 let p = x * y;
622 let ep = x.mul_add(y, -p);
623 let s = sr + p;
624 let bb = s - sr;
625 let es = (sr - (s - bb)) + (p - bb);
626 sr = s;
627 cr += ep + es;
628 }
629 let mut total = sr + cr;
631 for l in 0..FMA_LANES {
632 total += sum[l] + comp[l];
633 }
634 total
635}
636
637fn fast_av_rowmajor_into(x_all: &[f64], v: &[f64], n: usize, p: usize, out: &mut [f64]) {
641 assert_eq!(x_all.len(), n * p, "fast_av_rowmajor_into: x_all length");
642 assert_eq!(v.len(), p, "fast_av_rowmajor_into: v length");
643 assert_eq!(out.len(), n, "fast_av_rowmajor_into: out length");
644 if kernel_should_parallelize(n, p) {
645 use rayon::prelude::*;
646 out.par_chunks_mut(AV_PAR_CHUNK_ROWS)
647 .enumerate()
648 .for_each(|(c, chunk)| {
649 let base = c * AV_PAR_CHUNK_ROWS;
650 for (k, o) in chunk.iter_mut().enumerate() {
651 let i = base + k;
652 *o = fma_dot(&x_all[i * p..i * p + p], v);
653 }
654 });
655 } else {
656 for (i, o) in out.iter_mut().enumerate() {
657 *o = fma_dot(&x_all[i * p..i * p + p], v);
658 }
659 }
660}
661
662fn pairwise_sum_into(parts: &[Vec<f64>], out: &mut [f64]) {
664 match parts.len() {
665 0 => out.fill(0.0),
666 1 => out.copy_from_slice(&parts[0]),
667 _ => {
668 let mid = parts.len() / 2;
669 let p = out.len();
670 let mut left = vec![0.0f64; p];
671 let mut right = vec![0.0f64; p];
672 pairwise_sum_into(&parts[..mid], &mut left);
673 pairwise_sum_into(&parts[mid..], &mut right);
674 for ((o, &l), &r) in out.iter_mut().zip(left.iter()).zip(right.iter()) {
675 *o = l + r;
676 }
677 }
678 }
679}
680
681fn fast_atv_rowmajor_into(x_all: &[f64], v: &[f64], n: usize, p: usize, out: &mut [f64]) {
689 assert_eq!(x_all.len(), n * p, "fast_atv_rowmajor_into: x_all length");
690 assert_eq!(v.len(), n, "fast_atv_rowmajor_into: v length");
691 assert_eq!(out.len(), p, "fast_atv_rowmajor_into: out length");
692 let nblocks = n.div_ceil(ATV_BLOCK_ROWS);
693
694 let block_partial = |b: usize| -> Vec<f64> {
695 let start = b * ATV_BLOCK_ROWS;
696 let end = (start + ATV_BLOCK_ROWS).min(n);
697 let mut acc = vec![0.0f64; p];
698 for i in start..end {
699 let vi = v[i];
700 let row = &x_all[i * p..i * p + p];
701 for (a, &xij) in acc.iter_mut().zip(row.iter()) {
702 *a = xij.mul_add(vi, *a);
703 }
704 }
705 acc
706 };
707
708 let partials: Vec<Vec<f64>> = if kernel_should_parallelize(n, p) {
709 use rayon::prelude::*;
710 (0..nblocks).into_par_iter().map(block_partial).collect()
711 } else {
712 (0..nblocks).map(block_partial).collect()
713 };
714
715 pairwise_sum_into(&partials, out);
716}
717
718#[inline]
721pub fn fast_av<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
722 a: &ArrayBase<S1, Ix2>,
723 v: &ArrayBase<S2, Ix1>,
724) -> Array1<f64> {
725 if let Some(out) =
726 crate::gpu_hook::gpu_dispatch().and_then(|d| d.try_fast_av(a.view(), v.view()))
727 {
728 return out;
729 }
730 fast_av_impl(a, v)
731}
732
733#[inline]
734fn fast_av_impl<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
735 a: &ArrayBase<S1, Ix2>,
736 v: &ArrayBase<S2, Ix1>,
737) -> Array1<f64> {
738 use faer::linalg::matmul::matmul;
739 use faer::{Accum, Mat};
740
741 let (n, p) = a.dim();
742 assert_eq!(p, v.len(), "A cols must match v length");
743
744 if let (Some(x_all), Some(vs)) = (a.as_slice(), v.as_slice())
748 && n != 0
749 && p != 0
750 {
751 let mut out = Array1::<f64>::zeros(n);
752 fast_av_rowmajor_into(
753 x_all,
754 vs,
755 n,
756 p,
757 out.as_slice_mut().expect("fresh Array1 is contiguous"),
758 );
759 return out;
760 }
761
762 if !should_use_faer_matmul(n, 1, p) {
763 return a.dot(v);
764 }
765
766 let mut result = Mat::<f64>::zeros(n, 1);
767
768 let aview = FaerArrayView::new(a);
769 let vview = FaerColView::new(v);
770 let a_ref = aview.as_ref();
771 let v_ref = vview.as_ref();
772
773 let par = matmul_parallelism(n, 1, p);
774 matmul(result.as_mut(), Accum::Replace, a_ref, v_ref, 1.0, par);
775
776 let mut out = Array1::<f64>::zeros(n);
777 for i in 0..n {
778 out[i] = result[(i, 0)];
779 }
780 out
781}
782
783#[inline]
786pub fn fast_av_into<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
787 a: &ArrayBase<S1, Ix2>,
788 v: &ArrayBase<S2, Ix1>,
789 out: &mut Array1<f64>,
790) {
791 fast_av_into_impl(a, v, out);
792}
793
794#[inline]
795fn fast_av_into_impl<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
796 a: &ArrayBase<S1, Ix2>,
797 v: &ArrayBase<S2, Ix1>,
798 out: &mut Array1<f64>,
799) {
800 use faer::Accum;
801 use faer::linalg::matmul::matmul;
802
803 let (n, p) = a.dim();
804 assert_eq!(v.len(), p, "vector length must match A cols");
805 assert_eq!(out.len(), n, "output length must match A rows");
806
807 if let (Some(x_all), Some(vs)) = (a.as_slice(), v.as_slice())
808 && n != 0
809 && p != 0
810 && let Some(out_s) = out.as_slice_mut()
811 {
812 fast_av_rowmajor_into(x_all, vs, n, p, out_s);
813 return;
814 }
815
816 if !should_use_faer_matmul(n, 1, p) {
817 out.assign(&a.dot(v));
818 return;
819 }
820
821 let mut outview = array1_to_col_matmut(out);
822
823 let aview = FaerArrayView::new(a);
824 let vview = FaerColView::new(v);
825 let a_ref = aview.as_ref();
826 let v_ref = vview.as_ref();
827 let par = matmul_parallelism(n, 1, p);
828 matmul(outview.as_mut(), Accum::Replace, a_ref, v_ref, 1.0, par);
829}
830
831#[inline]
838pub fn fast_av_view_into<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
839 a: &ArrayBase<S1, Ix2>,
840 v: &ArrayBase<S2, Ix1>,
841 out: ArrayViewMut1<'_, f64>,
842) {
843 fast_av_view_into_impl(a, v, out);
844}
845
846#[inline]
847fn fast_av_view_into_impl<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
848 a: &ArrayBase<S1, Ix2>,
849 v: &ArrayBase<S2, Ix1>,
850 mut out: ArrayViewMut1<'_, f64>,
851) {
852 use faer::Accum;
853 use faer::linalg::matmul::matmul;
854
855 let (n, p) = a.dim();
856 assert_eq!(v.len(), p, "vector length must match A cols");
857 assert_eq!(out.len(), n, "output length must match A rows");
858
859 if let (Some(x_all), Some(vs)) = (a.as_slice(), v.as_slice())
860 && n != 0
861 && p != 0
862 && let Some(out_s) = out.as_slice_mut()
863 {
864 fast_av_rowmajor_into(x_all, vs, n, p, out_s);
865 return;
866 }
867
868 if !should_use_faer_matmul(n, 1, p) {
869 let prod = a.dot(v);
870 out.assign(&prod);
871 return;
872 }
873
874 let len = out.len();
875 let stride = out.strides()[0];
876 let outview = unsafe {
880 MatMut::from_raw_parts_mut(
881 out.as_mut_ptr(),
882 len,
883 1,
884 stride,
885 0, )
887 };
888
889 let aview = FaerArrayView::new(a);
890 let vview = FaerColView::new(v);
891 let a_ref = aview.as_ref();
892 let v_ref = vview.as_ref();
893 let par = matmul_parallelism(n, 1, p);
894 matmul(outview, Accum::Replace, a_ref, v_ref, 1.0, par);
895}
896
897#[inline]
900pub fn fast_atv<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
901 a: &ArrayBase<S1, Ix2>,
902 v: &ArrayBase<S2, Ix1>,
903) -> Array1<f64> {
904 if let Some(out) =
905 crate::gpu_hook::gpu_dispatch().and_then(|d| d.try_fast_atv(a.view(), v.view()))
906 {
907 return out;
908 }
909 fast_atv_impl(a, v)
910}
911
912#[inline]
913fn fast_atv_impl<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
914 a: &ArrayBase<S1, Ix2>,
915 v: &ArrayBase<S2, Ix1>,
916) -> Array1<f64> {
917 use faer::Accum;
918 use faer::linalg::matmul::matmul;
919
920 let (n, p) = a.dim();
921 assert_eq!(n, v.len(), "A rows must match v length");
922
923 if let (Some(x_all), Some(vs)) = (a.as_slice(), v.as_slice())
927 && n != 0
928 && p != 0
929 {
930 let mut out = Array1::<f64>::zeros(p);
931 fast_atv_rowmajor_into(
932 x_all,
933 vs,
934 n,
935 p,
936 out.as_slice_mut().expect("fresh Array1 is contiguous"),
937 );
938 return out;
939 }
940
941 if !should_use_faer_matmul(p, 1, n) {
943 return a.t().dot(v);
944 }
945
946 let mut out = Array1::<f64>::zeros(p);
947 let mut outview = array1_to_col_matmut(&mut out);
948
949 let aview = FaerArrayView::new(a);
950 let vview = FaerColView::new(v);
951 let a_ref = aview.as_ref();
952 let v_ref = vview.as_ref();
953
954 let par = matmul_parallelism(p, 1, n);
956 matmul(
957 outview.as_mut(),
958 Accum::Replace,
959 a_ref.transpose(),
960 v_ref,
961 1.0,
962 par,
963 );
964
965 out
966}
967
968#[inline]
971pub fn fast_atv_into<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
972 a: &ArrayBase<S1, Ix2>,
973 v: &ArrayBase<S2, Ix1>,
974 out: &mut Array1<f64>,
975) {
976 fast_atv_into_impl(a, v, out);
977}
978
979#[inline]
980fn fast_atv_into_impl<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
981 a: &ArrayBase<S1, Ix2>,
982 v: &ArrayBase<S2, Ix1>,
983 out: &mut Array1<f64>,
984) {
985 use faer::Accum;
986 use faer::linalg::matmul::matmul;
987
988 let (n, p) = a.dim();
989 assert_eq!(v.len(), n, "vector length must match A rows");
990 assert_eq!(out.len(), p, "output length must match A cols");
991
992 if let (Some(x_all), Some(vs)) = (a.as_slice(), v.as_slice())
993 && n != 0
994 && p != 0
995 && let Some(out_s) = out.as_slice_mut()
996 {
997 fast_atv_rowmajor_into(x_all, vs, n, p, out_s);
998 return;
999 }
1000
1001 if !should_use_faer_matmul(p, 1, n) {
1002 out.assign(&a.t().dot(v));
1003 return;
1004 }
1005
1006 let mut outview = array1_to_col_matmut(out);
1007
1008 let aview = FaerArrayView::new(a);
1009 let vview = FaerColView::new(v);
1010 let a_ref = aview.as_ref();
1011 let v_ref = vview.as_ref();
1012 let par = matmul_parallelism(p, 1, n);
1013 matmul(
1014 outview.as_mut(),
1015 Accum::Replace,
1016 a_ref.transpose(),
1017 v_ref,
1018 1.0,
1019 par,
1020 );
1021}
1022
1023#[inline]
1025pub fn fast_xt_diag_x<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
1026 x: &ArrayBase<S1, Ix2>,
1027 w: &ArrayBase<S2, Ix1>,
1028) -> Array2<f64> {
1029 assert_eq!(
1030 x.nrows(),
1031 w.len(),
1032 "fast_xt_diag_x row/weight length mismatch"
1033 );
1034 if let Some(out) =
1035 crate::gpu_hook::gpu_dispatch().and_then(|d| d.try_fast_xt_diag_x(x.view(), w.view()))
1036 {
1037 return out;
1038 }
1039 let p = x.ncols();
1040 fast_xt_diag_x_with_parallelism(x, w, matmul_parallelism(p, p, x.nrows()))
1041}
1042
1043#[inline]
1046pub fn fast_xt_diag_x_with_parallelism<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
1047 x: &ArrayBase<S1, Ix2>,
1048 w: &ArrayBase<S2, Ix1>,
1049 par: Par,
1050) -> Array2<f64> {
1051 assert_eq!(
1052 x.nrows(),
1053 w.len(),
1054 "fast_xt_diag_x_with_parallelism row/weight length mismatch"
1055 );
1056 fast_xt_diag_x_with_parallelism_impl(x, w, par)
1057}
1058
1059#[inline]
1060fn fast_xt_diag_x_with_parallelism_impl<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
1061 x: &ArrayBase<S1, Ix2>,
1062 w: &ArrayBase<S2, Ix1>,
1063 par: Par,
1064) -> Array2<f64> {
1065 use ndarray::ShapeBuilder;
1066
1067 let p = x.ncols();
1068 let mut result = Array2::<f64>::zeros((p, p).f());
1071 stream_weighted_crossprod_into(
1072 x,
1073 w,
1074 &mut result,
1075 CrossprodStructure::SymmetricLower,
1076 CrossprodAccum::Replace,
1077 par,
1078 );
1079 result
1080}
1081
1082#[derive(Clone, Copy, PartialEq, Eq, Debug)]
1084pub enum CrossprodStructure {
1085 Full,
1087 SymmetricLower,
1091}
1092
1093#[derive(Clone, Copy, PartialEq, Eq, Debug)]
1095pub enum CrossprodAccum {
1096 Replace,
1098 Add,
1100}
1101
1102pub fn stream_weighted_crossprod_into<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
1121 x: &ArrayBase<S1, Ix2>,
1122 w: &ArrayBase<S2, Ix1>,
1123 out: &mut Array2<f64>,
1124 structure: CrossprodStructure,
1125 accum: CrossprodAccum,
1126 par: Par,
1127) {
1128 use faer::Accum;
1129 use faer::linalg::matmul::matmul;
1130 use faer::linalg::matmul::triangular::{BlockStructure, matmul as tri_matmul};
1131 use ndarray::s;
1132
1133 let (n, p) = x.dim();
1134 assert_eq!(n, w.len(), "X rows must match W length");
1135 assert_eq!(out.nrows(), p, "output rows must match X cols");
1136 assert_eq!(out.ncols(), p, "output cols must match X cols");
1137 if p == 0 {
1138 return;
1139 }
1140 if n == 0 {
1141 if accum == CrossprodAccum::Replace {
1142 out.fill(0.0);
1143 }
1144 return;
1145 }
1146
1147 if !should_use_faer_matmul(p, p, n) {
1148 let w_x = Array2::from_shape_fn((n, p), |(i, j)| w[i] * x[[i, j]]);
1150 let gram = x.t().dot(&w_x);
1151 match accum {
1152 CrossprodAccum::Replace => out.assign(&gram),
1153 CrossprodAccum::Add => *out += &gram,
1154 }
1155 return;
1156 }
1157
1158 const TARGET_BYTES: usize = 8 * 1024 * 1024;
1160 const MIN_ROWS: usize = 512;
1161 const MAX_ROWS: usize = 131_072;
1162 let chunk_rows = (TARGET_BYTES / (p.max(1) * 8))
1163 .clamp(MIN_ROWS, MAX_ROWS)
1164 .min(n);
1165
1166 if accum == CrossprodAccum::Replace {
1171 out.fill(0.0);
1172 }
1173
1174 let mut wx_chunk = Array2::<f64>::zeros((chunk_rows, p));
1180
1181 let x_is_row_major = x.is_standard_layout();
1182 let w_slice_opt = w.as_slice();
1183
1184 {
1187 let mut out_view = array2_to_matmut(out);
1188 for start in (0..n).step_by(chunk_rows) {
1189 let rows = (n - start).min(chunk_rows);
1190 {
1191 let chunk_slice = wx_chunk
1192 .as_slice_mut()
1193 .expect("row-major chunk is contiguous");
1194 if x_is_row_major && let (Some(x_all), Some(w_all)) = (x.as_slice(), w_slice_opt) {
1195 for local in 0..rows {
1196 let src = start + local;
1197 let wi = w_all[src];
1198 let src_off = src * p;
1199 let dst_off = local * p;
1200 let src_row = &x_all[src_off..src_off + p];
1201 let dst_row = &mut chunk_slice[dst_off..dst_off + p];
1202 for col in 0..p {
1203 dst_row[col] = src_row[col] * wi;
1204 }
1205 }
1206 } else {
1207 let x_slice = x.slice(s![start..start + rows, ..]);
1208 for local in 0..rows {
1209 let wi = w[start + local];
1210 let xrow = x_slice.row(local);
1211 let dst_off = local * p;
1212 let dst_row = &mut chunk_slice[dst_off..dst_off + p];
1213 for (col, xij) in xrow.iter().enumerate() {
1214 dst_row[col] = xij * wi;
1215 }
1216 }
1217 }
1218 }
1219 let x_slice = x.slice(s![start..start + rows, ..]);
1220 let wx_slice = wx_chunk.slice(s![0..rows, ..]);
1221 let x_view = FaerArrayView::new(&x_slice);
1222 let wx_view = FaerArrayView::new(&wx_slice);
1223 match structure {
1224 CrossprodStructure::SymmetricLower => {
1225 tri_matmul(
1229 out_view.as_mut(),
1230 BlockStructure::TriangularLower,
1231 Accum::Add,
1232 x_view.as_ref().transpose(),
1233 BlockStructure::Rectangular,
1234 wx_view.as_ref(),
1235 BlockStructure::Rectangular,
1236 1.0,
1237 par,
1238 );
1239 }
1240 CrossprodStructure::Full => {
1241 matmul(
1242 out_view.as_mut(),
1243 Accum::Add,
1244 x_view.as_ref().transpose(),
1245 wx_view.as_ref(),
1246 1.0,
1247 par,
1248 );
1249 }
1250 }
1251 }
1252 }
1253
1254 if structure == CrossprodStructure::SymmetricLower {
1255 for i in 0..p {
1257 for j in (i + 1)..p {
1258 out[[i, j]] = out[[j, i]];
1259 }
1260 }
1261 }
1262}
1263
1264#[inline]
1266pub fn fast_xt_diag_y<S1: Data<Elem = f64>, S2: Data<Elem = f64>, S3: Data<Elem = f64>>(
1267 x: &ArrayBase<S1, Ix2>,
1268 w: &ArrayBase<S2, Ix1>,
1269 y: &ArrayBase<S3, Ix2>,
1270) -> Array2<f64> {
1271 assert_eq!(x.nrows(), y.nrows(), "fast_xt_diag_y X/Y row mismatch");
1272 assert_eq!(
1273 y.nrows(),
1274 w.len(),
1275 "fast_xt_diag_y row/weight length mismatch"
1276 );
1277 if let Some(out) = crate::gpu_hook::gpu_dispatch()
1278 .and_then(|d| d.try_fast_xt_diag_y(x.view(), w.view(), y.view()))
1279 {
1280 return out;
1281 }
1282 fast_xt_diag_y_impl(x, w, y)
1283}
1284
1285#[inline]
1286fn fast_xt_diag_y_impl<S1: Data<Elem = f64>, S2: Data<Elem = f64>, S3: Data<Elem = f64>>(
1287 x: &ArrayBase<S1, Ix2>,
1288 w: &ArrayBase<S2, Ix1>,
1289 y: &ArrayBase<S3, Ix2>,
1290) -> Array2<f64> {
1291 use faer::Accum;
1292 use faer::linalg::matmul::matmul;
1293 use ndarray::{ShapeBuilder, s};
1294
1295 let (n, q) = y.dim();
1296 let px = x.ncols();
1297 assert_eq!(n, w.len(), "Y rows must match W length");
1298 assert_eq!(n, x.nrows(), "X rows must match Y rows");
1299 if n == 0 || px == 0 || q == 0 {
1300 return Array2::<f64>::zeros((px, q));
1301 }
1302 if !should_use_faer_matmul(px, q, n) {
1303 let w_y = Array2::from_shape_fn((n, q), |(i, j)| w[i] * y[[i, j]]);
1304 return x.t().dot(&w_y);
1305 }
1306
1307 const TARGET_BYTES: usize = 8 * 1024 * 1024;
1309 const MIN_ROWS: usize = 512;
1310 const MAX_ROWS: usize = 131_072;
1311 let total_cols = px + q;
1312 let chunk_rows = (TARGET_BYTES / (total_cols.max(1) * 8))
1313 .clamp(MIN_ROWS, MAX_ROWS)
1314 .min(n);
1315
1316 let mut result = Array2::<f64>::zeros((px, q).f());
1317 let mut wy_chunk = Array2::<f64>::zeros((chunk_rows, q));
1320
1321 let y_is_row_major = y.is_standard_layout();
1322 let w_slice_opt = w.as_slice();
1323
1324 {
1325 let mut out_view = array2_to_matmut(&mut result);
1326
1327 for start in (0..n).step_by(chunk_rows) {
1328 let rows = (n - start).min(chunk_rows);
1329 {
1330 let chunk_slice = wy_chunk
1331 .as_slice_mut()
1332 .expect("row-major chunk is contiguous");
1333 if y_is_row_major && let (Some(y_all), Some(w_all)) = (y.as_slice(), w_slice_opt) {
1334 for local in 0..rows {
1335 let src = start + local;
1336 let wi = w_all[src];
1337 let src_off = src * q;
1338 let dst_off = local * q;
1339 let src_row = &y_all[src_off..src_off + q];
1340 let dst_row = &mut chunk_slice[dst_off..dst_off + q];
1341 for col in 0..q {
1342 dst_row[col] = src_row[col] * wi;
1343 }
1344 }
1345 } else {
1346 let y_slice = y.slice(s![start..start + rows, ..]);
1347 for local in 0..rows {
1348 let wi = w[start + local];
1349 let yrow = y_slice.row(local);
1350 let dst_off = local * q;
1351 let dst_row = &mut chunk_slice[dst_off..dst_off + q];
1352 for (col, yij) in yrow.iter().enumerate() {
1353 dst_row[col] = yij * wi;
1354 }
1355 }
1356 }
1357 }
1358 let x_slice = x.slice(s![start..start + rows, ..]);
1359 let wy_slice = wy_chunk.slice(s![0..rows, ..]);
1360 let x_view = FaerArrayView::new(&x_slice);
1361 let wy_view = FaerArrayView::new(&wy_slice);
1362 let par = matmul_parallelism(px, q, rows);
1363 matmul(
1364 out_view.as_mut(),
1365 Accum::Add,
1366 x_view.as_ref().transpose(),
1367 wy_view.as_ref(),
1368 1.0,
1369 par,
1370 );
1371 }
1372 }
1373
1374 result
1375}
1376
1377pub fn fast_joint_hessian_2x2<
1383 S1: Data<Elem = f64>,
1384 S2: Data<Elem = f64>,
1385 S3: Data<Elem = f64>,
1386 S4: Data<Elem = f64>,
1387 S5: Data<Elem = f64>,
1388>(
1389 x_a: &ArrayBase<S1, Ix2>,
1390 x_b: &ArrayBase<S2, Ix2>,
1391 w_aa: &ArrayBase<S3, Ix1>,
1392 w_ab: &ArrayBase<S4, Ix1>,
1393 w_bb: &ArrayBase<S5, Ix1>,
1394) -> Array2<f64> {
1395 if let Some(out) = crate::gpu_hook::gpu_dispatch().and_then(|d| {
1396 d.try_fast_joint_hessian_2x2(
1397 x_a.view(),
1398 x_b.view(),
1399 w_aa.view(),
1400 w_ab.view(),
1401 w_bb.view(),
1402 )
1403 }) {
1404 return out;
1405 }
1406 fast_joint_hessian_2x2_impl(x_a, x_b, w_aa, w_ab, w_bb)
1407}
1408
1409#[inline]
1410fn fast_joint_hessian_2x2_impl<
1411 S1: Data<Elem = f64>,
1412 S2: Data<Elem = f64>,
1413 S3: Data<Elem = f64>,
1414 S4: Data<Elem = f64>,
1415 S5: Data<Elem = f64>,
1416>(
1417 x_a: &ArrayBase<S1, Ix2>,
1418 x_b: &ArrayBase<S2, Ix2>,
1419 w_aa: &ArrayBase<S3, Ix1>,
1420 w_ab: &ArrayBase<S4, Ix1>,
1421 w_bb: &ArrayBase<S5, Ix1>,
1422) -> Array2<f64> {
1423 use faer::Accum;
1424 use faer::linalg::matmul::matmul;
1425 use ndarray::{ShapeBuilder, s};
1426
1427 let n = x_a.nrows();
1428 let pa = x_a.ncols();
1429 let pb = x_b.ncols();
1430 let total = pa + pb;
1431 assert_eq!(n, x_b.nrows());
1432 assert_eq!(n, w_aa.len());
1433 assert_eq!(n, w_ab.len());
1434 assert_eq!(n, w_bb.len());
1435
1436 if n == 0 || total == 0 {
1437 return Array2::<f64>::zeros((total, total));
1438 }
1439
1440 if !should_use_faer_matmul(pa.max(pb), pa.max(pb), n) {
1442 let waa_xa = Array2::from_shape_fn((n, pa), |(i, j)| w_aa[i] * x_a[[i, j]]);
1443 let wab_xb = Array2::from_shape_fn((n, pb), |(i, j)| w_ab[i] * x_b[[i, j]]);
1444 let wbb_xb = Array2::from_shape_fn((n, pb), |(i, j)| w_bb[i] * x_b[[i, j]]);
1445 let mut out = Array2::<f64>::zeros((total, total));
1446 out.slice_mut(s![..pa, ..pa]).assign(&x_a.t().dot(&waa_xa));
1447 out.slice_mut(s![..pa, pa..]).assign(&x_a.t().dot(&wab_xb));
1448 out.slice_mut(s![pa.., pa..]).assign(&x_b.t().dot(&wbb_xb));
1449 for i in 0..total {
1451 for j in 0..i {
1452 out[[i, j]] = out[[j, i]];
1453 }
1454 }
1455 return out;
1456 }
1457
1458 const TARGET_BYTES: usize = 8 * 1024 * 1024;
1459 const MIN_ROWS: usize = 512;
1460 const MAX_ROWS: usize = 131_072;
1461 let cols_needed = pa + 2 * pb;
1463 let chunk_rows = (TARGET_BYTES / (cols_needed.max(1) * 8))
1464 .clamp(MIN_ROWS, MAX_ROWS)
1465 .min(n);
1466
1467 let mut out = Array2::<f64>::zeros((total, total).f());
1468 let mut waa_xa_buf = Array2::<f64>::zeros((chunk_rows, pa));
1473 let mut wab_xb_buf = Array2::<f64>::zeros((chunk_rows, pb));
1474 let mut wbb_xb_buf = Array2::<f64>::zeros((chunk_rows, pb));
1475
1476 let xa_is_row_major = x_a.is_standard_layout();
1477 let xb_is_row_major = x_b.is_standard_layout();
1478 let waa_slice_opt = w_aa.as_slice();
1479 let wab_slice_opt = w_ab.as_slice();
1480 let wbb_slice_opt = w_bb.as_slice();
1481
1482 {
1483 let mut out_mat = array2_to_matmut(&mut out);
1484
1485 for start in (0..n).step_by(chunk_rows) {
1486 let rows = (n - start).min(chunk_rows);
1487 let xa_slice = x_a.slice(s![start..start + rows, ..]);
1488 let xb_slice = x_b.slice(s![start..start + rows, ..]);
1489
1490 {
1492 let waa_chunk = waa_xa_buf
1493 .as_slice_mut()
1494 .expect("row-major waa chunk is contiguous");
1495 let wab_chunk = wab_xb_buf
1496 .as_slice_mut()
1497 .expect("row-major wab chunk is contiguous");
1498 let wbb_chunk = wbb_xb_buf
1499 .as_slice_mut()
1500 .expect("row-major wbb chunk is contiguous");
1501
1502 if xa_is_row_major
1503 && xb_is_row_major
1504 && let (Some(xa_all), Some(xb_all)) = (x_a.as_slice(), x_b.as_slice())
1505 && let (Some(waa_all), Some(wab_all), Some(wbb_all)) =
1506 (waa_slice_opt, wab_slice_opt, wbb_slice_opt)
1507 {
1508 for local in 0..rows {
1509 let i = start + local;
1510 let waa_i = waa_all[i];
1511 let wab_i = wab_all[i];
1512 let wbb_i = wbb_all[i];
1513 let xa_off = i * pa;
1514 let xa_row = &xa_all[xa_off..xa_off + pa];
1515 let xb_off = i * pb;
1516 let xb_row = &xb_all[xb_off..xb_off + pb];
1517 let waa_off = local * pa;
1518 let wab_off = local * pb;
1519 let wbb_off = local * pb;
1520 let waa_row = &mut waa_chunk[waa_off..waa_off + pa];
1521 for col in 0..pa {
1522 waa_row[col] = xa_row[col] * waa_i;
1523 }
1524 let wab_row = &mut wab_chunk[wab_off..wab_off + pb];
1525 let wbb_row = &mut wbb_chunk[wbb_off..wbb_off + pb];
1526 for col in 0..pb {
1527 let xij = xb_row[col];
1528 wab_row[col] = xij * wab_i;
1529 wbb_row[col] = xij * wbb_i;
1530 }
1531 }
1532 } else {
1533 for local in 0..rows {
1534 let i = start + local;
1535 let waa_i = w_aa[i];
1536 let wab_i = w_ab[i];
1537 let wbb_i = w_bb[i];
1538 let waa_off = local * pa;
1539 let wab_off = local * pb;
1540 let wbb_off = local * pb;
1541 let waa_row = &mut waa_chunk[waa_off..waa_off + pa];
1542 let xa_row = xa_slice.row(local);
1543 for (col, xij) in xa_row.iter().enumerate() {
1544 waa_row[col] = xij * waa_i;
1545 }
1546 let wab_row = &mut wab_chunk[wab_off..wab_off + pb];
1547 let wbb_row = &mut wbb_chunk[wbb_off..wbb_off + pb];
1548 let xb_row = xb_slice.row(local);
1549 for (col, xij) in xb_row.iter().enumerate() {
1550 wab_row[col] = xij * wab_i;
1551 wbb_row[col] = xij * wbb_i;
1552 }
1553 }
1554 }
1555 }
1556
1557 let xa_view = FaerArrayView::new(&xa_slice);
1558 let xb_view = FaerArrayView::new(&xb_slice);
1559 let waa_xa_slice = waa_xa_buf.slice(s![0..rows, ..]);
1560 let wab_xb_slice = wab_xb_buf.slice(s![0..rows, ..]);
1561 let wbb_xb_slice = wbb_xb_buf.slice(s![0..rows, ..]);
1562 let waa_xa_view = FaerArrayView::new(&waa_xa_slice);
1563 let wab_xb_view = FaerArrayView::new(&wab_xb_slice);
1564 let wbb_xb_view = FaerArrayView::new(&wbb_xb_slice);
1565
1566 matmul(
1568 out_mat.rb_mut().submatrix_mut(0, 0, pa, pa),
1569 Accum::Add,
1570 xa_view.as_ref().transpose(),
1571 waa_xa_view.as_ref(),
1572 1.0,
1573 matmul_parallelism(pa, pa, rows),
1574 );
1575 matmul(
1577 out_mat.rb_mut().submatrix_mut(0, pa, pa, pb),
1578 Accum::Add,
1579 xa_view.as_ref().transpose(),
1580 wab_xb_view.as_ref(),
1581 1.0,
1582 matmul_parallelism(pa, pb, rows),
1583 );
1584 matmul(
1586 out_mat.rb_mut().submatrix_mut(pa, pa, pb, pb),
1587 Accum::Add,
1588 xb_view.as_ref().transpose(),
1589 wbb_xb_view.as_ref(),
1590 1.0,
1591 matmul_parallelism(pb, pb, rows),
1592 );
1593 }
1594 } for i in 0..total {
1597 for j in 0..i {
1598 out[[i, j]] = out[[j, i]];
1599 }
1600 }
1601 out
1602}
1603
1604fn mat_to_array(mat: MatRef<'_, f64>) -> Array2<f64> {
1605 let nrows = mat.nrows();
1606 let ncols = mat.ncols();
1607 let mut out = Array2::<f64>::zeros((nrows, ncols));
1608 if nrows == 0 || ncols == 0 {
1609 return out;
1610 }
1611 if let Some(out_slice) = out.as_slice_memory_order_mut() {
1614 for i in 0..nrows {
1616 let row_start = i * ncols;
1617 for j in 0..ncols {
1618 out_slice[row_start + j] = mat[(i, j)];
1619 }
1620 }
1621 } else {
1622 for j in 0..ncols {
1623 for i in 0..nrows {
1624 out[[i, j]] = mat[(i, j)];
1625 }
1626 }
1627 }
1628 out
1629}
1630
1631#[inline]
1634pub fn fast_ab_into<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
1635 a: &ArrayBase<S1, Ix2>,
1636 b: &ArrayBase<S2, Ix2>,
1637 out: &mut Array2<f64>,
1638) {
1639 fast_ab_into_impl(a, b, out);
1640}
1641
1642#[inline]
1643fn fast_ab_into_impl<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
1644 a: &ArrayBase<S1, Ix2>,
1645 b: &ArrayBase<S2, Ix2>,
1646 out: &mut Array2<f64>,
1647) {
1648 use faer::Accum;
1649 use faer::linalg::matmul::matmul;
1650
1651 let (n, p) = a.dim();
1652 let (p_b, q) = b.dim();
1653 assert_eq!(p, p_b, "A and B must have compatible inner dimensions");
1654 assert_eq!(out.dim(), (n, q), "output dimensions must match A*B result");
1655
1656 if !should_use_faer_matmul(n, q, p) {
1657 out.assign(&a.dot(b));
1658 return;
1659 }
1660
1661 let aview = FaerArrayView::new(a);
1662 let bview = FaerArrayView::new(b);
1663 let a_ref = aview.as_ref();
1664 let b_ref = bview.as_ref();
1665
1666 let par = matmul_parallelism(n, q, p);
1667 let mut outview = array2_to_matmut(out);
1668 matmul(outview.as_mut(), Accum::Replace, a_ref, b_ref, 1.0, par);
1669}
1670
1671fn diag_to_array(diag: DiagRef<'_, f64>) -> Array1<f64> {
1672 let mat = diag.column_vector().as_mat();
1673 let mut out = Array1::<f64>::zeros(mat.nrows());
1674 for i in 0..mat.nrows() {
1675 out[i] = mat[(i, 0)];
1676 }
1677 out
1678}
1679
1680pub struct FaerArrayView<'a> {
1681 ptr: *const f64,
1682 rows: usize,
1683 cols: usize,
1684 row_stride: isize,
1685 col_stride: isize,
1686 owned: Option<Array2<f64>>,
1687 marker: PhantomData<&'a f64>,
1688}
1689
1690impl<'a> FaerArrayView<'a> {
1691 #[inline]
1692 pub fn new<S: Data<Elem = f64>>(array: &'a ArrayBase<S, Ix2>) -> Self {
1693 let (rows, cols) = array.dim();
1694 let strides = array.strides();
1695 if strides[0] <= 0 || strides[1] <= 0 {
1699 let owned = array.to_owned();
1700 let owned_strides = owned.strides();
1701 return Self {
1702 ptr: owned.as_ptr(),
1703 rows,
1704 cols,
1705 row_stride: owned_strides[0],
1706 col_stride: owned_strides[1],
1707 owned: Some(owned),
1708 marker: PhantomData,
1709 };
1710 }
1711
1712 Self {
1713 ptr: array.as_ptr(),
1714 rows,
1715 cols,
1716 row_stride: strides[0],
1717 col_stride: strides[1],
1718 owned: None,
1719 marker: PhantomData,
1720 }
1721 }
1722
1723 #[inline]
1724 pub fn as_ref(&self) -> MatRef<'_, f64> {
1725 let (ptr, rows, cols, row_stride, col_stride) = if let Some(owned) = &self.owned {
1726 let strides = owned.strides();
1727 (
1728 owned.as_ptr(),
1729 owned.nrows(),
1730 owned.ncols(),
1731 strides[0],
1732 strides[1],
1733 )
1734 } else {
1735 (
1736 self.ptr,
1737 self.rows,
1738 self.cols,
1739 self.row_stride,
1740 self.col_stride,
1741 )
1742 };
1743 unsafe { MatRef::from_raw_parts(ptr, rows, cols, row_stride, col_stride) }
1747 }
1748}
1749
1750pub struct FaerColView<'a> {
1751 ptr: *const f64,
1752 len: usize,
1753 stride: isize,
1754 owned: Option<Array1<f64>>,
1755 marker: PhantomData<&'a f64>,
1756}
1757
1758impl<'a> FaerColView<'a> {
1759 #[inline]
1760 pub fn new<S: Data<Elem = f64>>(array: &'a ArrayBase<S, Ix1>) -> Self {
1761 let len = array.len();
1762 let stride = array.strides()[0];
1763 if stride <= 0 {
1764 let owned = array.to_owned();
1765 return Self {
1766 ptr: owned.as_ptr(),
1767 len,
1768 stride: 1,
1769 owned: Some(owned),
1770 marker: PhantomData,
1771 };
1772 }
1773 Self {
1774 ptr: array.as_ptr(),
1775 len,
1776 stride,
1777 owned: None,
1778 marker: PhantomData,
1779 }
1780 }
1781
1782 #[inline]
1783 pub fn as_ref(&self) -> MatRef<'_, f64> {
1784 let (ptr, len, stride) = if let Some(owned) = &self.owned {
1785 (owned.as_ptr(), owned.len(), 1)
1786 } else {
1787 (self.ptr, self.len, self.stride)
1788 };
1789 unsafe { MatRef::from_raw_parts(ptr, len, 1, stride, 0) }
1793 }
1794}
1795
1796pub trait FaerSvd {
1797 fn svd(
1798 &self,
1799 compute_u: bool,
1800 computevt: bool,
1801 ) -> Result<(Option<Array2<f64>>, Array1<f64>, Option<Array2<f64>>), FaerLinalgError>;
1802}
1803
1804impl<S: Data<Elem = f64>> FaerSvd for ArrayBase<S, Ix2> {
1805 fn svd(
1806 &self,
1807 compute_u: bool,
1808 computevt: bool,
1809 ) -> Result<(Option<Array2<f64>>, Array1<f64>, Option<Array2<f64>>), FaerLinalgError> {
1810 let faerview = FaerArrayView::new(self);
1811 let faer_mat = faerview.as_ref();
1812 if !compute_u && !computevt {
1813 let (rows, cols) = faer_mat.shape();
1814 let mut singular = Diag::<f64>::zeros(rows.min(cols));
1815 let par = get_global_parallelism();
1816 let mut mem = MemBuffer::new(svd::svd_scratch::<f64>(
1817 rows,
1818 cols,
1819 ComputeSvdVectors::No,
1820 ComputeSvdVectors::No,
1821 par,
1822 Default::default(),
1823 ));
1824 let stack = MemStack::new(&mut mem);
1825 svd::svd(
1826 faer_mat,
1827 singular.as_mut(),
1828 None,
1829 None,
1830 par,
1831 stack,
1832 Default::default(),
1833 )
1834 .map_err(|_| FaerLinalgError::SvdNoConvergence {
1835 context: "faer SVD singular values only",
1836 })?;
1837 let singularvalues = diag_to_array(singular.as_ref());
1838 return Ok((None, singularvalues, None));
1839 }
1840
1841 let (rows, cols) = faer_mat.shape();
1842 let rank = rows.min(cols);
1843 let compute_u_flag = if compute_u {
1844 ComputeSvdVectors::Thin
1845 } else {
1846 ComputeSvdVectors::No
1847 };
1848 let computev_flag = if computevt {
1849 ComputeSvdVectors::Thin
1850 } else {
1851 ComputeSvdVectors::No
1852 };
1853
1854 let mut singular = Diag::<f64>::zeros(rows.min(cols));
1855 let mut u_storage = compute_u.then(|| Mat::<f64>::zeros(rows, rank));
1856 let mut v_storage = computevt.then(|| Mat::<f64>::zeros(cols, rank));
1857
1858 let par = get_global_parallelism();
1859 let mut mem = MemBuffer::new(svd::svd_scratch::<f64>(
1860 rows,
1861 cols,
1862 compute_u_flag,
1863 computev_flag,
1864 par,
1865 Default::default(),
1866 ));
1867 let stack = MemStack::new(&mut mem);
1868
1869 svd::svd(
1870 faer_mat.as_ref(),
1871 singular.as_mut(),
1872 u_storage.as_mut().map(|mat| mat.as_mut()),
1873 v_storage.as_mut().map(|mat| mat.as_mut()),
1874 par,
1875 stack,
1876 Default::default(),
1877 )
1878 .map_err(|_| FaerLinalgError::SvdNoConvergence {
1879 context: "faer SVD with vectors",
1880 })?;
1881
1882 let singularvalues = diag_to_array(singular.as_ref());
1883 let u_opt = u_storage.map(|mat| mat_to_array(mat.as_ref()));
1884 let vt_opt = v_storage.map(|mat| {
1885 let mat_ref = mat.as_ref();
1886 let mut out = Array2::<f64>::zeros((mat_ref.ncols(), mat_ref.nrows()));
1887 for j in 0..mat_ref.nrows() {
1888 for i in 0..mat_ref.ncols() {
1889 out[[i, j]] = mat_ref[(j, i)];
1890 }
1891 }
1892 out
1893 });
1894
1895 Ok((u_opt, singularvalues, vt_opt))
1896 }
1897}
1898
1899pub trait FaerEigh {
1900 fn eigh(&self, side: Side) -> Result<(Array1<f64>, Array2<f64>), FaerLinalgError>;
1901}
1902
1903pub fn strict_symmetric_eigh<S: Data<Elem = f64>>(
1910 matrix: &ArrayBase<S, Ix2>,
1911 side: Side,
1912) -> Result<(Array1<f64>, Array2<f64>), FaerLinalgError> {
1913 let owned = matrix.to_owned();
1914 if owned.nrows() == 0 || owned.nrows() != owned.ncols() {
1915 return Err(FaerLinalgError::StrictSelfAdjointEigenInvalidInput {
1916 reason: format!(
1917 "expected non-empty square matrix, got {}x{}",
1918 owned.nrows(),
1919 owned.ncols()
1920 ),
1921 });
1922 }
1923 crate::utils::validate_finite_symmetric_matrix(
1924 &owned,
1925 "strict self-adjoint eigendecomposition",
1926 )
1927 .map_err(
1928 |error| FaerLinalgError::StrictSelfAdjointEigenInvalidInput {
1929 reason: error.to_string(),
1930 },
1931 )?;
1932 let view = FaerArrayView::new(&owned);
1933 let eigen = catch_unwind(AssertUnwindSafe(|| view.as_ref().self_adjoint_eigen(side)))
1934 .map_err(|_| FaerLinalgError::FactorizationFailed {
1935 context: "strict self-adjoint eigendecomposition panic boundary",
1936 })?
1937 .map_err(FaerLinalgError::SelfAdjointEigen)?;
1938 let values = diag_to_array(eigen.S());
1939 let vectors = mat_to_array(eigen.U());
1940 if values.iter().any(|value| !value.is_finite())
1941 || vectors.iter().any(|value| !value.is_finite())
1942 {
1943 return Err(FaerLinalgError::SelfAdjointEigenNonFiniteInput {
1944 context: "strict self-adjoint eigendecomposition output validation",
1945 });
1946 }
1947 Ok((values, vectors))
1948}
1949
1950impl<S: Data<Elem = f64>> FaerEigh for ArrayBase<S, Ix2> {
1951 fn eigh(&self, side: Side) -> Result<(Array1<f64>, Array2<f64>), FaerLinalgError> {
1952 fn try_eigh(
1953 matrix: &Array2<f64>,
1954 side: Side,
1955 ) -> Result<(Array1<f64>, Array2<f64>), FaerLinalgError> {
1956 let faerview = FaerArrayView::new(matrix);
1957 let eigen = catch_unwind(AssertUnwindSafe(|| {
1958 faerview.as_ref().self_adjoint_eigen(side)
1959 }))
1960 .map_err(|_| FaerLinalgError::FactorizationFailed {
1961 context: "self-adjoint eigendecomposition panic boundary",
1962 })?
1963 .map_err(FaerLinalgError::SelfAdjointEigen)?;
1964 let values = diag_to_array(eigen.S());
1965 let vectors = mat_to_array(eigen.U());
1966 Ok((values, vectors))
1967 }
1968
1969 let owned = self.to_owned();
1970 if owned.nrows() != owned.ncols() {
1971 return Err(FaerLinalgError::FactorizationFailed {
1972 context: "self-adjoint eigendecomposition non-square input",
1973 });
1974 }
1975 if owned.nrows() == 0 {
1976 return Ok((Array1::zeros(0), Array2::zeros((0, 0))));
1977 }
1978 if owned.iter().any(|value| !value.is_finite()) {
1979 return Err(FaerLinalgError::SelfAdjointEigenNonFiniteInput {
1980 context: "self-adjoint eigendecomposition input validation",
1981 });
1982 }
1983 if let Ok((evals, evecs)) = try_eigh(&owned, side)
1984 && evals.iter().all(|value| value.is_finite())
1985 && evecs.iter().all(|value| value.is_finite())
1986 {
1987 return Ok((evals, evecs));
1988 }
1989
1990 let mut repaired = owned.clone();
1991 crate::matrix::symmetrize_in_place(&mut repaired);
1992
1993 let scale = repaired
1994 .iter()
1995 .fold(0.0_f64, |acc, &value| acc.max(value.abs()))
1996 .max(1.0);
1997 let scaled = repaired.mapv(|value| value / scale);
1998 const JITTER_SCHEDULE: [f64; 6] = [0.0, 1e-12, 1e-10, 1e-8, 1e-6, 1e-4];
2004 let jitter_schedule = JITTER_SCHEDULE;
2005 let mut last_error = FaerLinalgError::FactorizationFailed {
2006 context: "self-adjoint eigendecomposition repair attempts",
2007 };
2008
2009 for &jitter in &jitter_schedule {
2010 let mut candidate = scaled.clone();
2011 if jitter > 0.0 {
2012 let n = candidate.nrows();
2013 for i in 0..n {
2014 candidate[[i, i]] += jitter;
2015 }
2016 }
2017
2018 match try_eigh(&candidate, side) {
2019 Ok((mut evals, evecs))
2020 if evals.iter().all(|value| value.is_finite())
2021 && evecs.iter().all(|value| value.is_finite()) =>
2022 {
2023 for value in &mut evals {
2024 *value = (*value - jitter) * scale;
2025 }
2026 return Ok((evals, evecs));
2027 }
2028 Ok((_, _)) => {
2029 last_error = FaerLinalgError::SelfAdjointEigenNonFiniteInput {
2030 context: "self-adjoint eigendecomposition repaired output validation",
2031 };
2032 }
2033 Err(err) => {
2034 last_error = err;
2035 }
2036 }
2037 }
2038
2039 Err(last_error)
2040 }
2041}
2042
2043pub struct FaerCholeskyFactor {
2044 factor: solvers::Llt<f64>,
2045}
2046
2047impl FaerCholeskyFactor {
2048 pub fn solvevec(&self, rhs: &Array1<f64>) -> Array1<f64> {
2049 let mut rhs = rhs.to_owned();
2050 let mut rhsview = array1_to_col_matmut(&mut rhs);
2051 self.factor.solve_in_place(rhsview.as_mut());
2052 rhs
2053 }
2054
2055 pub fn solve_mat_in_place(&self, rhs: &mut Array2<f64>) {
2056 let mut rhsview = array2_to_matmut(rhs);
2057 self.factor.solve_in_place(rhsview.as_mut());
2058 }
2059
2060 pub fn solve_mat_into<S: Data<Elem = f64>>(
2061 &self,
2062 rhs: &ArrayBase<S, Ix2>,
2063 out: &mut Array2<f64>,
2064 ) {
2065 if out.dim() != rhs.dim() {
2066 *out = Array2::<f64>::zeros(rhs.dim());
2067 }
2068 out.assign(rhs);
2069 self.solve_mat_in_place(out);
2070 }
2071
2072 pub fn solve_mat(&self, rhs: &Array2<f64>) -> Array2<f64> {
2073 let mut out = Array2::<f64>::zeros(rhs.dim());
2074 self.solve_mat_into(rhs, &mut out);
2075 out
2076 }
2077
2078 pub fn diag(&self) -> Array1<f64> {
2079 diag_to_array(self.factor.L().diagonal())
2080 }
2081
2082 pub fn lower_triangular(&self) -> Array2<f64> {
2083 mat_to_array(self.factor.L())
2084 }
2085}
2086
2087impl crate::matrix::FactorizedSystem for FaerCholeskyFactor {
2088 fn solve(&self, rhs: &Array1<f64>) -> Result<Array1<f64>, String> {
2089 let out = self.solvevec(rhs);
2090 if out.iter().all(|value| value.is_finite()) {
2091 Ok(out)
2092 } else {
2093 Err("strict Cholesky solve produced non-finite values".to_string())
2094 }
2095 }
2096
2097 fn solvemulti(&self, rhs: &Array2<f64>) -> Result<Array2<f64>, String> {
2098 let out = self.solve_mat(rhs);
2099 if out.iter().all(|value| value.is_finite()) {
2100 Ok(out)
2101 } else {
2102 Err("strict Cholesky multi-solve produced non-finite values".to_string())
2103 }
2104 }
2105
2106 fn logdet(&self) -> f64 {
2107 cholesky_factor_logdet(self.factor.L())
2108 }
2109}
2110
2111pub trait FaerCholesky {
2112 fn cholesky(&self, side: Side) -> Result<FaerCholeskyFactor, FaerLinalgError>;
2113}
2114
2115impl<S: Data<Elem = f64>> FaerCholesky for ArrayBase<S, Ix2> {
2116 fn cholesky(&self, side: Side) -> Result<FaerCholeskyFactor, FaerLinalgError> {
2117 let faerview = FaerArrayView::new(self);
2118 let factor = faerview
2119 .as_ref()
2120 .llt(side)
2121 .map_err(FaerLinalgError::Cholesky)?;
2122 Ok(FaerCholeskyFactor { factor })
2123 }
2124}
2125
2126pub trait FaerQr {
2127 fn qr(&self) -> Result<(Array2<f64>, Array2<f64>), FaerLinalgError>;
2128}
2129
2130impl<S: Data<Elem = f64>> FaerQr for ArrayBase<S, Ix2> {
2131 fn qr(&self) -> Result<(Array2<f64>, Array2<f64>), FaerLinalgError> {
2132 let faerview = FaerArrayView::new(self);
2133 let qr = faerview.as_ref().qr();
2134 let q = qr.compute_thin_Q();
2135 let r = qr.thin_R();
2136 Ok((mat_to_array(q.as_ref()), mat_to_array(r)))
2137 }
2138}
2139
2140pub fn rrqr_nullspace_basis<S: Data<Elem = f64>>(
2159 a: &ArrayBase<S, Ix2>,
2160 rank_alpha: f64,
2161) -> Result<(Array2<f64>, usize), FaerLinalgError> {
2162 rrqr_nullspace_basis_inner(a, RrqrRankCutoff::RelativeAlpha(rank_alpha))
2163}
2164
2165#[derive(Debug, Clone, Copy)]
2168enum RrqrRankCutoff {
2169 RelativeAlpha(f64),
2173 Absolute(f64),
2179}
2180
2181pub fn rrqr_nullspace_basis_with_cutoff<S: Data<Elem = f64>>(
2194 a: &ArrayBase<S, Ix2>,
2195 cutoff: f64,
2196) -> Result<(Array2<f64>, usize), FaerLinalgError> {
2197 rrqr_nullspace_basis_inner(a, RrqrRankCutoff::Absolute(cutoff))
2198}
2199
2200fn rrqr_nullspace_basis_inner<S: Data<Elem = f64>>(
2201 a: &ArrayBase<S, Ix2>,
2202 cutoff: RrqrRankCutoff,
2203) -> Result<(Array2<f64>, usize), FaerLinalgError> {
2204 let faerview = FaerArrayView::new(a);
2205 let qr = faerview.as_ref().col_piv_qr();
2206 let r = qr.thin_R();
2207 let diag_len = r.nrows().min(r.ncols());
2208 let leading_diag = if diag_len > 0 { r[(0, 0)].abs() } else { 0.0 };
2209 let tol = match cutoff {
2210 RrqrRankCutoff::RelativeAlpha(rank_alpha) => {
2211 rank_alpha
2212 * f64::EPSILON
2213 * (a.nrows().max(a.ncols()).max(1) as f64)
2214 * leading_diag.max(1.0)
2215 }
2216 RrqrRankCutoff::Absolute(tol) => tol,
2217 };
2218 let rank = (0..diag_len).filter(|&i| r[(i, i)].abs() > tol).count();
2219 let z = if rank >= a.nrows() {
2220 Array2::<f64>::zeros((a.nrows(), 0))
2221 } else if rank == 0 {
2222 Array2::<f64>::eye(a.nrows())
2226 } else {
2227 let nullity = a.nrows() - rank;
2228 let mut selector = Mat::<f64>::zeros(a.nrows(), nullity);
2229 for j in 0..nullity {
2230 selector[(rank + j, j)] = 1.0;
2231 }
2232 let par = get_global_parallelism();
2233 faer::linalg::householder::apply_block_householder_sequence_on_the_left_in_place_with_conj(
2234 qr.Q_basis(),
2235 qr.Q_coeff(),
2236 Conj::No,
2237 selector.as_mut(),
2238 par,
2239 MemStack::new(&mut MemBuffer::new(
2240 faer::linalg::householder::apply_block_householder_sequence_on_the_left_in_place_scratch::<f64>(
2241 a.nrows(),
2242 qr.Q_coeff().nrows(),
2243 nullity,
2244 ),
2245 )),
2246 );
2247 mat_to_array(selector.as_ref())
2248 };
2249 Ok((z, rank))
2250}
2251
2252#[inline]
2253pub const fn default_rrqr_rank_alpha() -> f64 {
2254 RRQR_RANK_ALPHA
2255}
2256
2257pub struct RrqrWithPermutation {
2268 pub rank: usize,
2269 pub column_permutation: Vec<usize>,
2270 pub leading_diag_abs: f64,
2271 pub rank_tol: f64,
2272}
2273
2274pub fn rrqr_with_permutation<S: Data<Elem = f64>>(
2283 a: &ArrayBase<S, Ix2>,
2284 rank_alpha: f64,
2285) -> Result<RrqrWithPermutation, FaerLinalgError> {
2286 if a.nrows() == 0 {
2287 return Err(FaerLinalgError::FactorizationFailed {
2288 context: "rrqr_with_permutation: input has zero rows",
2289 });
2290 }
2291 let faerview = FaerArrayView::new(a);
2292 let qr = faerview.as_ref().col_piv_qr();
2293 let r = qr.thin_R();
2294 let diag_len = r.nrows().min(r.ncols());
2295 let leading_diag = if diag_len > 0 { r[(0, 0)].abs() } else { 0.0 };
2296 let tol = rank_alpha
2297 * f64::EPSILON
2298 * (a.nrows().max(a.ncols()).max(1) as f64)
2299 * leading_diag.max(1.0);
2300 let rank = (0..diag_len).filter(|&i| r[(i, i)].abs() > tol).count();
2301 let (forward, _inverse) = qr.P().arrays();
2302 let column_permutation: Vec<usize> = forward.iter().copied().map(|idx| idx.unbound()).collect();
2303 Ok(RrqrWithPermutation {
2304 rank,
2305 column_permutation,
2306 leading_diag_abs: leading_diag,
2307 rank_tol: tol,
2308 })
2309}
2310
2311pub struct RrqrFromGram {
2320 pub rank: usize,
2321 pub column_permutation: Vec<usize>,
2322 pub rank_tol: f64,
2323 pub leading_diag_abs: f64,
2328 pub verdict_margin: f64,
2331}
2332
2333pub fn rrqr_from_gram_with_permutation<S: Data<Elem = f64>>(
2369 gram: &ArrayBase<S, Ix2>,
2370 m_rows: usize,
2371 rank_alpha: f64,
2372) -> Result<RrqrFromGram, FaerLinalgError> {
2373 let p = gram.ncols();
2374 if p == 0 {
2375 return Ok(RrqrFromGram {
2376 rank: 0,
2377 column_permutation: Vec::new(),
2378 rank_tol: 0.0,
2379 leading_diag_abs: 0.0,
2380 verdict_margin: 0.0,
2381 });
2382 }
2383 if gram.nrows() != p {
2384 return Err(FaerLinalgError::FactorizationFailed {
2385 context: "rrqr_from_gram_with_permutation: Gram is not square",
2386 });
2387 }
2388 let (evals, evecs) = gram.eigh(Side::Lower)?;
2397 let mut f = Array2::<f64>::zeros((p, p));
2398 for k in 0..p {
2399 let scale = evals[k].max(0.0).sqrt();
2400 if scale == 0.0 {
2401 continue;
2402 }
2403 for i in 0..p {
2404 f[[k, i]] = scale * evecs[[i, k]];
2405 }
2406 }
2407 let faer_f = FaerArrayView::new(&f);
2411 let qr = faer_f.as_ref().col_piv_qr();
2412 let r = qr.thin_R();
2413 let diag_len = r.nrows().min(r.ncols());
2414 let pivots: Vec<f64> = (0..diag_len).map(|i| r[(i, i)].abs()).collect();
2415 let leading_diag = pivots.first().copied().unwrap_or(0.0);
2416 let (forward, _inverse) = qr.P().arrays();
2417 let column_permutation: Vec<usize> = forward.iter().copied().map(|idx| idx.unbound()).collect();
2418 let tol = rank_alpha * f64::EPSILON * (m_rows.max(p).max(1) as f64) * leading_diag.max(1.0);
2422 let rank = pivots.iter().filter(|&&v| v > tol).count();
2423 let min_kept = pivots[..rank].iter().copied().fold(f64::INFINITY, f64::min);
2424 let max_dropped = pivots[rank..].iter().copied().fold(0.0f64, f64::max);
2425 let kept_margin = if rank == 0 {
2429 f64::INFINITY
2430 } else {
2431 min_kept / tol
2432 };
2433 let dropped_margin = if rank == diag_len {
2434 f64::INFINITY
2435 } else {
2436 tol / max_dropped.max(f64::MIN_POSITIVE)
2437 };
2438 let gram_precision_floor = f64::EPSILON.sqrt() * leading_diag.max(1.0);
2460 let kept_floor_margin = if rank == 0 {
2461 f64::INFINITY
2462 } else {
2463 min_kept / gram_precision_floor.max(f64::MIN_POSITIVE)
2464 };
2465 let verdict_margin = kept_margin.min(dropped_margin).min(kept_floor_margin);
2466 Ok(RrqrFromGram {
2467 rank,
2468 column_permutation,
2469 rank_tol: tol,
2470 leading_diag_abs: leading_diag,
2471 verdict_margin,
2472 })
2473}
2474
2475#[cfg(test)]
2476mod tests {
2477 use super::*;
2478 use ndarray::{array, s};
2479
2480 const JOINT_GRAM_RRQR_TRUST_MARGIN_FOR_TEST: f64 = 1.0e3;
2484
2485 #[test]
2486 fn rrqr_nullspace_basis_is_orthonormal_and_annihilates_transpose() {
2487 let a = array![[1.0, 0.0], [1.0, 0.0], [0.0, 2.0], [0.0, 0.0],];
2488 let (z, rank) =
2489 rrqr_nullspace_basis(&a, default_rrqr_rank_alpha()).expect("RRQR should succeed");
2490 assert_eq!(rank, 2);
2491 assert_eq!(z.nrows(), 4);
2492 assert_eq!(z.ncols(), 2);
2493
2494 let gram = z.t().dot(&z);
2495 let ident = Array2::<f64>::eye(z.ncols());
2496 let gram_err = (&gram - &ident)
2497 .iter()
2498 .fold(0.0_f64, |acc, &v| acc.max(v.abs()));
2499 assert!(gram_err < 1e-10, "Z is not orthonormal: {gram_err:e}");
2500
2501 let residual = a.t().dot(&z);
2502 let resid_max = residual.iter().fold(0.0_f64, |acc, &v| acc.max(v.abs()));
2503 assert!(resid_max < 1e-10, "A^T Z residual too large: {resid_max:e}");
2504 }
2505
2506 #[test]
2507 fn rrqr_with_permutation_attributes_redundant_column() {
2508 let a = array![
2512 [1.0, 0.0, 1.0],
2513 [1.0, 0.0, 1.0],
2514 [0.0, 2.0, 0.0],
2515 [0.0, 0.0, 0.0],
2516 ];
2517 let result =
2518 rrqr_with_permutation(&a, default_rrqr_rank_alpha()).expect("RRQR should succeed");
2519 assert_eq!(result.rank, 2);
2520 assert_eq!(result.column_permutation.len(), 3);
2521 let demoted = result.column_permutation[result.rank..].to_vec();
2522 assert!(
2523 demoted.contains(&2) || demoted.contains(&0),
2524 "demoted suffix should include one of the aliased columns (0 or 2), got {demoted:?}"
2525 );
2526 let mut sorted = result.column_permutation.clone();
2527 sorted.sort();
2528 assert_eq!(
2529 sorted,
2530 vec![0, 1, 2],
2531 "permutation must be a valid bijection on 0..n"
2532 );
2533 }
2534
2535 #[test]
2536 fn rrqr_with_permutation_full_rank_returns_identity_like_order() {
2537 let a = array![[1.0, 0.0], [0.0, 2.0], [0.0, 0.0]];
2538 let result =
2539 rrqr_with_permutation(&a, default_rrqr_rank_alpha()).expect("RRQR should succeed");
2540 assert_eq!(result.rank, 2);
2541 let mut sorted = result.column_permutation.clone();
2542 sorted.sort();
2543 assert_eq!(sorted, vec![0, 1]);
2544 }
2545
2546 #[test]
2547 fn rrqr_with_permutation_rejects_zero_rows() {
2548 let a = Array2::<f64>::zeros((0, 3));
2549 assert!(rrqr_with_permutation(&a, default_rrqr_rank_alpha()).is_err());
2550 }
2551
2552 #[test]
2553 fn rrqr_nullspace_basis_square_zero_matrix_is_finite_identity() {
2554 let a = Array2::<f64>::zeros((3, 3));
2557 let (z, rank) =
2558 rrqr_nullspace_basis(&a, default_rrqr_rank_alpha()).expect("RRQR should succeed");
2559 assert_eq!(rank, 0);
2560 assert_eq!(z.dim(), (3, 3));
2561 assert!(
2562 z.iter().all(|v| v.is_finite()),
2563 "square zero matrix produced a non-finite null basis: {z:?}"
2564 );
2565 let gram = z.t().dot(&z);
2566 let ident = Array2::<f64>::eye(3);
2567 let gram_err = (&gram - &ident)
2568 .iter()
2569 .fold(0.0_f64, |acc, &v| acc.max(v.abs()));
2570 assert!(gram_err < 1e-10, "Z is not orthonormal: {gram_err:e}");
2571 }
2572
2573 #[test]
2574 fn rrqr_nullspace_basis_detectszero_rank_matrix() {
2575 let a = Array2::<f64>::zeros((5, 2));
2576 let (z, rank) =
2577 rrqr_nullspace_basis(&a, default_rrqr_rank_alpha()).expect("RRQR should succeed");
2578 assert_eq!(rank, 0);
2579 assert_eq!(z.dim(), (5, 5));
2580 let ident = Array2::<f64>::eye(5);
2581 let max_err = (&z.slice(s![.., ..5]).to_owned() - &ident)
2582 .iter()
2583 .fold(0.0_f64, |acc, &v| acc.max(v.abs()));
2584 assert!(max_err < 1e-10, "zero matrix should yield identity basis");
2585 }
2586
2587 #[test]
2596 fn eigh_on_nan_matrix_rejects_non_finite_input() {
2597 let mat = array![
2598 [1.0, 0.0, 0.0, 0.0],
2599 [0.0, 2.0, 0.0, 0.0],
2600 [0.0, 0.0, 3.0, f64::NAN],
2601 [0.0, 0.0, f64::NAN, 4.0]
2602 ];
2603 let err = mat
2604 .eigh(Side::Lower)
2605 .expect_err("non-finite symmetric input must be rejected");
2606 assert!(matches!(
2607 err,
2608 FaerLinalgError::SelfAdjointEigenNonFiniteInput { .. }
2609 ));
2610 }
2611
2612 #[test]
2613 fn fast_ata_matches_full_gemm_above_threshold() {
2614 let n = 200;
2617 let p = 40;
2618 let a: Array2<f64> = Array2::from_shape_fn((n, p), |(i, j)| {
2619 ((i * 7 + j * 3) as f64).sin() + 0.1 * j as f64
2620 });
2621 let expected = a.t().dot(&a);
2622 let got = fast_ata(&a);
2623 let max_err = (&got - &expected)
2624 .iter()
2625 .fold(0.0_f64, |acc, &v| acc.max(v.abs()));
2626 assert!(max_err < 1e-10, "fast_ata mismatch: {max_err:e}");
2627 for i in 0..p {
2629 for j in 0..p {
2630 assert!((got[[i, j]] - got[[j, i]]).abs() < 1e-12);
2631 }
2632 }
2633 }
2634
2635 #[test]
2636 fn fast_xt_diag_x_matches_naive_above_threshold() {
2637 let n = 400;
2638 let p = 36;
2639 let x: Array2<f64> =
2640 Array2::from_shape_fn((n, p), |(i, j)| (i as f64 * 0.1).cos() + j as f64 * 0.05);
2641 let w: Array1<f64> = Array1::from_shape_fn(n, |i| (i as f64 * 0.03).sin());
2642 let wx = Array2::from_shape_fn((n, p), |(i, j)| w[i] * x[[i, j]]);
2644 let expected = x.t().dot(&wx);
2645 let got = fast_xt_diag_x(&x, &w);
2646 let max_err = (&got - &expected)
2647 .iter()
2648 .fold(0.0_f64, |acc, &v| acc.max(v.abs()));
2649 assert!(max_err < 1e-9, "fast_xt_diag_x mismatch: {max_err:e}");
2650 for i in 0..p {
2651 for j in 0..p {
2652 assert!((got[[i, j]] - got[[j, i]]).abs() < 1e-12);
2653 }
2654 }
2655 }
2656
2657 #[test]
2658 fn stream_weighted_crossprod_full_and_triangular_parity_with_negative_weights() {
2659 for &(n, p) in &[(900usize, 40usize), (8usize, 3usize)] {
2668 let x: Array2<f64> =
2669 Array2::from_shape_fn((n, p), |(i, j)| (i as f64 * 0.07).cos() + j as f64 * 0.013);
2670 let w: Array1<f64> =
2673 Array1::from_shape_fn(n, |i| (i as f64 * 0.11).sin() - 0.25 * (i % 3) as f64);
2674 assert!(
2675 w.iter().any(|&v| v < 0.0),
2676 "weight vector must contain negatives to test sign preservation"
2677 );
2678
2679 let wx = Array2::from_shape_fn((n, p), |(i, j)| w[i] * x[[i, j]]);
2681 let expected = x.t().dot(&wx);
2682
2683 let par = matmul_parallelism(p, p, n);
2684
2685 let mut full = Array2::<f64>::ones((p, p));
2687 stream_weighted_crossprod_into(
2688 &x,
2689 &w,
2690 &mut full,
2691 CrossprodStructure::Full,
2692 CrossprodAccum::Replace,
2693 par,
2694 );
2695
2696 let mut tri = Array2::<f64>::from_elem((p, p), -7.0);
2700 stream_weighted_crossprod_into(
2701 &x,
2702 &w,
2703 &mut tri,
2704 CrossprodStructure::SymmetricLower,
2705 CrossprodAccum::Replace,
2706 par,
2707 );
2708
2709 let full_err = (&full - &expected)
2710 .iter()
2711 .fold(0.0_f64, |a, &v| a.max(v.abs()));
2712 let tri_err = (&tri - &expected)
2713 .iter()
2714 .fold(0.0_f64, |a, &v| a.max(v.abs()));
2715 assert!(
2716 full_err < 1e-9,
2717 "full kernel mismatch (n={n}, p={p}): {full_err:e}"
2718 );
2719 assert!(
2720 tri_err < 1e-9,
2721 "triangular kernel mismatch (n={n}, p={p}): {tri_err:e}"
2722 );
2723
2724 for i in 0..p {
2727 for j in 0..p {
2728 assert!(
2729 (full[[i, j]] - tri[[i, j]]).abs() < 1e-12,
2730 "full vs triangular disagree at ({i},{j})"
2731 );
2732 assert!(
2733 (tri[[i, j]] - tri[[j, i]]).abs() < 1e-12,
2734 "triangular output not symmetric at ({i},{j})"
2735 );
2736 }
2737 }
2738
2739 let base = Array2::<f64>::from_elem((p, p), 1.5);
2742 let mut add_full = base.clone();
2743 stream_weighted_crossprod_into(
2744 &x,
2745 &w,
2746 &mut add_full,
2747 CrossprodStructure::Full,
2748 CrossprodAccum::Add,
2749 par,
2750 );
2751 let mut add_tri = base.clone();
2752 stream_weighted_crossprod_into(
2753 &x,
2754 &w,
2755 &mut add_tri,
2756 CrossprodStructure::SymmetricLower,
2757 CrossprodAccum::Add,
2758 par,
2759 );
2760 let expected_add = &base + &expected;
2761 let add_full_err = (&add_full - &expected_add)
2762 .iter()
2763 .fold(0.0_f64, |a, &v| a.max(v.abs()));
2764 let add_tri_err = (&add_tri - &expected_add)
2765 .iter()
2766 .fold(0.0_f64, |a, &v| a.max(v.abs()));
2767 assert!(
2768 add_full_err < 1e-9,
2769 "full Add mismatch (n={n}, p={p}): {add_full_err:e}"
2770 );
2771 assert!(
2772 add_tri_err < 1e-9,
2773 "triangular Add mismatch (n={n}, p={p}): {add_tri_err:e}"
2774 );
2775
2776 let returned = fast_xt_diag_x(&x, &w);
2779 let returned_err = (&returned - &full)
2780 .iter()
2781 .fold(0.0_f64, |a, &v| a.max(v.abs()));
2782 assert!(
2783 returned_err < 1e-12,
2784 "return adapter vs stream-into adapter disagree (n={n}, p={p}): {returned_err:e}"
2785 );
2786 }
2787 }
2788
2789 #[test]
2790 fn eigh_succeeds_on_same_structure_without_nan() {
2791 let mat = array![[1.0, 0.5, 0.1], [0.5, 2.0, 0.3], [0.1, 0.3, 1.5]];
2793 let (evals, _) = mat
2794 .eigh(Side::Lower)
2795 .expect("eigh should succeed on a well-conditioned finite matrix");
2796 assert!(
2797 evals.iter().all(|&v| v.is_finite()),
2798 "all eigenvalues should be finite"
2799 );
2800 }
2801
2802 #[test]
2811 fn gram_rrqr_flags_low_margin_on_exact_collinearity_so_caller_falls_back() {
2812 let n = 48usize;
2815 let x: Vec<f64> = (0..n)
2816 .map(|i| -1.0 + 2.0 * (i as f64) / (n as f64 - 1.0))
2817 .collect();
2818 let mut a = Array2::<f64>::zeros((n, 4));
2819 for i in 0..n {
2820 a[[i, 0]] = 1.0;
2821 a[[i, 1]] = x[i];
2822 a[[i, 2]] = x[i];
2823 a[[i, 3]] = x[i] * x[i];
2824 }
2825 let alpha = default_rrqr_rank_alpha();
2826
2827 let tall = rrqr_with_permutation(&a, alpha).expect("tall RRQR should succeed");
2830 assert_eq!(tall.rank, 3, "tall RRQR must demote the exact alias");
2831
2832 let unit = Array1::<f64>::ones(n);
2843 let gram = fast_xt_diag_x_with_parallelism(&a, &unit, faer::get_global_parallelism());
2844 let gram_rrqr =
2845 rrqr_from_gram_with_permutation(&gram, n, alpha).expect("Gram RRQR should succeed");
2846 let ok =
2847 gram_rrqr.rank == 3 || gram_rrqr.verdict_margin < JOINT_GRAM_RRQR_TRUST_MARGIN_FOR_TEST;
2848 assert!(
2849 ok,
2850 "gam#933: Gram RRQR must either find correct rank=3 OR signal low margin \
2851 (< {:.0e}) to force the tall fallback; got rank={} margin={:.3e}",
2852 JOINT_GRAM_RRQR_TRUST_MARGIN_FOR_TEST, gram_rrqr.rank, gram_rrqr.verdict_margin,
2853 );
2854 }
2855
2856 #[test]
2861 fn gram_rrqr_keeps_high_margin_on_full_rank_design() {
2862 let n = 200usize;
2863 let p = 5usize;
2864 let mut a = Array2::<f64>::zeros((n, p));
2865 for i in 0..n {
2867 let t = (i as f64) / (n as f64 - 1.0);
2868 a[[i, 0]] = 1.0;
2869 a[[i, 1]] = t;
2870 a[[i, 2]] = t * t;
2871 a[[i, 3]] = t * t * t;
2872 a[[i, 4]] = (t * 6.0).sin();
2873 }
2874 let alpha = default_rrqr_rank_alpha();
2875 let unit = Array1::<f64>::ones(n);
2876 let gram = fast_xt_diag_x_with_parallelism(&a, &unit, faer::get_global_parallelism());
2877 let gram_rrqr =
2878 rrqr_from_gram_with_permutation(&gram, n, alpha).expect("Gram RRQR should succeed");
2879 assert_eq!(gram_rrqr.rank, p, "full-rank design must keep all columns");
2880 assert!(
2881 gram_rrqr.verdict_margin >= JOINT_GRAM_RRQR_TRUST_MARGIN_FOR_TEST,
2882 "full-rank design must keep a high margin (fast Gram path); got {:.3e}",
2883 gram_rrqr.verdict_margin,
2884 );
2885 }
2886
2887 fn max_abs_diff(a: &Array2<f64>, b: &Array2<f64>) -> f64 {
2890 assert_eq!(a.dim(), b.dim(), "shape mismatch in max_abs_diff");
2891 a.iter()
2892 .zip(b.iter())
2893 .fold(0.0_f64, |acc, (&x, &y)| acc.max((x - y).abs()))
2894 }
2895
2896 fn max_abs_diff_1d(a: &Array1<f64>, b: &Array1<f64>) -> f64 {
2897 assert_eq!(a.len(), b.len(), "len mismatch in max_abs_diff_1d");
2898 a.iter()
2899 .zip(b.iter())
2900 .fold(0.0_f64, |acc, (&x, &y)| acc.max((x - y).abs()))
2901 }
2902
2903 #[test]
2905 fn fast_ab_small_matches_ndarray_dot() {
2906 let a = array![[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]];
2907 let b = array![[7.0, 8.0], [9.0, 10.0], [11.0, 12.0]];
2908 let got = fast_ab(&a, &b);
2909 let want = a.dot(&b);
2910 assert!(max_abs_diff(&got, &want) < 1e-12, "fast_ab small mismatch");
2911 assert_eq!(got.dim(), (2, 2));
2912 }
2913
2914 #[test]
2916 fn fast_ab_large_matches_ndarray_dot() {
2917 let n = 50usize;
2918 let p = 40usize;
2919 let q = 35usize;
2920 let mut a = Array2::<f64>::zeros((n, p));
2921 let mut b = Array2::<f64>::zeros((p, q));
2922 let mut state = 0xDEAD_BEEF_1234_5678u64;
2923 let next = |s: &mut u64| -> f64 {
2924 *s ^= *s << 13;
2925 *s ^= *s >> 7;
2926 *s ^= *s << 17;
2927 ((*s >> 11) as f64 / ((1u64 << 53) as f64)) - 0.5
2928 };
2929 for v in a.iter_mut() {
2930 *v = next(&mut state);
2931 }
2932 for v in b.iter_mut() {
2933 *v = next(&mut state);
2934 }
2935 let got = fast_ab(&a, &b);
2936 let want = a.dot(&b);
2937 assert!(max_abs_diff(&got, &want) < 1e-9, "fast_ab large mismatch");
2938 }
2939
2940 #[test]
2942 fn fast_atb_small_matches_ndarray_dot() {
2943 let a = array![[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]];
2944 let b = array![[7.0, 8.0, 9.0], [10.0, 11.0, 12.0], [13.0, 14.0, 15.0]];
2945 let got = fast_atb(&a, &b);
2946 let want = a.t().dot(&b);
2947 assert!(max_abs_diff(&got, &want) < 1e-12, "fast_atb small mismatch");
2948 assert_eq!(got.dim(), (2, 3));
2949 }
2950
2951 #[test]
2953 fn fast_atb_large_matches_ndarray_dot() {
2954 let n = 50usize;
2955 let p = 40usize;
2956 let q = 35usize;
2957 let mut a = Array2::<f64>::zeros((n, p));
2958 let mut b = Array2::<f64>::zeros((n, q));
2959 let mut state = 0xCAFE_BABE_9876_5432u64;
2960 let next = |s: &mut u64| -> f64 {
2961 *s ^= *s << 13;
2962 *s ^= *s >> 7;
2963 *s ^= *s << 17;
2964 ((*s >> 11) as f64 / ((1u64 << 53) as f64)) - 0.5
2965 };
2966 for v in a.iter_mut() {
2967 *v = next(&mut state);
2968 }
2969 for v in b.iter_mut() {
2970 *v = next(&mut state);
2971 }
2972 let got = fast_atb(&a, &b);
2973 let want = a.t().dot(&b);
2974 assert!(max_abs_diff(&got, &want) < 1e-9, "fast_atb large mismatch");
2975 }
2976
2977 #[test]
2979 fn fast_abt_small_matches_ndarray_dot() {
2980 let a = array![[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]];
2981 let b = array![[7.0, 8.0, 9.0], [10.0, 11.0, 12.0]];
2982 let got = fast_abt(&a, &b);
2983 let want = a.dot(&b.t());
2984 assert!(max_abs_diff(&got, &want) < 1e-12, "fast_abt small mismatch");
2985 assert_eq!(got.dim(), (2, 2));
2986 }
2987
2988 #[test]
2990 fn fast_av_small_matches_ndarray_dot() {
2991 let a = array![[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]];
2992 let v = array![1.0, -1.0, 2.0];
2993 let got = fast_av(&a, &v);
2994 let want = a.dot(&v);
2995 assert!(
2996 max_abs_diff_1d(&got, &want) < 1e-12,
2997 "fast_av small mismatch"
2998 );
2999 assert!((got[0] - 5.0).abs() < 1e-12, "fast_av[0] should be 5");
3001 assert!((got[1] - 11.0).abs() < 1e-12, "fast_av[1] should be 11");
3003 }
3004
3005 #[test]
3007 fn fast_av_large_matches_ndarray_dot() {
3008 let n = 50usize;
3009 let p = 40usize;
3010 let mut a = Array2::<f64>::zeros((n, p));
3011 let mut v = Array1::<f64>::zeros(p);
3012 let mut state = 0xFEED_FACE_ABCD_EF01u64;
3013 let next = |s: &mut u64| -> f64 {
3014 *s ^= *s << 13;
3015 *s ^= *s >> 7;
3016 *s ^= *s << 17;
3017 ((*s >> 11) as f64 / ((1u64 << 53) as f64)) - 0.5
3018 };
3019 for v in a.iter_mut() {
3020 *v = next(&mut state);
3021 }
3022 for x in v.iter_mut() {
3023 *x = next(&mut state);
3024 }
3025 let got = fast_av(&a, &v);
3026 let want = a.dot(&v);
3027 assert!(
3028 max_abs_diff_1d(&got, &want) < 1e-9,
3029 "fast_av large mismatch"
3030 );
3031 }
3032
3033 #[test]
3035 fn fast_atv_small_matches_ndarray_dot() {
3036 let a = array![[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]];
3037 let v = array![1.0, 0.0, -1.0];
3038 let got = fast_atv(&a, &v);
3039 let want = a.t().dot(&v);
3040 assert!(
3042 max_abs_diff_1d(&got, &want) < 1e-12,
3043 "fast_atv small mismatch"
3044 );
3045 assert!((got[0] - (-4.0)).abs() < 1e-12, "fast_atv[0]");
3046 assert!((got[1] - (-4.0)).abs() < 1e-12, "fast_atv[1]");
3047 }
3048
3049 #[test]
3051 fn fast_atv_large_matches_ndarray_dot() {
3052 let n = 50usize;
3053 let p = 40usize;
3054 let mut a = Array2::<f64>::zeros((n, p));
3055 let mut v = Array1::<f64>::zeros(n);
3056 let mut state = 0x1234_ABCD_5678_EF90u64;
3057 let next = |s: &mut u64| -> f64 {
3058 *s ^= *s << 13;
3059 *s ^= *s >> 7;
3060 *s ^= *s << 17;
3061 ((*s >> 11) as f64 / ((1u64 << 53) as f64)) - 0.5
3062 };
3063 for x in a.iter_mut() {
3064 *x = next(&mut state);
3065 }
3066 for x in v.iter_mut() {
3067 *x = next(&mut state);
3068 }
3069 let got = fast_atv(&a, &v);
3070 let want = a.t().dot(&v);
3071 assert!(
3072 max_abs_diff_1d(&got, &want) < 1e-9,
3073 "fast_atv large mismatch"
3074 );
3075 }
3076
3077 #[test]
3080 fn fast_xt_diag_y_small_matches_manual() {
3081 let x = array![[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]];
3082 let d = array![2.0, 0.5, 1.0];
3083 let y = array![[7.0, 8.0, 9.0], [10.0, 11.0, 12.0], [13.0, 14.0, 15.0]];
3084 let got = fast_xt_diag_y(&x, &d, &y);
3085 let diag_y = {
3087 let mut dy = Array2::<f64>::zeros(y.dim());
3088 for i in 0..3 {
3089 for j in 0..3 {
3090 dy[[i, j]] = d[i] * y[[i, j]];
3091 }
3092 }
3093 dy
3094 };
3095 let want = x.t().dot(&diag_y);
3096 assert!(
3097 max_abs_diff(&got, &want) < 1e-12,
3098 "fast_xt_diag_y small mismatch"
3099 );
3100 assert_eq!(got.dim(), (2, 3));
3101 }
3102
3103 #[inline]
3110 fn two_prod(a: f64, b: f64) -> (f64, f64) {
3111 let p = a * b;
3112 let e = a.mul_add(b, -p);
3113 (p, e)
3114 }
3115
3116 #[inline]
3117 fn two_sum(a: f64, b: f64) -> (f64, f64) {
3118 let s = a + b;
3119 let bb = s - a;
3120 let e = (a - (s - bb)) + (b - bb);
3121 (s, e)
3122 }
3123
3124 fn grow_expansion(e: &mut Vec<f64>, mut q: f64) {
3126 for h in e.iter_mut() {
3127 let (s, err) = two_sum(*h, q);
3128 *h = err;
3129 q = s;
3130 }
3131 if q != 0.0 {
3132 e.push(q);
3133 }
3134 }
3135
3136 fn exact_dot(a: &[f64], b: &[f64]) -> f64 {
3141 let mut e: Vec<f64> = Vec::new();
3142 for (&x, &y) in a.iter().zip(b.iter()) {
3143 let (p, ep) = two_prod(x, y);
3144 grow_expansion(&mut e, p);
3145 grow_expansion(&mut e, ep);
3146 }
3147 e.iter().fold(0.0f64, |acc, &c| acc + c)
3150 }
3151
3152 fn dd_dot(a: &[f64], b: &[f64]) -> f64 {
3156 let (mut s, mut c) = (0.0f64, 0.0f64);
3157 for (&x, &y) in a.iter().zip(b.iter()) {
3158 let (p, ep) = two_prod(x, y);
3159 let (s2, es) = two_sum(s, p);
3160 s = s2;
3161 c += ep + es;
3162 }
3163 s + c
3164 }
3165
3166 fn naive_dot(a: &[f64], b: &[f64]) -> f64 {
3167 let mut acc = 0.0f64;
3168 for (&x, &y) in a.iter().zip(b.iter()) {
3169 acc += x * y;
3170 }
3171 acc
3172 }
3173
3174 fn ill_conditioned_pair(len: usize, seed: u64) -> (Vec<f64>, Vec<f64>) {
3177 let mut s = seed | 1;
3178 let mut next = || {
3179 s ^= s << 13;
3180 s ^= s >> 7;
3181 s ^= s << 17;
3182 (s >> 11) as f64 / ((1u64 << 53) as f64) - 0.5
3183 };
3184 let mut a = Vec::with_capacity(len);
3185 let mut b = Vec::with_capacity(len);
3186 for i in 0..len {
3187 let scale = 10f64.powi((i % 17) as i32 - 8);
3189 let sign = if i % 2 == 0 { 1.0 } else { -1.0 };
3190 a.push(sign * next() * scale);
3191 b.push(next() * scale);
3192 }
3193 (a, b)
3194 }
3195
3196 #[test]
3199 fn fma_dot_beats_naive_accuracy() {
3200 let mut fma_total = 0.0f64;
3201 let mut naive_total = 0.0f64;
3202 let mut strict_wins = 0;
3203 for seed in 0..64u64 {
3204 let len = 200 + (seed as usize % 57);
3205 let (a, b) = ill_conditioned_pair(len, 0x9E37_79B9 ^ seed.wrapping_mul(2654435761));
3206 let truth = exact_dot(&a, &b);
3207 let fe = (super::fma_dot(&a, &b) - truth).abs();
3208 let ne = (naive_dot(&a, &b) - truth).abs();
3209 let floor = 8.0 * f64::EPSILON * truth.abs();
3213 assert!(
3214 fe <= ne * (1.0 + 1e-6) + floor,
3215 "fma_dot worse than naive: seed={seed} fma_err={fe:.3e} naive_err={ne:.3e}",
3216 );
3217 if fe < ne {
3218 strict_wins += 1;
3219 }
3220 fma_total += fe;
3221 naive_total += ne;
3222 }
3223 assert!(
3224 fma_total < naive_total,
3225 "fma_dot aggregate error {fma_total:.3e} not below naive {naive_total:.3e}",
3226 );
3227 assert!(
3228 strict_wins >= 40,
3229 "expected fma_dot to strictly win the majority; only {strict_wins}/64",
3230 );
3231 }
3232
3233 #[test]
3236 fn fast_atv_blocked_beats_naive_accuracy() {
3237 let n = 200_003usize;
3238 let p = 3usize;
3239 let mut s = 0xD1B5_4A32u64;
3240 let mut next = || {
3241 s ^= s << 13;
3242 s ^= s >> 7;
3243 s ^= s << 17;
3244 (s >> 11) as f64 / ((1u64 << 53) as f64) - 0.5
3245 };
3246 let mut x = Array2::<f64>::zeros((n, p));
3247 let mut v = Array1::<f64>::zeros(n);
3248 for i in 0..n {
3249 let scale = 10f64.powi((i % 17) as i32 - 8);
3250 v[i] = if i % 2 == 0 { scale } else { -scale } * next();
3251 for j in 0..p {
3252 x[[i, j]] = next() * scale;
3253 }
3254 }
3255 let got = fast_atv(&x, &v);
3256 for j in 0..p {
3258 let col: Vec<f64> = (0..n).map(|i| x[[i, j]]).collect();
3259 let vv: Vec<f64> = v.to_vec();
3260 let truth = dd_dot(&col, &vv);
3261 let naive = naive_dot(&col, &vv);
3262 let ge = (got[j] - truth).abs();
3263 let ne = (naive - truth).abs();
3264 assert!(
3265 ge <= ne + f64::MIN_POSITIVE,
3266 "col {j}: blocked err {ge:.3e} exceeds naive {ne:.3e}",
3267 );
3268 }
3269 }
3270
3271 #[test]
3274 fn fast_av_strided_input_matches_ndarray() {
3275 let mut base = Array2::<f64>::zeros((40, 60));
3276 let mut s = 0x0BAD_F00Du64;
3277 let mut next = || {
3278 s ^= s << 13;
3279 s ^= s >> 7;
3280 s ^= s << 17;
3281 (s >> 11) as f64 / ((1u64 << 53) as f64) - 0.5
3282 };
3283 for x in base.iter_mut() {
3284 *x = next();
3285 }
3286 let a = base.t();
3288 let mut v = Array1::<f64>::zeros(40);
3289 for x in v.iter_mut() {
3290 *x = next();
3291 }
3292 let got = fast_av(&a, &v);
3293 let want = a.dot(&v);
3294 assert!(
3295 max_abs_diff_1d(&got, &want) < 1e-11,
3296 "strided fast_av mismatch (fallback path)",
3297 );
3298 }
3299
3300 #[test]
3312 fn faer_sequential_scope_sets_seq_inside_and_restores_after() {
3313 let baseline = faer::get_global_parallelism();
3314 faer::set_global_parallelism(Par::rayon(4));
3317 assert_eq!(
3318 faer::get_global_parallelism(),
3319 Par::rayon(4),
3320 "baseline must be the parallel policy we just set",
3321 );
3322
3323 {
3324 let faer_seq_guard = FaerSequentialScope::enter();
3325 assert_eq!(
3326 faer::get_global_parallelism(),
3327 Par::Seq,
3328 "faer must be pinned to Par::Seq inside the scope",
3329 );
3330
3331 {
3333 let faer_seq_inner_guard = FaerSequentialScope::enter();
3334 assert_eq!(
3335 faer::get_global_parallelism(),
3336 Par::Seq,
3337 "nested scope stays Par::Seq",
3338 );
3339 drop(faer_seq_inner_guard);
3340 }
3341 assert_eq!(
3342 faer::get_global_parallelism(),
3343 Par::Seq,
3344 "inner drop must not restore while outer scope is still live",
3345 );
3346 drop(faer_seq_guard);
3347 }
3348
3349 assert_eq!(
3350 faer::get_global_parallelism(),
3351 Par::rayon(4),
3352 "outermost drop must restore the pre-scope parallelism policy",
3353 );
3354
3355 let observed = with_faer_sequential(|| faer::get_global_parallelism());
3357 assert_eq!(
3358 observed,
3359 Par::Seq,
3360 "with_faer_sequential runs body under Seq"
3361 );
3362 assert_eq!(
3363 faer::get_global_parallelism(),
3364 Par::rayon(4),
3365 "with_faer_sequential restores after the body returns",
3366 );
3367
3368 faer::set_global_parallelism(baseline);
3370 }
3371}