mod drm;
mod intel;
mod kfd;
mod metal;
mod nvidia;
mod oneapi;
mod rocm;
mod vulkan;
#[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",
})
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
#[non_exhaustive]
pub struct GpuInfo {
pub name: String,
pub vendor: Vendor,
pub total_bytes: u64,
pub free_bytes: Option<u64>,
pub used_bytes: Option<u64>,
pub arch_target: Option<ArchTarget>,
}
impl std::fmt::Display for GpuInfo {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{} ({}", self.name, self.vendor)?;
if let Some(arch) = self.arch_target {
write!(f, ", {arch}")?;
}
write!(f, "): {:.1} GiB total", gib(self.total_bytes))?;
if let Some(free) = self.free_bytes {
write!(f, ", {:.1} GiB free", gib(free))?;
}
Ok(())
}
}
#[allow(clippy::cast_precision_loss)] fn gib(bytes: u64) -> f64 {
bytes as f64 / (1024.0 * 1024.0 * 1024.0)
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct ComputeCapability {
pub major: u32,
pub minor: u32,
}
impl ComputeCapability {
#[must_use]
pub const fn new(major: u32, minor: u32) -> Self {
Self { major, minor }
}
}
impl std::fmt::Display for ComputeCapability {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}.{}", self.major, self.minor)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct GfxTarget {
pub major: u32,
pub minor: u32,
pub step: u32,
}
impl GfxTarget {
#[must_use]
pub const fn new(major: u32, minor: u32, step: u32) -> Self {
Self { major, minor, step }
}
}
impl std::fmt::Display for GfxTarget {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "gfx{}{:x}{:x}", self.major, self.minor, self.step)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
#[non_exhaustive]
pub enum IntelArch {
XeLp,
XeHpg,
XeHpc,
XeLpg,
Xe2,
}
impl std::fmt::Display for IntelArch {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str(match self {
IntelArch::XeLp => "xe-lp",
IntelArch::XeHpg => "xe-hpg",
IntelArch::XeHpc => "xe-hpc",
IntelArch::XeLpg => "xe-lpg",
IntelArch::Xe2 => "xe2",
})
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct AppleFamily {
pub generation: u32,
}
impl AppleFamily {
#[must_use]
pub const fn new(generation: u32) -> Self {
Self { generation }
}
}
impl std::fmt::Display for AppleFamily {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "apple{}", self.generation)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
#[non_exhaustive]
pub enum ArchTarget {
Gfx(GfxTarget),
Sm(ComputeCapability),
Xe(IntelArch),
Apple(AppleFamily),
}
impl ArchTarget {
#[must_use]
pub const fn gfx(self) -> Option<GfxTarget> {
match self {
Self::Gfx(target) => Some(target),
_ => None,
}
}
#[must_use]
pub const fn sm(self) -> Option<ComputeCapability> {
match self {
Self::Sm(capability) => Some(capability),
_ => None,
}
}
#[must_use]
pub const fn xe(self) -> Option<IntelArch> {
match self {
Self::Xe(arch) => Some(arch),
_ => None,
}
}
#[must_use]
pub const fn apple(self) -> Option<AppleFamily> {
match self {
Self::Apple(family) => Some(family),
_ => None,
}
}
}
impl std::fmt::Display for ArchTarget {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Gfx(target) => write!(f, "{target}"),
Self::Sm(capability) => {
write!(f, "sm_{}{}", capability.major, capability.minor)
}
Self::Xe(arch) => write!(f, "{arch}"),
Self::Apple(family) => write!(f, "{family}"),
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct CudaVersion {
pub major: u32,
pub minor: u32,
}
impl CudaVersion {
#[must_use]
pub const fn new(major: u32, minor: u32) -> Self {
Self { major, minor }
}
}
impl std::fmt::Display for CudaVersion {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}.{}", self.major, self.minor)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub struct CudaHost {
pub compute_capability: ComputeCapability,
pub driver_version: CudaVersion,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct RocmVersion {
pub major: u32,
pub minor: u32,
pub patch: u32,
}
impl RocmVersion {
#[must_use]
pub const fn new(major: u32, minor: u32, patch: u32) -> Self {
Self {
major,
minor,
patch,
}
}
}
impl std::fmt::Display for RocmVersion {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}.{}.{}", self.major, self.minor, self.patch)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub struct RocmHost {
pub version: RocmVersion,
}
fn parse_dotted_version(text: &str) -> Option<(u32, u32, u32)> {
let version = text.trim().split(['-', '+']).next()?;
let mut parts = version.split('.');
let major = parts.next()?.trim().parse().ok()?;
let minor = parts.next()?.trim().parse().ok()?;
let patch = match parts.next() {
Some(patch) => patch.trim().parse().ok()?,
None => 0,
};
Some((major, minor, patch))
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct OneApiVersion {
pub major: u32,
pub minor: u32,
pub patch: u32,
}
impl OneApiVersion {
#[must_use]
pub const fn new(major: u32, minor: u32, patch: u32) -> Self {
Self {
major,
minor,
patch,
}
}
}
impl std::fmt::Display for OneApiVersion {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}.{}.{}", self.major, self.minor, self.patch)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub struct OneApiHost {
pub version: OneApiVersion,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct VulkanVersion {
pub major: u32,
pub minor: u32,
pub patch: u32,
}
impl VulkanVersion {
#[must_use]
pub const fn new(major: u32, minor: u32, patch: u32) -> Self {
Self {
major,
minor,
patch,
}
}
}
impl std::fmt::Display for VulkanVersion {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}.{}.{}", self.major, self.minor, self.patch)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub struct VulkanHost {
pub api_version: VulkanVersion,
}
#[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
}
#[must_use]
pub fn cuda_host() -> Option<CudaHost> {
nvidia::cuda_host()
}
#[must_use]
pub fn rocm_host() -> Option<RocmHost> {
rocm::host()
}
#[must_use]
pub fn oneapi_host() -> Option<OneApiHost> {
oneapi::host()
}
#[must_use]
pub fn vulkan_host() -> Option<VulkanHost> {
vulkan::host()
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn parses_dotted_versions_in_both_shapes() {
assert_eq!(parse_dotted_version("6.2.4-123"), Some((6, 2, 4)));
assert_eq!(parse_dotted_version("2024.2"), Some((2024, 2, 0)));
assert_eq!(parse_dotted_version(" 5.7.1 "), Some((5, 7, 1)));
assert_eq!(parse_dotted_version("6"), None, "a bare major is not one");
assert_eq!(parse_dotted_version("latest"), None);
assert_eq!(parse_dotted_version(""), None);
}
#[test]
fn oneapi_version_renders_with_patch() {
assert_eq!(OneApiVersion::new(2024, 2, 1).to_string(), "2024.2.1");
assert_eq!(OneApiVersion::new(2025, 0, 0).to_string(), "2025.0.0");
}
#[test]
fn oneapi_host_is_stable_across_calls() {
assert_eq!(oneapi_host(), oneapi_host());
}
#[test]
fn rocm_version_renders_with_patch() {
assert_eq!(RocmVersion::new(6, 2, 4).to_string(), "6.2.4");
assert_eq!(RocmVersion::new(6, 2, 0).to_string(), "6.2.0");
}
#[test]
fn rocm_host_is_stable_across_calls() {
assert_eq!(rocm_host(), rocm_host());
}
#[test]
fn vulkan_version_renders_with_patch() {
assert_eq!(VulkanVersion::new(1, 3, 280).to_string(), "1.3.280");
assert_eq!(VulkanVersion::new(1, 0, 0).to_string(), "1.0.0");
assert_eq!(VulkanVersion::new(1, 10, 5).to_string(), "1.10.5");
}
#[test]
fn vulkan_host_is_stable_across_calls() {
assert_eq!(vulkan_host(), vulkan_host());
}
#[test]
fn vulkan_host_carries_only_an_api_version() {
let host = VulkanHost {
api_version: VulkanVersion::new(1, 3, 280),
};
let copied = host;
assert_eq!(copied, host);
assert_eq!(copied.api_version, VulkanVersion::new(1, 3, 280));
assert_ne!(
host,
VulkanHost {
api_version: VulkanVersion::new(1, 2, 0)
}
);
}
#[test]
fn detect_never_panics() {
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),
arch_target: None,
};
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,
arch_target: None,
};
let shown = gpu.to_string();
assert!(shown.contains("8.0 GiB total"));
assert!(!shown.contains("free"));
}
#[test]
fn display_includes_gfx_target_when_present() {
let gpu = GpuInfo {
name: "AMD cyan_skillfish".to_string(),
vendor: Vendor::Amd,
total_bytes: 15 * 1024 * 1024 * 1024,
free_bytes: None,
used_bytes: None,
arch_target: Some(ArchTarget::Gfx(GfxTarget::new(10, 1, 3))),
};
assert!(gpu.to_string().contains("(AMD, gfx1013)"));
}
#[test]
fn display_includes_compute_capability_when_present() {
let gpu = GpuInfo {
name: "NVIDIA GeForce RTX 4090".to_string(),
vendor: Vendor::Nvidia,
total_bytes: 24 * 1024 * 1024 * 1024,
free_bytes: None,
used_bytes: None,
arch_target: Some(ArchTarget::Sm(ComputeCapability::new(8, 9))),
};
assert!(gpu.to_string().contains("(NVIDIA, sm_89)"));
}
#[test]
fn arch_target_unwraps_only_its_own_vendor() {
let amd = ArchTarget::Gfx(GfxTarget::new(10, 1, 3));
assert_eq!(amd.gfx(), Some(GfxTarget::new(10, 1, 3)));
assert_eq!(amd.sm(), None);
let nvidia = ArchTarget::Sm(ComputeCapability::new(8, 9));
assert_eq!(nvidia.sm(), Some(ComputeCapability::new(8, 9)));
assert_eq!(nvidia.gfx(), None);
}
#[test]
fn arch_target_accessors_are_exclusive_across_all_vendors() {
let targets = [
ArchTarget::Gfx(GfxTarget::new(10, 1, 3)),
ArchTarget::Sm(ComputeCapability::new(8, 9)),
ArchTarget::Xe(IntelArch::XeHpg),
ArchTarget::Apple(AppleFamily::new(8)),
];
for target in targets {
let hits = [
target.gfx().is_some(),
target.sm().is_some(),
target.xe().is_some(),
target.apple().is_some(),
];
assert_eq!(
hits.iter().filter(|hit| **hit).count(),
1,
"{target} must answer exactly one accessor",
);
}
}
#[test]
fn intel_and_apple_targets_render_by_family() {
assert_eq!(ArchTarget::Xe(IntelArch::XeHpg).to_string(), "xe-hpg");
assert_eq!(ArchTarget::Xe(IntelArch::Xe2).to_string(), "xe2");
assert_eq!(ArchTarget::Apple(AppleFamily::new(8)).to_string(), "apple8");
}
#[test]
fn apple_families_are_ordered() {
assert!(AppleFamily::new(9) > AppleFamily::new(8));
assert!(AppleFamily::new(8) > AppleFamily::new(7));
}
#[test]
fn arch_target_renders_per_vendor_toolchain() {
assert_eq!(
ArchTarget::Sm(ComputeCapability::new(8, 9)).to_string(),
"sm_89"
);
assert_eq!(
ArchTarget::Gfx(GfxTarget::new(10, 1, 3)).to_string(),
"gfx1013"
);
}
#[test]
fn gfx_target_renders_as_offload_arch() {
assert_eq!(GfxTarget::new(10, 1, 3).to_string(), "gfx1013");
assert_eq!(GfxTarget::new(10, 3, 0).to_string(), "gfx1030");
assert_eq!(GfxTarget::new(11, 0, 0).to_string(), "gfx1100");
assert_eq!(GfxTarget::new(9, 0, 10).to_string(), "gfx90a");
assert_eq!(GfxTarget::new(9, 4, 2).to_string(), "gfx942");
}
#[test]
fn gfx_targets_order_major_first() {
assert!(GfxTarget::new(11, 0, 0) > GfxTarget::new(10, 3, 0));
assert!(GfxTarget::new(10, 3, 0) > GfxTarget::new(10, 1, 3));
assert!(GfxTarget::new(10, 1, 3) > GfxTarget::new(10, 1, 0));
}
#[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);
assert!((gib(3 * 1024 * 1024 * 1024 / 2) - 1.5).abs() < f64::EPSILON);
}
#[test]
fn display_rounds_to_one_decimal_place() {
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,
arch_target: None,
};
assert!(gpu.to_string().contains("25.2 GiB total"));
}
#[test]
fn detect_results_have_consistent_memory_fields() {
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");
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));
assert!(ComputeCapability::new(9, 0) > ComputeCapability::new(8, 10));
}
#[test]
fn cuda_host_is_environment_dependent_but_coherent() {
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,
arch_target: None,
};
assert_eq!(base.clone(), base);
let mut other = base.clone();
other.vendor = Vendor::Amd;
assert_ne!(base, other);
}
}