[llvm] [SROA] Unfold gep of index phi (PR #83087)

[Artem Belevich via llvm-commits]

================
@@ -4023,64 +4023,88 @@ class AggLoadStoreRewriter : public InstVisitor<AggLoadStoreRewriter, bool> {
     Visited.insert(NSelI);
     enqueueUsers(*NSelI);
 
-    LLVM_DEBUG(dbgs() << "\n          to: " << *NTrue
-                      << "\n              " << *NFalse
-                      << "\n              " << *NSel << '\n');
+    LLVM_DEBUG(dbgs() << "\n          to: " << *NTrue << "\n              "
+                      << *NFalse << "\n              " << *NSel << '\n');
 
     return true;
   }
 
-  // Fold gep (phi ptr1, ptr2) => phi gep(ptr1), gep(ptr2)
+  // Fold gep (phi ptr1, ptr2), idx
+  //   => phi ((gep ptr1, idx), (gep ptr2, idx))
+  // and  gep ptr, (phi idx1, idx2)
+  //   => phi ((gep ptr, idx1), (gep ptr, idx2))
   bool foldGEPPhi(GetElementPtrInst &GEPI) {
-    if (!GEPI.hasAllConstantIndices())
+    // Check whether the GEP has exactly one phi operand and all indices
+    // will become constant after the transform.
+    PHINode *Phi = dyn_cast<PHINode>(GEPI.getPointerOperand());
+    // To prevent infinitely expanding recursive phis, bail if the GEP pointer
+    // operand is the phi and any of its incoming values is not an alloca or a
+    // constant.
+    if (Phi && any_of(Phi->operands(), [](Value *V) {
+          return isa<Instruction>(V) && !isa<AllocaInst>(V);
+        })) {
       return false;
+    }
+    for (Value *Op : GEPI.indices()) {
+      if (auto *SI = dyn_cast<PHINode>(Op)) {
+        if (Phi)
+          return false;
+
+        Phi = SI;
+        if (!all_of(Phi->incoming_values(),
+                    [](Value *V) { return isa<ConstantInt>(V); }))
+          return false;
----------------
Artem-B wrote:

There may be merit in folding if we see any constant values. It should still be useful for SROA, but there's obviously a trade-off vs having the non-const GEPs around.

https://github.com/llvm/llvm-project/pull/83087