Add auto-bitcasts for AMX

sayantn · sayantn · commit 74e65ab9aca8 · 2025-05-20T23:15:56.000+05:30
diff --git a/compiler/rustc_codegen_llvm/src/abi.rs b/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::*;
@@ -412,7 +413,35 @@ impl<'ll, 'tcx> FnAbiLlvmExt<'ll, 'tcx> for FnAbi<'tcx, Ty<'tcx>> {
             let expected_return_ty = cx.get_return_type(fn_ty);
             let expected_argument_tys = cx.func_params_types(fn_ty);
 
-            let equate_ty = |rust_ty, llvm_ty| rust_ty == llvm_ty;
+            let equate_ty = |rust_ty, llvm_ty| {
+                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
+                        }
+                    }
+                    _ => false,
+                }
+            };
 
             if actual_argument_tys.len() != expected_argument_tys.len() {
                 todo!("A very friendly error msg")
diff --git a/compiler/rustc_codegen_llvm/src/builder.rs b/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,51 @@ 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);
+
+        let intrinsic = llvm::Intrinsic::lookup(b"llvm.x86.cast.vector.to.tile".as_ref()).unwrap();
+        let (fn_ty, f) = self.cx.get_or_declare_intrinsic(intrinsic, &[vector_type]);
+        unsafe {
+            llvm::LLVMBuildCallWithOperandBundles(
+                self.llbuilder,
+                fn_ty,
+                f,
+                [val].as_ptr().cast(),
+                1,
+                [].as_ptr(),
+                0,
+                c"".as_ptr(),
+            )
+        }
+    }
+
+    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);
+
+        let intrinsic = llvm::Intrinsic::lookup(b"llvm.x86.cast.tile.to.vector").unwrap();
+        let (fn_ty, f) = self.cx.get_or_declare_intrinsic(intrinsic, &[vector_type]);
+        unsafe {
+            llvm::LLVMBuildCallWithOperandBundles(
+                self.llbuilder,
+                fn_ty,
+                f,
+                [val].as_ptr().cast(),
+                1,
+                [].as_ptr(),
+                0,
+                c"".as_ptr(),
+            )
+        }
+    }
+
     pub(crate) fn ret_void(&mut self) {
         llvm::LLVMBuildRetVoid(self.llbuilder);
     }
@@ -349,7 +394,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 +426,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) {
@@ -1404,7 +1451,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 {
@@ -1434,8 +1481,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 {
@@ -1596,7 +1645,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,
@@ -1627,7 +1676,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
                 }
@@ -1708,6 +1761,31 @@ impl<'a, 'll, 'tcx> Builder<'a, 'll, 'tcx> {
         self.call(self.type_func(&[src_ty], dest_ty), None, None, f, &[val], None, None)
     }
 
+    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,
@@ -1737,7 +1815,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 {
@@ -1770,8 +1848,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.
diff --git a/compiler/rustc_codegen_llvm/src/type_.rs b/compiler/rustc_codegen_llvm/src/type_.rs
@@ -175,7 +175,6 @@ impl<'ll, CX: Borrow<SCx<'ll>>> GenericCx<'ll, CX> {
         unsafe { llvm::LLVMX86AMXTypeInContext(self.llcx()) }
     }
 
-    #[expect(unused)]
     pub(crate) fn get_or_declare_intrinsic(
         &self,
         intrinsic: llvm::Intrinsic,
diff --git a/compiler/rustc_monomorphize/src/mono_checks/abi_check.rs b/compiler/rustc_monomorphize/src/mono_checks/abi_check.rs
@@ -49,6 +49,11 @@ fn do_check_simd_vector_abi<'tcx>(
             // Find the first feature that provides at least this vector size.
             let feature = match feature_def.iter().find(|(bits, _)| size.bits() <= *bits) {
                 Some((_, feature)) => feature,
+                None if matches!(&*tcx.sess.target.arch, "x86" | "x86_64")
+                    && size.bits() == 8192 =>
+                {
+                    "amx-tile"
+                }
                 None => {
                     let (span, _hir_id) = loc();
                     tcx.dcx().emit_err(errors::AbiErrorUnsupportedVectorType {
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")];
 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
`@@ -175,7 +175,6 @@ impl<'ll, CX: Borrow<SCx<'ll>>> GenericCx<'ll, CX> {`
`175`	`175`	`unsafe { llvm::LLVMX86AMXTypeInContext(self.llcx()) }`
`176`	`176`	`}`
`177`	`177`
`178`		`- #[expect(unused)]`
`179`	`178`	`pub(crate) fn get_or_declare_intrinsic(`
`180`	`179`	`&self,`
`181`	`180`	`intrinsic: llvm::Intrinsic,`