Add linking support for x86amx

Author: sayantn

compiler/rustc_codegen_llvm/src/abi.rs

@@ -1,10 +1,10 @@
 use std::borrow::Borrow;
-use std::cmp;
+use std::{cmp, iter};
 
-use itertools::zip_eq;
 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::*;
@@ -362,15 +362,15 @@ fn match_intrinsic_signature<'ll>(
         );
     }
 
-    if rust_return_ty != llvm_return_ty {
+    if !equate_ty(cx, 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
-        zip_eq(rust_argument_tys, llvm_argument_tys).enumerate()
+        iter::zip(rust_argument_tys, llvm_argument_tys).enumerate()
     {
-        if rust_argument_ty != llvm_argument_ty {
+        if !equate_ty(cx, 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,6 +380,30 @@ 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;
+
+        return vector_size_bits == 8192;
+    }
+    return 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 {

compiler/rustc_codegen_llvm/src/builder.rs

@@ -27,7 +27,7 @@ use rustc_span::Span;
 use rustc_target::callconv::FnAbi;
 use rustc_target::spec::{HasTargetSpec, SanitizerSet, Target};
 use smallvec::SmallVec;
-use tracing::{debug, instrument};
+use tracing::{debug, instrument, warn};
 
 use crate::abi::FnAbiLlvmExt;
 use crate::attributes;
@@ -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, false);
         let funclet_bundle = funclet.map(|funclet| funclet.bundle());
         let mut bundles: SmallVec<[_; 2]> = SmallVec::new();
         if let Some(funclet_bundle) = funclet_bundle {
@@ -79,7 +79,7 @@ impl<'a, 'll> SBuilder<'a, 'll> {
                 self.llbuilder,
                 llty,
                 llfn,
-                args.as_ptr() as *const &llvm::Value,
+                args.as_ptr(),
                 args.len() as c_uint,
                 bundles.as_ptr(),
                 bundles.len() as c_uint,
@@ -101,6 +101,22 @@ 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.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 +365,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, fn_abi.is_some());
         let funclet_bundle = funclet.map(|funclet| funclet.bundle());
         let mut bundles: SmallVec<[_; 2]> = SmallVec::new();
         if let Some(funclet_bundle) = funclet_bundle {
@@ -381,8 +397,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 +1366,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, fn_abi.is_some());
         let funclet_bundle = funclet.map(|funclet| funclet.bundle());
         let mut bundles: SmallVec<[_; 2]> = SmallVec::new();
         if let Some(funclet_bundle) = funclet_bundle {
@@ -1378,8 +1396,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,45 +1560,47 @@ 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,
         llfn: &'ll Value,
         args: &'b [&'ll Value],
+        has_fnabi: bool,
     ) -> Cow<'b, [&'ll Value]> {
-        assert!(
-            self.cx.type_kind(fn_ty) == TypeKind::Function,
-            "builder::{typ} not passed a function, but {fn_ty:?}"
+        assert_eq!(
+            self.cx.type_kind(fn_ty),
+            TypeKind::Function,
+            "{typ} not passed a function, but {fn_ty:?}"
         );
 
         let param_tys = self.cx.func_params_types(fn_ty);
 
-        let all_args_match = iter::zip(&param_tys, args.iter().map(|&v| self.cx.val_ty(v)))
-            .all(|(expected_ty, actual_ty)| *expected_ty == actual_ty);
+        let mut casted_args = Cow::Borrowed(args);
 
-        if all_args_match {
-            return Cow::Borrowed(args);
-        }
+        for (i, (expected_ty, &actual_val)) in iter::zip(param_tys, args).enumerate() {
+            let actual_ty = self.cx.val_ty(actual_val);
 
-        let casted_args: Vec<_> = iter::zip(param_tys, args)
-            .enumerate()
-            .map(|(i, (expected_ty, &actual_val))| {
-                let actual_ty = self.cx.val_ty(actual_val);
-                if expected_ty != actual_ty {
-                    debug!(
-                        "type mismatch in function call of {:?}. \
-                            Expected {:?} for param {}, got {:?}; injecting bitcast",
-                        llfn, expected_ty, i, actual_ty
+            if expected_ty != actual_ty {
+                warn!(
+                    "type mismatch in function call of {llfn:?}. \
+                        Expected {expected_ty:?} for param {i}, got {actual_ty:?}; injecting bitcast",
+                );
+
+                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.bitcast(actual_val, expected_ty)
+                    self.cast_vector_to_tile(actual_val)
                 } else {
-                    actual_val
+                    self.bitcast(actual_val, expected_ty)
                 }
-            })
-            .collect();
+            }
+        }
 
-        Cow::Owned(casted_args)
+        casted_args
     }
 
     pub(crate) fn va_arg(&mut self, list: &'ll Value, ty: &'ll Type) -> &'ll Value {
@@ -1591,7 +1613,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, false);
 
         unsafe {
             llvm::LLVMBuildCall2(
@@ -1692,6 +1714,30 @@ 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 {
+            warn!(
+                "Type mismatch in function call of {llfn:?}. \
+                    Expected {expected_ty:?} for return type, got {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 +1767,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, fn_abi.is_some());
         let funclet_bundle = funclet.map(|funclet| funclet.bundle());
         let mut bundles: SmallVec<[_; 2]> = SmallVec::new();
         if let Some(funclet_bundle) = funclet_bundle {
@@ -1754,8 +1800,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.