open-standards/wai/wai-rs/tests/int_kernels_geometry.rs

//! Geometry at the edges of what the fast integer kernels lay out, run through the
//! public API under the scalar reference kernels AND under each fast-kernel policy
//! (serial, explicit multi-threaded tilings, automatic):
//!
//! * a 1×1 input upsampled by strides from 1 up to `usize::MAX`, through the
//!   transposed convolutions of int_hyper (`run_net`), int_synth
//!   (`decode_from_spec`) and the JPEG-AI synthesis (`conv_transpose`);
//! * a regular convolution whose stride reaches or exceeds its padded input
//!   (`run_net`, `conv_grouped`);
//! * padding far beyond the kernel, with output padding restoring a small output;
//! * an input of two or more rows and columns whose stride and padding make the
//!   reference's own size arithmetic overflow (int_hyper, int_synth, JPEG-AI).
//!
//! The reference computes each of the first three in a handful of operations. Every
//! policy must return the same integers (checked against a hand computation), and
//! the fast path must not overflow, or allocate or loop in proportion to the stride
//! or the padding, where the reference does not. In the last, the reference panics
//! (or, for int_synth with overflow checks off, returns), and every policy must do
//! the same. The int_audio 1-D transposed convolution has the
//! same stride check beside its kernel (`int_audio` unit tests). CI runs this file
//! beside `int_kernels_vectors` on every `wai-conformance` leg, with overflow checks
//! on and off.
#![cfg(feature = "neural_int")]

use serde_json::json;
use wai::int_kernels::{with_policy, Policy};

const POLICIES: [Policy; 5] = [Policy::reference(), Policy::fast(1), Policy::fast(2), Policy::fast(5), Policy::fast(0)];

/// A 3×3 kernel `[[1, 2, 3], [4, 5, 6], [7, 8, 9]]`.
const KERNEL: [i64; 9] = [1, 2, 3, 4, 5, 6, 7, 8, 9];

/// Run under every policy; all results must be identical. Returns the result.
fn same_under_every_policy<T: PartialEq + std::fmt::Debug>(what: &str, f: impl Fn() -> T) -> T {
    let want = with_policy(POLICIES[0], &f);
    for p in &POLICIES[1..] {
        let got = with_policy(*p, &f);
        assert!(got == want, "{what}: result under {p:?} differs from the reference kernels ({got:?} vs {want:?})");
    }
    want
}

/// Whether this build traps on integer overflow. Detected by performing one, since
/// `cfg!(debug_assertions)` only usually coincides with `overflow-checks`.
#[cfg(target_pointer_width = "64")]
fn overflow_checks() -> bool {
    std::panic::catch_unwind(|| std::hint::black_box(i64::MAX) + std::hint::black_box(1)).is_err()
}

/// Run under every policy, catching a panic. The outcome (the value, or that the
/// call panicked) must be the same under every policy. Returns the reference
/// kernels' outcome, `None` for a panic. The panic message is not compared: a fast
/// policy may fail at another element first, or on a worker thread.
fn same_outcome_under_every_policy<T: PartialEq + std::fmt::Debug>(what: &str, f: impl Fn() -> T) -> Option<T> {
    let run = |p: Policy| std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| with_policy(p, &f))).ok();
    let want = run(POLICIES[0]);
    for p in &POLICIES[1..] {
        let got = run(*p);
        assert!(got == want, "{what}: outcome under {p:?} differs from the reference kernels' ({got:?} vs {want:?}; None = panicked)");
    }
    want
}

/// Strides from ordinary to the largest a `usize` holds. 2^29 fits the `u32` stride
/// field of a WHM1 model; from 2^32 (64-bit targets) the product of the two strides
/// no longer fits in 64 bits.
///
/// 2^29 comes first, deliberately. A kernel that allocates one tap table per
/// (row, column) phase — as the fast transposed kernel did before it built tables
/// only for phases holding output — fails there at once (the table count exceeds
/// what any allocation can hold: "capacity overflow"). Placed after 2^16 it would
/// first commit memory for 2^32 tables (hundreds of GiB) and drive the machine into
/// swap before failing.
fn strides() -> Vec<usize> {
    #[allow(unused_mut)] // pushed to on 64-bit targets only
    let mut s = vec![1 << 29, 1, 2, 3, 1 << 13, 1 << 16, usize::MAX];
    #[cfg(target_pointer_width = "64")]
    s.push(1 << 32);
    s
}

/// A one-layer int_hyper model.
fn hyper_model(layer: serde_json::Value) -> wai::int_hyper_synth::IntModel {
    wai::int_hyper_synth::IntModel::from_json(&json!({ "afrac": 0, "gf": 0, "shift": 0, "layers": [layer] }))
}

/// An int_synth payload: a 1×1 output image from one int8 latent (C = 1, S = 1).
fn synth_payload(latent: i8) -> Vec<u8> {
    let mut p = Vec::new();
    p.extend(1u32.to_le_bytes()); // H
    p.extend(1u32.to_le_bytes()); // W
    p.extend(1u16.to_le_bytes()); // C
    p.extend([0u8, 0]); // S_log2, reserved
    p.push(latent as u8);
    p
}

/// A one-layer int_synth spec: 1 → 3 channels, 3×3 kernel `1..=27` (channel `oc`
/// holds `9·oc + KERNEL`), shift 0, no ReLU.
fn synth_spec(stride: usize, pad: usize, out_pad: usize, bias: [i32; 3]) -> Vec<u8> {
    let weight: Vec<i32> = (1..=27).collect();
    json!({ "layers": [{
        "in_ch": 1, "out_ch": 3, "k": 3, "stride": stride, "pad": pad, "out_pad": out_pad,
        "weight": weight, "bias": bias, "shift": 0, "relu": false
    }]})
    .to_string()
    .into_bytes()
}

#[test]
fn a_1x1_input_upsampled_by_any_stride() {
    // k = 3, pad 1: the output is 1×1 and receives only the centre tap, 5·x.
    for s in strides() {
        let got = same_under_every_policy(&format!("JPEG-AI conv_transpose, stride {s}"), || {
            wai::int_transformer::conv_transpose(&[5], 1, 1, 1, &KERNEL, 1, 3, s, 1, 0, &[1], &[0], 0)
        });
        assert_eq!(got, (vec![25], 1, 1), "JPEG-AI conv_transpose, stride {s}");

        let model = hyper_model(json!({
            "type": "convT", "out_ch": 1, "k": 3, "stride": s, "pad": 1, "opad": 0,
            "weight": KERNEL, "mult": [1], "bias": [0]
        }));
        let got = same_under_every_policy(&format!("int_hyper convT, stride {s}"), || {
            wai::int_hyper_synth::run_net(&model, &[5], 1, 1, 1)
        });
        assert_eq!(got, (vec![25], 1, 1, 1), "int_hyper convT, stride {s}");

        let (payload, spec) = (synth_payload(5), synth_spec(s, 1, 0, [0, 0, 0]));
        let got = same_under_every_policy(&format!("int_synth, stride {s}"), || {
            wai::int_synth::decode_from_spec(&payload, &spec).map(|img| img.rgb)
        });
        // centre taps 5, 14, 23 of the three output channels, times 5
        assert_eq!(got, Ok(vec![25, 70, 115]), "int_synth, stride {s}");
    }
}

#[test]
fn a_regular_conv_whose_stride_exceeds_its_padded_input() {
    // 5×5 input 1..=25, k = 3. Pad 1 (7×7 padded): from stride 5 on, one output,
    // the taps that meet x[0], x[1], x[5], x[6]: 5·1 + 6·2 + 8·6 + 9·7. Pad 2 (9×9
    // padded): from stride 7 on, one output, only tap (2, 2) meets x[0]: 9·1. At
    // stride usize::MAX, pad 2 puts two leading column phases before the first input
    // column, whose position must be found without overflowing.
    let x: Vec<i64> = (1..=25).collect();
    for (pad, want, first) in [(1usize, 5 + 6 * 2 + 8 * 6 + 9 * 7, 5usize), (2, 9, 7)] {
        #[allow(unused_mut)] // pushed to on 64-bit targets only
        let mut ss = vec![first, 7, 8, 9, 1 << 13, 1 << 16, 1 << 29, usize::MAX];
        #[cfg(target_pointer_width = "64")]
        ss.push(1 << 32);
        for s in ss {
            let got = same_under_every_policy(&format!("JPEG-AI conv_grouped, pad {pad}, stride {s}"), || {
                wai::int_transformer::conv_grouped(&x, 1, 5, 5, &KERNEL, 1, 3, s, pad, 1, &[1], &[0], 0)
            });
            assert_eq!(got, (vec![want], 1, 1), "JPEG-AI conv_grouped, pad {pad}, stride {s}");

            let model = hyper_model(json!({
                "type": "conv", "out_ch": 1, "k": 3, "stride": s, "pad": pad,
                "weight": KERNEL, "mult": [1], "bias": [0]
            }));
            let got = same_under_every_policy(&format!("int_hyper conv, pad {pad}, stride {s}"), || {
                wai::int_hyper_synth::run_net(&model, &x, 1, 5, 5)
            });
            assert_eq!(got, (vec![want], 1, 1, 1), "int_hyper conv, pad {pad}, stride {s}");
        }
    }
}

#[test]
fn padding_far_beyond_the_kernel() {
    let p = 1usize << 12;
    // Transposed: pad 2^12 with output padding 2^13 − 2 restores a 1×1 output that
    // no tap reaches (it sits 2^12 rows into the uncropped plane), so it is the bias.
    let (payload, spec) = (synth_payload(5), synth_spec(1, p, 2 * p - 2, [10, 20, 30]));
    let got = same_under_every_policy("int_synth, pad 2^12", || {
        wai::int_synth::decode_from_spec(&payload, &spec).map(|img| img.rgb)
    });
    assert_eq!(got, Ok(vec![10, 20, 30]));
    // Regular: a 1×1 input, 1×1 kernel, pad 2^12, stride 2^13 — a 2×2 output whose
    // samples all fall in the padding, so each is the bias.
    let got = same_under_every_policy("JPEG-AI conv_grouped, pad 2^12", || {
        wai::int_transformer::conv_grouped(&[5], 1, 1, 1, &[3], 1, 1, 2 * p, p, 1, &[1], &[7], 0)
    });
    assert_eq!(got, (vec![7; 4], 2, 2));
    let model = hyper_model(json!({
        "type": "conv", "out_ch": 1, "k": 1, "stride": 2 * p, "pad": p, "weight": [3], "mult": [1], "bias": [7]
    }));
    let got = same_under_every_policy("int_hyper conv, pad 2^12", || wai::int_hyper_synth::run_net(&model, &[5], 1, 1, 1));
    assert_eq!(got, (vec![7; 4], 1, 2, 2));
}

/// An int_synth payload: an `h`×`w` output image from a 2×2 int8 latent (C = 1,
/// S = 2), latents in row-major order.
#[cfg(target_pointer_width = "64")]
fn synth_payload_2x2(h: u32, w: u32, latents: [i8; 4]) -> Vec<u8> {
    let mut p = Vec::new();
    p.extend(h.to_le_bytes());
    p.extend(w.to_le_bytes());
    p.extend(1u16.to_le_bytes()); // C
    p.extend([1u8, 0]); // S_log2 = 1, reserved
    p.extend(latents.map(|v| v as u8));
    p
}

// Transposed convolutions of an input with two or more rows and columns, whose
// stride and padding make the reference's own size arithmetic overflow. The
// int_hyper, int_audio and JPEG-AI references size an uncropped plane
// `cout·(ho + 2·pad)·(wo + 2·pad)` in unchecked `usize` arithmetic; the int_synth
// reference computes the output size in unchecked `i64`. With overflow checks on,
// the reference panics on that arithmetic. With them off it wraps and, for the
// uncropped plane, then indexes out of the wrapped buffer. A fast policy must have
// the same outcome: it may not return a value where the reference panics. A 1×1
// input cannot show this, since `(h − 1)·stride` is 0. The int_audio 1-D form is
// tested beside its kernel.

/// A 3×3 input, and a 2×2 one with fields that fit a WHM1 model's `u32`.
const X: [i64; 9] = [1, 2, 3, 4, 5, 6, 7, 8, 9];
/// Stride 2^32 − 1 and pad 2^31 − 1: a 2×2 output from a 2×2 input, k = 1 (or a
/// 4×4 one, k = 3), under an uncropped plane of 2^32 × 2^32. On a 64-bit target
/// the plane size overflows; on a 32-bit one the stride arithmetic already does.
const S32: usize = u32::MAX as usize;
const P32: usize = (u32::MAX >> 1) as usize;

#[test]
fn jpegai_conv_transpose_where_the_reference_plane_size_overflows() {
    // 3×3 input, k = 1, stride = pad = 2^62: a 1×1 output, but an uncropped plane
    // of (2^63 + 1)² elements. The reference panics in both profiles.
    #[cfg(target_pointer_width = "64")]
    {
        let s = 1usize << 62;
        let got = same_outcome_under_every_policy("JPEG-AI conv_transpose, 3×3, stride = pad = 2^62", || {
            wai::int_transformer::conv_transpose(&X, 1, 3, 3, &[7], 1, 1, s, s, 0, &[1], &[0], 0)
        });
        assert_eq!(got, None, "the reference panics");
    }
    let got = same_outcome_under_every_policy("JPEG-AI conv_transpose, 2×2, u32 stride and pad", || {
        wai::int_transformer::conv_transpose(&X[..4], 1, 2, 2, &[7], 1, 1, S32, P32, 0, &[1], &[0], 0)
    });
    assert_eq!(got, None, "u32 fields: the reference panics");
}

#[test]
fn int_hyper_convt_where_the_reference_plane_size_overflows() {
    // As above, through the int_hyper decoder.
    let convt = |s: usize, p: usize| {
        hyper_model(json!({
            "type": "convT", "out_ch": 1, "k": 1, "stride": s, "pad": p, "opad": 0,
            "weight": [7], "mult": [1], "bias": [0]
        }))
    };
    #[cfg(target_pointer_width = "64")]
    {
        let model = convt(1 << 62, 1 << 62);
        let got = same_outcome_under_every_policy("int_hyper convT, 3×3, stride = pad = 2^62", || {
            wai::int_hyper_synth::run_net(&model, &X, 1, 3, 3)
        });
        assert_eq!(got, None, "the reference panics");
    }
    let model = convt(S32, P32);
    let got = same_outcome_under_every_policy("int_hyper convT, 2×2, u32 stride and pad", || {
        wai::int_hyper_synth::run_net(&model, &X[..4], 1, 2, 2)
    });
    assert_eq!(got, None, "u32 fields: the reference panics");
}

#[cfg(target_pointer_width = "64")]
#[test]
fn int_synth_where_the_reference_output_size_overflows_i64() {
    // 2×2 latent, k = 3, stride = pad = out_pad = 2^62: a 3×3 output. The reference
    // computes (2 − 1)·2^62 + 3 + 2^62 − 2·2^62 in `i64`, which overflows (2^63 + 3)
    // although every `usize` step fits. With overflow checks on it panics; with them
    // off the wrapped sum is still 3 and the decode returns. Only latent (1, 1)
    // reaches the output, through every tap.
    let s = 1usize << 62;
    let (payload, spec) = (synth_payload_2x2(3, 3, [1, 2, 3, 4]), synth_spec(s, s, s, [0, 0, 0]));
    let got = same_outcome_under_every_policy("int_synth, 2×2 latent, stride = pad = out_pad = 2^62", || {
        wai::int_synth::decode_from_spec(&payload, &spec).map(|img| img.rgb)
    });
    if overflow_checks() {
        assert_eq!(got, None, "the reference panics on its i64 output size");
    } else {
        // rgb[(3y + x)·3 + c] = 4 · weight[c][y][x] = 4 · (9c + 3y + x + 1)
        let want: Vec<u8> = (0..9).flat_map(|px| (0..3).map(move |c| (4 * (9 * c + px + 1)) as u8)).collect();
        assert_eq!(got, Some(Ok(want)), "the wrapped output size is 3");
    }

    // The u32 fields: no `i64` step overflows, and int_synth's reference allocates
    // no uncropped plane, so it returns. A 4×4 output that no tap reaches: the bias.
    let (payload, spec) = (synth_payload_2x2(4, 4, [1, 2, 3, 4]), synth_spec(S32, P32, 0, [10, 20, 30]));
    let got = same_outcome_under_every_policy("int_synth, 2×2 latent, u32 stride and pad", || {
        wai::int_synth::decode_from_spec(&payload, &spec).map(|img| img.rgb)
    });
    assert_eq!(got, Some(Ok([10u8, 20, 30].repeat(16))), "u32 fields: the bias alone");
}