[BOLT][AArch64] Add support for long absolute LLD thunks/veneers (#113408)

Absolute thunks generated by LLD reference function addresses recorded
as data in code. Since they are generated by the linker, they don't have
relocations associated with them and thus the addresses are left
undetected. Use pattern matching to detect such thunks and handle them
in VeneerElimination pass.

Author: maksfb

bolt/include/bolt/Core/MCPlusBuilder.h

@@ -1536,6 +1536,14 @@ class MCPlusBuilder {
     llvm_unreachable("not implemented");
   }
 
+  /// Match function \p BF to a long veneer for absolute code. Return true if
+  /// the match was successful and populate \p TargetAddress with an address of
+  /// the function veneer jumps to.
+  virtual bool matchAbsLongVeneer(const BinaryFunction &BF,
+                                  uint64_t &TargetAddress) const {
+    llvm_unreachable("not implemented");
+  }
+
   virtual bool matchAdrpAddPair(const MCInst &Adrp, const MCInst &Add) const {
     llvm_unreachable("not implemented");
     return false;

bolt/lib/Passes/VeneerElimination.cpp

@@ -29,30 +29,47 @@ static llvm::cl::opt<bool>
 namespace llvm {
 namespace bolt {
 
+static bool isPossibleVeneer(const BinaryFunction &BF) {
+  return BF.isAArch64Veneer() || BF.getOneName().starts_with("__AArch64");
+}
+
 Error VeneerElimination::runOnFunctions(BinaryContext &BC) {
   if (!opts::EliminateVeneers || !BC.isAArch64())
     return Error::success();
 
-  std::map<uint64_t, BinaryFunction> &BFs = BC.getBinaryFunctions();
   std::unordered_map<const MCSymbol *, const MCSymbol *> VeneerDestinations;
-  uint64_t VeneersCount = 0;
-  for (auto &It : BFs) {
-    BinaryFunction &VeneerFunction = It.second;
-    if (!VeneerFunction.isAArch64Veneer())
+  uint64_t NumEliminatedVeneers = 0;
+  for (BinaryFunction &BF : llvm::make_second_range(BC.getBinaryFunctions())) {
+    if (!isPossibleVeneer(BF))
       continue;
 
-    VeneersCount++;
-    VeneerFunction.setPseudo(true);
-    MCInst &FirstInstruction = *(VeneerFunction.begin()->begin());
-    const MCSymbol *VeneerTargetSymbol =
-        BC.MIB->getTargetSymbol(FirstInstruction, 1);
-    assert(VeneerTargetSymbol && "Expecting target symbol for instruction");
-    for (const MCSymbol *Symbol : VeneerFunction.getSymbols())
+    if (BF.isIgnored())
+      continue;
+
+    const MCSymbol *VeneerTargetSymbol = 0;
+    uint64_t TargetAddress;
+    if (BC.MIB->matchAbsLongVeneer(BF, TargetAddress)) {
+      if (BinaryFunction *TargetBF =
+              BC.getBinaryFunctionAtAddress(TargetAddress))
+        VeneerTargetSymbol = TargetBF->getSymbol();
+    } else {
+      MCInst &FirstInstruction = *(BF.begin()->begin());
+      if (BC.MIB->hasAnnotation(FirstInstruction, "AArch64Veneer"))
+        VeneerTargetSymbol = BC.MIB->getTargetSymbol(FirstInstruction, 1);
+    }
+
+    if (!VeneerTargetSymbol)
+      continue;
+
+    for (const MCSymbol *Symbol : BF.getSymbols())
       VeneerDestinations[Symbol] = VeneerTargetSymbol;
+
+    NumEliminatedVeneers++;
+    BF.setPseudo(true);
   }
 
   BC.outs() << "BOLT-INFO: number of removed linker-inserted veneers: "
-            << VeneersCount << "\n";
+            << NumEliminatedVeneers << '\n';
 
   // Handle veneers to veneers in case they occur
   for (auto &Entry : VeneerDestinations) {
@@ -65,9 +82,8 @@ Error VeneerElimination::runOnFunctions(BinaryContext &BC) {
   }
 
   uint64_t VeneerCallers = 0;
-  for (auto &It : BFs) {
-    BinaryFunction &Function = It.second;
-    for (BinaryBasicBlock &BB : Function) {
+  for (BinaryFunction &BF : llvm::make_second_range(BC.getBinaryFunctions())) {
+    for (BinaryBasicBlock &BB : BF) {
       for (MCInst &Instr : BB) {
         if (!BC.MIB->isCall(Instr) || BC.MIB->isIndirectCall(Instr))
           continue;

bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp

@@ -15,13 +15,16 @@
 #include "MCTargetDesc/AArch64MCExpr.h"
 #include "MCTargetDesc/AArch64MCTargetDesc.h"
 #include "Utils/AArch64BaseInfo.h"
+#include "bolt/Core/BinaryBasicBlock.h"
+#include "bolt/Core/BinaryFunction.h"
 #include "bolt/Core/MCPlusBuilder.h"
 #include "llvm/BinaryFormat/ELF.h"
 #include "llvm/MC/MCContext.h"
 #include "llvm/MC/MCFixupKindInfo.h"
 #include "llvm/MC/MCInstBuilder.h"
 #include "llvm/MC/MCInstrInfo.h"
 #include "llvm/MC/MCRegisterInfo.h"
+#include "llvm/Support/DataExtractor.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/ErrorHandling.h"
 
@@ -1320,6 +1323,67 @@ class AArch64MCPlusBuilder : public MCPlusBuilder {
     return 3;
   }
 
+  /// Match the following pattern:
+  ///
+  ///   LDR x16, .L1
+  ///   BR  x16
+  /// L1:
+  ///   .quad Target
+  ///
+  /// Populate \p TargetAddress with the Target value on successful match.
+  bool matchAbsLongVeneer(const BinaryFunction &BF,
+                          uint64_t &TargetAddress) const override {
+    if (BF.size() != 1 || BF.getMaxSize() < 16)
+      return false;
+
+    if (!BF.hasConstantIsland())
+      return false;
+
+    const BinaryBasicBlock &BB = BF.front();
+    if (BB.size() != 2)
+      return false;
+
+    const MCInst &LDRInst = BB.getInstructionAtIndex(0);
+    if (LDRInst.getOpcode() != AArch64::LDRXl)
+      return false;
+
+    if (!LDRInst.getOperand(0).isReg() ||
+        LDRInst.getOperand(0).getReg() != AArch64::X16)
+      return false;
+
+    const MCSymbol *TargetSym = getTargetSymbol(LDRInst, 1);
+    if (!TargetSym)
+      return false;
+
+    const MCInst &BRInst = BB.getInstructionAtIndex(1);
+    if (BRInst.getOpcode() != AArch64::BR)
+      return false;
+    if (!BRInst.getOperand(0).isReg() ||
+        BRInst.getOperand(0).getReg() != AArch64::X16)
+      return false;
+
+    const BinaryFunction::IslandInfo &IInfo = BF.getIslandInfo();
+    if (IInfo.HasDynamicRelocations)
+      return false;
+
+    auto Iter = IInfo.Offsets.find(8);
+    if (Iter == IInfo.Offsets.end() || Iter->second != TargetSym)
+      return false;
+
+    // Extract the absolute value stored inside the island.
+    StringRef SectionContents = BF.getOriginSection()->getContents();
+    StringRef FunctionContents = SectionContents.substr(
+        BF.getAddress() - BF.getOriginSection()->getAddress(), BF.getMaxSize());
+
+    const BinaryContext &BC = BF.getBinaryContext();
+    DataExtractor DE(FunctionContents, BC.AsmInfo->isLittleEndian(),
+                     BC.AsmInfo->getCodePointerSize());
+    uint64_t Offset = 8;
+    TargetAddress = DE.getAddress(&Offset);
+
+    return true;
+  }
+
   bool matchAdrpAddPair(const MCInst &Adrp, const MCInst &Add) const override {
     if (!isADRP(Adrp) || !isAddXri(Add))
       return false;

bolt/test/AArch64/veneer-lld-abs.s

@@ -0,0 +1,51 @@
+## Check that llvm-bolt correctly recognizes long absolute thunks generated
+## by LLD.
+
+# RUN: llvm-mc -filetype=obj -triple aarch64-unknown-unknown %s -o %t.o
+# RUN: %clang %cflags -fno-PIC -no-pie %t.o -o %t.exe -nostdlib \
+# RUN:    -fuse-ld=lld -Wl,-q
+# RUN: llvm-objdump -d %t.exe | FileCheck --check-prefix=CHECK-INPUT %s
+# RUN: llvm-objcopy --remove-section .rela.mytext %t.exe
+# RUN: llvm-bolt %t.exe -o %t.bolt --elim-link-veneers=true --lite=0
+# RUN: llvm-objdump -d -j .text  %t.bolt | \
+# RUN:   FileCheck --check-prefix=CHECK-OUTPUT %s
+
+.text
+.balign 4
+.global foo
+.type foo, %function
+foo:
+  adrp x1, foo
+  ret
+.size foo, .-foo
+
+.section ".mytext", "ax"
+.balign 4
+
+.global __AArch64AbsLongThunk_foo
+.type __AArch64AbsLongThunk_foo, %function
+__AArch64AbsLongThunk_foo:
+  ldr x16, .L1
+  br x16
+# CHECK-INPUT-LABEL: <__AArch64AbsLongThunk_foo>:
+# CHECK-INPUT-NEXT:    ldr
+# CHECK-INPUT-NEXT:    br
+.L1:
+  .quad foo
+.size __AArch64AbsLongThunk_foo, .-__AArch64AbsLongThunk_foo
+
+## Check that the thunk was removed from .text and _start() calls foo()
+## directly.
+
+# CHECK-OUTPUT-NOT: __AArch64AbsLongThunk_foo
+
+.global _start
+.type _start, %function
+_start:
+# CHECK-INPUT-LABEL:  <_start>:
+# CHECK-OUTPUT-LABEL: <_start>:
+  bl __AArch64AbsLongThunk_foo
+# CHECK-INPUT-NEXT:     bl {{.*}} <__AArch64AbsLongThunk_foo>
+# CHECK-OUTPUT-NEXT:    bl {{.*}} <foo>
+  ret
+.size _start, .-_start