Merge pull request #373 from AdaWorldAPI/feat/simd-narrow-ladder-20261024

simd: saturating narrows down the 512-bit ladder (16x32 -> 32x16 -> 64x8)

Author: AdaWorldAPI

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

@@ -31,8 +31,9 @@
 //! over every pair of 16 IEEE edge values: NaN, ±0, ±inf, subnormal); `0x73x` / `0x74x` / `0x750` the strided 16-byte register match (all sixteen
 //! bytes; strides 16 / 20 / 32 / 512 from first offsets 0..15, care bits in
 //! bytes 12..16); `0xFxx` the 16-lane byte vectors `I8x16` / `U8x16` (wrapping
-//! `add`/`sub`, signed/unsigned `min`/`max`, the load/store round trips); `0x10xx` the saturating
-//! `I32x16 → I8x16 / U8x16` narrows and the `F32x16 ⇄ I32x16` casts on in-range lanes. `main.rs` (native / qemu) and
+//! `add`/`sub`, signed/unsigned `min`/`max`, the load/store round trips); `0x10xx`
+//! the saturating narrows down the 512-bit ladder (`I32x16 → I16x32 → I8x64 /
+//! U8x64`) and the `F32x16` ↔ `I32x16` casts. `main.rs` (native / qemu) and
 //! `selfcheck()` (the wasm cdylib export, driven by `run.mjs`) both call
 //! [`run`].
 
@@ -49,7 +50,7 @@ use ndarray::simd::{
     ne_i32_to_mask_under, ne_u32_to_mask, ne_u32_to_mask_under, ne_u64_to_mask, ne_u8_to_mask,
     ternary_match_strided16_to_mask, ternary_match_strided_to_mask, ternary_match_u32_to_mask,
     ternary_match_u32_to_mask_under, ternary_match_u64_to_mask, ternary_match_u64_to_mask_under, ternlog, F32x16,
-    F64x8, I32x16, I8x16, KeyRunCarry, MortonDir, U32x16, U64x8, U8x16,
+    F64x8, I16x32, I32x16, I8x16, KeyRunCarry, MortonDir, U32x16, U64x8, U8x16,
 };
 
 /// Number of check groups [`run`] executes (for the log line only).
@@ -67,7 +68,7 @@ pub fn run() -> u32 {
         check_ternlog_all_tables, check_u64x8_algebra, check_i32x16_compare, check_predicates_to_mask,
         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_and_casts,
+        check_f64x8_compare, check_i8x16_u8x16_lanes, check_care_match16, check_narrow_ladder,
     ];
     for g in groups {
         if let Err(code) = g() {
@@ -1845,59 +1846,93 @@ fn check_gather_scatter_group() -> Result<(), u32> {
     Ok(())
 }
 
-/// `I32x16::narrow_saturate_i8` / `narrow_saturate_u8` against `clamp`, and the
-/// `F32x16::cast_i32` / `I32x16::cast_f32` casts against `as` on in-range lanes.
-fn check_narrow_and_casts() -> Result<(), u32> {
-    let edges: [i32; 16] =
-        [i32::MIN, -1_000_000, -256, -129, -128, -127, -1, 0, 1, 126, 127, 128, 254, 255, 256, i32::MAX];
-    let mut rng = SplitMix64(0x0C0D_EC08_0000_0001);
-    let mut vecs: Vec<[i32; 16]> = vec![edges];
-    for k in 0..512 {
-        vecs.push(core::array::from_fn(|_| {
-            let r = rng.next();
-            // Half the vectors draw from the whole i32 range, half from the saturation band.
-            if k % 2 == 0 {
-                r as i32
+/// The saturating narrows down the 512-bit ladder (16×32 → 32×16 → 64×8) and the in-range
+/// f32 ↔ i32 casts, against `clamp` / `as` references.
+fn check_narrow_ladder() -> Result<(), u32> {
+    let mut rng = SplitMix64(0x1000_0000_0017);
+    let edges32 =
+        [i32::MIN, i32::MIN + 1, -32769, -32768, -32767, -256, -129, -128, -127, -1, 0, 1, 127, 128, 255, 256];
+    let edges16 = [i16::MIN, -32767, -300, -256, -129, -128, -127, -1, 0, 1, 127, 128, 254, 255, 256, i16::MAX];
+    let (mut sat, mut kept) = (0usize, 0usize);
+    for round in 0..64u32 {
+        // Round 0 is the edge set; later rounds mix random words with values near the seams.
+        let a: [i32; 32] = core::array::from_fn(|i| match (round, rng.next()) {
+            (0, _) => edges32[i % 16].wrapping_add(if i >= 16 { 1 } else { 0 }),
+            (_, r) if r % 3 == 0 => r as i32,
+            (_, r) => (r as i32) >> (16 + (r >> 40) % 16),
+        });
+        let lo = I32x16::from_array(a[..16].try_into().unwrap());
+        let hi = I32x16::from_array(a[16..].try_into().unwrap());
+        let n16 = lo.narrow_saturating_i16(hi).to_array();
+        for i in 0..32 {
+            let want = a[i].clamp(i16::MIN as i32, i16::MAX as i32) as i16;
+            if n16[i] != want {
+                return Err(0x1000 + i as u32);
+            }
+            if want as i32 == a[i] {
+                kept += 1
             } else {
-                (r % 601) as i32 - 300
+                sat += 1
             }
-        }));
-    }
-    let (mut below, mut above) = (0usize, 0usize);
-    for v in &vecs {
-        let x = I32x16::from_array(*v);
-        let want_i8: [i8; 16] = core::array::from_fn(|i| v[i].clamp(-128, 127) as i8);
-        let want_u8: [u8; 16] = core::array::from_fn(|i| v[i].clamp(0, 255) as u8);
-        if x.narrow_saturate_i8().to_array() != want_i8 {
-            return Err(0x1001);
         }
-        if x.narrow_saturate_u8().to_array() != want_u8 {
-            return Err(0x1002);
+        // Chained to 8 bits: two saturating steps equal one clamp to the narrow range.
+        let n16v = I16x32::from_array(n16);
+        let back: [i16; 32] = core::array::from_fn(|i| n16[31 - i]);
+        let (i8s, u8s) = (
+            n16v.narrow_saturating_i8(I16x32::from_array(back))
+                .to_array(),
+            n16v.narrow_saturating_u8(I16x32::from_array(back))
+                .to_array(),
+        );
+        for i in 0..64 {
+            let src = if i < 32 {
+                a[i]
+            } else {
+                a[63 - i].clamp(i16::MIN as i32, i16::MAX as i32)
+            };
+            if i8s[i] != src.clamp(-128, 127) as i8 {
+                return Err(0x1040 + i as u32);
+            }
+            if u8s[i] != src.clamp(0, 255) as u8 {
+                return Err(0x1080 + i as u32);
+            }
         }
-        below += v.iter().filter(|&&a| a < -128).count();
-        above += v.iter().filter(|&&a| a > 255).count();
-        let want_f: [u32; 16] = core::array::from_fn(|i| (v[i] as f32).to_bits());
-        if x.cast_f32().to_array().map(f32::to_bits) != want_f {
-            return Err(0x1003);
+    }
+    // The i16 edge set directly, both halves.
+    let e = I16x32::from_array(core::array::from_fn(|i| edges16[i % 16]));
+    let f = I16x32::from_array(core::array::from_fn(|i| edges16[15 - i % 16]));
+    let (i8s, u8s) = (e.narrow_saturating_i8(f).to_array(), e.narrow_saturating_u8(f).to_array());
+    for i in 0..64 {
+        let v = if i < 32 {
+            edges16[i % 16]
+        } else {
+            edges16[15 - (i - 32) % 16]
+        };
+        if i8s[i] != v.clamp(-128, 127) as i8 || u8s[i] != v.clamp(0, 255) as u8 {
+            return Err(0x10C0 + i as u32);
         }
     }
-    // Anti-vacuity: the corpus must actually saturate in both directions.
-    if below < 1000 || above < 1000 {
-        return Err(0x10FF);
+    // Anti-vacuity: the rounds must both saturate and pass lanes through.
+    if sat < 64 || kept < 64 {
+        return Err(0x10F0);
     }
-    // In-range f32 → i32 truncation, including negative fractions, ±0 and the ±2^31 edges.
-    let fedges: [f32; 16] = [
-        -0.0, 0.0, 0.5, -0.5, 0.99999994, -0.99999994, 1.5, -1.5, 127.5, -128.5, 255.9, -2147483520.0, 2147483520.0,
-        8388607.5, -8388607.5, 1e9,
-    ];
-    let mut fvecs: Vec<[f32; 16]> = vec![fedges];
-    for _ in 0..512 {
-        fvecs.push(core::array::from_fn(|_| ((rng.next() >> 11) as f64 / (1u64 << 53) as f64 * 4.0e9 - 2.0e9) as f32));
-    }
-    for v in &fvecs {
-        let want: [i32; 16] = core::array::from_fn(|i| v[i] as i32);
-        if F32x16::from_array(*v).cast_i32().to_array() != want {
-            return Err(0x1004);
+    // Casts, in range only (out-of-range and NaN differ by backend by design).
+    for _ in 0..64 {
+        let v: [i32; 16] = core::array::from_fn(|_| (rng.next() as i32) >> (rng.next() % 31));
+        let f = I32x16::from_array(v).cast_f32().to_array();
+        for i in 0..16 {
+            if f[i].to_bits() != (v[i] as f32).to_bits() {
+                return Err(0x10F1);
+            }
+        }
+        let x: [f32; 16] = core::array::from_fn(|_| {
+            f32::from_bits((rng.next() as u32 & 0x807F_FFFF) | (((rng.next() % 31) as u32 + 112) << 23))
+        });
+        let t = F32x16::from_array(x).cast_i32().to_array();
+        for i in 0..16 {
+            if t[i] != x[i] as i32 {
+                return Err(0x10F2);
+            }
         }
     }
     Ok(())

src/simd.rs

@@ -429,6 +429,10 @@ pub use crate::simd_neon::{
     batch_packed_i4_16, i8x16, i8x32, palette_lookup_u8x8, prefetch_read_t0, prefetch_read_t1, prefetch_read_t2, u8x16,
     u8x8, I8x16, I8x32, U8x16, U8x8,
 };
+// The 64×8 and 32×16 signed rungs of the ladder (64×8, 32×16, 16×32). `simd_neon.rs` has
+// carried them since #337 but no arm re-exported them, so aarch64 had neither type.
+#[cfg(all(target_arch = "aarch64", not(feature = "nightly-simd")))]
+pub use crate::simd_neon::{i16x32, i8x64, I16x32, I8x64};
 // U16x8 on aarch64 comes from simd_neon (backed by uint16x8_t)
 #[cfg(all(target_arch = "aarch64", not(feature = "nightly-simd")))]
 pub use crate::simd_neon::{u16x8, U16x8};

src/simd_avx2.rs

@@ -1389,6 +1389,67 @@ impl I8x64 {
 }
 
 impl I16x32 {
+    /// Narrow two 32×i16 vectors to one 64×i8 vector with signed saturation: lanes 0..32
+    /// come from `self`, lanes 32..64 from `hi`, each clamped to `-128..=127`.
+    ///
+    /// ```
+    /// use ndarray::simd::I16x32;
+    /// let mut a = [0i16; 32];
+    /// a[..8].copy_from_slice(&[i16::MIN, -129, -128, -1, 1, 127, 128, i16::MAX]);
+    /// let v = I16x32::from_array(a).narrow_saturating_i8(I16x32::splat(-300)).to_array();
+    /// assert_eq!(v[..8], [-128, -128, -128, -1, 1, 127, 127, 127]);
+    /// assert!(v[32..].iter().all(|&x| x == -128));
+    /// ```
+    #[inline(always)]
+    pub fn narrow_saturating_i8(self, hi: Self) -> I8x64 {
+        use core::arch::x86_64::*;
+        let mut out = [0i8; 64];
+        // SAFETY: SSE2 baseline. The loads read 16-byte quarters of the two 64-byte arrays and
+        // the stores write 16-byte quarters of the 64-byte `out`.
+        unsafe {
+            let o = out.as_mut_ptr() as *mut __m128i;
+            for (h, src) in [&self.0, &hi.0].into_iter().enumerate() {
+                let p = src.as_ptr() as *const __m128i;
+                _mm_storeu_si128(o.add(2 * h), _mm_packs_epi16(_mm_loadu_si128(p), _mm_loadu_si128(p.add(1))));
+                _mm_storeu_si128(
+                    o.add(2 * h + 1),
+                    _mm_packs_epi16(_mm_loadu_si128(p.add(2)), _mm_loadu_si128(p.add(3))),
+                );
+            }
+        }
+        I8x64(out)
+    }
+
+    /// Narrow two 32×i16 vectors to one 64×u8 vector with saturation: lanes 0..32 come
+    /// from `self`, lanes 32..64 from `hi`, each clamped to `0..=255`.
+    ///
+    /// ```
+    /// use ndarray::simd::I16x32;
+    /// let mut a = [0i16; 32];
+    /// a[..8].copy_from_slice(&[i16::MIN, -1, 0, 1, 254, 255, 256, i16::MAX]);
+    /// let v = I16x32::from_array(a).narrow_saturating_u8(I16x32::splat(300)).to_array();
+    /// assert_eq!(v[..8], [0, 0, 0, 1, 254, 255, 255, 255]);
+    /// assert!(v[32..].iter().all(|&x| x == 255));
+    /// ```
+    #[inline(always)]
+    pub fn narrow_saturating_u8(self, hi: Self) -> U8x64 {
+        use core::arch::x86_64::*;
+        let mut out = [0u8; 64];
+        // SAFETY: as in `narrow_saturating_i8`; `packuswb` saturates signed i16 to 0..=255.
+        unsafe {
+            let o = out.as_mut_ptr() as *mut __m128i;
+            for (h, src) in [&self.0, &hi.0].into_iter().enumerate() {
+                let p = src.as_ptr() as *const __m128i;
+                _mm_storeu_si128(o.add(2 * h), _mm_packus_epi16(_mm_loadu_si128(p), _mm_loadu_si128(p.add(1))));
+                _mm_storeu_si128(
+                    o.add(2 * h + 1),
+                    _mm_packus_epi16(_mm_loadu_si128(p.add(2)), _mm_loadu_si128(p.add(3))),
+                );
+            }
+        }
+        U8x64(out)
+    }
+
     #[inline(always)]
     pub fn zero() -> Self {
         Self([0i16; 32])
@@ -2632,65 +2693,35 @@ impl U16x32 {
 }
 
 impl I32x16 {
-    /// The two saturating `i32 → i16` packs (`packssdw`) of lanes 0..8 and 8..16.
-    #[inline(always)]
-    fn packs_i16(self) -> (core::arch::x86_64::__m128i, core::arch::x86_64::__m128i) {
-        use core::arch::x86_64::*;
-        // SAFETY: SSE2 is part of the x86_64 baseline; the four 16-byte unaligned loads
-        // read the 64-byte array in bounds.
-        unsafe {
-            let p = self.0.as_ptr() as *const __m128i;
-            (
-                _mm_packs_epi32(_mm_loadu_si128(p), _mm_loadu_si128(p.add(1))),
-                _mm_packs_epi32(_mm_loadu_si128(p.add(2)), _mm_loadu_si128(p.add(3))),
-            )
-        }
-    }
-
-    /// Narrow each lane to `i8` with signed saturation: lanes below -128 become -128 and
-    /// lanes above 127 become 127; every other lane keeps its value.
+    /// 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`.
     ///
     /// ```
     /// use ndarray::simd::I32x16;
-    /// let v = I32x16::from_array([-1000, -129, -128, -1, 0, 1, 127, 128, 1000, 5, -5, 64, -64, 100, -100, i32::MIN]);
-    /// assert_eq!(
-    ///     v.narrow_saturate_i8().to_array(),
-    ///     [-128, -128, -128, -1, 0, 1, 127, 127, 127, 5, -5, 64, -64, 100, -100, -128]
-    /// );
+    /// let lo = I32x16::from_array([i32::MIN, -32769, -32768, -1, 0, 1, 32767, 32768, i32::MAX, 5, -5, 7, -7, 9, -9, 0]);
+    /// let hi = I32x16::splat(70_000);
+    /// let v = lo.narrow_saturating_i16(hi).to_array();
+    /// assert_eq!(v[..9], [-32768, -32768, -32768, -1, 0, 1, 32767, 32767, 32767]);
+    /// assert!(v[16..].iter().all(|&x| x == 32767));
     /// ```
     #[inline(always)]
-    pub fn narrow_saturate_i8(self) -> crate::simd_avx512::I8x16 {
+    pub fn narrow_saturating_i16(self, hi: Self) -> I16x32 {
         use core::arch::x86_64::*;
-        let (lo, hi) = self.packs_i16();
-        let mut out = [0i8; 16];
-        // SAFETY: SSE2 baseline; `packsswb` is register-only and the store writes exactly the
-        // 16 bytes of `out`.
-        unsafe { _mm_storeu_si128(out.as_mut_ptr() as *mut __m128i, _mm_packs_epi16(lo, hi)) };
-        crate::simd_avx512::I8x16(out)
-    }
-
-    /// Narrow each lane to `u8` with saturation: lanes below 0 become 0 and lanes above 255
-    /// become 255; every other lane keeps its value.
-    ///
-    /// ```
-    /// use ndarray::simd::I32x16;
-    /// let v = I32x16::from_array([-1000, -1, 0, 1, 127, 128, 255, 256, 1000, 5, 64, 200, 254, i32::MIN, i32::MAX, 0]);
-    /// assert_eq!(
-    ///     v.narrow_saturate_u8().to_array(),
-    ///     [0, 0, 0, 1, 127, 128, 255, 255, 255, 5, 64, 200, 254, 0, 255, 0]
-    /// );
-    /// ```
-    #[inline(always)]
-    pub fn narrow_saturate_u8(self) -> crate::simd_avx512::U8x16 {
-        use core::arch::x86_64::*;
-        // Saturating to i16 and then to u8 is the same as clamping to 0..=255, because both
-        // steps are monotone and 0..=255 lies inside the i16 range.
-        let (lo, hi) = self.packs_i16();
-        let mut out = [0u8; 16];
-        // SAFETY: SSE2 baseline; `packuswb` is register-only and the store writes exactly the
-        // 16 bytes of `out`.
-        unsafe { _mm_storeu_si128(out.as_mut_ptr() as *mut __m128i, _mm_packus_epi16(lo, hi)) };
-        crate::simd_avx512::U8x16(out)
+        let mut out = [0i16; 32];
+        // SAFETY: SSE2 is part of the x86_64 baseline. The loads read 16-byte quarters of the
+        // two 64-byte arrays (k < 4) and the stores write 16-byte quarters of the 64-byte `out`.
+        unsafe {
+            let o = out.as_mut_ptr() as *mut __m128i;
+            for (h, src) in [&self.0, &hi.0].into_iter().enumerate() {
+                let p = src.as_ptr() as *const __m128i;
+                _mm_storeu_si128(o.add(2 * h), _mm_packs_epi32(_mm_loadu_si128(p), _mm_loadu_si128(p.add(1))));
+                _mm_storeu_si128(
+                    o.add(2 * h + 1),
+                    _mm_packs_epi32(_mm_loadu_si128(p.add(2)), _mm_loadu_si128(p.add(3))),
+                );
+            }
+        }
+        I16x32(out)
     }
 
     /// The two 256-bit halves of the 64-byte-aligned array, loaded once.