wai-quantum 0.6.0

//! The fabric's arithmetic on a neural processor, and what names that member:
//! the limb-product model, the layout of a batch, the exact assembly of each
//! product, and the implementation id. Nothing here touches a device, so all of
//! it is testable on any host; the member itself is `quantum_fabric_npu`.
//!
//! **Why the products are exact.** A neural processor computes in fp16, and does
//! not document how it accumulates. So each 32-bit operand is split into eight
//! 4-bit limbs (the top one signed, `−8..=7`, the rest `0..=15`), and the
//! processor computes, for each of the 15 shifts `s`, the limb sum
//! `Σ_{i+j=s} a_i·b_j`: the coefficients of the product of two polynomials in
//! 16. Every limb product is at most 225 in size and every sum of at most eight
//! of them at most 1800, below 2048, so every value the processor forms is an
//! integer fp16 holds exactly, whatever order it adds in. The host then
//! assembles `a·b = Σ_s z_s·16^s` exactly in 64 bits, and the gate's arithmetic
//! (`fx`, the row sums, the 32-bit lane and the range report) is
//! [`lanes::apply_words`](super::lanes::apply_words)'s.
//!
//! The model is one multiply and one sum over the last axis
//! ([`limb_product_model`]); the layout puts operand `a`'s limbs in `x[p, s, j]`
//! and `b`'s limb `s − j` in `y[p, s, j]` (zero outside), so
//! `z[p, s] = Σ_j x[p, s, j]·y[p, s, j]` is the coefficient `s`.
//!
//! **What the processor is credited with.** The member exists only where the
//! runtime's compute plan places every operation of the model on the neural
//! processor; a run the runtime would place on the CPU is refused, not counted.

use super::FabricClass;

/// Products per batch: the model's fixed first dimension.
pub const BATCH: usize = 16384;
/// Limbs per operand.
pub const LIMBS: usize = 8;
/// Coefficients of a product: `2·LIMBS − 1`.
pub const SHIFTS: usize = 2 * LIMBS - 1;

/// Version of the host-side layout: the limb split, the `x`/`y` expansion, the
/// assembly, and the batch. Bumped on any change; it is part of the
/// implementation id.
pub const HOST_LAYOUT: u32 = 1;

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

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

/// The class a neural-processor member has.
pub const CLASS: FabricClass = FabricClass::Npu;

/// Protobuf wire types.
const VARINT: u32 = 0;
const LEN: u32 = 2;

fn varint(o: &mut Vec<u8>, mut v: u64) {
    loop {
        let b = (v & 0x7f) as u8;
        v >>= 7;
        if v == 0 {
            o.push(b);
            return;
        }
        o.push(b | 0x80);
    }
}

fn key(o: &mut Vec<u8>, field: u32, wire: u32) {
    varint(o, u64::from(field << 3 | wire));
}

fn uint(o: &mut Vec<u8>, field: u32, v: u64) {
    key(o, field, VARINT);
    varint(o, v);
}

/// A signed int64 field (two's complement, as protobuf's int64 is written).
fn int(o: &mut Vec<u8>, field: u32, v: i64) {
    key(o, field, VARINT);
    varint(o, v as u64);
}

fn bytes(o: &mut Vec<u8>, field: u32, b: &[u8]) {
    key(o, field, LEN);
    varint(o, b.len() as u64);
    o.extend_from_slice(b);
}

fn msg(o: &mut Vec<u8>, field: u32, f: impl FnOnce(&mut Vec<u8>)) {
    let mut inner = Vec::new();
    f(&mut inner);
    bytes(o, field, &inner);
}

// The model description's field numbers (its Model, FeatureTypes and MIL schemas).
const ARRAY_FLOAT16: u64 = 65552;
/// The operation set the program is written in.
const OPSET: &[u8] = b"CoreML8"; // allow-name: the runtime's operation-set identifier
const MIL_BOOL: u64 = 1;
const MIL_STRING: u64 = 2;
const MIL_FLOAT16: u64 = 10;
const MIL_INT32: u64 = 23;

/// A MIL tensor type of `dtype` and `shape` (rank 0 for a scalar).
fn tensor_type(o: &mut Vec<u8>, dtype: u64, shape: &[u64]) {
    // ValueType.tensorType = 1
    msg(o, 1, |t| {
        uint(t, 1, dtype); // dataType
        if !shape.is_empty() {
            int(t, 2, shape.len() as i64); // rank
        }
        for &d in shape {
            // dimensions = 3 { constant = 1 { size = 1 } }
            msg(t, 3, |dim| msg(dim, 1, |c| uint(c, 1, d)));
        }
    });
}

/// A feature description: `name`, a multi-array of fp16 and `shape`.
fn feature(o: &mut Vec<u8>, field: u32, name: &str, shape: &[u64]) {
    msg(o, field, |f| {
        bytes(f, 1, name.as_bytes());
        // type = 3 { multiArrayType = 5 { shape = 1 (repeated), dataType = 2 } }
        msg(f, 3, |t| {
            msg(t, 5, |a| {
                for &d in shape {
                    int(a, 1, d as i64);
                }
                uint(a, 2, ARRAY_FLOAT16);
            })
        });
    });
}

/// A function input `name`: a MIL fp16 tensor of `shape`.
fn named_input(o: &mut Vec<u8>, name: &str, shape: &[u64]) {
    // Function.inputs = 1 { name = 1, type = 2 }
    msg(o, 1, |n| {
        bytes(n, 1, name.as_bytes());
        msg(n, 2, |v| tensor_type(v, MIL_FLOAT16, shape));
    });
}

/// An operation's `name` attribute.
fn name_attribute(o: &mut Vec<u8>, name: &str) {
    // attributes = 5 (map<string, Value>) { key = 1, value = 2 }
    msg(o, 5, |e| {
        bytes(e, 1, b"name");
        msg(e, 2, |v| {
            msg(v, 2, |t| tensor_type(t, MIL_STRING, &[])); // Value.type
            // immediateValue = 3 { tensor = 1 { strings = 4 { values = 1 } } }
            msg(v, 3, |iv| msg(iv, 1, |tv| msg(tv, 4, |s| bytes(s, 1, name.as_bytes()))));
        });
    });
}

/// An operation input binding: parameter `param` reads variable `var`.
fn input_binding(o: &mut Vec<u8>, param: &str, var: &str) {
    // inputs = 2 (map<string, Argument>) { key, value = Argument { arguments = 1 { name = 1 } } }
    msg(o, 2, |e| {
        bytes(e, 1, param.as_bytes());
        msg(e, 2, |arg| msg(arg, 1, |b| bytes(b, 1, var.as_bytes())));
    });
}

/// An operation output `name` of `dtype` and `shape`.
fn op_output(o: &mut Vec<u8>, name: &str, dtype: u64, shape: &[u64]) {
    // outputs = 3 { name = 1, type = 2 }
    msg(o, 3, |n| {
        bytes(n, 1, name.as_bytes());
        msg(n, 2, |v| tensor_type(v, dtype, shape));
    });
}

/// A `const` operation `name` holding `value` (written by `write`, a
/// TensorValue) of `dtype` and `shape`.
fn const_op(o: &mut Vec<u8>, name: &str, dtype: u64, shape: &[u64], write: impl FnOnce(&mut Vec<u8>)) {
    // Block.operations = 3
    msg(o, 3, |op| {
        bytes(op, 1, b"const");
        op_output(op, name, dtype, shape);
        msg(op, 5, |e| {
            bytes(e, 1, b"val");
            msg(e, 2, |v| {
                msg(v, 2, |t| tensor_type(t, dtype, shape));
                msg(v, 3, |iv| msg(iv, 1, write));
            });
        });
        name_attribute(op, name);
    });
}

/// The model: `z[b, s] = Σ_j x[b, s, j] · y[b, s, j]`, all fp16, for
/// `batch` products of `shifts` limb sums of `limbs` terms.
pub fn limb_product_model(batch: u64, shifts: u64, limbs: u64) -> Vec<u8> {
    let full = [batch, shifts, limbs];
    let out = [batch, shifts];
    let mut m = Vec::new();
    uint(&mut m, 1, 9); // specificationVersion
    msg(&mut m, 2, |d| {
        feature(d, 1, "x", &full);
        feature(d, 1, "y", &full);
        feature(d, 10, "z", &out);
    });
    // mlProgram = 502 { version = 1, functions = 2 (map<string, Function>) }
    msg(&mut m, 502, |p| {
        int(p, 1, 1);
        msg(p, 2, |e| {
            bytes(e, 1, b"main");
            msg(e, 2, |f| {
                named_input(f, "x", &full);
                named_input(f, "y", &full);
                bytes(f, 2, OPSET); // opset
                // block_specializations = 3 (map<string, Block>)
                msg(f, 3, |e| {
                    bytes(e, 1, OPSET);
                    msg(e, 2, |b| {
                        bytes(b, 2, b"z"); // Block.outputs
                        // mul
                        msg(b, 3, |op| {
                            bytes(op, 1, b"mul");
                            input_binding(op, "x", "x");
                            input_binding(op, "y", "y");
                            op_output(op, "mul_0", MIL_FLOAT16, &full);
                            name_attribute(op, "mul_0");
                        });
                        const_op(b, "z_axes_0", MIL_INT32, &[1], |tv| msg(tv, 2, |ints| int(ints, 1, -1)));
                        const_op(b, "z_keep_dims_0", MIL_BOOL, &[], |tv| msg(tv, 3, |bools| uint(bools, 1, 0)));
                        // reduce_sum
                        msg(b, 3, |op| {
                            bytes(op, 1, b"reduce_sum");
                            input_binding(op, "x", "mul_0");
                            input_binding(op, "axes", "z_axes_0");
                            input_binding(op, "keep_dims", "z_keep_dims_0");
                            op_output(op, "z", MIL_FLOAT16, &out);
                            name_attribute(op, "z");
                        });
                    });
                });
            });
        });
    });
    m
}


/// The element type of a portable model's tensors.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub enum Elem {
    /// fp16: every value the model forms is an integer of at most 1800 in size,
    /// which fp16 holds exactly.
    F16,
    /// int32: exact on any processor that computes it in integers without
    /// wrapping; a coefficient is at most 1800 in size.
    I32,
}

// ONNX's field numbers (onnx.proto) and element types.
const ONNX_FLOAT16: u64 = 10;
const ONNX_INT32: u64 = 6;
const ONNX_INT64: u64 = 7;
const ONNX_ATTR_INT: u64 = 2;

/// An ONNX value description: `name`, a tensor of `elem` and `shape`.
fn onnx_value(o: &mut Vec<u8>, field: u32, name: &str, elem: u64, shape: &[u64]) {
    msg(o, field, |v| {
        bytes(v, 1, name.as_bytes());
        // type = 2 { tensor_type = 1 { elem_type = 1, shape = 2 { dim = 1 { dim_value = 1 } } } }
        msg(v, 2, |t| {
            msg(t, 1, |tt| {
                uint(tt, 1, elem);
                msg(tt, 2, |sh| {
                    for &d in shape {
                        msg(sh, 1, |dim| int(dim, 1, d as i64));
                    }
                });
            })
        });
    });
}

/// The same program in the portable model format (ONNX, opset 17):
/// `z = ReduceSum(Mul(x, y), axes = [-1], keepdims = 0)`, `x` and `y`
/// `[batch, shifts, limbs]`, all of `elem`. Neural-processor runtimes that
/// read this format take it as it is; the inputs are [`expand`]'s (as fp16
/// bits) or [`expand_i32`]'s.
pub fn onnx_limb_model(batch: u64, shifts: u64, limbs: u64, elem: Elem) -> Vec<u8> {
    let e = match elem {
        Elem::F16 => ONNX_FLOAT16,
        Elem::I32 => ONNX_INT32,
    };
    let full = [batch, shifts, limbs];
    let out = [batch, shifts];
    let mut m = Vec::new();
    int(&mut m, 1, 8); // ir_version
    bytes(&mut m, 2, b"wai-quantum"); // producer_name
    // graph = 7
    msg(&mut m, 7, |g| {
        // node = 1: Mul
        msg(g, 1, |n| {
            bytes(n, 1, b"x");
            bytes(n, 1, b"y");
            bytes(n, 2, b"p");
            bytes(n, 3, b"mul");
            bytes(n, 4, b"Mul");
        });
        // node = 1: ReduceSum over the last axis, the axes as an input
        msg(g, 1, |n| {
            bytes(n, 1, b"p");
            bytes(n, 1, b"axes");
            bytes(n, 2, b"z");
            bytes(n, 3, b"sum");
            bytes(n, 4, b"ReduceSum");
            msg(n, 5, |a| {
                bytes(a, 1, b"keepdims");
                int(a, 3, 0);
                uint(a, 20, ONNX_ATTR_INT);
            });
        });
        bytes(g, 2, b"limb_products");
        // initializer = 5: axes = [-1], int64
        msg(g, 5, |t| {
            int(t, 1, 1); // dims
            uint(t, 2, ONNX_INT64); // data_type
            int(t, 7, -1); // int64_data
            bytes(t, 8, b"axes"); // name
        });
        onnx_value(g, 11, "x", e, &full);
        onnx_value(g, 11, "y", e, &full);
        onnx_value(g, 12, "z", e, &out);
    });
    // opset_import = 8 { domain = 1 "", version = 2 }
    msg(&mut m, 8, |op| {
        bytes(op, 1, b"");
        int(op, 2, 17);
    });
    m
}

/// [`expand`] with int32 values: `x[p, s, j]` and `y[p, s, j]` as the limbs
/// themselves.
pub fn expand_i32(pairs: &[(i32, i32)], x: &mut [i32], y: &mut [i32]) {
    assert!(pairs.len() <= BATCH && x.len() == BATCH * SHIFTS * LIMBS && y.len() == x.len());
    x.fill(0);
    y.fill(0);
    for (p, &(a, b)) in pairs.iter().enumerate() {
        let (la, lb) = (limbs(a), limbs(b));
        for s in 0..SHIFTS {
            for j in 0..LIMBS {
                if s >= j && s - j < LIMBS {
                    let at = (p * SHIFTS + s) * LIMBS + j;
                    x[at] = la[j];
                    y[at] = lb[s - j];
                }
            }
        }
    }
}

/// `"npu-limbs4/<16 hex>"` over the model's bytes and the host layout.
pub fn implementation_id_of(model: &[u8]) -> String {
    let mut o = DOMAIN_IMPLEMENTATION.to_vec();
    o.extend_from_slice(&(model.len() as u64).to_be_bytes());
    o.extend_from_slice(model);
    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}")
}

/// The model this member runs.
pub fn model() -> Vec<u8> {
    limb_product_model(BATCH as u64, SHIFTS as u64, LIMBS as u64)
}

pub fn implementation_id() -> String {
    implementation_id_of(&model())
}

/// The eight 4-bit limbs of `v`, least significant first, the top one signed.
pub fn limbs(v: i32) -> [i32; LIMBS] {
    let u = v as u32;
    let mut out = [0i32; LIMBS];
    for (i, l) in out.iter_mut().enumerate() {
        *l = ((u >> (4 * i)) & 15) as i32;
    }
    if out[LIMBS - 1] >= 8 {
        out[LIMBS - 1] -= 16;
    }
    out
}

/// The fp16 bits of an integer `v` with `|v| ≤ 2048`, exactly.
pub fn f16_of(v: i32) -> u16 {
    debug_assert!(v.abs() <= 2048);
    if v == 0 {
        return 0;
    }
    let sign: u16 = if v < 0 { 0x8000 } else { 0 };
    let m = v.unsigned_abs();
    let e = 31 - m.leading_zeros(); // floor(log2 m), 0..=11
    // Significand bits below the leading one, left-aligned in 10 bits (2048,
    // the one value with e = 11, has none).
    let below = m - (1 << e);
    let frac = if e <= 10 { below << (10 - e) } else { below >> (e - 10) } as u16;
    sign | (((e + 15) as u16) << 10) | frac
}

/// The integer an fp16 value holds, or `None` if it is not an integer of at most
/// 2048 in size.
pub fn int_of_f16(bits: u16) -> Option<i64> {
    let sign = if bits & 0x8000 != 0 { -1 } else { 1 };
    let e = i32::from((bits >> 10) & 0x1f);
    let frac = i64::from(bits & 0x3ff);
    if e == 0 {
        return (frac == 0).then_some(0);
    }
    if e == 31 {
        return None;
    }
    let p = e - 15; // the exponent
    if !(0..=11).contains(&p) {
        return None;
    }
    let m = (1024 + frac) << p; // value · 2^10
    (m & 1023 == 0).then_some(sign * (m >> 10))
}

/// Writes the inputs of one batch for `pairs` (at most [`BATCH`]; the rest is
/// zero): `x[p, s, j] = limb j of a`, `y[p, s, j] = limb s − j of b`, as fp16
/// bits, row-major.
pub fn expand(pairs: &[(i32, i32)], x: &mut [u16], y: &mut [u16]) {
    assert!(pairs.len() <= BATCH && x.len() == BATCH * SHIFTS * LIMBS && y.len() == x.len());
    x.fill(0);
    y.fill(0);
    for (p, &(a, b)) in pairs.iter().enumerate() {
        let (la, lb) = (limbs(a), limbs(b));
        for s in 0..SHIFTS {
            for j in 0..LIMBS {
                if s >= j && s - j < LIMBS {
                    let at = (p * SHIFTS + s) * LIMBS + j;
                    x[at] = f16_of(la[j]);
                    y[at] = f16_of(lb[s - j]);
                }
            }
        }
    }
}

/// The product `a·b` from its 15 coefficients: `Σ_s z_s·16^s`, exact in 64 bits.
pub fn assemble(z: &[i64]) -> i64 {
    z.iter().rev().fold(0i64, |acc, &c| acc.wrapping_mul(16).wrapping_add(c))
}

/// The model's arithmetic in integers, on the host: what a processor that
/// computes it exactly returns. For the tests, and for nothing else.
pub fn replay(x: &[u16], y: &[u16]) -> Vec<i64> {
    x.chunks(LIMBS).zip(y.chunks(LIMBS)).map(|(a, b)| a.iter().zip(b).map(|(&p, &q)| int_of_f16(p).unwrap_or(0) * int_of_f16(q).unwrap_or(0)).sum()).collect()
}

/// The model with one operation's type replaced: for the mutation tests.
pub fn mutate_model(op_from: &str, op_to: &str) -> Vec<u8> {
    let m = model();
    let (from, to) = (op_from.as_bytes(), op_to.as_bytes());
    assert_eq!(from.len(), to.len(), "a mutant keeps the operation name's length");
    let at = m.windows(from.len()).position(|w| w == from).expect("the operation is in the model");
    let mut out = m;
    out[at..at + to.len()].copy_from_slice(to);
    out
}

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

    #[test]
    fn small_integers_round_trip_through_fp16_exactly() {
        for v in -2048..=2048 {
            assert_eq!(int_of_f16(f16_of(v)), Some(i64::from(v)), "{v}");
        }
        assert_eq!(f16_of(1), 0x3c00);
        assert_eq!(f16_of(-2), 0xc000);
        assert_eq!(f16_of(1800), 0x6708);
        assert_eq!(int_of_f16(0x3800), None, "one half is no integer");
        assert_eq!(int_of_f16(0x7c00), None, "infinity is no integer");
    }

    #[test]
    fn limbs_reassemble_and_every_value_stays_exact_in_fp16() {
        for v in [0, 1, -1, 7, -8, i32::MAX, i32::MIN, 0x7654_3210, -0x0f0f_0f0f] {
            let l = limbs(v);
            assert_eq!(l.iter().rev().fold(0i64, |a, &x| a * 16 + i64::from(x)), i64::from(v), "{v}");
            assert!(l[..LIMBS - 1].iter().all(|x| (0..=15).contains(x)) && (-8..=7).contains(&l[LIMBS - 1]));
        }
        // The largest coefficient a batch can form.
        assert!(LIMBS as i64 * 15 * 15 <= 2048);
    }

    #[test]
    fn a_batch_replayed_assembles_every_product_exactly() {
        let mut pairs = lanes::probe_pairs(4096, 0x11);
        pairs.truncate(BATCH);
        let (mut x, mut y) = (vec![0u16; BATCH * SHIFTS * LIMBS], vec![0u16; BATCH * SHIFTS * LIMBS]);
        expand(&pairs, &mut x, &mut y);
        let z = replay(&x, &y);
        assert!(z.iter().all(|c| c.abs() <= 1800));
        for (p, &(a, b)) in pairs.iter().enumerate() {
            let got = assemble(&z[p * SHIFTS..(p + 1) * SHIFTS]);
            assert_eq!(got, i64::from(a) * i64::from(b), "({a}, {b})");
        }
    }

    #[test]
    fn the_portable_model_names_its_program_and_its_element_type() {
        for (elem, code) in [(Elem::F16, ONNX_FLOAT16), (Elem::I32, ONNX_INT32)] {
            let m = onnx_limb_model(BATCH as u64, SHIFTS as u64, LIMBS as u64, elem);
            assert_eq!(&m[..2], &[0x08, 8], "ir_version first");
            for op in [&b"Mul"[..], b"ReduceSum", b"keepdims", b"axes"] {
                assert!(m.windows(op.len()).any(|w| w == op));
            }
            // elem_type = code, inside each value description.
            assert!(m.windows(2).any(|w| w == [0x08, code as u8]));
        }
        assert_ne!(onnx_limb_model(4, 15, 8, Elem::F16), onnx_limb_model(4, 15, 8, Elem::I32));
        // The int32 layout carries the limbs themselves.
        let pairs = [(i32::MIN, -1), (0x7654_3210, 0x0123_4567)];
        let n = BATCH * SHIFTS * LIMBS;
        let (mut x, mut y) = (vec![0i32; n], vec![0i32; n]);
        expand_i32(&pairs, &mut x, &mut y);
        for (p, &(a, b)) in pairs.iter().enumerate() {
            let z: Vec<i64> = (0..SHIFTS).map(|s| (0..LIMBS).map(|j| i64::from(x[(p * SHIFTS + s) * LIMBS + j]) * i64::from(y[(p * SHIFTS + s) * LIMBS + j])).sum()).collect();
            assert_eq!(assemble(&z), i64::from(a) * i64::from(b));
        }
    }

    #[test]
    fn the_model_is_the_program_it_names_and_its_mutants_differ() {
        let m = model();
        assert_eq!(m[0], 0x08, "field 1, specificationVersion, comes first");
        assert_eq!(m[1], 9);
        for op in [&b"mul"[..], b"reduce_sum", OPSET] {
            assert!(m.windows(op.len()).any(|w| w == op));
        }
        let id = implementation_id();
        assert!(id.starts_with("npu-limbs4/") && id.len() == "npu-limbs4/".len() + 16);
        assert_ne!(implementation_id_of(&mutate_model("mul", "add")), id);
        assert_ne!(implementation_id_of(&mutate_model("reduce_sum", "reduce_max")), id);
    }
}