[llvm] [AMDGPU] Fix incorrect grid_dims constant folding for reqd_work_group_size (PR #211285)
Arseniy Obolenskiy via llvm-commits
llvm-commits at lists.llvm.org
Wed Jul 22 08:10:26 PDT 2026
https://github.com/aobolensk created https://github.com/llvm/llvm-project/pull/211285
reqd_work_group_size(X, Y, 1) does not guarantee the dispatch runs with work_dim == 2, since the spec still allows enqueuing with work_dim == 3 and a unit-size Z dimension. Only fold hidden_grid_dims to a constant when Z != 1 (work_dim must be 3)
Otherwise tighten the range instead of assuming an exact value
>From 5a690619261334ee1eee6f30712a84a75a6932c1 Mon Sep 17 00:00:00 2001
From: Arseniy Obolenskiy <arseniy.obolenskiy at amd.com>
Date: Wed, 22 Jul 2026 17:07:48 +0200
Subject: [PATCH] [AMDGPU] Fix incorrect grid_dims constant folding for
reqd_work_group_size
reqd_work_group_size(X, Y, 1) does not guarantee the dispatch runs with work_dim == 2, since the spec still allows enqueuing with work_dim == 3 and a unit-size Z dimension. Only fold hidden_grid_dims to a constant when Z != 1 (work_dim must be 3)
Otherwise tighten the range instead of assuming an exact value.
---
.../AMDGPU/AMDGPULowerKernelAttributes.cpp | 18 +++++++++--------
.../CodeGen/AMDGPU/implicit-arg-v5-opt.ll | 20 ++++++++++++++-----
2 files changed, 25 insertions(+), 13 deletions(-)
diff --git a/llvm/lib/Target/AMDGPU/AMDGPULowerKernelAttributes.cpp b/llvm/lib/Target/AMDGPU/AMDGPULowerKernelAttributes.cpp
index 32dbbfe6b65c8..2b4a4872d3bbe 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPULowerKernelAttributes.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPULowerKernelAttributes.cpp
@@ -124,9 +124,8 @@ static bool annotateGridDimsLoadWithRangeMD(LoadInst *Load,
if (!Ty || Ty->getBitWidth() < 3)
return false;
- if (KnownNumGridDims != 0) {
- Load->replaceAllUsesWith(
- ConstantInt::get(Load->getType(), KnownNumGridDims));
+ if (KnownNumGridDims == 3) {
+ Load->replaceAllUsesWith(ConstantInt::get(Load->getType(), 3));
return true;
}
@@ -134,15 +133,18 @@ static bool annotateGridDimsLoadWithRangeMD(LoadInst *Load,
if (Load->hasMetadata(LLVMContext::MD_range))
return false;
+ unsigned LowerBound = KnownNumGridDims == 2 ? 2 : 1;
MDBuilder MDB(Load->getContext());
- MDNode *Range =
- MDB.createRange(APInt(Ty->getBitWidth(), 1), APInt(Ty->getBitWidth(), 4));
+ MDNode *Range = MDB.createRange(APInt(Ty->getBitWidth(), LowerBound),
+ APInt(Ty->getBitWidth(), 4));
Load->setMetadata(LLVMContext::MD_range, Range);
return true;
}
-/// Compute the number of grid dimensions based on !reqd_work_group_size
-/// metadata
+/// Compute the known number of grid dimensions based on !reqd_work_group_size
+/// metadata. Returns 3 if the grid is known to be exactly 3-D, 2 if it is
+/// known to be at least 2-D, or 0 if nothing more than the default [1, 3]
+/// range can be deduced.
static unsigned computeNumGridDims(const MDNode *ReqdWorkGroupSize) {
ConstantInt *KnownZ =
mdconst::extract<ConstantInt>(ReqdWorkGroupSize->getOperand(2));
@@ -154,7 +156,7 @@ static unsigned computeNumGridDims(const MDNode *ReqdWorkGroupSize) {
if (KnownY->getZExtValue() != 1)
return 2;
- return 1;
+ return 0;
}
static bool processUse(CallInst *CI, bool IsV5OrAbove) {
diff --git a/llvm/test/CodeGen/AMDGPU/implicit-arg-v5-opt.ll b/llvm/test/CodeGen/AMDGPU/implicit-arg-v5-opt.ll
index f334411cdf68a..f1ee94828089e 100644
--- a/llvm/test/CodeGen/AMDGPU/implicit-arg-v5-opt.ll
+++ b/llvm/test/CodeGen/AMDGPU/implicit-arg-v5-opt.ll
@@ -390,7 +390,10 @@ define i32 @get_grid_dims_i32() #2 {
define i16 @get_grid_dims_reqd_work_group_size_1d() #3 !reqd_work_group_size !2 {
; GCN-LABEL: @get_grid_dims_reqd_work_group_size_1d(
-; GCN-NEXT: ret i16 1
+; GCN-NEXT: [[IMPLICITARG_PTR:%.*]] = tail call dereferenceable(256) ptr addrspace(4) @llvm.amdgcn.implicitarg.ptr()
+; GCN-NEXT: [[GEP_GRID_DIMS:%.*]] = getelementptr inbounds nuw i8, ptr addrspace(4) [[IMPLICITARG_PTR]], i64 64
+; GCN-NEXT: [[GRID_DIMS:%.*]] = load i16, ptr addrspace(4) [[GEP_GRID_DIMS]], align 4, !range [[RNG5]]
+; GCN-NEXT: ret i16 [[GRID_DIMS]]
;
%implicitarg.ptr = tail call ptr addrspace(4) @llvm.amdgcn.implicitarg.ptr()
%gep.grid.dims = getelementptr inbounds i8, ptr addrspace(4) %implicitarg.ptr, i64 64
@@ -400,7 +403,10 @@ define i16 @get_grid_dims_reqd_work_group_size_1d() #3 !reqd_work_group_size !2
define i16 @get_grid_dims_reqd_work_group_size_2d() #4 !reqd_work_group_size !3 {
; GCN-LABEL: @get_grid_dims_reqd_work_group_size_2d(
-; GCN-NEXT: ret i16 2
+; GCN-NEXT: [[IMPLICITARG_PTR:%.*]] = tail call dereferenceable(256) ptr addrspace(4) @llvm.amdgcn.implicitarg.ptr()
+; GCN-NEXT: [[GEP_GRID_DIMS:%.*]] = getelementptr inbounds nuw i8, ptr addrspace(4) [[IMPLICITARG_PTR]], i64 64
+; GCN-NEXT: [[GRID_DIMS:%.*]] = load i16, ptr addrspace(4) [[GEP_GRID_DIMS]], align 4, !range [[RNG10:![0-9]+]]
+; GCN-NEXT: ret i16 [[GRID_DIMS]]
;
%implicitarg.ptr = tail call ptr addrspace(4) @llvm.amdgcn.implicitarg.ptr()
%gep.grid.dims = getelementptr inbounds i8, ptr addrspace(4) %implicitarg.ptr, i64 64
@@ -410,7 +416,10 @@ define i16 @get_grid_dims_reqd_work_group_size_2d() #4 !reqd_work_group_size !3
define i16 @get_grid_dims_reqd_work_group_size_2d_weird() #5 !reqd_work_group_size !5 {
; GCN-LABEL: @get_grid_dims_reqd_work_group_size_2d_weird(
-; GCN-NEXT: ret i16 2
+; GCN-NEXT: [[IMPLICITARG_PTR:%.*]] = tail call dereferenceable(256) ptr addrspace(4) @llvm.amdgcn.implicitarg.ptr()
+; GCN-NEXT: [[GEP_GRID_DIMS:%.*]] = getelementptr inbounds nuw i8, ptr addrspace(4) [[IMPLICITARG_PTR]], i64 64
+; GCN-NEXT: [[GRID_DIMS:%.*]] = load i16, ptr addrspace(4) [[GEP_GRID_DIMS]], align 4, !range [[RNG10]]
+; GCN-NEXT: ret i16 [[GRID_DIMS]]
;
%implicitarg.ptr = tail call ptr addrspace(4) @llvm.amdgcn.implicitarg.ptr()
%gep.grid.dims = getelementptr inbounds i8, ptr addrspace(4) %implicitarg.ptr, i64 64
@@ -455,7 +464,7 @@ define i16 @get_grid_dims_existing_range() #2 {
; GCN-LABEL: @get_grid_dims_existing_range(
; GCN-NEXT: [[IMPLICITARG_PTR:%.*]] = tail call dereferenceable(256) ptr addrspace(4) @llvm.amdgcn.implicitarg.ptr()
; GCN-NEXT: [[GEP_GRID_DIMS:%.*]] = getelementptr inbounds nuw i8, ptr addrspace(4) [[IMPLICITARG_PTR]], i64 64
-; GCN-NEXT: [[GRID_DIMS:%.*]] = load i16, ptr addrspace(4) [[GEP_GRID_DIMS]], align 4, !range [[RNG12:![0-9]+]]
+; GCN-NEXT: [[GRID_DIMS:%.*]] = load i16, ptr addrspace(4) [[GEP_GRID_DIMS]], align 4, !range [[RNG13:![0-9]+]]
; GCN-NEXT: ret i16 [[GRID_DIMS]]
;
%implicitarg.ptr = tail call ptr addrspace(4) @llvm.amdgcn.implicitarg.ptr()
@@ -493,5 +502,6 @@ attributes #6 = { nounwind "amdgpu-flat-work-group-size"="256,256" }
; GCN: [[RNG5]] = !{i16 1, i16 4}
; GCN: [[RNG6]] = !{i8 1, i8 4}
; GCN: [[RNG7]] = !{i3 1, i3 -4}
-; GCN: [[RNG12]] = !{i16 1, i16 2}
+; GCN: [[RNG10]] = !{i16 2, i16 4}
+; GCN: [[RNG13]] = !{i16 1, i16 2}
;.
More information about the llvm-commits
mailing list