bench_latency/
bench_latency.rs

1// Common-model embedding latency for RLX (all-MiniLM-L6-v2), matching the
2// PyTorch/ONNX-Runtime harness: synthetic batch x seq inputs, p50 of forward.
3// Run: MINILM_DIR=/tmp/minilm6 cargo run -p rlx-embed --release --example bench_latency
4use rlx_embed::{Pooling, RlxEmbed};
5use rlx_runtime::Device;
6use std::path::Path;
7use std::time::Instant;
8
9fn main() -> anyhow::Result<()> {
10    let dir = std::env::var("MINILM_DIR").unwrap_or_else(|_| "/tmp/minilm6".into());
11    let devname = std::env::var("DEVICE").unwrap_or_else(|_| "cpu".into());
12    let device = match devname.as_str() {
13        "metal" => Device::Metal,
14        "mlx" => Device::Mlx,
15        "gpu" | "wgpu" => Device::Gpu,
16        "vulkan" => Device::Vulkan,
17        "ane" | "coreml" => Device::Ane,
18        _ => Device::Cpu,
19    };
20    let mut model = RlxEmbed::from_dir_on(Path::new(&dir), Pooling::Mean, device)?;
21    let label = format!("rlx-{devname}");
22    let seq: usize = std::env::var("SEQ")
23        .ok()
24        .and_then(|s| s.parse().ok())
25        .unwrap_or(128);
26    let (runs, warmup) = (50usize, 10usize);
27    println!("framework,batch,p50_ms");
28    for &b in &[1usize, 4, 8, 16, 32] {
29        let n = b * seq;
30        let ids: Vec<f32> = (0..n).map(|i| (i * 131 % 30000) as f32).collect();
31        let mask = vec![1.0f32; n];
32        let tt = vec![0.0f32; n];
33        let pos: Vec<f32> = (0..b).flat_map(|_| (0..seq).map(|i| i as f32)).collect();
34        let inputs = [
35            ("input_ids", ids.as_slice()),
36            ("attention_mask", mask.as_slice()),
37            ("token_type_ids", tt.as_slice()),
38            ("position_ids", pos.as_slice()),
39        ];
40        for _ in 0..warmup {
41            let _ = model.forward(&inputs, b, seq)?;
42        }
43        let mut t: Vec<f64> = Vec::with_capacity(runs);
44        for _ in 0..runs {
45            let s = Instant::now();
46            let _ = model.forward(&inputs, b, seq)?;
47            t.push(s.elapsed().as_secs_f64() * 1e3);
48        }
49        t.sort_by(|a, c| a.partial_cmp(c).unwrap());
50        let p50 = t[t.len() / 2];
51        eprintln!("{label}  b={b:>3}  p50={p50:7.2} ms");
52        println!("{label},{b},{p50:.2}");
53    }
54    Ok(())
55}
bench_latency/bench_latency.rs

bench_latency/
bench_latency.rs
