[Mlir-commits] [mlir] [MLIR][NVVM] Add "exclusive" support in tcgen05 alloc/dealloc Ops (PR #219451)
Dharuni R Acharya
llvmlistbot at llvm.org
Sun Aug 30 21:04:22 PDT 2026
https://github.com/DharuniRAcharya updated https://github.com/llvm/llvm-project/pull/219451
>From e0936c15863a7a8852e03433c1aeaa52da6a102a Mon Sep 17 00:00:00 2001
From: DharuniRAcharya <dharunira at nvidia.com>
Date: Fri, 28 Aug 2026 12:11:15 +0000
Subject: [PATCH 1/2] [MLIR][NVVM] Add "exclusive" support in tcgen05
alloc/dealloc Ops
This patch adds exclusive support in tcgen05.alloc/dealloc MLIR ops.
Lit tests are added to verify the lowering to the intrinsics.
Signed-off-by: DharuniRAcharya <dharunira at nvidia.com>
---
mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td | 27 ++++++++++++++-----
mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp | 6 +++--
.../Target/LLVMIR/nvvm/tcgen05-alloc.mlir | 18 +++++++++++++
3 files changed, 43 insertions(+), 8 deletions(-)
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
}
>From d526b4718e715902c6db2cb0f56d761e02679e39 Mon Sep 17 00:00:00 2001
From: DharuniRAcharya <dharunira at nvidia.com>
Date: Mon, 31 Aug 2026 04:04:00 +0000
Subject: [PATCH 2/2] Address comments
---
mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td | 30 +++--------------
mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp | 37 ++++++++-------------
2 files changed, 17 insertions(+), 50 deletions(-)
diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
index 83754f7296cfe..ccec2ed50f07a 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
@@ -5556,7 +5556,8 @@ def Tcgen05WaitKindAttr :
let assemblyFormat = "`<` $value `>`";
}
-def NVVM_Tcgen05AllocOp : NVVM_Op<"tcgen05.alloc", [NVVMRequiresSMf<[100, 101, 110]>]> {
+def NVVM_Tcgen05AllocOp : NVVM_VoidIntrinsicOp<"tcgen05.alloc",
+ [NVVMRequiresSMf<[100, 101, 110]>]> {
let summary = "Tcgen05 alloc operation";
let description = [{
The `tcgen05.alloc` Op allocates tensor core memory for
@@ -5582,21 +5583,10 @@ def NVVM_Tcgen05AllocOp : NVVM_Op<"tcgen05.alloc", [NVVMRequiresSMf<[100, 101, 1
$addr `,` $nCols (`,` `group` `=` $group^)?
(`exclusive` `=` $isExclusive^)? attr-dict `:` type(operands)
}];
-
- let extraClassDeclaration = [{
- static llvm::Intrinsic::ID
- getIntrinsicIDAndArgs(Operation &op, LLVM::ModuleTranslation &mt,
- llvm::SmallVector<llvm::Value *> &args);
- }];
- string llvmBuilder = [{
- llvm::SmallVector<llvm::Value *> args;
- auto id = NVVM::Tcgen05AllocOp::getIntrinsicIDAndArgs(
- *op, moduleTranslation, args);
- createIntrinsicCall(builder, id, builder.getVoidTy(), args);
- }];
}
-def NVVM_Tcgen05DeallocOp : NVVM_Op<"tcgen05.dealloc", [NVVMRequiresSMf<[100, 101, 110]>]> {
+def NVVM_Tcgen05DeallocOp : NVVM_VoidIntrinsicOp<"tcgen05.dealloc",
+ [NVVMRequiresSMf<[100, 101, 110]>]> {
let summary = "Tcgen05 dealloc operation";
let description = [{
The `tcgen05.dealloc` Op de-allocates the tensor core memory
@@ -5619,18 +5609,6 @@ def NVVM_Tcgen05DeallocOp : NVVM_Op<"tcgen05.dealloc", [NVVMRequiresSMf<[100, 10
$taddr `,` $nCols (`,` `group` `=` $group^)?
(`exclusive` `=` $isExclusive^)? attr-dict `:` type(operands)
}];
-
- let extraClassDeclaration = [{
- static llvm::Intrinsic::ID
- getIntrinsicIDAndArgs(Operation &op, LLVM::ModuleTranslation &mt,
- llvm::SmallVector<llvm::Value *> &args);
- }];
- string llvmBuilder = [{
- llvm::SmallVector<llvm::Value *> args;
- auto id = NVVM::Tcgen05DeallocOp::getIntrinsicIDAndArgs(
- *op, moduleTranslation, args);
- createIntrinsicCall(builder, id, args);
- }];
}
def NVVM_Tcgen05RelinquishAllocPermitOp : NVVM_Op<"tcgen05.relinquish_alloc_permit",
diff --git a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
index 0da13bae8b5b1..c5cb1b6d72698 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
@@ -5504,40 +5504,29 @@ NVVM::IDArgPair ConvertS2F6x2ToBF16x2Op::getIntrinsicIDAndArgs(
return {ids[idx], std::move(args)};
}
-llvm::Intrinsic::ID
-Tcgen05AllocOp::getIntrinsicIDAndArgs(Operation &op,
- LLVM::ModuleTranslation &mt,
- llvm::SmallVector<llvm::Value *> &args) {
+mlir::NVVM::IDArgPair Tcgen05AllocOp::getIntrinsicIDAndArgs(
+ Operation &op, LLVM::ModuleTranslation &mt, llvm::IRBuilderBase &builder) {
auto curOp = cast<NVVM::Tcgen05AllocOp>(op);
bool is2CTAMode = curOp.getGroup() == CTAGroupKind::CTA_2;
llvm::Intrinsic::ID id = is2CTAMode ? llvm::Intrinsic::nvvm_tcgen05_alloc_cg2
: llvm::Intrinsic::nvvm_tcgen05_alloc_cg1;
- // Fill the Intrinsic Args
- args.push_back(mt.lookupValue(curOp.getAddr()));
- args.push_back(mt.lookupValue(curOp.getNCols()));
- args.push_back(
- llvm::ConstantInt::getBool(mt.getLLVMContext(), curOp.getIsExclusive()));
-
- return id;
+ return {id,
+ {mt.lookupValue(curOp.getAddr()), mt.lookupValue(curOp.getNCols()),
+ builder.getInt1(curOp.getIsExclusive())}};
}
-llvm::Intrinsic::ID Tcgen05DeallocOp::getIntrinsicIDAndArgs(
- Operation &op, LLVM::ModuleTranslation &mt,
- llvm::SmallVector<llvm::Value *> &args) {
+mlir::NVVM::IDArgPair Tcgen05DeallocOp::getIntrinsicIDAndArgs(
+ Operation &op, LLVM::ModuleTranslation &mt, llvm::IRBuilderBase &builder) {
auto curOp = cast<NVVM::Tcgen05DeallocOp>(op);
- auto id = (curOp.getGroup() == CTAGroupKind::CTA_1)
- ? llvm::Intrinsic::nvvm_tcgen05_dealloc_cg1
- : llvm::Intrinsic::nvvm_tcgen05_dealloc_cg2;
+ llvm::Intrinsic::ID id = (curOp.getGroup() == CTAGroupKind::CTA_1)
+ ? llvm::Intrinsic::nvvm_tcgen05_dealloc_cg1
+ : llvm::Intrinsic::nvvm_tcgen05_dealloc_cg2;
- // Fill the Intrinsic Args
- args.push_back(mt.lookupValue(curOp.getTaddr()));
- args.push_back(mt.lookupValue(curOp.getNCols()));
- args.push_back(
- llvm::ConstantInt::getBool(mt.getLLVMContext(), curOp.getIsExclusive()));
-
- return id;
+ return {id,
+ {mt.lookupValue(curOp.getTaddr()), mt.lookupValue(curOp.getNCols()),
+ builder.getInt1(curOp.getIsExclusive())}};
}
llvm::Intrinsic::ID
More information about the Mlir-commits
mailing list