[JITLink] Add initial Aarch64 support

[sgraenitz]

Summary

Set up basic infrastructure for 64-bit ARM architecture support in JITLink. It allows for loading a minimal object file and resolving a single relocation. Advanced features like GOT and PLT handling or relaxations were intentionally left out for the moment.

This patch follows the idea to keep implementations for ARM (32-bit) and Aaarch64 (64-bit) separate, because:

• it might be easier to share code with the MachO "arm64" JITLink backend
• LLVM has individual targets for ARM and Aaarch64 as well

[sgraenitz]

This is an early draft and it will require polishing here and there. Please don't hesitate to propose improvements.

[sgraenitz]

Basic polishing

[lhames]

This looks great to me. Thank you very much Stefan!

This patch follows the idea to keep implementations for ARM (32-bit) and Aaarch64 (64-bit) separate, because:
it might be easier to share code with the MachO "arm64" JITLink backend
LLVM has individual targets for ARM and Aaarch64 as well

Yes -- This is the right approach. We'll want to unify the AArch64 edge set the same way that we did for x86-64.

llvm/lib/ExecutionEngine/JITLink/ELF_aarch64.cpp

88–98

There is a lot of redundancy in these addRelocations functions. Do you think we could add a generic function (or functions) to the ELFLinkGraphBuilder so that clients can just write:

for_all_relas([](const Rela &R) {
  // Handle rela
});

?

I don't think this is something that we should do for this review, but it should be done either before or after it lands so that future backends don't keep copying this code.

[sgraenitz]

llvm/lib/ExecutionEngine/JITLink/ELF_aarch64.cpp

88–98

Yes, I agree. Shouldn't be too much effort. I will post something soon.

llvm/include/llvm/ExecutionEngine/JITLink/ELF_aarch64.h

+//===--- ELF_aarch64.h - JIT link functions for ELF/aarch64 --*- C++ -*----===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+//===----------------------------------------------------------------------===//
+//
+// jit-link functions for ELF/aarch64.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_EXECUTIONENGINE_JITLINK_ELF_AARCH64_H
+#define LLVM_EXECUTIONENGINE_JITLINK_ELF_AARCH64_H
+
+#include "llvm/ExecutionEngine/JITLink/JITLink.h"
+
+namespace llvm {
+namespace jitlink {
+
+/// Create a LinkGraph from an ELF/aarch64 relocatable object
+///
+/// Note: The graph does not take ownership of the underlying buffer, nor copy
+/// its contents. The caller is responsible for ensuring that the object buffer
+/// outlives the graph.
+Expected<std::unique_ptr<LinkGraph>>
+createLinkGraphFromELFObject_aarch64(MemoryBufferRef ObjectBuffer);
+
+/// jit-link the given object buffer, which must be a ELF aarch64 relocatable
+/// object file.
+void link_ELF_aarch64(std::unique_ptr<LinkGraph> G,
+                      std::unique_ptr<JITLinkContext> Ctx);
+
+} // end namespace jitlink
+} // end namespace llvm
+
+#endif // LLVM_EXECUTIONENGINE_JITLINK_ELF_AARCH64_H

llvm/include/llvm/ExecutionEngine/JITLink/aarch64.h

+//=== aarch64.h - Generic JITLink aarch64 edge kinds, utilities -*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Generic utilities for graphs representing aarch64 objects.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_EXECUTIONENGINE_JITLINK_AARCH64_H
+#define LLVM_EXECUTIONENGINE_JITLINK_AARCH64_H
+
+#include "llvm/ExecutionEngine/JITLink/JITLink.h"
+
+namespace llvm {
+namespace jitlink {
+namespace aarch64 {
+
+/// Represets aarch64 fixups
+enum EdgeKind_aarch64 : Edge::Kind {
+
+  /// Set a CALL immediate field to bits [27:2] of X = Target - Fixup + Addend
+  R_AARCH64_CALL26 = Edge::FirstRelocation,
+
+};
+
+/// Returns a string name for the given aarch64 edge. For debugging purposes
+/// only
+const char *getEdgeKindName(Edge::Kind K);
+
+} // namespace aarch64
+} // namespace jitlink
+} // namespace llvm
+
+#endif // LLVM_EXECUTIONENGINE_JITLINK_AARCH64_H

llvm/lib/ExecutionEngine/JITLink/CMakeLists.txt

 add_llvm_component_library(LLVMJITLink
   EHFrameSupport.cpp
   JITLink.cpp
   JITLinkGeneric.cpp
   JITLinkMemoryManager.cpp
 
   # Formats:
 
   # MachO
   MachO.cpp
   MachO_arm64.cpp
   MachO_x86_64.cpp
   MachOLinkGraphBuilder.cpp
 
   # ELF
-
   ELF.cpp
   ELFLinkGraphBuilder.cpp
+  ELF_aarch64.cpp
   ELF_riscv.cpp
   ELF_x86_64.cpp
 
   # Architectures:
+  aarch64.cpp
   riscv.cpp
   x86_64.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${LLVM_MAIN_INCLUDE_DIR}/llvm/ExecutionEngine/JITLink
 
   DEPENDS
   intrinsics_gen
 
   LINK_COMPONENTS
   BinaryFormat
   Object
   OrcTargetProcess
   Support
   )
- 
 target_link_libraries(LLVMJITLink
   PRIVATE
   LLVMObject
   LLVMOrcTargetProcess
   LLVMSupport
 )

llvm/lib/ExecutionEngine/JITLink/ELF.cpp

 //===-------------- ELF.cpp - JIT linker function for ELF -------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // ELF jit-link function.
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/ExecutionEngine/JITLink/ELF.h"
 
 #include "llvm/BinaryFormat/ELF.h"
+#include "llvm/ExecutionEngine/JITLink/ELF_aarch64.h"
 #include "llvm/ExecutionEngine/JITLink/ELF_riscv.h"
 #include "llvm/ExecutionEngine/JITLink/ELF_x86_64.h"
 #include "llvm/Object/ELF.h"
 #include "llvm/Support/Endian.h"
 #include "llvm/Support/Format.h"
 #include "llvm/Support/MemoryBuffer.h"
 #include <cstring>
 
createLinkGraphFromELFObject(MemoryBufferRef ObjectBuffer) {
   if (memcmp(Buffer.data(), ELF::ElfMagic, strlen(ELF::ElfMagic)) != 0)
     return make_error<JITLinkError>("ELF magic not valid");
 
   Expected<uint16_t> TargetMachineArch = readTargetMachineArch(Buffer);
   if (!TargetMachineArch)
     return TargetMachineArch.takeError();
 
   switch (*TargetMachineArch) {
+  case ELF::EM_AARCH64:
+    return createLinkGraphFromELFObject_aarch64(ObjectBuffer);
   case ELF::EM_RISCV:
     return createLinkGraphFromELFObject_riscv(ObjectBuffer);
   case ELF::EM_X86_64:
     return createLinkGraphFromELFObject_x86_64(ObjectBuffer);
   default:
     return make_error<JITLinkError>(
         "Unsupported target machine architecture in ELF object " +
         ObjectBuffer.getBufferIdentifier());
   }
 }
 
 void link_ELF(std::unique_ptr<LinkGraph> G,
               std::unique_ptr<JITLinkContext> Ctx) {
   switch (G->getTargetTriple().getArch()) {
+  case Triple::aarch64:
+    link_ELF_aarch64(std::move(G), std::move(Ctx));
+    return;
   case Triple::riscv32:
   case Triple::riscv64:
     link_ELF_riscv(std::move(G), std::move(Ctx));
     return;
   case Triple::x86_64:
     link_ELF_x86_64(std::move(G), std::move(Ctx));
     return;
   default:

llvm/lib/ExecutionEngine/JITLink/ELF_aarch64.cpp

+//===----- ELF_aarch64.cpp - JIT linker implementation for ELF/aarch64 ----===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// ELF/aarch64 jit-link implementation.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/ExecutionEngine/JITLink/ELF_aarch64.h"
+#include "ELFLinkGraphBuilder.h"
+#include "JITLinkGeneric.h"
+#include "llvm/BinaryFormat/ELF.h"
+#include "llvm/ExecutionEngine/JITLink/aarch64.h"
+#include "llvm/Object/ELFObjectFile.h"
+
+#define DEBUG_TYPE "jitlink"
+
+using namespace llvm;
+using namespace llvm::jitlink;
+
+namespace llvm {
+namespace jitlink {
+
+class ELFJITLinker_aarch64 : public JITLinker<ELFJITLinker_aarch64> {
+  friend class JITLinker<ELFJITLinker_aarch64>;
+
+public:
+  ELFJITLinker_aarch64(std::unique_ptr<JITLinkContext> Ctx,
+                       std::unique_ptr<LinkGraph> G,
+                       PassConfiguration PassConfig)
+      : JITLinker(std::move(Ctx), std::move(G), std::move(PassConfig)) {}
+
+private:
+  Error applyFixup(LinkGraph &G, Block &B, const Edge &E) const {
+    using namespace aarch64;
+    using namespace llvm::support;
+
+    char *BlockWorkingMem = B.getAlreadyMutableContent().data();
+    char *FixupPtr = BlockWorkingMem + E.getOffset();
+    JITTargetAddress FixupAddress = B.getAddress() + E.getOffset();
+    switch (E.getKind()) {
+    case aarch64::R_AARCH64_CALL26: {
+      assert((FixupAddress & 0x3) == 0 && "Call-inst is not 32-bit aligned");
+      int64_t Value = E.getTarget().getAddress() - FixupAddress + E.getAddend();
+
+      if (static_cast<uint64_t>(Value) & 0x3)
+        return make_error<JITLinkError>("Call target is not 32-bit aligned");
+
+      if (!fitsRangeSignedInt<27>(Value))
+        return makeTargetOutOfRangeError(G, B, E);
+
+      uint32_t RawInstr = *(little32_t *)FixupPtr;
+      assert((RawInstr & 0x7fffffff) == 0x14000000 &&
+             "RawInstr isn't a B or BR immediate instruction");
+      uint32_t Imm = (static_cast<uint32_t>(Value) & ((1 << 28) - 1)) >> 2;
+      uint32_t FixedInstr = RawInstr | Imm;
+      *(little32_t *)FixupPtr = FixedInstr;
+      break;
+    }
+    }
+    return Error::success();
+  }
+
+  template <uint8_t Bits> static bool fitsRangeSignedInt(int64_t Value) {
+    return Value >= -(1 << Bits) && Value < (1 << Bits);
+  }
+};
+
+template <typename ELFT>
+class ELFLinkGraphBuilder_aarch64 : public ELFLinkGraphBuilder<ELFT> {
+private:
+  static Expected<aarch64::EdgeKind_aarch64>
+  getRelocationKind(const uint32_t Type) {
+    using namespace aarch64;
+    switch (Type) {
+    case ELF::R_AARCH64_CALL26:
+      return EdgeKind_aarch64::R_AARCH64_CALL26;
+    }
+
+    return make_error<JITLinkError>("Unsupported aarch64 relocation:" +
+                                    formatv("{0:d}", Type));
+  }
+
+  Error addRelocations() override {
+    using Base = ELFLinkGraphBuilder<ELFT>;
+    LLVM_DEBUG(dbgs() << "Adding relocations\n");
+
+    // Iterate sections and only process the interesting ones.
+    for (auto &SecRef : Base::Sections) {
+      if (SecRef.sh_type != ELF::SHT_RELA && SecRef.sh_type != ELF::SHT_REL)
+        continue;
+      auto RelSectName = Base::Obj.getSectionName(SecRef);
+      if (!RelSectName)
+        return RelSectName.takeError();
+
+      LLVM_DEBUG({
+        dbgs() << "Adding relocations from section " << *RelSectName << "\n";
+      });
+
+      auto UpdateSect = Base::Obj.getSection(SecRef.sh_info);
+      if (!UpdateSect)
+        return UpdateSect.takeError();
+
+      auto UpdateSectName = Base::Obj.getSectionName(**UpdateSect);
+      if (!UpdateSectName)
+        return UpdateSectName.takeError();
+
+      // Don't process relocations for debug sections.
+      if (Base::isDwarfSection(*UpdateSectName)) {
+        LLVM_DEBUG({
+          dbgs() << "  Target is dwarf section " << *UpdateSectName
+                 << ". Skipping.\n";
+        });
+        continue;
+      }
+      LLVM_DEBUG(dbgs() << "  For target section " << *UpdateSectName << "\n");
+
+      auto *JITSection = Base::G->findSectionByName(*UpdateSectName);
+      if (!JITSection)
+        return make_error<llvm::StringError>(
+            "Refencing a section that wasn't added to graph" + *UpdateSectName,
+            llvm::inconvertibleErrorCode());
+
+      auto Relocations = Base::Obj.relas(SecRef);
+      if (!Relocations)
+        return Relocations.takeError();
+
+      for (const auto &Rela : *Relocations) {
+        auto Type = Rela.getType(false);
+
+        LLVM_DEBUG({
+          dbgs() << "Relocation Type: " << Type << "\n"
+                 << "Name: " << Base::Obj.getRelocationTypeName(Type) << "\n";
+        });
+
+        auto SymbolIndex = Rela.getSymbol(false);
+        auto Symbol = Base::Obj.getRelocationSymbol(Rela, Base::SymTabSec);
+        if (!Symbol)
+          return Symbol.takeError();
+
+        auto BlockToFix = *(JITSection->blocks().begin());
+        auto *TargetSymbol = Base::getGraphSymbol(SymbolIndex);
+
+        if (!TargetSymbol) {
+          return make_error<llvm::StringError>(
+              "Could not find symbol at given index, did you add it to "
+              "JITSymbolTable? index: " +
+                  std::to_string(SymbolIndex) + ", shndx: " +
+                  std::to_string((*Symbol)->st_shndx) + " Size of table: " +
+                  std::to_string(Base::GraphSymbols.size()),
+              llvm::inconvertibleErrorCode());
+        }
+        int64_t Addend = Rela.r_addend;
+        JITTargetAddress FixupAddress = (*UpdateSect)->sh_addr + Rela.r_offset;
+
+        LLVM_DEBUG({
+          dbgs() << "Processing relocation at "
+                 << format("0x%016" PRIx64, FixupAddress) << "\n";
+        });
+        auto Kind = getRelocationKind(Type);
+        if (!Kind)
+          return Kind.takeError();
+
+        BlockToFix->addEdge(*Kind, FixupAddress - BlockToFix->getAddress(),
+                            *TargetSymbol, Addend);
+      }
+    }
+    return Error::success();
+  }
+
+public:
+  ELFLinkGraphBuilder_aarch64(StringRef FileName,
+                              const object::ELFFile<ELFT> &Obj, const Triple T)
+      : ELFLinkGraphBuilder<ELFT>(Obj, std::move(T), FileName,
+                                  aarch64::getEdgeKindName) {}
+};
+
+Expected<std::unique_ptr<LinkGraph>>
+createLinkGraphFromELFObject_aarch64(MemoryBufferRef ObjectBuffer) {
+  LLVM_DEBUG({
+    dbgs() << "Building jitlink graph for new input "
+           << ObjectBuffer.getBufferIdentifier() << "...\n";
+  });
+
+  auto ELFObj = object::ObjectFile::createELFObjectFile(ObjectBuffer);
+  if (!ELFObj)
+    return ELFObj.takeError();
+
+  assert((*ELFObj)->getArch() == Triple::aarch64 &&
+         "Only AArch64 (little endian) is supported for now");
+
+  auto &ELFObjFile = cast<object::ELFObjectFile<object::ELF64LE>>(**ELFObj);
+  return ELFLinkGraphBuilder_aarch64<object::ELF64LE>((*ELFObj)->getFileName(),
+                                                      ELFObjFile.getELFFile(),
+                                                      (*ELFObj)->makeTriple())
+      .buildGraph();
+}
+
+void link_ELF_aarch64(std::unique_ptr<LinkGraph> G,
+                      std::unique_ptr<JITLinkContext> Ctx) {
+  PassConfiguration Config;
+  const Triple &TT = G->getTargetTriple();
+  if (Ctx->shouldAddDefaultTargetPasses(TT)) {
+    if (auto MarkLive = Ctx->getMarkLivePass(TT))
+      Config.PrePrunePasses.push_back(std::move(MarkLive));
+    else
+      Config.PrePrunePasses.push_back(markAllSymbolsLive);
+  }
+  if (auto Err = Ctx->modifyPassConfig(*G, Config))
+    return Ctx->notifyFailed(std::move(Err));
+
+  ELFJITLinker_aarch64::link(std::move(Ctx), std::move(G), std::move(Config));
+}
+
+} // namespace jitlink
+} // namespace llvm

llvm/lib/ExecutionEngine/JITLink/aarch64.cpp

+//===---- aarch64.cpp - Generic JITLink aarch64 edge kinds, utilities -----===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Generic utilities for graphs representing aarch64 objects.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/ExecutionEngine/JITLink/aarch64.h"
+
+#define DEBUG_TYPE "jitlink"
+
+namespace llvm {
+namespace jitlink {
+namespace aarch64 {
+
+const char *getEdgeKindName(Edge::Kind K) {
+  switch (K) {
+  case R_AARCH64_CALL26:
+    return "R_AARCH64_CALL26";
+  }
+  return getGenericEdgeKindName(K);
+}
+} // namespace aarch64
+} // namespace jitlink
+} // namespace llvm

llvm/test/ExecutionEngine/JITLink/AArch64/ELF_aarch64_minimal.s

+# RUN: rm -rf %t && mkdir -p %t
+# RUN: llvm-mc -triple=aarch64-unknown-linux-gnu -relax-relocations=false -position-independent -filetype=obj -o %t/aarch64_reloc.o %s
+# RUN: llvm-jitlink -noexec %t/aarch64_reloc.o
+
+  .text
+  .globl  sub1
+  .p2align  2
+  .type sub1,@function
+sub1:
+  sub sp, sp, #16
+  str w0, [sp, #12]
+  ldr w8, [sp, #12]
+  subs  w0, w8, #1
+  add sp, sp, #16
+  ret
+
+  .size sub1, .-sub1
+
+  .globl  main
+  .p2align  2
+  .type main,@function
+main:
+  sub sp, sp, #32
+  stp  x29, x30, [sp, #16]
+  add x29, sp, #16
+  stur  wzr, [x29, #-4]
+  str w0, [sp, #8]
+  str x1, [sp]
+  ldr w0, [sp, #8]
+  bl  sub1
+  ldp  x29, x30, [sp, #16]
+  add sp, sp, #32
+  ret
+
+  .size main, .-main