gpu_probe/
lib.rs

1// SPDX-License-Identifier: Apache-2.0
2//! Cross-platform GPU memory (VRAM) detection with **no vendor SDKs**.
3//!
4//! `gpu_probe` reports the GPUs visible on the host and how much memory each
5//! has, using only facilities the OS or driver already ship:
6//!
7//! - **NVIDIA** (Linux, Windows): NVML (`libnvidia-ml`) via `nvml-wrapper`,
8//!   loaded at runtime. The CUDA toolkit is not required and nothing links at
9//!   build time. Behind the default `nvidia` feature.
10//! - **AMD & Intel** (Linux): DRM sysfs under `/sys/class/drm`. Discrete cards
11//!   report dedicated VRAM; integrated GPUs report the shared system-memory
12//!   ceiling (see [`GpuInfo::total_bytes`]).
13//! - **Apple/macOS**: `system_profiler` + `sysctl` (Apple Silicon reports
14//!   unified memory).
15//!
16//! Detection is best-effort: [`detect`] returns an empty `Vec` when no GPU is
17//! found or the platform is unsupported — never an error.
18//!
19//! ```no_run
20//! for gpu in gpu_probe::detect() {
21//!     println!("{gpu}");
22//! }
23//! ```
24
25mod drm;
26mod metal;
27mod nvidia;
28
29/// GPU hardware vendor.
30#[derive(Debug, Clone, Copy, PartialEq, Eq)]
31#[non_exhaustive]
32pub enum Vendor {
33    Nvidia,
34    Amd,
35    Intel,
36    Apple,
37    Unknown,
38}
39
40impl std::fmt::Display for Vendor {
41    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
42        f.write_str(match self {
43            Vendor::Nvidia => "NVIDIA",
44            Vendor::Amd => "AMD",
45            Vendor::Intel => "Intel",
46            Vendor::Apple => "Apple",
47            Vendor::Unknown => "Unknown",
48        })
49    }
50}
51
52/// A single detected GPU and its memory.
53#[derive(Debug, Clone, PartialEq, Eq)]
54#[non_exhaustive]
55pub struct GpuInfo {
56    /// Human-readable name (e.g. `"NVIDIA GeForce RTX 4090"`).
57    pub name: String,
58    /// Hardware vendor.
59    pub vendor: Vendor,
60    /// Total memory in bytes. For discrete GPUs this is dedicated VRAM; for
61    /// integrated/unified GPUs (Intel iGPUs, AMD APUs, Apple Silicon) it is the
62    /// shared system-memory ceiling available to the GPU, not a dedicated pool.
63    pub total_bytes: u64,
64    /// Free device memory in bytes, when known.
65    pub free_bytes: Option<u64>,
66    /// Used device memory in bytes, when known.
67    pub used_bytes: Option<u64>,
68}
69
70impl std::fmt::Display for GpuInfo {
71    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
72        write!(
73            f,
74            "{} ({}): {:.1} GiB total",
75            self.name,
76            self.vendor,
77            gib(self.total_bytes)
78        )?;
79        if let Some(free) = self.free_bytes {
80            write!(f, ", {:.1} GiB free", gib(free))?;
81        }
82        Ok(())
83    }
84}
85
86#[allow(clippy::cast_precision_loss)] // display-only; the imprecision is cosmetic
87fn gib(bytes: u64) -> f64 {
88    bytes as f64 / (1024.0 * 1024.0 * 1024.0)
89}
90
91/// CUDA compute capability, e.g. `8.6` for `sm_86`.
92///
93/// Ordered `major` first, so a host can be checked against a minimum:
94///
95/// ```
96/// use gpu_probe::ComputeCapability;
97/// assert!(ComputeCapability::new(8, 6) >= ComputeCapability::new(8, 0));
98/// assert!(ComputeCapability::new(9, 0) >= ComputeCapability::new(8, 9));
99/// ```
100///
101/// Constructible so callers can express such a requirement.
102#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
103pub struct ComputeCapability {
104    /// Major version — the `8` in `8.6`.
105    pub major: u32,
106    /// Minor version — the `6` in `8.6`.
107    pub minor: u32,
108}
109
110impl ComputeCapability {
111    /// Create a compute capability from its major and minor parts.
112    #[must_use]
113    pub const fn new(major: u32, minor: u32) -> Self {
114        Self { major, minor }
115    }
116}
117
118impl std::fmt::Display for ComputeCapability {
119    /// Renders as `8.6`, matching `nvidia-smi`'s `compute_cap`.
120    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
121        write!(f, "{}.{}", self.major, self.minor)
122    }
123}
124
125/// A CUDA version, e.g. `12.9`.
126///
127/// Ordered `major` first, so a host can be checked against a minimum:
128///
129/// ```
130/// use gpu_probe::CudaVersion;
131/// assert!(CudaVersion::new(12, 9) >= CudaVersion::new(12, 0));
132/// ```
133///
134/// Constructible so callers can express such a requirement.
135#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
136pub struct CudaVersion {
137    /// Major version — the `12` in `12.9`.
138    pub major: u32,
139    /// Minor version — the `9` in `12.9`.
140    pub minor: u32,
141}
142
143impl CudaVersion {
144    /// Create a CUDA version from its major and minor parts.
145    #[must_use]
146    pub const fn new(major: u32, minor: u32) -> Self {
147        Self { major, minor }
148    }
149}
150
151impl std::fmt::Display for CudaVersion {
152    /// Renders as `12.9`.
153    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
154        write!(f, "{}.{}", self.major, self.minor)
155    }
156}
157
158/// Host-wide CUDA properties reported by the NVIDIA driver.
159///
160/// These describe the host and its driver rather than any one GPU, which is why
161/// they are separate from the per-GPU [`GpuInfo`]. Consumers typically use them
162/// to select a prebuilt artifact compatible with the host.
163#[derive(Debug, Clone, Copy, PartialEq, Eq)]
164#[non_exhaustive]
165pub struct CudaHost {
166    /// Compute capability of device 0.
167    pub compute_capability: ComputeCapability,
168    /// Version of the installed CUDA driver.
169    pub driver_version: CudaVersion,
170}
171
172/// Detect all GPUs visible on the host.
173///
174/// Best-effort: spawns only read-only platform queries (NVML, `system_profiler`,
175/// `sysctl`) and reads sysfs. Returns an empty `Vec` on unsupported platforms
176/// or when no GPU is found.
177#[must_use]
178pub fn detect() -> Vec<GpuInfo> {
179    let mut gpus = Vec::new();
180    gpus.extend(nvidia::detect());
181    gpus.extend(drm::detect());
182    gpus.extend(metal::detect());
183    gpus
184}
185
186/// Host-wide CUDA properties, or `None` when NVML is unavailable — no NVIDIA
187/// driver, the `nvidia` feature disabled, no device, or a driver reporting
188/// values that aren't usable.
189///
190/// Shares the one process-wide NVML handle with [`detect`], so calling this on
191/// a timer does not accumulate resources.
192///
193/// ```no_run
194/// use gpu_probe::ComputeCapability;
195///
196/// if let Some(cuda) = gpu_probe::cuda_host() {
197///     println!("sm_{}{} on CUDA {}",
198///         cuda.compute_capability.major,
199///         cuda.compute_capability.minor,
200///         cuda.driver_version);
201///
202///     if cuda.compute_capability >= ComputeCapability::new(8, 0) {
203///         // pick an Ampere-or-newer build
204///     }
205/// }
206/// ```
207#[must_use]
208pub fn cuda_host() -> Option<CudaHost> {
209    nvidia::cuda_host()
210}
211
212#[cfg(test)]
213mod tests {
214    use super::*;
215
216    #[test]
217    fn detect_never_panics() {
218        // Environment-dependent (may be empty on headless CI); exercise the
219        // full path plus the Display impl without asserting a GPU exists.
220        for gpu in detect() {
221            assert!(!gpu.name.is_empty());
222            let _ = gpu.to_string();
223        }
224    }
225
226    #[test]
227    fn display_includes_free_when_present() {
228        let gpu = GpuInfo {
229            name: "Test GPU".to_string(),
230            vendor: Vendor::Nvidia,
231            total_bytes: 24 * 1024 * 1024 * 1024,
232            free_bytes: Some(12 * 1024 * 1024 * 1024),
233            used_bytes: Some(12 * 1024 * 1024 * 1024),
234        };
235        let shown = gpu.to_string();
236        assert!(shown.contains("NVIDIA"));
237        assert!(shown.contains("24.0 GiB total"));
238        assert!(shown.contains("12.0 GiB free"));
239    }
240
241    #[test]
242    fn display_omits_free_when_absent() {
243        let gpu = GpuInfo {
244            name: "AMD GPU (card0)".to_string(),
245            vendor: Vendor::Amd,
246            total_bytes: 8 * 1024 * 1024 * 1024,
247            free_bytes: None,
248            used_bytes: None,
249        };
250        let shown = gpu.to_string();
251        assert!(shown.contains("8.0 GiB total"));
252        assert!(!shown.contains("free"));
253    }
254
255    #[test]
256    fn vendor_display_covers_every_variant() {
257        assert_eq!(Vendor::Nvidia.to_string(), "NVIDIA");
258        assert_eq!(Vendor::Amd.to_string(), "AMD");
259        assert_eq!(Vendor::Intel.to_string(), "Intel");
260        assert_eq!(Vendor::Apple.to_string(), "Apple");
261        assert_eq!(Vendor::Unknown.to_string(), "Unknown");
262    }
263
264    #[test]
265    fn gib_converts_using_binary_units() {
266        assert!((gib(0) - 0.0).abs() < f64::EPSILON);
267        assert!((gib(1024 * 1024 * 1024) - 1.0).abs() < f64::EPSILON);
268        // 1.5 GiB exercises the fractional path the Display rounds to one place.
269        assert!((gib(3 * 1024 * 1024 * 1024 / 2) - 1.5).abs() < f64::EPSILON);
270    }
271
272    #[test]
273    fn display_rounds_to_one_decimal_place() {
274        // 25 GiB + 256 MiB -> 25.25 GiB, which "{:.1}" renders as "25.2".
275        let gpu = GpuInfo {
276            name: "Rounding".to_string(),
277            vendor: Vendor::Nvidia,
278            total_bytes: 25 * 1024 * 1024 * 1024 + 256 * 1024 * 1024,
279            free_bytes: None,
280            used_bytes: None,
281        };
282        assert!(gpu.to_string().contains("25.2 GiB total"));
283    }
284
285    #[test]
286    fn detect_results_have_consistent_memory_fields() {
287        // Environment-dependent; asserts invariants only for whatever is present.
288        for gpu in detect() {
289            assert!(!gpu.name.is_empty());
290            if let Some(free) = gpu.free_bytes {
291                assert!(free <= gpu.total_bytes, "free must not exceed total");
292            }
293            if let (Some(free), Some(used)) = (gpu.free_bytes, gpu.used_bytes) {
294                assert!(
295                    free.saturating_add(used) <= gpu.total_bytes.saturating_add(used),
296                    "free/used must be coherent",
297                );
298            }
299        }
300    }
301
302    #[test]
303    fn versions_display_as_major_dot_minor() {
304        assert_eq!(ComputeCapability::new(8, 6).to_string(), "8.6");
305        assert_eq!(CudaVersion::new(12, 9).to_string(), "12.9");
306        // A two-digit minor stays unambiguous — the reason these aren't packed
307        // into a single integer.
308        assert_eq!(ComputeCapability::new(8, 10).to_string(), "8.10");
309    }
310
311    #[test]
312    fn versions_order_by_major_then_minor() {
313        assert!(ComputeCapability::new(8, 6) > ComputeCapability::new(8, 0));
314        assert!(ComputeCapability::new(9, 0) > ComputeCapability::new(8, 9));
315        assert_eq!(ComputeCapability::new(8, 6), ComputeCapability::new(8, 6));
316        assert!(CudaVersion::new(12, 9) > CudaVersion::new(12, 0));
317        assert!(CudaVersion::new(13, 0) > CudaVersion::new(12, 9));
318        // Packing as `major * 10 + minor` would collide here: 8.10 and 9.0
319        // both pack to 90, which is why the parts are kept separate.
320        assert!(ComputeCapability::new(9, 0) > ComputeCapability::new(8, 10));
321    }
322
323    #[test]
324    fn cuda_host_is_environment_dependent_but_coherent() {
325        // No NVIDIA driver is a valid, passing environment.
326        if let Some(cuda) = cuda_host() {
327            assert!(
328                cuda.compute_capability.major > 0,
329                "a real device has a nonzero major capability",
330            );
331            assert!(cuda.driver_version.major > 0, "a real driver has a version");
332            assert_eq!(
333                cuda_host(),
334                Some(cuda),
335                "host/driver properties must be stable across calls",
336            );
337        }
338    }
339
340    #[test]
341    fn gpu_info_equality_compares_all_fields() {
342        let base = GpuInfo {
343            name: "G".to_string(),
344            vendor: Vendor::Intel,
345            total_bytes: 16 * 1024 * 1024 * 1024,
346            free_bytes: None,
347            used_bytes: None,
348        };
349        assert_eq!(base.clone(), base);
350        let mut other = base.clone();
351        other.vendor = Vendor::Amd;
352        assert_ne!(base, other);
353    }
354}
gpu_probe/lib.rs

gpu_probe/
lib.rs
