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, ArrayView1, ArrayViewMut1, Data, Ix1, Ix2};
11use std::marker::PhantomData;
12use std::panic::{AssertUnwindSafe, catch_unwind};
13use std::sync::atomic::{AtomicU64, Ordering};
14use thiserror::Error;
15
16pub fn symmetric_matvec_into(
24 matrix: &Array2<f64>,
25 vector: &[f64],
26 output: &mut [f64],
27) -> Result<(), String> {
28 let n = matrix.nrows();
29 if matrix.ncols() != n || vector.len() != n || output.len() != n {
30 return Err(format!(
31 "symmetric matvec shape mismatch: matrix={:?}, vector={}, output={}",
32 matrix.dim(),
33 vector.len(),
34 output.len()
35 ));
36 }
37 fast_av_standard_view_into(
38 matrix,
39 &ArrayView1::from(vector),
40 ArrayViewMut1::from(output),
41 );
42 Ok(())
43}
44
45const RRQR_RANK_ALPHA: f64 = 100.0;
46
47thread_local! {
48 static NESTED_PARALLEL_DEPTH: std::cell::Cell<usize> = const { std::cell::Cell::new(0) };
49}
50
51struct NestedParallelGuard;
52
53impl NestedParallelGuard {
54 #[inline]
55 fn enter() -> Self {
56 NESTED_PARALLEL_DEPTH.with(|depth| depth.set(depth.get().saturating_add(1)));
57 Self
58 }
59}
60
61impl Drop for NestedParallelGuard {
62 #[inline]
63 fn drop(&mut self) {
64 NESTED_PARALLEL_DEPTH.with(|depth| depth.set(depth.get().saturating_sub(1)));
65 }
66}
67
68#[inline]
78pub fn with_nested_parallel<T>(body: impl FnOnce() -> T) -> T {
79 let guard = NestedParallelGuard::enter();
80 let out = body();
81 drop(guard);
82 out
83}
84
85#[inline]
88pub fn in_nested_parallel_region() -> bool {
89 NESTED_PARALLEL_DEPTH.with(|depth| depth.get() > 0)
90}
91
92static EIGH_CALLS: AtomicU64 = AtomicU64::new(0);
101static EIGH_NANOS: AtomicU64 = AtomicU64::new(0);
102static EIGH_SEQ_CALLS: AtomicU64 = AtomicU64::new(0);
107static EIGH_MAX_DIM: AtomicU64 = AtomicU64::new(0);
110
111thread_local! {
126 static EIGH_THREAD: std::cell::Cell<EighCensus> = const {
127 std::cell::Cell::new(EighCensus {
128 calls: 0,
129 sequential_calls: 0,
130 max_dim: 0,
131 nanos: 0,
132 })
133 };
134}
135
136fn record_thread_eigh(sequential: bool, dim: u64, nanos: u64) {
138 EIGH_THREAD.with(|cell| {
139 let mut census = cell.get();
140 census.calls += 1;
141 if sequential {
142 census.sequential_calls += 1;
143 }
144 census.max_dim = census.max_dim.max(dim);
145 census.nanos += nanos;
146 cell.set(census);
147 });
148}
149
150#[derive(Debug, Clone, Copy, PartialEq, Eq)]
152pub struct EighCensus {
153 pub calls: u64,
155 pub sequential_calls: u64,
157 pub max_dim: u64,
159 pub nanos: u64,
161}
162
163#[derive(Debug, Clone, Copy, PartialEq, Eq)]
202pub struct ParallelismSnapshot {
203 pub rayon_current_num_threads: usize,
206 pub faer_global_sequential: bool,
211 pub faer_global_degree: usize,
215 pub faer_sequential_scope_depth: usize,
220 pub process_available_parallelism: Option<usize>,
226}
227
228impl ParallelismSnapshot {
229 pub fn capture() -> Self {
231 Self::from_parts(
232 get_global_parallelism(),
233 rayon::current_num_threads(),
234 faer_sequential_scope_depth(),
235 std::thread::available_parallelism().ok().map(|n| n.get()),
236 )
237 }
238
239 pub fn from_parts(
244 faer_global: Par,
245 rayon_current_num_threads: usize,
246 faer_sequential_scope_depth: usize,
247 process_available_parallelism: Option<usize>,
248 ) -> Self {
249 Self {
250 rayon_current_num_threads,
251 faer_global_sequential: faer_global == Par::Seq,
252 faer_global_degree: faer_global.degree(),
253 faer_sequential_scope_depth,
254 process_available_parallelism,
255 }
256 }
257
258 pub fn inconsistency(&self) -> Option<String> {
266 if self.rayon_current_num_threads == 0 {
267 return Some("rayon reports a pool of zero threads".to_string());
268 }
269 if self.faer_global_degree == 0 {
270 return Some("faer's global parallelism has degree zero".to_string());
271 }
272 if self.faer_global_sequential && self.faer_global_degree != 1 {
273 return Some(format!(
274 "faer is sequential but reports degree {}",
275 self.faer_global_degree
276 ));
277 }
278 if self.faer_sequential_scope_depth > 0 && !self.faer_global_sequential {
279 return Some(format!(
280 "{} live FaerSequentialScope guard(s) but faer's global parallelism is not sequential",
281 self.faer_sequential_scope_depth
282 ));
283 }
284 if self.process_available_parallelism == Some(0) {
285 return Some("this process reports zero available cores".to_string());
286 }
287 None
288 }
289}
290
291impl std::fmt::Display for ParallelismSnapshot {
292 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
293 write!(
294 f,
295 "rayon_current_num_threads={} | faer_global_sequential={} | \
296 faer_global_degree={} | faer_sequential_scope_depth={} | \
297 process_available_parallelism={}",
298 self.rayon_current_num_threads,
299 self.faer_global_sequential,
300 self.faer_global_degree,
301 self.faer_sequential_scope_depth,
302 match self.process_available_parallelism {
303 Some(cores) => cores.to_string(),
304 None => "unavailable".to_string(),
305 },
306 )
307 }
308}
309
310#[inline]
318pub fn effective_global_parallelism() -> Par {
319 if in_nested_parallel_region() {
320 Par::Seq
321 } else {
322 get_global_parallelism()
323 }
324}
325
326static FAER_SEQ_STATE: std::sync::Mutex<FaerSeqState> = std::sync::Mutex::new(FaerSeqState {
348 depth: 0,
349 saved: None,
350});
351
352struct FaerSeqState {
353 depth: usize,
354 saved: Option<Par>,
355}
356
357#[must_use = "the sequential scope only holds while the guard is alive"]
366pub struct FaerSequentialScope {
367 _private: (),
368}
369
370impl FaerSequentialScope {
371 pub fn enter() -> Self {
373 let mut state = FAER_SEQ_STATE
374 .lock()
375 .unwrap_or_else(std::sync::PoisonError::into_inner);
376 if state.depth == 0 {
377 state.saved = Some(get_global_parallelism());
378 faer::set_global_parallelism(Par::Seq);
379 }
380 state.depth += 1;
381 Self { _private: () }
382 }
383}
384
385impl Drop for FaerSequentialScope {
386 fn drop(&mut self) {
387 let mut state = FAER_SEQ_STATE
388 .lock()
389 .unwrap_or_else(std::sync::PoisonError::into_inner);
390 state.depth -= 1;
391 if state.depth == 0 {
392 if let Some(par) = state.saved.take() {
393 faer::set_global_parallelism(par);
394 }
395 }
396 }
397}
398
399pub fn faer_sequential_scope_depth() -> usize {
405 FAER_SEQ_STATE
406 .lock()
407 .unwrap_or_else(std::sync::PoisonError::into_inner)
408 .depth
409}
410
411#[inline]
416pub fn with_faer_sequential<T>(body: impl FnOnce() -> T) -> T {
417 let faer_seq_guard = FaerSequentialScope::enter();
418 let out = body();
419 drop(faer_seq_guard);
420 out
421}
422
423#[derive(Debug, Error)]
424pub enum FaerLinalgError {
425 #[error("Factorization failed in {context}")]
426 FactorizationFailed { context: &'static str },
427 #[error("SVD failed to converge in {context}")]
428 SvdNoConvergence { context: &'static str },
429 #[error("Self-adjoint eigendecomposition input contains non-finite values in {context}")]
430 SelfAdjointEigenNonFiniteInput { context: &'static str },
431 #[error("Strict self-adjoint eigendecomposition rejected its input: {reason}")]
432 StrictSelfAdjointEigenInvalidInput { reason: String },
433 #[error("Self-adjoint eigendecomposition failed: {0:?}")]
434 SelfAdjointEigen(solvers::EvdError),
435 #[error("Cholesky factorization failed: {0:?}")]
436 Cholesky(solvers::LltError),
437 #[error("LDLT factorization failed: {0:?}")]
438 Ldlt(solvers::LdltError),
439}
440
441pub enum FaerSymmetricFactor {
442 Llt(FaerLlt<f64>),
443 Ldlt(FaerLdlt<f64>),
444 Lblt(FaerLblt<f64>),
445}
446
447#[inline]
448pub fn cholesky_factor_logdet(factor: MatRef<'_, f64>) -> f64 {
449 2.0 * diagonal_log_sum(factor.diagonal())
450}
451
452#[inline]
453fn diagonal_log_sum(diagonal: DiagRef<'_, f64>) -> f64 {
454 diagonal
455 .column_vector()
456 .iter()
457 .map(|&x| x.ln())
458 .sum::<f64>()
459}
460
461impl FaerSymmetricFactor {
462 #[inline]
464 pub fn n(&self) -> usize {
465 use faer::linalg::solvers::ShapeCore;
466 match self {
467 FaerSymmetricFactor::Llt(f) => f.nrows(),
468 FaerSymmetricFactor::Ldlt(f) => f.nrows(),
469 FaerSymmetricFactor::Lblt(f) => f.nrows(),
470 }
471 }
472
473 #[inline]
474 pub fn solve(&self, rhs: MatRef<'_, f64>) -> Mat<f64> {
475 match self {
476 FaerSymmetricFactor::Llt(f) => f.solve(rhs),
477 FaerSymmetricFactor::Ldlt(f) => f.solve(rhs),
478 FaerSymmetricFactor::Lblt(f) => f.solve(rhs),
479 }
480 }
481
482 #[inline]
483 pub fn solve_in_place(&self, rhs: MatMut<'_, f64>) {
484 match self {
485 FaerSymmetricFactor::Llt(f) => f.solve_in_place(rhs),
486 FaerSymmetricFactor::Ldlt(f) => f.solve_in_place(rhs),
487 FaerSymmetricFactor::Lblt(f) => f.solve_in_place(rhs),
488 }
489 }
490}
491
492impl crate::matrix::FactorizedSystem for FaerSymmetricFactor {
493 fn solve(&self, rhs: &Array1<f64>) -> Result<Array1<f64>, String> {
494 let mut out = rhs.clone();
495 let mut out_mat = array1_to_col_matmut(&mut out);
496 self.solve_in_place(out_mat.as_mut());
497 if !out.iter().all(|v| v.is_finite()) {
498 return Err("symmetric factor solve produced non-finite values".to_string());
499 }
500 Ok(out)
501 }
502
503 fn solvemulti(&self, rhs: &Array2<f64>) -> Result<Array2<f64>, String> {
504 let mut out = Array2::<f64>::zeros(rhs.raw_dim());
505 for j in 0..rhs.ncols() {
506 for i in 0..rhs.nrows() {
507 out[[i, j]] = rhs[[i, j]];
508 }
509 }
510 let mut out_mat = array2_to_matmut(&mut out);
511 self.solve_in_place(out_mat.as_mut());
512 if !out.iter().all(|v| v.is_finite()) {
513 return Err("symmetric factor multi-solve produced non-finite values".to_string());
514 }
515 Ok(out)
516 }
517
518 fn logdet(&self) -> f64 {
519 match self {
520 FaerSymmetricFactor::Llt(f) => cholesky_factor_logdet(f.L()),
521 FaerSymmetricFactor::Ldlt(f) => diagonal_log_sum(f.D()),
522 FaerSymmetricFactor::Lblt(..) => {
523 f64::NAN
527 }
528 }
529 }
530}
531
532#[inline]
534pub fn factorize_symmetricwith_fallback(
535 matrix: MatRef<'_, f64>,
536 side: Side,
537) -> Result<FaerSymmetricFactor, FaerLinalgError> {
538 if let Ok(llt) = FaerLlt::new(matrix, side) {
539 return Ok(FaerSymmetricFactor::Llt(llt));
540 }
541 let ldlt_err = match FaerLdlt::new(matrix, side) {
542 Ok(ldlt) => return Ok(FaerSymmetricFactor::Ldlt(ldlt)),
543 Err(err) => err,
544 };
545 let lblt = catch_unwind(AssertUnwindSafe(|| FaerLblt::new(matrix, side)))
546 .map_err(|_| FaerLinalgError::Ldlt(ldlt_err))?;
547 Ok(FaerSymmetricFactor::Lblt(lblt))
548}
549
550#[inline]
551const fn should_use_faer_matmul(m: usize, n: usize, k: usize) -> bool {
552 const MIN_DIM: usize = 32;
556 const MIN_FLOP_SCALE: usize = 64 * 64;
557 (m >= MIN_DIM || n >= MIN_DIM || k >= MIN_DIM)
558 && m.saturating_mul(n).saturating_mul(k) >= MIN_FLOP_SCALE
559}
560
561#[inline]
562pub fn matmul_parallelism(m: usize, n: usize, k: usize) -> Par {
563 const PAR_MIN_FLOP_SCALE: usize = 2_000_000;
567 const PAR_MIN_LONG_DIM: usize = 256;
568 let flop_scale = m.saturating_mul(n).saturating_mul(k);
569 let long_dim = m.max(n).max(k);
570 if flop_scale >= PAR_MIN_FLOP_SCALE && long_dim >= PAR_MIN_LONG_DIM {
571 effective_global_parallelism()
575 } else {
576 Par::Seq
577 }
578}
579
580#[inline]
581pub fn array2_to_matmut(array: &mut Array2<f64>) -> MatMut<'_, f64> {
582 let (rows, cols) = array.dim();
583 let strides = array.strides();
584
585 let s0 = strides[0];
592 let s1 = strides[1];
593
594 unsafe { MatMut::from_raw_parts_mut(array.as_mut_ptr(), rows, cols, s0, s1) }
598}
599
600pub fn array2_to_nested_vec(array: &Array2<f64>) -> Vec<Vec<f64>> {
603 array.rows().into_iter().map(|row| row.to_vec()).collect()
604}
605
606#[inline]
607pub fn array1_to_col_matmut(array: &mut Array1<f64>) -> MatMut<'_, f64> {
608 let len = array.len();
609 let stride = array.strides()[0];
610 unsafe {
614 MatMut::from_raw_parts_mut(
615 array.as_mut_ptr(),
616 len,
617 1,
618 stride,
619 0, )
621 }
622}
623
624#[inline]
631pub fn fast_ata<S: Data<Elem = f64>>(a: &ArrayBase<S, Ix2>) -> Array2<f64> {
632 let p = a.ncols();
633 let mut out = Array2::<f64>::zeros((p, p));
634 fast_ata_into(a, &mut out);
635 out
636}
637
638#[inline]
641pub fn fast_ata_into<S: Data<Elem = f64>>(a: &ArrayBase<S, Ix2>, out: &mut Array2<f64>) {
642 use faer::Accum;
643 use faer::linalg::matmul::triangular::{BlockStructure, matmul as tri_matmul};
644
645 let (n, p) = a.dim();
646 assert_eq!(out.nrows(), p, "output rows must match p");
647 assert_eq!(out.ncols(), p, "output cols must match p");
648
649 if !should_use_faer_matmul(p, p, n) {
650 out.assign(&a.t().dot(a));
651 return;
652 }
653
654 let mut outview = array2_to_matmut(out);
655
656 let aview = FaerArrayView::new(a);
657 let a_ref = aview.as_ref();
658 let a_t = a_ref.transpose();
659 let par = matmul_parallelism(p, p, n);
660 tri_matmul(
661 outview.as_mut(),
662 BlockStructure::TriangularLower,
663 Accum::Replace,
664 a_t,
665 BlockStructure::Rectangular,
666 a_ref,
667 BlockStructure::Rectangular,
668 1.0,
669 par,
670 );
671 for i in 0..p {
673 for j in (i + 1)..p {
674 out[[i, j]] = out[[j, i]];
675 }
676 }
677}
678
679#[inline]
683pub fn fast_atb<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
684 a: &ArrayBase<S1, Ix2>,
685 b: &ArrayBase<S2, Ix2>,
686) -> Array2<f64> {
687 if let Some(out) =
688 crate::gpu_hook::gpu_dispatch().and_then(|d| d.try_fast_atb(a.view(), b.view()))
689 {
690 return out;
691 }
692 let (n_a, p) = a.dim();
693 let q = b.ncols();
694 fast_atb_with_parallelism(a, b, matmul_parallelism(p, q, n_a))
695}
696
697#[inline]
700pub fn fast_atb_with_parallelism<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
701 a: &ArrayBase<S1, Ix2>,
702 b: &ArrayBase<S2, Ix2>,
703 par: Par,
704) -> Array2<f64> {
705 use faer::linalg::matmul::matmul;
706 use faer::{Accum, Mat};
707
708 let (n_a, p) = a.dim();
709 let (n_b, q) = b.dim();
710 assert_eq!(n_a, n_b, "A and B must have same number of rows");
711
712 if !should_use_faer_matmul(p, q, n_a) {
714 return a.t().dot(b);
715 }
716
717 let mut result = Mat::<f64>::zeros(p, q);
718
719 let aview = FaerArrayView::new(a);
720 let bview = FaerArrayView::new(b);
721 let a_ref = aview.as_ref();
722 let b_ref = bview.as_ref();
723
724 matmul(
726 result.as_mut(),
727 Accum::Replace,
728 a_ref.transpose(),
729 b_ref,
730 1.0,
731 par,
732 );
733
734 mat_to_array(result.as_ref())
735}
736
737#[inline]
740pub fn fast_abt<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
741 a: &ArrayBase<S1, Ix2>,
742 b: &ArrayBase<S2, Ix2>,
743) -> Array2<f64> {
744 use faer::linalg::matmul::matmul;
745 use faer::{Accum, Mat};
746
747 let (m, k_a) = a.dim();
748 let (n, k_b) = b.dim();
749 assert_eq!(
750 k_a, k_b,
751 "A and B must have same number of columns for A·Bᵀ"
752 );
753
754 if !should_use_faer_matmul(m, n, k_a) {
755 return a.dot(&b.t());
756 }
757
758 let mut result = Mat::<f64>::zeros(m, n);
759 let aview = FaerArrayView::new(a);
760 let bview = FaerArrayView::new(b);
761 let par = matmul_parallelism(m, n, k_a);
762 matmul(
763 result.as_mut(),
764 Accum::Replace,
765 aview.as_ref(),
766 bview.as_ref().transpose(),
767 1.0,
768 par,
769 );
770 mat_to_array(result.as_ref())
771}
772
773#[inline]
777pub fn fast_ab<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
778 a: &ArrayBase<S1, Ix2>,
779 b: &ArrayBase<S2, Ix2>,
780) -> Array2<f64> {
781 if let Some(out) =
782 crate::gpu_hook::gpu_dispatch().and_then(|d| d.try_fast_ab(a.view(), b.view()))
783 {
784 return out;
785 }
786 let n = a.nrows();
787 let q = b.ncols();
788 let mut out = Array2::<f64>::zeros((n, q));
789 fast_ab_into(a, b, &mut out);
790 out
791}
792
793const FMA_LANES: usize = 8;
827
828const KERNEL_PAR_MIN_FLOP: usize = 1 << 18; const AV_PAR_MAX_CHUNK_ROWS: usize = 1024;
837
838const ATV_BLOCK_ROWS: usize = 512;
843
844#[inline]
845fn kernel_should_parallelize(n: usize, p: usize) -> bool {
846 !in_nested_parallel_region()
847 && n.saturating_mul(p) >= KERNEL_PAR_MIN_FLOP
848 && rayon::current_num_threads() > 1
849}
850
851#[inline]
852fn av_parallel_chunk_rows(p: usize) -> usize {
853 KERNEL_PAR_MIN_FLOP
854 .div_ceil(p.max(1))
855 .clamp(64, AV_PAR_MAX_CHUNK_ROWS)
856}
857
858#[inline(always)]
873fn fma_dot_body(a: &[f64], b: &[f64]) -> f64 {
874 assert_eq!(a.len(), b.len(), "fma_dot: operand length mismatch");
875 let mut sum = [0.0f64; FMA_LANES];
876 let mut comp = [0.0f64; FMA_LANES];
877 let mut ca = a.chunks_exact(FMA_LANES);
878 let mut cb = b.chunks_exact(FMA_LANES);
879 for (xa, xb) in ca.by_ref().zip(cb.by_ref()) {
880 for l in 0..FMA_LANES {
881 let x = xa[l];
882 let y = xb[l];
883 let p = x * y;
885 let ep = x.mul_add(y, -p);
886 let s = sum[l] + p;
888 let bb = s - sum[l];
889 let es = (sum[l] - (s - bb)) + (p - bb);
890 sum[l] = s;
891 comp[l] += ep + es;
892 }
893 }
894 let mut sr = 0.0f64;
896 let mut cr = 0.0f64;
897 for (&x, &y) in ca.remainder().iter().zip(cb.remainder().iter()) {
898 let p = x * y;
899 let ep = x.mul_add(y, -p);
900 let s = sr + p;
901 let bb = s - sr;
902 let es = (sr - (s - bb)) + (p - bb);
903 sr = s;
904 cr += ep + es;
905 }
906 let mut total = sr + cr;
908 for l in 0..FMA_LANES {
909 total += sum[l] + comp[l];
910 }
911 total
912}
913
914#[cfg(target_arch = "x86_64")]
920#[inline]
921fn fma_avx2_available() -> bool {
922 std::arch::is_x86_feature_detected!("fma") && std::arch::is_x86_feature_detected!("avx2")
923}
924
925#[cfg(target_arch = "x86_64")]
928#[target_feature(enable = "fma,avx2")]
929fn fma_dot_fma_avx2(a: &[f64], b: &[f64]) -> f64 {
930 fma_dot_body(a, b)
931}
932
933#[inline]
938fn fma_dot(a: &[f64], b: &[f64]) -> f64 {
939 #[cfg(target_arch = "x86_64")]
940 if fma_avx2_available() {
941 return unsafe { fma_dot_fma_avx2(a, b) };
944 }
945 fma_dot_body(a, b)
946}
947
948fn fast_av_rowmajor_into(x_all: &[f64], v: &[f64], n: usize, p: usize, out: &mut [f64]) {
952 assert_eq!(x_all.len(), n * p, "fast_av_rowmajor_into: x_all length");
953 assert_eq!(v.len(), p, "fast_av_rowmajor_into: v length");
954 assert_eq!(out.len(), n, "fast_av_rowmajor_into: out length");
955 if kernel_should_parallelize(n, p) {
956 use rayon::prelude::*;
957 let chunk_rows = av_parallel_chunk_rows(p);
958 out.par_chunks_mut(chunk_rows)
959 .enumerate()
960 .for_each(|(c, chunk)| {
961 let base = c * chunk_rows;
962 for (k, o) in chunk.iter_mut().enumerate() {
963 let i = base + k;
964 *o = fma_dot(&x_all[i * p..i * p + p], v);
965 }
966 });
967 } else {
968 for (i, o) in out.iter_mut().enumerate() {
969 *o = fma_dot(&x_all[i * p..i * p + p], v);
970 }
971 }
972}
973
974#[inline(always)]
983fn standard_fma_dot_body(a: &[f64], b: &[f64]) -> f64 {
984 assert_eq!(
985 a.len(),
986 b.len(),
987 "standard_fma_dot: operand length mismatch"
988 );
989 let mut sum = [0.0_f64; FMA_LANES];
990 let mut ca = a.chunks_exact(FMA_LANES);
991 let mut cb = b.chunks_exact(FMA_LANES);
992 for (xa, xb) in ca.by_ref().zip(cb.by_ref()) {
993 for lane in 0..FMA_LANES {
994 sum[lane] = xa[lane].mul_add(xb[lane], sum[lane]);
995 }
996 }
997 let mut remainder = 0.0;
998 for (&x, &y) in ca.remainder().iter().zip(cb.remainder().iter()) {
999 remainder = x.mul_add(y, remainder);
1000 }
1001 let pair01 = sum[0] + sum[1];
1002 let pair23 = sum[2] + sum[3];
1003 let pair45 = sum[4] + sum[5];
1004 let pair67 = sum[6] + sum[7];
1005 remainder + (pair01 + pair23) + (pair45 + pair67)
1006}
1007
1008#[cfg(target_arch = "x86_64")]
1010#[target_feature(enable = "fma,avx2")]
1011fn standard_fma_dot_fma_avx2(a: &[f64], b: &[f64]) -> f64 {
1012 standard_fma_dot_body(a, b)
1013}
1014
1015#[inline]
1018fn standard_fma_dot(a: &[f64], b: &[f64]) -> f64 {
1019 #[cfg(target_arch = "x86_64")]
1020 if fma_avx2_available() {
1021 return unsafe { standard_fma_dot_fma_avx2(a, b) };
1024 }
1025 standard_fma_dot_body(a, b)
1026}
1027
1028fn standard_av_rowmajor_into(x_all: &[f64], v: &[f64], n: usize, p: usize, out: &mut [f64]) {
1029 assert_eq!(
1030 x_all.len(),
1031 n * p,
1032 "standard_av_rowmajor_into: matrix length"
1033 );
1034 assert_eq!(v.len(), p, "standard_av_rowmajor_into: vector length");
1035 assert_eq!(out.len(), n, "standard_av_rowmajor_into: output length");
1036 if kernel_should_parallelize(n, p) {
1037 use rayon::prelude::*;
1038 let chunk_rows = av_parallel_chunk_rows(p);
1039 out.par_chunks_mut(chunk_rows)
1040 .enumerate()
1041 .for_each(|(chunk_index, chunk)| {
1042 let base = chunk_index * chunk_rows;
1043 for (offset, output) in chunk.iter_mut().enumerate() {
1044 let row = base + offset;
1045 *output = standard_fma_dot(&x_all[row * p..row * p + p], v);
1046 }
1047 });
1048 } else {
1049 for (row, output) in out.iter_mut().enumerate() {
1050 *output = standard_fma_dot(&x_all[row * p..row * p + p], v);
1051 }
1052 }
1053}
1054
1055fn pairwise_sum_into(parts: &[Vec<f64>], out: &mut [f64]) {
1057 match parts.len() {
1058 0 => out.fill(0.0),
1059 1 => out.copy_from_slice(&parts[0]),
1060 _ => {
1061 let mid = parts.len() / 2;
1062 let p = out.len();
1063 let mut left = vec![0.0f64; p];
1064 let mut right = vec![0.0f64; p];
1065 pairwise_sum_into(&parts[..mid], &mut left);
1066 pairwise_sum_into(&parts[mid..], &mut right);
1067 for ((o, &l), &r) in out.iter_mut().zip(left.iter()).zip(right.iter()) {
1068 *o = l + r;
1069 }
1070 }
1071 }
1072}
1073
1074fn fast_atv_rowmajor_into(x_all: &[f64], v: &[f64], n: usize, p: usize, out: &mut [f64]) {
1082 assert_eq!(x_all.len(), n * p, "fast_atv_rowmajor_into: x_all length");
1083 assert_eq!(v.len(), n, "fast_atv_rowmajor_into: v length");
1084 assert_eq!(out.len(), p, "fast_atv_rowmajor_into: out length");
1085 let nblocks = n.div_ceil(ATV_BLOCK_ROWS);
1086
1087 let block_partial = |b: usize| -> Vec<f64> {
1088 let start = b * ATV_BLOCK_ROWS;
1089 let end = (start + ATV_BLOCK_ROWS).min(n);
1090 let mut acc = vec![0.0f64; p];
1091 atv_block_accumulate(&x_all[start * p..end * p], &v[start..end], &mut acc);
1092 acc
1093 };
1094
1095 let partials: Vec<Vec<f64>> = if kernel_should_parallelize(n, p) {
1096 use rayon::prelude::*;
1097 (0..nblocks).into_par_iter().map(block_partial).collect()
1098 } else {
1099 (0..nblocks).map(block_partial).collect()
1100 };
1101
1102 pairwise_sum_into(&partials, out);
1103}
1104
1105#[inline(always)]
1109fn atv_block_accumulate_body(rows: &[f64], v: &[f64], acc: &mut [f64]) {
1110 let p = acc.len();
1111 assert_eq!(rows.len(), v.len() * p, "atv_block_accumulate: block length");
1112 for (&vi, row) in v.iter().zip(rows.chunks_exact(p)) {
1113 for (a, &xij) in acc.iter_mut().zip(row.iter()) {
1114 *a = xij.mul_add(vi, *a);
1115 }
1116 }
1117}
1118
1119#[cfg(target_arch = "x86_64")]
1121#[target_feature(enable = "fma,avx2")]
1122fn atv_block_accumulate_fma_avx2(rows: &[f64], v: &[f64], acc: &mut [f64]) {
1123 atv_block_accumulate_body(rows, v, acc)
1124}
1125
1126#[inline]
1129fn atv_block_accumulate(rows: &[f64], v: &[f64], acc: &mut [f64]) {
1130 #[cfg(target_arch = "x86_64")]
1131 if fma_avx2_available() {
1132 return unsafe { atv_block_accumulate_fma_avx2(rows, v, acc) };
1135 }
1136 atv_block_accumulate_body(rows, v, acc)
1137}
1138
1139pub(crate) fn fma_axpy_into(alpha: f64, x: &[f64], y: &mut [f64]) {
1146 #[cfg(target_arch = "x86_64")]
1147 if fma_avx2_available() {
1148 return unsafe { fma_axpy_into_fma_avx2(alpha, x, y) };
1151 }
1152 fma_axpy_into_body(alpha, x, y)
1153}
1154
1155#[inline(always)]
1156fn fma_axpy_into_body(alpha: f64, x: &[f64], y: &mut [f64]) {
1157 assert_eq!(x.len(), y.len(), "fma_axpy_into: operand length mismatch");
1158 for (yi, &xi) in y.iter_mut().zip(x.iter()) {
1159 *yi = alpha.mul_add(xi, *yi);
1160 }
1161}
1162
1163#[cfg(target_arch = "x86_64")]
1165#[target_feature(enable = "fma,avx2")]
1166fn fma_axpy_into_fma_avx2(alpha: f64, x: &[f64], y: &mut [f64]) {
1167 fma_axpy_into_body(alpha, x, y)
1168}
1169
1170#[inline]
1173pub fn fast_av<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
1174 a: &ArrayBase<S1, Ix2>,
1175 v: &ArrayBase<S2, Ix1>,
1176) -> Array1<f64> {
1177 if let Some(out) =
1178 crate::gpu_hook::gpu_dispatch().and_then(|d| d.try_fast_av(a.view(), v.view()))
1179 {
1180 return out;
1181 }
1182 fast_av_impl(a, v)
1183}
1184
1185#[inline]
1186fn fast_av_impl<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
1187 a: &ArrayBase<S1, Ix2>,
1188 v: &ArrayBase<S2, Ix1>,
1189) -> Array1<f64> {
1190 use faer::linalg::matmul::matmul;
1191 use faer::{Accum, Mat};
1192
1193 let (n, p) = a.dim();
1194 assert_eq!(p, v.len(), "A cols must match v length");
1195
1196 if let (Some(x_all), Some(vs)) = (a.as_slice(), v.as_slice())
1200 && n != 0
1201 && p != 0
1202 {
1203 let mut out = Array1::<f64>::zeros(n);
1204 fast_av_rowmajor_into(
1205 x_all,
1206 vs,
1207 n,
1208 p,
1209 out.as_slice_mut().expect("fresh Array1 is contiguous"),
1210 );
1211 return out;
1212 }
1213
1214 if !should_use_faer_matmul(n, 1, p) {
1215 return a.dot(v);
1216 }
1217
1218 let mut result = Mat::<f64>::zeros(n, 1);
1219
1220 let aview = FaerArrayView::new(a);
1221 let vview = FaerColView::new(v);
1222 let a_ref = aview.as_ref();
1223 let v_ref = vview.as_ref();
1224
1225 let par = matmul_parallelism(n, 1, p);
1226 matmul(result.as_mut(), Accum::Replace, a_ref, v_ref, 1.0, par);
1227
1228 let mut out = Array1::<f64>::zeros(n);
1229 for i in 0..n {
1230 out[i] = result[(i, 0)];
1231 }
1232 out
1233}
1234
1235#[inline]
1238pub fn fast_av_into<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
1239 a: &ArrayBase<S1, Ix2>,
1240 v: &ArrayBase<S2, Ix1>,
1241 out: &mut Array1<f64>,
1242) {
1243 fast_av_into_impl(a, v, out);
1244}
1245
1246#[inline]
1247fn fast_av_into_impl<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
1248 a: &ArrayBase<S1, Ix2>,
1249 v: &ArrayBase<S2, Ix1>,
1250 out: &mut Array1<f64>,
1251) {
1252 use faer::Accum;
1253 use faer::linalg::matmul::matmul;
1254
1255 let (n, p) = a.dim();
1256 assert_eq!(v.len(), p, "vector length must match A cols");
1257 assert_eq!(out.len(), n, "output length must match A rows");
1258
1259 if let (Some(x_all), Some(vs)) = (a.as_slice(), v.as_slice())
1260 && n != 0
1261 && p != 0
1262 && let Some(out_s) = out.as_slice_mut()
1263 {
1264 fast_av_rowmajor_into(x_all, vs, n, p, out_s);
1265 return;
1266 }
1267
1268 if !should_use_faer_matmul(n, 1, p) {
1269 out.assign(&a.dot(v));
1270 return;
1271 }
1272
1273 let mut outview = array1_to_col_matmut(out);
1274
1275 let aview = FaerArrayView::new(a);
1276 let vview = FaerColView::new(v);
1277 let a_ref = aview.as_ref();
1278 let v_ref = vview.as_ref();
1279 let par = matmul_parallelism(n, 1, p);
1280 matmul(outview.as_mut(), Accum::Replace, a_ref, v_ref, 1.0, par);
1281}
1282
1283#[inline]
1290pub fn fast_av_view_into<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
1291 a: &ArrayBase<S1, Ix2>,
1292 v: &ArrayBase<S2, Ix1>,
1293 out: ArrayViewMut1<'_, f64>,
1294) {
1295 fast_av_view_into_impl(a, v, out);
1296}
1297
1298pub fn fast_av_standard_view_into<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
1306 a: &ArrayBase<S1, Ix2>,
1307 v: &ArrayBase<S2, Ix1>,
1308 mut out: ArrayViewMut1<'_, f64>,
1309) {
1310 use faer::Accum;
1311 use faer::linalg::matmul::matmul;
1312
1313 let (n, p) = a.dim();
1314 assert_eq!(v.len(), p, "vector length must match A cols");
1315 assert_eq!(out.len(), n, "output length must match A rows");
1316 if let (Some(x_all), Some(vs), Some(out_slice)) =
1317 (a.as_slice(), v.as_slice(), out.as_slice_mut())
1318 && n != 0
1319 && p != 0
1320 {
1321 standard_av_rowmajor_into(x_all, vs, n, p, out_slice);
1322 return;
1323 }
1324 if !should_use_faer_matmul(n, 1, p) {
1325 out.assign(&a.dot(v));
1326 return;
1327 }
1328
1329 let len = out.len();
1330 let stride = out.strides()[0];
1331 let outview = unsafe { MatMut::from_raw_parts_mut(out.as_mut_ptr(), len, 1, stride, 0) };
1334 let aview = FaerArrayView::new(a);
1335 let vview = FaerColView::new(v);
1336 matmul(
1337 outview,
1338 Accum::Replace,
1339 aview.as_ref(),
1340 vview.as_ref(),
1341 1.0,
1342 matmul_parallelism(n, 1, p),
1343 );
1344}
1345
1346#[inline]
1347fn fast_av_view_into_impl<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
1348 a: &ArrayBase<S1, Ix2>,
1349 v: &ArrayBase<S2, Ix1>,
1350 mut out: ArrayViewMut1<'_, f64>,
1351) {
1352 use faer::Accum;
1353 use faer::linalg::matmul::matmul;
1354
1355 let (n, p) = a.dim();
1356 assert_eq!(v.len(), p, "vector length must match A cols");
1357 assert_eq!(out.len(), n, "output length must match A rows");
1358
1359 if let (Some(x_all), Some(vs)) = (a.as_slice(), v.as_slice())
1360 && n != 0
1361 && p != 0
1362 && let Some(out_s) = out.as_slice_mut()
1363 {
1364 fast_av_rowmajor_into(x_all, vs, n, p, out_s);
1365 return;
1366 }
1367
1368 if !should_use_faer_matmul(n, 1, p) {
1369 let prod = a.dot(v);
1370 out.assign(&prod);
1371 return;
1372 }
1373
1374 let len = out.len();
1375 let stride = out.strides()[0];
1376 let outview = unsafe {
1380 MatMut::from_raw_parts_mut(
1381 out.as_mut_ptr(),
1382 len,
1383 1,
1384 stride,
1385 0, )
1387 };
1388
1389 let aview = FaerArrayView::new(a);
1390 let vview = FaerColView::new(v);
1391 let a_ref = aview.as_ref();
1392 let v_ref = vview.as_ref();
1393 let par = matmul_parallelism(n, 1, p);
1394 matmul(outview, Accum::Replace, a_ref, v_ref, 1.0, par);
1395}
1396
1397#[inline]
1400pub fn fast_atv<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
1401 a: &ArrayBase<S1, Ix2>,
1402 v: &ArrayBase<S2, Ix1>,
1403) -> Array1<f64> {
1404 if let Some(out) =
1405 crate::gpu_hook::gpu_dispatch().and_then(|d| d.try_fast_atv(a.view(), v.view()))
1406 {
1407 return out;
1408 }
1409 fast_atv_impl(a, v)
1410}
1411
1412#[inline]
1413fn fast_atv_impl<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
1414 a: &ArrayBase<S1, Ix2>,
1415 v: &ArrayBase<S2, Ix1>,
1416) -> Array1<f64> {
1417 use faer::Accum;
1418 use faer::linalg::matmul::matmul;
1419
1420 let (n, p) = a.dim();
1421 assert_eq!(n, v.len(), "A rows must match v length");
1422
1423 if let (Some(x_all), Some(vs)) = (a.as_slice(), v.as_slice())
1427 && n != 0
1428 && p != 0
1429 {
1430 let mut out = Array1::<f64>::zeros(p);
1431 fast_atv_rowmajor_into(
1432 x_all,
1433 vs,
1434 n,
1435 p,
1436 out.as_slice_mut().expect("fresh Array1 is contiguous"),
1437 );
1438 return out;
1439 }
1440
1441 if !should_use_faer_matmul(p, 1, n) {
1443 return a.t().dot(v);
1444 }
1445
1446 let mut out = Array1::<f64>::zeros(p);
1447 let mut outview = array1_to_col_matmut(&mut out);
1448
1449 let aview = FaerArrayView::new(a);
1450 let vview = FaerColView::new(v);
1451 let a_ref = aview.as_ref();
1452 let v_ref = vview.as_ref();
1453
1454 let par = matmul_parallelism(p, 1, n);
1456 matmul(
1457 outview.as_mut(),
1458 Accum::Replace,
1459 a_ref.transpose(),
1460 v_ref,
1461 1.0,
1462 par,
1463 );
1464
1465 out
1466}
1467
1468#[inline]
1471pub fn fast_atv_into<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
1472 a: &ArrayBase<S1, Ix2>,
1473 v: &ArrayBase<S2, Ix1>,
1474 out: &mut Array1<f64>,
1475) {
1476 fast_atv_into_impl(a, v, out);
1477}
1478
1479#[inline]
1480fn fast_atv_into_impl<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
1481 a: &ArrayBase<S1, Ix2>,
1482 v: &ArrayBase<S2, Ix1>,
1483 out: &mut Array1<f64>,
1484) {
1485 use faer::Accum;
1486 use faer::linalg::matmul::matmul;
1487
1488 let (n, p) = a.dim();
1489 assert_eq!(v.len(), n, "vector length must match A rows");
1490 assert_eq!(out.len(), p, "output length must match A cols");
1491
1492 if let (Some(x_all), Some(vs)) = (a.as_slice(), v.as_slice())
1493 && n != 0
1494 && p != 0
1495 && let Some(out_s) = out.as_slice_mut()
1496 {
1497 fast_atv_rowmajor_into(x_all, vs, n, p, out_s);
1498 return;
1499 }
1500
1501 if !should_use_faer_matmul(p, 1, n) {
1502 out.assign(&a.t().dot(v));
1503 return;
1504 }
1505
1506 let mut outview = array1_to_col_matmut(out);
1507
1508 let aview = FaerArrayView::new(a);
1509 let vview = FaerColView::new(v);
1510 let a_ref = aview.as_ref();
1511 let v_ref = vview.as_ref();
1512 let par = matmul_parallelism(p, 1, n);
1513 matmul(
1514 outview.as_mut(),
1515 Accum::Replace,
1516 a_ref.transpose(),
1517 v_ref,
1518 1.0,
1519 par,
1520 );
1521}
1522
1523#[inline]
1525pub fn fast_xt_diag_x<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
1526 x: &ArrayBase<S1, Ix2>,
1527 w: &ArrayBase<S2, Ix1>,
1528) -> Array2<f64> {
1529 assert_eq!(
1530 x.nrows(),
1531 w.len(),
1532 "fast_xt_diag_x row/weight length mismatch"
1533 );
1534 if let Some(out) =
1535 crate::gpu_hook::gpu_dispatch().and_then(|d| d.try_fast_xt_diag_x(x.view(), w.view()))
1536 {
1537 return out;
1538 }
1539 let p = x.ncols();
1540 fast_xt_diag_x_with_parallelism(x, w, matmul_parallelism(p, p, x.nrows()))
1541}
1542
1543#[inline]
1546pub fn fast_xt_diag_x_with_parallelism<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
1547 x: &ArrayBase<S1, Ix2>,
1548 w: &ArrayBase<S2, Ix1>,
1549 par: Par,
1550) -> Array2<f64> {
1551 assert_eq!(
1552 x.nrows(),
1553 w.len(),
1554 "fast_xt_diag_x_with_parallelism row/weight length mismatch"
1555 );
1556 fast_xt_diag_x_with_parallelism_impl(x, w, par)
1557}
1558
1559#[inline]
1560fn fast_xt_diag_x_with_parallelism_impl<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
1561 x: &ArrayBase<S1, Ix2>,
1562 w: &ArrayBase<S2, Ix1>,
1563 par: Par,
1564) -> Array2<f64> {
1565 use ndarray::ShapeBuilder;
1566
1567 let p = x.ncols();
1568 let mut result = Array2::<f64>::zeros((p, p).f());
1571 stream_weighted_crossprod_into(
1572 x,
1573 w,
1574 &mut result,
1575 CrossprodStructure::SymmetricLower,
1576 CrossprodAccum::Replace,
1577 par,
1578 );
1579 result
1580}
1581
1582#[derive(Clone, Copy, PartialEq, Eq, Debug)]
1584pub enum CrossprodStructure {
1585 Full,
1587 SymmetricLower,
1591}
1592
1593#[derive(Clone, Copy, PartialEq, Eq, Debug)]
1595pub enum CrossprodAccum {
1596 Replace,
1598 Add,
1600}
1601
1602#[inline]
1622fn streaming_chunk_rows(cols: usize, n: usize) -> usize {
1623 const TARGET_BYTES: usize = gam_runtime::resource::LIBRARY_ROW_CHUNK_TARGET_BYTES;
1631 const MIN_ROWS: usize = 512;
1632 const MAX_ROWS: usize = 131_072;
1633 (TARGET_BYTES / (cols.max(1) * std::mem::size_of::<f64>()))
1634 .clamp(MIN_ROWS, MAX_ROWS)
1635 .min(n)
1636}
1637
1638pub fn stream_weighted_crossprod_into<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
1657 x: &ArrayBase<S1, Ix2>,
1658 w: &ArrayBase<S2, Ix1>,
1659 out: &mut Array2<f64>,
1660 structure: CrossprodStructure,
1661 accum: CrossprodAccum,
1662 par: Par,
1663) {
1664 use faer::Accum;
1665 use faer::linalg::matmul::matmul;
1666 use faer::linalg::matmul::triangular::{BlockStructure, matmul as tri_matmul};
1667 use ndarray::s;
1668
1669 let (n, p) = x.dim();
1670 assert_eq!(n, w.len(), "X rows must match W length");
1671 assert_eq!(out.nrows(), p, "output rows must match X cols");
1672 assert_eq!(out.ncols(), p, "output cols must match X cols");
1673 if p == 0 {
1674 return;
1675 }
1676 if n == 0 {
1677 if accum == CrossprodAccum::Replace {
1678 out.fill(0.0);
1679 }
1680 return;
1681 }
1682
1683 if !should_use_faer_matmul(p, p, n) {
1684 let w_x = Array2::from_shape_fn((n, p), |(i, j)| w[i] * x[[i, j]]);
1686 let gram = x.t().dot(&w_x);
1687 match accum {
1688 CrossprodAccum::Replace => out.assign(&gram),
1689 CrossprodAccum::Add => *out += &gram,
1690 }
1691 return;
1692 }
1693
1694 let chunk_rows = streaming_chunk_rows(p, n);
1696
1697 if accum == CrossprodAccum::Replace {
1702 out.fill(0.0);
1703 }
1704
1705 let mut wx_chunk = Array2::<f64>::zeros((chunk_rows, p));
1711
1712 let x_is_row_major = x.is_standard_layout();
1713 let w_slice_opt = w.as_slice();
1714
1715 {
1718 let mut out_view = array2_to_matmut(out);
1719 for start in (0..n).step_by(chunk_rows) {
1720 let rows = (n - start).min(chunk_rows);
1721 {
1722 let chunk_slice = wx_chunk
1723 .as_slice_mut()
1724 .expect("row-major chunk is contiguous");
1725 if x_is_row_major && let (Some(x_all), Some(w_all)) = (x.as_slice(), w_slice_opt) {
1726 for local in 0..rows {
1727 let src = start + local;
1728 let wi = w_all[src];
1729 let src_off = src * p;
1730 let dst_off = local * p;
1731 let src_row = &x_all[src_off..src_off + p];
1732 let dst_row = &mut chunk_slice[dst_off..dst_off + p];
1733 for col in 0..p {
1734 dst_row[col] = src_row[col] * wi;
1735 }
1736 }
1737 } else {
1738 let x_slice = x.slice(s![start..start + rows, ..]);
1739 for local in 0..rows {
1740 let wi = w[start + local];
1741 let xrow = x_slice.row(local);
1742 let dst_off = local * p;
1743 let dst_row = &mut chunk_slice[dst_off..dst_off + p];
1744 for (col, xij) in xrow.iter().enumerate() {
1745 dst_row[col] = xij * wi;
1746 }
1747 }
1748 }
1749 }
1750 let x_slice = x.slice(s![start..start + rows, ..]);
1751 let wx_slice = wx_chunk.slice(s![0..rows, ..]);
1752 let x_view = FaerArrayView::new(&x_slice);
1753 let wx_view = FaerArrayView::new(&wx_slice);
1754 match structure {
1755 CrossprodStructure::SymmetricLower => {
1756 tri_matmul(
1760 out_view.as_mut(),
1761 BlockStructure::TriangularLower,
1762 Accum::Add,
1763 x_view.as_ref().transpose(),
1764 BlockStructure::Rectangular,
1765 wx_view.as_ref(),
1766 BlockStructure::Rectangular,
1767 1.0,
1768 par,
1769 );
1770 }
1771 CrossprodStructure::Full => {
1772 matmul(
1773 out_view.as_mut(),
1774 Accum::Add,
1775 x_view.as_ref().transpose(),
1776 wx_view.as_ref(),
1777 1.0,
1778 par,
1779 );
1780 }
1781 }
1782 }
1783 }
1784
1785 if structure == CrossprodStructure::SymmetricLower {
1786 for i in 0..p {
1788 for j in (i + 1)..p {
1789 out[[i, j]] = out[[j, i]];
1790 }
1791 }
1792 }
1793}
1794
1795#[inline]
1797pub fn fast_xt_diag_y<S1: Data<Elem = f64>, S2: Data<Elem = f64>, S3: Data<Elem = f64>>(
1798 x: &ArrayBase<S1, Ix2>,
1799 w: &ArrayBase<S2, Ix1>,
1800 y: &ArrayBase<S3, Ix2>,
1801) -> Array2<f64> {
1802 assert_eq!(x.nrows(), y.nrows(), "fast_xt_diag_y X/Y row mismatch");
1803 assert_eq!(
1804 y.nrows(),
1805 w.len(),
1806 "fast_xt_diag_y row/weight length mismatch"
1807 );
1808 if let Some(out) = crate::gpu_hook::gpu_dispatch()
1809 .and_then(|d| d.try_fast_xt_diag_y(x.view(), w.view(), y.view()))
1810 {
1811 return out;
1812 }
1813 fast_xt_diag_y_impl(x, w, y)
1814}
1815
1816#[inline]
1817fn fast_xt_diag_y_impl<S1: Data<Elem = f64>, S2: Data<Elem = f64>, S3: Data<Elem = f64>>(
1818 x: &ArrayBase<S1, Ix2>,
1819 w: &ArrayBase<S2, Ix1>,
1820 y: &ArrayBase<S3, Ix2>,
1821) -> Array2<f64> {
1822 use faer::Accum;
1823 use faer::linalg::matmul::matmul;
1824 use ndarray::{ShapeBuilder, s};
1825
1826 let (n, q) = y.dim();
1827 let px = x.ncols();
1828 assert_eq!(n, w.len(), "Y rows must match W length");
1829 assert_eq!(n, x.nrows(), "X rows must match Y rows");
1830 if n == 0 || px == 0 || q == 0 {
1831 return Array2::<f64>::zeros((px, q));
1832 }
1833 if !should_use_faer_matmul(px, q, n) {
1834 let w_y = Array2::from_shape_fn((n, q), |(i, j)| w[i] * y[[i, j]]);
1835 return x.t().dot(&w_y);
1836 }
1837
1838 let total_cols = px + q;
1840 let chunk_rows = streaming_chunk_rows(total_cols, n);
1841
1842 let mut result = Array2::<f64>::zeros((px, q).f());
1843 let mut wy_chunk = Array2::<f64>::zeros((chunk_rows, q));
1846
1847 let y_is_row_major = y.is_standard_layout();
1848 let w_slice_opt = w.as_slice();
1849
1850 {
1851 let mut out_view = array2_to_matmut(&mut result);
1852
1853 for start in (0..n).step_by(chunk_rows) {
1854 let rows = (n - start).min(chunk_rows);
1855 {
1856 let chunk_slice = wy_chunk
1857 .as_slice_mut()
1858 .expect("row-major chunk is contiguous");
1859 if y_is_row_major && let (Some(y_all), Some(w_all)) = (y.as_slice(), w_slice_opt) {
1860 for local in 0..rows {
1861 let src = start + local;
1862 let wi = w_all[src];
1863 let src_off = src * q;
1864 let dst_off = local * q;
1865 let src_row = &y_all[src_off..src_off + q];
1866 let dst_row = &mut chunk_slice[dst_off..dst_off + q];
1867 for col in 0..q {
1868 dst_row[col] = src_row[col] * wi;
1869 }
1870 }
1871 } else {
1872 let y_slice = y.slice(s![start..start + rows, ..]);
1873 for local in 0..rows {
1874 let wi = w[start + local];
1875 let yrow = y_slice.row(local);
1876 let dst_off = local * q;
1877 let dst_row = &mut chunk_slice[dst_off..dst_off + q];
1878 for (col, yij) in yrow.iter().enumerate() {
1879 dst_row[col] = yij * wi;
1880 }
1881 }
1882 }
1883 }
1884 let x_slice = x.slice(s![start..start + rows, ..]);
1885 let wy_slice = wy_chunk.slice(s![0..rows, ..]);
1886 let x_view = FaerArrayView::new(&x_slice);
1887 let wy_view = FaerArrayView::new(&wy_slice);
1888 let par = matmul_parallelism(px, q, rows);
1889 matmul(
1890 out_view.as_mut(),
1891 Accum::Add,
1892 x_view.as_ref().transpose(),
1893 wy_view.as_ref(),
1894 1.0,
1895 par,
1896 );
1897 }
1898 }
1899
1900 result
1901}
1902
1903pub fn fast_joint_hessian_2x2<
1909 S1: Data<Elem = f64>,
1910 S2: Data<Elem = f64>,
1911 S3: Data<Elem = f64>,
1912 S4: Data<Elem = f64>,
1913 S5: Data<Elem = f64>,
1914>(
1915 x_a: &ArrayBase<S1, Ix2>,
1916 x_b: &ArrayBase<S2, Ix2>,
1917 w_aa: &ArrayBase<S3, Ix1>,
1918 w_ab: &ArrayBase<S4, Ix1>,
1919 w_bb: &ArrayBase<S5, Ix1>,
1920) -> Array2<f64> {
1921 if let Some(out) = crate::gpu_hook::gpu_dispatch().and_then(|d| {
1922 d.try_fast_joint_hessian_2x2(
1923 x_a.view(),
1924 x_b.view(),
1925 w_aa.view(),
1926 w_ab.view(),
1927 w_bb.view(),
1928 )
1929 }) {
1930 return out;
1931 }
1932 fast_joint_hessian_2x2_impl(x_a, x_b, w_aa, w_ab, w_bb)
1933}
1934
1935#[inline]
1936fn fast_joint_hessian_2x2_impl<
1937 S1: Data<Elem = f64>,
1938 S2: Data<Elem = f64>,
1939 S3: Data<Elem = f64>,
1940 S4: Data<Elem = f64>,
1941 S5: Data<Elem = f64>,
1942>(
1943 x_a: &ArrayBase<S1, Ix2>,
1944 x_b: &ArrayBase<S2, Ix2>,
1945 w_aa: &ArrayBase<S3, Ix1>,
1946 w_ab: &ArrayBase<S4, Ix1>,
1947 w_bb: &ArrayBase<S5, Ix1>,
1948) -> Array2<f64> {
1949 use faer::Accum;
1950 use faer::linalg::matmul::matmul;
1951 use ndarray::{ShapeBuilder, s};
1952
1953 let n = x_a.nrows();
1954 let pa = x_a.ncols();
1955 let pb = x_b.ncols();
1956 let total = pa + pb;
1957 assert_eq!(n, x_b.nrows());
1958 assert_eq!(n, w_aa.len());
1959 assert_eq!(n, w_ab.len());
1960 assert_eq!(n, w_bb.len());
1961
1962 if n == 0 || total == 0 {
1963 return Array2::<f64>::zeros((total, total));
1964 }
1965
1966 if !should_use_faer_matmul(pa.max(pb), pa.max(pb), n) {
1968 let waa_xa = Array2::from_shape_fn((n, pa), |(i, j)| w_aa[i] * x_a[[i, j]]);
1969 let wab_xb = Array2::from_shape_fn((n, pb), |(i, j)| w_ab[i] * x_b[[i, j]]);
1970 let wbb_xb = Array2::from_shape_fn((n, pb), |(i, j)| w_bb[i] * x_b[[i, j]]);
1971 let mut out = Array2::<f64>::zeros((total, total));
1972 out.slice_mut(s![..pa, ..pa]).assign(&x_a.t().dot(&waa_xa));
1973 out.slice_mut(s![..pa, pa..]).assign(&x_a.t().dot(&wab_xb));
1974 out.slice_mut(s![pa.., pa..]).assign(&x_b.t().dot(&wbb_xb));
1975 for i in 0..total {
1977 for j in 0..i {
1978 out[[i, j]] = out[[j, i]];
1979 }
1980 }
1981 return out;
1982 }
1983
1984 let cols_needed = pa + 2 * pb;
1986 let chunk_rows = streaming_chunk_rows(cols_needed, n);
1987
1988 let mut out = Array2::<f64>::zeros((total, total).f());
1989 let mut waa_xa_buf = Array2::<f64>::zeros((chunk_rows, pa));
1994 let mut wab_xb_buf = Array2::<f64>::zeros((chunk_rows, pb));
1995 let mut wbb_xb_buf = Array2::<f64>::zeros((chunk_rows, pb));
1996
1997 let xa_is_row_major = x_a.is_standard_layout();
1998 let xb_is_row_major = x_b.is_standard_layout();
1999 let waa_slice_opt = w_aa.as_slice();
2000 let wab_slice_opt = w_ab.as_slice();
2001 let wbb_slice_opt = w_bb.as_slice();
2002
2003 {
2004 let mut out_mat = array2_to_matmut(&mut out);
2005
2006 for start in (0..n).step_by(chunk_rows) {
2007 let rows = (n - start).min(chunk_rows);
2008 let xa_slice = x_a.slice(s![start..start + rows, ..]);
2009 let xb_slice = x_b.slice(s![start..start + rows, ..]);
2010
2011 {
2013 let waa_chunk = waa_xa_buf
2014 .as_slice_mut()
2015 .expect("row-major waa chunk is contiguous");
2016 let wab_chunk = wab_xb_buf
2017 .as_slice_mut()
2018 .expect("row-major wab chunk is contiguous");
2019 let wbb_chunk = wbb_xb_buf
2020 .as_slice_mut()
2021 .expect("row-major wbb chunk is contiguous");
2022
2023 if xa_is_row_major
2024 && xb_is_row_major
2025 && let (Some(xa_all), Some(xb_all)) = (x_a.as_slice(), x_b.as_slice())
2026 && let (Some(waa_all), Some(wab_all), Some(wbb_all)) =
2027 (waa_slice_opt, wab_slice_opt, wbb_slice_opt)
2028 {
2029 for local in 0..rows {
2030 let i = start + local;
2031 let waa_i = waa_all[i];
2032 let wab_i = wab_all[i];
2033 let wbb_i = wbb_all[i];
2034 let xa_off = i * pa;
2035 let xa_row = &xa_all[xa_off..xa_off + pa];
2036 let xb_off = i * pb;
2037 let xb_row = &xb_all[xb_off..xb_off + pb];
2038 let waa_off = local * pa;
2039 let wab_off = local * pb;
2040 let wbb_off = local * pb;
2041 let waa_row = &mut waa_chunk[waa_off..waa_off + pa];
2042 for col in 0..pa {
2043 waa_row[col] = xa_row[col] * waa_i;
2044 }
2045 let wab_row = &mut wab_chunk[wab_off..wab_off + pb];
2046 let wbb_row = &mut wbb_chunk[wbb_off..wbb_off + pb];
2047 for col in 0..pb {
2048 let xij = xb_row[col];
2049 wab_row[col] = xij * wab_i;
2050 wbb_row[col] = xij * wbb_i;
2051 }
2052 }
2053 } else {
2054 for local in 0..rows {
2055 let i = start + local;
2056 let waa_i = w_aa[i];
2057 let wab_i = w_ab[i];
2058 let wbb_i = w_bb[i];
2059 let waa_off = local * pa;
2060 let wab_off = local * pb;
2061 let wbb_off = local * pb;
2062 let waa_row = &mut waa_chunk[waa_off..waa_off + pa];
2063 let xa_row = xa_slice.row(local);
2064 for (col, xij) in xa_row.iter().enumerate() {
2065 waa_row[col] = xij * waa_i;
2066 }
2067 let wab_row = &mut wab_chunk[wab_off..wab_off + pb];
2068 let wbb_row = &mut wbb_chunk[wbb_off..wbb_off + pb];
2069 let xb_row = xb_slice.row(local);
2070 for (col, xij) in xb_row.iter().enumerate() {
2071 wab_row[col] = xij * wab_i;
2072 wbb_row[col] = xij * wbb_i;
2073 }
2074 }
2075 }
2076 }
2077
2078 let xa_view = FaerArrayView::new(&xa_slice);
2079 let xb_view = FaerArrayView::new(&xb_slice);
2080 let waa_xa_slice = waa_xa_buf.slice(s![0..rows, ..]);
2081 let wab_xb_slice = wab_xb_buf.slice(s![0..rows, ..]);
2082 let wbb_xb_slice = wbb_xb_buf.slice(s![0..rows, ..]);
2083 let waa_xa_view = FaerArrayView::new(&waa_xa_slice);
2084 let wab_xb_view = FaerArrayView::new(&wab_xb_slice);
2085 let wbb_xb_view = FaerArrayView::new(&wbb_xb_slice);
2086
2087 matmul(
2089 out_mat.rb_mut().submatrix_mut(0, 0, pa, pa),
2090 Accum::Add,
2091 xa_view.as_ref().transpose(),
2092 waa_xa_view.as_ref(),
2093 1.0,
2094 matmul_parallelism(pa, pa, rows),
2095 );
2096 matmul(
2098 out_mat.rb_mut().submatrix_mut(0, pa, pa, pb),
2099 Accum::Add,
2100 xa_view.as_ref().transpose(),
2101 wab_xb_view.as_ref(),
2102 1.0,
2103 matmul_parallelism(pa, pb, rows),
2104 );
2105 matmul(
2107 out_mat.rb_mut().submatrix_mut(pa, pa, pb, pb),
2108 Accum::Add,
2109 xb_view.as_ref().transpose(),
2110 wbb_xb_view.as_ref(),
2111 1.0,
2112 matmul_parallelism(pb, pb, rows),
2113 );
2114 }
2115 } for i in 0..total {
2118 for j in 0..i {
2119 out[[i, j]] = out[[j, i]];
2120 }
2121 }
2122 out
2123}
2124
2125fn mat_to_array(mat: MatRef<'_, f64>) -> Array2<f64> {
2126 let nrows = mat.nrows();
2127 let ncols = mat.ncols();
2128 let mut out = Array2::<f64>::zeros((nrows, ncols));
2129 if nrows == 0 || ncols == 0 {
2130 return out;
2131 }
2132 if let Some(out_slice) = out.as_slice_memory_order_mut() {
2135 for i in 0..nrows {
2137 let row_start = i * ncols;
2138 for j in 0..ncols {
2139 out_slice[row_start + j] = mat[(i, j)];
2140 }
2141 }
2142 } else {
2143 for j in 0..ncols {
2144 for i in 0..nrows {
2145 out[[i, j]] = mat[(i, j)];
2146 }
2147 }
2148 }
2149 out
2150}
2151
2152#[inline]
2155pub fn fast_ab_into<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
2156 a: &ArrayBase<S1, Ix2>,
2157 b: &ArrayBase<S2, Ix2>,
2158 out: &mut Array2<f64>,
2159) {
2160 fast_ab_into_impl(a, b, out);
2161}
2162
2163#[inline]
2164fn fast_ab_into_impl<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
2165 a: &ArrayBase<S1, Ix2>,
2166 b: &ArrayBase<S2, Ix2>,
2167 out: &mut Array2<f64>,
2168) {
2169 use faer::Accum;
2170 use faer::linalg::matmul::matmul;
2171
2172 let (n, p) = a.dim();
2173 let (p_b, q) = b.dim();
2174 assert_eq!(p, p_b, "A and B must have compatible inner dimensions");
2175 assert_eq!(out.dim(), (n, q), "output dimensions must match A*B result");
2176
2177 if !should_use_faer_matmul(n, q, p) {
2178 out.assign(&a.dot(b));
2179 return;
2180 }
2181
2182 let aview = FaerArrayView::new(a);
2183 let bview = FaerArrayView::new(b);
2184 let a_ref = aview.as_ref();
2185 let b_ref = bview.as_ref();
2186
2187 let par = matmul_parallelism(n, q, p);
2188 let mut outview = array2_to_matmut(out);
2189 matmul(outview.as_mut(), Accum::Replace, a_ref, b_ref, 1.0, par);
2190}
2191
2192fn diag_to_array(diag: DiagRef<'_, f64>) -> Array1<f64> {
2193 let mat = diag.column_vector().as_mat();
2194 let mut out = Array1::<f64>::zeros(mat.nrows());
2195 for i in 0..mat.nrows() {
2196 out[i] = mat[(i, 0)];
2197 }
2198 out
2199}
2200
2201pub struct FaerArrayView<'a> {
2202 ptr: *const f64,
2203 rows: usize,
2204 cols: usize,
2205 row_stride: isize,
2206 col_stride: isize,
2207 owned: Option<Array2<f64>>,
2208 marker: PhantomData<&'a f64>,
2209}
2210
2211impl<'a> FaerArrayView<'a> {
2212 #[inline]
2213 pub fn new<S: Data<Elem = f64>>(array: &'a ArrayBase<S, Ix2>) -> Self {
2214 let (rows, cols) = array.dim();
2215 let strides = array.strides();
2216 if strides[0] <= 0 || strides[1] <= 0 {
2220 let owned = array.to_owned();
2221 let owned_strides = owned.strides();
2222 return Self {
2223 ptr: owned.as_ptr(),
2224 rows,
2225 cols,
2226 row_stride: owned_strides[0],
2227 col_stride: owned_strides[1],
2228 owned: Some(owned),
2229 marker: PhantomData,
2230 };
2231 }
2232
2233 Self {
2234 ptr: array.as_ptr(),
2235 rows,
2236 cols,
2237 row_stride: strides[0],
2238 col_stride: strides[1],
2239 owned: None,
2240 marker: PhantomData,
2241 }
2242 }
2243
2244 #[inline]
2245 pub fn as_ref(&self) -> MatRef<'_, f64> {
2246 let (ptr, rows, cols, row_stride, col_stride) = if let Some(owned) = &self.owned {
2247 let strides = owned.strides();
2248 (
2249 owned.as_ptr(),
2250 owned.nrows(),
2251 owned.ncols(),
2252 strides[0],
2253 strides[1],
2254 )
2255 } else {
2256 (
2257 self.ptr,
2258 self.rows,
2259 self.cols,
2260 self.row_stride,
2261 self.col_stride,
2262 )
2263 };
2264 unsafe { MatRef::from_raw_parts(ptr, rows, cols, row_stride, col_stride) }
2268 }
2269}
2270
2271pub struct FaerColView<'a> {
2272 ptr: *const f64,
2273 len: usize,
2274 stride: isize,
2275 owned: Option<Array1<f64>>,
2276 marker: PhantomData<&'a f64>,
2277}
2278
2279impl<'a> FaerColView<'a> {
2280 #[inline]
2281 pub fn new<S: Data<Elem = f64>>(array: &'a ArrayBase<S, Ix1>) -> Self {
2282 let len = array.len();
2283 let stride = array.strides()[0];
2284 if stride <= 0 {
2285 let owned = array.to_owned();
2286 return Self {
2287 ptr: owned.as_ptr(),
2288 len,
2289 stride: 1,
2290 owned: Some(owned),
2291 marker: PhantomData,
2292 };
2293 }
2294 Self {
2295 ptr: array.as_ptr(),
2296 len,
2297 stride,
2298 owned: None,
2299 marker: PhantomData,
2300 }
2301 }
2302
2303 #[inline]
2304 pub fn as_ref(&self) -> MatRef<'_, f64> {
2305 let (ptr, len, stride) = if let Some(owned) = &self.owned {
2306 (owned.as_ptr(), owned.len(), 1)
2307 } else {
2308 (self.ptr, self.len, self.stride)
2309 };
2310 unsafe { MatRef::from_raw_parts(ptr, len, 1, stride, 0) }
2314 }
2315}
2316
2317pub trait FaerSvd {
2318 fn svd(
2319 &self,
2320 compute_u: bool,
2321 computevt: bool,
2322 ) -> Result<(Option<Array2<f64>>, Array1<f64>, Option<Array2<f64>>), FaerLinalgError>;
2323}
2324
2325impl<S: Data<Elem = f64>> FaerSvd for ArrayBase<S, Ix2> {
2326 fn svd(
2327 &self,
2328 compute_u: bool,
2329 computevt: bool,
2330 ) -> Result<(Option<Array2<f64>>, Array1<f64>, Option<Array2<f64>>), FaerLinalgError> {
2331 let faerview = FaerArrayView::new(self);
2332 let faer_mat = faerview.as_ref();
2333 if !compute_u && !computevt {
2334 let (rows, cols) = faer_mat.shape();
2335 let mut singular = Diag::<f64>::zeros(rows.min(cols));
2336 let par = get_global_parallelism();
2337 let mut mem = MemBuffer::new(svd::svd_scratch::<f64>(
2338 rows,
2339 cols,
2340 ComputeSvdVectors::No,
2341 ComputeSvdVectors::No,
2342 par,
2343 Default::default(),
2344 ));
2345 let stack = MemStack::new(&mut mem);
2346 svd::svd(
2347 faer_mat,
2348 singular.as_mut(),
2349 None,
2350 None,
2351 par,
2352 stack,
2353 Default::default(),
2354 )
2355 .map_err(|_| FaerLinalgError::SvdNoConvergence {
2356 context: "faer SVD singular values only",
2357 })?;
2358 let singularvalues = diag_to_array(singular.as_ref());
2359 return Ok((None, singularvalues, None));
2360 }
2361
2362 let (rows, cols) = faer_mat.shape();
2363 let rank = rows.min(cols);
2364 let compute_u_flag = if compute_u {
2365 ComputeSvdVectors::Thin
2366 } else {
2367 ComputeSvdVectors::No
2368 };
2369 let computev_flag = if computevt {
2370 ComputeSvdVectors::Thin
2371 } else {
2372 ComputeSvdVectors::No
2373 };
2374
2375 let mut singular = Diag::<f64>::zeros(rows.min(cols));
2376 let mut u_storage = compute_u.then(|| Mat::<f64>::zeros(rows, rank));
2377 let mut v_storage = computevt.then(|| Mat::<f64>::zeros(cols, rank));
2378
2379 let par = get_global_parallelism();
2380 let mut mem = MemBuffer::new(svd::svd_scratch::<f64>(
2381 rows,
2382 cols,
2383 compute_u_flag,
2384 computev_flag,
2385 par,
2386 Default::default(),
2387 ));
2388 let stack = MemStack::new(&mut mem);
2389
2390 svd::svd(
2391 faer_mat.as_ref(),
2392 singular.as_mut(),
2393 u_storage.as_mut().map(|mat| mat.as_mut()),
2394 v_storage.as_mut().map(|mat| mat.as_mut()),
2395 par,
2396 stack,
2397 Default::default(),
2398 )
2399 .map_err(|_| FaerLinalgError::SvdNoConvergence {
2400 context: "faer SVD with vectors",
2401 })?;
2402
2403 let singularvalues = diag_to_array(singular.as_ref());
2404 let u_opt = u_storage.map(|mat| mat_to_array(mat.as_ref()));
2405 let vt_opt = v_storage.map(|mat| {
2406 let mat_ref = mat.as_ref();
2407 let mut out = Array2::<f64>::zeros((mat_ref.ncols(), mat_ref.nrows()));
2408 for j in 0..mat_ref.nrows() {
2409 for i in 0..mat_ref.ncols() {
2410 out[[i, j]] = mat_ref[(j, i)];
2411 }
2412 }
2413 out
2414 });
2415
2416 Ok((u_opt, singularvalues, vt_opt))
2417 }
2418}
2419
2420pub trait FaerEigh {
2421 fn eigh(&self, side: Side) -> Result<(Array1<f64>, Array2<f64>), FaerLinalgError>;
2422}
2423
2424pub fn strict_symmetric_eigh<S: Data<Elem = f64>>(
2431 matrix: &ArrayBase<S, Ix2>,
2432 side: Side,
2433) -> Result<(Array1<f64>, Array2<f64>), FaerLinalgError> {
2434 let owned = matrix.to_owned();
2435 if owned.nrows() == 0 || owned.nrows() != owned.ncols() {
2436 return Err(FaerLinalgError::StrictSelfAdjointEigenInvalidInput {
2437 reason: format!(
2438 "expected non-empty square matrix, got {}x{}",
2439 owned.nrows(),
2440 owned.ncols()
2441 ),
2442 });
2443 }
2444 crate::utils::validate_finite_symmetric_matrix(
2445 &owned,
2446 "strict self-adjoint eigendecomposition",
2447 )
2448 .map_err(
2449 |error| FaerLinalgError::StrictSelfAdjointEigenInvalidInput {
2450 reason: error.to_string(),
2451 },
2452 )?;
2453 let view = FaerArrayView::new(&owned);
2454 let eigen = catch_unwind(AssertUnwindSafe(|| view.as_ref().self_adjoint_eigen(side)))
2455 .map_err(|_| FaerLinalgError::FactorizationFailed {
2456 context: "strict self-adjoint eigendecomposition panic boundary",
2457 })?
2458 .map_err(FaerLinalgError::SelfAdjointEigen)?;
2459 let values = diag_to_array(eigen.S());
2460 let vectors = mat_to_array(eigen.U());
2461 if values.iter().any(|value| !value.is_finite())
2462 || vectors.iter().any(|value| !value.is_finite())
2463 {
2464 return Err(FaerLinalgError::SelfAdjointEigenNonFiniteInput {
2465 context: "strict self-adjoint eigendecomposition output validation",
2466 });
2467 }
2468 Ok((values, vectors))
2469}
2470
2471impl<S: Data<Elem = f64>> FaerEigh for ArrayBase<S, Ix2> {
2472 fn eigh(&self, side: Side) -> Result<(Array1<f64>, Array2<f64>), FaerLinalgError> {
2473 fn try_eigh(
2474 matrix: &Array2<f64>,
2475 side: Side,
2476 ) -> Result<(Array1<f64>, Array2<f64>), FaerLinalgError> {
2477 let faerview = FaerArrayView::new(matrix);
2478 let eigh_started = std::time::Instant::now();
2497 let eigh_par = get_global_parallelism();
2498 let eigen = catch_unwind(AssertUnwindSafe(|| {
2499 faerview.as_ref().self_adjoint_eigen(side)
2500 }))
2501 .map_err(|_| FaerLinalgError::FactorizationFailed {
2502 context: "self-adjoint eigendecomposition panic boundary",
2503 })?
2504 .map_err(FaerLinalgError::SelfAdjointEigen)?;
2505 let eigh_elapsed = eigh_started.elapsed();
2506 let eigh_calls = EIGH_CALLS.fetch_add(1, Ordering::Relaxed) + 1;
2507 if eigh_par == Par::Seq {
2508 EIGH_SEQ_CALLS.fetch_add(1, Ordering::Relaxed);
2509 }
2510 EIGH_MAX_DIM.fetch_max(matrix.nrows() as u64, Ordering::Relaxed);
2511 let eigh_nanos_total = EIGH_NANOS
2512 .fetch_add(eigh_elapsed.as_nanos() as u64, Ordering::Relaxed)
2513 + eigh_elapsed.as_nanos() as u64;
2514 record_thread_eigh(
2515 eigh_par == Par::Seq,
2516 matrix.nrows() as u64,
2517 eigh_elapsed.as_nanos() as u64,
2518 );
2519 log::debug!(
2520 "[eigh] dim={} elapsed={:.3}s faer_global_parallelism={:?} \
2521 calls_so_far={eigh_calls} cumulative={:.3}s",
2522 matrix.nrows(),
2523 eigh_elapsed.as_secs_f64(),
2524 eigh_par,
2525 eigh_nanos_total as f64 / 1e9,
2526 );
2527 let values = diag_to_array(eigen.S());
2528 let vectors = mat_to_array(eigen.U());
2529 Ok((values, vectors))
2530 }
2531
2532 let owned = self.to_owned();
2533 if owned.nrows() != owned.ncols() {
2534 return Err(FaerLinalgError::FactorizationFailed {
2535 context: "self-adjoint eigendecomposition non-square input",
2536 });
2537 }
2538 if owned.nrows() == 0 {
2539 return Ok((Array1::zeros(0), Array2::zeros((0, 0))));
2540 }
2541 if owned.iter().any(|value| !value.is_finite()) {
2542 return Err(FaerLinalgError::SelfAdjointEigenNonFiniteInput {
2543 context: "self-adjoint eigendecomposition input validation",
2544 });
2545 }
2546 if let Ok((evals, evecs)) = try_eigh(&owned, side)
2547 && evals.iter().all(|value| value.is_finite())
2548 && evecs.iter().all(|value| value.is_finite())
2549 {
2550 return Ok((evals, evecs));
2551 }
2552
2553 let mut repaired = owned.clone();
2554 crate::matrix::symmetrize_in_place(&mut repaired);
2555
2556 let scale = repaired
2557 .iter()
2558 .fold(0.0_f64, |acc, &value| acc.max(value.abs()))
2559 .max(1.0);
2560 let scaled = repaired.mapv(|value| value / scale);
2561 const JITTER_SCHEDULE: [f64; 6] = [0.0, 1e-12, 1e-10, 1e-8, 1e-6, 1e-4];
2567 let jitter_schedule = JITTER_SCHEDULE;
2568 let mut last_error = FaerLinalgError::FactorizationFailed {
2569 context: "self-adjoint eigendecomposition repair attempts",
2570 };
2571
2572 for &jitter in &jitter_schedule {
2573 let mut candidate = scaled.clone();
2574 if jitter > 0.0 {
2575 let n = candidate.nrows();
2576 for i in 0..n {
2577 candidate[[i, i]] += jitter;
2578 }
2579 }
2580
2581 match try_eigh(&candidate, side) {
2582 Ok((mut evals, evecs))
2583 if evals.iter().all(|value| value.is_finite())
2584 && evecs.iter().all(|value| value.is_finite()) =>
2585 {
2586 for value in &mut evals {
2587 *value = (*value - jitter) * scale;
2588 }
2589 return Ok((evals, evecs));
2590 }
2591 Ok((_, _)) => {
2592 last_error = FaerLinalgError::SelfAdjointEigenNonFiniteInput {
2593 context: "self-adjoint eigendecomposition repaired output validation",
2594 };
2595 }
2596 Err(err) => {
2597 last_error = err;
2598 }
2599 }
2600 }
2601
2602 Err(last_error)
2603 }
2604}
2605
2606pub struct FaerCholeskyFactor {
2607 factor: solvers::Llt<f64>,
2608}
2609
2610impl FaerCholeskyFactor {
2611 pub fn solvevec(&self, rhs: &Array1<f64>) -> Array1<f64> {
2612 let mut rhs = rhs.to_owned();
2613 let mut rhsview = array1_to_col_matmut(&mut rhs);
2614 self.factor.solve_in_place(rhsview.as_mut());
2615 rhs
2616 }
2617
2618 pub fn solve_mat_in_place(&self, rhs: &mut Array2<f64>) {
2619 let mut rhsview = array2_to_matmut(rhs);
2620 self.factor.solve_in_place(rhsview.as_mut());
2621 }
2622
2623 pub fn solve_mat_into<S: Data<Elem = f64>>(
2624 &self,
2625 rhs: &ArrayBase<S, Ix2>,
2626 out: &mut Array2<f64>,
2627 ) {
2628 if out.dim() != rhs.dim() {
2629 *out = Array2::<f64>::zeros(rhs.dim());
2630 }
2631 out.assign(rhs);
2632 self.solve_mat_in_place(out);
2633 }
2634
2635 pub fn solve_mat(&self, rhs: &Array2<f64>) -> Array2<f64> {
2636 let mut out = Array2::<f64>::zeros(rhs.dim());
2637 self.solve_mat_into(rhs, &mut out);
2638 out
2639 }
2640
2641 pub fn diag(&self) -> Array1<f64> {
2642 diag_to_array(self.factor.L().diagonal())
2643 }
2644
2645 pub fn lower_triangular(&self) -> Array2<f64> {
2646 mat_to_array(self.factor.L())
2647 }
2648}
2649
2650impl crate::matrix::FactorizedSystem for FaerCholeskyFactor {
2651 fn solve(&self, rhs: &Array1<f64>) -> Result<Array1<f64>, String> {
2652 let out = self.solvevec(rhs);
2653 if out.iter().all(|value| value.is_finite()) {
2654 Ok(out)
2655 } else {
2656 Err("strict Cholesky solve produced non-finite values".to_string())
2657 }
2658 }
2659
2660 fn solvemulti(&self, rhs: &Array2<f64>) -> Result<Array2<f64>, String> {
2661 let out = self.solve_mat(rhs);
2662 if out.iter().all(|value| value.is_finite()) {
2663 Ok(out)
2664 } else {
2665 Err("strict Cholesky multi-solve produced non-finite values".to_string())
2666 }
2667 }
2668
2669 fn logdet(&self) -> f64 {
2670 cholesky_factor_logdet(self.factor.L())
2671 }
2672}
2673
2674pub trait FaerCholesky {
2675 fn cholesky(&self, side: Side) -> Result<FaerCholeskyFactor, FaerLinalgError>;
2676}
2677
2678impl<S: Data<Elem = f64>> FaerCholesky for ArrayBase<S, Ix2> {
2679 fn cholesky(&self, side: Side) -> Result<FaerCholeskyFactor, FaerLinalgError> {
2680 let faerview = FaerArrayView::new(self);
2681 let factor = faerview
2682 .as_ref()
2683 .llt(side)
2684 .map_err(FaerLinalgError::Cholesky)?;
2685 Ok(FaerCholeskyFactor { factor })
2686 }
2687}
2688
2689pub trait FaerQr {
2690 fn qr(&self) -> Result<(Array2<f64>, Array2<f64>), FaerLinalgError>;
2691}
2692
2693impl<S: Data<Elem = f64>> FaerQr for ArrayBase<S, Ix2> {
2694 fn qr(&self) -> Result<(Array2<f64>, Array2<f64>), FaerLinalgError> {
2695 let faerview = FaerArrayView::new(self);
2696 let qr = faerview.as_ref().qr();
2697 let q = qr.compute_thin_Q();
2698 let r = qr.thin_R();
2699 Ok((mat_to_array(q.as_ref()), mat_to_array(r)))
2700 }
2701}
2702
2703pub fn rrqr_nullspace_basis<S: Data<Elem = f64>>(
2722 a: &ArrayBase<S, Ix2>,
2723 rank_alpha: f64,
2724) -> Result<(Array2<f64>, usize), FaerLinalgError> {
2725 rrqr_nullspace_basis_inner(a, RrqrRankCutoff::RelativeAlpha(rank_alpha))
2726}
2727
2728#[derive(Debug, Clone, Copy)]
2731enum RrqrRankCutoff {
2732 RelativeAlpha(f64),
2736 Absolute(f64),
2742}
2743
2744pub fn rrqr_nullspace_basis_with_cutoff<S: Data<Elem = f64>>(
2757 a: &ArrayBase<S, Ix2>,
2758 cutoff: f64,
2759) -> Result<(Array2<f64>, usize), FaerLinalgError> {
2760 rrqr_nullspace_basis_inner(a, RrqrRankCutoff::Absolute(cutoff))
2761}
2762
2763fn rrqr_nullspace_basis_inner<S: Data<Elem = f64>>(
2764 a: &ArrayBase<S, Ix2>,
2765 cutoff: RrqrRankCutoff,
2766) -> Result<(Array2<f64>, usize), FaerLinalgError> {
2767 let faerview = FaerArrayView::new(a);
2768 let qr = faerview.as_ref().col_piv_qr();
2769 let r = qr.thin_R();
2770 let diag_len = r.nrows().min(r.ncols());
2771 let leading_diag = if diag_len > 0 { r[(0, 0)].abs() } else { 0.0 };
2772 let tol = match cutoff {
2773 RrqrRankCutoff::RelativeAlpha(rank_alpha) => {
2774 rank_alpha
2775 * f64::EPSILON
2776 * (a.nrows().max(a.ncols()).max(1) as f64)
2777 * leading_diag.max(1.0)
2778 }
2779 RrqrRankCutoff::Absolute(tol) => tol,
2780 };
2781 let rank = (0..diag_len).filter(|&i| r[(i, i)].abs() > tol).count();
2782 let z = if rank >= a.nrows() {
2783 Array2::<f64>::zeros((a.nrows(), 0))
2784 } else if rank == 0 {
2785 Array2::<f64>::eye(a.nrows())
2789 } else {
2790 let nullity = a.nrows() - rank;
2791 let mut selector = Mat::<f64>::zeros(a.nrows(), nullity);
2792 for j in 0..nullity {
2793 selector[(rank + j, j)] = 1.0;
2794 }
2795 let par = get_global_parallelism();
2796 faer::linalg::householder::apply_block_householder_sequence_on_the_left_in_place_with_conj(
2797 qr.Q_basis(),
2798 qr.Q_coeff(),
2799 Conj::No,
2800 selector.as_mut(),
2801 par,
2802 MemStack::new(&mut MemBuffer::new(
2803 faer::linalg::householder::apply_block_householder_sequence_on_the_left_in_place_scratch::<f64>(
2804 a.nrows(),
2805 qr.Q_coeff().nrows(),
2806 nullity,
2807 ),
2808 )),
2809 );
2810 mat_to_array(selector.as_ref())
2811 };
2812 Ok((z, rank))
2813}
2814
2815#[inline]
2816pub const fn default_rrqr_rank_alpha() -> f64 {
2817 RRQR_RANK_ALPHA
2818}
2819
2820pub struct RrqrWithPermutation {
2831 pub rank: usize,
2832 pub column_permutation: Vec<usize>,
2833 pub leading_diag_abs: f64,
2834 pub rank_tol: f64,
2835}
2836
2837pub fn rrqr_with_permutation<S: Data<Elem = f64>>(
2846 a: &ArrayBase<S, Ix2>,
2847 rank_alpha: f64,
2848) -> Result<RrqrWithPermutation, FaerLinalgError> {
2849 if a.nrows() == 0 {
2850 return Err(FaerLinalgError::FactorizationFailed {
2851 context: "rrqr_with_permutation: input has zero rows",
2852 });
2853 }
2854 let faerview = FaerArrayView::new(a);
2855 let qr = faerview.as_ref().col_piv_qr();
2856 let r = qr.thin_R();
2857 let diag_len = r.nrows().min(r.ncols());
2858 let leading_diag = if diag_len > 0 { r[(0, 0)].abs() } else { 0.0 };
2859 let tol = rank_alpha
2860 * f64::EPSILON
2861 * (a.nrows().max(a.ncols()).max(1) as f64)
2862 * leading_diag.max(1.0);
2863 let rank = (0..diag_len).filter(|&i| r[(i, i)].abs() > tol).count();
2864 let (forward, _inverse) = qr.P().arrays();
2865 let column_permutation: Vec<usize> = forward.iter().copied().map(|idx| idx.unbound()).collect();
2866 Ok(RrqrWithPermutation {
2867 rank,
2868 column_permutation,
2869 leading_diag_abs: leading_diag,
2870 rank_tol: tol,
2871 })
2872}
2873
2874pub struct RrqrFromGram {
2883 pub rank: usize,
2884 pub column_permutation: Vec<usize>,
2885 pub rank_tol: f64,
2886 pub leading_diag_abs: f64,
2891 pub verdict_margin: f64,
2894}
2895
2896pub fn rrqr_from_gram_with_permutation<S: Data<Elem = f64>>(
2932 gram: &ArrayBase<S, Ix2>,
2933 m_rows: usize,
2934 rank_alpha: f64,
2935) -> Result<RrqrFromGram, FaerLinalgError> {
2936 let p = gram.ncols();
2937 if p == 0 {
2938 return Ok(RrqrFromGram {
2939 rank: 0,
2940 column_permutation: Vec::new(),
2941 rank_tol: 0.0,
2942 leading_diag_abs: 0.0,
2943 verdict_margin: 0.0,
2944 });
2945 }
2946 if gram.nrows() != p {
2947 return Err(FaerLinalgError::FactorizationFailed {
2948 context: "rrqr_from_gram_with_permutation: Gram is not square",
2949 });
2950 }
2951 let (evals, evecs) = gram.eigh(Side::Lower)?;
2960 let mut f = Array2::<f64>::zeros((p, p));
2961 for k in 0..p {
2962 let scale = evals[k].max(0.0).sqrt();
2963 if scale == 0.0 {
2964 continue;
2965 }
2966 for i in 0..p {
2967 f[[k, i]] = scale * evecs[[i, k]];
2968 }
2969 }
2970 let faer_f = FaerArrayView::new(&f);
2974 let qr = faer_f.as_ref().col_piv_qr();
2975 let r = qr.thin_R();
2976 let diag_len = r.nrows().min(r.ncols());
2977 let pivots: Vec<f64> = (0..diag_len).map(|i| r[(i, i)].abs()).collect();
2978 let leading_diag = pivots.first().copied().unwrap_or(0.0);
2979 let (forward, _inverse) = qr.P().arrays();
2980 let column_permutation: Vec<usize> = forward.iter().copied().map(|idx| idx.unbound()).collect();
2981 let tol = rank_alpha * f64::EPSILON * (m_rows.max(p).max(1) as f64) * leading_diag.max(1.0);
2985 let rank = pivots.iter().filter(|&&v| v > tol).count();
2986 let min_kept = pivots[..rank].iter().copied().fold(f64::INFINITY, f64::min);
2987 let max_dropped = pivots[rank..].iter().copied().fold(0.0f64, f64::max);
2988 let kept_margin = if rank == 0 {
2992 f64::INFINITY
2993 } else {
2994 min_kept / tol
2995 };
2996 let dropped_margin = if rank == diag_len {
2997 f64::INFINITY
2998 } else {
2999 tol / max_dropped.max(f64::MIN_POSITIVE)
3000 };
3001 let gram_precision_floor = f64::EPSILON.sqrt() * leading_diag.max(1.0);
3023 let kept_floor_margin = if rank == 0 {
3024 f64::INFINITY
3025 } else {
3026 min_kept / gram_precision_floor.max(f64::MIN_POSITIVE)
3027 };
3028 let verdict_margin = kept_margin.min(dropped_margin).min(kept_floor_margin);
3029 Ok(RrqrFromGram {
3030 rank,
3031 column_permutation,
3032 rank_tol: tol,
3033 leading_diag_abs: leading_diag,
3034 verdict_margin,
3035 })
3036}
3037
3038#[cfg(test)]
3039mod tests {
3040 use super::*;
3041 use ndarray::{array, s};
3042
3043 const JOINT_GRAM_RRQR_TRUST_MARGIN_FOR_TEST: f64 = 1.0e3;
3047
3048 #[test]
3049 fn rrqr_nullspace_basis_is_orthonormal_and_annihilates_transpose() {
3050 let a = array![[1.0, 0.0], [1.0, 0.0], [0.0, 2.0], [0.0, 0.0],];
3051 let (z, rank) =
3052 rrqr_nullspace_basis(&a, default_rrqr_rank_alpha()).expect("RRQR should succeed");
3053 assert_eq!(rank, 2);
3054 assert_eq!(z.nrows(), 4);
3055 assert_eq!(z.ncols(), 2);
3056
3057 let gram = z.t().dot(&z);
3058 let ident = Array2::<f64>::eye(z.ncols());
3059 let gram_err = (&gram - &ident)
3060 .iter()
3061 .fold(0.0_f64, |acc, &v| acc.max(v.abs()));
3062 assert!(gram_err < 1e-10, "Z is not orthonormal: {gram_err:e}");
3063
3064 let residual = a.t().dot(&z);
3065 let resid_max = residual.iter().fold(0.0_f64, |acc, &v| acc.max(v.abs()));
3066 assert!(resid_max < 1e-10, "A^T Z residual too large: {resid_max:e}");
3067 }
3068
3069 #[test]
3070 fn rrqr_with_permutation_attributes_redundant_column() {
3071 let a = array![
3075 [1.0, 0.0, 1.0],
3076 [1.0, 0.0, 1.0],
3077 [0.0, 2.0, 0.0],
3078 [0.0, 0.0, 0.0],
3079 ];
3080 let result =
3081 rrqr_with_permutation(&a, default_rrqr_rank_alpha()).expect("RRQR should succeed");
3082 assert_eq!(result.rank, 2);
3083 assert_eq!(result.column_permutation.len(), 3);
3084 let demoted = result.column_permutation[result.rank..].to_vec();
3085 assert!(
3086 demoted.contains(&2) || demoted.contains(&0),
3087 "demoted suffix should include one of the aliased columns (0 or 2), got {demoted:?}"
3088 );
3089 let mut sorted = result.column_permutation.clone();
3090 sorted.sort();
3091 assert_eq!(
3092 sorted,
3093 vec![0, 1, 2],
3094 "permutation must be a valid bijection on 0..n"
3095 );
3096 }
3097
3098 #[test]
3107 fn rrqr_with_permutation_pivots_the_larger_norm_column_first() {
3108 let a = array![[1.0, 0.0], [0.0, 2.0], [0.0, 0.0]];
3109 let result =
3110 rrqr_with_permutation(&a, default_rrqr_rank_alpha()).expect("RRQR should succeed");
3111 assert_eq!(result.rank, 2);
3112 let perm = result.column_permutation.clone();
3113
3114 let mut sorted = perm.clone();
3115 sorted.sort();
3116 assert_eq!(
3117 sorted,
3118 vec![0, 1],
3119 "permutation must be a bijection on 0..n, got {perm:?}"
3120 );
3121
3122 assert_eq!(
3124 perm,
3125 vec![1, 0],
3126 "column-pivoted QR must take the larger-norm column (1, norm 2) \
3127 before the smaller (0, norm 1), got {perm:?}"
3128 );
3129
3130 let norms: Vec<f64> = perm
3133 .iter()
3134 .map(|&j| a.column(j).iter().map(|value| value * value).sum::<f64>().sqrt())
3135 .collect();
3136 for window in norms.windows(2) {
3137 assert!(
3138 window[0] >= window[1],
3139 "pivoted column norms must be non-increasing, got {norms:?}"
3140 );
3141 }
3142 }
3143
3144 #[test]
3145 fn rrqr_with_permutation_rejects_zero_rows() {
3146 let a = Array2::<f64>::zeros((0, 3));
3147 assert!(rrqr_with_permutation(&a, default_rrqr_rank_alpha()).is_err());
3148 }
3149
3150 #[test]
3151 fn rrqr_nullspace_basis_square_zero_matrix_is_finite_identity() {
3152 let a = Array2::<f64>::zeros((3, 3));
3155 let (z, rank) =
3156 rrqr_nullspace_basis(&a, default_rrqr_rank_alpha()).expect("RRQR should succeed");
3157 assert_eq!(rank, 0);
3158 assert_eq!(z.dim(), (3, 3));
3159 assert!(
3160 z.iter().all(|v| v.is_finite()),
3161 "square zero matrix produced a non-finite null basis: {z:?}"
3162 );
3163 let gram = z.t().dot(&z);
3164 let ident = Array2::<f64>::eye(3);
3165 let gram_err = (&gram - &ident)
3166 .iter()
3167 .fold(0.0_f64, |acc, &v| acc.max(v.abs()));
3168 assert!(gram_err < 1e-10, "Z is not orthonormal: {gram_err:e}");
3169 }
3170
3171 #[test]
3172 fn rrqr_nullspace_basis_detectszero_rank_matrix() {
3173 let a = Array2::<f64>::zeros((5, 2));
3174 let (z, rank) =
3175 rrqr_nullspace_basis(&a, default_rrqr_rank_alpha()).expect("RRQR should succeed");
3176 assert_eq!(rank, 0);
3177 assert_eq!(z.dim(), (5, 5));
3178 let ident = Array2::<f64>::eye(5);
3179 let max_err = (&z.slice(s![.., ..5]).to_owned() - &ident)
3180 .iter()
3181 .fold(0.0_f64, |acc, &v| acc.max(v.abs()));
3182 assert!(max_err < 1e-10, "zero matrix should yield identity basis");
3183 }
3184
3185 #[test]
3194 fn eigh_on_nan_matrix_rejects_non_finite_input() {
3195 let mat = array![
3196 [1.0, 0.0, 0.0, 0.0],
3197 [0.0, 2.0, 0.0, 0.0],
3198 [0.0, 0.0, 3.0, f64::NAN],
3199 [0.0, 0.0, f64::NAN, 4.0]
3200 ];
3201 let err = mat
3202 .eigh(Side::Lower)
3203 .expect_err("non-finite symmetric input must be rejected");
3204 assert!(matches!(
3205 err,
3206 FaerLinalgError::SelfAdjointEigenNonFiniteInput { .. }
3207 ));
3208 }
3209
3210 #[test]
3211 fn fast_ata_matches_full_gemm_above_threshold() {
3212 let n = 200;
3215 let p = 40;
3216 let a: Array2<f64> = Array2::from_shape_fn((n, p), |(i, j)| {
3217 ((i * 7 + j * 3) as f64).sin() + 0.1 * j as f64
3218 });
3219 let expected = a.t().dot(&a);
3220 let got = fast_ata(&a);
3221 let max_err = (&got - &expected)
3222 .iter()
3223 .fold(0.0_f64, |acc, &v| acc.max(v.abs()));
3224 assert!(max_err < 1e-10, "fast_ata mismatch: {max_err:e}");
3225 for i in 0..p {
3227 for j in 0..p {
3228 assert!((got[[i, j]] - got[[j, i]]).abs() < 1e-12);
3229 }
3230 }
3231 }
3232
3233 #[test]
3234 fn fast_xt_diag_x_matches_naive_above_threshold() {
3235 let n = 400;
3236 let p = 36;
3237 let x: Array2<f64> =
3238 Array2::from_shape_fn((n, p), |(i, j)| (i as f64 * 0.1).cos() + j as f64 * 0.05);
3239 let w: Array1<f64> = Array1::from_shape_fn(n, |i| (i as f64 * 0.03).sin());
3240 let wx = Array2::from_shape_fn((n, p), |(i, j)| w[i] * x[[i, j]]);
3242 let expected = x.t().dot(&wx);
3243 let got = fast_xt_diag_x(&x, &w);
3244 let max_err = (&got - &expected)
3245 .iter()
3246 .fold(0.0_f64, |acc, &v| acc.max(v.abs()));
3247 assert!(max_err < 1e-9, "fast_xt_diag_x mismatch: {max_err:e}");
3248 for i in 0..p {
3249 for j in 0..p {
3250 assert!((got[[i, j]] - got[[j, i]]).abs() < 1e-12);
3251 }
3252 }
3253 }
3254
3255 #[test]
3256 fn stream_weighted_crossprod_full_and_triangular_parity_with_negative_weights() {
3257 for &(n, p) in &[(900usize, 40usize), (8usize, 3usize)] {
3266 let x: Array2<f64> =
3267 Array2::from_shape_fn((n, p), |(i, j)| (i as f64 * 0.07).cos() + j as f64 * 0.013);
3268 let w: Array1<f64> =
3271 Array1::from_shape_fn(n, |i| (i as f64 * 0.11).sin() - 0.25 * (i % 3) as f64);
3272 assert!(
3273 w.iter().any(|&v| v < 0.0),
3274 "weight vector must contain negatives to test sign preservation"
3275 );
3276
3277 let wx = Array2::from_shape_fn((n, p), |(i, j)| w[i] * x[[i, j]]);
3279 let expected = x.t().dot(&wx);
3280
3281 let par = matmul_parallelism(p, p, n);
3282
3283 let mut full = Array2::<f64>::ones((p, p));
3285 stream_weighted_crossprod_into(
3286 &x,
3287 &w,
3288 &mut full,
3289 CrossprodStructure::Full,
3290 CrossprodAccum::Replace,
3291 par,
3292 );
3293
3294 let mut tri = Array2::<f64>::from_elem((p, p), -7.0);
3298 stream_weighted_crossprod_into(
3299 &x,
3300 &w,
3301 &mut tri,
3302 CrossprodStructure::SymmetricLower,
3303 CrossprodAccum::Replace,
3304 par,
3305 );
3306
3307 let full_err = (&full - &expected)
3308 .iter()
3309 .fold(0.0_f64, |a, &v| a.max(v.abs()));
3310 let tri_err = (&tri - &expected)
3311 .iter()
3312 .fold(0.0_f64, |a, &v| a.max(v.abs()));
3313 assert!(
3314 full_err < 1e-9,
3315 "full kernel mismatch (n={n}, p={p}): {full_err:e}"
3316 );
3317 assert!(
3318 tri_err < 1e-9,
3319 "triangular kernel mismatch (n={n}, p={p}): {tri_err:e}"
3320 );
3321
3322 for i in 0..p {
3325 for j in 0..p {
3326 assert!(
3327 (full[[i, j]] - tri[[i, j]]).abs() < 1e-12,
3328 "full vs triangular disagree at ({i},{j})"
3329 );
3330 assert!(
3331 (tri[[i, j]] - tri[[j, i]]).abs() < 1e-12,
3332 "triangular output not symmetric at ({i},{j})"
3333 );
3334 }
3335 }
3336
3337 let base = Array2::<f64>::from_elem((p, p), 1.5);
3340 let mut add_full = base.clone();
3341 stream_weighted_crossprod_into(
3342 &x,
3343 &w,
3344 &mut add_full,
3345 CrossprodStructure::Full,
3346 CrossprodAccum::Add,
3347 par,
3348 );
3349 let mut add_tri = base.clone();
3350 stream_weighted_crossprod_into(
3351 &x,
3352 &w,
3353 &mut add_tri,
3354 CrossprodStructure::SymmetricLower,
3355 CrossprodAccum::Add,
3356 par,
3357 );
3358 let expected_add = &base + &expected;
3359 let add_full_err = (&add_full - &expected_add)
3360 .iter()
3361 .fold(0.0_f64, |a, &v| a.max(v.abs()));
3362 let add_tri_err = (&add_tri - &expected_add)
3363 .iter()
3364 .fold(0.0_f64, |a, &v| a.max(v.abs()));
3365 assert!(
3366 add_full_err < 1e-9,
3367 "full Add mismatch (n={n}, p={p}): {add_full_err:e}"
3368 );
3369 assert!(
3370 add_tri_err < 1e-9,
3371 "triangular Add mismatch (n={n}, p={p}): {add_tri_err:e}"
3372 );
3373
3374 let returned = fast_xt_diag_x(&x, &w);
3377 let returned_err = (&returned - &full)
3378 .iter()
3379 .fold(0.0_f64, |a, &v| a.max(v.abs()));
3380 assert!(
3381 returned_err < 1e-12,
3382 "return adapter vs stream-into adapter disagree (n={n}, p={p}): {returned_err:e}"
3383 );
3384 }
3385 }
3386
3387 #[test]
3388 fn eigh_succeeds_on_same_structure_without_nan() {
3389 let mat = array![[1.0, 0.5, 0.1], [0.5, 2.0, 0.3], [0.1, 0.3, 1.5]];
3391 let (evals, _) = mat
3392 .eigh(Side::Lower)
3393 .expect("eigh should succeed on a well-conditioned finite matrix");
3394 assert!(
3395 evals.iter().all(|&v| v.is_finite()),
3396 "all eigenvalues should be finite"
3397 );
3398 }
3399
3400 #[test]
3409 fn gram_rrqr_flags_low_margin_on_exact_collinearity_so_caller_falls_back() {
3410 let n = 48usize;
3413 let x: Vec<f64> = (0..n)
3414 .map(|i| -1.0 + 2.0 * (i as f64) / (n as f64 - 1.0))
3415 .collect();
3416 let mut a = Array2::<f64>::zeros((n, 4));
3417 for i in 0..n {
3418 a[[i, 0]] = 1.0;
3419 a[[i, 1]] = x[i];
3420 a[[i, 2]] = x[i];
3421 a[[i, 3]] = x[i] * x[i];
3422 }
3423 let alpha = default_rrqr_rank_alpha();
3424
3425 let tall = rrqr_with_permutation(&a, alpha).expect("tall RRQR should succeed");
3428 assert_eq!(tall.rank, 3, "tall RRQR must demote the exact alias");
3429
3430 let unit = Array1::<f64>::ones(n);
3441 let gram = fast_xt_diag_x_with_parallelism(&a, &unit, faer::get_global_parallelism());
3442 let gram_rrqr =
3443 rrqr_from_gram_with_permutation(&gram, n, alpha).expect("Gram RRQR should succeed");
3444 let ok =
3445 gram_rrqr.rank == 3 || gram_rrqr.verdict_margin < JOINT_GRAM_RRQR_TRUST_MARGIN_FOR_TEST;
3446 assert!(
3447 ok,
3448 "gam#933: Gram RRQR must either find correct rank=3 OR signal low margin \
3449 (< {:.0e}) to force the tall fallback; got rank={} margin={:.3e}",
3450 JOINT_GRAM_RRQR_TRUST_MARGIN_FOR_TEST, gram_rrqr.rank, gram_rrqr.verdict_margin,
3451 );
3452 }
3453
3454 #[test]
3459 fn gram_rrqr_keeps_high_margin_on_full_rank_design() {
3460 let n = 200usize;
3461 let p = 5usize;
3462 let mut a = Array2::<f64>::zeros((n, p));
3463 for i in 0..n {
3465 let t = (i as f64) / (n as f64 - 1.0);
3466 a[[i, 0]] = 1.0;
3467 a[[i, 1]] = t;
3468 a[[i, 2]] = t * t;
3469 a[[i, 3]] = t * t * t;
3470 a[[i, 4]] = (t * 6.0).sin();
3471 }
3472 let alpha = default_rrqr_rank_alpha();
3473 let unit = Array1::<f64>::ones(n);
3474 let gram = fast_xt_diag_x_with_parallelism(&a, &unit, faer::get_global_parallelism());
3475 let gram_rrqr =
3476 rrqr_from_gram_with_permutation(&gram, n, alpha).expect("Gram RRQR should succeed");
3477 assert_eq!(gram_rrqr.rank, p, "full-rank design must keep all columns");
3478 assert!(
3479 gram_rrqr.verdict_margin >= JOINT_GRAM_RRQR_TRUST_MARGIN_FOR_TEST,
3480 "full-rank design must keep a high margin (fast Gram path); got {:.3e}",
3481 gram_rrqr.verdict_margin,
3482 );
3483 }
3484
3485 fn max_abs_diff(a: &Array2<f64>, b: &Array2<f64>) -> f64 {
3488 assert_eq!(a.dim(), b.dim(), "shape mismatch in max_abs_diff");
3489 a.iter()
3490 .zip(b.iter())
3491 .fold(0.0_f64, |acc, (&x, &y)| acc.max((x - y).abs()))
3492 }
3493
3494 fn max_abs_diff_1d(a: &Array1<f64>, b: &Array1<f64>) -> f64 {
3495 assert_eq!(a.len(), b.len(), "len mismatch in max_abs_diff_1d");
3496 a.iter()
3497 .zip(b.iter())
3498 .fold(0.0_f64, |acc, (&x, &y)| acc.max((x - y).abs()))
3499 }
3500
3501 #[test]
3503 fn fast_ab_small_matches_ndarray_dot() {
3504 let a = array![[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]];
3505 let b = array![[7.0, 8.0], [9.0, 10.0], [11.0, 12.0]];
3506 let got = fast_ab(&a, &b);
3507 let want = a.dot(&b);
3508 assert!(max_abs_diff(&got, &want) < 1e-12, "fast_ab small mismatch");
3509 assert_eq!(got.dim(), (2, 2));
3510 }
3511
3512 #[test]
3514 fn fast_ab_large_matches_ndarray_dot() {
3515 let n = 50usize;
3516 let p = 40usize;
3517 let q = 35usize;
3518 let mut a = Array2::<f64>::zeros((n, p));
3519 let mut b = Array2::<f64>::zeros((p, q));
3520 let mut state = 0xDEAD_BEEF_1234_5678u64;
3521 let next = |s: &mut u64| -> f64 {
3522 *s ^= *s << 13;
3523 *s ^= *s >> 7;
3524 *s ^= *s << 17;
3525 ((*s >> 11) as f64 / ((1u64 << 53) as f64)) - 0.5
3526 };
3527 for v in a.iter_mut() {
3528 *v = next(&mut state);
3529 }
3530 for v in b.iter_mut() {
3531 *v = next(&mut state);
3532 }
3533 let got = fast_ab(&a, &b);
3534 let want = a.dot(&b);
3535 assert!(max_abs_diff(&got, &want) < 1e-9, "fast_ab large mismatch");
3536 }
3537
3538 #[test]
3540 fn fast_atb_small_matches_ndarray_dot() {
3541 let a = array![[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]];
3542 let b = array![[7.0, 8.0, 9.0], [10.0, 11.0, 12.0], [13.0, 14.0, 15.0]];
3543 let got = fast_atb(&a, &b);
3544 let want = a.t().dot(&b);
3545 assert!(max_abs_diff(&got, &want) < 1e-12, "fast_atb small mismatch");
3546 assert_eq!(got.dim(), (2, 3));
3547 }
3548
3549 #[test]
3551 fn fast_atb_large_matches_ndarray_dot() {
3552 let n = 50usize;
3553 let p = 40usize;
3554 let q = 35usize;
3555 let mut a = Array2::<f64>::zeros((n, p));
3556 let mut b = Array2::<f64>::zeros((n, q));
3557 let mut state = 0xCAFE_BABE_9876_5432u64;
3558 let next = |s: &mut u64| -> f64 {
3559 *s ^= *s << 13;
3560 *s ^= *s >> 7;
3561 *s ^= *s << 17;
3562 ((*s >> 11) as f64 / ((1u64 << 53) as f64)) - 0.5
3563 };
3564 for v in a.iter_mut() {
3565 *v = next(&mut state);
3566 }
3567 for v in b.iter_mut() {
3568 *v = next(&mut state);
3569 }
3570 let got = fast_atb(&a, &b);
3571 let want = a.t().dot(&b);
3572 assert!(max_abs_diff(&got, &want) < 1e-9, "fast_atb large mismatch");
3573 }
3574
3575 #[test]
3577 fn fast_abt_small_matches_ndarray_dot() {
3578 let a = array![[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]];
3579 let b = array![[7.0, 8.0, 9.0], [10.0, 11.0, 12.0]];
3580 let got = fast_abt(&a, &b);
3581 let want = a.dot(&b.t());
3582 assert!(max_abs_diff(&got, &want) < 1e-12, "fast_abt small mismatch");
3583 assert_eq!(got.dim(), (2, 2));
3584 }
3585
3586 #[test]
3588 fn fast_av_small_matches_ndarray_dot() {
3589 let a = array![[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]];
3590 let v = array![1.0, -1.0, 2.0];
3591 let got = fast_av(&a, &v);
3592 let want = a.dot(&v);
3593 assert!(
3594 max_abs_diff_1d(&got, &want) < 1e-12,
3595 "fast_av small mismatch"
3596 );
3597 assert!((got[0] - 5.0).abs() < 1e-12, "fast_av[0] should be 5");
3599 assert!((got[1] - 11.0).abs() < 1e-12, "fast_av[1] should be 11");
3601 }
3602
3603 #[test]
3605 fn fast_av_large_matches_ndarray_dot() {
3606 let n = 50usize;
3607 let p = 40usize;
3608 let mut a = Array2::<f64>::zeros((n, p));
3609 let mut v = Array1::<f64>::zeros(p);
3610 let mut state = 0xFEED_FACE_ABCD_EF01u64;
3611 let next = |s: &mut u64| -> f64 {
3612 *s ^= *s << 13;
3613 *s ^= *s >> 7;
3614 *s ^= *s << 17;
3615 ((*s >> 11) as f64 / ((1u64 << 53) as f64)) - 0.5
3616 };
3617 for v in a.iter_mut() {
3618 *v = next(&mut state);
3619 }
3620 for x in v.iter_mut() {
3621 *x = next(&mut state);
3622 }
3623 let got = fast_av(&a, &v);
3624 let want = a.dot(&v);
3625 assert!(
3626 max_abs_diff_1d(&got, &want) < 1e-9,
3627 "fast_av large mismatch"
3628 );
3629 }
3630
3631 #[test]
3632 fn standard_fma_av_matches_ndarray_dot() {
3633 let n = 73usize;
3634 let p = 257usize;
3635 let a = Array2::from_shape_fn((n, p), |(i, j)| {
3636 ((i + 3 * j + 1) as f64).sin() / (j + 1) as f64
3637 });
3638 let v = Array1::from_shape_fn(p, |j| ((2 * j + 1) as f64).cos());
3639 let want = a.dot(&v);
3640 let mut got = Array1::<f64>::zeros(n);
3641 fast_av_standard_view_into(&a, &v, got.view_mut());
3642 assert!(
3643 max_abs_diff_1d(&got, &want) < 1e-12,
3644 "standard-FMA matrix-vector product mismatch"
3645 );
3646 }
3647
3648 #[test]
3650 fn fast_atv_small_matches_ndarray_dot() {
3651 let a = array![[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]];
3652 let v = array![1.0, 0.0, -1.0];
3653 let got = fast_atv(&a, &v);
3654 let want = a.t().dot(&v);
3655 assert!(
3657 max_abs_diff_1d(&got, &want) < 1e-12,
3658 "fast_atv small mismatch"
3659 );
3660 assert!((got[0] - (-4.0)).abs() < 1e-12, "fast_atv[0]");
3661 assert!((got[1] - (-4.0)).abs() < 1e-12, "fast_atv[1]");
3662 }
3663
3664 #[test]
3666 fn fast_atv_large_matches_ndarray_dot() {
3667 let n = 50usize;
3668 let p = 40usize;
3669 let mut a = Array2::<f64>::zeros((n, p));
3670 let mut v = Array1::<f64>::zeros(n);
3671 let mut state = 0x1234_ABCD_5678_EF90u64;
3672 let next = |s: &mut u64| -> f64 {
3673 *s ^= *s << 13;
3674 *s ^= *s >> 7;
3675 *s ^= *s << 17;
3676 ((*s >> 11) as f64 / ((1u64 << 53) as f64)) - 0.5
3677 };
3678 for x in a.iter_mut() {
3679 *x = next(&mut state);
3680 }
3681 for x in v.iter_mut() {
3682 *x = next(&mut state);
3683 }
3684 let got = fast_atv(&a, &v);
3685 let want = a.t().dot(&v);
3686 assert!(
3687 max_abs_diff_1d(&got, &want) < 1e-9,
3688 "fast_atv large mismatch"
3689 );
3690 }
3691
3692 #[test]
3695 fn fast_xt_diag_y_small_matches_manual() {
3696 let x = array![[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]];
3697 let d = array![2.0, 0.5, 1.0];
3698 let y = array![[7.0, 8.0, 9.0], [10.0, 11.0, 12.0], [13.0, 14.0, 15.0]];
3699 let got = fast_xt_diag_y(&x, &d, &y);
3700 let diag_y = {
3702 let mut dy = Array2::<f64>::zeros(y.dim());
3703 for i in 0..3 {
3704 for j in 0..3 {
3705 dy[[i, j]] = d[i] * y[[i, j]];
3706 }
3707 }
3708 dy
3709 };
3710 let want = x.t().dot(&diag_y);
3711 assert!(
3712 max_abs_diff(&got, &want) < 1e-12,
3713 "fast_xt_diag_y small mismatch"
3714 );
3715 assert_eq!(got.dim(), (2, 3));
3716 }
3717
3718 #[inline]
3725 fn two_prod(a: f64, b: f64) -> (f64, f64) {
3726 let p = a * b;
3727 let e = a.mul_add(b, -p);
3728 (p, e)
3729 }
3730
3731 #[inline]
3732 fn two_sum(a: f64, b: f64) -> (f64, f64) {
3733 let s = a + b;
3734 let bb = s - a;
3735 let e = (a - (s - bb)) + (b - bb);
3736 (s, e)
3737 }
3738
3739 fn grow_expansion(e: &mut Vec<f64>, mut q: f64) {
3741 for h in e.iter_mut() {
3742 let (s, err) = two_sum(*h, q);
3743 *h = err;
3744 q = s;
3745 }
3746 if q != 0.0 {
3747 e.push(q);
3748 }
3749 }
3750
3751 fn exact_dot(a: &[f64], b: &[f64]) -> f64 {
3756 let mut e: Vec<f64> = Vec::new();
3757 for (&x, &y) in a.iter().zip(b.iter()) {
3758 let (p, ep) = two_prod(x, y);
3759 grow_expansion(&mut e, p);
3760 grow_expansion(&mut e, ep);
3761 }
3762 e.iter().fold(0.0f64, |acc, &c| acc + c)
3765 }
3766
3767 fn dd_dot(a: &[f64], b: &[f64]) -> f64 {
3771 let (mut s, mut c) = (0.0f64, 0.0f64);
3772 for (&x, &y) in a.iter().zip(b.iter()) {
3773 let (p, ep) = two_prod(x, y);
3774 let (s2, es) = two_sum(s, p);
3775 s = s2;
3776 c += ep + es;
3777 }
3778 s + c
3779 }
3780
3781 fn naive_dot(a: &[f64], b: &[f64]) -> f64 {
3782 let mut acc = 0.0f64;
3783 for (&x, &y) in a.iter().zip(b.iter()) {
3784 acc += x * y;
3785 }
3786 acc
3787 }
3788
3789 fn ill_conditioned_pair(len: usize, seed: u64) -> (Vec<f64>, Vec<f64>) {
3792 let mut s = seed | 1;
3793 let mut next = || {
3794 s ^= s << 13;
3795 s ^= s >> 7;
3796 s ^= s << 17;
3797 (s >> 11) as f64 / ((1u64 << 53) as f64) - 0.5
3798 };
3799 let mut a = Vec::with_capacity(len);
3800 let mut b = Vec::with_capacity(len);
3801 for i in 0..len {
3802 let scale = 10f64.powi((i % 17) as i32 - 8);
3804 let sign = if i % 2 == 0 { 1.0 } else { -1.0 };
3805 a.push(sign * next() * scale);
3806 b.push(next() * scale);
3807 }
3808 (a, b)
3809 }
3810
3811 #[cfg(target_arch = "x86_64")]
3819 #[test]
3820 fn fma_avx2_kernel_variants_are_bit_identical_to_the_baseline_bodies() {
3821 assert!(
3822 super::fma_avx2_available(),
3823 "this machine reports no fma/avx2: the variant path cannot be exercised here"
3824 );
3825 for seed in 0..64u64 {
3826 let len = 200 + (seed as usize % 57);
3827 let (a, b) = ill_conditioned_pair(len, 0x9E37_79B9 ^ seed.wrapping_mul(2654435761));
3828 let (dot_v, std_v) = unsafe {
3831 (
3832 super::fma_dot_fma_avx2(&a, &b),
3833 super::standard_fma_dot_fma_avx2(&a, &b),
3834 )
3835 };
3836 assert_eq!(dot_v.to_bits(), super::fma_dot_body(&a, &b).to_bits(), "fma_dot seed={seed}");
3837 assert_eq!(
3838 std_v.to_bits(),
3839 super::standard_fma_dot_body(&a, &b).to_bits(),
3840 "standard_fma_dot seed={seed}"
3841 );
3842 let p = 7;
3845 let rows: Vec<f64> = (0..len * p).map(|k| a[k % len] * (1.0 + (k % 3) as f64)).collect();
3846 let mut acc_body = vec![0.0f64; p];
3847 let mut acc_var = vec![0.0f64; p];
3848 super::atv_block_accumulate_body(&rows, &b, &mut acc_body);
3849 unsafe { super::atv_block_accumulate_fma_avx2(&rows, &b, &mut acc_var) };
3851 assert_eq!(
3852 acc_var.iter().map(|v| v.to_bits()).collect::<Vec<_>>(),
3853 acc_body.iter().map(|v| v.to_bits()).collect::<Vec<_>>(),
3854 "atv block seed={seed}"
3855 );
3856 let mut y_body = b.clone();
3857 let mut y_var = b.clone();
3858 super::fma_axpy_into_body(a[0], &a, &mut y_body);
3859 unsafe { super::fma_axpy_into_fma_avx2(a[0], &a, &mut y_var) };
3861 assert_eq!(
3862 y_var.iter().map(|v| v.to_bits()).collect::<Vec<_>>(),
3863 y_body.iter().map(|v| v.to_bits()).collect::<Vec<_>>(),
3864 "axpy seed={seed}"
3865 );
3866 }
3867 }
3868
3869 #[test]
3872 fn fma_dot_beats_naive_accuracy() {
3873 let mut fma_total = 0.0f64;
3874 let mut naive_total = 0.0f64;
3875 let mut strict_wins = 0;
3876 for seed in 0..64u64 {
3877 let len = 200 + (seed as usize % 57);
3878 let (a, b) = ill_conditioned_pair(len, 0x9E37_79B9 ^ seed.wrapping_mul(2654435761));
3879 let truth = exact_dot(&a, &b);
3880 let fe = (super::fma_dot(&a, &b) - truth).abs();
3881 let ne = (naive_dot(&a, &b) - truth).abs();
3882 let floor = 8.0 * f64::EPSILON * truth.abs();
3886 assert!(
3887 fe <= ne * (1.0 + 1e-6) + floor,
3888 "fma_dot worse than naive: seed={seed} fma_err={fe:.3e} naive_err={ne:.3e}",
3889 );
3890 if fe < ne {
3891 strict_wins += 1;
3892 }
3893 fma_total += fe;
3894 naive_total += ne;
3895 }
3896 assert!(
3897 fma_total < naive_total,
3898 "fma_dot aggregate error {fma_total:.3e} not below naive {naive_total:.3e}",
3899 );
3900 assert!(
3901 strict_wins >= 40,
3902 "expected fma_dot to strictly win the majority; only {strict_wins}/64",
3903 );
3904 }
3905
3906 #[test]
3941 fn fast_atv_blocked_beats_naive_accuracy() {
3942 let n = 200_003usize;
3943 let p = 8usize;
3947 let mut s = 0xD1B5_4A32u64;
3948 let mut next = || {
3949 s ^= s << 13;
3950 s ^= s >> 7;
3951 s ^= s << 17;
3952 (s >> 11) as f64 / ((1u64 << 53) as f64) - 0.5
3953 };
3954 let mut x = Array2::<f64>::zeros((n, p));
3955 let mut v = Array1::<f64>::zeros(n);
3956 for i in 0..n {
3957 let scale = 10f64.powi((i % 17) as i32 - 8);
3958 v[i] = if i % 2 == 0 { scale } else { -scale } * next();
3959 for j in 0..p {
3960 x[[i, j]] = next() * scale;
3961 }
3962 }
3963 let got = fast_atv(&x, &v);
3964 let vv: Vec<f64> = v.to_vec();
3965
3966 let mut table: Vec<(usize, f64, f64, f64)> = Vec::with_capacity(p);
3969 for j in 0..p {
3970 let col: Vec<f64> = (0..n).map(|i| x[[i, j]]).collect();
3971 let truth = dd_dot(&col, &vv);
3972 let naive = naive_dot(&col, &vv);
3973 table.push((j, truth, (got[j] - truth).abs(), (naive - truth).abs()));
3974 }
3975 let report: String = table
3976 .iter()
3977 .map(|&(j, truth, ge, ne)| {
3978 format!(" col {j}: truth={truth:.6e} blocked_err={ge:.3e} naive_err={ne:.3e}\n")
3979 })
3980 .collect();
3981
3982 let mut blocked_total = 0.0f64;
3983 let mut naive_total = 0.0f64;
3984 for &(j, truth, ge, ne) in &table {
3985 assert!(
3986 ne > 0.0,
3987 "col {j}: naive baseline error is exactly 0.0, so this column \
3988 cannot discriminate the two reductions - the fixture is no \
3989 longer ill-conditioned\n{report}",
3990 );
3991 assert!(
4006 ge <= 64.0 * f64::EPSILON * truth.abs(),
4007 "col {j}: blocked err {ge:.3e} exceeds naive {ne:.3e}\n{report}",
4008 );
4009 blocked_total += ge;
4010 naive_total += ne;
4011 }
4012 assert!(
4013 2.0 * blocked_total < naive_total,
4014 "blocked aggregate error {blocked_total:.3e} is not at least 2x \
4015 below naive {naive_total:.3e}; a 391-block pairwise reduction \
4016 should be roughly sqrt(391) = 20x better\n{report}",
4017 );
4018 }
4019
4020 #[test]
4023 fn fast_av_strided_input_matches_ndarray() {
4024 let mut base = Array2::<f64>::zeros((40, 60));
4025 let mut s = 0x0BAD_F00Du64;
4026 let mut next = || {
4027 s ^= s << 13;
4028 s ^= s >> 7;
4029 s ^= s << 17;
4030 (s >> 11) as f64 / ((1u64 << 53) as f64) - 0.5
4031 };
4032 for x in base.iter_mut() {
4033 *x = next();
4034 }
4035 let a = base.t();
4037 let mut v = Array1::<f64>::zeros(40);
4038 for x in v.iter_mut() {
4039 *x = next();
4040 }
4041 let got = fast_av(&a, &v);
4042 let want = a.dot(&v);
4043 assert!(
4044 max_abs_diff_1d(&got, &want) < 1e-11,
4045 "strided fast_av mismatch (fallback path)",
4046 );
4047 }
4048
4049 #[test]
4065 fn faer_sequential_scope_sets_seq_inside_and_restores_after() {
4066 crate::test_support::with_global_parallelism_serialized(|| {
4067 let baseline = faer::get_global_parallelism();
4068 faer::set_global_parallelism(Par::rayon(4));
4071 assert_eq!(
4072 faer::get_global_parallelism(),
4073 Par::rayon(4),
4074 "baseline must be the parallel policy we just set",
4075 );
4076
4077 {
4078 let faer_seq_guard = FaerSequentialScope::enter();
4079 assert_eq!(
4080 faer::get_global_parallelism(),
4081 Par::Seq,
4082 "faer must be pinned to Par::Seq inside the scope",
4083 );
4084
4085 {
4087 let faer_seq_inner_guard = FaerSequentialScope::enter();
4088 assert_eq!(
4089 faer::get_global_parallelism(),
4090 Par::Seq,
4091 "nested scope stays Par::Seq",
4092 );
4093 drop(faer_seq_inner_guard);
4094 }
4095 assert_eq!(
4096 faer::get_global_parallelism(),
4097 Par::Seq,
4098 "inner drop must not restore while outer scope is still live",
4099 );
4100 drop(faer_seq_guard);
4101 }
4102
4103 assert_eq!(
4104 faer::get_global_parallelism(),
4105 Par::rayon(4),
4106 "outermost drop must restore the pre-scope parallelism policy",
4107 );
4108
4109 let observed = with_faer_sequential(|| faer::get_global_parallelism());
4111 assert_eq!(
4112 observed,
4113 Par::Seq,
4114 "with_faer_sequential runs body under Seq"
4115 );
4116 assert_eq!(
4117 faer::get_global_parallelism(),
4118 Par::rayon(4),
4119 "with_faer_sequential restores after the body returns",
4120 );
4121
4122 faer::set_global_parallelism(baseline);
4124 });
4125 }
4126}
4127
4128#[cfg(test)]
4133mod parallelism_snapshot_2738_tests {
4134 use super::*;
4135
4136 #[test]
4137 fn captured_snapshot_is_self_consistent() {
4138 let snapshot =
4144 crate::test_support::with_global_parallelism_serialized(ParallelismSnapshot::capture);
4145 assert!(
4146 snapshot.inconsistency().is_none(),
4147 "the live thread configuration disagrees with itself: {} ({snapshot})",
4148 snapshot.inconsistency().unwrap_or_default(),
4149 );
4150 }
4151
4152 #[test]
4158 fn inconsistent_configurations_are_reported() {
4159 let contradictory = ParallelismSnapshot::from_parts(Par::rayon(4), 4, 1, Some(4));
4163 assert!(
4164 contradictory.inconsistency().is_some(),
4165 "a live FaerSequentialScope with non-sequential faer must be flagged: \
4166 {contradictory}",
4167 );
4168
4169 assert!(
4173 ParallelismSnapshot::from_parts(Par::Seq, 0, 0, Some(1))
4174 .inconsistency()
4175 .is_some(),
4176 "a zero-wide rayon pool must be flagged",
4177 );
4178 assert!(
4179 ParallelismSnapshot::from_parts(Par::Seq, 1, 0, Some(0))
4180 .inconsistency()
4181 .is_some(),
4182 "zero cores available to the process must be flagged",
4183 );
4184
4185 assert!(
4188 ParallelismSnapshot::from_parts(Par::rayon(4), 4, 0, Some(8))
4189 .inconsistency()
4190 .is_none(),
4191 "a wide pool with no sequential scope is consistent",
4192 );
4193 assert!(
4194 ParallelismSnapshot::from_parts(Par::Seq, 4, 2, None)
4195 .inconsistency()
4196 .is_none(),
4197 "a pinned scope on a wide pool is consistent, and an unavailable core \
4198 count is not itself an inconsistency",
4199 );
4200 }
4201
4202 #[test]
4205 fn a_sequential_pin_changes_the_snapshot() {
4206 let pinned = crate::test_support::with_global_parallelism_serialized(|| {
4207 with_faer_sequential(ParallelismSnapshot::capture)
4208 });
4209 assert!(
4210 pinned.faer_global_sequential,
4211 "inside a FaerSequentialScope the snapshot must report faer sequential: \
4212 {pinned}",
4213 );
4214 assert_eq!(
4215 pinned.faer_global_degree, 1,
4216 "a sequential pin is one thread of numerics: {pinned}",
4217 );
4218 assert!(
4219 pinned.faer_sequential_scope_depth >= 1,
4220 "the scope that did the pinning must be visible in the depth: {pinned}",
4221 );
4222 assert!(
4223 pinned.inconsistency().is_none(),
4224 "a pinned snapshot must still be self-consistent: {pinned}",
4225 );
4226 assert_eq!(
4229 pinned.rayon_current_num_threads,
4230 rayon::current_num_threads(),
4231 "the pin must not be mistaken for a narrower rayon pool",
4232 );
4233 }
4234
4235 #[test]
4239 fn rendering_carries_every_field() {
4240 let snapshot = ParallelismSnapshot::from_parts(Par::rayon(3), 5, 2, Some(7));
4241 let rendered = snapshot.to_string();
4242 for field in [
4243 "rayon_current_num_threads=5",
4244 "faer_global_sequential=false",
4245 "faer_global_degree=3",
4246 "faer_sequential_scope_depth=2",
4247 "process_available_parallelism=7",
4248 ] {
4249 assert!(
4250 rendered.contains(field),
4251 "the rendered snapshot dropped `{field}`: {rendered}",
4252 );
4253 }
4254
4255 let unavailable = ParallelismSnapshot::from_parts(Par::Seq, 1, 0, None);
4256 assert!(
4257 unavailable.to_string().contains("unavailable"),
4258 "a missing core count must say so rather than render as a number: \
4259 {unavailable}",
4260 );
4261 }
4262
4263}
4264
4265#[cfg(test)]
4266mod eigh_ordering_contract_tests {
4267 use super::*;
4268 use ndarray::Array2;
4269
4270 #[test]
4296 fn eigh_returns_eigenvalues_in_ascending_order() {
4297 fn hashed_unit(seed: u64) -> f64 {
4298 let mut z = seed.wrapping_add(0x9E37_79B9_7F4A_7C15);
4299 z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
4300 z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
4301 z ^= z >> 31;
4302 ((z >> 11) as f64 / (1u64 << 53) as f64) * 2.0 - 1.0
4303 }
4304
4305 for &n in &[1_usize, 2, 3, 5, 6, 9, 17, 36] {
4310 for seed in 0..8_u64 {
4311 for &scale in &[1.0_f64, 1.0e-9, 1.0e9] {
4312 let mut m = Array2::<f64>::zeros((n, n));
4313 let mut k = seed.wrapping_mul(1_000_003).wrapping_add(n as u64);
4314 for i in 0..n {
4315 for j in 0..=i {
4316 k = k.wrapping_add(0x1234_5678);
4317 let value = hashed_unit(k) * scale;
4318 m[[i, j]] = value;
4319 m[[j, i]] = value;
4320 }
4321 }
4322 let (values, _) = m.eigh(Side::Lower).expect("eigendecomposition");
4323 assert_eq!(values.len(), n, "n={n}: one eigenvalue per dimension");
4324 for w in 1..n {
4325 assert!(
4326 values[w - 1] <= values[w],
4327 "n={n} seed={seed} scale={scale:e}: eigenvalues are NOT ascending at \
4328 index {w} ({:e} then {:e}). `cluster_stable_eigh` scans for RUNS of \
4329 equal eigenvalues and would silently stop finding degenerate clusters.",
4330 values[w - 1],
4331 values[w]
4332 );
4333 }
4334 }
4335 }
4336 }
4337 }
4338
4339 #[test]
4344 fn equal_eigenvalues_are_returned_adjacent() {
4345 let d = ndarray::arr1(&[2.0_f64, 7.0, 2.0, 7.0, 2.0]);
4348 let n = d.len();
4349 let v = ndarray::arr1(&[1.0_f64, -2.0, 3.0, -4.0, 5.0]);
4352 let vtv: f64 = v.iter().map(|x| x * x).sum();
4353 let mut q = Array2::<f64>::zeros((n, n));
4354 for i in 0..n {
4355 q[[i, i]] = 1.0;
4356 }
4357 for i in 0..n {
4358 for j in 0..n {
4359 q[[i, j]] -= 2.0 * v[i] * v[j] / vtv;
4360 }
4361 }
4362 let mut a = Array2::<f64>::zeros((n, n));
4363 for i in 0..n {
4364 for j in 0..n {
4365 let mut acc = 0.0;
4366 for k in 0..n {
4367 acc += q[[i, k]] * d[k] * q[[j, k]];
4368 }
4369 a[[i, j]] = acc;
4370 }
4371 }
4372 let (values, _) = a.eigh(Side::Lower).expect("eigendecomposition");
4373 let low = values.iter().filter(|v| (**v - 2.0).abs() < 1.0e-9).count();
4377 let high = values.iter().filter(|v| (**v - 7.0).abs() < 1.0e-9).count();
4378 assert_eq!(low, 3, "planted multiplicity 3 at lambda=2, got {values:?}");
4379 assert_eq!(high, 2, "planted multiplicity 2 at lambda=7, got {values:?}");
4380 for w in 0..3 {
4381 assert!(
4382 (values[w] - 2.0).abs() < 1.0e-9,
4383 "the three lambda=2 eigenvalues must occupy indices 0..3 contiguously, \
4384 or `cluster_stable_eigh`'s run scan splits the cluster: {values:?}"
4385 );
4386 }
4387 for w in 3..5 {
4388 assert!(
4389 (values[w] - 7.0).abs() < 1.0e-9,
4390 "the two lambda=7 eigenvalues must occupy indices 3..5 contiguously: {values:?}"
4391 );
4392 }
4393 }
4394}