[Mlir-commits] [mlir] [MLIR][NVVM] Add "exclusive" support in tcgen05 alloc/dealloc Ops (PR #219451)
Dharuni R Acharya
llvmlistbot at llvm.org
Fri Aug 28 05:15:22 PDT 2026
https://github.com/DharuniRAcharya created https://github.com/llvm/llvm-project/pull/219451
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
>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] [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
}
More information about the Mlir-commits
mailing list