wai-quantum 0.6.0

//! What names the tensor-engine member, and the arithmetic it relies on: the
//! gate's products as exact limb convolutions on a tensor processor's matrix
//! engine, which multiplies bf16 values and accumulates in fp32. No device is
//! touched here, so all of it is testable on any host; the member itself is
//! `quantum_fabric_tensor_engine`, and the kernel is [`HELPER`].
//!
//! **Why it is bit-exact on any such engine.** Each 32-bit value is split into
//! four limbs of [`Variant::bits`] (8) bits, the lower three in `[0, 255]` and
//! the top one signed in `[-128, 127]`; bf16 holds every one of them exactly.
//! A product `e·a` is `Σ c_s·2^(8s)` with `c_s = Σ_{i+j=s} e_i·a_j`, seven
//! coefficients of at most four limb products each, so every product the
//! engine forms is an integer at most `2^16` in size and every coefficient
//! one below `2^18`: fp32 holds each partial sum exactly, in whatever order the
//! engine adds them. The weights put each of a gate's sixteen products per pair
//! (four terms in each of four rows) in its own seven columns, so one matrix
//! multiply forms them all. The host checks every coefficient is an integer
//! below `2^24`, assembles each product exactly in 64 bits, and applies the
//! fixed-point arithmetic of [`lanes::apply_words`](super::lanes::apply_words):
//! `fx = (product + 2^29) >> 30`, each row the sum of four, the range checked
//! and the low 32 bits kept.
//!
//! So the engine does every multiplication and the host none; the host does
//! the carries, the rounding and the sums, as the neural-processor member's
//! host does.

/// The kernel and its driver, run by the member with the matrix engine's
/// toolchain's Python.
pub const HELPER: &str = include_str!("tensor_engine_helper.py");

/// The implementation name the id starts with.
pub const IMPLEMENTATION: &str = "tensor-engine-limbs8";

/// Pairs per matrix multiply: the moving matrix's columns.
pub const TILE: usize = 512;

/// The contraction: four components of a pair, four limbs each.
pub const K: usize = 16;

/// The outputs per pair: sixteen products, seven coefficients each.
pub const PRODUCTS: usize = 16;
pub const COEFFICIENTS: usize = 7;
pub const M: usize = PRODUCTS * COEFFICIENTS;

/// The largest coefficient magnitude fp32 holds exactly.
pub const EXACT: f32 = 16_777_216.0;

/// The sixteen products, as `(entry, component)`: rows 0r, 0i, 1r and 1i,
/// four terms each, in the order [`lanes::exact_row`](super::lanes) adds them.
pub const TERMS: [(usize, usize); PRODUCTS] =
    [(0, 0), (1, 1), (2, 2), (3, 3), (0, 1), (1, 0), (2, 3), (3, 2), (4, 0), (5, 1), (6, 2), (7, 3), (4, 1), (5, 0), (6, 3), (7, 2)];

/// Version of the host layout: the weights' and the limbs' layout, the
/// assembly and the host's range check. Part of the implementation id.
pub const HOST_LAYOUT: u32 = 1;

/// 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,
    /// Limbs of nine bits, which bf16 cannot hold: the engine rounds them,
    /// and the products are wrong. The exactness argument rests on the limb
    /// size, and this mutant shows the engine is where it would fail.
    Radix9,
}

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

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

    /// Bits per limb.
    pub fn bits(self) -> u32 {
        if self == Variant::Radix9 { 9 } else { 8 }
    }

    pub fn label(self) -> &'static str {
        match self {
            Variant::Conforming => "conforming",
            Variant::M1 => "m1",
            Variant::M2 => "m2",
            Variant::M3 => "m3",
            Variant::Radix9 => "radix9",
        }
    }
}

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

/// `"tensor-engine-limbs8/<16 hex>"` over the kernel's text, the variant's
/// constants and limb size, the tile and the host layout.
pub fn implementation_id_of(v: Variant) -> String {
    let (rnd, sh, last_add) = v.constants();
    let mut o = b"wai:quantum-fabric-tensor-engine\x01".to_vec();
    o.extend_from_slice(&(HELPER.len() as u64).to_be_bytes());
    o.extend_from_slice(HELPER.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(&v.bits().to_be_bytes());
    o.extend_from_slice(&(TILE as u64).to_be_bytes());
    o.extend_from_slice(&HOST_LAYOUT.to_be_bytes());
    format!("{IMPLEMENTATION}/{}", hex(&blake3::hash(&o).as_bytes()[..8]))
}

pub fn implementation_id() -> String {
    implementation_id_of(Variant::Conforming)
}

/// `x`'s four limbs of `bits` bits: the lower three unsigned, the top one the
/// signed rest, so that `x = Σ limb_i·2^(bits·i)`.
pub fn limbs(x: i32, bits: u32) -> [i32; 4] {
    let mask = (1i64 << bits) - 1;
    let v = i64::from(x);
    [(v & mask) as i32, ((v >> bits) & mask) as i32, ((v >> (2 * bits)) & mask) as i32, (v >> (3 * bits)) as i32]
}

/// `v` as bf16 bits, rounded to nearest with ties to even, as a conversion to
/// the engine's input type rounds. Exact for every integer of at most 256 in
/// size.
pub fn bf16(v: i32) -> u16 {
    let bits = (v as f32).to_bits();
    let round = 0x7fff + ((bits >> 16) & 1);
    ((bits + round) >> 16) as u16
}

/// The value bf16 bits stand for.
pub fn from_bf16(b: u16) -> f32 {
    f32::from_bits(u32::from(b) << 16)
}

/// The stationary matrix for gate entries `m`: `K × M` bf16 bits, row-major.
/// Row `c·4 + j` is limb `j` of component `c`; column `p·7 + s` coefficient
/// `s` of product `p`, which takes limb `s − j` of its entry where it exists.
pub fn weights(m: &[i32; 8], bits: u32) -> Vec<u16> {
    let mut w = vec![0u16; K * M];
    let el: Vec<[i32; 4]> = m.iter().map(|&e| limbs(e, bits)).collect();
    for (p, &(k, c)) in TERMS.iter().enumerate() {
        for j in 0..4 {
            for (i, &l) in el[k].iter().enumerate() {
                w[(c * 4 + j) * M + p * COEFFICIENTS + i + j] = bf16(l);
            }
        }
    }
    w
}

/// Pairs per probe request: one per product slot.
pub const PROBE_PAIRS: usize = PRODUCTS;

/// The stationary matrix for the arithmetic probe: product `p`'s weights are
/// the limbs of `bs[p]`, taken against component `p % 4`. With the probe's
/// moving matrix ([`probe_moving`]), product `q` of row `q / 4` is
/// `a_q · bs[q]`: sixteen independent pairs in one multiply.
pub fn probe_weights(bs: &[i32; PROBE_PAIRS], bits: u32) -> Vec<u16> {
    let mut w = vec![0u16; K * M];
    for (p, &b) in bs.iter().enumerate() {
        let c = p % 4;
        for j in 0..4 {
            for (i, &l) in limbs(b, bits).iter().enumerate() {
                w[(c * 4 + j) * M + p * COEFFICIENTS + i + j] = bf16(l);
            }
        }
    }
    w
}

/// The probe's cases: pair `q`'s `a` at row `q / 4`, component `q % 4`.
pub fn probe_moving(a: &[i32; PROBE_PAIRS]) -> Vec<[i32; 4]> {
    (0..PROBE_PAIRS / 4).map(|r| [a[4 * r], a[4 * r + 1], a[4 * r + 2], a[4 * r + 3]]).collect()
}

/// The moving matrix for `cases` (each a pair's four components): `K × n`
/// bf16 bits, row-major, `n` the cases rounded up to a whole number of
/// [`TILE`]s, the rest zeros.
pub fn moving(cases: &[[i32; 4]], bits: u32) -> (Vec<u16>, usize) {
    let n = cases.len().div_ceil(TILE).max(1) * TILE;
    let mut a = vec![0u16; K * n];
    for (q, case) in cases.iter().enumerate() {
        for (c, &x) in case.iter().enumerate() {
            for (j, &l) in limbs(x, bits).iter().enumerate() {
                a[(c * 4 + j) * n + q] = bf16(l);
            }
        }
    }
    (a, n)
}

/// Why the engine's coefficients are not a product's.
#[derive(Debug, PartialEq, Eq)]
pub struct Inexact {
    pub product: usize,
    pub case: usize,
}

/// The sixteen products of each of the first `count` cases, from the engine's
/// coefficients (`M × n` f32, row-major): each coefficient must be an integer
/// below [`EXACT`] in size, and each product is `Σ c_s·2^(bits·s)`, exact in
/// 64 bits.
pub fn assemble(coefficients: &[f32], n: usize, count: usize, bits: u32) -> Result<Vec<[i64; PRODUCTS]>, Inexact> {
    assert_eq!(coefficients.len(), M * n);
    let mut out = Vec::with_capacity(count);
    for q in 0..count {
        let mut ps = [0i64; PRODUCTS];
        for (p, slot) in ps.iter_mut().enumerate() {
            let mut v = 0i64;
            for s in 0..COEFFICIENTS {
                let c = coefficients[(p * COEFFICIENTS + s) * n + q];
                if c.is_nan() || c.abs() >= EXACT || c.fract() != 0.0 {
                    return Err(Inexact { product: p, case: q });
                }
                v += (c as i64) << (bits * s as u32);
            }
            *slot = v;
        }
        out.push(ps);
    }
    Ok(out)
}

/// A pair's four rows from its sixteen products, with the variant's
/// constants.
pub fn rows(p: &[i64; PRODUCTS], rnd: i64, sh: u32, last_add: bool) -> [i64; 4] {
    let fx = |v: i64| (v + rnd) >> sh;
    let last = if last_add { fx(p[3]) } else { -fx(p[3]) };
    [
        fx(p[0]) - fx(p[1]) + fx(p[2]) + last,
        fx(p[4]) + fx(p[5]) + fx(p[6]) + fx(p[7]),
        fx(p[8]) - fx(p[9]) + fx(p[10]) - fx(p[11]),
        fx(p[12]) + fx(p[13]) + fx(p[14]) + fx(p[15]),
    ]
}

/// What a tensor-engine member is, as its id hashes it.
#[derive(Clone, Debug, PartialEq, Eq)]
#[non_exhaustive]
pub struct DeviceIdentity {
    /// Where the kernel runs, as the driver reports it: `"simulator"` or a
    /// device's name.
    pub target: String,
    /// The driver's report of its packages and versions.
    pub versions: String,
}

impl DeviceIdentity {
    pub fn new(target: String, versions: String) -> DeviceIdentity {
        DeviceIdentity { target, versions }
    }

    /// `"local/tensor-engine-<d12>"` over the target and the versions.
    pub fn member_id(&self) -> String {
        let mut o = b"wai:quantum-fabric-tensor-engine-member\x01".to_vec();
        o.extend_from_slice(&(self.target.len() as u64).to_be_bytes());
        o.extend_from_slice(self.target.as_bytes());
        o.extend_from_slice(&(self.versions.len() as u64).to_be_bytes());
        o.extend_from_slice(self.versions.as_bytes());
        format!("local/tensor-engine-{}", hex(&blake3::hash(&o).as_bytes()[..6]))
    }

    /// The key a runner declares: `"tensor-engine-<target>"`.
    pub fn key(&self) -> String {
        format!("tensor-engine-{}", self.target)
    }
}

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

    /// The engine's arithmetic, as it is: bf16 inputs, products and sums in
    /// f32, in the given order.
    fn engine(w: &[u16], a: &[u16], n: usize) -> Vec<f32> {
        let mut c = vec![0f32; M * n];
        for col in 0..M {
            for q in 0..n {
                let mut acc = 0f32;
                for k in 0..K {
                    acc += from_bf16(w[k * M + col]) * from_bf16(a[k * n + q]);
                }
                c[col * n + q] = acc;
            }
        }
        c
    }

    #[test]
    fn limbs_rebuild_every_value_and_bf16_holds_them() {
        for x in [0, 1, -1, i32::MAX, i32::MIN, 0x1234_5678, -0x0765_4321, 1 << 29, -(1 << 29) - 1] {
            for bits in [8, 9] {
                let l = limbs(x, bits);
                let back: i64 = l.iter().enumerate().map(|(i, &v)| i64::from(v) << (bits * i as u32)).sum();
                assert_eq!(back, i64::from(x), "{x} in {bits}-bit limbs");
            }
            for v in limbs(x, 8) {
                assert_eq!(from_bf16(bf16(v)), v as f32, "{v} in bf16");
            }
        }
        // Nine-bit limbs are past what bf16 holds.
        assert_ne!(from_bf16(bf16(257)), 257.0);
        assert_eq!(from_bf16(bf16(256)), 256.0);
    }

    #[test]
    fn the_engines_arithmetic_gives_every_product_exactly() {
        let mut s = 0x7e57u64;
        let mut next = || {
            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)) as u32 as i32
        };
        let m: [i32; 8] = [i32::MAX, i32::MIN, 1 << 30, -(1 << 30), 0x1234_5678, next(), next(), -1];
        let mut cases: Vec<[i32; 4]> = (0..300).map(|_| [next(), next(), next(), next()]).collect();
        cases[0] = [i32::MIN, i32::MAX, 0, -1];
        let (a, n) = moving(&cases, 8);
        assert_eq!(n, TILE);
        let c = engine(&weights(&m, 8), &a, n);
        let ps = assemble(&c, n, cases.len(), 8).expect("exact");
        for (q, case) in cases.iter().enumerate() {
            for (p, &(k, comp)) in TERMS.iter().enumerate() {
                assert_eq!(ps[q][p], i64::from(m[k]) * i64::from(case[comp]), "case {q} product {p}");
            }
        }
        // The probe's layout: sixteen independent pairs in one multiply.
        let bs: [i32; PROBE_PAIRS] = std::array::from_fn(|k| if k == 0 { i32::MIN } else { next() });
        let az: [i32; PROBE_PAIRS] = std::array::from_fn(|k| if k == 0 { i32::MIN } else { next() });
        let (a, n) = moving(&probe_moving(&az), 8);
        let ps = assemble(&engine(&probe_weights(&bs, 8), &a, n), n, PROBE_PAIRS / 4, 8).unwrap();
        for q in 0..PROBE_PAIRS {
            assert_eq!(ps[q / 4][q], i64::from(az[q]) * i64::from(bs[q]), "probe pair {q}");
        }
        // The rows are the gate's.
        let one = 1i32 << 30;
        let id = [one, 0, 0, 0, 0, 0, one, 0];
        let (a, n) = moving(&[[5, -6, 7, -8]], 8);
        let ps = assemble(&engine(&weights(&id, 8), &a, n), n, 1, 8).unwrap();
        assert_eq!(rows(&ps[0], 1 << 29, 30, false), [5, -6, 7, -8]);
    }

    #[test]
    fn nine_bit_limbs_come_back_wrong_from_the_engine() {
        let m = [0x1234_5ff7, 0, 0, 0, 0, 0, 0, 0];
        let cases = [[0x7fff_ffff, 0, 0, 0]];
        let (a, n) = moving(&cases, 9);
        let got = match assemble(&engine(&weights(&m, 9), &a, n), n, 1, 9) {
            Ok(ps) => ps[0][0],
            Err(_) => return,
        };
        assert_ne!(got, i64::from(m[0]) * i64::from(cases[0][0]));
    }

    #[test]
    fn ids_follow_the_variant_and_the_identity() {
        let ids: std::collections::BTreeSet<String> = [Variant::Conforming, Variant::M1, Variant::M2, Variant::M3, Variant::Radix9].into_iter().map(implementation_id_of).collect();
        assert_eq!(ids.len(), 5);
        let a = DeviceIdentity::new("simulator".into(), "x 1".into());
        assert_eq!(a.key(), "tensor-engine-simulator");
        assert!(a.member_id().starts_with("local/tensor-engine-") && a.member_id().as_str() > "local/cpu-x86_64");
        assert_ne!(a.member_id(), DeviceIdentity::new("simulator".into(), "x 2".into()).member_id());
        assert_ne!(a.member_id(), DeviceIdentity::new("device".into(), "x 1".into()).member_id());
        assert!(HELPER.contains("nc_matmul") && HELPER.contains("wq-tensor-engine 1"));
    }
}