rlx_runtime/
device_bench.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3// SPDX-License-Identifier: MIT OR Apache-2.0
4
5//! Warm-all and micro-benchmark helpers for backend ranking.
6
7use rlx_driver::Device;
8use rlx_ir::Tick;
9
10use crate::device_parse::device_label;
11use crate::graph_devices::GraphDevices;
12
13/// One backend timed after compile + repeated execution.
14#[derive(Debug, Clone, PartialEq)]
15pub struct DeviceBenchResult {
16    pub device: Device,
17    pub label: &'static str,
18    pub compile_ns: u64,
19    pub median_exec_ns: u64,
20}
21
22/// Compile every viable backend once.
23pub fn warm_all(runner: &mut GraphDevices) -> Result<Vec<Device>, String> {
24    let devices: Vec<Device> = runner.devices().to_vec();
25    for device in &devices {
26        runner.compile(*device)?;
27    }
28    Ok(devices)
29}
30
31/// Compile (if needed) and time each backend; returns results sorted fastest first.
32pub fn benchmark_devices(
33    runner: &mut GraphDevices,
34    inputs: &[(&str, &[f32])],
35    runs: usize,
36) -> Result<Vec<DeviceBenchResult>, String> {
37    let mut results = Vec::new();
38    let devices: Vec<Device> = runner.devices().to_vec();
39    for device in devices {
40        let t0 = Tick::now();
41        runner.compile(device)?;
42        let compile_ns = Tick::now().elapsed_ns(t0);
43        let mut samples = Vec::with_capacity(runs.max(1));
44        for _ in 0..runs.max(1) {
45            let t1 = Tick::now();
46            runner.run(device, inputs)?;
47            samples.push(Tick::now().elapsed_ns(t1));
48        }
49        samples.sort_unstable();
50        let median_exec_ns = samples[samples.len() / 2];
51        results.push(DeviceBenchResult {
52            device,
53            label: device_label(device),
54            compile_ns,
55            median_exec_ns,
56        });
57    }
58    results.sort_by_key(|r| r.median_exec_ns);
59    Ok(results)
60}
rlx_runtime/device_bench.rs

rlx_runtime/
device_bench.rs
