[llvm] [SROA] Unfold gep of index select (PR #80983)

[Artem Belevich via llvm-commits]

================
@@ -3937,30 +3937,63 @@ class AggLoadStoreRewriter : public InstVisitor<AggLoadStoreRewriter, bool> {
     return false;
   }
 
-  // Fold gep (select cond, ptr1, ptr2) => select cond, gep(ptr1), gep(ptr2)
+  // Fold gep (select cond, ptr1, ptr2), idx
+  //   => select cond, gep(ptr1, idx), gep(ptr2, idx)
+  // and  gep ptr, (select cond, idx1, idx2)
+  //   => select cond, gep(ptr, idx1), gep(ptr, idx2)
   bool foldGEPSelect(GetElementPtrInst &GEPI) {
-    if (!GEPI.hasAllConstantIndices())
-      return false;
+    // Check whether the GEP has exactly one select operand and all indices
+    // will become constant after the transform.
+    auto IsValidOp = [](Value *Op) {
+      return Op->getType()->isPointerTy() || isa<ConstantInt>(Op);
+    };
+
+    SelectInst *Sel = nullptr;
+    for (Value *Op : GEPI.operands()) {
+      if (auto *SI = dyn_cast<SelectInst>(Op)) {
----------------
Artem-B wrote:

> I believe only the "all constant" case is useful. 

Counter-example: If we had `gep(select(1,select(2,3)))`, it could be further split into `gep(const)` when we'd visited the split geps after the first iteration. 

`isa<ConstantInt>` is too conservative. I think the correct check here is whether LLVM can compute a known value. for each branch of select, as the select branches may be other ops, yet could be evaluated to an int value. 

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