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

//! The Metal member's tensor-operation kernel on the devices of the machine
//! that runs this: every product the gate needs formed by the shading
//! language's tensor operation, from offset-binary 8-bit limbs in bf16
//! accumulated in f32, held to the same tests as the other kernels
//! (`metal_member`).
//!
//! The machine declares its devices in `WAI_QUANTUM_METAL_EXPECT`, as for
//! `metal_member`: their keys, comma-separated and sorted, or `none`. Every
//! declared device with the tensor operation (GPU family 1007 or later, under
//! operating system 26 or later) must run the kernel and pass its probe; one
//! without must be refused as unsupported. Each test prints what it established and, when
//! `WAI_QUANTUM_METAL_REPORT` names a file, appends the same lines to it.
//!
//! ```text
//! WAI_QUANTUM_METAL_EXPECT=gpu-msl-1009 \
//!   cargo test --features quantum_fabric_metal --test metal_tensor_ops -- --test-threads=1
//! ```
#![cfg(not(target_arch = "wasm32"))]

use std::io::Write;
use std::sync::OnceLock;
use wai_quantum::ed25519_dalek::SigningKey;
use wai_quantum::quantum::{BaseGate, Circuit, Gate, Gateset, MAX_QUBITS, QuantumError};
use wai_quantum::quantum_fabric::metal::{self, TENSOR_OPS_MUTANTS};
use wai_quantum::quantum_fabric::lanes::{self, GateWords, Lane};
use wai_quantum::quantum_fabric::plan::{Program, Step};
use wai_quantum::quantum_fabric::wgsl::{Detector, Mutant};
use wai_quantum::quantum_fabric::{
    Corpus, CpuReference, FABRIC_DEFAULT, Fabric, FabricClass, FabricReceipt, MemberError, Parity, Presence, REFERENCE, apply_kernel, detect_local, place, ready,
};
use wai_quantum::quantum_fabric_metal::{Discovery, MetalError, MetalMember, discover};

/// The CPU's corpus digests: random-v1 (seed 7, 200 cases) and fabric-v1
/// (seed 7, 200 + 24 cases), per gateset. A bit-exact member reproduces them.
const PINS: [(Gateset, &str, &str); 2] = [
    (Gateset::V1, "d18454c639acfd9c7b66450f4ded1aac4ab46cc10a7bedbcf4958a96e204fb36", "41c726a0ac80f7defef92e15ecf893bdb7a3076b5e7f670fa0b7142e5c033df7"),
    (Gateset::V2, "3dbde276c1095bf827faec71d5980e050f8a55a5ba7d97cd3bbdaec3f43a13b2", "9a84b55af68f1a8b07390b234cf562dd8d85cef3185a697ab85690c8790a7809"),
];

const RANDOM_V1: Corpus = Corpus::Random { seed: 7, cases: 200 };

static DISCOVERY: OnceLock<Discovery> = OnceLock::new();
static TENSOR_OPS: OnceLock<Vec<MetalMember>> = OnceLock::new();

fn discovery() -> &'static Discovery {
    DISCOVERY.get_or_init(discover)
}

/// Prints `line` and appends it to the file `WAI_QUANTUM_METAL_REPORT` names.
fn report(line: &str) {
    println!("{line}");
    if let Some(path) = std::env::var_os("WAI_QUANTUM_METAL_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");
    }
}

/// The declared device set, as written, if it is exactly what was found;
/// otherwise why not, with what was found and refused and how to declare it.
fn declared(d: &Discovery) -> Result<String, String> {
    let mut found: Vec<&str> = d.found.iter().map(|m| m.key()).collect();
    found.sort_unstable();
    let found_list = if found.is_empty() { "none".to_string() } else { found.join(",") };
    let refused: Vec<String> = d.refused.iter().map(|r| format!("{} ({:?}): {}", r.key, r.device, r.error)).collect();
    let refused_list = if refused.is_empty() { "none".to_string() } else { refused.join("; ") };
    let unavailable = d.unavailable.as_deref().unwrap_or("no");
    let how = format!("found: {found_list}; refused: {refused_list}; unavailable: {unavailable}. Declare the devices this machine must have: WAI_QUANTUM_METAL_EXPECT={found_list}");
    let expect = match std::env::var("WAI_QUANTUM_METAL_EXPECT") {
        Ok(v) if !v.is_empty() => v,
        _ => return Err(format!("WAI_QUANTUM_METAL_EXPECT is not set. {how}")),
    };
    let keys: Vec<&str> = if expect == "none" { Vec::new() } else { expect.split(',').collect() };
    let well_formed = keys.windows(2).all(|w| w[0] <= w[1]) && keys.iter().all(|k| !k.is_empty() && *k != "none" && k.trim() == *k);
    if !well_formed {
        return Err(format!("WAI_QUANTUM_METAL_EXPECT={expect:?} is not a sorted, comma-separated list of keys, nor none. {how}"));
    }
    for k in &keys {
        if !found.contains(k)
            && let Some(r) = d.refused.iter().find(|r| r.key == *k)
        {
            return Err(format!("the declared device {k} was refused: {}. {how}", r.error));
        }
    }
    if found != keys {
        return Err(format!("the devices found are not the declared ones ({expect}). {how}"));
    }
    Ok(expect)
}

/// The tensor-operation members: one per declared device with the operation, once
/// the declared devices are exactly the ones found.
fn members() -> &'static [MetalMember] {
    let d = discovery();
    if let Err(why) = declared(d) {
        panic!("{why}");
    }
    TENSOR_OPS.get_or_init(|| {
        d.found.iter().filter(|m| m.identity().has_tensor_ops()).map(|m| m.tensor_ops().unwrap_or_else(|e| panic!("{}: the tensor-operation kernel: {e}", m.key()))).collect()
    })
}

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

fn next(s: &mut u64) -> u64 {
    *s = s.wrapping_add(0x9E37_79B9_7F4A_7C15);
    let mut z = *s;
    z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
    z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
    z ^ (z >> 31)
}

/// A value in `[-bound, bound]`.
fn within(s: &mut u64, bound: i64) -> i32 {
    ((next(s) % (2 * bound as u64 + 1)) as i64 - bound) as i32
}

/// A random circuit of `ops` gates on `n` qubits: every base gate, up to three
/// distinct controls, `P(1..=32)`.
fn random_circuit(n: u8, ops: usize, seed: u64, gs: Gateset) -> Circuit {
    let bases = [BaseGate::I, BaseGate::X, BaseGate::Y, BaseGate::Z, BaseGate::H, BaseGate::S, BaseGate::Sdg, BaseGate::T, BaseGate::Tdg, BaseGate::P];
    let mut s = seed;
    let mut c = Circuit::with_gateset(n, gs);
    for _ in 0..ops {
        let base = bases[(next(&mut s) % bases.len() as u64) as usize];
        let target = (next(&mut s) % u64::from(n)) as u8;
        let mut controls = Vec::new();
        for _ in 0..next(&mut s) % 4 {
            let q = (next(&mut s) % u64::from(n)) as u8;
            if q != target && !controls.contains(&q) {
                controls.push(q);
            }
        }
        let param = if base == BaseGate::P { 1 + (next(&mut s) % 32) as u16 } else { 0 };
        c.ops.push(Gate { base, controls, target, param });
    }
    c
}

/// Five circuits the native simulator refuses, in each gateset.
fn refused_circuits() -> Vec<Circuit> {
    let gate = |base, controls: Vec<u8>, target, param| Gate { base, controls, target, param };
    let bad = [
        (1, vec![gate(BaseGate::H, vec![], 0, 0), gate(BaseGate::P, vec![], 0, 0)]),
        (1, vec![gate(BaseGate::P, vec![], 0, 33)]),
        (1, vec![gate(BaseGate::X, vec![], 3, 0)]),
        (2, vec![gate(BaseGate::X, vec![0], 0, 0)]),
        (2, vec![gate(BaseGate::X, vec![5], 1, 0)]),
    ];
    let mut out = Vec::new();
    for gs in Gateset::ALL {
        for (n, ops) in bad.clone() {
            let mut c = Circuit::with_gateset(n, gs);
            c.ops = ops;
            out.push(c);
        }
    }
    out
}

/// `Z` on qubit 0: on `[(0, 0), (i32::MIN, 0)]` its exact result is `2^31`,
/// which leaves 32 bits.
fn z_words() -> GateWords {
    lanes::gate_words(&Gate { base: BaseGate::Z, controls: vec![], target: 0, param: 0 }, Gateset::V2)
}

const Z_STATE: [Lane; 2] = [[0, 0], [i32::MIN, 0]];

#[test]
fn c0_devices_are_declared() {
    let d = discovery();
    match declared(d) {
        Ok(expect) => report(&format!("metal tensor-ops devices: as declared ({expect})")),
        Err(why) => panic!("{why}"),
    }
    for m in &d.found {
        if !m.identity().has_tensor_ops() {
            assert!(matches!(m.tensor_ops(), Err(MetalError::Unsupported(_))), "{}", m.key());
            report(&format!("metal tensor-ops unsupported: key={} family {} os {}", m.key(), m.identity().family, m.identity().os));
        }
    }
    for t in members() {
        assert_eq!(t.class(), FabricClass::Gpu);
        assert!(t.id().starts_with("local/gpu-msl-"), "{}", t.id());
        assert_eq!(t.implementation(), metal::tensor_ops_implementation_id());
        assert_ne!(t.implementation(), metal::implementation_id());
        assert_eq!(t.identity().member_id(), t.id());
        let lanes64 = d.found.iter().find(|m| m.id() == t.id()).expect("the device's member");
        assert_eq!((lanes64.implementation(), lanes64.max_qubits()), (metal::implementation_id().as_str(), t.max_qubits()));
        assert!(t.max_qubits() >= 20, "{}: holds only {} qubits", t.key(), t.max_qubits());
        assert!(t.detail().starts_with(lanes64.detail()) && t.detail().contains("tensor operation"), "{}", t.detail());
        report(&format!("metal tensor-ops member: key={} id={} implementation={} {}", t.key(), t.id(), t.implementation(), t.detail()));
    }
}

#[test]
fn c1_probe_matches_lanes() {
    for m in members() {
        let pairs = lanes::probe_pairs(1 << 16, 0x0cda);
        let out = m.probe(&pairs).unwrap_or_else(|e| panic!("{}: {e}", m.key()));
        assert_eq!(lanes::first_mismatch(&pairs, &out), None, "{}", m.key());
        report(&format!("metal tensor-ops probe: key={} {} pairs agree", m.key(), pairs.len()));
    }
}

#[test]
fn c2_kernel_matches_run_words() {
    for m in members() {
        let mut s = 0x6b65_726e_656c_u64;
        let (mut flagged, mut clean) = (0, 0);
        for case in 0..2000 {
            let n = 1 + (next(&mut s) % 12) as u32;
            // Three regimes: the whole 32-bit range (nearly always out of
            // range), small states and entries (never), and states and
            // entries up to 2^30 (either).
            let mode = case % 3;
            let value = |s: &mut u64| -> i32 {
                match mode {
                    0 => next(s) as i32,
                    1 => within(s, 1 << 22),
                    _ => within(s, 1 << 30),
                }
            };
            let entry = |s: &mut u64| -> u32 {
                match mode {
                    0 => next(s) as u32,
                    1 => within(s, 1 << 29) as u32,
                    _ => within(s, 1 << 30) as u32,
                }
            };
            let state: Vec<Lane> = (0..1usize << n).map(|_| [value(&mut s), value(&mut s)]).collect();
            let gates: Vec<GateWords> = (0..1 + next(&mut s) % 8)
                .map(|_| {
                    let t = (next(&mut s) % u64::from(n)) as u32;
                    let ctrl = (0..n).filter(|&b| b != t && next(&mut s).is_multiple_of(4)).fold(0u32, |acc, b| acc | 1 << b);
                    let mut w = [0u32; 8];
                    w.iter_mut().for_each(|x| *x = entry(&mut s));
                    GateWords { m: w, tbit: 1 << t, ctrl }
                })
                .collect();
            let want = lanes::run_words(&state, &gates).expect("valid shapes");
            let got = m.run_lanes(&state, &gates).unwrap_or_else(|e| panic!("{}: case {case}: {e}", m.key()));
            assert_eq!(got.out_of_range, want.out_of_range, "{}: case {case}, n = {n}", m.key());
            assert!(got.amps == want.amps, "{}: case {case}, n = {n}: the lanes differ", m.key());
            if want.out_of_range.is_some() {
                flagged += 1;
            } else {
                clean += 1;
            }
        }
        assert!(flagged >= 100 && clean >= 100, "{}: {flagged} flagged, {clean} clean", m.key());
        report(&format!("metal tensor-ops kernel: key={} ok", m.key()));
    }
}

#[test]
fn c3_certificates_reproduce_the_pins() {
    let ms = members();
    for (gs, random_pin, fabric_pin) in PINS {
        let cpu = lanes::certify_corpus(gs, "local/test", RANDOM_V1);
        assert_eq!(hex(&cpu.corpus), random_pin, "{gs:?}");
        let other = Corpus::Random { seed: 8, cases: 300 };
        let cpu_other = lanes::certify_corpus(gs, "local/test", other);
        for m in ms {
            for (corpus, cases, pin) in [(RANDOM_V1, 200, random_pin), (FABRIC_DEFAULT, 224, fabric_pin)] {
                let c = m.certify(gs, corpus).unwrap_or_else(|e| panic!("{} {gs:?} {}: {e}", m.key(), corpus.name()));
                assert_eq!(c.parity, Parity::BitExact, "{} {gs:?} {}", m.key(), corpus.name());
                assert_eq!(c.cases, cases);
                assert_eq!(hex(&c.corpus), pin, "{} {gs:?} {}", m.key(), corpus.name());
                assert_eq!((c.kernel.as_str(), c.implementation.as_str(), c.member.as_str()), (apply_kernel(gs), m.implementation(), m.id()));
                assert_eq!(c.generator, corpus.generator(gs));
                report(&format!(
                    "metal tensor-ops certified: key={} member={} kernel={} corpus={} implementation={} parity=bit-exact cases={} digest={}",
                    m.key(),
                    m.id(),
                    c.kernel,
                    corpus.name(),
                    c.implementation,
                    c.cases,
                    hex(&c.corpus)
                ));
            }
            // A corpus with no pin: the member and the CPU agree on it.
            let c = m.certify(gs, other).unwrap_or_else(|e| panic!("{} {gs:?} seed 8: {e}", m.key()));
            assert_eq!((c.parity, c.cases), (Parity::BitExact, 300), "{} {gs:?}", m.key());
            assert_eq!(c.corpus, cpu_other.corpus, "{} {gs:?}", m.key());
        }
    }
}

#[test]
fn c4_wide_circuits_match_native() {
    for m in members() {
        let widest = m.max_qubits().min(22);
        let mut s = 0x0077_6964_6500_u64;
        for (k, n) in [12u8, 16, 18, 20, widest].into_iter().enumerate() {
            for gs in Gateset::ALL {
                let c = random_circuit(n, 40, 0x1000 + k as u64, gs);
                let start = (next(&mut s) as usize) % (1usize << n);
                let want = c.simulate_from(start).expect("the native run");
                let got = m.simulate_from(&c, start).unwrap_or_else(|e| panic!("{}: n = {n} {gs:?}: {e}", m.key()));
                assert!(got == want, "{}: n = {n} {gs:?}: the states differ", m.key());
            }
        }
        // The transforms the conformance corpus carries.
        for n in [3u8, 5, 8, 12] {
            let qft = Circuit::qft(n);
            assert!(m.simulate(&qft).unwrap() == qft.simulate().unwrap(), "{}: QFT-{n}", m.key());
        }
        report(&format!("metal tensor-ops wide: key={} up to {widest} qubits match native", m.key()));
    }
}

#[test]
fn c5_refusals_are_typed() {
    for m in members() {
        for c in refused_circuits() {
            let native = c.simulate_from(0).unwrap_err();
            assert_eq!(m.simulate_from(&c, 0).unwrap_err(), MetalError::Circuit(native), "{c:?}");
        }
        assert_eq!(m.simulate(&Circuit::new(MAX_QUBITS + 1)).unwrap_err(), MetalError::Circuit(QuantumError::TooManyQubits(MAX_QUBITS + 1)));
        let small = m.with_max_qubits(4);
        assert_eq!(small.simulate(&Circuit::new(5)).unwrap_err(), MetalError::TooWide { n_qubits: 5, max: 4 });
        let wide = Program { n_qubits: 5, start: 0, gates: Vec::new(), steps: Vec::new() };
        assert_eq!(small.run_program(&wide).unwrap_err(), MetalError::TooWide { n_qubits: 5, max: 4 });
        // A result that leaves 32 bits is reported at its op, and the lane
        // holds its low 32 bits.
        let r = m.run_lanes(&Z_STATE, &[z_words()]).unwrap();
        assert_eq!(r.out_of_range, Some(0), "{}", m.key());
        assert_eq!(r.amps[1][0] as u32, 0x8000_0000);
        // Malformed words, a state the host cannot lay out and a tile step are
        // host defects, never results.
        let bad = GateWords { tbit: 3, ..z_words() };
        assert!(matches!(m.run_lanes(&[[0, 0]; 4], &[bad]), Err(MetalError::Internal(_))));
        assert!(matches!(m.run_lanes(&[[0, 0]; 3], &[z_words()]), Err(MetalError::Internal(_))));
        let tiled = Program { n_qubits: 2, start: 0, gates: Vec::new(), steps: vec![Step::Tile { first: 0, count: 1, size: 2 }] };
        assert!(matches!(m.run_program(&tiled), Err(MetalError::Internal(_))));
        // The refusals that may still fall back to another member, do.
        assert!(MemberError::from(MetalError::TooWide { n_qubits: 5, max: 4 }).falls_back());
        report(&format!("metal tensor-ops refusals: key={} typed", m.key()));
    }
}

/// Whether a member running `mutant`'s text is caught by the mutant's
/// detector.
fn caught(mm: &MetalMember, mutant: &Mutant, key: &str) -> bool {
    match mutant.detector {
        Detector::Probe => {
            let edges = lanes::probe_pairs(0, 0);
            let out = mm.probe(&edges).expect("the mutant's probe runs");
            let differing = edges.iter().zip(&out).filter(|&(&(a, b), o)| *o != (lanes::mul_wide(a, b), lanes::fxmul_wide(a, b))).count();
            println!("metal tensor-ops mutant {} on {key}: {differing} of 484 edge pairs differ", mutant.name);
            differing >= 1
        }
        Detector::Certificate => Gateset::ALL.into_iter().all(|gs| match mm.certify(gs, RANDOM_V1) {
            Err(e) => {
                println!("metal tensor-ops mutant {} on {key}: {gs:?} certification refused: {e}", mutant.name);
                true
            }
            Ok(c) => {
                let cpu = lanes::certify_corpus(gs, "local/test", RANDOM_V1);
                c.parity != Parity::BitExact || c.corpus != cpu.corpus
            }
        }),
        Detector::Range => mm.run_lanes(&Z_STATE, &[z_words()]).map(|r| r.out_of_range) != Ok(Some(0)),
        other => panic!("no tensor-operation mutant uses the {other:?} detector"),
    }
}

#[test]
fn c6_mutants_are_killed() {
    assert_eq!(lanes::probe_pairs(0, 0).len(), 484);
    for m in members() {
        let mut killed = 0;
        for mutant in &TENSOR_OPS_MUTANTS {
            let mm = m.with_tensor_ops_source(&metal::tensor_ops_mutate(mutant)).unwrap_or_else(|e| panic!("{}: mutant {} does not compile: {e}", m.key(), mutant.name));
            assert_ne!(mm.implementation(), m.implementation());
            assert!(caught(&mm, mutant, m.key()), "{}: mutant {} survived", m.key(), mutant.name);
            killed += 1;
        }
        assert_eq!(killed, TENSOR_OPS_MUTANTS.len());
        report(&format!("metal tensor-ops mutants: key={} killed {killed} of {}", m.key(), TENSOR_OPS_MUTANTS.len()));
    }
}

#[test]
fn c7_receipt_placement_and_the_fabric() {
    for m in members() {
        let member = m.member(FABRIC_DEFAULT).unwrap_or_else(|e| panic!("{}: {e}", m.key()));
        assert_eq!((member.id.as_str(), member.class, member.presence), (m.id(), FabricClass::Gpu, Presence::Detected));
        assert_eq!(member.kernels.len(), 2);
        for (entry, gs) in member.kernels.iter().zip(Gateset::ALL) {
            assert_eq!((entry.kernel.as_str(), entry.implementation.as_str()), (apply_kernel(gs), m.implementation()));
            assert!(entry.energy.is_none());
            assert_eq!(entry.certificate.parity, Parity::BitExact, "{gs:?}");
        }
        // A receipt over a member's run.
        let qft = Circuit::qft(5);
        let state = m.simulate(&qft).unwrap();
        assert_eq!(state, qft.simulate().unwrap());
        let entry = member.kernels.iter().find(|k| k.kernel == apply_kernel(qft.gateset)).unwrap();
        let signer = SigningKey::from_bytes(&[7; 32]);
        let receipt = FabricReceipt::seal(&signer, "local/test", state.statevector_hash(), &member, entry, 0, None);
        assert!(receipt.verify());
        assert!(receipt.certificate_matches(&entry.certificate));
        // Placement: beside the CPU reference with no energy figures, the
        // reference; alone, the member.
        let local = detect_local();
        for gs in Gateset::ALL {
            let both = [local.clone(), member.clone()];
            let (placed, k) = place(&both, apply_kernel(gs), &Parity::BitExact).unwrap();
            assert_eq!((placed.id.as_str(), k.implementation.as_str()), (local.id.as_str(), REFERENCE));
            let (placed, k) = place(std::slice::from_ref(&member), apply_kernel(gs), &Parity::BitExact).unwrap();
            assert_eq!((placed.id.as_str(), k.implementation.as_str()), (m.id(), m.implementation()));
        }
        // On the fabric: the member alone runs the circuit; a pairing with
        // another member's entries is refused.
        let mut fabric = Fabric::new();
        fabric.add(member.clone(), m);
        let run = ready(fabric.run(&qft, 0, &Parity::BitExact, &signer, "local/test")).expect("resolves").expect("runs");
        assert_eq!(run.state, state);
        assert_eq!(run.receipt.member_id, m.id());
        assert!(run.receipt.verify());
        let mut mismatched = Fabric::new();
        mismatched.add(local.clone(), m);
        assert!(ready(mismatched.run(&qft, 0, &Parity::BitExact, &signer, "local/test")).expect("resolves").is_err());
        let mut both = Fabric::new();
        both.add(local.clone(), CpuReference).add(member.clone(), m);
        let run = ready(both.run(&qft, 0, &Parity::BitExact, &signer, "local/test")).expect("resolves").expect("runs");
        assert_eq!(run.state, state);
        report(&format!("metal tensor-ops receipt: key={} ok", m.key()));
    }
}

#[test]
fn c8_repeated_runs_agree() {
    for m in members() {
        let c = random_circuit(16, 60, 0xabc, Gateset::V2);
        let first = m.simulate(&c).unwrap();
        for _ in 0..5 {
            assert!(m.simulate(&c).unwrap() == first, "{}", m.key());
        }
        report(&format!("metal tensor-ops repeat: key={} ok", m.key()));
    }
}