1 files changed, 133 insertions, 10 deletions

diff --git a/contrib/llvm-project/llvm/lib/CodeGen/BranchRelaxation.cpp b/contrib/llvm-project/llvm/lib/CodeGen/BranchRelaxation.cpp
index 05494f1ddc67..f50eb5e1730a 100644
--- a/contrib/llvm-project/llvm/lib/CodeGen/BranchRelaxation.cpp
+++ b/contrib/llvm-project/llvm/lib/CodeGen/BranchRelaxation.cpp
@@ -26,6 +26,7 @@
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/Format.h"
 #include "llvm/Support/raw_ostream.h"
+#include "llvm/Target/TargetMachine.h"
 #include <cassert>
 #include <cstdint>
 #include <iterator>
@@ -78,12 +79,19 @@ class BranchRelaxation : public MachineFunctionPass {
   };
 
   SmallVector<BasicBlockInfo, 16> BlockInfo;
+
+  // The basic block after which trampolines are inserted. This is the last
+  // basic block that isn't in the cold section.
+  MachineBasicBlock *TrampolineInsertionPoint = nullptr;
+  SmallDenseSet<std::pair<MachineBasicBlock *, MachineBasicBlock *>>
+      RelaxedUnconditionals;
   std::unique_ptr<RegScavenger> RS;
   LivePhysRegs LiveRegs;
 
   MachineFunction *MF = nullptr;
   const TargetRegisterInfo *TRI = nullptr;
   const TargetInstrInfo *TII = nullptr;
+  const TargetMachine *TM = nullptr;
 
   bool relaxBranchInstructions();
   void scanFunction();
@@ -142,7 +150,8 @@ void BranchRelaxation::verify() {
       if (MI.getOpcode() == TargetOpcode::FAULTING_OP)
         continue;
       MachineBasicBlock *DestBB = TII->getBranchDestBlock(MI);
-      assert(isBlockInRange(MI, *DestBB));
+      assert(isBlockInRange(MI, *DestBB) ||
+             RelaxedUnconditionals.contains({&MBB, DestBB}));
     }
   }
 #endif
@@ -165,15 +174,28 @@ void BranchRelaxation::scanFunction() {
   BlockInfo.clear();
   BlockInfo.resize(MF->getNumBlockIDs());
 
+  TrampolineInsertionPoint = nullptr;
+  RelaxedUnconditionals.clear();
+
   // First thing, compute the size of all basic blocks, and see if the function
   // has any inline assembly in it. If so, we have to be conservative about
   // alignment assumptions, as we don't know for sure the size of any
-  // instructions in the inline assembly.
-  for (MachineBasicBlock &MBB : *MF)
+  // instructions in the inline assembly. At the same time, place the
+  // trampoline insertion point at the end of the hot portion of the function.
+  for (MachineBasicBlock &MBB : *MF) {
     BlockInfo[MBB.getNumber()].Size = computeBlockSize(MBB);
 
+    if (MBB.getSectionID() != MBBSectionID::ColdSectionID)
+      TrampolineInsertionPoint = &MBB;
+  }
+
   // Compute block offsets and known bits.
   adjustBlockOffsets(*MF->begin());
+
+  if (TrampolineInsertionPoint == nullptr) {
+    LLVM_DEBUG(dbgs() << "  No suitable trampoline insertion point found in "
+                      << MF->getName() << ".\n");
+  }
 }
 
 /// computeBlockSize - Compute the size for MBB.
@@ -232,6 +254,11 @@ BranchRelaxation::createNewBlockAfter(MachineBasicBlock &OrigMBB,
   MachineBasicBlock *NewBB = MF->CreateMachineBasicBlock(BB);
   MF->insert(++OrigMBB.getIterator(), NewBB);
 
+  // Place the new block in the same section as OrigBB
+  NewBB->setSectionID(OrigMBB.getSectionID());
+  NewBB->setIsEndSection(OrigMBB.isEndSection());
+  OrigMBB.setIsEndSection(false);
+
   // Insert an entry into BlockInfo to align it properly with the block numbers.
   BlockInfo.insert(BlockInfo.begin() + NewBB->getNumber(), BasicBlockInfo());
 
@@ -241,8 +268,9 @@ BranchRelaxation::createNewBlockAfter(MachineBasicBlock &OrigMBB,
 /// Split the basic block containing MI into two blocks, which are joined by
 /// an unconditional branch.  Update data structures and renumber blocks to
 /// account for this change and returns the newly created block.
-MachineBasicBlock *BranchRelaxation::splitBlockBeforeInstr(MachineInstr &MI,
-                                                           MachineBasicBlock *DestBB) {
+MachineBasicBlock *
+BranchRelaxation::splitBlockBeforeInstr(MachineInstr &MI,
+                                        MachineBasicBlock *DestBB) {
   MachineBasicBlock *OrigBB = MI.getParent();
 
   // Create a new MBB for the code after the OrigBB.
@@ -250,6 +278,11 @@ MachineBasicBlock *BranchRelaxation::splitBlockBeforeInstr(MachineInstr &MI,
       MF->CreateMachineBasicBlock(OrigBB->getBasicBlock());
   MF->insert(++OrigBB->getIterator(), NewBB);
 
+  // Place the new block in the same section as OrigBB.
+  NewBB->setSectionID(OrigBB->getSectionID());
+  NewBB->setIsEndSection(OrigBB->isEndSection());
+  OrigBB->setIsEndSection(false);
+
   // Splice the instructions starting with MI over to NewBB.
   NewBB->splice(NewBB->end(), OrigBB, MI.getIterator(), OrigBB->end());
 
@@ -300,7 +333,12 @@ bool BranchRelaxation::isBlockInRange(
   int64_t BrOffset = getInstrOffset(MI);
   int64_t DestOffset = BlockInfo[DestBB.getNumber()].Offset;
 
-  if (TII->isBranchOffsetInRange(MI.getOpcode(), DestOffset - BrOffset))
+  const MachineBasicBlock *SrcBB = MI.getParent();
+
+  if (TII->isBranchOffsetInRange(MI.getOpcode(),
+                                 SrcBB->getSectionID() != DestBB.getSectionID()
+                                     ? TM->getMaxCodeSize()
+                                     : DestOffset - BrOffset))
     return true;
 
   LLVM_DEBUG(dbgs() << "Out of range branch to destination "
@@ -358,6 +396,50 @@ bool BranchRelaxation::fixupConditionalBranch(MachineInstr &MI) {
   assert(!Fail && "branches to be relaxed must be analyzable");
   (void)Fail;
 
+  // Since cross-section conditional branches to the cold section are rarely
+  // taken, try to avoid inverting the condition. Instead, add a "trampoline
+  // branch", which unconditionally branches to the branch destination. Place
+  // the trampoline branch at the end of the function and retarget the
+  // conditional branch to the trampoline.
+  // tbz L1
+  // =>
+  // tbz L1Trampoline
+  // ...
+  // L1Trampoline: b  L1
+  if (MBB->getSectionID() != TBB->getSectionID() &&
+      TBB->getSectionID() == MBBSectionID::ColdSectionID &&
+      TrampolineInsertionPoint != nullptr) {
+    // If the insertion point is out of range, we can't put a trampoline there.
+    NewBB =
+        createNewBlockAfter(*TrampolineInsertionPoint, MBB->getBasicBlock());
+
+    if (isBlockInRange(MI, *NewBB)) {
+      LLVM_DEBUG(dbgs() << "  Retarget destination to trampoline at "
+                        << NewBB->back());
+
+      insertUncondBranch(NewBB, TBB);
+
+      // Update the successor lists to include the trampoline.
+      MBB->replaceSuccessor(TBB, NewBB);
+      NewBB->addSuccessor(TBB);
+
+      // Replace branch in the current (MBB) block.
+      removeBranch(MBB);
+      insertBranch(MBB, NewBB, FBB, Cond);
+
+      TrampolineInsertionPoint = NewBB;
+      finalizeBlockChanges(MBB, NewBB);
+      return true;
+    }
+
+    LLVM_DEBUG(
+        dbgs() << "  Trampoline insertion point out of range for Bcc from "
+               << printMBBReference(*MBB) << " to " << printMBBReference(*TBB)
+               << ".\n");
+    TrampolineInsertionPoint->setIsEndSection(NewBB->isEndSection());
+    MF->erase(NewBB);
+  }
+
   // Add an unconditional branch to the destination and invert the branch
   // condition to jump over it:
   // tbz L1
@@ -462,7 +544,10 @@ bool BranchRelaxation::fixupUnconditionalBranch(MachineInstr &MI) {
   int64_t DestOffset = BlockInfo[DestBB->getNumber()].Offset;
   int64_t SrcOffset = getInstrOffset(MI);
 
-  assert(!TII->isBranchOffsetInRange(MI.getOpcode(), DestOffset - SrcOffset));
+  assert(!TII->isBranchOffsetInRange(
+      MI.getOpcode(), MBB->getSectionID() != DestBB->getSectionID()
+                          ? TM->getMaxCodeSize()
+                          : DestOffset - SrcOffset));
 
   BlockInfo[MBB->getNumber()].Size -= OldBrSize;
 
@@ -482,6 +567,8 @@ bool BranchRelaxation::fixupUnconditionalBranch(MachineInstr &MI) {
     BranchBB->sortUniqueLiveIns();
     BranchBB->addSuccessor(DestBB);
     MBB->replaceSuccessor(DestBB, BranchBB);
+    if (TrampolineInsertionPoint == MBB)
+      TrampolineInsertionPoint = BranchBB;
   }
 
   DebugLoc DL = MI.getDebugLoc();
@@ -492,15 +579,41 @@ bool BranchRelaxation::fixupUnconditionalBranch(MachineInstr &MI) {
   // be erased.
   MachineBasicBlock *RestoreBB = createNewBlockAfter(MF->back(),
                                                      DestBB->getBasicBlock());
+  std::prev(RestoreBB->getIterator())
+      ->setIsEndSection(RestoreBB->isEndSection());
+  RestoreBB->setIsEndSection(false);
 
   TII->insertIndirectBranch(*BranchBB, *DestBB, *RestoreBB, DL,
-                            DestOffset - SrcOffset, RS.get());
+                            BranchBB->getSectionID() != DestBB->getSectionID()
+                                ? TM->getMaxCodeSize()
+                                : DestOffset - SrcOffset,
+                            RS.get());
 
   BlockInfo[BranchBB->getNumber()].Size = computeBlockSize(*BranchBB);
   adjustBlockOffsets(*MBB);
 
-  // If RestoreBB is required, try to place just before DestBB.
+  // If RestoreBB is required, place it appropriately.
   if (!RestoreBB->empty()) {
+    // If the jump is Cold -> Hot, don't place the restore block (which is
+    // cold) in the middle of the function. Place it at the end.
+    if (MBB->getSectionID() == MBBSectionID::ColdSectionID &&
+        DestBB->getSectionID() != MBBSectionID::ColdSectionID) {
+      MachineBasicBlock *NewBB = createNewBlockAfter(*TrampolineInsertionPoint);
+      TII->insertUnconditionalBranch(*NewBB, DestBB, DebugLoc());
+      BlockInfo[NewBB->getNumber()].Size = computeBlockSize(*NewBB);
+
+      // New trampolines should be inserted after NewBB.
+      TrampolineInsertionPoint = NewBB;
+
+      // Retarget the unconditional branch to the trampoline block.
+      BranchBB->replaceSuccessor(DestBB, NewBB);
+      NewBB->addSuccessor(DestBB);
+
+      DestBB = NewBB;
+    }
+
+    // In all other cases, try to place just before DestBB.
+
     // TODO: For multiple far branches to the same destination, there are
     // chances that some restore blocks could be shared if they clobber the
     // same registers and share the same restore sequence. So far, those
@@ -525,9 +638,16 @@ bool BranchRelaxation::fixupUnconditionalBranch(MachineInstr &MI) {
     BlockInfo[RestoreBB->getNumber()].Size = computeBlockSize(*RestoreBB);
     // Update the offset starting from the previous block.
     adjustBlockOffsets(*PrevBB);
+
+    // Fix up section information for RestoreBB and DestBB
+    RestoreBB->setSectionID(DestBB->getSectionID());
+    RestoreBB->setIsBeginSection(DestBB->isBeginSection());
+    DestBB->setIsBeginSection(false);
+    RelaxedUnconditionals.insert({BranchBB, RestoreBB});
   } else {
     // Remove restore block if it's not required.
     MF->erase(RestoreBB);
+    RelaxedUnconditionals.insert({BranchBB, DestBB});
   }
 
   return true;
@@ -553,7 +673,8 @@ bool BranchRelaxation::relaxBranchInstructions() {
       // Unconditional branch destination might be unanalyzable, assume these
       // are OK.
       if (MachineBasicBlock *DestBB = TII->getBranchDestBlock(*Last)) {
-        if (!isBlockInRange(*Last, *DestBB)) {
+        if (!isBlockInRange(*Last, *DestBB) && !TII->isTailCall(*Last) &&
+            !RelaxedUnconditionals.contains({&MBB, DestBB})) {
           fixupUnconditionalBranch(*Last);
           ++NumUnconditionalRelaxed;
           Changed = true;
@@ -607,6 +728,7 @@ bool BranchRelaxation::runOnMachineFunction(MachineFunction &mf) {
 
   const TargetSubtargetInfo &ST = MF->getSubtarget();
   TII = ST.getInstrInfo();
+  TM = &MF->getTarget();
 
   TRI = ST.getRegisterInfo();
   if (TRI->trackLivenessAfterRegAlloc(*MF))
@@ -632,6 +754,7 @@ bool BranchRelaxation::runOnMachineFunction(MachineFunction &mf) {
   LLVM_DEBUG(dbgs() << "  Basic blocks after relaxation\n\n"; dumpBBs());
 
   BlockInfo.clear();
+  RelaxedUnconditionals.clear();
 
   return MadeChange;
 }