[llvm] [SROA] Allow as zext<i1> index when unfolding GEP select (PR #146929)
Alex MacLean via llvm-commits
llvm-commits at lists.llvm.org
Thu Jul 3 20:25:10 PDT 2025
https://github.com/AlexMaclean updated https://github.com/llvm/llvm-project/pull/146929
>From 65dad0e6bb7979c26eddb6f605896c687799543e Mon Sep 17 00:00:00 2001
From: Alex Maclean <amaclean at nvidia.com>
Date: Thu, 3 Jul 2025 03:56:33 +0000
Subject: [PATCH 1/2] [SROA] Allow as zext<i1> index when unfolding GEP select
---
llvm/lib/Transforms/Scalar/SROA.cpp | 31 +++++++++++++++++++------
llvm/test/Transforms/SROA/select-gep.ll | 18 ++++++++++++++
2 files changed, 42 insertions(+), 7 deletions(-)
diff --git a/llvm/lib/Transforms/Scalar/SROA.cpp b/llvm/lib/Transforms/Scalar/SROA.cpp
index 42d1d9a437bb2..70b4552190a4e 100644
--- a/llvm/lib/Transforms/Scalar/SROA.cpp
+++ b/llvm/lib/Transforms/Scalar/SROA.cpp
@@ -4070,18 +4070,27 @@ class AggLoadStoreRewriter : public InstVisitor<AggLoadStoreRewriter, bool> {
// => select cond, gep(ptr1, idx), gep(ptr2, idx)
// and gep ptr, (select cond, idx1, idx2)
// => select cond, gep(ptr, idx1), gep(ptr, idx2)
+ // We also allow for i1 zext indices, which are equivalent to selects.
bool unfoldGEPSelect(GetElementPtrInst &GEPI) {
// Check whether the GEP has exactly one select operand and all indices
// will become constant after the transform.
- SelectInst *Sel = dyn_cast<SelectInst>(GEPI.getPointerOperand());
+ Instruction *Sel = dyn_cast<SelectInst>(GEPI.getPointerOperand());
for (Value *Op : GEPI.indices()) {
if (auto *SI = dyn_cast<SelectInst>(Op)) {
if (Sel)
return false;
Sel = SI;
- if (!isa<ConstantInt>(Sel->getTrueValue()) ||
- !isa<ConstantInt>(Sel->getFalseValue()))
+ if (!isa<ConstantInt>(SI->getTrueValue()) ||
+ !isa<ConstantInt>(SI->getFalseValue()))
+ return false;
+ continue;
+ }
+ if (auto *ZI = dyn_cast<ZExtInst>(Op)) {
+ if (Sel)
+ return false;
+ Sel = ZI;
+ if (!ZI->getSrcTy()->isIntegerTy(1))
return false;
continue;
}
@@ -4107,8 +4116,16 @@ class AggLoadStoreRewriter : public InstVisitor<AggLoadStoreRewriter, bool> {
return NewOps;
};
- Value *True = Sel->getTrueValue();
- Value *False = Sel->getFalseValue();
+ Value *Cond, *True, *False;
+ if (auto *SI = dyn_cast<SelectInst>(Sel)) {
+ Cond = SI->getCondition();
+ True = SI->getTrueValue();
+ False = SI->getFalseValue();
+ } else {
+ Cond = Sel->getOperand(0);
+ True = ConstantInt::get(Sel->getType(), 1);
+ False = ConstantInt::get(Sel->getType(), 0);
+ }
SmallVector<Value *> TrueOps = GetNewOps(True);
SmallVector<Value *> FalseOps = GetNewOps(False);
@@ -4123,8 +4140,8 @@ class AggLoadStoreRewriter : public InstVisitor<AggLoadStoreRewriter, bool> {
IRB.CreateGEP(Ty, FalseOps[0], ArrayRef(FalseOps).drop_front(),
False->getName() + ".sroa.gep", NW);
- Value *NSel = IRB.CreateSelect(Sel->getCondition(), NTrue, NFalse,
- Sel->getName() + ".sroa.sel");
+ Value *NSel =
+ IRB.CreateSelect(Cond, NTrue, NFalse, Sel->getName() + ".sroa.sel");
Visited.erase(&GEPI);
GEPI.replaceAllUsesWith(NSel);
GEPI.eraseFromParent();
diff --git a/llvm/test/Transforms/SROA/select-gep.ll b/llvm/test/Transforms/SROA/select-gep.ll
index 1342a2ca4ea2b..56c2b2fd87bd8 100644
--- a/llvm/test/Transforms/SROA/select-gep.ll
+++ b/llvm/test/Transforms/SROA/select-gep.ll
@@ -201,6 +201,24 @@ define i32 @test_select_idx_mem2reg(i1 %c) {
ret i32 %res
}
+; Test gep with a select-like zext index unfolding on an alloca that is
+; splittable and promotable.
+define i64 @test_select_like_zext_idx_mem2reg(i1 %c) {
+; CHECK-LABEL: @test_select_like_zext_idx_mem2reg(
+; CHECK-NEXT: [[IDX:%.*]] = zext i1 [[C:%.*]] to i64
+; CHECK-NEXT: [[RES:%.*]] = select i1 [[C]], i64 2, i64 1
+; CHECK-NEXT: ret i64 [[RES]]
+;
+ %alloca = alloca [2 x i64], align 8
+ store i64 1, ptr %alloca
+ %gep1 = getelementptr inbounds i64, ptr %alloca, i64 1
+ store i64 2, ptr %gep1
+ %idx = zext i1 %c to i64
+ %gep2 = getelementptr inbounds i64, ptr %alloca, i64 %idx
+ %res = load i64, ptr %gep2
+ ret i64 %res
+}
+
; Test gep of index select unfolding on an alloca that escaped, and as such
; is not splittable or promotable.
; FIXME: Ideally, no transform would take place in this case.
>From 83d8478dd5665f2f0b9588a76137163d2717def9 Mon Sep 17 00:00:00 2001
From: Alex Maclean <amaclean at nvidia.com>
Date: Fri, 4 Jul 2025 03:24:49 +0000
Subject: [PATCH 2/2] address comments
---
llvm/test/Transforms/SROA/select-gep.ll | 23 +++++++++++++++++++++++
1 file changed, 23 insertions(+)
diff --git a/llvm/test/Transforms/SROA/select-gep.ll b/llvm/test/Transforms/SROA/select-gep.ll
index 56c2b2fd87bd8..b48b0f77aa991 100644
--- a/llvm/test/Transforms/SROA/select-gep.ll
+++ b/llvm/test/Transforms/SROA/select-gep.ll
@@ -219,6 +219,29 @@ define i64 @test_select_like_zext_idx_mem2reg(i1 %c) {
ret i64 %res
}
+; Test gep with a zext index that is not equivalent to a select. No unfolding
+; or promotion should take place.
+define i64 @test_zext_unlike_select_idx_mem2reg(i8 %c) {
+; CHECK-LABEL: @test_zext_unlike_select_idx_mem2reg(
+; CHECK-NEXT: [[ALLOCA:%.*]] = alloca [2 x i64], align 8
+; CHECK-NEXT: store i64 1, ptr [[ALLOCA]], align 4
+; CHECK-NEXT: [[GEP1:%.*]] = getelementptr inbounds i64, ptr [[ALLOCA]], i64 1
+; CHECK-NEXT: store i64 2, ptr [[GEP1]], align 4
+; CHECK-NEXT: [[IDX:%.*]] = zext i8 [[C:%.*]] to i64
+; CHECK-NEXT: [[GEP2:%.*]] = getelementptr inbounds i64, ptr [[ALLOCA]], i64 [[IDX]]
+; CHECK-NEXT: [[RES:%.*]] = load i64, ptr [[GEP2]], align 4
+; CHECK-NEXT: ret i64 [[RES]]
+;
+ %alloca = alloca [2 x i64], align 8
+ store i64 1, ptr %alloca
+ %gep1 = getelementptr inbounds i64, ptr %alloca, i64 1
+ store i64 2, ptr %gep1
+ %idx = zext i8 %c to i64
+ %gep2 = getelementptr inbounds i64, ptr %alloca, i64 %idx
+ %res = load i64, ptr %gep2
+ ret i64 %res
+}
+
; Test gep of index select unfolding on an alloca that escaped, and as such
; is not splittable or promotable.
; FIXME: Ideally, no transform would take place in this case.
More information about the llvm-commits
mailing list