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

Author: sayantn

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::*;
@@ -331,20 +332,37 @@ impl<'ll, 'tcx> FnAbiLlvmExt<'ll, 'tcx> for FnAbi<'tcx, Ty<'tcx>> {
         let args =
             if self.c_variadic { &self.args[..self.fixed_count as usize] } else { &self.args };
 
+        let adjust_ty = |ty| {
+            // todo: rectify this to be more selective (help wanted)
+            let probably_unadjusted = self.conv == Conv::C && !self.can_unwind && !self.c_variadic;
+            let probably_amx_intrinsic = probably_unadjusted && cx.tcx.sess.target.arch == "x86_64";
+            // Change type to `x86amx` from `i32x256` for x86_64 AMX intrinsics
+            if probably_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 +406,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/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;
 

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> {

compiler/rustc_target/src/target_features.rs

@@ -786,7 +786,7 @@ pub fn all_rust_features() -> impl Iterator<Item = (&'static str, Stability)> {
 // certain size to have their "proper" ABI on each architecture.
 // Note that they must be kept sorted by vector size.
 const X86_FEATURES_FOR_CORRECT_VECTOR_ABI: &'static [(u64, &'static str)] =
-    &[(128, "sse"), (256, "avx"), (512, "avx512f")]; // FIXME: might need changes for AVX10.
+    &[(128, "sse"), (256, "avx"), (512, "avx512f"), (8192, "amx-tile")];
 const AARCH64_FEATURES_FOR_CORRECT_VECTOR_ABI: &'static [(u64, &'static str)] = &[(128, "neon")];
 
 // We might want to add "helium" too.