add shim for _mm512_maddubs_epi16

folkertdev · folkertdev · commit 349e2767a25a · 2025-11-17T22:31:20.000+01:00
diff --git a/src/tools/miri/src/shims/x86/avx2.rs b/src/tools/miri/src/shims/x86/avx2.rs
@@ -6,7 +6,7 @@ use rustc_target::callconv::FnAbi;
 
 use super::{
     ShiftOp, horizontal_bin_op, mask_load, mask_store, mpsadbw, packssdw, packsswb, packusdw,
-    packuswb, pmulhrsw, psadbw, psign, shift_simd_by_scalar, shift_simd_by_simd,
+    packuswb, pmaddbw, pmulhrsw, psadbw, psign, shift_simd_by_scalar, shift_simd_by_simd,
 };
 use crate::*;
 
@@ -102,39 +102,11 @@ pub(super) trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
                 }
             }
             // Used to implement the _mm256_maddubs_epi16 function.
-            // Multiplies packed 8-bit unsigned integers from `left` and packed
-            // signed 8-bit integers from `right` into 16-bit signed integers. Then,
-            // the saturating sum of the products with indices `2*i` and `2*i+1`
-            // produces the output at index `i`.
             "pmadd.ub.sw" => {
                 let [left, right] =
                     this.check_shim_sig_lenient(abi, CanonAbi::C, link_name, args)?;
 
-                let (left, left_len) = this.project_to_simd(left)?;
-                let (right, right_len) = this.project_to_simd(right)?;
-                let (dest, dest_len) = this.project_to_simd(dest)?;
-
-                assert_eq!(left_len, right_len);
-                assert_eq!(dest_len.strict_mul(2), left_len);
-
-                for i in 0..dest_len {
-                    let j1 = i.strict_mul(2);
-                    let left1 = this.read_scalar(&this.project_index(&left, j1)?)?.to_u8()?;
-                    let right1 = this.read_scalar(&this.project_index(&right, j1)?)?.to_i8()?;
-
-                    let j2 = j1.strict_add(1);
-                    let left2 = this.read_scalar(&this.project_index(&left, j2)?)?.to_u8()?;
-                    let right2 = this.read_scalar(&this.project_index(&right, j2)?)?.to_i8()?;
-
-                    let dest = this.project_index(&dest, i)?;
-
-                    // Multiplication of a u8 and an i8 into an i16 cannot overflow.
-                    let mul1 = i16::from(left1).strict_mul(right1.into());
-                    let mul2 = i16::from(left2).strict_mul(right2.into());
-                    let res = mul1.saturating_add(mul2);
-
-                    this.write_scalar(Scalar::from_i16(res), &dest)?;
-                }
+                pmaddbw(this, left, right, dest)?;
             }
             // Used to implement the _mm_maskload_epi32, _mm_maskload_epi64,
             // _mm256_maskload_epi32 and _mm256_maskload_epi64 functions.
diff --git a/src/tools/miri/src/shims/x86/avx512.rs b/src/tools/miri/src/shims/x86/avx512.rs
@@ -3,7 +3,7 @@ use rustc_middle::ty::Ty;
 use rustc_span::Symbol;
 use rustc_target::callconv::FnAbi;
 
-use super::psadbw;
+use super::{pmaddbw, psadbw};
 use crate::*;
 
 impl<'tcx> EvalContextExt<'tcx> for crate::MiriInterpCx<'tcx> {}
@@ -88,6 +88,13 @@ pub(super) trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
 
                 psadbw(this, left, right, dest)?
             }
+            // Used to implement the _mm512_maddubs_epi16 function.
+            "pmaddubs.w.512" => {
+                let [left, right] =
+                    this.check_shim_sig_lenient(abi, CanonAbi::C, link_name, args)?;
+
+                pmaddbw(this, left, right, dest)?;
+            }
             _ => return interp_ok(EmulateItemResult::NotSupported),
         }
         interp_ok(EmulateItemResult::NeedsReturn)
diff --git a/src/tools/miri/src/shims/x86/mod.rs b/src/tools/miri/src/shims/x86/mod.rs
@@ -1086,6 +1086,52 @@ fn psadbw<'tcx>(
     interp_ok(())
 }
 
+/// Multiplies packed 8-bit unsigned integers from `left` and packed
+/// signed 8-bit integers from `right` into 16-bit signed integers. Then,
+/// the saturating sum of the products with indices `2*i` and `2*i+1`
+/// produces the output at index `i`.
+///
+/// <https://www.intel.com/content/www/us/en/docs/intrinsics-guide/index.html#text=_mm_maddubs_epi16>
+/// <https://www.intel.com/content/www/us/en/docs/intrinsics-guide/index.html#text=_mm256_maddubs_epi16>
+/// <https://www.intel.com/content/www/us/en/docs/intrinsics-guide/index.html#text=_mm512_maddubs_epi16>
+fn pmaddbw<'tcx>(
+    ecx: &mut crate::MiriInterpCx<'tcx>,
+    left: &OpTy<'tcx>,
+    right: &OpTy<'tcx>,
+    dest: &MPlaceTy<'tcx>,
+) -> InterpResult<'tcx, ()> {
+    let (left, left_len) = ecx.project_to_simd(left)?;
+    let (right, right_len) = ecx.project_to_simd(right)?;
+    let (dest, dest_len) = ecx.project_to_simd(dest)?;
+
+    // fn pmaddubsw128(a: u8x16, b: i8x16) -> i16x8;
+    // fn pmaddubsw(   a: u8x32, b: i8x32) -> i16x16;
+    // fn vpmaddubsw(  a: u8x64, b: i8x64) -> i16x32;
+    assert_eq!(left_len, right_len);
+    assert_eq!(dest_len.strict_mul(2), left_len);
+
+    for i in 0..dest_len {
+        let j1 = i.strict_mul(2);
+        let left1 = ecx.read_scalar(&ecx.project_index(&left, j1)?)?.to_u8()?;
+        let right1 = ecx.read_scalar(&ecx.project_index(&right, j1)?)?.to_i8()?;
+
+        let j2 = j1.strict_add(1);
+        let left2 = ecx.read_scalar(&ecx.project_index(&left, j2)?)?.to_u8()?;
+        let right2 = ecx.read_scalar(&ecx.project_index(&right, j2)?)?.to_i8()?;
+
+        let dest = ecx.project_index(&dest, i)?;
+
+        // Multiplication of a u8 and an i8 into an i16 cannot overflow.
+        let mul1 = i16::from(left1).strict_mul(right1.into());
+        let mul2 = i16::from(left2).strict_mul(right2.into());
+        let res = mul1.saturating_add(mul2);
+
+        ecx.write_scalar(Scalar::from_i16(res), &dest)?;
+    }
+
+    interp_ok(())
+}
+
 /// Multiplies packed 16-bit signed integer values, truncates the 32-bit
 /// product to the 18 most significant bits by right-shifting, and then
 /// divides the 18-bit value by 2 (rounding to nearest) by first adding
diff --git a/src/tools/miri/src/shims/x86/ssse3.rs b/src/tools/miri/src/shims/x86/ssse3.rs
@@ -4,7 +4,7 @@ use rustc_middle::ty::Ty;
 use rustc_span::Symbol;
 use rustc_target::callconv::FnAbi;
 
-use super::{horizontal_bin_op, pmulhrsw, psign};
+use super::{horizontal_bin_op, pmaddbw, pmulhrsw, psign};
 use crate::*;
 
 impl<'tcx> EvalContextExt<'tcx> for crate::MiriInterpCx<'tcx> {}
@@ -67,40 +67,11 @@ pub(super) trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
                 horizontal_bin_op(this, which, /*saturating*/ true, left, right, dest)?;
             }
             // Used to implement the _mm_maddubs_epi16 function.
-            // Multiplies packed 8-bit unsigned integers from `left` and packed
-            // signed 8-bit integers from `right` into 16-bit signed integers. Then,
-            // the saturating sum of the products with indices `2*i` and `2*i+1`
-            // produces the output at index `i`.
-            // https://www.intel.com/content/www/us/en/docs/intrinsics-guide/index.html#text=_mm_maddubs_epi16
             "pmadd.ub.sw.128" => {
                 let [left, right] =
                     this.check_shim_sig_lenient(abi, CanonAbi::C, link_name, args)?;
 
-                let (left, left_len) = this.project_to_simd(left)?;
-                let (right, right_len) = this.project_to_simd(right)?;
-                let (dest, dest_len) = this.project_to_simd(dest)?;
-
-                assert_eq!(left_len, right_len);
-                assert_eq!(dest_len.strict_mul(2), left_len);
-
-                for i in 0..dest_len {
-                    let j1 = i.strict_mul(2);
-                    let left1 = this.read_scalar(&this.project_index(&left, j1)?)?.to_u8()?;
-                    let right1 = this.read_scalar(&this.project_index(&right, j1)?)?.to_i8()?;
-
-                    let j2 = j1.strict_add(1);
-                    let left2 = this.read_scalar(&this.project_index(&left, j2)?)?.to_u8()?;
-                    let right2 = this.read_scalar(&this.project_index(&right, j2)?)?.to_i8()?;
-
-                    let dest = this.project_index(&dest, i)?;
-
-                    // Multiplication of a u8 and an i8 into an i16 cannot overflow.
-                    let mul1 = i16::from(left1).strict_mul(right1.into());
-                    let mul2 = i16::from(left2).strict_mul(right2.into());
-                    let res = mul1.saturating_add(mul2);
-
-                    this.write_scalar(Scalar::from_i16(res), &dest)?;
-                }
+                pmaddbw(this, left, right, dest)?;
             }
             // Used to implement the _mm_mulhrs_epi16 function.
             // Multiplies packed 16-bit signed integer values, truncates the 32-bit
diff --git a/src/tools/miri/tests/pass/shims/x86/intrinsics-x86-avx512.rs b/src/tools/miri/tests/pass/shims/x86/intrinsics-x86-avx512.rs
@@ -55,6 +55,48 @@ unsafe fn test_avx512() {
         assert_eq_m512i(r, e);
     }
     test_mm512_sad_epu8();
+
+    #[target_feature(enable = "avx512bw")]
+    unsafe fn test_mm512_maddubs_epi16() {
+        // `a` is interpreted as `u8x16`, but `_mm512_set_epi8` expects `i8`, so we have to cast.
+        #[rustfmt::skip]
+        let a = _mm512_set_epi8(
+            255u8 as i8, 255u8 as i8,  60,  50, 100, 100, 255u8 as i8, 200u8 as i8,
+            255u8 as i8, 200u8 as i8, 200u8 as i8, 100,  60,  50,  20,  10,
+
+            255u8 as i8, 255u8 as i8,  60,  50, 100, 100, 255u8 as i8, 200u8 as i8,
+            255u8 as i8, 200u8 as i8, 200u8 as i8, 100,  60,  50,  20,  10,
+
+            255u8 as i8, 255u8 as i8,  60,  50, 100, 100, 255u8 as i8, 200u8 as i8,
+            255u8 as i8, 200u8 as i8, 200u8 as i8, 100,  60,  50,  20,  10,
+
+            255u8 as i8, 255u8 as i8,  60,  50, 100, 100, 255u8 as i8, 200u8 as i8,
+            255u8 as i8, 200u8 as i8, 200u8 as i8, 100,  60,  50,  20,  10,
+        );
+
+        let b = _mm512_set_epi8(
+            64, 64, -2, 1, 100, 100, -128, -128, //
+            127, 127, -1, 1, 2, 2, 1, 1, //
+            64, 64, -2, 1, 100, 100, -128, -128, //
+            127, 127, -1, 1, 2, 2, 1, 1, //
+            64, 64, -2, 1, 100, 100, -128, -128, //
+            127, 127, -1, 1, 2, 2, 1, 1, //
+            64, 64, -2, 1, 100, 100, -128, -128, //
+            127, 127, -1, 1, 2, 2, 1, 1, //
+        );
+
+        let r = _mm512_maddubs_epi16(a, b);
+
+        let e = _mm512_set_epi16(
+            32640, -70, 20000, -32768, 32767, -100, 220, 30, //
+            32640, -70, 20000, -32768, 32767, -100, 220, 30, //
+            32640, -70, 20000, -32768, 32767, -100, 220, 30, //
+            32640, -70, 20000, -32768, 32767, -100, 220, 30, //
+        );
+
+        assert_eq_m512i(r, e);
+    }
+    test_mm512_maddubs_epi16();
 }
 
 // Some of the constants in the tests below are just bit patterns. They should not
