[SimplifyCFG] Enable MergeBlockIntoPredecessor with two successors

Author: HighW4y2H3ll

llvm/lib/Transforms/Utils/BasicBlockUtils.cpp

@@ -207,6 +207,10 @@ bool llvm::MergeBlockIntoPredecessor(BasicBlock *BB, DomTreeUpdater *DTU,
   if (PredecessorWithTwoSuccessors) {
     if (!(PredBB_BI = dyn_cast<BranchInst>(PTI)))
       return false;
+    // Can't merge if any of the instructions has side effects
+    for (Instruction &candid : *BB)
+      if (candid.mayHaveSideEffects())
+        return false;
     BranchInst *BB_JmpI = dyn_cast<BranchInst>(BB->getTerminator());
     if (!BB_JmpI || !BB_JmpI->isUnconditional())
       return false;
@@ -283,6 +287,36 @@ bool llvm::MergeBlockIntoPredecessor(BasicBlock *BB, DomTreeUpdater *DTU,
   if (MSSAU)
     MSSAU->moveAllAfterMergeBlocks(BB, PredBB, Start);
 
+  // PHI nodes that references the Pred shall be replaced with a SelectInst
+  for (BasicBlock *Succ : successors(PTI)) {
+    if (Succ == BB)
+      continue;
+    for (PHINode &PhiSucc : Succ->phis()) {
+      if (llvm::is_contained(PhiSucc.blocks(), PredBB)) {
+        if (PredBB_BI->isConditional()) {
+          Value *Cond = PredBB_BI->getCondition();
+          Value *OrigPhi0 = PhiSucc.removeIncomingValue(PredBB, false);
+          Value *newPhi = OrigPhi0;
+
+          if (PhiSucc.getNumOperands() > 0) {
+            if (llvm::is_contained(PhiSucc.blocks(), BB)) {
+              Value *OrigPhi1 = PhiSucc.removeIncomingValue(BB, false);
+              // Swap branch given the conditions
+              if (PredBB_BI->getSuccessor(0) == Succ) {
+                newPhi = SelectInst::Create(Cond, OrigPhi0, OrigPhi1, "", PTI->getIterator());
+              } else {
+                newPhi = SelectInst::Create(Cond, OrigPhi1, OrigPhi0, "", PTI->getIterator());
+              }
+              PhiSucc.addIncoming(newPhi, PredBB);
+            }
+          }
+
+          PhiSucc.addIncoming(newPhi, PredBB);
+        }
+      }
+    }
+  }
+
   // Make all PHI nodes that referred to BB now refer to Pred as their
   // source...
   BB->replaceAllUsesWith(PredBB);

llvm/lib/Transforms/Utils/SimplifyCFG.cpp

@@ -8339,7 +8339,7 @@ bool SimplifyCFGOpt::simplifyOnce(BasicBlock *BB) {
   // Merge basic blocks into their predecessor if there is only one distinct
   // pred, and if there is only one distinct successor of the predecessor, and
   // if there are no PHI nodes.
-  if (MergeBlockIntoPredecessor(BB, DTU))
+  if (MergeBlockIntoPredecessor(BB, DTU, nullptr, nullptr, nullptr, true))
     return true;
 
   if (SinkCommon && Options.SinkCommonInsts)

llvm/test/Transforms/SimplifyCFG/two-succ-merge-block.ll

@@ -0,0 +1,58 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --tool ./opt --version 5
+; RUN: opt < %s -passes=simplifycfg -S | FileCheck %s
+
+define i1 @_Z7compareRK1SS1_(ptr %a, ptr %b) {
+; CHECK-LABEL: @_Z7compareRK1SS1_(
+; CHECK-SAME: ptr [[A:%.*]], ptr [[B:%.*]]) {
+; CHECK-NEXT: entry:
+; CHECK-NEXT:   %0 = load i32, ptr %a, align 4, !tbaa !3
+; CHECK-NEXT:   %1 = load i32, ptr %b, align 4, !tbaa !3
+; CHECK-NEXT:   %cmp.i = icmp slt i32 %0, %1
+; CHECK-NEXT:   %cmp.i19 = icmp eq i32 %0, %1
+; CHECK-NEXT:   %y = getelementptr inbounds nuw i8, ptr %a, i64 4
+; CHECK-NEXT:   %2 = load i32, ptr %y, align 4, !tbaa !8
+; CHECK-NEXT:   %y14 = getelementptr inbounds nuw i8, ptr %b, i64 4
+; CHECK-NEXT:   %3 = load i32, ptr %y14, align 4, !tbaa !8
+; CHECK-NEXT:   %cmp = icmp slt i32 %2, %3
+; CHECK-NEXT:   %4 = select i1 %cmp.i19, i1 %cmp, i1 false
+; CHECK-NEXT:   %5 = select i1 %cmp.i, i1 true, i1 %4
+; CHECK-NEXT:   br label %lor.end
+; CHECK-EMPTY:
+; CHECK-NEXT: lor.end:                                          ; preds = %entry
+; CHECK-NEXT:   ret i1 %5
+entry:
+  %0 = load i32, ptr %a, align 4, !tbaa !3
+  %1 = load i32, ptr %b, align 4, !tbaa !3
+  %cmp.i = icmp slt i32 %0, %1
+  br i1 %cmp.i, label %lor.end, label %lor.rhs
+
+lor.rhs:                                          ; preds = %entry
+  %cmp.i19 = icmp eq i32 %0, %1
+  br i1 %cmp.i19, label %land.rhs, label %lor.end
+
+land.rhs:                                         ; preds = %lor.rhs
+  %y = getelementptr inbounds nuw i8, ptr %a, i64 4
+  %2 = load i32, ptr %y, align 4, !tbaa !8
+  %y14 = getelementptr inbounds nuw i8, ptr %b, i64 4
+  %3 = load i32, ptr %y14, align 4, !tbaa !8
+  %cmp = icmp slt i32 %2, %3
+  br label %lor.end
+
+lor.end:                                          ; preds = %lor.rhs, %land.rhs, %entry
+  %4 = phi i1 [ true, %entry ], [ false, %lor.rhs ], [ %cmp, %land.rhs ]
+  ret i1 %4
+}
+
+!llvm.module.flags = !{!0, !1}
+!llvm.ident = !{!2}
+
+!0 = !{i32 1, !"wchar_size", i32 4}
+!1 = !{i32 7, !"uwtable", i32 2}
+!2 = !{!"clang"}
+!3 = !{!4, !5, i64 0}
+!4 = !{!"_ZTS1S", !5, i64 0, !5, i64 4}
+!5 = !{!"int", !6, i64 0}
+!6 = !{!"omnipotent char", !7, i64 0}
+!7 = !{!"Simple C++ TBAA"}
+!8 = !{!4, !5, i64 4}
+!9 = !{!5, !5, i64 0}