wai-quantum 0.6.0

//! What names the npu-cores member, and the line layout it sends: the
//! kernel one compute core of a neural processor's tile array runs
//! ([`KERNEL`], `npu_cores.cc`), the program around it ([`DESIGN`], built by
//! the array's open toolchain), the host side ([`HOST`]), the implementation
//! id, a device's identity and the mutants a conforming device must be caught
//! running. No device is touched here, so all of it is testable on any host;
//! the member itself is `quantum_fabric_npu_cores`.
//!
//! **Why it is bit-exact.** The kernel is the same transcription of
//! [`lanes::exact_row`](super::lanes) as the other members': each operand a
//! 32-bit value, each `fx(a, b) = (a·b + 2^29) >> 30` an exact 64-bit
//! product, an add and an arithmetic shift, each row the sum of four. The core
//! writes the rows whole; the host checks each row's range and keeps its low
//! 32 bits, as [`lanes::apply_words`](super::lanes::apply_words) does. The
//! member's probe checks the core's products and shifts against
//! [`lanes::mul_wide`](super::lanes::mul_wide) and
//! [`lanes::fxmul_wide`](super::lanes::fxmul_wide) before the device is a
//! member, and its certificates check whole runs against the CPU's pinned
//! corpus digests.

/// The kernel, exactly as it is compiled for the core.
pub const KERNEL: &str = include_str!("npu_cores.cc");

/// The program around the kernel, run by the toolchain's Python to build it.
pub const DESIGN: &str = include_str!("npu_cores_design.py");

/// The host side, compiled against the device runtime and loaded.
pub const HOST: &str = include_str!("npu_cores_host.cpp");

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

/// Cases per line.
pub const LINE: usize = 256;

/// Lines per dispatch.
pub const LINES: usize = 4;

/// Cases per dispatch.
pub const CASES: usize = LINE * LINES;

/// Words before a line's cases: word 0 the mode, words 8..15 the gate's
/// entries.
pub const HEADER: usize = 16;

/// Words per case in, and out (four 64-bit values).
pub const CASE_IN: usize = 4;
pub const CASE_OUT: usize = 8;

/// The modes a line's header names.
pub const MODE_ROWS: i32 = 0;
pub const MODE_PROBE: i32 = 1;

/// Version of the host layout: the line format above, the gather of a gate's
/// pairs, the 64-bit rows 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. Each is its
/// own compiled program.
#[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) -> (i64, u32, 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),
        }
    }

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

fn put(o: &mut Vec<u8>, bytes: &[u8]) {
    o.extend_from_slice(&(bytes.len() as u64).to_be_bytes());
    o.extend_from_slice(bytes);
}

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

/// `"npu-cores-lanes64/<16 hex>"`: the first 8 bytes of
/// BLAKE3(`"wai:quantum-fabric-npu-cores\x01"` ‖ put([`KERNEL`]) ‖
/// put([`DESIGN`]) ‖ put([`HOST`]) ‖ the variant's constants ‖ [`LINE`] ‖
/// [`LINES`] ‖ [`HOST_LAYOUT`]), the numbers big-endian. The toolchain that
/// builds the program is the device's and is in the member id instead.
pub fn implementation_id_of(v: Variant) -> String {
    let (rnd, sh, last_add) = v.constants();
    let mut o = b"wai:quantum-fabric-npu-cores\x01".to_vec();
    put(&mut o, KERNEL.as_bytes());
    put(&mut o, DESIGN.as_bytes());
    put(&mut o, HOST.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(&(LINE as u64).to_be_bytes());
    o.extend_from_slice(&(LINES 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)
}

/// The toolchain's name for a device's tile-array generation, by its PCI
/// device id, for the ones this member's program is written for.
pub fn generation(device: u16) -> Option<&'static str> {
    match device {
        // The first generation, and the second (512-bit cores, eight columns).
        0x1502 => Some("npu1"),
        0x17f0 => Some("npu2"),
        _ => None,
    }
}

/// What an npu-cores member is, as its id hashes it. The member fills it from
/// the device and the toolchain; the fields are plain values here so that the
/// id is testable without a device.
#[derive(Clone, Debug, PartialEq, Eq)]
#[non_exhaustive]
pub struct DeviceIdentity {
    /// PCI vendor and device ids, and the revision.
    pub vendor: u16,
    pub device: u16,
    pub revision: u8,
    /// The device's name as its runtime gives it.
    pub name: String,
    /// The device's firmware version, as its driver reports it.
    pub firmware: String,
    /// The toolchain that built the program, as it reports itself (one line
    /// per package).
    pub toolchain: String,
}

impl DeviceIdentity {
    /// `"local/npu-cores-<d12>"`: the first 12 hex digits of
    /// BLAKE3(`"wai:quantum-fabric-npu-cores-member\x01"` ‖ vendor ‖ device ‖
    /// revision ‖ put(name) ‖ put(firmware) ‖ put(toolchain)), the numbers
    /// big-endian.
    pub fn member_id(&self) -> String {
        let mut o = b"wai:quantum-fabric-npu-cores-member\x01".to_vec();
        o.extend_from_slice(&self.vendor.to_be_bytes());
        o.extend_from_slice(&self.device.to_be_bytes());
        o.push(self.revision);
        put(&mut o, self.name.as_bytes());
        put(&mut o, self.firmware.as_bytes());
        put(&mut o, self.toolchain.as_bytes());
        format!("local/npu-cores-{}", hex(&blake3::hash(&o).as_bytes()[..6]))
    }

    /// The key a runner declares in CI: `"npu-cores-<vendor>-<device>"`, the
    /// PCI ids in hex.
    pub fn key(&self) -> String {
        format!("npu-cores-{:04x}-{:04x}", self.vendor, self.device)
    }
}

/// One dispatch's input: [`LINES`] lines, each a header naming `mode` and
/// `m`, then [`LINE`] cases from `cases` (zeros past its end). At most
/// [`CASES`] cases.
pub fn pack(mode: i32, m: &[i32; 8], cases: &[[i32; 4]]) -> Vec<i32> {
    assert!(cases.len() <= CASES, "{} cases; a dispatch holds {CASES}", cases.len());
    let mut v = Vec::with_capacity(LINES * (HEADER + CASE_IN * LINE));
    for l in 0..LINES {
        let mut head = [0i32; HEADER];
        head[0] = mode;
        head[8..16].copy_from_slice(m);
        v.extend_from_slice(&head);
        for k in 0..LINE {
            v.extend_from_slice(&cases.get(l * LINE + k).copied().unwrap_or([0; 4]));
        }
    }
    v
}

/// The words one dispatch's output holds.
pub const OUT_WORDS: usize = CASES * CASE_OUT;

/// A dispatch's output as each case's four 64-bit values.
pub fn unpack(out: &[i32]) -> Vec<[i64; 4]> {
    assert_eq!(out.len(), OUT_WORDS);
    out.as_chunks::<8>()
        .0
        .iter()
        .map(|w| {
            let v = |k: usize| (u64::from(w[2 * k] as u32) | (u64::from(w[2 * k + 1] as u32) << 32)) as i64;
            [v(0), v(1), v(2), v(3)]
        })
        .collect()
}

#[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("npu-cores-lanes64/") && implementation_id().len() == "npu-cores-lanes64/".len() + 16);
        for v in Variant::MUTANTS {
            assert_ne!(v.constants(), Variant::Conforming.constants());
        }
        assert!(KERNEL.contains("extern \"C\" void wq_line(") && DESIGN.contains("\"wq_line\"") && HOST.contains("wq_npu_run"));
    }

    #[test]
    fn every_field_of_the_identity_is_in_the_member_id() {
        let base = DeviceIdentity { vendor: 0x1022, device: 0x17f0, revision: 0x11, name: "a device".into(), firmware: "1.1.2.65".into(), toolchain: "a 1\nb 2".into() };
        let id = base.member_id();
        assert!(id.starts_with("local/npu-cores-") && id.len() == "local/npu-cores-".len() + 12, "{id}");
        assert!(id.as_str() > "local/cpu-x86_64", "sorts after the CPU reference");
        for v in [
            DeviceIdentity { vendor: 0x1023, ..base.clone() },
            DeviceIdentity { device: 0x1502, ..base.clone() },
            DeviceIdentity { revision: 0x10, ..base.clone() },
            DeviceIdentity { name: "another".into(), ..base.clone() },
            DeviceIdentity { firmware: "1.1.2.66".into(), ..base.clone() },
            DeviceIdentity { toolchain: "a 1\nb 3".into(), ..base.clone() },
        ] {
            assert_ne!(v.member_id(), id, "{v:?}");
        }
        assert_eq!(base.key(), "npu-cores-1022-17f0");
        assert_eq!((generation(0x17f0), generation(0x1502), generation(0x1234)), (Some("npu2"), Some("npu1"), None));
    }

    #[test]
    fn a_dispatch_is_lines_of_a_header_and_cases() {
        let m = [1, 2, 3, 4, 5, 6, 7, 8];
        let cases: Vec<[i32; 4]> = (0..LINE as i32 + 3).map(|k| [k, -k, 2 * k, -2 * k]).collect();
        let v = pack(MODE_ROWS, &m, &cases);
        assert_eq!(v.len(), LINES * (HEADER + CASE_IN * LINE));
        let stride = HEADER + CASE_IN * LINE;
        for l in 0..LINES {
            assert_eq!(v[l * stride], MODE_ROWS);
            assert_eq!(&v[l * stride + 8..l * stride + 16], &m);
        }
        // The second line starts with case LINE; past the cases, zeros.
        assert_eq!(&v[stride + HEADER..stride + HEADER + 4], &[LINE as i32, -(LINE as i32), 2 * LINE as i32, -2 * LINE as i32]);
        assert_eq!(&v[stride + HEADER + 4 * 3..stride + HEADER + 4 * 4], &[0, 0, 0, 0]);
        assert_eq!(pack(MODE_PROBE, &m, &[])[0], MODE_PROBE);
        let mut out = vec![0i32; OUT_WORDS];
        out[0] = -1;
        out[1] = 0x7fff_ffff;
        out[2] = 5;
        let u = unpack(&out);
        assert_eq!(u.len(), CASES);
        assert_eq!(u[0], [i64::MAX, 5, 0, 0]);
    }
}