open-standards/wai/wai-quantum/examples/fabric_timing.rs

//! The measurements `auto` places a state vector by
//! (`engine::fabric_measured`), taken again on this host: whole jobs of the
//! measured circuit on the CPU and on each member found here, their runs
//! alternating, at each fusion width and qubit count asked, printed as the
//! table's rows.
//!
//! ```sh
//! cargo run --release --example fabric_timing \
//!     --features quantum_engine,quantum_fabric,quantum_fabric_cuda,quantum_fabric_opencl -- \
//!     --runs 5 --qubits 18,20,22,24,26 --fusion 2,3 [--members key,key]
//! ```
//!
//! Every placed job's answers are checked against the CPU's, byte for byte,
//! before its time is kept. A size whose state would not fit in half the
//! host's free memory, or past what a member holds, is left out and said so.

use serde_json::Value;
use std::time::Instant;
use wai_quantum::engine::{catalog, fabric, fabric_measured, receipt, run};

fn arg(name: &str) -> Option<String> {
    let args: Vec<String> = std::env::args().collect();
    args.iter().position(|a| a == name).and_then(|i| args.get(i + 1).cloned())
}

fn list<T: std::str::FromStr>(name: &str, default: &str) -> Vec<T> {
    arg(name).unwrap_or_else(|| default.to_string()).split(',').filter(|s| !s.is_empty()).map(|s| s.parse().unwrap_or_else(|_| panic!("{name}: {s}"))).collect()
}

/// `MemAvailable` from the kernel, in bytes, or 8 GiB where it cannot be read.
fn available() -> u64 {
    std::fs::read_to_string("/proc/meminfo")
        .ok()
        .and_then(|m| m.lines().find(|l| l.starts_with("MemAvailable")).and_then(|l| l.split_whitespace().nth(1)).and_then(|k| k.parse::<u64>().ok()))
        .map_or(8 << 30, |kib| kib << 10)
}

fn load() -> String {
    std::fs::read_to_string("/proc/loadavg").map(|l| l.split_whitespace().take(3).collect::<Vec<_>>().join(" ")).unwrap_or_else(|_| "unread".into())
}

/// The answers, without where the state ran and how long it took.
fn answers(out: &Value) -> Vec<u8> {
    let mut o = out.clone();
    if let Some(m) = o.as_object_mut() {
        m.remove("fabric");
        m.remove("elapsed_ms");
    }
    receipt::canonical(&o)
}

/// One timed job, and its answers.
fn timed(v: &Value, cap: u64) -> (f64, Vec<u8>) {
    let t = Instant::now();
    let p = catalog::plan(v, cap).unwrap_or_else(|e| panic!("plan: {e}"));
    let out = run::run(v, &p).unwrap_or_else(|e| panic!("run: {e}"));
    (t.elapsed().as_secs_f64() * 1e3, answers(&out))
}

fn spread(v: &[f64]) -> (f64, f64, f64) {
    let mut s = v.to_vec();
    s.sort_by(f64::total_cmp);
    let median = if s.len() % 2 == 1 { s[s.len() / 2] } else { (s[s.len() / 2 - 1] + s[s.len() / 2]) / 2.0 };
    (s[0], median, s[s.len() - 1])
}

fn main() {
    let runs: usize = arg("--runs").map_or(5, |r| r.parse().expect("--runs"));
    let qubits: Vec<u32> = list("--qubits", "18,20,22,24,26");
    let fusions: Vec<u32> = list("--fusion", "3");
    let only: Vec<String> = list("--members", "");
    let cap = available() / 2;
    let t = Instant::now();
    let listed = fabric::list(cap);
    let discovery_ms = t.elapsed().as_secs_f64() * 1e3;
    let host = fabric::host();
    println!("host {host} discovery_ms {discovery_ms:.0} load {} memory_available {} MiB", load(), available() >> 20);
    let members: Vec<(String, String, u64)> = listed["members"]
        .as_array()
        .into_iter()
        .flatten()
        .filter(|m| m["key"] != "cpu")
        .map(|m| (m["member"].as_str().unwrap_or("").to_string(), m["key"].as_str().unwrap_or("").to_string(), m["detail"]["max_qubits"].as_u64().unwrap_or(0)))
        .filter(|(_, key, _)| only.is_empty() || only.contains(key))
        .collect();
    for (id, key, max) in &members {
        println!("member {key} {id} up to {max} qubits");
    }
    let mut rows = Vec::new();
    let mut warmed = false;
    for &fusion in &fusions {
        for &n in &qubits {
            // The CPU's state, a member's copy back and the job's own work,
            // with room to spare.
            if (1u64 << n) * 16 * 4 > cap {
                println!("skip fusion={fusion} qubits={n}: the state does not fit in half the free memory");
                continue;
            }
            let here: Vec<&(String, String, u64)> = members.iter().filter(|m| m.2 >= u64::from(n)).collect();
            if here.is_empty() {
                println!("skip fusion={fusion} qubits={n}: no member holds it");
                continue;
            }
            let cpu_job = fabric_measured::job(n, fusion, "cpu");
            let launches = fabric::launches(&cpu_job).expect("the job plans");
            if !warmed {
                timed(&cpu_job, cap);
                for m in &here {
                    timed(&fabric_measured::job(n, fusion, &m.0), cap);
                }
                warmed = true;
            }
            let before = load();
            let mut cpu = Vec::new();
            let mut per: Vec<Vec<f64>> = vec![Vec::new(); here.len()];
            for _ in 0..runs {
                let (ms, want) = timed(&cpu_job, cap);
                cpu.push(ms);
                for (i, m) in here.iter().enumerate() {
                    let (ms, got) = timed(&fabric_measured::job(n, fusion, &m.0), cap);
                    assert_eq!(got, want, "{}: fusion {fusion}, {n} qubits: the answers are not the CPU's", m.1);
                    per[i].push(ms);
                }
            }
            let after = load();
            let c = spread(&cpu);
            for (i, m) in here.iter().enumerate() {
                let g = spread(&per[i]);
                println!(
                    "row key={} fusion={fusion} qubits={n} launches={launches} runs={runs} cpu={:.1}/{:.1}/{:.1} member={:.1}/{:.1}/{:.1} faster={} load {before} -> {after}",
                    m.1,
                    c.0,
                    c.1,
                    c.2,
                    g.0,
                    g.1,
                    g.2,
                    g.2 < c.0
                );
                rows.push(format!(
                    "    Row {{ host: {host:?}, member: {:?}, key: {:?}, fusion: {fusion}, qubits: {n}, launches: {launches}, runs: {runs}, cpu_ms: Spread {{ min: {:.1}, median: {:.1}, max: {:.1} }}, member_ms: Spread {{ min: {:.1}, median: {:.1}, max: {:.1} }} }},",
                    m.0, m.1, c.0, c.1, c.2, g.0, g.1, g.2
                ));
            }
        }
    }
    println!("rows:");
    for r in rows {
        println!("{r}");
    }
}