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

//! The graph vectors of `ops-conformance/` (`vectors/` and `graphs/`), run on a device and
//! held to the committed expectations and to the CPU interpreter.

#![allow(dead_code)]

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_gpu::{Gpu, GpuError, GpuSequence, Width, blocking, support};

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

/// 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>,
}

fn dec(v: &Value) -> i64 {
    v.as_dec_i64().expect("an i64")
}

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

/// The vectors of one file.
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(dec).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
}

/// What one adapter did with the corpus.
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct Tally {
    /// Vectors run, and their decodes, at each width.
    pub vectors32: u32,
    pub decodes32: u32,
    pub vectors64: u32,
    pub decodes64: u32,
    /// The same `acc64` vectors run again with narrowing off: every (transposed)
    /// convolution on the 64-bit kernels.
    pub vectors64w: u32,
    pub decodes64w: u32,
    /// Operators the `acc64` runs gave the narrow kernels (`prepare::narrow_plan`), and
    /// the vectors holding one.
    pub narrow_ops: u32,
    pub narrow_vectors: u32,
    /// Of those decodes, the ones whose expectation is a refusal.
    pub refusals: u32,
    /// Graphs with an operator outside wai-ops/1-i32: run on the CPU, not here.
    pub cpu: u32,
    /// Vectors whose load is refused: nothing for a device to run.
    pub load_refused: u32,
    /// Decodes whose device result differs from the expectation, with what differed.
    pub mismatches: Vec<String>,
    /// Device errors.
    pub errors: Vec<String>,
}

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

fn cpu_of(l: &Loaded, d: &Decode, seq: &mut Sequence) -> Expect {
    match wai_ops::interp::run(l, &d.inputs, &d.scalars, seq) {
        Ok(t) => Expect::Outputs(t),
        Err(r) => Expect::Refuse(r.code.as_str().to_string()),
    }
}

/// Run every vector on `gpu` at every width its graph admits (acc64 always, with
/// narrowing on and then off; acc32 too where its certificate meets it), stopping a
/// width's sequence at its first mismatch. `limit` caps the mismatches recorded. The
/// device's narrowing is on when this returns.
pub fn run_all(gpu: &Gpu, vs: &[Vector], limit: usize) -> Tally {
    let mut t = Tally::default();
    for v in vs {
        let l = match wai_ops::graph::load(&v.graph, Profile::Acc64) {
            Ok(l) => l,
            Err(r) => {
                assert_eq!(r.code.as_str(), v.load, "{}: the load the vector records", v.id);
                t.load_refused += 1;
                continue;
            }
        };
        assert_eq!(v.load, "ok", "{}", v.id);
        if support(&l).is_err() {
            t.cpu += 1;
            continue;
        }
        let narrow = wai_ops_gpu::prepare::narrow_plan(l.graph()).iter().filter(|c| c.is_some()).count() as u32;
        t.narrow_ops += narrow;
        t.narrow_vectors += (narrow > 0) as u32;
        let mut widths = vec![(Width::Acc64, true), (Width::Acc64, false)];
        if l.meets(Profile::Acc32) {
            widths.push((Width::Acc32, true));
        }
        for (w, narrowing) in widths {
            gpu.set_narrowing(narrowing);
            let mut seq = GpuSequence::new();
            let mut cpu = Sequence::new();
            match (w, narrowing) {
                (Width::Acc32, _) => t.vectors32 += 1,
                (Width::Acc64, true) => t.vectors64 += 1,
                (Width::Acc64, false) => t.vectors64w += 1,
            }
            for (k, d) in v.decodes.iter().enumerate() {
                let want = cpu_of(&l, d, &mut cpu);
                assert_eq!(want, d.expect, "{}: decode {k}: the CPU interpreter's result is the committed one", v.id);
                match (w, narrowing) {
                    (Width::Acc32, _) => t.decodes32 += 1,
                    (Width::Acc64, true) => t.decodes64 += 1,
                    (Width::Acc64, false) => t.decodes64w += 1,
                }
                if matches!(d.expect, Expect::Refuse(_)) && narrowing {
                    t.refusals += 1;
                }
                match got_of(blocking::run_at(gpu, &l, w, &d.inputs, &d.scalars, &mut seq)) {
                    Ok(got) if got == d.expect => {}
                    Ok(got) => {
                        if t.mismatches.len() < limit {
                            let n = if narrowing { "" } else { " (narrowing off)" };
                            t.mismatches.push(format!("{} ({}) decode {k} at {}{n}: {}", v.id, v.file, w.as_str(), diff(&d.expect, &got)));
                        }
                        break;
                    }
                    Err(e) => {
                        t.errors.push(format!("{} decode {k} at {}: {e}", v.id, w.as_str()));
                        break;
                    }
                }
            }
        }
    }
    gpu.set_narrowing(true);
    t
}

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(),
    }
}