#[cfg(feature = "cuda")]
use crate::driver::{cuda_available, device_count, CudaContext};
use crate::GpuError;
#[derive(Debug, Clone)]
pub struct CudaDeviceInfo {
pub index: u32,
pub name: String,
pub total_memory: u64,
}
impl CudaDeviceInfo {
#[cfg(feature = "cuda")]
#[allow(clippy::cast_possible_wrap)]
pub fn query(device_index: u32) -> Result<Self, GpuError> {
let ctx = CudaContext::new(device_index as i32)?;
let name = ctx.device_name()?;
let total_memory = ctx.total_memory()? as u64;
Ok(Self { index: device_index, name, total_memory })
}
#[cfg(not(feature = "cuda"))]
pub fn query(_device_index: u32) -> Result<Self, GpuError> {
Err(GpuError::CudaNotAvailable("cuda feature not enabled".to_string()))
}
#[cfg(feature = "cuda")]
pub fn enumerate() -> Result<Vec<Self>, GpuError> {
let count = device_count()?;
let mut devices = Vec::with_capacity(count);
for i in 0..count {
devices.push(Self::query(i as u32)?);
}
Ok(devices)
}
#[cfg(not(feature = "cuda"))]
pub fn enumerate() -> Result<Vec<Self>, GpuError> {
Err(GpuError::CudaNotAvailable("cuda feature not enabled".to_string()))
}
#[must_use]
pub fn total_memory_mb(&self) -> u64 {
self.total_memory / (1024 * 1024)
}
#[must_use]
pub fn total_memory_gb(&self) -> f64 {
self.total_memory as f64 / (1024.0 * 1024.0 * 1024.0)
}
}
impl std::fmt::Display for CudaDeviceInfo {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "[{}] {} ({:.1} GB)", self.index, self.name, self.total_memory_gb())
}
}
#[derive(Debug, Clone, Copy)]
pub struct CudaMemoryInfo {
pub free: u64,
pub total: u64,
}
impl CudaMemoryInfo {
#[cfg(feature = "cuda")]
pub fn query(ctx: &CudaContext) -> Result<Self, GpuError> {
let (free, total) = ctx.memory_info()?;
Ok(Self { free: free as u64, total: total as u64 })
}
#[must_use]
pub fn used(&self) -> u64 {
self.total.saturating_sub(self.free)
}
#[must_use]
pub fn free_mb(&self) -> u64 {
self.free / (1024 * 1024)
}
#[must_use]
pub fn total_mb(&self) -> u64 {
self.total / (1024 * 1024)
}
#[must_use]
pub fn used_mb(&self) -> u64 {
self.used() / (1024 * 1024)
}
#[must_use]
pub fn usage_percent(&self) -> f64 {
if self.total == 0 {
0.0
} else {
(self.used() as f64 / self.total as f64) * 100.0
}
}
}
impl std::fmt::Display for CudaMemoryInfo {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{} / {} MB ({:.1}% used)", self.used_mb(), self.total_mb(), self.usage_percent())
}
}
#[must_use]
pub fn cuda_monitoring_available() -> bool {
#[cfg(feature = "cuda")]
{
cuda_available()
}
#[cfg(not(feature = "cuda"))]
{
false
}
}
pub fn cuda_device_count() -> Result<usize, GpuError> {
#[cfg(feature = "cuda")]
{
device_count()
}
#[cfg(not(feature = "cuda"))]
{
Err(GpuError::CudaNotAvailable("cuda feature not enabled".to_string()))
}
}
#[cfg(test)]
mod tests;