Add dr-inference-engine and route every model session through it

One crate names the runtime, the providers and the devices; dr-face and
dr-segment ask it for a session by role. It hands ort an API table once
per process — from a libonnxruntime it dlopens when the app names a
directory holding one, otherwise from tract — so the Rust build stays
free of C on every target and a package can install the runtime as a
file (docs/inference.md §3).

Sessions live in a registry behind a Model handle that holds the bytes,
not the session: every use refreshes a timestamp and a reaper unloads
whatever sat idle past the decay. A scan that runs the detector on each
image never lets it go idle; a click in the develop view lets the
segmenter go after thirty seconds; a handle used after that reloads,
and reloads on a higher rung if a compiled engine has landed meanwhile.

The probe walks the platform's ladder by building strict sessions and
timing them against the CPU provider, caches the choice against a
fingerprint of the runtime, driver, hardware and models, and compiles
engines for the selected rung in the background, smallest model first.
Nothing in this commit turns the native path on: the apps still run on
tract until they call init with a runtime directory.
This commit is contained in:
2026-09-19 16:02:37 +02:00
parent caf21bea64
commit d15c41e699
16 changed files with 1368 additions and 90 deletions
+129
View File
@@ -0,0 +1,129 @@
//! One session builder per rung (docs/inference.md §2, §7, §9).
use ort::session::builder::GraphOptimizationLevel;
use ort::session::Session;
use crate::{Config, Role, Rung};
/// Build a session for `bytes` on `rung`.
///
/// `strict` is the probe's flag: with it, a provider that would hand any
/// node to the CPU fails the build instead, so "the session built" means
/// "the provider took the graph" and not "the provider registered" (§4).
pub fn build(
rung: Rung,
role: Role,
bytes: &[u8],
cfg: &Config,
strict: bool,
) -> ort::Result<Session> {
let mut b = Session::builder()?
.with_optimization_level(GraphOptimizationLevel::Level3)?
.with_intra_threads(threads(cfg))?;
if strict {
b = b.with_config_entry("session.disable_cpu_ep_fallback", "1")?;
}
// A Hexagon session loads the compiled context when there is one and
// compiles it from the model when there is not; the engine thread is
// what makes the second case rare (§6).
let context = (rung == Rung::Hexagon).then(|| crate::engines::context_path(cfg, bytes));
let ready = context.as_ref().is_some_and(|p| p.is_file());
b = providers(
b,
rung,
role,
cfg,
if ready { None } else { context.as_deref() },
)?;
match (ready, context) {
(true, Some(path)) => b.commit_from_file(path),
_ => b.commit_from_memory(bytes),
}
}
/// The intra-op pool: what the config says, else the cores less two for
/// the compositor and the decoder (§9). tract ignores it.
fn threads(cfg: &Config) -> usize {
if cfg.threads > 0 {
return cfg.threads;
}
std::thread::available_parallelism()
.map(|n| n.get().saturating_sub(2).max(1))
.unwrap_or(1)
}
#[cfg(not(target_os = "android"))]
fn providers(
b: ort::session::builder::SessionBuilder,
rung: Rung,
role: Role,
cfg: &Config,
_generate_context: Option<&std::path::Path>,
) -> ort::Result<ort::session::builder::SessionBuilder> {
use ort::ep;
match rung {
Rung::Cpu => Ok(b),
Rung::Cuda => {
Ok(b.with_execution_providers([ep::CUDA::default().build().error_on_failure()])?)
}
Rung::TensorRt => {
let cache = cfg.cache_dir.join("tensorrt");
let _ = std::fs::create_dir_all(&cache);
let cache = cache.to_string_lossy().into_owned();
// fp16 for everything but the embedder, whose comparability
// across devices is worth more than its 0.2 ms (§7). The
// workspace cap keeps the develop view's tiles on the card
// (NFR-RES-2). CUDA behind it takes any node TensorRT declines.
Ok(b.with_execution_providers([
ep::TensorRT::default()
.with_fp16(role != Role::Embedder)
.with_engine_cache(true)
.with_engine_cache_path(&cache)
.with_timing_cache(true)
.with_timing_cache_path(&cache)
.with_max_workspace_size(512 << 20)
.build()
.error_on_failure(),
ep::CUDA::default().build(),
])?)
}
Rung::Hexagon => unreachable!("the Hexagon rung is not on a desktop ladder"),
}
}
#[cfg(target_os = "android")]
fn providers(
b: ort::session::builder::SessionBuilder,
rung: Rung,
_role: Role,
_cfg: &Config,
generate_context: Option<&std::path::Path>,
) -> ort::Result<ort::session::builder::SessionBuilder> {
use ort::ep;
match rung {
Rung::Cpu => Ok(b),
Rung::Hexagon => {
// The HTP compiles the graph once per device (0.8–1.7 s here).
// With `ep.context_enable` ONNX Runtime writes the compiled
// context beside the probe cache; the next session loads that
// file as its model and skips the compile (§5).
let mut b = b;
if let Some(ctx) = generate_context {
let _ = std::fs::create_dir_all(ctx.parent().unwrap());
b = b
.with_config_entry("ep.context_enable", "1")?
.with_config_entry("ep.context_file_path", ctx.to_string_lossy())?
.with_config_entry("ep.context_embed_mode", "0")?;
}
// Quantise/dequantise at the graph's edges stay on the NPU too,
// so a strict build is a whole-graph build.
Ok(b.with_execution_providers([ep::QNN::default()
.with_backend_path("libQnnHtp.so")
.with_performance_mode(ep::qnn::PerformanceMode::Burst)
.with_offload_graph_io_quantization(false)
.build()
.error_on_failure()])?)
}
Rung::Cuda | Rung::TensorRt => unreachable!("no NVIDIA rung on Android"),
}
}