[Mlir-commits] [llvm] [mlir] [LLVM][NVPTX] Add Rubin extensions for G2S Tensor intrinsics (PR #220029)

[Durgadoss R]

================
@@ -966,6 +1002,148 @@ multiclass NVPTX_CP_ASYNC_BULK_TENSOR_S2G_OVERRIDE_INTR<int dim, string mode,
   }
 }
 
+// TMA Copy from Global to Shared memory (Cluster) with Override Global Address.
+multiclass NVPTX_CP_ASYNC_BULK_TENSOR_G2S_OVERRIDE_INTR<int dim, string mode, string override> {
+  defvar base = TMA_TENSOR_OVERRIDE_BASE<dim, mode, override>;
+  defvar num_dims = TMA_DIMS_UTIL<dim, mode>.num_dims;
+  defvar suffix = base.override_suffix;
+  // The override_addr_dim / override_addr_dim_stride variants support only 32-bit multicast.
+  defvar supports_mc16 = base.is_override_addr;
+  defvar asm_operands = " [$dst], [$tmap, $override_addr, " #
+                        base.tensor_size_component # "{{" # base.dims_str #
+                        "}}], [$mbar]" # base.im2col_component;
+
+  defvar inst_name = "cp.async.bulk.tensor." # num_dims #
+                     "d.shared::cluster.global." # !subst("_", "::", mode) #
+                     ".mbarrier::complete_tx::bytes" # "$validate" # "$cg";
+
+  defvar ins_dag_base = !con(
+      (ins ADDR:$dst, ADDR:$mbar, B64:$tmap, B64:$override_addr),
+      base.tensor_size_dag, base.lower_stride_dag, base.upper_stride_dag,
+      base.dims_dag, base.im2col_dag);
+  defvar ins_tail = (ins TMAValidateDataFlags:$validate, CTAGroupFlags:$cg);
+
+  defvar ins_dag         = !con(ins_dag_base, ins_tail);
+  defvar ins_dag_ch      = !con(ins_dag_base, (ins B64:$ch), ins_tail);
+  defvar ins_dag_mc16    = !con(ins_dag_base, (ins B16:$mc), ins_tail);
+  defvar ins_dag_mc16_ch = !con(ins_dag_base, (ins B16:$mc), (ins B64:$ch), ins_tail);
+  defvar ins_dag_mc32    = !con(ins_dag_base, (ins B32:$mc), ins_tail);
+  defvar ins_dag_mc32_ch = !con(ins_dag_base, (ins B32:$mc), (ins B64:$ch), ins_tail);
+
+  defvar intr = !cast<Intrinsic>(
+      "int_nvvm_cp_async_bulk_tensor_g2s_" # mode # "_" # override # "_" #
+      num_dims # "d");
+
+  defvar intr_dag_base = !con(
+      (intr addr:$dst, addr:$mbar, B64:$tmap, B64:$override_addr),
+      !setdagop(base.tensor_size_dag, intr),
+      !setdagop(base.lower_stride_dag, intr),
+      !setdagop(base.upper_stride_dag, intr),
+      !setdagop(base.dims_dag, intr),
+      !setdagop(base.im2col_dag, intr));
+  defvar intr_tail = (intr tma_cta_group_imm_any:$cg, tma_validate_data_imm:$validate);
+
+  defvar intr_dag_no_mc16_ch      = !con(intr_dag_base, (intr (i16 srcvalue), (i64 srcvalue), 0,  0),  intr_tail);
+  defvar intr_dag_with_mc16       = !con(intr_dag_base, (intr B16:$mc, (i64 srcvalue), -1, 0),  intr_tail);
+  defvar intr_dag_no_mc16_with_ch = !con(intr_dag_base, (intr (i16 srcvalue), B64:$ch, 0, -1),  intr_tail);
+  defvar intr_dag_with_mc16_ch    = !con(intr_dag_base, (intr B16:$mc, B64:$ch, -1, -1), intr_tail);
+  defvar intr_dag_no_mc32_ch      = !con(intr_dag_base, (intr (i32 srcvalue), (i64 srcvalue), 0,  0),  intr_tail);
+  defvar intr_dag_with_mc32       = !con(intr_dag_base, (intr i32:$mc, (i64 srcvalue), -1, 0),  intr_tail);
+  defvar intr_dag_no_mc32_with_ch = !con(intr_dag_base, (intr (i32 srcvalue), B64:$ch, 0, -1),  intr_tail);
+  defvar intr_dag_with_mc32_ch    = !con(intr_dag_base, (intr i32:$mc, B64:$ch, -1, -1), intr_tail);
+
+  let Predicates = [hasRubinFamilySupport] in {
+    if supports_mc16 then {
+      def _NO_MC16_CH : NVPTXInst<(outs), ins_dag,
+                          inst_name # suffix # asm_operands # ";",
+                          [intr_dag_no_mc16_ch]>;
+      def _MC16 : NVPTXInst<(outs), ins_dag_mc16,
+                    inst_name # ".multicast::cluster" # suffix # asm_operands # ", $mc;",
+                    [intr_dag_with_mc16]>;
+      def _NO_MC16_WITH_CH : NVPTXInst<(outs), ins_dag_ch,
+                              inst_name # ".L2::cache_hint" # suffix # asm_operands # ", $ch;",
+                              [intr_dag_no_mc16_with_ch]>;
+      def _MC16_CH : NVPTXInst<(outs), ins_dag_mc16_ch,
+                      inst_name # ".multicast::cluster.L2::cache_hint" # suffix # asm_operands # ", $mc, $ch;",
+                      [intr_dag_with_mc16_ch]>;
+    }
+    def _NO_MC32_CH : NVPTXInst<(outs), ins_dag,
----------------
durga4github wrote:

I wonder if there is a way to use one opcode for both i16 or i32 mc-mask using some wrappers .. (like how we do SelectAddr). This can be a follow-up too.

https://github.com/llvm/llvm-project/pull/220029