Skip to main content

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, and AMD APUs their VRAM carveout plus GTT pool (see
13//!   [`GpuInfo::total_bytes`]). AMD cards additionally report their `gfx`
14//!   target from KFD sysfs — no `ROCm` install needed.
15//! - **Apple/macOS**: `system_profiler` + `sysctl` for the chip and its memory
16//!   ceiling, plus `vm_stat` for the used/free split (Apple Silicon reports
17//!   unified memory, so that split is system-wide).
18//!
19//! Host toolchain properties are reported separately from any one GPU:
20//! [`cuda_host`] for the CUDA driver, [`rocm_host`] for the `ROCm` install,
21//! [`oneapi_host`] for the Intel `oneAPI` install, and [`vulkan_host`] for the
22//! Vulkan runtime.
23//!
24//! Detection is best-effort: [`detect`] returns an empty `Vec` when no GPU is
25//! found or the platform is unsupported — never an error.
26//!
27//! ```no_run
28//! for gpu in gpu_probe::detect() {
29//!     println!("{gpu}");
30//! }
31//! ```
32
33mod drm;
34mod intel;
35mod kfd;
36mod metal;
37mod nvidia;
38mod oneapi;
39mod rocm;
40mod vulkan;
41
42/// GPU hardware vendor.
43#[derive(Debug, Clone, Copy, PartialEq, Eq)]
44#[non_exhaustive]
45pub enum Vendor {
46    Nvidia,
47    Amd,
48    Intel,
49    Apple,
50    Unknown,
51}
52
53impl std::fmt::Display for Vendor {
54    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
55        f.write_str(match self {
56            Vendor::Nvidia => "NVIDIA",
57            Vendor::Amd => "AMD",
58            Vendor::Intel => "Intel",
59            Vendor::Apple => "Apple",
60            Vendor::Unknown => "Unknown",
61        })
62    }
63}
64
65/// A single detected GPU and its memory.
66#[derive(Debug, Clone, PartialEq, Eq)]
67#[non_exhaustive]
68pub struct GpuInfo {
69    /// Human-readable name (e.g. `"NVIDIA GeForce RTX 4090"`).
70    pub name: String,
71    /// Hardware vendor.
72    pub vendor: Vendor,
73    /// Total memory in bytes. For discrete GPUs this is dedicated VRAM; for
74    /// integrated/unified GPUs (Intel iGPUs, AMD APUs, Apple Silicon) it is the
75    /// shared system-memory ceiling available to the GPU, not a dedicated pool.
76    ///
77    /// An AMD APU reports its BIOS VRAM carveout plus the GTT pool the driver
78    /// allocates from — the latter sized by the kernel's `ttm.pages_limit` —
79    /// since the carveout alone is far below what the part can actually hand
80    /// out (512 MiB of 14.5 GiB on one such part).
81    pub total_bytes: u64,
82    /// Free device memory in bytes, when known.
83    pub free_bytes: Option<u64>,
84    /// Used device memory in bytes, when known.
85    pub used_bytes: Option<u64>,
86    /// The architecture a prebuilt artifact must target to run on this GPU:
87    /// [`ArchTarget::Gfx`] on AMD, from KFD sysfs, and [`ArchTarget::Sm`] on
88    /// NVIDIA, from NVML.
89    ///
90    /// `None` when neither driver reports one — an Apple or Intel GPU, an AMD
91    /// card on a kernel without KFD, or the `nvidia` feature disabled. The
92    /// NVIDIA value also appears on [`CudaHost`], which reports device 0's
93    /// alongside the host driver version; this field is per-GPU.
94    pub arch_target: Option<ArchTarget>,
95}
96
97impl std::fmt::Display for GpuInfo {
98    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
99        write!(f, "{} ({}", self.name, self.vendor)?;
100        if let Some(arch) = self.arch_target {
101            write!(f, ", {arch}")?;
102        }
103        write!(f, "): {:.1} GiB total", gib(self.total_bytes))?;
104        if let Some(free) = self.free_bytes {
105            write!(f, ", {:.1} GiB free", gib(free))?;
106        }
107        Ok(())
108    }
109}
110
111#[allow(clippy::cast_precision_loss)] // display-only; the imprecision is cosmetic
112fn gib(bytes: u64) -> f64 {
113    bytes as f64 / (1024.0 * 1024.0 * 1024.0)
114}
115
116/// CUDA compute capability, e.g. `8.6` for `sm_86`.
117///
118/// Ordered `major` first, so a host can be checked against a minimum:
119///
120/// ```
121/// use gpu_probe::ComputeCapability;
122/// assert!(ComputeCapability::new(8, 6) >= ComputeCapability::new(8, 0));
123/// assert!(ComputeCapability::new(9, 0) >= ComputeCapability::new(8, 9));
124/// ```
125///
126/// Constructible so callers can express such a requirement.
127#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
128pub struct ComputeCapability {
129    /// Major version — the `8` in `8.6`.
130    pub major: u32,
131    /// Minor version — the `6` in `8.6`.
132    pub minor: u32,
133}
134
135impl ComputeCapability {
136    /// Create a compute capability from its major and minor parts.
137    #[must_use]
138    pub const fn new(major: u32, minor: u32) -> Self {
139        Self { major, minor }
140    }
141}
142
143impl std::fmt::Display for ComputeCapability {
144    /// Renders as `8.6`, matching `nvidia-smi`'s `compute_cap`.
145    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
146        write!(f, "{}.{}", self.major, self.minor)
147    }
148}
149
150/// AMD GPU architecture target, e.g. `gfx1013` — the identifier a `ROCm`/HIP
151/// code object is built for (`--offload-arch=gfx1013`).
152///
153/// The AMD counterpart of [`ComputeCapability`], and read the same way: to pick
154/// a prebuilt artifact the host can actually run. Ordered `major` first, so a
155/// host can be checked against a minimum:
156///
157/// ```
158/// use gpu_probe::GfxTarget;
159/// assert!(GfxTarget::new(10, 3, 0) >= GfxTarget::new(10, 1, 3));
160/// assert!(GfxTarget::new(11, 0, 0) >= GfxTarget::new(10, 3, 0));
161/// ```
162///
163/// Constructible so callers can express such a requirement.
164#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
165pub struct GfxTarget {
166    /// Major version — the `10` in `gfx1013`.
167    pub major: u32,
168    /// Minor version — the `1` in `gfx1013`.
169    pub minor: u32,
170    /// Stepping — the `3` in `gfx1013`, and the `a` in `gfx90a`.
171    pub step: u32,
172}
173
174impl GfxTarget {
175    /// Create a target from its major, minor, and stepping parts.
176    #[must_use]
177    pub const fn new(major: u32, minor: u32, step: u32) -> Self {
178        Self { major, minor, step }
179    }
180}
181
182impl std::fmt::Display for GfxTarget {
183    /// Renders as `gfx1013`, matching `--offload-arch`. Minor and stepping are
184    /// single hex digits there, so `9.0.10` renders as `gfx90a`.
185    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
186        write!(f, "gfx{}{:x}{:x}", self.major, self.minor, self.step)
187    }
188}
189
190/// Intel GPU architecture family, e.g. [`IntelArch::XeHpg`] for an Arc A-series
191/// card.
192///
193/// Coarser than its AMD and NVIDIA counterparts by necessity. Neither `i915`
194/// nor `xe` publishes an architecture anywhere readable, so this is derived
195/// from the PCI device id, which identifies the family reliably but not the
196/// exact product. The `ocloc -device` value for an ahead-of-time build (`dg2`,
197/// `acm-g10`, …) is more specific than this; treat it as "which generation is
198/// this" rather than a literal compiler argument.
199///
200/// Deliberately not ordered: "newer" across integrated and discrete lines is
201/// not a total order worth implying.
202#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
203#[non_exhaustive]
204pub enum IntelArch {
205    /// Xe-LP — Tiger Lake, Rocket Lake, Alder Lake, Raptor Lake, DG1.
206    XeLp,
207    /// Xe-HPG — DG2, sold as Arc A-series (Alchemist).
208    XeHpg,
209    /// Xe-HPC — Ponte Vecchio, sold as Data Center GPU Max.
210    XeHpc,
211    /// Xe-LPG — Meteor Lake and Arrow Lake integrated graphics.
212    XeLpg,
213    /// Xe2 — Lunar Lake integrated graphics, and Arc B-series (Battlemage).
214    Xe2,
215}
216
217impl std::fmt::Display for IntelArch {
218    /// Renders as the lowercase family name — `xe-hpg`, `xe2`.
219    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
220        f.write_str(match self {
221            IntelArch::XeLp => "xe-lp",
222            IntelArch::XeHpg => "xe-hpg",
223            IntelArch::XeHpc => "xe-hpc",
224            IntelArch::XeLpg => "xe-lpg",
225            IntelArch::Xe2 => "xe2",
226        })
227    }
228}
229
230/// Apple GPU family, e.g. `apple8` for an M2.
231///
232/// The Metal feature tier a shader can be compiled against
233/// (`MTLGPUFamily.apple8`). Ordered, so a minimum can be expressed:
234///
235/// ```
236/// use gpu_probe::AppleFamily;
237/// assert!(AppleFamily::new(9) >= AppleFamily::new(8));
238/// ```
239///
240/// Unlike a `gfx` target or a compute capability, this does not select a build
241/// artifact — a `.metallib` is not per-family — so it reads as a capability
242/// tier. It is derived from the chip name `system_profiler` reports, since
243/// querying it properly means linking Metal.
244#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
245pub struct AppleFamily {
246    /// Family generation — the `8` in `apple8`.
247    pub generation: u32,
248}
249
250impl AppleFamily {
251    /// Create a family from its generation number.
252    #[must_use]
253    pub const fn new(generation: u32) -> Self {
254        Self { generation }
255    }
256}
257
258impl std::fmt::Display for AppleFamily {
259    /// Renders as `apple8`, matching the `MTLGPUFamily` name.
260    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
261        write!(f, "apple{}", self.generation)
262    }
263}
264
265/// The architecture a prebuilt GPU artifact must target.
266///
267/// Each vendor names this differently but uses it the same way — to select a
268/// build the device can actually run — so one field carries whichever form
269/// applies. A GPU has at most one, which the type enforces.
270///
271/// ```
272/// use gpu_probe::{ArchTarget, GfxTarget};
273///
274/// let target = ArchTarget::Gfx(GfxTarget::new(10, 1, 3));
275/// assert_eq!(target.to_string(), "gfx1013");
276/// assert_eq!(target.gfx(), Some(GfxTarget::new(10, 1, 3)));
277/// assert_eq!(target.sm(), None);
278/// ```
279///
280/// Deliberately not `Ord`: comparing an AMD target against an NVIDIA one has no
281/// meaning. Order within a vendor by matching out [`GfxTarget`] or
282/// [`ComputeCapability`], both of which are ordered.
283#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
284#[non_exhaustive]
285pub enum ArchTarget {
286    /// AMD: the `--offload-arch` value a `ROCm`/HIP code object is built for.
287    Gfx(GfxTarget),
288    /// NVIDIA: the compute capability a CUDA artifact is built for.
289    Sm(ComputeCapability),
290    /// Intel: the GPU architecture family, from the PCI device id.
291    Xe(IntelArch),
292    /// Apple: the Metal GPU family. A capability tier rather than a build
293    /// target — see [`AppleFamily`].
294    Apple(AppleFamily),
295}
296
297impl ArchTarget {
298    /// The AMD target, or `None` when this names another vendor's.
299    #[must_use]
300    pub const fn gfx(self) -> Option<GfxTarget> {
301        match self {
302            Self::Gfx(target) => Some(target),
303            _ => None,
304        }
305    }
306
307    /// The NVIDIA compute capability, or `None` when this names another
308    /// vendor's.
309    #[must_use]
310    pub const fn sm(self) -> Option<ComputeCapability> {
311        match self {
312            Self::Sm(capability) => Some(capability),
313            _ => None,
314        }
315    }
316
317    /// The Intel architecture family, or `None` when this names another
318    /// vendor's.
319    #[must_use]
320    pub const fn xe(self) -> Option<IntelArch> {
321        match self {
322            Self::Xe(arch) => Some(arch),
323            _ => None,
324        }
325    }
326
327    /// The Apple GPU family, or `None` when this names another vendor's.
328    #[must_use]
329    pub const fn apple(self) -> Option<AppleFamily> {
330        match self {
331            Self::Apple(family) => Some(family),
332            _ => None,
333        }
334    }
335}
336
337impl std::fmt::Display for ArchTarget {
338    /// Renders in the form each vendor's toolchain expects: `gfx1013` for
339    /// `--offload-arch`, `sm_89` for CUDA.
340    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
341        match self {
342            Self::Gfx(target) => write!(f, "{target}"),
343            Self::Sm(capability) => {
344                write!(f, "sm_{}{}", capability.major, capability.minor)
345            }
346            Self::Xe(arch) => write!(f, "{arch}"),
347            Self::Apple(family) => write!(f, "{family}"),
348        }
349    }
350}
351
352/// A CUDA version, e.g. `12.9`.
353///
354/// Ordered `major` first, so a host can be checked against a minimum:
355///
356/// ```
357/// use gpu_probe::CudaVersion;
358/// assert!(CudaVersion::new(12, 9) >= CudaVersion::new(12, 0));
359/// ```
360///
361/// Constructible so callers can express such a requirement.
362#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
363pub struct CudaVersion {
364    /// Major version — the `12` in `12.9`.
365    pub major: u32,
366    /// Minor version — the `9` in `12.9`.
367    pub minor: u32,
368}
369
370impl CudaVersion {
371    /// Create a CUDA version from its major and minor parts.
372    #[must_use]
373    pub const fn new(major: u32, minor: u32) -> Self {
374        Self { major, minor }
375    }
376}
377
378impl std::fmt::Display for CudaVersion {
379    /// Renders as `12.9`.
380    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
381        write!(f, "{}.{}", self.major, self.minor)
382    }
383}
384
385/// Host-wide CUDA properties reported by the NVIDIA driver.
386///
387/// These describe the host and its driver rather than any one GPU, which is why
388/// they are separate from the per-GPU [`GpuInfo`]. Consumers typically use them
389/// to select a prebuilt artifact compatible with the host.
390#[derive(Debug, Clone, Copy, PartialEq, Eq)]
391#[non_exhaustive]
392pub struct CudaHost {
393    /// Compute capability of device 0.
394    pub compute_capability: ComputeCapability,
395    /// Version of the installed CUDA driver.
396    pub driver_version: CudaVersion,
397}
398
399/// A `ROCm` release version, e.g. `6.2.4`.
400///
401/// Ordered `major` first, so a host can be checked against a minimum:
402///
403/// ```
404/// use gpu_probe::RocmVersion;
405/// assert!(RocmVersion::new(6, 2, 4) >= RocmVersion::new(6, 0, 0));
406/// ```
407///
408/// Constructible so callers can express such a requirement.
409#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
410pub struct RocmVersion {
411    /// Major version — the `6` in `6.2.4`.
412    pub major: u32,
413    /// Minor version — the `2` in `6.2.4`.
414    pub minor: u32,
415    /// Patch version — the `4` in `6.2.4`.
416    pub patch: u32,
417}
418
419impl RocmVersion {
420    /// Create a version from its major, minor, and patch parts.
421    #[must_use]
422    pub const fn new(major: u32, minor: u32, patch: u32) -> Self {
423        Self {
424            major,
425            minor,
426            patch,
427        }
428    }
429}
430
431impl std::fmt::Display for RocmVersion {
432    /// Renders as `6.2.4`, matching the `.info/version` file it comes from.
433    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
434        write!(f, "{}.{}.{}", self.major, self.minor, self.patch)
435    }
436}
437
438/// The host's `ROCm` installation.
439///
440/// The AMD counterpart of [`CudaHost`], but a narrower one: there is no
441/// driver-side version to report, so this describes the userspace install only.
442/// See [`rocm_host`] for what its absence does and does not imply.
443#[derive(Debug, Clone, Copy, PartialEq, Eq)]
444#[non_exhaustive]
445pub struct RocmHost {
446    /// Installed `ROCm` release.
447    pub version: RocmVersion,
448}
449
450/// Parse a dotted version — `6.2.4`, `2024.2` — into major, minor, and patch.
451///
452/// A trailing build suffix (`6.2.4-123`) is dropped: it identifies a package
453/// build, not the release. Patch defaults to `0`, since some releases ship only
454/// `major.minor`. Shared by the `ROCm` and `oneAPI` probes, which read the same
455/// shape of version out of different places.
456fn parse_dotted_version(text: &str) -> Option<(u32, u32, u32)> {
457    let version = text.trim().split(['-', '+']).next()?;
458    let mut parts = version.split('.');
459    let major = parts.next()?.trim().parse().ok()?;
460    let minor = parts.next()?.trim().parse().ok()?;
461    let patch = match parts.next() {
462        Some(patch) => patch.trim().parse().ok()?,
463        None => 0,
464    };
465    Some((major, minor, patch))
466}
467
468/// An Intel `oneAPI` toolkit version, e.g. `2024.2.1`.
469///
470/// Ordered `major` first — which for `oneAPI` is the release year — so a host
471/// can be checked against a minimum:
472///
473/// ```
474/// use gpu_probe::OneApiVersion;
475/// assert!(OneApiVersion::new(2025, 0, 0) >= OneApiVersion::new(2024, 2, 0));
476/// ```
477///
478/// Constructible so callers can express such a requirement.
479#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
480pub struct OneApiVersion {
481    /// Major version — the release year, the `2024` in `2024.2.1`.
482    pub major: u32,
483    /// Minor version — the `2` in `2024.2.1`.
484    pub minor: u32,
485    /// Patch version — the `1` in `2024.2.1`.
486    pub patch: u32,
487}
488
489impl OneApiVersion {
490    /// Create a version from its major, minor, and patch parts.
491    #[must_use]
492    pub const fn new(major: u32, minor: u32, patch: u32) -> Self {
493        Self {
494            major,
495            minor,
496            patch,
497        }
498    }
499}
500
501impl std::fmt::Display for OneApiVersion {
502    /// Renders as `2024.2.1`, matching the install directory it comes from.
503    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
504        write!(f, "{}.{}.{}", self.major, self.minor, self.patch)
505    }
506}
507
508/// The host's Intel `oneAPI` installation.
509///
510/// The Intel counterpart of [`RocmHost`], and equally narrow: a userspace
511/// install, with no driver version behind it. See [`oneapi_host`].
512#[derive(Debug, Clone, Copy, PartialEq, Eq)]
513#[non_exhaustive]
514pub struct OneApiHost {
515    /// Installed `oneAPI` toolkit release.
516    pub version: OneApiVersion,
517}
518
519/// A Vulkan API version, e.g. `1.3.280`.
520///
521/// Ordered `major` first, so a host can be checked against a minimum:
522///
523/// ```
524/// use gpu_probe::VulkanVersion;
525/// assert!(VulkanVersion::new(1, 3, 280) >= VulkanVersion::new(1, 2, 0));
526/// ```
527///
528/// Constructible so callers can express such a requirement.
529///
530/// Gate on `major`/`minor`. Vulkan's patch number is the spec header revision
531/// and carries no feature guarantee — a 1.4.354 driver and a 1.4.357 loader
532/// are both Vulkan 1.4 — so a patch-sensitive comparison rejects builds that
533/// would have run.
534#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
535pub struct VulkanVersion {
536    /// Major version — the `1` in `1.3.280`.
537    pub major: u32,
538    /// Minor version — the `3` in `1.3.280`.
539    pub minor: u32,
540    /// Patch version — the `280` in `1.3.280`. The spec header revision, not
541    /// a feature level.
542    pub patch: u32,
543}
544
545impl VulkanVersion {
546    /// Create a version from its major, minor, and patch parts.
547    #[must_use]
548    pub const fn new(major: u32, minor: u32, patch: u32) -> Self {
549        Self {
550            major,
551            minor,
552            patch,
553        }
554    }
555}
556
557impl std::fmt::Display for VulkanVersion {
558    /// Renders as `1.3.280`.
559    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
560        write!(f, "{}.{}.{}", self.major, self.minor, self.patch)
561    }
562}
563
564/// The host's Vulkan runtime.
565///
566/// Reported from the loader and the installed ICD manifests, so — unlike
567/// [`RocmHost`] and [`OneApiHost`] — this describes a runtime that is actually
568/// present rather than a toolkit that may be. There is no architecture field:
569/// SPIR-V is portable and driver-compiled, so a Vulkan build has no per-GPU
570/// target to match.
571///
572/// Host-level, not per-device: see [`api_version`](Self::api_version) for what
573/// the number does and does not promise.
574#[derive(Debug, Clone, Copy, PartialEq, Eq)]
575#[non_exhaustive]
576pub struct VulkanHost {
577    /// Highest API version any installed driver advertises in its ICD
578    /// manifest.
579    ///
580    /// A driver's own static declaration on disk, which makes it two things it
581    /// is easy to mistake it for:
582    ///
583    /// - **Not per-device.** With two drivers installed — an AMD iGPU beside
584    ///   an NVIDIA dGPU, say — this is the higher of the two and may describe
585    ///   neither card. A specific device's version comes from
586    ///   `vkGetPhysicalDeviceProperties`, which means linking the loader and
587    ///   creating an instance; this crate deliberately does neither.
588    /// - **Not the loader's instance version.** That is what `vulkaninfo` and
589    ///   `vkEnumerateInstanceVersion` report, and it is usually the newer of
590    ///   the two, so the numbers routinely disagree. The driver's is the one
591    ///   that binds in practice: loaders track the current headers while
592    ///   drivers implement features on their own schedule.
593    ///
594    /// Compare on `major`/`minor` only — see [`VulkanVersion`].
595    pub api_version: VulkanVersion,
596}
597
598/// Detect all GPUs visible on the host.
599///
600/// Best-effort: spawns only read-only platform queries (NVML, `system_profiler`,
601/// `sysctl`, `vm_stat`) and reads sysfs. Returns an empty `Vec` on unsupported
602/// platforms or when no GPU is found.
603#[must_use]
604pub fn detect() -> Vec<GpuInfo> {
605    let mut gpus = Vec::new();
606    gpus.extend(nvidia::detect());
607    gpus.extend(drm::detect());
608    gpus.extend(metal::detect());
609    gpus
610}
611
612/// Host-wide CUDA properties, or `None` when NVML is unavailable — no NVIDIA
613/// driver, the `nvidia` feature disabled, no device, or a driver reporting
614/// values that aren't usable.
615///
616/// Shares the one process-wide NVML handle with [`detect`], so calling this on
617/// a timer does not accumulate resources.
618///
619/// ```no_run
620/// use gpu_probe::ComputeCapability;
621///
622/// if let Some(cuda) = gpu_probe::cuda_host() {
623///     println!("sm_{}{} on CUDA {}",
624///         cuda.compute_capability.major,
625///         cuda.compute_capability.minor,
626///         cuda.driver_version);
627///
628///     if cuda.compute_capability >= ComputeCapability::new(8, 0) {
629///         // pick an Ampere-or-newer build
630///     }
631/// }
632/// ```
633#[must_use]
634pub fn cuda_host() -> Option<CudaHost> {
635    nvidia::cuda_host()
636}
637
638/// The host's `ROCm` installation, or `None` when `ROCm` is not installed.
639///
640/// Read from `$ROCM_PATH/.info/version`, falling back to `/opt/rocm` — the
641/// plain text file the `rocm-core` package writes. Nothing is linked or
642/// executed, so this costs one file read.
643///
644/// `Some` is the signal that `ROCm` is installed and the host can run its
645/// builds. `None` is weaker: the install was not found at the prefixes above,
646/// which a distro shipping `ROCm` into `/usr` — or a container carrying only
647/// the runtime libraries — will trigger despite working. Treat `Some` as proof
648/// and `None` as "probably not, worth confirming".
649///
650/// `None` does not mean the GPU is unusable for compute: the kernel side is a
651/// separate component, and what a build has to target is
652/// [`GpuInfo::arch_target`], reported with no `ROCm` installed at all.
653///
654/// ```no_run
655/// use gpu_probe::RocmVersion;
656///
657/// if let Some(rocm) = gpu_probe::rocm_host()
658///     && rocm.version >= RocmVersion::new(6, 0, 0)
659/// {
660///     // pick a `ROCm` 6 build
661/// }
662/// ```
663#[must_use]
664pub fn rocm_host() -> Option<RocmHost> {
665    rocm::host()
666}
667
668/// The host's Intel `oneAPI` installation, or `None` when it is not installed.
669///
670/// Read from the component layout under `$ONEAPI_ROOT`, falling back to
671/// `/opt/intel/oneapi`. Nothing is linked or executed.
672///
673/// Narrower than it looks: this reports the **toolkit**, not the GPU runtime.
674/// A host running compute through a distro-packaged Level Zero driver with no
675/// toolkit installed reports `None`, because reading that runtime's version
676/// requires linking it rather than reading a file. So `Some` proves the
677/// toolkit is present, while `None` is the weakest negative of the three
678/// probes — it does not rule out a usable Level Zero runtime.
679///
680/// ```no_run
681/// use gpu_probe::OneApiVersion;
682///
683/// if let Some(oneapi) = gpu_probe::oneapi_host()
684///     && oneapi.version >= OneApiVersion::new(2024, 0, 0)
685/// {
686///     // pick a oneAPI 2024-or-newer build
687/// }
688/// ```
689#[must_use]
690pub fn oneapi_host() -> Option<OneApiHost> {
691    oneapi::host()
692}
693
694/// The host's Vulkan runtime, or `None` when no loader is installed.
695///
696/// Read from `libvulkan.so.1` plus the ICD manifests under
697/// `/usr/share/vulkan/icd.d`. Nothing is linked or executed, so this costs a
698/// handful of file reads.
699///
700/// The version is the highest any installed *driver* advertises — neither the
701/// loader's instance version nor any one GPU's, both of which would require
702/// calling into the loader. See [`VulkanHost::api_version`].
703///
704/// ```no_run
705/// use gpu_probe::VulkanVersion;
706///
707/// if let Some(vulkan) = gpu_probe::vulkan_host()
708///     && vulkan.api_version >= VulkanVersion::new(1, 2, 0)
709/// {
710///     // pick a Vulkan 1.2-or-newer build
711/// }
712/// ```
713#[must_use]
714pub fn vulkan_host() -> Option<VulkanHost> {
715    vulkan::host()
716}
717
718#[cfg(test)]
719mod tests {
720    use super::*;
721
722    #[test]
723    fn parses_dotted_versions_in_both_shapes() {
724        assert_eq!(parse_dotted_version("6.2.4-123"), Some((6, 2, 4)));
725        assert_eq!(parse_dotted_version("2024.2"), Some((2024, 2, 0)));
726        assert_eq!(parse_dotted_version("  5.7.1  "), Some((5, 7, 1)));
727        assert_eq!(parse_dotted_version("6"), None, "a bare major is not one");
728        assert_eq!(parse_dotted_version("latest"), None);
729        assert_eq!(parse_dotted_version(""), None);
730    }
731
732    #[test]
733    fn oneapi_version_renders_with_patch() {
734        assert_eq!(OneApiVersion::new(2024, 2, 1).to_string(), "2024.2.1");
735        assert_eq!(OneApiVersion::new(2025, 0, 0).to_string(), "2025.0.0");
736    }
737
738    #[test]
739    fn oneapi_host_is_stable_across_calls() {
740        // Environment-dependent: most hosts have no oneAPI, which is a valid,
741        // passing environment. A filesystem read must not vary between calls.
742        assert_eq!(oneapi_host(), oneapi_host());
743    }
744
745    #[test]
746    fn rocm_version_renders_with_patch() {
747        assert_eq!(RocmVersion::new(6, 2, 4).to_string(), "6.2.4");
748        assert_eq!(RocmVersion::new(6, 2, 0).to_string(), "6.2.0");
749    }
750
751    #[test]
752    fn rocm_host_is_stable_across_calls() {
753        // Environment-dependent: most hosts have no ROCm, which is a valid,
754        // passing environment. A filesystem read must not vary between calls.
755        assert_eq!(rocm_host(), rocm_host());
756    }
757
758    #[test]
759    fn vulkan_version_renders_with_patch() {
760        assert_eq!(VulkanVersion::new(1, 3, 280).to_string(), "1.3.280");
761        assert_eq!(VulkanVersion::new(1, 0, 0).to_string(), "1.0.0");
762        // A three-digit patch is the norm for Vulkan headers, and must not be
763        // packed or truncated the way `1.3` alone would be.
764        assert_eq!(VulkanVersion::new(1, 10, 5).to_string(), "1.10.5");
765    }
766
767    #[test]
768    fn vulkan_host_is_stable_across_calls() {
769        // Environment-dependent: a host with no loader is a valid, passing
770        // environment. A filesystem read must not vary between calls.
771        assert_eq!(vulkan_host(), vulkan_host());
772    }
773
774    #[test]
775    fn vulkan_host_carries_only_an_api_version() {
776        // No `arch_target` counterpart, deliberately: SPIR-V is portable and
777        // driver-compiled, so there is no per-GPU target to match. This pins
778        // that shape, and the `Copy`/`Eq` derives callers rely on.
779        let host = VulkanHost {
780            api_version: VulkanVersion::new(1, 3, 280),
781        };
782        let copied = host;
783        assert_eq!(copied, host);
784        assert_eq!(copied.api_version, VulkanVersion::new(1, 3, 280));
785        assert_ne!(
786            host,
787            VulkanHost {
788                api_version: VulkanVersion::new(1, 2, 0)
789            }
790        );
791    }
792
793    #[test]
794    fn detect_never_panics() {
795        // Environment-dependent (may be empty on headless CI); exercise the
796        // full path plus the Display impl without asserting a GPU exists.
797        for gpu in detect() {
798            assert!(!gpu.name.is_empty());
799            let _ = gpu.to_string();
800        }
801    }
802
803    #[test]
804    fn display_includes_free_when_present() {
805        let gpu = GpuInfo {
806            name: "Test GPU".to_string(),
807            vendor: Vendor::Nvidia,
808            total_bytes: 24 * 1024 * 1024 * 1024,
809            free_bytes: Some(12 * 1024 * 1024 * 1024),
810            used_bytes: Some(12 * 1024 * 1024 * 1024),
811            arch_target: None,
812        };
813        let shown = gpu.to_string();
814        assert!(shown.contains("NVIDIA"));
815        assert!(shown.contains("24.0 GiB total"));
816        assert!(shown.contains("12.0 GiB free"));
817    }
818
819    #[test]
820    fn display_omits_free_when_absent() {
821        let gpu = GpuInfo {
822            name: "AMD GPU (card0)".to_string(),
823            vendor: Vendor::Amd,
824            total_bytes: 8 * 1024 * 1024 * 1024,
825            free_bytes: None,
826            used_bytes: None,
827            arch_target: None,
828        };
829        let shown = gpu.to_string();
830        assert!(shown.contains("8.0 GiB total"));
831        assert!(!shown.contains("free"));
832    }
833
834    #[test]
835    fn display_includes_gfx_target_when_present() {
836        let gpu = GpuInfo {
837            name: "AMD cyan_skillfish".to_string(),
838            vendor: Vendor::Amd,
839            total_bytes: 15 * 1024 * 1024 * 1024,
840            free_bytes: None,
841            used_bytes: None,
842            arch_target: Some(ArchTarget::Gfx(GfxTarget::new(10, 1, 3))),
843        };
844        assert!(gpu.to_string().contains("(AMD, gfx1013)"));
845    }
846
847    #[test]
848    fn display_includes_compute_capability_when_present() {
849        let gpu = GpuInfo {
850            name: "NVIDIA GeForce RTX 4090".to_string(),
851            vendor: Vendor::Nvidia,
852            total_bytes: 24 * 1024 * 1024 * 1024,
853            free_bytes: None,
854            used_bytes: None,
855            arch_target: Some(ArchTarget::Sm(ComputeCapability::new(8, 9))),
856        };
857        assert!(gpu.to_string().contains("(NVIDIA, sm_89)"));
858    }
859
860    #[test]
861    fn arch_target_unwraps_only_its_own_vendor() {
862        let amd = ArchTarget::Gfx(GfxTarget::new(10, 1, 3));
863        assert_eq!(amd.gfx(), Some(GfxTarget::new(10, 1, 3)));
864        assert_eq!(amd.sm(), None);
865
866        let nvidia = ArchTarget::Sm(ComputeCapability::new(8, 9));
867        assert_eq!(nvidia.sm(), Some(ComputeCapability::new(8, 9)));
868        assert_eq!(nvidia.gfx(), None);
869    }
870
871    #[test]
872    fn arch_target_accessors_are_exclusive_across_all_vendors() {
873        let targets = [
874            ArchTarget::Gfx(GfxTarget::new(10, 1, 3)),
875            ArchTarget::Sm(ComputeCapability::new(8, 9)),
876            ArchTarget::Xe(IntelArch::XeHpg),
877            ArchTarget::Apple(AppleFamily::new(8)),
878        ];
879        for target in targets {
880            let hits = [
881                target.gfx().is_some(),
882                target.sm().is_some(),
883                target.xe().is_some(),
884                target.apple().is_some(),
885            ];
886            assert_eq!(
887                hits.iter().filter(|hit| **hit).count(),
888                1,
889                "{target} must answer exactly one accessor",
890            );
891        }
892    }
893
894    #[test]
895    fn intel_and_apple_targets_render_by_family() {
896        assert_eq!(ArchTarget::Xe(IntelArch::XeHpg).to_string(), "xe-hpg");
897        assert_eq!(ArchTarget::Xe(IntelArch::Xe2).to_string(), "xe2");
898        assert_eq!(ArchTarget::Apple(AppleFamily::new(8)).to_string(), "apple8");
899    }
900
901    #[test]
902    fn apple_families_are_ordered() {
903        assert!(AppleFamily::new(9) > AppleFamily::new(8));
904        assert!(AppleFamily::new(8) > AppleFamily::new(7));
905    }
906
907    #[test]
908    fn arch_target_renders_per_vendor_toolchain() {
909        // `sm_89`, not the bare `8.9` `ComputeCapability` renders on its own,
910        // which would read as a version number in this position.
911        assert_eq!(
912            ArchTarget::Sm(ComputeCapability::new(8, 9)).to_string(),
913            "sm_89"
914        );
915        assert_eq!(
916            ArchTarget::Gfx(GfxTarget::new(10, 1, 3)).to_string(),
917            "gfx1013"
918        );
919    }
920
921    #[test]
922    fn gfx_target_renders_as_offload_arch() {
923        assert_eq!(GfxTarget::new(10, 1, 3).to_string(), "gfx1013");
924        assert_eq!(GfxTarget::new(10, 3, 0).to_string(), "gfx1030");
925        assert_eq!(GfxTarget::new(11, 0, 0).to_string(), "gfx1100");
926        // Stepping 10 is the `a` in `gfx90a`, not a literal "10".
927        assert_eq!(GfxTarget::new(9, 0, 10).to_string(), "gfx90a");
928        assert_eq!(GfxTarget::new(9, 4, 2).to_string(), "gfx942");
929    }
930
931    #[test]
932    fn gfx_targets_order_major_first() {
933        assert!(GfxTarget::new(11, 0, 0) > GfxTarget::new(10, 3, 0));
934        assert!(GfxTarget::new(10, 3, 0) > GfxTarget::new(10, 1, 3));
935        assert!(GfxTarget::new(10, 1, 3) > GfxTarget::new(10, 1, 0));
936    }
937
938    #[test]
939    fn vendor_display_covers_every_variant() {
940        assert_eq!(Vendor::Nvidia.to_string(), "NVIDIA");
941        assert_eq!(Vendor::Amd.to_string(), "AMD");
942        assert_eq!(Vendor::Intel.to_string(), "Intel");
943        assert_eq!(Vendor::Apple.to_string(), "Apple");
944        assert_eq!(Vendor::Unknown.to_string(), "Unknown");
945    }
946
947    #[test]
948    fn gib_converts_using_binary_units() {
949        assert!((gib(0) - 0.0).abs() < f64::EPSILON);
950        assert!((gib(1024 * 1024 * 1024) - 1.0).abs() < f64::EPSILON);
951        // 1.5 GiB exercises the fractional path the Display rounds to one place.
952        assert!((gib(3 * 1024 * 1024 * 1024 / 2) - 1.5).abs() < f64::EPSILON);
953    }
954
955    #[test]
956    fn display_rounds_to_one_decimal_place() {
957        // 25 GiB + 256 MiB -> 25.25 GiB, which "{:.1}" renders as "25.2".
958        let gpu = GpuInfo {
959            name: "Rounding".to_string(),
960            vendor: Vendor::Nvidia,
961            total_bytes: 25 * 1024 * 1024 * 1024 + 256 * 1024 * 1024,
962            free_bytes: None,
963            used_bytes: None,
964            arch_target: None,
965        };
966        assert!(gpu.to_string().contains("25.2 GiB total"));
967    }
968
969    #[test]
970    fn detect_results_have_consistent_memory_fields() {
971        // Environment-dependent; asserts invariants only for whatever is present.
972        for gpu in detect() {
973            assert!(!gpu.name.is_empty());
974            if let Some(free) = gpu.free_bytes {
975                assert!(free <= gpu.total_bytes, "free must not exceed total");
976            }
977            if let (Some(free), Some(used)) = (gpu.free_bytes, gpu.used_bytes) {
978                assert!(
979                    free.saturating_add(used) <= gpu.total_bytes.saturating_add(used),
980                    "free/used must be coherent",
981                );
982            }
983        }
984    }
985
986    #[test]
987    fn versions_display_as_major_dot_minor() {
988        assert_eq!(ComputeCapability::new(8, 6).to_string(), "8.6");
989        assert_eq!(CudaVersion::new(12, 9).to_string(), "12.9");
990        // A two-digit minor stays unambiguous — the reason these aren't packed
991        // into a single integer.
992        assert_eq!(ComputeCapability::new(8, 10).to_string(), "8.10");
993    }
994
995    #[test]
996    fn versions_order_by_major_then_minor() {
997        assert!(ComputeCapability::new(8, 6) > ComputeCapability::new(8, 0));
998        assert!(ComputeCapability::new(9, 0) > ComputeCapability::new(8, 9));
999        assert_eq!(ComputeCapability::new(8, 6), ComputeCapability::new(8, 6));
1000        assert!(CudaVersion::new(12, 9) > CudaVersion::new(12, 0));
1001        assert!(CudaVersion::new(13, 0) > CudaVersion::new(12, 9));
1002        // Packing as `major * 10 + minor` would collide here: 8.10 and 9.0
1003        // both pack to 90, which is why the parts are kept separate.
1004        assert!(ComputeCapability::new(9, 0) > ComputeCapability::new(8, 10));
1005    }
1006
1007    #[test]
1008    fn cuda_host_is_environment_dependent_but_coherent() {
1009        // No NVIDIA driver is a valid, passing environment.
1010        if let Some(cuda) = cuda_host() {
1011            assert!(
1012                cuda.compute_capability.major > 0,
1013                "a real device has a nonzero major capability",
1014            );
1015            assert!(cuda.driver_version.major > 0, "a real driver has a version");
1016            assert_eq!(
1017                cuda_host(),
1018                Some(cuda),
1019                "host/driver properties must be stable across calls",
1020            );
1021        }
1022    }
1023
1024    #[test]
1025    fn gpu_info_equality_compares_all_fields() {
1026        let base = GpuInfo {
1027            name: "G".to_string(),
1028            vendor: Vendor::Intel,
1029            total_bytes: 16 * 1024 * 1024 * 1024,
1030            free_bytes: None,
1031            used_bytes: None,
1032            arch_target: None,
1033        };
1034        assert_eq!(base.clone(), base);
1035        let mut other = base.clone();
1036        other.vendor = Vendor::Amd;
1037        assert_ne!(base, other);
1038    }
1039}