Add intrinsic to create an integer bitmask from the MSB of integer vectors

Author: gnzlbg

src/librustc_codegen_llvm/intrinsic.rs

@@ -1303,6 +1303,52 @@ fn generic_simd_intrinsic(
         return Ok(bx.select(m_i1s, args[1].immediate(), args[2].immediate()));
     }
 
+    if name == "simd_bitmask" {
+        // The `fn simd_bitmask(vector) -> unsigned integer` intrinsic takes a
+        // vector mask and returns an unsigned integer containing the most
+        // significant bit (MSB) of each lane.
+        use rustc_target::abi::HasDataLayout;
+
+        // If the vector has less than 8 lanes, an u8 is returned with zeroed
+        // trailing bits.
+        let expected_int_bits = in_len.max(8);
+        match ret_ty.sty {
+           ty::Uint(i) if i.bit_width() == Some(expected_int_bits) => (),
+            _ => return_error!(
+                "bitmask `{}`, expected `u{}`",
+                ret_ty, expected_int_bits
+            ),
+        }
+
+        // Integer vector <i{in_bitwidth} x in_len>:
+        let (i_xn, in_elem_bitwidth) = match in_elem.sty {
+            ty::Int(i) => (
+                args[0].immediate(),
+                i.bit_width().unwrap_or(bx.data_layout().pointer_size.bits() as _)
+            ),
+            ty::Uint(i) => (
+                args[0].immediate(),
+                i.bit_width().unwrap_or(bx.data_layout().pointer_size.bits() as _)
+            ),
+            _ => return_error!(
+                "vector argument `{}`'s element type `{}`, expected integer element type",
+                in_ty, in_elem
+            ),
+        };
+
+        // Shift the MSB to the right by "in_elem_bitwidth - 1" into the first bit position.
+        let shift_indices = vec![
+            bx.cx.const_int(bx.type_ix(in_elem_bitwidth as _), (in_elem_bitwidth - 1) as _); in_len
+        ];
+        let i_xn_msb = bx.lshr(i_xn, bx.const_vector(shift_indices.as_slice()));
+        // Truncate vector to an <i1 x N>
+        let i1xn = bx.trunc(i_xn_msb, bx.type_vector(bx.type_i1(), in_len as _));
+        // Bitcast <i1 x N> to iN:
+        let i_ = bx.bitcast(i1xn, bx.type_ix(in_len as _));
+        // Zero-extend iN to the bitmask type:
+        return Ok(bx.zext(i_, bx.type_ix(expected_int_bits as _)));
+    }
+
     fn simd_simple_float_intrinsic(
         name: &str,
         in_elem: &::rustc::ty::TyS,

src/librustc_typeck/check/intrinsic.rs

@@ -425,6 +425,7 @@ pub fn check_platform_intrinsic_type<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>,
         "simd_insert" => (2, vec![param(0), tcx.types.u32, param(1)], param(0)),
         "simd_extract" => (2, vec![param(0), tcx.types.u32], param(1)),
         "simd_cast" => (2, vec![param(0)], param(1)),
+        "simd_bitmask" => (2, vec![param(0)], param(1)),
         "simd_select" |
         "simd_select_bitmask" => (2, vec![param(0), param(1), param(1)], param(1)),
         "simd_reduce_all" | "simd_reduce_any" => (1, vec![param(0)], tcx.types.bool),

src/test/codegen/simd-intrinsic-generic-bitmask.rs

@@ -0,0 +1,57 @@
+// compile-flags: -C no-prepopulate-passes
+// ignore-tidy-linelength
+
+#![crate_type = "lib"]
+
+#![feature(repr_simd, platform_intrinsics)]
+#![allow(non_camel_case_types)]
+
+#[repr(simd)]
+#[derive(Copy, Clone)]
+pub struct u32x2(u32, u32);
+
+#[repr(simd)]
+#[derive(Copy, Clone)]
+pub struct i32x2(i32, i32);
+
+#[repr(simd)]
+#[derive(Copy, Clone)]
+pub struct i8x16(
+    i8, i8, i8, i8, i8, i8, i8, i8,
+    i8, i8, i8, i8, i8, i8, i8, i8,
+);
+
+
+extern "platform-intrinsic" {
+    fn simd_bitmask<T, U>(x: T) -> U;
+}
+
+// CHECK-LABEL: @bitmask_int
+#[no_mangle]
+pub unsafe fn bitmask_int(x: i32x2) -> u8 {
+    // CHECK: [[A:%[0-9]+]] = lshr <2 x i32> %{{[0-9]+}}, <i32 31, i32 31>
+    // CHECK: [[B:%[0-9]+]] = trunc <2 x i32> [[A]] to <2 x i1>
+    // CHECK: [[C:%[0-9]+]] = bitcast <2 x i1> [[B]] to i2
+    // CHECK: %{{[0-9]+}} = zext i2 [[C]] to i8
+    simd_bitmask(x)
+}
+
+// CHECK-LABEL: @bitmask_uint
+#[no_mangle]
+pub unsafe fn bitmask_uint(x: u32x2) -> u8 {
+    // CHECK: [[A:%[0-9]+]] = lshr <2 x i32> %{{[0-9]+}}, <i32 31, i32 31>
+    // CHECK: [[B:%[0-9]+]] = trunc <2 x i32> [[A]] to <2 x i1>
+    // CHECK: [[C:%[0-9]+]] = bitcast <2 x i1> [[B]] to i2
+    // CHECK: %{{[0-9]+}} = zext i2 [[C]] to i8
+    simd_bitmask(x)
+}
+
+// CHECK-LABEL: @bitmask_int16
+#[no_mangle]
+pub unsafe fn bitmask_int16(x: i8x16) -> u16 {
+    // CHECK: [[A:%[0-9]+]] = lshr <16 x i8> %{{[0-9]+}}, <i8 7, i8 7, i8 7, i8 7, i8 7, i8 7, i8 7, i8 7, i8 7, i8 7, i8 7, i8 7, i8 7, i8 7, i8 7, i8 7>
+    // CHECK: [[B:%[0-9]+]] = trunc <16 x i8> [[A]] to <16 x i1>
+    // CHECK: %{{[0-9]+}} = bitcast <16 x i1> [[B]] to i16
+    // CHECK-NOT: zext
+    simd_bitmask(x)
+}

src/test/run-pass/simd/simd-intrinsic-generic-bitmask.rs

@@ -0,0 +1,61 @@
+// run-pass
+#![allow(non_camel_case_types)]
+
+// ignore-emscripten
+
+// Test that the simd_bitmask intrinsic produces correct results.
+
+#![feature(repr_simd, platform_intrinsics)]
+#[allow(non_camel_case_types)]
+
+#[repr(simd)]
+#[derive(Copy, Clone, PartialEq, Debug)]
+struct u32x4(pub u32, pub u32, pub u32, pub u32);
+
+#[repr(simd)]
+#[derive(Copy, Clone, PartialEq, Debug)]
+struct u8x4(pub u8, pub u8, pub u8, pub u8);
+
+#[repr(simd)]
+#[derive(Copy, Clone, PartialEq, Debug)]
+struct Tx4<T>(pub T, pub T, pub T, pub T);
+
+extern "platform-intrinsic" {
+    fn simd_bitmask<T, U>(x: T) -> U;
+}
+
+fn main() {
+    let z = u32x4(0, 0, 0, 0);
+    let ez = 0_u8;
+
+    let o = u32x4(!0, !0, !0, !0);
+    let eo = 0b_1111_u8;
+
+    let m0 = u32x4(!0, 0, !0, 0);
+    let e0 = 0b_0000_0101_u8;
+
+    // Check that the MSB is extracted:
+    let m = u8x4(0b_1000_0000, 0b_0100_0001, 0b_1100_0001, 0b_1111_1111);
+    let e = 0b_1101;
+
+    // Check usize / isize
+    let msize: Tx4<usize> = Tx4(usize::max_value(), 0, usize::max_value(), usize::max_value());
+
+    unsafe {
+        let r: u8 = simd_bitmask(z);
+        assert_eq!(r, ez);
+
+        let r: u8 = simd_bitmask(o);
+        assert_eq!(r, eo);
+
+        let r: u8 = simd_bitmask(m0);
+        assert_eq!(r, e0);
+
+        let r: u8 = simd_bitmask(m);
+        assert_eq!(r, e);
+
+        let r: u8 = simd_bitmask(msize);
+        assert_eq!(r, e);
+
+    }
+}

src/test/ui/simd-intrinsic/simd-intrinsic-generic-bitmask.rs

@@ -0,0 +1,90 @@
+// Test that the simd_bitmask intrinsic produces ok-ish error
+// messages when misused.
+
+#![feature(repr_simd, platform_intrinsics)]
+#![allow(non_camel_case_types)]
+
+#[repr(simd)]
+#[derive(Copy, Clone)]
+pub struct u32x2(pub u32, pub u32);
+
+#[repr(simd)]
+#[derive(Copy, Clone)]
+pub struct u32x4(pub u32, pub u32, pub u32, pub u32);
+
+#[repr(simd)]
+#[derive(Copy, Clone)]
+struct u8x8(
+    pub u8, pub u8, pub u8, pub u8, pub u8, pub u8, pub u8, pub u8,
+);
+
+#[repr(simd)]
+#[derive(Copy, Clone)]
+struct u8x16(
+    pub u8, pub u8, pub u8, pub u8, pub u8, pub u8, pub u8, pub u8,
+    pub u8, pub u8, pub u8, pub u8, pub u8, pub u8, pub u8, pub u8,
+);
+
+#[repr(simd)]
+#[derive(Copy, Clone)]
+struct u8x32(
+    pub u8, pub u8, pub u8, pub u8, pub u8, pub u8, pub u8, pub u8,
+    pub u8, pub u8, pub u8, pub u8, pub u8, pub u8, pub u8, pub u8,
+    pub u8, pub u8, pub u8, pub u8, pub u8, pub u8, pub u8, pub u8,
+    pub u8, pub u8, pub u8, pub u8, pub u8, pub u8, pub u8, pub u8,
+);
+
+#[repr(simd)]
+#[derive(Copy, Clone)]
+struct u8x64(
+    pub u8, pub u8, pub u8, pub u8, pub u8, pub u8, pub u8, pub u8,
+    pub u8, pub u8, pub u8, pub u8, pub u8, pub u8, pub u8, pub u8,
+    pub u8, pub u8, pub u8, pub u8, pub u8, pub u8, pub u8, pub u8,
+    pub u8, pub u8, pub u8, pub u8, pub u8, pub u8, pub u8, pub u8,
+    pub u8, pub u8, pub u8, pub u8, pub u8, pub u8, pub u8, pub u8,
+    pub u8, pub u8, pub u8, pub u8, pub u8, pub u8, pub u8, pub u8,
+    pub u8, pub u8, pub u8, pub u8, pub u8, pub u8, pub u8, pub u8,
+    pub u8, pub u8, pub u8, pub u8, pub u8, pub u8, pub u8, pub u8,
+);
+
+extern "platform-intrinsic" {
+    fn simd_bitmask<T, U>(x: T) -> U;
+}
+
+fn main() {
+    let m2 = u32x2(0, 0);
+    let m4 = u32x4(0, 0, 0, 0);
+    let m8 = u8x8(0, 0, 0, 0, 0, 0, 0, 0);
+    let m16 = u8x16(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0);
+    let m32 = u8x32(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+                    0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0);
+    let m64 = u8x64(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+                    0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+                    0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+                    0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0);
+
+    unsafe {
+        let _: u8 = simd_bitmask(m2);
+        let _: u8 = simd_bitmask(m4);
+        let _: u8 = simd_bitmask(m8);
+        let _: u16 = simd_bitmask(m16);
+        let _: u32 = simd_bitmask(m32);
+        let _: u64 = simd_bitmask(m64);
+
+        let _: u16 = simd_bitmask(m2);
+        //~^ ERROR bitmask `u16`, expected `u8`
+
+        let _: u16 = simd_bitmask(m8);
+        //~^ ERROR bitmask `u16`, expected `u8`
+
+        let _: u32 = simd_bitmask(m16);
+        //~^ ERROR bitmask `u32`, expected `u16`
+
+        let _: u64 = simd_bitmask(m32);
+        //~^ ERROR bitmask `u64`, expected `u32`
+
+        let _: u128 = simd_bitmask(m64);
+        //~^ ERROR bitmask `u128`, expected `u64`
+
+   }
+}

src/test/ui/simd-intrinsic/simd-intrinsic-generic-bitmask.stderr

@@ -0,0 +1,33 @@
+error[E0511]: invalid monomorphization of `simd_bitmask` intrinsic: bitmask `u16`, expected `u8`
+  --> $DIR/simd-intrinsic-generic-bitmask.rs:74:22
+   |
+LL |         let _: u16 = simd_bitmask(m2);
+   |                      ^^^^^^^^^^^^^^^^
+
+error[E0511]: invalid monomorphization of `simd_bitmask` intrinsic: bitmask `u16`, expected `u8`
+  --> $DIR/simd-intrinsic-generic-bitmask.rs:77:22
+   |
+LL |         let _: u16 = simd_bitmask(m8);
+   |                      ^^^^^^^^^^^^^^^^
+
+error[E0511]: invalid monomorphization of `simd_bitmask` intrinsic: bitmask `u32`, expected `u16`
+  --> $DIR/simd-intrinsic-generic-bitmask.rs:80:22
+   |
+LL |         let _: u32 = simd_bitmask(m16);
+   |                      ^^^^^^^^^^^^^^^^^
+
+error[E0511]: invalid monomorphization of `simd_bitmask` intrinsic: bitmask `u64`, expected `u32`
+  --> $DIR/simd-intrinsic-generic-bitmask.rs:83:22
+   |
+LL |         let _: u64 = simd_bitmask(m32);
+   |                      ^^^^^^^^^^^^^^^^^
+
+error[E0511]: invalid monomorphization of `simd_bitmask` intrinsic: bitmask `u128`, expected `u64`
+  --> $DIR/simd-intrinsic-generic-bitmask.rs:86:23
+   |
+LL |         let _: u128 = simd_bitmask(m64);
+   |                       ^^^^^^^^^^^^^^^^^
+
+error: aborting due to 5 previous errors
+
+For more information about this error, try `rustc --explain E0511`.