Skip to main content

probe/
accel.rs

1#[derive(Debug, Clone, Copy, PartialEq, Eq)]
2pub enum Accel {
3    NvidiaCuda,
4    AppleVideoToolbox,
5    AmdVaapi,
6    IntelQsv,
7    Cpu,
8}
9
10#[cfg(target_os = "linux")]
11pub fn probe_accel() -> Accel {
12    linux::probe()
13}
14
15#[cfg(target_os = "macos")]
16pub fn probe_accel() -> Accel {
17    Accel::AppleVideoToolbox
18}
19
20#[cfg(target_os = "windows")]
21pub fn probe_accel() -> Accel {
22    windows::probe()
23}
24
25#[cfg(not(any(target_os = "linux", target_os = "macos", target_os = "windows")))]
26pub fn probe_accel() -> Accel {
27    Accel::Cpu
28}
29
30#[cfg(target_os = "linux")]
31mod linux {
32    use super::Accel;
33    use std::path::Path;
34
35    pub fn probe() -> Accel {
36        if nvidia_present() {
37            return Accel::NvidiaCuda;
38        }
39        pick(&render_node_vendors())
40    }
41
42    pub(super) fn pick(vendors: &[u32]) -> Accel {
43        if vendors.contains(&0x1002) {
44            return Accel::AmdVaapi;
45        }
46        if vendors.contains(&0x8086) {
47            return Accel::IntelQsv;
48        }
49        Accel::Cpu
50    }
51
52    fn nvidia_present() -> bool {
53        Path::new("/dev/nvidia0").exists() || nvidia_smi_on_path()
54    }
55
56    fn nvidia_smi_on_path() -> bool {
57        let Some(path_var) = std::env::var_os("PATH") else {
58            return false;
59        };
60        std::env::split_paths(&path_var).any(|dir| dir.join("nvidia-smi").is_file())
61    }
62
63    fn render_node_vendors() -> Vec<u32> {
64        let Ok(entries) = std::fs::read_dir("/dev/dri") else {
65            return Vec::new();
66        };
67        let mut names: Vec<String> = entries
68            .flatten()
69            .filter_map(|entry| entry.file_name().into_string().ok())
70            .filter(|name| name.starts_with("renderD"))
71            .collect();
72        names.sort();
73        names
74            .into_iter()
75            .filter_map(|name| {
76                let contents =
77                    std::fs::read_to_string(format!("/sys/class/drm/{name}/device/vendor")).ok()?;
78                u32::from_str_radix(contents.trim().trim_start_matches("0x"), 16).ok()
79            })
80            .collect()
81    }
82}
83
84#[cfg(target_os = "windows")]
85mod windows {
86    use super::Accel;
87
88    pub fn probe() -> Accel {
89        if nvidia_smi_on_path() {
90            Accel::NvidiaCuda
91        } else {
92            Accel::Cpu
93        }
94    }
95
96    fn nvidia_smi_on_path() -> bool {
97        let Some(path_var) = std::env::var_os("PATH") else {
98            return false;
99        };
100        std::env::split_paths(&path_var).any(|dir| dir.join("nvidia-smi.exe").is_file())
101    }
102}
103
104#[cfg(test)]
105mod tests {
106    use super::*;
107
108    #[test]
109    fn probe_accel_never_panics() {
110        std::hint::black_box(probe_accel());
111    }
112
113    #[cfg(target_os = "linux")]
114    #[test]
115    fn the_discrete_amd_wins_over_an_intel_igpu_whichever_node_comes_first() {
116        assert_eq!(linux::pick(&[0x8086, 0x1002]), Accel::AmdVaapi);
117        assert_eq!(linux::pick(&[0x1002, 0x8086]), Accel::AmdVaapi);
118    }
119
120    #[cfg(target_os = "linux")]
121    #[test]
122    fn a_lone_intel_node_is_quicksync_and_nothing_is_cpu() {
123        assert_eq!(linux::pick(&[0x8086]), Accel::IntelQsv);
124        assert_eq!(linux::pick(&[]), Accel::Cpu);
125        assert_eq!(linux::pick(&[0xffff]), Accel::Cpu);
126    }
127}