Remove the unnecessary bitcasts in check_call

Add support fot `bf16` linking

Author: sayantn

compiler/rustc_codegen_llvm/src/abi.rs

@@ -20,7 +20,7 @@ use smallvec::SmallVec;
 
 use crate::attributes::{self, llfn_attrs_from_instance};
 use crate::builder::Builder;
-use crate::context::CodegenCx;
+use crate::context::{CodegenCx, GenericCx, SCx};
 use crate::llvm::{self, Attribute, AttributePlace};
 use crate::type_::Type;
 use crate::type_of::LayoutLlvmExt;
@@ -362,15 +362,15 @@ fn match_intrinsic_signature<'ll>(
         );
     }
 
-    if !equate_ty(cx, rust_return_ty, llvm_return_ty) {
+    if !cx.equate_ty(rust_return_ty, llvm_return_ty) {
         error!(
             "Intrinsic signature mismatch: could not match `{rust_return_ty:?}` (found) with  {llvm_return_ty:?} (expected) as return type for `{fn_name}`"
         );
     }
     for (idx, (&rust_argument_ty, llvm_argument_ty)) in
         iter::zip(rust_argument_tys, llvm_argument_tys).enumerate()
     {
-        if !equate_ty(cx, rust_argument_ty, llvm_argument_ty) {
+        if !cx.equate_ty(rust_argument_ty, llvm_argument_ty) {
             error!(
                 "Intrinsic signature mismatch: could not match `{rust_return_ty:?}` (found) with  {llvm_return_ty:?} (expected) as argument {idx} for `{fn_name}`"
             );
@@ -380,28 +380,53 @@ fn match_intrinsic_signature<'ll>(
     fn_ty
 }
 
-fn equate_ty<'ll>(cx: &CodegenCx<'ll, '_>, rust_ty: &'ll Type, llvm_ty: &'ll Type) -> bool {
-    if rust_ty == llvm_ty {
-        return true;
-    }
-    if cx.type_kind(llvm_ty) == TypeKind::X86_AMX && cx.type_kind(rust_ty) == TypeKind::Vector {
-        let element_count = cx.vector_length(rust_ty);
-        let element_ty = cx.element_type(rust_ty);
-
-        let element_size_bits = match cx.type_kind(element_ty) {
-            TypeKind::Half => 16,
-            TypeKind::Float => 32,
-            TypeKind::Double => 64,
-            TypeKind::FP128 => 128,
-            TypeKind::Integer => cx.int_width(element_ty),
-            TypeKind::Pointer => cx.int_width(cx.isize_ty),
-            _ => bug!("Vector element type `{element_ty:?}` not one of integer, float or pointer"),
-        };
-        let vector_size_bits = element_size_bits * element_count as u64;
+impl<'ll, CX: Borrow<SCx<'ll>>> GenericCx<'ll, CX> {
+    pub(crate) fn equate_ty(&self, rust_ty: &'ll Type, llvm_ty: &'ll Type) -> bool {
+        if rust_ty == llvm_ty {
+            return true;
+        }
+
+        match self.type_kind(llvm_ty) {
+            TypeKind::X86_AMX if self.type_kind(rust_ty) == TypeKind::Vector => {
+                let element_count = self.vector_length(rust_ty);
+                let element_ty = self.element_type(rust_ty);
+
+                let element_size_bits = match self.type_kind(element_ty) {
+                    TypeKind::Half => 16,
+                    TypeKind::Float => 32,
+                    TypeKind::Double => 64,
+                    TypeKind::FP128 => 128,
+                    TypeKind::Integer => self.int_width(element_ty),
+                    TypeKind::Pointer => self.int_width(self.isize_ty()),
+                    _ => bug!(
+                        "Vector element type `{element_ty:?}` not one of integer, float or pointer"
+                    ),
+                };
+                let vector_size_bits = element_size_bits * element_count as u64;
+
+                vector_size_bits == 8192
+            }
+            TypeKind::BFloat => rust_ty == self.type_i16(),
+            TypeKind::Vector if self.type_kind(rust_ty) == TypeKind::Vector => {
+                let llvm_element_count = self.vector_length(llvm_ty);
+                let rust_element_count = self.vector_length(rust_ty);
 
-        return vector_size_bits == 8192;
+                if llvm_element_count != rust_element_count {
+                    return false;
+                }
+
+                let llvm_element_ty = self.element_type(llvm_ty);
+                let rust_element_ty = self.element_type(rust_ty);
+
+                if llvm_element_ty == self.type_bf16() {
+                    rust_element_ty == self.type_i16()
+                } else {
+                    false
+                }
+            }
+            _ => false,
+        }
     }
-    return false;
 }
 
 impl<'ll, 'tcx> FnAbiLlvmExt<'ll, 'tcx> for FnAbi<'tcx, Ty<'tcx>> {

compiler/rustc_codegen_llvm/src/builder.rs

@@ -1582,17 +1582,18 @@ impl<'a, 'll, CX: Borrow<SCx<'ll>>> GenericBuilder<'a, 'll, CX> {
             let actual_ty = self.cx.val_ty(actual_val);
 
             if expected_ty != actual_ty {
-                warn!(
+                assert!(
+                    self.cx.equate_ty(actual_ty, expected_ty),
                     "type mismatch in function call of {llfn:?}. \
-                        Expected {expected_ty:?} for param {i}, got {actual_ty:?}; injecting bitcast",
+                            Can't match {expected_ty:?} (expected) as parameter {i} with {actual_ty:?} (found)",
+                );
+
+                assert!(
+                    has_fnabi,
+                    "Should inject auto-bitcasts in function call of {llfn:?}, but not able to get Rust signature"
                 );
 
                 casted_args.to_mut()[i] = if self.cx.type_kind(expected_ty) == TypeKind::X86_AMX {
-                    // we can't do `cast_return` in without `FnAbi`
-                    assert!(
-                        has_fnabi,
-                        "Found `x86amx` for parameter {i} in function call of {llfn:?}, but not able to get Rust return type"
-                    );
                     self.cast_vector_to_tile(actual_val)
                 } else {
                     self.bitcast(actual_val, expected_ty)
@@ -1609,11 +1610,20 @@ impl<'a, 'll, CX: Borrow<SCx<'ll>>> GenericBuilder<'a, 'll, CX> {
 
     pub(crate) fn simple_call(
         &mut self,
+        name: &[u8],
         fn_ty: &'ll Type,
-        llfn: &'ll Value,
         args: &[&'ll Value],
     ) -> &'ll Value {
-        let args = self.cast_arguments("simple call", fn_ty, llfn, args, false);
+        let llfn = unsafe {
+            llvm::LLVMRustGetOrInsertFunction(
+                self.cx.llmod(),
+                name.as_c_char_ptr(),
+                name.len(),
+                fn_ty,
+            )
+        };
+
+        let args = self.cast_arguments("simple_call", fn_ty, llfn, args, false);
 
         unsafe {
             llvm::LLVMBuildCall2(
@@ -1646,15 +1656,8 @@ impl<'a, 'll, CX: Borrow<SCx<'ll>>> GenericBuilder<'a, 'll, CX> {
             assert!(type_params.is_empty());
             base_name.as_bytes()
         };
-        let llfn = unsafe {
-            llvm::LLVMRustGetOrInsertFunction(
-                self.cx.llmod(),
-                full_name.as_c_char_ptr(),
-                full_name.len(),
-                fn_ty,
-            )
-        };
-        self.simple_call(fn_ty, llfn, args)
+
+        self.simple_call(full_name, fn_ty, args)
     }
 }
 
@@ -1707,11 +1710,7 @@ impl<'a, 'll, 'tcx> Builder<'a, 'll, 'tcx> {
         let fn_ty =
             self.type_func(&[self.type_ptr(), self.type_ptr(), self.type_isize()], llreturn_ty);
 
-        let llfn = self
-            .get_function("memcmp")
-            .unwrap_or_else(|| self.declare_cfn("memcmp", llvm::UnnamedAddr::No, fn_ty));
-
-        self.simple_call(fn_ty, llfn, &[ptr1, ptr2, num])
+        self.simple_call(b"memcmp", fn_ty, &[ptr1, ptr2, num])
     }
 
     fn cast_return(

compiler/rustc_codegen_llvm/src/llvm/ffi.rs

@@ -1050,6 +1050,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

@@ -147,6 +147,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> {