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]>
666 lines
32 KiB
Rust
666 lines
32 KiB
Rust
//! Integer-exact synthesis of a LEARNED 1-D audio codec — the decode side of a
|
||
//! `NeuralIntegerExact` audio codec, the audio analog of [`crate::int_hyper_synth`].
|
||
//! Pure i64, no float, no `ort`: the reconstructed int16 PCM is byte-identical on
|
||
//! every machine (*sample-equivalence*, a portable PCM hash).
|
||
//!
|
||
//! A decode through a FLOAT network lets ~1e-5 round-off make the waveform — and
|
||
//! its hash — diverge between machines. Here the synthesis `g_s` is a stack of integer 1-D
|
||
//! transposed convolutions (per-output-channel `(mult,shift)` requant) + integer
|
||
//! IGDN ([`crate::int_transform::igdn`]); the latent is recovered by the integer
|
||
//! rANS ([`crate::rans`]) under a per-channel factorized prior. Every operation is
|
||
//! i64 + an arithmetic shift, so two sinks honoring the spec produce the same PCM.
|
||
//!
|
||
//! It runs the model the export tool emits (`tools/wai_int_audio_export.py`): a
|
||
//! compact binary (`WAM1`) with int12 transposed-conv weights as `i16` and IGDN
|
||
//! `gamma` as `i32`.
|
||
|
||
use serde_json::Value;
|
||
|
||
use crate::int_transform::igdn;
|
||
use crate::rans::RansDecoder;
|
||
|
||
/// One layer of the integer synthesis `g_s`. `ConvT1d` weight is `[cin·cout·k]`
|
||
/// (layout `ic, oc, kx`); `Igdn` carries a full `c×c` gamma (`2^gf`) + beta
|
||
/// (`2^(2·afrac)`). Widened to i64 for the decode (arithmetic stays inside i64).
|
||
pub enum ALayer {
|
||
ConvT1d { out_ch: usize, k: usize, stride: usize, pad: usize, opad: usize, weight: Vec<i64>, mult: Vec<i64>, bias: Vec<i64> },
|
||
Igdn { c: usize, gamma: Vec<i64>, beta: Vec<i64> },
|
||
}
|
||
|
||
/// A parsed integer 1-D audio synthesis model (`WAM1`).
|
||
pub struct AudioModel {
|
||
pub afrac: u32,
|
||
pub gf: u32,
|
||
pub shift: u32,
|
||
pub layers: Vec<ALayer>,
|
||
}
|
||
|
||
/// 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 AudioModel {
|
||
/// Parse the export tool's compact binary model: magic `WAM1`, then
|
||
/// `afrac/gf/shift/n_layers` (u32) and per layer a `u8` tag — 0 convT1d, 2 igdn
|
||
/// — with int12 weights as `i16`, IGDN `gamma` as `i32`, the rest `i64`, and
|
||
/// nothing after the last layer.
|
||
pub fn from_bytes(b: &[u8]) -> Result<Self, &'static str> {
|
||
if b.get(0..4) != Some(b"WAM1") {
|
||
return Err("bad model magic (want WAM1)");
|
||
}
|
||
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, and reserving for it would let four damaged bytes ask for a
|
||
// large allocation.
|
||
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()?);
|
||
ALayer::ConvT1d { out_ch: out, k, stride, pad, opad,
|
||
weight: c.i16v(count(&[inc, out, k])?)?, mult: c.i64v(out)?, bias: c.i64v(out)? }
|
||
}
|
||
2 => {
|
||
let cc = c.u32()?;
|
||
ALayer::Igdn { c: cc, gamma: c.i32v(count(&[cc, cc])?)?, beta: c.i64v(cc)? }
|
||
}
|
||
_ => return Err("bad layer tag"),
|
||
});
|
||
}
|
||
// Nothing follows the last layer (`extensions/integer-payloads.md` §9).
|
||
if c.p != b.len() {
|
||
return Err("bytes follow the last layer");
|
||
}
|
||
Ok(Self { afrac, gf, shift, layers })
|
||
}
|
||
}
|
||
|
||
/// One integer 1-D transposed conv with per-output-channel requant — see
|
||
/// [`reference::conv_t_1d`] for the definition. Runs the fast kernel
|
||
/// ([`crate::int_kernels`], as a 1-row 2-D transposed conv) unless the current
|
||
/// [`crate::int_kernels::Policy`] selects the reference; both give the same integer
|
||
/// for every output element.
|
||
#[allow(clippy::too_many_arguments)]
|
||
fn conv_t_1d(
|
||
x: &[i64], cin: usize, t: usize, weight: &[i64], cout: usize, k: usize,
|
||
stride: usize, pad: usize, opad: usize, mult: &[i64], bias: &[i64], shift: u32,
|
||
) -> (Vec<i64>, 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, 1, t, cout, 1, k, 1, stride, 0, pad, 0, 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.wo);
|
||
}
|
||
reference::conv_t_1d(x, cin, t, weight, cout, k, stride, pad, opad, mult, bias, shift)
|
||
}
|
||
|
||
/// The scalar reference kernel the fast path is checked against, kept verbatim.
|
||
#[doc(hidden)]
|
||
pub mod reference {
|
||
/// One integer 1-D transposed conv with per-output-channel requant
|
||
/// `((raw·mult)>>shift)+bias`. `x` is `[cin, t]`; `weight` is `[cin][cout][k]`.
|
||
/// Returns `([cout, to], to)`. Pure i64 (i128 only for the requant multiply) —
|
||
/// mirrors the export byte-for-byte.
|
||
#[allow(clippy::too_many_arguments)]
|
||
pub fn conv_t_1d(
|
||
x: &[i64], cin: usize, t: usize, weight: &[i64], cout: usize, k: usize,
|
||
stride: usize, pad: usize, opad: usize, mult: &[i64], bias: &[i64], shift: u32,
|
||
) -> (Vec<i64>, usize) {
|
||
let to = (t - 1) * stride + k + opad - 2 * pad;
|
||
let tp = to + 2 * pad;
|
||
let mut raw = vec![0i64; cout * tp];
|
||
for ic in 0..cin {
|
||
let xc = &x[ic * t..ic * t + t];
|
||
for oc in 0..cout {
|
||
let wk = &weight[(ic * cout + oc) * k..(ic * cout + oc) * k + k];
|
||
let plane = oc * tp;
|
||
for (it, &v) in xc.iter().enumerate() {
|
||
if v == 0 {
|
||
continue;
|
||
}
|
||
let base = plane + it * stride;
|
||
for kx in 0..k {
|
||
raw[base + kx] += wk[kx] * v;
|
||
}
|
||
}
|
||
}
|
||
}
|
||
let mut out = vec![0i64; cout * to];
|
||
for oc in 0..cout {
|
||
let (m, b) = (mult[oc] as i128, bias[oc]);
|
||
for ot in 0..to {
|
||
let r = raw[oc * tp + ot + pad] as i128;
|
||
out[oc * to + ot] = ((r * m) >> shift) as i64 + b;
|
||
}
|
||
}
|
||
(out, to)
|
||
}
|
||
}
|
||
|
||
fn i64s(v: &Value) -> Vec<i64> {
|
||
v.as_array().map(|a| a.iter().filter_map(|n| n.as_i64()).collect()).unwrap_or_default()
|
||
}
|
||
|
||
use crate::int_hyper_synth::{refused, Refused, RefusalKind, MAX_AFRAC, MAX_TENSOR};
|
||
|
||
/// A clip's coded form, typed: what a `WIA1` payload carries
|
||
/// (`extensions/integer-payloads.md` §5), and what the tool's JSON form carries, read
|
||
/// by [`AudioJson`]. `cdfs` holds one table per latent channel, all of one length.
|
||
#[derive(Debug, Clone, Copy)]
|
||
pub struct AudioStream<'a> {
|
||
/// The fixed-point scale the stream was coded at; it must equal the model's.
|
||
pub afrac: u32,
|
||
pub n_ch: usize,
|
||
pub lat_len: usize,
|
||
/// Samples presented: the first `t_out` of the synthesis.
|
||
pub t_out: usize,
|
||
/// The symbol offset every channel's table is coded at (the JSON form's `lo`).
|
||
pub offset: i32,
|
||
pub median_q: &'a [i64],
|
||
pub cdfs: &'a [Vec<i32>],
|
||
pub rans: &'a [u8],
|
||
}
|
||
|
||
/// The tool's JSON bitstream (`tools/wai_int_audio_export.py`'s
|
||
/// `audio_bitstream.json`), read into the typed form. A converter input and a
|
||
/// debugging aid; the payload is `WIA1`. What the old JSON decoder left implicit is
|
||
/// explicit here: its `afrac` is read (the old decoder ignored it) and must be the
|
||
/// model's; every channel's table must have the first table's length (the old decoder
|
||
/// read every table at that length).
|
||
#[derive(Debug, Clone)]
|
||
pub struct AudioJson {
|
||
pub afrac: u32,
|
||
pub n_ch: usize,
|
||
pub lat_len: usize,
|
||
pub t_out: usize,
|
||
pub offset: i32,
|
||
pub median_q: Vec<i64>,
|
||
pub cdfs: Vec<Vec<i32>>,
|
||
pub rans: Vec<u8>,
|
||
}
|
||
|
||
impl AudioJson {
|
||
/// Read the tool's JSON bitstream.
|
||
pub fn from_value(bs: &Value) -> Result<Self, &'static str> {
|
||
let uz = |k: &'static str| bs[k].as_u64().and_then(|v| usize::try_from(v).ok()).ok_or(k);
|
||
let (n_ch, lat_len, t_out) = (uz("n_ch")?, uz("lat_len")?, uz("t_out")?);
|
||
let afrac = bs["afrac"].as_u64().and_then(|v| u32::try_from(v).ok()).ok_or("afrac")?;
|
||
let offset = bs["lo"].as_i64().and_then(|v| i32::try_from(v).ok()).ok_or("lo")?;
|
||
let median_q = i64s(&bs["median_q"]);
|
||
let cdfs = bs["cdfs"].as_array().ok_or("cdfs")?.iter()
|
||
.map(|row| row.as_array().and_then(|r| r.iter().map(|n| n.as_i64().and_then(|v| i32::try_from(v).ok())).collect::<Option<Vec<i32>>>()))
|
||
.collect::<Option<Vec<_>>>().ok_or("cdfs holds a value outside i32")?;
|
||
let first = cdfs.first().ok_or("empty cdfs")?.len();
|
||
if cdfs.iter().any(|c| c.len() != first) {
|
||
return Err("the channels' tables are not of one length");
|
||
}
|
||
let rans = b64_decode(bs["rans_b64"].as_str().ok_or("rans_b64")?)?;
|
||
Ok(AudioJson { afrac, n_ch, lat_len, t_out, offset, median_q, cdfs, rans })
|
||
}
|
||
|
||
/// The typed form, borrowing this one's tables.
|
||
pub fn stream(&self) -> AudioStream<'_> {
|
||
AudioStream {
|
||
afrac: self.afrac, n_ch: self.n_ch, lat_len: self.lat_len, t_out: self.t_out, offset: self.offset,
|
||
median_q: &self.median_q, cdfs: &self.cdfs, rans: &self.rans,
|
||
}
|
||
}
|
||
}
|
||
|
||
/// The model's output `[c, t]` on a `[c, t]` latent, planned from its layers' geometry
|
||
/// and parameter lengths alone, before anything is allocated
|
||
/// (`extensions/integer-payloads.md` §7). Layer by layer in order: the layer fits the
|
||
/// running channel count (`ModelMismatch`: output channels, kernel and stride
|
||
/// non-zero, `c·out·k` weights, `out` multipliers and biases, a requantisation shift
|
||
/// under 128; an inverse GDN of the running channel count, with `c·c` and `c`
|
||
/// parameters and shifts under 128); the uncropped length `(t − 1)·stride + k + opad`
|
||
/// and the buffer of `out` channels of it are within [`MAX_TENSOR`] (`TooLarge`); and
|
||
/// the crop `2·pad` leaves an output (`ModelMismatch`). Every size is formed in `u64`
|
||
/// with checked arithmetic, a size past `u64` being over the cap, so every target
|
||
/// refuses alike. Returns the output `[c, t]`, the most elements one layer holds, and
|
||
/// the multiply-accumulates the layers comprise (`extensions/integer-payloads.md` §7):
|
||
/// a transposed convolution's every input times `out·k`, an inverse GDN's every
|
||
/// element times `c`, summed in `u128`.
|
||
fn plan(model: &AudioModel, c: usize, t: usize) -> Result<(usize, usize, u64, u128), Refused> {
|
||
use RefusalKind::*;
|
||
const CAP: u64 = MAX_TENSOR as u64;
|
||
let (mut c, mut t) = (c as u64, t as u64);
|
||
// The most elements the reference kernels hold while one layer runs, its input
|
||
// included: a transposed convolution's input, uncropped buffer and output; an
|
||
// inverse GDN's input, output and one position's squares (`c` values of 16 bytes).
|
||
let mut live = c * t;
|
||
let mut macs: u128 = 0;
|
||
for layer in &model.layers {
|
||
match layer {
|
||
ALayer::ConvT1d { 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 {
|
||
return Err(refused(ModelMismatch, "a layer has zero output channels, kernel or stride"));
|
||
}
|
||
if c.checked_mul(cout).and_then(|v| v.checked_mul(k)) != Some(weight.len() as u64)
|
||
|| mult.len() as u64 != cout || bias.len() as u64 != cout || model.shift > 127
|
||
{
|
||
return Err(refused(ModelMismatch, "a layer's parameters do not match the running channel count"));
|
||
}
|
||
let full = (t - 1).checked_mul(s).and_then(|v| v.checked_add(k)).and_then(|v| v.checked_add(op));
|
||
let full = full.filter(|&n| n <= CAP).ok_or(refused(TooLarge, "a layer's uncropped length exceeds 2^26"))?;
|
||
if cout.checked_mul(full).is_none_or(|n| n > CAP) {
|
||
return Err(refused(TooLarge, "a layer would form a buffer over 2^26 elements"));
|
||
}
|
||
let crop = p.checked_mul(2).unwrap_or(u64::MAX);
|
||
if full <= crop {
|
||
return Err(refused(ModelMismatch, "a layer's padding leaves no output"));
|
||
}
|
||
live = live.max(c * t + cout * full + cout * (full - crop));
|
||
macs = macs.saturating_add((c * t) as u128 * (cout * k) as u128);
|
||
(c, t) = (cout, full - crop);
|
||
}
|
||
ALayer::Igdn { c: cc, gamma, beta } => {
|
||
if *cc as u64 != c || c.checked_mul(c) != Some(gamma.len() as u64) || beta.len() as u64 != c
|
||
|| model.gf > 127 || model.afrac > 127
|
||
{
|
||
return Err(refused(ModelMismatch, "an IGDN layer does not match the running channel count"));
|
||
}
|
||
live = live.max(2 * c * t + 2 * c);
|
||
macs = macs.saturating_add((c * t) as u128 * c as u128);
|
||
}
|
||
}
|
||
}
|
||
// Every running size is at most 2^26 here, so it fits usize on every target.
|
||
Ok((c as usize, t as usize, live, macs))
|
||
}
|
||
|
||
/// What [`decode_typed`] checks before it decodes a symbol, and the buffers it then
|
||
/// forms. Returns the synthesis's length, the most bytes the decode holds at once
|
||
/// on the reference kernels beyond the payload and the parsed model (the latent and
|
||
/// each layer's buffers, `i64` each, then the synthesis's output, `i64`, with the
|
||
/// PCM, `i16`), and the synthesis's multiply-accumulates on the latent.
|
||
///
|
||
/// The latent's length is the one `t_out` needs and no longer: the synthesis of
|
||
/// `lat_len` latent samples gives at least `t_out` samples (`GeometryMismatch`
|
||
/// otherwise), and that of `lat_len − 1` gives fewer (`GeometryMismatch`), so a
|
||
/// payload cannot declare a large latent to present a short clip. A length for which
|
||
/// the model leaves no output gives none.
|
||
pub(crate) fn audio_plan(model: &AudioModel, s: &AudioStream<'_>) -> Result<(usize, u64, u128), Refused> {
|
||
use RefusalKind::*;
|
||
let afrac = model.afrac;
|
||
if !(1..=MAX_AFRAC).contains(&afrac) {
|
||
return Err(refused(ModelMismatch, "the model's fixed-point scale is outside 1..=30"));
|
||
}
|
||
if s.afrac != afrac {
|
||
return Err(refused(ModelMismatch, "the stream's afrac is not the model's"));
|
||
}
|
||
let (c, l) = (s.n_ch, s.lat_len);
|
||
if c == 0 || l == 0 || s.t_out == 0 || s.median_q.len() != c || s.cdfs.len() != c {
|
||
return Err(refused(FieldRange, "a dimension is zero, or the medians or tables are not one per channel"));
|
||
}
|
||
if s.median_q.iter().any(|m| m.unsigned_abs() > 1 << 47) {
|
||
return Err(refused(FieldRange, "a median's magnitude exceeds 2^47"));
|
||
}
|
||
let len = s.cdfs[0].len();
|
||
let ok = |t: &Vec<i32>| t.len() == len && (3..=4097).contains(&len) && t[0] == 0 && t[len - 1] == 1 << 16
|
||
&& t.windows(2).all(|p| p[0] < p[1]);
|
||
if !s.cdfs.iter().all(ok) {
|
||
return Err(refused(CdfMalformed, "a table is not a CDF from 0 to 65536 of the channels' one length"));
|
||
}
|
||
if c.checked_mul(l).is_none_or(|n| n > MAX_TENSOR) || s.t_out > MAX_TENSOR {
|
||
return Err(refused(TooLarge, "the latent or the clip exceeds 2^26 elements"));
|
||
}
|
||
let (oc, t, live, macs) = plan(model, c, l)?;
|
||
if oc != 1 {
|
||
return Err(refused(ModelMismatch, "the model does not give one channel"));
|
||
}
|
||
if t < s.t_out {
|
||
return Err(refused(GeometryMismatch, "the synthesis gives fewer samples than t_out"));
|
||
}
|
||
let shorter = if l > 1 { plan(model, c, l - 1).map_or(0, |(_, t, ..)| t) } else { 0 };
|
||
if shorter >= s.t_out {
|
||
return Err(refused(GeometryMismatch, "the latent is longer than t_out needs: one sample fewer gives t_out samples"));
|
||
}
|
||
let working = (8 * live).max(8 * t as u64 + 2 * s.t_out as u64);
|
||
Ok((t, working, macs))
|
||
}
|
||
|
||
/// One layer's arithmetic on the values `x` holds, checked before it runs, as
|
||
/// [`crate::int_hyper_synth::run_net_checked`] checks a 2-D layer: every product and
|
||
/// partial sum and the requantised value of a convolution fit `i64`; IGDN's terms,
|
||
/// norm and root fit their types.
|
||
fn layer_in_range(model: &AudioModel, layer: &ALayer, x: &[i64], c: usize) -> Result<(), Refused> {
|
||
use crate::int_kernels::max_abs;
|
||
const I64: u128 = i64::MAX as u128;
|
||
match layer {
|
||
ALayer::ConvT1d { k, weight, mult, bias, .. } => {
|
||
if model.shift > 127 {
|
||
return Err(refused(RefusalKind::ModelMismatch, "requantisation shift of 128 or more"));
|
||
}
|
||
let terms = (c as u128).saturating_mul(*k as u128);
|
||
let acc = terms.saturating_mul(max_abs(weight) as u128).saturating_mul(max_abs(x) as u128);
|
||
if acc > I64 {
|
||
return Err(refused(RefusalKind::FieldRange, "a latent could take an accumulation outside i64"));
|
||
}
|
||
let q = ((acc * max_abs(mult) as u128) >> model.shift) + 1;
|
||
if q.saturating_add(max_abs(bias) as u128) > I64 {
|
||
return Err(refused(RefusalKind::FieldRange, "a latent could take a requantised value outside i64"));
|
||
}
|
||
Ok(())
|
||
}
|
||
ALayer::Igdn { gamma, beta, .. } => crate::int_hyper_synth::igdn_check(x, c, gamma, beta, model.gf, model.afrac)
|
||
.map_err(|(why, values)| refused(if values { RefusalKind::FieldRange } else { RefusalKind::ModelMismatch }, why)),
|
||
}
|
||
}
|
||
|
||
/// Decode a typed clip → int16 PCM, mono, byte-exact on every machine, refusing
|
||
/// rather than panicking or overflowing on any input:
|
||
///
|
||
/// 1. the stream's `afrac` is the model's, in `1..=30`; every table is well formed;
|
||
/// the model's layers fit the latent's shape, give one channel, and give at least
|
||
/// `t_out` samples from a latent no longer than that needs, with no buffer over
|
||
/// [`MAX_TENSOR`] elements — all before anything is allocated ([`audio_plan`]);
|
||
/// 2. per channel, rANS-decode the symbol `s`; the latent is `s·2^afrac + median_q`
|
||
/// ([`crate::int_entropy::dequantize`]), and a stream that runs out before its last
|
||
/// symbol is refused;
|
||
/// 3. integer synthesis `g_s`, each layer's arithmetic checked before it runs;
|
||
/// 4. the first `t_out` samples, `(x·32767 + 2^(afrac−1)) >> afrac` clamped to `i16`.
|
||
pub fn decode_typed(model: &AudioModel, s: &AudioStream<'_>) -> Result<Vec<i16>, Refused> {
|
||
use crate::int_entropy::dequantize;
|
||
use RefusalKind::*;
|
||
audio_plan(model, s)?;
|
||
let (afrac, c, l, len) = (model.afrac, s.n_ch, s.lat_len, s.cdfs[0].len());
|
||
|
||
// 1. factorized rANS-decode the per-channel latent (channel-major), rebuild y_q.
|
||
// The latent grows a symbol at a time, and the first symbol that runs the stream
|
||
// out refuses the payload before its value is used.
|
||
let mut y = Vec::new();
|
||
let mut dec = RansDecoder::new(s.rans).map_err(|_| refused(FieldRange, "the latent stream is not a rANS stream"))?;
|
||
for ch in 0..c {
|
||
for _ in 0..l {
|
||
let sym = dec.decode(&s.cdfs[ch], len, s.offset);
|
||
if dec.is_exhausted() {
|
||
return Err(refused(StreamExhausted, "the latent stream ended before its last symbol"));
|
||
}
|
||
crate::int_hyper_synth::push_capped(&mut y, dequantize(sym, afrac, s.median_q[ch]).ok_or(refused(FieldRange, "a latent value leaves i64"))?, c * l);
|
||
}
|
||
}
|
||
if !dec.is_canonical_end() {
|
||
return Err(refused(StreamNoncanonical, "the latent stream does not end canonically after its last symbol"));
|
||
}
|
||
// 2. integer synthesis → waveform, each layer checked on the values it is given.
|
||
let (mut x, mut cc, mut tt) = (y, c, l);
|
||
for layer in &model.layers {
|
||
layer_in_range(model, layer, &x, cc)?;
|
||
(x, cc, tt) = run_layer(model, layer, x, cc, tt);
|
||
}
|
||
// 3. waveform → int16 PCM (round-half-up, clamp): the first t_out samples.
|
||
let half = 1i64 << (afrac - 1);
|
||
let mut pcm = Vec::with_capacity(s.t_out);
|
||
for &v in &x[..s.t_out] {
|
||
let num = v.checked_mul(32767).and_then(|n| n.checked_add(half))
|
||
.ok_or(refused(FieldRange, "a synthesised sample leaves i64 when scaled"))?;
|
||
pcm.push((num >> afrac).clamp(-32768, 32767) as i16);
|
||
}
|
||
Ok(pcm)
|
||
}
|
||
|
||
/// One layer of the synthesis on `x` (`[c, t]`): the new activations and shape.
|
||
fn run_layer(model: &AudioModel, layer: &ALayer, x: Vec<i64>, c: usize, t: usize) -> (Vec<i64>, usize, usize) {
|
||
match layer {
|
||
ALayer::ConvT1d { out_ch, k, stride, pad, opad, weight, mult, bias } => {
|
||
let (nx, to) = conv_t_1d(&x, c, t, weight, *out_ch, *k, *stride, *pad, *opad, mult, bias, model.shift);
|
||
(nx, *out_ch, to)
|
||
}
|
||
ALayer::Igdn { gamma, beta, .. } => (igdn(&x, c, t, model.afrac, gamma, beta, model.gf), c, t),
|
||
}
|
||
}
|
||
|
||
/// Decode the tool's JSON bitstream (`audio_bitstream.json`) → int16 PCM:
|
||
/// [`AudioJson`] then [`decode_typed`]. A debugging and converter path; the payload
|
||
/// is `WIA1`. Byte-identical on every machine.
|
||
pub fn decode(model: &AudioModel, bs: &Value) -> Result<Vec<i16>, &'static str> {
|
||
decode_typed(model, &AudioJson::from_value(bs)?.stream()).map_err(|r| r.detail)
|
||
}
|
||
|
||
/// FNV-1a 64-bit over the little-endian int16 PCM bytes — the dependency-free
|
||
/// checksum the export tool and this decoder share to prove byte-for-byte agreement
|
||
/// across languages (sample-equivalence).
|
||
pub fn pcm_fnv1a(pcm: &[i16]) -> u64 {
|
||
let mut hash: u64 = 14695981039346656037;
|
||
for &s in pcm {
|
||
for &b in &s.to_le_bytes() {
|
||
hash = (hash ^ b as u64).wrapping_mul(1099511628211);
|
||
}
|
||
}
|
||
hash
|
||
}
|
||
|
||
/// Dependency-free base64 decode (standard alphabet, no padding required) — the
|
||
/// rANS bytes ride in JSON as base64, same as `int_video`.
|
||
fn b64_decode(s: &str) -> Result<Vec<u8>, &'static str> {
|
||
fn val(c: u8) -> Result<u8, &'static str> {
|
||
Ok(match c {
|
||
b'A'..=b'Z' => c - b'A',
|
||
b'a'..=b'z' => c - b'a' + 26,
|
||
b'0'..=b'9' => c - b'0' + 52,
|
||
b'+' => 62,
|
||
b'/' => 63,
|
||
_ => return Err("bad base64 char"),
|
||
})
|
||
}
|
||
let s: Vec<u8> = s.bytes().filter(|&c| c != b'=' && !c.is_ascii_whitespace()).collect();
|
||
let mut out = Vec::with_capacity(s.len() * 3 / 4);
|
||
for chunk in s.chunks(4) {
|
||
let mut acc = 0u32;
|
||
let mut bits = 0;
|
||
for &c in chunk {
|
||
acc = (acc << 6) | val(c)? as u32;
|
||
bits += 6;
|
||
}
|
||
// each full byte sits in the high `bits-8k` positions
|
||
for k in 0..bits / 8 {
|
||
out.push((acc >> (bits - 8 * (k + 1))) as u8);
|
||
}
|
||
}
|
||
Ok(out)
|
||
}
|
||
|
||
/// The fast 1-D transposed conv against the scalar reference, element by element.
|
||
#[cfg(test)]
|
||
mod plan_tests {
|
||
use super::*;
|
||
|
||
/// The committed synthesis's planned multiply-accumulates on the committed latent are
|
||
/// its layers' terms at the lengths the reference kernel gives (run here on a zero
|
||
/// latent): a transposed convolution's every input times `out·k`, an inverse GDN's
|
||
/// every element times `c`.
|
||
#[test]
|
||
fn the_planned_macs_are_every_layers_terms() {
|
||
let dir = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("tests/vectors/int_audio");
|
||
let model = AudioModel::from_bytes(&std::fs::read(dir.join("audio_model.bin")).unwrap()).unwrap();
|
||
let wia1 = std::fs::read(dir.join("audio_bitstream.wia1")).unwrap();
|
||
let a = crate::int_wire::parse_wia1(&wia1).unwrap();
|
||
let (c0, l) = (a.n_ch as usize, a.lat_len as usize);
|
||
let (.., macs) = plan(&model, c0, l).unwrap();
|
||
let (mut c, mut t, mut terms) = (c0, l, 0u128);
|
||
for layer in &model.layers {
|
||
match layer {
|
||
ALayer::ConvT1d { out_ch, k, stride, pad, opad, weight, mult, bias } => {
|
||
let (_, to) = reference::conv_t_1d(&vec![0; c * t], c, t, weight, *out_ch, *k, *stride, *pad, *opad, mult, bias, model.shift);
|
||
terms += (c * t) as u128 * (out_ch * k) as u128;
|
||
(c, t) = (*out_ch, to);
|
||
}
|
||
ALayer::Igdn { .. } => terms += (c * t) as u128 * c as u128,
|
||
}
|
||
}
|
||
assert!(terms > 0);
|
||
assert_eq!(macs, terms);
|
||
}
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod kernel_equivalence {
|
||
use super::*;
|
||
use crate::int_kernels::testkit::{Rng, FAST};
|
||
use crate::int_kernels::with_policy;
|
||
|
||
#[test]
|
||
fn conv_t_1d_equals_reference() {
|
||
let mut g = Rng(0xA0D_0001);
|
||
let mut ran = 0;
|
||
for case in 0..500 {
|
||
let k = 1 + g.below(9);
|
||
let stride = 1 + g.below(5);
|
||
let pad = g.below(k);
|
||
let opad = if stride > 1 { g.below(stride) } else { 0 };
|
||
let t = 1 + g.below(40);
|
||
let (cin, cout) = (1 + g.below(6), 1 + g.below(6));
|
||
if (t - 1) * stride + k + opad <= 2 * pad {
|
||
continue;
|
||
}
|
||
let (xr, wr) = [(1 << 15, 2047), (1 << 40, 1 << 10)][case % 2];
|
||
let x = g.vec(cin * t, xr, 30);
|
||
let wt = g.vec(cin * cout * k, wr, 5);
|
||
let mult: Vec<i64> = (0..cout).map(|_| 1 + g.signed(1 << 16).abs()).collect();
|
||
let bias = g.vec(cout, 1 << 16, 0);
|
||
let want = reference::conv_t_1d(&x, cin, t, &wt, cout, k, stride, pad, opad, &mult, &bias, 24);
|
||
for p in FAST {
|
||
let got = with_policy(p, || conv_t_1d(&x, cin, t, &wt, cout, k, stride, pad, opad, &mult, &bias, 24));
|
||
assert_eq!(got, want, "conv_t_1d {p:?} k={k} s={stride} p={pad} op={opad} t={t}");
|
||
}
|
||
ran += 1;
|
||
}
|
||
assert!(ran > 250, "too few valid cases ({ran})");
|
||
}
|
||
|
||
/// Overflow only at a position the reference accumulates into and then crops:
|
||
/// a length-1 input, kernel 3, pad 1 keeps only the middle of the reference's
|
||
/// 3-sample raw buffer (`2·2^40·1 = 2^41`); both outer samples are
|
||
/// `2·2^40·2^22 = 2^63` and overflow `i64`. The fast path must panic where the
|
||
/// reference does and otherwise give the same 2^41.
|
||
#[test]
|
||
fn overflow_only_at_a_cropped_border_behaves_like_reference() {
|
||
use crate::int_kernels::testkit::overflow_checks;
|
||
use crate::int_kernels::Policy;
|
||
let (cin, cout, k) = (2usize, 1usize, 3usize);
|
||
let x = vec![1i64 << 40; cin];
|
||
let wt: Vec<i64> = (0..cin * cout).flat_map(|_| [1i64 << 22, 1, 1 << 22]).collect();
|
||
let (mult, bias) = (vec![1i64; cout], vec![0i64; cout]);
|
||
let run = |p: Policy| {
|
||
std::panic::catch_unwind(|| with_policy(p, || conv_t_1d(&x, cin, 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 sample in a checked build");
|
||
} else {
|
||
assert_eq!(want, Some((vec![1i64 << 41], 1)));
|
||
}
|
||
for p in FAST {
|
||
assert_eq!(run(p), want, "fast must match the reference, panic included ({p:?})");
|
||
}
|
||
}
|
||
|
||
/// A one-sample input upsampled by any stride is one output sample (k = 3,
|
||
/// pad 1): a handful of operations for the reference. The fast path must return
|
||
/// the same sample without overflowing, or allocating or looping per stride step.
|
||
#[test]
|
||
fn a_one_sample_input_upsampled_by_any_stride() {
|
||
use crate::int_kernels::Policy;
|
||
let (x, wt) = ([5i64], [1i64, 2, 3]);
|
||
// usize::MAX first: a kernel allocating a table per phase fails there at once
|
||
// ("capacity overflow") instead of committing memory for 2^29 tables later.
|
||
for s in [usize::MAX, 1, 2, 1 << 13, 1 << 16, 1 << 29] {
|
||
let run = |p: Policy| with_policy(p, || conv_t_1d(&x, 1, 1, &wt, 1, 3, s, 1, 0, &[1], &[0], 0));
|
||
let want = run(Policy::reference());
|
||
assert_eq!(want, (vec![10], 1), "stride {s}: the middle of 5·[1, 2, 3]");
|
||
for p in FAST {
|
||
assert_eq!(run(p), want, "stride {s} {p:?}");
|
||
}
|
||
}
|
||
}
|
||
|
||
/// Several input samples upsampled by a stride and cropped by padding so large
|
||
/// that the reference's uncropped buffer `cout · (to + 2·pad)` cannot be sized:
|
||
/// three samples, kernel 1, stride = pad = 2^62, one output sample. With two
|
||
/// output channels the size overflows `usize` (the reference panics with
|
||
/// overflow checks on, and with them off indexes out of the wrapped buffer);
|
||
/// with one it is 2^63 + 1 samples, more than any allocation can hold ("capacity
|
||
/// overflow"). The reference panics in every build, so every policy must, and
|
||
/// the fast kernel must not be entered.
|
||
#[cfg(target_pointer_width = "64")]
|
||
#[test]
|
||
fn several_samples_whose_uncropped_buffer_cannot_be_sized() {
|
||
use crate::int_kernels::testkit::fast_conv_runs;
|
||
use crate::int_kernels::Policy;
|
||
let s = 1usize << 62;
|
||
for cout in [2usize, 1] {
|
||
let (x, wt) = ([1i64, 2, 3], vec![7i64; cout]);
|
||
let (mult, bias) = (vec![1i64; cout], vec![0i64; cout]);
|
||
let run = |p: Policy| {
|
||
fast_conv_runs(|| std::panic::catch_unwind(|| with_policy(p, || conv_t_1d(&x, 1, 3, &wt, cout, 1, s, s, 0, &mult, &bias, 0))).ok())
|
||
};
|
||
assert_eq!(run(Policy::reference()), (None, 0), "{cout} channels: the reference panics");
|
||
for p in FAST {
|
||
assert_eq!(run(p), (None, 0), "{cout} channels {p:?}");
|
||
}
|
||
}
|
||
}
|
||
|
||
/// Geometry whose padded input would be more than `PAD_FACTOR` (16) 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 samples are the reference's. At
|
||
/// exactly the factor the fast kernel runs. A one-sample input, kernel 1, stride
|
||
/// 1, cropped by `pad` with output padding `2·pad` restoring one output sample
|
||
/// that no tap reaches: the padded input is `cin · (1 + pad)` samples against
|
||
/// `cin + cin + 1` the reference touches (input, weights, output).
|
||
#[test]
|
||
fn padding_beyond_the_pad_factor_runs_the_reference() {
|
||
use crate::int_kernels::testkit::fast_conv_runs;
|
||
// (channels, pad, expected fast-kernel runs): 1 · 48 = 16 · 3 and
|
||
// 2 · 40 = 16 · 5 at the factor; 1 · 49 and 2 · 41 beyond it.
|
||
for (cin, pad, runs) in [(1usize, 47usize, 1u64), (1, 48, 0), (2, 39, 1), (2, 40, 0)] {
|
||
let (x, wt) = (vec![5i64; cin], vec![3i64; cin]);
|
||
let want = reference::conv_t_1d(&x, cin, 1, &wt, 1, 1, 1, pad, 2 * pad, &[3], &[7], 1);
|
||
assert_eq!(want, (vec![7], 1), "reference, {cin} channels, pad {pad}");
|
||
for p in FAST {
|
||
let got = fast_conv_runs(|| with_policy(p, || conv_t_1d(&x, cin, 1, &wt, 1, 1, 1, pad, 2 * pad, &[3], &[7], 1)));
|
||
assert_eq!(got, (want.clone(), runs), "{cin} channels, pad {pad} {p:?}");
|
||
}
|
||
}
|
||
}
|
||
}
|