open-standards/wai/wai-ops-npu/tests/support/corpus.rs

//! The graph vectors of `ops-conformance/` (`vectors/` and `graphs/`), run with their
//! arithmetic on a device and held to the committed expectations and to the CPU
//! interpreter; and the device tests' shared reporting.

#![allow(dead_code)]

use std::io::Write;
use std::path::PathBuf;

use wai_ops::json::{self, Value, unhex};
use wai_ops::vectors::read_tensor;
use wai_ops::{Loaded, Profile, Sequence, Tensor};
use wai_ops_npu::{Npu, Plan, Stats, discover, plan};

pub fn conformance() -> PathBuf {
    PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../ops-conformance")
}

/// A report line: printed, and appended to `WAI_OPS_NPU_REPORT` when it is set.
pub fn report(line: &str) {
    println!("{line}");
    if let Ok(p) = std::env::var("WAI_OPS_NPU_REPORT") {
        let mut f = std::fs::OpenOptions::new().create(true).append(true).open(&p).unwrap_or_else(|e| panic!("{p}: {e}"));
        writeln!(f, "{line}").unwrap();
    }
}

/// The device this host has, held to `WAI_OPS_NPU_EXPECT` (its key, or `none`); or
/// `None` with SKIPPED reported, which fails under `WAI_OPS_NPU_REQUIRE`.
pub fn device(test: &str) -> Option<Npu> {
    let d = discover();
    for r in &d.refused {
        report(&format!("ops-npu device refused: key={} {} reason={}", r.key, r.device, r.error));
    }
    let mut keys: Vec<String> = d.found.iter().map(|n| n.key()).collect();
    keys.sort();
    if let Ok(want) = std::env::var("WAI_OPS_NPU_EXPECT") {
        let got = if keys.is_empty() { "none".to_string() } else { keys.join(",") };
        assert_eq!(got, want, "the devices found are not the ones WAI_OPS_NPU_EXPECT declares");
        report(&format!("ops-npu devices: as declared ({got})"));
    }
    let mut found = d.found;
    if found.is_empty() {
        let why = d.unavailable.clone().unwrap_or_else(|| format!("{} device(s) refused, above", d.refused.len()));
        report(&format!("ops-npu {test}: SKIPPED: no device ran it: {why}"));
        assert!(std::env::var_os("WAI_OPS_NPU_REQUIRE").is_none(), "WAI_OPS_NPU_REQUIRE is set and no device ran the {test}");
        return None;
    }
    Some(found.remove(0))
}

/// Every graph-vector file: `vectors/*.jsonl`, then `graphs/*.jsonl`, each sorted by name.
pub fn files() -> Vec<(String, String)> {
    let mut out = Vec::new();
    for dir in ["vectors", "graphs"] {
        let d = conformance().join(dir);
        let mut names: Vec<String> = std::fs::read_dir(&d)
            .unwrap_or_else(|e| panic!("{}: {e}", d.display()))
            .filter_map(|e| e.ok().map(|e| e.file_name().to_string_lossy().into_owned()))
            .filter(|n| n.ends_with(".jsonl"))
            .collect();
        names.sort();
        for n in names {
            let text = std::fs::read_to_string(d.join(&n)).unwrap();
            out.push((format!("{dir}/{n}"), text));
        }
    }
    assert!(out.len() > 30, "the corpus was not found");
    out
}

/// A vector's expectation for one decode: the outputs, or a refusal code.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum Expect {
    Outputs(Vec<Tensor>),
    Refuse(String),
}

pub struct Decode {
    pub inputs: Vec<Tensor>,
    pub scalars: Vec<i64>,
    pub expect: Expect,
}

pub struct Vector {
    pub file: String,
    pub id: String,
    pub graph: Vec<u8>,
    pub load: String,
    pub decodes: Vec<Decode>,
}

pub fn vectors() -> Vec<Vector> {
    files().iter().flat_map(|(file, text)| parse(file, text)).collect()
}

pub fn parse(file: &str, text: &str) -> Vec<Vector> {
    let mut out = Vec::new();
    for line in text.split_terminator('\n') {
        let v = json::parse(line.as_bytes()).expect("a canonical line");
        let id = v.get("id").and_then(Value::as_str).unwrap().to_string();
        let graph = unhex(v.get("graph").and_then(Value::as_str).unwrap()).unwrap();
        let load = v.get("load").and_then(Value::as_str).unwrap().to_string();
        let decodes = v
            .get("decodes")
            .and_then(Value::as_arr)
            .unwrap()
            .iter()
            .map(|d| {
                let inputs = d.get("inputs").and_then(Value::as_arr).unwrap().iter().map(|t| read_tensor(t).unwrap()).collect();
                let scalars = d.get("scalars").and_then(Value::as_arr).unwrap().iter().map(|x| x.as_dec_i64().expect("an i64")).collect();
                let e = d.get("expect").unwrap();
                let expect = match e.get("refuse") {
                    Some(c) => Expect::Refuse(c.as_str().unwrap().to_string()),
                    None => Expect::Outputs(e.get("outputs").and_then(Value::as_arr).unwrap().iter().map(|t| read_tensor(t).unwrap()).collect()),
                };
                Decode { inputs, scalars, expect }
            })
            .collect();
        out.push(Vector { file: file.to_string(), id, graph, load, decodes });
    }
    out
}

/// A loaded vector and its plan.
pub struct Ready<'a> {
    pub v: &'a Vector,
    pub l: Loaded,
    pub plan: Plan,
}

/// Every vector that loads, with its plan; and the count refused at load.
pub fn ready(vs: &[Vector]) -> (Vec<Ready<'_>>, u32) {
    let mut out = Vec::new();
    let mut refused = 0;
    for v in vs {
        match wai_ops::graph::load(&v.graph, Profile::Acc64) {
            Ok(l) => {
                assert_eq!(v.load, "ok", "{}", v.id);
                let p = plan(&l).unwrap_or_else(|e| panic!("{}: {e}", v.id));
                out.push(Ready { v, l, plan: p });
            }
            Err(r) => {
                assert_eq!(r.code.as_str(), v.load, "{}: the load the vector records", v.id);
                refused += 1;
            }
        }
    }
    (out, refused)
}

fn got_of(r: Result<Vec<Tensor>, wai_ops::Refusal>) -> Expect {
    match r {
        Ok(t) => Expect::Outputs(t),
        Err(r) => Expect::Refuse(r.code.as_str().to_string()),
    }
}

/// What one device did with a set of vectors.
#[derive(Clone, Debug, Default)]
pub struct Tally {
    pub vectors: u32,
    pub decodes: u32,
    pub refusals: u32,
    /// Vectors whose graph is in the subset but moves data only: nothing for the core.
    pub movement_only: u32,
    pub stats: Stats,
    /// Vectors with a decode whose result differs, and the first `limit` of them.
    pub mismatched: u32,
    pub mismatches: Vec<String>,
}

/// Run every decode of `rs` with its arithmetic on `npu`, against the committed
/// expectation and the interpreter, stopping a vector at its first mismatch. `limit`
/// caps the mismatches recorded.
pub fn run_all(npu: &Npu, rs: &[&Ready<'_>], limit: usize) -> Tally {
    let mut t = Tally::default();
    for r in rs {
        t.vectors += 1;
        let mut seq = Sequence::new();
        let mut cpu = Sequence::new();
        let before = t.stats.npu_ops;
        for (k, d) in r.v.decodes.iter().enumerate() {
            let want = got_of(wai_ops::interp::run(&r.l, &d.inputs, &d.scalars, &mut cpu));
            assert_eq!(want, d.expect, "{}: decode {k}: the CPU interpreter's result is the committed one", r.v.id);
            t.decodes += 1;
            if matches!(d.expect, Expect::Refuse(_)) {
                t.refusals += 1;
            }
            let (res, st) = wai_ops_npu::run(npu, &r.l, &r.plan, &d.inputs, &d.scalars, &mut seq);
            let erred = !st.errors.is_empty();
            t.stats.add(&st);
            if erred {
                break;
            }
            let got = got_of(res);
            if got != d.expect {
                t.mismatched += 1;
                if t.mismatches.len() < limit {
                    t.mismatches.push(format!("{} ({}) decode {k}: {}", r.v.id, r.v.file, diff(&d.expect, &got)));
                }
                break;
            }
        }
        if t.stats.npu_ops == before {
            t.movement_only += 1;
        }
    }
    t
}

pub fn diff(want: &Expect, got: &Expect) -> String {
    match (want, got) {
        (Expect::Outputs(a), Expect::Outputs(b)) => {
            for (k, (x, y)) in a.iter().zip(b).enumerate() {
                if (x.c, x.h, x.w) != (y.c, y.h, y.w) {
                    return format!("output {k} shape {:?} vs {:?}", (x.c, x.h, x.w), (y.c, y.h, y.w));
                }
                if let Some(i) = (0..x.data.len()).find(|&i| x.data[i] != y.data[i]) {
                    return format!("output {k} element {i}: want {} got {}", x.data[i], y.data[i]);
                }
            }
            "output count".into()
        }
        _ => format!("want {want:?} got {got:?}").chars().take(300).collect(),
    }
}