rustytorch/crates/specialized/rtx-polygraph/src/cache.rs

use crate::fusion::FusionOpportunity;
use crate::ir::IRNode;
use rustc_hash::FxHasher;
use std::collections::HashMap;
use std::hash::{Hash, Hasher};
use std::time::Duration;

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct CacheKey {
    hash: u64,
}

impl CacheKey {
    pub fn from_node(node: &IRNode) -> Self {
        let mut hasher = FxHasher::default();

        // Hash operation type and configuration
        node.node_type().hash(&mut hasher);
        node.output_dtypes().hash(&mut hasher);

        // Hash output shapes
        for shape in node.output_shapes() {
            shape.dims().hash(&mut hasher);
        }

        // Hash input count (but not specific IDs for reusability)
        node.inputs().len().hash(&mut hasher);

        Self {
            hash: hasher.finish(),
        }
    }

    pub fn from_fusion(fusion: &FusionOpportunity) -> Self {
        let mut hasher = FxHasher::default();

        // Hash fusion type
        fusion.fusion_type().hash(&mut hasher);

        // Hash participating node signatures
        for node in fusion.participating_nodes() {
            let node_key = Self::from_node(node);
            node_key.hash.hash(&mut hasher);
        }

        Self {
            hash: hasher.finish(),
        }
    }

    pub fn hash(&self) -> String {
        format!("{:016x}", self.hash)
    }
}

#[derive(Debug, Clone)]
pub struct CachedKernel {
    name: String,
    binary: Vec<u8>,
    compile_time: Duration,
}

impl CachedKernel {
    pub fn new(name: String, binary: Vec<u8>, compile_time: Duration) -> Self {
        Self {
            name,
            binary,
            compile_time,
        }
    }

    pub fn name(&self) -> &str {
        &self.name
    }

    pub fn binary(&self) -> &[u8] {
        &self.binary
    }

    pub fn compile_time(&self) -> Duration {
        self.compile_time
    }
}

#[derive(Debug, Clone, Default)]
pub struct CacheStatistics {
    hits: u64,
    misses: u64,
}

impl CacheStatistics {
    pub fn hits(&self) -> u64 {
        self.hits
    }

    pub fn misses(&self) -> u64 {
        self.misses
    }

    pub fn total(&self) -> u64 {
        self.hits + self.misses
    }

    pub fn hit_rate(&self) -> f64 {
        if self.total() == 0 {
            0.0
        } else {
            self.hits as f64 / self.total() as f64
        }
    }
}

pub struct KernelCache {
    cache: HashMap<CacheKey, CachedKernel>,
    access_order: Vec<CacheKey>, // For LRU eviction
    capacity: Option<usize>,
    stats: CacheStatistics,
}

impl KernelCache {
    pub fn new() -> Self {
        Self {
            cache: HashMap::new(),
            access_order: Vec::new(),
            capacity: None,
            stats: CacheStatistics::default(),
        }
    }

    pub fn with_capacity(capacity: usize) -> Self {
        Self {
            cache: HashMap::new(),
            access_order: Vec::new(),
            capacity: Some(capacity),
            stats: CacheStatistics::default(),
        }
    }

    pub fn get(&mut self, key: &CacheKey) -> Option<&CachedKernel> {
        if let Some(kernel) = self.cache.get(key) {
            self.stats.hits += 1;

            // Update LRU order
            if let Some(pos) = self.access_order.iter().position(|k| k == key) {
                self.access_order.remove(pos);
            }
            self.access_order.push(key.clone());

            Some(kernel)
        } else {
            self.stats.misses += 1;
            None
        }
    }

    pub fn insert(&mut self, key: CacheKey, kernel: CachedKernel) {
        // Check if we need to evict
        if let Some(capacity) = self.capacity {
            while self.cache.len() >= capacity {
                if let Some(lru_key) = self.access_order.first().cloned() {
                    self.cache.remove(&lru_key);
                    self.access_order.remove(0);
                } else {
                    break;
                }
            }
        }

        // Insert new entry
        self.cache.insert(key.clone(), kernel);

        // Update access order
        if let Some(pos) = self.access_order.iter().position(|k| k == &key) {
            self.access_order.remove(pos);
        }
        self.access_order.push(key);
    }

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

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

    pub fn statistics(&self) -> &CacheStatistics {
        &self.stats
    }

    pub fn clear(&mut self) {
        self.cache.clear();
        self.access_order.clear();
    }
}

impl Default for KernelCache {
    fn default() -> Self {
        Self::new()
    }
}