wai-quantum 0.6.0

//! The fabric's DSP vector-unit member: the gate's fixed-point arithmetic in the
//! digital signal processor's vector unit, in 128-byte vectors.
//!
//! **Why it is bit-exact.** The kernel ([`KERNEL_TEXT`]) is the same
//! transcription of [`lanes::apply_words`] as the other members', in 32-bit
//! lanes holding each 64-bit value as two words: each product is the vector
//! unit's exact 32×32 multiply (the 32-bit operand times the other's low half,
//! unsigned, plus its high half, signed, sixteen bits up), the rounding constant
//! is added with its carry into the high word, the shift is arithmetic over all
//! 64 bits, and each row is the sum of four, every carry and borrow kept. The
//! host gathers the pairs a gate selects, the kernel forms the four rows of
//! thirty-two pairs at a time, and the host checks each row's range and keeps
//! its low 32 bits, as [`lanes::apply_words`] does.
//!
//! **Where it runs.** The kernels are compiled only in a build for the signal
//! processor with its vector unit in 128-byte mode ([`BUILT_FOR_UNIT`]); their
//! instructions exist nowhere else, so a kernel that ran, ran there. Any other
//! build finds no member and says why. Under an emulator, the runner declares it
//! (`WAI_QUANTUM_DSP_EMULATOR`) and every certificate says so.
//!
//! Three mutants, compiled in from the same kernel text with other constants,
//! must each be caught.

use crate::quantum::{Amp, Circuit, Gateset, MAX_QUBITS, ONE, QuantumError, StateVector};
use crate::quantum_fabric::lanes::{self, GateWords, Lane};
use crate::quantum_fabric::plan::{self, Config, LaneRun, Program, Step};
use crate::quantum_fabric::{Certificate, Certification, Corpus, FabricClass, KernelEntry, Member, MemberError, Presence, Provenance, RunFuture, RunMeter, Runner, apply_kernel};

/// Whether this build is for the signal processor with its vector unit in
/// 128-byte mode.
pub const BUILT_FOR_UNIT: bool = cfg!(all(target_arch = "hexagon", target_feature = "hvx-length128b")); // allow-name: the toolchain's architecture and target-feature names

#[cfg(all(target_arch = "hexagon", target_feature = "hvx-length128b"))] // allow-name: the toolchain's architecture and target-feature names
#[path = "quantum_fabric/dsp.rs"]
mod kernel;

/// Any other build: no member is ever opened ([`discover`] finds none), so
/// nothing here runs.
#[cfg(not(all(target_arch = "hexagon", target_feature = "hvx-length128b")))] // allow-name: the toolchain's architecture and target-feature names
mod kernel {
    pub const LANES: usize = 32;

    /// # Safety
    ///
    /// Never called: no member exists in this build.
    pub unsafe fn rows<const RND: i32, const SH: i32, const LAST_ADD: bool>(_: &[i32; 8], _: [&[i32]; 4], _: [&mut [i32]; 4], _: [&mut [i32]; 4]) {
        unreachable!("no vector unit in this build")
    }

    /// # Safety
    ///
    /// Never called: no member exists in this build.
    pub unsafe fn probe<const RND: i32, const SH: i32>(_: &[i32], _: &[i32], _: [&mut [i32]; 2], _: [&mut [i32]; 2]) {
        unreachable!("no vector unit in this build")
    }
}

/// One probe case: [`lanes::mul_wide`] and [`lanes::fxmul_wide`] as
/// `(high, low)` words.
pub type ProbeOutput = ((u32, u32), (u32, u32));

/// Why the vector unit is not a member, or a run has no state.
#[derive(Debug, PartialEq, Eq)]
#[non_exhaustive]
pub enum DspError {
    Unavailable(String),
    Circuit(QuantumError),
    TooWide { n_qubits: u8, max: u8 },
    OutOfRange { gate: u32 },
    Internal(&'static str),
    Probe { a: i32, b: i32 },
}

impl std::fmt::Display for DspError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            DspError::Unavailable(why) => write!(f, "no DSP vector unit: {why}"),
            DspError::Circuit(e) => write!(f, "{e}"),
            DspError::TooWide { n_qubits, max } => write!(f, "{n_qubits} qubits; this member holds at most {max}"),
            DspError::OutOfRange { gate } => write!(f, "the result of op {gate} leaves 32 bits"),
            DspError::Internal(why) => write!(f, "host defect: {why}"),
            DspError::Probe { a, b } => write!(f, "the kernel's arithmetic differs from the CPU's at ({a}, {b})"),
        }
    }
}

impl std::error::Error for DspError {}

impl From<DspError> for MemberError {
    fn from(e: DspError) -> MemberError {
        match e {
            DspError::TooWide { n_qubits, max } => MemberError::TooWide { n_qubits, max },
            DspError::OutOfRange { gate } => MemberError::OutOfRange { gate },
            DspError::Circuit(e) => MemberError::Circuit(e),
            other => MemberError::Other(other.to_string()),
        }
    }
}

/// The kernels a member runs: the conforming one, or a mutant.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub enum Variant {
    Conforming,
    /// The rounding constant is half what it should be.
    M1,
    /// The product is shifted by one bit too few.
    M2,
    /// Row 0's real part adds its last term instead of subtracting it.
    M3,
}

impl Variant {
    pub const MUTANTS: [Variant; 3] = [Variant::M1, Variant::M2, Variant::M3];

    /// The kernel's constants: the rounding constant, the shift, and whether
    /// row 0's last real term is added.
    pub fn constants(self) -> (i32, i32, bool) {
        match self {
            Variant::Conforming => (1 << 29, 30, false),
            Variant::M1 => (1 << 28, 30, false),
            Variant::M2 => (1 << 29, 29, false),
            Variant::M3 => (1 << 29, 30, true),
        }
    }
}

/// The text the kernels are compiled from.
pub const KERNEL_TEXT: &str = include_str!("quantum_fabric/dsp.rs");

/// The implementation name the id starts with.
pub const IMPLEMENTATION: &str = "dsp-lanes64";

/// Version of the host layout: the gather of a gate's pairs into four
/// columns padded with zeros to whole vectors, the rows' high and low words
/// joined into 64 bits, and the host's range check. Part of the
/// implementation id.
pub const HOST_LAYOUT: u32 = 1;

/// `"dsp-lanes64/<16 hex>"` over the kernel text, a variant's constants and
/// the layout.
pub fn implementation_id_of(v: Variant) -> String {
    let (rnd, sh, last_add) = v.constants();
    let mut o = b"wai:quantum-fabric-dsp\x01".to_vec();
    o.extend_from_slice(&(KERNEL_TEXT.len() as u64).to_be_bytes());
    o.extend_from_slice(KERNEL_TEXT.as_bytes());
    o.extend_from_slice(&rnd.to_be_bytes());
    o.extend_from_slice(&sh.to_be_bytes());
    o.push(u8::from(last_add));
    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(Variant::Conforming)
}

/// Whether this build runs on the signal processor's vector unit.
pub fn has_vector_unit() -> Result<(), String> {
    if BUILT_FOR_UNIT {
        Ok(())
    } else {
        Err(format!("this build ({}) is not for the signal processor with its vector unit in 128-byte mode", std::env::consts::ARCH))
    }
}

/// What [`discover`] found.
#[derive(Debug)]
pub struct Discovery {
    pub found: Vec<DspMember>,
    pub refused: Vec<(String, DspError)>,
    pub unavailable: Option<String>,
}

/// The vector unit, probed in this process, or why there is none.
pub fn discover() -> Discovery {
    let mut d = Discovery { found: Vec::new(), refused: Vec::new(), unavailable: None };
    if let Err(why) = has_vector_unit() {
        d.unavailable = Some(why);
        return d;
    }
    match DspMember::open(Variant::Conforming) {
        Ok(m) => d.found.push(m),
        Err(e) => d.refused.push((KEY.into(), e)),
    }
    d
}

/// The key a runner declares.
pub const KEY: &str = "dsp-vector";

/// The most qubits a run holds.
pub const MAX_DSP_QUBITS: u8 = if 24 < MAX_QUBITS { 24 } else { MAX_QUBITS };

/// The DSP's vector unit as a fabric member.
#[derive(Clone, Debug)]
pub struct DspMember {
    variant: Variant,
    id: String,
    implementation: String,
    detail: String,
    provenance: Provenance,
    max_qubits: u8,
}

/// `v` padded with zeros to whole vectors.
fn padded(v: &[i32]) -> Vec<i32> {
    let mut p = v.to_vec();
    p.resize(v.len().next_multiple_of(kernel::LANES), 0);
    p
}

fn join(hi: i32, lo: i32) -> i64 {
    (i64::from(hi) << 32) | i64::from(lo as u32)
}

impl DspMember {
    fn open(variant: Variant) -> Result<DspMember, DspError> {
        let emulator = std::env::var("WAI_QUANTUM_DSP_EMULATOR").ok().filter(|s| !s.trim().is_empty());
        let provenance = emulator.as_ref().map_or(Provenance::Silicon, |e| Provenance::Emulator(e.trim().to_string()));
        let mut o = b"wai:quantum-fabric-dsp-member\x01".to_vec();
        o.extend_from_slice(std::env::consts::ARCH.as_bytes());
        o.extend_from_slice(std::env::consts::OS.as_bytes());
        o.extend_from_slice(provenance.label().as_bytes());
        let id = format!("local/dsp-vector-{}", blake3::hash(&o).as_bytes()[..6].iter().map(|b| format!("{b:02x}")).collect::<String>());
        let member = DspMember {
            variant,
            id,
            implementation: implementation_id_of(variant),
            detail: format!("dsp; the signal processor's vector unit, 128-byte vectors, 64-bit values as word pairs; {}; up to {MAX_DSP_QUBITS} qubits; unmetered", provenance.label()),
            provenance,
            max_qubits: MAX_DSP_QUBITS,
        };
        if variant == Variant::Conforming {
            let pairs = lanes::probe_pairs(4096, 0x5eed);
            let out = member.probe(&pairs)?;
            if let Some((a, b)) = lanes::first_mismatch(&pairs, &out) {
                return Err(DspError::Probe { a, b });
            }
        }
        Ok(member)
    }

    /// The same processor running a mutant: for the mutation tests.
    #[doc(hidden)]
    pub fn with_variant(&self, v: Variant) -> DspMember {
        DspMember { variant: v, implementation: implementation_id_of(v), ..self.clone() }
    }

    #[doc(hidden)]
    pub fn with_max_qubits(&self, max_qubits: u8) -> DspMember {
        DspMember { max_qubits: max_qubits.min(self.max_qubits), ..self.clone() }
    }

    /// The kernel this member runs: the conforming one, or a mutant.
    pub fn variant(&self) -> Variant {
        self.variant
    }

    pub fn id(&self) -> &str {
        &self.id
    }

    pub fn key(&self) -> &str {
        KEY
    }

    pub fn class(&self) -> FabricClass {
        FabricClass::Dsp
    }

    pub fn detail(&self) -> &str {
        &self.detail
    }

    pub fn implementation(&self) -> &str {
        &self.implementation
    }

    pub fn provenance(&self) -> &Provenance {
        &self.provenance
    }

    pub fn max_qubits(&self) -> u8 {
        self.max_qubits
    }

    pub fn config(&self) -> Config {
        Config::PER_GATE
    }

    /// The kernel's [`lanes::mul_wide`] and [`lanes::fxmul_wide`] of each pair.
    pub fn probe(&self, pairs: &[(i32, i32)]) -> Result<Vec<ProbeOutput>, DspError> {
        {
            let a = padded(&pairs.iter().map(|p| p.0).collect::<Vec<i32>>());
            let b = padded(&pairs.iter().map(|p| p.1).collect::<Vec<i32>>());
            let n = a.len();
            let (mut ph, mut pl, mut fh, mut fl) = (vec![0i32; n], vec![0i32; n], vec![0i32; n], vec![0i32; n]);
            // SAFETY: this build targets the vector unit; every slice holds n
            // values, a whole number of vectors.
            unsafe {
                let probe = match self.variant {
                    Variant::M1 => kernel::probe::<{ 1 << 28 }, 30>,
                    Variant::M2 => kernel::probe::<{ 1 << 29 }, 29>,
                    Variant::Conforming | Variant::M3 => kernel::probe::<{ 1 << 29 }, 30>,
                };
                probe(&a, &b, [&mut ph[..], &mut pl[..]], [&mut fh[..], &mut fl[..]]);
            }
            let split = |v: i64| ((v as u64 >> 32) as u32, v as u64 as u32);
            Ok((0..pairs.len()).map(|k| (split(join(ph[k], pl[k])), split(join(fh[k], fl[k])))).collect())
        }
    }

    /// The four rows of each pair, from the kernel.
    fn rows(&self, m: &[i32; 8], cols: &[Vec<i32>; 4]) -> [Vec<i64>; 4] {
        let len = cols[0].len();
        let cols = [padded(&cols[0]), padded(&cols[1]), padded(&cols[2]), padded(&cols[3])];
        let n = cols[0].len();
        let mut hi = [vec![0i32; n], vec![0i32; n], vec![0i32; n], vec![0i32; n]];
        let mut lo = hi.clone();
        {
            let [h0, h1, h2, h3] = &mut hi;
            let [l0, l1, l2, l3] = &mut lo;
            let c = [&cols[0][..], &cols[1][..], &cols[2][..], &cols[3][..]];
            // SAFETY: this build targets the vector unit; every slice holds n
            // values, a whole number of vectors.
            unsafe {
                let rows = match self.variant {
                    Variant::Conforming => kernel::rows::<{ 1 << 29 }, 30, false>,
                    Variant::M1 => kernel::rows::<{ 1 << 28 }, 30, false>,
                    Variant::M2 => kernel::rows::<{ 1 << 29 }, 29, false>,
                    Variant::M3 => kernel::rows::<{ 1 << 29 }, 30, true>,
                };
                rows(m, c, [&mut h0[..], &mut h1[..], &mut h2[..], &mut h3[..]], [&mut l0[..], &mut l1[..], &mut l2[..], &mut l3[..]]);
            }
        }
        let row = |r: usize| (0..len).map(|k| join(hi[r][k], lo[r][k])).collect::<Vec<i64>>();
        [row(0), row(1), row(2), row(3)]
    }

    /// Gates in order from `state`, as [`lanes::run_words`] runs them.
    pub fn run_lanes(&self, state: &[Lane], gates: &[GateWords]) -> Result<LaneRun, DspError> {
        let len = state.len();
        if len < 2 || !len.is_power_of_two() {
            return Err(DspError::Internal("a state that is not 2^n lanes, n at least 1"));
        }
        let n = len.trailing_zeros() as u8;
        if n > self.max_qubits {
            return Err(DspError::TooWide { n_qubits: n, max: self.max_qubits });
        }
        let planned: Vec<(&GateWords, u32)> = gates.iter().zip(0u32..).collect();
        self.execute(state.to_vec(), &planned)
    }

    fn execute(&self, mut amps: Vec<Lane>, gates: &[(&GateWords, u32)]) -> Result<LaneRun, DspError> {
        let len = amps.len();
        for (w, _) in gates {
            let (tbit, ctrl) = (w.tbit as usize, w.ctrl as usize);
            if tbit == 0 || !tbit.is_power_of_two() || tbit > len / 2 || ctrl & tbit != 0 || ctrl & !(len - 1) != 0 {
                return Err(DspError::Internal("malformed gate words for this state"));
            }
        }
        {
            let mut report: Option<u32> = None;
            for &(w, index) in gates {
                let m: [i32; 8] = w.m.map(|x| x as i32);
                let (tbit, ctrl) = (w.tbit as usize, w.ctrl as usize);
                let mut at = Vec::new();
                let mut cols: [Vec<i32>; 4] = Default::default();
                for p in 0..len / 2 {
                    let low = p & (tbit - 1);
                    let i = ((p - low) << 1) | low;
                    if i & ctrl != ctrl {
                        continue;
                    }
                    at.push(i);
                    let (a0, a1) = (amps[i], amps[i | tbit]);
                    cols[0].push(a0[0]);
                    cols[1].push(a0[1]);
                    cols[2].push(a1[0]);
                    cols[3].push(a1[1]);
                }
                let [o0r, o0i, o1r, o1i] = self.rows(&m, &cols);
                for (k, &i) in at.iter().enumerate() {
                    let rows = [o0r[k], o0i[k], o1r[k], o1i[k]];
                    if rows.iter().any(|&v| i32::try_from(v).is_err()) {
                        report = Some(report.map_or(index, |r| r.min(index)));
                    }
                    let low = |v: i64| v as u64 as u32 as i32;
                    amps[i] = [low(rows[0]), low(rows[1])];
                    amps[i | tbit] = [low(rows[2]), low(rows[3])];
                }
            }
            Ok(LaneRun { amps, out_of_range: report })
        }
    }

    pub fn run_program(&self, p: &Program) -> Result<LaneRun, DspError> {
        let n = p.n_qubits;
        if n == 0 {
            return Ok(LaneRun { amps: vec![[ONE as i32, 0]], out_of_range: None });
        }
        if n > self.max_qubits {
            return Err(DspError::TooWide { n_qubits: n, max: self.max_qubits });
        }
        if p.start >= 1usize << n {
            return Err(DspError::Internal("the start is outside the state"));
        }
        let mut gates = Vec::with_capacity(p.steps.len());
        for step in &p.steps {
            let Step::Gate { gate } = *step else { return Err(DspError::Internal("a tile step; this member runs gate steps only")) };
            let g = p.gates.get(gate as usize).ok_or(DspError::Internal("a step names a gate the program does not hold"))?;
            if !g.is_well_formed() {
                return Err(DspError::Internal("malformed gate words for this state"));
            }
            gates.push((&g.words, g.index));
        }
        let mut amps = vec![[0i32, 0]; 1usize << n];
        amps[p.start] = [ONE as i32, 0];
        self.execute(amps, &gates)
    }

    pub fn simulate_from(&self, c: &Circuit, start: usize) -> Result<StateVector, DspError> {
        let p = plan::compile_with(c, start, &self.config()).map_err(DspError::Circuit)?;
        if p.n_qubits > self.max_qubits {
            return Err(DspError::TooWide { n_qubits: p.n_qubits, max: self.max_qubits });
        }
        let run = self.run_program(&p)?;
        if let Some(gate) = run.out_of_range {
            return Err(DspError::OutOfRange { gate });
        }
        Ok(StateVector { n_qubits: p.n_qubits, amps: run.amps.iter().map(|l| Amp { re: i64::from(l[0]), im: i64::from(l[1]) }).collect() })
    }

    pub fn simulate(&self, c: &Circuit) -> Result<StateVector, DspError> {
        self.simulate_from(c, 0)
    }

    /// The certificate of [`apply_kernel`]`(gs)` on this member, with its
    /// provenance.
    pub fn certify(&self, gs: Gateset, corpus: Corpus) -> Result<Certificate, DspError> {
        let mut cert = Certification::new();
        for c in corpus.circuits(gs) {
            let reference = lanes::reference_output(&c).map_err(DspError::Circuit)?;
            let run = self.run_program(&plan::compile_with(&c, 0, &self.config()).map_err(DspError::Circuit)?)?;
            if let Some(gate) = run.out_of_range {
                return Err(DspError::OutOfRange { gate });
            }
            cert.add(&lanes::circuit_bytes(&c), &reference, &lanes::output(&run.amps));
        }
        Ok(cert.finish(apply_kernel(gs), &self.implementation, &self.id, &corpus.generator(gs)).with_provenance(self.provenance.clone()))
    }

    pub fn member(&self, corpus: Corpus) -> Result<Member, DspError> {
        let mut kernels = Vec::new();
        for gs in Gateset::ALL {
            kernels.push(KernelEntry { kernel: apply_kernel(gs).into(), implementation: self.implementation.clone(), certificate: self.certify(gs, corpus)?, energy: None });
        }
        Ok(Member { id: self.id.clone(), class: FabricClass::Dsp, presence: Presence::Detected, kernels, detail: self.detail.clone() })
    }
}

impl Runner for DspMember {
    fn member_id(&self) -> String {
        self.id.clone()
    }

    fn run<'r>(&'r self, implementation: &'r str, c: &'r Circuit, start: usize, _meter: &'r mut dyn RunMeter) -> RunFuture<'r> {
        Box::pin(async move {
            if implementation != self.implementation {
                return Err(MemberError::Other(format!("member {} runs implementation {}, not {implementation:?}", self.id, self.implementation)));
            }
            Ok((self.simulate_from(c, start)?, 0, None))
        })
    }
}

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

    #[test]
    fn every_variant_has_its_own_implementation_id() {
        let ids: std::collections::BTreeSet<String> = [Variant::Conforming, Variant::M1, Variant::M2, Variant::M3].into_iter().map(implementation_id_of).collect();
        assert_eq!(ids.len(), 4);
        assert!(implementation_id().starts_with("dsp-lanes64/"));
        assert!(KERNEL_TEXT.contains("Q6_W_vmpye_VwVuh") && KERNEL_TEXT.contains("Q6_W_vmpyoacc_WVwVh"));
    }

    #[test]
    fn the_constants_are_the_reference_arithmetic() {
        assert_eq!(Variant::Conforming.constants(), (1 << 29, 30, false));
        for v in Variant::MUTANTS {
            assert_ne!(v.constants(), Variant::Conforming.constants());
        }
    }

    #[test]
    fn the_member_id_sorts_after_the_cpu_reference() {
        assert!("local/dsp-vector-000000000000" > "local/cpu-aarch64");
        assert!("local/dsp-vector-000000000000" > "local/cpu-x86_64");
    }

    #[test]
    fn elsewhere_the_member_is_absent_and_says_why() {
        if !BUILT_FOR_UNIT {
            let d = discover();
            assert!(d.found.is_empty() && d.refused.is_empty());
            assert!(d.unavailable.as_deref().is_some_and(|why| why.contains("is not for the signal processor")), "{:?}", d.unavailable);
        }
    }
}