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

//! Compiled engines: what a rung builds once per device, and the thread that
//! builds them before anyone asks (docs/dev/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()
                );
            }
        }
    }
}