[llvm] [SROA] Allow as zext<i1> index when unfolding GEP select (PR #146929)

[Alex MacLean via llvm-commits]

https://github.com/AlexMaclean created https://github.com/llvm/llvm-project/pull/146929

A zero-extension from an i1 is equivalent to a select with constant 0 and 1 values. Add this case when rewriting gep(select) -> select(gep) to expose more opportunities for SROA.

>From 60fd6bbaa5b8108c652a3739085b37b82d6186fb 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] [SROA] Allow as zext<i1> index when unfolding GEP select

---
 llvm/lib/Transforms/Scalar/SROA.cpp     | 29 ++++++++++++++++++++-----
 llvm/test/Transforms/SROA/select-gep.ll | 18 +++++++++++++++
 2 files changed, 41 insertions(+), 6 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/SROA.cpp b/llvm/lib/Transforms/Scalar/SROA.cpp
index 42d1d9a437bb2..dd2e83ca5ecec 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,7 +4140,7 @@ 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,
+    Value *NSel = IRB.CreateSelect(Cond, NTrue, NFalse,
                                    Sel->getName() + ".sroa.sel");
     Visited.erase(&GEPI);
     GEPI.replaceAllUsesWith(NSel);
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.