rustorch 0.6.29

//! GPU Synchronization Primitives
//! GPU同期プリミティブ
//!
//! Advanced synchronization mechanisms for multi-GPU operations including
//! barriers, events, streams, and cross-GPU synchronization.

use crate::error::{RusTorchError, RusTorchResult};
use std::collections::HashMap;
use std::sync::{Arc, Barrier, Condvar, Mutex};
use std::time::{Duration, Instant};

/// GPU event for synchronization
#[derive(Debug, Clone)]
pub struct GpuEvent {
    /// Event ID
    pub id: u64,
    /// GPU device ID
    pub device_id: usize,
    /// Creation timestamp
    pub created_at: Instant,
    /// Completion status
    pub completed: Arc<Mutex<bool>>,
    /// Completion notifier
    pub completion_notifier: Arc<(Mutex<bool>, Condvar)>,
}

/// GPU stream for asynchronous operations
#[derive(Debug)]
pub struct GpuStream {
    /// Stream ID
    pub id: u64,
    /// Device ID
    pub device_id: usize,
    /// Priority level
    pub priority: StreamPriority,
    /// Associated events
    pub events: Vec<GpuEvent>,
    /// Operation queue
    pub operation_queue: Arc<Mutex<Vec<StreamOperation>>>,
}

/// Stream priority levels
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum StreamPriority {
    Low = 0,
    Normal = 1,
    High = 2,
    Critical = 3,
}

/// Stream operations
#[derive(Debug, Clone)]
pub enum StreamOperation {
    /// Memory copy operation
    MemoryCopy {
        src_device: usize,
        dst_device: usize,
        size: usize,
    },
    /// Compute kernel execution
    KernelExecution {
        kernel_name: String,
        device_id: usize,
        grid_size: (u32, u32, u32),
        block_size: (u32, u32, u32),
    },
    /// Synchronization barrier
    Barrier { group_name: String },
    /// Event recording
    EventRecord { event_id: u64 },
    /// Wait for event
    EventWait { event_id: u64 },
}

/// Multi-GPU barrier for synchronization
pub struct MultiGpuBarrier {
    /// Number of participating GPUs
    num_gpus: usize,
    /// Per-GPU barriers
    gpu_barriers: HashMap<usize, Arc<Barrier>>,
    /// Global completion counter
    completion_counter: Arc<Mutex<usize>>,
    /// Completion condition variable
    completion_cv: Arc<Condvar>,
    /// Timeout duration
    timeout: Duration,
}

impl MultiGpuBarrier {
    /// Create new multi-GPU barrier
    pub fn new(gpu_ids: Vec<usize>, timeout: Duration) -> Self {
        let num_gpus = gpu_ids.len();
        let mut gpu_barriers = HashMap::new();

        for gpu_id in gpu_ids {
            gpu_barriers.insert(gpu_id, Arc::new(Barrier::new(1)));
        }

        Self {
            num_gpus,
            gpu_barriers,
            completion_counter: Arc::new(Mutex::new(0)),
            completion_cv: Arc::new(Condvar::new()),
            timeout,
        }
    }

    /// Wait for all GPUs to reach barrier
    pub fn wait(&self, gpu_id: usize) -> RusTorchResult<()> {
        let start_time = Instant::now();

        // Local GPU barrier
        if let Some(barrier) = self.gpu_barriers.get(&gpu_id) {
            barrier.wait();
        }

        // Global completion tracking
        {
            let mut counter = self.completion_counter.lock().unwrap();
            *counter += 1;

            if *counter >= self.num_gpus {
                // All GPUs reached barrier
                self.completion_cv.notify_all();
                *counter = 0; // Reset for next barrier
                return Ok(());
            }
        }

        // Wait for other GPUs with timeout
        let cv = &*self.completion_cv;
        let mut completed = self.completion_counter.lock().unwrap();

        loop {
            let elapsed = start_time.elapsed();
            if elapsed >= self.timeout {
                return Err(RusTorchError::gpu(format!(
                    "Multi-GPU barrier timeout after {:?}",
                    elapsed
                )));
            }

            let remaining = self.timeout - elapsed;
            let (_guard, timeout_result) = cv.wait_timeout(completed, remaining).unwrap();
            completed = self.completion_counter.lock().unwrap();

            if timeout_result.timed_out() {
                return Err(RusTorchError::gpu("Multi-GPU barrier wait timeout"));
            }

            // Check if all GPUs completed
            if *completed >= self.num_gpus {
                break;
            }
        }

        Ok(())
    }

    /// Reset barrier state
    pub fn reset(&self) {
        let mut counter = self.completion_counter.lock().unwrap();
        *counter = 0;
    }
}

/// Stream manager for coordinating GPU operations
pub struct StreamManager {
    /// Active streams per GPU
    streams: HashMap<usize, Vec<GpuStream>>,
    /// Event registry
    events: HashMap<u64, GpuEvent>,
    /// Next stream ID
    next_stream_id: Arc<Mutex<u64>>,
    /// Next event ID
    next_event_id: Arc<Mutex<u64>>,
}

impl StreamManager {
    /// Create new stream manager
    pub fn new() -> Self {
        Self {
            streams: HashMap::new(),
            events: HashMap::new(),
            next_stream_id: Arc::new(Mutex::new(0)),
            next_event_id: Arc::new(Mutex::new(0)),
        }
    }

    /// Create new stream on device
    pub fn create_stream(
        &mut self,
        device_id: usize,
        priority: StreamPriority,
    ) -> RusTorchResult<u64> {
        let mut stream_id_guard = self.next_stream_id.lock().unwrap();
        let stream_id = *stream_id_guard;
        *stream_id_guard += 1;
        drop(stream_id_guard);

        let stream = GpuStream {
            id: stream_id,
            device_id,
            priority,
            events: Vec::new(),
            operation_queue: Arc::new(Mutex::new(Vec::new())),
        };

        self.streams
            .entry(device_id)
            .or_insert_with(Vec::new)
            .push(stream);
        Ok(stream_id)
    }

    /// Create new event
    pub fn create_event(&mut self, device_id: usize) -> RusTorchResult<u64> {
        let mut event_id_guard = self.next_event_id.lock().unwrap();
        let event_id = *event_id_guard;
        *event_id_guard += 1;
        drop(event_id_guard);

        let event = GpuEvent {
            id: event_id,
            device_id,
            created_at: Instant::now(),
            completed: Arc::new(Mutex::new(false)),
            completion_notifier: Arc::new((Mutex::new(false), Condvar::new())),
        };

        self.events.insert(event_id, event);
        Ok(event_id)
    }

    /// Record event on stream
    pub fn record_event(&mut self, stream_id: u64, event_id: u64) -> RusTorchResult<()> {
        if let Some(event) = self.events.get(&event_id) {
            // Mark event as completed
            let mut completed = event.completed.lock().unwrap();
            *completed = true;

            // Notify waiters
            let (lock, cv) = &*event.completion_notifier;
            let mut notified = lock.lock().unwrap();
            *notified = true;
            cv.notify_all();
        }

        Ok(())
    }

    /// Wait for event completion
    pub fn wait_event(&self, event_id: u64, timeout: Option<Duration>) -> RusTorchResult<()> {
        if let Some(event) = self.events.get(&event_id) {
            let (lock, cv) = &*event.completion_notifier;
            let notified = lock.lock().unwrap();

            if let Some(timeout_duration) = timeout {
                let (_notified, timeout_result) =
                    cv.wait_timeout(notified, timeout_duration).unwrap();

                if timeout_result.timed_out() {
                    return Err(RusTorchError::gpu(format!(
                        "Event {} wait timeout",
                        event_id
                    )));
                }
            } else {
                let _notified = cv.wait(notified).unwrap();
            }
        }

        Ok(())
    }

    /// Synchronize all streams on device
    pub fn synchronize_device(&self, device_id: usize) -> RusTorchResult<()> {
        if let Some(streams) = self.streams.get(&device_id) {
            for stream in streams {
                // Wait for all operations in stream to complete
                let queue = stream.operation_queue.lock().unwrap();
                for operation in queue.iter() {
                    match operation {
                        StreamOperation::EventWait { event_id } => {
                            self.wait_event(*event_id, Some(Duration::from_secs(30)))?;
                        }
                        _ => {
                            // Other operations are assumed to be synchronous in this implementation
                        }
                    }
                }
            }
        }

        Ok(())
    }

    /// Query event status
    pub fn query_event(&self, event_id: u64) -> bool {
        if let Some(event) = self.events.get(&event_id) {
            let completed = event.completed.lock().unwrap();
            *completed
        } else {
            false
        }
    }
}

impl Default for StreamManager {
    fn default() -> Self {
        Self::new()
    }
}