122 lines
3.4 KiB
Rust
122 lines
3.4 KiB
Rust
//! Stream management and scheduling
|
|
//!
|
|
//! This module provides stream abstractions for asynchronous GPU execution
|
|
//! with dependency tracking and sub-microsecond scheduling overhead.
|
|
|
|
use crate::error::{Result, RuntimeError};
|
|
#[cfg(feature = "cuda")]
|
|
use cudarc::driver::{CudaContext, CudaStream};
|
|
use std::sync::Arc;
|
|
|
|
// Stub types when CUDA is not available
|
|
#[cfg(not(feature = "cuda"))]
|
|
pub struct CudaContext;
|
|
#[cfg(not(feature = "cuda"))]
|
|
pub struct CudaStream;
|
|
|
|
/// Stream identifier
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
|
|
pub struct StreamId(pub u32);
|
|
|
|
/// GPU stream for asynchronous execution
|
|
#[derive(Clone)]
|
|
pub struct Stream {
|
|
/// Stream ID
|
|
pub id: StreamId,
|
|
/// The underlying CUDA stream (if using CUDA backend)
|
|
cuda_stream: Option<Arc<CudaStream>>,
|
|
}
|
|
|
|
impl Stream {
|
|
/// Create a new stream with just an ID (for non-CUDA backends)
|
|
#[inline]
|
|
pub fn new(id: StreamId) -> Self {
|
|
Self {
|
|
id,
|
|
cuda_stream: None,
|
|
}
|
|
}
|
|
|
|
/// Create a new CUDA stream from a CudaContext
|
|
#[cfg(feature = "cuda")]
|
|
pub fn new_cuda(context: Arc<CudaContext>) -> Result<Self> {
|
|
// For simplicity, use the default stream
|
|
let cuda_stream = context.default_stream();
|
|
Ok(Self {
|
|
id: StreamId(0), // Default stream ID
|
|
cuda_stream: Some(cuda_stream),
|
|
})
|
|
}
|
|
|
|
#[cfg(not(feature = "cuda"))]
|
|
pub fn new_cuda(_context: Arc<CudaContext>) -> Result<Self> {
|
|
Err(RuntimeError::InvalidOperation(
|
|
"CUDA stream creation requires cuda feature".to_string(),
|
|
))
|
|
}
|
|
|
|
/// Create a Stream from an existing CudaStream
|
|
#[cfg(feature = "cuda")]
|
|
pub fn from_cuda_stream(stream: Arc<CudaStream>) -> Self {
|
|
Self {
|
|
id: StreamId(0), // Default ID for external streams
|
|
cuda_stream: Some(stream),
|
|
}
|
|
}
|
|
|
|
#[cfg(not(feature = "cuda"))]
|
|
pub fn from_cuda_stream(_stream: Arc<CudaStream>) -> Self {
|
|
Self {
|
|
id: StreamId(0),
|
|
cuda_stream: None,
|
|
}
|
|
}
|
|
|
|
/// Get the raw CUDA stream handle for kernel launching
|
|
#[cfg(feature = "cuda")]
|
|
pub fn raw_stream(&self) -> Result<&Arc<CudaStream>> {
|
|
self.cuda_stream
|
|
.as_ref()
|
|
.ok_or(RuntimeError::InvalidOperation(
|
|
"raw_stream: Stream does not have CUDA backend".to_string(),
|
|
))
|
|
}
|
|
|
|
#[cfg(not(feature = "cuda"))]
|
|
pub fn raw_stream(&self) -> Result<&Arc<CudaStream>> {
|
|
Err(RuntimeError::InvalidOperation(
|
|
"raw_stream requires cuda feature".to_string(),
|
|
))
|
|
}
|
|
|
|
/// Check if this is a CUDA stream
|
|
#[inline]
|
|
pub fn is_cuda(&self) -> bool {
|
|
self.cuda_stream.is_some()
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_stream_creation() {
|
|
let stream = Stream::new(StreamId(0));
|
|
assert_eq!(stream.id, StreamId(0));
|
|
assert!(!stream.is_cuda());
|
|
}
|
|
|
|
#[test]
|
|
#[cfg(feature = "cuda")]
|
|
fn test_cuda_stream_creation() {
|
|
// This test requires CUDA to be available
|
|
// CudaContext::new returns Arc<CudaContext> in cudarc
|
|
if let Ok(context) = CudaContext::new(0) {
|
|
let stream = Stream::new_cuda(context).expect("Failed to create CUDA stream");
|
|
assert!(stream.is_cuda());
|
|
assert!(stream.raw_stream().is_ok());
|
|
}
|
|
}
|
|
}
|