DarkRoom/core/dr-inference-engine/src/hardware.rs

//! Which GPUs this device has, as far as choosing a runtime needs to know
//! (docs/dev/inference.md §3.2).
//!
//! A runtime carries one vendor's providers — Intel's build has OpenVINO,
//! the `onnxruntime-gpu` wheel CUDA and TensorRT, a ROCm build MIGraphX,
//! Microsoft's WebGPU build the generic rung — and only one runtime loads
//! per process. These checks are what lets `api` load the one that fits
//! when a device has several installed. They read files, never a driver:
//! a wrong answer costs a slower rung, which the probe still measures, and
//! a driver call at start-up could cost the launch.

/// What a runtime's providers are scored against.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct Gpus {
    pub nvidia: bool,
    /// An AMD GPU with the ROCm kernel interface, which MIGraphX needs.
    pub amd_rocm: bool,
    pub intel: bool,
    pub qualcomm: bool,
}

/// The score of a runtime whose vendor rung matches the device's GPU.
/// Nothing beats it, so the search stops there.
pub const PERFECT: u32 = 3;

impl Gpus {
    /// How well a runtime offering `providers` fits this device. The vendor
    /// rungs score above OpenVINO because a machine with an Intel iGPU and
    /// an NVIDIA or AMD card wants the card; the generic rung scores above
    /// a CPU-only build because it carries the same CPU provider and might
    /// beat it.
    pub fn score(&self, providers: &[String]) -> u32 {
        providers
            .iter()
            .map(|p| match p.as_str() {
                "TensorrtExecutionProvider" | "CUDAExecutionProvider" if self.nvidia => PERFECT,
                "MIGraphXExecutionProvider" if self.amd_rocm => PERFECT,
                "QNNExecutionProvider" if self.qualcomm => PERFECT,
                "CoreMLExecutionProvider" => PERFECT,
                "OpenVINOExecutionProvider" if self.intel => 2,
                "WebGpuExecutionProvider" => 1,
                _ => 0,
            })
            .max()
            .unwrap_or(0)
    }
}

#[cfg(target_os = "linux")]
pub fn detect() -> Gpus {
    use std::path::Path;
    // Every DRM card's PCI vendor: an Intel iGPU is `0x8086` whether or
    // not its compute driver is installed, which the probe finds out.
    let vendors: Vec<String> = std::fs::read_dir("/sys/class/drm")
        .into_iter()
        .flatten()
        .filter_map(|e| e.ok())
        .filter(|e| {
            let name = e.file_name();
            let name = name.to_string_lossy();
            name.starts_with("card") && !name.contains('-')
        })
        .filter_map(|e| std::fs::read_to_string(e.path().join("device/vendor")).ok())
        .map(|v| v.trim().to_string())
        .collect();
    Gpus {
        nvidia: Path::new("/proc/driver/nvidia/version").exists(),
        amd_rocm: Path::new("/dev/kfd").exists(),
        intel: vendors.iter().any(|v| v == "0x8086"),
        qualcomm: false,
    }
}

#[cfg(target_os = "windows")]
pub fn detect() -> Gpus {
    use std::path::PathBuf;
    let root = std::env::var_os("SystemRoot")
        .map(PathBuf::from)
        .unwrap_or_else(|| PathBuf::from(r"C:\Windows"));
    let system32 = root.join("System32");
    // Intel's DCH graphics driver, integrated and Arc alike, installs
    // from `iigd_dch.inf`; its package directory is the evidence.
    let intel = std::fs::read_dir(system32.join(r"DriverStore\FileRepository"))
        .into_iter()
        .flatten()
        .filter_map(|e| e.ok())
        .any(|e| e.file_name().to_string_lossy().starts_with("iigd_dch"));
    Gpus {
        nvidia: system32.join("nvcuda.dll").exists(),
        amd_rocm: false,
        intel,
        qualcomm: false,
    }
}

#[cfg(target_os = "android")]
pub fn detect() -> Gpus {
    // Fail-safe: only a device that names another vendor is not Qualcomm.
    // `ro.soc.manufacturer` exists from Android 12, and a property or file
    // the app cannot read reads as nothing; nothing keeps the QNN build
    // first, as 0.22 had it, where a Qualcomm device mistaken for another
    // would trade its NPU for the generic rung. Qualcomm's FastRPC library,
    // which the Hexagon path loads anyway, overrules a name.
    let soc = crate::probe::system_property("ro.soc.manufacturer");
    let fastrpc = [
        "/vendor/lib64/libcdsprpc.so",
        "/system/vendor/lib64/libcdsprpc.so",
    ]
    .iter()
    .any(|p| std::path::Path::new(p).exists());
    Gpus {
        qualcomm: qualcomm_soc(&soc) || fastrpc,
        ..Gpus::default()
    }
}

/// Whether `ro.soc.manufacturer` leaves the device Qualcomm's: it says so,
/// or it says nothing.
#[cfg(any(target_os = "android", test))]
fn qualcomm_soc(manufacturer: &str) -> bool {
    let m = manufacturer.trim();
    m.is_empty() || m.eq_ignore_ascii_case("QTI") || m.eq_ignore_ascii_case("Qualcomm")
}

#[cfg(not(any(target_os = "linux", target_os = "windows", target_os = "android")))]
pub fn detect() -> Gpus {
    Gpus::default()
}

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

    fn offers(p: &[&str]) -> Vec<String> {
        p.iter().map(|s| s.to_string()).collect()
    }

    #[test]
    fn only_a_named_other_vendor_is_not_qualcomm() {
        assert!(qualcomm_soc("QTI"));
        assert!(qualcomm_soc("Qualcomm"));
        // Unreadable, or older than Android 12: the QNN build stays first.
        assert!(qualcomm_soc(""));
        assert!(!qualcomm_soc("Mediatek"));
        assert!(!qualcomm_soc("Google"));
        assert!(!qualcomm_soc("Samsung"));
    }

    #[test]
    fn the_card_beats_the_integrated_gpu_and_both_beat_the_generic_rung() {
        let cpu = offers(&["CPUExecutionProvider"]);
        let nvidia = offers(&[
            "TensorrtExecutionProvider",
            "CUDAExecutionProvider",
            "CPUExecutionProvider",
        ]);
        let intel = offers(&["OpenVINOExecutionProvider", "CPUExecutionProvider"]);
        let webgpu = offers(&["WebGpuExecutionProvider", "CPUExecutionProvider"]);
        let laptop = Gpus {
            nvidia: true,
            intel: true,
            ..Gpus::default()
        };
        assert!(laptop.score(&nvidia) > laptop.score(&intel));
        assert!(laptop.score(&intel) > laptop.score(&webgpu));
        assert!(laptop.score(&webgpu) > laptop.score(&cpu));
        // No Intel GPU: Intel's build is worth no more than a CPU build to
        // this device, and the generic rung is worth more.
        let amd_on_windows = Gpus::default();
        assert_eq!(amd_on_windows.score(&intel), amd_on_windows.score(&cpu));
        assert!(amd_on_windows.score(&webgpu) > amd_on_windows.score(&intel));
        // A ROCm build on a machine without ROCm is a CPU build.
        let rocm = offers(&["MIGraphXExecutionProvider", "CPUExecutionProvider"]);
        assert_eq!(amd_on_windows.score(&rocm), 0);
    }
}