benchmark/
benchmark.rs

1//! REVE RLX inference benchmark.
2//!
3//! Usage:
4//!   benchmark [--device cpu|metal|mlx] [--config PATH] [--weights PATH] \
5//!       <n_chans> <n_times> <warmup> <repeats>
6//!
7//! Prints JSON to stdout:
8//!   {"times_ms": [..], "backend": "rlx-cpu", "ok": true, "output_dim": 512}
9
10use std::path::PathBuf;
11use std::time::Instant;
12
13use clap::Parser;
14use reve_rs::rlx::ReveEncoder;
15
16#[derive(Parser, Debug)]
17#[command(about = "REVE RLX inference benchmark")]
18struct Args {
19    #[arg(long, default_value = "cpu")]
20    device: String,
21
22    #[arg(long)]
23    config: Option<PathBuf>,
24
25    #[arg(long)]
26    weights: Option<PathBuf>,
27
28    n_chans: usize,
29    n_times: usize,
30    warmup: usize,
31    repeats: usize,
32}
33
34fn locate_paths(config: Option<PathBuf>, weights: Option<PathBuf>) -> anyhow::Result<(PathBuf, PathBuf)> {
35    let manifest = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
36    let config = config.unwrap_or_else(|| manifest.join("data/config.json"));
37    let weights = weights.unwrap_or_else(|| manifest.join("data/model.safetensors"));
38    anyhow::ensure!(config.exists(), "config not found: {}", config.display());
39    anyhow::ensure!(weights.exists(), "weights not found: {}", weights.display());
40    Ok((config, weights))
41}
42
43fn parse_device(s: &str) -> anyhow::Result<rlx::Device> {
44    Ok(match s.to_lowercase().as_str() {
45        "cpu" => rlx::Device::Cpu,
46        "metal" => rlx::Device::Metal,
47        "mlx" => rlx::Device::Mlx,
48        other => anyhow::bail!("unsupported device '{other}' (use cpu, metal, or mlx)"),
49    })
50}
51
52fn backend_name(dev: rlx::Device) -> &'static str {
53    match dev {
54        rlx::Device::Cpu => "rlx-cpu",
55        rlx::Device::Metal => "rlx-metal",
56        rlx::Device::Mlx => "rlx-mlx",
57        _ => "rlx-other",
58    }
59}
60
61fn main() -> anyhow::Result<()> {
62    let args = Args::parse();
63    let dev = parse_device(&args.device)?;
64    let (config, weights) = locate_paths(args.config, args.weights)?;
65
66    let (mut enc, _) = ReveEncoder::load(&config, &weights, dev)?;
67
68    let n_chans = args.n_chans;
69    let n_times = args.n_times;
70    let signal = vec![0.0f32; n_chans * n_times];
71    let positions = vec![0.0f32; n_chans * 3];
72
73    for _ in 0..args.warmup {
74        let out = enc.run_one(signal.clone(), positions.clone(), n_chans, n_times)?;
75        anyhow::ensure!(
76            out.output.iter().all(|v| v.is_finite()),
77            "warmup produced non-finite values"
78        );
79    }
80
81    let mut times = Vec::with_capacity(args.repeats);
82    let mut last_dim = 0usize;
83    for _ in 0..args.repeats {
84        let t0 = Instant::now();
85        let out = enc.run_one(signal.clone(), positions.clone(), n_chans, n_times)?;
86        anyhow::ensure!(
87            out.output.iter().all(|v| v.is_finite()),
88            "run produced non-finite values"
89        );
90        last_dim = out.output.len();
91        times.push(t0.elapsed().as_secs_f64() * 1000.0);
92    }
93
94    let times_str: Vec<String> = times.iter().map(|t| format!("{t:.4}")).collect();
95    println!(
96        "{{\"times_ms\": [{}], \"backend\": \"{}\", \"ok\": true, \"output_dim\": {}}}",
97        times_str.join(", "),
98        backend_name(dev),
99        last_dim,
100    );
101    Ok(())
102}
benchmark/benchmark.rs

benchmark/
benchmark.rs
