Add bypass for x86amx

Author: sayantn

compiler/rustc_codegen_llvm/src/abi.rs

@@ -361,6 +361,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

@@ -1650,6 +1650,15 @@ impl<'a, 'll, 'tcx> Builder<'a, 'll, 'tcx> {
         }
 
         match self.type_kind(llvm_ty) {
+            TypeKind::X86_AMX => {
+                let base_name = if is_argument {
+                    "llvm.x86.cast.vector.to.tile"
+                } else {
+                    "llvm.x86.cast.tile.to.vector"
+                };
+
+                self.call_intrinsic(base_name, &[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

@@ -927,6 +927,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);
@@ -1039,6 +1040,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

@@ -1085,6 +1085,9 @@ unsafe extern "C" {
     pub(crate) fn LLVMTokenTypeInContext(C: &Context) -> &Type;
     pub(crate) fn LLVMMetadataTypeInContext(C: &Context) -> &Type;
 
+    // X86-specific type for AMX
+    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

@@ -181,6 +181,10 @@ impl<'ll, CX: Borrow<SCx<'ll>>> GenericCx<'ll, CX> {
     pub(crate) fn type_bf16(&self) -> &'ll Type {
         unsafe { llvm::LLVMBFloatTypeInContext(self.llcx()) }
     }
+
+    pub(crate) fn type_x86amx(&self) -> &'ll Type {
+        unsafe { llvm::LLVMX86AMXTypeInContext(self.llcx()) }
+    }
 }
 
 impl<'ll, CX: Borrow<SCx<'ll>>> BaseTypeCodegenMethods for GenericCx<'ll, CX> {