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

//! The fabric's GPU members on one circuit, for the cross-simulator corpus
//! (`wai/quantum-cross-simulator/`, its `gpu` tier).
//!
//! ```text
//! fabric_state PROGRAM.qasm OUT_PREFIX     # one JSON line on standard output
//! ```
//!
//! The program is read as the bit-exact method reads it (its op-log, its
//! gateset) and run on the CPU reference and on every member this build finds:
//! the WebGPU member on each adapter of the platform's GPU interfaces
//! (`quantum_fabric_webgpu`), the Metal member on each device
//! (`quantum_fabric_metal`), the CUDA member on each device
//! (`quantum_fabric_cuda`). Each state that comes back is written to
//! `OUT_PREFIX.<key>` as little-endian `i64` pairs at scale `2^30`, qubit 0 the
//! lowest bit of the index, beside its hash. A device found and not placed, a
//! member's refusal of the circuit, and a family this build leaves out are
//! listed with why, so the harness can say what could not run.

use serde_json::{Value, json};
use std::io::Write as _;
use wai_quantum::quantum::StateVector;

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

fn write_state(path: &str, sv: &StateVector) -> 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 a in &sv.amps {
        w.write_all(&a.re.to_le_bytes()).map_err(|e| e.to_string())?;
        w.write_all(&a.im.to_le_bytes()).map_err(|e| e.to_string())?;
    }
    w.flush().map_err(|e| e.to_string())
}

/// One member's run: its state written beside its hash, or its refusal.
fn ran(family: &str, key: &str, implementation: &str, max_qubits: u32, n: u32, run: impl FnOnce() -> Result<StateVector, String>, prefix: &str) -> Value {
    let mut v = json!({ "family": family, "key": key, "implementation": implementation, "max_qubits": max_qubits });
    if n > max_qubits {
        v["refused"] = json!(format!("{n} qubits; this member holds at most {max_qubits}"));
        return v;
    }
    let t = std::time::Instant::now();
    match run() {
        Ok(sv) => {
            let path = format!("{prefix}.{key}");
            match write_state(&path, &sv) {
                Ok(()) => {
                    v["seconds"] = json!(t.elapsed().as_secs_f64());
                    v["state_out"] = json!(path);
                    v["state_hash"] = json!(hex(&sv.statevector_hash()));
                }
                Err(e) => v["refused"] = json!(e),
            }
        }
        Err(e) => v["refused"] = json!(e),
    }
    v
}

fn main() {
    let args: Vec<String> = std::env::args().collect();
    if args.len() != 3 {
        eprintln!("usage: fabric_state PROGRAM.qasm OUT_PREFIX");
        std::process::exit(2);
    }
    let (src_path, prefix) = (&args[1], &args[2]);
    let answer = (|| -> Result<Value, String> {
        let src = std::fs::read_to_string(src_path).map_err(|e| format!("{src_path}: {e}"))?;
        let program = wai_quantum::qasm::read(&src).map_err(|d| d.to_string())?.program;
        let (c, _) = program.dyadic_with_report()?;
        let n = program.n_qubits;
        let mut members: Vec<Value> = Vec::new();
        let mut absent: Vec<Value> = Vec::new();
        members.push(ran("cpu", "cpu", "reference", 64, n, || c.simulate().map_err(|e| e.to_string()), prefix));
        #[cfg(feature = "quantum_fabric_webgpu")]
        {
            use wai_quantum::quantum_fabric_webgpu::blocking;
            match blocking::discover() {
                Ok(d) => {
                    for m in &d.found {
                        members.push(ran("webgpu", m.key(), m.implementation(), u32::from(m.max_qubits()), n, || blocking::simulate_from(m, &c, 0).map_err(|e| e.to_string()), prefix));
                    }
                    for r in &d.refused {
                        absent.push(json!({ "family": "webgpu", "key": r.key, "device": r.adapter, "why": r.error.to_string() }));
                    }
                    if d.found.is_empty() && d.refused.is_empty() {
                        absent.push(json!({ "family": "webgpu", "why": "the platform's GPU interfaces list no adapter" }));
                    }
                }
                Err(e) => absent.push(json!({ "family": "webgpu", "why": e.to_string() })),
            }
        }
        #[cfg(not(feature = "quantum_fabric_webgpu"))]
        absent.push(json!({ "family": "webgpu", "why": "this build leaves the WebGPU member out (feature quantum_fabric_webgpu)" }));
        #[cfg(feature = "quantum_fabric_metal")]
        {
            let d = wai_quantum::quantum_fabric_metal::discover();
            for m in &d.found {
                members.push(ran("metal", m.key(), m.implementation(), u32::from(m.max_qubits()), n, || m.simulate(&c).map_err(|e| e.to_string()), prefix));
            }
            for r in &d.refused {
                absent.push(json!({ "family": "metal", "key": r.key, "device": r.device, "why": r.error.to_string() }));
            }
            if let Some(why) = &d.unavailable {
                absent.push(json!({ "family": "metal", "why": why }));
            }
        }
        #[cfg(not(feature = "quantum_fabric_metal"))]
        absent.push(json!({ "family": "metal", "why": "this build leaves the Metal member out (feature quantum_fabric_metal)" }));
        #[cfg(feature = "quantum_fabric_cuda")]
        {
            let d = wai_quantum::quantum_fabric_cuda::discover();
            for m in &d.found {
                members.push(ran("cuda", m.key(), m.implementation(), u32::from(m.max_qubits()), n, || m.simulate(&c).map_err(|e| e.to_string()), prefix));
            }
            for r in &d.refused {
                absent.push(json!({ "family": "cuda", "key": r.key, "device": r.device, "why": r.error.to_string() }));
            }
            if let Some(why) = &d.unavailable {
                absent.push(json!({ "family": "cuda", "why": why }));
            }
        }
        #[cfg(not(feature = "quantum_fabric_cuda"))]
        absent.push(json!({ "family": "cuda", "why": "this build leaves the CUDA member out (feature quantum_fabric_cuda)" }));
        Ok(json!({ "n": n, "n_ops": c.ops.len(), "members": members, "absent": absent }))
    })();
    match answer {
        Ok(v) => println!("{}", json!({ "ok": true, "answer": v })),
        Err(e) => println!("{}", json!({ "ok": false, "refused": e })),
    }
}