[AArch64][CodeGen] Avoid inverting hot branches during relaxation

[dhoekwater]

Summary

Current behavior for relaxing out-of-range conditional branches
is to invert the conditional and insert a fallthrough unconditional
branch to the original destination. This approach biases the branch
predictor in the wrong direction, which can degrading performance.

Machine function splitting introduces many rarely-taken cross-section
conditional branches, which are improperly relaxed. Avoid inverting
these branches; instead, retarget them to trampolines at the end of the
function. Doing so increases the runtime cost of jumping to cold code
but eliminates the misprediction cost of jumping to hot code.

[dhoekwater]

Add exhaustive tests

[dhoekwater]

Remove debug lines

[dhoekwater]

Rebase onto up-to-date parent

[dhoekwater]

Rebase parent onto main and add clearer comments

[dhoekwater]

removed reviewers because I initially added way too many.

[arsenm]

llvm/lib/CodeGen/BranchRelaxation.cpp

184

this just scanned through the function, couldn't you track this on that first loop?

403–427

Split this into a helper function?

[dhoekwater]

Move trampoline initialization into an existing loop

[dhoekwater]

llvm/lib/CodeGen/BranchRelaxation.cpp

403–427

insertUncondBranch, removeBranch, insertBranch, and finalizeBlockChanges are function-local lambdas, so moving this all into a helper function would require moving those to helper functions as well.

I think abstracting away BlockInfo maintenance into helper functions could be a good idea, but doing so in this change would probably be overly ambitious.

[dhoekwater]

Remove the closed dependency from the commit message

[arsenm]

LGTM with some nits

llvm/lib/CodeGen/BranchRelaxation.cpp

85

null initialize and remove the one in scanFunction?

184

Need a << '\n'

726

redundant with clearing this at the start of the scan

[dhoekwater]

Make the comments more obvious and remove an unnecessary assignment

[dhoekwater]

A bit more info on the motivation for this patch:

Machine function splitting (MFS) places super cold basic blocks in a different ELF section (the .test.split section), so any hot branches targeting the cold basic blocks must be relaxed. The current method of relaxing a conditional branch is to invert the condition, which biases the branch predictor in the opposite direction. Since these branches are almost never taken, the bias has a significant impact on runtime performance.

This patch mitigates the impact of MFS on the branch predictor by avoiding inverting the conditional branch. Instead, if the spot right after the last hot block is in range, we can 1) place a "trampoline," a basic block that only consists of an unconditional to the destination, at this location and 2) retarget the conditional branch to the trampoline. This increases the cost of actually taking the branch to the cold section, but it is almost never taken; doing so significantly improves performance when the branch is not taken.

llvm/lib/CodeGen/BranchRelaxation.cpp

85

We either need the null assignment in scanFunction or runOnMachineFunction since just null initializing it will keep the insertion point from the previous function. I'll keep the one in scanFunction, since it doesn't require that the previous call to runOnMachineFunction executed successfully.

[dhoekwater]

Rebase onto main

[dhoekwater]

Rebase onto main to hopefully resolve some spurious test failures

[dhoekwater]

Relanded in e223e4567722661f1b32aa052cd13d9f47b896d9.

llvm/lib/CodeGen/BranchRelaxation.cpp

unsigned postOffset(const MachineBasicBlock &MBB) const {
 
       // The alignment of this MBB is larger than the function's alignment, so we
       // can't tell whether or not it will insert nops. Assume that it will.
       return alignTo(PO, Alignment) + Alignment.value() - ParentAlign.value();
     }
   };
 
   SmallVector<BasicBlockInfo, 16> BlockInfo;
+  MachineBasicBlock *TrampolineInsertionPoint;
   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();
 

 }
 #endif
 
 /// scanFunction - Do the initial scan of the function, building up
 /// information about each block.
 void BranchRelaxation::scanFunction() {
   BlockInfo.clear();
   BlockInfo.resize(MF->getNumBlockIDs());
+  TrampolineInsertionPoint = nullptr;
 
   // 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)
     BlockInfo[MBB.getNumber()].Size = computeBlockSize(MBB);
 
   // Compute block offsets and known bits.
   adjustBlockOffsets(*MF->begin());
+
+  // Place the trampoline insertion point at the end of the hot portion of the
+  // function.
+  for (MachineBasicBlock &MBB : reverse(*MF)) {
+    if (MBB.getSectionID() != MBBSectionID::ColdSectionID) {
+      TrampolineInsertionPoint = &MBB;
+      break;
+    }
+  }
+  if (TrampolineInsertionPoint == nullptr) {
+    LLVM_DEBUG(dbgs() << "  No suitable trampoline insertion point found in "
+                      << MF->getName());
+  }
 }
 
 /// computeBlockSize - Compute the size for MBB.
 uint64_t BranchRelaxation::computeBlockSize(const MachineBasicBlock &MBB) const {
   uint64_t Size = 0;
   for (const MachineInstr &MI : MBB)
     Size += TII->getInstSizeInBytes(MI);
   return Size;
auto finalizeBlockChanges = [&](MachineBasicBlock *MBB,
     if (NewBB && TRI->trackLivenessAfterRegAlloc(*MF))
       computeAndAddLiveIns(LiveRegs, *NewBB);
   };
 
   bool Fail = TII->analyzeBranch(*MBB, TBB, FBB, Cond);
   assert(!Fail && "branches to be relaxed must be analyzable");
   (void)Fail;
 
+  // For cross-section branches whose destination is in the cold section, 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
   // =>
   // tbnz L2
   // b   L1
   // L2:
 
while (relaxBranchInstructions())
     MadeChange = true;
 
   // After a while, this might be made debug-only, but it is not expensive.
   verify();
 
   LLVM_DEBUG(dbgs() << "  Basic blocks after relaxation\n\n"; dumpBBs());
 
   BlockInfo.clear();
+  TrampolineInsertionPoint = nullptr;
 
   return MadeChange;
 }