[BOLT] Move from RuntimeDyld to JITLink

RuntimeDyld has been deprecated in favor of JITLink. [1] This patch
replaces all uses of RuntimeDyld in BOLT with JITLink.

Care has been taken to minimize the impact on the code structure in
order to ease the inspection of this (rather large) changeset. Since
BOLT relied on the RuntimeDyld API in multiple places, this wasn't
always possible though and I'll explain the changes in code structure
first.

Design note: BOLT uses a JIT linker to perform what essentially is
static linking. No linked code is ever executed; the result of linking
is simply written back to an executable file. For this reason, I
restricted myself to the use of the core JITLink library and avoided ORC
as much as possible.

RuntimeDyld contains methods for loading objects (loadObject) and symbol
lookup (getSymbol). Since JITLink doesn't provide a class with a similar
interface, the BOLTLinker abstract class was added to implement it. It
was added to Core since both the Rewrite and RuntimeLibs libraries make
use of it. Wherever a RuntimeDyld object was used before, it was
replaced with a BOLTLinker object.

There is one major difference between the RuntimeDyld and BOLTLinker
interfaces: in JITLink, section allocation and the application of fixups
(relocation) happens in a single call (jitlink::link). That is, there is
no separate method like finalizeWithMemoryManagerLocking in RuntimeDyld.
BOLT used to remap sections between allocating (loadObject) and linking
them (finalizeWithMemoryManagerLocking). This doesn't work anymore with
JITLink. Instead, BOLTLinker::loadObject accepts a callback that is
called before fixups are applied which is used to remap sections.

The actual implementation of the BOLTLinker interface lives in the
JITLinkLinker class in the Rewrite library. It's the only part of the
BOLT code that should directly interact with the JITLink API.

For loading object, JITLinkLinker first creates a LinkGraph
(jitlink::createLinkGraphFromObject) and then links it (jitlink::link).
For the latter, it uses a custom JITLinkContext with the following
properties:
- Use BOLT's ExecutableFileMemoryManager. This one was updated to
  implement the JITLinkMemoryManager interface. Since BOLT never
  executes code, its finalization step is a no-op.
- Pass config: don't use the default target passes since they modify
  DWARF sections in a way that seems incompatible with BOLT. Also run a
  custom pre-prune pass that makes sure sections without symbols are not
  pruned by JITLink.
- Implement symbol lookup. This used to be implemented by
  BOLTSymbolResolver.
- Call the section mapper callback before the final linking step.
- Copy symbol values when the LinkGraph is resolved. Symbols are stored
  inside JITLinkLinker to ensure that later objects (i.e.,
  instrumentation libraries) can find them. This functionality used to
  be provided by RuntimeDyld but I did not find a way to use JITLink
  directly for this.

Some more minor points of interest:
- BinarySection::SectionID: JITLink doesn't have something equivalent to
  RuntimeDyld's Section IDs. Instead, sections can only be referred to
  by name. Hence, SectionID was updated to a string.
- There seem to be no tests for Mach-O. I've tested a small hello-world
  style binary but not more than that.
- On Mach-O, JITLink "normalizes" section names to include the segment
  name. I had to parse the section name back from this manually which
  feels slightly hacky.

[1] https://reviews.llvm.org/D145686#4222642

Reviewed By: rafauler

Differential Revision: https://reviews.llvm.org/D147544

1 files changed, 131 insertions, 20 deletions

diff --git a/bolt/lib/Rewrite/ExecutableFileMemoryManager.cpp b/bolt/lib/Rewrite/ExecutableFileMemoryManager.cpp
index bc0dd2f..2718389 100644
--- a/bolt/lib/Rewrite/ExecutableFileMemoryManager.cpp
+++ b/bolt/lib/Rewrite/ExecutableFileMemoryManager.cpp
@@ -7,7 +7,9 @@
 //===----------------------------------------------------------------------===//
 
 #include "bolt/Rewrite/ExecutableFileMemoryManager.h"
+#include "bolt/Rewrite/JITLinkLinker.h"
 #include "bolt/Rewrite/RewriteInstance.h"
+#include "llvm/ExecutionEngine/JITLink/JITLink.h"
 #include "llvm/Support/MemAlloc.h"
 
 #undef  DEBUG_TYPE
@@ -21,30 +23,105 @@ namespace llvm {
 
 namespace bolt {
 
-uint8_t *ExecutableFileMemoryManager::allocateSection(
-    uintptr_t Size, unsigned Alignment, unsigned SectionID,
-    StringRef SectionName, bool IsCode, bool IsReadOnly) {
-  uint8_t *Ret = static_cast<uint8_t *>(llvm::allocate_buffer(Size, Alignment));
-  AllocatedSections.push_back(AllocInfo{Ret, Size, Alignment});
+namespace {
 
-  // A Size of 1 might mean an empty section for which RuntimeDyld decided to
-  // allocate 1 byte. In this case, the allocation will never be initialized
-  // causing non-deterministic output section contents.
-  if (Size == 1)
-    *Ret = 0;
+SmallVector<jitlink::Section *> orderedSections(jitlink::LinkGraph &G) {
+  SmallVector<jitlink::Section *> Sections(
+      llvm::map_range(G.sections(), [](auto &S) { return &S; }));
+  llvm::sort(Sections, [](const auto *LHS, const auto *RHS) {
+    return LHS->getOrdinal() < RHS->getOrdinal();
+  });
+  return Sections;
+}
+
+size_t sectionAlignment(const jitlink::Section &Section) {
+  assert(!Section.empty() && "Cannot get alignment for empty section");
+  return JITLinkLinker::orderedBlocks(Section).front()->getAlignment();
+}
+
+StringRef sectionName(const jitlink::Section &Section,
+                      const BinaryContext &BC) {
+  auto Name = Section.getName();
+
+  if (BC.isMachO()) {
+    // JITLink "normalizes" section names as "SegmentName,SectionName" on
+    // Mach-O. BOLT internally refers to sections just by the section name so
+    // strip-off the segment name.
+    auto SegmentEnd = Name.find(',');
+    assert(SegmentEnd != StringRef::npos && "Mach-O segment not found");
+    Name = Name.substr(SegmentEnd + 1);
+  }
+
+  return Name;
+}
+
+struct SectionAllocInfo {
+  void *Address;
+  size_t Size;
+  size_t Alignment;
+};
+
+struct AllocInfo {
+  SmallVector<SectionAllocInfo, 8> AllocatedSections;
+
+  ~AllocInfo() {
+    for (auto &Section : AllocatedSections)
+      deallocate_buffer(Section.Address, Section.Size, Section.Alignment);
+  }
+
+  SectionAllocInfo allocateSection(const jitlink::Section &Section) {
+    auto Size = JITLinkLinker::sectionSize(Section);
+    auto Alignment = sectionAlignment(Section);
+    auto *Buf = allocate_buffer(Size, Alignment);
+    SectionAllocInfo Alloc{Buf, Size, Alignment};
+    AllocatedSections.push_back(Alloc);
+    return Alloc;
+  }
+};
+
+struct BOLTInFlightAlloc : ExecutableFileMemoryManager::InFlightAlloc {
+  // Even though this is passed using a raw pointer in FinalizedAlloc, we keep
+  // it in a unique_ptr as long as possible to enjoy automatic cleanup when
+  // something goes wrong.
+  std::unique_ptr<AllocInfo> Alloc;
+
+public:
+  BOLTInFlightAlloc(std::unique_ptr<AllocInfo> Alloc)
+      : Alloc(std::move(Alloc)) {}
+
+  virtual void abandon(OnAbandonedFunction OnAbandoned) override {
+    llvm_unreachable("unexpected abandoned allocation");
+  }
+
+  virtual void finalize(OnFinalizedFunction OnFinalized) override {
+    OnFinalized(ExecutableFileMemoryManager::FinalizedAlloc(
+        orc::ExecutorAddr::fromPtr(Alloc.release())));
+  }
+};
+
+} // anonymous namespace
+
+void ExecutableFileMemoryManager::updateSection(
+    const jitlink::Section &JLSection, uint8_t *Contents, size_t Size,
+    size_t Alignment) {
+  auto SectionID = JLSection.getName();
+  auto SectionName = sectionName(JLSection, BC);
+  auto Prot = JLSection.getMemProt();
+  auto IsCode = (Prot & orc::MemProt::Exec) != orc::MemProt::None;
+  auto IsReadOnly = (Prot & orc::MemProt::Write) == orc::MemProt::None;
 
   // Register a debug section as a note section.
   if (!ObjectsLoaded && RewriteInstance::isDebugSection(SectionName)) {
     BinarySection &Section =
-        BC.registerOrUpdateNoteSection(SectionName, Ret, Size, Alignment);
+        BC.registerOrUpdateNoteSection(SectionName, Contents, Size, Alignment);
     Section.setSectionID(SectionID);
     assert(!Section.isAllocatable() && "note sections cannot be allocatable");
-    return Ret;
+    return;
   }
 
   if (!IsCode && (SectionName == ".strtab" || SectionName == ".symtab" ||
                   SectionName == "" || SectionName.startswith(".rela.")))
-    return Ret;
+    return;
 
   SmallVector<char, 256> Buf;
   if (ObjectsLoaded > 0) {
@@ -70,7 +147,7 @@ uint8_t *ExecutableFileMemoryManager::allocateSection(
            "Original section must exist and be allocatable.");
 
     Section = &OrgSection.get();
-    Section->updateContents(Ret, Size);
+    Section->updateContents(Contents, Size);
   } else {
     // If the input contains a section with the section name, rename it in the
     // output file to avoid the section name conflict and emit the new section
@@ -87,7 +164,7 @@ uint8_t *ExecutableFileMemoryManager::allocateSection(
     // sections in the input file.
     BinarySection &NewSection = BC.registerOrUpdateSection(
         UsePrefix ? NewSecPrefix + SectionName : SectionName, ELF::SHT_PROGBITS,
-        BinarySection::getFlags(IsReadOnly, IsCode, true), Ret, Size,
+        BinarySection::getFlags(IsReadOnly, IsCode, true), Contents, Size,
         Alignment);
     if (UsePrefix)
       NewSection.setOutputName(SectionName);
@@ -100,17 +177,51 @@ uint8_t *ExecutableFileMemoryManager::allocateSection(
            << " section : " << Section->getOutputName() << " ("
            << Section->getName() << ")"
            << " with size " << Size << ", alignment " << Alignment << " at "
-           << Ret << ", ID = " << SectionID << "\n";
+           << Contents << ", ID = " << SectionID << "\n";
   });
 
   Section->setSectionID(SectionID);
+}
+
+void ExecutableFileMemoryManager::allocate(const jitlink::JITLinkDylib *JD,
+                                           jitlink::LinkGraph &G,
+                                           OnAllocatedFunction OnAllocated) {
+  auto Alloc = std::make_unique<AllocInfo>();
+
+  for (auto *Section : orderedSections(G)) {
+    if (Section->empty())
+      continue;
+
+    auto SectionAlloc = Alloc->allocateSection(*Section);
+    updateSection(*Section, static_cast<uint8_t *>(SectionAlloc.Address),
+                  SectionAlloc.Size, SectionAlloc.Alignment);
+
+    size_t CurrentOffset = 0;
+    auto *Buf = static_cast<char *>(SectionAlloc.Address);
+    for (auto *Block : JITLinkLinker::orderedBlocks(*Section)) {
+      CurrentOffset = jitlink::alignToBlock(CurrentOffset, *Block);
+      auto BlockSize = Block->getSize();
+      auto *BlockBuf = Buf + CurrentOffset;
+
+      if (Block->isZeroFill())
+        std::memset(BlockBuf, 0, BlockSize);
+      else
+        std::memcpy(BlockBuf, Block->getContent().data(), BlockSize);
+
+      Block->setMutableContent({BlockBuf, Block->getSize()});
+      CurrentOffset += BlockSize;
+    }
+  }
 
-  return Ret;
+  OnAllocated(std::make_unique<BOLTInFlightAlloc>(std::move(Alloc)));
 }
 
-ExecutableFileMemoryManager::~ExecutableFileMemoryManager() {
-  for (const AllocInfo &AI : AllocatedSections)
-    llvm::deallocate_buffer(AI.Address, AI.Size, AI.Alignment);
+void ExecutableFileMemoryManager::deallocate(
+    std::vector<FinalizedAlloc> Allocs, OnDeallocatedFunction OnDeallocated) {
+  for (auto &Alloc : Allocs)
+    delete Alloc.release().toPtr<AllocInfo *>();
+
+  OnDeallocated(Error::success());
 }
 
 } // namespace bolt