Auto merge of #3204 - RalfJung:simd, r=RalfJung

add new SIMD intrinsics

Author: bors

src/tools/miri/src/shims/intrinsics/simd.rs

@@ -1,6 +1,7 @@
 use rustc_apfloat::{Float, Round};
 use rustc_middle::ty::layout::{HasParamEnv, LayoutOf};
 use rustc_middle::{mir, ty, ty::FloatTy};
+use rustc_span::{sym, Symbol};
 use rustc_target::abi::{Endian, HasDataLayout};
 
 use crate::*;
@@ -25,7 +26,12 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
             | "floor"
             | "round"
             | "trunc"
-            | "fsqrt" => {
+            | "fsqrt"
+            | "ctlz"
+            | "cttz"
+            | "bswap"
+            | "bitreverse"
+            => {
                 let [op] = check_arg_count(args)?;
                 let (op, op_len) = this.operand_to_simd(op)?;
                 let (dest, dest_len) = this.place_to_simd(dest)?;
@@ -38,6 +44,7 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
                     Abs,
                     Sqrt,
                     Round(rustc_apfloat::Round),
+                    Numeric(Symbol),
                 }
                 let which = match intrinsic_name {
                     "neg" => Op::MirOp(mir::UnOp::Neg),
@@ -47,6 +54,10 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
                     "floor" => Op::Round(rustc_apfloat::Round::TowardNegative),
                     "round" => Op::Round(rustc_apfloat::Round::NearestTiesToAway),
                     "trunc" => Op::Round(rustc_apfloat::Round::TowardZero),
+                    "ctlz" => Op::Numeric(sym::ctlz),
+                    "cttz" => Op::Numeric(sym::cttz),
+                    "bswap" => Op::Numeric(sym::bswap),
+                    "bitreverse" => Op::Numeric(sym::bitreverse),
                     _ => unreachable!(),
                 };
 
@@ -101,6 +112,20 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
                                 }
                             }
                         }
+                        Op::Numeric(name) => {
+                            assert!(op.layout.ty.is_integral());
+                            let size = op.layout.size;
+                            let bits = op.to_scalar().to_bits(size).unwrap();
+                            let extra = 128u128.checked_sub(u128::from(size.bits())).unwrap();
+                            let bits_out = match name {
+                                sym::ctlz => u128::from(bits.leading_zeros()).checked_sub(extra).unwrap(),
+                                sym::cttz => u128::from((bits << extra).trailing_zeros()).checked_sub(extra).unwrap(),
+                                sym::bswap => (bits << extra).swap_bytes(),
+                                sym::bitreverse => (bits << extra).reverse_bits(),
+                                _ => unreachable!(),
+                            };
+                            Scalar::from_uint(bits_out, size)
+                        }
                     };
                     this.write_scalar(val, &dest)?;
                 }
@@ -126,7 +151,8 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
             | "fmin"
             | "saturating_add"
             | "saturating_sub"
-            | "arith_offset" => {
+            | "arith_offset"
+            => {
                 use mir::BinOp;
 
                 let [left, right] = check_arg_count(args)?;
@@ -386,16 +412,25 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
                 let (dest, dest_len) = this.place_to_simd(dest)?;
                 let bitmask_len = dest_len.max(8);
 
-                assert!(mask.layout.ty.is_integral());
                 assert!(bitmask_len <= 64);
                 assert_eq!(bitmask_len, mask.layout.size.bits());
                 assert_eq!(dest_len, yes_len);
                 assert_eq!(dest_len, no_len);
                 let dest_len = u32::try_from(dest_len).unwrap();
                 let bitmask_len = u32::try_from(bitmask_len).unwrap();
 
-                let mask: u64 =
-                    this.read_scalar(mask)?.to_bits(mask.layout.size)?.try_into().unwrap();
+                // The mask can be a single integer or an array.
+                let mask: u64 = match mask.layout.ty.kind() {
+                    ty::Int(..) | ty::Uint(..) =>
+                        this.read_scalar(mask)?.to_bits(mask.layout.size)?.try_into().unwrap(),
+                    ty::Array(elem, _) if matches!(elem.kind(), ty::Uint(ty::UintTy::U8)) => {
+                        let mask_ty = this.machine.layouts.uint(mask.layout.size).unwrap();
+                        let mask = mask.transmute(mask_ty, this)?;
+                        this.read_scalar(&mask)?.to_bits(mask_ty.size)?.try_into().unwrap()
+                    }
+                    _ => bug!("simd_select_bitmask: invalid mask type {}", mask.layout.ty),
+                };
+
                 for i in 0..dest_len {
                     let mask = mask
                         & 1u64

src/tools/miri/tests/pass/portable-simd.rs

@@ -197,6 +197,24 @@ fn simd_ops_i32() {
     assert_eq!(b.reduce_or(), -1);
     assert_eq!(a.reduce_xor(), 0);
     assert_eq!(b.reduce_xor(), -4);
+
+    assert_eq!(b.leading_zeros(), u32x4::from_array([31, 30, 30, 0]));
+    assert_eq!(b.trailing_zeros(), u32x4::from_array([0, 1, 0, 2]));
+    assert_eq!(b.leading_ones(), u32x4::from_array([0, 0, 0, 30]));
+    assert_eq!(b.trailing_ones(), u32x4::from_array([1, 0, 2, 0]));
+    assert_eq!(
+        b.swap_bytes(),
+        i32x4::from_array([0x01000000, 0x02000000, 0x03000000, 0xfcffffffu32 as i32])
+    );
+    assert_eq!(
+        b.reverse_bits(),
+        i32x4::from_array([
+            0x80000000u32 as i32,
+            0x40000000,
+            0xc0000000u32 as i32,
+            0x3fffffffu32 as i32
+        ])
+    );
 }
 
 fn simd_mask() {
@@ -247,6 +265,22 @@ fn simd_mask() {
             assert_eq!(bitmask2, [0b0001]);
         }
     }
+
+    // This used to cause an ICE.
+    let bitmask = u8x8::from_array([0b01000101, 0, 0, 0, 0, 0, 0, 0]);
+    assert_eq!(
+        mask32x8::from_bitmask_vector(bitmask),
+        mask32x8::from_array([true, false, true, false, false, false, true, false]),
+    );
+    let bitmask =
+        u8x16::from_array([0b01000101, 0b11110000, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]);
+    assert_eq!(
+        mask32x16::from_bitmask_vector(bitmask),
+        mask32x16::from_array([
+            true, false, true, false, false, false, true, false, false, false, false, false, true,
+            true, true, true,
+        ]),
+    );
 }
 
 fn simd_cast() {