[SimplifyCFG] Teach FoldBranchToCommonDest() to preserve DomTree, part 1

... for conditional branch case

llvm/include/llvm/Transforms/Utils/Local.h

@@ -201,7 +201,8 @@ bool FlattenCFG(BasicBlock *BB, AAResults *AA = nullptr);
 /// If this basic block is ONLY a setcc and a branch, and if a predecessor
 /// branches to us and one of our successors, fold the setcc into the
 /// predecessor and use logical operations to pick the right destination.
-bool FoldBranchToCommonDest(BranchInst *BI, MemorySSAUpdater *MSSAU = nullptr,
+bool FoldBranchToCommonDest(BranchInst *BI, llvm::DomTreeUpdater *DTU = nullptr,
+                            MemorySSAUpdater *MSSAU = nullptr,
                             const TargetTransformInfo *TTI = nullptr,
                             unsigned BonusInstThreshold = 1);
 

llvm/lib/Transforms/Utils/LoopSimplify.cpp

@@ -681,7 +681,7 @@ static bool simplifyOneLoop(Loop *L, SmallVectorImpl<Loop *> &Worklist,
       // The block has now been cleared of all instructions except for
       // a comparison and a conditional branch. SimplifyCFG may be able
       // to fold it now.
-      if (!FoldBranchToCommonDest(BI, MSSAU))
+      if (!FoldBranchToCommonDest(BI, /*DTU=*/nullptr, MSSAU))
         continue;
 
       // Success. The block is now dead, so remove it from the loop,

llvm/lib/Transforms/Utils/SimplifyCFG.cpp

@@ -2743,7 +2743,8 @@ static bool extractPredSuccWeights(BranchInst *PBI, BranchInst *BI,
 /// If this basic block is simple enough, and if a predecessor branches to us
 /// and one of our successors, fold the block into the predecessor and use
 /// logical operations to pick the right destination.
-bool llvm::FoldBranchToCommonDest(BranchInst *BI, MemorySSAUpdater *MSSAU,
+bool llvm::FoldBranchToCommonDest(BranchInst *BI, DomTreeUpdater *DTU,
+                                  MemorySSAUpdater *MSSAU,
                                   const TargetTransformInfo *TTI,
                                   unsigned BonusInstThreshold) {
   BasicBlock *BB = BI->getParent();
@@ -2907,6 +2908,8 @@ bool llvm::FoldBranchToCommonDest(BranchInst *BI, MemorySSAUpdater *MSSAU,
     LLVM_DEBUG(dbgs() << "FOLDING BRANCH TO COMMON DEST:\n" << *PBI << *BB);
     Changed = true;
 
+    SmallVector<DominatorTree::UpdateType, 3> Updates;
+
     IRBuilder<> Builder(PBI);
     // The builder is used to create instructions to eliminate the branch in BB.
     // If BB's terminator has !annotation metadata, add it to the new
@@ -3055,6 +3058,9 @@ bool llvm::FoldBranchToCommonDest(BranchInst *BI, MemorySSAUpdater *MSSAU,
         setBranchWeights(PBI, MDWeights[0], MDWeights[1]);
       } else
         PBI->setMetadata(LLVMContext::MD_prof, nullptr);
+
+      Updates.push_back({DominatorTree::Delete, PredBlock, BB});
+      Updates.push_back({DominatorTree::Insert, PredBlock, UniqueSucc});
     } else {
       // Update PHI nodes in the common successors.
       for (unsigned i = 0, e = PHIs.size(); i != e; ++i) {
@@ -3102,6 +3108,9 @@ bool llvm::FoldBranchToCommonDest(BranchInst *BI, MemorySSAUpdater *MSSAU,
       PBI = New_PBI;
     }
 
+    if (DTU)
+      DTU->applyUpdatesPermissive(Updates);
+
     // If BI was a loop latch, it may have had associated loop metadata.
     // We need to copy it to the new latch, that is, PBI.
     if (MDNode *LoopMD = BI->getMetadata(LLVMContext::MD_loop))
@@ -6194,7 +6203,8 @@ bool SimplifyCFGOpt::simplifyUncondBranch(BranchInst *BI,
   // branches to us and our successor, fold the comparison into the
   // predecessor and use logical operations to update the incoming value
   // for PHI nodes in common successor.
-  if (FoldBranchToCommonDest(BI, nullptr, &TTI, Options.BonusInstThreshold))
+  if (FoldBranchToCommonDest(BI, DTU, /*MSSAU=*/nullptr, &TTI,
+                             Options.BonusInstThreshold))
     return requestResimplify();
   return false;
 }
@@ -6257,7 +6267,8 @@ bool SimplifyCFGOpt::simplifyCondBranch(BranchInst *BI, IRBuilder<> &Builder) {
   // If this basic block is ONLY a compare and a branch, and if a predecessor
   // branches to us and one of our successors, fold the comparison into the
   // predecessor and use logical operations to pick the right destination.
-  if (FoldBranchToCommonDest(BI, nullptr, &TTI, Options.BonusInstThreshold))
+  if (FoldBranchToCommonDest(BI, DTU, /*MSSAU=*/nullptr, &TTI,
+                             Options.BonusInstThreshold))
     return requestResimplify();
 
   // We have a conditional branch to two blocks that are only reachable

llvm/test/Transforms/LoopVectorize/if-pred-non-void.ll

@@ -1,4 +1,4 @@
-; RUN: opt -S -force-vector-width=2 -force-vector-interleave=1 -loop-vectorize -verify-loop-info -simplifycfg < %s | FileCheck %s
+; RUN: opt -S -force-vector-width=2 -force-vector-interleave=1 -loop-vectorize -verify-loop-info -simplifycfg -simplifycfg-require-and-preserve-domtree=1 < %s | FileCheck %s
 ; RUN: opt -S -force-vector-width=1 -force-vector-interleave=2 -loop-vectorize -verify-loop-info < %s | FileCheck %s --check-prefix=UNROLL-NO-VF
 
 target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"

llvm/test/Transforms/SimplifyCFG/2006-12-08-Ptr-ICmp-Branch.ll

@@ -1,4 +1,4 @@
-; RUN: opt < %s -simplifycfg | llvm-dis
+; RUN: opt < %s -simplifycfg -simplifycfg-require-and-preserve-domtree=1 | llvm-dis
 ; END.
 
 ; ModuleID = '2006-12-08-Ptr-ICmp-Branch.ll'

llvm/test/Transforms/SimplifyCFG/ARM/branch-fold-threshold.ll

@@ -1,7 +1,7 @@
 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
-; RUN: opt -mtriple=thumbv8m.main %s -simplifycfg -S | FileCheck %s --check-prefix=CHECK --check-prefix=THUMB
-; RUN: opt -mtriple=thumbv8a %s -simplifycfg -S | FileCheck %s --check-prefix=CHECK --check-prefix=THUMB
-; RUN: opt -mtriple=armv8a %s -simplifycfg -S | FileCheck %s --check-prefix=CHECK --check-prefix=ARM
+; RUN: opt -mtriple=thumbv8m.main %s -simplifycfg -simplifycfg-require-and-preserve-domtree=1 -S | FileCheck %s --check-prefix=CHECK --check-prefix=THUMB
+; RUN: opt -mtriple=thumbv8a %s -simplifycfg -simplifycfg-require-and-preserve-domtree=1 -S | FileCheck %s --check-prefix=CHECK --check-prefix=THUMB
+; RUN: opt -mtriple=armv8a %s -simplifycfg -simplifycfg-require-and-preserve-domtree=1 -S | FileCheck %s --check-prefix=CHECK --check-prefix=ARM
 
 define i32 @foo(i32 %a, i32 %b, i32 %c, i32 %d, i32* %input) {
 ; CHECK-LABEL: @foo(

llvm/test/Transforms/SimplifyCFG/annotations.ll

@@ -1,4 +1,4 @@
-; RUN: opt -simplifycfg -S %s | FileCheck --match-full-lines %s
+; RUN: opt -simplifycfg -simplifycfg-require-and-preserve-domtree=1 -S %s | FileCheck --match-full-lines %s
 
 ; The branch in %cont has !annotation metadata. Make sure generated AND
 ; has !annotation metadata.

llvm/test/Transforms/SimplifyCFG/basictest.ll

@@ -1,6 +1,6 @@
 ; Test CFG simplify removal of branch instructions.
 ;
-; RUN: opt < %s -simplifycfg -S | FileCheck %s
+; RUN: opt < %s -simplifycfg -simplifycfg-require-and-preserve-domtree=1 -S | FileCheck %s
 ; RUN: opt < %s -passes=simplify-cfg -S | FileCheck %s
 
 define void @test1() {

llvm/test/Transforms/SimplifyCFG/branch-fold-dbg.ll

@@ -1,4 +1,4 @@
-; RUN: opt -simplifycfg -S < %s | FileCheck %s
+; RUN: opt -simplifycfg -simplifycfg-require-and-preserve-domtree=1 -S < %s | FileCheck %s
 
 %0 = type { i32*, i32* }
 

llvm/test/Transforms/SimplifyCFG/branch-fold-threshold.ll

@@ -1,6 +1,6 @@
-; RUN: opt %s -simplifycfg -S | FileCheck %s --check-prefix=NORMAL
-; RUN: opt %s -simplifycfg -S -bonus-inst-threshold=2 | FileCheck %s --check-prefix=AGGRESSIVE
-; RUN: opt %s -simplifycfg -S -bonus-inst-threshold=4 | FileCheck %s --check-prefix=WAYAGGRESSIVE
+; RUN: opt %s -simplifycfg -simplifycfg-require-and-preserve-domtree=1 -S | FileCheck %s --check-prefix=NORMAL
+; RUN: opt %s -simplifycfg -simplifycfg-require-and-preserve-domtree=1 -S -bonus-inst-threshold=2 | FileCheck %s --check-prefix=AGGRESSIVE
+; RUN: opt %s -simplifycfg -simplifycfg-require-and-preserve-domtree=1 -S -bonus-inst-threshold=4 | FileCheck %s --check-prefix=WAYAGGRESSIVE
 ; RUN: opt %s -passes=simplify-cfg -S | FileCheck %s --check-prefix=NORMAL
 ; RUN: opt %s -passes='simplify-cfg<bonus-inst-threshold=2>' -S | FileCheck %s --check-prefix=AGGRESSIVE
 ; RUN: opt %s -passes='simplify-cfg<bonus-inst-threshold=4>' -S | FileCheck %s --check-prefix=WAYAGGRESSIVE

llvm/test/Transforms/SimplifyCFG/common-dest-folding.ll

@@ -1,4 +1,4 @@
-; RUN: opt < %s -simplifycfg -S | FileCheck %s
+; RUN: opt < %s -simplifycfg -simplifycfg-require-and-preserve-domtree=1 -S | FileCheck %s
 
 ;CHECK: @foo
 ;CHECK: and i32 %c1, %k

llvm/test/Transforms/SimplifyCFG/fold-branch-to-common-dest.ll

@@ -1,5 +1,5 @@
 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
-; RUN: opt < %s -S -simplifycfg -bonus-inst-threshold=10 | FileCheck %s
+; RUN: opt < %s -S -simplifycfg -simplifycfg-require-and-preserve-domtree=1 -bonus-inst-threshold=10 | FileCheck %s
 
 declare void @sideeffect0()
 declare void @sideeffect1()

llvm/test/Transforms/SimplifyCFG/fold-debug-location.ll

@@ -1,4 +1,4 @@
-; RUN: opt -S -simplifycfg < %s | FileCheck %s --match-full-lines
+; RUN: opt -S -simplifycfg -simplifycfg-require-and-preserve-domtree=1 < %s | FileCheck %s --match-full-lines
 
 ; Make sure we reset the debug location when folding instructions.
 ; CHECK: [[VAL:%.*]] = and i32 %c2, %k