Skip to main content

startup_bench/
startup_bench.rs

1//! Fresh-process startup and compilation costs. Use an empty XDG_CACHE_HOME
2//! for each invocation; this does not flush the operating system's page cache.
3//! Pass --compiler-only to measure CPU compilation without opening a GPU.
4use hrx::{Result, Stream, loom::Specialization};
5use serde_json::{Value, json};
6use std::time::Instant;
7
8fn memory() -> Result<Value> {
9    let status = std::fs::read_to_string("/proc/self/status")?;
10    let kib = |key: &str| -> Result<u64> {
11        status
12            .lines()
13            .find_map(|line| line.strip_prefix(key))
14            .and_then(|value| value.split_whitespace().next())
15            .and_then(|value| value.parse().ok())
16            .ok_or_else(|| hrx::Error::Message(format!("missing {key} in process status")))
17    };
18    Ok(json!({"rss_kib": kib("VmRSS:")?, "peak_rss_kib": kib("VmHWM:")?}))
19}
20
21fn main() -> Result<()> {
22    let compiler_only = match std::env::args().nth(1).as_deref() {
23        None => false,
24        Some("--compiler-only") => true,
25        Some(_) => {
26            return Err(hrx::Error::Message(
27                "usage: startup_bench [--compiler-only]".into(),
28            ));
29        }
30    };
31    let initial = memory()?;
32    let start = Instant::now();
33    let stream = if compiler_only {
34        None
35    } else {
36        Some(Stream::open()?)
37    };
38    let open_ms = stream
39        .as_ref()
40        .map(|_| start.elapsed().as_secs_f64() * 1000.0);
41    let after_open = if compiler_only { None } else { Some(memory()?) };
42    let start = Instant::now();
43    let compiler = match &stream {
44        Some(stream) => hrx::loom::Compiler::for_stream(None, stream)?,
45        None => hrx::loom::Compiler::resolve(None)?,
46    };
47    let compiler_ms = start.elapsed().as_secs_f64() * 1000.0;
48    let after_compiler = memory()?;
49    let module = compiler.module(include_str!("../tests/kernels/euler.loom"));
50    let mut spec = Specialization::new("krea2_euler");
51    spec.set_config("krea2.euler.grid_x", "1");
52    spec.set_config("krea2.euler.grid_y", "1");
53    let start = Instant::now();
54    let artifact = module.compile(&spec)?;
55    let first_compile_ms = start.elapsed().as_secs_f64() * 1000.0;
56    let after_compile = memory()?;
57    let start = Instant::now();
58    let cached = module.compile(&spec)?;
59    let cache_hit_ms = start.elapsed().as_secs_f64() * 1000.0;
60    assert_eq!(artifact.bytes(), cached.bytes());
61    println!(
62        "{}",
63        json!({"compiler_only": compiler_only, "open_ms": open_ms, "compiler_ms": compiler_ms,
64            "first_compile_ms": first_compile_ms, "cache_hit_ms": cache_hit_ms,
65            "initial": initial, "after_open": after_open,
66            "after_compiler": after_compiler, "after_compile": after_compile,
67            "compiler_identity": artifact.compiler_identity()})
68    );
69    Ok(())
70}