[llvm] [BOLT][AArch64] Don't change layout in PatchEntries (PR #71278)

[Vladislav Khmelevsky via llvm-commits]

https://github.com/yota9 updated https://github.com/llvm/llvm-project/pull/71278

>From 3373eb48cbc5df75da484d4b373177bd2f3d0d58 Mon Sep 17 00:00:00 2001
From: Vladislav Khmelevsky <och95 at yandex.ru>
Date: Sat, 4 Nov 2023 15:00:31 +0400
Subject: [PATCH] [BOLT][AArch64] Don't change layout in PatchEntries

Due to LongJmp pass that is executed before PatchEntries we can't ignore
the function here since it would change pre-calculated output layout.
The test reloc-26 relied on the wrong behavior, rewritten to unittest.
This is also attemp to fix #70771
---
 bolt/lib/Passes/PatchEntries.cpp      | 11 ++++++
 bolt/test/AArch64/reloc-call26.s      | 51 ---------------------------
 bolt/unittests/Core/BinaryContext.cpp | 39 ++++++++++++++++++++
 3 files changed, 50 insertions(+), 51 deletions(-)
 delete mode 100644 bolt/test/AArch64/reloc-call26.s

diff --git a/bolt/lib/Passes/PatchEntries.cpp b/bolt/lib/Passes/PatchEntries.cpp
index 02a044d8b2f69f5..ee7512d89962f69 100644
--- a/bolt/lib/Passes/PatchEntries.cpp
+++ b/bolt/lib/Passes/PatchEntries.cpp
@@ -98,6 +98,17 @@ void PatchEntries::runOnFunctions(BinaryContext &BC) {
     });
 
     if (!Success) {
+      // We can't change output layout for AArch64 due to LongJmp pass
+      if (BC.isAArch64()) {
+        if (opts::ForcePatch) {
+          errs() << "BOLT-ERROR: unable to patch entries in " << Function
+                 << "\n";
+          exit(1);
+        }
+
+        continue;
+      }
+
       // If the original function entries cannot be patched, then we cannot
       // safely emit new function body.
       errs() << "BOLT-WARNING: failed to patch entries in " << Function
diff --git a/bolt/test/AArch64/reloc-call26.s b/bolt/test/AArch64/reloc-call26.s
deleted file mode 100644
index 42e4f7f2b43786f..000000000000000
--- a/bolt/test/AArch64/reloc-call26.s
+++ /dev/null
@@ -1,51 +0,0 @@
-## This test checks processing of R_AARCH64_CALL26 relocation
-## when option `--funcs` is enabled
-
-## We want to test on relocations against functions with both higher
-## and lower addresses. The '--force-patch' option is used to prevent
-## the functions func1 and func2 from being optimized, so that their
-## addresses can remain unchanged. Therefore, the relocations can be
-## updated via encodeValueAArch64 and the address order in the output
-## binary is func1 < _start < func2.
-
-# REQUIRES: system-linux
-
-# RUN: llvm-mc -filetype=obj -triple aarch64-unknown-unknown \
-# RUN:   %s -o %t.o
-# RUN: %clang %cflags %t.o -o %t.exe -Wl,-q
-# RUN: llvm-bolt %t.exe -o %t.bolt --funcs=func1,func2 \
-# RUN:   --force-patch 2>&1 | FileCheck %s -check-prefix=CHECK-BOLT
-# RUN: llvm-objdump -d --disassemble-symbols='_start' %t.bolt | \
-# RUN:   FileCheck %s
-# RUN: llvm-nm --numeric-sort --extern-only %t.bolt  | FileCheck \
-# RUN:   %s -check-prefix=CHECK-FUNC-ORDER
-
-# CHECK-BOLT: BOLT-WARNING: failed to patch entries in func1. The function will not be optimized.
-# CHECK-BOLT: BOLT-WARNING: failed to patch entries in func2. The function will not be optimized.
-# CHECK: {{.*}} bl {{.*}} <func1>
-# CHECK: {{.*}} bl {{.*}} <func2>
-
-# CHECK-FUNC-ORDER: {{.*}} func1
-# CHECK-FUNC-ORDER-NEXT: {{.*}} _start
-# CHECK-FUNC-ORDER-NEXT: {{.*}} func2
-
-  .text
-  .align 4
-  .global func1
-  .type func1, %function
-func1:
-  ret
-  .size func1, .-func1
-  .global _start
-  .type _start, %function
-_start:
-  bl func1
-  bl func2
-  mov     w8, #93
-  svc     #0
-  .size _start, .-_start
-  .global func2
-  .type func2, %function
-func2:
-  ret
-  .size func2, .-func2
diff --git a/bolt/unittests/Core/BinaryContext.cpp b/bolt/unittests/Core/BinaryContext.cpp
index bac264141251ccd..e184d2fcaadc0ce 100644
--- a/bolt/unittests/Core/BinaryContext.cpp
+++ b/bolt/unittests/Core/BinaryContext.cpp
@@ -62,6 +62,45 @@ INSTANTIATE_TEST_SUITE_P(X86, BinaryContextTester,
 INSTANTIATE_TEST_SUITE_P(AArch64, BinaryContextTester,
                          ::testing::Values(Triple::aarch64));
 
+TEST_P(BinaryContextTester, FlushPendingRelocCALL26) {
+  if (GetParam() != Triple::aarch64)
+    GTEST_SKIP();
+
+  // This test checks that encodeValueAArch64 used by flushPendingRelocations
+  // returns right encoded values for CALL26 relocation for both backward and
+  // forward branches.
+  // The offsets layout is:
+  // 4:  func1
+  // 8:  bl func1
+  // 12: bl func2
+  // 16: func2
+
+  char Data[20] = {};
+  BinarySection &BS = BC->registerOrUpdateSection(
+      ".text", ELF::SHT_PROGBITS, ELF::SHF_EXECINSTR | ELF::SHF_ALLOC,
+      (uint8_t *)Data, sizeof(Data), 4);
+  MCSymbol *RelSymbol1 = BC->getOrCreateGlobalSymbol(4, "Func1");
+  ASSERT_TRUE(RelSymbol1);
+  BS.addRelocation(8, RelSymbol1, ELF::R_AARCH64_CALL26, 0, 0, true);
+  MCSymbol *RelSymbol2 = BC->getOrCreateGlobalSymbol(16, "Func2");
+  ASSERT_TRUE(RelSymbol2);
+  BS.addRelocation(12, RelSymbol2, ELF::R_AARCH64_CALL26, 0, 0, true);
+
+  std::error_code EC;
+  SmallVector<char> Vect(sizeof(Data));
+  raw_svector_ostream OS(Vect);
+
+  BS.flushPendingRelocations(OS, [&](const MCSymbol *S) {
+    return S == RelSymbol1 ? 4 : S == RelSymbol2 ? 16 : 0;
+  });
+
+  const char Func1Call[4] = {255, 255, 255, 151};
+  const char Func2Call[4] = {1, 0, 0, 148};
+
+  EXPECT_FALSE(memcmp(Func1Call, &Vect[8], 4)) << "Wrong backward call value\n";
+  EXPECT_FALSE(memcmp(Func2Call, &Vect[12], 4)) << "Wrong forward call value\n";
+}
+
 #endif
 
 TEST_P(BinaryContextTester, BaseAddress) {