weirflow 0.1.0

//! Direct-resident IFDS batch facade.
//!
//! This facade is for repeated IFDS solves where the invariant exploded CSR
//! should be built directly into resident backend buffers and retained across
//! seed queries. It composes direct resident preparation with the existing
//! resident scratch-reuse solver; callers do not need to materialize host CSR
//! byte buffers first.

mod cache;
mod graph_key;
mod sizing;
mod solve;
#[cfg(test)]
mod tests;

use crate::ifds_gpu::ResidentPreparedIfdsCsr;
use crate::ifds_resident_batch::{ResidentIfdsBatch, ResidentIfdsBatchStats};
use crate::ifds_resident_direct_prepare_scratch::DirectResidentIfdsPrepareScratch;
use crate::resident_cache_lru::ResidentCacheLru;
use std::collections::HashMap;

pub use graph_key::{DirectResidentIfdsGraph, DirectResidentIfdsGraphKey};
pub(crate) use sizing::graph_hash;

/// Runtime counters for [`DirectResidentIfdsBatch`].
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
pub struct DirectResidentIfdsBatchStats {
    /// Direct resident CSR cache hits.
    pub hits: u64,
    /// Direct resident CSR cache misses.
    pub misses: u64,
    /// Direct resident CSR prepare operations.
    pub direct_prepares: u64,
    /// Cumulative direct resident CSR bytes prepared into GPU-resident buffers.
    pub direct_prepare_bytes: u64,
    /// Direct resident graph cache misses that prepared cold graph state.
    pub resident_graph_cold_uploads: u64,
    /// Direct resident graph cache hits that reused warm graph state.
    pub resident_graph_warm_reuses: u64,
    /// Direct resident graph bytes prepared by cold misses.
    pub resident_graph_upload_bytes: u64,
    /// Direct resident graph upload bytes avoided by warm cache hits.
    pub resident_graph_avoided_upload_bytes: u64,
    /// Current retained direct-resident CSR entries.
    pub entries: usize,
    /// Current resident CSR bytes retained by the direct cache.
    pub retained_bytes: usize,
    /// Direct resident CSR entries evicted to respect the retention bound.
    pub evictions: u64,
    /// Resident IFDS solve batches executed after direct graph cache lookup or
    /// prepare.
    pub resident_solve_batches: u64,
    /// Seed-set rows submitted through resident IFDS solve batches.
    pub resident_solve_queries: u64,
    /// Non-empty seed-set rows submitted through resident IFDS solve batches.
    pub resident_solve_non_empty_queries: u64,
    /// Successful resident IFDS solve dispatches executed by the underlying
    /// solver.
    pub resident_dispatches: u64,
    /// Host-to-device seed payload bytes uploaded by direct resident solves.
    pub seed_upload_bytes: u64,
    /// Device-to-host result bytes read back by direct resident solves.
    pub result_readback_bytes: u64,
    /// Underlying resident scratch/cache solver counters.
    pub resident_solver: ResidentIfdsBatchStats,
}

/// Batch facade for direct-resident IFDS CSR preparation and solving.
pub struct DirectResidentIfdsBatch<R> {
    entries: HashMap<DirectIfdsKey, DirectIfdsEntry<R>>,
    lru: ResidentCacheLru<DirectIfdsKey>,
    max_entries: Option<usize>,
    max_retained_bytes: Option<usize>,
    retained_bytes: usize,
    clock: u64,
    solver: ResidentIfdsBatch<R>,
    prepare_scratch: DirectResidentIfdsPrepareScratch,
    prepare_resource_scratch: Vec<R>,
    single_result_scratch: Vec<Vec<u32>>,
    hits: u64,
    misses: u64,
    direct_prepares: u64,
    direct_prepare_bytes: u64,
    resident_graph_reuse: vyre::ResidentGraphReuseTelemetry,
    evictions: u64,
    resident_solve_batches: u64,
    resident_solve_queries: u64,
    resident_solve_non_empty_queries: u64,
}

#[derive(Clone, Debug, Eq, Hash, PartialEq)]
pub(crate) struct DirectIfdsKey {
    backend_id: &'static str,
    backend_version: &'static str,
    num_procs: u32,
    blocks_per_proc: u32,
    facts_per_proc: u32,
    graph_hash: u64,
}

struct DirectIfdsEntry<R> {
    retained_bytes: usize,
    last_seen: u64,
    graph: ResidentPreparedIfdsCsr<R>,
}

impl<R> Default for DirectResidentIfdsBatch<R>
where
    R: Clone,
{
    fn default() -> Self {
        Self {
            entries: HashMap::new(),
            lru: ResidentCacheLru::default(),
            max_entries: None,
            max_retained_bytes: None,
            retained_bytes: 0,
            clock: 0,
            solver: ResidentIfdsBatch::new(),
            prepare_scratch: DirectResidentIfdsPrepareScratch::default(),
            prepare_resource_scratch: Vec::new(),
            single_result_scratch: Vec::new(),
            hits: 0,
            misses: 0,
            direct_prepares: 0,
            direct_prepare_bytes: 0,
            resident_graph_reuse: vyre::ResidentGraphReuseTelemetry::default(),
            evictions: 0,
            resident_solve_batches: 0,
            resident_solve_queries: 0,
            resident_solve_non_empty_queries: 0,
        }
    }
}

impl<R> DirectResidentIfdsBatch<R>
where
    R: Clone,
{
    /// Create an empty direct-resident IFDS batch facade.
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }

    /// Create a direct-resident IFDS batch facade that retains at most
    /// `max_entries` resident CSR graphs.
    #[must_use]
    pub fn with_max_entries(max_entries: usize) -> Self {
        Self {
            max_entries: Some(max_entries.max(1)),
            ..Self::default()
        }
    }

    /// Create a direct-resident IFDS batch facade with explicit cache-entry
    /// and retained-CSR byte bounds.
    #[must_use]
    pub fn with_memory_budget(max_entries: usize, max_retained_bytes: usize) -> Self {
        Self {
            max_entries: Some(max_entries.max(1)),
            max_retained_bytes: Some(max_retained_bytes),
            ..Self::default()
        }
    }

    /// Inspect runtime counters.
    #[must_use]
    pub fn stats(&self) -> DirectResidentIfdsBatchStats {
        let resident_solver = self.solver.stats();
        DirectResidentIfdsBatchStats {
            hits: self.hits,
            misses: self.misses,
            direct_prepares: self.direct_prepares,
            direct_prepare_bytes: self.direct_prepare_bytes,
            resident_graph_cold_uploads: self.resident_graph_reuse.cold_uploads,
            resident_graph_warm_reuses: self.resident_graph_reuse.warm_reuses,
            resident_graph_upload_bytes: self.resident_graph_reuse.upload_bytes,
            resident_graph_avoided_upload_bytes: self.resident_graph_reuse.avoided_upload_bytes,
            entries: self.entries.len(),
            retained_bytes: self.retained_bytes,
            evictions: self.evictions,
            resident_solve_batches: self.resident_solve_batches,
            resident_solve_queries: self.resident_solve_queries,
            resident_solve_non_empty_queries: self.resident_solve_non_empty_queries,
            resident_dispatches: resident_solver.resident_dispatches,
            seed_upload_bytes: resident_solver.seed_upload_bytes,
            result_readback_bytes: resident_solver.result_readback_bytes,
            resident_solver,
        }
    }

    /// Backend-neutral resident graph reuse telemetry for this direct batch.
    #[must_use]
    pub fn resident_graph_reuse_telemetry(&self) -> vyre::ResidentGraphReuseTelemetry {
        self.stats().resident_graph_reuse_telemetry()
    }

    fn record_cold_graph_prepare(&mut self, retained_bytes: usize) -> Result<(), String> {
        let retained_bytes = u64::try_from(retained_bytes).map_err(|error| {
            format!(
                "weir direct resident IFDS retained graph byte count does not fit u64: {error}. Fix: shard the IFDS graph before direct resident telemetry accounting."
            )
        })?;
        self.resident_graph_reuse
            .record_cold_upload(retained_bytes)
            .map_err(|error| error.to_string())?;
        Ok(())
    }

    fn record_warm_graph_reuse(&mut self, retained_bytes: usize) -> Result<(), String> {
        let retained_bytes = u64::try_from(retained_bytes).map_err(|error| {
            format!(
                "weir direct resident IFDS retained graph byte count does not fit u64: {error}. Fix: shard the IFDS graph before direct resident telemetry accounting."
            )
        })?;
        self.resident_graph_reuse
            .record_warm_reuse(retained_bytes)
            .map_err(|error| error.to_string())?;
        Ok(())
    }

    fn record_resident_solve_batch(&mut self, seed_sets: &[&[(u32, u32, u32)]]) -> Result<(), String> {
        self.resident_solve_batches = self
            .resident_solve_batches
            .checked_add(1)
            .ok_or_else(|| {
                "weir direct resident IFDS solve-batch counter overflowed u64. Fix: rebuild the direct resident batch facade before reuse."
                    .to_string()
            })?;
        self.resident_solve_queries = self
            .resident_solve_queries
            .checked_add(u64::try_from(seed_sets.len()).map_err(|error| {
                format!(
                    "weir direct resident IFDS query count does not fit u64: {error}. Fix: shard the IFDS query batch before resident solve."
                )
            })?)
            .ok_or_else(|| {
                "weir direct resident IFDS query counter overflowed u64. Fix: rebuild the direct resident batch facade before reuse."
                    .to_string()
            })?;
        let non_empty = seed_sets
            .iter()
            .filter(|seed_facts| !seed_facts.is_empty())
            .count();
        self.resident_solve_non_empty_queries = self
            .resident_solve_non_empty_queries
            .checked_add(u64::try_from(non_empty).map_err(|error| {
                format!(
                    "weir direct resident IFDS non-empty query count does not fit u64: {error}. Fix: shard the IFDS query batch before resident solve."
                )
            })?)
            .ok_or_else(|| {
                "weir direct resident IFDS non-empty query counter overflowed u64. Fix: rebuild the direct resident batch facade before reuse."
                    .to_string()
            })?;
        Ok(())
    }
}

impl DirectResidentIfdsBatchStats {
    /// Backend-neutral resident graph reuse telemetry for this direct batch snapshot.
    #[must_use]
    pub const fn resident_graph_reuse_telemetry(self) -> vyre::ResidentGraphReuseTelemetry {
        vyre::ResidentGraphReuseTelemetry::from_counters(
            self.resident_graph_cold_uploads,
            self.resident_graph_warm_reuses,
            self.resident_graph_upload_bytes,
            self.resident_graph_avoided_upload_bytes,
        )
    }
}