open-standards/wai/wai-pack-gpu/tests/kernels.rs

//! Each kernel against its reference on every adapter that opens: the motion search
//! against packaging §5.5's search as the text states it, the SSE against the plain sum,
//! BLAKE3 against the reference crate on the test vectors' lengths. Every case is seeded,
//! so a difference names its case. With no adapter, each test prints SKIPPED and passes,
//! unless WAI_PACK_GPU_REQUIRE is set.

mod support;

use support::*;

#[test]
fn the_motion_search_is_the_texts() {
    for g in gpus() {
        let mut rng = Rng(0x5eed_0001);
        let mut n = 0;
        for &(h, w, block, range, prec) in MOTION_CASES {
            for (kind, src, rf) in frame_pairs(&mut rng, h, w) {
                let got = g.motion_frame(&src, &rf, h, w, block, range, prec).expect("the kernel runs");
                let want = search(&src, &rf, h, w, block, range, prec);
                assert_eq!(got, want, "{}: {kind} {h}x{w} block {block} range {range} prec {prec}", g.adapter().label());
                n += 1;
            }
        }
        eprintln!("{}: motion search, {n} frames, every vector the text's", g.adapter().label());
    }
}

/// What the kernel does not take is refused before anything runs.
#[test]
fn the_motion_search_refuses_what_it_does_not_take() {
    for g in gpus() {
        let f = vec![0u8; 4 * 4 * 3];
        let un = |r: Result<Vec<(i16, i16)>, wai_pack_gpu::GpuError>| matches!(r, Err(wai_pack_gpu::GpuError::Unsupported(_)));
        assert!(un(g.motion_frame(&f, &f, 4, 4, 4, 1, 3)), "precision 3");
        assert!(un(g.motion_frame(&f, &f, 4, 4, 4, 128, 2)), "range 128");
        assert!(un(g.motion_frame(&f, &f, 4, 5, 4, 1, 2)), "a frame that is not h × w × 3");
        assert!(un(g.motion_frame(&f, &f, 4, 4, 0, 1, 2)), "block 0");
        assert!(un(g.motion_frame(&f, &f, 4, 4, 2370, 1, 2)), "a block whose cost can pass 2^32 - 1");
    }
}

#[test]
fn the_sse_is_the_sum() {
    for g in gpus() {
        let mut rng = Rng(0x5eed_0002);
        for n in [0usize, 1, 3, 4, 5, 4095, 4096, 4097, 65_537, 1_048_579] {
            let (a, b) = (rng.bytes(n), rng.bytes(n));
            assert_eq!(g.sse(&a, &b).unwrap(), sse(&a, &b), "{}: {n} random bytes", g.adapter().label());
        }
        // The largest partial: every byte 0 against 255, 4,096 to a thread.
        let (a, b) = (vec![0u8; 3 * 4096 + 7], vec![255u8; 3 * 4096 + 7]);
        assert_eq!(g.sse(&a, &b).unwrap(), sse(&a, &b), "{}: the largest squares", g.adapter().label());
    }
}

#[test]
fn blake3_is_the_reference() {
    for g in gpus() {
        let inputs: Vec<Vec<u8>> = BLAKE3_LENGTHS.iter().map(|&n| blake3_input(n)).collect();
        let refs: Vec<&[u8]> = inputs.iter().map(|v| v.as_slice()).collect();
        // All at once, in one dispatch, and each alone.
        let many = g.blake3_many(&refs).unwrap();
        for (m, h) in refs.iter().zip(&many) {
            assert_eq!(h, blake3::hash(m).as_bytes(), "{}: {} bytes, in a batch", g.adapter().label(), m.len());
            assert_eq!(&g.blake3_many(&[m]).unwrap()[0], blake3::hash(m).as_bytes(), "{}: {} bytes alone", g.adapter().label(), m.len());
        }
        // Random messages at random lengths, hashed in batches of mixed sizes.
        let mut rng = Rng(0x5eed_0003);
        let msgs: Vec<Vec<u8>> = (0..200).map(|_| {
            let n = rng.below(20_000) as usize;
            rng.bytes(n)
        }).collect();
        let refs: Vec<&[u8]> = msgs.iter().map(|v| v.as_slice()).collect();
        for (m, h) in refs.iter().zip(g.blake3_many(&refs).unwrap()) {
            assert_eq!(&h, blake3::hash(m).as_bytes(), "{}: a random message of {} bytes", g.adapter().label(), m.len());
        }
        eprintln!("{}: BLAKE3, {} lengths and 200 random messages, every hash the reference crate's", g.adapter().label(), BLAKE3_LENGTHS.len());
    }
}