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

sayantn · sayantn · commit 9eddf137d66c · 2025-05-08T21:21:00.000+05:30
diff --git a/compiler/rustc_codegen_gcc/src/type_of.rs b/compiler/rustc_codegen_gcc/src/type_of.rs
@@ -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)
diff --git a/compiler/rustc_codegen_llvm/src/abi.rs b/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,7 @@ 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]) -> &'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 +326,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]) -> &'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 };
 
+        // todo(sayantn): a better way is to look at the `link_name` instead of the function name, because function name can be "faked" using `#[export_name]`
+        let llvm_intrinsic = name.starts_with(b"llvm.")
+            && !self.c_variadic
+            && self.conv == Conv::C
+            && !self.can_unwind;
+        let amx_intrinsic =
+            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 +408,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 {
diff --git a/compiler/rustc_codegen_llvm/src/builder.rs b/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 {
diff --git a/compiler/rustc_codegen_llvm/src/declare.rs b/compiler/rustc_codegen_llvm/src/declare.rs
@@ -158,7 +158,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()),
         );
         fn_abi.apply_attrs_llfn(self, llfn, instance);
 
diff --git a/compiler/rustc_codegen_llvm/src/intrinsic.rs b/compiler/rustc_codegen_llvm/src/intrinsic.rs
@@ -1090,7 +1090,7 @@ fn gen_fn<'a, 'll, 'tcx>(
     codegen: &mut dyn FnMut(Builder<'a, 'll, 'tcx>),
 ) -> (&'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 llty = fn_abi.llvm_type(cx, name.as_ref());
     let llfn = cx.declare_fn(name, fn_abi, None);
     cx.set_frame_pointer_type(llfn);
     cx.apply_target_cpu_attr(llfn);
diff --git a/compiler/rustc_codegen_llvm/src/llvm/ffi.rs b/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;
 
diff --git a/compiler/rustc_codegen_llvm/src/type_.rs b/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,12 @@ 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 {
+        fn_abi.llvm_type(self, llvm::get_value_name(fn_ptr))
     }
     fn fn_ptr_backend_type(&self, fn_abi: &FnAbi<'tcx, Ty<'tcx>>) -> &'ll Type {
         fn_abi.ptr_to_llvm_type(self)
diff --git a/compiler/rustc_codegen_ssa/src/mir/block.rs b/compiler/rustc_codegen_ssa/src/mir/block.rs
@@ -187,7 +187,7 @@ impl<'a, 'tcx> TerminatorCodegenHelper<'tcx> {
 
         // If there is a cleanup block and the function we're calling can unwind, then
         // do an invoke, otherwise do a call.
-        let fn_ty = bx.fn_decl_backend_type(fn_abi);
+        let fn_ty = bx.fn_decl_backend_type(fn_abi, fn_ptr);
 
         let fn_attrs = if bx.tcx().def_kind(fx.instance.def_id()).has_codegen_attrs() {
             Some(bx.tcx().codegen_fn_attrs(fx.instance.def_id()))
@@ -1806,7 +1806,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
         if is_call_from_compiler_builtins_to_upstream_monomorphization(bx.tcx(), instance) {
             bx.abort();
         } else {
-            let fn_ty = bx.fn_decl_backend_type(fn_abi);
+            let fn_ty = bx.fn_decl_backend_type(fn_abi, fn_ptr);
 
             let llret = bx.call(fn_ty, None, Some(fn_abi), fn_ptr, &[], funclet.as_ref(), None);
             bx.apply_attrs_to_cleanup_callsite(llret);
diff --git a/compiler/rustc_codegen_ssa/src/mir/rvalue.rs b/compiler/rustc_codegen_ssa/src/mir/rvalue.rs
@@ -779,7 +779,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                     };
                     let fn_ptr = bx.get_fn_addr(instance);
                     let fn_abi = bx.fn_abi_of_instance(instance, ty::List::empty());
-                    let fn_ty = bx.fn_decl_backend_type(fn_abi);
+                    let fn_ty = bx.fn_decl_backend_type(fn_abi, fn_ptr);
                     let fn_attrs = if bx.tcx().def_kind(instance.def_id()).has_codegen_attrs() {
                         Some(bx.tcx().codegen_fn_attrs(instance.def_id()))
                     } else {
diff --git a/compiler/rustc_codegen_ssa/src/size_of_val.rs b/compiler/rustc_codegen_ssa/src/size_of_val.rs
@@ -67,7 +67,7 @@ pub fn size_and_align_of_dst<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
             // Generate the call. Cannot use `do_call` since we don't have a MIR terminator so we
             // can't create a `TerminationCodegenHelper`. (But we are in good company, this code is
             // duplicated plenty of times.)
-            let fn_ty = bx.fn_decl_backend_type(fn_abi);
+            let fn_ty = bx.fn_decl_backend_type(fn_abi, llfn);
 
             bx.call(
                 fn_ty,
diff --git a/compiler/rustc_codegen_ssa/src/traits/type_.rs b/compiler/rustc_codegen_ssa/src/traits/type_.rs
@@ -96,7 +96,11 @@ pub trait LayoutTypeCodegenMethods<'tcx>: BackendTypes {
     /// such as when it's stack-allocated or when it's being loaded or stored.
     fn backend_type(&self, layout: TyAndLayout<'tcx>) -> Self::Type;
     fn cast_backend_type(&self, ty: &CastTarget) -> Self::Type;
-    fn fn_decl_backend_type(&self, fn_abi: &FnAbi<'tcx, Ty<'tcx>>) -> Self::Type;
+    fn fn_decl_backend_type(
+        &self,
+        fn_abi: &FnAbi<'tcx, Ty<'tcx>>,
+        fn_ptr: Self::Value,
+    ) -> Self::Type;
     fn fn_ptr_backend_type(&self, fn_abi: &FnAbi<'tcx, Ty<'tcx>>) -> Self::Type;
     fn reg_backend_type(&self, ty: &Reg) -> Self::Type;
     /// The backend type used for a rust type when it's in an SSA register.
diff --git a/compiler/rustc_target/src/target_features.rs b/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.

Original file line number	Diff line number	Diff line change
`@@ -1435,7 +1435,12 @@ impl<'a, 'll, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'll, 'tcx> {`
`1435`	`1435`	`if let Some(fn_abi) = fn_abi {`
`1436`	`1436`	`fn_abi.apply_attrs_callsite(self, call);`
`1437`	`1437`	`}`
`1438`		`- call`
	`1438`	`+`
	`1439`	`+ if self.cx.type_kind(self.cx.val_ty(call)) == TypeKind::X86_AMX {`
	`1440`	`+ self.bitcast(call, self.cx.type_vector(self.cx.type_i32(), 256))`
	`1441`	`+ } else {`
	`1442`	`+ call`
	`1443`	`+ }`
`1439`	`1444`	`}`
`1440`	`1445`
`1441`	`1446`	`fn zext(&mut self, val: &'ll Value, dest_ty: &'ll Type) -> &'ll Value {`
Original file line number	Diff line number	Diff line change
`@@ -154,6 +154,10 @@ impl<'ll, CX: Borrow<SCx<'ll>>> GenericCx<'ll, CX> {`
`154`	`154`	`)`
`155`	`155`	`}`
`156`	`156`	`}`
	`157`	`+`
	`158`	`+ pub(crate) fn type_x86amx(&self) -> &'ll Type {`
	`159`	`+ unsafe { llvm::LLVMX86AMXTypeInContext(self.llcx()) }`
	`160`	`+ }`
`157`	`161`	`}`
`158`	`162`
`159`	`163`	`impl<'ll, CX: Borrow<SCx<'ll>>> BaseTypeCodegenMethods for GenericCx<'ll, CX> {`
`@@ -284,8 +288,12 @@ impl<'ll, 'tcx> LayoutTypeCodegenMethods<'tcx> for CodegenCx<'ll, 'tcx> {`
`284`	`288`	`fn cast_backend_type(&self, ty: &CastTarget) -> &'ll Type {`
`285`	`289`	`ty.llvm_type(self)`
`286`	`290`	`}`
`287`		`- fn fn_decl_backend_type(&self, fn_abi: &FnAbi<'tcx, Ty<'tcx>>) -> &'ll Type {`
`288`		`- fn_abi.llvm_type(self)`
	`291`	`+ fn fn_decl_backend_type(`
	`292`	`+ &self,`
	`293`	`+ fn_abi: &FnAbi<'tcx, Ty<'tcx>>,`
	`294`	`+ fn_ptr: &'ll Value,`
	`295`	`+ ) -> &'ll Type {`
	`296`	`+ fn_abi.llvm_type(self, llvm::get_value_name(fn_ptr))`
`289`	`297`	`}`
`290`	`298`	`fn fn_ptr_backend_type(&self, fn_abi: &FnAbi<'tcx, Ty<'tcx>>) -> &'ll Type {`
`291`	`299`	`fn_abi.ptr_to_llvm_type(self)`