[llvm] [BOLT] Relax conditional tail calls in non-simple functions (PR #214400)
YongKang Zhu via llvm-commits
llvm-commits at lists.llvm.org
Fri Aug 7 17:52:56 PDT 2026
https://github.com/yozhu updated https://github.com/llvm/llvm-project/pull/214400
>From 8e0c96847e03718a197410a334415aa01005c979 Mon Sep 17 00:00:00 2001
From: YongKang Zhu <yongzhu at fb.com>
Date: Wed, 5 Aug 2026 21:30:56 -0700
Subject: [PATCH 1/2] [BOLT] Relax conditional tail calls in non-simple
functions
In compact code model relaxLocalBranches() skipped non-simple functions
that may have conditional tail calls targeting symbols in other far away
functions, which could fail during JITLink because the targets are out
of range.
Fix this by having relaxLocalBranches() process non-simple functions,
for which trampolines are added at the end of function body using
unconditional branches (128MB range) to reach far away targets.
---
bolt/lib/Passes/LongJmp.cpp | 150 +++++++++++++-----
.../AArch64/compact-code-model-nonsimple.s | 87 ++++++++++
2 files changed, 196 insertions(+), 41 deletions(-)
create mode 100644 bolt/test/AArch64/compact-code-model-nonsimple.s
diff --git a/bolt/lib/Passes/LongJmp.cpp b/bolt/lib/Passes/LongJmp.cpp
index 289e4877b2bdd..681c47d540f41 100644
--- a/bolt/lib/Passes/LongJmp.cpp
+++ b/bolt/lib/Passes/LongJmp.cpp
@@ -665,13 +665,24 @@ Error LongJmpPass::relax(BinaryFunction &Func, bool &Modified) {
return Error::success();
}
+// Set when relaxLocalBranches() fails to relax a branch. We don't exit directly
+// from relaxLocalBranches() since it runs on a thread pool, and exiting from a
+// worker thread would run the ThreadPool destructor on that same thread. So
+// we're finishing all parallel jobs and checking the flag after it.
+static std::atomic<bool> PassFailed{false};
+
void LongJmpPass::relaxLocalBranches(BinaryFunction &BF,
const BranchLivenessInfo *BLI) {
+ if (PassFailed)
+ return;
+
BinaryContext &BC = BF.getBinaryContext();
auto &MIB = BC.MIB;
- // Quick path.
- if (!BF.isSplit() && BF.estimateSize() < ShortestJumpSpan)
+ // Quick path. Only valid for simple functions, where all branch targets are
+ // basic blocks of the function itself. A non-simple function may branch to a
+ // symbol outside of it that ends up out of range.
+ if (BF.isSimple() && !BF.isSplit() && BF.estimateSize() < ShortestJumpSpan)
return;
auto isBranchOffsetInRange = [&](const MCInst &Inst, int64_t Offset) {
@@ -712,11 +723,14 @@ void LongJmpPass::relaxLocalBranches(BinaryFunction &BF,
DenseMap<const BinaryBasicBlock *, BinaryBasicBlock *> FragmentTrampolines;
// Create a trampoline code after \p BB or at the end of the fragment if BB
- // is nullptr. \p Offset reflects the size delta of BB caused by splitting
- // unconditional branches, or replacing a branch with a longer instruction
- // sequence. It is used to update the output addresses of basic blocks
- // following the trampoline.
+ // is nullptr. The trampoline branches to \p TargetSym. If \p TargetBB is
+ // set, it is added as a successor and registered in FragmentTrampolines.
+ // \p Offset reflects the size delta of BB caused by splitting unconditional
+ // branches, or replacing a branch with a longer instruction sequence. It is
+ // used to update the output addresses of basic blocks following the
+ // trampoline.
auto addTrampolineAfter = [&](BinaryBasicBlock *BB,
+ const MCSymbol *TargetSym,
BinaryBasicBlock *TargetBB, uint64_t Count,
uint64_t Offset = 0) {
FunctionTrampolines.emplace_back(BB ? BB : FF.back(),
@@ -730,10 +744,11 @@ void LongJmpPass::relaxLocalBranches(BinaryFunction &BF,
MCInst Inst;
{
auto L = BC.scopeLock();
- MIB->createUncondBranch(Inst, TargetBB->getLabel(), BC.Ctx.get());
+ MIB->createUncondBranch(Inst, TargetSym, BC.Ctx.get());
}
TrampolineBB->addInstruction(Inst);
- TrampolineBB->addSuccessor(TargetBB, Count);
+ if (TargetBB)
+ TrampolineBB->addSuccessor(TargetBB, Count);
TrampolineBB->setExecutionCount(Count);
const uint64_t TrampolineAddress =
BB ? BB->getOutputEndAddress() : FragmentSize;
@@ -751,7 +766,7 @@ void LongJmpPass::relaxLocalBranches(BinaryFunction &BF,
BB->setOutputEndAddress(BB->getOutputEndAddress() + Offset);
};
- if (!FragmentTrampolines.lookup(TargetBB))
+ if (TargetBB && !FragmentTrampolines.lookup(TargetBB))
FragmentTrampolines[TargetBB] = TrampolineBB;
if (!Offset)
@@ -785,22 +800,26 @@ void LongJmpPass::relaxLocalBranches(BinaryFunction &BF,
};
// Pre-populate trampolines by splitting unconditional branches from the
- // containing basic block.
- for (BinaryBasicBlock *BB : FF) {
- MCInst *Inst = BB->getLastNonPseudoInstr();
- if (!Inst || !MIB->isUnconditionalBranch(*Inst))
- continue;
+ // containing basic block. Skip for non-simple functions because inserting
+ // blocks after existing BBs would shift code and break jump table offsets.
+ if (BF.isSimple()) {
+ for (BinaryBasicBlock *BB : FF) {
+ MCInst *Inst = BB->getLastNonPseudoInstr();
+ if (!Inst || !MIB->isUnconditionalBranch(*Inst))
+ continue;
- const MCSymbol *TargetSymbol = MIB->getTargetSymbol(*Inst);
- BB->eraseInstruction(BB->findInstruction(Inst));
+ const MCSymbol *TargetSymbol = MIB->getTargetSymbol(*Inst);
+ BB->eraseInstruction(BB->findInstruction(Inst));
- BinaryBasicBlock::BinaryBranchInfo BI;
- BinaryBasicBlock *TargetBB = BB->getSuccessor(TargetSymbol, BI);
+ BinaryBasicBlock::BinaryBranchInfo BI;
+ BinaryBasicBlock *TargetBB = BB->getSuccessor(TargetSymbol, BI);
- // Erasing the unconditional branch shrinks BB by one instruction.
- BinaryBasicBlock *TrampolineBB =
- addTrampolineAfter(BB, TargetBB, BI.Count, /*Offset=*/-4);
- BB->replaceSuccessor(TargetBB, TrampolineBB, BI.Count);
+ // Erasing the unconditional branch shrinks BB by one instruction.
+ BinaryBasicBlock *TrampolineBB =
+ addTrampolineAfter(BB, TargetBB->getLabel(), TargetBB, BI.Count,
+ /*Offset=*/-4);
+ BB->replaceSuccessor(TargetBB, TrampolineBB, BI.Count);
+ }
}
/// Relax the branch \p Inst in basic block \p BB that targets \p TargetBB.
@@ -832,7 +851,8 @@ void LongJmpPass::relaxLocalBranches(BinaryFunction &BF,
// case we will need further relaxation.
const int64_t OffsetToEnd = FragmentSize - InstAddress;
if (Count == 0 && isBranchOffsetInRange(Inst, OffsetToEnd)) {
- TrampolineBB = addTrampolineAfter(nullptr, TargetBB, Count);
+ TrampolineBB =
+ addTrampolineAfter(nullptr, TargetBB->getLabel(), TargetBB, Count);
BB->replaceSuccessor(TargetBB, TrampolineBB, Count);
auto L = BC.scopeLock();
MIB->replaceBranchTarget(Inst, TrampolineBB->getLabel(), BC.Ctx.get());
@@ -852,7 +872,7 @@ void LongJmpPass::relaxLocalBranches(BinaryFunction &BF,
if (ShouldReverseBranch && !IsReversibleBranch) {
const uint64_t NextCount = BB->getBranchInfo(*NextBB).Count;
BinaryBasicBlock *FallThrough =
- addTrampolineAfter(BB, NextBB, NextCount);
+ addTrampolineAfter(BB, NextBB->getLabel(), NextBB, NextCount);
BB->replaceSuccessor(NextBB, FallThrough, NextCount);
}
@@ -870,17 +890,22 @@ void LongJmpPass::relaxLocalBranches(BinaryFunction &BF,
const uint64_t NewBBSize = BB->estimateSize();
// Create a trampoline basic block for the original taken target.
- TrampolineBB =
- addTrampolineAfter(BB, TargetBB, Count, NewBBSize - OldBBSize);
+ TrampolineBB = addTrampolineAfter(BB, TargetBB->getLabel(), TargetBB,
+ Count, NewBBSize - OldBBSize);
} else {
// Create a trampoline basic block for the taken target of the branch.
- TrampolineBB = addTrampolineAfter(BB, TargetBB, Count);
+ TrampolineBB =
+ addTrampolineAfter(BB, TargetBB->getLabel(), TargetBB, Count);
auto L = BC.scopeLock();
MIB->replaceBranchTarget(Inst, TrampolineBB->getLabel(), BC.Ctx.get());
}
BB->replaceSuccessor(TargetBB, TrampolineBB, Count);
};
+ // For non-simple functions, branch targets may be different functions,
+ // so we track trampolines by symbol rather than by basic block.
+ DenseMap<const MCSymbol *, BinaryBasicBlock *> SymbolTrampolines;
+
bool MayNeedRelaxation;
uint64_t NumIterations = 0;
do {
@@ -907,18 +932,59 @@ void LongJmpPass::relaxLocalBranches(BinaryFunction &BF,
continue;
const MCSymbol *TargetSymbol = MIB->getTargetSymbol(Inst);
- BinaryBasicBlock *TargetBB = BB->getSuccessor(TargetSymbol);
- assert(TargetBB &&
- "Basic block target expected for conditional branch.");
- // Check if the relaxation is needed.
- if (TargetBB->getFragmentNum() == FF.getFragmentNum() &&
- isBlockInRange(Inst, InstAddress, *TargetBB))
- continue;
-
- relaxBranch(BB, Inst, InstAddress, TargetBB);
-
- MayNeedRelaxation = true;
+ if (BF.isSimple()) {
+ BinaryBasicBlock *TargetBB = BB->getSuccessor(TargetSymbol);
+ assert(TargetBB &&
+ "Basic block target expected for conditional branch.");
+
+ // Check if the relaxation is needed.
+ if (TargetBB->getFragmentNum() == FF.getFragmentNum() &&
+ isBlockInRange(Inst, InstAddress, *TargetBB))
+ continue;
+
+ relaxBranch(BB, Inst, InstAddress, TargetBB);
+ MayNeedRelaxation = true;
+ } else {
+ // Skip if the target is within this function.
+ if (BF.getBasicBlockForLabel(TargetSymbol))
+ continue;
+
+ // Try to reuse an existing trampoline for this symbol.
+ BinaryBasicBlock *TrampolineBB =
+ SymbolTrampolines.lookup(TargetSymbol);
+ if (TrampolineBB &&
+ isBlockInRange(Inst, InstAddress, *TrampolineBB)) {
+ auto L = BC.scopeLock();
+ MIB->replaceBranchTarget(Inst, TrampolineBB->getLabel(),
+ BC.Ctx.get());
+ continue;
+ }
+
+ // Create a trampoline at the end of the function. Since the layout
+ // of a non-simple function has to be preserved, the end of the
+ // function is the only place where we can put it.
+ const int64_t OffsetToEnd = FragmentSize - InstAddress;
+ if (!isBranchOffsetInRange(Inst, OffsetToEnd)) {
+ auto L = BC.scopeLock();
+ BC.errs() << "BOLT-ERROR: cannot relax branch in non-simple "
+ "function "
+ << BF << ": a trampoline at the end of the function is "
+ << OffsetToEnd << " bytes away, out of reach for a "
+ << BitsAvailable << "-bit branch\n";
+ BC.printInstruction(BC.errs(), Inst);
+ PassFailed = true;
+ return;
+ }
+
+ TrampolineBB = addTrampolineAfter(/*BB=*/nullptr, TargetSymbol,
+ /*TargetBB=*/nullptr,
+ /*Count=*/0);
+ SymbolTrampolines[TargetSymbol] = TrampolineBB;
+ auto L = BC.scopeLock();
+ MIB->replaceBranchTarget(Inst, TrampolineBB->getLabel(),
+ BC.Ctx.get());
+ }
}
}
@@ -989,14 +1055,16 @@ Error LongJmpPass::runOnFunctions(BinaryContext &BC) {
};
ParallelUtilities::PredicateTy SkipPredicate =
- [&](const BinaryFunction &BF) {
- return !BC.shouldEmit(BF) || !BF.isSimple();
- };
+ [&](const BinaryFunction &BF) { return !BC.shouldEmit(BF); };
ParallelUtilities::runOnEachFunction(
BC, ParallelUtilities::SchedulingPolicy::SP_INST_LINEAR, WorkFun,
SkipPredicate, "RelaxLocalBranches");
+ // The error has already been reported by relaxLocalBranches().
+ if (PassFailed)
+ return createFatalBOLTError("");
+
return Error::success();
}
diff --git a/bolt/test/AArch64/compact-code-model-nonsimple.s b/bolt/test/AArch64/compact-code-model-nonsimple.s
new file mode 100644
index 0000000000000..eff53f8c0f307
--- /dev/null
+++ b/bolt/test/AArch64/compact-code-model-nonsimple.s
@@ -0,0 +1,87 @@
+## Check that llvm-bolt relaxes conditional tail calls in non-simple functions
+## for compact code model. Without the relaxation, the branches below are out
+## of range after reordering and llvm-bolt fails with JITLink error.
+
+# REQUIRES: system-linux
+
+# RUN: llvm-mc -filetype=obj -triple aarch64-unknown-unknown %s -o %t.o
+# RUN: llvm-strip --strip-unneeded %t.o
+# RUN: %clang %cflags %t.o -o %t.exe -Wl,-q -static
+# RUN: echo nonsimple > %t.order
+# RUN: echo large_function >> %t.order
+# RUN: echo _start >> %t.order
+# RUN: llvm-bolt %t.exe -o %t.bolt --compact-code-model --keep-nops \
+# RUN: --function-order=%t.order --print-cfg --print-only=nonsimple \
+# RUN: | FileCheck %s --check-prefix=CHECK-CFG
+# RUN: llvm-objdump -d --disassemble-symbols=nonsimple %t.bolt | FileCheck %s
+
+ .text
+ .globl _start
+ .type _start, %function
+_start:
+ .cfi_startproc
+ bl nonsimple
+ ret x30
+ .cfi_endproc
+.size _start, .-_start
+
+## 64KB of code placed between "nonsimple" and "cold_target" by the order file,
+## which puts "cold_target" beyond the +-32KB reach of the tbz below.
+ .globl large_function
+ .type large_function, %function
+large_function:
+ .cfi_startproc
+ .rept 16000
+ nop
+ .endr
+ ret x30
+ .cfi_endproc
+.size large_function, .-large_function
+
+## Non-simple function ("br x16" has unknown control flow) with two conditional
+## tail calls to the same target.
+ .globl nonsimple
+ .type nonsimple, %function
+nonsimple:
+ .cfi_startproc
+ cmp x0, #1
+ b.eq cold_target
+ tbz x0, #0, cold_target
+ ldr x16, [sp]
+ br x16
+ .cfi_endproc
+.size nonsimple, .-nonsimple
+
+ .globl cold_target
+ .type cold_target, %function
+cold_target:
+ .cfi_startproc
+ mov x0, #1
+ ret x30
+ .cfi_endproc
+.size cold_target, .-cold_target
+
+## Force relocation mode.
+ .reloc 0, R_AARCH64_NONE
+
+## Verify that the function under test is really non-simple. If this ever
+## starts printing "IsSimple : 1", the test no longer covers the non-simple
+## relaxation path.
+# CHECK-CFG: Binary Function "nonsimple" after building cfg
+# CHECK-CFG: IsSimple : 0
+
+## Both branches should be retargeted to a single trampoline appended after the
+## function body, which in turn branches to the original target. The layout of
+## the function itself must be preserved.
+##
+## Note that the tbz reuses the trampoline created for the b.eq, as both target
+## the same symbol: the address captured from the b.eq must match the one used
+## by the tbz, and no second trampoline may be emitted.
+# CHECK: <nonsimple>:
+# CHECK: cmp x0, #0x1
+# CHECK-NEXT: b.eq 0x[[TRAMP:[0-9a-f]+]] <nonsimple+0x{{[0-9a-f]+}}>
+# CHECK-NEXT: tbz {{.*}}, 0x[[TRAMP]] <nonsimple+0x{{[0-9a-f]+}}>
+# CHECK-NEXT: ldr x16, [sp]
+# CHECK-NEXT: br x16
+# CHECK-NEXT: [[TRAMP]]: {{.*}} b 0x{{[0-9a-f]+}} <cold_target>
+# CHECK-NOT: b 0x{{[0-9a-f]+}} <cold_target>
>From 0f45472eff11bc0a569010499a4da9f20f19c754 Mon Sep 17 00:00:00 2001
From: YongKang Zhu <yongzhu at fb.com>
Date: Fri, 7 Aug 2026 17:52:07 -0700
Subject: [PATCH 2/2] address review feedback
---
bolt/include/bolt/Passes/LongJmp.h | 5 +++--
bolt/lib/Passes/LongJmp.cpp | 33 ++++++++++++++----------------
2 files changed, 18 insertions(+), 20 deletions(-)
diff --git a/bolt/include/bolt/Passes/LongJmp.h b/bolt/include/bolt/Passes/LongJmp.h
index 0a7c4d5f33334..4a288d3bde4e3 100644
--- a/bolt/include/bolt/Passes/LongJmp.h
+++ b/bolt/include/bolt/Passes/LongJmp.h
@@ -75,8 +75,9 @@ class LongJmpPass : public BinaryFunctionPass {
/// Relax all internal function branches including those between fragments.
/// Assume that fragments are placed in different sections but are within
- /// 128MB of each other.
- void relaxLocalBranches(BinaryFunction &BF,
+ /// 128MB of each other. Return false and report an error if a branch cannot
+ /// be relaxed.
+ bool relaxLocalBranches(BinaryFunction &BF,
const BranchLivenessInfo *BLI = nullptr);
/// -- Layout estimation methods --
diff --git a/bolt/lib/Passes/LongJmp.cpp b/bolt/lib/Passes/LongJmp.cpp
index 681c47d540f41..38ad4ed52f339 100644
--- a/bolt/lib/Passes/LongJmp.cpp
+++ b/bolt/lib/Passes/LongJmp.cpp
@@ -665,17 +665,8 @@ Error LongJmpPass::relax(BinaryFunction &Func, bool &Modified) {
return Error::success();
}
-// Set when relaxLocalBranches() fails to relax a branch. We don't exit directly
-// from relaxLocalBranches() since it runs on a thread pool, and exiting from a
-// worker thread would run the ThreadPool destructor on that same thread. So
-// we're finishing all parallel jobs and checking the flag after it.
-static std::atomic<bool> PassFailed{false};
-
-void LongJmpPass::relaxLocalBranches(BinaryFunction &BF,
+bool LongJmpPass::relaxLocalBranches(BinaryFunction &BF,
const BranchLivenessInfo *BLI) {
- if (PassFailed)
- return;
-
BinaryContext &BC = BF.getBinaryContext();
auto &MIB = BC.MIB;
@@ -683,7 +674,7 @@ void LongJmpPass::relaxLocalBranches(BinaryFunction &BF,
// basic blocks of the function itself. A non-simple function may branch to a
// symbol outside of it that ends up out of range.
if (BF.isSimple() && !BF.isSplit() && BF.estimateSize() < ShortestJumpSpan)
- return;
+ return true;
auto isBranchOffsetInRange = [&](const MCInst &Inst, int64_t Offset) {
const unsigned Bits = MIB->getPCRelEncodingSize(Inst);
@@ -800,8 +791,9 @@ void LongJmpPass::relaxLocalBranches(BinaryFunction &BF,
};
// Pre-populate trampolines by splitting unconditional branches from the
- // containing basic block. Skip for non-simple functions because inserting
- // blocks after existing BBs would shift code and break jump table offsets.
+ // containing basic block. Skip for non-simple functions: this creates
+ // trampolines for targets inside the function, while in a non-simple
+ // function we only relax branches to targets outside of it.
if (BF.isSimple()) {
for (BinaryBasicBlock *BB : FF) {
MCInst *Inst = BB->getLastNonPseudoInstr();
@@ -973,8 +965,7 @@ void LongJmpPass::relaxLocalBranches(BinaryFunction &BF,
<< OffsetToEnd << " bytes away, out of reach for a "
<< BitsAvailable << "-bit branch\n";
BC.printInstruction(BC.errs(), Inst);
- PassFailed = true;
- return;
+ return false;
}
TrampolineBB = addTrampolineAfter(/*BB=*/nullptr, TargetSymbol,
@@ -1019,6 +1010,8 @@ void LongJmpPass::relaxLocalBranches(BinaryFunction &BF,
/*UpdateLayout*/ true, /*UpdateCFI*/ true,
/*RecomputeLPs*/ false);
}
+
+ return true;
}
Error LongJmpPass::runOnFunctions(BinaryContext &BC) {
@@ -1050,8 +1043,12 @@ Error LongJmpPass::runOnFunctions(BinaryContext &BC) {
BC.outs()
<< "BOLT-INFO: relaxing branches for compact code model (<128MB)\n";
+ std::atomic<bool> HasFatal{false};
ParallelUtilities::WorkFuncTy WorkFun = [&](BinaryFunction &BF) {
- relaxLocalBranches(BF, getBranchLiveness(BF));
+ if (HasFatal)
+ return;
+ if (!relaxLocalBranches(BF, getBranchLiveness(BF)))
+ HasFatal = true;
};
ParallelUtilities::PredicateTy SkipPredicate =
@@ -1062,8 +1059,8 @@ Error LongJmpPass::runOnFunctions(BinaryContext &BC) {
SkipPredicate, "RelaxLocalBranches");
// The error has already been reported by relaxLocalBranches().
- if (PassFailed)
- return createFatalBOLTError("");
+ if (HasFatal)
+ return createFatalBOLTError("branch relaxation failure");
return Error::success();
}
More information about the llvm-commits
mailing list