[llvm] [BOLT] Optimize basic block loops to avoid n^2 loop (PR #156243)

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-bolt

Author: Mark Rousskov (Mark-Simulacrum)

<details>
<summary>Changes</summary>

This improves BOLT runtime when optimizing rustc_driver.so from 15 minutes to 7 minutes (or 49 minutes to 37 minutes of userspace time).

---
Full diff: https://github.com/llvm/llvm-project/pull/156243.diff


1 Files Affected:

- (modified) bolt/lib/Core/BinaryFunction.cpp (+13-1) 


``````````diff
diff --git a/bolt/lib/Core/BinaryFunction.cpp b/bolt/lib/Core/BinaryFunction.cpp
index 6cac2d0cca2cb..a86e204cae974 100644
--- a/bolt/lib/Core/BinaryFunction.cpp
+++ b/bolt/lib/Core/BinaryFunction.cpp
@@ -3591,6 +3591,18 @@ void BinaryFunction::fixBranches() {
   auto &MIB = BC.MIB;
   MCContext *Ctx = BC.Ctx.get();
 
+  // Caches `FunctionLayout::nextBasicBlock(IgnoreSplits = false)`.
+  // nextBasicBlock uses linear search to find the next block, so the loop
+  // below becomes O(n^2). This avoids that.
+  DenseMap<BinaryBasicBlock *, BinaryBasicBlock *> nextBasicBlock(
+      Layout.block_size());
+  for (size_t i = 0; i + 1 < Layout.block_size(); i++) {
+    auto current = Layout.block_begin() + i;
+    auto next = Layout.block_begin() + i + 1;
+    if (next != Layout.getFragment((*current)->getFragmentNum()).end())
+      nextBasicBlock.insert(std::pair(*current, *next));
+  }
+
   for (BinaryBasicBlock *BB : BasicBlocks) {
     const MCSymbol *TBB = nullptr;
     const MCSymbol *FBB = nullptr;
@@ -3605,7 +3617,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

``````````

</details>


https://github.com/llvm/llvm-project/pull/156243