[llvm] 3fdbee1 - [NVPTX] Constant fold clusterDim when reqnctapercluster is specified (#195967)
via llvm-commits
llvm-commits at lists.llvm.org
Tue May 19 11:46:12 PDT 2026
Author: Chengjun
Date: 2026-05-19T11:46:07-07:00
New Revision: 3fdbee11679ea85ffad8ef67eff10f9df1401290
URL: https://github.com/llvm/llvm-project/commit/3fdbee11679ea85ffad8ef67eff10f9df1401290
DIFF: https://github.com/llvm/llvm-project/commit/3fdbee11679ea85ffad8ef67eff10f9df1401290.diff
LOG: [NVPTX] Constant fold clusterDim when reqnctapercluster is specified (#195967)
This is a follow-up of https://github.com/llvm/llvm-project/pull/191575.
Currently, NVPTX cannot fold the `cluster_nctaid.x/y/z` and
`cluster_nctarank` intrinsic calls into const values when
`reqnctapercluster` is specified, which prevents the code from further
optimization.
Therefore, in this change, we extend the `NVVMIntrRange` pass to:
- Tighten `cluster_nctaid.x/y/z` intrinsic calls to one value range,
which can be const folded in later InstCombine pass
- Tighten `cluster_nctarank` intrinsic calls to one value range when
`cluster_dim` is specified
- Tighten `cluster_ctaid.x/y/z` range attributes to use per-dimension
`cluster_dim` bounds
Added:
llvm/test/CodeGen/NVPTX/reqnctapercluster-const-fold.ll
Modified:
llvm/lib/Target/NVPTX/NVVMIntrRange.cpp
llvm/test/CodeGen/NVPTX/intr-range.ll
Removed:
################################################################################
diff --git a/llvm/lib/Target/NVPTX/NVVMIntrRange.cpp b/llvm/lib/Target/NVPTX/NVVMIntrRange.cpp
index b95cf2fab0885..ed1d880f612c1 100644
--- a/llvm/lib/Target/NVPTX/NVVMIntrRange.cpp
+++ b/llvm/lib/Target/NVPTX/NVVMIntrRange.cpp
@@ -68,18 +68,17 @@ static bool runNVVMIntrRange(Function &F) {
auto ReqNTID = getReqNTID(F);
const auto OverallMaxNTID = getOverallMaxNTID(F);
- const auto OverallClusterRank = getOverallClusterRank(F);
+ auto ClusterDim = getClusterDim(F);
+ const auto MaxClusterRank = getMaxClusterRank(F);
// If this function lacks any range information, do nothing.
- if (!(!ReqNTID.empty() || OverallMaxNTID || OverallClusterRank))
+ if (ReqNTID.empty() && !OverallMaxNTID && ClusterDim.empty() &&
+ !MaxClusterRank)
return false;
const unsigned MaxNTID =
OverallMaxNTID.value_or(std::numeric_limits<unsigned>::max());
- const unsigned FunctionClusterRank =
- OverallClusterRank.value_or(std::numeric_limits<unsigned>::max());
-
// When reqntid is specified, block dimensions are exact compile-time
// constants. Otherwise, use maxntid (capped at hardware limits) as upper
// bounds.
@@ -93,11 +92,29 @@ static bool runNVVMIntrRange(Function &F) {
std::min(64u, MaxNTID)};
}
- // We conservatively use the maximum grid size as an upper bound for the
- // cluster rank.
- const Vector3 MaxClusterRank{std::min(0x7fffffffu, FunctionClusterRank),
- std::min(0xffffu, FunctionClusterRank),
- std::min(0xffffu, FunctionClusterRank)};
+ const bool HasClusterInfo = !ClusterDim.empty() || MaxClusterRank;
+
+ // When cluster_dim is specified, cluster dimensions are exact compile-time
+ // constants. Otherwise, use maxclusterrank (capped at hardware limits) as
+ // upper bounds.
+ Vector3 MinClusterDim, MaxClusterDim;
+ uint64_t MinClusterSize, MaxClusterSize;
+ if (!ClusterDim.empty()) {
+ ClusterDim.resize(3, 1);
+ MinClusterDim =
+ MaxClusterDim = {ClusterDim[0], ClusterDim[1], ClusterDim[2]};
+ MinClusterSize = MaxClusterSize =
+ ClusterDim[0] * ClusterDim[1] * ClusterDim[2];
+ } else {
+ const unsigned MaxNctaPerCluster =
+ MaxClusterRank.value_or(std::numeric_limits<unsigned>::max());
+ MinClusterDim = {1, 1, 1};
+ MaxClusterDim = {std::min(0x7fffffffu, MaxNctaPerCluster),
+ std::min(0xffffu, MaxNctaPerCluster),
+ std::min(0xffffu, MaxNctaPerCluster)};
+ MinClusterSize = 1;
+ MaxClusterSize = MaxNctaPerCluster;
+ }
const auto ProcessIntrinsic = [&](IntrinsicInst *II) -> bool {
switch (II->getIntrinsicID()) {
@@ -118,32 +135,30 @@ static bool runNVVMIntrRange(Function &F) {
case Intrinsic::nvvm_read_ptx_sreg_ntid_z:
return addRangeAttr(MinBlockDim.Z, MaxBlockDim.Z + 1, II);
- // Cluster size
+ // Cluster size: use single-value ranges when cluster_dim is specified;
+ // InstCombine will fold cluster_nctaid.* / cluster_nctarank to constants
+ // later.
case Intrinsic::nvvm_read_ptx_sreg_cluster_ctaid_x:
- return addRangeAttr(0, MaxClusterRank.X, II);
+ return addRangeAttr(0, MaxClusterDim.X, II);
case Intrinsic::nvvm_read_ptx_sreg_cluster_ctaid_y:
- return addRangeAttr(0, MaxClusterRank.Y, II);
+ return addRangeAttr(0, MaxClusterDim.Y, II);
case Intrinsic::nvvm_read_ptx_sreg_cluster_ctaid_z:
- return addRangeAttr(0, MaxClusterRank.Z, II);
+ return addRangeAttr(0, MaxClusterDim.Z, II);
case Intrinsic::nvvm_read_ptx_sreg_cluster_nctaid_x:
- return addRangeAttr(1, MaxClusterRank.X + 1, II);
+ return addRangeAttr(MinClusterDim.X, MaxClusterDim.X + 1, II);
case Intrinsic::nvvm_read_ptx_sreg_cluster_nctaid_y:
- return addRangeAttr(1, MaxClusterRank.Y + 1, II);
+ return addRangeAttr(MinClusterDim.Y, MaxClusterDim.Y + 1, II);
case Intrinsic::nvvm_read_ptx_sreg_cluster_nctaid_z:
- return addRangeAttr(1, MaxClusterRank.Z + 1, II);
+ return addRangeAttr(MinClusterDim.Z, MaxClusterDim.Z + 1, II);
case Intrinsic::nvvm_read_ptx_sreg_cluster_ctarank:
- if (OverallClusterRank)
- return addRangeAttr(0, FunctionClusterRank, II);
- break;
+ return HasClusterInfo && addRangeAttr(0, MaxClusterSize, II);
case Intrinsic::nvvm_read_ptx_sreg_cluster_nctarank:
- if (OverallClusterRank)
- return addRangeAttr(1, FunctionClusterRank + 1, II);
- break;
+ return HasClusterInfo &&
+ addRangeAttr(MinClusterSize, MaxClusterSize + 1, II);
default:
return false;
}
- return false;
};
// Go through the calls in this function.
diff --git a/llvm/test/CodeGen/NVPTX/intr-range.ll b/llvm/test/CodeGen/NVPTX/intr-range.ll
index 59c1d3d3e39d8..cd35fd1ddfc0c 100644
--- a/llvm/test/CodeGen/NVPTX/intr-range.ll
+++ b/llvm/test/CodeGen/NVPTX/intr-range.ll
@@ -108,12 +108,12 @@ define ptx_kernel i32 @test_cluster_ctaid() "nvvm.maxclusterrank"="8" {
define ptx_kernel i32 @test_cluster_dim() "nvvm.cluster_dim"="4,4,1" {
; CHECK-LABEL: define ptx_kernel i32 @test_cluster_dim(
; CHECK-SAME: ) #[[ATTR4:[0-9]+]] {
-; CHECK-NEXT: [[TMP1:%.*]] = call range(i32 0, 16) i32 @llvm.nvvm.read.ptx.sreg.cluster.ctaid.x()
-; CHECK-NEXT: [[TMP2:%.*]] = call range(i32 0, 16) i32 @llvm.nvvm.read.ptx.sreg.cluster.ctaid.y()
-; CHECK-NEXT: [[TMP3:%.*]] = call range(i32 0, 16) i32 @llvm.nvvm.read.ptx.sreg.cluster.ctaid.z()
-; CHECK-NEXT: [[TMP5:%.*]] = call range(i32 1, 17) i32 @llvm.nvvm.read.ptx.sreg.cluster.nctaid.x()
-; CHECK-NEXT: [[TMP6:%.*]] = call range(i32 1, 17) i32 @llvm.nvvm.read.ptx.sreg.cluster.nctaid.y()
-; CHECK-NEXT: [[TMP7:%.*]] = call range(i32 1, 17) i32 @llvm.nvvm.read.ptx.sreg.cluster.nctaid.z()
+; CHECK-NEXT: [[TMP1:%.*]] = call range(i32 0, 4) i32 @llvm.nvvm.read.ptx.sreg.cluster.ctaid.x()
+; CHECK-NEXT: [[TMP2:%.*]] = call range(i32 0, 4) i32 @llvm.nvvm.read.ptx.sreg.cluster.ctaid.y()
+; CHECK-NEXT: [[TMP3:%.*]] = call range(i32 0, 1) i32 @llvm.nvvm.read.ptx.sreg.cluster.ctaid.z()
+; CHECK-NEXT: [[TMP5:%.*]] = call range(i32 4, 5) i32 @llvm.nvvm.read.ptx.sreg.cluster.nctaid.x()
+; CHECK-NEXT: [[TMP6:%.*]] = call range(i32 4, 5) i32 @llvm.nvvm.read.ptx.sreg.cluster.nctaid.y()
+; CHECK-NEXT: [[TMP7:%.*]] = call range(i32 1, 2) i32 @llvm.nvvm.read.ptx.sreg.cluster.nctaid.z()
; CHECK-NEXT: [[TMP9:%.*]] = add i32 [[TMP1]], [[TMP2]]
; CHECK-NEXT: [[TMP10:%.*]] = add i32 [[TMP9]], [[TMP3]]
; CHECK-NEXT: [[TMP12:%.*]] = add i32 [[TMP10]], [[TMP5]]
diff --git a/llvm/test/CodeGen/NVPTX/reqnctapercluster-const-fold.ll b/llvm/test/CodeGen/NVPTX/reqnctapercluster-const-fold.ll
new file mode 100644
index 0000000000000..eb94456d94f16
--- /dev/null
+++ b/llvm/test/CodeGen/NVPTX/reqnctapercluster-const-fold.ll
@@ -0,0 +1,98 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 6
+; RUN: opt < %s -S -mtriple=nvptx-nvidia-cuda -mcpu=sm_100 -passes=nvvm-intr-range | FileCheck %s
+
+; When .reqnctapercluster specifies 3D dimensions, cluster_nctaid.x/y/z and
+; cluster_nctarank should be replaced with constants, and cluster_ctaid.x/y/z
+; should get per-dimension ranges.
+define ptx_kernel i32 @test_cluster_dim_3d() "nvvm.cluster_dim"="4,2,1" {
+; CHECK-LABEL: define ptx_kernel i32 @test_cluster_dim_3d(
+; CHECK-SAME: ) #[[ATTR0:[0-9]+]] {
+; CHECK-NEXT: [[CTAID_X:%.*]] = call range(i32 0, 4) i32 @llvm.nvvm.read.ptx.sreg.cluster.ctaid.x()
+; CHECK-NEXT: [[CTAID_Y:%.*]] = call range(i32 0, 2) i32 @llvm.nvvm.read.ptx.sreg.cluster.ctaid.y()
+; CHECK-NEXT: [[CTAID_Z:%.*]] = call range(i32 0, 1) i32 @llvm.nvvm.read.ptx.sreg.cluster.ctaid.z()
+; CHECK-NEXT: [[NCTAID_X:%.*]] = call range(i32 4, 5) i32 @llvm.nvvm.read.ptx.sreg.cluster.nctaid.x()
+; CHECK-NEXT: [[NCTAID_Y:%.*]] = call range(i32 2, 3) i32 @llvm.nvvm.read.ptx.sreg.cluster.nctaid.y()
+; CHECK-NEXT: [[NCTAID_Z:%.*]] = call range(i32 1, 2) i32 @llvm.nvvm.read.ptx.sreg.cluster.nctaid.z()
+; CHECK-NEXT: [[RANK:%.*]] = call range(i32 0, 8) i32 @llvm.nvvm.read.ptx.sreg.cluster.ctarank()
+; CHECK-NEXT: [[NRANK:%.*]] = call range(i32 8, 9) i32 @llvm.nvvm.read.ptx.sreg.cluster.nctarank()
+; CHECK-NEXT: [[A:%.*]] = add i32 [[CTAID_X]], [[CTAID_Y]]
+; CHECK-NEXT: [[B:%.*]] = add i32 [[A]], [[CTAID_Z]]
+; CHECK-NEXT: [[C:%.*]] = add i32 [[B]], [[NCTAID_X]]
+; CHECK-NEXT: [[D:%.*]] = add i32 [[C]], [[NCTAID_Y]]
+; CHECK-NEXT: [[E:%.*]] = add i32 [[D]], [[NCTAID_Z]]
+; CHECK-NEXT: [[F:%.*]] = add i32 [[E]], [[RANK]]
+; CHECK-NEXT: [[G:%.*]] = add i32 [[F]], [[NRANK]]
+; CHECK-NEXT: ret i32 [[G]]
+;
+ %ctaid.x = call i32 @llvm.nvvm.read.ptx.sreg.cluster.ctaid.x()
+ %ctaid.y = call i32 @llvm.nvvm.read.ptx.sreg.cluster.ctaid.y()
+ %ctaid.z = call i32 @llvm.nvvm.read.ptx.sreg.cluster.ctaid.z()
+ %nctaid.x = call i32 @llvm.nvvm.read.ptx.sreg.cluster.nctaid.x()
+ %nctaid.y = call i32 @llvm.nvvm.read.ptx.sreg.cluster.nctaid.y()
+ %nctaid.z = call i32 @llvm.nvvm.read.ptx.sreg.cluster.nctaid.z()
+ %rank = call i32 @llvm.nvvm.read.ptx.sreg.cluster.ctarank()
+ %nrank = call i32 @llvm.nvvm.read.ptx.sreg.cluster.nctarank()
+ %a = add i32 %ctaid.x, %ctaid.y
+ %b = add i32 %a, %ctaid.z
+ %c = add i32 %b, %nctaid.x
+ %d = add i32 %c, %nctaid.y
+ %e = add i32 %d, %nctaid.z
+ %f = add i32 %e, %rank
+ %g = add i32 %f, %nrank
+ ret i32 %g
+}
+
+; When .reqnctapercluster specifies only 1D, y and z default to 1.
+define ptx_kernel i32 @test_cluster_dim_1d() "nvvm.cluster_dim"="4" {
+; CHECK-LABEL: define ptx_kernel i32 @test_cluster_dim_1d(
+; CHECK-SAME: ) #[[ATTR1:[0-9]+]] {
+; CHECK-NEXT: [[CTAID_X:%.*]] = call range(i32 0, 4) i32 @llvm.nvvm.read.ptx.sreg.cluster.ctaid.x()
+; CHECK-NEXT: [[CTAID_Y:%.*]] = call range(i32 0, 1) i32 @llvm.nvvm.read.ptx.sreg.cluster.ctaid.y()
+; CHECK-NEXT: [[CTAID_Z:%.*]] = call range(i32 0, 1) i32 @llvm.nvvm.read.ptx.sreg.cluster.ctaid.z()
+; CHECK-NEXT: [[NCTAID_X:%.*]] = call range(i32 4, 5) i32 @llvm.nvvm.read.ptx.sreg.cluster.nctaid.x()
+; CHECK-NEXT: [[NCTAID_Y:%.*]] = call range(i32 1, 2) i32 @llvm.nvvm.read.ptx.sreg.cluster.nctaid.y()
+; CHECK-NEXT: [[NCTAID_Z:%.*]] = call range(i32 1, 2) i32 @llvm.nvvm.read.ptx.sreg.cluster.nctaid.z()
+; CHECK-NEXT: [[A:%.*]] = add i32 [[CTAID_X]], [[CTAID_Y]]
+; CHECK-NEXT: [[B:%.*]] = add i32 [[A]], [[CTAID_Z]]
+; CHECK-NEXT: [[C:%.*]] = add i32 [[B]], [[NCTAID_X]]
+; CHECK-NEXT: [[D:%.*]] = add i32 [[C]], [[NCTAID_Y]]
+; CHECK-NEXT: [[E:%.*]] = add i32 [[D]], [[NCTAID_Z]]
+; CHECK-NEXT: ret i32 [[E]]
+;
+ %ctaid.x = call i32 @llvm.nvvm.read.ptx.sreg.cluster.ctaid.x()
+ %ctaid.y = call i32 @llvm.nvvm.read.ptx.sreg.cluster.ctaid.y()
+ %ctaid.z = call i32 @llvm.nvvm.read.ptx.sreg.cluster.ctaid.z()
+ %nctaid.x = call i32 @llvm.nvvm.read.ptx.sreg.cluster.nctaid.x()
+ %nctaid.y = call i32 @llvm.nvvm.read.ptx.sreg.cluster.nctaid.y()
+ %nctaid.z = call i32 @llvm.nvvm.read.ptx.sreg.cluster.nctaid.z()
+ %a = add i32 %ctaid.x, %ctaid.y
+ %b = add i32 %a, %ctaid.z
+ %c = add i32 %b, %nctaid.x
+ %d = add i32 %c, %nctaid.y
+ %e = add i32 %d, %nctaid.z
+ ret i32 %e
+}
+
+; maxclusterrank is still an upper bound, not an exact cluster shape.
+define ptx_kernel i32 @test_maxclusterrank() "nvvm.maxclusterrank"="8" {
+; CHECK-LABEL: define ptx_kernel i32 @test_maxclusterrank(
+; CHECK-SAME: ) #[[ATTR2:[0-9]+]] {
+; CHECK-NEXT: [[NCTAID_X:%.*]] = call range(i32 1, 9) i32 @llvm.nvvm.read.ptx.sreg.cluster.nctaid.x()
+; CHECK-NEXT: [[NRANK:%.*]] = call range(i32 1, 9) i32 @llvm.nvvm.read.ptx.sreg.cluster.nctarank()
+; CHECK-NEXT: [[A:%.*]] = add i32 [[NCTAID_X]], [[NRANK]]
+; CHECK-NEXT: ret i32 [[A]]
+;
+ %nctaid.x = call i32 @llvm.nvvm.read.ptx.sreg.cluster.nctaid.x()
+ %nrank = call i32 @llvm.nvvm.read.ptx.sreg.cluster.nctarank()
+ %a = add i32 %nctaid.x, %nrank
+ ret i32 %a
+}
+
+declare i32 @llvm.nvvm.read.ptx.sreg.cluster.ctaid.x()
+declare i32 @llvm.nvvm.read.ptx.sreg.cluster.ctaid.y()
+declare i32 @llvm.nvvm.read.ptx.sreg.cluster.ctaid.z()
+declare i32 @llvm.nvvm.read.ptx.sreg.cluster.nctaid.x()
+declare i32 @llvm.nvvm.read.ptx.sreg.cluster.nctaid.y()
+declare i32 @llvm.nvvm.read.ptx.sreg.cluster.nctaid.z()
+declare i32 @llvm.nvvm.read.ptx.sreg.cluster.ctarank()
+declare i32 @llvm.nvvm.read.ptx.sreg.cluster.nctarank()
More information about the llvm-commits
mailing list