rustytorch/crates/core/rtx-fusion/src/stream.rs

//! Operation Stream for queuing and managing tensor operations
//!
//! The `OperationStream` is the core data structure for kernel fusion. It queues
//! tensor operations as they are created, tracks dependencies between them, and
//! provides the operation sequence to the fusion analyzer.

use crate::config::FusionConfig;
use crate::kernel::{DType, StreamOpKind, StreamOperation, TensorId};
use parking_lot::RwLock;
use std::collections::{HashMap, HashSet, VecDeque};
use std::sync::Arc;

/// Dependency information for a tensor
#[derive(Debug, Clone)]
pub struct TensorDependency {
    /// The operation that produces this tensor
    pub producer: Option<usize>,
    /// Operations that consume this tensor
    pub consumers: Vec<usize>,
    /// Reference count (number of pending uses)
    pub ref_count: usize,
    /// Whether this tensor has been materialized
    pub materialized: bool,
}

impl TensorDependency {
    fn new() -> Self {
        Self {
            producer: None,
            consumers: Vec::new(),
            ref_count: 0,
            materialized: false,
        }
    }

    fn as_input() -> Self {
        Self {
            producer: None,
            consumers: Vec::new(),
            ref_count: 1,
            materialized: true, // Inputs are already materialized
        }
    }
}

/// A queue of pending operations for a single device
#[derive(Debug)]
pub struct OperationStream {
    /// Pending operations in submission order
    operations: Vec<StreamOperation>,

    /// Dependency tracking for tensors
    dependencies: HashMap<TensorId, TensorDependency>,

    /// Set of external input tensor IDs (not produced by stream operations)
    external_inputs: HashSet<TensorId>,

    /// Configuration
    config: FusionConfig,

    /// Whether the stream needs to be flushed
    needs_flush: bool,
}

impl OperationStream {
    /// Create a new empty operation stream
    pub fn new(config: FusionConfig) -> Self {
        Self {
            operations: Vec::new(),
            dependencies: HashMap::new(),
            external_inputs: HashSet::new(),
            config,
            needs_flush: false,
        }
    }

    /// Create a stream with default configuration
    pub fn with_defaults() -> Self {
        Self::new(FusionConfig::default())
    }

    /// Record a new operation in the stream
    pub fn record(&mut self, op: StreamOperation) {
        let op_idx = self.operations.len();

        // Track dependencies for inputs
        for &input_id in &op.inputs {
            let dep = self
                .dependencies
                .entry(input_id)
                .or_insert_with(TensorDependency::as_input);

            dep.consumers.push(op_idx);
            dep.ref_count += 1;

            // If this input wasn't produced by the stream, it's external
            if dep.producer.is_none() {
                self.external_inputs.insert(input_id);
            }
        }

        // Track the output tensor
        let output_id = op.output;
        let mut output_dep = TensorDependency::new();
        output_dep.producer = Some(op_idx);
        self.dependencies.insert(output_id, output_dep);

        // Check if this is a sync point
        if op.op.is_sync_point() {
            self.needs_flush = true;
        }

        self.operations.push(op);

        // Auto-flush if we've reached the limit
        if self.operations.len() >= self.config.max_pending_ops {
            self.needs_flush = true;
        }
    }

    /// Register an external input tensor (already materialized)
    pub fn register_input(&mut self, tensor_id: TensorId) {
        self.external_inputs.insert(tensor_id);
        self.dependencies
            .entry(tensor_id)
            .or_insert_with(TensorDependency::as_input);
    }

    /// Mark a tensor as materialized (its data is available)
    pub fn mark_materialized(&mut self, tensor_id: TensorId) {
        if let Some(dep) = self.dependencies.get_mut(&tensor_id) {
            dep.materialized = true;
        }
    }

    /// Decrement the reference count for a tensor (called when tensor is consumed)
    pub fn decrement_ref(&mut self, tensor_id: TensorId) {
        if let Some(dep) = self.dependencies.get_mut(&tensor_id) {
            dep.ref_count = dep.ref_count.saturating_sub(1);
        }
    }

    /// Check if a tensor is the last use (ref_count == 1)
    pub fn is_last_use(&self, tensor_id: TensorId) -> bool {
        self.dependencies
            .get(&tensor_id)
            .map_or(false, |dep| dep.ref_count == 1)
    }

    /// Check if the stream needs to be flushed
    pub fn needs_flush(&self) -> bool {
        self.needs_flush
    }

    /// Get the number of pending operations
    pub fn len(&self) -> usize {
        self.operations.len()
    }

    /// Check if the stream is empty
    pub fn is_empty(&self) -> bool {
        self.operations.is_empty()
    }

    /// Get all pending operations
    pub fn operations(&self) -> &[StreamOperation] {
        &self.operations
    }

    /// Get external input tensor IDs
    pub fn external_inputs(&self) -> &HashSet<TensorId> {
        &self.external_inputs
    }

    /// Get the dependency information for a tensor
    pub fn get_dependency(&self, tensor_id: TensorId) -> Option<&TensorDependency> {
        self.dependencies.get(&tensor_id)
    }

    /// Take all pending operations (clears the stream)
    pub fn take_operations(&mut self) -> Vec<StreamOperation> {
        self.needs_flush = false;
        std::mem::take(&mut self.operations)
    }

    /// Clear the stream
    pub fn clear(&mut self) {
        self.operations.clear();
        self.dependencies.clear();
        self.external_inputs.clear();
        self.needs_flush = false;
    }

    /// Get fuseable operation chains
    ///
    /// Returns groups of operation indices that can be fused together.
    /// Each group is a contiguous sequence of fuseable operations.
    pub fn find_fuseable_chains(&self) -> Vec<Vec<usize>> {
        if !self.config.enabled || self.operations.is_empty() {
            return Vec::new();
        }

        let mut chains = Vec::new();
        let mut current_chain: Vec<usize> = Vec::new();
        let mut chain_outputs: HashSet<TensorId> = HashSet::new();

        for (idx, op) in self.operations.iter().enumerate() {
            if op.op.is_fuseable() {
                // Check if we can extend the current chain
                let can_extend = current_chain.is_empty()
                    || op.inputs.iter().any(|input| chain_outputs.contains(input));

                if can_extend {
                    // Extend the current chain
                    current_chain.push(idx);
                    chain_outputs.insert(op.output);

                    // Check max length
                    if current_chain.len() >= self.config.max_fusion_ops {
                        if current_chain.len() >= self.config.min_fusion_ops {
                            chains.push(current_chain.clone());
                        }
                        current_chain.clear();
                        chain_outputs.clear();
                    }
                } else {
                    // Start a new chain
                    if current_chain.len() >= self.config.min_fusion_ops {
                        chains.push(current_chain.clone());
                    }
                    current_chain.clear();
                    chain_outputs.clear();
                    current_chain.push(idx);
                    chain_outputs.insert(op.output);
                }
            } else {
                // Non-fuseable operation breaks the chain
                if current_chain.len() >= self.config.min_fusion_ops {
                    chains.push(current_chain.clone());
                }
                current_chain.clear();
                chain_outputs.clear();
            }
        }

        // Don't forget the last chain
        if current_chain.len() >= self.config.min_fusion_ops {
            chains.push(current_chain);
        }

        chains
    }

    /// Build a dependency graph for the operations
    ///
    /// Returns a mapping from operation index to its input operation indices.
    pub fn build_dependency_graph(&self) -> HashMap<usize, Vec<usize>> {
        let mut graph: HashMap<usize, Vec<usize>> = HashMap::new();

        // Build a map from tensor ID to producing operation
        let mut tensor_producers: HashMap<TensorId, usize> = HashMap::new();
        for (idx, op) in self.operations.iter().enumerate() {
            tensor_producers.insert(op.output, idx);
        }

        // Build the dependency graph
        for (idx, op) in self.operations.iter().enumerate() {
            let mut deps = Vec::new();
            for &input_id in &op.inputs {
                if let Some(&producer_idx) = tensor_producers.get(&input_id) {
                    if producer_idx < idx {
                        deps.push(producer_idx);
                    }
                }
            }
            graph.insert(idx, deps);
        }

        graph
    }

    /// Compute topological order of operations
    pub fn topological_order(&self) -> Vec<usize> {
        if self.operations.is_empty() {
            return Vec::new();
        }

        let dep_graph = self.build_dependency_graph();
        let mut in_degree: Vec<usize> = vec![0; self.operations.len()];
        let mut queue: VecDeque<usize> = VecDeque::new();
        let mut result = Vec::with_capacity(self.operations.len());

        // Calculate in-degrees
        for deps in dep_graph.values() {
            for &dep in deps {
                if dep < in_degree.len() {
                    in_degree[dep] += 1;
                }
            }
        }

        // Find nodes with no dependencies
        for (idx, &degree) in in_degree.iter().enumerate() {
            if degree == 0 && dep_graph.contains_key(&idx) {
                queue.push_back(idx);
            }
        }

        // Process nodes
        while let Some(idx) = queue.pop_front() {
            result.push(idx);

            if let Some(consumers) = dep_graph.get(&idx) {
                for &consumer in consumers {
                    if consumer < in_degree.len() {
                        in_degree[consumer] = in_degree[consumer].saturating_sub(1);
                        if in_degree[consumer] == 0 {
                            queue.push_back(consumer);
                        }
                    }
                }
            }
        }

        // If we couldn't order all nodes, fall back to submission order
        if result.len() != self.operations.len() {
            (0..self.operations.len()).collect()
        } else {
            result
        }
    }
}

/// Thread-safe operation stream with per-device queues
pub struct DeviceStreams {
    /// Streams indexed by device ID
    streams: RwLock<HashMap<u64, Arc<RwLock<OperationStream>>>>,
    /// Default configuration for new streams
    default_config: FusionConfig,
}

impl DeviceStreams {
    /// Create a new device stream manager
    pub fn new(config: FusionConfig) -> Self {
        Self {
            streams: RwLock::new(HashMap::new()),
            default_config: config,
        }
    }

    /// Get or create a stream for a device
    pub fn get_or_create(&self, device_id: u64) -> Arc<RwLock<OperationStream>> {
        // Fast path: read lock
        {
            let streams = self.streams.read();
            if let Some(stream) = streams.get(&device_id) {
                return Arc::clone(stream);
            }
        }

        // Slow path: write lock
        let mut streams = self.streams.write();
        streams
            .entry(device_id)
            .or_insert_with(|| {
                Arc::new(RwLock::new(OperationStream::new(
                    self.default_config.clone(),
                )))
            })
            .clone()
    }

    /// Clear all streams
    pub fn clear_all(&self) {
        let streams = self.streams.read();
        for stream in streams.values() {
            stream.write().clear();
        }
    }
}

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

    #[test]
    fn test_empty_stream() {
        let stream = OperationStream::with_defaults();
        assert!(stream.is_empty());
        assert_eq!(stream.len(), 0);
        assert!(!stream.needs_flush());
    }

    #[test]
    fn test_record_operation() {
        let mut stream = OperationStream::with_defaults();

        let input1 = TensorId::new();
        let input2 = TensorId::new();
        let output = TensorId::new();

        stream.register_input(input1);
        stream.register_input(input2);

        stream.record(StreamOperation::new(
            StreamOpKind::Add,
            vec![input1, input2],
            output,
            vec![1024],
            DType::F32,
        ));

        assert_eq!(stream.len(), 1);
        assert!(!stream.needs_flush()); // Add is fuseable
    }

    #[test]
    fn test_sync_point_triggers_flush() {
        let mut stream = OperationStream::with_defaults();

        let input = TensorId::new();
        let output = TensorId::new();

        stream.register_input(input);

        stream.record(StreamOperation::new(
            StreamOpKind::MatMul,
            vec![input],
            output,
            vec![1024, 1024],
            DType::F32,
        ));

        assert!(stream.needs_flush()); // MatMul is a sync point
    }

    #[test]
    fn test_find_fuseable_chains() {
        let mut stream = OperationStream::new(FusionConfig::default().with_min_fusion_ops(2));

        let a = TensorId::new();
        let b = TensorId::new();
        let c = TensorId::new();
        let d = TensorId::new();
        let e = TensorId::new();

        stream.register_input(a);
        stream.register_input(b);

        // Chain: add -> mul -> relu
        stream.record(StreamOperation::new(
            StreamOpKind::Add,
            vec![a, b],
            c,
            vec![1024],
            DType::F32,
        ));

        stream.record(StreamOperation::new(
            StreamOpKind::Mul,
            vec![c, a],
            d,
            vec![1024],
            DType::F32,
        ));

        stream.record(StreamOperation::new(
            StreamOpKind::ReLU,
            vec![d],
            e,
            vec![1024],
            DType::F32,
        ));

        let chains = stream.find_fuseable_chains();
        assert_eq!(chains.len(), 1);
        assert_eq!(chains[0].len(), 3);
    }

    #[test]
    fn test_dependency_graph() {
        let mut stream = OperationStream::with_defaults();

        let a = TensorId::new();
        let b = TensorId::new();
        let c = TensorId::new();

        stream.register_input(a);

        stream.record(StreamOperation::new(
            StreamOpKind::Neg,
            vec![a],
            b,
            vec![1024],
            DType::F32,
        ));

        stream.record(StreamOperation::new(
            StreamOpKind::Exp,
            vec![b],
            c,
            vec![1024],
            DType::F32,
        ));

        let graph = stream.build_dependency_graph();
        assert!(graph.get(&0).unwrap().is_empty()); // First op has no deps
        assert_eq!(graph.get(&1).unwrap(), &vec![0]); // Second op depends on first
    }

    #[test]
    fn test_max_pending_triggers_flush() {
        let config = FusionConfig::default().with_max_pending_ops(3);
        let mut stream = OperationStream::new(config);

        let a = TensorId::new();
        stream.register_input(a);

        for i in 0..3 {
            let output = TensorId::new();
            stream.record(StreamOperation::new(
                StreamOpKind::Neg,
                vec![a],
                output,
                vec![1024],
                DType::F32,
            ));
        }

        assert!(stream.needs_flush());
    }
}