[DFAJumpThreading] Handle select unfolding when user phi is not a dir… (#109511)

…ect successor

Previously the code assumed that the select instruction is defined in a
block that is a direct predecessor of the block where the PHINode uses
it. So, we were hitting an assertion when we tried to access the def
block as an incoming block for the user phi node.

This patch handles that case by using the correct end block and creating
a new phi node that aggregates both the values of the select in that end
block, and then using that new unfolded phi to overwrite the original
user phi node.

Fixes #106083

Change-Id: Ie471994cca232318f74a6e6438efa21e561c2dc0

2 files changed, 152 insertions, 17 deletions

diff --git a/llvm/lib/Transforms/Scalar/DFAJumpThreading.cpp b/llvm/lib/Transforms/Scalar/DFAJumpThreading.cpp
index ef9c264..0e2b5c9 100644
--- a/llvm/lib/Transforms/Scalar/DFAJumpThreading.cpp
+++ b/llvm/lib/Transforms/Scalar/DFAJumpThreading.cpp
@@ -194,7 +194,6 @@ void unfold(DomTreeUpdater *DTU, LoopInfo *LI, SelectInstToUnfold SIToUnfold,
   SelectInst *SI = SIToUnfold.getInst();
   PHINode *SIUse = SIToUnfold.getUse();
   BasicBlock *StartBlock = SI->getParent();
-  BasicBlock *EndBlock = SIUse->getParent();
   BranchInst *StartBlockTerm =
       dyn_cast<BranchInst>(StartBlock->getTerminator());
 
@@ -202,6 +201,7 @@ void unfold(DomTreeUpdater *DTU, LoopInfo *LI, SelectInstToUnfold SIToUnfold,
   assert(SI->hasOneUse());
 
   if (StartBlockTerm->isUnconditional()) {
+    BasicBlock *EndBlock = StartBlock->getUniqueSuccessor();
     // Arbitrarily choose the 'false' side for a new input value to the PHI.
     BasicBlock *NewBlock = BasicBlock::Create(
         SI->getContext(), Twine(SI->getName(), ".si.unfold.false"),
@@ -223,32 +223,44 @@ void unfold(DomTreeUpdater *DTU, LoopInfo *LI, SelectInstToUnfold SIToUnfold,
                                       NewBlock->getFirstInsertionPt());
     NewPhi->addIncoming(SIOp2, StartBlock);
 
-    if (auto *OpSi = dyn_cast<SelectInst>(SIOp1))
-      NewSIsToUnfold->push_back(SelectInstToUnfold(OpSi, SIUse));
-    if (auto *OpSi = dyn_cast<SelectInst>(SIOp2))
-      NewSIsToUnfold->push_back(SelectInstToUnfold(OpSi, NewPhi));
-
-    // Update the phi node of SI.
-    for (unsigned Idx = 0; Idx < SIUse->getNumIncomingValues(); ++Idx) {
-      if (SIUse->getIncomingBlock(Idx) == StartBlock)
-        SIUse->setIncomingValue(Idx, SIOp1);
+    // Update any other PHI nodes in EndBlock.
+    for (PHINode &Phi : EndBlock->phis()) {
+      if (SIUse == &Phi)
+        continue;
+      Phi.addIncoming(Phi.getIncomingValueForBlock(StartBlock), NewBlock);
     }
-    SIUse->addIncoming(NewPhi, NewBlock);
 
-    // Update any other PHI nodes in EndBlock.
-    for (auto II = EndBlock->begin(); PHINode *Phi = dyn_cast<PHINode>(II);
-         ++II) {
-      if (Phi != SIUse)
-        Phi->addIncoming(Phi->getIncomingValueForBlock(StartBlock), NewBlock);
+    // Update the phi node of SI, which is its only use.
+    if (EndBlock == SIUse->getParent()) {
+      SIUse->addIncoming(NewPhi, NewBlock);
+      SIUse->replaceUsesOfWith(SI, SIOp1);
+    } else {
+      PHINode *EndPhi = PHINode::Create(SIUse->getType(), pred_size(EndBlock),
+                                        Twine(SI->getName(), ".si.unfold.phi"),
+                                        EndBlock->getFirstInsertionPt());
+      for (BasicBlock *Pred : predecessors(EndBlock)) {
+        if (Pred != StartBlock && Pred != NewBlock)
+          EndPhi->addIncoming(EndPhi, Pred);
+      }
+
+      EndPhi->addIncoming(SIOp1, StartBlock);
+      EndPhi->addIncoming(NewPhi, NewBlock);
+      SIUse->replaceUsesOfWith(SI, EndPhi);
+      SIUse = EndPhi;
     }
 
-    StartBlockTerm->eraseFromParent();
+    if (auto *OpSi = dyn_cast<SelectInst>(SIOp1))
+      NewSIsToUnfold->push_back(SelectInstToUnfold(OpSi, SIUse));
+    if (auto *OpSi = dyn_cast<SelectInst>(SIOp2))
+      NewSIsToUnfold->push_back(SelectInstToUnfold(OpSi, NewPhi));
 
     // Insert the real conditional branch based on the original condition.
+    StartBlockTerm->eraseFromParent();
     BranchInst::Create(EndBlock, NewBlock, SI->getCondition(), StartBlock);
     DTU->applyUpdates({{DominatorTree::Insert, StartBlock, EndBlock},
                        {DominatorTree::Insert, StartBlock, NewBlock}});
   } else {
+    BasicBlock *EndBlock = SIUse->getParent();
     BasicBlock *NewBlockT = BasicBlock::Create(
         SI->getContext(), Twine(SI->getName(), ".si.unfold.true"),
         EndBlock->getParent(), EndBlock);
diff --git a/llvm/test/Transforms/DFAJumpThreading/dfa-jump-threading-transform.ll b/llvm/test/Transforms/DFAJumpThreading/dfa-jump-threading-transform.ll
index c38f81d..cba1ba8 100644
--- a/llvm/test/Transforms/DFAJumpThreading/dfa-jump-threading-transform.ll
+++ b/llvm/test/Transforms/DFAJumpThreading/dfa-jump-threading-transform.ll
@@ -300,3 +300,126 @@ define void @self-reference() {
 end:
   ret void
 }
+
+define void @pr106083_invalidBBarg_fold(i1 %cmp1, i1 %cmp2, i1 %not, ptr %d) {
+; CHECK-LABEL: @pr106083_invalidBBarg_fold(
+; CHECK-NEXT:  bb:
+; CHECK-NEXT:    br i1 [[CMP1:%.*]], label [[BB1:%.*]], label [[SEL_SI_UNFOLD_FALSE:%.*]]
+; CHECK:       sel.si.unfold.false:
+; CHECK-NEXT:    [[DOTSI_UNFOLD_PHI1:%.*]] = phi i32 [ 1, [[BB:%.*]] ]
+; CHECK-NEXT:    br label [[BB1]]
+; CHECK:       BB1:
+; CHECK-NEXT:    [[I:%.*]] = phi i16 [ 0, [[BB1_BACKEDGE:%.*]] ], [ 0, [[BB]] ], [ 1, [[BB7:%.*]] ], [ 0, [[SEL_SI_UNFOLD_FALSE]] ], [ 1, [[BB7_JT0:%.*]] ]
+; CHECK-NEXT:    [[SEL_SI_UNFOLD_PHI:%.*]] = phi i32 [ [[SEL_SI_UNFOLD_PHI]], [[BB1_BACKEDGE]] ], [ [[SEL_SI_UNFOLD_PHI]], [[BB7]] ], [ 0, [[BB]] ], [ [[DOTSI_UNFOLD_PHI1]], [[SEL_SI_UNFOLD_FALSE]] ], [ [[SEL_SI_UNFOLD_PHI]], [[BB7_JT0]] ]
+; CHECK-NEXT:    br i1 [[NOT:%.*]], label [[BB7_JT0]], label [[BB2:%.*]]
+; CHECK:       BB2:
+; CHECK-NEXT:    store i16 0, ptr [[D:%.*]], align 2
+; CHECK-NEXT:    br i1 [[CMP2:%.*]], label [[BB7]], label [[SPEC_SELECT_SI_UNFOLD_FALSE_JT0:%.*]]
+; CHECK:       spec.select.si.unfold.false:
+; CHECK-NEXT:    br label [[BB7]]
+; CHECK:       spec.select.si.unfold.false.jt0:
+; CHECK-NEXT:    [[DOTSI_UNFOLD_PHI_JT0:%.*]] = phi i32 [ 0, [[BB2]] ]
+; CHECK-NEXT:    br label [[BB7_JT0]]
+; CHECK:       BB7:
+; CHECK-NEXT:    [[D_PROMOTED4:%.*]] = phi i16 [ 1, [[BB2]] ], [ 1, [[SPEC_SELECT_SI_UNFOLD_FALSE:%.*]] ]
+; CHECK-NEXT:    [[_3:%.*]] = phi i32 [ [[SEL_SI_UNFOLD_PHI]], [[BB2]] ], [ poison, [[SPEC_SELECT_SI_UNFOLD_FALSE]] ]
+; CHECK-NEXT:    switch i32 [[_3]], label [[BB1_BACKEDGE]] [
+; CHECK-NEXT:      i32 0, label [[BB1]]
+; CHECK-NEXT:      i32 1, label [[BB8:%.*]]
+; CHECK-NEXT:    ]
+; CHECK:       BB7.jt0:
+; CHECK-NEXT:    [[D_PROMOTED4_JT0:%.*]] = phi i16 [ 0, [[BB1]] ], [ 1, [[SPEC_SELECT_SI_UNFOLD_FALSE_JT0]] ]
+; CHECK-NEXT:    [[_3_JT0:%.*]] = phi i32 [ 0, [[BB1]] ], [ [[DOTSI_UNFOLD_PHI_JT0]], [[SPEC_SELECT_SI_UNFOLD_FALSE_JT0]] ]
+; CHECK-NEXT:    br label [[BB1]]
+; CHECK:       BB1.backedge:
+; CHECK-NEXT:    br label [[BB1]]
+; CHECK:       BB8:
+; CHECK-NEXT:    ret void
+;
+bb:
+  %sel = select i1 %cmp1, i32 0, i32 1
+  br label %BB1
+
+BB1:                                              ; preds = %BB1.backedge, %BB7, %bb
+  %i = phi i16 [ 0, %BB1.backedge ], [ 0, %bb ], [ 1, %BB7 ]
+  br i1 %not, label %BB7, label %BB2
+
+BB2:                                              ; preds = %BB1
+  store i16 0, ptr %d, align 2
+  %spec.select = select i1 %cmp2, i32 %sel, i32 0
+  br label %BB7
+
+BB7:                                              ; preds = %BB2, %BB1
+  %d.promoted4 = phi i16 [ 0, %BB1 ], [ 1, %BB2 ]
+  %_3 = phi i32 [ 0, %BB1 ], [ %spec.select, %BB2 ]
+  switch i32 %_3, label %BB1.backedge [
+  i32 0, label %BB1
+  i32 1, label %BB8
+  ]
+
+BB1.backedge:                                     ; preds = %BB7
+  br label %BB1
+
+BB8:                                              ; preds = %BB7
+  ret void
+}
+
+define void @pr106083_select_dead_uses(i1 %cmp1, i1 %not, ptr %p) {
+; CHECK-LABEL: @pr106083_select_dead_uses(
+; CHECK-NEXT:  bb:
+; CHECK-NEXT:    br i1 [[CMP1:%.*]], label [[DOTLOOPEXIT6:%.*]], label [[SPEC_SELECT_SI_UNFOLD_FALSE:%.*]]
+; CHECK:       spec.select.si.unfold.false:
+; CHECK-NEXT:    [[DOTSI_UNFOLD_PHI1:%.*]] = phi i32 [ 1, [[BB:%.*]] ]
+; CHECK-NEXT:    br label [[DOTLOOPEXIT6]]
+; CHECK:       .loopexit6:
+; CHECK-NEXT:    [[SPEC_SELECT_SI_UNFOLD_PHI:%.*]] = phi i32 [ [[SPEC_SELECT_SI_UNFOLD_PHI]], [[SELECT_UNFOLD:%.*]] ], [ 0, [[BB]] ], [ [[DOTSI_UNFOLD_PHI1]], [[SPEC_SELECT_SI_UNFOLD_FALSE]] ]
+; CHECK-NEXT:    br i1 [[NOT:%.*]], label [[SELECT_UNFOLD_JT0:%.*]], label [[BB1:%.*]]
+; CHECK:       bb1:
+; CHECK-NEXT:    [[I:%.*]] = load i32, ptr [[P:%.*]], align 4
+; CHECK-NEXT:    [[NOT2:%.*]] = icmp eq i32 0, 0
+; CHECK-NEXT:    br i1 [[NOT2]], label [[SELECT_UNFOLD]], label [[SPEC_SELECT7_SI_UNFOLD_FALSE_JT0:%.*]]
+; CHECK:       spec.select7.si.unfold.false:
+; CHECK-NEXT:    br label [[SELECT_UNFOLD]]
+; CHECK:       spec.select7.si.unfold.false.jt0:
+; CHECK-NEXT:    [[DOTSI_UNFOLD_PHI_JT0:%.*]] = phi i32 [ 0, [[BB1]] ]
+; CHECK-NEXT:    br label [[SELECT_UNFOLD_JT0]]
+; CHECK:       select.unfold:
+; CHECK-NEXT:    [[_2:%.*]] = phi i32 [ [[SPEC_SELECT_SI_UNFOLD_PHI]], [[BB1]] ], [ poison, [[SPEC_SELECT7_SI_UNFOLD_FALSE:%.*]] ]
+; CHECK-NEXT:    switch i32 [[_2]], label [[BB2:%.*]] [
+; CHECK-NEXT:      i32 0, label [[DOTPREHEADER_PREHEADER:%.*]]
+; CHECK-NEXT:      i32 1, label [[DOTLOOPEXIT6]]
+; CHECK-NEXT:    ]
+; CHECK:       select.unfold.jt0:
+; CHECK-NEXT:    [[_2_JT0:%.*]] = phi i32 [ 0, [[DOTLOOPEXIT6]] ], [ [[DOTSI_UNFOLD_PHI_JT0]], [[SPEC_SELECT7_SI_UNFOLD_FALSE_JT0]] ]
+; CHECK-NEXT:    br label [[DOTPREHEADER_PREHEADER]]
+; CHECK:       .preheader.preheader:
+; CHECK-NEXT:    ret void
+; CHECK:       bb2:
+; CHECK-NEXT:    unreachable
+;
+bb:
+  %spec.select = select i1 %cmp1, i32 0, i32 1
+  br label %.loopexit6
+
+.loopexit6:                                       ; preds = %select.unfold, %bb
+  br i1 %not, label %select.unfold, label %bb1
+
+bb1:                                              ; preds = %.loopexit6
+  %i = load i32, ptr %p, align 4
+  %not2 = icmp eq i32 0, 0
+  %spec.select7 = select i1 %not2, i32 %spec.select, i32 0
+  br label %select.unfold
+
+select.unfold:                                    ; preds = %bb1, %.loopexit6
+  %_2 = phi i32 [ 0, %.loopexit6 ], [ %spec.select7, %bb1 ]
+  switch i32 %_2, label %bb2 [
+  i32 0, label %.preheader.preheader
+  i32 1, label %.loopexit6
+  ]
+
+.preheader.preheader:                             ; preds = %select.unfold
+  ret void
+
+bb2:                                              ; preds = %select.unfold
+  unreachable
+}