fix(linalg): use a lane-wise intrinsic for i32x8 multiplication on x86

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);
+    }
 }