rustytorch/crates/production/rtx-inference/src/inference_graph.rs

//! Inference CUDA Graph Capture State Machine
//!
//! Tracks warmup, capture, and replay phases for decode-step CUDA graph
//! optimization. The state machine is pure logic — no GPU resources are held
//! here. GPU resources (stream handles, graph IDs) are managed by
//! [`CudaGraphManager`] and wired in by the caller.
//!
//! # Phase Transitions
//!
//! ```text
//! Warmup (step_count < warmup_steps)
//!   ──► Capture (step_count == warmup_steps && !capture_attempted)
//!   ──► Replay (capture_attempted == true)
//! ```
//!
//! When `enabled` is `false` the machine stays in `Warmup` indefinitely,
//! which means normal (non-graph) execution is used for every step.

/// The current execution mode as determined by the graph capture state machine.
#[derive(Debug, Clone, PartialEq)]
pub enum StepMode {
    /// Normal execution — kernels are launched individually. Used during the
    /// initial warmup phase to prime caches and stabilize shapes.
    Warmup,
    /// The next step should be captured into a CUDA Graph. Called exactly once
    /// (at `step_count == warmup_steps`), then transitions to `Replay`.
    Capture,
    /// The captured graph is replayed instead of re-launching individual kernels.
    Replay,
}

/// State machine for inference-time CUDA Graph capture and replay.
///
/// # Example
///
/// ```rust
/// use rtx_inference::inference_graph::{InferenceGraphCapture, StepMode};
///
/// let mut capture = InferenceGraphCapture::new(3).enabled(true);
/// assert_eq!(capture.step_mode(), StepMode::Warmup);
///
/// // Advance through warmup steps
/// capture.advance(); // step 1
/// capture.advance(); // step 2
/// capture.advance(); // step 3 — still warmup (count is now 3, threshold is 3)
/// assert_eq!(capture.step_mode(), StepMode::Capture);
///
/// // Record a successful capture
/// capture.record_capture(42, 8, 1);
/// assert!(capture.is_captured());
/// assert_eq!(capture.step_mode(), StepMode::Replay);
/// ```
#[derive(Debug)]
pub struct InferenceGraphCapture {
    /// Captured CUDA Graph ID, or `None` if not yet captured.
    graph_id: Option<u64>,
    /// Whether a capture has already been attempted (prevents re-capture).
    capture_attempted: bool,
    /// Number of warmup steps before capture is attempted.
    warmup_steps: usize,
    /// Total number of steps advanced so far.
    step_count: usize,
    /// Whether CUDA Graph capture is enabled at all.
    enabled: bool,
    /// Batch size used during capture, for shape validation.
    captured_batch_size: Option<usize>,
    /// Sequence step length used during capture, for shape validation.
    captured_seq_step: Option<usize>,
}

impl InferenceGraphCapture {
    /// Create a new state machine with the given warmup step count.
    ///
    /// Graph capture is **disabled** by default. Call `.enabled(true)` to
    /// activate it.
    #[must_use]
    pub fn new(warmup_steps: usize) -> Self {
        Self {
            graph_id: None,
            capture_attempted: false,
            warmup_steps,
            step_count: 0,
            enabled: false,
            captured_batch_size: None,
            captured_seq_step: None,
        }
    }

    /// Builder method to set whether graph capture is enabled.
    ///
    /// When `enabled` is `false`, [`step_mode`](Self::step_mode) always
    /// returns [`StepMode::Warmup`].
    #[must_use]
    pub fn enabled(mut self, enabled: bool) -> Self {
        self.enabled = enabled;
        self
    }

    /// Returns `true` if a graph has been successfully captured.
    #[must_use]
    pub fn is_captured(&self) -> bool {
        self.graph_id.is_some()
    }

    /// Returns the total number of steps that have been advanced.
    #[must_use]
    pub fn step_count(&self) -> usize {
        self.step_count
    }

    /// Determine the execution mode for the current step.
    ///
    /// - [`StepMode::Warmup`] — if capture is disabled, or `step_count < warmup_steps`.
    /// - [`StepMode::Capture`] — if `step_count == warmup_steps` and no capture has
    ///   been attempted yet.
    /// - [`StepMode::Replay`] — if a capture has been attempted (successfully or not).
    #[must_use]
    pub fn step_mode(&self) -> StepMode {
        if !self.enabled {
            return StepMode::Warmup;
        }
        if self.step_count < self.warmup_steps {
            StepMode::Warmup
        } else if self.step_count == self.warmup_steps && !self.capture_attempted {
            StepMode::Capture
        } else {
            StepMode::Replay
        }
    }

    /// Advance the internal step counter by one and mark `capture_attempted`
    /// once we are past the warmup threshold.
    pub fn advance(&mut self) {
        self.step_count += 1;
        if self.step_count > self.warmup_steps {
            self.capture_attempted = true;
        }
    }

    /// Record a successful graph capture.
    ///
    /// # Arguments
    ///
    /// * `graph_id` — the ID returned by `CudaGraphManager::end_capture`.
    /// * `batch_size` — batch dimension used during capture.
    /// * `seq_step` — sequence-step dimension used during capture.
    pub fn record_capture(&mut self, graph_id: u64, batch_size: usize, seq_step: usize) {
        self.graph_id = Some(graph_id);
        self.captured_batch_size = Some(batch_size);
        self.captured_seq_step = Some(seq_step);
        self.capture_attempted = true;
    }

    /// Check whether the provided shapes match the shapes used at capture time.
    ///
    /// Returns `true` if the shapes match (or if no capture has occurred yet),
    /// `false` if they differ and the graph cannot be replayed safely.
    #[must_use]
    pub fn check_static_shape(&self, batch_size: usize, seq_step: usize) -> bool {
        match (self.captured_batch_size, self.captured_seq_step) {
            (Some(b), Some(s)) => b == batch_size && s == seq_step,
            _ => true,
        }
    }

    /// Return the captured graph ID, or `None` if no capture has occurred.
    #[must_use]
    pub fn captured_graph_id(&self) -> Option<u64> {
        self.graph_id
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    // -----------------------------------------------------------------------
    // step_mode tests
    // -----------------------------------------------------------------------

    #[test]
    fn test_step_mode_warmup_until_threshold() {
        let mut cap = InferenceGraphCapture::new(3).enabled(true);
        // Steps 0, 1, 2 are all Warmup
        for _ in 0..3 {
            assert_eq!(cap.step_mode(), StepMode::Warmup);
            cap.advance();
        }
        // step_count is now 3 — at threshold; should be Capture, not Warmup
        assert_ne!(cap.step_mode(), StepMode::Warmup);
    }

    #[test]
    fn test_step_mode_capture_at_threshold() {
        let mut cap = InferenceGraphCapture::new(3).enabled(true);
        // Advance through exactly warmup_steps steps
        for _ in 0..3 {
            cap.advance();
        }
        assert_eq!(cap.step_mode(), StepMode::Capture);
    }

    #[test]
    fn test_step_mode_replay_after_capture() {
        let mut cap = InferenceGraphCapture::new(2).enabled(true);
        cap.advance();
        cap.advance();
        assert_eq!(cap.step_mode(), StepMode::Capture);

        cap.record_capture(1, 4, 1);
        assert_eq!(cap.step_mode(), StepMode::Replay);
    }

    // -----------------------------------------------------------------------
    // advance tests
    // -----------------------------------------------------------------------

    #[test]
    fn test_advance_increments_step_count() {
        let mut cap = InferenceGraphCapture::new(5).enabled(true);
        assert_eq!(cap.step_count(), 0);
        cap.advance();
        assert_eq!(cap.step_count(), 1);
        cap.advance();
        assert_eq!(cap.step_count(), 2);
    }

    // -----------------------------------------------------------------------
    // check_static_shape tests
    // -----------------------------------------------------------------------

    #[test]
    fn test_static_shape_check_passes_same_shape() {
        let mut cap = InferenceGraphCapture::new(1).enabled(true);
        cap.advance();
        cap.record_capture(7, 8, 1);
        assert!(cap.check_static_shape(8, 1));
    }

    #[test]
    fn test_static_shape_check_fails_different_batch() {
        let mut cap = InferenceGraphCapture::new(1).enabled(true);
        cap.advance();
        cap.record_capture(7, 8, 1);
        // Different batch size
        assert!(!cap.check_static_shape(4, 1));
    }

    // -----------------------------------------------------------------------
    // is_captured tests
    // -----------------------------------------------------------------------

    #[test]
    fn test_is_captured_false_before_record() {
        let cap = InferenceGraphCapture::new(3).enabled(true);
        assert!(!cap.is_captured());
    }

    // -----------------------------------------------------------------------
    // disabled tests
    // -----------------------------------------------------------------------

    #[test]
    fn test_inference_graph_capture_disabled() {
        let mut cap = InferenceGraphCapture::new(1).enabled(false);
        // Even after advancing past the warmup threshold the mode stays Warmup
        cap.advance();
        cap.advance();
        assert_eq!(cap.step_mode(), StepMode::Warmup);
        assert!(!cap.is_captured());
    }
}