diff --git a/lib/CodeGen/IndirectBrExpandPass.cpp b/lib/CodeGen/IndirectBrExpandPass.cpp index 4473a139d3a..e4606daba35 100644 --- a/lib/CodeGen/IndirectBrExpandPass.cpp +++ b/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,6 +61,10 @@ public: initializeIndirectBrExpandPassPass(*PassRegistry::getPassRegistry()); } + void getAnalysisUsage(AnalysisUsage &AU) const override { + AU.addPreserved(); + } + bool runOnFunction(Function &F) override; }; @@ -66,8 +72,11 @@ public: 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 DTU; + if (auto *DTWP = getAnalysisIfAvailable()) + DTU.emplace(DTWP->getDomTree(), DomTreeUpdater::UpdateStrategy::Lazy); + SmallVector 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 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 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(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 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; } diff --git a/test/CodeGen/X86/opt-pipeline.ll b/test/CodeGen/X86/opt-pipeline.ll index b851eea60b0..9abf6077776 100644 --- a/test/CodeGen/X86/opt-pipeline.ll +++ b/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