[SROA] Allow as zext<i1> index when unfolding GEP select (#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.

Author: AlexMaclean

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();

llvm/test/Transforms/SROA/select-gep.ll

@@ -201,6 +201,47 @@ 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 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.