dtourolle/DarkRoom
1
0
Fork 0
Code Issues 39 Pull Requests Actions Packages Projects Releases 20 Wiki Activity

Load the runtime that fits the GPU, not the first one found

Browse Source 
A runtime carries one vendor's providers, only one loads per process,
and a device can now hold several: the package's OpenVINO or WebGPU
build, a CUDA build the user fetched, the distribution's ROCm build.
`api::install` opens each it finds, lists its providers with
GetAvailableProviders, and installs the one scoring highest against the
GPUs `hardware::detect` reads from files — a vendor rung on its own
vendor's GPU above OpenVINO on an Intel one above the generic WebGPU
rung above a CPU-only build. Equal scores keep the old first-found
order, and DARKROOM_ORT_DIR still wins outright. The losers stay mapped
rather than unloaded.

The Linux fingerprint now names the OpenCL drivers too, so installing
Intel's re-probes. `ladder` takes DARKROOM_ORT_DIRS to show the choice.
This commit is contained in:
dtourolle
2026-10-04 21:00:15 -04:00
parent 87c405eb46
commit 85dee4375b
6 changed files with 317 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files
core/dr-inference-engine/examples/ladder.rs
+10
View File
@@ -6,6 +6,9 @@
//! cargo run --release -p dr-inference-engine --features native,tract \
//! --example ladder -- CACHE_DIR models/face/scrfd_500m_640.onnx [ROLE=MODEL.onnx ...]
//!
//! `DARKROOM_ORT_DIRS=a:b:c` offers several runtimes, as the app's search
//! list does, and shows which the engine chose for this device's GPU.
//!
//! A bare path is a `Detector`; `denoiser=…`, `scene=…`, `inpainter=…`,
//! `landmarks=…` (any `Role`, lower case) says otherwise, so a device can
//! show each role taking its own form (inference.md §1.5). Each is opened
@@ -31,9 +34,16 @@ fn main() {
std::process::exit(2);
}
// DARKROOM_ORT_DIRS lists several, colon-separated, as the app's search
// does: the engine loads the one that fits the GPU (§3.2).
let runtime_dirs: Vec<PathBuf> = std::env::var_os("DARKROOM_ORT_DIR")
.map(PathBuf::from)
.into_iter()
.chain(
std::env::var_os("DARKROOM_ORT_DIRS")
.map(|v| std::env::split_paths(&v).collect::<Vec<_>>())
.unwrap_or_default(),
)
.collect();
let started = Instant::now();
dr_inference_engine::init(dr_inference_engine::Config {
core/dr-inference-engine/src/api.rs
+141 -28
View File
@@ -51,17 +51,22 @@ pub fn ensure_installed() {
}
}
/// Look for `libonnxruntime` in `dirs`, in order, and hand `ort` the first
/// table that loads; otherwise tract. Once per process.
/// Find every `libonnxruntime` in `dirs`, hand `ort` the table of the one
/// that best fits this device's GPUs, and fall to tract if none loads.
/// Once per process.
///
/// Best fit, not first found (§3.2): a device can hold several runtimes —
/// the package's OpenVINO build, a CUDA build the user fetched, the
/// distribution's ROCm build — and each carries one vendor's providers.
/// Between equals, the earlier directory wins, as it always has.
/// `DARKROOM_ORT_DIR`, when it loads, wins outright: it is how a person
/// says which runtime they mean.
pub fn install(dirs: &[PathBuf]) -> Runtime {
RUNTIME
.get_or_init(|| {
#[cfg(feature = "native")]
for dir in dirs {
match load_native(dir) {
Ok(rt) => return rt,
Err(e) => log::info!("inference: no runtime in {}: {e}", dir.display()),
}
if let Some(rt) = install_best(dirs) {
return rt;
}
#[cfg(not(feature = "native"))]
let _ = dirs;
@@ -70,6 +75,86 @@ pub fn install(dirs: &[PathBuf]) -> Runtime {
.clone()
}
/// A runtime opened to read its providers, not yet handed to `ort`.
#[cfg(feature = "native")]
struct Found {
lib: libloading::Library,
api: *const ort_sys::OrtApi,
path: PathBuf,
version: String,
providers: Vec<String>,
}
#[cfg(feature = "native")]
fn install_best(dirs: &[PathBuf]) -> Option<Runtime> {
let named = std::env::var_os("DARKROOM_ORT_DIR").map(PathBuf::from);
let gpus = crate::hardware::detect();
let mut found: Vec<Found> = Vec::new();
let mut seen = std::collections::HashSet::new();
for dir in dirs {
match open_native(dir) {
Ok(f) => {
// `bin/../lib/darkroom` and `/usr/lib/darkroom` are one file.
if !seen.insert(std::fs::canonicalize(&f.path).unwrap_or(f.path.clone())) {
std::mem::forget(f.lib);
continue;
}
log::info!(
"inference: ONNX Runtime {} at {} offers {}",
f.version,
f.path.display(),
f.providers.join(", ")
);
if named.as_deref() == Some(dir.as_path()) {
found.clear();
found.push(f);
break;
}
found.push(f);
}
Err(e) => log::info!("inference: no runtime in {}: {e}", dir.display()),
}
}
let best = (0..found.len())
.max_by_key(|&i| (gpus.score(&found[i].providers), std::cmp::Reverse(i)))?;
let chosen = found.swap_remove(best);
// The others stay mapped. Unloading a C++ runtime after its static
// constructors ran is a crash at exit waiting to happen, and an
// unused mapping costs address space, not memory.
for other in found {
std::mem::forget(other.lib);
}
log::info!("inference: chose {} for {gpus:?}", chosen.path.display());
// SAFETY: the table came from this library's `OrtGetApiBase`, and the
// library is leaked below, so every pointer in the copy stays valid for
// the life of the process.
if !ort::set_api(unsafe { (*chosen.api).clone() }) {
log::warn!("inference: an API table was already installed");
std::mem::forget(chosen.lib);
return None;
}
std::mem::forget(chosen.lib);
// Qualcomm's DSP loader finds the Hexagon skel through this variable,
// and only through it; the runtime's own directory is where the APK
// put it. Harmless anywhere else.
#[cfg(target_os = "android")]
if let Some(dir) = chosen.path.parent().filter(|d| !d.as_os_str().is_empty()) {
std::env::set_var("ADSP_LIBRARY_PATH", dir);
}
log::info!(
"inference: ONNX Runtime {} from {}",
chosen.version,
chosen.path.display()
);
Some(Runtime::OnnxRuntime {
path: chosen.path,
version: chosen.version,
})
}
#[cfg(feature = "tract")]
fn install_tract() -> Runtime {
let _ = ort::set_api(ort_tract::api());
@@ -85,8 +170,12 @@ fn install_tract() -> Runtime {
Runtime::Tract
}
/// Open the runtime in `dir` and read what it offers. `dir` may also name
/// the library itself — Android has two runtimes and one directory, so the
/// second goes by its file name — and an empty path is the bare name
/// through the system loader, which on Android is the APK's own copy.
#[cfg(feature = "native")]
fn load_native(dir: &std::path::Path) -> Result<Runtime, String> {
fn open_native(dir: &std::path::Path) -> Result<Found, String> {
let name = if cfg!(target_os = "windows") {
"onnxruntime.dll"
} else if cfg!(any(target_os = "macos", target_os = "ios")) {
@@ -94,18 +183,21 @@ fn load_native(dir: &std::path::Path) -> Result<Runtime, String> {
} else {
"libonnxruntime.so"
};
// An empty dir means the bare name: the system loader's search, which on
// Android includes the APK's own native libraries.
let is_library = dir.file_name().and_then(|n| n.to_str()).is_some_and(|n| {
n.contains("onnxruntime")
&& (n.ends_with(".so") || n.ends_with(".dll") || n.ends_with(".dylib"))
});
let path = if dir.as_os_str().is_empty() {
PathBuf::from(name)
} else if is_library {
dir.to_path_buf()
} else {
find_library(dir, name).ok_or("not present")?
};
// SAFETY: the library's initialisers are ONNX Runtime's own; the symbol
// is the documented entry point with the documented signature; the table
// is copied out and the library handle is leaked, so every pointer in
// the copy stays valid for the life of the process.
// is the documented entry point with the documented signature. The
// table pointer is valid while `lib` is, which the caller keeps.
unsafe {
let lib = libloading::Library::new(&path).map_err(|e| e.to_string())?;
let get_base: libloading::Symbol<
@@ -125,24 +217,45 @@ fn load_native(dir: &std::path::Path) -> Result<Runtime, String> {
ort_sys::ORT_API_VERSION
));
}
if !ort::set_api((*api).clone()) {
return Err("an API table was already installed".into());
}
std::mem::forget(lib);
// Qualcomm's DSP loader finds the Hexagon skel through this variable,
// and only through it; the runtime's own directory is where the APK
// put it. Harmless anywhere else.
#[cfg(target_os = "android")]
if !dir.as_os_str().is_empty() {
std::env::set_var("ADSP_LIBRARY_PATH", dir);
}
log::info!("inference: ONNX Runtime {version} from {}", path.display());
Ok(Runtime::OnnxRuntime { path, version })
let providers = available_providers(api);
Ok(Found {
lib,
api,
path,
version,
providers,
})
}
}
/// The providers compiled into the runtime behind `api` — not the ones this
/// device can run, which is the probe's question.
///
/// # Safety
/// `api` must be a live table from `GetApi`.
#[cfg(feature = "native")]
unsafe fn available_providers(api: *const ort_sys::OrtApi) -> Vec<String> {
let mut list: *mut *mut std::ffi::c_char = std::ptr::null_mut();
let mut n: std::ffi::c_int = 0;
let status = ((*api).GetAvailableProviders)(&mut list, &mut n);
if !status.0.is_null() {
((*api).ReleaseStatus)(status.0);
return Vec::new();
}
let names = (0..n.max(0) as usize)
.map(|i| {
std::ffi::CStr::from_ptr(*list.add(i))
.to_string_lossy()
.into_owned()
})
.collect();
let status = ((*api).ReleaseAvailableProviders)(list, n);
if !status.0.is_null() {
((*api).ReleaseStatus)(status.0);
}
names
}
/// `libonnxruntime.so` in `dir`, or a versioned spelling of it —
/// `libonnxruntime.so.1.30.0` is what the Python wheel ships, and a package
/// that installs only the versioned file is not wrong.
core/dr-inference-engine/src/hardware.rs
+141
View File
@@ -0,0 +1,141 @@
//! 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,
}
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 => 3,
"MIGraphXExecutionProvider" if self.amd_rocm => 3,
"QNNExecutionProvider" if self.qualcomm => 3,
"CoreMLExecutionProvider" => 3,
"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 {
let soc = crate::probe::system_property("ro.soc.manufacturer");
Gpus {
qualcomm: soc.eq_ignore_ascii_case("QTI") || soc.eq_ignore_ascii_case("Qualcomm"),
..Gpus::default()
}
}
#[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 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);
}
}
core/dr-inference-engine/src/lib.rs
+1
View File
@@ -21,6 +21,7 @@ use serde::{Deserialize, Serialize};
mod api;
mod engines;
mod hardware;
mod probe;
mod session;
core/dr-inference-engine/src/probe.rs
+23 -7
View File
@@ -391,19 +391,35 @@ fn providers_beside(runtime: &Path) -> String {
#[cfg(target_os = "linux")]
fn device_identity() -> String {
// The NVIDIA driver's version line, or the ROCm release the AMD stack
// came from (`rocm-core` writes it; the kernel driver has no version
// of its own). Absent means neither.
// The NVIDIA driver's version line, the ROCm release the AMD stack came
// from (`rocm-core` writes it; the kernel driver has no version of its
// own), and the OpenCL drivers registered — OpenVINO reaches the GPU
// through one, and installing Intel's is what makes the Iris Xe a rung.
let mut parts = Vec::new();
if let Some(line) = std::fs::read_to_string("/proc/driver/nvidia/version")
.ok()
.and_then(|s| s.lines().next().map(str::to_string))
{
return line;
parts.push(line);
}
if let Ok(rocm) = std::fs::read_to_string("/opt/rocm/.info/version") {
return format!("rocm {}", rocm.trim());
parts.push(format!("rocm {}", rocm.trim()));
}
let mut icds: Vec<String> = std::fs::read_dir("/etc/OpenCL/vendors")
.into_iter()
.flatten()
.filter_map(|e| e.ok())
.map(|e| e.file_name().to_string_lossy().into_owned())
.collect();
icds.sort();
if !icds.is_empty() {
parts.push(format!("opencl {}", icds.join(" ")));
}
if parts.is_empty() {
"no nvidia driver, no rocm, no opencl".into()
} else {
parts.join("; ")
}
"no nvidia driver, no rocm".into()
}
#[cfg(target_os = "android")]
@@ -418,7 +434,7 @@ fn device_identity() -> String {
}
#[cfg(target_os = "android")]
fn system_property(name: &str) -> String {
pub(crate) fn system_property(name: &str) -> String {
extern "C" {
fn __system_property_get(
name: *const std::ffi::c_char,