weirflow 0.1.0

GPU-first dataflow analysis primitives for Vyre and Santh compiler pipelines.
Documentation
use super::{
    entry::FixedPointResidentGraphCacheEntry, retained_graph_bytes,
    stats::FixedPointResidentGraphCacheStats,
};
use crate::fixed_point_graph::FixedPointForwardGraph;
use crate::fixed_point_resident::FixedPointResidentGraph;
use crate::resident_cache_identity::{
    ResidentGraphCacheIdentity, ResidentGraphCacheMissEvidence, ResidentGraphCacheMissReason,
};
use crate::resident_cache_lru::ResidentCacheLru;
use std::collections::HashMap;
use vyre::ResidentGraphReuseTelemetry;
use vyre_primitives::bitset::bitset_words;

/// Backend-resident graph cache keyed by backend identity and normalized graph
/// layout.
#[derive(Clone, Debug, Default, PartialEq)]
pub struct FixedPointResidentGraphCache {
    entries: HashMap<ResidentGraphCacheIdentity, FixedPointResidentGraphCacheEntry>,
    lru: ResidentCacheLru<ResidentGraphCacheIdentity>,
    max_retained_bytes: Option<usize>,
    retained_bytes: usize,
    clock: u64,
    hits: u64,
    misses: u64,
    resident_uploads: u64,
    resident_graph_reuse: ResidentGraphReuseTelemetry,
    evictions: u64,
    last_miss_reason: Option<ResidentGraphCacheMissReason>,
}

impl FixedPointResidentGraphCache {
    /// Create an empty resident graph cache.
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }

    /// Create a resident graph cache with a soft upper bound on retained bytes.
    ///
    /// The newest graph is always retained so one oversized graph may exceed
    /// the limit; older least-recently-used entries are evicted first.
    #[must_use]
    pub fn with_max_retained_bytes(max_retained_bytes: usize) -> Self {
        Self {
            max_retained_bytes: Some(max_retained_bytes),
            ..Self::default()
        }
    }

    /// Return a resident graph for `graph`, uploading it only on a cache miss.
    pub fn get_or_upload(
        &mut self,
        backend: &dyn vyre::VyreBackend,
        graph: &FixedPointForwardGraph,
    ) -> Result<&FixedPointResidentGraph, vyre::BackendError> {
        let backend_id = backend.id();
        let backend_version = backend.version();
        let layout_hash = graph.stable_layout_hash();
        let identity = ResidentGraphCacheIdentity::fixed_point(
            backend_id,
            backend_version,
            graph.kind(),
            layout_hash,
            graph.node_count(),
            graph.edge_count(),
            bitset_words(graph.node_count()),
        );
        if let Some(entry) = self.entries.get(&identity) {
            let retained_bytes = entry.retained_bytes;
            if entry.graph.node_count() != graph.node_count()
                || entry.graph.edge_count() != graph.edge_count()
                || entry.graph.kind() != graph.kind()
                || entry.graph.stable_layout_hash() != layout_hash
            {
                return Err(vyre::BackendError::InvalidProgram {
                    fix: format!(
                        "Fix: resident fixed-point graph cache entry for backend `{}` version `{}` has layout_hash={} nodes={} edges={} kind={:?} but requested layout_hash={} nodes={} edges={} kind={:?}; rebuild the resident graph cache before dispatch.",
                        backend_id,
                        backend_version,
                        entry.graph.stable_layout_hash(),
                        entry.graph.node_count(),
                        entry.graph.edge_count(),
                        entry.graph.kind(),
                        layout_hash,
                        graph.node_count(),
                        graph.edge_count(),
                        graph.kind()
                    ),
                });
            }
            self.last_miss_reason = None;
            self.hits = self
                .hits
                .checked_add(1)
                .ok_or_else(|| resident_graph_cache_error("Fix: resident graph cache hit counter overflowed u64; rebuild the resident graph cache before reuse."))?;
            let retained_bytes_u64 = u64::try_from(retained_bytes).map_err(|error| {
                resident_graph_cache_error(format!(
                    "Fix: resident fixed-point graph retained byte count does not fit u64 during hit accounting: {error}; shard the graph before resident reuse."
                ))
            })?;
            self.resident_graph_reuse
                .record_warm_reuse(retained_bytes_u64)
                .map_err(|error| resident_graph_cache_error(error.to_string()))?;
            let last_seen = self.next_cache_tick()?;
            self.record_lru(identity.clone(), last_seen)?;
            let entry = self.entries.get_mut(&identity).ok_or_else(|| {
                resident_graph_cache_error("Fix: resident graph cache key disappeared after LRU refresh; rebuild the resident graph cache because this indicates internal cache state corruption.")
            })?;
            entry.last_seen = last_seen;
            return Ok(&entry.graph);
        }
        self.last_miss_reason = Some(self.miss_reason_for_identity(&identity));
        self.misses = self
            .misses
            .checked_add(1)
            .ok_or_else(|| resident_graph_cache_error("Fix: resident graph cache miss counter overflowed u64; rebuild the resident graph cache before reuse."))?;
        let retained_bytes = retained_graph_bytes(graph)?;
        self.evict_until_room(backend, retained_bytes)?;
        let resident = FixedPointResidentGraph::upload(backend, graph)?;
        self.resident_uploads = self
            .resident_uploads
            .checked_add(1)
            .ok_or_else(|| resident_graph_cache_error("Fix: resident graph cache upload counter overflowed u64; rebuild the resident graph cache before reuse."))?;
        let retained_bytes_u64 = u64::try_from(retained_bytes).map_err(|error| {
            resident_graph_cache_error(format!(
                "Fix: resident fixed-point graph retained byte count does not fit u64: {error}; shard the graph before resident upload."
            ))
        })?;
        self.resident_graph_reuse
            .record_cold_upload(retained_bytes_u64)
            .map_err(|error| resident_graph_cache_error(error.to_string()))?;
        let last_seen = self.next_cache_tick()?;
        self.entries.insert(
            identity.clone(),
            FixedPointResidentGraphCacheEntry {
                retained_bytes,
                last_seen,
                graph: resident,
            },
        );
        self.record_lru(identity.clone(), last_seen)?;
        self.retained_bytes = self
            .retained_bytes
            .checked_add(retained_bytes)
            .ok_or_else(|| resident_graph_cache_error("Fix: resident graph cache retained byte accounting overflowed usize; reduce retained graph budget or rebuild the resident graph cache."))?;
        Ok(&self
            .entries
            .get(&identity)
            .ok_or_else(|| resident_graph_cache_error("Fix: resident graph cache key disappeared after upload; rebuild the resident graph cache because this indicates internal cache state corruption."))?
            .graph)
    }

    /// Number of resident graph entries retained by the cache.
    #[must_use]
    pub fn len(&self) -> usize {
        self.entries.len()
    }

    /// Whether the cache is empty.
    #[must_use]
    pub fn is_empty(&self) -> bool {
        self.entries.is_empty()
    }

    /// Return point-in-time resident graph cache counters.
    #[must_use]
    pub fn stats(&self) -> FixedPointResidentGraphCacheStats {
        FixedPointResidentGraphCacheStats {
            hits: self.hits,
            misses: self.misses,
            resident_uploads: self.resident_uploads,
            resident_upload_bytes: self.resident_graph_reuse.upload_bytes,
            resident_avoided_upload_bytes: self.resident_graph_reuse.avoided_upload_bytes,
            evictions: self.evictions,
            retained_bytes: self.retained_bytes,
            entries: self.entries.len(),
        }
    }

    /// Deterministic cache identity snapshot for diagnostics and release
    /// evidence.
    #[must_use]
    pub fn resident_identities(&self) -> Vec<ResidentGraphCacheIdentity> {
        let mut identities = self.entries.keys().cloned().collect::<Vec<_>>();
        identities.sort_by(ResidentGraphCacheIdentity::cmp_stable);
        identities
    }

    /// Explain why `requested` would miss the current resident graph cache.
    #[must_use]
    pub fn miss_reason_for_identity(
        &self,
        requested: &ResidentGraphCacheIdentity,
    ) -> ResidentGraphCacheMissReason {
        ResidentGraphCacheMissReason::classify(ResidentGraphCacheMissEvidence::from_identities(
            self.entries.keys(),
            requested,
        ))
    }

    /// Reason attached to the most recent cache miss. Hits clear this field.
    #[must_use]
    pub const fn last_miss_reason(&self) -> Option<ResidentGraphCacheMissReason> {
        self.last_miss_reason
    }

    fn evict_until_room(
        &mut self,
        backend: &dyn vyre::VyreBackend,
        incoming_bytes: usize,
    ) -> Result<(), vyre::BackendError> {
        let Some(max_retained_bytes) = self.max_retained_bytes else {
            return Ok(());
        };
        while !self.entries.is_empty()
            && self
                .retained_bytes
                .checked_add(incoming_bytes)
                .is_none_or(|total| total > max_retained_bytes)
        {
            let Some(key) = self.pop_lru_key() else {
                return Ok(());
            };
            let entry = self
                .entries
                .remove(&key)
                .ok_or_else(|| resident_graph_cache_error("Fix: resident graph cache eviction key disappeared before removal; rebuild the resident graph cache because this indicates internal cache state corruption."))?;
            self.retained_bytes = self
                .retained_bytes
                .checked_sub(entry.retained_bytes)
                .ok_or_else(|| resident_graph_cache_error("Fix: resident graph cache retained byte accounting underflowed during eviction; rebuild the resident graph cache before reuse."))?;
            self.evictions = self
                .evictions
                .checked_add(1)
                .ok_or_else(|| resident_graph_cache_error("Fix: resident graph cache eviction counter overflowed u64; rebuild the resident graph cache before reuse."))?;
            entry.graph.free(backend)?;
        }
        Ok(())
    }

    /// Free every resident graph owned by this cache.
    pub fn free_all(&mut self, backend: &dyn vyre::VyreBackend) -> Result<(), vyre::BackendError> {
        let mut first_error = None;
        for (_, entry) in self.entries.drain() {
            self.retained_bytes =
                self.retained_bytes
                    .checked_sub(entry.retained_bytes)
                    .ok_or_else(|| vyre::BackendError::InvalidProgram {
                        fix: "Fix: resident graph cache retained byte accounting underflowed while freeing all entries; rebuild the cache before reuse.".to_string(),
                    })?;
            if let Err(error) = entry.graph.free(backend) {
                first_error.get_or_insert(error);
            }
        }
        self.lru.clear();
        if let Some(error) = first_error {
            Err(error)
        } else {
            Ok(())
        }
    }

    fn next_cache_tick(&mut self) -> Result<u64, vyre::BackendError> {
        self.clock = self
            .clock
            .checked_add(1)
            .ok_or_else(|| resident_graph_cache_error("Fix: resident graph cache logical clock overflowed u64; rebuild the resident graph cache before reuse."))?;
        Ok(self.clock)
    }

    fn pop_lru_key(&mut self) -> Option<ResidentGraphCacheIdentity> {
        self.lru.pop_valid(
            |key| self.entries.get(key).map(|entry| entry.last_seen),
            || self.entries.keys().next().cloned(),
        )
    }

    fn record_lru(
        &mut self,
        key: ResidentGraphCacheIdentity,
        last_seen: u64,
    ) -> Result<(), vyre::BackendError> {
        self.lru
            .record(
                key,
                last_seen,
                self.entries
                    .iter()
                    .map(|(key, entry)| (key.clone(), entry.last_seen)),
                "resident graph cache",
            )
            .map_err(resident_graph_cache_error)
    }

    #[cfg(test)]
    pub(super) fn lru_len_for_tests(&self) -> usize {
        self.lru.len()
    }
}

fn resident_graph_cache_error(fix: impl Into<String>) -> vyre::BackendError {
    vyre::BackendError::InvalidProgram { fix: fix.into() }
}