1398 lines
56 KiB
Rust
1398 lines
56 KiB
Rust
//! Device abstraction layer
|
|
//!
|
|
//! This module provides unified abstractions for GPU devices across different backends
|
|
//! (CUDA, ROCm, Metal). The design prioritizes performance, safety, and ease of use.
|
|
//!
|
|
//! # Architecture
|
|
//!
|
|
//! - **Device**: Main abstraction representing a GPU device
|
|
//! - **Stream**: Asynchronous execution stream for non-blocking operations
|
|
//! - **Event**: Synchronization primitive for inter-stream coordination
|
|
//! - **Context**: Device context managing resources and state
|
|
//!
|
|
//! # Safety
|
|
//!
|
|
//! All device operations are carefully designed to prevent resource leaks and ensure
|
|
//! proper cleanup even in error conditions.
|
|
|
|
use crate::error::{RuntimeError, Result};
|
|
use crate::allocator::{GpuAllocator, DevicePtr};
|
|
use parking_lot::{RwLock, Mutex};
|
|
use std::collections::{BTreeMap, HashMap};
|
|
use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
|
|
use std::sync::Arc;
|
|
use std::fmt;
|
|
use tracing::{debug, info, trace};
|
|
|
|
// CUDA integration - using simple pointer handles for now
|
|
// TODO: Use proper cudarc types once API is clarified
|
|
|
|
/// Device identifier
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
|
|
pub struct DeviceId(pub u32);
|
|
|
|
impl fmt::Display for DeviceId {
|
|
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
|
write!(f, "Device({})", self.0)
|
|
}
|
|
}
|
|
|
|
/// Stream identifier for asynchronous execution
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
|
|
pub struct StreamId(pub u32);
|
|
|
|
impl fmt::Display for StreamId {
|
|
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
|
write!(f, "Stream({})", self.0)
|
|
}
|
|
}
|
|
|
|
/// Event identifier for synchronization
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
|
|
pub struct EventId(pub u32);
|
|
|
|
impl fmt::Display for EventId {
|
|
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
|
write!(f, "Event({})", self.0)
|
|
}
|
|
}
|
|
|
|
/// Stream priority levels for critical path optimization
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
|
#[derive(Default)]
|
|
pub enum StreamPriority {
|
|
/// Highest priority for critical execution paths
|
|
High,
|
|
/// Normal priority for regular operations
|
|
#[default]
|
|
Normal,
|
|
/// Lowest priority for background tasks
|
|
Low,
|
|
}
|
|
|
|
|
|
/// GPU backend type
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
|
pub enum BackendType {
|
|
/// NVIDIA CUDA backend
|
|
Cuda,
|
|
/// AMD ROCm backend
|
|
Rocm,
|
|
/// Apple Metal backend
|
|
Metal,
|
|
/// CPU fallback backend
|
|
Cpu,
|
|
}
|
|
|
|
impl fmt::Display for BackendType {
|
|
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
|
match self {
|
|
BackendType::Cuda => write!(f, "CUDA"),
|
|
BackendType::Rocm => write!(f, "ROCm"),
|
|
BackendType::Metal => write!(f, "Metal"),
|
|
BackendType::Cpu => write!(f, "CPU"),
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Device properties and capabilities
|
|
#[derive(Debug, Clone)]
|
|
pub struct DeviceProperties {
|
|
/// Device name
|
|
pub name: String,
|
|
/// Backend type
|
|
pub backend: BackendType,
|
|
/// Total memory in bytes
|
|
pub total_memory: usize,
|
|
/// Available memory in bytes
|
|
pub available_memory: usize,
|
|
/// Compute capability major version
|
|
pub major: u32,
|
|
/// Compute capability minor version
|
|
pub minor: u32,
|
|
/// Number of multiprocessors
|
|
pub multiprocessor_count: u32,
|
|
/// Maximum threads per block
|
|
pub max_threads_per_block: u32,
|
|
/// Maximum shared memory per block
|
|
pub max_shared_memory_per_block: usize,
|
|
/// Supports unified memory
|
|
pub unified_memory: bool,
|
|
/// PCI bus ID
|
|
pub pci_bus_id: String,
|
|
}
|
|
|
|
/// Stream execution statistics
|
|
#[derive(Debug, Clone, Default)]
|
|
pub struct StreamStats {
|
|
/// Total kernels executed
|
|
pub kernels_executed: u64,
|
|
/// Total execution time in microseconds
|
|
pub total_execution_time_us: u64,
|
|
/// Number of synchronization events
|
|
pub sync_events: u64,
|
|
/// Number of memory operations
|
|
pub memory_operations: u64,
|
|
}
|
|
|
|
/// GPU stream for asynchronous execution
|
|
pub struct Stream {
|
|
/// Stream ID
|
|
pub id: StreamId,
|
|
/// Parent device
|
|
pub device_id: DeviceId,
|
|
/// Backend-specific handle (opaque)
|
|
handle: u64,
|
|
/// Stream statistics
|
|
stats: Arc<Mutex<StreamStats>>,
|
|
/// Active flag
|
|
active: AtomicU64,
|
|
/// Stream priority
|
|
priority: StreamPriority,
|
|
/// Stream callbacks
|
|
callbacks: Arc<Mutex<Vec<Box<dyn FnOnce() + Send + 'static>>>>,
|
|
}
|
|
|
|
/// Event for synchronization between streams
|
|
pub struct Event {
|
|
/// Event ID
|
|
pub id: EventId,
|
|
/// Parent device
|
|
pub device_id: DeviceId,
|
|
/// Backend-specific handle (opaque)
|
|
handle: u64,
|
|
/// Creation timestamp
|
|
created_at: AtomicU64,
|
|
/// Completion status
|
|
completed: AtomicU64,
|
|
}
|
|
|
|
/// Device context managing resources
|
|
pub struct DeviceContext {
|
|
/// Device ID
|
|
pub device_id: DeviceId,
|
|
/// Active streams
|
|
streams: RwLock<HashMap<StreamId, Arc<Stream>>>,
|
|
/// Active events
|
|
events: RwLock<HashMap<EventId, Arc<Event>>>,
|
|
/// Next stream ID
|
|
next_stream_id: AtomicU32,
|
|
/// Next event ID
|
|
next_event_id: AtomicU32,
|
|
/// Memory allocator for this device
|
|
allocator: Arc<GpuAllocator>,
|
|
/// Mock CUDA device initialization flag
|
|
cuda_device_initialized: AtomicU64, // Use as boolean (0 = false, 1 = true)
|
|
}
|
|
|
|
/// GPU device abstraction
|
|
pub struct Device {
|
|
/// Device ID
|
|
pub id: DeviceId,
|
|
/// Device properties
|
|
pub properties: DeviceProperties,
|
|
/// Device context
|
|
context: Arc<DeviceContext>,
|
|
}
|
|
|
|
impl Device {
|
|
/// Create a new device
|
|
pub fn new(id: DeviceId, properties: DeviceProperties) -> Result<Self> {
|
|
let context = Arc::new(DeviceContext {
|
|
device_id: id,
|
|
streams: RwLock::new(HashMap::new()),
|
|
events: RwLock::new(HashMap::new()),
|
|
next_stream_id: AtomicU32::new(0),
|
|
next_event_id: AtomicU32::new(0),
|
|
allocator: Arc::new(GpuAllocator::new()?),
|
|
cuda_device_initialized: AtomicU64::new(0),
|
|
});
|
|
|
|
Ok(Self {
|
|
id,
|
|
properties,
|
|
context,
|
|
})
|
|
}
|
|
|
|
/// Get device properties
|
|
pub fn properties(&self) -> &DeviceProperties {
|
|
&self.properties
|
|
}
|
|
|
|
/// Create a new stream on this device
|
|
pub fn create_stream(&self) -> Result<Arc<Stream>> {
|
|
self.create_priority_stream(StreamPriority::Normal)
|
|
}
|
|
|
|
/// Create a new stream with specified priority
|
|
pub fn create_priority_stream(&self, priority: StreamPriority) -> Result<Arc<Stream>> {
|
|
let stream_id = StreamId(self.context.next_stream_id.fetch_add(1, Ordering::SeqCst));
|
|
let handle = self.allocate_stream_handle_with_priority(priority)?;
|
|
|
|
let stream = Arc::new(Stream {
|
|
id: stream_id,
|
|
device_id: self.id,
|
|
handle,
|
|
stats: Arc::new(Mutex::new(StreamStats::default())),
|
|
active: AtomicU64::new(1),
|
|
priority,
|
|
callbacks: Arc::new(Mutex::new(Vec::new())),
|
|
});
|
|
|
|
self.context.streams.write().insert(stream_id, stream.clone());
|
|
info!("Created stream {} on device {} with priority {:?}", stream_id, self.id, priority);
|
|
Ok(stream)
|
|
}
|
|
|
|
/// Create a new event on this device
|
|
pub fn create_event(&self) -> Result<Arc<Event>> {
|
|
let event_id = EventId(self.context.next_event_id.fetch_add(1, Ordering::SeqCst));
|
|
let handle = self.allocate_event_handle()?;
|
|
|
|
let event = Arc::new(Event {
|
|
id: event_id,
|
|
device_id: self.id,
|
|
handle,
|
|
created_at: AtomicU64::new(current_time_us()),
|
|
completed: AtomicU64::new(0),
|
|
});
|
|
|
|
self.context.events.write().insert(event_id, event.clone());
|
|
info!("Created event {} on device {}", event_id, self.id);
|
|
Ok(event)
|
|
}
|
|
|
|
/// Allocate device memory
|
|
pub fn allocate(&self, size: usize) -> Result<DevicePtr> {
|
|
self.context.allocator.allocate(size)
|
|
}
|
|
|
|
/// Free device memory
|
|
pub fn free(&self, ptr: DevicePtr) -> Result<()> {
|
|
self.context.allocator.free(ptr)
|
|
}
|
|
|
|
/// Get memory statistics
|
|
pub fn memory_stats(&self) -> crate::allocator::AllocationStats {
|
|
self.context.allocator.stats()
|
|
}
|
|
|
|
/// Synchronize the device (wait for all operations to complete)
|
|
pub fn synchronize(&self) -> Result<()> {
|
|
debug!("Synchronizing device {}", self.id);
|
|
// In real implementation, this would call cuDeviceSynchronize() or equivalent
|
|
self.device_synchronize_impl()
|
|
}
|
|
|
|
/// Get the number of active streams
|
|
pub fn active_stream_count(&self) -> usize {
|
|
self.context.streams.read().len()
|
|
}
|
|
|
|
/// Get the number of active events
|
|
pub fn active_event_count(&self) -> usize {
|
|
self.context.events.read().len()
|
|
}
|
|
|
|
/// Clean up completed resources
|
|
pub fn cleanup(&self) -> Result<usize> {
|
|
let mut cleaned = 0;
|
|
|
|
// Clean up completed streams
|
|
let mut streams = self.context.streams.write();
|
|
streams.retain(|_, stream| {
|
|
if stream.active.load(Ordering::SeqCst) == 0 {
|
|
cleaned += 1;
|
|
false
|
|
} else {
|
|
true
|
|
}
|
|
});
|
|
|
|
// Clean up completed events
|
|
let mut events = self.context.events.write();
|
|
events.retain(|_, event| {
|
|
if event.completed.load(Ordering::SeqCst) > 0 {
|
|
cleaned += 1;
|
|
false
|
|
} else {
|
|
true
|
|
}
|
|
});
|
|
|
|
if cleaned > 0 {
|
|
debug!("Cleaned up {} resources on device {}", cleaned, self.id);
|
|
}
|
|
|
|
Ok(cleaned)
|
|
}
|
|
|
|
// Private implementation methods
|
|
|
|
/// CUDA device initialization
|
|
fn ensure_cuda_initialized(&self) -> Result<()> {
|
|
// Use compare_exchange to ensure only one thread initializes
|
|
if self.context.cuda_device_initialized.compare_exchange(
|
|
0, 1, Ordering::SeqCst, Ordering::SeqCst
|
|
).is_ok() {
|
|
debug!("Initializing CUDA device {}", self.id.0);
|
|
|
|
// Initialize CUDA device and set current context
|
|
#[cfg(feature = "cuda")]
|
|
{
|
|
use crate::cuda_backend;
|
|
unsafe {
|
|
cuda_backend::cuda_device_init(self.id.0 as i32)?;
|
|
cuda_backend::cuda_context_set_current(self.id.0 as i32)?;
|
|
}
|
|
}
|
|
#[cfg(not(feature = "cuda"))]
|
|
{
|
|
warn!("CUDA not available - using CPU fallback for device {}", self.id.0);
|
|
}
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
fn allocate_stream_handle(&self) -> Result<u64> {
|
|
self.allocate_stream_handle_with_priority(StreamPriority::Normal)
|
|
}
|
|
|
|
fn allocate_stream_handle_with_priority(&self, priority: StreamPriority) -> Result<u64> {
|
|
self.ensure_cuda_initialized()?;
|
|
debug!("Creating CUDA stream on device {} with priority {:?}", self.id, priority);
|
|
|
|
// Use real CUDA stream creation from cuda_backend
|
|
use crate::cuda_backend;
|
|
|
|
let cuda_stream = unsafe {
|
|
cuda_backend::cuda_stream_create_with_priority(self.id.0 as i32, priority)?
|
|
};
|
|
|
|
// Convert stream handle to u64 for storage
|
|
let handle = Box::into_raw(Box::new(cuda_stream)) as u64;
|
|
Ok(handle)
|
|
}
|
|
|
|
fn allocate_event_handle(&self) -> Result<u64> {
|
|
self.ensure_cuda_initialized()?;
|
|
debug!("Creating CUDA event on device {}", self.id);
|
|
|
|
// Use real CUDA event creation from cuda_backend
|
|
use crate::cuda_backend;
|
|
|
|
let cuda_event = unsafe {
|
|
cuda_backend::cuda_event_create(self.id.0 as i32)?
|
|
};
|
|
|
|
// Convert event handle to u64 for storage
|
|
let handle = Box::into_raw(Box::new(cuda_event)) as u64;
|
|
Ok(handle)
|
|
}
|
|
|
|
fn device_synchronize_impl(&self) -> Result<()> {
|
|
self.ensure_cuda_initialized()?;
|
|
debug!("Synchronizing CUDA device {}", self.id);
|
|
|
|
// Use real CUDA device synchronization
|
|
use crate::cuda_backend;
|
|
unsafe { cuda_backend::cuda_device_synchronize()? };
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Create a mock device for testing
|
|
pub fn mock_device(device_id: u32) -> Result<Self> {
|
|
let props = DeviceProperties {
|
|
name: "Mock GPU".to_string(),
|
|
backend: BackendType::Cuda,
|
|
total_memory: 8 * 1024 * 1024 * 1024,
|
|
available_memory: 7 * 1024 * 1024 * 1024,
|
|
major: 8,
|
|
minor: 0,
|
|
multiprocessor_count: 80,
|
|
max_threads_per_block: 1024,
|
|
max_shared_memory_per_block: 48 * 1024,
|
|
unified_memory: false,
|
|
pci_bus_id: format!("0000:0{}:00.0", device_id),
|
|
};
|
|
|
|
Self::new(DeviceId(device_id), props)
|
|
}
|
|
}
|
|
|
|
impl Stream {
|
|
/// Get stream statistics
|
|
pub fn stats(&self) -> StreamStats {
|
|
self.stats.lock().clone()
|
|
}
|
|
|
|
/// Get stream priority
|
|
pub fn priority(&self) -> StreamPriority {
|
|
self.priority
|
|
}
|
|
|
|
/// Add callback to be executed after stream completion
|
|
pub fn add_callback(&self, callback: Box<dyn FnOnce() + Send + 'static>) -> Result<()> {
|
|
self.callbacks.lock().push(callback);
|
|
Ok(())
|
|
}
|
|
|
|
/// Query stream completion status without blocking
|
|
pub fn query(&self) -> Result<bool> {
|
|
use crate::cuda_backend;
|
|
|
|
// Get CUDA stream handle
|
|
let cuda_stream = unsafe { &*(self.handle as *const crate::cuda_backend::CudaStreamHandle) };
|
|
|
|
// Query CUDA stream
|
|
unsafe { cuda_backend::cuda_stream_query(cuda_stream) }
|
|
}
|
|
|
|
/// CUDA stream operations
|
|
fn cuda_stream_op(&self, operation: &str) -> Result<()> {
|
|
trace!("{} operation on stream {}", operation, self.id);
|
|
#[cfg(feature = "cuda")]
|
|
{
|
|
// In a real implementation, this would perform the actual CUDA operation
|
|
// For now, just ensure the stream is valid
|
|
self.ensure_cuda_initialized()?;
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
/// Launch a kernel on this stream
|
|
pub fn launch_kernel(&self, kernel_name: &str, grid_size: (u32, u32, u32), block_size: (u32, u32, u32)) -> Result<()> {
|
|
trace!("Launching kernel '{}' on stream {} with grid {:?}, block {:?}",
|
|
kernel_name, self.id, grid_size, block_size);
|
|
|
|
// Update statistics
|
|
let mut stats = self.stats.lock();
|
|
stats.kernels_executed += 1;
|
|
|
|
#[cfg(feature = "cuda")]
|
|
{
|
|
use crate::cuda_backend;
|
|
// Launch kernel using CUDA backend
|
|
unsafe {
|
|
let stream = self.get_cuda_stream()?;
|
|
cuda_backend::cuda_kernel_launch(
|
|
kernel_name,
|
|
grid_size,
|
|
block_size,
|
|
stream
|
|
)?;
|
|
}
|
|
}
|
|
#[cfg(not(feature = "cuda"))]
|
|
{
|
|
trace!("CUDA not available - kernel '{}' execution simulated", kernel_name);
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Copy memory asynchronously on this stream
|
|
pub fn copy_async(&self, dst: DevicePtr, src: DevicePtr, size: usize) -> Result<()> {
|
|
trace!("Async copy {} bytes from {:?} to {:?} on stream {}",
|
|
size, src, dst, self.id);
|
|
|
|
// Validate pointers
|
|
if src.is_null() || dst.is_null() {
|
|
return Err(RuntimeError::device_error(
|
|
self.device_id.0,
|
|
"Cannot copy from/to null pointer".to_string()
|
|
));
|
|
}
|
|
|
|
use crate::cuda_backend;
|
|
|
|
// Get CUDA stream handle
|
|
let cuda_stream = unsafe { &*(self.handle as *const crate::cuda_backend::CudaStreamHandle) };
|
|
|
|
// Perform real CUDA async memory copy
|
|
unsafe {
|
|
cuda_backend::cuda_memcpy_device_to_device_async(dst, src, size, cuda_stream)?
|
|
};
|
|
|
|
let mut stats = self.stats.lock();
|
|
stats.memory_operations += 1;
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Record an event on this stream
|
|
pub fn record_event(&self, event: &Event) -> Result<()> {
|
|
if event.device_id != self.device_id {
|
|
return Err(RuntimeError::device_error(
|
|
self.device_id.0,
|
|
format!("Event {} belongs to device {}, cannot record on stream {} of device {}",
|
|
event.id, event.device_id, self.id, self.device_id)
|
|
));
|
|
}
|
|
|
|
trace!("Recording event {} on stream {}", event.id, self.id);
|
|
|
|
use crate::cuda_backend;
|
|
|
|
// Get handles
|
|
let cuda_stream = unsafe { &*(self.handle as *const crate::cuda_backend::CudaStreamHandle) };
|
|
let cuda_event = unsafe { &*(event.handle as *const crate::cuda_backend::CudaEventHandle) };
|
|
|
|
// Record event on CUDA stream
|
|
unsafe { cuda_backend::cuda_event_record(cuda_event, cuda_stream)? };
|
|
|
|
let mut stats = self.stats.lock();
|
|
stats.sync_events += 1;
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Wait for an event on this stream
|
|
pub fn wait_event(&self, event: &Event) -> Result<()> {
|
|
if event.device_id != self.device_id {
|
|
return Err(RuntimeError::device_error(
|
|
self.device_id.0,
|
|
format!("Event {} belongs to device {}, cannot wait on stream {} of device {}",
|
|
event.id, event.device_id, self.id, self.device_id)
|
|
));
|
|
}
|
|
|
|
trace!("Waiting for event {} on stream {}", event.id, self.id);
|
|
|
|
use crate::cuda_backend;
|
|
|
|
// Get handles
|
|
let cuda_stream = unsafe { &*(self.handle as *const crate::cuda_backend::CudaStreamHandle) };
|
|
let cuda_event = unsafe { &*(event.handle as *const crate::cuda_backend::CudaEventHandle) };
|
|
|
|
// Wait for event on CUDA stream
|
|
unsafe { cuda_backend::cuda_stream_wait_event(cuda_stream, cuda_event)? };
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Synchronize this stream (wait for completion)
|
|
pub fn synchronize(&self) -> Result<()> {
|
|
debug!("Synchronizing stream {}", self.id);
|
|
|
|
use crate::cuda_backend;
|
|
|
|
// Get CUDA stream handle
|
|
let cuda_stream = unsafe { &*(self.handle as *const crate::cuda_backend::CudaStreamHandle) };
|
|
|
|
// Synchronize CUDA stream
|
|
unsafe { cuda_backend::cuda_stream_synchronize(cuda_stream)? };
|
|
|
|
// Execute any pending callbacks
|
|
let mut callbacks = self.callbacks.lock();
|
|
for callback in callbacks.drain(..) {
|
|
callback();
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Check if stream is active
|
|
pub fn is_active(&self) -> bool {
|
|
self.active.load(Ordering::SeqCst) > 0
|
|
}
|
|
|
|
/// Mark stream as inactive (internal use)
|
|
pub fn deactivate(&self) {
|
|
self.active.store(0, Ordering::SeqCst);
|
|
}
|
|
|
|
/// Get the raw stream handle (for internal use by backends)
|
|
pub fn raw_handle(&self) -> u64 {
|
|
self.handle
|
|
}
|
|
|
|
/// Create a simple stream with default values (for testing/simplified usage)
|
|
pub fn new(id: StreamId) -> Self {
|
|
use std::sync::atomic::AtomicU64;
|
|
use parking_lot::Mutex;
|
|
|
|
Self {
|
|
id,
|
|
device_id: DeviceId(0), // Default device
|
|
handle: 0, // Default handle
|
|
stats: Arc::new(Mutex::new(StreamStats {
|
|
kernels_executed: 0,
|
|
total_execution_time_us: 0,
|
|
sync_events: 0,
|
|
memory_operations: 0,
|
|
})),
|
|
active: AtomicU64::new(0),
|
|
priority: StreamPriority::Normal,
|
|
callbacks: Arc::new(Mutex::new(Vec::new())),
|
|
}
|
|
}
|
|
}
|
|
|
|
// Manual Debug implementation for Stream
|
|
impl std::fmt::Debug for Stream {
|
|
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
|
use std::sync::atomic::Ordering;
|
|
|
|
f.debug_struct("Stream")
|
|
.field("id", &self.id)
|
|
.field("device_id", &self.device_id)
|
|
.field("handle", &self.handle)
|
|
.field("stats", &self.stats)
|
|
.field("active", &self.active.load(Ordering::SeqCst))
|
|
.field("priority", &self.priority)
|
|
.field("callbacks", &format!("{} callbacks", self.callbacks.lock().len()))
|
|
.finish()
|
|
}
|
|
}
|
|
|
|
impl Event {
|
|
/// Check if event has completed
|
|
pub fn is_completed(&self) -> bool {
|
|
// First check our cached completion status
|
|
if self.completed.load(Ordering::SeqCst) > 0 {
|
|
return true;
|
|
}
|
|
|
|
// Query real CUDA event completion
|
|
self.query().unwrap_or(false)
|
|
}
|
|
|
|
/// Query event completion status (non-blocking)
|
|
pub fn query(&self) -> Result<bool> {
|
|
use crate::cuda_backend;
|
|
|
|
// Get CUDA event handle
|
|
let cuda_event = unsafe { &*(self.handle as *const crate::cuda_backend::CudaEventHandle) };
|
|
|
|
// Query CUDA event
|
|
let is_completed = unsafe { cuda_backend::cuda_event_query(cuda_event)? };
|
|
|
|
if is_completed {
|
|
self.completed.store(current_time_us(), Ordering::SeqCst);
|
|
}
|
|
|
|
Ok(is_completed)
|
|
}
|
|
|
|
/// Wait for event completion (blocking)
|
|
pub fn wait(&self) -> Result<()> {
|
|
debug!("Waiting for event {} completion", self.id);
|
|
|
|
use crate::cuda_backend;
|
|
|
|
// Get CUDA event handle
|
|
let cuda_event = unsafe { &*(self.handle as *const crate::cuda_backend::CudaEventHandle) };
|
|
|
|
// Synchronize CUDA event (blocking)
|
|
unsafe { cuda_backend::cuda_event_synchronize(cuda_event)? };
|
|
|
|
// Mark as completed
|
|
self.completed.store(current_time_us(), Ordering::SeqCst);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Calculate elapsed time between events in milliseconds
|
|
pub fn elapsed_time_ms(&self, start_event: &Event) -> Result<f32> {
|
|
use crate::cuda_backend;
|
|
|
|
// Get CUDA event handles
|
|
let start_cuda_event = unsafe { &*(start_event.handle as *const crate::cuda_backend::CudaEventHandle) };
|
|
let end_cuda_event = unsafe { &*(self.handle as *const crate::cuda_backend::CudaEventHandle) };
|
|
|
|
// Calculate elapsed time using CUDA events
|
|
unsafe { cuda_backend::cuda_event_elapsed_time(start_cuda_event, end_cuda_event) }
|
|
}
|
|
|
|
/// Get elapsed time since event creation
|
|
pub fn elapsed_time_us(&self) -> u64 {
|
|
let current = current_time_us();
|
|
let created = self.created_at.load(Ordering::SeqCst);
|
|
current.saturating_sub(created)
|
|
}
|
|
}
|
|
|
|
/// Get current time in microseconds (mock implementation)
|
|
fn current_time_us() -> u64 {
|
|
use std::time::{SystemTime, UNIX_EPOCH};
|
|
SystemTime::now()
|
|
.duration_since(UNIX_EPOCH)
|
|
.unwrap()
|
|
.as_micros() as u64
|
|
}
|
|
|
|
/// Discover available devices
|
|
pub fn discover_devices(devices: &mut BTreeMap<DeviceId, Device>) -> Result<usize> {
|
|
// Use real CUDA device discovery
|
|
crate::cuda_backend::discover_cuda_devices(devices)
|
|
.map_err(|e| RuntimeError::device_error(0, format!("CUDA device discovery failed: {}", e)))
|
|
}
|
|
|
|
// Comprehensive failing tests (TDD approach)
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use std::thread;
|
|
use std::time::Duration;
|
|
|
|
// Basic device tests
|
|
|
|
#[test]
|
|
fn test_device_creation() {
|
|
let props = DeviceProperties {
|
|
name: "Test GPU".to_string(),
|
|
backend: BackendType::Cuda,
|
|
total_memory: 8 * 1024 * 1024 * 1024,
|
|
available_memory: 7 * 1024 * 1024 * 1024,
|
|
major: 8,
|
|
minor: 0,
|
|
multiprocessor_count: 80,
|
|
max_threads_per_block: 1024,
|
|
max_shared_memory_per_block: 48 * 1024,
|
|
unified_memory: false,
|
|
pci_bus_id: "0000:01:00.0".to_string(),
|
|
};
|
|
|
|
let device = Device::new(DeviceId(0), props).expect("Device creation should succeed");
|
|
assert_eq!(device.id, DeviceId(0));
|
|
assert_eq!(device.properties.name, "Test GPU");
|
|
assert_eq!(device.properties.backend, BackendType::Cuda);
|
|
assert_eq!(device.active_stream_count(), 0);
|
|
assert_eq!(device.active_event_count(), 0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_device_discovery() {
|
|
let mut devices = BTreeMap::new();
|
|
let count = discover_devices(&mut devices).expect("Device discovery should succeed");
|
|
assert_eq!(count, 2);
|
|
assert!(devices.contains_key(&DeviceId(0)));
|
|
assert!(devices.contains_key(&DeviceId(1)));
|
|
|
|
let device_0 = devices.get(&DeviceId(0)).unwrap();
|
|
assert_eq!(device_0.properties.name, "RTX 5090");
|
|
assert_eq!(device_0.properties.backend, BackendType::Cuda);
|
|
}
|
|
|
|
// Stream tests (should fail initially due to unimplemented methods)
|
|
|
|
#[test]
|
|
fn test_stream_creation() {
|
|
let props = DeviceProperties {
|
|
name: "Test GPU".to_string(),
|
|
backend: BackendType::Cuda,
|
|
total_memory: 8 * 1024 * 1024 * 1024,
|
|
available_memory: 7 * 1024 * 1024 * 1024,
|
|
major: 8,
|
|
minor: 0,
|
|
multiprocessor_count: 80,
|
|
max_threads_per_block: 1024,
|
|
max_shared_memory_per_block: 48 * 1024,
|
|
unified_memory: false,
|
|
pci_bus_id: "0000:01:00.0".to_string(),
|
|
};
|
|
|
|
let device = Device::new(DeviceId(0), props).unwrap();
|
|
let stream = device.create_stream().expect("Stream creation should succeed");
|
|
|
|
assert_eq!(stream.device_id, DeviceId(0));
|
|
assert_eq!(stream.id, StreamId(0));
|
|
assert!(stream.is_active());
|
|
assert_eq!(device.active_stream_count(), 1);
|
|
}
|
|
|
|
#[test]
|
|
fn test_multiple_streams() {
|
|
let props = DeviceProperties {
|
|
name: "Test GPU".to_string(),
|
|
backend: BackendType::Cuda,
|
|
total_memory: 8 * 1024 * 1024 * 1024,
|
|
available_memory: 7 * 1024 * 1024 * 1024,
|
|
major: 8,
|
|
minor: 0,
|
|
multiprocessor_count: 80,
|
|
max_threads_per_block: 1024,
|
|
max_shared_memory_per_block: 48 * 1024,
|
|
unified_memory: false,
|
|
pci_bus_id: "0000:01:00.0".to_string(),
|
|
};
|
|
|
|
let device = Device::new(DeviceId(0), props).unwrap();
|
|
|
|
let stream_1 = device.create_stream().unwrap();
|
|
let stream_2 = device.create_stream().unwrap();
|
|
let stream_3 = device.create_stream().unwrap();
|
|
|
|
assert_eq!(device.active_stream_count(), 3);
|
|
assert_eq!(stream_1.id, StreamId(0));
|
|
assert_eq!(stream_2.id, StreamId(1));
|
|
assert_eq!(stream_3.id, StreamId(2));
|
|
|
|
// All streams should belong to the same device
|
|
assert_eq!(stream_1.device_id, DeviceId(0));
|
|
assert_eq!(stream_2.device_id, DeviceId(0));
|
|
assert_eq!(stream_3.device_id, DeviceId(0));
|
|
}
|
|
|
|
#[test]
|
|
fn test_stream_kernel_launch() {
|
|
let props = DeviceProperties {
|
|
name: "Test GPU".to_string(),
|
|
backend: BackendType::Cuda,
|
|
total_memory: 8 * 1024 * 1024 * 1024,
|
|
available_memory: 7 * 1024 * 1024 * 1024,
|
|
major: 8,
|
|
minor: 0,
|
|
multiprocessor_count: 80,
|
|
max_threads_per_block: 1024,
|
|
max_shared_memory_per_block: 48 * 1024,
|
|
unified_memory: false,
|
|
pci_bus_id: "0000:01:00.0".to_string(),
|
|
};
|
|
|
|
let device = Device::new(DeviceId(0), props).unwrap();
|
|
let stream = device.create_stream().unwrap();
|
|
|
|
stream.launch_kernel("test_kernel", (1, 1, 1), (256, 1, 1))
|
|
.expect("Kernel launch should succeed");
|
|
|
|
let stats = stream.stats();
|
|
assert_eq!(stats.kernels_executed, 1);
|
|
}
|
|
|
|
// Event tests (should fail initially)
|
|
|
|
#[test]
|
|
fn test_event_creation() {
|
|
let props = DeviceProperties {
|
|
name: "Test GPU".to_string(),
|
|
backend: BackendType::Cuda,
|
|
total_memory: 8 * 1024 * 1024 * 1024,
|
|
available_memory: 7 * 1024 * 1024 * 1024,
|
|
major: 8,
|
|
minor: 0,
|
|
multiprocessor_count: 80,
|
|
max_threads_per_block: 1024,
|
|
max_shared_memory_per_block: 48 * 1024,
|
|
unified_memory: false,
|
|
pci_bus_id: "0000:01:00.0".to_string(),
|
|
};
|
|
|
|
let device = Device::new(DeviceId(0), props).unwrap();
|
|
let event = device.create_event().expect("Event creation should succeed");
|
|
|
|
assert_eq!(event.device_id, DeviceId(0));
|
|
assert_eq!(event.id, EventId(0));
|
|
assert!(!event.is_completed());
|
|
assert_eq!(device.active_event_count(), 1);
|
|
}
|
|
|
|
#[test]
|
|
fn test_stream_event_interaction() {
|
|
let props = DeviceProperties {
|
|
name: "Test GPU".to_string(),
|
|
backend: BackendType::Cuda,
|
|
total_memory: 8 * 1024 * 1024 * 1024,
|
|
available_memory: 7 * 1024 * 1024 * 1024,
|
|
major: 8,
|
|
minor: 0,
|
|
multiprocessor_count: 80,
|
|
max_threads_per_block: 1024,
|
|
max_shared_memory_per_block: 48 * 1024,
|
|
unified_memory: false,
|
|
pci_bus_id: "0000:01:00.0".to_string(),
|
|
};
|
|
|
|
let device = Device::new(DeviceId(0), props).unwrap();
|
|
let stream = device.create_stream().unwrap();
|
|
let event = device.create_event().unwrap();
|
|
|
|
// Record event on stream
|
|
stream.record_event(&event).expect("Event recording should succeed");
|
|
|
|
// Wait for event on another stream
|
|
let stream2 = device.create_stream().unwrap();
|
|
stream2.wait_event(&event).expect("Event waiting should succeed");
|
|
}
|
|
|
|
#[test]
|
|
fn test_cross_device_event_error() {
|
|
let props1 = DeviceProperties {
|
|
name: "GPU 1".to_string(),
|
|
backend: BackendType::Cuda,
|
|
total_memory: 8 * 1024 * 1024 * 1024,
|
|
available_memory: 7 * 1024 * 1024 * 1024,
|
|
major: 8,
|
|
minor: 0,
|
|
multiprocessor_count: 80,
|
|
max_threads_per_block: 1024,
|
|
max_shared_memory_per_block: 48 * 1024,
|
|
unified_memory: false,
|
|
pci_bus_id: "0000:01:00.0".to_string(),
|
|
};
|
|
|
|
let props2 = DeviceProperties {
|
|
name: "GPU 2".to_string(),
|
|
backend: BackendType::Cuda,
|
|
total_memory: 8 * 1024 * 1024 * 1024,
|
|
available_memory: 7 * 1024 * 1024 * 1024,
|
|
major: 8,
|
|
minor: 0,
|
|
multiprocessor_count: 80,
|
|
max_threads_per_block: 1024,
|
|
max_shared_memory_per_block: 48 * 1024,
|
|
unified_memory: false,
|
|
pci_bus_id: "0000:02:00.0".to_string(),
|
|
};
|
|
|
|
let device1 = Device::new(DeviceId(0), props1).unwrap();
|
|
let device2 = Device::new(DeviceId(1), props2).unwrap();
|
|
|
|
// This should work without panicking
|
|
let _stream1 = device1.create_stream();
|
|
let _event2 = device2.create_event();
|
|
|
|
// Cross-device operations should return errors, not panic
|
|
// We can't test this fully until stream creation is implemented
|
|
}
|
|
|
|
// Memory allocation tests
|
|
|
|
#[test]
|
|
fn test_device_memory_allocation() {
|
|
let props = DeviceProperties {
|
|
name: "Test GPU".to_string(),
|
|
backend: BackendType::Cuda,
|
|
total_memory: 8 * 1024 * 1024 * 1024,
|
|
available_memory: 7 * 1024 * 1024 * 1024,
|
|
major: 8,
|
|
minor: 0,
|
|
multiprocessor_count: 80,
|
|
max_threads_per_block: 1024,
|
|
max_shared_memory_per_block: 48 * 1024,
|
|
unified_memory: false,
|
|
pci_bus_id: "0000:01:00.0".to_string(),
|
|
};
|
|
|
|
let device = Device::new(DeviceId(0), props).unwrap();
|
|
|
|
// Test allocation
|
|
let ptr = device.allocate(1024).expect("Memory allocation should succeed");
|
|
assert!(!ptr.is_null());
|
|
|
|
// Test stats tracking
|
|
let stats = device.memory_stats();
|
|
assert!(stats.active_allocations > 0);
|
|
assert!(stats.total_allocated > 0);
|
|
|
|
// Test deallocation
|
|
device.free(ptr).expect("Memory deallocation should succeed");
|
|
}
|
|
|
|
#[test]
|
|
fn test_device_synchronization() {
|
|
let props = DeviceProperties {
|
|
name: "Test GPU".to_string(),
|
|
backend: BackendType::Cuda,
|
|
total_memory: 8 * 1024 * 1024 * 1024,
|
|
available_memory: 7 * 1024 * 1024 * 1024,
|
|
major: 8,
|
|
minor: 0,
|
|
multiprocessor_count: 80,
|
|
max_threads_per_block: 1024,
|
|
max_shared_memory_per_block: 48 * 1024,
|
|
unified_memory: false,
|
|
pci_bus_id: "0000:01:00.0".to_string(),
|
|
};
|
|
|
|
let device = Device::new(DeviceId(0), props).unwrap();
|
|
device.synchronize().expect("Device synchronization should succeed");
|
|
}
|
|
|
|
// Cleanup and resource management tests
|
|
|
|
#[test]
|
|
fn test_resource_cleanup() {
|
|
let props = DeviceProperties {
|
|
name: "Test GPU".to_string(),
|
|
backend: BackendType::Cuda,
|
|
total_memory: 8 * 1024 * 1024 * 1024,
|
|
available_memory: 7 * 1024 * 1024 * 1024,
|
|
major: 8,
|
|
minor: 0,
|
|
multiprocessor_count: 80,
|
|
max_threads_per_block: 1024,
|
|
max_shared_memory_per_block: 48 * 1024,
|
|
unified_memory: false,
|
|
pci_bus_id: "0000:01:00.0".to_string(),
|
|
};
|
|
|
|
let device = Device::new(DeviceId(0), props).unwrap();
|
|
|
|
// Create some resources
|
|
let stream = device.create_stream().unwrap();
|
|
let event = device.create_event().unwrap();
|
|
|
|
assert_eq!(device.active_stream_count(), 1);
|
|
assert_eq!(device.active_event_count(), 1);
|
|
|
|
// Deactivate resources
|
|
stream.deactivate();
|
|
|
|
// Cleanup should remove inactive resources
|
|
let cleaned = device.cleanup().expect("Cleanup should succeed");
|
|
assert!(cleaned > 0);
|
|
}
|
|
|
|
// Performance and stress tests
|
|
|
|
#[test]
|
|
fn test_many_streams() {
|
|
let props = DeviceProperties {
|
|
name: "Test GPU".to_string(),
|
|
backend: BackendType::Cuda,
|
|
total_memory: 8 * 1024 * 1024 * 1024,
|
|
available_memory: 7 * 1024 * 1024 * 1024,
|
|
major: 8,
|
|
minor: 0,
|
|
multiprocessor_count: 80,
|
|
max_threads_per_block: 1024,
|
|
max_shared_memory_per_block: 48 * 1024,
|
|
unified_memory: false,
|
|
pci_bus_id: "0000:01:00.0".to_string(),
|
|
};
|
|
|
|
let device = Device::new(DeviceId(0), props).unwrap();
|
|
let mut streams = Vec::new();
|
|
|
|
// Create many streams
|
|
for _ in 0..100 {
|
|
streams.push(device.create_stream().unwrap());
|
|
}
|
|
|
|
assert_eq!(device.active_stream_count(), 100);
|
|
|
|
// Launch kernels on all streams
|
|
for stream in &streams {
|
|
stream.launch_kernel("test_kernel", (1, 1, 1), (256, 1, 1)).unwrap();
|
|
}
|
|
|
|
// Synchronize all streams
|
|
for stream in &streams {
|
|
stream.synchronize().unwrap();
|
|
}
|
|
}
|
|
|
|
// TDD tests for real CUDA streams (should fail initially)
|
|
|
|
#[test]
|
|
fn test_real_cuda_stream_creation() {
|
|
let mut devices = BTreeMap::new();
|
|
discover_devices(&mut devices).expect("Should discover CUDA devices");
|
|
|
|
let device = devices.values().next().expect("Should have at least one device");
|
|
let stream = device.create_stream().expect("Real CUDA stream creation should succeed");
|
|
|
|
assert!(stream.is_active());
|
|
assert_eq!(stream.device_id, device.id);
|
|
|
|
// Stream should have real CUDA handle (not mock)
|
|
assert_ne!(stream.handle, 0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_real_cuda_stream_synchronization() {
|
|
let mut devices = BTreeMap::new();
|
|
discover_devices(&mut devices).expect("Should discover CUDA devices");
|
|
|
|
let device = devices.values().next().expect("Should have at least one device");
|
|
let stream = device.create_stream().expect("Stream creation should succeed");
|
|
|
|
// Real synchronization should work
|
|
stream.synchronize().expect("Real CUDA stream synchronization should succeed");
|
|
}
|
|
|
|
#[test]
|
|
fn test_real_cuda_stream_priority() {
|
|
let mut devices = BTreeMap::new();
|
|
discover_devices(&mut devices).expect("Should discover CUDA devices");
|
|
|
|
let device = devices.values().next().expect("Should have at least one device");
|
|
|
|
// Create high priority stream for critical paths
|
|
let high_priority_stream = device.create_priority_stream(StreamPriority::High)
|
|
.expect("High priority stream creation should succeed");
|
|
|
|
// Create low priority stream for background tasks
|
|
let low_priority_stream = device.create_priority_stream(StreamPriority::Low)
|
|
.expect("Low priority stream creation should succeed");
|
|
|
|
// Priority streams should have different handles
|
|
assert_ne!(high_priority_stream.handle, low_priority_stream.handle);
|
|
assert_eq!(high_priority_stream.priority(), StreamPriority::High);
|
|
assert_eq!(low_priority_stream.priority(), StreamPriority::Low);
|
|
}
|
|
|
|
#[test]
|
|
fn test_real_cuda_stream_callbacks() {
|
|
use std::sync::atomic::{AtomicBool, Ordering};
|
|
use std::sync::Arc;
|
|
|
|
let mut devices = BTreeMap::new();
|
|
discover_devices(&mut devices).expect("Should discover CUDA devices");
|
|
|
|
let device = devices.values().next().expect("Should have at least one device");
|
|
let stream = device.create_stream().expect("Stream creation should succeed");
|
|
|
|
// Test callback execution
|
|
let callback_executed = Arc::new(AtomicBool::new(false));
|
|
let callback_flag = callback_executed.clone();
|
|
|
|
stream.add_callback(Box::new(move || {
|
|
callback_flag.store(true, Ordering::SeqCst);
|
|
})).expect("Adding callback should succeed");
|
|
|
|
// Synchronize to ensure callback is executed
|
|
stream.synchronize().expect("Stream sync should succeed");
|
|
|
|
// Callback should have been executed
|
|
assert!(callback_executed.load(Ordering::SeqCst));
|
|
}
|
|
|
|
#[test]
|
|
fn test_real_cuda_memory_copy_async() {
|
|
let mut devices = BTreeMap::new();
|
|
discover_devices(&mut devices).expect("Should discover CUDA devices");
|
|
|
|
let device = devices.values().next().expect("Should have at least one device");
|
|
let stream = device.create_stream().expect("Stream creation should succeed");
|
|
|
|
// Allocate source and destination memory
|
|
let src_ptr = device.allocate(1024).expect("Source allocation should succeed");
|
|
let dst_ptr = device.allocate(1024).expect("Destination allocation should succeed");
|
|
|
|
// Real async copy should work
|
|
stream.copy_async(dst_ptr, src_ptr, 1024)
|
|
.expect("Real async copy should succeed");
|
|
|
|
// Synchronize to ensure copy completes
|
|
stream.synchronize().expect("Stream sync should succeed");
|
|
|
|
// Clean up
|
|
device.free(src_ptr).expect("Source deallocation should succeed");
|
|
device.free(dst_ptr).expect("Destination deallocation should succeed");
|
|
}
|
|
|
|
#[test]
|
|
fn test_real_cuda_stream_query_status() {
|
|
let mut devices = BTreeMap::new();
|
|
discover_devices(&mut devices).expect("Should discover CUDA devices");
|
|
|
|
let device = devices.values().next().expect("Should have at least one device");
|
|
let stream = device.create_stream().expect("Stream creation should succeed");
|
|
|
|
// Query should work on real stream
|
|
let is_ready = stream.query().expect("Stream query should succeed");
|
|
|
|
// Empty stream should be ready
|
|
assert!(is_ready);
|
|
}
|
|
|
|
// TDD tests for real CUDA events (should fail initially)
|
|
|
|
#[test]
|
|
fn test_real_cuda_event_creation() {
|
|
let mut devices = BTreeMap::new();
|
|
discover_devices(&mut devices).expect("Should discover CUDA devices");
|
|
|
|
let device = devices.values().next().expect("Should have at least one device");
|
|
let event = device.create_event().expect("Real CUDA event creation should succeed");
|
|
|
|
assert_eq!(event.device_id, device.id);
|
|
// Event should have real CUDA handle (not mock)
|
|
assert_ne!(event.handle, 0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_real_cuda_event_timing() {
|
|
let mut devices = BTreeMap::new();
|
|
discover_devices(&mut devices).expect("Should discover CUDA devices");
|
|
|
|
let device = devices.values().next().expect("Should have at least one device");
|
|
|
|
// Create events for timing
|
|
let start_event = device.create_event().expect("Start event creation should succeed");
|
|
let end_event = device.create_event().expect("End event creation should succeed");
|
|
|
|
let stream = device.create_stream().expect("Stream creation should succeed");
|
|
|
|
// Record start event
|
|
stream.record_event(&start_event).expect("Recording start event should succeed");
|
|
|
|
// Simulate some work (memory allocation/deallocation)
|
|
let ptr = device.allocate(1024 * 1024).expect("Allocation should succeed");
|
|
device.free(ptr).expect("Deallocation should succeed");
|
|
|
|
// Record end event
|
|
stream.record_event(&end_event).expect("Recording end event should succeed");
|
|
|
|
// Wait for completion
|
|
stream.synchronize().expect("Stream sync should succeed");
|
|
|
|
// Calculate elapsed time
|
|
let elapsed_ms = end_event.elapsed_time_ms(&start_event)
|
|
.expect("Elapsed time calculation should succeed");
|
|
|
|
// Should have some measurable time
|
|
assert!(elapsed_ms >= 0.0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_real_cuda_event_record_and_wait() {
|
|
let mut devices = BTreeMap::new();
|
|
discover_devices(&mut devices).expect("Should discover CUDA devices");
|
|
|
|
let device = devices.values().next().expect("Should have at least one device");
|
|
|
|
let stream1 = device.create_stream().expect("Stream1 creation should succeed");
|
|
let stream2 = device.create_stream().expect("Stream2 creation should succeed");
|
|
let event = device.create_event().expect("Event creation should succeed");
|
|
|
|
// Record event on stream1
|
|
stream1.record_event(&event).expect("Recording event should succeed");
|
|
|
|
// Wait for event on stream2 (inter-stream dependency)
|
|
stream2.wait_event(&event).expect("Waiting for event should succeed");
|
|
|
|
// Both streams should synchronize without error
|
|
stream1.synchronize().expect("Stream1 sync should succeed");
|
|
stream2.synchronize().expect("Stream2 sync should succeed");
|
|
}
|
|
|
|
#[test]
|
|
fn test_real_cuda_event_query() {
|
|
let mut devices = BTreeMap::new();
|
|
discover_devices(&mut devices).expect("Should discover CUDA devices");
|
|
|
|
let device = devices.values().next().expect("Should have at least one device");
|
|
let event = device.create_event().expect("Event creation should succeed");
|
|
let stream = device.create_stream().expect("Stream creation should succeed");
|
|
|
|
// Record event on stream
|
|
stream.record_event(&event).expect("Recording event should succeed");
|
|
|
|
// Query event status - might be completed or not
|
|
let is_completed = event.query().expect("Event query should succeed");
|
|
|
|
// After synchronization, event should definitely be completed
|
|
stream.synchronize().expect("Stream sync should succeed");
|
|
let is_completed_after_sync = event.query().expect("Event query after sync should succeed");
|
|
assert!(is_completed_after_sync);
|
|
}
|
|
|
|
// TDD tests for CUDA device synchronization
|
|
|
|
#[test]
|
|
fn test_real_cuda_device_synchronize() {
|
|
let mut devices = BTreeMap::new();
|
|
discover_devices(&mut devices).expect("Should discover CUDA devices");
|
|
|
|
let device = devices.values().next().expect("Should have at least one device");
|
|
|
|
// Create multiple streams with work
|
|
let stream1 = device.create_stream().expect("Stream1 creation should succeed");
|
|
let stream2 = device.create_stream().expect("Stream2 creation should succeed");
|
|
|
|
// Launch work on both streams
|
|
let ptr1 = device.allocate(1024).expect("Allocation1 should succeed");
|
|
let ptr2 = device.allocate(1024).expect("Allocation2 should succeed");
|
|
|
|
stream1.copy_async(ptr1, ptr1, 1024).expect("Copy on stream1 should succeed");
|
|
stream2.copy_async(ptr2, ptr2, 1024).expect("Copy on stream2 should succeed");
|
|
|
|
// Device sync should wait for all streams
|
|
device.synchronize().expect("Real CUDA device synchronization should succeed");
|
|
|
|
// Clean up
|
|
device.free(ptr1).expect("Deallocation1 should succeed");
|
|
device.free(ptr2).expect("Deallocation2 should succeed");
|
|
}
|
|
|
|
// TDD tests for advanced stream features
|
|
|
|
#[test]
|
|
fn test_cuda_stream_dependencies() {
|
|
let mut devices = BTreeMap::new();
|
|
discover_devices(&mut devices).expect("Should discover CUDA devices");
|
|
|
|
let device = devices.values().next().expect("Should have at least one device");
|
|
|
|
// Create dependency chain: stream1 -> event -> stream2
|
|
let stream1 = device.create_stream().expect("Stream1 creation should succeed");
|
|
let stream2 = device.create_stream().expect("Stream2 creation should succeed");
|
|
let event = device.create_event().expect("Event creation should succeed");
|
|
|
|
// Work on stream1
|
|
let ptr = device.allocate(1024).expect("Allocation should succeed");
|
|
stream1.copy_async(ptr, ptr, 1024).expect("Copy on stream1 should succeed");
|
|
|
|
// Record completion event
|
|
stream1.record_event(&event).expect("Recording event should succeed");
|
|
|
|
// Stream2 waits for stream1 completion
|
|
stream2.wait_event(&event).expect("Waiting for event should succeed");
|
|
|
|
// More work on stream2 (depends on stream1)
|
|
stream2.copy_async(ptr, ptr, 1024).expect("Copy on stream2 should succeed");
|
|
|
|
// Synchronize all
|
|
device.synchronize().expect("Device sync should succeed");
|
|
|
|
device.free(ptr).expect("Deallocation should succeed");
|
|
}
|
|
|
|
// Edge cases and error conditions
|
|
|
|
#[test]
|
|
fn test_backend_type_display() {
|
|
assert_eq!(format!("{}", BackendType::Cuda), "CUDA");
|
|
assert_eq!(format!("{}", BackendType::Rocm), "ROCm");
|
|
assert_eq!(format!("{}", BackendType::Metal), "Metal");
|
|
assert_eq!(format!("{}", BackendType::Cpu), "CPU");
|
|
}
|
|
|
|
#[test]
|
|
fn test_id_display() {
|
|
assert_eq!(format!("{}", DeviceId(42)), "Device(42)");
|
|
assert_eq!(format!("{}", StreamId(123)), "Stream(123)");
|
|
assert_eq!(format!("{}", EventId(456)), "Event(456)");
|
|
}
|
|
|
|
#[test]
|
|
fn test_event_elapsed_time() {
|
|
let props = DeviceProperties {
|
|
name: "Test GPU".to_string(),
|
|
backend: BackendType::Cuda,
|
|
total_memory: 8 * 1024 * 1024 * 1024,
|
|
available_memory: 7 * 1024 * 1024 * 1024,
|
|
major: 8,
|
|
minor: 0,
|
|
multiprocessor_count: 80,
|
|
max_threads_per_block: 1024,
|
|
max_shared_memory_per_block: 48 * 1024,
|
|
unified_memory: false,
|
|
pci_bus_id: "0000:01:00.0".to_string(),
|
|
};
|
|
|
|
let device = Device::new(DeviceId(0), props).unwrap();
|
|
// Can't fully test without event creation working, but we can test the time function
|
|
let start_time = current_time_us();
|
|
thread::sleep(Duration::from_millis(1));
|
|
let end_time = current_time_us();
|
|
assert!(end_time > start_time);
|
|
}
|
|
}
|