[llvm] Speed up branch relaxation from O(n^2) to O(n) (PR #96250)

[Daniel Hoekwater via llvm-commits]

https://github.com/dhoekwater updated https://github.com/llvm/llvm-project/pull/96250

>From 0dfe53c32080b7c93a69f3f8c4326a1f23f21ebc Mon Sep 17 00:00:00 2001
From: Daniel Hoekwater <hoekwater at google.com>
Date: Thu, 20 Jun 2024 23:06:29 +0000
Subject: [PATCH] [BranchRelaxation] Remove quadratic behavior in relaxation

Currently, we recompute block offsets after each relaxation. This causes
the complexity to be O(n^2) in the number of instructions, inflating compile
time.

If we instead recompute block offsets after each iteration of the outer
loop, the complexity is O(n). Recomputing offsets in the outer loop will
cause some out-of-range branches to be missed in the inner loop, but
they will be relaxed in the next iteration of the outer loop.

This change may alter the order that branches are relaxed, changing the
order of trampolines at the end of the function.

This change may also  introduce unnecessary relaxations for an
architecture where the relaxed branch is smaller than the unrelaxed
branch, but AFAIK there is no such architecture.
---
 llvm/lib/CodeGen/BranchRelaxation.cpp | 41 +++++++++++++++++++++------
 1 file changed, 32 insertions(+), 9 deletions(-)

diff --git a/llvm/lib/CodeGen/BranchRelaxation.cpp b/llvm/lib/CodeGen/BranchRelaxation.cpp
index f50eb5e1730a3..520b4a22678d3 100644
--- a/llvm/lib/CodeGen/BranchRelaxation.cpp
+++ b/llvm/lib/CodeGen/BranchRelaxation.cpp
@@ -103,6 +103,8 @@ class BranchRelaxation : public MachineFunctionPass {
   MachineBasicBlock *splitBlockBeforeInstr(MachineInstr &MI,
                                            MachineBasicBlock *DestBB);
   void adjustBlockOffsets(MachineBasicBlock &Start);
+  void adjustBlockOffsets(MachineBasicBlock &Start,
+                          MachineFunction::iterator End);
   bool isBlockInRange(const MachineInstr &MI, const MachineBasicBlock &BB) const;
 
   bool fixupConditionalBranch(MachineInstr &MI);
@@ -227,9 +229,14 @@ unsigned BranchRelaxation::getInstrOffset(const MachineInstr &MI) const {
 }
 
 void BranchRelaxation::adjustBlockOffsets(MachineBasicBlock &Start) {
+  adjustBlockOffsets(Start, MF->end());
+}
+
+void BranchRelaxation::adjustBlockOffsets(MachineBasicBlock &Start,
+                                          MachineFunction::iterator End) {
   unsigned PrevNum = Start.getNumber();
   for (auto &MBB :
-       make_range(std::next(MachineFunction::iterator(Start)), MF->end())) {
+       make_range(std::next(MachineFunction::iterator(Start)), End)) {
     unsigned Num = MBB.getNumber();
     // Get the offset and known bits at the end of the layout predecessor.
     // Include the alignment of the current block.
@@ -314,8 +321,8 @@ BranchRelaxation::splitBlockBeforeInstr(MachineInstr &MI,
   // block, it may contain a tablejump.
   BlockInfo[NewBB->getNumber()].Size = computeBlockSize(*NewBB);
 
-  // All BBOffsets following these blocks must be modified.
-  adjustBlockOffsets(*OrigBB);
+  // Update the offset of the new block.
+  adjustBlockOffsets(*OrigBB, std::next(NewBB->getIterator()));
 
   // Need to fix live-in lists if we track liveness.
   if (TRI->trackLivenessAfterRegAlloc(*MF))
@@ -384,11 +391,17 @@ bool BranchRelaxation::fixupConditionalBranch(MachineInstr &MI) {
 
   auto finalizeBlockChanges = [&](MachineBasicBlock *MBB,
                                   MachineBasicBlock *NewBB) {
-    // Keep the block offsets up to date.
-    adjustBlockOffsets(*MBB);
+    if (NewBB == nullptr)
+      return;
+
+    // Keep the block offsets approximately up to date. While they will be
+    // slight underestimates, we will update them appropriately in the next
+    // pass.
+    adjustBlockOffsets(*std::prev(NewBB->getIterator()),
+                       std::next(NewBB->getIterator()));
 
     // Need to fix live-in lists if we track liveness.
-    if (NewBB && TRI->trackLivenessAfterRegAlloc(*MF))
+    if (TRI->trackLivenessAfterRegAlloc(*MF))
       computeAndAddLiveIns(LiveRegs, *NewBB);
   };
 
@@ -438,6 +451,7 @@ bool BranchRelaxation::fixupConditionalBranch(MachineInstr &MI) {
                << ".\n");
     TrampolineInsertionPoint->setIsEndSection(NewBB->isEndSection());
     MF->erase(NewBB);
+    NewBB = nullptr;
   }
 
   // Add an unconditional branch to the destination and invert the branch
@@ -590,7 +604,7 @@ bool BranchRelaxation::fixupUnconditionalBranch(MachineInstr &MI) {
                             RS.get());
 
   BlockInfo[BranchBB->getNumber()].Size = computeBlockSize(*BranchBB);
-  adjustBlockOffsets(*MBB);
+  adjustBlockOffsets(*MBB, std::next(BranchBB->getIterator()));
 
   // If RestoreBB is required, place it appropriately.
   if (!RestoreBB->empty()) {
@@ -601,6 +615,8 @@ bool BranchRelaxation::fixupUnconditionalBranch(MachineInstr &MI) {
       MachineBasicBlock *NewBB = createNewBlockAfter(*TrampolineInsertionPoint);
       TII->insertUnconditionalBranch(*NewBB, DestBB, DebugLoc());
       BlockInfo[NewBB->getNumber()].Size = computeBlockSize(*NewBB);
+      adjustBlockOffsets(*TrampolineInsertionPoint,
+                         std::next(NewBB->getIterator()));
 
       // New trampolines should be inserted after NewBB.
       TrampolineInsertionPoint = NewBB;
@@ -636,8 +652,8 @@ bool BranchRelaxation::fixupUnconditionalBranch(MachineInstr &MI) {
       computeAndAddLiveIns(LiveRegs, *RestoreBB);
     // Compute the restore block size.
     BlockInfo[RestoreBB->getNumber()].Size = computeBlockSize(*RestoreBB);
-    // Update the offset starting from the previous block.
-    adjustBlockOffsets(*PrevBB);
+    // Update the estimated offset for the restore block.
+    adjustBlockOffsets(*PrevBB, DestBB->getIterator());
 
     // Fix up section information for RestoreBB and DestBB
     RestoreBB->setSectionID(DestBB->getSectionID());
@@ -718,6 +734,13 @@ bool BranchRelaxation::relaxBranchInstructions() {
     }
   }
 
+  // If we relaxed a branch, subsequent block offsets may be underestimated.
+  // Recompute them to guarantee that we relax any branches that we pushed out
+  // of range.
+  if (Changed) {
+    adjustBlockOffsets(MF->front());
+  }
+
   return Changed;
 }