dtourolle
c73743394f
The macOS ladder was the CPU provider alone, with CoreML listed as a gap. It is now CoreML, then the CPU, then tract — unmeasured, since nobody here has a Mac, and safe to ship unmeasured because the probe's clock rejects a CoreML slower than the CPU and `attempt` refuses one that crashes. - `Rung::CoreMl`, a compiling rung like TensorRT: an ML Program with every compute unit allowed, falling back to the CPU until each model's program is built. The embedder stays on the CPU, as on the Hexagon (§7). - The cache is one directory per model and runtime version. CoreML keys a model committed from memory on its input and node names, not its weights (ONNX Runtime 1.29, coreml_execution_provider.cc), so two exports of one architecture would otherwise share a program. - The fingerprint on macOS is the chip and the OS release, which ships CoreML. - The desktop looks for the runtime in the bundle's Contents/Frameworks and Homebrew's prefixes; fetch-desktop-runtime.sh on a Mac downloads ONNX Runtime 1.29.0 for Apple silicon, which carries CoreML. docs/dev/macos.md says what exists, how to build it, and which log lines to ask a Mac user for.
258 lines
11 KiB
Rust
258 lines
11 KiB
Rust
//! One session builder per rung (docs/dev/inference.md §2, §7, §9).
|
||
|
||
use ort::session::Session;
|
||
|
||
use crate::{Config, Role, Rung};
|
||
|
||
/// Build a session for `bytes` on `rung`.
|
||
///
|
||
/// Not strict about the CPU: `session.disable_cpu_ep_fallback` was tried as
|
||
/// the probe's proof that a provider took the graph, and it refuses the
|
||
/// Hexagon over the ten quantise/dequantise nodes at the graph's edges that
|
||
/// QNN declines by policy and that cost microseconds. The probe's proof is
|
||
/// its clock instead (§4): a provider that hands real work to the CPU is
|
||
/// slower than the CPU floor and rejected by the same measurement.
|
||
pub fn build(rung: Rung, role: Role, bytes: &[u8], cfg: &Config) -> ort::Result<Session> {
|
||
// No optimisation level named. ONNX Runtime's default is already its
|
||
// fullest, and on tract any level but "disabled" means `into_optimized`,
|
||
// whose optimiser divides by zero inside yolo26n-seg (tract-data
|
||
// `stack_tensors`) — a panic across the C API, which is an abort. The
|
||
// app never asked tract for that and does not start now.
|
||
let mut b = Session::builder()?.with_intra_threads(threads(cfg))?;
|
||
if crate::api::runtime().is_native() {
|
||
b = with_runtime_log(b)?;
|
||
}
|
||
// 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());
|
||
// What the rung keeps for this model: the context the Hexagon is to
|
||
// write, or the directory CoreML compiles into.
|
||
let per_model = match rung {
|
||
Rung::CoreMl => Some(crate::engines::coreml_dir(cfg, bytes)),
|
||
_ if ready => None,
|
||
_ => context.clone(),
|
||
};
|
||
b = providers(b, rung, role, cfg, per_model.as_deref())?;
|
||
match (ready, context) {
|
||
(true, Some(path)) => b.commit_from_file(path),
|
||
_ => b.commit_from_memory(bytes),
|
||
}
|
||
}
|
||
|
||
/// Send the runtime's own messages for this session to `log`, under the
|
||
/// target `onnxruntime`, instead of to ONNX Runtime's stdio logger.
|
||
///
|
||
/// Its stderr is nowhere once the app is launched from a menu, and what a
|
||
/// provider says while it partitions a graph — how many nodes it took, which
|
||
/// operator it declined, the library it failed to load — is most of what a
|
||
/// failed rung tells you (docs/dev/inference.md §4). The level follows the
|
||
/// filter: warnings always, `debug` adds the runtime's info lines (the
|
||
/// partition counts), `trace` its verbose ones (every node placement).
|
||
fn with_runtime_log(
|
||
b: ort::session::builder::SessionBuilder,
|
||
) -> ort::Result<ort::session::builder::SessionBuilder> {
|
||
use ort::logging::LogLevel;
|
||
let level = if log::log_enabled!(target: "onnxruntime", log::Level::Trace) {
|
||
LogLevel::Verbose
|
||
} else if log::log_enabled!(target: "onnxruntime", log::Level::Debug) {
|
||
LogLevel::Info
|
||
} else {
|
||
LogLevel::Warning
|
||
};
|
||
let forward = |level: LogLevel, _category: &str, _id: &str, location: &str, message: &str| {
|
||
let level = match level {
|
||
LogLevel::Verbose => log::Level::Trace,
|
||
LogLevel::Info => log::Level::Debug,
|
||
LogLevel::Warning => log::Level::Warn,
|
||
LogLevel::Error | LogLevel::Fatal => log::Level::Error,
|
||
};
|
||
log::log!(target: "onnxruntime", level, "{message} ({location})");
|
||
};
|
||
Ok(b.with_logger(std::sync::Arc::new(forward))?
|
||
.with_log_level(level)?)
|
||
}
|
||
|
||
/// 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,
|
||
per_model: Option<&std::path::Path>,
|
||
) -> ort::Result<ort::session::builder::SessionBuilder> {
|
||
use ort::ep;
|
||
match rung {
|
||
Rung::Cpu => Ok(b),
|
||
Rung::CoreMl => coreml(b, per_model),
|
||
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::MiGraphX => {
|
||
// fp16 on the same terms as TensorRT (§7). MIGraphX compiles a
|
||
// program per graph — 20–60 s here — and keeps it in the cache
|
||
// directory, keyed on the graph, the GPU and its own version
|
||
// but not the precision: hence one directory per precision.
|
||
// The CPU takes any node it declines.
|
||
let fp16 = role != Role::Embedder;
|
||
let cache = cfg
|
||
.cache_dir
|
||
.join("migraphx")
|
||
.join(if fp16 { "fp16" } else { "f32" });
|
||
let _ = std::fs::create_dir_all(&cache);
|
||
let mut b = b;
|
||
migraphx(&mut b, fp16, &cache)?;
|
||
Ok(b)
|
||
}
|
||
Rung::Hexagon => unreachable!("the Hexagon rung is not on a desktop ladder"),
|
||
}
|
||
}
|
||
|
||
/// CoreML, compiling an ML Program — the format with the operators these
|
||
/// graphs use and the one that reaches the Neural Engine — into `cache`.
|
||
///
|
||
/// The option names are those ONNX Runtime 1.29 reads from the generic
|
||
/// key/value map (`coreml_options.cc`), which is what `ort`'s builder
|
||
/// fills. The cache is per model because of how CoreML keys it: a model
|
||
/// committed from memory, as every session here is, has no path, and the
|
||
/// key falls back to a hash of the graph's input and node names — not its
|
||
/// weights. Two exports of one architecture would share a program. The
|
||
/// directory `engines::coreml_dir` names is the hash of the bytes.
|
||
///
|
||
/// Every compute unit is allowed, so CoreML may place a graph on the
|
||
/// Neural Engine, the GPU or the CPU; the probe's clock judges the result.
|
||
#[cfg(target_os = "macos")]
|
||
fn coreml(
|
||
b: ort::session::builder::SessionBuilder,
|
||
cache: Option<&std::path::Path>,
|
||
) -> ort::Result<ort::session::builder::SessionBuilder> {
|
||
use ort::ep::{self, coreml};
|
||
let mut ep = ep::CoreML::default()
|
||
.with_model_format(coreml::ModelFormat::MLProgram)
|
||
.with_compute_units(coreml::ComputeUnits::All);
|
||
if let Some(dir) = cache {
|
||
let _ = std::fs::create_dir_all(dir);
|
||
ep = ep.with_model_cache_dir(dir.to_string_lossy());
|
||
}
|
||
Ok(b.with_execution_providers([ep.build().error_on_failure()])?)
|
||
}
|
||
|
||
#[cfg(not(any(target_os = "android", target_os = "macos")))]
|
||
fn coreml(
|
||
_b: ort::session::builder::SessionBuilder,
|
||
_cache: Option<&std::path::Path>,
|
||
) -> ort::Result<ort::session::builder::SessionBuilder> {
|
||
unreachable!("the CoreML rung is on the macOS ladder only")
|
||
}
|
||
|
||
/// Register MIGraphX through ONNX Runtime's generic key/value entry point.
|
||
///
|
||
/// `ort`'s own builder (`ep::MIGraphX`) fills the legacy
|
||
/// `OrtMIGraphXProviderOptions`, and 1.29 reads that struct for its
|
||
/// precision flags and nothing else — the compiled-program cache directory
|
||
/// is only a key in the generic map (`migraphx_model_cache_dir`), and
|
||
/// without it every session is a full compile. Registration through the
|
||
/// generic entry point needs no `ort` feature: it is one call on the API
|
||
/// table, which is why the crate's `ort` dependency names no AMD feature.
|
||
#[cfg(not(target_os = "android"))]
|
||
fn migraphx(
|
||
b: &mut ort::session::builder::SessionBuilder,
|
||
fp16: bool,
|
||
cache: &std::path::Path,
|
||
) -> ort::Result<()> {
|
||
use ort::AsPointer;
|
||
use std::ffi::CString;
|
||
let keys = [c"migraphx_fp16_enable", c"migraphx_model_cache_dir"];
|
||
let values = [
|
||
CString::new(if fp16 { "1" } else { "0" }).unwrap(),
|
||
CString::new(cache.to_string_lossy().as_bytes())
|
||
.map_err(|e| ort::Error::new(e.to_string()))?,
|
||
];
|
||
let key_ptrs: Vec<_> = keys.iter().map(|k| k.as_ptr()).collect();
|
||
let value_ptrs: Vec<_> = values.iter().map(|v| v.as_ptr()).collect();
|
||
// SAFETY: the documented C call over arrays that outlive it; the
|
||
// runtime copies the strings into its own options map before returning.
|
||
unsafe {
|
||
let status = (ort::api().SessionOptionsAppendExecutionProvider)(
|
||
b.ptr_mut(),
|
||
c"MIGraphX".as_ptr(),
|
||
key_ptrs.as_ptr(),
|
||
value_ptrs.as_ptr(),
|
||
keys.len(),
|
||
);
|
||
ort::Error::result_from_status(status)
|
||
}
|
||
}
|
||
|
||
#[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 | Rung::MiGraphX | Rung::CoreMl => {
|
||
unreachable!("no desktop rung on Android")
|
||
}
|
||
}
|
||
}
|