simd: TFLite/gemmlowp fixed-point requantization, bit-exact on every backend

crates/simd-masking-parity/src/lib.rs

@@ -35,7 +35,8 @@
 //! the saturating narrows down the 512-bit ladder (`I32x16 → I16x32 → I8x64 /
 //! U8x64`) and the `F32x16` ↔ `I32x16` casts; `0x11xx` `F32x16` / `F64x8`
 //! `mul_add` rounding once, bit for bit against an exact integer reference (no
-//! `mul_add`, so no libm in the binary). `main.rs` (native / qemu) and
+//! `mul_add`, so no libm in the binary); `0x12xx` the `I32x16` gemmlowp requantization
+//! methods against an independent `i128` reference. `main.rs` (native / qemu) and
 //! `selfcheck()` (the wasm cdylib export, driven by `run.mjs`) both call
 //! [`run`].
 
@@ -56,7 +57,7 @@ use ndarray::simd::{
 };
 
 /// Number of check groups [`run`] executes (for the log line only).
-pub const CHECKS: usize = 18;
+pub const CHECKS: usize = 19;
 
 /// The wasm export: identical to [`run`], `extern "C"` so `run.mjs` can call it.
 #[no_mangle]
@@ -71,7 +72,7 @@ pub fn run() -> u32 {
         check_mask_algebra, check_care_match, check_masked_reductions, check_blend, check_morton_shift,
         check_predicates_under, check_set_range, check_unsigned_compare_to_mask, check_gather_scatter_group,
         check_f64x8_compare, check_i8x16_u8x16_lanes, check_care_match16, check_narrow_ladder,
-        check_mul_add_single_rounding,
+        check_mul_add_single_rounding, check_fixedpoint_requant,
     ];
     for g in groups {
         if let Err(code) = g() {
@@ -2077,6 +2078,161 @@ fn check_mul_add_single_rounding() -> Result<(), u32> {
     Ok(())
 }
 
+/// `I32x16::saturating_rounding_doubling_high_mul`, `rounding_divide_by_pot` and
+/// `multiply_by_quantized_multiplier` against an independent `i128` reference
+/// written from the definitions (no call into the scalar functions of
+/// `simd_fixedpoint`): srdhm is `ab / 2^31` rounded half toward +inf with the
+/// `MIN * MIN` saturation, rdpot is `x / 2^e` rounded half away from zero, and
+/// the multiplier is their composition. Anti-vacuity: the rounding must differ
+/// from truncation on some lanes, the saturating case must be hit, and exact ties
+/// of both signs must occur.
+fn check_fixedpoint_requant() -> Result<(), u32> {
+    fn o_srdhm(a: i32, b: i32) -> i32 {
+        let p = a as i128 * b as i128;
+        let q = (p + (1i128 << 30)).div_euclid(1i128 << 31);
+        q.clamp(i32::MIN as i128, i32::MAX as i128) as i32
+    }
+    fn o_rdpot(x: i32, e: u32) -> i32 {
+        let d = 1i128 << e;
+        let m = (x as i128).abs();
+        let r = (2 * m + d).div_euclid(2 * d);
+        (if x < 0 { -r } else { r }) as i32
+    }
+    fn o_mbqm(x: i32, m: i32, s: i32) -> i32 {
+        let left = s.max(0) as u32;
+        let right = (-s).max(0) as u32;
+        let xs = ((x as i128) << left) as i32; // wraps like the realization's contract
+        o_rdpot(o_srdhm(xs, m), right)
+    }
+
+    let mut rng = SplitMix64(0x1200_0000_0019);
+    const EDGE: [i32; 12] =
+        [i32::MIN, i32::MIN + 1, -1, 0, 1, i32::MAX - 1, i32::MAX, 1 << 30, -(1 << 30), (1 << 30) + 1, -(1 << 30) - 1, 1 << 15];
+    const EXPS: [u32; 4] = [0, 1, 30, 31];
+    let (mut trunc_differs_srdhm, mut trunc_differs_rdpot) = (0usize, 0usize);
+    let (mut min_min, mut tie_pos, mut tie_neg) = (0usize, 0usize, 0usize);
+
+    for round in 0..512u32 {
+        // ---- M1: srdhm ----
+        let mut a = [0i32; 16];
+        let mut b = [0i32; 16];
+        for i in 0..16 {
+            match (round as usize + i) % 4 {
+                0 => {
+                    a[i] = rng.next() as i32;
+                    b[i] = rng.next() as i32;
+                }
+                1 => {
+                    a[i] = EDGE[(rng.next() % 12) as usize];
+                    b[i] = EDGE[(rng.next() % 12) as usize];
+                }
+                2 => {
+                    // exact tie: odd * odd * 2^30
+                    let f = |r: &mut SplitMix64| {
+                        let m = ((r.next() % 0x7FFF) as i32 | 1) << 15;
+                        if r.next() & 1 == 0 { m } else { -m }
+                    };
+                    a[i] = f(&mut rng);
+                    b[i] = f(&mut rng);
+                }
+                _ => {
+                    a[i] = if i % 2 == 0 { i32::MIN } else { rng.next() as i32 };
+                    b[i] = if i % 4 == 0 { i32::MIN } else { rng.next() as i32 };
+                }
+            }
+        }
+        let got = I32x16::from_array(a).saturating_rounding_doubling_high_mul(I32x16::from_array(b)).to_array();
+        for i in 0..16 {
+            let want = o_srdhm(a[i], b[i]);
+            if got[i] != want {
+                return Err(0x1200 + i as u32);
+            }
+            let p = a[i] as i128 * b[i] as i128;
+            if want as i128 != p / (1i128 << 31) {
+                trunc_differs_srdhm += 1;
+            }
+            if a[i] == i32::MIN && b[i] == i32::MIN {
+                min_min += 1;
+            }
+            if p.rem_euclid(1i128 << 31) == 1i128 << 30 {
+                if p >= 0 { tie_pos += 1 } else { tie_neg += 1 }
+            }
+        }
+
+        // ---- M2: rounding_divide_by_pot ----
+        let mut x = [0i32; 16];
+        let mut e = [0i32; 16];
+        for i in 0..16 {
+            let ex = match (round as usize + i) % 4 {
+                0 => (rng.next() % 32) as u32,
+                1 => EXPS[(rng.next() % 4) as usize],
+                2 => EXPS[i % 4],
+                _ => 1 + (rng.next() % 31) as u32,
+            };
+            let r = rng.next() as i32;
+            x[i] = match (round as usize + i) % 3 {
+                0 => r,
+                1 if ex >= 1 => ((r >> ex) << ex) | (1i32 << (ex - 1)), // exact tie
+                1 => r,
+                _ => EDGE[(rng.next() % 12) as usize],
+            };
+            e[i] = ex as i32;
+        }
+        let got = I32x16::from_array(x).rounding_divide_by_pot(I32x16::from_array(e)).to_array();
+        for i in 0..16 {
+            let want = o_rdpot(x[i], e[i] as u32);
+            if got[i] != want {
+                return Err(0x1210 + i as u32);
+            }
+            if want as i128 != (x[i] as i128) / (1i128 << e[i]) {
+                trunc_differs_rdpot += 1;
+            }
+        }
+
+        // ---- M3: multiply_by_quantized_multiplier ----
+        let mut mx = [0i32; 16];
+        let mut mm = [0i32; 16];
+        let mut ms = [0i32; 16];
+        for i in 0..16 {
+            let s = match (round as usize + i) % 4 {
+                0 => (rng.next() % 62) as i32 - 31,
+                1 => -31 + (rng.next() % 4) as i32,
+                2 => 27 + (rng.next() % 4) as i32,
+                _ => (rng.next() % 5) as i32 - 2,
+            };
+            ms[i] = s;
+            mm[i] = match (round as usize + i) % 3 {
+                0 => (rng.next() >> 33) as i32 | (1 << 30), // normalised, as quantize_multiplier yields
+                1 => (rng.next() >> 33) as i32,
+                _ => [0, 1, 1 << 30, i32::MAX][(rng.next() % 4) as usize],
+            };
+            // keep x << max(s, 0) inside i32 (the documented contract)
+            mx[i] = (rng.next() as i32) >> (s.max(0) + 1);
+        }
+        let got = I32x16::from_array(mx)
+            .multiply_by_quantized_multiplier(I32x16::from_array(mm), I32x16::from_array(ms))
+            .to_array();
+        for i in 0..16 {
+            if got[i] != o_mbqm(mx[i], mm[i], ms[i]) {
+                return Err(0x1220 + i as u32);
+            }
+        }
+    }
+    if trunc_differs_srdhm < 256 {
+        return Err(0x12F0);
+    }
+    if trunc_differs_rdpot < 256 {
+        return Err(0x12F1);
+    }
+    if min_min == 0 {
+        return Err(0x12F2);
+    }
+    if tie_pos == 0 || tie_neg == 0 {
+        return Err(0x12F3);
+    }
+    Ok(())
+}
+
 /// Exact `a * b + c`, rounded once to nearest-even: the reference for the
 /// `0x11xx` group. It holds the exact sum in a 4608-bit integer (bit 0 is
 /// `2^ORACLE_BASE`, below every product of two subnormals), so it shares no

src/lib.rs

@@ -490,6 +490,10 @@ pub mod simd_half;
 #[cfg(feature = "std")]
 pub(crate) mod simd_fma;
 
+/// TFLite / gemmlowp fixed-point requantization (scalar reference + batch epilogue).
+#[cfg(feature = "std")]
+pub(crate) mod simd_fixedpoint;
+
 /// Runtime SIMD dispatch — release-binary distribution path.
 ///
 /// Gated under `--features runtime-dispatch`. Mutually exclusive with

src/simd.rs

@@ -997,6 +997,17 @@ pub use crate::simd_ops::{
     scale_f32_inplace, sub_f32, sub_f32_inplace,
 };
 
+// Fixed-point requantization, bit-exact to TFLite-micro / gemmlowp (double rounding,
+// TFLM's default build). The `I32x16` methods `saturating_rounding_doubling_high_mul`,
+// `rounding_divide_by_pot` and `multiply_by_quantized_multiplier` are the vector form.
+pub use crate::simd_fixedpoint::{
+    multiply_by_quantized_multiplier_i32, multiply_by_quantized_multiplier_i64,
+    multiply_by_quantized_multiplier_single_rounding_i32, requantize_i32_to_i8, rounding_divide_by_pot_i32, srdhm_i16,
+    srdhm_i32,
+};
+#[cfg(feature = "std")]
+pub use crate::simd_fixedpoint::{quantize_multiplier, quantize_multiplier_single_rounding};
+
 // ChaCha20 keystream is NO LONGER an `ndarray::simd` surface: the AdaWorldAPI
 // `chacha20` fork (`vendor/chacha20/`) carries the accelerated backend, riding
 // the `U32x16` ARX lane above, and is `[patch]`ed transitively under the

src/simd_avx.rs

@@ -271,6 +271,32 @@ impl Shr<Self> for U64x8 {
 // ---------------------------------------------------------------------------
 
 impl I32x16 {
+    /// `SaturatingRoundingDoublingHighMul` per lane (TFLM `fixedpoint.h`).
+    /// AVX without AVX2 has no 256-bit integer multiply, so this arm goes
+    /// lane by lane through [`crate::simd_fixedpoint::srdhm_i32`].
+    #[inline]
+    pub fn saturating_rounding_doubling_high_mul(self, b: Self) -> Self {
+        let (a, b) = (self.to_array(), b.to_array());
+        Self::from_array(core::array::from_fn(|i| crate::simd_fixedpoint::srdhm_i32(a[i], b[i])))
+    }
+
+    /// `RoundingDivideByPOT` per lane, per-lane exponent in `0..=31`.
+    #[inline]
+    pub fn rounding_divide_by_pot(self, exponent: Self) -> Self {
+        let (x, e) = (self.to_array(), exponent.to_array());
+        Self::from_array(core::array::from_fn(|i| crate::simd_fixedpoint::rounding_divide_by_pot_i32(x[i], e[i])))
+    }
+
+    /// `MultiplyByQuantizedMultiplier` (double rounding) per lane, with
+    /// per-lane multiplier and shift.
+    #[inline]
+    pub fn multiply_by_quantized_multiplier(self, multiplier: Self, shift: Self) -> Self {
+        let (x, m, s) = (self.to_array(), multiplier.to_array(), shift.to_array());
+        Self::from_array(core::array::from_fn(|i| {
+            crate::simd_fixedpoint::multiply_by_quantized_multiplier_i32(x[i], m[i], s[i])
+        }))
+    }
+
     /// Four 128-bit quarters of the 16 lanes.
     #[inline(always)]
     fn quarters(self) -> [__m128i; 4] {

src/simd_avx2.rs

@@ -2682,6 +2682,140 @@ impl U16x32 {
     }
 }
 
+#[cfg(not(all(target_feature = "avx", not(target_feature = "avx2"))))]
+impl I32x16 {
+    /// Store two 256-bit halves back into a fresh array (see `avx2_halves`).
+    #[inline(always)]
+    fn from_avx2_halves(lo: __m256i, hi: __m256i) -> Self {
+        let mut o = [0i32; 16];
+        // SAFETY: two 32-byte `storeu` stores into the 64-byte local array; no alignment needed.
+        unsafe {
+            let p = o.as_mut_ptr() as *mut __m256i;
+            _mm256_storeu_si256(p, lo);
+            _mm256_storeu_si256(p.add(1), hi);
+        }
+        Self(o)
+    }
+
+    /// `SaturatingRoundingDoublingHighMul` on one 256-bit half (exactness argument: see
+    /// [`Self::saturating_rounding_doubling_high_mul`]).
+    #[inline(always)]
+    fn srdhm_half(a: __m256i, b: __m256i) -> __m256i {
+        // SAFETY: AVX2 is a run-time precondition of this arm (see `avx2_halves`); register-only.
+        unsafe {
+            let nudge = _mm256_set1_epi64x(1 << 30);
+            let even = _mm256_srli_epi64(_mm256_add_epi64(_mm256_mul_epi32(a, b), nudge), 31);
+            let odd = _mm256_slli_epi64(
+                _mm256_add_epi64(_mm256_mul_epi32(_mm256_srli_epi64(a, 32), _mm256_srli_epi64(b, 32)), nudge),
+                1,
+            );
+            let r = _mm256_blend_epi32::<0xAA>(even, odd);
+            let min = _mm256_set1_epi32(i32::MIN);
+            let ovf = _mm256_and_si256(_mm256_cmpeq_epi32(a, min), _mm256_cmpeq_epi32(b, min));
+            _mm256_blendv_epi8(r, _mm256_set1_epi32(i32::MAX), ovf)
+        }
+    }
+
+    /// `RoundingDivideByPOT` on one 256-bit half, per-lane exponent.
+    #[inline(always)]
+    fn rdpot_half(x: __m256i, e: __m256i) -> __m256i {
+        // SAFETY: AVX2 run-time precondition (see `avx2_halves`); register-only. `vpsllvd` /
+        // `vpsravd` define all counts, so an out-of-contract exponent is unspecified, not UB.
+        unsafe {
+            let one = _mm256_set1_epi32(1);
+            let mask = _mm256_sub_epi32(_mm256_sllv_epi32(one, e), one);
+            let rem = _mm256_and_si256(x, mask);
+            let thr = _mm256_add_epi32(_mm256_srli_epi32(mask, 1), _mm256_srli_epi32(x, 31));
+            let q = _mm256_srav_epi32(x, e);
+            // `cmpgt` lanes are -1 where rem > thr, so subtracting adds 1 there.
+            _mm256_sub_epi32(q, _mm256_cmpgt_epi32(rem, thr))
+        }
+    }
+
+    /// Saturating rounding doubling high multiply (gemmlowp / TFLM
+    /// `SaturatingRoundingDoublingHighMul<i32>`): per lane, `round(a * b / 2^31)`
+    /// with the half-way case rounded toward `+inf`, i.e. the C++ definition
+    /// `((a*b + nudge) / 2^31)` with `nudge = 1<<30` (product `>= 0`) or `1 - (1<<30)`
+    /// (product `< 0`) and **truncating** division. `i32::MIN * i32::MIN` saturates to
+    /// `i32::MAX`.
+    ///
+    /// ```
+    /// use ndarray::simd::I32x16;
+    /// let a = I32x16::from_array([3, -3, -1, 1, 0, i32::MIN, i32::MAX, 0, 0, 0, 0, 0, 0, 0, 0, 0]);
+    /// let r = a.saturating_rounding_doubling_high_mul(I32x16::splat(1 << 30)).to_array();
+    /// // 1.5 -> 2, -1.5 -> -1 (tie toward +inf), -0.5 -> 0, 0.5 -> 1.
+    /// assert_eq!(r[..7], [2, -1, 0, 1, 0, -1_073_741_824, 1_073_741_824]);
+    /// let m = I32x16::splat(i32::MIN).saturating_rounding_doubling_high_mul(I32x16::splat(i32::MIN));
+    /// assert!(m.to_array().iter().all(|&x| x == i32::MAX));
+    /// ```
+    #[inline(always)]
+    pub fn saturating_rounding_doubling_high_mul(self, b: Self) -> Self {
+        // Exact for both product signs: C++ truncating division after the sign-dependent nudge
+        // equals the flooring shift of `a*b + 2^30` (taking bits 31..=62 as the i32), computed
+        // with `vpmuldq` on even lanes and on odd lanes after a 32-bit shift down; `MIN*MIN`
+        // (2^31, not representable) is blended to `MAX`.
+        let (alo, ahi) = self.avx2_halves();
+        let (blo, bhi) = b.avx2_halves();
+        Self::from_avx2_halves(Self::srdhm_half(alo, blo), Self::srdhm_half(ahi, bhi))
+    }
+
+    /// Rounding arithmetic right shift (gemmlowp / TFLM `RoundingDivideByPOT`) with a
+    /// per-lane exponent: `x / 2^e` rounded to nearest, **ties away from zero**.
+    /// `exponent` lanes must be in `0..=31`; other values are a contract violation
+    /// (the result is unspecified, but no undefined behaviour).
+    ///
+    /// ```
+    /// use ndarray::simd::I32x16;
+    /// let x = I32x16::from_array([5, -5, 3, -3, 2, -2, 1, -1, 0, 0, 0, 0, 0, 0, 0, 0]);
+    /// let r = x.rounding_divide_by_pot(I32x16::splat(1)).to_array();
+    /// assert_eq!(r[..8], [3, -3, 2, -2, 1, -1, 1, -1]); // negative ties go away from zero
+    /// let e = I32x16::from_array([0, 1, 2, 3, 4, 5, 0, 0, 0, 0, 0, 0, 0, 0, 0, 31]);
+    /// let r = I32x16::splat(-10).rounding_divide_by_pot(e).to_array();
+    /// assert_eq!(r[..6], [-10, -5, -3, -1, -1, 0]);
+    /// assert_eq!(r[15], 0);
+    /// let m = I32x16::splat(i32::MIN).rounding_divide_by_pot(I32x16::splat(31)).to_array();
+    /// assert_eq!(m[0], -1);
+    /// ```
+    #[inline(always)]
+    pub fn rounding_divide_by_pot(self, exponent: Self) -> Self {
+        let (xlo, xhi) = self.avx2_halves();
+        let (elo, ehi) = exponent.avx2_halves();
+        Self::from_avx2_halves(Self::rdpot_half(xlo, elo), Self::rdpot_half(xhi, ehi))
+    }
+
+    /// Per-lane fixed-point requantization (gemmlowp / TFLM `MultiplyByQuantizedMultiplier`,
+    /// double rounding): with `left = max(shift, 0)`, `right = max(-shift, 0)`,
+    /// `rounding_divide_by_pot(srdhm(x << left, multiplier), right)`.
+    ///
+    /// Contract: `multiplier >= 0`, `-31 <= shift <= 30`, and `x << left` must not
+    /// overflow (TFLM asserts; this wraps).
+    ///
+    /// ```
+    /// use ndarray::simd::I32x16;
+    /// let x = I32x16::from_array([100, -100, 7, -7, 7, -7, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0]);
+    /// let s = I32x16::from_array([0, 0, 0, 0, -1, -1, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0]);
+    /// let r = x.multiply_by_quantized_multiplier(I32x16::splat(1 << 30), s).to_array();
+    /// assert_eq!(r[..7], [50, -50, 4, -3, 2, -2, 6]);
+    /// ```
+    #[inline(always)]
+    pub fn multiply_by_quantized_multiplier(self, multiplier: Self, shift: Self) -> Self {
+        let (xlo, xhi) = self.avx2_halves();
+        let (mlo, mhi) = multiplier.avx2_halves();
+        let (slo, shi) = shift.avx2_halves();
+        // SAFETY: AVX2 run-time precondition (see `avx2_halves`); register-only (`vpsllvd` wraps,
+        // like the scalar `wrapping_mul(1 << left)`).
+        unsafe {
+            let zero = _mm256_setzero_si256();
+            let part = |x: __m256i, m: __m256i, s: __m256i| {
+                let left = _mm256_max_epi32(s, zero);
+                let right = _mm256_max_epi32(_mm256_sub_epi32(zero, s), zero);
+                Self::rdpot_half(Self::srdhm_half(_mm256_sllv_epi32(x, left), m), right)
+            };
+            Self::from_avx2_halves(part(xlo, mlo, slo), part(xhi, mhi, shi))
+        }
+    }
+}
+
 impl I32x16 {
     /// Narrow two 16×i32 vectors to one 32×i16 vector with signed saturation: lanes 0..16
     /// come from `self`, lanes 16..32 from `hi`, each clamped to `i16::MIN..=i16::MAX`.