open-standards/wai/wai-quantum-cli/examples/cross_simulator.rs

//! The engine's side of the cross-simulator corpus
//! (`wai/quantum-cross-simulator/`): one job in, one answer out,
//! through the same calls the command line and the job service make.
//!
//! ```text
//! cross_simulator JOB.json      # the answer as JSON on standard output
//! ```
//!
//! Each job runs in a process of its own, so its peak memory is its own: on
//! Linux the answer carries it (`peak_rss_bytes`, the system's `VmHWM`), and on
//! macOS the harness reads it when the process ends. A job names its `kind`:
//!
//! - `state`: the program's whole state, written to `state_out` with qubit 0 as
//!   the lowest bit of the index, each gate under many controls held as a job
//!   holds it: whole on `statevector`, in standard gates on `exact` unless
//!   `controls` is `whole` (`standard` writes standard gates for both). `method`
//!   `exact` writes `i64` pairs at scale `2^30` (with the state's hash, its
//!   op-log, which an independent reader reruns, and the global phase the
//!   job's result states); `statevector` writes `f64` pairs, at fusion width
//!   `fusion`.
//! - `job`: `request` through [`wai_quantum::engine::job::plan_and_run`], the
//!   engine's whole path: its plan, its run and its answers, unsealed.
//! - `stabilizers`: the program on the tableau, and `⟨g⟩` of each Pauli string
//!   of `paulis`, each as `XYZ_` letters with qubit 0 first. A state is fixed by
//!   `n` independent commuting strings that all read `+1`.
//! - `qec_model`, `qec_sample`: a circuit in the error-correction text, its
//!   detector error model in the flat form, or how often each detector and
//!   observable fired over `shots`.

use serde_json::{Value, json};
use std::io::Write as _;
use std::time::Instant;
use wai_quantum::engine::job;
use wai_quantum::qasm;
use wai_quantum::quantum_frame as frame;
use wai_quantum::quantum_sv::StateVector;
use wai_quantum::quantum_tableau::Tableau;

const GIB: u64 = 1 << 30;

fn write_pairs<T: Copy>(path: &str, pairs: impl Iterator<Item = (T, T)>, bytes: impl Fn(T) -> [u8; 8]) -> Result<(), String> {
    let f = std::fs::File::create(path).map_err(|e| format!("{path}: {e}"))?;
    let mut w = std::io::BufWriter::with_capacity(1 << 20, f);
    for (re, im) in pairs {
        w.write_all(&bytes(re)).map_err(|e| e.to_string())?;
        w.write_all(&bytes(im)).map_err(|e| e.to_string())?;
    }
    w.flush().map_err(|e| e.to_string())
}

fn hex(b: &[u8]) -> String {
    b.iter().map(|x| format!("{x:02x}")).collect()
}

fn text<'a>(j: &'a Value, k: &str) -> Result<&'a str, String> {
    j.get(k).and_then(Value::as_str).ok_or_else(|| format!("the job has no \"{k}\""))
}

fn state(j: &Value) -> Result<Value, String> {
    let method = text(j, "method")?;
    // As a job runs it: a gate under many controls held whole on the state
    // vector, and on the bit-exact method only when asked ("controls": "whole");
    // "controls": "standard" writes standard gates for both.
    let whole = match j.get("controls").and_then(Value::as_str) {
        Some("standard") => false,
        Some("whole") => true,
        _ => method == "statevector",
    };
    let src = text(j, "qasm")?;
    let program = if whole { qasm::read_whole(src) } else { qasm::read(src) }.map_err(|d| d.to_string())?.program;
    let n = program.n_qubits;
    let out = j.get("state_out").and_then(Value::as_str);
    match method {
        "exact" => {
            let (c, report) = program.dyadic_with_report()?;
            let t = Instant::now();
            let sv = c.simulate().map_err(|e| e.to_string())?;
            let seconds = t.elapsed().as_secs_f64();
            if let Some(path) = out {
                write_pairs(path, sv.amps.iter().map(|a| (a.re, a.im)), i64::to_le_bytes)?;
            }
            let ops: Vec<Value> = c.ops.iter().map(|g| json!([g.base as u8, g.controls, g.target, g.param])).collect();
            Ok(json!({ "n": n, "seconds": seconds, "state_hash": hex(&sv.statevector_hash()), "ops": if j.get("ops").and_then(Value::as_bool).unwrap_or(true) { Value::Array(ops) } else { Value::Null }, "n_ops": c.ops.len(), "global_phase": report.global_phase_radians() }))
        }
        "statevector" => {
            let gates = program.sv_ops();
            let width = j.get("fusion").and_then(Value::as_u64).unwrap_or(3) as usize;
            let mut sv = StateVector::zero(n).map_err(|e| format!("{e:?}"))?;
            let t = Instant::now();
            sv.run_ops(&gates, width).map_err(|e| format!("{e:?}"))?;
            let seconds = t.elapsed().as_secs_f64();
            if let Some(path) = out {
                // Qubit 0 is the state vector's highest bit; the file has it lowest.
                let reverse = |x: usize| if n == 0 { 0 } else { x.reverse_bits() >> (usize::BITS - n) };
                write_pairs(path, (0..sv.re.len()).map(|x| (sv.re[reverse(x)], sv.im[reverse(x)])), f64::to_le_bytes)?;
            }
            Ok(json!({ "n": n, "seconds": seconds, "n_ops": gates.len(), "fusion": width, "norm2": sv.norm2() }))
        }
        m => Err(format!("state: no method {m}")),
    }
}

fn stabilizers(j: &Value) -> Result<Value, String> {
    let program = qasm::read(text(j, "qasm")?).map_err(|d| d.to_string())?.program;
    let ops = program.clifford()?;
    let t = Instant::now();
    let tableau = Tableau::run(program.n_qubits, &ops).map_err(|e| e.to_string())?;
    let run = t.elapsed().as_secs_f64();
    let dimension = tableau.support().dimension();
    let mut values = Vec::new();
    let t = Instant::now();
    for s in j.get("paulis").and_then(Value::as_array).into_iter().flatten() {
        let s = s.as_str().ok_or("a Pauli string is text")?;
        let string: Vec<(u32, char)> = s.chars().enumerate().filter(|&(_, c)| c != '_' && c != 'I').map(|(q, c)| (q as u32, c)).collect();
        values.push(i64::from(tableau.expectation(&string).map_err(|e| e.to_string())?));
    }
    Ok(json!({ "n": program.n_qubits, "seconds": run, "expectation_seconds": t.elapsed().as_secs_f64(), "n_ops": ops.len(), "values": values, "outcome_dimension": dimension }))
}

fn qec_model(j: &Value) -> Result<Value, String> {
    let c = frame::parse(text(j, "text")?).map_err(|e| e.to_string())?;
    let t = Instant::now();
    let model = frame::error_model(&c).map_err(|e| e.to_string())?;
    let seconds = t.elapsed().as_secs_f64();
    Ok(json!({ "qubits": c.qubits, "detectors": model.detectors, "observables": model.observables, "seconds": seconds, "model": model.to_text() }))
}

fn qec_sample(j: &Value) -> Result<Value, String> {
    let c = frame::parse(text(j, "text")?).map_err(|e| e.to_string())?;
    let shots = j.get("shots").and_then(Value::as_u64).ok_or("the job has no \"shots\"")? as usize;
    let seed = j.get("seed").and_then(Value::as_u64).unwrap_or(1);
    let t = Instant::now();
    let d = frame::sample(&c, shots, seed);
    let seconds = t.elapsed().as_secs_f64();
    let detectors: Vec<u64> = (0..d.detectors).map(|k| d.count(k)).collect();
    let mut observables = vec![0u64; d.observables];
    for s in 0..shots {
        let f = d.flips(s);
        for (k, o) in observables.iter_mut().enumerate() {
            *o += f >> k & 1;
        }
    }
    Ok(json!({ "shots": shots, "seconds": seconds, "detectors": detectors, "observables": observables }))
}

fn run_job(j: &Value) -> Result<Value, String> {
    match text(j, "kind")? {
        "state" => state(j),
        "job" => {
            let memory = j.get("memory_gb").and_then(Value::as_u64).unwrap_or(16) * GIB;
            let request = j.get("request").ok_or("the job has no \"request\"")?;
            let t = Instant::now();
            let (plan, result) = job::plan_and_run(request, memory)?;
            Ok(json!({ "seconds": t.elapsed().as_secs_f64(), "method": plan.method, "result": result }))
        }
        "stabilizers" => stabilizers(j),
        "qec_model" => qec_model(j),
        "qec_sample" => qec_sample(j),
        k => Err(format!("no job kind {k}")),
    }
}

/// This process's peak resident set, where the system reports it for this
/// process alone: `VmHWM` on Linux. (`ru_maxrss` there carries what the parent
/// held when it forked; on macOS the harness reads `ru_maxrss` itself.)
fn peak_rss_bytes() -> Option<u64> {
    let status = std::fs::read_to_string("/proc/self/status").ok()?;
    let line = status.lines().find(|l| l.starts_with("VmHWM:"))?;
    line.split_whitespace().nth(1)?.parse::<u64>().ok().map(|kib| kib * 1024)
}

fn main() {
    let path = std::env::args().nth(1).unwrap_or_else(|| {
        eprintln!("usage: cross_simulator JOB.json");
        std::process::exit(2)
    });
    let answer = std::fs::read_to_string(&path)
        .map_err(|e| format!("{path}: {e}"))
        .and_then(|t| serde_json::from_str::<Value>(&t).map_err(|e| format!("{path}: {e}")))
        .and_then(|j| run_job(&j));
    // A refusal is an answer: the harness records it beside the job.
    let peak = peak_rss_bytes();
    match answer {
        Ok(v) => println!("{}", json!({ "ok": true, "answer": v, "peak_rss_bytes": peak })),
        Err(e) => println!("{}", json!({ "ok": false, "refused": e, "peak_rss_bytes": peak })),
    }
}