open-standards/wai/wai-ops/tests/trace.rs

//! Hand-built graphs at the edge of what `acc64` certifies, run observed: every value
//! each operator forms under a named intermediate (§8.5, Table 8.5-A), and every
//! producer result, lies in the interval the bounds pass records for it.
//!
//! Each graph loads from its bytes under `acc64` (load step 8 included), and is run on
//! inputs at the ends of their declared ranges — all at one end, all at the other, the
//! two ends alternating in both phases — and on pseudo-random inputs within them.

use wai_ops::bounds::{self, Within};
use wai_ops::graph::{self, BiasKind, CAP, Graph, InputRec, Op, Params, Round, Rq, ScalarRec, TensorRec, opcode::*};
use wai_ops::interp::{Sequence, Tensor, run_observed};
use wai_ops::{Profile, plan};

fn rng(s: &mut u64) -> u64 {
    *s = s.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
    *s >> 11
}

/// A one-operator graph: inputs with records `ins` (declared sizes `hw`, or none), one
/// output.
fn graph1(ins: Vec<TensorRec>, out: TensorRec, code: u8, p: Params, scalars: Vec<ScalarRec>, hw: (u32, u32)) -> Graph {
    let n = ins.len() as u32;
    let mut tensors = ins;
    tensors.push(out);
    Graph {
        max_elems: CAP as u32,
        scalars,
        inputs: (0..n).map(|t| InputRec { tensor: t, max_h: hw.0, max_w: hw.1, kind: 0, src: 0, init: 0 }).collect(),
        outputs: vec![n],
        tensors,
        ops: vec![Op { code, ins: (0..n).collect(), outs: vec![n], p }],
    }
}

/// Inputs for one run: pattern 0 all at `lo`, 1 all at `hi`, 2 and 3 alternating, and
/// above 3 pseudo-random within `[lo, hi]`.
fn fill(c: u32, h: u32, w: u32, lo: i64, hi: i64, pattern: u32, s: &mut u64) -> Tensor {
    let n = (c * h * w) as usize;
    let data = (0..n)
        .map(|k| match pattern {
            0 => lo,
            1 => hi,
            2 => if k % 2 == 0 { lo } else { hi },
            3 => if k % 2 == 0 { hi } else { lo },
            _ => {
                let span = (hi as i128 - lo as i128) as u128 + 1;
                (lo as i128 + (rng(s) as u128 % span) as i128) as i64
            }
        })
        .collect();
    Tensor::new(c, h, w, data)
}

/// Load `g` under `acc64`, then run it on every pattern for its inputs' records and
/// `shapes`, observed; every value must lie in its interval.
fn check(name: &str, g: Graph, shapes: &[(u32, u32)], scalars: &[i64]) {
    let l = graph::load(&g.to_bytes(), Profile::Acc64).unwrap_or_else(|r| panic!("{name}: refused at load: {r:?}"));
    let sizes = plan::load_plan(&l).unwrap();
    let b = bounds::pass(&l, sizes.as_deref()).unwrap();
    let mut s = 0x9E3779B97F4A7C15u64;
    let mut seen = 0u64;
    for pattern in 0..12 {
        let inputs: Vec<Tensor> = l.inputs.iter().zip(shapes).map(|(i, &(h, w))| {
            let t = &l.tensors[i.tensor as usize];
            fill(t.channels, h, w, t.lo, t.hi, pattern, &mut s)
        }).collect();
        let mut obs = Within::new(&b);
        run_observed(&l, &inputs, scalars, &mut Sequence::new(), &mut obs).unwrap_or_else(|r| panic!("{name}: pattern {pattern} refused: {r:?}"));
        assert!(obs.outside.is_empty(), "{name}: pattern {pattern}: outside: {:?}", obs.outside);
        assert!(obs.unrecorded.is_empty(), "{name}: pattern {pattern}: unrecorded: {:?}", obs.unrecorded);
        seen += obs.seen;
    }
    assert!(seen > 0, "{name}");
}

fn rq(round: Round, bias: BiasKind, b: Vec<i64>, m: Option<Vec<i64>>, shift: u8) -> Rq {
    let (m_kind, m_bytes, mult) = match m {
        None => (0, 0, vec![]),
        Some(v) => (2, 8, v),
    };
    let b_bytes = if bias == BiasKind::None { 0 } else { 8 };
    Rq { round, bias_kind: bias, m_kind, s_kind: 1, m_bytes, b_bytes, mult, shift: vec![shift], bias: if bias == BiasKind::None { vec![] } else { b } }
}

#[test]
fn convolutions_at_the_edge() {
    // 18 terms of 2^20 × 2^38: 2^62.2, under i64; products by ±2^30, under i128.
    let w: Vec<i64> = (0..36).map(|k| if k % 3 == 0 { -(1 << 20) } else { 1 << 20 }).collect();
    for round in [Round::Floor, Round::HalfUp, Round::HalfEven] {
        let r = rq(round, BiasKind::Post, vec![i32::MAX as i64, i32::MIN as i64], Some(vec![1 << 30, -(1 << 30)]), 40);
        let p = Params::Conv { cout: 2, kh: 3, kw: 3, sh: 1, sw: 1, ph: 1, pw: 1, groups: 1, w_bytes: 4, rq: r, w: w.clone() };
        check("conv", graph1(vec![TensorRec::check(2, -(1 << 38), 1 << 38)], TensorRec::derived(2), CONV, p, vec![], (0, 0)), &[(4, 5)], &[]);
        let r = rq(round, BiasKind::Pre, vec![-(1 << 61), 1 << 61], None, 63);
        let p = Params::ConvT { cout: 2, kh: 3, kw: 3, sh: 2, sw: 2, ph: 1, pw: 1, oph: 1, opw: 1, w_bytes: 4, rq: r, w: w.clone() };
        check("conv_transpose", graph1(vec![TensorRec::check(2, -(1 << 38), 1 << 38)], TensorRec::derived(2), CONV_TRANSPOSE, p, vec![], (0, 0)), &[(3, 4)], &[]);
    }
}

#[test]
fn requantisations_and_products_at_the_edge() {
    let r = rq(Round::HalfUp, BiasKind::Pre, vec![i64::MIN / 2, i64::MAX / 2], Some(vec![i64::MAX, i64::MIN]), 127);
    check("rescale", graph1(vec![TensorRec::check(2, i64::MIN / 2, i64::MAX / 2)], TensorRec::derived(2), RESCALE, Params::Rescale(r), vec![], (0, 0)), &[(2, 3)], &[]);
    let p = Params::Chscale { scalar: 0, n_rows: 2, round: Round::HalfEven, m_bytes: 8, shift: 64, m: vec![i64::MAX, i64::MIN, -1, 1] };
    check("chscale", graph1(vec![TensorRec::check(2, i64::MIN, i64::MAX)], TensorRec::derived(2), CHSCALE, p.clone(), vec![ScalarRec { lo: 0, hi: 1 }], (0, 0)), &[(2, 2)], &[0]);
    check("chscale", graph1(vec![TensorRec::check(2, i64::MIN, i64::MAX)], TensorRec::derived(2), CHSCALE, p, vec![ScalarRec { lo: 0, hi: 1 }], (0, 0)), &[(2, 2)], &[1]);
    let r = rq(Round::Floor, BiasKind::None, vec![], Some(vec![1 << 40, -(1 << 40)]), 70);
    check("mul", graph1(vec![TensorRec::check(2, -(1 << 31), 1 << 31), TensorRec::check(1, i32::MIN as i64, i32::MAX as i64)], TensorRec::derived(2), MUL, Params::Mul(r), vec![], (0, 0)), &[(2, 2), (2, 2)], &[]);
    let gr = Rq { round: Round::HalfUp, bias_kind: BiasKind::None, m_kind: 0, s_kind: 1, m_bytes: 0, b_bytes: 0, mult: vec![], shift: vec![16], bias: vec![] };
    check("gram", graph1(vec![TensorRec::check(4, -(1 << 28), 1 << 28), TensorRec::check(4, -(1 << 28), 1 << 28)], TensorRec::derived(4), GRAM, Params::Gram { groups: 2, rq: gr.clone() }, vec![], (3, 3)), &[(3, 3), (3, 3)], &[]);
    check("apply", graph1(vec![TensorRec::check(4, -(1 << 30), 1 << 30), TensorRec::check(4, -(1 << 31), 1 << 31)], TensorRec::derived(4), APPLY, Params::Apply { groups: 2, rq: gr }, vec![], (0, 0)), &[(1, 2), (3, 3)], &[]);
}

#[test]
fn elementwise_at_the_edge() {
    check("leaky", graph1(vec![TensorRec::check(1, -(1 << 31), i64::MAX)], TensorRec::derived(1), LEAKY, Params::Leaky { a: 1 << 31, f: 63 }, vec![], (0, 0)), &[(3, 3)], &[]);
    check("leaky", graph1(vec![TensorRec::check(1, i32::MIN as i64, 5)], TensorRec::derived(1), LEAKY, Params::Leaky { a: i32::MIN as i64, f: 0 }, vec![], (0, 0)), &[(3, 3)], &[]);
    let mut t = vec![0i16; 65536];
    for (k, v) in t.iter_mut().enumerate().skip(1) {
        *v = ((k as i64 * 7919) % 65536 - 32768) as i16;
    }
    check("lut16", graph1(vec![TensorRec::check(1, -32767, 32767)], TensorRec::derived(1), LUT16, Params::Lut16 { e_in: -128, e_out: -128, t }, vec![], (0, 0)), &[(4, 4)], &[]);
    check("abs", graph1(vec![TensorRec::check(1, i64::MIN + 1, i64::MAX)], TensorRec::derived(1), ABS, Params::None, vec![], (0, 0)), &[(2, 2)], &[]);
    check("chunk_add", graph1(vec![TensorRec::check(2, i64::MIN / 2, i64::MAX / 2)], TensorRec::derived(1), CHUNK_ADD, Params::None, vec![], (0, 0)), &[(2, 2)], &[]);
}

#[test]
fn normalisations_and_scalar_functions_at_the_edge() {
    // igdn: γ·x² of 2^31 × 2^62 = 2^93, under i128; r of the final norm under 2^47.
    let p = Params::Igdn { gf: 0, frac: 30, g_bytes: 4, gamma: vec![i32::MAX as i64, -1, 7, i32::MIN as i64], beta: vec![i64::MAX, i64::MIN] };
    check("igdn", graph1(vec![TensorRec::check(2, -(1 << 31), 1 << 31)], TensorRec::derived(2), IGDN, p, vec![], (0, 0)), &[(2, 3)], &[]);
    for eps in [1u64, 1 << 32, u64::MAX] {
        let p = Params::LayerNorm { eps, g_bytes: 8, gamma: vec![i32::MAX as i64, i32::MIN as i64, 1], beta: vec![i64::MAX / 4, i64::MIN / 4, 0] };
        check("layernorm_c", graph1(vec![TensorRec::check(3, -(1 << 40), 1 << 40)], TensorRec::derived(3), LAYERNORM_C, p, vec![], (0, 0)), &[(2, 2)], &[]);
        check("l2norm", graph1(vec![TensorRec::check(2, -(1 << 40), 1 << 40)], TensorRec::derived(2), L2NORM, Params::L2norm { eps }, vec![], (3, 3)), &[(3, 3)], &[]);
    }
    // The widest spread acc64 allows: exp.negx = h − l = 2^63 − 1. One more, and the
    // set is refused at load (below).
    check("softmax", graph1(vec![TensorRec::check(2, -(1 << 62), (1 << 62) - 1)], TensorRec::derived(2), SOFTMAX, Params::None, vec![], (2, 5)), &[(2, 5)], &[]);
    let wide = graph1(vec![TensorRec::check(2, -(1 << 62), 1 << 62)], TensorRec::derived(2), SOFTMAX, Params::None, vec![], (2, 5));
    assert!(graph::load(&wide.to_bytes(), Profile::Acc64).is_err(), "exp.negx = 2^63 does not fit i64");
    check("softmax", graph1(vec![TensorRec::check(1, -5, 7)], TensorRec::derived(1), SOFTMAX, Params::None, vec![], (1, 4)), &[(1, 4)], &[]);
    check("sigmoid", graph1(vec![TensorRec::check(1, i64::MIN + 1, i64::MAX)], TensorRec::derived(1), SIGMOID, Params::None, vec![], (0, 0)), &[(3, 3)], &[]);
    check("sigmoid", graph1(vec![TensorRec::check(1, -200000, 200000)], TensorRec::derived(1), SIGMOID, Params::None, vec![], (0, 0)), &[(3, 3)], &[]);
    // α = −2^63 gives ⌊α·(−2^16) / 2^16⌋ = 2^63, past i64: refused at load.
    let wide = graph1(vec![TensorRec::check(1, -1, 0)], TensorRec::derived(1), ELU, Params::Elu { alpha: i64::MIN }, vec![], (0, 0));
    assert!(graph::load(&wide.to_bytes(), Profile::Acc64).is_err(), "elu's producer result reaches 2^63");
    for alpha in [i64::MAX, i64::MIN + 1, 65536, -3] {
        check("elu", graph1(vec![TensorRec::check(1, i64::MIN + 1, i64::MAX)], TensorRec::derived(1), ELU, Params::Elu { alpha }, vec![], (0, 0)), &[(3, 3)], &[]);
        check("elu", graph1(vec![TensorRec::check(1, -300000, 10)], TensorRec::derived(1), ELU, Params::Elu { alpha }, vec![], (0, 0)), &[(3, 3)], &[]);
    }
}