open-standards/wai/wai-ops-npu/tests/probe.rs

//! The device's arithmetic, measured against §2 before any graph: which of the unit's
//! rounding modes is which rounding of §2 and at which shifts, what its saturation
//! does, how wide a 32-bit accumulator lane is, and that the kernel's requantisation,
//! as every operator forms it, is exact at every shift and rounding.
//!
//! With no device each test prints `ops-npu <probe>: SKIPPED`, and fails under
//! `WAI_OPS_NPU_REQUIRE`. Each result is a line `ops-npu <probe>: key=…`, appended to
//! `WAI_OPS_NPU_REPORT` when it is set.

#[path = "support/corpus.rs"]
mod corpus;
#[path = "support/rq.rs"]
mod rq;

use corpus::report;
use wai_ops::graph::Round;
use wai_ops_npu::lines::{Epi, Init, MAX_N, Width, mac_line, probe_line};
use wai_ops_npu::probe::{ROUNDS, Rng, SAT_NONE, SAT_SATURATE, UNIT_MODES, cases, round_name, sat32, spec};

/// Shifts written as ranges: `0..59, 61`.
fn ranges(v: &[u32]) -> String {
    let mut out: Vec<String> = Vec::new();
    let mut k = 0;
    while k < v.len() {
        let mut j = k;
        while j + 1 < v.len() && v[j + 1] == v[j] + 1 {
            j += 1;
        }
        out.push(if j == k { v[k].to_string() } else { format!("{}..{}", v[k], v[j]) });
        k = j + 1;
    }
    if out.is_empty() { "none".into() } else { out.join(", ") }
}

#[test]
fn the_requantisation_is_exact_at_every_shift_and_rounding() {
    let Some(npu) = corpus::device("requantisation probe") else { return };
    let (cases, bad, first) = rq::run_counted(&npu, 4, 0x7271_7072_6f62);
    report(&format!("ops-npu requantisation: key={} {cases} cases over shifts 0..63 and {:?}, {bad} differ{}", npu.key(), ROUNDS.map(round_name), if bad > 0 { format!("; the first: {first}") } else { String::new() }));
    assert!(cases > 500_000, "too few cases: {cases}");
    assert_eq!(bad, 0, "the kernel's requantisation differs from §2: {first}");
}

#[test]
fn the_units_rounding_modes_against_section_2() {
    let Some(npu) = corpus::device("rounding probe") else { return };
    let mut rng = Rng::new(0x726f_756e_6473);
    // The same cases for every mode: per shift, MAX_N of them.
    let per_shift: Vec<Vec<(i32, i32)>> = (0..=63u32).map(|s| cases(s, MAX_N, &mut rng)).collect();
    let mut lines = Vec::new();
    for &(mode, _) in &UNIT_MODES {
        for s in 0..=63u32 {
            let (t, m): (Vec<i32>, Vec<i32>) = per_shift[s as usize].iter().copied().unzip();
            lines.push(probe_line(mode, SAT_SATURATE, s as i32, &t, &m));
        }
    }
    let outs = npu.dispatch(&lines).expect("the probe's dispatches");
    let mut exact_modes = Vec::new();
    for (mi, &(mode, name)) in UNIT_MODES.iter().enumerate() {
        let mut desc = Vec::new();
        for r in ROUNDS {
            // Shifts at which the mode equals R_r on every case whose result fits 32 bits.
            let mut eq = Vec::new();
            for s in 0..=63u32 {
                let o = &outs[mi * 64 + s as usize];
                let ok = per_shift[s as usize].iter().enumerate().all(|(i, &(t, m))| {
                    let q = spec(t, m, s, r);
                    i32::try_from(q).is_err() || o[i] as i128 == q
                });
                if ok {
                    eq.push(s);
                }
            }
            if !eq.is_empty() {
                desc.push(format!("{} at shifts {}", round_name(r), ranges(&eq)));
            }
            if eq.len() == 64 || eq == (0..=59).collect::<Vec<u32>>() {
                exact_modes.push((name, r, eq.len()));
            }
        }
        // What the unit gives at the shifts it does not take: a left shift by 64 − s?
        let left: Vec<u32> = (60..=63u32)
            .filter(|&s| {
                let o = &outs[mi * 64 + s as usize];
                per_shift[s as usize].iter().enumerate().all(|(i, &(t, m))| o[i] == sat32((t as i128 * m as i128) << (64 - s)))
            })
            .collect();
        report(&format!(
            "ops-npu unit rounding: key={} mode {name} ({mode}): {}; at shifts {} it gives t·m shifted left by 64 − s, saturated",
            npu.key(),
            if desc.is_empty() { "equals no rounding of §2 at any shift but 0".into() } else { desc.join("; ") },
            ranges(&left)
        ));
    }
    // The three modes the kernel uses are exact at shifts 0..59 and at no larger shift.
    for (name, r) in [("floor", Round::Floor), ("positive_inf", Round::HalfUp), ("conv_even", Round::HalfEven)] {
        assert!(exact_modes.iter().any(|&(n, rr, len)| n == name && rr == r && len == 60), "{name} is not {} at exactly the shifts 0..59: {exact_modes:?}", round_name(r));
    }
    // And no mode is exact at shifts 60..63, which is why the kernel takes them elsewhere.
    assert!(exact_modes.iter().all(|&(_, _, len)| len < 64), "a mode is exact at every shift: {exact_modes:?}");
}

#[test]
fn saturation_and_the_accumulator_lanes() {
    let Some(npu) = corpus::device("lane probe") else { return };
    // q past 32 bits: shift 0 of products past i32, under each saturation mode.
    let mut rng = Rng::new(0x7361_7475_7261);
    let (t, m): (Vec<i32>, Vec<i32>) = (0..MAX_N).map(|_| (rng.i32() | 0x10000, rng.i32() | 0x10000)).unzip();
    let outs = npu.dispatch(&[probe_line(0, SAT_SATURATE, 0, &t, &m), probe_line(0, SAT_NONE, 0, &t, &m)]).unwrap();
    let past = |i: usize| i32::try_from(t[i] as i64 * m[i] as i64).is_err();
    let n_past = (0..MAX_N).filter(|&i| past(i)).count();
    let sat_ok = (0..MAX_N).all(|i| outs[0][i] == sat32(t[i] as i128 * m[i] as i128));
    let wrap_ok = (0..MAX_N).all(|i| outs[1][i] == (t[i] as i64 * m[i] as i64) as i32);
    report(&format!("ops-npu saturation: key={} {n_past} of {MAX_N} products past 32 bits: the saturating mode clamps them ({sat_ok}); the non-saturating mode keeps their low 32 bits ({wrap_ok})", npu.key()));
    assert!(n_past > 100 && sat_ok);

    // Sums past 32 bits in a 32-bit accumulator lane (8- and 16-bit operands), and past
    // 64 bits in a 64-bit lane (32-bit operands): each lane starts at `init` and adds k
    // equal products; what comes back is the lane as the unit narrows it to 32 bits.
    for (width, v, init) in [(Width::W8, -128i64, i32::MAX - 10), (Width::W16, -32768, 0), (Width::W32, i32::MIN as i64, 0)] {
        let mut res = Vec::new();
        for k in [1usize, 2, 3, 4] {
            let w = vec![v; k];
            let x = vec![vec![v; 64]; k];
            let o = &npu.dispatch(&[mac_line(width, &w, &x, &Init::Scalar(init), false, Epi::IDENTITY).unwrap()]).unwrap()[0];
            let exact = init as i128 + k as i128 * v as i128 * v as i128;
            let kind = if o[0] as i128 == exact {
                "exact"
            } else if o[0] == exact as i64 as i32 {
                "its low 32 bits"
            } else if o[0] == sat32(exact) {
                "saturated to 32 bits"
            } else if exact != exact as i64 as i128 && o[0] == sat32(exact as i64 as i128) {
                "wrapped at 64 bits, then saturated"
            } else {
                "other"
            };
            res.push(format!("{init} + {k}×{}: {} ({kind})", v * v, o[0]));
        }
        report(&format!("ops-npu accumulator: key={} {}: {}", npu.key(), width.as_str(), res.join("; ")));
    }
}