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