Add autocasts for bf16 and bf16xN

Author: sayantn

compiler/rustc_codegen_llvm/src/abi.rs

@@ -370,6 +370,8 @@ impl<'ll, CX: Borrow<SCx<'ll>>> GenericCx<'ll, CX> {
         }
 
         match self.type_kind(llvm_ty) {
+            TypeKind::BFloat => rust_ty == self.type_i16(),
+
             // 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)
@@ -387,11 +389,18 @@ impl<'ll, CX: Borrow<SCx<'ll>>> GenericCx<'ll, CX> {
                     },
                 )
             }
-            TypeKind::Vector if self.element_type(llvm_ty) == self.type_i1() => {
+            TypeKind::Vector => {
                 let element_count = self.vector_length(llvm_ty) as u64;
-                let int_width = element_count.next_power_of_two().max(8);
+                let llvm_element_ty = self.element_type(llvm_ty);
 
-                rust_ty == self.type_ix(int_width)
+                if llvm_element_ty == self.type_bf16() {
+                    rust_ty == self.type_vector(self.type_i16(), element_count)
+                } else if llvm_element_ty == self.type_i1() {
+                    let int_width = element_count.next_power_of_two().max(8);
+                    rust_ty == self.type_ix(int_width)
+                } else {
+                    false
+                }
             }
             _ => false,
         }

compiler/rustc_codegen_llvm/src/builder.rs

@@ -1761,7 +1761,7 @@ impl<'a, 'll, 'tcx> Builder<'a, 'll, 'tcx> {
                     self.zext_i1_vector_to_int(val, src_ty, dest_ty)
                 }
             }
-            _ => unreachable!(),
+            _ => self.bitcast(val, dest_ty), // for `bf16(xN)` <-> `u16(xN)`
         }
     }
 

compiler/rustc_codegen_llvm/src/llvm/ffi.rs

@@ -963,6 +963,9 @@ unsafe extern "C" {
     pub(crate) fn LLVMDoubleTypeInContext(C: &Context) -> &Type;
     pub(crate) fn LLVMFP128TypeInContext(C: &Context) -> &Type;
 
+    // Operations on non-IEEE real types
+    pub(crate) fn LLVMBFloatTypeInContext(C: &Context) -> &Type;
+
     // Operations on function types
     pub(crate) fn LLVMFunctionType<'a>(
         ReturnType: &'a Type,

compiler/rustc_codegen_llvm/src/type_.rs

@@ -174,6 +174,10 @@ impl<'ll, CX: Borrow<SCx<'ll>>> GenericCx<'ll, CX> {
             )
         }
     }
+
+    pub(crate) fn type_bf16(&self) -> &'ll Type {
+        unsafe { llvm::LLVMBFloatTypeInContext(self.llcx()) }
+    }
 }
 
 impl<'ll, CX: Borrow<SCx<'ll>>> BaseTypeCodegenMethods for GenericCx<'ll, CX> {
@@ -247,7 +251,7 @@ impl<'ll, CX: Borrow<SCx<'ll>>> BaseTypeCodegenMethods for GenericCx<'ll, CX> {
 
     fn float_width(&self, ty: &'ll Type) -> usize {
         match self.type_kind(ty) {
-            TypeKind::Half => 16,
+            TypeKind::Half | TypeKind::BFloat => 16,
             TypeKind::Float => 32,
             TypeKind::Double => 64,
             TypeKind::X86_FP80 => 80,

tests/codegen-llvm/inject-autocast.rs

@@ -4,7 +4,7 @@
 #![feature(link_llvm_intrinsics, abi_unadjusted, repr_simd, simd_ffi, portable_simd, f16)]
 #![crate_type = "lib"]
 
-use std::simd::i64x2;
+use std::simd::{f32x4, i16x8, i64x2};
 
 #[repr(simd)]
 pub struct Tile([i8; 1024]);
@@ -36,6 +36,19 @@ pub unsafe fn struct_with_i1_vector_autocast(a: i64x2, b: i64x2) -> (u8, u8) {
     foo(a, b)
 }
 
+// CHECK-LABEL: @bf16_vector_autocast
+#[no_mangle]
+pub unsafe fn bf16_vector_autocast(a: f32x4) -> i16x8 {
+    extern "unadjusted" {
+        #[link_name = "llvm.x86.vcvtneps2bf16128"]
+        fn foo(a: f32x4) -> i16x8;
+    }
+
+    // CHECK: [[A:%[0-9]+]] = call <8 x bfloat> @llvm.x86.vcvtneps2bf16128(<4 x float> {{.*}})
+    // CHECK: bitcast <8 x bfloat> [[A]] to <8 x i16>
+    foo(a)
+}
+
 // CHECK-LABEL: @struct_autocast
 #[no_mangle]
 pub unsafe fn struct_autocast(key_metadata: u32, key: i64x2) -> Bar {
@@ -77,6 +90,8 @@ pub unsafe fn i1_vector_autocast(a: f16x8) -> u8 {
 
 // CHECK: declare { <2 x i1>, <2 x i1> } @llvm.x86.avx512.vp2intersect.q.128(<2 x i64>, <2 x i64>)
 
+// CHECK: declare <8 x bfloat> @llvm.x86.vcvtneps2bf16128(<4 x float>)
+
 // 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 <8 x i1> @llvm.x86.avx512fp16.fpclass.ph.128(<8 x half>, i32 immarg)