himada-dispatch 0.1.1

Adaptive SIMD dispatch for Himada — auto-selects fastest kernel at runtime
//! Auto-tuning cache — persist kernel selection results to `~/.cache/himada/`.
//!
//! On first run, dispatch benchmarks all kernels and caches the winner.
//! On subsequent runs, the cached selection is loaded immediately,
//! skipping the benchmark entirely.

use serde::{Serialize, Deserialize};
use std::path::PathBuf;
use std::sync::Mutex;

/// Cached dispatch configuration.
#[derive(Serialize, Deserialize)]
pub struct DispatchCache {
    pub entries: Vec<CachedEntry>,
}

#[derive(Serialize, Deserialize)]
pub struct CachedEntry {
    pub dispatch_name: String,
    pub kernel_name: String,
    pub cpu_name: String,
}

impl DispatchCache {
    pub fn new() -> Self {
        Self { entries: Vec::new() }
    }

    fn cache_dir() -> PathBuf {
        if let Ok(dir) = std::env::var("HIMADA_CACHE_DIR") {
            return PathBuf::from(dir);
        }
        let base = dirs::cache_dir().unwrap_or_else(|| PathBuf::from("."));
        base.join("himada")
    }

    fn cache_path() -> PathBuf {
        Self::cache_dir().join("dispatch_cache.json")
    }

    pub fn load() -> Self {
        let path = Self::cache_path();
        if !path.exists() {
            return Self::new();
        }
        std::fs::read_to_string(&path)
            .ok()
            .and_then(|s| serde_json::from_str(&s).ok())
            .unwrap_or_else(Self::new)
    }

    pub fn save(&self) {
        if let Ok(_) = std::fs::create_dir_all(Self::cache_dir()) {
            if let Ok(s) = serde_json::to_string_pretty(self) {
                let _ = std::fs::write(Self::cache_path(), s);
            }
        }
    }

    pub fn lookup(&self, dispatch_name: &str, cpu_name: &str) -> Option<&str> {
        for entry in &self.entries {
            if entry.dispatch_name == dispatch_name && entry.cpu_name == cpu_name {
                return Some(&entry.kernel_name);
            }
        }
        None
    }

    pub fn insert(&mut self, dispatch_name: String, kernel_name: String, cpu_name: String) {
        self.entries.retain(|e| e.dispatch_name != dispatch_name);
        self.entries.push(CachedEntry { dispatch_name, kernel_name, cpu_name });
    }
}

static CACHE_LOCK: Mutex<()> = Mutex::new(());

/// Get the CPU brand string for cache key.
pub fn cpu_brand_string() -> String {
    #[cfg(target_arch = "x86_64")]
    {
        let mut brand = String::new();
        unsafe {
            for leaf in [0x80000002u32, 0x80000003, 0x80000004] {
                let result = std::arch::x86_64::__cpuid(leaf);
                brand.extend(std::str::from_utf8(&[
                    result.eax as u8, (result.eax >> 8) as u8, (result.eax >> 16) as u8, (result.eax >> 24) as u8,
                    result.ebx as u8, (result.ebx >> 8) as u8, (result.ebx >> 16) as u8, (result.ebx >> 24) as u8,
                    result.ecx as u8, (result.ecx >> 8) as u8, (result.ecx >> 16) as u8, (result.ecx >> 24) as u8,
                    result.edx as u8, (result.edx >> 8) as u8, (result.edx >> 16) as u8, (result.edx >> 24) as u8,
                ]).unwrap_or(""));
            }
        }
        brand.trim().to_string()
    }
    #[cfg(target_arch = "aarch64")]
    {
        "arm64".to_string()
    }
    #[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
    {
        "unknown".to_string()
    }
}

/// Try to load a cached kernel selection for this dispatch.
/// Returns `Some(kernel_name)` if found.
pub fn load_cached_selection(dispatch_name: &str) -> Option<String> {
    let _lock = CACHE_LOCK.lock().ok()?;
    let cache = DispatchCache::load();
    let cpu_name = cpu_brand_string();
    cache.lookup(dispatch_name, &cpu_name).map(|s| s.to_string())
}

/// Save a kernel selection to the cache.
pub fn save_cached_selection(dispatch_name: &str, kernel_name: &str) {
    let _lock = CACHE_LOCK.lock().ok();
    let mut cache = DispatchCache::load();
    let cpu_name = cpu_brand_string();
    cache.insert(dispatch_name.to_string(), kernel_name.to_string(), cpu_name);
    cache.save();
}