open-standards/wai/wai-quantum/tests/cuda_sv.rs

//! The dense f64 state vector on the CUDA devices of the machine that runs
//! this: bit for bit the CPU's, or the certificate says not.
//!
//! The machine declares its devices in `WAI_QUANTUM_CUDA_EXPECT` (comma-
//! separated keys, or `none`), as for `cuda_member`; `WAI_QUANTUM_CUDA_REPORT`
//! names a file the established lines are appended to.
#![cfg(not(target_arch = "wasm32"))]

use std::io::Write;
use wai_quantum::gates::Gate;
use wai_quantum::quantum_fabric::sv64::{self, MUTANTS, SV_DEFAULT, SV_DEFAULT_PIN, SV_WIDE, SV_WIDE_PIN, SvCorpus};
use wai_quantum::quantum_fabric::{Certification, Parity};
use wai_quantum::quantum_fabric_cuda::{CudaMember, discover};
use wai_quantum::quantum_sv::StateVector;

fn report(line: &str) {
    println!("{line}");
    if let Some(path) = std::env::var_os("WAI_QUANTUM_CUDA_REPORT") {
        let mut f = std::fs::OpenOptions::new().create(true).append(true).open(&path).expect("the report file opens");
        writeln!(f, "{line}").expect("the report line is written");
    }
}

fn members() -> Vec<CudaMember> {
    let d = discover();
    let mut found: Vec<String> = d.found.iter().map(|m| m.key().to_string()).collect();
    found.sort();
    let found_list = if found.is_empty() { "none".to_string() } else { found.join(",") };
    let expect = std::env::var("WAI_QUANTUM_CUDA_EXPECT").unwrap_or_else(|_| panic!("WAI_QUANTUM_CUDA_EXPECT is not set; found {found_list}"));
    assert_eq!(found_list, expect, "the devices found are not the declared ones");
    d.found
}

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

/// The CPU certifying itself on `corpus`: the digest a bit-exact device
/// reproduces.
fn cpu_digest(corpus: SvCorpus) -> String {
    let mut cert = Certification::new();
    for case in corpus.cases() {
        let r = case.reference().unwrap();
        cert.add(&case.bytes(), &r, &r);
    }
    hex(&cert.finish(sv64::KERNEL_SV_RUN, "cpu", "cpu", &corpus.generator()).corpus)
}

/// The CPU certifying itself on [`SV_WIDE`]'s wide cases.
fn cpu_wide_digest() -> String {
    let mut cert = Certification::new();
    for case in SV_WIDE.wide_cases() {
        let r = case.reference().unwrap();
        cert.add(&case.bytes(), &r, &r);
    }
    hex(&cert.finish(sv64::KERNEL_SV_RUN_WIDE, "cpu", "cpu", &SV_WIDE.wide_generator()).corpus)
}

/// Whether a wide certificate reproduces the CPU's pinned digest.
fn wide_exact(w: &wai_quantum::quantum_fabric::Certificate) -> bool {
    w.parity == Parity::BitExact && hex(&w.corpus) == SV_WIDE_PIN
}

#[test]
fn s0_the_state_vector_is_the_cpus_bit_for_bit() {
    let cpu = cpu_digest(SV_DEFAULT);
    assert_eq!(cpu, SV_DEFAULT_PIN, "this machine's CPU does not reproduce the pin");
    assert_eq!(cpu_wide_digest(), SV_WIDE_PIN, "this machine's CPU does not reproduce the wide corpus's pin");
    for m in members() {
        let sv = m.sv().unwrap_or_else(|e| panic!("{}: {e}", m.key()));
        let c = sv.certify(SV_DEFAULT).unwrap_or_else(|e| panic!("{}: {e}", m.key()));
        assert_eq!(c.parity, Parity::BitExact, "{}", m.key());
        assert_eq!(c.cases, SV_DEFAULT.cases as u64);
        assert_eq!(hex(&c.corpus), cpu, "{}: the digest is not the CPU's", m.key());
        report(&format!("cuda sv certified: key={} member={} kernel={} implementation={} parity=bit-exact cases={} digest={}", m.key(), m.id(), c.kernel, c.implementation, c.cases, hex(&c.corpus)));
        // The wider kernels, on their own corpus: bit-exact, or said not to
        // be (such a member is placed only for gates on up to three qubits).
        let w = sv.certify_wide(SV_WIDE).unwrap_or_else(|e| panic!("{}: {e}", m.key()));
        assert_eq!(w.cases, SV_WIDE.cases as u64);
        let exact = wide_exact(&w);
        report(&format!("cuda sv wide certified: key={} member={} kernel={} implementation={} parity={} cases={} digest={}", m.key(), m.id(), w.kernel, w.implementation, if exact { "bit-exact" } else { "not-bit-exact" }, w.cases, hex(&w.corpus)));
    }
}

#[test]
fn s1_wide_states_match_the_cpu() {
    for m in members() {
        let sv = m.sv().unwrap();
        for (n, width) in [(16u32, 0usize), (20, 3)] {
            let mut gates = Vec::new();
            let mut a = 0.1f64;
            for layer in 0..6u32 {
                for q in 0..n {
                    a += 0.731;
                    gates.push(match (q + layer) % 5 {
                        0 => Gate::rx(q, a),
                        1 => Gate::ry(q, a),
                        2 => Gate::rz(q, a),
                        3 => Gate::sqrt_w(q),
                        _ => Gate::t(q),
                    });
                }
                let mut q = layer % 2;
                while q + 1 < n {
                    gates.push(if q % 3 == 0 { Gate::fsim(q, q + 1, a, 0.5 * a) } else { Gate::cz(q, q + 1) });
                    q += 2;
                }
                gates.push(Gate::iswap(0, n - 1));
            }
            let mut cpu = StateVector::zero(n).unwrap();
            cpu.run(&gates, width).unwrap();
            let (re, im) = sv.run(n, &gates, width).unwrap();
            let same = |a: &[f64], b: &[f64]| a.iter().zip(b).all(|(x, y)| x.to_bits() == y.to_bits());
            assert!(same(&re, &cpu.re) && same(&im, &cpu.im), "{}: n={n} width={width}", m.key());
        }
        report(&format!("cuda sv wide: key={} 16 and 20 qubits match the CPU bit for bit", m.key()));
    }
}

#[test]
fn s2_mutants_are_caught() {
    let small = SvCorpus { cases: 18, ..SV_DEFAULT };
    let cpu = cpu_digest(small);
    for m in members() {
        let mut killed = 0;
        for mutant in &MUTANTS {
            let sv = m.sv_with_source(&sv64::mutate(mutant)).unwrap_or_else(|e| panic!("mutant {}: {e}", mutant.name));
            let c = sv.certify(small).unwrap();
            assert!(c.parity != Parity::BitExact || hex(&c.corpus) != cpu, "{}: mutant {} survived", m.key(), mutant.name);
            killed += 1;
        }
        report(&format!("cuda sv mutants: key={} killed {killed} of {}", m.key(), MUTANTS.len()));
    }
}

/// Ground truth beyond the pinned corpus, on the resident path: random
/// circuits at every fusion width the kernels take (0 to 6), run from
/// `|0…0⟩` and, as a partition runs, from a state given (the CPU's state after
/// the first half of the gates), each the CPU's state bit for bit.
#[test]
fn s3_random_circuits_match_the_cpu_from_zero_and_from_a_state() {
    let same = |a: &[f64], b: &[f64]| a.len() == b.len() && a.iter().zip(b).all(|(x, y)| x.to_bits() == y.to_bits());
    for m in members() {
        let sv = m.sv().unwrap();
        let wide = wide_exact(&sv.certify_wide(SV_WIDE).unwrap());
        let mut runs = 0;
        for case in sv64::random_cases(2026, 70, 16, 60) {
            let launches = sv64::plan(case.n, &case.gates, case.width).unwrap();
            // A member whose wider kernels did not certify is never placed
            // for them: its narrow kernels alone are held to the CPU.
            if !wide && launches.iter().any(|l| sv64::launch_bits(l).len() as u32 > sv64::NARROW) {
                continue;
            }
            let mut cpu = StateVector::zero(case.n).unwrap();
            cpu.run(&case.gates, case.width).unwrap();
            let (re, im) = sv.run_from_zero(case.n, &launches).unwrap_or_else(|e| panic!("{}: {e}", m.key()));
            assert!(same(&re, &cpu.re) && same(&im, &cpu.im), "{}: from zero, n={} width={}", m.key(), case.n, case.width);
            let half = case.gates.len() / 2;
            let mut cpu = StateVector::zero(case.n).unwrap();
            cpu.run(&case.gates[..half], case.width).unwrap();
            let rest = sv64::plan(case.n, &case.gates[half..], case.width).unwrap();
            let (re, im) = sv.run_launches(&cpu.re, &cpu.im, &rest).unwrap();
            cpu.run(&case.gates[half..], case.width).unwrap();
            assert!(same(&re, &cpu.re) && same(&im, &cpu.im), "{}: from a state, n={} width={}", m.key(), case.n, case.width);
            runs += 2;
        }
        report(&format!("cuda sv random: key={} {runs} resident runs on {} match the CPU bit for bit", m.key(), if wide { "every kernel" } else { "the narrow kernels" }));
    }
}