rustytorch/crates/training/rtx-distributed/src/collective_fusion.rs

//! Collective Operation Fusion
//!
//! This module provides fusion of collective operations for efficiency:
//! - Batches multiple small collectives into larger operations
//! - Reduces kernel launch overhead
//! - Optimizes network utilization
//! - Supports various fusion strategies
//!
//! Fusion can improve throughput by 10-50% for models with many
//! small parameter tensors.

use crate::error::{DistributedError, Result};
use parking_lot::{Mutex, RwLock};
use serde::{Deserialize, Serialize};
use std::collections::{HashMap, VecDeque};
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::time::{Duration, Instant};

// =============================================================================
// Configuration
// =============================================================================

/// Fusion strategy
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum FusionStrategy {
    /// No fusion - execute immediately
    None,
    /// Size-based fusion - fuse until buffer is full
    SizeBased,
    /// Count-based fusion - fuse until N operations
    CountBased,
    /// Time-based fusion - fuse until timeout
    TimeBased,
    /// Adaptive fusion based on throughput
    Adaptive,
}

/// Collective operation type
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum CollectiveType {
    /// AllReduce
    AllReduce,
    /// AllGather
    AllGather,
    /// ReduceScatter
    ReduceScatter,
    /// Broadcast
    Broadcast,
    /// Reduce
    Reduce,
    /// Scatter
    Scatter,
    /// Gather
    Gather,
}

/// Configuration for collective fusion
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FusionConfig {
    /// Fusion strategy
    pub strategy: FusionStrategy,
    /// Maximum fusion buffer size (bytes)
    pub max_buffer_size: usize,
    /// Maximum number of operations to fuse
    pub max_ops_per_fusion: usize,
    /// Fusion timeout (ms)
    pub fusion_timeout_ms: u64,
    /// Minimum size to consider for fusion (bytes)
    pub min_fusion_size: usize,
    /// Enable per-collective-type buffers
    pub per_type_buffers: bool,
    /// Alignment for fused buffers
    pub alignment: usize,
    /// Enable overlap of fusion with computation
    pub overlap_fusion: bool,
}

impl Default for FusionConfig {
    fn default() -> Self {
        Self {
            strategy: FusionStrategy::SizeBased,
            max_buffer_size: 25 * 1024 * 1024, // 25MB
            max_ops_per_fusion: 64,
            fusion_timeout_ms: 10,
            min_fusion_size: 256,
            per_type_buffers: true,
            alignment: 256,
            overlap_fusion: true,
        }
    }
}

// =============================================================================
// Fusion Operation
// =============================================================================

/// A single operation to be fused
#[derive(Debug, Clone)]
pub struct FusionOp {
    /// Operation ID
    pub id: u64,
    /// Collective type
    pub collective_type: CollectiveType,
    /// Parameter name
    pub name: String,
    /// Data size in bytes
    pub size_bytes: usize,
    /// Offset in fusion buffer
    pub buffer_offset: Option<usize>,
    /// Creation time
    pub created_at: Instant,
    /// Priority (lower = higher priority)
    pub priority: i32,
    /// Is ready for execution
    pub ready: bool,
}

impl FusionOp {
    /// Create a new fusion operation
    pub fn new(id: u64, collective_type: CollectiveType, name: String, size_bytes: usize) -> Self {
        Self {
            id,
            collective_type,
            name,
            size_bytes,
            buffer_offset: None,
            created_at: Instant::now(),
            priority: 0,
            ready: true,
        }
    }

    /// Get age of this operation
    pub fn age(&self) -> Duration {
        self.created_at.elapsed()
    }
}

// =============================================================================
// Fusion Group
// =============================================================================

/// A group of fused operations
#[derive(Debug, Clone)]
pub struct FusionGroup {
    /// Group ID
    pub id: u64,
    /// Collective type
    pub collective_type: CollectiveType,
    /// Operations in this group
    pub ops: Vec<FusionOp>,
    /// Total size in bytes
    pub total_size: usize,
    /// Fused buffer data
    pub buffer: Vec<u8>,
    /// Creation time
    pub created_at: Instant,
    /// Is group sealed (no more additions)
    pub sealed: bool,
}

impl FusionGroup {
    /// Create a new fusion group
    pub fn new(id: u64, collective_type: CollectiveType) -> Self {
        Self {
            id,
            collective_type,
            ops: Vec::new(),
            total_size: 0,
            buffer: Vec::new(),
            created_at: Instant::now(),
            sealed: false,
        }
    }

    /// Add an operation to the group
    pub fn add_op(&mut self, mut op: FusionOp, alignment: usize) -> bool {
        if self.sealed {
            return false;
        }

        if op.collective_type != self.collective_type {
            return false;
        }

        // Calculate aligned offset
        let aligned_offset = (self.total_size + alignment - 1) & !(alignment - 1);
        let padding = aligned_offset - self.total_size;

        op.buffer_offset = Some(aligned_offset);
        self.total_size = aligned_offset + op.size_bytes;
        self.ops.push(op);

        // Add padding to buffer
        self.buffer.extend(std::iter::repeat_n(0u8, padding));

        true
    }

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

    /// Check if group is empty
    pub fn is_empty(&self) -> bool {
        self.ops.is_empty()
    }

    /// Seal the group (prevent further additions)
    pub fn seal(&mut self) {
        self.sealed = true;
    }

    /// Get age of this group
    pub fn age(&self) -> Duration {
        self.created_at.elapsed()
    }

    /// Get operation by name
    pub fn get_op(&self, name: &str) -> Option<&FusionOp> {
        self.ops.iter().find(|op| op.name == name)
    }
}

// =============================================================================
// Fusion Buffer
// =============================================================================

/// Buffer for accumulating operations to fuse
pub struct FusionBuffer {
    /// Configuration
    config: FusionConfig,
    /// Current fusion group
    current_group: Mutex<Option<FusionGroup>>,
    /// Next group ID
    next_group_id: AtomicU64,
    /// Next operation ID
    next_op_id: AtomicU64,
    /// Collective type this buffer handles
    collective_type: CollectiveType,
    /// Ready groups waiting to execute
    ready_groups: Mutex<VecDeque<FusionGroup>>,
    /// Statistics
    stats: RwLock<FusionBufferStats>,
}

impl FusionBuffer {
    /// Create a new fusion buffer
    pub fn new(config: FusionConfig, collective_type: CollectiveType) -> Self {
        Self {
            config,
            current_group: Mutex::new(None),
            next_group_id: AtomicU64::new(1),
            next_op_id: AtomicU64::new(1),
            collective_type,
            ready_groups: Mutex::new(VecDeque::new()),
            stats: RwLock::new(FusionBufferStats::default()),
        }
    }

    /// Submit an operation for fusion
    pub fn submit(&self, name: String, size_bytes: usize) -> u64 {
        let op_id = self.next_op_id.fetch_add(1, Ordering::SeqCst);
        let op = FusionOp::new(op_id, self.collective_type, name, size_bytes);

        let mut current = self.current_group.lock();

        // Check if we need a new group
        let need_new_group = match &*current {
            None => true,
            Some(group) => self.should_flush_group(group, size_bytes),
        };

        if need_new_group {
            // Flush current group if exists
            if let Some(mut group) = current.take() {
                if !group.is_empty() {
                    group.seal();
                    self.ready_groups.lock().push_back(group);

                    let mut stats = self.stats.write();
                    stats.groups_flushed += 1;
                }
            }

            // Create new group
            let group_id = self.next_group_id.fetch_add(1, Ordering::SeqCst);
            *current = Some(FusionGroup::new(group_id, self.collective_type));
        }

        // Add to current group
        if let Some(group) = current.as_mut() {
            group.add_op(op, self.config.alignment);

            let mut stats = self.stats.write();
            stats.ops_submitted += 1;
            stats.bytes_submitted += size_bytes;
        }

        op_id
    }

    /// Check if group should be flushed
    fn should_flush_group(&self, group: &FusionGroup, new_size: usize) -> bool {
        match self.config.strategy {
            FusionStrategy::None => true,
            FusionStrategy::SizeBased => group.total_size + new_size > self.config.max_buffer_size,
            FusionStrategy::CountBased => group.num_ops() >= self.config.max_ops_per_fusion,
            FusionStrategy::TimeBased => {
                group.age() >= Duration::from_millis(self.config.fusion_timeout_ms)
            }
            FusionStrategy::Adaptive => {
                group.total_size + new_size > self.config.max_buffer_size
                    || group.num_ops() >= self.config.max_ops_per_fusion
                    || group.age() >= Duration::from_millis(self.config.fusion_timeout_ms)
            }
        }
    }

    /// Force flush the current group
    pub fn flush(&self) -> Option<FusionGroup> {
        let mut current = self.current_group.lock();

        if let Some(mut group) = current.take() {
            if !group.is_empty() {
                group.seal();

                let mut stats = self.stats.write();
                stats.groups_flushed += 1;
                stats.ops_fused += group.num_ops();
                stats.bytes_fused += group.total_size;

                return Some(group);
            }
        }

        None
    }

    /// Get next ready group
    pub fn pop_ready(&self) -> Option<FusionGroup> {
        self.ready_groups.lock().pop_front()
    }

    /// Check if there are ready groups
    pub fn has_ready(&self) -> bool {
        !self.ready_groups.lock().is_empty()
    }

    /// Get number of ready groups
    pub fn ready_count(&self) -> usize {
        self.ready_groups.lock().len()
    }

    /// Get current group size
    pub fn current_size(&self) -> usize {
        self.current_group
            .lock()
            .as_ref()
            .map_or(0, |g| g.total_size)
    }

    /// Get current operation count
    pub fn current_ops(&self) -> usize {
        self.current_group
            .lock()
            .as_ref()
            .map_or(0, FusionGroup::num_ops)
    }

    /// Get statistics
    pub fn stats(&self) -> FusionBufferStats {
        self.stats.read().clone()
    }

    /// Get collective type
    pub fn collective_type(&self) -> CollectiveType {
        self.collective_type
    }
}

/// Statistics for fusion buffer
#[derive(Debug, Default, Clone)]
pub struct FusionBufferStats {
    /// Operations submitted
    pub ops_submitted: usize,
    /// Bytes submitted
    pub bytes_submitted: usize,
    /// Operations fused
    pub ops_fused: usize,
    /// Bytes fused
    pub bytes_fused: usize,
    /// Groups flushed
    pub groups_flushed: usize,
}

impl FusionBufferStats {
    /// Get average ops per fusion
    pub fn avg_ops_per_fusion(&self) -> f32 {
        if self.groups_flushed == 0 {
            0.0
        } else {
            self.ops_fused as f32 / self.groups_flushed as f32
        }
    }

    /// Get average bytes per fusion
    pub fn avg_bytes_per_fusion(&self) -> f32 {
        if self.groups_flushed == 0 {
            0.0
        } else {
            self.bytes_fused as f32 / self.groups_flushed as f32
        }
    }
}

// =============================================================================
// Collective Fusion Manager
// =============================================================================

/// Manages fusion of collective operations
pub struct CollectiveFusionManager {
    /// Configuration
    config: FusionConfig,
    /// Buffers per collective type
    buffers: HashMap<CollectiveType, Arc<FusionBuffer>>,
    /// Unified buffer (when per_type_buffers is false)
    unified_buffer: Option<Arc<FusionBuffer>>,
    /// Is fusion enabled
    enabled: AtomicBool,
    /// Statistics
    stats: RwLock<FusionManagerStats>,
}

impl CollectiveFusionManager {
    /// Create a new fusion manager
    pub fn new(config: FusionConfig) -> Self {
        let mut buffers = HashMap::new();

        if config.per_type_buffers {
            // Create buffer for each collective type
            for collective_type in [
                CollectiveType::AllReduce,
                CollectiveType::AllGather,
                CollectiveType::ReduceScatter,
                CollectiveType::Broadcast,
            ] {
                buffers.insert(
                    collective_type,
                    Arc::new(FusionBuffer::new(config.clone(), collective_type)),
                );
            }
        }

        let unified_buffer = if !config.per_type_buffers {
            Some(Arc::new(FusionBuffer::new(
                config.clone(),
                CollectiveType::AllReduce,
            )))
        } else {
            None
        };

        Self {
            config,
            buffers,
            unified_buffer,
            enabled: AtomicBool::new(true),
            stats: RwLock::new(FusionManagerStats::default()),
        }
    }

    /// Submit an operation for fusion
    pub fn submit(
        &self,
        collective_type: CollectiveType,
        name: String,
        size_bytes: usize,
    ) -> Result<u64> {
        if !self.enabled.load(Ordering::SeqCst) {
            return Err(DistributedError::runtime("Fusion manager is disabled"));
        }

        // Skip fusion for large operations
        if size_bytes > self.config.max_buffer_size {
            let mut stats = self.stats.write();
            stats.skipped_large += 1;
            // Return fake ID for immediate execution
            return Ok(0);
        }

        // Skip fusion for very small operations
        if size_bytes < self.config.min_fusion_size {
            let mut stats = self.stats.write();
            stats.skipped_small += 1;
            return Ok(0);
        }

        let buffer = self.get_buffer(collective_type)?;
        let id = buffer.submit(name, size_bytes);

        let mut stats = self.stats.write();
        stats.ops_received += 1;

        Ok(id)
    }

    /// Get buffer for collective type
    fn get_buffer(&self, collective_type: CollectiveType) -> Result<Arc<FusionBuffer>> {
        if let Some(ref unified) = self.unified_buffer {
            return Ok(unified.clone());
        }

        self.buffers.get(&collective_type).cloned().ok_or_else(|| {
            DistributedError::runtime(format!(
                "No buffer for collective type {:?}",
                collective_type
            ))
        })
    }

    /// Flush all buffers
    pub fn flush_all(&self) -> Vec<FusionGroup> {
        let mut groups = Vec::new();

        if let Some(ref unified) = self.unified_buffer {
            if let Some(group) = unified.flush() {
                groups.push(group);
            }
        } else {
            for buffer in self.buffers.values() {
                if let Some(group) = buffer.flush() {
                    groups.push(group);
                }
            }
        }

        let mut stats = self.stats.write();
        stats.flush_count += 1;

        groups
    }

    /// Flush specific buffer
    pub fn flush(&self, collective_type: CollectiveType) -> Option<FusionGroup> {
        if let Ok(buffer) = self.get_buffer(collective_type) {
            buffer.flush()
        } else {
            None
        }
    }

    /// Get all ready groups
    pub fn get_ready_groups(&self) -> Vec<FusionGroup> {
        let mut groups = Vec::new();

        if let Some(ref unified) = self.unified_buffer {
            while let Some(group) = unified.pop_ready() {
                groups.push(group);
            }
        } else {
            for buffer in self.buffers.values() {
                while let Some(group) = buffer.pop_ready() {
                    groups.push(group);
                }
            }
        }

        groups
    }

    /// Check if any buffer has ready groups
    pub fn has_ready(&self) -> bool {
        if let Some(ref unified) = self.unified_buffer {
            return unified.has_ready();
        }

        self.buffers.values().any(|b| b.has_ready())
    }

    /// Enable fusion
    pub fn enable(&self) {
        self.enabled.store(true, Ordering::SeqCst);
    }

    /// Disable fusion
    pub fn disable(&self) {
        self.enabled.store(false, Ordering::SeqCst);
    }

    /// Check if enabled
    pub fn is_enabled(&self) -> bool {
        self.enabled.load(Ordering::SeqCst)
    }

    /// Get statistics
    pub fn stats(&self) -> FusionManagerStats {
        self.stats.read().clone()
    }

    /// Get buffer statistics
    pub fn buffer_stats(&self, collective_type: CollectiveType) -> Option<FusionBufferStats> {
        self.get_buffer(collective_type).ok().map(|b| b.stats())
    }

    /// Get configuration
    pub fn config(&self) -> &FusionConfig {
        &self.config
    }

    /// Get total pending size across all buffers
    pub fn total_pending_size(&self) -> usize {
        if let Some(ref unified) = self.unified_buffer {
            return unified.current_size();
        }

        self.buffers.values().map(|b| b.current_size()).sum()
    }

    /// Get total pending ops across all buffers
    pub fn total_pending_ops(&self) -> usize {
        if let Some(ref unified) = self.unified_buffer {
            return unified.current_ops();
        }

        self.buffers.values().map(|b| b.current_ops()).sum()
    }
}

/// Statistics for fusion manager
#[derive(Debug, Default, Clone)]
pub struct FusionManagerStats {
    /// Operations received
    pub ops_received: usize,
    /// Flush operations
    pub flush_count: usize,
    /// Skipped (too large)
    pub skipped_large: usize,
    /// Skipped (too small)
    pub skipped_small: usize,
}

// =============================================================================
// Thread-Safe Wrappers
// =============================================================================

/// Thread-safe shared fusion manager
pub type SharedFusionManager = Arc<CollectiveFusionManager>;

/// Create a shared fusion manager
pub fn shared_fusion_manager(config: FusionConfig) -> SharedFusionManager {
    Arc::new(CollectiveFusionManager::new(config))
}

// =============================================================================
// Tests
// =============================================================================

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

    #[test]
    fn test_fusion_config_default() {
        let config = FusionConfig::default();
        assert_eq!(config.strategy, FusionStrategy::SizeBased);
        assert_eq!(config.max_buffer_size, 25 * 1024 * 1024);
        assert!(config.per_type_buffers);
    }

    #[test]
    fn test_fusion_op_creation() {
        let op = FusionOp::new(1, CollectiveType::AllReduce, "grad".to_string(), 1024);
        assert_eq!(op.id, 1);
        assert_eq!(op.collective_type, CollectiveType::AllReduce);
        assert_eq!(op.size_bytes, 1024);
        assert!(op.buffer_offset.is_none());
    }

    #[test]
    fn test_fusion_op_age() {
        let op = FusionOp::new(1, CollectiveType::AllReduce, "grad".to_string(), 1024);
        std::thread::sleep(Duration::from_millis(5));
        assert!(op.age() >= Duration::from_millis(5));
    }

    #[test]
    fn test_fusion_group_creation() {
        let group = FusionGroup::new(1, CollectiveType::AllReduce);
        assert_eq!(group.id, 1);
        assert!(group.is_empty());
        assert!(!group.sealed);
    }

    #[test]
    fn test_fusion_group_add_op() {
        let mut group = FusionGroup::new(1, CollectiveType::AllReduce);

        let op = FusionOp::new(1, CollectiveType::AllReduce, "grad1".to_string(), 1024);
        assert!(group.add_op(op, 256));

        assert_eq!(group.num_ops(), 1);
        assert_eq!(group.total_size, 1024);
    }

    #[test]
    fn test_fusion_group_alignment() {
        let mut group = FusionGroup::new(1, CollectiveType::AllReduce);

        // First op: 100 bytes
        let op1 = FusionOp::new(1, CollectiveType::AllReduce, "g1".to_string(), 100);
        group.add_op(op1, 256);

        // Second op should be at aligned offset
        let op2 = FusionOp::new(2, CollectiveType::AllReduce, "g2".to_string(), 50);
        group.add_op(op2, 256);

        let second_offset = group.ops[1].buffer_offset.unwrap();
        assert_eq!(second_offset, 256); // Aligned to 256
    }

    #[test]
    fn test_fusion_group_seal() {
        let mut group = FusionGroup::new(1, CollectiveType::AllReduce);

        let op1 = FusionOp::new(1, CollectiveType::AllReduce, "g1".to_string(), 100);
        assert!(group.add_op(op1, 256));

        group.seal();
        assert!(group.sealed);

        // Cannot add after seal
        let op2 = FusionOp::new(2, CollectiveType::AllReduce, "g2".to_string(), 100);
        assert!(!group.add_op(op2, 256));
    }

    #[test]
    fn test_fusion_group_wrong_type() {
        let mut group = FusionGroup::new(1, CollectiveType::AllReduce);

        // Try to add AllGather to AllReduce group
        let op = FusionOp::new(1, CollectiveType::AllGather, "g".to_string(), 100);
        assert!(!group.add_op(op, 256));
    }

    #[test]
    fn test_fusion_buffer_creation() {
        let config = FusionConfig::default();
        let buffer = FusionBuffer::new(config, CollectiveType::AllReduce);

        assert_eq!(buffer.collective_type(), CollectiveType::AllReduce);
        assert_eq!(buffer.current_size(), 0);
        assert_eq!(buffer.current_ops(), 0);
    }

    #[test]
    fn test_fusion_buffer_submit() {
        let config = FusionConfig::default();
        let buffer = FusionBuffer::new(config, CollectiveType::AllReduce);

        let id = buffer.submit("grad1".to_string(), 1024);
        assert!(id > 0);
        assert_eq!(buffer.current_ops(), 1);
    }

    #[test]
    fn test_fusion_buffer_multiple_submit() {
        let config = FusionConfig::default();
        let buffer = FusionBuffer::new(config, CollectiveType::AllReduce);

        buffer.submit("grad1".to_string(), 1024);
        buffer.submit("grad2".to_string(), 2048);
        buffer.submit("grad3".to_string(), 512);

        assert_eq!(buffer.current_ops(), 3);
    }

    #[test]
    fn test_fusion_buffer_flush() {
        let config = FusionConfig::default();
        let buffer = FusionBuffer::new(config, CollectiveType::AllReduce);

        buffer.submit("grad1".to_string(), 1024);
        buffer.submit("grad2".to_string(), 2048);

        let group = buffer.flush().unwrap();
        assert_eq!(group.num_ops(), 2);
        assert!(group.sealed);

        // Buffer should be empty now
        assert_eq!(buffer.current_ops(), 0);
    }

    #[test]
    fn test_fusion_buffer_size_trigger() {
        let config = FusionConfig {
            strategy: FusionStrategy::SizeBased,
            max_buffer_size: 1000,
            ..Default::default()
        };
        let buffer = FusionBuffer::new(config, CollectiveType::AllReduce);

        // First op fits
        buffer.submit("grad1".to_string(), 500);
        assert_eq!(buffer.ready_count(), 0);

        // Second op exceeds limit, triggers flush
        buffer.submit("grad2".to_string(), 600);
        assert_eq!(buffer.ready_count(), 1);
    }

    #[test]
    fn test_fusion_buffer_count_trigger() {
        let config = FusionConfig {
            strategy: FusionStrategy::CountBased,
            max_ops_per_fusion: 3,
            ..Default::default()
        };
        let buffer = FusionBuffer::new(config, CollectiveType::AllReduce);

        buffer.submit("g1".to_string(), 100);
        buffer.submit("g2".to_string(), 100);
        assert_eq!(buffer.ready_count(), 0);

        buffer.submit("g3".to_string(), 100);
        buffer.submit("g4".to_string(), 100); // Triggers flush
        assert_eq!(buffer.ready_count(), 1);
    }

    #[test]
    fn test_fusion_buffer_stats() {
        let config = FusionConfig::default();
        let buffer = FusionBuffer::new(config, CollectiveType::AllReduce);

        buffer.submit("g1".to_string(), 1000);
        buffer.submit("g2".to_string(), 2000);
        buffer.flush();

        let stats = buffer.stats();
        assert_eq!(stats.ops_submitted, 2);
        assert_eq!(stats.bytes_submitted, 3000);
        assert_eq!(stats.groups_flushed, 1);
    }

    #[test]
    fn test_fusion_manager_creation() {
        let config = FusionConfig::default();
        let manager = CollectiveFusionManager::new(config);

        assert!(manager.is_enabled());
        assert_eq!(manager.total_pending_ops(), 0);
    }

    #[test]
    fn test_fusion_manager_submit() {
        let config = FusionConfig::default();
        let manager = CollectiveFusionManager::new(config);

        let id = manager
            .submit(CollectiveType::AllReduce, "grad".to_string(), 1024)
            .unwrap();
        assert!(id > 0);
        assert_eq!(manager.total_pending_ops(), 1);
    }

    #[test]
    fn test_fusion_manager_multiple_types() {
        let config = FusionConfig {
            per_type_buffers: true,
            ..Default::default()
        };
        let manager = CollectiveFusionManager::new(config);

        manager
            .submit(CollectiveType::AllReduce, "ar".to_string(), 1024)
            .unwrap();
        manager
            .submit(CollectiveType::AllGather, "ag".to_string(), 2048)
            .unwrap();

        // Should have ops in different buffers
        let ar_stats = manager.buffer_stats(CollectiveType::AllReduce).unwrap();
        let ag_stats = manager.buffer_stats(CollectiveType::AllGather).unwrap();

        assert_eq!(ar_stats.ops_submitted, 1);
        assert_eq!(ag_stats.ops_submitted, 1);
    }

    #[test]
    fn test_fusion_manager_flush_all() {
        let config = FusionConfig::default();
        let manager = CollectiveFusionManager::new(config);

        manager
            .submit(CollectiveType::AllReduce, "g1".to_string(), 1024)
            .unwrap();
        manager
            .submit(CollectiveType::AllReduce, "g2".to_string(), 2048)
            .unwrap();

        let groups = manager.flush_all();
        assert_eq!(groups.len(), 1);
        assert_eq!(groups[0].num_ops(), 2);
    }

    #[test]
    fn test_fusion_manager_skip_large() {
        let config = FusionConfig {
            max_buffer_size: 1000,
            ..Default::default()
        };
        let manager = CollectiveFusionManager::new(config);

        // Operation larger than buffer
        let id = manager
            .submit(CollectiveType::AllReduce, "big".to_string(), 5000)
            .unwrap();
        assert_eq!(id, 0); // Skipped

        let stats = manager.stats();
        assert_eq!(stats.skipped_large, 1);
    }

    #[test]
    fn test_fusion_manager_skip_small() {
        let config = FusionConfig {
            min_fusion_size: 500,
            ..Default::default()
        };
        let manager = CollectiveFusionManager::new(config);

        // Operation smaller than minimum
        let id = manager
            .submit(CollectiveType::AllReduce, "tiny".to_string(), 100)
            .unwrap();
        assert_eq!(id, 0); // Skipped

        let stats = manager.stats();
        assert_eq!(stats.skipped_small, 1);
    }

    #[test]
    fn test_fusion_manager_enable_disable() {
        let config = FusionConfig::default();
        let manager = CollectiveFusionManager::new(config);

        assert!(manager.is_enabled());

        manager.disable();
        assert!(!manager.is_enabled());

        let result = manager.submit(CollectiveType::AllReduce, "g".to_string(), 1024);
        assert!(result.is_err());

        manager.enable();
        assert!(manager.is_enabled());
    }

    #[test]
    fn test_fusion_manager_ready_groups() {
        let config = FusionConfig {
            strategy: FusionStrategy::SizeBased,
            max_buffer_size: 500,
            ..Default::default()
        };
        let manager = CollectiveFusionManager::new(config);

        // Trigger automatic flush
        manager
            .submit(CollectiveType::AllReduce, "g1".to_string(), 400)
            .unwrap();
        manager
            .submit(CollectiveType::AllReduce, "g2".to_string(), 400)
            .unwrap();

        assert!(manager.has_ready());

        let groups = manager.get_ready_groups();
        assert!(!groups.is_empty());
    }

    #[test]
    fn test_shared_fusion_manager() {
        let config = FusionConfig::default();
        let manager = shared_fusion_manager(config);

        manager
            .submit(CollectiveType::AllReduce, "g".to_string(), 1024)
            .unwrap();
        assert_eq!(manager.total_pending_ops(), 1);
    }

    #[test]
    fn test_fusion_buffer_stats_averages() {
        let stats = FusionBufferStats {
            ops_fused: 100,
            bytes_fused: 50000,
            groups_flushed: 10,
            ..Default::default()
        };

        assert!((stats.avg_ops_per_fusion() - 10.0).abs() < 0.01);
        assert!((stats.avg_bytes_per_fusion() - 5000.0).abs() < 0.01);
    }

    #[test]
    fn test_fusion_group_get_op() {
        let mut group = FusionGroup::new(1, CollectiveType::AllReduce);

        let op = FusionOp::new(1, CollectiveType::AllReduce, "my_grad".to_string(), 1024);
        group.add_op(op, 256);

        assert!(group.get_op("my_grad").is_some());
        assert!(group.get_op("other").is_none());
    }

    #[test]
    fn test_strategy_none() {
        let config = FusionConfig {
            strategy: FusionStrategy::None,
            ..Default::default()
        };
        let buffer = FusionBuffer::new(config, CollectiveType::AllReduce);

        // Each submit should trigger immediate flush to ready queue
        buffer.submit("g1".to_string(), 100);
        assert_eq!(buffer.ready_count(), 0); // First op creates group

        buffer.submit("g2".to_string(), 100);
        assert_eq!(buffer.ready_count(), 1); // Triggers flush
    }

    #[test]
    fn test_unified_buffer() {
        let config = FusionConfig {
            per_type_buffers: false,
            min_fusion_size: 0, // Allow any size for this test
            ..Default::default()
        };
        let manager = CollectiveFusionManager::new(config);

        // All types go to same buffer
        manager
            .submit(CollectiveType::AllReduce, "ar".to_string(), 100)
            .unwrap();
        manager
            .submit(CollectiveType::AllGather, "ag".to_string(), 100)
            .unwrap();

        assert_eq!(manager.total_pending_ops(), 2);

        let groups = manager.flush_all();
        assert_eq!(groups.len(), 1);
        assert_eq!(groups[0].num_ops(), 2);
    }
}