wai-quantum 0.6.0

//! The gate as a portable tensor program (StableHLO, as MLIR text): the input
//! of compilers for tensor processors, GPUs and CPUs. Every value is a 64-bit
//! integer that never overflows: `|e·a| < 2^62` and a row of four
//! `fx = (e·a + 2^29) >> 30` stays below 2^35, so each compiler must give the
//! exact rows whatever it does with them.
//!
//! - `rows(e0, …, e7, a0r, a0i, a1r, a1i)`: the gate's entries as scalars, a
//!   batch of `n` pairs' columns as 32-bit integers; the four rows, 64-bit.
//! - `probe(a, b)`: each pair's `a·b` and `fx(a, b)`, 64-bit.
//!
//! [`npy`] writes inputs and expected outputs in the array format such tools
//! read.

use std::fmt::Write as _;

/// The kernel's constants: the rounding constant, the shift, and whether row
/// 0's last real term is added.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Constants {
    pub round: i64,
    pub shift: i64,
    pub last_add: bool,
}

pub const CONFORMING: Constants = Constants { round: 1 << 29, shift: 30, last_add: false };

struct Body {
    text: String,
    next: usize,
    n: usize,
}

impl Body {
    fn ty(&self) -> String {
        format!("tensor<{}xi64>", self.n)
    }
    fn op(&mut self, rhs: String) -> String {
        let v = format!("%v{}", self.next);
        self.next += 1;
        let _ = writeln!(self.text, "    {v} = {rhs}");
        v
    }
    fn constant(&mut self, c: i64) -> String {
        let t = self.ty();
        self.op(format!("stablehlo.constant dense<{c}> : {t}"))
    }
    fn binary(&mut self, name: &str, a: &str, b: &str) -> String {
        let t = self.ty();
        self.op(format!("stablehlo.{name} {a}, {b} : {t}"))
    }
    fn widen(&mut self, a: &str) -> String {
        let (n, t) = (self.n, self.ty());
        self.op(format!("stablehlo.convert {a} : (tensor<{n}xi32>) -> {t}"))
    }
    fn splat(&mut self, scalar: &str) -> String {
        let t = self.ty();
        self.op(format!("stablehlo.broadcast_in_dim {scalar}, dims = [] : (tensor<i64>) -> {t}"))
    }
    /// `(e·a + round) >> shift`.
    fn fx(&mut self, e: &str, a: &str, round: &str, shift: &str) -> String {
        let p = self.binary("multiply", e, a);
        let r = self.binary("add", &p, round);
        self.binary("shift_right_arithmetic", &r, shift)
    }
}

/// The module: `rows` and `probe` over batches of `n` pairs.
pub fn module(n: usize, c: Constants) -> String {
    let t = format!("tensor<{n}xi64>");
    let ti = format!("tensor<{n}xi32>");
    let mut out = String::new();
    let _ = writeln!(out, "// The gate's rows and the arithmetic probe, exact in 64-bit integers.");
    let _ = writeln!(out, "module @wai_quantum_gate {{");
    // rows
    let mut b = Body { text: String::new(), next: 0, n };
    let round = b.constant(c.round);
    let shift = b.constant(c.shift);
    let e: Vec<String> = (0..8).map(|k| b.splat(&format!("%e{k}"))).collect();
    let a: Vec<String> = ["%a0r", "%a0i", "%a1r", "%a1i"].iter().map(|x| b.widen(x)).collect();
    let last = if c.last_add { "add" } else { "subtract" };
    let terms: [[(usize, usize, &str); 4]; 4] = [
        [(0, 0, "add"), (1, 1, "subtract"), (2, 2, "add"), (3, 3, last)],
        [(0, 1, "add"), (1, 0, "add"), (2, 3, "add"), (3, 2, "add")],
        [(4, 0, "add"), (5, 1, "subtract"), (6, 2, "add"), (7, 3, "subtract")],
        [(4, 1, "add"), (5, 0, "add"), (6, 3, "add"), (7, 2, "add")],
    ];
    let mut rows = Vec::new();
    for row in terms {
        let mut acc = b.fx(&e[row[0].0], &a[row[0].1], &round, &shift);
        for &(k, op, how) in &row[1..] {
            let f = b.fx(&e[k], &a[op], &round, &shift);
            acc = b.binary(how, &acc, &f);
        }
        rows.push(acc);
    }
    let params = (0..8).map(|k| format!("%e{k}: tensor<i64>")).chain(["%a0r", "%a0i", "%a1r", "%a1i"].iter().map(|x| format!("{x}: {ti}"))).collect::<Vec<_>>().join(", ");
    let _ = writeln!(out, "  func.func @rows({params}) -> ({t}, {t}, {t}, {t}) {{");
    out.push_str(&b.text);
    let _ = writeln!(out, "    return {} : {t}, {t}, {t}, {t}", rows.join(", "));
    let _ = writeln!(out, "  }}");
    // probe
    let mut b = Body { text: String::new(), next: 0, n };
    let round = b.constant(c.round);
    let shift = b.constant(c.shift);
    let (wa, wb) = (b.widen("%a"), b.widen("%b"));
    let prod = b.binary("multiply", &wa, &wb);
    let fx = b.fx(&wb, &wa, &round, &shift);
    let _ = writeln!(out, "  func.func @probe(%a: {ti}, %b: {ti}) -> ({t}, {t}) {{");
    out.push_str(&b.text);
    let _ = writeln!(out, "    return {prod}, {fx} : {t}, {t}");
    let _ = writeln!(out, "  }}");
    let _ = writeln!(out, "}}");
    out
}

/// The gate's products on a matrix unit, as a second module: one
/// `dot_general` of the gate's limb weights (`16 × 112`, bf16) and a batch's
/// limbs (`16 × n`, bf16), accumulated and returned in f32. The layout and the
/// argument that every coefficient is exact on any unit that multiplies bf16
/// and accumulates in f32 are [`tensor_engine`](super::tensor_engine)'s; a
/// compiler that maps this to a matrix unit must give the exact coefficients.
pub fn limb_module(n: usize) -> String {
    use super::tensor_engine::{K, M};
    let (w, a, c) = (format!("tensor<{K}x{M}xbf16>"), format!("tensor<{K}x{n}xbf16>"), format!("tensor<{M}x{n}xf32>"));
    format!(
        "// The gate's products as limb coefficients: one bf16 matrix multiply, f32 results.\n\
         module @wai_quantum_limbs {{\n\
         \x20 func.func @coefficients(%w: {w}, %a: {a}) -> {c} {{\n\
         \x20   %c = stablehlo.dot_general %w, %a, contracting_dims = [0] x [0] : ({w}, {a}) -> {c}\n\
         \x20   return %c : {c}\n\
         \x20 }}\n\
         }}\n"
    )
}

/// The exact coefficients of [`limb_module`] for limbs `w` (`16 × 112`) and
/// `a` (`16 × n`), given as the integers the limbs are (not as bf16 holds
/// them): `112 × n`, row-major.
pub fn reference_coefficients(w: &[i32], a: &[i32], n: usize) -> Vec<f32> {
    use super::tensor_engine::{K, M};
    let mut c = vec![0f32; M * n];
    for col in 0..M {
        for q in 0..n {
            let v: i64 = (0..K).map(|k| i64::from(w[k * M + col]) * i64::from(a[k * n + q])).sum();
            c[col * n + q] = v as f32;
        }
    }
    c
}

/// The rows the module must give, from the CPU's arithmetic.
pub fn reference_rows(m: &[i64; 8], cols: &[Vec<i32>; 4], c: Constants) -> [Vec<i64>; 4] {
    let fx = |e: i64, a: i32| (e * i64::from(a) + c.round) >> c.shift;
    let last = if c.last_add { 1 } else { -1 };
    let mut out: [Vec<i64>; 4] = Default::default();
    for (((&a0r, &a0i), &a1r), &a1i) in cols[0].iter().zip(&cols[1]).zip(&cols[2]).zip(&cols[3]) {
        out[0].push(fx(m[0], a0r) - fx(m[1], a0i) + fx(m[2], a1r) + last * fx(m[3], a1i));
        out[1].push(fx(m[0], a0i) + fx(m[1], a0r) + fx(m[2], a1i) + fx(m[3], a1r));
        out[2].push(fx(m[4], a0r) - fx(m[5], a0i) + fx(m[6], a1r) - fx(m[7], a1i));
        out[3].push(fx(m[4], a0i) + fx(m[5], a0r) + fx(m[6], a1i) + fx(m[7], a1r));
    }
    out
}

/// A little-endian array in the `.npy` format (version 1.0): `descr` `<i4`
/// or `<i8`, `shape` empty for a scalar.
pub fn npy(descr: &str, shape: &[usize], data: &[u8]) -> Vec<u8> {
    let dims = match shape {
        [] => "()".to_string(),
        [n] => format!("({n},)"),
        _ => format!("({})", shape.iter().map(|d| d.to_string()).collect::<Vec<_>>().join(", ")),
    };
    let mut header = format!("{{'descr': '{descr}', 'fortran_order': False, 'shape': {dims}, }}");
    // Magic (6) + version (2) + length (2) + header, padded to 64 bytes, ending in '\n'.
    while (10 + header.len() + 1) % 64 != 0 {
        header.push(' ');
    }
    header.push('\n');
    let mut out = b"\x93NUMPY\x01\x00".to_vec();
    out.extend_from_slice(&(header.len() as u16).to_le_bytes());
    out.extend_from_slice(header.as_bytes());
    out.extend_from_slice(data);
    out
}

pub fn npy_i32(v: &[i32]) -> Vec<u8> {
    npy("<i4", &[v.len()], &v.iter().flat_map(|x| x.to_le_bytes()).collect::<Vec<u8>>())
}

pub fn npy_i64(v: &[i64]) -> Vec<u8> {
    npy("<i8", &[v.len()], &v.iter().flat_map(|x| x.to_le_bytes()).collect::<Vec<u8>>())
}

pub fn npy_scalar_i64(v: i64) -> Vec<u8> {
    npy("<i8", &[], &v.to_le_bytes())
}

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

    #[test]
    fn the_module_has_both_functions_and_its_mutants_differ() {
        let m = module(64, CONFORMING);
        assert!(m.contains("func.func @rows(") && m.contains("func.func @probe("));
        assert_eq!(m.matches("stablehlo.shift_right_arithmetic").count(), 16 + 1);
        let l = limb_module(512);
        assert!(l.contains("func.func @coefficients(%w: tensor<16x112xbf16>, %a: tensor<16x512xbf16>) -> tensor<112x512xf32>"), "{l}");
        assert!(l.contains("stablehlo.dot_general %w, %a, contracting_dims = [0] x [0]"));
        for c in [Constants { round: 1 << 28, ..CONFORMING }, Constants { shift: 29, ..CONFORMING }, Constants { last_add: true, ..CONFORMING }] {
            assert_ne!(module(64, c), m);
        }
    }

    #[test]
    fn npy_headers_are_aligned_and_reference_rows_are_the_gate() {
        for v in [npy_i32(&[1, 2, 3]), npy_i64(&[7; 512]), npy_scalar_i64(-5)] {
            let hl = u16::from_le_bytes([v[8], v[9]]) as usize;
            assert_eq!((10 + hl) % 64, 0);
            assert_eq!(v[10 + hl - 1], b'\n');
        }
        // One pair, the identity gate (1.0 = 2^30): the rows are the inputs.
        let one = 1i64 << 30;
        let m = [one, 0, 0, 0, 0, 0, one, 0];
        let r = reference_rows(&m, &[vec![5], vec![-6], vec![7], vec![-8]], CONFORMING);
        assert_eq!(r.map(|x| x[0]), [5, -6, 7, -8]);
    }
}