[llvm] [AArch64][CodeGen] Fix trampoline basic block offset (PR #202716)
via llvm-commits
llvm-commits at lists.llvm.org
Fri Jun 12 06:27:09 PDT 2026
https://github.com/eleviant updated https://github.com/llvm/llvm-project/pull/202716
>From 8aae23350db518e2806b627d89ffe3639c3fd675 Mon Sep 17 00:00:00 2001
From: Evgeny Leviant <eleviant at accesssoftek.com>
Date: Tue, 9 Jun 2026 19:24:42 +0200
Subject: [PATCH 1/2] [AArch64][CodeGen] Fix trampoline basic block offset
Trampoline basic blocks are initially created with an offset of zero.
As a result, `isBlockInRange()` may operate on incorrect block offset,
potentially causing unnecessary conditional branch inversions or the
insertion of redundant trampolines.
---
llvm/lib/CodeGen/BranchRelaxation.cpp | 4 +-
.../AArch64/branch-relax-cross-section.mir | 3 +-
.../test/CodeGen/AArch64/branch-relax-tbz.mir | 80 +++++++++++++++++++
3 files changed, 84 insertions(+), 3 deletions(-)
create mode 100644 llvm/test/CodeGen/AArch64/branch-relax-tbz.mir
diff --git a/llvm/lib/CodeGen/BranchRelaxation.cpp b/llvm/lib/CodeGen/BranchRelaxation.cpp
index 657588cb48b8e..d458fcae57847 100644
--- a/llvm/lib/CodeGen/BranchRelaxation.cpp
+++ b/llvm/lib/CodeGen/BranchRelaxation.cpp
@@ -280,7 +280,9 @@ BranchRelaxation::createNewBlockAfter(MachineBasicBlock &OrigMBB,
OrigMBB.setIsEndSection(false);
// Insert an entry into BlockInfo to align it properly with the block numbers.
- BlockInfo.insert(BlockInfo.begin() + NewBB->getNumber(), BasicBlockInfo());
+ auto It = BlockInfo.insert(BlockInfo.begin() + NewBB->getNumber(),
+ BasicBlockInfo());
+ It->Offset = BlockInfo[OrigMBB.getNumber()].postOffset(*NewBB);
return NewBB;
}
diff --git a/llvm/test/CodeGen/AArch64/branch-relax-cross-section.mir b/llvm/test/CodeGen/AArch64/branch-relax-cross-section.mir
index cc54e494e4e6b..4d116d3b3f70f 100644
--- a/llvm/test/CodeGen/AArch64/branch-relax-cross-section.mir
+++ b/llvm/test/CodeGen/AArch64/branch-relax-cross-section.mir
@@ -303,11 +303,10 @@ body: |
; COM: viable insertion location.
; CHECK: bb.0 (%ir-block.1):
; CHECK-NEXT: successors: %bb.1
- ; CHECK-SAME: , %bb.3
+ ; CHECK-SAME: , %bb.2
; CHECK: CBNZW
; CHECK-SAME: %bb.1
; CHECK-NEXT: B
- ; CHECK-SAME: %bb.3
; CHECK: bb.1.hot:
; CHECK: TCRETURNdi
; CHECK: bb.2.cold (bbsections Cold):
diff --git a/llvm/test/CodeGen/AArch64/branch-relax-tbz.mir b/llvm/test/CodeGen/AArch64/branch-relax-tbz.mir
new file mode 100644
index 0000000000000..f1879824368c9
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/branch-relax-tbz.mir
@@ -0,0 +1,80 @@
+# Specify short tbz offset so we can check that isBlockInRange works correctly
+# for newly created trampolines.
+# RUN: llc -mtriple=aarch64-none-linux-gnu -run-pass branch-relaxation -aarch64-tbz-offset-bits=4 %s -o - | FileCheck %s
+
+# CHECK: body:
+# CHECK: TBZW killed renamable $w0, 0, %[[TBB1:.*]]
+# CHECK: bb.1.hot_block:
+# CHECK: [[TBB1]]
+# CHECK: successors: %[[TBB2:.*]]({{.*}})
+# CHECK: [[TBB2]]
+# CHECK: successors: %[[BB_COLD:.*]]({{.*}})
+# CHECK: B %[[BB_COLD]]
+# CHECK: [[BB_COLD]].cold_block
+
+--- |
+ declare i32 @bar()
+ declare i32 @baz()
+ declare i32 @qux()
+
+ 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
+ }
+
+...
+---
+name: tbz_hot_to_cold
+tracksRegLiveness: true
+liveins:
+ - { reg: '$w0', virtual-reg: '' }
+stack:
+ - { id: 0, name: '', type: spill-slot, offset: -16, size: 8, alignment: 16,
+ stack-id: default, callee-saved-register: '$lr', callee-saved-restored: true,
+ debug-info-variable: '', debug-info-expression: '', debug-info-location: '' }
+body: |
+ 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)
+ HINT 0
+ HINT 0
+ HINT 0
+ HINT 0
+ HINT 0
+ HINT 0
+ HINT 0
+ HINT 0
+ HINT 0
+ HINT 0
+ HINT 0
+ HINT 0
+ HINT 0
+ HINT 0
+ HINT 0
+ HINT 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
+
+...
>From 37b2d7623a62e658acafcf312682c0c9f7ce950e Mon Sep 17 00:00:00 2001
From: Evgeny Leviant <eleviant at accesssoftek.com>
Date: Fri, 12 Jun 2026 15:26:33 +0200
Subject: [PATCH 2/2] Move offset calculation from updateOffsetAndLiveness to
createNewBlockAfter
---
llvm/lib/CodeGen/BranchRelaxation.cpp | 27 +++++++++++++--------------
1 file changed, 13 insertions(+), 14 deletions(-)
diff --git a/llvm/lib/CodeGen/BranchRelaxation.cpp b/llvm/lib/CodeGen/BranchRelaxation.cpp
index d458fcae57847..a4f3fc68d03f0 100644
--- a/llvm/lib/CodeGen/BranchRelaxation.cpp
+++ b/llvm/lib/CodeGen/BranchRelaxation.cpp
@@ -106,6 +106,8 @@ class BranchRelaxation {
MachineBasicBlock *splitBlockBeforeInstr(MachineInstr &MI,
MachineBasicBlock *DestBB);
void adjustBlockOffsets(MachineBasicBlock &Start);
+ // Computes basic block offsets for blocks in the range (Start, End),
+ // i.e. beginning with the block immediately following Start.
void adjustBlockOffsets(MachineBasicBlock &Start,
MachineFunction::iterator End);
bool isBlockInRange(const MachineInstr &MI,
@@ -280,9 +282,12 @@ BranchRelaxation::createNewBlockAfter(MachineBasicBlock &OrigMBB,
OrigMBB.setIsEndSection(false);
// Insert an entry into BlockInfo to align it properly with the block numbers.
- auto It = BlockInfo.insert(BlockInfo.begin() + NewBB->getNumber(),
- BasicBlockInfo());
- It->Offset = BlockInfo[OrigMBB.getNumber()].postOffset(*NewBB);
+ BlockInfo.insert(BlockInfo.begin() + NewBB->getNumber(), BasicBlockInfo());
+
+ // Keep the block offsets approximately up to date. While they will be
+ // slight underestimates, we will update them appropriately in the next
+ // scan through the function.
+ adjustBlockOffsets(OrigMBB, std::next(NewBB->getIterator()));
return NewBB;
}
@@ -405,14 +410,8 @@ bool BranchRelaxation::fixupConditionalBranch(MachineInstr &MI) {
};
// Populate the block offset and live-ins for a new basic block.
- auto updateOffsetAndLiveness = [&](MachineBasicBlock *NewBB) {
- assert(NewBB != nullptr && "can't populate offset for nullptr");
-
- // Keep the block offsets approximately up to date. While they will be
- // slight underestimates, we will update them appropriately in the next
- // scan through the function.
- adjustBlockOffsets(*std::prev(NewBB->getIterator()),
- std::next(NewBB->getIterator()));
+ auto updateLiveness = [&](MachineBasicBlock *NewBB) {
+ assert(NewBB != nullptr && "can't update liveness for nullptr");
// Need to fix live-in lists if we track liveness.
if (TRI->trackLivenessAfterRegAlloc(*MF))
@@ -455,7 +454,7 @@ bool BranchRelaxation::fixupConditionalBranch(MachineInstr &MI) {
insertBranch(MBB, NewBB, FBB, Cond);
TrampolineInsertionPoint = NewBB;
- updateOffsetAndLiveness(NewBB);
+ updateLiveness(NewBB);
return true;
}
@@ -518,7 +517,7 @@ bool BranchRelaxation::fixupConditionalBranch(MachineInstr &MI) {
// Do it here since if there's no split, no update is needed.
MBB->replaceSuccessor(FBB, NewBB);
NewBB->addSuccessor(FBB);
- updateOffsetAndLiveness(NewBB);
+ updateLiveness(NewBB);
}
// We now have an appropriate fall-through block in place (either naturally
@@ -571,7 +570,7 @@ bool BranchRelaxation::fixupConditionalBranch(MachineInstr &MI) {
removeBranch(MBB);
insertBranch(MBB, NewBB, FBB, Cond);
- updateOffsetAndLiveness(NewBB);
+ updateLiveness(NewBB);
return true;
}
More information about the llvm-commits
mailing list