rlx-cpu 0.2.17

// RLX — versatile ML compiler + runtime.
// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
// SPDX-License-Identifier: MIT OR Apache-2.0

//! Direct BLAS FFI — zero abstraction overhead.
//!
//! Calls cblas_sgemm directly without going through ndarray, faer, or any
//! wrapper. This is the same approach that gave a 2× speedup
//! over Burn's NdArray backend.
//!
//! Whether a real CBLAS is linked is decided by `build.rs`, which sets the
//! `rlx_cpu_blas` cfg (used throughout this file) only when it actually links
//! one: the `blas` feature must be on AND the target must have a BLAS —
//! Accelerate on Apple, OpenBLAS/MKL on x86_64, or wherever `OPENBLAS_LIB_DIR`
//! points. On targets with no BLAS (aarch64 Linux — Raspberry Pi / cross /
//! QEMU — and wasm), or with `--no-default-features`, `rlx_cpu_blas` is unset
//! and the extern is replaced by a portable scalar/SIMD gemm with the same
//! calling convention, so every consumer here keeps working (slower, correct).

/// Which CBLAS/LAPACK vendor is linked into this build, if any.
///
/// - `"accelerate"` — Apple Accelerate (macOS / iOS)
/// - `"openblas"` — OpenBLAS
/// - `"mkl"` — Intel oneMKL
/// - `"cblas"` — generic CBLAS (rare)
/// - `None` — portable scalar/SIMD fallback (`--no-default-features` or no system BLAS)
#[must_use]
pub fn linked_blas() -> Option<&'static str> {
    if cfg!(rlx_cpu_blas_accelerate) {
        Some("accelerate")
    } else if cfg!(rlx_cpu_blas_openblas) {
        Some("openblas")
    } else if cfg!(rlx_cpu_blas_mkl) {
        Some("mkl")
    } else if cfg!(rlx_cpu_blas) {
        Some("cblas")
    } else {
        None
    }
}

#[cfg(rlx_cpu_blas)]
unsafe extern "C" {
    #[link_name = "cblas_sgemm"]
    fn cblas_sgemm_raw(
        order: i32,
        transa: i32,
        transb: i32,
        m: i32,
        n: i32,
        k: i32,
        alpha: f32,
        a: *const f32,
        lda: i32,
        b: *const f32,
        ldb: i32,
        beta: f32,
        c: *mut f32,
        ldc: i32,
    );

    fn cblas_sgemv(
        order: i32,
        trans: i32,
        m: i32,
        n: i32,
        alpha: f32,
        a: *const f32,
        lda: i32,
        x: *const f32,
        incx: i32,
        beta: f32,
        y: *mut f32,
        incy: i32,
    );

    fn cblas_sger(
        order: i32,
        m: i32,
        n: i32,
        alpha: f32,
        x: *const f32,
        incx: i32,
        y: *const f32,
        incy: i32,
        a: *mut f32,
        lda: i32,
    );

    fn cblas_sscal(n: i32, alpha: f32, x: *mut f32, incx: i32);
}

/// Empty-dim-guarded wrapper over vendor `cblas_sgemm`. Vendor BLAS rejects a
/// zero contraction (`lda=0` when `K=0`) with "Parameter 9 had an invalid
/// value", but an empty matmul is well-defined: `K=0` ⇒ the `A·B` term is a sum
/// over nothing (0), so `C = beta·C`; `M==0`/`N==0` ⇒ `C` is empty (nothing to
/// write). Handle those here without calling BLAS. This is reachable now that
/// the importer emits genuinely empty tensors (0-length dims) instead of
/// promoting them to `[1]` — e.g. supertonic's ConstantOfShape-derived operand
/// hit `cblas_sgemm K=0` and crashed the whole CPU run. The scalar fallback
/// (`cfg(not(rlx_cpu_blas))`) already degenerates correctly for `K=0`, so this
/// wrapper only guards the vendor path. Signature matches the old extern, so
/// every `cblas_sgemm(...)` call site is covered unchanged.
#[cfg(rlx_cpu_blas)]
#[allow(non_snake_case, clippy::too_many_arguments)]
#[inline]
unsafe fn cblas_sgemm(
    order: i32,
    transa: i32,
    transb: i32,
    m: i32,
    n: i32,
    k: i32,
    alpha: f32,
    a: *const f32,
    lda: i32,
    b: *const f32,
    ldb: i32,
    beta: f32,
    c: *mut f32,
    ldc: i32,
) {
    if m <= 0 || n <= 0 {
        return;
    }
    if k <= 0 {
        let (mm, nn, ldc) = (m as usize, n as usize, ldc as usize);
        for i in 0..mm {
            for j in 0..nn {
                let cp = unsafe { c.add(i * ldc + j) };
                unsafe { *cp *= beta };
            }
        }
        return;
    }
    unsafe {
        cblas_sgemm_raw(
            order, transa, transb, m, n, k, alpha, a, lda, b, ldb, beta, c, ldc,
        );
    }
}

/// Cap BLAS's own thread pool when Rayon owns outer parallelism.
///
/// Nested OpenBLAS/MKL × Rayon oversubscription (N² threads) is catastrophic
/// for CNN training: wall time jumps from ~1–2 min/epoch to tens of minutes.
/// Called once when the Rayon pool starts (and from `RLX_FAST_CONV`). Honours
/// an explicit user override via `OPENBLAS_NUM_THREADS` / `OMP_NUM_THREADS` /
/// `MKL_NUM_THREADS` / `VECLIB_MAXIMUM_THREADS`.
pub fn limit_inner_threads() {
    use std::sync::Once;
    static ONCE: Once = Once::new();
    ONCE.call_once(|| {
        let user_set = [
            "OPENBLAS_NUM_THREADS",
            "OMP_NUM_THREADS",
            "MKL_NUM_THREADS",
            "VECLIB_MAXIMUM_THREADS",
        ]
        .iter()
        .any(|k| std::env::var_os(k).is_some());
        if user_set {
            return;
        }
        set_blas_num_threads(1);
    });
}

/// Set the vendor BLAS thread pool size (OpenBLAS / MKL / Accelerate).
///
/// Prefer keeping this at 1 while Rayon owns outer loops; briefly raise it
/// for a single huge GEMM on the main thread (see [`sgemm_auto`]).
/// No-ops when the requested size already matches the last set value.
pub fn set_blas_num_threads(n: i32) {
    let n = n.max(1);
    use std::sync::atomic::{AtomicI32, Ordering};
    // Sentinel `0` (not a valid thread count) so the *first* pin actually calls
    // into the vendor BLAS. Starting at `1` wrongly assumes OpenBLAS/MKL boot
    // single-threaded — the OpenBLAS pthread build defaults to all cores, so a
    // no-op first call left it oversubscribed (N rayon workers × N BLAS threads).
    static LAST: AtomicI32 = AtomicI32::new(0);
    if LAST.swap(n, Ordering::SeqCst) == n {
        return;
    }
    #[cfg(rlx_cpu_blas_accelerate)]
    {
        // SAFETY: env read by Accelerate at first use / next GEMM.
        unsafe {
            std::env::set_var("VECLIB_MAXIMUM_THREADS", n.to_string());
        }
    }
    #[cfg(rlx_cpu_blas_openblas)]
    {
        unsafe extern "C" {
            fn openblas_set_num_threads(n: i32);
        }
        unsafe {
            openblas_set_num_threads(n);
        }
    }
    #[cfg(rlx_cpu_blas_mkl)]
    {
        unsafe extern "C" {
            fn MKL_Set_Num_Threads(n: i32);
        }
        unsafe {
            MKL_Set_Num_Threads(n);
        }
    }
    let _ = n; // no-op when no BLAS linked
}

/// Ensure BLAS is at `n` threads (idempotent). Used so consecutive huge
/// GEMMs avoid OpenBLAS set/teardown thrash, while Rayon entry points pin
/// back to 1.
#[inline]
pub fn ensure_blas_threads(n: i32) {
    set_blas_num_threads(n);
}

/// Run `f` with a temporary BLAS thread count, then restore to 1.
///
/// Only safe off the Rayon pool (nested MT BLAS × Rayon oversubscribes).
#[inline]
pub fn with_blas_threads<R>(n: i32, f: impl FnOnce() -> R) -> R {
    if n <= 1 || rayon::current_thread_index().is_some() {
        return f();
    }
    set_blas_num_threads(n);
    let out = f();
    set_blas_num_threads(1);
    out
}

#[cfg(not(rlx_cpu_blas))]
#[allow(non_snake_case, clippy::too_many_arguments)]
#[inline]
unsafe fn cblas_sgemm(
    _order: i32,
    transa: i32,
    transb: i32,
    m: i32,
    n: i32,
    k: i32,
    alpha: f32,
    a: *const f32,
    lda: i32,
    b: *const f32,
    ldb: i32,
    beta: f32,
    c: *mut f32,
    ldc: i32,
) {
    // Row-major fallback. _order is ignored — every call site in this file
    // passes ROW_MAJOR (101). Supports both NoTrans/Trans on either operand
    // and arbitrary positive lda/ldb/ldc.
    let m = m as usize;
    let n = n as usize;
    let k = k as usize;
    let lda = lda as usize;
    let ldb = ldb as usize;
    let ldc = ldc as usize;
    let trans_a = transa != NO_TRANS;
    let trans_b = transb != NO_TRANS;

    // Fast path: the common NoTrans×NoTrans case, reordered i→p→j so the
    // inner j-loop is unit-stride in both B and C. The compiler
    // auto-vectorizes it (NEON / AVX) and it stays cache-resident — several×
    // to an order of magnitude over the naive p-innermost loop (whose B
    // access strides by ldb) on a BLAS-less Pi-class CPU. This is a
    // different summation order than the scalar path below, but every GEMM
    // caller here already tolerates the vendor-BLAS order too. Transposed
    // cases (rare in-tree) fall through to the general scalar loop.
    if !trans_a && !trans_b {
        for i in 0..m {
            let crow = unsafe { c.add(i * ldc) };
            if beta == 0.0 {
                for j in 0..n {
                    unsafe { *crow.add(j) = 0.0 };
                }
            } else if beta != 1.0 {
                for j in 0..n {
                    unsafe { *crow.add(j) *= beta };
                }
            }
            for p in 0..k {
                let aip = alpha * unsafe { *a.add(i * lda + p) };
                if aip == 0.0 {
                    continue;
                }
                let brow = unsafe { b.add(p * ldb) };
                for j in 0..n {
                    unsafe { *crow.add(j) += aip * *brow.add(j) };
                }
            }
        }
        return;
    }

    for i in 0..m {
        for j in 0..n {
            let mut acc: f32 = 0.0;
            for p in 0..k {
                let av = if trans_a {
                    unsafe { *a.add(p * lda + i) }
                } else {
                    unsafe { *a.add(i * lda + p) }
                };
                let bv = if trans_b {
                    unsafe { *b.add(j * ldb + p) }
                } else {
                    unsafe { *b.add(p * ldb + j) }
                };
                acc += av * bv;
            }
            let cp = unsafe { c.add(i * ldc + j) };
            unsafe {
                *cp = alpha * acc + beta * *cp;
            }
        }
    }
}

#[cfg(not(rlx_cpu_blas))]
#[allow(non_snake_case, clippy::too_many_arguments)]
#[inline]
unsafe fn cblas_sgemv(
    _order: i32,
    trans: i32,
    m: i32,
    n: i32,
    alpha: f32,
    a: *const f32,
    lda: i32,
    x: *const f32,
    _incx: i32,
    beta: f32,
    y: *mut f32,
    _incy: i32,
) {
    let m = m as usize;
    let n = n as usize;
    let lda = lda as usize;
    let trans_a = trans != NO_TRANS;
    for i in 0..m {
        let mut acc = 0f32;
        for j in 0..n {
            let av = if trans_a {
                unsafe { *a.add(j * lda + i) }
            } else {
                unsafe { *a.add(i * lda + j) }
            };
            acc += av * unsafe { *x.add(j) };
        }
        let yp = unsafe { y.add(i) };
        unsafe {
            *yp = alpha * acc + beta * *yp;
        }
    }
}

#[cfg(not(rlx_cpu_blas))]
#[allow(non_snake_case, clippy::too_many_arguments)]
#[inline]
unsafe fn cblas_sger(
    _order: i32,
    m: i32,
    n: i32,
    alpha: f32,
    x: *const f32,
    _incx: i32,
    y: *const f32,
    _incy: i32,
    a: *mut f32,
    lda: i32,
) {
    let m = m as usize;
    let n = n as usize;
    let lda = lda as usize;
    for i in 0..m {
        let xi = unsafe { *x.add(i) };
        for j in 0..n {
            let yj = unsafe { *y.add(j) };
            let ap = unsafe { a.add(i * lda + j) };
            unsafe {
                *ap += alpha * xi * yj;
            }
        }
    }
}

#[cfg(not(rlx_cpu_blas))]
#[inline]
unsafe fn cblas_sscal(n: i32, alpha: f32, x: *mut f32, _incx: i32) {
    for i in 0..n as usize {
        let xp = unsafe { x.add(i) };
        unsafe {
            *xp *= alpha;
        }
    }
}

// ── f64 BLAS / LAPACK ─────────────────────────────────────────
//
// Accelerate's vecLib (macOS) and OpenBLAS (Linux/Win) both export
// these. They follow the exact same calling conventions as the f32
// variants — only the element type changes.
//
// `dgesv_` is the LAPACK Fortran ABI name (column-major, trailing
// underscore). It does an in-place LU factorization of `a` and then
// solves `a · x = b`, overwriting `b` with `x`. Row-major callers
// must transpose A in/out (or transpose b's leading-dim convention)
// — see the `dgesv` wrapper below for the row-major adapter.

#[cfg(rlx_cpu_blas)]
unsafe extern "C" {
    fn cblas_dgemm(
        order: i32,
        transa: i32,
        transb: i32,
        m: i32,
        n: i32,
        k: i32,
        alpha: f64,
        a: *const f64,
        lda: i32,
        b: *const f64,
        ldb: i32,
        beta: f64,
        c: *mut f64,
        ldc: i32,
    );

    /// LAPACK column-major dgesv:
    ///   A · X = B,  A: [n, n],  B: [n, nrhs],  pivot: [n] (i32).
    /// Overwrites A with its LU factors and B with the solution X.
    /// `info_out`: 0 = success, k>0 = U[k-1,k-1] is exactly zero.
    #[link_name = "dgesv_"]
    fn lapack_dgesv(
        n: *const i32,
        nrhs: *const i32,
        a: *mut f64,
        lda: *const i32,
        ipiv: *mut i32,
        b: *mut f64,
        ldb: *const i32,
        info_out: *mut i32,
    );

    /// Single-precision twin of dgesv. Same shape contract.
    #[link_name = "sgesv_"]
    fn lapack_sgesv(
        n: *const i32,
        nrhs: *const i32,
        a: *mut f32,
        lda: *const i32,
        ipiv: *mut i32,
        b: *mut f32,
        ldb: *const i32,
        info_out: *mut i32,
    );

    /// dpotrf — Cholesky factorization of an SPD matrix.
    /// `uplo` is a single byte: b'U' or b'L' (passed as i8).
    #[link_name = "dpotrf_"]
    fn lapack_dpotrf(
        uplo: *const i8,
        n: *const i32,
        a: *mut f64,
        lda: *const i32,
        info_out: *mut i32,
    );

    /// dgetrf — LU factorization with partial pivoting (det / logdet).
    #[link_name = "dgetrf_"]
    fn lapack_dgetrf(
        m: *const i32,
        n: *const i32,
        a: *mut f64,
        lda: *const i32,
        ipiv: *mut i32,
        info_out: *mut i32,
    );

    /// dsyevd — symmetric eigendecomp via divide-and-conquer.
    /// `jobz`: b'N' (eigenvalues only) or b'V' (also eigenvectors).
    #[link_name = "dsyevd_"]
    fn lapack_dsyevd(
        jobz: *const i8,
        uplo: *const i8,
        n: *const i32,
        a: *mut f64,
        lda: *const i32,
        w: *mut f64,
        work: *mut f64,
        lwork: *const i32,
        iwork: *mut i32,
        liwork: *const i32,
        info_out: *mut i32,
    );

    /// dgeqrf — QR factorization. On output, A's upper triangle holds R;
    /// elementary reflectors stored below the diagonal + scalar factors `tau`.
    #[link_name = "dgeqrf_"]
    fn lapack_dgeqrf(
        m: *const i32,
        n: *const i32,
        a: *mut f64,
        lda: *const i32,
        tau: *mut f64,
        work: *mut f64,
        lwork: *const i32,
        info_out: *mut i32,
    );

    /// dorgqr — generate Q from the elementary reflectors produced by dgeqrf.
    #[link_name = "dorgqr_"]
    fn lapack_dorgqr(
        m: *const i32,
        n: *const i32,
        k: *const i32,
        a: *mut f64,
        lda: *const i32,
        tau: *const f64,
        work: *mut f64,
        lwork: *const i32,
        info_out: *mut i32,
    );

    /// dgesvd — singular value decomposition. `jobu`/`jobvt`: b'A' (all),
    /// b'S' (singular vectors only), b'O' (overwrite A), b'N' (none).
    #[link_name = "dgesvd_"]
    fn lapack_dgesvd(
        jobu: *const i8,
        jobvt: *const i8,
        m: *const i32,
        n: *const i32,
        a: *mut f64,
        lda: *const i32,
        s: *mut f64,
        u: *mut f64,
        ldu: *const i32,
        vt: *mut f64,
        ldvt: *const i32,
        work: *mut f64,
        lwork: *const i32,
        info_out: *mut i32,
    );

    /// dgesdd — SVD via divide-and-conquer (the driver NumPy's `linalg.pinv`
    /// / `linalg.svd` use). Differs from `dgesvd` on ill-conditioned/degenerate
    /// inputs, so matching NumPy bit-for-bit requires *this* routine.
    #[link_name = "dgesdd_"]
    fn lapack_dgesdd(
        jobz: *const i8,
        m: *const i32,
        n: *const i32,
        a: *mut f64,
        lda: *const i32,
        s: *mut f64,
        u: *mut f64,
        ldu: *const i32,
        vt: *mut f64,
        ldvt: *const i32,
        work: *mut f64,
        lwork: *const i32,
        iwork: *mut i32,
        info_out: *mut i32,
    );

    /// dgelsd — minimum-norm least-squares via SVD (the driver SciPy's
    /// `linalg.lstsq` uses by default). Solves `min‖A·X − B‖`.
    #[link_name = "dgelsd_"]
    fn lapack_dgelsd(
        m: *const i32,
        n: *const i32,
        nrhs: *const i32,
        a: *mut f64,
        lda: *const i32,
        b: *mut f64,
        ldb: *const i32,
        s: *mut f64,
        rcond: *const f64,
        rank: *mut i32,
        work: *mut f64,
        lwork: *const i32,
        iwork: *mut i32,
        info_out: *mut i32,
    );

    /// cblas_dtrsm — BLAS-3 row-major-friendly triangular solve.
    /// `op(A) · X = α·B`  (or  `X · op(A) = α·B` for `side=Right`).
    /// A is triangular, B overwritten with X.
    fn cblas_dtrsm(
        order: i32,
        side: i32,
        uplo: i32,
        transa: i32,
        diag: i32,
        m: i32,
        n: i32,
        alpha: f64,
        a: *const f64,
        lda: i32,
        b: *mut f64,
        ldb: i32,
    );
}

#[cfg(not(rlx_cpu_blas))]
#[allow(non_snake_case, clippy::too_many_arguments)]
#[inline]
unsafe fn cblas_dgemm(
    _order: i32,
    transa: i32,
    transb: i32,
    m: i32,
    n: i32,
    k: i32,
    alpha: f64,
    a: *const f64,
    lda: i32,
    b: *const f64,
    ldb: i32,
    beta: f64,
    c: *mut f64,
    ldc: i32,
) {
    let m = m as usize;
    let n = n as usize;
    let k = k as usize;
    let lda = lda as usize;
    let ldb = ldb as usize;
    let ldc = ldc as usize;
    let trans_a = transa != NO_TRANS;
    let trans_b = transb != NO_TRANS;
    for i in 0..m {
        for j in 0..n {
            let mut acc: f64 = 0.0;
            for p in 0..k {
                let av = if trans_a {
                    unsafe { *a.add(p * lda + i) }
                } else {
                    unsafe { *a.add(i * lda + p) }
                };
                let bv = if trans_b {
                    unsafe { *b.add(j * ldb + p) }
                } else {
                    unsafe { *b.add(p * ldb + j) }
                };
                acc += av * bv;
            }
            let cp = unsafe { c.add(i * ldc + j) };
            unsafe {
                *cp = alpha * acc + beta * *cp;
            }
        }
    }
}

/// Pure-Rust LU + solve fallback for builds without BLAS/LAPACK.
/// Partial pivoting; column-major in/out to match LAPACK's ABI.
/// Returns 0 on success, k+1 if U[k,k] is zero (singular).
#[cfg(not(rlx_cpu_blas))]
#[allow(non_snake_case, clippy::too_many_arguments)]
unsafe fn lapack_dgesv(
    n: *const i32,
    nrhs: *const i32,
    a: *mut f64,
    lda: *const i32,
    ipiv: *mut i32,
    b: *mut f64,
    ldb: *const i32,
    info_out: *mut i32,
) {
    let nn = unsafe { *n } as usize;
    let nrhs = unsafe { *nrhs } as usize;
    let lda = unsafe { *lda } as usize;
    let ldb = unsafe { *ldb } as usize;
    // Column-major access helper: a[i, j] = a[j*lda + i]
    let aij = |a: *mut f64, i: usize, j: usize| unsafe { a.add(j * lda + i) };
    for k in 0..nn {
        // Pivot: row with max |a[k..,k]|
        let mut piv = k;
        let mut max_abs = unsafe { *aij(a, k, k) }.abs();
        for i in (k + 1)..nn {
            let v = unsafe { *aij(a, i, k) }.abs();
            if v > max_abs {
                max_abs = v;
                piv = i;
            }
        }
        unsafe {
            *ipiv.add(k) = (piv + 1) as i32;
        }
        if max_abs == 0.0 {
            unsafe {
                *info_out = (k + 1) as i32;
            }
            return;
        }
        // Swap rows piv and k in A
        if piv != k {
            for j in 0..nn {
                let p1 = aij(a, k, j);
                let p2 = aij(a, piv, j);
                unsafe {
                    std::ptr::swap(p1, p2);
                }
            }
            for j in 0..nrhs {
                let p1 = unsafe { b.add(j * ldb + k) };
                let p2 = unsafe { b.add(j * ldb + piv) };
                unsafe {
                    std::ptr::swap(p1, p2);
                }
            }
        }
        // Eliminate
        let akk = unsafe { *aij(a, k, k) };
        for i in (k + 1)..nn {
            let factor = unsafe { *aij(a, i, k) } / akk;
            unsafe {
                *aij(a, i, k) = factor;
            }
            for j in (k + 1)..nn {
                let v = unsafe { *aij(a, i, j) } - factor * unsafe { *aij(a, k, j) };
                unsafe {
                    *aij(a, i, j) = v;
                }
            }
            for j in 0..nrhs {
                let v = unsafe { *b.add(j * ldb + i) } - factor * unsafe { *b.add(j * ldb + k) };
                unsafe {
                    *b.add(j * ldb + i) = v;
                }
            }
        }
    }
    // Back-substitute U·x = y
    for j in 0..nrhs {
        for i in (0..nn).rev() {
            let mut sum = unsafe { *b.add(j * ldb + i) };
            for k in (i + 1)..nn {
                sum -= unsafe { *aij(a, i, k) } * unsafe { *b.add(j * ldb + k) };
            }
            unsafe {
                *b.add(j * ldb + i) = sum / *aij(a, i, i);
            }
        }
    }
    unsafe {
        *info_out = 0;
    }
}

/// f32 twin of `lapack_dgesv` for no-blas builds. Same algorithm,
/// f32 instead of f64.
#[cfg(not(rlx_cpu_blas))]
#[allow(non_snake_case, clippy::too_many_arguments)]
unsafe fn lapack_sgesv(
    n: *const i32,
    nrhs: *const i32,
    a: *mut f32,
    lda: *const i32,
    ipiv: *mut i32,
    b: *mut f32,
    ldb: *const i32,
    info_out: *mut i32,
) {
    let nn = unsafe { *n } as usize;
    let nrhs = unsafe { *nrhs } as usize;
    let lda = unsafe { *lda } as usize;
    let ldb = unsafe { *ldb } as usize;
    let aij = |a: *mut f32, i: usize, j: usize| unsafe { a.add(j * lda + i) };
    for k in 0..nn {
        let mut piv = k;
        let mut max_abs = unsafe { *aij(a, k, k) }.abs();
        for i in (k + 1)..nn {
            let v = unsafe { *aij(a, i, k) }.abs();
            if v > max_abs {
                max_abs = v;
                piv = i;
            }
        }
        unsafe {
            *ipiv.add(k) = (piv + 1) as i32;
        }
        if max_abs == 0.0 {
            unsafe {
                *info_out = (k + 1) as i32;
            }
            return;
        }
        if piv != k {
            for j in 0..nn {
                let p1 = aij(a, k, j);
                let p2 = aij(a, piv, j);
                unsafe {
                    std::ptr::swap(p1, p2);
                }
            }
            for j in 0..nrhs {
                let p1 = unsafe { b.add(j * ldb + k) };
                let p2 = unsafe { b.add(j * ldb + piv) };
                unsafe {
                    std::ptr::swap(p1, p2);
                }
            }
        }
        let akk = unsafe { *aij(a, k, k) };
        for i in (k + 1)..nn {
            let factor = unsafe { *aij(a, i, k) } / akk;
            unsafe {
                *aij(a, i, k) = factor;
            }
            for j in (k + 1)..nn {
                let v = unsafe { *aij(a, i, j) } - factor * unsafe { *aij(a, k, j) };
                unsafe {
                    *aij(a, i, j) = v;
                }
            }
            for j in 0..nrhs {
                let v = unsafe { *b.add(j * ldb + i) } - factor * unsafe { *b.add(j * ldb + k) };
                unsafe {
                    *b.add(j * ldb + i) = v;
                }
            }
        }
    }
    for j in 0..nrhs {
        for i in (0..nn).rev() {
            let mut sum = unsafe { *b.add(j * ldb + i) };
            for k in (i + 1)..nn {
                sum -= unsafe { *aij(a, i, k) } * unsafe { *b.add(j * ldb + k) };
            }
            unsafe {
                *b.add(j * ldb + i) = sum / *aij(a, i, i);
            }
        }
    }
    unsafe {
        *info_out = 0;
    }
}

/// f64 GEMM. C = A @ B, all row-major, A: [m, k], B: [k, n], C: [m, n].
#[inline]
pub fn dgemm(a: &[f64], b: &[f64], c: &mut [f64], m: usize, k: usize, n: usize) {
    unsafe {
        cblas_dgemm(
            ROW_MAJOR,
            NO_TRANS,
            NO_TRANS,
            m as i32,
            n as i32,
            k as i32,
            1.0,
            a.as_ptr(),
            k as i32,
            b.as_ptr(),
            n as i32,
            0.0,
            c.as_mut_ptr(),
            n as i32,
        );
    }
}

/// Solve `A · x = b` in-place. Row-major caller's API: `a` is `[n, n]`
/// row-major and `b` is `[n]` (single-RHS) or `[n, nrhs]` row-major.
/// On return, `b` holds the solution. `a` is overwritten with LU
/// factors (caller's copy is destroyed).
///
/// Implementation detail: LAPACK's `dgesv_` is column-major, so we
/// transpose `A` in place (square ⇒ cheap), call `dgesv_` with the
/// natural column-major `B` interpretation, and transpose `B` back if
/// `nrhs > 1`. For `nrhs = 1` the column-major and row-major layouts
/// of a vector are identical — no transpose on B needed.
///
/// Returns 0 on success; k > 0 means `U[k-1, k-1]` was exactly zero
/// (singular system).
pub fn dgesv(a: &mut [f64], b: &mut [f64], n: usize, nrhs: usize) -> i32 {
    assert_eq!(a.len(), n * n, "dgesv: A must be n×n");
    assert_eq!(b.len(), n * nrhs, "dgesv: B must be n×nrhs");
    // Row→col-major: in-place transpose of the square A.
    for i in 0..n {
        for j in (i + 1)..n {
            a.swap(i * n + j, j * n + i);
        }
    }
    // Same for B if nrhs > 1.
    if nrhs > 1 {
        let mut tmp = vec![0f64; n * nrhs];
        for i in 0..n {
            for j in 0..nrhs {
                tmp[j * n + i] = b[i * nrhs + j];
            }
        }
        b.copy_from_slice(&tmp);
    }
    let mut ipiv = vec![0i32; n];
    let mut info: i32 = 0;
    let nn = n as i32;
    let nrhs_i = nrhs as i32;
    unsafe {
        lapack_dgesv(
            &nn,
            &nrhs_i,
            a.as_mut_ptr(),
            &nn,
            ipiv.as_mut_ptr(),
            b.as_mut_ptr(),
            &nn,
            &mut info,
        );
    }
    // Col-major B back to row-major.
    if nrhs > 1 && info == 0 {
        let mut tmp = vec![0f64; n * nrhs];
        for j in 0..nrhs {
            for i in 0..n {
                tmp[i * nrhs + j] = b[j * n + i];
            }
        }
        b.copy_from_slice(&tmp);
    }
    info
}

/// f32 twin of `dgesv`. Same row-major caller's API + same return
/// code semantics. Uses LAPACK's `sgesv_` under the hood.
pub fn sgesv(a: &mut [f32], b: &mut [f32], n: usize, nrhs: usize) -> i32 {
    assert_eq!(a.len(), n * n, "sgesv: A must be n×n");
    assert_eq!(b.len(), n * nrhs, "sgesv: B must be n×nrhs");
    for i in 0..n {
        for j in (i + 1)..n {
            a.swap(i * n + j, j * n + i);
        }
    }
    if nrhs > 1 {
        let mut tmp = vec![0f32; n * nrhs];
        for i in 0..n {
            for j in 0..nrhs {
                tmp[j * n + i] = b[i * nrhs + j];
            }
        }
        b.copy_from_slice(&tmp);
    }
    let mut ipiv = vec![0i32; n];
    let mut info: i32 = 0;
    let nn = n as i32;
    let nrhs_i = nrhs as i32;
    unsafe {
        lapack_sgesv(
            &nn,
            &nrhs_i,
            a.as_mut_ptr(),
            &nn,
            ipiv.as_mut_ptr(),
            b.as_mut_ptr(),
            &nn,
            &mut info,
        );
    }
    if nrhs > 1 && info == 0 {
        let mut tmp = vec![0f32; n * nrhs];
        for j in 0..nrhs {
            for i in 0..n {
                tmp[i * nrhs + j] = b[j * n + i];
            }
        }
        b.copy_from_slice(&tmp);
    }
    info
}

const ROW_MAJOR: i32 = 101;
const NO_TRANS: i32 = 111;
const TRANS: i32 = 112;
const CBLAS_LEFT: i32 = 141;
#[allow(dead_code)]
const CBLAS_RIGHT: i32 = 142;
const CBLAS_UPPER: i32 = 121;
const CBLAS_LOWER: i32 = 122;
const CBLAS_NON_UNIT: i32 = 131;
#[allow(dead_code)]
const CBLAS_UNIT: i32 = 132;

// ── Pure-Rust LAPACK fallback (no linked BLAS) ───────────────────
//
// When `rlx_cpu_blas` is unset — the `blas` feature is off, or the
// target has no linked BLAS/LAPACK (a bare aarch64 / Raspberry Pi with
// no OpenBLAS, wasm, …) — the linalg ops route through these instead of
// panicking. They are dependency-free reference implementations that
// match the column-major Fortran ABI the row-major wrappers above
// expect, so those wrappers stay backend-agnostic. Correct but not
// tuned; the linked-BLAS path remains the fast default where present.
//
// Numerics: Cholesky (Banachiewicz), LU (partial-pivot), symmetric
// eigendecomposition (cyclic Jacobi), QR (Householder + `dorg2r`), thin
// SVD (one-sided Jacobi); least-squares and `dtrsm` build on those.

/// Cholesky of a packed column-major `n×n` SPD matrix (`lda = n`).
/// `upper` ⇒ factor `Uᵀ·U` into the upper triangle, else `L·Lᵀ` into the
/// lower. Returns 0, or `k>0` if the leading minor of order `k` is not
/// positive-definite. Mirrors LAPACK `dpotrf`.
#[cfg(not(rlx_cpu_blas))]
fn cholesky_colmajor(a: &mut [f64], n: usize, upper: bool) -> i32 {
    let at = |i: usize, j: usize| i + j * n; // column-major, lda = n
    if upper {
        for j in 0..n {
            for i in 0..=j {
                let mut sum = a[at(i, j)];
                for l in 0..i {
                    sum -= a[at(l, i)] * a[at(l, j)];
                }
                if i < j {
                    let d = a[at(i, i)];
                    a[at(i, j)] = sum / d;
                } else {
                    if sum <= 0.0 {
                        return (j + 1) as i32;
                    }
                    a[at(j, j)] = sum.sqrt();
                }
            }
        }
    } else {
        for j in 0..n {
            for i in j..n {
                let mut sum = a[at(i, j)];
                for l in 0..j {
                    sum -= a[at(i, l)] * a[at(j, l)];
                }
                if i == j {
                    if sum <= 0.0 {
                        return (j + 1) as i32;
                    }
                    a[at(j, j)] = sum.sqrt();
                } else {
                    let d = a[at(j, j)];
                    a[at(i, j)] = sum / d;
                }
            }
        }
    }
    0
}

/// LU with partial pivoting of a packed column-major `m×n` matrix
/// (`lda = m`): unit-lower `L` below the diagonal, `U` on and above,
/// 1-based row pivots in `ipiv`. Returns 0, or the 1-based index of the
/// first zero pivot. Mirrors LAPACK `dgetrf`.
#[cfg(not(rlx_cpu_blas))]
fn lu_colmajor(a: &mut [f64], m: usize, n: usize, ipiv: &mut [i32]) -> i32 {
    let at = |i: usize, j: usize| i + j * m;
    let mut info = 0;
    for j in 0..m.min(n) {
        let mut p = j;
        let mut maxv = a[at(j, j)].abs();
        for i in (j + 1)..m {
            let v = a[at(i, j)].abs();
            if v > maxv {
                maxv = v;
                p = i;
            }
        }
        ipiv[j] = (p + 1) as i32;
        if a[at(p, j)] != 0.0 {
            if p != j {
                for c in 0..n {
                    a.swap(at(j, c), at(p, c));
                }
            }
            let d = a[at(j, j)];
            for i in (j + 1)..m {
                a[at(i, j)] /= d;
            }
        } else if info == 0 {
            info = (j + 1) as i32;
        }
        for c in (j + 1)..n {
            let ajc = a[at(j, c)];
            if ajc != 0.0 {
                for i in (j + 1)..m {
                    a[at(i, c)] -= a[at(i, j)] * ajc;
                }
            }
        }
    }
    info
}

/// Symmetric eigendecomposition via cyclic Jacobi. `s` is a row-major
/// `n×n` symmetric matrix (consumed). Returns `(eigenvalues ascending,
/// eigenvectors)`, where eigenvector `j` is column `j` of the row-major
/// `n×n` result (`evec[i*n + j]`). Orthonormal by construction.
#[cfg(not(rlx_cpu_blas))]
fn sym_eig_jacobi(mut s: Vec<f64>, n: usize) -> (Vec<f64>, Vec<f64>) {
    let mut v = vec![0f64; n * n];
    for i in 0..n {
        v[i * n + i] = 1.0;
    }
    if n <= 1 {
        let w = if n == 1 { vec![s[0]] } else { vec![] };
        return (w, v);
    }
    for _sweep in 0..100 {
        let mut off = 0.0;
        for p in 0..n {
            for q in (p + 1)..n {
                off += s[p * n + q] * s[p * n + q];
            }
        }
        if off <= 1e-300 {
            break;
        }
        let mut rotated = false;
        for p in 0..n {
            for q in (p + 1)..n {
                let apq = s[p * n + q];
                let app = s[p * n + p];
                let aqq = s[q * n + q];
                if apq.abs() <= f64::EPSILON * (app.abs() + aqq.abs()).max(f64::MIN_POSITIVE) {
                    continue;
                }
                rotated = true;
                // Rotation angle that annihilates s[p][q]: solves
                // t² + 2θt − 1 = 0 for the smaller |t| (stable form).
                let theta = (aqq - app) / (2.0 * apq);
                let t = if theta == 0.0 {
                    1.0
                } else {
                    theta.signum() / (theta.abs() + (theta * theta + 1.0).sqrt())
                };
                let c = 1.0 / (t * t + 1.0).sqrt();
                let sn = t * c;
                // s ← Gᵀ s G: column update (s·G) then row update (Gᵀ·).
                for i in 0..n {
                    let sip = s[i * n + p];
                    let siq = s[i * n + q];
                    s[i * n + p] = c * sip - sn * siq;
                    s[i * n + q] = sn * sip + c * siq;
                }
                for i in 0..n {
                    let spi = s[p * n + i];
                    let sqi = s[q * n + i];
                    s[p * n + i] = c * spi - sn * sqi;
                    s[q * n + i] = sn * spi + c * sqi;
                }
                // Accumulate eigenvectors: V ← V·G.
                for i in 0..n {
                    let vip = v[i * n + p];
                    let viq = v[i * n + q];
                    v[i * n + p] = c * vip - sn * viq;
                    v[i * n + q] = sn * vip + c * viq;
                }
            }
        }
        if !rotated {
            break;
        }
    }
    let evals: Vec<f64> = (0..n).map(|i| s[i * n + i]).collect();
    let mut idx: Vec<usize> = (0..n).collect();
    idx.sort_by(|&a, &b| {
        evals[a]
            .partial_cmp(&evals[b])
            .unwrap_or(std::cmp::Ordering::Equal)
    });
    let w: Vec<f64> = idx.iter().map(|&i| evals[i]).collect();
    let mut vout = vec![0f64; n * n];
    for (newj, &oldj) in idx.iter().enumerate() {
        for i in 0..n {
            vout[i * n + newj] = v[i * n + oldj];
        }
    }
    (w, vout)
}

/// Thin SVD `A = U·diag(s)·Vᵀ` of a row-major `m×n` matrix via one-sided
/// Jacobi. Returns `(U row-major m×k, s descending, Vᵀ row-major k×n)`,
/// `k = min(m,n)`; singular values ≥ 0, `U`/`V` orthonormal.
#[cfg(not(rlx_cpu_blas))]
fn svd_thin_rowmajor(am: &[f64], m: usize, n: usize) -> (Vec<f64>, Vec<f64>, Vec<f64>) {
    let k = m.min(n);
    // One-sided Jacobi orthogonalizes columns, so work on the tall side.
    let (rows, cols, transposed) = if m >= n { (m, n, false) } else { (n, m, true) };
    let mut w = vec![0f64; rows * cols];
    if transposed {
        for i in 0..rows {
            for j in 0..cols {
                w[i * cols + j] = am[j * n + i];
            }
        }
    } else {
        w.copy_from_slice(am);
    }
    let mut vv = vec![0f64; cols * cols];
    for i in 0..cols {
        vv[i * cols + i] = 1.0;
    }
    for _ in 0..60 {
        let mut changed = false;
        for p in 0..cols {
            for q in (p + 1)..cols {
                let (mut alpha, mut beta, mut gamma) = (0.0, 0.0, 0.0);
                for r in 0..rows {
                    let wp = w[r * cols + p];
                    let wq = w[r * cols + q];
                    alpha += wp * wp;
                    beta += wq * wq;
                    gamma += wp * wq;
                }
                if gamma.abs() <= 1e-15 * (alpha * beta).sqrt() {
                    continue;
                }
                changed = true;
                let zeta = (beta - alpha) / (2.0 * gamma);
                let t = if zeta == 0.0 {
                    1.0
                } else {
                    zeta.signum() / (zeta.abs() + (1.0 + zeta * zeta).sqrt())
                };
                let c = 1.0 / (1.0 + t * t).sqrt();
                let sn = c * t;
                for r in 0..rows {
                    let wp = w[r * cols + p];
                    let wq = w[r * cols + q];
                    w[r * cols + p] = c * wp - sn * wq;
                    w[r * cols + q] = sn * wp + c * wq;
                }
                for r in 0..cols {
                    let vp = vv[r * cols + p];
                    let vq = vv[r * cols + q];
                    vv[r * cols + p] = c * vp - sn * vq;
                    vv[r * cols + q] = sn * vp + c * vq;
                }
            }
        }
        if !changed {
            break;
        }
    }
    let mut sigma = vec![0f64; cols];
    for c in 0..cols {
        let mut nrm = 0.0;
        for r in 0..rows {
            nrm += w[r * cols + c] * w[r * cols + c];
        }
        sigma[c] = nrm.sqrt();
    }
    let mut idx: Vec<usize> = (0..cols).collect();
    idx.sort_by(|&a, &b| {
        sigma[b]
            .partial_cmp(&sigma[a])
            .unwrap_or(std::cmp::Ordering::Equal)
    });
    let mut uw = vec![0f64; rows * cols]; // U in W-space (rows×cols)
    let mut vp = vec![0f64; cols * cols];
    let mut ss = vec![0f64; cols];
    for (newj, &oldj) in idx.iter().enumerate() {
        let sv = sigma[oldj];
        ss[newj] = sv;
        for r in 0..rows {
            uw[r * cols + newj] = if sv > 1e-300 {
                w[r * cols + oldj] / sv
            } else {
                0.0
            };
        }
        for r in 0..cols {
            vp[r * cols + newj] = vv[r * cols + oldj];
        }
    }
    let mut u_out = vec![0f64; m * k];
    let mut vt_out = vec![0f64; k * n];
    let s_out = ss[0..k].to_vec();
    if transposed {
        // W = Aᵀ = U_w·Σ·V_permᵀ ⇒ A = V_perm·Σ·U_wᵀ.
        for r in 0..m {
            for cc in 0..k {
                u_out[r * k + cc] = vp[r * cols + cc];
            }
        }
        for i in 0..k {
            for j in 0..n {
                vt_out[i * n + j] = uw[j * cols + i];
            }
        }
    } else {
        // W = A = U_w·Σ·V_permᵀ.
        for r in 0..m {
            for cc in 0..k {
                u_out[r * k + cc] = uw[r * cols + cc];
            }
        }
        for i in 0..k {
            for j in 0..n {
                vt_out[i * n + j] = vp[j * cols + i];
            }
        }
    }
    (u_out, s_out, vt_out)
}

/// Householder QR of a packed column-major `m×n` matrix (`lda = m`),
/// LAPACK `dgeqrf` layout: `R` in the upper triangle, reflector `v_j`
/// below the diagonal of column `j`, `tau[j]` its scalar.
#[cfg(not(rlx_cpu_blas))]
fn qr_householder_colmajor(a: &mut [f64], m: usize, n: usize, tau: &mut [f64]) {
    let at = |i: usize, j: usize| i + j * m;
    for j in 0..m.min(n) {
        let alpha = a[at(j, j)];
        let mut xn2 = 0.0;
        for i in (j + 1)..m {
            xn2 += a[at(i, j)] * a[at(i, j)];
        }
        if xn2 == 0.0 {
            tau[j] = 0.0;
            continue;
        }
        let norm = (alpha * alpha + xn2).sqrt();
        let beta = if alpha >= 0.0 { -norm } else { norm };
        let tj = (beta - alpha) / beta;
        let inv = 1.0 / (alpha - beta);
        for i in (j + 1)..m {
            a[at(i, j)] *= inv;
        }
        tau[j] = tj;
        a[at(j, j)] = beta; // R(j,j)
        for c in (j + 1)..n {
            let mut wsum = a[at(j, c)];
            for i in (j + 1)..m {
                wsum += a[at(i, j)] * a[at(i, c)];
            }
            let f = tj * wsum;
            a[at(j, c)] -= f;
            for i in (j + 1)..m {
                a[at(i, c)] -= f * a[at(i, j)];
            }
        }
    }
}

/// Form the `m×n` matrix `Q` (first `n` columns of `H_0···H_{k-1}`) in
/// place from the reflectors + `tau` left by [`qr_householder_colmajor`].
/// LAPACK `dorgqr` / `dorg2r`, column-major.
#[cfg(not(rlx_cpu_blas))]
fn form_q_colmajor(a: &mut [f64], m: usize, n: usize, k: usize, tau: &[f64]) {
    let at = |i: usize, j: usize| i + j * m;
    for j in k..n {
        for i in 0..m {
            a[at(i, j)] = 0.0;
        }
        if j < m {
            a[at(j, j)] = 1.0;
        }
    }
    for j in (0..k).rev() {
        let tj = tau[j];
        if j < n {
            a[at(j, j)] = 1.0;
            for c in (j + 1)..n {
                let mut wsum = a[at(j, c)];
                for i in (j + 1)..m {
                    wsum += a[at(i, j)] * a[at(i, c)];
                }
                let f = tj * wsum;
                a[at(j, c)] -= f;
                for i in (j + 1)..m {
                    a[at(i, c)] -= f * a[at(i, j)];
                }
            }
        }
        for i in (j + 1)..m {
            a[at(i, j)] *= -tj;
        }
        a[at(j, j)] = 1.0 - tj;
        for i in 0..j {
            a[at(i, j)] = 0.0;
        }
    }
}

#[cfg(not(rlx_cpu_blas))]
#[allow(non_snake_case, clippy::too_many_arguments)]
unsafe fn lapack_dpotrf(
    uplo: *const i8,
    n: *const i32,
    a: *mut f64,
    lda: *const i32,
    info: *mut i32,
) {
    let (uplo, n, lda) = unsafe { (*uplo as u8, *n as usize, *lda as usize) };
    let mut buf = vec![0f64; n * n];
    unsafe {
        for j in 0..n {
            for i in 0..n {
                buf[i + j * n] = *a.add(i + j * lda);
            }
        }
    }
    let rc = cholesky_colmajor(&mut buf, n, uplo == b'U');
    unsafe {
        *info = rc;
        if rc == 0 {
            for j in 0..n {
                for i in 0..n {
                    *a.add(i + j * lda) = buf[i + j * n];
                }
            }
        }
    }
}
#[cfg(not(rlx_cpu_blas))]
#[allow(non_snake_case, clippy::too_many_arguments)]
unsafe fn lapack_dgetrf(
    m: *const i32,
    n: *const i32,
    a: *mut f64,
    lda: *const i32,
    ipiv: *mut i32,
    info: *mut i32,
) {
    let (m, n, lda) = unsafe { (*m as usize, *n as usize, *lda as usize) };
    let mut buf = vec![0f64; m * n];
    unsafe {
        for j in 0..n {
            for i in 0..m {
                buf[i + j * m] = *a.add(i + j * lda);
            }
        }
    }
    let mut ip = vec![0i32; m.min(n)];
    let rc = lu_colmajor(&mut buf, m, n, &mut ip);
    unsafe {
        for j in 0..n {
            for i in 0..m {
                *a.add(i + j * lda) = buf[i + j * m];
            }
        }
        for (i, &p) in ip.iter().enumerate() {
            *ipiv.add(i) = p;
        }
        *info = rc;
    }
}
#[cfg(not(rlx_cpu_blas))]
#[allow(non_snake_case, clippy::too_many_arguments)]
unsafe fn lapack_dsyevd(
    jobz: *const i8,
    uplo: *const i8,
    n: *const i32,
    a: *mut f64,
    lda: *const i32,
    w: *mut f64,
    work: *mut f64,
    lwork: *const i32,
    iwork: *mut i32,
    _liwork: *const i32,
    info: *mut i32,
) {
    let (jobz, uplo, n, lda, lwork) =
        unsafe { (*jobz as u8, *uplo as u8, *n as usize, *lda as usize, *lwork) };
    if lwork == -1 {
        unsafe {
            *work = 1.0;
            *iwork = 1;
            *info = 0;
        }
        return;
    }
    let mut full = vec![0f64; n * n];
    unsafe {
        for j in 0..n {
            for i in 0..n {
                let v = if uplo == b'U' {
                    if i <= j {
                        *a.add(i + j * lda)
                    } else {
                        *a.add(j + i * lda)
                    }
                } else if i >= j {
                    *a.add(i + j * lda)
                } else {
                    *a.add(j + i * lda)
                };
                full[i * n + j] = v;
            }
        }
    }
    let (evals, evecs) = sym_eig_jacobi(full, n);
    unsafe {
        for (i, &e) in evals.iter().enumerate() {
            *w.add(i) = e;
        }
        if jobz == b'V' {
            for j in 0..n {
                for i in 0..n {
                    *a.add(i + j * lda) = evecs[i * n + j];
                }
            }
        }
        *info = 0;
    }
}
#[cfg(not(rlx_cpu_blas))]
#[allow(non_snake_case, clippy::too_many_arguments)]
unsafe fn lapack_dgeqrf(
    m: *const i32,
    n: *const i32,
    a: *mut f64,
    lda: *const i32,
    tau: *mut f64,
    work: *mut f64,
    lwork: *const i32,
    info: *mut i32,
) {
    let (m, n, lda, lwork) = unsafe { (*m as usize, *n as usize, *lda as usize, *lwork) };
    if lwork == -1 {
        unsafe {
            *work = n.max(1) as f64;
            *info = 0;
        }
        return;
    }
    let mut buf = vec![0f64; m * n];
    unsafe {
        for j in 0..n {
            for i in 0..m {
                buf[i + j * m] = *a.add(i + j * lda);
            }
        }
    }
    let k = m.min(n);
    let mut tv = vec![0f64; k];
    qr_householder_colmajor(&mut buf, m, n, &mut tv);
    unsafe {
        for j in 0..n {
            for i in 0..m {
                *a.add(i + j * lda) = buf[i + j * m];
            }
        }
        for (i, &t) in tv.iter().enumerate() {
            *tau.add(i) = t;
        }
        *info = 0;
    }
}
#[cfg(not(rlx_cpu_blas))]
#[allow(non_snake_case, clippy::too_many_arguments)]
unsafe fn lapack_dorgqr(
    m: *const i32,
    n: *const i32,
    k: *const i32,
    a: *mut f64,
    lda: *const i32,
    tau: *const f64,
    work: *mut f64,
    lwork: *const i32,
    info: *mut i32,
) {
    let (m, n, k, lda, lwork) =
        unsafe { (*m as usize, *n as usize, *k as usize, *lda as usize, *lwork) };
    if lwork == -1 {
        unsafe {
            *work = n.max(1) as f64;
            *info = 0;
        }
        return;
    }
    let mut buf = vec![0f64; m * n];
    let mut tv = vec![0f64; k];
    unsafe {
        for j in 0..n {
            for i in 0..m {
                buf[i + j * m] = *a.add(i + j * lda);
            }
        }
        for (i, t) in tv.iter_mut().enumerate() {
            *t = *tau.add(i);
        }
    }
    form_q_colmajor(&mut buf, m, n, k, &tv);
    unsafe {
        for j in 0..n {
            for i in 0..m {
                *a.add(i + j * lda) = buf[i + j * m];
            }
        }
        *info = 0;
    }
}
#[cfg(not(rlx_cpu_blas))]
#[allow(non_snake_case, clippy::too_many_arguments)]
unsafe fn lapack_dgesvd(
    _jobu: *const i8,
    _jobvt: *const i8,
    m: *const i32,
    n: *const i32,
    a: *mut f64,
    lda: *const i32,
    s: *mut f64,
    u: *mut f64,
    ldu: *const i32,
    vt: *mut f64,
    ldvt: *const i32,
    work: *mut f64,
    lwork: *const i32,
    info: *mut i32,
) {
    // The wrappers always request thin `S`/`S` mode, which is what this
    // computes; `_jobu`/`_jobvt` are accepted for ABI parity.
    let (m, n, lda, ldu, ldvt, lwork) = unsafe {
        (
            *m as usize,
            *n as usize,
            *lda as usize,
            *ldu as usize,
            *ldvt as usize,
            *lwork,
        )
    };
    if lwork == -1 {
        unsafe {
            *work = 1.0;
            *info = 0;
        }
        return;
    }
    let k = m.min(n);
    let mut am = vec![0f64; m * n];
    unsafe {
        for i in 0..m {
            for j in 0..n {
                am[i * n + j] = *a.add(i + j * lda);
            }
        }
    }
    let (u_rm, s_v, vt_rm) = svd_thin_rowmajor(&am, m, n);
    unsafe {
        for (i, &sv) in s_v.iter().enumerate() {
            *s.add(i) = sv;
        }
        for j in 0..k {
            for i in 0..m {
                *u.add(i + j * ldu) = u_rm[i * k + j];
            }
        }
        for j in 0..n {
            for i in 0..k {
                *vt.add(i + j * ldvt) = vt_rm[i * n + j];
            }
        }
        *info = 0;
    }
}
#[cfg(not(rlx_cpu_blas))]
#[allow(non_snake_case, clippy::too_many_arguments)]
unsafe fn lapack_dgesdd(
    _jobz: *const i8,
    m: *const i32,
    n: *const i32,
    a: *mut f64,
    lda: *const i32,
    s: *mut f64,
    u: *mut f64,
    ldu: *const i32,
    vt: *mut f64,
    ldvt: *const i32,
    work: *mut f64,
    lwork: *const i32,
    _iwork: *mut i32,
    info: *mut i32,
) {
    // Divide-and-conquer SVD driver; the fallback shares the one-sided
    // Jacobi thin-SVD used by `dgesvd` (same I/O contract, thin mode).
    let (m, n, lda, ldu, ldvt, lwork) = unsafe {
        (
            *m as usize,
            *n as usize,
            *lda as usize,
            *ldu as usize,
            *ldvt as usize,
            *lwork,
        )
    };
    if lwork == -1 {
        unsafe {
            *work = 1.0;
            *info = 0;
        }
        return;
    }
    let k = m.min(n);
    let mut am = vec![0f64; m * n];
    unsafe {
        for i in 0..m {
            for j in 0..n {
                am[i * n + j] = *a.add(i + j * lda);
            }
        }
    }
    let (u_rm, s_v, vt_rm) = svd_thin_rowmajor(&am, m, n);
    unsafe {
        for (i, &sv) in s_v.iter().enumerate() {
            *s.add(i) = sv;
        }
        for j in 0..k {
            for i in 0..m {
                *u.add(i + j * ldu) = u_rm[i * k + j];
            }
        }
        for j in 0..n {
            for i in 0..k {
                *vt.add(i + j * ldvt) = vt_rm[i * n + j];
            }
        }
        *info = 0;
    }
}
#[cfg(not(rlx_cpu_blas))]
#[allow(non_snake_case, clippy::too_many_arguments)]
unsafe fn lapack_dgelsd(
    m: *const i32,
    n: *const i32,
    nrhs: *const i32,
    a: *mut f64,
    lda: *const i32,
    b: *mut f64,
    ldb: *const i32,
    s: *mut f64,
    rcond: *const f64,
    rank: *mut i32,
    work: *mut f64,
    lwork: *const i32,
    _iwork: *mut i32,
    info: *mut i32,
) {
    // Minimum-norm least squares via the thin SVD: X = V·Σ⁺·Uᵀ·B, with
    // singular values ≤ rcond·σ_max (machine eps when rcond < 0) dropped.
    let (m, n, nrhs, lda, ldb, lwork, rcond) = unsafe {
        (
            *m as usize,
            *n as usize,
            *nrhs as usize,
            *lda as usize,
            *ldb as usize,
            *lwork,
            *rcond,
        )
    };
    if lwork == -1 {
        unsafe {
            *work = 1.0;
            *info = 0;
        }
        return;
    }
    let k = m.min(n);
    let mut am = vec![0f64; m * n];
    unsafe {
        for i in 0..m {
            for j in 0..n {
                am[i * n + j] = *a.add(i + j * lda);
            }
        }
    }
    let (u_rm, s_v, vt_rm) = svd_thin_rowmajor(&am, m, n);
    let smax = if k > 0 { s_v[0] } else { 0.0 };
    let eff = if rcond < 0.0 { f64::EPSILON } else { rcond };
    let thresh = eff * smax;
    let mut rk = 0i32;
    for &sv in &s_v {
        if sv > thresh {
            rk += 1;
        }
    }
    // Solution written row-major first, then marshalled to B (col-major).
    let mut xout = vec![0f64; n * nrhs];
    unsafe {
        for c in 0..nrhs {
            let mut z = vec![0f64; k];
            for i in 0..k {
                if s_v[i] > thresh {
                    let mut y = 0.0;
                    for l in 0..m {
                        y += u_rm[l * k + i] * *b.add(l + c * ldb);
                    }
                    z[i] = y / s_v[i];
                }
            }
            for r in 0..n {
                let mut xr = 0.0;
                for i in 0..k {
                    xr += vt_rm[i * n + r] * z[i];
                }
                xout[r * nrhs + c] = xr;
            }
        }
        for c in 0..nrhs {
            for r in 0..n {
                *b.add(r + c * ldb) = xout[r * nrhs + c];
            }
        }
        for (i, &sv) in s_v.iter().enumerate() {
            *s.add(i) = sv;
        }
        *rank = rk;
        *info = 0;
    }
}
#[cfg(not(rlx_cpu_blas))]
#[allow(non_snake_case, clippy::too_many_arguments)]
unsafe fn cblas_dtrsm(
    order: i32,
    side: i32,
    uplo: i32,
    transa: i32,
    diag: i32,
    m: i32,
    n: i32,
    alpha: f64,
    a: *const f64,
    lda: i32,
    b: *mut f64,
    ldb: i32,
) {
    // Solve op(A)·X = alpha·B in place. Covers the sole in-tree caller
    // (`dtrsm_lower_or_upper`): row-major, left side, both triangles /
    // transpose / unit-diagonal. Col-major or right-side are never used.
    debug_assert!(
        order == ROW_MAJOR && side == CBLAS_LEFT,
        "rlx-cpu dtrsm fallback: only row-major left-side is exercised"
    );
    let _ = (order, side);
    let m = m as usize;
    let n = n as usize;
    let lda = lda as usize;
    let ldb = ldb as usize;
    let unit = diag == CBLAS_UNIT;
    let upper = uplo == CBLAS_UPPER;
    let trans = transa != NO_TRANS;
    let op_lower = upper == trans; // op(A) lower ⇔ (A upper) XOR-not (transpose)
    let aat = |i: usize, l: usize| unsafe {
        if trans {
            *a.add(l * lda + i)
        } else {
            *a.add(i * lda + l)
        }
    };
    unsafe {
        if alpha != 1.0 {
            for i in 0..m {
                for j in 0..n {
                    *b.add(i * ldb + j) *= alpha;
                }
            }
        }
        if op_lower {
            for i in 0..m {
                for j in 0..n {
                    let mut val = *b.add(i * ldb + j);
                    for l in 0..i {
                        val -= aat(i, l) * *b.add(l * ldb + j);
                    }
                    if !unit {
                        val /= aat(i, i);
                    }
                    *b.add(i * ldb + j) = val;
                }
            }
        } else {
            for i in (0..m).rev() {
                for j in 0..n {
                    let mut val = *b.add(i * ldb + j);
                    for l in (i + 1)..m {
                        val -= aat(i, l) * *b.add(l * ldb + j);
                    }
                    if !unit {
                        val /= aat(i, i);
                    }
                    *b.add(i * ldb + j) = val;
                }
            }
        }
    }
}

// ── Linalg row-major wrappers ────────────────────────────────────

/// In-place Cholesky factorization of an SPD matrix `A`. The caller
/// passes `a` row-major, n×n. On output:
///   - if `lower`: lower triangle of `a` holds `L` such that `L·Lᵀ = A`,
///     upper triangle is zeroed.
///   - if `!lower`: upper triangle holds `U` such that `Uᵀ·U = A`,
///     lower triangle is zeroed.
/// Returns 0 on success; k>0 means the leading minor of order k is
/// not positive-definite (A is not SPD).
///
/// LAPACK's `dpotrf_` is column-major. For symmetric A, the col-major
/// bytes equal the row-major bytes (`A = Aᵀ`). After dpotrf with
/// UPLO='U' (col-major upper), the factor U lives in col-major's
/// upper triangle — which is row-major's lower triangle = L. Same
/// trick in reverse for the upper-row-major case.
pub fn dpotrf(a: &mut [f64], n: usize, lower: bool) -> i32 {
    assert_eq!(a.len(), n * n, "dpotrf: A must be n×n");
    let uplo: i8 = if lower { b'U' as i8 } else { b'L' as i8 };
    let nn = n as i32;
    let mut info: i32 = 0;
    unsafe {
        lapack_dpotrf(&uplo, &nn, a.as_mut_ptr(), &nn, &mut info);
    }
    if info != 0 {
        return info;
    }
    // Mask the unused triangle in the row-major view.
    if lower {
        for i in 0..n {
            for j in (i + 1)..n {
                a[i * n + j] = 0.0;
            }
        }
    } else {
        for i in 1..n {
            for j in 0..i {
                a[i * n + j] = 0.0;
            }
        }
    }
    info
}

/// LU-factorize a row-major n×n matrix and return `(log|det|, sign, det)`.
/// Row-major bytes are col-major `Aᵀ`, and `det(Aᵀ) = det(A)`, so this is exact.
/// Overwrites `a` with the LU factors. Singular ⇒ `(-inf, 0, 0)`.
pub fn lu_slogdet(a: &mut [f64], n: usize) -> (f64, f64, f64) {
    assert_eq!(a.len(), n * n, "lu_slogdet: A must be n×n");
    let mut ipiv = vec![0i32; n];
    let mut info: i32 = 0;
    let nn = n as i32;
    unsafe {
        lapack_dgetrf(&nn, &nn, a.as_mut_ptr(), &nn, ipiv.as_mut_ptr(), &mut info);
    }
    if info > 0 {
        return (f64::NEG_INFINITY, 0.0, 0.0); // U has a zero pivot → singular.
    }
    let (mut logabs, mut sign, mut det) = (0.0f64, 1.0f64, 1.0f64);
    for i in 0..n {
        let d = a[i * n + i]; // U_ii (col-major diagonal = row-major diagonal).
        det *= d;
        logabs += d.abs().ln();
        if d < 0.0 {
            sign = -sign;
        }
        if ipiv[i] != (i as i32 + 1) {
            // 1-based pivot; a row swap flips the sign.
            sign = -sign;
            det = -det;
        }
    }
    (logabs, sign, det)
}

/// In-place symmetric eigendecomposition. `a` is row-major n×n
/// symmetric. On output: `w` holds eigenvalues (length n, ascending);
/// `a` is overwritten with eigenvectors as columns (col-major) =
/// rows (row-major view, since the matrix's transpose interpretation
/// is the same set of orthonormal eigenvectors). Symmetric → no
/// transpose dance needed.
///
/// Returns 0 on success; k>0 means the algorithm failed to converge.
pub fn dsyevd(a: &mut [f64], w: &mut [f64], n: usize) -> i32 {
    assert_eq!(a.len(), n * n);
    assert_eq!(w.len(), n);
    let jobz: i8 = b'V' as i8;
    let uplo: i8 = b'U' as i8;
    let nn = n as i32;
    let mut info: i32 = 0;
    // dsyevd workspace (per LAPACK manual):
    //   lwork  ≥ 1 + 6n + 2n²       (jobz='V')
    //   liwork ≥ 3 + 5n             (jobz='V')
    let lwork = (1 + 6 * n + 2 * n * n) as i32;
    let liwork = (3 + 5 * n) as i32;
    let mut work = vec![0f64; lwork.max(1) as usize];
    let mut iwork = vec![0i32; liwork.max(1) as usize];
    unsafe {
        lapack_dsyevd(
            &jobz,
            &uplo,
            &nn,
            a.as_mut_ptr(),
            &nn,
            w.as_mut_ptr(),
            work.as_mut_ptr(),
            &lwork,
            iwork.as_mut_ptr(),
            &liwork,
            &mut info,
        );
    }
    info
}

/// QR factorization. Inputs:
///   - `a`: m×n row-major matrix (overwritten by the factorization)
///   - `q_out`: m×k row-major, k = min(m, n) — receives Q
///   - `r_out`: k×n row-major — receives R
/// Returns 0 on success; <0 means LAPACK got a bad arg.
pub fn dgeqrf_full(a: &mut [f64], m: usize, n: usize, q_out: &mut [f64], r_out: &mut [f64]) -> i32 {
    assert_eq!(a.len(), m * n, "dgeqrf: A must be m×n");
    let k = m.min(n);
    assert_eq!(q_out.len(), m * k, "Q must be m×min(m,n)");
    assert_eq!(r_out.len(), k * n, "R must be min(m,n)×n");

    // Row→col-major: transpose A in place.
    let mut a_col = transpose_to_col(a, m, n);
    let mut tau = vec![0f64; k];
    let mm = m as i32;
    let nn = n as i32;
    let kk = k as i32;
    let lwork = (n.max(1)) as i32;
    let mut work = vec![0f64; lwork.max(1) as usize];
    let mut info: i32 = 0;
    unsafe {
        lapack_dgeqrf(
            &mm,
            &nn,
            a_col.as_mut_ptr(),
            &mm,
            tau.as_mut_ptr(),
            work.as_mut_ptr(),
            &lwork,
            &mut info,
        );
    }
    if info != 0 {
        return info;
    }

    // Extract R from the upper triangle of the (k×n) leading block of
    // a_col (col-major). Row-major R has shape [k, n]; entry (i,j)
    // = a_col[j * m + i] for i ≤ j, else 0.
    for i in 0..k {
        for j in 0..n {
            let v = if i <= j { a_col[j * m + i] } else { 0.0 };
            r_out[i * n + j] = v;
        }
    }

    // Form Q in-place (col-major, m×k) using dorgqr.
    let mut work2 = vec![0f64; lwork.max(1) as usize];
    let mut info2: i32 = 0;
    unsafe {
        lapack_dorgqr(
            &mm,
            &kk,
            &kk,
            a_col.as_mut_ptr(),
            &mm,
            tau.as_ptr(),
            work2.as_mut_ptr(),
            &lwork,
            &mut info2,
        );
    }
    if info2 != 0 {
        return info2;
    }

    // Transpose Q (col-major m×k) → row-major m×k.
    for i in 0..m {
        for j in 0..k {
            q_out[i * k + j] = a_col[j * m + i];
        }
    }
    0
}

/// SVD: `a = U · diag(s) · V^T`. Inputs:
///   - `a`: m×n row-major (destroyed on return)
///   - `s`: length min(m, n) — singular values, descending
///   - `u`: m×min(m,n) row-major
///   - `vt`: min(m,n)×n row-major (V transposed)
/// Returns 0 on success; >0 means superdiagonals failed to converge.
pub fn dgesvd_thin(
    a: &mut [f64],
    m: usize,
    n: usize,
    s: &mut [f64],
    u: &mut [f64],
    vt: &mut [f64],
) -> i32 {
    assert_eq!(a.len(), m * n);
    let k = m.min(n);
    assert_eq!(s.len(), k);
    assert_eq!(u.len(), m * k);
    assert_eq!(vt.len(), k * n);

    let mut a_col = transpose_to_col(a, m, n);
    let mut u_col = vec![0f64; m * k];
    let mut vt_col = vec![0f64; k * n];
    let jobu = b'S' as i8;
    let jobvt = b'S' as i8;
    let mm = m as i32;
    let nn = n as i32;
    let ldu = m as i32;
    let ldvt = k as i32;
    // dgesvd workspace: lwork ≥ max(3*min(m,n) + max(m,n), 5*min(m,n))
    let lwork = (((3 * k + m.max(n)).max(5 * k)) as i32).max(1);
    let mut work = vec![0f64; lwork as usize];
    let mut info: i32 = 0;
    unsafe {
        lapack_dgesvd(
            &jobu,
            &jobvt,
            &mm,
            &nn,
            a_col.as_mut_ptr(),
            &mm,
            s.as_mut_ptr(),
            u_col.as_mut_ptr(),
            &ldu,
            vt_col.as_mut_ptr(),
            &ldvt,
            work.as_mut_ptr(),
            &lwork,
            &mut info,
        );
    }
    if info != 0 {
        return info;
    }
    // Transpose U (col-major m×k) → row-major m×k.
    for i in 0..m {
        for j in 0..k {
            u[i * k + j] = u_col[j * m + i];
        }
    }
    // Transpose V^T (col-major k×n) → row-major k×n.
    for i in 0..k {
        for j in 0..n {
            vt[i * n + j] = vt_col[j * k + i];
        }
    }
    0
}

/// Thin SVD via LAPACK **`dgesdd`** (divide-and-conquer) — the driver NumPy's
/// `linalg.svd`/`linalg.pinv` use. Same I/O contract as [`dgesvd_thin`]; use
/// this when bit-exact NumPy parity on ill-conditioned inputs is required.
/// Thin QR of a row-major `m×n` matrix: `A = Q·R`, `Q` `[m,k]` orthonormal,
/// `R` `[k,n]` upper-triangular (`k = min(m,n)`). `a` is consumed (converted to
/// col-major internally). Returns LAPACK `info` (0 = ok).
pub fn qr_thin(a: &[f64], m: usize, n: usize, q: &mut [f64], r: &mut [f64]) -> i32 {
    let k = m.min(n);
    assert_eq!(a.len(), m * n);
    assert_eq!(q.len(), m * k);
    assert_eq!(r.len(), k * n);
    let mut a_col = transpose_to_col(a, m, n); // col-major [m, n]
    let (mm, nn, kk) = (m as i32, n as i32, k as i32);
    let mut tau = vec![0f64; k];
    let mut info: i32 = 0;
    let m1: i32 = -1;
    let mut wq = [0f64; 1];
    unsafe {
        lapack_dgeqrf(
            &mm,
            &nn,
            a_col.as_mut_ptr(),
            &mm,
            tau.as_mut_ptr(),
            wq.as_mut_ptr(),
            &m1,
            &mut info,
        );
    }
    let lwork = (wq[0] as i32).max(1);
    let mut work = vec![0f64; lwork as usize];
    unsafe {
        lapack_dgeqrf(
            &mm,
            &nn,
            a_col.as_mut_ptr(),
            &mm,
            tau.as_mut_ptr(),
            work.as_mut_ptr(),
            &lwork,
            &mut info,
        );
    }
    if info != 0 {
        return info;
    }
    // R = upper triangle (k×n) of the factored matrix → row-major.
    for i in 0..k {
        for j in 0..n {
            r[i * n + j] = if i <= j { a_col[i + j * m] } else { 0.0 };
        }
    }
    // Form Q [m,k] in a_col (col-major, first k columns).
    let mut wq2 = [0f64; 1];
    unsafe {
        lapack_dorgqr(
            &mm,
            &kk,
            &kk,
            a_col.as_mut_ptr(),
            &mm,
            tau.as_ptr(),
            wq2.as_mut_ptr(),
            &m1,
            &mut info,
        );
    }
    let lwork2 = (wq2[0] as i32).max(1);
    let mut work2 = vec![0f64; lwork2 as usize];
    unsafe {
        lapack_dorgqr(
            &mm,
            &kk,
            &kk,
            a_col.as_mut_ptr(),
            &mm,
            tau.as_ptr(),
            work2.as_mut_ptr(),
            &lwork2,
            &mut info,
        );
    }
    if info != 0 {
        return info;
    }
    for i in 0..m {
        for j in 0..k {
            q[i * k + j] = a_col[i + j * m];
        }
    }
    info
}

pub fn dgesdd_thin(
    a: &mut [f64],
    m: usize,
    n: usize,
    s: &mut [f64],
    u: &mut [f64],
    vt: &mut [f64],
) -> i32 {
    assert_eq!(a.len(), m * n);
    let k = m.min(n);
    assert_eq!(s.len(), k);
    assert_eq!(u.len(), m * k);
    assert_eq!(vt.len(), k * n);

    let mut a_col = transpose_to_col(a, m, n);
    let mut u_col = vec![0f64; m * k];
    let mut vt_col = vec![0f64; k * n];
    let jobz = b'S' as i8;
    let mm = m as i32;
    let nn = n as i32;
    let ldu = m as i32;
    let ldvt = k as i32;
    let mut iwork = vec![0i32; 8 * k];
    let mut info: i32 = 0;
    // workspace query
    let mut wq = [0f64; 1];
    let m1: i32 = -1;
    unsafe {
        lapack_dgesdd(
            &jobz,
            &mm,
            &nn,
            a_col.as_mut_ptr(),
            &mm,
            s.as_mut_ptr(),
            u_col.as_mut_ptr(),
            &ldu,
            vt_col.as_mut_ptr(),
            &ldvt,
            wq.as_mut_ptr(),
            &m1,
            iwork.as_mut_ptr(),
            &mut info,
        );
    }
    if info != 0 {
        return info;
    }
    let lwork = wq[0] as i32;
    let mut work = vec![0f64; lwork.max(1) as usize];
    unsafe {
        lapack_dgesdd(
            &jobz,
            &mm,
            &nn,
            a_col.as_mut_ptr(),
            &mm,
            s.as_mut_ptr(),
            u_col.as_mut_ptr(),
            &ldu,
            vt_col.as_mut_ptr(),
            &ldvt,
            work.as_mut_ptr(),
            &lwork,
            iwork.as_mut_ptr(),
            &mut info,
        );
    }
    if info != 0 {
        return info;
    }
    for i in 0..m {
        for j in 0..k {
            u[i * k + j] = u_col[j * m + i];
        }
    }
    for i in 0..k {
        for j in 0..n {
            vt[i * n + j] = vt_col[j * k + i];
        }
    }
    0
}

/// Minimum-norm least-squares `min‖A·X − B‖` via LAPACK **`dgelsd`** (the driver
/// SciPy's `linalg.lstsq` uses). `a` is m×n row-major, `b` is m×nrhs row-major;
/// the solution `X` (n×nrhs) is written row-major into `x`. `rcond < 0` uses
/// machine precision (SciPy's default). Returns 0 on success.
pub fn dgelsd_solve(
    a: &[f64],
    b: &[f64],
    m: usize,
    n: usize,
    nrhs: usize,
    rcond: f64,
    x: &mut [f64],
) -> i32 {
    assert_eq!(a.len(), m * n);
    assert_eq!(b.len(), m * nrhs);
    assert_eq!(x.len(), n * nrhs);
    let k = m.min(n);
    let ldb = m.max(n);
    let mut a_col = transpose_to_col(a, m, n);
    // B column-major with leading dim ldb=max(m,n); first m rows filled.
    let mut b_col = vec![0f64; ldb * nrhs];
    for i in 0..m {
        for j in 0..nrhs {
            b_col[i + j * ldb] = b[i * nrhs + j];
        }
    }
    let mm = m as i32;
    let nn = n as i32;
    let nrhs_i = nrhs as i32;
    let ldb_i = ldb as i32;
    let mut s = vec![0f64; k];
    let rc = rcond;
    let mut rank: i32 = 0;
    let mut info: i32 = 0;
    let nlvl = (k as f64).log2().floor() as usize + 1;
    let mut iwork = vec![0i32; (3 * k * nlvl + 11 * k).max(1)];
    let mut wq = [0f64; 1];
    let m1: i32 = -1;
    unsafe {
        lapack_dgelsd(
            &mm,
            &nn,
            &nrhs_i,
            a_col.as_mut_ptr(),
            &mm,
            b_col.as_mut_ptr(),
            &ldb_i,
            s.as_mut_ptr(),
            &rc,
            &mut rank,
            wq.as_mut_ptr(),
            &m1,
            iwork.as_mut_ptr(),
            &mut info,
        );
    }
    if info != 0 {
        return info;
    }
    let lwork = wq[0] as i32;
    let mut work = vec![0f64; lwork.max(1) as usize];
    unsafe {
        lapack_dgelsd(
            &mm,
            &nn,
            &nrhs_i,
            a_col.as_mut_ptr(),
            &mm,
            b_col.as_mut_ptr(),
            &ldb_i,
            s.as_mut_ptr(),
            &rc,
            &mut rank,
            work.as_mut_ptr(),
            &lwork,
            iwork.as_mut_ptr(),
            &mut info,
        );
    }
    if info != 0 {
        return info;
    }
    // Solution X is in the first n rows of b_col (col-major, ld=ldb) → row-major.
    for i in 0..n {
        for j in 0..nrhs {
            x[i * nrhs + j] = b_col[i + j * ldb];
        }
    }
    0
}

/// Solve `op(A) · X = B` (or `X · op(A) = B` for `right`) where A is
/// triangular. Row-major callers throughout. Overwrites B with X.
///
///   - `a`: n×n row-major triangular
///   - `b`: m×n_rhs row-major (B = X · A) or n_rhs×n (A · X = B)?
///
/// Concretely: solve `A · X = B`:
///   - `a`: n×n row-major triangular (lower or upper per `lower`)
///   - `b`: n×nrhs row-major. Overwritten with X (n×nrhs).
pub fn dtrsm_lower_or_upper(
    a: &[f64],
    b: &mut [f64],
    n: usize,
    nrhs: usize,
    lower: bool,
    transpose_a: bool,
) {
    assert_eq!(a.len(), n * n);
    assert_eq!(b.len(), n * nrhs);
    unsafe {
        cblas_dtrsm(
            ROW_MAJOR,
            CBLAS_LEFT,
            if lower { CBLAS_LOWER } else { CBLAS_UPPER },
            if transpose_a { TRANS } else { NO_TRANS },
            CBLAS_NON_UNIT,
            n as i32,
            nrhs as i32,
            1.0,
            a.as_ptr(),
            n as i32,
            b.as_mut_ptr(),
            nrhs as i32,
        );
    }
}

fn transpose_to_col(a_row: &[f64], m: usize, n: usize) -> Vec<f64> {
    let mut out = vec![0f64; m * n];
    for i in 0..m {
        for j in 0..n {
            out[j * m + i] = a_row[i * n + j];
        }
    }
    out
}

/// C = alpha * A @ B + beta * C
///
/// A: [m, k] row-major
/// Opt-in high-precision matmul accumulation (`RLX_CPU_MATMUL_F64_ACCUM=1`).
///
/// The default f32 GEMM accumulates the K-reduction in f32; for a long K (or a
/// dequantized low-precision operand fed into the matmul) that rounding is the
/// dominant avoidable error. When this flag is set, the affected GEMM entries
/// promote to f64, accumulate via vendor `dgemm`, and narrow back — turning the
/// K-reduction into an f64 sum. This is a *precision* mode (validated/parity
/// runs), NOT a perf mode: it is DEFAULT-OFF so the vendor-BLAS fast path — the
/// per-hardware peak — is completely untouched unless a user opts in. Cached
/// once (set at process start, as for a validated pipeline).
#[inline]
pub(crate) fn f64_accum_enabled() -> bool {
    use std::sync::OnceLock;
    static ON: OnceLock<bool> = OnceLock::new();
    *ON.get_or_init(|| rlx_ir::env_registry::flag("RLX_CPU_MATMUL_F64_ACCUM"))
}

/// `C = A@B` (or `C += A@B` when `accumulate`) with the K-reduction done in f64.
/// Promote → vendor `dgemm` → narrow; the accumulate add is also f64 so an
/// existing f32 `C` is folded in without a second rounding of the running sum.
fn dgemm_f32_precise(
    a: &[f32],
    b: &[f32],
    c: &mut [f32],
    m: usize,
    k: usize,
    n: usize,
    accumulate: bool,
) {
    if m == 0 || n == 0 {
        return;
    }
    let a64: Vec<f64> = a.iter().map(|&x| x as f64).collect();
    let b64: Vec<f64> = b.iter().map(|&x| x as f64).collect();
    let mut c64 = vec![0f64; m * n];
    dgemm(&a64, &b64, &mut c64, m, k, n);
    if accumulate {
        for (dst, &v) in c.iter_mut().zip(c64.iter()) {
            *dst = (*dst as f64 + v) as f32;
        }
    } else {
        for (dst, &v) in c.iter_mut().zip(c64.iter()) {
            *dst = v as f32;
        }
    }
}

/// Report a GEMM whose destination overlaps an operand, when
/// `RLX_BLAS_ALIAS_CHECK=1`.
///
/// Every BLAS call site is supposed to detour through a scratch buffer when an
/// operand lands inside the destination — arena slot reuse makes that routine,
/// not exotic. A site that forgets streams an operand it is simultaneously
/// overwriting, which corrupts whichever rows the kernel reached first: a
/// nondeterministic band of output that `RLX_ARENA_NO_REUSE=1` makes disappear
/// and no amount of staring at the graph explains. Auditing call sites by hand
/// does not scale, so this catches them centrally and names the shape.
#[inline]
pub(crate) fn alias_check(
    site: &str,
    a: &[f32],
    b: &[f32],
    c: &[f32],
    m: usize,
    k: usize,
    n: usize,
) {
    if !rlx_ir::env::flag("RLX_BLAS_ALIAS_CHECK") {
        return;
    }
    let span = |x: &[f32]| {
        let p = x.as_ptr() as usize;
        (p, p + std::mem::size_of_val(x))
    };
    let (c0, c1) = span(c);
    for (name, operand) in [("a", a), ("b", b)] {
        let (o0, o1) = span(operand);
        if o0 < c1 && c0 < o1 {
            eprintln!(
                "[blas-alias] {site}: operand {name} [{o0:#x},{o1:#x}) overlaps dest \
                 [{c0:#x},{c1:#x}) — m={m} k={k} n={n}"
            );
        }
    }
}

/// B: [k, n] row-major
/// C: [m, n] row-major
#[inline]
pub fn sgemm(a: &[f32], b: &[f32], c: &mut [f32], m: usize, k: usize, n: usize) {
    alias_check("sgemm", a, b, c, m, k, n);
    if f64_accum_enabled() {
        dgemm_f32_precise(a, b, c, m, k, n, false);
        return;
    }
    unsafe {
        cblas_sgemm(
            ROW_MAJOR,
            NO_TRANS,
            NO_TRANS,
            m as i32,
            n as i32,
            k as i32,
            1.0,
            a.as_ptr(),
            k as i32,
            b.as_ptr(),
            n as i32,
            0.0,
            c.as_mut_ptr(),
            n as i32,
        );
    }
}

/// C += alpha * A @ B (accumulate into existing C)
#[inline]
pub fn sgemm_accumulate(a: &[f32], b: &[f32], c: &mut [f32], m: usize, k: usize, n: usize) {
    alias_check("sgemm_accumulate", a, b, c, m, k, n);
    if f64_accum_enabled() {
        dgemm_f32_precise(a, b, c, m, k, n, true);
        return;
    }
    unsafe {
        cblas_sgemm(
            ROW_MAJOR,
            NO_TRANS,
            NO_TRANS,
            m as i32,
            n as i32,
            k as i32,
            1.0,
            a.as_ptr(),
            k as i32,
            b.as_ptr(),
            n as i32,
            1.0,
            c.as_mut_ptr(),
            n as i32,
        );
    }
}

/// Like [`sgemm_accumulate`] with the same Rayon-split dispatch as [`sgemm_auto`].
#[inline]
pub fn sgemm_accumulate_auto(a: &[f32], b: &[f32], c: &mut [f32], m: usize, k: usize, n: usize) {
    alias_check("sgemm_accumulate_auto", a, b, c, m, k, n);
    // Precision mode wins over the parallel dispatch: the f64-accum path is
    // serial (vendor `dgemm`), so route straight to it and skip `par_sgemm`.
    if f64_accum_enabled() {
        sgemm_accumulate(a, b, c, m, k, n);
        return;
    }
    // Defined below with `par_sgemm`; call via the public name once both are
    // in scope (same module — order-independent). Prefer parallel when large.
    if prefer_par_sgemm(m, k, n) {
        crate::blas::par_sgemm_accumulate(a, b, c, m, k, n);
    } else {
        sgemm_accumulate(a, b, c, m, k, n);
    }
}

/// C = A @ B^T (B transposed — useful for SDPA Q@K^T)
#[inline]
pub fn sgemm_bt(a: &[f32], b: &[f32], c: &mut [f32], m: usize, k: usize, n: usize, alpha: f32) {
    alias_check("sgemm_bt", a, b, c, m, k, n);
    unsafe {
        cblas_sgemm(
            ROW_MAJOR,
            NO_TRANS,
            TRANS,
            m as i32,
            n as i32,
            k as i32,
            alpha,
            a.as_ptr(),
            k as i32,
            b.as_ptr(),
            k as i32, // B is [n, k], transposed to [k, n]
            0.0,
            c.as_mut_ptr(),
            n as i32,
        );
    }
}

/// `y = alpha * A @ x + beta * y` with `A` row-major `[n, k]` (`lda = k`).
#[inline]
pub fn sgemv_nn(a: &[f32], x: &[f32], y: &mut [f32], n: usize, k: usize, alpha: f32, beta: f32) {
    unsafe {
        cblas_sgemv(
            ROW_MAJOR,
            NO_TRANS,
            n as i32,
            k as i32,
            alpha,
            a.as_ptr(),
            k as i32,
            x.as_ptr(),
            1,
            beta,
            y.as_mut_ptr(),
            1,
        );
    }
}

/// `y = alpha * A^T @ x + beta * y` with `A` stored row-major `[n,n]`.
#[inline]
pub fn sgemv_at(a: &[f32], x: &[f32], y: &mut [f32], n: usize, alpha: f32, beta: f32) {
    unsafe {
        cblas_sgemv(
            ROW_MAJOR,
            TRANS,
            n as i32,
            n as i32,
            alpha,
            a.as_ptr(),
            n as i32,
            x.as_ptr(),
            1,
            beta,
            y.as_mut_ptr(),
            1,
        );
    }
}

/// Rank-1 update: `A += alpha * x @ y^T`, `A` row-major `[n,n]`.
#[inline]
pub fn sger(a: &mut [f32], x: &[f32], y: &[f32], n: usize, alpha: f32) {
    unsafe {
        cblas_sger(
            ROW_MAJOR,
            n as i32,
            n as i32,
            alpha,
            x.as_ptr(),
            1,
            y.as_ptr(),
            1,
            a.as_mut_ptr(),
            n as i32,
        );
    }
}

/// In-place scale: `x *= alpha`.
#[inline]
pub fn sscal(x: &mut [f32], alpha: f32) {
    if x.is_empty() {
        return;
    }
    unsafe {
        cblas_sscal(x.len() as i32, alpha, x.as_mut_ptr(), 1);
    }
}

/// C = A @ B with custom strides (for reading interleaved data).
/// lda = stride between rows of A, ldc = stride between rows of C.
#[inline]
pub fn sgemm_strided(
    a: &[f32],
    b: &[f32],
    c: &mut [f32],
    m: usize,
    k: usize,
    n: usize,
    lda: usize,
    ldc: usize,
) {
    alias_check("sgemm_strided", a, b, c, m, k, n);
    unsafe {
        cblas_sgemm(
            ROW_MAJOR,
            NO_TRANS,
            NO_TRANS,
            m as i32,
            n as i32,
            k as i32,
            1.0,
            a.as_ptr(),
            lda as i32,
            b.as_ptr(),
            n as i32,
            0.0,
            c.as_mut_ptr(),
            ldc as i32,
        );
    }
}

/// NEON-vectorized bias addition in-place.
#[cfg(target_arch = "aarch64")]
pub fn bias_add(data: &mut [f32], bias: &[f32], m: usize, n: usize) {
    use std::arch::aarch64::*;
    let chunks = n / 4;
    unsafe {
        for row in 0..m {
            let base = row * n;
            for c in 0..chunks {
                let off = base + c * 4;
                let v = vld1q_f32(data.as_ptr().add(off));
                let b = vld1q_f32(bias.as_ptr().add(c * 4));
                vst1q_f32(data.as_mut_ptr().add(off), vaddq_f32(v, b));
            }
            for i in (chunks * 4)..n {
                data[base + i] += bias[i];
            }
        }
    }
}

#[cfg(not(target_arch = "aarch64"))]
pub fn bias_add(data: &mut [f32], bias: &[f32], m: usize, n: usize) {
    for row in 0..m {
        let base = row * n;
        for i in 0..n {
            data[base + i] += bias[i];
        }
    }
}

/// General sgemm with full control over transposition and strides.
/// C = alpha * op(A) @ op(B) + beta * C
/// op(X) = X if trans=false, X^T if trans=true
///
/// For row-major:
/// - NoTrans A: M×K with lda ≥ K
/// - Trans A:   stored K×M with lda ≥ M
/// - NoTrans B: K×N with ldb ≥ N
/// - Trans B:   stored N×K with ldb ≥ K
///
/// # Safety
/// Caller must ensure pointers are valid and output region is writable.
#[inline]
pub unsafe fn sgemm_general(
    a: *const f32,
    b: *const f32,
    c: *mut f32,
    m: usize,
    n: usize,
    k: usize,
    alpha: f32,
    beta: f32,
    lda: usize,
    ldb: usize,
    ldc: usize,
    trans_a: bool,
    trans_b: bool,
) {
    unsafe {
        cblas_sgemm(
            ROW_MAJOR,
            if trans_a { TRANS } else { NO_TRANS },
            if trans_b { TRANS } else { NO_TRANS },
            m as i32,
            n as i32,
            k as i32,
            alpha,
            a,
            lda as i32,
            b,
            ldb as i32,
            beta,
            c,
            ldc as i32,
        );
    }
}

/// sgemm + bias addition in one call. C = A @ B + bias (broadcast per row).
/// Auto-dispatches: NEON for tiny matrices, BLAS for everything else.
#[inline]
pub fn sgemm_bias(a: &[f32], b: &[f32], bias: &[f32], c: &mut [f32], m: usize, k: usize, n: usize) {
    alias_check("sgemm_bias", a, b, c, m, k, n);
    // Cost model decides: NEON when BLAS overhead dominates the compute.
    if m <= 8 && crate::cost::hw_model().prefer_neon_sgemm(m, k, n) {
        crate::kernels::neon_sgemm_bias_small(a, b, bias, c, m, k, n);
    } else if prefer_par_sgemm(m, k, n) {
        par_sgemm_bias(a, b, bias, c, m, k, n);
    } else {
        sgemm(a, b, c, m, k, n);
        bias_add(c, bias, m, n);
    }
}

/// sgemm with a generic epilogue closure (plan #1).
///
/// `C = epilogue(A @ B)`. The closure runs element-wise after the
/// matmul. Generics + `#[inline]` make the closure body monomorphize
/// per call site, so adding a new fusion type (matmul+scale, matmul+
/// clamp, matmul+gelu, etc.) is a new closure at the call site —
/// not a new `Thunk` variant + a new dispatch arm + a new boxed
/// function. Borrowed from MAX's `elementwise_lambda_fn` parameter
/// pattern on matmul; the Rust spelling is `impl Fn(f32) -> f32`.
///
/// For shapes where NEON beats BLAS, this routes through the NEON
/// path and applies the epilogue inline. Otherwise it does sgemm
/// then a single elementwise pass — still one less round-trip
/// through memory than sgemm + bias_add + activation.
#[inline]
pub fn sgemm_epilogue<E: Fn(f32) -> f32>(
    a: &[f32],
    b: &[f32],
    c: &mut [f32],
    m: usize,
    k: usize,
    n: usize,
    epilogue: E,
) {
    alias_check("sgemm_epilogue", a, b, c, m, k, n);
    sgemm(a, b, c, m, k, n);
    for v in c.iter_mut() {
        *v = epilogue(*v);
    }
}

/// sgemm + per-row bias + arbitrary post-activation in one call.
/// `C[i, j] = activation(A[i,:] @ B[:,j] + bias[j])`.
///
/// Mirrors `sgemm_bias` but the activation can be any closure.
/// `FusedMatMulBiasAct` could be reimplemented on top of this
/// instead of carrying its own dispatch logic.
#[inline]
pub fn sgemm_bias_epilogue<E: Fn(f32) -> f32>(
    a: &[f32],
    b: &[f32],
    bias: &[f32],
    c: &mut [f32],
    m: usize,
    k: usize,
    n: usize,
    activation: E,
) {
    alias_check("sgemm_bias_epilogue", a, b, c, m, k, n);
    sgemm(a, b, c, m, k, n);
    // Fuse bias + activation in one pass over C.
    for i in 0..m {
        let row = &mut c[i * n..(i + 1) * n];
        for (j, v) in row.iter_mut().enumerate() {
            *v = activation(*v + bias[j]);
        }
    }
}

/// sgemm with auto-dispatch: NEON for tiny m, Rayon-split BLAS for
/// medium/large shapes, sequential BLAS otherwise.
///
/// `limit_inner_threads()` pins OpenBLAS/MKL/Accelerate to 1 thread so
/// Rayon owns outer parallelism. Large DiT/CNN GEMMs go through [`par_sgemm`].
#[inline]
pub fn sgemm_auto(a: &[f32], b: &[f32], c: &mut [f32], m: usize, k: usize, n: usize) {
    alias_check("sgemm_auto", a, b, c, m, k, n);
    // `parity-gemm` feature swaps in the same Rust `gemm` crate that
    // candle uses, yielding bit-exact reduction order. Useful for
    // parity tests and reproducibility-critical workloads; loses AMX.
    //
    // The `return` reads as needless to clippy because every path below is
    // `cfg(not(feature = "parity-gemm"))`, making this the function tail under
    // this feature. Keep it explicit: an unguarded path added below must not
    // silently run a second GEMM over `c`.
    #[cfg(feature = "parity-gemm")]
    #[allow(clippy::needless_return)]
    {
        sgemm_via_gemm_crate(a, b, c, m, k, n);
        return;
    }
    // Direct ARM SME2 path (Apple M4+), opt-in via `RLX_CPU_SME=1`. Compiled
    // only under `amx-sme` on Apple; `dispatch_enabled()` also checks the chip
    // truly has SME2 at runtime. Placed after the parity-gemm early-return so
    // bit-exact-candle intent still wins when both features are on.
    #[cfg(all(rlx_cpu_amx_sme, not(feature = "parity-gemm")))]
    {
        if crate::intrinsics::apple_amx::sme::dispatch_enabled()
            && crate::intrinsics::apple_amx::sme::worth_sme(m, k, n)
        {
            crate::intrinsics::apple_amx::sme::sme_sgemm(a, b, c, m, k, n);
            return;
        }
    }
    // Native SME bf16 path (Apple M4+, `amx-sme`), opt-in via `RLX_CPU_SME_BF16=1`.
    // Our own bf16 GEMM (no vendor call); lossy f32→bf16 downcast so opt-in.
    #[cfg(all(rlx_cpu_amx_sme, not(feature = "parity-gemm")))]
    {
        if crate::intrinsics::apple_amx::sme::bf16_dispatch_enabled()
            && crate::intrinsics::apple_amx::sme::worth_sme(m, k, n)
        {
            crate::intrinsics::apple_amx::sme::sme_sgemm_bf16(a, b, c, m, k, n);
            return;
        }
    }
    // BNNS bf16 low-precision path (Apple, `amx-bnns`), opt-in via
    // `RLX_CPU_BNNS_BF16=1`. Lossy (f32→bf16 downcast) so strictly opt-in; if
    // BNNS rejects the shape it returns false and we fall through to Accelerate.
    #[cfg(all(rlx_cpu_amx_bnns, not(feature = "parity-gemm")))]
    {
        // f16 (`F16F32`, 10 mantissa bits) — more precise than bf16 at the same
        // half bandwidth; opt-in via `RLX_CPU_BNNS_F16=1`.
        if crate::intrinsics::apple_amx::bnns::dispatch_enabled_f16()
            && crate::intrinsics::apple_amx::bnns::matmul_f32_via_f16(a, b, c, m, k, n)
        {
            return;
        }
        if crate::intrinsics::apple_amx::bnns::dispatch_enabled()
            && crate::intrinsics::apple_amx::bnns::matmul_f32_via_bf16(a, b, c, m, k, n)
        {
            return;
        }
    }
    // `neon_sgemm_small` has a fixed `acc[8]` accumulator so m must be
    // ≤ 8. The cost model prefers NEON whenever its FLOP rate beats
    // BLAS-plus-overhead, which can fire for m up to thousands at
    // small k·n — we'd OOB on `acc[i]`. Cap explicitly.
    #[cfg(not(feature = "parity-gemm"))]
    if m <= 8 && crate::cost::hw_model().prefer_neon_sgemm(m, k, n) {
        crate::kernels::neon_sgemm_small(a, b, c, m, k, n);
        return;
    }
    #[cfg(not(feature = "parity-gemm"))]
    {
        let flops = (m as u64).saturating_mul(k as u64).saturating_mul(n as u64);
        // Prefer Rayon-split single-thread BLAS. Vendor MT BLAS for these DiT
        // shapes was slower on OpenBLAS (NFE=8 ~23s vs ~15s) and sticky MT
        // oversubscribed Rayon elementwise (BinaryFull ~5×). Keep
        // `with_blas_threads` for explicit callers.
        let _ = flops;
        if prefer_par_sgemm(m, k, n) {
            par_sgemm(a, b, c, m, k, n);
        } else {
            sgemm(a, b, c, m, k, n);
        }
    }
}

/// Enough FLOPs + a wide enough split axis to pay for Rayon hand-off.
#[inline]
fn prefer_par_sgemm(m: usize, k: usize, n: usize) -> bool {
    let workers = crate::pool::num_threads();
    if workers <= 1 {
        return false;
    }
    let flops = (m as u64).saturating_mul(k as u64).saturating_mul(n as u64);
    if flops < 1_000_000 {
        return false;
    }
    // Split the wider output axis; need ≥8 cols/rows per worker.
    m.max(n) >= workers.saturating_mul(8)
}

/// Parallel sgemm: split across Rayon workers, each calling single-threaded
/// CBLAS (see [`limit_inner_threads`]).
///
/// For large square-ish GEMMs (DiT FFN: m≈1k, n≈2–3k, k≈1k) a 2-D tile grid
/// beats pure column/row splits — each worker only streams an A-panel and a
/// B-panel instead of re-reading the full shared operand.
pub fn par_sgemm(a: &[f32], b: &[f32], c: &mut [f32], m: usize, k: usize, n: usize) {
    alias_check("par_sgemm", a, b, c, m, k, n);
    let workers = crate::pool::num_threads();
    if !prefer_par_sgemm(m, k, n) {
        sgemm(a, b, c, m, k, n);
        return;
    }
    let a_addr = a.as_ptr() as usize;
    let b_addr = b.as_ptr() as usize;
    let c_addr = c.as_mut_ptr() as usize;
    let flops = (m as u64).saturating_mul(k as u64).saturating_mul(n as u64);

    // 2-D tiling when both output axes are wide enough for a grid.
    if m >= 256 && n >= 256 && flops >= 200_000_000 {
        let tiles_m = ((workers as f64).sqrt().round() as usize).clamp(2, 8);
        let tiles_n = (workers / tiles_m).max(2);
        let n_tiles = tiles_m * tiles_n;
        let m_chunk = m.div_ceil(tiles_m);
        let n_chunk = n.div_ceil(tiles_n);
        crate::pool::par_for(n_tiles, 1, &|off, cnt| {
            for t in off..off + cnt {
                let tm = t / tiles_n;
                let tn = t % tiles_n;
                let m0 = tm * m_chunk;
                let n0 = tn * n_chunk;
                if m0 >= m || n0 >= n {
                    continue;
                }
                let local_m = (m0 + m_chunk).min(m) - m0;
                let local_n = (n0 + n_chunk).min(n) - n0;
                unsafe {
                    cblas_sgemm(
                        101,
                        111,
                        111,
                        local_m as i32,
                        local_n as i32,
                        k as i32,
                        1.0,
                        (a_addr as *const f32).add(m0 * k),
                        k as i32,
                        (b_addr as *const f32).add(n0),
                        n as i32,
                        0.0,
                        (c_addr as *mut f32).add(m0 * n + n0),
                        n as i32,
                    );
                }
            }
        });
        return;
    }

    if n >= m {
        // Split columns of C / B.
        let chunk = n.div_ceil(workers);
        crate::pool::par_for(workers, 1, &|off, cnt| {
            for w in off..off + cnt {
                let n_start = w * chunk;
                if n_start >= n {
                    continue;
                }
                let n_end = (n_start + chunk).min(n);
                let local_n = n_end - n_start;
                unsafe {
                    cblas_sgemm(
                        101,
                        111,
                        111,
                        m as i32,
                        local_n as i32,
                        k as i32,
                        1.0,
                        a_addr as *const f32,
                        k as i32,
                        (b_addr as *const f32).add(n_start),
                        n as i32,
                        0.0,
                        (c_addr as *mut f32).add(n_start),
                        n as i32,
                    );
                }
            }
        });
    } else {
        // Split rows of C / A.
        let chunk = m.div_ceil(workers);
        crate::pool::par_for(workers, 1, &|off, cnt| {
            for w in off..off + cnt {
                let m_start = w * chunk;
                if m_start >= m {
                    continue;
                }
                let m_end = (m_start + chunk).min(m);
                let local_m = m_end - m_start;
                unsafe {
                    cblas_sgemm(
                        101,
                        111,
                        111,
                        local_m as i32,
                        n as i32,
                        k as i32,
                        1.0,
                        (a_addr as *const f32).add(m_start * k),
                        k as i32,
                        b_addr as *const f32,
                        n as i32,
                        0.0,
                        (c_addr as *mut f32).add(m_start * n),
                        n as i32,
                    );
                }
            }
        });
    }
}

/// Parallel accumulate GEMM: same split as [`par_sgemm`] but β = 1 (add into C).
pub fn par_sgemm_accumulate(a: &[f32], b: &[f32], c: &mut [f32], m: usize, k: usize, n: usize) {
    alias_check("par_sgemm_accumulate", a, b, c, m, k, n);
    let workers = crate::pool::num_threads();
    if !prefer_par_sgemm(m, k, n) {
        sgemm_accumulate(a, b, c, m, k, n);
        return;
    }
    let a_addr = a.as_ptr() as usize;
    let b_addr = b.as_ptr() as usize;
    let c_addr = c.as_mut_ptr() as usize;
    let flops = (m as u64).saturating_mul(k as u64).saturating_mul(n as u64);

    if m >= 256 && n >= 256 && flops >= 200_000_000 {
        let tiles_m = ((workers as f64).sqrt().round() as usize).clamp(2, 8);
        let tiles_n = (workers / tiles_m).max(2);
        let n_tiles = tiles_m * tiles_n;
        let m_chunk = m.div_ceil(tiles_m);
        let n_chunk = n.div_ceil(tiles_n);
        crate::pool::par_for(n_tiles, 1, &|off, cnt| {
            for t in off..off + cnt {
                let tm = t / tiles_n;
                let tn = t % tiles_n;
                let m0 = tm * m_chunk;
                let n0 = tn * n_chunk;
                if m0 >= m || n0 >= n {
                    continue;
                }
                let local_m = (m0 + m_chunk).min(m) - m0;
                let local_n = (n0 + n_chunk).min(n) - n0;
                unsafe {
                    cblas_sgemm(
                        101,
                        111,
                        111,
                        local_m as i32,
                        local_n as i32,
                        k as i32,
                        1.0,
                        (a_addr as *const f32).add(m0 * k),
                        k as i32,
                        (b_addr as *const f32).add(n0),
                        n as i32,
                        1.0,
                        (c_addr as *mut f32).add(m0 * n + n0),
                        n as i32,
                    );
                }
            }
        });
        return;
    }

    if n >= m {
        let chunk = n.div_ceil(workers);
        crate::pool::par_for(workers, 1, &|off, cnt| {
            for w in off..off + cnt {
                let n_start = w * chunk;
                if n_start >= n {
                    continue;
                }
                let n_end = (n_start + chunk).min(n);
                let local_n = n_end - n_start;
                unsafe {
                    cblas_sgemm(
                        101,
                        111,
                        111,
                        m as i32,
                        local_n as i32,
                        k as i32,
                        1.0,
                        a_addr as *const f32,
                        k as i32,
                        (b_addr as *const f32).add(n_start),
                        n as i32,
                        1.0,
                        (c_addr as *mut f32).add(n_start),
                        n as i32,
                    );
                }
            }
        });
    } else {
        let chunk = m.div_ceil(workers);
        crate::pool::par_for(workers, 1, &|off, cnt| {
            for w in off..off + cnt {
                let m_start = w * chunk;
                if m_start >= m {
                    continue;
                }
                let m_end = (m_start + chunk).min(m);
                let local_m = m_end - m_start;
                unsafe {
                    cblas_sgemm(
                        101,
                        111,
                        111,
                        local_m as i32,
                        n as i32,
                        k as i32,
                        1.0,
                        (a_addr as *const f32).add(m_start * k),
                        k as i32,
                        b_addr as *const f32,
                        n as i32,
                        1.0,
                        (c_addr as *mut f32).add(m_start * n),
                        n as i32,
                    );
                }
            }
        });
    }
}

/// Bit-exact CPU sgemm via the same `gemm` crate candle uses for its
/// CPU backend. Row-major `[m, k] @ [k, n] = [m, n]` overwrite (β = 0).
#[cfg(feature = "parity-gemm")]
fn sgemm_via_gemm_crate(a: &[f32], b: &[f32], c: &mut [f32], m: usize, k: usize, n: usize) {
    use gemm::{Parallelism, gemm};
    let cfg = crate::config::RuntimeConfig::global();
    let workers = cfg.pool_workers + 1;
    let par = if workers > 1 {
        Parallelism::Rayon(workers)
    } else {
        Parallelism::None
    };
    // Row-major strides (cs=1, rs=#cols).
    unsafe {
        gemm(
            m,
            n,
            k,
            c.as_mut_ptr(),
            1,          // dst_cs
            n as isize, // dst_rs
            false,      // read_dst (β=0)
            a.as_ptr(),
            1,          // lhs_cs
            k as isize, // lhs_rs
            b.as_ptr(),
            1,          // rhs_cs
            n as isize, // rhs_rs
            0.0,        // alpha (zero out)
            1.0,        // beta
            false,      // conj_dst
            false,      // conj_lhs
            false,      // conj_rhs
            par,
        );
    }
}

/// Parallelized sgemm + bias: splits across the wider of m/n via Rayon,
/// each slice using single-threaded CBLAS ([`limit_inner_threads`]).
///
/// For each worker thread: computes a row or column slice of C.
pub fn par_sgemm_bias(
    a: &[f32],
    b: &[f32],
    bias: &[f32],
    c: &mut [f32],
    m: usize,
    k: usize,
    n: usize,
) {
    alias_check("par_sgemm_bias", a, b, c, m, k, n);
    let workers = crate::pool::num_threads();
    if !prefer_par_sgemm(m, k, n) {
        sgemm_bias(a, b, bias, c, m, k, n);
        return;
    }

    let a_addr = a.as_ptr() as usize;
    let b_addr = b.as_ptr() as usize;
    let bias_addr = bias.as_ptr() as usize;
    let c_addr = c.as_mut_ptr() as usize;

    if n >= m {
        let chunk = n.div_ceil(workers);
        crate::pool::par_for(workers, 1, &|off, cnt| {
            for w in off..off + cnt {
                let n_start = w * chunk;
                if n_start >= n {
                    continue;
                }
                let n_end = (n_start + chunk).min(n);
                let local_n = n_end - n_start;
                unsafe {
                    cblas_sgemm(
                        101,
                        111,
                        111,
                        m as i32,
                        local_n as i32,
                        k as i32,
                        1.0,
                        a_addr as *const f32,
                        k as i32,
                        (b_addr as *const f32).add(n_start),
                        n as i32,
                        0.0,
                        (c_addr as *mut f32).add(n_start),
                        n as i32,
                    );
                    let local_bias =
                        std::slice::from_raw_parts((bias_addr as *const f32).add(n_start), local_n);
                    let local_c =
                        std::slice::from_raw_parts_mut((c_addr as *mut f32).add(n_start), m * n);
                    for row in 0..m {
                        let base = row * n;
                        for i in 0..local_n {
                            local_c[base + i] += local_bias[i];
                        }
                    }
                }
            }
        });
    } else {
        let chunk = m.div_ceil(workers);
        crate::pool::par_for(workers, 1, &|off, cnt| {
            for w in off..off + cnt {
                let m_start = w * chunk;
                if m_start >= m {
                    continue;
                }
                let m_end = (m_start + chunk).min(m);
                let local_m = m_end - m_start;
                unsafe {
                    cblas_sgemm(
                        101,
                        111,
                        111,
                        local_m as i32,
                        n as i32,
                        k as i32,
                        1.0,
                        (a_addr as *const f32).add(m_start * k),
                        k as i32,
                        b_addr as *const f32,
                        n as i32,
                        0.0,
                        (c_addr as *mut f32).add(m_start * n),
                        n as i32,
                    );
                    let bias_sl = std::slice::from_raw_parts(bias_addr as *const f32, n);
                    let local_c = std::slice::from_raw_parts_mut(
                        (c_addr as *mut f32).add(m_start * n),
                        local_m * n,
                    );
                    for row in 0..local_m {
                        let base = row * n;
                        for i in 0..n {
                            local_c[base + i] += bias_sl[i];
                        }
                    }
                }
            }
        });
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn sgemm_identity() {
        // A = [[1,0],[0,1]], B = [[3,4],[5,6]]
        let a = [1.0, 0.0, 0.0, 1.0f32];
        let b = [3.0, 4.0, 5.0, 6.0f32];
        let mut c = [0.0f32; 4];
        sgemm(&a, &b, &mut c, 2, 2, 2);
        assert_eq!(c, [3.0, 4.0, 5.0, 6.0]);
    }

    #[test]
    fn sgemm_rectangular() {
        // [2,3] @ [3,2] = [2,2]
        let a = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0f32]; // [2,3]
        let b = [1.0, 0.0, 0.0, 1.0, 1.0, 0.0f32]; // [3,2]
        let mut c = [0.0f32; 4];
        sgemm(&a, &b, &mut c, 2, 3, 2);
        // [1*1+2*0+3*1, 1*0+2*1+3*0] = [4, 2]
        // [4*1+5*0+6*1, 4*0+5*1+6*0] = [10, 5]
        assert_eq!(c, [4.0, 2.0, 10.0, 5.0]);
    }

    #[test]
    fn sgemm_bias_test() {
        let a = [1.0, 0.0, 0.0, 1.0f32]; // identity [2,2]
        let b = [3.0, 4.0, 5.0, 6.0f32]; // [2,2]
        let bias = [10.0, 20.0f32]; // [2]
        let mut c = [0.0f32; 4];
        sgemm_bias(&a, &b, &bias, &mut c, 2, 2, 2);
        assert_eq!(c, [13.0, 24.0, 15.0, 26.0]);
    }

    #[test]
    fn dgemm_identity() {
        let a = [1.0, 0.0, 0.0, 1.0f64];
        let b = [3.0, 4.0, 5.0, 6.0f64];
        let mut c = [0.0f64; 4];
        dgemm(&a, &b, &mut c, 2, 2, 2);
        assert_eq!(c, [3.0, 4.0, 5.0, 6.0]);
    }

    /// `RLX_CPU_MATMUL_F64_ACCUM` precise path (#6): the f64-accumulating GEMM
    /// must reproduce an independent f64 reference bit-exactly after narrowing.
    #[test]
    fn dgemm_f32_precise_matches_f64_reference() {
        // A few shapes incl. a long K (where f32 accumulation drifts most).
        for &(m, k, n) in &[(2usize, 3usize, 4usize), (3, 512, 2), (1, 4096, 1)] {
            let mut s: u32 = 0x1234_5678 ^ (k as u32);
            let mut rng = || {
                s ^= s << 13;
                s ^= s >> 17;
                s ^= s << 5;
                (s as f32 / u32::MAX as f32) - 0.5
            };
            let a: Vec<f32> = (0..m * k).map(|_| rng()).collect();
            let b: Vec<f32> = (0..k * n).map(|_| rng()).collect();
            // Independent f64 oracle, narrowed to f32.
            let mut want = vec![0f32; m * n];
            for i in 0..m {
                for j in 0..n {
                    let mut acc = 0f64;
                    for p in 0..k {
                        acc += a[i * k + p] as f64 * b[p * n + j] as f64;
                    }
                    want[i * n + j] = acc as f32;
                }
            }
            let mut got = vec![0f32; m * n];
            dgemm_f32_precise(&a, &b, &mut got, m, k, n, false);
            assert_eq!(got, want, "precise GEMM != f64 oracle at ({m},{k},{n})");
        }
    }

    /// Catastrophic cancellation: the true result (1.0) is representable in f32
    /// but a left-to-right f32 partial sum annihilates it (`2^24 + 1 - 2^24`).
    /// The f64-accum path recovers it; this is the precision the knob buys.
    #[test]
    fn dgemm_f32_precise_survives_cancellation() {
        let big = 16_777_216.0f32; // 2^24, where f32 ULP == 2
        let a = [big, 1.0, -big]; // [1, 3]
        let b = [1.0f32, 1.0, 1.0]; // [3, 1]
        let mut c = [0.0f32];
        dgemm_f32_precise(&a, &b, &mut c, 1, 3, 1, false);
        assert_eq!(c[0], 1.0, "f64 accumulation should recover the exact 1.0");
        // Accumulate mode folds an existing C in at f64 precision too.
        let mut c2 = [10.0f32];
        dgemm_f32_precise(&a, &b, &mut c2, 1, 3, 1, true);
        assert_eq!(c2[0], 11.0);
    }

    #[test]
    fn dgesv_2x2_known_solution() {
        // A = [[2, 1],
        //      [1, 3]],   b = [5, 10]
        // Solution: x = [1, 3]  (verified: 2·1 + 1·3 = 5; 1·1 + 3·3 = 10)
        let mut a = [2.0, 1.0, 1.0, 3.0_f64];
        let mut b = [5.0, 10.0_f64];
        let info = dgesv(&mut a, &mut b, 2, 1);
        assert_eq!(info, 0, "dgesv signaled singular: info={info}");
        let want = [1.0, 3.0_f64];
        for (i, (g, w)) in b.iter().zip(want.iter()).enumerate() {
            assert!((g - w).abs() < 1e-12, "x[{i}] = {g}, expected {w}");
        }
    }

    #[test]
    fn dgesv_3x3_general() {
        // A = [[ 4, -1,  0],
        //      [-1,  4, -1],
        //      [ 0, -1,  4]]    (1-D Laplacian-ish; symmetric pos-def)
        // b = [1, 0, -1]
        // Expected: solve via reference. We just check Ax ≈ b after.
        let a_orig = [4.0, -1.0, 0.0, -1.0, 4.0, -1.0, 0.0, -1.0, 4.0_f64];
        let mut a = a_orig;
        let mut b = [1.0, 0.0, -1.0_f64];
        let info = dgesv(&mut a, &mut b, 3, 1);
        assert_eq!(info, 0);
        // Verify: Ax ≈ b_original.
        let mut residual = [0.0_f64; 3];
        for i in 0..3 {
            for j in 0..3 {
                residual[i] += a_orig[i * 3 + j] * b[j];
            }
        }
        let want_b = [1.0, 0.0, -1.0_f64];
        for i in 0..3 {
            assert!(
                (residual[i] - want_b[i]).abs() < 1e-12,
                "residual[{i}] = {} vs {}",
                residual[i],
                want_b[i]
            );
        }
    }

    #[test]
    fn sgemm_bt_test() {
        // Q@K^T: Q=[2,3], K=[2,3], result=[2,2]
        let q = [1.0, 0.0, 0.0, 0.0, 1.0, 0.0f32]; // [2,3]
        let k = [1.0, 0.0, 0.0, 0.0, 1.0, 0.0f32]; // [2,3]
        let mut scores = [0.0f32; 4];
        sgemm_bt(&q, &k, &mut scores, 2, 3, 2, 1.0);
        // Q@K^T = [[1,0],[0,1]] (dot products of rows)
        assert_eq!(scores, [1.0, 0.0, 0.0, 1.0]);
    }

    /// Plan #1: epilogue closure runs after the matmul. Identity
    /// closure must match plain sgemm; relu closure must match the
    /// hand-fused reference.
    #[test]
    fn sgemm_epilogue_matches_post_pass() {
        let a = [1.0f32, -2.0, 3.0, -4.0]; // [2, 2]
        let b = [1.0f32, 0.0, 0.0, 1.0]; // [2, 2]  (identity)
        let mut c1 = [0f32; 4];
        let mut c2 = [0f32; 4];
        // identity epilogue == plain sgemm
        sgemm(&a, &b, &mut c1, 2, 2, 2);
        sgemm_epilogue(&a, &b, &mut c2, 2, 2, 2, |x| x);
        assert_eq!(c1, c2);

        // relu epilogue zeros negative outputs
        let mut c3 = [0f32; 4];
        sgemm_epilogue(&a, &b, &mut c3, 2, 2, 2, |x| x.max(0.0));
        assert_eq!(c3, [1.0, 0.0, 3.0, 0.0]);
    }

    #[test]
    fn sgemm_bias_epilogue_matches_reference() {
        let a = [1.0f32, 2.0, 3.0, 4.0]; // [2, 2]
        let b = [1.0f32, 0.0, 0.0, 1.0]; // [2, 2]
        let bias = [10.0f32, 100.0];
        // Reference: A@B + bias, then activation
        // A@B = [[1, 2], [3, 4]]; +bias = [[11, 102], [13, 104]]; relu = same
        let mut c = [0f32; 4];
        sgemm_bias_epilogue(&a, &b, &bias, &mut c, 2, 2, 2, |x| x.max(0.0));
        assert_eq!(c, [11.0, 102.0, 13.0, 104.0]);
    }
}

/// bf16 bit-pattern → f32 (bf16 is the high 16 bits of an f32).
/// (Only the non-aarch64 SAXPY fallback uses this scalar form; NEON inlines it.)
#[inline]
#[allow(dead_code)]
fn bf16_bits_to_f32(bits: u16) -> f32 {
    f32::from_bits((bits as u32) << 16)
}

/// GEMM `C[m,n] = A[m,k] · Bᵀ` with a **B-transposed BF16** right-hand: `B` is
/// `[N, K]` row-major (each output row `n` is a CONTIGUOUS `K`-vector — the raw HF
/// `[vocab, hidden]` LM-head layout, so no transpose on load). `A`/`C` are f32.
/// Dequants `B` on the fly (f32 accumulate), reading HALF the weight bytes of an
/// f32 GEMM. Parallel over outputs `n`; each is an independent dot product
/// `C[m,n] = Σ_k A[m,k]·bf16(B[n,k])` — `A` stays hot in cache, `B[n,:]` streams.
/// The inner dot uses a 4-lane manual accumulator so the compiler can vectorize
/// the FMA (the bf16→f32 widen is a single shift). Precision-approximate (bf16
/// weights), NOT bit-exact to an f32 head.
/// `acc[nn] += scale · bf16(b[nn])` over a contiguous BF16 row — the SAXPY inner of
/// [`sgemm_bf16_rhs`]. NEON on aarch64: widen 8 `u16→u32`, shift left 16 (bf16 IS
/// the high half of an f32), FMA into the accumulator, so the dequant vectorizes.
#[inline]
fn saxpy_bf16(scale: f32, b: &[u16], acc: &mut [f32]) {
    #[cfg(target_arch = "aarch64")]
    unsafe {
        use std::arch::aarch64::*;
        let n = acc.len();
        let sv = vdupq_n_f32(scale);
        let mut nn = 0;
        while nn + 8 <= n {
            let bh = vld1q_u16(b.as_ptr().add(nn));
            let lo = vreinterpretq_f32_u32(vshlq_n_u32::<16>(vmovl_u16(vget_low_u16(bh))));
            let hi = vreinterpretq_f32_u32(vshlq_n_u32::<16>(vmovl_high_u16(bh)));
            let a0 = vld1q_f32(acc.as_ptr().add(nn));
            let a1 = vld1q_f32(acc.as_ptr().add(nn + 4));
            vst1q_f32(acc.as_mut_ptr().add(nn), vfmaq_f32(a0, sv, lo));
            vst1q_f32(acc.as_mut_ptr().add(nn + 4), vfmaq_f32(a1, sv, hi));
            nn += 8;
        }
        while nn < n {
            acc[nn] += scale * f32::from_bits((b[nn] as u32) << 16);
            nn += 1;
        }
    }
    #[cfg(not(target_arch = "aarch64"))]
    for (v, &bw) in acc.iter_mut().zip(b) {
        *v += scale * bf16_bits_to_f32(bw);
    }
}

/// GEMM `C[m,n] = A[m,k] · B[k,n]` with a **BF16** right-hand `B` in the standard
/// `[K, N]` (row-major) layout — the SAME tensor every backend's native matmul
/// consumes, so this is portable (not a CPU-only transposed reinterpret). `A`/`C`
/// are f32; `B` is dequant-on-the-fly (f32 accumulate), reading HALF the weight
/// bytes of an f32 GEMM. For the head GEMV (`m==1`) it splits `K` across threads —
/// each streams a CONTIGUOUS `B` slab (rows `k0..k1`) as a NEON SAXPY into a private
/// `[N]` partial, then reduces. NOTE: on Apple the f32 path runs on the AMX
/// coprocessor (Accelerate) and still wins; this helps non-AMX CPUs + is the
/// reference the GPU/ANE backends' native BF16 matmuls match. Precision-approximate.
pub fn sgemm_bf16_rhs(a: &[f32], b: &[u16], c: &mut [f32], m: usize, k: usize, n: usize) {
    use rayon::prelude::*;
    if m == 1 {
        let nthreads = rayon::current_num_threads().max(1);
        let kchunk = k.div_ceil(nthreads).max(1);
        let partials: Vec<Vec<f32>> = (0..nthreads)
            .into_par_iter()
            .map(|t| {
                let (k0, k1) = ((t * kchunk).min(k), ((t + 1) * kchunk).min(k));
                let mut acc = vec![0f32; n];
                for kk in k0..k1 {
                    saxpy_bf16(a[kk], &b[kk * n..kk * n + n], &mut acc);
                }
                acc
            })
            .collect();
        c.par_iter_mut().enumerate().for_each(|(j, cj)| {
            *cj = partials.iter().map(|p| p[j]).sum();
        });
    } else {
        c.par_chunks_mut(n).enumerate().for_each(|(i, crow)| {
            crow.iter_mut().for_each(|v| *v = 0.0);
            for kk in 0..k {
                saxpy_bf16(a[i * k + kk], &b[kk * n..kk * n + n], crow);
            }
        });
    }
}

/// `acc[nn] += scale · f16(b[nn])` over a contiguous IEEE-half row — the SAXPY
/// inner of [`sgemm_f16_rhs`]. f16→f32 widen accumulates in f32 (exact promotion).
#[inline]
fn saxpy_f16(scale: f32, b: &[u16], acc: &mut [f32]) {
    for (v, &bw) in acc.iter_mut().zip(b) {
        *v += scale * half::f16::from_bits(bw).to_f32();
    }
}

/// GEMM `C[m,n] = A[m,k](f32) · B[k,n](F16)`, the IEEE-half twin of
/// [`sgemm_bf16_rhs`]. `B` is a `k*n` `u16` (f16-bit) buffer in the standard
/// `[K, N]` row-major layout every backend's native matmul consumes; `A`/`C` are
/// f32; `B` is dequant-on-the-fly with f32 accumulate (so bit-parity with an f32
/// GEMM over the same f16-valued weights). Fixes generic `Op::MatMul` reading an
/// F16 weight as f32 garbage (there was only a BF16 path before).
pub fn sgemm_f16_rhs(a: &[f32], b: &[u16], c: &mut [f32], m: usize, k: usize, n: usize) {
    use rayon::prelude::*;
    if m == 1 {
        let nthreads = rayon::current_num_threads().max(1);
        let kchunk = k.div_ceil(nthreads).max(1);
        let partials: Vec<Vec<f32>> = (0..nthreads)
            .into_par_iter()
            .map(|t| {
                let (k0, k1) = ((t * kchunk).min(k), ((t + 1) * kchunk).min(k));
                let mut acc = vec![0f32; n];
                for kk in k0..k1 {
                    saxpy_f16(a[kk], &b[kk * n..kk * n + n], &mut acc);
                }
                acc
            })
            .collect();
        c.par_iter_mut().enumerate().for_each(|(j, cj)| {
            *cj = partials.iter().map(|p| p[j]).sum();
        });
    } else {
        c.par_chunks_mut(n).enumerate().for_each(|(i, crow)| {
            crow.iter_mut().for_each(|v| *v = 0.0);
            for kk in 0..k {
                saxpy_f16(a[i * k + kk], &b[kk * n..kk * n + n], crow);
            }
        });
    }
}

#[cfg(test)]
mod bf16_gemm_tests {
    use super::*;
    #[test]
    fn sgemm_bf16_rhs_matches_f32_rounded() {
        // B is [K, N] (standard): out[i,j] = Σ_k a[i,k]·bf16(b[k,j]).
        for &(m, k, n) in &[(1usize, 128usize, 64usize), (3, 96, 40)] {
            let a: Vec<f32> = (0..m * k)
                .map(|i| ((i * 7 % 13) as f32 - 6.0) * 0.1)
                .collect();
            let bf: Vec<f32> = (0..k * n)
                .map(|i| ((i * 5 % 11) as f32 - 5.0) * 0.1)
                .collect();
            let b16: Vec<u16> = bf.iter().map(|&x| (x.to_bits() >> 16) as u16).collect();
            let brnd: Vec<f32> = b16
                .iter()
                .map(|&w| f32::from_bits((w as u32) << 16))
                .collect();
            let mut cref = vec![0f32; m * n];
            sgemm(&a, &brnd, &mut cref, m, k, n);
            let mut cbf = vec![0f32; m * n];
            sgemm_bf16_rhs(&a, &b16, &mut cbf, m, k, n);
            let worst = cref
                .iter()
                .zip(&cbf)
                .map(|(x, y)| (x - y).abs())
                .fold(0.0, f32::max);
            assert!(worst < 1e-4, "bf16 gemm ({m}x{k}x{n}) mismatch {worst:.2e}");
        }
    }
}

#[cfg(test)]
mod bf16_bench {
    use super::*;
    use std::time::Instant;
    #[test]
    #[ignore = "benchmark; run with --ignored --nocapture"]
    fn bench_head_bf16_vs_f32() {
        let (m, k, n) = (1usize, 7168usize, 163840usize); // real LM-head GEMV
        let a: Vec<f32> = (0..m * k)
            .map(|i| ((i % 97) as f32 - 48.0) * 0.01)
            .collect();
        let bf: Vec<f32> = (0..k * n)
            .map(|i| ((i % 101) as f32 - 50.0) * 0.01)
            .collect();
        let b16: Vec<u16> = bf.iter().map(|&x| (x.to_bits() >> 16) as u16).collect(); // [K,N]
        let bf32_kn: Vec<f32> = b16
            .iter()
            .map(|&w| f32::from_bits((w as u32) << 16))
            .collect();
        let mut c = vec![0f32; m * n];
        for _ in 0..2 {
            sgemm(&a, &bf32_kn, &mut c, m, k, n);
        } // warm
        let t = Instant::now();
        for _ in 0..5 {
            sgemm(&a, &bf32_kn, &mut c, m, k, n);
        }
        let f32_ms = t.elapsed().as_secs_f64() * 1000.0 / 5.0;
        for _ in 0..2 {
            sgemm_bf16_rhs(&a, &b16, &mut c, m, k, n);
        }
        let t = Instant::now();
        for _ in 0..5 {
            sgemm_bf16_rhs(&a, &b16, &mut c, m, k, n);
        }
        let bf16_ms = t.elapsed().as_secs_f64() * 1000.0 / 5.0;
        eprintln!(
            "HEAD GEMV [1,{k}]x[{k},{n}]: cblas-f32 {f32_ms:.1}ms  bf16-neon {bf16_ms:.1}ms  ({:.2}x)",
            f32_ms / bf16_ms
        );
    }
}

#[cfg(test)]
mod linalg_fallback_tests {
    use super::*;

    // ── Linalg wrappers: reference-reconstruction tests ──────────────
    //
    // These exercise the public row-major wrappers, so they validate BOTH
    // the linked-BLAS/LAPACK path (default) AND the dependency-free
    // fallback (`--no-default-features`, `rlx_cpu_blas` unset). Both must
    // agree with the independent reference here.

    fn matmul_rm(a: &[f64], b: &[f64], m: usize, k: usize, n: usize) -> Vec<f64> {
        let mut c = vec![0f64; m * n];
        for i in 0..m {
            for j in 0..n {
                let mut s = 0.0;
                for p in 0..k {
                    s += a[i * k + p] * b[p * n + j];
                }
                c[i * n + j] = s;
            }
        }
        c
    }

    fn transpose_rm(a: &[f64], r: usize, c: usize) -> Vec<f64> {
        let mut t = vec![0f64; r * c];
        for i in 0..r {
            for j in 0..c {
                t[j * r + i] = a[i * c + j];
            }
        }
        t
    }

    #[test]
    fn dpotrf_lower_reconstructs_spd() {
        // A = Mᵀ·M is SPD; Cholesky L must satisfy L·Lᵀ = A.
        let m = [2.0, 0.0, 1.0, 0.5, 1.0, 0.0, 0.0, 0.3, 1.5];
        let mt = transpose_rm(&m, 3, 3);
        let a = matmul_rm(&mt, &m, 3, 3, 3);
        let mut lf = a.clone();
        let info = dpotrf(&mut lf, 3, true);
        assert_eq!(info, 0, "dpotrf flagged non-SPD");
        let lt = transpose_rm(&lf, 3, 3);
        let recon = matmul_rm(&lf, &lt, 3, 3, 3);
        for i in 0..9 {
            assert!(
                (recon[i] - a[i]).abs() < 1e-9,
                "L·Lᵀ[{i}]={} A={}",
                recon[i],
                a[i]
            );
        }
    }

    #[test]
    fn dpotrf_rejects_non_spd() {
        // Indefinite: leading 2×2 minor is fine but full matrix isn't SPD.
        let mut a = [1.0, 2.0, 2.0, 1.0];
        let info = dpotrf(&mut a, 2, true);
        assert!(info > 0, "expected non-SPD signal, got {info}");
    }

    #[test]
    fn dsyevd_eigenpairs_ascending() {
        // 1-D Laplacian-ish tridiagonal; eigenvalues 2−√2, 2, 2+√2.
        let a = [2.0, 1.0, 0.0, 1.0, 2.0, 1.0, 0.0, 1.0, 2.0];
        let a_orig = a;
        let mut af = a;
        let mut w = [0.0; 3];
        let info = dsyevd(&mut af, &mut w, 3);
        assert_eq!(info, 0);
        let want = [2.0 - 2f64.sqrt(), 2.0, 2.0 + 2f64.sqrt()];
        for i in 0..3 {
            if i > 0 {
                assert!(w[i] >= w[i - 1] - 1e-12, "eigenvalues not ascending");
            }
            assert!(
                (w[i] - want[i]).abs() < 1e-9,
                "w[{i}]={} want {}",
                w[i],
                want[i]
            );
        }
        // Eigenvector j is column j of the col-major output, i.e. the
        // CONTIGUOUS triple `af[j*n + i]` (== `spd.rs`'s `evecs[k*n+i]`
        // convention), NOT the strided `af[i*n + j]`. Validate both
        // A·v = w·v (per-column, sign-invariant) and the full spectral
        // reconstruction Σ_j w_j·v_j·v_jᵀ = A (fully sign-invariant, so
        // it holds regardless of each backend's eigenvector-sign choice).
        let mut recon = [0.0f64; 9];
        for j in 0..3 {
            let v: Vec<f64> = (0..3).map(|i| af[j * 3 + i]).collect();
            for i in 0..3 {
                let av: f64 = (0..3).map(|l| a_orig[i * 3 + l] * v[l]).sum();
                assert!((av - w[j] * v[i]).abs() < 1e-9, "A·v≠w·v j={j} i={i}");
                for c in 0..3 {
                    recon[i * 3 + c] += w[j] * v[i] * v[c];
                }
            }
        }
        for idx in 0..9 {
            assert!(
                (recon[idx] - a_orig[idx]).abs() < 1e-9,
                "spectral recon[{idx}]={} A={}",
                recon[idx],
                a_orig[idx]
            );
        }
    }

    #[test]
    fn dgeqrf_full_reconstructs_and_orthonormal() {
        let (m, n, k) = (4usize, 3usize, 3usize);
        let a = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 10.0, 1.0, 0.0, 1.0];
        let mut af = a;
        let mut q = vec![0.0; m * k];
        let mut r = vec![0.0; k * n];
        let info = dgeqrf_full(&mut af, m, n, &mut q, &mut r);
        assert_eq!(info, 0);
        let recon = matmul_rm(&q, &r, m, k, n);
        for i in 0..m * n {
            assert!((recon[i] - a[i]).abs() < 1e-9, "Q·R[{i}]");
        }
        let qt = transpose_rm(&q, m, k);
        let qtq = matmul_rm(&qt, &q, k, m, k);
        for i in 0..k {
            for j in 0..k {
                let e = if i == j { 1.0 } else { 0.0 };
                assert!((qtq[i * k + j] - e).abs() < 1e-9, "Qᵀ·Q[{i},{j}]");
            }
        }
    }

    #[test]
    fn dgesvd_thin_reconstructs_both_shapes() {
        for &(m, n) in &[(4usize, 2usize), (2usize, 4usize), (3, 3)] {
            let k = m.min(n);
            let a: Vec<f64> = (0..m * n)
                .map(|i| ((i * 7 % 11) as f64) - 5.0 + 0.25 * i as f64)
                .collect();
            let mut af = a.clone();
            let mut s = vec![0.0; k];
            let mut u = vec![0.0; m * k];
            let mut vt = vec![0.0; k * n];
            let info = dgesvd_thin(&mut af, m, n, &mut s, &mut u, &mut vt);
            assert_eq!(info, 0, "svd info m={m} n={n}");
            for i in 1..k {
                assert!(s[i] <= s[i - 1] + 1e-12, "s not descending");
            }
            let mut us = vec![0.0; m * k];
            for i in 0..m {
                for j in 0..k {
                    us[i * k + j] = u[i * k + j] * s[j];
                }
            }
            let recon = matmul_rm(&us, &vt, m, k, n);
            for i in 0..m * n {
                assert!(
                    (recon[i] - a[i]).abs() < 1e-8,
                    "SVD recon m={m} n={n} [{i}] got {} want {}",
                    recon[i],
                    a[i]
                );
            }
        }
    }

    #[test]
    fn dgesdd_thin_reconstructs() {
        let (m, n, k) = (5usize, 3usize, 3usize);
        let a: Vec<f64> = (0..m * n).map(|i| (i as f64).sin() * 3.0 + 1.0).collect();
        let mut af = a.clone();
        let mut s = vec![0.0; k];
        let mut u = vec![0.0; m * k];
        let mut vt = vec![0.0; k * n];
        let info = dgesdd_thin(&mut af, m, n, &mut s, &mut u, &mut vt);
        assert_eq!(info, 0);
        let mut us = vec![0.0; m * k];
        for i in 0..m {
            for j in 0..k {
                us[i * k + j] = u[i * k + j] * s[j];
            }
        }
        let recon = matmul_rm(&us, &vt, m, k, n);
        for i in 0..m * n {
            assert!((recon[i] - a[i]).abs() < 1e-8, "dgesdd recon[{i}]");
        }
    }

    #[test]
    fn dgelsd_overdetermined_lstsq() {
        // Fit y = 2x + 1; design columns [x, 1], exact solution [2, 1].
        let xs = [0.0, 1.0, 2.0, 3.0];
        let mut a = vec![0.0; 4 * 2];
        let mut b = vec![0.0; 4];
        for (i, &x) in xs.iter().enumerate() {
            a[i * 2] = x;
            a[i * 2 + 1] = 1.0;
            b[i] = 2.0 * x + 1.0;
        }
        let mut x = vec![0.0; 2];
        let info = dgelsd_solve(&a, &b, 4, 2, 1, -1.0, &mut x);
        assert_eq!(info, 0);
        assert!(
            (x[0] - 2.0).abs() < 1e-9 && (x[1] - 1.0).abs() < 1e-9,
            "lstsq x={x:?}"
        );
    }

    #[test]
    fn lu_slogdet_diagonal_known() {
        let mut a = [2.0, 0.0, 0.0, 0.0, 3.0, 0.0, 0.0, 0.0, 4.0];
        let (logabs, sign, det) = lu_slogdet(&mut a, 3);
        assert!((det - 24.0).abs() < 1e-9, "det={det}");
        assert!((sign - 1.0).abs() < 1e-12);
        assert!((logabs - 24f64.ln()).abs() < 1e-9);
    }

    #[test]
    fn dtrsm_lower_solve_roundtrips() {
        let l = [2.0, 0.0, 0.0, 1.0, 3.0, 0.0, 0.5, -1.0, 4.0]; // 3×3 lower
        let xtrue = [1.0, 2.0, 3.0, -1.0, 0.5, 0.5]; // 3×2
        let mut b = matmul_rm(&l, &xtrue, 3, 3, 2); // B = L·X
        dtrsm_lower_or_upper(&l, &mut b, 3, 2, true, false);
        for i in 0..6 {
            assert!(
                (b[i] - xtrue[i]).abs() < 1e-9,
                "dtrsm X[{i}]={} want {}",
                b[i],
                xtrue[i]
            );
        }
    }

    /// The no-BLAS `sgemm` fast path (i→p→j) must agree with a naive `ijk`
    /// reference on a non-trivial shape with non-unit `alpha`/`beta`, and
    /// (when built `--no-default-features`) be markedly faster. Correctness
    /// runs always; the timing print is informational.
    #[test]
    fn sgemm_fallback_ikj_matches_and_timing() {
        let (m, k, n) = (96usize, 128usize, 112usize);
        let mut s: u32 = 0x9e37_79b9;
        let mut rng = || {
            s ^= s << 13;
            s ^= s >> 17;
            s ^= s << 5;
            (s as f32 / u32::MAX as f32) - 0.5
        };
        let a: Vec<f32> = (0..m * k).map(|_| rng()).collect();
        let b: Vec<f32> = (0..k * n).map(|_| rng()).collect();
        // Naive ijk reference: C = A·B.
        let mut want = vec![0f32; m * n];
        for i in 0..m {
            for j in 0..n {
                let mut acc = 0f32;
                for p in 0..k {
                    acc += a[i * k + p] * b[p * n + j];
                }
                want[i * n + j] = acc;
            }
        }
        let mut got = vec![0f32; m * n];
        sgemm(&a, &b, &mut got, m, k, n);
        // Same math, possibly a different summation order → allow a small
        // relative slack rather than bit-exactness.
        for idx in 0..m * n {
            let tol = 1e-4 * (1.0 + want[idx].abs());
            assert!(
                (got[idx] - want[idx]).abs() <= tol,
                "sgemm[{idx}]={} want {}",
                got[idx],
                want[idx]
            );
        }

        // Informational A/B (only meaningful in a no-BLAS build; on a
        // vendor-BLAS build `sgemm` dispatches to Accelerate/OpenBLAS/MKL).
        // Times the fast `sgemm` fallback vs the naive p-innermost loop on
        // the same shape so the reorder's win is visible directly.
        use std::time::Instant;
        let (bm, bk, bn) = (256usize, 256usize, 256usize);
        let ba: Vec<f32> = (0..bm * bk).map(|_| rng()).collect();
        let bb: Vec<f32> = (0..bk * bn).map(|_| rng()).collect();
        let mut bc = vec![0f32; bm * bn];
        let flop = 2.0 * bm as f64 * bk as f64 * bn as f64;
        let naive = |a: &[f32], b: &[f32], c: &mut [f32]| {
            for i in 0..bm {
                for j in 0..bn {
                    let mut acc = 0f32;
                    for p in 0..bk {
                        acc += a[i * bk + p] * b[p * bn + j];
                    }
                    c[i * bn + j] = acc;
                }
            }
        };
        for _ in 0..2 {
            sgemm(&ba, &bb, &mut bc, bm, bk, bn);
            naive(&ba, &bb, &mut bc);
        }
        let iters = 20;
        let tf = Instant::now();
        for _ in 0..iters {
            sgemm(&ba, &bb, &mut bc, bm, bk, bn);
        }
        let fast_ms = tf.elapsed().as_secs_f64() * 1e3 / iters as f64;
        let tn = Instant::now();
        for _ in 0..iters {
            naive(&ba, &bb, &mut bc);
        }
        let naive_ms = tn.elapsed().as_secs_f64() * 1e3 / iters as f64;
        eprintln!(
            "SGEMM [{bm}x{bk}x{bn}]: fast {fast_ms:.2}ms ({:.2} GFLOP/s)  naive-ijk {naive_ms:.2}ms ({:.2} GFLOP/s)  speedup {:.1}x",
            flop / (fast_ms * 1e-3) / 1e9,
            flop / (naive_ms * 1e-3) / 1e9,
            naive_ms / fast_ms
        );
    }
}

/// Widen a packed half buffer (`u16` bits, BF16 or IEEE-F16) to f32.
///
/// A half tensor is stored packed at 2 B/elem in some arena slots and widened
/// to 4 B in others; every f32 kernel that reads such a slot must go through
/// this, or it reinterprets two halves as one f32 (silent NaN/garbage).
pub fn widen_half_vec(src: &[u16], bf16: bool) -> Vec<f32> {
    if bf16 {
        src.iter()
            .map(|&w| f32::from_bits((w as u32) << 16))
            .collect()
    } else {
        src.iter()
            .map(|&w| half::f16::from_bits(w).to_f32())
            .collect()
    }
}

/// Round f32 back into a packed half buffer — inverse of [`widen_half_vec`].
pub fn round_half_into(src: &[f32], dst: &mut [u16], bf16: bool) {
    if bf16 {
        for (d, &v) in dst.iter_mut().zip(src) {
            // round-to-nearest-even on the truncated mantissa, matching the
            // Cast lowering (plain truncation drifts on long accumulations).
            let bits = v.to_bits();
            let round = ((bits >> 16) & 1) + 0x7fff;
            *d = ((bits.wrapping_add(round)) >> 16) as u16;
        }
    } else {
        for (d, &v) in dst.iter_mut().zip(src) {
            *d = half::f16::from_f32(v).to_bits();
        }
    }
}