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

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

@@ -31,7 +31,8 @@
 //! 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). `main.rs` (native / qemu) and
+//! `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
 //! `selfcheck()` (the wasm cdylib export, driven by `run.mjs`) both call
 //! [`run`].
 
@@ -40,20 +41,19 @@ use ndarray::simd::{
     eq_u32_via_to_mask, eq_u64_to_mask, eq_u8_to_mask, ge_i32_to_mask, ge_i32_to_mask_under, ge_u64_to_mask,
     ge_u8_to_mask, gt_i32_to_mask, gt_i32_to_mask_under, gt_u64_to_mask, gt_u8_to_mask, le_i32_to_mask,
     le_i32_to_mask_under, le_u64_to_mask, le_u8_to_mask, lt_i32_to_mask, lt_i32_to_mask_under, lt_u64_to_mask,
-    lt_u8_to_mask, mask_all, mask_and, mask_and_assign, mask_andnot, mask_andnot_assign, mask_any,
-    mask_gather_u32, mask_gather_u32_under, mask_not,
-    mask_not_assign, mask_or, mask_or_assign, mask_scatter_or_u32, mask_set_range, mask_shift_morton, mask_ternlog,
-    mask_ternlog_any, mask_ternlog_assign, mask_ternlog_popcount, mask_xor, mask_xor_assign, masked_group_sum_i32,
-    masked_group_sum_i32_via, masked_key_run_count_u32, masked_max_i32, masked_min_i32, masked_strided_group_sum,
-    masked_sum_i32, masked_sum_wrapping_add_i32, ne_i32_to_mask, 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, F64x8, I32x16, I8x16,
-    KeyRunCarry, MortonDir, U32x16, U64x8, U8x16,
+    lt_u8_to_mask, mask_all, mask_and, mask_and_assign, mask_andnot, mask_andnot_assign, mask_any, mask_gather_u32,
+    mask_gather_u32_under, mask_not, mask_not_assign, mask_or, mask_or_assign, mask_scatter_or_u32, mask_set_range,
+    mask_shift_morton, mask_ternlog, mask_ternlog_any, mask_ternlog_assign, mask_ternlog_popcount, mask_xor,
+    mask_xor_assign, masked_group_sum_i32, masked_group_sum_i32_via, masked_key_run_count_u32, masked_max_i32,
+    masked_min_i32, masked_strided_group_sum, masked_sum_i32, masked_sum_wrapping_add_i32, ne_i32_to_mask,
+    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,
 };
 
 /// Number of check groups [`run`] executes (for the log line only).
-pub const CHECKS: usize = 16;
+pub const CHECKS: usize = 17;
 
 /// The wasm export: identical to [`run`], `extern "C"` so `run.mjs` can call it.
 #[no_mangle]
@@ -67,7 +67,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_f64x8_compare, check_i8x16_u8x16_lanes, check_care_match16, check_narrow_and_casts,
     ];
     for g in groups {
         if let Err(code) = g() {
@@ -1844,3 +1844,61 @@ 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
+            } else {
+                (r % 601) as i32 - 300
+            }
+        }));
+    }
+    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);
+        }
+        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);
+        }
+    }
+    // Anti-vacuity: the corpus must actually saturate in both directions.
+    if below < 1000 || above < 1000 {
+        return Err(0x10FF);
+    }
+    // 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);
+        }
+    }
+    Ok(())
+}

src/simd_avx2.rs

@@ -2632,6 +2632,67 @@ 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.
+    ///
+    /// ```
+    /// 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]
+    /// );
+    /// ```
+    #[inline(always)]
+    pub fn narrow_saturate_i8(self) -> crate::simd_avx512::I8x16 {
+        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)
+    }
+
     /// The two 256-bit halves of the 64-byte-aligned array, loaded once.
     #[inline(always)]
     fn avx2_halves(self) -> (__m256i, __m256i) {

src/simd_avx512.rs

@@ -1007,6 +1007,49 @@ impl I32x16 {
     pub fn cast_f32(self) -> F32x16 {
         F32x16(unsafe { _mm512_cvtepi32_ps(self.0) })
     }
+
+    /// 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.
+    ///
+    /// ```
+    /// 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]
+    /// );
+    /// ```
+    #[inline(always)]
+    pub fn narrow_saturate_i8(self) -> I8x16 {
+        let mut out = [0i8; 16];
+        // SAFETY: avx512f is compiled in on this arm; `vpmovsdb` is register-only and the
+        // unaligned store writes exactly the 16 bytes of `out`.
+        unsafe { _mm_storeu_si128(out.as_mut_ptr() as *mut __m128i, _mm512_cvtsepi32_epi8(self.0)) };
+        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) -> U8x16 {
+        let mut out = [0u8; 16];
+        // SAFETY: as in `narrow_saturate_i8`. `vpmovusdb` reads its lanes as unsigned, so
+        // negative lanes are raised to 0 first.
+        unsafe {
+            let v = _mm512_max_epi32(self.0, _mm512_setzero_si512());
+            _mm_storeu_si128(out.as_mut_ptr() as *mut __m128i, _mm512_cvtusepi32_epi8(v));
+        }
+        U8x16(out)
+    }
 }
 
 impl_bin_op!(I32x16, Add, add, _mm512_add_epi32);

src/simd_neon.rs

@@ -3213,6 +3213,51 @@ fn quad_mask4(cmp: uint32x4_t) -> u16 {
 impl I32x16 {
     pub const LANES: usize = 16;
 
+    /// The two saturating `i32 → i16` narrows (`sqxtn` ×4) of lanes 0..8 and 8..16.
+    #[inline(always)]
+    fn narrow_i16(self) -> (int16x8_t, int16x8_t) {
+        let [a, b, c, d] = self.0;
+        // SAFETY: NEON baseline; register-only ops.
+        unsafe { (vcombine_s16(vqmovn_s32(a.0), vqmovn_s32(b.0)), vcombine_s16(vqmovn_s32(c.0), vqmovn_s32(d.0))) }
+    }
+
+    /// 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.
+    ///
+    /// ```
+    /// 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]
+    /// );
+    /// ```
+    #[inline(always)]
+    pub fn narrow_saturate_i8(self) -> I8x16 {
+        let (lo, hi) = self.narrow_i16();
+        // SAFETY: NEON baseline; register-only ops (`sqxtn`).
+        I8x16(unsafe { vcombine_s8(vqmovn_s16(lo), vqmovn_s16(hi)) })
+    }
+
+    /// 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) -> U8x16 {
+        // Saturating to i16 and then to u8 (`sqxtun`) is the same as clamping to 0..=255.
+        let (lo, hi) = self.narrow_i16();
+        // SAFETY: NEON baseline; register-only ops (`sqxtun`).
+        U8x16(unsafe { vcombine_u8(vqmovun_s16(lo), vqmovun_s16(hi)) })
+    }
+
     /// Broadcast (`vdupq_n_s32` ×4).
     #[inline(always)]
     pub fn splat(v: i32) -> Self {

src/simd_nightly/f32_types.rs

@@ -27,6 +27,14 @@ pub struct F32x16(pub core_f32x16);
 impl F32x16 {
     pub const LANES: usize = 16;
 
+    /// Truncating cast f32→i32 (equivalent to `portable_simd .cast::<i32>()`). Lanes outside
+    /// the i32 range and NaN follow `as i32` here (saturate, NaN to 0); the x86 backends give
+    /// `i32::MIN` for them instead, so only in-range lanes are identical across backends.
+    #[inline(always)]
+    pub fn cast_i32(self) -> I32x16 {
+        I32x16(self.0.cast::<i32>())
+    }
+
     // ── Constructors ──────────────────────────────────────────────
 
     /// Broadcast `v` to all 16 lanes.

src/simd_nightly/i_word_types.rs

@@ -315,6 +315,52 @@ pub struct I32x16(pub i32x16);
 impl I32x16 {
     pub const LANES: usize = 16;
 
+    /// 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.
+    ///
+    /// ```
+    /// 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]
+    /// );
+    /// ```
+    #[inline(always)]
+    pub fn narrow_saturate_i8(self) -> super::w1a_types::I8x16 {
+        super::w1a_types::I8x16(
+            self.0
+                .simd_clamp(i32x16::splat(-128), i32x16::splat(127))
+                .cast::<i8>(),
+        )
+    }
+
+    /// 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) -> super::w1a_types::U8x16 {
+        super::w1a_types::U8x16(
+            self.0
+                .simd_clamp(i32x16::splat(0), i32x16::splat(255))
+                .cast::<u8>(),
+        )
+    }
+
+    /// Cast i32→f32 (equivalent to `portable_simd .cast::<f32>()`).
+    #[inline(always)]
+    pub fn cast_f32(self) -> super::f32_types::F32x16 {
+        super::f32_types::F32x16(self.0.cast::<f32>())
+    }
+
     #[inline(always)]
     pub fn splat(v: i32) -> Self {
         Self(i32x16::splat(v))

src/simd_scalar.rs

@@ -942,6 +942,38 @@ impl U16x32 {
 
 // Extra methods for I32x16 that float types have via the macro
 impl I32x16 {
+    /// 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.
+    ///
+    /// ```
+    /// 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]
+    /// );
+    /// ```
+    #[inline(always)]
+    pub fn narrow_saturate_i8(self) -> I8x16 {
+        I8x16(core::array::from_fn(|i| self.0[i].clamp(-128, 127) as i8))
+    }
+
+    /// 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) -> U8x16 {
+        U8x16(core::array::from_fn(|i| self.0[i].clamp(0, 255) as u8))
+    }
+
     #[inline(always)]
     pub fn reduce_min(self) -> i32 {
         *self.0.iter().min().unwrap_or(&0)