From 7e88942d2548704b35ac92bd327b38f9c6a991d2 Mon Sep 17 00:00:00 2001 From: Roman Lebedev Date: Thu, 28 Jan 2021 01:24:19 +0300 Subject: [PATCH] [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 | 62 +++++++++++++++++++++-- llvm/test/CodeGen/X86/opt-pipeline.ll | 1 - 2 files changed, 57 insertions(+), 6 deletions(-) diff --git a/llvm/lib/CodeGen/IndirectBrExpandPass.cpp b/llvm/lib/CodeGen/IndirectBrExpandPass.cpp index 4473a139d3ad47..e4606daba35290 100644 --- a/llvm/lib/CodeGen/IndirectBrExpandPass.cpp +++ b/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,6 +61,10 @@ class IndirectBrExpandPass : public FunctionPass { initializeIndirectBrExpandPassPass(*PassRegistry::getPassRegistry()); } + void getAnalysisUsage(AnalysisUsage &AU) const override { + AU.addPreserved(); + } + bool runOnFunction(Function &F) override; }; @@ -66,8 +72,11 @@ class IndirectBrExpandPass : public FunctionPass { 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/llvm/test/CodeGen/X86/opt-pipeline.ll b/llvm/test/CodeGen/X86/opt-pipeline.ll index b851eea60b0a89..9abf60777763d3 100644 --- a/llvm/test/CodeGen/X86/opt-pipeline.ll +++ b/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