[InstCombine] PHI-aware aggregate reconstruction: properly handle dup…

…licate predecessors

While it may seem like we can just "deduplicate" the case where
some basic block happens to be a predecessor more than once,
which happens for e.g. switches, that is not correct thing to do.
We must actually add a PHI operand for each predecessor.

This was initially reported to me by David Major
as a clang crash during gecko build for android.

llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp

@@ -931,25 +931,32 @@ Instruction *InstCombinerImpl::foldAggregateConstructionIntoAggregateReuse(
 
   // Arbitrary predecessor count limit.
   static const int PredCountLimit = 64;
-  // Don't bother if there are too many predecessors.
-  if (UseBB->hasNPredecessorsOrMore(PredCountLimit + 1))
-    return nullptr;
+
+  // Cache the (non-uniqified!) list of predecessors in a vector,
+  // checking the limit at the same time for efficiency.
+  SmallVector<BasicBlock *, 4> Preds; // May have duplicates!
+  for (BasicBlock *Pred : predecessors(UseBB)) {
+    // Don't bother if there are too many predecessors.
+    if (Preds.size() >= PredCountLimit) // FIXME: only count duplicates once?
+      return nullptr;
+    Preds.emplace_back(Pred);
+  }
 
   // For each predecessor, what is the source aggregate,
   // from which all the elements were originally extracted from?
   // Note that we want for the map to have stable iteration order!
   SmallMapVector<BasicBlock *, Value *, 4> SourceAggregates;
-  for (BasicBlock *PredBB : predecessors(UseBB)) {
+  for (BasicBlock *Pred : Preds) {
     std::pair<decltype(SourceAggregates)::iterator, bool> IV =
-        SourceAggregates.insert({PredBB, nullptr});
+        SourceAggregates.insert({Pred, nullptr});
     // Did we already evaluate this predecessor?
     if (!IV.second)
       continue;
 
     // Let's hope that when coming from predecessor Pred, all elements of the
     // 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, PredBB);
+    SourceAggregate = FindCommonSourceAggregate(UseBB, Pred);
     if (Describe(SourceAggregate) != AggregateDescription::Found)
       return nullptr; // Give up.
     IV.first->second = *SourceAggregate;
@@ -958,13 +965,14 @@ Instruction *InstCombinerImpl::foldAggregateConstructionIntoAggregateReuse(
   // All good! Now we just need to thread the source aggregates here.
   // Note that we have to insert the new PHI here, ourselves, because we can't
   // rely on InstCombinerImpl::run() inserting it into the right basic block.
+  // Note that the same block can be a predecessor more than once,
+  // and we need to preserve that invariant for the PHI node.
   BuilderTy::InsertPointGuard Guard(Builder);
   Builder.SetInsertPoint(UseBB->getFirstNonPHI());
-  auto *PHI = Builder.CreatePHI(AggTy, SourceAggregates.size(),
-                                OrigIVI.getName() + ".merged");
-  for (const std::pair<BasicBlock *, Value *> &SourceAggregate :
-       SourceAggregates)
-    PHI->addIncoming(SourceAggregate.second, SourceAggregate.first);
+  auto *PHI =
+      Builder.CreatePHI(AggTy, Preds.size(), OrigIVI.getName() + ".merged");
+  for (BasicBlock *Pred : Preds)
+    PHI->addIncoming(SourceAggregates[Pred], Pred);
 
   ++NumAggregateReconstructionsSimplified;
   OrigIVI.replaceAllUsesWith(PHI);

llvm/test/Transforms/InstCombine/phi-aware-aggregate-reconstruction.ll

@@ -420,3 +420,60 @@ end:
   %i8 = insertvalue { i32, i32 } %i7, i32 %i3, 1
   ret { i32, i32 } %i8
 }
+
+; Most basic test - diamond structure, but with a switch, which results in multiple duplicate predecessors
+define { i32, i32 } @test8({ i32, i32 } %agg_left, { i32, i32 } %agg_right, i1 %c, i32 %val_left, i32 %val_right) {
+; CHECK-LABEL: @test8(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br i1 [[C:%.*]], label [[LEFT:%.*]], label [[RIGHT:%.*]]
+; CHECK:       left:
+; CHECK-NEXT:    call void @foo()
+; CHECK-NEXT:    switch i32 [[VAL_LEFT:%.*]], label [[IMPOSSIBLE:%.*]] [
+; CHECK-NEXT:    i32 -42, label [[END:%.*]]
+; CHECK-NEXT:    i32 42, label [[END]]
+; CHECK-NEXT:    ]
+; CHECK:       right:
+; CHECK-NEXT:    call void @bar()
+; CHECK-NEXT:    switch i32 [[VAL_RIGHT:%.*]], label [[IMPOSSIBLE]] [
+; CHECK-NEXT:    i32 42, label [[END]]
+; CHECK-NEXT:    i32 -42, label [[END]]
+; CHECK-NEXT:    ]
+; CHECK:       impossible:
+; CHECK-NEXT:    unreachable
+; CHECK:       end:
+; CHECK-NEXT:    [[I8_MERGED:%.*]] = phi { i32, i32 } [ [[AGG_RIGHT:%.*]], [[RIGHT]] ], [ [[AGG_RIGHT]], [[RIGHT]] ], [ [[AGG_LEFT:%.*]], [[LEFT]] ], [ [[AGG_LEFT]], [[LEFT]] ]
+; CHECK-NEXT:    call void @baz()
+; CHECK-NEXT:    ret { i32, i32 } [[I8_MERGED]]
+;
+entry:
+  br i1 %c, label %left, label %right
+
+left:
+  %i0 = extractvalue { i32, i32 } %agg_left, 0
+  %i2 = extractvalue { i32, i32 } %agg_left, 1
+  call void @foo()
+  switch i32 %val_left, label %impossible [
+  i32 -42, label %end
+  i32 42, label %end
+  ]
+
+right:
+  %i3 = extractvalue { i32, i32 } %agg_right, 0
+  %i4 = extractvalue { i32, i32 } %agg_right, 1
+  call void @bar()
+  switch i32 %val_right, label %impossible [
+  i32 42, label %end
+  i32 -42, label %end
+  ]
+
+impossible:
+  unreachable
+
+end:
+  %i5 = phi i32 [ %i0, %left ], [ %i0, %left ], [ %i3, %right ], [ %i3, %right ]
+  %i6 = phi i32 [ %i2, %left ], [ %i2, %left ], [ %i4, %right ], [ %i4, %right ]
+  call void @baz()
+  %i7 = insertvalue { i32, i32 } undef, i32 %i5, 0
+  %i8 = insertvalue { i32, i32 } %i7, i32 %i6, 1
+  ret { i32, i32 } %i8
+}