dtourolle
031ba7b77d
A border fill acquires the filler once a tile, and each acquire hashed the 28 MB model twice — 60 ms a tile, a third of the tile's run on a throttled TensorRT. The Model keeps its hash from open.
122 lines
4.2 KiB
Rust
122 lines
4.2 KiB
Rust
//! Compiled engines: what a rung builds once per device, and the thread that
|
|
//! builds them before anyone asks (docs/inference.md §5, §6).
|
|
//!
|
|
//! TensorRT keeps its own engine cache keyed by graph hash; QNN writes a
|
|
//! context model. Both are opaque to this crate, which tracks only *that* a
|
|
//! model compiled — by the hash of its bytes — so [`crate::open`] can tell a
|
|
//! request whether to expect the rung or its fallback.
|
|
|
|
use std::path::PathBuf;
|
|
|
|
use crate::{state, Config, Form, Rung};
|
|
|
|
enum Source {
|
|
File(PathBuf),
|
|
Bytes(&'static [u8]),
|
|
}
|
|
|
|
/// 64-bit FNV-1a. A cache key, not a checksum: two model files that collide
|
|
/// here would have to also be the same size and the same role, and the cost
|
|
/// of that is a rebuilt engine.
|
|
pub fn hash(bytes: &[u8]) -> u64 {
|
|
let mut h = 0xcbf2_9ce4_8422_2325u64;
|
|
for &b in bytes {
|
|
h ^= b as u64;
|
|
h = h.wrapping_mul(0x0000_0100_0000_01b3);
|
|
}
|
|
h
|
|
}
|
|
|
|
/// The cache entry for `bytes` compiled on `rung`.
|
|
pub fn key(rung: Rung, bytes: &[u8]) -> String {
|
|
key_of(rung, hash(bytes))
|
|
}
|
|
|
|
/// The same, from a hash already taken.
|
|
pub fn key_of(rung: Rung, hash: u64) -> String {
|
|
format!("{}:{:016x}", rung.label(), hash)
|
|
}
|
|
|
|
/// Where QNN's compiled context for `bytes` lives.
|
|
pub fn context_path(cfg: &Config, bytes: &[u8]) -> PathBuf {
|
|
cfg.cache_dir
|
|
.join("qnn")
|
|
.join(format!("{:016x}_ctx.onnx", hash(bytes)))
|
|
}
|
|
|
|
/// After the probe: compile every configured model the selected rung can
|
|
/// take, smallest first, recording each as it lands.
|
|
pub fn run() {
|
|
let (rung, cfg) = {
|
|
let s = state().lock().unwrap();
|
|
(crate::current_rung(&s), s.config.clone())
|
|
};
|
|
if !rung.compiles() {
|
|
return;
|
|
}
|
|
|
|
// Smallest first, so the detector — the one that runs per image — is
|
|
// ready soonest (§6 step 3).
|
|
let mut jobs: Vec<(crate::Role, Source, u64)> = cfg
|
|
.models
|
|
.iter()
|
|
.filter(|(role, _)| rung.serves(*role))
|
|
.filter_map(|(role, path)| {
|
|
let (path, form) = crate::resolve_model(*role, path);
|
|
(form == rung.form(*role)).then(|| {
|
|
let size = std::fs::metadata(&path).map(|m| m.len()).unwrap_or(0);
|
|
(*role, Source::File(path), size)
|
|
})
|
|
})
|
|
.chain(cfg.embedded.iter().filter_map(|(role, bytes)| {
|
|
// An embedded model has no int8 sibling to offer a rung that
|
|
// wants one; it runs on that rung's fallback.
|
|
(rung.serves(*role) && rung.form(*role) == Form::F32).then_some((
|
|
*role,
|
|
Source::Bytes(bytes),
|
|
bytes.len() as u64,
|
|
))
|
|
}))
|
|
.collect();
|
|
jobs.sort_by_key(|j| j.2);
|
|
state().lock().unwrap().wanted = jobs.len();
|
|
|
|
for (role, source, _) in jobs {
|
|
let (bytes, name) = match &source {
|
|
Source::File(path) => match std::fs::read(path) {
|
|
Ok(b) => (b, path.display().to_string()),
|
|
Err(_) => continue,
|
|
},
|
|
Source::Bytes(b) => (b.to_vec(), format!("embedded {role:?}")),
|
|
};
|
|
let key = key(rung, &bytes);
|
|
if state().lock().unwrap().cache.compiled.contains(&key) {
|
|
continue;
|
|
}
|
|
log::info!("inference: compiling {name} for {}", rung.label());
|
|
let started = std::time::Instant::now();
|
|
match crate::session::build(rung, role, &bytes, &cfg) {
|
|
Ok(session) => {
|
|
drop(session);
|
|
let mut s = state().lock().unwrap();
|
|
s.cache.compiled.insert(key);
|
|
crate::probe::write_cache(&s.config, &s.cache);
|
|
log::info!(
|
|
"inference: {name} ready on {} in {:.1} s",
|
|
rung.label(),
|
|
started.elapsed().as_secs_f64()
|
|
);
|
|
}
|
|
Err(e) => {
|
|
// This model stays on the fallback; the others still get
|
|
// their engine. A corrected model file changes the hash and
|
|
// is retried.
|
|
log::warn!(
|
|
"inference: {name} will not compile for {}: {e}",
|
|
rung.label()
|
|
);
|
|
}
|
|
}
|
|
}
|
|
}
|