himada-dispatch 0.1.1

use std::collections::VecDeque;
use std::path::Path;
use std::time::{Duration, Instant};

pub use himada_core::HardwareDNA;
pub use himada_macros::{himada_adapt, himada_optimize};
use himada_core::profile;
use serde::{Deserialize, Serialize};

pub mod adapt;
pub mod alloc;
pub mod ffi;
pub mod gpu;
pub mod graph;
pub mod jit;
pub mod kernels;
pub mod pad;
pub mod plugin;
pub mod plugin_macros;
pub mod precision;
pub mod shape;
pub mod trace;
#[cfg(feature = "vulkan")]
pub mod vulkan;

/// Benchmark harness for a kernel function type `F`.
pub trait FnBench: Copy + Send + Sync + 'static {
    /// Full benchmark — pick the fastest kernel from `candidates`.
    fn bench(kernels: &[KernelInfo<Self>], candidates: &[usize]) -> usize;
    /// Quick timing of a single kernel (small data, few iterations).
    fn time_func(kernels: &[KernelInfo<Self>], idx: usize, iters: u32) -> Duration;
}

/// Errors that can occur during dispatch operations.
#[derive(Debug)]
pub enum HwdnaError {
    /// No kernels are supported on this hardware.
    NoSupportedKernels,
    /// Mutex poisoned (another thread panicked while holding the lock).
    MutexPoisoned,
}

impl std::fmt::Display for HwdnaError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            HwdnaError::NoSupportedKernels => write!(f, "no supported kernels available"),
            HwdnaError::MutexPoisoned => write!(f, "mutex poisoned"),
        }
    }
}

impl std::error::Error for HwdnaError {}

/// Result type for Himada operations.
pub type HwdnaResult<T> = Result<T, HwdnaError>;

/// Ring buffer of timing observations for online learning.
#[derive(Clone, Serialize, Deserialize)]
pub struct TimingHistory {
    samples: VecDeque<Duration>,
    max_len: usize,
}

impl TimingHistory {
    pub fn new(max_len: usize) -> Self {
        Self { samples: VecDeque::with_capacity(max_len), max_len }
    }

    pub fn record(&mut self, elapsed: Duration) {
        if self.samples.len() >= self.max_len {
            self.samples.pop_front();
        }
        self.samples.push_back(elapsed);
    }

    pub fn avg(&self) -> Duration {
        if self.samples.is_empty() {
            return Duration::MAX;
        }
        let sum: Duration = self.samples.iter().sum();
        sum / self.samples.len() as u32
    }

    pub fn min(&self) -> Duration {
        self.samples.iter().copied().min().unwrap_or(Duration::MAX)
    }

    /// Trend: positive means slowing down, negative means speeding up.
    /// Returns a fraction of the average: e.g., 0.1 means 10% slower.
    pub fn trend(&self) -> f64 {
        if self.samples.len() < 4 {
            return 0.0;
        }
        // Exponential Moving Average: weight recent observations higher
        // EMA_t = alpha * value_t + (1-alpha) * EMA_{t-1}
        let alpha = 0.3;
        let samples: Vec<f64> = self.samples.iter().map(|d| d.as_secs_f64()).collect();
        let mut ema = samples[0];
        for &s in &samples[1..] {
            ema = alpha * s + (1.0 - alpha) * ema;
        }
        // Trend = (EMA - overall_avg) / overall_avg
        let avg: f64 = samples.iter().sum::<f64>() / samples.len() as f64;
        if avg == 0.0 { return 0.0; }
        (ema - avg) / avg
    }

    pub fn len(&self) -> usize {
        self.samples.len()
    }

    pub fn is_empty(&self) -> bool {
        self.samples.is_empty()
    }
}

/// A type-erased dispatch step for pipeline optimization.
pub trait Step: Send + Sync {
    fn op_name(&self) -> &str;
    fn select_step(&mut self) -> Result<(), HwdnaError>;
    fn selected_kernel_name(&self) -> String;
    fn observe_step(&mut self, elapsed: Duration);
    fn rebench_step(&mut self) -> Result<usize, HwdnaError>;
}

/// Generate `try_compute(…) -> Result`, `compute(…) -> T`, and `compute_many(…) -> Vec<T>` for Dispatch<T>.
macro_rules! make_dispatch_compute {
    ($kernel:ty, ($($arg:ident: $argty:ty),*) -> $ret:ty) => {
        impl Dispatch<$kernel> {
            pub fn try_compute(&mut self, $($arg: $argty),*) -> Result<$ret, HwdnaError> {
                let idx = self.select()?;
                Ok((self.supported[idx].func)($($arg),*))
            }
            #[track_caller]
            pub fn compute(&mut self, $($arg: $argty),*) -> $ret {
                self.try_compute($($arg),*)
                    .expect("Dispatch::compute: no supported kernels")
            }
        }
    };
    // With batch — single arg slice
    ($kernel:ty, ($arg:ident: $argty:ty) -> $ret:ty, batch = single) => {
        impl Dispatch<$kernel> {
            pub fn try_compute(&mut self, $arg: $argty) -> Result<$ret, HwdnaError> {
                let idx = self.select()?;
                Ok((self.supported[idx].func)($arg))
            }
            #[track_caller]
            pub fn compute(&mut self, $arg: $argty) -> $ret {
                self.try_compute($arg)
                    .expect("Dispatch::compute: no supported kernels")
            }
            pub fn compute_many(&mut self, inputs: &[$argty]) -> Vec<$ret> {
                let idx = self.select().expect("compute_many: no supported kernels");
                let f = self.supported[idx].func;
                inputs.iter().map(|a| f(*a)).collect()
            }
        }
    };
    // With batch — two arg slice
    ($kernel:ty, ($a:ident: $aty:ty, $b:ident: $bty:ty) -> $ret:ty, batch = pair) => {
        impl Dispatch<$kernel> {
            pub fn try_compute(&mut self, $a: $aty, $b: $bty) -> Result<$ret, HwdnaError> {
                let idx = self.select()?;
                Ok((self.supported[idx].func)($a, $b))
            }
            #[track_caller]
            pub fn compute(&mut self, $a: $aty, $b: $bty) -> $ret {
                self.try_compute($a, $b)
                    .expect("Dispatch::compute: no supported kernels")
            }
            pub fn compute_many(&mut self, inputs: &[($aty, $bty)]) -> Vec<$ret> {
                let idx = self.select().expect("compute_many: no supported kernels");
                let f = self.supported[idx].func;
                inputs.iter().map(|&(a, b)| f(a, b)).collect()
            }
        }
    };
}

#[derive(Clone)]
pub struct KernelInfo<F> {
    pub name: &'static str,
    pub func: F,
    pub is_supported: fn(&HardwareDNA) -> bool,
    pub thermal_priority: u8,
}

pub struct Dispatch<F> {
    pub(crate) name: String,
    pub(crate) supported: Vec<KernelInfo<F>>,
    pub(crate) selected: Option<usize>,
    pub(crate) dna: HardwareDNA,
    pub(crate) timing_history: TimingHistory,
    pub(crate) observation_count: u64,
}

impl<F> Dispatch<F> {
    pub fn new(name: &str, kernels: Vec<KernelInfo<F>>) -> Self {
        let dna = profile::load_or_collect();
        Self {
            name: name.into(),
            supported: kernels,
            selected: None,
            dna,
            timing_history: TimingHistory::new(64),
            observation_count: 0,
        }
    }

    pub fn selected_name(&self) -> &str {
        self.selected
            .map(|i| self.supported[i].name)
            .unwrap_or("unselected")
    }

    pub fn available_indices(&self) -> Vec<usize> {
        (0..self.supported.len())
            .filter(|&i| (self.supported[i].is_supported)(&self.dna))
            .collect()
    }

    pub fn to_json(&self) -> serde_json::Value {
        serde_json::json!({
            "name": self.name,
            "selected": self.selected_name(),
            "observations": self.observation_count,
        })
    }
}

impl<F: FnBench> Dispatch<F> {
    pub fn select(&mut self) -> Result<usize, HwdnaError> {
        if let Some(idx) = self.selected {
            return Ok(idx);
        }
        // Build-time dispatch (Cargo features force-scalar/sse/avx2/neon)
        if let Some(idx) = self.try_build_time_select() {
            self.selected = Some(idx);
            return Ok(idx);
        }
        // Auto-tuning cache — skip benchmark if we already know the best kernel
        #[cfg(not(any(test, debug_assertions)))]
        {
            let dispatch_name = &self.name;
            if let Some(cached_name) = crate::alloc::load_cached_selection(dispatch_name) {
                for (idx, k) in self.supported.iter().enumerate() {
                    if k.name == cached_name && (k.is_supported)(&self.dna) {
                        self.selected = Some(idx);
                        return Ok(idx);
                    }
                }
            }
        }
        let available: Vec<usize> = self.available_indices();
        if available.is_empty() {
            return Err(HwdnaError::NoSupportedKernels);
        }
        let idx = if available.len() == 1 {
            available[0]
        } else {
            F::bench(&self.supported, &available)
        };
        #[cfg(not(any(
            himada_force_scalar,
            himada_force_sse,
            himada_force_avx2,
            himada_force_neon,
        )))]
        eprintln!("[himada::{}] selected: {}", self.name, self.supported[idx].name);
        self.selected = Some(idx);
        #[cfg(not(any(test, debug_assertions)))]
        {
            crate::alloc::save_cached_selection(&self.name, self.supported[idx].name);
        }
        Ok(idx)
    }

    pub fn reselect(&mut self) -> Result<usize, HwdnaError> {
        self.selected = None;
        self.select()
    }

    /// Record an observation of actual compute call duration.
    /// Used by online learning to track real-world performance.
    pub fn observe(&mut self, elapsed: Duration) {
        self.timing_history.record(elapsed);
        self.observation_count += 1;
    }

    /// Force re-benchmark all kernels and re-select the best.
    /// Returns the index of the selected kernel.
    pub fn rebench(&mut self) -> Result<usize, HwdnaError> {
        self.selected = None;
        self.select()
    }

    /// Trend of timing observations: positive = slowing down, negative = speeding up.
    pub fn timing_trend(&self) -> f64 {
        self.timing_history.trend()
    }

    /// Average observed execution time.
    pub fn timing_avg(&self) -> Duration {
        self.timing_history.avg()
    }
}

// SAFETY: All fields are Send + Sync (F: FnBench, Vec, String, Option<usize>, HardwareDNA).
unsafe impl<F: FnBench> Send for Dispatch<F> {}
unsafe impl<F: FnBench> Sync for Dispatch<F> {}

impl<F: FnBench> Step for Dispatch<F> {
    fn op_name(&self) -> &str { &self.name }
    fn select_step(&mut self) -> Result<(), HwdnaError> { self.select().map(|_| ()) }
    fn selected_kernel_name(&self) -> String { self.selected_name().to_string() }
    fn observe_step(&mut self, elapsed: Duration) { self.observe(elapsed); }
    fn rebench_step(&mut self) -> Result<usize, HwdnaError> { self.rebench() }
}

impl<F: FnBench> Dispatch<F> {
    /// Build-time dispatch: when a Cargo feature `force-scalar`, `force-sse`,
    /// `force-avx2`, or `force-neon` is enabled, skip runtime benchmarking
    /// and directly select the matching kernel (if supported on this platform).
    fn try_build_time_select(&mut self) -> Option<usize> {
        #[cfg(any(
            himada_force_scalar,
            himada_force_sse,
            himada_force_avx2,
            himada_force_neon,
        ))]
        {
            let target = {
                #[cfg(himada_force_scalar)] { "scalar" }
                #[cfg(himada_force_sse)] { "SSE" }
                #[cfg(himada_force_avx2)] { "AVX2" }
                #[cfg(himada_force_neon)] { "NEON" }
            };
            for (idx, k) in self.supported.iter().enumerate() {
                if k.name == target && (k.is_supported)(&self.dna) {
                    return Some(idx);
                }
            }
        }
        None
    }
}

/// Builder for [`Dispatch`].
///
/// # Example
/// ```
/// use himada_dispatch::{DispatchBuilder, DotKernel, KernelInfo};
/// use himada_core::HardwareDNA;
/// let d = DispatchBuilder::<DotKernel>::new("my_dot")
///     .with("scalar", |a, b| a.iter().zip(b).map(|(x, y)| x * y).sum(),
///           |_: &HardwareDNA| true, 1)
///     .build();
/// ```
pub struct DispatchBuilder<F> {
    name: String,
    kernels: Vec<KernelInfo<F>>,
}

impl<F> DispatchBuilder<F> {
    pub fn new(name: &str) -> Self {
        Self { name: name.into(), kernels: Vec::new() }
    }
    pub fn with(mut self, name: &'static str, func: F, is_supported: fn(&HardwareDNA) -> bool, thermal_priority: u8) -> Self {
        self.kernels.push(KernelInfo { name, func, is_supported, thermal_priority });
        self
    }
    pub fn build(self) -> Dispatch<F> {
        Dispatch::new(&self.name, self.kernels)
    }
}

/// Create a [`DispatchBuilder`] with one-line syntax.
///
/// ```
/// use himada_dispatch::{dispatch_builder, DotKernel, KernelInfo};
/// use himada_core::HardwareDNA;
/// let d = dispatch_builder!("dot", DotKernel)
///     .with("scalar", |a, b| a.iter().zip(b).map(|(x, y)| x * y).sum(),
///           |_: &HardwareDNA| true, 1)
///     .build();
/// ```
#[macro_export]
macro_rules! dispatch_builder {
    ($name:expr, $ty:ty) => {
        $crate::DispatchBuilder::<$ty>::new($name)
    };
}

/// Create a [`Pipeline`] with concise step-declaration syntax.
///
/// Each step is written as `<operation>("<name>")`. See [`pipeline_add_step!`]
/// for the full list of supported operations.
///
/// # Example
/// ```
/// use himada_dispatch::{pipeline, Pipeline};
/// let mut p = pipeline! {
///     dot("embed_query");
///     dot_f32("embed_doc");
///     matmul("attention");
///     softmax("normalize");
/// };
/// ```
#[macro_export]
macro_rules! pipeline {
    ( $( $op:ident ( $name:expr ); )* $(;)? ) => {{
        let mut __p = $crate::Pipeline::new();
        $(
            $crate::pipeline_add_step!($op, __p, $name);
        )*
        __p
    }};
}

/// Map an operation name token to the corresponding `Pipeline::add_*` method.
#[macro_export]
macro_rules! pipeline_add_step {
    (dot, $p:expr, $name:expr) => { $p.add_dot($name); };
    (dot_f32, $p:expr, $name:expr) => { $p.add_dot_f32($name); };
    (reduce_sum, $p:expr, $name:expr) => { $p.add_reduce_sum($name); };
    (reduce_sum_f32, $p:expr, $name:expr) => { $p.add_reduce_sum_f32($name); };
    (softmax, $p:expr, $name:expr) => { $p.add_softmax($name); };
    (softmax_f32, $p:expr, $name:expr) => { $p.add_softmax_f32($name); };
    (matmul, $p:expr, $name:expr) => { $p.add_matmul($name); };
    (matmul_f32, $p:expr, $name:expr) => { $p.add_matmul_f32($name); };
    (binary_op, $p:expr, $name:expr) => { $p.add_binary_op($name); };
    (binary_op_f32, $p:expr, $name:expr) => { $p.add_binary_op_f32($name); };
    (unary_op, $p:expr, $name:expr) => { $p.add_unary_op($name); };
    (unary_op_f32, $p:expr, $name:expr) => { $p.add_unary_op_f32($name); };
    (clamp, $p:expr, $name:expr) => { $p.add_clamp($name); };
    (clamp_f32, $p:expr, $name:expr) => { $p.add_clamp_f32($name); };
    (dot_i32, $p:expr, $name:expr) => { $p.add_dot_i32($name); };
    (dot_i64, $p:expr, $name:expr) => { $p.add_dot_i64($name); };
    (reduce_i32, $p:expr, $name:expr) => { $p.add_reduce_i32($name); };
    (reduce_i64, $p:expr, $name:expr) => { $p.add_reduce_i64($name); };
    (sort_i32, $p:expr, $name:expr) => { $p.add_sort_i32($name); };
    (sort_f64, $p:expr, $name:expr) => { $p.add_sort_f64($name); };
    (sort_f32, $p:expr, $name:expr) => { $p.add_sort_f32($name); };
    (gemv_f64, $p:expr, $name:expr) => { $p.add_gemv_f64($name); };
    (gemv_f32, $p:expr, $name:expr) => { $p.add_gemv_f32($name); };
    (layer_norm_f64, $p:expr, $name:expr) => { $p.add_layer_norm_f64($name); };
    (layer_norm_f32, $p:expr, $name:expr) => { $p.add_layer_norm_f32($name); };
    (cumsum_f64, $p:expr, $name:expr) => { $p.add_cumsum_f64($name); };
    (cumsum_f32, $p:expr, $name:expr) => { $p.add_cumsum_f32($name); };
    (cumsum_i32, $p:expr, $name:expr) => { $p.add_cumsum_i32($name); };
}

// ---------------------------------------------------------------------------
// Pipeline — group of dispatch steps optimizable together
// ---------------------------------------------------------------------------

/// A pipeline of dispatch operations that can be optimized as a group.
/// Each step can be added, configured, and then the whole pipeline
/// can be optimized (kernel-selected) at once.
pub struct Pipeline {
    dots: Vec<Dispatch<DotKernel>>,
    dots_f32: Vec<Dispatch<DotF32Kernel>>,
    reduces: Vec<Dispatch<ReduceKernel>>,
    reduces_f32: Vec<Dispatch<ReduceSumF32Kernel>>,
    softmaxes: Vec<Dispatch<SoftmaxKernel>>,
    softmaxes_f32: Vec<Dispatch<UnaryOpF32Kernel>>,
    matmuls: Vec<Dispatch<MatMulKernel>>,
    matmuls_f32: Vec<Dispatch<MatMulF32Kernel>>,
    binary_ops: Vec<Dispatch<BinaryOpKernel>>,
    binary_ops_f32: Vec<Dispatch<AddF32Kernel>>,
    unary_ops: Vec<Dispatch<NegateKernel>>,
    unary_ops_f32: Vec<Dispatch<UnaryOpF32Kernel>>,
    clamps: Vec<Dispatch<ClampKernel>>,
    clamps_f32: Vec<Dispatch<ClampF32Kernel>>,
    dot_i32: Vec<Dispatch<DotI32Kernel>>,
    dot_i64: Vec<Dispatch<DotI64Kernel>>,
    reduce_i32: Vec<Dispatch<ReduceI32Kernel>>,
    reduce_i64: Vec<Dispatch<ReduceI64Kernel>>,
    sort_i32: Vec<Dispatch<SortI32Kernel>>,
    sort_f64: Vec<Dispatch<SortF64Kernel>>,
    sort_f32: Vec<Dispatch<SortF32Kernel>>,
    gemv_f64: Vec<Dispatch<GemvF64Kernel>>,
    gemv_f32: Vec<Dispatch<GemvF32Kernel>>,
    layer_norm_f64: Vec<Dispatch<LayerNormF64Kernel>>,
    layer_norm_f32: Vec<Dispatch<LayerNormF32Kernel>>,
    cumsum_f64: Vec<Dispatch<CumsumF64Kernel>>,
    cumsum_f32: Vec<Dispatch<CumsumF32Kernel>>,
    cumsum_i32: Vec<Dispatch<CumsumI32Kernel>>,
}

impl Pipeline {
    pub fn new() -> Self {
        Self {
            dots: Vec::new(),
            dots_f32: Vec::new(),
            reduces: Vec::new(),
            reduces_f32: Vec::new(),
            softmaxes: Vec::new(),
            softmaxes_f32: Vec::new(),
            matmuls: Vec::new(),
            matmuls_f32: Vec::new(),
            binary_ops: Vec::new(),
            binary_ops_f32: Vec::new(),
            unary_ops: Vec::new(),
            unary_ops_f32: Vec::new(),
            clamps: Vec::new(),
            clamps_f32: Vec::new(),
            dot_i32: Vec::new(),
            dot_i64: Vec::new(),
            reduce_i32: Vec::new(),
            reduce_i64: Vec::new(),
            sort_i32: Vec::new(),
            sort_f64: Vec::new(),
            sort_f32: Vec::new(),
            gemv_f64: Vec::new(),
            gemv_f32: Vec::new(),
            layer_norm_f64: Vec::new(),
            layer_norm_f32: Vec::new(),
            cumsum_f64: Vec::new(),
            cumsum_f32: Vec::new(),
            cumsum_i32: Vec::new(),
        }
    }

    pub fn add_dot(&mut self, name: &str) -> &mut Dispatch<DotKernel> {
        self.dots.push(Dispatch::new(name, Vec::new()));
        self.dots.last_mut().unwrap()
    }

    pub fn add_dot_f32(&mut self, name: &str) -> &mut Dispatch<DotF32Kernel> {
        self.dots_f32.push(Dispatch::new(name, Vec::new()));
        self.dots_f32.last_mut().unwrap()
    }

    pub fn add_reduce_sum(&mut self, name: &str) -> &mut Dispatch<ReduceKernel> {
        self.reduces.push(Dispatch::new(name, Vec::new()));
        self.reduces.last_mut().unwrap()
    }

    pub fn add_reduce_sum_f32(&mut self, name: &str) -> &mut Dispatch<ReduceSumF32Kernel> {
        self.reduces_f32.push(Dispatch::new(name, Vec::new()));
        self.reduces_f32.last_mut().unwrap()
    }

    pub fn add_softmax(&mut self, name: &str) -> &mut Dispatch<SoftmaxKernel> {
        self.softmaxes.push(Dispatch::new(name, Vec::new()));
        self.softmaxes.last_mut().unwrap()
    }

    pub fn add_softmax_f32(&mut self, name: &str) -> &mut Dispatch<UnaryOpF32Kernel> {
        self.softmaxes_f32.push(Dispatch::new(name, Vec::new()));
        self.softmaxes_f32.last_mut().unwrap()
    }

    pub fn add_matmul(&mut self, name: &str) -> &mut Dispatch<MatMulKernel> {
        self.matmuls.push(Dispatch::new(name, Vec::new()));
        self.matmuls.last_mut().unwrap()
    }

    pub fn add_matmul_f32(&mut self, name: &str) -> &mut Dispatch<MatMulF32Kernel> {
        self.matmuls_f32.push(Dispatch::new(name, Vec::new()));
        self.matmuls_f32.last_mut().unwrap()
    }

    pub fn add_binary_op(&mut self, name: &str) -> &mut Dispatch<BinaryOpKernel> {
        self.binary_ops.push(Dispatch::new(name, Vec::new()));
        self.binary_ops.last_mut().unwrap()
    }

    pub fn add_binary_op_f32(&mut self, name: &str) -> &mut Dispatch<AddF32Kernel> {
        self.binary_ops_f32.push(Dispatch::new(name, Vec::new()));
        self.binary_ops_f32.last_mut().unwrap()
    }

    pub fn add_unary_op(&mut self, name: &str) -> &mut Dispatch<NegateKernel> {
        self.unary_ops.push(Dispatch::new(name, Vec::new()));
        self.unary_ops.last_mut().unwrap()
    }

    pub fn add_unary_op_f32(&mut self, name: &str) -> &mut Dispatch<UnaryOpF32Kernel> {
        self.unary_ops_f32.push(Dispatch::new(name, Vec::new()));
        self.unary_ops_f32.last_mut().unwrap()
    }

    pub fn add_clamp(&mut self, name: &str) -> &mut Dispatch<ClampKernel> {
        self.clamps.push(Dispatch::new(name, Vec::new()));
        self.clamps.last_mut().unwrap()
    }

    pub fn add_clamp_f32(&mut self, name: &str) -> &mut Dispatch<ClampF32Kernel> {
        self.clamps_f32.push(Dispatch::new(name, Vec::new()));
        self.clamps_f32.last_mut().unwrap()
    }

    pub fn add_dot_i32(&mut self, name: &str) -> &mut Dispatch<DotI32Kernel> {
        self.dot_i32.push(Dispatch::new(name, Vec::new()));
        self.dot_i32.last_mut().unwrap()
    }

    pub fn add_dot_i64(&mut self, name: &str) -> &mut Dispatch<DotI64Kernel> {
        self.dot_i64.push(Dispatch::new(name, Vec::new()));
        self.dot_i64.last_mut().unwrap()
    }

    pub fn add_reduce_i32(&mut self, name: &str) -> &mut Dispatch<ReduceI32Kernel> {
        self.reduce_i32.push(Dispatch::new(name, Vec::new()));
        self.reduce_i32.last_mut().unwrap()
    }

    pub fn add_reduce_i64(&mut self, name: &str) -> &mut Dispatch<ReduceI64Kernel> {
        self.reduce_i64.push(Dispatch::new(name, Vec::new()));
        self.reduce_i64.last_mut().unwrap()
    }

    pub fn add_sort_i32(&mut self, name: &str) -> &mut Dispatch<SortI32Kernel> {
        self.sort_i32.push(Dispatch::new(name, Vec::new()));
        self.sort_i32.last_mut().unwrap()
    }

    pub fn add_sort_f64(&mut self, name: &str) -> &mut Dispatch<SortF64Kernel> {
        self.sort_f64.push(Dispatch::new(name, Vec::new()));
        self.sort_f64.last_mut().unwrap()
    }

    pub fn add_sort_f32(&mut self, name: &str) -> &mut Dispatch<SortF32Kernel> {
        self.sort_f32.push(Dispatch::new(name, Vec::new()));
        self.sort_f32.last_mut().unwrap()
    }

    pub fn add_gemv_f64(&mut self, name: &str) -> &mut Dispatch<GemvF64Kernel> {
        self.gemv_f64.push(Dispatch::new(name, Vec::new()));
        self.gemv_f64.last_mut().unwrap()
    }

    pub fn add_gemv_f32(&mut self, name: &str) -> &mut Dispatch<GemvF32Kernel> {
        self.gemv_f32.push(Dispatch::new(name, Vec::new()));
        self.gemv_f32.last_mut().unwrap()
    }

    pub fn add_layer_norm_f64(&mut self, name: &str) -> &mut Dispatch<LayerNormF64Kernel> {
        self.layer_norm_f64.push(Dispatch::new(name, Vec::new()));
        self.layer_norm_f64.last_mut().unwrap()
    }

    pub fn add_layer_norm_f32(&mut self, name: &str) -> &mut Dispatch<LayerNormF32Kernel> {
        self.layer_norm_f32.push(Dispatch::new(name, Vec::new()));
        self.layer_norm_f32.last_mut().unwrap()
    }

    pub fn add_cumsum_f64(&mut self, name: &str) -> &mut Dispatch<CumsumF64Kernel> {
        self.cumsum_f64.push(Dispatch::new(name, Vec::new()));
        self.cumsum_f64.last_mut().unwrap()
    }

    pub fn add_cumsum_f32(&mut self, name: &str) -> &mut Dispatch<CumsumF32Kernel> {
        self.cumsum_f32.push(Dispatch::new(name, Vec::new()));
        self.cumsum_f32.last_mut().unwrap()
    }

    pub fn add_cumsum_i32(&mut self, name: &str) -> &mut Dispatch<CumsumI32Kernel> {
        self.cumsum_i32.push(Dispatch::new(name, Vec::new()));
        self.cumsum_i32.last_mut().unwrap()
    }

    pub fn dot_count(&self) -> usize { self.dots.len() }
    pub fn dot_f32_count(&self) -> usize { self.dots_f32.len() }
    pub fn reduce_count(&self) -> usize { self.reduces.len() }
    pub fn reduce_f32_count(&self) -> usize { self.reduces_f32.len() }
    pub fn softmax_count(&self) -> usize { self.softmaxes.len() }
    pub fn softmax_f32_count(&self) -> usize { self.softmaxes_f32.len() }
    pub fn matmul_count(&self) -> usize { self.matmuls.len() }
    pub fn matmul_f32_count(&self) -> usize { self.matmuls_f32.len() }
    pub fn binary_op_count(&self) -> usize { self.binary_ops.len() }
    pub fn binary_op_f32_count(&self) -> usize { self.binary_ops_f32.len() }
    pub fn unary_op_count(&self) -> usize { self.unary_ops.len() }
    pub fn unary_op_f32_count(&self) -> usize { self.unary_ops_f32.len() }
    pub fn clamp_count(&self) -> usize { self.clamps.len() }
    pub fn clamp_f32_count(&self) -> usize { self.clamps_f32.len() }
    pub fn dot_i32_count(&self) -> usize { self.dot_i32.len() }
    pub fn dot_i64_count(&self) -> usize { self.dot_i64.len() }
    pub fn reduce_i32_count(&self) -> usize { self.reduce_i32.len() }
    pub fn reduce_i64_count(&self) -> usize { self.reduce_i64.len() }
    pub fn sort_i32_count(&self) -> usize { self.sort_i32.len() }
    pub fn sort_f64_count(&self) -> usize { self.sort_f64.len() }
    pub fn sort_f32_count(&self) -> usize { self.sort_f32.len() }
    pub fn gemv_f64_count(&self) -> usize { self.gemv_f64.len() }
    pub fn gemv_f32_count(&self) -> usize { self.gemv_f32.len() }
    pub fn layer_norm_f64_count(&self) -> usize { self.layer_norm_f64.len() }
    pub fn layer_norm_f32_count(&self) -> usize { self.layer_norm_f32.len() }
    pub fn cumsum_f64_count(&self) -> usize { self.cumsum_f64.len() }
    pub fn cumsum_f32_count(&self) -> usize { self.cumsum_f32.len() }
    pub fn cumsum_i32_count(&self) -> usize { self.cumsum_i32.len() }

    pub fn optimize(&mut self) -> Result<usize, HwdnaError> {
        let mut count = 0;
        for d in &mut self.dots { if d.select().is_ok() { count += 1; } }
        for d in &mut self.dots_f32 { if d.select().is_ok() { count += 1; } }
        for d in &mut self.reduces { if d.select().is_ok() { count += 1; } }
        for d in &mut self.reduces_f32 { if d.select().is_ok() { count += 1; } }
        for d in &mut self.softmaxes { if d.select().is_ok() { count += 1; } }
        for d in &mut self.softmaxes_f32 { if d.select().is_ok() { count += 1; } }
        for d in &mut self.matmuls { if d.select().is_ok() { count += 1; } }
        for d in &mut self.matmuls_f32 { if d.select().is_ok() { count += 1; } }
        for d in &mut self.binary_ops { if d.select().is_ok() { count += 1; } }
        for d in &mut self.binary_ops_f32 { if d.select().is_ok() { count += 1; } }
        for d in &mut self.unary_ops { if d.select().is_ok() { count += 1; } }
        for d in &mut self.unary_ops_f32 { if d.select().is_ok() { count += 1; } }
        for d in &mut self.clamps { if d.select().is_ok() { count += 1; } }
        for d in &mut self.clamps_f32 { if d.select().is_ok() { count += 1; } }
        for d in &mut self.dot_i32 { if d.select().is_ok() { count += 1; } }
        for d in &mut self.dot_i64 { if d.select().is_ok() { count += 1; } }
        for d in &mut self.reduce_i32 { if d.select().is_ok() { count += 1; } }
        for d in &mut self.reduce_i64 { if d.select().is_ok() { count += 1; } }
        for d in &mut self.sort_i32 { if d.select().is_ok() { count += 1; } }
        for d in &mut self.sort_f64 { if d.select().is_ok() { count += 1; } }
        for d in &mut self.sort_f32 { if d.select().is_ok() { count += 1; } }
        for d in &mut self.gemv_f64 { if d.select().is_ok() { count += 1; } }
        for d in &mut self.gemv_f32 { if d.select().is_ok() { count += 1; } }
        for d in &mut self.layer_norm_f64 { if d.select().is_ok() { count += 1; } }
        for d in &mut self.layer_norm_f32 { if d.select().is_ok() { count += 1; } }
        for d in &mut self.cumsum_f64 { if d.select().is_ok() { count += 1; } }
        for d in &mut self.cumsum_f32 { if d.select().is_ok() { count += 1; } }
        for d in &mut self.cumsum_i32 { if d.select().is_ok() { count += 1; } }
        Ok(count)
    }

    pub fn execute(&mut self) -> Result<usize, HwdnaError> {
        self.optimize()
    }

    pub fn to_json(&self) -> serde_json::Value {
        serde_json::json!({
            "dots": self.dots.iter().map(|d| d.to_json()).collect::<Vec<_>>(),
            "dots_f32": self.dots_f32.iter().map(|d| d.to_json()).collect::<Vec<_>>(),
            "reduces": self.reduces.iter().map(|d| d.to_json()).collect::<Vec<_>>(),
            "reduces_f32": self.reduces_f32.iter().map(|d| d.to_json()).collect::<Vec<_>>(),
            "softmaxes": self.softmaxes.iter().map(|d| d.to_json()).collect::<Vec<_>>(),
            "softmaxes_f32": self.softmaxes_f32.iter().map(|d| d.to_json()).collect::<Vec<_>>(),
            "matmuls": self.matmuls.iter().map(|d| d.to_json()).collect::<Vec<_>>(),
            "matmuls_f32": self.matmuls_f32.iter().map(|d| d.to_json()).collect::<Vec<_>>(),
            "binary_ops": self.binary_ops.iter().map(|d| d.to_json()).collect::<Vec<_>>(),
            "binary_ops_f32": self.binary_ops_f32.iter().map(|d| d.to_json()).collect::<Vec<_>>(),
            "unary_ops": self.unary_ops.iter().map(|d| d.to_json()).collect::<Vec<_>>(),
            "unary_ops_f32": self.unary_ops_f32.iter().map(|d| d.to_json()).collect::<Vec<_>>(),
            "clamps": self.clamps.iter().map(|d| d.to_json()).collect::<Vec<_>>(),
            "clamps_f32": self.clamps_f32.iter().map(|d| d.to_json()).collect::<Vec<_>>(),
            "dot_i32": self.dot_i32.iter().map(|d| d.to_json()).collect::<Vec<_>>(),
            "dot_i64": self.dot_i64.iter().map(|d| d.to_json()).collect::<Vec<_>>(),
            "reduce_i32": self.reduce_i32.iter().map(|d| d.to_json()).collect::<Vec<_>>(),
            "reduce_i64": self.reduce_i64.iter().map(|d| d.to_json()).collect::<Vec<_>>(),
            "sort_i32": self.sort_i32.iter().map(|d| d.to_json()).collect::<Vec<_>>(),
            "sort_f64": self.sort_f64.iter().map(|d| d.to_json()).collect::<Vec<_>>(),
            "sort_f32": self.sort_f32.iter().map(|d| d.to_json()).collect::<Vec<_>>(),
            "gemv_f64": self.gemv_f64.iter().map(|d| d.to_json()).collect::<Vec<_>>(),
            "gemv_f32": self.gemv_f32.iter().map(|d| d.to_json()).collect::<Vec<_>>(),
            "layer_norm_f64": self.layer_norm_f64.iter().map(|d| d.to_json()).collect::<Vec<_>>(),
            "layer_norm_f32": self.layer_norm_f32.iter().map(|d| d.to_json()).collect::<Vec<_>>(),
            "cumsum_f64": self.cumsum_f64.iter().map(|d| d.to_json()).collect::<Vec<_>>(),
            "cumsum_f32": self.cumsum_f32.iter().map(|d| d.to_json()).collect::<Vec<_>>(),
            "cumsum_i32": self.cumsum_i32.iter().map(|d| d.to_json()).collect::<Vec<_>>(),
        })
    }

    pub fn save_json(&self, path: &Path) -> Result<(), Box<dyn std::error::Error>> {
        let file = std::fs::File::create(path)?;
        serde_json::to_writer_pretty(file, &self.to_json())?;
        Ok(())
    }

    pub fn dots_mut(&mut self) -> impl Iterator<Item = &mut Dispatch<DotKernel>> {
        self.dots.iter_mut()
    }

    pub fn dots_f32_mut(&mut self) -> impl Iterator<Item = &mut Dispatch<DotF32Kernel>> {
        self.dots_f32.iter_mut()
    }

    pub fn reduces_mut(&mut self) -> impl Iterator<Item = &mut Dispatch<ReduceKernel>> {
        self.reduces.iter_mut()
    }

    pub fn reduces_f32_mut(&mut self) -> impl Iterator<Item = &mut Dispatch<ReduceSumF32Kernel>> {
        self.reduces_f32.iter_mut()
    }

    pub fn softmaxes_mut(&mut self) -> impl Iterator<Item = &mut Dispatch<SoftmaxKernel>> {
        self.softmaxes.iter_mut()
    }

    pub fn softmaxes_f32_mut(&mut self) -> impl Iterator<Item = &mut Dispatch<UnaryOpF32Kernel>> {
        self.softmaxes_f32.iter_mut()
    }

    pub fn matmuls_mut(&mut self) -> impl Iterator<Item = &mut Dispatch<MatMulKernel>> {
        self.matmuls.iter_mut()
    }

    pub fn matmuls_f32_mut(&mut self) -> impl Iterator<Item = &mut Dispatch<MatMulF32Kernel>> {
        self.matmuls_f32.iter_mut()
    }

    pub fn binary_ops_mut(&mut self) -> impl Iterator<Item = &mut Dispatch<BinaryOpKernel>> {
        self.binary_ops.iter_mut()
    }

    pub fn binary_ops_f32_mut(&mut self) -> impl Iterator<Item = &mut Dispatch<AddF32Kernel>> {
        self.binary_ops_f32.iter_mut()
    }

    pub fn unary_ops_mut(&mut self) -> impl Iterator<Item = &mut Dispatch<NegateKernel>> {
        self.unary_ops.iter_mut()
    }

    pub fn unary_ops_f32_mut(&mut self) -> impl Iterator<Item = &mut Dispatch<UnaryOpF32Kernel>> {
        self.unary_ops_f32.iter_mut()
    }

    pub fn clamps_mut(&mut self) -> impl Iterator<Item = &mut Dispatch<ClampKernel>> {
        self.clamps.iter_mut()
    }

    pub fn clamps_f32_mut(&mut self) -> impl Iterator<Item = &mut Dispatch<ClampF32Kernel>> {
        self.clamps_f32.iter_mut()
    }

    pub fn dot_i32_mut(&mut self) -> impl Iterator<Item = &mut Dispatch<DotI32Kernel>> {
        self.dot_i32.iter_mut()
    }

    pub fn dot_i64_mut(&mut self) -> impl Iterator<Item = &mut Dispatch<DotI64Kernel>> {
        self.dot_i64.iter_mut()
    }

    pub fn reduce_i32_mut(&mut self) -> impl Iterator<Item = &mut Dispatch<ReduceI32Kernel>> {
        self.reduce_i32.iter_mut()
    }

    pub fn reduce_i64_mut(&mut self) -> impl Iterator<Item = &mut Dispatch<ReduceI64Kernel>> {
        self.reduce_i64.iter_mut()
    }

    pub fn sort_i32_mut(&mut self) -> impl Iterator<Item = &mut Dispatch<SortI32Kernel>> {
        self.sort_i32.iter_mut()
    }

    pub fn sort_f64_mut(&mut self) -> impl Iterator<Item = &mut Dispatch<SortF64Kernel>> {
        self.sort_f64.iter_mut()
    }

    pub fn sort_f32_mut(&mut self) -> impl Iterator<Item = &mut Dispatch<SortF32Kernel>> {
        self.sort_f32.iter_mut()
    }

    pub fn gemv_f64_mut(&mut self) -> impl Iterator<Item = &mut Dispatch<GemvF64Kernel>> {
        self.gemv_f64.iter_mut()
    }

    pub fn gemv_f32_mut(&mut self) -> impl Iterator<Item = &mut Dispatch<GemvF32Kernel>> {
        self.gemv_f32.iter_mut()
    }

    pub fn layer_norm_f64_mut(&mut self) -> impl Iterator<Item = &mut Dispatch<LayerNormF64Kernel>> {
        self.layer_norm_f64.iter_mut()
    }

    pub fn layer_norm_f32_mut(&mut self) -> impl Iterator<Item = &mut Dispatch<LayerNormF32Kernel>> {
        self.layer_norm_f32.iter_mut()
    }

    pub fn cumsum_f64_mut(&mut self) -> impl Iterator<Item = &mut Dispatch<CumsumF64Kernel>> {
        self.cumsum_f64.iter_mut()
    }

    pub fn cumsum_f32_mut(&mut self) -> impl Iterator<Item = &mut Dispatch<CumsumF32Kernel>> {
        self.cumsum_f32.iter_mut()
    }

    pub fn cumsum_i32_mut(&mut self) -> impl Iterator<Item = &mut Dispatch<CumsumI32Kernel>> {
        self.cumsum_i32.iter_mut()
    }
}

// ---------------------------------------------------------------------------
// Backend selection for heterogeneous dispatch
// ---------------------------------------------------------------------------

/// Compute backend.
#[derive(Clone, Copy, PartialEq, Eq, Debug, Serialize, Deserialize)]
pub enum Backend {
    /// CPU-based execution (default).
    Cpu,
    /// Auto-select between CPU and GPU.
    Auto,
    /// Apple Metal GPU.
    Metal,
    /// NVIDIA CUDA GPU.
    Cuda,
}

/// A dispatch that can run on CPU or GPU, auto-selecting the fastest backend.
pub struct HeterogeneousDispatch<F> {
    cpu: Dispatch<F>,
    gpu: Option<gpu::GpuRuntime>,
    preferred: Backend,
    threshold: usize,
}

impl<F: FnBench> HeterogeneousDispatch<F> {
    pub fn new(cpu: Dispatch<F>) -> Self {
        let gpu_runtime = gpu::GpuRuntime::new();
        Self {
            cpu,
            gpu: if gpu_runtime.is_available() { Some(gpu_runtime) } else { None },
            preferred: Backend::Auto,
            threshold: 1024,
        }
    }

    pub fn set_preferred_backend(&mut self, backend: Backend) {
        self.preferred = backend;
    }

    pub fn preferred_backend(&self) -> Backend {
        self.preferred
    }

    pub fn cpu_dispatch(&self) -> &Dispatch<F> {
        &self.cpu
    }

    pub fn cpu_dispatch_mut(&mut self) -> &mut Dispatch<F> {
        &mut self.cpu
    }

    pub fn gpu_runtime(&self) -> Option<&gpu::GpuRuntime> {
        self.gpu.as_ref()
    }

    pub fn gpu_runtime_mut(&mut self) -> Option<&mut gpu::GpuRuntime> {
        self.gpu.as_mut()
    }

    pub fn set_threshold(&mut self, threshold: usize) {
        self.threshold = threshold;
    }

    pub fn threshold(&self) -> usize {
        self.threshold
    }

    pub fn should_use_gpu(&self, batch_size: usize) -> bool {
        match self.preferred {
            Backend::Auto => self.gpu.is_some() && batch_size >= self.threshold,
            Backend::Cpu => false,
            Backend::Metal | Backend::Cuda => self.gpu.is_some(),
        }
    }
}

unsafe impl<F: FnBench> Send for HeterogeneousDispatch<F> {}
unsafe impl<F: FnBench> Sync for HeterogeneousDispatch<F> {}

// ---------------------------------------------------------------------------
// Compile-time dispatch — fixed kernel selected at compile time
// ---------------------------------------------------------------------------

/// A dispatch that skips runtime selection and uses a fixed kernel
/// determined by `KERNEL_IDX` at compile time.
/// The index is validated at runtime via `assert!` in `new()`.
///
/// # Example
/// ```
/// use himada_dispatch::{DispatchFixed, DotKernel, KernelInfo};
/// use himada_core::HardwareDNA;
/// let kernels = vec![
///     KernelInfo { name: "scalar", func: |a, b| a.iter().zip(b).map(|(x, y)| x * y).sum(), is_supported: |_: &HardwareDNA| true, thermal_priority: 1 },
///     KernelInfo { name: "neon", func: |a, b| a.iter().zip(b).map(|(x, y)| x * y).sum(), is_supported: |_: &HardwareDNA| true, thermal_priority: 2 },
/// ];
/// let fixed = DispatchFixed::<DotKernel, 0>::new(&kernels);
/// assert_eq!(fixed.kernel_name(), "scalar");
/// ```
pub struct DispatchFixed<F, const KERNEL_IDX: usize> {
    kernel: KernelInfo<F>,
}

impl<F: Clone, const KERNEL_IDX: usize> DispatchFixed<F, KERNEL_IDX> {
    /// Create a `DispatchFixed` from a slice of kernels.
    /// Panics if `KERNEL_IDX >= kernels.len()`.
    pub fn new(kernels: &[KernelInfo<F>]) -> Self {
        assert!(
            KERNEL_IDX < kernels.len(),
            "DispatchFixed::new: index {} out of bounds (kernel count = {})",
            KERNEL_IDX,
            kernels.len()
        );
        Self { kernel: kernels[KERNEL_IDX].clone() }
    }

    pub fn kernel_func(&self) -> F {
        self.kernel.func.clone()
    }

    pub fn kernel_name(&self) -> &'static str {
        self.kernel.name
    }

    pub fn thermal_priority(&self) -> u8 {
        self.kernel.thermal_priority
    }

    /// The inner `KernelInfo`.
    pub fn kernel_info(&self) -> &KernelInfo<F> {
        &self.kernel
    }
}

/// Add typed compute methods to DispatchFixed for a specific kernel type.
macro_rules! make_dispatch_fixed_compute {
    ($kernel:ty, ($($arg:ident: $argty:ty),*) -> $ret:ty) => {
        impl<const KERNEL_IDX: usize> DispatchFixed<$kernel, KERNEL_IDX> {
            pub fn try_compute(&self, $($arg: $argty),*) -> Result<$ret, HwdnaError> {
                Ok((self.kernel.func)($($arg),*))
            }
            #[track_caller]
            pub fn compute(&self, $($arg: $argty),*) -> $ret {
                (self.kernel.func)($($arg),*)
            }
        }
    };
}

make_dispatch_fixed_compute!(DotKernel, (a: &[f64], b: &[f64]) -> f64);
make_dispatch_fixed_compute!(ReduceKernel, (a: &[f64]) -> f64);
make_dispatch_fixed_compute!(SoftmaxKernel, (input: &[f64], output: &mut [f64]) -> ());
make_dispatch_fixed_compute!(MatMulKernel, (a: &[f64], b: &[f64], c: &mut [f64], n: usize) -> ());
make_dispatch_fixed_compute!(DotF32Kernel, (a: &[f32], b: &[f32]) -> f32);
make_dispatch_fixed_compute!(BinaryOpKernel, (a: &[f64], b: &[f64], c: &mut [f64]) -> ());
make_dispatch_fixed_compute!(ClampKernel, (a: &[f64], lo: f64, hi: f64, b: &mut [f64]) -> ());
make_dispatch_fixed_compute!(AddF32Kernel, (a: &[f32], b: &[f32], c: &mut [f32]) -> ());
make_dispatch_fixed_compute!(MemchrKernel, (byte: u8, data: &[u8]) -> Option<usize>);
make_dispatch_fixed_compute!(MatMulF32Kernel, (a: &[f32], b: &[f32], c: &mut [f32], n: usize) -> ());
make_dispatch_fixed_compute!(DotI32Kernel, (a: &[i32], b: &[i32]) -> i32);
make_dispatch_fixed_compute!(GemvF64Kernel, (alpha: f64, a: &[f64], x: &[f64], beta: f64, y: &mut [f64], rows: usize, cols: usize) -> ());
make_dispatch_fixed_compute!(LayerNormF64Kernel, (input: &[f64], gamma: &[f64], beta: &[f64], epsilon: f64, output: &mut [f64]) -> ());

// ---------------------------------------------------------------------------
// Dot product: fn(&[f64], &[f64]) -> f64
// ---------------------------------------------------------------------------

/// Dot product kernel: `fn(&[f64], &[f64]) -> f64`
pub type DotKernel = fn(&[f64], &[f64]) -> f64;
make_dispatch_compute!(DotKernel, (a: &[f64], b: &[f64]) -> f64, batch = pair);

impl FnBench for DotKernel {
    fn bench(kernels: &[KernelInfo<Self>], candidates: &[usize]) -> usize {
        const N: usize = 1_000_000;
        let a = vec![1.0; N];
        let b = vec![2.0; N];
        let mut best = candidates[0];
        let mut best_time = Duration::MAX;
        for &idx in candidates {
            let start = Instant::now();
            for _ in 0..100 {
                (kernels[idx].func)(&a, &b);
            }
            let elapsed = start.elapsed();
            if elapsed < best_time {
                best_time = elapsed;
                best = idx;
            }
        }
        best
    }

    fn time_func(kernels: &[KernelInfo<Self>], idx: usize, iters: u32) -> Duration {
        const SAMPLE: usize = 10_000;
        let a = vec![1.0; SAMPLE];
        let b = vec![2.0; SAMPLE];
        let start = Instant::now();
        for _ in 0..iters {
            (kernels[idx].func)(&a, &b);
        }
        start.elapsed()
    }
}

// ---------------------------------------------------------------------------
// Matrix multiply: fn(&[f64], &[f64], &mut [f64], usize)
//   flat N×N matrices, last arg = N
// ---------------------------------------------------------------------------

/// Matrix multiply kernel: `fn(&[f64], &[f64], &mut [f64], usize)`
pub type MatMulKernel = fn(&[f64], &[f64], &mut [f64], usize);

impl FnBench for MatMulKernel {
    fn bench(kernels: &[KernelInfo<Self>], candidates: &[usize]) -> usize {
        const N: usize = 512;
        let a = vec![1.0; N * N];
        let b = vec![2.0; N * N];
        let mut c = vec![0.0; N * N];
        let mut best = candidates[0];
        let mut best_time = Duration::MAX;
        for &idx in candidates {
            let start = Instant::now();
            for _ in 0..5 {
                (kernels[idx].func)(&a, &b, &mut c, N);
                c.fill(0.0);
            }
            let elapsed = start.elapsed();
            if elapsed < best_time {
                best_time = elapsed;
                best = idx;
            }
        }
        best
    }

    fn time_func(kernels: &[KernelInfo<Self>], idx: usize, iters: u32) -> Duration {
        const N: usize = 64;
        let a = vec![1.0; N * N];
        let b = vec![2.0; N * N];
        let mut c = vec![0.0; N * N];
        let start = Instant::now();
        for _ in 0..iters {
            c.fill(0.0);
            (kernels[idx].func)(&a, &b, &mut c, N);
        }
        start.elapsed()
    }
}

make_dispatch_compute!(MatMulKernel, (a: &[f64], b: &[f64], c: &mut [f64], n: usize) -> ());

// ---------------------------------------------------------------------------
// Reduce: fn(&[f64]) -> f64
//   shared by reduce_sum & reduce_max
// ---------------------------------------------------------------------------

/// Reduce kernel (sum or max): `fn(&[f64]) -> f64`
pub type ReduceKernel = fn(&[f64]) -> f64;
make_dispatch_compute!(ReduceKernel, (a: &[f64]) -> f64, batch = single);

impl FnBench for ReduceKernel {
    fn bench(kernels: &[KernelInfo<Self>], candidates: &[usize]) -> usize {
        const N: usize = 1_000_000;
        let a: Vec<f64> = (0..N).map(|i| (i % 100) as f64).collect();
        let mut best = candidates[0];
        let mut best_time = Duration::MAX;
        for &idx in candidates {
            let start = Instant::now();
            for _ in 0..100 {
                std::hint::black_box((kernels[idx].func)(&a));
            }
            let elapsed = start.elapsed();
            if elapsed < best_time {
                best_time = elapsed;
                best = idx;
            }
        }
        best
    }

    fn time_func(kernels: &[KernelInfo<Self>], idx: usize, iters: u32) -> Duration {
        const SAMPLE: usize = 10_000;
        let a: Vec<f64> = (0..SAMPLE).map(|i| (i % 100) as f64).collect();
        let start = Instant::now();
        for _ in 0..iters {
            std::hint::black_box((kernels[idx].func)(&a));
        }
        start.elapsed()
    }
}

// ---------------------------------------------------------------------------
// Softmax: fn(&[f64], &mut [f64])
// ---------------------------------------------------------------------------
/// Softmax kernel: `fn(&[f64], &mut [f64])`
pub type SoftmaxKernel = fn(&[f64], &mut [f64]);
make_dispatch_compute!(SoftmaxKernel, (input: &[f64], output: &mut [f64]) -> ());

impl FnBench for SoftmaxKernel {
    fn bench(kernels: &[KernelInfo<Self>], candidates: &[usize]) -> usize {
        const N: usize = 100_000;
        let input: Vec<f64> = (0..N).map(|i| (i % 100) as f64).collect();
        let mut output = vec![0.0; N];
        let mut best = candidates[0];
        let mut best_time = Duration::MAX;
        for &idx in candidates {
            let start = Instant::now();
            for _ in 0..20 {
                output.fill(0.0);
                (kernels[idx].func)(&input, &mut output);
                std::hint::black_box(&output[0]);
            }
            let elapsed = start.elapsed();
            if elapsed < best_time {
                best_time = elapsed;
                best = idx;
            }
        }
        best
    }

    fn time_func(kernels: &[KernelInfo<Self>], idx: usize, iters: u32) -> Duration {
        const N: usize = 1_000;
        let input: Vec<f64> = (0..N).map(|i| (i % 100) as f64).collect();
        let mut output = vec![0.0; N];
        let start = Instant::now();
        for _ in 0..iters {
            output.fill(0.0);
            (kernels[idx].func)(&input, &mut output);
            std::hint::black_box(&output[0]);
        }
        start.elapsed()
    }
}

// ---------------------------------------------------------------------------
// Memchr: fn(u8, &[u8]) -> Option<usize>
// ---------------------------------------------------------------------------
/// Memchr kernel (find byte in slice): `fn(u8, &[u8]) -> Option<usize>`
pub type MemchrKernel = fn(u8, &[u8]) -> Option<usize>;
make_dispatch_compute!(MemchrKernel, (byte: u8, data: &[u8]) -> Option<usize>);

impl FnBench for MemchrKernel {
    fn bench(kernels: &[KernelInfo<Self>], candidates: &[usize]) -> usize {
        const N: usize = 1_000_000;
        let data: Vec<u8> = (0..N).map(|i| (i % 251) as u8).collect();
        let target = 42u8;
        let mut best = candidates[0];
        let mut best_time = Duration::MAX;
        for &idx in candidates {
            let start = Instant::now();
            for _ in 0..100 {
                std::hint::black_box((kernels[idx].func)(target, &data));
            }
            let elapsed = start.elapsed();
            if elapsed < best_time {
                best_time = elapsed;
                best = idx;
            }
        }
        best
    }

    fn time_func(kernels: &[KernelInfo<Self>], idx: usize, iters: u32) -> Duration {
        const N: usize = 10_000;
        let data: Vec<u8> = (0..N).map(|i| (i % 251) as u8).collect();
        let target = 42u8;
        let start = Instant::now();
        for _ in 0..iters {
            std::hint::black_box((kernels[idx].func)(target, &data));
        }
        start.elapsed()
    }
}

// ---------------------------------------------------------------------------
// Shared: fn(&[f64]) -> f64 — reduce_sum, reduce_max, abs_max
// ---------------------------------------------------------------------------

/// AbsMax convenience alias — same type as ReduceKernel
pub type AbsMaxKernel = fn(&[f64]) -> f64;

// ---------------------------------------------------------------------------
// Argmax: fn(&[f64]) -> usize  (unique return type)
// ---------------------------------------------------------------------------

/// Argmax kernel (f64): `fn(&[f64]) -> usize`
pub type ArgmaxKernel = fn(&[f64]) -> usize;

impl FnBench for ArgmaxKernel {
    fn bench(kernels: &[KernelInfo<Self>], candidates: &[usize]) -> usize {
        const N: usize = 1_000_000;
        let a: Vec<f64> = (0..N).map(|i| (i % 1000) as f64).collect();
        let mut best = candidates[0];
        let mut best_time = Duration::MAX;
        for &idx in candidates {
            let start = Instant::now();
            for _ in 0..100 {
                std::hint::black_box((kernels[idx].func)(&a));
            }
            let elapsed = start.elapsed();
            if elapsed < best_time {
                best_time = elapsed;
                best = idx;
            }
        }
        best
    }

    fn time_func(kernels: &[KernelInfo<Self>], idx: usize, iters: u32) -> Duration {
        const SAMPLE: usize = 10_000;
        let a: Vec<f64> = (0..SAMPLE).map(|i| (i % 100) as f64).collect();
        let start = Instant::now();
        for _ in 0..iters {
            std::hint::black_box((kernels[idx].func)(&a));
        }
        start.elapsed()
    }
}

make_dispatch_compute!(ArgmaxKernel, (a: &[f64]) -> usize, batch = single);

// ---------------------------------------------------------------------------
// ArgmaxF32: fn(&[f32]) -> usize
// ---------------------------------------------------------------------------

/// Argmax kernel (f32): `fn(&[f32]) -> usize`
pub type ArgmaxF32Kernel = fn(&[f32]) -> usize;

impl FnBench for ArgmaxF32Kernel {
    fn bench(kernels: &[KernelInfo<Self>], candidates: &[usize]) -> usize {
        const N: usize = 1_000_000;
        let a = vec![1.0f32; N];
        let mut best = candidates[0];
        let mut best_time = Duration::MAX;
        for &idx in candidates {
            let start = Instant::now();
            for _ in 0..100 {
                std::hint::black_box((kernels[idx].func)(&a));
            }
            let elapsed = start.elapsed();
            if elapsed < best_time {
                best_time = elapsed;
                best = idx;
            }
        }
        best
    }
    fn time_func(kernels: &[KernelInfo<Self>], idx: usize, iters: u32) -> Duration {
        const SAMPLE: usize = 10_000;
        let a = vec![1.0f32; SAMPLE];
        let start = Instant::now();
        for _ in 0..iters {
            std::hint::black_box((kernels[idx].func)(&a));
        }
        start.elapsed()
    }
}

make_dispatch_compute!(ArgmaxF32Kernel, (a: &[f32]) -> usize, batch = single);

/// Euclidean distance convenience alias — same type as DotKernel
pub type EuclideanDistanceKernel = fn(&[f64], &[f64]) -> f64;

/// Cosine similarity convenience alias — same type as DotKernel
pub type CosineSimilarityKernel = fn(&[f64], &[f64]) -> f64;

// ---------------------------------------------------------------------------
// BinaryOp: fn(&[f64], &[f64], &mut [f64]) — hadamard, add, sub, mul
// ---------------------------------------------------------------------------

/// Binary element-wise operation kernel: `fn(&[f64], &[f64], &mut [f64])`
pub type BinaryOpKernel = fn(&[f64], &[f64], &mut [f64]);

impl FnBench for BinaryOpKernel {
    fn bench(kernels: &[KernelInfo<Self>], candidates: &[usize]) -> usize {
        const N: usize = 1_000_000;
        let a = vec![1.0; N];
        let b = vec![2.0; N];
        let mut c = vec![0.0; N];
        let mut best = candidates[0];
        let mut best_time = Duration::MAX;
        for &idx in candidates {
            let start = Instant::now();
            for _ in 0..100 {
                (kernels[idx].func)(&a, &b, &mut c);
            }
            let elapsed = start.elapsed();
            if elapsed < best_time {
                best_time = elapsed;
                best = idx;
            }
        }
        best
    }

    fn time_func(kernels: &[KernelInfo<Self>], idx: usize, iters: u32) -> Duration {
        const SAMPLE: usize = 10_000;
        let a = vec![1.0; SAMPLE];
        let b = vec![2.0; SAMPLE];
        let mut c = vec![0.0; SAMPLE];
        let start = Instant::now();
        for _ in 0..iters {
            (kernels[idx].func)(&a, &b, &mut c);
        }
        start.elapsed()
    }
}

make_dispatch_compute!(BinaryOpKernel, (a: &[f64], b: &[f64], c: &mut [f64]) -> ());

/// Negate convenience alias — same type as SoftmaxKernel
pub type NegateKernel = fn(&[f64], &mut [f64]);

// ---------------------------------------------------------------------------
// Clamp: fn(&[f64], f64, f64, &mut [f64])
// ---------------------------------------------------------------------------

/// Clamp kernel: `fn(&[f64], f64, f64, &mut [f64])`
pub type ClampKernel = fn(&[f64], f64, f64, &mut [f64]);

impl FnBench for ClampKernel {
    fn bench(kernels: &[KernelInfo<Self>], candidates: &[usize]) -> usize {
        const N: usize = 1_000_000;
        let a: Vec<f64> = (0..N).map(|i| (i % 200) as f64 - 100.0).collect();
        let mut b = vec![0.0; N];
        let mut best = candidates[0];
        let mut best_time = Duration::MAX;
        for &idx in candidates {
            let start = Instant::now();
            for _ in 0..100 {
                (kernels[idx].func)(&a, -20.0, 20.0, &mut b);
            }
            let elapsed = start.elapsed();
            if elapsed < best_time {
                best_time = elapsed;
                best = idx;
            }
        }
        best
    }

    fn time_func(kernels: &[KernelInfo<Self>], idx: usize, iters: u32) -> Duration {
        const SAMPLE: usize = 10_000;
        let a: Vec<f64> = (0..SAMPLE).map(|i| (i % 200) as f64 - 100.0).collect();
        let mut b = vec![0.0; SAMPLE];
        let start = Instant::now();
        for _ in 0..iters {
            (kernels[idx].func)(&a, -20.0, 20.0, &mut b);
        }
        start.elapsed()
    }
}

make_dispatch_compute!(ClampKernel, (a: &[f64], lo: f64, hi: f64, b: &mut [f64]) -> ());

// ---------------------------------------------------------------------------
// DotF32: fn(&[f32], &[f32]) -> f32
// ---------------------------------------------------------------------------

/// Dot product kernel (f32): `fn(&[f32], &[f32]) -> f32`
pub type DotF32Kernel = fn(&[f32], &[f32]) -> f32;

impl FnBench for DotF32Kernel {
    fn bench(kernels: &[KernelInfo<Self>], candidates: &[usize]) -> usize {
        const N: usize = 1_000_000;
        let a = vec![1.0f32; N];
        let b = vec![2.0f32; N];
        let mut best = candidates[0];
        let mut best_time = Duration::MAX;
        for &idx in candidates {
            let start = Instant::now();
            for _ in 0..100 {
                (kernels[idx].func)(&a, &b);
            }
            let elapsed = start.elapsed();
            if elapsed < best_time {
                best_time = elapsed;
                best = idx;
            }
        }
        best
    }

    fn time_func(kernels: &[KernelInfo<Self>], idx: usize, iters: u32) -> Duration {
        const SAMPLE: usize = 10_000;
        let a = vec![1.0f32; SAMPLE];
        let b = vec![2.0f32; SAMPLE];
        let start = Instant::now();
        for _ in 0..iters {
            (kernels[idx].func)(&a, &b);
        }
        start.elapsed()
    }
}

make_dispatch_compute!(DotF32Kernel, (a: &[f32], b: &[f32]) -> f32, batch = pair);

// ---------------------------------------------------------------------------
// ReduceSumF32: fn(&[f32]) -> f32  (shared: reduce_sum_f32, reduce_max_f32)
// ---------------------------------------------------------------------------

/// Reduce sum kernel (f32): `fn(&[f32]) -> f32`
pub type ReduceSumF32Kernel = fn(&[f32]) -> f32;

impl FnBench for ReduceSumF32Kernel {
    fn bench(kernels: &[KernelInfo<Self>], candidates: &[usize]) -> usize {
        const N: usize = 1_000_000;
        let a: Vec<f32> = (0..N).map(|i| (i % 100) as f32).collect();
        let mut best = candidates[0];
        let mut best_time = Duration::MAX;
        for &idx in candidates {
            let start = Instant::now();
            for _ in 0..100 {
                std::hint::black_box((kernels[idx].func)(&a));
            }
            let elapsed = start.elapsed();
            if elapsed < best_time {
                best_time = elapsed;
                best = idx;
            }
        }
        best
    }

    fn time_func(kernels: &[KernelInfo<Self>], idx: usize, iters: u32) -> Duration {
        const SAMPLE: usize = 10_000;
        let a: Vec<f32> = (0..SAMPLE).map(|i| (i % 100) as f32).collect();
        let start = Instant::now();
        for _ in 0..iters {
            std::hint::black_box((kernels[idx].func)(&a));
        }
        start.elapsed()
    }
}

make_dispatch_compute!(ReduceSumF32Kernel, (a: &[f32]) -> f32, batch = single);

/// Same type as ReduceSumF32Kernel — convenience alias for max operations
/// Reduce max kernel (f32): `fn(&[f32]) -> f32`
pub type ReduceMaxF32Kernel = fn(&[f32]) -> f32;

// ---------------------------------------------------------------------------
// AddF32: fn(&[f32], &[f32], &mut [f32])  (shared: add, sub, mul)
// ---------------------------------------------------------------------------

/// Add kernel (f32): `fn(&[f32], &[f32], &mut [f32])`
pub type AddF32Kernel = fn(&[f32], &[f32], &mut [f32]);

impl FnBench for AddF32Kernel {
    fn bench(kernels: &[KernelInfo<Self>], candidates: &[usize]) -> usize {
        const N: usize = 1_000_000;
        let a = vec![1.0f32; N];
        let b = vec![2.0f32; N];
        let mut c = vec![0.0f32; N];
        let mut best = candidates[0];
        let mut best_time = Duration::MAX;
        for &idx in candidates {
            let start = Instant::now();
            for _ in 0..100 {
                (kernels[idx].func)(&a, &b, &mut c);
            }
            let elapsed = start.elapsed();
            if elapsed < best_time {
                best_time = elapsed;
                best = idx;
            }
        }
        best
    }

    fn time_func(kernels: &[KernelInfo<Self>], idx: usize, iters: u32) -> Duration {
        const SAMPLE: usize = 10_000;
        let a = vec![1.0f32; SAMPLE];
        let b = vec![2.0f32; SAMPLE];
        let mut c = vec![0.0f32; SAMPLE];
        let start = Instant::now();
        for _ in 0..iters {
            (kernels[idx].func)(&a, &b, &mut c);
        }
        start.elapsed()
    }
}

make_dispatch_compute!(AddF32Kernel, (a: &[f32], b: &[f32], c: &mut [f32]) -> ());

/// Same type as AddF32Kernel — convenience aliases
/// Subtract kernel (f32): `fn(&[f32], &[f32], &mut [f32])`
pub type SubF32Kernel = fn(&[f32], &[f32], &mut [f32]);

/// Multiply kernel (f32): `fn(&[f32], &[f32], &mut [f32])`
pub type MulF32Kernel = fn(&[f32], &[f32], &mut [f32]);
// ---------------------------------------------------------------------------
// UnaryOpF32: fn(&[f32], &mut [f32])
// Shared by: negate_f32, softmax_f32
// ---------------------------------------------------------------------------

/// Unary element-wise operation kernel (f32): `fn(&[f32], &mut [f32])`
pub type UnaryOpF32Kernel = fn(&[f32], &mut [f32]);

impl FnBench for UnaryOpF32Kernel {
    fn bench(kernels: &[KernelInfo<Self>], candidates: &[usize]) -> usize {
        const N: usize = 1_000_000;
        let a = vec![1.0f32; N];
        let mut b = vec![0.0f32; N];
        let mut best = candidates[0];
        let mut best_time = Duration::MAX;
        for &idx in candidates {
            let start = Instant::now();
            for _ in 0..100 {
                (kernels[idx].func)(&a, &mut b);
                std::hint::black_box(&b[0]);
            }
            let elapsed = start.elapsed();
            if elapsed < best_time {
                best_time = elapsed;
                best = idx;
            }
        }
        best
    }

    fn time_func(kernels: &[KernelInfo<Self>], idx: usize, iters: u32) -> Duration {
        const SAMPLE: usize = 10_000;
        let a = vec![1.0f32; SAMPLE];
        let mut b = vec![0.0f32; SAMPLE];
        let start = Instant::now();
        for _ in 0..iters {
            (kernels[idx].func)(&a, &mut b);
            std::hint::black_box(&b[0]);
        }
        start.elapsed()
    }
}

make_dispatch_compute!(UnaryOpF32Kernel, (a: &[f32], b: &mut [f32]) -> ());

/// Negate kernel (f32): alias for `UnaryOpF32Kernel`
pub type NegateF32Kernel = UnaryOpF32Kernel;

// ---------------------------------------------------------------------------
// ClampF32: fn(&[f32], f32, f32, &mut [f32])
// ---------------------------------------------------------------------------

/// Clamp kernel (f32): `fn(&[f32], f32, f32, &mut [f32])`
pub type ClampF32Kernel = fn(&[f32], f32, f32, &mut [f32]);

impl FnBench for ClampF32Kernel {
    fn bench(kernels: &[KernelInfo<Self>], candidates: &[usize]) -> usize {
        const N: usize = 1_000_000;
        let a: Vec<f32> = (0..N).map(|i| (i % 200) as f32 - 100.0).collect();
        let mut b = vec![0.0f32; N];
        let mut best = candidates[0];
        let mut best_time = Duration::MAX;
        for &idx in candidates {
            let start = Instant::now();
            for _ in 0..100 {
                (kernels[idx].func)(&a, -20.0, 20.0, &mut b);
            }
            let elapsed = start.elapsed();
            if elapsed < best_time {
                best_time = elapsed;
                best = idx;
            }
        }
        best
    }

    fn time_func(kernels: &[KernelInfo<Self>], idx: usize, iters: u32) -> Duration {
        const SAMPLE: usize = 10_000;
        let a: Vec<f32> = (0..SAMPLE).map(|i| (i % 200) as f32 - 100.0).collect();
        let mut b = vec![0.0f32; SAMPLE];
        let start = Instant::now();
        for _ in 0..iters {
            (kernels[idx].func)(&a, -20.0, 20.0, &mut b);
        }
        start.elapsed()
    }
}

make_dispatch_compute!(ClampF32Kernel, (a: &[f32], lo: f32, hi: f32, b: &mut [f32]) -> ());

// ---------------------------------------------------------------------------
// Fused kernels — use DotKernel type (same signature), registered as extra kernels
// ---------------------------------------------------------------------------

// ---------------------------------------------------------------------------
// MatMulF32: fn(&[f32], &[f32], &mut [f32], usize)
// ---------------------------------------------------------------------------

/// Matrix multiply kernel (f32): `fn(&[f32], &[f32], &mut [f32], usize)`
pub type MatMulF32Kernel = fn(&[f32], &[f32], &mut [f32], usize);

impl FnBench for MatMulF32Kernel {
    fn bench(kernels: &[KernelInfo<Self>], candidates: &[usize]) -> usize {
        const N: usize = 512;
        let a = vec![1.0f32; N * N];
        let b = vec![2.0f32; N * N];
        let mut c = vec![0.0f32; N * N];
        let mut best = candidates[0];
        let mut best_time = Duration::MAX;
        for &idx in candidates {
            let start = Instant::now();
            for _ in 0..5 {
                (kernels[idx].func)(&a, &b, &mut c, N);
                c.fill(0.0);
            }
            let elapsed = start.elapsed();
            if elapsed < best_time {
                best_time = elapsed;
                best = idx;
            }
        }
        best
    }
    fn time_func(kernels: &[KernelInfo<Self>], idx: usize, iters: u32) -> Duration {
        const N: usize = 64;
        let a = vec![1.0f32; N * N];
        let b = vec![2.0f32; N * N];
        let mut c = vec![0.0f32; N * N];
        let start = Instant::now();
        for _ in 0..iters {
            c.fill(0.0);
            (kernels[idx].func)(&a, &b, &mut c, N);
        }
        start.elapsed()
    }
}

make_dispatch_compute!(MatMulF32Kernel, (a: &[f32], b: &[f32], c: &mut [f32], n: usize) -> ());

// ---------------------------------------------------------------------------
// MatMulBiasRelu: fn(&[f64], &[f64], &[f64], &mut [f64], usize, usize, usize)
// ---------------------------------------------------------------------------

/// MatMul with bias and ReLU activation kernel: `fn(&[f64], &[f64], &[f64], &mut [f64], usize, usize, usize)`
pub type MatMulBiasReluKernel = fn(&[f64], &[f64], &[f64], &mut [f64], usize, usize, usize);

impl FnBench for MatMulBiasReluKernel {
    fn bench(kernels: &[KernelInfo<Self>], candidates: &[usize]) -> usize {
        const M: usize = 64; const K: usize = 64; const N: usize = 64;
        let a = vec![1.0f64; M * K];
        let w = vec![2.0f64; K * N];
        let bias = vec![0.5f64; N];
        let mut c = vec![0.0f64; M * N];
        let mut best = candidates[0];
        let mut best_time = Duration::MAX;
        for &idx in candidates {
            let start = Instant::now();
            for _ in 0..5 {
                (kernels[idx].func)(&a, &w, &bias, &mut c, M, K, N);
                c.fill(0.0);
            }
            let elapsed = start.elapsed();
            if elapsed < best_time {
                best_time = elapsed;
                best = idx;
            }
        }
        best
    }
    fn time_func(kernels: &[KernelInfo<Self>], idx: usize, iters: u32) -> Duration {
        const M: usize = 32; const K: usize = 32; const N: usize = 32;
        let a = vec![1.0f64; M * K];
        let w = vec![2.0f64; K * N];
        let bias = vec![0.5f64; N];
        let mut c = vec![0.0f64; M * N];
        let start = Instant::now();
        for _ in 0..iters {
            c.fill(0.0);
            (kernels[idx].func)(&a, &w, &bias, &mut c, M, K, N);
        }
        start.elapsed()
    }
}

make_dispatch_compute!(MatMulBiasReluKernel, (a: &[f64], w: &[f64], bias: &[f64], c: &mut [f64], m: usize, k: usize, n: usize) -> ());

// ---------------------------------------------------------------------------
// Int32 kernels
// ---------------------------------------------------------------------------

/// Dot product kernel (i32): `fn(&[i32], &[i32]) -> i32`
pub type DotI32Kernel = fn(&[i32], &[i32]) -> i32;
impl FnBench for DotI32Kernel {
    fn bench(kernels: &[KernelInfo<Self>], candidates: &[usize]) -> usize {
        const N: usize = 1_000_000;
        let a: Vec<i32> = (0..N).map(|i| (i % 100) as i32).collect();
        let b: Vec<i32> = (0..N).map(|i| (i % 200) as i32).collect();
        let mut best = candidates[0];
        let mut best_time = Duration::MAX;
        for &idx in candidates {
            let start = Instant::now();
            for _ in 0..100 {
                (kernels[idx].func)(&a, &b);
            }
            let elapsed = start.elapsed();
            if elapsed < best_time { best_time = elapsed; best = idx; }
        }
        best
    }
    fn time_func(kernels: &[KernelInfo<Self>], idx: usize, iters: u32) -> Duration {
        const SAMPLE: usize = 10_000;
        let a: Vec<i32> = (0..SAMPLE).map(|i| (i % 100) as i32).collect();
        let b: Vec<i32> = (0..SAMPLE).map(|i| (i % 200) as i32).collect();
        let start = Instant::now();
        for _ in 0..iters { (kernels[idx].func)(&a, &b); }
        start.elapsed()
    }
}
make_dispatch_compute!(DotI32Kernel, (a: &[i32], b: &[i32]) -> i32, batch = pair);

/// Reduce sum kernel (i32): `fn(&[i32]) -> i32`
pub type ReduceI32Kernel = fn(&[i32]) -> i32;
impl FnBench for ReduceI32Kernel {
    fn bench(kernels: &[KernelInfo<Self>], candidates: &[usize]) -> usize {
        const N: usize = 1_000_000;
        let a: Vec<i32> = (0..N).map(|i| (i % 100) as i32).collect();
        let mut best = candidates[0];
        let mut best_time = Duration::MAX;
        for &idx in candidates {
            let start = Instant::now();
            for _ in 0..100 { std::hint::black_box((kernels[idx].func)(&a)); }
            let elapsed = start.elapsed();
            if elapsed < best_time { best_time = elapsed; best = idx; }
        }
        best
    }
    fn time_func(kernels: &[KernelInfo<Self>], idx: usize, iters: u32) -> Duration {
        const SAMPLE: usize = 10_000;
        let a: Vec<i32> = (0..SAMPLE).map(|i| (i % 100) as i32).collect();
        let start = Instant::now();
        for _ in 0..iters { std::hint::black_box((kernels[idx].func)(&a)); }
        start.elapsed()
    }
}
make_dispatch_compute!(ReduceI32Kernel, (a: &[i32]) -> i32, batch = single);
/// Reduce max kernel (i32): alias for `ReduceI32Kernel`
pub type ReduceMaxI32Kernel = ReduceI32Kernel;
/// Abs max kernel (i32): alias for `ReduceI32Kernel`
pub type AbsMaxI32Kernel = ReduceI32Kernel;

/// Argmax kernel (i32): `fn(&[i32]) -> usize`
pub type ArgmaxI32Kernel = fn(&[i32]) -> usize;
impl FnBench for ArgmaxI32Kernel {
    fn bench(kernels: &[KernelInfo<Self>], candidates: &[usize]) -> usize {
        const N: usize = 1_000_000;
        let a: Vec<i32> = (0..N).map(|i| (i % 1000) as i32).collect();
        let mut best = candidates[0];
        let mut best_time = Duration::MAX;
        for &idx in candidates {
            let start = Instant::now();
            for _ in 0..100 { std::hint::black_box((kernels[idx].func)(&a)); }
            let elapsed = start.elapsed();
            if elapsed < best_time { best_time = elapsed; best = idx; }
        }
        best
    }
    fn time_func(kernels: &[KernelInfo<Self>], idx: usize, iters: u32) -> Duration {
        const SAMPLE: usize = 10_000;
        let a: Vec<i32> = (0..SAMPLE).map(|i| (i % 100) as i32).collect();
        let start = Instant::now();
        for _ in 0..iters { std::hint::black_box((kernels[idx].func)(&a)); }
        start.elapsed()
    }
}
make_dispatch_compute!(ArgmaxI32Kernel, (a: &[i32]) -> usize, batch = single);

/// Binary element-wise operation kernel (i32): `fn(&[i32], &[i32], &mut [i32])`
pub type BinaryI32Kernel = fn(&[i32], &[i32], &mut [i32]);
impl FnBench for BinaryI32Kernel {
    fn bench(kernels: &[KernelInfo<Self>], candidates: &[usize]) -> usize {
        const N: usize = 1_000_000;
        let a: Vec<i32> = (0..N).map(|i| (i % 100) as i32).collect();
        let b: Vec<i32> = (0..N).map(|i| (i % 200) as i32).collect();
        let mut c = vec![0i32; N];
        let mut best = candidates[0];
        let mut best_time = Duration::MAX;
        for &idx in candidates {
            let start = Instant::now();
            for _ in 0..100 { (kernels[idx].func)(&a, &b, &mut c); }
            let elapsed = start.elapsed();
            if elapsed < best_time { best_time = elapsed; best = idx; }
        }
        best
    }
    fn time_func(kernels: &[KernelInfo<Self>], idx: usize, iters: u32) -> Duration {
        const SAMPLE: usize = 10_000;
        let a: Vec<i32> = (0..SAMPLE).map(|i| (i % 100) as i32).collect();
        let b: Vec<i32> = (0..SAMPLE).map(|i| (i % 200) as i32).collect();
        let mut c = vec![0i32; SAMPLE];
        let start = Instant::now();
        for _ in 0..iters { (kernels[idx].func)(&a, &b, &mut c); }
        start.elapsed()
    }
}
make_dispatch_compute!(BinaryI32Kernel, (a: &[i32], b: &[i32], c: &mut [i32]) -> ());

/// Unary element-wise operation kernel (i32): `fn(&[i32], &mut [i32])`
pub type UnaryI32Kernel = fn(&[i32], &mut [i32]);
impl FnBench for UnaryI32Kernel {
    fn bench(kernels: &[KernelInfo<Self>], candidates: &[usize]) -> usize {
        const N: usize = 1_000_000;
        let a: Vec<i32> = (0..N).map(|i| (i % 100) as i32).collect();
        let mut b = vec![0i32; N];
        let mut best = candidates[0];
        let mut best_time = Duration::MAX;
        for &idx in candidates {
            let start = Instant::now();
            for _ in 0..100 { (kernels[idx].func)(&a, &mut b); std::hint::black_box(&b[0]); }
            let elapsed = start.elapsed();
            if elapsed < best_time { best_time = elapsed; best = idx; }
        }
        best
    }
    fn time_func(kernels: &[KernelInfo<Self>], idx: usize, iters: u32) -> Duration {
        const SAMPLE: usize = 10_000;
        let a: Vec<i32> = (0..SAMPLE).map(|i| (i % 100) as i32).collect();
        let mut b = vec![0i32; SAMPLE];
        let start = Instant::now();
        for _ in 0..iters { (kernels[idx].func)(&a, &mut b); std::hint::black_box(&b[0]); }
        start.elapsed()
    }
}
make_dispatch_compute!(UnaryI32Kernel, (a: &[i32], b: &mut [i32]) -> ());

/// Clamp kernel (i32): `fn(&[i32], i32, i32, &mut [i32])`
pub type ClampI32Kernel = fn(&[i32], i32, i32, &mut [i32]);
impl FnBench for ClampI32Kernel {
    fn bench(kernels: &[KernelInfo<Self>], candidates: &[usize]) -> usize {
        const N: usize = 1_000_000;
        let a: Vec<i32> = (0..N).map(|i| (i % 200) as i32 - 100).collect();
        let mut b = vec![0i32; N];
        let mut best = candidates[0];
        let mut best_time = Duration::MAX;
        for &idx in candidates {
            let start = Instant::now();
            for _ in 0..100 { (kernels[idx].func)(&a, -20, 20, &mut b); }
            let elapsed = start.elapsed();
            if elapsed < best_time { best_time = elapsed; best = idx; }
        }
        best
    }
    fn time_func(kernels: &[KernelInfo<Self>], idx: usize, iters: u32) -> Duration {
        const SAMPLE: usize = 10_000;
        let a: Vec<i32> = (0..SAMPLE).map(|i| (i % 200) as i32 - 100).collect();
        let mut b = vec![0i32; SAMPLE];
        let start = Instant::now();
        for _ in 0..iters { (kernels[idx].func)(&a, -20, 20, &mut b); }
        start.elapsed()
    }
}
make_dispatch_compute!(ClampI32Kernel, (a: &[i32], lo: i32, hi: i32, b: &mut [i32]) -> ());

// ---------------------------------------------------------------------------
// Int64 kernels
// ---------------------------------------------------------------------------

/// Dot product kernel (i64): `fn(&[i64], &[i64]) -> i64`
pub type DotI64Kernel = fn(&[i64], &[i64]) -> i64;
impl FnBench for DotI64Kernel {
    fn bench(kernels: &[KernelInfo<Self>], candidates: &[usize]) -> usize {
        const N: usize = 1_000_000;
        let a: Vec<i64> = (0..N).map(|i| (i % 100) as i64).collect();
        let b: Vec<i64> = (0..N).map(|i| (i % 200) as i64).collect();
        let mut best = candidates[0];
        let mut best_time = Duration::MAX;
        for &idx in candidates {
            let start = Instant::now();
            for _ in 0..100 { (kernels[idx].func)(&a, &b); }
            let elapsed = start.elapsed();
            if elapsed < best_time { best_time = elapsed; best = idx; }
        }
        best
    }
    fn time_func(kernels: &[KernelInfo<Self>], idx: usize, iters: u32) -> Duration {
        const SAMPLE: usize = 10_000;
        let a: Vec<i64> = (0..SAMPLE).map(|i| (i % 100) as i64).collect();
        let b: Vec<i64> = (0..SAMPLE).map(|i| (i % 200) as i64).collect();
        let start = Instant::now();
        for _ in 0..iters { (kernels[idx].func)(&a, &b); }
        start.elapsed()
    }
}
make_dispatch_compute!(DotI64Kernel, (a: &[i64], b: &[i64]) -> i64, batch = pair);

/// Reduce sum kernel (i64): `fn(&[i64]) -> i64`
pub type ReduceI64Kernel = fn(&[i64]) -> i64;
impl FnBench for ReduceI64Kernel {
    fn bench(kernels: &[KernelInfo<Self>], candidates: &[usize]) -> usize {
        const N: usize = 1_000_000;
        let a: Vec<i64> = (0..N).map(|i| (i % 100) as i64).collect();
        let mut best = candidates[0]; let mut best_time = Duration::MAX;
        for &idx in candidates {
            let start = Instant::now();
            for _ in 0..100 { std::hint::black_box((kernels[idx].func)(&a)); }
            let elapsed = start.elapsed();
            if elapsed < best_time { best_time = elapsed; best = idx; }
        }
        best
    }
    fn time_func(kernels: &[KernelInfo<Self>], idx: usize, iters: u32) -> Duration {
        const SAMPLE: usize = 10_000;
        let a: Vec<i64> = (0..SAMPLE).map(|i| (i % 100) as i64).collect();
        let start = Instant::now();
        for _ in 0..iters { std::hint::black_box((kernels[idx].func)(&a)); }
        start.elapsed()
    }
}
make_dispatch_compute!(ReduceI64Kernel, (a: &[i64]) -> i64, batch = single);
/// Reduce max kernel (i64): alias for `ReduceI64Kernel`
pub type ReduceMaxI64Kernel = ReduceI64Kernel;
/// Abs max kernel (i64): alias for `ReduceI64Kernel`
pub type AbsMaxI64Kernel = ReduceI64Kernel;

/// Argmax kernel (i64): `fn(&[i64]) -> usize`
pub type ArgmaxI64Kernel = fn(&[i64]) -> usize;
impl FnBench for ArgmaxI64Kernel {
    fn bench(kernels: &[KernelInfo<Self>], candidates: &[usize]) -> usize {
        const N: usize = 1_000_000;
        let a: Vec<i64> = (0..N).map(|i| (i % 1000) as i64).collect();
        let mut best = candidates[0]; let mut best_time = Duration::MAX;
        for &idx in candidates {
            let start = Instant::now();
            for _ in 0..100 { std::hint::black_box((kernels[idx].func)(&a)); }
            let elapsed = start.elapsed();
            if elapsed < best_time { best_time = elapsed; best = idx; }
        }
        best
    }
    fn time_func(kernels: &[KernelInfo<Self>], idx: usize, iters: u32) -> Duration {
        const SAMPLE: usize = 10_000;
        let a: Vec<i64> = (0..SAMPLE).map(|i| (i % 100) as i64).collect();
        let start = Instant::now();
        for _ in 0..iters { std::hint::black_box((kernels[idx].func)(&a)); }
        start.elapsed()
    }
}
make_dispatch_compute!(ArgmaxI64Kernel, (a: &[i64]) -> usize, batch = single);

/// Binary element-wise operation kernel (i64): `fn(&[i64], &[i64], &mut [i64])`
pub type BinaryI64Kernel = fn(&[i64], &[i64], &mut [i64]);
impl FnBench for BinaryI64Kernel {
    fn bench(kernels: &[KernelInfo<Self>], candidates: &[usize]) -> usize {
        const N: usize = 1_000_000;
        let a: Vec<i64> = (0..N).map(|i| (i % 100) as i64).collect();
        let b: Vec<i64> = (0..N).map(|i| (i % 200) as i64).collect();
        let mut c = vec![0i64; N];
        let mut best = candidates[0]; let mut best_time = Duration::MAX;
        for &idx in candidates {
            let start = Instant::now();
            for _ in 0..100 { (kernels[idx].func)(&a, &b, &mut c); }
            let elapsed = start.elapsed();
            if elapsed < best_time { best_time = elapsed; best = idx; }
        }
        best
    }
    fn time_func(kernels: &[KernelInfo<Self>], idx: usize, iters: u32) -> Duration {
        const SAMPLE: usize = 10_000;
        let a: Vec<i64> = (0..SAMPLE).map(|i| (i % 100) as i64).collect();
        let b: Vec<i64> = (0..SAMPLE).map(|i| (i % 200) as i64).collect();
        let mut c = vec![0i64; SAMPLE];
        let start = Instant::now();
        for _ in 0..iters { (kernels[idx].func)(&a, &b, &mut c); }
        start.elapsed()
    }
}
make_dispatch_compute!(BinaryI64Kernel, (a: &[i64], b: &[i64], c: &mut [i64]) -> ());

/// Unary element-wise operation kernel (i64): `fn(&[i64], &mut [i64])`
pub type UnaryI64Kernel = fn(&[i64], &mut [i64]);
impl FnBench for UnaryI64Kernel {
    fn bench(kernels: &[KernelInfo<Self>], candidates: &[usize]) -> usize {
        const N: usize = 1_000_000;
        let a: Vec<i64> = (0..N).map(|i| (i % 100) as i64).collect();
        let mut b = vec![0i64; N];
        let mut best = candidates[0]; let mut best_time = Duration::MAX;
        for &idx in candidates {
            let start = Instant::now();
            for _ in 0..100 { (kernels[idx].func)(&a, &mut b); std::hint::black_box(&b[0]); }
            let elapsed = start.elapsed();
            if elapsed < best_time { best_time = elapsed; best = idx; }
        }
        best
    }
    fn time_func(kernels: &[KernelInfo<Self>], idx: usize, iters: u32) -> Duration {
        const SAMPLE: usize = 10_000;
        let a: Vec<i64> = (0..SAMPLE).map(|i| (i % 100) as i64).collect();
        let mut b = vec![0i64; SAMPLE];
        let start = Instant::now();
        for _ in 0..iters { (kernels[idx].func)(&a, &mut b); std::hint::black_box(&b[0]); }
        start.elapsed()
    }
}
make_dispatch_compute!(UnaryI64Kernel, (a: &[i64], b: &mut [i64]) -> ());

/// Clamp kernel (i64): `fn(&[i64], i64, i64, &mut [i64])`
pub type ClampI64Kernel = fn(&[i64], i64, i64, &mut [i64]);
impl FnBench for ClampI64Kernel {
    fn bench(kernels: &[KernelInfo<Self>], candidates: &[usize]) -> usize {
        const N: usize = 1_000_000;
        let a: Vec<i64> = (0..N).map(|i| (i % 200) as i64 - 100).collect();
        let mut b = vec![0i64; N];
        let mut best = candidates[0]; let mut best_time = Duration::MAX;
        for &idx in candidates {
            let start = Instant::now();
            for _ in 0..100 { (kernels[idx].func)(&a, -20, 20, &mut b); }
            let elapsed = start.elapsed();
            if elapsed < best_time { best_time = elapsed; best = idx; }
        }
        best
    }
    fn time_func(kernels: &[KernelInfo<Self>], idx: usize, iters: u32) -> Duration {
        const SAMPLE: usize = 10_000;
        let a: Vec<i64> = (0..SAMPLE).map(|i| (i % 200) as i64 - 100).collect();
        let mut b = vec![0i64; SAMPLE];
        let start = Instant::now();
        for _ in 0..iters { (kernels[idx].func)(&a, -20, 20, &mut b); }
        start.elapsed()
    }
}
make_dispatch_compute!(ClampI64Kernel, (a: &[i64], lo: i64, hi: i64, b: &mut [i64]) -> ());

// ---------------------------------------------------------------------------
// Int16 kernels
// ---------------------------------------------------------------------------

/// Dot product kernel (i16): `fn(&[i16], &[i16]) -> i16`
pub type DotI16Kernel = fn(&[i16], &[i16]) -> i16;
impl FnBench for DotI16Kernel {
    fn bench(kernels: &[KernelInfo<Self>], candidates: &[usize]) -> usize {
        const N: usize = 1_000_000;
        let a: Vec<i16> = (0..N).map(|i| (i % 100) as i16).collect();
        let b: Vec<i16> = (0..N).map(|i| (i % 200) as i16).collect();
        let mut best = candidates[0]; let mut best_time = Duration::MAX;
        for &idx in candidates {
            let start = Instant::now();
            for _ in 0..100 { (kernels[idx].func)(&a, &b); }
            let elapsed = start.elapsed();
            if elapsed < best_time { best_time = elapsed; best = idx; }
        }
        best
    }
    fn time_func(kernels: &[KernelInfo<Self>], idx: usize, iters: u32) -> Duration {
        const SAMPLE: usize = 10_000;
        let a: Vec<i16> = (0..SAMPLE).map(|i| (i % 100) as i16).collect();
        let b: Vec<i16> = (0..SAMPLE).map(|i| (i % 200) as i16).collect();
        let start = Instant::now();
        for _ in 0..iters { (kernels[idx].func)(&a, &b); }
        start.elapsed()
    }
}
make_dispatch_compute!(DotI16Kernel, (a: &[i16], b: &[i16]) -> i16, batch = pair);

/// Reduce sum kernel (i16): `fn(&[i16]) -> i16`
pub type ReduceI16Kernel = fn(&[i16]) -> i16;
impl FnBench for ReduceI16Kernel {
    fn bench(kernels: &[KernelInfo<Self>], candidates: &[usize]) -> usize {
        const N: usize = 1_000_000;
        let a: Vec<i16> = (0..N).map(|i| (i % 100) as i16).collect();
        let mut best = candidates[0]; let mut best_time = Duration::MAX;
        for &idx in candidates {
            let start = Instant::now();
            for _ in 0..100 { std::hint::black_box((kernels[idx].func)(&a)); }
            let elapsed = start.elapsed();
            if elapsed < best_time { best_time = elapsed; best = idx; }
        }
        best
    }
    fn time_func(kernels: &[KernelInfo<Self>], idx: usize, iters: u32) -> Duration {
        const SAMPLE: usize = 10_000;
        let a: Vec<i16> = (0..SAMPLE).map(|i| (i % 100) as i16).collect();
        let start = Instant::now();
        for _ in 0..iters { std::hint::black_box((kernels[idx].func)(&a)); }
        start.elapsed()
    }
}
make_dispatch_compute!(ReduceI16Kernel, (a: &[i16]) -> i16, batch = single);
/// Reduce max kernel (i16): alias for `ReduceI16Kernel`
pub type ReduceMaxI16Kernel = ReduceI16Kernel;
/// Abs max kernel (i16): alias for `ReduceI16Kernel`
pub type AbsMaxI16Kernel = ReduceI16Kernel;

/// Argmax kernel (i16): `fn(&[i16]) -> usize`
pub type ArgmaxI16Kernel = fn(&[i16]) -> usize;
impl FnBench for ArgmaxI16Kernel {
    fn bench(kernels: &[KernelInfo<Self>], candidates: &[usize]) -> usize {
        const N: usize = 1_000_000;
        let a: Vec<i16> = (0..N).map(|i| (i % 1000) as i16).collect();
        let mut best = candidates[0]; let mut best_time = Duration::MAX;
        for &idx in candidates {
            let start = Instant::now();
            for _ in 0..100 { std::hint::black_box((kernels[idx].func)(&a)); }
            let elapsed = start.elapsed();
            if elapsed < best_time { best_time = elapsed; best = idx; }
        }
        best
    }
    fn time_func(kernels: &[KernelInfo<Self>], idx: usize, iters: u32) -> Duration {
        const SAMPLE: usize = 10_000;
        let a: Vec<i16> = (0..SAMPLE).map(|i| (i % 100) as i16).collect();
        let start = Instant::now();
        for _ in 0..iters { std::hint::black_box((kernels[idx].func)(&a)); }
        start.elapsed()
    }
}
make_dispatch_compute!(ArgmaxI16Kernel, (a: &[i16]) -> usize, batch = single);

/// Binary element-wise operation kernel (i16): `fn(&[i16], &[i16], &mut [i16])`
pub type BinaryI16Kernel = fn(&[i16], &[i16], &mut [i16]);
impl FnBench for BinaryI16Kernel {
    fn bench(kernels: &[KernelInfo<Self>], candidates: &[usize]) -> usize {
        const N: usize = 1_000_000;
        let a: Vec<i16> = (0..N).map(|i| (i % 100) as i16).collect();
        let b: Vec<i16> = (0..N).map(|i| (i % 200) as i16).collect();
        let mut c = vec![0i16; N];
        let mut best = candidates[0]; let mut best_time = Duration::MAX;
        for &idx in candidates {
            let start = Instant::now();
            for _ in 0..100 { (kernels[idx].func)(&a, &b, &mut c); }
            let elapsed = start.elapsed();
            if elapsed < best_time { best_time = elapsed; best = idx; }
        }
        best
    }
    fn time_func(kernels: &[KernelInfo<Self>], idx: usize, iters: u32) -> Duration {
        const SAMPLE: usize = 10_000;
        let a: Vec<i16> = (0..SAMPLE).map(|i| (i % 100) as i16).collect();
        let b: Vec<i16> = (0..SAMPLE).map(|i| (i % 200) as i16).collect();
        let mut c = vec![0i16; SAMPLE];
        let start = Instant::now();
        for _ in 0..iters { (kernels[idx].func)(&a, &b, &mut c); }
        start.elapsed()
    }
}
make_dispatch_compute!(BinaryI16Kernel, (a: &[i16], b: &[i16], c: &mut [i16]) -> ());

/// Unary element-wise operation kernel (i16): `fn(&[i16], &mut [i16])`
pub type UnaryI16Kernel = fn(&[i16], &mut [i16]);
impl FnBench for UnaryI16Kernel {
    fn bench(kernels: &[KernelInfo<Self>], candidates: &[usize]) -> usize {
        const N: usize = 1_000_000;
        let a: Vec<i16> = (0..N).map(|i| (i % 100) as i16).collect();
        let mut b = vec![0i16; N];
        let mut best = candidates[0]; let mut best_time = Duration::MAX;
        for &idx in candidates {
            let start = Instant::now();
            for _ in 0..100 { (kernels[idx].func)(&a, &mut b); std::hint::black_box(&b[0]); }
            let elapsed = start.elapsed();
            if elapsed < best_time { best_time = elapsed; best = idx; }
        }
        best
    }
    fn time_func(kernels: &[KernelInfo<Self>], idx: usize, iters: u32) -> Duration {
        const SAMPLE: usize = 10_000;
        let a: Vec<i16> = (0..SAMPLE).map(|i| (i % 100) as i16).collect();
        let mut b = vec![0i16; SAMPLE];
        let start = Instant::now();
        for _ in 0..iters { (kernels[idx].func)(&a, &mut b); std::hint::black_box(&b[0]); }
        start.elapsed()
    }
}
make_dispatch_compute!(UnaryI16Kernel, (a: &[i16], b: &mut [i16]) -> ());

/// Clamp kernel (i16): `fn(&[i16], i16, i16, &mut [i16])`
pub type ClampI16Kernel = fn(&[i16], i16, i16, &mut [i16]);
impl FnBench for ClampI16Kernel {
    fn bench(kernels: &[KernelInfo<Self>], candidates: &[usize]) -> usize {
        const N: usize = 1_000_000;
        let a: Vec<i16> = (0..N).map(|i| (i % 200) as i16 - 100).collect();
        let mut b = vec![0i16; N];
        let mut best = candidates[0]; let mut best_time = Duration::MAX;
        for &idx in candidates {
            let start = Instant::now();
            for _ in 0..100 { (kernels[idx].func)(&a, -20, 20, &mut b); }
            let elapsed = start.elapsed();
            if elapsed < best_time { best_time = elapsed; best = idx; }
        }
        best
    }
    fn time_func(kernels: &[KernelInfo<Self>], idx: usize, iters: u32) -> Duration {
        const SAMPLE: usize = 10_000;
        let a: Vec<i16> = (0..SAMPLE).map(|i| (i % 200) as i16 - 100).collect();
        let mut b = vec![0i16; SAMPLE];
        let start = Instant::now();
        for _ in 0..iters { (kernels[idx].func)(&a, -20, 20, &mut b); }
        start.elapsed()
    }
}
make_dispatch_compute!(ClampI16Kernel, (a: &[i16], lo: i16, hi: i16, b: &mut [i16]) -> ());

// ---------------------------------------------------------------------------
// Int8 kernels
// ---------------------------------------------------------------------------

/// Dot product kernel (i8): `fn(&[i8], &[i8]) -> i8`
pub type DotI8Kernel = fn(&[i8], &[i8]) -> i8;
impl FnBench for DotI8Kernel {
    fn bench(kernels: &[KernelInfo<Self>], candidates: &[usize]) -> usize {
        const N: usize = 1_000_000;
        let a: Vec<i8> = (0..N).map(|i| (i % 100) as i8).collect();
        let b: Vec<i8> = (0..N).map(|i| (i % 200) as i8).collect();
        let mut best = candidates[0]; let mut best_time = Duration::MAX;
        for &idx in candidates {
            let start = Instant::now();
            for _ in 0..100 { (kernels[idx].func)(&a, &b); }
            let elapsed = start.elapsed();
            if elapsed < best_time { best_time = elapsed; best = idx; }
        }
        best
    }
    fn time_func(kernels: &[KernelInfo<Self>], idx: usize, iters: u32) -> Duration {
        const SAMPLE: usize = 10_000;
        let a: Vec<i8> = (0..SAMPLE).map(|i| (i % 100) as i8).collect();
        let b: Vec<i8> = (0..SAMPLE).map(|i| (i % 200) as i8).collect();
        let start = Instant::now();
        for _ in 0..iters { (kernels[idx].func)(&a, &b); }
        start.elapsed()
    }
}
make_dispatch_compute!(DotI8Kernel, (a: &[i8], b: &[i8]) -> i8, batch = pair);

/// Reduce sum kernel (i8): `fn(&[i8]) -> i8`
pub type ReduceI8Kernel = fn(&[i8]) -> i8;
impl FnBench for ReduceI8Kernel {
    fn bench(kernels: &[KernelInfo<Self>], candidates: &[usize]) -> usize {
        const N: usize = 1_000_000;
        let a: Vec<i8> = (0..N).map(|i| (i % 100) as i8).collect();
        let mut best = candidates[0]; let mut best_time = Duration::MAX;
        for &idx in candidates {
            let start = Instant::now();
            for _ in 0..100 { std::hint::black_box((kernels[idx].func)(&a)); }
            let elapsed = start.elapsed();
            if elapsed < best_time { best_time = elapsed; best = idx; }
        }
        best
    }
    fn time_func(kernels: &[KernelInfo<Self>], idx: usize, iters: u32) -> Duration {
        const SAMPLE: usize = 10_000;
        let a: Vec<i8> = (0..SAMPLE).map(|i| (i % 100) as i8).collect();
        let start = Instant::now();
        for _ in 0..iters { std::hint::black_box((kernels[idx].func)(&a)); }
        start.elapsed()
    }
}
make_dispatch_compute!(ReduceI8Kernel, (a: &[i8]) -> i8, batch = single);
/// Reduce max kernel (i8): alias for `ReduceI8Kernel`
pub type ReduceMaxI8Kernel = ReduceI8Kernel;
/// Abs max kernel (i8): alias for `ReduceI8Kernel`
pub type AbsMaxI8Kernel = ReduceI8Kernel;

/// Argmax kernel (i8): `fn(&[i8]) -> usize`
pub type ArgmaxI8Kernel = fn(&[i8]) -> usize;
impl FnBench for ArgmaxI8Kernel {
    fn bench(kernels: &[KernelInfo<Self>], candidates: &[usize]) -> usize {
        const N: usize = 1_000_000;
        let a: Vec<i8> = (0..N).map(|i| (i % 1000) as i8).collect();
        let mut best = candidates[0]; let mut best_time = Duration::MAX;
        for &idx in candidates {
            let start = Instant::now();
            for _ in 0..100 { std::hint::black_box((kernels[idx].func)(&a)); }
            let elapsed = start.elapsed();
            if elapsed < best_time { best_time = elapsed; best = idx; }
        }
        best
    }
    fn time_func(kernels: &[KernelInfo<Self>], idx: usize, iters: u32) -> Duration {
        const SAMPLE: usize = 10_000;
        let a: Vec<i8> = (0..SAMPLE).map(|i| (i % 100) as i8).collect();
        let start = Instant::now();
        for _ in 0..iters { std::hint::black_box((kernels[idx].func)(&a)); }
        start.elapsed()
    }
}
make_dispatch_compute!(ArgmaxI8Kernel, (a: &[i8]) -> usize, batch = single);

/// Binary element-wise operation kernel (i8): `fn(&[i8], &[i8], &mut [i8])`
pub type BinaryI8Kernel = fn(&[i8], &[i8], &mut [i8]);
impl FnBench for BinaryI8Kernel {
    fn bench(kernels: &[KernelInfo<Self>], candidates: &[usize]) -> usize {
        const N: usize = 1_000_000;
        let a: Vec<i8> = (0..N).map(|i| (i % 100) as i8).collect();
        let b: Vec<i8> = (0..N).map(|i| (i % 200) as i8).collect();
        let mut c = vec![0i8; N];
        let mut best = candidates[0]; let mut best_time = Duration::MAX;
        for &idx in candidates {
            let start = Instant::now();
            for _ in 0..100 { (kernels[idx].func)(&a, &b, &mut c); }
            let elapsed = start.elapsed();
            if elapsed < best_time { best_time = elapsed; best = idx; }
        }
        best
    }
    fn time_func(kernels: &[KernelInfo<Self>], idx: usize, iters: u32) -> Duration {
        const SAMPLE: usize = 10_000;
        let a: Vec<i8> = (0..SAMPLE).map(|i| (i % 100) as i8).collect();
        let b: Vec<i8> = (0..SAMPLE).map(|i| (i % 200) as i8).collect();
        let mut c = vec![0i8; SAMPLE];
        let start = Instant::now();
        for _ in 0..iters { (kernels[idx].func)(&a, &b, &mut c); }
        start.elapsed()
    }
}
make_dispatch_compute!(BinaryI8Kernel, (a: &[i8], b: &[i8], c: &mut [i8]) -> ());

/// Unary element-wise operation kernel (i8): `fn(&[i8], &mut [i8])`
pub type UnaryI8Kernel = fn(&[i8], &mut [i8]);
impl FnBench for UnaryI8Kernel {
    fn bench(kernels: &[KernelInfo<Self>], candidates: &[usize]) -> usize {
        const N: usize = 1_000_000;
        let a: Vec<i8> = (0..N).map(|i| (i % 100) as i8).collect();
        let mut b = vec![0i8; N];
        let mut best = candidates[0]; let mut best_time = Duration::MAX;
        for &idx in candidates {
            let start = Instant::now();
            for _ in 0..100 { (kernels[idx].func)(&a, &mut b); std::hint::black_box(&b[0]); }
            let elapsed = start.elapsed();
            if elapsed < best_time { best_time = elapsed; best = idx; }
        }
        best
    }
    fn time_func(kernels: &[KernelInfo<Self>], idx: usize, iters: u32) -> Duration {
        const SAMPLE: usize = 10_000;
        let a: Vec<i8> = (0..SAMPLE).map(|i| (i % 100) as i8).collect();
        let mut b = vec![0i8; SAMPLE];
        let start = Instant::now();
        for _ in 0..iters { (kernels[idx].func)(&a, &mut b); std::hint::black_box(&b[0]); }
        start.elapsed()
    }
}
make_dispatch_compute!(UnaryI8Kernel, (a: &[i8], b: &mut [i8]) -> ());

/// Clamp kernel (i8): `fn(&[i8], i8, i8, &mut [i8])`
pub type ClampI8Kernel = fn(&[i8], i8, i8, &mut [i8]);
impl FnBench for ClampI8Kernel {
    fn bench(kernels: &[KernelInfo<Self>], candidates: &[usize]) -> usize {
        const N: usize = 1_000_000;
        let a: Vec<i8> = (0..N).map(|i| (i % 200) as i8 - 100).collect();
        let mut b = vec![0i8; N];
        let mut best = candidates[0]; let mut best_time = Duration::MAX;
        for &idx in candidates {
            let start = Instant::now();
            for _ in 0..100 { (kernels[idx].func)(&a, -20, 20, &mut b); }
            let elapsed = start.elapsed();
            if elapsed < best_time { best_time = elapsed; best = idx; }
        }
        best
    }
    fn time_func(kernels: &[KernelInfo<Self>], idx: usize, iters: u32) -> Duration {
        const SAMPLE: usize = 10_000;
        let a: Vec<i8> = (0..SAMPLE).map(|i| (i % 200) as i8 - 100).collect();
        let mut b = vec![0i8; SAMPLE];
        let start = Instant::now();
        for _ in 0..iters { (kernels[idx].func)(&a, -20, 20, &mut b); }
        start.elapsed()
    }
}
make_dispatch_compute!(ClampI8Kernel, (a: &[i8], lo: i8, hi: i8, b: &mut [i8]) -> ());

// ---------------------------------------------------------------------------
// Trig batch kernels (use existing SoftmaxKernel/UnaryOpF32Kernel types)
// ---------------------------------------------------------------------------

/// Trig batch kernel (f64): `fn(&[f64], &mut [f64])`
pub type TrigBatchKernel = fn(&[f64], &mut [f64]); // = SoftmaxKernel
/// Trig batch kernel (f32): `fn(&[f32], &mut [f32])`
pub type TrigBatchF32Kernel = fn(&[f32], &mut [f32]); // = UnaryOpF32Kernel

// ---------------------------------------------------------------------------
// Sort kernels
// ---------------------------------------------------------------------------

/// Sort kernel (i32): `fn(&mut [i32])`
pub type SortI32Kernel = fn(&mut [i32]);
impl FnBench for SortI32Kernel {
    fn bench(kernels: &[KernelInfo<Self>], candidates: &[usize]) -> usize {
        const N: usize = 100_000;
        let a: Vec<i32> = (0..N).map(|i| (i * 7) as i32).collect();
        let mut best = candidates[0]; let mut best_time = Duration::MAX;
        for &idx in candidates {
            let mut copy = a.clone();
            let start = Instant::now();
            for _ in 0..10 { copy.copy_from_slice(&a); (kernels[idx].func)(&mut copy); }
            let elapsed = start.elapsed();
            if elapsed < best_time { best_time = elapsed; best = idx; }
        }
        best
    }
    fn time_func(kernels: &[KernelInfo<Self>], idx: usize, iters: u32) -> Duration {
        const N: usize = 10_000;
        let a: Vec<i32> = (0..N).map(|i| (i * 7) as i32).collect();
        let mut copy = a.clone();
        let start = Instant::now();
        for _ in 0..iters { copy.copy_from_slice(&a); (kernels[idx].func)(&mut copy); }
        start.elapsed()
    }
}
make_dispatch_compute!(SortI32Kernel, (a: &mut [i32]) -> ());

/// Sort kernel (f64): `fn(&mut [f64])`
pub type SortF64Kernel = fn(&mut [f64]);
impl FnBench for SortF64Kernel {
    fn bench(kernels: &[KernelInfo<Self>], candidates: &[usize]) -> usize {
        const N: usize = 100_000;
        let a: Vec<f64> = (0..N).map(|i| (i * 7) as f64).collect();
        let mut best = candidates[0]; let mut best_time = Duration::MAX;
        for &idx in candidates {
            let mut copy = a.clone();
            let start = Instant::now();
            for _ in 0..10 { copy.copy_from_slice(&a); (kernels[idx].func)(&mut copy); }
            let elapsed = start.elapsed();
            if elapsed < best_time { best_time = elapsed; best = idx; }
        }
        best
    }
    fn time_func(kernels: &[KernelInfo<Self>], idx: usize, iters: u32) -> Duration {
        const N: usize = 10_000;
        let a: Vec<f64> = (0..N).map(|i| (i * 7) as f64).collect();
        let mut copy = a.clone();
        let start = Instant::now();
        for _ in 0..iters { copy.copy_from_slice(&a); (kernels[idx].func)(&mut copy); }
        start.elapsed()
    }
}
make_dispatch_compute!(SortF64Kernel, (a: &mut [f64]) -> ());

/// Sort kernel (f32): `fn(&mut [f32])`
pub type SortF32Kernel = fn(&mut [f32]);
impl FnBench for SortF32Kernel {
    fn bench(kernels: &[KernelInfo<Self>], candidates: &[usize]) -> usize {
        const N: usize = 100_000;
        let a: Vec<f32> = (0..N).map(|i| (i * 7) as f32).collect();
        let mut best = candidates[0]; let mut best_time = Duration::MAX;
        for &idx in candidates {
            let mut copy = a.clone();
            let start = Instant::now();
            for _ in 0..10 { copy.copy_from_slice(&a); (kernels[idx].func)(&mut copy); }
            let elapsed = start.elapsed();
            if elapsed < best_time { best_time = elapsed; best = idx; }
        }
        best
    }
    fn time_func(kernels: &[KernelInfo<Self>], idx: usize, iters: u32) -> Duration {
        const N: usize = 10_000;
        let a: Vec<f32> = (0..N).map(|i| (i * 7) as f32).collect();
        let mut copy = a.clone();
        let start = Instant::now();
        for _ in 0..iters { copy.copy_from_slice(&a); (kernels[idx].func)(&mut copy); }
        start.elapsed()
    }
}
make_dispatch_compute!(SortF32Kernel, (a: &mut [f32]) -> ());

// ---------------------------------------------------------------------------
// GEMV kernels
// ---------------------------------------------------------------------------

/// GEMV kernel (f64): `fn(f64, &[f64], &[f64], f64, &mut [f64], usize, usize)`
pub type GemvF64Kernel = fn(f64, &[f64], &[f64], f64, &mut [f64], usize, usize);
impl FnBench for GemvF64Kernel {
    fn bench(kernels: &[KernelInfo<Self>], candidates: &[usize]) -> usize {
        const M: usize = 1000; const K: usize = 1000;
        let a = vec![1.0; M * K];
        let x = vec![2.0; K];
        let mut y = vec![0.0; M];
        let mut best = candidates[0]; let mut best_time = Duration::MAX;
        for &idx in candidates {
            y.fill(0.0);
            let start = Instant::now();
            for _ in 0..10 { (kernels[idx].func)(1.0, &a, &x, 0.0, &mut y, M, K); }
            let elapsed = start.elapsed();
            if elapsed < best_time { best_time = elapsed; best = idx; }
        }
        best
    }
    fn time_func(kernels: &[KernelInfo<Self>], idx: usize, iters: u32) -> Duration {
        const M: usize = 100; const K: usize = 100;
        let a = vec![1.0; M * K];
        let x = vec![2.0; K];
        let mut y = vec![0.0; M];
        let start = Instant::now();
        for _ in 0..iters { (kernels[idx].func)(1.0, &a, &x, 0.0, &mut y, M, K); y.fill(0.0); }
        start.elapsed()
    }
}
make_dispatch_compute!(GemvF64Kernel, (alpha: f64, a: &[f64], x: &[f64], beta: f64, y: &mut [f64], rows: usize, cols: usize) -> ());

/// GEMV kernel (f32): `fn(f32, &[f32], &[f32], f32, &mut [f32], usize, usize)`
pub type GemvF32Kernel = fn(f32, &[f32], &[f32], f32, &mut [f32], usize, usize);
impl FnBench for GemvF32Kernel {
    fn bench(kernels: &[KernelInfo<Self>], candidates: &[usize]) -> usize {
        const M: usize = 1000; const K: usize = 1000;
        let a = vec![1.0f32; M * K];
        let x = vec![2.0f32; K];
        let mut y = vec![0.0f32; M];
        let mut best = candidates[0]; let mut best_time = Duration::MAX;
        for &idx in candidates {
            y.fill(0.0);
            let start = Instant::now();
            for _ in 0..10 { (kernels[idx].func)(1.0, &a, &x, 0.0, &mut y, M, K); }
            let elapsed = start.elapsed();
            if elapsed < best_time { best_time = elapsed; best = idx; }
        }
        best
    }
    fn time_func(kernels: &[KernelInfo<Self>], idx: usize, iters: u32) -> Duration {
        const M: usize = 100; const K: usize = 100;
        let a = vec![1.0f32; M * K];
        let x = vec![2.0f32; K];
        let mut y = vec![0.0f32; M];
        let start = Instant::now();
        for _ in 0..iters { (kernels[idx].func)(1.0, &a, &x, 0.0, &mut y, M, K); y.fill(0.0); }
        start.elapsed()
    }
}
make_dispatch_compute!(GemvF32Kernel, (alpha: f32, a: &[f32], x: &[f32], beta: f32, y: &mut [f32], rows: usize, cols: usize) -> ());

// ---------------------------------------------------------------------------
// Conv1D kernels (share BinaryOpKernel / AddF32Kernel types)
// ---------------------------------------------------------------------------

/// 1D convolution kernel (f64): `fn(&[f64], &[f64], &mut [f64])`
pub type Conv1dF64Kernel = fn(&[f64], &[f64], &mut [f64]); // = BinaryOpKernel
/// 1D convolution kernel (f32): `fn(&[f32], &[f32], &mut [f32])`
pub type Conv1dF32Kernel = fn(&[f32], &[f32], &mut [f32]); // = AddF32Kernel

// ---------------------------------------------------------------------------
// LayerNorm kernels
// ---------------------------------------------------------------------------

/// Layer normalization kernel (f64): `fn(&[f64], &[f64], &[f64], f64, &mut [f64])`
pub type LayerNormF64Kernel = fn(&[f64], &[f64], &[f64], f64, &mut [f64]);
impl FnBench for LayerNormF64Kernel {
    fn bench(kernels: &[KernelInfo<Self>], candidates: &[usize]) -> usize {
        const N: usize = 100_000;
        let input: Vec<f64> = (0..N).map(|i| (i % 100) as f64).collect();
        let gamma: Vec<f64> = (0..N).map(|i| (i % 10) as f64 * 0.1).collect();
        let beta: Vec<f64> = (0..N).map(|i| (i % 10) as f64 * 0.05).collect();
        let mut output = vec![0.0; N];
        let mut best = candidates[0]; let mut best_time = Duration::MAX;
        for &idx in candidates {
            let start = Instant::now();
            for _ in 0..20 { (kernels[idx].func)(&input, &gamma, &beta, 1e-8, &mut output); }
            let elapsed = start.elapsed();
            if elapsed < best_time { best_time = elapsed; best = idx; }
        }
        best
    }
    fn time_func(kernels: &[KernelInfo<Self>], idx: usize, iters: u32) -> Duration {
        const N: usize = 1_000;
        let input: Vec<f64> = (0..N).map(|i| (i % 100) as f64).collect();
        let gamma: Vec<f64> = (0..N).map(|i| (i % 10) as f64 * 0.1).collect();
        let beta: Vec<f64> = (0..N).map(|i| (i % 10) as f64 * 0.05).collect();
        let mut output = vec![0.0; N];
        let start = Instant::now();
        for _ in 0..iters { (kernels[idx].func)(&input, &gamma, &beta, 1e-8, &mut output); }
        start.elapsed()
    }
}
make_dispatch_compute!(LayerNormF64Kernel, (input: &[f64], gamma: &[f64], beta: &[f64], epsilon: f64, output: &mut [f64]) -> ());

/// Layer normalization kernel (f32): `fn(&[f32], &[f32], &[f32], f32, &mut [f32])`
pub type LayerNormF32Kernel = fn(&[f32], &[f32], &[f32], f32, &mut [f32]);
impl FnBench for LayerNormF32Kernel {
    fn bench(kernels: &[KernelInfo<Self>], candidates: &[usize]) -> usize {
        const N: usize = 100_000;
        let input: Vec<f32> = (0..N).map(|i| (i % 100) as f32).collect();
        let gamma: Vec<f32> = (0..N).map(|i| (i % 10) as f32 * 0.1).collect();
        let beta: Vec<f32> = (0..N).map(|i| (i % 10) as f32 * 0.05).collect();
        let mut output = vec![0.0; N];
        let mut best = candidates[0]; let mut best_time = Duration::MAX;
        for &idx in candidates {
            let start = Instant::now();
            for _ in 0..20 { (kernels[idx].func)(&input, &gamma, &beta, 1e-8, &mut output); }
            let elapsed = start.elapsed();
            if elapsed < best_time { best_time = elapsed; best = idx; }
        }
        best
    }
    fn time_func(kernels: &[KernelInfo<Self>], idx: usize, iters: u32) -> Duration {
        const N: usize = 1_000;
        let input: Vec<f32> = (0..N).map(|i| (i % 100) as f32).collect();
        let gamma: Vec<f32> = (0..N).map(|i| (i % 10) as f32 * 0.1).collect();
        let beta: Vec<f32> = (0..N).map(|i| (i % 10) as f32 * 0.05).collect();
        let mut output = vec![0.0; N];
        let start = Instant::now();
        for _ in 0..iters { (kernels[idx].func)(&input, &gamma, &beta, 1e-8, &mut output); }
        start.elapsed()
    }
}
make_dispatch_compute!(LayerNormF32Kernel, (input: &[f32], gamma: &[f32], beta: &[f32], epsilon: f32, output: &mut [f32]) -> ());

// ---------------------------------------------------------------------------
// Cumsum kernels (share existing types)
// ---------------------------------------------------------------------------

/// Cumulative sum kernel (f64): `fn(&[f64], &mut [f64])`
pub type CumsumF64Kernel = fn(&[f64], &mut [f64]); // = SoftmaxKernel
/// Cumulative sum kernel (f32): `fn(&[f32], &mut [f32])`
pub type CumsumF32Kernel = fn(&[f32], &mut [f32]); // = UnaryOpF32Kernel
/// Cumulative sum kernel (i32): `fn(&[i32], &mut [i32])`
pub type CumsumI32Kernel = fn(&[i32], &mut [i32]); // = UnaryI32Kernel

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

    macro_rules! assert_close_f64 {
        ($a:expr, $b:expr) => {
            let diff = ($a - $b).abs();
            assert!(diff < 1e-12, "left={}, right={}, diff={}", $a, $b, diff);
        };
    }

    macro_rules! assert_close_f32 {
        ($a:expr, $b:expr) => {
            let diff = ($a - $b).abs();
            assert!(diff < 1e-6, "left={}, right={}, diff={}", $a, $b, diff);
        };
    }

    // -----------------------------------------------------------------------
    // Dot product (f64)
    // -----------------------------------------------------------------------
    #[test]
    fn test_dot_f64_scalar_vs_neon() {
        let a: Vec<f64> = (0..100).map(|i| i as f64).collect();
        let b: Vec<f64> = (100..200).map(|i| i as f64).collect();
        let expected = dot_scalar(&a, &b);
        let neon = dot_neon(&a, &b);
        assert_close_f64!(expected, neon);
    }

    #[test]
    fn test_dot_f64_empty() {
        let a: Vec<f64> = vec![];
        let b: Vec<f64> = vec![];
        assert_close_f64!(dot_scalar(&a, &b), 0.0);
        assert_close_f64!(dot_neon(&a, &b), 0.0);
    }

    #[test]
    fn test_dot_f64_single() {
        let a = vec![3.0];
        let b = vec![4.0];
        assert_close_f64!(dot_scalar(&a, &b), 12.0);
        assert_close_f64!(dot_neon(&a, &b), 12.0);
    }

    // -----------------------------------------------------------------------
    // Reduce sum (f64)
    // -----------------------------------------------------------------------
    #[test]
    fn test_reduce_sum_f64_scalar_vs_neon() {
        let a: Vec<f64> = (0..200).map(|i| (i % 7) as f64).collect();
        let expected = reduce_sum_scalar(&a);
        assert_close_f64!(expected, reduce_sum_neon(&a));
    }

    #[test]
    fn test_reduce_sum_f64_empty() {
        assert_close_f64!(reduce_sum_scalar(&[]), 0.0);
        assert_close_f64!(reduce_sum_neon(&[]), 0.0);
    }

    // -----------------------------------------------------------------------
    // Reduce max (f64)
    // -----------------------------------------------------------------------
    #[test]
    fn test_reduce_max_f64_scalar_vs_neon() {
        let a: Vec<f64> = vec![1.0, 5.0, 3.0, 9.0, 2.0];
        assert_close_f64!(reduce_max_scalar(&a), reduce_max_neon(&a));
    }

    #[test]
    fn test_reduce_max_f64_negative() {
        let a: Vec<f64> = vec![-5.0, -3.0, -10.0, -1.0];
        assert_close_f64!(reduce_max_scalar(&a), reduce_max_neon(&a));
    }

    // -----------------------------------------------------------------------
    // Abs max (f64)
    // -----------------------------------------------------------------------
    #[test]
    fn test_abs_max_f64_scalar_vs_neon() {
        let a: Vec<f64> = vec![-5.0, 3.0, -10.0, 7.0];
        assert_close_f64!(abs_max_f64_scalar(&a), abs_max_f64_neon(&a));
    }

    #[test]
    fn test_abs_max_f64_empty() {
        assert_close_f64!(abs_max_f64_scalar(&[]), 0.0);
        assert_close_f64!(abs_max_f64_neon(&[]), 0.0);
    }

    // -----------------------------------------------------------------------
    // Argmax (f64)
    // -----------------------------------------------------------------------
    #[test]
    fn test_argmax_f64_scalar_vs_neon() {
        let a: Vec<f64> = vec![1.0, 5.0, 3.0, 9.0, 2.0];
        assert_eq!(argmax_f64_scalar(&a), argmax_f64_neon(&a));
    }

    #[test]
    fn test_argmax_f64_first() {
        let a = vec![10.0, 1.0, 2.0, 3.0];
        assert_eq!(argmax_f64_scalar(&a), 0);
        assert_eq!(argmax_f64_neon(&a), 0);
    }

    #[test]
    fn test_argmax_f64_last() {
        let a = vec![1.0, 2.0, 3.0, 10.0];
        assert_eq!(argmax_f64_scalar(&a), 3);
        assert_eq!(argmax_f64_neon(&a), 3);
    }

    #[test]
    fn test_argmax_f64_empty() {
        assert_eq!(argmax_f64_scalar(&[]), 0);
        assert_eq!(argmax_f64_neon(&[]), 0);
    }

    // -----------------------------------------------------------------------
    // Argmax (f32) — SSE/AVX2/NEON
    // -----------------------------------------------------------------------
    #[test]
    fn test_argmax_f32_scalar_vs_neon() {
        let a: Vec<f32> = vec![1.0, 5.0, 3.0, 9.0, 2.0];
        assert_eq!(argmax_f32_scalar(&a), argmax_f32_neon(&a));
    }

    #[test]
    fn test_argmax_f32_first() {
        let a = vec![10.0, 1.0, 2.0, 3.0];
        assert_eq!(argmax_f32_scalar(&a), 0);
        assert_eq!(argmax_f32_neon(&a), 0);
    }

    #[test]
    fn test_argmax_f32_last() {
        let a = vec![1.0, 2.0, 3.0, 10.0];
        assert_eq!(argmax_f32_scalar(&a), 3);
        assert_eq!(argmax_f32_neon(&a), 3);
    }

    #[test]
    fn test_argmax_f32_empty() {
        assert_eq!(argmax_f32_scalar(&[]), 0);
        assert_eq!(argmax_f32_neon(&[]), 0);
    }

    // -----------------------------------------------------------------------
    // Euclidean distance (f64)
    // -----------------------------------------------------------------------
    #[test]
    fn test_euclidean_f64_scalar_vs_neon() {
        let a: Vec<f64> = (0..50).map(|i| i as f64).collect();
        let b: Vec<f64> = (50..100).map(|i| i as f64).collect();
        assert_close_f64!(
            euclidean_distance_f64_scalar(&a, &b),
            euclidean_distance_f64_neon(&a, &b)
        );
    }

    #[test]
    fn test_euclidean_f64_empty() {
        assert_close_f64!(euclidean_distance_f64_scalar(&[], &[]), 0.0);
        assert_close_f64!(euclidean_distance_f64_neon(&[], &[]), 0.0);
    }

    #[test]
    fn test_euclidean_f64_identical() {
        let a: Vec<f64> = vec![1.0, 2.0, 3.0];
        assert_close_f64!(euclidean_distance_f64_scalar(&a, &a), 0.0);
        assert_close_f64!(euclidean_distance_f64_neon(&a, &a), 0.0);
    }

    // -----------------------------------------------------------------------
    // Cosine similarity (f64)
    // -----------------------------------------------------------------------
    #[test]
    fn test_cosine_similarity() {
        let a = vec![1.0, 0.0];
        let b = vec![0.0, 1.0];
        assert_close_f64!(cosine_similarity_f64_scalar(&a, &b), 0.0);
    }

    #[test]
    fn test_cosine_similarity_parallel() {
        let a = vec![2.0, 4.0];
        let b = vec![1.0, 2.0];
        assert_close_f64!(cosine_similarity_f64_scalar(&a, &b), 1.0);
    }

    // -----------------------------------------------------------------------
    // Negate (f64)
    // -----------------------------------------------------------------------
    #[test]
    fn test_negate_f64_scalar_vs_neon() {
        let a: Vec<f64> = vec![1.0, -2.0, 3.0, -4.0];
        let mut c1 = vec![0.0; 4];
        let mut c2 = vec![0.0; 4];
        negate_f64_scalar(&a, &mut c1);
        negate_f64_neon(&a, &mut c2);
        assert_eq!(c1, c2);
    }

    // -----------------------------------------------------------------------
    // Clamp (f64)
    // -----------------------------------------------------------------------
    #[test]
    fn test_clamp_f64_scalar_vs_neon() {
        let a: Vec<f64> = vec![-10.0, -5.0, 0.0, 5.0, 10.0];
        let mut c1 = vec![0.0; 5];
        let mut c2 = vec![0.0; 5];
        clamp_f64_scalar(&a, -3.0, 3.0, &mut c1);
        clamp_f64_neon(&a, -3.0, 3.0, &mut c2);
        assert_eq!(c1, c2);
    }

    // -----------------------------------------------------------------------
    // Binary ops (f64): hadamard, add, sub, mul
    // -----------------------------------------------------------------------
    #[test]
    fn test_hadamard_f64() {
        let a = vec![1.0, 2.0, 3.0];
        let b = vec![4.0, 5.0, 6.0];
        let mut c = vec![0.0; 3];
        hadamard_product_f64_scalar(&a, &b, &mut c);
        assert_close_f64!(c[0], 4.0);
        assert_close_f64!(c[1], 10.0);
        assert_close_f64!(c[2], 18.0);
    }

    #[test]
    fn test_add_f64() {
        let a = vec![1.0, 2.0];
        let b = vec![3.0, 4.0];
        let mut c = vec![0.0; 2];
        add_f64_scalar(&a, &b, &mut c);
        assert_close_f64!(c[0], 4.0);
        assert_close_f64!(c[1], 6.0);
    }

    #[test]
    fn test_sub_f64() {
        let a = vec![5.0, 8.0];
        let b = vec![3.0, 2.0];
        let mut c = vec![0.0; 2];
        sub_f64_scalar(&a, &b, &mut c);
        assert_close_f64!(c[0], 2.0);
        assert_close_f64!(c[1], 6.0);
    }

    #[test]
    fn test_mul_f64() {
        let a = vec![2.0, 3.0];
        let b = vec![4.0, 5.0];
        let mut c = vec![0.0; 2];
        mul_f64_scalar(&a, &b, &mut c);
        assert_close_f64!(c[0], 8.0);
        assert_close_f64!(c[1], 15.0);
    }

    // -----------------------------------------------------------------------
    // Matmul
    // -----------------------------------------------------------------------
    #[test]
    fn test_matmul_f64_scalar_vs_neon() {
        const N: usize = 8;
        let a: Vec<f64> = (0..N * N).map(|i| (i % 5) as f64).collect();
        let b: Vec<f64> = (0..N * N).map(|i| (i % 7) as f64).collect();
        let mut c1 = vec![0.0; N * N];
        let mut c2 = vec![0.0; N * N];
        matmul_scalar(&a, &b, &mut c1, N);
        matmul_neon(&a, &b, &mut c2, N);
        for i in 0..N * N {
            assert_close_f64!(c1[i], c2[i]);
        }
    }

    #[test]
    fn test_matmul_cache_tiled_vs_scalar() {
        const N: usize = 64;
        let a: Vec<f64> = (0..N * N).map(|i| (i % 5) as f64).collect();
        let b: Vec<f64> = (0..N * N).map(|i| (i % 7) as f64).collect();
        let mut c1 = vec![0.0; N * N];
        let mut c2 = vec![0.0; N * N];
        matmul_scalar(&a, &b, &mut c1, N);
        matmul_cache_tiled(&a, &b, &mut c2, N);
        for i in 0..N * N {
            assert_close_f64!(c1[i], c2[i]);
        }
    }

    // -----------------------------------------------------------------------
    // Memchr
    // -----------------------------------------------------------------------
    #[test]
    fn test_memchr_basic() {
        let data = b"hello world";
        assert_eq!(memchr_scalar(b'w', data), Some(6));
        assert_eq!(memchr_scalar(b'z', data), None);
    }

    #[test]
    fn test_memchr_neon() {
        let data: Vec<u8> = (0..200).collect();
        assert_eq!(memchr_scalar(42, &data), memchr_neon(42, &data));
        assert_eq!(memchr_scalar(255, &data), memchr_neon(255, &data));
    }

    // -----------------------------------------------------------------------
    // Softmax
    // -----------------------------------------------------------------------
    #[test]
    fn test_softmax_f64_scalar_vs_neon() {
        let input: Vec<f64> = vec![1.0, 2.0, 3.0, 4.0, 5.0];
        let mut c1 = vec![0.0; 5];
        let mut c2 = vec![0.0; 5];
        softmax_scalar(&input, &mut c1);
        softmax_neon(&input, &mut c2);
        let sum1: f64 = c1.iter().sum();
        let sum2: f64 = c2.iter().sum();
        assert_close_f64!(sum1, 1.0);
        assert_close_f64!(sum2, 1.0);
        for i in 0..5 {
            assert_close_f64!(c1[i], c2[i]);
        }
    }

    // -----------------------------------------------------------------------
    // Dispatch selection
    // -----------------------------------------------------------------------
    #[test]
    fn test_dispatch_select_dot() {
        let mut d = Dispatch::<DotKernel>::new("test_dot", vec![
            KernelInfo { name: "scalar", func: dot_scalar, is_supported: |_| true, thermal_priority: 1 },
        ]);
        assert!(d.select().is_ok());
        assert_eq!(d.selected_name(), "scalar");
    }

    #[test]
    fn test_dispatch_no_supported() {
        let mut d = Dispatch::<DotKernel>::new("test_empty", vec![
            KernelInfo { name: "never", func: dot_scalar, is_supported: |_| false, thermal_priority: 1 },
        ]);
        assert!(matches!(d.select(), Err(HwdnaError::NoSupportedKernels)));
    }

    #[test]
    fn test_dispatch_compute_dot() {
        let mut d = Dispatch::<DotKernel>::new("test_compute", vec![
            KernelInfo { name: "scalar", func: dot_scalar, is_supported: |_| true, thermal_priority: 1 },
        ]);
        let a = vec![1.0, 2.0];
        let b = vec![3.0, 4.0];
        assert_close_f64!(d.compute(&a, &b), 11.0);
    }

    // -----------------------------------------------------------------------
    // f32 kernels
    // -----------------------------------------------------------------------
    #[test]
    fn test_dot_f32_scalar_vs_neon() {
        let a: Vec<f32> = (0..100).map(|i| i as f32).collect();
        let b: Vec<f32> = (100..200).map(|i| i as f32).collect();
        assert_close_f32!(dot_f32_scalar(&a, &b), dot_f32_neon(&a, &b));
    }

    #[test]
    fn test_reduce_sum_f32_scalar_vs_neon() {
        let a: Vec<f32> = (0..200).map(|i| (i % 7) as f32).collect();
        assert_close_f32!(reduce_sum_f32_scalar(&a), reduce_sum_f32_neon(&a));
    }

    #[test]
    fn test_reduce_max_f32_scalar_vs_neon() {
        let a: Vec<f32> = vec![1.0, 5.0, 3.0, 9.0, 2.0];
        assert_close_f32!(reduce_max_f32_scalar(&a), reduce_max_f32_neon(&a));
    }

    #[test]
    fn test_abs_max_f32_scalar_vs_neon() {
        let a: Vec<f32> = vec![-5.0, 3.0, -10.0, 7.0];
        assert_close_f32!(abs_max_f32_scalar(&a), abs_max_f32_neon(&a));
    }

    #[test]
    fn test_negate_f32_scalar_vs_neon() {
        let a: Vec<f32> = vec![1.0, -2.0, 3.0];
        let mut c1 = vec![0.0; 3];
        let mut c2 = vec![0.0; 3];
        negate_f32_scalar(&a, &mut c1);
        negate_f32_neon(&a, &mut c2);
        assert_eq!(c1, c2);
    }

    #[test]
    fn test_clamp_f32_scalar_vs_neon() {
        let a: Vec<f32> = vec![-10.0, -5.0, 0.0, 5.0, 10.0];
        let mut c1 = vec![0.0; 5];
        let mut c2 = vec![0.0; 5];
        clamp_f32_scalar(&a, -3.0, 3.0, &mut c1);
        clamp_f32_neon(&a, -3.0, 3.0, &mut c2);
        assert_eq!(c1, c2);
    }

    #[test]
    fn test_euclidean_f32_scalar_vs_neon() {
        let a: Vec<f32> = (0..50).map(|i| i as f32).collect();
        let b: Vec<f32> = (50..100).map(|i| i as f32).collect();
        assert_close_f32!(
            euclidean_distance_f32_scalar(&a, &b),
            euclidean_distance_f32_neon(&a, &b)
        );
    }

    #[test]
    fn test_add_f32_scalar_vs_neon() {
        let a: Vec<f32> = (0..20).map(|i| i as f32).collect();
        let b: Vec<f32> = (20..40).map(|i| i as f32).collect();
        let mut c1 = vec![0.0; 20];
        let mut c2 = vec![0.0; 20];
        add_f32_scalar(&a, &b, &mut c1);
        add_f32_neon(&a, &b, &mut c2);
        for i in 0..20 {
            assert_close_f32!(c1[i], c2[i]);
        }
    }

    #[cfg_attr(not(target_arch = "aarch64"), ignore)]
    #[test]
    fn test_dispatch_select_dot_neon_on_arm() {
        let mut d = Dispatch::<DotKernel>::new("test_neon", vec![
            KernelInfo { name: "scalar", func: dot_scalar, is_supported: |_| true, thermal_priority: 1 },
            KernelInfo { name: "NEON", func: dot_neon, is_supported: |_| true, thermal_priority: 2 },
        ]);
        assert!(d.select().is_ok());
    }

    // -----------------------------------------------------------------------
    // SVE kernels (aarch64 only)
    // -----------------------------------------------------------------------
    #[cfg(target_arch = "aarch64")]
    #[test]
    fn test_dot_f64_scalar_vs_sve() {
        let a: Vec<f64> = (0..100).map(|i| i as f64).collect();
        let b: Vec<f64> = (100..200).map(|i| i as f64).collect();
        let expected = dot_scalar(&a, &b);
        let sve = dot_f64_sve(&a, &b);
        assert_close_f64!(expected, sve);
    }

    #[cfg(target_arch = "aarch64")]
    #[test]
    fn test_dot_f64_sve_empty() {
        assert_close_f64!(dot_f64_sve(&[], &[]), 0.0);
    }

    #[cfg(target_arch = "aarch64")]
    #[test]
    fn test_dot_f32_scalar_vs_sve() {
        let a: Vec<f32> = (0..100).map(|i| i as f32).collect();
        let b: Vec<f32> = (100..200).map(|i| i as f32).collect();
        assert_close_f32!(dot_f32_scalar(&a, &b), dot_f32_sve(&a, &b));
    }

    #[cfg(target_arch = "aarch64")]
    #[test]
    fn test_reduce_sum_f64_scalar_vs_sve() {
        let a: Vec<f64> = (0..200).map(|i| (i % 7) as f64).collect();
        assert_close_f64!(reduce_sum_scalar(&a), reduce_sum_f64_sve(&a));
    }

    #[cfg(target_arch = "aarch64")]
    #[test]
    fn test_reduce_sum_f32_scalar_vs_sve() {
        let a: Vec<f32> = (0..200).map(|i| (i % 7) as f32).collect();
        assert_close_f32!(reduce_sum_f32_scalar(&a), reduce_sum_f32_sve(&a));
    }

    #[cfg(target_arch = "aarch64")]
    #[test]
    fn test_reduce_max_f64_scalar_vs_sve() {
        let a: Vec<f64> = vec![1.0, 5.0, 3.0, 9.0, 2.0];
        assert_close_f64!(reduce_max_scalar(&a), reduce_max_f64_sve(&a));
    }

    #[cfg(target_arch = "aarch64")]
    #[test]
    fn test_abs_max_f64_scalar_vs_sve() {
        let a: Vec<f64> = vec![-5.0, 3.0, -10.0, 7.0];
        assert_close_f64!(abs_max_f64_scalar(&a), abs_max_f64_sve(&a));
    }

    #[cfg(target_arch = "aarch64")]
    #[test]
    fn test_argmax_f64_scalar_vs_sve() {
        let a: Vec<f64> = vec![1.0, 5.0, 3.0, 9.0, 2.0];
        assert_eq!(argmax_f64_scalar(&a), argmax_f64_sve(&a));
    }

    #[cfg(target_arch = "aarch64")]
    #[test]
    fn test_hadamard_f64_scalar_vs_sve() {
        let a: Vec<f64> = (0..20).map(|i| i as f64).collect();
        let b: Vec<f64> = (20..40).map(|i| i as f64).collect();
        let mut c1 = vec![0.0; 20];
        let mut c2 = vec![0.0; 20];
        hadamard_product_f64_scalar(&a, &b, &mut c1);
        hadamard_product_f64_sve(&a, &b, &mut c2);
        for i in 0..20 {
            assert_close_f64!(c1[i], c2[i]);
        }
    }

    #[cfg(target_arch = "aarch64")]
    #[test]
    fn test_matmul_f64_scalar_vs_sve() {
        const N: usize = 8;
        let a: Vec<f64> = (0..N * N).map(|i| (i % 5) as f64).collect();
        let b: Vec<f64> = (0..N * N).map(|i| (i % 7) as f64).collect();
        let mut c1 = vec![0.0; N * N];
        let mut c2 = vec![0.0; N * N];
        matmul_scalar(&a, &b, &mut c1, N);
        matmul_f64_sve(&a, &b, &mut c2, N);
        for i in 0..N * N {
            assert_close_f64!(c1[i], c2[i]);
        }
    }

    #[cfg(target_arch = "aarch64")]
    #[test]
    fn test_softmax_f64_scalar_vs_sve() {
        let input: Vec<f64> = vec![1.0, 2.0, 3.0, 4.0, 5.0];
        let mut c1 = vec![0.0; 5];
        let mut c2 = vec![0.0; 5];
        softmax_scalar(&input, &mut c1);
        softmax_f64_sve(&input, &mut c2);
        for i in 0..5 {
            assert_close_f64!(c1[i], c2[i]);
        }
    }

    #[cfg(target_arch = "aarch64")]
    #[test]
    fn test_clamp_f64_scalar_vs_sve() {
        let a: Vec<f64> = vec![-10.0, -5.0, 0.0, 5.0, 10.0];
        let mut c1 = vec![0.0; 5];
        let mut c2 = vec![0.0; 5];
        clamp_f64_scalar(&a, -3.0, 3.0, &mut c1);
        clamp_f64_sve(&a, -3.0, 3.0, &mut c2);
        assert_eq!(c1, c2);
    }

    #[cfg(target_arch = "aarch64")]
    #[test]
    fn test_negate_f64_scalar_vs_sve() {
        let a: Vec<f64> = vec![1.0, -2.0, 3.0, -4.0];
        let mut c1 = vec![0.0; 4];
        let mut c2 = vec![0.0; 4];
        negate_f64_scalar(&a, &mut c1);
        negate_f64_sve(&a, &mut c2);
        assert_eq!(c1, c2);
    }

    #[cfg(target_arch = "aarch64")]
    #[test]
    fn test_euclidean_f64_scalar_vs_sve() {
        let a: Vec<f64> = (0..50).map(|i| i as f64).collect();
        let b: Vec<f64> = (50..100).map(|i| i as f64).collect();
        assert_close_f64!(
            euclidean_distance_f64_scalar(&a, &b),
            euclidean_distance_f64_sve(&a, &b)
        );
    }

    #[cfg(target_arch = "aarch64")]
    #[test]
    fn test_cosine_similarity_f64_scalar_vs_sve() {
        let a: Vec<f64> = vec![1.0, 0.0];
        let b: Vec<f64> = vec![0.0, 1.0];
        assert_close_f64!(cosine_similarity_f64_scalar(&a, &b), cosine_similarity_f64_sve(&a, &b));
    }

    #[cfg(target_arch = "aarch64")]
    #[test]
    fn test_memchr_sve() {
        let data: Vec<u8> = (0..200).collect();
        assert_eq!(memchr_scalar(42, &data), memchr_sve(42, &data));
        assert_eq!(memchr_scalar(255, &data), memchr_sve(255, &data));
    }

    // -----------------------------------------------------------------------
    // Online Learning: observe, rebench, TimingHistory
    // -----------------------------------------------------------------------
    #[test]
    fn test_timing_history_basic() {
        let mut th = TimingHistory::new(4);
        assert!(th.is_empty());
        th.record(Duration::from_micros(100));
        th.record(Duration::from_micros(200));
        assert_eq!(th.len(), 2);
        assert_eq!(th.avg(), Duration::from_micros(150));
        assert_eq!(th.min(), Duration::from_micros(100));
    }

    #[test]
    fn test_timing_history_ring() {
        let mut th = TimingHistory::new(3);
        for i in 0..5 {
            th.record(Duration::from_micros(i * 100));
        }
        assert_eq!(th.len(), 3);
        assert_eq!(th.avg(), Duration::from_micros(300));
        assert_eq!(th.min(), Duration::from_micros(200));
    }

    #[test]
    fn test_timing_history_trend() {
        let mut th = TimingHistory::new(10);
        for i in 1..=8 {
            th.record(Duration::from_micros(i * 100));
        }
        let trend = th.trend();
        assert!(trend > 0.0, "expected positive trend, got {}", trend);
    }

    #[test]
    fn test_dispatch_observe() {
        let mut d = Dispatch::<DotKernel>::new("test_obs", vec![
            KernelInfo { name: "scalar", func: dot_scalar, is_supported: |_| true, thermal_priority: 1 },
        ]);
        d.observe(Duration::from_micros(500));
        d.observe(Duration::from_micros(600));
        assert_eq!(d.observation_count, 2);
        assert_eq!(d.timing_avg(), Duration::from_micros(550));
    }

    #[test]
    fn test_dispatch_rebench() {
        let mut d = Dispatch::<DotKernel>::new("test_rebench", vec![
            KernelInfo { name: "scalar", func: dot_scalar, is_supported: |_| true, thermal_priority: 1 },
        ]);
        assert!(d.rebench().is_ok());
        assert_eq!(d.selected_name(), "scalar");
    }

    // -----------------------------------------------------------------------
    // Pipeline
    // -----------------------------------------------------------------------
    #[test]
    fn test_pipeline_add_and_optimize() {
        let mut pipe = Pipeline::new();
        {
            let dot = pipe.add_dot("dot1");
            dot.supported.push(KernelInfo {
                name: "scalar",
                func: dot_scalar,
                is_supported: |_| true,
                thermal_priority: 1,
            });
        }
        {
            let softmax = pipe.add_softmax("sm1");
            softmax.supported.push(KernelInfo {
                name: "scalar",
                func: softmax_scalar,
                is_supported: |_| true,
                thermal_priority: 1,
            });
        }
        let selected = pipe.optimize().unwrap();
        assert_eq!(pipe.dot_count(), 1);
        assert_eq!(pipe.softmax_count(), 1);
        assert_eq!(pipe.reduce_count(), 0);
        assert!(selected >= 2);
        assert_eq!(pipe.dots[0].selected_name(), "scalar");
    }

    #[test]
    fn test_pipeline_execute() {
        let mut pipe = Pipeline::new();
        {
            let dot = pipe.add_dot("dot_ex");
            dot.supported.push(KernelInfo {
                name: "scalar",
                func: dot_scalar,
                is_supported: |_| true,
                thermal_priority: 1,
            });
        }
        let count = pipe.execute().unwrap();
        assert!(count >= 1);
    }

    // -----------------------------------------------------------------------
    // Heterogeneous dispatch
    // -----------------------------------------------------------------------
    #[test]
    fn test_heterogeneous_dispatch() {
        let d = Dispatch::<DotKernel>::new("het_test", vec![
            KernelInfo { name: "scalar", func: dot_scalar, is_supported: |_| true, thermal_priority: 1 },
        ]);
        let mut het = HeterogeneousDispatch::new(d);
        assert_eq!(het.preferred_backend(), Backend::Auto);
        het.set_preferred_backend(Backend::Metal);
        assert_eq!(het.preferred_backend(), Backend::Metal);
        assert_eq!(het.cpu_dispatch().selected_name(), "unselected");
    }

    // -----------------------------------------------------------------------
    // Compile-time dispatch (DispatchFixed)
    // -----------------------------------------------------------------------
    #[test]
    fn test_dispatch_fixed() {
        let kernels = vec![
            KernelInfo { name: "scalar", func: dot_scalar as DotKernel, is_supported: |_| true, thermal_priority: 1 },
            KernelInfo { name: "neon", func: dot_neon as DotKernel, is_supported: |_| true, thermal_priority: 2 },
        ];
        let fixed = DispatchFixed::<DotKernel, 0>::new(&kernels);
        assert_eq!(fixed.kernel_name(), "scalar");
        assert_eq!(fixed.thermal_priority(), 1);

        let fixed2 = DispatchFixed::<DotKernel, 1>::new(&kernels);
        assert_eq!(fixed2.kernel_name(), "neon");
    }

    #[test]
    #[should_panic(expected = "out of bounds")]
    fn test_dispatch_fixed_oob() {
        let kernels = vec![
            KernelInfo { name: "scalar", func: dot_scalar as DotKernel, is_supported: |_| true, thermal_priority: 1 },
        ];
        let _ = DispatchFixed::<DotKernel, 5>::new(&kernels);
    }

    #[test]
    fn test_dispatch_fixed_compute() {
        let kernels = vec![
            KernelInfo { name: "scalar", func: dot_scalar as DotKernel, is_supported: |_| true, thermal_priority: 1 },
        ];
        let fixed = DispatchFixed::<DotKernel, 0>::new(&kernels);
        let a = vec![1.0, 2.0];
        let b = vec![3.0, 4.0];
        let result = (fixed.kernel_func())(&a, &b);
        assert_close_f64!(result, 11.0);
    }

    // -----------------------------------------------------------------------
    // Step trait
    // -----------------------------------------------------------------------
    #[test]
    fn test_step_trait() {
        let mut d = Dispatch::<DotKernel>::new("step_test", vec![
            KernelInfo { name: "scalar", func: dot_scalar, is_supported: |_| true, thermal_priority: 1 },
        ]);
        assert_eq!(d.op_name(), "step_test");
        assert_eq!(d.selected_kernel_name(), "unselected");
        assert!(d.select_step().is_ok());
        assert_eq!(d.selected_kernel_name(), "scalar");
        d.observe_step(Duration::from_micros(100));
        assert_eq!(d.timing_avg(), Duration::from_micros(100));
    }

    // -----------------------------------------------------------------------
    // Int32 operations
    // -----------------------------------------------------------------------
    #[test]
    fn test_dot_i32_scalar_vs_neon() {
        let a: Vec<i32> = (0..100).map(|i| i as i32).collect();
        let b: Vec<i32> = (100..200).map(|i| i as i32).collect();
        let expected = dot_i32_scalar(&a, &b);
        let neon = dot_i32_neon(&a, &b);
        assert_eq!(expected, neon);
    }

    #[test]
    fn test_dot_i32_empty() {
        assert_eq!(dot_i32_scalar(&[], &[]), 0);
        assert_eq!(dot_i32_neon(&[], &[]), 0);
    }

    #[test]
    fn test_reduce_sum_i32_scalar_vs_neon() {
        let a: Vec<i32> = (0..200).map(|i| (i % 7) as i32).collect();
        assert_eq!(reduce_sum_i32_scalar(&a), reduce_sum_i32_neon(&a));
    }

    #[test]
    fn test_reduce_max_i32_scalar_vs_neon() {
        let a: Vec<i32> = vec![1, 5, 3, 9, 2];
        assert_eq!(reduce_max_i32_scalar(&a), reduce_max_i32_neon(&a));
    }

    #[test]
    fn test_reduce_max_i32_negative() {
        let a: Vec<i32> = vec![-5, -3, -10, -1];
        assert_eq!(reduce_max_i32_scalar(&a), reduce_max_i32_neon(&a));
    }

    #[test]
    fn test_abs_max_i32_scalar_vs_neon() {
        let a: Vec<i32> = vec![-5, 3, -10, 8, -2];
        assert_eq!(abs_max_i32_scalar(&a), abs_max_i32_neon(&a));
    }

    #[test]
    fn test_abs_max_i32_empty() {
        assert_eq!(abs_max_i32_scalar(&[]), 0);
        assert_eq!(abs_max_i32_neon(&[]), 0);
    }

    #[test]
    fn test_argmax_i32_scalar_vs_neon() {
        let a: Vec<i32> = vec![10, 30, 20, 50, 40];
        assert_eq!(argmax_i32_scalar(&a), argmax_i32_neon(&a));
    }

    #[test]
    fn test_argmax_i32_empty() {
        assert_eq!(argmax_i32_scalar(&[]), 0);
        assert_eq!(argmax_i32_neon(&[]), 0);
    }

    #[test]
    fn test_hadamard_i32_scalar_vs_neon() {
        let a: Vec<i32> = (0..100).map(|i| i as i32).collect();
        let b: Vec<i32> = (0..100).map(|i| (i * 2) as i32).collect();
        let mut c1 = vec![0i32; 100];
        let mut c2 = vec![0i32; 100];
        hadamard_product_i32_scalar(&a, &b, &mut c1);
        hadamard_product_i32_neon(&a, &b, &mut c2);
        assert_eq!(c1, c2);
    }

    #[test]
    fn test_add_i32_scalar_vs_neon() {
        let a: Vec<i32> = (0..100).map(|i| i as i32).collect();
        let b: Vec<i32> = (0..100).map(|i| (i * 2) as i32).collect();
        let mut c1 = vec![0i32; 100];
        let mut c2 = vec![0i32; 100];
        add_i32_scalar(&a, &b, &mut c1);
        add_i32_neon(&a, &b, &mut c2);
        assert_eq!(c1, c2);
    }

    #[test]
    fn test_sub_i32_scalar_vs_neon() {
        let a: Vec<i32> = (0..100).map(|i| i as i32).collect();
        let b: Vec<i32> = (0..100).map(|i| (i * 2) as i32).collect();
        let mut c1 = vec![0i32; 100];
        let mut c2 = vec![0i32; 100];
        sub_i32_scalar(&a, &b, &mut c1);
        sub_i32_neon(&a, &b, &mut c2);
        assert_eq!(c1, c2);
    }

    #[test]
    fn test_mul_i32_scalar_vs_neon() {
        let a: Vec<i32> = (0..100).map(|i| i as i32).collect();
        let b: Vec<i32> = (0..100).map(|i| (i * 2) as i32).collect();
        let mut c1 = vec![0i32; 100];
        let mut c2 = vec![0i32; 100];
        mul_i32_scalar(&a, &b, &mut c1);
        mul_i32_neon(&a, &b, &mut c2);
        assert_eq!(c1, c2);
    }

    // -----------------------------------------------------------------------
    // Int64 operations
    // -----------------------------------------------------------------------
    #[test]
    fn test_dot_i64_scalar_vs_neon() {
        let a: Vec<i64> = (0..100).map(|i| i as i64).collect();
        let b: Vec<i64> = (100..200).map(|i| i as i64).collect();
        assert_eq!(dot_i64_scalar(&a, &b), dot_i64_neon(&a, &b));
    }

    #[test]
    fn test_reduce_sum_i64_scalar_vs_neon() {
        let a: Vec<i64> = (0..200).map(|i| (i % 7) as i64).collect();
        assert_eq!(reduce_sum_i64_scalar(&a), reduce_sum_i64_neon(&a));
    }

    #[test]
    fn test_reduce_max_i64_scalar_vs_neon() {
        let a: Vec<i64> = vec![1, 5, 3, 9, 2];
        assert_eq!(reduce_max_i64_scalar(&a), reduce_max_i64_neon(&a));
    }

    #[test]
    fn test_abs_max_i64_scalar_vs_neon() {
        let a: Vec<i64> = vec![-5, 3, -10, 8, -2];
        assert_eq!(abs_max_i64_scalar(&a), abs_max_i64_neon(&a));
    }

    #[test]
    fn test_argmax_i64_scalar_vs_neon() {
        let a: Vec<i64> = vec![10, 30, 20, 50, 40];
        assert_eq!(argmax_i64_scalar(&a), argmax_i64_neon(&a));
    }

    #[test]
    fn test_add_i64_scalar_vs_neon() {
        let a: Vec<i64> = (0..100).map(|i| i as i64).collect();
        let b: Vec<i64> = (0..100).map(|i| (i * 2) as i64).collect();
        let mut c1 = vec![0i64; 100];
        let mut c2 = vec![0i64; 100];
        add_i64_scalar(&a, &b, &mut c1);
        add_i64_neon(&a, &b, &mut c2);
        assert_eq!(c1, c2);
    }

    // -----------------------------------------------------------------------
    // Trig operations
    // -----------------------------------------------------------------------
    #[test]
    fn test_sin_batch_f64_scalar_vs_neon() {
        let a: Vec<f64> = (0..50).map(|i| (i as f64) * 0.5).collect();
        let mut c1 = vec![0.0; 50];
        let mut c2 = vec![0.0; 50];
        sin_batch_f64_scalar(&a, &mut c1);
        sin_batch_f64_neon(&a, &mut c2);
        for i in 0..50 { assert_close_f64!(c1[i], c2[i]); }
    }

    #[test]
    fn test_cos_batch_f64_scalar_vs_neon() {
        let a: Vec<f64> = (0..50).map(|i| (i as f64) * 0.5).collect();
        let mut c1 = vec![0.0; 50];
        let mut c2 = vec![0.0; 50];
        cos_batch_f64_scalar(&a, &mut c1);
        cos_batch_f64_neon(&a, &mut c2);
        for i in 0..50 { assert_close_f64!(c1[i], c2[i]); }
    }

    #[test]
    fn test_tan_batch_f64_scalar_vs_neon() {
        let a: Vec<f64> = (0..20).map(|i| (i as f64) * 0.3).collect();
        let mut c1 = vec![0.0; 20];
        let mut c2 = vec![0.0; 20];
        tan_batch_f64_scalar(&a, &mut c1);
        tan_batch_f64_neon(&a, &mut c2);
        for i in 0..20 { assert_close_f64!(c1[i], c2[i]); }
    }

    #[test]
    fn test_sin_batch_f32_scalar_vs_neon() {
        let a: Vec<f32> = (0..50).map(|i| (i as f32) * 0.5).collect();
        let mut c1 = vec![0.0; 50];
        let mut c2 = vec![0.0; 50];
        sin_batch_f32_scalar(&a, &mut c1);
        sin_batch_f32_neon(&a, &mut c2);
        for i in 0..50 { assert_close_f32!(c1[i], c2[i]); }
    }

    #[test]
    fn test_cos_batch_f32_scalar_vs_neon() {
        let a: Vec<f32> = (0..50).map(|i| (i as f32) * 0.5).collect();
        let mut c1 = vec![0.0; 50];
        let mut c2 = vec![0.0; 50];
        cos_batch_f32_scalar(&a, &mut c1);
        cos_batch_f32_neon(&a, &mut c2);
        for i in 0..50 { assert_close_f32!(c1[i], c2[i]); }
    }

    // -----------------------------------------------------------------------
    // GEMV operations
    // -----------------------------------------------------------------------
    #[test]
    fn test_gemv_f64_scalar_vs_neon() {
        let a: Vec<f64> = vec![1.0, 2.0, 3.0, 4.0]; // 2x2 row-major
        let x = vec![2.0, 3.0];
        let mut y1 = vec![0.0; 2];
        let mut y2 = vec![0.0; 2];
        gemv_f64_scalar(1.0, &a, &x, 0.0, &mut y1, 2, 2);
        gemv_f64_neon(1.0, &a, &x, 0.0, &mut y2, 2, 2);
        assert_close_f64!(y1[0], y2[0]);
        assert_close_f64!(y1[1], y2[1]);
    }

    // -----------------------------------------------------------------------
    // Conv1D operations
    // -----------------------------------------------------------------------
    #[test]
    fn test_conv1d_f64_scalar_vs_neon() {
        let input: Vec<f64> = vec![1.0, 2.0, 3.0, 4.0, 5.0];
        let kernel = vec![0.5, 1.0, 0.5];
        let mut c1 = vec![0.0; 3];
        let mut c2 = vec![0.0; 3];
        conv1d_f64_scalar(&input, &kernel, &mut c1);
        conv1d_f64_neon(&input, &kernel, &mut c2);
        assert_close_f64!(c1[0], c2[0]);
        assert_close_f64!(c1[1], c2[1]);
        assert_close_f64!(c1[2], c2[2]);
    }

    // -----------------------------------------------------------------------
    // LayerNorm operations
    // -----------------------------------------------------------------------
    #[test]
    fn test_layer_norm_f64_scalar_vs_neon() {
        let input: Vec<f64> = vec![1.0, 2.0, 3.0, 4.0, 5.0];
        let gamma: Vec<f64> = vec![1.0; 5];
        let beta: Vec<f64> = vec![0.0; 5];
        let mut c1 = vec![0.0; 5];
        let mut c2 = vec![0.0; 5];
        layer_norm_f64_scalar(&input, &gamma, &beta, 1e-8, &mut c1);
        layer_norm_f64_neon(&input, &gamma, &beta, 1e-8, &mut c2);
        for i in 0..5 { assert_close_f64!(c1[i], c2[i]); }
    }

    #[test]
    fn test_layer_norm_f32_scalar_vs_neon() {
        let input: Vec<f32> = vec![1.0, 2.0, 3.0, 4.0, 5.0];
        let gamma: Vec<f32> = vec![1.0; 5];
        let beta: Vec<f32> = vec![0.0; 5];
        let mut c1 = vec![0.0; 5];
        let mut c2 = vec![0.0; 5];
        layer_norm_f32_scalar(&input, &gamma, &beta, 1e-8, &mut c1);
        layer_norm_f32_neon(&input, &gamma, &beta, 1e-8, &mut c2);
        for i in 0..5 { assert_close_f32!(c1[i], c2[i]); }
    }

    // -----------------------------------------------------------------------
    // Cumsum operations
    // -----------------------------------------------------------------------
    #[test]
    fn test_cumsum_f64_scalar_vs_neon() {
        let input: Vec<f64> = vec![1.0, 2.0, 3.0, 4.0, 5.0];
        let mut c1 = vec![0.0; 5];
        let mut c2 = vec![0.0; 5];
        cumsum_f64_scalar(&input, &mut c1);
        cumsum_f64_neon(&input, &mut c2);
        for i in 0..5 { assert_close_f64!(c1[i], c2[i]); }
    }

    #[test]
    fn test_cumsum_f32_scalar_vs_neon() {
        let input: Vec<f32> = vec![1.0, 2.0, 3.0, 4.0, 5.0];
        let mut c1 = vec![0.0; 5];
        let mut c2 = vec![0.0; 5];
        cumsum_f32_scalar(&input, &mut c1);
        cumsum_f32_neon(&input, &mut c2);
        for i in 0..5 { assert_close_f32!(c1[i], c2[i]); }
    }

    #[test]
    fn test_cumsum_i32_scalar_vs_neon() {
        let input: Vec<i32> = vec![1, 2, 3, 4, 5];
        let mut c1 = vec![0i32; 5];
        let mut c2 = vec![0i32; 5];
        cumsum_i32_scalar(&input, &mut c1);
        cumsum_i32_neon(&input, &mut c2);
        assert_eq!(c1, c2);
    }
}

#[cfg(test)]
#[path = "proptests.rs"]
mod proptests;