diff --git a/rust/lance-linalg/src/simd/i32.rs b/rust/lance-linalg/src/simd/i32.rs index 4cfaef0bf72..6e0812928db 100644 --- a/rust/lance-linalg/src/simd/i32.rs +++ b/rust/lance-linalg/src/simd/i32.rs @@ -15,16 +15,24 @@ use std::mem::transmute; use super::SIMD; +/// 8 of 32-bit `i32` values. Use 256-bit SIMD if possible. +/// +/// The x86_64 arm reaches AVX and AVX2 intrinsics with no `#[target_feature]` +/// gate of its own, so callers must already be inside an AVX2-checked context. +/// `x86_64-unknown-linux-gnu` is pinned to `target-cpu=x86-64-v2` +/// (`.cargo/config.toml`), which is below AVX. #[allow(non_camel_case_types)] #[cfg(target_arch = "x86_64")] #[derive(Clone, Copy)] pub struct i32x8(pub(crate) __m256i); +/// 8 of 32-bit `i32` values. Use 256-bit SIMD if possible. #[allow(non_camel_case_types)] #[cfg(target_arch = "aarch64")] #[derive(Clone, Copy)] pub struct i32x8(int32x4x2_t); +/// 8 of 32-bit `i32` values. Use 256-bit SIMD if possible. #[allow(non_camel_case_types)] #[cfg(target_arch = "loongarch64")] #[derive(Clone, Copy)] @@ -302,11 +310,24 @@ impl SubAssign for i32x8 { impl Mul for i32x8 { type Output = Self; + /// Lane-wise product, keeping the low 32 bits of each result. + /// + /// `mul` wraps on overflow rather than panicking the way scalar `i32 * i32` + /// does in a debug build, and all three arms agree on that: `vpmulld`, + /// `vmulq_s32` and `lasx_xvmul_w` each discard the high half. This is a + /// statement about `mul` alone — `reduce_sum` sums in scalar `i32` on x86_64 + /// and loongarch64 (so it panics on overflow in a debug build) but reduces + /// in-register on aarch64, where it wraps. + /// + /// Picking a widening variant here is a silent wrong answer, not a compile + /// error: `_mm256_mul_epi32` (`vpmuldq`) multiplies only the even 32-bit + /// lanes and writes four 64-bit results, so `[1, 2, ..., 8]` squared came + /// back as `[1, 0, 9, 0, 25, 0, 49, 0]`. #[inline] fn mul(self, rhs: Self) -> Self::Output { #[cfg(target_arch = "x86_64")] unsafe { - Self(_mm256_mul_epi32(self.0, rhs.0)) + Self(_mm256_mullo_epi32(self.0, rhs.0)) } #[cfg(target_arch = "aarch64")] unsafe { @@ -325,9 +346,35 @@ impl Mul for i32x8 { #[cfg(test)] mod tests { use super::*; + use rstest::rstest; #[test] fn test_slice_conversion_rejects_short_input() { assert!(std::panic::catch_unwind(|| i32x8::from(&[0; 7][..])).is_err()); } + + /// Lane-wise, low-32-bits multiplication is what all three arms promise, so + /// this runs everywhere: only the x86 feature check is arch-gated, matching + /// `f32.rs`'s and `f64.rs`'s test modules. + /// + /// Every case below has to produce a different answer under the widening + /// `vpmuldq` this file used to call. All-zero *inputs* would not: `vpmuldq` + /// returns zeros for those too. + #[rstest] + #[case::squares([1, 2, 3, 4, 5, 6, 7, 8], [1, 2, 3, 4, 5, 6, 7, 8], [1, 4, 9, 16, 25, 36, 49, 64])] + #[case::mixed_signs([-3, 7, -3, 7, -3, 7, -3, 7], [7, -3, 7, -3, 7, -3, 7, -3], [-21; 8])] + #[case::wraps_to_low_32_bits([65536; 8], [65536; 8], [0; 8])] + fn mul_is_lane_wise(#[case] lhs: [i32; 8], #[case] rhs: [i32; 8], #[case] expected: [i32; 8]) { + // `load_unaligned` / `store_unaligned` are AVX and `mul` is AVX2, and + // none of them is `#[target_feature]`-gated, so a pre-Haswell host would + // SIGILL. The `qemu-pre-haswell` CI job runs exactly that. + #[cfg(target_arch = "x86_64")] + if !std::is_x86_feature_detected!("avx2") { + return; + } + + let product = i32x8::from(&lhs) * i32x8::from(&rhs); + + assert_eq!(product.as_array(), expected); + } }