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

llvmlistbot at llvm.org llvmlistbot at llvm.org
Thu Aug 13 04:54:02 PDT 2026


llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-llvm

Author: Dharuni R Acharya (DharuniRAcharya)

<details>
<summary>Changes</summary>

This patch adds `exclusive` intrinsic variants for `tcgen05.alloc` and `tcgen05.dealloc`, 
which allow exclusive ownership to the allocation.

- `tcgen05.alloc{.exclusive}.cta_group.sync.aligned{.shared::cta}.b32`
- `tcgen05.dealloc{.exclusive}.cta_group.sync.aligned.b32`

PTX ISA Reference: https://docs.nvidia.com/cuda/developer-preview/13.4/parallel-thread-execution/index.html#tcgen05-instructions-tcgen05-alloc-dealloc-relinquish-alloc-permit

---

Patch is 31.01 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/216016.diff


11 Files Affected:

- (modified) llvm/docs/NVPTXUsage.md (+22-16) 
- (modified) llvm/include/llvm/IR/IntrinsicsNVVM.td (+13-17) 
- (modified) llvm/lib/IR/AutoUpgrade.cpp (+21) 
- (modified) llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp (+12) 
- (modified) llvm/lib/Target/NVPTX/NVPTXIntrinsics.td (+73-18) 
- (modified) llvm/test/Assembler/auto_upgrade_nvvm_intrinsics.ll (+18) 
- (added) llvm/test/CodeGen/NVPTX/tcgen05-alloc-dealloc-exclusive.ll (+165) 
- (modified) llvm/test/CodeGen/NVPTX/tcgen05-alloc.ll (+8-8) 
- (modified) mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td (+1-1) 
- (modified) mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp (+4-9) 
- (modified) mlir/test/Target/LLVMIR/nvvm/tcgen05-alloc.mlir (+4-4) 


``````````diff
diff --git a/llvm/docs/NVPTXUsage.md b/llvm/docs/NVPTXUsage.md
index 2541d504c6e10..d576f82b36b1b 100644
--- a/llvm/docs/NVPTXUsage.md
+++ b/llvm/docs/NVPTXUsage.md
@@ -2616,24 +2616,27 @@ For more information on tensor-memory load/store instructions, refer to [PTX ISA
 ##### Syntax:
 
 ```llvm
-declare void @llvm.nvvm.tcgen05.alloc.cg1(ptr %dst, i32 %ncols)
-declare void @llvm.nvvm.tcgen05.alloc.cg2(ptr %dst, i32 %ncols)
-declare void @llvm.nvvm.tcgen05.alloc.shared.cg1(ptr addrspace(3) %dst, i32 %ncols)
-declare void @llvm.nvvm.tcgen05.alloc.shared.cg2(ptr addrspace(3) %dst, i32 %ncols)
+declare void @llvm.nvvm.tcgen05.alloc.{cg1,cg2}.p0(ptr %dst, i32 %ncols)
+declare void @llvm.nvvm.tcgen05.alloc.{cg1,cg2}.p3(ptr addrspace(3) %dst, i32 %ncols)
+
+declare void @llvm.nvvm.tcgen05.alloc.exclusive.{cg1,cg2}.p0(ptr %dst, i32 %ncols)
+declare void @llvm.nvvm.tcgen05.alloc.exclusive.{cg1,cg2}.p3(ptr addrspace(3) %dst, i32 %ncols)
 ```
 
 ##### Overview:
 
 The '`@llvm.nvvm.tcgen05.alloc.*`' intrinsics correspond to the
-`tcgen05.alloc.cta_group*.sync.aligned.b32` family of PTX instructions.
-The `tcgen05.alloc` is a potentially blocking instruction which dynamically
-allocates the specified number of columns in the Tensor Memory and writes the
-address of the allocated Tensor Memory into shared memory at the location
-specified by `%dst`. The 32-bit operand `%ncols` specifies the number of
-columns to be allocated and it must be a power-of-two. The `.shared` variant
-explicitly uses shared memory address space for the `%dst` operand. The
-`.cg1` and `.cg2` variants generate `cta_group::1` and `cta_group::2`
-variants of the instruction respectively.
+`tcgen05.alloc{.exclusive}.cta_group*.sync.aligned.b32` family of PTX 
+instructions. The `tcgen05.alloc` is a potentially blocking instruction which
+dynamically allocates the specified number of columns in the Tensor Memory 
+and writes the address of the allocated Tensor Memory into shared memory at the
+location specified by `%dst`. The 32-bit operand `%ncols` specifies the number
+of columns to be allocated and it must be a power-of-two for non-exclusive
+allocations and it must be a multiple of 32 for exclusive allocations. The
+overloaded pointer argument may use generic or shared memory address space;
+a shared pointer emits the `.shared::cta` qualifier. The `.cg1` and `.cg2`
+variants generate `cta_group::1` and `cta_group::2` variants of the instruction
+respectively. The `.exclusive` variants claim ownership of the allocation permit.
 
 For more information, refer to the [PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#tensor-memory-allocation-and-management-instructions).
 
@@ -2642,19 +2645,22 @@ For more information, refer to the [PTX ISA](https://docs.nvidia.com/cuda/parall
 ##### Syntax:
 
 ```llvm
-declare void @llvm.nvvm.tcgen05.dealloc.cg1(ptr addrspace(6) %tmem_addr, i32 %ncols)
-declare void @llvm.nvvm.tcgen05.dealloc.cg2(ptr addrspace(6) %tmem_addr, i32 %ncols)
+declare void @llvm.nvvm.tcgen05.dealloc.{cg1,cg2}(ptr addrspace(6) %tmem_addr, i32 %ncols)
+
+declare void @llvm.nvvm.tcgen05.dealloc.exclusive.{cg1,cg2}(ptr addrspace(6) %tmem_addr, i32 %ncols)
 ```
 
 ##### Overview:
 
 The '`@llvm.nvvm.tcgen05.dealloc.*`' intrinsics correspond to the
-`tcgen05.dealloc.*` set of PTX instructions. The `tcgen05.dealloc`
+`tcgen05.dealloc{.exclusive}.*` set of PTX instructions. The `tcgen05.dealloc`
 instructions deallocates the Tensor Memory specified by the Tensor Memory
 address `%tmem_addr`. The operand `%tmem_addr` must point to a previous
 Tensor Memory allocation. The 32-bit operand `%ncols` specifies the number
 of columns to be de-allocated. The `.cg1` and `.cg2` variants generate
 `cta_group::1` and `cta_group::2` variants of the instruction respectively.
+Memory must be deallocated with `.exclusive` if and only if it was allocated
+with `.exclusive`.
 
 For more information, refer to the [PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#tensor-memory-allocation-and-management-instructions).
 
diff --git a/llvm/include/llvm/IR/IntrinsicsNVVM.td b/llvm/include/llvm/IR/IntrinsicsNVVM.td
index e6761a852c93f..79052a5099caf 100644
--- a/llvm/include/llvm/IR/IntrinsicsNVVM.td
+++ b/llvm/include/llvm/IR/IntrinsicsNVVM.td
@@ -3082,23 +3082,19 @@ def int_nvvm_griddepcontrol_wait : Intrinsic<[], [], [IntrNoMem, IntrHasSideEffe
 // Tcgen05 alloc/dealloc related intrinsics
 
 foreach cta_group = ["cg1", "cg2"] in {
-  def int_nvvm_tcgen05_alloc_ # cta_group : Intrinsic<[],
-    [llvm_ptr_ty,        // dst_ptr
-     llvm_i32_ty] ,      // num_columns
-    [IntrConvergent, IntrInaccessibleMemOrArgMemOnly,
-     WriteOnly<ArgIndex<0>>, NoCapture<ArgIndex<0>>]>;
-
-  def int_nvvm_tcgen05_alloc_shared_ # cta_group : Intrinsic<[],
-    [llvm_shared_ptr_ty, // dst_ptr
-     llvm_i32_ty],       // num_columns
-    [IntrConvergent, IntrInaccessibleMemOrArgMemOnly,
-     WriteOnly<ArgIndex<0>>, NoCapture<ArgIndex<0>>]>;
-
-  def int_nvvm_tcgen05_dealloc_ # cta_group : Intrinsic<[],
-    [llvm_tmem_ptr_ty,   // tmem_addr
-     llvm_i32_ty],       // num_columns
-    [IntrConvergent, IntrArgMemOnly,
-     NoCapture<ArgIndex<0>>]>;
+  foreach exclusive = ["", "exclusive_"] in {
+    def int_nvvm_tcgen05_alloc_ # exclusive # cta_group : Intrinsic<[],
+      [llvm_anyptr_ty,     // dst_ptr
+       llvm_i32_ty],       // num_columns
+      [IntrConvergent, IntrInaccessibleMemOrArgMemOnly,
+       WriteOnly<ArgIndex<0>>, NoCapture<ArgIndex<0>>]>;
+
+    def int_nvvm_tcgen05_dealloc_ # exclusive # cta_group : Intrinsic<[],
+      [llvm_tmem_ptr_ty,   // tmem_addr
+       llvm_i32_ty],       // num_columns
+      [IntrConvergent, IntrArgMemOnly,
+       NoCapture<ArgIndex<0>>]>;
+  }
 
   def int_nvvm_tcgen05_relinq_alloc_permit_ # cta_group : Intrinsic<[], [],
     [IntrConvergent, IntrInaccessibleMemOnly]>;
diff --git a/llvm/lib/IR/AutoUpgrade.cpp b/llvm/lib/IR/AutoUpgrade.cpp
index 5b6f50df5d4d0..22a148c5512d7 100644
--- a/llvm/lib/IR/AutoUpgrade.cpp
+++ b/llvm/lib/IR/AutoUpgrade.cpp
@@ -1279,6 +1279,17 @@ shouldUpgradeNVPTXTcgen05CommitSharedIntrinsic(Function *F, StringRef Name) {
   return Intrinsic::not_intrinsic;
 }
 
+static Intrinsic::ID
+shouldUpgradeNVPTXTcgen05AllocSharedIntrinsic(StringRef Name) {
+  if (!Name.consume_front("tcgen05.alloc.shared."))
+    return Intrinsic::not_intrinsic;
+
+  return StringSwitch<Intrinsic::ID>(Name)
+      .Case("cg1", Intrinsic::nvvm_tcgen05_alloc_cg1)
+      .Case("cg2", Intrinsic::nvvm_tcgen05_alloc_cg2)
+      .Default(Intrinsic::not_intrinsic);
+}
+
 static Intrinsic::ID shouldUpgradeNVPTXBF16Intrinsic(StringRef Name) {
   if (Name.consume_front("fma.rn."))
     return StringSwitch<Intrinsic::ID>(Name)
@@ -1811,6 +1822,16 @@ static bool upgradeIntrinsicFunction1(Function *F, Function *&NewFn,
         return true;
       }
 
+      // Upgrade tcgen05.alloc shared variants to anyptr intrinsics.
+      IID = shouldUpgradeNVPTXTcgen05AllocSharedIntrinsic(Name);
+      if (IID != Intrinsic::not_intrinsic) {
+        rename(F);
+        NewFn = Intrinsic::getOrInsertDeclaration(
+            F->getParent(), IID, F->getReturnType(),
+            F->getFunctionType()->params());
+        return true;
+      }
+
       // Upgrade TMA copy G2S Intrinsics
       IID = shouldUpgradeNVPTXTMAG2SIntrinsics(F, Name);
       if (IID != Intrinsic::not_intrinsic) {
diff --git a/llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp b/llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp
index e788b0e44041f..be3078ed835fd 100644
--- a/llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp
+++ b/llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp
@@ -5546,6 +5546,18 @@ void NVPTXTargetLowering::getTgtMemIntrinsic(
     Infos.push_back(Info);
     return;
   }
+  case Intrinsic::nvvm_tcgen05_alloc_cg1:
+  case Intrinsic::nvvm_tcgen05_alloc_cg2:
+  case Intrinsic::nvvm_tcgen05_alloc_exclusive_cg1:
+  case Intrinsic::nvvm_tcgen05_alloc_exclusive_cg2:
+    Info.opc = ISD::INTRINSIC_VOID;
+    Info.memVT = MVT::i32;
+    Info.ptrVal = I.getArgOperand(0);
+    Info.offset = 0;
+    Info.flags = MachineMemOperand::MOStore;
+    Info.align = Align(4);
+    Infos.push_back(Info);
+    return;
   }
 }
 
diff --git a/llvm/lib/Target/NVPTX/NVPTXIntrinsics.td b/llvm/lib/Target/NVPTX/NVPTXIntrinsics.td
index 72b614d84e3e3..a1cfecd4dfe20 100644
--- a/llvm/lib/Target/NVPTX/NVPTXIntrinsics.td
+++ b/llvm/lib/Target/NVPTX/NVPTXIntrinsics.td
@@ -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 #
+                    ".sync.aligned.shared::cta.b32",
+                    [(SharedPat addr:$dst, B32:$ncols)]>;
+  }
 }
 
-defm TCGEN05_ALLOC_CG1 : TCGEN05_ALLOC_INTR<"", "1", int_nvvm_tcgen05_alloc_cg1>;
-defm TCGEN05_ALLOC_CG2 : TCGEN05_ALLOC_INTR<"", "2", int_nvvm_tcgen05_alloc_cg2>;
-
-defm TCGEN05_ALLOC_S64_CG1 : TCGEN05_ALLOC_INTR<".shared::cta", "1", int_nvvm_tcgen05_alloc_shared_cg1>;
-defm TCGEN05_ALLOC_S64_CG2 : TCGEN05_ALLOC_INTR<".shared::cta", "2", int_nvvm_tcgen05_alloc_shared_cg2>;
-
-multiclass TCGEN05_DEALLOC_INTR<string num, Intrinsic Intr> {
-  def "" : BasicNVPTXInst<(outs),
-             (ins B32:$tmem_addr, B32:$ncols),
-             "tcgen05.dealloc.cta_group::" # num # ".sync.aligned.b32",
-             [(Intr B32:$tmem_addr, B32:$ncols)]>;
+defm TCGEN05_ALLOC_CG1
+    : TCGEN05_ALLOC_INTR<"1", TCGEN05_ALLOC_CG1_PAT_GENERIC,
+                         TCGEN05_ALLOC_CG1_PAT_SHARED>;
+defm TCGEN05_ALLOC_CG2
+    : TCGEN05_ALLOC_INTR<"2", TCGEN05_ALLOC_CG2_PAT_GENERIC,
+                         TCGEN05_ALLOC_CG2_PAT_SHARED>;
+
+defm TCGEN05_ALLOC_EXCLUSIVE_CG1
+    : TCGEN05_ALLOC_INTR<"1", TCGEN05_ALLOC_EXCLUSIVE_CG1_PAT_GENERIC,
+                         TCGEN05_ALLOC_EXCLUSIVE_CG1_PAT_SHARED,
+                         /*IsExclusive=*/1>;
+defm TCGEN05_ALLOC_EXCLUSIVE_CG2
+    : TCGEN05_ALLOC_INTR<"2", TCGEN05_ALLOC_EXCLUSIVE_CG2_PAT_GENERIC,
+                         TCGEN05_ALLOC_EXCLUSIVE_CG2_PAT_SHARED,
+                         /*IsExclusive=*/1>;
+
+multiclass TCGEN05_DEALLOC_INTR<string num, Intrinsic Intr,
+                                bit IsExclusive = 0> {
+  defvar exclusive = !if(IsExclusive, ".exclusive", "");
+  defvar Pred = !if(IsExclusive,
+                    [hasTcgen05InstSupport, PTX94],
+                    [hasTcgen05InstSupport]);
+  let isConvergent = true, Predicates = Pred in {
+    def "" : BasicNVPTXInst<(outs),
+               (ins B32:$tmem_addr, B32:$ncols),
+               "tcgen05.dealloc" # exclusive # ".cta_group::" # num #
+               ".sync.aligned.b32",
+               [(Intr B32:$tmem_addr, B32:$ncols)]>;
+  }
 }
 defm TCGEN05_DEALLOC_CG1: TCGEN05_DEALLOC_INTR<"1", int_nvvm_tcgen05_dealloc_cg1>;
 defm TCGEN05_DEALLOC_CG2: TCGEN05_DEALLOC_INTR<"2", int_nvvm_tcgen05_dealloc_cg2>;
 
+defm TCGEN05_DEALLOC_EXCLUSIVE_CG1
+    : TCGEN05_DEALLOC_INTR<"1", int_nvvm_tcgen05_dealloc_exclusive_cg1,
+                           /*IsExclusive=*/1>;
+defm TCGEN05_DEALLOC_EXCLUSIVE_CG2
+    : TCGEN05_DEALLOC_INTR<"2", int_nvvm_tcgen05_dealloc_exclusive_cg2,
+                           /*IsExclusive=*/1>;
+
+let isConvergent = true in {
+let Predicates = [hasTcgen05InstSupport] in {
 multiclass TCGEN05_RELINQ_PERMIT_INTR<string num, Intrinsic Intr> {
   def "" : NullaryInst<"tcgen05.relinquish_alloc_permit.cta_group::" # num # ".sync.aligned", Intr>;
 }
diff --git a/llvm/test/Assembler/auto_upgrade_nvvm_intrinsics.ll b/llvm/test/Assembler/auto_upgrade_nvvm_intrinsics.ll
index 244ced4825517..5813e3bda1c49 100644
--- a/llvm/test/Assembler/auto_upgrade_nvvm_intrinsics.ll
+++ b/llvm/test/Assembler/auto_upgrade_nvvm_intrinsics.ll
@@ -124,6 +124,11 @@ declare void @llvm.nvvm.tcgen05.commit.mc.cg2(ptr, i16)
 declare void @llvm.nvvm.tcgen05.commit.mc.shared.cg1(ptr addrspace(3), i16)
 declare void @llvm.nvvm.tcgen05.commit.mc.shared.cg2(ptr addrspace(3), i16)
 
+declare void @llvm.nvvm.tcgen05.alloc.cg1(ptr, i32)
+declare void @llvm.nvvm.tcgen05.alloc.cg2(ptr, i32)
+declare void @llvm.nvvm.tcgen05.alloc.shared.cg1(ptr addrspace(3), i32)
+declare void @llvm.nvvm.tcgen05.alloc.shared.cg2(ptr addrspace(3), i32)
+
 declare void @llvm.nvvm.cp.async.bulk.tensor.g2s.tile.1d(ptr addrspace(3) %d, ptr addrspace(3) %bar, ptr %tm, i32 %d0, i16 %mc, i64 %ch, i1 %f1, i1 %f2);
 declare void @llvm.nvvm.cp.async.bulk.tensor.g2s.tile.2d(ptr addrspace(3) %d, ptr addrspace(3) %bar, ptr %tm, i32 %d0, i32 %d1, i16 %mc, i64 %ch, i1 %f1, i1 %f2);
 declare void @llvm.nvvm.cp.async.bulk.tensor.g2s.tile.3d(ptr addrspace(3) %d, ptr addrspace(3) %bar, ptr %tm, i32 %d0, i32 %d1, i32 %d2, i16 %mc, i64 %ch, i1 %f1, i1 %f2);
@@ -503,6 +508,19 @@ define void @nvvm_tcgen05_commit_intrinsics(ptr %bar, ptr addrspace(3) %bar_shar
   ret void
 }
 
+; CHECK-LABEL: @nvvm_tcgen05_alloc_intrinsics
+define void @nvvm_tcgen05_alloc_intrinsics(ptr %dst, ptr addrspace(3) %dst_shared, i32 %ncols) {
+; CHECK: call void @llvm.nvvm.tcgen05.alloc.cg1.p0(ptr %dst, i32 %ncols)
+; CHECK: call void @llvm.nvvm.tcgen05.alloc.cg2.p0(ptr %dst, i32 %ncols)
+; CHECK: call void @llvm.nvvm.tcgen05.alloc.cg1.p3(ptr addrspace(3) %dst_shared, i32 %ncols)
+; CHECK: call void @llvm.nvvm.tcgen05.alloc.cg2.p3(ptr addrspace(3) %dst_shared, i32 %ncols)
+  call void @llvm.nvvm.tcgen05.alloc.cg1(ptr %dst, i32 %ncols)
+  call void @llvm.nvvm.tcgen05.alloc.cg2(ptr %dst, i32 %ncols)
+  call void @llvm.nvvm.tcgen05.alloc.shared.cg1(ptr addrspace(3) %dst_shared, i32 %ncols)
+  call void @llvm.nvvm.tcgen05.alloc.shared.cg2(ptr addrspace(3) %dst_shared, i32 %ncols)
+  ret void
+}
+
 ; CHECK-LABEL: @nvvm_tcgen05_commit_mc_16b_intrinsics
 define void @nvvm_tcgen05_commit_mc_16b_intrinsics(ptr %bar, ptr addrspace(3) %bar_shared, i16 %cta_mask) {
 ; CHECK: call void @llvm.nvvm.tcgen05.commit.mc.cg1.p0.i16(ptr %bar, i16 %cta_mask)
diff --git a/llvm/test/CodeGen/NVPTX/tcgen05-alloc-dealloc-exclusive.ll b/llvm/test/CodeGen/NVPTX/tcgen05-alloc-dealloc-exclusive.ll
new file mode 100644
index 0000000000000..deb477ce41e54
--- /dev/null
+++ b/llvm/test/CodeGen/NVPTX/tcgen05-alloc-dealloc-exclusive.ll
@@ -0,0 +1,165 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 5
+; RUN: llc < %s -march=nvptx64 -mcpu=sm_100f -mattr=+ptx94 | FileCheck --check-prefixes=CHECK_PTX64 %s
+; RUN: llc < %s -march=nvptx64 -mcpu=sm_100f -mattr=+ptx94 -target-abi=shortptr | FileCheck --check-prefixes=CHECK_PTX64_SHARED32 %s
+; RUN: llc < %s -march=nvptx64 -mcpu=sm_110f -mattr=+ptx94 | FileCheck --check-prefixes=CHECK_PTX64 %s
+; RUN: %if ptxas-sm_100f && ptxas-isa-9.4 %{ llc < %s -march=nvptx64 -mcpu=sm_100f -mattr=+ptx94 | %ptxas-verify -arch=sm_100f %}
+; RUN: %if ptxas-sm_100f && ptxas-isa-9.4 %{ llc < %s -march=nvptx64 -mcpu=sm_100f -mattr=+ptx94 -target-abi=shortptr | %ptxas-verify -arch=sm_100f %}
+; RUN: %if ptxas-sm_110f && ptxas-isa-9.4 %{ llc < %s -march=nvptx64 -mcpu=sm_110f -mattr=+ptx94 | %ptxas-verify -arch=sm_110f %}
+
+declare void @llvm.nvvm.tcgen05.alloc.exclusive.cg1.p0(ptr %addr, i32 %ncols)
+declare void @llvm.nvvm.tcgen05.alloc.exclusive.cg2.p0(ptr %addr, i32 %ncols)
+declare void @llvm.nvvm.tcgen05.alloc.exclusive.cg1.p3(ptr addrspace(3) %addr, i32 %ncols)
+declare void @llvm.nvvm.tcgen05.alloc.exclusive.cg2.p3(ptr addrspace(3) %addr, i32 %ncols)
+
+define void @test_tcgen05_alloc_exclusive_cg1(ptr %addr, i32 %ncols) {
+; CHECK_PTX64-LABEL: test_tcgen05_alloc_exclusive_cg1(
+; CHECK_PTX64:       {
+; CHECK_PTX64-NEXT:    .reg .b32 %r<2>;
+; CHECK_PTX64-NEXT:    .reg .b64 %rd<2>;
+; CHECK_PTX64-EMPTY:
+; CHECK_PTX64-NEXT:  // %bb.0:
+; CHECK_PTX64-NEXT:    ld.param::func.b64 %rd1, [test_tcgen05_alloc_exclusive_cg1_param_0];
+; CHECK_PTX64-NEXT:    ld.param::func.b32 %r1, [test_tcgen05_alloc_exclusive_cg1_param_1];
+; CHECK_PTX64-NEXT:    tcgen05.alloc.exclusive.cta_group::1.sync.aligned.b32 [%rd1], %r1;
+; CHECK_PTX64-NEXT:    ret;
+;
+; CHECK_PTX64_SHARED32-LABEL: test_tcgen05_alloc_exclusive_cg1(
+; CHECK_PTX64_SHARED32:       {
+; CHECK_PTX64_SHARED32-NEXT:    .reg .b32 %r<2>;
+; CHECK_PTX64_SHARED32-NEXT:    .reg .b64 %rd<2>;
+; CHECK_PTX64_SHARED32-EMPTY:
+; CHECK_PTX64_SHARED32-NEXT:  // %bb.0:
+; CHECK_PTX64_SHARED32-NEXT:    ld.param::func.b64 %rd1, [test_tcgen05_alloc_exclusive_cg1_param_0];
+; CHECK_PTX64_SHARED32-NEXT:    ld.param::func.b32 %r1, [test_tcgen05_alloc_exclusive_cg1_param_1];
+; CHECK_PTX64_SHARED32-NEXT:    tcgen05.alloc.exclusive.cta_group::1.sync.aligned.b32 [%rd1], %r1;
+; CHECK_PTX64_SHARED32-NEXT:    ret;
+  call void @llvm.nvvm.tcgen05.alloc.exclusive.cg1.p0(ptr %addr, i32 %ncols)
+  ret void
+}
+
+define void @test_tcgen05_alloc_exclusive_cg2(ptr %addr, i32 %ncols) {
+; CHECK_PTX64-LABEL: test_tcgen05_alloc_exclusive_cg2(
+; CHECK_PTX64:       {
+; CHECK_PTX64-NEXT:    .reg .b32 %r<2>;
+; CHECK_PTX64-NEXT:    .reg .b64 %rd<2>;
+; CHECK_PTX64-EMPTY:
+; CHECK_PTX64-NEXT:  // %bb.0:
+; CHECK_PTX64-NEXT:    ld.param::func.b64 %rd1, [test_tcgen05_alloc_exclusive_cg2_param_0];
+; CHECK_PTX64-NEXT:    ld.param::func.b32 %r1, [test_tcgen05_alloc_exclusive_cg2_param_1];
+; CHECK_PTX64-NEXT:    tcgen05.alloc.exclusive.cta_group::2.sync.aligned.b32 [%rd1], %r1;
+; CHECK_PTX64-NEXT:    ret;
+;
+; CHECK_PTX64_SHARED32-LABEL: test_tcgen05_alloc_exclusive_cg2(
+; CHECK_PTX64_SHARED32:       {
+; CHECK_PTX64_SHARED32-NEXT:    .reg .b32 %r<2>;
+; CHECK_PTX64_SHARED32-NEXT:    .reg .b64 %rd<2>;
+; CHECK_PTX64_SHARED32-EMPTY:
+; CHECK_PTX64_SHARED32-NEXT:  // %bb.0:
+; CHECK_PTX64_SHARED32-NEXT:    ld.param::func.b64 %rd1, [test_tcgen05_alloc_exclusive_cg2_param_0];
+; CHECK_PTX64_SHARED32-NEXT:    ld.param::func.b32 %r1, [test_tcgen05_alloc_exclusive_cg2_param_1];
+; CHECK_PTX64_SHARED32-NEXT:    tcgen05.alloc.exclusive.cta_group::2.sync.aligned.b32 [%rd1], %r1;
+; CHECK_PTX64_SHARED32-NEXT:    ret;
+  call void @llvm.nvvm.tcgen05.alloc.exclusive.cg2.p0(ptr %addr, i32 %ncols)
+  ret void
+}
+
+define void @test_tcgen05_alloc_exclusive_shared_cg1(ptr addrspace(3) %addr, i32 %ncols) {
+; CHECK_PTX64-LABEL: test_tcgen05_alloc_exclusive_shared_cg1(
+; CHECK_PTX64:       {
+; CHECK_PTX64-NEXT:    .reg .b32 %r<2>;
+; CHECK_PTX64-NEXT:    .reg .b64 %rd<2>;
+; CHECK_PTX64-EMPTY:
+; CHECK_PTX64-NEXT:  // %bb.0:
+; CHECK_PTX64-NEXT:    ld.param::func.b64 %rd1, [test_tcgen05_alloc_exclusive_shared_cg1_param_0];
+; CHECK_PTX64-NEXT:    ld.param::func.b32 %r1, [test_tcgen05_alloc_exclusive_shared_cg1_param_1];
+; CHECK_PTX64-NEXT:    tcgen05.alloc.exclusive.cta_group::1.sync.aligned.shared::cta.b32 [%rd1], %r1;
+; CHECK_PTX64-NEXT:    ret;
+;
+; CHECK_PTX64_SHARED32-LABEL: test_tcgen...
[truncated]

``````````

</details>


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


More information about the Mlir-commits mailing list