Skip to main content

gam_linalg/
faer_ndarray.rs

1use dyn_stack::{MemBuffer, MemStack};
2use faer::diag::{Diag, DiagRef};
3use faer::linalg::solvers::{self, Solve};
4pub use faer::linalg::solvers::{
5    Lblt as FaerLblt, Ldlt as FaerLdlt, Llt as FaerLlt, Solve as FaerSolve,
6};
7use faer::linalg::svd::{self, ComputeSvdVectors};
8use faer::prelude::ReborrowMut;
9use faer::{Conj, Mat, MatMut, MatRef, Par, Side, Unbind, get_global_parallelism};
10use ndarray::{Array1, Array2, ArrayBase, ArrayViewMut1, Data, Ix1, Ix2};
11use std::marker::PhantomData;
12use std::panic::{AssertUnwindSafe, catch_unwind};
13use thiserror::Error;
14
15const RRQR_RANK_ALPHA: f64 = 100.0;
16
17thread_local! {
18    static NESTED_PARALLEL_DEPTH: std::cell::Cell<usize> = const { std::cell::Cell::new(0) };
19}
20
21struct NestedParallelGuard;
22
23impl NestedParallelGuard {
24    #[inline]
25    fn enter() -> Self {
26        NESTED_PARALLEL_DEPTH.with(|depth| depth.set(depth.get().saturating_add(1)));
27        Self
28    }
29}
30
31impl Drop for NestedParallelGuard {
32    #[inline]
33    fn drop(&mut self) {
34        NESTED_PARALLEL_DEPTH.with(|depth| depth.set(depth.get().saturating_sub(1)));
35    }
36}
37
38/// Run `body` with the current thread marked as inside a data-parallel row
39/// region, so any faer GEMM it issues (directly or transitively) pins to
40/// `Par::Seq` via [`effective_global_parallelism`] instead of re-fanning the
41/// global Rayon pool. The guard is held for exactly the duration of `body` and
42/// dropped on return — including early `?` returns from inside `body`, since the
43/// guard lives in this function's frame.
44///
45/// Call this from the per-chunk/per-row closure of an `into_par_iter` whose body
46/// performs GEMM, to prevent the Rayon-pool × faer-pool oversubscription.
47#[inline]
48pub fn with_nested_parallel<T>(body: impl FnOnce() -> T) -> T {
49    let guard = NestedParallelGuard::enter();
50    let out = body();
51    drop(guard);
52    out
53}
54
55/// `true` when the current thread is inside at least one [`NestedParallelGuard`]
56/// scope, i.e. a parallel row reduction is already in flight on this thread.
57#[inline]
58pub fn in_nested_parallel_region() -> bool {
59    NESTED_PARALLEL_DEPTH.with(|depth| depth.get() > 0)
60}
61
62/// faer parallelism policy that respects nested data-parallel regions: returns
63/// faer's global policy at the top level, but `Par::Seq` once a
64/// [`NestedParallelGuard`] is active so a GEMM issued from inside a parallel row
65/// fan-out does not multiply the live thread count against the outer pool.
66///
67/// Use this in place of `faer::get_global_parallelism()` for any matmul that can
68/// be reached from inside a row-parallel closure.
69#[inline]
70pub fn effective_global_parallelism() -> Par {
71    if in_nested_parallel_region() {
72        Par::Seq
73    } else {
74        get_global_parallelism()
75    }
76}
77
78/// Process-global depth counter + saved parallelism for [`FaerSequentialScope`].
79///
80/// The `effective_global_parallelism` / [`NestedParallelGuard`] pair only pins
81/// the codebase's OWN `matmul` calls to `Par::Seq`; it CANNOT reach faer's
82/// high-level factorization/solve entry points (`Llt::new`, `Solve::solve`, SVD,
83/// col-pivoted QR), which read `faer::get_global_parallelism()` internally and
84/// have no per-call parallelism argument. When such a solver runs from inside a
85/// Rayon worker (e.g. the topology race fans candidate fits into per-candidate
86/// pools via `run_topology_race_parallel`), faer's default `Par::rayon(0)`
87/// dispatches the factorization through its `spindle` barrier pool, which
88/// `rayon::scope`-spawns as many tasks as the pool has threads and waits for all
89/// of them at a barrier. Under thread oversubscription those worker slots are
90/// already occupied by the outer fan-out, so the barrier never completes and the
91/// fit parks at 0% CPU — the #2074 K=1 `sae_manifold_fit` deadlock.
92///
93/// [`FaerSequentialScope`] closes that hole by pinning faer's PROCESS-GLOBAL
94/// parallelism to `Par::Seq` around the nested solve, so every faer solver it
95/// reaches stays single-threaded and never spawns a nested barrier pool. The
96/// codebase engineers its faer reductions to be parallelism-invariant
97/// (`tests_parallelism_invariance_1557` asserts byte-identical `Par::Seq` vs
98/// `Par::rayon` output), so collapsing to sequential is bit-for-bit neutral.
99static FAER_SEQ_STATE: std::sync::Mutex<FaerSeqState> = std::sync::Mutex::new(FaerSeqState {
100    depth: 0,
101    saved: None,
102});
103
104struct FaerSeqState {
105    depth: usize,
106    saved: Option<Par>,
107}
108
109/// RAII guard that pins faer's process-global parallelism to [`Par::Seq`] for its
110/// lifetime and restores the previous setting when the LAST live guard drops.
111///
112/// The guard is depth-counted across threads: overlapping guards (e.g. several
113/// topology-race candidates fitting concurrently) all observe `Par::Seq`, and the
114/// prior policy is restored exactly once, when the outermost guard exits. Setting
115/// and restoring happen under the state mutex so the `depth == 0` transition is
116/// atomic with the `set_global_parallelism` call.
117#[must_use = "the sequential scope only holds while the guard is alive"]
118pub struct FaerSequentialScope {
119    _private: (),
120}
121
122impl FaerSequentialScope {
123    /// Enter the scope, forcing faer to `Par::Seq` on the `0 -> 1` transition.
124    pub fn enter() -> Self {
125        let mut state = FAER_SEQ_STATE
126            .lock()
127            .unwrap_or_else(std::sync::PoisonError::into_inner);
128        if state.depth == 0 {
129            state.saved = Some(get_global_parallelism());
130            faer::set_global_parallelism(Par::Seq);
131        }
132        state.depth += 1;
133        Self { _private: () }
134    }
135}
136
137impl Drop for FaerSequentialScope {
138    fn drop(&mut self) {
139        let mut state = FAER_SEQ_STATE
140            .lock()
141            .unwrap_or_else(std::sync::PoisonError::into_inner);
142        state.depth -= 1;
143        if state.depth == 0 {
144            if let Some(par) = state.saved.take() {
145                faer::set_global_parallelism(par);
146            }
147        }
148    }
149}
150
151/// Run `body` with faer pinned to `Par::Seq` (see [`FaerSequentialScope`]). Use
152/// this to wrap a fit/solve that runs inside a Rayon worker so faer's high-level
153/// solvers never fan a nested `spindle` barrier pool into an already-saturated
154/// Rayon pool.
155#[inline]
156pub fn with_faer_sequential<T>(body: impl FnOnce() -> T) -> T {
157    let faer_seq_guard = FaerSequentialScope::enter();
158    let out = body();
159    drop(faer_seq_guard);
160    out
161}
162
163#[derive(Debug, Error)]
164pub enum FaerLinalgError {
165    #[error("Factorization failed in {context}")]
166    FactorizationFailed { context: &'static str },
167    #[error("SVD failed to converge in {context}")]
168    SvdNoConvergence { context: &'static str },
169    #[error("Self-adjoint eigendecomposition input contains non-finite values in {context}")]
170    SelfAdjointEigenNonFiniteInput { context: &'static str },
171    #[error("Self-adjoint eigendecomposition failed: {0:?}")]
172    SelfAdjointEigen(solvers::EvdError),
173    #[error("Cholesky factorization failed: {0:?}")]
174    Cholesky(solvers::LltError),
175    #[error("LDLT factorization failed: {0:?}")]
176    Ldlt(solvers::LdltError),
177}
178
179pub enum FaerSymmetricFactor {
180    Llt(FaerLlt<f64>),
181    Ldlt(FaerLdlt<f64>),
182    Lblt(FaerLblt<f64>),
183}
184
185#[inline]
186pub fn cholesky_factor_logdet(factor: MatRef<'_, f64>) -> f64 {
187    2.0 * diagonal_log_sum(factor.diagonal())
188}
189
190#[inline]
191fn diagonal_log_sum(diagonal: DiagRef<'_, f64>) -> f64 {
192    diagonal
193        .column_vector()
194        .iter()
195        .map(|&x| x.ln())
196        .sum::<f64>()
197}
198
199impl FaerSymmetricFactor {
200    /// Returns the dimension of the factorized square matrix.
201    #[inline]
202    pub fn n(&self) -> usize {
203        use faer::linalg::solvers::ShapeCore;
204        match self {
205            FaerSymmetricFactor::Llt(f) => f.nrows(),
206            FaerSymmetricFactor::Ldlt(f) => f.nrows(),
207            FaerSymmetricFactor::Lblt(f) => f.nrows(),
208        }
209    }
210
211    #[inline]
212    pub fn solve(&self, rhs: MatRef<'_, f64>) -> Mat<f64> {
213        match self {
214            FaerSymmetricFactor::Llt(f) => f.solve(rhs),
215            FaerSymmetricFactor::Ldlt(f) => f.solve(rhs),
216            FaerSymmetricFactor::Lblt(f) => f.solve(rhs),
217        }
218    }
219
220    #[inline]
221    pub fn solve_in_place(&self, rhs: MatMut<'_, f64>) {
222        match self {
223            FaerSymmetricFactor::Llt(f) => f.solve_in_place(rhs),
224            FaerSymmetricFactor::Ldlt(f) => f.solve_in_place(rhs),
225            FaerSymmetricFactor::Lblt(f) => f.solve_in_place(rhs),
226        }
227    }
228}
229
230impl crate::matrix::FactorizedSystem for FaerSymmetricFactor {
231    fn solve(&self, rhs: &Array1<f64>) -> Result<Array1<f64>, String> {
232        let mut out = rhs.clone();
233        let mut out_mat = array1_to_col_matmut(&mut out);
234        self.solve_in_place(out_mat.as_mut());
235        if !out.iter().all(|v| v.is_finite()) {
236            return Err("symmetric factor solve produced non-finite values".to_string());
237        }
238        Ok(out)
239    }
240
241    fn solvemulti(&self, rhs: &Array2<f64>) -> Result<Array2<f64>, String> {
242        let mut out = Array2::<f64>::zeros(rhs.raw_dim());
243        for j in 0..rhs.ncols() {
244            for i in 0..rhs.nrows() {
245                out[[i, j]] = rhs[[i, j]];
246            }
247        }
248        let mut out_mat = array2_to_matmut(&mut out);
249        self.solve_in_place(out_mat.as_mut());
250        if !out.iter().all(|v| v.is_finite()) {
251            return Err("symmetric factor multi-solve produced non-finite values".to_string());
252        }
253        Ok(out)
254    }
255
256    fn logdet(&self) -> f64 {
257        match self {
258            FaerSymmetricFactor::Llt(f) => cholesky_factor_logdet(f.L()),
259            FaerSymmetricFactor::Ldlt(f) => diagonal_log_sum(f.D()),
260            FaerSymmetricFactor::Lblt(..) => {
261                // lblt doesn't easily expose diagonal determinant. Fallback to sparse or other representations if needed, but typically Lblt is indefinite!
262                // Actually faer doesn't easily expose lblt logdet since it has 2x2 blocks.
263                // For our ML systems, if we dropped to LBLT, the matrix was indefinite and logdet is ill-defined (or complex).
264                f64::NAN
265            }
266        }
267    }
268}
269
270/// Factorize a symmetric system with LLT -> LDLT -> LBLT fallback.
271#[inline]
272pub fn factorize_symmetricwith_fallback(
273    matrix: MatRef<'_, f64>,
274    side: Side,
275) -> Result<FaerSymmetricFactor, FaerLinalgError> {
276    if let Ok(llt) = FaerLlt::new(matrix, side) {
277        return Ok(FaerSymmetricFactor::Llt(llt));
278    }
279    let ldlt_err = match FaerLdlt::new(matrix, side) {
280        Ok(ldlt) => return Ok(FaerSymmetricFactor::Ldlt(ldlt)),
281        Err(err) => err,
282    };
283    let lblt = catch_unwind(AssertUnwindSafe(|| FaerLblt::new(matrix, side)))
284        .map_err(|_| FaerLinalgError::Ldlt(ldlt_err))?;
285    Ok(FaerSymmetricFactor::Lblt(lblt))
286}
287
288#[inline]
289const fn should_use_faer_matmul(m: usize, n: usize, k: usize) -> bool {
290    // Small, centralized dispatch policy:
291    // - stay on ndarray for tiny products to avoid setup overhead,
292    // - switch to faer GEMM/GEMV for moderate+ sizes.
293    const MIN_DIM: usize = 32;
294    const MIN_FLOP_SCALE: usize = 64 * 64;
295    (m >= MIN_DIM || n >= MIN_DIM || k >= MIN_DIM)
296        && m.saturating_mul(n).saturating_mul(k) >= MIN_FLOP_SCALE
297}
298
299#[inline]
300pub fn matmul_parallelism(m: usize, n: usize, k: usize) -> Par {
301    // Prefer a work-based policy over per-dimension thresholds.
302    // Tall/skinny products (e.g. N x p with large N, modest p) should still
303    // parallelize when total work is high.
304    const PAR_MIN_FLOP_SCALE: usize = 2_000_000;
305    const PAR_MIN_LONG_DIM: usize = 256;
306    let flop_scale = m.saturating_mul(n).saturating_mul(k);
307    let long_dim = m.max(n).max(k);
308    if flop_scale >= PAR_MIN_FLOP_SCALE && long_dim >= PAR_MIN_LONG_DIM {
309        // `effective_global_parallelism` collapses to `Par::Seq` when this GEMM
310        // is reached from inside a `NestedParallelGuard` row region, preventing
311        // the Rayon-pool × faer-pool multiplicative oversubscription.
312        effective_global_parallelism()
313    } else {
314        Par::Seq
315    }
316}
317
318#[inline]
319pub fn array2_to_matmut(array: &mut Array2<f64>) -> MatMut<'_, f64> {
320    let (rows, cols) = array.dim();
321    let strides = array.strides();
322
323    // Check if we can get a pointer.
324    // If the array is contiguous (either C or F order), or simply sliced with strides,
325    // faer can handle it as long as we pass the pointer and strides.
326    // However, as_mut_ptr() requires a mutable reference.
327    // ndarray's as_ptr/as_mut_ptr works for both layouts.
328
329    let s0 = strides[0];
330    let s1 = strides[1];
331
332    // SAFETY: array.as_mut_ptr() is ndarray's logical (0, 0) pointer, and
333    // ndarray's dimensions plus signed element strides describe every initialized
334    // element of this uniquely borrowed Array2 for the returned MatMut lifetime.
335    unsafe { MatMut::from_raw_parts_mut(array.as_mut_ptr(), rows, cols, s0, s1) }
336}
337
338/// Convert an ndarray matrix into row-major nested vectors for serialized
339/// payloads without exposing storage-layout assumptions to callers.
340pub fn array2_to_nested_vec(array: &Array2<f64>) -> Vec<Vec<f64>> {
341    array.rows().into_iter().map(|row| row.to_vec()).collect()
342}
343
344#[inline]
345pub fn array1_to_col_matmut(array: &mut Array1<f64>) -> MatMut<'_, f64> {
346    let len = array.len();
347    let stride = array.strides()[0];
348    // SAFETY: array.as_mut_ptr() is ndarray's logical first-element pointer, and
349    // len plus the signed element stride describe every initialized element of
350    // this uniquely borrowed Array1 for the returned len×1 MatMut lifetime.
351    unsafe {
352        MatMut::from_raw_parts_mut(
353            array.as_mut_ptr(),
354            len,
355            1,
356            stride,
357            0, // col stride irrelevant for 1 column
358        )
359    }
360}
361
362/// Compute A^T * A using faer's SIMD-optimized GEMM.
363/// This is MUCH faster than ndarray's .t().dot() for matrices where n > ~100.
364///
365/// For a matrix A of shape (n, p), this computes the (p, p) result.
366/// Uses a zero-copy view for positive-stride layouts and copies only layouts
367/// with non-positive strides.
368#[inline]
369pub fn fast_ata<S: Data<Elem = f64>>(a: &ArrayBase<S, Ix2>) -> Array2<f64> {
370    let p = a.ncols();
371    let mut out = Array2::<f64>::zeros((p, p));
372    fast_ata_into(a, &mut out);
373    out
374}
375
376/// Compute A^T * A into a pre-allocated output buffer.
377/// `out` must be shaped (p, p) where A is (n, p).
378#[inline]
379pub fn fast_ata_into<S: Data<Elem = f64>>(a: &ArrayBase<S, Ix2>, out: &mut Array2<f64>) {
380    use faer::Accum;
381    use faer::linalg::matmul::triangular::{BlockStructure, matmul as tri_matmul};
382
383    let (n, p) = a.dim();
384    assert_eq!(out.nrows(), p, "output rows must match p");
385    assert_eq!(out.ncols(), p, "output cols must match p");
386
387    if !should_use_faer_matmul(p, p, n) {
388        out.assign(&a.t().dot(a));
389        return;
390    }
391
392    let mut outview = array2_to_matmut(out);
393
394    let aview = FaerArrayView::new(a);
395    let a_ref = aview.as_ref();
396    let a_t = a_ref.transpose();
397    let par = matmul_parallelism(p, p, n);
398    tri_matmul(
399        outview.as_mut(),
400        BlockStructure::TriangularLower,
401        Accum::Replace,
402        a_t,
403        BlockStructure::Rectangular,
404        a_ref,
405        BlockStructure::Rectangular,
406        1.0,
407        par,
408    );
409    // Mirror lower triangle to upper to populate the full symmetric output.
410    for i in 0..p {
411        for j in (i + 1)..p {
412            out[[i, j]] = out[[j, i]];
413        }
414    }
415}
416
417/// Compute A^T * B using faer's SIMD-optimized GEMM.
418/// For A of shape (n, p) and B of shape (n, q), this computes the (p, q) result.
419/// Uses zero-copy views when possible.
420#[inline]
421pub fn fast_atb<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
422    a: &ArrayBase<S1, Ix2>,
423    b: &ArrayBase<S2, Ix2>,
424) -> Array2<f64> {
425    if let Some(out) =
426        crate::gpu_hook::gpu_dispatch().and_then(|d| d.try_fast_atb(a.view(), b.view()))
427    {
428        return out;
429    }
430    let (n_a, p) = a.dim();
431    let q = b.ncols();
432    fast_atb_with_parallelism(a, b, matmul_parallelism(p, q, n_a))
433}
434
435/// Compute A^T * B with an explicit faer parallelism policy for callers that
436/// are already running independent products in an outer Rayon task.
437#[inline]
438pub fn fast_atb_with_parallelism<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
439    a: &ArrayBase<S1, Ix2>,
440    b: &ArrayBase<S2, Ix2>,
441    par: Par,
442) -> Array2<f64> {
443    use faer::linalg::matmul::matmul;
444    use faer::{Accum, Mat};
445
446    let (n_a, p) = a.dim();
447    let (n_b, q) = b.dim();
448    assert_eq!(n_a, n_b, "A and B must have same number of rows");
449
450    // For very small matrices, ndarray might be faster due to less overhead
451    if !should_use_faer_matmul(p, q, n_a) {
452        return a.t().dot(b);
453    }
454
455    let mut result = Mat::<f64>::zeros(p, q);
456
457    let aview = FaerArrayView::new(a);
458    let bview = FaerArrayView::new(b);
459    let a_ref = aview.as_ref();
460    let b_ref = bview.as_ref();
461
462    // dst = A^T * B
463    matmul(
464        result.as_mut(),
465        Accum::Replace,
466        a_ref.transpose(),
467        b_ref,
468        1.0,
469        par,
470    );
471
472    mat_to_array(result.as_ref())
473}
474
475/// Compute A * B^T using faer's SIMD-optimized GEMM.
476/// For A of shape (m, k) and B of shape (n, k), this computes the (m, n) result.
477#[inline]
478pub fn fast_abt<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
479    a: &ArrayBase<S1, Ix2>,
480    b: &ArrayBase<S2, Ix2>,
481) -> Array2<f64> {
482    use faer::linalg::matmul::matmul;
483    use faer::{Accum, Mat};
484
485    let (m, k_a) = a.dim();
486    let (n, k_b) = b.dim();
487    assert_eq!(
488        k_a, k_b,
489        "A and B must have same number of columns for A·Bᵀ"
490    );
491
492    if !should_use_faer_matmul(m, n, k_a) {
493        return a.dot(&b.t());
494    }
495
496    let mut result = Mat::<f64>::zeros(m, n);
497    let aview = FaerArrayView::new(a);
498    let bview = FaerArrayView::new(b);
499    let par = matmul_parallelism(m, n, k_a);
500    matmul(
501        result.as_mut(),
502        Accum::Replace,
503        aview.as_ref(),
504        bview.as_ref().transpose(),
505        1.0,
506        par,
507    );
508    mat_to_array(result.as_ref())
509}
510
511/// Compute A * B using faer's SIMD-optimized GEMM.
512/// For A of shape (n, p) and B of shape (p, q), this computes the (n, q) result.
513/// Uses zero-copy views when possible.
514#[inline]
515pub fn fast_ab<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
516    a: &ArrayBase<S1, Ix2>,
517    b: &ArrayBase<S2, Ix2>,
518) -> Array2<f64> {
519    if let Some(out) =
520        crate::gpu_hook::gpu_dispatch().and_then(|d| d.try_fast_ab(a.view(), b.view()))
521    {
522        return out;
523    }
524    let n = a.nrows();
525    let q = b.ncols();
526    let mut out = Array2::<f64>::zeros((n, q));
527    fast_ab_into(a, b, &mut out);
528    out
529}
530
531// ────────────────────────────────────────────────────────────────────────
532// Compensated / blocked SIMD reduction kernels for the GEMV hot paths.
533//
534// `fast_av` (η = Xβ) and `fast_atv` (Xᵀr — e.g. the penalized-likelihood
535// gradient and REML score) are reduction-bound: every output entry is a sum
536// of products over a long axis. faer's generic GEMM serves them as degenerate
537// single-RHS-column matmuls, whose blocking/setup cost is poorly amortized by
538// one column. For the dominant row-major-contiguous case we use tight hand
539// kernels that are simultaneously
540//   * faster — several independent FMA accumulators expose the
541//     instruction-level parallelism the backend lowers to packed AVX
542//     `vfmadd` lanes, and the row work fans out across the Rayon pool; and
543//   * more accurate — `f64::mul_add` fuses each product into its accumulator
544//     with a single rounding (no rounded intermediate product), the lanes
545//     reduce as a small pairwise tree, and the long Xᵀr reduction is split
546//     into fixed-size row blocks whose partials are combined pairwise,
547//     turning the naive O(n·ε) error growth into ~O((block + log(n/block))·ε).
548//
549// Non-contiguous / non-row-major operands fall back to the faer path, so the
550// numerics only change (improve) on the common standard-layout inputs.
551// ────────────────────────────────────────────────────────────────────────
552
553/// Number of independent FMA accumulator lanes. Eight lanes keep two 256-bit
554/// (`f64x4`) FMA pipelines fed and set the partial-pairwise leaf width.
555const FMA_LANES: usize = 8;
556
557/// FLOP-scale (n·p) below which the kernels stay serial; at or above it, and
558/// only when not already inside a parallel row region, the row loop fans out
559/// across the Rayon pool.
560const KERNEL_PAR_MIN_FLOP: usize = 1 << 18; // 262_144
561
562/// Rows per row-block in [`fast_av_rowmajor_into`]'s parallel fan-out; large
563/// enough to amortize Rayon task overhead over many short row dots.
564const AV_PAR_CHUNK_ROWS: usize = 1024;
565
566/// Rows per reduction block in [`fast_atv_rowmajor_into`]; each block sums its
567/// rows into a private length-p partial and the partials combine pairwise, so
568/// the long-axis rounding error grows with the block size plus the log of the
569/// block count rather than with `n`.
570const ATV_BLOCK_ROWS: usize = 512;
571
572#[inline]
573fn kernel_should_parallelize(n: usize, p: usize) -> bool {
574    !in_nested_parallel_region()
575        && n.saturating_mul(p) >= KERNEL_PAR_MIN_FLOP
576        && rayon::current_num_threads() > 1
577}
578
579/// Compensated dot product (the Ogita–Rump–Oishi *Dot2* error-free transform)
580/// of two equal-length contiguous slices, evaluated over [`FMA_LANES`]
581/// independent compensated accumulators.
582///
583/// For each term the product is split into its rounded value plus the *exact*
584/// product error via `mul_add` (`two_prod`), and added into the running sum via
585/// a branchless `two_sum`, with both rounding errors folded into a
586/// per-lane compensation. The result carries roughly twice the working
587/// precision: its error-vs-truth is bounded by `u·|result| + O(n·u²)·|x|ᵀ|y|`
588/// versus the naive recurrence's `O(n·u)·|x|ᵀ|y|`, i.e. strictly — often by
589/// many orders of magnitude — more accurate. The eight independent lanes keep
590/// the FMA pipelines saturated, and on the GEMV hot paths the extra arithmetic
591/// is hidden under the memory traffic of streaming `X`, so accuracy rises with
592/// no throughput cost.
593#[inline(always)]
594fn fma_dot(a: &[f64], b: &[f64]) -> f64 {
595    assert_eq!(a.len(), b.len(), "fma_dot: operand length mismatch");
596    let mut sum = [0.0f64; FMA_LANES];
597    let mut comp = [0.0f64; FMA_LANES];
598    let mut ca = a.chunks_exact(FMA_LANES);
599    let mut cb = b.chunks_exact(FMA_LANES);
600    for (xa, xb) in ca.by_ref().zip(cb.by_ref()) {
601        for l in 0..FMA_LANES {
602            let x = xa[l];
603            let y = xb[l];
604            // two_prod: p = round(x·y), ep = exact error x·y − p.
605            let p = x * y;
606            let ep = x.mul_add(y, -p);
607            // two_sum: s = round(sum + p), es = exact error.
608            let s = sum[l] + p;
609            let bb = s - sum[l];
610            let es = (sum[l] - (s - bb)) + (p - bb);
611            sum[l] = s;
612            comp[l] += ep + es;
613        }
614    }
615    // Compensated remainder lane (length < FMA_LANES).
616    let mut sr = 0.0f64;
617    let mut cr = 0.0f64;
618    for (&x, &y) in ca.remainder().iter().zip(cb.remainder().iter()) {
619        let p = x * y;
620        let ep = x.mul_add(y, -p);
621        let s = sr + p;
622        let bb = s - sr;
623        let es = (sr - (s - bb)) + (p - bb);
624        sr = s;
625        cr += ep + es;
626    }
627    // Fold each lane's compensation back in, then reduce the (few) lanes.
628    let mut total = sr + cr;
629    for l in 0..FMA_LANES {
630        total += sum[l] + comp[l];
631    }
632    total
633}
634
635/// `out[i] = Σ_j X[i,j]·v[j]` for row-major-contiguous `x_all` (len `n·p`) and
636/// `v` (len `p`). Each output row is an independent [`fma_dot`]; rows fan out
637/// in chunks across the Rayon pool when the work is large.
638fn fast_av_rowmajor_into(x_all: &[f64], v: &[f64], n: usize, p: usize, out: &mut [f64]) {
639    assert_eq!(x_all.len(), n * p, "fast_av_rowmajor_into: x_all length");
640    assert_eq!(v.len(), p, "fast_av_rowmajor_into: v length");
641    assert_eq!(out.len(), n, "fast_av_rowmajor_into: out length");
642    if kernel_should_parallelize(n, p) {
643        use rayon::prelude::*;
644        out.par_chunks_mut(AV_PAR_CHUNK_ROWS)
645            .enumerate()
646            .for_each(|(c, chunk)| {
647                let base = c * AV_PAR_CHUNK_ROWS;
648                for (k, o) in chunk.iter_mut().enumerate() {
649                    let i = base + k;
650                    *o = fma_dot(&x_all[i * p..i * p + p], v);
651                }
652            });
653    } else {
654        for (i, o) in out.iter_mut().enumerate() {
655            *o = fma_dot(&x_all[i * p..i * p + p], v);
656        }
657    }
658}
659
660/// Pairwise (tree) sum of equal-length partial vectors into `out`.
661fn pairwise_sum_into(parts: &[Vec<f64>], out: &mut [f64]) {
662    match parts.len() {
663        0 => out.fill(0.0),
664        1 => out.copy_from_slice(&parts[0]),
665        _ => {
666            let mid = parts.len() / 2;
667            let p = out.len();
668            let mut left = vec![0.0f64; p];
669            let mut right = vec![0.0f64; p];
670            pairwise_sum_into(&parts[..mid], &mut left);
671            pairwise_sum_into(&parts[mid..], &mut right);
672            for ((o, &l), &r) in out.iter_mut().zip(left.iter()).zip(right.iter()) {
673                *o = l + r;
674            }
675        }
676    }
677}
678
679/// `out[j] = Σ_i v[i]·X[i,j]` for row-major-contiguous `x_all` (len `n·p`).
680///
681/// Rows are grouped into [`ATV_BLOCK_ROWS`] blocks; each block FMA-accumulates
682/// its rows into a private partial vector (fused `v[i]·X[i,j]`), and the block
683/// partials are combined pairwise. This blocked/pairwise reduction is both
684/// better-conditioned than a single running sum over all `n` rows and trivially
685/// parallel across blocks.
686fn fast_atv_rowmajor_into(x_all: &[f64], v: &[f64], n: usize, p: usize, out: &mut [f64]) {
687    assert_eq!(x_all.len(), n * p, "fast_atv_rowmajor_into: x_all length");
688    assert_eq!(v.len(), n, "fast_atv_rowmajor_into: v length");
689    assert_eq!(out.len(), p, "fast_atv_rowmajor_into: out length");
690    let nblocks = n.div_ceil(ATV_BLOCK_ROWS);
691
692    let block_partial = |b: usize| -> Vec<f64> {
693        let start = b * ATV_BLOCK_ROWS;
694        let end = (start + ATV_BLOCK_ROWS).min(n);
695        let mut acc = vec![0.0f64; p];
696        for i in start..end {
697            let vi = v[i];
698            let row = &x_all[i * p..i * p + p];
699            for (a, &xij) in acc.iter_mut().zip(row.iter()) {
700                *a = xij.mul_add(vi, *a);
701            }
702        }
703        acc
704    };
705
706    let partials: Vec<Vec<f64>> = if kernel_should_parallelize(n, p) {
707        use rayon::prelude::*;
708        (0..nblocks).into_par_iter().map(block_partial).collect()
709    } else {
710        (0..nblocks).map(block_partial).collect()
711    };
712
713    pairwise_sum_into(&partials, out);
714}
715
716/// Compute A * v using faer's SIMD-optimized GEMV.
717/// For A of shape (n, p) and v of shape (p,), this computes the (n,) result.
718#[inline]
719pub fn fast_av<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
720    a: &ArrayBase<S1, Ix2>,
721    v: &ArrayBase<S2, Ix1>,
722) -> Array1<f64> {
723    if let Some(out) =
724        crate::gpu_hook::gpu_dispatch().and_then(|d| d.try_fast_av(a.view(), v.view()))
725    {
726        return out;
727    }
728    fast_av_impl(a, v)
729}
730
731#[inline]
732fn fast_av_impl<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
733    a: &ArrayBase<S1, Ix2>,
734    v: &ArrayBase<S2, Ix1>,
735) -> Array1<f64> {
736    use faer::linalg::matmul::matmul;
737    use faer::{Accum, Mat};
738
739    let (n, p) = a.dim();
740    assert_eq!(p, v.len(), "A cols must match v length");
741
742    // Row-major-contiguous fast path: tight multi-lane FMA dot per row, both
743    // faster (ILP / Rayon fan-out) and more accurate (fused products, pairwise
744    // lane reduction) than the degenerate single-column faer GEMV.
745    if let (Some(x_all), Some(vs)) = (a.as_slice(), v.as_slice())
746        && n != 0
747        && p != 0
748    {
749        let mut out = Array1::<f64>::zeros(n);
750        fast_av_rowmajor_into(
751            x_all,
752            vs,
753            n,
754            p,
755            out.as_slice_mut().expect("fresh Array1 is contiguous"),
756        );
757        return out;
758    }
759
760    if !should_use_faer_matmul(n, 1, p) {
761        return a.dot(v);
762    }
763
764    let mut result = Mat::<f64>::zeros(n, 1);
765
766    let aview = FaerArrayView::new(a);
767    let vview = FaerColView::new(v);
768    let a_ref = aview.as_ref();
769    let v_ref = vview.as_ref();
770
771    let par = matmul_parallelism(n, 1, p);
772    matmul(result.as_mut(), Accum::Replace, a_ref, v_ref, 1.0, par);
773
774    let mut out = Array1::<f64>::zeros(n);
775    for i in 0..n {
776        out[i] = result[(i, 0)];
777    }
778    out
779}
780
781/// Compute A * v into a pre-allocated output buffer.
782/// `out` must be length n where A is (n, p) and v is length p.
783#[inline]
784pub fn fast_av_into<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
785    a: &ArrayBase<S1, Ix2>,
786    v: &ArrayBase<S2, Ix1>,
787    out: &mut Array1<f64>,
788) {
789    fast_av_into_impl(a, v, out);
790}
791
792#[inline]
793fn fast_av_into_impl<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
794    a: &ArrayBase<S1, Ix2>,
795    v: &ArrayBase<S2, Ix1>,
796    out: &mut Array1<f64>,
797) {
798    use faer::Accum;
799    use faer::linalg::matmul::matmul;
800
801    let (n, p) = a.dim();
802    assert_eq!(v.len(), p, "vector length must match A cols");
803    assert_eq!(out.len(), n, "output length must match A rows");
804
805    if let (Some(x_all), Some(vs)) = (a.as_slice(), v.as_slice())
806        && n != 0
807        && p != 0
808        && let Some(out_s) = out.as_slice_mut()
809    {
810        fast_av_rowmajor_into(x_all, vs, n, p, out_s);
811        return;
812    }
813
814    if !should_use_faer_matmul(n, 1, p) {
815        out.assign(&a.dot(v));
816        return;
817    }
818
819    let mut outview = array1_to_col_matmut(out);
820
821    let aview = FaerArrayView::new(a);
822    let vview = FaerColView::new(v);
823    let a_ref = aview.as_ref();
824    let v_ref = vview.as_ref();
825    let par = matmul_parallelism(n, 1, p);
826    matmul(outview.as_mut(), Accum::Replace, a_ref, v_ref, 1.0, par);
827}
828
829/// Compute A * v into a pre-allocated `ArrayViewMut1` slice. Like
830/// [`fast_av_into`] but accepts a writable slice rather than `&mut Array1`,
831/// so callers can write directly into a sub-range of a larger buffer
832/// without intermediate allocation.
833///
834/// `out` must have length n where A is (n, p) and v is length p.
835#[inline]
836pub fn fast_av_view_into<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
837    a: &ArrayBase<S1, Ix2>,
838    v: &ArrayBase<S2, Ix1>,
839    out: ArrayViewMut1<'_, f64>,
840) {
841    fast_av_view_into_impl(a, v, out);
842}
843
844#[inline]
845fn fast_av_view_into_impl<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
846    a: &ArrayBase<S1, Ix2>,
847    v: &ArrayBase<S2, Ix1>,
848    mut out: ArrayViewMut1<'_, f64>,
849) {
850    use faer::Accum;
851    use faer::linalg::matmul::matmul;
852
853    let (n, p) = a.dim();
854    assert_eq!(v.len(), p, "vector length must match A cols");
855    assert_eq!(out.len(), n, "output length must match A rows");
856
857    if let (Some(x_all), Some(vs)) = (a.as_slice(), v.as_slice())
858        && n != 0
859        && p != 0
860        && let Some(out_s) = out.as_slice_mut()
861    {
862        fast_av_rowmajor_into(x_all, vs, n, p, out_s);
863        return;
864    }
865
866    if !should_use_faer_matmul(n, 1, p) {
867        let prod = a.dot(v);
868        out.assign(&prod);
869        return;
870    }
871
872    let len = out.len();
873    let stride = out.strides()[0];
874    // SAFETY: out.as_mut_ptr() is ndarray's logical first-element pointer, and
875    // len plus the signed element stride describe every initialized element of
876    // this uniquely borrowed view for the returned len×1 MatMut lifetime.
877    let outview = unsafe {
878        MatMut::from_raw_parts_mut(
879            out.as_mut_ptr(),
880            len,
881            1,
882            stride,
883            0, // col stride irrelevant for 1 column
884        )
885    };
886
887    let aview = FaerArrayView::new(a);
888    let vview = FaerColView::new(v);
889    let a_ref = aview.as_ref();
890    let v_ref = vview.as_ref();
891    let par = matmul_parallelism(n, 1, p);
892    matmul(outview, Accum::Replace, a_ref, v_ref, 1.0, par);
893}
894
895/// Compute A^T * v using faer's SIMD-optimized GEMV.
896/// For A of shape (n, p) and v of shape (n,), this computes the (p,) result.
897#[inline]
898pub fn fast_atv<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
899    a: &ArrayBase<S1, Ix2>,
900    v: &ArrayBase<S2, Ix1>,
901) -> Array1<f64> {
902    if let Some(out) =
903        crate::gpu_hook::gpu_dispatch().and_then(|d| d.try_fast_atv(a.view(), v.view()))
904    {
905        return out;
906    }
907    fast_atv_impl(a, v)
908}
909
910#[inline]
911fn fast_atv_impl<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
912    a: &ArrayBase<S1, Ix2>,
913    v: &ArrayBase<S2, Ix1>,
914) -> Array1<f64> {
915    use faer::Accum;
916    use faer::linalg::matmul::matmul;
917
918    let (n, p) = a.dim();
919    assert_eq!(n, v.len(), "A rows must match v length");
920
921    // Row-major-contiguous fast path: blocked + pairwise FMA reduction over the
922    // long n-axis. Lower error-vs-truth than a single running sum and parallel
923    // across row blocks.
924    if let (Some(x_all), Some(vs)) = (a.as_slice(), v.as_slice())
925        && n != 0
926        && p != 0
927    {
928        let mut out = Array1::<f64>::zeros(p);
929        fast_atv_rowmajor_into(
930            x_all,
931            vs,
932            n,
933            p,
934            out.as_slice_mut().expect("fresh Array1 is contiguous"),
935        );
936        return out;
937    }
938
939    // For very small arrays, ndarray might be faster
940    if !should_use_faer_matmul(p, 1, n) {
941        return a.t().dot(v);
942    }
943
944    let mut out = Array1::<f64>::zeros(p);
945    let mut outview = array1_to_col_matmut(&mut out);
946
947    let aview = FaerArrayView::new(a);
948    let vview = FaerColView::new(v);
949    let a_ref = aview.as_ref();
950    let v_ref = vview.as_ref();
951
952    // dst = A^T * v (treating v as n×1 matrix)
953    let par = matmul_parallelism(p, 1, n);
954    matmul(
955        outview.as_mut(),
956        Accum::Replace,
957        a_ref.transpose(),
958        v_ref,
959        1.0,
960        par,
961    );
962
963    out
964}
965
966/// Compute A^T * v into a pre-allocated output buffer.
967/// `out` must be length p where A is (n, p) and v is length n.
968#[inline]
969pub fn fast_atv_into<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
970    a: &ArrayBase<S1, Ix2>,
971    v: &ArrayBase<S2, Ix1>,
972    out: &mut Array1<f64>,
973) {
974    fast_atv_into_impl(a, v, out);
975}
976
977#[inline]
978fn fast_atv_into_impl<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
979    a: &ArrayBase<S1, Ix2>,
980    v: &ArrayBase<S2, Ix1>,
981    out: &mut Array1<f64>,
982) {
983    use faer::Accum;
984    use faer::linalg::matmul::matmul;
985
986    let (n, p) = a.dim();
987    assert_eq!(v.len(), n, "vector length must match A rows");
988    assert_eq!(out.len(), p, "output length must match A cols");
989
990    if let (Some(x_all), Some(vs)) = (a.as_slice(), v.as_slice())
991        && n != 0
992        && p != 0
993        && let Some(out_s) = out.as_slice_mut()
994    {
995        fast_atv_rowmajor_into(x_all, vs, n, p, out_s);
996        return;
997    }
998
999    if !should_use_faer_matmul(p, 1, n) {
1000        out.assign(&a.t().dot(v));
1001        return;
1002    }
1003
1004    let mut outview = array1_to_col_matmut(out);
1005
1006    let aview = FaerArrayView::new(a);
1007    let vview = FaerColView::new(v);
1008    let a_ref = aview.as_ref();
1009    let v_ref = vview.as_ref();
1010    let par = matmul_parallelism(p, 1, n);
1011    matmul(
1012        outview.as_mut(),
1013        Accum::Replace,
1014        a_ref.transpose(),
1015        v_ref,
1016        1.0,
1017        par,
1018    );
1019}
1020
1021/// Compute A^T * diag(W) * A using streaming chunks to avoid O(n*p) allocation.
1022#[inline]
1023pub fn fast_xt_diag_x<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
1024    x: &ArrayBase<S1, Ix2>,
1025    w: &ArrayBase<S2, Ix1>,
1026) -> Array2<f64> {
1027    assert_eq!(
1028        x.nrows(),
1029        w.len(),
1030        "fast_xt_diag_x row/weight length mismatch"
1031    );
1032    if let Some(out) =
1033        crate::gpu_hook::gpu_dispatch().and_then(|d| d.try_fast_xt_diag_x(x.view(), w.view()))
1034    {
1035        return out;
1036    }
1037    let p = x.ncols();
1038    fast_xt_diag_x_with_parallelism(x, w, matmul_parallelism(p, p, x.nrows()))
1039}
1040
1041/// Compute A^T * diag(W) * A with an explicit faer parallelism policy for
1042/// callers that parallelize multiple independent Hessian blocks externally.
1043#[inline]
1044pub fn fast_xt_diag_x_with_parallelism<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
1045    x: &ArrayBase<S1, Ix2>,
1046    w: &ArrayBase<S2, Ix1>,
1047    par: Par,
1048) -> Array2<f64> {
1049    assert_eq!(
1050        x.nrows(),
1051        w.len(),
1052        "fast_xt_diag_x_with_parallelism row/weight length mismatch"
1053    );
1054    fast_xt_diag_x_with_parallelism_impl(x, w, par)
1055}
1056
1057#[inline]
1058fn fast_xt_diag_x_with_parallelism_impl<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
1059    x: &ArrayBase<S1, Ix2>,
1060    w: &ArrayBase<S2, Ix1>,
1061    par: Par,
1062) -> Array2<f64> {
1063    use ndarray::ShapeBuilder;
1064
1065    let p = x.ncols();
1066    // F-order result so the symmetric lower-triangle accumulation writes
1067    // column-contiguously; the kernel mirrors to a full symmetric matrix.
1068    let mut result = Array2::<f64>::zeros((p, p).f());
1069    stream_weighted_crossprod_into(
1070        x,
1071        w,
1072        &mut result,
1073        CrossprodStructure::SymmetricLower,
1074        CrossprodAccum::Replace,
1075        par,
1076    );
1077    result
1078}
1079
1080/// Output packaging for [`stream_weighted_crossprod_into`].
1081#[derive(Clone, Copy, PartialEq, Eq, Debug)]
1082pub enum CrossprodStructure {
1083    /// Compute every entry of the (symmetric) Gram via full GEMM.
1084    Full,
1085    /// Accumulate only the lower triangle via triangular matmul (~50% fewer
1086    /// FLOPs), then mirror once into the upper triangle for a full symmetric
1087    /// result. Mathematically identical output to [`Full`](Self::Full).
1088    SymmetricLower,
1089}
1090
1091/// Accumulation policy for [`stream_weighted_crossprod_into`].
1092#[derive(Clone, Copy, PartialEq, Eq, Debug)]
1093pub enum CrossprodAccum {
1094    /// Overwrite `out` with `Xᵀ·diag(W)·X`, ignoring prior contents.
1095    Replace,
1096    /// Add `Xᵀ·diag(W)·X` into the existing contents of `out`.
1097    Add,
1098}
1099
1100/// Shared dense weighted-Gram kernel: accumulate `Xᵀ·diag(W)·X` into `out`.
1101///
1102/// This is the single tuned implementation of the chunked row-scaling +
1103/// matmul strategy; the matrix-returning (`fast_xt_diag_x*`) entry points and
1104/// stream-in callers share it so that performance tuning, negative-weight
1105/// handling, chunk sizing, and layout fixes land in exactly one place.
1106///
1107/// Computes the product as `Xᵀ·(W·X)` to preserve the sign of `W`: the prior
1108/// `sqrt(max(0, w))`-then-Gram form clipped negative weights to zero, which
1109/// corrupted observed-Hessian assembly when any block carried heavy residuals
1110/// (e.g. under the logb σ link).
1111///
1112/// Peak working-set allocation is `chunk_rows × p × 8` bytes (~8 MB) rather
1113/// than `n × p × 8` bytes for a materialized `W·X`.
1114///
1115/// `out` must be `p × p`. With [`CrossprodStructure::SymmetricLower`] the
1116/// lower triangle is accumulated and then mirrored, so on return `out` holds
1117/// the full symmetric matrix regardless of `structure`.
1118pub fn stream_weighted_crossprod_into<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
1119    x: &ArrayBase<S1, Ix2>,
1120    w: &ArrayBase<S2, Ix1>,
1121    out: &mut Array2<f64>,
1122    structure: CrossprodStructure,
1123    accum: CrossprodAccum,
1124    par: Par,
1125) {
1126    use faer::Accum;
1127    use faer::linalg::matmul::matmul;
1128    use faer::linalg::matmul::triangular::{BlockStructure, matmul as tri_matmul};
1129    use ndarray::s;
1130
1131    let (n, p) = x.dim();
1132    assert_eq!(n, w.len(), "X rows must match W length");
1133    assert_eq!(out.nrows(), p, "output rows must match X cols");
1134    assert_eq!(out.ncols(), p, "output cols must match X cols");
1135    if p == 0 {
1136        return;
1137    }
1138    if n == 0 {
1139        if accum == CrossprodAccum::Replace {
1140            out.fill(0.0);
1141        }
1142        return;
1143    }
1144
1145    if !should_use_faer_matmul(p, p, n) {
1146        // Tiny products: ndarray's own GEMM avoids faer setup overhead.
1147        let w_x = Array2::from_shape_fn((n, p), |(i, j)| w[i] * x[[i, j]]);
1148        let gram = x.t().dot(&w_x);
1149        match accum {
1150            CrossprodAccum::Replace => out.assign(&gram),
1151            CrossprodAccum::Add => *out += &gram,
1152        }
1153        return;
1154    }
1155
1156    // Streaming chunked: peak allocation is chunk_rows × p instead of n × p.
1157    const TARGET_BYTES: usize = 8 * 1024 * 1024;
1158    const MIN_ROWS: usize = 512;
1159    const MAX_ROWS: usize = 131_072;
1160    let chunk_rows = (TARGET_BYTES / (p.max(1) * 8))
1161        .clamp(MIN_ROWS, MAX_ROWS)
1162        .min(n);
1163
1164    // Triangular accumulation requires a zero baseline in the lower triangle
1165    // because each chunk's `Accum::Add` lands there; for a Replace request we
1166    // zero up front and add every chunk, for an Add request the caller's
1167    // contents are preserved and every chunk adds on top.
1168    if accum == CrossprodAccum::Replace {
1169        out.fill(0.0);
1170    }
1171
1172    // Row-major wx_chunk so the per-row scaling loop has stride-1 writes
1173    // alongside stride-1 reads from a row-major X. An F-order wx_chunk would
1174    // force strided writes by `chunk_rows`, breaking vectorization and cache
1175    // locality on the per-PIRLS-iter Hessian assembly. faer's matmul handles
1176    // either layout via FaerArrayView.
1177    let mut wx_chunk = Array2::<f64>::zeros((chunk_rows, p));
1178
1179    let x_is_row_major = x.is_standard_layout();
1180    let w_slice_opt = w.as_slice();
1181
1182    // Scope the faer mutable view so its borrow on `out` ends before the
1183    // symmetric mirror step.
1184    {
1185        let mut out_view = array2_to_matmut(out);
1186        for start in (0..n).step_by(chunk_rows) {
1187            let rows = (n - start).min(chunk_rows);
1188            {
1189                let chunk_slice = wx_chunk
1190                    .as_slice_mut()
1191                    .expect("row-major chunk is contiguous");
1192                if x_is_row_major && let (Some(x_all), Some(w_all)) = (x.as_slice(), w_slice_opt) {
1193                    for local in 0..rows {
1194                        let src = start + local;
1195                        let wi = w_all[src];
1196                        let src_off = src * p;
1197                        let dst_off = local * p;
1198                        let src_row = &x_all[src_off..src_off + p];
1199                        let dst_row = &mut chunk_slice[dst_off..dst_off + p];
1200                        for col in 0..p {
1201                            dst_row[col] = src_row[col] * wi;
1202                        }
1203                    }
1204                } else {
1205                    let x_slice = x.slice(s![start..start + rows, ..]);
1206                    for local in 0..rows {
1207                        let wi = w[start + local];
1208                        let xrow = x_slice.row(local);
1209                        let dst_off = local * p;
1210                        let dst_row = &mut chunk_slice[dst_off..dst_off + p];
1211                        for (col, xij) in xrow.iter().enumerate() {
1212                            dst_row[col] = xij * wi;
1213                        }
1214                    }
1215                }
1216            }
1217            let x_slice = x.slice(s![start..start + rows, ..]);
1218            let wx_slice = wx_chunk.slice(s![0..rows, ..]);
1219            let x_view = FaerArrayView::new(&x_slice);
1220            let wx_view = FaerArrayView::new(&wx_slice);
1221            match structure {
1222                CrossprodStructure::SymmetricLower => {
1223                    // X^T diag(W) X is symmetric; accumulate the lower triangle
1224                    // only, then mirror once after the chunk loop. ~50% fewer
1225                    // FLOPs vs. full GEMM.
1226                    tri_matmul(
1227                        out_view.as_mut(),
1228                        BlockStructure::TriangularLower,
1229                        Accum::Add,
1230                        x_view.as_ref().transpose(),
1231                        BlockStructure::Rectangular,
1232                        wx_view.as_ref(),
1233                        BlockStructure::Rectangular,
1234                        1.0,
1235                        par,
1236                    );
1237                }
1238                CrossprodStructure::Full => {
1239                    matmul(
1240                        out_view.as_mut(),
1241                        Accum::Add,
1242                        x_view.as_ref().transpose(),
1243                        wx_view.as_ref(),
1244                        1.0,
1245                        par,
1246                    );
1247                }
1248            }
1249        }
1250    }
1251
1252    if structure == CrossprodStructure::SymmetricLower {
1253        // Mirror lower triangle to upper for a full symmetric output.
1254        for i in 0..p {
1255            for j in (i + 1)..p {
1256                out[[i, j]] = out[[j, i]];
1257            }
1258        }
1259    }
1260}
1261
1262/// Compute A^T * diag(W) * B using streaming chunks.
1263#[inline]
1264pub fn fast_xt_diag_y<S1: Data<Elem = f64>, S2: Data<Elem = f64>, S3: Data<Elem = f64>>(
1265    x: &ArrayBase<S1, Ix2>,
1266    w: &ArrayBase<S2, Ix1>,
1267    y: &ArrayBase<S3, Ix2>,
1268) -> Array2<f64> {
1269    assert_eq!(x.nrows(), y.nrows(), "fast_xt_diag_y X/Y row mismatch");
1270    assert_eq!(
1271        y.nrows(),
1272        w.len(),
1273        "fast_xt_diag_y row/weight length mismatch"
1274    );
1275    if let Some(out) = crate::gpu_hook::gpu_dispatch()
1276        .and_then(|d| d.try_fast_xt_diag_y(x.view(), w.view(), y.view()))
1277    {
1278        return out;
1279    }
1280    fast_xt_diag_y_impl(x, w, y)
1281}
1282
1283#[inline]
1284fn fast_xt_diag_y_impl<S1: Data<Elem = f64>, S2: Data<Elem = f64>, S3: Data<Elem = f64>>(
1285    x: &ArrayBase<S1, Ix2>,
1286    w: &ArrayBase<S2, Ix1>,
1287    y: &ArrayBase<S3, Ix2>,
1288) -> Array2<f64> {
1289    use faer::Accum;
1290    use faer::linalg::matmul::matmul;
1291    use ndarray::{ShapeBuilder, s};
1292
1293    let (n, q) = y.dim();
1294    let px = x.ncols();
1295    assert_eq!(n, w.len(), "Y rows must match W length");
1296    assert_eq!(n, x.nrows(), "X rows must match Y rows");
1297    if n == 0 || px == 0 || q == 0 {
1298        return Array2::<f64>::zeros((px, q));
1299    }
1300    if !should_use_faer_matmul(px, q, n) {
1301        let w_y = Array2::from_shape_fn((n, q), |(i, j)| w[i] * y[[i, j]]);
1302        return x.t().dot(&w_y);
1303    }
1304
1305    // Streaming: only allocate chunk_rows × q for the weighted Y slice.
1306    const TARGET_BYTES: usize = 8 * 1024 * 1024;
1307    const MIN_ROWS: usize = 512;
1308    const MAX_ROWS: usize = 131_072;
1309    let total_cols = px + q;
1310    let chunk_rows = (TARGET_BYTES / (total_cols.max(1) * 8))
1311        .clamp(MIN_ROWS, MAX_ROWS)
1312        .min(n);
1313
1314    let mut result = Array2::<f64>::zeros((px, q).f());
1315    // Row-major wy_chunk — same rationale as fast_xt_diag_x: stride-1
1316    // writes alongside stride-1 reads from a row-major Y.
1317    let mut wy_chunk = Array2::<f64>::zeros((chunk_rows, q));
1318
1319    let y_is_row_major = y.is_standard_layout();
1320    let w_slice_opt = w.as_slice();
1321
1322    {
1323        let mut out_view = array2_to_matmut(&mut result);
1324
1325        for start in (0..n).step_by(chunk_rows) {
1326            let rows = (n - start).min(chunk_rows);
1327            {
1328                let chunk_slice = wy_chunk
1329                    .as_slice_mut()
1330                    .expect("row-major chunk is contiguous");
1331                if y_is_row_major && let (Some(y_all), Some(w_all)) = (y.as_slice(), w_slice_opt) {
1332                    for local in 0..rows {
1333                        let src = start + local;
1334                        let wi = w_all[src];
1335                        let src_off = src * q;
1336                        let dst_off = local * q;
1337                        let src_row = &y_all[src_off..src_off + q];
1338                        let dst_row = &mut chunk_slice[dst_off..dst_off + q];
1339                        for col in 0..q {
1340                            dst_row[col] = src_row[col] * wi;
1341                        }
1342                    }
1343                } else {
1344                    let y_slice = y.slice(s![start..start + rows, ..]);
1345                    for local in 0..rows {
1346                        let wi = w[start + local];
1347                        let yrow = y_slice.row(local);
1348                        let dst_off = local * q;
1349                        let dst_row = &mut chunk_slice[dst_off..dst_off + q];
1350                        for (col, yij) in yrow.iter().enumerate() {
1351                            dst_row[col] = yij * wi;
1352                        }
1353                    }
1354                }
1355            }
1356            let x_slice = x.slice(s![start..start + rows, ..]);
1357            let wy_slice = wy_chunk.slice(s![0..rows, ..]);
1358            let x_view = FaerArrayView::new(&x_slice);
1359            let wy_view = FaerArrayView::new(&wy_slice);
1360            let par = matmul_parallelism(px, q, rows);
1361            matmul(
1362                out_view.as_mut(),
1363                Accum::Add,
1364                x_view.as_ref().transpose(),
1365                wy_view.as_ref(),
1366                1.0,
1367                par,
1368            );
1369        }
1370    }
1371
1372    result
1373}
1374
1375/// Compute the 2×2 block joint Hessian in a single streaming pass:
1376///   [X_a^T diag(w_aa) X_a,   X_a^T diag(w_ab) X_b]
1377///   [X_b^T diag(w_ab) X_a,   X_b^T diag(w_bb) X_b]
1378///
1379/// This reads X_a and X_b once per chunk instead of twice (saving 50% bandwidth).
1380pub fn fast_joint_hessian_2x2<
1381    S1: Data<Elem = f64>,
1382    S2: Data<Elem = f64>,
1383    S3: Data<Elem = f64>,
1384    S4: Data<Elem = f64>,
1385    S5: Data<Elem = f64>,
1386>(
1387    x_a: &ArrayBase<S1, Ix2>,
1388    x_b: &ArrayBase<S2, Ix2>,
1389    w_aa: &ArrayBase<S3, Ix1>,
1390    w_ab: &ArrayBase<S4, Ix1>,
1391    w_bb: &ArrayBase<S5, Ix1>,
1392) -> Array2<f64> {
1393    if let Some(out) = crate::gpu_hook::gpu_dispatch().and_then(|d| {
1394        d.try_fast_joint_hessian_2x2(
1395            x_a.view(),
1396            x_b.view(),
1397            w_aa.view(),
1398            w_ab.view(),
1399            w_bb.view(),
1400        )
1401    }) {
1402        return out;
1403    }
1404    fast_joint_hessian_2x2_impl(x_a, x_b, w_aa, w_ab, w_bb)
1405}
1406
1407#[inline]
1408fn fast_joint_hessian_2x2_impl<
1409    S1: Data<Elem = f64>,
1410    S2: Data<Elem = f64>,
1411    S3: Data<Elem = f64>,
1412    S4: Data<Elem = f64>,
1413    S5: Data<Elem = f64>,
1414>(
1415    x_a: &ArrayBase<S1, Ix2>,
1416    x_b: &ArrayBase<S2, Ix2>,
1417    w_aa: &ArrayBase<S3, Ix1>,
1418    w_ab: &ArrayBase<S4, Ix1>,
1419    w_bb: &ArrayBase<S5, Ix1>,
1420) -> Array2<f64> {
1421    use faer::Accum;
1422    use faer::linalg::matmul::matmul;
1423    use ndarray::{ShapeBuilder, s};
1424
1425    let n = x_a.nrows();
1426    let pa = x_a.ncols();
1427    let pb = x_b.ncols();
1428    let total = pa + pb;
1429    assert_eq!(n, x_b.nrows());
1430    assert_eq!(n, w_aa.len());
1431    assert_eq!(n, w_ab.len());
1432    assert_eq!(n, w_bb.len());
1433
1434    if n == 0 || total == 0 {
1435        return Array2::<f64>::zeros((total, total));
1436    }
1437
1438    // For small problems, fall back to separate computations
1439    if !should_use_faer_matmul(pa.max(pb), pa.max(pb), n) {
1440        let waa_xa = Array2::from_shape_fn((n, pa), |(i, j)| w_aa[i] * x_a[[i, j]]);
1441        let wab_xb = Array2::from_shape_fn((n, pb), |(i, j)| w_ab[i] * x_b[[i, j]]);
1442        let wbb_xb = Array2::from_shape_fn((n, pb), |(i, j)| w_bb[i] * x_b[[i, j]]);
1443        let mut out = Array2::<f64>::zeros((total, total));
1444        out.slice_mut(s![..pa, ..pa]).assign(&x_a.t().dot(&waa_xa));
1445        out.slice_mut(s![..pa, pa..]).assign(&x_a.t().dot(&wab_xb));
1446        out.slice_mut(s![pa.., pa..]).assign(&x_b.t().dot(&wbb_xb));
1447        // Mirror upper to lower
1448        for i in 0..total {
1449            for j in 0..i {
1450                out[[i, j]] = out[[j, i]];
1451            }
1452        }
1453        return out;
1454    }
1455
1456    const TARGET_BYTES: usize = 8 * 1024 * 1024;
1457    const MIN_ROWS: usize = 512;
1458    const MAX_ROWS: usize = 131_072;
1459    // Need buffers for: waa_xa(chunk×pa) + wab_xb(chunk×pb) + wbb_xb(chunk×pb)
1460    let cols_needed = pa + 2 * pb;
1461    let chunk_rows = (TARGET_BYTES / (cols_needed.max(1) * 8))
1462        .clamp(MIN_ROWS, MAX_ROWS)
1463        .min(n);
1464
1465    let mut out = Array2::<f64>::zeros((total, total).f());
1466    // Row-major weighted buffers so the per-row scale loops have stride-1
1467    // writes (the previous F-order layout strided writes by chunk_rows
1468    // across `pa` / `pb`, gutting vectorization on the per-PIRLS-iter
1469    // joint Hessian assembly). faer's matmul handles either layout.
1470    let mut waa_xa_buf = Array2::<f64>::zeros((chunk_rows, pa));
1471    let mut wab_xb_buf = Array2::<f64>::zeros((chunk_rows, pb));
1472    let mut wbb_xb_buf = Array2::<f64>::zeros((chunk_rows, pb));
1473
1474    let xa_is_row_major = x_a.is_standard_layout();
1475    let xb_is_row_major = x_b.is_standard_layout();
1476    let waa_slice_opt = w_aa.as_slice();
1477    let wab_slice_opt = w_ab.as_slice();
1478    let wbb_slice_opt = w_bb.as_slice();
1479
1480    {
1481        let mut out_mat = array2_to_matmut(&mut out);
1482
1483        for start in (0..n).step_by(chunk_rows) {
1484            let rows = (n - start).min(chunk_rows);
1485            let xa_slice = x_a.slice(s![start..start + rows, ..]);
1486            let xb_slice = x_b.slice(s![start..start + rows, ..]);
1487
1488            // Weight X_a and X_b in a single pass through this chunk.
1489            {
1490                let waa_chunk = waa_xa_buf
1491                    .as_slice_mut()
1492                    .expect("row-major waa chunk is contiguous");
1493                let wab_chunk = wab_xb_buf
1494                    .as_slice_mut()
1495                    .expect("row-major wab chunk is contiguous");
1496                let wbb_chunk = wbb_xb_buf
1497                    .as_slice_mut()
1498                    .expect("row-major wbb chunk is contiguous");
1499
1500                if xa_is_row_major
1501                    && xb_is_row_major
1502                    && let (Some(xa_all), Some(xb_all)) = (x_a.as_slice(), x_b.as_slice())
1503                    && let (Some(waa_all), Some(wab_all), Some(wbb_all)) =
1504                        (waa_slice_opt, wab_slice_opt, wbb_slice_opt)
1505                {
1506                    for local in 0..rows {
1507                        let i = start + local;
1508                        let waa_i = waa_all[i];
1509                        let wab_i = wab_all[i];
1510                        let wbb_i = wbb_all[i];
1511                        let xa_off = i * pa;
1512                        let xa_row = &xa_all[xa_off..xa_off + pa];
1513                        let xb_off = i * pb;
1514                        let xb_row = &xb_all[xb_off..xb_off + pb];
1515                        let waa_off = local * pa;
1516                        let wab_off = local * pb;
1517                        let wbb_off = local * pb;
1518                        let waa_row = &mut waa_chunk[waa_off..waa_off + pa];
1519                        for col in 0..pa {
1520                            waa_row[col] = xa_row[col] * waa_i;
1521                        }
1522                        let wab_row = &mut wab_chunk[wab_off..wab_off + pb];
1523                        let wbb_row = &mut wbb_chunk[wbb_off..wbb_off + pb];
1524                        for col in 0..pb {
1525                            let xij = xb_row[col];
1526                            wab_row[col] = xij * wab_i;
1527                            wbb_row[col] = xij * wbb_i;
1528                        }
1529                    }
1530                } else {
1531                    for local in 0..rows {
1532                        let i = start + local;
1533                        let waa_i = w_aa[i];
1534                        let wab_i = w_ab[i];
1535                        let wbb_i = w_bb[i];
1536                        let waa_off = local * pa;
1537                        let wab_off = local * pb;
1538                        let wbb_off = local * pb;
1539                        let waa_row = &mut waa_chunk[waa_off..waa_off + pa];
1540                        let xa_row = xa_slice.row(local);
1541                        for (col, xij) in xa_row.iter().enumerate() {
1542                            waa_row[col] = xij * waa_i;
1543                        }
1544                        let wab_row = &mut wab_chunk[wab_off..wab_off + pb];
1545                        let wbb_row = &mut wbb_chunk[wbb_off..wbb_off + pb];
1546                        let xb_row = xb_slice.row(local);
1547                        for (col, xij) in xb_row.iter().enumerate() {
1548                            wab_row[col] = xij * wab_i;
1549                            wbb_row[col] = xij * wbb_i;
1550                        }
1551                    }
1552                }
1553            }
1554
1555            let xa_view = FaerArrayView::new(&xa_slice);
1556            let xb_view = FaerArrayView::new(&xb_slice);
1557            let waa_xa_slice = waa_xa_buf.slice(s![0..rows, ..]);
1558            let wab_xb_slice = wab_xb_buf.slice(s![0..rows, ..]);
1559            let wbb_xb_slice = wbb_xb_buf.slice(s![0..rows, ..]);
1560            let waa_xa_view = FaerArrayView::new(&waa_xa_slice);
1561            let wab_xb_view = FaerArrayView::new(&wab_xb_slice);
1562            let wbb_xb_view = FaerArrayView::new(&wbb_xb_slice);
1563
1564            // Block [0..pa, 0..pa]: X_a^T diag(w_aa) X_a
1565            matmul(
1566                out_mat.rb_mut().submatrix_mut(0, 0, pa, pa),
1567                Accum::Add,
1568                xa_view.as_ref().transpose(),
1569                waa_xa_view.as_ref(),
1570                1.0,
1571                matmul_parallelism(pa, pa, rows),
1572            );
1573            // Block [0..pa, pa..total]: X_a^T diag(w_ab) X_b
1574            matmul(
1575                out_mat.rb_mut().submatrix_mut(0, pa, pa, pb),
1576                Accum::Add,
1577                xa_view.as_ref().transpose(),
1578                wab_xb_view.as_ref(),
1579                1.0,
1580                matmul_parallelism(pa, pb, rows),
1581            );
1582            // Block [pa..total, pa..total]: X_b^T diag(w_bb) X_b
1583            matmul(
1584                out_mat.rb_mut().submatrix_mut(pa, pa, pb, pb),
1585                Accum::Add,
1586                xb_view.as_ref().transpose(),
1587                wbb_xb_view.as_ref(),
1588                1.0,
1589                matmul_parallelism(pb, pb, rows),
1590            );
1591        }
1592    } // out_mat dropped
1593    // Mirror upper triangle to lower
1594    for i in 0..total {
1595        for j in 0..i {
1596            out[[i, j]] = out[[j, i]];
1597        }
1598    }
1599    out
1600}
1601
1602fn mat_to_array(mat: MatRef<'_, f64>) -> Array2<f64> {
1603    let nrows = mat.nrows();
1604    let ncols = mat.ncols();
1605    let mut out = Array2::<f64>::zeros((nrows, ncols));
1606    if nrows == 0 || ncols == 0 {
1607        return out;
1608    }
1609    // ndarray is row-major by default. Write row-by-row for best cache behavior
1610    // on the output side.
1611    if let Some(out_slice) = out.as_slice_memory_order_mut() {
1612        // Row-major: out_slice[i * ncols + j] = mat[(i, j)]
1613        for i in 0..nrows {
1614            let row_start = i * ncols;
1615            for j in 0..ncols {
1616                out_slice[row_start + j] = mat[(i, j)];
1617            }
1618        }
1619    } else {
1620        for j in 0..ncols {
1621            for i in 0..nrows {
1622                out[[i, j]] = mat[(i, j)];
1623            }
1624        }
1625    }
1626    out
1627}
1628
1629/// Write faer matmul result A*B directly into a pre-allocated ndarray Array2.
1630/// Avoids the intermediate faer::Mat allocation and mat_to_array copy.
1631#[inline]
1632pub fn fast_ab_into<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
1633    a: &ArrayBase<S1, Ix2>,
1634    b: &ArrayBase<S2, Ix2>,
1635    out: &mut Array2<f64>,
1636) {
1637    fast_ab_into_impl(a, b, out);
1638}
1639
1640#[inline]
1641fn fast_ab_into_impl<S1: Data<Elem = f64>, S2: Data<Elem = f64>>(
1642    a: &ArrayBase<S1, Ix2>,
1643    b: &ArrayBase<S2, Ix2>,
1644    out: &mut Array2<f64>,
1645) {
1646    use faer::Accum;
1647    use faer::linalg::matmul::matmul;
1648
1649    let (n, p) = a.dim();
1650    let (p_b, q) = b.dim();
1651    assert_eq!(p, p_b, "A and B must have compatible inner dimensions");
1652    assert_eq!(out.dim(), (n, q), "output dimensions must match A*B result");
1653
1654    if !should_use_faer_matmul(n, q, p) {
1655        out.assign(&a.dot(b));
1656        return;
1657    }
1658
1659    let aview = FaerArrayView::new(a);
1660    let bview = FaerArrayView::new(b);
1661    let a_ref = aview.as_ref();
1662    let b_ref = bview.as_ref();
1663
1664    let par = matmul_parallelism(n, q, p);
1665    let mut outview = array2_to_matmut(out);
1666    matmul(outview.as_mut(), Accum::Replace, a_ref, b_ref, 1.0, par);
1667}
1668
1669fn diag_to_array(diag: DiagRef<'_, f64>) -> Array1<f64> {
1670    let mat = diag.column_vector().as_mat();
1671    let mut out = Array1::<f64>::zeros(mat.nrows());
1672    for i in 0..mat.nrows() {
1673        out[i] = mat[(i, 0)];
1674    }
1675    out
1676}
1677
1678pub struct FaerArrayView<'a> {
1679    ptr: *const f64,
1680    rows: usize,
1681    cols: usize,
1682    row_stride: isize,
1683    col_stride: isize,
1684    owned: Option<Array2<f64>>,
1685    marker: PhantomData<&'a f64>,
1686}
1687
1688impl<'a> FaerArrayView<'a> {
1689    #[inline]
1690    pub fn new<S: Data<Elem = f64>>(array: &'a ArrayBase<S, Ix2>) -> Self {
1691        let (rows, cols) = array.dim();
1692        let strides = array.strides();
1693        // Guard against layouts that can alias or reverse memory traversal (e.g.
1694        // negative/zero strides). These can violate assumptions in faer kernels.
1695        // For such layouts we materialize a compact owned copy.
1696        if strides[0] <= 0 || strides[1] <= 0 {
1697            let owned = array.to_owned();
1698            let owned_strides = owned.strides();
1699            return Self {
1700                ptr: owned.as_ptr(),
1701                rows,
1702                cols,
1703                row_stride: owned_strides[0],
1704                col_stride: owned_strides[1],
1705                owned: Some(owned),
1706                marker: PhantomData,
1707            };
1708        }
1709
1710        Self {
1711            ptr: array.as_ptr(),
1712            rows,
1713            cols,
1714            row_stride: strides[0],
1715            col_stride: strides[1],
1716            owned: None,
1717            marker: PhantomData,
1718        }
1719    }
1720
1721    #[inline]
1722    pub fn as_ref(&self) -> MatRef<'_, f64> {
1723        let (ptr, rows, cols, row_stride, col_stride) = if let Some(owned) = &self.owned {
1724            let strides = owned.strides();
1725            (
1726                owned.as_ptr(),
1727                owned.nrows(),
1728                owned.ncols(),
1729                strides[0],
1730                strides[1],
1731            )
1732        } else {
1733            (
1734                self.ptr,
1735                self.rows,
1736                self.cols,
1737                self.row_stride,
1738                self.col_stride,
1739            )
1740        };
1741        // SAFETY: ptr/shape/strides come from either a live ndarray view
1742        // (positive strides, validated bounds/alignment) or the owned
1743        // compact copy held inside this wrapper — no mutable aliasing.
1744        unsafe { MatRef::from_raw_parts(ptr, rows, cols, row_stride, col_stride) }
1745    }
1746}
1747
1748pub struct FaerColView<'a> {
1749    ptr: *const f64,
1750    len: usize,
1751    stride: isize,
1752    owned: Option<Array1<f64>>,
1753    marker: PhantomData<&'a f64>,
1754}
1755
1756impl<'a> FaerColView<'a> {
1757    #[inline]
1758    pub fn new<S: Data<Elem = f64>>(array: &'a ArrayBase<S, Ix1>) -> Self {
1759        let len = array.len();
1760        let stride = array.strides()[0];
1761        if stride <= 0 {
1762            let owned = array.to_owned();
1763            return Self {
1764                ptr: owned.as_ptr(),
1765                len,
1766                stride: 1,
1767                owned: Some(owned),
1768                marker: PhantomData,
1769            };
1770        }
1771        Self {
1772            ptr: array.as_ptr(),
1773            len,
1774            stride,
1775            owned: None,
1776            marker: PhantomData,
1777        }
1778    }
1779
1780    #[inline]
1781    pub fn as_ref(&self) -> MatRef<'_, f64> {
1782        let (ptr, len, stride) = if let Some(owned) = &self.owned {
1783            (owned.as_ptr(), owned.len(), 1)
1784        } else {
1785            (self.ptr, self.len, self.stride)
1786        };
1787        // SAFETY: ptr/len/stride come from either a live ndarray column
1788        // (positive stride, validated bounds/alignment) or the owned
1789        // compact copy; ncols=1 so the 0 col-stride is unused.
1790        unsafe { MatRef::from_raw_parts(ptr, len, 1, stride, 0) }
1791    }
1792}
1793
1794pub trait FaerSvd {
1795    fn svd(
1796        &self,
1797        compute_u: bool,
1798        computevt: bool,
1799    ) -> Result<(Option<Array2<f64>>, Array1<f64>, Option<Array2<f64>>), FaerLinalgError>;
1800}
1801
1802impl<S: Data<Elem = f64>> FaerSvd for ArrayBase<S, Ix2> {
1803    fn svd(
1804        &self,
1805        compute_u: bool,
1806        computevt: bool,
1807    ) -> Result<(Option<Array2<f64>>, Array1<f64>, Option<Array2<f64>>), FaerLinalgError> {
1808        let faerview = FaerArrayView::new(self);
1809        let faer_mat = faerview.as_ref();
1810        if !compute_u && !computevt {
1811            let (rows, cols) = faer_mat.shape();
1812            let mut singular = Diag::<f64>::zeros(rows.min(cols));
1813            let par = get_global_parallelism();
1814            let mut mem = MemBuffer::new(svd::svd_scratch::<f64>(
1815                rows,
1816                cols,
1817                ComputeSvdVectors::No,
1818                ComputeSvdVectors::No,
1819                par,
1820                Default::default(),
1821            ));
1822            let stack = MemStack::new(&mut mem);
1823            svd::svd(
1824                faer_mat,
1825                singular.as_mut(),
1826                None,
1827                None,
1828                par,
1829                stack,
1830                Default::default(),
1831            )
1832            .map_err(|_| FaerLinalgError::SvdNoConvergence {
1833                context: "faer SVD singular values only",
1834            })?;
1835            let singularvalues = diag_to_array(singular.as_ref());
1836            return Ok((None, singularvalues, None));
1837        }
1838
1839        let (rows, cols) = faer_mat.shape();
1840        let rank = rows.min(cols);
1841        let compute_u_flag = if compute_u {
1842            ComputeSvdVectors::Thin
1843        } else {
1844            ComputeSvdVectors::No
1845        };
1846        let computev_flag = if computevt {
1847            ComputeSvdVectors::Thin
1848        } else {
1849            ComputeSvdVectors::No
1850        };
1851
1852        let mut singular = Diag::<f64>::zeros(rows.min(cols));
1853        let mut u_storage = compute_u.then(|| Mat::<f64>::zeros(rows, rank));
1854        let mut v_storage = computevt.then(|| Mat::<f64>::zeros(cols, rank));
1855
1856        let par = get_global_parallelism();
1857        let mut mem = MemBuffer::new(svd::svd_scratch::<f64>(
1858            rows,
1859            cols,
1860            compute_u_flag,
1861            computev_flag,
1862            par,
1863            Default::default(),
1864        ));
1865        let stack = MemStack::new(&mut mem);
1866
1867        svd::svd(
1868            faer_mat.as_ref(),
1869            singular.as_mut(),
1870            u_storage.as_mut().map(|mat| mat.as_mut()),
1871            v_storage.as_mut().map(|mat| mat.as_mut()),
1872            par,
1873            stack,
1874            Default::default(),
1875        )
1876        .map_err(|_| FaerLinalgError::SvdNoConvergence {
1877            context: "faer SVD with vectors",
1878        })?;
1879
1880        let singularvalues = diag_to_array(singular.as_ref());
1881        let u_opt = u_storage.map(|mat| mat_to_array(mat.as_ref()));
1882        let vt_opt = v_storage.map(|mat| {
1883            let mat_ref = mat.as_ref();
1884            let mut out = Array2::<f64>::zeros((mat_ref.ncols(), mat_ref.nrows()));
1885            for j in 0..mat_ref.nrows() {
1886                for i in 0..mat_ref.ncols() {
1887                    out[[i, j]] = mat_ref[(j, i)];
1888                }
1889            }
1890            out
1891        });
1892
1893        Ok((u_opt, singularvalues, vt_opt))
1894    }
1895}
1896
1897pub trait FaerEigh {
1898    fn eigh(&self, side: Side) -> Result<(Array1<f64>, Array2<f64>), FaerLinalgError>;
1899}
1900
1901impl<S: Data<Elem = f64>> FaerEigh for ArrayBase<S, Ix2> {
1902    fn eigh(&self, side: Side) -> Result<(Array1<f64>, Array2<f64>), FaerLinalgError> {
1903        fn try_eigh(
1904            matrix: &Array2<f64>,
1905            side: Side,
1906        ) -> Result<(Array1<f64>, Array2<f64>), FaerLinalgError> {
1907            let faerview = FaerArrayView::new(matrix);
1908            let eigen = catch_unwind(AssertUnwindSafe(|| {
1909                faerview.as_ref().self_adjoint_eigen(side)
1910            }))
1911            .map_err(|_| FaerLinalgError::FactorizationFailed {
1912                context: "self-adjoint eigendecomposition panic boundary",
1913            })?
1914            .map_err(FaerLinalgError::SelfAdjointEigen)?;
1915            let values = diag_to_array(eigen.S());
1916            let vectors = mat_to_array(eigen.U());
1917            Ok((values, vectors))
1918        }
1919
1920        let owned = self.to_owned();
1921        if owned.nrows() != owned.ncols() {
1922            return Err(FaerLinalgError::FactorizationFailed {
1923                context: "self-adjoint eigendecomposition non-square input",
1924            });
1925        }
1926        if owned.nrows() == 0 {
1927            return Ok((Array1::zeros(0), Array2::zeros((0, 0))));
1928        }
1929        if owned.iter().any(|value| !value.is_finite()) {
1930            return Err(FaerLinalgError::SelfAdjointEigenNonFiniteInput {
1931                context: "self-adjoint eigendecomposition input validation",
1932            });
1933        }
1934        if let Ok((evals, evecs)) = try_eigh(&owned, side)
1935            && evals.iter().all(|value| value.is_finite())
1936            && evecs.iter().all(|value| value.is_finite())
1937        {
1938            return Ok((evals, evecs));
1939        }
1940
1941        let mut repaired = owned.clone();
1942        crate::matrix::symmetrize_in_place(&mut repaired);
1943
1944        let scale = repaired
1945            .iter()
1946            .fold(0.0_f64, |acc, &value| acc.max(value.abs()))
1947            .max(1.0);
1948        let scaled = repaired.mapv(|value| value / scale);
1949        // Relative diagonal-jitter ladder for the eigendecomposition repair: the
1950        // matrix is pre-scaled to unit max-abs, so these are fractions of its
1951        // scale. We try the unperturbed matrix first, then escalate the ridge by
1952        // two decades per attempt until the factorization yields all-finite
1953        // eigenpairs, accepting the smallest jitter that succeeds.
1954        const JITTER_SCHEDULE: [f64; 6] = [0.0, 1e-12, 1e-10, 1e-8, 1e-6, 1e-4];
1955        let jitter_schedule = JITTER_SCHEDULE;
1956        let mut last_error = FaerLinalgError::FactorizationFailed {
1957            context: "self-adjoint eigendecomposition repair attempts",
1958        };
1959
1960        for &jitter in &jitter_schedule {
1961            let mut candidate = scaled.clone();
1962            if jitter > 0.0 {
1963                let n = candidate.nrows();
1964                for i in 0..n {
1965                    candidate[[i, i]] += jitter;
1966                }
1967            }
1968
1969            match try_eigh(&candidate, side) {
1970                Ok((mut evals, evecs))
1971                    if evals.iter().all(|value| value.is_finite())
1972                        && evecs.iter().all(|value| value.is_finite()) =>
1973                {
1974                    for value in &mut evals {
1975                        *value = (*value - jitter) * scale;
1976                    }
1977                    return Ok((evals, evecs));
1978                }
1979                Ok((_, _)) => {
1980                    last_error = FaerLinalgError::SelfAdjointEigenNonFiniteInput {
1981                        context: "self-adjoint eigendecomposition repaired output validation",
1982                    };
1983                }
1984                Err(err) => {
1985                    last_error = err;
1986                }
1987            }
1988        }
1989
1990        Err(last_error)
1991    }
1992}
1993
1994pub struct FaerCholeskyFactor {
1995    factor: solvers::Llt<f64>,
1996}
1997
1998impl FaerCholeskyFactor {
1999    pub fn solvevec(&self, rhs: &Array1<f64>) -> Array1<f64> {
2000        let mut rhs = rhs.to_owned();
2001        let mut rhsview = array1_to_col_matmut(&mut rhs);
2002        self.factor.solve_in_place(rhsview.as_mut());
2003        rhs
2004    }
2005
2006    pub fn solve_mat_in_place(&self, rhs: &mut Array2<f64>) {
2007        let mut rhsview = array2_to_matmut(rhs);
2008        self.factor.solve_in_place(rhsview.as_mut());
2009    }
2010
2011    pub fn solve_mat_into<S: Data<Elem = f64>>(
2012        &self,
2013        rhs: &ArrayBase<S, Ix2>,
2014        out: &mut Array2<f64>,
2015    ) {
2016        if out.dim() != rhs.dim() {
2017            *out = Array2::<f64>::zeros(rhs.dim());
2018        }
2019        out.assign(rhs);
2020        self.solve_mat_in_place(out);
2021    }
2022
2023    pub fn solve_mat(&self, rhs: &Array2<f64>) -> Array2<f64> {
2024        let mut out = Array2::<f64>::zeros(rhs.dim());
2025        self.solve_mat_into(rhs, &mut out);
2026        out
2027    }
2028
2029    pub fn diag(&self) -> Array1<f64> {
2030        diag_to_array(self.factor.L().diagonal())
2031    }
2032
2033    pub fn lower_triangular(&self) -> Array2<f64> {
2034        mat_to_array(self.factor.L())
2035    }
2036}
2037
2038pub trait FaerCholesky {
2039    fn cholesky(&self, side: Side) -> Result<FaerCholeskyFactor, FaerLinalgError>;
2040}
2041
2042impl<S: Data<Elem = f64>> FaerCholesky for ArrayBase<S, Ix2> {
2043    fn cholesky(&self, side: Side) -> Result<FaerCholeskyFactor, FaerLinalgError> {
2044        let faerview = FaerArrayView::new(self);
2045        let factor = faerview
2046            .as_ref()
2047            .llt(side)
2048            .map_err(FaerLinalgError::Cholesky)?;
2049        Ok(FaerCholeskyFactor { factor })
2050    }
2051}
2052
2053pub trait FaerQr {
2054    fn qr(&self) -> Result<(Array2<f64>, Array2<f64>), FaerLinalgError>;
2055}
2056
2057impl<S: Data<Elem = f64>> FaerQr for ArrayBase<S, Ix2> {
2058    fn qr(&self) -> Result<(Array2<f64>, Array2<f64>), FaerLinalgError> {
2059        let faerview = FaerArrayView::new(self);
2060        let qr = faerview.as_ref().qr();
2061        let q = qr.compute_thin_Q();
2062        let r = qr.thin_R();
2063        Ok((mat_to_array(q.as_ref()), mat_to_array(r)))
2064    }
2065}
2066
2067/// Compute an orthonormal basis for `null(a^T)` using column-pivoted QR on `a`.
2068///
2069/// This is intended for tall/skinny matrices where `a ∈ R^{m×n}` with `m >= n`.
2070/// If `A P^T = Q R`, then the trailing `m-rank(A)` columns of `Q` span
2071/// `null(A^T)`.
2072///
2073/// The trailing columns of `Q` are reconstructed by applying the stored
2074/// Householder reflector sequence to canonical basis vectors. When `A` is
2075/// numerically rank zero (e.g. an entirely unpenalized block penalty in a
2076/// parametric-only GLM), *every* reflector is degenerate — the Householder
2077/// vector of a zero column has zero norm, so faer's coefficients become
2078/// non-finite and the reconstructed basis is filled with `NaN`. Mathematically
2079/// a rank-zero `m×n` matrix has `null(A^T) = R^m`, whose canonical orthonormal
2080/// basis is the identity, so we return `I_m` directly instead of routing through
2081/// the (undefined) reflectors. This keeps every downstream consumer — REML
2082/// null-space log-determinants, identifiability audits — finite and exact for
2083/// the fully-unpenalized case. For `rank >= 1` at least one well-defined
2084/// reflector seeds the block, and the reconstruction stays finite.
2085pub fn rrqr_nullspace_basis<S: Data<Elem = f64>>(
2086    a: &ArrayBase<S, Ix2>,
2087    rank_alpha: f64,
2088) -> Result<(Array2<f64>, usize), FaerLinalgError> {
2089    let faerview = FaerArrayView::new(a);
2090    let qr = faerview.as_ref().col_piv_qr();
2091    let r = qr.thin_R();
2092    let diag_len = r.nrows().min(r.ncols());
2093    let leading_diag = if diag_len > 0 { r[(0, 0)].abs() } else { 0.0 };
2094    let tol = rank_alpha
2095        * f64::EPSILON
2096        * (a.nrows().max(a.ncols()).max(1) as f64)
2097        * leading_diag.max(1.0);
2098    let rank = (0..diag_len).filter(|&i| r[(i, i)].abs() > tol).count();
2099    let z = if rank >= a.nrows() {
2100        Array2::<f64>::zeros((a.nrows(), 0))
2101    } else if rank == 0 {
2102        // Numerically rank-zero input: the whole space is the null space.
2103        // Return the canonical orthonormal basis directly; the Householder
2104        // reflectors of a zero matrix are degenerate and would yield NaN.
2105        Array2::<f64>::eye(a.nrows())
2106    } else {
2107        let nullity = a.nrows() - rank;
2108        let mut selector = Mat::<f64>::zeros(a.nrows(), nullity);
2109        for j in 0..nullity {
2110            selector[(rank + j, j)] = 1.0;
2111        }
2112        let par = get_global_parallelism();
2113        faer::linalg::householder::apply_block_householder_sequence_on_the_left_in_place_with_conj(
2114            qr.Q_basis(),
2115            qr.Q_coeff(),
2116            Conj::No,
2117            selector.as_mut(),
2118            par,
2119            MemStack::new(&mut MemBuffer::new(
2120                faer::linalg::householder::apply_block_householder_sequence_on_the_left_in_place_scratch::<f64>(
2121                    a.nrows(),
2122                    qr.Q_coeff().nrows(),
2123                    nullity,
2124                ),
2125            )),
2126        );
2127        mat_to_array(selector.as_ref())
2128    };
2129    Ok((z, rank))
2130}
2131
2132#[inline]
2133pub const fn default_rrqr_rank_alpha() -> f64 {
2134    RRQR_RANK_ALPHA
2135}
2136
2137/// Result of a column-pivoted QR with rank detection and column permutation.
2138///
2139/// `A · P = Q · R` where the permutation `P` is exposed as the forward index
2140/// array: column `j` of `A · P` corresponds to original column
2141/// `column_permutation[j]` of `A`. With rank `r < min(m, n)`, the trailing
2142/// `min(m, n) - r` entries of `column_permutation` name the columns that the
2143/// pivoted QR demoted past the rank threshold — i.e., the columns identified
2144/// as redundant. Identifiability auditors (`identifiability::audit`)
2145/// use that suffix to attribute `DroppedColumn` entries to specific original
2146/// columns.
2147pub struct RrqrWithPermutation {
2148    pub rank: usize,
2149    pub column_permutation: Vec<usize>,
2150    pub leading_diag_abs: f64,
2151    pub rank_tol: f64,
2152}
2153
2154/// Column-pivoted rank-revealing QR returning the rank, the column permutation,
2155/// and the rank-detection tolerance. Use this when callers need to name which
2156/// columns the pivoted QR demoted past the rank threshold.
2157///
2158/// The rank cutoff matches [`rrqr_nullspace_basis`]: a column-pivoted QR is
2159/// computed on `a`; columns with `|R[i, i]| > tol` count toward the rank,
2160/// where `tol = rank_alpha · eps · max(m, n, 1) · max(|R[0, 0]|, 1)`. Returns
2161/// `Err` when `a` has zero rows.
2162pub fn rrqr_with_permutation<S: Data<Elem = f64>>(
2163    a: &ArrayBase<S, Ix2>,
2164    rank_alpha: f64,
2165) -> Result<RrqrWithPermutation, FaerLinalgError> {
2166    if a.nrows() == 0 {
2167        return Err(FaerLinalgError::FactorizationFailed {
2168            context: "rrqr_with_permutation: input has zero rows",
2169        });
2170    }
2171    let faerview = FaerArrayView::new(a);
2172    let qr = faerview.as_ref().col_piv_qr();
2173    let r = qr.thin_R();
2174    let diag_len = r.nrows().min(r.ncols());
2175    let leading_diag = if diag_len > 0 { r[(0, 0)].abs() } else { 0.0 };
2176    let tol = rank_alpha
2177        * f64::EPSILON
2178        * (a.nrows().max(a.ncols()).max(1) as f64)
2179        * leading_diag.max(1.0);
2180    let rank = (0..diag_len).filter(|&i| r[(i, i)].abs() > tol).count();
2181    let (forward, _inverse) = qr.P().arrays();
2182    let column_permutation: Vec<usize> = forward.iter().copied().map(|idx| idx.unbound()).collect();
2183    Ok(RrqrWithPermutation {
2184        rank,
2185        column_permutation,
2186        leading_diag_abs: leading_diag,
2187        rank_tol: tol,
2188    })
2189}
2190
2191/// Result of a Gram-driven column-pivoted RRQR (see
2192/// [`rrqr_from_gram_with_permutation`]). Carries the same rank / permutation /
2193/// tolerance as [`RrqrWithPermutation`], plus a `verdict_margin` that measures
2194/// how unambiguous the rank cut is — the ratio between the smallest *kept*
2195/// pivot and the rank tolerance. A large margin means squaring the design into
2196/// a Gram could not have flipped any rank decision; a small margin means the
2197/// verdict sits near the cliff and the caller should re-confirm on the full
2198/// (un-squared) design to stay bit-exact.
2199pub struct RrqrFromGram {
2200    pub rank: usize,
2201    pub column_permutation: Vec<usize>,
2202    pub rank_tol: f64,
2203    /// Leading pivot magnitude `|R[0,0]|` of the square-root factor — equal to
2204    /// the largest column norm of the original tall design (col-piv QR pivots the
2205    /// largest-norm column first), so it matches the tall path's
2206    /// `RrqrWithPermutation::leading_diag_abs`.
2207    pub leading_diag_abs: f64,
2208    /// `min_kept_pivot / rank_tol` (∞ when full rank with no kept pivot below
2209    /// tol, i.e. every pivot is comfortably above; `0` when rank is 0).
2210    pub verdict_margin: f64,
2211}
2212
2213/// Column-pivoted rank-revealing QR computed from the design's `p × p` Gram
2214/// `G = AᵀA` (or penalty-augmented `AᵀA + SᵀS`) instead of from the tall
2215/// `m × p` design itself.
2216///
2217/// # Why this is exact (in exact arithmetic)
2218///
2219/// Column-pivoted QR selects, at each step, the not-yet-pivoted column with the
2220/// largest residual norm, where the residual is the part orthogonal to the
2221/// already-chosen columns. Those residual norms — and the resulting pivot
2222/// sequence, the diagonal magnitudes `|R[i,i]|`, and hence the rank cut — are a
2223/// function of the column *inner products* only, i.e. of the Gram `G`. Running
2224/// col-piv QR on the Cholesky factor `R₀` of `G` (`R₀ᵀR₀ = G`, `R₀` is `p × p`)
2225/// reproduces the identical pivot order and identical `|R[i,i]|` as col-piv QR
2226/// on the original `m × p` matrix, because both see the same column geometry.
2227/// This is the standard "pivoted QR depends only on the Gram" identity and lets
2228/// the joint identifiability rank verdict run in `O(p³)` instead of streaming
2229/// all `m ≈ 2·10⁵` rows again.
2230///
2231/// # Tolerance
2232///
2233/// The rank cutoff must match what the tall-matrix [`rrqr_with_permutation`]
2234/// would have used, so the caller passes `m_rows` (the row count of the
2235/// original tall design, including any appended penalty rows). The tolerance is
2236/// `rank_alpha · eps · max(m_rows, p) · max(|R[0,0]|, 1)` — bit-identical to the
2237/// tall path, since `|R[0,0]|` (the leading pivot magnitude = largest column
2238/// norm) is the same in both factorizations.
2239///
2240/// # Finite-precision guard
2241///
2242/// Forming `G = AᵀA` squares the condition number, so a rank decision that sits
2243/// right at the tolerance cliff could in principle flip. The returned
2244/// `verdict_margin` lets the caller detect that case and fall back to the exact
2245/// tall RRQR; in the overwhelmingly common well-separated case (full column
2246/// rank, smallest pivot orders of magnitude above tol) the margin is huge and
2247/// no fallback is needed.
2248pub fn rrqr_from_gram_with_permutation<S: Data<Elem = f64>>(
2249    gram: &ArrayBase<S, Ix2>,
2250    m_rows: usize,
2251    rank_alpha: f64,
2252) -> Result<RrqrFromGram, FaerLinalgError> {
2253    let p = gram.ncols();
2254    if p == 0 {
2255        return Ok(RrqrFromGram {
2256            rank: 0,
2257            column_permutation: Vec::new(),
2258            rank_tol: 0.0,
2259            leading_diag_abs: 0.0,
2260            verdict_margin: 0.0,
2261        });
2262    }
2263    if gram.nrows() != p {
2264        return Err(FaerLinalgError::FactorizationFailed {
2265            context: "rrqr_from_gram_with_permutation: Gram is not square",
2266        });
2267    }
2268    // Symmetric square-root factor F (p×p) with FᵀF = G. The Gram is PSD by
2269    // construction (AᵀA), so its eigendecomposition G = V·diag(λ)·Vᵀ gives the
2270    // factor F = diag(√λ₊)·Vᵀ (rows indexed by eigenpair, columns by original
2271    // design column). Any factor with FᵀF = G reproduces the same column
2272    // geometry, which is all col-piv QR consumes — we use the eigen square root
2273    // rather than a bare Cholesky because Cholesky fails on the numerically
2274    // semidefinite Gram that is exactly the rank-deficient case we must classify.
2275    // Tiny-negative eigenvalues from finite precision are clamped to zero.
2276    let (evals, evecs) = gram.eigh(Side::Lower)?;
2277    let mut f = Array2::<f64>::zeros((p, p));
2278    for k in 0..p {
2279        let scale = evals[k].max(0.0).sqrt();
2280        if scale == 0.0 {
2281            continue;
2282        }
2283        for i in 0..p {
2284            f[[k, i]] = scale * evecs[[i, k]];
2285        }
2286    }
2287    // Single col-piv QR on F. Its pivot order, per-pivot |R[i,i]| magnitudes,
2288    // and leading pivot equal those of col-piv QR on the original tall design
2289    // (FᵀF = G), so this reproduces the exact tall-path geometry.
2290    let faer_f = FaerArrayView::new(&f);
2291    let qr = faer_f.as_ref().col_piv_qr();
2292    let r = qr.thin_R();
2293    let diag_len = r.nrows().min(r.ncols());
2294    let pivots: Vec<f64> = (0..diag_len).map(|i| r[(i, i)].abs()).collect();
2295    let leading_diag = pivots.first().copied().unwrap_or(0.0);
2296    let (forward, _inverse) = qr.P().arrays();
2297    let column_permutation: Vec<usize> = forward.iter().copied().map(|idx| idx.unbound()).collect();
2298    // Re-scale the tolerance from F's `max(p, p)=p` row dimension to the
2299    // original tall design's `max(m_rows, p)`, keeping the rank cut bit-
2300    // identical to what the tall [`rrqr_with_permutation`] would have produced.
2301    let tol = rank_alpha * f64::EPSILON * (m_rows.max(p).max(1) as f64) * leading_diag.max(1.0);
2302    let rank = pivots.iter().filter(|&&v| v > tol).count();
2303    let min_kept = pivots[..rank].iter().copied().fold(f64::INFINITY, f64::min);
2304    let max_dropped = pivots[rank..].iter().copied().fold(0.0f64, f64::max);
2305    // Margin: how far the verdict is from the cliff. Use the smaller of
2306    // (min_kept / tol) and (tol / max_dropped) so a near-tol dropped pivot also
2307    // shrinks the margin. A margin ≫ 1 means no rank decision could flip.
2308    let kept_margin = if rank == 0 {
2309        f64::INFINITY
2310    } else {
2311        min_kept / tol
2312    };
2313    let dropped_margin = if rank == diag_len {
2314        f64::INFINITY
2315    } else {
2316        tol / max_dropped.max(f64::MIN_POSITIVE)
2317    };
2318    // Gram-squaring precision floor. Forming `G = XᵀX` collapses the bottom half
2319    // of the spectrum: a true singular value below `√ε · σ_max` is lost in the
2320    // rounding of `G` (its squared value `σ² < ε·σ_max²` underflows the Gram's
2321    // representable range), and the eigen-square-root then RESURRECTS it as a
2322    // SPURIOUS pivot of magnitude `≈ √(ε·σ_max²) = √ε · σ_max` — orders of
2323    // magnitude ABOVE the true σ and above `tol`. That artefact makes col-piv QR
2324    // on `F` KEEP a column the tall (un-squared) QR would demote: an EXACTLY
2325    // collinear alias (true σ = 0, so `σ² = 0` floored at `≈ ε·σ_max²`) shows up
2326    // as a kept pivot near `√ε · leading`, over-ranking the design and dropping
2327    // nothing (gam#933: a callback-owned column aliased with a higher-priority
2328    // anchor was never demoted, so the reduction never ran and the MAP-uniqueness
2329    // check then fired on the raw collinear joint design). `min_kept / tol` does
2330    // NOT catch this — the spurious pivot sits comfortably above `tol`, so the
2331    // existing margin reports a falsely-confident verdict. The honest test is
2332    // whether the smallest KEPT pivot is itself near the Gram precision floor
2333    // `√ε · leading`: if so, the Gram path cannot distinguish it from a true zero
2334    // and the verdict MUST be re-confirmed on the full-precision tall design.
2335    // Encode that as a third margin term `min_kept / (√ε · leading)` so a kept
2336    // pivot in the floor regime shrinks `verdict_margin` below the caller's
2337    // fallback threshold; for a genuinely full-rank design every kept pivot is
2338    // `≫ √ε · leading` and this term is large, leaving the fast path intact.
2339    let gram_precision_floor = f64::EPSILON.sqrt() * leading_diag.max(1.0);
2340    let kept_floor_margin = if rank == 0 {
2341        f64::INFINITY
2342    } else {
2343        min_kept / gram_precision_floor.max(f64::MIN_POSITIVE)
2344    };
2345    let verdict_margin = kept_margin.min(dropped_margin).min(kept_floor_margin);
2346    Ok(RrqrFromGram {
2347        rank,
2348        column_permutation,
2349        rank_tol: tol,
2350        leading_diag_abs: leading_diag,
2351        verdict_margin,
2352    })
2353}
2354
2355#[cfg(test)]
2356mod tests {
2357    use super::*;
2358    use ndarray::{array, s};
2359
2360    /// Local mirror of the audit's `JOINT_GRAM_RRQR_MIN_VERDICT_MARGIN` fallback
2361    /// threshold, used only by the regression tests below to assert the verdict
2362    /// margin lands on the correct side of the cliff. Kept in sync by value (1e3).
2363    const JOINT_GRAM_RRQR_TRUST_MARGIN_FOR_TEST: f64 = 1.0e3;
2364
2365    #[test]
2366    fn rrqr_nullspace_basis_is_orthonormal_and_annihilates_transpose() {
2367        let a = array![[1.0, 0.0], [1.0, 0.0], [0.0, 2.0], [0.0, 0.0],];
2368        let (z, rank) =
2369            rrqr_nullspace_basis(&a, default_rrqr_rank_alpha()).expect("RRQR should succeed");
2370        assert_eq!(rank, 2);
2371        assert_eq!(z.nrows(), 4);
2372        assert_eq!(z.ncols(), 2);
2373
2374        let gram = z.t().dot(&z);
2375        let ident = Array2::<f64>::eye(z.ncols());
2376        let gram_err = (&gram - &ident)
2377            .iter()
2378            .fold(0.0_f64, |acc, &v| acc.max(v.abs()));
2379        assert!(gram_err < 1e-10, "Z is not orthonormal: {gram_err:e}");
2380
2381        let residual = a.t().dot(&z);
2382        let resid_max = residual.iter().fold(0.0_f64, |acc, &v| acc.max(v.abs()));
2383        assert!(resid_max < 1e-10, "A^T Z residual too large: {resid_max:e}");
2384    }
2385
2386    #[test]
2387    fn rrqr_with_permutation_attributes_redundant_column() {
2388        // 3 columns, column 2 is a duplicate of column 0 → rank 2, column 2
2389        // is the redundant one that the pivoted QR should demote past the
2390        // rank threshold. (Column 1 contributes a different direction.)
2391        let a = array![
2392            [1.0, 0.0, 1.0],
2393            [1.0, 0.0, 1.0],
2394            [0.0, 2.0, 0.0],
2395            [0.0, 0.0, 0.0],
2396        ];
2397        let result =
2398            rrqr_with_permutation(&a, default_rrqr_rank_alpha()).expect("RRQR should succeed");
2399        assert_eq!(result.rank, 2);
2400        assert_eq!(result.column_permutation.len(), 3);
2401        let demoted = result.column_permutation[result.rank..].to_vec();
2402        assert!(
2403            demoted.contains(&2) || demoted.contains(&0),
2404            "demoted suffix should include one of the aliased columns (0 or 2), got {demoted:?}"
2405        );
2406        let mut sorted = result.column_permutation.clone();
2407        sorted.sort();
2408        assert_eq!(
2409            sorted,
2410            vec![0, 1, 2],
2411            "permutation must be a valid bijection on 0..n"
2412        );
2413    }
2414
2415    #[test]
2416    fn rrqr_with_permutation_full_rank_returns_identity_like_order() {
2417        let a = array![[1.0, 0.0], [0.0, 2.0], [0.0, 0.0]];
2418        let result =
2419            rrqr_with_permutation(&a, default_rrqr_rank_alpha()).expect("RRQR should succeed");
2420        assert_eq!(result.rank, 2);
2421        let mut sorted = result.column_permutation.clone();
2422        sorted.sort();
2423        assert_eq!(sorted, vec![0, 1]);
2424    }
2425
2426    #[test]
2427    fn rrqr_with_permutation_rejects_zero_rows() {
2428        let a = Array2::<f64>::zeros((0, 3));
2429        assert!(rrqr_with_permutation(&a, default_rrqr_rank_alpha()).is_err());
2430    }
2431
2432    #[test]
2433    fn rrqr_nullspace_basis_square_zero_matrix_is_finite_identity() {
2434        // Square zero matrix (the parametric-only penalty case): null(A^T) is
2435        // the whole space, so the basis must be a finite orthonormal 3x3 set.
2436        let a = Array2::<f64>::zeros((3, 3));
2437        let (z, rank) =
2438            rrqr_nullspace_basis(&a, default_rrqr_rank_alpha()).expect("RRQR should succeed");
2439        assert_eq!(rank, 0);
2440        assert_eq!(z.dim(), (3, 3));
2441        assert!(
2442            z.iter().all(|v| v.is_finite()),
2443            "square zero matrix produced a non-finite null basis: {z:?}"
2444        );
2445        let gram = z.t().dot(&z);
2446        let ident = Array2::<f64>::eye(3);
2447        let gram_err = (&gram - &ident)
2448            .iter()
2449            .fold(0.0_f64, |acc, &v| acc.max(v.abs()));
2450        assert!(gram_err < 1e-10, "Z is not orthonormal: {gram_err:e}");
2451    }
2452
2453    #[test]
2454    fn rrqr_nullspace_basis_detectszero_rank_matrix() {
2455        let a = Array2::<f64>::zeros((5, 2));
2456        let (z, rank) =
2457            rrqr_nullspace_basis(&a, default_rrqr_rank_alpha()).expect("RRQR should succeed");
2458        assert_eq!(rank, 0);
2459        assert_eq!(z.dim(), (5, 5));
2460        let ident = Array2::<f64>::eye(5);
2461        let max_err = (&z.slice(s![.., ..5]).to_owned() - &ident)
2462            .iter()
2463            .fold(0.0_f64, |acc, &v| acc.max(v.abs()));
2464        assert!(max_err < 1e-10, "zero matrix should yield identity basis");
2465    }
2466
2467    //
2468    // Eigendecomposition NoConvergence on pathological matrices
2469    //
2470    // These tests lock down the hardened contract for FaerEigh::eigh:
2471    // non-finite input must be rejected explicitly, while finite symmetric
2472    // matrices still produce finite spectra.
2473    //
2474
2475    #[test]
2476    fn eigh_on_nan_matrix_rejects_non_finite_input() {
2477        let mat = array![
2478            [1.0, 0.0, 0.0, 0.0],
2479            [0.0, 2.0, 0.0, 0.0],
2480            [0.0, 0.0, 3.0, f64::NAN],
2481            [0.0, 0.0, f64::NAN, 4.0]
2482        ];
2483        let err = mat
2484            .eigh(Side::Lower)
2485            .expect_err("non-finite symmetric input must be rejected");
2486        assert!(matches!(
2487            err,
2488            FaerLinalgError::SelfAdjointEigenNonFiniteInput { .. }
2489        ));
2490    }
2491
2492    #[test]
2493    fn fast_ata_matches_full_gemm_above_threshold() {
2494        // Pick (n, p) large enough to trigger the faer triangular path
2495        // (should_use_faer_matmul threshold is MIN_DIM=32, MIN_FLOP_SCALE=64*64).
2496        let n = 200;
2497        let p = 40;
2498        let a: Array2<f64> = Array2::from_shape_fn((n, p), |(i, j)| {
2499            ((i * 7 + j * 3) as f64).sin() + 0.1 * j as f64
2500        });
2501        let expected = a.t().dot(&a);
2502        let got = fast_ata(&a);
2503        let max_err = (&got - &expected)
2504            .iter()
2505            .fold(0.0_f64, |acc, &v| acc.max(v.abs()));
2506        assert!(max_err < 1e-10, "fast_ata mismatch: {max_err:e}");
2507        // Output must be fully populated and symmetric.
2508        for i in 0..p {
2509            for j in 0..p {
2510                assert!((got[[i, j]] - got[[j, i]]).abs() < 1e-12);
2511            }
2512        }
2513    }
2514
2515    #[test]
2516    fn fast_xt_diag_x_matches_naive_above_threshold() {
2517        let n = 400;
2518        let p = 36;
2519        let x: Array2<f64> =
2520            Array2::from_shape_fn((n, p), |(i, j)| (i as f64 * 0.1).cos() + j as f64 * 0.05);
2521        let w: Array1<f64> = Array1::from_shape_fn(n, |i| (i as f64 * 0.03).sin());
2522        // Naive reference: X^T diag(w) X.
2523        let wx = Array2::from_shape_fn((n, p), |(i, j)| w[i] * x[[i, j]]);
2524        let expected = x.t().dot(&wx);
2525        let got = fast_xt_diag_x(&x, &w);
2526        let max_err = (&got - &expected)
2527            .iter()
2528            .fold(0.0_f64, |acc, &v| acc.max(v.abs()));
2529        assert!(max_err < 1e-9, "fast_xt_diag_x mismatch: {max_err:e}");
2530        for i in 0..p {
2531            for j in 0..p {
2532                assert!((got[[i, j]] - got[[j, i]]).abs() < 1e-12);
2533            }
2534        }
2535    }
2536
2537    #[test]
2538    fn stream_weighted_crossprod_full_and_triangular_parity_with_negative_weights() {
2539        // The stream-in and matrix-returning `fast_xt_diag_x*` packaging modes
2540        // share one kernel. Both packaging modes — and both accumulation
2541        // modes — must reproduce the naive `Xᵀ·diag(w)·X` reference, including signed
2542        // (negative) weights, which the pre-unification sqrt-clip form
2543        // silently corrupted.
2544        //
2545        // Exercise both the streaming faer path (n large enough to clear
2546        // `should_use_faer_matmul`) and the tiny ndarray fallback (small n,p).
2547        for &(n, p) in &[(900usize, 40usize), (8usize, 3usize)] {
2548            let x: Array2<f64> =
2549                Array2::from_shape_fn((n, p), |(i, j)| (i as f64 * 0.07).cos() + j as f64 * 0.013);
2550            // Weights span both signs and zero so negative-weight handling and
2551            // sign preservation are genuinely tested.
2552            let w: Array1<f64> =
2553                Array1::from_shape_fn(n, |i| (i as f64 * 0.11).sin() - 0.25 * (i % 3) as f64);
2554            assert!(
2555                w.iter().any(|&v| v < 0.0),
2556                "weight vector must contain negatives to test sign preservation"
2557            );
2558
2559            // Naive reference: Xᵀ diag(w) X with signed weights.
2560            let wx = Array2::from_shape_fn((n, p), |(i, j)| w[i] * x[[i, j]]);
2561            let expected = x.t().dot(&wx);
2562
2563            let par = matmul_parallelism(p, p, n);
2564
2565            // Full output, Replace.
2566            let mut full = Array2::<f64>::ones((p, p));
2567            stream_weighted_crossprod_into(
2568                &x,
2569                &w,
2570                &mut full,
2571                CrossprodStructure::Full,
2572                CrossprodAccum::Replace,
2573                par,
2574            );
2575
2576            // Triangular+mirror output, Replace. Seed with garbage to prove
2577            // Replace clears prior contents (incl. the upper triangle, which
2578            // the triangular path only reaches via the mirror).
2579            let mut tri = Array2::<f64>::from_elem((p, p), -7.0);
2580            stream_weighted_crossprod_into(
2581                &x,
2582                &w,
2583                &mut tri,
2584                CrossprodStructure::SymmetricLower,
2585                CrossprodAccum::Replace,
2586                par,
2587            );
2588
2589            let full_err = (&full - &expected)
2590                .iter()
2591                .fold(0.0_f64, |a, &v| a.max(v.abs()));
2592            let tri_err = (&tri - &expected)
2593                .iter()
2594                .fold(0.0_f64, |a, &v| a.max(v.abs()));
2595            assert!(
2596                full_err < 1e-9,
2597                "full kernel mismatch (n={n}, p={p}): {full_err:e}"
2598            );
2599            assert!(
2600                tri_err < 1e-9,
2601                "triangular kernel mismatch (n={n}, p={p}): {tri_err:e}"
2602            );
2603
2604            // Full and triangular packaging must agree elementwise, and both
2605            // must be exactly symmetric.
2606            for i in 0..p {
2607                for j in 0..p {
2608                    assert!(
2609                        (full[[i, j]] - tri[[i, j]]).abs() < 1e-12,
2610                        "full vs triangular disagree at ({i},{j})"
2611                    );
2612                    assert!(
2613                        (tri[[i, j]] - tri[[j, i]]).abs() < 1e-12,
2614                        "triangular output not symmetric at ({i},{j})"
2615                    );
2616                }
2617            }
2618
2619            // Accumulation parity: Add into a pre-filled buffer must equal the
2620            // prior contents plus the Gram, for both structures.
2621            let base = Array2::<f64>::from_elem((p, p), 1.5);
2622            let mut add_full = base.clone();
2623            stream_weighted_crossprod_into(
2624                &x,
2625                &w,
2626                &mut add_full,
2627                CrossprodStructure::Full,
2628                CrossprodAccum::Add,
2629                par,
2630            );
2631            let mut add_tri = base.clone();
2632            stream_weighted_crossprod_into(
2633                &x,
2634                &w,
2635                &mut add_tri,
2636                CrossprodStructure::SymmetricLower,
2637                CrossprodAccum::Add,
2638                par,
2639            );
2640            let expected_add = &base + &expected;
2641            let add_full_err = (&add_full - &expected_add)
2642                .iter()
2643                .fold(0.0_f64, |a, &v| a.max(v.abs()));
2644            let add_tri_err = (&add_tri - &expected_add)
2645                .iter()
2646                .fold(0.0_f64, |a, &v| a.max(v.abs()));
2647            assert!(
2648                add_full_err < 1e-9,
2649                "full Add mismatch (n={n}, p={p}): {add_full_err:e}"
2650            );
2651            assert!(
2652                add_tri_err < 1e-9,
2653                "triangular Add mismatch (n={n}, p={p}): {add_tri_err:e}"
2654            );
2655
2656            // The matrix.rs adapter (Full + Replace into a zeroed buffer) must
2657            // match the faer_ndarray return-style adapter bit-for-functionally.
2658            let returned = fast_xt_diag_x(&x, &w);
2659            let returned_err = (&returned - &full)
2660                .iter()
2661                .fold(0.0_f64, |a, &v| a.max(v.abs()));
2662            assert!(
2663                returned_err < 1e-12,
2664                "return adapter vs stream-into adapter disagree (n={n}, p={p}): {returned_err:e}"
2665            );
2666        }
2667    }
2668
2669    #[test]
2670    fn eigh_succeeds_on_same_structure_without_nan() {
2671        // Control: the same matrix with finite values produces finite eigenvalues.
2672        let mat = array![[1.0, 0.5, 0.1], [0.5, 2.0, 0.3], [0.1, 0.3, 1.5]];
2673        let (evals, _) = mat
2674            .eigh(Side::Lower)
2675            .expect("eigh should succeed on a well-conditioned finite matrix");
2676        assert!(
2677            evals.iter().all(|&v| v.is_finite()),
2678            "all eigenvalues should be finite"
2679        );
2680    }
2681
2682    /// gam#933 regression: the Gram-squared RRQR must NOT silently over-rank an
2683    /// EXACTLY collinear design. The invariant is: either the Gram path finds the
2684    /// correct rank (3) by itself — because the precision-floor logic demotes the
2685    /// spurious near-zero pivot before it reaches the kept set — OR, if it
2686    /// over-ranks (reports 4), the `verdict_margin` must collapse below the
2687    /// caller's fallback threshold so the full-precision tall path is used
2688    /// instead. Both outcomes prevent the original gam#933 bug (silent rank=4
2689    /// with high-confidence margin that the caller trusts without verification).
2690    #[test]
2691    fn gram_rrqr_flags_low_margin_on_exact_collinearity_so_caller_falls_back() {
2692        // Joint design [1, x | x, x²] with x ∈ [-1, 1]: columns 1 and 2 are an
2693        // EXACT duplicate (the #933 callback-owned alias), so the true rank is 3.
2694        let n = 48usize;
2695        let x: Vec<f64> = (0..n)
2696            .map(|i| -1.0 + 2.0 * (i as f64) / (n as f64 - 1.0))
2697            .collect();
2698        let mut a = Array2::<f64>::zeros((n, 4));
2699        for i in 0..n {
2700            a[[i, 0]] = 1.0;
2701            a[[i, 1]] = x[i];
2702            a[[i, 2]] = x[i];
2703            a[[i, 3]] = x[i] * x[i];
2704        }
2705        let alpha = default_rrqr_rank_alpha();
2706
2707        // The tall (un-squared) RRQR is the full-precision reference: it must see
2708        // rank 3 and demote one of the duplicate x columns.
2709        let tall = rrqr_with_permutation(&a, alpha).expect("tall RRQR should succeed");
2710        assert_eq!(tall.rank, 3, "tall RRQR must demote the exact alias");
2711
2712        // The Gram-squared RRQR must satisfy the gam#933 invariant:
2713        //   rank == 3 (correct result)  OR  verdict_margin < threshold (force fallback)
2714        //
2715        // The precision-floor margin term was designed to catch the case where
2716        // squaring the spectrum resurrects a spurious kept pivot near √ε·σ_max.
2717        // When the eigen-square-root approach correctly demotes that pivot
2718        // (yielding rank=3 without spurious kept columns), the margin is
2719        // legitimately high — trusting the Gram result is then safe and correct.
2720        // When it over-ranks (rank=4), the floor margin must be low so the
2721        // caller falls back to the tall RRQR and gets the right answer.
2722        let unit = Array1::<f64>::ones(n);
2723        let gram = fast_xt_diag_x_with_parallelism(&a, &unit, faer::get_global_parallelism());
2724        let gram_rrqr =
2725            rrqr_from_gram_with_permutation(&gram, n, alpha).expect("Gram RRQR should succeed");
2726        let ok =
2727            gram_rrqr.rank == 3 || gram_rrqr.verdict_margin < JOINT_GRAM_RRQR_TRUST_MARGIN_FOR_TEST;
2728        assert!(
2729            ok,
2730            "gam#933: Gram RRQR must either find correct rank=3 OR signal low margin \
2731             (< {:.0e}) to force the tall fallback; got rank={} margin={:.3e}",
2732            JOINT_GRAM_RRQR_TRUST_MARGIN_FOR_TEST, gram_rrqr.rank, gram_rrqr.verdict_margin,
2733        );
2734    }
2735
2736    /// Companion to the regression above: a genuinely full-rank, moderately
2737    /// conditioned design must keep a LARGE Gram verdict margin so the fast Gram
2738    /// path is retained (the precision-floor term must not trip on real, small-
2739    /// but-nonzero singular values).
2740    #[test]
2741    fn gram_rrqr_keeps_high_margin_on_full_rank_design() {
2742        let n = 200usize;
2743        let p = 5usize;
2744        let mut a = Array2::<f64>::zeros((n, p));
2745        // Deterministic, well-separated columns (distinct low-order polynomials).
2746        for i in 0..n {
2747            let t = (i as f64) / (n as f64 - 1.0);
2748            a[[i, 0]] = 1.0;
2749            a[[i, 1]] = t;
2750            a[[i, 2]] = t * t;
2751            a[[i, 3]] = t * t * t;
2752            a[[i, 4]] = (t * 6.0).sin();
2753        }
2754        let alpha = default_rrqr_rank_alpha();
2755        let unit = Array1::<f64>::ones(n);
2756        let gram = fast_xt_diag_x_with_parallelism(&a, &unit, faer::get_global_parallelism());
2757        let gram_rrqr =
2758            rrqr_from_gram_with_permutation(&gram, n, alpha).expect("Gram RRQR should succeed");
2759        assert_eq!(gram_rrqr.rank, p, "full-rank design must keep all columns");
2760        assert!(
2761            gram_rrqr.verdict_margin >= JOINT_GRAM_RRQR_TRUST_MARGIN_FOR_TEST,
2762            "full-rank design must keep a high margin (fast Gram path); got {:.3e}",
2763            gram_rrqr.verdict_margin,
2764        );
2765    }
2766
2767    // ── fast_ab / fast_atb / fast_abt / fast_av / fast_atv / fast_xt_diag_y ──
2768
2769    fn max_abs_diff(a: &Array2<f64>, b: &Array2<f64>) -> f64 {
2770        assert_eq!(a.dim(), b.dim(), "shape mismatch in max_abs_diff");
2771        a.iter()
2772            .zip(b.iter())
2773            .fold(0.0_f64, |acc, (&x, &y)| acc.max((x - y).abs()))
2774    }
2775
2776    fn max_abs_diff_1d(a: &Array1<f64>, b: &Array1<f64>) -> f64 {
2777        assert_eq!(a.len(), b.len(), "len mismatch in max_abs_diff_1d");
2778        a.iter()
2779            .zip(b.iter())
2780            .fold(0.0_f64, |acc, (&x, &y)| acc.max((x - y).abs()))
2781    }
2782
2783    /// `fast_ab(A, B)` matches `A.dot(&B)` for small (ndarray-path) matrices.
2784    #[test]
2785    fn fast_ab_small_matches_ndarray_dot() {
2786        let a = array![[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]];
2787        let b = array![[7.0, 8.0], [9.0, 10.0], [11.0, 12.0]];
2788        let got = fast_ab(&a, &b);
2789        let want = a.dot(&b);
2790        assert!(max_abs_diff(&got, &want) < 1e-12, "fast_ab small mismatch");
2791        assert_eq!(got.dim(), (2, 2));
2792    }
2793
2794    /// `fast_ab` on larger matrices (faer path) agrees with ndarray dot.
2795    #[test]
2796    fn fast_ab_large_matches_ndarray_dot() {
2797        let n = 50usize;
2798        let p = 40usize;
2799        let q = 35usize;
2800        let mut a = Array2::<f64>::zeros((n, p));
2801        let mut b = Array2::<f64>::zeros((p, q));
2802        let mut state = 0xDEAD_BEEF_1234_5678u64;
2803        let next = |s: &mut u64| -> f64 {
2804            *s ^= *s << 13;
2805            *s ^= *s >> 7;
2806            *s ^= *s << 17;
2807            ((*s >> 11) as f64 / ((1u64 << 53) as f64)) - 0.5
2808        };
2809        for v in a.iter_mut() {
2810            *v = next(&mut state);
2811        }
2812        for v in b.iter_mut() {
2813            *v = next(&mut state);
2814        }
2815        let got = fast_ab(&a, &b);
2816        let want = a.dot(&b);
2817        assert!(max_abs_diff(&got, &want) < 1e-9, "fast_ab large mismatch");
2818    }
2819
2820    /// `fast_atb(A, B)` = A^T * B for small matrices (ndarray path).
2821    #[test]
2822    fn fast_atb_small_matches_ndarray_dot() {
2823        let a = array![[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]];
2824        let b = array![[7.0, 8.0, 9.0], [10.0, 11.0, 12.0], [13.0, 14.0, 15.0]];
2825        let got = fast_atb(&a, &b);
2826        let want = a.t().dot(&b);
2827        assert!(max_abs_diff(&got, &want) < 1e-12, "fast_atb small mismatch");
2828        assert_eq!(got.dim(), (2, 3));
2829    }
2830
2831    /// `fast_atb` on larger matrices (faer path) agrees with ndarray.
2832    #[test]
2833    fn fast_atb_large_matches_ndarray_dot() {
2834        let n = 50usize;
2835        let p = 40usize;
2836        let q = 35usize;
2837        let mut a = Array2::<f64>::zeros((n, p));
2838        let mut b = Array2::<f64>::zeros((n, q));
2839        let mut state = 0xCAFE_BABE_9876_5432u64;
2840        let next = |s: &mut u64| -> f64 {
2841            *s ^= *s << 13;
2842            *s ^= *s >> 7;
2843            *s ^= *s << 17;
2844            ((*s >> 11) as f64 / ((1u64 << 53) as f64)) - 0.5
2845        };
2846        for v in a.iter_mut() {
2847            *v = next(&mut state);
2848        }
2849        for v in b.iter_mut() {
2850            *v = next(&mut state);
2851        }
2852        let got = fast_atb(&a, &b);
2853        let want = a.t().dot(&b);
2854        assert!(max_abs_diff(&got, &want) < 1e-9, "fast_atb large mismatch");
2855    }
2856
2857    /// `fast_abt(A, B)` = A * B^T for small matrices (ndarray path).
2858    #[test]
2859    fn fast_abt_small_matches_ndarray_dot() {
2860        let a = array![[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]];
2861        let b = array![[7.0, 8.0, 9.0], [10.0, 11.0, 12.0]];
2862        let got = fast_abt(&a, &b);
2863        let want = a.dot(&b.t());
2864        assert!(max_abs_diff(&got, &want) < 1e-12, "fast_abt small mismatch");
2865        assert_eq!(got.dim(), (2, 2));
2866    }
2867
2868    /// `fast_av(A, v)` = A * v for small (ndarray path) and larger (faer path).
2869    #[test]
2870    fn fast_av_small_matches_ndarray_dot() {
2871        let a = array![[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]];
2872        let v = array![1.0, -1.0, 2.0];
2873        let got = fast_av(&a, &v);
2874        let want = a.dot(&v);
2875        assert!(
2876            max_abs_diff_1d(&got, &want) < 1e-12,
2877            "fast_av small mismatch"
2878        );
2879        // 1*1 + 2*(-1) + 3*2 = 1-2+6 = 5
2880        assert!((got[0] - 5.0).abs() < 1e-12, "fast_av[0] should be 5");
2881        // 4*1 + 5*(-1) + 6*2 = 4-5+12 = 11
2882        assert!((got[1] - 11.0).abs() < 1e-12, "fast_av[1] should be 11");
2883    }
2884
2885    /// `fast_av` on larger matrices (faer path) agrees with ndarray.
2886    #[test]
2887    fn fast_av_large_matches_ndarray_dot() {
2888        let n = 50usize;
2889        let p = 40usize;
2890        let mut a = Array2::<f64>::zeros((n, p));
2891        let mut v = Array1::<f64>::zeros(p);
2892        let mut state = 0xFEED_FACE_ABCD_EF01u64;
2893        let next = |s: &mut u64| -> f64 {
2894            *s ^= *s << 13;
2895            *s ^= *s >> 7;
2896            *s ^= *s << 17;
2897            ((*s >> 11) as f64 / ((1u64 << 53) as f64)) - 0.5
2898        };
2899        for v in a.iter_mut() {
2900            *v = next(&mut state);
2901        }
2902        for x in v.iter_mut() {
2903            *x = next(&mut state);
2904        }
2905        let got = fast_av(&a, &v);
2906        let want = a.dot(&v);
2907        assert!(
2908            max_abs_diff_1d(&got, &want) < 1e-9,
2909            "fast_av large mismatch"
2910        );
2911    }
2912
2913    /// `fast_atv(A, v)` = A^T * v for small matrices (ndarray path).
2914    #[test]
2915    fn fast_atv_small_matches_ndarray_dot() {
2916        let a = array![[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]];
2917        let v = array![1.0, 0.0, -1.0];
2918        let got = fast_atv(&a, &v);
2919        let want = a.t().dot(&v);
2920        // A^T * v = [1*1+3*0+5*(-1), 2*1+4*0+6*(-1)] = [-4, -4]
2921        assert!(
2922            max_abs_diff_1d(&got, &want) < 1e-12,
2923            "fast_atv small mismatch"
2924        );
2925        assert!((got[0] - (-4.0)).abs() < 1e-12, "fast_atv[0]");
2926        assert!((got[1] - (-4.0)).abs() < 1e-12, "fast_atv[1]");
2927    }
2928
2929    /// `fast_atv` on larger matrices (faer path) agrees with ndarray.
2930    #[test]
2931    fn fast_atv_large_matches_ndarray_dot() {
2932        let n = 50usize;
2933        let p = 40usize;
2934        let mut a = Array2::<f64>::zeros((n, p));
2935        let mut v = Array1::<f64>::zeros(n);
2936        let mut state = 0x1234_ABCD_5678_EF90u64;
2937        let next = |s: &mut u64| -> f64 {
2938            *s ^= *s << 13;
2939            *s ^= *s >> 7;
2940            *s ^= *s << 17;
2941            ((*s >> 11) as f64 / ((1u64 << 53) as f64)) - 0.5
2942        };
2943        for x in a.iter_mut() {
2944            *x = next(&mut state);
2945        }
2946        for x in v.iter_mut() {
2947            *x = next(&mut state);
2948        }
2949        let got = fast_atv(&a, &v);
2950        let want = a.t().dot(&v);
2951        assert!(
2952            max_abs_diff_1d(&got, &want) < 1e-9,
2953            "fast_atv large mismatch"
2954        );
2955    }
2956
2957    /// `fast_xt_diag_y(X, d, Y)` = X^T * diag(d) * Y, verified against
2958    /// a manual triple-product for small inputs.
2959    #[test]
2960    fn fast_xt_diag_y_small_matches_manual() {
2961        let x = array![[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]];
2962        let d = array![2.0, 0.5, 1.0];
2963        let y = array![[7.0, 8.0, 9.0], [10.0, 11.0, 12.0], [13.0, 14.0, 15.0]];
2964        let got = fast_xt_diag_y(&x, &d, &y);
2965        // Manual: X^T * diag(d) * Y
2966        let diag_y = {
2967            let mut dy = Array2::<f64>::zeros(y.dim());
2968            for i in 0..3 {
2969                for j in 0..3 {
2970                    dy[[i, j]] = d[i] * y[[i, j]];
2971                }
2972            }
2973            dy
2974        };
2975        let want = x.t().dot(&diag_y);
2976        assert!(
2977            max_abs_diff(&got, &want) < 1e-12,
2978            "fast_xt_diag_y small mismatch"
2979        );
2980        assert_eq!(got.dim(), (2, 3));
2981    }
2982
2983    // ── Compensated-reduction accuracy oracle ────────────────────────────
2984    //
2985    // Truth is an error-free (exact-expansion / double-double) reference. We
2986    // assert the production GEMV kernels are pointwise no less accurate than —
2987    // and in aggregate strictly better than — a naive sequential sum.
2988
2989    #[inline]
2990    fn two_prod(a: f64, b: f64) -> (f64, f64) {
2991        let p = a * b;
2992        let e = a.mul_add(b, -p);
2993        (p, e)
2994    }
2995
2996    #[inline]
2997    fn two_sum(a: f64, b: f64) -> (f64, f64) {
2998        let s = a + b;
2999        let bb = s - a;
3000        let e = (a - (s - bb)) + (b - bb);
3001        (s, e)
3002    }
3003
3004    /// Shewchuk grow-expansion: add `q` to the non-overlapping expansion `e`.
3005    fn grow_expansion(e: &mut Vec<f64>, mut q: f64) {
3006        for h in e.iter_mut() {
3007            let (s, err) = two_sum(*h, q);
3008            *h = err;
3009            q = s;
3010        }
3011        if q != 0.0 {
3012            e.push(q);
3013        }
3014    }
3015
3016    /// Exact dot product (correctly rounded to `f64`) via an error-free
3017    /// expansion of every `two_prod` component. O(n²) — for short reference
3018    /// vectors only — but a true gold standard, strictly more precise than any
3019    /// double-precision accumulator under test.
3020    fn exact_dot(a: &[f64], b: &[f64]) -> f64 {
3021        let mut e: Vec<f64> = Vec::new();
3022        for (&x, &y) in a.iter().zip(b.iter()) {
3023            let (p, ep) = two_prod(x, y);
3024            grow_expansion(&mut e, p);
3025            grow_expansion(&mut e, ep);
3026        }
3027        // Components are non-overlapping and ascending in magnitude; summing
3028        // smallest-first yields the correctly rounded total.
3029        e.iter().fold(0.0f64, |acc, &c| acc + c)
3030    }
3031
3032    /// High-precision reference dot via compensated (double-double) summation.
3033    /// Cheap (O(n)) — used where naive's error is enormous so ~2u precision is
3034    /// already far more accurate than the baseline under test.
3035    fn dd_dot(a: &[f64], b: &[f64]) -> f64 {
3036        let (mut s, mut c) = (0.0f64, 0.0f64);
3037        for (&x, &y) in a.iter().zip(b.iter()) {
3038            let (p, ep) = two_prod(x, y);
3039            let (s2, es) = two_sum(s, p);
3040            s = s2;
3041            c += ep + es;
3042        }
3043        s + c
3044    }
3045
3046    fn naive_dot(a: &[f64], b: &[f64]) -> f64 {
3047        let mut acc = 0.0f64;
3048        for (&x, &y) in a.iter().zip(b.iter()) {
3049            acc += x * y;
3050        }
3051        acc
3052    }
3053
3054    /// Catastrophic-cancellation generator: large opposing terms plus small
3055    /// ones, so the naive running sum loses many bits to cancellation.
3056    fn ill_conditioned_pair(len: usize, seed: u64) -> (Vec<f64>, Vec<f64>) {
3057        let mut s = seed | 1;
3058        let mut next = || {
3059            s ^= s << 13;
3060            s ^= s >> 7;
3061            s ^= s << 17;
3062            (s >> 11) as f64 / ((1u64 << 53) as f64) - 0.5
3063        };
3064        let mut a = Vec::with_capacity(len);
3065        let mut b = Vec::with_capacity(len);
3066        for i in 0..len {
3067            // Span ~16 orders of magnitude with alternating signs.
3068            let scale = 10f64.powi((i % 17) as i32 - 8);
3069            let sign = if i % 2 == 0 { 1.0 } else { -1.0 };
3070            a.push(sign * next() * scale);
3071            b.push(next() * scale);
3072        }
3073        (a, b)
3074    }
3075
3076    /// `fma_dot` (compensated Dot2) error-vs-truth never exceeds the naive
3077    /// sum's and is strictly lower on the ill-conditioned ensemble in aggregate.
3078    #[test]
3079    fn fma_dot_beats_naive_accuracy() {
3080        let mut fma_total = 0.0f64;
3081        let mut naive_total = 0.0f64;
3082        let mut strict_wins = 0;
3083        for seed in 0..64u64 {
3084            let len = 200 + (seed as usize % 57);
3085            let (a, b) = ill_conditioned_pair(len, 0x9E37_79B9 ^ seed.wrapping_mul(2654435761));
3086            let truth = exact_dot(&a, &b);
3087            let fe = (super::fma_dot(&a, &b) - truth).abs();
3088            let ne = (naive_dot(&a, &b) - truth).abs();
3089            // Compensated (Dot2) summation is pointwise no less accurate than
3090            // the naive recurrence. The floor term tolerates a few-ulp tie when
3091            // both already sit at the round-to-nearest limit (well-conditioned).
3092            let floor = 8.0 * f64::EPSILON * truth.abs();
3093            assert!(
3094                fe <= ne * (1.0 + 1e-6) + floor,
3095                "fma_dot worse than naive: seed={seed} fma_err={fe:.3e} naive_err={ne:.3e}",
3096            );
3097            if fe < ne {
3098                strict_wins += 1;
3099            }
3100            fma_total += fe;
3101            naive_total += ne;
3102        }
3103        assert!(
3104            fma_total < naive_total,
3105            "fma_dot aggregate error {fma_total:.3e} not below naive {naive_total:.3e}",
3106        );
3107        assert!(
3108            strict_wins >= 40,
3109            "expected fma_dot to strictly win the majority; only {strict_wins}/64",
3110        );
3111    }
3112
3113    /// `fast_atv` blocked+pairwise reduction is strictly more accurate than a
3114    /// naive running column-sum on a long, ill-conditioned `n`-axis.
3115    #[test]
3116    fn fast_atv_blocked_beats_naive_accuracy() {
3117        let n = 200_003usize;
3118        let p = 3usize;
3119        let mut s = 0xD1B5_4A32u64;
3120        let mut next = || {
3121            s ^= s << 13;
3122            s ^= s >> 7;
3123            s ^= s << 17;
3124            (s >> 11) as f64 / ((1u64 << 53) as f64) - 0.5
3125        };
3126        let mut x = Array2::<f64>::zeros((n, p));
3127        let mut v = Array1::<f64>::zeros(n);
3128        for i in 0..n {
3129            let scale = 10f64.powi((i % 17) as i32 - 8);
3130            v[i] = if i % 2 == 0 { scale } else { -scale } * next();
3131            for j in 0..p {
3132                x[[i, j]] = next() * scale;
3133            }
3134        }
3135        let got = fast_atv(&x, &v);
3136        // Per-column truth and naive baseline.
3137        for j in 0..p {
3138            let col: Vec<f64> = (0..n).map(|i| x[[i, j]]).collect();
3139            let vv: Vec<f64> = v.to_vec();
3140            let truth = dd_dot(&col, &vv);
3141            let naive = naive_dot(&col, &vv);
3142            let ge = (got[j] - truth).abs();
3143            let ne = (naive - truth).abs();
3144            assert!(
3145                ge <= ne + f64::MIN_POSITIVE,
3146                "col {j}: blocked err {ge:.3e} exceeds naive {ne:.3e}",
3147            );
3148        }
3149    }
3150
3151    /// Non-contiguous (transposed-view) operands take the faer fallback and
3152    /// still match ndarray, proving the kernel gate is layout-safe.
3153    #[test]
3154    fn fast_av_strided_input_matches_ndarray() {
3155        let mut base = Array2::<f64>::zeros((40, 60));
3156        let mut s = 0x0BAD_F00Du64;
3157        let mut next = || {
3158            s ^= s << 13;
3159            s ^= s >> 7;
3160            s ^= s << 17;
3161            (s >> 11) as f64 / ((1u64 << 53) as f64) - 0.5
3162        };
3163        for x in base.iter_mut() {
3164            *x = next();
3165        }
3166        // A transposed view of `base` is (60, 40), non-row-major-contiguous.
3167        let a = base.t();
3168        let mut v = Array1::<f64>::zeros(40);
3169        for x in v.iter_mut() {
3170            *x = next();
3171        }
3172        let got = fast_av(&a, &v);
3173        let want = a.dot(&v);
3174        assert!(
3175            max_abs_diff_1d(&got, &want) < 1e-11,
3176            "strided fast_av mismatch (fallback path)",
3177        );
3178    }
3179
3180    // ── FaerSequentialScope (#2074) ───────────────────────────────────────────
3181    //
3182    // The guard pins faer's process-global parallelism to `Par::Seq` for its
3183    // lifetime and restores the prior policy when the outermost guard drops.
3184    // This is the primitive that closes the K=1 `sae_manifold_fit` deadlock: a
3185    // faer high-level solver reached from inside a topology-race Rayon worker
3186    // would otherwise fan a nested `spindle` barrier pool and park at 0% CPU.
3187    //
3188    // These tests mutate the process-global faer setting, so they save/restore a
3189    // known baseline and run serially under one `#[test]` to avoid racing the
3190    // other global-parallelism tests in this binary.
3191    #[test]
3192    fn faer_sequential_scope_sets_seq_inside_and_restores_after() {
3193        let baseline = faer::get_global_parallelism();
3194        // Establish a definitely-parallel baseline so the "restores" assertion is
3195        // meaningful (not vacuously Seq already).
3196        faer::set_global_parallelism(Par::rayon(4));
3197        assert_eq!(
3198            faer::get_global_parallelism(),
3199            Par::rayon(4),
3200            "baseline must be the parallel policy we just set",
3201        );
3202
3203        {
3204            let faer_seq_guard = FaerSequentialScope::enter();
3205            assert_eq!(
3206                faer::get_global_parallelism(),
3207                Par::Seq,
3208                "faer must be pinned to Par::Seq inside the scope",
3209            );
3210
3211            // Nested guard: still Seq, and the inner drop must NOT restore early.
3212            {
3213                let faer_seq_inner_guard = FaerSequentialScope::enter();
3214                assert_eq!(
3215                    faer::get_global_parallelism(),
3216                    Par::Seq,
3217                    "nested scope stays Par::Seq",
3218                );
3219                drop(faer_seq_inner_guard);
3220            }
3221            assert_eq!(
3222                faer::get_global_parallelism(),
3223                Par::Seq,
3224                "inner drop must not restore while outer scope is still live",
3225            );
3226            drop(faer_seq_guard);
3227        }
3228
3229        assert_eq!(
3230            faer::get_global_parallelism(),
3231            Par::rayon(4),
3232            "outermost drop must restore the pre-scope parallelism policy",
3233        );
3234
3235        // The convenience wrapper behaves identically and returns the body value.
3236        let observed = with_faer_sequential(|| faer::get_global_parallelism());
3237        assert_eq!(
3238            observed,
3239            Par::Seq,
3240            "with_faer_sequential runs body under Seq"
3241        );
3242        assert_eq!(
3243            faer::get_global_parallelism(),
3244            Par::rayon(4),
3245            "with_faer_sequential restores after the body returns",
3246        );
3247
3248        // Restore the binary-wide baseline.
3249        faer::set_global_parallelism(baseline);
3250    }
3251}