himada-dispatch 0.1.1

//! Shape specialization — pick different kernels based on input size.
//!
//! Some kernels are faster for small batches (e.g., direct scalar),
//! others for large batches (e.g., SIMD, GPU). Shape specialization
//! automatically selects the best kernel for the given size.

use std::time::Duration;

/// A size threshold with associated kernel index.
#[derive(Clone)]
pub struct SizeRule {
    /// Minimum input size for this rule to apply.
    pub min_size: usize,
    /// Kernel index to use.
    pub kernel_idx: usize,
}

/// A dispatch that selects kernels based on input size.
pub struct SizeAwareDispatch {
    /// Kernel name for display.
    pub name: String,
    /// Rules sorted by min_size ascending.
    pub rules: Vec<SizeRule>,
    /// Default kernel index (for very small inputs).
    pub default_idx: usize,
}

impl SizeAwareDispatch {
    pub fn new(name: &str, default_idx: usize) -> Self {
        Self {
            name: name.into(),
            rules: Vec::new(),
            default_idx,
        }
    }

    pub fn add_rule(&mut self, min_size: usize, kernel_idx: usize) {
        self.rules.push(SizeRule { min_size, kernel_idx });
        self.rules.sort_by_key(|r| r.min_size);
    }

    pub fn select(&self, size: usize) -> usize {
        for rule in self.rules.iter().rev() {
            if size >= rule.min_size {
                return rule.kernel_idx;
            }
        }
        self.default_idx
    }
}

/// Adaptive threshold tuner — adjusts GPU/CPU thresholds based on observations.
pub struct ThresholdTuner {
    cpu_times: Vec<Duration>,
    gpu_times: Vec<Duration>,
    current_threshold: usize,
    min_threshold: usize,
    max_threshold: usize,
    adaptation_rate: f64,
}

impl ThresholdTuner {
    pub fn new(initial: usize) -> Self {
        Self {
            cpu_times: Vec::with_capacity(32),
            gpu_times: Vec::with_capacity(32),
            current_threshold: initial,
            min_threshold: 64,
            max_threshold: 1_000_000,
            adaptation_rate: 0.1,
        }
    }

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

    pub fn observe_cpu(&mut self, elapsed: Duration) {
        self.cpu_times.push(elapsed);
        if self.cpu_times.len() > 64 { self.cpu_times.remove(0); }
    }

    pub fn observe_gpu(&mut self, elapsed: Duration) {
        self.gpu_times.push(elapsed);
        if self.gpu_times.len() > 64 { self.gpu_times.remove(0); }
    }

    /// Adjust threshold based on observed CPU vs GPU performance.
    /// If GPU is faster on average, lower the threshold (use GPU for more sizes).
    /// If CPU is faster on average, raise the threshold.
    pub fn adapt(&mut self) {
        if self.cpu_times.len() < 4 || self.gpu_times.len() < 4 {
            return;
        }
        let cpu_avg: Duration = self.cpu_times.iter().sum::<Duration>() / self.cpu_times.len() as u32;
        let gpu_avg: Duration = self.gpu_times.iter().sum::<Duration>() / self.gpu_times.len() as u32;

        let ratio = cpu_avg.as_secs_f64() / gpu_avg.as_secs_f64();
        if ratio > 1.2 {
            // GPU is faster by >20% — lower threshold to use GPU more
            self.current_threshold = ((self.current_threshold as f64) * (1.0 - self.adaptation_rate)) as usize;
            self.current_threshold = self.current_threshold.max(self.min_threshold);
        } else if ratio < 0.8 {
            // CPU is faster — raise threshold to use CPU more
            self.current_threshold = ((self.current_threshold as f64) * (1.0 + self.adaptation_rate)) as usize;
            self.current_threshold = self.current_threshold.max(self.min_threshold).min(self.max_threshold);
        }
    }
}