open-standards/wai/wai-pack-gpu/tests/support/mod.rs

//! What the kernel tests share: the adapters, the references the kernels are held to, and
//! the cases.
//!
//! The references are transcribed here from the extension text (packaging §5.5, §6.6),
//! apart from the crate's kernels and from wai-rs, whose encoder wai-pack's differential
//! test holds the device to; BLAKE3's is the reference crate.

#![allow(dead_code)]

use wai_pack_gpu::{Discovery, Gpu, Sources};

/// The adapters that opened. With none, a test prints SKIPPED and passes, unless
/// `WAI_PACK_GPU_REQUIRE` is set: then an absent adapter fails, never skips.
pub fn gpus_with(src: &Sources, label: &str) -> Vec<Gpu> {
    let Discovery { found, refused } = wai_pack_gpu::discover_with(src, label);
    for (a, why) in &refused {
        eprintln!("adapter refused: {}: {why}", a.label());
    }
    if found.is_empty() {
        assert!(std::env::var_os("WAI_PACK_GPU_REQUIRE").is_none(), "WAI_PACK_GPU_REQUIRE is set and no adapter opened");
        eprintln!("SKIPPED: no GPU adapter opened");
    }
    found
}

pub fn gpus() -> Vec<Gpu> {
    gpus_with(&Sources::shipped(), "shipped")
}

/// A seeded generator (xorshift64*), so every run sees the same cases.
pub struct Rng(pub u64);

impl Rng {
    pub fn next(&mut self) -> u64 {
        self.0 ^= self.0 >> 12;
        self.0 ^= self.0 << 25;
        self.0 ^= self.0 >> 27;
        self.0.wrapping_mul(0x2545_F491_4F6C_DD1D)
    }
    pub fn below(&mut self, n: u64) -> u64 {
        self.next() % n
    }
    pub fn bytes(&mut self, n: usize) -> Vec<u8> {
        (0..n).map(|_| self.next() as u8).collect()
    }
}

/// integer-payloads §3's prediction: whole-pel, or half-pel bilinear with integer
/// rounding, every coordinate clamped to the frame.
#[allow(clippy::too_many_arguments)]
pub fn predict(rf: &[u8], h: usize, w: usize, y: i64, x: i64, dy: i64, dx: i64, c: usize, prec: i64) -> i64 {
    let p = |yy: i64, xx: i64| rf[((yy.clamp(0, h as i64 - 1) as usize) * w + xx.clamp(0, w as i64 - 1) as usize) * 3 + c] as i64;
    if prec == 1 {
        return p(y + dy, x + dx);
    }
    let (iy, fy, ix, fx) = (y + dy.div_euclid(2), dy.rem_euclid(2), x + dx.div_euclid(2), dx.rem_euclid(2));
    match (fy, fx) {
        (0, 0) => p(iy, ix),
        (0, _) => (p(iy, ix) + p(iy, ix + 1) + 1) / 2,
        (_, 0) => (p(iy, ix) + p(iy + 1, ix) + 1) / 2,
        _ => (p(iy, ix) + p(iy, ix + 1) + p(iy + 1, ix) + p(iy + 1, ix + 1) + 2) / 4,
    }
}

/// packaging §5.5's search, as the text states it: for each block in raster order, every
/// candidate's full cost, and the least by cost, then by |dy| + |dx|, then the earliest.
pub fn search(src: &[u8], rf: &[u8], h: usize, w: usize, block: usize, range: i64, prec: i64) -> Vec<(i16, i16)> {
    let mut out = Vec::new();
    for by in 0..h.div_ceil(block) {
        for bx in 0..w.div_ceil(block) {
            let mut best: Option<(u64, i64, i64, i64)> = None;
            for dy in -range..=range {
                for dx in -range..=range {
                    let mut cost = 0u64;
                    for y in by * block..((by + 1) * block).min(h) {
                        for x in bx * block..((bx + 1) * block).min(w) {
                            for c in 0..3 {
                                cost += (src[(y * w + x) * 3 + c] as i64 - predict(rf, h, w, y as i64, x as i64, dy, dx, c, prec)).unsigned_abs();
                            }
                        }
                    }
                    let key = (cost, dy.abs() + dx.abs());
                    if best.is_none_or(|b| key < (b.0, b.1)) {
                        best = Some((cost, key.1, dy, dx));
                    }
                }
            }
            let b = best.expect("one candidate at least");
            out.push((b.2 as i16, b.3 as i16));
        }
    }
    out
}

/// `Σ (a − b)²`.
pub fn sse(a: &[u8], b: &[u8]) -> u64 {
    a.iter().zip(b).map(|(&x, &y)| (x as i64 - y as i64).pow(2) as u64).sum()
}

/// A frame pair of each kind the search meets: noise (costs all apart), a smooth field
/// moved by a fraction of a pel (near ties), a flat field (every candidate ties, so the
/// tie-break alone decides), and hard edges near the frame's borders (clamping).
pub fn frame_pairs(rng: &mut Rng, h: usize, w: usize) -> Vec<(&'static str, Vec<u8>, Vec<u8>)> {
    let n = h * w * 3;
    let noise_a = rng.bytes(n);
    let noise_b = rng.bytes(n);
    let field = |t: usize| -> Vec<u8> {
        (0..h * w)
            .flat_map(|i| {
                let (y, x) = (i / w, i % w);
                let v = ((x * 7 + y * 5 + t * 3) % 256) as u8;
                [v, v / 2 + (y % 7) as u8, 255 - v]
            })
            .collect()
    };
    let flat = vec![97u8; n];
    let flat2 = vec![101u8; n];
    let edges = |t: usize| -> Vec<u8> {
        (0..h * w)
            .flat_map(|i| {
                let (y, x) = (i / w, i % w);
                let on = (x + t) % 5 < 2 || y + t < 3 || x + 2 >= w;
                if on { [250, 10, 128] } else { [5, 240, 64] }
            })
            .collect()
    };
    // The noisy frame predicted from a blurred copy of itself, so costs differ by little.
    let mut near = noise_a.clone();
    for v in near.iter_mut() {
        *v = (*v / 8) * 8 + 3;
    }
    // Columns alternating in the reference and moved by one in the source: inside the
    // frame a vector one column left and one right match alike, at the same |dy| + |dx|,
    // so the earlier vector decides.
    let stripes = |t: usize| -> Vec<u8> {
        (0..h * w).flat_map(|i| if (i % w + t).is_multiple_of(2) { [40, 200, 90] } else { [200, 40, 160] }).collect()
    };
    vec![
        ("noise", noise_a.clone(), noise_b),
        ("near", noise_a, near),
        ("field", field(1), field(0)),
        ("flat", flat, flat2),
        ("edges", edges(1), edges(0)),
        ("stripes", stripes(1), stripes(0)),
    ]
}

/// The frame sizes and search parameters the motion cases cover: frames that are not a
/// whole number of blocks, a frame smaller than a block, both precisions, and ranges
/// from 0 to past the frame.
pub const MOTION_CASES: &[(usize, usize, usize, i64, i64)] = &[
    // (h, w, block, range, prec)
    (24, 40, 8, 3, 1),
    (24, 40, 8, 3, 2),
    (19, 23, 4, 2, 2),
    (33, 17, 16, 8, 2),
    (9, 13, 16, 5, 1),
    (12, 12, 4, 0, 2),
    (16, 16, 32, 1, 2),
    (36, 64, 16, 14, 2),
    (5, 7, 2, 9, 2),
];

/// The message lengths of the BLAKE3 test vectors, and a few past them.
pub const BLAKE3_LENGTHS: &[usize] = &[0, 1, 63, 64, 65, 127, 128, 129, 1023, 1024, 1025, 2048, 2049, 3072, 3073, 4096, 4097, 5120, 5121, 6144, 6145, 7168, 7169, 8192, 8193, 16384, 31744, 102_400, 1_000_003];

/// The BLAKE3 test vectors' input: byte i is i mod 251.
pub fn blake3_input(n: usize) -> Vec<u8> {
    (0..n).map(|i| (i % 251) as u8).collect()
}