287 lines
8.2 KiB
Rust
287 lines
8.2 KiB
Rust
//! CUDA device management.
|
|
//!
|
|
//! Provides the [`CudaDevice`] type for managing CUDA contexts and streams.
|
|
|
|
#[cfg(feature = "cuda")]
|
|
use crate::CudaBackend;
|
|
#[cfg(feature = "cuda")]
|
|
use cudarc::cublas::CudaBlas;
|
|
#[cfg(feature = "cuda")]
|
|
use cudarc::driver::{CudaContext, CudaStream};
|
|
#[cfg(feature = "cuda")]
|
|
use once_cell::sync::{Lazy, OnceCell};
|
|
#[cfg(feature = "cuda")]
|
|
use parking_lot::RwLock;
|
|
#[cfg(feature = "cuda")]
|
|
use rtx_backend::{DeviceId, DeviceOps};
|
|
#[cfg(feature = "cuda")]
|
|
use std::collections::HashMap;
|
|
#[cfg(feature = "cuda")]
|
|
use std::sync::Arc;
|
|
|
|
use crate::{CudaError, CudaResult};
|
|
|
|
#[cfg(feature = "cuda")]
|
|
/// Global cache of CUDA contexts for device reuse.
|
|
static CONTEXT_CACHE: Lazy<RwLock<HashMap<usize, Arc<CudaContext>>>> =
|
|
Lazy::new(|| RwLock::new(HashMap::new()));
|
|
|
|
#[cfg(feature = "cuda")]
|
|
/// CUDA device wrapper with context and stream management.
|
|
///
|
|
/// Each `CudaDevice` represents a specific GPU and holds:
|
|
/// - A CUDA context (shared across the device)
|
|
/// - A default stream for operations
|
|
/// - Device properties (memory, compute capability)
|
|
/// - Lazy-initialized cuBLAS handle
|
|
#[derive(Clone)]
|
|
pub struct CudaDevice {
|
|
/// Device index (0, 1, 2, ...)
|
|
pub(crate) index: usize,
|
|
/// CUDA context handle
|
|
pub(crate) context: Arc<CudaContext>,
|
|
/// Default stream for operations
|
|
pub(crate) stream: Arc<CudaStream>,
|
|
/// Total memory in bytes
|
|
pub(crate) memory_total: usize,
|
|
/// Compute capability (major, minor)
|
|
pub(crate) compute_capability: (u32, u32),
|
|
/// cuBLAS handle (lazy-initialized, shared across clones)
|
|
pub(crate) cublas: Arc<OnceCell<CudaBlas>>,
|
|
}
|
|
|
|
#[cfg(not(feature = "cuda"))]
|
|
/// CUDA device wrapper (stub when CUDA is not available).
|
|
#[derive(Clone)]
|
|
pub struct CudaDevice {
|
|
/// Device index (0, 1, 2, ...)
|
|
pub(crate) index: usize,
|
|
}
|
|
|
|
#[cfg(feature = "cuda")]
|
|
impl CudaDevice {
|
|
/// Create a new CUDA device.
|
|
///
|
|
/// # Arguments
|
|
/// - `index`: GPU device index (0 for first GPU)
|
|
///
|
|
/// # Returns
|
|
/// CUDA device if available, error otherwise.
|
|
pub fn new(index: usize) -> CudaResult<Self> {
|
|
// Check cache first
|
|
{
|
|
let cache = CONTEXT_CACHE.read();
|
|
if let Some(context) = cache.get(&index) {
|
|
let stream = context.default_stream();
|
|
return Ok(Self {
|
|
index,
|
|
context: context.clone(),
|
|
stream,
|
|
memory_total: 24 * 1024 * 1024 * 1024, // 24GB default
|
|
compute_capability: (9, 0), // sm_90 (RTX 5090)
|
|
cublas: Arc::new(OnceCell::new()),
|
|
});
|
|
}
|
|
}
|
|
|
|
// Initialize CUDA
|
|
cudarc::driver::result::init()
|
|
.map_err(|e| CudaError::Initialization(format!("Failed to init CUDA: {:?}", e)))?;
|
|
|
|
// Get device count
|
|
let count = CudaContext::device_count().map_err(|e| {
|
|
CudaError::Initialization(format!("Failed to get device count: {:?}", e))
|
|
})?;
|
|
|
|
if index >= count as usize {
|
|
return Err(CudaError::DeviceNotFound(index, count as usize));
|
|
}
|
|
|
|
// Create context (CudaContext::new returns Arc<CudaContext>)
|
|
let context = CudaContext::new(index)
|
|
.map_err(|e| CudaError::Initialization(format!("Failed to create context: {:?}", e)))?;
|
|
|
|
// Get default stream
|
|
let stream = context.default_stream();
|
|
|
|
// Cache the context
|
|
{
|
|
let mut cache = CONTEXT_CACHE.write();
|
|
cache.insert(index, context.clone());
|
|
}
|
|
|
|
// Get device properties
|
|
let memory_total = context
|
|
.attribute(
|
|
cudarc::driver::sys::CUdevice_attribute::CU_DEVICE_ATTRIBUTE_TOTAL_CONSTANT_MEMORY,
|
|
)
|
|
.unwrap_or(24_i32 * 1024 * 1024) as usize
|
|
* 1024; // 24GB default
|
|
|
|
// Get compute capability
|
|
let compute_capability = context
|
|
.compute_capability()
|
|
.map(|(major, minor)| (major as u32, minor as u32))
|
|
.unwrap_or((9, 0));
|
|
|
|
Ok(Self {
|
|
index,
|
|
context,
|
|
stream,
|
|
memory_total,
|
|
compute_capability,
|
|
cublas: Arc::new(OnceCell::new()),
|
|
})
|
|
}
|
|
|
|
/// Get the CUDA context.
|
|
pub fn context(&self) -> &Arc<CudaContext> {
|
|
&self.context
|
|
}
|
|
|
|
/// Get the default CUDA stream.
|
|
pub fn stream(&self) -> &Arc<CudaStream> {
|
|
&self.stream
|
|
}
|
|
|
|
/// Get or create a cuBLAS handle for this device.
|
|
///
|
|
/// The handle is lazily initialized and cached for reuse.
|
|
pub fn cublas(&self) -> CudaResult<&CudaBlas> {
|
|
self.cublas.get_or_try_init(|| {
|
|
CudaBlas::new(self.stream.clone())
|
|
.map_err(|e| CudaError::CuBlas(format!("Failed to create cuBLAS handle: {:?}", e)))
|
|
})
|
|
}
|
|
|
|
/// Create a new CUDA stream on this device.
|
|
pub fn create_stream(&self) -> CudaResult<Arc<CudaStream>> {
|
|
self.context
|
|
.new_stream()
|
|
.map_err(|e| CudaError::StreamCreation(format!("{:?}", e)))
|
|
}
|
|
|
|
/// Get device name.
|
|
pub fn name(&self) -> String {
|
|
self.context
|
|
.name()
|
|
.unwrap_or_else(|_| format!("CUDA Device {}", self.index))
|
|
}
|
|
|
|
/// Synchronize all operations on this device.
|
|
pub fn synchronize(&self) {
|
|
let _ = self.stream.synchronize();
|
|
}
|
|
}
|
|
|
|
#[cfg(not(feature = "cuda"))]
|
|
impl CudaDevice {
|
|
/// Create a new CUDA device (stub - always fails without CUDA feature).
|
|
pub fn new(_index: usize) -> CudaResult<Self> {
|
|
Err(CudaError::Initialization(
|
|
"CUDA support not compiled. Enable the 'cuda' feature.".to_string(),
|
|
))
|
|
}
|
|
|
|
/// Get device name (stub).
|
|
pub fn name(&self) -> String {
|
|
format!("CUDA Device {} (unavailable)", self.index)
|
|
}
|
|
|
|
/// Synchronize all operations on this device (stub - no-op).
|
|
pub fn synchronize(&self) {
|
|
// No-op
|
|
}
|
|
}
|
|
|
|
#[cfg(feature = "cuda")]
|
|
impl std::fmt::Debug for CudaDevice {
|
|
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
|
f.debug_struct("CudaDevice")
|
|
.field("index", &self.index)
|
|
.field("name", &self.name())
|
|
.field("memory_total", &self.memory_total)
|
|
.field("compute_capability", &self.compute_capability)
|
|
.finish()
|
|
}
|
|
}
|
|
|
|
#[cfg(not(feature = "cuda"))]
|
|
impl std::fmt::Debug for CudaDevice {
|
|
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
|
f.debug_struct("CudaDevice")
|
|
.field("index", &self.index)
|
|
.field("name", &"unavailable")
|
|
.finish()
|
|
}
|
|
}
|
|
|
|
impl Default for CudaDevice {
|
|
fn default() -> Self {
|
|
Self::new(0).expect("No CUDA device available")
|
|
}
|
|
}
|
|
|
|
impl PartialEq for CudaDevice {
|
|
fn eq(&self, other: &Self) -> bool {
|
|
self.index == other.index
|
|
}
|
|
}
|
|
|
|
impl Eq for CudaDevice {}
|
|
|
|
impl std::hash::Hash for CudaDevice {
|
|
fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
|
|
self.index.hash(state);
|
|
}
|
|
}
|
|
|
|
#[cfg(feature = "cuda")]
|
|
impl DeviceOps<CudaBackend> for CudaDevice {
|
|
fn id(&self) -> DeviceId {
|
|
DeviceId::Cuda(self.index)
|
|
}
|
|
|
|
fn memory_capacity(&self) -> usize {
|
|
self.memory_total
|
|
}
|
|
|
|
fn memory_available(&self) -> usize {
|
|
// Query available memory from CUDA (simplified)
|
|
self.memory_total / 2
|
|
}
|
|
|
|
fn compute_capability(&self) -> Option<(u32, u32)> {
|
|
Some(self.compute_capability)
|
|
}
|
|
|
|
fn synchronize(&self) {
|
|
let _ = self.stream.synchronize();
|
|
}
|
|
|
|
fn is_available(&self) -> bool {
|
|
CudaContext::device_count()
|
|
.map(|c| self.index < c as usize)
|
|
.unwrap_or(false)
|
|
}
|
|
}
|
|
|
|
// Note: DeviceOps impl for non-cuda builds is not provided because
|
|
// CudaBackend doesn't implement Backend without the cuda feature.
|
|
|
|
#[cfg(all(test, feature = "cuda"))]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_device_creation() {
|
|
// This test requires CUDA hardware
|
|
if let Ok(device) = CudaDevice::new(0) {
|
|
assert_eq!(device.index, 0);
|
|
assert!(device.memory_total > 0);
|
|
#[cfg(feature = "cuda")]
|
|
assert!(device.is_available());
|
|
}
|
|
}
|
|
}
|