Add linking support for x86amx and bf16 (only scalar and vector)

Author: sayantn

compiler/rustc_codegen_llvm/src/abi.rs

@@ -5,6 +5,7 @@ use std::iter::zip;
 use libc::c_uint;
 use rustc_abi::{BackendRepr, HasDataLayout, Primitive, Reg, RegKind, Size};
 use rustc_codegen_ssa::MemFlags;
+use rustc_codegen_ssa::common::TypeKind;
 use rustc_codegen_ssa::mir::operand::{OperandRef, OperandValue};
 use rustc_codegen_ssa::mir::place::{PlaceRef, PlaceValue};
 use rustc_codegen_ssa::traits::*;
@@ -327,6 +328,47 @@ pub(crate) trait FnAbiLlvmExt<'ll, 'tcx> {
     fn apply_attrs_callsite(&self, bx: &mut Builder<'_, 'll, 'tcx>, callsite: &'ll Value);
 }
 
+fn equate_ty<'ll>(cx: &CodegenCx<'ll, '_>, rust_ty: &'ll Type, llvm_ty: &'ll Type) -> bool {
+    if rust_ty == llvm_ty {
+        return true;
+    }
+    match cx.type_kind(llvm_ty) {
+        TypeKind::X86_AMX => {
+            // we will insert casts from/to x86amx in callsite, so this is fine
+            if 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"
+                    ),
+                };
+                element_size_bits * element_count as u64 == 8192
+            } else {
+                false
+            }
+        }
+        TypeKind::BFloat => rust_ty == cx.type_i16(),
+        TypeKind::Vector => {
+            let element_count = cx.vector_length(llvm_ty);
+            let element_ty = cx.element_type(llvm_ty);
+
+            if element_ty == cx.type_bf16() {
+                rust_ty == cx.type_vector(cx.type_i16(), element_count as u64)
+            } else {
+                false
+            }
+        }
+        _ => false,
+    }
+}
+
 impl<'ll, 'tcx> FnAbiLlvmExt<'ll, 'tcx> for FnAbi<'tcx, Ty<'tcx>> {
     fn llvm_return_type(&self, cx: &CodegenCx<'ll, 'tcx>) -> &'ll Type {
         match &self.ret.mode {
@@ -419,7 +461,7 @@ impl<'ll, 'tcx> FnAbiLlvmExt<'ll, 'tcx> for FnAbi<'tcx, Ty<'tcx>> {
                 ));
             }
 
-            if actual_return_ty != expected_return_ty {
+            if !equate_ty(cx, actual_return_ty, expected_return_ty) {
                 cx.tcx.dcx().fatal(format!(
                     "Intrinsic signature mismatch: expected {expected_return_ty:?} as return type for `{}`, found {actual_return_ty:?}",
                     str::from_utf8(name).unwrap()
@@ -428,7 +470,7 @@ impl<'ll, 'tcx> FnAbiLlvmExt<'ll, 'tcx> for FnAbi<'tcx, Ty<'tcx>> {
             for (idx, (actual_argument_ty, expected_argument_ty)) in
                 zip(actual_argument_tys, expected_argument_tys).enumerate()
             {
-                if actual_argument_ty != expected_argument_ty {
+                if !equate_ty(cx, actual_argument_ty, expected_argument_ty) {
                     cx.tcx.dcx().fatal(format!(
                         "Intrinsic signature mismatch: expected {expected_argument_ty:?} as argument {idx} for `{}`, found {actual_argument_ty:?}",
                         str::from_utf8(name).unwrap()

compiler/rustc_codegen_llvm/src/builder.rs

@@ -67,7 +67,7 @@ impl<'a, 'll> SBuilder<'a, 'll> {
     ) -> &'ll Value {
         debug!("call {:?} with args ({:?})", llfn, args);
 
-        let args = self.check_call("call", llty, llfn, args);
+        let args = self.cast_arguments("call", llty, llfn, args);
         let funclet_bundle = funclet.map(|funclet| funclet.bundle());
         let mut bundles: SmallVec<[_; 2]> = SmallVec::new();
         if let Some(funclet_bundle) = funclet_bundle {
@@ -101,6 +101,24 @@ impl<'a, 'll, CX: Borrow<SCx<'ll>>> GenericBuilder<'a, 'll, CX> {
         unsafe { llvm::LLVMBuildBitCast(self.llbuilder, val, dest_ty, UNNAMED) }
     }
 
+    pub(crate) fn cast_vector_to_tile(&mut self, val: &'ll Value) -> &'ll Value {
+        let vector_type = self.cx.val_ty(val);
+
+        assert!(self.cx.type_kind(vector_type) == TypeKind::Vector);
+        self.call_intrinsic("llvm.x86.cast.vector.to.tile", &[vector_type], &[val])
+    }
+
+    pub(crate) fn cast_tile_to_vector(
+        &mut self,
+        val: &'ll Value,
+        vector_type: &'ll Type,
+    ) -> &'ll Value {
+        assert!(self.cx.val_ty(val) == self.cx.type_x86amx());
+        assert!(self.cx.type_kind(vector_type) == TypeKind::Vector);
+
+        self.call_intrinsic("llvm.x86.cast.tile.to.vector", &[vector_type], &[val])
+    }
+
     pub(crate) fn ret_void(&mut self) {
         llvm::LLVMBuildRetVoid(self.llbuilder);
     }
@@ -349,7 +367,7 @@ impl<'a, 'll, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'll, 'tcx> {
     ) -> &'ll Value {
         debug!("invoke {:?} with args ({:?})", llfn, args);
 
-        let args = self.check_call("invoke", llty, llfn, args);
+        let args = self.cast_arguments("invoke", llty, llfn, args);
         let funclet_bundle = funclet.map(|funclet| funclet.bundle());
         let mut bundles: SmallVec<[_; 2]> = SmallVec::new();
         if let Some(funclet_bundle) = funclet_bundle {
@@ -381,8 +399,10 @@ impl<'a, 'll, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'll, 'tcx> {
         };
         if let Some(fn_abi) = fn_abi {
             fn_abi.apply_attrs_callsite(self, invoke);
+            self.cast_return(fn_abi, llfn, invoke)
+        } else {
+            invoke
         }
-        invoke
     }
 
     fn unreachable(&mut self) {
@@ -1348,7 +1368,7 @@ impl<'a, 'll, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'll, 'tcx> {
     ) -> &'ll Value {
         debug!("call {:?} with args ({:?})", llfn, args);
 
-        let args = self.check_call("call", llty, llfn, args);
+        let args = self.cast_arguments("call", llty, llfn, args);
         let funclet_bundle = funclet.map(|funclet| funclet.bundle());
         let mut bundles: SmallVec<[_; 2]> = SmallVec::new();
         if let Some(funclet_bundle) = funclet_bundle {
@@ -1378,8 +1398,10 @@ impl<'a, 'll, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'll, 'tcx> {
         };
         if let Some(fn_abi) = fn_abi {
             fn_abi.apply_attrs_callsite(self, call);
+            self.cast_return(fn_abi, llfn, call)
+        } else {
+            call
         }
-        call
     }
 
     fn zext(&mut self, val: &'ll Value, dest_ty: &'ll Type) -> &'ll Value {
@@ -1540,7 +1562,7 @@ impl<'a, 'll, CX: Borrow<SCx<'ll>>> GenericBuilder<'a, 'll, CX> {
         ret.expect("LLVM does not have support for catchret")
     }
 
-    fn check_call<'b>(
+    fn cast_arguments<'b>(
         &mut self,
         typ: &str,
         fn_ty: &'ll Type,
@@ -1571,7 +1593,11 @@ impl<'a, 'll, CX: Borrow<SCx<'ll>>> GenericBuilder<'a, 'll, CX> {
                             Expected {:?} for param {}, got {:?}; injecting bitcast",
                         llfn, expected_ty, i, actual_ty
                     );
-                    self.bitcast(actual_val, expected_ty)
+                    if self.cx.type_kind(expected_ty) == TypeKind::X86_AMX {
+                        self.cast_vector_to_tile(actual_val)
+                    } else {
+                        self.bitcast(actual_val, expected_ty)
+                    }
                 } else {
                     actual_val
                 }
@@ -1591,7 +1617,7 @@ impl<'a, 'll, CX: Borrow<SCx<'ll>>> GenericBuilder<'a, 'll, CX> {
         llfn: &'ll Value,
         args: &[&'ll Value],
     ) -> &'ll Value {
-        let args = self.check_call("simple call", fn_ty, llfn, args);
+        let args = self.cast_arguments("simple call", fn_ty, llfn, args);
 
         unsafe {
             llvm::LLVMBuildCall2(
@@ -1692,6 +1718,31 @@ impl<'a, 'll, 'tcx> Builder<'a, 'll, 'tcx> {
         self.simple_call(fn_ty, llfn, &[ptr1, ptr2, num])
     }
 
+    fn cast_return(
+        &mut self,
+        fn_abi: &FnAbi<'tcx, Ty<'tcx>>,
+        llfn: &'ll Value,
+        ret: &'ll Value,
+    ) -> &'ll Value {
+        let expected_ty = fn_abi.llvm_return_type(self.cx);
+        let actual_ty = self.cx.val_ty(ret);
+
+        if expected_ty != actual_ty {
+            debug!(
+                "type mismatch in function call of {:?}. \
+                    Expected {:?} for return value, got {:?}; injecting bitcast",
+                llfn, expected_ty, actual_ty
+            );
+            if self.cx.type_kind(actual_ty) == TypeKind::X86_AMX {
+                self.cast_tile_to_vector(ret, expected_ty)
+            } else {
+                self.bitcast(ret, expected_ty)
+            }
+        } else {
+            ret
+        }
+    }
+
     pub(crate) fn landing_pad(
         &mut self,
         ty: &'ll Type,
@@ -1721,7 +1772,7 @@ impl<'a, 'll, 'tcx> Builder<'a, 'll, 'tcx> {
     ) -> &'ll Value {
         debug!("invoke {:?} with args ({:?})", llfn, args);
 
-        let args = self.check_call("callbr", llty, llfn, args);
+        let args = self.cast_arguments("callbr", llty, llfn, args);
         let funclet_bundle = funclet.map(|funclet| funclet.bundle());
         let mut bundles: SmallVec<[_; 2]> = SmallVec::new();
         if let Some(funclet_bundle) = funclet_bundle {
@@ -1754,8 +1805,10 @@ impl<'a, 'll, 'tcx> Builder<'a, 'll, 'tcx> {
         };
         if let Some(fn_abi) = fn_abi {
             fn_abi.apply_attrs_callsite(self, callbr);
+            self.cast_return(fn_abi, llfn, callbr)
+        } else {
+            callbr
         }
-        callbr
     }
 
     // Emits CFI pointer type membership tests.