Merge pull request #4764 from usamoi/pshufb

add avx512 pshufb

Author: RalfJung

src/tools/miri/src/shims/x86/avx2.rs

@@ -6,7 +6,7 @@ use rustc_target::callconv::FnAbi;
 
 use super::{
  ShiftOp, horizontal_bin_op, mpsadbw, packssdw, packsswb, packusdw, packuswb, permute, pmaddbw,
- pmulhrsw, psadbw, psign, shift_simd_by_scalar,
+ pmulhrsw, psadbw, pshufb, psign, shift_simd_by_scalar,
 };
 use crate::*;
 
@@ -189,28 +189,7 @@ pub(super) trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
  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!(dest_len, left_len);
- assert_eq!(dest_len, right_len);
-
- for i in 0..dest_len {
- let right = this.read_scalar(&this.project_index(&right, i)?)?.to_u8()?;
- let dest = this.project_index(&dest, i)?;
-
- let res = if right & 0x80 == 0 {
- // Shuffle each 128-bit (16-byte) block independently.
- let j = u64::from(right % 16).strict_add(i & !15);
- this.read_scalar(&this.project_index(&left, j)?)?
- } else {
- // If the highest bit in `right` is 1, write zero.
- Scalar::from_u8(0)
- };
-
- this.write_scalar(res, &dest)?;
- }
+ pshufb(this, left, right, dest)?;
  }
  // Used to implement the _mm256_sign_epi{8,16,32} functions.
  // Negates elements from `left` when the corresponding element in

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::{permute, pmaddbw, psadbw};
+use super::{permute, pmaddbw, psadbw, pshufb};
 use crate::*;
 
 impl<'tcx> EvalContextExt<'tcx> for crate::MiriInterpCx<'tcx> {}
@@ -102,6 +102,13 @@ pub(super) trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
 
  permute(this, left, right, dest)?;
  }
+ // Used to implement the _mm512_shuffle_epi8 intrinsic.
+ "pshuf.b.512" => {
+ let [left, right] =
+ this.check_shim_sig_lenient(abi, CanonAbi::C, link_name, args)?;
+
+ pshufb(this, left, right, dest)?;
+ }
  _ => return interp_ok(EmulateItemResult::NotSupported),
  }
  interp_ok(EmulateItemResult::NeedsReturn)

src/tools/miri/src/shims/x86/mod.rs

@@ -1155,6 +1155,51 @@ fn pclmulqdq<'tcx>(
  interp_ok(())
 }
 
+/// Shuffles bytes from `left` using `right` as pattern. Each 16-byte block is shuffled independently.
+///
+/// `left` and `right` are both vectors of type `len` x i8.
+///
+/// If the highest bit of a byte in `right` is not set, the corresponding byte in `dest` is taken
+/// from the current 16-byte block of `left` at the position indicated by the lowest 4 bits of this
+/// byte in `right`. If the highest bit of a byte in `right` is set, the corresponding byte in
+/// `dest` is set to `0`.
+///
+/// <https://www.intel.com/content/www/us/en/docs/intrinsics-guide/index.html#text=_mm_shuffle_epi8>
+/// <https://www.intel.com/content/www/us/en/docs/intrinsics-guide/index.html#text=_mm256_shuffle_epi8>
+/// <https://www.intel.com/content/www/us/en/docs/intrinsics-guide/index.html#text=_mm512_shuffle_epi8>
+fn pshufb<'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)?;
+
+ assert_eq!(dest_len, left_len);
+ assert_eq!(dest_len, right_len);
+
+ for i in 0..dest_len {
+ let right = ecx.read_scalar(&ecx.project_index(&right, i)?)?.to_u8()?;
+ let dest = ecx.project_index(&dest, i)?;
+
+ let res = if right & 0x80 == 0 {
+ // Shuffle each 128-bit (16-byte) block independently.
+ let block_offset = i & !15; // round down to previous multiple of 16
+ let j = block_offset.strict_add((right % 16).into());
+ ecx.read_scalar(&ecx.project_index(&left, j)?)?
+ } else {
+ // If the highest bit in `right` is 1, write zero.
+ Scalar::from_u8(0)
+ };
+
+ ecx.write_scalar(res, &dest)?;
+ }
+
+ interp_ok(())
+}
+
 /// Packs two N-bit integer vectors to a single N/2-bit integers.
 ///
 /// The conversion from N-bit to N/2-bit should be provided by `f`.

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, pmaddbw, pmulhrsw, psign};
+use super::{horizontal_bin_op, pmaddbw, pmulhrsw, pshufb, psign};
 use crate::*;
 
 impl<'tcx> EvalContextExt<'tcx> for crate::MiriInterpCx<'tcx> {}
@@ -29,27 +29,7 @@ pub(super) trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
  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!(dest_len, left_len);
- assert_eq!(dest_len, right_len);
-
- for i in 0..dest_len {
- let right = this.read_scalar(&this.project_index(&right, i)?)?.to_u8()?;
- let dest = this.project_index(&dest, i)?;
-
- let res = if right & 0x80 == 0 {
- let j = right % 16; // index wraps around
- this.read_scalar(&this.project_index(&left, j.into())?)?
- } else {
- // If the highest bit in `right` is 1, write zero.
- Scalar::from_u8(0)
- };
-
- this.write_scalar(res, &dest)?;
- }
+ pshufb(this, left, right, dest)?;
  }
  // Used to implement the _mm_h{adds,subs}_epi16 functions.
  // Horizontally add / subtract with saturation adjacent 16-bit

src/tools/miri/tests/pass/shims/x86/intrinsics-x86-avx512.rs

@@ -143,6 +143,30 @@ unsafe fn test_avx512() {
  assert_eq_m512i(r, e);
  }
  test_mm512_permutexvar_epi32();
+
+ #[target_feature(enable = "avx512bw")]
+ unsafe fn test_mm512_shuffle_epi8() {
+ #[rustfmt::skip]
+ let a = _mm512_set_epi8(0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15,
+ 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31,
+ 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47,
+ 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63);
+ #[rustfmt::skip]
+ let b = _mm512_set_epi8(-1, 127, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
+ -1, 127, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
+ -1, 127, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
+ -1, 127, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1);
+ let r = _mm512_shuffle_epi8(a, b);
+ // `_mm512_set_epi8` sets the bytes in inverse order (?!?), so the indices in `b` seem to
+ // index from the *back* of the corresponding 16-byte block in `a`.
+ #[rustfmt::skip]
+ let e = _mm512_set_epi8(0, 0, 14, 14, 14, 14, 14, 14, 14, 14, 14, 14, 14, 14, 14, 14,
+ 0, 16, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30,
+ 0, 32, 46, 46, 46, 46, 46, 46, 46, 46, 46, 46, 46, 46, 46, 46,
+ 0, 48, 62, 62, 62, 62, 62, 62, 62, 62, 62, 62, 62, 62, 62, 62);
+ assert_eq_m512i(r, e);
+ }
+ test_mm512_shuffle_epi8();
 }
 
 // Some of the constants in the tests below are just bit patterns. They should not