[llvm] [BOLT] support AArch64 JUMP26 createRelocation (PR #83531)

[via llvm-commits]

https://github.com/linsinan1995 updated https://github.com/llvm/llvm-project/pull/83531

>From ff39b0722320fdd4ff7ffe26843ebba467fb07d2 Mon Sep 17 00:00:00 2001
From: Sinan Lin <sinan.lin at linux.alibaba.com>
Date: Mon, 19 Feb 2024 14:33:46 +0800
Subject: [PATCH] [BOLT] support AArch64 JUMP26 createRelocation

Add R_AARCH64_JUMP26 implementation for createRelocation, which
could significantly reduce the number of failed scan-refs cases.
---
 bolt/lib/Core/Relocation.cpp                  |  8 ++++
 .../Target/AArch64/AArch64MCPlusBuilder.cpp   |  3 ++
 bolt/unittests/Core/BinaryContext.cpp         | 43 +++++++++++++++++++
 3 files changed, 54 insertions(+)

diff --git a/bolt/lib/Core/Relocation.cpp b/bolt/lib/Core/Relocation.cpp
index ded497f3c06118..cbf95a7db08b52 100644
--- a/bolt/lib/Core/Relocation.cpp
+++ b/bolt/lib/Core/Relocation.cpp
@@ -381,6 +381,14 @@ static uint64_t encodeValueAArch64(uint64_t Type, uint64_t Value, uint64_t PC) {
     // OP 1001_01 goes in bits 31:26 of BL.
     Value = ((Value >> 2) & 0x3ffffff) | 0x94000000ULL;
     break;
+  case ELF::R_AARCH64_JUMP26:
+    Value -= PC;
+    assert(isInt<28>(Value) &&
+           "only PC +/- 128MB is allowed for direct branch");
+    // Immediate goes in bits 25:0 of B.
+    // OP 0001_01 goes in bits 31:26 of B.
+    Value = ((Value >> 2) & 0x3ffffff) | 0x14000000ULL;
+    break;
   }
   return Value;
 }
diff --git a/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp b/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
index c1c09c70ab01da..9efb428a6e1bae 100644
--- a/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
+++ b/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
@@ -1629,6 +1629,9 @@ class AArch64MCPlusBuilder : public MCPlusBuilder {
     uint64_t RelType;
     if (Fixup.getKind() == MCFixupKind(AArch64::fixup_aarch64_pcrel_call26))
       RelType = ELF::R_AARCH64_CALL26;
+    else if (Fixup.getKind() ==
+             MCFixupKind(AArch64::fixup_aarch64_pcrel_branch26))
+      RelType = ELF::R_AARCH64_JUMP26;
     else if (FKI.Flags & MCFixupKindInfo::FKF_IsPCRel) {
       switch (FKI.TargetSize) {
       default:
diff --git a/bolt/unittests/Core/BinaryContext.cpp b/bolt/unittests/Core/BinaryContext.cpp
index 94ee65e63a1dcf..99568b62ca8aab 100644
--- a/bolt/unittests/Core/BinaryContext.cpp
+++ b/bolt/unittests/Core/BinaryContext.cpp
@@ -104,6 +104,49 @@ TEST_P(BinaryContextTester, FlushPendingRelocCALL26) {
   EXPECT_FALSE(memcmp(Func2Call, &Vect[12], 4)) << "Wrong forward call value\n";
 }
 
+TEST_P(BinaryContextTester, FlushPendingRelocJUMP26) {
+  if (GetParam() != Triple::aarch64)
+    GTEST_SKIP();
+
+  // This test checks that encodeValueAArch64 used by flushPendingRelocations
+  // returns correctly encoded values for R_AARCH64_JUMP26 relocation for both
+  // backward and forward branches.
+  //
+  // The offsets layout is:
+  // 4:  func1
+  // 8:  b func1
+  // 12: b func2
+  // 16: func2
+
+  uint64_t Size = 20;
+  char *Data = new char[Size];
+  BinarySection &BS = BC->registerOrUpdateSection(
+      ".text", ELF::SHT_PROGBITS, ELF::SHF_EXECINSTR | ELF::SHF_ALLOC,
+      (uint8_t *)Data, Size, 4);
+  MCSymbol *RelSymbol1 = BC->getOrCreateGlobalSymbol(4, "Func1");
+  ASSERT_TRUE(RelSymbol1);
+  BS.addRelocation(8, RelSymbol1, ELF::R_AARCH64_JUMP26, 0, 0, true);
+  MCSymbol *RelSymbol2 = BC->getOrCreateGlobalSymbol(16, "Func2");
+  ASSERT_TRUE(RelSymbol2);
+  BS.addRelocation(12, RelSymbol2, ELF::R_AARCH64_JUMP26, 0, 0, true);
+
+  std::error_code EC;
+  SmallVector<char> Vect(Size);
+  raw_svector_ostream OS(Vect);
+
+  BS.flushPendingRelocations(OS, [&](const MCSymbol *S) {
+    return S == RelSymbol1 ? 4 : S == RelSymbol2 ? 16 : 0;
+  });
+
+  const uint8_t Func1Call[4] = {255, 255, 255, 23};
+  const uint8_t Func2Call[4] = {1, 0, 0, 20};
+
+  EXPECT_FALSE(memcmp(Func1Call, &Vect[8], 4))
+      << "Wrong backward branch value\n";
+  EXPECT_FALSE(memcmp(Func2Call, &Vect[12], 4))
+      << "Wrong forward branch value\n";
+}
+
 #endif
 
 TEST_P(BinaryContextTester, BaseAddress) {