open-standards/wai/tools/quantum-meter/src/main.rs

//! Calibrates the fabric's CPU reference and a WebGPU member with this
//! process's own energy counters, and places a circuit by the figures.
//!
//!   quantum-meter --counter-ppm <n> [--qubits <n>] [--seconds <s>] [--gateset 1|2] [--expect <adapter key>]
//!
//! `--counter-ppm` is the operator's evaluation of the counters' own relative
//! uncertainty, in parts per million. The platform publishes none, so there
//! is no default.
//!
//! One circuit on `--qubits` qubits (20) is run by each member, again and
//! again inside one window of `--seconds` seconds (3) at least: a window long
//! enough for the counters to resolve, where a single run is not. Each
//! member's line gives the window's readings and the figure they earn, in
//! nanojoules per operation (one pair of amplitudes under one gate), or says
//! why the window has none. With a figure on both, the last line is the member
//! the fabric places that circuit on.
//!
//! Exit 0 when it ran, with or without figures; 1 when a member could not be
//! found or run; 2 for a bad argument.

use quantum_meter::{Declared, ProcessRunMeter, Refused, Runs, Span, label, probe_load, span_of_runs};
use std::process::ExitCode;
use std::time::Duration;
use wai_quantum::quantum::{BaseGate, Circuit, Gate, Gateset};
use wai_quantum::quantum_energy::EnergyClass;
use wai_quantum::quantum_fabric::{EnergyCalibration, FABRIC_DEFAULT, Member, Parity, REFERENCE, apply_kernel, detect_local, pair_updates, place_for};
use wai_quantum::quantum_fabric_webgpu::blocking;

struct Args {
    declared: Declared,
    qubits: u8,
    at_least: Duration,
    gateset: Gateset,
    expect: Option<String>,
}

fn usage() -> ExitCode {
    eprintln!("usage: quantum-meter --counter-ppm <n> [--qubits <n>] [--seconds <s>] [--gateset 1|2] [--expect <adapter key>]");
    ExitCode::from(2)
}

fn args() -> Option<Args> {
    let (mut ppm, mut qubits, mut seconds, mut gateset, mut expect) = (None, 20u8, 3u64, Gateset::V2, None);
    let mut it = std::env::args().skip(1);
    while let Some(flag) = it.next() {
        let value = it.next()?;
        match flag.as_str() {
            "--counter-ppm" => ppm = Some(value.parse::<u32>().ok().filter(|&p| p > 0)?),
            "--qubits" => qubits = value.parse::<u8>().ok().filter(|n| (2..=24).contains(n))?,
            "--seconds" => seconds = value.parse::<u64>().ok().filter(|s| (1..=600).contains(s))?,
            "--gateset" => {
                gateset = match value.as_str() {
                    "1" => Gateset::V1,
                    "2" => Gateset::V2,
                    _ => return None,
                }
            }
            "--expect" => expect = Some(value),
            _ => return None,
        }
    }
    Some(Args { declared: Declared { counter_ppm: ppm? }, qubits, at_least: Duration::from_secs(seconds), gateset, expect })
}

/// The circuit every window runs: on each qubit a Hadamard, then a chain of
/// controlled X, then T and a dyadic phase on each, then a Hadamard on each.
/// Valid in either gateset; `5n − 1` gates.
fn circuit(n: u8, gs: Gateset) -> Circuit {
    let mut c = Circuit::with_gateset(n, gs);
    let one = |base, target, param| Gate { base, controls: Vec::new(), target, param };
    (0..n).for_each(|q| c.ops.push(one(BaseGate::H, q, 0)));
    (0..n - 1).for_each(|q| c.ops.push(Gate { base: BaseGate::X, controls: vec![q], target: q + 1, param: 0 }));
    for q in 0..n {
        c.ops.push(one(BaseGate::T, q, 0));
        c.ops.push(one(BaseGate::P, q, 1 + u16::from(q % 6)));
    }
    (0..n).for_each(|q| c.ops.push(one(BaseGate::H, q, 0)));
    c
}

/// One member's window as a line, and the figure it earns.
fn report(member: &str, kernel: &str, runs: &Runs, last: Option<&Result<Span, Refused>>, declared: Declared) -> Option<(u64, EnergyClass)> {
    let mut calibration = EnergyCalibration::new();
    calibration.add(runs.ops, runs.joules_micro, runs.class.as_ref());
    let figure = calibration.finish();
    let read = match last {
        Some(Ok(s)) => format!("cpu={} uJ gpu={} uJ gpu_time={} ms", s.cpu_uj, s.gpu_uj, s.gpu_ns / 1_000_000),
        _ => "no readings".to_string(),
    };
    let earned = match (&figure, last) {
        (Some((nanojoules, class)), _) => {
            let ppm = class.uncertainty().map_or(0, |u| u.relative_ppm);
            format!("{nanojoules} nJ/op, {}, {ppm} ppm", class.label())
        }
        (None, Some(Err(why))) => format!("unmetered: {why}"),
        // Read, and still no figure: the labelling rule says why, or the
        // fabric's resolution floor refused the label.
        (None, Some(Ok(s))) => match label(s, declared) {
            Err(why) => format!("unmetered: {why}"),
            Ok(_) => "unmetered: the reading is below the resolution floor".to_string(),
        },
        (None, None) => "unmetered".to_string(),
    };
    println!("quantum-meter: {member} kernel={kernel} runs={} ops={} wall={} ms {read} -> {earned}", runs.runs, runs.ops, runs.wall.as_millis());
    figure
}

fn main() -> ExitCode {
    let Some(args) = args() else { return usage() };
    let kernel = apply_kernel(args.gateset);
    let wide = circuit(args.qubits, args.gateset);
    let ops = pair_updates(&wide);
    println!("quantum-meter: circuit qubits={} gates={} ops={ops} kernel={kernel}", wide.n_qubits, wide.ops.len());

    let discovery = match blocking::discover() {
        Ok(d) => d,
        Err(e) => {
            eprintln!("quantum-meter: discovery failed: {e:?}");
            return ExitCode::FAILURE;
        }
    };
    for r in &discovery.refused {
        println!("quantum-meter: adapter {} ({}) is not a member: {:?}", r.key, r.adapter, r.error);
    }
    let found = match &args.expect {
        Some(key) => discovery.found.iter().find(|m| m.key() == key),
        None => discovery.found.first(),
    };
    let Some(gpu) = found else {
        eprintln!("quantum-meter: no WebGPU member{}", args.expect.as_ref().map(|k| format!(" with key {k}")).unwrap_or_default());
        return ExitCode::FAILURE;
    };
    println!("quantum-meter: adapter {} member={} holds {} qubits", gpu.key(), gpu.id(), gpu.max_qubits());

    // The first meter in a process probes the GPU counter with this load: a
    // run of the member it will meter.
    let mut meter = ProcessRunMeter::new(
        probe_load(|| {
            let _ = blocking::simulate_from(gpu, &wide, 0);
        }),
        args.declared,
    );
    println!("quantum-meter: counters live={} gpu={}", meter.is_live(), meter.gpu_liveness());

    let mut gpu_member = match blocking::member(gpu, FABRIC_DEFAULT) {
        Ok(m) => m,
        Err(e) => {
            eprintln!("quantum-meter: {} could not be certified: {e:?}", gpu.id());
            return ExitCode::FAILURE;
        }
    };
    let mut cpu_member = detect_local();

    // An unmetered window still runs, so each member is exercised either way.
    let window = |meter: &mut ProcessRunMeter, run: &mut dyn FnMut() -> Result<u64, String>| match span_of_runs(meter, args.at_least, &mut *run) {
        Ok(runs) => Ok(runs),
        // The meter would not start: run once without it.
        Err(_) if !meter.is_live() => run().map(|ops| Runs { runs: 1, ops, joules_micro: 0, class: None, wall: Duration::ZERO }),
        Err(e) => Err(e),
    };

    let gpu_runs = window(&mut meter, &mut || {
        let state = blocking::simulate_from(gpu, &wide, 0).map_err(|e| format!("{e:?}"))?;
        std::hint::black_box(state);
        Ok(ops)
    });
    let gpu_runs = match gpu_runs {
        Ok(r) => r,
        Err(e) => {
            eprintln!("quantum-meter: {} failed its run: {e}", gpu.id());
            return ExitCode::FAILURE;
        }
    };
    let gpu_figure = report(&gpu_member.id, kernel, &gpu_runs, meter.last(), args.declared);

    let cpu_runs = window(&mut meter, &mut || {
        let state = wide.simulate_from(0).map_err(|e| format!("{e:?}"))?;
        std::hint::black_box(state);
        Ok(ops)
    });
    let cpu_runs = match cpu_runs {
        Ok(r) => r,
        Err(e) => {
            eprintln!("quantum-meter: the CPU reference failed its run: {e}");
            return ExitCode::FAILURE;
        }
    };
    let cpu_figure = report(&cpu_member.id, kernel, &cpu_runs, meter.last(), args.declared);

    // Figures go on both members or on neither: a figure on one alone would
    // place it whatever it costs.
    let (Some(gpu_figure), Some(cpu_figure)) = (gpu_figure, cpu_figure) else {
        println!("quantum-meter: no placement by energy: a figure is needed on both members");
        return ExitCode::SUCCESS;
    };
    let set = |member: &mut Member, implementation: &str, figure: &(u64, EnergyClass)| {
        member.kernels.iter_mut().filter(|k| k.kernel == kernel && k.implementation == implementation).for_each(|k| k.energy = Some(figure.clone()));
    };
    set(&mut gpu_member, gpu.implementation(), &gpu_figure);
    set(&mut cpu_member, REFERENCE, &cpu_figure);
    let members = [cpu_member, gpu_member];
    let predict = |_: &Member, k: &wai_quantum::quantum_fabric::KernelEntry| k.energy.as_ref().map(|(nanojoules, _)| nanojoules.saturating_mul(ops).saturating_mul(1_000));
    match place_for(&members, kernel, &Parity::BitExact, predict) {
        Some((member, entry)) => println!(
            "quantum-meter: placed {} ({}) for {} qubits: {} nJ/op against {} nJ/op on the other member",
            member.id,
            entry.implementation,
            wide.n_qubits,
            entry.energy.as_ref().map_or(0, |e| e.0),
            if member.id == members[0].id { gpu_figure.0 } else { cpu_figure.0 }
        ),
        None => println!("quantum-meter: no member is certified for {kernel}"),
    }
    ExitCode::SUCCESS
}