wai-quantum 0.6.0

//! The fabric's tensor-engine member: the gate's products on a tensor
//! processor's matrix engine, as exact limb convolutions in one bf16 matrix
//! multiply per tile of pairs, accumulated in fp32. See
//! [`quantum_fabric::tensor_engine`](crate::quantum_fabric::tensor_engine)
//! for the layout and the argument that it is bit-exact on any engine that
//! multiplies bf16 and accumulates in fp32.
//!
//! - **Discovery** ([`discover`]): the kernel ([`tensor_engine::HELPER`]) is
//!   run as a long-lived helper by the Python the host names
//!   (`WAI_QUANTUM_TENSOR_ENGINE_PYTHON`, an environment with the engine's
//!   kernel toolchain). The helper says where the kernel runs (its target)
//!   and what it runs with; its arithmetic is then probed against [`lanes`]
//!   in this process. The target is the toolchain's simulator unless the
//!   host names a device platform (`WAI_QUANTUM_TENSOR_ENGINE_TARGET`, for
//!   instance `trn1`), where the kernel is compiled for the device and runs
//!   on it, and the provenance is silicon. A host with no such Python finds no
//!   member and says why.
//! - **Running**: the host gathers the pairs a gate selects and sends their
//!   limbs and the gate's limb weights; the engine returns every product's
//!   coefficients; the host checks each is an exact integer, assembles the
//!   products and applies the fixed-point arithmetic, checking each row's
//!   range and keeping its low 32 bits.
//!
//! # Honest boundaries
//!
//! - **The target is the helper's.** Today it is the kernel toolchain's
//!   simulator, whose arithmetic is a model of the engine's, not the engine:
//!   every certificate's provenance says so (`simulator:<versions>`), and
//!   placement ranks it after silicon. The same kernel and the same argument
//!   are what a device run would use.
//! - **The engine multiplies; the host assembles**, as the neural-processor
//!   member's host does.
//! - **Unmetered**, as every member is by default.

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::tensor_engine::{self, DeviceIdentity, Variant};
use crate::quantum_fabric::{Certificate, Certification, Corpus, FabricClass, KernelEntry, Member, MemberError, Presence, Provenance, RunFuture, RunMeter, Runner, apply_kernel};
use std::io::{BufRead, BufReader, Read, Write};
use std::path::PathBuf;
use std::process::{Child, ChildStdin, ChildStdout, Command, Stdio};
use std::sync::{Arc, Mutex};

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

/// Why the engine is not a member, or a run has no state.
#[derive(Debug, PartialEq, Eq)]
#[non_exhaustive]
pub enum TensorEngineError {
    /// No Python named, or it has no kernel toolchain, and why.
    Unavailable(String),
    /// The helper failed or said why it could not run a request.
    Device(String),
    /// A coefficient came back that is not an exact integer below 2^24.
    Inexact { product: usize, case: usize },
    /// Exactly what the CPU reference refuses.
    Circuit(QuantumError),
    /// More qubits than this member holds.
    TooWide { n_qubits: u8, max: u8 },
    /// The reference's result leaves 32 bits at this op; the member returns
    /// no state.
    OutOfRange { gate: u32 },
    /// A program the host cannot lay out: a defect of the host, never a result.
    Internal(&'static str),
    /// The engine's products disagreed with the CPU's on this pair.
    Probe { a: i32, b: i32 },
}

impl std::fmt::Display for TensorEngineError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            TensorEngineError::Unavailable(why) => write!(f, "no matrix engine to drive: {why}"),
            TensorEngineError::Device(why) => write!(f, "the engine's driver: {why}"),
            TensorEngineError::Inexact { product, case } => write!(f, "coefficient of product {product} for case {case} is not an exact integer below 2^24"),
            TensorEngineError::Circuit(e) => write!(f, "{e}"),
            TensorEngineError::TooWide { n_qubits, max } => write!(f, "{n_qubits} qubits; this member holds at most {max}"),
            TensorEngineError::OutOfRange { gate } => write!(f, "the result of op {gate} leaves 32 bits"),
            TensorEngineError::Internal(why) => write!(f, "host defect: {why}"),
            TensorEngineError::Probe { a, b } => write!(f, "the engine's products differ from the CPU's at ({a}, {b})"),
        }
    }
}

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

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

/// The request tag the helper checks.
const MAGIC: u32 = 0x4554_5157;

/// The helper process: its pipes, used one request at a time.
struct Helper {
    child: Child,
    stdin: ChildStdin,
    stdout: BufReader<ChildStdout>,
}

impl Drop for Helper {
    fn drop(&mut self) {
        let _ = self.child.kill();
        let _ = self.child.wait();
    }
}

/// The helper, started with `python`, and its first line: the target and the
/// versions it reports.
fn start(python: &PathBuf) -> Result<(Helper, String, String), String> {
    let dir = std::env::var_os("WAI_QUANTUM_TENSOR_ENGINE_CACHE").map(PathBuf::from).unwrap_or_else(|| std::env::temp_dir().join("wai-quantum-tensor-engine"));
    std::fs::create_dir_all(&dir).map_err(|e| format!("{}: {e}", dir.display()))?;
    let hash: String = blake3::hash(tensor_engine::HELPER.as_bytes()).as_bytes()[..6].iter().map(|b| format!("{b:02x}")).collect();
    let script = dir.join(format!("helper-{hash}.py"));
    if !script.exists() {
        let staging = dir.join(format!("helper-{hash}.py.{}", std::process::id()));
        std::fs::write(&staging, tensor_engine::HELPER).map_err(|e| format!("{}: {e}", staging.display()))?;
        std::fs::rename(&staging, &script).map_err(|e| format!("{}: {e}", script.display()))?;
    }
    // The simulator unless the host names a device platform.
    let target = std::env::var("WAI_QUANTUM_TENSOR_ENGINE_TARGET").ok().filter(|t| !t.is_empty()).unwrap_or_else(|| "simulator".into());
    if !target.bytes().all(|b| b.is_ascii_alphanumeric() || b == b'-') {
        return Err(format!("WAI_QUANTUM_TENSOR_ENGINE_TARGET={target:?} is not a platform name"));
    }
    let mut child = Command::new(python)
        .arg("-I")
        .arg(&script)
        .args(["--target", &target])
        .stdin(Stdio::piped())
        .stdout(Stdio::piped())
        .stderr(Stdio::null())
        .spawn()
        .map_err(|e| format!("{}: {e}", python.display()))?;
    let stdin = child.stdin.take().ok_or("no pipe to the helper")?;
    let mut stdout = BufReader::new(child.stdout.take().ok_or("no pipe from the helper")?);
    let mut line = String::new();
    stdout.read_line(&mut line).map_err(|e| format!("the helper's first line: {e}"))?;
    let mut words = line.split_whitespace();
    if (words.next(), words.next()) != (Some("wq-tensor-engine"), Some("1")) {
        let _ = child.kill();
        return Err(format!("{} did not start the kernel's driver (it said {:?}); does it have the kernel toolchain?", python.display(), line.trim()));
    }
    let target = words.next().unwrap_or("").to_string();
    let versions = words.collect::<Vec<_>>().join(" ");
    if target.is_empty() || !target.bytes().all(|b| b.is_ascii_alphanumeric() || b == b'-') {
        let _ = child.kill();
        return Err(format!("the helper named no target ({:?})", line.trim()));
    }
    Ok((Helper { child, stdin, stdout }, target, versions))
}

impl Helper {
    /// One request: the coefficients of `M × n` outputs.
    fn coefficients(&mut self, weights: &[u16], moving: &[u16], n: usize) -> Result<Vec<f32>, TensorEngineError> {
        let io = |e: std::io::Error| TensorEngineError::Device(format!("the pipe to the helper: {e}"));
        let mut req = Vec::with_capacity(8 + 2 * (weights.len() + moving.len()));
        req.extend_from_slice(&MAGIC.to_le_bytes());
        req.extend_from_slice(&(n as u32).to_le_bytes());
        for w in weights.iter().chain(moving) {
            req.extend_from_slice(&w.to_le_bytes());
        }
        self.stdin.write_all(&req).map_err(io)?;
        self.stdin.flush().map_err(io)?;
        let mut status = [0u8; 4];
        self.stdout.read_exact(&mut status).map_err(io)?;
        if u32::from_le_bytes(status) != 0 {
            let mut len = [0u8; 4];
            self.stdout.read_exact(&mut len).map_err(io)?;
            let mut msg = vec![0u8; u32::from_le_bytes(len).min(1 << 16) as usize];
            self.stdout.read_exact(&mut msg).map_err(io)?;
            return Err(TensorEngineError::Device(String::from_utf8_lossy(&msg).into_owned()));
        }
        let mut bytes = vec![0u8; tensor_engine::M * n * 4];
        self.stdout.read_exact(&mut bytes).map_err(io)?;
        Ok(bytes.as_chunks::<4>().0.iter().map(|b| f32::from_le_bytes(*b)).collect())
    }
}

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

/// The matrix engine the named Python drives, 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 };
    let Some(python) = std::env::var_os("WAI_QUANTUM_TENSOR_ENGINE_PYTHON").filter(|s| !s.is_empty()).map(PathBuf::from) else {
        d.unavailable = Some("no kernel toolchain named (WAI_QUANTUM_TENSOR_ENGINE_PYTHON, the Python of an environment with the matrix engine's kernel toolchain)".into());
        return d;
    };
    match start(&python) {
        Err(why) => d.unavailable = Some(why),
        Ok((helper, target, versions)) => {
            let identity = DeviceIdentity::new(target, versions);
            let key = identity.key();
            match TensorEngineMember::open(helper, identity) {
                Ok(m) => d.found.push(m),
                Err(e) => d.refused.push((key, e)),
            }
        }
    }
    d
}

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

/// A tensor processor's matrix engine as a fabric member.
#[derive(Clone)]
pub struct TensorEngineMember {
    helper: Arc<Mutex<Helper>>,
    variant: Variant,
    identity: DeviceIdentity,
    provenance: Provenance,
    id: String,
    key: String,
    implementation: String,
    detail: String,
    max_qubits: u8,
}

impl std::fmt::Debug for TensorEngineMember {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("TensorEngineMember").field("id", &self.id).field("key", &self.key).field("implementation", &self.implementation).finish()
    }
}

impl TensorEngineMember {
    fn open(helper: Helper, identity: DeviceIdentity) -> Result<TensorEngineMember, TensorEngineError> {
        // Anything but a device the helper names is a simulation of the engine.
        let provenance = if identity.target == "simulator" { Provenance::Simulator(identity.versions.clone()) } else { Provenance::Silicon };
        let member = TensorEngineMember {
            helper: Arc::new(Mutex::new(helper)),
            variant: Variant::Conforming,
            provenance: provenance.clone(),
            id: identity.member_id(),
            key: identity.key(),
            implementation: tensor_engine::implementation_id(),
            detail: format!(
                "tensor; the matrix engine, bf16 limb products accumulated in fp32, {} pairs a multiply; target {}; {}; up to {MAX_TENSOR_ENGINE_QUBITS} qubits; {}; unmetered",
                tensor_engine::TILE,
                identity.target,
                identity.versions,
                provenance.label()
            ),
            identity,
            max_qubits: MAX_TENSOR_ENGINE_QUBITS,
        };
        let pairs = lanes::probe_pairs(4096, 0x5eed);
        let out = member.probe(&pairs)?;
        if let Some((a, b)) = lanes::first_mismatch(&pairs, &out) {
            return Err(TensorEngineError::Probe { a, b });
        }
        Ok(member)
    }

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

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

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

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

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

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

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

    pub fn variant(&self) -> Variant {
        self.variant
    }

    pub fn identity(&self) -> &DeviceIdentity {
        &self.identity
    }

    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 engine's sixteen products of each case, for stationary weights `w`.
    fn run(&self, w: &[u16], cases: &[[i32; 4]]) -> Result<Vec<[i64; tensor_engine::PRODUCTS]>, TensorEngineError> {
        let bits = self.variant.bits();
        let (a, n) = tensor_engine::moving(cases, bits);
        let c = self.helper.lock().unwrap_or_else(|p| p.into_inner()).coefficients(w, &a, n)?;
        tensor_engine::assemble(&c, n, cases.len(), bits).map_err(|e| TensorEngineError::Inexact { product: e.product, case: e.case })
    }

    /// Each case's sixteen products under gate entries `m`, from the engine.
    fn products(&self, m: &[i32; 8], cases: &[[i32; 4]]) -> Result<Vec<[i64; tensor_engine::PRODUCTS]>, TensorEngineError> {
        self.run(&tensor_engine::weights(m, self.variant.bits()), cases)
    }

    /// The engine's [`lanes::mul_wide`] of each pair, and the host's
    /// [`lanes::fxmul_wide`] of that product: sixteen pairs a multiply
    /// ([`tensor_engine::probe_weights`]).
    pub fn probe(&self, pairs: &[(i32, i32)]) -> Result<Vec<ProbeOutput>, TensorEngineError> {
        let (rnd, sh, _) = self.variant.constants();
        let split = |v: i64| ((v as u64 >> 32) as u32, v as u64 as u32);
        let mut out = Vec::with_capacity(pairs.len());
        for chunk in pairs.chunks(tensor_engine::PROBE_PAIRS) {
            let a: [i32; tensor_engine::PROBE_PAIRS] = std::array::from_fn(|q| chunk.get(q).map_or(0, |p| p.0));
            let b: [i32; tensor_engine::PROBE_PAIRS] = std::array::from_fn(|q| chunk.get(q).map_or(0, |p| p.1));
            let ps = self.run(&tensor_engine::probe_weights(&b, self.variant.bits()), &tensor_engine::probe_moving(&a))?;
            for q in 0..chunk.len() {
                let v = ps[q / 4][q];
                out.push((split(v), split((v + rnd) >> sh)));
            }
        }
        Ok(out)
    }

    /// Gates in order from `state`, as [`lanes::run_words`] runs them.
    pub fn run_lanes(&self, state: &[Lane], gates: &[GateWords]) -> Result<LaneRun, TensorEngineError> {
        let len = state.len();
        if len < 2 || !len.is_power_of_two() {
            return Err(TensorEngineError::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(TensorEngineError::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, TensorEngineError> {
        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(TensorEngineError::Internal("malformed gate words for this state"));
            }
        }
        let (rnd, sh, last_add) = self.variant.constants();
        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 cases = Vec::new();
            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]);
                cases.push([a0[0], a0[1], a1[0], a1[1]]);
            }
            let products = self.products(&m, &cases)?;
            for (&i, p) in at.iter().zip(&products) {
                let r = tensor_engine::rows(p, rnd, sh, last_add);
                if r.iter().any(|&v| i32::try_from(v).is_err()) {
                    report = Some(report.map_or(index, |x| x.min(index)));
                }
                let low = |v: i64| v as u64 as u32 as i32;
                amps[i] = [low(r[0]), low(r[1])];
                amps[i | tbit] = [low(r[2]), low(r[3])];
            }
        }
        Ok(LaneRun { amps, out_of_range: report })
    }

    pub fn run_program(&self, p: &Program) -> Result<LaneRun, TensorEngineError> {
        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(TensorEngineError::TooWide { n_qubits: n, max: self.max_qubits });
        }
        if p.start >= 1usize << n {
            return Err(TensorEngineError::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(TensorEngineError::Internal("a tile step; this member runs gate steps only")) };
            let g = p.gates.get(gate as usize).ok_or(TensorEngineError::Internal("a step names a gate the program does not hold"))?;
            if !g.is_well_formed() {
                return Err(TensorEngineError::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, TensorEngineError> {
        let p = plan::compile_with(c, start, &self.config()).map_err(TensorEngineError::Circuit)?;
        if p.n_qubits > self.max_qubits {
            return Err(TensorEngineError::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(TensorEngineError::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, TensorEngineError> {
        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, TensorEngineError> {
        let mut cert = Certification::new();
        for c in corpus.circuits(gs) {
            let reference = lanes::reference_output(&c).map_err(TensorEngineError::Circuit)?;
            let run = self.run_program(&plan::compile_with(&c, 0, &self.config()).map_err(TensorEngineError::Circuit)?)?;
            if let Some(gate) = run.out_of_range {
                return Err(TensorEngineError::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, TensorEngineError> {
        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::Tensor, presence: Presence::Detected, kernels, detail: self.detail.clone() })
    }
}

impl Runner for TensorEngineMember {
    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 the_refusals_that_may_fall_back_do() {
        assert!(MemberError::from(TensorEngineError::OutOfRange { gate: 3 }).falls_back());
        assert!(!MemberError::from(TensorEngineError::Probe { a: 1, b: 2 }).falls_back());
        assert!(!MemberError::from(TensorEngineError::Inexact { product: 0, case: 0 }).falls_back());
    }

    #[test]
    fn a_python_that_is_not_the_driver_is_refused_with_its_words() {
        // `true` reads its arguments, prints nothing and exits: no driver line.
        let Ok(t) = which("true") else { return };
        let err = start(&t).err().expect("refused");
        assert!(err.contains("did not start the kernel's driver"), "{err}");
    }

    fn which(name: &str) -> Result<PathBuf, ()> {
        for d in std::env::var_os("PATH").map(|p| std::env::split_paths(&p).collect::<Vec<_>>()).unwrap_or_default() {
            let c = d.join(name);
            if c.is_file() {
                return Ok(c);
            }
        }
        Err(())
    }
}