inferencelayer 0.2.5

//! Prepacked f32 GEMM for the encoder's linear layers: `C[M,N] = A[M,K] · W[N,K]ᵀ (+ bias)`.
//!
//! # Why this exists
//!
//! The `gemm` crate (0.19) repacks its weight matrix on **every call** on aarch64 and exposes no
//! prepack API — `faer` has none either. The pack cost scales with the *weight*, not the batch, so
//! it is invisible at large M and ruinous at small M. Measured on this crate's own shapes
//! (`gemm_pack_probe`, K=N=768):
//!
//! ```text
//!   M:       16     32     64    128    256    512
//!   gemm:   181    365    355    481    509    595   GF/s
//! ```
//!
//! An encoder weight is immutable for the life of the process, so that pack belongs at load time.
//! [`PackedWeight`] does it once; [`gemm_packed`] then streams the panels straight into the
//! microkernel. Small-M shapes are not a corner case here — they are most of the workload (a short
//! text, an entity-label prompt, a reranked pair), and they were paying the full toll.
//!
//! # The microkernel
//!
//! Outer-product form, MR=8 rows × NR=8 columns, accumulated in 16 NEON registers. Both operands
//! are packed K-major, so each k-step is four vector loads (two of A, two of W) feeding sixteen
//! `vfmaq_laneq_f32` — 2 FLOP/byte, versus 0.25 for the inner-product/dot form, which is why a
//! GEMM built out of [`crate::simd::dot`] would still be bandwidth-bound no matter how well the dot
//! vectorizes. `vfmaq_laneq_f32` also gets us the fused multiply-add that safe Rust cannot emit
//! (see [`crate::simd`] — Rust does not enable FP contraction, and that alone is ~2×).
//!
//! # Determinism
//!
//! The k-loop runs in a fixed order and K is never split across threads; only the M×N tiles are
//! parallel, and they are disjoint. So the result is bit-reproducible run to run — required,
//! because the CPU path is the oracle the GPU kernels are gated against.
//!
//! On x86_64 the same prepacked path runs an AVX2+FMA twin of the NEON microkernel
//! (runtime-detected — pre-AVX2 hardware falls to the portable kernel, correct but slow). The
//! `gemm`-crate dispatch in `encoder_cpu::linear_from` remains for the encoder linears off
//! aarch64; the direct `PackedWeight` callers (GLiNER family, vision) now get real SIMD on both
//! ISAs. The portable kernel doubles as the reference the SIMD twins are tested against.

use rayon::prelude::*;

/// Rows of A per microkernel tile.
const MR: usize = 8;
/// Columns of W per microkernel tile — one panel.
const NR: usize = 8;

/// A `[N, K]` row-major weight, packed once into K-major panels of [`NR`] columns.
///
/// Layout: `p[panel * (k * NR) + kk * NR + j]` is `w[(panel * NR + j) * k + kk]`. The tail panel is
/// zero-padded to a full [`NR`], so the microkernel never needs a ragged-N variant; the *store*
/// masks instead, which costs nothing in the inner loop.
pub struct PackedWeight {
    p: Vec<f32>,
    n: usize,
    k: usize,
}

impl PackedWeight {
    /// Pack a `[n, k]` row-major weight (HF layout: `out_features × in_features`).
    pub fn new(w: &[f32], n: usize, k: usize) -> Self {
        assert_eq!(w.len(), n * k, "weight is not [n, k]");
        let panels = n.div_ceil(NR);
        let mut p = vec![0f32; panels * k * NR];
        for panel in 0..panels {
            let base = panel * k * NR;
            for j in 0..NR {
                let row = panel * NR + j;
                if row >= n {
                    break; // zero-padded tail
                }
                let src = &w[row * k..(row + 1) * k];
                for (kk, &v) in src.iter().enumerate() {
                    p[base + kk * NR + j] = v;
                }
            }
        }
        Self { p, n, k }
    }

    /// Output width.
    pub fn n(&self) -> usize {
        self.n
    }

    /// Contraction depth.
    pub fn k(&self) -> usize {
        self.k
    }

    /// Recover the original `[n, k]` row-major weight from the panels — exact (packing is a pure
    /// permutation). GPU loaders use this to upload from an already-built CPU model instead of
    /// re-reading and re-folding the checkpoint.
    pub fn unpack(&self) -> Vec<f32> {
        let mut w = vec![0f32; self.n * self.k];
        for panel in 0..self.n.div_ceil(NR) {
            let base = panel * self.k * NR;
            for j in 0..NR {
                let row = panel * NR + j;
                if row >= self.n {
                    break;
                }
                for kk in 0..self.k {
                    w[row * self.k + kk] = self.p[base + kk * NR + j];
                }
            }
        }
        w
    }
}

/// `c[1, N] = a[1, K] · wᵀ (+ bias)` — the m == 1 case without the padding rows.
///
/// [`gemm_packed`] evaluates a full `MR × NR` tile whatever `m` is, so a single row computes
/// `MR − 1` rows of zero padding. **Deleting that arithmetic is worth almost nothing**, and it is
/// worth recording why: at m = 1 the operation is MEMORY-BOUND — it streams the whole weight
/// matrix either way — so the padded FLOPs ride along free on a kernel already starved for weight
/// bandwidth. Interleaved A/B at every decode shape in the engine (`tests/gemv_ab.rs`):
/// **0.94x–1.14x, a wash**.
///
/// What this function is actually for is TASK CONTROL: it runs serially below `SERIAL_MACS` and
/// sizes tasks by work above it, where `gemm_packed` fans every call out to rayon (plus a parallel
/// pass just to pack a mostly-zero A). Back-to-back small calls are where that bites — Pocket
/// TTS's frame went 135 → 26 ms once its m = 1 projections stopped dispatching. So this is called
/// directly where that pattern holds, and `gemm_packed` deliberately does NOT reroute m = 1
/// callers into it.
///
/// **Bit-identical to `gemm_packed(.., m = 1, ..)`**: same panel decomposition, same
/// bias-seeded accumulators, same ascending-k order per output, same fused multiply-add — it is
/// the microkernel's row 0 with the padding rows deleted, not a different summation.
pub fn gemv_packed(c: &mut [f32], a: &[f32], w: &PackedWeight, bias: Option<&[f32]>) {
    let (n, k) = (w.n, w.k);
    assert_eq!(a.len(), k, "lhs is not [1, k]");
    assert!(c.len() >= n, "output too small");
    if let Some(b) = bias {
        assert_eq!(b.len(), n, "bias width");
    }
    // Task sizing by WORK, not by task count. A panel is only `NR·k` MACs, so splitting purely
    // by thread count hands rayon tasks too small to pay for themselves — at Whisper's 512×512
    // that was 64 tasks of ~4k MACs and it ran SLOWER than the padded GEMM it replaced. Give each
    // task a floor of `TASK_MACS`, and skip the pool entirely below `SERIAL_MACS`.
    const SERIAL_MACS: usize = 1 << 18;
    const TASK_MACS: usize = 1 << 16;
    let panels = n.div_ceil(NR);
    let per_task = if n * k <= SERIAL_MACS {
        panels
    } else {
        let by_work = TASK_MACS.div_ceil(NR * k.max(1));
        let by_threads = panels.div_ceil(4 * rayon::current_num_threads().max(1));
        by_work.max(by_threads).max(1).min(panels)
    };
    c[..n]
        .par_chunks_mut(per_task * NR)
        .enumerate()
        .for_each(|(task, out)| {
            for (p, oc) in out.chunks_mut(NR).enumerate() {
                let panel = task * per_task + p;
                let wp = &w.p[panel * k * NR..(panel + 1) * k * NR];
                let mut acc = [0f32; NR];
                if let Some(b) = bias {
                    let j0 = panel * NR;
                    acc[..oc.len()].copy_from_slice(&b[j0..j0 + oc.len()]);
                }
                for kk in 0..k {
                    let av = a[kk];
                    let wrow = &wp[kk * NR..(kk + 1) * NR];
                    for j in 0..NR {
                        acc[j] = av.mul_add(wrow[j], acc[j]);
                    }
                }
                oc.copy_from_slice(&acc[..oc.len()]);
            }
        });
}

/// `c[M, N] = a[M, K] · wᵀ (+ bias)`, with `w` prepacked.
///
/// Default path: A packs ONCE in a parallel pass, then the compute tasks tile BOTH dimensions —
/// (row-block × panel-chunk). The old row-block-only split put `m/8` tasks on the pool (24 for
/// a 192-token call: 1.5 waves over 16 threads, a third of the machine idle in the tail) and
/// re-packed A per task. Each 2D task owns a DISJOINT `[MR rows × panel-chunk cols]` tile of
/// `c` and the per-output k-order is unchanged, so the result is bit-identical to the 1D path
/// (`OSFKB_CPU_GEMM_2D=0` — the A/B control) and stays deterministic.
///
/// At `m == 1` this spends seven eighths of the microkernel on zero-padded rows, which MEASURES
/// AS A WASH (the shape is memory-bound) — but it also fans every call out to the pool. Callers
/// making many small m = 1 calls back to back want [`gemv_packed`] directly; see its docs.
pub fn gemm_packed(c: &mut [f32], a: &[f32], w: &PackedWeight, m: usize, bias: Option<&[f32]>) {
    let (n, k) = (w.n, w.k);
    assert_eq!(a.len(), m * k, "lhs is not [m, k]");
    assert!(c.len() >= m * n, "output too small");
    if let Some(b) = bias {
        assert_eq!(b.len(), n, "bias width");
    }
    let panels = n.div_ceil(NR);
    let use_2d = std::env::var("OSFKB_CPU_GEMM_2D").ok().as_deref() != Some("0");

    if use_2d {
        let nblocks = m.div_ceil(MR);
        // Pass 1: pack A K-major per row block, in parallel, once. (The 1D path re-packed
        // inside every task; sharing one packed copy also lets the 2D tasks start anywhere.)
        let mut apack = vec![0f32; nblocks * k * MR];
        apack
            .par_chunks_mut(k * MR)
            .enumerate()
            .for_each(|(blk, ap)| {
                let m0 = blk * MR;
                let rows = MR.min(m - m0);
                for i in 0..rows {
                    let src = &a[(m0 + i) * k..(m0 + i + 1) * k];
                    for (kk, &v) in src.iter().enumerate() {
                        ap[kk * MR + i] = v;
                    }
                }
            });
        // Pass 2: (block × panel-chunk) tasks — enough to fill the pool ~3 waves deep without
        // drowning it in rayon overhead.
        let threads = rayon::current_num_threads().max(1);
        let nchunks = (3 * threads).div_ceil(nblocks).clamp(1, panels);
        let chunk = panels.div_ceil(nchunks);
        let cptr = SendPtr(c.as_mut_ptr());
        // PANEL-MAJOR task order: concurrent cores work the SAME panel-chunk across different
        // row-blocks, so the W panels they stream stay hot in the shared cache — block-major
        // order had every core pulling a different W region simultaneously (the single-core
        // kernel measures ~110 GF/s while 16 threads deliver ~5.6 cores' worth: the multicore
        // wall is W traffic, not the microkernel — 16×6 vs 8×8 measured +3-13% single-core).
        // `OSFKB_CPU_GEMM_PM=0` pins block-major for A/B.
        let panel_major = std::env::var("OSFKB_CPU_GEMM_PM").ok().as_deref() != Some("0");
        (0..nblocks * nchunks).into_par_iter().for_each(|task| {
            let (blk, ch) = if panel_major {
                (task % nblocks, task / nblocks)
            } else {
                (task / nchunks, task % nchunks)
            };
            let (p0, p1) = (ch * chunk, ((ch + 1) * chunk).min(panels));
            if p0 >= p1 {
                return;
            }
            let m0 = blk * MR;
            let rows = MR.min(m - m0);
            let ap = &apack[blk * k * MR..(blk + 1) * k * MR];
            let mut acc = [0f32; MR * NR];
            for panel in p0..p1 {
                let wp = &w.p[panel * k * NR..(panel + 1) * k * NR];
                let j0 = panel * NR;
                let cols = NR.min(n - j0);
                let mut bvec = [0f32; NR];
                if let Some(b) = bias {
                    bvec[..cols].copy_from_slice(&b[j0..j0 + cols]);
                }
                kernel(&mut acc, ap, wp, k, &bvec);
                // SAFETY: task (blk, ch) exclusively owns rows [m0, m0+rows) × cols [j0, j0+cols)
                // of `c` — blocks partition the rows, chunks partition the panels, so no two
                // tasks overlap; `m0+i < m` and `j0+cols <= n` bound every write inside `c[..m*n]`.
                unsafe {
                    let base = cptr.get();
                    for i in 0..rows {
                        std::ptr::copy_nonoverlapping(
                            acc.as_ptr().add(i * NR),
                            base.add((m0 + i) * n + j0),
                            cols,
                        );
                    }
                }
            }
        });
        return;
    }

    c[..m * n]
        .par_chunks_mut(MR * n)
        .enumerate()
        .for_each(|(blk, cblk)| {
            let m0 = blk * MR;
            let rows = cblk.len() / n; // the last block is ragged
            // Pack this row block K-major: [k][MR]. Costs MR·K, and buys the inner loop a pair of
            // contiguous vector loads per k instead of MR loads strided by K (which would touch MR
            // distinct cache lines every step).
            let mut ap = vec![0f32; k * MR];
            for i in 0..rows {
                let src = &a[(m0 + i) * k..(m0 + i + 1) * k];
                for (kk, &v) in src.iter().enumerate() {
                    ap[kk * MR + i] = v;
                }
            }
            let mut acc = [0f32; MR * NR];
            for panel in 0..panels {
                let wp = &w.p[panel * k * NR..(panel + 1) * k * NR];
                let j0 = panel * NR;
                let cols = NR.min(n - j0);
                // The bias SEEDS the accumulators rather than being added on the way out. Adding it
                // in the store loop cost a branch and two bounds checks per output element — 1.5 M of
                // them for one MLP GEMM — where seeding is two vector loads per panel, once.
                let mut bvec = [0f32; NR];
                if let Some(b) = bias {
                    bvec[..cols].copy_from_slice(&b[j0..j0 + cols]);
                }
                kernel(&mut acc, &ap, wp, k, &bvec);
                for i in 0..rows {
                    cblk[i * n + j0..i * n + j0 + cols]
                        .copy_from_slice(&acc[i * NR..i * NR + cols]);
                }
            }
        });
}

/// INT8 prepacked weight for the aarch64 `i8mm` path (`OSFKB_CPU_I8=1` + runtime-detected
/// `smmla`): per-OUTPUT-ROW absmax quantization at pack time, panels laid out for the 2×8
/// operand shape `vmmlaq_s32` consumes — one instruction is a 2×2×8 int8 matmul-accumulate,
/// 16 MACs where the f32 FMA does 4. This is llama.cpp's proven Arm int8 shape, adapted to
/// this module's prepack-once contract. NOT parity-class (dynamic per-row activation quant):
/// opt-in behind its own cosine gate, exactly like the GPU's f16a band.
pub struct PackedWeightI8 {
    /// `[n_panels][k8][NR][8]` i8: panel = NR(8) output rows; within a k-chunk of 8, row pairs
    /// sit as the 2×8 blocks smmla wants.
    q: Vec<i8>,
    /// Per-output-row dequant scale (absmax/127).
    scales: Vec<f32>,
    n: usize,
    k: usize,
}

impl PackedWeightI8 {
    /// Quantize + pack a `[n, k]` row-major f32 weight. `k` is padded to a multiple of 8 with
    /// zeros (zero int8 contributes nothing).
    pub fn new(w: &[f32], n: usize, k: usize) -> Self {
        assert_eq!(w.len(), n * k, "weight is not [n, k]");
        let k8 = k.div_ceil(8);
        let panels = n.div_ceil(NR);
        let mut scales = vec![0f32; n];
        let mut q = vec![0i8; panels * k8 * NR * 8];
        for row in 0..n {
            let src = &w[row * k..(row + 1) * k];
            let absmax = src.iter().fold(0f32, |m, v| m.max(v.abs()));
            let s = if absmax > 0.0 { absmax / 127.0 } else { 1.0 };
            scales[row] = s;
            let inv = 1.0 / s;
            let (panel, j) = (row / NR, row % NR);
            for (kk, &v) in src.iter().enumerate() {
                let c = kk / 8;
                let l = kk % 8;
                q[((panel * k8 + c) * NR + j) * 8 + l] =
                    (v * inv).round().clamp(-127.0, 127.0) as i8;
            }
        }
        Self { q, scales, n, k }
    }
}

/// `c[M, N] = dequant( qa[M, K] · qwᵀ ) (+ bias)` on the i8mm units. The activation is
/// quantized PER ROW on the fly (absmax/127 — the dynamic half of the contract); accumulation
/// is exact int32, dequant is one f32 multiply per output. K padded to 8 like the weights.
///
/// # Safety-free API: the `i8mm` requirement is enforced by [`i8mm_available`] at the caller.
pub fn gemm_i8(c: &mut [f32], a: &[f32], w: &PackedWeightI8, m: usize, bias: Option<&[f32]>) {
    let (n, k) = (w.n, w.k);
    assert_eq!(a.len(), m * k, "lhs is not [m, k]");
    assert!(c.len() >= m * n, "output too small");
    let k8 = k.div_ceil(8);
    let kb32 = k.div_ceil(32);
    let panels = n.div_ceil(NR);
    // Quantize A once, packed as [row-pair][k8][2][8] for the smmla A operand. Scales are per
    // (row, 32-k-block) — Q8_0-class: the per-ROW variant measured 0.998 end-to-end cosine on
    // a 12-layer encode (compounding), below the 0.999 contract; per-block holds it.
    let mpairs = m.div_ceil(2);
    let mut qa = vec![0i8; mpairs * k8 * 16];
    let mut sa = vec![0f32; m * kb32];
    qa.par_chunks_mut(k8 * 16)
        .zip(sa.par_chunks_mut(2 * kb32))
        .enumerate()
        .for_each(|(p, (qp, sp))| {
            for r in 0..2usize {
                let row = p * 2 + r;
                if row >= m {
                    break;
                }
                let src = &a[row * k..(row + 1) * k];
                for blk in 0..kb32 {
                    let lo = blk * 32;
                    let hi = (lo + 32).min(k);
                    let absmax = src[lo..hi].iter().fold(0f32, |mx, v| mx.max(v.abs()));
                    let sc = if absmax > 0.0 { absmax / 127.0 } else { 1.0 };
                    sp[r * kb32 + blk] = sc;
                    let inv = 1.0 / sc;
                    for kk in lo..hi {
                        qp[(kk / 8) * 16 + r * 8 + (kk % 8)] =
                            (src[kk] * inv).round().clamp(-127.0, 127.0) as i8;
                    }
                }
            }
        });
    let cptr = SendPtr(c.as_mut_ptr());
    (0..mpairs * panels).into_par_iter().for_each(|task| {
        let mp = task / panels;
        let panel = task % panels;
        let rows = 2.min(m - mp * 2);
        let cols = NR.min(n - panel * NR);
        // Per 32-k window: exact int32 accumulate, then one f32 fold per (row, col) scaled by
        // the row-block activation scale — the weight scale applies once at the end.
        let mut facc = [0f32; 16];
        for blk in 0..kb32 {
            let c0 = blk * 4;
            let c1 = (c0 + 4).min(k8);
            let ap = &qa[(mp * k8 + c0) * 16..(mp * k8 + c1) * 16];
            let wp = &w.q[(panel * k8 + c0) * NR * 8..(panel * k8 + c1) * NR * 8];
            let mut acc = [0i32; 16];
            #[cfg(target_arch = "aarch64")]
            // SAFETY: caller gated on i8mm_available(); slices sized exactly by construction.
            unsafe {
                kernel_smmla(&mut acc, ap, wp, c1 - c0)
            };
            #[cfg(not(target_arch = "aarch64"))]
            kernel_i8_portable(&mut acc, ap, wp, c1 - c0);
            for r in 0..rows {
                let sblk = sa[(mp * 2 + r) * kb32 + blk];
                for j in 0..NR {
                    facc[r * NR + j] += acc[r * NR + j] as f32 * sblk;
                }
            }
        }
        // SAFETY: task (mp, panel) exclusively owns rows [2mp, 2mp+rows) × cols
        // [panel·NR, +cols) of `c` — the task grid partitions both dimensions.
        unsafe {
            let base = cptr.get();
            for r in 0..rows {
                let row = mp * 2 + r;
                for j in 0..cols {
                    let col = panel * NR + j;
                    let mut v = facc[r * NR + j] * w.scales[col];
                    if let Some(b) = bias {
                        v += b[col];
                    }
                    *base.add(row * n + col) = v;
                }
            }
        }
    });
}

/// Whether this host has the int8 matrix-multiply extension (runtime-detected, cached).
pub fn i8mm_available() -> bool {
    #[cfg(target_arch = "aarch64")]
    {
        static V: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
        return *V.get_or_init(|| std::arch::is_aarch64_feature_detected!("i8mm"));
    }
    #[cfg(not(target_arch = "aarch64"))]
    false
}

/// F16-STORAGE prepacked weight for the aarch64 FEAT_FHM path (`OSFKB_CPU_F16=1` +
/// runtime-detected FMLAL): the f32 panel layout with every element narrowed to IEEE binary16.
/// The widening `fmlal`/`fmlal2` instructions multiply f16 pairs into f32 accumulators at the
/// SAME 4-wide MAC rate as the f32 FMA kernel — but the streamed bytes halve, and the 07-17
/// terminal finding established that the multicore wall is exactly those bytes (16 threads
/// deliver ~5.6 cores' worth of the ~110 GF/s single-core kernel). Prototype measured
/// 1.30–1.40× at the bandwidth-bound encoder shapes, 0.98× at the one cache-resident shape.
/// NOT parity-class (storage rounding): opt-in behind its own cosine gate (`tests/cpu_f16.rs`),
/// exactly like the GPU's f16a band and the int8 arm.
pub struct PackedWeightF16 {
    /// `[n_panels][k][NR]` f16 bits — [`PackedWeight`]'s layout, halved.
    p: Vec<u16>,
    n: usize,
    k: usize,
}

impl PackedWeightF16 {
    /// Narrow + pack a `[n, k]` row-major f32 weight (round-to-nearest-even, same rounding the
    /// activation-side `fcvtn` applies).
    pub fn new(w: &[f32], n: usize, k: usize) -> Self {
        assert_eq!(w.len(), n * k, "weight is not [n, k]");
        let panels = n.div_ceil(NR);
        let mut p = vec![0u16; panels * k * NR];
        for panel in 0..panels {
            let base = panel * k * NR;
            for j in 0..NR {
                let row = panel * NR + j;
                if row >= n {
                    break; // zero-padded tail
                }
                let src = &w[row * k..(row + 1) * k];
                for (kk, &v) in src.iter().enumerate() {
                    p[base + kk * NR + j] = f32_to_f16(v);
                }
            }
        }
        Self { p, n, k }
    }
}

/// `c[M, N] = a[M, K] · wᵀ (+ bias)` through f16 storage. A is transpose-packed like
/// [`gemm_packed`]'s pass 1 and narrowed in the same parallel pass (hardware `fcvtn` — base
/// A64, no feature gate); the FMLAL microkernel widens both streams back into f32
/// accumulators. Task grid, panel-major order and disjoint-tile stores mirror [`gemm_packed`];
/// the per-output k-order is fixed, so the result is bitwise deterministic across runs and
/// thread counts.
pub fn gemm_f16(c: &mut [f32], a: &[f32], w: &PackedWeightF16, m: usize, bias: Option<&[f32]>) {
    let (n, k) = (w.n, w.k);
    assert_eq!(a.len(), m * k, "lhs is not [m, k]");
    assert!(c.len() >= m * n, "output too small");
    if let Some(b) = bias {
        assert_eq!(b.len(), n, "bias width");
    }
    let panels = n.div_ceil(NR);
    let nblocks = m.div_ceil(MR);
    // A-prep mirrors gemm_packed's pass 1 (transpose K-major per block) with the narrow fused
    // on the way out: one contiguous `fcvtn` sweep per block from a REUSED per-worker scratch —
    // no per-block allocation, no extra full pass over A. (A first cut allocated + zeroed the
    // scratch per block and paid ~10% of the whole GEMM for it.)
    let mut apack = vec![0u16; nblocks * k * MR];
    apack.par_chunks_mut(k * MR).enumerate().for_each_init(
        || vec![0f32; k * MR],
        |tmp, (blk, ap)| {
            let m0 = blk * MR;
            let rows = MR.min(m - m0);
            if rows < MR {
                tmp.iter_mut().for_each(|v| *v = 0.0); // ragged tail: unwritten lanes stay 0
            }
            for i in 0..rows {
                let src = &a[(m0 + i) * k..(m0 + i + 1) * k];
                for (kk, &v) in src.iter().enumerate() {
                    tmp[kk * MR + i] = v;
                }
            }
            f32_to_f16_slice(ap, tmp);
        },
    );
    let threads = rayon::current_num_threads().max(1);
    // Chunk like the f32 path by default; `OSFKB_CPU_F16_FLAT=1` pins the one-panel-per-task
    // grid the prototype used (finer stealing granularity — the pending quiet-window A/B).
    let flat = std::env::var("OSFKB_CPU_F16_FLAT").ok().as_deref() == Some("1");
    let nchunks = if flat {
        panels
    } else {
        (3 * threads).div_ceil(nblocks).clamp(1, panels)
    };
    let chunk = panels.div_ceil(nchunks);
    let cptr = SendPtr(c.as_mut_ptr());
    (0..nblocks * nchunks).into_par_iter().for_each(|task| {
        // Panel-major task order — the measured winner for the f32 path (shared-cache W reuse);
        // the f16 stream is half the bytes but the same shape of traffic.
        let (blk, ch) = (task % nblocks, task / nblocks);
        let (p0, p1) = (ch * chunk, ((ch + 1) * chunk).min(panels));
        if p0 >= p1 {
            return;
        }
        let m0 = blk * MR;
        let rows = MR.min(m - m0);
        let ap = &apack[blk * k * MR..(blk + 1) * k * MR];
        let mut acc = [0f32; MR * NR];
        for panel in p0..p1 {
            let wp = &w.p[panel * k * NR..(panel + 1) * k * NR];
            let j0 = panel * NR;
            let cols = NR.min(n - j0);
            let mut bvec = [0f32; NR];
            if let Some(b) = bias {
                bvec[..cols].copy_from_slice(&b[j0..j0 + cols]);
            }
            #[cfg(target_arch = "aarch64")]
            // SAFETY: caller gated on fhm_available(); slices sized exactly by construction.
            unsafe {
                kernel_fmlal(&mut acc, ap, wp, k, &bvec)
            };
            #[cfg(not(target_arch = "aarch64"))]
            kernel_f16_portable(&mut acc, ap, wp, k, &bvec);
            // SAFETY: task (blk, ch) exclusively owns rows [m0, m0+rows) × cols [j0, j0+cols)
            // of `c` — identical tiling to gemm_packed's 2D path.
            unsafe {
                let base = cptr.get();
                for i in 0..rows {
                    std::ptr::copy_nonoverlapping(
                        acc.as_ptr().add(i * NR),
                        base.add((m0 + i) * n + j0),
                        cols,
                    );
                }
            }
        }
    });
}

/// Whether this host has the f16 widening multiply-accumulate extension (FEAT_FHM —
/// `fmlal`/`fmlal2`), runtime-detected and cached.
pub fn fhm_available() -> bool {
    #[cfg(target_arch = "aarch64")]
    {
        static V: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
        return *V.get_or_init(|| std::arch::is_aarch64_feature_detected!("fhm"));
    }
    #[cfg(not(target_arch = "aarch64"))]
    false
}

/// Scalar f32 → IEEE binary16 bits, round-to-nearest-even — bit-identical to the hardware
/// `fcvtn` the slice path uses (both are IEEE-754 RNE; `tests/cpu_f16.rs` asserts the twins).
pub fn f32_to_f16(x: f32) -> u16 {
    let b = x.to_bits();
    let sign = ((b >> 16) & 0x8000) as u16;
    let exp = (b >> 23) & 0xff;
    let man = b & 0x007f_ffff;
    if exp == 0xff {
        // Inf stays inf; NaN keeps its quiet bit and at least one payload bit.
        return sign | 0x7c00 | ((man >> 13) as u16) | u16::from(man != 0);
    }
    let e = exp as i32 - 127 + 15; // rebias
    if e >= 31 {
        return sign | 0x7c00; // overflow → inf
    }
    if e <= 0 {
        if e < -10 {
            return sign; // underflow → signed zero
        }
        // Subnormal: shift the 24-bit significand down, RNE via the +half−1+lsb trick.
        let sig = man | 0x0080_0000;
        let shift = (14 - e) as u32;
        let half = 1u32 << (shift - 1);
        let rounded = sig + half - 1 + ((sig >> shift) & 1);
        return sign | (rounded >> shift) as u16;
    }
    // Normal: drop 13 mantissa bits with RNE; a carry ripples into the exponent (and on to
    // inf) by construction of the packed layout.
    let mut v = ((e as u32) << 10) | (man >> 13);
    let rem = man & 0x1fff;
    if rem > 0x1000 || (rem == 0x1000 && (v & 1) == 1) {
        v += 1;
    }
    if v >= 0x7c00 {
        return sign | 0x7c00;
    }
    sign | v as u16
}

/// IEEE binary16 bits → f32 (exact — every finite f16 is representable).
pub fn f16_to_f32(bits: u16) -> f32 {
    let sign = u32::from(bits >> 15) << 31;
    let exp = u32::from(bits >> 10) & 0x1f;
    let man = u32::from(bits) & 0x3ff;
    if exp == 0x1f {
        return f32::from_bits(sign | 0x7f80_0000 | (man << 13));
    }
    if exp == 0 {
        if man == 0 {
            return f32::from_bits(sign);
        }
        // Subnormal (value = man · 2⁻²⁴): renormalize — the top set bit becomes the implicit 1.
        let shift = man.leading_zeros() - 21;
        let man = (man << shift) & 0x3ff;
        let exp = 127 - 15 + 1 - shift;
        return f32::from_bits(sign | (exp << 23) | (man << 13));
    }
    f32::from_bits(sign | ((exp + 127 - 15) << 23) | (man << 13))
}

/// Narrow a whole f32 slice to f16 bits, RNE. aarch64 runs 8 lanes per `fcvtn`/`fcvtn2` pair
/// (base A64 — no feature gate); the scalar twin serves the tail and every other arch, and is
/// bit-identical (both IEEE RNE).
pub fn f32_to_f16_slice(dst: &mut [u16], src: &[f32]) {
    assert_eq!(dst.len(), src.len());
    // The mutation is inside the aarch64-only asm block; x86 clippy sees it unused.
    #[cfg_attr(not(target_arch = "aarch64"), allow(unused_mut))]
    let mut i = 0;
    #[cfg(target_arch = "aarch64")]
    {
        while i + 8 <= src.len() {
            // SAFETY: 8 f32 reads and 8 u16 writes at offset i, bounds checked by the loop.
            unsafe {
                core::arch::asm!(
                    "ld1 {{v0.4s, v1.4s}}, [{p}]",
                    "fcvtn v2.4h, v0.4s",
                    "fcvtn2 v2.8h, v1.4s",
                    "st1 {{v2.8h}}, [{q}]",
                    p = in(reg) src.as_ptr().add(i),
                    q = in(reg) dst.as_mut_ptr().add(i),
                    out("v0") _, out("v1") _, out("v2") _,
                    options(nostack)
                );
            }
            i += 8;
        }
    }
    for j in i..src.len() {
        dst[j] = f32_to_f16(src[j]);
    }
}

/// The FMLAL microkernel: 8 A-rows × 8 W-cols over the whole k. Per k step: ONE 8×f16 W load,
/// ONE 8×f16 A load, then 16 by-element widening FMAs (`fmlal` covers cols 0–3, `fmlal2`
/// cols 4–7 — the "2" reads the UPPER half of the W register; both spell the source `.4h`).
/// Same 64 MACs/step as the f32 kernel from HALF the loaded bytes. The by-element operand must
/// sit in v0–v15 (`vreg_low16`) — an ISA restriction on half-precision indexed forms.
#[cfg(target_arch = "aarch64")]
#[target_feature(enable = "fhm", enable = "fp16")]
unsafe fn kernel_fmlal(
    acc: &mut [f32; MR * NR],
    ap: &[u16],
    wp: &[u16],
    k: usize,
    bias: &[f32; NR],
) {
    use std::arch::aarch64::*;
    unsafe {
        let blo = vld1q_f32(bias.as_ptr());
        let bhi = vld1q_f32(bias.as_ptr().add(4));
        let mut c: [float32x4_t; 16] = [blo; 16];
        for i in 0..MR {
            c[i * 2 + 1] = bhi;
        }
        let (pa, pw) = (ap.as_ptr(), wp.as_ptr());
        for kk in 0..k {
            let w = vld1q_u16(pw.add(kk * NR));
            let a = vld1q_u16(pa.add(kk * MR));
            core::arch::asm!(
                "fmlal  {c0:v}.4s, {w:v}.4h, {a:v}.h[0]",
                "fmlal2 {c1:v}.4s, {w:v}.4h, {a:v}.h[0]",
                "fmlal  {c2:v}.4s, {w:v}.4h, {a:v}.h[1]",
                "fmlal2 {c3:v}.4s, {w:v}.4h, {a:v}.h[1]",
                "fmlal  {c4:v}.4s, {w:v}.4h, {a:v}.h[2]",
                "fmlal2 {c5:v}.4s, {w:v}.4h, {a:v}.h[2]",
                "fmlal  {c6:v}.4s, {w:v}.4h, {a:v}.h[3]",
                "fmlal2 {c7:v}.4s, {w:v}.4h, {a:v}.h[3]",
                c0 = inout(vreg) c[0], c1 = inout(vreg) c[1],
                c2 = inout(vreg) c[2], c3 = inout(vreg) c[3],
                c4 = inout(vreg) c[4], c5 = inout(vreg) c[5],
                c6 = inout(vreg) c[6], c7 = inout(vreg) c[7],
                w = in(vreg) w, a = in(vreg_low16) a,
                options(pure, nomem, nostack)
            );
            core::arch::asm!(
                "fmlal  {c0:v}.4s, {w:v}.4h, {a:v}.h[4]",
                "fmlal2 {c1:v}.4s, {w:v}.4h, {a:v}.h[4]",
                "fmlal  {c2:v}.4s, {w:v}.4h, {a:v}.h[5]",
                "fmlal2 {c3:v}.4s, {w:v}.4h, {a:v}.h[5]",
                "fmlal  {c4:v}.4s, {w:v}.4h, {a:v}.h[6]",
                "fmlal2 {c5:v}.4s, {w:v}.4h, {a:v}.h[6]",
                "fmlal  {c6:v}.4s, {w:v}.4h, {a:v}.h[7]",
                "fmlal2 {c7:v}.4s, {w:v}.4h, {a:v}.h[7]",
                c0 = inout(vreg) c[8], c1 = inout(vreg) c[9],
                c2 = inout(vreg) c[10], c3 = inout(vreg) c[11],
                c4 = inout(vreg) c[12], c5 = inout(vreg) c[13],
                c6 = inout(vreg) c[14], c7 = inout(vreg) c[15],
                w = in(vreg) w, a = in(vreg_low16) a,
                options(pure, nomem, nostack)
            );
        }
        for i in 0..MR {
            vst1q_f32(acc.as_mut_ptr().add(i * NR), c[i * 2]);
            vst1q_f32(acc.as_mut_ptr().add(i * NR + 4), c[i * 2 + 1]);
        }
    }
}

/// Reference twin of [`kernel_fmlal`] for non-aarch64 builds: same widen-to-f32 semantics,
/// same accumulation order per output.
#[cfg(not(target_arch = "aarch64"))]
fn kernel_f16_portable(
    acc: &mut [f32; MR * NR],
    ap: &[u16],
    wp: &[u16],
    k: usize,
    bias: &[f32; NR],
) {
    for i in 0..MR {
        for j in 0..NR {
            acc[i * NR + j] = bias[j];
        }
    }
    for kk in 0..k {
        for i in 0..MR {
            let av = f16_to_f32(ap[kk * MR + i]);
            for j in 0..NR {
                acc[i * NR + j] += av * f16_to_f32(wp[kk * NR + j]);
            }
        }
    }
}

/// The smmla microkernel: 2 A-rows × 8 W-rows over the whole k, 4 accumulators of 2×2 int32.
/// Per k-chunk of 8: one A register (2×8 i8) against four W registers (each 2 output rows × 8).
#[cfg(target_arch = "aarch64")]
#[target_feature(enable = "i8mm")]
unsafe fn kernel_smmla(acc: &mut [i32; 16], ap: &[i8], wp: &[i8], k8: usize) {
    use std::arch::aarch64::*;
    unsafe {
        let mut c0 = vdupq_n_s32(0); // rows 0-1 × cols 0-1
        let mut c1 = vdupq_n_s32(0); // rows 0-1 × cols 2-3
        let mut c2 = vdupq_n_s32(0); // rows 0-1 × cols 4-5
        let mut c3 = vdupq_n_s32(0); // rows 0-1 × cols 6-7
        let (pa, pw) = (ap.as_ptr(), wp.as_ptr());
        for c in 0..k8 {
            let a = vld1q_s8(pa.add(c * 16)); // [2 rows][8 k]
            let w0 = vld1q_s8(pw.add(c * 64)); // cols 0-1
            let w1 = vld1q_s8(pw.add(c * 64 + 16)); // cols 2-3
            let w2 = vld1q_s8(pw.add(c * 64 + 32));
            let w3 = vld1q_s8(pw.add(c * 64 + 48));
            // `vmmlaq_s32` is nightly-gated (stdarch_neon_i8mm); the instruction itself is
            // stable via asm. Each smmla: 2×2 int32 += A[2×8 i8] · B[2×8 i8]ᵀ — 16 MACs.
            core::arch::asm!(
                "smmla {c0:v}.4s, {a:v}.16b, {w0:v}.16b",
                "smmla {c1:v}.4s, {a:v}.16b, {w1:v}.16b",
                "smmla {c2:v}.4s, {a:v}.16b, {w2:v}.16b",
                "smmla {c3:v}.4s, {a:v}.16b, {w3:v}.16b",
                c0 = inout(vreg) c0,
                c1 = inout(vreg) c1,
                c2 = inout(vreg) c2,
                c3 = inout(vreg) c3,
                a = in(vreg) a,
                w0 = in(vreg) w0,
                w1 = in(vreg) w1,
                w2 = in(vreg) w2,
                w3 = in(vreg) w3,
                options(pure, nomem, nostack)
            );
        }
        // vmmlaq_s32 lane order: [r0c0, r0c1, r1c0, r1c1] per 2×2 block.
        let mut lanes = [0i32; 4];
        for (bi, cc) in [c0, c1, c2, c3].into_iter().enumerate() {
            vst1q_s32(lanes.as_mut_ptr(), cc);
            acc[bi * 2] = lanes[0];
            acc[bi * 2 + 1] = lanes[1];
            acc[NR + bi * 2] = lanes[2];
            acc[NR + bi * 2 + 1] = lanes[3];
        }
    }
}

/// Portable twin of the smmla kernel (exact same integer math) — the reference it is tested
/// against, and the non-aarch64 path (int8 GEMM is opt-in; off-arm hosts simply keep f32).
#[cfg_attr(target_arch = "aarch64", allow(dead_code))]
fn kernel_i8_portable(acc: &mut [i32; 16], ap: &[i8], wp: &[i8], k8: usize) {
    for c in 0..k8 {
        for r in 0..2usize {
            for j in 0..NR {
                let mut s = 0i32;
                for l in 0..8 {
                    s += ap[c * 16 + r * 8 + l] as i32 * wp[(c * NR + j) * 8 + l] as i32;
                }
                acc[r * NR + j] += s;
            }
        }
    }
}

/// A `*mut f32` that rayon may carry across threads. Soundness lives at the USE site: the 2D
/// tasks write provably disjoint tiles (see the SAFETY comment there). The accessor exists so
/// closures capture the WRAPPER — edition-2024 disjoint capture would otherwise reach through
/// and capture the raw field itself, which is not `Sync`.
struct SendPtr(*mut f32);
impl SendPtr {
    fn get(&self) -> *mut f32 {
        self.0
    }
}
unsafe impl Send for SendPtr {}
unsafe impl Sync for SendPtr {}

/// `acc[MR, NR] = bias + ap[K, MR]ᵀ · wp[K, NR]` — the register-blocked inner loop.
#[inline]
fn kernel(acc: &mut [f32; MR * NR], ap: &[f32], wp: &[f32], k: usize, bias: &[f32; NR]) {
    #[cfg(target_arch = "aarch64")]
    {
        // SAFETY: `ap` is [k, MR] and `wp` is [k, NR], both indexed strictly below `k * MR` / `k *
        // NR`; `acc` is exactly MR*NR. NEON is baseline on aarch64.
        unsafe { kernel_neon(acc, ap, wp, k, bias) }
    }
    #[cfg(target_arch = "x86_64")]
    {
        // AVX2+FMA is runtime-detected (not baseline on x86_64); the verdict is cached. Same
        // per-accumulator k-order as the portable twin, so determinism is preserved — only FMA
        // contraction moves bits, exactly as NEON's `vfmaq` does on aarch64.
        use std::sync::OnceLock;
        static AVX2: OnceLock<bool> = OnceLock::new();
        if *AVX2.get_or_init(|| {
            std::arch::is_x86_feature_detected!("avx2")
                && std::arch::is_x86_feature_detected!("fma")
        }) {
            // SAFETY: bounds as above; the detector just proved avx2+fma.
            unsafe { kernel_avx2(acc, ap, wp, k, bias) }
        } else {
            kernel_portable(acc, ap, wp, k, bias);
        }
    }
    #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
    kernel_portable(acc, ap, wp, k, bias);
}

/// x86_64 twin of [`kernel_neon`]: MR=8 rows × NR=8 columns, NR spanning ONE `__m256`, eight
/// accumulators seeded with the bias, `_mm256_fmadd_ps` against a broadcast A lane per row. Same
/// outer-product form and the same 2 FLOP/byte; ~11 of 16 ymm registers live.
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx2", enable = "fma")]
unsafe fn kernel_avx2(
    acc: &mut [f32; MR * NR],
    ap: &[f32],
    wp: &[f32],
    k: usize,
    bias: &[f32; NR],
) {
    use std::arch::x86_64::*;
    unsafe {
        let b = _mm256_loadu_ps(bias.as_ptr());
        let mut c = [b; MR];
        let (pa, pw) = (ap.as_ptr(), wp.as_ptr());
        for kk in 0..k {
            let w = _mm256_loadu_ps(pw.add(kk * NR));
            let a = pa.add(kk * MR);
            c[0] = _mm256_fmadd_ps(_mm256_set1_ps(*a), w, c[0]);
            c[1] = _mm256_fmadd_ps(_mm256_set1_ps(*a.add(1)), w, c[1]);
            c[2] = _mm256_fmadd_ps(_mm256_set1_ps(*a.add(2)), w, c[2]);
            c[3] = _mm256_fmadd_ps(_mm256_set1_ps(*a.add(3)), w, c[3]);
            c[4] = _mm256_fmadd_ps(_mm256_set1_ps(*a.add(4)), w, c[4]);
            c[5] = _mm256_fmadd_ps(_mm256_set1_ps(*a.add(5)), w, c[5]);
            c[6] = _mm256_fmadd_ps(_mm256_set1_ps(*a.add(6)), w, c[6]);
            c[7] = _mm256_fmadd_ps(_mm256_set1_ps(*a.add(7)), w, c[7]);
        }
        for i in 0..MR {
            _mm256_storeu_ps(acc.as_mut_ptr().add(i * NR), c[i]);
        }
    }
}

#[cfg(target_arch = "aarch64")]
#[inline]
unsafe fn kernel_neon(
    acc: &mut [f32; MR * NR],
    ap: &[f32],
    wp: &[f32],
    k: usize,
    bias: &[f32; NR],
) {
    use std::arch::aarch64::*;
    unsafe {
        // 16 accumulators: 8 rows × 2 vectors of 4 columns. Plus 2 A vectors and 2 W vectors in
        // flight = 20 of the 32 NEON registers, which leaves the scheduler room to run ahead.
        // Seeded with the bias — every row of the tile gets the same one.
        let (b0, b1) = (vld1q_f32(bias.as_ptr()), vld1q_f32(bias.as_ptr().add(4)));
        let mut c = [
            b0, b1, b0, b1, b0, b1, b0, b1, b0, b1, b0, b1, b0, b1, b0, b1,
        ];
        let (pa, pw) = (ap.as_ptr(), wp.as_ptr());
        for kk in 0..k {
            let w0 = vld1q_f32(pw.add(kk * NR));
            let w1 = vld1q_f32(pw.add(kk * NR + 4));
            let a0 = vld1q_f32(pa.add(kk * MR)); // rows 0..4
            let a1 = vld1q_f32(pa.add(kk * MR + 4)); // rows 4..8

            // `vfmaq_laneq_f32(acc, b, a, LANE)` = acc + b · a[LANE]: one instruction broadcasts the
            // row's scalar and fuses the multiply-add. LANE must be a const, hence the unrolling.
            c[0] = vfmaq_laneq_f32(c[0], w0, a0, 0);
            c[1] = vfmaq_laneq_f32(c[1], w1, a0, 0);
            c[2] = vfmaq_laneq_f32(c[2], w0, a0, 1);
            c[3] = vfmaq_laneq_f32(c[3], w1, a0, 1);
            c[4] = vfmaq_laneq_f32(c[4], w0, a0, 2);
            c[5] = vfmaq_laneq_f32(c[5], w1, a0, 2);
            c[6] = vfmaq_laneq_f32(c[6], w0, a0, 3);
            c[7] = vfmaq_laneq_f32(c[7], w1, a0, 3);
            c[8] = vfmaq_laneq_f32(c[8], w0, a1, 0);
            c[9] = vfmaq_laneq_f32(c[9], w1, a1, 0);
            c[10] = vfmaq_laneq_f32(c[10], w0, a1, 1);
            c[11] = vfmaq_laneq_f32(c[11], w1, a1, 1);
            c[12] = vfmaq_laneq_f32(c[12], w0, a1, 2);
            c[13] = vfmaq_laneq_f32(c[13], w1, a1, 2);
            c[14] = vfmaq_laneq_f32(c[14], w0, a1, 3);
            c[15] = vfmaq_laneq_f32(c[15], w1, a1, 3);
        }
        for i in 0..MR {
            vst1q_f32(acc.as_mut_ptr().add(i * NR), c[i * 2]);
            vst1q_f32(acc.as_mut_ptr().add(i * NR + 4), c[i * 2 + 1]);
        }
    }
}

/// Same arithmetic, same k-order, no intrinsics — the reference the NEON kernel is tested against.
#[cfg_attr(target_arch = "aarch64", allow(dead_code))]
fn kernel_portable(acc: &mut [f32; MR * NR], ap: &[f32], wp: &[f32], k: usize, bias: &[f32; NR]) {
    for i in 0..MR {
        acc[i * NR..(i + 1) * NR].copy_from_slice(bias);
    }
    for kk in 0..k {
        for i in 0..MR {
            let av = ap[kk * MR + i];
            for j in 0..NR {
                acc[i * NR + j] += av * wp[kk * NR + j];
            }
        }
    }
}

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

    fn naive(a: &[f32], w: &[f32], m: usize, n: usize, k: usize, bias: Option<&[f32]>) -> Vec<f32> {
        let mut c = vec![0f32; m * n];
        for i in 0..m {
            for j in 0..n {
                let mut s = bias.map_or(0.0, |b| b[j]);
                for kk in 0..k {
                    s += a[i * k + kk] * w[j * k + kk];
                }
                c[i * n + j] = s;
            }
        }
        c
    }

    /// The GEMV path exists only to delete the 7 padding rows the microkernel computes at m == 1,
    /// so it must agree with `gemm_packed` to the BIT — not merely to a tolerance. Anything less
    /// and it would be a second numeric path, which is what parity gates elsewhere in the engine
    /// are written to forbid. Shapes cover ragged n, the serial/parallel threshold, and bias.
    #[test]
    fn gemv_is_bit_identical_to_the_m1_gemm() {
        for (n, k) in [
            (1usize, 1usize),
            (5, 7),
            (8, 16),
            (17, 33),
            (512, 512),   // below the serial threshold (the flow head)
            (3072, 1024), // above it (the backbone's fused qkv)
            (4096, 1024),
        ] {
            let a: Vec<f32> = (0..k)
                .map(|i| ((i * 37 % 101) as f32 - 50.0) / 50.0)
                .collect();
            let w: Vec<f32> = (0..n * k)
                .map(|i| ((i * 61 % 197) as f32 - 98.0) / 98.0)
                .collect();
            let bias: Vec<f32> = (0..n).map(|i| (i % 13) as f32 / 13.0).collect();
            let packed = PackedWeight::new(&w, n, k);
            for b in [None, Some(bias.as_slice())] {
                let mut want = vec![0f32; n];
                gemm_packed(&mut want, &a, &packed, 1, b);
                let mut got = vec![0f32; n];
                gemv_packed(&mut got, &a, &packed, b);
                assert_eq!(
                    want.iter().map(|v| v.to_bits()).collect::<Vec<_>>(),
                    got.iter().map(|v| v.to_bits()).collect::<Vec<_>>(),
                    "gemv diverged from the m=1 gemm at n={n} k={k} bias={}",
                    b.is_some()
                );
            }
        }
    }

    #[test]
    fn matches_a_naive_gemm_across_ragged_shapes() {
        // Ragged M (not a multiple of MR=8) and ragged N (not a multiple of NR=8) are the cases a
        // packed kernel gets wrong: the tail panel is zero-padded, so a bad store would write the
        // padding into `c`, or read past `n`.
        for (m, n, k) in [
            (1usize, 1usize, 1usize),
            (3, 5, 7),
            (8, 8, 16),
            (9, 17, 33),
            (64, 768, 768),
            (13, 3072, 768),
        ] {
            let a: Vec<f32> = (0..m * k)
                .map(|i| ((i % 37) as f32 - 18.0) * 0.031)
                .collect();
            let w: Vec<f32> = (0..n * k)
                .map(|i| ((i % 41) as f32 - 20.0) * 0.017)
                .collect();
            let bias: Vec<f32> = (0..n).map(|i| (i % 7) as f32 * 0.05).collect();
            for b in [None, Some(&bias[..])] {
                let want = naive(&a, &w, m, n, k, b);
                let packed = PackedWeight::new(&w, n, k);
                let mut got = vec![f32::NAN; m * n];
                gemm_packed(&mut got, &a, &packed, m, b);
                let dmax = got
                    .iter()
                    .zip(&want)
                    .map(|(x, y)| (x - y).abs())
                    .fold(0f32, f32::max);
                assert!(
                    dmax < 1e-4,
                    "m={m} n={n} k={k} bias={}: max |Δ| {dmax:.3e}",
                    b.is_some()
                );
            }
        }
    }

    #[test]
    fn is_deterministic() {
        let (m, n, k) = (37usize, 129usize, 71usize);
        let a: Vec<f32> = (0..m * k)
            .map(|i| ((i % 37) as f32 - 18.0) * 0.031)
            .collect();
        let w: Vec<f32> = (0..n * k)
            .map(|i| ((i % 41) as f32 - 20.0) * 0.017)
            .collect();
        let packed = PackedWeight::new(&w, n, k);
        let run = || {
            let mut c = vec![0f32; m * n];
            gemm_packed(&mut c, &a, &packed, m, None);
            c
        };
        assert_eq!(
            run(),
            run(),
            "the CPU path is the GPU's oracle: it must not move"
        );
    }
}

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

    /// Prepacked vs the `gemm` crate at the encoder's production shapes, INTERLEAVED min-of-N.
    ///
    /// The prepack is hoisted out of the timed region for us and left inside it for `gemm`, because
    /// that is the honest comparison: our weight really is packed once at load, and `gemm` really
    /// does repack on every call. That asymmetry IS the optimization.
    #[test]
    #[ignore = "perf probe: cargo test --release cpu_gemm::bench -- --ignored --nocapture"]
    fn beats_the_gemm_crate() {
        for (k, n, tag) in [
            (768usize, 768usize, "qkv/o "),
            (768, 3072, "mlp-up"),
            (3072, 768, "mlp-dn"),
        ] {
            let w: Vec<f32> = (0..k * n).map(|i| ((i % 19) as f32 - 9.0) * 0.01).collect();
            let packed = PackedWeight::new(&w, n, k);
            let bias: Vec<f32> = (0..n).map(|i| (i % 7) as f32 * 0.05).collect();
            eprintln!("\n  [{tag}] K={k} N={n}");
            for m in [16usize, 32, 64, 128, 256, 512] {
                let a: Vec<f32> = (0..m * k)
                    .map(|i| ((i % 23) as f32 - 11.0) * 0.01)
                    .collect();
                let mut y = vec![0f32; m * n];
                let (mut ours, mut theirs) = (f64::MAX, f64::MAX);
                for _ in 0..15 {
                    let t0 = std::time::Instant::now();
                    gemm_packed(&mut y, &a, &packed, m, Some(&bias));
                    ours = ours.min(t0.elapsed().as_secs_f64());

                    let t0 = std::time::Instant::now();
                    // SAFETY: row-major [m,k] × [n,k]ᵀ → [m,n]; buffers are ours and correctly sized.
                    unsafe {
                        gemm::gemm(
                            m,
                            n,
                            k,
                            y.as_mut_ptr(),
                            1,
                            n as isize,
                            false,
                            a.as_ptr(),
                            1,
                            k as isize,
                            w.as_ptr(),
                            k as isize,
                            1,
                            0.0,
                            1.0,
                            false,
                            false,
                            false,
                            gemm::Parallelism::Rayon(rayon::current_num_threads()),
                        );
                    }
                    theirs = theirs.min(t0.elapsed().as_secs_f64());
                }
                let gf = |t: f64| 2.0 * (m * n * k) as f64 / t / 1e9;
                eprintln!(
                    "    M={m:4}  packed {:7.1} GF/s | gemm {:7.1} GF/s  ->  {:.2}x",
                    gf(ours),
                    gf(theirs),
                    theirs / ours,
                );
            }
        }
    }
}