[Mlir-commits] [mlir] b08b219 - [MLIR][NVVM] Add "blocksareclusters" kernel attribute support (#154519)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Aug 21 23:02:25 PDT 2025
Author: Rajat Bajpai
Date: 2025-08-22T11:32:21+05:30
New Revision: b08b219650ea220a75f8988522a4ea41f92707b7
URL: https://github.com/llvm/llvm-project/commit/b08b219650ea220a75f8988522a4ea41f92707b7
DIFF: https://github.com/llvm/llvm-project/commit/b08b219650ea220a75f8988522a4ea41f92707b7.diff
LOG: [MLIR][NVVM] Add "blocksareclusters" kernel attribute support (#154519)
This change adds "nvvm.blocksareclusters" kernel attribute support in NVVM Dialect/MLIR.
Added:
Modified:
mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
mlir/lib/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.cpp
mlir/test/Target/LLVMIR/nvvmir-invalid.mlir
mlir/test/Target/LLVMIR/nvvmir.mlir
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
index 9661df5d11790..880c268ae6614 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
@@ -83,6 +83,15 @@ def NVVM_Dialect : Dialect {
/// are grid constants.
static StringRef getGridConstantAttrName() { return "nvvm.grid_constant"; }
+ /// Get the name of the attribute used to annotate the `.blocksareclusters`
+ /// PTX directive for kernel functions.
+ /// This attribute implies that the grid launch configuration for the
+ /// corresponding kernel function is specifying the number of clusters
+ /// instead of the number of thread blocks. This attribute is only
+ /// allowed for kernel functions and requires nvvm.reqntid and
+ /// nvvm.cluster_dim attributes.
+ static StringRef getBlocksAreClustersAttrName() { return "nvvm.blocksareclusters"; }
+
/// Verify an attribute from this dialect on the argument at 'argIndex' for
/// the region at 'regionIndex' on the given operation. Returns failure if
/// the verification failed, success otherwise. This hook may optionally be
diff --git a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
index 042103cad83ca..c2032aa7f6443 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
@@ -1925,19 +1925,31 @@ LogicalResult NVVMDialect::verifyOperationAttribute(Operation *op,
attrName == NVVMDialect::getReqntidAttrName() ||
attrName == NVVMDialect::getClusterDimAttrName()) {
auto values = llvm::dyn_cast<DenseI32ArrayAttr>(attr.getValue());
- if (!values || values.empty() || values.size() > 3)
+ if (!values || values.empty() || values.size() > 3) {
return op->emitError()
<< "'" << attrName
<< "' attribute must be integer array with maximum 3 index";
+ }
}
// If minctasm / maxnreg / cluster_max_blocks exist, it must be an integer
// attribute
if (attrName == NVVMDialect::getMinctasmAttrName() ||
attrName == NVVMDialect::getMaxnregAttrName() ||
attrName == NVVMDialect::getClusterMaxBlocksAttrName()) {
- if (!llvm::dyn_cast<IntegerAttr>(attr.getValue()))
+ if (!llvm::dyn_cast<IntegerAttr>(attr.getValue())) {
return op->emitError()
<< "'" << attrName << "' attribute must be integer constant";
+ }
+ }
+ // blocksareclusters must be used along with reqntid and cluster_dim
+ if (attrName == NVVMDialect::getBlocksAreClustersAttrName()) {
+ if (!op->hasAttr(NVVMDialect::getReqntidAttrName()) ||
+ !op->hasAttr(NVVMDialect::getClusterDimAttrName())) {
+ return op->emitError()
+ << "'" << attrName << "' attribute must be used along with "
+ << "'" << NVVMDialect::getReqntidAttrName() << "' and "
+ << "'" << NVVMDialect::getClusterDimAttrName() << "'";
+ }
}
return success();
diff --git a/mlir/lib/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.cpp
index e67cfed983255..a20701ce75bc0 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.cpp
@@ -468,7 +468,11 @@ class NVVMDialectLLVMIRTranslationInterface
} else if (attribute.getName() ==
NVVM::NVVMDialect::getKernelFuncAttrName()) {
llvmFunc->setCallingConv(llvm::CallingConv::PTX_Kernel);
+ } else if (attribute.getName() ==
+ NVVM::NVVMDialect::getBlocksAreClustersAttrName()) {
+ llvmFunc->addFnAttr("nvvm.blocksareclusters");
}
+
return success();
}
diff --git a/mlir/test/Target/LLVMIR/nvvmir-invalid.mlir b/mlir/test/Target/LLVMIR/nvvmir-invalid.mlir
index 703b8b38b3d87..06f2f52dfc0c6 100644
--- a/mlir/test/Target/LLVMIR/nvvmir-invalid.mlir
+++ b/mlir/test/Target/LLVMIR/nvvmir-invalid.mlir
@@ -56,6 +56,22 @@ llvm.func @kernel_func() attributes {nvvm.kernel, nvvm.maxntid = array<i32: 3, 4
// -----
+// expected-error @below {{'"nvvm.blocksareclusters"' attribute must be used along with 'nvvm.reqntid' and 'nvvm.cluster_dim'}}
+llvm.func @kernel_func() attributes {nvvm.kernel, nvvm.blocksareclusters,
+ nvvm.cluster_dim = array<i32: 3, 5, 7>} {
+ llvm.return
+}
+
+// -----
+
+// expected-error @below {{'"nvvm.blocksareclusters"' attribute must be used along with 'nvvm.reqntid' and 'nvvm.cluster_dim'}}
+llvm.func @kernel_func() attributes {nvvm.kernel, nvvm.blocksareclusters,
+ nvvm.reqntid = array<i32: 1, 23, 32>} {
+ llvm.return
+}
+
+// -----
+
llvm.func @nvvm_fence_proxy_acquire(%addr : !llvm.ptr, %size : i32) {
// expected-error @below {{'nvvm.fence.proxy.acquire' op uni-directional proxies only support generic for from_proxy attribute}}
nvvm.fence.proxy.acquire #nvvm.mem_scope<cta> %addr, %size from_proxy=#nvvm.proxy_kind<tensormap> to_proxy=#nvvm.proxy_kind<generic>
diff --git a/mlir/test/Target/LLVMIR/nvvmir.mlir b/mlir/test/Target/LLVMIR/nvvmir.mlir
index 16191d925959b..5f0bda3aa2dfe 100644
--- a/mlir/test/Target/LLVMIR/nvvmir.mlir
+++ b/mlir/test/Target/LLVMIR/nvvmir.mlir
@@ -692,7 +692,16 @@ llvm.func @kernel_func() attributes {nvvm.kernel, nvvm.maxntid = array<i32: 1, 2
// CHECK: define ptx_kernel void @kernel_func() #[[ATTR0:[0-9]+]]
// CHECK: attributes #[[ATTR0]] = { "nvvm.maxnreg"="32" "nvvm.maxntid"="1,23,32" "nvvm.minctasm"="16" }
+// -----
+llvm.func @kernel_func() attributes {nvvm.kernel, nvvm.blocksareclusters,
+ nvvm.reqntid = array<i32: 1, 23, 32>,
+ nvvm.cluster_dim = array<i32: 3, 5, 7>} {
+ llvm.return
+}
+
+// CHECK: define ptx_kernel void @kernel_func() #[[ATTR0:[0-9]+]]
+// CHECK: attributes #[[ATTR0]] = { "nvvm.blocksareclusters" "nvvm.cluster_dim"="3,5,7" "nvvm.reqntid"="1,23,32" }
// -----
// CHECK: define ptx_kernel void @kernel_func
// CHECK: !nvvm.annotations =
More information about the Mlir-commits
mailing list