[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