foxhunt/crates/ml-ensemble/src/cuda_streams.rs

//! CUDA stream pool for multi-stream ensemble inference.
//!
//! Provides parallel execution of multiple models on different CUDA streams.
//! Falls back to no-op on CPU devices.
//!
//! On CUDA devices, each stream can execute kernels concurrently with other streams,
//! enabling multiple model forward passes to overlap. This is the foundation for
//! `StreamAwareEnsemble` (Task 14) which assigns each model to a separate stream.

use ml_core::device::MlDevice;
use crate::MLError;

use std::sync::Arc;

#[cfg(feature = "cuda")]
type CudaStream = cudarc::driver::CudaStream;

#[cfg(feature = "cuda")]
type CudaEvent = cudarc::driver::CudaEvent;

/// Pool of CUDA streams for parallel model inference.
///
/// On CUDA devices, creates multiple streams via `CudaContext::new_stream()`.
/// On CPU, acts as a no-op (all methods succeed without doing anything).
///
/// # Usage
///
/// ```ignore
/// let pool = CudaStreamPool::new(&device, 4)?;
/// // Each model gets assigned stream index = model_index % pool.count()
/// // Sync all streams after inference:
/// pool.sync_all()?;
/// ```
pub struct CudaStreamPool {
    /// Number of streams (0 for CPU)
    count: usize,
    /// Whether this pool is on a CUDA device
    is_cuda: bool,
    /// The device this pool is associated with
    device: MlDevice,
    /// Forked CUDA streams (empty on CPU)
    #[cfg(feature = "cuda")]
    streams: Vec<Arc<CudaStream>>,
    #[cfg(not(feature = "cuda"))]
    streams: Vec<()>,
    /// Per-stream events for event-based sync (empty on CPU).
    /// Used by `sync_all()` to avoid N blocking `cuStreamSynchronize` calls.
    /// Instead: record event on each stream, have stream[0] wait on all events,
    /// then sync stream[0] once — 1 blocking call instead of N.
    #[cfg(feature = "cuda")]
    events: Vec<CudaEvent>,
}

impl std::fmt::Debug for CudaStreamPool {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("CudaStreamPool")
            .field("count", &self.count)
            .field("is_cuda", &self.is_cuda)
            .field("device", &self.device)
            .finish()
    }
}

impl CudaStreamPool {
    /// Create a new stream pool.
    ///
    /// On CUDA: creates `count` streams via the device context.
    /// On CPU: creates a no-op pool that passes through all operations.
    pub fn new(device: &MlDevice, count: usize) -> Result<Self, MLError> {
        #[cfg(feature = "cuda")]
        {
            if let MlDevice::Cuda { context, .. } = device {
                let mut streams = Vec::with_capacity(count);
                let mut events = Vec::with_capacity(count);
                for i in 0..count {
                    let stream = context.new_stream().map_err(|e| {
                        MLError::DeviceError(format!(
                            "Failed to create CUDA stream {i}/{count}: {e}"
                        ))
                    })?;
                    // Disable timing on sync events (pure synchronization, no profiling overhead).
                    let event = context.new_event(None).map_err(|e| {
                        MLError::DeviceError(format!(
                            "Failed to create CUDA event {i}/{count}: {e}"
                        ))
                    })?;
                    streams.push(stream);
                    events.push(event);
                }

                return Ok(Self {
                    count,
                    is_cuda: true,
                    device: device.clone(),
                    streams,
                    events,
                });
            }
        }

        let _ = count; // suppress unused warning on non-cuda

        // CPU device: no-op pool
        Ok(Self {
            count: 0,
            is_cuda: false,
            device: device.clone(),
            streams: Vec::new(),
            #[cfg(feature = "cuda")]
            events: Vec::new(),
        })
    }

    /// Number of streams in the pool.
    pub fn count(&self) -> usize {
        self.count
    }

    /// Whether this pool has CUDA streams.
    pub fn is_cuda(&self) -> bool {
        self.is_cuda
    }

    /// Get the device associated with this pool.
    pub fn device(&self) -> &MlDevice {
        &self.device
    }

    /// Get a CUDA stream by index.
    ///
    /// Returns `None` on CPU or if index is out of bounds.
    /// Typical usage: `pool.get_stream(model_index % pool.count())`.
    #[cfg(feature = "cuda")]
    pub fn get_stream(&self, index: usize) -> Option<&Arc<CudaStream>> {
        self.streams.get(index)
    }

    /// Get a CUDA stream by index (no-op on non-CUDA builds).
    #[cfg(not(feature = "cuda"))]
    pub fn get_stream(&self, _index: usize) -> Option<&()> {
        None
    }

    /// Synchronize all streams (wait for completion).
    ///
    /// On CPU: no-op.
    /// On CUDA: uses event-based synchronization to avoid N blocking host calls.
    /// Records an event on each stream, has stream[0] wait on all events, then
    /// synchronizes stream[0] — 1 blocking `cuStreamSynchronize` instead of N.
    pub fn sync_all(&self) -> Result<(), MLError> {
        #[cfg(feature = "cuda")]
        {
            if self.streams.len() <= 1 {
                // 0 or 1 streams: direct sync is already optimal.
                if let Some(s) = self.streams.first() {
                    s.synchronize().map_err(|e| {
                        MLError::DeviceError(format!("CUDA stream 0 sync failed: {e}"))
                    })?;
                }
                return Ok(());
            }

            // Record an event on each stream capturing all enqueued work.
            for (i, (stream, event)) in self.streams.iter().zip(self.events.iter()).enumerate() {
                event.record(stream).map_err(|e| {
                    MLError::DeviceError(format!(
                        "CUDA event record on stream {i} failed: {e}"
                    ))
                })?;
            }

            // Have stream[0] wait on events from streams[1..N].
            // stream[0]'s own event is already implicitly ordered.
            for (i, event) in self.events.iter().enumerate().skip(1) {
                if let Some(s0) = self.streams.first() {
                    s0.wait(event).map_err(|e| {
                        MLError::DeviceError(format!(
                            "CUDA stream 0 wait on event {i} failed: {e}"
                        ))
                    })?;
                }
            }

            // Single blocking sync on stream[0] — completes when all work is done.
            if let Some(s0) = self.streams.first() {
                s0.synchronize().map_err(|e| {
                    MLError::DeviceError(format!(
                        "CUDA stream 0 sync after event join failed: {e}"
                    ))
                })?;
            }
        }
        Ok(())
    }

    /// Synchronize a single stream by index.
    ///
    /// On CPU: no-op.
    /// On CUDA: synchronizes the specified stream.
    /// Returns `Ok(())` if index is out of bounds (no-op for safety).
    pub fn sync_stream(&self, index: usize) -> Result<(), MLError> {
        #[cfg(feature = "cuda")]
        if let Some(stream) = self.streams.get(index) {
            stream.synchronize().map_err(|e| {
                MLError::DeviceError(format!(
                    "CUDA stream {index} sync failed: {e}"
                ))
            })?;
        }
        let _ = index; // suppress unused on non-cuda
        Ok(())
    }
}

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

    #[test]
    fn test_stream_pool_cpu_noop() {
        let device = MlDevice::Cpu;
        let pool = CudaStreamPool::new(&device, 4)
            .expect("CPU pool should succeed");
        assert_eq!(pool.count(), 0);
        assert!(!pool.is_cuda());
        pool.sync_all().expect("CPU sync should succeed");
    }

    #[test]
    fn test_stream_pool_cpu_sync_stream_noop() {
        let device = MlDevice::Cpu;
        let pool = CudaStreamPool::new(&device, 4)
            .expect("CPU pool should succeed");
        pool.sync_stream(0).expect("sync_stream on CPU should succeed");
        pool.sync_stream(100).expect("out-of-bounds sync_stream should be noop");
    }

    #[test]
    fn test_stream_pool_cpu_zero_count() {
        let device = MlDevice::Cpu;
        let pool = CudaStreamPool::new(&device, 0)
            .expect("Zero-count CPU pool should succeed");
        assert_eq!(pool.count(), 0);
        pool.sync_all().expect("sync_all on empty pool should be noop");
    }

    #[test]
    fn test_stream_pool_cpu_no_streams() {
        let device = MlDevice::Cpu;
        let pool = CudaStreamPool::new(&device, 4)
            .expect("CPU pool should succeed");
        assert!(pool.get_stream(0).is_none());
        assert!(pool.get_stream(3).is_none());
        assert!(pool.get_stream(4).is_none());
    }
}