[CodeGen] IndirectBrExpandPass: preserve Dominator Tree, if available

This fully de-pessimizes the common case of no indirectbr's,
(where we don't actually need to do anything to preserve domtree)
and avoids domtree recomputation in the case there were indirectbr's.

Note that two indirectbr's could have a common successor, and not all
successors of an indirectbr's are meant to survive the expansion.

Though, the code assumes that an indirectbr's doesn't have
duplicate successors, those *should* have been deduplicated
by simplifycfg or something already.

llvm/lib/CodeGen/IndirectBrExpandPass.cpp

@@ -28,9 +28,11 @@
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/Sequence.h"
 #include "llvm/ADT/SmallVector.h"
+#include "llvm/Analysis/DomTreeUpdater.h"
 #include "llvm/CodeGen/TargetPassConfig.h"
 #include "llvm/CodeGen/TargetSubtargetInfo.h"
 #include "llvm/IR/BasicBlock.h"
+#include "llvm/IR/Dominators.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/InstIterator.h"
@@ -59,15 +61,22 @@ class IndirectBrExpandPass : public FunctionPass {
     initializeIndirectBrExpandPassPass(*PassRegistry::getPassRegistry());
   }
 
+  void getAnalysisUsage(AnalysisUsage &AU) const override {
+    AU.addPreserved<DominatorTreeWrapperPass>();
+  }
+
   bool runOnFunction(Function &F) override;
 };
 
 } // end anonymous namespace
 
 char IndirectBrExpandPass::ID = 0;
 
-INITIALIZE_PASS(IndirectBrExpandPass, DEBUG_TYPE,
-                "Expand indirectbr instructions", false, false)
+INITIALIZE_PASS_BEGIN(IndirectBrExpandPass, DEBUG_TYPE,
+                      "Expand indirectbr instructions", false, false)
+INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
+INITIALIZE_PASS_END(IndirectBrExpandPass, DEBUG_TYPE,
+                    "Expand indirectbr instructions", false, false)
 
 FunctionPass *llvm::createIndirectBrExpandPass() {
   return new IndirectBrExpandPass();
@@ -85,6 +94,10 @@ bool IndirectBrExpandPass::runOnFunction(Function &F) {
     return false;
   TLI = STI.getTargetLowering();
 
+  Optional<DomTreeUpdater> DTU;
+  if (auto *DTWP = getAnalysisIfAvailable<DominatorTreeWrapperPass>())
+    DTU.emplace(DTWP->getDomTree(), DomTreeUpdater::UpdateStrategy::Lazy);
+
   SmallVector<IndirectBrInst *, 1> IndirectBrs;
 
   // Set of all potential successors for indirectbr instructions.
@@ -158,10 +171,22 @@ bool IndirectBrExpandPass::runOnFunction(Function &F) {
   if (BBs.empty()) {
     // There are no blocks whose address is taken, so any indirectbr instruction
     // cannot get a valid input and we can replace all of them with unreachable.
+    SmallVector<DominatorTree::UpdateType, 8> Updates;
+    if (DTU)
+      Updates.reserve(IndirectBrSuccs.size());
     for (auto *IBr : IndirectBrs) {
+      if (DTU) {
+        for (BasicBlock *SuccBB : IBr->successors())
+          Updates.push_back({DominatorTree::Delete, IBr->getParent(), SuccBB});
+      }
       (void)new UnreachableInst(F.getContext(), IBr);
       IBr->eraseFromParent();
     }
+    if (DTU) {
+      assert(Updates.size() == IndirectBrSuccs.size() &&
+             "Got unexpected update count.");
+      DTU->applyUpdates(Updates);
+    }
     return true;
   }
 
@@ -183,12 +208,22 @@ bool IndirectBrExpandPass::runOnFunction(Function &F) {
         Twine(IBr->getAddress()->getName()) + ".switch_cast", IBr);
   };
 
+  SmallVector<DominatorTree::UpdateType, 8> Updates;
+
   if (IndirectBrs.size() == 1) {
     // If we only have one indirectbr, we can just directly replace it within
     // its block.
-    SwitchBB = IndirectBrs[0]->getParent();
-    SwitchValue = GetSwitchValue(IndirectBrs[0]);
-    IndirectBrs[0]->eraseFromParent();
+    IndirectBrInst *IBr = IndirectBrs[0];
+    SwitchBB = IBr->getParent();
+    SwitchValue = GetSwitchValue(IBr);
+    if (DTU) {
+      Updates.reserve(IndirectBrSuccs.size());
+      for (BasicBlock *SuccBB : IBr->successors())
+        Updates.push_back({DominatorTree::Delete, IBr->getParent(), SuccBB});
+      assert(Updates.size() == IndirectBrSuccs.size() &&
+             "Got unexpected update count.");
+    }
+    IBr->eraseFromParent();
   } else {
     // Otherwise we need to create a new block to hold the switch across BBs,
     // jump to that block instead of each indirectbr, and phi together the
@@ -200,9 +235,16 @@ bool IndirectBrExpandPass::runOnFunction(Function &F) {
 
     // Now replace the indirectbr instructions with direct branches to the
     // switch block and fill out the PHI operands.
+    if (DTU)
+      Updates.reserve(IndirectBrs.size() + 2 * IndirectBrSuccs.size());
     for (auto *IBr : IndirectBrs) {
       SwitchPN->addIncoming(GetSwitchValue(IBr), IBr->getParent());
       BranchInst::Create(SwitchBB, IBr);
+      if (DTU) {
+        Updates.push_back({DominatorTree::Insert, IBr->getParent(), SwitchBB});
+        for (BasicBlock *SuccBB : IBr->successors())
+          Updates.push_back({DominatorTree::Delete, IBr->getParent(), SuccBB});
+      }
       IBr->eraseFromParent();
     }
   }
@@ -215,5 +257,15 @@ bool IndirectBrExpandPass::runOnFunction(Function &F) {
   for (int i : llvm::seq<int>(1, BBs.size()))
     SI->addCase(ConstantInt::get(CommonITy, i + 1), BBs[i]);
 
+  if (DTU) {
+    // If there were multiple indirectbr's, they may have common successors,
+    // but in the dominator tree, we only track unique edges.
+    SmallPtrSet<BasicBlock *, 8> UniqueSuccessors(BBs.begin(), BBs.end());
+    Updates.reserve(Updates.size() + UniqueSuccessors.size());
+    for (BasicBlock *BB : UniqueSuccessors)
+      Updates.push_back({DominatorTree::Insert, SwitchBB, BB});
+    DTU->applyUpdates(Updates);
+  }
+
   return true;
 }

llvm/test/CodeGen/X86/opt-pipeline.ll

@@ -60,7 +60,6 @@
 ; CHECK-NEXT:       Interleaved Access Pass
 ; CHECK-NEXT:       X86 Partial Reduction
 ; CHECK-NEXT:       Expand indirectbr instructions
-; CHECK-NEXT:       Dominator Tree Construction
 ; CHECK-NEXT:       Natural Loop Information
 ; CHECK-NEXT:       CodeGen Prepare
 ; CHECK-NEXT:     Rewrite Symbols