[AArch64] [BranchRelaxation] Optimize for hot code size in AArch64 branch relaxation

[dhoekwater]

Summary

On AArch64, it is safe to let the linker handle relaxation of
unconditional branches; in most cases, the destination is within range,
and the linker doesn't need to do anything. If the linker does insert
fixup code, it clobbers the x16 inter-procedural register, so x16 must
be available across the branch before linking. If x16 isn't available,
but some other register is, we can relax the branch either by spilling
x16 OR using the free register for a manually-inserted indirect branch.

This patch builds on D145211. While that patch is for correctness, this
one is for performance of the common case. As noted in
https://reviews.llvm.org/D145211#4537173, we can trust the linker to
relax cross-section unconditional branches across which x16 is
available.

Programs that use machine function splitting care most about the
performance of hot code at the expense of the performance of cold code,
so we prioritize minimizing hot code size.

Here's a breakdown of the cases:

Hot -> Cold [x16 is free across the branch]
  Do nothing; let the linker relax the branch.

Cold -> Hot [x16 is free across the branch]
  Do nothing; let the linker relax the branch.

Hot -> Cold [x16 used across the branch, but there is a free register]
  Spill x16; let the linker relax the branch.

  Spilling requires fewer instructions than manually inserting an
  indirect branch.

Cold -> Hot [x16 used across the branch, but there is a free register]
  Manually insert an indirect branch.

  Spilling would require adding a restore block in the hot section.

 Hot -> Cold [No free regs]
   Spill x16; let the linker relax the branch.

 Cold -> Hot [No free regs]
   Spill x16 and put the restore block at the end of the hot function; let the linker relax the branch.
   Ex:
     [Hot section]
     func.hot:
       ... hot code...
     func.restore:
       ... restore x16 ...
       B func.hot

     [Cold section]
       func.cold:
       ... spill x16 ...
       B func.restore

   Putting the restore block at the end of the function instead of
   just before the destination increases the cost of executing the
   store, but it avoids putting cold code in the middle of hot code.
   Since the restore is very rarely taken, this is a worthwhile
   tradeoff.

[dhoekwater]

Here's a breakdown of the cases:

Hot -> Cold [x16 is free across the branch]
  Do nothing; let the linker relax the branch.

Cold -> Hot [x16 is free across the branch]
  Do nothing; let the linker relax the branch.

Hot -> Cold [x16 not free, but others are]
  Spill x16; let the linker relax the branch.

Cold -> Hot [x16 not free, but others are]
  Manually insert an indirect branch.

Hot -> Cold [No free regs]
  Spill x16; let the linker relax the branch.

Cold -> Hot [No free regs]
  Spill x16 and put the restore block at the end of the hot function; let the linker relax the branch.
  Ex:
    [Hot section]
    func.hot:
      ... hot code...
    func.restore:
      ... restore x16 ...
      B func.hot

    [Cold section]
      func.cold:
      ... spill x16 ...
      B func.restore

[dhoekwater]

Rebase onto up-to-date head

[dhoekwater]

removed reviewers because I initially added way too many.

[mingmingl]

In D156767#4549141, @dhoekwater wrote:

Here's a breakdown of the cases:

Hot -> Cold [x16 is free across the branch]
  Do nothing; let the linker relax the branch.

Cold -> Hot [x16 is free across the branch]
  Do nothing; let the linker relax the branch.

Hot -> Cold [x16 not free, but others are]
  Spill x16; let the linker relax the branch.

Cold -> Hot [x16 not free, but others are]
  Manually insert an indirect branch.

Hot -> Cold [No free regs]
  Spill x16; let the linker relax the branch.

Cold -> Hot [No free regs]
  Spill x16 and put the restore block at the end of the hot function; let the linker relax the branch.
  Ex:
    [Hot section]
    func.hot:
      ... hot code...
    func.restore:
      ... restore x16 ...
      B func.hot

    [Cold section]
      func.cold:
      ... spill x16 ...
      B func.restore

The change overall looks reasonable to me and the breakdown is well written.

To make the motivation slightly more clearer, maybe describe what's the current status without this patch (hot blocks not optimized, etc), and how this patch helps when cold blocks are split by function-splitting.

In D156767#4549141, @dhoekwater wrote:

Here's a breakdown of the cases:

Hot -> Cold [x16 is free across the branch]
  Do nothing; let the linker relax the branch.

Cold -> Hot [x16 is free across the branch]
  Do nothing; let the linker relax the branch.

Hot -> Cold [x16 not free, but others are]
  Spill x16; let the linker relax the branch.

Cold -> Hot [x16 not free, but others are]
  Manually insert an indirect branch.

Hot -> Cold [No free regs]
  Spill x16; let the linker relax the branch.

Cold -> Hot [No free regs]
  Spill x16 and put the restore block at the end of the hot function; let the linker relax the branch.
  Ex:
    [Hot section]
    func.hot:
      ... hot code...
    func.restore:
      ... restore x16 ...
      B func.hot

    [Cold section]
      func.cold:
      ... spill x16 ...
      B func.restore

llvm/test/CodeGen/AArch64/branch-relax-b.ll

134

nit: add a brief comment about which case (in the breakdown list) is being tested for posterity.

If I read correctly, bb.3 is the cold block (that would be split out by MFS), and the restore block .LBB2_4 is the restore block and put at the end of the function. For my information, is x16 spilled as the 4th or 6th case?

llvm/test/CodeGen/AArch64/branch-relax-cross-section.mir

31

nit: maybe add test cases in an NFC (https://llvm.org/docs/Lexicon.html#n) and manually put cold blocks at the end of the IR function. This way, how the IR function looks like after function-splitting is easier to read, and the diff of prioritized hot code is clearer.

70–72

for my information, how does branch-relaxation pass tell cold: is a cold block for test case all_used_cold_to_hot? It seems entry will branch to cold unconditionally but I may miss something.

[dhoekwater]

Without this patch, out-of-range branches are correctly relaxed, but
they are relaxed suboptimally. Notably, most (~99%) of out-of-range
unconditional branches have X16 available across the branch, where we
can safely defer to the linker rather than inserting 2 extra
instructions. Because machine function splitting (MFS) makes
unconditional branches to cold blocks out of range, enabling MFS would
incur a big code size overhead due to cross-section branch relaxation.
This patch lets us mitigate that overhead in the vast majority of cases.

[dhoekwater]

llvm/test/CodeGen/AArch64/branch-relax-b.ll

134

Added a comment explaining the case a bit further. In this test, machine function splitting isn't enabled at all, so the function is made up of entirely hot blocks. For this case, the only thing that matters is that we're branching _from_ the hot section, so we don't need to use function splitting in this test.

llvm/test/CodeGen/AArch64/branch-relax-cross-section.mir

70–72

This part of the test (the LLVM IR) is only present so llc is happy; the brunt of the test details are specified in the MIR sections below. For this specific test case, it is specified in

218    bb.2.iftrue (bbsections Cold):

[dhoekwater]

Remove the closed dependency from the commit message

[dhoekwater]

Remove the closed dependency from the commit message

[dhoekwater]

Accidentally updated the wrong patch

[arsenm]

llvm/lib/CodeGen/BranchRelaxation.cpp

597

Broken MBB check, MBB cannot be null

[dhoekwater]

Remove a redundant null check

[dhoekwater]

llvm/lib/CodeGen/BranchRelaxation.cpp

597

Ah, debugging change, I'm not sure how this made it in. Thanks for the feedback.

[mingmingl]

llvm/test/CodeGen/AArch64/branch-relax-cross-section.mir

248–259

I should have missed something here.. Where does bb.4 and bb.5 come from for 'x16_used_cold_to_hot'? Testing the IR (https://gcc.godbolt.org/z/esjchq3PK) gives 3 blocks..

Relatedly, I think submitting test cases as NFC and rebase on top of that should display the diff clearer.

[dhoekwater]

llvm/test/CodeGen/AArch64/branch-relax-cross-section.mir

248–259

Where does bb.4 and bb.5 come from for 'x16_used_cold_to_hot'?

bb.3, bb.4, and bb.5 are all machine basic blocks created by the Branch Relaxation pass in the fixupUnconditionalBranch function.

• bb.3 is a block that contains an unconditional branch that will be safely relaxed by the linker.
• bb.4 is a block that restores the contents of X16 that have been written to the stack in bb.3.entry and clobbered by the linker. For correctness, X16 must be restored before bb.2.cold executes.
• bb.5 is a block that contains an indirect branch to bb.1.hot. The indirect branch 1) loads the address of bb.1 into an available register by loading the page address then adding the page offset and 2) branches to the register containing the address.

This test in particular is checking that, when a cold -> hot unconditional branch is relaxed, it eagerly inserts an indirect branch using a free register, if there is one.

[dhoekwater]

llvm/test/CodeGen/AArch64/branch-relax-cross-section.mir

248–259

Relatedly, I think submitting test cases as NFC and rebase on top of that should display the diff clearer.

In the NFC, do I add the MIR for the tests but not the CHECKs? Or do I add checks for the wrong behavior?

[dhoekwater]

llvm/test/CodeGen/AArch64/branch-relax-cross-section.mir

248–259

Tests added in an NFC.

[dhoekwater]

Remove redundant assignments to TrampolineInsertionPoint

[mingmingl]

thanks! The implementation makes sense to me. I'm not very familiar with this pass though, so maybe we could wait a little bit more for owners.

llvm/lib/CodeGen/BranchRelaxation.cpp

595

nit: move this after the cold->hot handling ( if (MBB->getSectionID() == MBBSectionID::ColdSectionID && DestBB->getSectionID() != MBBSectionID::ColdSectionID)), or maybe update it accordingly.

[dhoekwater]

In D156767#4632641, @mingmingl wrote:

thanks! The implementation makes sense to me. I'm not very familiar with this pass though, so maybe we could wait a little bit more for owners.

I think MaskRay is the only reviewer who's contributed to both BranchRelaxation and AArch64 code. @MaskRay, could you please take a look?

[mingmingl]

lgtm and approve. Maybe give other reviewers one or two days (after the holiday) before commit this.

[dhoekwater]

In D156767#4639060, @mingmingl wrote:

lgtm and approve. Maybe give other reviewers one or two days (after the holiday) before commit this.

Sounds good, I'll plan to commit it close to EOD on Wednesday or early Thursday barring any further feedback.

[MaskRay]

It seems that this patch should link to D145211. This patch will address this comment https://reviews.llvm.org/D145211#4537173

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

285–286

This sentence sounds like the linker needs to do work, but in the majority of cases I think the linker won't need to insert a range extension thunk.

Perhaps:

// If X16 is unused, we can rely on linker range extension thunk in case NewDestBB cannot be reached by a single B (e.g. cross-section).

287

Can we use another variable so that we can use const Register Reg = AArch64::X16; above and remove Reg = AArch64::X16; on L312?

[MaskRay]

In D156767#4549141, @dhoekwater wrote:

Here's a breakdown of the cases:
[...]

It'd be great to fold this information to the summary.

Hot -> Cold [x16 not free, but others are]

but there is an unused register

[dhoekwater]

Clarify some comments and add a temporary variable to make indirect branch insertion cleaner

[dhoekwater]

In D156767#4639219, @MaskRay wrote:

It'd be great to fold this information to the summary.
...
but there is an unused register

Folded into the summary and updated phrasing.

[MaskRay]

llvm/lib/CodeGen/BranchRelaxation.cpp

605

End complete sentences with a period.

608

period

[dhoekwater]

Properly punctuate comments

llvm/lib/CodeGen/BranchRelaxation.cpp

struct BasicBlockInfo {
     }
   };
 
   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;
 
for (MachineBasicBlock &MBB : *MF) {
     for (MachineBasicBlock::iterator J = MBB.getFirstTerminator();
          J != MBB.end(); J = std::next(J)) {
       MachineInstr &MI = *J;
       if (!MI.isConditionalBranch() && !MI.isUnconditionalBranch())
         continue;
       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
 }
 
 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
 /// print block size and offset information - debugging
 LLVM_DUMP_METHOD void BranchRelaxation::dumpBBs() {
   for (auto &MBB : *MF) {
     const BasicBlockInfo &BBI = BlockInfo[MBB.getNumber()];
     dbgs() << format("%%bb.%u\toffset=%08x\t", MBB.getNumber(), BBI.Offset)
            << format("size=%#x\n", BBI.Size);
   }
 }
 #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;
+  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. 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->empty()) {
     for (const MachineBasicBlock *Succ : MBB->successors()) {
       for (const MachineBasicBlock::RegisterMaskPair &LiveIn : Succ->liveins())
         BranchBB->addLiveIn(LiveIn);
     }
 
     BranchBB->sortUniqueLiveIns();
     BranchBB->addSuccessor(DestBB);
     MBB->replaceSuccessor(DestBB, BranchBB);
+    if (TrampolineInsertionPoint == MBB)
+      TrampolineInsertionPoint = BranchBB;
   }
 
   DebugLoc DL = MI.getDebugLoc();
   MI.eraseFromParent();
 
   // Create the optional restore block and, initially, place it at the end of
   // function. That block will be placed later if it's used; otherwise, it will
   // 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,
                             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
     // restore blocks are just duplicated for each far branch.
     assert(!DestBB->isEntryBlock());
     MachineBasicBlock *PrevBB = &*std::prev(DestBB->getIterator());
     // Fall through only if PrevBB has no unconditional branch as one of its
     // terminators.
if (!RestoreBB->empty()) {
     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;
 }
 
 bool BranchRelaxation::relaxBranchInstructions() {
   bool Changed = false;
 
for (MachineBasicBlock &MBB : *MF) {
     // conditional branch, this will end up changing the branch destination of
     // it to be over the newly inserted indirect branch block, which may avoid
     // the need to try expanding the conditional branch first, saving an extra
     // jump.
     if (Last->isUnconditionalBranch()) {
       // Unconditional branch destination might be unanalyzable, assume these
       // are OK.
       if (MachineBasicBlock *DestBB = TII->getBranchDestBlock(*Last)) {
-        if (!isBlockInRange(*Last, *DestBB) && !TII->isTailCall(*Last)) {
+        if (!isBlockInRange(*Last, *DestBB) && !TII->isTailCall(*Last) &&
+            !RelaxedUnconditionals.contains({&MBB, DestBB})) {
           fixupUnconditionalBranch(*Last);
           ++NumUnconditionalRelaxed;
           Changed = true;
         }
       }
     }
 
     // Loop over the conditional branches.
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();
+  RelaxedUnconditionals.clear();
 
   return MadeChange;
 }

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

BuildMI(MBB, MBB.end(), DL, get(AArch64::ADRP), Reg)
         .addSym(DestBB.getSymbol(), AArch64II::MO_PAGE);
     BuildMI(MBB, MBB.end(), DL, get(AArch64::ADDXri), Reg)
         .addReg(Reg)
         .addSym(DestBB.getSymbol(), AArch64II::MO_PAGEOFF | AArch64II::MO_NC)
         .addImm(0);
     BuildMI(MBB, MBB.end(), DL, get(AArch64::BR)).addReg(Reg);
   };
 
   RS->enterBasicBlockEnd(MBB);
-  Register Reg = RS->FindUnusedReg(&AArch64::GPR64RegClass);
-
-  // If there's a free register, manually insert the indirect branch using it.
-  if (Reg != AArch64::NoRegister) {
-    buildIndirectBranch(Reg, NewDestBB);
+  // If X16 is unused, we can rely on the linker to insert a range extension
+  // thunk if NewDestBB is out of range of a single B instruction.
+  constexpr Register Reg = AArch64::X16;
+  if (!RS->isRegUsed(Reg)) {
+    insertUnconditionalBranch(MBB, &NewDestBB, DL);
     RS->setRegUsed(Reg);
     return;
   }
 
-  // Otherwise, spill and use X16. This briefly moves the stack pointer, making
-  // it incompatible with red zones.
+  // If there's a free register and it's worth inflating the code size,
+  // manually insert the indirect branch.
+  Register Scavenged = RS->FindUnusedReg(&AArch64::GPR64RegClass);
+  if (Scavenged != AArch64::NoRegister &&
+      MBB.getSectionID() == MBBSectionID::ColdSectionID) {
+    buildIndirectBranch(Scavenged, NewDestBB);
+    RS->setRegUsed(Scavenged);
+    return;
+  }
+
+  // Note: Spilling X16 briefly moves the stack pointer, making it incompatible
+  // with red zones.
   AArch64FunctionInfo *AFI = MBB.getParent()->getInfo<AArch64FunctionInfo>();
   if (!AFI || AFI->hasRedZone().value_or(true))
     report_fatal_error(
         "Unable to insert indirect branch inside function that has red zone");
 
-  Reg = AArch64::X16;
+  // Otherwise, spill X16 and defer range extension to the linker.
   BuildMI(MBB, MBB.end(), DL, get(AArch64::STRXpre))
       .addReg(AArch64::SP, RegState::Define)
       .addReg(Reg)
       .addReg(AArch64::SP)
       .addImm(-16);
 
-  buildIndirectBranch(Reg, RestoreBB);
+  BuildMI(MBB, MBB.end(), DL, get(AArch64::B)).addMBB(&RestoreBB);
 
   BuildMI(RestoreBB, RestoreBB.end(), DL, get(AArch64::LDRXpost))
       .addReg(AArch64::SP, RegState::Define)
       .addReg(Reg, RegState::Define)
       .addReg(AArch64::SP)
       .addImm(16);
 }
 

llvm/test/CodeGen/AArch64/branch-relax-b.ll

 ; RUN: llc < %s -mtriple=aarch64-none-linux-gnu --verify-machineinstrs -aarch64-b-offset-bits=9 -aarch64-tbz-offset-bits=6 -aarch64-cbz-offset-bits=6 -aarch64-bcc-offset-bits=6 | FileCheck %s
 
 define void @relax_b_nospill(i1 zeroext %0) {
 ; CHECK-LABEL: relax_b_nospill:
 ; CHECK:       // %bb.0:                               // %entry
 ; CHECK-NEXT:    tbnz w0,
 ; CHECK-SAME:                 LBB0_1
 ; CHECK-NEXT:  // %bb.3:                               // %entry
-; CHECK-NEXT:          adrp    [[SCAVENGED_REGISTER:x[0-9]+]], .LBB0_2
-; CHECK-NEXT:          add     [[SCAVENGED_REGISTER]], [[SCAVENGED_REGISTER]], :lo12:.LBB0_2
-; CHECK-NEXT:          br      [[SCAVENGED_REGISTER]]
+; CHECK-NEXT:          b      .LBB0_2
 ; CHECK-NEXT:  .LBB0_1:                                // %iftrue
 ; CHECK-NEXT:          //APP
 ; CHECK-NEXT:          .zero   2048
 ; CHECK-NEXT:          //NO_APP
 ; CHECK-NEXT:          ret
 ; CHECK-NEXT:  .LBB0_2:                                // %iffalse
 ; CHECK-NEXT:          //APP
 ; CHECK-NEXT:          .zero   8

 ; CHECK:                  //APP
 ; CHECK-COUNT-29:         mov     {{x[0-9]+}},
 ; CHECK-NOT:              mov     {{x[0-9]+}},
 ; CHECK-NEXT:             //NO_APP
 ; CHECK-NEXT:             b.eq    .LBB1_1
 ; CHECK-NEXT:     // %bb.4:                               // %entry
 ; CHECK-NEXT:             str     [[SPILL_REGISTER:x[0-9]+]], [sp,
 ; CHECK-SAME:                                                       -16]!
-; CHECK-NEXT:             adrp    [[SPILL_REGISTER]], .LBB1_5
-; CHECK-NEXT:             add     [[SPILL_REGISTER]], [[SPILL_REGISTER]], :lo12:.LBB1_5
-; CHECK-NEXT:             br      [[SPILL_REGISTER]]
+; CHECK-NEXT:             b       .LBB1_5
 ; CHECK-NEXT:     .LBB1_1:                                // %iftrue
 ; CHECK-NEXT:             //APP
 ; CHECK-NEXT:             .zero   2048
 ; CHECK-NEXT:             //NO_APP
 ; CHECK-NEXT:             b       .LBB1_3
 ; CHECK-NEXT:     .LBB1_5:                                // %iffalse
 ; CHECK-NEXT:             ldr     [[SPILL_REGISTER]], [sp], 
 ; CHECK-SAME:                                                        16
iffalse:
   call void asm sideeffect "# reg use $0", "{x24}"(i64 %x24)
   call void asm sideeffect "# reg use $0", "{x25}"(i64 %x25)
   call void asm sideeffect "# reg use $0", "{x26}"(i64 %x26)
   call void asm sideeffect "# reg use $0", "{x27}"(i64 %x27)
   call void asm sideeffect "# reg use $0", "{x28}"(i64 %x28)
   ret void
 }
 
 define void @relax_b_x16_taken() {
 ; CHECK-LABEL:    relax_b_x16_taken:                      // @relax_b_x16_taken
-; COM: Pre-commit to record the behavior of relaxing an unconditional
-; COM: branch across which x16 is taken.
+; COM: Since the source of the out-of-range branch is hot and x16 is
+; COM: taken, it makes sense to spill x16 and let the linker insert
+; COM: fixup code for this branch rather than inflating the hot code
+; COM: size by eagerly relaxing the unconditional branch.
 ; CHECK:          // %bb.0:                               // %entry
 ; CHECK-NEXT:             //APP
 ; CHECK-NEXT:             mov     x16, #1
 ; CHECK-NEXT:             //NO_APP
 ; CHECK-NEXT:             cbnz    x16, .LBB2_1
 ; CHECK-NEXT:     // %bb.3:                               // %entry
-; CHECK-NEXT:             adrp    [[SCAVENGED_REGISTER2:x[0-9]+]], .LBB2_2
-; CHECK-NEXT:             add     [[SCAVENGED_REGISTER2]], [[SCAVENGED_REGISTER2]], :lo12:.LBB2_2
-; CHECK-NEXT:             br      [[SCAVENGED_REGISTER2]]
+; CHECK-NEXT:             str     [[SPILL_REGISTER]], [sp,
+; CHECK-SAME:                                                       -16]!
+; CHECK-NEXT:             b       .LBB2_4
 ; CHECK-NEXT:     .LBB2_1:                                // %iftrue
 ; CHECK-NEXT:             //APP
 ; CHECK-NEXT:             .zero   2048
 ; CHECK-NEXT:             //NO_APP
 ; CHECK-NEXT:             ret
-; CHECK-NEXT:     .LBB2_2:                                // %iffalse
+; CHECK-NEXT:     .LBB2_4:                                // %iffalse
+; CHECK-NEXT:             ldr     [[SPILL_REGISTER]], [sp], 
+; CHECK-SAME:                                                        16
+; CHECK-NEXT:     // %bb.2:                               // %iffalse
 ; CHECK-NEXT:             //APP
 ; CHECK-NEXT:             // reg use x16
 ; CHECK-NEXT:             //NO_APP
 ; CHECK-NEXT:             ret
 entry:
   %x16 = call i64 asm sideeffect "mov x16, 1", "={x16}"()
 
   %cmp = icmp eq i64 %x16, 0
   br i1 %cmp, label %iffalse, label %iftrue
 
 iftrue:
   call void asm sideeffect ".space 2048", ""()
   ret void
 
 iffalse:
   call void asm sideeffect "# reg use $0", "{x16}"(i64 %x16)
   ret void
 }
 
 declare i32 @bar()
-declare i32 @baz()
+declare i32 @baz()
 

llvm/test/CodeGen/AArch64/branch-relax-cross-section.mir

--- |
   }
 
   define void @tbz_hot_to_cold(i1 zeroext %0) {
     br i1 %0, label %hot_block, label %cold_block
   
   hot_block:                                        ; preds = %1
     %2 = call i32 @baz()
     br label %end
   
   end:                                              ; preds = %cold_block, %hot_block
     %3 = tail call i32 @qux()
     ret void
   
   cold_block:                                       ; preds = %1
     %4 = call i32 @bar()
     br label %end
   }
cold_block:                                       ; preds = %1
     %2 = call i32 @baz()
     br label %end
   
   end:                                              ; preds = %hot_block, %cold_block
     %3 = tail call i32 @qux()
     ret void
   
   hot_block:                                        ; preds = %1
     %4 = call i32 @bar()
     br label %end
   }
 
   define void @tbz_tramp_pushed_oob(i1 zeroext %0, i1 zeroext %1) {
   entry:
     %x16 = call i64 asm sideeffect "mov x16, 1", "={x16}"()
     br i1 %0, label %unrelaxable, label %cold
 
   unrelaxable:                                      ; preds = %entry
     br i1 %1, label %end, label %cold
body:             |
   ; CHECK-SAME:                  , %bb.3
   ; CHECK:  TBZW
   ; CHECK-SAME: %bb.3
   ; CHECK:  bb.1.hot_block:
   ; CHECK:    TCRETURNdi @qux, 0, csr_aarch64_aapcs, implicit $sp
   ; CHECK:  bb.3 (%ir-block.1):
   ; CHECK-NEXT:    successors: %bb.2
   ; CHECK-NEXT: {{  $}}
   ; CHECK-NEXT:    B %bb.2
   ; CHECK-NEXT: {{  $}}
   ; CHECK-NEXT:  bb.2.cold_block (bbsections Cold):
   ; CHECK:    TCRETURNdi @qux, 0, csr_aarch64_aapcs, implicit $sp
   bb.0 (%ir-block.1):
     successors: %bb.1, %bb.2
     liveins: $w0, $lr
 
     early-clobber $sp = frame-setup STRXpre killed $lr, $sp, -16 :: (store (s64) into %stack.0)
     TBZW killed renamable $w0, 0, %bb.2
 
   bb.1.hot_block:
     BL @baz, csr_aarch64_aapcs, implicit-def dead $lr, implicit $sp, implicit-def $sp, implicit-def dead $w0
     early-clobber $sp, $lr = frame-destroy LDRXpost $sp, 16 :: (load (s64) from %stack.0)
     TCRETURNdi @qux, 0, csr_aarch64_aapcs, implicit $sp
 
   bb.2.cold_block (bbsections Cold):
     BL @bar, csr_aarch64_aapcs, implicit-def dead $lr, implicit $sp, implicit-def $sp, implicit-def dead $w0
     early-clobber $sp, $lr = frame-destroy LDRXpost $sp, 16 :: (load (s64) from %stack.0)
     TCRETURNdi @qux, 0, csr_aarch64_aapcs, implicit $sp
body:             |
   ; INDIRECT-NEXT:    B{{ }}
   ; INDIRECT-SAME:           %[[TRAMP1]]
   ; INDIRECT:      bb.1.unrelaxable:
   ; INDIRECT-NEXT:   successors: %bb.2
   ; INDIRECT-SAME:                    , %[[TRAMP2:bb.[0-9]+]]
   ; INDIRECT:        TBNZW
   ; INDIRECT-SAME:         %bb.2
   ; INDIRECT:      [[TRAMP2]]
-  ; INDIRECT-NEXT:   successors: %bb.3
+  ; INDIRECT-NEXT:   successors: %bb.6
   ; INDIRECT:      bb.2.end:
   ; INDIRECT:        TCRETURNdi
   ; INDIRECT:      [[TRAMP1]].entry:
-  ; INDIRECT:        successors: %bb.3
-  ; INDIRECT-NOT:  bbsections Cold
+  ; INDIRECT-NEXT:   successors: %[[TRAMP1_SPILL:bb.[0-9]+]]
+  ; INDIRECT:      [[TRAMP1_SPILL]].entry:
+  ; INDIRECT-NEXT:   successors: %[[TRAMP1_RESTORE:bb.[0-9]+]]
+  ; INDIRECT:        early-clobber $sp = STRXpre $[[SPILL_REGISTER:x[0-9]+]], $sp, -16
+  ; INDIRECT-NEXT:   B %[[TRAMP1_RESTORE:bb.[0-9]+]]
+  ; INDIRECT:      [[TRAMP1_RESTORE]].cold (bbsections Cold):
+  ; INDIRECT-NEXT:   successors: %bb.3
+  ; INDIRECT-NEXT:   {{ $}}
+  ; INDIRECT-NEXT:   early-clobber $sp, $[[SPILL_REGISTER]] = LDRXpost $sp, 16
   ; INDIRECT:      bb.3.cold (bbsections Cold):
   ; INDIRECT:        TCRETURNdi
 
   bb.0.entry (%ir-block.entry):
     successors: %bb.1, %bb.3
     liveins: $w0, $w1, $lr
 
     early-clobber $sp = frame-setup STRXpre killed $lr, $sp, -16 :: (store (s64) into %stack.0)

 
 name:            x16_used_cold_to_hot
 tracksRegLiveness: true
 liveins:         []
 machineFunctionInfo:
   hasRedZone:      false
 body:             |
   ; INDIRECT-LABEL: name: x16_used_cold_to_hot
-  ; COM: Pre-commit to record the behavior of relaxing a "cold-to-hot"
-  ; COM: unconditional branch across which x16 is taken but there is
-  ; COM: still a free register.
+  ; COM: Check that unconditional branches from the cold section to
+  ; COM: the hot section manually insert indirect branches if x16
+  ; COM: isn't available but there is still a free register.
   ; INDIRECT:       bb.0.entry:
   ; INDIRECT-NEXT:    successors: %bb.1
   ; INDIRECT-SAME:                     , %bb.3
   ; INDIRECT:         TBZW killed renamable $w8, 0, %bb.1
   ; INDIRECT-NEXT:    {{ $}}
   ; INDIRECT-NEXT:  bb.3.entry:
-  ; INDIRECT-NEXT:    successors: %bb.2
+  ; INDIRECT-NEXT:    successors: %bb.4
   ; INDIRECT-NEXT:    liveins: $x16
   ; INDIRECT-NEXT:    {{ $}}
-  ; INDIRECT-NEXT:    $[[SCAVENGED_REGISTER3:x[0-9]+]] = ADRP target-flags(aarch64-page) <mcsymbol x16_used_cold_to_hot.cold>
-  ; INDIRECT-NEXT:    $[[SCAVENGED_REGISTER3]] = ADDXri $[[SCAVENGED_REGISTER3]], target-flags(aarch64-pageoff, aarch64-nc) <mcsymbol x16_used_cold_to_hot.cold>, 0
-  ; INDIRECT-NEXT:    BR $[[SCAVENGED_REGISTER3]]
+  ; INDIRECT-NEXT:    early-clobber $sp = STRXpre $[[SPILL_REGISTER]], $sp, -16
+  ; INDIRECT-NEXT:    B %bb.4
   ; INDIRECT:       bb.1.hot:
   ; INDIRECT-NEXT:    liveins: $x16
   ; INDIRECT:           killed $x16
   ; INDIRECT:         RET undef $lr
-  ; INDIRECT:       bb.2.cold (bbsections Cold):
+  ; INDIRECT:       bb.4.cold (bbsections Cold):
+  ; INDIRECT-NEXT:    successors: %bb.2
+  ; INDIRECT-NEXT:    {{ $}}
+  ; INDIRECT-NEXT:    early-clobber $sp, $[[SPILL_REGISTER]] = LDRXpost $sp, 16
+  ; INDIRECT-NEXT:    {{ $}}
+  ; INDIRECT-NEXT:  bb.2.cold (bbsections Cold):
   ; INDIRECT-NEXT:    successors: %bb.5
   ; INDIRECT-NEXT:    liveins: $x16
   ; INDIRECT-NEXT:    {{ $}}
   ; INDIRECT-NEXT:    INLINEASM &".space 4", 1 /* sideeffect attdialect */
   ; INDIRECT-NEXT:    {{ $}}
   ; INDIRECT-NEXT:  bb.5.cold (bbsections Cold):
   ; INDIRECT-NEXT:    successors: %bb.1
   ; INDIRECT-NEXT:    liveins: $x16
   ; INDIRECT-NEXT:    {{ $}}
-  ; INDIRECT-NEXT:    $[[SCAVENGED_REGISTER4:x[0-9]+]] = ADRP target-flags(aarch64-page) <mcsymbol .LBB5_1>
-  ; INDIRECT-NEXT:    $[[SCAVENGED_REGISTER4]] = ADDXri $[[SCAVENGED_REGISTER4]], target-flags(aarch64-pageoff, aarch64-nc) <mcsymbol .LBB5_1>, 0
-  ; INDIRECT-NEXT:    BR $[[SCAVENGED_REGISTER4]]
+  ; INDIRECT-NEXT:    $[[SCAVENGED_REGISTER:x[0-9]+]] = ADRP target-flags(aarch64-page) <mcsymbol .LBB5_1>
+  ; INDIRECT-NEXT:    $[[SCAVENGED_REGISTER]] = ADDXri $[[SCAVENGED_REGISTER]], target-flags(aarch64-pageoff, aarch64-nc) <mcsymbol .LBB5_1>, 0
+  ; INDIRECT-NEXT:    BR $[[SCAVENGED_REGISTER]]
 
   bb.0.entry:
     successors: %bb.1, %bb.2
 
     $sp = frame-setup SUBXri $sp, 16, 0
     INLINEASM &"mov x16, 1", 1 /* sideeffect attdialect */, 10 /* regdef */, implicit-def $x16
     dead renamable $x8 = SUBSXri $x16, 0, 0, implicit-def $nzcv
     renamable $w8 = CSINCWr $wzr, $wzr, 1, implicit killed $nzcv
- { id: 9, name: '', type: spill-slot, offset: -80, size: 8, alignment: 8,
       debug-info-variable: '', debug-info-expression: '', debug-info-location: '' }
   - { id: 10, name: '', type: spill-slot, offset: -96, size: 8, alignment: 16, 
       stack-id: default, callee-saved-register: '$fp', callee-saved-restored: true, 
       debug-info-variable: '', debug-info-expression: '', debug-info-location: '' }
 machineFunctionInfo:
   hasRedZone:      false
 body:             |
   ; INDIRECT-LABEL: name: all_used_cold_to_hot
-  ; COM: Pre-commit to record the behavior of relaxing a "cold-to-hot"
-  ; COM: unconditional branch across which there are no free registers.
+  ; COM: Check that unconditional branches from the cold section to
+  ; COM: the hot section spill x16 and defer indirect branch
+  ; COM: insertion to the linker if there are no free general-purpose
+  ; COM: registers.
   ; INDIRECT:      bb.0.entry:
   ; INDIRECT-NEXT:     successors: %bb.3
   ; INDIRECT-NEXT:     liveins: $fp, $x27, $x28, $x25, $x26, $x23, $x24, $x21, $x22, $x19, $x20
   ; INDIRECT-COUNT-29: INLINEASM &"mov
   ; INDIRECT-NEXT:     {{ $}}
   ; INDIRECT:      bb.3.entry:
   ; INDIRECT-NEXT:     successors: %bb.2
   ; INDIRECT-NEXT:     liveins: $fp,  $x0,  $x1,  $x2,  $x3,  $x4,  $x5,  $x6,  $x7,  $x8,
   ; INDIRECT-SAME:              $x9,  $x10, $x11, $x12, $x13, $x14, $x15, $x17, $x18, $x19,
   ; INDIRECT-SAME:              $x20, $x21, $x22, $x23, $x24, $x25, $x26, $x27, $x28
   ; INDIRECT-NEXT:     {{ $}}
-  ; INDIRECT-NEXT:    $[[SCAVENGED_REGISTER5:x[0-9]+]] = ADRP target-flags(aarch64-page) <mcsymbol all_used_cold_to_hot.cold>
-  ; INDIRECT-NEXT:    $[[SCAVENGED_REGISTER5]] = ADDXri $[[SCAVENGED_REGISTER5]], target-flags(aarch64-pageoff, aarch64-nc) <mcsymbol all_used_cold_to_hot.cold>, 0
-  ; INDIRECT-NEXT:    BR $[[SCAVENGED_REGISTER5]]
-  ; INDIRECT-NEXT:     {{ $}}
-  ; INDIRECT-NEXT:    bb.6.exit:
-  ; INDIRECT-NEXT:    successors: %bb.1
-  ; INDIRECT-NEXT:    liveins: $x0,  $x1,  $x2,  $x3,  $x4,  $x5,  $x6,  $x7,  $x8,  $x9,
-  ; INDIRECT-SAME:             $x10, $x11, $x12, $x13, $x14, $x15, $x17, $x18, $x19, $x20,
-  ; INDIRECT-SAME:             $x21, $x22, $x23, $x24, $x25, $x26, $x27, $x28
-  ; INDIRECT-NEXT:     {{ $}}
-  ; INDIRECT-NEXT:    early-clobber $sp, $[[SPILL_REGISTER:x[0-9]+]] = LDRXpost $sp, 16
+  ; INDIRECT-NEXT:     B %bb.2
   ; INDIRECT-NEXT:     {{ $}}
   ; INDIRECT-NEXT:   bb.1.exit:
   ; INDIRECT-NEXT:     liveins: $x0,  $x1,  $x2,  $x3,  $x4,  $x5,  $x6,  $x7,  $x8,  $x9,
   ; INDIRECT-SAME:              $x10, $x11, $x12, $x13, $x14, $x15, $x16, $x17, $x18, $x19,
   ; INDIRECT-SAME:              $x20, $x21, $x22, $x23, $x24, $x25, $x26, $x27, $x28, $fp
   ; INDIRECT-NEXT:     {{ $}}
   ; INDIRECT-COUNT-30: INLINEASM &"# reg use $0", 1 /* sideeffect attdialect */, 9 /* reguse */, killed
   ; INDIRECT:          RET undef $lr
   ; INDIRECT-NEXT:     {{ $}}
+  ; INDIRECT-NEXT:   bb.6.exit:
+  ; INDIRECT-NEXT:     successors: %bb.7(0x80000000)
+  ; INDIRECT-NEXT:     {{ $}}
+  ; INDIRECT-NEXT:     early-clobber $sp, $[[SPILL_REGISTER]] = LDRXpost $sp, 16
+  ; INDIRECT-NEXT:     {{ $}}
+  ; INDIRECT-NEXT:   bb.7.exit:
+  ; INDIRECT-NEXT:     successors: %bb.1(0x80000000)
+  ; INDIRECT-NEXT:     {{ $}}
+  ; INDIRECT-NEXT:     B %bb.1
+  ; INDIRECT-NEXT:     {{ $}}
   ; INDIRECT-NEXT:   bb.2.cold (bbsections Cold):
   ; INDIRECT-NEXT:     successors: %bb.5
   ; INDIRECT-NEXT:     liveins: $x0,  $x1,  $x2,  $x3,  $x4,  $x5,  $x6,  $x7,  $x8,  $x9,
   ; INDIRECT-SAME:              $x10, $x11, $x12, $x13, $x14, $x15, $x17, $x18, $x19, $x20,
   ; INDIRECT-SAME:              $x21, $x22, $x23, $x24, $x25, $x26, $x27, $x28, $fp
   ; INDIRECT-NEXT:     {{ $}}
   ; INDIRECT-NEXT:     INLINEASM &"mov x16, 1", 1 /* sideeffect attdialect */, 10 /* regdef */, implicit-def $x16
   ; INDIRECT-NEXT:     {{ $}}
   ; INDIRECT-NEXT:   bb.5.cold (bbsections Cold):
   ; INDIRECT-NEXT:     successors: %bb.6
   ; INDIRECT-NEXT:     liveins: $fp,  $x0,  $x1,  $x2,  $x3,  $x4,  $x5,  $x6,  $x7,  $x8,
   ; INDIRECT-SAME:              $x9,  $x10, $x11, $x12, $x13, $x14, $x15, $x16, $x17, $x18,
   ; INDIRECT-SAME:              $x19, $x20, $x21, $x22, $x23, $x24, $x25, $x26, $x27, $x28
   ; INDIRECT-NEXT:     {{ $}}
   ; INDIRECT-NEXT:     early-clobber $sp = STRXpre $[[SPILL_REGISTER]], $sp, -16
-  ; INDIRECT-NEXT:     $[[SPILL_REGISTER]] = ADRP target-flags(aarch64-page) <mcsymbol .LBB6_6>
-  ; INDIRECT-NEXT:     $[[SPILL_REGISTER]] = ADDXri $[[SPILL_REGISTER]], target-flags(aarch64-pageoff, aarch64-nc) <mcsymbol .LBB6_6>, 0
-  ; INDIRECT-NEXT:     BR $[[SPILL_REGISTER]]
+  ; INDIRECT-NEXT:     B %bb.6
 
   bb.0.entry:
     successors: %bb.2
     liveins: $fp, $x27, $x28, $x25, $x26, $x23, $x24, $x21, $x22, $x19, $x20
 
     $sp = frame-setup SUBXri $sp, 112, 0
     frame-setup STRXui killed $fp, $sp, 2 :: (store (s64) into %stack.10)
     frame-setup STPXi killed $x28, killed $x27, $sp, 4 :: (store (s64) into %stack.9), (store (s64) into %stack.8)