inferencelayer 0.2.3

//! 2D convolution foundation for the Docling ports (RT-DETRv2 layout + TableFormer) — CPU.
//!
//! One general path: **im2col → [`cpu_gemm::gemm_packed`]**, NHWC activations, single image
//! (the document pipeline is strictly per-page). The direct-loop `Conv2d` in `parakeet.rs` is
//! correct but scalar; ResNet-50-class backbones at 640×640 are ~tens of GFLOPs, which only the
//! packed GEMM microkernels (NEON/AVX2, deterministic fixed k-order) can carry.
//!
//! Layout contract:
//! * activations are NHWC `[h, w, c]` — im2col patch extraction is then a contiguous per-pixel
//!   memcpy per kernel tap, and the 1×1 fast path is a plain GEMM on the buffer as-is;
//! * torch weights `[oc, ic, kh, kw]` are permuted ONCE at load to `[oc, kh·kw·ic]` row-major
//!   (matching the im2col column order) and packed via [`PackedWeight`].
//!
//! Scratch: im2col materializes the full `[oh·ow, kh·kw·ic]` matrix (~60 MB worst case at the
//! stem). That is deliberate for the correctness milestone — `gemm_packed` parallelizes
//! internally and stays bit-identical run-to-run; band-tiling the scratch is an A4 perf lever,
//! not a correctness one.
//!
//! BatchNorm never exists at inference: [`fold_bn`] composes it into the conv weights at load
//! (same recipe as `diarize.rs`), so the parity bar vs torch's separate conv+BN is the usual
//! f32-reassociation tolerance, gated in `tests/conv2d_parity.rs` against oracle tensors
//! exported by `tests/fixtures/export_docling_layout.py` (the `unit.*` family).

use rayon::prelude::*;

use crate::cpu_gemm::{PackedWeight, gemm_packed};

/// A folded (bias-carrying, BN-free) 2D convolution, weights prepacked for the GEMM.
pub struct Conv2d {
    w: PackedWeight,
    bias: Vec<f32>,
    ic: usize,
    kh: usize,
    kw: usize,
    stride: usize,
    pad: usize,
}

impl Conv2d {
    /// Build from a torch-layout weight `[oc, ic, kh, kw]` (+ optional bias), permuting to the
    /// im2col order `[oc][kh][kw][ic]` and packing once.
    pub fn from_torch(
        w: &[f32],
        bias: Option<&[f32]>,
        oc: usize,
        ic: usize,
        kh: usize,
        kw: usize,
        stride: usize,
        pad: usize,
    ) -> Self {
        assert_eq!(w.len(), oc * ic * kh * kw, "weight is not [oc, ic, kh, kw]");
        let k = kh * kw * ic;
        let mut perm = vec![0f32; oc * k];
        for o in 0..oc {
            for i in 0..ic {
                for y in 0..kh {
                    for x in 0..kw {
                        perm[o * k + (y * kw + x) * ic + i] = w[((o * ic + i) * kh + y) * kw + x];
                    }
                }
            }
        }
        let bias = bias.map(<[f32]>::to_vec).unwrap_or_else(|| vec![0.0; oc]);
        assert_eq!(bias.len(), oc);
        Self {
            w: PackedWeight::new(&perm, oc, k),
            bias,
            ic,
            kh,
            kw,
            stride,
            pad,
        }
    }

    /// Output channels.
    pub fn oc(&self) -> usize {
        self.w.n()
    }

    /// Everything a GPU twin needs to upload this conv: the packed weight (unpack for a flat
    /// `[oc, kh·kw·ic]`), bias, and the geometry. Mirrors the whisper_gpu-from-cpu pattern.
    pub(crate) fn gpu_parts(&self) -> (&PackedWeight, &[f32], [usize; 5]) {
        (
            &self.w,
            &self.bias,
            [self.ic, self.kh, self.kw, self.stride, self.pad],
        )
    }

    /// NHWC forward: `x` is `[h, w, ic]`; returns (`[oh, ow, oc]`, oh, ow).
    pub fn forward(&self, x: &[f32], h: usize, w: usize) -> (Vec<f32>, usize, usize) {
        assert_eq!(x.len(), h * w * self.ic, "input is not [h, w, ic]");
        let (kh, kw, s, p, ic) = (self.kh, self.kw, self.stride, self.pad, self.ic);
        let oh = (h + 2 * p - kh) / s + 1;
        let ow = (w + 2 * p - kw) / s + 1;
        let oc = self.oc();
        let m = oh * ow;
        let mut out = vec![0f32; m * oc];

        if kh == 1 && kw == 1 && p == 0 {
            if s == 1 {
                // 1×1 stride-1: the NHWC buffer IS the [h·w, ic] GEMM input. This is the
                // majority of ResNet bottleneck convs and every CCFM lateral conv.
                gemm_packed(&mut out, x, &self.w, m, Some(&self.bias));
            } else {
                // 1×1 strided (classic ResNet downsample): gather the strided pixel rows.
                let mut a = vec![0f32; m * ic];
                for oy in 0..oh {
                    for ox in 0..ow {
                        let src = ((oy * s) * w + ox * s) * ic;
                        let dst = (oy * ow + ox) * ic;
                        a[dst..dst + ic].copy_from_slice(&x[src..src + ic]);
                    }
                }
                gemm_packed(&mut out, &a, &self.w, m, Some(&self.bias));
            }
            return (out, oh, ow);
        }

        // General: im2col, PARALLEL over output rows (the packing was ~half the conv wall time
        // single-threaded — the profile, not intuition, picked this). Two inner shapes:
        // * stride 1: per (row, dy) the in-bounds taps of ALL kernel columns read one contiguous
        //   src run, so the dx loop advances a src cursor instead of recomputing indices;
        // * strided: per-pixel tap copies (the stem's s2 convs — small share of the budget).
        // Zero-padding falls out of the pre-zeroed per-row chunk.
        let k = kh * kw * ic;
        let mut a = vec![0f32; m * k];
        a.par_chunks_mut(ow * k).enumerate().for_each(|(oy, arow)| {
            let iy0 = (oy * s) as isize - p as isize;
            for dy in 0..kh {
                let iy = iy0 + dy as isize;
                if iy < 0 || iy >= h as isize {
                    continue;
                }
                let srow = &x[iy as usize * w * ic..(iy as usize + 1) * w * ic];
                for ox in 0..ow {
                    let row = ox * k + dy * kw * ic;
                    let ix0 = (ox * s) as isize - p as isize;
                    for dx in 0..kw {
                        let ix = ix0 + dx as isize;
                        if ix < 0 || ix >= w as isize {
                            continue;
                        }
                        let src = ix as usize * ic;
                        arow[row + dx * ic..row + (dx + 1) * ic]
                            .copy_from_slice(&srow[src..src + ic]);
                    }
                }
            }
        });
        gemm_packed(&mut out, &a, &self.w, m, Some(&self.bias));
        (out, oh, ow)
    }
}

/// Fold an eval-mode BatchNorm2d into torch-layout conv weights, BEFORE [`Conv2d::from_torch`].
///
/// `w'[o,..] = w[o,..] · γ[o]/√(σ²[o]+ε)`, `b'[o] = (b[o] − μ[o]) · γ[o]/√(σ²[o]+ε) + β[o]`.
/// Composing the affine into the weights only reassociates f32 math — the parity gate's usual
/// tolerance covers it (proven against the torch BN oracle in `conv2d_parity.rs`).
pub fn fold_bn(
    w: &[f32],
    bias: Option<&[f32]>,
    oc: usize,
    gamma: &[f32],
    beta: &[f32],
    mean: &[f32],
    var: &[f32],
    eps: f32,
) -> (Vec<f32>, Vec<f32>) {
    assert_eq!(w.len() % oc, 0);
    let per = w.len() / oc;
    let mut wf = vec![0f32; w.len()];
    let mut bf = vec![0f32; oc];
    for o in 0..oc {
        let scale = gamma[o] / (var[o] + eps).sqrt();
        for i in 0..per {
            wf[o * per + i] = w[o * per + i] * scale;
        }
        let b0 = bias.map_or(0.0, |b| b[o]);
        bf[o] = (b0 - mean[o]) * scale + beta[o];
    }
    (wf, bf)
}

/// NHWC max-pool, torch semantics: padding taps are −∞ (never win), floor output size.
pub fn max_pool2d(
    x: &[f32],
    h: usize,
    w: usize,
    c: usize,
    k: usize,
    s: usize,
    p: usize,
) -> (Vec<f32>, usize, usize) {
    let oh = (h + 2 * p - k) / s + 1;
    let ow = (w + 2 * p - k) / s + 1;
    let mut out = vec![f32::NEG_INFINITY; oh * ow * c];
    for oy in 0..oh {
        for ox in 0..ow {
            let dst = (oy * ow + ox) * c;
            let iy0 = (oy * s) as isize - p as isize;
            let ix0 = (ox * s) as isize - p as isize;
            for dy in 0..k {
                let iy = iy0 + dy as isize;
                if iy < 0 || iy >= h as isize {
                    continue;
                }
                for dx in 0..k {
                    let ix = ix0 + dx as isize;
                    if ix < 0 || ix >= w as isize {
                        continue;
                    }
                    let src = (iy as usize * w + ix as usize) * c;
                    for ch in 0..c {
                        let v = x[src + ch];
                        if v > out[dst + ch] {
                            out[dst + ch] = v;
                        }
                    }
                }
            }
        }
    }
    (out, oh, ow)
}

/// NHWC adaptive average pool, torch's start/end index formula
/// (`start = ⌊i·H/OH⌋`, `end = ⌈(i+1)·H/OH⌉`) — exact for both even and uneven ratios.
pub fn adaptive_avg_pool2d(
    x: &[f32],
    h: usize,
    w: usize,
    c: usize,
    oh: usize,
    ow: usize,
) -> Vec<f32> {
    let mut out = vec![0f32; oh * ow * c];
    for oy in 0..oh {
        let y0 = oy * h / oh;
        let y1 = ((oy + 1) * h).div_ceil(oh);
        for ox in 0..ow {
            let x0 = ox * w / ow;
            let x1 = ((ox + 1) * w).div_ceil(ow);
            let inv = 1.0 / ((y1 - y0) * (x1 - x0)) as f32;
            let dst = (oy * ow + ox) * c;
            for iy in y0..y1 {
                for ix in x0..x1 {
                    let src = (iy * w + ix) * c;
                    for ch in 0..c {
                        out[dst + ch] += x[src + ch];
                    }
                }
            }
            for ch in 0..c {
                out[dst + ch] *= inv;
            }
        }
    }
    out
}

/// Bilinear grid-sample of ONE point per output row, torch semantics `align_corners=False`,
/// `padding_mode="zeros"`: unnormalize `x = ((gx+1)·W − 1)/2`, out-of-range taps contribute 0.
/// `grid` is `[(gx, gy); n]` in [−1, 1]-ish; returns `[n, c]`.
///
/// This is the core read of RT-DETRv2's multi-scale deformable attention (and is unit-gated on
/// out-of-bounds points, which zeros-padding makes load-bearing there).
pub fn grid_sample_bilinear(
    x: &[f32],
    h: usize,
    w: usize,
    c: usize,
    grid: &[(f32, f32)],
) -> Vec<f32> {
    let mut out = vec![0f32; grid.len() * c];
    for (n, &(gx, gy)) in grid.iter().enumerate() {
        let fx = ((gx + 1.0) * w as f32 - 1.0) * 0.5;
        let fy = ((gy + 1.0) * h as f32 - 1.0) * 0.5;
        let x0 = fx.floor();
        let y0 = fy.floor();
        let (tx, ty) = (fx - x0, fy - y0);
        let dst = n * c;
        // (tap, weight) over the 2×2 neighborhood; out-of-bounds taps are skipped (= zero).
        for (dy, wy) in [(0i64, 1.0 - ty), (1, ty)] {
            let iy = y0 as i64 + dy;
            if iy < 0 || iy >= h as i64 || wy == 0.0 {
                continue;
            }
            for (dx, wx) in [(0i64, 1.0 - tx), (1, tx)] {
                let ix = x0 as i64 + dx;
                if ix < 0 || ix >= w as i64 || wx == 0.0 {
                    continue;
                }
                let src = (iy as usize * w + ix as usize) * c;
                let wgt = wy * wx;
                for ch in 0..c {
                    out[dst + ch] += x[src + ch] * wgt;
                }
            }
        }
    }
    out
}