fuel-cpu-backend 0.22.4

// SPDX-License-Identifier: MIT OR Apache-2.0
//! Typed byte-shaped kernels — Phase 7.5 B5.
//!
//! These kernels operate on [`CpuStorageBytes`] (bytes-based CPU
//! storage). They take typed slices via `bytemuck::cast_slice` /
//! `as_slice<T>` / `as_slice_mut<T>`, do the per-element work, and
//! return.
//!
//! These are the per-T monomorphic units that the dispatch wrapper
//! in `fuel_dispatch::dispatch::cpu_wrappers` calls after extracting
//! the `CpuStorageBytes` from a `BackendStorage::Cpu(...)` variant.
//!
//! ## Status
//!
//! B5 shipped the proof-of-concept (`add_f32`); Phase C grows the
//! coverage matrix family-by-family. Today's additions:
//!
//! - elementwise binary `f32`: `add`, `sub`, `mul`, `div`
//! - elementwise unary `f32`: `relu`, `neg`, `sqr`, `sqrt`, `recip`,
//!   `abs`, `tanh`

use crate::byte_storage::CpuStorageBytes;
use crate::chassis::reduction::{Max, Mean, Min, Sum, reduce, reduce_to};
use fuel_ir::{Error, Layout, Result};

/// Verify two byte buffers have matching lengths (used by kernels
/// outside the unary/binary chassis families — softmax, affine,
/// clamp, powi, etc.). The chassis functions in `crate::chassis::*`
/// carry their own length-check logic.
fn check_lens_2(name: &str, a: usize, b: usize) -> Result<()> {
    if a != b {
        return Err(Error::Msg(format!(
            "{name}: byte length mismatch (input={a}, output={b})",
        ))
        .bt());
    }
    Ok(())
}

// =============================================================================
// Elementwise binary kernels (f32)
// =============================================================================

/// Generate a public binary-kernel entry point as a thin thunk over
/// the chassis. Op marker `$Op` lives in [`crate::chassis::binary`]
/// (e.g. `Add`, `Mul`, `Pow`); per-precision math lives there once
/// and the four dtype impls fall out of the chassis's blanket
/// `BinaryOp<T>` impls.
///
/// Pre-chassis surface used closure-based macros (`binary_kernel!`,
/// `binary_f32_kernel!`) carrying the op math inline at each call
/// site. The thunk macro replaces both; per-function docstrings
/// from the old form moved to the op-marker definitions in the
/// chassis module.
macro_rules! binary_thunk {
    ($name:ident, $T:ty, $Op:ident) => {
        pub fn $name(
            lhs: &CpuStorageBytes,
            rhs: &CpuStorageBytes,
            out: &mut CpuStorageBytes,
        ) -> Result<()> {
            crate::chassis::binary::binary::<$T, crate::chassis::binary::$Op>(
                stringify!($name),
                lhs,
                rhs,
                out,
            )
        }
    };
}

binary_thunk!(add_f32, f32, Add);
binary_thunk!(sub_f32, f32, Sub);
binary_thunk!(mul_f32, f32, Mul);
binary_thunk!(div_f32, f32, Div);

binary_thunk!(add_f64, f64, Add);
binary_thunk!(sub_f64, f64, Sub);
binary_thunk!(mul_f64, f64, Mul);
binary_thunk!(div_f64, f64, Div);

// =============================================================================
// Elementwise unary kernels (f32)
// =============================================================================

/// Generate a public unary-kernel entry point as a thin thunk over
/// the chassis. Op marker `$Op` lives in [`crate::chassis::unary`]
/// (e.g. `Relu`, `Sqr`, `GeluTanh`); per-precision math lives there
/// once and the four dtype impls fall out of the chassis's blanket
/// `UnaryOp<T>` impls.
///
/// Pre-chassis surface used closure-based macros (`unary_kernel!`,
/// `unary_f32_kernel!`) carrying the op math inline at each call
/// site. The thunk macro replaces both; per-function docstrings
/// from the old form moved to the op-marker definitions in the
/// chassis module.
macro_rules! unary_thunk {
    ($name:ident, $T:ty, $Op:ident) => {
        pub fn $name(input: &CpuStorageBytes, output: &mut CpuStorageBytes) -> Result<()> {
            crate::chassis::unary::unary::<$T, crate::chassis::unary::$Op>(
                stringify!($name),
                input,
                output,
            )
        }
    };
}

// Base unary set — Relu / Neg / Sqr / Sqrt / Recip / Abs / Tanh /
// Exp / Log / Sin / Cos / Sigmoid / Silu / Step — across f32 + f64.
// The bf16 / f16 mirrors live in their own dtype-grouped section
// below (the chassis's blanket impls route them through f32).
unary_thunk!(relu_f32, f32, Relu);
unary_thunk!(neg_f32, f32, Neg);
unary_thunk!(sqr_f32, f32, Sqr);
unary_thunk!(sqrt_f32, f32, Sqrt);
unary_thunk!(recip_f32, f32, Recip);
unary_thunk!(abs_f32, f32, Abs);
unary_thunk!(tanh_f32, f32, Tanh);
unary_thunk!(exp_f32, f32, Exp);
unary_thunk!(log_f32, f32, Log);
unary_thunk!(sin_f32, f32, Sin);
unary_thunk!(cos_f32, f32, Cos);
unary_thunk!(sigmoid_f32, f32, Sigmoid);
unary_thunk!(silu_f32, f32, Silu);
unary_thunk!(step_f32, f32, Step);

unary_thunk!(relu_f64, f64, Relu);
unary_thunk!(neg_f64, f64, Neg);
unary_thunk!(sqr_f64, f64, Sqr);
unary_thunk!(sqrt_f64, f64, Sqrt);
unary_thunk!(recip_f64, f64, Recip);
unary_thunk!(abs_f64, f64, Abs);
unary_thunk!(tanh_f64, f64, Tanh);
unary_thunk!(exp_f64, f64, Exp);
unary_thunk!(log_f64, f64, Log);
unary_thunk!(sin_f64, f64, Sin);
unary_thunk!(cos_f64, f64, Cos);
unary_thunk!(sigmoid_f64, f64, Sigmoid);
unary_thunk!(silu_f64, f64, Silu);
unary_thunk!(step_f64, f64, Step);

// f64 versions of the binary extrema (Maximum/Minimum). The f32
// versions sit later in the file (in the scalar/clamp/extrema
// section); placing the f64 mirrors here keeps the elementwise
// arithmetic block contiguous.
binary_thunk!(maximum_f64, f64, Maximum);
binary_thunk!(minimum_f64, f64, Minimum);

// =============================================================================
// bf16 / f16 elementwise — via-f32 round-trip
// =============================================================================
//
// `bf16` and `f16` lack native transcendentals (no `.exp()`,
// `.sin()`, etc. on `half::{bf16, f16}`), and even arithmetic via
// the native `Add`/`Sub`/`Mul`/`Div` impls truncates intermediate
// values. The standard CPU pattern — what NumPy/PyTorch do for
// these dtypes — is to widen each element to `f32`, do the work,
// and narrow back. Performance is sub-optimal (3 conversions per
// op) but correctness-first; vendor backends override this once
// they're wired.

binary_thunk!(add_bf16, half::bf16, Add);
binary_thunk!(sub_bf16, half::bf16, Sub);
binary_thunk!(mul_bf16, half::bf16, Mul);
binary_thunk!(div_bf16, half::bf16, Div);
binary_thunk!(maximum_bf16, half::bf16, Maximum);
binary_thunk!(minimum_bf16, half::bf16, Minimum);

unary_thunk!(relu_bf16, half::bf16, Relu);
unary_thunk!(neg_bf16, half::bf16, Neg);
unary_thunk!(sqr_bf16, half::bf16, Sqr);
unary_thunk!(sqrt_bf16, half::bf16, Sqrt);
unary_thunk!(recip_bf16, half::bf16, Recip);
unary_thunk!(abs_bf16, half::bf16, Abs);
unary_thunk!(tanh_bf16, half::bf16, Tanh);
unary_thunk!(exp_bf16, half::bf16, Exp);
unary_thunk!(log_bf16, half::bf16, Log);
unary_thunk!(sin_bf16, half::bf16, Sin);
unary_thunk!(cos_bf16, half::bf16, Cos);
unary_thunk!(sigmoid_bf16, half::bf16, Sigmoid);
unary_thunk!(silu_bf16, half::bf16, Silu);
unary_thunk!(step_bf16, half::bf16, Step);
unary_thunk!(gelu_bf16, half::bf16, GeluTanh);

// f16 mirrors of the bf16 set above. Identical patterns; only the
// concrete type differs.

binary_thunk!(add_f16, half::f16, Add);
binary_thunk!(sub_f16, half::f16, Sub);
binary_thunk!(mul_f16, half::f16, Mul);
binary_thunk!(div_f16, half::f16, Div);
binary_thunk!(maximum_f16, half::f16, Maximum);
binary_thunk!(minimum_f16, half::f16, Minimum);

unary_thunk!(relu_f16, half::f16, Relu);
unary_thunk!(neg_f16, half::f16, Neg);
unary_thunk!(sqr_f16, half::f16, Sqr);
unary_thunk!(sqrt_f16, half::f16, Sqrt);
unary_thunk!(recip_f16, half::f16, Recip);
unary_thunk!(abs_f16, half::f16, Abs);
unary_thunk!(tanh_f16, half::f16, Tanh);
unary_thunk!(exp_f16, half::f16, Exp);
unary_thunk!(log_f16, half::f16, Log);
unary_thunk!(sin_f16, half::f16, Sin);
unary_thunk!(cos_f16, half::f16, Cos);
unary_thunk!(sigmoid_f16, half::f16, Sigmoid);
unary_thunk!(silu_f16, half::f16, Silu);
unary_thunk!(step_f16, half::f16, Step);
unary_thunk!(gelu_f16, half::f16, GeluTanh);

// =============================================================================
// Elementwise comparison kernels — typed input, U8 output
// =============================================================================
//
// Comparison ops break the same-dtype-as-input convention the binary/
// unary kernels follow: the output is always a `U8` mask (`1` where
// the predicate holds, `0` otherwise). The wrapper signature is still
// `(lhs: &CpuStorageBytes, rhs: &CpuStorageBytes, out: &mut CpuStorageBytes)`
// — same as `binary_kernel!` — but the kernel internally casts the
// inputs to `&[T]` and the output to `&mut [u8]`. This mirrors the
// ArgMax pattern (typed input, U32 output) and is what the binding-
// table dispatch keys on `[T, T, U8]` resolve to.

/// Generate a public compare-kernel entry point as a thin thunk
/// over the chassis. Op marker `$Op` lives in
/// [`crate::chassis::compare`] (`Eq`, `Ne`, `Lt`, `Le`, `Gt`,
/// `Ge`); per-precision predicate lives there once and the four
/// dtype impls fall out of the chassis's blanket `CompareOp<T>`
/// impls (which also handle bool → u8 conversion).
///
/// Output is `&mut [u8]`; the caller pre-allocates it sized as
/// one byte per element (not byte-equal to input lengths because
/// inputs are wider than u8). The chassis function validates
/// element counts (not byte counts) accordingly.
macro_rules! compare_thunk {
    ($name:ident, $T:ty, $Op:ident) => {
        pub fn $name(
            lhs: &CpuStorageBytes,
            rhs: &CpuStorageBytes,
            out: &mut CpuStorageBytes,
        ) -> Result<()> {
            crate::chassis::compare::compare::<$T, crate::chassis::compare::$Op>(
                stringify!($name),
                lhs,
                rhs,
                out,
            )
        }
    };
}

compare_thunk!(eq_f32_u8, f32, Eq);
compare_thunk!(eq_f64_u8, f64, Eq);
compare_thunk!(eq_bf16_u8, half::bf16, Eq);
compare_thunk!(eq_f16_u8, half::f16, Eq);

compare_thunk!(ne_f32_u8, f32, Ne);
compare_thunk!(ne_f64_u8, f64, Ne);
compare_thunk!(ne_bf16_u8, half::bf16, Ne);
compare_thunk!(ne_f16_u8, half::f16, Ne);

compare_thunk!(lt_f32_u8, f32, Lt);
compare_thunk!(lt_f64_u8, f64, Lt);
compare_thunk!(lt_bf16_u8, half::bf16, Lt);
compare_thunk!(lt_f16_u8, half::f16, Lt);

compare_thunk!(le_f32_u8, f32, Le);
compare_thunk!(le_f64_u8, f64, Le);
compare_thunk!(le_bf16_u8, half::bf16, Le);
compare_thunk!(le_f16_u8, half::f16, Le);

compare_thunk!(gt_f32_u8, f32, Gt);
compare_thunk!(gt_f64_u8, f64, Gt);
compare_thunk!(gt_bf16_u8, half::bf16, Gt);
compare_thunk!(gt_f16_u8, half::f16, Gt);

compare_thunk!(ge_f32_u8, f32, Ge);
compare_thunk!(ge_f64_u8, f64, Ge);
compare_thunk!(ge_bf16_u8, half::bf16, Ge);
compare_thunk!(ge_f16_u8, half::f16, Ge);

// =============================================================================
// Ternary select (Op::Where) — U8 cond + T × T → T
// =============================================================================
//
// Where is the natural consumer of the U8 masks the comparison family
// produces: `out[i] = if cond[i] != 0 { a[i] } else { b[i] }`. Cond is
// always U8; the other two operands and the output share dtype T. The
// kernel takes 4 byte buffers (cond, a, b, out) — the new
// `where_kernel!` macro generalizes the binary_kernel pattern to this
// shape.

macro_rules! where_kernel {
    ($name:ident, $T:ty, $doc:literal) => {
        #[doc = $doc]
        pub fn $name(
            cond: &CpuStorageBytes,
            a: &CpuStorageBytes,
            b: &CpuStorageBytes,
            out: &mut CpuStorageBytes,
        ) -> Result<()> {
            let cond_view: &[u8] = cond.as_slice()?;
            let a_view: &[$T] = a.as_slice()?;
            let b_view: &[$T] = b.as_slice()?;
            let out_view: &mut [$T] = out.as_slice_mut()?;
            if cond_view.len() != a_view.len()
                || cond_view.len() != b_view.len()
                || cond_view.len() != out_view.len()
            {
                return Err(Error::Msg(format!(
                    "{}: element count mismatch (cond={}, a={}, b={}, out={})",
                    stringify!($name),
                    cond_view.len(),
                    a_view.len(),
                    b_view.len(),
                    out_view.len(),
                ))
                .bt());
            }
            for (i, slot) in out_view.iter_mut().enumerate() {
                *slot = if cond_view[i] != 0 {
                    a_view[i]
                } else {
                    b_view[i]
                };
            }
            Ok(())
        }
    };
}

where_kernel!(
    where_f32,
    f32,
    "Ternary select on `f32`: `out[i] = if cond[i] != 0 { a[i] } else { b[i] }`. Cond is `U8`."
);
where_kernel!(where_f64, f64, "Ternary select on `f64`.");
where_kernel!(where_bf16, half::bf16, "Ternary select on `bf16`.");
where_kernel!(where_f16, half::f16, "Ternary select on `f16`.");

// =============================================================================
// Rounding family (Floor / Ceil / Round) — same-dtype unary kernels
// =============================================================================
//
// Rounding ops are non-differentiable almost everywhere; the GradientRule
// drops their gradient silently. Forward is straightforward: `f32::floor`,
// `f32::ceil`, `f32::round_ties_even` (IEEE roundeven / banker's rounding,
// stable since Rust 1.77). The bf16/f16 variants widen to f32 first since
// the half crate doesn't expose roundeven natively.

unary_thunk!(floor_f32, f32, Floor);
unary_thunk!(floor_f64, f64, Floor);
unary_thunk!(floor_bf16, half::bf16, Floor);
unary_thunk!(floor_f16, half::f16, Floor);

unary_thunk!(ceil_f32, f32, Ceil);
unary_thunk!(ceil_f64, f64, Ceil);
unary_thunk!(ceil_bf16, half::bf16, Ceil);
unary_thunk!(ceil_f16, half::f16, Ceil);

// Round uses **banker's rounding** (IEEE 754 roundeven / round-half-to-
// even), via stable `f32::round_ties_even` / `f64::round_ties_even`
// (Rust 1.77+). This is the convention NumPy's `np.round` uses by
// default and matches PyTorch's `torch.round`. It differs from
// `f32::round` (half-away-from-zero) at exact halves: 0.5 → 0 (not 1),
// 1.5 → 2, 2.5 → 2 (not 3), 3.5 → 4, etc.
unary_thunk!(round_f32, f32, Round);
unary_thunk!(round_f64, f64, Round);
unary_thunk!(round_bf16, half::bf16, Round);
unary_thunk!(round_f16, half::f16, Round);

unary_thunk!(sign_f32, f32, Sign);
unary_thunk!(sign_f64, f64, Sign);
unary_thunk!(sign_bf16, half::bf16, Sign);
unary_thunk!(sign_f16, half::f16, Sign);

unary_thunk!(erf_f32, f32, Erf);
unary_thunk!(erf_f64, f64, Erf);
unary_thunk!(erf_bf16, half::bf16, Erf);
unary_thunk!(erf_f16, half::f16, Erf);

unary_thunk!(gelu_erf_f32, f32, GeluErf);
unary_thunk!(gelu_erf_f64, f64, GeluErf);
unary_thunk!(gelu_erf_bf16, half::bf16, GeluErf);
unary_thunk!(gelu_erf_f16, half::f16, GeluErf);

// Binary pow with real exponent. `f32::powf` / `f64::powf` are
// stdlib; bf16/f16 widen to f32. NaN follows IEEE-754 (e.g.
// `pow(-2, 0.5) = NaN`).
binary_thunk!(pow_f32, f32, Pow);
binary_thunk!(pow_f64, f64, Pow);
binary_thunk!(pow_bf16, half::bf16, Pow);
binary_thunk!(pow_f16, half::f16, Pow);

// Reciprocal square root. `1 / sqrt(x)` directly; `Float::sqrt`
// followed by reciprocal is bit-equivalent on IEEE-754, but expressing
// it as one op (rather than Sqrt + Recip) saves one kernel launch and
// matches the RMSNorm pattern that drives the demand. Negative inputs
// give NaN per IEEE.
unary_thunk!(rsqrt_f32, f32, Rsqrt);
unary_thunk!(rsqrt_f64, f64, Rsqrt);
unary_thunk!(rsqrt_bf16, half::bf16, Rsqrt);
unary_thunk!(rsqrt_f16, half::f16, Rsqrt);

// Remainder, PyTorch convention: `a - floor(a/b) * b`. Sign of result
// matches sign of divisor — distinct from `f32::%` (C99 fmod, sign of
// dividend) and from `f32::rem_euclid` (always non-negative). Picked
// to match `torch.remainder` for cross-framework compatibility.
binary_thunk!(rem_f32, f32, Rem);
binary_thunk!(rem_f64, f64, Rem);
binary_thunk!(rem_bf16, half::bf16, Rem);
binary_thunk!(rem_f16, half::f16, Rem);

// =============================================================================
// Flip / Roll along one dim — dtype-agnostic byte reorder
// =============================================================================
//
// Both ops walk a flat (outer × dim × inner) view of the input and
// re-permute the `dim` axis. Bytes-only — `dtype_size` carries the
// per-element width so one kernel covers every dtype.

/// Flip the order of elements along one dim.
/// `out[outer, j, inner] = in[outer, dim_size - 1 - j, inner]`.
pub fn flip_cpu(
    input: &CpuStorageBytes,
    out: &mut CpuStorageBytes,
    outer: usize,
    dim_size: usize,
    inner: usize,
    dtype_size: usize,
) -> Result<()> {
    let row_bytes = inner * dtype_size;
    let in_b = input.bytes();
    let needed = outer * dim_size * row_bytes;
    if in_b.len() != needed || out.len_bytes() != needed {
        return Err(Error::Msg(format!(
            "flip_cpu: byte length mismatch (in={}, out={}, expected={needed})",
            in_b.len(),
            out.len_bytes(),
        ))
        .bt());
    }
    let mut out_buf = vec![0u8; needed];
    for o in 0..outer {
        for j in 0..dim_size {
            let src_j = dim_size - 1 - j;
            let in_off = (o * dim_size + src_j) * row_bytes;
            let out_off = (o * dim_size + j) * row_bytes;
            out_buf[out_off..out_off + row_bytes]
                .copy_from_slice(&in_b[in_off..in_off + row_bytes]);
        }
    }
    out.bytes_mut().copy_from_slice(&out_buf);
    Ok(())
}

/// Cyclic shift along one dim by `shift` positions.
/// `out[outer, j, inner] = in[outer, (j - shift) mod dim_size, inner]`.
/// Positive `shift` moves elements to higher indices; negative the
/// opposite. Always wraps.
pub fn roll_cpu(
    input: &CpuStorageBytes,
    out: &mut CpuStorageBytes,
    outer: usize,
    dim_size: usize,
    inner: usize,
    shift: i64,
    dtype_size: usize,
) -> Result<()> {
    let row_bytes = inner * dtype_size;
    let in_b = input.bytes();
    let needed = outer * dim_size * row_bytes;
    if in_b.len() != needed || out.len_bytes() != needed {
        return Err(Error::Msg(format!(
            "roll_cpu: byte length mismatch (in={}, out={}, expected={needed})",
            in_b.len(),
            out.len_bytes(),
        ))
        .bt());
    }
    if dim_size == 0 {
        return Ok(());
    }
    let d = dim_size as i64;
    // Normalize `shift` into [0, dim_size) once.
    let s = ((shift % d) + d) % d; // Python-style modulo
    let mut out_buf = vec![0u8; needed];
    for o in 0..outer {
        for j in 0..dim_size {
            // out[j] = in[(j - s) mod d]
            let src_j = (j as i64 - s).rem_euclid(d) as usize;
            let in_off = (o * dim_size + src_j) * row_bytes;
            let out_off = (o * dim_size + j) * row_bytes;
            out_buf[out_off..out_off + row_bytes]
                .copy_from_slice(&in_b[in_off..in_off + row_bytes]);
        }
    }
    out.bytes_mut().copy_from_slice(&out_buf);
    Ok(())
}

// =============================================================================
// Pad — multi-dim, dtype-agnostic (Constant mode)
// =============================================================================
//
// Two-pass: (1) fill the entire output with the dtype-pre-encoded
// `fill_bytes` pattern; (2) copy the input region into the unpadded
// slot. Walks input multi-indices via a mixed-radix counter — clean
// for arbitrary rank without recursion or per-dtype specialization.

/// Multi-dim Pad-Constant. `fill_bytes` is one element's worth of
/// bytes in the output dtype; the kernel just byte-copies that
/// pattern into every padded slot, then overwrites the input region
/// with input bytes (also dtype-agnostic).
pub fn pad_const_cpu(
    input: &CpuStorageBytes,
    out: &mut CpuStorageBytes,
    in_shape: &[usize],
    out_shape: &[usize],
    padding: &[(usize, usize)],
    dtype_size: usize,
    fill_bytes: &[u8],
) -> Result<()> {
    let rank = in_shape.len();
    if out_shape.len() != rank || padding.len() != rank {
        return Err(Error::Msg(format!(
            "pad_const_cpu: rank mismatch (in_shape={rank}, out_shape={}, padding={})",
            out_shape.len(),
            padding.len(),
        ))
        .bt());
    }
    if fill_bytes.len() != dtype_size {
        return Err(Error::Msg(format!(
            "pad_const_cpu: fill_bytes length {} != dtype_size {dtype_size}",
            fill_bytes.len(),
        ))
        .bt());
    }
    let in_elem: usize = in_shape.iter().product();
    let out_elem: usize = out_shape.iter().product();
    let in_b = input.bytes();
    if in_b.len() != in_elem * dtype_size || out.len_bytes() != out_elem * dtype_size {
        return Err(Error::Msg(format!(
            "pad_const_cpu: byte length mismatch (in={}, expected={}; out={}, expected={})",
            in_b.len(),
            in_elem * dtype_size,
            out.len_bytes(),
            out_elem * dtype_size,
        ))
        .bt());
    }
    let out_buf = out.bytes_mut();
    // Pass 1: fill every output element with the constant.
    for i in 0..out_elem {
        let off = i * dtype_size;
        out_buf[off..off + dtype_size].copy_from_slice(fill_bytes);
    }
    // Pass 2: copy input into the unpadded region. Walk input
    // multi-indices via a mixed-radix counter; translate each to an
    // output flat index by adding the per-axis `before` offset.
    let in_strides = compute_row_major_strides(in_shape);
    let out_strides = compute_row_major_strides(out_shape);
    let mut idx = vec![0usize; rank];
    for _ in 0..in_elem {
        let in_flat: usize = idx.iter().zip(&in_strides).map(|(&i, &s)| i * s).sum();
        let out_flat: usize = idx
            .iter()
            .zip(&out_strides)
            .zip(padding.iter())
            .map(|((&i, &s), &(b, _))| (i + b) * s)
            .sum();
        let in_off = in_flat * dtype_size;
        let out_off = out_flat * dtype_size;
        out_buf[out_off..out_off + dtype_size].copy_from_slice(&in_b[in_off..in_off + dtype_size]);
        // Advance multi-index (last axis varies fastest).
        for k in (0..rank).rev() {
            idx[k] += 1;
            if idx[k] < in_shape[k] {
                break;
            }
            idx[k] = 0;
        }
    }
    Ok(())
}

fn compute_row_major_strides(shape: &[usize]) -> Vec<usize> {
    let n = shape.len();
    let mut s = vec![1usize; n];
    for i in (0..n.saturating_sub(1)).rev() {
        s[i] = s[i + 1] * shape[i + 1];
    }
    s
}

/// Reflect padding (without repeating the edge). Per-axis mapping:
/// for output axis `j` and input dim size `n`, let `i = j - before`.
/// - `i < 0`: out = in[-i]
/// - `0 <= i < n`: out = `in[i]`
/// - `i >= n`: out = in[2*(n-1) - i]
///
/// Caller must validate `before <= n-1` and `after <= n-1` for every
/// axis (the dispatch wrapper does this); otherwise the reflection
/// would run off the other side and produce garbage indices.
pub fn pad_reflect_cpu(
    input: &CpuStorageBytes,
    out: &mut CpuStorageBytes,
    in_shape: &[usize],
    out_shape: &[usize],
    padding: &[(usize, usize)],
    dtype_size: usize,
) -> Result<()> {
    pad_walk_cpu(
        input,
        out,
        in_shape,
        out_shape,
        padding,
        dtype_size,
        reflect_index,
    )
}

/// Replicate (edge-repeat) padding. Per-axis mapping:
/// - `i < 0`:    out = `in[0]`
/// - `0..n`:     out = `in[i]`
/// - `i >= n`:   out = in[n-1]
pub fn pad_replicate_cpu(
    input: &CpuStorageBytes,
    out: &mut CpuStorageBytes,
    in_shape: &[usize],
    out_shape: &[usize],
    padding: &[(usize, usize)],
    dtype_size: usize,
) -> Result<()> {
    pad_walk_cpu(
        input,
        out,
        in_shape,
        out_shape,
        padding,
        dtype_size,
        replicate_index,
    )
}

fn reflect_index(i: i64, n: usize) -> usize {
    if i < 0 {
        (-i) as usize
    } else if (i as usize) >= n {
        2 * (n - 1) - (i as usize)
    } else {
        i as usize
    }
}

fn replicate_index(i: i64, n: usize) -> usize {
    if i < 0 {
        0
    } else if (i as usize) >= n {
        n - 1
    } else {
        i as usize
    }
}

/// Generic multi-dim Pad walker. Walks every output position; for
/// each position computes the corresponding input position via the
/// per-axis `map_index` function and copies bytes. Used by both
/// Reflect and Replicate kernels (Constant has its own simpler
/// two-pass shape since the padded slots don't index into input).
fn pad_walk_cpu<F>(
    input: &CpuStorageBytes,
    out: &mut CpuStorageBytes,
    in_shape: &[usize],
    out_shape: &[usize],
    padding: &[(usize, usize)],
    dtype_size: usize,
    map_index: F,
) -> Result<()>
where
    F: Fn(i64, usize) -> usize,
{
    let rank = in_shape.len();
    if out_shape.len() != rank || padding.len() != rank {
        return Err(Error::Msg(format!(
            "pad_walk_cpu: rank mismatch (in_shape={rank}, out_shape={}, padding={})",
            out_shape.len(),
            padding.len(),
        ))
        .bt());
    }
    let in_elem: usize = in_shape.iter().product();
    let out_elem: usize = out_shape.iter().product();
    let in_b = input.bytes();
    if in_b.len() != in_elem * dtype_size || out.len_bytes() != out_elem * dtype_size {
        return Err(Error::Msg(format!(
            "pad_walk_cpu: byte length mismatch (in={}, expected={}; out={}, expected={})",
            in_b.len(),
            in_elem * dtype_size,
            out.len_bytes(),
            out_elem * dtype_size,
        ))
        .bt());
    }
    let in_strides = compute_row_major_strides(in_shape);
    let out_strides = compute_row_major_strides(out_shape);
    let out_buf = out.bytes_mut();
    let mut out_idx = vec![0usize; rank];
    for _ in 0..out_elem {
        // Compute input flat offset via per-axis mapping.
        let mut in_off: usize = 0;
        for k in 0..rank {
            let i = out_idx[k] as i64 - padding[k].0 as i64;
            let m = map_index(i, in_shape[k]);
            in_off += m * in_strides[k];
        }
        let out_flat: usize = out_idx.iter().zip(&out_strides).map(|(&i, &s)| i * s).sum();
        let in_byte = in_off * dtype_size;
        let out_byte = out_flat * dtype_size;
        out_buf[out_byte..out_byte + dtype_size]
            .copy_from_slice(&in_b[in_byte..in_byte + dtype_size]);
        for k in (0..rank).rev() {
            out_idx[k] += 1;
            if out_idx[k] < out_shape[k] {
                break;
            }
            out_idx[k] = 0;
        }
    }
    Ok(())
}

// =============================================================================
// Pad backward — gradient accumulation per mode (per-dtype since adds)
// =============================================================================
//
// For each output position, the forward maps to some input position.
// Backward: walk output positions, accumulate grad_out[j] into
// grad_in[map(j)]. Constant-mode backward is the special case where
// padded positions don't index into input — they're skipped.
//
// Per-dtype because accumulation is typed addition. bf16/f16 widen
// the accumulator to f32 to avoid precision loss when many output
// positions accumulate into the same input slot (e.g. all 4 corners
// of a Replicate-padded image map to the same input corner).

macro_rules! pad_backward_kernel {
    ($name:ident, $T:ty, $zero:expr, $widen:expr, $narrow:expr, $doc:literal) => {
        #[doc = $doc]
        pub fn $name(
            grad_out: &CpuStorageBytes,
            grad_in: &mut CpuStorageBytes,
            in_shape: &[usize],
            out_shape: &[usize],
            padding: &[(usize, usize)],
            mode_tag: u8,
        ) -> Result<()> {
            let rank = in_shape.len();
            if out_shape.len() != rank || padding.len() != rank {
                return Err(Error::Msg(format!("{}: rank mismatch", stringify!($name),)).bt());
            }
            let go: &[$T] = grad_out.as_slice()?;
            let gi: &mut [$T] = grad_in.as_slice_mut()?;
            let in_elem: usize = in_shape.iter().product();
            let out_elem: usize = out_shape.iter().product();
            if go.len() != out_elem || gi.len() != in_elem {
                return Err(Error::Msg(format!(
                    "{}: element count mismatch (grad_out={}, expected={out_elem}; \
                     grad_in={}, expected={in_elem})",
                    stringify!($name),
                    go.len(),
                    gi.len(),
                ))
                .bt());
            }
            let widen: fn($T) -> f64 = $widen;
            let narrow: fn(f64) -> $T = $narrow;
            let in_strides = compute_row_major_strides(in_shape);
            let out_strides = compute_row_major_strides(out_shape);
            // Accumulator scratch widened to f64 (one slot per input
            // element). Initialized to zero; final pass narrows back
            // to T.
            let _ = $zero;
            let mut acc: Vec<f64> = vec![0.0; in_elem];
            let mut out_idx = vec![0usize; rank];
            for _ in 0..out_elem {
                // For each axis, compute either the input index this
                // output position maps to, OR signal "skip" (Constant
                // padded slot).
                let mut in_off: usize = 0;
                let mut skip = false;
                for k in 0..rank {
                    let i = out_idx[k] as i64 - padding[k].0 as i64;
                    let n = in_shape[k];
                    let m = match mode_tag {
                        0 => {
                            // Constant: skip if outside input region.
                            if i < 0 || (i as usize) >= n {
                                skip = true;
                                break;
                            }
                            i as usize
                        }
                        1 => reflect_index(i, n),
                        2 => replicate_index(i, n),
                        _ => {
                            return Err(Error::Msg(format!(
                                "{}: unknown mode_tag {mode_tag}",
                                stringify!($name),
                            ))
                            .bt());
                        }
                    };
                    in_off += m * in_strides[k];
                }
                if !skip {
                    let out_flat: usize =
                        out_idx.iter().zip(&out_strides).map(|(&i, &s)| i * s).sum();
                    acc[in_off] += widen(go[out_flat]);
                }
                for k in (0..rank).rev() {
                    out_idx[k] += 1;
                    if out_idx[k] < out_shape[k] {
                        break;
                    }
                    out_idx[k] = 0;
                }
            }
            for (slot, v) in gi.iter_mut().zip(acc.iter()) {
                *slot = narrow(*v);
            }
            Ok(())
        }
    };
}

pad_backward_kernel!(
    pad_backward_f32,
    f32,
    0.0_f32,
    |x: f32| x as f64,
    |x: f64| x as f32,
    "Pad backward (Constant/Reflect/Replicate) on `f32`."

);
pad_backward_kernel!(
    pad_backward_f64,
    f64,
    0.0_f64,
    |x: f64| x,
    |x: f64| x,
    "Pad backward on `f64`."

);
pad_backward_kernel!(
    pad_backward_bf16,
    half::bf16,
    half::bf16::ZERO,
    |x: half::bf16| x.to_f32() as f64,
    |x: f64| half::bf16::from_f32(x as f32),
    "Pad backward on `bf16` (acc widened to f64)."

);
pad_backward_kernel!(
    pad_backward_f16,
    half::f16,
    half::f16::ZERO,
    |x: half::f16| x.to_f32() as f64,
    |x: f64| half::f16::from_f32(x as f32),
    "Pad backward on `f16` (acc widened to f64)."

);

// =============================================================================
// CumSum — running prefix-sum along one dim, per-dtype
// =============================================================================
//
// `out[outer, j, inner] = sum_{k=0..=j} in[outer, k, inner]`. Per-dtype
// because the accumulation needs typed `+` (unlike Flip/Roll which only
// reorder bytes). The bf16/f16 paths widen each accumulation step to f32
// to avoid catastrophic precision loss on long axes.

macro_rules! cumsum_kernel {
    ($name:ident, $T:ty, $zero:expr, $doc:literal) => {
        #[doc = $doc]
        pub fn $name(
            input: &CpuStorageBytes,
            out: &mut CpuStorageBytes,
            outer: usize,
            dim_size: usize,
            inner: usize,
        ) -> Result<()> {
            let inv: &[$T] = input.as_slice()?;
            let outv: &mut [$T] = out.as_slice_mut()?;
            let needed = outer * dim_size * inner;
            if inv.len() != needed || outv.len() != needed {
                return Err(Error::Msg(format!(
                    "{}: element count mismatch (in={}, out={}, expected={needed})",
                    stringify!($name),
                    inv.len(),
                    outv.len(),
                ))
                .bt());
            }
            for o in 0..outer {
                for i in 0..inner {
                    let mut acc = $zero;
                    for j in 0..dim_size {
                        let pos = (o * dim_size + j) * inner + i;
                        acc += inv[pos];
                        outv[pos] = acc;
                    }
                }
            }
            Ok(())
        }
    };
}

cumsum_kernel!(
    cumsum_f32,
    f32,
    0.0_f32,
    "Running prefix-sum on `f32` along one dim."
);
cumsum_kernel!(
    cumsum_f64,
    f64,
    0.0_f64,
    "Running prefix-sum on `f64` along one dim."
);

// bf16/f16: widen accumulator to f32 to avoid precision loss across long axes.
pub fn cumsum_bf16(
    input: &CpuStorageBytes,
    out: &mut CpuStorageBytes,
    outer: usize,
    dim_size: usize,
    inner: usize,
) -> Result<()> {
    let inv: &[half::bf16] = input.as_slice()?;
    let outv: &mut [half::bf16] = out.as_slice_mut()?;
    let needed = outer * dim_size * inner;
    if inv.len() != needed || outv.len() != needed {
        return Err(Error::Msg(format!(
            "cumsum_bf16: element count mismatch (in={}, out={}, expected={needed})",
            inv.len(),
            outv.len(),
        ))
        .bt());
    }
    for o in 0..outer {
        for i in 0..inner {
            let mut acc: f32 = 0.0;
            for j in 0..dim_size {
                let pos = (o * dim_size + j) * inner + i;
                acc += inv[pos].to_f32();
                outv[pos] = half::bf16::from_f32(acc);
            }
        }
    }
    Ok(())
}

pub fn cumsum_f16(
    input: &CpuStorageBytes,
    out: &mut CpuStorageBytes,
    outer: usize,
    dim_size: usize,
    inner: usize,
) -> Result<()> {
    let inv: &[half::f16] = input.as_slice()?;
    let outv: &mut [half::f16] = out.as_slice_mut()?;
    let needed = outer * dim_size * inner;
    if inv.len() != needed || outv.len() != needed {
        return Err(Error::Msg(format!(
            "cumsum_f16: element count mismatch (in={}, out={}, expected={needed})",
            inv.len(),
            outv.len(),
        ))
        .bt());
    }
    for o in 0..outer {
        for i in 0..inner {
            let mut acc: f32 = 0.0;
            for j in 0..dim_size {
                let pos = (o * dim_size + j) * inner + i;
                acc += inv[pos].to_f32();
                outv[pos] = half::f16::from_f32(acc);
            }
        }
    }
    Ok(())
}

// =============================================================================
// Triu / Tril (dtype-agnostic byte-level masking)
// =============================================================================

/// Triangular masking shared by Triu and Tril. Keeps `x[..., i, j]`
/// when `keep_upper && j as i64 >= i as i64 + diagonal`
/// or `!keep_upper && j as i64 <= i as i64 + diagonal`; zeros
/// everywhere else. Dtype-agnostic at the byte level: the caller
/// supplies `dtype_size_in_bytes` and the kernel walks bytes.
pub fn triangular_cpu(
    input: &CpuStorageBytes,
    out: &mut CpuStorageBytes,
    batch_count: usize,
    rows: usize,
    cols: usize,
    diagonal: i64,
    keep_upper: bool,
    dtype_size: usize,
) -> Result<()> {
    let needed = batch_count * rows * cols * dtype_size;
    let inv = input.bytes();
    let outv = out.bytes_mut();
    if inv.len() != needed || outv.len() != needed {
        return Err(Error::Msg(format!(
            "triangular_cpu: byte length mismatch (in={}, out={}, expected={needed})",
            inv.len(),
            outv.len(),
        ))
        .bt());
    }
    // Zero the output first; then selectively copy kept positions.
    // Dtype-agnostic zero == all-zero bytes for every IEEE-754 / integer
    // dtype Fuel supports.
    outv.fill(0);
    for b in 0..batch_count {
        let base = b * rows * cols * dtype_size;
        for i in 0..rows {
            for j in 0..cols {
                let keep = if keep_upper {
                    (j as i64) >= (i as i64) + diagonal
                } else {
                    (j as i64) <= (i as i64) + diagonal
                };
                if keep {
                    let off = base + (i * cols + j) * dtype_size;
                    outv[off..off + dtype_size].copy_from_slice(&inv[off..off + dtype_size]);
                }
            }
        }
    }
    Ok(())
}

// =============================================================================
// LogSoftmax (per-dtype, last dim)
// =============================================================================

macro_rules! log_softmax_last_dim_kernel {
    ($name:ident, $T:ty, $to_f32:expr, $from_f32:expr, $doc:literal) => {
        #[doc = $doc]
        pub fn $name(
            input: &CpuStorageBytes,
            out: &mut CpuStorageBytes,
            outer: usize,
            last_dim: usize,
        ) -> Result<()> {
            let inv: &[$T] = input.as_slice()?;
            let outv: &mut [$T] = out.as_slice_mut()?;
            let needed = outer * last_dim;
            if inv.len() != needed || outv.len() != needed {
                return Err(Error::Msg(format!(
                    "{}: element count mismatch (in={}, out={}, expected={needed})",
                    stringify!($name),
                    inv.len(),
                    outv.len(),
                ))
                .bt());
            }
            for r in 0..outer {
                let start = r * last_dim;
                // Row max for numerical stability.
                let mut row_max = f32::NEG_INFINITY;
                for i in 0..last_dim {
                    let v: f32 = $to_f32(inv[start + i]);
                    if v > row_max {
                        row_max = v;
                    }
                }
                // sum exp(x - max).
                let mut sum = 0.0f32;
                for i in 0..last_dim {
                    sum += ((($to_f32)(inv[start + i])) - row_max).exp();
                }
                let log_sum = sum.ln();
                for i in 0..last_dim {
                    let v: f32 = $to_f32(inv[start + i]);
                    outv[start + i] = $from_f32(v - row_max - log_sum);
                }
            }
            Ok(())
        }
    };
}

#[inline]
fn id_f32(v: f32) -> f32 {
    v
}
#[inline]
fn id_f32_back(v: f32) -> f32 {
    v
}
#[inline]
fn bf16_to_f32(v: half::bf16) -> f32 {
    v.to_f32()
}
#[inline]
fn f32_to_bf16(v: f32) -> half::bf16 {
    half::bf16::from_f32(v)
}
#[inline]
fn f16_to_f32(v: half::f16) -> f32 {
    v.to_f32()
}
#[inline]
fn f32_to_f16(v: f32) -> half::f16 {
    half::f16::from_f32(v)
}

log_softmax_last_dim_kernel!(
    log_softmax_last_dim_f32,
    f32,
    id_f32,
    id_f32_back,
    "Numerically-stable log-softmax along the last dim — f32."
);

// f64: keep an f64 accumulator path for max range/precision.
pub fn log_softmax_last_dim_f64(
    input: &CpuStorageBytes,
    out: &mut CpuStorageBytes,
    outer: usize,
    last_dim: usize,
) -> Result<()> {
    let inv: &[f64] = input.as_slice()?;
    let outv: &mut [f64] = out.as_slice_mut()?;
    let needed = outer * last_dim;
    if inv.len() != needed || outv.len() != needed {
        return Err(Error::Msg(format!(
            "log_softmax_last_dim_f64: element count mismatch (in={}, out={}, expected={needed})",
            inv.len(),
            outv.len(),
        ))
        .bt());
    }
    for r in 0..outer {
        let start = r * last_dim;
        let mut row_max = f64::NEG_INFINITY;
        for i in 0..last_dim {
            if inv[start + i] > row_max {
                row_max = inv[start + i];
            }
        }
        let mut sum = 0.0f64;
        for i in 0..last_dim {
            sum += (inv[start + i] - row_max).exp();
        }
        let log_sum = sum.ln();
        for i in 0..last_dim {
            outv[start + i] = inv[start + i] - row_max - log_sum;
        }
    }
    Ok(())
}

log_softmax_last_dim_kernel!(
    log_softmax_last_dim_bf16,
    half::bf16,
    bf16_to_f32,
    f32_to_bf16,
    "Numerically-stable log-softmax along the last dim — bf16 (f32 accumulator)."
);
log_softmax_last_dim_kernel!(
    log_softmax_last_dim_f16,
    half::f16,
    f16_to_f32,
    f32_to_f16,
    "Numerically-stable log-softmax along the last dim — f16 (f32 accumulator)."
);

// =============================================================================
// LogSoftmax backward (per-dtype, last dim)
// =============================================================================

macro_rules! log_softmax_last_dim_backward_kernel {
    ($name:ident, $T:ty, $to_f32:expr, $from_f32:expr, $doc:literal) => {
        #[doc = $doc]
        pub fn $name(
            y: &CpuStorageBytes, // forward output
            g: &CpuStorageBytes, // upstream gradient
            out: &mut CpuStorageBytes,
            outer: usize,
            last_dim: usize,
        ) -> Result<()> {
            let yv: &[$T] = y.as_slice()?;
            let gv: &[$T] = g.as_slice()?;
            let outv: &mut [$T] = out.as_slice_mut()?;
            let needed = outer * last_dim;
            if yv.len() != needed || gv.len() != needed || outv.len() != needed {
                return Err(Error::Msg(format!(
                    "{}: element count mismatch (y={}, g={}, out={}, expected={needed})",
                    stringify!($name),
                    yv.len(),
                    gv.len(),
                    outv.len(),
                ))
                .bt());
            }
            for r in 0..outer {
                let start = r * last_dim;
                let mut g_sum = 0.0f32;
                for i in 0..last_dim {
                    g_sum += $to_f32(gv[start + i]);
                }
                for i in 0..last_dim {
                    let g_i: f32 = $to_f32(gv[start + i]);
                    let y_i: f32 = $to_f32(yv[start + i]);
                    outv[start + i] = $from_f32(g_i - y_i.exp() * g_sum);
                }
            }
            Ok(())
        }
    };
}

log_softmax_last_dim_backward_kernel!(
    log_softmax_last_dim_backward_f32,
    f32,
    id_f32,
    id_f32_back,
    "LogSoftmax-last-dim backward — f32."
);

pub fn log_softmax_last_dim_backward_f64(
    y: &CpuStorageBytes,
    g: &CpuStorageBytes,
    out: &mut CpuStorageBytes,
    outer: usize,
    last_dim: usize,
) -> Result<()> {
    let yv: &[f64] = y.as_slice()?;
    let gv: &[f64] = g.as_slice()?;
    let outv: &mut [f64] = out.as_slice_mut()?;
    let needed = outer * last_dim;
    if yv.len() != needed || gv.len() != needed || outv.len() != needed {
        return Err(Error::Msg(format!(
            "log_softmax_last_dim_backward_f64: element count mismatch (y={}, g={}, out={}, expected={needed})",
            yv.len(), gv.len(), outv.len(),
        )).bt());
    }
    for r in 0..outer {
        let start = r * last_dim;
        let mut g_sum = 0.0f64;
        for i in 0..last_dim {
            g_sum += gv[start + i];
        }
        for i in 0..last_dim {
            outv[start + i] = gv[start + i] - yv[start + i].exp() * g_sum;
        }
    }
    Ok(())
}

log_softmax_last_dim_backward_kernel!(
    log_softmax_last_dim_backward_bf16,
    half::bf16,
    bf16_to_f32,
    f32_to_bf16,
    "LogSoftmax-last-dim backward — bf16 (f32 accumulator)."
);
log_softmax_last_dim_backward_kernel!(
    log_softmax_last_dim_backward_f16,
    half::f16,
    f16_to_f32,
    f32_to_f16,
    "LogSoftmax-last-dim backward — f16 (f32 accumulator)."
);

// =============================================================================
// MaskedFill (dtype-agnostic byte-level)
// =============================================================================

/// MaskedFill: copy bytes from `input` to `out` where `mask` is 0;
/// write `fill_bytes` (one element's worth) where `mask` is nonzero.
/// All three tensors have the same logical element count; the input
/// and output are `count * dtype_size` bytes; the mask is `count`
/// bytes (U8).
pub fn masked_fill_cpu(
    input: &CpuStorageBytes,
    mask: &CpuStorageBytes,
    out: &mut CpuStorageBytes,
    fill_bytes: &[u8],
    dtype_size: usize,
) -> Result<()> {
    let inv = input.bytes();
    let mv = mask.bytes();
    let outv = out.bytes_mut();
    if fill_bytes.len() != dtype_size {
        return Err(Error::Msg(format!(
            "masked_fill_cpu: fill_bytes.len() ({}) != dtype_size ({dtype_size})",
            fill_bytes.len(),
        ))
        .bt());
    }
    if inv.len() != outv.len() {
        return Err(Error::Msg(format!(
            "masked_fill_cpu: byte length mismatch (in={}, out={})",
            inv.len(),
            outv.len(),
        ))
        .bt());
    }
    let count = inv.len() / dtype_size;
    if mv.len() != count {
        return Err(Error::Msg(format!(
            "masked_fill_cpu: mask byte length ({}) != element count ({count})",
            mv.len(),
        ))
        .bt());
    }
    #[expect(
        clippy::needless_range_loop,
        reason = "loop bound is a declared element count, not an array len, and `i` also computes the byte offset `i * dtype_size`; enumerate could change the iteration count -- a wrong-number bug in kernel code"
    )]
    for i in 0..count {
        let off = i * dtype_size;
        if mv[i] != 0 {
            outv[off..off + dtype_size].copy_from_slice(fill_bytes);
        } else {
            outv[off..off + dtype_size].copy_from_slice(&inv[off..off + dtype_size]);
        }
    }
    Ok(())
}

// =============================================================================
// Contiguize (dtype-agnostic, byte-level)
// =============================================================================

/// Materialize a contiguous-row-major buffer from a (potentially
/// strided / offset / broadcast) input. The output is a freshly
/// allocated [`CpuStorageBytes`] holding `layout.shape().elem_count() *
/// dtype_size` bytes; element `i` of the output corresponds to
/// the i-th element produced by `layout`'s strided iteration over
/// the input.
///
/// Dtype-agnostic: only `dtype_size` matters; the kernel copies
/// that many bytes per element. Broadcast layouts (stride 0)
/// transparently replicate source elements.
///
/// Used by the pipelined executor's auto-Contiguize pass before
/// kernels that require contiguous input (currently every kernel,
/// since today's kernels assume contiguous f32 walks).
pub fn contiguize_cpu(
    input: &CpuStorageBytes,
    layout: &Layout,
    dtype_size: usize,
) -> Result<CpuStorageBytes> {
    let elem_count = layout.shape().elem_count();
    let total_bytes = elem_count.checked_mul(dtype_size).ok_or_else(|| {
        Error::Msg("contiguize_cpu: elem_count * dtype_size overflow".to_string()).bt()
    })?;
    let mut out = CpuStorageBytes::from_zero_bytes(total_bytes);
    if elem_count == 0 {
        return Ok(out);
    }
    let in_bytes = input.bytes();
    let out_bytes = out.bytes_mut();
    for (out_i, src_elem_off) in layout.strided_index().enumerate() {
        let src_byte_off = src_elem_off.checked_mul(dtype_size).ok_or_else(|| {
            Error::Msg("contiguize_cpu: src byte offset overflow".to_string()).bt()
        })?;
        let dst_byte_off = out_i * dtype_size;
        if src_byte_off + dtype_size > in_bytes.len() {
            return Err(Error::Msg(format!(
                "contiguize_cpu: layout points past input bytes \
                 (src_byte={src_byte_off}, dtype_size={dtype_size}, input_bytes={})",
                in_bytes.len(),
            ))
            .bt());
        }
        out_bytes[dst_byte_off..dst_byte_off + dtype_size]
            .copy_from_slice(&in_bytes[src_byte_off..src_byte_off + dtype_size]);
    }
    Ok(out)
}

// =============================================================================
// Index-add / scatter-add (f32 base + src, U32 indices)
// =============================================================================

/// Index-add along one dim with a rank-1 `u32` index tensor.
/// Inputs:
/// - `base` with shape `[outer, base_dim, inner]` (the destination's
///   prior contents; the kernel copies these into `out` first).
/// - `indices` with shape `[n_indices]` (`u32`). Each
///   `indices[i] ∈ [0, base_dim)` is the destination row.
/// - `src` with shape `[outer, n_indices, inner]`.
///
/// Updates: `out[outer, indices[i], inner] += src[outer, i, inner]`
/// for each `i ∈ 0..n_indices`. Out-of-bounds indices return a
/// typed Error.
pub fn index_add_f32(
    base: &CpuStorageBytes,
    indices: &CpuStorageBytes,
    src: &CpuStorageBytes,
    out: &mut CpuStorageBytes,
    outer_count: usize,
    base_dim_size: usize,
    n_indices: usize,
    inner_count: usize,
) -> Result<()> {
    let elem = std::mem::size_of::<f32>();
    let need_base = outer_count
        .saturating_mul(base_dim_size)
        .saturating_mul(inner_count)
        .saturating_mul(elem);
    let need_idx = n_indices.saturating_mul(std::mem::size_of::<u32>());
    let need_src = outer_count
        .saturating_mul(n_indices)
        .saturating_mul(inner_count)
        .saturating_mul(elem);
    if base.len_bytes() != need_base || out.len_bytes() != need_base {
        return Err(Error::Msg(format!(
            "index_add_f32: base bytes={} or out bytes={} doesn't match outer={outer_count} × base_dim={base_dim_size} × inner={inner_count} × {elem}",
            base.len_bytes(), out.len_bytes(),
        ))
        .bt());
    }
    if indices.len_bytes() != need_idx {
        return Err(Error::Msg(format!(
            "index_add_f32: indices bytes={} doesn't match n_indices={n_indices} × 4",
            indices.len_bytes(),
        ))
        .bt());
    }
    if src.len_bytes() != need_src {
        return Err(Error::Msg(format!(
            "index_add_f32: src bytes={} doesn't match outer={outer_count} × n={n_indices} × inner={inner_count} × {elem}",
            src.len_bytes(),
        ))
        .bt());
    }
    // Copy base into out as the starting point.
    out.bytes_mut().copy_from_slice(base.bytes());
    if n_indices == 0 {
        return Ok(());
    }
    let idx_view: &[u32] = indices.as_slice()?;
    let src_view: &[f32] = src.as_slice()?;
    let out_view: &mut [f32] = out.as_slice_mut()?;
    #[expect(
        clippy::needless_range_loop,
        reason = "loop bound is the declared index count, not an array len, and `i` is reported as the out-of-bounds position in the error path; enumerate could change the iteration count"
    )]
    for i in 0..n_indices {
        let target = idx_view[i] as usize;
        if target >= base_dim_size {
            return Err(Error::Msg(format!(
                "index_add_f32: index {target} at position {i} out of bounds for base dim {base_dim_size}",
            ))
            .bt());
        }
        for outer in 0..outer_count {
            let src_off = (outer * n_indices + i) * inner_count;
            let dst_off = (outer * base_dim_size + target) * inner_count;
            for inner in 0..inner_count {
                out_view[dst_off + inner] += src_view[src_off + inner];
            }
        }
    }
    Ok(())
}

/// Generate an IndexAdd kernel parameterized over native arithmetic
/// type `$T` (f32 / f64). Accumulates in-place using `$T`'s `+=`.
macro_rules! index_add_native_kernel {
    ($name:ident, $T:ty, $T_size:expr) => {
        pub fn $name(
            base: &CpuStorageBytes,
            indices: &CpuStorageBytes,
            src: &CpuStorageBytes,
            out: &mut CpuStorageBytes,
            outer_count: usize,
            base_dim_size: usize,
            n_indices: usize,
            inner_count: usize,
        ) -> Result<()> {
            let elem = $T_size;
            let need_base = outer_count
                .saturating_mul(base_dim_size)
                .saturating_mul(inner_count)
                .saturating_mul(elem);
            let need_idx = n_indices.saturating_mul(std::mem::size_of::<u32>());
            let need_src = outer_count
                .saturating_mul(n_indices)
                .saturating_mul(inner_count)
                .saturating_mul(elem);
            if base.len_bytes() != need_base || out.len_bytes() != need_base {
                return Err(Error::Msg(format!(
                    "{}: base/out bytes don't match outer={outer_count} × base_dim={base_dim_size} × inner={inner_count} × {elem}",
                    stringify!($name),
                ))
                .bt());
            }
            if indices.len_bytes() != need_idx {
                return Err(Error::Msg(format!(
                    "{}: indices bytes={} doesn't match n_indices={n_indices} × 4",
                    stringify!($name), indices.len_bytes(),
                ))
                .bt());
            }
            if src.len_bytes() != need_src {
                return Err(Error::Msg(format!(
                    "{}: src bytes={} doesn't match outer={outer_count} × n={n_indices} × inner={inner_count} × {elem}",
                    stringify!($name), src.len_bytes(),
                ))
                .bt());
            }
            out.bytes_mut().copy_from_slice(base.bytes());
            if n_indices == 0 {
                return Ok(());
            }
            let idx_view: &[u32] = indices.as_slice()?;
            let src_view: &[$T] = src.as_slice()?;
            let out_view: &mut [$T] = out.as_slice_mut()?;
            for i in 0..n_indices {
                let target = idx_view[i] as usize;
                if target >= base_dim_size {
                    return Err(Error::Msg(format!(
                        "{}: index {target} at position {i} out of bounds for base dim {base_dim_size}",
                        stringify!($name),
                    ))
                    .bt());
                }
                for outer in 0..outer_count {
                    let src_off = (outer * n_indices + i) * inner_count;
                    let dst_off = (outer * base_dim_size + target) * inner_count;
                    for inner in 0..inner_count {
                        out_view[dst_off + inner] += src_view[src_off + inner];
                    }
                }
            }
            Ok(())
        }
    };
}

index_add_native_kernel!(index_add_f64, f64, std::mem::size_of::<f64>());

/// IndexAdd kernel for half-float types (`bf16` / `f16`).
/// Accumulates in f32 (widen → +=  → narrow back).
macro_rules! index_add_half_kernel {
    ($name:ident, $T:ty, $T_size:expr) => {
        pub fn $name(
            base: &CpuStorageBytes,
            indices: &CpuStorageBytes,
            src: &CpuStorageBytes,
            out: &mut CpuStorageBytes,
            outer_count: usize,
            base_dim_size: usize,
            n_indices: usize,
            inner_count: usize,
        ) -> Result<()> {
            let elem = $T_size;
            let need_base = outer_count
                .saturating_mul(base_dim_size)
                .saturating_mul(inner_count)
                .saturating_mul(elem);
            let need_idx = n_indices.saturating_mul(std::mem::size_of::<u32>());
            let need_src = outer_count
                .saturating_mul(n_indices)
                .saturating_mul(inner_count)
                .saturating_mul(elem);
            if base.len_bytes() != need_base || out.len_bytes() != need_base {
                return Err(Error::Msg(format!(
                    "{}: base/out bytes don't match",
                    stringify!($name),
                ))
                .bt());
            }
            if indices.len_bytes() != need_idx {
                return Err(
                    Error::Msg(format!("{}: indices bytes mismatch", stringify!($name),)).bt(),
                );
            }
            if src.len_bytes() != need_src {
                return Err(Error::Msg(format!("{}: src bytes mismatch", stringify!($name),)).bt());
            }
            out.bytes_mut().copy_from_slice(base.bytes());
            if n_indices == 0 {
                return Ok(());
            }
            let idx_view: &[u32] = indices.as_slice()?;
            let src_view: &[$T] = src.as_slice()?;
            let out_view: &mut [$T] = out.as_slice_mut()?;
            for i in 0..n_indices {
                let target = idx_view[i] as usize;
                if target >= base_dim_size {
                    return Err(Error::Msg(format!(
                        "{}: index {target} OOB for {base_dim_size}",
                        stringify!($name),
                    ))
                    .bt());
                }
                for outer in 0..outer_count {
                    let src_off = (outer * n_indices + i) * inner_count;
                    let dst_off = (outer * base_dim_size + target) * inner_count;
                    for inner in 0..inner_count {
                        let acc =
                            out_view[dst_off + inner].to_f32() + src_view[src_off + inner].to_f32();
                        out_view[dst_off + inner] = <$T>::from_f32(acc);
                    }
                }
            }
            Ok(())
        }
    };
}

index_add_half_kernel!(
    index_add_bf16,
    half::bf16,
    std::mem::size_of::<half::bf16>()
);
index_add_half_kernel!(index_add_f16, half::f16, std::mem::size_of::<half::f16>());

/// Generate a ScatterAdd kernel parameterized over native arithmetic
/// type `$T` (f32 / f64). Same shape as `scatter_add_f32`.
macro_rules! scatter_add_native_kernel {
    ($name:ident, $T:ty, $T_size:expr) => {
        pub fn $name(
            base: &CpuStorageBytes,
            indices: &CpuStorageBytes,
            src: &CpuStorageBytes,
            out: &mut CpuStorageBytes,
            base_shape: &[usize],
            src_shape: &[usize],
            dim: usize,
        ) -> Result<()> {
            if base_shape.len() != src_shape.len() {
                return Err(Error::Msg(format!(
                    "{}: base rank ({}) != src rank ({})",
                    stringify!($name),
                    base_shape.len(),
                    src_shape.len(),
                ))
                .bt());
            }
            let rank = base_shape.len();
            if dim >= rank {
                return Err(Error::Msg(format!(
                    "{}: dim {dim} out of range for rank {rank}",
                    stringify!($name),
                ))
                .bt());
            }
            for d in 0..rank {
                if d != dim && base_shape[d] != src_shape[d] {
                    return Err(Error::Msg(format!(
                        "{}: base/src disagree at dim {d}",
                        stringify!($name),
                    ))
                    .bt());
                }
            }
            let elem = $T_size;
            let base_total: usize = base_shape.iter().product();
            let src_total: usize = src_shape.iter().product();
            if base.len_bytes() != base_total.saturating_mul(elem)
                || out.len_bytes() != base_total.saturating_mul(elem)
            {
                return Err(
                    Error::Msg(format!("{}: base/out bytes mismatch", stringify!($name),)).bt(),
                );
            }
            if indices.len_bytes() != src_total.saturating_mul(std::mem::size_of::<u32>()) {
                return Err(
                    Error::Msg(format!("{}: indices bytes mismatch", stringify!($name),)).bt(),
                );
            }
            if src.len_bytes() != src_total.saturating_mul(elem) {
                return Err(Error::Msg(format!("{}: src bytes mismatch", stringify!($name),)).bt());
            }
            out.bytes_mut().copy_from_slice(base.bytes());
            if src_total == 0 {
                return Ok(());
            }
            let idx_view: &[u32] = indices.as_slice()?;
            let src_view: &[$T] = src.as_slice()?;
            let out_view: &mut [$T] = out.as_slice_mut()?;
            let mut base_strides = vec![0usize; rank];
            let mut s = 1;
            for d in (0..rank).rev() {
                base_strides[d] = s;
                s *= base_shape[d];
            }
            let mut multi = vec![0usize; rank];
            for f in 0..src_total {
                let mut rem = f;
                for d in (0..rank).rev() {
                    multi[d] = rem % src_shape[d];
                    rem /= src_shape[d];
                }
                let dst_dim_idx = idx_view[f] as usize;
                if dst_dim_idx >= base_shape[dim] {
                    return Err(Error::Msg(format!(
                        "{}: index {dst_dim_idx} OOB",
                        stringify!($name),
                    ))
                    .bt());
                }
                let mut dst_flat = 0;
                for d in 0..rank {
                    let coord = if d == dim { dst_dim_idx } else { multi[d] };
                    dst_flat += coord * base_strides[d];
                }
                out_view[dst_flat] += src_view[f];
            }
            Ok(())
        }
    };
}

scatter_add_native_kernel!(scatter_add_f64, f64, std::mem::size_of::<f64>());

/// ScatterAdd for half-float types — accumulates in f32 then narrows.
macro_rules! scatter_add_half_kernel {
    ($name:ident, $T:ty, $T_size:expr) => {
        pub fn $name(
            base: &CpuStorageBytes,
            indices: &CpuStorageBytes,
            src: &CpuStorageBytes,
            out: &mut CpuStorageBytes,
            base_shape: &[usize],
            src_shape: &[usize],
            dim: usize,
        ) -> Result<()> {
            if base_shape.len() != src_shape.len() {
                return Err(Error::Msg(format!("{}: rank mismatch", stringify!($name),)).bt());
            }
            let rank = base_shape.len();
            if dim >= rank {
                return Err(Error::Msg(format!("{}: dim OOB", stringify!($name),)).bt());
            }
            for d in 0..rank {
                if d != dim && base_shape[d] != src_shape[d] {
                    return Err(Error::Msg(format!(
                        "{}: shape mismatch at dim {d}",
                        stringify!($name),
                    ))
                    .bt());
                }
            }
            let elem = $T_size;
            let base_total: usize = base_shape.iter().product();
            let src_total: usize = src_shape.iter().product();
            if base.len_bytes() != base_total.saturating_mul(elem)
                || out.len_bytes() != base_total.saturating_mul(elem)
            {
                return Err(
                    Error::Msg(format!("{}: base/out bytes mismatch", stringify!($name),)).bt(),
                );
            }
            if indices.len_bytes() != src_total.saturating_mul(std::mem::size_of::<u32>()) {
                return Err(
                    Error::Msg(format!("{}: indices bytes mismatch", stringify!($name),)).bt(),
                );
            }
            if src.len_bytes() != src_total.saturating_mul(elem) {
                return Err(Error::Msg(format!("{}: src bytes mismatch", stringify!($name),)).bt());
            }
            out.bytes_mut().copy_from_slice(base.bytes());
            if src_total == 0 {
                return Ok(());
            }
            let idx_view: &[u32] = indices.as_slice()?;
            let src_view: &[$T] = src.as_slice()?;
            let out_view: &mut [$T] = out.as_slice_mut()?;
            let mut base_strides = vec![0usize; rank];
            let mut s = 1;
            for d in (0..rank).rev() {
                base_strides[d] = s;
                s *= base_shape[d];
            }
            let mut multi = vec![0usize; rank];
            for f in 0..src_total {
                let mut rem = f;
                for d in (0..rank).rev() {
                    multi[d] = rem % src_shape[d];
                    rem /= src_shape[d];
                }
                let dst_dim_idx = idx_view[f] as usize;
                if dst_dim_idx >= base_shape[dim] {
                    return Err(Error::Msg(format!("{}: index OOB", stringify!($name),)).bt());
                }
                let mut dst_flat = 0;
                for d in 0..rank {
                    let coord = if d == dim { dst_dim_idx } else { multi[d] };
                    dst_flat += coord * base_strides[d];
                }
                let acc = out_view[dst_flat].to_f32() + src_view[f].to_f32();
                out_view[dst_flat] = <$T>::from_f32(acc);
            }
            Ok(())
        }
    };
}

scatter_add_half_kernel!(
    scatter_add_bf16,
    half::bf16,
    std::mem::size_of::<half::bf16>()
);
scatter_add_half_kernel!(scatter_add_f16, half::f16, std::mem::size_of::<half::f16>());

/// N-dimensional scatter-add — the functional inverse of
/// [`gather_f32`]. `base_shape` and `src_shape` agree on every
/// dim except `dim`. For each src/indices position `p`, the
/// destination multi-index is the same as `p` except `dim`'s
/// coord is replaced by `indices[p]`, and `src[p]` is added into
/// `out` at that destination.
pub fn scatter_add_f32(
    base: &CpuStorageBytes,
    indices: &CpuStorageBytes,
    src: &CpuStorageBytes,
    out: &mut CpuStorageBytes,
    base_shape: &[usize],
    src_shape: &[usize],
    dim: usize,
) -> Result<()> {
    if base_shape.len() != src_shape.len() {
        return Err(Error::Msg(format!(
            "scatter_add_f32: base rank ({}) != src rank ({})",
            base_shape.len(),
            src_shape.len(),
        ))
        .bt());
    }
    let rank = base_shape.len();
    if dim >= rank {
        return Err(Error::Msg(format!(
            "scatter_add_f32: dim {dim} out of range for rank {rank}",
        ))
        .bt());
    }
    for d in 0..rank {
        if d != dim && base_shape[d] != src_shape[d] {
            return Err(Error::Msg(format!(
                "scatter_add_f32: base and src disagree at dim {d} (base={}, src={}); only `dim`={dim} may differ",
                base_shape[d], src_shape[d],
            ))
            .bt());
        }
    }
    let elem = std::mem::size_of::<f32>();
    let base_total: usize = base_shape.iter().product();
    let src_total: usize = src_shape.iter().product();
    if base.len_bytes() != base_total.saturating_mul(elem)
        || out.len_bytes() != base_total.saturating_mul(elem)
    {
        return Err(Error::Msg(format!(
            "scatter_add_f32: base/out bytes don't match shape {base_shape:?} (f32)",
        ))
        .bt());
    }
    if indices.len_bytes() != src_total.saturating_mul(std::mem::size_of::<u32>()) {
        return Err(Error::Msg(format!(
            "scatter_add_f32: indices bytes don't match src shape {src_shape:?} (u32)",
        ))
        .bt());
    }
    if src.len_bytes() != src_total.saturating_mul(elem) {
        return Err(Error::Msg(format!(
            "scatter_add_f32: src bytes don't match shape {src_shape:?} (f32)",
        ))
        .bt());
    }
    // Copy base into out as the starting point.
    out.bytes_mut().copy_from_slice(base.bytes());
    if src_total == 0 {
        return Ok(());
    }
    let idx_view: &[u32] = indices.as_slice()?;
    let src_view: &[f32] = src.as_slice()?;
    let out_view: &mut [f32] = out.as_slice_mut()?;
    // Base strides (row-major).
    let mut base_strides = vec![0usize; rank];
    let mut s = 1;
    for d in (0..rank).rev() {
        base_strides[d] = s;
        s *= base_shape[d];
    }
    let mut multi = vec![0usize; rank];
    for f in 0..src_total {
        // Decode src multi-index from f (row-major over src_shape).
        let mut rem = f;
        for d in (0..rank).rev() {
            multi[d] = rem % src_shape[d];
            rem /= src_shape[d];
        }
        let dst_dim_idx = idx_view[f] as usize;
        if dst_dim_idx >= base_shape[dim] {
            return Err(Error::Msg(format!(
                "scatter_add_f32: index {dst_dim_idx} at position {f} out of bounds for base dim {} = {}",
                dim, base_shape[dim],
            ))
            .bt());
        }
        // Compose destination flat index in base.
        let mut dst_flat = 0;
        for d in 0..rank {
            let coord = if d == dim { dst_dim_idx } else { multi[d] };
            dst_flat += coord * base_strides[d];
        }
        out_view[dst_flat] += src_view[f];
    }
    Ok(())
}

// =============================================================================
// Index select (f32 source, U32 indices)
// =============================================================================

/// Pick slices from a source tensor along one axis using `u32`
/// indices. Dtype-agnostic: copies `inner_count * dtype_size`
/// bytes per source slab.
///
/// Each output element at `(outer, j, inner)` reads from source
/// at `(outer, indices[j], inner)`. Out-of-bounds indices return
/// a typed error rather than reading garbage.
pub fn index_select_cpu(
    source: &CpuStorageBytes,
    indices: &CpuStorageBytes,
    out: &mut CpuStorageBytes,
    outer_count: usize,
    source_dim_size: usize,
    n_indices: usize,
    inner_count: usize,
    dtype_size: usize,
) -> Result<()> {
    if dtype_size == 0 {
        return Err(Error::Msg("index_select_cpu: dtype_size must be > 0".to_string()).bt());
    }
    let stride_bytes = inner_count.saturating_mul(dtype_size);
    let need_src = outer_count
        .saturating_mul(source_dim_size)
        .saturating_mul(stride_bytes);
    let need_idx = n_indices.saturating_mul(std::mem::size_of::<u32>());
    let need_out = outer_count
        .saturating_mul(n_indices)
        .saturating_mul(stride_bytes);
    if source.len_bytes() != need_src {
        return Err(Error::Msg(format!(
            "index_select_cpu: source bytes={} doesn't match outer={outer_count} × dim={source_dim_size} × inner={inner_count} × dtype_size={dtype_size}",
            source.len_bytes(),
        ))
        .bt());
    }
    if indices.len_bytes() != need_idx {
        return Err(Error::Msg(format!(
            "index_select_cpu: indices bytes={} doesn't match n_indices={n_indices} × 4",
            indices.len_bytes(),
        ))
        .bt());
    }
    if out.len_bytes() != need_out {
        return Err(Error::Msg(format!(
            "index_select_cpu: out bytes={} doesn't match outer={outer_count} × n={n_indices} × inner={inner_count} × dtype_size={dtype_size}",
            out.len_bytes(),
        ))
        .bt());
    }
    let src_bytes = source.bytes();
    let idx_view: &[u32] = indices.as_slice()?;
    let out_bytes = out.bytes_mut();
    #[expect(
        clippy::needless_range_loop,
        reason = "loop bound is the declared index count, not an array len, and `j` is reported as the out-of-bounds position in the error path; enumerate could change the iteration count"
    )]
    for j in 0..n_indices {
        let i = idx_view[j] as usize;
        if i >= source_dim_size {
            return Err(Error::Msg(format!(
                "index_select_cpu: index {i} at position {j} out of bounds for source dim {source_dim_size}",
            ))
            .bt());
        }
        for outer in 0..outer_count {
            let src_off = (outer * source_dim_size + i) * stride_bytes;
            let dst_off = (outer * n_indices + j) * stride_bytes;
            out_bytes[dst_off..dst_off + stride_bytes]
                .copy_from_slice(&src_bytes[src_off..src_off + stride_bytes]);
        }
    }
    Ok(())
}

/// Backward-compat shim — same shape as the prior f32-only kernel.
pub fn index_select_f32(
    source: &CpuStorageBytes,
    indices: &CpuStorageBytes,
    out: &mut CpuStorageBytes,
    outer_count: usize,
    source_dim_size: usize,
    n_indices: usize,
    inner_count: usize,
) -> Result<()> {
    index_select_cpu(
        source,
        indices,
        out,
        outer_count,
        source_dim_size,
        n_indices,
        inner_count,
        std::mem::size_of::<f32>(),
    )
}

// =============================================================================
// Rotary position embedding (f32, rotate_half convention)
// =============================================================================

/// Fused rotary position embedding. Inputs `(x, cos, sin)`:
///
/// - `x` is laid out as `[outer_count, seq, head_dim]` row-major
///   (the leading "..." dims fold into `outer_count`).
/// - `cos` and `sin` are `[seq, head_dim]`, broadcasting across
///   the outer dims.
/// - `head_dim` must be even; `h = head_dim / 2`.
///
/// rotate_half formula:
/// ```text
///   out[..., s, i]     = x[..., s, i]     * cos[s, i]     - x[..., s, i+h] * sin[s, i]
///   out[..., s, i+h]   = x[..., s, i+h]   * cos[s, i+h]   + x[..., s, i]   * sin[s, i+h]
/// ```
/// for `i ∈ 0..h`. Replaces a 9-op decomposition (slice + neg +
/// concat + broadcast_mul + add + ...) with a single fused kernel.
pub fn rope_f32(
    x: &CpuStorageBytes,
    cos: &CpuStorageBytes,
    sin: &CpuStorageBytes,
    out: &mut CpuStorageBytes,
    outer_count: usize,
    seq: usize,
    head_dim: usize,
) -> Result<()> {
    if !head_dim.is_multiple_of(2) {
        return Err(Error::Msg(format!("rope_f32: head_dim ({head_dim}) must be even",)).bt());
    }
    let elem = std::mem::size_of::<f32>();
    let need_x = outer_count
        .saturating_mul(seq)
        .saturating_mul(head_dim)
        .saturating_mul(elem);
    let need_cs = seq.saturating_mul(head_dim).saturating_mul(elem);
    if x.len_bytes() != need_x {
        return Err(Error::Msg(format!(
            "rope_f32: x bytes={} doesn't match outer={outer_count} × seq={seq} × head_dim={head_dim} × {elem}",
            x.len_bytes(),
        ))
        .bt());
    }
    if cos.len_bytes() != need_cs {
        return Err(Error::Msg(format!(
            "rope_f32: cos bytes={} doesn't match seq={seq} × head_dim={head_dim} × {elem}",
            cos.len_bytes(),
        ))
        .bt());
    }
    if sin.len_bytes() != need_cs {
        return Err(Error::Msg(format!(
            "rope_f32: sin bytes={} doesn't match seq={seq} × head_dim={head_dim} × {elem}",
            sin.len_bytes(),
        ))
        .bt());
    }
    if out.len_bytes() != need_x {
        return Err(Error::Msg(format!(
            "rope_f32: out bytes={} doesn't match x shape",
            out.len_bytes(),
        ))
        .bt());
    }
    if seq == 0 || head_dim == 0 {
        return Ok(());
    }
    let x_view: &[f32] = x.as_slice()?;
    let cos_view: &[f32] = cos.as_slice()?;
    let sin_view: &[f32] = sin.as_slice()?;
    let out_view: &mut [f32] = out.as_slice_mut()?;
    let h = head_dim / 2;
    for outer in 0..outer_count {
        for s in 0..seq {
            let x_row_off = (outer * seq + s) * head_dim;
            let cs_row_off = s * head_dim;
            for i in 0..h {
                let x_lo_off = x_row_off + i;
                let x_hi_off = x_row_off + i + h;
                let cs_lo_off = cs_row_off + i;
                let cs_hi_off = cs_row_off + i + h;
                let x_lo = x_view[x_lo_off];
                let x_hi = x_view[x_hi_off];
                let cos_lo = cos_view[cs_lo_off];
                let cos_hi = cos_view[cs_hi_off];
                let sin_lo = sin_view[cs_lo_off];
                let sin_hi = sin_view[cs_hi_off];
                out_view[x_lo_off] = x_lo * cos_lo - x_hi * sin_lo;
                out_view[x_hi_off] = x_hi * cos_hi + x_lo * sin_hi;
            }
        }
    }
    Ok(())
}

/// Generate a half-float Rope kernel parameterized over `$T`.
/// Each element is widened to f32 for the rotate_half computation,
/// then narrowed back when written to the output.
macro_rules! rope_half {
    ($name:ident, $T:ty) => {
        pub fn $name(
            x: &CpuStorageBytes,
            cos: &CpuStorageBytes,
            sin: &CpuStorageBytes,
            out: &mut CpuStorageBytes,
            outer_count: usize,
            seq: usize,
            head_dim: usize,
        ) -> Result<()> {
            if head_dim % 2 != 0 {
                return Err(Error::Msg(format!(
                    "{}: head_dim ({head_dim}) must be even",
                    stringify!($name),
                ))
                .bt());
            }
            let elem = std::mem::size_of::<$T>();
            let need_x = outer_count
                .saturating_mul(seq)
                .saturating_mul(head_dim)
                .saturating_mul(elem);
            let need_cs = seq.saturating_mul(head_dim).saturating_mul(elem);
            if x.len_bytes() != need_x {
                return Err(Error::Msg(format!(
                    "{}: x bytes={} doesn't match outer={outer_count} × seq={seq} × head_dim={head_dim} × {elem}",
                    stringify!($name), x.len_bytes(),
                ))
                .bt());
            }
            if cos.len_bytes() != need_cs {
                return Err(Error::Msg(format!(
                    "{}: cos bytes={} doesn't match seq × head_dim × {elem}",
                    stringify!($name), cos.len_bytes(),
                ))
                .bt());
            }
            if sin.len_bytes() != need_cs {
                return Err(Error::Msg(format!(
                    "{}: sin bytes={} doesn't match seq × head_dim × {elem}",
                    stringify!($name), sin.len_bytes(),
                ))
                .bt());
            }
            if out.len_bytes() != need_x {
                return Err(Error::Msg(format!(
                    "{}: out bytes={} doesn't match x shape",
                    stringify!($name), out.len_bytes(),
                ))
                .bt());
            }
            if seq == 0 || head_dim == 0 {
                return Ok(());
            }
            let x_view: &[$T] = x.as_slice()?;
            let cos_view: &[$T] = cos.as_slice()?;
            let sin_view: &[$T] = sin.as_slice()?;
            let out_view: &mut [$T] = out.as_slice_mut()?;
            let h = head_dim / 2;
            for outer in 0..outer_count {
                for s in 0..seq {
                    let x_row_off = (outer * seq + s) * head_dim;
                    let cs_row_off = s * head_dim;
                    for i in 0..h {
                        let x_lo_off = x_row_off + i;
                        let x_hi_off = x_row_off + i + h;
                        let cs_lo_off = cs_row_off + i;
                        let cs_hi_off = cs_row_off + i + h;
                        let x_lo = x_view[x_lo_off].to_f32();
                        let x_hi = x_view[x_hi_off].to_f32();
                        let cos_lo = cos_view[cs_lo_off].to_f32();
                        let cos_hi = cos_view[cs_hi_off].to_f32();
                        let sin_lo = sin_view[cs_lo_off].to_f32();
                        let sin_hi = sin_view[cs_hi_off].to_f32();
                        out_view[x_lo_off] = <$T>::from_f32(x_lo * cos_lo - x_hi * sin_lo);
                        out_view[x_hi_off] = <$T>::from_f32(x_hi * cos_hi + x_lo * sin_hi);
                    }
                }
            }
            Ok(())
        }
    };
}

rope_half!(rope_bf16, half::bf16);
rope_half!(rope_f16, half::f16);

/// `f64` Rope — native arithmetic, same rotate_half formula as
/// `rope_f32`.
pub fn rope_f64(
    x: &CpuStorageBytes,
    cos: &CpuStorageBytes,
    sin: &CpuStorageBytes,
    out: &mut CpuStorageBytes,
    outer_count: usize,
    seq: usize,
    head_dim: usize,
) -> Result<()> {
    if !head_dim.is_multiple_of(2) {
        return Err(Error::Msg(format!("rope_f64: head_dim ({head_dim}) must be even",)).bt());
    }
    let elem = std::mem::size_of::<f64>();
    let need_x = outer_count
        .saturating_mul(seq)
        .saturating_mul(head_dim)
        .saturating_mul(elem);
    let need_cs = seq.saturating_mul(head_dim).saturating_mul(elem);
    if x.len_bytes() != need_x {
        return Err(Error::Msg(format!(
            "rope_f64: x bytes={} doesn't match shape (outer={outer_count} × seq={seq} × head_dim={head_dim})",
            x.len_bytes(),
        ))
        .bt());
    }
    if cos.len_bytes() != need_cs || sin.len_bytes() != need_cs {
        return Err(Error::Msg(format!(
            "rope_f64: cos/sin bytes don't match (seq × head_dim × {elem})",
        ))
        .bt());
    }
    if out.len_bytes() != need_x {
        return Err(Error::Msg(format!(
            "rope_f64: out bytes={} doesn't match x shape",
            out.len_bytes(),
        ))
        .bt());
    }
    if seq == 0 || head_dim == 0 {
        return Ok(());
    }
    let x_view: &[f64] = x.as_slice()?;
    let cos_view: &[f64] = cos.as_slice()?;
    let sin_view: &[f64] = sin.as_slice()?;
    let out_view: &mut [f64] = out.as_slice_mut()?;
    let h = head_dim / 2;
    for outer in 0..outer_count {
        for s in 0..seq {
            let x_row_off = (outer * seq + s) * head_dim;
            let cs_row_off = s * head_dim;
            for i in 0..h {
                let x_lo = x_view[x_row_off + i];
                let x_hi = x_view[x_row_off + i + h];
                let cos_lo = cos_view[cs_row_off + i];
                let cos_hi = cos_view[cs_row_off + i + h];
                let sin_lo = sin_view[cs_row_off + i];
                let sin_hi = sin_view[cs_row_off + i + h];
                out_view[x_row_off + i] = x_lo * cos_lo - x_hi * sin_lo;
                out_view[x_row_off + i + h] = x_hi * cos_hi + x_lo * sin_hi;
            }
        }
    }
    Ok(())
}

// =============================================================================
// Gather (f32 source, U32 indices, same-rank)
// =============================================================================

/// N-dimensional gather along `dim`. Dtype-agnostic byte-level
/// version — copies `dtype_size` bytes per output element.
pub fn gather_cpu(
    source: &CpuStorageBytes,
    indices: &CpuStorageBytes,
    out: &mut CpuStorageBytes,
    source_shape: &[usize],
    output_shape: &[usize],
    dim: usize,
    dtype_size: usize,
) -> Result<()> {
    if dtype_size == 0 {
        return Err(Error::Msg("gather_cpu: dtype_size must be > 0".to_string()).bt());
    }
    if source_shape.len() != output_shape.len() {
        return Err(Error::Msg(format!(
            "gather_cpu: source rank ({}) != output rank ({})",
            source_shape.len(),
            output_shape.len(),
        ))
        .bt());
    }
    let rank = source_shape.len();
    if dim >= rank {
        return Err(Error::Msg(format!(
            "gather_cpu: dim {dim} out of range for rank {rank}",
        ))
        .bt());
    }
    for d in 0..rank {
        if d != dim && source_shape[d] != output_shape[d] {
            return Err(Error::Msg(format!(
                "gather_cpu: source and output disagree at dim {d} \
                 (source={}, output={}); only `dim`={dim} may differ",
                source_shape[d], output_shape[d],
            ))
            .bt());
        }
    }
    let source_total: usize = source_shape.iter().product();
    let output_total: usize = output_shape.iter().product();
    if source.len_bytes() != source_total.saturating_mul(dtype_size) {
        return Err(Error::Msg(format!(
            "gather_cpu: source bytes={} doesn't match shape {source_shape:?} (dtype_size={dtype_size})",
            source.len_bytes(),
        ))
        .bt());
    }
    if indices.len_bytes() != output_total.saturating_mul(std::mem::size_of::<u32>()) {
        return Err(Error::Msg(format!(
            "gather_cpu: indices bytes={} doesn't match output shape {output_shape:?} (u32)",
            indices.len_bytes(),
        ))
        .bt());
    }
    if out.len_bytes() != output_total.saturating_mul(dtype_size) {
        return Err(Error::Msg(format!(
            "gather_cpu: out bytes={} doesn't match output shape {output_shape:?} (dtype_size={dtype_size})",
            out.len_bytes(),
        ))
        .bt());
    }
    if output_total == 0 {
        return Ok(());
    }
    let src_bytes = source.bytes();
    let idx_view: &[u32] = indices.as_slice()?;
    let out_bytes = out.bytes_mut();
    let mut src_strides = vec![0usize; rank];
    let mut s = 1;
    for d in (0..rank).rev() {
        src_strides[d] = s;
        s *= source_shape[d];
    }
    let mut multi = vec![0usize; rank];
    #[expect(
        clippy::needless_range_loop,
        reason = "flat index drives explicit multi-index reconstruction (rem = f), not a sequential scan; enumerate would obscure the flat<->multi-index mapping and risk an iteration-count mismatch"
    )]
    for f in 0..output_total {
        let mut rem = f;
        for d in (0..rank).rev() {
            multi[d] = rem % output_shape[d];
            rem /= output_shape[d];
        }
        let src_dim_idx = idx_view[f] as usize;
        if src_dim_idx >= source_shape[dim] {
            return Err(Error::Msg(format!(
                "gather_cpu: index {src_dim_idx} at output position {f} out of bounds for source dim {} = {}",
                dim, source_shape[dim],
            ))
            .bt());
        }
        let mut src_flat = 0;
        for d in 0..rank {
            let coord = if d == dim { src_dim_idx } else { multi[d] };
            src_flat += coord * src_strides[d];
        }
        let src_off = src_flat * dtype_size;
        let dst_off = f * dtype_size;
        out_bytes[dst_off..dst_off + dtype_size]
            .copy_from_slice(&src_bytes[src_off..src_off + dtype_size]);
    }
    Ok(())
}

/// Backward-compat shim — calls [`gather_cpu`] with `f32`'s dtype size.
pub fn gather_f32(
    source: &CpuStorageBytes,
    indices: &CpuStorageBytes,
    out: &mut CpuStorageBytes,
    source_shape: &[usize],
    output_shape: &[usize],
    dim: usize,
) -> Result<()> {
    gather_cpu(
        source,
        indices,
        out,
        source_shape,
        output_shape,
        dim,
        std::mem::size_of::<f32>(),
    )
}

// =============================================================================
// Softmax (f32)
// =============================================================================

/// Softmax along the last dim, numerically stable. For each of the
/// `outer_count` rows of `last_dim` elements, computes
/// `out[i] = exp(x[i] - max_row) / sum(exp(x - max_row))`. The
/// `max_row` subtraction prevents overflow in `exp`.
///
/// Edge cases:
///   - `last_dim == 0`: no work; returns `Ok(())`.
///   - All-`-inf` row: would divide 0/0; returns NaN per IEEE-754.
///     Caller is expected to ensure inputs are finite.
pub fn softmax_last_dim_f32(
    input: &CpuStorageBytes,
    out: &mut CpuStorageBytes,
    outer_count: usize,
    last_dim: usize,
) -> Result<()> {
    check_lens_2("softmax_last_dim_f32", input.len_bytes(), out.len_bytes())?;
    let elem = std::mem::size_of::<f32>();
    let need = outer_count.saturating_mul(last_dim).saturating_mul(elem);
    if input.len_bytes() != need {
        return Err(Error::Msg(format!(
            "softmax_last_dim_f32: input bytes={} doesn't match outer={outer_count} × last={last_dim} × {elem}",
            input.len_bytes(),
        ))
        .bt());
    }
    if last_dim == 0 {
        return Ok(());
    }
    let in_view: &[f32] = input.as_slice()?;
    let out_view: &mut [f32] = out.as_slice_mut()?;
    for row in 0..outer_count {
        let off = row * last_dim;
        let row_in = &in_view[off..off + last_dim];
        // Find row max for numerical stability.
        let mut row_max = row_in[0];
        for &v in &row_in[1..] {
            if v > row_max {
                row_max = v;
            }
        }
        // Compute exp(x - max) and accumulate sum.
        let mut sum = 0.0_f32;
        for j in 0..last_dim {
            let e = (row_in[j] - row_max).exp();
            out_view[off + j] = e;
            sum += e;
        }
        // Normalize.
        let inv = 1.0 / sum;
        for j in 0..last_dim {
            out_view[off + j] *= inv;
        }
    }
    Ok(())
}

// =============================================================================
// RMS / Layer norm along the last dim (f32, no affine)
// =============================================================================

/// Helper: validate that input/output bytes match `outer_count *
/// last_dim * sizeof::<f32>()`. Returns Ok if all three agree.
fn check_norm_lens(
    name: &str,
    input: &CpuStorageBytes,
    out: &CpuStorageBytes,
    outer_count: usize,
    last_dim: usize,
) -> Result<()> {
    let elem = std::mem::size_of::<f32>();
    let need = outer_count.saturating_mul(last_dim).saturating_mul(elem);
    if input.len_bytes() != need || out.len_bytes() != need {
        return Err(Error::Msg(format!(
            "{name}: bytes mismatch (input={}, out={}, expected outer={outer_count} × last={last_dim} × {elem} = {need})",
            input.len_bytes(), out.len_bytes(),
        ))
        .bt());
    }
    Ok(())
}

/// RMS normalization along the last dim, no affine params:
/// `out[i] = x[i] / sqrt(mean(x²) + eps)` per row.
/// `eps` is `f64` for graph-API consistency; converted to `f32`
/// before use in the f32 kernel.
pub fn rms_norm_last_dim_f32(
    input: &CpuStorageBytes,
    out: &mut CpuStorageBytes,
    outer_count: usize,
    last_dim: usize,
    eps: f64,
) -> Result<()> {
    check_norm_lens("rms_norm_last_dim_f32", input, out, outer_count, last_dim)?;
    if last_dim == 0 {
        return Ok(());
    }
    let in_view: &[f32] = input.as_slice()?;
    let out_view: &mut [f32] = out.as_slice_mut()?;
    let eps32 = eps as f32;
    let inv_n = 1.0_f32 / last_dim as f32;
    for row in 0..outer_count {
        let off = row * last_dim;
        let row_in = &in_view[off..off + last_dim];
        let mut sum_sq = 0.0_f32;
        for &v in row_in {
            sum_sq += v * v;
        }
        let mean_sq = sum_sq * inv_n;
        let rms_inv = 1.0_f32 / (mean_sq + eps32).sqrt();
        for j in 0..last_dim {
            out_view[off + j] = row_in[j] * rms_inv;
        }
    }
    Ok(())
}

/// Generate a half-float SoftmaxLastDim kernel parameterized
/// over `$T` (`half::bf16` / `half::f16`). All arithmetic happens
/// in f32 (numerically stable max-subtract + exp + sum +
/// normalize); narrowing back to half-float happens only on the
/// final write.
macro_rules! softmax_last_dim_half {
    ($name:ident, $T:ty) => {
        pub fn $name(
            input: &CpuStorageBytes,
            out: &mut CpuStorageBytes,
            outer_count: usize,
            last_dim: usize,
        ) -> Result<()> {
            check_lens_2(stringify!($name), input.len_bytes(), out.len_bytes())?;
            let elem = std::mem::size_of::<$T>();
            let need = outer_count.saturating_mul(last_dim).saturating_mul(elem);
            if input.len_bytes() != need {
                return Err(Error::Msg(format!(
                    "{}: input bytes={} doesn't match outer={outer_count} × last={last_dim} × {elem}",
                    stringify!($name), input.len_bytes(),
                ))
                .bt());
            }
            if last_dim == 0 {
                return Ok(());
            }
            let in_view: &[$T] = input.as_slice()?;
            let out_view: &mut [$T] = out.as_slice_mut()?;
            // Per-row scratch buffer for exp(x - max) values.
            let mut exps = vec![0.0_f32; last_dim];
            for row in 0..outer_count {
                let off = row * last_dim;
                let mut row_max = in_view[off].to_f32();
                for j in 1..last_dim {
                    let v = in_view[off + j].to_f32();
                    if v > row_max {
                        row_max = v;
                    }
                }
                let mut sum = 0.0_f32;
                for j in 0..last_dim {
                    let e = (in_view[off + j].to_f32() - row_max).exp();
                    exps[j] = e;
                    sum += e;
                }
                let inv = 1.0_f32 / sum;
                for j in 0..last_dim {
                    out_view[off + j] = <$T>::from_f32(exps[j] * inv);
                }
            }
            Ok(())
        }
    };
}

softmax_last_dim_half!(softmax_last_dim_bf16, half::bf16);
softmax_last_dim_half!(softmax_last_dim_f16, half::f16);

/// Softmax along the last dim for `f64`. Same numerically-stable
/// algorithm as `softmax_last_dim_f32`; native f64 arithmetic.
pub fn softmax_last_dim_f64(
    input: &CpuStorageBytes,
    out: &mut CpuStorageBytes,
    outer_count: usize,
    last_dim: usize,
) -> Result<()> {
    check_lens_2("softmax_last_dim_f64", input.len_bytes(), out.len_bytes())?;
    let elem = std::mem::size_of::<f64>();
    let need = outer_count.saturating_mul(last_dim).saturating_mul(elem);
    if input.len_bytes() != need {
        return Err(Error::Msg(format!(
            "softmax_last_dim_f64: input bytes={} doesn't match outer={outer_count} × last={last_dim} × {elem}",
            input.len_bytes(),
        ))
        .bt());
    }
    if last_dim == 0 {
        return Ok(());
    }
    let in_view: &[f64] = input.as_slice()?;
    let out_view: &mut [f64] = out.as_slice_mut()?;
    for row in 0..outer_count {
        let off = row * last_dim;
        let row_in = &in_view[off..off + last_dim];
        let mut row_max = row_in[0];
        for &v in &row_in[1..] {
            if v > row_max {
                row_max = v;
            }
        }
        let mut sum = 0.0_f64;
        for j in 0..last_dim {
            let e = (row_in[j] - row_max).exp();
            out_view[off + j] = e;
            sum += e;
        }
        let inv = 1.0_f64 / sum;
        for j in 0..last_dim {
            out_view[off + j] *= inv;
        }
    }
    Ok(())
}

/// Layer normalization along the last dim, no affine params:
/// `out[i] = (x[i] - mean(x)) / sqrt(var(x) + eps)` per row, with
/// `var = mean((x - mean)²)`.
pub fn layer_norm_last_dim_f32(
    input: &CpuStorageBytes,
    out: &mut CpuStorageBytes,
    outer_count: usize,
    last_dim: usize,
    eps: f64,
) -> Result<()> {
    check_norm_lens("layer_norm_last_dim_f32", input, out, outer_count, last_dim)?;
    if last_dim == 0 {
        return Ok(());
    }
    let in_view: &[f32] = input.as_slice()?;
    let out_view: &mut [f32] = out.as_slice_mut()?;
    let eps32 = eps as f32;
    let inv_n = 1.0_f32 / last_dim as f32;
    for row in 0..outer_count {
        let off = row * last_dim;
        let row_in = &in_view[off..off + last_dim];
        let mut sum = 0.0_f32;
        for &v in row_in {
            sum += v;
        }
        let mean = sum * inv_n;
        let mut sum_sq = 0.0_f32;
        for &v in row_in {
            let d = v - mean;
            sum_sq += d * d;
        }
        let var = sum_sq * inv_n;
        let inv_std = 1.0_f32 / (var + eps32).sqrt();
        for j in 0..last_dim {
            out_view[off + j] = (row_in[j] - mean) * inv_std;
        }
    }
    Ok(())
}

/// Generate a half-float RmsNormLastDim kernel parameterized
/// over `$T`. All arithmetic happens in f32 (sum-of-squares,
/// reciprocal sqrt); narrowing back to the half-float type
/// happens only on the final write.
macro_rules! rms_norm_last_dim_half {
    ($name:ident, $T:ty) => {
        pub fn $name(
            input: &CpuStorageBytes,
            out: &mut CpuStorageBytes,
            outer_count: usize,
            last_dim: usize,
            eps: f64,
        ) -> Result<()> {
            check_norm_lens_typed::<$T>(stringify!($name), input, out, outer_count, last_dim)?;
            if last_dim == 0 {
                return Ok(());
            }
            let in_view: &[$T] = input.as_slice()?;
            let out_view: &mut [$T] = out.as_slice_mut()?;
            let eps32 = eps as f32;
            let inv_n = 1.0_f32 / last_dim as f32;
            for row in 0..outer_count {
                let off = row * last_dim;
                let mut sum_sq = 0.0_f32;
                for j in 0..last_dim {
                    let v = in_view[off + j].to_f32();
                    sum_sq += v * v;
                }
                let mean_sq = sum_sq * inv_n;
                let rms_inv = 1.0_f32 / (mean_sq + eps32).sqrt();
                for j in 0..last_dim {
                    out_view[off + j] = <$T>::from_f32(in_view[off + j].to_f32() * rms_inv);
                }
            }
            Ok(())
        }
    };
}

rms_norm_last_dim_half!(rms_norm_last_dim_bf16, half::bf16);
rms_norm_last_dim_half!(rms_norm_last_dim_f16, half::f16);

/// `f64` RMS norm — native arithmetic, same algorithm as the
/// f32 version.
pub fn rms_norm_last_dim_f64(
    input: &CpuStorageBytes,
    out: &mut CpuStorageBytes,
    outer_count: usize,
    last_dim: usize,
    eps: f64,
) -> Result<()> {
    check_norm_lens_typed::<f64>("rms_norm_last_dim_f64", input, out, outer_count, last_dim)?;
    if last_dim == 0 {
        return Ok(());
    }
    let in_view: &[f64] = input.as_slice()?;
    let out_view: &mut [f64] = out.as_slice_mut()?;
    let inv_n = 1.0_f64 / last_dim as f64;
    for row in 0..outer_count {
        let off = row * last_dim;
        let row_in = &in_view[off..off + last_dim];
        let mut sum_sq = 0.0_f64;
        for &v in row_in {
            sum_sq += v * v;
        }
        let mean_sq = sum_sq * inv_n;
        let rms_inv = 1.0_f64 / (mean_sq + eps).sqrt();
        for j in 0..last_dim {
            out_view[off + j] = row_in[j] * rms_inv;
        }
    }
    Ok(())
}

// `check_norm_lens` (used by f32 norm helpers) is dtype-specific
// — it currently bakes in `size_of::<f32>()`. The bf16/f16 norm
// kernels above can't call it. Verify shape locally:
fn check_norm_lens_typed<T>(
    name: &str,
    input: &CpuStorageBytes,
    out: &CpuStorageBytes,
    outer_count: usize,
    last_dim: usize,
) -> Result<()> {
    let elem = std::mem::size_of::<T>();
    let need = outer_count.saturating_mul(last_dim).saturating_mul(elem);
    if input.len_bytes() != need || out.len_bytes() != need {
        return Err(Error::Msg(format!(
            "{name}: bytes mismatch (input={}, out={}, expected outer={outer_count} × last={last_dim} × {elem})",
            input.len_bytes(),
            out.len_bytes(),
        ))
        .bt());
    }
    Ok(())
}

/// Generate a half-float LayerNormLastDim kernel (no affine
/// params) parameterized over `$T`. Two-pass per row in f32.
macro_rules! layer_norm_last_dim_half {
    ($name:ident, $T:ty) => {
        pub fn $name(
            input: &CpuStorageBytes,
            out: &mut CpuStorageBytes,
            outer_count: usize,
            last_dim: usize,
            eps: f64,
        ) -> Result<()> {
            check_norm_lens_typed::<$T>(stringify!($name), input, out, outer_count, last_dim)?;
            if last_dim == 0 {
                return Ok(());
            }
            let in_view: &[$T] = input.as_slice()?;
            let out_view: &mut [$T] = out.as_slice_mut()?;
            let eps32 = eps as f32;
            let inv_n = 1.0_f32 / last_dim as f32;
            for row in 0..outer_count {
                let off = row * last_dim;
                let mut sum = 0.0_f32;
                for j in 0..last_dim {
                    sum += in_view[off + j].to_f32();
                }
                let mean = sum * inv_n;
                let mut sum_sq = 0.0_f32;
                for j in 0..last_dim {
                    let d = in_view[off + j].to_f32() - mean;
                    sum_sq += d * d;
                }
                let var = sum_sq * inv_n;
                let inv_std = 1.0_f32 / (var + eps32).sqrt();
                for j in 0..last_dim {
                    out_view[off + j] =
                        <$T>::from_f32((in_view[off + j].to_f32() - mean) * inv_std);
                }
            }
            Ok(())
        }
    };
}

layer_norm_last_dim_half!(layer_norm_last_dim_bf16, half::bf16);
layer_norm_last_dim_half!(layer_norm_last_dim_f16, half::f16);

/// `f64` layer norm — native arithmetic, same algorithm as the
/// f32 version.
pub fn layer_norm_last_dim_f64(
    input: &CpuStorageBytes,
    out: &mut CpuStorageBytes,
    outer_count: usize,
    last_dim: usize,
    eps: f64,
) -> Result<()> {
    check_norm_lens_typed::<f64>("layer_norm_last_dim_f64", input, out, outer_count, last_dim)?;
    if last_dim == 0 {
        return Ok(());
    }
    let in_view: &[f64] = input.as_slice()?;
    let out_view: &mut [f64] = out.as_slice_mut()?;
    let inv_n = 1.0_f64 / last_dim as f64;
    for row in 0..outer_count {
        let off = row * last_dim;
        let row_in = &in_view[off..off + last_dim];
        let mut sum = 0.0_f64;
        for &v in row_in {
            sum += v;
        }
        let mean = sum * inv_n;
        let mut sum_sq = 0.0_f64;
        for &v in row_in {
            let d = v - mean;
            sum_sq += d * d;
        }
        let var = sum_sq * inv_n;
        let inv_std = 1.0_f64 / (var + eps).sqrt();
        for j in 0..last_dim {
            out_view[off + j] = (row_in[j] - mean) * inv_std;
        }
    }
    Ok(())
}

// Patch the bf16/f16 RmsNorm kernels to use `check_norm_lens_typed`
// — the version above mistakenly called `check_norm_lens` (which
// hardcodes f32). Replace those calls below.

// =============================================================================
// Concat (dtype-agnostic, byte-level)
// =============================================================================

/// Concatenate N inputs along one dim. Dtype-agnostic: the kernel
/// memcpys `inner_count * dtype_size` bytes per slab, so it works
/// uniformly for f32 / f64 / bf16 / f16 / u32 / etc.
///
/// Output layout: `[outer_count, total_dim, inner_count]` row-major
/// where `total_dim = sum(input_dim_sizes)`. Each input contributes
/// a `[outer_count, input_dim_sizes[i], inner_count]` slab.
pub fn concat_cpu(
    inputs: &[&CpuStorageBytes],
    out: &mut CpuStorageBytes,
    outer_count: usize,
    input_dim_sizes: &[usize],
    inner_count: usize,
    dtype_size: usize,
) -> Result<()> {
    if inputs.len() != input_dim_sizes.len() {
        return Err(Error::Msg(format!(
            "concat_cpu: inputs count ({}) != input_dim_sizes len ({})",
            inputs.len(),
            input_dim_sizes.len(),
        ))
        .bt());
    }
    if inputs.is_empty() {
        return Err(Error::Msg("concat_cpu: at least one input required".to_string()).bt());
    }
    if dtype_size == 0 {
        return Err(Error::Msg("concat_cpu: dtype_size must be > 0".to_string()).bt());
    }
    let total_dim: usize = input_dim_sizes.iter().sum();
    let stride_bytes = inner_count.saturating_mul(dtype_size);
    let need_out = outer_count
        .saturating_mul(total_dim)
        .saturating_mul(stride_bytes);
    if out.len_bytes() != need_out {
        return Err(Error::Msg(format!(
            "concat_cpu: out bytes={} doesn't match outer={outer_count} × total_dim={total_dim} × inner={inner_count} × dtype_size={dtype_size}",
            out.len_bytes(),
        ))
        .bt());
    }
    for (i, input) in inputs.iter().enumerate() {
        let need_in = outer_count
            .saturating_mul(input_dim_sizes[i])
            .saturating_mul(stride_bytes);
        if input.len_bytes() != need_in {
            return Err(Error::Msg(format!(
                "concat_cpu: input[{i}] bytes={} doesn't match outer={outer_count} × dim={} × inner={inner_count} × dtype_size={dtype_size}",
                input.len_bytes(),
                input_dim_sizes[i],
            ))
            .bt());
        }
    }
    let out_bytes = out.bytes_mut();
    let mut dim_offset = 0usize;
    for (i, input) in inputs.iter().enumerate() {
        let in_bytes = input.bytes();
        let d_i = input_dim_sizes[i];
        for outer in 0..outer_count {
            for dim_pos in 0..d_i {
                let src_off = (outer * d_i + dim_pos) * stride_bytes;
                let dst_off = (outer * total_dim + dim_offset + dim_pos) * stride_bytes;
                out_bytes[dst_off..dst_off + stride_bytes]
                    .copy_from_slice(&in_bytes[src_off..src_off + stride_bytes]);
            }
        }
        dim_offset += d_i;
    }
    Ok(())
}

/// Backward-compat shim: forwards to [`concat_cpu`] with `f32`'s
/// dtype size. Existing callers (the dispatch wrapper and tests)
/// keep working unchanged.
pub fn concat_f32(
    inputs: &[&CpuStorageBytes],
    out: &mut CpuStorageBytes,
    outer_count: usize,
    input_dim_sizes: &[usize],
    inner_count: usize,
) -> Result<()> {
    concat_cpu(
        inputs,
        out,
        outer_count,
        input_dim_sizes,
        inner_count,
        std::mem::size_of::<f32>(),
    )
}

// =============================================================================
// Scalar / clamp / pow / extrema (f32)
// =============================================================================

/// Affine transformation: `out[i] = mul * input[i] + add`. The
/// pipelined executor maps `Op::AddScalar(c)` as `mul=1, add=c`
/// and `Op::MulScalar(c)` as `mul=c, add=0`, so this single kernel
/// covers both.
pub fn affine_f32(
    input: &CpuStorageBytes,
    out: &mut CpuStorageBytes,
    mul: f32,
    add: f32,
) -> Result<()> {
    check_lens_2("affine_f32", input.len_bytes(), out.len_bytes())?;
    let in_view: &[f32] = input.as_slice()?;
    let out_view: &mut [f32] = out.as_slice_mut()?;
    for (i, slot) in out_view.iter_mut().enumerate() {
        *slot = mul * in_view[i] + add;
    }
    Ok(())
}

// =============================================================================
// In-place elementwise unary kernels — Phase 3e of in-place ops
// =============================================================================
//
// Each kernel mutates the target buffer in place: `out[i] = f(out[i])`.
// The op math reuses the chassis's `UnaryOpCore` impls so the in-place
// path stays in lockstep with the non-inplace cousins (same numerics
// bit-for-bit; new ops added to the chassis are one-line additions here).
// Each kernel takes only `&mut CpuStorageBytes` — the executor's
// `WorkItemKind::InplaceKernel` arm acquires the target's write lock and
// passes it as the sole buffer.

macro_rules! unary_inplace_thunk {
    ($name:ident, $T:ty, $Op:ident) => {
        pub fn $name(out: &mut CpuStorageBytes) -> Result<()> {
            let view: &mut [$T] = out.as_slice_mut()?;
            for slot in view.iter_mut() {
                *slot = <crate::chassis::unary::$Op as crate::chassis::unary::UnaryOp<$T>>::apply(
                    *slot,
                );
            }
            Ok(())
        }
    };
}

unary_inplace_thunk!(relu_inplace_f32, f32, Relu);
unary_inplace_thunk!(silu_inplace_f32, f32, Silu);
unary_inplace_thunk!(gelu_inplace_f32, f32, GeluTanh);
unary_inplace_thunk!(tanh_inplace_f32, f32, Tanh);
unary_inplace_thunk!(sigmoid_inplace_f32, f32, Sigmoid);

unary_inplace_thunk!(relu_inplace_f64, f64, Relu);
unary_inplace_thunk!(silu_inplace_f64, f64, Silu);
unary_inplace_thunk!(gelu_inplace_f64, f64, GeluTanh);
unary_inplace_thunk!(tanh_inplace_f64, f64, Tanh);
unary_inplace_thunk!(sigmoid_inplace_f64, f64, Sigmoid);

unary_inplace_thunk!(relu_inplace_bf16, half::bf16, Relu);
unary_inplace_thunk!(silu_inplace_bf16, half::bf16, Silu);
unary_inplace_thunk!(gelu_inplace_bf16, half::bf16, GeluTanh);
unary_inplace_thunk!(tanh_inplace_bf16, half::bf16, Tanh);
unary_inplace_thunk!(sigmoid_inplace_bf16, half::bf16, Sigmoid);

unary_inplace_thunk!(relu_inplace_f16, half::f16, Relu);
unary_inplace_thunk!(silu_inplace_f16, half::f16, Silu);
unary_inplace_thunk!(gelu_inplace_f16, half::f16, GeluTanh);
unary_inplace_thunk!(tanh_inplace_f16, half::f16, Tanh);
unary_inplace_thunk!(sigmoid_inplace_f16, half::f16, Sigmoid);

// In-place unary op family expansion (2026-05-30). Each (op × dtype)
// reuses the chassis's `UnaryOp<T>::apply` blanket impl, identical
// numerics to the non-inplace cousin.
unary_inplace_thunk!(neg_inplace_f32, f32, Neg);
unary_inplace_thunk!(neg_inplace_f64, f64, Neg);
unary_inplace_thunk!(neg_inplace_bf16, half::bf16, Neg);
unary_inplace_thunk!(neg_inplace_f16, half::f16, Neg);

unary_inplace_thunk!(abs_inplace_f32, f32, Abs);
unary_inplace_thunk!(abs_inplace_f64, f64, Abs);
unary_inplace_thunk!(abs_inplace_bf16, half::bf16, Abs);
unary_inplace_thunk!(abs_inplace_f16, half::f16, Abs);

unary_inplace_thunk!(sqr_inplace_f32, f32, Sqr);
unary_inplace_thunk!(sqr_inplace_f64, f64, Sqr);
unary_inplace_thunk!(sqr_inplace_bf16, half::bf16, Sqr);
unary_inplace_thunk!(sqr_inplace_f16, half::f16, Sqr);

unary_inplace_thunk!(sqrt_inplace_f32, f32, Sqrt);
unary_inplace_thunk!(sqrt_inplace_f64, f64, Sqrt);
unary_inplace_thunk!(sqrt_inplace_bf16, half::bf16, Sqrt);
unary_inplace_thunk!(sqrt_inplace_f16, half::f16, Sqrt);

unary_inplace_thunk!(rsqrt_inplace_f32, f32, Rsqrt);
unary_inplace_thunk!(rsqrt_inplace_f64, f64, Rsqrt);
unary_inplace_thunk!(rsqrt_inplace_bf16, half::bf16, Rsqrt);
unary_inplace_thunk!(rsqrt_inplace_f16, half::f16, Rsqrt);

unary_inplace_thunk!(recip_inplace_f32, f32, Recip);
unary_inplace_thunk!(recip_inplace_f64, f64, Recip);
unary_inplace_thunk!(recip_inplace_bf16, half::bf16, Recip);
unary_inplace_thunk!(recip_inplace_f16, half::f16, Recip);

unary_inplace_thunk!(exp_inplace_f32, f32, Exp);
unary_inplace_thunk!(exp_inplace_f64, f64, Exp);
unary_inplace_thunk!(exp_inplace_bf16, half::bf16, Exp);
unary_inplace_thunk!(exp_inplace_f16, half::f16, Exp);

unary_inplace_thunk!(log_inplace_f32, f32, Log);
unary_inplace_thunk!(log_inplace_f64, f64, Log);
unary_inplace_thunk!(log_inplace_bf16, half::bf16, Log);
unary_inplace_thunk!(log_inplace_f16, half::f16, Log);

unary_inplace_thunk!(sin_inplace_f32, f32, Sin);
unary_inplace_thunk!(sin_inplace_f64, f64, Sin);
unary_inplace_thunk!(sin_inplace_bf16, half::bf16, Sin);
unary_inplace_thunk!(sin_inplace_f16, half::f16, Sin);

unary_inplace_thunk!(cos_inplace_f32, f32, Cos);
unary_inplace_thunk!(cos_inplace_f64, f64, Cos);
unary_inplace_thunk!(cos_inplace_bf16, half::bf16, Cos);
unary_inplace_thunk!(cos_inplace_f16, half::f16, Cos);

unary_inplace_thunk!(sign_inplace_f32, f32, Sign);
unary_inplace_thunk!(sign_inplace_f64, f64, Sign);
unary_inplace_thunk!(sign_inplace_bf16, half::bf16, Sign);
unary_inplace_thunk!(sign_inplace_f16, half::f16, Sign);

unary_inplace_thunk!(floor_inplace_f32, f32, Floor);
unary_inplace_thunk!(floor_inplace_f64, f64, Floor);
unary_inplace_thunk!(floor_inplace_bf16, half::bf16, Floor);
unary_inplace_thunk!(floor_inplace_f16, half::f16, Floor);

unary_inplace_thunk!(ceil_inplace_f32, f32, Ceil);
unary_inplace_thunk!(ceil_inplace_f64, f64, Ceil);
unary_inplace_thunk!(ceil_inplace_bf16, half::bf16, Ceil);
unary_inplace_thunk!(ceil_inplace_f16, half::f16, Ceil);

unary_inplace_thunk!(round_inplace_f32, f32, Round);
unary_inplace_thunk!(round_inplace_f64, f64, Round);
unary_inplace_thunk!(round_inplace_bf16, half::bf16, Round);
unary_inplace_thunk!(round_inplace_f16, half::f16, Round);

unary_inplace_thunk!(erf_inplace_f32, f32, Erf);
unary_inplace_thunk!(erf_inplace_f64, f64, Erf);
unary_inplace_thunk!(erf_inplace_bf16, half::bf16, Erf);
unary_inplace_thunk!(erf_inplace_f16, half::f16, Erf);

unary_inplace_thunk!(gelu_erf_inplace_f32, f32, GeluErf);
unary_inplace_thunk!(gelu_erf_inplace_f64, f64, GeluErf);
unary_inplace_thunk!(gelu_erf_inplace_bf16, half::bf16, GeluErf);
unary_inplace_thunk!(gelu_erf_inplace_f16, half::f16, GeluErf);

/// In-place affine: mutates `out[i] = mul * out[i] + add`. The
/// caller passes `out` as both the input + output through the
/// PipelinedExecutor's `WorkItemKind::InplaceKernel` arm. Mirrors
/// `affine_f32`'s formula; no separate input buffer. Used by the
/// `Op::Fused(INPLACE_AFFINE, _)` dispatch path.
pub fn affine_inplace_f32(out: &mut CpuStorageBytes, mul: f32, add: f32) -> Result<()> {
    let out_view: &mut [f32] = out.as_slice_mut()?;
    for slot in out_view.iter_mut() {
        *slot = mul * *slot + add;
    }
    Ok(())
}

/// In-place affine for f64 — same shape as [`affine_inplace_f32`].
pub fn affine_inplace_f64(out: &mut CpuStorageBytes, mul: f64, add: f64) -> Result<()> {
    let out_view: &mut [f64] = out.as_slice_mut()?;
    for slot in out_view.iter_mut() {
        *slot = mul * *slot + add;
    }
    Ok(())
}

/// In-place affine for bf16 — accumulator pivots through f32 (matches
/// the non-inplace `affine_*` family's precision convention).
pub fn affine_inplace_bf16(out: &mut CpuStorageBytes, mul: f64, add: f64) -> Result<()> {
    let out_view: &mut [half::bf16] = out.as_slice_mut()?;
    let mul_f32 = mul as f32;
    let add_f32 = add as f32;
    for slot in out_view.iter_mut() {
        *slot = half::bf16::from_f32(mul_f32 * slot.to_f32() + add_f32);
    }
    Ok(())
}

/// In-place affine for f16 — accumulator pivots through f32 (matches
/// the non-inplace `affine_*` family's precision convention).
pub fn affine_inplace_f16(out: &mut CpuStorageBytes, mul: f64, add: f64) -> Result<()> {
    let out_view: &mut [half::f16] = out.as_slice_mut()?;
    let mul_f32 = mul as f32;
    let add_f32 = add as f32;
    for slot in out_view.iter_mut() {
        *slot = half::f16::from_f32(mul_f32 * slot.to_f32() + add_f32);
    }
    Ok(())
}

// In-place clamp + powi — Phase 3e successor for the scalar-param
// in-place op family (2026-05-30). Same shape contract as
// `affine_inplace_*`: single buffer, scalar params come in via
// `OpParams::Clamp` / `OpParams::PowI`.

pub fn clamp_inplace_f32(out: &mut CpuStorageBytes, min: f32, max: f32) -> Result<()> {
    if min > max {
        return Err(Error::Msg(format!("clamp_inplace_f32: min ({min}) > max ({max})")).bt());
    }
    let view: &mut [f32] = out.as_slice_mut()?;
    for slot in view.iter_mut() {
        *slot = slot.clamp(min, max);
    }
    Ok(())
}

pub fn clamp_inplace_f64(out: &mut CpuStorageBytes, min: f64, max: f64) -> Result<()> {
    if min > max {
        return Err(Error::Msg(format!("clamp_inplace_f64: min ({min}) > max ({max})")).bt());
    }
    let view: &mut [f64] = out.as_slice_mut()?;
    for slot in view.iter_mut() {
        *slot = slot.clamp(min, max);
    }
    Ok(())
}

/// In-place clamp for bf16 — pivots through f32 like the non-inplace
/// half-precision affine kernels.
pub fn clamp_inplace_bf16(out: &mut CpuStorageBytes, min: f64, max: f64) -> Result<()> {
    if min > max {
        return Err(Error::Msg(format!("clamp_inplace_bf16: min ({min}) > max ({max})")).bt());
    }
    let view: &mut [half::bf16] = out.as_slice_mut()?;
    let (min_f32, max_f32) = (min as f32, max as f32);
    for slot in view.iter_mut() {
        *slot = half::bf16::from_f32(slot.to_f32().clamp(min_f32, max_f32));
    }
    Ok(())
}

pub fn clamp_inplace_f16(out: &mut CpuStorageBytes, min: f64, max: f64) -> Result<()> {
    if min > max {
        return Err(Error::Msg(format!("clamp_inplace_f16: min ({min}) > max ({max})")).bt());
    }
    let view: &mut [half::f16] = out.as_slice_mut()?;
    let (min_f32, max_f32) = (min as f32, max as f32);
    for slot in view.iter_mut() {
        *slot = half::f16::from_f32(slot.to_f32().clamp(min_f32, max_f32));
    }
    Ok(())
}

pub fn powi_inplace_f32(out: &mut CpuStorageBytes, exp: i32) -> Result<()> {
    let view: &mut [f32] = out.as_slice_mut()?;
    for slot in view.iter_mut() {
        *slot = slot.powi(exp);
    }
    Ok(())
}

pub fn powi_inplace_f64(out: &mut CpuStorageBytes, exp: i32) -> Result<()> {
    let view: &mut [f64] = out.as_slice_mut()?;
    for slot in view.iter_mut() {
        *slot = slot.powi(exp);
    }
    Ok(())
}

pub fn powi_inplace_bf16(out: &mut CpuStorageBytes, exp: i32) -> Result<()> {
    let view: &mut [half::bf16] = out.as_slice_mut()?;
    for slot in view.iter_mut() {
        *slot = half::bf16::from_f32(slot.to_f32().powi(exp));
    }
    Ok(())
}

pub fn powi_inplace_f16(out: &mut CpuStorageBytes, exp: i32) -> Result<()> {
    let view: &mut [half::f16] = out.as_slice_mut()?;
    for slot in view.iter_mut() {
        *slot = half::f16::from_f32(slot.to_f32().powi(exp));
    }
    Ok(())
}

/// Element-wise clamp: `out[i] = clamp(input[i], min, max)`.
pub fn clamp_f32(
    input: &CpuStorageBytes,
    out: &mut CpuStorageBytes,
    min: f32,
    max: f32,
) -> Result<()> {
    check_lens_2("clamp_f32", input.len_bytes(), out.len_bytes())?;
    if min > max {
        return Err(Error::Msg(format!("clamp_f32: min ({min}) > max ({max})")).bt());
    }
    let in_view: &[f32] = input.as_slice()?;
    let out_view: &mut [f32] = out.as_slice_mut()?;
    for (i, slot) in out_view.iter_mut().enumerate() {
        *slot = in_view[i].clamp(min, max);
    }
    Ok(())
}

/// Element-wise integer power: `out[i] = input[i].powi(exp)`.
pub fn powi_f32(input: &CpuStorageBytes, out: &mut CpuStorageBytes, exp: i32) -> Result<()> {
    check_lens_2("powi_f32", input.len_bytes(), out.len_bytes())?;
    let in_view: &[f32] = input.as_slice()?;
    let out_view: &mut [f32] = out.as_slice_mut()?;
    for (i, slot) in out_view.iter_mut().enumerate() {
        *slot = in_view[i].powi(exp);
    }
    Ok(())
}

binary_thunk!(maximum_f32, f32, Maximum);
binary_thunk!(minimum_f32, f32, Minimum);

// Native-arithmetic Affine / Clamp / PowI for f64.
pub fn affine_f64(
    input: &CpuStorageBytes,
    out: &mut CpuStorageBytes,
    mul: f64,
    add: f64,
) -> Result<()> {
    check_lens_2("affine_f64", input.len_bytes(), out.len_bytes())?;
    let in_view: &[f64] = input.as_slice()?;
    let out_view: &mut [f64] = out.as_slice_mut()?;
    for (i, slot) in out_view.iter_mut().enumerate() {
        *slot = mul * in_view[i] + add;
    }
    Ok(())
}

pub fn clamp_f64(
    input: &CpuStorageBytes,
    out: &mut CpuStorageBytes,
    min: f64,
    max: f64,
) -> Result<()> {
    check_lens_2("clamp_f64", input.len_bytes(), out.len_bytes())?;
    if min > max {
        return Err(Error::Msg(format!("clamp_f64: min ({min}) > max ({max})")).bt());
    }
    let in_view: &[f64] = input.as_slice()?;
    let out_view: &mut [f64] = out.as_slice_mut()?;
    for (i, slot) in out_view.iter_mut().enumerate() {
        *slot = in_view[i].clamp(min, max);
    }
    Ok(())
}

pub fn powi_f64(input: &CpuStorageBytes, out: &mut CpuStorageBytes, exp: i32) -> Result<()> {
    check_lens_2("powi_f64", input.len_bytes(), out.len_bytes())?;
    let in_view: &[f64] = input.as_slice()?;
    let out_view: &mut [f64] = out.as_slice_mut()?;
    for (i, slot) in out_view.iter_mut().enumerate() {
        *slot = in_view[i].powi(exp);
    }
    Ok(())
}

// Half-float Affine / Clamp / PowI via f32 round-trip.
macro_rules! affine_half_kernel {
    ($name:ident, $T:ty) => {
        pub fn $name(
            input: &CpuStorageBytes,
            out: &mut CpuStorageBytes,
            mul: f32,
            add: f32,
        ) -> Result<()> {
            check_lens_2(stringify!($name), input.len_bytes(), out.len_bytes())?;
            let in_view: &[$T] = input.as_slice()?;
            let out_view: &mut [$T] = out.as_slice_mut()?;
            for (i, slot) in out_view.iter_mut().enumerate() {
                *slot = <$T>::from_f32(mul * in_view[i].to_f32() + add);
            }
            Ok(())
        }
    };
}

affine_half_kernel!(affine_bf16, half::bf16);
affine_half_kernel!(affine_f16, half::f16);

macro_rules! clamp_half_kernel {
    ($name:ident, $T:ty) => {
        pub fn $name(
            input: &CpuStorageBytes,
            out: &mut CpuStorageBytes,
            min: f32,
            max: f32,
        ) -> Result<()> {
            check_lens_2(stringify!($name), input.len_bytes(), out.len_bytes())?;
            if min > max {
                return Err(Error::Msg(format!(
                    "{}: min ({min}) > max ({max})",
                    stringify!($name),
                ))
                .bt());
            }
            let in_view: &[$T] = input.as_slice()?;
            let out_view: &mut [$T] = out.as_slice_mut()?;
            for (i, slot) in out_view.iter_mut().enumerate() {
                *slot = <$T>::from_f32(in_view[i].to_f32().clamp(min, max));
            }
            Ok(())
        }
    };
}

clamp_half_kernel!(clamp_bf16, half::bf16);
clamp_half_kernel!(clamp_f16, half::f16);

macro_rules! powi_half_kernel {
    ($name:ident, $T:ty) => {
        pub fn $name(input: &CpuStorageBytes, out: &mut CpuStorageBytes, exp: i32) -> Result<()> {
            check_lens_2(stringify!($name), input.len_bytes(), out.len_bytes())?;
            let in_view: &[$T] = input.as_slice()?;
            let out_view: &mut [$T] = out.as_slice_mut()?;
            for (i, slot) in out_view.iter_mut().enumerate() {
                *slot = <$T>::from_f32(in_view[i].to_f32().powi(exp));
            }
            Ok(())
        }
    };
}

powi_half_kernel!(powi_bf16, half::bf16);
powi_half_kernel!(powi_f16, half::f16);

// PowI backward — `grad_x = exp * x^(exp-1) * upstream`. Two inputs
// `(x, upstream)`, single pass per element. Replaces the 3-node
// primitive decomposition (PowI(exp-1) → MulScalar → Mul) the
// pre-PowIBackward autograd emitted. Per-dtype.

/// `grad_x[i] = exp * x[i]^(exp-1) * upstream[i]` for f32.
pub fn powi_backward_f32(
    x: &CpuStorageBytes,
    upstream: &CpuStorageBytes,
    out: &mut CpuStorageBytes,
    exp: i32,
) -> Result<()> {
    check_lens_2("powi_backward_f32", x.len_bytes(), out.len_bytes())?;
    check_lens_2("powi_backward_f32", upstream.len_bytes(), out.len_bytes())?;
    let x_view: &[f32] = x.as_slice()?;
    let up_view: &[f32] = upstream.as_slice()?;
    let out_view: &mut [f32] = out.as_slice_mut()?;
    let coeff = exp as f32;
    let nm1 = exp - 1;
    for (i, slot) in out_view.iter_mut().enumerate() {
        *slot = coeff * x_view[i].powi(nm1) * up_view[i];
    }
    Ok(())
}

/// `grad_x[i] = exp * x[i]^(exp-1) * upstream[i]` for f64.
pub fn powi_backward_f64(
    x: &CpuStorageBytes,
    upstream: &CpuStorageBytes,
    out: &mut CpuStorageBytes,
    exp: i32,
) -> Result<()> {
    check_lens_2("powi_backward_f64", x.len_bytes(), out.len_bytes())?;
    check_lens_2("powi_backward_f64", upstream.len_bytes(), out.len_bytes())?;
    let x_view: &[f64] = x.as_slice()?;
    let up_view: &[f64] = upstream.as_slice()?;
    let out_view: &mut [f64] = out.as_slice_mut()?;
    let coeff = exp as f64;
    let nm1 = exp - 1;
    for (i, slot) in out_view.iter_mut().enumerate() {
        *slot = coeff * x_view[i].powi(nm1) * up_view[i];
    }
    Ok(())
}

macro_rules! powi_backward_half_kernel {
    ($name:ident, $T:ty) => {
        pub fn $name(
            x: &CpuStorageBytes,
            upstream: &CpuStorageBytes,
            out: &mut CpuStorageBytes,
            exp: i32,
        ) -> Result<()> {
            check_lens_2(stringify!($name), x.len_bytes(), out.len_bytes())?;
            check_lens_2(stringify!($name), upstream.len_bytes(), out.len_bytes())?;
            let x_view: &[$T] = x.as_slice()?;
            let up_view: &[$T] = upstream.as_slice()?;
            let out_view: &mut [$T] = out.as_slice_mut()?;
            let coeff = exp as f32;
            let nm1 = exp - 1;
            for (i, slot) in out_view.iter_mut().enumerate() {
                let v = coeff * x_view[i].to_f32().powi(nm1) * up_view[i].to_f32();
                *slot = <$T>::from_f32(v);
            }
            Ok(())
        }
    };
}

powi_backward_half_kernel!(powi_backward_bf16, half::bf16);
powi_backward_half_kernel!(powi_backward_f16, half::f16);

// ArgMax/ArgMin extensions to f64/bf16/f16. The existing
// `argextremum_dim_f32` operates on f32 directly; for other
// dtypes we widen to f32 for comparison (uniform NaN handling).
macro_rules! argextremum_dim_via_f32 {
    ($name:ident, $T:ty, $T_size:expr, $is_better:expr, $init:expr) => {
        pub fn $name(
            input: &CpuStorageBytes,
            output: &mut CpuStorageBytes,
            input_shape: &[usize],
            dim: usize,
        ) -> Result<()> {
            if dim >= input_shape.len() {
                return Err(Error::Msg(format!(
                    "{}: dim {dim} out of range for rank {}",
                    stringify!($name),
                    input_shape.len(),
                ))
                .bt());
            }
            let total_input: usize = input_shape.iter().product();
            if input.len_bytes() != total_input.saturating_mul($T_size) {
                return Err(Error::Msg(format!(
                    "{}: input bytes={} doesn't match shape {input_shape:?}",
                    stringify!($name),
                    input.len_bytes(),
                ))
                .bt());
            }
            let outer_count: usize = input_shape[..dim].iter().product();
            let dim_size = input_shape[dim];
            let inner_count: usize = input_shape[dim + 1..].iter().product();
            let output_count = outer_count * inner_count;
            if output.len_bytes() != output_count.saturating_mul(std::mem::size_of::<u32>()) {
                return Err(Error::Msg(format!(
                    "{}: output bytes={} doesn't match",
                    stringify!($name),
                    output.len_bytes(),
                ))
                .bt());
            }
            if dim_size == 0 {
                return Err(
                    Error::Msg(format!("{}: dim {dim} has size 0", stringify!($name),)).bt(),
                );
            }
            let in_view: &[$T] = input.as_slice()?;
            let out_view: &mut [u32] = output.as_slice_mut()?;
            let is_better: fn(f32, f32) -> bool = $is_better;
            for outer in 0..outer_count {
                for inner in 0..inner_count {
                    let mut best_val: f32 = $init;
                    let mut best_idx = 0u32;
                    for d in 0..dim_size {
                        let off = (outer * dim_size + d) * inner_count + inner;
                        let v = in_view[off].to_f32();
                        if d == 0 {
                            best_val = v;
                            best_idx = 0;
                        } else if is_better(v, best_val) {
                            best_val = v;
                            best_idx = d as u32;
                        }
                    }
                    out_view[outer * inner_count + inner] = best_idx;
                }
            }
            Ok(())
        }
    };
}

// f64 has a native to_f32 method via the `as` cast — wrap it.
trait ToF32Ext {
    fn to_f32(self) -> f32;
}
impl ToF32Ext for f64 {
    fn to_f32(self) -> f32 {
        self as f32
    }
}

argextremum_dim_via_f32!(
    argmax_dim_f64,
    f64,
    std::mem::size_of::<f64>(),
    |new: f32, best: f32| new > best,
    f32::NEG_INFINITY

);
argextremum_dim_via_f32!(
    argmin_dim_f64,
    f64,
    std::mem::size_of::<f64>(),
    |new: f32, best: f32| new < best,
    f32::INFINITY

);
argextremum_dim_via_f32!(
    argmax_dim_bf16,
    half::bf16,
    std::mem::size_of::<half::bf16>(),
    |new: f32, best: f32| new > best,
    f32::NEG_INFINITY

);
argextremum_dim_via_f32!(
    argmin_dim_bf16,
    half::bf16,
    std::mem::size_of::<half::bf16>(),
    |new: f32, best: f32| new < best,
    f32::INFINITY

);
argextremum_dim_via_f32!(
    argmax_dim_f16,
    half::f16,
    std::mem::size_of::<half::f16>(),
    |new: f32, best: f32| new > best,
    f32::NEG_INFINITY

);
argextremum_dim_via_f32!(
    argmin_dim_f16,
    half::f16,
    std::mem::size_of::<half::f16>(),
    |new: f32, best: f32| new < best,
    f32::INFINITY

);

// =============================================================================
// Dtype conversion (Cast)
// =============================================================================

/// Generate a typed dtype-conversion kernel of the form
/// `out[i] = convert(input[i])`. Validates that the input byte
/// length is a multiple of the source type's size and that the
/// output byte length matches `elem_count * size_of::<TOut>()`.
/// Output is pre-allocated by the caller.
macro_rules! cast_kernel {
    ($name:ident, $TIn:ty, $TOut:ty, $convert:expr, $doc:literal) => {
        #[doc = $doc]
        pub fn $name(input: &CpuStorageBytes, out: &mut CpuStorageBytes) -> Result<()> {
            let in_size = std::mem::size_of::<$TIn>();
            let out_size = std::mem::size_of::<$TOut>();
            if in_size == 0 || input.len_bytes() % in_size != 0 {
                return Err(Error::Msg(format!(
                    "{}: input bytes {} not a multiple of {} ({})",
                    stringify!($name),
                    input.len_bytes(),
                    in_size,
                    stringify!($TIn),
                ))
                .bt());
            }
            let elem_count = input.len_bytes() / in_size;
            let want_out = elem_count.saturating_mul(out_size);
            if out.len_bytes() != want_out {
                return Err(Error::Msg(format!(
                    "{}: output bytes {} doesn't match input elem count {} \
                     × {} ({}) = {}",
                    stringify!($name),
                    out.len_bytes(),
                    elem_count,
                    out_size,
                    stringify!($TOut),
                    want_out,
                ))
                .bt());
            }
            let in_view: &[$TIn] = input.as_slice()?;
            let out_view: &mut [$TOut] = out.as_slice_mut()?;
            let convert: fn($TIn) -> $TOut = $convert;
            for (i, slot) in out_view.iter_mut().enumerate() {
                *slot = convert(in_view[i]);
            }
            Ok(())
        }
    };
}

cast_kernel!(
    cast_f32_to_f64,
    f32,
    f64,
    |x: f32| x as f64,
    "Convert `f32` → `f64`. Lossless widening."

);
cast_kernel!(
    cast_f64_to_f32,
    f64,
    f32,
    |x: f64| x as f32,
    "Convert `f64` → `f32`. Lossy narrowing per IEEE-754 rounding."

);
cast_kernel!(
    cast_f32_to_bf16,
    f32,
    half::bf16,
    half::bf16::from_f32,
    "Convert `f32` → `bf16`. Lossy narrowing — keeps the f32 exponent and the top mantissa bits."
);
cast_kernel!(
    cast_bf16_to_f32,
    half::bf16,
    f32,
    |x: half::bf16| x.to_f32(),
    "Convert `bf16` → `f32`. Lossless widening (bf16 is a strict subset of f32)."

);
cast_kernel!(
    cast_f32_to_f16,
    f32,
    half::f16,
    half::f16::from_f32,
    "Convert `f32` → `f16`. Lossy narrowing — clips to f16 range with NaN/inf preserved."
);
cast_kernel!(
    cast_f16_to_f32,
    half::f16,
    f32,
    |x: half::f16| x.to_f32(),
    "Convert `f16` → `f32`. Lossless widening within f16's representable range."

);

// F8E4M3 ↔ {f32, f16, bf16} — mirrors baracuda alpha.29's CastSubBytePlan
// surface (see fuel-cuda-backend/src/baracuda/cast.rs). `float8::F8E4M3`
// doesn't impl `bytemuck::Pod`, so we can't use the generic
// `cast_kernel!` macro (which requires `Pod` on both src + dst element
// types). Instead, we handle F8E4M3 as raw `u8` bytes via `from_bits` /
// `to_bits` round-trips. F16↔F8E4M3 and BF16↔F8E4M3 pivot through f32
// because `float8::F8E4M3` only exposes `from_f32` / `to_f32`; the f32
// pivot leg is lossless for both F16 and BF16 (each is a strict subset
// of f32's representable values).

macro_rules! cast_kernel_to_fp8 {
    ($name:ident, $TIn:ty, $convert:expr, $doc:literal) => {
        #[doc = $doc]
        pub fn $name(input: &CpuStorageBytes, out: &mut CpuStorageBytes) -> Result<()> {
            let in_size = std::mem::size_of::<$TIn>();
            if input.len_bytes() % in_size != 0 {
                return Err(Error::Msg(format!(
                    "{}: input bytes {} not a multiple of {} size {}",
                    stringify!($name),
                    input.len_bytes(),
                    stringify!($TIn),
                    in_size,
                ))
                .bt());
            }
            let elem_count = input.len_bytes() / in_size;
            if out.len_bytes() != elem_count {
                return Err(Error::Msg(format!(
                    "{}: output bytes {} doesn't match input elem count {} (F8E4M3 is 1 byte/elem)",
                    stringify!($name),
                    out.len_bytes(),
                    elem_count,
                ))
                .bt());
            }
            let in_view: &[$TIn] = input.as_slice()?;
            let out_bytes: &mut [u8] = out.bytes_mut();
            let convert: fn($TIn) -> float8::F8E4M3 = $convert;
            for (i, slot) in out_bytes.iter_mut().enumerate() {
                *slot = convert(in_view[i]).to_bits();
            }
            Ok(())
        }
    };
}

macro_rules! cast_kernel_from_fp8 {
    ($name:ident, $TOut:ty, $convert:expr, $doc:literal) => {
        #[doc = $doc]
        pub fn $name(input: &CpuStorageBytes, out: &mut CpuStorageBytes) -> Result<()> {
            let out_size = std::mem::size_of::<$TOut>();
            let elem_count = input.len_bytes();
            let want_out = elem_count.saturating_mul(out_size);
            if out.len_bytes() != want_out {
                return Err(Error::Msg(format!(
                    "{}: output bytes {} doesn't match input elem count {} × {} ({}) = {}",
                    stringify!($name),
                    out.len_bytes(),
                    elem_count,
                    out_size,
                    stringify!($TOut),
                    want_out,
                ))
                .bt());
            }
            let in_bytes: &[u8] = input.bytes();
            let out_view: &mut [$TOut] = out.as_slice_mut()?;
            let convert: fn(float8::F8E4M3) -> $TOut = $convert;
            for (i, slot) in out_view.iter_mut().enumerate() {
                *slot = convert(float8::F8E4M3::from_bits(in_bytes[i]));
            }
            Ok(())
        }
    };
}

cast_kernel_to_fp8!(
    cast_f32_to_f8e4m3,
    f32,
    float8::F8E4M3::from_f32,
    "Convert `f32` → `F8E4M3`. Lossy narrowing per NV/OCP FP8 E4M3."
);
cast_kernel_from_fp8!(
    cast_f8e4m3_to_f32,
    f32,
    |x: float8::F8E4M3| x.to_f32(),
    "Convert `F8E4M3` → `f32`. Lossless widening (FP8 is a strict subset of f32)."

);
cast_kernel_to_fp8!(
    cast_bf16_to_f8e4m3,
    half::bf16,
    |x: half::bf16| float8::F8E4M3::from_f32(x.to_f32()),
    "Convert `bf16` → `F8E4M3` via f32. Lossy narrowing on the f8 leg only."

);
cast_kernel_from_fp8!(
    cast_f8e4m3_to_bf16,
    half::bf16,
    |x: float8::F8E4M3| half::bf16::from_f32(x.to_f32()),
    "Convert `F8E4M3` → `bf16` via f32. Lossless within F8E4M3's range."

);
cast_kernel_to_fp8!(
    cast_f16_to_f8e4m3,
    half::f16,
    |x: half::f16| float8::F8E4M3::from_f32(x.to_f32()),
    "Convert `f16` → `F8E4M3` via f32. Lossy narrowing on the f8 leg only."

);
cast_kernel_from_fp8!(
    cast_f8e4m3_to_f16,
    half::f16,
    |x: float8::F8E4M3| half::f16::from_f32(x.to_f32()),
    "Convert `F8E4M3` → `f16` via f32. Lossless within F8E4M3's range."

);

// Cast matrix completion: every ordered pair of the 11 real numeric
// dtypes {F32,F64,F16,BF16,F8E4M3,U8,I8,U32,I16,I32,I64} now has a
// registered CPU Cast kernel (identity pairs excluded — the optimizer
// elides them). Conversions follow the standard idioms: int/f64<->f32
// via `as`; ->half via `from_f32`; half/f8->int via `to_f32() as`;
// half<->half and anything touching F8E4M3 pivot through f32.
cast_kernel!(
    cast_f32_to_u8,
    f32,
    u8,
    |x: f32| x as u8,
    "Convert `f32` -> `u8`."

);
cast_kernel!(
    cast_f32_to_i8,
    f32,
    i8,
    |x: f32| x as i8,
    "Convert `f32` -> `i8`."

);
cast_kernel!(
    cast_f32_to_u32,
    f32,
    u32,
    |x: f32| x as u32,
    "Convert `f32` -> `u32`."

);
cast_kernel!(
    cast_f32_to_i16,
    f32,
    i16,
    |x: f32| x as i16,
    "Convert `f32` -> `i16`."

);
cast_kernel!(
    cast_f32_to_i32,
    f32,
    i32,
    |x: f32| x as i32,
    "Convert `f32` -> `i32`."

);
cast_kernel!(
    cast_f32_to_i64,
    f32,
    i64,
    |x: f32| x as i64,
    "Convert `f32` -> `i64`."

);
cast_kernel!(
    cast_f64_to_f16,
    f64,
    half::f16,
    |x: f64| half::f16::from_f32(x as f32),
    "Convert `f64` -> `f16`."

);
cast_kernel!(
    cast_f64_to_bf16,
    f64,
    half::bf16,
    |x: f64| half::bf16::from_f32(x as f32),
    "Convert `f64` -> `bf16`."

);
cast_kernel_to_fp8!(
    cast_f64_to_f8e4m3,
    f64,
    |x: f64| float8::F8E4M3::from_f32(x as f32),
    "Convert `f64` -> `f8e4m3`."

);
cast_kernel!(
    cast_f64_to_u8,
    f64,
    u8,
    |x: f64| x as u8,
    "Convert `f64` -> `u8`."

);
cast_kernel!(
    cast_f64_to_i8,
    f64,
    i8,
    |x: f64| x as i8,
    "Convert `f64` -> `i8`."

);
cast_kernel!(
    cast_f64_to_u32,
    f64,
    u32,
    |x: f64| x as u32,
    "Convert `f64` -> `u32`."

);
cast_kernel!(
    cast_f64_to_i16,
    f64,
    i16,
    |x: f64| x as i16,
    "Convert `f64` -> `i16`."

);
cast_kernel!(
    cast_f64_to_i32,
    f64,
    i32,
    |x: f64| x as i32,
    "Convert `f64` -> `i32`."

);
cast_kernel!(
    cast_f64_to_i64,
    f64,
    i64,
    |x: f64| x as i64,
    "Convert `f64` -> `i64`."

);
cast_kernel!(
    cast_f16_to_f64,
    half::f16,
    f64,
    |x: half::f16| x.to_f32() as f64,
    "Convert `f16` -> `f64`."

);
cast_kernel!(
    cast_f16_to_bf16,
    half::f16,
    half::bf16,
    |x: half::f16| half::bf16::from_f32(x.to_f32()),
    "Convert `f16` -> `bf16`."

);
cast_kernel!(
    cast_f16_to_u8,
    half::f16,
    u8,
    |x: half::f16| x.to_f32() as u8,
    "Convert `f16` -> `u8`."

);
cast_kernel!(
    cast_f16_to_i8,
    half::f16,
    i8,
    |x: half::f16| x.to_f32() as i8,
    "Convert `f16` -> `i8`."

);
cast_kernel!(
    cast_f16_to_u32,
    half::f16,
    u32,
    |x: half::f16| x.to_f32() as u32,
    "Convert `f16` -> `u32`."

);
cast_kernel!(
    cast_f16_to_i16,
    half::f16,
    i16,
    |x: half::f16| x.to_f32() as i16,
    "Convert `f16` -> `i16`."

);
cast_kernel!(
    cast_f16_to_i32,
    half::f16,
    i32,
    |x: half::f16| x.to_f32() as i32,
    "Convert `f16` -> `i32`."

);
cast_kernel!(
    cast_f16_to_i64,
    half::f16,
    i64,
    |x: half::f16| x.to_f32() as i64,
    "Convert `f16` -> `i64`."

);
cast_kernel!(
    cast_bf16_to_f64,
    half::bf16,
    f64,
    |x: half::bf16| x.to_f32() as f64,
    "Convert `bf16` -> `f64`."

);
cast_kernel!(
    cast_bf16_to_f16,
    half::bf16,
    half::f16,
    |x: half::bf16| half::f16::from_f32(x.to_f32()),
    "Convert `bf16` -> `f16`."

);
cast_kernel!(
    cast_bf16_to_u8,
    half::bf16,
    u8,
    |x: half::bf16| x.to_f32() as u8,
    "Convert `bf16` -> `u8`."

);
cast_kernel!(
    cast_bf16_to_i8,
    half::bf16,
    i8,
    |x: half::bf16| x.to_f32() as i8,
    "Convert `bf16` -> `i8`."

);
cast_kernel!(
    cast_bf16_to_u32,
    half::bf16,
    u32,
    |x: half::bf16| x.to_f32() as u32,
    "Convert `bf16` -> `u32`."

);
cast_kernel!(
    cast_bf16_to_i16,
    half::bf16,
    i16,
    |x: half::bf16| x.to_f32() as i16,
    "Convert `bf16` -> `i16`."

);
cast_kernel!(
    cast_bf16_to_i32,
    half::bf16,
    i32,
    |x: half::bf16| x.to_f32() as i32,
    "Convert `bf16` -> `i32`."

);
cast_kernel!(
    cast_bf16_to_i64,
    half::bf16,
    i64,
    |x: half::bf16| x.to_f32() as i64,
    "Convert `bf16` -> `i64`."

);
cast_kernel_from_fp8!(
    cast_f8e4m3_to_f64,
    f64,
    |x: float8::F8E4M3| x.to_f32() as f64,
    "Convert `f8e4m3` -> `f64`."

);
cast_kernel_from_fp8!(
    cast_f8e4m3_to_u8,
    u8,
    |x: float8::F8E4M3| x.to_f32() as u8,
    "Convert `f8e4m3` -> `u8`."

);
cast_kernel_from_fp8!(
    cast_f8e4m3_to_i8,
    i8,
    |x: float8::F8E4M3| x.to_f32() as i8,
    "Convert `f8e4m3` -> `i8`."

);
cast_kernel_from_fp8!(
    cast_f8e4m3_to_u32,
    u32,
    |x: float8::F8E4M3| x.to_f32() as u32,
    "Convert `f8e4m3` -> `u32`."

);
cast_kernel_from_fp8!(
    cast_f8e4m3_to_i16,
    i16,
    |x: float8::F8E4M3| x.to_f32() as i16,
    "Convert `f8e4m3` -> `i16`."

);
cast_kernel_from_fp8!(
    cast_f8e4m3_to_i32,
    i32,
    |x: float8::F8E4M3| x.to_f32() as i32,
    "Convert `f8e4m3` -> `i32`."

);
cast_kernel_from_fp8!(
    cast_f8e4m3_to_i64,
    i64,
    |x: float8::F8E4M3| x.to_f32() as i64,
    "Convert `f8e4m3` -> `i64`."

);
cast_kernel!(
    cast_u8_to_f32,
    u8,
    f32,
    |x: u8| x as f32,
    "Convert `u8` -> `f32`."

);
cast_kernel!(
    cast_u8_to_f64,
    u8,
    f64,
    |x: u8| x as f64,
    "Convert `u8` -> `f64`."

);
cast_kernel!(
    cast_u8_to_f16,
    u8,
    half::f16,
    |x: u8| half::f16::from_f32(x as f32),
    "Convert `u8` -> `f16`."

);
cast_kernel!(
    cast_u8_to_bf16,
    u8,
    half::bf16,
    |x: u8| half::bf16::from_f32(x as f32),
    "Convert `u8` -> `bf16`."

);
cast_kernel_to_fp8!(
    cast_u8_to_f8e4m3,
    u8,
    |x: u8| float8::F8E4M3::from_f32(x as f32),
    "Convert `u8` -> `f8e4m3`."

);
cast_kernel!(
    cast_u8_to_i8,
    u8,
    i8,
    |x: u8| x as i8,
    "Convert `u8` -> `i8`."

);
cast_kernel!(
    cast_u8_to_u32,
    u8,
    u32,
    |x: u8| x as u32,
    "Convert `u8` -> `u32`."

);
cast_kernel!(
    cast_u8_to_i16,
    u8,
    i16,
    |x: u8| x as i16,
    "Convert `u8` -> `i16`."

);
cast_kernel!(
    cast_u8_to_i32,
    u8,
    i32,
    |x: u8| x as i32,
    "Convert `u8` -> `i32`."

);
cast_kernel!(
    cast_u8_to_i64,
    u8,
    i64,
    |x: u8| x as i64,
    "Convert `u8` -> `i64`."

);
cast_kernel!(
    cast_i8_to_f32,
    i8,
    f32,
    |x: i8| x as f32,
    "Convert `i8` -> `f32`."

);
cast_kernel!(
    cast_i8_to_f64,
    i8,
    f64,
    |x: i8| x as f64,
    "Convert `i8` -> `f64`."

);
cast_kernel!(
    cast_i8_to_f16,
    i8,
    half::f16,
    |x: i8| half::f16::from_f32(x as f32),
    "Convert `i8` -> `f16`."

);
cast_kernel!(
    cast_i8_to_bf16,
    i8,
    half::bf16,
    |x: i8| half::bf16::from_f32(x as f32),
    "Convert `i8` -> `bf16`."

);
cast_kernel_to_fp8!(
    cast_i8_to_f8e4m3,
    i8,
    |x: i8| float8::F8E4M3::from_f32(x as f32),
    "Convert `i8` -> `f8e4m3`."

);
cast_kernel!(
    cast_i8_to_u8,
    i8,
    u8,
    |x: i8| x as u8,
    "Convert `i8` -> `u8`."

);
cast_kernel!(
    cast_i8_to_u32,
    i8,
    u32,
    |x: i8| x as u32,
    "Convert `i8` -> `u32`."

);
cast_kernel!(
    cast_i8_to_i16,
    i8,
    i16,
    |x: i8| x as i16,
    "Convert `i8` -> `i16`."

);
cast_kernel!(
    cast_i8_to_i32,
    i8,
    i32,
    |x: i8| x as i32,
    "Convert `i8` -> `i32`."

);
cast_kernel!(
    cast_i8_to_i64,
    i8,
    i64,
    |x: i8| x as i64,
    "Convert `i8` -> `i64`."

);
cast_kernel!(
    cast_u32_to_f32,
    u32,
    f32,
    |x: u32| x as f32,
    "Convert `u32` -> `f32`."

);
cast_kernel!(
    cast_u32_to_f64,
    u32,
    f64,
    |x: u32| x as f64,
    "Convert `u32` -> `f64`."

);
cast_kernel!(
    cast_u32_to_f16,
    u32,
    half::f16,
    |x: u32| half::f16::from_f32(x as f32),
    "Convert `u32` -> `f16`."

);
cast_kernel!(
    cast_u32_to_bf16,
    u32,
    half::bf16,
    |x: u32| half::bf16::from_f32(x as f32),
    "Convert `u32` -> `bf16`."

);
cast_kernel_to_fp8!(
    cast_u32_to_f8e4m3,
    u32,
    |x: u32| float8::F8E4M3::from_f32(x as f32),
    "Convert `u32` -> `f8e4m3`."

);
cast_kernel!(
    cast_u32_to_u8,
    u32,
    u8,
    |x: u32| x as u8,
    "Convert `u32` -> `u8`."

);
cast_kernel!(
    cast_u32_to_i8,
    u32,
    i8,
    |x: u32| x as i8,
    "Convert `u32` -> `i8`."

);
cast_kernel!(
    cast_u32_to_i16,
    u32,
    i16,
    |x: u32| x as i16,
    "Convert `u32` -> `i16`."

);
cast_kernel!(
    cast_u32_to_i32,
    u32,
    i32,
    |x: u32| x as i32,
    "Convert `u32` -> `i32`."

);
cast_kernel!(
    cast_u32_to_i64,
    u32,
    i64,
    |x: u32| x as i64,
    "Convert `u32` -> `i64`."

);
cast_kernel!(
    cast_i16_to_f32,
    i16,
    f32,
    |x: i16| x as f32,
    "Convert `i16` -> `f32`."

);
cast_kernel!(
    cast_i16_to_f64,
    i16,
    f64,
    |x: i16| x as f64,
    "Convert `i16` -> `f64`."

);
cast_kernel!(
    cast_i16_to_f16,
    i16,
    half::f16,
    |x: i16| half::f16::from_f32(x as f32),
    "Convert `i16` -> `f16`."

);
cast_kernel!(
    cast_i16_to_bf16,
    i16,
    half::bf16,
    |x: i16| half::bf16::from_f32(x as f32),
    "Convert `i16` -> `bf16`."

);
cast_kernel_to_fp8!(
    cast_i16_to_f8e4m3,
    i16,
    |x: i16| float8::F8E4M3::from_f32(x as f32),
    "Convert `i16` -> `f8e4m3`."

);
cast_kernel!(
    cast_i16_to_u8,
    i16,
    u8,
    |x: i16| x as u8,
    "Convert `i16` -> `u8`."

);
cast_kernel!(
    cast_i16_to_i8,
    i16,
    i8,
    |x: i16| x as i8,
    "Convert `i16` -> `i8`."

);
cast_kernel!(
    cast_i16_to_u32,
    i16,
    u32,
    |x: i16| x as u32,
    "Convert `i16` -> `u32`."

);
cast_kernel!(
    cast_i16_to_i32,
    i16,
    i32,
    |x: i16| x as i32,
    "Convert `i16` -> `i32`."

);
cast_kernel!(
    cast_i16_to_i64,
    i16,
    i64,
    |x: i16| x as i64,
    "Convert `i16` -> `i64`."

);
cast_kernel!(
    cast_i32_to_f32,
    i32,
    f32,
    |x: i32| x as f32,
    "Convert `i32` -> `f32`."

);
cast_kernel!(
    cast_i32_to_f64,
    i32,
    f64,
    |x: i32| x as f64,
    "Convert `i32` -> `f64`."

);
cast_kernel!(
    cast_i32_to_f16,
    i32,
    half::f16,
    |x: i32| half::f16::from_f32(x as f32),
    "Convert `i32` -> `f16`."

);
cast_kernel!(
    cast_i32_to_bf16,
    i32,
    half::bf16,
    |x: i32| half::bf16::from_f32(x as f32),
    "Convert `i32` -> `bf16`."

);
cast_kernel_to_fp8!(
    cast_i32_to_f8e4m3,
    i32,
    |x: i32| float8::F8E4M3::from_f32(x as f32),
    "Convert `i32` -> `f8e4m3`."

);
cast_kernel!(
    cast_i32_to_u8,
    i32,
    u8,
    |x: i32| x as u8,
    "Convert `i32` -> `u8`."

);
cast_kernel!(
    cast_i32_to_i8,
    i32,
    i8,
    |x: i32| x as i8,
    "Convert `i32` -> `i8`."

);
cast_kernel!(
    cast_i32_to_u32,
    i32,
    u32,
    |x: i32| x as u32,
    "Convert `i32` -> `u32`."

);
cast_kernel!(
    cast_i32_to_i16,
    i32,
    i16,
    |x: i32| x as i16,
    "Convert `i32` -> `i16`."

);
cast_kernel!(
    cast_i32_to_i64,
    i32,
    i64,
    |x: i32| x as i64,
    "Convert `i32` -> `i64`."

);
cast_kernel!(
    cast_i64_to_f32,
    i64,
    f32,
    |x: i64| x as f32,
    "Convert `i64` -> `f32`."

);
cast_kernel!(
    cast_i64_to_f64,
    i64,
    f64,
    |x: i64| x as f64,
    "Convert `i64` -> `f64`."

);
cast_kernel!(
    cast_i64_to_f16,
    i64,
    half::f16,
    |x: i64| half::f16::from_f32(x as f32),
    "Convert `i64` -> `f16`."

);
cast_kernel!(
    cast_i64_to_bf16,
    i64,
    half::bf16,
    |x: i64| half::bf16::from_f32(x as f32),
    "Convert `i64` -> `bf16`."

);
cast_kernel_to_fp8!(
    cast_i64_to_f8e4m3,
    i64,
    |x: i64| float8::F8E4M3::from_f32(x as f32),
    "Convert `i64` -> `f8e4m3`."

);
cast_kernel!(
    cast_i64_to_u8,
    i64,
    u8,
    |x: i64| x as u8,
    "Convert `i64` -> `u8`."

);
cast_kernel!(
    cast_i64_to_i8,
    i64,
    i8,
    |x: i64| x as i8,
    "Convert `i64` -> `i8`."

);
cast_kernel!(
    cast_i64_to_u32,
    i64,
    u32,
    |x: i64| x as u32,
    "Convert `i64` -> `u32`."

);
cast_kernel!(
    cast_i64_to_i16,
    i64,
    i16,
    |x: i64| x as i16,
    "Convert `i64` -> `i16`."

);
cast_kernel!(
    cast_i64_to_i32,
    i64,
    i32,
    |x: i64| x as i32,
    "Convert `i64` -> `i32`."

);

// ---------------------------------------------------------------------------
// Bool cast pairs (GAP-168).  `DType::Bool` is one byte per element carrying
// 0 = false / 1 = true, so every kernel below uses `u8` as the CARRIER type
// (`bool` is not `bytemuck::Pod` — it has invalid bit patterns — while the
// storage byte is unconstrained until we normalize it).
//
// Semantics match PyTorch's `.bool()` / `.to(dtype)`:
//   X -> Bool : `x != 0`, so `-0.0` is FALSE and NaN is TRUE.
//   Bool -> X : `false` -> 0, `true` -> 1.
// The `!= 0` on the Bool-source side is a DEFENSIVE NORMALIZATION: Fuel only
// ever writes 0/1 into Bool storage, but a non-canonical byte reaching here
// resolves to `true` rather than propagating a raw integer into a float.
// This is what makes Bool DISTINGUISHABLE from U8 rather than a reinterpret:
// `U8(5) -> Bool` is 1, not 5.

cast_kernel!(
    cast_f32_to_bool,
    f32,
    u8,
    |x: f32| (x != 0.0) as u8,
    "Convert `f32` -> `bool`. Nonzero (incl. NaN) is true."

);
cast_kernel!(
    cast_f64_to_bool,
    f64,
    u8,
    |x: f64| (x != 0.0) as u8,
    "Convert `f64` -> `bool`. Nonzero (incl. NaN) is true."

);
cast_kernel!(
    cast_f16_to_bool,
    half::f16,
    u8,
    |x: half::f16| (x.to_f32() != 0.0) as u8,
    "Convert `f16` -> `bool`. Nonzero (incl. NaN) is true."

);
cast_kernel!(
    cast_bf16_to_bool,
    half::bf16,
    u8,
    |x: half::bf16| (x.to_f32() != 0.0) as u8,
    "Convert `bf16` -> `bool`. Nonzero (incl. NaN) is true."

);
cast_kernel_from_fp8!(
    cast_f8e4m3_to_bool,
    u8,
    |x: float8::F8E4M3| (x.to_f32() != 0.0) as u8,
    "Convert `F8E4M3` -> `bool` via f32. Nonzero (incl. NaN) is true."

);
cast_kernel!(
    cast_u8_to_bool,
    u8,
    u8,
    |x: u8| (x != 0) as u8,
    "Convert `u8` -> `bool`. Nonzero is true."

);
cast_kernel!(
    cast_i8_to_bool,
    i8,
    u8,
    |x: i8| (x != 0) as u8,
    "Convert `i8` -> `bool`. Nonzero is true."

);
cast_kernel!(
    cast_u32_to_bool,
    u32,
    u8,
    |x: u32| (x != 0) as u8,
    "Convert `u32` -> `bool`. Nonzero is true."

);
cast_kernel!(
    cast_i16_to_bool,
    i16,
    u8,
    |x: i16| (x != 0) as u8,
    "Convert `i16` -> `bool`. Nonzero is true."

);
cast_kernel!(
    cast_i32_to_bool,
    i32,
    u8,
    |x: i32| (x != 0) as u8,
    "Convert `i32` -> `bool`. Nonzero is true."

);
cast_kernel!(
    cast_i64_to_bool,
    i64,
    u8,
    |x: i64| (x != 0) as u8,
    "Convert `i64` -> `bool`. Nonzero is true."

);
cast_kernel!(
    cast_bool_to_f32,
    u8,
    f32,
    |b: u8| (b != 0) as u8 as f32,
    "Convert `bool` -> `f32`. false -> 0, true -> 1."

);
cast_kernel!(
    cast_bool_to_f64,
    u8,
    f64,
    |b: u8| (b != 0) as u8 as f64,
    "Convert `bool` -> `f64`. false -> 0, true -> 1."

);
cast_kernel!(
    cast_bool_to_f16,
    u8,
    half::f16,
    |b: u8| half::f16::from_f32((b != 0) as u8 as f32),
    "Convert `bool` -> `f16`. false -> 0, true -> 1."

);
cast_kernel!(
    cast_bool_to_bf16,
    u8,
    half::bf16,
    |b: u8| half::bf16::from_f32((b != 0) as u8 as f32),
    "Convert `bool` -> `bf16`. false -> 0, true -> 1."

);
cast_kernel_to_fp8!(
    cast_bool_to_f8e4m3,
    u8,
    |b: u8| float8::F8E4M3::from_f32((b != 0) as u8 as f32),
    "Convert `bool` -> `F8E4M3` via f32. false -> 0.0, true -> 1.0."

);
cast_kernel!(
    cast_bool_to_u8,
    u8,
    u8,
    |b: u8| (b != 0) as u8,
    "Convert `bool` -> `u8`. false -> 0, true -> 1."

);
cast_kernel!(
    cast_bool_to_i8,
    u8,
    i8,
    |b: u8| (b != 0) as u8 as i8,
    "Convert `bool` -> `i8`. false -> 0, true -> 1."

);
cast_kernel!(
    cast_bool_to_u32,
    u8,
    u32,
    |b: u8| (b != 0) as u8 as u32,
    "Convert `bool` -> `u32`. false -> 0, true -> 1."

);
cast_kernel!(
    cast_bool_to_i16,
    u8,
    i16,
    |b: u8| (b != 0) as u8 as i16,
    "Convert `bool` -> `i16`. false -> 0, true -> 1."

);
cast_kernel!(
    cast_bool_to_i32,
    u8,
    i32,
    |b: u8| (b != 0) as u8 as i32,
    "Convert `bool` -> `i32`. false -> 0, true -> 1."

);
cast_kernel!(
    cast_bool_to_i64,
    u8,
    i64,
    |b: u8| (b != 0) as u8 as i64,
    "Convert `bool` -> `i64`. false -> 0, true -> 1."

);

// =============================================================================
// Matrix multiplication (f32)
// =============================================================================

/// Batched row-major `f32` matrix multiply with optional GQA-style
/// batch broadcasting. For each output batch index, the kernel
/// runs the textbook (i, k, j) triple loop on the corresponding
/// `[m, k] @ [k, n]` slice.
///
/// Per-axis the batch dims either match (`lhs_batch_dims[i] ==
/// rhs_batch_dims[i]`) or follow GQA-style divisibility
/// (`lhs_dim > rhs_dim && lhs_dim % rhs_dim == 0`). The kernel
/// maps each lhs batch slot to the corresponding rhs slot via
/// `rhs_axis_idx = lhs_axis_idx / n_rep_axis`. Equal-batch and
/// rank-2 cases fall out as `n_rep == 1` everywhere.
///
/// Inputs are assumed contiguous in row-major order (the pipelined
/// executor's auto-Contiguize pass guarantees this). Output is
/// pre-allocated and overwritten by this kernel.
///
/// This is correctness-first; vendor BLAS backends will eclipse
/// it on performance once they're wired into the unified path.
pub fn matmul_f32(
    lhs: &CpuStorageBytes,
    rhs: &CpuStorageBytes,
    out: &mut CpuStorageBytes,
    lhs_batch_dims: &[usize],
    rhs_batch_dims: &[usize],
    m: usize,
    n: usize,
    k: usize,
) -> Result<()> {
    // Static M: compute every row (`m_compute == m_capacity == m`).
    matmul_f32_capacity(lhs, rhs, out, lhs_batch_dims, rhs_batch_dims, m, m, n, k)
}

/// Data-determined-M matmul (the dropless-MoE / capacity-buffer atom). The
/// LHS and OUTPUT are **capacity buffers** with `m_capacity` rows each; only
/// the first `m_compute` rows (`m_compute <= m_capacity`) are computed, the
/// rest of the output left zero. This is what lets a sparse-MoE expert FFN
/// process exactly its `count_e` routed tokens (a runtime `Extent` resolved
/// at execute) inside a statically-`m_capacity`-sized buffer — the compute
/// saving that makes sparse dispatch worth it. Row `i`'s LHS/OUTPUT offsets
/// use `m_capacity` for striding (buffer layout), so the untouched tail rows
/// keep their (zeroed) capacity slots. `matmul_f32` is the `m_compute ==
/// m_capacity` special case.
pub fn matmul_f32_capacity(
    lhs: &CpuStorageBytes,
    rhs: &CpuStorageBytes,
    out: &mut CpuStorageBytes,
    lhs_batch_dims: &[usize],
    rhs_batch_dims: &[usize],
    m_capacity: usize,
    m_compute: usize,
    n: usize,
    k: usize,
) -> Result<()> {
    if m_compute > m_capacity {
        return Err(Error::Msg(format!(
            "matmul_f32_capacity: m_compute ({m_compute}) must not exceed \
             m_capacity ({m_capacity})",
        ))
        .bt());
    }
    let m = m_capacity;
    if lhs_batch_dims.len() != rhs_batch_dims.len() {
        return Err(Error::Msg(format!(
            "matmul_f32: batch ranks must match (lhs={}, rhs={}); fuel-graph's \
             auto-broadcast equalizes them at graph construction time",
            lhs_batch_dims.len(),
            rhs_batch_dims.len(),
        ))
        .bt());
    }
    let batch_rank = lhs_batch_dims.len();
    // Per-axis n_rep validation: matched (n_rep=1) or GQA (lhs is a
    // multiple of rhs).
    let mut n_rep: Vec<usize> = Vec::with_capacity(batch_rank);
    for i in 0..batch_rank {
        let la = lhs_batch_dims[i];
        let ra = rhs_batch_dims[i];
        if la == ra {
            n_rep.push(1);
        } else if ra > 0 && la > ra && la.is_multiple_of(ra) {
            n_rep.push(la / ra);
        } else {
            return Err(Error::Msg(format!(
                "matmul_f32: batch dim {i} disallowed combination (lhs={la}, rhs={ra}); \
                 must be equal or GQA-divisible (lhs > rhs && lhs % rhs == 0)",
            ))
            .bt());
        }
    }
    let elem = std::mem::size_of::<f32>();
    let lhs_per_batch = m.saturating_mul(k);
    let rhs_per_batch = k.saturating_mul(n);
    let out_per_batch = m.saturating_mul(n);
    let lhs_batch_count: usize = lhs_batch_dims.iter().product::<usize>().max(1);
    let rhs_batch_count: usize = rhs_batch_dims.iter().product::<usize>().max(1);
    let need_lhs = lhs_batch_count
        .saturating_mul(lhs_per_batch)
        .saturating_mul(elem);
    let need_rhs = rhs_batch_count
        .saturating_mul(rhs_per_batch)
        .saturating_mul(elem);
    let need_out = lhs_batch_count
        .saturating_mul(out_per_batch)
        .saturating_mul(elem);
    if lhs.len_bytes() != need_lhs {
        return Err(Error::Msg(format!(
            "matmul_f32: lhs bytes={} doesn't match shape {:?} + [{m}, {k}] (f32)",
            lhs.len_bytes(),
            lhs_batch_dims,
        ))
        .bt());
    }
    if rhs.len_bytes() != need_rhs {
        return Err(Error::Msg(format!(
            "matmul_f32: rhs bytes={} doesn't match shape {:?} + [{k}, {n}] (f32)",
            rhs.len_bytes(),
            rhs_batch_dims,
        ))
        .bt());
    }
    if out.len_bytes() != need_out {
        return Err(Error::Msg(format!(
            "matmul_f32: out bytes={} doesn't match shape {:?} + [{m}, {n}] (f32)",
            out.len_bytes(),
            lhs_batch_dims,
        ))
        .bt());
    }
    let lhs_view: &[f32] = lhs.as_slice()?;
    let rhs_view: &[f32] = rhs.as_slice()?;
    let out_view: &mut [f32] = out.as_slice_mut()?;
    for slot in out_view.iter_mut() {
        *slot = 0.0;
    }
    // Decode lhs's flat batch index into a multi-index, map per-axis
    // to rhs's multi-index, encode rhs's flat batch index. Reuses
    // two scratch buffers across iterations.
    let mut lhs_multi = vec![0usize; batch_rank];
    let mut rhs_multi = vec![0usize; batch_rank];
    for b in 0..lhs_batch_count {
        // Decode b into lhs_multi (row-major over lhs_batch_dims).
        let mut rem = b;
        for d in (0..batch_rank).rev() {
            let s = lhs_batch_dims[d];
            lhs_multi[d] = rem % s;
            rem /= s;
        }
        // Per-axis GQA mapping.
        for d in 0..batch_rank {
            rhs_multi[d] = lhs_multi[d] / n_rep[d];
        }
        // Encode rhs's flat batch index.
        let mut rhs_b = 0usize;
        for d in 0..batch_rank {
            rhs_b = rhs_b * rhs_batch_dims[d] + rhs_multi[d];
        }
        let lhs_off = b * lhs_per_batch;
        let rhs_off = rhs_b * rhs_per_batch;
        let out_off = b * out_per_batch;
        // Compute only the first `m_compute` rows; the tail
        // (`m_compute..m_capacity`) stays zeroed — the capacity slots for
        // the rows this pass's runtime extent did not fill.
        for i in 0..m_compute {
            for kk in 0..k {
                let a = lhs_view[lhs_off + i * k + kk];
                let rhs_row_off = rhs_off + kk * n;
                let out_row_off = out_off + i * n;
                for j in 0..n {
                    out_view[out_row_off + j] += a * rhs_view[rhs_row_off + j];
                }
            }
        }
    }
    Ok(())
}

/// Batched row-major matmul for half-float types (bf16/f16) with
/// f32 accumulation. The kernel widens each input element to f32,
/// accumulates the inner-product sum in f32, then narrows the
/// final output back to the half-float type. Matches the standard
/// "f16 ops with f32 accumulator" pattern that cuBLAS / TPU /
/// cuDNN use to keep matmul numerically stable on half floats.
macro_rules! matmul_half_kernel {
    ($name:ident, $T:ty, $type_name:literal) => {
        pub fn $name(
            lhs: &CpuStorageBytes,
            rhs: &CpuStorageBytes,
            out: &mut CpuStorageBytes,
            lhs_batch_dims: &[usize],
            rhs_batch_dims: &[usize],
            m: usize,
            n: usize,
            k: usize,
        ) -> Result<()> {
            if lhs_batch_dims.len() != rhs_batch_dims.len() {
                return Err(Error::Msg(format!(
                    "{}: batch ranks must match (lhs={}, rhs={})",
                    $type_name,
                    lhs_batch_dims.len(),
                    rhs_batch_dims.len(),
                ))
                .bt());
            }
            let batch_rank = lhs_batch_dims.len();
            let mut n_rep: Vec<usize> = Vec::with_capacity(batch_rank);
            for i in 0..batch_rank {
                let la = lhs_batch_dims[i];
                let ra = rhs_batch_dims[i];
                if la == ra {
                    n_rep.push(1);
                } else if ra > 0 && la > ra && la % ra == 0 {
                    n_rep.push(la / ra);
                } else {
                    return Err(Error::Msg(format!(
                        "{}: batch dim {i} disallowed combination (lhs={la}, rhs={ra})",
                        $type_name,
                    ))
                    .bt());
                }
            }
            let elem = std::mem::size_of::<$T>();
            let lhs_per_batch = m.saturating_mul(k);
            let rhs_per_batch = k.saturating_mul(n);
            let out_per_batch = m.saturating_mul(n);
            let lhs_batch_count: usize = lhs_batch_dims.iter().product::<usize>().max(1);
            let rhs_batch_count: usize = rhs_batch_dims.iter().product::<usize>().max(1);
            let need_lhs = lhs_batch_count
                .saturating_mul(lhs_per_batch)
                .saturating_mul(elem);
            let need_rhs = rhs_batch_count
                .saturating_mul(rhs_per_batch)
                .saturating_mul(elem);
            let need_out = lhs_batch_count
                .saturating_mul(out_per_batch)
                .saturating_mul(elem);
            if lhs.len_bytes() != need_lhs {
                return Err(Error::Msg(format!(
                    "{}: lhs bytes={} doesn't match shape {:?} + [{m}, {k}]",
                    $type_name,
                    lhs.len_bytes(),
                    lhs_batch_dims,
                ))
                .bt());
            }
            if rhs.len_bytes() != need_rhs {
                return Err(Error::Msg(format!(
                    "{}: rhs bytes={} doesn't match shape {:?} + [{k}, {n}]",
                    $type_name,
                    rhs.len_bytes(),
                    rhs_batch_dims,
                ))
                .bt());
            }
            if out.len_bytes() != need_out {
                return Err(Error::Msg(format!(
                    "{}: out bytes={} doesn't match shape {:?} + [{m}, {n}]",
                    $type_name,
                    out.len_bytes(),
                    lhs_batch_dims,
                ))
                .bt());
            }
            let lhs_view: &[$T] = lhs.as_slice()?;
            let rhs_view: &[$T] = rhs.as_slice()?;
            let out_view: &mut [$T] = out.as_slice_mut()?;
            let mut lhs_multi = vec![0usize; batch_rank];
            let mut rhs_multi = vec![0usize; batch_rank];
            // Per-batch f32 accumulator buffer (reused across batches).
            let mut acc = vec![0.0_f32; out_per_batch];
            for b in 0..lhs_batch_count {
                let mut rem = b;
                for d in (0..batch_rank).rev() {
                    let s = lhs_batch_dims[d];
                    lhs_multi[d] = rem % s;
                    rem /= s;
                }
                for d in 0..batch_rank {
                    rhs_multi[d] = lhs_multi[d] / n_rep[d];
                }
                let mut rhs_b = 0usize;
                for d in 0..batch_rank {
                    rhs_b = rhs_b * rhs_batch_dims[d] + rhs_multi[d];
                }
                let lhs_off = b * lhs_per_batch;
                let rhs_off = rhs_b * rhs_per_batch;
                let out_off = b * out_per_batch;
                for slot in acc.iter_mut() {
                    *slot = 0.0;
                }
                for i in 0..m {
                    for kk in 0..k {
                        let a = lhs_view[lhs_off + i * k + kk].to_f32();
                        let rhs_row_off = rhs_off + kk * n;
                        let acc_row_off = i * n;
                        for j in 0..n {
                            acc[acc_row_off + j] += a * rhs_view[rhs_row_off + j].to_f32();
                        }
                    }
                }
                for (slot, &v) in out_view[out_off..out_off + out_per_batch]
                    .iter_mut()
                    .zip(&acc)
                {
                    *slot = <$T>::from_f32(v);
                }
            }
            Ok(())
        }
    };
}

matmul_half_kernel!(matmul_bf16, half::bf16, "matmul_bf16");
matmul_half_kernel!(matmul_f16, half::f16, "matmul_f16");

/// Batched row-major matmul for 8-bit integer types (i8/u8) with i32
/// accumulation and saturating cast on store. Mirrors baracuda's
/// `gemm_{s8,u8}_rrr_sm80_run` contract: `alpha = 1`, `beta = 0`,
/// identity epilogue, output type = input type. This is the CPU sibling
/// of the baracuda Int8 tensor-core kernel; it's correctness-first
/// (textbook (i, k, j) triple loop) and not a performance target.
macro_rules! matmul_int_kernel {
    ($name:ident, $T:ty, $type_name:literal) => {
        pub fn $name(
            lhs: &CpuStorageBytes,
            rhs: &CpuStorageBytes,
            out: &mut CpuStorageBytes,
            lhs_batch_dims: &[usize],
            rhs_batch_dims: &[usize],
            m: usize,
            n: usize,
            k: usize,
        ) -> Result<()> {
            if lhs_batch_dims.len() != rhs_batch_dims.len() {
                return Err(Error::Msg(format!(
                    "{}: batch ranks must match (lhs={}, rhs={})",
                    $type_name,
                    lhs_batch_dims.len(),
                    rhs_batch_dims.len(),
                ))
                .bt());
            }
            let batch_rank = lhs_batch_dims.len();
            let mut n_rep: Vec<usize> = Vec::with_capacity(batch_rank);
            for i in 0..batch_rank {
                let la = lhs_batch_dims[i];
                let ra = rhs_batch_dims[i];
                if la == ra {
                    n_rep.push(1);
                } else if ra > 0 && la > ra && la % ra == 0 {
                    n_rep.push(la / ra);
                } else {
                    return Err(Error::Msg(format!(
                        "{}: batch dim {i} disallowed combination (lhs={la}, rhs={ra})",
                        $type_name,
                    ))
                    .bt());
                }
            }
            let elem = std::mem::size_of::<$T>();
            let lhs_per_batch = m.saturating_mul(k);
            let rhs_per_batch = k.saturating_mul(n);
            let out_per_batch = m.saturating_mul(n);
            let lhs_batch_count: usize = lhs_batch_dims.iter().product::<usize>().max(1);
            let rhs_batch_count: usize = rhs_batch_dims.iter().product::<usize>().max(1);
            let need_lhs = lhs_batch_count
                .saturating_mul(lhs_per_batch)
                .saturating_mul(elem);
            let need_rhs = rhs_batch_count
                .saturating_mul(rhs_per_batch)
                .saturating_mul(elem);
            let need_out = lhs_batch_count
                .saturating_mul(out_per_batch)
                .saturating_mul(elem);
            if lhs.len_bytes() != need_lhs {
                return Err(Error::Msg(format!(
                    "{}: lhs bytes={} doesn't match shape {:?} + [{m}, {k}]",
                    $type_name,
                    lhs.len_bytes(),
                    lhs_batch_dims,
                ))
                .bt());
            }
            if rhs.len_bytes() != need_rhs {
                return Err(Error::Msg(format!(
                    "{}: rhs bytes={} doesn't match shape {:?} + [{k}, {n}]",
                    $type_name,
                    rhs.len_bytes(),
                    rhs_batch_dims,
                ))
                .bt());
            }
            if out.len_bytes() != need_out {
                return Err(Error::Msg(format!(
                    "{}: out bytes={} doesn't match shape {:?} + [{m}, {n}]",
                    $type_name,
                    out.len_bytes(),
                    lhs_batch_dims,
                ))
                .bt());
            }
            let lhs_view: &[$T] = lhs.as_slice()?;
            let rhs_view: &[$T] = rhs.as_slice()?;
            let out_view: &mut [$T] = out.as_slice_mut()?;
            let mut lhs_multi = vec![0usize; batch_rank];
            let mut rhs_multi = vec![0usize; batch_rank];
            // Per-batch i32 accumulator buffer (reused across batches).
            // i32 holds 128×128 i8 products without overflow up to K ≈
            // 2^31 / 128^2 ≈ 131k, which dwarfs realistic GEMM K.
            let mut acc = vec![0_i32; out_per_batch];
            const T_MIN: i32 = <$T>::MIN as i32;
            const T_MAX: i32 = <$T>::MAX as i32;
            for b in 0..lhs_batch_count {
                let mut rem = b;
                for d in (0..batch_rank).rev() {
                    let s = lhs_batch_dims[d];
                    lhs_multi[d] = rem % s;
                    rem /= s;
                }
                for d in 0..batch_rank {
                    rhs_multi[d] = lhs_multi[d] / n_rep[d];
                }
                let mut rhs_b = 0usize;
                for d in 0..batch_rank {
                    rhs_b = rhs_b * rhs_batch_dims[d] + rhs_multi[d];
                }
                let lhs_off = b * lhs_per_batch;
                let rhs_off = rhs_b * rhs_per_batch;
                let out_off = b * out_per_batch;
                for slot in acc.iter_mut() {
                    *slot = 0;
                }
                for i in 0..m {
                    for kk in 0..k {
                        let a = lhs_view[lhs_off + i * k + kk] as i32;
                        let rhs_row_off = rhs_off + kk * n;
                        let acc_row_off = i * n;
                        for j in 0..n {
                            acc[acc_row_off + j] += a * (rhs_view[rhs_row_off + j] as i32);
                        }
                    }
                }
                // Saturating cast i32 → T on store. Matches the CUDA
                // path's `.satfinite.s32.s8.s8.s32` MMA epilogue.
                for (slot, &v) in out_view[out_off..out_off + out_per_batch]
                    .iter_mut()
                    .zip(&acc)
                {
                    *slot = v.clamp(T_MIN, T_MAX) as $T;
                }
            }
            Ok(())
        }
    };
}

matmul_int_kernel!(matmul_i8, i8, "matmul_i8");
matmul_int_kernel!(matmul_u8, u8, "matmul_u8");

/// Batched row-major `f64` matrix multiply — a direct mirror of
/// [`matmul_f32`] with f64 element type and accumulator. Same
/// per-axis matched/GQA contract; same (i, k, j) inner loop.
pub fn matmul_f64(
    lhs: &CpuStorageBytes,
    rhs: &CpuStorageBytes,
    out: &mut CpuStorageBytes,
    lhs_batch_dims: &[usize],
    rhs_batch_dims: &[usize],
    m: usize,
    n: usize,
    k: usize,
) -> Result<()> {
    if lhs_batch_dims.len() != rhs_batch_dims.len() {
        return Err(Error::Msg(format!(
            "matmul_f64: batch ranks must match (lhs={}, rhs={})",
            lhs_batch_dims.len(),
            rhs_batch_dims.len(),
        ))
        .bt());
    }
    let batch_rank = lhs_batch_dims.len();
    let mut n_rep: Vec<usize> = Vec::with_capacity(batch_rank);
    for i in 0..batch_rank {
        let la = lhs_batch_dims[i];
        let ra = rhs_batch_dims[i];
        if la == ra {
            n_rep.push(1);
        } else if ra > 0 && la > ra && la.is_multiple_of(ra) {
            n_rep.push(la / ra);
        } else {
            return Err(Error::Msg(format!(
                "matmul_f64: batch dim {i} disallowed combination (lhs={la}, rhs={ra})",
            ))
            .bt());
        }
    }
    let elem = std::mem::size_of::<f64>();
    let lhs_per_batch = m.saturating_mul(k);
    let rhs_per_batch = k.saturating_mul(n);
    let out_per_batch = m.saturating_mul(n);
    let lhs_batch_count: usize = lhs_batch_dims.iter().product::<usize>().max(1);
    let rhs_batch_count: usize = rhs_batch_dims.iter().product::<usize>().max(1);
    let need_lhs = lhs_batch_count
        .saturating_mul(lhs_per_batch)
        .saturating_mul(elem);
    let need_rhs = rhs_batch_count
        .saturating_mul(rhs_per_batch)
        .saturating_mul(elem);
    let need_out = lhs_batch_count
        .saturating_mul(out_per_batch)
        .saturating_mul(elem);
    if lhs.len_bytes() != need_lhs {
        return Err(Error::Msg(format!(
            "matmul_f64: lhs bytes={} doesn't match shape {:?} + [{m}, {k}] (f64)",
            lhs.len_bytes(),
            lhs_batch_dims,
        ))
        .bt());
    }
    if rhs.len_bytes() != need_rhs {
        return Err(Error::Msg(format!(
            "matmul_f64: rhs bytes={} doesn't match shape {:?} + [{k}, {n}] (f64)",
            rhs.len_bytes(),
            rhs_batch_dims,
        ))
        .bt());
    }
    if out.len_bytes() != need_out {
        return Err(Error::Msg(format!(
            "matmul_f64: out bytes={} doesn't match shape {:?} + [{m}, {n}] (f64)",
            out.len_bytes(),
            lhs_batch_dims,
        ))
        .bt());
    }
    let lhs_view: &[f64] = lhs.as_slice()?;
    let rhs_view: &[f64] = rhs.as_slice()?;
    let out_view: &mut [f64] = out.as_slice_mut()?;
    for slot in out_view.iter_mut() {
        *slot = 0.0;
    }
    let mut lhs_multi = vec![0usize; batch_rank];
    let mut rhs_multi = vec![0usize; batch_rank];
    for b in 0..lhs_batch_count {
        let mut rem = b;
        for d in (0..batch_rank).rev() {
            let s = lhs_batch_dims[d];
            lhs_multi[d] = rem % s;
            rem /= s;
        }
        for d in 0..batch_rank {
            rhs_multi[d] = lhs_multi[d] / n_rep[d];
        }
        let mut rhs_b = 0usize;
        for d in 0..batch_rank {
            rhs_b = rhs_b * rhs_batch_dims[d] + rhs_multi[d];
        }
        let lhs_off = b * lhs_per_batch;
        let rhs_off = rhs_b * rhs_per_batch;
        let out_off = b * out_per_batch;
        for i in 0..m {
            for kk in 0..k {
                let a = lhs_view[lhs_off + i * k + kk];
                let rhs_row_off = rhs_off + kk * n;
                let out_row_off = out_off + i * n;
                for j in 0..n {
                    out_view[out_row_off + j] += a * rhs_view[rhs_row_off + j];
                }
            }
        }
    }
    Ok(())
}

// =============================================================================
// Quantized matmul (Q4_0 / Q8_0 / Q4_K_M weights, F32 activations)
// =============================================================================

/// Generic batched quantized-matmul kernel parameterized over the
/// quantized block type `T: GgmlType`. Activations are F32 with
/// shape `[batch, m, k]`; weights `[n, k / block_size]` blocks
/// (laid out row-major over `n`, then over `k`-blocks). Output is
/// F32 with shape `[batch, m, n]`.
///
/// fuel-quantized's `matmul` takes (m, k, n) + slices and computes
/// the inner products via SIMD-friendly per-column dot products.
/// This kernel iterates batches and calls into it.
/// Reinterpret a `&[u8]` byte stream as `&[T]` where `T` is a GGML
/// block type (`T: GgmlType` implies `bytemuck::Pod`, so the cast is
/// checked rather than asserted). Used by the quantized matmul
/// kernels. The borrow keeps the returned slice tied to `bytes`, which
/// is what makes this sound (GAP-336: fuel-quantized's former
/// `as_t_slice` took an owned `Cow` and returned a dangling slice).
/// Returns `Err` on length or alignment mismatch.
fn block_slice_from_bytes<'a, T: fuel_quantized::GgmlType>(
    name: &str,
    bytes: &'a [u8],
) -> Result<&'a [T]> {
    let size = std::mem::size_of::<T>();
    if size == 0 {
        return Err(Error::Msg(format!("{name}: zero-sized block type")).bt());
    }
    if !bytes.len().is_multiple_of(size) {
        return Err(Error::Msg(format!(
            "{name}: byte length {} not a multiple of block size {size}",
            bytes.len(),
        ))
        .bt());
    }
    // `T: GgmlType` implies `T: Pod`, so this is a checked reinterpretation:
    // it fails (never UB) on a misaligned pointer, and the length multiple was
    // validated above.
    bytemuck::try_cast_slice::<u8, T>(bytes).map_err(|_| {
        let align = std::mem::align_of::<T>();
        Error::Msg(format!(
            "{name}: byte pointer not aligned to block type's alignment {align}",
        ))
        .bt()
    })
}

fn qmatmul_generic_f32<T: fuel_quantized::GgmlType>(
    name: &str,
    activations: &CpuStorageBytes,
    weight_bytes: &CpuStorageBytes,
    out: &mut CpuStorageBytes,
    batch_count: usize,
    m: usize,
    n: usize,
    k: usize,
) -> Result<()> {
    let elements_per_block = T::BLCK_SIZE;
    if !k.is_multiple_of(elements_per_block) {
        return Err(Error::Msg(format!(
            "{name}: k={k} must be a multiple of block size {elements_per_block}",
        ))
        .bt());
    }
    let blocks_per_row = k / elements_per_block;
    let total_blocks = n.saturating_mul(blocks_per_row);
    let bytes_per_block = std::mem::size_of::<T>();
    let need_w = total_blocks.saturating_mul(bytes_per_block);
    let need_a = batch_count
        .saturating_mul(m)
        .saturating_mul(k)
        .saturating_mul(std::mem::size_of::<f32>());
    let need_out = batch_count
        .saturating_mul(m)
        .saturating_mul(n)
        .saturating_mul(std::mem::size_of::<f32>());
    if activations.len_bytes() != need_a {
        return Err(Error::Msg(format!(
            "{name}: activations bytes={} doesn't match batch={batch_count} × m={m} × k={k} × 4",
            activations.len_bytes(),
        ))
        .bt());
    }
    if weight_bytes.len_bytes() != need_w {
        return Err(Error::Msg(format!(
            "{name}: weight bytes={} doesn't match n={n} × k/block_size={blocks_per_row} × {bytes_per_block}",
            weight_bytes.len_bytes(),
        ))
        .bt());
    }
    if out.len_bytes() != need_out {
        return Err(Error::Msg(format!(
            "{name}: out bytes={} doesn't match batch={batch_count} × m={m} × n={n} × 4",
            out.len_bytes(),
        ))
        .bt());
    }
    if batch_count == 0 || m == 0 || n == 0 {
        return Ok(());
    }
    let act_view: &[f32] = activations.as_slice()?;
    // Reinterpret the U32-typed weight bytes as `&[T]` blocks via
    // a safety-wrapper around `from_raw_parts`. The GGML block
    // types are `#[repr(C)]` POD layouts so the cast is sound.
    let weight_view: &[T] = block_slice_from_bytes::<T>(name, weight_bytes.bytes())?;
    let out_view: &mut [f32] = out.as_slice_mut()?;
    let act_per_batch = m * k;
    let out_per_batch = m * n;
    for b in 0..batch_count {
        let act_off = b * act_per_batch;
        let out_off = b * out_per_batch;
        fuel_quantized::matmul::<T>(
            (m, k, n),
            &act_view[act_off..act_off + act_per_batch],
            weight_view,
            &mut out_view[out_off..out_off + out_per_batch],
        )
        .map_err(|e| Error::Msg(format!("{name}: {e}")).bt())?;
    }
    Ok(())
}

/// Q4_0 quantized matmul. Activations F32, weights `[n,
/// k/32]` `BlockQ4_0`s, output F32.
pub fn qmatmul_q4_0_f32(
    activations: &CpuStorageBytes,
    weight_bytes: &CpuStorageBytes,
    out: &mut CpuStorageBytes,
    batch_count: usize,
    m: usize,
    n: usize,
    k: usize,
) -> Result<()> {
    qmatmul_generic_f32::<fuel_quantized::BlockQ4_0>(
        "qmatmul_q4_0_f32",
        activations,
        weight_bytes,
        out,
        batch_count,
        m,
        n,
        k,
    )
}

/// Q8_0 quantized matmul.
pub fn qmatmul_q8_0_f32(
    activations: &CpuStorageBytes,
    weight_bytes: &CpuStorageBytes,
    out: &mut CpuStorageBytes,
    batch_count: usize,
    m: usize,
    n: usize,
    k: usize,
) -> Result<()> {
    qmatmul_generic_f32::<fuel_quantized::BlockQ8_0>(
        "qmatmul_q8_0_f32",
        activations,
        weight_bytes,
        out,
        batch_count,
        m,
        n,
        k,
    )
}

/// Q4_K_M (256-element super-block) quantized matmul.
pub fn qmatmul_q4_k_m_f32(
    activations: &CpuStorageBytes,
    weight_bytes: &CpuStorageBytes,
    out: &mut CpuStorageBytes,
    batch_count: usize,
    m: usize,
    n: usize,
    k: usize,
) -> Result<()> {
    qmatmul_generic_f32::<fuel_quantized::BlockQ4K>(
        "qmatmul_q4_k_m_f32",
        activations,
        weight_bytes,
        out,
        batch_count,
        m,
        n,
        k,
    )
}

macro_rules! qmatmul_thin_wrapper {
    ($name:ident, $blk:ty, $kname:expr) => {
        pub fn $name(
            activations: &CpuStorageBytes,
            weight_bytes: &CpuStorageBytes,
            out: &mut CpuStorageBytes,
            batch_count: usize,
            m: usize,
            n: usize,
            k: usize,
        ) -> Result<()> {
            qmatmul_generic_f32::<$blk>(
                $kname,
                activations,
                weight_bytes,
                out,
                batch_count,
                m,
                n,
                k,
            )
        }
    };
}

qmatmul_thin_wrapper!(
    qmatmul_q4_1_f32,
    fuel_quantized::BlockQ4_1,
    "qmatmul_q4_1_f32"
);
qmatmul_thin_wrapper!(
    qmatmul_q5_0_f32,
    fuel_quantized::BlockQ5_0,
    "qmatmul_q5_0_f32"
);
qmatmul_thin_wrapper!(
    qmatmul_q5_1_f32,
    fuel_quantized::BlockQ5_1,
    "qmatmul_q5_1_f32"
);
qmatmul_thin_wrapper!(
    qmatmul_q8_1_f32,
    fuel_quantized::BlockQ8_1,
    "qmatmul_q8_1_f32"
);
qmatmul_thin_wrapper!(qmatmul_q2k_f32, fuel_quantized::BlockQ2K, "qmatmul_q2k_f32");
qmatmul_thin_wrapper!(qmatmul_q3k_f32, fuel_quantized::BlockQ3K, "qmatmul_q3k_f32");
qmatmul_thin_wrapper!(qmatmul_q5k_f32, fuel_quantized::BlockQ5K, "qmatmul_q5k_f32");
qmatmul_thin_wrapper!(qmatmul_q6k_f32, fuel_quantized::BlockQ6K, "qmatmul_q6k_f32");

// =============================================================================
// 2D Convolution — multi-dtype (f64 native, bf16/f16 via f32 acc)
// =============================================================================

/// Generate a Conv2D kernel where input/weight/output are `$T`
/// (native arithmetic). Used for f64 (and f32 has its own
/// hand-written version).
macro_rules! conv2d_native_kernel {
    ($name:ident, $T:ty, $T_size:expr, $zero:expr) => {
        pub fn $name(
            x: &CpuStorageBytes,
            weight: &CpuStorageBytes,
            bias: Option<&CpuStorageBytes>,
            out: &mut CpuStorageBytes,
            x_shape: [usize; 4],
            w_shape: [usize; 4],
            out_shape: [usize; 4],
            stride: (usize, usize),
            padding: (usize, usize),
            dilation: (usize, usize),
            groups: usize,
        ) -> Result<()> {
            let [n, cin, h_in, w_in] = x_shape;
            let [cout, cin_per_group, kh, kw] = w_shape;
            let [n_out, cout_out, h_out, w_out] = out_shape;
            if n != n_out || cout != cout_out
                || groups == 0 || cin % groups != 0 || cout % groups != 0
                || cin / groups != cin_per_group
            {
                return Err(Error::Msg(format!(
                    "{}: shape contract violation (x={x_shape:?}, w={w_shape:?}, out={out_shape:?}, groups={groups})",
                    stringify!($name),
                ))
                .bt());
            }
            let cout_per_group = cout / groups;
            let elem = $T_size;
            if x.len_bytes() != n * cin * h_in * w_in * elem
                || weight.len_bytes() != cout * cin_per_group * kh * kw * elem
                || out.len_bytes() != n * cout * h_out * w_out * elem
            {
                return Err(Error::Msg(format!(
                    "{}: bytes mismatch", stringify!($name),
                ))
                .bt());
            }
            if let Some(b) = bias {
                if b.len_bytes() != cout * elem {
                    return Err(Error::Msg(format!(
                        "{}: bias bytes mismatch", stringify!($name),
                    ))
                    .bt());
                }
            }
            let x_view: &[$T] = x.as_slice()?;
            let w_view: &[$T] = weight.as_slice()?;
            let bias_view: Option<&[$T]> = match bias {
                Some(b) => Some(b.as_slice()?),
                None => None,
            };
            let out_view: &mut [$T] = out.as_slice_mut()?;
            let (sh, sw) = stride;
            let (ph, pw) = padding;
            let (dh, dw) = dilation;
            for b_idx in 0..n {
                for co in 0..cout {
                    let group = co / cout_per_group;
                    let ci_offset = group * cin_per_group;
                    let bias_val = bias_view.map(|bv| bv[co]).unwrap_or($zero);
                    for oh in 0..h_out {
                        for ow in 0..w_out {
                            let mut acc = bias_val;
                            for ci in 0..cin_per_group {
                                for kh_i in 0..kh {
                                    let in_h = (oh * sh + kh_i * dh) as isize - ph as isize;
                                    if in_h < 0 || in_h as usize >= h_in { continue; }
                                    let in_h = in_h as usize;
                                    for kw_i in 0..kw {
                                        let in_w = (ow * sw + kw_i * dw) as isize - pw as isize;
                                        if in_w < 0 || in_w as usize >= w_in { continue; }
                                        let in_w = in_w as usize;
                                        let x_idx = ((b_idx * cin + (ci_offset + ci)) * h_in + in_h) * w_in + in_w;
                                        let w_idx = ((co * cin_per_group + ci) * kh + kh_i) * kw + kw_i;
                                        acc += x_view[x_idx] * w_view[w_idx];
                                    }
                                }
                            }
                            let out_idx = ((b_idx * cout + co) * h_out + oh) * w_out + ow;
                            out_view[out_idx] = acc;
                        }
                    }
                }
            }
            Ok(())
        }
    };
}

conv2d_native_kernel!(conv2d_f64, f64, std::mem::size_of::<f64>(), 0.0_f64);

/// Conv2D for half-float types — accumulates each output position
/// in f32, narrows at the end. Inputs (x, weight, bias) all in T.
macro_rules! conv2d_half_kernel {
    ($name:ident, $T:ty) => {
        pub fn $name(
            x: &CpuStorageBytes,
            weight: &CpuStorageBytes,
            bias: Option<&CpuStorageBytes>,
            out: &mut CpuStorageBytes,
            x_shape: [usize; 4],
            w_shape: [usize; 4],
            out_shape: [usize; 4],
            stride: (usize, usize),
            padding: (usize, usize),
            dilation: (usize, usize),
            groups: usize,
        ) -> Result<()> {
            let [n, cin, h_in, w_in] = x_shape;
            let [cout, cin_per_group, kh, kw] = w_shape;
            let [n_out, cout_out, h_out, w_out] = out_shape;
            if n != n_out
                || cout != cout_out
                || groups == 0
                || cin % groups != 0
                || cout % groups != 0
                || cin / groups != cin_per_group
            {
                return Err(Error::Msg(
                    format!("{}: shape contract violation", stringify!($name),),
                )
                .bt());
            }
            let cout_per_group = cout / groups;
            let elem = std::mem::size_of::<$T>();
            if x.len_bytes() != n * cin * h_in * w_in * elem
                || weight.len_bytes() != cout * cin_per_group * kh * kw * elem
                || out.len_bytes() != n * cout * h_out * w_out * elem
            {
                return Err(Error::Msg(format!("{}: bytes mismatch", stringify!($name),)).bt());
            }
            if let Some(b) = bias {
                if b.len_bytes() != cout * elem {
                    return Err(
                        Error::Msg(format!("{}: bias bytes mismatch", stringify!($name),)).bt(),
                    );
                }
            }
            let x_view: &[$T] = x.as_slice()?;
            let w_view: &[$T] = weight.as_slice()?;
            let bias_view: Option<&[$T]> = match bias {
                Some(b) => Some(b.as_slice()?),
                None => None,
            };
            let out_view: &mut [$T] = out.as_slice_mut()?;
            let (sh, sw) = stride;
            let (ph, pw) = padding;
            let (dh, dw) = dilation;
            for b_idx in 0..n {
                for co in 0..cout {
                    let group = co / cout_per_group;
                    let ci_offset = group * cin_per_group;
                    let bias_val = bias_view.map(|bv| bv[co].to_f32()).unwrap_or(0.0_f32);
                    for oh in 0..h_out {
                        for ow in 0..w_out {
                            let mut acc: f32 = bias_val;
                            for ci in 0..cin_per_group {
                                for kh_i in 0..kh {
                                    let in_h = (oh * sh + kh_i * dh) as isize - ph as isize;
                                    if in_h < 0 || in_h as usize >= h_in {
                                        continue;
                                    }
                                    let in_h = in_h as usize;
                                    for kw_i in 0..kw {
                                        let in_w = (ow * sw + kw_i * dw) as isize - pw as isize;
                                        if in_w < 0 || in_w as usize >= w_in {
                                            continue;
                                        }
                                        let in_w = in_w as usize;
                                        let x_idx =
                                            ((b_idx * cin + (ci_offset + ci)) * h_in + in_h) * w_in
                                                + in_w;
                                        let w_idx =
                                            ((co * cin_per_group + ci) * kh + kh_i) * kw + kw_i;
                                        acc += x_view[x_idx].to_f32() * w_view[w_idx].to_f32();
                                    }
                                }
                            }
                            let out_idx = ((b_idx * cout + co) * h_out + oh) * w_out + ow;
                            out_view[out_idx] = <$T>::from_f32(acc);
                        }
                    }
                }
            }
            Ok(())
        }
    };
}

conv2d_half_kernel!(conv2d_bf16, half::bf16);
conv2d_half_kernel!(conv2d_f16, half::f16);

// =============================================================================
// 2D Transposed Convolution — multi-dtype
// =============================================================================
//
// Shapes:
//   x:      [N, Cin, H_in, W_in]
//   weight: [Cin, Cout/groups, Kh, Kw]   (transposed channel order vs Conv2D)
//   bias:   optional [Cout]
//   out:    [N, Cout, H_out, W_out]
// where:
//   H_out = (H_in - 1) * stride.0 - 2*pad.0 + dil.0*(Kh - 1) + out_pad.0 + 1
//   W_out = (W_in - 1) * stride.1 - 2*pad.1 + dil.1*(Kw - 1) + out_pad.1 + 1
//
// Strategy: zero output, optionally seed with broadcast bias, then for
// every input element scatter-accumulate its kernel-shaped contribution
// into the output. For half-floats we accumulate into a parallel
// `Vec<f32>` buffer and narrow at the end (same f32-accumulator
// pattern Conv2D-half uses).

/// ConvTranspose2D for native arithmetic ($T = f32 or f64).
macro_rules! conv_transpose2d_native_kernel {
    ($name:ident, $T:ty, $T_size:expr, $zero:expr) => {
        pub fn $name(
            x: &CpuStorageBytes,
            weight: &CpuStorageBytes,
            bias: Option<&CpuStorageBytes>,
            out: &mut CpuStorageBytes,
            x_shape: [usize; 4],
            w_shape: [usize; 4],
            out_shape: [usize; 4],
            stride: (usize, usize),
            padding: (usize, usize),
            dilation: (usize, usize),
            groups: usize,
        ) -> Result<()> {
            let [n, cin, h_in, w_in] = x_shape;
            let [cin_w, cout_per_group, kh, kw] = w_shape;
            let [n_out, cout, h_out, w_out] = out_shape;
            if n != n_out
                || groups == 0
                || cin % groups != 0
                || cout % groups != 0
                || cin != cin_w
                || cout / groups != cout_per_group
            {
                return Err(Error::Msg(format!(
                    "{}: shape contract violation (x={x_shape:?}, w={w_shape:?}, out={out_shape:?}, groups={groups})",
                    stringify!($name),
                ))
                .bt());
            }
            let elem = $T_size;
            if x.len_bytes() != n * cin * h_in * w_in * elem
                || weight.len_bytes() != cin * cout_per_group * kh * kw * elem
                || out.len_bytes() != n * cout * h_out * w_out * elem
            {
                return Err(Error::Msg(format!(
                    "{}: bytes mismatch", stringify!($name),
                ))
                .bt());
            }
            if let Some(b) = bias {
                if b.len_bytes() != cout * elem {
                    return Err(Error::Msg(format!(
                        "{}: bias bytes mismatch", stringify!($name),
                    ))
                    .bt());
                }
            }
            let x_view: &[$T] = x.as_slice()?;
            let w_view: &[$T] = weight.as_slice()?;
            let bias_view: Option<&[$T]> = match bias {
                Some(b) => Some(b.as_slice()?),
                None => None,
            };
            let out_view: &mut [$T] = out.as_slice_mut()?;
            // Initialize output: bias broadcast across spatial dims, or zero.
            for n_i in 0..n {
                for co in 0..cout {
                    let bias_val = bias_view.map(|bv| bv[co]).unwrap_or($zero);
                    for oh in 0..h_out {
                        for ow in 0..w_out {
                            let idx = ((n_i * cout + co) * h_out + oh) * w_out + ow;
                            out_view[idx] = bias_val;
                        }
                    }
                }
            }
            let (sh, sw) = stride;
            let (ph, pw) = padding;
            let (dh, dw) = dilation;
            let cin_per_group = cin / groups;
            for n_i in 0..n {
                for g in 0..groups {
                    for ci_local in 0..cin_per_group {
                        let ci = g * cin_per_group + ci_local;
                        for hi in 0..h_in {
                            for wi in 0..w_in {
                                let val = x_view[((n_i * cin + ci) * h_in + hi) * w_in + wi];
                                if val == $zero { continue; }
                                for co_local in 0..cout_per_group {
                                    let co = g * cout_per_group + co_local;
                                    for kh_i in 0..kh {
                                        let oh_signed = (hi * sh) as isize + (kh_i * dh) as isize - ph as isize;
                                        if oh_signed < 0 || oh_signed as usize >= h_out { continue; }
                                        let oh = oh_signed as usize;
                                        for kw_i in 0..kw {
                                            let ow_signed = (wi * sw) as isize + (kw_i * dw) as isize - pw as isize;
                                            if ow_signed < 0 || ow_signed as usize >= w_out { continue; }
                                            let ow = ow_signed as usize;
                                            let w_idx = ((ci * cout_per_group + co_local) * kh + kh_i) * kw + kw_i;
                                            let out_idx = ((n_i * cout + co) * h_out + oh) * w_out + ow;
                                            out_view[out_idx] += val * w_view[w_idx];
                                        }
                                    }
                                }
                            }
                        }
                    }
                }
            }
            Ok(())
        }
    };
}

conv_transpose2d_native_kernel!(
    conv_transpose2d_f32,
    f32,
    std::mem::size_of::<f32>(),
    0.0_f32
);
conv_transpose2d_native_kernel!(
    conv_transpose2d_f64,
    f64,
    std::mem::size_of::<f64>(),
    0.0_f64
);

/// ConvTranspose2D for half-floats — accumulates into a parallel
/// `Vec<f32>` buffer, narrows back at the end. Same f32-accumulator
/// pattern as conv2d_half_kernel!.
macro_rules! conv_transpose2d_half_kernel {
    ($name:ident, $T:ty) => {
        pub fn $name(
            x: &CpuStorageBytes,
            weight: &CpuStorageBytes,
            bias: Option<&CpuStorageBytes>,
            out: &mut CpuStorageBytes,
            x_shape: [usize; 4],
            w_shape: [usize; 4],
            out_shape: [usize; 4],
            stride: (usize, usize),
            padding: (usize, usize),
            dilation: (usize, usize),
            groups: usize,
        ) -> Result<()> {
            let [n, cin, h_in, w_in] = x_shape;
            let [cin_w, cout_per_group, kh, kw] = w_shape;
            let [n_out, cout, h_out, w_out] = out_shape;
            if n != n_out
                || groups == 0
                || cin % groups != 0
                || cout % groups != 0
                || cin != cin_w
                || cout / groups != cout_per_group
            {
                return Err(Error::Msg(
                    format!("{}: shape contract violation", stringify!($name),),
                )
                .bt());
            }
            let elem = std::mem::size_of::<$T>();
            if x.len_bytes() != n * cin * h_in * w_in * elem
                || weight.len_bytes() != cin * cout_per_group * kh * kw * elem
                || out.len_bytes() != n * cout * h_out * w_out * elem
            {
                return Err(Error::Msg(format!("{}: bytes mismatch", stringify!($name),)).bt());
            }
            if let Some(b) = bias {
                if b.len_bytes() != cout * elem {
                    return Err(
                        Error::Msg(format!("{}: bias bytes mismatch", stringify!($name),)).bt(),
                    );
                }
            }
            let x_view: &[$T] = x.as_slice()?;
            let w_view: &[$T] = weight.as_slice()?;
            let bias_view: Option<&[$T]> = match bias {
                Some(b) => Some(b.as_slice()?),
                None => None,
            };
            let out_view: &mut [$T] = out.as_slice_mut()?;
            // f32 accumulator buffer; seeded with bias broadcast.
            let total = n * cout * h_out * w_out;
            let mut acc = vec![0.0_f32; total];
            for n_i in 0..n {
                for co in 0..cout {
                    let bias_val = bias_view.map(|bv| bv[co].to_f32()).unwrap_or(0.0_f32);
                    for oh in 0..h_out {
                        for ow in 0..w_out {
                            let idx = ((n_i * cout + co) * h_out + oh) * w_out + ow;
                            acc[idx] = bias_val;
                        }
                    }
                }
            }
            let (sh, sw) = stride;
            let (ph, pw) = padding;
            let (dh, dw) = dilation;
            let cin_per_group = cin / groups;
            for n_i in 0..n {
                for g in 0..groups {
                    for ci_local in 0..cin_per_group {
                        let ci = g * cin_per_group + ci_local;
                        for hi in 0..h_in {
                            for wi in 0..w_in {
                                let val =
                                    x_view[((n_i * cin + ci) * h_in + hi) * w_in + wi].to_f32();
                                if val == 0.0_f32 {
                                    continue;
                                }
                                for co_local in 0..cout_per_group {
                                    let co = g * cout_per_group + co_local;
                                    for kh_i in 0..kh {
                                        let oh_signed =
                                            (hi * sh) as isize + (kh_i * dh) as isize - ph as isize;
                                        if oh_signed < 0 || oh_signed as usize >= h_out {
                                            continue;
                                        }
                                        let oh = oh_signed as usize;
                                        for kw_i in 0..kw {
                                            let ow_signed = (wi * sw) as isize
                                                + (kw_i * dw) as isize
                                                - pw as isize;
                                            if ow_signed < 0 || ow_signed as usize >= w_out {
                                                continue;
                                            }
                                            let ow = ow_signed as usize;
                                            let w_idx =
                                                ((ci * cout_per_group + co_local) * kh + kh_i) * kw
                                                    + kw_i;
                                            let out_idx =
                                                ((n_i * cout + co) * h_out + oh) * w_out + ow;
                                            acc[out_idx] += val * w_view[w_idx].to_f32();
                                        }
                                    }
                                }
                            }
                        }
                    }
                }
            }
            for (dst, src) in out_view.iter_mut().zip(acc.iter()) {
                *dst = <$T>::from_f32(*src);
            }
            Ok(())
        }
    };
}

conv_transpose2d_half_kernel!(conv_transpose2d_bf16, half::bf16);
conv_transpose2d_half_kernel!(conv_transpose2d_f16, half::f16);

// =============================================================================
// ReduceSumTo / ReduceMaxTo — sum / max to a broadcast-compatible
// target shape. All 8 entries (2 ops × 4 dtypes) are thunks over
// [`reduce_to`].
// =============================================================================
//
// Shape contract (preserved verbatim from the pre-chassis macros):
// `output_shape` is left-padded with 1s if its rank < input rank.
// For each padded axis: if the value matches the input's size, the
// axis carries through; if 1, the axis is reduced away; any other
// value is a contract violation. Half-float accumulation runs in
// f32 (the trait's `Acc = f32` impl for `Sum<bf16>` / `Sum<f16>` /
// `Max<bf16>` / `Max<f16>`).

/// Sum-reduce `f32` to a broadcast-compatible target shape.
pub fn reduce_sum_to_f32(
    input: &CpuStorageBytes,
    output: &mut CpuStorageBytes,
    input_shape: &[usize],
    output_shape: &[usize],
) -> Result<()> {
    reduce_to::<f32, Sum>(
        "reduce_sum_to_f32",
        input,
        output,
        input_shape,
        output_shape,
    )
}

/// Sum-reduce `f64` to a broadcast-compatible target shape.
pub fn reduce_sum_to_f64(
    input: &CpuStorageBytes,
    output: &mut CpuStorageBytes,
    input_shape: &[usize],
    output_shape: &[usize],
) -> Result<()> {
    reduce_to::<f64, Sum>(
        "reduce_sum_to_f64",
        input,
        output,
        input_shape,
        output_shape,
    )
}

/// Sum-reduce `bf16` to a broadcast-compatible target shape.
/// Accumulator runs in f32.
pub fn reduce_sum_to_bf16(
    input: &CpuStorageBytes,
    output: &mut CpuStorageBytes,
    input_shape: &[usize],
    output_shape: &[usize],
) -> Result<()> {
    reduce_to::<half::bf16, Sum>(
        "reduce_sum_to_bf16",
        input,
        output,
        input_shape,
        output_shape,
    )
}

/// Sum-reduce `f16` to a broadcast-compatible target shape.
/// Accumulator runs in f32.
pub fn reduce_sum_to_f16(
    input: &CpuStorageBytes,
    output: &mut CpuStorageBytes,
    input_shape: &[usize],
    output_shape: &[usize],
) -> Result<()> {
    reduce_to::<half::f16, Sum>(
        "reduce_sum_to_f16",
        input,
        output,
        input_shape,
        output_shape,
    )
}

/// Max-reduce `f32` to a broadcast-compatible target shape.
pub fn reduce_max_to_f32(
    input: &CpuStorageBytes,
    output: &mut CpuStorageBytes,
    input_shape: &[usize],
    output_shape: &[usize],
) -> Result<()> {
    reduce_to::<f32, Max>(
        "reduce_max_to_f32",
        input,
        output,
        input_shape,
        output_shape,
    )
}

/// Max-reduce `f64` to a broadcast-compatible target shape.
pub fn reduce_max_to_f64(
    input: &CpuStorageBytes,
    output: &mut CpuStorageBytes,
    input_shape: &[usize],
    output_shape: &[usize],
) -> Result<()> {
    reduce_to::<f64, Max>(
        "reduce_max_to_f64",
        input,
        output,
        input_shape,
        output_shape,
    )
}

/// Max-reduce `bf16` to a broadcast-compatible target shape.
/// Extremum runs in f32, narrows back to bf16.
pub fn reduce_max_to_bf16(
    input: &CpuStorageBytes,
    output: &mut CpuStorageBytes,
    input_shape: &[usize],
    output_shape: &[usize],
) -> Result<()> {
    reduce_to::<half::bf16, Max>(
        "reduce_max_to_bf16",
        input,
        output,
        input_shape,
        output_shape,
    )
}

/// Max-reduce `f16` to a broadcast-compatible target shape.
pub fn reduce_max_to_f16(
    input: &CpuStorageBytes,
    output: &mut CpuStorageBytes,
    input_shape: &[usize],
    output_shape: &[usize],
) -> Result<()> {
    reduce_to::<half::f16, Max>(
        "reduce_max_to_f16",
        input,
        output,
        input_shape,
        output_shape,
    )
}

// =============================================================================
// FusedLinear — matmul + bias-add (3-input)
// =============================================================================
//
// Inputs:
//   a:    [..., M, K]
//   b:    [..., K, N]
//   bias: [N]   (broadcast across all leading dims of the matmul output)
// Output: [..., M, N] where out[..., i, j] = bias[j] + sum_k a[..., i, k] * b[..., k, j]
//
// Shape semantics match `matmul_*` (per-axis GQA broadcasting on the
// batch prefixes). The fused form simply seeds the accumulator with
// bias[j] before the inner product.

fn fused_linear_check<T>(
    name: &str,
    lhs: &CpuStorageBytes,
    rhs: &CpuStorageBytes,
    bias: &CpuStorageBytes,
    out: &CpuStorageBytes,
    lhs_batch_dims: &[usize],
    rhs_batch_dims: &[usize],
    m: usize,
    n: usize,
    k: usize,
) -> Result<Vec<usize>> {
    if lhs_batch_dims.len() != rhs_batch_dims.len() {
        return Err(Error::Msg(format!(
            "{name}: batch ranks must match (lhs={}, rhs={})",
            lhs_batch_dims.len(),
            rhs_batch_dims.len(),
        ))
        .bt());
    }
    let batch_rank = lhs_batch_dims.len();
    let mut n_rep: Vec<usize> = Vec::with_capacity(batch_rank);
    for i in 0..batch_rank {
        let la = lhs_batch_dims[i];
        let ra = rhs_batch_dims[i];
        if la == ra {
            n_rep.push(1);
        } else if ra > 0 && la > ra && la.is_multiple_of(ra) {
            n_rep.push(la / ra);
        } else {
            return Err(Error::Msg(format!(
                "{name}: batch dim {i} disallowed (lhs={la}, rhs={ra})",
            ))
            .bt());
        }
    }
    let elem = std::mem::size_of::<T>();
    let lhs_per = m.saturating_mul(k);
    let rhs_per = k.saturating_mul(n);
    let out_per = m.saturating_mul(n);
    let lhs_count: usize = lhs_batch_dims.iter().product::<usize>().max(1);
    let rhs_count: usize = rhs_batch_dims.iter().product::<usize>().max(1);
    if lhs.len_bytes() != lhs_count.saturating_mul(lhs_per).saturating_mul(elem)
        || rhs.len_bytes() != rhs_count.saturating_mul(rhs_per).saturating_mul(elem)
        || out.len_bytes() != lhs_count.saturating_mul(out_per).saturating_mul(elem)
        || bias.len_bytes() != n.saturating_mul(elem)
    {
        return Err(Error::Msg(format!("{name}: bytes mismatch",)).bt());
    }
    Ok(n_rep)
}

/// FusedLinear for native arithmetic.
macro_rules! fused_linear_native_kernel {
    ($name:ident, $T:ty) => {
        pub fn $name(
            lhs: &CpuStorageBytes,
            rhs: &CpuStorageBytes,
            bias: &CpuStorageBytes,
            out: &mut CpuStorageBytes,
            lhs_batch_dims: &[usize],
            rhs_batch_dims: &[usize],
            m: usize,
            n: usize,
            k: usize,
        ) -> Result<()> {
            let n_rep = fused_linear_check::<$T>(
                stringify!($name),
                lhs,
                rhs,
                bias,
                out,
                lhs_batch_dims,
                rhs_batch_dims,
                m,
                n,
                k,
            )?;
            let lhs_view: &[$T] = lhs.as_slice()?;
            let rhs_view: &[$T] = rhs.as_slice()?;
            let bias_view: &[$T] = bias.as_slice()?;
            let out_view: &mut [$T] = out.as_slice_mut()?;
            let batch_rank = lhs_batch_dims.len();
            let lhs_per = m * k;
            let rhs_per = k * n;
            let out_per = m * n;
            let lhs_count: usize = lhs_batch_dims.iter().product::<usize>().max(1);
            let mut lhs_multi = vec![0_usize; batch_rank];
            let mut rhs_multi = vec![0_usize; batch_rank];
            for b in 0..lhs_count {
                let mut rem = b;
                for d in (0..batch_rank).rev() {
                    let s = lhs_batch_dims[d];
                    lhs_multi[d] = rem % s;
                    rem /= s;
                }
                for d in 0..batch_rank {
                    rhs_multi[d] = lhs_multi[d] / n_rep[d];
                }
                let mut rhs_b = 0_usize;
                for d in 0..batch_rank {
                    rhs_b = rhs_b * rhs_batch_dims[d] + rhs_multi[d];
                }
                let lhs_off = b * lhs_per;
                let rhs_off = rhs_b * rhs_per;
                let out_off = b * out_per;
                // Seed each output row with the broadcast bias.
                for i in 0..m {
                    let row_off = out_off + i * n;
                    out_view[row_off..row_off + n].copy_from_slice(bias_view);
                }
                for i in 0..m {
                    for kk in 0..k {
                        let a = lhs_view[lhs_off + i * k + kk];
                        let rhs_row_off = rhs_off + kk * n;
                        let out_row_off = out_off + i * n;
                        for j in 0..n {
                            out_view[out_row_off + j] += a * rhs_view[rhs_row_off + j];
                        }
                    }
                }
            }
            Ok(())
        }
    };
}

fused_linear_native_kernel!(fused_linear_f32, f32);
fused_linear_native_kernel!(fused_linear_f64, f64);

/// FusedLinear for half-floats — accumulate in f32, narrow at end.
macro_rules! fused_linear_half_kernel {
    ($name:ident, $T:ty) => {
        pub fn $name(
            lhs: &CpuStorageBytes,
            rhs: &CpuStorageBytes,
            bias: &CpuStorageBytes,
            out: &mut CpuStorageBytes,
            lhs_batch_dims: &[usize],
            rhs_batch_dims: &[usize],
            m: usize,
            n: usize,
            k: usize,
        ) -> Result<()> {
            let n_rep = fused_linear_check::<$T>(
                stringify!($name),
                lhs,
                rhs,
                bias,
                out,
                lhs_batch_dims,
                rhs_batch_dims,
                m,
                n,
                k,
            )?;
            let lhs_view: &[$T] = lhs.as_slice()?;
            let rhs_view: &[$T] = rhs.as_slice()?;
            let bias_view: &[$T] = bias.as_slice()?;
            let out_view: &mut [$T] = out.as_slice_mut()?;
            let bias_f32: Vec<f32> = bias_view.iter().map(|v| v.to_f32()).collect();
            let batch_rank = lhs_batch_dims.len();
            let lhs_per = m * k;
            let rhs_per = k * n;
            let out_per = m * n;
            let lhs_count: usize = lhs_batch_dims.iter().product::<usize>().max(1);
            let mut lhs_multi = vec![0_usize; batch_rank];
            let mut rhs_multi = vec![0_usize; batch_rank];
            let mut row_acc = vec![0.0_f32; n];
            for b in 0..lhs_count {
                let mut rem = b;
                for d in (0..batch_rank).rev() {
                    let s = lhs_batch_dims[d];
                    lhs_multi[d] = rem % s;
                    rem /= s;
                }
                for d in 0..batch_rank {
                    rhs_multi[d] = lhs_multi[d] / n_rep[d];
                }
                let mut rhs_b = 0_usize;
                for d in 0..batch_rank {
                    rhs_b = rhs_b * rhs_batch_dims[d] + rhs_multi[d];
                }
                let lhs_off = b * lhs_per;
                let rhs_off = rhs_b * rhs_per;
                let out_off = b * out_per;
                for i in 0..m {
                    row_acc.copy_from_slice(&bias_f32);
                    for kk in 0..k {
                        let a = lhs_view[lhs_off + i * k + kk].to_f32();
                        let rhs_row_off = rhs_off + kk * n;
                        for j in 0..n {
                            row_acc[j] += a * rhs_view[rhs_row_off + j].to_f32();
                        }
                    }
                    let out_row_off = out_off + i * n;
                    for j in 0..n {
                        out_view[out_row_off + j] = <$T>::from_f32(row_acc[j]);
                    }
                }
            }
            Ok(())
        }
    };
}

fused_linear_half_kernel!(fused_linear_bf16, half::bf16);
fused_linear_half_kernel!(fused_linear_f16, half::f16);

// =============================================================================
// FusedSoftmaxCrossEntropy — fused softmax + NLL + reduction
// =============================================================================
//
// The textbook stable form computed row-by-row:
//
//   row_max = max(logits_row)
//   sum_exp = sum(exp(logits_row[v] - row_max) for v in 0..V)
//   nll_row = -(logits_row[target] - row_max - log(sum_exp))
//
// `ignore_index` rows contribute 0 to the loss sum and 0 to the
// denominator. For `Reduction::Mean` the denominator is the count of
// non-ignored rows (so an all-ignored input yields 0.0 — matches
// PyTorch with `reduction='mean'` plus `ignore_index`).
//
// Targets are I64 (PyTorch / baracuda convention). The kernel
// validates each index is in `0..vocab` unless it equals the
// `ignore_index` sentinel.
//
// Memory: a per-row accumulator is the only allocation beyond the
// inputs/output. The `[N, V]` softmax + log-softmax that the
// primitive composition path materializes is never built.

/// Reduction-mode discriminant for [`fused_softmax_cross_entropy_f32`].
/// Mirrors `fuel_graph::registry::Reduction`; the wrapper in
/// fuel-storage translates between them. Kept primitive so this leaf
/// crate doesn't pull in fuel-graph just for one enum.
///
/// Encoding:
/// - `0` → Mean
/// - `1` → Sum
/// - `2` → None
pub const REDUCTION_MEAN: u8 = 0;
pub const REDUCTION_SUM: u8 = 1;
pub const REDUCTION_NONE: u8 = 2;

/// Shared shape-validation for the FusedSoftmaxCrossEntropy dtype
/// variants. Output dtype is always F32 regardless of logits dtype
/// (the declared FSCE output dtype — losses accumulate in F64 and
/// narrow to F32 at the end). Output buffer size depends on the
/// reduction tag:
///   - `REDUCTION_MEAN` / `REDUCTION_SUM`: 4 bytes (one f32 scalar).
///   - `REDUCTION_NONE`: `n_rows * 4` bytes.
fn fused_softmax_cross_entropy_check_shapes(
    name: &str,
    logits_elem_bytes: usize,
    logits_bytes: usize,
    targets_bytes: usize,
    out_bytes: usize,
    n_rows: usize,
    vocab: usize,
    reduction_tag: u8,
) -> Result<()> {
    let logits_need = n_rows
        .saturating_mul(vocab)
        .saturating_mul(logits_elem_bytes);
    if logits_bytes != logits_need {
        return Err(Error::Msg(format!(
            "{name}: logits bytes={logits_bytes} doesn't match \
             n_rows={n_rows} × vocab={vocab} × {logits_elem_bytes} = {logits_need}",
        ))
        .bt());
    }
    let targets_need = n_rows.saturating_mul(std::mem::size_of::<i64>());
    if targets_bytes != targets_need {
        return Err(Error::Msg(format!(
            "{name}: targets bytes={targets_bytes} doesn't match \
             n_rows={n_rows} × {} = {targets_need}",
            std::mem::size_of::<i64>(),
        ))
        .bt());
    }
    let out_need = match reduction_tag {
        REDUCTION_MEAN | REDUCTION_SUM => std::mem::size_of::<f32>(),
        REDUCTION_NONE => n_rows.saturating_mul(std::mem::size_of::<f32>()),
        other => {
            return Err(Error::Msg(format!("{name}: unknown reduction tag {other}",)).bt());
        }
    };
    if out_bytes != out_need {
        return Err(Error::Msg(format!(
            "{name}: out bytes={out_bytes} doesn't match \
             expected {out_need} for reduction_tag={reduction_tag}",
        ))
        .bt());
    }
    Ok(())
}

/// FusedSoftmaxCrossEntropy CPU kernel template. `T` is the logits
/// element dtype; output is always F32 (the FSCE declared dtype,
/// matching PyTorch and the baracuda kernel). All math runs through
/// f64 via `$to_f64` for log-sum-exp stability — the same precision
/// pattern the F32 v1 used internally (which already accumulated in
/// f64). F64 is lossless; F32/BF16/F16 promote on read.
///
/// Rows whose target equals `ignore_index` contribute 0 to the loss
/// sum and 0 to the Mean denominator. For Mean with zero valid rows
/// the output is 0.0 (no division by zero — matches PyTorch's
/// `nll_loss` behavior under `reduction='mean'` with all rows masked).
macro_rules! fused_softmax_cross_entropy_kernel {
    ($name:ident, $T:ty, $to_f64:expr) => {
        pub fn $name(
            logits: &CpuStorageBytes,
            targets: &CpuStorageBytes,
            out: &mut CpuStorageBytes,
            n_rows: usize,
            vocab: usize,
            reduction_tag: u8,
            ignore_index: i64,
        ) -> Result<()> {
            fused_softmax_cross_entropy_check_shapes(
                stringify!($name),
                std::mem::size_of::<$T>(),
                logits.len_bytes(),
                targets.len_bytes(),
                out.len_bytes(),
                n_rows,
                vocab,
                reduction_tag,
            )?;
            let out_view: &mut [f32] = out.as_slice_mut()?;
            if vocab == 0 {
                // Pathological: no classes. Define loss as 0.0 — there's
                // nothing to learn and producing NaN here would mask real
                // bugs upstream.
                out_view.fill(0.0);
                return Ok(());
            }
            let logits_view: &[$T] = logits.as_slice()?;
            let targets_view: &[i64] = targets.as_slice()?;

            if reduction_tag == REDUCTION_NONE {
                out_view.fill(0.0);
            }
            let mut loss_sum: f64 = 0.0;
            let mut n_valid: usize = 0;
            for row_idx in 0..n_rows {
                let target = targets_view[row_idx];
                if target == ignore_index {
                    // Per-row output stays 0.0 (None) / no contribution to sum.
                    continue;
                }
                if target < 0 || (target as i128) >= vocab as i128 {
                    return Err(Error::Msg(format!(
                        "{}: target {target} at row {row_idx} \
                         is out of range [0, {vocab}) (ignore_index={ignore_index})",
                        stringify!($name),
                    ))
                    .bt());
                }
                let target_idx = target as usize;
                let row_off = row_idx * vocab;
                let row_slice = &logits_view[row_off..row_off + vocab];
                // Row max in T (comparison is order-preserving). Promote
                // to f64 before the shift so half precisions get the
                // full f64 dynamic range under exp().
                let mut row_max: $T = row_slice[0];
                for &v in &row_slice[1..] {
                    if v > row_max {
                        row_max = v;
                    }
                }
                let row_max_f64: f64 = $to_f64(row_max);
                let mut sum_exp: f64 = 0.0;
                for &v in row_slice.iter() {
                    sum_exp += ($to_f64(v) - row_max_f64).exp();
                }
                let log_sum_exp = sum_exp.ln();
                let nll = -($to_f64(row_slice[target_idx]) - row_max_f64 - log_sum_exp);
                match reduction_tag {
                    REDUCTION_NONE => {
                        out_view[row_idx] = nll as f32;
                    }
                    _ => {
                        loss_sum += nll;
                        n_valid += 1;
                    }
                }
            }
            match reduction_tag {
                REDUCTION_MEAN => {
                    out_view[0] = if n_valid == 0 {
                        0.0
                    } else {
                        (loss_sum / n_valid as f64) as f32
                    };
                }
                REDUCTION_SUM => {
                    out_view[0] = loss_sum as f32;
                }
                REDUCTION_NONE => {}
                _ => unreachable!("reduction_tag validated above"),
            }
            Ok(())
        }
    };
}

fused_softmax_cross_entropy_kernel!(fused_softmax_cross_entropy_f32, f32, |v: f32| v as f64);
fused_softmax_cross_entropy_kernel!(fused_softmax_cross_entropy_f64, f64, |v: f64| v);
fused_softmax_cross_entropy_kernel!(
    fused_softmax_cross_entropy_bf16,
    half::bf16,
    |v: half::bf16| v.to_f64()

);
fused_softmax_cross_entropy_kernel!(
    fused_softmax_cross_entropy_f16,
    half::f16,
    |v: half::f16| v.to_f64()

);

// =============================================================================
// CausalConv1d — depthwise 1-D conv + bias + optional fused SiLU
// =============================================================================
//
// The Mamba-1 / Mamba-2 prefill convolution. The textbook depthwise
// recurrence computed in scalar Rust:
//
//   for b in 0..batch, c in 0..channels, t in 0..seq_out:
//     y = bias[c]
//     for k in 0..kernel:
//       y += weight[c, 0, k] * x[b, c, t + k]
//     out[b, c, t] = if use_silu { y * sigmoid(y) } else { y }
//
// Caller is responsible for left-padding x with `kernel - 1` zeros
// (matches Mamba's existing prefill convention — see mamba2.rs:552).
// The kernel does NOT pad internally; it just walks `seq_out` output
// positions and reads `kernel` consecutive elements per position.
//
// Memory: zero allocations beyond inputs/output. Per-output-element
// work is `O(kernel)` FMAs — the inner loop is bandwidth-bound for
// typical kernel sizes (4 for Mamba).

/// Shared shape-validation for the four CausalConv1d dtype variants.
fn causal_conv1d_check_shapes(
    name: &str,
    elem_bytes: usize,
    x_bytes: usize,
    weight_bytes: usize,
    bias_bytes: usize,
    out_bytes: usize,
    batch: usize,
    channels: usize,
    seq_in: usize,
    seq_out: usize,
    kernel: usize,
) -> Result<()> {
    if seq_in != seq_out + kernel - 1 {
        return Err(Error::Msg(format!(
            "{name}: seq_in={seq_in} must equal seq_out={seq_out} + kernel-1={}",
            kernel - 1,
        ))
        .bt());
    }
    let x_need = batch
        .saturating_mul(channels)
        .saturating_mul(seq_in)
        .saturating_mul(elem_bytes);
    let w_need = channels.saturating_mul(kernel).saturating_mul(elem_bytes);
    let b_need = channels.saturating_mul(elem_bytes);
    let out_need = batch
        .saturating_mul(channels)
        .saturating_mul(seq_out)
        .saturating_mul(elem_bytes);
    if x_bytes != x_need {
        return Err(Error::Msg(format!(
            "{name}: x bytes={x_bytes} doesn't match \
             batch={batch} × channels={channels} × seq_in={seq_in} × {elem_bytes} = {x_need}",
        ))
        .bt());
    }
    if weight_bytes != w_need {
        return Err(Error::Msg(format!(
            "{name}: weight bytes={weight_bytes} doesn't match \
             channels={channels} × 1 × kernel={kernel} × {elem_bytes} = {w_need}",
        ))
        .bt());
    }
    if bias_bytes != b_need {
        return Err(Error::Msg(format!(
            "{name}: bias bytes={bias_bytes} doesn't match \
             channels={channels} × {elem_bytes} = {b_need}",
        ))
        .bt());
    }
    if out_bytes != out_need {
        return Err(Error::Msg(format!(
            "{name}: out bytes={out_bytes} doesn't match \
             batch={batch} × channels={channels} × seq_out={seq_out} × {elem_bytes} = {out_need}",
        ))
        .bt());
    }
    Ok(())
}

/// CausalConv1d kernel for native float arithmetic (F32 / F64).
/// Accumulator type matches the element type — there's no precision
/// gain from F64 over F64. SiLU computed in the same type.
macro_rules! causal_conv1d_native_kernel {
    ($name:ident, $T:ty) => {
        /// CausalConv1d CPU kernel. Output buffer size is
        /// `batch * channels * seq_out * sizeof::<T>()` bytes.
        pub fn $name(
            x: &CpuStorageBytes,
            weight: &CpuStorageBytes,
            bias: &CpuStorageBytes,
            out: &mut CpuStorageBytes,
            batch: usize,
            channels: usize,
            seq_in: usize,
            seq_out: usize,
            kernel: usize,
            use_silu: bool,
        ) -> Result<()> {
            causal_conv1d_check_shapes(
                stringify!($name),
                std::mem::size_of::<$T>(),
                x.len_bytes(),
                weight.len_bytes(),
                bias.len_bytes(),
                out.len_bytes(),
                batch,
                channels,
                seq_in,
                seq_out,
                kernel,
            )?;
            if seq_out == 0 || channels == 0 || batch == 0 {
                return Ok(());
            }
            let x_view: &[$T] = x.as_slice()?;
            let w_view: &[$T] = weight.as_slice()?;
            let b_view: &[$T] = bias.as_slice()?;
            let out_view: &mut [$T] = out.as_slice_mut()?;
            for b in 0..batch {
                for c in 0..channels {
                    let x_row_off = (b * channels + c) * seq_in;
                    let w_row_off = c * kernel;
                    let out_row_off = (b * channels + c) * seq_out;
                    let bias_c = b_view[c];
                    for t in 0..seq_out {
                        let mut acc = bias_c;
                        for k in 0..kernel {
                            acc += w_view[w_row_off + k] * x_view[x_row_off + t + k];
                        }
                        out_view[out_row_off + t] = if use_silu {
                            // SiLU(x) = x * sigmoid(x) = x / (1 + exp(-x))
                            // `1.0` alone makes the literal default to f64, and
                            // `f32: From<f64>` does not hold — rustc falls back to
                            // inferring f32, which is CORRECT here but is a
                            // future-compat lint. `1.0_f32` satisfies both
                            // instantiations exactly (`f32: From<f32>` and
                            // `f64: From<f32>`), and 1.0 is exactly representable
                            // in both, so this changes the lint and not the math.
                            acc / (<$T>::from(1.0_f32) + (-acc).exp())
                        } else {
                            acc
                        };
                    }
                }
            }
            Ok(())
        }
    };
}

causal_conv1d_native_kernel!(causal_conv1d_f32, f32);
causal_conv1d_native_kernel!(causal_conv1d_f64, f64);

/// CausalConv1d kernel for half-precision floats (F16 / BF16) —
/// accumulate in F32, narrow on store. Matches the precision
/// pattern used by fused_linear_half_kernel and the rest of the
/// Mamba trio.
macro_rules! causal_conv1d_half_kernel {
    ($name:ident, $T:ty) => {
        pub fn $name(
            x: &CpuStorageBytes,
            weight: &CpuStorageBytes,
            bias: &CpuStorageBytes,
            out: &mut CpuStorageBytes,
            batch: usize,
            channels: usize,
            seq_in: usize,
            seq_out: usize,
            kernel: usize,
            use_silu: bool,
        ) -> Result<()> {
            causal_conv1d_check_shapes(
                stringify!($name),
                std::mem::size_of::<$T>(),
                x.len_bytes(),
                weight.len_bytes(),
                bias.len_bytes(),
                out.len_bytes(),
                batch,
                channels,
                seq_in,
                seq_out,
                kernel,
            )?;
            if seq_out == 0 || channels == 0 || batch == 0 {
                return Ok(());
            }
            let x_view: &[$T] = x.as_slice()?;
            let w_view: &[$T] = weight.as_slice()?;
            let b_view: &[$T] = bias.as_slice()?;
            let out_view: &mut [$T] = out.as_slice_mut()?;
            for b in 0..batch {
                for c in 0..channels {
                    let x_row_off = (b * channels + c) * seq_in;
                    let w_row_off = c * kernel;
                    let out_row_off = (b * channels + c) * seq_out;
                    let bias_c: f32 = b_view[c].to_f32();
                    for t in 0..seq_out {
                        let mut acc: f32 = bias_c;
                        for k in 0..kernel {
                            acc +=
                                w_view[w_row_off + k].to_f32() * x_view[x_row_off + t + k].to_f32();
                        }
                        let stored = if use_silu {
                            acc / (1.0 + (-acc).exp())
                        } else {
                            acc
                        };
                        out_view[out_row_off + t] = <$T>::from_f32(stored);
                    }
                }
            }
            Ok(())
        }
    };
}

causal_conv1d_half_kernel!(causal_conv1d_bf16, half::bf16);
causal_conv1d_half_kernel!(causal_conv1d_f16, half::f16);

// =============================================================================
// SelectiveScan — Mamba-1's selective state-space scan (forward)
// =============================================================================
//
// Per-timestep recurrence (per (batch b, dim i)):
//
//   d = softplus(delta[b,t,i]) if delta_softplus else delta[b,t,i]
//   for j in 0..dstate:
//     h[b,i,j] = exp(d * a[i,j]) * h[b,i,j] + d * b[b,t,j] * u[b,t,i]
//   y[b,t,i] = sum_j(h[b,i,j] * c[b,t,j])
//
// h is the recurrent hidden state, [batch, dim, dstate], allocated
// internally and zero-initialized at the start of the scan.
//
// Algorithmic complexity: O(batch · seqlen · dim · dstate) FMAs.

/// Shared shape-validation for the SelectiveScan dtype variants.
/// Multi-output (Option C, 2026-06-01): the output buffer is always
/// the bundled `[y_bytes ; last_state_bytes]` concatenation.
fn selective_scan_check_shapes(
    name: &str,
    elem_bytes: usize,
    u_bytes: usize,
    delta_bytes: usize,
    a_bytes: usize,
    b_bytes: usize,
    c_bytes: usize,
    out_bytes: usize,
    batch: usize,
    seqlen: usize,
    dim: usize,
    dstate: usize,
) -> Result<()> {
    let u_need = batch
        .saturating_mul(seqlen)
        .saturating_mul(dim)
        .saturating_mul(elem_bytes);
    let delta_need = u_need;
    let a_need = dim.saturating_mul(dstate).saturating_mul(elem_bytes);
    let b_need = batch
        .saturating_mul(seqlen)
        .saturating_mul(dstate)
        .saturating_mul(elem_bytes);
    let c_need = b_need;
    let last_state_need = batch
        .saturating_mul(dim)
        .saturating_mul(dstate)
        .saturating_mul(elem_bytes);
    let out_need = u_need.saturating_add(last_state_need);
    let check = |name_arg: &str, got: usize, need: usize| -> Result<()> {
        if got != need {
            Err(Error::Msg(format!("{name}: {name_arg} bytes={got}, expected {need}",)).bt())
        } else {
            Ok(())
        }
    };
    check("u", u_bytes, u_need)?;
    check("delta", delta_bytes, delta_need)?;
    check("a", a_bytes, a_need)?;
    check("b", b_bytes, b_need)?;
    check("c", c_bytes, c_need)?;
    check("out", out_bytes, out_need)?;
    Ok(())
}

/// SelectiveScan kernel template. `T` is the element / accumulator
/// dtype boundary — the state recurrence accumulates in F64
/// internally regardless of T (matches bnb / Mamba implementations
/// that up-cast to fp32+ for stability), then narrows to T on the
/// y store. Conversion through f32 keeps F16/BF16 paths efficient
/// without losing the F64 accumulator precision.
macro_rules! selective_scan_kernel {
    ($name:ident, $T:ty, $to_f64:expr, $from_f64:expr) => {
        /// SelectiveScan CPU kernel (multi-output, Option C). The
        /// `out` buffer is the bundled `[y_bytes ; last_state_bytes]`
        /// concatenation: slot 0 = `y [batch, seqlen, dim]` followed
        /// by slot 1 = `last_state [batch, dim, dstate]`, both in
        /// the same `T` dtype.
        pub fn $name(
            u: &CpuStorageBytes,
            delta: &CpuStorageBytes,
            a: &CpuStorageBytes,
            b: &CpuStorageBytes,
            c: &CpuStorageBytes,
            out: &mut CpuStorageBytes,
            batch: usize,
            seqlen: usize,
            dim: usize,
            dstate: usize,
            delta_softplus: bool,
        ) -> Result<()> {
            selective_scan_check_shapes(
                stringify!($name),
                std::mem::size_of::<$T>(),
                u.len_bytes(),
                delta.len_bytes(),
                a.len_bytes(),
                b.len_bytes(),
                c.len_bytes(),
                out.len_bytes(),
                batch,
                seqlen,
                dim,
                dstate,
            )?;
            if batch == 0 || seqlen == 0 || dim == 0 || dstate == 0 {
                let out_view: &mut [$T] = out.as_slice_mut()?;
                out_view.fill(<$T>::default());
                return Ok(());
            }
            let u_view: &[$T] = u.as_slice()?;
            let delta_view: &[$T] = delta.as_slice()?;
            let a_view: &[$T] = a.as_slice()?;
            let b_view: &[$T] = b.as_slice()?;
            let c_view: &[$T] = c.as_slice()?;
            let out_view: &mut [$T] = out.as_slice_mut()?;

            // Bundled output: slot 0 = y [batch, seqlen, dim],
            // slot 1 = last_state [batch, dim, dstate], same dtype.
            let y_elems = batch * seqlen * dim;
            let last_state_elems = batch * dim * dstate;
            debug_assert_eq!(out_view.len(), y_elems + last_state_elems);
            let (y_out, last_state_out) = out_view.split_at_mut(y_elems);

            // Recurrent hidden state — F64 accumulator regardless of T.
            // Layout: [batch][dim][dstate].
            let mut h = vec![0.0_f64; batch * dim * dstate];

            for bi in 0..batch {
                for t in 0..seqlen {
                    let u_row_off = (bi * seqlen + t) * dim;
                    let delta_row_off = u_row_off;
                    let b_row_off = (bi * seqlen + t) * dstate;
                    let c_row_off = b_row_off;
                    let y_row_off = u_row_off;
                    for i in 0..dim {
                        let raw_d = $to_f64(delta_view[delta_row_off + i]);
                        let d = if delta_softplus {
                            // softplus(x) = ln(1 + exp(x)). Numerically stable
                            // form: max(x, 0) + ln(1 + exp(-|x|)).
                            let abs_x = raw_d.abs();
                            raw_d.max(0.0) + (1.0 + (-abs_x).exp()).ln()
                        } else {
                            raw_d
                        };
                        let u_val = $to_f64(u_view[u_row_off + i]);
                        let h_off = (bi * dim + i) * dstate;
                        let a_row_off = i * dstate;
                        let mut y_acc: f64 = 0.0;
                        for j in 0..dstate {
                            let a_val = $to_f64(a_view[a_row_off + j]);
                            let b_val = $to_f64(b_view[b_row_off + j]);
                            let h_new = (d * a_val).exp() * h[h_off + j] + d * b_val * u_val;
                            h[h_off + j] = h_new;
                            y_acc += h_new * $to_f64(c_view[c_row_off + j]);
                        }
                        y_out[y_row_off + i] = $from_f64(y_acc);
                    }
                }
            }
            // Narrow the final hidden state into slot 1.
            debug_assert_eq!(last_state_out.len(), h.len());
            for (slot, &h_val) in last_state_out.iter_mut().zip(h.iter()) {
                *slot = $from_f64(h_val);
            }
            Ok(())
        }
    };
}

selective_scan_kernel!(selective_scan_f32, f32, |v: f32| v as f64, |v: f64| v

    as f32);
selective_scan_kernel!(selective_scan_f64, f64, |v: f64| v, |v: f64| v);
selective_scan_kernel!(
    selective_scan_bf16,
    half::bf16,
    |v: half::bf16| v.to_f64(),
    |v: f64| half::bf16::from_f32(v as f32)

);
selective_scan_kernel!(
    selective_scan_f16,
    half::f16,
    |v: half::f16| v.to_f64(),
    |v: f64| half::f16::from_f32(v as f32)

);

// =============================================================================
// NonZeroIndices — data-determined nonzero-index extraction (the keystone
// primitive for data-dependent dynamic shapes)
// =============================================================================

/// NonZeroIndices CPU kernel (multi-output, Option C). One input `x`
/// (read flat, contiguous — `capacity` elements); one bundled output
/// `[indices [capacity] U32 ; count [1] U32]` (both slots U32, so the
/// whole bundle reads as `capacity + 1` `u32`s and `split_at_mut` divides
/// it). Writes the ascending flat indices of `x`'s nonzero elements into
/// slot 0 (first `count` entries) and the runtime `count` into slot 1;
/// the tail `indices[count..]` is left as the bundle allocator
/// initialized it (zeros). `capacity == x.elem_count()` is the static
/// worst-case bound (every element nonzero).
macro_rules! nonzero_indices_kernel {
    ($name:ident, $T:ty, $is_nonzero:expr) => {
        pub fn $name(
            x: &CpuStorageBytes,
            out: &mut CpuStorageBytes,
            capacity: usize,
        ) -> Result<()> {
            let x_view: &[$T] = x.as_slice()?;
            if x_view.len() != capacity {
                return Err(Error::Msg(format!(
                    concat!(stringify!($name), ": input len {} != capacity {}"),
                    x_view.len(),
                    capacity,
                ))
                .bt());
            }
            let out_view: &mut [u32] = out.as_slice_mut()?;
            if out_view.len() != capacity + 1 {
                return Err(Error::Msg(format!(
                    concat!(
                        stringify!($name),
                        ": bundled output len {} != capacity+1 ({})",
                    ),
                    out_view.len(),
                    capacity + 1,
                ))
                .bt());
            }
            let (indices_out, count_out) = out_view.split_at_mut(capacity);
            let mut count: u32 = 0;
            for (i, v) in x_view.iter().enumerate() {
                if ($is_nonzero)(v) {
                    indices_out[count as usize] = i as u32;
                    count += 1;
                }
            }
            count_out[0] = count;
            Ok(())
        }
    };
}

nonzero_indices_kernel!(nonzero_indices_f32, f32, |v: &f32| *v != 0.0);
nonzero_indices_kernel!(nonzero_indices_u32, u32, |v: &u32| *v != 0);

// =============================================================================
// SsdChunkScan — Mamba-2's State-Space Duality chunked scan (forward)
// =============================================================================
//
// The CPU kernel runs a straight sequential scan over all `seqlen`
// tokens (`chunk_size` is a validated-but-inert GPU parallelism knob
// on CPU — see the `ssd_chunk_scan_kernel!` doc comment below for the
// full rationale). The mathematical result is identical for any
// `chunk_size ∈ [1, seqlen]` that divides seqlen.
//
// Per-timestep recurrence (per (batch b, head h, head_dim i)):
//
//   d = dt[b,t,h] (scalar)
//   a_h = a[h] (scalar log A)
//   for j in 0..state_dim:
//     h_state[b,h,i,j] = exp(d * a_h) * h_state[b,h,i,j]
//                      + d * b[b,t,h,j] * x[b,t,h,i]
//   y[b,t,h,i] = sum_j(h_state[b,h,i,j] * c[b,t,h,j])
//
// h_state is the per-head, per-head-dim recurrent state of shape
// [batch, heads, head_dim, state_dim], allocated internally and
// zero-initialized at the start of the scan.

/// Shared shape-validation for the SsdChunkScan dtype variants.
/// Multi-output (Option C, 2026-06-01): the output buffer is the
/// bundled `[y_bytes ; last_state_bytes]` concatenation.
fn ssd_chunk_scan_check_shapes(
    name: &str,
    elem_bytes: usize,
    x_bytes: usize,
    dt_bytes: usize,
    a_bytes: usize,
    b_bytes: usize,
    c_bytes: usize,
    out_bytes: usize,
    batch: usize,
    seqlen: usize,
    heads: usize,
    head_dim: usize,
    state_dim: usize,
) -> Result<()> {
    let x_need = batch
        .saturating_mul(seqlen)
        .saturating_mul(heads)
        .saturating_mul(head_dim)
        .saturating_mul(elem_bytes);
    let dt_need = batch
        .saturating_mul(seqlen)
        .saturating_mul(heads)
        .saturating_mul(elem_bytes);
    let a_need = heads.saturating_mul(elem_bytes);
    let b_need = batch
        .saturating_mul(seqlen)
        .saturating_mul(heads)
        .saturating_mul(state_dim)
        .saturating_mul(elem_bytes);
    let c_need = b_need;
    let last_state_need = batch
        .saturating_mul(heads)
        .saturating_mul(head_dim)
        .saturating_mul(state_dim)
        .saturating_mul(elem_bytes);
    let out_need = x_need.saturating_add(last_state_need);
    let check = |name_arg: &str, got: usize, need: usize| -> Result<()> {
        if got != need {
            Err(Error::Msg(format!("{name}: {name_arg} bytes={got}, expected {need}",)).bt())
        } else {
            Ok(())
        }
    };
    check("x", x_bytes, x_need)?;
    check("dt", dt_bytes, dt_need)?;
    check("a", a_bytes, a_need)?;
    check("b", b_bytes, b_need)?;
    check("c", c_bytes, c_need)?;
    check("out", out_bytes, out_need)?;
    Ok(())
}

/// SsdChunkScan kernel template. Same shape as SelectiveScan: F64
/// hidden-state accumulator regardless of element type, narrow to T
/// on output store.
///
/// ## On `chunk_size` and CPU correctness
///
/// `chunk_size` is the SSD block size — a GPU parallelization knob
/// that controls how many tokens are processed in parallel per
/// block. The Mamba-2 chunked algorithm rearranges the sequential
/// scan into block matrix ops (intra-chunk diagonal + inter-chunk
/// decay propagation) that GPUs can execute in parallel, but the
/// **mathematical result is identical** to a straight sequential
/// scan over all `seqlen` tokens.
///
/// On CPU we're sequential anyway, so we run a single fused scan
/// across the full sequence with hidden state `h` threaded across
/// every timestep. `chunk_size` is accepted (validated as positive
/// and dividing seqlen) but the inner-loop structure doesn't depend
/// on it — it's there for ABI compat with the GPU path. Any
/// `chunk_size ∈ [1, seqlen]` that divides seqlen produces the same
/// answer.
macro_rules! ssd_chunk_scan_kernel {
    ($name:ident, $T:ty, $to_f64:expr, $from_f64:expr) => {
        /// SsdChunkScan CPU kernel (multi-output, Option C). The
        /// `out` buffer is the bundled `[y_bytes ; last_state_bytes]`
        /// concatenation: slot 0 = `y [batch, seqlen, heads,
        /// head_dim]` followed by slot 1 = `last_state [batch,
        /// heads, head_dim, state_dim]`, both in the same `T` dtype.
        pub fn $name(
            x: &CpuStorageBytes,
            dt: &CpuStorageBytes,
            a: &CpuStorageBytes,
            b: &CpuStorageBytes,
            c: &CpuStorageBytes,
            out: &mut CpuStorageBytes,
            batch: usize,
            seqlen: usize,
            heads: usize,
            head_dim: usize,
            state_dim: usize,
            chunk_size: usize,
        ) -> Result<()> {
            if chunk_size == 0 {
                return Err(Error::Msg(format!(
                    "{}: chunk_size must be positive, got 0",
                    stringify!($name),
                ))
                .bt());
            }
            if seqlen != 0 && seqlen % chunk_size != 0 {
                return Err(Error::Msg(format!(
                    "{}: seqlen={seqlen} must be a multiple of chunk_size={chunk_size}",
                    stringify!($name),
                ))
                .bt());
            }
            ssd_chunk_scan_check_shapes(
                stringify!($name),
                std::mem::size_of::<$T>(),
                x.len_bytes(),
                dt.len_bytes(),
                a.len_bytes(),
                b.len_bytes(),
                c.len_bytes(),
                out.len_bytes(),
                batch,
                seqlen,
                heads,
                head_dim,
                state_dim,
            )?;
            if batch == 0 || seqlen == 0 || heads == 0 || head_dim == 0 || state_dim == 0 {
                let out_view: &mut [$T] = out.as_slice_mut()?;
                out_view.fill(<$T>::default());
                return Ok(());
            }
            let x_view: &[$T] = x.as_slice()?;
            let dt_view: &[$T] = dt.as_slice()?;
            let a_view: &[$T] = a.as_slice()?;
            let b_view: &[$T] = b.as_slice()?;
            let c_view: &[$T] = c.as_slice()?;
            let out_view: &mut [$T] = out.as_slice_mut()?;

            // Bundled output split: slot 0 = y, slot 1 = last_state.
            let y_elems = batch * seqlen * heads * head_dim;
            let last_state_elems = batch * heads * head_dim * state_dim;
            debug_assert_eq!(out_view.len(), y_elems + last_state_elems);
            let (y_out, last_state_out) = out_view.split_at_mut(y_elems);

            let mut h = vec![0.0_f64; batch * heads * head_dim * state_dim];

            for bi in 0..batch {
                for t in 0..seqlen {
                    for hi in 0..heads {
                        let dt_val = $to_f64(dt_view[(bi * seqlen + t) * heads + hi]);
                        let a_val = $to_f64(a_view[hi]);
                        let exp_d_a = (dt_val * a_val).exp();
                        let x_off = ((bi * seqlen + t) * heads + hi) * head_dim;
                        let b_off = ((bi * seqlen + t) * heads + hi) * state_dim;
                        let c_off = b_off;
                        let y_off = x_off;
                        for i in 0..head_dim {
                            let x_val = $to_f64(x_view[x_off + i]);
                            let h_off = ((bi * heads + hi) * head_dim + i) * state_dim;
                            let mut y_acc: f64 = 0.0;
                            for j in 0..state_dim {
                                let b_val = $to_f64(b_view[b_off + j]);
                                let h_new = exp_d_a * h[h_off + j] + dt_val * b_val * x_val;
                                h[h_off + j] = h_new;
                                y_acc += h_new * $to_f64(c_view[c_off + j]);
                            }
                            y_out[y_off + i] = $from_f64(y_acc);
                        }
                    }
                }
            }
            // Narrow the final hidden state into slot 1.
            debug_assert_eq!(last_state_out.len(), h.len());
            for (slot, &h_val) in last_state_out.iter_mut().zip(h.iter()) {
                *slot = $from_f64(h_val);
            }
            Ok(())
        }
    };
}

ssd_chunk_scan_kernel!(ssd_chunk_scan_f32, f32, |v: f32| v as f64, |v: f64| v

    as f32);
ssd_chunk_scan_kernel!(ssd_chunk_scan_f64, f64, |v: f64| v, |v: f64| v);
ssd_chunk_scan_kernel!(
    ssd_chunk_scan_bf16,
    half::bf16,
    |v: half::bf16| v.to_f64(),
    |v: f64| half::bf16::from_f32(v as f32)

);
ssd_chunk_scan_kernel!(
    ssd_chunk_scan_f16,
    half::f16,
    |v: half::f16| v.to_f64(),
    |v: f64| half::f16::from_f32(v as f32)

);

// =============================================================================
// Nf4Matmul — bitsandbytes 4-bit NormalFloat quantized matmul
// =============================================================================
//
// Inputs:
//   activations: [batch, m, k]      (T = F32/F16/BF16)
//   w_packed:    [n, k/2]           (u8 — two 4-bit codes per byte)
//   absmax:      [n, k/block_size]  (f32 — per-output-row, per-block)
//
// Output: [batch, m, n] (T)
//
// Per output position (b, i, j):
//   acc = 0
//   for kk in 0..k:
//     byte = w_packed[j, kk/2]
//     code = (byte >> (4 * (kk % 2))) & 0x0f   // even-k: low nibble; odd: high
//     dequant = NF4_LUT[code] * absmax[j, kk/block_size]
//     acc += activations[b, i, kk] * dequant
//   out[b, i, j] = acc
//
// Accumulate in f32 regardless of T (matches bitsandbytes' fp16/bf16
// path which also up-casts to f32 for the inner product); narrow on
// store. The NF4 LUT is baked in here — same 16 values as
// bitsandbytes' standard NF4 quantization curve.

/// The 16 NormalFloat values used by NF4 quantization. Matches
/// bitsandbytes' implementation byte-for-byte. These are the
/// inverse-CDF quantiles of the standard normal N(0,1) that minimize
/// the expected quantization error for normally-distributed weights.
pub const NF4_LUT: [f32; 16] = [
    -1.0,
    -0.696_192_8,
    -0.525_073_05,
    -0.394_917_5,
    -0.284_441_38,
    -0.184_773_43,
    -0.091_050_036,
    0.0,
    0.079_580_3,
    0.160_930_2,
    0.246_112_3,
    0.337_915_24,
    0.440_709_83,
    0.562_617,
    0.722_956_84,
    1.0,
];

/// Shared shape-validation for the three Nf4Matmul dtype variants.
/// Returns the per-element byte size of `T` so callers can verify
/// activation/output sizes without re-templating the check.
fn nf4_matmul_check_shapes(
    name: &str,
    elem_bytes: usize,
    activations_bytes: usize,
    w_packed_bytes: usize,
    absmax_bytes: usize,
    out_bytes: usize,
    batch: usize,
    m: usize,
    n: usize,
    k: usize,
    block_size: usize,
) -> Result<()> {
    if k == 0 || !k.is_multiple_of(2) {
        return Err(Error::Msg(format!("{name}: k={k} must be even and non-zero",)).bt());
    }
    if block_size == 0 || !k.is_multiple_of(block_size) {
        return Err(Error::Msg(format!(
            "{name}: k={k} must be a multiple of block_size={block_size}",
        ))
        .bt());
    }
    let n_blocks = k / block_size;
    let a_need = batch
        .saturating_mul(m)
        .saturating_mul(k)
        .saturating_mul(elem_bytes);
    let w_need = n.saturating_mul(k / 2);
    let abs_need = n
        .saturating_mul(n_blocks)
        .saturating_mul(std::mem::size_of::<f32>());
    let out_need = batch
        .saturating_mul(m)
        .saturating_mul(n)
        .saturating_mul(elem_bytes);
    if activations_bytes != a_need {
        return Err(Error::Msg(format!(
            "{name}: activations bytes={activations_bytes}, expected {a_need}",
        ))
        .bt());
    }
    if w_packed_bytes != w_need {
        return Err(Error::Msg(format!(
            "{name}: w_packed bytes={w_packed_bytes}, expected {w_need} (n={n} × k/2={})",
            k / 2,
        ))
        .bt());
    }
    if absmax_bytes != abs_need {
        return Err(Error::Msg(format!(
            "{name}: absmax bytes={absmax_bytes}, expected {abs_need} (n={n} × n_blocks={n_blocks} × 4)",
        )).bt());
    }
    if out_bytes != out_need {
        return Err(Error::Msg(format!(
            "{name}: out bytes={out_bytes}, expected {out_need}",
        ))
        .bt());
    }
    Ok(())
}

/// Macro for the per-dtype Nf4Matmul kernel body. `T` is the
/// activation/output type; conversion to/from f32 goes through the
/// `from_f32`/`to_f32` helpers expected on T (these match the
/// `half::f16` / `half::bf16` API; for the f32 case we use identity
/// closures).
macro_rules! nf4_matmul_native {
    ($name:ident, $T:ty, $to_f32:expr, $from_f32:expr) => {
        pub fn $name(
            activations: &CpuStorageBytes,
            w_packed: &CpuStorageBytes,
            absmax: &CpuStorageBytes,
            out: &mut CpuStorageBytes,
            batch: usize,
            m: usize,
            n: usize,
            k: usize,
            block_size: usize,
        ) -> Result<()> {
            nf4_matmul_check_shapes(
                stringify!($name),
                std::mem::size_of::<$T>(),
                activations.len_bytes(),
                w_packed.len_bytes(),
                absmax.len_bytes(),
                out.len_bytes(),
                batch,
                m,
                n,
                k,
                block_size,
            )?;
            if batch == 0 || m == 0 || n == 0 || k == 0 {
                let out_view: &mut [$T] = out.as_slice_mut()?;
                out_view.fill(<$T>::default());
                return Ok(());
            }
            let act_view: &[$T] = activations.as_slice()?;
            let w_view: &[u8] = w_packed.as_slice()?;
            let abs_view: &[f32] = absmax.as_slice()?;
            let out_view: &mut [$T] = out.as_slice_mut()?;
            let n_blocks = k / block_size;
            let k_packed = k / 2;
            for bi in 0..batch {
                for i in 0..m {
                    let a_row_off = (bi * m + i) * k;
                    let out_row_off = (bi * m + i) * n;
                    for j in 0..n {
                        let w_row_off = j * k_packed;
                        let abs_row_off = j * n_blocks;
                        let mut acc: f32 = 0.0;
                        for kk in 0..k {
                            let byte = w_view[w_row_off + (kk / 2)];
                            let nibble = if kk % 2 == 0 { byte & 0x0f } else { byte >> 4 };
                            let lut_val = NF4_LUT[nibble as usize];
                            let scale = abs_view[abs_row_off + (kk / block_size)];
                            let dequant = lut_val * scale;
                            let a_val: f32 = $to_f32(act_view[a_row_off + kk]);
                            acc += a_val * dequant;
                        }
                        out_view[out_row_off + j] = $from_f32(acc);
                    }
                }
            }
            Ok(())
        }
    };
}

nf4_matmul_native!(nf4_matmul_f32, f32, |v: f32| v, |v: f32| v);
nf4_matmul_native!(
    nf4_matmul_f16,
    half::f16,
    |v: half::f16| v.to_f32(),
    |v: f32| half::f16::from_f32(v)

);
nf4_matmul_native!(
    nf4_matmul_bf16,
    half::bf16,
    |v: half::bf16| v.to_f32(),
    |v: f32| half::bf16::from_f32(v)

);

// =============================================================================
// FlashAttn — naive multi-head SDPA (math definition)
// =============================================================================
//
// This is the math-definition oracle, not a tiled FlashAttention-2.
// On CPU the win from tiling is marginal compared to GPU, so we keep
// the simpler O(Sq*Sk*D) form per head and let backends ship a tiled
// kernel when one is worth the maintenance cost.
//
// Inputs:
//   q:               [B, Hq,  Sq, D]
//   k, v:            [B, Hkv, Sk, D]   (GQA: Hq must be a multiple of Hkv)
//   alibi_slopes:    [Hq] (optional)
// Output: same shape as q.

#[inline]
fn flash_attn_admissible(
    qi: usize,
    kj: usize,
    causal: bool,
    window_left: Option<usize>,
    window_right: Option<usize>,
) -> bool {
    if causal && kj > qi {
        return false;
    }
    if let Some(w) = window_left
        && kj + w < qi
    {
        return false;
    }
    if let Some(w) = window_right
        && kj > qi + w
    {
        return false;
    }
    true
}

/// FlashAttn kernel for native arithmetic ($T = f32 or f64).
macro_rules! flash_attn_native_kernel {
    ($name:ident, $T:ty, $T_zero:expr) => {
        pub fn $name(
            q: &CpuStorageBytes,
            k: &CpuStorageBytes,
            v: &CpuStorageBytes,
            alibi_slopes: Option<&CpuStorageBytes>,
            out: &mut CpuStorageBytes,
            b: usize,
            hq: usize,
            hkv: usize,
            sq: usize,
            sk: usize,
            d: usize,
            k_len: usize,
            softmax_scale: f32,
            causal: bool,
            window_left: Option<usize>,
            window_right: Option<usize>,
            softcap: Option<f32>,
        ) -> Result<()> {
            if hq % hkv != 0 || hkv == 0 {
                return Err(Error::Msg(format!(
                    "{}: Hq={hq} must be a positive multiple of Hkv={hkv}",
                    stringify!($name),
                ))
                .bt());
            }
            let elem = std::mem::size_of::<$T>();
            if q.len_bytes() != b * hq * sq * d * elem
                || k.len_bytes() != b * hkv * sk * d * elem
                || v.len_bytes() != b * hkv * sk * d * elem
                || out.len_bytes() != b * hq * sq * d * elem
            {
                return Err(Error::Msg(format!("{}: bytes mismatch", stringify!($name),)).bt());
            }
            if let Some(a) = alibi_slopes {
                if a.len_bytes() != hq * elem {
                    return Err(Error::Msg(format!(
                        "{}: alibi_slopes must be [{hq}] {} bytes, got {}",
                        stringify!($name),
                        hq * elem,
                        a.len_bytes(),
                    ))
                    .bt());
                }
            }
            let q_view: &[$T] = q.as_slice()?;
            let k_view: &[$T] = k.as_slice()?;
            let v_view: &[$T] = v.as_slice()?;
            let out_view: &mut [$T] = out.as_slice_mut()?;
            let alibi_view: Option<&[$T]> = match alibi_slopes {
                Some(a) => Some(a.as_slice()?),
                None => None,
            };
            let groups = hq / hkv;
            let q_h_stride = sq * d;
            let q_b_stride = hq * q_h_stride;
            let k_h_stride = sk * d;
            let k_b_stride = hkv * k_h_stride;
            let o_h_stride = sq * d;
            let o_b_stride = hq * o_h_stride;
            // `k_len` (logical attended length) must fit the physical K
            // extent `sk` (capacity): the strides above use `sk`, the
            // attention loops below use `k_len`. Bottom-right causal
            // alignment: query row `qi` is at absolute position
            // `qi + causal_offset` (= `qi + cached_len` in decode).
            if k_len > sk {
                return Err(Error::Msg(format!(
                    "{}: k_len ({k_len}) exceeds K extent sk ({sk})",
                    stringify!($name),
                ))
                .bt());
            }
            let causal_offset = k_len.saturating_sub(sq);
            let scale = softmax_scale as $T;
            // Zero output up front so masked rows stay zero.
            for slot in out_view.iter_mut() {
                *slot = $T_zero;
            }
            for bi in 0..b {
                for hi in 0..hq {
                    let kv_h = hi / groups;
                    let q_off = bi * q_b_stride + hi * q_h_stride;
                    let k_off = bi * k_b_stride + kv_h * k_h_stride;
                    let v_off = k_off;
                    let o_off = bi * o_b_stride + hi * o_h_stride;
                    let alibi_h: Option<$T> = alibi_view.map(|a| a[hi]);
                    for qi in 0..sq {
                        // Absolute query position (bottom-right causal).
                        let aq = qi + causal_offset;
                        // Build admissible scores over the LOGICAL prefix
                        // `0..k_len` of the (capacity) K/V buffers.
                        let mut scores = vec![$T_zero; k_len];
                        let mut admissible = vec![false; k_len];
                        let mut max_score: $T = <$T>::NEG_INFINITY;
                        for kj in 0..k_len {
                            if !flash_attn_admissible(aq, kj, causal, window_left, window_right) {
                                continue;
                            }
                            admissible[kj] = true;
                            let mut acc: $T = $T_zero;
                            let q_row = &q_view[q_off + qi * d..q_off + (qi + 1) * d];
                            let k_row = &k_view[k_off + kj * d..k_off + (kj + 1) * d];
                            for (qx, kx) in q_row.iter().zip(k_row.iter()) {
                                acc += (*qx) * (*kx);
                            }
                            let mut s = acc * scale;
                            if let Some(c) = softcap {
                                let cc = c as $T;
                                s = (s / cc).tanh() * cc;
                            }
                            if let Some(slope) = alibi_h {
                                let delta = (kj as f32 - aq as f32) as $T;
                                s += slope * delta;
                            }
                            scores[kj] = s;
                            if s > max_score {
                                max_score = s;
                            }
                        }
                        if !max_score.is_finite() {
                            continue;
                        }
                        let mut sum: $T = $T_zero;
                        for (s, ad) in scores.iter_mut().zip(admissible.iter()) {
                            if *ad {
                                *s = (*s - max_score).exp();
                                sum += *s;
                            } else {
                                *s = $T_zero;
                            }
                        }
                        if sum == $T_zero {
                            continue;
                        }
                        let inv_sum = (1.0 as $T) / sum;
                        for kj in 0..k_len {
                            if !admissible[kj] {
                                continue;
                            }
                            let p_ij = scores[kj] * inv_sum;
                            if p_ij == $T_zero {
                                continue;
                            }
                            let v_row = &v_view[v_off + kj * d..v_off + (kj + 1) * d];
                            for (od, vd) in out_view[o_off + qi * d..o_off + (qi + 1) * d]
                                .iter_mut()
                                .zip(v_row.iter())
                            {
                                *od += p_ij * (*vd);
                            }
                        }
                    }
                }
            }
            Ok(())
        }
    };
}

flash_attn_native_kernel!(flash_attn_f32, f32, 0.0_f32);
flash_attn_native_kernel!(flash_attn_f64, f64, 0.0_f64);

/// FlashAttn for half-floats — accumulates the dot products and
/// softmax math in f32, narrows back to T for the output.
macro_rules! flash_attn_half_kernel {
    ($name:ident, $T:ty) => {
        pub fn $name(
            q: &CpuStorageBytes,
            k: &CpuStorageBytes,
            v: &CpuStorageBytes,
            alibi_slopes: Option<&CpuStorageBytes>,
            out: &mut CpuStorageBytes,
            b: usize,
            hq: usize,
            hkv: usize,
            sq: usize,
            sk: usize,
            d: usize,
            k_len: usize,
            softmax_scale: f32,
            causal: bool,
            window_left: Option<usize>,
            window_right: Option<usize>,
            softcap: Option<f32>,
        ) -> Result<()> {
            if hq % hkv != 0 || hkv == 0 {
                return Err(Error::Msg(format!(
                    "{}: Hq={hq} must be a positive multiple of Hkv={hkv}",
                    stringify!($name),
                ))
                .bt());
            }
            let elem = std::mem::size_of::<$T>();
            if q.len_bytes() != b * hq * sq * d * elem
                || k.len_bytes() != b * hkv * sk * d * elem
                || v.len_bytes() != b * hkv * sk * d * elem
                || out.len_bytes() != b * hq * sq * d * elem
            {
                return Err(Error::Msg(format!("{}: bytes mismatch", stringify!($name),)).bt());
            }
            if let Some(a) = alibi_slopes {
                if a.len_bytes() != hq * elem {
                    return Err(Error::Msg(format!(
                        "{}: alibi_slopes bytes mismatch",
                        stringify!($name),
                    ))
                    .bt());
                }
            }
            let q_view: &[$T] = q.as_slice()?;
            let k_view: &[$T] = k.as_slice()?;
            let v_view: &[$T] = v.as_slice()?;
            let out_view: &mut [$T] = out.as_slice_mut()?;
            let alibi_view: Option<&[$T]> = match alibi_slopes {
                Some(a) => Some(a.as_slice()?),
                None => None,
            };
            let groups = hq / hkv;
            let q_h_stride = sq * d;
            let q_b_stride = hq * q_h_stride;
            let k_h_stride = sk * d;
            let k_b_stride = hkv * k_h_stride;
            let o_h_stride = sq * d;
            let o_b_stride = hq * o_h_stride;
            // `k_len` ≤ physical extent `sk`; strides use `sk`, the
            // attention loops use `k_len`. Bottom-right causal: query
            // row `qi` sits at absolute position `qi + causal_offset`.
            if k_len > sk {
                return Err(Error::Msg(format!(
                    "{}: k_len ({k_len}) exceeds K extent sk ({sk})",
                    stringify!($name),
                ))
                .bt());
            }
            let causal_offset = k_len.saturating_sub(sq);
            for slot in out_view.iter_mut() {
                *slot = <$T>::from_f32(0.0);
            }
            for bi in 0..b {
                for hi in 0..hq {
                    let kv_h = hi / groups;
                    let q_off = bi * q_b_stride + hi * q_h_stride;
                    let k_off = bi * k_b_stride + kv_h * k_h_stride;
                    let v_off = k_off;
                    let o_off = bi * o_b_stride + hi * o_h_stride;
                    let alibi_h: Option<f32> = alibi_view.map(|a| a[hi].to_f32());
                    for qi in 0..sq {
                        let aq = qi + causal_offset;
                        let mut scores = vec![0.0_f32; k_len];
                        let mut admissible = vec![false; k_len];
                        let mut max_score = f32::NEG_INFINITY;
                        for kj in 0..k_len {
                            if !flash_attn_admissible(aq, kj, causal, window_left, window_right) {
                                continue;
                            }
                            admissible[kj] = true;
                            let mut acc = 0.0_f32;
                            let q_row = &q_view[q_off + qi * d..q_off + (qi + 1) * d];
                            let k_row = &k_view[k_off + kj * d..k_off + (kj + 1) * d];
                            for (qx, kx) in q_row.iter().zip(k_row.iter()) {
                                acc += qx.to_f32() * kx.to_f32();
                            }
                            let mut s = acc * softmax_scale;
                            if let Some(c) = softcap {
                                s = (s / c).tanh() * c;
                            }
                            if let Some(slope) = alibi_h {
                                let delta = kj as f32 - aq as f32;
                                s += slope * delta;
                            }
                            scores[kj] = s;
                            if s > max_score {
                                max_score = s;
                            }
                        }
                        if !max_score.is_finite() {
                            continue;
                        }
                        let mut sum = 0.0_f32;
                        for (s, ad) in scores.iter_mut().zip(admissible.iter()) {
                            if *ad {
                                *s = (*s - max_score).exp();
                                sum += *s;
                            } else {
                                *s = 0.0;
                            }
                        }
                        if sum == 0.0 {
                            continue;
                        }
                        let inv_sum = 1.0_f32 / sum;
                        let mut row_acc = vec![0.0_f32; d];
                        for kj in 0..k_len {
                            if !admissible[kj] {
                                continue;
                            }
                            let p_ij = scores[kj] * inv_sum;
                            if p_ij == 0.0 {
                                continue;
                            }
                            let v_row = &v_view[v_off + kj * d..v_off + (kj + 1) * d];
                            for (od, vd) in row_acc.iter_mut().zip(v_row.iter()) {
                                *od += p_ij * vd.to_f32();
                            }
                        }
                        for (slot, val) in out_view[o_off + qi * d..o_off + (qi + 1) * d]
                            .iter_mut()
                            .zip(row_acc.iter())
                        {
                            *slot = <$T>::from_f32(*val);
                        }
                    }
                }
            }
            Ok(())
        }
    };
}

flash_attn_half_kernel!(flash_attn_bf16, half::bf16);
flash_attn_half_kernel!(flash_attn_f16, half::f16);

// =============================================================================
// FlashAttn backward — recompute approach (matches fuel_reference_backend's
// attention_flash_backward math, byte-storage flavor).
// =============================================================================
//
// One inner function `flash_attn_backward_*_inner` computes dQ, dK, dV
// jointly into pre-allocated Vec<T> buffers. Three byte wrappers
// (_q / _k / _v) each call the inner function and copy the
// requested gradient into their byte output. Three calls compute
// the same backward state three times — acceptable for math-
// definition CPU; GPU kernels can fuse the recompute.

/// Inner FlashAttn-backward kernel for native arithmetic types
/// (`$T = f32 | f64`). Writes into pre-allocated dq/dk/dv slices.
macro_rules! flash_attn_backward_native_inner {
    ($name:ident, $T:ty, $T_zero:expr) => {
        fn $name(
            q_view: &[$T],
            k_view: &[$T],
            v_view: &[$T],
            do_view: &[$T],
            alibi_view: Option<&[$T]>,
            dq: &mut [$T],
            dk: &mut [$T],
            dv: &mut [$T],
            b: usize,
            hq: usize,
            hkv: usize,
            sq: usize,
            sk: usize,
            d: usize,
            softmax_scale: f32,
            causal: bool,
            window_left: Option<usize>,
            window_right: Option<usize>,
            softcap: Option<f32>,
        ) {
            let groups = hq / hkv;
            let q_h_stride = sq * d;
            let q_b_stride = hq * q_h_stride;
            let k_h_stride = sk * d;
            let k_b_stride = hkv * k_h_stride;
            let scale = softmax_scale as $T;
            let softcap_pair = softcap.map(|c| {
                let c_t = c as $T;
                (c_t, (1.0 as $T) / c_t)
            });

            for slot in dq.iter_mut() {
                *slot = $T_zero;
            }
            for slot in dk.iter_mut() {
                *slot = $T_zero;
            }
            for slot in dv.iter_mut() {
                *slot = $T_zero;
            }

            for bi in 0..b {
                for hi in 0..hq {
                    let kv_h = hi / groups;
                    let q_off = bi * q_b_stride + hi * q_h_stride;
                    let k_off = bi * k_b_stride + kv_h * k_h_stride;
                    let v_off = k_off;
                    let do_off = q_off;
                    let dq_off = q_off;
                    let dk_off = k_off;
                    let dv_off = k_off;
                    let alibi_h: Option<$T> = alibi_view.map(|a| a[hi]);
                    for qi in 0..sq {
                        let mut p_row = vec![$T_zero; sk];
                        let mut s_row = vec![<$T>::NEG_INFINITY; sk];
                        let mut s_pre_softcap = vec![<$T>::NEG_INFINITY; sk];
                        let mut max_s = <$T>::NEG_INFINITY;
                        for kj in 0..sk {
                            if !crate::byte_kernels::flash_attn_admissible(
                                qi,
                                kj,
                                causal,
                                window_left,
                                window_right,
                            ) {
                                continue;
                            }
                            let q_row = &q_view[q_off + qi * d..q_off + (qi + 1) * d];
                            let k_row = &k_view[k_off + kj * d..k_off + (kj + 1) * d];
                            let mut acc: $T = $T_zero;
                            for (qx, kx) in q_row.iter().zip(k_row.iter()) {
                                acc += (*qx) * (*kx);
                            }
                            let mut s = acc * scale;
                            s_pre_softcap[kj] = s;
                            if let Some((c, inv_c)) = softcap_pair {
                                s = (s * inv_c).tanh() * c;
                            }
                            if let Some(slope) = alibi_h {
                                let delta = (kj as f32 - qi as f32) as $T;
                                s += slope * delta;
                            }
                            s_row[kj] = s;
                            if s > max_s {
                                max_s = s;
                            }
                        }
                        if !max_s.is_finite() {
                            continue;
                        }
                        let mut sum_exp: $T = $T_zero;
                        for kj in 0..sk {
                            if s_row[kj].is_finite() {
                                let e = (s_row[kj] - max_s).exp();
                                p_row[kj] = e;
                                sum_exp += e;
                            }
                        }
                        if sum_exp == $T_zero {
                            continue;
                        }
                        let inv_sum = (1.0 as $T) / sum_exp;
                        for kj in 0..sk {
                            p_row[kj] *= inv_sum;
                        }

                        let do_row = &do_view[do_off + qi * d..do_off + (qi + 1) * d];
                        for kj in 0..sk {
                            let p_ij = p_row[kj];
                            if p_ij == $T_zero {
                                continue;
                            }
                            let dst_off = dv_off + kj * d;
                            for (dvd, &dod) in
                                dv[dst_off..dst_off + d].iter_mut().zip(do_row.iter())
                            {
                                *dvd += p_ij * dod;
                            }
                        }

                        let mut dp = vec![$T_zero; sk];
                        for kj in 0..sk {
                            let v_row = &v_view[v_off + kj * d..v_off + (kj + 1) * d];
                            let mut acc: $T = $T_zero;
                            for (dod, &vd) in do_row.iter().zip(v_row.iter()) {
                                acc += (*dod) * vd;
                            }
                            dp[kj] = acc;
                        }

                        let mut row_dot: $T = $T_zero;
                        for kj in 0..sk {
                            row_dot += p_row[kj] * dp[kj];
                        }
                        let mut ds = vec![$T_zero; sk];
                        for kj in 0..sk {
                            ds[kj] = (dp[kj] - row_dot) * p_row[kj];
                            if let Some((_c, inv_c)) = softcap_pair {
                                let pre = s_pre_softcap[kj];
                                if pre.is_finite() {
                                    let t = (pre * inv_c).tanh();
                                    let dtanh = (1.0 as $T) - t * t;
                                    ds[kj] *= dtanh;
                                }
                            }
                        }

                        let dq_row_off = dq_off + qi * d;
                        for kj in 0..sk {
                            let ds_ij = ds[kj] * scale;
                            if ds_ij == $T_zero {
                                continue;
                            }
                            let k_row = &k_view[k_off + kj * d..k_off + (kj + 1) * d];
                            for (dqd, &kd) in
                                dq[dq_row_off..dq_row_off + d].iter_mut().zip(k_row.iter())
                            {
                                *dqd += ds_ij * kd;
                            }
                        }
                        let q_row = &q_view[q_off + qi * d..q_off + (qi + 1) * d];
                        for kj in 0..sk {
                            let ds_ij = ds[kj] * scale;
                            if ds_ij == $T_zero {
                                continue;
                            }
                            let dst_off = dk_off + kj * d;
                            for (dkd, &qd) in dk[dst_off..dst_off + d].iter_mut().zip(q_row.iter())
                            {
                                *dkd += ds_ij * qd;
                            }
                        }
                    }
                }
            }
        }
    };
}

flash_attn_backward_native_inner!(flash_attn_backward_f32_inner, f32, 0.0_f32);
flash_attn_backward_native_inner!(flash_attn_backward_f64_inner, f64, 0.0_f64);

/// Inner FlashAttn-backward kernel for half-precision types (math at f32,
/// narrow on store).
macro_rules! flash_attn_backward_half_inner {
    ($name:ident, $T:ty) => {
        fn $name(
            q_view: &[$T],
            k_view: &[$T],
            v_view: &[$T],
            do_view: &[$T],
            alibi_view: Option<&[$T]>,
            dq: &mut [$T],
            dk: &mut [$T],
            dv: &mut [$T],
            b: usize,
            hq: usize,
            hkv: usize,
            sq: usize,
            sk: usize,
            d: usize,
            softmax_scale: f32,
            causal: bool,
            window_left: Option<usize>,
            window_right: Option<usize>,
            softcap: Option<f32>,
        ) {
            let groups = hq / hkv;
            let q_h_stride = sq * d;
            let q_b_stride = hq * q_h_stride;
            let k_h_stride = sk * d;
            let k_b_stride = hkv * k_h_stride;

            // Accumulate gradients in f32 then narrow on store.
            let mut dq_f32 = vec![0.0_f32; dq.len()];
            let mut dk_f32 = vec![0.0_f32; dk.len()];
            let mut dv_f32 = vec![0.0_f32; dv.len()];

            for bi in 0..b {
                for hi in 0..hq {
                    let kv_h = hi / groups;
                    let q_off = bi * q_b_stride + hi * q_h_stride;
                    let k_off = bi * k_b_stride + kv_h * k_h_stride;
                    let v_off = k_off;
                    let do_off = q_off;
                    let dq_off = q_off;
                    let dk_off = k_off;
                    let dv_off = k_off;
                    let alibi_h: Option<f32> = alibi_view.map(|a| a[hi].to_f32());
                    for qi in 0..sq {
                        let mut p_row = vec![0.0_f32; sk];
                        let mut s_row = vec![f32::NEG_INFINITY; sk];
                        let mut s_pre_softcap = vec![f32::NEG_INFINITY; sk];
                        let mut max_s = f32::NEG_INFINITY;
                        for kj in 0..sk {
                            if !crate::byte_kernels::flash_attn_admissible(
                                qi,
                                kj,
                                causal,
                                window_left,
                                window_right,
                            ) {
                                continue;
                            }
                            let q_row = &q_view[q_off + qi * d..q_off + (qi + 1) * d];
                            let k_row = &k_view[k_off + kj * d..k_off + (kj + 1) * d];
                            let mut acc = 0.0_f32;
                            for (qx, kx) in q_row.iter().zip(k_row.iter()) {
                                acc += qx.to_f32() * kx.to_f32();
                            }
                            let mut s = acc * softmax_scale;
                            s_pre_softcap[kj] = s;
                            if let Some(c) = softcap {
                                s = (s / c).tanh() * c;
                            }
                            if let Some(slope) = alibi_h {
                                let delta = kj as f32 - qi as f32;
                                s += slope * delta;
                            }
                            s_row[kj] = s;
                            if s > max_s {
                                max_s = s;
                            }
                        }
                        if !max_s.is_finite() {
                            continue;
                        }
                        let mut sum_exp = 0.0_f32;
                        for kj in 0..sk {
                            if s_row[kj].is_finite() {
                                let e = (s_row[kj] - max_s).exp();
                                p_row[kj] = e;
                                sum_exp += e;
                            }
                        }
                        if sum_exp == 0.0 {
                            continue;
                        }
                        let inv_sum = 1.0_f32 / sum_exp;
                        for kj in 0..sk {
                            p_row[kj] *= inv_sum;
                        }

                        let do_row = &do_view[do_off + qi * d..do_off + (qi + 1) * d];
                        for kj in 0..sk {
                            let p_ij = p_row[kj];
                            if p_ij == 0.0 {
                                continue;
                            }
                            let dst_off = dv_off + kj * d;
                            for (i, &dod) in do_row.iter().enumerate() {
                                dv_f32[dst_off + i] += p_ij * dod.to_f32();
                            }
                        }

                        let mut dp = vec![0.0_f32; sk];
                        for kj in 0..sk {
                            let v_row = &v_view[v_off + kj * d..v_off + (kj + 1) * d];
                            let mut acc = 0.0_f32;
                            for (dod, &vd) in do_row.iter().zip(v_row.iter()) {
                                acc += dod.to_f32() * vd.to_f32();
                            }
                            dp[kj] = acc;
                        }

                        let mut row_dot = 0.0_f32;
                        for kj in 0..sk {
                            row_dot += p_row[kj] * dp[kj];
                        }
                        let mut ds = vec![0.0_f32; sk];
                        for kj in 0..sk {
                            ds[kj] = (dp[kj] - row_dot) * p_row[kj];
                            if let Some(c) = softcap {
                                let inv_c = 1.0 / c;
                                let pre = s_pre_softcap[kj];
                                if pre.is_finite() {
                                    let t = (pre * inv_c).tanh();
                                    let dtanh = 1.0 - t * t;
                                    ds[kj] *= dtanh;
                                }
                            }
                        }

                        let dq_row_off = dq_off + qi * d;
                        for kj in 0..sk {
                            let ds_ij = ds[kj] * softmax_scale;
                            if ds_ij == 0.0 {
                                continue;
                            }
                            let k_row = &k_view[k_off + kj * d..k_off + (kj + 1) * d];
                            for (i, &kd) in k_row.iter().enumerate() {
                                dq_f32[dq_row_off + i] += ds_ij * kd.to_f32();
                            }
                        }
                        let q_row = &q_view[q_off + qi * d..q_off + (qi + 1) * d];
                        for kj in 0..sk {
                            let ds_ij = ds[kj] * softmax_scale;
                            if ds_ij == 0.0 {
                                continue;
                            }
                            let dst_off = dk_off + kj * d;
                            for (i, &qd) in q_row.iter().enumerate() {
                                dk_f32[dst_off + i] += ds_ij * qd.to_f32();
                            }
                        }
                    }
                }
            }

            for (slot, &acc) in dq.iter_mut().zip(dq_f32.iter()) {
                *slot = <$T>::from_f32(acc);
            }
            for (slot, &acc) in dk.iter_mut().zip(dk_f32.iter()) {
                *slot = <$T>::from_f32(acc);
            }
            for (slot, &acc) in dv.iter_mut().zip(dv_f32.iter()) {
                *slot = <$T>::from_f32(acc);
            }
        }
    };
}

flash_attn_backward_half_inner!(flash_attn_backward_bf16_inner, half::bf16);
flash_attn_backward_half_inner!(flash_attn_backward_f16_inner, half::f16);

/// Selector for which gradient (`Q`, `K`, or `V`) a byte wrapper writes
/// into its single output buffer. The inner function always computes
/// all three; the wrapper picks one.
#[derive(Copy, Clone)]
pub enum FaBackwardWhich {
    Q,
    K,
    V,
}

/// FlashAttn backward byte wrapper for native arithmetic types
/// (f32 / f64). Computes all three gradients then writes only the
/// requested one into `out` — same 3× cost on the math-definition CPU
/// path that the OpKind split implies.
macro_rules! flash_attn_backward_native_wrapper {
    ($name:ident, $T:ty, $T_zero:expr, $inner:ident) => {
        pub fn $name(
            q: &CpuStorageBytes,
            k: &CpuStorageBytes,
            v: &CpuStorageBytes,
            do_grad: &CpuStorageBytes,
            alibi_slopes: Option<&CpuStorageBytes>,
            out: &mut CpuStorageBytes,
            which: FaBackwardWhich,
            b: usize,
            hq: usize,
            hkv: usize,
            sq: usize,
            sk: usize,
            d: usize,
            softmax_scale: f32,
            causal: bool,
            window_left: Option<usize>,
            window_right: Option<usize>,
            softcap: Option<f32>,
        ) -> Result<()> {
            if hq % hkv != 0 || hkv == 0 {
                return Err(Error::Msg(format!(
                    "{}: Hq={hq} must be a positive multiple of Hkv={hkv}",
                    stringify!($name),
                ))
                .bt());
            }
            let elem = std::mem::size_of::<$T>();
            let q_n = b * hq * sq * d;
            let kv_n = b * hkv * sk * d;
            if q.len_bytes() != q_n * elem
                || k.len_bytes() != kv_n * elem
                || v.len_bytes() != kv_n * elem
                || do_grad.len_bytes() != q_n * elem
            {
                return Err(Error::Msg(format!("{}: bytes mismatch", stringify!($name),)).bt());
            }
            let need_out = match which {
                FaBackwardWhich::Q => q_n * elem,
                FaBackwardWhich::K => kv_n * elem,
                FaBackwardWhich::V => kv_n * elem,
            };
            if out.len_bytes() != need_out {
                return Err(Error::Msg(format!(
                    "{}: out bytes {} != required {}",
                    stringify!($name),
                    out.len_bytes(),
                    need_out,
                ))
                .bt());
            }
            let q_view: &[$T] = q.as_slice()?;
            let k_view: &[$T] = k.as_slice()?;
            let v_view: &[$T] = v.as_slice()?;
            let do_view: &[$T] = do_grad.as_slice()?;
            let alibi_view: Option<&[$T]> = match alibi_slopes {
                Some(a) => Some(a.as_slice()?),
                None => None,
            };
            let mut dq = vec![$T_zero; q_n];
            let mut dk = vec![$T_zero; kv_n];
            let mut dv = vec![$T_zero; kv_n];
            $inner(
                q_view,
                k_view,
                v_view,
                do_view,
                alibi_view,
                &mut dq,
                &mut dk,
                &mut dv,
                b,
                hq,
                hkv,
                sq,
                sk,
                d,
                softmax_scale,
                causal,
                window_left,
                window_right,
                softcap,
            );
            let out_view: &mut [$T] = out.as_slice_mut()?;
            let src = match which {
                FaBackwardWhich::Q => &dq[..],
                FaBackwardWhich::K => &dk[..],
                FaBackwardWhich::V => &dv[..],
            };
            out_view.copy_from_slice(src);
            Ok(())
        }
    };
}

flash_attn_backward_native_wrapper!(
    flash_attn_backward_f32,
    f32,
    0.0_f32,
    flash_attn_backward_f32_inner
);
flash_attn_backward_native_wrapper!(
    flash_attn_backward_f64,
    f64,
    0.0_f64,
    flash_attn_backward_f64_inner
);

/// FlashAttn backward byte wrapper for half-precision types.
macro_rules! flash_attn_backward_half_wrapper {
    ($name:ident, $T:ty, $inner:ident) => {
        pub fn $name(
            q: &CpuStorageBytes,
            k: &CpuStorageBytes,
            v: &CpuStorageBytes,
            do_grad: &CpuStorageBytes,
            alibi_slopes: Option<&CpuStorageBytes>,
            out: &mut CpuStorageBytes,
            which: FaBackwardWhich,
            b: usize,
            hq: usize,
            hkv: usize,
            sq: usize,
            sk: usize,
            d: usize,
            softmax_scale: f32,
            causal: bool,
            window_left: Option<usize>,
            window_right: Option<usize>,
            softcap: Option<f32>,
        ) -> Result<()> {
            if hq % hkv != 0 || hkv == 0 {
                return Err(Error::Msg(format!(
                    "{}: Hq={hq} must be a positive multiple of Hkv={hkv}",
                    stringify!($name),
                ))
                .bt());
            }
            let elem = std::mem::size_of::<$T>();
            let q_n = b * hq * sq * d;
            let kv_n = b * hkv * sk * d;
            if q.len_bytes() != q_n * elem
                || k.len_bytes() != kv_n * elem
                || v.len_bytes() != kv_n * elem
                || do_grad.len_bytes() != q_n * elem
            {
                return Err(Error::Msg(format!("{}: bytes mismatch", stringify!($name),)).bt());
            }
            let need_out = match which {
                FaBackwardWhich::Q => q_n * elem,
                FaBackwardWhich::K => kv_n * elem,
                FaBackwardWhich::V => kv_n * elem,
            };
            if out.len_bytes() != need_out {
                return Err(Error::Msg(format!(
                    "{}: out bytes {} != required {}",
                    stringify!($name),
                    out.len_bytes(),
                    need_out,
                ))
                .bt());
            }
            let q_view: &[$T] = q.as_slice()?;
            let k_view: &[$T] = k.as_slice()?;
            let v_view: &[$T] = v.as_slice()?;
            let do_view: &[$T] = do_grad.as_slice()?;
            let alibi_view: Option<&[$T]> = match alibi_slopes {
                Some(a) => Some(a.as_slice()?),
                None => None,
            };
            let mut dq = vec![<$T>::from_f32(0.0); q_n];
            let mut dk = vec![<$T>::from_f32(0.0); kv_n];
            let mut dv = vec![<$T>::from_f32(0.0); kv_n];
            $inner(
                q_view,
                k_view,
                v_view,
                do_view,
                alibi_view,
                &mut dq,
                &mut dk,
                &mut dv,
                b,
                hq,
                hkv,
                sq,
                sk,
                d,
                softmax_scale,
                causal,
                window_left,
                window_right,
                softcap,
            );
            let out_view: &mut [$T] = out.as_slice_mut()?;
            let src = match which {
                FaBackwardWhich::Q => &dq[..],
                FaBackwardWhich::K => &dk[..],
                FaBackwardWhich::V => &dv[..],
            };
            out_view.copy_from_slice(src);
            Ok(())
        }
    };
}

flash_attn_backward_half_wrapper!(
    flash_attn_backward_bf16,
    half::bf16,
    flash_attn_backward_bf16_inner
);
flash_attn_backward_half_wrapper!(
    flash_attn_backward_f16,
    half::f16,
    flash_attn_backward_f16_inner
);

// =============================================================================
// PagedAttn — paged-KV-cache attention (naive)
// =============================================================================
//
// Inputs:
//   q:             [B, Hq, Sq, D]
//   k_cache:       [num_blocks, block_size, Hkv, D]
//   v_cache:       [num_blocks, block_size, Hkv, D]
//   block_table:   [B, max_num_blocks_per_seq] (U32)  logical → physical block
//   context_lens:  [B] (U32)                          true context length per seq
//   alibi_slopes:  [Hq] (optional)
// Output: [B, Hq, Sq, D] (same shape as q).
//
// Causal masking is implicit: query at position `q_pos = ctx_len - Sq + sq`
// (absolute position in the sequence) admits keys at `kj <= q_pos`.

/// PagedAttn for native arithmetic.
macro_rules! paged_attn_native_kernel {
    ($name:ident, $T:ty, $T_zero:expr) => {
        pub fn $name(
            q: &CpuStorageBytes,
            k_cache: &CpuStorageBytes,
            v_cache: &CpuStorageBytes,
            block_table: &CpuStorageBytes,
            context_lens: &CpuStorageBytes,
            alibi_slopes: Option<&CpuStorageBytes>,
            out: &mut CpuStorageBytes,
            b: usize,
            hq: usize,
            hkv: usize,
            sq: usize,
            d: usize,
            block_size: usize,
            max_blocks_per_seq: usize,
            num_blocks: usize,
            softmax_scale: f32,
            softcap: Option<f32>,
        ) -> Result<()> {
            if hq % hkv != 0 || hkv == 0 || block_size == 0 {
                return Err(Error::Msg(format!(
                    "{}: contract violation (Hq={hq}, Hkv={hkv}, block_size={block_size})",
                    stringify!($name),
                ))
                .bt());
            }
            let elem = std::mem::size_of::<$T>();
            let u32_elem = std::mem::size_of::<u32>();
            if q.len_bytes() != b * hq * sq * d * elem
                || k_cache.len_bytes() != num_blocks * block_size * hkv * d * elem
                || v_cache.len_bytes() != num_blocks * block_size * hkv * d * elem
                || block_table.len_bytes() != b * max_blocks_per_seq * u32_elem
                || context_lens.len_bytes() != b * u32_elem
                || out.len_bytes() != b * hq * sq * d * elem
            {
                return Err(Error::Msg(format!("{}: bytes mismatch", stringify!($name),)).bt());
            }
            if let Some(a) = alibi_slopes {
                if a.len_bytes() != hq * elem {
                    return Err(Error::Msg(format!(
                        "{}: alibi_slopes bytes mismatch",
                        stringify!($name),
                    ))
                    .bt());
                }
            }
            let q_view: &[$T] = q.as_slice()?;
            let k_view: &[$T] = k_cache.as_slice()?;
            let v_view: &[$T] = v_cache.as_slice()?;
            let bt_view: &[u32] = block_table.as_slice()?;
            let cl_view: &[u32] = context_lens.as_slice()?;
            let out_view: &mut [$T] = out.as_slice_mut()?;
            let alibi_view: Option<&[$T]> = match alibi_slopes {
                Some(a) => Some(a.as_slice()?),
                None => None,
            };
            let groups = hq / hkv;
            let q_h_stride = sq * d;
            let q_b_stride = hq * q_h_stride;
            let o_h_stride = sq * d;
            let o_b_stride = hq * o_h_stride;
            // KV cache strides for [num_blocks, block_size, Hkv, D]
            let kv_block_stride = block_size * hkv * d;
            let kv_slot_stride = hkv * d;
            let kv_head_stride = d;
            let scale = softmax_scale as $T;
            for slot in out_view.iter_mut() {
                *slot = $T_zero;
            }
            for bi in 0..b {
                let ctx_len = cl_view[bi] as usize;
                if ctx_len == 0 {
                    continue;
                }
                if ctx_len > max_blocks_per_seq * block_size {
                    return Err(Error::Msg(format!(
                        "{}: context_lens[{bi}]={ctx_len} > capacity {} block_size*max_blocks",
                        stringify!($name),
                        max_blocks_per_seq * block_size,
                    ))
                    .bt());
                }
                let bt_off = bi * max_blocks_per_seq;
                for hi in 0..hq {
                    let kv_h = hi / groups;
                    let q_off = bi * q_b_stride + hi * q_h_stride;
                    let o_off = bi * o_b_stride + hi * o_h_stride;
                    let alibi_h: Option<$T> = alibi_view.map(|a| a[hi]);
                    for qi in 0..sq {
                        let q_pos_abs = ctx_len + qi - sq;
                        // Build admissible scores over [0, ctx_len).
                        let mut scores = vec![$T_zero; ctx_len];
                        let mut admissible = vec![false; ctx_len];
                        let mut max_score: $T = <$T>::NEG_INFINITY;
                        for kj in 0..ctx_len {
                            // Implicit causal mask.
                            if kj > q_pos_abs {
                                continue;
                            }
                            let logical_block = kj / block_size;
                            let block_off = kj % block_size;
                            let physical_block = bt_view[bt_off + logical_block] as usize;
                            if physical_block >= num_blocks {
                                return Err(Error::Msg(format!(
                                    "{}: block_table[{bi}, {logical_block}]={physical_block} \
                                     out of range (num_blocks={num_blocks})",
                                    stringify!($name),
                                ))
                                .bt());
                            }
                            admissible[kj] = true;
                            let k_off = physical_block * kv_block_stride
                                + block_off * kv_slot_stride
                                + kv_h * kv_head_stride;
                            let v_off = k_off;
                            let mut acc: $T = $T_zero;
                            let q_row = &q_view[q_off + qi * d..q_off + (qi + 1) * d];
                            let k_row = &k_view[k_off..k_off + d];
                            for (qx, kx) in q_row.iter().zip(k_row.iter()) {
                                acc += (*qx) * (*kx);
                            }
                            let mut s = acc * scale;
                            if let Some(c) = softcap {
                                let cc = c as $T;
                                s = (s / cc).tanh() * cc;
                            }
                            if let Some(slope) = alibi_h {
                                let delta = (kj as f32 - q_pos_abs as f32) as $T;
                                s += slope * delta;
                            }
                            scores[kj] = s;
                            if s > max_score {
                                max_score = s;
                            }
                            // The v_off variable is just a name alias for k_off in
                            // this naive impl since v_cache shares the block layout.
                            let _ = v_off;
                        }
                        if !max_score.is_finite() {
                            continue;
                        }
                        let mut sum: $T = $T_zero;
                        for (s, ad) in scores.iter_mut().zip(admissible.iter()) {
                            if *ad {
                                *s = (*s - max_score).exp();
                                sum += *s;
                            } else {
                                *s = $T_zero;
                            }
                        }
                        if sum == $T_zero {
                            continue;
                        }
                        let inv_sum = (1.0 as $T) / sum;
                        for kj in 0..ctx_len {
                            if !admissible[kj] {
                                continue;
                            }
                            let p_ij = scores[kj] * inv_sum;
                            if p_ij == $T_zero {
                                continue;
                            }
                            let logical_block = kj / block_size;
                            let block_off = kj % block_size;
                            let physical_block = bt_view[bt_off + logical_block] as usize;
                            let v_off = physical_block * kv_block_stride
                                + block_off * kv_slot_stride
                                + kv_h * kv_head_stride;
                            let v_row = &v_view[v_off..v_off + d];
                            for (od, vd) in out_view[o_off + qi * d..o_off + (qi + 1) * d]
                                .iter_mut()
                                .zip(v_row.iter())
                            {
                                *od += p_ij * (*vd);
                            }
                        }
                    }
                }
            }
            Ok(())
        }
    };
}

paged_attn_native_kernel!(paged_attn_f32, f32, 0.0_f32);
paged_attn_native_kernel!(paged_attn_f64, f64, 0.0_f64);

/// PagedAttn for half-floats — f32 accumulator.
macro_rules! paged_attn_half_kernel {
    ($name:ident, $T:ty) => {
        pub fn $name(
            q: &CpuStorageBytes,
            k_cache: &CpuStorageBytes,
            v_cache: &CpuStorageBytes,
            block_table: &CpuStorageBytes,
            context_lens: &CpuStorageBytes,
            alibi_slopes: Option<&CpuStorageBytes>,
            out: &mut CpuStorageBytes,
            b: usize,
            hq: usize,
            hkv: usize,
            sq: usize,
            d: usize,
            block_size: usize,
            max_blocks_per_seq: usize,
            num_blocks: usize,
            softmax_scale: f32,
            softcap: Option<f32>,
        ) -> Result<()> {
            if hq % hkv != 0 || hkv == 0 || block_size == 0 {
                return Err(Error::Msg(format!("{}: contract violation", stringify!($name),)).bt());
            }
            let elem = std::mem::size_of::<$T>();
            let u32_elem = std::mem::size_of::<u32>();
            if q.len_bytes() != b * hq * sq * d * elem
                || k_cache.len_bytes() != num_blocks * block_size * hkv * d * elem
                || v_cache.len_bytes() != num_blocks * block_size * hkv * d * elem
                || block_table.len_bytes() != b * max_blocks_per_seq * u32_elem
                || context_lens.len_bytes() != b * u32_elem
                || out.len_bytes() != b * hq * sq * d * elem
            {
                return Err(Error::Msg(format!("{}: bytes mismatch", stringify!($name),)).bt());
            }
            if let Some(a) = alibi_slopes {
                if a.len_bytes() != hq * elem {
                    return Err(Error::Msg(format!(
                        "{}: alibi_slopes bytes mismatch",
                        stringify!($name),
                    ))
                    .bt());
                }
            }
            let q_view: &[$T] = q.as_slice()?;
            let k_view: &[$T] = k_cache.as_slice()?;
            let v_view: &[$T] = v_cache.as_slice()?;
            let bt_view: &[u32] = block_table.as_slice()?;
            let cl_view: &[u32] = context_lens.as_slice()?;
            let out_view: &mut [$T] = out.as_slice_mut()?;
            let alibi_view: Option<&[$T]> = match alibi_slopes {
                Some(a) => Some(a.as_slice()?),
                None => None,
            };
            let groups = hq / hkv;
            let q_h_stride = sq * d;
            let q_b_stride = hq * q_h_stride;
            let o_h_stride = sq * d;
            let o_b_stride = hq * o_h_stride;
            let kv_block_stride = block_size * hkv * d;
            let kv_slot_stride = hkv * d;
            let kv_head_stride = d;
            for slot in out_view.iter_mut() {
                *slot = <$T>::from_f32(0.0);
            }
            for bi in 0..b {
                let ctx_len = cl_view[bi] as usize;
                if ctx_len == 0 {
                    continue;
                }
                if ctx_len > max_blocks_per_seq * block_size {
                    return Err(Error::Msg(format!(
                        "{}: ctx_len out of capacity",
                        stringify!($name),
                    ))
                    .bt());
                }
                let bt_off = bi * max_blocks_per_seq;
                for hi in 0..hq {
                    let kv_h = hi / groups;
                    let q_off = bi * q_b_stride + hi * q_h_stride;
                    let o_off = bi * o_b_stride + hi * o_h_stride;
                    let alibi_h: Option<f32> = alibi_view.map(|a| a[hi].to_f32());
                    for qi in 0..sq {
                        let q_pos_abs = ctx_len + qi - sq;
                        let mut scores = vec![0.0_f32; ctx_len];
                        let mut admissible = vec![false; ctx_len];
                        let mut max_score = f32::NEG_INFINITY;
                        for kj in 0..ctx_len {
                            if kj > q_pos_abs {
                                continue;
                            }
                            let logical_block = kj / block_size;
                            let block_off = kj % block_size;
                            let physical_block = bt_view[bt_off + logical_block] as usize;
                            if physical_block >= num_blocks {
                                return Err(Error::Msg(format!(
                                    "{}: block_table out of range",
                                    stringify!($name),
                                ))
                                .bt());
                            }
                            admissible[kj] = true;
                            let k_off = physical_block * kv_block_stride
                                + block_off * kv_slot_stride
                                + kv_h * kv_head_stride;
                            let mut acc = 0.0_f32;
                            let q_row = &q_view[q_off + qi * d..q_off + (qi + 1) * d];
                            let k_row = &k_view[k_off..k_off + d];
                            for (qx, kx) in q_row.iter().zip(k_row.iter()) {
                                acc += qx.to_f32() * kx.to_f32();
                            }
                            let mut s = acc * softmax_scale;
                            if let Some(c) = softcap {
                                s = (s / c).tanh() * c;
                            }
                            if let Some(slope) = alibi_h {
                                let delta = kj as f32 - q_pos_abs as f32;
                                s += slope * delta;
                            }
                            scores[kj] = s;
                            if s > max_score {
                                max_score = s;
                            }
                        }
                        if !max_score.is_finite() {
                            continue;
                        }
                        let mut sum = 0.0_f32;
                        for (s, ad) in scores.iter_mut().zip(admissible.iter()) {
                            if *ad {
                                *s = (*s - max_score).exp();
                                sum += *s;
                            } else {
                                *s = 0.0;
                            }
                        }
                        if sum == 0.0 {
                            continue;
                        }
                        let inv_sum = 1.0_f32 / sum;
                        let mut row_acc = vec![0.0_f32; d];
                        for kj in 0..ctx_len {
                            if !admissible[kj] {
                                continue;
                            }
                            let p_ij = scores[kj] * inv_sum;
                            if p_ij == 0.0 {
                                continue;
                            }
                            let logical_block = kj / block_size;
                            let block_off = kj % block_size;
                            let physical_block = bt_view[bt_off + logical_block] as usize;
                            let v_off = physical_block * kv_block_stride
                                + block_off * kv_slot_stride
                                + kv_h * kv_head_stride;
                            let v_row = &v_view[v_off..v_off + d];
                            for (od, vd) in row_acc.iter_mut().zip(v_row.iter()) {
                                *od += p_ij * vd.to_f32();
                            }
                        }
                        for (slot, val) in out_view[o_off + qi * d..o_off + (qi + 1) * d]
                            .iter_mut()
                            .zip(row_acc.iter())
                        {
                            *slot = <$T>::from_f32(*val);
                        }
                    }
                }
            }
            Ok(())
        }
    };
}

paged_attn_half_kernel!(paged_attn_bf16, half::bf16);
paged_attn_half_kernel!(paged_attn_f16, half::f16);

// =============================================================================
// 2D Convolution (f32)
// =============================================================================

/// Direct (no-im2col) 2D convolution forward pass on `f32`.
///
/// Shapes:
///   x:      [N, Cin,                Hin, Win]
///   weight: [Cout, Cin/groups,      Kh,  Kw ]
///   bias:   optional `[Cout]`
///   out:    [N, Cout,               Hout, Wout]
///
/// Out-of-bounds reads (from padding) yield 0. Groups partition Cin
/// and Cout into `groups` even chunks; output channel `co` reads
/// from input channels `(co / Co_per_group) * Ci_per_group ..`.
///
/// Correctness-first; vendor backends (cuDNN, MKL-DNN) will
/// dramatically outperform this once they're wired.
pub fn conv2d_f32(
    x: &CpuStorageBytes,
    weight: &CpuStorageBytes,
    bias: Option<&CpuStorageBytes>,
    out: &mut CpuStorageBytes,
    x_shape: [usize; 4],
    w_shape: [usize; 4],
    out_shape: [usize; 4],
    stride: (usize, usize),
    padding: (usize, usize),
    dilation: (usize, usize),
    groups: usize,
) -> Result<()> {
    let [n, cin, h_in, w_in] = x_shape;
    let [cout, cin_per_group, kh, kw] = w_shape;
    let [n_out, cout_out, h_out, w_out] = out_shape;
    if n != n_out {
        return Err(Error::Msg(format!(
            "conv2d_f32: input batch {n} != output batch {n_out}"
        ))
        .bt());
    }
    if cout != cout_out {
        return Err(Error::Msg(format!(
            "conv2d_f32: weight Cout {cout} != output Cout {cout_out}"
        ))
        .bt());
    }
    if groups == 0 || cin % groups != 0 || cout % groups != 0 {
        return Err(Error::Msg(format!(
            "conv2d_f32: groups={groups} must divide Cin={cin} and Cout={cout}"
        ))
        .bt());
    }
    if cin / groups != cin_per_group {
        return Err(Error::Msg(format!(
            "conv2d_f32: weight expects Cin/group={cin_per_group}, but \
             Cin/groups = {cin}/{groups} = {}",
            cin / groups,
        ))
        .bt());
    }
    let cout_per_group = cout / groups;
    let elem = std::mem::size_of::<f32>();
    let need_x = n
        .saturating_mul(cin)
        .saturating_mul(h_in)
        .saturating_mul(w_in)
        .saturating_mul(elem);
    let need_w = cout
        .saturating_mul(cin_per_group)
        .saturating_mul(kh)
        .saturating_mul(kw)
        .saturating_mul(elem);
    let need_out = n
        .saturating_mul(cout)
        .saturating_mul(h_out)
        .saturating_mul(w_out)
        .saturating_mul(elem);
    if x.len_bytes() != need_x {
        return Err(Error::Msg(format!(
            "conv2d_f32: x bytes={} doesn't match shape {:?}",
            x.len_bytes(),
            x_shape,
        ))
        .bt());
    }
    if weight.len_bytes() != need_w {
        return Err(Error::Msg(format!(
            "conv2d_f32: weight bytes={} doesn't match shape {:?}",
            weight.len_bytes(),
            w_shape,
        ))
        .bt());
    }
    if out.len_bytes() != need_out {
        return Err(Error::Msg(format!(
            "conv2d_f32: out bytes={} doesn't match shape {:?}",
            out.len_bytes(),
            out_shape,
        ))
        .bt());
    }
    if let Some(b) = bias
        && b.len_bytes() != cout * elem
    {
        return Err(Error::Msg(format!(
            "conv2d_f32: bias bytes={} doesn't match Cout={cout} (f32)",
            b.len_bytes(),
        ))
        .bt());
    }
    let x_view: &[f32] = x.as_slice()?;
    let w_view: &[f32] = weight.as_slice()?;
    let bias_view: Option<&[f32]> = match bias {
        Some(b) => Some(b.as_slice()?),
        None => None,
    };
    let out_view: &mut [f32] = out.as_slice_mut()?;
    let (sh, sw) = stride;
    let (ph, pw) = padding;
    let (dh, dw) = dilation;

    // Output element [b, co, oh, ow] = bias[co] + sum over (ci_inner, kh_i, kw_i) of
    //   weight[co, ci_inner, kh_i, kw_i] * x[b, group_offset + ci_inner, in_h, in_w]
    // where in_h = oh * sh + kh_i * dh - ph, in_w = ow * sw + kw_i * dw - pw.
    for b in 0..n {
        for co in 0..cout {
            let group = co / cout_per_group;
            let ci_offset = group * cin_per_group;
            let bias_val = bias_view.map(|bv| bv[co]).unwrap_or(0.0);
            for oh in 0..h_out {
                for ow in 0..w_out {
                    let mut acc: f32 = bias_val;
                    for ci in 0..cin_per_group {
                        for kh_i in 0..kh {
                            let in_h = (oh * sh + kh_i * dh) as isize - ph as isize;
                            if in_h < 0 || in_h as usize >= h_in {
                                continue;
                            }
                            let in_h = in_h as usize;
                            for kw_i in 0..kw {
                                let in_w = (ow * sw + kw_i * dw) as isize - pw as isize;
                                if in_w < 0 || in_w as usize >= w_in {
                                    continue;
                                }
                                let in_w = in_w as usize;
                                let x_idx =
                                    ((b * cin + (ci_offset + ci)) * h_in + in_h) * w_in + in_w;
                                let w_idx = ((co * cin_per_group + ci) * kh + kh_i) * kw + kw_i;
                                acc += x_view[x_idx] * w_view[w_idx];
                            }
                        }
                    }
                    let out_idx = ((b * cout + co) * h_out + oh) * w_out + ow;
                    out_view[out_idx] = acc;
                }
            }
        }
    }
    Ok(())
}

// =============================================================================
// Reduction kernels — Sum / Max / Min / Mean across f32 / f64 / bf16 / f16
// =============================================================================
//
// All sixteen entry points (4 ops × 4 dtypes) are 1-line thunks over
// the trait-chassis in [`crate::chassis::reduction`]. The trait
// carries the loop / shape / index math once per kernel family; the
// per-(op, dtype) impl carries only what changes per dtype — most
// notably the f32 accumulator promotion for bf16 / f16 sums and
// means (encoded as the trait's `Acc` associated type so a future
// contributor adding a new low-precision dtype cannot forget it).

/// Sum-reduce `f32`. Output is the input with the listed `reduce_dims`
/// removed; each remaining slot holds the sum over the reduced axes.
pub fn sum_reduce_f32(
    input: &CpuStorageBytes,
    output: &mut CpuStorageBytes,
    input_shape: &[usize],
    reduce_dims: &[usize],
) -> Result<()> {
    reduce::<f32, Sum>("sum_reduce_f32", input, output, input_shape, reduce_dims)
}

/// Mean-reduce `f32` — sum-reduce divided by the product of
/// reduced-dim sizes.
pub fn mean_reduce_f32(
    input: &CpuStorageBytes,
    output: &mut CpuStorageBytes,
    input_shape: &[usize],
    reduce_dims: &[usize],
) -> Result<()> {
    reduce::<f32, Mean>("mean_reduce_f32", input, output, input_shape, reduce_dims)
}

/// Max-reduce `f32`. Output slots initialize to `-inf`.
pub fn max_reduce_f32(
    input: &CpuStorageBytes,
    output: &mut CpuStorageBytes,
    input_shape: &[usize],
    reduce_dims: &[usize],
) -> Result<()> {
    reduce::<f32, Max>("max_reduce_f32", input, output, input_shape, reduce_dims)
}

/// Min-reduce `f32`. Output slots initialize to `+inf`.
pub fn min_reduce_f32(
    input: &CpuStorageBytes,
    output: &mut CpuStorageBytes,
    input_shape: &[usize],
    reduce_dims: &[usize],
) -> Result<()> {
    reduce::<f32, Min>("min_reduce_f32", input, output, input_shape, reduce_dims)
}

/// Sum-reduce `f64`.
pub fn sum_reduce_f64(
    input: &CpuStorageBytes,
    output: &mut CpuStorageBytes,
    input_shape: &[usize],
    reduce_dims: &[usize],
) -> Result<()> {
    reduce::<f64, Sum>("sum_reduce_f64", input, output, input_shape, reduce_dims)
}

/// Mean-reduce `f64`.
pub fn mean_reduce_f64(
    input: &CpuStorageBytes,
    output: &mut CpuStorageBytes,
    input_shape: &[usize],
    reduce_dims: &[usize],
) -> Result<()> {
    reduce::<f64, Mean>("mean_reduce_f64", input, output, input_shape, reduce_dims)
}

/// Max-reduce `f64`.
pub fn max_reduce_f64(
    input: &CpuStorageBytes,
    output: &mut CpuStorageBytes,
    input_shape: &[usize],
    reduce_dims: &[usize],
) -> Result<()> {
    reduce::<f64, Max>("max_reduce_f64", input, output, input_shape, reduce_dims)
}

/// Min-reduce `f64`.
pub fn min_reduce_f64(
    input: &CpuStorageBytes,
    output: &mut CpuStorageBytes,
    input_shape: &[usize],
    reduce_dims: &[usize],
) -> Result<()> {
    reduce::<f64, Min>("min_reduce_f64", input, output, input_shape, reduce_dims)
}

/// Sum-reduce `bf16`. Accumulator runs in f32 (load-bearing
/// invariant: summing many bf16 values in bf16 loses precision
/// rapidly; an f32 accumulator preserves full f32 precision up to
/// ~16M elements).
pub fn sum_reduce_bf16(
    input: &CpuStorageBytes,
    output: &mut CpuStorageBytes,
    input_shape: &[usize],
    reduce_dims: &[usize],
) -> Result<()> {
    reduce::<half::bf16, Sum>("sum_reduce_bf16", input, output, input_shape, reduce_dims)
}

/// Mean-reduce `bf16` — f32 accumulator + divide by count.
pub fn mean_reduce_bf16(
    input: &CpuStorageBytes,
    output: &mut CpuStorageBytes,
    input_shape: &[usize],
    reduce_dims: &[usize],
) -> Result<()> {
    reduce::<half::bf16, Mean>("mean_reduce_bf16", input, output, input_shape, reduce_dims)
}

/// Max-reduce `bf16`. Extremum runs in f32 space for uniform NaN
/// handling, narrows back to bf16.
pub fn max_reduce_bf16(
    input: &CpuStorageBytes,
    output: &mut CpuStorageBytes,
    input_shape: &[usize],
    reduce_dims: &[usize],
) -> Result<()> {
    reduce::<half::bf16, Max>("max_reduce_bf16", input, output, input_shape, reduce_dims)
}

/// Min-reduce `bf16`.
pub fn min_reduce_bf16(
    input: &CpuStorageBytes,
    output: &mut CpuStorageBytes,
    input_shape: &[usize],
    reduce_dims: &[usize],
) -> Result<()> {
    reduce::<half::bf16, Min>("min_reduce_bf16", input, output, input_shape, reduce_dims)
}

/// Sum-reduce `f16`. Accumulator runs in f32.
pub fn sum_reduce_f16(
    input: &CpuStorageBytes,
    output: &mut CpuStorageBytes,
    input_shape: &[usize],
    reduce_dims: &[usize],
) -> Result<()> {
    reduce::<half::f16, Sum>("sum_reduce_f16", input, output, input_shape, reduce_dims)
}

/// Mean-reduce `f16`.
pub fn mean_reduce_f16(
    input: &CpuStorageBytes,
    output: &mut CpuStorageBytes,
    input_shape: &[usize],
    reduce_dims: &[usize],
) -> Result<()> {
    reduce::<half::f16, Mean>("mean_reduce_f16", input, output, input_shape, reduce_dims)
}

/// Max-reduce `f16`.
pub fn max_reduce_f16(
    input: &CpuStorageBytes,
    output: &mut CpuStorageBytes,
    input_shape: &[usize],
    reduce_dims: &[usize],
) -> Result<()> {
    reduce::<half::f16, Max>("max_reduce_f16", input, output, input_shape, reduce_dims)
}

/// Min-reduce `f16`.
pub fn min_reduce_f16(
    input: &CpuStorageBytes,
    output: &mut CpuStorageBytes,
    input_shape: &[usize],
    reduce_dims: &[usize],
) -> Result<()> {
    reduce::<half::f16, Min>("min_reduce_f16", input, output, input_shape, reduce_dims)
}

// =============================================================================
// Argmax / argmin along one dim (f32) — distinct from the value-
// reducing kernels above because the output dtype is U32 (the
// index of the extremum within each row), not the input dtype.
// =============================================================================

/// Argmax / argmin along a single dim — same shape contract as
/// the value-reducing kernels (output's `dim` is removed) but the
/// output dtype is `u32` carrying the index of the extremum within
/// each row of `dim`.
///
/// Tie-breaking: returns the first index that achieves the
/// extremum (lowest index on ties). NaN handling propagates per
/// IEEE-754 — comparisons against NaN return false, so a NaN
/// row's argmax is whichever non-NaN slot it first encountered
/// (or 0 if every value is NaN).
fn argextremum_dim_f32(
    name: &str,
    input: &CpuStorageBytes,
    output: &mut CpuStorageBytes,
    input_shape: &[usize],
    dim: usize,
    is_better: fn(f32, f32) -> bool,
    init: f32,
) -> Result<()> {
    if dim >= input_shape.len() {
        return Err(Error::Msg(format!(
            "{name}: dim {dim} out of range for rank {}",
            input_shape.len(),
        ))
        .bt());
    }
    let f32_size = std::mem::size_of::<f32>();
    let u32_size = std::mem::size_of::<u32>();
    let total_input: usize = input_shape.iter().product();
    if input.len_bytes() != total_input.saturating_mul(f32_size) {
        return Err(Error::Msg(format!(
            "{name}: input bytes={} doesn't match shape {input_shape:?} (f32)",
            input.len_bytes(),
        ))
        .bt());
    }
    let outer_count: usize = input_shape[..dim].iter().product();
    let dim_size = input_shape[dim];
    let inner_count: usize = input_shape[dim + 1..].iter().product();
    let output_count = outer_count * inner_count;
    if output.len_bytes() != output_count.saturating_mul(u32_size) {
        return Err(Error::Msg(format!(
            "{name}: output bytes={} doesn't match (input shape - dim {dim}) × {u32_size}",
            output.len_bytes(),
        ))
        .bt());
    }
    if dim_size == 0 {
        return Err(Error::Msg(format!(
            "{name}: dim {dim} has size 0 — argmax/argmin undefined",
        ))
        .bt());
    }
    let in_view: &[f32] = input.as_slice()?;
    let out_view: &mut [u32] = output.as_slice_mut()?;
    for outer in 0..outer_count {
        for inner in 0..inner_count {
            let mut best_val = init;
            let mut best_idx = 0u32;
            for d in 0..dim_size {
                let off = (outer * dim_size + d) * inner_count + inner;
                let v = in_view[off];
                // Initial slot: take the first valid value seen.
                if d == 0 {
                    best_val = v;
                    best_idx = 0;
                } else if is_better(v, best_val) {
                    best_val = v;
                    best_idx = d as u32;
                }
            }
            out_view[outer * inner_count + inner] = best_idx;
        }
    }
    Ok(())
}

/// Argmax along one dim — `out[i] = argmax over dim of input`.
pub fn argmax_dim_f32(
    input: &CpuStorageBytes,
    output: &mut CpuStorageBytes,
    input_shape: &[usize],
    dim: usize,
) -> Result<()> {
    argextremum_dim_f32(
        "argmax_dim_f32",
        input,
        output,
        input_shape,
        dim,
        |new, best| new > best,
        f32::NEG_INFINITY,
    )
}

/// Argmin along one dim — `out[i] = argmin over dim of input`.
pub fn argmin_dim_f32(
    input: &CpuStorageBytes,
    output: &mut CpuStorageBytes,
    input_shape: &[usize],
    dim: usize,
) -> Result<()> {
    argextremum_dim_f32(
        "argmin_dim_f32",
        input,
        output,
        input_shape,
        dim,
        |new, best| new < best,
        f32::INFINITY,
    )
}

// GELU (tanh approximation) — see chassis::unary::GeluTanh. The f64
// variant uses a 16-digit coefficient for √(2/π); the f32 variant
// uses a 7-digit coefficient. bf16 / f16 route through f32 via the
// chassis's blanket bf16 / f16 impls (see gelu_bf16 / gelu_f16
// thunks alongside the half-float unary set earlier in this file).
unary_thunk!(gelu_f32, f32, GeluTanh);
unary_thunk!(gelu_f64, f64, GeluTanh);

// =============================================================================
// Phase 7.6 step 6 follow-up — backward-helper byte-level kernels.
//
// These mirror the generic-typed reference implementations in
// fuel-reference-backend::ops::* and let the four backward helpers
// register through the byte-level KernelBindingTable + FusedKernelRegistry
// path. Before this commit, fuel-graph-executor's trait-method dispatch
// (`self.backend.softmax_last_dim_backward(...)`) was the only CPU path.
// =============================================================================

/// Softmax-last-dim backward — F32. Inputs `(y, g)` both
/// `[outer × last_dim]`. Output: `y * (g - sum(y * g, last_dim))`
/// per row.
pub fn softmax_last_dim_backward_f32(
    y: &CpuStorageBytes,
    g: &CpuStorageBytes,
    out: &mut CpuStorageBytes,
    outer_count: usize,
    last_dim: usize,
) -> Result<()> {
    let yv: &[f32] = y.as_slice()?;
    let gv: &[f32] = g.as_slice()?;
    let outv: &mut [f32] = out.as_slice_mut()?;
    let needed = outer_count * last_dim;
    if yv.len() != needed || gv.len() != needed || outv.len() != needed {
        return Err(Error::Msg(format!(
            "softmax_last_dim_backward_f32: element count mismatch \
             (y={}, g={}, out={}, expected={needed})",
            yv.len(),
            gv.len(),
            outv.len(),
        ))
        .bt());
    }
    for r in 0..outer_count {
        let off = r * last_dim;
        let mut dot = 0.0_f32;
        for i in 0..last_dim {
            dot += yv[off + i] * gv[off + i];
        }
        for i in 0..last_dim {
            outv[off + i] = yv[off + i] * (gv[off + i] - dot);
        }
    }
    Ok(())
}

/// Softmax-last-dim backward — F64.
pub fn softmax_last_dim_backward_f64(
    y: &CpuStorageBytes,
    g: &CpuStorageBytes,
    out: &mut CpuStorageBytes,
    outer_count: usize,
    last_dim: usize,
) -> Result<()> {
    let yv: &[f64] = y.as_slice()?;
    let gv: &[f64] = g.as_slice()?;
    let outv: &mut [f64] = out.as_slice_mut()?;
    let needed = outer_count * last_dim;
    if yv.len() != needed || gv.len() != needed || outv.len() != needed {
        return Err(Error::Msg(format!(
            "softmax_last_dim_backward_f64: element count mismatch \
             (y={}, g={}, out={}, expected={needed})",
            yv.len(),
            gv.len(),
            outv.len(),
        ))
        .bt());
    }
    for r in 0..outer_count {
        let off = r * last_dim;
        let mut dot = 0.0_f64;
        for i in 0..last_dim {
            dot += yv[off + i] * gv[off + i];
        }
        for i in 0..last_dim {
            outv[off + i] = yv[off + i] * (gv[off + i] - dot);
        }
    }
    Ok(())
}

macro_rules! softmax_last_dim_backward_half {
    ($name:ident, $T:ty) => {
        pub fn $name(
            y: &CpuStorageBytes,
            g: &CpuStorageBytes,
            out: &mut CpuStorageBytes,
            outer_count: usize,
            last_dim: usize,
        ) -> Result<()> {
            let yv: &[$T] = y.as_slice()?;
            let gv: &[$T] = g.as_slice()?;
            let outv: &mut [$T] = out.as_slice_mut()?;
            let needed = outer_count * last_dim;
            if yv.len() != needed || gv.len() != needed || outv.len() != needed {
                return Err(Error::Msg(format!(
                    "{}: element count mismatch (y={}, g={}, out={}, expected={needed})",
                    stringify!($name),
                    yv.len(),
                    gv.len(),
                    outv.len(),
                ))
                .bt());
            }
            for r in 0..outer_count {
                let off = r * last_dim;
                let mut dot = 0.0_f32;
                for i in 0..last_dim {
                    dot += yv[off + i].to_f32() * gv[off + i].to_f32();
                }
                for i in 0..last_dim {
                    let val = yv[off + i].to_f32() * (gv[off + i].to_f32() - dot);
                    outv[off + i] = <$T>::from_f32(val);
                }
            }
            Ok(())
        }
    };
}

softmax_last_dim_backward_half!(softmax_last_dim_backward_bf16, half::bf16);
softmax_last_dim_backward_half!(softmax_last_dim_backward_f16, half::f16);

/// LayerNorm-last-dim backward — F32. Inputs `(x, g)`; both
/// `[outer × last_dim]`. Carries `eps` to recompute the same
/// statistics the forward used. Formula:
///   `grad_x_i = rstd · (g_i - mean(g) - y_i · mean(g · y))`
/// where `y_i = (x_i - μ) · rstd`, `rstd = 1/sqrt(var + eps)`.
pub fn layer_norm_last_dim_backward_f32(
    x: &CpuStorageBytes,
    g: &CpuStorageBytes,
    out: &mut CpuStorageBytes,
    outer_count: usize,
    last_dim: usize,
    eps: f64,
) -> Result<()> {
    let xv: &[f32] = x.as_slice()?;
    let gv: &[f32] = g.as_slice()?;
    let outv: &mut [f32] = out.as_slice_mut()?;
    let needed = outer_count * last_dim;
    if xv.len() != needed || gv.len() != needed || outv.len() != needed {
        return Err(Error::Msg(format!(
            "layer_norm_last_dim_backward_f32: element count mismatch \
             (x={}, g={}, out={}, expected={needed})",
            xv.len(),
            gv.len(),
            outv.len(),
        ))
        .bt());
    }
    let n = last_dim as f32;
    let eps_t = eps as f32;
    for r in 0..outer_count {
        let off = r * last_dim;
        let mut mean = 0.0_f32;
        for i in 0..last_dim {
            mean += xv[off + i];
        }
        mean /= n;
        let mut var = 0.0_f32;
        for i in 0..last_dim {
            let d = xv[off + i] - mean;
            var += d * d;
        }
        var /= n;
        let rstd = 1.0_f32 / (var + eps_t).sqrt();
        let mut mean_g = 0.0_f32;
        let mut mean_g_y = 0.0_f32;
        for i in 0..last_dim {
            let yi = (xv[off + i] - mean) * rstd;
            mean_g += gv[off + i];
            mean_g_y += gv[off + i] * yi;
        }
        mean_g /= n;
        mean_g_y /= n;
        for i in 0..last_dim {
            let yi = (xv[off + i] - mean) * rstd;
            outv[off + i] = rstd * (gv[off + i] - mean_g - yi * mean_g_y);
        }
    }
    Ok(())
}

/// LayerNorm-last-dim backward — F64.
pub fn layer_norm_last_dim_backward_f64(
    x: &CpuStorageBytes,
    g: &CpuStorageBytes,
    out: &mut CpuStorageBytes,
    outer_count: usize,
    last_dim: usize,
    eps: f64,
) -> Result<()> {
    let xv: &[f64] = x.as_slice()?;
    let gv: &[f64] = g.as_slice()?;
    let outv: &mut [f64] = out.as_slice_mut()?;
    let needed = outer_count * last_dim;
    if xv.len() != needed || gv.len() != needed || outv.len() != needed {
        return Err(Error::Msg(format!(
            "layer_norm_last_dim_backward_f64: element count mismatch \
             (x={}, g={}, out={}, expected={needed})",
            xv.len(),
            gv.len(),
            outv.len(),
        ))
        .bt());
    }
    let n = last_dim as f64;
    for r in 0..outer_count {
        let off = r * last_dim;
        let mut mean = 0.0_f64;
        for i in 0..last_dim {
            mean += xv[off + i];
        }
        mean /= n;
        let mut var = 0.0_f64;
        for i in 0..last_dim {
            let d = xv[off + i] - mean;
            var += d * d;
        }
        var /= n;
        let rstd = 1.0_f64 / (var + eps).sqrt();
        let mut mean_g = 0.0_f64;
        let mut mean_g_y = 0.0_f64;
        for i in 0..last_dim {
            let yi = (xv[off + i] - mean) * rstd;
            mean_g += gv[off + i];
            mean_g_y += gv[off + i] * yi;
        }
        mean_g /= n;
        mean_g_y /= n;
        for i in 0..last_dim {
            let yi = (xv[off + i] - mean) * rstd;
            outv[off + i] = rstd * (gv[off + i] - mean_g - yi * mean_g_y);
        }
    }
    Ok(())
}

macro_rules! layer_norm_last_dim_backward_half {
    ($name:ident, $T:ty) => {
        pub fn $name(
            x: &CpuStorageBytes,
            g: &CpuStorageBytes,
            out: &mut CpuStorageBytes,
            outer_count: usize,
            last_dim: usize,
            eps: f64,
        ) -> Result<()> {
            let xv: &[$T] = x.as_slice()?;
            let gv: &[$T] = g.as_slice()?;
            let outv: &mut [$T] = out.as_slice_mut()?;
            let needed = outer_count * last_dim;
            if xv.len() != needed || gv.len() != needed || outv.len() != needed {
                return Err(Error::Msg(format!(
                    "{}: element count mismatch (x={}, g={}, out={}, expected={needed})",
                    stringify!($name),
                    xv.len(),
                    gv.len(),
                    outv.len(),
                ))
                .bt());
            }
            let n = last_dim as f32;
            let eps_t = eps as f32;
            for r in 0..outer_count {
                let off = r * last_dim;
                let mut mean = 0.0_f32;
                for i in 0..last_dim {
                    mean += xv[off + i].to_f32();
                }
                mean /= n;
                let mut var = 0.0_f32;
                for i in 0..last_dim {
                    let d = xv[off + i].to_f32() - mean;
                    var += d * d;
                }
                var /= n;
                let rstd = 1.0_f32 / (var + eps_t).sqrt();
                let mut mean_g = 0.0_f32;
                let mut mean_g_y = 0.0_f32;
                for i in 0..last_dim {
                    let yi = (xv[off + i].to_f32() - mean) * rstd;
                    mean_g += gv[off + i].to_f32();
                    mean_g_y += gv[off + i].to_f32() * yi;
                }
                mean_g /= n;
                mean_g_y /= n;
                for i in 0..last_dim {
                    let yi = (xv[off + i].to_f32() - mean) * rstd;
                    let val = rstd * (gv[off + i].to_f32() - mean_g - yi * mean_g_y);
                    outv[off + i] = <$T>::from_f32(val);
                }
            }
            Ok(())
        }
    };
}

layer_norm_last_dim_backward_half!(layer_norm_last_dim_backward_bf16, half::bf16);
layer_norm_last_dim_backward_half!(layer_norm_last_dim_backward_f16, half::f16);

/// RmsNorm-last-dim backward — F32. Inputs `(x, g_y)`; both
/// `[outer × last_dim]`. Closed form:
///   `grad_x_j = r_rms · (g_y_j - x_j · s / (n · (mean_sq + eps)))`
/// where `r_rms = 1/sqrt(mean_sq + eps)`, `s = sum(g_y · x)`.
pub fn rms_norm_last_dim_backward_f32(
    x: &CpuStorageBytes,
    g_y: &CpuStorageBytes,
    out: &mut CpuStorageBytes,
    outer_count: usize,
    last_dim: usize,
    eps: f64,
) -> Result<()> {
    let xv: &[f32] = x.as_slice()?;
    let gv: &[f32] = g_y.as_slice()?;
    let outv: &mut [f32] = out.as_slice_mut()?;
    let needed = outer_count * last_dim;
    if xv.len() != needed || gv.len() != needed || outv.len() != needed {
        return Err(Error::Msg(format!(
            "rms_norm_last_dim_backward_f32: element count mismatch \
             (x={}, g_y={}, out={}, expected={needed})",
            xv.len(),
            gv.len(),
            outv.len(),
        ))
        .bt());
    }
    let n = last_dim as f32;
    let eps_t = eps as f32;
    for r in 0..outer_count {
        let off = r * last_dim;
        let mut sum_sq = 0.0_f32;
        let mut sum_gx = 0.0_f32;
        for i in 0..last_dim {
            sum_sq += xv[off + i] * xv[off + i];
            sum_gx += gv[off + i] * xv[off + i];
        }
        let mean_sq = sum_sq / n;
        let denom_sq = mean_sq + eps_t;
        let r_rms = 1.0_f32 / denom_sq.sqrt();
        let coeff = sum_gx / (n * denom_sq);
        for i in 0..last_dim {
            outv[off + i] = r_rms * (gv[off + i] - xv[off + i] * coeff);
        }
    }
    Ok(())
}

/// RmsNorm-last-dim backward — F64.
pub fn rms_norm_last_dim_backward_f64(
    x: &CpuStorageBytes,
    g_y: &CpuStorageBytes,
    out: &mut CpuStorageBytes,
    outer_count: usize,
    last_dim: usize,
    eps: f64,
) -> Result<()> {
    let xv: &[f64] = x.as_slice()?;
    let gv: &[f64] = g_y.as_slice()?;
    let outv: &mut [f64] = out.as_slice_mut()?;
    let needed = outer_count * last_dim;
    if xv.len() != needed || gv.len() != needed || outv.len() != needed {
        return Err(Error::Msg(format!(
            "rms_norm_last_dim_backward_f64: element count mismatch \
             (x={}, g_y={}, out={}, expected={needed})",
            xv.len(),
            gv.len(),
            outv.len(),
        ))
        .bt());
    }
    let n = last_dim as f64;
    for r in 0..outer_count {
        let off = r * last_dim;
        let mut sum_sq = 0.0_f64;
        let mut sum_gx = 0.0_f64;
        for i in 0..last_dim {
            sum_sq += xv[off + i] * xv[off + i];
            sum_gx += gv[off + i] * xv[off + i];
        }
        let mean_sq = sum_sq / n;
        let denom_sq = mean_sq + eps;
        let r_rms = 1.0_f64 / denom_sq.sqrt();
        let coeff = sum_gx / (n * denom_sq);
        for i in 0..last_dim {
            outv[off + i] = r_rms * (gv[off + i] - xv[off + i] * coeff);
        }
    }
    Ok(())
}

macro_rules! rms_norm_last_dim_backward_half {
    ($name:ident, $T:ty) => {
        pub fn $name(
            x: &CpuStorageBytes,
            g_y: &CpuStorageBytes,
            out: &mut CpuStorageBytes,
            outer_count: usize,
            last_dim: usize,
            eps: f64,
        ) -> Result<()> {
            let xv: &[$T] = x.as_slice()?;
            let gv: &[$T] = g_y.as_slice()?;
            let outv: &mut [$T] = out.as_slice_mut()?;
            let needed = outer_count * last_dim;
            if xv.len() != needed || gv.len() != needed || outv.len() != needed {
                return Err(Error::Msg(format!(
                    "{}: element count mismatch (x={}, g_y={}, out={}, expected={needed})",
                    stringify!($name),
                    xv.len(),
                    gv.len(),
                    outv.len(),
                ))
                .bt());
            }
            let n = last_dim as f32;
            let eps_t = eps as f32;
            for r in 0..outer_count {
                let off = r * last_dim;
                let mut sum_sq = 0.0_f32;
                let mut sum_gx = 0.0_f32;
                for i in 0..last_dim {
                    let xi = xv[off + i].to_f32();
                    let gi = gv[off + i].to_f32();
                    sum_sq += xi * xi;
                    sum_gx += gi * xi;
                }
                let mean_sq = sum_sq / n;
                let denom_sq = mean_sq + eps_t;
                let r_rms = 1.0_f32 / denom_sq.sqrt();
                let coeff = sum_gx / (n * denom_sq);
                for i in 0..last_dim {
                    let xi = xv[off + i].to_f32();
                    let gi = gv[off + i].to_f32();
                    outv[off + i] = <$T>::from_f32(r_rms * (gi - xi * coeff));
                }
            }
            Ok(())
        }
    };
}

rms_norm_last_dim_backward_half!(rms_norm_last_dim_backward_bf16, half::bf16);
rms_norm_last_dim_backward_half!(rms_norm_last_dim_backward_f16, half::f16);

/// ReduceMaxTo backward — F32. Inputs `(x, upstream)`. `input_shape`
/// is x's shape; `output_shape` is upstream's shape (= the forward
/// target_shape). Routes upstream to argmax positions, splitting
/// equally on ties.
///
/// Algorithm:
/// 1. recompute the forward max via reduce_max_to (uses input_shape
///    + output_shape to determine which axes reduce);
/// 2. broadcast max to input shape, build a U8 mask where x == max;
/// 3. reduce the mask to output_shape via sum to get tie counts;
/// 4. clamp counts to >= 1 (defensive guard against empty windows);
/// 5. scaled_upstream = upstream / count;
/// 6. broadcast scaled_upstream to input_shape, gate by mask.
///
/// Implementation borrows the same broadcast-and-reduce helpers as
/// the typed reference op; the byte-level path materializes
/// intermediate buffers on the stack rather than RefTensor allocs.
pub fn reduce_max_to_backward_f32(
    x: &CpuStorageBytes,
    upstream: &CpuStorageBytes,
    out: &mut CpuStorageBytes,
    input_shape: &[usize],
    output_shape: &[usize],
) -> Result<()> {
    let xv: &[f32] = x.as_slice()?;
    let uv: &[f32] = upstream.as_slice()?;
    let outv: &mut [f32] = out.as_slice_mut()?;
    let in_count: usize = input_shape.iter().product();
    let out_count: usize = output_shape.iter().product();
    if xv.len() != in_count || outv.len() != in_count || uv.len() != out_count {
        return Err(Error::Msg(format!(
            "reduce_max_to_backward_f32: shape mismatch (x={}, up={}, out={}, \
             expected x.out={in_count}, up={out_count})",
            xv.len(),
            uv.len(),
            outv.len(),
        ))
        .bt());
    }
    reduce_max_to_backward_impl(xv, uv, outv, input_shape, output_shape);
    Ok(())
}

/// ReduceMaxTo backward — F64.
pub fn reduce_max_to_backward_f64(
    x: &CpuStorageBytes,
    upstream: &CpuStorageBytes,
    out: &mut CpuStorageBytes,
    input_shape: &[usize],
    output_shape: &[usize],
) -> Result<()> {
    let xv: &[f64] = x.as_slice()?;
    let uv: &[f64] = upstream.as_slice()?;
    let outv: &mut [f64] = out.as_slice_mut()?;
    let in_count: usize = input_shape.iter().product();
    let out_count: usize = output_shape.iter().product();
    if xv.len() != in_count || outv.len() != in_count || uv.len() != out_count {
        return Err(Error::Msg(format!(
            "reduce_max_to_backward_f64: shape mismatch (x={}, up={}, out={}, \
             expected x.out={in_count}, up={out_count})",
            xv.len(),
            uv.len(),
            outv.len(),
        ))
        .bt());
    }
    reduce_max_to_backward_impl(xv, uv, outv, input_shape, output_shape);
    Ok(())
}

macro_rules! reduce_max_to_backward_half {
    ($name:ident, $T:ty) => {
        pub fn $name(
            x: &CpuStorageBytes,
            upstream: &CpuStorageBytes,
            out: &mut CpuStorageBytes,
            input_shape: &[usize],
            output_shape: &[usize],
        ) -> Result<()> {
            let xv: &[$T] = x.as_slice()?;
            let uv: &[$T] = upstream.as_slice()?;
            let outv: &mut [$T] = out.as_slice_mut()?;
            let in_count: usize = input_shape.iter().product();
            let out_count: usize = output_shape.iter().product();
            if xv.len() != in_count || outv.len() != in_count || uv.len() != out_count {
                return Err(Error::Msg(format!(
                    "{}: shape mismatch (x={}, up={}, out={}, expected x.out={in_count}, up={out_count})",
                    stringify!($name), xv.len(), uv.len(), outv.len(),
                )).bt());
            }
            // Promote to f32, run f32 kernel, narrow on output.
            let xv32: Vec<f32> = xv.iter().map(|v| v.to_f32()).collect();
            let uv32: Vec<f32> = uv.iter().map(|v| v.to_f32()).collect();
            let mut outv32 = vec![0.0_f32; in_count];
            reduce_max_to_backward_impl(&xv32, &uv32, &mut outv32, input_shape, output_shape);
            for i in 0..in_count {
                outv[i] = <$T>::from_f32(outv32[i]);
            }
            Ok(())
        }
    };
}

reduce_max_to_backward_half!(reduce_max_to_backward_bf16, half::bf16);
reduce_max_to_backward_half!(reduce_max_to_backward_f16, half::f16);

/// Shared implementation for `reduce_max_to_backward` over a generic
/// numeric type that supports the comparisons + arithmetic used.
/// Mirrors fuel-reference-backend::ops::reduce_max_to_backward step
/// by step; see that function for the algorithmic explanation.
fn reduce_max_to_backward_impl<T>(
    x: &[T],
    upstream: &[T],
    out: &mut [T],
    input_shape: &[usize],
    output_shape: &[usize],
) where
    T: Copy
        + PartialEq
        + PartialOrd
        + Default
        + std::ops::Add<Output = T>
        + std::ops::Mul<Output = T>
        + std::ops::Div<Output = T>
        + From<u8>,
{
    let in_count = x.len();
    let one = T::from(1u8);
    // Step 1: forward max via reduce_max_to.
    let mut max_y = vec![T::default(); upstream.len()];
    reduce_max_impl(x, &mut max_y, input_shape, output_shape);
    // Step 2: broadcast max_y to input shape (in-place into a buffer);
    //         build mask: 1 where x == broadcast(max_y), else 0.
    let mut max_b = vec![T::default(); in_count];
    broadcast_to_impl(&max_y, &mut max_b, output_shape, input_shape);
    let mut mask = vec![T::default(); in_count];
    for i in 0..in_count {
        if x[i] == max_b[i] {
            mask[i] = one;
        }
    }
    // Step 3: count = reduce_sum_to(mask).
    let mut count = vec![T::default(); upstream.len()];
    reduce_sum_to_impl(&mask, &mut count, input_shape, output_shape);
    // Step 4: clamp count to >= 1.
    for c in count.iter_mut() {
        if *c < one {
            *c = one;
        }
    }
    // Step 5: scaled = upstream / count.
    let mut scaled = vec![T::default(); upstream.len()];
    for i in 0..upstream.len() {
        scaled[i] = upstream[i] / count[i];
    }
    // Step 6: broadcast scaled to input shape, gate by mask.
    let mut scaled_b = vec![T::default(); in_count];
    broadcast_to_impl(&scaled, &mut scaled_b, output_shape, input_shape);
    for i in 0..in_count {
        out[i] = scaled_b[i] * mask[i];
    }
}

/// Internal: reduce-max-to. Walks each output position, reads every
/// input position that maps to it (per NumPy-style broadcast axis
/// alignment), keeps the max. Generic so the same logic works in
/// f32 (the half-precision path) and f64.
fn reduce_max_impl<T>(input: &[T], output: &mut [T], in_shape: &[usize], out_shape: &[usize])
where
    T: Copy + PartialOrd + Default,
{
    // Pad output_shape with leading 1s to match input rank.
    let in_rank = in_shape.len();
    let out_rank = out_shape.len();
    let pad = in_rank - out_rank;
    let mut padded_out: Vec<usize> = vec![1; in_rank];
    for (i, &s) in out_shape.iter().enumerate() {
        padded_out[pad + i] = s;
    }
    let in_count: usize = in_shape.iter().product();
    // Strides for output indexing into the in-shape lattice.
    let mut out_strides = vec![0_usize; in_rank];
    let mut acc = 1;
    for i in (0..in_rank).rev() {
        out_strides[i] = if padded_out[i] == 1 { 0 } else { acc };
        if padded_out[i] != 1 {
            acc *= padded_out[i];
        }
    }
    // Initialize: copy first contributing input to each output (or
    // mark as "not yet set"). Easiest: iterate in lex order and
    // update max via comparisons.
    let mut first_set = vec![false; output.len()];
    let mut idx = vec![0_usize; in_rank];
    #[expect(
        clippy::needless_range_loop,
        reason = "flat index drives explicit multi-index reconstruction (rem = flat), not a sequential scan; enumerate would obscure the flat<->multi-index mapping and risk an iteration-count mismatch"
    )]
    for flat in 0..in_count {
        // Reconstruct multi-index `idx` from flat.
        let rem = flat;
        let mut in_stride = 1;
        for i in (0..in_rank).rev() {
            idx[i] = (rem / in_stride) % in_shape[i];
            in_stride *= in_shape[i];
            // Simpler decomposition below; use direct.
        }
        // Compute out_flat via per-axis strides.
        let mut out_flat = 0;
        for i in 0..in_rank {
            out_flat += idx[i] * out_strides[i];
        }
        let v = input[flat];
        if !first_set[out_flat] {
            output[out_flat] = v;
            first_set[out_flat] = true;
        } else if v > output[out_flat] {
            output[out_flat] = v;
        }
    }
}

/// Internal: broadcast `src` (shape `src_shape`) into `dst` (shape
/// `dst_shape`). Right-aligns axes; size-1 axes broadcast.
fn broadcast_to_impl<T>(src: &[T], dst: &mut [T], src_shape: &[usize], dst_shape: &[usize])
where
    T: Copy,
{
    let dst_rank = dst_shape.len();
    let src_rank = src_shape.len();
    let pad = dst_rank - src_rank;
    let mut padded_src: Vec<usize> = vec![1; dst_rank];
    for (i, &s) in src_shape.iter().enumerate() {
        padded_src[pad + i] = s;
    }
    // src strides (where size==1, stride is 0 so the axis is replicated).
    let mut src_strides = vec![0_usize; dst_rank];
    let mut acc = 1;
    for i in (0..dst_rank).rev() {
        src_strides[i] = if padded_src[i] == 1 { 0 } else { acc };
        if padded_src[i] != 1 {
            acc *= padded_src[i];
        }
    }
    let dst_count: usize = dst_shape.iter().product();
    let mut idx = vec![0_usize; dst_rank];
    #[expect(
        clippy::needless_range_loop,
        reason = "flat index drives explicit multi-index reconstruction (rem = flat), not a sequential scan; enumerate would obscure the flat<->multi-index mapping and risk an iteration-count mismatch"
    )]
    for flat in 0..dst_count {
        let mut rem = flat;
        for i in (0..dst_rank).rev() {
            idx[i] = rem % dst_shape[i];
            rem /= dst_shape[i];
        }
        let mut src_flat = 0;
        for i in 0..dst_rank {
            src_flat += idx[i] * src_strides[i];
        }
        dst[flat] = src[src_flat];
    }
}

/// Internal: reduce-sum-to. Same structure as `reduce_max_impl` but
/// accumulates via `+` instead of taking max.
fn reduce_sum_to_impl<T>(input: &[T], output: &mut [T], in_shape: &[usize], out_shape: &[usize])
where
    T: Copy + Default + std::ops::Add<Output = T>,
{
    let in_rank = in_shape.len();
    let out_rank = out_shape.len();
    let pad = in_rank - out_rank;
    let mut padded_out: Vec<usize> = vec![1; in_rank];
    for (i, &s) in out_shape.iter().enumerate() {
        padded_out[pad + i] = s;
    }
    let in_count: usize = in_shape.iter().product();
    let mut out_strides = vec![0_usize; in_rank];
    let mut acc = 1;
    for i in (0..in_rank).rev() {
        out_strides[i] = if padded_out[i] == 1 { 0 } else { acc };
        if padded_out[i] != 1 {
            acc *= padded_out[i];
        }
    }
    for o in output.iter_mut() {
        *o = T::default();
    }
    let mut idx = vec![0_usize; in_rank];
    #[expect(
        clippy::needless_range_loop,
        reason = "flat index drives explicit multi-index reconstruction (rem = flat), not a sequential scan; enumerate would obscure the flat<->multi-index mapping and risk an iteration-count mismatch"
    )]
    for flat in 0..in_count {
        let mut rem = flat;
        for i in (0..in_rank).rev() {
            idx[i] = rem % in_shape[i];
            rem /= in_shape[i];
        }
        let mut out_flat = 0;
        for i in 0..in_rank {
            out_flat += idx[i] * out_strides[i];
        }
        output[out_flat] = output[out_flat] + input[flat];
    }
}

// =============================================================================
// WriteSlice — Phase 7.6 step 9c E.3.2 (KV-cache writes)
// =============================================================================

/// In-place rectangular scatter write: copy `source`'s bytes into a
/// slab of `dest` defined per-axis by `ranges`. `dest_shape` is the
/// destination's shape; `ranges[i] = (start, end)` is the half-open
/// slab range along axis `i`, so the source's shape along axis `i`
/// is implicitly `end - start`.
///
/// Dtype-agnostic at the byte level — `dtype_size` scales offsets.
/// The kernel walks every "non-innermost" coordinate of the slab in
/// row-major order and copies the innermost (last-axis) row in one
/// `copy_from_slice`. Bytes outside the slab are not touched.
pub fn write_slice_cpu(
    source: &CpuStorageBytes,
    dest: &mut CpuStorageBytes,
    dest_shape: &[usize],
    ranges: &[(usize, usize)],
    dtype_size: usize,
) -> Result<()> {
    if dtype_size == 0 {
        return Err(Error::Msg("write_slice_cpu: dtype_size must be > 0".to_string()).bt());
    }
    let rank = dest_shape.len();
    if rank == 0 {
        return Err(Error::Msg(
            "write_slice_cpu: dest rank must be >= 1 (scalar destinations have \
             no slab to write into)"
                .to_string(),
        )
        .bt());
    }
    if ranges.len() != rank {
        return Err(Error::Msg(format!(
            "write_slice_cpu: ranges.len() ({}) != dest rank ({rank})",
            ranges.len(),
        ))
        .bt());
    }
    // Per-axis validation + slab shape derivation.
    let mut slab_shape: Vec<usize> = Vec::with_capacity(rank);
    for (i, &(start, end)) in ranges.iter().enumerate() {
        if end < start || end > dest_shape[i] {
            return Err(Error::Msg(format!(
                "write_slice_cpu: ranges[{i}] = ({start}, {end}) invalid for \
                 dest_shape[{i}] = {}",
                dest_shape[i],
            ))
            .bt());
        }
        slab_shape.push(end - start);
    }
    // dest_shape[i] = 0 ⇒ slab dim must also be 0 (no-op).
    let slab_elems: usize = slab_shape.iter().copied().product();
    if slab_elems == 0 {
        return Ok(());
    }
    let dest_elems: usize = dest_shape.iter().copied().product();
    let dest_bytes_needed = dest_elems.saturating_mul(dtype_size);
    if dest.len_bytes() != dest_bytes_needed {
        return Err(Error::Msg(format!(
            "write_slice_cpu: dest bytes ({}) != dest_shape {:?} * dtype_size {dtype_size}",
            dest.len_bytes(),
            dest_shape,
        ))
        .bt());
    }
    let source_bytes_needed = slab_elems.saturating_mul(dtype_size);
    if source.len_bytes() != source_bytes_needed {
        return Err(Error::Msg(format!(
            "write_slice_cpu: source bytes ({}) != slab {:?} * dtype_size {dtype_size}",
            source.len_bytes(),
            slab_shape,
        ))
        .bt());
    }
    // Row-major strides for dest (contiguous, in elements).
    let mut dest_strides: Vec<usize> = vec![0; rank];
    if rank > 0 {
        dest_strides[rank - 1] = 1;
        for i in (0..rank - 1).rev() {
            dest_strides[i] = dest_strides[i + 1] * dest_shape[i + 1];
        }
    }
    // The innermost axis is copied as one contiguous span per outer
    // coordinate tuple. The span's element count is the slab's
    // innermost dim; the dest offset along the innermost axis is
    // `ranges[rank-1].0` (the slab's start element along the inner-
    // most axis).
    let inner_axis = rank - 1;
    let inner_slab_elems = slab_shape[inner_axis];
    let inner_slab_bytes = inner_slab_elems * dtype_size;
    let inner_dest_start_elems = ranges[inner_axis].0;
    let inner_dest_start_bytes = inner_dest_start_elems * dtype_size;

    let dest_buf = dest.bytes_mut();
    let src_buf = source.bytes();
    // Walk the outer (non-innermost) coordinates of the SLAB in
    // row-major order. For each outer tuple, compute (a) the source
    // byte offset (slab strides) and (b) the destination byte offset
    // (dest strides + `ranges[i].0` offset along each axis).
    if rank == 1 {
        // Special case — only the innermost axis exists; one span.
        let dst_off = inner_dest_start_bytes;
        dest_buf[dst_off..dst_off + inner_slab_bytes].copy_from_slice(&src_buf[..inner_slab_bytes]);
        return Ok(());
    }
    // outer_count = product of slab dims except the innermost.
    let outer_count: usize = slab_shape[..inner_axis].iter().product();
    let mut outer_idx = vec![0usize; inner_axis];
    for outer_n in 0..outer_count {
        // Reconstruct outer coords from `outer_n` (row-major).
        let mut rem = outer_n;
        for i in 0..inner_axis {
            let stride: usize = slab_shape[i + 1..inner_axis]
                .iter()
                .product::<usize>()
                .max(1);
            // Slab outer dims unroll left-to-right; compute coord
            // using slab shape from axis `i` outward (excluding
            // innermost, which is contiguous-copied).
            outer_idx[i] = rem / stride;
            rem %= stride;
        }
        // Source offset: row-major over slab_shape (rank-1).
        let mut src_elem = 0usize;
        let mut src_stride = inner_slab_elems;
        for i in (0..inner_axis).rev() {
            src_elem += outer_idx[i] * src_stride;
            src_stride *= slab_shape[i];
        }
        // Dest offset: row-major over dest_shape, with each axis
        // shifted by `ranges[i].0`.
        let mut dst_elem = inner_dest_start_elems;
        for i in 0..inner_axis {
            let dest_coord = ranges[i].0 + outer_idx[i];
            dst_elem += dest_coord * dest_strides[i];
        }
        let src_off = src_elem * dtype_size;
        let dst_off = dst_elem * dtype_size;
        dest_buf[dst_off..dst_off + inner_slab_bytes]
            .copy_from_slice(&src_buf[src_off..src_off + inner_slab_bytes]);
    }
    Ok(())
}

// =============================================================================
// WriteSliceRotating — Phase C (sliding-window KV caches)
// =============================================================================

/// In-place ring-buffer scatter write: copy `source`'s bytes into a
/// slab of `dest` where the `axis` axis wraps modulo `modulus`. The
/// dynamic write start on `axis` comes from `position_bytes` (the
/// first `u32` is read host-side). On other axes, `ranges` defines
/// the slab the same way as [`write_slice_cpu`].
///
/// With `wrapped_start = (position % modulus)` and
/// `slab_axis_len = ranges[axis].1 - ranges[axis].0`, the kernel
/// writes rows whose `axis` coord is
/// `(wrapped_start + r) % modulus` for `r in 0..slab_axis_len`.
/// When `wrapped_start + slab_axis_len > modulus` the write splits
/// across the ring boundary; both halves are issued as separate
/// `write_slice_cpu` calls.
///
/// **v1 constraint**: `axis == 0` (the rotating axis must be the
/// leading dim). This makes the source-side byte split a trivial
/// prefix/suffix on the source buffer. Mistral / Phi-3 sliding-
/// window caches store per-layer K/V as `[seq, n_kv_heads, head_dim]`
/// with seq as axis 0, so this constraint is fine. Strided multi-
/// axis splits are a follow-up if needed.
pub fn write_slice_rotating_cpu(
    source: &CpuStorageBytes,
    position_bytes: &CpuStorageBytes,
    dest: &mut CpuStorageBytes,
    dest_shape: &[usize],
    axis: usize,
    modulus: usize,
    ranges: &[(usize, usize)],
    dtype_size: usize,
) -> Result<()> {
    if dtype_size == 0 {
        return Err(
            Error::Msg("write_slice_rotating_cpu: dtype_size must be > 0".to_string()).bt(),
        );
    }
    let rank = dest_shape.len();
    if rank == 0 {
        return Err(
            Error::Msg("write_slice_rotating_cpu: dest rank must be >= 1".to_string()).bt(),
        );
    }
    if ranges.len() != rank {
        return Err(Error::Msg(format!(
            "write_slice_rotating_cpu: ranges.len() ({}) != dest rank ({rank})",
            ranges.len(),
        ))
        .bt());
    }
    if axis >= rank {
        return Err(Error::Msg(format!(
            "write_slice_rotating_cpu: axis {axis} out of bounds for rank {rank}",
        ))
        .bt());
    }
    if modulus == 0 || modulus > dest_shape[axis] {
        return Err(Error::Msg(format!(
            "write_slice_rotating_cpu: modulus {modulus} invalid for dest_shape[{axis}] = {}",
            dest_shape[axis],
        ))
        .bt());
    }
    if position_bytes.len_bytes() < 4 {
        return Err(Error::Msg(format!(
            "write_slice_rotating_cpu: position storage has {} bytes, need >= 4",
            position_bytes.len_bytes(),
        ))
        .bt());
    }
    let pos_buf = position_bytes.bytes();
    let position = u32::from_ne_bytes([pos_buf[0], pos_buf[1], pos_buf[2], pos_buf[3]]) as usize;
    let wrapped_start = position % modulus;

    let mut slab_shape: Vec<usize> = Vec::with_capacity(rank);
    for (i, &(start, end)) in ranges.iter().enumerate() {
        if end < start {
            return Err(Error::Msg(format!(
                "write_slice_rotating_cpu: ranges[{i}] = ({start}, {end}) has end < start",
            ))
            .bt());
        }
        let slab = end - start;
        if i == axis {
            if slab > modulus {
                return Err(Error::Msg(format!(
                    "write_slice_rotating_cpu: rotating-axis slab {slab} > modulus {modulus}",
                ))
                .bt());
            }
        } else if end > dest_shape[i] {
            return Err(Error::Msg(format!(
                "write_slice_rotating_cpu: ranges[{i}] = ({start}, {end}) > dest_shape[{i}] = {}",
                dest_shape[i],
            ))
            .bt());
        }
        slab_shape.push(slab);
    }
    let slab_elems: usize = slab_shape.iter().copied().product();
    if slab_elems == 0 {
        return Ok(());
    }
    let dest_elems: usize = dest_shape.iter().copied().product();
    let dest_bytes_needed = dest_elems.saturating_mul(dtype_size);
    if dest.len_bytes() != dest_bytes_needed {
        return Err(Error::Msg(format!(
            "write_slice_rotating_cpu: dest bytes ({}) != dest_shape {:?} * dtype_size {dtype_size}",
            dest.len_bytes(), dest_shape,
        ))
        .bt());
    }
    let source_bytes_needed = slab_elems.saturating_mul(dtype_size);
    if source.len_bytes() != source_bytes_needed {
        return Err(Error::Msg(format!(
            "write_slice_rotating_cpu: source bytes ({}) != slab {:?} * dtype_size {dtype_size}",
            source.len_bytes(),
            slab_shape,
        ))
        .bt());
    }
    // Split the write across the ring boundary on the rotating axis.
    let slab_axis_len = slab_shape[axis];
    let first_len = slab_axis_len.min(modulus - wrapped_start);
    let second_len = slab_axis_len - first_len;

    // Decompose the slab into outer/axis/inner tiles for the strided
    // source extract. With slab_shape = `[d0, …, slab_axis_len, …, dn]`
    // and `axis = K`:
    //   outer_count   = product of slab dims before axis K
    //   inner_per_row = product of slab dims after axis K
    //
    // The source bytes are `outer_count × slab_axis_len × inner_per_row × dtype_size`
    // in row-major order, so each "row" along axis K is a contiguous
    // run of `inner_per_row * dtype_size` bytes, and each outer-tuple
    // is a contiguous run of `slab_axis_len * inner_per_row * dtype_size`
    // bytes. For axis == 0 this reduces to a flat byte prefix/suffix
    // (outer_count == 1).
    let outer_count: usize = slab_shape[..axis].iter().copied().product();
    let inner_per_row: usize = slab_shape[axis + 1..].iter().copied().product();
    let row_bytes = inner_per_row * dtype_size;
    let stride_bytes = slab_axis_len * row_bytes;
    let src_buf = source.bytes();

    if first_len > 0 {
        let mut sub_ranges = ranges.to_vec();
        sub_ranges[axis] = (wrapped_start, wrapped_start + first_len);
        let first_chunk_row_bytes = first_len * row_bytes;
        let src_first =
            extract_strided_chunk(src_buf, outer_count, stride_bytes, 0, first_chunk_row_bytes);
        write_slice_cpu(&src_first, dest, dest_shape, &sub_ranges, dtype_size)?;
    }
    if second_len > 0 {
        let mut sub_ranges = ranges.to_vec();
        sub_ranges[axis] = (0, second_len);
        let second_chunk_row_bytes = second_len * row_bytes;
        // The second chunk starts at `first_len * row_bytes` into each
        // outer tuple.
        let chunk_offset_in_outer = first_len * row_bytes;
        let src_second = extract_strided_chunk(
            src_buf,
            outer_count,
            stride_bytes,
            chunk_offset_in_outer,
            second_chunk_row_bytes,
        );
        write_slice_cpu(&src_second, dest, dest_shape, &sub_ranges, dtype_size)?;
    }
    Ok(())
}

/// In-place scatter write with a **device-resident-style** start on
/// `axis` read from `offset_bytes` (the first `i64`, host-side here on
/// CPU). Backs `fuel_graph::Op::WriteSliceDoff` — the CUDA-graph-
/// capturable KV-cache append. No modulo wrap (unlike
/// [`write_slice_rotating_cpu`]): the write lands at
/// `[offset, offset + width)` on `axis`, where `width` is
/// `ranges[axis].1 - ranges[axis].0`. On other axes, `ranges` defines
/// the slab the same way as [`write_slice_cpu`].
///
/// Unlike the CUDA `_doff` kernel (which reads the start device-side
/// and does NOT clamp, delegating bounds to the caller), the CPU path
/// reads the start host-side and CAN — so it bounds-checks
/// `offset + width <= dest_shape[axis]` and returns a typed error on
/// overflow rather than corrupting memory.
pub fn write_slice_doff_cpu(
    source: &CpuStorageBytes,
    offset_bytes: &CpuStorageBytes,
    dest: &mut CpuStorageBytes,
    dest_shape: &[usize],
    axis: usize,
    ranges: &[(usize, usize)],
    dtype_size: usize,
) -> Result<()> {
    if dtype_size == 0 {
        return Err(Error::Msg("write_slice_doff_cpu: dtype_size must be > 0".to_string()).bt());
    }
    let rank = dest_shape.len();
    if rank == 0 {
        return Err(Error::Msg("write_slice_doff_cpu: dest rank must be >= 1".to_string()).bt());
    }
    if ranges.len() != rank {
        return Err(Error::Msg(format!(
            "write_slice_doff_cpu: ranges.len() ({}) != dest rank ({rank})",
            ranges.len(),
        ))
        .bt());
    }
    if axis >= rank {
        return Err(Error::Msg(format!(
            "write_slice_doff_cpu: axis {axis} out of bounds for rank {rank}",
        ))
        .bt());
    }
    if offset_bytes.len_bytes() < 8 {
        return Err(Error::Msg(format!(
            "write_slice_doff_cpu: offset storage has {} bytes, need >= 8 (I64)",
            offset_bytes.len_bytes(),
        ))
        .bt());
    }
    let off_buf = offset_bytes.bytes();
    let offset_i64 = i64::from_ne_bytes([
        off_buf[0], off_buf[1], off_buf[2], off_buf[3], off_buf[4], off_buf[5], off_buf[6],
        off_buf[7],
    ]);
    if offset_i64 < 0 {
        return Err(Error::Msg(format!(
            "write_slice_doff_cpu: device offset {offset_i64} is negative",
        ))
        .bt());
    }
    let offset = offset_i64 as usize;

    // The dynamic-axis width is the range span; other axes keep their
    // ranges. Build the resolved ranges (offset overrides axis start).
    let width = {
        let (start, end) = ranges[axis];
        if end < start {
            return Err(Error::Msg(format!(
                "write_slice_doff_cpu: ranges[axis={axis}] = ({start}, {end}) has end < start",
            ))
            .bt());
        }
        end - start
    };
    if offset + width > dest_shape[axis] {
        return Err(Error::Msg(format!(
            "write_slice_doff_cpu: offset {offset} + width {width} > dest_shape[{axis}] = {} \
             (write would overflow the buffer)",
            dest_shape[axis],
        ))
        .bt());
    }
    let mut resolved = ranges.to_vec();
    resolved[axis] = (offset, offset + width);
    write_slice_cpu(source, dest, dest_shape, &resolved, dtype_size)
}

/// Extract `outer_count` tiles of `chunk_row_bytes` from `src` into a
/// fresh contiguous CPU storage. Tile `t` is at byte range
/// `t * stride_bytes + offset_in_outer .. + chunk_row_bytes`. Used by
/// `write_slice_rotating_cpu` to gather the source chunk corresponding
/// to one half of a ring-boundary split for any rotating axis.
fn extract_strided_chunk(
    src: &[u8],
    outer_count: usize,
    stride_bytes: usize,
    offset_in_outer: usize,
    chunk_row_bytes: usize,
) -> CpuStorageBytes {
    let total = outer_count * chunk_row_bytes;
    let mut buf = vec![0_u8; total];
    for t in 0..outer_count {
        let src_off = t * stride_bytes + offset_in_outer;
        let dst_off = t * chunk_row_bytes;
        buf[dst_off..dst_off + chunk_row_bytes]
            .copy_from_slice(&src[src_off..src_off + chunk_row_bytes]);
    }
    CpuStorageBytes::from_bytes(&buf)
}

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

    /// Smoke: add two f32 storages elementwise.
    #[test]
    fn add_f32_round_trip() {
        let a = CpuStorageBytes::from_slice(&[1.0_f32, 2.0, 3.0, 4.0]);
        let b = CpuStorageBytes::from_slice(&[10.0_f32, 20.0, 30.0, 40.0]);
        let mut out = CpuStorageBytes::from_zero_bytes(16);

        add_f32(&a, &b, &mut out).expect("add");

        let result: &[f32] = out.as_slice().unwrap();
        assert_eq!(result, &[11.0, 22.0, 33.0, 44.0]);
    }

    /// Mismatched byte counts produce an error, not a panic.
    #[test]
    fn add_f32_errors_on_size_mismatch() {
        let a = CpuStorageBytes::from_slice(&[1.0_f32, 2.0]);
        let b = CpuStorageBytes::from_slice(&[10.0_f32, 20.0, 30.0]);
        let mut out = CpuStorageBytes::from_zero_bytes(8);

        let result = add_f32(&a, &b, &mut out);
        assert!(result.is_err(), "size mismatch must error");
    }

    #[test]
    fn sub_mul_div_f32_round_trip() {
        let a = CpuStorageBytes::from_slice(&[10.0_f32, 20.0, 30.0]);
        let b = CpuStorageBytes::from_slice(&[1.0_f32, 4.0, 5.0]);
        let mut out = CpuStorageBytes::from_zero_bytes(12);

        sub_f32(&a, &b, &mut out).expect("sub");
        assert_eq!(out.as_slice::<f32>().unwrap(), &[9.0, 16.0, 25.0]);

        mul_f32(&a, &b, &mut out).expect("mul");
        assert_eq!(out.as_slice::<f32>().unwrap(), &[10.0, 80.0, 150.0]);

        div_f32(&a, &b, &mut out).expect("div");
        assert_eq!(out.as_slice::<f32>().unwrap(), &[10.0, 5.0, 6.0]);
    }

    /// FlashAttn backward sanity: zero upstream → zero gradients;
    /// non-zero upstream → non-zero gradients with correct shapes.
    /// Exercises the dispatch chain (which-selector + buffer plumbing)
    /// and confirms the math runs end-to-end. Numerical correctness
    /// is enforced by the reference impl that the macro mirrors.
    #[test]
    fn flash_attn_backward_f32_sanity() {
        let b = 1usize;
        let hq = 1usize;
        let hkv = 1usize;
        let sq = 2usize;
        let sk = 3usize;
        let d = 4usize;
        let scale: f32 = 1.0 / (d as f32).sqrt();
        let q_n = b * hq * sq * d;
        let kv_n = b * hkv * sk * d;

        let q_host: Vec<f32> = (0..q_n).map(|i| (i as f32) * 0.1).collect();
        let k_host: Vec<f32> = (0..kv_n).map(|i| (i as f32 + 1.0) * 0.07).collect();
        let v_host: Vec<f32> = (0..kv_n).map(|i| (i as f32 + 2.0) * 0.05).collect();

        let q = CpuStorageBytes::from_slice(&q_host);
        let k = CpuStorageBytes::from_slice(&k_host);
        let v = CpuStorageBytes::from_slice(&v_host);

        // Case 1: zero dO → zero gradients.
        let zero_do = CpuStorageBytes::from_zero_bytes(q_n * 4);
        let mut dq_out = CpuStorageBytes::from_zero_bytes(q_n * 4);
        flash_attn_backward_f32(
            &q,
            &k,
            &v,
            &zero_do,
            None,
            &mut dq_out,
            FaBackwardWhich::Q,
            b,
            hq,
            hkv,
            sq,
            sk,
            d,
            scale,
            false,
            None,
            None,
            None,
        )
        .expect("dq with zero do");
        let dq: &[f32] = dq_out.as_slice().unwrap();
        for (i, &g) in dq.iter().enumerate() {
            assert_eq!(g, 0.0, "zero dO → dQ[{i}] must be 0, got {g}");
        }

        // Case 2: non-zero dO → non-zero gradients on all three of dQ/dK/dV.
        let do_host: Vec<f32> = (0..q_n).map(|i| (i as f32 + 0.5) * 0.03).collect();
        let do_grad = CpuStorageBytes::from_slice(&do_host);
        let mut dq_buf = CpuStorageBytes::from_zero_bytes(q_n * 4);
        let mut dk_buf = CpuStorageBytes::from_zero_bytes(kv_n * 4);
        let mut dv_buf = CpuStorageBytes::from_zero_bytes(kv_n * 4);
        flash_attn_backward_f32(
            &q,
            &k,
            &v,
            &do_grad,
            None,
            &mut dq_buf,
            FaBackwardWhich::Q,
            b,
            hq,
            hkv,
            sq,
            sk,
            d,
            scale,
            false,
            None,
            None,
            None,
        )
        .expect("dq nonzero");
        flash_attn_backward_f32(
            &q,
            &k,
            &v,
            &do_grad,
            None,
            &mut dk_buf,
            FaBackwardWhich::K,
            b,
            hq,
            hkv,
            sq,
            sk,
            d,
            scale,
            false,
            None,
            None,
            None,
        )
        .expect("dk nonzero");
        flash_attn_backward_f32(
            &q,
            &k,
            &v,
            &do_grad,
            None,
            &mut dv_buf,
            FaBackwardWhich::V,
            b,
            hq,
            hkv,
            sq,
            sk,
            d,
            scale,
            false,
            None,
            None,
            None,
        )
        .expect("dv nonzero");
        let dq: &[f32] = dq_buf.as_slice().unwrap();
        let dk: &[f32] = dk_buf.as_slice().unwrap();
        let dv: &[f32] = dv_buf.as_slice().unwrap();
        let dq_norm: f32 = dq.iter().map(|x| x * x).sum::<f32>().sqrt();
        let dk_norm: f32 = dk.iter().map(|x| x * x).sum::<f32>().sqrt();
        let dv_norm: f32 = dv.iter().map(|x| x * x).sum::<f32>().sqrt();
        assert!(dq_norm > 1e-6, "dQ should be non-zero, ||dQ||={dq_norm}");
        assert!(dk_norm > 1e-6, "dK should be non-zero, ||dK||={dk_norm}");
        assert!(dv_norm > 1e-6, "dV should be non-zero, ||dV||={dv_norm}");

        // Causal mask shouldn't blow anything up either.
        let mut dq_causal = CpuStorageBytes::from_zero_bytes(q_n * 4);
        flash_attn_backward_f32(
            &q,
            &k,
            &v,
            &do_grad,
            None,
            &mut dq_causal,
            FaBackwardWhich::Q,
            b,
            hq,
            hkv,
            sq,
            sk,
            d,
            scale,
            true,
            None,
            None,
            None,
        )
        .expect("dq causal");
        let dq_c: &[f32] = dq_causal.as_slice().unwrap();
        let dq_c_norm: f32 = dq_c.iter().map(|x| x * x).sum::<f32>().sqrt();
        assert!(
            dq_c_norm > 1e-6,
            "dQ causal should be non-zero, ||dQ||={dq_c_norm}"
        );
        // Causal masks out keys → typically changes the gradient.
        assert!(
            (dq_c_norm - dq_norm).abs() > 1e-6 || (dq_c_norm == dq_norm),
            "causal vs non-causal dQ should differ (or coincidentally match)",
        );
    }

    #[test]
    fn binary_f32_size_mismatch_errors() {
        let a = CpuStorageBytes::from_slice(&[1.0_f32, 2.0]);
        let b = CpuStorageBytes::from_slice(&[1.0_f32]);
        let mut out = CpuStorageBytes::from_zero_bytes(8);
        for f in [sub_f32, mul_f32, div_f32] {
            assert!(f(&a, &b, &mut out).is_err(), "size mismatch must error");
        }
    }

    #[test]
    fn relu_f32_clips_negatives() {
        let input = CpuStorageBytes::from_slice(&[-1.0_f32, 0.0, 0.5, -3.5, 7.25]);
        let mut out = CpuStorageBytes::from_zero_bytes(input.len_bytes());
        relu_f32(&input, &mut out).expect("relu");
        assert_eq!(out.as_slice::<f32>().unwrap(), &[0.0, 0.0, 0.5, 0.0, 7.25]);
    }

    /// Pinned NaN-semantics convention (2026-07-08, decisions log): Relu is
    /// NaN-propagating, torch parity (`torch.relu(nan) == nan`). Also checks
    /// the non-NaN sanity cases (`relu(-2)=0`, `relu(3)=3`) that motivated
    /// the convention pin.
    #[test]
    fn relu_f32_propagates_nan() {
        let input = CpuStorageBytes::from_slice(&[f32::NAN, -2.0, 3.0]);
        let mut out = CpuStorageBytes::from_zero_bytes(input.len_bytes());
        relu_f32(&input, &mut out).expect("relu");
        let r: &[f32] = out.as_slice().unwrap();
        assert!(r[0].is_nan(), "relu(NaN) must propagate NaN, got {}", r[0]);
        assert_eq!(r[1], 0.0, "relu(-2) must clip to 0");
        assert_eq!(r[2], 3.0, "relu(3) must pass through");
    }

    /// KISS-OPS-6.15-0002: `relu` is `select(x<0, 0, x)`, which PRESERVES
    /// `-0.0` — it is NOT `max(x, 0)`, which normalizes `-0.0` to `+0.0`.
    #[test]
    fn relu_f32_preserves_negative_zero() {
        let input = CpuStorageBytes::from_slice(&[-0.0_f32, 0.0_f32]);
        let mut out = CpuStorageBytes::from_zero_bytes(input.len_bytes());
        relu_f32(&input, &mut out).expect("relu");
        let r: &[f32] = out.as_slice().unwrap();
        assert!(
            r[0].is_sign_negative() && r[0] == 0.0,
            "relu(-0.0) must stay -0.0 (sign bit set), got bits {:#010x}",
            r[0].to_bits()
        );
        assert!(
            r[1].is_sign_positive() && r[1] == 0.0,
            "relu(+0.0) must stay +0.0"
        );
    }

    #[test]
    fn relu_f64_preserves_negative_zero() {
        let input = CpuStorageBytes::from_slice(&[-0.0_f64]);
        let mut out = CpuStorageBytes::from_zero_bytes(input.len_bytes());
        relu_f64(&input, &mut out).expect("relu f64");
        let r: &[f64] = out.as_slice().unwrap();
        assert!(
            r[0].is_sign_negative() && r[0] == 0.0,
            "relu(-0.0_f64) must stay -0.0, got bits {:#018x}",
            r[0].to_bits()
        );
    }

    #[test]
    fn unary_f32_round_trip() {
        let input = CpuStorageBytes::from_slice(&[1.0_f32, -2.0, 4.0, 0.5]);
        let mut out = CpuStorageBytes::from_zero_bytes(input.len_bytes());

        neg_f32(&input, &mut out).expect("neg");
        assert_eq!(out.as_slice::<f32>().unwrap(), &[-1.0, 2.0, -4.0, -0.5]);

        sqr_f32(&input, &mut out).expect("sqr");
        assert_eq!(out.as_slice::<f32>().unwrap(), &[1.0, 4.0, 16.0, 0.25]);

        let pos = CpuStorageBytes::from_slice(&[1.0_f32, 4.0, 9.0, 16.0]);
        let mut sqrt_out = CpuStorageBytes::from_zero_bytes(pos.len_bytes());
        sqrt_f32(&pos, &mut sqrt_out).expect("sqrt");
        assert_eq!(sqrt_out.as_slice::<f32>().unwrap(), &[1.0, 2.0, 3.0, 4.0]);

        recip_f32(&pos, &mut sqrt_out).expect("recip");
        assert_eq!(
            sqrt_out.as_slice::<f32>().unwrap(),
            &[1.0, 0.25, 1.0 / 9.0, 1.0 / 16.0]
        );

        abs_f32(&input, &mut out).expect("abs");
        assert_eq!(out.as_slice::<f32>().unwrap(), &[1.0, 2.0, 4.0, 0.5]);
    }

    #[test]
    fn tanh_f32_at_known_points() {
        let input = CpuStorageBytes::from_slice(&[0.0_f32, 1.0, -1.0]);
        let mut out = CpuStorageBytes::from_zero_bytes(input.len_bytes());
        tanh_f32(&input, &mut out).expect("tanh");
        let result: &[f32] = out.as_slice().unwrap();
        assert!((result[0] - 0.0).abs() < 1e-7);
        assert!((result[1] - 1.0_f32.tanh()).abs() < 1e-7);
        assert!((result[2] - (-1.0_f32).tanh()).abs() < 1e-7);
    }

    #[test]
    fn unary_f32_size_mismatch_errors() {
        let input = CpuStorageBytes::from_slice(&[1.0_f32, 2.0]);
        let mut out = CpuStorageBytes::from_zero_bytes(4); // wrong size
        for f in [
            relu_f32,
            neg_f32,
            sqr_f32,
            sqrt_f32,
            recip_f32,
            abs_f32,
            tanh_f32,
            exp_f32,
            log_f32,
            sin_f32,
            cos_f32,
            sigmoid_f32,
            silu_f32,
            step_f32,
            gelu_f32,
        ] {
            assert!(f(&input, &mut out).is_err(), "size mismatch must error");
        }
    }

    #[test]
    fn exp_log_round_trip() {
        let input = CpuStorageBytes::from_slice(&[1.0_f32, 2.0, 3.0]);
        let mut intermediate = CpuStorageBytes::from_zero_bytes(input.len_bytes());
        let mut out = CpuStorageBytes::from_zero_bytes(input.len_bytes());

        exp_f32(&input, &mut intermediate).expect("exp");
        log_f32(&intermediate, &mut out).expect("log");

        let result: &[f32] = out.as_slice().unwrap();
        for (got, want) in result.iter().zip(&[1.0_f32, 2.0, 3.0]) {
            assert!((got - want).abs() < 1e-5, "exp+log not identity");
        }
    }

    #[test]
    fn sin_cos_at_known_points() {
        let input = CpuStorageBytes::from_slice(&[0.0_f32, std::f32::consts::FRAC_PI_2]);
        let mut out = CpuStorageBytes::from_zero_bytes(input.len_bytes());

        sin_f32(&input, &mut out).expect("sin");
        let r: &[f32] = out.as_slice().unwrap();
        assert!(r[0].abs() < 1e-7);
        assert!((r[1] - 1.0).abs() < 1e-6);

        cos_f32(&input, &mut out).expect("cos");
        let r: &[f32] = out.as_slice().unwrap();
        assert!((r[0] - 1.0).abs() < 1e-6);
        assert!(r[1].abs() < 1e-6);
    }

    #[test]
    fn sigmoid_silu_at_known_points() {
        let input = CpuStorageBytes::from_slice(&[0.0_f32, 5.0, -5.0]);
        let mut out = CpuStorageBytes::from_zero_bytes(input.len_bytes());

        sigmoid_f32(&input, &mut out).expect("sigmoid");
        let r: &[f32] = out.as_slice().unwrap();
        assert!((r[0] - 0.5).abs() < 1e-7);
        assert!(r[1] > 0.99);
        assert!(r[2] < 0.01);

        silu_f32(&input, &mut out).expect("silu");
        let r: &[f32] = out.as_slice().unwrap();
        assert!(r[0].abs() < 1e-7); // 0 * sigmoid(0) = 0
    }

    #[test]
    fn step_clips_correctly() {
        let input = CpuStorageBytes::from_slice(&[-2.0_f32, 0.0, 0.5, 100.0]);
        let mut out = CpuStorageBytes::from_zero_bytes(input.len_bytes());
        step_f32(&input, &mut out).expect("step");
        assert_eq!(out.as_slice::<f32>().unwrap(), &[0.0, 0.0, 1.0, 1.0]);
    }

    #[test]
    fn qmatmul_q4_0_f32_zero_activations_yields_zero() {
        // Smallest valid shape: n=2, k=32 (one Q4_0 block per row).
        // With zero activations, the weight contents are irrelevant —
        // every dot product is 0. Sanity-checks the dispatch + size
        // validation without needing valid quantized data.
        let act = CpuStorageBytes::from_slice(&[0.0_f32; 32]);
        let block_size = std::mem::size_of::<fuel_quantized::BlockQ4_0>();
        // Must be a multiple of 4 for U32-aligned storage.
        assert!(block_size.is_multiple_of(2));
        let w = CpuStorageBytes::from_bytes(&vec![0u8; 2 * block_size]);
        let mut out = CpuStorageBytes::from_zero_bytes(2 * 4);
        qmatmul_q4_0_f32(&act, &w, &mut out, 1, 1, 2, 32).expect("qmatmul");
        let r: &[f32] = out.as_slice().unwrap();
        assert_eq!(r, &[0.0, 0.0]);
    }

    #[test]
    fn qmatmul_q4_0_f32_unit_weight_sums_activations() {
        // Construct a Q4_0 weight where every weight = 1.0 by
        // setting d (scale) = 1.0 and every nibble = 9 (so the
        // effective weight is 1 * (9 - 8) = 1).
        // Then A @ W^T computes per-row sum of activations.
        use half::f16;
        let block_size = std::mem::size_of::<fuel_quantized::BlockQ4_0>();
        let mut w_bytes = vec![0u8; 2 * block_size];
        for block_idx in 0..2 {
            let off = block_idx * block_size;
            // d = f16(1.0) — little-endian bytes
            let d_bytes = f16::from_f32(1.0).to_le_bytes();
            w_bytes[off..off + 2].copy_from_slice(&d_bytes);
            // Every nibble = 9 → packed byte = 0x99 (low nibble first)
            for i in 0..16 {
                w_bytes[off + 2 + i] = 0x99;
            }
        }
        let w = CpuStorageBytes::from_bytes(&w_bytes);

        // Activations: [1, 2, 3, ..., 32]; sum = 528.
        let act_vec: Vec<f32> = (1..=32).map(|x| x as f32).collect();
        let act = CpuStorageBytes::from_slice(&act_vec);
        let mut out = CpuStorageBytes::from_zero_bytes(2 * 4);
        qmatmul_q4_0_f32(&act, &w, &mut out, 1, 1, 2, 32).expect("qmatmul");
        let r: &[f32] = out.as_slice().unwrap();
        // Both output rows = sum over k of activations = 528.
        // Tolerance ~0.5 — fuel_quantized's matmul re-quantizes the
        // f32 activations to Q8_1 internally, introducing small
        // round-trip error.
        assert!((r[0] - 528.0).abs() < 0.5, "got {}, want 528", r[0]);
        assert!((r[1] - 528.0).abs() < 0.5, "got {}, want 528", r[1]);
    }

    #[test]
    fn qmatmul_q4_0_f32_rejects_bad_k() {
        // k=33 isn't a multiple of 32 — must error.
        let act = CpuStorageBytes::from_slice(&[0.0_f32; 33]);
        let w = CpuStorageBytes::from_bytes(&[0u8; 18]); // one block
        let mut out = CpuStorageBytes::from_zero_bytes(4);
        let r = qmatmul_q4_0_f32(&act, &w, &mut out, 1, 1, 1, 33);
        assert!(r.is_err(), "k must be a multiple of 32 for Q4_0");
    }

    #[test]
    fn softmax_last_dim_bf16_uniform_row() {
        let v: Vec<half::bf16> = [1.0_f32, 1.0, 1.0, 1.0]
            .iter()
            .map(|&x| half::bf16::from_f32(x))
            .collect();
        let input = CpuStorageBytes::from_slice(&v);
        let mut out = CpuStorageBytes::from_zero_bytes(4 * 2);
        softmax_last_dim_bf16(&input, &mut out, 1, 4).expect("softmax bf16");
        let r: &[half::bf16] = out.as_slice().unwrap();
        for x in r {
            assert!((x.to_f32() - 0.25).abs() < 0.01);
        }
    }

    #[test]
    fn softmax_last_dim_f16_sums_to_one() {
        let v: Vec<half::f16> = [1.0_f32, 2.0, 3.0]
            .iter()
            .map(|&x| half::f16::from_f32(x))
            .collect();
        let input = CpuStorageBytes::from_slice(&v);
        let mut out = CpuStorageBytes::from_zero_bytes(3 * 2);
        softmax_last_dim_f16(&input, &mut out, 1, 3).expect("softmax f16");
        let r: &[half::f16] = out.as_slice().unwrap();
        let sum: f32 = r.iter().map(|x| x.to_f32()).sum();
        assert!((sum - 1.0).abs() < 0.01);
    }

    #[test]
    fn rms_norm_last_dim_bf16_basic() {
        let v: Vec<half::bf16> = [3.0_f32, 4.0]
            .iter()
            .map(|&x| half::bf16::from_f32(x))
            .collect();
        let input = CpuStorageBytes::from_slice(&v);
        let mut out = CpuStorageBytes::from_zero_bytes(2 * 2);
        rms_norm_last_dim_bf16(&input, &mut out, 1, 2, 0.0).expect("rms_norm bf16");
        let r: &[half::bf16] = out.as_slice().unwrap();
        let rms = (12.5_f32).sqrt();
        // bf16 has ~3 decimal digits — accept small absolute error.
        assert!((r[0].to_f32() - 3.0 / rms).abs() < 0.05);
        assert!((r[1].to_f32() - 4.0 / rms).abs() < 0.05);
    }

    #[test]
    fn layer_norm_last_dim_f16_zero_mean() {
        let v: Vec<half::f16> = [1.0_f32, 2.0, 3.0]
            .iter()
            .map(|&x| half::f16::from_f32(x))
            .collect();
        let input = CpuStorageBytes::from_slice(&v);
        let mut out = CpuStorageBytes::from_zero_bytes(3 * 2);
        layer_norm_last_dim_f16(&input, &mut out, 1, 3, 0.0).expect("layer_norm f16");
        let r: &[half::f16] = out.as_slice().unwrap();
        // Output mean should be ≈ 0; output var should be ≈ 1.
        let result_f32: Vec<f32> = r.iter().map(|x| x.to_f32()).collect();
        let mean: f32 = result_f32.iter().sum::<f32>() / 3.0;
        assert!(mean.abs() < 0.01);
        let var: f32 = result_f32.iter().map(|v| (v - mean).powi(2)).sum::<f32>() / 3.0;
        assert!((var - 1.0).abs() < 0.05);
    }

    #[test]
    fn rope_bf16_pi_over_two_swaps_with_sign() {
        // x [1, 1, 4] = [1, 2, 3, 4]; cos=0, sin=1.
        // Expected: [-3, -4, 1, 2].
        let x_v: Vec<half::bf16> = [1.0_f32, 2.0, 3.0, 4.0]
            .iter()
            .map(|&x| half::bf16::from_f32(x))
            .collect();
        let zero_v: Vec<half::bf16> = [0.0_f32; 4]
            .iter()
            .map(|&x| half::bf16::from_f32(x))
            .collect();
        let one_v: Vec<half::bf16> = [1.0_f32; 4]
            .iter()
            .map(|&x| half::bf16::from_f32(x))
            .collect();
        let x = CpuStorageBytes::from_slice(&x_v);
        let cos = CpuStorageBytes::from_slice(&zero_v);
        let sin = CpuStorageBytes::from_slice(&one_v);
        let mut out = CpuStorageBytes::from_zero_bytes(4 * 2);
        rope_bf16(&x, &cos, &sin, &mut out, 1, 1, 4).expect("rope bf16");
        let r: &[half::bf16] = out.as_slice().unwrap();
        let result_f32: Vec<f32> = r.iter().map(|x| x.to_f32()).collect();
        assert_eq!(result_f32, vec![-3.0, -4.0, 1.0, 2.0]);
    }

    #[test]
    fn sum_reduce_bf16_along_one_dim() {
        let v: Vec<half::bf16> = [1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0]
            .iter()
            .map(|&x| half::bf16::from_f32(x))
            .collect();
        let input = CpuStorageBytes::from_slice(&v);
        let mut out = CpuStorageBytes::from_zero_bytes(2 * 2);
        sum_reduce_bf16(&input, &mut out, &[2, 3], &[1]).expect("sum bf16");
        let r: &[half::bf16] = out.as_slice().unwrap();
        let f32_out: Vec<f32> = r.iter().map(|x| x.to_f32()).collect();
        assert_eq!(f32_out, vec![6.0, 15.0]);
    }

    #[test]
    fn max_min_reduce_bf16() {
        let v: Vec<half::bf16> = [1.0_f32, -5.0, 3.0, 2.0, 0.0, -1.0]
            .iter()
            .map(|&x| half::bf16::from_f32(x))
            .collect();
        let input = CpuStorageBytes::from_slice(&v);
        let mut out = CpuStorageBytes::from_zero_bytes(2 * 2);
        max_reduce_bf16(&input, &mut out, &[2, 3], &[1]).expect("max bf16");
        let r: &[half::bf16] = out.as_slice().unwrap();
        let f32_out: Vec<f32> = r.iter().map(|x| x.to_f32()).collect();
        assert_eq!(f32_out, vec![3.0, 2.0]);

        min_reduce_bf16(&input, &mut out, &[2, 3], &[1]).expect("min bf16");
        let r: &[half::bf16] = out.as_slice().unwrap();
        let f32_out: Vec<f32> = r.iter().map(|x| x.to_f32()).collect();
        assert_eq!(f32_out, vec![-5.0, -1.0]);
    }

    #[test]
    fn mean_reduce_f16_divides_by_count() {
        let v: Vec<half::f16> = [2.0_f32, 4.0, 6.0, 8.0]
            .iter()
            .map(|&x| half::f16::from_f32(x))
            .collect();
        let input = CpuStorageBytes::from_slice(&v);
        let mut out = CpuStorageBytes::from_zero_bytes(2 * 2);
        mean_reduce_f16(&input, &mut out, &[2, 2], &[1]).expect("mean f16");
        let r: &[half::f16] = out.as_slice().unwrap();
        let f32_out: Vec<f32> = r.iter().map(|x| x.to_f32()).collect();
        assert_eq!(f32_out, vec![3.0, 7.0]);
    }

    #[test]
    fn matmul_bf16_2x3_times_3x2() {
        let lhs_v: Vec<half::bf16> = [1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0]
            .iter()
            .map(|&x| half::bf16::from_f32(x))
            .collect();
        let rhs_v: Vec<half::bf16> = [7.0_f32, 8.0, 9.0, 10.0, 11.0, 12.0]
            .iter()
            .map(|&x| half::bf16::from_f32(x))
            .collect();
        let lhs = CpuStorageBytes::from_slice(&lhs_v);
        let rhs = CpuStorageBytes::from_slice(&rhs_v);
        let mut out = CpuStorageBytes::from_zero_bytes(2 * 2 * 2);
        matmul_bf16(&lhs, &rhs, &mut out, &[], &[], 2, 2, 3).expect("matmul bf16");
        let r: &[half::bf16] = out.as_slice().unwrap();
        let f32_out: Vec<f32> = r.iter().map(|x| x.to_f32()).collect();
        // bf16 has ~3 decimal digits; the integer values 58, 64, 139,
        // 154 don't round-trip exactly through bf16 — accept ~1%
        // tolerance.
        let expected = [58.0_f32, 64.0, 139.0, 154.0];
        for (got, want) in f32_out.iter().zip(&expected) {
            assert!((got - want).abs() / want < 0.01, "got {got}, want {want}");
        }
    }

    #[test]
    fn matmul_f16_identity() {
        // Identity matmul on f16 — bytes round-trip exactly.
        let lhs_v: Vec<half::f16> = [1.0_f32, 2.0, 3.0, 4.0]
            .iter()
            .map(|&x| half::f16::from_f32(x))
            .collect();
        let rhs_v: Vec<half::f16> = [1.0_f32, 0.0, 0.0, 1.0]
            .iter()
            .map(|&x| half::f16::from_f32(x))
            .collect();
        let lhs = CpuStorageBytes::from_slice(&lhs_v);
        let rhs = CpuStorageBytes::from_slice(&rhs_v);
        let mut out = CpuStorageBytes::from_zero_bytes(2 * 2 * 2);
        matmul_f16(&lhs, &rhs, &mut out, &[], &[], 2, 2, 2).expect("matmul f16");
        let r: &[half::f16] = out.as_slice().unwrap();
        let f32_out: Vec<f32> = r.iter().map(|x| x.to_f32()).collect();
        assert_eq!(f32_out, vec![1.0, 2.0, 3.0, 4.0]);
    }

    #[test]
    fn matmul_bf16_batched_2x_2x2_times_2x2() {
        let lhs_v: Vec<half::bf16> = [1.0_f32, 2.0, 3.0, 4.0, 1.0, 0.0, 0.0, 1.0]
            .iter()
            .map(|&x| half::bf16::from_f32(x))
            .collect();
        let rhs_v: Vec<half::bf16> = [5.0_f32, 6.0, 7.0, 8.0, 10.0, 20.0, 30.0, 40.0]
            .iter()
            .map(|&x| half::bf16::from_f32(x))
            .collect();
        let lhs = CpuStorageBytes::from_slice(&lhs_v);
        let rhs = CpuStorageBytes::from_slice(&rhs_v);
        let mut out = CpuStorageBytes::from_zero_bytes(2 * 4 * 2);
        matmul_bf16(&lhs, &rhs, &mut out, &[2], &[2], 2, 2, 2).expect("matmul bf16");
        let r: &[half::bf16] = out.as_slice().unwrap();
        let f32_out: Vec<f32> = r.iter().map(|x| x.to_f32()).collect();
        // Expected (in f32 reference): [19, 22, 43, 50, 10, 20, 30, 40]
        let expected = [19.0_f32, 22.0, 43.0, 50.0, 10.0, 20.0, 30.0, 40.0];
        for (got, want) in f32_out.iter().zip(&expected) {
            // Small powers-of-two and integers up to 50 round-trip
            // through bf16 exactly.
            assert!((got - want).abs() < 0.5, "got {got}, want {want}");
        }
    }

    #[test]
    fn add_bf16_round_trips_through_f32() {
        let a_vec = vec![half::bf16::from_f32(1.0), half::bf16::from_f32(2.0)];
        let b_vec = vec![half::bf16::from_f32(10.0), half::bf16::from_f32(20.0)];
        let a = CpuStorageBytes::from_slice(&a_vec);
        let b = CpuStorageBytes::from_slice(&b_vec);
        let mut out = CpuStorageBytes::from_zero_bytes(2 * 2);
        add_bf16(&a, &b, &mut out).expect("add bf16");
        let result: &[half::bf16] = out.as_slice().unwrap();
        // bf16 has ~3 decimal digits; small integers round-trip exactly.
        assert_eq!(result[0].to_f32(), 11.0);
        assert_eq!(result[1].to_f32(), 22.0);
    }

    #[test]
    fn relu_bf16_clips_negatives() {
        let v: Vec<half::bf16> = [-1.0_f32, 0.0, 0.5, -3.5, 7.25]
            .iter()
            .map(|&x| half::bf16::from_f32(x))
            .collect();
        let input = CpuStorageBytes::from_slice(&v);
        let mut out = CpuStorageBytes::from_zero_bytes(input.len_bytes());
        relu_bf16(&input, &mut out).expect("relu bf16");
        let r: &[half::bf16] = out.as_slice().unwrap();
        let result_f32: Vec<f32> = r.iter().map(|x| x.to_f32()).collect();
        assert_eq!(result_f32, vec![0.0, 0.0, 0.5, 0.0, 7.25]);
    }

    /// Pinned NaN-semantics convention (2026-07-08): Relu is NaN-propagating
    /// (torch parity), across the bf16-via-f32 round-trip too.
    #[test]
    fn relu_bf16_propagates_nan() {
        let v: Vec<half::bf16> = [f32::NAN, -2.0, 3.0]
            .iter()
            .map(|&x| half::bf16::from_f32(x))
            .collect();
        let input = CpuStorageBytes::from_slice(&v);
        let mut out = CpuStorageBytes::from_zero_bytes(input.len_bytes());
        relu_bf16(&input, &mut out).expect("relu bf16");
        let r: &[half::bf16] = out.as_slice().unwrap();
        assert!(
            r[0].is_nan(),
            "relu(NaN) must propagate NaN, got {:?}",
            r[0]
        );
        assert_eq!(r[1].to_f32(), 0.0, "relu(-2) must clip to 0");
        assert_eq!(r[2].to_f32(), 3.0, "relu(3) must pass through");
    }

    #[test]
    fn exp_log_bf16_round_trip_within_precision() {
        // bf16 has ~3 decimal digits of mantissa precision. Test
        // values chosen so exp + log doesn't drift much.
        let v: Vec<half::bf16> = [1.0_f32, 2.0, 3.0]
            .iter()
            .map(|&x| half::bf16::from_f32(x))
            .collect();
        let input = CpuStorageBytes::from_slice(&v);
        let mut intermediate = CpuStorageBytes::from_zero_bytes(input.len_bytes());
        let mut out = CpuStorageBytes::from_zero_bytes(input.len_bytes());
        exp_bf16(&input, &mut intermediate).expect("exp");
        log_bf16(&intermediate, &mut out).expect("log");
        let r: &[half::bf16] = out.as_slice().unwrap();
        for (got, want) in r.iter().zip(&[1.0_f32, 2.0, 3.0]) {
            assert!(
                (got.to_f32() - want).abs() < 0.05,
                "bf16 exp+log lost too much: {} vs {}",
                got.to_f32(),
                want,
            );
        }
    }

    #[test]
    fn add_f16_round_trips_through_f32() {
        let a_vec = vec![half::f16::from_f32(1.0), half::f16::from_f32(2.5)];
        let b_vec = vec![half::f16::from_f32(0.5), half::f16::from_f32(-1.0)];
        let a = CpuStorageBytes::from_slice(&a_vec);
        let b = CpuStorageBytes::from_slice(&b_vec);
        let mut out = CpuStorageBytes::from_zero_bytes(2 * 2);
        add_f16(&a, &b, &mut out).expect("add f16");
        let result: &[half::f16] = out.as_slice().unwrap();
        assert_eq!(result[0].to_f32(), 1.5);
        assert_eq!(result[1].to_f32(), 1.5);
    }

    #[test]
    fn relu_f16_clips_negatives() {
        let v: Vec<half::f16> = [-2.0_f32, 0.0, 4.0, -0.5]
            .iter()
            .map(|&x| half::f16::from_f32(x))
            .collect();
        let input = CpuStorageBytes::from_slice(&v);
        let mut out = CpuStorageBytes::from_zero_bytes(input.len_bytes());
        relu_f16(&input, &mut out).expect("relu f16");
        let r: &[half::f16] = out.as_slice().unwrap();
        let result_f32: Vec<f32> = r.iter().map(|x| x.to_f32()).collect();
        assert_eq!(result_f32, vec![0.0, 0.0, 4.0, 0.0]);
    }

    #[test]
    fn sigmoid_bf16_at_zero_is_half() {
        let v = vec![half::bf16::from_f32(0.0_f32)];
        let input = CpuStorageBytes::from_slice(&v);
        let mut out = CpuStorageBytes::from_zero_bytes(input.len_bytes());
        sigmoid_bf16(&input, &mut out).expect("sigmoid bf16");
        let r: &[half::bf16] = out.as_slice().unwrap();
        assert!((r[0].to_f32() - 0.5).abs() < 0.01);
    }

    #[test]
    fn gelu_bf16_at_zero_is_zero() {
        let v = vec![half::bf16::from_f32(0.0_f32)];
        let input = CpuStorageBytes::from_slice(&v);
        let mut out = CpuStorageBytes::from_zero_bytes(input.len_bytes());
        gelu_bf16(&input, &mut out).expect("gelu_tanh bf16");
        let r: &[half::bf16] = out.as_slice().unwrap();
        assert!(r[0].to_f32().abs() < 0.001);
    }

    #[test]
    fn sum_reduce_f64_along_one_dim() {
        let input = CpuStorageBytes::from_slice(&[1.0_f64, 2.0, 3.0, 4.0, 5.0, 6.0]);
        let mut out = CpuStorageBytes::from_zero_bytes(2 * 8);
        sum_reduce_f64(&input, &mut out, &[2, 3], &[1]).expect("sum");
        assert_eq!(out.as_slice::<f64>().unwrap(), &[6.0_f64, 15.0]);
    }

    #[test]
    fn max_min_reduce_f64_basic() {
        let input = CpuStorageBytes::from_slice(&[1.0_f64, -5.0, 3.0, 2.0, 0.0, -1.0]);
        let mut out = CpuStorageBytes::from_zero_bytes(2 * 8);
        max_reduce_f64(&input, &mut out, &[2, 3], &[1]).expect("max");
        assert_eq!(out.as_slice::<f64>().unwrap(), &[3.0_f64, 2.0]);
        min_reduce_f64(&input, &mut out, &[2, 3], &[1]).expect("min");
        assert_eq!(out.as_slice::<f64>().unwrap(), &[-5.0_f64, -1.0]);
    }

    #[test]
    fn mean_reduce_f64_divides_by_count() {
        let input = CpuStorageBytes::from_slice(&[2.0_f64, 4.0, 6.0, 8.0]);
        let mut out = CpuStorageBytes::from_zero_bytes(2 * 8);
        mean_reduce_f64(&input, &mut out, &[2, 2], &[1]).expect("mean");
        assert_eq!(out.as_slice::<f64>().unwrap(), &[3.0_f64, 7.0]);
    }

    #[test]
    fn matmul_f64_2x3_times_3x2() {
        let lhs = CpuStorageBytes::from_slice(&[1.0_f64, 2.0, 3.0, 4.0, 5.0, 6.0]);
        let rhs = CpuStorageBytes::from_slice(&[7.0_f64, 8.0, 9.0, 10.0, 11.0, 12.0]);
        let mut out = CpuStorageBytes::from_zero_bytes(2 * 2 * 8);
        matmul_f64(&lhs, &rhs, &mut out, &[], &[], 2, 2, 3).expect("matmul");
        assert_eq!(
            out.as_slice::<f64>().unwrap(),
            &[58.0_f64, 64.0, 139.0, 154.0]
        );
    }

    #[test]
    fn matmul_f64_batched_2x_2x2_times_2x2() {
        let lhs = CpuStorageBytes::from_slice(&[1.0_f64, 2.0, 3.0, 4.0, 1.0, 0.0, 0.0, 1.0]);
        let rhs = CpuStorageBytes::from_slice(&[5.0_f64, 6.0, 7.0, 8.0, 10.0, 20.0, 30.0, 40.0]);
        let mut out = CpuStorageBytes::from_zero_bytes(2 * 4 * 8);
        matmul_f64(&lhs, &rhs, &mut out, &[2], &[2], 2, 2, 2).expect("matmul");
        assert_eq!(
            out.as_slice::<f64>().unwrap(),
            &[19.0_f64, 22.0, 43.0, 50.0, 10.0, 20.0, 30.0, 40.0]
        );
    }

    #[test]
    fn add_f64_basic() {
        let a = CpuStorageBytes::from_slice(&[1.0_f64, 2.0, 3.0]);
        let b = CpuStorageBytes::from_slice(&[10.0_f64, 20.0, 30.0]);
        let mut out = CpuStorageBytes::from_zero_bytes(3 * 8);
        add_f64(&a, &b, &mut out).expect("add f64");
        assert_eq!(out.as_slice::<f64>().unwrap(), &[11.0, 22.0, 33.0]);
    }

    #[test]
    fn relu_f64_clips_negatives() {
        let input = CpuStorageBytes::from_slice(&[-1.0_f64, 0.0, 0.5, -3.5, 7.25]);
        let mut out = CpuStorageBytes::from_zero_bytes(input.len_bytes());
        relu_f64(&input, &mut out).expect("relu f64");
        assert_eq!(out.as_slice::<f64>().unwrap(), &[0.0, 0.0, 0.5, 0.0, 7.25]);
    }

    #[test]
    fn unary_f64_round_trip_sample() {
        // Smoke-test several f64 unaries on one input.
        let input = CpuStorageBytes::from_slice(&[1.0_f64, -2.0, 4.0]);
        let mut out = CpuStorageBytes::from_zero_bytes(input.len_bytes());

        neg_f64(&input, &mut out).expect("neg f64");
        assert_eq!(out.as_slice::<f64>().unwrap(), &[-1.0, 2.0, -4.0]);

        sqr_f64(&input, &mut out).expect("sqr f64");
        assert_eq!(out.as_slice::<f64>().unwrap(), &[1.0, 4.0, 16.0]);

        let pos = CpuStorageBytes::from_slice(&[1.0_f64, 4.0, 9.0]);
        let mut out2 = CpuStorageBytes::from_zero_bytes(pos.len_bytes());
        sqrt_f64(&pos, &mut out2).expect("sqrt f64");
        assert_eq!(out2.as_slice::<f64>().unwrap(), &[1.0, 2.0, 3.0]);
    }

    #[test]
    fn maximum_minimum_f64_basic() {
        let lhs = CpuStorageBytes::from_slice(&[1.0_f64, 5.0, -3.0]);
        let rhs = CpuStorageBytes::from_slice(&[2.0_f64, 1.0, -1.0]);
        let mut out = CpuStorageBytes::from_zero_bytes(lhs.len_bytes());
        maximum_f64(&lhs, &rhs, &mut out).expect("max f64");
        assert_eq!(out.as_slice::<f64>().unwrap(), &[2.0, 5.0, -1.0]);
        minimum_f64(&lhs, &rhs, &mut out).expect("min f64");
        assert_eq!(out.as_slice::<f64>().unwrap(), &[1.0, 1.0, -3.0]);
    }

    #[test]
    fn gelu_f64_at_known_points() {
        let input = CpuStorageBytes::from_slice(&[0.0_f64, 1.0, -1.0]);
        let mut out = CpuStorageBytes::from_zero_bytes(input.len_bytes());
        gelu_f64(&input, &mut out).expect("gelu_tanh f64");
        let r: &[f64] = out.as_slice().unwrap();
        // gelu_tanh(0) = 0
        assert!(r[0].abs() < 1e-12);
        // gelu_tanh(1) ≈ 0.8412 (tanh approx); f64 should be ~12 digits accurate
        assert!((r[1] - 0.841_192).abs() < 1e-3);
        assert!((r[2] - (-0.158_808)).abs() < 1e-3);
    }

    #[test]
    fn argmax_dim_f32_basic() {
        // input [2, 3]: row 0 = [1, 5, 2], row 1 = [9, 0, 4]
        // argmax along dim=1 → output [2] with [1, 0]
        let input = CpuStorageBytes::from_slice(&[1.0_f32, 5.0, 2.0, 9.0, 0.0, 4.0]);
        let mut out = CpuStorageBytes::from_zero_bytes(2 * 4);
        argmax_dim_f32(&input, &mut out, &[2, 3], 1).expect("argmax");
        assert_eq!(out.as_slice::<u32>().unwrap(), &[1u32, 0]);
    }

    #[test]
    fn argmin_dim_f32_basic() {
        let input = CpuStorageBytes::from_slice(&[1.0_f32, 5.0, 2.0, 9.0, 0.0, 4.0]);
        let mut out = CpuStorageBytes::from_zero_bytes(2 * 4);
        argmin_dim_f32(&input, &mut out, &[2, 3], 1).expect("argmin");
        // Row 0 min at index 0 (= 1.0); row 1 min at index 1 (= 0.0)
        assert_eq!(out.as_slice::<u32>().unwrap(), &[0u32, 1]);
    }

    #[test]
    fn argmax_dim_f32_outer_dim() {
        // input [3, 2]:
        //   [[1, 4],
        //    [3, 2],
        //    [0, 9]]
        // argmax along dim=0 → output [2]
        // col 0: max at index 1 (= 3); col 1: max at index 2 (= 9)
        let input = CpuStorageBytes::from_slice(&[1.0_f32, 4.0, 3.0, 2.0, 0.0, 9.0]);
        let mut out = CpuStorageBytes::from_zero_bytes(2 * 4);
        argmax_dim_f32(&input, &mut out, &[3, 2], 0).expect("argmax");
        assert_eq!(out.as_slice::<u32>().unwrap(), &[1u32, 2]);
    }

    #[test]
    fn argmax_dim_f32_first_index_on_tie() {
        // All equal — argmax must return index 0 (first occurrence).
        let input = CpuStorageBytes::from_slice(&[5.0_f32, 5.0, 5.0, 5.0]);
        let mut out = CpuStorageBytes::from_zero_bytes(4);
        argmax_dim_f32(&input, &mut out, &[4], 0).expect("argmax");
        assert_eq!(out.as_slice::<u32>().unwrap(), &[0u32]);
    }

    #[test]
    fn argmax_dim_f32_rejects_zero_dim() {
        let input = CpuStorageBytes::from_zero_bytes(0);
        let mut out = CpuStorageBytes::from_zero_bytes(0);
        let r = argmax_dim_f32(&input, &mut out, &[0], 0);
        assert!(r.is_err());
    }

    #[test]
    fn index_add_f32_simple_accumulate() {
        // base [3] = [10, 20, 30]; indices [2] = [0, 0]; src [2] = [1, 2]
        // → out = [10 + 1 + 2, 20, 30] = [13, 20, 30]
        let base = CpuStorageBytes::from_slice(&[10.0_f32, 20.0, 30.0]);
        let indices = CpuStorageBytes::from_slice(&[0u32, 0]);
        let src = CpuStorageBytes::from_slice(&[1.0_f32, 2.0]);
        let mut out = CpuStorageBytes::from_zero_bytes(3 * 4);
        index_add_f32(&base, &indices, &src, &mut out, 1, 3, 2, 1).expect("index_add");
        assert_eq!(out.as_slice::<f32>().unwrap(), &[13.0, 20.0, 30.0]);
    }

    #[test]
    fn index_add_f32_along_inner_dim() {
        // base [2, 4]; indices [3] = [1, 3, 1]; src [2, 3].
        // For each row, accumulate src cols at base cols indices.
        let base =
            CpuStorageBytes::from_slice(&[10.0_f32, 20.0, 30.0, 40.0, 50.0, 60.0, 70.0, 80.0]);
        let indices = CpuStorageBytes::from_slice(&[1u32, 3, 1]);
        let src = CpuStorageBytes::from_slice(&[1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0]);
        let mut out = CpuStorageBytes::from_zero_bytes(2 * 4 * 4);
        // outer=2, base_dim=4, n_indices=3, inner=1
        index_add_f32(&base, &indices, &src, &mut out, 2, 4, 3, 1).expect("index_add");
        // row 0: base [10, 20, 30, 40]. src [1, 2, 3] at indices [1, 3, 1]:
        //   col 0: 10
        //   col 1: 20 + 1 + 3 = 24
        //   col 2: 30
        //   col 3: 40 + 2 = 42
        // row 1: base [50, 60, 70, 80]. src [4, 5, 6] at indices [1, 3, 1]:
        //   col 0: 50
        //   col 1: 60 + 4 + 6 = 70
        //   col 2: 70
        //   col 3: 80 + 5 = 85
        assert_eq!(
            out.as_slice::<f32>().unwrap(),
            &[10.0, 24.0, 30.0, 42.0, 50.0, 70.0, 70.0, 85.0]
        );
    }

    #[test]
    fn index_add_f32_rejects_out_of_bounds_index() {
        let base = CpuStorageBytes::from_slice(&[1.0_f32, 2.0, 3.0]);
        let indices = CpuStorageBytes::from_slice(&[5u32]);
        let src = CpuStorageBytes::from_slice(&[10.0_f32]);
        let mut out = CpuStorageBytes::from_zero_bytes(3 * 4);
        let r = index_add_f32(&base, &indices, &src, &mut out, 1, 3, 1, 1);
        assert!(r.is_err());
    }

    #[test]
    fn scatter_add_f32_along_outer_dim() {
        // base [3, 2] = zeros; indices [2, 2]:
        //   [[0, 1],
        //    [2, 0]]
        // src [2, 2] = [[1, 2], [3, 4]]
        // dim=0:
        //   src[0, 0]=1 → out[indices[0,0]=0, 0] += 1 → out[0, 0] += 1
        //   src[0, 1]=2 → out[indices[0,1]=1, 1] += 2 → out[1, 1] += 2
        //   src[1, 0]=3 → out[indices[1,0]=2, 0] += 3 → out[2, 0] += 3
        //   src[1, 1]=4 → out[indices[1,1]=0, 1] += 4 → out[0, 1] += 4
        // → [[1, 4], [0, 2], [3, 0]]
        let base = CpuStorageBytes::from_slice(&[0.0_f32; 6]);
        let indices = CpuStorageBytes::from_slice(&[0u32, 1, 2, 0]);
        let src = CpuStorageBytes::from_slice(&[1.0_f32, 2.0, 3.0, 4.0]);
        let mut out = CpuStorageBytes::from_zero_bytes(3 * 2 * 4);
        scatter_add_f32(&base, &indices, &src, &mut out, &[3, 2], &[2, 2], 0).expect("scatter_add");
        assert_eq!(
            out.as_slice::<f32>().unwrap(),
            &[1.0, 4.0, 0.0, 2.0, 3.0, 0.0]
        );
    }

    #[test]
    fn scatter_add_f32_starts_from_base() {
        // Same as above but base is already nonzero — verifies that
        // out copies base before accumulating.
        let base = CpuStorageBytes::from_slice(&[100.0_f32, 200.0, 300.0]);
        let indices = CpuStorageBytes::from_slice(&[0u32, 0, 2]);
        let src = CpuStorageBytes::from_slice(&[1.0_f32, 2.0, 3.0]);
        let mut out = CpuStorageBytes::from_zero_bytes(3 * 4);
        scatter_add_f32(&base, &indices, &src, &mut out, &[3], &[3], 0).expect("scatter_add");
        // out[0] = 100 + 1 + 2 = 103
        // out[1] = 200 (untouched)
        // out[2] = 300 + 3 = 303
        assert_eq!(out.as_slice::<f32>().unwrap(), &[103.0, 200.0, 303.0]);
    }

    #[test]
    fn scatter_add_f32_rejects_shape_mismatch() {
        let base = CpuStorageBytes::from_slice(&[0.0_f32; 4]);
        let indices = CpuStorageBytes::from_slice(&[0u32, 1]);
        let src = CpuStorageBytes::from_slice(&[1.0_f32, 2.0]);
        let mut out = CpuStorageBytes::from_zero_bytes(4 * 4);
        // base shape claims [2, 2] but src claims [3] — rank mismatch.
        let r = scatter_add_f32(&base, &indices, &src, &mut out, &[2, 2], &[3], 0);
        assert!(r.is_err());
    }

    #[test]
    fn rope_f32_identity_when_cos_one_sin_zero() {
        // cos=1, sin=0 everywhere → out == x.
        let x = CpuStorageBytes::from_slice(&[1.0_f32, 2.0, 3.0, 4.0]);
        let cos = CpuStorageBytes::from_slice(&[1.0_f32, 1.0, 1.0, 1.0]);
        let sin = CpuStorageBytes::from_slice(&[0.0_f32, 0.0, 0.0, 0.0]);
        let mut out = CpuStorageBytes::from_zero_bytes(4 * 4);
        // outer=1, seq=1, head_dim=4, h=2
        rope_f32(&x, &cos, &sin, &mut out, 1, 1, 4).expect("rope");
        assert_eq!(out.as_slice::<f32>().unwrap(), &[1.0, 2.0, 3.0, 4.0]);
    }

    #[test]
    fn rope_f32_pi_over_two_swaps_with_sign() {
        // cos=0, sin=1 (i.e. θ = π/2 everywhere). Then:
        //   out[i]   = x[i]*0 - x[i+h]*1 = -x[i+h]
        //   out[i+h] = x[i+h]*0 + x[i]*1 = x[i]
        // So [a, b, c, d] (head_dim=4, h=2) → [-c, -d, a, b].
        let x = CpuStorageBytes::from_slice(&[1.0_f32, 2.0, 3.0, 4.0]);
        let cos = CpuStorageBytes::from_slice(&[0.0_f32, 0.0, 0.0, 0.0]);
        let sin = CpuStorageBytes::from_slice(&[1.0_f32, 1.0, 1.0, 1.0]);
        let mut out = CpuStorageBytes::from_zero_bytes(4 * 4);
        rope_f32(&x, &cos, &sin, &mut out, 1, 1, 4).expect("rope");
        assert_eq!(out.as_slice::<f32>().unwrap(), &[-3.0, -4.0, 1.0, 2.0]);
    }

    #[test]
    fn rope_f32_broadcasts_cos_sin_over_outer() {
        // outer=2, seq=1, head_dim=2, h=1. Shared cos=[0, 0], sin=[1, 1].
        // x outer 0: [1, 2] → [-2, 1]
        // x outer 1: [10, 20] → [-20, 10]
        let x = CpuStorageBytes::from_slice(&[1.0_f32, 2.0, 10.0, 20.0]);
        let cos = CpuStorageBytes::from_slice(&[0.0_f32, 0.0]);
        let sin = CpuStorageBytes::from_slice(&[1.0_f32, 1.0]);
        let mut out = CpuStorageBytes::from_zero_bytes(4 * 4);
        rope_f32(&x, &cos, &sin, &mut out, 2, 1, 2).expect("rope");
        assert_eq!(out.as_slice::<f32>().unwrap(), &[-2.0, 1.0, -20.0, 10.0]);
    }

    #[test]
    fn rope_f32_rejects_odd_head_dim() {
        let x = CpuStorageBytes::from_slice(&[1.0_f32, 2.0, 3.0]);
        let cos = CpuStorageBytes::from_slice(&[1.0_f32, 1.0, 1.0]);
        let sin = CpuStorageBytes::from_slice(&[0.0_f32, 0.0, 0.0]);
        let mut out = CpuStorageBytes::from_zero_bytes(3 * 4);
        // head_dim=3 is odd
        let r = rope_f32(&x, &cos, &sin, &mut out, 1, 1, 3);
        assert!(r.is_err());
    }

    #[test]
    fn gather_f32_along_inner_dim() {
        // source [2, 4]:
        //   row 0: [10, 20, 30, 40]
        //   row 1: [50, 60, 70, 80]
        // indices [2, 3]:
        //   row 0: [0, 2, 1]
        //   row 1: [3, 0, 0]
        // gather dim=1:
        //   out[0, j] = source[0, indices[0, j]]
        //   out[1, j] = source[1, indices[1, j]]
        let source =
            CpuStorageBytes::from_slice(&[10.0_f32, 20.0, 30.0, 40.0, 50.0, 60.0, 70.0, 80.0]);
        let indices = CpuStorageBytes::from_slice(&[0u32, 2, 1, 3, 0, 0]);
        let mut out = CpuStorageBytes::from_zero_bytes(2 * 3 * 4);
        gather_f32(&source, &indices, &mut out, &[2, 4], &[2, 3], 1).expect("gather");
        assert_eq!(
            out.as_slice::<f32>().unwrap(),
            &[10.0, 30.0, 20.0, 80.0, 50.0, 50.0]
        );
    }

    #[test]
    fn gather_f32_along_outer_dim() {
        // source [3, 2]:
        //   row 0: [1, 2]
        //   row 1: [3, 4]
        //   row 2: [5, 6]
        // indices [4, 2]:
        //   [[0, 1],
        //    [1, 0],
        //    [2, 2],
        //    [0, 0]]
        // gather dim=0:
        //   out[i, j] = source[indices[i, j], j]
        let source = CpuStorageBytes::from_slice(&[1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0]);
        let indices = CpuStorageBytes::from_slice(&[0u32, 1, 1, 0, 2, 2, 0, 0]);
        let mut out = CpuStorageBytes::from_zero_bytes(4 * 2 * 4);
        gather_f32(&source, &indices, &mut out, &[3, 2], &[4, 2], 0).expect("gather");
        // row 0: source[0,0]=1, source[1,1]=4 → [1, 4]
        // row 1: source[1,0]=3, source[0,1]=2 → [3, 2]
        // row 2: source[2,0]=5, source[2,1]=6 → [5, 6]
        // row 3: source[0,0]=1, source[0,1]=2 → [1, 2]
        assert_eq!(
            out.as_slice::<f32>().unwrap(),
            &[1.0, 4.0, 3.0, 2.0, 5.0, 6.0, 1.0, 2.0]
        );
    }

    #[test]
    fn gather_f32_rejects_out_of_bounds_index() {
        let source = CpuStorageBytes::from_slice(&[1.0_f32, 2.0, 3.0]);
        let indices = CpuStorageBytes::from_slice(&[0u32, 5]); // 5 OOB for dim 3
        let mut out = CpuStorageBytes::from_zero_bytes(2 * 4);
        let r = gather_f32(&source, &indices, &mut out, &[3], &[2], 0);
        assert!(r.is_err());
    }

    #[test]
    fn gather_f32_rejects_rank_mismatch() {
        let source = CpuStorageBytes::from_slice(&[1.0_f32, 2.0, 3.0]);
        let indices = CpuStorageBytes::from_slice(&[0u32, 1]);
        let mut out = CpuStorageBytes::from_zero_bytes(2 * 4);
        // source rank=1, output rank=2 → must error
        let r = gather_f32(&source, &indices, &mut out, &[3], &[1, 2], 0);
        assert!(r.is_err());
    }

    #[test]
    fn index_select_f32_embedding_lookup() {
        // Embedding table [4, 3]: rows 0..3 are different.
        // Indices [u32; 3] = [2, 0, 2]; output [3, 3] picks
        // rows 2, 0, 2 in that order.
        let table = CpuStorageBytes::from_slice(&[
            10.0_f32, 11.0, 12.0, // row 0
            20.0, 21.0, 22.0, // row 1
            30.0, 31.0, 32.0, // row 2
            40.0, 41.0, 42.0, // row 3
        ]);
        let indices = CpuStorageBytes::from_slice(&[2u32, 0, 2]);
        let mut out = CpuStorageBytes::from_zero_bytes(3 * 3 * 4);
        index_select_f32(&table, &indices, &mut out, 1, 4, 3, 3).expect("index_select");
        assert_eq!(
            out.as_slice::<f32>().unwrap(),
            &[30.0, 31.0, 32.0, 10.0, 11.0, 12.0, 30.0, 31.0, 32.0,]
        );
    }

    #[test]
    fn index_select_f32_inner_dim() {
        // source [2, 3]: outer=2, dim_size=3, inner=1.
        // Pick along dim=1 with indices [2, 0]: output [2, 2].
        let source = CpuStorageBytes::from_slice(&[1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0]);
        let indices = CpuStorageBytes::from_slice(&[2u32, 0]);
        let mut out = CpuStorageBytes::from_zero_bytes(2 * 2 * 4);
        index_select_f32(&source, &indices, &mut out, 2, 3, 2, 1).expect("index_select");
        // outer 0: row [1, 2, 3], pick (2, 0) → [3, 1]
        // outer 1: row [4, 5, 6], pick (2, 0) → [6, 4]
        assert_eq!(out.as_slice::<f32>().unwrap(), &[3.0, 1.0, 6.0, 4.0]);
    }

    #[test]
    fn index_select_f32_rejects_out_of_bounds_index() {
        let source = CpuStorageBytes::from_slice(&[1.0_f32, 2.0, 3.0]);
        let indices = CpuStorageBytes::from_slice(&[5u32]); // out of bounds for dim 3
        let mut out = CpuStorageBytes::from_zero_bytes(4);
        let r = index_select_f32(&source, &indices, &mut out, 1, 3, 1, 1);
        assert!(r.is_err());
    }

    #[test]
    fn rms_norm_last_dim_f32_basic() {
        // Input [3, 4]: rms = sqrt(mean(9, 16)) = sqrt(12.5)
        // out = [3 / sqrt(12.5), 4 / sqrt(12.5)]
        let input = CpuStorageBytes::from_slice(&[3.0_f32, 4.0]);
        let mut out = CpuStorageBytes::from_zero_bytes(2 * 4);
        rms_norm_last_dim_f32(&input, &mut out, 1, 2, 0.0).expect("rms_norm");
        let result: &[f32] = out.as_slice().unwrap();
        let rms = (12.5_f32).sqrt();
        assert!((result[0] - 3.0 / rms).abs() < 1e-6);
        assert!((result[1] - 4.0 / rms).abs() < 1e-6);
    }

    #[test]
    fn rms_norm_last_dim_f32_eps_prevents_zero_division() {
        // All-zero row → mean_sq = 0; without eps, would divide by 0.
        let input = CpuStorageBytes::from_slice(&[0.0_f32, 0.0, 0.0]);
        let mut out = CpuStorageBytes::from_zero_bytes(3 * 4);
        rms_norm_last_dim_f32(&input, &mut out, 1, 3, 1e-5).expect("rms_norm");
        let result: &[f32] = out.as_slice().unwrap();
        for v in result {
            assert!(v.is_finite() && *v == 0.0);
        }
    }

    #[test]
    fn layer_norm_last_dim_f32_zero_mean_unit_var() {
        // Input [1, 2, 3]: mean = 2; var = (1 + 0 + 1) / 3 = 2/3
        // inv_std = 1/sqrt(2/3 + eps)
        let input = CpuStorageBytes::from_slice(&[1.0_f32, 2.0, 3.0]);
        let mut out = CpuStorageBytes::from_zero_bytes(3 * 4);
        layer_norm_last_dim_f32(&input, &mut out, 1, 3, 0.0).expect("layer_norm");
        let result: &[f32] = out.as_slice().unwrap();
        // Output mean should be 0; output stddev should be 1.
        let out_sum: f32 = result.iter().sum();
        assert!(out_sum.abs() < 1e-6);
        let out_mean = out_sum / 3.0;
        let out_var: f32 = result.iter().map(|v| (v - out_mean).powi(2)).sum::<f32>() / 3.0;
        assert!((out_var - 1.0).abs() < 1e-6);
    }

    #[test]
    fn layer_norm_last_dim_f32_two_rows_independent() {
        // Two rows of 3; each row's output should have mean ≈ 0.
        let input = CpuStorageBytes::from_slice(&[1.0_f32, 2.0, 3.0, 10.0, 20.0, 30.0]);
        let mut out = CpuStorageBytes::from_zero_bytes(6 * 4);
        layer_norm_last_dim_f32(&input, &mut out, 2, 3, 0.0).expect("layer_norm");
        let result: &[f32] = out.as_slice().unwrap();
        let r0_mean: f32 = result[..3].iter().sum::<f32>() / 3.0;
        let r1_mean: f32 = result[3..].iter().sum::<f32>() / 3.0;
        assert!(r0_mean.abs() < 1e-6);
        assert!(r1_mean.abs() < 1e-6);
    }

    #[test]
    fn softmax_last_dim_f32_uniform_row() {
        // [1, 1, 1, 1] → uniform softmax = 0.25 each
        let input = CpuStorageBytes::from_slice(&[1.0_f32, 1.0, 1.0, 1.0]);
        let mut out = CpuStorageBytes::from_zero_bytes(4 * 4);
        softmax_last_dim_f32(&input, &mut out, 1, 4).expect("softmax");
        let result: &[f32] = out.as_slice().unwrap();
        for v in result {
            assert!((v - 0.25).abs() < 1e-7);
        }
    }

    #[test]
    fn softmax_last_dim_f32_sums_to_one_per_row() {
        // Two rows of 3; arbitrary values; each row should sum to 1.
        let input = CpuStorageBytes::from_slice(&[1.0_f32, 2.0, 3.0, -1.0, 0.0, 1.0]);
        let mut out = CpuStorageBytes::from_zero_bytes(6 * 4);
        softmax_last_dim_f32(&input, &mut out, 2, 3).expect("softmax");
        let result: &[f32] = out.as_slice().unwrap();
        let row0_sum: f32 = result[..3].iter().sum();
        let row1_sum: f32 = result[3..].iter().sum();
        assert!((row0_sum - 1.0).abs() < 1e-6);
        assert!((row1_sum - 1.0).abs() < 1e-6);
    }

    #[test]
    fn softmax_last_dim_f32_numerical_stability_at_large_values() {
        // [1000, 1001, 1002] — without max subtraction, exp would
        // overflow. Stable softmax should still produce finite output.
        let input = CpuStorageBytes::from_slice(&[1000.0_f32, 1001.0, 1002.0]);
        let mut out = CpuStorageBytes::from_zero_bytes(3 * 4);
        softmax_last_dim_f32(&input, &mut out, 1, 3).expect("softmax");
        let result: &[f32] = out.as_slice().unwrap();
        for v in result {
            assert!(v.is_finite(), "softmax must not overflow at large inputs");
        }
        let sum: f32 = result.iter().sum();
        assert!((sum - 1.0).abs() < 1e-6);
    }

    #[test]
    fn softmax_last_dim_f32_known_values() {
        // [0, 1] → softmax = [1/(1+e), e/(1+e)]
        let input = CpuStorageBytes::from_slice(&[0.0_f32, 1.0]);
        let mut out = CpuStorageBytes::from_zero_bytes(2 * 4);
        softmax_last_dim_f32(&input, &mut out, 1, 2).expect("softmax");
        let result: &[f32] = out.as_slice().unwrap();
        let e = std::f32::consts::E;
        let expected = [1.0 / (1.0 + e), e / (1.0 + e)];
        for (got, want) in result.iter().zip(&expected) {
            assert!((got - want).abs() < 1e-6, "got {got}, want {want}");
        }
    }

    #[test]
    fn softmax_last_dim_f32_size_mismatch_errors() {
        let input = CpuStorageBytes::from_slice(&[1.0_f32, 2.0, 3.0]); // 3 elements
        let mut out = CpuStorageBytes::from_zero_bytes(3 * 4);
        // Claim 2 rows of 4 = 8 elements but only 3 — must error.
        let r = softmax_last_dim_f32(&input, &mut out, 2, 4);
        assert!(r.is_err());
    }

    #[test]
    fn concat_f32_along_inner_dim() {
        // Two [2, 3] tensors concatenated along dim 1 → [2, 6].
        // a = [[1, 2, 3], [4, 5, 6]]
        // b = [[7, 8, 9], [10, 11, 12]]
        // concat dim 1 = [[1, 2, 3, 7, 8, 9], [4, 5, 6, 10, 11, 12]]
        let a = CpuStorageBytes::from_slice(&[1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0]);
        let b = CpuStorageBytes::from_slice(&[7.0_f32, 8.0, 9.0, 10.0, 11.0, 12.0]);
        let mut out = CpuStorageBytes::from_zero_bytes(12 * 4);

        // dim=1: outer=2 (rows), inner=1 (1D inside each cell)
        concat_f32(&[&a, &b], &mut out, 2, &[3, 3], 1).expect("concat");
        assert_eq!(
            out.as_slice::<f32>().unwrap(),
            &[
                1.0, 2.0, 3.0, 7.0, 8.0, 9.0, 4.0, 5.0, 6.0, 10.0, 11.0, 12.0
            ]
        );
    }

    #[test]
    fn concat_f32_along_outer_dim() {
        // Two [2, 3] tensors concatenated along dim 0 → [4, 3].
        // For dim=0: outer_count=1 (no dims before), inner_count=3.
        // Input dim sizes are 2 and 2.
        // Output: [a-row-0, a-row-1, b-row-0, b-row-1].
        let a = CpuStorageBytes::from_slice(&[1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0]);
        let b = CpuStorageBytes::from_slice(&[7.0_f32, 8.0, 9.0, 10.0, 11.0, 12.0]);
        let mut out = CpuStorageBytes::from_zero_bytes(12 * 4);

        concat_f32(&[&a, &b], &mut out, 1, &[2, 2], 3).expect("concat");
        assert_eq!(
            out.as_slice::<f32>().unwrap(),
            &[
                1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0
            ]
        );
    }

    #[test]
    fn concat_f32_three_inputs() {
        // Three [2] tensors concatenated → [6].
        let a = CpuStorageBytes::from_slice(&[1.0_f32, 2.0]);
        let b = CpuStorageBytes::from_slice(&[3.0_f32, 4.0]);
        let c = CpuStorageBytes::from_slice(&[5.0_f32, 6.0]);
        let mut out = CpuStorageBytes::from_zero_bytes(6 * 4);
        concat_f32(&[&a, &b, &c], &mut out, 1, &[2, 2, 2], 1).expect("concat");
        assert_eq!(
            out.as_slice::<f32>().unwrap(),
            &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]
        );
    }

    #[test]
    fn concat_f32_size_mismatch_errors() {
        let a = CpuStorageBytes::from_slice(&[1.0_f32, 2.0]);
        let b = CpuStorageBytes::from_slice(&[3.0_f32, 4.0]);
        let mut out = CpuStorageBytes::from_zero_bytes(8); // expects 4 elements but only 8 bytes
        let r = concat_f32(&[&a, &b], &mut out, 1, &[3, 2], 1); // claims a has 3 but it has 2
        assert!(r.is_err());
    }

    #[test]
    fn affine_f32_basic() {
        let input = CpuStorageBytes::from_slice(&[1.0_f32, 2.0, -3.0]);
        let mut out = CpuStorageBytes::from_zero_bytes(input.len_bytes());
        // y = 2 * x + 1
        affine_f32(&input, &mut out, 2.0, 1.0).expect("affine");
        assert_eq!(out.as_slice::<f32>().unwrap(), &[3.0, 5.0, -5.0]);
    }

    #[test]
    fn affine_f32_handles_addscalar_and_mulscalar() {
        let input = CpuStorageBytes::from_slice(&[1.0_f32, 2.0, 3.0]);
        let mut out = CpuStorageBytes::from_zero_bytes(input.len_bytes());

        // AddScalar(10): mul=1, add=10
        affine_f32(&input, &mut out, 1.0, 10.0).expect("addscalar");
        assert_eq!(out.as_slice::<f32>().unwrap(), &[11.0, 12.0, 13.0]);

        // MulScalar(3): mul=3, add=0
        affine_f32(&input, &mut out, 3.0, 0.0).expect("mulscalar");
        assert_eq!(out.as_slice::<f32>().unwrap(), &[3.0, 6.0, 9.0]);
    }

    #[test]
    fn clamp_f32_basic() {
        let input = CpuStorageBytes::from_slice(&[-5.0_f32, -1.0, 0.5, 3.0, 100.0]);
        let mut out = CpuStorageBytes::from_zero_bytes(input.len_bytes());
        clamp_f32(&input, &mut out, -2.0, 2.0).expect("clamp");
        assert_eq!(out.as_slice::<f32>().unwrap(), &[-2.0, -1.0, 0.5, 2.0, 2.0]);
    }

    #[test]
    fn clamp_f32_rejects_inverted_bounds() {
        let input = CpuStorageBytes::from_slice(&[1.0_f32]);
        let mut out = CpuStorageBytes::from_zero_bytes(input.len_bytes());
        let r = clamp_f32(&input, &mut out, 5.0, 1.0);
        assert!(r.is_err());
    }

    #[test]
    fn powi_f32_basic() {
        let input = CpuStorageBytes::from_slice(&[1.0_f32, 2.0, 3.0, 4.0]);
        let mut out = CpuStorageBytes::from_zero_bytes(input.len_bytes());
        // exp = 3
        powi_f32(&input, &mut out, 3).expect("powi");
        assert_eq!(out.as_slice::<f32>().unwrap(), &[1.0, 8.0, 27.0, 64.0]);
        // exp = 0 → 1.0 everywhere
        powi_f32(&input, &mut out, 0).expect("powi");
        assert_eq!(out.as_slice::<f32>().unwrap(), &[1.0, 1.0, 1.0, 1.0]);
        // exp = -1 → reciprocal
        let pos = CpuStorageBytes::from_slice(&[2.0_f32, 4.0, 8.0]);
        let mut out2 = CpuStorageBytes::from_zero_bytes(pos.len_bytes());
        powi_f32(&pos, &mut out2, -1).expect("powi");
        assert_eq!(out2.as_slice::<f32>().unwrap(), &[0.5, 0.25, 0.125]);
    }

    #[test]
    fn maximum_minimum_f32_basic() {
        let lhs = CpuStorageBytes::from_slice(&[1.0_f32, 5.0, -3.0]);
        let rhs = CpuStorageBytes::from_slice(&[2.0_f32, 1.0, -1.0]);
        let mut out = CpuStorageBytes::from_zero_bytes(lhs.len_bytes());

        maximum_f32(&lhs, &rhs, &mut out).expect("max");
        assert_eq!(out.as_slice::<f32>().unwrap(), &[2.0, 5.0, -1.0]);

        minimum_f32(&lhs, &rhs, &mut out).expect("min");
        assert_eq!(out.as_slice::<f32>().unwrap(), &[1.0, 1.0, -3.0]);
    }

    /// Pinned NaN-semantics convention (2026-07-08, decisions log): Maximum
    /// and Minimum are NaN-propagating (torch parity — `torch.maximum` /
    /// `torch.minimum` return NaN if *either* operand is NaN), not the prior
    /// `f32::max`/`f32::min` "NaN-as-missing" (non-NaN-operand-wins)
    /// behavior. Covers NaN in lhs, in rhs, in both, plus non-NaN sanity.
    #[test]
    fn maximum_minimum_f32_propagates_nan() {
        let nan = f32::NAN;
        let lhs = CpuStorageBytes::from_slice(&[nan, 2.0, nan, 1.0]);
        let rhs = CpuStorageBytes::from_slice(&[3.0, nan, nan, 2.0]);
        let mut out = CpuStorageBytes::from_zero_bytes(lhs.len_bytes());

        maximum_f32(&lhs, &rhs, &mut out).expect("max");
        let r: &[f32] = out.as_slice().unwrap();
        assert!(r[0].is_nan(), "maximum(NaN, 3) must be NaN, got {}", r[0]);
        assert!(r[1].is_nan(), "maximum(2, NaN) must be NaN, got {}", r[1]);
        assert!(r[2].is_nan(), "maximum(NaN, NaN) must be NaN, got {}", r[2]);
        assert_eq!(r[3], 2.0, "maximum(1, 2) non-NaN sanity");

        minimum_f32(&lhs, &rhs, &mut out).expect("min");
        let r: &[f32] = out.as_slice().unwrap();
        assert!(r[0].is_nan(), "minimum(NaN, 3) must be NaN, got {}", r[0]);
        assert!(r[1].is_nan(), "minimum(2, NaN) must be NaN, got {}", r[1]);
        assert!(r[2].is_nan(), "minimum(NaN, NaN) must be NaN, got {}", r[2]);
        assert_eq!(r[3], 1.0, "minimum(1, 2) non-NaN sanity");
    }

    /// bf16 mirror of `maximum_minimum_f32_propagates_nan` — the
    /// widen-to-f32 round-trip must not scrub NaN on narrow-back.
    #[test]
    fn maximum_minimum_bf16_propagates_nan() {
        let nan = half::bf16::from_f32(f32::NAN);
        let two = half::bf16::from_f32(2.0);
        let three = half::bf16::from_f32(3.0);
        let one = half::bf16::from_f32(1.0);
        let lhs = CpuStorageBytes::from_slice(&[nan, two, nan, one]);
        let rhs = CpuStorageBytes::from_slice(&[three, nan, nan, two]);
        let mut out = CpuStorageBytes::from_zero_bytes(lhs.len_bytes());

        maximum_bf16(&lhs, &rhs, &mut out).expect("max bf16");
        let r: &[half::bf16] = out.as_slice().unwrap();
        assert!(r[0].is_nan(), "maximum(NaN, 3) must be NaN");
        assert!(r[1].is_nan(), "maximum(2, NaN) must be NaN");
        assert!(r[2].is_nan(), "maximum(NaN, NaN) must be NaN");
        assert_eq!(r[3].to_f32(), 2.0, "maximum(1, 2) non-NaN sanity");

        minimum_bf16(&lhs, &rhs, &mut out).expect("min bf16");
        let r: &[half::bf16] = out.as_slice().unwrap();
        assert!(r[0].is_nan(), "minimum(NaN, 3) must be NaN");
        assert!(r[1].is_nan(), "minimum(2, NaN) must be NaN");
        assert!(r[2].is_nan(), "minimum(NaN, NaN) must be NaN");
        assert_eq!(r[3].to_f32(), 1.0, "minimum(1, 2) non-NaN sanity");
    }

    #[test]
    fn conv2d_f32_identity_3x3_kernel() {
        // 1×1 input ch, 1×1 output ch, 3×3 kernel with center 1 and rest 0
        // → output equals input.
        let x = CpuStorageBytes::from_slice(&[1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0]);
        let weight =
            CpuStorageBytes::from_slice(&[0.0_f32, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0]);
        let mut out = CpuStorageBytes::from_zero_bytes(9 * 4);
        conv2d_f32(
            &x,
            &weight,
            None,
            &mut out,
            [1, 1, 3, 3],
            [1, 1, 3, 3],
            [1, 1, 3, 3],
            (1, 1),
            (1, 1),
            (1, 1),
            1,
        )
        .expect("conv");
        assert_eq!(
            out.as_slice::<f32>().unwrap(),
            &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0]
        );
    }

    #[test]
    fn conv2d_f32_2x2_sum_kernel_no_padding() {
        // Input 1×1×3×3; kernel 2×2 of all-ones; no padding; stride 1.
        // Output shape 1×1×2×2; each output is sum of a 2×2 window.
        let x = CpuStorageBytes::from_slice(&[1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0]);
        let weight = CpuStorageBytes::from_slice(&[1.0_f32, 1.0, 1.0, 1.0]);
        let mut out = CpuStorageBytes::from_zero_bytes(4 * 4);
        conv2d_f32(
            &x,
            &weight,
            None,
            &mut out,
            [1, 1, 3, 3],
            [1, 1, 2, 2],
            [1, 1, 2, 2],
            (1, 1),
            (0, 0),
            (1, 1),
            1,
        )
        .expect("conv");
        // Window at (0,0): 1+2+4+5 = 12
        // Window at (0,1): 2+3+5+6 = 16
        // Window at (1,0): 4+5+7+8 = 24
        // Window at (1,1): 5+6+8+9 = 28
        assert_eq!(out.as_slice::<f32>().unwrap(), &[12.0, 16.0, 24.0, 28.0]);
    }

    #[test]
    fn conv2d_f32_with_bias() {
        // Same as the 2×2 sum kernel test, plus a bias of 100 — every
        // output gets +100.
        let x = CpuStorageBytes::from_slice(&[1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0]);
        let weight = CpuStorageBytes::from_slice(&[1.0_f32, 1.0, 1.0, 1.0]);
        let bias = CpuStorageBytes::from_slice(&[100.0_f32]);
        let mut out = CpuStorageBytes::from_zero_bytes(4 * 4);
        conv2d_f32(
            &x,
            &weight,
            Some(&bias),
            &mut out,
            [1, 1, 3, 3],
            [1, 1, 2, 2],
            [1, 1, 2, 2],
            (1, 1),
            (0, 0),
            (1, 1),
            1,
        )
        .expect("conv");
        assert_eq!(
            out.as_slice::<f32>().unwrap(),
            &[112.0, 116.0, 124.0, 128.0]
        );
    }

    #[test]
    fn conv2d_f32_padding_yields_zero_outside() {
        // 1×1×2×2 input, 2×2 kernel of ones, padding 1 → out shape 1×1×3×3.
        // Output[0,0] reads only x[0,0] (rest is padded zeros) → 1.
        let x = CpuStorageBytes::from_slice(&[1.0_f32, 2.0, 3.0, 4.0]);
        let weight = CpuStorageBytes::from_slice(&[1.0_f32, 1.0, 1.0, 1.0]);
        let mut out = CpuStorageBytes::from_zero_bytes(9 * 4);
        conv2d_f32(
            &x,
            &weight,
            None,
            &mut out,
            [1, 1, 2, 2],
            [1, 1, 2, 2],
            [1, 1, 3, 3],
            (1, 1),
            (1, 1),
            (1, 1),
            1,
        )
        .expect("conv");
        // Output positions:
        //   [0,0]: only x[0,0]=1
        //   [0,1]: x[0,0]+x[0,1]=1+2=3
        //   [0,2]: only x[0,1]=2
        //   [1,0]: x[0,0]+x[1,0]=1+3=4
        //   [1,1]: 1+2+3+4=10
        //   [1,2]: x[0,1]+x[1,1]=2+4=6
        //   [2,0]: only x[1,0]=3
        //   [2,1]: x[1,0]+x[1,1]=3+4=7
        //   [2,2]: only x[1,1]=4
        assert_eq!(
            out.as_slice::<f32>().unwrap(),
            &[1.0, 3.0, 2.0, 4.0, 10.0, 6.0, 3.0, 7.0, 4.0]
        );
    }

    #[test]
    fn conv2d_f32_depthwise_groups_equal_cin() {
        // 2 channels in, 2 channels out, groups=2 (depthwise). Each
        // output channel reads only its corresponding input channel.
        // x[0, ch=0] = [[1, 2], [3, 4]]; ch=1 = [[10, 20], [30, 40]]
        // weight: ch0 has all-ones 2x2; ch1 has identity
        let x = CpuStorageBytes::from_slice(&[
            1.0_f32, 2.0, 3.0, 4.0, // ch 0
            10.0, 20.0, 30.0, 40.0, // ch 1
        ]);
        // weight shape [Cout, Cin/groups=1, 2, 2] = [2, 1, 2, 2]
        let weight = CpuStorageBytes::from_slice(&[
            1.0_f32, 1.0, 1.0, 1.0, // co 0 sees ci 0
            1.0, 0.0, 0.0, 0.0, // co 1 sees ci 1, but only top-left
        ]);
        let mut out = CpuStorageBytes::from_zero_bytes(2 * 4); // [1, 2, 1, 1]
        conv2d_f32(
            &x,
            &weight,
            None,
            &mut out,
            [1, 2, 2, 2],
            [2, 1, 2, 2],
            [1, 2, 1, 1],
            (1, 1),
            (0, 0),
            (1, 1),
            2,
        )
        .expect("conv");
        // co 0: 1+2+3+4 = 10
        // co 1: 10 (only top-left of ch 1)
        assert_eq!(out.as_slice::<f32>().unwrap(), &[10.0, 10.0]);
    }

    #[test]
    fn cast_f32_to_f64_round_trip() {
        let input = CpuStorageBytes::from_slice(&[1.5_f32, -2.25, 100.0]);
        let mut out = CpuStorageBytes::from_zero_bytes(3 * 8);
        cast_f32_to_f64(&input, &mut out).expect("cast");
        assert_eq!(out.as_slice::<f64>().unwrap(), &[1.5_f64, -2.25, 100.0]);
    }

    #[test]
    fn cast_f64_to_f32_lossy() {
        let input = CpuStorageBytes::from_slice(&[1.5_f64, -2.25, 100.0]);
        let mut out = CpuStorageBytes::from_zero_bytes(3 * 4);
        cast_f64_to_f32(&input, &mut out).expect("cast");
        assert_eq!(out.as_slice::<f32>().unwrap(), &[1.5_f32, -2.25, 100.0]);
    }

    #[test]
    fn cast_bf16_round_trip_via_f32() {
        let input_f32 = CpuStorageBytes::from_slice(&[1.0_f32, 2.0, -3.0, 0.5]);
        let mut bf16_buf = CpuStorageBytes::from_zero_bytes(4 * 2);
        cast_f32_to_bf16(&input_f32, &mut bf16_buf).expect("to_bf16");

        let mut back_f32 = CpuStorageBytes::from_zero_bytes(4 * 4);
        cast_bf16_to_f32(&bf16_buf, &mut back_f32).expect("to_f32");

        let result: &[f32] = back_f32.as_slice().unwrap();
        // bf16 has ~3 decimal digits of precision; these inputs round-trip exactly.
        assert_eq!(result, &[1.0_f32, 2.0, -3.0, 0.5]);
    }

    #[test]
    fn cast_f16_round_trip_via_f32() {
        let input_f32 = CpuStorageBytes::from_slice(&[1.0_f32, 2.0, -3.0, 0.5]);
        let mut f16_buf = CpuStorageBytes::from_zero_bytes(4 * 2);
        cast_f32_to_f16(&input_f32, &mut f16_buf).expect("to_f16");

        let mut back_f32 = CpuStorageBytes::from_zero_bytes(4 * 4);
        cast_f16_to_f32(&f16_buf, &mut back_f32).expect("to_f32");

        let result: &[f32] = back_f32.as_slice().unwrap();
        // f16 has ~3 decimal digits of precision in this range.
        for (got, want) in result.iter().zip(&[1.0_f32, 2.0, -3.0, 0.5]) {
            assert!(
                (got - want).abs() < 1e-3,
                "f16 round trip lost too much: {got} vs {want}"
            );
        }
    }

    #[test]
    fn cast_size_mismatch_errors() {
        let input = CpuStorageBytes::from_slice(&[1.0_f32, 2.0]);
        let mut out_wrong = CpuStorageBytes::from_zero_bytes(8); // should be 16 (2 f64)
        let r = cast_f32_to_f64(&input, &mut out_wrong);
        assert!(r.is_err());
    }

    #[test]
    fn matmul_f32_2x3_times_3x2() {
        let lhs = CpuStorageBytes::from_slice(&[1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0]);
        let rhs = CpuStorageBytes::from_slice(&[7.0_f32, 8.0, 9.0, 10.0, 11.0, 12.0]);
        let mut out = CpuStorageBytes::from_zero_bytes(16);

        matmul_f32(&lhs, &rhs, &mut out, &[], &[], 2, 2, 3).expect("matmul");
        assert_eq!(out.as_slice::<f32>().unwrap(), &[58.0, 64.0, 139.0, 154.0]);
    }

    #[test]
    fn matmul_f32_identity_returns_input() {
        let lhs = CpuStorageBytes::from_slice(&[1.0_f32, 2.0, 3.0, 4.0]);
        let rhs = CpuStorageBytes::from_slice(&[1.0_f32, 0.0, 0.0, 1.0]);
        let mut out = CpuStorageBytes::from_zero_bytes(16);

        matmul_f32(&lhs, &rhs, &mut out, &[], &[], 2, 2, 2).expect("matmul");
        assert_eq!(out.as_slice::<f32>().unwrap(), &[1.0, 2.0, 3.0, 4.0]);
    }

    #[test]
    fn matmul_f32_inner_product_1x3_times_3x1() {
        let lhs = CpuStorageBytes::from_slice(&[1.0_f32, 2.0, 3.0]);
        let rhs = CpuStorageBytes::from_slice(&[4.0_f32, 5.0, 6.0]);
        let mut out = CpuStorageBytes::from_zero_bytes(4);

        matmul_f32(&lhs, &rhs, &mut out, &[], &[], 1, 1, 3).expect("matmul");
        assert_eq!(out.as_slice::<f32>().unwrap(), &[32.0]);
    }

    #[test]
    fn matmul_f32_size_mismatch_errors() {
        let lhs = CpuStorageBytes::from_slice(&[1.0_f32, 2.0, 3.0, 4.0]);
        let rhs = CpuStorageBytes::from_slice(&[1.0_f32, 2.0]);
        let mut out = CpuStorageBytes::from_zero_bytes(16);
        let r = matmul_f32(&lhs, &rhs, &mut out, &[], &[], 2, 2, 2);
        assert!(r.is_err(), "matmul must error on shape/byte mismatch");
    }

    #[test]
    fn matmul_f32_batched_2x_2x2_times_2x2() {
        let lhs = CpuStorageBytes::from_slice(&[1.0_f32, 2.0, 3.0, 4.0, 1.0, 0.0, 0.0, 1.0]);
        let rhs = CpuStorageBytes::from_slice(&[5.0_f32, 6.0, 7.0, 8.0, 10.0, 20.0, 30.0, 40.0]);
        let mut out = CpuStorageBytes::from_zero_bytes(2 * 4 * 4);

        matmul_f32(&lhs, &rhs, &mut out, &[2], &[2], 2, 2, 2).expect("matmul");
        assert_eq!(
            out.as_slice::<f32>().unwrap(),
            &[19.0, 22.0, 43.0, 50.0, 10.0, 20.0, 30.0, 40.0]
        );
    }

    /// GQA-style: lhs has 4 batch heads, rhs has 2; n_rep = 2.
    /// lhs heads {0, 1} share rhs head 0; lhs heads {2, 3} share rhs head 1.
    #[test]
    fn matmul_f32_gqa_4x_vs_2x() {
        // lhs shape [4, 1, 2]: heads 0..3 are [[1,2]], [[3,4]], [[5,6]], [[7,8]]
        let lhs = CpuStorageBytes::from_slice(&[1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]);
        // rhs shape [2, 2, 1]: heads 0,1 are [[1],[0]], [[0],[1]]
        let rhs = CpuStorageBytes::from_slice(&[1.0_f32, 0.0, 0.0, 1.0]);
        // output shape [4, 1, 1]:
        //   head 0: [[1, 2]] @ [[1], [0]] = [[1]]
        //   head 1: [[3, 4]] @ [[1], [0]] = [[3]]   (still rhs 0)
        //   head 2: [[5, 6]] @ [[0], [1]] = [[6]]   (rhs 1)
        //   head 3: [[7, 8]] @ [[0], [1]] = [[8]]   (rhs 1)
        let mut out = CpuStorageBytes::from_zero_bytes(4 * 4);

        matmul_f32(&lhs, &rhs, &mut out, &[4], &[2], 1, 1, 2).expect("gqa matmul");
        assert_eq!(out.as_slice::<f32>().unwrap(), &[1.0, 3.0, 6.0, 8.0]);
    }

    #[test]
    fn matmul_f32_gqa_rejects_non_divisible() {
        let lhs = CpuStorageBytes::from_slice(&[1.0_f32; 6]); // [3, 1, 2]
        let rhs = CpuStorageBytes::from_slice(&[1.0_f32; 4]); // [2, 2, 1]
        let mut out = CpuStorageBytes::from_zero_bytes(3 * 4);
        // 3 is NOT a multiple of 2 — must error.
        let r = matmul_f32(&lhs, &rhs, &mut out, &[3], &[2], 1, 1, 2);
        assert!(r.is_err());
    }

    #[test]
    fn contiguize_cpu_no_op_on_contiguous_input() {
        let input = CpuStorageBytes::from_slice(&[1.0_f32, 2.0, 3.0, 4.0]);
        let layout = Layout::contiguous(fuel_ir::Shape::from_dims(&[2, 2]));
        let out = contiguize_cpu(&input, &layout, 4).expect("contiguize");
        assert_eq!(out.as_slice::<f32>().unwrap(), &[1.0, 2.0, 3.0, 4.0]);
    }

    #[test]
    fn contiguize_cpu_realizes_transpose() {
        // shape [2, 3], stored row-major: 1 2 3 / 4 5 6
        let input = CpuStorageBytes::from_slice(&[1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0]);
        // Transposed view: shape [3, 2], strides [1, 3]
        let layout = Layout::new(
            fuel_ir::Shape::from_dims(&[3, 2]),
            fuel_ir::StrideVec::from_slice(&[1_isize, 3]),
            0,
        );
        let out = contiguize_cpu(&input, &layout, 4).expect("contiguize");
        // Transposed: column 0 of source (1, 4) becomes row 0; col 1 (2, 5) row 1; col 2 (3, 6) row 2
        assert_eq!(
            out.as_slice::<f32>().unwrap(),
            &[1.0, 4.0, 2.0, 5.0, 3.0, 6.0]
        );
    }

    #[test]
    fn contiguize_cpu_replicates_broadcast() {
        // shape [3] broadcast to [2, 3] — leading dim has stride 0
        let input = CpuStorageBytes::from_slice(&[10.0_f32, 20.0, 30.0]);
        let layout = Layout::new(
            fuel_ir::Shape::from_dims(&[2, 3]),
            fuel_ir::StrideVec::from_slice(&[0_isize, 1]),
            0,
        );
        let out = contiguize_cpu(&input, &layout, 4).expect("contiguize");
        assert_eq!(
            out.as_slice::<f32>().unwrap(),
            &[10.0, 20.0, 30.0, 10.0, 20.0, 30.0]
        );
    }

    #[test]
    fn contiguize_cpu_offset_layout() {
        // input has 5 elements; layout views the last 3 with offset 2
        let input = CpuStorageBytes::from_slice(&[100.0_f32, 200.0, 1.0, 2.0, 3.0]);
        let layout = Layout::new(
            fuel_ir::Shape::from_dims(&[3]),
            fuel_ir::StrideVec::from_slice(&[1_isize]),
            2,
        );
        let out = contiguize_cpu(&input, &layout, 4).expect("contiguize");
        assert_eq!(out.as_slice::<f32>().unwrap(), &[1.0, 2.0, 3.0]);
    }

    #[test]
    fn sum_reduce_along_one_dim() {
        // shape [2, 3], reduce dim 1 → output shape [2]
        let input = CpuStorageBytes::from_slice(&[1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0]);
        let mut out = CpuStorageBytes::from_zero_bytes(8); // 2 f32s

        sum_reduce_f32(&input, &mut out, &[2, 3], &[1]).expect("sum_reduce");
        // row 0: 1+2+3=6; row 1: 4+5+6=15
        assert_eq!(out.as_slice::<f32>().unwrap(), &[6.0, 15.0]);
    }

    #[test]
    fn sum_reduce_all_dims() {
        let input = CpuStorageBytes::from_slice(&[1.0_f32, 2.0, 3.0, 4.0]);
        let mut out = CpuStorageBytes::from_zero_bytes(4); // 1 f32

        sum_reduce_f32(&input, &mut out, &[2, 2], &[0, 1]).expect("sum_reduce_all");
        assert_eq!(out.as_slice::<f32>().unwrap(), &[10.0]);
    }

    #[test]
    fn sum_reduce_inner_dim_of_rank_3() {
        // shape [2, 2, 3], reduce dim 1 → output shape [2, 3]
        // Layout (row-major):
        //   batch=0: [[1,2,3],[4,5,6]]   → reduce dim 1 → [5,7,9]
        //   batch=1: [[7,8,9],[10,11,12]] → reduce dim 1 → [17,19,21]
        let data: Vec<f32> = (1..=12).map(|x| x as f32).collect();
        let input = CpuStorageBytes::from_slice(&data);
        let mut out = CpuStorageBytes::from_zero_bytes(24); // 6 f32s

        sum_reduce_f32(&input, &mut out, &[2, 2, 3], &[1]).expect("sum_reduce");
        assert_eq!(
            out.as_slice::<f32>().unwrap(),
            &[5.0, 7.0, 9.0, 17.0, 19.0, 21.0]
        );
    }

    #[test]
    fn mean_reduce_divides_by_reduced_dim_count() {
        let input = CpuStorageBytes::from_slice(&[2.0_f32, 4.0, 6.0, 8.0]);
        let mut out = CpuStorageBytes::from_zero_bytes(8); // [2] output

        mean_reduce_f32(&input, &mut out, &[2, 2], &[1]).expect("mean_reduce");
        // row 0 mean = (2+4)/2 = 3; row 1 mean = (6+8)/2 = 7
        assert_eq!(out.as_slice::<f32>().unwrap(), &[3.0, 7.0]);
    }

    #[test]
    fn max_reduce_f32_basic() {
        let input = CpuStorageBytes::from_slice(&[1.0_f32, -5.0, 3.0, 2.0, 0.0, -1.0]);
        let mut out = CpuStorageBytes::from_zero_bytes(8); // [2] output

        max_reduce_f32(&input, &mut out, &[2, 3], &[1]).expect("max_reduce");
        assert_eq!(out.as_slice::<f32>().unwrap(), &[3.0, 2.0]);
    }

    #[test]
    fn min_reduce_f32_basic() {
        let input = CpuStorageBytes::from_slice(&[1.0_f32, -5.0, 3.0, 2.0, 0.0, -1.0]);
        let mut out = CpuStorageBytes::from_zero_bytes(8); // [2] output

        min_reduce_f32(&input, &mut out, &[2, 3], &[1]).expect("min_reduce");
        assert_eq!(out.as_slice::<f32>().unwrap(), &[-5.0, -1.0]);
    }

    #[test]
    fn reduce_errors_on_shape_mismatch() {
        let input = CpuStorageBytes::from_slice(&[1.0_f32, 2.0, 3.0]); // 3 elems
        let mut out = CpuStorageBytes::from_zero_bytes(8);
        // shape says 4 elems but storage has 3 — must error.
        let r = sum_reduce_f32(&input, &mut out, &[2, 2], &[1]);
        assert!(r.is_err());
    }

    #[test]
    fn reduce_errors_on_unsorted_dims() {
        let input = CpuStorageBytes::from_slice(&[1.0_f32, 2.0, 3.0, 4.0]);
        let mut out = CpuStorageBytes::from_zero_bytes(4);
        let r = sum_reduce_f32(&input, &mut out, &[2, 2], &[1, 0]); // unsorted
        assert!(r.is_err());
    }

    #[test]
    fn reduce_errors_on_out_of_range_dim() {
        let input = CpuStorageBytes::from_slice(&[1.0_f32, 2.0]);
        let mut out = CpuStorageBytes::from_zero_bytes(4);
        let r = sum_reduce_f32(&input, &mut out, &[2], &[5]); // dim 5 doesn't exist
        assert!(r.is_err());
    }

    #[test]
    fn gelu_at_known_points() {
        let input = CpuStorageBytes::from_slice(&[0.0_f32, 1.0, -1.0]);
        let mut out = CpuStorageBytes::from_zero_bytes(input.len_bytes());
        gelu_f32(&input, &mut out).expect("gelu_tanh");
        let r: &[f32] = out.as_slice().unwrap();
        // gelu_tanh(0) = 0
        assert!(r[0].abs() < 1e-6);
        // gelu_tanh(1) ≈ 0.8412 (tanh approx)
        assert!((r[1] - 0.841_192).abs() < 1e-3);
        // gelu_tanh(-1) ≈ -0.1588 (tanh approx)
        assert!((r[2] - (-0.158_808)).abs() < 1e-3);
    }

    // ---- write_slice_cpu (Phase E.3.2) -------------------------------------

    /// KV-cache append shape: write source [1, 3, 2] into dest [4, 3, 2]
    /// at axis-0 row 2. Result: dest row 2 is replaced; rows 0/1/3 are
    /// untouched.
    #[test]
    fn write_slice_cpu_kv_cache_append_pattern() {
        // dest: 4 rows of 6 f32 each = 24 elements. Init to 0.0.
        let mut dest = CpuStorageBytes::from_slice(&[0.0_f32; 24]);
        let source = CpuStorageBytes::from_slice(&[1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0]);
        write_slice_cpu(
            &source,
            &mut dest,
            &[4, 3, 2],
            &[(2, 3), (0, 3), (0, 2)],
            std::mem::size_of::<f32>(),
        )
        .expect("write_slice");
        let r: &[f32] = dest.as_slice().unwrap();
        let expected = [
            0.0, 0.0, 0.0, 0.0, 0.0, 0.0, // row 0 untouched
            0.0, 0.0, 0.0, 0.0, 0.0, 0.0, // row 1 untouched
            1.0, 2.0, 3.0, 4.0, 5.0, 6.0, // row 2 = source
            0.0, 0.0, 0.0, 0.0, 0.0, 0.0, // row 3 untouched
        ];
        assert_eq!(r, &expected);
    }

    /// Interior write: dest [3, 4] receives source [2, 2] at center.
    /// Verifies non-full slab along both axes.
    #[test]
    fn write_slice_cpu_interior_2d() {
        let mut dest = CpuStorageBytes::from_slice(&[
            10.0_f32, 11.0, 12.0, 13.0, 14.0, 15.0, 16.0, 17.0, 18.0, 19.0, 20.0, 21.0,
        ]);
        let source = CpuStorageBytes::from_slice(&[100.0_f32, 101.0, 102.0, 103.0]);
        write_slice_cpu(
            &source,
            &mut dest,
            &[3, 4],
            &[(0, 2), (1, 3)],
            std::mem::size_of::<f32>(),
        )
        .expect("write_slice");
        let r: &[f32] = dest.as_slice().unwrap();
        let expected = [
            10.0, 100.0, 101.0, 13.0, 14.0, 102.0, 103.0, 17.0, 18.0, 19.0, 20.0, 21.0,
        ];
        assert_eq!(r, &expected);
    }

    /// 1-D write: special-cased single contiguous span.
    #[test]
    fn write_slice_cpu_1d() {
        let mut dest = CpuStorageBytes::from_slice(&[0.0_f32, 1.0, 2.0, 3.0, 4.0]);
        let source = CpuStorageBytes::from_slice(&[100.0_f32, 200.0]);
        write_slice_cpu(
            &source,
            &mut dest,
            &[5],
            &[(2, 4)],
            std::mem::size_of::<f32>(),
        )
        .expect("write_slice");
        let r: &[f32] = dest.as_slice().unwrap();
        assert_eq!(r, &[0.0, 1.0, 100.0, 200.0, 4.0]);
    }

    /// Rejects out-of-range slab.
    #[test]
    fn write_slice_cpu_rejects_out_of_range() {
        let mut dest = CpuStorageBytes::from_slice(&[0.0_f32; 12]);
        let source = CpuStorageBytes::from_slice(&[1.0_f32, 2.0, 3.0]);
        let err = write_slice_cpu(
            &source,
            &mut dest,
            &[4, 3],
            &[(4, 5), (0, 3)], // axis 0 range past dest dim
            std::mem::size_of::<f32>(),
        );
        assert!(err.is_err());
    }

    /// Rejects rank mismatch.
    #[test]
    fn write_slice_cpu_rejects_rank_mismatch() {
        let mut dest = CpuStorageBytes::from_slice(&[0.0_f32; 12]);
        let source = CpuStorageBytes::from_slice(&[1.0_f32; 3]);
        let err = write_slice_cpu(
            &source,
            &mut dest,
            &[4, 3],
            &[(0, 1)], // rank 1 vs dest rank 2
            std::mem::size_of::<f32>(),
        );
        assert!(err.is_err());
    }

    /// f64 + bf16 dtype-agnostic byte-level copies.
    #[test]
    fn write_slice_cpu_f64_byte_correct() {
        let mut dest = CpuStorageBytes::from_slice(&[0.0_f64; 6]);
        let source = CpuStorageBytes::from_slice(&[1.5_f64, 2.5]);
        write_slice_cpu(
            &source,
            &mut dest,
            &[3, 2],
            &[(1, 2), (0, 2)],
            std::mem::size_of::<f64>(),
        )
        .expect("write_slice");
        let r: &[f64] = dest.as_slice().unwrap();
        assert_eq!(r, &[0.0, 0.0, 1.5, 2.5, 0.0, 0.0]);
    }

    // ---- write_slice_rotating_cpu (Phase C) --------------------------------

    /// Sliding-window write before the wrap: position < modulus and
    /// `wrapped_start + slab < modulus` so no boundary split.
    #[test]
    fn write_slice_rotating_cpu_within_window() {
        // modulus = 4, dest [4, 2] (axis 0 rotates), write a 1-row slab
        // at position 1 — should land at row 1 only.
        let mut dest = CpuStorageBytes::from_slice(&[0.0_f32; 8]);
        let source = CpuStorageBytes::from_slice(&[7.0_f32, 8.0]);
        let position = CpuStorageBytes::from_slice(&[1_u32]);
        write_slice_rotating_cpu(
            &source,
            &position,
            &mut dest,
            &[4, 2],
            /* axis */ 0,
            /* modulus */ 4,
            &[(0, 1), (0, 2)],
            std::mem::size_of::<f32>(),
        )
        .expect("write_slice_rotating");
        let r: &[f32] = dest.as_slice().unwrap();
        assert_eq!(r, &[0.0, 0.0, 7.0, 8.0, 0.0, 0.0, 0.0, 0.0]);
    }

    /// Sliding-window write past modulus: position == modulus wraps to 0.
    #[test]
    fn write_slice_rotating_cpu_wraps_position() {
        let mut dest = CpuStorageBytes::from_slice(&[0.0_f32; 8]);
        let source = CpuStorageBytes::from_slice(&[7.0_f32, 8.0]);
        let position = CpuStorageBytes::from_slice(&[4_u32]); // == modulus
        write_slice_rotating_cpu(
            &source,
            &position,
            &mut dest,
            &[4, 2],
            0,
            4,
            &[(0, 1), (0, 2)],
            std::mem::size_of::<f32>(),
        )
        .expect("write_slice_rotating");
        let r: &[f32] = dest.as_slice().unwrap();
        // position 4 % 4 = 0 → writes to row 0.
        assert_eq!(r, &[7.0, 8.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]);
    }

    /// Boundary split: position == 3, slab = 2 rows, modulus = 4.
    /// First row at index 3, second wraps to index 0.
    #[test]
    fn write_slice_rotating_cpu_splits_across_boundary() {
        let mut dest = CpuStorageBytes::from_slice(&[0.0_f32; 8]);
        let source = CpuStorageBytes::from_slice(&[10.0_f32, 11.0, 20.0, 21.0]);
        let position = CpuStorageBytes::from_slice(&[3_u32]);
        write_slice_rotating_cpu(
            &source,
            &position,
            &mut dest,
            &[4, 2],
            0,
            4,
            &[(0, 2), (0, 2)],
            std::mem::size_of::<f32>(),
        )
        .expect("write_slice_rotating");
        let r: &[f32] = dest.as_slice().unwrap();
        // row 3 = (10, 11), row 0 = (20, 21).
        assert_eq!(r, &[20.0, 21.0, 0.0, 0.0, 0.0, 0.0, 10.0, 11.0]);
    }

    /// modulus less than dest_shape[axis]: ring lives in a leading
    /// prefix of the destination. Rows past `modulus` should never be
    /// written.
    #[test]
    fn write_slice_rotating_cpu_modulus_smaller_than_dest() {
        // dest [6, 2], modulus = 4 → only rows 0..4 participate.
        // position 5 % 4 = 1, slab 1 row → writes row 1.
        let mut dest = CpuStorageBytes::from_slice(&[0.0_f32; 12]);
        let source = CpuStorageBytes::from_slice(&[7.0_f32, 8.0]);
        let position = CpuStorageBytes::from_slice(&[5_u32]);
        write_slice_rotating_cpu(
            &source,
            &position,
            &mut dest,
            &[6, 2],
            0,
            4,
            &[(0, 1), (0, 2)],
            std::mem::size_of::<f32>(),
        )
        .expect("write_slice_rotating");
        let r: &[f32] = dest.as_slice().unwrap();
        assert_eq!(
            r,
            &[0.0, 0.0, 7.0, 8.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
        );
    }

    /// Rotating axis == 1 (the inner-of-rank-2 dim). dest `[2, 4]`,
    /// modulus 4, slab `[2, 1]` at position 1 → writes column 1 of
    /// both rows.
    #[test]
    fn write_slice_rotating_cpu_inner_axis_within_window() {
        // dest layout (row-major): [r0c0 r0c1 r0c2 r0c3 r1c0 r1c1 r1c2 r1c3]
        let mut dest = CpuStorageBytes::from_slice(&[0.0_f32; 8]);
        let source = CpuStorageBytes::from_slice(&[7.0_f32, 8.0]); // 2 rows × 1 col
        let position = CpuStorageBytes::from_slice(&[1_u32]);
        write_slice_rotating_cpu(
            &source,
            &position,
            &mut dest,
            &[2, 4],
            /* axis */ 1,
            /* modulus */ 4,
            &[(0, 2), (0, 1)],
            std::mem::size_of::<f32>(),
        )
        .expect("write_slice_rotating axis 1");
        let r: &[f32] = dest.as_slice().unwrap();
        // wrapped_start = 1 % 4 = 1 on axis 1 → writes column 1.
        assert_eq!(r, &[0.0, 7.0, 0.0, 0.0, 0.0, 8.0, 0.0, 0.0]);
    }

    /// Boundary split on axis 1: position 3, slab `[2, 2]`, modulus 4.
    /// First column lands at index 3, second wraps to index 0 (per row).
    #[test]
    fn write_slice_rotating_cpu_inner_axis_splits_across_boundary() {
        let mut dest = CpuStorageBytes::from_slice(&[0.0_f32; 8]);
        // source `[2, 2]` row-major: row0 = (10, 20), row1 = (30, 40).
        let source = CpuStorageBytes::from_slice(&[10.0_f32, 20.0, 30.0, 40.0]);
        let position = CpuStorageBytes::from_slice(&[3_u32]);
        write_slice_rotating_cpu(
            &source,
            &position,
            &mut dest,
            &[2, 4],
            1,
            4,
            &[(0, 2), (0, 2)],
            std::mem::size_of::<f32>(),
        )
        .expect("write_slice_rotating axis 1 boundary split");
        let r: &[f32] = dest.as_slice().unwrap();
        // wrapped_start = 3 on axis 1. row0 col 3 = 10, row0 col 0 (wrap) = 20.
        // row1 col 3 = 30, row1 col 0 (wrap) = 40.
        assert_eq!(r, &[20.0, 0.0, 0.0, 10.0, 40.0, 0.0, 0.0, 30.0]);
    }

    /// Middle-axis split (axis 1 of rank-3 dest). dest `[2, 4, 2]`,
    /// modulus 4 on axis 1, slab `[2, 2, 2]` at position 3. Each
    /// outer index (axis 0) gets a strided gather.
    #[test]
    fn write_slice_rotating_cpu_middle_axis_splits() {
        // dest shape [2, 4, 2] → 16 elems zero-initialized.
        let mut dest = CpuStorageBytes::from_slice(&[0.0_f32; 16]);
        // source shape [2, 2, 2] = 8 elems, distinguishable values
        let source = CpuStorageBytes::from_slice(&[
            1.0_f32, 2.0, // outer 0 axis 0
            3.0, 4.0, // outer 0 axis 1
            10.0, 20.0, // outer 1 axis 0
            30.0, 40.0, // outer 1 axis 1
        ]);
        let position = CpuStorageBytes::from_slice(&[3_u32]);
        write_slice_rotating_cpu(
            &source,
            &position,
            &mut dest,
            &[2, 4, 2],
            /* axis */ 1,
            /* modulus */ 4,
            &[(0, 2), (0, 2), (0, 2)],
            std::mem::size_of::<f32>(),
        )
        .expect("write_slice_rotating middle-axis split");
        let r: &[f32] = dest.as_slice().unwrap();
        // wrapped_start = 3. On axis 1: chunk0 (rows in src axis-1 [0..1]) →
        // dest axis-1 [3..4]; chunk1 (rows in src axis-1 [1..2]) → dest axis-1 [0..1].
        // For outer=0: dest[0,0,:] = (3.0, 4.0) (chunk1), dest[0,3,:] = (1.0, 2.0) (chunk0)
        // For outer=1: dest[1,0,:] = (30.0, 40.0), dest[1,3,:] = (10.0, 20.0)
        // dest layout: outer × 4 axis-1 rows × 2 inner cols, row-major.
        assert_eq!(
            r,
            &[
                3.0, 4.0, // dest[0,0]
                0.0, 0.0, // dest[0,1]
                0.0, 0.0, // dest[0,2]
                1.0, 2.0, // dest[0,3]
                30.0, 40.0, // dest[1,0]
                0.0, 0.0, // dest[1,1]
                0.0, 0.0, // dest[1,2]
                10.0, 20.0, // dest[1,3]
            ],
        );
    }

    /// Slab width > modulus must error (can't overfill the ring).
    #[test]
    fn write_slice_rotating_cpu_rejects_slab_larger_than_modulus() {
        let mut dest = CpuStorageBytes::from_zero_bytes(16);
        let source = CpuStorageBytes::from_slice(&[1.0_f32; 10]);
        let position = CpuStorageBytes::from_slice(&[0_u32]);
        let r = write_slice_rotating_cpu(
            &source,
            &position,
            &mut dest,
            &[4, 2],
            0,
            /* modulus */ 4,
            &[(0, 5), (0, 2)], // slab on axis 0 is 5, > modulus 4
            std::mem::size_of::<f32>(),
        );
        assert!(r.is_err());
    }

    // ---- F8E4M3 casts -------------------------------------------------------
    //
    // Mirrors the live-CUDA pattern in
    // fuel-storage/tests/baracuda_cast_live.rs: exercise round-trips on
    // values that are exactly representable in F8E4M3 (1 sign bit, 4 exp,
    // 3 mantissa). All powers of two from 2^-9 up to 2^8 with mantissa=0
    // are exact; we use the small-integer subset for readable assertions.

    /// GAP-275: pin the f8e4m3 cast BOUNDARY, which nothing covered.
    ///
    /// ⚠️ THE EXISTING ROUND-TRIP TEST FEEDS ONLY THE EXACTLY-REPRESENTABLE
    /// INTERIOR — `0, ±0.5, ±1, ±2, 4` — a maximum magnitude of 4.0 against
    /// E4M3's max finite of 448. So overflow, saturation, infinity, NaN,
    /// subnormals and signed zero were all UNOBSERVED, and none of the 22 f8
    /// test names covered them.
    ///
    /// Every value below was MEASURED before being asserted; this pins correct
    /// behaviour rather than encoding a guess. All are conformant to OCP
    /// E4M3FN (the finite variant — see
    /// `f8e4m3_maps_to_the_ocp_finite_variant_not_fnuz`).
    ///
    /// ⚠️ WHY IT IS WORTH PINNING WHEN THERE IS NO DEFECT: `float8` 0.7.0's
    /// `is_infinite` predicate is broken for both f8 types (it returns true for
    /// `0x7E`, which is the FINITE 448), and `F8E5M2::INFINITY` is a finite
    /// normal. Fuel calls neither predicate today, so exposure is zero — but a
    /// future `HostBuffer::F8E5M2` arm would inherit the upstream bug AND its
    /// tests would pass, because nothing here exercises the boundary. This is
    /// the coverage that makes that trigger visible instead of silent.
    ///
    /// ⚠️ **This pins UPSTREAM behaviour, which is the point.** The kernel is
    /// generated by `cast_kernel_to_fp8!` and delegates straight to
    /// `float8::F8E4M3::from_f32` -- there is no Fuel arithmetic here to be
    /// wrong. So the test is a CHANGE-DETECTOR ON THE DEPENDENCY: a `float8`
    /// bump that starts mapping overflow to NaN, or drops the sign on zero,
    /// reds here instead of silently changing what fuel stores.
    ///
    /// Born-red observed by sabotage (overflow -> `0x7F` instead of saturating):
    /// `encoding 449: expected 0x7E, got 0x7F`, with a `Compiling` artifact so
    /// it was not a stale binary; restored byte-identically and re-verified.
    #[test]
    fn cast_f32_f8e4m3_boundary_saturates_and_preserves_nan() {
        // (input, expected stored byte, expected value back, why)
        let cases: &[(f32, u8, f32, &str)] = &[
            (448.0, 0x7E, 448.0, "max finite is representable exactly"),
            (
                449.0,
                0x7E,
                448.0,
                "overflow SATURATES -- E4M3FN has no infinity",
            ),
            (
                1e30,
                0x7E,
                448.0,
                "far overflow saturates, does not wrap or NaN",
            ),
            (-448.0, 0xFE, -448.0, "negative max finite, sign preserved"),
            (
                -1e30,
                0xFE,
                -448.0,
                "negative overflow saturates symmetrically",
            ),
            (
                f32::INFINITY,
                0x7E,
                448.0,
                "+inf saturates: E4M3FN cannot encode it",
            ),
            (
                f32::NEG_INFINITY,
                0xFE,
                -448.0,
                "-inf saturates symmetrically",
            ),
            (0.015625, 0x08, 0.015625, "smallest normal 2^-6"),
            (0.001953125, 0x01, 0.001953125, "subnormal 2^-9 survives"),
            (
                1e-30,
                0x00,
                0.0,
                "underflow flushes to +0, not to a subnormal",
            ),
            (-0.0, 0x80, -0.0, "signed zero KEEPS its sign bit"),
        ];
        // An emptied or trimmed table would make every assertion below
        // unreachable and the test would still report ok. A BOUND rather than an
        // exact count, so adding rows is free and removing them is not.
        assert!(
            cases.len() >= 10,
            "boundary table has {} rows -- it must keep the overflow, infinity, subnormal and signed-zero cases, or this test passes without asserting them",
            cases.len()
        );
        let mut checked = 0usize;
        for &(input, want_bits, want_back, why) in cases {
            let src = CpuStorageBytes::from_slice(&[input]);
            let mut fp8 = CpuStorageBytes::from_zero_bytes(1);
            cast_f32_to_f8e4m3(&src, &mut fp8).expect("f32 -> f8e4m3");
            let got_bits: &[u8] = fp8.as_slice().unwrap();
            assert_eq!(
                got_bits[0], want_bits,
                "encoding {input}: expected 0x{want_bits:02X}, got 0x{:02X} ({why})",
                got_bits[0]
            );
            let mut back = CpuStorageBytes::from_zero_bytes(4);
            cast_f8e4m3_to_f32(&fp8, &mut back).expect("f8e4m3 -> f32");
            let got: &[f32] = back.as_slice().unwrap();
            assert_eq!(
                got[0].to_bits(),
                want_back.to_bits(),
                "decoding {input}: expected {want_back}, got {} ({why})",
                got[0]
            );
            checked += 1;
        }
        assert_eq!(
            checked,
            cases.len(),
            "the boundary loop did not run to completion"
        );

        // NaN separately: `assert_eq!` on NaN is always false, so it needs its
        // own shape rather than a row above.
        let src = CpuStorageBytes::from_slice(&[f32::NAN]);
        let mut fp8 = CpuStorageBytes::from_zero_bytes(1);
        cast_f32_to_f8e4m3(&src, &mut fp8).expect("f32 -> f8e4m3");
        let got_bits: &[u8] = fp8.as_slice().unwrap();
        assert_eq!(
            got_bits[0], 0x7F,
            "NaN must encode to E4M3's canonical NaN 0x7F, got 0x{:02X}",
            got_bits[0]
        );
        let mut back = CpuStorageBytes::from_zero_bytes(4);
        cast_f8e4m3_to_f32(&fp8, &mut back).expect("f8e4m3 -> f32");
        let got: &[f32] = back.as_slice().unwrap();
        assert!(
            got[0].is_nan(),
            "NaN must survive the round trip, got {}",
            got[0]
        );
    }

    #[test]
    fn cast_f32_f8e4m3_round_trip_exact_for_representable() {
        // 0, ±0.5, ±1, ±2, ±4 — all exact in F8E4M3.
        let input: Vec<f32> = vec![0.0, 0.5, -0.5, 1.0, -1.0, 2.0, -2.0, 4.0];
        let src = CpuStorageBytes::from_slice(&input);
        let mut fp8 = CpuStorageBytes::from_zero_bytes(input.len());
        cast_f32_to_f8e4m3(&src, &mut fp8).expect("f32 → f8e4m3");
        assert_eq!(fp8.len_bytes(), input.len(), "F8E4M3 is 1 byte/elem");

        let mut back = CpuStorageBytes::from_zero_bytes(input.len() * 4);
        cast_f8e4m3_to_f32(&fp8, &mut back).expect("f8e4m3 → f32");
        let got: &[f32] = back.as_slice().unwrap();
        for (i, (&want, &g)) in input.iter().zip(got.iter()).enumerate() {
            assert_eq!(want, g, "round-trip mismatch at {i}: want {want}, got {g}");
        }
    }

    #[test]
    fn cast_f8e4m3_through_bf16_round_trip() {
        let input: Vec<f32> = vec![0.0, 0.5, 1.0, 2.0];
        let src = CpuStorageBytes::from_slice(&input);
        let mut fp8 = CpuStorageBytes::from_zero_bytes(input.len());
        cast_f32_to_f8e4m3(&src, &mut fp8).expect("f32 → f8e4m3");

        let mut bf = CpuStorageBytes::from_zero_bytes(input.len() * 2);
        cast_f8e4m3_to_bf16(&fp8, &mut bf).expect("f8e4m3 → bf16");

        let mut fp8_back = CpuStorageBytes::from_zero_bytes(input.len());
        cast_bf16_to_f8e4m3(&bf, &mut fp8_back).expect("bf16 → f8e4m3");
        assert_eq!(
            fp8.bytes(),
            fp8_back.bytes(),
            "F8E4M3 → BF16 → F8E4M3 must be bit-stable for representable inputs",
        );
    }

    #[test]
    fn cast_f8e4m3_through_f16_round_trip() {
        let input: Vec<f32> = vec![0.0, 0.5, 1.0, 2.0];
        let src = CpuStorageBytes::from_slice(&input);
        let mut fp8 = CpuStorageBytes::from_zero_bytes(input.len());
        cast_f32_to_f8e4m3(&src, &mut fp8).expect("f32 → f8e4m3");

        let mut f16 = CpuStorageBytes::from_zero_bytes(input.len() * 2);
        cast_f8e4m3_to_f16(&fp8, &mut f16).expect("f8e4m3 → f16");

        let mut fp8_back = CpuStorageBytes::from_zero_bytes(input.len());
        cast_f16_to_f8e4m3(&f16, &mut fp8_back).expect("f16 → f8e4m3");
        assert_eq!(
            fp8.bytes(),
            fp8_back.bytes(),
            "F8E4M3 → F16 → F8E4M3 must be bit-stable for representable inputs",
        );
    }

    #[test]
    fn cast_f8e4m3_size_mismatch_errors() {
        let src = CpuStorageBytes::from_slice(&[1.0_f32, 2.0, 3.0]);
        let mut wrong_size = CpuStorageBytes::from_zero_bytes(5); // want 3
        assert!(cast_f32_to_f8e4m3(&src, &mut wrong_size).is_err());

        let fp8 = CpuStorageBytes::from_zero_bytes(3);
        let mut wrong_out = CpuStorageBytes::from_zero_bytes(7); // want 12
        assert!(cast_f8e4m3_to_f32(&fp8, &mut wrong_out).is_err());
    }

    /// End-to-end exercise of representative pairs from the newly-completed
    /// cast matrix (int↔float, float→int truncation, int→half, int↔int, and
    /// the F8E4M3 integer edges). Each pair runs through its byte kernel and
    /// the output bytes are asserted directly.
    #[test]
    fn cast_matrix_new_pairs_end_to_end() {
        // U32 → F32: 5u32 → 5.0f32 (widening, exact).
        {
            let src = CpuStorageBytes::from_slice(&[5_u32, 0, 4_294_967_295]);
            let mut out = CpuStorageBytes::from_zero_bytes(3 * 4);
            cast_u32_to_f32(&src, &mut out).expect("u32 → f32");
            let got: &[f32] = out.as_slice().unwrap();
            assert_eq!(got, &[5.0_f32, 0.0, 4_294_967_295.0]);
        }
        // F32 → I64: 3.9 truncates toward zero to 3; -3.9 → -3.
        {
            let src = CpuStorageBytes::from_slice(&[3.9_f32, -3.9, 0.0, 7.0]);
            let mut out = CpuStorageBytes::from_zero_bytes(4 * 8);
            cast_f32_to_i64(&src, &mut out).expect("f32 → i64");
            assert_eq!(out.as_slice::<i64>().unwrap(), &[3_i64, -3, 0, 7]);
        }
        // I64 → F16: small integers are exact in f16; round-trip via to_f32.
        {
            let src = CpuStorageBytes::from_slice(&[0_i64, 1, -2, 256]);
            let mut out = CpuStorageBytes::from_zero_bytes(4 * 2);
            cast_i64_to_f16(&src, &mut out).expect("i64 → f16");
            let got: &[half::f16] = out.as_slice().unwrap();
            let as_f32: Vec<f32> = got.iter().map(|h| h.to_f32()).collect();
            assert_eq!(as_f32, vec![0.0_f32, 1.0, -2.0, 256.0]);
        }
        // U8 → F32: 255u8 → 255.0f32.
        {
            let src = CpuStorageBytes::from_slice(&[0_u8, 1, 128, 255]);
            let mut out = CpuStorageBytes::from_zero_bytes(4 * 4);
            cast_u8_to_f32(&src, &mut out).expect("u8 → f32");
            assert_eq!(
                out.as_slice::<f32>().unwrap(),
                &[0.0_f32, 1.0, 128.0, 255.0]
            );
        }
        // I32 → I16: narrowing truncates the high bits (wrapping `as`).
        {
            let src = CpuStorageBytes::from_slice(&[1_i32, -1, 32767, 70000]);
            let mut out = CpuStorageBytes::from_zero_bytes(4 * 2);
            cast_i32_to_i16(&src, &mut out).expect("i32 → i16");
            assert_eq!(
                out.as_slice::<i16>().unwrap(),
                &[1_i16, -1, 32767, 70000_i32 as i16]
            );
        }
        // I16 → F8E4M3 → I16: 4 is exactly representable in F8E4M3 (mant=0).
        {
            let src = CpuStorageBytes::from_slice(&[0_i16, 1, 2, 4]);
            let mut fp8 = CpuStorageBytes::from_zero_bytes(4);
            cast_i16_to_f8e4m3(&src, &mut fp8).expect("i16 → f8e4m3");
            let mut back = CpuStorageBytes::from_zero_bytes(4 * 2);
            cast_f8e4m3_to_i16(&fp8, &mut back).expect("f8e4m3 → i16");
            assert_eq!(back.as_slice::<i16>().unwrap(), &[0_i16, 1, 2, 4]);
        }
        // BF16 → I32 via f32: 6.75 truncates to 6.
        {
            let f32_src = CpuStorageBytes::from_slice(&[6.75_f32, -1.5, 100.0]);
            let mut bf = CpuStorageBytes::from_zero_bytes(3 * 2);
            cast_f32_to_bf16(&f32_src, &mut bf).expect("f32 → bf16");
            let mut out = CpuStorageBytes::from_zero_bytes(3 * 4);
            cast_bf16_to_i32(&bf, &mut out).expect("bf16 → i32");
            assert_eq!(out.as_slice::<i32>().unwrap(), &[6_i32, -1, 100]);
        }
    }

    // ---- Integer MatMul (i8 / u8) -------------------------------------------

    /// 2×3 @ 3×2 → 2×2 with hand-computed expected values. Mirrors the
    /// baracuda `gemm_s8` shape exactly; no batching.
    #[test]
    fn matmul_i8_small() {
        // A = [[1, 2, 3], [4, 5, 6]]
        // B = [[1, 2], [3, 4], [5, 6]]
        // A @ B = [[1+6+15, 2+8+18], [4+15+30, 8+20+36]] = [[22, 28], [49, 64]]
        let a = CpuStorageBytes::from_slice(&[1_i8, 2, 3, 4, 5, 6]);
        let b = CpuStorageBytes::from_slice(&[1_i8, 2, 3, 4, 5, 6]);
        let mut out = CpuStorageBytes::from_zero_bytes(4);
        matmul_i8(&a, &b, &mut out, &[], &[], 2, 2, 3).expect("matmul_i8");
        assert_eq!(out.as_slice::<i8>().unwrap(), &[22, 28, 49, 64]);
    }

    /// Negative inputs accumulate correctly through i32.
    #[test]
    fn matmul_i8_with_negatives() {
        // [[-1, 2], [3, -4]] @ [[5, -6], [-7, 8]] =
        //   [[-1*5 + 2*-7, -1*-6 + 2*8], [3*5 + -4*-7, 3*-6 + -4*8]]
        // = [[-19, 22], [43, -50]]
        let a = CpuStorageBytes::from_slice(&[-1_i8, 2, 3, -4]);
        let b = CpuStorageBytes::from_slice(&[5_i8, -6, -7, 8]);
        let mut out = CpuStorageBytes::from_zero_bytes(4);
        matmul_i8(&a, &b, &mut out, &[], &[], 2, 2, 2).expect("matmul_i8");
        assert_eq!(out.as_slice::<i8>().unwrap(), &[-19, 22, 43, -50]);
    }

    /// Accumulator overflow saturates to i8::MAX rather than wrapping.
    /// 127 × 127 × 4 = 64516, far above i8::MAX (127); must clamp.
    #[test]
    fn matmul_i8_saturates_on_overflow() {
        // 1×4 @ 4×1: a single dot product that saturates positive.
        let a = CpuStorageBytes::from_slice(&[127_i8, 127, 127, 127]);
        let b = CpuStorageBytes::from_slice(&[127_i8, 127, 127, 127]);
        let mut out = CpuStorageBytes::from_zero_bytes(1);
        matmul_i8(&a, &b, &mut out, &[], &[], 1, 1, 4).expect("matmul_i8");
        assert_eq!(out.as_slice::<i8>().unwrap(), &[i8::MAX]);

        // And saturates negative: -128 × 127 × 4 ≈ -65k → i8::MIN.
        let a = CpuStorageBytes::from_slice(&[-128_i8, -128, -128, -128]);
        let b = CpuStorageBytes::from_slice(&[127_i8, 127, 127, 127]);
        let mut out = CpuStorageBytes::from_zero_bytes(1);
        matmul_i8(&a, &b, &mut out, &[], &[], 1, 1, 4).expect("matmul_i8");
        assert_eq!(out.as_slice::<i8>().unwrap(), &[i8::MIN]);
    }

    /// u8 path: same shape as the i8 small test, all positive.
    #[test]
    fn matmul_u8_small() {
        let a = CpuStorageBytes::from_slice(&[1_u8, 2, 3, 4, 5, 6]);
        let b = CpuStorageBytes::from_slice(&[1_u8, 2, 3, 4, 5, 6]);
        let mut out = CpuStorageBytes::from_zero_bytes(4);
        matmul_u8(&a, &b, &mut out, &[], &[], 2, 2, 3).expect("matmul_u8");
        assert_eq!(out.as_slice::<u8>().unwrap(), &[22, 28, 49, 64]);
    }

    /// u8 saturates at 255 on overflow (no negative direction).
    #[test]
    fn matmul_u8_saturates_on_overflow() {
        let a = CpuStorageBytes::from_slice(&[255_u8, 255, 255, 255]);
        let b = CpuStorageBytes::from_slice(&[255_u8, 255, 255, 255]);
        let mut out = CpuStorageBytes::from_zero_bytes(1);
        matmul_u8(&a, &b, &mut out, &[], &[], 1, 1, 4).expect("matmul_u8");
        assert_eq!(out.as_slice::<u8>().unwrap(), &[u8::MAX]);
    }

    /// Size-mismatch on any operand errors instead of panicking.
    #[test]
    fn matmul_i8_size_mismatch_errors() {
        let a = CpuStorageBytes::from_slice(&[1_i8, 2, 3]); // 1×3
        let b = CpuStorageBytes::from_slice(&[1_i8, 2, 3, 4]); // 4×1 — K mismatch
        let mut out = CpuStorageBytes::from_zero_bytes(1);
        assert!(matmul_i8(&a, &b, &mut out, &[], &[], 1, 1, 3).is_err());
    }

    /// FusedSoftmaxCrossEntropy: hand-computed loss for [2, 4] logits.
    /// row 0: logits=[1, 2, 3, 4], target=1
    ///   row_max=4, sum_exp=e(-3)+e(-2)+e(-1)+e(0)=0.04979+0.13534+0.36788+1=1.55301
    ///   log_sum_exp=0.44019; nll=-(2-4-0.44019)=2.44019
    /// row 1: logits=[0, 0, 0, 0], target=3
    ///   row_max=0, sum_exp=4, log_sum_exp=ln(4)=1.38629
    ///   nll=-(0-0-1.38629)=1.38629
    /// mean = (2.44019 + 1.38629) / 2 ≈ 1.91324
    #[test]
    fn fused_softmax_cross_entropy_f32_mean() {
        let logits = CpuStorageBytes::from_slice(&[1.0_f32, 2.0, 3.0, 4.0, 0.0, 0.0, 0.0, 0.0]);
        let targets = CpuStorageBytes::from_slice(&[1_i64, 3]);
        let mut out = CpuStorageBytes::from_zero_bytes(4);
        fused_softmax_cross_entropy_f32(&logits, &targets, &mut out, 2, 4, REDUCTION_MEAN, -100)
            .expect("fsce");
        let loss = out.as_slice::<f32>().unwrap()[0];
        let expected = (2.440_189_8_f32 + 1.386_294_4_f32) / 2.0;
        assert!(
            (loss - expected).abs() < 1e-5,
            "got {loss}, expected {expected}",
        );
    }

    /// Sum reduction: scalar = 2.44019 + 1.38629 ≈ 3.82648.
    #[test]
    fn fused_softmax_cross_entropy_f32_sum() {
        let logits = CpuStorageBytes::from_slice(&[1.0_f32, 2.0, 3.0, 4.0, 0.0, 0.0, 0.0, 0.0]);
        let targets = CpuStorageBytes::from_slice(&[1_i64, 3]);
        let mut out = CpuStorageBytes::from_zero_bytes(4);
        fused_softmax_cross_entropy_f32(&logits, &targets, &mut out, 2, 4, REDUCTION_SUM, -100)
            .expect("fsce");
        let loss = out.as_slice::<f32>().unwrap()[0];
        let expected = 2.440_189_8_f32 + 1.386_294_4_f32;
        assert!(
            (loss - expected).abs() < 1e-5,
            "got {loss}, expected {expected}",
        );
    }

    /// None reduction: per-row losses [2.44019, 1.38629].
    #[test]
    fn fused_softmax_cross_entropy_f32_none() {
        let logits = CpuStorageBytes::from_slice(&[1.0_f32, 2.0, 3.0, 4.0, 0.0, 0.0, 0.0, 0.0]);
        let targets = CpuStorageBytes::from_slice(&[1_i64, 3]);
        let mut out = CpuStorageBytes::from_zero_bytes(2 * 4);
        fused_softmax_cross_entropy_f32(&logits, &targets, &mut out, 2, 4, REDUCTION_NONE, -100)
            .expect("fsce");
        let per_row = out.as_slice::<f32>().unwrap();
        assert!(
            (per_row[0] - 2.440_189_8).abs() < 1e-5,
            "row 0: got {}",
            per_row[0]
        );
        assert!(
            (per_row[1] - 1.386_294_4).abs() < 1e-5,
            "row 1: got {}",
            per_row[1]
        );
    }

    /// ignore_index drops a row from both numerator and denominator.
    /// Same logits as above but with target row 1 set to ignore_index.
    /// Mean over 1 valid row = 2.44019.
    #[test]
    fn fused_softmax_cross_entropy_f32_ignore_index_mean() {
        let logits = CpuStorageBytes::from_slice(&[1.0_f32, 2.0, 3.0, 4.0, 0.0, 0.0, 0.0, 0.0]);
        let targets = CpuStorageBytes::from_slice(&[1_i64, -100]);
        let mut out = CpuStorageBytes::from_zero_bytes(4);
        fused_softmax_cross_entropy_f32(&logits, &targets, &mut out, 2, 4, REDUCTION_MEAN, -100)
            .expect("fsce");
        let loss = out.as_slice::<f32>().unwrap()[0];
        let expected = 2.440_189_8_f32;
        assert!(
            (loss - expected).abs() < 1e-5,
            "got {loss}, expected {expected}",
        );
    }

    /// All rows ignored — Mean returns 0.0 (no division by zero).
    #[test]
    fn fused_softmax_cross_entropy_f32_all_ignored_mean_is_zero() {
        let logits = CpuStorageBytes::from_slice(&[1.0_f32, 2.0, 3.0, 4.0]);
        let targets = CpuStorageBytes::from_slice(&[-100_i64]);
        let mut out = CpuStorageBytes::from_zero_bytes(4);
        fused_softmax_cross_entropy_f32(&logits, &targets, &mut out, 1, 4, REDUCTION_MEAN, -100)
            .expect("fsce");
        assert_eq!(out.as_slice::<f32>().unwrap()[0], 0.0);
    }

    /// Out-of-range target (not equal to ignore_index) errors.
    #[test]
    fn fused_softmax_cross_entropy_f32_target_out_of_range_errors() {
        let logits = CpuStorageBytes::from_slice(&[1.0_f32, 2.0, 3.0, 4.0]);
        let targets = CpuStorageBytes::from_slice(&[7_i64]); // 7 ≥ vocab=4
        let mut out = CpuStorageBytes::from_zero_bytes(4);
        let r = fused_softmax_cross_entropy_f32(
            &logits,
            &targets,
            &mut out,
            1,
            4,
            REDUCTION_MEAN,
            -100,
        );
        assert!(r.is_err());
    }

    /// F64 sibling: same expected mean as the F32 test, all logits
    /// are exactly representable in F64 so the result is bit-identical
    /// up to the F32 output narrowing.
    #[test]
    fn fused_softmax_cross_entropy_f64_mean() {
        let logits = CpuStorageBytes::from_slice(&[1.0_f64, 2.0, 3.0, 4.0, 0.0, 0.0, 0.0, 0.0]);
        let targets = CpuStorageBytes::from_slice(&[1_i64, 3]);
        let mut out = CpuStorageBytes::from_zero_bytes(4);
        fused_softmax_cross_entropy_f64(&logits, &targets, &mut out, 2, 4, REDUCTION_MEAN, -100)
            .expect("fsce f64");
        let loss = out.as_slice::<f32>().unwrap()[0];
        let expected = (2.440_189_8_f32 + 1.386_294_4_f32) / 2.0;
        assert!(
            (loss - expected).abs() < 1e-5,
            "got {loss}, expected {expected}",
        );
    }

    /// BF16 sibling: same inputs are exact in BF16 (1.0/2.0/3.0/4.0/0.0
    /// all have <8-bit mantissa). 5e-2 tolerance matches the
    /// SelectiveScan precedent — the f64 accumulator + f32 output
    /// narrowing means the actual error here is much tighter.
    #[test]
    fn fused_softmax_cross_entropy_bf16_mean() {
        let logits = CpuStorageBytes::from_slice(&[
            half::bf16::from_f32(1.0),
            half::bf16::from_f32(2.0),
            half::bf16::from_f32(3.0),
            half::bf16::from_f32(4.0),
            half::bf16::from_f32(0.0),
            half::bf16::from_f32(0.0),
            half::bf16::from_f32(0.0),
            half::bf16::from_f32(0.0),
        ]);
        let targets = CpuStorageBytes::from_slice(&[1_i64, 3]);
        let mut out = CpuStorageBytes::from_zero_bytes(4);
        fused_softmax_cross_entropy_bf16(&logits, &targets, &mut out, 2, 4, REDUCTION_MEAN, -100)
            .expect("fsce bf16");
        let loss = out.as_slice::<f32>().unwrap()[0];
        let expected = (2.440_189_8_f32 + 1.386_294_4_f32) / 2.0;
        assert!(
            (loss - expected).abs() < 5e-2,
            "got {loss}, expected {expected}",
        );
    }

    /// F16 sibling: 1e-2 tolerance matches the SelectiveScan precedent.
    #[test]
    fn fused_softmax_cross_entropy_f16_mean() {
        let logits = CpuStorageBytes::from_slice(&[
            half::f16::from_f32(1.0),
            half::f16::from_f32(2.0),
            half::f16::from_f32(3.0),
            half::f16::from_f32(4.0),
            half::f16::from_f32(0.0),
            half::f16::from_f32(0.0),
            half::f16::from_f32(0.0),
            half::f16::from_f32(0.0),
        ]);
        let targets = CpuStorageBytes::from_slice(&[1_i64, 3]);
        let mut out = CpuStorageBytes::from_zero_bytes(4);
        fused_softmax_cross_entropy_f16(&logits, &targets, &mut out, 2, 4, REDUCTION_MEAN, -100)
            .expect("fsce f16");
        let loss = out.as_slice::<f32>().unwrap()[0];
        let expected = (2.440_189_8_f32 + 1.386_294_4_f32) / 2.0;
        assert!(
            (loss - expected).abs() < 1e-2,
            "got {loss}, expected {expected}",
        );
    }

    /// CausalConv1d: hand-computed result for a single-channel,
    /// batch=1, kernel=3, seq_out=2 case. Caller pre-pads x with
    /// 2 zeros on the left.
    ///   x (pre-padded) = [0, 0, 1, 2]  (seq_in = 4)
    ///   weight         = [w0=0.5, w1=1.0, w2=2.0]
    ///   bias           = [b=0.1]
    ///   out[0] = 0.1 + 0.5*0 + 1.0*0 + 2.0*1 = 2.1
    ///   out[1] = 0.1 + 0.5*0 + 1.0*1 + 2.0*2 = 5.1
    #[test]
    fn causal_conv1d_f32_no_silu_basic() {
        let x = CpuStorageBytes::from_slice(&[0.0_f32, 0.0, 1.0, 2.0]);
        let w = CpuStorageBytes::from_slice(&[0.5_f32, 1.0, 2.0]);
        let b = CpuStorageBytes::from_slice(&[0.1_f32]);
        let mut out = CpuStorageBytes::from_zero_bytes(2 * 4);
        causal_conv1d_f32(&x, &w, &b, &mut out, 1, 1, 4, 2, 3, false).expect("causal_conv1d");
        let result: &[f32] = out.as_slice().unwrap();
        assert!((result[0] - 2.1).abs() < 1e-5, "out[0]={}", result[0]);
        assert!((result[1] - 5.1).abs() < 1e-5, "out[1]={}", result[1]);
    }

    /// CausalConv1d with SiLU: out = y * sigmoid(y) for the raw outputs
    /// computed in the no-silu test.
    #[test]
    fn causal_conv1d_f32_with_silu() {
        let x = CpuStorageBytes::from_slice(&[0.0_f32, 0.0, 1.0, 2.0]);
        let w = CpuStorageBytes::from_slice(&[0.5_f32, 1.0, 2.0]);
        let b = CpuStorageBytes::from_slice(&[0.1_f32]);
        let mut out = CpuStorageBytes::from_zero_bytes(2 * 4);
        causal_conv1d_f32(&x, &w, &b, &mut out, 1, 1, 4, 2, 3, true).expect("causal_conv1d");
        let result: &[f32] = out.as_slice().unwrap();
        let expected0 = 2.1_f32 / (1.0 + (-2.1_f32).exp());
        let expected1 = 5.1_f32 / (1.0 + (-5.1_f32).exp());
        assert!(
            (result[0] - expected0).abs() < 1e-5,
            "out[0]={}, expected {expected0}",
            result[0]
        );
        assert!(
            (result[1] - expected1).abs() < 1e-5,
            "out[1]={}, expected {expected1}",
            result[1]
        );
    }

    /// CausalConv1d: two channels share x layout but have different
    /// weights — verify the depthwise per-channel separation.
    #[test]
    fn causal_conv1d_f32_two_channels_depthwise() {
        // batch=1, channels=2, kernel=2, seq_out=3, so seq_in=4.
        // x = [c0: [0, 1, 2, 3], c1: [0, 10, 20, 30]]
        let x = CpuStorageBytes::from_slice(&[
            0.0_f32, 1.0, 2.0, 3.0, // channel 0
            0.0, 10.0, 20.0, 30.0, // channel 1
        ]);
        // weight c0 = [1.0, 1.0], c1 = [2.0, 3.0]
        let w = CpuStorageBytes::from_slice(&[
            1.0_f32, 1.0, // channel 0
            2.0, 3.0, // channel 1
        ]);
        // bias: c0 = 0.0, c1 = 0.0
        let b = CpuStorageBytes::from_slice(&[0.0_f32, 0.0]);
        let mut out = CpuStorageBytes::from_zero_bytes(2 * 3 * 4);
        causal_conv1d_f32(&x, &w, &b, &mut out, 1, 2, 4, 3, 2, false).expect("causal_conv1d");
        let result: &[f32] = out.as_slice().unwrap();
        // channel 0: out[t] = 1*x[t] + 1*x[t+1]
        //   t=0: 1*0 + 1*1 = 1
        //   t=1: 1*1 + 1*2 = 3
        //   t=2: 1*2 + 1*3 = 5
        // channel 1: out[t] = 2*x[t] + 3*x[t+1]
        //   t=0: 2*0 + 3*10 = 30
        //   t=1: 2*10 + 3*20 = 80
        //   t=2: 2*20 + 3*30 = 130
        assert_eq!(result, &[1.0, 3.0, 5.0, 30.0, 80.0, 130.0]);
    }

    /// CausalConv1d: kernel=1 degenerate case is pointwise (out[t] =
    /// weight * x[t] + bias). seq_in == seq_out.
    #[test]
    fn causal_conv1d_f32_kernel_1_is_pointwise() {
        let x = CpuStorageBytes::from_slice(&[1.0_f32, 2.0, 3.0]);
        let w = CpuStorageBytes::from_slice(&[2.0_f32]);
        let b = CpuStorageBytes::from_slice(&[1.0_f32]);
        let mut out = CpuStorageBytes::from_zero_bytes(3 * 4);
        causal_conv1d_f32(&x, &w, &b, &mut out, 1, 1, 3, 3, 1, false).expect("causal_conv1d");
        let result: &[f32] = out.as_slice().unwrap();
        assert_eq!(result, &[3.0, 5.0, 7.0]);
    }

    /// CausalConv1d F64 sanity — same math as `causal_conv1d_f32_no_silu_basic`
    /// but with F64 inputs. Verifies the native F64 path.
    #[test]
    fn causal_conv1d_f64_no_silu_basic() {
        let x = CpuStorageBytes::from_slice(&[0.0_f64, 0.0, 1.0, 2.0]);
        let w = CpuStorageBytes::from_slice(&[0.5_f64, 1.0, 2.0]);
        let b = CpuStorageBytes::from_slice(&[0.1_f64]);
        let mut out = CpuStorageBytes::from_zero_bytes(2 * 8);
        causal_conv1d_f64(&x, &w, &b, &mut out, 1, 1, 4, 2, 3, false).expect("causal_conv1d_f64");
        let r: &[f64] = out.as_slice().unwrap();
        assert!((r[0] - 2.1).abs() < 1e-10);
        assert!((r[1] - 5.1).abs() < 1e-10);
    }

    /// CausalConv1d BF16 sanity. Loose tolerance for the 7-bit mantissa.
    #[test]
    fn causal_conv1d_bf16_no_silu_basic() {
        let xf = [0.0_f32, 0.0, 1.0, 2.0];
        let wf = [0.5_f32, 1.0, 2.0];
        let bf = [0.1_f32];
        let xb: Vec<half::bf16> = xf.iter().map(|&v| half::bf16::from_f32(v)).collect();
        let wb: Vec<half::bf16> = wf.iter().map(|&v| half::bf16::from_f32(v)).collect();
        let bb: Vec<half::bf16> = bf.iter().map(|&v| half::bf16::from_f32(v)).collect();
        let x = CpuStorageBytes::from_slice(&xb);
        let w = CpuStorageBytes::from_slice(&wb);
        let b = CpuStorageBytes::from_slice(&bb);
        let mut out = CpuStorageBytes::from_zero_bytes(2 * 2);
        causal_conv1d_bf16(&x, &w, &b, &mut out, 1, 1, 4, 2, 3, false).expect("causal_conv1d_bf16");
        let r: &[half::bf16] = out.as_slice().unwrap();
        assert!((r[0].to_f32() - 2.1).abs() < 5e-2);
        assert!((r[1].to_f32() - 5.1).abs() < 5e-2);
    }

    /// CausalConv1d F16 sanity with SiLU. Tighter tolerance than BF16.
    #[test]
    fn causal_conv1d_f16_with_silu() {
        let xf = [0.0_f32, 0.0, 1.0, 2.0];
        let wf = [0.5_f32, 1.0, 2.0];
        let bf = [0.1_f32];
        let xb: Vec<half::f16> = xf.iter().map(|&v| half::f16::from_f32(v)).collect();
        let wb: Vec<half::f16> = wf.iter().map(|&v| half::f16::from_f32(v)).collect();
        let bb: Vec<half::f16> = bf.iter().map(|&v| half::f16::from_f32(v)).collect();
        let x = CpuStorageBytes::from_slice(&xb);
        let w = CpuStorageBytes::from_slice(&wb);
        let b = CpuStorageBytes::from_slice(&bb);
        let mut out = CpuStorageBytes::from_zero_bytes(2 * 2);
        causal_conv1d_f16(&x, &w, &b, &mut out, 1, 1, 4, 2, 3, true).expect("causal_conv1d_f16");
        let r: &[half::f16] = out.as_slice().unwrap();
        let e0 = 2.1_f32 / (1.0 + (-2.1_f32).exp());
        let e1 = 5.1_f32 / (1.0 + (-5.1_f32).exp());
        assert!((r[0].to_f32() - e0).abs() < 1e-2);
        assert!((r[1].to_f32() - e1).abs() < 1e-2);
    }

    /// CausalConv1d: bad shapes error rather than panicking.
    #[test]
    fn causal_conv1d_f32_bad_shapes_error() {
        let x = CpuStorageBytes::from_slice(&[0.0_f32; 4]);
        let w = CpuStorageBytes::from_slice(&[0.5_f32, 1.0, 2.0]);
        let b = CpuStorageBytes::from_slice(&[0.1_f32]);
        let mut out = CpuStorageBytes::from_zero_bytes(2 * 4);
        // seq_in (4) != seq_out (3) + kernel - 1 (2) → 4 != 5 → error
        let r = causal_conv1d_f32(&x, &w, &b, &mut out, 1, 1, 4, 3, 3, false);
        assert!(r.is_err());
    }

    /// SelectiveScan: minimal seqlen=1 case verifies the single-step
    /// recurrence. With h initialized to zero, after one step:
    ///   d  = delta[0,0,0] = 1.0
    ///   for j in 0..1:
    ///     h[0,0,0] = exp(1.0 * a[0,0]) * 0 + 1.0 * b[0,0,0] * u[0,0,0]
    ///              = 1.0 * 2.0 * 3.0 = 6.0
    ///   y[0,0,0] = h[0,0,0] * c[0,0,0] = 6.0 * 0.5 = 3.0
    /// (a, b, c, u all picked so the math is hand-verifiable)
    #[test]
    fn selective_scan_f32_single_step_seqlen_1() {
        let u = CpuStorageBytes::from_slice(&[3.0_f32]); // [1, 1, 1]
        let delta = CpuStorageBytes::from_slice(&[1.0_f32]); // [1, 1, 1]
        let a = CpuStorageBytes::from_slice(&[-1.0_f32]); // [1, 1] — note: exp(-1) used below
        let b = CpuStorageBytes::from_slice(&[2.0_f32]); // [1, 1, 1]
        let c = CpuStorageBytes::from_slice(&[0.5_f32]); // [1, 1, 1]
        let mut out = CpuStorageBytes::from_zero_bytes(8);
        selective_scan_f32(&u, &delta, &a, &b, &c, &mut out, 1, 1, 1, 1, false)
            .expect("selective_scan");
        let result: &[f32] = out.as_slice().unwrap();
        // h[0,0,0] = exp(1.0 * -1.0) * 0 + 1.0 * 2.0 * 3.0 = 6.0
        // y[0,0,0] = 6.0 * 0.5 = 3.0
        assert!((result[0] - 3.0).abs() < 1e-5, "got {}", result[0]);
    }

    /// SelectiveScan: two timesteps verify the state is threaded
    /// across steps. Same a/b/c/u as the single-step test repeated;
    /// the second step picks up h=6.0 from the first.
    #[test]
    fn selective_scan_f32_two_steps_state_threading() {
        let u = CpuStorageBytes::from_slice(&[3.0_f32, 3.0]); // [1, 2, 1]
        let delta = CpuStorageBytes::from_slice(&[1.0_f32, 1.0]); // [1, 2, 1]
        let a = CpuStorageBytes::from_slice(&[-1.0_f32]); // [1, 1]
        let b = CpuStorageBytes::from_slice(&[2.0_f32, 2.0]); // [1, 2, 1]
        let c = CpuStorageBytes::from_slice(&[0.5_f32, 0.5]); // [1, 2, 1]
        let mut out = CpuStorageBytes::from_zero_bytes(12);
        selective_scan_f32(&u, &delta, &a, &b, &c, &mut out, 1, 2, 1, 1, false)
            .expect("selective_scan");
        let result: &[f32] = out.as_slice().unwrap();
        // step 0: h = exp(-1) * 0 + 1*2*3 = 6, y = 6 * 0.5 = 3.0
        // step 1: h = exp(-1) * 6 + 1*2*3 = 6/e + 6, y = (6/e + 6) * 0.5
        let h_step1 = (-1.0_f32).exp() * 6.0 + 6.0;
        let y_step1 = h_step1 * 0.5;
        assert!((result[0] - 3.0).abs() < 1e-5, "step 0: {}", result[0]);
        assert!(
            (result[1] - y_step1).abs() < 1e-5,
            "step 1: {} expected {y_step1}",
            result[1]
        );
    }

    /// SelectiveScan: delta_softplus toggle applies softplus(delta).
    /// softplus(0) = ln(1+1) = ln 2 ≈ 0.693.
    #[test]
    fn selective_scan_f32_delta_softplus() {
        let u = CpuStorageBytes::from_slice(&[1.0_f32]);
        let delta = CpuStorageBytes::from_slice(&[0.0_f32]); // softplus(0) ≈ 0.693
        let a = CpuStorageBytes::from_slice(&[0.0_f32]); // exp(0.693 * 0) = 1
        let b = CpuStorageBytes::from_slice(&[1.0_f32]);
        let c = CpuStorageBytes::from_slice(&[1.0_f32]);
        let mut out = CpuStorageBytes::from_zero_bytes(8);
        selective_scan_f32(&u, &delta, &a, &b, &c, &mut out, 1, 1, 1, 1, true)
            .expect("selective_scan");
        // d = softplus(0) = ln 2
        // h = exp(ln2 * 0) * 0 + ln2 * 1 * 1 = ln 2
        // y = ln 2 * 1 = ln 2
        let result: &[f32] = out.as_slice().unwrap();
        let expected = 2.0_f32.ln();
        assert!(
            (result[0] - expected).abs() < 1e-5,
            "got {} expected {expected}",
            result[0]
        );
    }

    /// SelectiveScan F64 sanity — mirrors the F32 single-step test.
    #[test]
    fn selective_scan_f64_single_step() {
        let u = CpuStorageBytes::from_slice(&[3.0_f64]);
        let delta = CpuStorageBytes::from_slice(&[1.0_f64]);
        let a = CpuStorageBytes::from_slice(&[-1.0_f64]);
        let b = CpuStorageBytes::from_slice(&[2.0_f64]);
        let c = CpuStorageBytes::from_slice(&[0.5_f64]);
        let mut out = CpuStorageBytes::from_zero_bytes(16);
        selective_scan_f64(&u, &delta, &a, &b, &c, &mut out, 1, 1, 1, 1, false)
            .expect("selective_scan_f64");
        let r: &[f64] = out.as_slice().unwrap();
        assert!((r[0] - 3.0).abs() < 1e-12);
    }

    /// SelectiveScan BF16 sanity. Loose tolerance for 7-bit mantissa
    /// (the F64 internal accumulator keeps the recurrence stable;
    /// rounding hits on the y store).
    #[test]
    fn selective_scan_bf16_single_step() {
        let u = CpuStorageBytes::from_slice(&[half::bf16::from_f32(3.0)]);
        let delta = CpuStorageBytes::from_slice(&[half::bf16::from_f32(1.0)]);
        let a = CpuStorageBytes::from_slice(&[half::bf16::from_f32(-1.0)]);
        let b = CpuStorageBytes::from_slice(&[half::bf16::from_f32(2.0)]);
        let c = CpuStorageBytes::from_slice(&[half::bf16::from_f32(0.5)]);
        let mut out = CpuStorageBytes::from_zero_bytes(4);
        selective_scan_bf16(&u, &delta, &a, &b, &c, &mut out, 1, 1, 1, 1, false)
            .expect("selective_scan_bf16");
        let r: &[half::bf16] = out.as_slice().unwrap();
        assert!((r[0].to_f32() - 3.0).abs() < 5e-2);
    }

    /// SelectiveScan F16 sanity. Tighter than BF16 (10-bit mantissa).
    #[test]
    fn selective_scan_f16_single_step() {
        let u = CpuStorageBytes::from_slice(&[half::f16::from_f32(3.0)]);
        let delta = CpuStorageBytes::from_slice(&[half::f16::from_f32(1.0)]);
        let a = CpuStorageBytes::from_slice(&[half::f16::from_f32(-1.0)]);
        let b = CpuStorageBytes::from_slice(&[half::f16::from_f32(2.0)]);
        let c = CpuStorageBytes::from_slice(&[half::f16::from_f32(0.5)]);
        let mut out = CpuStorageBytes::from_zero_bytes(4);
        selective_scan_f16(&u, &delta, &a, &b, &c, &mut out, 1, 1, 1, 1, false)
            .expect("selective_scan_f16");
        let r: &[half::f16] = out.as_slice().unwrap();
        assert!((r[0].to_f32() - 3.0).abs() < 1e-2);
    }

    /// SelectiveScan: bad shapes error rather than panicking.
    #[test]
    fn selective_scan_f32_bad_shapes_error() {
        let u = CpuStorageBytes::from_slice(&[1.0_f32]);
        let delta = CpuStorageBytes::from_slice(&[1.0_f32]);
        let a = CpuStorageBytes::from_slice(&[1.0_f32]);
        let b = CpuStorageBytes::from_slice(&[1.0_f32]);
        let c = CpuStorageBytes::from_slice(&[1.0_f32]);
        let mut out = CpuStorageBytes::from_zero_bytes(16);
        // batch=2 but only 1-element inputs → error
        let r = selective_scan_f32(&u, &delta, &a, &b, &c, &mut out, 2, 1, 1, 1, false);
        assert!(r.is_err());
    }

    /// SsdChunkScan: minimal degenerate case — batch=heads=head_dim=
    /// state_dim=1, seqlen=chunk_size=1. Same scale as SelectiveScan's
    /// single-step test.
    ///   d = dt[0,0,0] = 1.0
    ///   a_h = a[0] = -1.0
    ///   h[0,0,0,0] = exp(1*-1)*0 + 1*2*3 = 6.0   (b=2, x=3)
    ///   y[0,0,0,0] = 6.0 * 0.5 = 3.0             (c=0.5)
    #[test]
    fn ssd_chunk_scan_f32_minimal() {
        let x = CpuStorageBytes::from_slice(&[3.0_f32]);
        let dt = CpuStorageBytes::from_slice(&[1.0_f32]);
        let a = CpuStorageBytes::from_slice(&[-1.0_f32]);
        let b = CpuStorageBytes::from_slice(&[2.0_f32]);
        let c = CpuStorageBytes::from_slice(&[0.5_f32]);
        let mut out = CpuStorageBytes::from_zero_bytes(8);
        ssd_chunk_scan_f32(&x, &dt, &a, &b, &c, &mut out, 1, 1, 1, 1, 1, 1)
            .expect("ssd_chunk_scan");
        let result: &[f32] = out.as_slice().unwrap();
        assert!((result[0] - 3.0).abs() < 1e-5, "got {}", result[0]);
    }

    /// SsdChunkScan: two timesteps verify state threading per head.
    /// Same trick as SelectiveScan's two-step test with the extra
    /// heads dimension fixed at 1.
    #[test]
    fn ssd_chunk_scan_f32_two_steps_state_threading() {
        // [batch=1, seqlen=2, heads=1, head_dim=1]
        let x = CpuStorageBytes::from_slice(&[3.0_f32, 3.0]);
        let dt = CpuStorageBytes::from_slice(&[1.0_f32, 1.0]);
        let a = CpuStorageBytes::from_slice(&[-1.0_f32]);
        let b = CpuStorageBytes::from_slice(&[2.0_f32, 2.0]);
        let c = CpuStorageBytes::from_slice(&[0.5_f32, 0.5]);
        let mut out = CpuStorageBytes::from_zero_bytes(12);
        ssd_chunk_scan_f32(&x, &dt, &a, &b, &c, &mut out, 1, 2, 1, 1, 1, 2)
            .expect("ssd_chunk_scan");
        let result: &[f32] = out.as_slice().unwrap();
        let h_step1 = (-1.0_f32).exp() * 6.0 + 6.0;
        let y_step1 = h_step1 * 0.5;
        assert!((result[0] - 3.0).abs() < 1e-5, "step 0: {}", result[0]);
        assert!(
            (result[1] - y_step1).abs() < 1e-5,
            "step 1: {} expected {y_step1}",
            result[1]
        );
    }

    /// SsdChunkScan: two heads with distinct per-head `a` values
    /// verifies the per-head log-A is applied correctly to each head.
    #[test]
    fn ssd_chunk_scan_f32_two_heads_distinct_a() {
        // [batch=1, seqlen=1, heads=2, head_dim=1]
        let x = CpuStorageBytes::from_slice(&[3.0_f32, 5.0]);
        let dt = CpuStorageBytes::from_slice(&[1.0_f32, 1.0]);
        let a = CpuStorageBytes::from_slice(&[-1.0_f32, -2.0]);
        let b = CpuStorageBytes::from_slice(&[2.0_f32, 4.0]);
        let c = CpuStorageBytes::from_slice(&[0.5_f32, 1.0]);
        let mut out = CpuStorageBytes::from_zero_bytes(16);
        ssd_chunk_scan_f32(&x, &dt, &a, &b, &c, &mut out, 1, 1, 2, 1, 1, 1)
            .expect("ssd_chunk_scan");
        let result: &[f32] = out.as_slice().unwrap();
        // head 0: h = exp(-1)*0 + 1*2*3 = 6; y = 6 * 0.5 = 3
        // head 1: h = exp(-2)*0 + 1*4*5 = 20; y = 20 * 1.0 = 20
        assert!((result[0] - 3.0).abs() < 1e-5, "head 0: {}", result[0]);
        assert!((result[1] - 20.0).abs() < 1e-5, "head 1: {}", result[1]);
    }

    /// SsdChunkScan: multi-chunk request (chunk_size < seqlen)
    /// produces the same result as single-chunk. Mathematical
    /// equivalence: the sequential scan answer doesn't depend on
    /// chunk_size; it's a GPU-parallelism knob. Same numbers as
    /// `ssd_chunk_scan_f32_two_steps_state_threading` but called
    /// with `chunk_size = 1` (instead of 2) — must give identical
    /// output.
    #[test]
    fn ssd_chunk_scan_f32_chunk_size_is_invariant() {
        let x = CpuStorageBytes::from_slice(&[3.0_f32, 3.0]);
        let dt = CpuStorageBytes::from_slice(&[1.0_f32, 1.0]);
        let a = CpuStorageBytes::from_slice(&[-1.0_f32]);
        let b = CpuStorageBytes::from_slice(&[2.0_f32, 2.0]);
        let c = CpuStorageBytes::from_slice(&[0.5_f32, 0.5]);
        // chunk_size = 1 (two chunks of 1 token each)
        let mut out_cs1 = CpuStorageBytes::from_zero_bytes(12);
        ssd_chunk_scan_f32(&x, &dt, &a, &b, &c, &mut out_cs1, 1, 2, 1, 1, 1, 1)
            .expect("ssd_chunk_scan cs=1");
        // chunk_size = 2 (single chunk of 2 tokens)
        let mut out_cs2 = CpuStorageBytes::from_zero_bytes(12);
        ssd_chunk_scan_f32(&x, &dt, &a, &b, &c, &mut out_cs2, 1, 2, 1, 1, 1, 2)
            .expect("ssd_chunk_scan cs=2");
        let r1: &[f32] = out_cs1.as_slice().unwrap();
        let r2: &[f32] = out_cs2.as_slice().unwrap();
        assert_eq!(r1, r2, "chunk_size should not affect the math");
        // Verify the expected value (same as two_steps_state_threading).
        assert!((r1[0] - 3.0).abs() < 1e-5);
        let h_step1 = (-1.0_f32).exp() * 6.0 + 6.0;
        assert!((r1[1] - h_step1 * 0.5).abs() < 1e-5);
    }

    /// SsdChunkScan: chunk_size that doesn't divide seqlen errors.
    #[test]
    fn ssd_chunk_scan_f32_chunk_size_must_divide_seqlen() {
        let x = CpuStorageBytes::from_slice(&[1.0_f32; 4]);
        let dt = CpuStorageBytes::from_slice(&[1.0_f32; 4]);
        let a = CpuStorageBytes::from_slice(&[1.0_f32]);
        let b = CpuStorageBytes::from_slice(&[1.0_f32; 4]);
        let c = CpuStorageBytes::from_slice(&[1.0_f32; 4]);
        let mut out = CpuStorageBytes::from_zero_bytes(20);
        // seqlen=4, chunk_size=3 → 4 % 3 != 0 → error
        let r = ssd_chunk_scan_f32(&x, &dt, &a, &b, &c, &mut out, 1, 4, 1, 1, 1, 3);
        assert!(r.is_err());
    }

    /// SsdChunkScan: chunk_size = 0 errors.
    #[test]
    fn ssd_chunk_scan_f32_chunk_size_zero_errors() {
        let x = CpuStorageBytes::from_slice(&[1.0_f32]);
        let dt = CpuStorageBytes::from_slice(&[1.0_f32]);
        let a = CpuStorageBytes::from_slice(&[1.0_f32]);
        let b = CpuStorageBytes::from_slice(&[1.0_f32]);
        let c = CpuStorageBytes::from_slice(&[1.0_f32]);
        let mut out = CpuStorageBytes::from_zero_bytes(8);
        let r = ssd_chunk_scan_f32(&x, &dt, &a, &b, &c, &mut out, 1, 1, 1, 1, 1, 0);
        assert!(r.is_err());
    }

    /// Nf4Matmul: hand-computed F32 sanity case. Single output
    /// position, k=2 (one packed byte), block_size=2 (one absmax
    /// scale). Verifies the LUT lookup + dequant + dot product math.
    ///
    ///   activations = [a0, a1] = [1.0, 2.0]
    ///   w_packed = [byte] where byte = code1<<4 | code0
    ///     pick code0=8 (NF4_LUT[8] ≈ 0.0796), code1=15 (LUT[15] = 1.0)
    ///     byte = (15<<4) | 8 = 248
    ///   absmax = [3.0]  (single block, scale = 3.0)
    ///   dequant_w[0] = 0.0796 * 3.0 ≈ 0.2387
    ///   dequant_w[1] = 1.0 * 3.0 = 3.0
    ///   y = 1.0 * 0.2387 + 2.0 * 3.0 ≈ 6.2387
    #[test]
    fn nf4_matmul_f32_hand_computed() {
        let activations = CpuStorageBytes::from_slice(&[1.0_f32, 2.0]); // [1,1,2]
        let w_packed = CpuStorageBytes::from_slice(&[248_u8]); // [1, 1]
        let absmax = CpuStorageBytes::from_slice(&[3.0_f32]); // [1, 1]
        let mut out = CpuStorageBytes::from_zero_bytes(4);
        nf4_matmul_f32(&activations, &w_packed, &absmax, &mut out, 1, 1, 1, 2, 2)
            .expect("nf4_matmul_f32");
        let result: &[f32] = out.as_slice().unwrap();
        let expected = 1.0 * NF4_LUT[8] * 3.0 + 2.0 * NF4_LUT[15] * 3.0;
        assert!(
            (result[0] - expected).abs() < 1e-5,
            "got {} expected {expected}",
            result[0]
        );
    }

    /// Nf4Matmul: same math as the F32 test, F16 activations/output.
    /// F32 accumulator path means precision matches modulo F16 rounding.
    #[test]
    fn nf4_matmul_f16_hand_computed() {
        let a_f16: Vec<half::f16> = [1.0_f32, 2.0]
            .iter()
            .map(|&v| half::f16::from_f32(v))
            .collect();
        let activations = CpuStorageBytes::from_slice(&a_f16);
        let w_packed = CpuStorageBytes::from_slice(&[248_u8]);
        let absmax = CpuStorageBytes::from_slice(&[3.0_f32]);
        let mut out = CpuStorageBytes::from_zero_bytes(2);
        nf4_matmul_f16(&activations, &w_packed, &absmax, &mut out, 1, 1, 1, 2, 2)
            .expect("nf4_matmul_f16");
        let result: &[half::f16] = out.as_slice().unwrap();
        let expected = 1.0_f32 * NF4_LUT[8] * 3.0 + 2.0 * NF4_LUT[15] * 3.0;
        assert!(
            (result[0].to_f32() - expected).abs() < 1e-2,
            "got {} expected {expected}",
            result[0].to_f32(),
        );
    }

    /// Nf4Matmul: BF16 activations/output, same math.
    #[test]
    fn nf4_matmul_bf16_hand_computed() {
        let a_bf16: Vec<half::bf16> = [1.0_f32, 2.0]
            .iter()
            .map(|&v| half::bf16::from_f32(v))
            .collect();
        let activations = CpuStorageBytes::from_slice(&a_bf16);
        let w_packed = CpuStorageBytes::from_slice(&[248_u8]);
        let absmax = CpuStorageBytes::from_slice(&[3.0_f32]);
        let mut out = CpuStorageBytes::from_zero_bytes(2);
        nf4_matmul_bf16(&activations, &w_packed, &absmax, &mut out, 1, 1, 1, 2, 2)
            .expect("nf4_matmul_bf16");
        let result: &[half::bf16] = out.as_slice().unwrap();
        let expected = 1.0_f32 * NF4_LUT[8] * 3.0 + 2.0 * NF4_LUT[15] * 3.0;
        assert!(
            (result[0].to_f32() - expected).abs() < 5e-2, // BF16 has only 7 mantissa bits
            "got {} expected {expected}",
            result[0].to_f32(),
        );
    }

    /// Nf4Matmul: m=1, n=2, k=4, block_size=2 — verifies per-output-
    /// row distinct weights + per-block absmax scaling.
    #[test]
    fn nf4_matmul_f32_two_outputs_two_blocks() {
        // activations [1, 2, 2, 4]: m=1, k=4
        let activations = CpuStorageBytes::from_slice(&[1.0_f32, 2.0, 2.0, 4.0]);
        // w_packed: [n=2, k/2=2] = 4 bytes
        // Output 0: codes [7, 15, 7, 15] (LUT: 0.0, 1.0, 0.0, 1.0)
        //   bytes: lo=7|hi=15 = (15<<4)|7 = 247; same for second byte
        // Output 1: codes [15, 7, 15, 7] (LUT: 1.0, 0.0, 1.0, 0.0)
        //   bytes: lo=15|hi=7 = (7<<4)|15 = 127
        let w_packed = CpuStorageBytes::from_slice(&[247_u8, 247, 127, 127]);
        // absmax: [n=2, k/block_size=2] = 4 floats
        // Output 0: [1.0, 2.0] (block 0 scale 1.0, block 1 scale 2.0)
        // Output 1: [10.0, 20.0]
        let absmax = CpuStorageBytes::from_slice(&[1.0_f32, 2.0, 10.0, 20.0]);
        let mut out = CpuStorageBytes::from_zero_bytes(2 * 4);
        nf4_matmul_f32(&activations, &w_packed, &absmax, &mut out, 1, 1, 2, 4, 2)
            .expect("nf4_matmul_f32");
        let result: &[f32] = out.as_slice().unwrap();
        // Output 0:
        //   k=0: code=7, scale=1.0, dequant=0.0  → contrib = 1.0 * 0.0 = 0
        //   k=1: code=15, scale=1.0, dequant=1.0 → contrib = 2.0 * 1.0 = 2
        //   k=2: code=7, scale=2.0, dequant=0.0  → contrib = 2.0 * 0.0 = 0
        //   k=3: code=15, scale=2.0, dequant=2.0 → contrib = 4.0 * 2.0 = 8
        //   y[0] = 10.0
        assert!(
            (result[0] - 10.0).abs() < 1e-5,
            "output 0: got {}",
            result[0]
        );
        // Output 1:
        //   k=0: code=15, scale=10.0, dequant=10.0 → 1.0 * 10.0 = 10
        //   k=1: code=7, scale=10.0, dequant=0.0   → 2.0 * 0.0 = 0
        //   k=2: code=15, scale=20.0, dequant=20.0 → 2.0 * 20.0 = 40
        //   k=3: code=7, scale=20.0, dequant=0.0   → 4.0 * 0.0 = 0
        //   y[1] = 50.0
        assert!(
            (result[1] - 50.0).abs() < 1e-5,
            "output 1: got {}",
            result[1]
        );
    }

    /// Nf4Matmul: validates the NF4 LUT contents — first 4 values
    /// should match the standard NormalFloat curve.
    #[test]
    fn nf4_lut_known_values() {
        assert_eq!(NF4_LUT.len(), 16);
        assert!((NF4_LUT[0] - (-1.0)).abs() < 1e-6, "LUT[0] should be -1.0");
        assert!((NF4_LUT[7] - 0.0).abs() < 1e-6, "LUT[7] should be 0.0");
        assert!((NF4_LUT[15] - 1.0).abs() < 1e-6, "LUT[15] should be 1.0");
        // Values should be monotonically increasing.
        for i in 0..15 {
            assert!(
                NF4_LUT[i] < NF4_LUT[i + 1],
                "LUT not monotonic at index {i}"
            );
        }
    }

    /// Nf4Matmul: bad shape errors — k odd, k not multiple of
    /// block_size, byte-count mismatches.
    #[test]
    fn nf4_matmul_f32_bad_shapes_error() {
        let activations = CpuStorageBytes::from_slice(&[1.0_f32; 3]);
        let w_packed = CpuStorageBytes::from_slice(&[0_u8; 2]);
        let absmax = CpuStorageBytes::from_slice(&[1.0_f32]);
        let mut out = CpuStorageBytes::from_zero_bytes(4);
        // k=3 is odd
        let r = nf4_matmul_f32(&activations, &w_packed, &absmax, &mut out, 1, 1, 1, 3, 1);
        assert!(r.is_err());

        // k=4 but block_size=3 — 4 % 3 != 0
        let r = nf4_matmul_f32(&activations, &w_packed, &absmax, &mut out, 1, 1, 1, 4, 3);
        assert!(r.is_err());
    }

    /// SsdChunkScan F64 sanity — mirrors the F32 minimal test.
    #[test]
    fn ssd_chunk_scan_f64_minimal() {
        let x = CpuStorageBytes::from_slice(&[3.0_f64]);
        let dt = CpuStorageBytes::from_slice(&[1.0_f64]);
        let a = CpuStorageBytes::from_slice(&[-1.0_f64]);
        let b = CpuStorageBytes::from_slice(&[2.0_f64]);
        let c = CpuStorageBytes::from_slice(&[0.5_f64]);
        let mut out = CpuStorageBytes::from_zero_bytes(16);
        ssd_chunk_scan_f64(&x, &dt, &a, &b, &c, &mut out, 1, 1, 1, 1, 1, 1)
            .expect("ssd_chunk_scan_f64");
        let r: &[f64] = out.as_slice().unwrap();
        assert!((r[0] - 3.0).abs() < 1e-12);
    }

    /// SsdChunkScan BF16 sanity. Loose tolerance for 7-bit mantissa.
    #[test]
    fn ssd_chunk_scan_bf16_minimal() {
        let x = CpuStorageBytes::from_slice(&[half::bf16::from_f32(3.0)]);
        let dt = CpuStorageBytes::from_slice(&[half::bf16::from_f32(1.0)]);
        let a = CpuStorageBytes::from_slice(&[half::bf16::from_f32(-1.0)]);
        let b = CpuStorageBytes::from_slice(&[half::bf16::from_f32(2.0)]);
        let c = CpuStorageBytes::from_slice(&[half::bf16::from_f32(0.5)]);
        let mut out = CpuStorageBytes::from_zero_bytes(4);
        ssd_chunk_scan_bf16(&x, &dt, &a, &b, &c, &mut out, 1, 1, 1, 1, 1, 1)
            .expect("ssd_chunk_scan_bf16");
        let r: &[half::bf16] = out.as_slice().unwrap();
        assert!((r[0].to_f32() - 3.0).abs() < 5e-2);
    }

    /// SsdChunkScan F16 sanity. 10-bit mantissa.
    #[test]
    fn ssd_chunk_scan_f16_minimal() {
        let x = CpuStorageBytes::from_slice(&[half::f16::from_f32(3.0)]);
        let dt = CpuStorageBytes::from_slice(&[half::f16::from_f32(1.0)]);
        let a = CpuStorageBytes::from_slice(&[half::f16::from_f32(-1.0)]);
        let b = CpuStorageBytes::from_slice(&[half::f16::from_f32(2.0)]);
        let c = CpuStorageBytes::from_slice(&[half::f16::from_f32(0.5)]);
        let mut out = CpuStorageBytes::from_zero_bytes(4);
        ssd_chunk_scan_f16(&x, &dt, &a, &b, &c, &mut out, 1, 1, 1, 1, 1, 1)
            .expect("ssd_chunk_scan_f16");
        let r: &[half::f16] = out.as_slice().unwrap();
        assert!((r[0].to_f32() - 3.0).abs() < 1e-2);
    }

    /// SsdChunkScan: bad shapes error rather than panicking.
    #[test]
    fn ssd_chunk_scan_f32_bad_shapes_error() {
        let x = CpuStorageBytes::from_slice(&[1.0_f32]);
        let dt = CpuStorageBytes::from_slice(&[1.0_f32]);
        let a = CpuStorageBytes::from_slice(&[1.0_f32]);
        let b = CpuStorageBytes::from_slice(&[1.0_f32]);
        let c = CpuStorageBytes::from_slice(&[1.0_f32]);
        let mut out = CpuStorageBytes::from_zero_bytes(16);
        // heads=2 but only 1-element a → byte count mismatch
        let r = ssd_chunk_scan_f32(&x, &dt, &a, &b, &c, &mut out, 1, 1, 2, 1, 1, 1);
        assert!(r.is_err());
    }

    // -----------------------------------------------------------------
    // In-place unary dtype expansion — one representative (op, dtype)
    // pair per non-f32 dtype the chassis serves. All 20 (op × dtype)
    // entries collapse onto the same `unary_inplace_thunk!` macro that
    // delegates to `UnaryOp<T>::apply`, so a per-dtype smoke is
    // sufficient at this layer (executor-level coverage lives in
    // fuel-storage's `op_relu_inplace_cpu_mutates_target_storage`).
    // -----------------------------------------------------------------

    #[test]
    fn relu_inplace_f64_round_trip() {
        let mut out = CpuStorageBytes::from_slice(&[-1.0_f64, 0.0, 1.0, 2.0]);
        relu_inplace_f64(&mut out).expect("relu inplace f64");
        assert_eq!(out.as_slice::<f64>().unwrap(), &[0.0, 0.0, 1.0, 2.0]);
    }

    #[test]
    fn silu_inplace_f64_matches_functional() {
        let mut got = CpuStorageBytes::from_slice(&[1.0_f64, -2.0, 0.5]);
        silu_inplace_f64(&mut got).expect("silu inplace f64");
        let r = got.as_slice::<f64>().unwrap();
        // SiLU(x) = x · sigmoid(x); compare against direct math.
        for (i, &x) in [1.0_f64, -2.0, 0.5].iter().enumerate() {
            let want = x / (1.0 + (-x).exp());
            assert!((r[i] - want).abs() < 1e-12, "got {} want {}", r[i], want);
        }
    }

    #[test]
    fn relu_inplace_bf16_round_trip() {
        use half::bf16;
        let mut out = CpuStorageBytes::from_slice(&[
            bf16::from_f32(-1.0),
            bf16::from_f32(0.0),
            bf16::from_f32(1.0),
            bf16::from_f32(2.0),
        ]);
        relu_inplace_bf16(&mut out).expect("relu inplace bf16");
        let r = out.as_slice::<bf16>().unwrap();
        assert_eq!(r[0], bf16::from_f32(0.0));
        assert_eq!(r[1], bf16::from_f32(0.0));
        assert_eq!(r[2], bf16::from_f32(1.0));
        assert_eq!(r[3], bf16::from_f32(2.0));
    }

    #[test]
    fn tanh_inplace_bf16_via_f32_pivot() {
        use half::bf16;
        let mut got = CpuStorageBytes::from_slice(&[
            bf16::from_f32(0.0),
            bf16::from_f32(1.0),
            bf16::from_f32(-1.0),
        ]);
        tanh_inplace_bf16(&mut got).expect("tanh inplace bf16");
        let r = got.as_slice::<bf16>().unwrap();
        // bf16 round-trip through f32 loses a lot of precision; check
        // each slot matches the expected bf16(tanh(f32(x))) bitwise.
        for (i, &x) in [0.0_f32, 1.0, -1.0].iter().enumerate() {
            assert_eq!(r[i], bf16::from_f32(x.tanh()));
        }
    }

    #[test]
    fn sigmoid_inplace_f16_via_f32_pivot() {
        use half::f16;
        let mut got = CpuStorageBytes::from_slice(&[
            f16::from_f32(0.0),
            f16::from_f32(2.0),
            f16::from_f32(-2.0),
        ]);
        sigmoid_inplace_f16(&mut got).expect("sigmoid inplace f16");
        let r = got.as_slice::<f16>().unwrap();
        for (i, &x) in [0.0_f32, 2.0, -2.0].iter().enumerate() {
            let want = f16::from_f32(1.0 / (1.0 + (-x).exp()));
            assert_eq!(r[i], want);
        }
    }

    // ---- Expanded in-place op family (2026-05-30) ----

    #[test]
    fn neg_inplace_f32_round_trip() {
        let mut out = CpuStorageBytes::from_slice(&[1.0_f32, -2.0, 3.0, -4.0]);
        neg_inplace_f32(&mut out).expect("neg inplace f32");
        assert_eq!(out.as_slice::<f32>().unwrap(), &[-1.0, 2.0, -3.0, 4.0]);
    }

    #[test]
    fn sqr_inplace_f64_round_trip() {
        let mut out = CpuStorageBytes::from_slice(&[2.0_f64, -3.0, 4.0]);
        sqr_inplace_f64(&mut out).expect("sqr inplace f64");
        assert_eq!(out.as_slice::<f64>().unwrap(), &[4.0, 9.0, 16.0]);
    }

    #[test]
    fn exp_inplace_bf16_via_f32_pivot() {
        use half::bf16;
        let mut out = CpuStorageBytes::from_slice(&[bf16::from_f32(0.0), bf16::from_f32(1.0)]);
        exp_inplace_bf16(&mut out).expect("exp inplace bf16");
        let r = out.as_slice::<bf16>().unwrap();
        assert!((r[0].to_f32() - 1.0).abs() < 1e-2);
        assert!((r[1].to_f32() - std::f32::consts::E).abs() < 1e-2);
    }

    #[test]
    fn floor_inplace_f32_round_trip() {
        let mut out = CpuStorageBytes::from_slice(&[1.7_f32, -2.3, 3.0]);
        floor_inplace_f32(&mut out).expect("floor inplace f32");
        assert_eq!(out.as_slice::<f32>().unwrap(), &[1.0, -3.0, 3.0]);
    }

    #[test]
    fn clamp_inplace_f32_round_trip() {
        let mut out = CpuStorageBytes::from_slice(&[-5.0_f32, 0.0, 5.0, 2.5]);
        clamp_inplace_f32(&mut out, -1.0, 1.0).expect("clamp inplace f32");
        assert_eq!(out.as_slice::<f32>().unwrap(), &[-1.0, 0.0, 1.0, 1.0]);
    }

    #[test]
    fn clamp_inplace_f32_rejects_bad_bounds() {
        let mut out = CpuStorageBytes::from_slice(&[1.0_f32]);
        assert!(
            clamp_inplace_f32(&mut out, 2.0, 1.0).is_err(),
            "min > max must error"
        );
    }

    #[test]
    fn powi_inplace_f64_round_trip() {
        let mut out = CpuStorageBytes::from_slice(&[2.0_f64, -3.0, 4.0]);
        powi_inplace_f64(&mut out, 3).expect("powi inplace f64");
        assert_eq!(out.as_slice::<f64>().unwrap(), &[8.0, -27.0, 64.0]);
    }

    #[test]
    fn powi_inplace_f16_via_f32_pivot() {
        use half::f16;
        let mut out = CpuStorageBytes::from_slice(&[f16::from_f32(2.0), f16::from_f32(3.0)]);
        powi_inplace_f16(&mut out, 2).expect("powi inplace f16");
        let r = out.as_slice::<f16>().unwrap();
        assert_eq!(r[0], f16::from_f32(4.0));
        assert_eq!(r[1], f16::from_f32(9.0));
    }

    /// The dropless-MoE atom: `matmul_f32_capacity` computes only
    /// `m_compute` rows of an `m_capacity`-row buffer, leaving the tail
    /// zeroed. LHS is a capacity buffer `[m_capacity, k]` whose first
    /// `m_compute` rows are real and whose tail is garbage that must NOT
    /// affect the output.
    #[test]
    fn matmul_f32_capacity_computes_only_m_compute_rows() {
        // rhs [k=2, n=3]
        let rhs = CpuStorageBytes::from_slice(&[1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0]);
        // lhs capacity [m_capacity=4, k=2]: rows 0,1 real; rows 2,3 garbage
        // (huge values that would blow up the output if computed).
        let lhs = CpuStorageBytes::from_slice(&[
            1.0_f32, 0.0, // row 0
            0.0, 1.0, // row 1
            9.9e30, 9.9e30, // row 2 (garbage — must be ignored)
            9.9e30, 9.9e30, // row 3 (garbage)
        ]);
        let mut out = CpuStorageBytes::from_slice(&[7.0_f32; 4 * 3]); // pre-dirtied
        matmul_f32_capacity(&lhs, &rhs, &mut out, &[], &[], 4, 2, 3, 2).expect("capacity matmul");
        let got = out.as_slice::<f32>().unwrap();
        // row 0 = rhs row 0 = [1,2,3]; row 1 = rhs row 1 = [4,5,6];
        // rows 2,3 = zero (not computed — the tail is zero-filled, NOT the
        // garbage product).
        assert_eq!(
            got,
            &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
            "only the first m_compute=2 rows computed; tail zeroed",
        );
    }

    /// `matmul_f32` (the static delegate) is byte-identical to computing
    /// every row via the capacity kernel with `m_compute == m_capacity`.
    #[test]
    fn matmul_f32_delegates_to_capacity_all_rows() {
        let lhs = CpuStorageBytes::from_slice(&[1.0_f32, 2.0, 3.0, 4.0]); // [2,2]
        let rhs = CpuStorageBytes::from_slice(&[5.0_f32, 6.0, 7.0, 8.0]); // [2,2]
        let mut out = CpuStorageBytes::from_slice(&[0.0_f32; 4]);
        matmul_f32(&lhs, &rhs, &mut out, &[], &[], 2, 2, 2).expect("static matmul");
        // [[1,2],[3,4]] @ [[5,6],[7,8]] = [[19,22],[43,50]]
        assert_eq!(out.as_slice::<f32>().unwrap(), &[19.0, 22.0, 43.0, 50.0]);
    }
}