67ac8c544b
Each decoder-level equivalence test now runs int_kernels::equivalence_policies: the reference, the fast kernels on every tiling, and each kernel set the processor runs. tests/int_ops_conformance.rs decodes every graph vector of ops-conformance/ (vectors/ and graphs/) through the interpreter with the fast kernels behind it, under each of those policies, and holds every decode to the integers or refusal the vector states. Co-Authored-By: Claude Opus 5.5 <[email protected]>
319 lines
17 KiB
Rust
319 lines
17 KiB
Rust
//! 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 (a `WHM1` layer through its lowering), int_synth
|
||
//! (`decode_from_spec`) and the JPEG-AI synthesis (`conv_transpose`);
|
||
//! * a regular convolution whose stride reaches or exceeds its padded input
|
||
//! (a `WHM1` layer through its lowering, `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,
|
||
//! and every policy must do the same. int_synth decodes through its lowering
|
||
//! (integer-operators §13.1), and the int_hyper networks through theirs (`whm1/1`,
|
||
//! §13.2): their fields are 32 bits and their run-time plan caps every size it forms at
|
||
//! 2^26, so where a case is past either, every policy refuses it, where the scalar kernel
|
||
//! the int_hyper decodes once ran panicked or computed. int_audio decodes through its
|
||
//! lowering `wam1/1` too; its 1-D transposed convolution's edge geometry is tested
|
||
//! through that graph under every policy (`int_audio` unit tests). CI runs this file
|
||
//! beside `int_kernels_vectors` on every `wai-conformance` leg, with overflow checks
|
||
//! on and off.
|
||
#![cfg(all(feature = "neural_int", feature = "legacy_int"))]
|
||
|
||
use serde_json::json;
|
||
use wai::int_kernels::{with_policy, Policy};
|
||
|
||
static POLICIES: std::sync::LazyLock<Vec<Policy>> = std::sync::LazyLock::new(wai::int_kernels::equivalence_policies);
|
||
|
||
/// 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
|
||
}
|
||
|
||
/// 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": 12, "gf": 16, "shift": 0, "layers": [layer] }))
|
||
}
|
||
|
||
/// A one-layer int_hyper model run as the int_hyper decodes run their networks: lowered by
|
||
/// `whm1/1`, loaded under `acc64`, and run by the integer operator set with this crate's
|
||
/// fast kernels behind it (`int_ops`). The output and its shape `(c, h, w)`, or the
|
||
/// refusal's code: a field past 32 bits, or a size the run-time plan caps.
|
||
fn hyper_run(model: &wai::int_hyper_synth::IntModel, x: &[i64], c: usize, h: usize, w: usize) -> Result<(Vec<i64>, usize, usize, usize), wai_ops::Code> {
|
||
let net = wai::int_ops::whm1_net(model).map_err(|r| r.code)?;
|
||
let l = wai_ops::lower::whm1_net_loaded(net, false, wai_ops::Profile::Acc64).map_err(|r| r.code)?;
|
||
let t = wai::int_ops::run1(&l, wai_ops::Tensor::new(c as u32, h as u32, w as u32, x.to_vec()), &[]).map_err(|r| r.code)?;
|
||
Ok((t.data, t.c as usize, t.h as usize, t.w as usize))
|
||
}
|
||
|
||
/// 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}"), || hyper_run(&model, &[5], 1, 1, 1));
|
||
if s <= u32::MAX as usize {
|
||
assert_eq!(got, Ok((vec![25], 1, 1, 1)), "int_hyper convT, stride {s}");
|
||
} else {
|
||
// A WHM1 field is 32 bits: a model past it is not a WHM1 network.
|
||
assert_eq!(got, Err(wai_ops::Code::ModelMismatch), "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)
|
||
});
|
||
if s <= u32::MAX as usize {
|
||
// centre taps 5, 14, 23 of the three output channels, times 5
|
||
assert_eq!(got, Ok(vec![25, 70, 115]), "int_synth, stride {s}");
|
||
} else {
|
||
// int_synth decodes through its lowering, whose `WHG1` fields are 32 bits: a
|
||
// stride past them is refused, a planned break.
|
||
assert!(got.is_err(), "int_synth, stride {s}: a stride past 32 bits is refused");
|
||
}
|
||
}
|
||
}
|
||
|
||
#[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}"), || hyper_run(&model, &x, 1, 5, 5));
|
||
if s <= u32::MAX as usize {
|
||
assert_eq!(got, Ok((vec![want], 1, 1, 1)), "int_hyper conv, pad {pad}, stride {s}");
|
||
} else {
|
||
assert_eq!(got, Err(wai_ops::Code::ModelMismatch), "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). int_synth decodes
|
||
// through its lowering, whose run-time plan caps the uncropped output a transposed
|
||
// convolution forms at 2^26 elements (integer-operators §5.3); this one is
|
||
// 3·8193·8193, so every policy refuses it before anything is allocated, a planned
|
||
// break. With pad 2^9 and output padding 2^10 − 2 the uncropped output, 3·1025·1025,
|
||
// is within the cap, and the decode gives 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!(got.is_err(), "int_synth, pad 2^12: the uncropped output is past the cap");
|
||
let q = 1usize << 9;
|
||
let (payload, spec) = (synth_payload(5), synth_spec(1, q, 2 * q - 2, [10, 20, 30]));
|
||
let got = same_under_every_policy("int_synth, pad 2^9", || {
|
||
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", || hyper_run(&model, &[5], 1, 1, 1));
|
||
assert_eq!(got, Ok((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. With overflow
|
||
// checks on, the reference panics on that arithmetic. With them off it wraps and then
|
||
// indexes out of the wrapped buffer. (int_synth decodes through its lowering, whose
|
||
// plan forms every size in checked 64-bit arithmetic and caps it: it refuses these,
|
||
// below.) 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 through its lowering's graph (`int_audio` unit tests).
|
||
|
||
/// 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 a WHM1 layer's lowering, as the int_hyper decodes run it. The
|
||
// scalar kernel they once ran panicked on both; through the operator set, every policy
|
||
// refuses both before anything is allocated: 2^62 is past a WHM1 field's 32 bits, and
|
||
// the u32 stride and pad give an uncropped side, (2 − 1)·(2^32 − 1) + 1, past the cap.
|
||
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", || hyper_run(&model, &X, 1, 3, 3));
|
||
assert_eq!(got, Some(Err(wai_ops::Code::ModelMismatch)), "refused, not panicked");
|
||
}
|
||
let model = convt(S32, P32);
|
||
let got = same_outcome_under_every_policy("int_hyper convT, 2×2, u32 stride and pad", || hyper_run(&model, &X[..4], 1, 2, 2));
|
||
assert!(matches!(got, Some(Err(_))), "u32 fields: refused by the cap, not panicked: {got:?}");
|
||
}
|
||
|
||
#[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 direct
|
||
// decode computed (2 − 1)·2^62 + 3 + 2^62 − 2·2^62 in `i64`, which overflows
|
||
// (2^63 + 3) although every `usize` step fits: it panicked with overflow checks on
|
||
// and, with them off, returned the wrapped size's output. int_synth decodes through
|
||
// its lowering (integer-operators §13.1), whose `WHG1` fields are 32 bits: in every
|
||
// build and under every policy the decode is refused, not panicked. A planned break.
|
||
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)
|
||
});
|
||
assert!(matches!(got, Some(Err(_))), "refused, not panicked: {got:?}");
|
||
|
||
// The u32 fields fit `WHG1`, but the uncropped side (2 − 1)·(2^32 − 1) + 3 is past
|
||
// the run-time plan's cap of 2^26: refused before anything is allocated (the direct
|
||
// decode returned the bias), a planned break.
|
||
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!(matches!(got, Some(Err(_))), "u32 fields: refused by the cap: {got:?}");
|
||
}
|