inferencelayer 0.2.10

//! Engine arm of the layout-model A/B (torch twin: `tests/fixtures/bench_layout_torch.py`).
//! Loads the pinned heron checkpoint, times `detect()` (preprocess + forward + postprocess) on
//! the oracle's page PNG, prints `layout_detect_ms,<min-of-iters>` — kernel time only, so the
//! driver can interleave arms per PERF_CAMPAIGN.md.
//! Usage: layout-bench [iters=5] [cpu|gpu] [model_dir]

use std::time::Instant;

fn main() -> anyhow::Result<()> {
    let mut args = std::env::args().skip(1);
    let iters: usize = args.next().and_then(|s| s.parse().ok()).unwrap_or(5);
    let arm = args.next().unwrap_or_else(|| "cpu".into());
    let dir = args.next().unwrap_or_else(|| {
        "/Users/dlo/code/reasoninglayer-service/models/docling/layout-heron".into()
    });

    let fixtures = std::path::PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("tests/fixtures");
    let meta: serde_json::Value =
        serde_json::from_slice(&std::fs::read(fixtures.join("docling_layout_oracle.json"))?)?;
    let page = meta["full_stage_page"].as_u64().unwrap().to_string();
    let png_b64 = meta["pages"][&page]["png_b64"].as_str().unwrap();
    let png = b64(png_b64);

    let (rgb, w, h) = inferencelayer::vision::decode_rgb8(&png)?;
    let model = inferencelayer::rtdetr::RtDetr::load(std::path::Path::new(&dir))?;

    let mut best = f64::INFINITY;
    if arm == "gpu" {
        let ctx = inferencelayer::GpuCtx::new()?;
        let gpu = inferencelayer::rtdetr_gpu::RtDetrGpu::new(&ctx, model)?;
        let dets = gpu.detect(&ctx, &rgb, w, h)?; // warmup + sanity
        eprintln!("# {} detections on the fixture page (gpu)", dets.len());
        for _ in 0..iters {
            let t = Instant::now();
            let d = gpu.detect(&ctx, &rgb, w, h)?;
            let dt = t.elapsed().as_secs_f64() * 1e3;
            std::hint::black_box(&d);
            best = best.min(dt);
        }
    } else {
        let dets = model.detect(&rgb, w, h); // warmup + sanity
        eprintln!("# {} detections on the fixture page (cpu)", dets.len());
        for _ in 0..iters {
            let t = Instant::now();
            let d = model.detect(&rgb, w, h);
            let dt = t.elapsed().as_secs_f64() * 1e3;
            std::hint::black_box(&d);
            best = best.min(dt);
        }
    }
    println!("layout_detect_ms,{best:.2}");
    Ok(())
}

fn b64(s: &str) -> Vec<u8> {
    let val = |c: u8| -> i32 {
        match c {
            b'A'..=b'Z' => (c - b'A') as i32,
            b'a'..=b'z' => (c - b'a') as i32 + 26,
            b'0'..=b'9' => (c - b'0') as i32 + 52,
            b'+' => 62,
            b'/' => 63,
            _ => -1,
        }
    };
    let mut out = Vec::new();
    let (mut acc, mut bits) = (0u32, 0u32);
    for &c in s.as_bytes() {
        let v = val(c);
        if v < 0 {
            continue;
        }
        acc = (acc << 6) | v as u32;
        bits += 6;
        if bits >= 8 {
            bits -= 8;
            out.push((acc >> bits) as u8);
        }
    }
    out
}