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

//! Throughput measurement, shared by the desktop and on-device entry
//! points.
//!
//! The question this exists to answer: the Adreno in a Quest 3 is not
//! expected to expose `VK_KHR_cooperative_matrix`, so meganeura will fall
//! back to its register-tiled scalar matmul. Whether the encoder then
//! lands at 20 ms or 200 ms per frame decides the input resolution, the
//! inference rate, and whether the render loop has to be decoupled from
//! inference (it almost certainly does).
//!
//! Everything here reports achieved GFLOP/s next to the theoretical MAC
//! count, so a number that looks fast can be checked against what the
//! device could possibly do.

use std::sync::Arc;
use std::time::Instant;

use meganeura::graph::Op;
use meganeura::train::{Mode, SessionConfig};
use meganeura::{Graph, Session};

use crate::dinov3::Config;

/// One timing result.
#[derive(Debug, Clone)]
pub struct Timing {
    pub label: String,
    /// Minimum retained sample. Useful diagnostically, but not the primary
    /// result: selecting the minimum biases comparisons on a noisy device.
    pub best_ms: f64,
    pub mean_ms: f64,
    pub median_ms: f64,
    pub p25_ms: f64,
    pub p75_ms: f64,
    pub max_ms: f64,
    /// Every retained synchronized wall-time observation.
    pub samples_ms: Vec<f64>,
    /// Multiply-accumulates per iteration, for the GFLOP/s figure.
    pub macs: u64,
}

impl Timing {
    /// Median achieved GFLOP/s, counting a multiply-accumulate as two
    /// operations.
    pub fn gflops(&self) -> f64 {
        (self.macs as f64 * 2.0) / (self.median_ms / 1000.0) / 1e9
    }

    /// One self-contained JSON record suitable for retaining from logcat.
    pub fn json_line(&self) -> String {
        let label = self.label.replace('\\', "\\\\").replace('"', "\\\"");
        let samples = self
            .samples_ms
            .iter()
            .map(|v| format!("{v:.6}"))
            .collect::<Vec<_>>()
            .join(",");
        format!(
            "{{\"schema_version\":1,\"kind\":\"dinovision_benchmark\",\
             \"label\":\"{label}\",\"macs\":{},\"median_ms\":{:.6},\
             \"p25_ms\":{:.6},\"p75_ms\":{:.6},\"min_ms\":{:.6},\
             \"max_ms\":{:.6},\"mean_ms\":{:.6},\"samples_ms\":[{samples}]}}",
            self.macs,
            self.median_ms,
            self.p25_ms,
            self.p75_ms,
            self.best_ms,
            self.max_ms,
            self.mean_ms,
        )
    }
}

impl std::fmt::Display for Timing {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "{:<28} median {:>8.2} ms  IQR [{:>7.2}, {:>7.2}]  min/max [{:>7.2}, {:>7.2}]  {:>7.1} GFLOP/s",
            self.label,
            self.median_ms,
            self.p25_ms,
            self.p75_ms,
            self.best_ms,
            self.max_ms,
            self.gflops()
        )
    }
}

/// Log the device identity and whether cooperative matrix is usable.
///
/// This is the single most important line of the whole benchmark: it
/// determines which matmul kernel every subsequent number came from.
pub fn describe_device(gpu: &blade_graphics::Context) {
    let info = gpu.device_information();
    log::info!("device       : {}", info.device_name);
    log::info!("driver       : {} {}", info.driver_name, info.driver_info);

    let caps = meganeura::runtime::auto_tune(gpu, 0).coop_caps;
    if caps.f16_tile == 0 && caps.f32_tile == 0 {
        log::warn!("coop matrix  : NOT AVAILABLE — expect the register-tiled scalar matmul");
    } else {
        log::info!(
            "coop matrix  : available (f32 tile {}, f16 tile {})",
            caps.f32_tile,
            caps.f16_tile
        );
    }
}

/// Fill every parameter with small deterministic pseudo-random values.
///
/// Timing is not sensitive to weight *values*, but leaving parameters at
/// zero would make every activation zero, and a benchmark whose data is
/// entirely one value invites doubt about data-dependent fast paths.
/// Cheap insurance.
pub fn fill_parameters(session: &mut Session, graph: &Graph) {
    let mut state = 0x2545_F491_4F6C_DD1Du64;
    let mut scratch = Vec::new();
    for node in graph.nodes() {
        let Op::Parameter { name } = &node.op else {
            continue;
        };
        scratch.clear();
        scratch.reserve(node.ty.num_elements());
        for _ in 0..node.ty.num_elements() {
            // xorshift64*, inlined to avoid a dependency.
            state ^= state >> 12;
            state ^= state << 25;
            state ^= state >> 27;
            let bits = state.wrapping_mul(0x2545_F491_4F6C_DD1D) >> 40;
            scratch.push((bits as f32 / (1u32 << 24) as f32 - 0.5) * 0.05);
        }
        session.set_parameter(name, &scratch);
    }
}

pub(crate) fn time_session(session: &mut Session, label: &str, macs: u64, iters: usize) -> Timing {
    // Five is the paper protocol's minimum. The environment override exists
    // for diagnostic sweeps, and is recorded alongside the raw samples by
    // the calling artifact script.
    let warmups = std::env::var("DINOVISION_WARMUPS")
        .ok()
        .and_then(|s| s.parse::<usize>().ok())
        .unwrap_or(5);
    for _ in 0..warmups {
        session.step();
        session.wait();
    }

    let mut samples_ms = Vec::with_capacity(iters.max(1));
    for _ in 0..iters.max(1) {
        let start = Instant::now();
        session.step();
        session.wait();
        samples_ms.push(start.elapsed().as_secs_f64() * 1000.0);
    }

    let mut sorted = samples_ms.clone();
    sorted.sort_by(f64::total_cmp);
    let quantile = |q: f64| {
        let at = q * (sorted.len() - 1) as f64;
        let lo = at.floor() as usize;
        let hi = at.ceil() as usize;
        sorted[lo] + (sorted[hi] - sorted[lo]) * (at - lo as f64)
    };
    Timing {
        label: label.to_string(),
        best_ms: sorted[0],
        mean_ms: samples_ms.iter().sum::<f64>() / samples_ms.len() as f64,
        median_ms: quantile(0.5),
        p25_ms: quantile(0.25),
        p75_ms: quantile(0.75),
        max_ms: sorted[sorted.len() - 1],
        samples_ms,
        macs,
    }
}

/// Time a bare `[m, k] @ [k, n]` matmul.
///
/// Run at the shapes the encoder actually uses, so the result is an upper
/// bound on what the encoder could reach rather than a peak-throughput
/// number from an unrepresentatively large matrix.
pub fn matmul_throughput(
    gpu: Arc<blade_graphics::Context>,
    m: usize,
    k: usize,
    n: usize,
    iters: usize,
) -> Timing {
    let mut g = Graph::new();
    let a = g.input("a", &[m, k]);
    let b = g.parameter("b", &[k, n]);
    let y = g.matmul(a, b);
    g.set_outputs(vec![y]);

    let (mut session, _) = meganeura::train::build(
        &g,
        SessionConfig {
            mode: Mode::Inference,
            gpu: Some(gpu),
            ..Default::default()
        },
    );
    fill_parameters(&mut session, &g);
    session.set_input("a", &vec![0.01f32; m * k]);

    time_session(
        &mut session,
        &format!("matmul {m}x{k}x{n}"),
        (m * k * n) as u64,
        iters,
    )
}

/// Sweep matmul shapes to find what the kernel is actually short of.
///
/// This diagnostic separates three plausible sources of lost utilization;
/// no operation-share or f16 conclusion is inferred from it without a
/// matched end-to-end run:
///
/// * **Tail waste.** `m` is 201 — 196 patches, CLS, 4 registers — against
///   `BM = 64`, so the last tile row is 9/64 useful and 21% of every
///   dispatch is padding. Padding `m` to 256 does 27% more arithmetic; if
///   the wall time does not move, we were paying for it already.
/// * **Occupancy.** Tiles are 64×64 with `KTILE = 32`, so a workgroup holds
///   16.4 KB of shared memory and few stay resident. Widening `n` at fixed
///   cost per workgroup raises the workgroup count; if throughput climbs
///   with it, the device is starved rather than saturated. This is also the
///   one thing f16 would genuinely fix, by halving the LDS footprint.
/// * **Prologue cost.** `k = 384` is 12 k-tiles, so the per-tile load and
///   double barrier amortize over very few iterations. Deepening `k` at the
///   same output size isolates that.
///
/// The large square shapes at the end give the kernel's ceiling with every
/// one of those effects removed, which is the number f16 arithmetic would
/// have to improve on to be worth the work.
pub fn matmul_shapes(gpu: Arc<blade_graphics::Context>, iters: usize) -> Vec<Timing> {
    let mut results = Vec::new();
    for (m, k, n, note) in [
        (201, 384, 1536, "encoder MLP up, as it runs"),
        (256, 384, 1536, "m padded to a whole tile: +27% MACs"),
        (201, 384, 384, "24 workgroups"),
        (201, 384, 1536, "96 workgroups"),
        (201, 384, 6144, "384 workgroups"),
        (201, 1536, 1536, "4x the k-tiles per workgroup"),
        (512, 512, 512, "square, no tail"),
        (1024, 1024, 1024, "square, no tail"),
        (2048, 2048, 2048, "square, kernel ceiling"),
    ] {
        let t = matmul_throughput(gpu.clone(), m, k, n, iters);
        log::info!("{t}   ({note})");
        results.push(t);
    }

    // Probe the tail-waste hypothesis. Every token-major matmul has `m = tokens`, and
    // 201 tokens occupy four 64-row tiles — 256 rows, of which 55 are
    // padding we already pay for. A 240² input gives a 15×15 grid, so 225
    // patches plus 5 prefix tokens is 230 rows: still four tiles, still the
    // same dispatch, but 15% more spatial resolution reaching the decoder.
    //
    // Not entirely free — attention scores are O(tokens²) and they are ~8%
    // of the layer — so measure it rather than assert it.
    // Bracket 240 with repeated 224 cells so time drift is visible rather than
    // silently attributed to the middle shape.
    for size in [224, 240, 224] {
        let c = Config::vits16().at_resolution(size).with_layers(3);
        let mut t = encoder_forward(gpu.clone(), &c, iters);
        t.label = format!("dinov3 @{size} 3L");
        log::info!(
            "{t}   (grid {}x{}, {} tokens)",
            c.grid(),
            c.grid(),
            c.num_tokens()
        );
        results.push(t);
    }
    results
}

/// Build an encoder session on the shared context.
///
/// Weights are left to the caller: [`crate::weights::load_encoder`] for
/// real features, [`fill_parameters`] for pure timing.
pub fn build_encoder_session(
    gpu: Arc<blade_graphics::Context>,
    config: &Config,
    cache: Option<&std::path::Path>,
) -> (Session, Graph) {
    let mut g = Graph::new();
    let out = crate::dinov3::build_encoder(&mut g, config);
    g.set_outputs(vec![out]);

    let build_start = Instant::now();
    let (session, _) = meganeura::train::build(
        &g,
        SessionConfig {
            mode: Mode::Inference,
            gpu: Some(gpu),
            cache,
            ..Default::default()
        },
    );
    // Worth reporting separately: on a mobile CPU, graph optimization plus
    // WGSL-to-SPIR-V for every kernel is slow enough to matter at startup,
    // and it is exactly what the plan cache is meant to remove.
    log::info!(
        "session build: {:.1} s (cache {})",
        build_start.elapsed().as_secs_f64(),
        if cache.is_some() { "on" } else { "off" }
    );
    (session, g)
}

/// Time a full DINOv3 forward pass with synthetic weights.
pub fn encoder_forward(gpu: Arc<blade_graphics::Context>, config: &Config, iters: usize) -> Timing {
    let (mut session, g) = build_encoder_session(gpu, config, None);
    fill_parameters(&mut session, &g);
    session.set_input(
        "patches",
        &vec![0.1f32; config.num_patches() * config.patch_dim()],
    );

    time_session(
        &mut session,
        &format!("dinov3 vits16 @{}", config.image_size),
        config.forward_macs(),
        iters,
    )
}

/// Time the deployed joined encoder-to-RGB graph, including decoder and all
/// graph-internal layout transformations. Camera conversion, CPU patchifying,
/// output readback, and rendering are outside this isolated measurement.
pub fn roundtrip_forward(
    gpu: Arc<blade_graphics::Context>,
    config: &Config,
    submission_chunks: usize,
    iters: usize,
) -> Timing {
    let mut graph = Graph::new();
    let encoder_output = crate::dinov3::build_encoder(&mut graph, config);
    let reconstruction = crate::decoder::attach_to_encoder(&mut graph, config, encoder_output);
    graph.set_outputs(vec![reconstruction]);
    let (mut session, _) = meganeura::train::build(
        &graph,
        SessionConfig {
            mode: Mode::Inference,
            gpu: Some(gpu),
            ..Default::default()
        },
    );
    session.set_submission_chunks(submission_chunks);
    fill_parameters(&mut session, &graph);
    session.set_input(
        "patches",
        &vec![0.1f32; config.num_patches() * config.patch_dim()],
    );
    time_session(
        &mut session,
        &format!(
            "roundtrip @{} {}L x{}",
            config.image_size, config.num_hidden_layers, submission_chunks
        ),
        config.forward_macs() + crate::decoder::forward_macs(config),
        iters,
    )
}

/// The full benchmark: device identity, matmul throughput at the shapes
/// the encoder uses, then the encoder itself across candidate input
/// resolutions.
pub fn run_all(gpu: Arc<blade_graphics::Context>, iters: usize) -> Vec<Timing> {
    describe_device(&gpu);

    let base = Config::vits16();
    let tokens = base.num_tokens();
    let hidden = base.hidden_size;
    let mut results = Vec::new();

    // The three matmul shapes that dominate a ViT-S forward pass.
    for (m, k, n) in [
        (tokens, hidden, hidden),                 // Q/K/V/output projection
        (tokens, hidden, base.intermediate_size), // MLP up
        (tokens, base.intermediate_size, hidden), // MLP down
    ] {
        let t = matmul_throughput(gpu.clone(), m, k, n, iters);
        log::info!("{t}");
        results.push(t);
    }

    // 224 is the reference resolution (14×14 features); 256 buys a 16×16
    // grid for ~30% more work. Anything larger is unlikely to fit a
    // real-time budget, which is what these numbers will confirm or deny.
    for size in [224, 256] {
        let config = base.clone().at_resolution(size);
        let t = encoder_forward(gpu.clone(), &config, iters);
        log::info!(
            "{t}   (grid {}x{}, {} tokens)",
            config.grid(),
            config.grid(),
            config.num_tokens()
        );
        results.push(t);
    }

    // ViT-S/16 is the smallest DINOv3 there is, so cheaper means fewer
    // layers or fewer tokens rather than a smaller model. Measure both
    // knobs, since which one to spend is a quality judgement that wants
    // real numbers under it.
    for layers in [6, 3] {
        let c = base.clone().at_resolution(224).with_layers(layers);
        let mut t = encoder_forward(gpu.clone(), &c, iters);
        t.label = format!("dinov3 @224 {layers}L");
        log::info!("{t}   ({layers} of 12 layers)");
        results.push(t);
    }
    for size in [160, 128] {
        let c = base.clone().at_resolution(size);
        let mut t = encoder_forward(gpu.clone(), &c, iters);
        t.label = format!("dinov3 vits16 @{size}");
        log::info!("{t}   (grid {}x{})", c.grid(), c.grid());
        results.push(t);
    }

    // The actual reconstruction workload is a joined three-layer encoder and
    // decoder. Keep isolated chunk costs beside the live-XR sweep so queue
    // fairness is not presented without its throughput price.
    let reconstruction = base.clone().at_resolution(224).with_layers(3);
    // The order is fixed before corrected device data and deliberately not
    // monotonic, matching the live sweep's protection against time drift.
    for chunks in [4, 16, 1, 12, 2, 8] {
        let t = roundtrip_forward(gpu.clone(), &reconstruction, chunks, iters);
        log::info!("{t}   (joined encoder + decoder, {chunks} submission chunks)");
        results.push(t);
    }

    // f16 weight storage halves weight traffic. Whether that shows up
    // depends on whether the device is short of bandwidth or of ALU, which
    // is exactly the thing worth measuring rather than reasoning about.
    let f16 = base.clone().at_resolution(224).with_f16_weights(true);
    let mut t = encoder_forward(gpu.clone(), &f16, iters);
    t.label = "dinov3 vits16 @224 f16w".to_string();
    log::info!("{t}   (f16 weight storage, f32 arithmetic)");
    results.push(t);

    // Chunked submission should cost a little throughput — several submits
    // instead of one — in exchange for letting a co-tenant interleave. This
    // measures the price when nothing else is on the queue.
    let chunked = base.clone().at_resolution(224);
    let (mut session, g) = build_encoder_session(gpu, &chunked, None);
    session.set_submission_chunks(12);
    fill_parameters(&mut session, &g);
    session.set_input(
        "patches",
        &vec![0.1f32; chunked.num_patches() * chunked.patch_dim()],
    );
    let t = time_session(
        &mut session,
        "dinov3 vits16 @224 x12",
        chunked.forward_macs(),
        iters,
    );
    log::info!("{t}   (12 submissions instead of 1)");
    results.push(t);

    results
}