oar-ocr-core 0.10.2

//! SIMD-accelerated kernels for the pipeline's two hot CPU loops: per-pixel
//! image normalization and per-timestep CTC argmax decode. (Conv/matmul work
//! lives in ONNX Runtime.)
//!
//! Each kernel ships as a scalar reference (`*_scalar`, always compiled) and a
//! SIMD path (`*_simd`, behind the `simd` feature) built on [`wide`] and
//! runtime-dispatched via [`multiversion`], so one binary selects AVX2/SSE/NEON
//! at runtime without `target-cpu=native`. Entry points pick the SIMD path when
//! the feature is on, else the scalar one.
//!
//! Normalization keeps the scalar `value * alpha + beta` form (plain
//! multiply-add, *not* FMA). Argmax preserves the last winning index and the
//! existing confidence reduction: SIMD retains the target's max behavior for
//! NaNs and signed zeros, whereas the scalar fallback uses `Iterator::max_by`.
//!
//! [`wide`]: https://docs.rs/wide
//! [`multiversion`]: https://docs.rs/multiversion

/// Writes a normalized image in CHW (channel-major) layout into `out`.
///
/// For each output channel `c`, plane `c` is `out[c * plane .. (c + 1) * plane]`
/// where `plane = width * height`, and
/// `out[c * plane + p] = rgb[p * 3 + src_channels[c]] as f32 * alpha[c] + beta[c]`.
///
/// `rgb` is the tightly packed interleaved RGB buffer (`image::RgbImage::as_raw`),
/// length `width * height * 3`. `out` must have length `3 * plane`.
#[inline]
pub fn normalize_chw_into(
    rgb: &[u8],
    width: usize,
    height: usize,
    src_channels: [usize; 3],
    alpha: &[f32; 3],
    beta: &[f32; 3],
    out: &mut [f32],
) {
    #[cfg(feature = "simd")]
    {
        normalize_chw_simd(rgb, width, height, src_channels, alpha, beta, out);
    }
    #[cfg(not(feature = "simd"))]
    {
        normalize_chw_scalar(rgb, width, height, src_channels, alpha, beta, out);
    }
}

/// Writes a normalized image in HWC (interleaved) layout into `out`.
///
/// `out[p * 3 + c] = rgb[p * 3 + src_channels[c]] as f32 * alpha[c] + beta[c]`.
/// `out` must have length `width * height * 3`.
#[inline]
pub fn normalize_hwc_into(
    rgb: &[u8],
    width: usize,
    height: usize,
    src_channels: [usize; 3],
    alpha: &[f32; 3],
    beta: &[f32; 3],
    out: &mut [f32],
) {
    // HWC is interleaved (and BGR swaps per triple), so it doesn't vectorize
    // cleanly; the scalar raw-slice path is used regardless of the `simd` feature.
    normalize_hwc_scalar(rgb, width, height, src_channels, alpha, beta, out);
}

/// Returns the `(index, value)` of the maximum element of `row`.
///
/// On ties the **last** maximal index is returned, matching
/// [`Iterator::max_by`] semantics (which the previous decode loop relied on).
/// Returns `None` for an empty row.
#[inline]
pub fn argmax(row: &[f32]) -> Option<(usize, f32)> {
    #[cfg(feature = "simd")]
    {
        argmax_simd(row)
    }
    #[cfg(not(feature = "simd"))]
    {
        argmax_scalar(row)
    }
}

// Scalar reference implementations: the active fallback when `simd` is off, and
// the oracle for the parity tests.
#[cfg(any(not(feature = "simd"), test))]
#[inline]
fn normalize_chw_scalar(
    rgb: &[u8],
    width: usize,
    height: usize,
    src_channels: [usize; 3],
    alpha: &[f32; 3],
    beta: &[f32; 3],
    out: &mut [f32],
) {
    let plane = width * height;
    for c in 0..3 {
        let (a, b, sc) = (alpha[c], beta[c], src_channels[c]);
        let dst = &mut out[c * plane..c * plane + plane];
        for (p, d) in dst.iter_mut().enumerate() {
            *d = rgb[p * 3 + sc] as f32 * a + b;
        }
    }
}

#[inline]
fn normalize_hwc_scalar(
    rgb: &[u8],
    width: usize,
    height: usize,
    src_channels: [usize; 3],
    alpha: &[f32; 3],
    beta: &[f32; 3],
    out: &mut [f32],
) {
    let plane = width * height;
    for p in 0..plane {
        let base = p * 3;
        for c in 0..3 {
            out[base + c] = rgb[base + src_channels[c]] as f32 * alpha[c] + beta[c];
        }
    }
}

// Fallback when `simd` is off; also the oracle in `argmax_matches_scalar_*`.
#[cfg(any(not(feature = "simd"), test))]
#[inline]
fn argmax_scalar(row: &[f32]) -> Option<(usize, f32)> {
    row.iter()
        .enumerate()
        .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
        .map(|(idx, &v)| (idx, v))
}

// SIMD implementations (feature = "simd").

#[cfg(feature = "simd")]
fn normalize_chw_simd(
    rgb: &[u8],
    width: usize,
    height: usize,
    src_channels: [usize; 3],
    alpha: &[f32; 3],
    beta: &[f32; 3],
    out: &mut [f32],
) {
    let plane = width * height;
    for c in 0..3 {
        let dst = &mut out[c * plane..c * plane + plane];
        normalize_plane_simd(rgb, src_channels[c], alpha[c], beta[c], dst);
    }
}

/// Fills one contiguous CHW plane: `dst[p] = rgb[p * 3 + sc] as f32 * a + b`.
///
/// The strided u8 gather is scalar (3-channel deinterleave has no portable SIMD
/// gather), but the `* a + b` arithmetic and the store run 8 lanes wide. Uses a
/// plain multiply-add (not `mul_add`) to stay bit-identical to the scalar path.
#[cfg(feature = "simd")]
#[multiversion::multiversion(targets("x86_64+avx2+fma", "x86_64+sse4.2", "aarch64+neon"))]
fn normalize_plane_simd(rgb: &[u8], sc: usize, a: f32, b: f32, dst: &mut [f32]) {
    use wide::f32x8;

    let va = f32x8::splat(a);
    let vb = f32x8::splat(b);
    let (chunks, remainder) = dst.as_chunks_mut::<8>();
    let mut p = 0usize;
    for chunk in chunks {
        let gathered = [
            rgb[p * 3 + sc] as f32,
            rgb[(p + 1) * 3 + sc] as f32,
            rgb[(p + 2) * 3 + sc] as f32,
            rgb[(p + 3) * 3 + sc] as f32,
            rgb[(p + 4) * 3 + sc] as f32,
            rgb[(p + 5) * 3 + sc] as f32,
            rgb[(p + 6) * 3 + sc] as f32,
            rgb[(p + 7) * 3 + sc] as f32,
        ];
        let v = f32x8::new(gathered) * va + vb;
        *chunk = v.to_array();
        p += 8;
    }
    for d in remainder {
        *d = rgb[p * 3 + sc] as f32 * a + b;
        p += 1;
    }
}

#[cfg(feature = "simd")]
#[multiversion::multiversion(targets("x86_64+avx2+fma", "x86_64+sse4.2", "aarch64+neon"))]
fn argmax_simd(row: &[f32]) -> Option<(usize, f32)> {
    use wide::{bytemuck::cast, f32x8, u32x8};

    if row.is_empty() {
        return None;
    }
    // CTC vocabularies fit in u32 lanes. Keep the old path for larger slices
    // rather than truncating indices or changing exceptional-float behavior.
    if row.len() > u32::MAX as usize {
        return argmax_two_pass_simd(row);
    }

    let mut values = f32x8::splat(f32::NEG_INFINITY);
    let mut nan_mask = f32x8::splat(0.0);
    let mut indices = u32x8::splat(0);
    let mut positions = u32x8::new([0, 1, 2, 3, 4, 5, 6, 7]);
    let (chunks, remainder) = row.as_chunks::<8>();
    for chunk in chunks {
        let next = f32x8::new(*chunk);
        nan_mask |= next.simd_ne(next);
        let replace: u32x8 = cast(next.simd_ge(values));
        // Comparison masks have all-zero or all-one lanes, as select requires.
        indices = replace.select(positions, indices);
        // Retain max's NaN/zero handling and the original confidence bits.
        values = values.max(next);
        positions += u32x8::splat(8);
    }

    let mut best = f32::NEG_INFINITY;
    let mut index = 0usize;
    for (value, lane_index) in values.to_array().into_iter().zip(indices.to_array()) {
        let lane_index = lane_index as usize;
        if value > best || (value == best && lane_index > index) {
            index = lane_index;
        }
        best = best.max(value);
    }
    let offset = chunks.len() * 8;
    let mut has_nan = nan_mask.any();
    for (i, &value) in remainder.iter().enumerate() {
        has_nan |= value.is_nan();
        if value >= best {
            index = offset + i;
        }
        best = best.max(value);
    }
    // NaN max/compare behavior and signed-zero selection differ across targets.
    // Preserve both the index and confidence via the original two-pass path.
    if has_nan || best == 0.0 {
        return argmax_two_pass_simd(row);
    }
    Some((index, best))
}

#[cfg(feature = "simd")]
fn argmax_two_pass_simd(row: &[f32]) -> Option<(usize, f32)> {
    if row.is_empty() {
        return None;
    }
    let max_val = reduce_max_simd(row);
    let mut idx = 0usize;
    for (i, &v) in row.iter().enumerate() {
        // `v >= max_val` holds only at the maximum (NaN-free input).
        if v >= max_val {
            idx = i;
        }
    }
    Some((idx, max_val))
}

#[cfg(feature = "simd")]
#[multiversion::multiversion(targets("x86_64+avx2+fma", "x86_64+sse4.2", "aarch64+neon"))]
fn reduce_max_simd(row: &[f32]) -> f32 {
    use wide::f32x8;

    let mut acc = f32x8::splat(f32::NEG_INFINITY);
    let (chunks, remainder) = row.as_chunks::<8>();
    for chunk in chunks {
        let v = f32x8::new(*chunk);
        acc = acc.max(v);
    }
    let mut best = f32::NEG_INFINITY;
    for lane in acc.to_array() {
        best = best.max(lane);
    }
    for &v in remainder {
        best = best.max(v);
    }
    best
}

// CRNN recognition normalize: `(v / 255.0 - 0.5) / 0.5` in BGR order.

/// BGR source-channel order for the CRNN tensor (destination channel 0 uses source B, etc.).
const CRNN_SRC: [usize; 3] = [2, 1, 0];

/// Fills a CRNN CHW input tensor from a resized RGB crop.
///
/// Each channel is normalized as `(v / 255.0 - 0.5) / 0.5` — the PaddleOCR CRNN
/// convention — and emitted in BGR channel order. `dst` is the full
/// `3 * img_h * tensor_width` tensor; columns `[resized_w, tensor_width)` of
/// every row keep their prior value (the caller's zero padding).
///
/// `rgb` is the tightly packed `resized_w * img_h * 3` interleaved buffer
/// (`image::RgbImage::as_raw`). The kernel is bit-identical to the scalar
/// reference: it keeps the exact `(v / 255 - 0.5) / 0.5` op order (a true
/// division, not a reciprocal multiply), so recognition input is unchanged.
#[inline]
pub fn normalize_crnn_chw_into(
    rgb: &[u8],
    resized_w: usize,
    img_h: usize,
    tensor_width: usize,
    dst: &mut [f32],
) {
    let plane = img_h * tensor_width;
    for (c, &sc) in CRNN_SRC.iter().enumerate() {
        let cbase = c * plane;
        for y in 0..img_h {
            let src_row = &rgb[y * resized_w * 3..y * resized_w * 3 + resized_w * 3];
            let row_off = cbase + y * tensor_width;
            let dst_row = &mut dst[row_off..row_off + resized_w];
            #[cfg(feature = "simd")]
            crnn_row_simd(src_row, sc, dst_row);
            #[cfg(not(feature = "simd"))]
            crnn_row_scalar(src_row, sc, dst_row);
        }
    }
}

// Fallback when `simd` is off. The CRNN parity test uses an independent
// full-tensor oracle rather than this per-row helper, so it is not needed in
// `simd` test builds.
#[cfg(not(feature = "simd"))]
#[inline]
fn crnn_row_scalar(src_row: &[u8], sc: usize, dst_row: &mut [f32]) {
    for (x, d) in dst_row.iter_mut().enumerate() {
        *d = (src_row[x * 3 + sc] as f32 / 255.0 - 0.5) / 0.5;
    }
}

#[cfg(feature = "simd")]
#[multiversion::multiversion(targets("x86_64+avx2+fma", "x86_64+sse4.2", "aarch64+neon"))]
fn crnn_row_simd(src_row: &[u8], sc: usize, dst_row: &mut [f32]) {
    use wide::f32x8;

    let (chunks, remainder) = dst_row.as_chunks_mut::<8>();
    let mut x = 0usize;
    for chunk in chunks {
        let g = [
            src_row[x * 3 + sc] as f32,
            src_row[(x + 1) * 3 + sc] as f32,
            src_row[(x + 2) * 3 + sc] as f32,
            src_row[(x + 3) * 3 + sc] as f32,
            src_row[(x + 4) * 3 + sc] as f32,
            src_row[(x + 5) * 3 + sc] as f32,
            src_row[(x + 6) * 3 + sc] as f32,
            src_row[(x + 7) * 3 + sc] as f32,
        ];
        // Exact `(v / 255.0 - 0.5) / 0.5`, lane-wise.
        let v = (f32x8::new(g) / 255.0 - 0.5) / 0.5;
        *chunk = v.to_array();
        x += 8;
    }
    for d in remainder {
        *d = (src_row[x * 3 + sc] as f32 / 255.0 - 0.5) / 0.5;
        x += 1;
    }
}

// UVDoc rectification output: CHW float planes -> interleaved RGB u8.

/// Converts a BGR-ordered CHW float plane-triple into an interleaved RGB `u8`
/// image, computing `(v * scale).clamp(0.0, 255.0) as u8` per element.
///
/// `c0`/`c1`/`c2` are the three contiguous channel planes (each `plane` long,
/// in the model's BGR order). `out` is `plane * 3` interleaved bytes, emitted as
/// RGB (`r <- c2`, `g <- c1`, `b <- c0`).
///
/// This is intentionally a tight contiguous scalar loop, *not* a hand-vectorized
/// kernel: the output is an interleaved (strided) RGB scatter, which an explicit
/// `wide` gather/scatter measured ~1.8x *slower* than letting the autovectorizer
/// handle this `as_chunks_mut::<3>()` form. The win over the previous
/// `put_pixel` + 4-D strided `ndarray` indexing comes from the contiguous
/// plane access alone (~1.8x faster than the old loop). Kept here next to the
/// other post-processing kernels for cohesion.
#[inline]
pub fn scale_clamp_bgr_planes_to_rgb(
    c0: &[f32],
    c1: &[f32],
    c2: &[f32],
    scale: f32,
    out: &mut [u8],
) {
    let (pixels, _) = out.as_chunks_mut::<3>();
    for (p, px) in pixels.iter_mut().enumerate() {
        px[0] = to_u8(c2[p], scale);
        px[1] = to_u8(c1[p], scale);
        px[2] = to_u8(c0[p], scale);
    }
}

#[inline]
fn to_u8(v: f32, scale: f32) -> u8 {
    (v * scale).clamp(0.0, 255.0) as u8
}

#[cfg(test)]
mod tests {
    use super::*;

    fn make_rgb(width: usize, height: usize) -> Vec<u8> {
        let n = width * height * 3;
        // Deterministic pseudo-random bytes covering the full range.
        (0..n).map(|i| ((i * 37 + 11) % 256) as u8).collect()
    }

    #[test]
    fn chw_simd_matches_scalar_rgb_and_bgr() {
        let (w, h) = (37, 19); // not a multiple of 8 -> exercises the tail
        let rgb = make_rgb(w, h);
        let alpha = [1.0 / 255.0, 0.5, 2.0];
        let beta = [-0.485, 0.1, -1.0];
        for src in [[0, 1, 2], [2, 1, 0]] {
            let mut a = vec![0.0f32; w * h * 3];
            let mut b = vec![0.0f32; w * h * 3];
            normalize_chw_scalar(&rgb, w, h, src, &alpha, &beta, &mut a);
            normalize_chw_into(&rgb, w, h, src, &alpha, &beta, &mut b);
            assert_eq!(
                a, b,
                "CHW SIMD must be bit-identical to scalar (src={src:?})"
            );
        }
    }

    #[test]
    fn hwc_matches_scalar_rgb_and_bgr() {
        let (w, h) = (23, 7);
        let rgb = make_rgb(w, h);
        let alpha = [1.0 / 255.0, 0.5, 2.0];
        let beta = [-0.485, 0.1, -1.0];
        for src in [[0, 1, 2], [2, 1, 0]] {
            let mut a = vec![0.0f32; w * h * 3];
            let mut b = vec![0.0f32; w * h * 3];
            normalize_hwc_scalar(&rgb, w, h, src, &alpha, &beta, &mut a);
            normalize_hwc_into(&rgb, w, h, src, &alpha, &beta, &mut b);
            assert_eq!(a, b);
        }
    }

    #[test]
    fn argmax_matches_scalar_including_ties() {
        let bits = |result: Option<(usize, f32)>| result.map(|(i, value)| (i, value.to_bits()));
        // Tail-exercising length, with deterministic values.
        let row: Vec<f32> = (0..101).map(|i| ((i * 17) % 13) as f32 * 0.5).collect();
        assert_eq!(argmax(&row), argmax_scalar(&row));

        // Explicit ties: the maximum 9.0 appears at indices 3 and 6 -> last wins.
        let tied = vec![1.0f32, 2.0, 5.0, 9.0, 4.0, 8.0, 9.0, 0.0];
        assert_eq!(argmax_scalar(&tied), Some((6, 9.0)));
        assert_eq!(argmax(&tied), Some((6, 9.0)));

        // Single element and empty.
        assert_eq!(argmax(&[42.0]), Some((0, 42.0)));
        assert_eq!(argmax(&[]), None);

        for len in [1, 7, 8, 9, 15, 16, 17, 31, 32, 33, 101] {
            let mut row = vec![-3.0; len];
            for i in [0, 6, 7, 8, 14, 15, 16, 23, 24, len - 1] {
                if i < len {
                    row[i] = 9.0;
                    assert_eq!(bits(argmax(&row)), bits(argmax_scalar(&row)));
                }
            }
            for zero in [-0.0, 0.0] {
                for (i, value) in row.iter_mut().enumerate() {
                    *value = if i.is_multiple_of(2) { zero } else { -zero };
                }
                #[cfg(feature = "simd")]
                let expected = argmax_two_pass_simd(&row);
                #[cfg(not(feature = "simd"))]
                let expected = argmax_scalar(&row);
                assert_eq!(bits(argmax(&row)), bits(expected), "zero ties, len={len}");
                assert_eq!(argmax(&row).unwrap().0, len - 1);
            }
            for value in [
                f32::NEG_INFINITY,
                f32::INFINITY,
                -0.0,
                0.0,
                f32::NAN,
                f32::from_bits(0xffc0_0123),
                f32::from_bits(0x7f80_0001),
            ] {
                row.fill(value);
                for i in 0..=len {
                    // Also cover mixed exceptional values at every lane/tail boundary.
                    if i > 0 {
                        row[i - 1] = if i % 2 == 0 { -0.0 } else { f32::NAN };
                    }
                    #[cfg(feature = "simd")]
                    let expected = argmax_two_pass_simd(&row);
                    #[cfg(not(feature = "simd"))]
                    let expected = argmax_scalar(&row);
                    assert_eq!(bits(argmax(&row)), bits(expected), "len={len}, i={i}");
                }
            }
        }
        // A 9, NaN, 5 lane sequence must match the target's two-pass reduction,
        // including payload and signaling NaNs.
        for nan in [
            f32::NAN,
            f32::from_bits(0xffc0_0123),
            f32::from_bits(0x7f80_0001),
        ] {
            for lane in 0..8 {
                let mut row = vec![f32::NEG_INFINITY; 25];
                row[lane] = 9.0;
                row[lane + 8] = nan;
                row[lane + 16] = 5.0;
                #[cfg(feature = "simd")]
                let expected = argmax_two_pass_simd(&row);
                #[cfg(not(feature = "simd"))]
                let expected = argmax_scalar(&row);
                assert_eq!(bits(argmax(&row)), bits(expected), "lane={lane}");
            }
            let mut row = vec![1.0; 17];
            row[0] = 9.0;
            row[16] = nan; // The only NaN is in the scalar tail.
            #[cfg(feature = "simd")]
            let expected = argmax_two_pass_simd(&row);
            #[cfg(not(feature = "simd"))]
            let expected = argmax_scalar(&row);
            assert_eq!(bits(argmax(&row)), bits(expected), "tail NaN");
        }
    }

    #[test]
    fn crnn_simd_matches_scalar_with_padding() {
        // resized_w not a multiple of 8, tensor_width > resized_w (padding).
        let (resized_w, img_h, tensor_width) = (21usize, 32usize, 40usize);
        let rgb = make_rgb(resized_w, img_h);

        let mut a = vec![0.0f32; 3 * img_h * tensor_width];
        let mut b = vec![0.0f32; 3 * img_h * tensor_width];
        // Scalar reference, emulating the original crnn.rs loop exactly.
        let plane = img_h * tensor_width;
        for y in 0..img_h {
            for x in 0..resized_w {
                let s = (y * resized_w + x) * 3;
                let idx = y * tensor_width + x;
                a[idx] = (rgb[s + 2] as f32 / 255.0 - 0.5) / 0.5;
                a[plane + idx] = (rgb[s + 1] as f32 / 255.0 - 0.5) / 0.5;
                a[2 * plane + idx] = (rgb[s] as f32 / 255.0 - 0.5) / 0.5;
            }
        }
        normalize_crnn_chw_into(&rgb, resized_w, img_h, tensor_width, &mut b);
        assert_eq!(
            a, b,
            "CRNN SIMD must be bit-identical to scalar (incl. padding)"
        );
    }

    #[test]
    fn uvdoc_conversion_matches_reference() {
        let plane = 8 * 7 + 3;
        // Values spanning below 0 and above 255 to exercise both clamp bounds.
        let c0: Vec<f32> = (0..plane).map(|i| (i as f32) * 1.9 - 5.0).collect();
        let c1: Vec<f32> = (0..plane).map(|i| (i as f32) * -2.3 + 260.0).collect();
        let c2: Vec<f32> = (0..plane).map(|i| (i as f32) * 0.7).collect();
        let scale = 1.0;

        // Independent reference matching the original uvdoc per-pixel arithmetic.
        let mut expected = vec![0u8; plane * 3];
        for p in 0..plane {
            expected[p * 3] = (c2[p] * scale).clamp(0.0, 255.0) as u8;
            expected[p * 3 + 1] = (c1[p] * scale).clamp(0.0, 255.0) as u8;
            expected[p * 3 + 2] = (c0[p] * scale).clamp(0.0, 255.0) as u8;
        }
        let mut got = vec![0u8; plane * 3];
        scale_clamp_bgr_planes_to_rgb(&c0, &c1, &c2, scale, &mut got);
        assert_eq!(expected, got);
    }
}