[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

Author: mtvec

bolt/include/bolt/Core/BinaryContext.h

@@ -59,7 +59,6 @@ using namespace object;
 namespace bolt {
 
 class BinaryFunction;
-class ExecutableFileMemoryManager;
 
 /// Information on loadable part of the file.
 struct SegmentInfo {
@@ -313,10 +312,6 @@ class BinaryContext {
       FilterIterator<binary_data_const_iterator>;
   using FilteredBinaryDataIterator = FilterIterator<binary_data_iterator>;
 
-  /// Memory manager for sections and segments. Used to communicate with ORC
-  /// among other things.
-  std::shared_ptr<ExecutableFileMemoryManager> EFMM;
-
   StringRef getFilename() const { return Filename; }
   void setFilename(StringRef Name) { Filename = std::string(Name); }
 

bolt/include/bolt/Core/BinarySection.h

@@ -90,7 +90,7 @@ class BinarySection {
   uint64_t OutputFileOffset{0};    // File offset in the rewritten binary file.
   StringRef OutputContents;        // Rewritten section contents.
   const uint64_t SectionNumber;    // Order in which the section was created.
-  unsigned SectionID{-1u};         // Unique ID used for address mapping.
+  std::string SectionID;           // Unique ID used for address mapping.
                                    // Set by ExecutableFileMemoryManager.
   uint32_t Index{0};               // Section index in the output file.
   mutable bool IsReordered{false}; // Have the contents been reordered?
@@ -430,18 +430,18 @@ class BinarySection {
   }
   uint64_t getOutputAddress() const { return OutputAddress; }
   uint64_t getOutputFileOffset() const { return OutputFileOffset; }
-  unsigned getSectionID() const {
+  StringRef getSectionID() const {
     assert(hasValidSectionID() && "trying to use uninitialized section id");
     return SectionID;
   }
-  bool hasValidSectionID() const { return SectionID != -1u; }
+  bool hasValidSectionID() const { return !SectionID.empty(); }
   bool hasValidIndex() { return Index != 0; }
   uint32_t getIndex() const { return Index; }
 
   // mutation
   void setOutputAddress(uint64_t Address) { OutputAddress = Address; }
   void setOutputFileOffset(uint64_t Offset) { OutputFileOffset = Offset; }
-  void setSectionID(unsigned ID) {
+  void setSectionID(StringRef ID) {
     assert(!hasValidSectionID() && "trying to set section id twice");
     SectionID = ID;
   }

bolt/include/bolt/Core/Linker.h

@@ -0,0 +1,48 @@
+//===- bolt/Core/Linker.h - BOLTLinker interface ----------------*- 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file contains the interface BOLT uses for linking.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef BOLT_CORE_LINKER_H
+#define BOLT_CORE_LINKER_H
+
+#include "llvm/ADT/StringRef.h"
+#include "llvm/Support/MemoryBufferRef.h"
+
+#include <cstdint>
+#include <functional>
+#include <optional>
+
+namespace llvm {
+namespace bolt {
+
+class BinarySection;
+
+class BOLTLinker {
+public:
+  using SectionMapper =
+      std::function<void(const BinarySection &Section, uint64_t Address)>;
+  using SectionsMapper = std::function<void(SectionMapper)>;
+
+  virtual ~BOLTLinker() = default;
+
+  /// Load and link \p Obj. \p MapSections will be called before the object is
+  /// linked to allow section addresses to be remapped. When called, the address
+  /// of a section can be changed by calling the passed SectionMapper.
+  virtual void loadObject(MemoryBufferRef Obj, SectionsMapper MapSections) = 0;
+
+  /// Return the address of a symbol or std::nullopt if it cannot be found.
+  virtual std::optional<uint64_t> lookupSymbol(StringRef Name) const = 0;
+};
+
+} // namespace bolt
+} // namespace llvm
+
+#endif // BOLT_CORE_LINKER_H

bolt/include/bolt/Rewrite/ExecutableFileMemoryManager.h

@@ -10,7 +10,7 @@
 #define BOLT_REWRITE_EXECUTABLE_FILE_MEMORY_MANAGER_H
 
 #include "llvm/ADT/StringRef.h"
-#include "llvm/ExecutionEngine/RuntimeDyld.h"
+#include "llvm/ExecutionEngine/JITLink/JITLinkMemoryManager.h"
 #include <cstdint>
 #include <string>
 
@@ -20,20 +20,12 @@ namespace bolt {
 class BinaryContext;
 
 /// Class responsible for allocating and managing code and data sections.
-class ExecutableFileMemoryManager : public RuntimeDyld::MemoryManager {
+class ExecutableFileMemoryManager : public jitlink::JITLinkMemoryManager {
 private:
-  uint8_t *allocateSection(uintptr_t Size, unsigned Alignment,
-                           unsigned SectionID, StringRef SectionName,
-                           bool IsCode, bool IsReadOnly);
-  BinaryContext &BC;
-  bool AllowStubs;
+  void updateSection(const jitlink::Section &Section, uint8_t *Contents,
+                     size_t Size, size_t Alignment);
 
-  struct AllocInfo {
-    uint8_t *Address;
-    size_t Size;
-    size_t Alignment;
-  };
-  SmallVector<AllocInfo, 8> AllocatedSections;
+  BinaryContext &BC;
 
   // All new sections will be identified by the following prefix.
   std::string NewSecPrefix;
@@ -50,48 +42,14 @@ class ExecutableFileMemoryManager : public RuntimeDyld::MemoryManager {
   // user-supplied objects into the main input executable.
   uint32_t ObjectsLoaded{0};
 
-  ExecutableFileMemoryManager(BinaryContext &BC, bool AllowStubs)
-      : BC(BC), AllowStubs(AllowStubs) {}
-
-  ~ExecutableFileMemoryManager();
-
-  uint8_t *allocateCodeSection(uintptr_t Size, unsigned Alignment,
-                               unsigned SectionID,
-                               StringRef SectionName) override {
-    return allocateSection(Size, Alignment, SectionID, SectionName,
-                           /*IsCode=*/true, true);
-  }
-
-  uint8_t *allocateDataSection(uintptr_t Size, unsigned Alignment,
-                               unsigned SectionID, StringRef SectionName,
-                               bool IsReadOnly) override {
-    return allocateSection(Size, Alignment, SectionID, SectionName,
-                           /*IsCode=*/false, IsReadOnly);
-  }
-
-  // Ignore TLS sections by treating them as a regular data section
-  TLSSection allocateTLSSection(uintptr_t Size, unsigned Alignment,
-                                unsigned SectionID,
-                                StringRef SectionName) override {
-    TLSSection Res;
-    Res.Offset = 0;
-    Res.InitializationImage = allocateDataSection(
-        Size, Alignment, SectionID, SectionName, /*IsReadOnly=*/false);
-    return Res;
-  }
-
-  bool allowStubAllocation() const override { return AllowStubs; }
+  ExecutableFileMemoryManager(BinaryContext &BC) : BC(BC) {}
 
-  /// Count processed objects and skip memory finalization.
-  bool finalizeMemory(std::string *ErrMsg) override {
-    ++ObjectsLoaded;
-    return false;
-  }
+  void allocate(const jitlink::JITLinkDylib *JD, jitlink::LinkGraph &G,
+                OnAllocatedFunction OnAllocated) override;
 
-  /// Ignore EH frames.
-  void registerEHFrames(uint8_t *Addr, uint64_t LoadAddr,
-                        size_t Size) override {}
-  void deregisterEHFrames() override {}
+  void deallocate(std::vector<FinalizedAlloc> Allocs,
+                  OnDeallocatedFunction OnDeallocated) override;
+  using JITLinkMemoryManager::deallocate;
 
   /// Section name management.
   void setNewSecPrefix(StringRef Prefix) { NewSecPrefix = Prefix; }

bolt/include/bolt/Rewrite/JITLinkLinker.h

@@ -0,0 +1,56 @@
+//===- bolt/Rewrite/JITLinkLinker.h - Linker using JITLink ------*- 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
+//
+//===----------------------------------------------------------------------===//
+//
+// BOLTLinker using JITLink.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef BOLT_REWRITE_JITLINK_LINKER_H
+#define BOLT_REWRITE_JITLINK_LINKER_H
+
+#include "bolt/Core/Linker.h"
+#include "bolt/Rewrite/ExecutableFileMemoryManager.h"
+#include "llvm/ExecutionEngine/JITLink/JITLinkDylib.h"
+
+#include <map>
+#include <memory>
+#include <vector>
+
+namespace llvm {
+namespace bolt {
+
+class BinaryContext;
+
+class JITLinkLinker : public BOLTLinker {
+private:
+  struct Context;
+  friend struct Context;
+
+  BinaryContext &BC;
+  std::unique_ptr<ExecutableFileMemoryManager> MM;
+  jitlink::JITLinkDylib Dylib{"main"};
+  std::vector<ExecutableFileMemoryManager::FinalizedAlloc> Allocs;
+  std::map<std::string, uint64_t> Symtab;
+
+public:
+  JITLinkLinker(BinaryContext &BC,
+                std::unique_ptr<ExecutableFileMemoryManager> MM);
+  ~JITLinkLinker();
+
+  void loadObject(MemoryBufferRef Obj, SectionsMapper MapSections) override;
+  std::optional<uint64_t> lookupSymbol(StringRef Name) const override;
+
+  static SmallVector<jitlink::Block *, 2>
+  orderedBlocks(const jitlink::Section &Section);
+  static size_t sectionSize(const jitlink::Section &Section);
+};
+
+} // namespace bolt
+} // namespace llvm
+
+#endif // BOLT_REWRITE_JITLINK_LINKER_H

bolt/include/bolt/Rewrite/MachORewriteInstance.h

@@ -13,13 +13,13 @@
 #ifndef BOLT_REWRITE_MACHO_REWRITE_INSTANCE_H
 #define BOLT_REWRITE_MACHO_REWRITE_INSTANCE_H
 
+#include "bolt/Core/Linker.h"
 #include "bolt/Utils/NameResolver.h"
 #include "llvm/Support/Error.h"
 #include <memory>
 
 namespace llvm {
 class ToolOutputFile;
-class RuntimeDyld;
 class raw_pwrite_stream;
 namespace object {
 class MachOObjectFile;
@@ -37,7 +37,7 @@ class MachORewriteInstance {
 
   NameResolver NR;
 
-  std::unique_ptr<RuntimeDyld> RTDyld;
+  std::unique_ptr<BOLTLinker> Linker;
 
   std::unique_ptr<ToolOutputFile> Out;
 
@@ -49,8 +49,9 @@ class MachORewriteInstance {
   static StringRef getNewSecPrefix() { return ".bolt.new"; }
   static StringRef getOrgSecPrefix() { return ".bolt.org"; }
 
-  void mapInstrumentationSection(StringRef SectionName);
-  void mapCodeSections();
+  void mapInstrumentationSection(StringRef SectionName,
+                                 BOLTLinker::SectionMapper MapSection);
+  void mapCodeSections(BOLTLinker::SectionMapper MapSection);
 
   void adjustCommandLineOptions();
   void readSpecialSections();

bolt/include/bolt/Rewrite/RewriteInstance.h

@@ -14,6 +14,7 @@
 #define BOLT_REWRITE_REWRITE_INSTANCE_H
 
 #include "bolt/Core/BinaryContext.h"
+#include "bolt/Core/Linker.h"
 #include "bolt/Utils/NameResolver.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/MC/StringTableBuilder.h"
@@ -206,13 +207,13 @@ class RewriteInstance {
   std::vector<BinarySection *> getCodeSections();
 
   /// Map all sections to their final addresses.
-  void mapFileSections(RuntimeDyld &RTDyld);
+  void mapFileSections(BOLTLinker::SectionMapper MapSection);
 
   /// Map code sections generated by BOLT.
-  void mapCodeSections(RuntimeDyld &RTDyld);
+  void mapCodeSections(BOLTLinker::SectionMapper MapSection);
 
   /// Map the rest of allocatable sections.
-  void mapAllocatableSections(RuntimeDyld &RTDyld);
+  void mapAllocatableSections(BOLTLinker::SectionMapper MapSection);
 
   /// Update output object's values based on the final \p Layout.
   void updateOutputValues(const MCAsmLayout &Layout);
@@ -459,7 +460,7 @@ class RewriteInstance {
   std::unique_ptr<CFIReaderWriter> CFIRdWrt;
 
   // Run ExecutionEngine linker with custom memory manager and symbol resolver.
-  std::unique_ptr<RuntimeDyld> RTDyld;
+  std::unique_ptr<BOLTLinker> Linker;
 
   /// Output file where we mix original code from the input binary and
   /// optimized code for selected functions.

bolt/include/bolt/RuntimeLibs/HugifyRuntimeLibrary.h

@@ -28,8 +28,8 @@ class HugifyRuntimeLibrary : public RuntimeLibrary {
 
   void emitBinary(BinaryContext &BC, MCStreamer &Streamer) final {}
 
-  void link(BinaryContext &BC, StringRef ToolPath, RuntimeDyld &RTDyld,
-            std::function<void(RuntimeDyld &)> OnLoad) final;
+  void link(BinaryContext &BC, StringRef ToolPath, BOLTLinker &Linker,
+            BOLTLinker::SectionsMapper MapSections) override;
 };
 
 } // namespace bolt

bolt/include/bolt/RuntimeLibs/InstrumentationRuntimeLibrary.h

@@ -35,8 +35,8 @@ class InstrumentationRuntimeLibrary : public RuntimeLibrary {
 
   void emitBinary(BinaryContext &BC, MCStreamer &Streamer) final;
 
-  void link(BinaryContext &BC, StringRef ToolPath, RuntimeDyld &RTDyld,
-            std::function<void(RuntimeDyld &)> OnLoad) final;
+  void link(BinaryContext &BC, StringRef ToolPath, BOLTLinker &Linker,
+            BOLTLinker::SectionsMapper MapSections) override;
 
 private:
   std::string buildTables(BinaryContext &BC);

bolt/include/bolt/RuntimeLibs/RuntimeLibrary.h

@@ -15,14 +15,14 @@
 #ifndef BOLT_RUNTIMELIBS_RUNTIME_LIBRARY_H
 #define BOLT_RUNTIMELIBS_RUNTIME_LIBRARY_H
 
+#include "bolt/Core/Linker.h"
 #include "llvm/ADT/StringRef.h"
 #include <functional>
 #include <vector>
 
 namespace llvm {
 
 class MCStreamer;
-class RuntimeDyld;
 
 namespace bolt {
 
@@ -51,8 +51,8 @@ class RuntimeLibrary {
   virtual void emitBinary(BinaryContext &BC, MCStreamer &Streamer) = 0;
 
   /// Link with the library code.
-  virtual void link(BinaryContext &BC, StringRef ToolPath, RuntimeDyld &RTDyld,
-                    std::function<void(RuntimeDyld &)> OnLoad) = 0;
+  virtual void link(BinaryContext &BC, StringRef ToolPath, BOLTLinker &Linker,
+                    BOLTLinker::SectionsMapper MapSections) = 0;
 
 protected:
   /// The fini and init address set by the runtime library.
@@ -63,7 +63,8 @@ class RuntimeLibrary {
   static std::string getLibPath(StringRef ToolPath, StringRef LibFileName);
 
   /// Load a static runtime library specified by \p LibPath.
-  static void loadLibrary(StringRef LibPath, RuntimeDyld &RTDyld);
+  static void loadLibrary(StringRef LibPath, BOLTLinker &Linker,
+                          BOLTLinker::SectionsMapper MapSections);
 };
 
 } // namespace bolt

bolt/lib/Core/BinaryContext.cpp

@@ -1975,7 +1975,9 @@ void BinaryContext::deregisterUnusedSections() {
   ErrorOr<BinarySection &> AbsSection = getUniqueSectionByName("<absolute>");
   for (auto SI = Sections.begin(); SI != Sections.end();) {
     BinarySection *Section = *SI;
-    if (Section->hasSectionRef() || Section->getOutputSize() ||
+    // We check getOutputData() instead of getOutputSize() because sometimes
+    // zero-sized .text.cold sections are allocated.
+    if (Section->hasSectionRef() || Section->getOutputData() ||
         (AbsSection && Section == &AbsSection.get())) {
       ++SI;
       continue;