Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
AdaWorldAPI merged 1 commit into from
Oct 24, 2026
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
135 changes: 85 additions & 50 deletions crates/simd-masking-parity/src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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`].

Expand All @@ -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).
Expand All @@ -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() {
Expand Down Expand Up @@ -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(())
Expand Down
4 changes: 4 additions & 0 deletions src/simd.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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};
Expand Down
137 changes: 84 additions & 53 deletions src/simd_avx2.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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])
Expand Down Expand Up @@ -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.
Expand Down
Loading
Loading