[JITLink][ORC] Major JITLinkMemoryManager refactor.

This commit substantially refactors the JITLinkMemoryManager API to: (1) add
asynchronous versions of key operations, (2) give memory manager implementations
full control over link graph address layout, (3) enable more efficient tracking
of allocated memory, and (4) support "allocation actions" and finalize-lifetime
memory.

Together these changes provide a more usable API, and enable more powerful and
efficient memory manager implementations.

To support these changes the JITLinkMemoryManager::Allocation inner class has
been split into two new classes: InFlightAllocation, and FinalizedAllocation.
The allocate method returns an InFlightAllocation that tracks memory (both
working and executor memory) prior to finalization. The finalize method returns
a FinalizedAllocation object, and the InFlightAllocation is discarded. Breaking
Allocation into InFlightAllocation and FinalizedAllocation allows
InFlightAllocation subclassses to be written more naturally, and FinalizedAlloc
to be implemented and used efficiently (see (3) below).

In addition to the memory manager changes this commit also introduces a new
MemProt type to represent memory protections (MemProt replaces use of
sys::Memory::ProtectionFlags in JITLink), and a new MemDeallocPolicy type that
can be used to indicate when a section should be deallocated (see (4) below).

Plugin/pass writers who were using sys::Memory::ProtectionFlags will have to
switch to MemProt -- this should be straightworward. Clients with out-of-tree
memory managers will need to update their implementations. Clients using
in-tree memory managers should mostly be able to ignore it.

Major features:

(1) More asynchrony:

The allocate and deallocate methods are now asynchronous by default, with
synchronous convenience wrappers supplied. The asynchronous versions allow
clients (including JITLink) to request and deallocate memory without blocking.

(2) Improved control over graph address layout:

Instead of a SegmentRequestMap, JITLinkMemoryManager::allocate now takes a
reference to the LinkGraph to be allocated. The memory manager is responsible
for calculating the memory requirements for the graph, and laying out the graph
(setting working and executor memory addresses) within the allocated memory.
This gives memory managers full control over JIT'd memory layout. For clients
that don't need or want this degree of control the new "BasicLayout" utility can
be used to get a segment-based view of the graph, similar to the one provided by
SegmentRequestMap. Once segment addresses are assigned the BasicLayout::apply
method can be used to automatically lay out the graph.

(3) Efficient tracking of allocated memory.

The FinalizedAlloc type is a wrapper for an ExecutorAddr and requires only
64-bits to store in the controller. The meaning of the address held by the
FinalizedAlloc is left up to the memory manager implementation, but the
FinalizedAlloc type enforces a requirement that deallocate be called on any
non-default values prior to destruction. The deallocate method takes a
vector<FinalizedAlloc>, allowing for bulk deallocation of many allocations in a
single call.

Memory manager implementations will typically store the address of some
allocation metadata in the executor in the FinalizedAlloc, as holding this
metadata in the executor is often cheaper and may allow for clean deallocation
even in failure cases where the connection with the controller is lost.

(4) Support for "allocation actions" and finalize-lifetime memory.

Allocation actions are pairs (finalize_act, deallocate_act) of JITTargetAddress
triples (fn, arg_buffer_addr, arg_buffer_size), that can be attached to a
finalize request. At finalization time, after memory protections have been
applied, each of the "finalize_act" elements will be called in order (skipping
any elements whose fn value is zero) as

((char*(*)(const char *, size_t))fn)((const char *)arg_buffer_addr,
                                     (size_t)arg_buffer_size);

At deallocation time the deallocate elements will be run in reverse order (again
skipping any elements where fn is zero).

The returned char * should be null to indicate success, or a non-null
heap-allocated string error message to indicate failure.

These actions allow finalization and deallocation to be extended to include
operations like registering and deregistering eh-frames, TLS sections,
initializer and deinitializers, and language metadata sections. Previously these
operations required separate callWrapper invocations. Compared to callWrapper
invocations, actions require no extra IPC/RPC, reducing costs and eliminating
a potential source of errors.

Finalize lifetime memory can be used to support finalize actions: Sections with
finalize lifetime should be destroyed by memory managers immediately after
finalization actions have been run. Finalize memory can be used to support
finalize actions (e.g. with extra-metadata, or synthesized finalize actions)
without incurring permanent memory overhead.

1 files changed, 72 insertions, 63 deletions

diff --git a/llvm/lib/ExecutionEngine/Orc/DebugObjectManagerPlugin.cpp b/llvm/lib/ExecutionEngine/Orc/DebugObjectManagerPlugin.cpp
index 36efc744..fcfe389 100644
--- a/llvm/lib/ExecutionEngine/Orc/DebugObjectManagerPlugin.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/DebugObjectManagerPlugin.cpp
@@ -1,10 +1,15 @@
-//===---- DebugObjectManagerPlugin.h - JITLink debug objects ---*- C++ -*-===//
+//===------- DebugObjectManagerPlugin.cpp - JITLink debug objects ---------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
+//
+// FIXME: Update Plugin to poke the debug object into a new JITLink section,
+//        rather than creating a new allocation.
+//
+//===----------------------------------------------------------------------===//
 
 #include "llvm/ExecutionEngine/Orc/DebugObjectManagerPlugin.h"
 
@@ -108,70 +113,77 @@ void ELFDebugObjectSection<ELFT>::dump(raw_ostream &OS, StringRef Name) {
   }
 }
 
-static constexpr sys::Memory::ProtectionFlags ReadOnly =
-    static_cast<sys::Memory::ProtectionFlags>(sys::Memory::MF_READ);
-
 enum class Requirement {
   // Request final target memory load-addresses for all sections.
   ReportFinalSectionLoadAddresses,
 };
 
-/// The plugin creates a debug object from JITLinkContext when JITLink starts
-/// processing the corresponding LinkGraph. It provides access to the pass
-/// configuration of the LinkGraph and calls the finalization function, once
-/// the resulting link artifact was emitted.
+/// The plugin creates a debug object from when JITLink starts processing the
+/// corresponding LinkGraph. It provides access to the pass configuration of
+/// the LinkGraph and calls the finalization function, once the resulting link
+/// artifact was emitted.
 ///
 class DebugObject {
 public:
-  DebugObject(JITLinkContext &Ctx, ExecutionSession &ES) : Ctx(Ctx), ES(ES) {}
+  DebugObject(JITLinkMemoryManager &MemMgr, const JITLinkDylib *JD,
+              ExecutionSession &ES)
+      : MemMgr(MemMgr), JD(JD), ES(ES) {}
 
   void set(Requirement Req) { Reqs.insert(Req); }
   bool has(Requirement Req) const { return Reqs.count(Req) > 0; }
 
-  using FinalizeContinuation = std::function<void(Expected<sys::MemoryBlock>)>;
+  using FinalizeContinuation = std::function<void(Expected<ExecutorAddrRange>)>;
+
   void finalizeAsync(FinalizeContinuation OnFinalize);
 
   virtual ~DebugObject() {
-    if (Alloc)
-      if (Error Err = Alloc->deallocate())
+    if (Alloc) {
+      std::vector<FinalizedAlloc> Allocs;
+      Allocs.push_back(std::move(Alloc));
+      if (Error Err = MemMgr.deallocate(std::move(Allocs)))
         ES.reportError(std::move(Err));
+    }
   }
 
   virtual void reportSectionTargetMemoryRange(StringRef Name,
                                               SectionRange TargetMem) {}
 
 protected:
-  using Allocation = JITLinkMemoryManager::Allocation;
+  using InFlightAlloc = JITLinkMemoryManager::InFlightAlloc;
+  using FinalizedAlloc = JITLinkMemoryManager::FinalizedAlloc;
 
-  virtual Expected<std::unique_ptr<Allocation>>
-  finalizeWorkingMemory(JITLinkContext &Ctx) = 0;
+  virtual Expected<SimpleSegmentAlloc> finalizeWorkingMemory() = 0;
+
+  JITLinkMemoryManager &MemMgr;
+  const JITLinkDylib *JD = nullptr;
 
 private:
-  JITLinkContext &Ctx;
   ExecutionSession &ES;
   std::set<Requirement> Reqs;
-  std::unique_ptr<Allocation> Alloc{nullptr};
+  FinalizedAlloc Alloc;
 };
 
 // Finalize working memory and take ownership of the resulting allocation. Start
 // copying memory over to the target and pass on the result once we're done.
 // Ownership of the allocation remains with us for the rest of our lifetime.
 void DebugObject::finalizeAsync(FinalizeContinuation OnFinalize) {
-  assert(Alloc == nullptr && "Cannot finalize more than once");
-
-  auto AllocOrErr = finalizeWorkingMemory(Ctx);
-  if (!AllocOrErr)
-    OnFinalize(AllocOrErr.takeError());
-  Alloc = std::move(*AllocOrErr);
-
-  Alloc->finalizeAsync([this, OnFinalize](Error Err) {
-    if (Err)
-      OnFinalize(std::move(Err));
-    else
-      OnFinalize(sys::MemoryBlock(
-          jitTargetAddressToPointer<void *>(Alloc->getTargetMemory(ReadOnly)),
-          Alloc->getWorkingMemory(ReadOnly).size()));
-  });
+  assert(!Alloc && "Cannot finalize more than once");
+
+  if (auto SimpleSegAlloc = finalizeWorkingMemory()) {
+    auto ROSeg = SimpleSegAlloc->getSegInfo(MemProt::Read);
+    ExecutorAddrRange DebugObjRange(ExecutorAddr(ROSeg.Addr),
+                                    ExecutorAddrDiff(ROSeg.WorkingMem.size()));
+    SimpleSegAlloc->finalize(
+        [this, DebugObjRange,
+         OnFinalize = std::move(OnFinalize)](Expected<FinalizedAlloc> FA) {
+          if (FA) {
+            Alloc = std::move(*FA);
+            OnFinalize(DebugObjRange);
+          } else
+            OnFinalize(FA.takeError());
+        });
+  } else
+    OnFinalize(SimpleSegAlloc.takeError());
 }
 
 /// The current implementation of ELFDebugObject replicates the approach used in
@@ -190,8 +202,7 @@ public:
   StringRef getBuffer() const { return Buffer->getMemBufferRef().getBuffer(); }
 
 protected:
-  Expected<std::unique_ptr<Allocation>>
-  finalizeWorkingMemory(JITLinkContext &Ctx) override;
+  Expected<SimpleSegmentAlloc> finalizeWorkingMemory() override;
 
   template <typename ELFT>
   Error recordSection(StringRef Name,
@@ -201,15 +212,16 @@ protected:
 private:
   template <typename ELFT>
   static Expected<std::unique_ptr<ELFDebugObject>>
-  CreateArchType(MemoryBufferRef Buffer, JITLinkContext &Ctx,
-                 ExecutionSession &ES);
+  CreateArchType(MemoryBufferRef Buffer, JITLinkMemoryManager &MemMgr,
+                 const JITLinkDylib *JD, ExecutionSession &ES);
 
   static std::unique_ptr<WritableMemoryBuffer>
   CopyBuffer(MemoryBufferRef Buffer, Error &Err);
 
   ELFDebugObject(std::unique_ptr<WritableMemoryBuffer> Buffer,
-                 JITLinkContext &Ctx, ExecutionSession &ES)
-      : DebugObject(Ctx, ES), Buffer(std::move(Buffer)) {
+                 JITLinkMemoryManager &MemMgr, const JITLinkDylib *JD,
+                 ExecutionSession &ES)
+      : DebugObject(MemMgr, JD, ES), Buffer(std::move(Buffer)) {
     set(Requirement::ReportFinalSectionLoadAddresses);
   }
 
@@ -244,13 +256,14 @@ ELFDebugObject::CopyBuffer(MemoryBufferRef Buffer, Error &Err) {
 
 template <typename ELFT>
 Expected<std::unique_ptr<ELFDebugObject>>
-ELFDebugObject::CreateArchType(MemoryBufferRef Buffer, JITLinkContext &Ctx,
-                               ExecutionSession &ES) {
+ELFDebugObject::CreateArchType(MemoryBufferRef Buffer,
+                               JITLinkMemoryManager &MemMgr,
+                               const JITLinkDylib *JD, ExecutionSession &ES) {
   using SectionHeader = typename ELFT::Shdr;
 
   Error Err = Error::success();
   std::unique_ptr<ELFDebugObject> DebugObj(
-      new ELFDebugObject(CopyBuffer(Buffer, Err), Ctx, ES));
+      new ELFDebugObject(CopyBuffer(Buffer, Err), MemMgr, JD, ES));
   if (Err)
     return std::move(Err);
 
@@ -299,23 +312,26 @@ ELFDebugObject::Create(MemoryBufferRef Buffer, JITLinkContext &Ctx,
 
   if (Class == ELF::ELFCLASS32) {
     if (Endian == ELF::ELFDATA2LSB)
-      return CreateArchType<ELF32LE>(Buffer, Ctx, ES);
+      return CreateArchType<ELF32LE>(Buffer, Ctx.getMemoryManager(),
+                                     Ctx.getJITLinkDylib(), ES);
     if (Endian == ELF::ELFDATA2MSB)
-      return CreateArchType<ELF32BE>(Buffer, Ctx, ES);
+      return CreateArchType<ELF32BE>(Buffer, Ctx.getMemoryManager(),
+                                     Ctx.getJITLinkDylib(), ES);
     return nullptr;
   }
   if (Class == ELF::ELFCLASS64) {
     if (Endian == ELF::ELFDATA2LSB)
-      return CreateArchType<ELF64LE>(Buffer, Ctx, ES);
+      return CreateArchType<ELF64LE>(Buffer, Ctx.getMemoryManager(),
+                                     Ctx.getJITLinkDylib(), ES);
     if (Endian == ELF::ELFDATA2MSB)
-      return CreateArchType<ELF64BE>(Buffer, Ctx, ES);
+      return CreateArchType<ELF64BE>(Buffer, Ctx.getMemoryManager(),
+                                     Ctx.getJITLinkDylib(), ES);
     return nullptr;
   }
   return nullptr;
 }
 
-Expected<std::unique_ptr<DebugObject::Allocation>>
-ELFDebugObject::finalizeWorkingMemory(JITLinkContext &Ctx) {
+Expected<SimpleSegmentAlloc> ELFDebugObject::finalizeWorkingMemory() {
   LLVM_DEBUG({
     dbgs() << "Section load-addresses in debug object for \""
            << Buffer->getBufferIdentifier() << "\":\n";
@@ -324,28 +340,21 @@ ELFDebugObject::finalizeWorkingMemory(JITLinkContext &Ctx) {
   });
 
   // TODO: This works, but what actual alignment requirements do we have?
-  unsigned Alignment = sys::Process::getPageSizeEstimate();
-  JITLinkMemoryManager &MemMgr = Ctx.getMemoryManager();
-  const JITLinkDylib *JD = Ctx.getJITLinkDylib();
+  unsigned PageSize = sys::Process::getPageSizeEstimate();
   size_t Size = Buffer->getBufferSize();
 
   // Allocate working memory for debug object in read-only segment.
-  JITLinkMemoryManager::SegmentsRequestMap SingleReadOnlySegment;
-  SingleReadOnlySegment[ReadOnly] =
-      JITLinkMemoryManager::SegmentRequest(Alignment, Size, 0);
-
-  auto AllocOrErr = MemMgr.allocate(JD, SingleReadOnlySegment);
-  if (!AllocOrErr)
-    return AllocOrErr.takeError();
+  auto Alloc = SimpleSegmentAlloc::Create(
+      MemMgr, JD, {{MemProt::Read, {Size, Align(PageSize)}}});
+  if (!Alloc)
+    return Alloc;
 
   // Initialize working memory with a copy of our object buffer.
-  // TODO: Use our buffer as working memory directly.
-  std::unique_ptr<Allocation> Alloc = std::move(*AllocOrErr);
-  MutableArrayRef<char> WorkingMem = Alloc->getWorkingMemory(ReadOnly);
-  memcpy(WorkingMem.data(), Buffer->getBufferStart(), Size);
+  auto SegInfo = Alloc->getSegInfo(MemProt::Read);
+  memcpy(SegInfo.WorkingMem.data(), Buffer->getBufferStart(), Size);
   Buffer.reset();
 
-  return std::move(Alloc);
+  return Alloc;
 }
 
 void ELFDebugObject::reportSectionTargetMemoryRange(StringRef Name,
@@ -447,7 +456,7 @@ Error DebugObjectManagerPlugin::notifyEmitted(
   std::future<MSVCPError> FinalizeErr = FinalizePromise.get_future();
 
   It->second->finalizeAsync(
-      [this, &FinalizePromise, &MR](Expected<sys::MemoryBlock> TargetMem) {
+      [this, &FinalizePromise, &MR](Expected<ExecutorAddrRange> TargetMem) {
         // Any failure here will fail materialization.
         if (!TargetMem) {
           FinalizePromise.set_value(TargetMem.takeError());