[Mlir-commits] [mlir] [MLIR][NVVM] Add "blocksareclusters" kernel attribute support (PR #154519)
Rajat Bajpai
llvmlistbot at llvm.org
Thu Aug 21 00:32:55 PDT 2025
https://github.com/rajatbajpai updated https://github.com/llvm/llvm-project/pull/154519
>From 4ed0cfd2487752f01e19f60bafc10238c9e7cee4 Mon Sep 17 00:00:00 2001
From: rbajpai <rbajpai at nvidia.com>
Date: Wed, 20 Aug 2025 16:02:14 +0530
Subject: [PATCH 1/2] [MLIR] Add "blocksareclusters" kernel attribute support
This change adds "nvvm.blocksareclusters" kernel attribute support.
---
mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td | 4 ++++
mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp | 9 +++++++++
.../Dialect/NVVM/NVVMToLLVMIRTranslation.cpp | 4 ++++
mlir/test/Target/LLVMIR/nvvmir-invalid.mlir | 16 ++++++++++++++++
mlir/test/Target/LLVMIR/nvvmir.mlir | 9 +++++++++
5 files changed, 42 insertions(+)
diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
index f9cd58de8915f..ad82b4aefee40 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
@@ -83,6 +83,10 @@ 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.
+ 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 dbcc738b4419f..0be2a1fcda78d 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
@@ -1922,6 +1922,15 @@ LogicalResult NVVMDialect::verifyOperationAttribute(Operation *op,
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 33398cfb92429..90cf9b5593054 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 c8ba91efbff4d..554042dee81ae 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 =
>From 04862b9d8b4c0d35631ecc15093bd76f83751707 Mon Sep 17 00:00:00 2001
From: rbajpai <rbajpai at nvidia.com>
Date: Thu, 21 Aug 2025 12:16:29 +0530
Subject: [PATCH 2/2] Addressed review comments
+ Added curly braces for multi line formatted statement.
+ Added small description about "blocksareclusters."
---
mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td | 5 +++++
mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp | 9 ++++++---
2 files changed, 11 insertions(+), 3 deletions(-)
diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
index ad82b4aefee40..68eb56a90a6ab 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
@@ -85,6 +85,11 @@ def NVVM_Dialect : Dialect {
/// 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
diff --git a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
index 0be2a1fcda78d..7cd5ceeff5a1b 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
@@ -1908,28 +1908,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()))
+ !op->hasAttr(NVVMDialect::getClusterDimAttrName())) {
return op->emitError()
<< "'" << attrName << "' attribute must be used along with "
<< "'" << NVVMDialect::getReqntidAttrName() << "' and "
<< "'" << NVVMDialect::getClusterDimAttrName() << "'";
+ }
}
return success();
More information about the Mlir-commits
mailing list