[InstCombine] Extend folding of aggregate construction to cases when source aggregates are partially available (#100828)

Function foldAggregateConstructionIntoAggregateReuse can fold
  insertvalue(phi(extractvalue(src1), extractvalue(src2)))
into
  phi(src1, src2)
when we can find source aggregates in all predecessors.

This patch extends it to handle following case
  insertvalue(phi(extractvalue(src1), elm2))
into
  phi(src1, insertvalue(elm2))
with the condition that the predecessor without source aggregate has
only one successor.

Author: weiguozhi

llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp

@@ -963,6 +963,9 @@ Instruction *InstCombinerImpl::foldAggregateConstructionIntoAggregateReuse(
     return AggregateDescription::Found;
   };
 
+  // If an aggregate element is defined in UseBB, we can't use it in PredBB.
+  bool EltDefinedInUseBB = false;
+
   // Given the value \p Elt that was being inserted into element \p EltIdx of an
   // aggregate AggTy, see if \p Elt was originally defined by an
   // appropriate extractvalue (same element index, same aggregate type).
@@ -972,8 +975,11 @@ Instruction *InstCombinerImpl::foldAggregateConstructionIntoAggregateReuse(
       [&](Instruction *Elt, unsigned EltIdx, std::optional<BasicBlock *> UseBB,
           std::optional<BasicBlock *> PredBB) -> std::optional<Value *> {
     // For now(?), only deal with, at most, a single level of PHI indirection.
-    if (UseBB && PredBB)
+    if (UseBB && PredBB) {
       Elt = dyn_cast<Instruction>(Elt->DoPHITranslation(*UseBB, *PredBB));
+      if (Elt && Elt->getParent() == *UseBB)
+        EltDefinedInUseBB = true;
+    }
     // FIXME: deal with multiple levels of PHI indirection?
 
     // Did we find an extraction?
@@ -1106,6 +1112,7 @@ Instruction *InstCombinerImpl::foldAggregateConstructionIntoAggregateReuse(
   // from which all the elements were originally extracted from?
   // Note that we want for the map to have stable iteration order!
   SmallDenseMap<BasicBlock *, Value *, 4> SourceAggregates;
+  bool FoundSrcAgg = false;
   for (BasicBlock *Pred : Preds) {
     std::pair<decltype(SourceAggregates)::iterator, bool> IV =
         SourceAggregates.insert({Pred, nullptr});
@@ -1117,9 +1124,68 @@ Instruction *InstCombinerImpl::foldAggregateConstructionIntoAggregateReuse(
     // aggregate produced by OrigIVI must have been originally extracted from
     // the same aggregate. Is that so? Can we find said original aggregate?
     SourceAggregate = FindCommonSourceAggregate(UseBB, Pred);
-    if (Describe(SourceAggregate) != AggregateDescription::Found)
-      return nullptr; // Give up.
-    IV.first->second = *SourceAggregate;
+    if (Describe(SourceAggregate) == AggregateDescription::Found) {
+      FoundSrcAgg = true;
+      IV.first->second = *SourceAggregate;
+    } else {
+      // If UseBB is the single successor of Pred, we can add InsertValue to
+      // Pred.
+      auto *BI = dyn_cast<BranchInst>(Pred->getTerminator());
+      if (!BI || !BI->isUnconditional())
+        return nullptr;
+    }
+  }
+
+  if (!FoundSrcAgg)
+    return nullptr;
+
+  // Do some sanity check if we need to add insertvalue into predecessors.
+  auto OrigBB = OrigIVI.getParent();
+  for (auto &It : SourceAggregates) {
+    if (Describe(It.second) == AggregateDescription::Found)
+      continue;
+
+    // Element is defined in UseBB, so it can't be used in predecessors.
+    if (EltDefinedInUseBB)
+      return nullptr;
+
+    // Do this transformation cross loop boundary may cause dead loop. So we
+    // should avoid this situation. But LoopInfo is not generally available, we
+    // must be conservative here.
+    // If OrigIVI is in UseBB and it's the only successor of PredBB, PredBB
+    // can't be in inner loop.
+    if (UseBB != OrigBB)
+      return nullptr;
+
+    // Avoid constructing constant aggregate because constant value may expose
+    // more optimizations.
+    bool ConstAgg = true;
+    for (auto Val : AggElts) {
+      Value *Elt = (*Val)->DoPHITranslation(UseBB, It.first);
+      if (!isa<Constant>(Elt)) {
+        ConstAgg = false;
+        break;
+      }
+    }
+    if (ConstAgg)
+      return nullptr;
+  }
+
+  // For predecessors without appropriate source aggregate, create one in the
+  // predecessor.
+  for (auto &It : SourceAggregates) {
+    if (Describe(It.second) == AggregateDescription::Found)
+      continue;
+
+    BasicBlock *Pred = It.first;
+    Builder.SetInsertPoint(Pred->getTerminator());
+    Value *V = PoisonValue::get(AggTy);
+    for (auto [Idx, Val] : enumerate(AggElts)) {
+      Value *Elt = (*Val)->DoPHITranslation(UseBB, Pred);
+      V = Builder.CreateInsertValue(V, Elt, Idx);
+    }
+
+    It.second = V;
   }
 
   // All good! Now we just need to thread the source aggregates here.