wai-quantum 0.5.0

//! The dense f64 state vector on an accelerator, bit for bit: what
//! [`crate::quantum_sv`] computes, planned as launches an accelerator runs.
//!
//! [`crate::quantum_sv::StateVector::run`] applies each gate with one of three
//! formulas, whichever thread count or memory layout computes it:
//!
//! - **a diagonal gate** multiplies each amplitude by its diagonal entry,
//!   leaving it as it is where the entry is exactly 1 (multiplying by 1 can
//!   flip the sign of a zero);
//! - **a gate on one qubit** applies the explicit 2×2 formula
//!   `m00·a − m00ᵢ·aᵢ + m01·b − m01ᵢ·bᵢ`, left to right;
//! - **a gate on two or more qubits** computes each output as a matrix row
//!   against the group's members in local-index order, accumulated from zero.
//!
//! [`plan`] turns a gate list into those launches ([`Launch`]), exactly as the
//! CPU decides them; [`replay`] runs them on the CPU, and must equal
//! [`StateVector::run`](crate::quantum_sv::StateVector::run) bit for bit, which
//! the tests check over a corpus. An accelerator kernel that performs the same
//! IEEE binary64 operations in the same order, with no fused multiply-add,
//! returns the same bits; its certificate says whether it did.
//!
//! The kernel for CUDA devices is [`SOURCE`] (`sv64.cu`).

use crate::gates::Gate;
use crate::linalg::C;
use crate::quantum_fabric::Output;
use crate::quantum_sv::{StateVector, SvError, fuse};

/// The kernel this module certifies: a gate list run on a dense f64 state.
pub const KERNEL_SV_RUN: &str = "wai.quantum.sv/run";

/// The CUDA kernels, `sv64.cu`, exactly as they are compiled.
pub const SOURCE: &str = include_str!("sv64.cu");

/// The implementation name the CUDA id starts with.
pub const IMPLEMENTATION: &str = "cuda-sv64";

/// Version of the host-side layout: separate real and imaginary arrays, one
/// launch per planned gate, the launch kinds and their parameters. Part of the
/// implementation id.
pub const HOST_LAYOUT: u32 = 1;

/// The largest gate a launch takes: `2^k ≤ MAX_D` (the kernel's `SV_MAX_D`).
pub const MAX_D: usize = 64;

/// The compile options: no fused multiply-add, so each operation rounds once.
pub const COMPILE_OPTIONS: [&str; 2] = ["--std=c++17", "--fmad=false"];

const DOMAIN_IMPLEMENTATION: &[u8] = b"wai:quantum-fabric-sv64\x01";

fn put(o: &mut Vec<u8>, bytes: &[u8]) {
    o.extend_from_slice(&(bytes.len() as u64).to_be_bytes());
    o.extend_from_slice(bytes);
}

/// `"cuda-sv64/<16 hex>"` over the kernel text, the options and the layout.
pub fn implementation_id_of(source: &str) -> String {
    let mut o = DOMAIN_IMPLEMENTATION.to_vec();
    put(&mut o, source.as_bytes());
    for option in COMPILE_OPTIONS {
        put(&mut o, option.as_bytes());
    }
    o.extend_from_slice(&HOST_LAYOUT.to_be_bytes());
    let h: String = blake3::hash(&o).as_bytes()[..8].iter().map(|b| format!("{b:02x}")).collect();
    format!("{IMPLEMENTATION}/{h}")
}

pub fn implementation_id() -> String {
    implementation_id_of(SOURCE)
}

/// The OpenCL kernels, `sv64.cl`: the same formulas, with floating-point
/// contraction off (`#pragma OPENCL FP_CONTRACT OFF`), for devices with
/// `cl_khr_fp64`.
pub const CL_SOURCE: &str = include_str!("sv64.cl");

/// The implementation name the OpenCL id starts with.
pub const CL_IMPLEMENTATION: &str = "opencl-sv64";

/// The OpenCL build options.
pub const CL_BUILD_OPTIONS: &str = "-cl-std=CL1.2";

const DOMAIN_CL_IMPLEMENTATION: &[u8] = b"wai:quantum-fabric-sv64-opencl\x01";

/// `"opencl-sv64/<16 hex>"` over the kernel text, the options and the layout.
pub fn cl_implementation_id_of(source: &str) -> String {
    let mut o = DOMAIN_CL_IMPLEMENTATION.to_vec();
    put(&mut o, source.as_bytes());
    put(&mut o, CL_BUILD_OPTIONS.as_bytes());
    o.extend_from_slice(&HOST_LAYOUT.to_be_bytes());
    let h: String = blake3::hash(&o).as_bytes()[..8].iter().map(|b| format!("{b:02x}")).collect();
    format!("{CL_IMPLEMENTATION}/{h}")
}

pub fn cl_implementation_id() -> String {
    cl_implementation_id_of(CL_SOURCE)
}

/// One planned gate.
#[derive(Clone, Debug, PartialEq)]
pub enum Launch {
    /// A diagonal gate: the index bit of each gate qubit (most significant
    /// first) and the diagonal.
    Diag { bitpos: Vec<u32>, diag: Vec<C> },
    /// A gate on one qubit at index bit `s.trailing_zeros()`: its matrix
    /// `[m00, m01, m10, m11]`.
    One { s: u64, m: [C; 4] },
    /// A gate on `k ≥ 2` qubits: the target bits ascending, each local index's
    /// offset, and the `2^k × 2^k` matrix, row-major.
    Many { sorted: Vec<u32>, offsets: Vec<u64>, m: Vec<C> },
}

/// The launches [`StateVector::run`](crate::quantum_sv::StateVector::run) on
/// `n` qubits with `width` performs for `gates`: the same checks, the same
/// fusion, and each gate's formula. A gate on more than `log2(MAX_D)` qubits
/// is refused.
pub fn plan(n: u32, gates: &[Gate], width: usize) -> Result<Vec<Launch>, SvError> {
    for (i, g) in gates.iter().enumerate() {
        let k = g.qubits.len();
        let mut q = g.qubits.clone();
        q.sort_unstable();
        q.dedup();
        if k == 0 || q.len() != k || g.qubits.iter().any(|&x| x >= n) || g.matrix.len() != 1 << (2 * k) {
            return Err(SvError::BadGate(i));
        }
    }
    let fused;
    let list = if width >= 2 {
        fused = fuse(gates, width);
        &fused
    } else {
        gates
    };
    let mut out = Vec::with_capacity(list.len());
    for (i, g) in list.iter().enumerate() {
        let k = g.qubits.len();
        if 1usize << k > MAX_D {
            return Err(SvError::BadGate(i));
        }
        let bitpos: Vec<u32> = g.qubits.iter().map(|&q| n - 1 - q).collect();
        let d = 1usize << k;
        if g.is_diagonal() {
            out.push(Launch::Diag { diag: (0..d).map(|i| g.matrix[i * d + i]).collect(), bitpos });
            continue;
        }
        let offsets: Vec<u64> = (0..d).map(|local| bitpos.iter().enumerate().fold(0u64, |o, (i, &p)| o | ((((local >> (k - 1 - i)) & 1) as u64) << p))).collect();
        if k == 1 {
            out.push(Launch::One { s: offsets[1], m: [g.matrix[0], g.matrix[1], g.matrix[2], g.matrix[3]] });
        } else {
            let mut sorted = bitpos;
            sorted.sort_unstable();
            out.push(Launch::Many { sorted, offsets, m: g.matrix.clone() });
        }
    }
    Ok(out)
}

/// `t` with a zero inserted at each position of `pos`, ascending.
fn spread(mut t: u64, pos: &[u32]) -> u64 {
    for &p in pos {
        let low = t & ((1u64 << p) - 1);
        t = ((t >> p) << (p + 1)) | low;
    }
    t
}

/// Runs `launches` on `(re, im)` on the CPU, as a kernel does: each formula
/// exactly as [`SOURCE`] writes it.
pub fn replay(re: &mut [f64], im: &mut [f64], launches: &[Launch]) {
    let len = re.len() as u64;
    for l in launches {
        match l {
            Launch::Diag { bitpos, diag } => {
                for x in 0..len {
                    let local = bitpos.iter().fold(0u64, |a, &p| (a << 1) | ((x >> p) & 1));
                    let c = diag[local as usize];
                    if c.re == 1.0 && c.im == 0.0 {
                        continue;
                    }
                    let (ar, ai) = (re[x as usize], im[x as usize]);
                    re[x as usize] = c.re * ar - c.im * ai;
                    im[x as usize] = c.re * ai + c.im * ar;
                }
            }
            Launch::One { s, m } => {
                let [m00, m01, m10, m11] = *m;
                for p in 0..len / 2 {
                    let low = p & (s - 1);
                    let i = (((p - low) << 1) | low) as usize;
                    let j = i + *s as usize;
                    let (ar, ai, br, bi) = (re[i], im[i], re[j], im[j]);
                    re[i] = m00.re * ar - m00.im * ai + m01.re * br - m01.im * bi;
                    im[i] = m00.re * ai + m00.im * ar + m01.re * bi + m01.im * br;
                    re[j] = m10.re * ar - m10.im * ai + m11.re * br - m11.im * bi;
                    im[j] = m10.re * ai + m10.im * ar + m11.re * bi + m11.im * br;
                }
            }
            Launch::Many { sorted, offsets, m } => {
                let d = offsets.len();
                let (mut vr, mut vi) = (vec![0.0; d], vec![0.0; d]);
                for t in 0..len >> sorted.len() {
                    let base = spread(t, sorted);
                    for (j, &o) in offsets.iter().enumerate() {
                        vr[j] = re[(base | o) as usize];
                        vi[j] = im[(base | o) as usize];
                    }
                    for (r, &o) in offsets.iter().enumerate() {
                        let (mut sr, mut si) = (0.0, 0.0);
                        for c in 0..d {
                            let e = m[r * d + c];
                            sr += e.re * vr[c] - e.im * vi[c];
                            si += e.re * vi[c] + e.im * vr[c];
                        }
                        re[(base | o) as usize] = sr;
                        im[(base | o) as usize] = si;
                    }
                }
            }
        }
    }
}

/// A corpus of gate lists: the state vector's named gates at angles drawn
/// from a seed, run at fusion widths 0, 2 and 3 in turn.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct SvCorpus {
    pub seed: u64,
    pub cases: usize,
    pub max_qubits: u32,
    pub depth: usize,
}

/// The default: 48 cases of 2 to 10 qubits, 30 gates each.
pub const SV_DEFAULT: SvCorpus = SvCorpus { seed: 7, cases: 48, max_qubits: 10, depth: 30 };

/// The corpus digest of [`SV_DEFAULT`] when the device's state is the CPU's on
/// every case: what a bit-exact certificate reproduces.
pub const SV_DEFAULT_PIN: &str = "b7d129d0593adf70dc164477664970c5c0b8c09f79f1a0573036bd60888dbcd1";

/// One case: the width, the fusion width and the gates.
#[derive(Clone, Debug)]
pub struct SvCase {
    pub n: u32,
    pub width: usize,
    pub gates: Vec<Gate>,
}

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)
}

impl SvCorpus {
    /// `"sv64-v1 seed=… cases=… qubits=… depth=…"`.
    pub fn generator(&self) -> String {
        format!("sv64-v1 seed={} cases={} qubits={} depth={}", self.seed, self.cases, self.max_qubits, self.depth)
    }

    pub fn cases(&self) -> Vec<SvCase> {
        let mut s = self.seed;
        (0..self.cases)
            .map(|i| {
                let n = 2 + (next(&mut s) % u64::from(self.max_qubits.max(2) - 1)) as u32;
                let width = [0usize, 2, 3][i % 3];
                let mut gates = Vec::with_capacity(self.depth);
                for _ in 0..self.depth {
                    let a = (next(&mut s) % u64::from(n)) as u32;
                    let mut b = (next(&mut s) % u64::from(n)) as u32;
                    if b == a {
                        b = (a + 1) % n;
                    }
                    let angle = (next(&mut s) >> 11) as f64 / (1u64 << 53) as f64 * 6.5 - 0.1;
                    gates.push(match next(&mut s) % 9 {
                        0 => Gate::rx(a, angle),
                        1 => Gate::ry(a, angle),
                        2 => Gate::rz(a, angle),
                        3 => Gate::sqrt_w(a),
                        4 => Gate::t(a),
                        5 => Gate::h(a),
                        6 => Gate::fsim(a, b, angle, 0.5 * angle),
                        7 => Gate::cz(a, b),
                        _ => Gate::iswap(a, b),
                    });
                }
                SvCase { n, width, gates }
            })
            .collect()
    }
}

impl SvCase {
    /// The case as bytes: its width, fusion width, and each gate's qubits and
    /// matrix bits.
    pub fn bytes(&self) -> Vec<u8> {
        let mut o = Vec::new();
        o.extend_from_slice(&self.n.to_le_bytes());
        o.extend_from_slice(&(self.width as u64).to_le_bytes());
        for g in &self.gates {
            o.extend_from_slice(&(g.qubits.len() as u64).to_le_bytes());
            for q in &g.qubits {
                o.extend_from_slice(&q.to_le_bytes());
            }
            for c in &g.matrix {
                o.extend_from_slice(&c.re.to_bits().to_le_bytes());
                o.extend_from_slice(&c.im.to_bits().to_le_bytes());
            }
        }
        o
    }

    /// The CPU's state after the case, as an output: real then imaginary
    /// parts, interleaved by amplitude.
    pub fn reference(&self) -> Result<Output, SvError> {
        let mut sv = StateVector::zero(self.n)?;
        sv.run(&self.gates, self.width)?;
        Ok(state_output(&sv.re, &sv.im))
    }
}

/// A state as a certificate output.
pub fn state_output(re: &[f64], im: &[f64]) -> Output {
    Output::Float(re.iter().zip(im).flat_map(|(&a, &b)| [a, b]).collect())
}

/// The mutants a conforming device must be caught running: each changes one
/// operation's order or one formula, so its bits differ from the CPU's.
pub const MUTANTS: [crate::quantum_fabric::wgsl::Mutant; 4] = [
    // The identity skip is lost: a zero's sign may flip.
    crate::quantum_fabric::wgsl::Mutant { name: "S1", find: "        if (cr == 1.0 && ci == 0.0) {\n            continue;\n        }\n", replace: "", detector: crate::quantum_fabric::wgsl::Detector::Certificate },
    // The one-qubit formula is reassociated.
    crate::quantum_fabric::wgsl::Mutant { name: "S2", find: "re[i] = m00r * ar - m00i * ai + m01r * br - m01i * bi;", replace: "re[i] = (m00r * ar - m00i * ai) + (m01r * br - m01i * bi);", detector: crate::quantum_fabric::wgsl::Detector::Certificate },
    // The row's real part is reassociated: (sr + a) − b, not sr + (a − b).
    // (Swapping an addition's operands would not do: IEEE addition commutes.)
    crate::quantum_fabric::wgsl::Mutant { name: "S3", find: "sr += er * vr[c] - ei * vi[c];", replace: "sr += er * vr[c];\n                sr -= ei * vi[c];", detector: crate::quantum_fabric::wgsl::Detector::Certificate },
    // The imaginary row term's sign is wrong.
    crate::quantum_fabric::wgsl::Mutant { name: "S4", find: "si += er * vi[c] + ei * vr[c];", replace: "si += er * vi[c] - ei * vr[c];", detector: crate::quantum_fabric::wgsl::Detector::Certificate },
];

/// [`SOURCE`] with a mutant applied. Panics unless its `find` occurs exactly
/// once.
pub fn mutate(m: &crate::quantum_fabric::wgsl::Mutant) -> String {
    assert_eq!(SOURCE.matches(m.find).count(), 1, "mutant {}", m.name);
    SOURCE.replacen(m.find, m.replace, 1)
}

/// [`CL_SOURCE`] with a mutant applied: the same edits, which the two texts
/// share.
pub fn mutate_cl(m: &crate::quantum_fabric::wgsl::Mutant) -> String {
    assert_eq!(CL_SOURCE.matches(m.find).count(), 1, "mutant {}", m.name);
    CL_SOURCE.replacen(m.find, m.replace, 1)
}

#[cfg(test)]
mod tests {
    use super::*;

    fn bits(o: &Output) -> Vec<u64> {
        match o {
            Output::Float(v) => v.iter().map(|x| x.to_bits()).collect(),
            _ => unreachable!(),
        }
    }

    #[test]
    fn the_replayed_plan_is_the_cpu_run_bit_for_bit() {
        let mut kinds = [0usize; 3];
        for case in SV_DEFAULT.cases() {
            let launches = plan(case.n, &case.gates, case.width).unwrap();
            for l in &launches {
                kinds[match l {
                    Launch::Diag { .. } => 0,
                    Launch::One { .. } => 1,
                    Launch::Many { .. } => 2,
                }] += 1;
            }
            let dim = 1usize << case.n;
            let (mut re, mut im) = (vec![0.0; dim], vec![0.0; dim]);
            re[0] = 1.0;
            replay(&mut re, &mut im, &launches);
            assert_eq!(bits(&state_output(&re, &im)), bits(&case.reference().unwrap()), "n={} width={}", case.n, case.width);
        }
        assert!(kinds.iter().all(|&k| k > 10), "every launch kind is exercised: {kinds:?}");
    }

    #[test]
    fn the_cpu_reproduces_the_pinned_digest() {
        let mut cert = crate::quantum_fabric::Certification::new();
        for case in SV_DEFAULT.cases() {
            let r = case.reference().unwrap();
            cert.add(&case.bytes(), &r, &r);
        }
        let c = cert.finish(KERNEL_SV_RUN, "cpu", "cpu", &SV_DEFAULT.generator());
        let hex: String = c.corpus.iter().map(|b| format!("{b:02x}")).collect();
        assert_eq!(hex, SV_DEFAULT_PIN);
    }

    #[test]
    fn a_replay_with_a_formula_changed_is_caught() {
        // The one-qubit formula, reassociated: the corpus sees the difference.
        let differs = SV_DEFAULT.cases().iter().any(|case| {
            let launches = plan(case.n, &case.gates, case.width).unwrap();
            let dim = 1usize << case.n;
            let (mut re, mut im) = (vec![0.0; dim], vec![0.0; dim]);
            re[0] = 1.0;
            for l in &launches {
                if let Launch::One { s, m } = l {
                    let [m00, m01, m10, m11] = *m;
                    for p in 0..dim as u64 / 2 {
                        let low = p & (s - 1);
                        let i = (((p - low) << 1) | low) as usize;
                        let j = i + *s as usize;
                        let (ar, ai, br, bi) = (re[i], im[i], re[j], im[j]);
                        re[i] = (m00.re * ar - m00.im * ai) + (m01.re * br - m01.im * bi);
                        im[i] = m00.re * ai + m00.im * ar + m01.re * bi + m01.im * br;
                        re[j] = m10.re * ar - m10.im * ai + m11.re * br - m11.im * bi;
                        im[j] = m10.re * ai + m10.im * ar + m11.re * bi + m11.im * br;
                    }
                } else {
                    replay(&mut re, &mut im, std::slice::from_ref(l));
                }
            }
            bits(&state_output(&re, &im)) != bits(&case.reference().unwrap())
        });
        assert!(differs, "the corpus cannot tell a reassociated formula from the reference");
    }

    #[test]
    fn plans_refuse_what_the_state_vector_refuses() {
        let bad = Gate { qubits: vec![0, 0], matrix: vec![C { re: 1.0, im: 0.0 }; 16] };
        assert_eq!(plan(2, &[bad], 0), Err(SvError::BadGate(0)));
        assert_eq!(plan(1, &[Gate::h(1)], 0), Err(SvError::BadGate(0)));
    }

    #[test]
    fn the_implementation_id_follows_the_kernel_and_every_mutant_edits_once() {
        let id = implementation_id();
        assert!(id.starts_with("cuda-sv64/") && id.len() == "cuda-sv64/".len() + 16);
        for m in &MUTANTS {
            assert_ne!(implementation_id_of(&mutate(m)), id, "{}", m.name);
        }
        assert!(SOURCE.contains("#define SV_MAX_D 64") && MAX_D == 64);
        let cl = cl_implementation_id();
        assert!(cl.starts_with("opencl-sv64/") && cl != id);
        for m in &MUTANTS {
            assert_ne!(cl_implementation_id_of(&mutate_cl(m)), cl, "{}", m.name);
        }
        assert!(CL_SOURCE.contains("#pragma OPENCL FP_CONTRACT OFF") && CL_SOURCE.contains("#define SV_MAX_D 64"));
    }
}