[llvm] [BOLT][AArch64] Account for hugify alignment in AArch64 long jump layout (PR #195272)

[Alexandros Lamprineas via llvm-commits]

https://github.com/labrinea updated https://github.com/llvm/llvm-project/pull/195272

>From 4428f40cf6e9045a19b9d1115ff73de849d35a63 Mon Sep 17 00:00:00 2001
From: Alexandros Lamprineas <alexandros.lamprineas at arm.com>
Date: Fri, 1 May 2026 14:36:37 +0100
Subject: [PATCH 1/2] [BOLT][AArch64] Account for hugify alignment in AArch64
 long jump layout.

When --hugify is used for a PIE, the final section allocation in
RewriteInstance::mapCodeSections aligns the address after the last
non-cold text section before laying out the following sections:

  for (BinarySection *Section : CodeSections) {
    Address = alignTo(Address, Section->getAlignment());
    Section->setOutputAddress(Address);
    Address += Section->getOutputSize();

    if (opts::Hugify && !BC->HasFixedLoadAddress &&
        Section->getName() == LastNonColdSectionName)
      Address = alignTo(Address, Section->getAlignment());
  }

The AArch64 long-jump pass doesn't model that gap in its tentative
layout, so a CBZ could be considered in range during stub insertion
and later become out of range when JITLink applied the final layout.

This patch mirrors the hugify alignment before assigning cold fragment
addresses so that range checks see the same hot-to-cold distance as
the final layout.

Assisted-by: Codex
---
 bolt/lib/Passes/LongJmp.cpp                   |  5 ++++
 .../long-jmp-hugify-fixup-out-of-range.s      | 27 ++++++++++++++-----
 2 files changed, 26 insertions(+), 6 deletions(-)

diff --git a/bolt/lib/Passes/LongJmp.cpp b/bolt/lib/Passes/LongJmp.cpp
index 7744cf08defa9..f085500fccbd3 100644
--- a/bolt/lib/Passes/LongJmp.cpp
+++ b/bolt/lib/Passes/LongJmp.cpp
@@ -367,6 +367,11 @@ LongJmpPass::tentativeLayoutRelocMode(const BinaryContext &BC,
   CurrentIndex = 0;
   bool ColdLayoutDone = false;
   auto runColdLayout = [&]() {
+    // Mirror the extra hugify alignment inserted by final section allocation
+    // after the last non-cold section. Account for it before assigning cold
+    // fragment addresses so range checks see the hot-to-cold gap.
+    if (opts::Hugify && !BC.HasFixedLoadAddress && !opts::HotFunctionsAtEnd)
+      DotAddress = alignTo(DotAddress, opts::AlignText);
     DotAddress = tentativeLayoutRelocColdPart(BC, SortedFunctions, DotAddress);
     ColdLayoutDone = true;
     if (opts::HotFunctionsAtEnd)
diff --git a/bolt/test/AArch64/long-jmp-hugify-fixup-out-of-range.s b/bolt/test/AArch64/long-jmp-hugify-fixup-out-of-range.s
index 3763583887bd4..b96950e2e6da0 100644
--- a/bolt/test/AArch64/long-jmp-hugify-fixup-out-of-range.s
+++ b/bolt/test/AArch64/long-jmp-hugify-fixup-out-of-range.s
@@ -1,13 +1,17 @@
-# Check that branches considered in-range during longjump
-# may go out of range at JITLink if hugify moves hot code.
+# The longjump pass may consider branch targets in range during tentative
+# layout and decide not to insert stubs for them. Later, final section
+# allocation may insert alignment padding after the last non-cold text section
+# when hugify is enabled. This moves the following cold section farther away,
+# resulting in relocation fixups going out of range at JITLink. Check that the
+# longjump pass accounts for this padding and inserts stubs when needed.
 
-# REQUIRES: system-linux, asserts
+# REQUIRES: system-linux, asserts, bolt-runtime, target=aarch64{{.*}}
 
 # RUN: %clang %cflags -Wl,-q %s -o %t
 # RUN: link_fdata --no-lbr %s %t %t.fdata
 # RUN: llvm-strip --strip-unneeded %t
-# RUN: not llvm-bolt %t -o %t.bolt --data %t.fdata -split-functions --hugify 2>&1 \
-# RUN:   | FileCheck %s
+# RUN: llvm-bolt %t -o %t.bolt --data %t.fdata -split-functions --hugify
+# RUN: llvm-objdump -d %t.bolt | FileCheck %s
 
   .globl foo
   .type foo, %function
@@ -34,4 +38,15 @@ _start:
 ## Force relocation mode.
 .reloc 0, R_AARCH64_NONE
 
-# CHECK: BOLT-ERROR: JITLink failed: In graph in-memory object file, section .text: relocation target {{0x[0-9a-f]+}} {{.*}} is out of range of CondBranch19PCRel fixup at address {{0x[0-9a-f]+}} {{.*}}
+# CHECK: Disassembly of section .text:
+
+# CHECK: <foo>:
+# CHECK-NEXT:            {{.*}} cbz x0, 0x[[ADDR0:[0-9a-f]+]] <{{.*}}>
+# CHECK-NEXT:            {{.*}} b 0x[[ADDR1:[0-9a-f]+]] <{{.*}}>
+# CHECK-NEXT: [[ADDR0]]: {{.*}} b 0x[[ADDR2:[0-9a-f]+]] <{{.*}}>
+
+# CHECK: Disassembly of section .text.cold:
+
+# CHECK: <foo.cold.0>:
+# CHECK-NEXT: [[ADDR1]]: {{.*}} mov x0, #0x1 // =1
+# CHECK-NEXT: [[ADDR2]]: {{.*}} ret

>From a6f841363415c3417539b8bd7d7a4ef76a86730c Mon Sep 17 00:00:00 2001
From: Alexandros Lamprineas <alexandros.lamprineas at arm.com>
Date: Fri, 1 May 2026 16:07:44 +0100
Subject: [PATCH 2/2] reverse branch and relevant check lines

---
 bolt/test/AArch64/long-jmp-hugify-fixup-out-of-range.s | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/bolt/test/AArch64/long-jmp-hugify-fixup-out-of-range.s b/bolt/test/AArch64/long-jmp-hugify-fixup-out-of-range.s
index b96950e2e6da0..03c35e962e99e 100644
--- a/bolt/test/AArch64/long-jmp-hugify-fixup-out-of-range.s
+++ b/bolt/test/AArch64/long-jmp-hugify-fixup-out-of-range.s
@@ -41,12 +41,12 @@ _start:
 # CHECK: Disassembly of section .text:
 
 # CHECK: <foo>:
-# CHECK-NEXT:            {{.*}} cbz x0, 0x[[ADDR0:[0-9a-f]+]] <{{.*}}>
+# CHECK-NEXT:            {{.*}} cbnz x0, 0x[[ADDR0:[0-9a-f]+]] <{{.*}}>
 # CHECK-NEXT:            {{.*}} b 0x[[ADDR1:[0-9a-f]+]] <{{.*}}>
 # CHECK-NEXT: [[ADDR0]]: {{.*}} b 0x[[ADDR2:[0-9a-f]+]] <{{.*}}>
 
 # CHECK: Disassembly of section .text.cold:
 
 # CHECK: <foo.cold.0>:
-# CHECK-NEXT: [[ADDR1]]: {{.*}} mov x0, #0x1 // =1
-# CHECK-NEXT: [[ADDR2]]: {{.*}} ret
+# CHECK-NEXT: [[ADDR2]]: {{.*}} mov x0, #0x1 // =1
+# CHECK-NEXT: [[ADDR1]]: {{.*}} ret