Skip to main content

entrenar/efficiency/device/
gpu.rs

1//! GPU information and detection.
2
3use serde::{Deserialize, Serialize};
4
5/// GPU information
6#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
7pub struct GpuInfo {
8    /// GPU name/model
9    pub name: String,
10    /// Video RAM in bytes
11    pub vram_bytes: u64,
12    /// CUDA compute capability (major, minor) if NVIDIA
13    pub compute_capability: Option<(u32, u32)>,
14    /// GPU index (for multi-GPU systems)
15    pub index: u32,
16}
17
18impl GpuInfo {
19    /// Create new GPU info
20    pub fn new(name: impl Into<String>, vram_bytes: u64) -> Self {
21        Self { name: name.into(), vram_bytes, compute_capability: None, index: 0 }
22    }
23
24    /// Set CUDA compute capability
25    pub fn with_compute_capability(mut self, major: u32, minor: u32) -> Self {
26        self.compute_capability = Some((major, minor));
27        self
28    }
29
30    /// Set GPU index
31    pub fn with_index(mut self, index: u32) -> Self {
32        self.index = index;
33        self
34    }
35
36    /// Check if GPU supports specific CUDA compute capability
37    pub fn supports_compute_capability(&self, major: u32, minor: u32) -> bool {
38        self.compute_capability.is_some_and(|(m, n)| m > major || (m == major && n >= minor))
39    }
40
41    /// Get VRAM in GB
42    pub fn vram_gb(&self) -> f64 {
43        self.vram_bytes as f64 / (1024.0 * 1024.0 * 1024.0)
44    }
45}