diff --git a/bolt/include/bolt/Core/FunctionLayout.h b/bolt/include/bolt/Core/FunctionLayout.h index ee4dd689b8dd6..f600d8ab0dabf 100644 --- a/bolt/include/bolt/Core/FunctionLayout.h +++ b/bolt/include/bolt/Core/FunctionLayout.h @@ -243,9 +243,18 @@ class FunctionLayout { /// Returns the basic block after the given basic block in the layout or /// nullptr if the last basic block is given. + /// + /// Note that this performs a linear search for BB. const BinaryBasicBlock *getBasicBlockAfter(const BinaryBasicBlock *BB, bool IgnoreSplits = true) const; + /// Returns a mapping from BB -> getBasicBlockAfter(BB). + /// + /// This should be preferred in loops that call getBasicBlockAfter without + /// changes to the function layout. Caching the results avoid n^2 lookup cost. + DenseMap + getBasicBlocksAfter(bool IgnoreSplits = true) const; + /// True if the layout contains at least two non-empty fragments. bool isSplit() const; diff --git a/bolt/lib/Core/BinaryFunction.cpp b/bolt/lib/Core/BinaryFunction.cpp index 6cac2d0cca2cb..35c00a8012426 100644 --- a/bolt/lib/Core/BinaryFunction.cpp +++ b/bolt/lib/Core/BinaryFunction.cpp @@ -3591,6 +3591,8 @@ void BinaryFunction::fixBranches() { auto &MIB = BC.MIB; MCContext *Ctx = BC.Ctx.get(); + auto NextBasicBlock = Layout.getBasicBlocksAfter(/* IgnoreSplits */ false); + for (BinaryBasicBlock *BB : BasicBlocks) { const MCSymbol *TBB = nullptr; const MCSymbol *FBB = nullptr; @@ -3605,7 +3607,7 @@ void BinaryFunction::fixBranches() { // Basic block that follows the current one in the final layout. const BinaryBasicBlock *const NextBB = - Layout.getBasicBlockAfter(BB, /*IgnoreSplits=*/false); + NextBasicBlock.lookup_or(BB, nullptr); if (BB->succ_size() == 1) { // __builtin_unreachable() could create a conditional branch that diff --git a/bolt/lib/Core/FunctionLayout.cpp b/bolt/lib/Core/FunctionLayout.cpp index 4498fc44da954..4f8d75585b4e4 100644 --- a/bolt/lib/Core/FunctionLayout.cpp +++ b/bolt/lib/Core/FunctionLayout.cpp @@ -241,6 +241,22 @@ FunctionLayout::getBasicBlockAfter(const BinaryBasicBlock *BB, return *BlockAfter; } +DenseMap +FunctionLayout::getBasicBlocksAfter(bool IgnoreSplits) const { + DenseMap NextBasicBlock(block_size()); + for (size_t i = 0; i + 1 < block_size(); i++) { + auto Current = block_begin() + i; + auto Next = block_begin() + i + 1; + + if (IgnoreSplits) { + NextBasicBlock.insert(std::pair(*Current, *Next)); + } else if (Next != getFragment((*Current)->getFragmentNum()).end()) { + NextBasicBlock.insert(std::pair(*Current, *Next)); + } + } + return NextBasicBlock; +} + bool FunctionLayout::isSplit() const { const unsigned NonEmptyFragCount = llvm::count_if( fragments(), [](const FunctionFragment &FF) { return !FF.empty(); });