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

//! Communication Overlap Optimization
//!
//! This module provides infrastructure for overlapping computation with communication
//! to maximize GPU utilization during distributed training.
//!
//! Key features:
//! - Bucketed gradient communication for reduced kernel launch overhead
//! - Stream-based async operations for compute/comm overlap
//! - Priority-based scheduling for optimal bandwidth utilization
//! - Automatic bucket size tuning based on hardware capabilities

use crate::comm::ReduceOp;
use crate::error::{DistributedError, Result};
use crate::group::ProcessGroup;
use crate::hardware_topology::HardwareTopology;
use crate::nvlink_p2p::P2PManager;
use parking_lot::RwLock;
use rtx_tensor::Tensor;
use serde::{Deserialize, Serialize};
use std::collections::VecDeque;
use std::sync::Arc;
use std::time::{Duration, Instant};

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

/// Configuration for communication overlap
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct OverlapConfig {
    /// Enable overlap optimization
    pub enabled: bool,
    /// Bucket size in bytes (default: 25MB, optimal for most networks)
    pub bucket_size_bytes: usize,
    /// Maximum number of concurrent communication operations
    pub max_concurrent_ops: usize,
    /// Enable priority scheduling (larger buckets first)
    pub priority_scheduling: bool,
    /// Minimum tensor size to include in bucketing (elements)
    pub min_tensor_elements: usize,
    /// Use separate streams for communication
    pub use_comm_stream: bool,
    /// Enable double buffering for gradient buckets
    pub double_buffering: bool,
    /// Auto-tune bucket sizes based on bandwidth
    pub auto_tune: bool,
    /// Minimum bucket size for auto-tuning (bytes)
    pub min_bucket_size: usize,
    /// Maximum bucket size for auto-tuning (bytes)
    pub max_bucket_size: usize,
    /// Target overlap ratio (0.0-1.0) for auto-tuning
    pub target_overlap_ratio: f32,
    /// Number of warmup iterations before auto-tuning
    pub auto_tune_warmup: usize,
}

impl Default for OverlapConfig {
    fn default() -> Self {
        Self {
            enabled: true,
            bucket_size_bytes: 25 * 1024 * 1024, // 25MB default bucket
            max_concurrent_ops: 2,
            priority_scheduling: true,
            min_tensor_elements: 1000,
            use_comm_stream: true,
            double_buffering: true,
            auto_tune: false,
            min_bucket_size: 1024 * 1024,       // 1MB minimum
            max_bucket_size: 256 * 1024 * 1024, // 256MB maximum
            target_overlap_ratio: 0.8,          // Target 80% overlap
            auto_tune_warmup: 10,               // 10 warmup iterations
        }
    }
}

impl OverlapConfig {
    /// Create config optimized for high-bandwidth networks (e.g., InfiniBand)
    pub fn high_bandwidth() -> Self {
        Self {
            bucket_size_bytes: 50 * 1024 * 1024, // 50MB for better bandwidth utilization
            max_concurrent_ops: 4,
            ..Default::default()
        }
    }

    /// Create config optimized for low-latency networks
    pub fn low_latency() -> Self {
        Self {
            bucket_size_bytes: 5 * 1024 * 1024, // 5MB for faster startup
            max_concurrent_ops: 8,
            ..Default::default()
        }
    }

    /// Create config with auto-tuning enabled
    pub fn with_auto_tune() -> Self {
        Self {
            auto_tune: true,
            auto_tune_warmup: 10,
            ..Default::default()
        }
    }

    /// Set bucket size
    pub fn with_bucket_size(mut self, size_bytes: usize) -> Self {
        self.bucket_size_bytes = size_bytes;
        self
    }

    /// Set maximum concurrent operations
    pub fn with_max_concurrent_ops(mut self, max_ops: usize) -> Self {
        self.max_concurrent_ops = max_ops;
        self
    }
}

// =============================================================================
// Gradient Bucket
// =============================================================================

/// A bucket containing multiple gradients for batched communication
#[derive(Debug)]
pub struct GradientBucket {
    /// Unique bucket identifier
    pub id: usize,
    /// Gradients in this bucket (name -> tensor)
    gradients: Vec<(String, Tensor)>,
    /// Total size in bytes
    pub size_bytes: usize,
    /// Total number of elements
    pub num_elements: usize,
    /// Whether this bucket is ready for communication
    pub ready: bool,
    /// Number of gradients expected before bucket is ready
    expected_count: usize,
    /// Current gradient count
    current_count: usize,
    /// Creation timestamp
    created_at: Instant,
    /// Communication priority (higher = sooner)
    pub priority: i32,
}

impl GradientBucket {
    /// Create a new empty bucket
    pub fn new(id: usize, expected_count: usize) -> Self {
        Self {
            id,
            gradients: Vec::with_capacity(expected_count),
            size_bytes: 0,
            num_elements: 0,
            ready: false,
            expected_count,
            current_count: 0,
            created_at: Instant::now(),
            priority: 0,
        }
    }

    /// Add a gradient to this bucket
    pub fn add_gradient(&mut self, name: String, gradient: Tensor) {
        let elem_size = gradient.numel();
        let byte_size = elem_size * std::mem::size_of::<f32>(); // Assuming f32

        self.num_elements += elem_size;
        self.size_bytes += byte_size;
        self.gradients.push((name, gradient));
        self.current_count += 1;

        if self.current_count >= self.expected_count {
            self.ready = true;
        }
    }

    /// Check if bucket can accept more gradients
    pub fn has_capacity(&self, max_bytes: usize) -> bool {
        self.size_bytes < max_bytes && !self.ready
    }

    /// Get all gradients in this bucket
    pub fn gradients(&self) -> &[(String, Tensor)] {
        &self.gradients
    }

    /// Take ownership of gradients
    pub fn take_gradients(&mut self) -> Vec<(String, Tensor)> {
        std::mem::take(&mut self.gradients)
    }

    /// Get time since bucket creation
    pub fn age(&self) -> Duration {
        self.created_at.elapsed()
    }
}

// =============================================================================
// Bucket Manager
// =============================================================================

/// State of a communication operation
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CommState {
    /// Waiting for gradients
    Pending,
    /// Ready for communication
    Ready,
    /// Communication in progress
    InProgress,
    /// Communication completed
    Completed,
    /// Communication failed
    Failed,
}

/// Tracks an in-flight communication operation
#[derive(Debug)]
pub struct CommOperation {
    /// Bucket being communicated
    pub bucket_id: usize,
    /// Current state
    pub state: CommState,
    /// Start time of communication
    pub start_time: Option<Instant>,
    /// Completion time
    pub end_time: Option<Instant>,
    /// Error message if failed
    pub error: Option<String>,
}

/// Manages gradient buckets and their communication
pub struct BucketManager {
    /// Configuration
    config: OverlapConfig,
    /// Active buckets (not yet ready)
    active_buckets: Vec<GradientBucket>,
    /// Ready buckets (waiting for communication)
    ready_queue: VecDeque<GradientBucket>,
    /// In-flight communication operations
    in_flight: Vec<CommOperation>,
    /// Completed bucket IDs
    completed: Vec<usize>,
    /// Next bucket ID
    next_bucket_id: usize,
    /// Statistics
    stats: BucketStats,
    /// Gradient name to bucket ID mapping
    gradient_bucket_map: std::collections::HashMap<String, usize>,
    /// Auto-tuner for bucket size optimization
    auto_tuner: Option<BucketAutoTuner>,
    /// Track compute time for overlap calculation
    last_compute_start: Option<Instant>,
    /// Total compute time in current iteration
    current_iteration_compute: Duration,
}

/// Statistics for bucket management
#[derive(Debug, Default, Clone)]
pub struct BucketStats {
    /// Total buckets created
    pub buckets_created: usize,
    /// Total buckets communicated
    pub buckets_communicated: usize,
    /// Total bytes communicated
    pub bytes_communicated: usize,
    /// Total communication time
    pub total_comm_time: Duration,
    /// Average bucket fill ratio
    pub avg_bucket_fill_ratio: f32,
    /// Number of priority inversions (smaller bucket sent before larger)
    pub priority_inversions: usize,
    /// Estimated bandwidth (bytes/sec)
    pub estimated_bandwidth: f64,
    /// Average compute time between buckets
    pub avg_compute_time: Duration,
    /// Achieved overlap ratio (0.0-1.0)
    pub overlap_ratio: f32,
}

// =============================================================================
// Auto-Tuner
// =============================================================================

/// Auto-tuner for bucket sizes and overlap scheduling
#[derive(Debug)]
pub struct BucketAutoTuner {
    /// Current bucket size
    current_bucket_size: usize,
    /// Minimum bucket size
    min_bucket_size: usize,
    /// Maximum bucket size
    max_bucket_size: usize,
    /// Target overlap ratio
    target_overlap: f32,
    /// History of (bucket_size, bandwidth, overlap_ratio) tuples
    history: Vec<TuningDataPoint>,
    /// Number of warmup iterations remaining
    warmup_remaining: usize,
    /// Best configuration found
    best_config: Option<(usize, f64)>, // (bucket_size, score)
}

#[derive(Debug, Clone)]
struct TuningDataPoint {
    bucket_size: usize,
    bandwidth: f64,
    overlap_ratio: f32,
    timestamp: Instant,
}

impl BucketAutoTuner {
    /// Create a new auto-tuner
    pub fn new(config: &OverlapConfig) -> Self {
        Self {
            current_bucket_size: config.bucket_size_bytes,
            min_bucket_size: config.min_bucket_size,
            max_bucket_size: config.max_bucket_size,
            target_overlap: config.target_overlap_ratio,
            history: Vec::new(),
            warmup_remaining: config.auto_tune_warmup,
            best_config: None,
        }
    }

    /// Record a measurement
    pub fn record(
        &mut self,
        bucket_size: usize,
        bytes: usize,
        duration: Duration,
        overlap_ratio: f32,
    ) {
        if self.warmup_remaining > 0 {
            self.warmup_remaining -= 1;
            return;
        }

        let bandwidth = if duration.as_secs_f64() > 0.0 {
            bytes as f64 / duration.as_secs_f64()
        } else {
            0.0
        };

        self.history.push(TuningDataPoint {
            bucket_size,
            bandwidth,
            overlap_ratio,
            timestamp: Instant::now(),
        });

        // Keep only recent history (last 100 data points)
        if self.history.len() > 100 {
            self.history.remove(0);
        }
    }

    /// Get the recommended bucket size
    pub fn recommend_bucket_size(&mut self) -> usize {
        if self.history.is_empty() {
            return self.current_bucket_size;
        }

        // Calculate score for recent data points
        // Score = bandwidth * overlap_ratio (we want high bandwidth AND high overlap)
        let mut size_scores: std::collections::HashMap<usize, (f64, usize)> =
            std::collections::HashMap::new();

        for point in &self.history {
            // Penalize if overlap is below target
            let overlap_factor = if point.overlap_ratio < self.target_overlap {
                point.overlap_ratio / self.target_overlap
            } else {
                1.0
            };

            let score = point.bandwidth * overlap_factor as f64;

            let entry = size_scores.entry(point.bucket_size).or_insert((0.0, 0));
            entry.0 += score;
            entry.1 += 1;
        }

        // Find best average score
        let best = size_scores
            .iter()
            .map(|(&size, &(sum, count))| (size, sum / count as f64))
            .max_by(|a, b| a.1.partial_cmp(&b.1).unwrap());

        if let Some((size, score)) = best {
            // Only update if significantly better
            let should_update = match self.best_config {
                Some((_, best_score)) => score > best_score * 1.1, // 10% improvement threshold
                None => true,
            };

            if should_update {
                self.best_config = Some((size, score));
                self.current_bucket_size = size;
            }
        }

        // Occasionally try exploring new sizes
        if self.should_explore() {
            self.explore_new_size()
        } else {
            self.current_bucket_size
        }
    }

    /// Determine if we should explore a new bucket size
    fn should_explore(&self) -> bool {
        // Explore every 20 iterations
        self.history.len() % 20 == 0
    }

    /// Explore a new bucket size
    fn explore_new_size(&mut self) -> usize {
        use rand::Rng;
        let mut rng = rand::thread_rng();

        // Choose exploration direction
        let direction = rng.gen_range(0..3);

        match direction {
            0 => {
                // Try smaller bucket
                (self.current_bucket_size / 2).max(self.min_bucket_size)
            }
            1 => {
                // Try larger bucket
                (self.current_bucket_size * 2).min(self.max_bucket_size)
            }
            _ => {
                // Random size within range
                rng.gen_range(self.min_bucket_size..=self.max_bucket_size)
            }
        }
    }

    /// Get current statistics
    pub fn stats(&self) -> AutoTunerStats {
        let recent_bandwidth: f64 = if !self.history.is_empty() {
            self.history.iter().map(|p| p.bandwidth).sum::<f64>() / self.history.len() as f64
        } else {
            0.0
        };

        let recent_overlap: f32 = if !self.history.is_empty() {
            self.history.iter().map(|p| p.overlap_ratio).sum::<f32>() / self.history.len() as f32
        } else {
            0.0
        };

        AutoTunerStats {
            current_bucket_size: self.current_bucket_size,
            best_bucket_size: self.best_config.map(|(s, _)| s),
            average_bandwidth: recent_bandwidth,
            average_overlap_ratio: recent_overlap,
            data_points: self.history.len(),
        }
    }
}

/// Statistics from the auto-tuner
#[derive(Debug, Clone)]
pub struct AutoTunerStats {
    /// Current bucket size being used
    pub current_bucket_size: usize,
    /// Best bucket size found so far
    pub best_bucket_size: Option<usize>,
    /// Average bandwidth achieved (bytes/sec)
    pub average_bandwidth: f64,
    /// Average overlap ratio achieved
    pub average_overlap_ratio: f32,
    /// Number of data points collected
    pub data_points: usize,
}

impl BucketManager {
    /// Create a new bucket manager
    pub fn new(config: OverlapConfig) -> Self {
        let auto_tuner = if config.auto_tune {
            Some(BucketAutoTuner::new(&config))
        } else {
            None
        };

        Self {
            config,
            active_buckets: Vec::new(),
            ready_queue: VecDeque::new(),
            in_flight: Vec::new(),
            completed: Vec::new(),
            next_bucket_id: 0,
            stats: BucketStats::default(),
            gradient_bucket_map: std::collections::HashMap::new(),
            auto_tuner,
            last_compute_start: None,
            current_iteration_compute: Duration::ZERO,
        }
    }

    /// Register a gradient with the bucket manager
    pub fn register_gradient(&mut self, name: &str, num_elements: usize) -> usize {
        let bucket_id = self.find_or_create_bucket(num_elements);
        self.gradient_bucket_map.insert(name.to_string(), bucket_id);
        bucket_id
    }

    /// Find an existing bucket with capacity or create a new one
    fn find_or_create_bucket(&mut self, _additional_elements: usize) -> usize {
        // Get current bucket size (possibly tuned)
        let bucket_size = self.current_bucket_size();

        // Try to find existing bucket with capacity
        for bucket in &self.active_buckets {
            if bucket.has_capacity(bucket_size) {
                return bucket.id;
            }
        }

        // Create new bucket
        let bucket_id = self.next_bucket_id;
        self.next_bucket_id += 1;

        let bucket = GradientBucket::new(bucket_id, 1);
        self.active_buckets.push(bucket);
        self.stats.buckets_created += 1;

        bucket_id
    }

    /// Get the current bucket size (tuned or configured)
    fn current_bucket_size(&mut self) -> usize {
        if let Some(ref mut tuner) = self.auto_tuner {
            tuner.recommend_bucket_size()
        } else {
            self.config.bucket_size_bytes
        }
    }

    /// Add a computed gradient to its bucket
    pub fn add_gradient(&mut self, name: &str, gradient: Tensor) -> Result<Option<usize>> {
        // Track compute time (time since last gradient was added)
        if let Some(start) = self.last_compute_start.take() {
            self.current_iteration_compute += start.elapsed();
        }
        self.last_compute_start = Some(Instant::now());

        let bucket_id = self.gradient_bucket_map.get(name).copied().ok_or_else(|| {
            DistributedError::configuration(format!(
                "Gradient '{}' not registered with bucket manager",
                name
            ))
        })?;

        // Get current bucket size threshold
        let bucket_size_threshold = self.current_bucket_size();

        // Find the bucket
        let bucket_idx = self.active_buckets.iter().position(|b| b.id == bucket_id);

        if let Some(idx) = bucket_idx {
            self.active_buckets[idx].add_gradient(name.to_string(), gradient);

            // Check if bucket is now ready
            if self.active_buckets[idx].ready
                || self.active_buckets[idx].size_bytes >= bucket_size_threshold
            {
                let mut bucket = self.active_buckets.remove(idx);
                bucket.ready = true;

                // Assign priority based on size (larger = higher priority)
                if self.config.priority_scheduling {
                    bucket.priority = bucket.size_bytes as i32;
                }

                self.ready_queue.push_back(bucket);
                return Ok(Some(bucket_id));
            }
        }

        Ok(None)
    }

    /// Get the next ready bucket for communication
    pub fn next_ready_bucket(&mut self) -> Option<GradientBucket> {
        if !self.config.priority_scheduling {
            return self.ready_queue.pop_front();
        }

        // Find highest priority bucket
        if self.ready_queue.is_empty() {
            return None;
        }

        let max_priority_idx = self
            .ready_queue
            .iter()
            .enumerate()
            .max_by_key(|(_, b)| b.priority)
            .map(|(i, _)| i)?;

        // Check for priority inversion
        if max_priority_idx != 0 {
            self.stats.priority_inversions += 1;
        }

        Some(self.ready_queue.remove(max_priority_idx).unwrap())
    }

    /// Mark a bucket as in-flight
    pub fn start_communication(&mut self, bucket_id: usize) {
        self.in_flight.push(CommOperation {
            bucket_id,
            state: CommState::InProgress,
            start_time: Some(Instant::now()),
            end_time: None,
            error: None,
        });
    }

    /// Mark a bucket as completed
    pub fn complete_communication(
        &mut self,
        bucket_id: usize,
        bytes_transferred: usize,
        success: bool,
        error: Option<String>,
    ) {
        if let Some(op) = self
            .in_flight
            .iter_mut()
            .find(|op| op.bucket_id == bucket_id)
        {
            op.end_time = Some(Instant::now());
            op.state = if success {
                CommState::Completed
            } else {
                CommState::Failed
            };
            op.error = error;

            if let (Some(start), Some(end)) = (op.start_time, op.end_time) {
                let comm_duration = end.duration_since(start);
                self.stats.total_comm_time += comm_duration;

                // Update bandwidth estimate
                if comm_duration.as_secs_f64() > 0.0 {
                    let bandwidth = bytes_transferred as f64 / comm_duration.as_secs_f64();
                    // Exponential moving average
                    let alpha = 0.1;
                    self.stats.estimated_bandwidth =
                        alpha * bandwidth + (1.0 - alpha) * self.stats.estimated_bandwidth;
                }

                // Feed data to auto-tuner
                if success {
                    let bucket_size = self.current_bucket_size();
                    let overlap_ratio = self.calculate_overlap_ratio(comm_duration);

                    if let Some(ref mut tuner) = self.auto_tuner {
                        tuner.record(bucket_size, bytes_transferred, comm_duration, overlap_ratio);
                    }

                    // Update stats
                    self.stats.overlap_ratio = overlap_ratio;
                    self.stats.bytes_communicated += bytes_transferred;
                }
            }

            if success {
                self.stats.buckets_communicated += 1;
                self.completed.push(bucket_id);
            }
        }
    }

    /// Calculate overlap ratio based on compute and comm times
    fn calculate_overlap_ratio(&self, comm_duration: Duration) -> f32 {
        let compute_time = self.current_iteration_compute.as_secs_f64();
        let comm_time = comm_duration.as_secs_f64();

        if compute_time <= 0.0 || comm_time <= 0.0 {
            return 0.0;
        }

        // Overlap ratio: how much of communication time overlaps with compute
        // Perfect overlap (1.0) = all comm happens during compute
        // No overlap (0.0) = comm blocks compute entirely
        let total_time = compute_time.max(comm_time);
        let sequential_time = compute_time + comm_time;
        let overlap_time = sequential_time - total_time;

        (overlap_time / comm_time).clamp(0.0, 1.0) as f32
    }

    /// Get number of in-flight operations
    pub fn in_flight_count(&self) -> usize {
        self.in_flight
            .iter()
            .filter(|op| op.state == CommState::InProgress)
            .count()
    }

    /// Check if we can start more communication
    pub fn can_start_comm(&self) -> bool {
        self.in_flight_count() < self.config.max_concurrent_ops && !self.ready_queue.is_empty()
    }

    /// Get statistics
    pub fn stats(&self) -> &BucketStats {
        &self.stats
    }

    /// Reset for next iteration
    pub fn reset(&mut self) {
        self.active_buckets.clear();
        self.ready_queue.clear();
        self.in_flight
            .retain(|op| op.state == CommState::InProgress);
        self.completed.clear();

        // Reset compute time tracking for new iteration
        self.last_compute_start = None;
        self.current_iteration_compute = Duration::ZERO;
    }

    /// Get auto-tuner statistics (if auto-tuning is enabled)
    pub fn auto_tuner_stats(&self) -> Option<AutoTunerStats> {
        self.auto_tuner.as_ref().map(BucketAutoTuner::stats)
    }

    /// Enable or disable auto-tuning at runtime
    pub fn set_auto_tune(&mut self, enabled: bool) {
        if enabled && self.auto_tuner.is_none() {
            self.auto_tuner = Some(BucketAutoTuner::new(&self.config));
        } else if !enabled {
            self.auto_tuner = None;
        }
        self.config.auto_tune = enabled;
    }

    /// Get the effective bucket size being used
    pub fn effective_bucket_size(&self) -> usize {
        if let Some(ref tuner) = self.auto_tuner {
            tuner.stats().current_bucket_size
        } else {
            self.config.bucket_size_bytes
        }
    }

    /// Flush all remaining buckets (mark as ready)
    pub fn flush(&mut self) {
        for bucket in self.active_buckets.drain(..) {
            self.ready_queue.push_back(bucket);
        }
    }
}

// =============================================================================
// Overlap Scheduler
// =============================================================================

/// Schedules compute and communication for optimal overlap
pub struct OverlapScheduler {
    /// Configuration
    config: OverlapConfig,
    /// Bucket manager
    bucket_manager: BucketManager,
    /// Process group for communication
    process_group: ProcessGroup,
    /// P2P manager for direct transfers (optional)
    p2p_manager: Option<Arc<P2PManager>>,
    /// Hardware topology for optimization decisions
    topology: Option<HardwareTopology>,
    /// Current iteration
    iteration: usize,
    /// Scheduler statistics
    stats: SchedulerStats,
}

/// Statistics for the overlap scheduler
#[derive(Debug, Default, Clone)]
pub struct SchedulerStats {
    /// Total iterations
    pub iterations: usize,
    /// Compute time (excluding communication)
    pub total_compute_time: Duration,
    /// Communication time (overlapped portions counted once)
    pub total_comm_time: Duration,
    /// Overlap efficiency (0-1, higher is better)
    pub overlap_efficiency: f32,
    /// Buckets processed per iteration
    pub avg_buckets_per_iter: f32,
}

impl OverlapScheduler {
    /// Create a new overlap scheduler
    pub fn new(config: OverlapConfig, process_group: ProcessGroup) -> Self {
        let bucket_manager = BucketManager::new(config.clone());

        Self {
            config,
            bucket_manager,
            process_group,
            p2p_manager: None,
            topology: None,
            iteration: 0,
            stats: SchedulerStats::default(),
        }
    }

    /// Set P2P manager for direct GPU transfers
    pub fn with_p2p_manager(mut self, p2p_manager: Arc<P2PManager>) -> Self {
        self.p2p_manager = Some(p2p_manager);
        self
    }

    /// Set hardware topology for optimization decisions
    pub fn with_topology(mut self, topology: HardwareTopology) -> Self {
        self.topology = Some(topology);
        self
    }

    /// Register model parameters for gradient bucketing
    pub fn register_parameters(&mut self, params: &[(String, usize)]) {
        for (name, num_elements) in params {
            if *num_elements >= self.config.min_tensor_elements {
                self.bucket_manager.register_gradient(name, *num_elements);
            }
        }
    }

    /// Called when a gradient is computed during backward pass
    pub async fn on_gradient_computed(&mut self, name: &str, gradient: Tensor) -> Result<()> {
        // Add gradient to bucket
        let ready_bucket = self.bucket_manager.add_gradient(name, gradient)?;

        // If bucket became ready, try to start communication
        if ready_bucket.is_some() && self.bucket_manager.can_start_comm() {
            self.process_ready_buckets().await?;
        }

        Ok(())
    }

    /// Process all ready buckets (start communication)
    async fn process_ready_buckets(&mut self) -> Result<()> {
        while self.bucket_manager.can_start_comm() {
            if let Some(mut bucket) = self.bucket_manager.next_ready_bucket() {
                let bucket_id = bucket.id;
                let bytes_transferred = bucket.size_bytes;
                self.bucket_manager.start_communication(bucket_id);

                // Perform AllReduce on bucket gradients
                let result = self.allreduce_bucket(&mut bucket).await;

                self.bucket_manager.complete_communication(
                    bucket_id,
                    bytes_transferred,
                    result.is_ok(),
                    result.err().map(|e| e.to_string()),
                );
            }
        }

        Ok(())
    }

    /// AllReduce a bucket of gradients
    async fn allreduce_bucket(&self, bucket: &mut GradientBucket) -> Result<()> {
        for (_, gradient) in &mut bucket.gradients {
            self.process_group
                .all_reduce(gradient, ReduceOp::Sum)
                .await?;

            // Average the gradient
            let world_size = self.process_group.world_size() as f32;
            *gradient = gradient.div_scalar(world_size)?;
        }

        Ok(())
    }

    /// Called at the end of backward pass to flush remaining gradients
    pub async fn flush_and_wait(&mut self) -> Result<()> {
        // Flush any remaining buckets
        self.bucket_manager.flush();

        // Process all remaining buckets
        while !self.bucket_manager.ready_queue.is_empty()
            || self.bucket_manager.in_flight_count() > 0
        {
            self.process_ready_buckets().await?;
        }

        Ok(())
    }

    /// Start a new iteration
    pub fn start_iteration(&mut self) {
        self.iteration += 1;
        self.bucket_manager.reset();
        self.stats.iterations += 1;
    }

    /// Get scheduler statistics
    pub fn stats(&self) -> &SchedulerStats {
        &self.stats
    }

    /// Get bucket manager statistics
    pub fn bucket_stats(&self) -> &BucketStats {
        self.bucket_manager.stats()
    }

    /// Get auto-tuner statistics (if enabled)
    pub fn auto_tuner_stats(&self) -> Option<AutoTunerStats> {
        self.bucket_manager.auto_tuner_stats()
    }

    /// Enable or disable auto-tuning
    pub fn set_auto_tune(&mut self, enabled: bool) {
        self.bucket_manager.set_auto_tune(enabled);
    }

    /// Get the current effective bucket size
    pub fn effective_bucket_size(&self) -> usize {
        self.bucket_manager.effective_bucket_size()
    }
}

// =============================================================================
// Thread-Safe Wrapper
// =============================================================================

/// Thread-safe wrapper for OverlapScheduler
pub type SharedOverlapScheduler = Arc<RwLock<OverlapScheduler>>;

/// Create a shared overlap scheduler
pub fn shared_overlap_scheduler(
    config: OverlapConfig,
    process_group: ProcessGroup,
) -> SharedOverlapScheduler {
    Arc::new(RwLock::new(OverlapScheduler::new(config, process_group)))
}

// =============================================================================
// Double Buffer for Compute/Comm Overlap
// =============================================================================

/// Double buffer for gradient tensors to enable compute/comm overlap
pub struct DoubleBuffer {
    /// Buffer A
    buffer_a: Vec<Tensor>,
    /// Buffer B
    buffer_b: Vec<Tensor>,
    /// Currently active buffer (0 = A, 1 = B)
    active: usize,
    /// Buffer size
    capacity: usize,
}

impl DoubleBuffer {
    /// Create a new double buffer
    pub fn new(capacity: usize) -> Self {
        Self {
            buffer_a: Vec::with_capacity(capacity),
            buffer_b: Vec::with_capacity(capacity),
            active: 0,
            capacity,
        }
    }

    /// Get mutable reference to active buffer (for compute)
    pub fn active_buffer(&mut self) -> &mut Vec<Tensor> {
        if self.active == 0 {
            &mut self.buffer_a
        } else {
            &mut self.buffer_b
        }
    }

    /// Get reference to inactive buffer (for communication)
    pub fn inactive_buffer(&self) -> &Vec<Tensor> {
        if self.active == 0 {
            &self.buffer_b
        } else {
            &self.buffer_a
        }
    }

    /// Swap buffers
    pub fn swap(&mut self) {
        self.active = 1 - self.active;
    }

    /// Clear active buffer
    pub fn clear_active(&mut self) {
        self.active_buffer().clear();
    }
}

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

#[cfg(test)]
mod tests {
    use super::*;
    use crate::backend::{Backend, BackendConfig};
    use rtx_tensor::Shape;

    #[test]
    fn test_overlap_config_default() {
        let config = OverlapConfig::default();
        assert!(config.enabled);
        assert_eq!(config.bucket_size_bytes, 25 * 1024 * 1024);
        assert_eq!(config.max_concurrent_ops, 2);
        assert!(config.priority_scheduling);
    }

    #[test]
    fn test_gradient_bucket() {
        let mut bucket = GradientBucket::new(0, 3);

        assert_eq!(bucket.id, 0);
        assert!(!bucket.ready);
        assert_eq!(bucket.current_count, 0);

        let tensor =
            Tensor::zeros(Shape::new(vec![100]).unwrap(), &rtx_tensor::Device::Cpu).unwrap();
        bucket.add_gradient("grad1".to_string(), tensor);

        assert_eq!(bucket.current_count, 1);
        assert_eq!(bucket.num_elements, 100);
        assert!(!bucket.ready);
    }

    #[test]
    fn test_bucket_manager_creation() {
        let config = OverlapConfig::default();
        let manager = BucketManager::new(config);

        assert_eq!(manager.stats.buckets_created, 0);
        assert!(manager.ready_queue.is_empty());
    }

    #[test]
    fn test_bucket_manager_register() {
        let config = OverlapConfig::default();
        let mut manager = BucketManager::new(config);

        let bucket_id = manager.register_gradient("layer1.weight", 1000);
        assert_eq!(bucket_id, 0);
        assert_eq!(manager.stats.buckets_created, 1);

        // Same bucket should be reused if has capacity
        let bucket_id2 = manager.register_gradient("layer1.bias", 100);
        assert_eq!(bucket_id2, 0);
    }

    #[tokio::test]
    async fn test_overlap_scheduler_creation() {
        let config = OverlapConfig::default();
        let backend_config = BackendConfig::cpu();
        let pg = ProcessGroup::new_with_config(Backend::Cpu, 1, 0, backend_config)
            .await
            .unwrap();

        let scheduler = OverlapScheduler::new(config, pg);
        assert_eq!(scheduler.iteration, 0);
        assert_eq!(scheduler.stats.iterations, 0);
    }

    #[test]
    fn test_double_buffer() {
        let mut buffer = DoubleBuffer::new(10);

        // Add to active buffer
        let tensor = Tensor::zeros(Shape::new(vec![5]).unwrap(), &rtx_tensor::Device::Cpu).unwrap();
        buffer.active_buffer().push(tensor);

        assert_eq!(buffer.active_buffer().len(), 1);
        assert_eq!(buffer.inactive_buffer().len(), 0);

        // Swap
        buffer.swap();

        assert_eq!(buffer.active_buffer().len(), 0);
        assert_eq!(buffer.inactive_buffer().len(), 1);
    }

    #[test]
    fn test_bucket_stats_default() {
        let stats = BucketStats::default();
        assert_eq!(stats.buckets_created, 0);
        assert_eq!(stats.buckets_communicated, 0);
        assert_eq!(stats.bytes_communicated, 0);
    }

    #[test]
    fn test_comm_state() {
        let op = CommOperation {
            bucket_id: 0,
            state: CommState::Pending,
            start_time: None,
            end_time: None,
            error: None,
        };

        assert_eq!(op.state, CommState::Pending);
        assert!(op.start_time.is_none());
    }

    #[test]
    fn test_bucket_priority() {
        let mut bucket = GradientBucket::new(0, 1);
        bucket.priority = 100;

        assert_eq!(bucket.priority, 100);
    }

    #[test]
    fn test_bucket_age() {
        let bucket = GradientBucket::new(0, 1);
        std::thread::sleep(std::time::Duration::from_millis(10));

        assert!(bucket.age() >= Duration::from_millis(10));
    }
}