[Mlir-commits] [mlir] [MLIR][NVVM] Add "exclusive" support in tcgen05 alloc/dealloc Ops (PR #219451)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Aug 28 05:15:59 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
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
More information about the Mlir-commits
mailing list