[llvm] [AMDGPU] Add more cases to the "this is a wave ID" recognizer (PR #177713)
Krzysztof Drewniak via llvm-commits
llvm-commits at lists.llvm.org
Mon Mar 9 08:29:44 PDT 2026
https://github.com/krzysz00 updated https://github.com/llvm/llvm-project/pull/177713
>From 1f659b636aa8272f99d42e42b160fdde7810ae86 Mon Sep 17 00:00:00 2001
From: Krzysztof Drewniak <Krzysztof.Drewniak at amd.com>
Date: Sat, 24 Jan 2026 00:27:37 +0000
Subject: [PATCH 1/4] [AMDGPU] Add more cases to the "this is a wave ID"
recognizer
Previously, we'd only catch things that were actually wave
IDs (workitem_id_x >> log2(wavefrontsize) or ands with a mask that
cleared off the within-lane bits ... and all that only in cases where
we had static assurances about the workgroup size) when the operation
was being performed directly on the ID. This commit adds a test for
the mask case and also extends the matcher to support
1. The intrinsic being cast to another integer type before being
shifted
2. A mask being applied to the intristic before the shift occurs, so
that the computation naturally written as (x % 128) / 64 which
optimizes to (x & 0x7f) >> 6 is properly recognized as wave-uniform.
I've noticed that there's a higher-level problem here. These sorts of
quasi-wave-IDs are often needed to compute, for example, the address
for an LDS DMA on gfx950 or will be used on gfx1250 when computing
parts of TDM descriptors. These values need to be in SGPRs, but the
VGPR-ness of the workitem ID causes all the address computation to be
forced into VGPRs until the final moment it's needed, wanting VGPRs
and putting readfristlanes into loops. However, this commit doesn't
try to address that issue.
---
llvm/lib/Target/AMDGPU/AMDGPUSubtarget.h | 2 +-
.../AMDGPU/AMDGPUTargetTransformInfo.cpp | 18 ++++---
.../AMDGPU/workitem-intrinsics.ll | 49 +++++++++++++++++++
3 files changed, 62 insertions(+), 7 deletions(-)
diff --git a/llvm/lib/Target/AMDGPU/AMDGPUSubtarget.h b/llvm/lib/Target/AMDGPU/AMDGPUSubtarget.h
index 4e1e54da5f4a9..24c392a5765b1 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUSubtarget.h
+++ b/llvm/lib/Target/AMDGPU/AMDGPUSubtarget.h
@@ -97,7 +97,7 @@ class AMDGPUSubtarget {
/// This information is currently only gathered from the !reqd_work_group_size
/// metadata on \p F, but this may be improved in the future.
bool hasWavefrontsEvenlySplittingXDim(const Function &F,
- bool REquiresUniformYZ = false) const;
+ bool RequiresUniformYZ = false) const;
/// \returns Subtarget's default pair of minimum/maximum number of waves per
/// execution unit for function \p F, or minimum/maximum number of waves per
diff --git a/llvm/lib/Target/AMDGPU/AMDGPUTargetTransformInfo.cpp b/llvm/lib/Target/AMDGPU/AMDGPUTargetTransformInfo.cpp
index d746ce65a6288..5e5e07213594b 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUTargetTransformInfo.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUTargetTransformInfo.cpp
@@ -1095,16 +1095,22 @@ bool GCNTTIImpl::isAlwaysUniform(const Value *V) const {
}
using namespace llvm::PatternMatch;
uint64_t C;
- if (match(V, m_LShr(m_Intrinsic<Intrinsic::amdgcn_workitem_id_x>(),
- m_ConstantInt(C))) ||
- match(V, m_AShr(m_Intrinsic<Intrinsic::amdgcn_workitem_id_x>(),
- m_ConstantInt(C)))) {
+ auto MatchTidXCall = m_Intrinsic<Intrinsic::amdgcn_workitem_id_x>();
+ ;
+ auto MaybeMaskedTidX =
+ m_CombineOr(m_c_And(MatchTidXCall, m_Value()), MatchTidXCall);
+ auto MaybeCastTidX = m_CastOrSelf(MaybeMaskedTidX);
+ auto MaybeMaskedCastTidX =
+ m_CombineOr(m_c_And(MaybeCastTidX, m_Value()), MaybeCastTidX);
+ if (match(V, m_LShr(MaybeMaskedCastTidX, m_ConstantInt(C))) ||
+ match(V, m_AShr(MaybeMaskedCastTidX, m_ConstantInt(C)))) {
return C >= ST->getWavefrontSizeLog2() && XDimDoesntResetWithinWaves;
}
Value *Mask;
- if (match(V, m_c_And(m_Intrinsic<Intrinsic::amdgcn_workitem_id_x>(),
- m_Value(Mask)))) {
+ if (match(V, m_c_And(
+ m_CastOrSelf(m_Intrinsic<Intrinsic::amdgcn_workitem_id_x>()),
+ m_Value(Mask)))) {
return computeKnownBits(Mask, DL).countMinTrailingZeros() >=
ST->getWavefrontSizeLog2() &&
XDimDoesntResetWithinWaves;
diff --git a/llvm/test/Analysis/UniformityAnalysis/AMDGPU/workitem-intrinsics.ll b/llvm/test/Analysis/UniformityAnalysis/AMDGPU/workitem-intrinsics.ll
index f5668cef5d63e..29edcdfd7c0c5 100644
--- a/llvm/test/Analysis/UniformityAnalysis/AMDGPU/workitem-intrinsics.ll
+++ b/llvm/test/Analysis/UniformityAnalysis/AMDGPU/workitem-intrinsics.ll
@@ -131,6 +131,54 @@ define amdgpu_kernel void @workitem_id_x_div_wavefront_size(ptr %o) #3 !reqd_wor
ret void
}
+; CHECK-LABEL: UniformityInfo for function 'workitem_id_x_masked'
+; CHECK: DIVERGENT: %id.x = call i32 @llvm.amdgcn.workitem.id.x()
+; CHECK-NOT: DIVERGENT
+define amdgpu_kernel void @workitem_id_x_masked(ptr %o) #3 !reqd_work_group_size !5 {
+ %id.x = call i32 @llvm.amdgcn.workitem.id.x()
+ %id.sg.shl.6 = and i32 %id.x, 192
+ store i32 %id.sg.shl.6, ptr %o
+ ret void
+}
+
+; CHECK-LABEL: UniformityInfo for function 'workitem_id_x_zext_masked'
+; CHECK: DIVERGENT: %id.x = call i32 @llvm.amdgcn.workitem.id.x()
+; CHECK: DIVERGENT: %id.x.zext = zext i32 %id.x to i64
+; CHECK-NOT: DIVERGENT
+define amdgpu_kernel void @workitem_id_x_zext_masked(ptr %o) #3 !reqd_work_group_size !5 {
+ %id.x = call i32 @llvm.amdgcn.workitem.id.x()
+ %id.x.zext = zext i32 %id.x to i64
+ %id.sg.shl.6 = and i64 %id.x.zext, 192
+ store i64 %id.sg.shl.6, ptr %o
+ ret void
+}
+
+; CHECK-LABEL: UniformityInfo for function 'workitem_id_x_div_wavefront_size_masked'
+; CHECK: DIVERGENT: %id.x = call i32 @llvm.amdgcn.workitem.id.x()
+; CHECK: DIVERGENT: %id.x.masked = and i32 %id.x, 127
+; CHECK-NOT: DIVERGENT
+define amdgpu_kernel void @workitem_id_x_div_wavefront_size_masked(ptr %o) #3 !reqd_work_group_size !5 {
+ %id.x = call i32 @llvm.amdgcn.workitem.id.x()
+ %id.x.masked = and i32 %id.x, 127
+ %id.sg = lshr i32 %id.x.masked, 6
+ store i32 %id.sg, ptr %o
+ ret void
+}
+
+; CHECK-LABEL: UniformityInfo for function 'workitem_id_x_div_wavefront_size_trunc_masked'
+; CHECK: DIVERGENT: %id.x = call i32 @llvm.amdgcn.workitem.id.x()
+; CHECK: DIVERGENT: %id.x.trunc = trunc i32 %id.x to i16
+; CHECK: DIVERGENT: %id.x.masked = and i16 %id.x.trunc, 127
+; CHECK-NOT: DIVERGENT
+define amdgpu_kernel void @workitem_id_x_div_wavefront_size_trunc_masked(ptr %o) #3 !reqd_work_group_size !5 {
+ %id.x = call i32 @llvm.amdgcn.workitem.id.x()
+ %id.x.trunc = trunc i32 %id.x to i16
+ %id.x.masked = and i16 %id.x.trunc, 127
+ %id.sg = lshr i16 %id.x.masked, 6
+ store i16 %id.sg, ptr %o
+ ret void
+}
+
; CHECK-LABEL: UniformityInfo for function 'workitem_id_y_uniform_in_subgroup'
; CHECK-NOT: DIVERGENT
define amdgpu_kernel void @workitem_id_y_uniform_in_subgroup(ptr %o) #3 !reqd_work_group_size !5 {
@@ -150,3 +198,4 @@ attributes #3 = { "target-cpu"="gfx900" "amdgpu-flat-work-group-size"="256,256"
!3 = !{i32 1, i32 1, i32 2}
!4 = !{i32 64, i32 1, i32 1}
!5 = !{i32 128, i32 2, i32 1}
+!6 = !{i32 256, i32 1, i32 1}
>From 80a981ee4de351b1f2ce9a4a39d4ccacc6704b0c Mon Sep 17 00:00:00 2001
From: Krzysztof Drewniak <Krzysztof.Drewniak at amd.com>
Date: Wed, 4 Feb 2026 19:35:49 +0000
Subject: [PATCH 2/4] Make tests more canonical
---
.../AMDGPU/AMDGPUTargetTransformInfo.cpp | 1 -
.../AMDGPU/workitem-intrinsics.ll | 20 +++++++++++++------
2 files changed, 14 insertions(+), 7 deletions(-)
diff --git a/llvm/lib/Target/AMDGPU/AMDGPUTargetTransformInfo.cpp b/llvm/lib/Target/AMDGPU/AMDGPUTargetTransformInfo.cpp
index 5e5e07213594b..8c762c9e30507 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUTargetTransformInfo.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUTargetTransformInfo.cpp
@@ -1096,7 +1096,6 @@ bool GCNTTIImpl::isAlwaysUniform(const Value *V) const {
using namespace llvm::PatternMatch;
uint64_t C;
auto MatchTidXCall = m_Intrinsic<Intrinsic::amdgcn_workitem_id_x>();
- ;
auto MaybeMaskedTidX =
m_CombineOr(m_c_And(MatchTidXCall, m_Value()), MatchTidXCall);
auto MaybeCastTidX = m_CastOrSelf(MaybeMaskedTidX);
diff --git a/llvm/test/Analysis/UniformityAnalysis/AMDGPU/workitem-intrinsics.ll b/llvm/test/Analysis/UniformityAnalysis/AMDGPU/workitem-intrinsics.ll
index 29edcdfd7c0c5..6434b76441fc1 100644
--- a/llvm/test/Analysis/UniformityAnalysis/AMDGPU/workitem-intrinsics.ll
+++ b/llvm/test/Analysis/UniformityAnalysis/AMDGPU/workitem-intrinsics.ll
@@ -143,13 +143,16 @@ define amdgpu_kernel void @workitem_id_x_masked(ptr %o) #3 !reqd_work_group_size
; CHECK-LABEL: UniformityInfo for function 'workitem_id_x_zext_masked'
; CHECK: DIVERGENT: %id.x = call i32 @llvm.amdgcn.workitem.id.x()
-; CHECK: DIVERGENT: %id.x.zext = zext i32 %id.x to i64
+; CHECK: DIVERGENT: %id.x.zext = zext nneg i32 %id.x to i64
+; CHECK: DIVERGENT: %my.out = getelementptr i64, ptr %o, i64 %id.x.zext
; CHECK-NOT: DIVERGENT
+; CHECK: DIVERGENT: store i64 %id.sg.shl.6, ptr %my.out
define amdgpu_kernel void @workitem_id_x_zext_masked(ptr %o) #3 !reqd_work_group_size !5 {
%id.x = call i32 @llvm.amdgcn.workitem.id.x()
- %id.x.zext = zext i32 %id.x to i64
+ %id.x.zext = zext nneg i32 %id.x to i64
+ %my.out = getelementptr i64, ptr %o, i64 %id.x.zext
%id.sg.shl.6 = and i64 %id.x.zext, 192
- store i64 %id.sg.shl.6, ptr %o
+ store i64 %id.sg.shl.6, ptr %my.out
ret void
}
@@ -167,15 +170,20 @@ define amdgpu_kernel void @workitem_id_x_div_wavefront_size_masked(ptr %o) #3 !r
; CHECK-LABEL: UniformityInfo for function 'workitem_id_x_div_wavefront_size_trunc_masked'
; CHECK: DIVERGENT: %id.x = call i32 @llvm.amdgcn.workitem.id.x()
-; CHECK: DIVERGENT: %id.x.trunc = trunc i32 %id.x to i16
+; CHECK: DIVERGENT: %id.x.trunc = trunc nuw nsw i32 %id.x to i16
; CHECK: DIVERGENT: %id.x.masked = and i16 %id.x.trunc, 127
+; CHECK: DIVERGENT: %offset = zext nneg i16 %id.x.masked to i64
+; CHECK: DIVERGENT: %my.out = getelementptr i16, ptr %o, i64 %offset
; CHECK-NOT: DIVERGENT
+; CHECK: DIVERGENT: store i16 %id.sg, ptr %my.out
define amdgpu_kernel void @workitem_id_x_div_wavefront_size_trunc_masked(ptr %o) #3 !reqd_work_group_size !5 {
%id.x = call i32 @llvm.amdgcn.workitem.id.x()
- %id.x.trunc = trunc i32 %id.x to i16
+ %id.x.trunc = trunc nuw nsw i32 %id.x to i16
%id.x.masked = and i16 %id.x.trunc, 127
+ %offset = zext nneg i16 %id.x.masked to i64
+ %my.out = getelementptr i16, ptr %o, i64 %offset
%id.sg = lshr i16 %id.x.masked, 6
- store i16 %id.sg, ptr %o
+ store i16 %id.sg, ptr %my.out
ret void
}
>From 460710397adc92cb0f812eacd65427a166bf50a9 Mon Sep 17 00:00:00 2001
From: Krzysztof Drewniak <Krzysztof.Drewniak at amd.com>
Date: Mon, 9 Feb 2026 17:15:32 +0000
Subject: [PATCH 3/4] Remove the ashr case as non-canononical
---
llvm/lib/Target/AMDGPU/AMDGPUTargetTransformInfo.cpp | 3 +--
1 file changed, 1 insertion(+), 2 deletions(-)
diff --git a/llvm/lib/Target/AMDGPU/AMDGPUTargetTransformInfo.cpp b/llvm/lib/Target/AMDGPU/AMDGPUTargetTransformInfo.cpp
index 8c762c9e30507..5feac978241d0 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUTargetTransformInfo.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUTargetTransformInfo.cpp
@@ -1101,8 +1101,7 @@ bool GCNTTIImpl::isAlwaysUniform(const Value *V) const {
auto MaybeCastTidX = m_CastOrSelf(MaybeMaskedTidX);
auto MaybeMaskedCastTidX =
m_CombineOr(m_c_And(MaybeCastTidX, m_Value()), MaybeCastTidX);
- if (match(V, m_LShr(MaybeMaskedCastTidX, m_ConstantInt(C))) ||
- match(V, m_AShr(MaybeMaskedCastTidX, m_ConstantInt(C)))) {
+ if (match(V, m_LShr(MaybeMaskedCastTidX, m_ConstantInt(C)))) {
return C >= ST->getWavefrontSizeLog2() && XDimDoesntResetWithinWaves;
}
>From 0255c59d7ca450b596c15f0c83010bf68c9d2fce Mon Sep 17 00:00:00 2001
From: Krzysztof Drewniak <Krzysztof.Drewniak at amd.com>
Date: Fri, 6 Mar 2026 18:46:30 +0000
Subject: [PATCH 4/4] Put ASHR back, people might be relying on it.
---
llvm/lib/Target/AMDGPU/AMDGPUTargetTransformInfo.cpp | 3 ++-
1 file changed, 2 insertions(+), 1 deletion(-)
diff --git a/llvm/lib/Target/AMDGPU/AMDGPUTargetTransformInfo.cpp b/llvm/lib/Target/AMDGPU/AMDGPUTargetTransformInfo.cpp
index 5feac978241d0..8c762c9e30507 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUTargetTransformInfo.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUTargetTransformInfo.cpp
@@ -1101,7 +1101,8 @@ bool GCNTTIImpl::isAlwaysUniform(const Value *V) const {
auto MaybeCastTidX = m_CastOrSelf(MaybeMaskedTidX);
auto MaybeMaskedCastTidX =
m_CombineOr(m_c_And(MaybeCastTidX, m_Value()), MaybeCastTidX);
- if (match(V, m_LShr(MaybeMaskedCastTidX, m_ConstantInt(C)))) {
+ if (match(V, m_LShr(MaybeMaskedCastTidX, m_ConstantInt(C))) ||
+ match(V, m_AShr(MaybeMaskedCastTidX, m_ConstantInt(C)))) {
return C >= ST->getWavefrontSizeLog2() && XDimDoesntResetWithinWaves;
}
More information about the llvm-commits
mailing list