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

@@ -1731,6 +1731,31 @@ impl<'a, 'll, 'tcx> Builder<'a, 'll, 'tcx> {
         }
 
         match self.type_kind(llvm_ty) {
+            TypeKind::X86_AMX => {
+                let vector_length = self.vector_length(rust_ty);
+                let element_ty = self.element_type(rust_ty);
+                let element_ty_str = match self.type_kind(element_ty) {
+                    TypeKind::Half => "f16",
+                    TypeKind::Float => "f32",
+                    TypeKind::Double => "f64",
+                    TypeKind::FP128 => "f128",
+                    TypeKind::Integer => &format!("i{}", self.int_width(element_ty)),
+                    TypeKind::Pointer => "p0",
+                    _ => bug!(
+                        "Vector element type `{element_ty:?}` not one of integer, float or pointer"
+                    ),
+                };
+
+                let base_name = if is_argument {
+                    "llvm.x86.cast.vector.to.tile"
+                } else {
+                    "llvm.x86.cast.tile.to.vector"
+                };
+                let llvm_intrinsic = format!("{base_name}.v{vector_length}{element_ty_str}");
+                let fn_ty = self.type_func(&[src_ty], dest_ty);
+                let f = self.declare_cfn(&llvm_intrinsic, llvm::UnnamedAddr::No, fn_ty);
+                self.call(fn_ty, None, None, f, &[val], None, None)
+            }
             TypeKind::Vector if self.element_type(llvm_ty) == self.type_i1() => {
                 if is_argument {
                     self.trunc_int_to_i1_vector(val, dest_ty)