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