Skip to main content

probe/
machine.rs

1use crate::accel::{Accel, probe_accel};
2use crate::gpu::{DeviceType, probe_gpu};
3
4#[derive(Debug, Clone, Copy, PartialEq, Eq)]
5pub enum Os {
6    Linux,
7    Windows,
8    MacOs,
9    Other,
10}
11
12impl Os {
13    pub const fn label(self) -> &'static str {
14        match self {
15            Os::Linux => "Linux",
16            Os::Windows => "Windows",
17            Os::MacOs => "macOS",
18            Os::Other => "This computer",
19        }
20    }
21}
22
23#[derive(Debug, Clone, Copy, PartialEq, Eq)]
24pub enum Vendor {
25    Nvidia,
26    Amd,
27    Intel,
28    Apple,
29    None,
30}
31
32impl Vendor {
33    pub const fn label(self) -> &'static str {
34        match self {
35            Vendor::Nvidia => "NVIDIA",
36            Vendor::Amd => "AMD",
37            Vendor::Intel => "Intel",
38            Vendor::Apple => "Apple Silicon",
39            Vendor::None => "no graphics card",
40        }
41    }
42}
43
44#[derive(Debug, Clone, Copy, PartialEq, Eq)]
45pub enum Lane {
46    Cuda,
47    Vulkan,
48    Metal,
49    DirectMl,
50    WebGpu,
51    CoreMl,
52    VideoToolbox,
53    Nvidia,
54    Vaapi,
55    Qsv,
56    Amf,
57    Processor,
58}
59
60impl Lane {
61    pub const fn label(self) -> &'static str {
62        match self {
63            Lane::Cuda => "CUDA",
64            Lane::Vulkan => "Vulkan",
65            Lane::Metal => "Metal",
66            Lane::DirectMl => "DirectML",
67            Lane::WebGpu => "WebGPU",
68            Lane::CoreMl => "Core ML",
69            Lane::VideoToolbox => "VideoToolbox",
70            Lane::Nvidia => "NVDEC and NVENC",
71            Lane::Vaapi => "VAAPI",
72            Lane::Qsv => "Quick Sync",
73            Lane::Amf => "AMF",
74            Lane::Processor => "the processor",
75        }
76    }
77
78    pub const fn accelerated(self) -> bool {
79        !matches!(self, Lane::Processor)
80    }
81}
82
83#[derive(Debug, Clone, PartialEq, Eq)]
84pub struct Machine {
85    pub os: Os,
86    pub release: Option<String>,
87    pub vendor: Vendor,
88    pub card: Option<String>,
89    pub cores: usize,
90    pub ram_bytes: u64,
91    pub vram_bytes: Option<u64>,
92    pub listening: Lane,
93    pub filing: Lane,
94    pub video: Lane,
95}
96
97impl Machine {
98    pub fn accelerated(&self) -> bool {
99        self.listening.accelerated() && self.video.accelerated()
100    }
101
102    pub fn headline(&self) -> String {
103        match self.vendor {
104            Vendor::None => format!("{}, processor only", self.os.label()),
105            Vendor::Apple => format!("{} on {}", self.os.label(), Vendor::Apple.label()),
106            vendor => format!(
107                "{} with {} {}",
108                self.os.label(),
109                vendor.label(),
110                self.listening.label()
111            ),
112        }
113    }
114}
115
116fn filing_lane(os: Os, vendor: Vendor) -> Lane {
117    match (os, vendor) {
118        (_, Vendor::None) => Lane::Processor,
119        (Os::MacOs, _) => Lane::CoreMl,
120        (Os::Windows, _) => Lane::DirectMl,
121        (Os::Linux, _) => Lane::WebGpu,
122        _ => Lane::Processor,
123    }
124}
125
126pub fn probe_machine() -> Machine {
127    let os = host_os();
128    let accel = probe_accel();
129    let vendor = match accel {
130        Accel::NvidiaCuda => Vendor::Nvidia,
131        Accel::AppleVideoToolbox => Vendor::Apple,
132        Accel::AmdVaapi => Vendor::Amd,
133        Accel::IntelQsv => Vendor::Intel,
134        Accel::Cpu => vendor_of_adapters(&probe_gpu().adapters),
135    };
136
137    let listening = match (os, vendor) {
138        (_, Vendor::None) => Lane::Processor,
139        (_, Vendor::Nvidia) => Lane::Cuda,
140        (Os::MacOs, _) => Lane::Metal,
141        _ => Lane::Vulkan,
142    };
143
144    let filing = filing_lane(os, vendor);
145
146    let video = match (os, vendor) {
147        (_, Vendor::None) => Lane::Processor,
148        (_, Vendor::Nvidia) => Lane::Nvidia,
149        (Os::MacOs, _) => Lane::VideoToolbox,
150        (Os::Windows, Vendor::Amd) => Lane::Amf,
151        (Os::Linux, Vendor::Amd) => Lane::Vaapi,
152        _ => Lane::Qsv,
153    };
154
155    Machine {
156        os,
157        release: os_release(),
158        vendor,
159        card: card_name(),
160        cores: std::thread::available_parallelism()
161            .map(std::num::NonZeroUsize::get)
162            .unwrap_or(1),
163        ram_bytes: ram_bytes(),
164        vram_bytes: vram_bytes(vendor),
165        listening,
166        filing,
167        video,
168    }
169}
170
171fn readable_bytes(bytes: u64) -> String {
172    let gb = bytes as f64 / 1_073_741_824.0;
173    if gb >= 1.0 {
174        format!("{gb:.0} GB")
175    } else {
176        format!("{:.0} MB", bytes as f64 / 1_048_576.0)
177    }
178}
179
180impl Machine {
181    pub fn spec_line(&self) -> String {
182        let mut parts = Vec::new();
183        if let Some(card) = &self.card {
184            parts.push(short_card(card));
185        }
186        match (self.vram_bytes, self.vendor) {
187            (Some(vram), _) => parts.push(readable_bytes(vram)),
188            (None, Vendor::Apple) => parts.push("unified memory".to_string()),
189            (None, Vendor::Intel) => parts.push("shared graphics memory".to_string()),
190            (None, _) => {}
191        }
192        parts.push(format!("{} cores", self.cores));
193        if self.ram_bytes > 0 {
194            parts.push(readable_bytes(self.ram_bytes));
195        }
196        parts.join(" ยท ")
197    }
198
199    pub fn title(&self) -> String {
200        match &self.release {
201            Some(release) => format!("{} system information", release),
202            None => format!("{} system information", self.os.label()),
203        }
204    }
205}
206
207#[cfg(target_os = "linux")]
208fn os_release() -> Option<String> {
209    let raw = std::fs::read_to_string("/etc/os-release").ok()?;
210    raw.lines()
211        .find_map(|line| line.strip_prefix("PRETTY_NAME="))
212        .map(|value| value.trim_matches('"').to_string())
213}
214
215#[cfg(target_os = "macos")]
216fn os_release() -> Option<String> {
217    let out = std::process::Command::new("sw_vers")
218        .arg("-productVersion")
219        .output()
220        .ok()?;
221    let version = String::from_utf8_lossy(&out.stdout).trim().to_string();
222    (!version.is_empty()).then(|| format!("macOS {version}"))
223}
224
225#[cfg(not(any(target_os = "linux", target_os = "macos")))]
226fn os_release() -> Option<String> {
227    None
228}
229
230#[cfg(target_os = "linux")]
231fn ram_bytes() -> u64 {
232    std::fs::read_to_string("/proc/meminfo")
233        .ok()
234        .and_then(|raw| {
235            raw.lines()
236                .find_map(|line| line.strip_prefix("MemTotal:"))
237                .and_then(|value| value.split_whitespace().next()?.parse::<u64>().ok())
238        })
239        .map(|kb| kb * 1024)
240        .unwrap_or(0)
241}
242
243#[cfg(target_os = "macos")]
244fn ram_bytes() -> u64 {
245    std::process::Command::new("sysctl")
246        .args(["-n", "hw.memsize"])
247        .output()
248        .ok()
249        .and_then(|out| {
250            String::from_utf8_lossy(&out.stdout)
251                .trim()
252                .parse::<u64>()
253                .ok()
254        })
255        .unwrap_or(0)
256}
257
258#[cfg(not(any(target_os = "linux", target_os = "macos")))]
259fn ram_bytes() -> u64 {
260    0
261}
262
263fn vram_bytes(vendor: Vendor) -> Option<u64> {
264    match vendor {
265        Vendor::Nvidia => nvidia_vram(),
266        Vendor::Amd => amd_vram(),
267        Vendor::Intel | Vendor::Apple | Vendor::None => None,
268    }
269}
270
271fn nvidia_vram() -> Option<u64> {
272    let out = std::process::Command::new("nvidia-smi")
273        .args(["--query-gpu=memory.total", "--format=csv,noheader,nounits"])
274        .output()
275        .ok()?;
276    let first = String::from_utf8_lossy(&out.stdout)
277        .lines()
278        .next()?
279        .trim()
280        .parse::<u64>()
281        .ok()?;
282    Some(first * 1_048_576)
283}
284
285#[cfg(target_os = "linux")]
286fn amd_vram() -> Option<u64> {
287    amd_vram_from(std::path::Path::new("/sys/class/drm"))
288}
289
290#[cfg(not(target_os = "linux"))]
291fn amd_vram() -> Option<u64> {
292    None
293}
294
295#[cfg(any(target_os = "linux", test))]
296fn amd_vram_from(drm_root: &std::path::Path) -> Option<u64> {
297    let entries = std::fs::read_dir(drm_root).ok()?;
298    entries
299        .flatten()
300        .filter(|entry| {
301            let name = entry.file_name();
302            let name = name.to_string_lossy();
303            name.starts_with("card") && !name.contains('-')
304        })
305        .filter_map(|entry| {
306            std::fs::read_to_string(entry.path().join("device/mem_info_vram_total")).ok()
307        })
308        .filter_map(|raw| raw.trim().parse::<u64>().ok())
309        .max()
310}
311
312const fn host_os() -> Os {
313    if cfg!(target_os = "linux") {
314        Os::Linux
315    } else if cfg!(target_os = "windows") {
316        Os::Windows
317    } else if cfg!(target_os = "macos") {
318        Os::MacOs
319    } else {
320        Os::Other
321    }
322}
323
324fn vendor_of_name(name: &str) -> Vendor {
325    let name = name.to_lowercase();
326    if name.contains("nvidia") || name.contains("geforce") || name.contains("rtx") {
327        return Vendor::Nvidia;
328    }
329    if name.contains("amd") || name.contains("radeon") {
330        return Vendor::Amd;
331    }
332    if name.contains("intel") || name.contains("arc") {
333        return Vendor::Intel;
334    }
335    if name.contains("apple") {
336        return Vendor::Apple;
337    }
338    Vendor::None
339}
340
341fn vendor_of_adapters(adapters: &[crate::gpu::AdapterInfo]) -> Vendor {
342    let candidates = adapters
343        .iter()
344        .filter(|adapter| adapter.device_type != DeviceType::Cpu);
345    let discrete = candidates
346        .clone()
347        .filter(|adapter| adapter.device_type == DeviceType::DiscreteGpu)
348        .map(|adapter| vendor_of_name(&adapter.name))
349        .find(|vendor| *vendor != Vendor::None);
350    if let Some(vendor) = discrete {
351        return vendor;
352    }
353    candidates
354        .map(|adapter| vendor_of_name(&adapter.name))
355        .find(|vendor| *vendor != Vendor::None)
356        .unwrap_or(Vendor::None)
357}
358
359const CARD_NOISE: [&str; 12] = [
360    "NVIDIA", "GeForce", "AMD", "ATI", "Radeon", "Intel", "Intel(R)", "Apple", "Laptop", "GPU",
361    "Graphics", "(TM)",
362];
363
364pub fn short_card(name: &str) -> String {
365    let kept: Vec<&str> = name
366        .split_whitespace()
367        .filter(|w| !CARD_NOISE.iter().any(|n| n.eq_ignore_ascii_case(w)))
368        .collect();
369    if kept.is_empty() {
370        name.trim().to_string()
371    } else {
372        kept.join(" ")
373    }
374}
375
376fn card_name() -> Option<String> {
377    let gpu = probe_gpu();
378    gpu.adapters
379        .iter()
380        .find(|adapter| adapter.device_type == DeviceType::DiscreteGpu)
381        .or_else(|| {
382            gpu.adapters
383                .iter()
384                .find(|adapter| adapter.device_type != DeviceType::Cpu)
385        })
386        .map(|adapter| adapter.name.clone())
387}
388
389#[cfg(test)]
390mod tests {
391    use super::*;
392
393    #[test]
394    fn no_gpu_is_processor_regardless_of_os() {
395        for os in [Os::Linux, Os::Windows, Os::MacOs, Os::Other] {
396            assert_eq!(filing_lane(os, Vendor::None), Lane::Processor);
397        }
398    }
399
400    #[test]
401    fn macos_gets_coreml_regardless_of_vendor() {
402        for vendor in [Vendor::Apple, Vendor::Amd, Vendor::Intel, Vendor::Nvidia] {
403            assert_eq!(filing_lane(Os::MacOs, vendor), Lane::CoreMl);
404        }
405    }
406
407    #[test]
408    fn windows_gets_directml_regardless_of_vendor() {
409        for vendor in [Vendor::Amd, Vendor::Intel, Vendor::Nvidia] {
410            assert_eq!(filing_lane(Os::Windows, vendor), Lane::DirectMl);
411        }
412    }
413
414    #[test]
415    fn linux_gets_webgpu_regardless_of_vendor() {
416        for vendor in [Vendor::Amd, Vendor::Intel, Vendor::Nvidia] {
417            assert_eq!(filing_lane(Os::Linux, vendor), Lane::WebGpu);
418        }
419    }
420
421    #[test]
422    fn an_unrecognized_os_refuses_rather_than_guessing() {
423        assert_eq!(filing_lane(Os::Other, Vendor::Nvidia), Lane::Processor);
424    }
425
426    fn adapter(name: &str, device_type: crate::gpu::DeviceType) -> crate::gpu::AdapterInfo {
427        crate::gpu::AdapterInfo {
428            name: name.to_string(),
429            backend: crate::gpu::Backend::Vulkan,
430            device_type,
431        }
432    }
433
434    #[test]
435    fn the_discrete_adapter_names_the_vendor_even_when_the_igpu_enumerates_first() {
436        let adapters = [
437            adapter("Intel(R) UHD Graphics 770", DeviceType::IntegratedGpu),
438            adapter("NVIDIA GeForce RTX 4080", DeviceType::DiscreteGpu),
439        ];
440        assert_eq!(vendor_of_adapters(&adapters), Vendor::Nvidia);
441
442        let flipped = [
443            adapter("AMD Radeon RX 7900", DeviceType::DiscreteGpu),
444            adapter("Intel(R) Iris Xe", DeviceType::IntegratedGpu),
445        ];
446        assert_eq!(vendor_of_adapters(&flipped), Vendor::Amd);
447    }
448
449    #[test]
450    fn with_no_discrete_adapter_the_integrated_one_still_names_the_vendor() {
451        let adapters = [adapter("Intel(R) Iris Xe", DeviceType::IntegratedGpu)];
452        assert_eq!(vendor_of_adapters(&adapters), Vendor::Intel);
453
454        let apple = [adapter("Apple M3 Max", DeviceType::IntegratedGpu)];
455        assert_eq!(vendor_of_adapters(&apple), Vendor::Apple);
456    }
457
458    #[test]
459    fn a_cpu_adapter_alone_is_no_vendor_at_all() {
460        let adapters = [adapter("llvmpipe (LLVM)", DeviceType::Cpu)];
461        assert_eq!(vendor_of_adapters(&adapters), Vendor::None);
462        assert_eq!(vendor_of_adapters(&[]), Vendor::None);
463    }
464
465    #[test]
466    fn amd_vram_reads_the_biggest_card_from_a_sysfs_shaped_tree() {
467        let root =
468            std::path::PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../../tmp/testscratch");
469        std::fs::create_dir_all(&root).unwrap();
470        let dir = tempfile::Builder::new()
471            .prefix("probe-amd-vram-")
472            .tempdir_in(&root)
473            .unwrap();
474        let drm = dir.path();
475        std::fs::create_dir_all(drm.join("card0/device")).unwrap();
476        std::fs::create_dir_all(drm.join("card1/device")).unwrap();
477        std::fs::create_dir_all(drm.join("card1-DP-1")).unwrap();
478        std::fs::write(drm.join("card0/device/mem_info_vram_total"), "536870912\n").unwrap();
479        std::fs::write(
480            drm.join("card1/device/mem_info_vram_total"),
481            "17163091968\n",
482        )
483        .unwrap();
484
485        assert_eq!(amd_vram_from(drm), Some(17_163_091_968));
486    }
487
488    #[test]
489    fn amd_vram_on_a_tree_with_no_amdgpu_files_is_none_not_a_panic() {
490        let root =
491            std::path::PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../../tmp/testscratch");
492        std::fs::create_dir_all(&root).unwrap();
493        let dir = tempfile::Builder::new()
494            .prefix("probe-amd-none-")
495            .tempdir_in(&root)
496            .unwrap();
497        std::fs::create_dir_all(dir.path().join("card0/device")).unwrap();
498        assert_eq!(amd_vram_from(dir.path()), None);
499        assert_eq!(amd_vram_from(&dir.path().join("missing")), None);
500    }
501
502    fn machine_with(vendor: Vendor, vram: Option<u64>) -> Machine {
503        Machine {
504            os: Os::Linux,
505            release: None,
506            vendor,
507            card: Some("Test Card".to_string()),
508            cores: 8,
509            ram_bytes: 16 * 1_073_741_824,
510            vram_bytes: vram,
511            listening: Lane::Vulkan,
512            filing: Lane::WebGpu,
513            video: Lane::Vaapi,
514        }
515    }
516
517    #[test]
518    fn the_card_reads_short_without_vendor_words() {
519        assert_eq!(short_card("NVIDIA GeForce RTX 4080 Laptop GPU"), "RTX 4080");
520        assert_eq!(short_card("AMD Radeon RX 7900 XTX"), "RX 7900 XTX");
521        assert_eq!(short_card("Apple M2 Max"), "M2 Max");
522        assert_eq!(short_card("Intel(R) Iris(R) Xe Graphics"), "Iris(R) Xe");
523        assert_eq!(short_card("GPU"), "GPU");
524    }
525
526    #[test]
527    fn the_spec_line_shows_real_vram_when_a_number_exists() {
528        let line = machine_with(Vendor::Amd, Some(16 * 1_073_741_824)).spec_line();
529        assert!(line.contains("16 GB"), "{line}");
530    }
531
532    #[test]
533    fn apple_reads_unified_and_intel_reads_shared_when_no_number_exists() {
534        let apple = machine_with(Vendor::Apple, None).spec_line();
535        assert!(apple.contains("unified memory"), "{apple}");
536
537        let intel = machine_with(Vendor::Intel, None).spec_line();
538        assert!(intel.contains("shared graphics memory"), "{intel}");
539
540        let none = machine_with(Vendor::None, None).spec_line();
541        assert!(
542            !none.contains("unified") && !none.contains("shared"),
543            "{none}"
544        );
545    }
546}