[Mlir-commits] [llvm] [mlir] [NVPTX] Add "exclusive" intrinsic variants for tcgen05.alloc/dealloc (PR #216016)

[Durgadoss R]

================
@@ -5735,30 +5735,85 @@ let Predicates = [SM90] in {
 def EXIT : NullaryInst<"exit", int_nvvm_exit>;
 
 // Tcgen05 intrinsics
-let isConvergent = true in {
-let Predicates = [hasTcgen05InstSupport] in {
-multiclass TCGEN05_ALLOC_INTR<string AS, string num, Intrinsic Intr> {
-  def "" : BasicNVPTXInst<(outs),
-             (ins ADDR:$dst, B32:$ncols),
-             "tcgen05.alloc.cta_group::" # num # ".sync.aligned" # AS # ".b32",
-             [(Intr addr:$dst, B32:$ncols)]>;
+multiclass TCGEN05_ALLOC_PATFRAGS<Intrinsic Intr> {
+  defvar frag_pat = (Intr node:$dst, node:$ncols);
+  def _GENERIC
+      : PatFrag<!setdagop(frag_pat, ops), frag_pat, AS_match.generic>;
+  def _SHARED
+      : PatFrag<!setdagop(frag_pat, ops), frag_pat, AS_match.shared>;
+}
+
+defm TCGEN05_ALLOC_CG1_PAT
+    : TCGEN05_ALLOC_PATFRAGS<int_nvvm_tcgen05_alloc_cg1>;
+defm TCGEN05_ALLOC_CG2_PAT
+    : TCGEN05_ALLOC_PATFRAGS<int_nvvm_tcgen05_alloc_cg2>;
+defm TCGEN05_ALLOC_EXCLUSIVE_CG1_PAT
+    : TCGEN05_ALLOC_PATFRAGS<int_nvvm_tcgen05_alloc_exclusive_cg1>;
+defm TCGEN05_ALLOC_EXCLUSIVE_CG2_PAT
+    : TCGEN05_ALLOC_PATFRAGS<int_nvvm_tcgen05_alloc_exclusive_cg2>;
+
+multiclass TCGEN05_ALLOC_INTR<string num, PatFrag GenericPat,
+                              PatFrag SharedPat,
+                              bit IsExclusive = 0> {
+  defvar exclusive = !if(IsExclusive, ".exclusive", "");
+  defvar Pred = !if(IsExclusive,
+                    [hasTcgen05InstSupport, PTX94],
+                    [hasTcgen05InstSupport]);
+
+  let isConvergent = true, Predicates = Pred in {
+    def _GENERIC : BasicNVPTXInst<(outs), (ins ADDR:$dst, B32:$ncols),
+                     "tcgen05.alloc" # exclusive # ".cta_group::" # num #
+                     ".sync.aligned.b32",
+                     [(GenericPat addr:$dst, B32:$ncols)]>;
+    def _SHARED : BasicNVPTXInst<(outs), (ins ADDR:$dst, B32:$ncols),
+                    "tcgen05.alloc" # exclusive # ".cta_group::" # num #
----------------
durga4github wrote:

ok, we should leverage the cta_group printer methods for all the `cta_group` in these intrinsics. But I think that can be a separate NFC PR. (so, for now, let us go with what we have here)

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