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