[Mlir-commits] [mlir] [MLIR][NVVM] Add "exclusive" support in tcgen05 alloc/dealloc Ops (PR #219451)

Fri Aug 28 05:16:00 PDT 2026

llvmorg-github-actions[bot] wrote:




@llvm/pr-subscribers-mlir-llvm

Author: Dharuni R Acharya (DharuniRAcharya)

<details>
<summary>Changes</summary>

This patch adds `exclusive` support in `tcgen05.alloc/dealloc` MLIR ops. 
Lit tests are added to verify the lowering to the intrinsics.

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 

---
Full diff: https://github.com/llvm/llvm-project/pull/219451.diff


3 Files Affected:

- (modified) mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td (+21-6) 
- (modified) mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp (+4-2) 
- (modified) mlir/test/Target/LLVMIR/nvvm/tcgen05-alloc.mlir (+18) 


``````````diff

diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
index 634918b1b21f2..83754f7296cfe 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
@@ -5562,17 +5562,25 @@ def NVVM_Tcgen05AllocOp : NVVM_Op<"tcgen05.alloc", [NVVMRequiresSMf<[100, 101, 1
     The `tcgen05.alloc` Op allocates tensor core memory for
     the amount specified by `nCols` and writes the destination
     address to the `addr` argument. The `nCols` operand specifies the
-    number of columns to be allocated and it must be a power-of-two.
+    number of columns to be allocated and it must be a power-of-two
+    for non-exclusive allocations and a multiple of 32 for exclusive
+    allocations. The `exclusive` attribute requests an exclusive
+    allocation and defaults to `false`. When `exclusive` is `true`, the
+    `.exclusive` variant is emitted, which claims ownership of the
+    allocation permit. No other allocation may exist at the same time
+    as an exclusive allocation.
     [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-memory-alloc-manage-instructions)
   }];
 
   let arguments = (ins
     AnyTypeOf<[LLVM_AnyPointer, LLVM_PointerShared]>:$addr,
     I32:$nCols,
-    DefaultValuedAttr<CTAGroupKindAttr, "CTAGroupKind::CTA_1">:$group);
+    DefaultValuedAttr<CTAGroupKindAttr, "CTAGroupKind::CTA_1">:$group,
+    DefaultValuedAttr<BoolAttr, "false">:$isExclusive);
 
   let assemblyFormat = [{
-    $addr `,` $nCols (`,` `group` `=` $group^)? attr-dict `:` type(operands)
+    $addr `,` $nCols (`,` `group` `=` $group^)?
+    (`exclusive` `=` $isExclusive^)? attr-dict `:` type(operands)
   }];
 
   let extraClassDeclaration = [{
@@ -5594,15 +5602,22 @@ def NVVM_Tcgen05DeallocOp : NVVM_Op<"tcgen05.dealloc", [NVVMRequiresSMf<[100, 10
     The `tcgen05.dealloc` Op de-allocates the tensor core memory
     specified by `tmemAddr`, which must be from a previous tensor
     memory allocation. The `nCols` operand specifies the number
-    of columns to be de-allocated, and it must be a power-of-two.
+    of columns to be de-allocated, and it must be a power-of-two
+    for non-exclusive deallocations and a multiple of 32 for exclusive
+    deallocations. The `exclusive` attribute defaults to `false`. When
+    `exclusive` is `true`, the `.exclusive` variant is emitted. Memory
+    must be deallocated with `exclusive` if and only if it was allocated
+    with `exclusive`.
     [For more information, see PTX ISA](https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-memory-alloc-manage-instructions)
   }];
 
   let arguments = (ins LLVM_PointerTensor:$taddr, I32:$nCols,
-    DefaultValuedAttr<CTAGroupKindAttr, "CTAGroupKind::CTA_1">:$group);
+    DefaultValuedAttr<CTAGroupKindAttr, "CTAGroupKind::CTA_1">:$group,
+    DefaultValuedAttr<BoolAttr, "false">:$isExclusive);
 
   let assemblyFormat = [{
-    $taddr `,` $nCols (`,` `group` `=` $group^)? attr-dict `:` type(operands)
+    $taddr `,` $nCols (`,` `group` `=` $group^)?
+    (`exclusive` `=` $isExclusive^)? attr-dict `:` type(operands)
   }];
 
   let extraClassDeclaration = [{
diff --git a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
index aa45bba341240..0da13bae8b5b1 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
@@ -5517,7 +5517,8 @@ Tcgen05AllocOp::getIntrinsicIDAndArgs(Operation &op,
   // Fill the Intrinsic Args
   args.push_back(mt.lookupValue(curOp.getAddr()));
   args.push_back(mt.lookupValue(curOp.getNCols()));
-  args.push_back(llvm::ConstantInt::getFalse(mt.getLLVMContext()));
+  args.push_back(
+      llvm::ConstantInt::getBool(mt.getLLVMContext(), curOp.getIsExclusive()));
 
   return id;
 }
@@ -5533,7 +5534,8 @@ llvm::Intrinsic::ID Tcgen05DeallocOp::getIntrinsicIDAndArgs(
   // Fill the Intrinsic Args
   args.push_back(mt.lookupValue(curOp.getTaddr()));
   args.push_back(mt.lookupValue(curOp.getNCols()));
-  args.push_back(llvm::ConstantInt::getFalse(mt.getLLVMContext()));
+  args.push_back(
+      llvm::ConstantInt::getBool(mt.getLLVMContext(), curOp.getIsExclusive()));
 
   return id;
 }
diff --git a/mlir/test/Target/LLVMIR/nvvm/tcgen05-alloc.mlir b/mlir/test/Target/LLVMIR/nvvm/tcgen05-alloc.mlir
index 981aa8cef22ae..0c45816c68817 100644
--- a/mlir/test/Target/LLVMIR/nvvm/tcgen05-alloc.mlir
+++ b/mlir/test/Target/LLVMIR/nvvm/tcgen05-alloc.mlir
@@ -7,6 +7,12 @@ llvm.func @llvm_nvvm_tcgen05_alloc(%addr : !llvm.ptr, %ncols : i32) {
 
   // CHECK-LLVM: call void @llvm.nvvm.tcgen05.alloc.cg2.p0(ptr %{{.*}}, i32 %{{.*}}, /* is_exclusive= */ i1 false)
   nvvm.tcgen05.alloc %addr, %ncols , group = <cta_2> : !llvm.ptr, i32
+
+  // CHECK-LLVM: call void @llvm.nvvm.tcgen05.alloc.cg1.p0(ptr %{{.*}}, i32 %{{.*}}, /* is_exclusive= */ i1 true)
+  nvvm.tcgen05.alloc %addr, %ncols exclusive = true : !llvm.ptr, i32
+
+  // CHECK-LLVM: call void @llvm.nvvm.tcgen05.alloc.cg2.p0(ptr %{{.*}}, i32 %{{.*}}, /* is_exclusive= */ i1 true)
+  nvvm.tcgen05.alloc %addr, %ncols, group = <cta_2> exclusive = true : !llvm.ptr, i32
   llvm.return
 }
 
@@ -17,6 +23,12 @@ llvm.func @llvm_nvvm_tcgen05_alloc_shared(%addr : !llvm.ptr<3>, %ncols : i32) {
 
   // CHECK-LLVM: call void @llvm.nvvm.tcgen05.alloc.cg2.p3(ptr addrspace(3) %{{.*}}, i32 %{{.*}}, /* is_exclusive= */ i1 false)
   nvvm.tcgen05.alloc %addr, %ncols , group = <cta_2> : !llvm.ptr<3>, i32
+
+  // CHECK-LLVM: call void @llvm.nvvm.tcgen05.alloc.cg1.p3(ptr addrspace(3) %{{.*}}, i32 %{{.*}}, /* is_exclusive= */ i1 true)
+  nvvm.tcgen05.alloc %addr, %ncols exclusive = true : !llvm.ptr<3>, i32
+
+  // CHECK-LLVM: call void @llvm.nvvm.tcgen05.alloc.cg2.p3(ptr addrspace(3) %{{.*}}, i32 %{{.*}}, /* is_exclusive= */ i1 true)
+  nvvm.tcgen05.alloc %addr, %ncols, group = <cta_2> exclusive = true : !llvm.ptr<3>, i32
   llvm.return
 }
 
@@ -27,6 +39,12 @@ llvm.func @llvm_nvvm_tcgen05_dealloc(%addr : !llvm.ptr<6>, %ncols : i32) {
 
   // CHECK-LLVM: call void @llvm.nvvm.tcgen05.dealloc.cg2(ptr addrspace(6) %{{.*}}, i32 %{{.*}}, /* is_exclusive= */ i1 false)
   nvvm.tcgen05.dealloc %addr, %ncols , group = <cta_2> : !llvm.ptr<6>, i32
+
+  // CHECK-LLVM: call void @llvm.nvvm.tcgen05.dealloc.cg1(ptr addrspace(6) %{{.*}}, i32 %{{.*}}, /* is_exclusive= */ i1 true)
+  nvvm.tcgen05.dealloc %addr, %ncols exclusive = true : !llvm.ptr<6>, i32
+
+  // CHECK-LLVM: call void @llvm.nvvm.tcgen05.dealloc.cg2(ptr addrspace(6) %{{.*}}, i32 %{{.*}}, /* is_exclusive= */ i1 true)
+  nvvm.tcgen05.dealloc %addr, %ncols, group = <cta_2> exclusive = true : !llvm.ptr<6>, i32
   llvm.return
 }
 

``````````

</details>


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