[BOLT] Avoid n^2 complexity in fixBranches(). NFCI

[maksfb]

Iterator implementation of PR #156243:

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

[llvmbot]

@llvm/pr-subscribers-bolt

Author: Maksim Panchenko (maksfb)

Changes

Iterator implementation of PR #156243:

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/160407.diff

3 Files Affected:

• (modified) bolt/include/bolt/Core/FunctionLayout.h (+16-5)
• (modified) bolt/lib/Core/BinaryFunction.cpp (+4-2)
• (modified) bolt/lib/Core/FunctionLayout.cpp (+13-7)

diff --git a/bolt/include/bolt/Core/FunctionLayout.h b/bolt/include/bolt/Core/FunctionLayout.h
index ee4dd689b8dd6..240d5138a093a 100644
--- a/bolt/include/bolt/Core/FunctionLayout.h
+++ b/bolt/include/bolt/Core/FunctionLayout.h
@@ -232,8 +232,24 @@ class FunctionLayout {
     return Blocks[Index];
   }
 
+  /// Return the basic block after the given basic block iterator in the layout
+  /// or nullptr if the last basic block iterator is given.
+  const BinaryBasicBlock *getBasicBlockAfter(block_const_iterator BlockIt,
+                                             bool IgnoreSplits = true) const;
+
+  /// Returns the basic block after the given basic block in the layout or
+  /// nullptr if the last basic block is given.
+  ///
+  /// Note: prefer the version that takes the iterator as this function uses
+  /// linear basic block lookup.
+  const BinaryBasicBlock *getBasicBlockAfter(const BinaryBasicBlock *BB,
+                                             bool IgnoreSplits = true) const;
+
   /// Returns the basic block after the given basic block in the layout or
   /// nullptr if the last basic block is given.
+  ///
+  /// Note: prefer the version that takes the iterator as this function uses
+  /// linear basic block lookup.
   BinaryBasicBlock *getBasicBlockAfter(const BinaryBasicBlock *const BB,
                                        const bool IgnoreSplits = true) {
     return const_cast<BinaryBasicBlock *>(
@@ -241,11 +257,6 @@ class FunctionLayout {
             BB, IgnoreSplits));
   }
 
-  /// Returns the basic block after the given basic block in the layout or
-  /// nullptr if the last basic block is given.
-  const BinaryBasicBlock *getBasicBlockAfter(const BinaryBasicBlock *BB,
-                                             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 578a87dc6c09d..7102054a366fa 100644
--- a/bolt/lib/Core/BinaryFunction.cpp
+++ b/bolt/lib/Core/BinaryFunction.cpp
@@ -3598,7 +3598,9 @@ void BinaryFunction::fixBranches() {
   auto &MIB = BC.MIB;
   MCContext *Ctx = BC.Ctx.get();
 
-  for (BinaryBasicBlock *BB : BasicBlocks) {
+  for (auto BBI = Layout.block_begin(), BBE = Layout.block_end(); BBI != BBE;
+       ++BBI) {
+    BinaryBasicBlock *BB = *BBI;
     const MCSymbol *TBB = nullptr;
     const MCSymbol *FBB = nullptr;
     MCInst *CondBranch = nullptr;
@@ -3612,7 +3614,7 @@ void BinaryFunction::fixBranches() {
 
     // Basic block that follows the current one in the final layout.
     const BinaryBasicBlock *const NextBB =
-        Layout.getBasicBlockAfter(BB, /*IgnoreSplits=*/false);
+        Layout.getBasicBlockAfter(BBI, /*IgnoreSplits*/ false);
 
     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..98ed6e1320b3e 100644
--- a/bolt/lib/Core/FunctionLayout.cpp
+++ b/bolt/lib/Core/FunctionLayout.cpp
@@ -224,23 +224,29 @@ void FunctionLayout::clear() {
 }
 
 const BinaryBasicBlock *
-FunctionLayout::getBasicBlockAfter(const BinaryBasicBlock *BB,
+FunctionLayout::getBasicBlockAfter(block_const_iterator BBIter,
                                    bool IgnoreSplits) const {
-  const block_const_iterator BBPos = find(blocks(), BB);
-  if (BBPos == block_end())
-    return nullptr;
-
-  const block_const_iterator BlockAfter = std::next(BBPos);
+  const block_const_iterator BlockAfter = std::next(BBIter);
   if (BlockAfter == block_end())
     return nullptr;
 
   if (!IgnoreSplits)
-    if (BlockAfter == getFragment(BB->getFragmentNum()).end())
+    if (BlockAfter == getFragment((*BBIter)->getFragmentNum()).end())
       return nullptr;
 
   return *BlockAfter;
 }
 
+const BinaryBasicBlock *
+FunctionLayout::getBasicBlockAfter(const BinaryBasicBlock *BB,
+                                   bool IgnoreSplits) const {
+  const block_const_iterator BBPos = find(blocks(), BB);
+  if (BBPos == block_end())
+    return nullptr;
+
+  return getBasicBlockAfter(BBPos, IgnoreSplits);
+}
+
 bool FunctionLayout::isSplit() const {
   const unsigned NonEmptyFragCount = llvm::count_if(
       fragments(), [](const FunctionFragment &FF) { return !FF.empty(); });

[maksfb]

Hmm... I had "Co-authored-by: Mark-Simulacrum mark.simulacrum@gmail.com" field in my PR. Not sure why it's not showing up.

Edit: it shows under "Commits".