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("Self-adjoint eigendecomposition failed: {0:?}")]
172 SelfAdjointEigen(solvers::EvdError),
173 #[error("Cholesky factorization failed: {0:?}")]
174 Cholesky(solvers::LltError),
175 #[error("LDLT factorization failed: {0:?}")]
176 Ldlt(solvers::LdltError),
177}
178
179pub enum FaerSymmetricFactor {
180 Llt(FaerLlt<f64>),
181 Ldlt(FaerLdlt<f64>),
182 Lblt(FaerLblt<f64>),
183}
184
185#[inline]
186pub fn cholesky_factor_logdet(factor: MatRef<'_, f64>) -> f64 {
187 2.0 * diagonal_log_sum(factor.diagonal())
188}
189
190#[inline]
191fn diagonal_log_sum(diagonal: DiagRef<'_, f64>) -> f64 {
192 diagonal
193 .column_vector()
194 .iter()
195 .map(|&x| x.ln())
196 .sum::<f64>()
197}
198
199impl FaerSymmetricFactor {
200 #[inline]
202 pub fn n(&self) -> usize {
203 use faer::linalg::solvers::ShapeCore;
204 match self {
205 FaerSymmetricFactor::Llt(f) => f.nrows(),
206 FaerSymmetricFactor::Ldlt(f) => f.nrows(),
207 FaerSymmetricFactor::Lblt(f) => f.nrows(),
208 }
209 }
210
211 #[inline]
212 pub fn solve(&self, rhs: MatRef<'_, f64>) -> Mat<f64> {
213 match self {
214 FaerSymmetricFactor::Llt(f) => f.solve(rhs),
215 FaerSymmetricFactor::Ldlt(f) => f.solve(rhs),
216 FaerSymmetricFactor::Lblt(f) => f.solve(rhs),
217 }
218 }
219
220 #[inline]
221 pub fn solve_in_place(&self, rhs: MatMut<'_, f64>) {
222 match self {
223 FaerSymmetricFactor::Llt(f) => f.solve_in_place(rhs),
224 FaerSymmetricFactor::Ldlt(f) => f.solve_in_place(rhs),
225 FaerSymmetricFactor::Lblt(f) => f.solve_in_place(rhs),
226 }
227 }
228}
229
230impl crate::matrix::FactorizedSystem for FaerSymmetricFactor {
231 fn solve(&self, rhs: &Array1<f64>) -> Result<Array1<f64>, String> {
232 let mut out = rhs.clone();
233 let mut out_mat = array1_to_col_matmut(&mut out);
234 self.solve_in_place(out_mat.as_mut());
235 if !out.iter().all(|v| v.is_finite()) {
236 return Err("symmetric factor solve produced non-finite values".to_string());
237 }
238 Ok(out)
239 }
240
241 fn solvemulti(&self, rhs: &Array2<f64>) -> Result<Array2<f64>, String> {
242 let mut out = Array2::<f64>::zeros(rhs.raw_dim());
243 for j in 0..rhs.ncols() {
244 for i in 0..rhs.nrows() {
245 out[[i, j]] = rhs[[i, j]];
246 }
247 }
248 let mut out_mat = array2_to_matmut(&mut out);
249 self.solve_in_place(out_mat.as_mut());
250 if !out.iter().all(|v| v.is_finite()) {
251 return Err("symmetric factor multi-solve produced non-finite values".to_string());
252 }
253 Ok(out)
254 }
255
256 fn logdet(&self) -> f64 {
257 match self {
258 FaerSymmetricFactor::Llt(f) => cholesky_factor_logdet(f.L()),
259 FaerSymmetricFactor::Ldlt(f) => diagonal_log_sum(f.D()),
260 FaerSymmetricFactor::Lblt(..) => {
261 f64::NAN
265 }
266 }
267 }
268}
269
270#[inline]
272pub fn factorize_symmetricwith_fallback(
273 matrix: MatRef<'_, f64>,
274 side: Side,
275) -> Result<FaerSymmetricFactor, FaerLinalgError> {
276 if let Ok(llt) = FaerLlt::new(matrix, side) {
277 return Ok(FaerSymmetricFactor::Llt(llt));
278 }
279 let ldlt_err = match FaerLdlt::new(matrix, side) {
280 Ok(ldlt) => return Ok(FaerSymmetricFactor::Ldlt(ldlt)),
281 Err(err) => err,
282 };
283 let lblt = catch_unwind(AssertUnwindSafe(|| FaerLblt::new(matrix, side)))
284 .map_err(|_| FaerLinalgError::Ldlt(ldlt_err))?;
285 Ok(FaerSymmetricFactor::Lblt(lblt))
286}
287
288#[inline]
289const fn should_use_faer_matmul(m: usize, n: usize, k: usize) -> bool {
290 const MIN_DIM: usize = 32;
294 const MIN_FLOP_SCALE: usize = 64 * 64;
295 (m >= MIN_DIM || n >= MIN_DIM || k >= MIN_DIM)
296 && m.saturating_mul(n).saturating_mul(k) >= MIN_FLOP_SCALE
297}
298
299#[inline]
300pub fn matmul_parallelism(m: usize, n: usize, k: usize) -> Par {
301 const PAR_MIN_FLOP_SCALE: usize = 2_000_000;
305 const PAR_MIN_LONG_DIM: usize = 256;
306 let flop_scale = m.saturating_mul(n).saturating_mul(k);
307 let long_dim = m.max(n).max(k);
308 if flop_scale >= PAR_MIN_FLOP_SCALE && long_dim >= PAR_MIN_LONG_DIM {
309 effective_global_parallelism()
313 } else {
314 Par::Seq
315 }
316}
317
318#[inline]
319pub fn array2_to_matmut(array: &mut Array2<f64>) -> MatMut<'_, f64> {
320 let (rows, cols) = array.dim();
321 let strides = array.strides();
322
323 let s0 = strides[0];
330 let s1 = strides[1];
331
332 unsafe { MatMut::from_raw_parts_mut(array.as_mut_ptr(), rows, cols, s0, s1) }
336}
337
338pub fn array2_to_nested_vec(array: &Array2<f64>) -> Vec<Vec<f64>> {
341 array.rows().into_iter().map(|row| row.to_vec()).collect()
342}
343
344#[inline]
345pub fn array1_to_col_matmut(array: &mut Array1<f64>) -> MatMut<'_, f64> {
346 let len = array.len();
347 let stride = array.strides()[0];
348 unsafe {
352 MatMut::from_raw_parts_mut(
353 array.as_mut_ptr(),
354 len,
355 1,
356 stride,
357 0, )
359 }
360}
361
362#[inline]
369pub fn fast_ata<S: Data<Elem = f64>>(a: &ArrayBase<S, Ix2>) -> Array2<f64> {
370 let p = a.ncols();
371 let mut out = Array2::<f64>::zeros((p, p));
372 fast_ata_into(a, &mut out);
373 out
374}
375
376#[inline]
379pub fn fast_ata_into<S: Data<Elem = f64>>(a: &ArrayBase<S, Ix2>, out: &mut Array2<f64>) {
380 use faer::Accum;
381 use faer::linalg::matmul::triangular::{BlockStructure, matmul as tri_matmul};
382
383 let (n, p) = a.dim();
384 assert_eq!(out.nrows(), p, "output rows must match p");
385 assert_eq!(out.ncols(), p, "output cols must match p");
386
387 if !should_use_faer_matmul(p, p, n) {
388 out.assign(&a.t().dot(a));
389 return;
390 }
391
392 let mut outview = array2_to_matmut(out);
393
394 let aview = FaerArrayView::new(a);
395 let a_ref = aview.as_ref();
396 let a_t = a_ref.transpose();
397 let par = matmul_parallelism(p, p, n);
398 tri_matmul(
399 outview.as_mut(),
400 BlockStructure::TriangularLower,
401 Accum::Replace,
402 a_t,
403 BlockStructure::Rectangular,
404 a_ref,
405 BlockStructure::Rectangular,
406 1.0,
407 par,
408 );
409 for i in 0..p {
411 for j in (i + 1)..p {
412 out[[i, j]] = out[[j, i]];
413 }
414 }
415}
416
417#[inline]
421pub fn fast_atb<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
422 a: &ArrayBase<S1, Ix2>,
423 b: &ArrayBase<S2, Ix2>,
424) -> Array2<f64> {
425 if let Some(out) =
426 crate::gpu_hook::gpu_dispatch().and_then(|d| d.try_fast_atb(a.view(), b.view()))
427 {
428 return out;
429 }
430 let (n_a, p) = a.dim();
431 let q = b.ncols();
432 fast_atb_with_parallelism(a, b, matmul_parallelism(p, q, n_a))
433}
434
435#[inline]
438pub fn fast_atb_with_parallelism<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
439 a: &ArrayBase<S1, Ix2>,
440 b: &ArrayBase<S2, Ix2>,
441 par: Par,
442) -> Array2<f64> {
443 use faer::linalg::matmul::matmul;
444 use faer::{Accum, Mat};
445
446 let (n_a, p) = a.dim();
447 let (n_b, q) = b.dim();
448 assert_eq!(n_a, n_b, "A and B must have same number of rows");
449
450 if !should_use_faer_matmul(p, q, n_a) {
452 return a.t().dot(b);
453 }
454
455 let mut result = Mat::<f64>::zeros(p, q);
456
457 let aview = FaerArrayView::new(a);
458 let bview = FaerArrayView::new(b);
459 let a_ref = aview.as_ref();
460 let b_ref = bview.as_ref();
461
462 matmul(
464 result.as_mut(),
465 Accum::Replace,
466 a_ref.transpose(),
467 b_ref,
468 1.0,
469 par,
470 );
471
472 mat_to_array(result.as_ref())
473}
474
475#[inline]
478pub fn fast_abt<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
479 a: &ArrayBase<S1, Ix2>,
480 b: &ArrayBase<S2, Ix2>,
481) -> Array2<f64> {
482 use faer::linalg::matmul::matmul;
483 use faer::{Accum, Mat};
484
485 let (m, k_a) = a.dim();
486 let (n, k_b) = b.dim();
487 assert_eq!(
488 k_a, k_b,
489 "A and B must have same number of columns for A·Bᵀ"
490 );
491
492 if !should_use_faer_matmul(m, n, k_a) {
493 return a.dot(&b.t());
494 }
495
496 let mut result = Mat::<f64>::zeros(m, n);
497 let aview = FaerArrayView::new(a);
498 let bview = FaerArrayView::new(b);
499 let par = matmul_parallelism(m, n, k_a);
500 matmul(
501 result.as_mut(),
502 Accum::Replace,
503 aview.as_ref(),
504 bview.as_ref().transpose(),
505 1.0,
506 par,
507 );
508 mat_to_array(result.as_ref())
509}
510
511#[inline]
515pub fn fast_ab<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
516 a: &ArrayBase<S1, Ix2>,
517 b: &ArrayBase<S2, Ix2>,
518) -> Array2<f64> {
519 if let Some(out) =
520 crate::gpu_hook::gpu_dispatch().and_then(|d| d.try_fast_ab(a.view(), b.view()))
521 {
522 return out;
523 }
524 let n = a.nrows();
525 let q = b.ncols();
526 let mut out = Array2::<f64>::zeros((n, q));
527 fast_ab_into(a, b, &mut out);
528 out
529}
530
531const FMA_LANES: usize = 8;
556
557const KERNEL_PAR_MIN_FLOP: usize = 1 << 18; const AV_PAR_CHUNK_ROWS: usize = 1024;
565
566const ATV_BLOCK_ROWS: usize = 512;
571
572#[inline]
573fn kernel_should_parallelize(n: usize, p: usize) -> bool {
574 !in_nested_parallel_region()
575 && n.saturating_mul(p) >= KERNEL_PAR_MIN_FLOP
576 && rayon::current_num_threads() > 1
577}
578
579#[inline(always)]
594fn fma_dot(a: &[f64], b: &[f64]) -> f64 {
595 assert_eq!(a.len(), b.len(), "fma_dot: operand length mismatch");
596 let mut sum = [0.0f64; FMA_LANES];
597 let mut comp = [0.0f64; FMA_LANES];
598 let mut ca = a.chunks_exact(FMA_LANES);
599 let mut cb = b.chunks_exact(FMA_LANES);
600 for (xa, xb) in ca.by_ref().zip(cb.by_ref()) {
601 for l in 0..FMA_LANES {
602 let x = xa[l];
603 let y = xb[l];
604 let p = x * y;
606 let ep = x.mul_add(y, -p);
607 let s = sum[l] + p;
609 let bb = s - sum[l];
610 let es = (sum[l] - (s - bb)) + (p - bb);
611 sum[l] = s;
612 comp[l] += ep + es;
613 }
614 }
615 let mut sr = 0.0f64;
617 let mut cr = 0.0f64;
618 for (&x, &y) in ca.remainder().iter().zip(cb.remainder().iter()) {
619 let p = x * y;
620 let ep = x.mul_add(y, -p);
621 let s = sr + p;
622 let bb = s - sr;
623 let es = (sr - (s - bb)) + (p - bb);
624 sr = s;
625 cr += ep + es;
626 }
627 let mut total = sr + cr;
629 for l in 0..FMA_LANES {
630 total += sum[l] + comp[l];
631 }
632 total
633}
634
635fn fast_av_rowmajor_into(x_all: &[f64], v: &[f64], n: usize, p: usize, out: &mut [f64]) {
639 assert_eq!(x_all.len(), n * p, "fast_av_rowmajor_into: x_all length");
640 assert_eq!(v.len(), p, "fast_av_rowmajor_into: v length");
641 assert_eq!(out.len(), n, "fast_av_rowmajor_into: out length");
642 if kernel_should_parallelize(n, p) {
643 use rayon::prelude::*;
644 out.par_chunks_mut(AV_PAR_CHUNK_ROWS)
645 .enumerate()
646 .for_each(|(c, chunk)| {
647 let base = c * AV_PAR_CHUNK_ROWS;
648 for (k, o) in chunk.iter_mut().enumerate() {
649 let i = base + k;
650 *o = fma_dot(&x_all[i * p..i * p + p], v);
651 }
652 });
653 } else {
654 for (i, o) in out.iter_mut().enumerate() {
655 *o = fma_dot(&x_all[i * p..i * p + p], v);
656 }
657 }
658}
659
660fn pairwise_sum_into(parts: &[Vec<f64>], out: &mut [f64]) {
662 match parts.len() {
663 0 => out.fill(0.0),
664 1 => out.copy_from_slice(&parts[0]),
665 _ => {
666 let mid = parts.len() / 2;
667 let p = out.len();
668 let mut left = vec![0.0f64; p];
669 let mut right = vec![0.0f64; p];
670 pairwise_sum_into(&parts[..mid], &mut left);
671 pairwise_sum_into(&parts[mid..], &mut right);
672 for ((o, &l), &r) in out.iter_mut().zip(left.iter()).zip(right.iter()) {
673 *o = l + r;
674 }
675 }
676 }
677}
678
679fn fast_atv_rowmajor_into(x_all: &[f64], v: &[f64], n: usize, p: usize, out: &mut [f64]) {
687 assert_eq!(x_all.len(), n * p, "fast_atv_rowmajor_into: x_all length");
688 assert_eq!(v.len(), n, "fast_atv_rowmajor_into: v length");
689 assert_eq!(out.len(), p, "fast_atv_rowmajor_into: out length");
690 let nblocks = n.div_ceil(ATV_BLOCK_ROWS);
691
692 let block_partial = |b: usize| -> Vec<f64> {
693 let start = b * ATV_BLOCK_ROWS;
694 let end = (start + ATV_BLOCK_ROWS).min(n);
695 let mut acc = vec![0.0f64; p];
696 for i in start..end {
697 let vi = v[i];
698 let row = &x_all[i * p..i * p + p];
699 for (a, &xij) in acc.iter_mut().zip(row.iter()) {
700 *a = xij.mul_add(vi, *a);
701 }
702 }
703 acc
704 };
705
706 let partials: Vec<Vec<f64>> = if kernel_should_parallelize(n, p) {
707 use rayon::prelude::*;
708 (0..nblocks).into_par_iter().map(block_partial).collect()
709 } else {
710 (0..nblocks).map(block_partial).collect()
711 };
712
713 pairwise_sum_into(&partials, out);
714}
715
716#[inline]
719pub fn fast_av<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
720 a: &ArrayBase<S1, Ix2>,
721 v: &ArrayBase<S2, Ix1>,
722) -> Array1<f64> {
723 if let Some(out) =
724 crate::gpu_hook::gpu_dispatch().and_then(|d| d.try_fast_av(a.view(), v.view()))
725 {
726 return out;
727 }
728 fast_av_impl(a, v)
729}
730
731#[inline]
732fn fast_av_impl<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
733 a: &ArrayBase<S1, Ix2>,
734 v: &ArrayBase<S2, Ix1>,
735) -> Array1<f64> {
736 use faer::linalg::matmul::matmul;
737 use faer::{Accum, Mat};
738
739 let (n, p) = a.dim();
740 assert_eq!(p, v.len(), "A cols must match v length");
741
742 if let (Some(x_all), Some(vs)) = (a.as_slice(), v.as_slice())
746 && n != 0
747 && p != 0
748 {
749 let mut out = Array1::<f64>::zeros(n);
750 fast_av_rowmajor_into(
751 x_all,
752 vs,
753 n,
754 p,
755 out.as_slice_mut().expect("fresh Array1 is contiguous"),
756 );
757 return out;
758 }
759
760 if !should_use_faer_matmul(n, 1, p) {
761 return a.dot(v);
762 }
763
764 let mut result = Mat::<f64>::zeros(n, 1);
765
766 let aview = FaerArrayView::new(a);
767 let vview = FaerColView::new(v);
768 let a_ref = aview.as_ref();
769 let v_ref = vview.as_ref();
770
771 let par = matmul_parallelism(n, 1, p);
772 matmul(result.as_mut(), Accum::Replace, a_ref, v_ref, 1.0, par);
773
774 let mut out = Array1::<f64>::zeros(n);
775 for i in 0..n {
776 out[i] = result[(i, 0)];
777 }
778 out
779}
780
781#[inline]
784pub fn fast_av_into<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
785 a: &ArrayBase<S1, Ix2>,
786 v: &ArrayBase<S2, Ix1>,
787 out: &mut Array1<f64>,
788) {
789 fast_av_into_impl(a, v, out);
790}
791
792#[inline]
793fn fast_av_into_impl<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
794 a: &ArrayBase<S1, Ix2>,
795 v: &ArrayBase<S2, Ix1>,
796 out: &mut Array1<f64>,
797) {
798 use faer::Accum;
799 use faer::linalg::matmul::matmul;
800
801 let (n, p) = a.dim();
802 assert_eq!(v.len(), p, "vector length must match A cols");
803 assert_eq!(out.len(), n, "output length must match A rows");
804
805 if let (Some(x_all), Some(vs)) = (a.as_slice(), v.as_slice())
806 && n != 0
807 && p != 0
808 && let Some(out_s) = out.as_slice_mut()
809 {
810 fast_av_rowmajor_into(x_all, vs, n, p, out_s);
811 return;
812 }
813
814 if !should_use_faer_matmul(n, 1, p) {
815 out.assign(&a.dot(v));
816 return;
817 }
818
819 let mut outview = array1_to_col_matmut(out);
820
821 let aview = FaerArrayView::new(a);
822 let vview = FaerColView::new(v);
823 let a_ref = aview.as_ref();
824 let v_ref = vview.as_ref();
825 let par = matmul_parallelism(n, 1, p);
826 matmul(outview.as_mut(), Accum::Replace, a_ref, v_ref, 1.0, par);
827}
828
829#[inline]
836pub fn fast_av_view_into<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
837 a: &ArrayBase<S1, Ix2>,
838 v: &ArrayBase<S2, Ix1>,
839 out: ArrayViewMut1<'_, f64>,
840) {
841 fast_av_view_into_impl(a, v, out);
842}
843
844#[inline]
845fn fast_av_view_into_impl<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
846 a: &ArrayBase<S1, Ix2>,
847 v: &ArrayBase<S2, Ix1>,
848 mut out: ArrayViewMut1<'_, f64>,
849) {
850 use faer::Accum;
851 use faer::linalg::matmul::matmul;
852
853 let (n, p) = a.dim();
854 assert_eq!(v.len(), p, "vector length must match A cols");
855 assert_eq!(out.len(), n, "output length must match A rows");
856
857 if let (Some(x_all), Some(vs)) = (a.as_slice(), v.as_slice())
858 && n != 0
859 && p != 0
860 && let Some(out_s) = out.as_slice_mut()
861 {
862 fast_av_rowmajor_into(x_all, vs, n, p, out_s);
863 return;
864 }
865
866 if !should_use_faer_matmul(n, 1, p) {
867 let prod = a.dot(v);
868 out.assign(&prod);
869 return;
870 }
871
872 let len = out.len();
873 let stride = out.strides()[0];
874 let outview = unsafe {
878 MatMut::from_raw_parts_mut(
879 out.as_mut_ptr(),
880 len,
881 1,
882 stride,
883 0, )
885 };
886
887 let aview = FaerArrayView::new(a);
888 let vview = FaerColView::new(v);
889 let a_ref = aview.as_ref();
890 let v_ref = vview.as_ref();
891 let par = matmul_parallelism(n, 1, p);
892 matmul(outview, Accum::Replace, a_ref, v_ref, 1.0, par);
893}
894
895#[inline]
898pub fn fast_atv<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
899 a: &ArrayBase<S1, Ix2>,
900 v: &ArrayBase<S2, Ix1>,
901) -> Array1<f64> {
902 if let Some(out) =
903 crate::gpu_hook::gpu_dispatch().and_then(|d| d.try_fast_atv(a.view(), v.view()))
904 {
905 return out;
906 }
907 fast_atv_impl(a, v)
908}
909
910#[inline]
911fn fast_atv_impl<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
912 a: &ArrayBase<S1, Ix2>,
913 v: &ArrayBase<S2, Ix1>,
914) -> Array1<f64> {
915 use faer::Accum;
916 use faer::linalg::matmul::matmul;
917
918 let (n, p) = a.dim();
919 assert_eq!(n, v.len(), "A rows must match v length");
920
921 if let (Some(x_all), Some(vs)) = (a.as_slice(), v.as_slice())
925 && n != 0
926 && p != 0
927 {
928 let mut out = Array1::<f64>::zeros(p);
929 fast_atv_rowmajor_into(
930 x_all,
931 vs,
932 n,
933 p,
934 out.as_slice_mut().expect("fresh Array1 is contiguous"),
935 );
936 return out;
937 }
938
939 if !should_use_faer_matmul(p, 1, n) {
941 return a.t().dot(v);
942 }
943
944 let mut out = Array1::<f64>::zeros(p);
945 let mut outview = array1_to_col_matmut(&mut out);
946
947 let aview = FaerArrayView::new(a);
948 let vview = FaerColView::new(v);
949 let a_ref = aview.as_ref();
950 let v_ref = vview.as_ref();
951
952 let par = matmul_parallelism(p, 1, n);
954 matmul(
955 outview.as_mut(),
956 Accum::Replace,
957 a_ref.transpose(),
958 v_ref,
959 1.0,
960 par,
961 );
962
963 out
964}
965
966#[inline]
969pub fn fast_atv_into<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
970 a: &ArrayBase<S1, Ix2>,
971 v: &ArrayBase<S2, Ix1>,
972 out: &mut Array1<f64>,
973) {
974 fast_atv_into_impl(a, v, out);
975}
976
977#[inline]
978fn fast_atv_into_impl<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
979 a: &ArrayBase<S1, Ix2>,
980 v: &ArrayBase<S2, Ix1>,
981 out: &mut Array1<f64>,
982) {
983 use faer::Accum;
984 use faer::linalg::matmul::matmul;
985
986 let (n, p) = a.dim();
987 assert_eq!(v.len(), n, "vector length must match A rows");
988 assert_eq!(out.len(), p, "output length must match A cols");
989
990 if let (Some(x_all), Some(vs)) = (a.as_slice(), v.as_slice())
991 && n != 0
992 && p != 0
993 && let Some(out_s) = out.as_slice_mut()
994 {
995 fast_atv_rowmajor_into(x_all, vs, n, p, out_s);
996 return;
997 }
998
999 if !should_use_faer_matmul(p, 1, n) {
1000 out.assign(&a.t().dot(v));
1001 return;
1002 }
1003
1004 let mut outview = array1_to_col_matmut(out);
1005
1006 let aview = FaerArrayView::new(a);
1007 let vview = FaerColView::new(v);
1008 let a_ref = aview.as_ref();
1009 let v_ref = vview.as_ref();
1010 let par = matmul_parallelism(p, 1, n);
1011 matmul(
1012 outview.as_mut(),
1013 Accum::Replace,
1014 a_ref.transpose(),
1015 v_ref,
1016 1.0,
1017 par,
1018 );
1019}
1020
1021#[inline]
1023pub fn fast_xt_diag_x<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
1024 x: &ArrayBase<S1, Ix2>,
1025 w: &ArrayBase<S2, Ix1>,
1026) -> Array2<f64> {
1027 assert_eq!(
1028 x.nrows(),
1029 w.len(),
1030 "fast_xt_diag_x row/weight length mismatch"
1031 );
1032 if let Some(out) =
1033 crate::gpu_hook::gpu_dispatch().and_then(|d| d.try_fast_xt_diag_x(x.view(), w.view()))
1034 {
1035 return out;
1036 }
1037 let p = x.ncols();
1038 fast_xt_diag_x_with_parallelism(x, w, matmul_parallelism(p, p, x.nrows()))
1039}
1040
1041#[inline]
1044pub fn fast_xt_diag_x_with_parallelism<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
1045 x: &ArrayBase<S1, Ix2>,
1046 w: &ArrayBase<S2, Ix1>,
1047 par: Par,
1048) -> Array2<f64> {
1049 assert_eq!(
1050 x.nrows(),
1051 w.len(),
1052 "fast_xt_diag_x_with_parallelism row/weight length mismatch"
1053 );
1054 fast_xt_diag_x_with_parallelism_impl(x, w, par)
1055}
1056
1057#[inline]
1058fn fast_xt_diag_x_with_parallelism_impl<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
1059 x: &ArrayBase<S1, Ix2>,
1060 w: &ArrayBase<S2, Ix1>,
1061 par: Par,
1062) -> Array2<f64> {
1063 use ndarray::ShapeBuilder;
1064
1065 let p = x.ncols();
1066 let mut result = Array2::<f64>::zeros((p, p).f());
1069 stream_weighted_crossprod_into(
1070 x,
1071 w,
1072 &mut result,
1073 CrossprodStructure::SymmetricLower,
1074 CrossprodAccum::Replace,
1075 par,
1076 );
1077 result
1078}
1079
1080#[derive(Clone, Copy, PartialEq, Eq, Debug)]
1082pub enum CrossprodStructure {
1083 Full,
1085 SymmetricLower,
1089}
1090
1091#[derive(Clone, Copy, PartialEq, Eq, Debug)]
1093pub enum CrossprodAccum {
1094 Replace,
1096 Add,
1098}
1099
1100pub fn stream_weighted_crossprod_into<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
1119 x: &ArrayBase<S1, Ix2>,
1120 w: &ArrayBase<S2, Ix1>,
1121 out: &mut Array2<f64>,
1122 structure: CrossprodStructure,
1123 accum: CrossprodAccum,
1124 par: Par,
1125) {
1126 use faer::Accum;
1127 use faer::linalg::matmul::matmul;
1128 use faer::linalg::matmul::triangular::{BlockStructure, matmul as tri_matmul};
1129 use ndarray::s;
1130
1131 let (n, p) = x.dim();
1132 assert_eq!(n, w.len(), "X rows must match W length");
1133 assert_eq!(out.nrows(), p, "output rows must match X cols");
1134 assert_eq!(out.ncols(), p, "output cols must match X cols");
1135 if p == 0 {
1136 return;
1137 }
1138 if n == 0 {
1139 if accum == CrossprodAccum::Replace {
1140 out.fill(0.0);
1141 }
1142 return;
1143 }
1144
1145 if !should_use_faer_matmul(p, p, n) {
1146 let w_x = Array2::from_shape_fn((n, p), |(i, j)| w[i] * x[[i, j]]);
1148 let gram = x.t().dot(&w_x);
1149 match accum {
1150 CrossprodAccum::Replace => out.assign(&gram),
1151 CrossprodAccum::Add => *out += &gram,
1152 }
1153 return;
1154 }
1155
1156 const TARGET_BYTES: usize = 8 * 1024 * 1024;
1158 const MIN_ROWS: usize = 512;
1159 const MAX_ROWS: usize = 131_072;
1160 let chunk_rows = (TARGET_BYTES / (p.max(1) * 8))
1161 .clamp(MIN_ROWS, MAX_ROWS)
1162 .min(n);
1163
1164 if accum == CrossprodAccum::Replace {
1169 out.fill(0.0);
1170 }
1171
1172 let mut wx_chunk = Array2::<f64>::zeros((chunk_rows, p));
1178
1179 let x_is_row_major = x.is_standard_layout();
1180 let w_slice_opt = w.as_slice();
1181
1182 {
1185 let mut out_view = array2_to_matmut(out);
1186 for start in (0..n).step_by(chunk_rows) {
1187 let rows = (n - start).min(chunk_rows);
1188 {
1189 let chunk_slice = wx_chunk
1190 .as_slice_mut()
1191 .expect("row-major chunk is contiguous");
1192 if x_is_row_major && let (Some(x_all), Some(w_all)) = (x.as_slice(), w_slice_opt) {
1193 for local in 0..rows {
1194 let src = start + local;
1195 let wi = w_all[src];
1196 let src_off = src * p;
1197 let dst_off = local * p;
1198 let src_row = &x_all[src_off..src_off + p];
1199 let dst_row = &mut chunk_slice[dst_off..dst_off + p];
1200 for col in 0..p {
1201 dst_row[col] = src_row[col] * wi;
1202 }
1203 }
1204 } else {
1205 let x_slice = x.slice(s![start..start + rows, ..]);
1206 for local in 0..rows {
1207 let wi = w[start + local];
1208 let xrow = x_slice.row(local);
1209 let dst_off = local * p;
1210 let dst_row = &mut chunk_slice[dst_off..dst_off + p];
1211 for (col, xij) in xrow.iter().enumerate() {
1212 dst_row[col] = xij * wi;
1213 }
1214 }
1215 }
1216 }
1217 let x_slice = x.slice(s![start..start + rows, ..]);
1218 let wx_slice = wx_chunk.slice(s![0..rows, ..]);
1219 let x_view = FaerArrayView::new(&x_slice);
1220 let wx_view = FaerArrayView::new(&wx_slice);
1221 match structure {
1222 CrossprodStructure::SymmetricLower => {
1223 tri_matmul(
1227 out_view.as_mut(),
1228 BlockStructure::TriangularLower,
1229 Accum::Add,
1230 x_view.as_ref().transpose(),
1231 BlockStructure::Rectangular,
1232 wx_view.as_ref(),
1233 BlockStructure::Rectangular,
1234 1.0,
1235 par,
1236 );
1237 }
1238 CrossprodStructure::Full => {
1239 matmul(
1240 out_view.as_mut(),
1241 Accum::Add,
1242 x_view.as_ref().transpose(),
1243 wx_view.as_ref(),
1244 1.0,
1245 par,
1246 );
1247 }
1248 }
1249 }
1250 }
1251
1252 if structure == CrossprodStructure::SymmetricLower {
1253 for i in 0..p {
1255 for j in (i + 1)..p {
1256 out[[i, j]] = out[[j, i]];
1257 }
1258 }
1259 }
1260}
1261
1262#[inline]
1264pub fn fast_xt_diag_y<S1: Data<Elem = f64>, S2: Data<Elem = f64>, S3: Data<Elem = f64>>(
1265 x: &ArrayBase<S1, Ix2>,
1266 w: &ArrayBase<S2, Ix1>,
1267 y: &ArrayBase<S3, Ix2>,
1268) -> Array2<f64> {
1269 assert_eq!(x.nrows(), y.nrows(), "fast_xt_diag_y X/Y row mismatch");
1270 assert_eq!(
1271 y.nrows(),
1272 w.len(),
1273 "fast_xt_diag_y row/weight length mismatch"
1274 );
1275 if let Some(out) = crate::gpu_hook::gpu_dispatch()
1276 .and_then(|d| d.try_fast_xt_diag_y(x.view(), w.view(), y.view()))
1277 {
1278 return out;
1279 }
1280 fast_xt_diag_y_impl(x, w, y)
1281}
1282
1283#[inline]
1284fn fast_xt_diag_y_impl<S1: Data<Elem = f64>, S2: Data<Elem = f64>, S3: Data<Elem = f64>>(
1285 x: &ArrayBase<S1, Ix2>,
1286 w: &ArrayBase<S2, Ix1>,
1287 y: &ArrayBase<S3, Ix2>,
1288) -> Array2<f64> {
1289 use faer::Accum;
1290 use faer::linalg::matmul::matmul;
1291 use ndarray::{ShapeBuilder, s};
1292
1293 let (n, q) = y.dim();
1294 let px = x.ncols();
1295 assert_eq!(n, w.len(), "Y rows must match W length");
1296 assert_eq!(n, x.nrows(), "X rows must match Y rows");
1297 if n == 0 || px == 0 || q == 0 {
1298 return Array2::<f64>::zeros((px, q));
1299 }
1300 if !should_use_faer_matmul(px, q, n) {
1301 let w_y = Array2::from_shape_fn((n, q), |(i, j)| w[i] * y[[i, j]]);
1302 return x.t().dot(&w_y);
1303 }
1304
1305 const TARGET_BYTES: usize = 8 * 1024 * 1024;
1307 const MIN_ROWS: usize = 512;
1308 const MAX_ROWS: usize = 131_072;
1309 let total_cols = px + q;
1310 let chunk_rows = (TARGET_BYTES / (total_cols.max(1) * 8))
1311 .clamp(MIN_ROWS, MAX_ROWS)
1312 .min(n);
1313
1314 let mut result = Array2::<f64>::zeros((px, q).f());
1315 let mut wy_chunk = Array2::<f64>::zeros((chunk_rows, q));
1318
1319 let y_is_row_major = y.is_standard_layout();
1320 let w_slice_opt = w.as_slice();
1321
1322 {
1323 let mut out_view = array2_to_matmut(&mut result);
1324
1325 for start in (0..n).step_by(chunk_rows) {
1326 let rows = (n - start).min(chunk_rows);
1327 {
1328 let chunk_slice = wy_chunk
1329 .as_slice_mut()
1330 .expect("row-major chunk is contiguous");
1331 if y_is_row_major && let (Some(y_all), Some(w_all)) = (y.as_slice(), w_slice_opt) {
1332 for local in 0..rows {
1333 let src = start + local;
1334 let wi = w_all[src];
1335 let src_off = src * q;
1336 let dst_off = local * q;
1337 let src_row = &y_all[src_off..src_off + q];
1338 let dst_row = &mut chunk_slice[dst_off..dst_off + q];
1339 for col in 0..q {
1340 dst_row[col] = src_row[col] * wi;
1341 }
1342 }
1343 } else {
1344 let y_slice = y.slice(s![start..start + rows, ..]);
1345 for local in 0..rows {
1346 let wi = w[start + local];
1347 let yrow = y_slice.row(local);
1348 let dst_off = local * q;
1349 let dst_row = &mut chunk_slice[dst_off..dst_off + q];
1350 for (col, yij) in yrow.iter().enumerate() {
1351 dst_row[col] = yij * wi;
1352 }
1353 }
1354 }
1355 }
1356 let x_slice = x.slice(s![start..start + rows, ..]);
1357 let wy_slice = wy_chunk.slice(s![0..rows, ..]);
1358 let x_view = FaerArrayView::new(&x_slice);
1359 let wy_view = FaerArrayView::new(&wy_slice);
1360 let par = matmul_parallelism(px, q, rows);
1361 matmul(
1362 out_view.as_mut(),
1363 Accum::Add,
1364 x_view.as_ref().transpose(),
1365 wy_view.as_ref(),
1366 1.0,
1367 par,
1368 );
1369 }
1370 }
1371
1372 result
1373}
1374
1375pub fn fast_joint_hessian_2x2<
1381 S1: Data<Elem = f64>,
1382 S2: Data<Elem = f64>,
1383 S3: Data<Elem = f64>,
1384 S4: Data<Elem = f64>,
1385 S5: Data<Elem = f64>,
1386>(
1387 x_a: &ArrayBase<S1, Ix2>,
1388 x_b: &ArrayBase<S2, Ix2>,
1389 w_aa: &ArrayBase<S3, Ix1>,
1390 w_ab: &ArrayBase<S4, Ix1>,
1391 w_bb: &ArrayBase<S5, Ix1>,
1392) -> Array2<f64> {
1393 if let Some(out) = crate::gpu_hook::gpu_dispatch().and_then(|d| {
1394 d.try_fast_joint_hessian_2x2(
1395 x_a.view(),
1396 x_b.view(),
1397 w_aa.view(),
1398 w_ab.view(),
1399 w_bb.view(),
1400 )
1401 }) {
1402 return out;
1403 }
1404 fast_joint_hessian_2x2_impl(x_a, x_b, w_aa, w_ab, w_bb)
1405}
1406
1407#[inline]
1408fn fast_joint_hessian_2x2_impl<
1409 S1: Data<Elem = f64>,
1410 S2: Data<Elem = f64>,
1411 S3: Data<Elem = f64>,
1412 S4: Data<Elem = f64>,
1413 S5: Data<Elem = f64>,
1414>(
1415 x_a: &ArrayBase<S1, Ix2>,
1416 x_b: &ArrayBase<S2, Ix2>,
1417 w_aa: &ArrayBase<S3, Ix1>,
1418 w_ab: &ArrayBase<S4, Ix1>,
1419 w_bb: &ArrayBase<S5, Ix1>,
1420) -> Array2<f64> {
1421 use faer::Accum;
1422 use faer::linalg::matmul::matmul;
1423 use ndarray::{ShapeBuilder, s};
1424
1425 let n = x_a.nrows();
1426 let pa = x_a.ncols();
1427 let pb = x_b.ncols();
1428 let total = pa + pb;
1429 assert_eq!(n, x_b.nrows());
1430 assert_eq!(n, w_aa.len());
1431 assert_eq!(n, w_ab.len());
1432 assert_eq!(n, w_bb.len());
1433
1434 if n == 0 || total == 0 {
1435 return Array2::<f64>::zeros((total, total));
1436 }
1437
1438 if !should_use_faer_matmul(pa.max(pb), pa.max(pb), n) {
1440 let waa_xa = Array2::from_shape_fn((n, pa), |(i, j)| w_aa[i] * x_a[[i, j]]);
1441 let wab_xb = Array2::from_shape_fn((n, pb), |(i, j)| w_ab[i] * x_b[[i, j]]);
1442 let wbb_xb = Array2::from_shape_fn((n, pb), |(i, j)| w_bb[i] * x_b[[i, j]]);
1443 let mut out = Array2::<f64>::zeros((total, total));
1444 out.slice_mut(s![..pa, ..pa]).assign(&x_a.t().dot(&waa_xa));
1445 out.slice_mut(s![..pa, pa..]).assign(&x_a.t().dot(&wab_xb));
1446 out.slice_mut(s![pa.., pa..]).assign(&x_b.t().dot(&wbb_xb));
1447 for i in 0..total {
1449 for j in 0..i {
1450 out[[i, j]] = out[[j, i]];
1451 }
1452 }
1453 return out;
1454 }
1455
1456 const TARGET_BYTES: usize = 8 * 1024 * 1024;
1457 const MIN_ROWS: usize = 512;
1458 const MAX_ROWS: usize = 131_072;
1459 let cols_needed = pa + 2 * pb;
1461 let chunk_rows = (TARGET_BYTES / (cols_needed.max(1) * 8))
1462 .clamp(MIN_ROWS, MAX_ROWS)
1463 .min(n);
1464
1465 let mut out = Array2::<f64>::zeros((total, total).f());
1466 let mut waa_xa_buf = Array2::<f64>::zeros((chunk_rows, pa));
1471 let mut wab_xb_buf = Array2::<f64>::zeros((chunk_rows, pb));
1472 let mut wbb_xb_buf = Array2::<f64>::zeros((chunk_rows, pb));
1473
1474 let xa_is_row_major = x_a.is_standard_layout();
1475 let xb_is_row_major = x_b.is_standard_layout();
1476 let waa_slice_opt = w_aa.as_slice();
1477 let wab_slice_opt = w_ab.as_slice();
1478 let wbb_slice_opt = w_bb.as_slice();
1479
1480 {
1481 let mut out_mat = array2_to_matmut(&mut out);
1482
1483 for start in (0..n).step_by(chunk_rows) {
1484 let rows = (n - start).min(chunk_rows);
1485 let xa_slice = x_a.slice(s![start..start + rows, ..]);
1486 let xb_slice = x_b.slice(s![start..start + rows, ..]);
1487
1488 {
1490 let waa_chunk = waa_xa_buf
1491 .as_slice_mut()
1492 .expect("row-major waa chunk is contiguous");
1493 let wab_chunk = wab_xb_buf
1494 .as_slice_mut()
1495 .expect("row-major wab chunk is contiguous");
1496 let wbb_chunk = wbb_xb_buf
1497 .as_slice_mut()
1498 .expect("row-major wbb chunk is contiguous");
1499
1500 if xa_is_row_major
1501 && xb_is_row_major
1502 && let (Some(xa_all), Some(xb_all)) = (x_a.as_slice(), x_b.as_slice())
1503 && let (Some(waa_all), Some(wab_all), Some(wbb_all)) =
1504 (waa_slice_opt, wab_slice_opt, wbb_slice_opt)
1505 {
1506 for local in 0..rows {
1507 let i = start + local;
1508 let waa_i = waa_all[i];
1509 let wab_i = wab_all[i];
1510 let wbb_i = wbb_all[i];
1511 let xa_off = i * pa;
1512 let xa_row = &xa_all[xa_off..xa_off + pa];
1513 let xb_off = i * pb;
1514 let xb_row = &xb_all[xb_off..xb_off + pb];
1515 let waa_off = local * pa;
1516 let wab_off = local * pb;
1517 let wbb_off = local * pb;
1518 let waa_row = &mut waa_chunk[waa_off..waa_off + pa];
1519 for col in 0..pa {
1520 waa_row[col] = xa_row[col] * waa_i;
1521 }
1522 let wab_row = &mut wab_chunk[wab_off..wab_off + pb];
1523 let wbb_row = &mut wbb_chunk[wbb_off..wbb_off + pb];
1524 for col in 0..pb {
1525 let xij = xb_row[col];
1526 wab_row[col] = xij * wab_i;
1527 wbb_row[col] = xij * wbb_i;
1528 }
1529 }
1530 } else {
1531 for local in 0..rows {
1532 let i = start + local;
1533 let waa_i = w_aa[i];
1534 let wab_i = w_ab[i];
1535 let wbb_i = w_bb[i];
1536 let waa_off = local * pa;
1537 let wab_off = local * pb;
1538 let wbb_off = local * pb;
1539 let waa_row = &mut waa_chunk[waa_off..waa_off + pa];
1540 let xa_row = xa_slice.row(local);
1541 for (col, xij) in xa_row.iter().enumerate() {
1542 waa_row[col] = xij * waa_i;
1543 }
1544 let wab_row = &mut wab_chunk[wab_off..wab_off + pb];
1545 let wbb_row = &mut wbb_chunk[wbb_off..wbb_off + pb];
1546 let xb_row = xb_slice.row(local);
1547 for (col, xij) in xb_row.iter().enumerate() {
1548 wab_row[col] = xij * wab_i;
1549 wbb_row[col] = xij * wbb_i;
1550 }
1551 }
1552 }
1553 }
1554
1555 let xa_view = FaerArrayView::new(&xa_slice);
1556 let xb_view = FaerArrayView::new(&xb_slice);
1557 let waa_xa_slice = waa_xa_buf.slice(s![0..rows, ..]);
1558 let wab_xb_slice = wab_xb_buf.slice(s![0..rows, ..]);
1559 let wbb_xb_slice = wbb_xb_buf.slice(s![0..rows, ..]);
1560 let waa_xa_view = FaerArrayView::new(&waa_xa_slice);
1561 let wab_xb_view = FaerArrayView::new(&wab_xb_slice);
1562 let wbb_xb_view = FaerArrayView::new(&wbb_xb_slice);
1563
1564 matmul(
1566 out_mat.rb_mut().submatrix_mut(0, 0, pa, pa),
1567 Accum::Add,
1568 xa_view.as_ref().transpose(),
1569 waa_xa_view.as_ref(),
1570 1.0,
1571 matmul_parallelism(pa, pa, rows),
1572 );
1573 matmul(
1575 out_mat.rb_mut().submatrix_mut(0, pa, pa, pb),
1576 Accum::Add,
1577 xa_view.as_ref().transpose(),
1578 wab_xb_view.as_ref(),
1579 1.0,
1580 matmul_parallelism(pa, pb, rows),
1581 );
1582 matmul(
1584 out_mat.rb_mut().submatrix_mut(pa, pa, pb, pb),
1585 Accum::Add,
1586 xb_view.as_ref().transpose(),
1587 wbb_xb_view.as_ref(),
1588 1.0,
1589 matmul_parallelism(pb, pb, rows),
1590 );
1591 }
1592 } for i in 0..total {
1595 for j in 0..i {
1596 out[[i, j]] = out[[j, i]];
1597 }
1598 }
1599 out
1600}
1601
1602fn mat_to_array(mat: MatRef<'_, f64>) -> Array2<f64> {
1603 let nrows = mat.nrows();
1604 let ncols = mat.ncols();
1605 let mut out = Array2::<f64>::zeros((nrows, ncols));
1606 if nrows == 0 || ncols == 0 {
1607 return out;
1608 }
1609 if let Some(out_slice) = out.as_slice_memory_order_mut() {
1612 for i in 0..nrows {
1614 let row_start = i * ncols;
1615 for j in 0..ncols {
1616 out_slice[row_start + j] = mat[(i, j)];
1617 }
1618 }
1619 } else {
1620 for j in 0..ncols {
1621 for i in 0..nrows {
1622 out[[i, j]] = mat[(i, j)];
1623 }
1624 }
1625 }
1626 out
1627}
1628
1629#[inline]
1632pub fn fast_ab_into<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
1633 a: &ArrayBase<S1, Ix2>,
1634 b: &ArrayBase<S2, Ix2>,
1635 out: &mut Array2<f64>,
1636) {
1637 fast_ab_into_impl(a, b, out);
1638}
1639
1640#[inline]
1641fn fast_ab_into_impl<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
1642 a: &ArrayBase<S1, Ix2>,
1643 b: &ArrayBase<S2, Ix2>,
1644 out: &mut Array2<f64>,
1645) {
1646 use faer::Accum;
1647 use faer::linalg::matmul::matmul;
1648
1649 let (n, p) = a.dim();
1650 let (p_b, q) = b.dim();
1651 assert_eq!(p, p_b, "A and B must have compatible inner dimensions");
1652 assert_eq!(out.dim(), (n, q), "output dimensions must match A*B result");
1653
1654 if !should_use_faer_matmul(n, q, p) {
1655 out.assign(&a.dot(b));
1656 return;
1657 }
1658
1659 let aview = FaerArrayView::new(a);
1660 let bview = FaerArrayView::new(b);
1661 let a_ref = aview.as_ref();
1662 let b_ref = bview.as_ref();
1663
1664 let par = matmul_parallelism(n, q, p);
1665 let mut outview = array2_to_matmut(out);
1666 matmul(outview.as_mut(), Accum::Replace, a_ref, b_ref, 1.0, par);
1667}
1668
1669fn diag_to_array(diag: DiagRef<'_, f64>) -> Array1<f64> {
1670 let mat = diag.column_vector().as_mat();
1671 let mut out = Array1::<f64>::zeros(mat.nrows());
1672 for i in 0..mat.nrows() {
1673 out[i] = mat[(i, 0)];
1674 }
1675 out
1676}
1677
1678pub struct FaerArrayView<'a> {
1679 ptr: *const f64,
1680 rows: usize,
1681 cols: usize,
1682 row_stride: isize,
1683 col_stride: isize,
1684 owned: Option<Array2<f64>>,
1685 marker: PhantomData<&'a f64>,
1686}
1687
1688impl<'a> FaerArrayView<'a> {
1689 #[inline]
1690 pub fn new<S: Data<Elem = f64>>(array: &'a ArrayBase<S, Ix2>) -> Self {
1691 let (rows, cols) = array.dim();
1692 let strides = array.strides();
1693 if strides[0] <= 0 || strides[1] <= 0 {
1697 let owned = array.to_owned();
1698 let owned_strides = owned.strides();
1699 return Self {
1700 ptr: owned.as_ptr(),
1701 rows,
1702 cols,
1703 row_stride: owned_strides[0],
1704 col_stride: owned_strides[1],
1705 owned: Some(owned),
1706 marker: PhantomData,
1707 };
1708 }
1709
1710 Self {
1711 ptr: array.as_ptr(),
1712 rows,
1713 cols,
1714 row_stride: strides[0],
1715 col_stride: strides[1],
1716 owned: None,
1717 marker: PhantomData,
1718 }
1719 }
1720
1721 #[inline]
1722 pub fn as_ref(&self) -> MatRef<'_, f64> {
1723 let (ptr, rows, cols, row_stride, col_stride) = if let Some(owned) = &self.owned {
1724 let strides = owned.strides();
1725 (
1726 owned.as_ptr(),
1727 owned.nrows(),
1728 owned.ncols(),
1729 strides[0],
1730 strides[1],
1731 )
1732 } else {
1733 (
1734 self.ptr,
1735 self.rows,
1736 self.cols,
1737 self.row_stride,
1738 self.col_stride,
1739 )
1740 };
1741 unsafe { MatRef::from_raw_parts(ptr, rows, cols, row_stride, col_stride) }
1745 }
1746}
1747
1748pub struct FaerColView<'a> {
1749 ptr: *const f64,
1750 len: usize,
1751 stride: isize,
1752 owned: Option<Array1<f64>>,
1753 marker: PhantomData<&'a f64>,
1754}
1755
1756impl<'a> FaerColView<'a> {
1757 #[inline]
1758 pub fn new<S: Data<Elem = f64>>(array: &'a ArrayBase<S, Ix1>) -> Self {
1759 let len = array.len();
1760 let stride = array.strides()[0];
1761 if stride <= 0 {
1762 let owned = array.to_owned();
1763 return Self {
1764 ptr: owned.as_ptr(),
1765 len,
1766 stride: 1,
1767 owned: Some(owned),
1768 marker: PhantomData,
1769 };
1770 }
1771 Self {
1772 ptr: array.as_ptr(),
1773 len,
1774 stride,
1775 owned: None,
1776 marker: PhantomData,
1777 }
1778 }
1779
1780 #[inline]
1781 pub fn as_ref(&self) -> MatRef<'_, f64> {
1782 let (ptr, len, stride) = if let Some(owned) = &self.owned {
1783 (owned.as_ptr(), owned.len(), 1)
1784 } else {
1785 (self.ptr, self.len, self.stride)
1786 };
1787 unsafe { MatRef::from_raw_parts(ptr, len, 1, stride, 0) }
1791 }
1792}
1793
1794pub trait FaerSvd {
1795 fn svd(
1796 &self,
1797 compute_u: bool,
1798 computevt: bool,
1799 ) -> Result<(Option<Array2<f64>>, Array1<f64>, Option<Array2<f64>>), FaerLinalgError>;
1800}
1801
1802impl<S: Data<Elem = f64>> FaerSvd for ArrayBase<S, Ix2> {
1803 fn svd(
1804 &self,
1805 compute_u: bool,
1806 computevt: bool,
1807 ) -> Result<(Option<Array2<f64>>, Array1<f64>, Option<Array2<f64>>), FaerLinalgError> {
1808 let faerview = FaerArrayView::new(self);
1809 let faer_mat = faerview.as_ref();
1810 if !compute_u && !computevt {
1811 let (rows, cols) = faer_mat.shape();
1812 let mut singular = Diag::<f64>::zeros(rows.min(cols));
1813 let par = get_global_parallelism();
1814 let mut mem = MemBuffer::new(svd::svd_scratch::<f64>(
1815 rows,
1816 cols,
1817 ComputeSvdVectors::No,
1818 ComputeSvdVectors::No,
1819 par,
1820 Default::default(),
1821 ));
1822 let stack = MemStack::new(&mut mem);
1823 svd::svd(
1824 faer_mat,
1825 singular.as_mut(),
1826 None,
1827 None,
1828 par,
1829 stack,
1830 Default::default(),
1831 )
1832 .map_err(|_| FaerLinalgError::SvdNoConvergence {
1833 context: "faer SVD singular values only",
1834 })?;
1835 let singularvalues = diag_to_array(singular.as_ref());
1836 return Ok((None, singularvalues, None));
1837 }
1838
1839 let (rows, cols) = faer_mat.shape();
1840 let rank = rows.min(cols);
1841 let compute_u_flag = if compute_u {
1842 ComputeSvdVectors::Thin
1843 } else {
1844 ComputeSvdVectors::No
1845 };
1846 let computev_flag = if computevt {
1847 ComputeSvdVectors::Thin
1848 } else {
1849 ComputeSvdVectors::No
1850 };
1851
1852 let mut singular = Diag::<f64>::zeros(rows.min(cols));
1853 let mut u_storage = compute_u.then(|| Mat::<f64>::zeros(rows, rank));
1854 let mut v_storage = computevt.then(|| Mat::<f64>::zeros(cols, rank));
1855
1856 let par = get_global_parallelism();
1857 let mut mem = MemBuffer::new(svd::svd_scratch::<f64>(
1858 rows,
1859 cols,
1860 compute_u_flag,
1861 computev_flag,
1862 par,
1863 Default::default(),
1864 ));
1865 let stack = MemStack::new(&mut mem);
1866
1867 svd::svd(
1868 faer_mat.as_ref(),
1869 singular.as_mut(),
1870 u_storage.as_mut().map(|mat| mat.as_mut()),
1871 v_storage.as_mut().map(|mat| mat.as_mut()),
1872 par,
1873 stack,
1874 Default::default(),
1875 )
1876 .map_err(|_| FaerLinalgError::SvdNoConvergence {
1877 context: "faer SVD with vectors",
1878 })?;
1879
1880 let singularvalues = diag_to_array(singular.as_ref());
1881 let u_opt = u_storage.map(|mat| mat_to_array(mat.as_ref()));
1882 let vt_opt = v_storage.map(|mat| {
1883 let mat_ref = mat.as_ref();
1884 let mut out = Array2::<f64>::zeros((mat_ref.ncols(), mat_ref.nrows()));
1885 for j in 0..mat_ref.nrows() {
1886 for i in 0..mat_ref.ncols() {
1887 out[[i, j]] = mat_ref[(j, i)];
1888 }
1889 }
1890 out
1891 });
1892
1893 Ok((u_opt, singularvalues, vt_opt))
1894 }
1895}
1896
1897pub trait FaerEigh {
1898 fn eigh(&self, side: Side) -> Result<(Array1<f64>, Array2<f64>), FaerLinalgError>;
1899}
1900
1901impl<S: Data<Elem = f64>> FaerEigh for ArrayBase<S, Ix2> {
1902 fn eigh(&self, side: Side) -> Result<(Array1<f64>, Array2<f64>), FaerLinalgError> {
1903 fn try_eigh(
1904 matrix: &Array2<f64>,
1905 side: Side,
1906 ) -> Result<(Array1<f64>, Array2<f64>), FaerLinalgError> {
1907 let faerview = FaerArrayView::new(matrix);
1908 let eigen = catch_unwind(AssertUnwindSafe(|| {
1909 faerview.as_ref().self_adjoint_eigen(side)
1910 }))
1911 .map_err(|_| FaerLinalgError::FactorizationFailed {
1912 context: "self-adjoint eigendecomposition panic boundary",
1913 })?
1914 .map_err(FaerLinalgError::SelfAdjointEigen)?;
1915 let values = diag_to_array(eigen.S());
1916 let vectors = mat_to_array(eigen.U());
1917 Ok((values, vectors))
1918 }
1919
1920 let owned = self.to_owned();
1921 if owned.nrows() != owned.ncols() {
1922 return Err(FaerLinalgError::FactorizationFailed {
1923 context: "self-adjoint eigendecomposition non-square input",
1924 });
1925 }
1926 if owned.nrows() == 0 {
1927 return Ok((Array1::zeros(0), Array2::zeros((0, 0))));
1928 }
1929 if owned.iter().any(|value| !value.is_finite()) {
1930 return Err(FaerLinalgError::SelfAdjointEigenNonFiniteInput {
1931 context: "self-adjoint eigendecomposition input validation",
1932 });
1933 }
1934 if let Ok((evals, evecs)) = try_eigh(&owned, side)
1935 && evals.iter().all(|value| value.is_finite())
1936 && evecs.iter().all(|value| value.is_finite())
1937 {
1938 return Ok((evals, evecs));
1939 }
1940
1941 let mut repaired = owned.clone();
1942 crate::matrix::symmetrize_in_place(&mut repaired);
1943
1944 let scale = repaired
1945 .iter()
1946 .fold(0.0_f64, |acc, &value| acc.max(value.abs()))
1947 .max(1.0);
1948 let scaled = repaired.mapv(|value| value / scale);
1949 const JITTER_SCHEDULE: [f64; 6] = [0.0, 1e-12, 1e-10, 1e-8, 1e-6, 1e-4];
1955 let jitter_schedule = JITTER_SCHEDULE;
1956 let mut last_error = FaerLinalgError::FactorizationFailed {
1957 context: "self-adjoint eigendecomposition repair attempts",
1958 };
1959
1960 for &jitter in &jitter_schedule {
1961 let mut candidate = scaled.clone();
1962 if jitter > 0.0 {
1963 let n = candidate.nrows();
1964 for i in 0..n {
1965 candidate[[i, i]] += jitter;
1966 }
1967 }
1968
1969 match try_eigh(&candidate, side) {
1970 Ok((mut evals, evecs))
1971 if evals.iter().all(|value| value.is_finite())
1972 && evecs.iter().all(|value| value.is_finite()) =>
1973 {
1974 for value in &mut evals {
1975 *value = (*value - jitter) * scale;
1976 }
1977 return Ok((evals, evecs));
1978 }
1979 Ok((_, _)) => {
1980 last_error = FaerLinalgError::SelfAdjointEigenNonFiniteInput {
1981 context: "self-adjoint eigendecomposition repaired output validation",
1982 };
1983 }
1984 Err(err) => {
1985 last_error = err;
1986 }
1987 }
1988 }
1989
1990 Err(last_error)
1991 }
1992}
1993
1994pub struct FaerCholeskyFactor {
1995 factor: solvers::Llt<f64>,
1996}
1997
1998impl FaerCholeskyFactor {
1999 pub fn solvevec(&self, rhs: &Array1<f64>) -> Array1<f64> {
2000 let mut rhs = rhs.to_owned();
2001 let mut rhsview = array1_to_col_matmut(&mut rhs);
2002 self.factor.solve_in_place(rhsview.as_mut());
2003 rhs
2004 }
2005
2006 pub fn solve_mat_in_place(&self, rhs: &mut Array2<f64>) {
2007 let mut rhsview = array2_to_matmut(rhs);
2008 self.factor.solve_in_place(rhsview.as_mut());
2009 }
2010
2011 pub fn solve_mat_into<S: Data<Elem = f64>>(
2012 &self,
2013 rhs: &ArrayBase<S, Ix2>,
2014 out: &mut Array2<f64>,
2015 ) {
2016 if out.dim() != rhs.dim() {
2017 *out = Array2::<f64>::zeros(rhs.dim());
2018 }
2019 out.assign(rhs);
2020 self.solve_mat_in_place(out);
2021 }
2022
2023 pub fn solve_mat(&self, rhs: &Array2<f64>) -> Array2<f64> {
2024 let mut out = Array2::<f64>::zeros(rhs.dim());
2025 self.solve_mat_into(rhs, &mut out);
2026 out
2027 }
2028
2029 pub fn diag(&self) -> Array1<f64> {
2030 diag_to_array(self.factor.L().diagonal())
2031 }
2032
2033 pub fn lower_triangular(&self) -> Array2<f64> {
2034 mat_to_array(self.factor.L())
2035 }
2036}
2037
2038pub trait FaerCholesky {
2039 fn cholesky(&self, side: Side) -> Result<FaerCholeskyFactor, FaerLinalgError>;
2040}
2041
2042impl<S: Data<Elem = f64>> FaerCholesky for ArrayBase<S, Ix2> {
2043 fn cholesky(&self, side: Side) -> Result<FaerCholeskyFactor, FaerLinalgError> {
2044 let faerview = FaerArrayView::new(self);
2045 let factor = faerview
2046 .as_ref()
2047 .llt(side)
2048 .map_err(FaerLinalgError::Cholesky)?;
2049 Ok(FaerCholeskyFactor { factor })
2050 }
2051}
2052
2053pub trait FaerQr {
2054 fn qr(&self) -> Result<(Array2<f64>, Array2<f64>), FaerLinalgError>;
2055}
2056
2057impl<S: Data<Elem = f64>> FaerQr for ArrayBase<S, Ix2> {
2058 fn qr(&self) -> Result<(Array2<f64>, Array2<f64>), FaerLinalgError> {
2059 let faerview = FaerArrayView::new(self);
2060 let qr = faerview.as_ref().qr();
2061 let q = qr.compute_thin_Q();
2062 let r = qr.thin_R();
2063 Ok((mat_to_array(q.as_ref()), mat_to_array(r)))
2064 }
2065}
2066
2067pub fn rrqr_nullspace_basis<S: Data<Elem = f64>>(
2086 a: &ArrayBase<S, Ix2>,
2087 rank_alpha: f64,
2088) -> Result<(Array2<f64>, usize), FaerLinalgError> {
2089 let faerview = FaerArrayView::new(a);
2090 let qr = faerview.as_ref().col_piv_qr();
2091 let r = qr.thin_R();
2092 let diag_len = r.nrows().min(r.ncols());
2093 let leading_diag = if diag_len > 0 { r[(0, 0)].abs() } else { 0.0 };
2094 let tol = rank_alpha
2095 * f64::EPSILON
2096 * (a.nrows().max(a.ncols()).max(1) as f64)
2097 * leading_diag.max(1.0);
2098 let rank = (0..diag_len).filter(|&i| r[(i, i)].abs() > tol).count();
2099 let z = if rank >= a.nrows() {
2100 Array2::<f64>::zeros((a.nrows(), 0))
2101 } else if rank == 0 {
2102 Array2::<f64>::eye(a.nrows())
2106 } else {
2107 let nullity = a.nrows() - rank;
2108 let mut selector = Mat::<f64>::zeros(a.nrows(), nullity);
2109 for j in 0..nullity {
2110 selector[(rank + j, j)] = 1.0;
2111 }
2112 let par = get_global_parallelism();
2113 faer::linalg::householder::apply_block_householder_sequence_on_the_left_in_place_with_conj(
2114 qr.Q_basis(),
2115 qr.Q_coeff(),
2116 Conj::No,
2117 selector.as_mut(),
2118 par,
2119 MemStack::new(&mut MemBuffer::new(
2120 faer::linalg::householder::apply_block_householder_sequence_on_the_left_in_place_scratch::<f64>(
2121 a.nrows(),
2122 qr.Q_coeff().nrows(),
2123 nullity,
2124 ),
2125 )),
2126 );
2127 mat_to_array(selector.as_ref())
2128 };
2129 Ok((z, rank))
2130}
2131
2132#[inline]
2133pub const fn default_rrqr_rank_alpha() -> f64 {
2134 RRQR_RANK_ALPHA
2135}
2136
2137pub struct RrqrWithPermutation {
2148 pub rank: usize,
2149 pub column_permutation: Vec<usize>,
2150 pub leading_diag_abs: f64,
2151 pub rank_tol: f64,
2152}
2153
2154pub fn rrqr_with_permutation<S: Data<Elem = f64>>(
2163 a: &ArrayBase<S, Ix2>,
2164 rank_alpha: f64,
2165) -> Result<RrqrWithPermutation, FaerLinalgError> {
2166 if a.nrows() == 0 {
2167 return Err(FaerLinalgError::FactorizationFailed {
2168 context: "rrqr_with_permutation: input has zero rows",
2169 });
2170 }
2171 let faerview = FaerArrayView::new(a);
2172 let qr = faerview.as_ref().col_piv_qr();
2173 let r = qr.thin_R();
2174 let diag_len = r.nrows().min(r.ncols());
2175 let leading_diag = if diag_len > 0 { r[(0, 0)].abs() } else { 0.0 };
2176 let tol = rank_alpha
2177 * f64::EPSILON
2178 * (a.nrows().max(a.ncols()).max(1) as f64)
2179 * leading_diag.max(1.0);
2180 let rank = (0..diag_len).filter(|&i| r[(i, i)].abs() > tol).count();
2181 let (forward, _inverse) = qr.P().arrays();
2182 let column_permutation: Vec<usize> = forward.iter().copied().map(|idx| idx.unbound()).collect();
2183 Ok(RrqrWithPermutation {
2184 rank,
2185 column_permutation,
2186 leading_diag_abs: leading_diag,
2187 rank_tol: tol,
2188 })
2189}
2190
2191pub struct RrqrFromGram {
2200 pub rank: usize,
2201 pub column_permutation: Vec<usize>,
2202 pub rank_tol: f64,
2203 pub leading_diag_abs: f64,
2208 pub verdict_margin: f64,
2211}
2212
2213pub fn rrqr_from_gram_with_permutation<S: Data<Elem = f64>>(
2249 gram: &ArrayBase<S, Ix2>,
2250 m_rows: usize,
2251 rank_alpha: f64,
2252) -> Result<RrqrFromGram, FaerLinalgError> {
2253 let p = gram.ncols();
2254 if p == 0 {
2255 return Ok(RrqrFromGram {
2256 rank: 0,
2257 column_permutation: Vec::new(),
2258 rank_tol: 0.0,
2259 leading_diag_abs: 0.0,
2260 verdict_margin: 0.0,
2261 });
2262 }
2263 if gram.nrows() != p {
2264 return Err(FaerLinalgError::FactorizationFailed {
2265 context: "rrqr_from_gram_with_permutation: Gram is not square",
2266 });
2267 }
2268 let (evals, evecs) = gram.eigh(Side::Lower)?;
2277 let mut f = Array2::<f64>::zeros((p, p));
2278 for k in 0..p {
2279 let scale = evals[k].max(0.0).sqrt();
2280 if scale == 0.0 {
2281 continue;
2282 }
2283 for i in 0..p {
2284 f[[k, i]] = scale * evecs[[i, k]];
2285 }
2286 }
2287 let faer_f = FaerArrayView::new(&f);
2291 let qr = faer_f.as_ref().col_piv_qr();
2292 let r = qr.thin_R();
2293 let diag_len = r.nrows().min(r.ncols());
2294 let pivots: Vec<f64> = (0..diag_len).map(|i| r[(i, i)].abs()).collect();
2295 let leading_diag = pivots.first().copied().unwrap_or(0.0);
2296 let (forward, _inverse) = qr.P().arrays();
2297 let column_permutation: Vec<usize> = forward.iter().copied().map(|idx| idx.unbound()).collect();
2298 let tol = rank_alpha * f64::EPSILON * (m_rows.max(p).max(1) as f64) * leading_diag.max(1.0);
2302 let rank = pivots.iter().filter(|&&v| v > tol).count();
2303 let min_kept = pivots[..rank].iter().copied().fold(f64::INFINITY, f64::min);
2304 let max_dropped = pivots[rank..].iter().copied().fold(0.0f64, f64::max);
2305 let kept_margin = if rank == 0 {
2309 f64::INFINITY
2310 } else {
2311 min_kept / tol
2312 };
2313 let dropped_margin = if rank == diag_len {
2314 f64::INFINITY
2315 } else {
2316 tol / max_dropped.max(f64::MIN_POSITIVE)
2317 };
2318 let gram_precision_floor = f64::EPSILON.sqrt() * leading_diag.max(1.0);
2340 let kept_floor_margin = if rank == 0 {
2341 f64::INFINITY
2342 } else {
2343 min_kept / gram_precision_floor.max(f64::MIN_POSITIVE)
2344 };
2345 let verdict_margin = kept_margin.min(dropped_margin).min(kept_floor_margin);
2346 Ok(RrqrFromGram {
2347 rank,
2348 column_permutation,
2349 rank_tol: tol,
2350 leading_diag_abs: leading_diag,
2351 verdict_margin,
2352 })
2353}
2354
2355#[cfg(test)]
2356mod tests {
2357 use super::*;
2358 use ndarray::{array, s};
2359
2360 const JOINT_GRAM_RRQR_TRUST_MARGIN_FOR_TEST: f64 = 1.0e3;
2364
2365 #[test]
2366 fn rrqr_nullspace_basis_is_orthonormal_and_annihilates_transpose() {
2367 let a = array![[1.0, 0.0], [1.0, 0.0], [0.0, 2.0], [0.0, 0.0],];
2368 let (z, rank) =
2369 rrqr_nullspace_basis(&a, default_rrqr_rank_alpha()).expect("RRQR should succeed");
2370 assert_eq!(rank, 2);
2371 assert_eq!(z.nrows(), 4);
2372 assert_eq!(z.ncols(), 2);
2373
2374 let gram = z.t().dot(&z);
2375 let ident = Array2::<f64>::eye(z.ncols());
2376 let gram_err = (&gram - &ident)
2377 .iter()
2378 .fold(0.0_f64, |acc, &v| acc.max(v.abs()));
2379 assert!(gram_err < 1e-10, "Z is not orthonormal: {gram_err:e}");
2380
2381 let residual = a.t().dot(&z);
2382 let resid_max = residual.iter().fold(0.0_f64, |acc, &v| acc.max(v.abs()));
2383 assert!(resid_max < 1e-10, "A^T Z residual too large: {resid_max:e}");
2384 }
2385
2386 #[test]
2387 fn rrqr_with_permutation_attributes_redundant_column() {
2388 let a = array![
2392 [1.0, 0.0, 1.0],
2393 [1.0, 0.0, 1.0],
2394 [0.0, 2.0, 0.0],
2395 [0.0, 0.0, 0.0],
2396 ];
2397 let result =
2398 rrqr_with_permutation(&a, default_rrqr_rank_alpha()).expect("RRQR should succeed");
2399 assert_eq!(result.rank, 2);
2400 assert_eq!(result.column_permutation.len(), 3);
2401 let demoted = result.column_permutation[result.rank..].to_vec();
2402 assert!(
2403 demoted.contains(&2) || demoted.contains(&0),
2404 "demoted suffix should include one of the aliased columns (0 or 2), got {demoted:?}"
2405 );
2406 let mut sorted = result.column_permutation.clone();
2407 sorted.sort();
2408 assert_eq!(
2409 sorted,
2410 vec![0, 1, 2],
2411 "permutation must be a valid bijection on 0..n"
2412 );
2413 }
2414
2415 #[test]
2416 fn rrqr_with_permutation_full_rank_returns_identity_like_order() {
2417 let a = array![[1.0, 0.0], [0.0, 2.0], [0.0, 0.0]];
2418 let result =
2419 rrqr_with_permutation(&a, default_rrqr_rank_alpha()).expect("RRQR should succeed");
2420 assert_eq!(result.rank, 2);
2421 let mut sorted = result.column_permutation.clone();
2422 sorted.sort();
2423 assert_eq!(sorted, vec![0, 1]);
2424 }
2425
2426 #[test]
2427 fn rrqr_with_permutation_rejects_zero_rows() {
2428 let a = Array2::<f64>::zeros((0, 3));
2429 assert!(rrqr_with_permutation(&a, default_rrqr_rank_alpha()).is_err());
2430 }
2431
2432 #[test]
2433 fn rrqr_nullspace_basis_square_zero_matrix_is_finite_identity() {
2434 let a = Array2::<f64>::zeros((3, 3));
2437 let (z, rank) =
2438 rrqr_nullspace_basis(&a, default_rrqr_rank_alpha()).expect("RRQR should succeed");
2439 assert_eq!(rank, 0);
2440 assert_eq!(z.dim(), (3, 3));
2441 assert!(
2442 z.iter().all(|v| v.is_finite()),
2443 "square zero matrix produced a non-finite null basis: {z:?}"
2444 );
2445 let gram = z.t().dot(&z);
2446 let ident = Array2::<f64>::eye(3);
2447 let gram_err = (&gram - &ident)
2448 .iter()
2449 .fold(0.0_f64, |acc, &v| acc.max(v.abs()));
2450 assert!(gram_err < 1e-10, "Z is not orthonormal: {gram_err:e}");
2451 }
2452
2453 #[test]
2454 fn rrqr_nullspace_basis_detectszero_rank_matrix() {
2455 let a = Array2::<f64>::zeros((5, 2));
2456 let (z, rank) =
2457 rrqr_nullspace_basis(&a, default_rrqr_rank_alpha()).expect("RRQR should succeed");
2458 assert_eq!(rank, 0);
2459 assert_eq!(z.dim(), (5, 5));
2460 let ident = Array2::<f64>::eye(5);
2461 let max_err = (&z.slice(s![.., ..5]).to_owned() - &ident)
2462 .iter()
2463 .fold(0.0_f64, |acc, &v| acc.max(v.abs()));
2464 assert!(max_err < 1e-10, "zero matrix should yield identity basis");
2465 }
2466
2467 #[test]
2476 fn eigh_on_nan_matrix_rejects_non_finite_input() {
2477 let mat = array![
2478 [1.0, 0.0, 0.0, 0.0],
2479 [0.0, 2.0, 0.0, 0.0],
2480 [0.0, 0.0, 3.0, f64::NAN],
2481 [0.0, 0.0, f64::NAN, 4.0]
2482 ];
2483 let err = mat
2484 .eigh(Side::Lower)
2485 .expect_err("non-finite symmetric input must be rejected");
2486 assert!(matches!(
2487 err,
2488 FaerLinalgError::SelfAdjointEigenNonFiniteInput { .. }
2489 ));
2490 }
2491
2492 #[test]
2493 fn fast_ata_matches_full_gemm_above_threshold() {
2494 let n = 200;
2497 let p = 40;
2498 let a: Array2<f64> = Array2::from_shape_fn((n, p), |(i, j)| {
2499 ((i * 7 + j * 3) as f64).sin() + 0.1 * j as f64
2500 });
2501 let expected = a.t().dot(&a);
2502 let got = fast_ata(&a);
2503 let max_err = (&got - &expected)
2504 .iter()
2505 .fold(0.0_f64, |acc, &v| acc.max(v.abs()));
2506 assert!(max_err < 1e-10, "fast_ata mismatch: {max_err:e}");
2507 for i in 0..p {
2509 for j in 0..p {
2510 assert!((got[[i, j]] - got[[j, i]]).abs() < 1e-12);
2511 }
2512 }
2513 }
2514
2515 #[test]
2516 fn fast_xt_diag_x_matches_naive_above_threshold() {
2517 let n = 400;
2518 let p = 36;
2519 let x: Array2<f64> =
2520 Array2::from_shape_fn((n, p), |(i, j)| (i as f64 * 0.1).cos() + j as f64 * 0.05);
2521 let w: Array1<f64> = Array1::from_shape_fn(n, |i| (i as f64 * 0.03).sin());
2522 let wx = Array2::from_shape_fn((n, p), |(i, j)| w[i] * x[[i, j]]);
2524 let expected = x.t().dot(&wx);
2525 let got = fast_xt_diag_x(&x, &w);
2526 let max_err = (&got - &expected)
2527 .iter()
2528 .fold(0.0_f64, |acc, &v| acc.max(v.abs()));
2529 assert!(max_err < 1e-9, "fast_xt_diag_x mismatch: {max_err:e}");
2530 for i in 0..p {
2531 for j in 0..p {
2532 assert!((got[[i, j]] - got[[j, i]]).abs() < 1e-12);
2533 }
2534 }
2535 }
2536
2537 #[test]
2538 fn stream_weighted_crossprod_full_and_triangular_parity_with_negative_weights() {
2539 for &(n, p) in &[(900usize, 40usize), (8usize, 3usize)] {
2548 let x: Array2<f64> =
2549 Array2::from_shape_fn((n, p), |(i, j)| (i as f64 * 0.07).cos() + j as f64 * 0.013);
2550 let w: Array1<f64> =
2553 Array1::from_shape_fn(n, |i| (i as f64 * 0.11).sin() - 0.25 * (i % 3) as f64);
2554 assert!(
2555 w.iter().any(|&v| v < 0.0),
2556 "weight vector must contain negatives to test sign preservation"
2557 );
2558
2559 let wx = Array2::from_shape_fn((n, p), |(i, j)| w[i] * x[[i, j]]);
2561 let expected = x.t().dot(&wx);
2562
2563 let par = matmul_parallelism(p, p, n);
2564
2565 let mut full = Array2::<f64>::ones((p, p));
2567 stream_weighted_crossprod_into(
2568 &x,
2569 &w,
2570 &mut full,
2571 CrossprodStructure::Full,
2572 CrossprodAccum::Replace,
2573 par,
2574 );
2575
2576 let mut tri = Array2::<f64>::from_elem((p, p), -7.0);
2580 stream_weighted_crossprod_into(
2581 &x,
2582 &w,
2583 &mut tri,
2584 CrossprodStructure::SymmetricLower,
2585 CrossprodAccum::Replace,
2586 par,
2587 );
2588
2589 let full_err = (&full - &expected)
2590 .iter()
2591 .fold(0.0_f64, |a, &v| a.max(v.abs()));
2592 let tri_err = (&tri - &expected)
2593 .iter()
2594 .fold(0.0_f64, |a, &v| a.max(v.abs()));
2595 assert!(
2596 full_err < 1e-9,
2597 "full kernel mismatch (n={n}, p={p}): {full_err:e}"
2598 );
2599 assert!(
2600 tri_err < 1e-9,
2601 "triangular kernel mismatch (n={n}, p={p}): {tri_err:e}"
2602 );
2603
2604 for i in 0..p {
2607 for j in 0..p {
2608 assert!(
2609 (full[[i, j]] - tri[[i, j]]).abs() < 1e-12,
2610 "full vs triangular disagree at ({i},{j})"
2611 );
2612 assert!(
2613 (tri[[i, j]] - tri[[j, i]]).abs() < 1e-12,
2614 "triangular output not symmetric at ({i},{j})"
2615 );
2616 }
2617 }
2618
2619 let base = Array2::<f64>::from_elem((p, p), 1.5);
2622 let mut add_full = base.clone();
2623 stream_weighted_crossprod_into(
2624 &x,
2625 &w,
2626 &mut add_full,
2627 CrossprodStructure::Full,
2628 CrossprodAccum::Add,
2629 par,
2630 );
2631 let mut add_tri = base.clone();
2632 stream_weighted_crossprod_into(
2633 &x,
2634 &w,
2635 &mut add_tri,
2636 CrossprodStructure::SymmetricLower,
2637 CrossprodAccum::Add,
2638 par,
2639 );
2640 let expected_add = &base + &expected;
2641 let add_full_err = (&add_full - &expected_add)
2642 .iter()
2643 .fold(0.0_f64, |a, &v| a.max(v.abs()));
2644 let add_tri_err = (&add_tri - &expected_add)
2645 .iter()
2646 .fold(0.0_f64, |a, &v| a.max(v.abs()));
2647 assert!(
2648 add_full_err < 1e-9,
2649 "full Add mismatch (n={n}, p={p}): {add_full_err:e}"
2650 );
2651 assert!(
2652 add_tri_err < 1e-9,
2653 "triangular Add mismatch (n={n}, p={p}): {add_tri_err:e}"
2654 );
2655
2656 let returned = fast_xt_diag_x(&x, &w);
2659 let returned_err = (&returned - &full)
2660 .iter()
2661 .fold(0.0_f64, |a, &v| a.max(v.abs()));
2662 assert!(
2663 returned_err < 1e-12,
2664 "return adapter vs stream-into adapter disagree (n={n}, p={p}): {returned_err:e}"
2665 );
2666 }
2667 }
2668
2669 #[test]
2670 fn eigh_succeeds_on_same_structure_without_nan() {
2671 let mat = array![[1.0, 0.5, 0.1], [0.5, 2.0, 0.3], [0.1, 0.3, 1.5]];
2673 let (evals, _) = mat
2674 .eigh(Side::Lower)
2675 .expect("eigh should succeed on a well-conditioned finite matrix");
2676 assert!(
2677 evals.iter().all(|&v| v.is_finite()),
2678 "all eigenvalues should be finite"
2679 );
2680 }
2681
2682 #[test]
2691 fn gram_rrqr_flags_low_margin_on_exact_collinearity_so_caller_falls_back() {
2692 let n = 48usize;
2695 let x: Vec<f64> = (0..n)
2696 .map(|i| -1.0 + 2.0 * (i as f64) / (n as f64 - 1.0))
2697 .collect();
2698 let mut a = Array2::<f64>::zeros((n, 4));
2699 for i in 0..n {
2700 a[[i, 0]] = 1.0;
2701 a[[i, 1]] = x[i];
2702 a[[i, 2]] = x[i];
2703 a[[i, 3]] = x[i] * x[i];
2704 }
2705 let alpha = default_rrqr_rank_alpha();
2706
2707 let tall = rrqr_with_permutation(&a, alpha).expect("tall RRQR should succeed");
2710 assert_eq!(tall.rank, 3, "tall RRQR must demote the exact alias");
2711
2712 let unit = Array1::<f64>::ones(n);
2723 let gram = fast_xt_diag_x_with_parallelism(&a, &unit, faer::get_global_parallelism());
2724 let gram_rrqr =
2725 rrqr_from_gram_with_permutation(&gram, n, alpha).expect("Gram RRQR should succeed");
2726 let ok =
2727 gram_rrqr.rank == 3 || gram_rrqr.verdict_margin < JOINT_GRAM_RRQR_TRUST_MARGIN_FOR_TEST;
2728 assert!(
2729 ok,
2730 "gam#933: Gram RRQR must either find correct rank=3 OR signal low margin \
2731 (< {:.0e}) to force the tall fallback; got rank={} margin={:.3e}",
2732 JOINT_GRAM_RRQR_TRUST_MARGIN_FOR_TEST, gram_rrqr.rank, gram_rrqr.verdict_margin,
2733 );
2734 }
2735
2736 #[test]
2741 fn gram_rrqr_keeps_high_margin_on_full_rank_design() {
2742 let n = 200usize;
2743 let p = 5usize;
2744 let mut a = Array2::<f64>::zeros((n, p));
2745 for i in 0..n {
2747 let t = (i as f64) / (n as f64 - 1.0);
2748 a[[i, 0]] = 1.0;
2749 a[[i, 1]] = t;
2750 a[[i, 2]] = t * t;
2751 a[[i, 3]] = t * t * t;
2752 a[[i, 4]] = (t * 6.0).sin();
2753 }
2754 let alpha = default_rrqr_rank_alpha();
2755 let unit = Array1::<f64>::ones(n);
2756 let gram = fast_xt_diag_x_with_parallelism(&a, &unit, faer::get_global_parallelism());
2757 let gram_rrqr =
2758 rrqr_from_gram_with_permutation(&gram, n, alpha).expect("Gram RRQR should succeed");
2759 assert_eq!(gram_rrqr.rank, p, "full-rank design must keep all columns");
2760 assert!(
2761 gram_rrqr.verdict_margin >= JOINT_GRAM_RRQR_TRUST_MARGIN_FOR_TEST,
2762 "full-rank design must keep a high margin (fast Gram path); got {:.3e}",
2763 gram_rrqr.verdict_margin,
2764 );
2765 }
2766
2767 fn max_abs_diff(a: &Array2<f64>, b: &Array2<f64>) -> f64 {
2770 assert_eq!(a.dim(), b.dim(), "shape mismatch in max_abs_diff");
2771 a.iter()
2772 .zip(b.iter())
2773 .fold(0.0_f64, |acc, (&x, &y)| acc.max((x - y).abs()))
2774 }
2775
2776 fn max_abs_diff_1d(a: &Array1<f64>, b: &Array1<f64>) -> f64 {
2777 assert_eq!(a.len(), b.len(), "len mismatch in max_abs_diff_1d");
2778 a.iter()
2779 .zip(b.iter())
2780 .fold(0.0_f64, |acc, (&x, &y)| acc.max((x - y).abs()))
2781 }
2782
2783 #[test]
2785 fn fast_ab_small_matches_ndarray_dot() {
2786 let a = array![[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]];
2787 let b = array![[7.0, 8.0], [9.0, 10.0], [11.0, 12.0]];
2788 let got = fast_ab(&a, &b);
2789 let want = a.dot(&b);
2790 assert!(max_abs_diff(&got, &want) < 1e-12, "fast_ab small mismatch");
2791 assert_eq!(got.dim(), (2, 2));
2792 }
2793
2794 #[test]
2796 fn fast_ab_large_matches_ndarray_dot() {
2797 let n = 50usize;
2798 let p = 40usize;
2799 let q = 35usize;
2800 let mut a = Array2::<f64>::zeros((n, p));
2801 let mut b = Array2::<f64>::zeros((p, q));
2802 let mut state = 0xDEAD_BEEF_1234_5678u64;
2803 let next = |s: &mut u64| -> f64 {
2804 *s ^= *s << 13;
2805 *s ^= *s >> 7;
2806 *s ^= *s << 17;
2807 ((*s >> 11) as f64 / ((1u64 << 53) as f64)) - 0.5
2808 };
2809 for v in a.iter_mut() {
2810 *v = next(&mut state);
2811 }
2812 for v in b.iter_mut() {
2813 *v = next(&mut state);
2814 }
2815 let got = fast_ab(&a, &b);
2816 let want = a.dot(&b);
2817 assert!(max_abs_diff(&got, &want) < 1e-9, "fast_ab large mismatch");
2818 }
2819
2820 #[test]
2822 fn fast_atb_small_matches_ndarray_dot() {
2823 let a = array![[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]];
2824 let b = array![[7.0, 8.0, 9.0], [10.0, 11.0, 12.0], [13.0, 14.0, 15.0]];
2825 let got = fast_atb(&a, &b);
2826 let want = a.t().dot(&b);
2827 assert!(max_abs_diff(&got, &want) < 1e-12, "fast_atb small mismatch");
2828 assert_eq!(got.dim(), (2, 3));
2829 }
2830
2831 #[test]
2833 fn fast_atb_large_matches_ndarray_dot() {
2834 let n = 50usize;
2835 let p = 40usize;
2836 let q = 35usize;
2837 let mut a = Array2::<f64>::zeros((n, p));
2838 let mut b = Array2::<f64>::zeros((n, q));
2839 let mut state = 0xCAFE_BABE_9876_5432u64;
2840 let next = |s: &mut u64| -> f64 {
2841 *s ^= *s << 13;
2842 *s ^= *s >> 7;
2843 *s ^= *s << 17;
2844 ((*s >> 11) as f64 / ((1u64 << 53) as f64)) - 0.5
2845 };
2846 for v in a.iter_mut() {
2847 *v = next(&mut state);
2848 }
2849 for v in b.iter_mut() {
2850 *v = next(&mut state);
2851 }
2852 let got = fast_atb(&a, &b);
2853 let want = a.t().dot(&b);
2854 assert!(max_abs_diff(&got, &want) < 1e-9, "fast_atb large mismatch");
2855 }
2856
2857 #[test]
2859 fn fast_abt_small_matches_ndarray_dot() {
2860 let a = array![[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]];
2861 let b = array![[7.0, 8.0, 9.0], [10.0, 11.0, 12.0]];
2862 let got = fast_abt(&a, &b);
2863 let want = a.dot(&b.t());
2864 assert!(max_abs_diff(&got, &want) < 1e-12, "fast_abt small mismatch");
2865 assert_eq!(got.dim(), (2, 2));
2866 }
2867
2868 #[test]
2870 fn fast_av_small_matches_ndarray_dot() {
2871 let a = array![[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]];
2872 let v = array![1.0, -1.0, 2.0];
2873 let got = fast_av(&a, &v);
2874 let want = a.dot(&v);
2875 assert!(
2876 max_abs_diff_1d(&got, &want) < 1e-12,
2877 "fast_av small mismatch"
2878 );
2879 assert!((got[0] - 5.0).abs() < 1e-12, "fast_av[0] should be 5");
2881 assert!((got[1] - 11.0).abs() < 1e-12, "fast_av[1] should be 11");
2883 }
2884
2885 #[test]
2887 fn fast_av_large_matches_ndarray_dot() {
2888 let n = 50usize;
2889 let p = 40usize;
2890 let mut a = Array2::<f64>::zeros((n, p));
2891 let mut v = Array1::<f64>::zeros(p);
2892 let mut state = 0xFEED_FACE_ABCD_EF01u64;
2893 let next = |s: &mut u64| -> f64 {
2894 *s ^= *s << 13;
2895 *s ^= *s >> 7;
2896 *s ^= *s << 17;
2897 ((*s >> 11) as f64 / ((1u64 << 53) as f64)) - 0.5
2898 };
2899 for v in a.iter_mut() {
2900 *v = next(&mut state);
2901 }
2902 for x in v.iter_mut() {
2903 *x = next(&mut state);
2904 }
2905 let got = fast_av(&a, &v);
2906 let want = a.dot(&v);
2907 assert!(
2908 max_abs_diff_1d(&got, &want) < 1e-9,
2909 "fast_av large mismatch"
2910 );
2911 }
2912
2913 #[test]
2915 fn fast_atv_small_matches_ndarray_dot() {
2916 let a = array![[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]];
2917 let v = array![1.0, 0.0, -1.0];
2918 let got = fast_atv(&a, &v);
2919 let want = a.t().dot(&v);
2920 assert!(
2922 max_abs_diff_1d(&got, &want) < 1e-12,
2923 "fast_atv small mismatch"
2924 );
2925 assert!((got[0] - (-4.0)).abs() < 1e-12, "fast_atv[0]");
2926 assert!((got[1] - (-4.0)).abs() < 1e-12, "fast_atv[1]");
2927 }
2928
2929 #[test]
2931 fn fast_atv_large_matches_ndarray_dot() {
2932 let n = 50usize;
2933 let p = 40usize;
2934 let mut a = Array2::<f64>::zeros((n, p));
2935 let mut v = Array1::<f64>::zeros(n);
2936 let mut state = 0x1234_ABCD_5678_EF90u64;
2937 let next = |s: &mut u64| -> f64 {
2938 *s ^= *s << 13;
2939 *s ^= *s >> 7;
2940 *s ^= *s << 17;
2941 ((*s >> 11) as f64 / ((1u64 << 53) as f64)) - 0.5
2942 };
2943 for x in a.iter_mut() {
2944 *x = next(&mut state);
2945 }
2946 for x in v.iter_mut() {
2947 *x = next(&mut state);
2948 }
2949 let got = fast_atv(&a, &v);
2950 let want = a.t().dot(&v);
2951 assert!(
2952 max_abs_diff_1d(&got, &want) < 1e-9,
2953 "fast_atv large mismatch"
2954 );
2955 }
2956
2957 #[test]
2960 fn fast_xt_diag_y_small_matches_manual() {
2961 let x = array![[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]];
2962 let d = array![2.0, 0.5, 1.0];
2963 let y = array![[7.0, 8.0, 9.0], [10.0, 11.0, 12.0], [13.0, 14.0, 15.0]];
2964 let got = fast_xt_diag_y(&x, &d, &y);
2965 let diag_y = {
2967 let mut dy = Array2::<f64>::zeros(y.dim());
2968 for i in 0..3 {
2969 for j in 0..3 {
2970 dy[[i, j]] = d[i] * y[[i, j]];
2971 }
2972 }
2973 dy
2974 };
2975 let want = x.t().dot(&diag_y);
2976 assert!(
2977 max_abs_diff(&got, &want) < 1e-12,
2978 "fast_xt_diag_y small mismatch"
2979 );
2980 assert_eq!(got.dim(), (2, 3));
2981 }
2982
2983 #[inline]
2990 fn two_prod(a: f64, b: f64) -> (f64, f64) {
2991 let p = a * b;
2992 let e = a.mul_add(b, -p);
2993 (p, e)
2994 }
2995
2996 #[inline]
2997 fn two_sum(a: f64, b: f64) -> (f64, f64) {
2998 let s = a + b;
2999 let bb = s - a;
3000 let e = (a - (s - bb)) + (b - bb);
3001 (s, e)
3002 }
3003
3004 fn grow_expansion(e: &mut Vec<f64>, mut q: f64) {
3006 for h in e.iter_mut() {
3007 let (s, err) = two_sum(*h, q);
3008 *h = err;
3009 q = s;
3010 }
3011 if q != 0.0 {
3012 e.push(q);
3013 }
3014 }
3015
3016 fn exact_dot(a: &[f64], b: &[f64]) -> f64 {
3021 let mut e: Vec<f64> = Vec::new();
3022 for (&x, &y) in a.iter().zip(b.iter()) {
3023 let (p, ep) = two_prod(x, y);
3024 grow_expansion(&mut e, p);
3025 grow_expansion(&mut e, ep);
3026 }
3027 e.iter().fold(0.0f64, |acc, &c| acc + c)
3030 }
3031
3032 fn dd_dot(a: &[f64], b: &[f64]) -> f64 {
3036 let (mut s, mut c) = (0.0f64, 0.0f64);
3037 for (&x, &y) in a.iter().zip(b.iter()) {
3038 let (p, ep) = two_prod(x, y);
3039 let (s2, es) = two_sum(s, p);
3040 s = s2;
3041 c += ep + es;
3042 }
3043 s + c
3044 }
3045
3046 fn naive_dot(a: &[f64], b: &[f64]) -> f64 {
3047 let mut acc = 0.0f64;
3048 for (&x, &y) in a.iter().zip(b.iter()) {
3049 acc += x * y;
3050 }
3051 acc
3052 }
3053
3054 fn ill_conditioned_pair(len: usize, seed: u64) -> (Vec<f64>, Vec<f64>) {
3057 let mut s = seed | 1;
3058 let mut next = || {
3059 s ^= s << 13;
3060 s ^= s >> 7;
3061 s ^= s << 17;
3062 (s >> 11) as f64 / ((1u64 << 53) as f64) - 0.5
3063 };
3064 let mut a = Vec::with_capacity(len);
3065 let mut b = Vec::with_capacity(len);
3066 for i in 0..len {
3067 let scale = 10f64.powi((i % 17) as i32 - 8);
3069 let sign = if i % 2 == 0 { 1.0 } else { -1.0 };
3070 a.push(sign * next() * scale);
3071 b.push(next() * scale);
3072 }
3073 (a, b)
3074 }
3075
3076 #[test]
3079 fn fma_dot_beats_naive_accuracy() {
3080 let mut fma_total = 0.0f64;
3081 let mut naive_total = 0.0f64;
3082 let mut strict_wins = 0;
3083 for seed in 0..64u64 {
3084 let len = 200 + (seed as usize % 57);
3085 let (a, b) = ill_conditioned_pair(len, 0x9E37_79B9 ^ seed.wrapping_mul(2654435761));
3086 let truth = exact_dot(&a, &b);
3087 let fe = (super::fma_dot(&a, &b) - truth).abs();
3088 let ne = (naive_dot(&a, &b) - truth).abs();
3089 let floor = 8.0 * f64::EPSILON * truth.abs();
3093 assert!(
3094 fe <= ne * (1.0 + 1e-6) + floor,
3095 "fma_dot worse than naive: seed={seed} fma_err={fe:.3e} naive_err={ne:.3e}",
3096 );
3097 if fe < ne {
3098 strict_wins += 1;
3099 }
3100 fma_total += fe;
3101 naive_total += ne;
3102 }
3103 assert!(
3104 fma_total < naive_total,
3105 "fma_dot aggregate error {fma_total:.3e} not below naive {naive_total:.3e}",
3106 );
3107 assert!(
3108 strict_wins >= 40,
3109 "expected fma_dot to strictly win the majority; only {strict_wins}/64",
3110 );
3111 }
3112
3113 #[test]
3116 fn fast_atv_blocked_beats_naive_accuracy() {
3117 let n = 200_003usize;
3118 let p = 3usize;
3119 let mut s = 0xD1B5_4A32u64;
3120 let mut next = || {
3121 s ^= s << 13;
3122 s ^= s >> 7;
3123 s ^= s << 17;
3124 (s >> 11) as f64 / ((1u64 << 53) as f64) - 0.5
3125 };
3126 let mut x = Array2::<f64>::zeros((n, p));
3127 let mut v = Array1::<f64>::zeros(n);
3128 for i in 0..n {
3129 let scale = 10f64.powi((i % 17) as i32 - 8);
3130 v[i] = if i % 2 == 0 { scale } else { -scale } * next();
3131 for j in 0..p {
3132 x[[i, j]] = next() * scale;
3133 }
3134 }
3135 let got = fast_atv(&x, &v);
3136 for j in 0..p {
3138 let col: Vec<f64> = (0..n).map(|i| x[[i, j]]).collect();
3139 let vv: Vec<f64> = v.to_vec();
3140 let truth = dd_dot(&col, &vv);
3141 let naive = naive_dot(&col, &vv);
3142 let ge = (got[j] - truth).abs();
3143 let ne = (naive - truth).abs();
3144 assert!(
3145 ge <= ne + f64::MIN_POSITIVE,
3146 "col {j}: blocked err {ge:.3e} exceeds naive {ne:.3e}",
3147 );
3148 }
3149 }
3150
3151 #[test]
3154 fn fast_av_strided_input_matches_ndarray() {
3155 let mut base = Array2::<f64>::zeros((40, 60));
3156 let mut s = 0x0BAD_F00Du64;
3157 let mut next = || {
3158 s ^= s << 13;
3159 s ^= s >> 7;
3160 s ^= s << 17;
3161 (s >> 11) as f64 / ((1u64 << 53) as f64) - 0.5
3162 };
3163 for x in base.iter_mut() {
3164 *x = next();
3165 }
3166 let a = base.t();
3168 let mut v = Array1::<f64>::zeros(40);
3169 for x in v.iter_mut() {
3170 *x = next();
3171 }
3172 let got = fast_av(&a, &v);
3173 let want = a.dot(&v);
3174 assert!(
3175 max_abs_diff_1d(&got, &want) < 1e-11,
3176 "strided fast_av mismatch (fallback path)",
3177 );
3178 }
3179
3180 #[test]
3192 fn faer_sequential_scope_sets_seq_inside_and_restores_after() {
3193 let baseline = faer::get_global_parallelism();
3194 faer::set_global_parallelism(Par::rayon(4));
3197 assert_eq!(
3198 faer::get_global_parallelism(),
3199 Par::rayon(4),
3200 "baseline must be the parallel policy we just set",
3201 );
3202
3203 {
3204 let faer_seq_guard = FaerSequentialScope::enter();
3205 assert_eq!(
3206 faer::get_global_parallelism(),
3207 Par::Seq,
3208 "faer must be pinned to Par::Seq inside the scope",
3209 );
3210
3211 {
3213 let faer_seq_inner_guard = FaerSequentialScope::enter();
3214 assert_eq!(
3215 faer::get_global_parallelism(),
3216 Par::Seq,
3217 "nested scope stays Par::Seq",
3218 );
3219 drop(faer_seq_inner_guard);
3220 }
3221 assert_eq!(
3222 faer::get_global_parallelism(),
3223 Par::Seq,
3224 "inner drop must not restore while outer scope is still live",
3225 );
3226 drop(faer_seq_guard);
3227 }
3228
3229 assert_eq!(
3230 faer::get_global_parallelism(),
3231 Par::rayon(4),
3232 "outermost drop must restore the pre-scope parallelism policy",
3233 );
3234
3235 let observed = with_faer_sequential(|| faer::get_global_parallelism());
3237 assert_eq!(
3238 observed,
3239 Par::Seq,
3240 "with_faer_sequential runs body under Seq"
3241 );
3242 assert_eq!(
3243 faer::get_global_parallelism(),
3244 Par::rayon(4),
3245 "with_faer_sequential restores after the body returns",
3246 );
3247
3248 faer::set_global_parallelism(baseline);
3250 }
3251}