Handle successor's PHI node correctly when flattening CFG merges two if-regions

[jaebaek]

Summary

FlattenCFG merges two 'if' basicblocks by inserting one basicblock
to another basicblock. The inserted basicblock can have a successor
that contains a PHI node whoes incoming basicblock is the inserted
basicblock. Since the existing code does not handle it, it becomes
a badref.

if (cond1)

statement

if (cond2)

statement

successor - contains PHI node whose predecessor is cond2

-->
if (cond1 || cond2)

statement

(BB for cond2 was deleted)
successor - contains PHI node whose predecessor is cond2 --> bad ref!

[kuhar]

llvm/lib/Transforms/Utils/FlattenCFG.cpp

457

nit: for (unsigned int i = 0; i < PBI->getNumSuccessors(); ++i)
->
for (BasicBlock *Succ : successors(PBI))

459

I think this should iterate over all the phi nodes in the successor basic block?
You can get that with: BB->phis()

460

LLVM prefers quick exists; I think it would be better to use:

if (!SomePhi)
  continue;

461

nit: for (unsigned int j = 0; j < SomePHI->getNumIncomingValues(); ++j)
-->
for (unsigned i = 0, e = SomePHI->getNumIncomingValues(); i != e; ++i)

[jaebaek]

Updated based on reviews. PTAL

[kuhar]

Seems fine to me. Thanks for the fixes.
Please wait some time and see if others are OK with this change too.

llvm/lib/Transforms/Utils/FlattenCFG.cpp

 
 /// If \param [in] BB has more than one predecessor that is a conditional
 /// branch, attempt to use parallel and/or for the branch condition. \returns
 /// true on success.
 ///
 /// Before:
 ///   ......
 ///   %cmp10 = fcmp une float %tmp1, %tmp2
-///   br i1 %cmp1, label %if.then, label %lor.rhs
+///   br i1 %cmp10, label %if.then, label %lor.rhs
 ///
 /// lor.rhs:
 ///   ......
 ///   %cmp11 = fcmp une float %tmp3, %tmp4
 ///   br i1 %cmp11, label %if.then, label %ifend
 ///
 /// if.end:  // the merge block
 ///   ......
bool FlattenCFGOpt::MergeIfRegion(BasicBlock *BB, IRBuilder<> &Builder) {
   Value *CC = PBI->getCondition();
   BasicBlock *SaveInsertBB = Builder.GetInsertBlock();
   BasicBlock::iterator SaveInsertPt = Builder.GetInsertPoint();
   Builder.SetInsertPoint(PBI);
   Value *NC = Builder.CreateOr(CInst1, CC);
   PBI->replaceUsesOfWith(CC, NC);
   Builder.SetInsertPoint(SaveInsertBB, SaveInsertPt);
 
+  // Handle PHI node to replace its predecessors to FirstEntryBlock.
+  for (BasicBlock *Succ : successors(PBI)) {
+    for (PHINode &Phi : Succ->phis()) {
+      for (unsigned i = 0, e = Phi.getNumIncomingValues(); i != e; ++i) {
+        if (Phi.getIncomingBlock(i) == SecondEntryBlock)
+          Phi.setIncomingBlock(i, FirstEntryBlock);
+      }
+    }
+  }
+
   // Remove IfTrue1
   if (IfTrue1 != FirstEntryBlock) {
     IfTrue1->dropAllReferences();
     IfTrue1->eraseFromParent();
   }
 
   // Remove IfFalse1
   if (IfFalse1 != FirstEntryBlock) {

llvm/test/Transforms/Util/flattencfg.ll

 bb1:                                               ; preds = %entry
   br label %bb0
 
 bb0:                                               ; preds = %bb1, %entry
   %1 = fcmp ult float %b, 1.000000e+00
   br i1 %1, label %bb4, label %bb3
 }
 
+; CHECK-LABEL: @test_not_crash3
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    %a_eq_0 = icmp eq i32 %a, 0
+; CHECK-NEXT:    %a_eq_1 = icmp eq i32 %a, 1
+; CHECK-NEXT:    [[COND:%[a-z0-9]+]] = or i1 %a_eq_0, %a_eq_1
+; CHECK-NEXT:    br i1 [[COND]], label %bb2, label %bb3
+; CHECK:       bb2:
+; CHECK-NEXT:    br label %bb3
+; CHECK:       bb3:
+; CHECK-NEXT:    %check_badref = phi i32 [ 17, %entry ], [ 11, %bb2 ]
+; CHECK-NEXT:    ret void
+define void @test_not_crash3(i32 %a) #0 {
+entry:
+  %a_eq_0 = icmp eq i32 %a, 0
+  br i1 %a_eq_0, label %bb0, label %bb1
+
+bb0:                                              ; preds = %entry
+  br label %bb1
+
+bb1:                                              ; preds = %bb0, %entry
+  %a_eq_1 = icmp eq i32 %a, 1
+  br i1 %a_eq_1, label %bb2, label %bb3
+
+bb2:                                              ; preds = %bb1
+  br label %bb3
+
+bb3:                                              ; preds = %bb2, %bb1
+  %check_badref = phi i32 [ 17, %bb1 ], [ 11, %bb2 ]
+  ret void
+}