open-standards/wai/wai-rs/src/int_synth.rs

//! Integer fixed-point synthesis — a learned-image decoder whose reconstruction
//! is **bit-identical on every machine, by construction**.
//!
//! This is the deterministic member of the neural family. The integer rANS
//! recovers latents exactly; the divergence in float decoders like
//! `bmshj2018` / `mbt2018_mean` is the **float ONNX synthesis**, whose round-off
//! depends on the kernels, accumulation order and fused operations the runtime
//! picks for each platform. Here the synthesis is a stack of integer
//! transposed convolutions with integer activations: every operation is `i64`
//! arithmetic + an arithmetic shift, so two sinks honoring the spec produce the
//! same bytes everywhere. No ONNX, no float, nothing implementation-defined.
//!
//! Because it pulls in no float runtime, this decoder lives at the crate root
//! behind the light `neural_int` feature (just `serde_json`), not inside the
//! `ort`-backed `neural` module. That keeps it buildable — and its
//! cross-hardware bit-exactness checkable — in a CI matrix with no ONNX
//! Runtime present. The `neural` module re-exposes it through the
//! `wai.neural.int_synth` capability.
//!
//! Quality depends on how the integer weights were made, which is a tooling step
//! outside the decode contract. The reference tools train in float, then quantise
//! after training, deriving each layer's shift from its measured activation range
//! (`tools/wai_int_synth_train.py`, `tools/wai_int_synth_real.py`); the trained
//! model in the conformance corpus was made that way. The *decode contract* — the
//! thing that makes cross-hardware bit-exactness real — lives here and is exact
//! whatever the weights.
//!
//! Wire format (little-endian), `extensions/integer-payloads.md` §13:
//!   u32 H, u32 W, u16 C, u8 S_log2, u8 _reserved   (12 B)
//!   C*S*S bytes   int8 latents, channel-major
//!
//! Sink file (next to the registered model path): `int_synth.json`
//!   { "layers": [ { in_ch, out_ch, k, stride, pad, out_pad,
//!                   weight:[i32; in*out*k*k] (layout ic,oc,ky,kx),
//!                   bias:[i32; out_ch], shift:u32, relu:bool }, … ] }
//! The final layer MUST output 3 channels at H×W; values are clamped to [0,255].
//!
//! The decode ([`decode_from_spec`]) is the graph `synth` of the set's lowering
//! `int_synth/1` (`extensions/integer-operators.md` §13.1), loaded under `acc64` with its
//! bounds certificate and run through the integer operator set ([`crate::int_ops`]),
//! whose interpreter lends the transposed convolutions to the fast kernels. The direct
//! decode it replaced is kept, test-only (feature `legacy_int`), as `legacy`: each
//! layer's transposed convolution, bias, shift and ReLU computed here by the scalar
//! kernel. The differential tests decode random sets and payloads both ways
//! (`tests/int_ops_differential.rs`).

use std::path::Path;

#[cfg(any(test, feature = "legacy_int"))]
use serde::Deserialize;

const HEADER: usize = 12;

/// Decode error for the integer synthesis path. Deliberately self-contained so
/// the decoder carries no dependency on the `ort`-backed `neural` module; the
/// `neural` dispatch maps this onto its own `DecodeError::InvalidPayload`.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum IntSynthError {
    InvalidPayload(String),
}

impl std::fmt::Display for IntSynthError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::InvalidPayload(e) => write!(f, "invalid payload: {e}"),
        }
    }
}
impl std::error::Error for IntSynthError {}

/// Decoded image: RGB row-major u8, H × W × 3 — the integer-exact output.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct IntImage {
    pub rgb: Vec<u8>,
    pub width: u32,
    pub height: u32,
}

#[cfg(any(test, feature = "legacy_int"))]
#[derive(Debug, Clone, Deserialize)]
struct LayerSpec {
    in_ch: usize,
    out_ch: usize,
    k: usize,
    stride: usize,
    pad: usize,
    #[serde(default)]
    out_pad: usize,
    /// Integer weights, layout `[in_ch][out_ch][k][k]` (transposed-conv order).
    weight: Vec<i32>,
    /// Integer bias, `[out_ch]`.
    bias: Vec<i32>,
    /// Arithmetic right-shift applied after accumulate+bias (requantize).
    #[serde(default)]
    shift: u32,
    /// Integer ReLU (clamp at 0) after requantize.
    #[serde(default)]
    relu: bool,
}

#[cfg(any(test, feature = "legacy_int"))]
#[derive(Debug, Clone, Deserialize)]
struct SynthSpec {
    layers: Vec<LayerSpec>,
}

/// The direct decode's transposed convolution, kept verbatim: the scalar kernel its fast
/// dispatch was checked against. That dispatch, and the guard that kept it to sizes this
/// kernel's own arithmetic forms, are gone: the decode runs through the integer operator
/// set, whose plan forms every size in checked arithmetic under its caps.
#[cfg(any(test, feature = "legacy_int"))]
mod reference {
    use super::{IntSynthError, LayerSpec};

    /// One integer transposed convolution + bias + requantize + optional ReLU.
    /// `input` is `[in_ch, hin, win]` row-major `i64`. Returns `([out_ch, hout,
    /// wout] i64, hout, wout)`. Pure integer arithmetic — the whole point.
    pub(super) fn conv_transpose2d(
        input: &[i64],
        cin: usize,
        hin: usize,
        win: usize,
        l: &LayerSpec,
    ) -> Result<(Vec<i64>, usize, usize), IntSynthError> {
        let (k, s, p, op) = (l.k, l.stride, l.pad, l.out_pad);
        let hout_i = (hin as i64 - 1) * s as i64 + k as i64 + op as i64 - 2 * p as i64;
        let wout_i = (win as i64 - 1) * s as i64 + k as i64 + op as i64 - 2 * p as i64;
        if hout_i <= 0 || wout_i <= 0 {
            return Err(IntSynthError::InvalidPayload(format!(
                "int_synth layer produces non-positive output {hout_i}x{wout_i}")));
        }
        let (hout, wout) = (hout_i as usize, wout_i as usize);
        if l.weight.len() != cin * l.out_ch * k * k {
            return Err(IntSynthError::InvalidPayload(format!(
                "int_synth weight len {} != {}*{}*{}*{}", l.weight.len(), cin, l.out_ch, k, k)));
        }
        if l.bias.len() != l.out_ch {
            return Err(IntSynthError::InvalidPayload(format!(
                "int_synth bias len {} != out_ch {}", l.bias.len(), l.out_ch)));
        }

        let mut out = vec![0i64; l.out_ch * hout * wout];
        for ic in 0..cin {
            for iy in 0..hin {
                for ix in 0..win {
                    let v = input[(ic * hin + iy) * win + ix];
                    if v == 0 {
                        continue;
                    }
                    for oc in 0..l.out_ch {
                        let wbase = (ic * l.out_ch + oc) * k * k;
                        let obase = oc * hout * wout;
                        for ky in 0..k {
                            let ty = iy * s + ky;
                            if ty < p {
                                continue;
                            }
                            let oy = ty - p;
                            if oy >= hout {
                                continue;
                            }
                            for kx in 0..k {
                                let tx = ix * s + kx;
                                if tx < p {
                                    continue;
                                }
                                let ox = tx - p;
                                if ox >= wout {
                                    continue;
                                }
                                out[obase + oy * wout + ox] +=
                                    l.weight[wbase + ky * k + kx] as i64 * v;
                            }
                        }
                    }
                }
            }
        }

        // Bias, requantize (arithmetic shift — deterministic), integer ReLU.
        for oc in 0..l.out_ch {
            let b = l.bias[oc] as i64;
            let base = oc * hout * wout;
            for v in &mut out[base..base + hout * wout] {
                let mut x = (*v + b) >> l.shift;
                if l.relu && x < 0 {
                    x = 0;
                }
                *v = x;
            }
        }
        Ok((out, hout, wout))
    }
}

/// Decode an int_synth payload against the `int_synth.json` sibling of
/// `model_path`. Byte-identical on every machine — the `replicate` contract.
///
/// Filesystem-only convenience: reads the spec and delegates to
/// [`decode_from_spec`], which is the entry point with no `std::fs` (the one a
/// WASM / no-filesystem sink calls with the spec bytes supplied by the host).
pub fn decode(payload: &[u8], model_path: &Path) -> Result<IntImage, IntSynthError> {
    let dir = model_path.parent()
        .ok_or_else(|| IntSynthError::InvalidPayload("model path has no parent".into()))?;
    let spec_path = dir.join("int_synth.json");
    let raw = std::fs::read(&spec_path)
        .map_err(|e| IntSynthError::InvalidPayload(format!("{}: {e}", spec_path.display())))?;
    decode_from_spec(payload, &raw)
}

/// [`decode_from_spec`] under the name `int_payload` and the browser crate use for
/// the byte-input entry of every integer decoder: the payload and the parameter
/// set's one file, `int_synth.json`. A pin over the set is verified before this is
/// called (`pin::verify_bytes`, `int_payload::resolve`), so it takes no pin.
pub fn decode_bytes(payload: &[u8], spec_json: &[u8]) -> Result<IntImage, IntSynthError> {
    decode_from_spec(payload, spec_json)
}

/// Decode an int_synth payload against the model spec supplied as bytes (the
/// `int_synth.json` contents). No filesystem, no float, no ONNX — the
/// integer-exact decode a browser / WASM sink runs directly, host-supplied
/// weights in, RGB out. Byte-identical on every machine.
///
/// It decodes through the integer operator set, in three steps a sink may also take
/// one by one: [`check_payload`], [`load_spec`] and [`decode_loaded`]. The payload is
/// checked first (`extensions/integer-payloads.md` §13): its header, a reserved byte of
/// 0, `C ≥ 1`, `C·S·S ≤ 2^26` and its exact length. Then the set is lowered
/// (integer-operators §13.1, whose checks include at least one layer, fields below
/// `2^32`, `out_ch`, `k` and `stride` at least 1 and a shift of at most 127) and loaded
/// under `acc64`: load steps 1–8, the last being the bounds certificate
/// ([`crate::int_ops`]). Then the graph runs on the latents, through the operator set's
/// run-time plan with its caps, and its output must be `[3, H, W]`. Every refusal is
/// [`IntSynthError::InvalidPayload`] (`payload_malformed`). A sink that keeps loaded
/// sets between decodes answers the middle step from its store (`int_payload`'s decode
/// of `wai.neural.int_synth`).
pub fn decode_from_spec(payload: &[u8], spec_json: &[u8]) -> Result<IntImage, IntSynthError> {
    check_payload(payload)?;
    let g = load_spec(spec_json)?;
    decode_loaded(payload, &g)
}

/// The payload's checks, the first step of [`decode_from_spec`].
pub fn check_payload(payload: &[u8]) -> Result<(), IntSynthError> {
    header(payload).map(|_| ())
}

fn refused(what: &str, r: wai_ops::Refusal) -> IntSynthError {
    IntSynthError::InvalidPayload(format!("int_synth {what}: {} at {:?}", r.code.as_str(), r.at))
}

/// The set `int_synth.json` lowered by `int_synth/1` and loaded under `acc64`, the
/// second step of [`decode_from_spec`].
pub fn load_spec(spec_json: &[u8]) -> Result<wai_ops::Loaded, IntSynthError> {
    let set = wai_ops::lower::int_synth(spec_json).map_err(|r| refused("int_synth.json does not lower", r))?;
    crate::int_ops::load(set.get("synth").expect("int_synth/1 lowers to the graph synth")).map_err(|r| refused("set does not load", r))
}

/// The decode with its set loaded ([`load_spec`]), the last step of
/// [`decode_from_spec`]: the payload's checks again, then the graph on the latents.
pub fn decode_loaded(payload: &[u8], g: &wai_ops::Loaded) -> Result<IntImage, IntSynthError> {
    let (h, w, c, s) = header(payload)?;
    let latents: Vec<i64> = payload[HEADER..].iter().map(|&b| b as i8 as i64).collect();
    let out = crate::int_ops::run1(g, wai_ops::Tensor::new(c, s, s, latents), &[]).map_err(|r| refused("decode", r))?;
    if (out.c, out.h as usize, out.w as usize) != (3, h, w) {
        return Err(IntSynthError::InvalidPayload(format!(
            "int_synth output [{},{},{}] != [3,{h},{w}]", out.c, out.h, out.w)));
    }
    // `clamp(0, 255)` ends the graph, so every value is a byte.
    let hw = h * w;
    let mut rgb = vec![0u8; hw * 3];
    for (k, &v) in out.data.iter().enumerate() {
        rgb[(k % hw) * 3 + k / hw] = v as u8;
    }
    Ok(IntImage { rgb, width: w as u32, height: h as u32 })
}

/// The payload's header (`extensions/integer-payloads.md` §13): `(H, W, C, S)`, after
/// its checks. The side is `2^S_log2`; `C·S·S` is formed without overflow and is at
/// most `2^26`, and the payload is exactly `12 + C·S·S` bytes.
fn header(payload: &[u8]) -> Result<(usize, usize, u32, u32), IntSynthError> {
    if payload.len() < HEADER {
        return Err(IntSynthError::InvalidPayload(format!(
            "int_synth payload {} B < {HEADER} B header", payload.len())));
    }
    let h = u32::from_le_bytes(payload[0..4].try_into().unwrap()) as usize;
    let w = u32::from_le_bytes(payload[4..8].try_into().unwrap()) as usize;
    let c = u16::from_le_bytes(payload[8..10].try_into().unwrap()) as usize;
    if payload[11] != 0 {
        return Err(IntSynthError::InvalidPayload("int_synth payload's reserved byte is not 0".into()));
    }
    if c == 0 {
        return Err(IntSynthError::InvalidPayload("int_synth payload has no latent channels".into()));
    }
    const CAP: usize = 1 << 26;
    let n = 1usize.checked_shl(payload[10] as u32)
        .and_then(|s| c.checked_mul(s)?.checked_mul(s).map(|n| (s, n)))
        .filter(|&(_, n)| n <= CAP);
    let Some((s, _)) = n.filter(|&(_, n)| HEADER.checked_add(n) == Some(payload.len())) else {
        return Err(IntSynthError::InvalidPayload(format!(
            "int_synth payload {} B is not {HEADER} + C·S·S with C = {c}, S = 2^{}, C·S·S ≤ 2^26", payload.len(), payload[10])));
    };
    Ok((h, w, c as u32, s as u32))
}

/// The direct decode, as the decoder ran before it decoded through the integer operator
/// set: the payload and set read here, each layer's transposed convolution, bias, shift
/// and ReLU computed by the scalar reference kernel. Test-only (`legacy_int`): the
/// differential tests decode both ways, and the two differ only on the planned breaks
/// of `extensions/integer-operators.md` §15.3.
#[cfg(any(test, feature = "legacy_int"))]
pub mod legacy {
    use super::*;

    /// The direct decode of [`super::decode_from_spec`]'s arguments.
    pub fn decode_from_spec(payload: &[u8], spec_json: &[u8]) -> Result<IntImage, IntSynthError> {
        if payload.len() < HEADER {
            return Err(IntSynthError::InvalidPayload(format!(
                "int_synth payload {} B < {HEADER} B header", payload.len())));
        }
        let h = u32::from_le_bytes(payload[0..4].try_into().unwrap()) as usize;
        let w = u32::from_le_bytes(payload[4..8].try_into().unwrap()) as usize;
        let c = u16::from_le_bytes(payload[8..10].try_into().unwrap()) as usize;
        // The side is 2^payload[10]. A shift of the word's width or more panics with
        // overflow checks on and masks without them, so the same bytes would decode two
        // ways; and C·S·S must be formed without overflow. Neither can match a payload's
        // length when it overflows, so both are refused here as a length mismatch.
        let n = 1usize.checked_shl(payload[10] as u32)
            .and_then(|s| c.checked_mul(s)?.checked_mul(s).map(|n| (s, n)));
        let Some((s, _)) = n.filter(|&(_, n)| HEADER.checked_add(n) == Some(payload.len())) else {
            return Err(IntSynthError::InvalidPayload(format!(
                "int_synth payload {} B is not {HEADER} + C·S·S with C = {c}, S = 2^{}", payload.len(), payload[10])));
        };

        let spec: SynthSpec = serde_json::from_slice(spec_json)
            .map_err(|e| IntSynthError::InvalidPayload(format!("int_synth.json parse: {e}")))?;

        // int8 latents → i64.
        let mut cur: Vec<i64> = payload[HEADER..].iter().map(|&b| b as i8 as i64).collect();
        let (mut cin, mut ch, mut cw) = (c, s, s);
        for (i, l) in spec.layers.iter().enumerate() {
            if l.in_ch != cin {
                return Err(IntSynthError::InvalidPayload(format!(
                    "int_synth layer {i} in_ch {} != current {cin}", l.in_ch)));
            }
            let (out, ho, wo) = reference::conv_transpose2d(&cur, cin, ch, cw, l)?;
            cur = out;
            cin = l.out_ch;
            ch = ho;
            cw = wo;
        }
        if cin != 3 || ch != h || cw != w {
            return Err(IntSynthError::InvalidPayload(format!(
                "int_synth output [{cin},{ch},{cw}] != [3,{h},{w}]")));
        }

        let mut rgb = vec![0u8; h * w * 3];
        let hw = h * w;
        for y in 0..h {
            for x in 0..w {
                let pi = y * w + x;
                for ch_i in 0..3 {
                    rgb[pi * 3 + ch_i] = cur[ch_i * hw + pi].clamp(0, 255) as u8;
                }
            }
        }
        Ok(IntImage { rgb, width: w as u32, height: h as u32 })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    // The direct decode's scalar kernel, against hand computation.
    use super::reference::conv_transpose2d;

    fn layer(out_ch: usize, in_ch: usize, k: usize, stride: usize, weight: Vec<i32>) -> LayerSpec {
        LayerSpec {
            in_ch, out_ch, k, stride, pad: 0, out_pad: 0,
            weight, bias: vec![0; out_ch], shift: 0, relu: false,
        }
    }

    #[test]
    fn transposed_conv_matches_hand_computation() {
        // 1 in-channel, 1 out-channel, 2×2 kernel [[1,2],[3,4]], stride 2.
        // input [[5]] (1×1) scatters to a 2×2 output = kernel * 5.
        let l = layer(1, 1, 2, 2, vec![1, 2, 3, 4]);
        let (out, ho, wo) = conv_transpose2d(&[5], 1, 1, 1, &l).unwrap();
        assert_eq!((ho, wo), (2, 2));
        assert_eq!(out, vec![5, 10, 15, 20]);
    }

    #[test]
    fn requantize_shift_and_relu_are_integer_exact() {
        // shift=1 halves (arithmetic), relu clamps negatives.
        let mut l = layer(1, 1, 1, 1, vec![1]);
        l.shift = 1;
        l.relu = true;
        l.bias = vec![0];
        // input -3 → weight*in = -3 → >>1 = -2 → relu → 0.
        let (out, _, _) = conv_transpose2d(&[-3], 1, 1, 1, &l).unwrap();
        assert_eq!(out, vec![0]);
        // input 7 → 7 → >>1 = 3.
        let (out, _, _) = conv_transpose2d(&[7], 1, 1, 1, &l).unwrap();
        assert_eq!(out, vec![3]);
    }

    #[test]
    fn decode_is_deterministic_by_construction() {
        // Two runs of the pure integer op on the same input are byte-identical
        // (the property that makes this a sound `replicate` member).
        let l = layer(2, 1, 3, 2, (0..18).collect());
        let input = [3, -4, 5];
        let a = conv_transpose2d(&input, 1, 1, 3, &l).unwrap();
        let b = conv_transpose2d(&input, 1, 1, 3, &l).unwrap();
        assert_eq!(a, b);
    }

    /// FULL decode path, no ONNX: a constructed 1-layer integer model + a tiny
    /// latent → exact RGB, and byte-identical across runs. This validates the
    /// integer-exact `replicate` member end to end (header parse, latent unpack,
    /// conv, requantize, clamp, RGB layout).
    #[test]
    fn end_to_end_decode_is_exact_and_deterministic() {
        let dir = std::env::temp_dir().join("wai_int_synth_e2e");
        std::fs::create_dir_all(&dir).unwrap();
        // 1 layer, in_ch=1 → out_ch=3 (RGB), 2×2 kernel, stride 2. Each output
        // channel has a constant kernel (1, 2, 3), so a 1×1 latent of 10 scatters
        // to a 2×2 image where every pixel is (10, 20, 30).
        let weight: Vec<i32> = (1..=3).flat_map(|oc| std::iter::repeat(oc).take(4)).collect();
        let spec = serde_json::json!({"layers": [{
            "in_ch": 1, "out_ch": 3, "k": 2, "stride": 2, "pad": 0, "out_pad": 0,
            "weight": weight, "bias": [0, 0, 0], "shift": 0, "relu": false,
        }]});
        std::fs::write(dir.join("int_synth.json"), serde_json::to_vec(&spec).unwrap()).unwrap();
        let model_path = dir.join("model");

        let mut payload = Vec::new();
        payload.extend_from_slice(&2u32.to_le_bytes()); // H
        payload.extend_from_slice(&2u32.to_le_bytes()); // W
        payload.extend_from_slice(&1u16.to_le_bytes()); // C
        payload.push(0); // S_log2 = 0  → S = 1
        payload.push(0); // reserved
        payload.push(10u8); // one int8 latent = 10

        let img = decode(&payload, &model_path).expect("decode");
        assert_eq!((img.width, img.height), (2, 2));
        let expected: Vec<u8> = (0..4).flat_map(|_| [10u8, 20, 30]).collect();
        assert_eq!(img.rgb, expected, "integer synthesis must hit the exact bytes");
        // Byte-identical across runs — the replicate property.
        assert_eq!(decode(&payload, &model_path).unwrap().rgb, img.rgb);
        let _ = std::fs::remove_dir_all(&dir);
    }
}