environment/training-source/source-snapshot/examples/bench.rs

//! Desktop benchmark runner.
//!
//! The same [`dinovision::bench`] code runs on the Quest through the
//! `android` crate; running it here first gives a reference point from a
//! GPU that *does* have cooperative matrix, which makes the Adreno numbers
//! interpretable rather than just small.
//!
//! ```text
//! cargo run --release --example bench -- [iters]
//! DINOVISION_BENCH=shapes cargo run --release --example bench -- [iters]
//! ```
//!
//! With `DINOVISION_BENCH=shapes`, runs only the matmul shape sweep — the diagnostic for
//! what the register-tiled kernel is short of. It takes a couple of minutes
//! instead of the full suite's twenty, which matters when the target is a
//! headset on a flaky wireless adb link.

fn main() {
    env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info")).init();

    let iters = std::env::args()
        .nth(1)
        .and_then(|s| s.parse().ok())
        .unwrap_or(20);
    // Selected by environment, not by argument: `std::env::args()` does not
    // survive the trip to an adb shell on Horizon OS, and silently running
    // the wrong benchmark is worse than an ugly interface. `RUST_LOG`
    // already proves the environment arrives intact.
    let shapes_only = std::env::var("DINOVISION_BENCH").as_deref() == Ok("shapes");

    let gpu = dinovision::init_context(None).expect("failed to initialize GPU context");
    let results = if shapes_only {
        dinovision::bench::describe_device(&gpu);
        dinovision::bench::matmul_shapes(gpu, iters)
    } else {
        dinovision::bench::run_all(gpu, iters)
    };

    println!("\n{:-<80}", "");
    for t in &results {
        println!("{t}");
        println!("DINOVISION_BENCH_JSON {}", t.json_line());
    }
}