weirflow 0.1.0

//! Resident IFDS batch facade.
//!
//! This module owns the orchestration between prepared host CSR layouts,
//! backend-resident CSR caching, and resident IFDS solves. The cache remains
//! responsible only for residency policy; this facade is the public path for
//! repeated IFDS queries that should not re-upload invariant graph buffers.

mod retained_scratch;
mod solve;

#[cfg(test)]
mod tests;

use crate::ifds_gpu::{ResidentIfdsParallelHostScratch, ResidentIfdsParallelScratch};
use crate::ifds_resident_cache::{ResidentIfdsCsrCache, ResidentIfdsCsrCacheStats};
use vyre::ResidentGraphReuseTelemetry;

/// Runtime counters for [`ResidentIfdsBatch`].
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
pub struct ResidentIfdsBatchStats {
    /// Resident CSR cache counters.
    pub csr_cache: ResidentIfdsCsrCacheStats,
    /// Parallel scratch allocations performed by the batch facade.
    pub scratch_allocations: u64,
    /// Solves that reused the retained parallel scratch allocation.
    pub scratch_reuses: u64,
    /// Parallel scratch allocations freed by resizing or [`ResidentIfdsBatch::free_all`].
    pub scratch_frees: u64,
    /// Successful resident IFDS solve dispatches executed through this facade.
    pub resident_dispatches: u64,
    /// Host-to-device seed payload bytes uploaded for resident IFDS solves.
    pub seed_upload_bytes: u64,
    /// Device-to-host decoded result bytes read back for resident IFDS solves.
    pub result_readback_bytes: u64,
    /// Resident CSR cache misses observed by this facade.
    pub resident_graph_cold_uploads: u64,
    /// Resident CSR cache hits observed by this facade.
    pub resident_graph_warm_reuses: u64,
    /// Resident CSR graph bytes uploaded by cold cache misses.
    pub resident_graph_upload_bytes: u64,
    /// Resident CSR graph upload bytes avoided by warm cache hits.
    pub resident_graph_avoided_upload_bytes: u64,
    /// Whether a reusable parallel scratch allocation is currently resident.
    pub scratch_resident: bool,
    /// Current resident parallel scratch bytes retained by this facade.
    pub scratch_bytes: usize,
}

impl ResidentIfdsBatchStats {
    /// Backend-neutral resident CSR graph reuse telemetry for this 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,
        )
    }
}

/// Reusable resident IFDS solver state.
pub struct ResidentIfdsBatch<R> {
    csr_cache: ResidentIfdsCsrCache<R>,
    parallel_scratch: Option<ResidentIfdsParallelScratch<R>>,
    parallel_host_scratch: ResidentIfdsParallelHostScratch,
    single_result_scratch: Vec<Vec<u32>>,
    max_parallel_scratch_bytes: Option<usize>,
    scratch_allocations: u64,
    scratch_reuses: u64,
    scratch_frees: u64,
    resident_dispatches: u64,
    seed_upload_bytes: u64,
    result_readback_bytes: u64,
    resident_graph_reuse: ResidentGraphReuseTelemetry,
    observed_csr_cache_reuse: ResidentGraphReuseTelemetry,
}

impl<R> Default for ResidentIfdsBatch<R> {
    fn default() -> Self {
        Self {
            csr_cache: ResidentIfdsCsrCache::default(),
            parallel_scratch: None,
            parallel_host_scratch: ResidentIfdsParallelHostScratch::default(),
            single_result_scratch: Vec::new(),
            max_parallel_scratch_bytes: None,
            scratch_allocations: 0,
            scratch_reuses: 0,
            scratch_frees: 0,
            resident_dispatches: 0,
            seed_upload_bytes: 0,
            result_readback_bytes: 0,
            resident_graph_reuse: ResidentGraphReuseTelemetry::default(),
            observed_csr_cache_reuse: ResidentGraphReuseTelemetry::default(),
        }
    }
}

impl<R> ResidentIfdsBatch<R>
where
    R: Clone,
{
    /// Create a resident IFDS batch facade with an unbounded CSR cache.
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }

    /// Create a resident IFDS batch facade with a soft retained-byte bound.
    #[must_use]
    pub fn with_max_retained_bytes(max_retained_bytes: usize) -> Self {
        Self {
            csr_cache: ResidentIfdsCsrCache::with_max_retained_bytes(max_retained_bytes),
            parallel_scratch: None,
            parallel_host_scratch: ResidentIfdsParallelHostScratch::default(),
            single_result_scratch: Vec::new(),
            max_parallel_scratch_bytes: None,
            scratch_allocations: 0,
            scratch_reuses: 0,
            scratch_frees: 0,
            resident_dispatches: 0,
            seed_upload_bytes: 0,
            result_readback_bytes: 0,
            resident_graph_reuse: ResidentGraphReuseTelemetry::default(),
            observed_csr_cache_reuse: ResidentGraphReuseTelemetry::default(),
        }
    }

    /// Create a resident IFDS batch facade with explicit retained-CSR and
    /// reusable parallel-scratch byte bounds.
    #[must_use]
    pub fn with_memory_budget(
        max_retained_bytes: usize,
        max_parallel_scratch_bytes: usize,
    ) -> Self {
        Self {
            csr_cache: ResidentIfdsCsrCache::with_max_retained_bytes(max_retained_bytes),
            parallel_scratch: None,
            parallel_host_scratch: ResidentIfdsParallelHostScratch::default(),
            single_result_scratch: Vec::new(),
            max_parallel_scratch_bytes: Some(max_parallel_scratch_bytes),
            scratch_allocations: 0,
            scratch_reuses: 0,
            scratch_frees: 0,
            resident_dispatches: 0,
            seed_upload_bytes: 0,
            result_readback_bytes: 0,
            resident_graph_reuse: ResidentGraphReuseTelemetry::default(),
            observed_csr_cache_reuse: ResidentGraphReuseTelemetry::default(),
        }
    }

    /// Inspect the resident CSR cache.
    #[must_use]
    pub fn csr_cache(&self) -> &ResidentIfdsCsrCache<R> {
        &self.csr_cache
    }

    /// Inspect resident CSR cache counters.
    #[must_use]
    pub fn csr_cache_stats(&self) -> ResidentIfdsCsrCacheStats {
        self.csr_cache.stats()
    }

    /// Inspect resident batch counters.
    #[must_use]
    pub fn stats(&self) -> ResidentIfdsBatchStats {
        ResidentIfdsBatchStats {
            csr_cache: self.csr_cache.stats(),
            scratch_allocations: self.scratch_allocations,
            scratch_reuses: self.scratch_reuses,
            scratch_frees: self.scratch_frees,
            resident_dispatches: self.resident_dispatches,
            seed_upload_bytes: self.seed_upload_bytes,
            result_readback_bytes: self.result_readback_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,
            scratch_resident: self.parallel_scratch.is_some(),
            scratch_bytes: self
                .parallel_scratch
                .as_ref()
                .map_or(0, ResidentIfdsParallelScratch::retained_bytes),
        }
    }

    fn record_graph_residency_from_cache_stats(&mut self) -> Result<(), String> {
        let csr_cache_reuse = self.csr_cache.stats().resident_graph_reuse_telemetry();
        let csr_cache_delta = csr_cache_reuse
            .checked_delta_since(self.observed_csr_cache_reuse)
            .map_err(|error| {
                format!("weir resident IFDS batch observed CSR cache telemetry regressed: {error}")
            })?;
        self.resident_graph_reuse = self
            .resident_graph_reuse
            .checked_add(csr_cache_delta)
            .map_err(|error| error.to_string())?;
        self.observed_csr_cache_reuse = csr_cache_reuse;
        Ok(())
    }

    fn record_resident_dispatch_io(
        &mut self,
        query_count: usize,
        total_seed_facts: usize,
        result_words: usize,
    ) -> Result<(), String> {
        self.resident_dispatches = self
            .resident_dispatches
            .checked_add(1)
            .ok_or_else(|| {
                "weir resident IFDS dispatch counter overflowed u64. Fix: rotate the batch facade or shard the workload before dispatch accounting saturates."
                    .to_string()
            })?;
        let seed_triple_count = total_seed_facts
            .checked_mul(3)
            .ok_or_else(|| {
                "weir resident IFDS seed triple count overflowed usize. Fix: shard the IFDS seed batch."
                    .to_string()
            })?
            .max(1);
        let seed_triple_bytes = seed_triple_count
            .checked_mul(std::mem::size_of::<u32>())
            .ok_or_else(|| {
                "weir resident IFDS seed triple byte count overflowed usize. Fix: shard the IFDS seed batch."
                    .to_string()
            })?;
        let seed_offset_count = query_count
            .checked_add(1)
            .ok_or_else(|| {
                "weir resident IFDS seed offset count overflowed usize. Fix: shard the IFDS query batch."
                    .to_string()
            })?;
        let seed_offset_bytes = seed_offset_count
            .checked_mul(std::mem::size_of::<u32>())
            .ok_or_else(|| {
                "weir resident IFDS seed offset byte count overflowed usize. Fix: shard the IFDS query batch."
                    .to_string()
            })?;
        let seed_upload_bytes = seed_triple_bytes
            .checked_add(seed_offset_bytes)
            .ok_or_else(|| {
                "weir resident IFDS seed upload byte count overflowed usize. Fix: shard the IFDS seed batch."
                    .to_string()
            })?;
        let seed_upload_bytes_u64 = u64::try_from(seed_upload_bytes).map_err(|error| {
            format!(
                "weir resident IFDS seed upload bytes do not fit u64: {error}. Fix: shard the IFDS seed batch."
            )
        })?;
        self.seed_upload_bytes = self
            .seed_upload_bytes
            .checked_add(seed_upload_bytes_u64)
            .ok_or_else(|| {
                "weir resident IFDS cumulative seed upload byte counter overflowed u64. Fix: rotate the batch facade or shard the workload."
                    .to_string()
            })?;
        let result_readback_bytes = result_words
            .checked_mul(std::mem::size_of::<u32>())
            .ok_or_else(|| {
                "weir resident IFDS result readback byte count overflowed usize. Fix: shard the IFDS result batch."
                    .to_string()
            })?;
        let result_readback_bytes_u64 = u64::try_from(result_readback_bytes).map_err(|error| {
            format!(
                "weir resident IFDS result readback bytes do not fit u64: {error}. Fix: shard the IFDS result batch."
            )
        })?;
        self.result_readback_bytes = self
            .result_readback_bytes
            .checked_add(result_readback_bytes_u64)
            .ok_or_else(|| {
                "weir resident IFDS cumulative result readback byte counter overflowed u64. Fix: rotate the batch facade or shard the workload."
                    .to_string()
            })?;
        Ok(())
    }
}