Add bypass for x86amx

Author: sayantn

compiler/rustc_codegen_llvm/src/abi.rs

@@ -368,6 +368,25 @@ impl<'ll, CX: Borrow<SCx<'ll>>> GenericCx<'ll, CX> {
         }
 
         match self.type_kind(llvm_ty) {
+            TypeKind::X86_AMX if self.type_kind(rust_ty) == TypeKind::Vector => {
+                let element_count = self.vector_length(rust_ty);
+                let element_ty = self.element_type(rust_ty);
+
+                let element_size_bits = match self.type_kind(element_ty) {
+                    TypeKind::Half => 16,
+                    TypeKind::Float => 32,
+                    TypeKind::Double => 64,
+                    TypeKind::FP128 => 128,
+                    TypeKind::Integer => self.int_width(element_ty),
+                    TypeKind::Pointer => self.int_width(self.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;
+
+                vector_size_bits == 8192
+            }
             TypeKind::BFloat => rust_ty == self.type_i16(),
             TypeKind::Vector => {
                 let llvm_element_count = self.vector_length(llvm_ty) as u64;

compiler/rustc_codegen_llvm/src/builder.rs

@@ -492,10 +492,10 @@ impl<'a, 'll, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'll, 'tcx> {
         let width = size.bits();
 
         if oop == OverflowOp::Sub && !signed {
-                    let sub = self.sub(lhs, rhs);
-                    let cmp = self.icmp(IntPredicate::IntULT, lhs, rhs);
-                    return (sub, cmp);
-                }
+            let sub = self.sub(lhs, rhs);
+            let cmp = self.icmp(IntPredicate::IntULT, lhs, rhs);
+            return (sub, cmp);
+        }
 
         let oop_str = match oop {
             OverflowOp::Add => "add",
@@ -1641,6 +1641,13 @@ impl<'a, 'll, 'tcx> Builder<'a, 'll, 'tcx> {
         }
 
         match self.type_kind(llvm_ty) {
+            TypeKind::X86_AMX => {
+                if is_argument {
+                    self.call_intrinsic("llvm.x86.cast.vector.to.tile", &[rust_ty], &[val])
+                } else {
+                    self.call_intrinsic("llvm.x86.cast.tile.to.vector", &[rust_ty], &[val])
+                }
+            }
             TypeKind::Vector if self.element_type(llvm_ty) == self.type_i1() => {
                 if is_argument {
                     self.trunc_int_to_i1_vector(val, dest_ty)

compiler/rustc_codegen_llvm/src/context.rs

@@ -925,6 +925,7 @@ impl<'ll> CodegenCx<'ll, '_> {
         let t_isize = self.type_isize();
         let t_metadata = self.type_metadata();
         let t_token = self.type_token();
+        let x86amx = self.type_x86amx();
 
         ifn!("llvm.wasm.get.exception", fn(t_token) -> ptr);
         ifn!("llvm.wasm.get.ehselector", fn(t_token) -> t_i32);
@@ -1038,6 +1039,9 @@ impl<'ll> CodegenCx<'ll, '_> {
         ifn!("llvm.masked.gather", fn(1, t_i32, same_width_vector(0, i1), 0) -> 0);
         ifn!("llvm.masked.scatter", fn(0, 1, t_i32, same_width_vector(0, i1)) -> void);
 
+        ifn!("llvm.x86.cast.vector.to.tile", fn(0) -> x86amx);
+        ifn!("llvm.x86.cast.tile.to.vector", fn(x86amx) -> 0);
+
         bug!("Unknown intrinsic: `{base_name}`")
     }
 

compiler/rustc_codegen_llvm/src/llvm/ffi.rs

@@ -1084,6 +1084,8 @@ unsafe extern "C" {
     pub(crate) fn LLVMTokenTypeInContext(C: &Context) -> &Type;
     pub(crate) fn LLVMMetadataTypeInContext(C: &Context) -> &Type;
 
+    pub(crate) fn LLVMX86AMXTypeInContext(C: &Context) -> &Type;
+
     // Operations on all values
     pub(crate) fn LLVMTypeOf(Val: &Value) -> &Type;
     pub(crate) fn LLVMGetValueName2(Val: &Value, Length: *mut size_t) -> *const c_char;

compiler/rustc_codegen_llvm/src/type_.rs

@@ -66,6 +66,10 @@ impl<'ll, CX: Borrow<SCx<'ll>>> GenericCx<'ll, CX> {
         unsafe { llvm::LLVMMetadataTypeInContext(self.llcx()) }
     }
 
+    pub(crate) fn type_x86amx(&self) -> &'ll Type {
+        unsafe { llvm::LLVMX86AMXTypeInContext(self.llcx()) }
+    }
+
     ///x Creates an integer type with the given number of bits, e.g., i24
     pub(crate) fn type_ix(&self, num_bits: u64) -> &'ll Type {
         unsafe { llvm::LLVMIntTypeInContext(self.llcx(), num_bits as c_uint) }