Add autocasts for i1 vectors

Author: sayantn

compiler/rustc_codegen_llvm/src/abi.rs

@@ -375,26 +375,31 @@ impl<'ll, CX: Borrow<SCx<'ll>>> GenericCx<'ll, CX> {
             return true;
         }
 
-        // Some LLVM intrinsics return **non-packed** structs, but they can't be mimicked from Rust
-        // due to auto field-alignment in non-packed structs (packed structs are represented in LLVM
-        // as, well, packed structs, so they won't match with those either)
-        if self.type_kind(llvm_ty) == TypeKind::Struct
-            && self.type_kind(rust_ty) == TypeKind::Struct
-        {
-            let rust_element_tys = self.struct_element_types(rust_ty);
-            let llvm_element_tys = self.struct_element_types(llvm_ty);
+        match self.type_kind(llvm_ty) {
+            // Some LLVM intrinsics return **non-packed** structs, but they can't be mimicked from Rust
+            // due to auto field-alignment in non-packed structs (packed structs are represented in LLVM
+            // as, well, packed structs, so they won't match with those either)
+            TypeKind::Struct if self.type_kind(rust_ty) == TypeKind::Struct => {
+                let rust_element_tys = self.struct_element_types(rust_ty);
+                let llvm_element_tys = self.struct_element_types(llvm_ty);
+
+                if rust_element_tys.len() != llvm_element_tys.len() {
+                    return false;
+                }
 
-            if rust_element_tys.len() != llvm_element_tys.len() {
-                return false;
+                iter::zip(rust_element_tys, llvm_element_tys).all(
+                    |(rust_element_ty, llvm_element_ty)| {
+                        self.equate_ty(rust_element_ty, llvm_element_ty)
+                    },
+                )
             }
+            TypeKind::Vector if self.element_type(llvm_ty) == self.type_i1() => {
+                let element_count = self.vector_length(llvm_ty) as u64;
+                let int_width = element_count.next_power_of_two().max(8);
 
-            iter::zip(rust_element_tys, llvm_element_tys).all(
-                |(rust_element_ty, llvm_element_ty)| {
-                    self.equate_ty(rust_element_ty, llvm_element_ty)
-                },
-            )
-        } else {
-            false
+                rust_ty == self.type_ix(int_width)
+            }
+            _ => false,
         }
     }
 }

compiler/rustc_codegen_llvm/src/builder.rs

@@ -1680,6 +1680,46 @@ impl<'a, 'll, CX: Borrow<SCx<'ll>>> GenericBuilder<'a, 'll, CX> {
     }
 }
 impl<'a, 'll, 'tcx> Builder<'a, 'll, 'tcx> {
+    fn trunc_int_to_i1_vector(&mut self, val: &'ll Value, dest_ty: &'ll Type) -> &'ll Value {
+        let vector_length = self.vector_length(dest_ty) as u64;
+        let int_width = vector_length.next_power_of_two().max(8);
+
+        let bitcasted = self.bitcast(val, self.type_vector(self.type_i1(), int_width));
+        if vector_length == int_width {
+            bitcasted
+        } else {
+            let shuffle_mask =
+                (0..vector_length).map(|i| self.const_i32(i as i32)).collect::<Vec<_>>();
+            self.shuffle_vector(bitcasted, bitcasted, self.const_vector(&shuffle_mask))
+        }
+    }
+
+    fn zext_i1_vector_to_int(
+        &mut self,
+        mut val: &'ll Value,
+        src_ty: &'ll Type,
+        dest_ty: &'ll Type,
+    ) -> &'ll Value {
+        let vector_length = self.vector_length(src_ty) as u64;
+        let int_width = vector_length.next_power_of_two().max(8);
+
+        if vector_length != int_width {
+            let shuffle_indices = match vector_length {
+                0 => unreachable!("zero length vectors are not allowed"),
+                1 => vec![0, 1, 1, 1, 1, 1, 1, 1],
+                2 => vec![0, 1, 2, 3, 2, 3, 2, 3],
+                3 => vec![0, 1, 2, 3, 4, 5, 3, 4],
+                4.. => (0..int_width as i32).collect(),
+            };
+            let shuffle_mask =
+                shuffle_indices.into_iter().map(|i| self.const_i32(i)).collect::<Vec<_>>();
+            val =
+                self.shuffle_vector(val, self.const_null(src_ty), self.const_vector(&shuffle_mask));
+        }
+
+        self.bitcast(val, dest_ty)
+    }
+
     fn autocast(
         &mut self,
         llfn: &'ll Value,
@@ -1708,6 +1748,13 @@ impl<'a, 'll, 'tcx> Builder<'a, 'll, 'tcx> {
                 }
                 ret
             }
+            TypeKind::Vector if self.element_type(llvm_ty) == self.type_i1() => {
+                if is_argument {
+                    self.trunc_int_to_i1_vector(val, dest_ty)
+                } else {
+                    self.zext_i1_vector_to_int(val, src_ty, dest_ty)
+                }
+            }
             _ => unreachable!(),
         }
     }

tests/codegen-llvm/inject-autocast.rs

@@ -1,7 +1,7 @@
 //@ compile-flags: -C opt-level=0
 //@ only-x86_64
 
-#![feature(link_llvm_intrinsics, abi_unadjusted, simd_ffi, portable_simd)]
+#![feature(link_llvm_intrinsics, abi_unadjusted, repr_simd, simd_ffi, portable_simd, f16)]
 #![crate_type = "lib"]
 
 use std::simd::i64x2;
@@ -10,6 +10,9 @@ use std::simd::i64x2;
 pub struct Bar(u32, i64x2, i64x2, i64x2, i64x2, i64x2, i64x2);
 // CHECK: %Bar = type <{ i32, <2 x i64>, <2 x i64>, <2 x i64>, <2 x i64>, <2 x i64>, <2 x i64> }>
 
+#[repr(simd)]
+pub struct f16x8([f16; 8]);
+
 // CHECK-LABEL: @struct_autocast
 #[no_mangle]
 pub unsafe fn struct_autocast(key_metadata: u32, key: i64x2) -> Bar {
@@ -36,4 +39,41 @@ pub unsafe fn struct_autocast(key_metadata: u32, key: i64x2) -> Bar {
     foo(key_metadata, key)
 }
 
+// CHECK-LABEL: @struct_with_i1_vector_autocast
+#[no_mangle]
+pub unsafe fn struct_with_i1_vector_autocast(a: i64x2, b: i64x2) -> (u8, u8) {
+    extern "unadjusted" {
+        #[link_name = "llvm.x86.avx512.vp2intersect.q.128"]
+        fn foo(a: i64x2, b: i64x2) -> (u8, u8);
+    }
+
+    // CHECK: [[A:%[0-9]+]] = call { <2 x i1>, <2 x i1> } @llvm.x86.avx512.vp2intersect.q.128(<2 x i64> {{.*}}, <2 x i64> {{.*}})
+    // CHECK: [[B:%[0-9]+]] = extractvalue { <2 x i1>, <2 x i1> } [[A]], 0
+    // CHECK: [[C:%[0-9]+]] = shufflevector <2 x i1> [[B]], <2 x i1> zeroinitializer, <8 x i32> <i32 0, i32 1, i32 2, i32 3, i32 2, i32 3, i32 2, i32 3>
+    // CHECK: [[D:%[0-9]+]] = bitcast <8 x i1> [[C]] to i8
+    // CHECK: [[E:%[0-9]+]] = insertvalue { i8, i8 } poison, i8 [[D]], 0
+    // CHECK: [[F:%[0-9]+]] = extractvalue { <2 x i1>, <2 x i1> } [[A]], 1
+    // CHECK: [[G:%[0-9]+]] = shufflevector <2 x i1> [[F]], <2 x i1> zeroinitializer, <8 x i32> <i32 0, i32 1, i32 2, i32 3, i32 2, i32 3, i32 2, i32 3>
+    // CHECK: [[H:%[0-9]+]] = bitcast <8 x i1> [[G]] to i8
+    // CHECK:                 insertvalue { i8, i8 } [[E]], i8 [[H]], 1
+    foo(a, b)
+}
+
+// CHECK-LABEL: @i1_vector_autocast
+#[no_mangle]
+pub unsafe fn i1_vector_autocast(a: f16x8) -> u8 {
+    extern "unadjusted" {
+        #[link_name = "llvm.x86.avx512fp16.fpclass.ph.128"]
+        fn foo(a: f16x8, b: i32) -> u8;
+    }
+
+    // CHECK: [[A:%[0-9]+]] = call <8 x i1> @llvm.x86.avx512fp16.fpclass.ph.128(<8 x half> {{.*}}, i32 1)
+    // CHECK: bitcast <8 x i1> [[A]] to i8
+    foo(a, 1)
+}
+
 // CHECK: declare { i32, <2 x i64>, <2 x i64>, <2 x i64>, <2 x i64>, <2 x i64>, <2 x i64> } @llvm.x86.encodekey128(i32, <2 x i64>)
+
+// CHECK: declare { <2 x i1>, <2 x i1> } @llvm.x86.avx512.vp2intersect.q.128(<2 x i64>, <2 x i64>)
+
+// CHECK: declare <8 x i1> @llvm.x86.avx512fp16.fpclass.ph.128(<8 x half>, i32 immarg)