weir/

fixed_point_resident_batch.rs

//! Resident fixed-point batch facade.
//!
//! This module composes the resident graph cache, borrowed resident plans, and
//! reusable resident frontier scratch so hot fixed-point analyses do not
//! repeatedly upload CSR layouts or allocate frontier buffers.

#![allow(clippy::too_many_arguments)]
mod analysis_methods;
#[cfg(test)]
mod tests;

use crate::fixed_point_execution_plan_cache::{
    FixedPointExecutionPlanCache, FixedPointExecutionPlanCacheStats,
};
use crate::fixed_point_graph::FixedPointAnalysisKind;
use crate::fixed_point_graph::FixedPointForwardGraph;
use crate::fixed_point_resident_cache::{
    FixedPointResidentGraphCache, FixedPointResidentGraphCacheStats,
};
use crate::fixed_point_resident_frontier::FixedPointResidentFrontierScratch;
use vyre::ResidentGraphReuseTelemetry;

/// Runtime counters for [`FixedPointResidentBatch`].
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
pub struct FixedPointResidentBatchStats {
    /// Resident graph cache counters.
    pub graph_cache: FixedPointResidentGraphCacheStats,
    /// Execution-plan cache counters.
    pub execution_plan_cache: FixedPointExecutionPlanCacheStats,
    /// Resident frontier allocation or resize events observed by this facade.
    pub frontier_allocations: u64,
    /// Solves that reused the existing resident frontier allocation.
    pub frontier_reuses: u64,
    /// Resident frontier clears performed by [`FixedPointResidentBatch::free_all`].
    pub frontier_clears: u64,
    /// Successful resident fixed-point dispatches executed through this facade.
    pub resident_dispatches: u64,
    /// Host-to-device frontier seed bytes uploaded through resident solves.
    pub frontier_upload_bytes: u64,
    /// Device-to-host result bytes read back through resident solves.
    pub result_readback_bytes: u64,
    /// Resident graph cache misses observed by this facade.
    pub resident_graph_cold_uploads: u64,
    /// Resident graph cache hits observed by this facade.
    pub resident_graph_warm_reuses: u64,
    /// Resident graph bytes uploaded by cold graph-cache misses.
    pub resident_graph_upload_bytes: u64,
    /// Resident graph upload bytes avoided by warm graph-cache hits.
    pub resident_graph_avoided_upload_bytes: u64,
    /// Whether the reusable resident frontier pair is currently allocated.
    pub frontier_resident: bool,
    /// Current resident frontier byte width.
    pub frontier_bytes: usize,
}

impl FixedPointResidentBatchStats {
    /// Backend-neutral resident graph reuse telemetry for this batch snapshot.
    #[must_use]
    pub const fn resident_graph_reuse_telemetry(self) -> ResidentGraphReuseTelemetry {
        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 fixed-point execution state.
#[derive(Clone, Debug, PartialEq)]
pub struct FixedPointResidentBatch {
    graph_cache: FixedPointResidentGraphCache,
    execution_plan_cache: FixedPointExecutionPlanCache,
    resident_frontier: FixedPointResidentFrontierScratch,
    host_scratch: crate::fixed_point_scratch::FixedPointScratch,
    max_frontier_bytes: Option<usize>,
    frontier_allocations: u64,
    frontier_reuses: u64,
    frontier_clears: u64,
    resident_dispatches: u64,
    frontier_upload_bytes: u64,
    result_readback_bytes: u64,
    resident_graph_reuse: ResidentGraphReuseTelemetry,
    observed_graph_cache_reuse: ResidentGraphReuseTelemetry,
}

impl Default for FixedPointResidentBatch {
    fn default() -> Self {
        Self {
            graph_cache: FixedPointResidentGraphCache::new(),
            execution_plan_cache: FixedPointExecutionPlanCache::new(),
            resident_frontier: FixedPointResidentFrontierScratch::default(),
            host_scratch: crate::fixed_point_scratch::FixedPointScratch::default(),
            max_frontier_bytes: None,
            frontier_allocations: 0,
            frontier_reuses: 0,
            frontier_clears: 0,
            resident_dispatches: 0,
            frontier_upload_bytes: 0,
            result_readback_bytes: 0,
            resident_graph_reuse: ResidentGraphReuseTelemetry::default(),
            observed_graph_cache_reuse: ResidentGraphReuseTelemetry::default(),
        }
    }
}

impl FixedPointResidentBatch {
    /// Create a resident fixed-point batch facade with an unbounded graph cache.
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }

    /// Create a resident fixed-point batch facade with a soft graph-cache byte bound.
    #[must_use]
    pub fn with_max_retained_bytes(max_retained_bytes: usize) -> Self {
        Self {
            graph_cache: FixedPointResidentGraphCache::with_max_retained_bytes(max_retained_bytes),
            ..Self::default()
        }
    }

    /// Create a resident fixed-point batch facade with explicit graph-cache
    /// byte and execution-plan entry bounds.
    #[must_use]
    pub fn with_cache_bounds(max_retained_bytes: usize, max_execution_plans: usize) -> Self {
        Self {
            graph_cache: FixedPointResidentGraphCache::with_max_retained_bytes(max_retained_bytes),
            execution_plan_cache: FixedPointExecutionPlanCache::with_max_entries(
                max_execution_plans,
            ),
            ..Self::default()
        }
    }

    /// Create a resident fixed-point batch facade with explicit graph-cache,
    /// execution-plan cache, and resident frontier byte bounds.
    #[must_use]
    pub fn with_memory_budget(
        max_retained_bytes: usize,
        max_execution_plans: usize,
        max_frontier_bytes: usize,
    ) -> Self {
        Self {
            graph_cache: FixedPointResidentGraphCache::with_max_retained_bytes(max_retained_bytes),
            execution_plan_cache: FixedPointExecutionPlanCache::with_max_entries(
                max_execution_plans,
            ),
            max_frontier_bytes: Some(max_frontier_bytes),
            ..Self::default()
        }
    }

    /// Inspect the resident graph cache.
    #[must_use]
    pub fn graph_cache(&self) -> &FixedPointResidentGraphCache {
        &self.graph_cache
    }

    /// Inspect resident graph cache counters.
    #[must_use]
    pub fn graph_cache_stats(&self) -> FixedPointResidentGraphCacheStats {
        self.graph_cache.stats()
    }

    /// Inspect execution-plan cache counters.
    #[must_use]
    pub fn execution_plan_cache_stats(&self) -> FixedPointExecutionPlanCacheStats {
        self.execution_plan_cache.stats()
    }

    /// Inspect resident batch counters.
    #[must_use]
    pub fn stats(&self) -> FixedPointResidentBatchStats {
        FixedPointResidentBatchStats {
            graph_cache: self.graph_cache.stats(),
            execution_plan_cache: self.execution_plan_cache.stats(),
            frontier_allocations: self.frontier_allocations,
            frontier_reuses: self.frontier_reuses,
            frontier_clears: self.frontier_clears,
            resident_dispatches: self.resident_dispatches,
            frontier_upload_bytes: self.frontier_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,
            frontier_resident: self.resident_frontier.is_allocated(),
            frontier_bytes: self.resident_frontier.byte_len(),
        }
    }

    /// Latest host-observed frontier density from the last fixed-point solve.
    #[must_use]
    pub fn frontier_density(&self) -> crate::fixed_point_scratch::FrontierDensityTelemetry {
        self.host_scratch.frontier_density()
    }

    /// Sparse/dense execution family recommended from the last solve's frontier density.
    #[must_use]
    pub fn recommended_frontier_execution_mode(
        &self,
    ) -> crate::fixed_point_scratch::FrontierExecutionMode {
        self.host_scratch
            .frontier_density()
            .recommended_execution_mode()
    }

    /// Build a scale-aware execution plan for a prepared graph using the last
    /// observed frontier density.
    pub fn execution_plan_for_prepared_graph(
        &self,
        graph: &FixedPointForwardGraph,
    ) -> Result<crate::fixed_point_execution_plan::FixedPointExecutionPlan, String> {
        crate::fixed_point_execution_plan::plan_prepared_graph(graph, self.frontier_density())
    }

    /// Return a cached scale-aware execution plan for a prepared graph and
    /// analysis family using the last observed frontier density.
    pub fn cached_execution_plan_for_prepared_graph(
        &mut self,
        backend: &dyn vyre::VyreBackend,
        kind: FixedPointAnalysisKind,
        graph: &FixedPointForwardGraph,
    ) -> Result<crate::fixed_point_execution_plan::FixedPointExecutionPlan, String> {
        self.execution_plan_cache
            .get_or_plan(backend, kind, graph, self.frontier_density())
    }
    /// Free retained resident frontier and graph resources.
    pub fn free_all(&mut self, backend: &dyn vyre::VyreBackend) -> Result<(), vyre::BackendError> {
        let frontier_result = if self.resident_frontier.is_allocated() {
            let result = self.resident_frontier.clear(backend);
            if result.is_ok() {
                self.frontier_clears = self
                    .frontier_clears
                    .checked_add(1)
                    .ok_or_else(|| {
                        vyre::BackendError::new(
                            "weir fixed-point resident batch frontier clear counter overflowed u64. Fix: rotate the batch facade before clearing more resident frontiers.",
                        )
                    })?;
            }
            result
        } else {
            Ok(())
        };
        let cache_result = self.graph_cache.free_all(backend);
        match (frontier_result, cache_result) {
            (Ok(()), Ok(())) => Ok(()),
            (Err(error), _) | (_, Err(error)) => Err(error),
        }
    }

    fn frontier_state(&self) -> (bool, usize) {
        (
            self.resident_frontier.is_allocated(),
            self.resident_frontier.byte_len(),
        )
    }

    fn record_frontier_reuse(&mut self, before: (bool, usize)) -> Result<(), String> {
        let after = self.frontier_state();
        if !after.0 {
            return Ok(());
        }
        if before.0 && before.1 == after.1 {
            self.frontier_reuses = self
                .frontier_reuses
                .checked_add(1)
                .ok_or_else(|| {
                    "weir fixed-point resident batch frontier reuse counter overflowed u64. Fix: rotate the batch facade before reusing more resident frontiers."
                        .to_string()
                })?;
        } else {
            self.frontier_allocations = self
                .frontier_allocations
                .checked_add(1)
                .ok_or_else(|| {
                    "weir fixed-point resident batch frontier allocation counter overflowed u64. Fix: rotate the batch facade before allocating more resident frontiers."
                        .to_string()
                })?;
        }
        Ok(())
    }

    fn record_execution_plan(
        &mut self,
        backend: &dyn vyre::VyreBackend,
        kind: FixedPointAnalysisKind,
        graph: &FixedPointForwardGraph,
    ) -> Result<(), String> {
        let _ =
            self.execution_plan_cache
                .get_or_plan(backend, kind, graph, self.frontier_density())?;
        self.record_graph_residency_from_cache_stats()?;
        Ok(())
    }

    fn record_graph_residency_from_cache_stats(&mut self) -> Result<(), String> {
        let graph_cache_reuse = self.graph_cache.stats().resident_graph_reuse_telemetry();
        let graph_cache_delta = graph_cache_reuse
            .checked_delta_since(self.observed_graph_cache_reuse)
            .map_err(|error| {
                format!(
                    "weir fixed-point resident batch observed graph cache telemetry regressed: {error}"
                )
            })?;
        self.resident_graph_reuse = self
            .resident_graph_reuse
            .checked_add(graph_cache_delta)
            .map_err(|error| error.to_string())?;
        self.observed_graph_cache_reuse = graph_cache_reuse;
        Ok(())
    }

    fn record_resident_dispatch_io(
        &mut self,
        seed_words: usize,
        result_words: usize,
    ) -> Result<(), String> {
        self.resident_dispatches = self
            .resident_dispatches
            .checked_add(1)
            .ok_or_else(|| {
                "weir fixed-point resident batch dispatch counter overflowed u64. Fix: rotate the batch facade before dispatch accounting saturates."
                    .to_string()
            })?;
        let logical_seed_bytes = seed_words
            .checked_mul(std::mem::size_of::<u32>())
            .ok_or_else(|| {
                "weir fixed-point resident seed byte count overflowed usize. Fix: shard the seed frontier before resident dispatch."
                    .to_string()
            })?;
        let seed_bytes = logical_seed_bytes.checked_mul(2).ok_or_else(|| {
            "weir fixed-point resident physical seed byte count overflowed usize. Fix: shard the seed frontier before resident dispatch."
                .to_string()
        })?;
        let seed_bytes_u64 = u64::try_from(seed_bytes).map_err(|error| {
            format!(
                "weir fixed-point resident seed byte count does not fit u64: {error}. Fix: shard the seed frontier before resident dispatch."
            )
        })?;
        self.frontier_upload_bytes = self
            .frontier_upload_bytes
            .checked_add(seed_bytes_u64)
            .ok_or_else(|| {
                "weir fixed-point resident frontier upload byte counter overflowed u64. Fix: rotate the batch facade or shard the workload."
                    .to_string()
            })?;
        let result_bytes = result_words
            .checked_mul(std::mem::size_of::<u32>())
            .ok_or_else(|| {
                "weir fixed-point resident result byte count overflowed usize. Fix: shard the result frontier before resident readback."
                    .to_string()
            })?;
        let result_bytes_u64 = u64::try_from(result_bytes).map_err(|error| {
            format!(
                "weir fixed-point resident result byte count does not fit u64: {error}. Fix: shard the result frontier before resident readback."
            )
        })?;
        self.result_readback_bytes = self
            .result_readback_bytes
            .checked_add(result_bytes_u64)
            .ok_or_else(|| {
                "weir fixed-point resident result readback byte counter overflowed u64. Fix: rotate the batch facade or shard the workload."
                    .to_string()
            })?;
        Ok(())
    }

    fn record_resident_changed_flag_sequence_window_io(
        &mut self,
        seed_words: usize,
        result_words: usize,
    ) -> Result<(), String> {
        let readback_words = result_words.checked_add(1).ok_or_else(|| {
            "weir fixed-point resident changed-flag sequence-window readback word count overflowed usize. Fix: shard the result frontier before resident readback."
                .to_string()
        })?;
        self.record_resident_dispatch_io(seed_words, readback_words)
    }

    fn require_frontier_budget(&self, graph: &FixedPointForwardGraph) -> Result<(), String> {
        let Some(max_frontier_bytes) = self.max_frontier_bytes else {
            return Ok(());
        };
        let node_count = usize::try_from(graph.node_count()).map_err(|error| {
            format!(
                "weir resident fixed-point frontier budget cannot represent node_count={} as usize: {error}. Fix: shard the graph before resident dispatch.",
                graph.node_count()
            )
        })?;
        let frontier_bytes = node_count
            .checked_add(31)
            .and_then(|bits| bits.checked_div(32))
            .and_then(|words| words.checked_mul(std::mem::size_of::<u32>()))
            .ok_or_else(|| {
                format!(
                    "weir resident fixed-point frontier budget overflowed for node_count={}. Fix: shard the graph before resident dispatch.",
                    graph.node_count()
                )
            })?;
        if frontier_bytes > max_frontier_bytes {
            return Err(format!(
                "weir resident fixed-point frontier requires {frontier_bytes} bytes but the configured device-memory budget allows {max_frontier_bytes}. Fix: increase max_frontier_bytes, reuse a smaller graph shard, or lower the analysis batch size before dispatch."
            ));
        }
        Ok(())
    }
}