Add auto-bitcasts from/to x86amx and i32x256 for AMX intrinsics

Author: sayantn

compiler/rustc_codegen_gcc/src/type_of.rs

@@ -1,6 +1,6 @@
 use std::fmt::Write;
 
-use gccjit::{Struct, Type};
+use gccjit::{RValue, Struct, Type};
 use rustc_abi as abi;
 use rustc_abi::Primitive::*;
 use rustc_abi::{
@@ -373,7 +373,11 @@ impl<'gcc, 'tcx> LayoutTypeCodegenMethods<'tcx> for CodegenCx<'gcc, 'tcx> {
         unimplemented!();
     }
 
-    fn fn_decl_backend_type(&self, fn_abi: &FnAbi<'tcx, Ty<'tcx>>) -> Type<'gcc> {
+    fn fn_decl_backend_type(
+        &self,
+        fn_abi: &FnAbi<'tcx, Ty<'tcx>>,
+        _fn_ptr: RValue<'gcc>,
+    ) -> Type<'gcc> {
         // FIXME(antoyo): Should we do something with `FnAbiGcc::fn_attributes`?
         let FnAbiGcc { return_type, arguments_type, is_c_variadic, .. } = fn_abi.gcc_type(self);
         self.context.new_function_pointer_type(None, return_type, &arguments_type, is_c_variadic)

compiler/rustc_codegen_llvm/src/abi.rs

@@ -4,6 +4,7 @@ use std::cmp;
 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::*;
@@ -308,7 +309,12 @@ impl<'ll, 'tcx> ArgAbiBuilderMethods<'tcx> for Builder<'_, 'll, 'tcx> {
 }
 
 pub(crate) trait FnAbiLlvmExt<'ll, 'tcx> {
-    fn llvm_type(&self, cx: &CodegenCx<'ll, 'tcx>) -> &'ll Type;
+    fn llvm_type(
+        &self,
+        cx: &CodegenCx<'ll, 'tcx>,
+        name: &[u8],
+        is_llvm_intrinsic: bool,
+    ) -> &'ll Type;
     fn ptr_to_llvm_type(&self, cx: &CodegenCx<'ll, 'tcx>) -> &'ll Type;
     fn llvm_cconv(&self, cx: &CodegenCx<'ll, 'tcx>) -> llvm::CallConv;
 
@@ -325,26 +331,45 @@ pub(crate) trait FnAbiLlvmExt<'ll, 'tcx> {
 }
 
 impl<'ll, 'tcx> FnAbiLlvmExt<'ll, 'tcx> for FnAbi<'tcx, Ty<'tcx>> {
-    fn llvm_type(&self, cx: &CodegenCx<'ll, 'tcx>) -> &'ll Type {
+    fn llvm_type(
+        &self,
+        cx: &CodegenCx<'ll, 'tcx>,
+        name: &[u8],
+        is_llvm_intrinsic: bool,
+    ) -> &'ll Type {
         // Ignore "extra" args from the call site for C variadic functions.
         // Only the "fixed" args are part of the LLVM function signature.
         let args =
             if self.c_variadic { &self.args[..self.fixed_count as usize] } else { &self.args };
 
+        let amx_intrinsic =
+            is_llvm_intrinsic && name.starts_with(b"llvm.x86.") && name.ends_with(b".internal");
+        let adjust_ty = |ty| {
+            // Change type to `x86amx` from `i32x256` for x86_64 AMX intrinsics
+            if amx_intrinsic && cx.type_kind(ty) == TypeKind::Vector && cx.vector_length(ty) == 256
+            {
+                let element_ty = cx.element_type(ty);
+                if cx.type_kind(element_ty) == TypeKind::Integer && cx.int_width(element_ty) == 32 {
+                    return cx.type_x86amx();
+                }
+            }
+            ty
+        };
+
         // This capacity calculation is approximate.
         let mut llargument_tys = Vec::with_capacity(
             self.args.len() + if let PassMode::Indirect { .. } = self.ret.mode { 1 } else { 0 },
         );
 
-        let llreturn_ty = match &self.ret.mode {
+        let llreturn_ty = adjust_ty(match &self.ret.mode {
             PassMode::Ignore => cx.type_void(),
             PassMode::Direct(_) | PassMode::Pair(..) => self.ret.layout.immediate_llvm_type(cx),
             PassMode::Cast { cast, pad_i32: _ } => cast.llvm_type(cx),
             PassMode::Indirect { .. } => {
                 llargument_tys.push(cx.type_ptr());
                 cx.type_void()
             }
-        };
+        });
 
         for arg in args {
             // Note that the exact number of arguments pushed here is carefully synchronized with
@@ -388,7 +413,7 @@ impl<'ll, 'tcx> FnAbiLlvmExt<'ll, 'tcx> for FnAbi<'tcx, Ty<'tcx>> {
                     cast.llvm_type(cx)
                 }
             };
-            llargument_tys.push(llarg_ty);
+            llargument_tys.push(adjust_ty(llarg_ty));
         }
 
         if self.c_variadic {

compiler/rustc_codegen_llvm/src/builder.rs

@@ -1435,7 +1435,12 @@ 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);
         }
-        call
+
+        if self.cx.type_kind(self.cx.val_ty(call)) == TypeKind::X86_AMX {
+            self.bitcast(call, self.cx.type_vector(self.cx.type_i32(), 256))
+        } else {
+            call
+        }
     }
 
     fn zext(&mut self, val: &'ll Value, dest_ty: &'ll Type) -> &'ll Value {

compiler/rustc_codegen_llvm/src/callee.rs

@@ -4,7 +4,7 @@
 //! and methods are represented as just a fn ptr and not a full
 //! closure.
 
-use rustc_codegen_ssa::common;
+use rustc_codegen_ssa::{base, common};
 use rustc_middle::ty::layout::{FnAbiOf, HasTyCtxt, HasTypingEnv};
 use rustc_middle::ty::{self, Instance, TypeVisitableExt};
 use tracing::debug;
@@ -36,6 +36,8 @@ pub(crate) fn get_fn<'ll, 'tcx>(cx: &CodegenCx<'ll, 'tcx>, instance: Instance<'t
         llfn
     } else {
         let instance_def_id = instance.def_id();
+        let is_llvm_intrinsic = base::is_llvm_intrinsic(tcx, instance_def_id);
+
         let llfn = if tcx.sess.target.arch == "x86"
             && let Some(dllimport) = crate::common::get_dllimport(tcx, instance_def_id, sym)
         {
@@ -53,6 +55,7 @@ pub(crate) fn get_fn<'ll, 'tcx>(cx: &CodegenCx<'ll, 'tcx>, instance: Instance<'t
                 ),
                 fn_abi,
                 Some(instance),
+                is_llvm_intrinsic,
             );
 
             // Fix for https://github.com/rust-lang/rust/issues/104453
@@ -69,7 +72,7 @@ pub(crate) fn get_fn<'ll, 'tcx>(cx: &CodegenCx<'ll, 'tcx>, instance: Instance<'t
             llvm::set_dllimport_storage_class(llfn);
             llfn
         } else {
-            cx.declare_fn(sym, fn_abi, Some(instance))
+            cx.declare_fn(sym, fn_abi, Some(instance), is_llvm_intrinsic)
         };
         debug!("get_fn: not casting pointer!");
 

compiler/rustc_codegen_llvm/src/consts.rs

@@ -191,7 +191,12 @@ fn check_and_apply_linkage<'ll, 'tcx>(
                 let fn_sig = sig.with(*header);
 
                 let fn_abi = cx.fn_abi_of_fn_ptr(fn_sig, ty::List::empty());
-                cx.declare_fn(sym, &fn_abi, None)
+                cx.declare_fn(
+                    sym,
+                    &fn_abi,
+                    None,
+                    rustc_codegen_ssa::base::is_llvm_intrinsic(cx.tcx, def_id),
+                )
             } else {
                 cx.declare_global(sym, cx.type_i8())
             }

compiler/rustc_codegen_llvm/src/declare.rs

@@ -147,6 +147,7 @@ impl<'ll, 'tcx> CodegenCx<'ll, 'tcx> {
         name: &str,
         fn_abi: &FnAbi<'tcx, Ty<'tcx>>,
         instance: Option<Instance<'tcx>>,
+        is_llvm_intrinsic: bool,
     ) -> &'ll Value {
         debug!("declare_rust_fn(name={:?}, fn_abi={:?})", name, fn_abi);
 
@@ -158,7 +159,7 @@ impl<'ll, 'tcx> CodegenCx<'ll, 'tcx> {
             fn_abi.llvm_cconv(self),
             llvm::UnnamedAddr::Global,
             llvm::Visibility::Default,
-            fn_abi.llvm_type(self),
+            fn_abi.llvm_type(self, name.as_ref(), is_llvm_intrinsic),
         );
         fn_abi.apply_attrs_llfn(self, llfn, instance);
 

compiler/rustc_codegen_llvm/src/intrinsic.rs

@@ -1088,10 +1088,11 @@ fn gen_fn<'a, 'll, 'tcx>(
     name: &str,
     rust_fn_sig: ty::PolyFnSig<'tcx>,
     codegen: &mut dyn FnMut(Builder<'a, 'll, 'tcx>),
+    is_llvm_intrinsic: bool,
 ) -> (&'ll Type, &'ll Value) {
     let fn_abi = cx.fn_abi_of_fn_ptr(rust_fn_sig, ty::List::empty());
-    let llty = fn_abi.llvm_type(cx);
-    let llfn = cx.declare_fn(name, fn_abi, None);
+    let llty = fn_abi.llvm_type(cx, name.as_ref(), is_llvm_intrinsic);
+    let llfn = cx.declare_fn(name, fn_abi, None, is_llvm_intrinsic);
     cx.set_frame_pointer_type(llfn);
     cx.apply_target_cpu_attr(llfn);
     // FIXME(eddyb) find a nicer way to do this.
@@ -1147,7 +1148,7 @@ fn get_rust_try_fn<'a, 'll, 'tcx>(
         hir::Safety::Unsafe,
         ExternAbi::Rust,
     ));
-    let rust_try = gen_fn(cx, "__rust_try", rust_fn_sig, codegen);
+    let rust_try = gen_fn(cx, "__rust_try", rust_fn_sig, codegen, false);
     cx.rust_try_fn.set(Some(rust_try));
     rust_try
 }

compiler/rustc_codegen_llvm/src/llvm/ffi.rs

@@ -1055,6 +1055,9 @@ unsafe extern "C" {
     pub(crate) fn LLVMPointerTypeInContext(C: &Context, AddressSpace: c_uint) -> &Type;
     pub(crate) fn LLVMVectorType(ElementType: &Type, ElementCount: c_uint) -> &Type;
 
+    // Special X86 Type for AMX
+    pub(crate) fn LLVMX86AMXTypeInContext(C: &Context) -> &Type;
+
     pub(crate) fn LLVMGetElementType(Ty: &Type) -> &Type;
     pub(crate) fn LLVMGetVectorSize(VectorTy: &Type) -> c_uint;
 
@@ -1177,6 +1180,7 @@ unsafe extern "C" {
 
     // Operations on functions
     pub(crate) fn LLVMSetFunctionCallConv(Fn: &Value, CC: c_uint);
+    pub(crate) fn LLVMGetIntrinsicID(Fn: &Value) -> c_uint;
 
     // Operations on parameters
     pub(crate) fn LLVMIsAArgument(Val: &Value) -> Option<&Value>;

compiler/rustc_codegen_llvm/src/mono_item.rs

@@ -53,7 +53,12 @@ impl<'tcx> PreDefineCodegenMethods<'tcx> for CodegenCx<'_, 'tcx> {
         assert!(!instance.args.has_infer());
 
         let fn_abi = self.fn_abi_of_instance(instance, ty::List::empty());
-        let lldecl = self.declare_fn(symbol_name, fn_abi, Some(instance));
+        let lldecl = self.declare_fn(
+            symbol_name,
+            fn_abi,
+            Some(instance),
+            rustc_codegen_ssa::base::is_llvm_intrinsic(self.tcx, instance.def_id()),
+        );
         llvm::set_linkage(lldecl, base::linkage_to_llvm(linkage));
         let attrs = self.tcx.codegen_fn_attrs(instance.def_id());
         base::set_link_section(lldecl, attrs);

compiler/rustc_codegen_llvm/src/type_.rs

@@ -154,6 +154,10 @@ impl<'ll, CX: Borrow<SCx<'ll>>> GenericCx<'ll, CX> {
             )
         }
     }
+
+    pub(crate) fn type_x86amx(&self) -> &'ll Type {
+        unsafe { llvm::LLVMX86AMXTypeInContext(self.llcx()) }
+    }
 }
 
 impl<'ll, CX: Borrow<SCx<'ll>>> BaseTypeCodegenMethods for GenericCx<'ll, CX> {
@@ -284,8 +288,14 @@ impl<'ll, 'tcx> LayoutTypeCodegenMethods<'tcx> for CodegenCx<'ll, 'tcx> {
     fn cast_backend_type(&self, ty: &CastTarget) -> &'ll Type {
         ty.llvm_type(self)
     }
-    fn fn_decl_backend_type(&self, fn_abi: &FnAbi<'tcx, Ty<'tcx>>) -> &'ll Type {
-        fn_abi.llvm_type(self)
+    fn fn_decl_backend_type(
+        &self,
+        fn_abi: &FnAbi<'tcx, Ty<'tcx>>,
+        fn_ptr: &'ll Value,
+    ) -> &'ll Type {
+        let intrinsic_id = unsafe { llvm::LLVMGetIntrinsicID(fn_ptr) };
+        // When the function is not an intrinsic, `Intrinsic::getIntrinsicID` returns `Intrinsic::not_intrinsic`, which is always defined to be 0
+        fn_abi.llvm_type(self, llvm::get_value_name(fn_ptr), intrinsic_id != 0)
     }
     fn fn_ptr_backend_type(&self, fn_abi: &FnAbi<'tcx, Ty<'tcx>>) -> &'ll Type {
         fn_abi.ptr_to_llvm_type(self)