gpu-probe 0.1.0

// SPDX-License-Identifier: Apache-2.0
//! Cross-platform GPU memory (VRAM) detection with **no vendor SDKs**.
//!
//! `gpu_probe` reports the GPUs visible on the host and how much memory each
//! has, using only facilities the OS or driver already ship:
//!
//! - **NVIDIA** (Linux, Windows): NVML (`libnvidia-ml`) via `nvml-wrapper`,
//!   loaded at runtime. The CUDA toolkit is not required and nothing links at
//!   build time. Behind the default `nvidia` feature.
//! - **AMD & Intel** (Linux): DRM sysfs under `/sys/class/drm`. Discrete cards
//!   report dedicated VRAM; integrated GPUs report the shared system-memory
//!   ceiling (see [`GpuInfo::total_bytes`]).
//! - **Apple/macOS**: `system_profiler` + `sysctl` (Apple Silicon reports
//!   unified memory).
//!
//! Detection is best-effort: [`detect`] returns an empty `Vec` when no GPU is
//! found or the platform is unsupported — never an error.
//!
//! ```no_run
//! for gpu in gpu_probe::detect() {
//!     println!("{gpu}");
//! }
//! ```

mod drm;
mod metal;
mod nvidia;

/// GPU hardware vendor.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub enum Vendor {
    Nvidia,
    Amd,
    Intel,
    Apple,
    Unknown,
}

impl std::fmt::Display for Vendor {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str(match self {
            Vendor::Nvidia => "NVIDIA",
            Vendor::Amd => "AMD",
            Vendor::Intel => "Intel",
            Vendor::Apple => "Apple",
            Vendor::Unknown => "Unknown",
        })
    }
}

/// A single detected GPU and its memory.
#[derive(Debug, Clone, PartialEq, Eq)]
#[non_exhaustive]
pub struct GpuInfo {
    /// Human-readable name (e.g. `"NVIDIA GeForce RTX 4090"`).
    pub name: String,
    /// Hardware vendor.
    pub vendor: Vendor,
    /// Total memory in bytes. For discrete GPUs this is dedicated VRAM; for
    /// integrated/unified GPUs (Intel iGPUs, AMD APUs, Apple Silicon) it is the
    /// shared system-memory ceiling available to the GPU, not a dedicated pool.
    pub total_bytes: u64,
    /// Free device memory in bytes, when known.
    pub free_bytes: Option<u64>,
    /// Used device memory in bytes, when known.
    pub used_bytes: Option<u64>,
}

impl std::fmt::Display for GpuInfo {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "{} ({}): {:.1} GiB total",
            self.name,
            self.vendor,
            gib(self.total_bytes)
        )?;
        if let Some(free) = self.free_bytes {
            write!(f, ", {:.1} GiB free", gib(free))?;
        }
        Ok(())
    }
}

#[allow(clippy::cast_precision_loss)] // display-only; the imprecision is cosmetic
fn gib(bytes: u64) -> f64 {
    bytes as f64 / (1024.0 * 1024.0 * 1024.0)
}

/// CUDA compute capability, e.g. `8.6` for `sm_86`.
///
/// Ordered `major` first, so a host can be checked against a minimum:
///
/// ```
/// use gpu_probe::ComputeCapability;
/// assert!(ComputeCapability::new(8, 6) >= ComputeCapability::new(8, 0));
/// assert!(ComputeCapability::new(9, 0) >= ComputeCapability::new(8, 9));
/// ```
///
/// Constructible so callers can express such a requirement.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct ComputeCapability {
    /// Major version — the `8` in `8.6`.
    pub major: u32,
    /// Minor version — the `6` in `8.6`.
    pub minor: u32,
}

impl ComputeCapability {
    /// Create a compute capability from its major and minor parts.
    #[must_use]
    pub const fn new(major: u32, minor: u32) -> Self {
        Self { major, minor }
    }
}

impl std::fmt::Display for ComputeCapability {
    /// Renders as `8.6`, matching `nvidia-smi`'s `compute_cap`.
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}.{}", self.major, self.minor)
    }
}

/// A CUDA version, e.g. `12.9`.
///
/// Ordered `major` first, so a host can be checked against a minimum:
///
/// ```
/// use gpu_probe::CudaVersion;
/// assert!(CudaVersion::new(12, 9) >= CudaVersion::new(12, 0));
/// ```
///
/// Constructible so callers can express such a requirement.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct CudaVersion {
    /// Major version — the `12` in `12.9`.
    pub major: u32,
    /// Minor version — the `9` in `12.9`.
    pub minor: u32,
}

impl CudaVersion {
    /// Create a CUDA version from its major and minor parts.
    #[must_use]
    pub const fn new(major: u32, minor: u32) -> Self {
        Self { major, minor }
    }
}

impl std::fmt::Display for CudaVersion {
    /// Renders as `12.9`.
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}.{}", self.major, self.minor)
    }
}

/// Host-wide CUDA properties reported by the NVIDIA driver.
///
/// These describe the host and its driver rather than any one GPU, which is why
/// they are separate from the per-GPU [`GpuInfo`]. Consumers typically use them
/// to select a prebuilt artifact compatible with the host.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub struct CudaHost {
    /// Compute capability of device 0.
    pub compute_capability: ComputeCapability,
    /// Version of the installed CUDA driver.
    pub driver_version: CudaVersion,
}

/// Detect all GPUs visible on the host.
///
/// Best-effort: spawns only read-only platform queries (NVML, `system_profiler`,
/// `sysctl`) and reads sysfs. Returns an empty `Vec` on unsupported platforms
/// or when no GPU is found.
#[must_use]
pub fn detect() -> Vec<GpuInfo> {
    let mut gpus = Vec::new();
    gpus.extend(nvidia::detect());
    gpus.extend(drm::detect());
    gpus.extend(metal::detect());
    gpus
}

/// Host-wide CUDA properties, or `None` when NVML is unavailable — no NVIDIA
/// driver, the `nvidia` feature disabled, no device, or a driver reporting
/// values that aren't usable.
///
/// Shares the one process-wide NVML handle with [`detect`], so calling this on
/// a timer does not accumulate resources.
///
/// ```no_run
/// use gpu_probe::ComputeCapability;
///
/// if let Some(cuda) = gpu_probe::cuda_host() {
///     println!("sm_{}{} on CUDA {}",
///         cuda.compute_capability.major,
///         cuda.compute_capability.minor,
///         cuda.driver_version);
///
///     if cuda.compute_capability >= ComputeCapability::new(8, 0) {
///         // pick an Ampere-or-newer build
///     }
/// }
/// ```
#[must_use]
pub fn cuda_host() -> Option<CudaHost> {
    nvidia::cuda_host()
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn detect_never_panics() {
        // Environment-dependent (may be empty on headless CI); exercise the
        // full path plus the Display impl without asserting a GPU exists.
        for gpu in detect() {
            assert!(!gpu.name.is_empty());
            let _ = gpu.to_string();
        }
    }

    #[test]
    fn display_includes_free_when_present() {
        let gpu = GpuInfo {
            name: "Test GPU".to_string(),
            vendor: Vendor::Nvidia,
            total_bytes: 24 * 1024 * 1024 * 1024,
            free_bytes: Some(12 * 1024 * 1024 * 1024),
            used_bytes: Some(12 * 1024 * 1024 * 1024),
        };
        let shown = gpu.to_string();
        assert!(shown.contains("NVIDIA"));
        assert!(shown.contains("24.0 GiB total"));
        assert!(shown.contains("12.0 GiB free"));
    }

    #[test]
    fn display_omits_free_when_absent() {
        let gpu = GpuInfo {
            name: "AMD GPU (card0)".to_string(),
            vendor: Vendor::Amd,
            total_bytes: 8 * 1024 * 1024 * 1024,
            free_bytes: None,
            used_bytes: None,
        };
        let shown = gpu.to_string();
        assert!(shown.contains("8.0 GiB total"));
        assert!(!shown.contains("free"));
    }

    #[test]
    fn vendor_display_covers_every_variant() {
        assert_eq!(Vendor::Nvidia.to_string(), "NVIDIA");
        assert_eq!(Vendor::Amd.to_string(), "AMD");
        assert_eq!(Vendor::Intel.to_string(), "Intel");
        assert_eq!(Vendor::Apple.to_string(), "Apple");
        assert_eq!(Vendor::Unknown.to_string(), "Unknown");
    }

    #[test]
    fn gib_converts_using_binary_units() {
        assert!((gib(0) - 0.0).abs() < f64::EPSILON);
        assert!((gib(1024 * 1024 * 1024) - 1.0).abs() < f64::EPSILON);
        // 1.5 GiB exercises the fractional path the Display rounds to one place.
        assert!((gib(3 * 1024 * 1024 * 1024 / 2) - 1.5).abs() < f64::EPSILON);
    }

    #[test]
    fn display_rounds_to_one_decimal_place() {
        // 25 GiB + 256 MiB -> 25.25 GiB, which "{:.1}" renders as "25.2".
        let gpu = GpuInfo {
            name: "Rounding".to_string(),
            vendor: Vendor::Nvidia,
            total_bytes: 25 * 1024 * 1024 * 1024 + 256 * 1024 * 1024,
            free_bytes: None,
            used_bytes: None,
        };
        assert!(gpu.to_string().contains("25.2 GiB total"));
    }

    #[test]
    fn detect_results_have_consistent_memory_fields() {
        // Environment-dependent; asserts invariants only for whatever is present.
        for gpu in detect() {
            assert!(!gpu.name.is_empty());
            if let Some(free) = gpu.free_bytes {
                assert!(free <= gpu.total_bytes, "free must not exceed total");
            }
            if let (Some(free), Some(used)) = (gpu.free_bytes, gpu.used_bytes) {
                assert!(
                    free.saturating_add(used) <= gpu.total_bytes.saturating_add(used),
                    "free/used must be coherent",
                );
            }
        }
    }

    #[test]
    fn versions_display_as_major_dot_minor() {
        assert_eq!(ComputeCapability::new(8, 6).to_string(), "8.6");
        assert_eq!(CudaVersion::new(12, 9).to_string(), "12.9");
        // A two-digit minor stays unambiguous — the reason these aren't packed
        // into a single integer.
        assert_eq!(ComputeCapability::new(8, 10).to_string(), "8.10");
    }

    #[test]
    fn versions_order_by_major_then_minor() {
        assert!(ComputeCapability::new(8, 6) > ComputeCapability::new(8, 0));
        assert!(ComputeCapability::new(9, 0) > ComputeCapability::new(8, 9));
        assert_eq!(ComputeCapability::new(8, 6), ComputeCapability::new(8, 6));
        assert!(CudaVersion::new(12, 9) > CudaVersion::new(12, 0));
        assert!(CudaVersion::new(13, 0) > CudaVersion::new(12, 9));
        // Packing as `major * 10 + minor` would collide here: 8.10 and 9.0
        // both pack to 90, which is why the parts are kept separate.
        assert!(ComputeCapability::new(9, 0) > ComputeCapability::new(8, 10));
    }

    #[test]
    fn cuda_host_is_environment_dependent_but_coherent() {
        // No NVIDIA driver is a valid, passing environment.
        if let Some(cuda) = cuda_host() {
            assert!(
                cuda.compute_capability.major > 0,
                "a real device has a nonzero major capability",
            );
            assert!(cuda.driver_version.major > 0, "a real driver has a version");
            assert_eq!(
                cuda_host(),
                Some(cuda),
                "host/driver properties must be stable across calls",
            );
        }
    }

    #[test]
    fn gpu_info_equality_compares_all_fields() {
        let base = GpuInfo {
            name: "G".to_string(),
            vendor: Vendor::Intel,
            total_bytes: 16 * 1024 * 1024 * 1024,
            free_bytes: None,
            used_bytes: None,
        };
        assert_eq!(base.clone(), base);
        let mut other = base.clone();
        other.vendor = Vendor::Amd;
        assert_ne!(base, other);
    }
}