c7daf86787
int_payload::DeclaredCost gains macs: the multiply-accumulates the decode's networks comprise, layer by layer, from the same plan that sizes their buffers (int_hyper_synth::plan_net, the audio synthesis's plan): a convolution's every output times c·k·k, a transposed convolution's every input times out·k·k (out·k in the audio synthesis), an inverse GDN's every element times c. It is exact, summed in u128, and 0 for a WIV1 with keyframe 0. A kernel may skip a zero term, so it performs at most this many. wai-web's intPayloadCost gives it as a BigInt (macs). integer-payloads §7 states the definition. Tests: int_hyper_synth's the_planned_macs_are_every_layers_terms (200 random networks, each term counted at the shapes apply_layer forms); int_audio's the_planned_macs_are_every_layers_terms (the committed synthesis, at the lengths the reference kernel gives); int_payload_cost (0 for keyframe 0, more for every network); wai-web's the_declared_cost_is_read_before_decoding. Co-Authored-By: Claude Opus 5.5 <[email protected]>
1781 lines
89 KiB
Rust
1781 lines
89 KiB
Rust
//! Integer-exact synthesis of a pretrained learned image codec (per-channel int12
|
||
//! weights + integer IGDN) — the decode side of a `NeuralIntegerExact`
|
||
//! conditional codec, run on the pretrained mbt2018_mean mean-scale hyperprior
|
||
//! exported to integers. Pure i64, no float, no `ort`: the same
|
||
//! discipline as `int_synth`/`int_entropy`, so the reconstruction is
|
||
//! byte-identical on every machine.
|
||
//!
|
||
//! It runs the model the export tool emits (`tools/wai_sota_int_synth_export.py`):
|
||
//! a stack of integer transposed-convolutions (per-output-channel `(mult,shift)`
|
||
//! requant) and inverse-GDN activations ([`crate::int_transform::igdn`]).
|
||
//!
|
||
//! Quality is image-dependent, so it is measured on a pinned image rather than
|
||
//! quoted. On the public-domain test photograph `tools/wai_sota_vectors.sh`
|
||
//! regenerates the vectors from (a 256×256 crop of scikit-image's `astronaut`),
|
||
//! measured on an arm64 macOS machine: the integer decode of the pretrained
|
||
//! `mbt2018_mean`'s full conditional bitstream scores 35.985 dB (8-bit PSNR of this
|
||
//! decoder's output against the source pixels, `tests/int_hyper_synth.rs`), against
|
||
//! 36.175 dB for the same model's float decode (the export tool's PSNR of its float
|
||
//! reconstruction) — a cost of 0.19 dB. CI (`wai-sota-vectors.yml`) requires that
|
||
//! cost to stay within 1 dB on every run. It is a relative bar on one image and does
|
||
//! not compare the model with any other codec. An earlier development figure of
|
||
//! 42.6 dB (43.25 dB float), for the integer synthesis alone, was measured on a
|
||
//! photograph that is not in the repository and cannot be reproduced from it.
|
||
|
||
use serde_json::Value;
|
||
|
||
use crate::int_entropy::det_build_index;
|
||
use crate::int_transform::igdn;
|
||
use crate::rans::RansDecoder;
|
||
|
||
fn u(v: &Value, k: &str) -> usize {
|
||
v[k].as_u64().unwrap_or(0) as usize
|
||
}
|
||
fn i64s(v: &Value) -> Vec<i64> {
|
||
v.as_array().map(|a| a.iter().filter_map(|n| n.as_i64()).collect()).unwrap_or_default()
|
||
}
|
||
|
||
/// One integer transposed conv + per-output-channel requant
|
||
/// `((raw·mult)>>shift) + bias`. `x` is `[cin, h·w]`; `weight` is `[cin][cout][k][k]`.
|
||
/// Returns `([cout, ho·wo], ho, wo)`. Runs [`reference::conv_t`] or the fast kernel
|
||
/// ([`crate::int_kernels`]) per the current [`crate::int_kernels::Policy`]; both give
|
||
/// the same integer for every output element.
|
||
#[allow(clippy::too_many_arguments)]
|
||
fn conv_t(
|
||
x: &[i64], cin: usize, h: usize, w: usize, weight: &[i64], cout: usize, k: usize,
|
||
stride: usize, pad: usize, opad: usize, mult: &[i64], bias: &[i64], shift: u32,
|
||
) -> (Vec<i64>, usize, usize) {
|
||
use crate::int_kernels::{conv_acc_i64, policy, ConvGeom, KernelPath};
|
||
if policy().path == KernelPath::Fast && mult.len() >= cout && bias.len() >= cout
|
||
&& let Some(g) = ConvGeom::transposed(cin, h, w, cout, k, k, stride, stride, pad, pad, opad, opad, x.len(), weight.len())
|
||
&& let Some(out) = conv_acc_i64(&g, x, weight, |oc, raw| ((raw as i128 * mult[oc] as i128) >> shift) as i64 + bias[oc])
|
||
{
|
||
return (out, g.ho, g.wo);
|
||
}
|
||
reference::conv_t(x, cin, h, w, weight, cout, k, stride, pad, opad, mult, bias, shift)
|
||
}
|
||
|
||
/// One integer **regular** conv (stride/pad, no output-pad) + per-output-channel
|
||
/// requant — see [`reference::conv`]. Dispatches like [`conv_t`].
|
||
#[allow(clippy::too_many_arguments)]
|
||
fn conv(
|
||
x: &[i64], cin: usize, h: usize, w: usize, weight: &[i64], cout: usize, k: usize,
|
||
stride: usize, pad: usize, mult: &[i64], bias: &[i64], shift: u32,
|
||
) -> (Vec<i64>, usize, usize) {
|
||
use crate::int_kernels::{conv_acc_i64, policy, ConvGeom, KernelPath};
|
||
if policy().path == KernelPath::Fast && mult.len() >= cout && bias.len() >= cout
|
||
&& let Some(g) = ConvGeom::conv(cin, h, w, cout, 1, k, k, stride, stride, pad, pad, x.len(), weight.len())
|
||
&& let Some(out) = conv_acc_i64(&g, x, weight, |oc, raw| ((raw as i128 * mult[oc] as i128) >> shift) as i64 + bias[oc])
|
||
{
|
||
return (out, g.ho, g.wo);
|
||
}
|
||
reference::conv(x, cin, h, w, weight, cout, k, stride, pad, mult, bias, shift)
|
||
}
|
||
|
||
/// The scalar reference kernels — the executable specification the fast kernels
|
||
/// in [`crate::int_kernels`] are checked against, kept verbatim. Selected for the
|
||
/// whole decode by `int_kernels::with_policy(Policy::reference(), ..)`.
|
||
#[doc(hidden)]
|
||
pub mod reference {
|
||
/// One integer transposed conv + per-output-channel requant `((raw·mult)>>shift)
|
||
/// + bias`. `x` is `[cin, h·w]`; `weight` is `[cin][cout][k][k]`. Returns
|
||
/// `([cout, ho·wo], ho, wo)`. Pure i64 (the accumulation stays well inside i64).
|
||
#[allow(clippy::too_many_arguments)]
|
||
pub fn conv_t(
|
||
x: &[i64], cin: usize, h: usize, w: usize, weight: &[i64], cout: usize, k: usize,
|
||
stride: usize, pad: usize, opad: usize, mult: &[i64], bias: &[i64], shift: u32,
|
||
) -> (Vec<i64>, usize, usize) {
|
||
let ho = (h - 1) * stride + k + opad - 2 * pad;
|
||
let wo = (w - 1) * stride + k + opad - 2 * pad;
|
||
let (hp, wp) = (ho + 2 * pad, wo + 2 * pad);
|
||
let kk = k * k;
|
||
let mut raw = vec![0i64; cout * hp * wp];
|
||
// Loop order ic → oc → spatial → kernel: for a fixed (ic, oc) the k×k weight
|
||
// block is loaded once and every scatter lands in the single `oc` output plane,
|
||
// so that plane (and the input row) stays hot in cache. Integer accumulation is
|
||
// order-independent, so this is byte-identical to the naive oc-inner order.
|
||
for ic in 0..cin {
|
||
let xc = &x[ic * h * w..ic * h * w + h * w];
|
||
for oc in 0..cout {
|
||
let wb = (ic * cout + oc) * kk;
|
||
let wk = &weight[wb..wb + kk];
|
||
let plane = oc * hp * wp;
|
||
for iy in 0..h {
|
||
let ry0 = iy * stride;
|
||
let xrow = &xc[iy * w..iy * w + w];
|
||
for (ix, &v) in xrow.iter().enumerate() {
|
||
if v == 0 {
|
||
continue;
|
||
}
|
||
let base = plane + ry0 * wp + ix * stride;
|
||
for ky in 0..k {
|
||
let dst = base + ky * wp;
|
||
let wky = ky * k;
|
||
for kx in 0..k {
|
||
raw[dst + kx] += wk[wky + kx] * v;
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
let mut out = vec![0i64; cout * ho * wo];
|
||
for oc in 0..cout {
|
||
let (m, b) = (mult[oc] as i128, bias[oc]);
|
||
for oy in 0..ho {
|
||
for ox in 0..wo {
|
||
let r = raw[(oc * hp + oy + pad) * wp + ox + pad] as i128;
|
||
out[(oc * ho + oy) * wo + ox] = ((r * m) >> shift) as i64 + b;
|
||
}
|
||
}
|
||
}
|
||
(out, ho, wo)
|
||
}
|
||
|
||
/// One integer **regular** conv (stride/pad, no output-pad) + per-output-channel
|
||
/// requant — the stride-1 `Conv2d` that closes the hyper-decoder (`M·3/2 → M·2`,
|
||
/// emitting σ‖μ). `weight` is `[cout][cin][k][k]`. Mirrors the export's
|
||
/// `conv_int` byte-for-byte. Pure i64 (i128 only for the requant multiply).
|
||
#[allow(clippy::too_many_arguments)]
|
||
pub fn conv(
|
||
x: &[i64], cin: usize, h: usize, w: usize, weight: &[i64], cout: usize, k: usize,
|
||
stride: usize, pad: usize, mult: &[i64], bias: &[i64], shift: u32,
|
||
) -> (Vec<i64>, usize, usize) {
|
||
let ho = (h + 2 * pad - k) / stride + 1;
|
||
let wo = (w + 2 * pad - k) / stride + 1;
|
||
let kk = k * k;
|
||
let mut out = vec![0i64; cout * ho * wo]; // raw sums first; requantized below
|
||
// Loop order oc → ic → kernel → spatial: the weight tap is a hoisted scalar
|
||
// and the inner `ox` loop streams one output row against one input row (great
|
||
// locality, vectorizable). For stride 1 the valid `ox` span is precomputed so
|
||
// the inner loop is branch-free. Integer accumulation is order-independent, so
|
||
// this is byte-identical to the naive output-stationary order.
|
||
for oc in 0..cout {
|
||
let plane = oc * ho * wo;
|
||
for ic in 0..cin {
|
||
let xb = ic * h * w;
|
||
let wb = (oc * cin + ic) * kk;
|
||
for ky in 0..k {
|
||
for kx in 0..k {
|
||
let wv = weight[wb + ky * k + kx];
|
||
if wv == 0 {
|
||
continue;
|
||
}
|
||
// valid output columns: input ix = ox·stride + kx − pad ∈ [0, w)
|
||
let (ox_lo, ox_hi) = if stride == 1 {
|
||
(pad.saturating_sub(kx), (w + pad).saturating_sub(kx).min(wo))
|
||
} else {
|
||
(0, wo)
|
||
};
|
||
for oy in 0..ho {
|
||
let iy = oy * stride + ky;
|
||
if iy < pad || iy >= pad + h {
|
||
continue;
|
||
}
|
||
let xrow = xb + (iy - pad) * w;
|
||
let orow = plane + oy * wo;
|
||
if stride == 1 {
|
||
// over [ox_lo, ox_hi): ox + kx ≥ pad, so ox+kx−pad ≥ 0.
|
||
for ox in ox_lo..ox_hi {
|
||
out[orow + ox] += wv * x[xrow + ox + kx - pad];
|
||
}
|
||
} else {
|
||
for ox in 0..wo {
|
||
let ix = ox * stride + kx;
|
||
if ix < pad || ix >= pad + w {
|
||
continue;
|
||
}
|
||
out[orow + ox] += wv * x[xrow + ix - pad];
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
for oc in 0..cout {
|
||
let (m, b) = (mult[oc] as i128, bias[oc]);
|
||
for v in &mut out[oc * ho * wo..(oc + 1) * ho * wo] {
|
||
*v = ((*v as i128 * m) >> shift) as i64 + b;
|
||
}
|
||
}
|
||
(out, ho, wo)
|
||
}
|
||
}
|
||
|
||
/// A parsed integer conv net — the typed form [`run_net`] consumes, built from
|
||
/// either the export's JSON ([`IntModel::from_json`]) or its **compact binary**
|
||
/// ([`IntModel::from_bytes`]): the binary stores the int12 weights as `i16` and
|
||
/// IGDN `gamma` as `i32`, roughly halving the model and making the in-browser
|
||
/// download practical. The decode is byte-identical whichever the source.
|
||
pub struct IntModel {
|
||
pub afrac: u32,
|
||
pub gf: u32,
|
||
pub shift: u32,
|
||
pub layers: Vec<IntLayer>,
|
||
}
|
||
|
||
/// One layer of an [`IntModel`]. Weights/params are widened to `i64` for the
|
||
/// decode (the arithmetic stays well inside i64); `in_ch` is not stored — the
|
||
/// running channel count threads through [`run_net`].
|
||
pub enum IntLayer {
|
||
ConvT { out_ch: usize, k: usize, stride: usize, pad: usize, opad: usize, weight: Vec<i64>, mult: Vec<i64>, bias: Vec<i64> },
|
||
Conv { out_ch: usize, k: usize, stride: usize, pad: usize, weight: Vec<i64>, mult: Vec<i64>, bias: Vec<i64> },
|
||
Igdn { gamma: Vec<i64>, beta: Vec<i64> },
|
||
Relu,
|
||
Leaky { slope_q: i64, slope_frac: u32 },
|
||
}
|
||
|
||
/// Little-endian cursor over the binary model, bounds-checked: every read is of bytes
|
||
/// that are present, with its end formed in checked arithmetic, so a damaged model is
|
||
/// refused and never makes the parser allocate more than the bytes it holds.
|
||
struct Cur<'a> {
|
||
b: &'a [u8],
|
||
p: usize,
|
||
}
|
||
impl<'a> Cur<'a> {
|
||
/// The next `n` elements of `size` bytes each.
|
||
fn take(&mut self, n: usize, size: usize) -> Result<&'a [u8], &'static str> {
|
||
let e = n.checked_mul(size).and_then(|len| self.p.checked_add(len)).ok_or("model truncated")?;
|
||
let s = self.b.get(self.p..e).ok_or("model truncated")?;
|
||
self.p = e;
|
||
Ok(s)
|
||
}
|
||
fn u32(&mut self) -> Result<usize, &'static str> {
|
||
Ok(u32::from_le_bytes(self.take(1, 4)?.try_into().unwrap()) as usize)
|
||
}
|
||
fn u8(&mut self) -> Result<u8, &'static str> {
|
||
Ok(self.take(1, 1)?[0])
|
||
}
|
||
fn i16v(&mut self, n: usize) -> Result<Vec<i64>, &'static str> {
|
||
Ok(self.take(n, 2)?.chunks_exact(2).map(|c| i16::from_le_bytes([c[0], c[1]]) as i64).collect())
|
||
}
|
||
fn i32v(&mut self, n: usize) -> Result<Vec<i64>, &'static str> {
|
||
Ok(self.take(n, 4)?.chunks_exact(4).map(|c| i32::from_le_bytes([c[0], c[1], c[2], c[3]]) as i64).collect())
|
||
}
|
||
fn i64v(&mut self, n: usize) -> Result<Vec<i64>, &'static str> {
|
||
Ok(self.take(n, 8)?.chunks_exact(8).map(|c| i64::from_le_bytes(c.try_into().unwrap())).collect())
|
||
}
|
||
}
|
||
|
||
/// The product of a layer's dimensions, refused where it overflows `usize`.
|
||
fn count(dims: &[usize]) -> Result<usize, &'static str> {
|
||
dims.iter().try_fold(1usize, |a, &d| a.checked_mul(d)).ok_or("layer size overflows")
|
||
}
|
||
|
||
impl IntModel {
|
||
/// [`Self::from_json`], refusing rather than panicking on any JSON. It reads what
|
||
/// `from_json` reads, an absent number as 0 and an absent list as empty, as that
|
||
/// does, but refuses a member of the wrong type, a size that does not fit (`u32`
|
||
/// for `afrac`, `gf`, `shift` and a shift, `usize` for the rest), and a layer type
|
||
/// it does not know. A layer's parameters are checked against the running shapes
|
||
/// when the network is planned ([`decode_conditional_checked`]).
|
||
pub fn try_from_json(v: &Value) -> Result<Self, &'static str> {
|
||
let word = |o: &Value, k: &'static str| match o.get(k) {
|
||
None => Ok(0),
|
||
Some(n) => n.as_u64().and_then(|n| u32::try_from(n).ok()).ok_or(k),
|
||
};
|
||
let size = |o: &Value, k: &'static str| match o.get(k) {
|
||
None => Ok(0),
|
||
Some(n) => n.as_u64().and_then(|n| usize::try_from(n).ok()).ok_or(k),
|
||
};
|
||
let ints = |o: &Value, k: &'static str| -> Result<Vec<i64>, &'static str> {
|
||
match o.get(k) {
|
||
None => Ok(Vec::new()),
|
||
Some(a) => a.as_array().ok_or(k)?.iter().map(Value::as_i64).collect::<Option<_>>().ok_or(k),
|
||
}
|
||
};
|
||
let layers = v.get("layers").and_then(Value::as_array).ok_or("layers")?.iter().map(|l| {
|
||
Ok(match l.get("type").and_then(Value::as_str).ok_or("a layer's type")? {
|
||
"convT" => IntLayer::ConvT {
|
||
out_ch: size(l, "out_ch")?, k: size(l, "k")?, stride: size(l, "stride")?, pad: size(l, "pad")?,
|
||
opad: size(l, "opad")?, weight: ints(l, "weight")?, mult: ints(l, "mult")?, bias: ints(l, "bias")?,
|
||
},
|
||
"conv" => IntLayer::Conv {
|
||
out_ch: size(l, "out_ch")?, k: size(l, "k")?, stride: size(l, "stride")?, pad: size(l, "pad")?,
|
||
weight: ints(l, "weight")?, mult: ints(l, "mult")?, bias: ints(l, "bias")?,
|
||
},
|
||
"igdn" => IntLayer::Igdn { gamma: ints(l, "gamma")?, beta: ints(l, "beta")? },
|
||
"relu" => IntLayer::Relu,
|
||
"leaky" => IntLayer::Leaky {
|
||
slope_q: match l.get("slope_q") { None => 655, Some(q) => q.as_i64().ok_or("slope_q")? },
|
||
slope_frac: word(l, "slope_frac")?,
|
||
},
|
||
_ => return Err("an unknown layer type"),
|
||
})
|
||
}).collect::<Result<Vec<_>, &'static str>>()?;
|
||
Ok(Self { afrac: word(v, "afrac")?, gf: word(v, "gf")?, shift: word(v, "shift")?, layers })
|
||
}
|
||
|
||
/// Parse the export tool's JSON model (`int_synth_sota_model.json` / `int_hs_model.json`).
|
||
/// Panics on a JSON the tool would not write; [`Self::try_from_json`] refuses
|
||
/// instead.
|
||
pub fn from_json(v: &Value) -> Self {
|
||
let layers = v["layers"]
|
||
.as_array()
|
||
.unwrap()
|
||
.iter()
|
||
.map(|l| match l["type"].as_str().unwrap() {
|
||
"convT" => IntLayer::ConvT {
|
||
out_ch: u(l, "out_ch"), k: u(l, "k"), stride: u(l, "stride"), pad: u(l, "pad"),
|
||
opad: u(l, "opad"), weight: i64s(&l["weight"]), mult: i64s(&l["mult"]), bias: i64s(&l["bias"]),
|
||
},
|
||
"conv" => IntLayer::Conv {
|
||
out_ch: u(l, "out_ch"), k: u(l, "k"), stride: u(l, "stride"), pad: u(l, "pad"),
|
||
weight: i64s(&l["weight"]), mult: i64s(&l["mult"]), bias: i64s(&l["bias"]),
|
||
},
|
||
"igdn" => IntLayer::Igdn { gamma: i64s(&l["gamma"]), beta: i64s(&l["beta"]) },
|
||
"relu" => IntLayer::Relu,
|
||
"leaky" => IntLayer::Leaky { slope_q: l["slope_q"].as_i64().unwrap_or(655), slope_frac: u(l, "slope_frac") as u32 },
|
||
other => panic!("unknown int_hyper_synth layer {other}"),
|
||
})
|
||
.collect();
|
||
Self { afrac: u(v, "afrac") as u32, gf: u(v, "gf") as u32, shift: u(v, "shift") as u32, layers }
|
||
}
|
||
|
||
/// Parse the export tool's **compact binary** model (`*_model.bin`): magic
|
||
/// `WHM1`, then `afrac/gf/shift/n_layers` (u32) and per layer a `u8` tag —
|
||
/// 0 convT, 1 conv, 2 igdn, 3 relu, 4 leaky — with int12 weights as `i16`,
|
||
/// IGDN `gamma` as `i32`, the rest `i64`, and nothing after the last layer.
|
||
/// Decodes identically to [`from_json`].
|
||
pub fn from_bytes(b: &[u8]) -> Result<Self, &'static str> {
|
||
if b.get(0..4) != Some(b"WHM1") {
|
||
return Err("bad model magic (want WHM1)");
|
||
}
|
||
let mut c = Cur { b, p: 4 };
|
||
let (afrac, gf, shift) = (c.u32()? as u32, c.u32()? as u32, c.u32()? as u32);
|
||
let n = c.u32()?;
|
||
// Each layer is at least its tag byte, so a count beyond the bytes left is a
|
||
// truncated model; reserving for it would let four damaged bytes ask for an
|
||
// allocation of hundreds of gigabytes.
|
||
let mut layers = Vec::with_capacity(n.min(b.len() - c.p));
|
||
for _ in 0..n {
|
||
layers.push(match c.u8()? {
|
||
0 => {
|
||
let inc = c.u32()?;
|
||
let out = c.u32()?;
|
||
let (k, stride, pad, opad) = (c.u32()?, c.u32()?, c.u32()?, c.u32()?);
|
||
IntLayer::ConvT { out_ch: out, k, stride, pad, opad,
|
||
weight: c.i16v(count(&[inc, out, k, k])?)?, mult: c.i64v(out)?, bias: c.i64v(out)? }
|
||
}
|
||
1 => {
|
||
let inc = c.u32()?;
|
||
let out = c.u32()?;
|
||
let (k, stride, pad) = (c.u32()?, c.u32()?, c.u32()?);
|
||
IntLayer::Conv { out_ch: out, k, stride, pad,
|
||
weight: c.i16v(count(&[out, inc, k, k])?)?, mult: c.i64v(out)?, bias: c.i64v(out)? }
|
||
}
|
||
2 => {
|
||
let cc = c.u32()?;
|
||
IntLayer::Igdn { gamma: c.i32v(count(&[cc, cc])?)?, beta: c.i64v(cc)? }
|
||
}
|
||
3 => IntLayer::Relu,
|
||
4 => IntLayer::Leaky { slope_q: c.i64v(1)?[0], slope_frac: c.u32()? as u32 },
|
||
_ => return Err("bad layer tag"),
|
||
});
|
||
}
|
||
// Nothing follows the last layer (`extensions/integer-payloads.md` §9), so one
|
||
// network has one file.
|
||
if c.p != b.len() {
|
||
return Err("bytes follow the last layer");
|
||
}
|
||
Ok(Self { afrac, gf, shift, layers })
|
||
}
|
||
}
|
||
|
||
/// Run an integer conv net (`model`: `afrac`/`gf`/`shift` + a `layers` stack of
|
||
/// `convT` / `conv` / `igdn` / `relu`) on `input` (`[c, h, w]` at `2^afrac`). Returns the
|
||
/// raw integer feature map `([c·h·w], c, h, w)` at `2^afrac`. Shared by both the
|
||
/// synthesis `g_s` (→ pixels via [`decode`]) and the hyper-decoder `h_s`
|
||
/// (→ σ‖μ). Pure i64; byte-identical on every machine.
|
||
pub fn run_net(model: &IntModel, input: &[i64], mut c: usize, mut h: usize, mut w: usize) -> (Vec<i64>, usize, usize, usize) {
|
||
let mut x = input.to_vec();
|
||
for layer in &model.layers {
|
||
(x, c, h, w) = apply_layer(model, layer, x, c, h, w);
|
||
}
|
||
(x, c, h, w)
|
||
}
|
||
|
||
/// One layer of [`run_net`] on `x` (`[c, h, w]`): the new activations and shape.
|
||
/// [`run_net_checked`] runs layers through this same function, so the two cannot
|
||
/// compute different integers.
|
||
fn apply_layer(model: &IntModel, layer: &IntLayer, mut x: Vec<i64>, c: usize, h: usize, w: usize) -> (Vec<i64>, usize, usize, usize) {
|
||
let (afrac, gf, shift) = (model.afrac, model.gf, model.shift);
|
||
match layer {
|
||
IntLayer::ConvT { out_ch, k, stride, pad, opad, weight, mult, bias } => {
|
||
let (nx, ho, wo) = conv_t(&x, c, h, w, weight, *out_ch, *k, *stride, *pad, *opad, mult, bias, shift);
|
||
(nx, *out_ch, ho, wo)
|
||
}
|
||
IntLayer::Conv { out_ch, k, stride, pad, weight, mult, bias } => {
|
||
let (nx, ho, wo) = conv(&x, c, h, w, weight, *out_ch, *k, *stride, *pad, mult, bias, shift);
|
||
(nx, *out_ch, ho, wo)
|
||
}
|
||
IntLayer::Igdn { gamma, beta } => (igdn(&x, c, h * w, afrac, gamma, beta, gf), c, h, w),
|
||
IntLayer::Relu => {
|
||
for v in &mut x {
|
||
if *v < 0 {
|
||
*v = 0;
|
||
}
|
||
}
|
||
(x, c, h, w)
|
||
}
|
||
IntLayer::Leaky { slope_q, slope_frac } => {
|
||
for v in &mut x {
|
||
if *v < 0 {
|
||
*v = (*v * *slope_q) >> *slope_frac;
|
||
}
|
||
}
|
||
(x, c, h, w)
|
||
}
|
||
}
|
||
}
|
||
|
||
/// Why [`run_net_checked`] declined to run a network: the index of the first layer
|
||
/// it would not run, the check that failed, and whether that check was about the
|
||
/// values presented (`out_of_range`: the layer is well formed, but these magnitudes
|
||
/// could overflow its arithmetic) or about the layer itself (malformed, or not
|
||
/// matching the input's shape).
|
||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||
pub struct NetRefusal {
|
||
pub layer: usize,
|
||
pub reason: &'static str,
|
||
pub out_of_range: bool,
|
||
}
|
||
|
||
impl std::fmt::Display for NetRefusal {
|
||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||
write!(f, "layer {}: {}", self.layer, self.reason)
|
||
}
|
||
}
|
||
|
||
/// [`run_net`], run only while every step stays exact.
|
||
///
|
||
/// Before each layer it checks, on the activations actually present, that:
|
||
///
|
||
/// - the geometry is well formed, and every size the layer computes fits `usize`
|
||
/// (a transposed convolution's uncropped plane included);
|
||
/// - the weight, requantisation and IGDN buffers have exactly the lengths the layer
|
||
/// reads for the running channel count;
|
||
/// - every shift is below its operand's width;
|
||
/// - no product, partial sum or requantised value can leave the type the reference
|
||
/// kernel forms it in. For a convolution with `terms = c·k·k` inputs per output,
|
||
/// `terms · max|w| · max|x|` bounds every product and partial sum (each output
|
||
/// element adds at most `terms` products, whatever their order) and must fit `i64`;
|
||
/// the requantised value is at most `(bound · max|mult|) >> shift` plus one (the
|
||
/// floor of a negative) plus `max|bias|`, and must fit `i64` too. For IGDN, with
|
||
/// `X = max|x|`, `G = max|gamma|`, `B = max|beta|`: `G·X²` and
|
||
/// `B + c·((G·X²) >> gf + 1)` must fit `i128`, its integer square root (which the
|
||
/// reference narrows to `i64`) must fit `i64`, and the output `(X·root) >> afrac`
|
||
/// plus one must fit `i64`. For the leaky ReLU, `max|x| · |slope|` must fit `i64`.
|
||
///
|
||
/// A layer that fails a check is not run, and the refusal says which and why. On
|
||
/// `Ok` the output is [`run_net`]'s, computed by the same code ([`apply_layer`], which
|
||
/// dispatches to the fast kernels or the reference as `run_net` does), and no step
|
||
/// overflowed, so a build with overflow checks and one without compute the same
|
||
/// integers and neither panics. The fast kernels, which narrow only under their own
|
||
/// bounds, form no value the reference does not (see [`crate::int_kernels`]).
|
||
///
|
||
/// The checks are one layer deep, on real values, so they refuse only activations
|
||
/// that come within a factor `terms` (or, for IGDN, `c`) of the type's range, not
|
||
/// a network whose worst case over all inputs would. A decoder that feeds values it
|
||
/// read from an untrusted stream into a network uses this, and refuses the stream on
|
||
/// `Err`.
|
||
pub fn run_net_checked(
|
||
model: &IntModel, input: &[i64], mut c: usize, mut h: usize, mut w: usize,
|
||
) -> Result<(Vec<i64>, usize, usize, usize), NetRefusal> {
|
||
let n = c.checked_mul(h).and_then(|v| v.checked_mul(w));
|
||
if c == 0 || h == 0 || w == 0 || n != Some(input.len()) {
|
||
return Err(NetRefusal { layer: 0, reason: "input length is not channels × height × width", out_of_range: false });
|
||
}
|
||
let mut x = input.to_vec();
|
||
for (i, layer) in model.layers.iter().enumerate() {
|
||
let out = check_layer(model, layer, &x, c, h, w)
|
||
.map_err(|(reason, out_of_range)| NetRefusal { layer: i, reason, out_of_range })?;
|
||
let (nx, nc, nh, nw) = apply_layer(model, layer, x, c, h, w);
|
||
debug_assert_eq!((nc, nh, nw), out, "check_layer and the layer disagree on the output shape");
|
||
(x, c, h, w) = (nx, nc, nh, nw);
|
||
}
|
||
Ok((x, c, h, w))
|
||
}
|
||
|
||
/// A check [`check_layer`] failed: its reason, and whether it concerned the values
|
||
/// presented (`true`) or the layer itself (`false`).
|
||
pub(crate) type Failed = (&'static str, bool);
|
||
const fn shape(reason: &'static str) -> Failed {
|
||
(reason, false)
|
||
}
|
||
const fn range(reason: &'static str) -> Failed {
|
||
(reason, true)
|
||
}
|
||
|
||
/// One layer's geometry on a `[c, h, w]` input, from its parameters alone: the output
|
||
/// shape, and the most elements the reference kernel holds in one buffer while
|
||
/// running it — for a transposed convolution the uncropped plane it accumulates into
|
||
/// before cropping, otherwise the output. Every size is formed with checked
|
||
/// arithmetic; a geometry that is malformed or overflows `usize` is refused. Shared
|
||
/// by [`check_layer`] and [`net_shape`], so the shape a caller plans for is the shape
|
||
/// the layer runs at.
|
||
fn layer_geometry(layer: &IntLayer, c: usize, h: usize, w: usize) -> Result<((usize, usize, usize), usize), Failed> {
|
||
let plane = |c: usize, h: usize, w: usize| c.checked_mul(h).and_then(|v| v.checked_mul(w));
|
||
match layer {
|
||
IntLayer::ConvT { out_ch, k, stride, pad, opad, .. } => {
|
||
let (cout, k, s, p, op) = (*out_ch, *k, *stride, *pad, *opad);
|
||
if cout == 0 || k == 0 || s == 0 {
|
||
return Err(shape("zero output channels, kernel or stride"));
|
||
}
|
||
// The reference sizes the uncropped plane (h-1)·s + k + op per side and
|
||
// crops 2·pad from it; both must be formed without overflow and leave a
|
||
// positive output.
|
||
let full = |n: usize| (n - 1).checked_mul(s)?.checked_add(k)?.checked_add(op);
|
||
let (fh, fw) = (full(h).ok_or(shape("output size overflows"))?, full(w).ok_or(shape("output size overflows"))?);
|
||
let crop = p.checked_mul(2).ok_or(shape("padding overflows"))?;
|
||
if fh <= crop || fw <= crop {
|
||
return Err(shape("padding leaves no output"));
|
||
}
|
||
let uncropped = plane(cout, fh, fw).ok_or(shape("uncropped plane overflows"))?;
|
||
Ok(((cout, fh - crop, fw - crop), uncropped))
|
||
}
|
||
IntLayer::Conv { out_ch, k, stride, pad, .. } => {
|
||
let (cout, k, s, p) = (*out_ch, *k, *stride, *pad);
|
||
if cout == 0 || k == 0 || s == 0 {
|
||
return Err(shape("zero output channels, kernel or stride"));
|
||
}
|
||
let padded = |n: usize| n.checked_add(p.checked_mul(2)?);
|
||
let (hp, wp) = (padded(h).ok_or(shape("padding overflows"))?, padded(w).ok_or(shape("padding overflows"))?);
|
||
if hp < k || wp < k {
|
||
return Err(shape("kernel larger than the padded input"));
|
||
}
|
||
let (ho, wo) = ((hp - k) / s + 1, (wp - k) / s + 1);
|
||
let out = plane(cout, ho, wo).ok_or(shape("output size overflows"))?;
|
||
Ok(((cout, ho, wo), out))
|
||
}
|
||
IntLayer::Igdn { .. } | IntLayer::Relu | IntLayer::Leaky { .. } => {
|
||
Ok(((c, h, w), plane(c, h, w).ok_or(shape("input size overflows"))?))
|
||
}
|
||
}
|
||
}
|
||
|
||
/// What a network does to a `[c, h, w]` input's shape, from its layer geometry alone
|
||
/// (see [`net_shape`]).
|
||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||
pub struct NetShape {
|
||
/// The output `[c, h, w]`.
|
||
pub out: (usize, usize, usize),
|
||
/// The most elements any one buffer holds while the network runs on the reference
|
||
/// kernels: the input, each layer's output, and each transposed convolution's
|
||
/// uncropped plane. (The fast kernels allocate no buffer the reference does not,
|
||
/// except a padded copy of a layer's input, which [`crate::int_kernels`] holds to a
|
||
/// fixed multiple of that layer's input, weights and output.)
|
||
pub peak: usize,
|
||
}
|
||
|
||
/// [`NetShape`] for `model` on a `[c, h, w]` input, walking the layers' geometry
|
||
/// without allocating an activation or reading a value. Refuses what
|
||
/// [`run_net_checked`] would refuse on geometry — an empty input, a malformed layer, a
|
||
/// size that overflows `usize` — naming the same layer and reason; the parameter-length
|
||
/// and value checks are `run_net_checked`'s, made before each layer runs.
|
||
///
|
||
/// A decoder that sizes a network's input from an untrusted header calls this first,
|
||
/// and refuses before running anything a header whose shapes the network does not
|
||
/// produce, or on which it would form a buffer larger than the decoder allows.
|
||
pub fn net_shape(model: &IntModel, mut c: usize, mut h: usize, mut w: usize) -> Result<NetShape, NetRefusal> {
|
||
let refuse = |layer: usize, (reason, out_of_range): Failed| NetRefusal { layer, reason, out_of_range };
|
||
if c == 0 || h == 0 || w == 0 {
|
||
return Err(refuse(0, shape("input length is not channels × height × width")));
|
||
}
|
||
let mut peak = c.checked_mul(h).and_then(|v| v.checked_mul(w))
|
||
.ok_or_else(|| refuse(0, shape("input length is not channels × height × width")))?;
|
||
for (i, layer) in model.layers.iter().enumerate() {
|
||
let (out, largest) = layer_geometry(layer, c, h, w).map_err(|f| refuse(i, f))?;
|
||
peak = peak.max(largest);
|
||
(c, h, w) = out;
|
||
}
|
||
Ok(NetShape { out: (c, h, w), peak })
|
||
}
|
||
|
||
/// The checks [`run_net_checked`] documents, for one layer on activations `x`
|
||
/// (`[c, h, w]`). Returns the layer's output shape.
|
||
fn check_layer(model: &IntModel, layer: &IntLayer, x: &[i64], c: usize, h: usize, w: usize) -> Result<(usize, usize, usize), Failed> {
|
||
use crate::int_kernels::max_abs;
|
||
let (out, _) = layer_geometry(layer, c, h, w)?;
|
||
const I64: u128 = i64::MAX as u128;
|
||
let xm = max_abs(x) as u128;
|
||
// One (transposed) convolution's arithmetic: every product and partial sum, then
|
||
// the requantised value, must fit i64 (the reference's accumulator and result).
|
||
let conv_range = |cout: usize, k: usize, weight: &[i64], mult: &[i64], bias: &[i64]| -> Result<(), Failed> {
|
||
let w_len = c.checked_mul(cout).and_then(|v| v.checked_mul(k)).and_then(|v| v.checked_mul(k));
|
||
if w_len != Some(weight.len()) || mult.len() != cout || bias.len() != cout {
|
||
return Err(shape("parameter lengths do not match the running channel count"));
|
||
}
|
||
if model.shift > 127 {
|
||
return Err(shape("requantisation shift of 128 or more"));
|
||
}
|
||
// c·k·k ≤ weight.len() (checked above), but saturate regardless.
|
||
let terms = (c as u128).saturating_mul(k as u128).saturating_mul(k as u128);
|
||
let acc = terms.saturating_mul(max_abs(weight) as u128).saturating_mul(xm);
|
||
if acc > I64 {
|
||
return Err(range("accumulation could leave i64"));
|
||
}
|
||
// acc ≤ 2^63 and |mult| ≤ 2^63, so the product fits u128.
|
||
let q = ((acc * max_abs(mult) as u128) >> model.shift) + 1;
|
||
if q.saturating_add(max_abs(bias) as u128) > I64 {
|
||
return Err(range("requantised value could leave i64"));
|
||
}
|
||
Ok(())
|
||
};
|
||
match layer {
|
||
IntLayer::ConvT { out_ch, k, weight, mult, bias, .. } | IntLayer::Conv { out_ch, k, weight, mult, bias, .. } => {
|
||
conv_range(*out_ch, *k, weight, mult, bias)?;
|
||
Ok(out)
|
||
}
|
||
IntLayer::Igdn { gamma, beta } => {
|
||
igdn_check(x, c, gamma, beta, model.gf, model.afrac)?;
|
||
Ok((c, h, w))
|
||
}
|
||
IntLayer::Relu => Ok((c, h, w)),
|
||
IntLayer::Leaky { slope_q, slope_frac } => {
|
||
if *slope_frac > 63 {
|
||
return Err(shape("leaky slope shift of 64 or more"));
|
||
}
|
||
if xm.saturating_mul(slope_q.unsigned_abs() as u128) > I64 {
|
||
return Err(range("leaky product could leave i64"));
|
||
}
|
||
Ok((c, h, w))
|
||
}
|
||
}
|
||
}
|
||
|
||
/// The IGDN checks [`run_net_checked`] documents, for activations `x` of `c` channels
|
||
/// at `2^afrac` against `gamma` (`c × c`, at `2^gf`) and `beta` (`c`). Shared with the
|
||
/// 1-D synthesis of `int_audio`, whose IGDN is the same [`igdn`].
|
||
pub(crate) fn igdn_check(x: &[i64], c: usize, gamma: &[i64], beta: &[i64], gf: u32, afrac: u32) -> Result<(), Failed> {
|
||
use crate::int_kernels::max_abs;
|
||
const I64: u128 = i64::MAX as u128;
|
||
const I128: u128 = i128::MAX as u128;
|
||
if c.checked_mul(c) != Some(gamma.len()) || beta.len() != c {
|
||
return Err(shape("IGDN parameter lengths do not match the channel count"));
|
||
}
|
||
if gf > 127 || afrac > 127 {
|
||
return Err(shape("IGDN shift of 128 or more"));
|
||
}
|
||
let xm = max_abs(x) as u128;
|
||
// x² ≤ 2^126 always fits; each term gamma·x², and the norm, must fit i128.
|
||
let x2 = xm * xm;
|
||
let term = (max_abs(gamma) as u128).checked_mul(x2).filter(|&t| t <= I128).ok_or(range("IGDN term could leave i128"))?;
|
||
let norm = (max_abs(beta) as u128).saturating_add((c as u128).saturating_mul((term >> gf) + 1));
|
||
if norm > I128 {
|
||
return Err(range("IGDN norm could leave i128"));
|
||
}
|
||
// The reference narrows the root to i64 (`isqrt(norm) as i64`), so it must
|
||
// fit; then |x · root| < 2^63 · 2^63 fits u128, and the result must fit i64.
|
||
let root = norm.isqrt();
|
||
if root > I64 {
|
||
return Err(range("IGDN root could leave i64"));
|
||
}
|
||
if ((xm * root) >> afrac) + 1 > I64 {
|
||
return Err(range("IGDN output could leave i64"));
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
/// Decode the integer synthesis `model` on the latent `input` (`[c, h, w]` at
|
||
/// `2^afrac`) → RGB, row-major **HWC** `u8`. Byte-identical on every machine.
|
||
pub fn decode(model: &IntModel, input: &[i64], c: usize, h: usize, w: usize) -> (Vec<u8>, usize, usize) {
|
||
let (x, c, h, w) = run_net(model, input, c, h, w);
|
||
(to_rgb8(&x, c, h, w, model.afrac), h, w)
|
||
}
|
||
|
||
/// The integer pixel quantisation that ends a synthesis: each value clamped to
|
||
/// `[0, 2^afrac]`, scaled to 8 bits rounding half up, CHW → row-major HWC `u8`.
|
||
/// `afrac` must be in `1..=54`, so that `2^afrac · 255 + 2^(afrac-1)` fits `i64`.
|
||
pub fn to_rgb8(x: &[i64], c: usize, h: usize, w: usize, afrac: u32) -> Vec<u8> {
|
||
let (one, half) = (1i64 << afrac, 1i64 << (afrac - 1));
|
||
let mut rgb = vec![0u8; h * w * c];
|
||
for cc in 0..c {
|
||
for p in 0..h * w {
|
||
let v = x[cc * h * w + p].clamp(0, one);
|
||
rgb[p * c + cc] = ((v * 255 + half) >> afrac).clamp(0, 255) as u8;
|
||
}
|
||
}
|
||
rgb
|
||
}
|
||
|
||
/// FNV-1a 64-bit of `bytes` — the dependency-free checksum the export tool and
|
||
/// this decoder share to prove byte-for-byte agreement across languages.
|
||
pub fn fnv1a(bytes: &[u8]) -> u64 {
|
||
let mut hash: u64 = 14695981039346656037;
|
||
for &b in bytes {
|
||
hash = (hash ^ b as u64).wrapping_mul(1099511628211);
|
||
}
|
||
hash
|
||
}
|
||
|
||
/// The conditional codec's **integer entropy model**: the per-scale-bucket
|
||
/// quantized Gaussian CDFs a real `GaussianConditional` builds, plus the integer
|
||
/// `scale_table` edges (`2^afrac`) that [`det_build_index`] buckets σ against.
|
||
/// Pure data — no float; the same tables the encoder used, so the decode is exact.
|
||
pub struct CondEntropy {
|
||
/// Ascending σ bucket edges at `2^afrac` (from the codec's `scale_table`).
|
||
pub edges: Vec<i64>,
|
||
/// One quantized CDF per bucket (`cdfs[b]` indexed by [`det_build_index`]).
|
||
pub cdfs: Vec<Vec<i32>>,
|
||
/// Per-bucket CDF length (the `[0, len-2)` codable span + escape slot).
|
||
pub cdf_lengths: Vec<usize>,
|
||
/// Per-bucket symbol offset (the residual that bucket's CDF is centered on).
|
||
pub offsets: Vec<i32>,
|
||
/// The hyper-latent `z`'s **factorized** prior: one quantized CDF per channel.
|
||
pub z_cdfs: Vec<Vec<i32>>,
|
||
pub z_cdf_lengths: Vec<usize>,
|
||
pub z_offsets: Vec<i32>,
|
||
/// Per-channel median as a fixed-point integer (`round(median·2^afrac)`), so
|
||
/// `z_q = s·2^afrac + median_q` rebuilds the latent with no float.
|
||
pub z_median_q: Vec<i64>,
|
||
}
|
||
|
||
fn cdf_rows(v: &Value) -> Vec<Vec<i32>> {
|
||
v.as_array()
|
||
.unwrap()
|
||
.iter()
|
||
.map(|row| row.as_array().unwrap().iter().map(|n| n.as_i64().unwrap() as i32).collect())
|
||
.collect()
|
||
}
|
||
fn usizes(v: &Value) -> Vec<usize> {
|
||
v.as_array().unwrap().iter().map(|n| n.as_u64().unwrap() as usize).collect()
|
||
}
|
||
fn i32s(v: &Value) -> Vec<i32> {
|
||
v.as_array().unwrap().iter().map(|n| n.as_i64().unwrap() as i32).collect()
|
||
}
|
||
|
||
impl CondEntropy {
|
||
/// Parse the entropy tables the export tool emits (`cond_entropy.json`).
|
||
pub fn from_value(v: &Value) -> Self {
|
||
Self {
|
||
edges: i64s(&v["scale_edges"]),
|
||
cdfs: cdf_rows(&v["cdfs"]),
|
||
cdf_lengths: usizes(&v["cdf_lengths"]),
|
||
offsets: i32s(&v["offsets"]),
|
||
z_cdfs: cdf_rows(&v["z_cdfs"]),
|
||
z_cdf_lengths: usizes(&v["z_cdf_lengths"]),
|
||
z_offsets: i32s(&v["z_offsets"]),
|
||
z_median_q: i64s(&v["z_median_q"]),
|
||
}
|
||
}
|
||
}
|
||
|
||
/// The codec's **bitstream** — both streams entropy-coded, fully integer:
|
||
/// `z_rans` is the factorized-prior rANS of the hyper-latent `z` (the σ/μ source);
|
||
/// `rans` is the Gaussian-conditional rANS of the `y` residuals.
|
||
pub struct CondBitstream {
|
||
pub z_rans: Vec<u8>,
|
||
pub n_z: usize,
|
||
pub zh: usize,
|
||
pub zw: usize,
|
||
pub m: usize,
|
||
pub sh: usize,
|
||
pub sw: usize,
|
||
pub rans: Vec<u8>,
|
||
}
|
||
|
||
fn bytes(v: &Value) -> Vec<u8> {
|
||
v.as_array().unwrap().iter().map(|n| n.as_u64().unwrap() as u8).collect()
|
||
}
|
||
|
||
impl CondBitstream {
|
||
/// [`Self::from_value`], refusing rather than panicking on any JSON: every member
|
||
/// present, a size an integer that fits `usize`, a stream an array of bytes. A
|
||
/// reader of JSON it did not write (the browser's debug decode) uses this, and
|
||
/// decodes with [`decode_conditional_checked`].
|
||
pub fn try_from_value(v: &Value) -> Result<Self, &'static str> {
|
||
let dim = |k: &'static str| v.get(k).and_then(Value::as_u64).and_then(|n| usize::try_from(n).ok()).ok_or(k);
|
||
let bytes = |k: &'static str| -> Result<Vec<u8>, &'static str> {
|
||
v.get(k).and_then(Value::as_array).ok_or(k)?.iter()
|
||
.map(|b| b.as_u64().and_then(|b| u8::try_from(b).ok())).collect::<Option<_>>().ok_or(k)
|
||
};
|
||
Ok(Self {
|
||
z_rans: bytes("z_rans")?, n_z: dim("n_z")?, zh: dim("zh")?, zw: dim("zw")?,
|
||
m: dim("m")?, sh: dim("sh")?, sw: dim("sw")?, rans: bytes("rans")?,
|
||
})
|
||
}
|
||
|
||
/// Parse the bitstream the export tool emits (`cond_bitstream.json`). Panics on a
|
||
/// JSON the tool would not write; [`Self::try_from_value`] refuses instead.
|
||
pub fn from_value(v: &Value) -> Self {
|
||
Self {
|
||
z_rans: bytes(&v["z_rans"]),
|
||
n_z: u(v, "n_z"),
|
||
zh: u(v, "zh"),
|
||
zw: u(v, "zw"),
|
||
m: u(v, "m"),
|
||
sh: u(v, "sh"),
|
||
sw: u(v, "sw"),
|
||
rans: bytes(&v["rans"]),
|
||
}
|
||
}
|
||
}
|
||
|
||
/// Decode a full conditional-codec **bitstream → image**, end to end and fully
|
||
/// integer — the assembled `NeuralIntegerExact` codec (BOTH streams entropy-coded):
|
||
///
|
||
/// 1. factorized rANS-decode `z`: per channel, `z_q = s·2^afrac + median_q` (no
|
||
/// float) → integer hyper-decoder `h_s` → σ‖μ (all i64);
|
||
/// 2. each `y` element: σ → bucket ([`det_build_index`], no float) → rANS-decode
|
||
/// the residual `r` → `y_hat = r·2^afrac + μ` (μ from the integer `h_s`);
|
||
/// 3. integer synthesis `g_s` ([`decode`]) → RGB (HWC `u8`).
|
||
///
|
||
/// Encoder and decoder agree on every bucket because σ comes from the **integer**
|
||
/// `h_s` — the integer-entropy contract. Byte-identical on every machine.
|
||
pub fn decode_conditional(
|
||
g_s: &IntModel, h_s: &IntModel, ent: &CondEntropy, bs: &CondBitstream,
|
||
) -> Result<(Vec<u8>, usize, usize), &'static str> {
|
||
let afrac = g_s.afrac;
|
||
let one = 1i64 << afrac;
|
||
// 1. factorized-decode z (per channel), rebuild z_q = s·2^afrac + median_q.
|
||
let mut z_int = vec![0i64; bs.n_z * bs.zh * bs.zw];
|
||
{
|
||
let mut dz = RansDecoder::new(&bs.z_rans)?;
|
||
let plane = bs.zh * bs.zw;
|
||
for ch in 0..bs.n_z {
|
||
for p in 0..plane {
|
||
let s = dz.decode(&ent.z_cdfs[ch], ent.z_cdf_lengths[ch], ent.z_offsets[ch]);
|
||
z_int[ch * plane + p] = s as i64 * one + ent.z_median_q[ch];
|
||
}
|
||
}
|
||
}
|
||
let (sm, oc, sh, sw) = run_net(h_s, &z_int, bs.n_z, bs.zh, bs.zw);
|
||
if oc != 2 * bs.m || sh != bs.sh || sw != bs.sw {
|
||
return Err("h_s output shape disagrees with the bitstream header");
|
||
}
|
||
let n = bs.m * sh * sw;
|
||
let (sigma, mu) = sm.split_at(n); // σ = first M channels, μ = next M
|
||
// 2. entropy-decode y, conditioned on the integer σ; y_hat = r·2^afrac + μ.
|
||
let mut dec = RansDecoder::new(&bs.rans)?;
|
||
let mut y_hat = vec![0i64; n];
|
||
for p in 0..n {
|
||
let b = det_build_index(sigma[p], &ent.edges);
|
||
let r = dec.decode(&ent.cdfs[b], ent.cdf_lengths[b], ent.offsets[b]);
|
||
y_hat[p] = r as i64 * one + mu[p];
|
||
}
|
||
// 3. integer synthesis g_s → pixels.
|
||
Ok(decode(g_s, &y_hat, bs.m, sh, sw))
|
||
}
|
||
|
||
/// The kind of rule a checked decode refused on. The integer-payload decoders
|
||
/// (`int_payload`, and `int_audio` and `int_splat`'s typed decoders) report these as
|
||
/// the refusal codes of `extensions/integer-payloads.md` §8.
|
||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||
pub enum RefusalKind {
|
||
/// A value the payload carries, or a latent its stream decodes, is outside the
|
||
/// range the decode is exact over.
|
||
FieldRange,
|
||
/// A tensor or an activation would exceed [`MAX_TENSOR`] elements.
|
||
TooLarge,
|
||
/// A CDF table breaks the table rules.
|
||
CdfMalformed,
|
||
/// A rANS stream ran out before its last symbol.
|
||
StreamExhausted,
|
||
/// A rANS stream does not end canonically after its last symbol.
|
||
StreamNoncanonical,
|
||
/// A shape the payload declares disagrees with the one the decode produces.
|
||
GeometryMismatch,
|
||
/// The parameter set does not fit the payload.
|
||
ModelMismatch,
|
||
}
|
||
|
||
/// A checked decode's refusal: the kind of rule and which.
|
||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||
pub struct Refused {
|
||
pub kind: RefusalKind,
|
||
pub detail: &'static str,
|
||
}
|
||
|
||
impl std::fmt::Display for Refused {
|
||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||
write!(f, "{:?}: {}", self.kind, self.detail)
|
||
}
|
||
}
|
||
|
||
pub(crate) const fn refused(kind: RefusalKind, detail: &'static str) -> Refused {
|
||
Refused { kind, detail }
|
||
}
|
||
|
||
/// The most elements any one buffer a checked decode forms may hold, the same on
|
||
/// every target (`extensions/integer-payloads.md` §7).
|
||
///
|
||
/// cbindgen:ignore
|
||
pub const MAX_TENSOR: usize = 1 << 26;
|
||
/// The largest fixed-point scale a checked decode accepts.
|
||
///
|
||
/// cbindgen:ignore
|
||
pub const MAX_AFRAC: u32 = 30;
|
||
|
||
/// A table of `extensions/integer-payloads.md` §2 over its first `len` entries:
|
||
/// `3 ≤ len ≤ 4097` and `len` within the row, `cdf[0] = 0`, `cdf[len − 1] = 2^16`,
|
||
/// strictly increasing over `len`.
|
||
fn table_ok(cdf: &[i32], len: usize) -> bool {
|
||
(3..=4097).contains(&len) && len <= cdf.len() && cdf[0] == 0 && cdf[len - 1] == 1 << 16
|
||
&& cdf[..len].windows(2).all(|p| p[0] < p[1])
|
||
}
|
||
|
||
/// The entropy model's consistency: a non-empty scale table with one Gaussian table,
|
||
/// length and offset per edge; one `z` table, length, offset and median per channel;
|
||
/// every table one the rANS decoder is total over.
|
||
fn check_entropy(t: &CondEntropy) -> Result<(), &'static str> {
|
||
let nb = t.edges.len();
|
||
if nb == 0 || t.cdfs.len() != nb || t.cdf_lengths.len() != nb || t.offsets.len() != nb {
|
||
return Err("one Gaussian table, length and offset per scale edge, and at least one edge");
|
||
}
|
||
let nz = t.z_cdfs.len();
|
||
if nz == 0 || t.z_cdf_lengths.len() != nz || t.z_offsets.len() != nz || t.z_median_q.len() != nz {
|
||
return Err("one z table, length, offset and median per channel");
|
||
}
|
||
let all = t.cdfs.iter().zip(&t.cdf_lengths).chain(t.z_cdfs.iter().zip(&t.z_cdf_lengths));
|
||
if all.into_iter().any(|(c, &l)| !table_ok(c, l)) {
|
||
return Err("a table is not a CDF from 0 to 65536, strictly increasing over its length");
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
impl CondEntropy {
|
||
/// Parse and check the entropy tables (`cond_entropy.json`, read from a set as
|
||
/// `entropy.json`): every member present with integers in range, the lengths
|
||
/// consistent and every table well formed. Refuses rather than panics.
|
||
pub fn try_from_value(v: &Value) -> Result<Self, &'static str> {
|
||
fn arr<'a>(v: &'a Value, k: &str) -> Result<&'a Vec<Value>, &'static str> {
|
||
v.get(k).and_then(Value::as_array).ok_or("an entropy member is missing or not an array")
|
||
}
|
||
fn ints<T: TryFrom<i64>>(v: &Value, k: &str) -> Result<Vec<T>, &'static str> {
|
||
arr(v, k)?.iter().map(|x| x.as_i64().and_then(|i| T::try_from(i).ok())).collect::<Option<_>>()
|
||
.ok_or("an entropy member holds a value out of range")
|
||
}
|
||
fn rows(v: &Value, k: &str) -> Result<Vec<Vec<i32>>, &'static str> {
|
||
arr(v, k)?.iter().map(|r| r.as_array().ok_or("an entropy table is not an array")
|
||
.and_then(|r| r.iter().map(|x| x.as_i64().and_then(|i| i32::try_from(i).ok())).collect::<Option<_>>()
|
||
.ok_or("an entropy table holds a value out of range"))).collect()
|
||
}
|
||
let t = CondEntropy {
|
||
edges: ints(v, "scale_edges")?,
|
||
cdfs: rows(v, "cdfs")?,
|
||
cdf_lengths: ints(v, "cdf_lengths")?,
|
||
offsets: ints(v, "offsets")?,
|
||
z_cdfs: rows(v, "z_cdfs")?,
|
||
z_cdf_lengths: ints(v, "z_cdf_lengths")?,
|
||
z_offsets: ints(v, "z_offsets")?,
|
||
z_median_q: ints(v, "z_median_q")?,
|
||
};
|
||
check_entropy(&t)?;
|
||
Ok(t)
|
||
}
|
||
}
|
||
|
||
/// A conditional bitstream's fields, borrowed: what a `WIH1` payload carries.
|
||
#[derive(Debug, Clone, Copy)]
|
||
pub struct CondStream<'a> {
|
||
pub m: usize,
|
||
pub sh: usize,
|
||
pub sw: usize,
|
||
pub n_z: usize,
|
||
pub zh: usize,
|
||
pub zw: usize,
|
||
pub z_rans: &'a [u8],
|
||
pub rans: &'a [u8],
|
||
}
|
||
|
||
impl CondBitstream {
|
||
/// The borrowed form [`decode_conditional_checked`] reads.
|
||
pub fn stream(&self) -> CondStream<'_> {
|
||
CondStream { m: self.m, sh: self.sh, sw: self.sw, n_z: self.n_z, zh: self.zh, zw: self.zw, z_rans: &self.z_rans, rans: &self.rans }
|
||
}
|
||
}
|
||
|
||
/// A network's output on a `[c, h, w]` input, planned from its layers' geometry and
|
||
/// parameter lengths alone, before anything is allocated or any value read
|
||
/// (`extensions/integer-payloads.md` §7). The input must be within [`MAX_TENSOR`]
|
||
/// (`TooLarge`). Then, layer by layer in order:
|
||
///
|
||
/// 1. the layer fits the running channel count (`ModelMismatch`): a convolution's
|
||
/// output channels, kernel and stride are non-zero, its weights number
|
||
/// `c·out·k·k` and its multipliers and biases `out`, and its requantisation shift is
|
||
/// under 128; an inverse GDN's `gamma` numbers `c·c` and its `beta` `c`, and its
|
||
/// shifts are under 128; a leaky ReLU's shift is under 64;
|
||
/// 2. every size the layer forms is within [`MAX_TENSOR`] (`TooLarge`): a
|
||
/// convolution's padded input side `n + 2·pad`, a transposed convolution's uncropped
|
||
/// side `(n − 1)·stride + k + opad` and uncropped plane, and the output;
|
||
/// 3. the layer leaves an output (`ModelMismatch`): a kernel no larger than the padded
|
||
/// input, a crop `2·pad` smaller than the uncropped side.
|
||
///
|
||
/// Every size is formed in `u64` with checked arithmetic, a size past `u64` being over
|
||
/// the cap, so a 32-bit sink refuses exactly what a 64-bit one does, with the same
|
||
/// kind; and once a network is planned, every size [`run_net_checked`] forms fits
|
||
/// `usize` on every target. The first rule that fails is the refusal.
|
||
pub(crate) fn plan_net(net: &IntModel, c: usize, h: usize, w: usize) -> Result<NetPlan, Refused> {
|
||
plan_layers(net, c, h, w)
|
||
}
|
||
|
||
/// What [`plan_net`] finds about a network on a `[c, h, w]` input, in elements.
|
||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||
pub(crate) struct NetPlan {
|
||
/// The output `[c, h, w]`.
|
||
pub out: (usize, usize, usize),
|
||
/// The largest single buffer the network forms: its input, a layer's output, a
|
||
/// transposed convolution's uncropped plane.
|
||
pub largest: u64,
|
||
/// The most elements the reference kernels hold at once while one layer runs, its
|
||
/// input included: a transposed convolution's input, uncropped plane and output; a
|
||
/// convolution's input and output; an inverse GDN's input, output and one
|
||
/// position's squares (`c` values of 16 bytes, counted as `2·c`); a ReLU's input,
|
||
/// rewritten in place. Each element is an `i64`.
|
||
pub live: u64,
|
||
/// The multiply-accumulates the layers comprise (`extensions/integer-payloads.md`
|
||
/// §7): a convolution's every output times `c·k·k` (every tap of its padded
|
||
/// input), a transposed convolution's every input times `out·k·k`, an inverse
|
||
/// GDN's every element times `c`. A kernel may skip a zero term, so it performs
|
||
/// at most this many. Summed in `u128`: a layer's terms are its weights times at
|
||
/// most `2^26` positions, so no set a sink can hold reaches the type's limit.
|
||
pub macs: u128,
|
||
}
|
||
|
||
fn plan_layers(net: &IntModel, c: usize, h: usize, w: usize) -> Result<NetPlan, Refused> {
|
||
use RefusalKind::{ModelMismatch, TooLarge};
|
||
const CAP: u64 = MAX_TENSOR as u64;
|
||
let big = |what: &'static str| refused(TooLarge, what);
|
||
let within = |v: Option<u64>, what: &'static str| v.filter(|&n| n <= CAP).ok_or(big(what));
|
||
let (mut c, mut h, mut w) = (c as u64, h as u64, w as u64);
|
||
if c == 0 || h == 0 || w == 0 {
|
||
return Err(refused(ModelMismatch, "a network's input is empty"));
|
||
}
|
||
let input = within(c.checked_mul(h).and_then(|v| v.checked_mul(w)), "a network's input exceeds 2^26 elements")?;
|
||
let (mut largest, mut live) = (input, input);
|
||
let mut macs: u128 = 0;
|
||
let params = |cout: u64, k: u64, weight: usize, mult: usize, bias: usize, c: u64| {
|
||
let n = c.checked_mul(cout).and_then(|v| v.checked_mul(k)).and_then(|v| v.checked_mul(k));
|
||
n == Some(weight as u64) && mult as u64 == cout && bias as u64 == cout
|
||
};
|
||
for layer in &net.layers {
|
||
match layer {
|
||
IntLayer::ConvT { out_ch, k, stride, pad, opad, weight, mult, bias } => {
|
||
let (cout, k, s, p, op) = (*out_ch as u64, *k as u64, *stride as u64, *pad as u64, *opad as u64);
|
||
if cout == 0 || k == 0 || s == 0 || !params(cout, k, weight.len(), mult.len(), bias.len(), c) || net.shift > 127 {
|
||
return Err(refused(ModelMismatch, "a transposed convolution does not fit the running channel count"));
|
||
}
|
||
let full = |n: u64| (n - 1).checked_mul(s).and_then(|v| v.checked_add(k)).and_then(|v| v.checked_add(op));
|
||
let (fh, fw) = (within(full(h), "a layer's uncropped side exceeds 2^26")?, within(full(w), "a layer's uncropped side exceeds 2^26")?);
|
||
let plane = within(cout.checked_mul(fh).and_then(|v| v.checked_mul(fw)), "a layer's uncropped plane exceeds 2^26 elements")?;
|
||
let crop = p.checked_mul(2).unwrap_or(u64::MAX);
|
||
if fh <= crop || fw <= crop {
|
||
return Err(refused(ModelMismatch, "a transposed convolution's padding leaves no output"));
|
||
}
|
||
let a = c * h * w;
|
||
macs = macs.saturating_add(a as u128 * (cout * k * k) as u128);
|
||
(c, h, w) = (cout, fh - crop, fw - crop);
|
||
let b = c * h * w;
|
||
largest = largest.max(plane).max(b);
|
||
live = live.max(a + plane + b);
|
||
}
|
||
IntLayer::Conv { out_ch, k, stride, pad, weight, mult, bias } => {
|
||
let (cout, k, s, p) = (*out_ch as u64, *k as u64, *stride as u64, *pad as u64);
|
||
if cout == 0 || k == 0 || s == 0 || !params(cout, k, weight.len(), mult.len(), bias.len(), c) || net.shift > 127 {
|
||
return Err(refused(ModelMismatch, "a convolution does not fit the running channel count"));
|
||
}
|
||
let padded = |n: u64| p.checked_mul(2).and_then(|v| v.checked_add(n));
|
||
let (hp, wp) = (within(padded(h), "a layer's padded side exceeds 2^26")?, within(padded(w), "a layer's padded side exceeds 2^26")?);
|
||
if hp < k || wp < k {
|
||
return Err(refused(ModelMismatch, "a convolution's kernel is larger than its padded input"));
|
||
}
|
||
let (ho, wo) = ((hp - k) / s + 1, (wp - k) / s + 1);
|
||
let b = within(cout.checked_mul(ho).and_then(|v| v.checked_mul(wo)), "a layer's output exceeds 2^26 elements")?;
|
||
live = live.max(c * h * w + b);
|
||
largest = largest.max(b);
|
||
macs = macs.saturating_add(b as u128 * (c * k * k) as u128);
|
||
(c, h, w) = (cout, ho, wo);
|
||
}
|
||
IntLayer::Igdn { gamma, beta } => {
|
||
if c.checked_mul(c) != Some(gamma.len() as u64) || beta.len() as u64 != c || net.gf > 127 || net.afrac > 127 {
|
||
return Err(refused(ModelMismatch, "an inverse GDN does not fit the running channel count"));
|
||
}
|
||
live = live.max(2 * c * h * w + 2 * c);
|
||
macs = macs.saturating_add((c * h * w) as u128 * c as u128);
|
||
}
|
||
IntLayer::Relu => {}
|
||
IntLayer::Leaky { slope_frac, .. } => {
|
||
if *slope_frac > 63 {
|
||
return Err(refused(ModelMismatch, "a leaky ReLU's shift is 64 or more"));
|
||
}
|
||
}
|
||
}
|
||
}
|
||
// Every running size is at most 2^26 here, so it fits usize on every target.
|
||
Ok(NetPlan { out: (c as usize, h as usize, w as usize), largest, live, macs })
|
||
}
|
||
|
||
/// [`run_net_checked`] with its refusal as a [`Refused`]: values that could overflow a
|
||
/// layer are the stream's (`FieldRange`); a layer that does not fit is the model's.
|
||
fn run_checked(net: &IntModel, x: &[i64], c: usize, h: usize, w: usize) -> Result<Vec<i64>, Refused> {
|
||
run_net_checked(net, x, c, h, w).map(|(o, ..)| o).map_err(|r| {
|
||
if r.out_of_range {
|
||
refused(RefusalKind::FieldRange, "a decoded latent could take a layer outside the exact integer range")
|
||
} else {
|
||
refused(RefusalKind::ModelMismatch, "a network's parameters do not fit the running shapes")
|
||
}
|
||
})
|
||
}
|
||
|
||
/// [`decode_conditional`], refusing instead of panicking or overflowing on any input:
|
||
/// the decode the `WIH1` payload (`wai.neural.int_hyper`, and the keyframe of
|
||
/// `wai.video.int_hyper`) runs. On a stream [`decode_conditional`] decodes, both give
|
||
/// the same image.
|
||
///
|
||
/// Before anything is allocated, in this order: the networks share one fixed-point
|
||
/// scale in `1..=`[`MAX_AFRAC`]; the entropy model is consistent and has a `z` table
|
||
/// for each of the payload's channels; the latents are within [`MAX_TENSOR`]; `h_s` on
|
||
/// the hyper-latent's shape, then `g_s` on `[m, sh, sw]`, are planned ([`plan_net`]: each
|
||
/// layer fits, forms nothing over the cap, and leaves an output); then `h_s` must give
|
||
/// `[2·m, sh, sw]` (`GeometryMismatch`) and `g_s` three channels. Then the hyper-latent
|
||
/// is decoded symbol by symbol (a symbol that runs the stream out is refused before
|
||
/// its value is rebuilt with [`dequantize`]), its stream must end canonically, `h_s`
|
||
/// runs through [`run_net_checked`], the latent is decoded the same way, and `g_s`
|
||
/// runs. Returns RGB HWC with its height and width.
|
||
pub fn decode_conditional_checked(
|
||
g_s: &IntModel, h_s: &IntModel, ent: &CondEntropy, bs: &CondStream<'_>,
|
||
) -> Result<(Vec<u8>, usize, usize), Refused> {
|
||
use crate::int_entropy::dequantize;
|
||
use RefusalKind::*;
|
||
let p = cond_plan(g_s, h_s, ent, bs)?;
|
||
let (afrac, n_z, zh, zw, m, sh, sw) = (g_s.afrac, bs.n_z, bs.zh, bs.zw, bs.m, bs.sh, bs.sw);
|
||
let (zplane, n, (oc, oh, ow)) = (zh * zw, m * sh * sw, p.gs.out);
|
||
// 1. factorized-decode z, rebuild z_q = s·2^afrac + median_q. The latent grows a
|
||
// symbol at a time, to its declared length and no further, and the first symbol
|
||
// that runs its stream out refuses the payload before its value is used.
|
||
let mut z = Vec::new();
|
||
let mut dz = RansDecoder::new(bs.z_rans).map_err(|_| refused(FieldRange, "the z stream is not a rANS stream"))?;
|
||
for ch in 0..n_z {
|
||
for _ in 0..zplane {
|
||
let s = dz.decode(&ent.z_cdfs[ch], ent.z_cdf_lengths[ch], ent.z_offsets[ch]);
|
||
if dz.is_exhausted() {
|
||
return Err(refused(StreamExhausted, "the z stream ended before its last symbol"));
|
||
}
|
||
push_capped(&mut z, dequantize(s, afrac, ent.z_median_q[ch]).ok_or(refused(FieldRange, "a hyper-latent value leaves i64"))?, n_z * zplane);
|
||
}
|
||
}
|
||
if !dz.is_canonical_end() {
|
||
return Err(refused(StreamNoncanonical, "the z stream does not end canonically after its last symbol"));
|
||
}
|
||
let sm = run_checked(h_s, &z, n_z, zh, zw)?;
|
||
drop(z);
|
||
let (sigma, mu) = sm.split_at(n);
|
||
// 2. y, conditioned on the integer σ: y_hat = r·2^afrac + μ.
|
||
let mut dec = RansDecoder::new(bs.rans).map_err(|_| refused(FieldRange, "the y stream is not a rANS stream"))?;
|
||
let mut y_hat = Vec::new();
|
||
for p in 0..n {
|
||
let b = det_build_index(sigma[p], &ent.edges);
|
||
let r = dec.decode(&ent.cdfs[b], ent.cdf_lengths[b], ent.offsets[b]);
|
||
if dec.is_exhausted() {
|
||
return Err(refused(StreamExhausted, "the y stream ended before its last symbol"));
|
||
}
|
||
push_capped(&mut y_hat, dequantize(r, afrac, mu[p]).ok_or(refused(FieldRange, "a latent value leaves i64"))?, n);
|
||
}
|
||
if !dec.is_canonical_end() {
|
||
return Err(refused(StreamNoncanonical, "the y stream does not end canonically after its last symbol"));
|
||
}
|
||
drop(sm);
|
||
// 3. integer synthesis g_s → pixels.
|
||
let x = run_checked(g_s, &y_hat, m, sh, sw)?;
|
||
drop(y_hat);
|
||
Ok((to_rgb8(&x, oc, oh, ow, afrac), oh, ow))
|
||
}
|
||
|
||
/// Append `x` to `v`, a buffer that will hold `total` elements, growing it by doubling
|
||
/// but never past `total`: a decode holds no more than it declared, and no buffer is
|
||
/// reserved ahead of the symbols that fill it.
|
||
pub(crate) fn push_capped<T>(v: &mut Vec<T>, x: T, total: usize) {
|
||
if v.len() == v.capacity() {
|
||
let want = v.capacity().saturating_mul(2).max(64).min(total).max(v.len() + 1);
|
||
v.reserve_exact(want - v.len());
|
||
}
|
||
v.push(x);
|
||
}
|
||
|
||
/// What [`decode_conditional_checked`] checks before it decodes a symbol, in its
|
||
/// order, and the buffers it then forms: the networks share one fixed-point scale in
|
||
/// `1..=`[`MAX_AFRAC`]; the entropy model is consistent and has a `z` table for each
|
||
/// of the payload's channels; the latents are within [`MAX_TENSOR`]; `h_s` on the
|
||
/// hyper-latent's shape, then `g_s` on `[m, sh, sw]`, are planned ([`plan_net`]); then `h_s`
|
||
/// must give `[2·m, sh, sw]` (`GeometryMismatch`) and `g_s` three channels.
|
||
pub(crate) fn cond_plan(g_s: &IntModel, h_s: &IntModel, ent: &CondEntropy, bs: &CondStream<'_>) -> Result<CondPlan, Refused> {
|
||
use RefusalKind::*;
|
||
let afrac = g_s.afrac;
|
||
if h_s.afrac != afrac || !(1..=MAX_AFRAC).contains(&afrac) {
|
||
return Err(refused(ModelMismatch, "g_s and h_s do not share one fixed-point scale in 1..=30"));
|
||
}
|
||
check_entropy(ent).map_err(|e| refused(ModelMismatch, e))?;
|
||
let (n_z, zh, zw, m, sh, sw) = (bs.n_z, bs.zh, bs.zw, bs.m, bs.sh, bs.sw);
|
||
if [n_z, zh, zw, m, sh, sw].contains(&0) {
|
||
return Err(refused(FieldRange, "a latent dimension is zero"));
|
||
}
|
||
if n_z > ent.z_cdfs.len() {
|
||
return Err(refused(ModelMismatch, "the hyper-latent has more channels than the entropy model has z tables"));
|
||
}
|
||
let size = |d: &[usize]| d.iter().try_fold(1usize, |a, &x| a.checked_mul(x)).filter(|&n| n <= MAX_TENSOR);
|
||
if size(&[n_z, zh, zw]).is_none() || size(&[2, m, sh, sw]).is_none() {
|
||
return Err(refused(TooLarge, "a latent exceeds 2^26 elements"));
|
||
}
|
||
// Both networks are planned, caps included, before either output shape is
|
||
// compared with the payload's.
|
||
let hs = plan_net(h_s, n_z, zh, zw)?;
|
||
let gs = plan_net(g_s, m, sh, sw)?;
|
||
if hs.out != (2 * m, sh, sw) {
|
||
return Err(refused(GeometryMismatch, "h_s does not give the latent shape the payload declares"));
|
||
}
|
||
if gs.out.0 != 3 {
|
||
return Err(refused(ModelMismatch, "g_s does not give three channels"));
|
||
}
|
||
Ok(CondPlan { hs, gs, z: (n_z * zh * zw) as u64, y: (m * sh * sw) as u64 })
|
||
}
|
||
|
||
/// The buffers [`decode_conditional_checked`] forms on a payload [`cond_plan`] admits.
|
||
#[derive(Debug, Clone, Copy)]
|
||
pub(crate) struct CondPlan {
|
||
pub hs: NetPlan,
|
||
pub gs: NetPlan,
|
||
/// The hyper-latent's and the latent's elements: the symbols each stream codes.
|
||
pub z: u64,
|
||
pub y: u64,
|
||
}
|
||
|
||
impl CondPlan {
|
||
/// The image's bytes (RGB8).
|
||
pub fn output_bytes(&self) -> u64 {
|
||
let (c, h, w) = self.gs.out;
|
||
(c * h * w) as u64
|
||
}
|
||
|
||
/// The most bytes the decode holds at once on the reference kernels, beyond the
|
||
/// payload and the parsed set: the hyper-latent while `h_s` runs; `h_s`'s output
|
||
/// (`[2·m, sh, sw]`) while the latent is decoded; the latent while `g_s` runs; and
|
||
/// `g_s`'s output while it becomes RGB8. Every value is an `i64`.
|
||
pub fn working_bytes(&self) -> u64 {
|
||
let gs_out = self.output_bytes();
|
||
[8 * (self.z + self.hs.live), 8 * (2 * self.y + self.y), 8 * (self.y + self.gs.live), 8 * gs_out + gs_out]
|
||
.into_iter().max().unwrap_or(0)
|
||
}
|
||
|
||
/// The networks' multiply-accumulates: `h_s`'s, then `g_s`'s ([`NetPlan::macs`]).
|
||
pub fn macs(&self) -> u128 {
|
||
self.hs.macs.saturating_add(self.gs.macs)
|
||
}
|
||
}
|
||
|
||
/// The fast kernels against the scalar reference, element by element. Operand
|
||
/// ranges cover the `i32` fast path, the `i64` path, and sparse (zero-heavy)
|
||
/// inputs; geometries cover strides, paddings, output paddings and kernels smaller
|
||
/// than the stride.
|
||
#[cfg(test)]
|
||
mod kernel_equivalence {
|
||
use super::*;
|
||
use crate::int_kernels::testkit::{Rng, FAST};
|
||
use crate::int_kernels::{with_policy, Policy};
|
||
|
||
/// (input range, weight range) per case: fits-i32, beyond-i32, mixed.
|
||
const RANGES: [(i64, i64); 3] = [(1 << 14, 2047), (1 << 40, 1 << 10), (1 << 20, 1 << 20)];
|
||
|
||
#[test]
|
||
fn transposed_conv_equals_reference() {
|
||
let mut g = Rng(0xC0FF_EE01);
|
||
let mut ran = 0;
|
||
for case in 0..600 {
|
||
let k = 1 + g.below(5);
|
||
let stride = 1 + g.below(3);
|
||
let pad = g.below(k);
|
||
let opad = if stride > 1 { g.below(stride) } else { 0 };
|
||
let (h, w) = (1 + g.below(7), 1 + g.below(9));
|
||
let (cin, cout) = (1 + g.below(6), 1 + g.below(9));
|
||
if (h - 1) * stride + k + opad <= 2 * pad || (w - 1) * stride + k + opad <= 2 * pad {
|
||
continue;
|
||
}
|
||
let (xr, wr) = RANGES[case % 3];
|
||
let x = g.vec(cin * h * w, xr, [0, 50, 90][case % 3]);
|
||
let wt = g.vec(cin * cout * k * k, wr, 10);
|
||
let mult: Vec<i64> = (0..cout).map(|_| 1 + g.signed(1 << 20).abs()).collect();
|
||
let bias = g.vec(cout, 1 << 20, 0);
|
||
let shift = 16 + g.below(8) as u32;
|
||
let want = reference::conv_t(&x, cin, h, w, &wt, cout, k, stride, pad, opad, &mult, &bias, shift);
|
||
for p in FAST {
|
||
let got = with_policy(p, || conv_t(&x, cin, h, w, &wt, cout, k, stride, pad, opad, &mult, &bias, shift));
|
||
assert_eq!(got, want, "conv_t {p:?} k={k} s={stride} p={pad} op={opad} {cin}x{h}x{w}->{cout}");
|
||
}
|
||
ran += 1;
|
||
}
|
||
assert!(ran > 300, "too few valid cases ({ran})");
|
||
}
|
||
|
||
#[test]
|
||
fn regular_conv_equals_reference() {
|
||
let mut g = Rng(0xC0FF_EE02);
|
||
let mut ran = 0;
|
||
for case in 0..600 {
|
||
let k = 1 + g.below(5);
|
||
let stride = 1 + g.below(3);
|
||
let pad = g.below(k);
|
||
let (h, w) = (1 + g.below(9), 1 + g.below(11));
|
||
let (cin, cout) = (1 + g.below(6), 1 + g.below(9));
|
||
if h + 2 * pad < k || w + 2 * pad < k {
|
||
continue;
|
||
}
|
||
let (xr, wr) = RANGES[case % 3];
|
||
let x = g.vec(cin * h * w, xr, [0, 50, 90][case % 3]);
|
||
let wt = g.vec(cin * cout * k * k, wr, 10);
|
||
let mult: Vec<i64> = (0..cout).map(|_| 1 + g.signed(1 << 20).abs()).collect();
|
||
let bias = g.vec(cout, 1 << 20, 0);
|
||
let shift = 16 + g.below(8) as u32;
|
||
let want = reference::conv(&x, cin, h, w, &wt, cout, k, stride, pad, &mult, &bias, shift);
|
||
for p in FAST {
|
||
let got = with_policy(p, || conv(&x, cin, h, w, &wt, cout, k, stride, pad, &mult, &bias, shift));
|
||
assert_eq!(got, want, "conv {p:?} k={k} s={stride} p={pad} {cin}x{h}x{w}->{cout}");
|
||
}
|
||
ran += 1;
|
||
}
|
||
assert!(ran > 300, "too few valid cases ({ran})");
|
||
}
|
||
|
||
/// Strides up to 16, beyond the kernel and the input. The transposed kernel's
|
||
/// row bands then hold fewer rows than the stride (a tap table only for each row
|
||
/// phase present) as well as more (rows wrapping onto the same phases), with
|
||
/// bands starting anywhere in a phase; many phases receive no tap at all. The
|
||
/// regular kernel's stride often exceeds its padded row.
|
||
#[test]
|
||
fn wide_strides_equal_reference() {
|
||
let mut g = Rng(0xC0FF_EE03);
|
||
let (mut ran_t, mut ran_c) = (0, 0);
|
||
for case in 0..300 {
|
||
let k = 1 + g.below(5);
|
||
let stride = 1 + g.below(16);
|
||
let pad = g.below(k);
|
||
let opad = g.below(stride);
|
||
let (h, w) = (1 + g.below(4), 1 + g.below(5));
|
||
let (cin, cout) = (1 + g.below(4), 1 + g.below(6));
|
||
let (xr, wr) = RANGES[case % 3];
|
||
let x = g.vec(cin * h * w, xr, [0, 50, 90][case % 3]);
|
||
let wt = g.vec(cin * cout * k * k, wr, 10);
|
||
let mult: Vec<i64> = (0..cout).map(|_| 1 + g.signed(1 << 20).abs()).collect();
|
||
let bias = g.vec(cout, 1 << 20, 0);
|
||
let shift = 16 + g.below(8) as u32;
|
||
if (h - 1) * stride + k + opad > 2 * pad && (w - 1) * stride + k + opad > 2 * pad {
|
||
let want = reference::conv_t(&x, cin, h, w, &wt, cout, k, stride, pad, opad, &mult, &bias, shift);
|
||
for p in FAST {
|
||
let got = with_policy(p, || conv_t(&x, cin, h, w, &wt, cout, k, stride, pad, opad, &mult, &bias, shift));
|
||
assert_eq!(got, want, "conv_t {p:?} k={k} s={stride} p={pad} op={opad} {cin}x{h}x{w}->{cout}");
|
||
}
|
||
ran_t += 1;
|
||
}
|
||
if h + 2 * pad >= k && w + 2 * pad >= k {
|
||
let want = reference::conv(&x, cin, h, w, &wt, cout, k, stride, pad, &mult, &bias, shift);
|
||
for p in FAST {
|
||
let got = with_policy(p, || conv(&x, cin, h, w, &wt, cout, k, stride, pad, &mult, &bias, shift));
|
||
assert_eq!(got, want, "conv {p:?} k={k} s={stride} p={pad} {cin}x{h}x{w}->{cout}");
|
||
}
|
||
ran_c += 1;
|
||
}
|
||
}
|
||
assert!(ran_t > 150 && ran_c > 150, "too few valid cases ({ran_t} transposed, {ran_c} regular)");
|
||
}
|
||
|
||
/// Operands whose `i64` accumulation overflows: the reference panics in an
|
||
/// overflow-checked build and wraps otherwise, and the fast kernel must do the
|
||
/// same — a panic where it panics, the identical wrapped integers where it wraps.
|
||
#[test]
|
||
fn overflowing_accumulation_behaves_like_reference() {
|
||
let (cin, h, w, cout, k) = (4usize, 3usize, 3usize, 2usize, 3usize);
|
||
let x = vec![i32::MAX as i64; cin * h * w];
|
||
let wt = vec![i32::MAX as i64; cin * cout * k * k];
|
||
let (mult, bias) = (vec![1i64; cout], vec![0i64; cout]);
|
||
let run = |p: Policy| {
|
||
std::panic::catch_unwind(|| {
|
||
with_policy(p, || conv_t(&x, cin, h, w, &wt, cout, k, 2, 1, 1, &mult, &bias, 0))
|
||
})
|
||
};
|
||
let want = run(Policy::reference()).ok();
|
||
for p in FAST {
|
||
let got = run(p).ok();
|
||
if crate::int_kernels::testkit::overflow_checks() {
|
||
assert!(want.is_none() && got.is_none(), "both must panic on overflow in a checked build ({p:?})");
|
||
} else {
|
||
assert!(want.is_some(), "an unchecked build wraps rather than panics");
|
||
assert_eq!(got, want, "both must wrap identically ({p:?})");
|
||
}
|
||
}
|
||
}
|
||
|
||
/// The same with rows wide enough (16 outputs) for the SIMD lanes to engage: the
|
||
/// lanes wrap instead of trapping, so they must not run where an accumulator can
|
||
/// overflow — in a checked build the fast path must still panic like the reference.
|
||
#[test]
|
||
fn overflowing_accumulation_in_simd_width_rows_behaves_like_reference() {
|
||
let (cin, h, w, cout, k) = (4usize, 3usize, 16usize, 4usize, 3usize);
|
||
let x = vec![i32::MAX as i64; cin * h * w];
|
||
let wt = vec![i32::MAX as i64; cin * cout * k * k];
|
||
let (mult, bias) = (vec![1i64; cout], vec![0i64; cout]);
|
||
let run = |p: Policy| {
|
||
std::panic::catch_unwind(|| with_policy(p, || conv(&x, cin, h, w, &wt, cout, k, 1, 1, &mult, &bias, 0))).ok()
|
||
};
|
||
let want = run(Policy::reference());
|
||
for p in FAST {
|
||
let got = run(p);
|
||
if crate::int_kernels::testkit::overflow_checks() {
|
||
assert!(want.is_none() && got.is_none(), "both must panic on overflow in a checked build ({p:?})");
|
||
} else {
|
||
assert_eq!(got, want, "both must wrap identically ({p:?})");
|
||
}
|
||
}
|
||
}
|
||
|
||
/// Overflow only at a position the reference accumulates into and then crops.
|
||
/// A 1×1 input with a 3×3 kernel, stride 1 and pad 1 keeps only the centre of
|
||
/// the reference's 3×3 raw buffer. The centre is `2·2^40·1 = 2^41`. Each corner
|
||
/// is `2·2^40·2^22 = 2^63` and overflows `i64` on its second term, which the
|
||
/// fast kernel never computes. So the reference panics in an overflow-checked
|
||
/// build and the fast path must too (these operands are outside its no-overflow
|
||
/// bound, so it hands them to the reference). Without the checks both give 2^41.
|
||
#[test]
|
||
fn overflow_only_at_a_cropped_border_behaves_like_reference() {
|
||
use crate::int_kernels::testkit::overflow_checks;
|
||
let (cin, cout, k) = (2usize, 1usize, 3usize);
|
||
let x = vec![1i64 << 40; cin];
|
||
let mut wt = vec![0i64; cin * cout * k * k];
|
||
for tap in wt.chunks_mut(k * k) {
|
||
for c in [0, 2, 6, 8] {
|
||
tap[c] = 1 << 22;
|
||
}
|
||
tap[4] = 1;
|
||
}
|
||
let (mult, bias) = (vec![1i64; cout], vec![0i64; cout]);
|
||
let run = |p: Policy| {
|
||
std::panic::catch_unwind(|| with_policy(p, || conv_t(&x, cin, 1, 1, &wt, cout, k, 1, 1, 0, &mult, &bias, 0))).ok()
|
||
};
|
||
let want = run(Policy::reference());
|
||
if overflow_checks() {
|
||
assert!(want.is_none(), "the reference overflows at a cropped corner in a checked build");
|
||
} else {
|
||
assert_eq!(want, Some((vec![1i64 << 41], 1, 1)), "the kept centre is unaffected by the wrapped corners");
|
||
}
|
||
for p in FAST {
|
||
assert_eq!(run(p), want, "fast must match the reference, panic included ({p:?})");
|
||
}
|
||
}
|
||
|
||
/// Malformed geometry goes to the reference, so it fails exactly as before.
|
||
#[test]
|
||
fn malformed_geometry_is_left_to_the_reference() {
|
||
let x = vec![1i64; 4];
|
||
let short_w = vec![1i64; 3]; // needs 1·1·3·3
|
||
let r = std::panic::catch_unwind(|| conv_t(&x, 1, 2, 2, &short_w, 1, 3, 1, 0, 0, &[1], &[0], 0));
|
||
let q = std::panic::catch_unwind(|| {
|
||
with_policy(Policy::reference(), || conv_t(&x, 1, 2, 2, &short_w, 1, 3, 1, 0, 0, &[1], &[0], 0))
|
||
});
|
||
assert_eq!(r.is_err(), q.is_err());
|
||
}
|
||
|
||
/// Geometry whose padded input would be more than `PAD_FACTOR` times what the
|
||
/// reference reads and writes is refused by the fast path, and the reference
|
||
/// runs: the fast kernel is not entered and the integers are the reference's.
|
||
/// At exactly the factor the fast kernel runs. The shapes and their sizes are
|
||
/// worked in `int_kernels`' `a_padded_input_beyond_pad_factor_is_left_to_the_reference`.
|
||
#[test]
|
||
fn padding_beyond_the_pad_factor_runs_the_reference() {
|
||
use crate::int_kernels::testkit::fast_conv_runs;
|
||
let (mult, bias, shift) = ([3i64], [7i64], 1u32);
|
||
// Regular, 1×1 kernel, pad 5: (channels, input side, stride, expected
|
||
// fast-kernel runs, output). Only output (1, 1) can read the input.
|
||
for (cin, h, stride, runs, out) in [
|
||
(1usize, 2usize, 6usize, 1u64, [7i64, 7, 7, ((10 * 4 * 3) >> 1) + 7]), // at the factor
|
||
(1, 2, 7, 0, [7; 4]), // beyond it
|
||
(3, 3, 7, 0, [7, 7, 7, (((10 * 9 + 20 * 18 + 30 * 27) * 3) >> 1) + 7]), // beyond it
|
||
] {
|
||
let x: Vec<i64> = (1..=(cin * h * h) as i64).collect();
|
||
let wt: Vec<i64> = (1..=cin as i64).map(|c| 10 * c).collect();
|
||
let want = reference::conv(&x, cin, h, h, &wt, 1, 1, stride, 5, &mult, &bias, shift);
|
||
assert_eq!(want, (out.to_vec(), 2, 2), "reference, {cin}×{h}² stride {stride}");
|
||
for p in FAST {
|
||
let got = fast_conv_runs(|| with_policy(p, || conv(&x, cin, h, h, &wt, 1, 1, stride, 5, &mult, &bias, shift)));
|
||
assert_eq!(got, (want.clone(), runs), "{cin}×{h}² stride {stride} {p:?}");
|
||
}
|
||
}
|
||
// Transposed, 1×1 input and kernel, stride 1: (channels, pad, output pad,
|
||
// expected fast-kernel runs, output side). No tap reaches the kept output.
|
||
for (cin, pad, opad, runs, side) in [(2usize, 6usize, 13usize, 1u64, 2usize), (2, 7, 15, 0, 2), (1, 6, 12, 0, 1)] {
|
||
let (x, wt) = (vec![5i64; cin], vec![3i64; cin]);
|
||
let want = reference::conv_t(&x, cin, 1, 1, &wt, 1, 1, 1, pad, opad, &mult, &bias, shift);
|
||
assert_eq!(want, (vec![7; side * side], side, side), "reference, pad {pad} opad {opad}");
|
||
for p in FAST {
|
||
let got = fast_conv_runs(|| with_policy(p, || conv_t(&x, cin, 1, 1, &wt, 1, 1, 1, pad, opad, &mult, &bias, shift)));
|
||
assert_eq!(got, (want.clone(), runs), "pad {pad} opad {opad} {p:?}");
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
/// [`run_net_checked`]: it computes [`run_net`]'s integers wherever it runs, runs
|
||
/// right up to each bound, and refuses — without panicking — past it and on
|
||
/// malformed layers, including an input on which `run_net` itself overflows.
|
||
#[cfg(test)]
|
||
mod checked {
|
||
use super::*;
|
||
use crate::int_kernels::testkit::Rng;
|
||
|
||
fn conv1(weight: Vec<i64>, mult: i64, bias: i64) -> IntLayer {
|
||
IntLayer::Conv { out_ch: 1, k: 1, stride: 1, pad: 0, weight, mult: vec![mult], bias: vec![bias] }
|
||
}
|
||
fn net(shift: u32, layers: Vec<IntLayer>) -> IntModel {
|
||
IntModel { afrac: 12, gf: 16, shift, layers }
|
||
}
|
||
|
||
/// A random well-formed net of every layer kind, at operand sizes the checks pass.
|
||
fn random_net(g: &mut Rng, cin: usize) -> (IntModel, usize) {
|
||
let mut layers = Vec::new();
|
||
let mut c = cin;
|
||
for _ in 0..(2 + g.below(4)) {
|
||
let cout = 1 + g.below(5);
|
||
let k = 1 + g.below(3) * 2;
|
||
let (weight, mult, bias) = (g.vec(c * cout * k * k, 2047, 10), (0..cout).map(|_| 1 + g.below(1 << 16) as i64).collect(), g.vec(cout, 1 << 12, 0));
|
||
layers.push(match g.below(5) {
|
||
0 | 1 => IntLayer::ConvT { out_ch: cout, k, stride: 1 + g.below(2), pad: k / 2, opad: 0, weight, mult, bias },
|
||
_ => IntLayer::Conv { out_ch: cout, k, stride: 1 + g.below(2), pad: k / 2, weight, mult, bias },
|
||
});
|
||
c = cout;
|
||
match g.below(4) {
|
||
0 => layers.push(IntLayer::Relu),
|
||
1 => layers.push(IntLayer::Leaky { slope_q: 655, slope_frac: 16 }),
|
||
2 => layers.push(IntLayer::Igdn {
|
||
gamma: (0..c * c).map(|_| g.below(1 << 12) as i64).collect(),
|
||
beta: (0..c).map(|_| (1 << 24) + g.below(1 << 20) as i64).collect(),
|
||
}),
|
||
_ => {}
|
||
}
|
||
}
|
||
(net(24, layers), c)
|
||
}
|
||
|
||
#[test]
|
||
fn where_it_runs_it_is_run_net() {
|
||
let mut g = Rng(0x5EED_0001);
|
||
let mut ran = 0;
|
||
for _ in 0..200 {
|
||
let cin = 1 + g.below(4);
|
||
let (model, _) = random_net(&mut g, cin);
|
||
let (h, w) = (1 + g.below(6), 1 + g.below(6));
|
||
let x = g.vec(cin * h * w, 1 << 16, 20);
|
||
if let Ok(got) = run_net_checked(&model, &x, cin, h, w) {
|
||
assert_eq!(got, run_net(&model, &x, cin, h, w));
|
||
ran += 1;
|
||
}
|
||
}
|
||
assert!(ran > 150, "too few nets passed the checks ({ran})");
|
||
}
|
||
|
||
/// `net_shape` plans the shapes the checked run runs at: the same output, and a
|
||
/// peak that is the largest buffer the reference kernels form on the way — each
|
||
/// activation, and each transposed convolution's uncropped plane (sized here from
|
||
/// the reference kernel's own formula).
|
||
#[test]
|
||
fn net_shape_is_the_shape_the_net_runs_at() {
|
||
let mut g = Rng(0x5EED_0002);
|
||
let mut ran = 0;
|
||
for _ in 0..200 {
|
||
let cin = 1 + g.below(4);
|
||
let (model, _) = random_net(&mut g, cin);
|
||
let (h, w) = (1 + g.below(6), 1 + g.below(6));
|
||
let x = g.vec(cin * h * w, 1 << 16, 20);
|
||
let Ok((out, oc, oh, ow)) = run_net_checked(&model, &x, cin, h, w) else { continue };
|
||
let plan = net_shape(&model, cin, h, w).expect("the checked run ran, so the geometry is well formed");
|
||
assert_eq!(plan.out, (oc, oh, ow));
|
||
assert_eq!(out.len(), oc * oh * ow);
|
||
let (mut a, mut c, mut hh, mut ww) = (x.clone(), cin, h, w);
|
||
let mut peak = a.len();
|
||
for layer in &model.layers {
|
||
if let IntLayer::ConvT { out_ch, k, stride, opad, .. } = layer {
|
||
peak = peak.max(out_ch * ((hh - 1) * stride + k + opad) * ((ww - 1) * stride + k + opad));
|
||
}
|
||
(a, c, hh, ww) = apply_layer(&model, layer, a, c, hh, ww);
|
||
peak = peak.max(a.len());
|
||
}
|
||
assert_eq!(plan.peak, peak);
|
||
ran += 1;
|
||
}
|
||
assert!(ran > 150, "too few nets passed the checks ({ran})");
|
||
}
|
||
|
||
/// `NetPlan::macs` is every layer's terms at the shapes the network runs at:
|
||
/// counted here one term at a time — a convolution's every output and every tap of
|
||
/// every input channel, a transposed convolution's every input and every tap of
|
||
/// every output channel, an inverse GDN's every element and every channel — over
|
||
/// the shapes `apply_layer` forms, not the plan's.
|
||
#[test]
|
||
fn the_planned_macs_are_every_layers_terms() {
|
||
let mut g = Rng(0x5EED_0003);
|
||
let mut ran = 0;
|
||
for _ in 0..200 {
|
||
let cin = 1 + g.below(4);
|
||
let (model, _) = random_net(&mut g, cin);
|
||
let (h, w) = (1 + g.below(6), 1 + g.below(6));
|
||
let x = g.vec(cin * h * w, 1 << 16, 20);
|
||
if run_net_checked(&model, &x, cin, h, w).is_err() {
|
||
continue;
|
||
}
|
||
let plan = plan_net(&model, cin, h, w).expect("the checked run ran, so the plan admits it");
|
||
let (mut a, mut c, mut hh, mut ww) = (x, cin, h, w);
|
||
let mut terms: u128 = 0;
|
||
for layer in &model.layers {
|
||
let (na, nc, nh, nw) = apply_layer(&model, layer, a, c, hh, ww);
|
||
match layer {
|
||
IntLayer::Conv { k, .. } => {
|
||
for _output in 0..nc * nh * nw {
|
||
for _tap in 0..c * k * k {
|
||
terms += 1;
|
||
}
|
||
}
|
||
}
|
||
IntLayer::ConvT { out_ch, k, .. } => {
|
||
for _input in 0..c * hh * ww {
|
||
for _tap in 0..out_ch * k * k {
|
||
terms += 1;
|
||
}
|
||
}
|
||
}
|
||
IntLayer::Igdn { .. } => {
|
||
for _element in 0..c * hh * ww {
|
||
for _channel in 0..c {
|
||
terms += 1;
|
||
}
|
||
}
|
||
}
|
||
IntLayer::Relu | IntLayer::Leaky { .. } => {}
|
||
}
|
||
(a, c, hh, ww) = (na, nc, nh, nw);
|
||
}
|
||
assert_eq!(plan.macs, terms);
|
||
ran += 1;
|
||
}
|
||
assert!(ran > 150, "too few nets passed the checks ({ran})");
|
||
}
|
||
|
||
#[test]
|
||
fn a_convolution_runs_up_to_the_accumulation_bound_and_no_further() {
|
||
// One term: the bound is |w|·|x|. Exactly i64::MAX runs; one past refuses.
|
||
let at = net(1, vec![conv1(vec![1], 1, 0)]);
|
||
assert_eq!(run_net_checked(&at, &[i64::MAX], 1, 1, 1), Ok(run_net(&at, &[i64::MAX], 1, 1, 1)));
|
||
let past = net(1, vec![conv1(vec![2], 1, 0)]);
|
||
let x = [(i64::MAX >> 1) + 1];
|
||
assert_eq!(run_net_checked(&past, &x, 1, 1, 1).unwrap_err(),
|
||
NetRefusal { layer: 0, reason: "accumulation could leave i64", out_of_range: true });
|
||
}
|
||
|
||
// It observes run_net's panic, which needs unwinding (wasm32 aborts on a panic).
|
||
#[cfg(panic = "unwind")]
|
||
#[test]
|
||
fn it_refuses_an_input_on_which_run_net_overflows() {
|
||
// Two terms of i64::MAX and 1 sum past i64::MAX: run_net panics with overflow
|
||
// checks on and wraps without them — two decoders. The checked run refuses.
|
||
let model = net(0, vec![IntLayer::Conv { out_ch: 1, k: 1, stride: 1, pad: 0, weight: vec![1, 1], mult: vec![1], bias: vec![0] }]);
|
||
let x = [i64::MAX, 1];
|
||
assert!(run_net_checked(&model, &x, 2, 1, 1).is_err());
|
||
let plain = std::panic::catch_unwind(|| run_net(&model, &x, 2, 1, 1));
|
||
assert_eq!(plain.is_err(), crate::int_kernels::testkit::overflow_checks(), "run_net's own behaviour here depends on the build");
|
||
}
|
||
|
||
#[test]
|
||
fn the_requantised_value_must_fit() {
|
||
// Accumulation fits (2^40), but ·2^30 with no shift leaves i64.
|
||
let model = net(0, vec![conv1(vec![1], 1 << 30, 0)]);
|
||
assert_eq!(run_net_checked(&model, &[1 << 40], 1, 1, 1).unwrap_err().reason, "requantised value could leave i64");
|
||
// The bias counts, and so does one for the floor of a negative value: the
|
||
// bound is conservative by that one, and refuses a value that could reach
|
||
// i64::MAX once the bias is added.
|
||
let edge = net(0, vec![conv1(vec![1], 1, 1)]);
|
||
assert_eq!(run_net_checked(&edge, &[i64::MAX - 1], 1, 1, 1).unwrap_err().reason, "requantised value could leave i64");
|
||
}
|
||
|
||
#[test]
|
||
fn igdn_and_leaky_have_their_own_bounds() {
|
||
let igdn = |gamma: i64| net(24, vec![IntLayer::Igdn { gamma: vec![gamma], beta: vec![1 << 24] }]);
|
||
// Moderate: runs and equals run_net.
|
||
let x = [123_456i64];
|
||
assert_eq!(run_net_checked(&igdn(1 << 14), &x, 1, 1, 1), Ok(run_net(&igdn(1 << 14), &x, 1, 1, 1)));
|
||
// gamma·x² beyond i128.
|
||
assert_eq!(run_net_checked(&igdn(i64::MAX), &[i64::MAX], 1, 1, 1).unwrap_err().reason, "IGDN term could leave i128");
|
||
// A large x (2^50) whose terms and norm fit i128 but whose normalised output,
|
||
// about 2^50 · 2^50 / 2^12, leaves i64.
|
||
assert_eq!(run_net_checked(&igdn(1 << 16), &[1 << 50], 1, 1, 1).unwrap_err().reason, "IGDN output could leave i64");
|
||
// A norm between 2^126 and 2^127 fits i128, but its root (about 2^63.2) does
|
||
// not fit the i64 the reference narrows it to.
|
||
let wide = IntModel { afrac: 12, gf: 0, shift: 24, layers: vec![IntLayer::Igdn { gamma: vec![5], beta: vec![1] }] };
|
||
assert_eq!(run_net_checked(&wide, &[1 << 62], 1, 1, 1).unwrap_err().reason, "IGDN root could leave i64");
|
||
let leaky = net(24, vec![IntLayer::Leaky { slope_q: 1 << 20, slope_frac: 16 }]);
|
||
assert_eq!(run_net_checked(&leaky, &[-(1 << 50)], 1, 1, 1).unwrap_err().reason, "leaky product could leave i64");
|
||
assert!(run_net_checked(&leaky, &[-(1 << 40)], 1, 1, 1).is_ok());
|
||
}
|
||
|
||
#[test]
|
||
fn malformed_layers_are_refused_not_run() {
|
||
let conv = |weight: Vec<i64>, k: usize, stride: usize, pad: usize| IntLayer::Conv { out_ch: 1, k, stride, pad, weight, mult: vec![1], bias: vec![0] };
|
||
let cases: [(IntModel, &str); 7] = [
|
||
(net(0, vec![conv(vec![1, 1], 1, 1, 0)]), "parameter lengths do not match the running channel count"),
|
||
(net(0, vec![conv(vec![1; 9], 3, 0, 1)]), "zero output channels, kernel or stride"),
|
||
(net(0, vec![conv(vec![1; 25], 5, 1, 0)]), "kernel larger than the padded input"),
|
||
(net(0, vec![IntLayer::ConvT { out_ch: 1, k: 1, stride: 1, pad: 1, opad: 0, weight: vec![1], mult: vec![1], bias: vec![0] }]),
|
||
"padding leaves no output"),
|
||
(net(128, vec![conv(vec![1], 1, 1, 0)]), "requantisation shift of 128 or more"),
|
||
(net(0, vec![IntLayer::Leaky { slope_q: 1, slope_frac: 64 }]), "leaky slope shift of 64 or more"),
|
||
(net(0, vec![IntLayer::Igdn { gamma: vec![1, 2], beta: vec![1] }]), "IGDN parameter lengths do not match the channel count"),
|
||
];
|
||
for (model, why) in cases {
|
||
let got = std::panic::catch_unwind(|| run_net_checked(&model, &[5], 1, 1, 1));
|
||
let refusal = got.expect("a refusal, not a panic").unwrap_err();
|
||
assert_eq!(refusal.reason, why);
|
||
assert!(!refusal.out_of_range, "{why}: a malformed layer is not a range refusal");
|
||
// net_shape refuses the geometry cases alike, and leaves the parameter
|
||
// checks to the run.
|
||
let geometric = ["zero output", "kernel larger", "padding leaves"].iter().any(|g| why.starts_with(g));
|
||
match net_shape(&model, 1, 1, 1) {
|
||
Err(r) => assert!(geometric && r == refusal, "{why}: net_shape refused with {r:?}"),
|
||
Ok(_) => assert!(!geometric, "{why}: net_shape accepted a malformed geometry"),
|
||
}
|
||
}
|
||
// An input whose length is not c·h·w.
|
||
assert!(run_net_checked(&net(0, vec![]), &[1, 2, 3], 2, 1, 1).is_err());
|
||
}
|
||
|
||
/// A damaged compact model is refused: a layer count or a layer size far beyond
|
||
/// the bytes present, or one whose product overflows `usize`, never panics and
|
||
/// never has the parser reserve for it. (Run on wasm32 too, where `usize` is 32
|
||
/// bits and these products do overflow.)
|
||
#[test]
|
||
fn a_damaged_compact_model_is_refused() {
|
||
let head = |n_layers: u32| [b"WHM1".as_slice(), &12u32.to_le_bytes(), &16u32.to_le_bytes(), &24u32.to_le_bytes(),
|
||
&n_layers.to_le_bytes()].concat();
|
||
let u32s = |v: &[u32]| v.iter().flat_map(|x| x.to_le_bytes()).collect::<Vec<u8>>();
|
||
let cases: [(Vec<u8>, &str); 4] = [
|
||
(head(u32::MAX), "model truncated"),
|
||
// conv, in 2^31 × out 2^31 × k 1 × 1: the weight count overflows 32-bit usize
|
||
// and is far past the bytes on 64-bit.
|
||
([head(1), vec![1], u32s(&[1 << 31, 1 << 31, 1, 1, 0])].concat(), ""),
|
||
// conv with a weight count that fits but whose byte length (·2) does not
|
||
// fit 32-bit usize, and far exceeds the bytes on 64-bit.
|
||
([head(1), vec![1], u32s(&[u32::MAX, 1, 1, 1, 0])].concat(), "model truncated"),
|
||
// IGDN with 2^31 channels: c² overflows 32-bit usize.
|
||
([head(1), vec![2], u32s(&[1 << 31])].concat(), ""),
|
||
];
|
||
for (bytes, why) in cases {
|
||
let got = IntModel::from_bytes(&bytes).err().expect("a damaged model is refused");
|
||
assert!(got == "model truncated" || got == "layer size overflows", "{got}");
|
||
assert!(why.is_empty() || got == why, "{got}, expected {why}");
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn the_committed_int_hyper_nets_run_unrefused() {
|
||
// The tiny int_hyper vector's decode-side nets, on its own z: the checked run
|
||
// gives run_net's integers (the vector's values sit far inside every bound).
|
||
let d = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("../wai-rs/tests/vectors/int_hyper");
|
||
let h_s = IntModel::from_bytes(&std::fs::read(d.join("tiny_h_s.bin")).unwrap()).unwrap();
|
||
let z: Vec<i64> = (0..8).map(|i| ((i % 5) - 2) << 12).collect();
|
||
assert_eq!(run_net_checked(&h_s, &z, 2, 2, 2), Ok(run_net(&h_s, &z, 2, 2, 2)));
|
||
}
|
||
}
|