[JITLink] Add JITLinkDylib type, thread through JITLinkMemoryManager APIs.

JITLinkDylib represents a target dylib for a JITLink link. By representing this
explicitly we can:
  - Enable JITLinkMemoryManagers to manage allocations on a per-dylib basis
    (e.g by maintaining a seperate allocation pool for each JITLinkDylib).
  - Enable new features and diagnostics that require information about the
    target dylib (not implemented in this patch).

include/llvm/ExecutionEngine/JITLink/JITLink.h

@@ -1270,9 +1270,15 @@ class JITLinkContext {
 public:
   using LookupMap = DenseMap<StringRef, SymbolLookupFlags>;
 
+  /// Create a JITLinkContext.
+  JITLinkContext(const JITLinkDylib *JD) : JD(JD) {}
+
   /// Destroy a JITLinkContext.
   virtual ~JITLinkContext();
 
+  /// Return the JITLinkDylib that this link is targeting, if any.
+  const JITLinkDylib *getJITLinkDylib() const { return JD; }
+
   /// Return the MemoryManager to be used for this link.
   virtual JITLinkMemoryManager &getMemoryManager() = 0;
 
@@ -1324,6 +1330,9 @@ public:
   /// Called by JITLink to modify the pass pipeline prior to linking.
   /// The default version performs no modification.
   virtual Error modifyPassConfig(const Triple &TT, PassConfiguration &Config);
+
+private:
+  const JITLinkDylib *JD = nullptr;
 };
 
 /// Marks all symbols in a graph live. This can be used as a default,

include/llvm/ExecutionEngine/JITLink/JITLinkMemoryManager.h

@@ -14,10 +14,11 @@
 #define LLVM_EXECUTIONENGINE_JITLINK_JITLINKMEMORYMANAGER_H
 
 #include "llvm/ADT/DenseMap.h"
+#include "llvm/ExecutionEngine/JITLink/JITLinkDylib.h"
 #include "llvm/ExecutionEngine/JITSymbol.h"
 #include "llvm/Support/Error.h"
-#include "llvm/Support/Memory.h"
 #include "llvm/Support/MSVCErrorWorkarounds.h"
+#include "llvm/Support/Memory.h"
 
 #include <cstdint>
 #include <future>
@@ -93,15 +94,25 @@ public:
   virtual ~JITLinkMemoryManager();
 
   /// Create an Allocation object.
+  ///
+  /// The JD argument represents the target JITLinkDylib, and can be used by
+  /// JITLinkMemoryManager implementers to manage per-dylib allocation pools
+  /// (e.g. one pre-reserved address space slab per dylib to ensure that all
+  /// allocations for the dylib are within a certain range). The JD argument
+  /// may be null (representing an allocation not associated with any
+  /// JITDylib.
+  ///
+  /// The request argument describes the segment sizes and permisssions being
+  /// requested.
   virtual Expected<std::unique_ptr<Allocation>>
-  allocate(const SegmentsRequestMap &Request) = 0;
+  allocate(const JITLinkDylib *JD, const SegmentsRequestMap &Request) = 0;
 };
 
 /// A JITLinkMemoryManager that allocates in-process memory.
 class InProcessMemoryManager : public JITLinkMemoryManager {
 public:
   Expected<std::unique_ptr<Allocation>>
-  allocate(const SegmentsRequestMap &Request) override;
+  allocate(const JITLinkDylib *JD, const SegmentsRequestMap &Request) override;
 };
 
 } // end namespace jitlink

include/llvm/ExecutionEngine/Orc/Core.h

@@ -17,6 +17,7 @@
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/FunctionExtras.h"
 #include "llvm/ADT/IntrusiveRefCntPtr.h"
+#include "llvm/ExecutionEngine/JITLink/JITLinkDylib.h"
 #include "llvm/ExecutionEngine/JITSymbol.h"
 #include "llvm/ExecutionEngine/Orc/SymbolStringPool.h"
 #include "llvm/ExecutionEngine/OrcV1Deprecation.h"
@@ -887,7 +888,8 @@ public:
 /// their addresses may be used as keys for resource management.
 /// JITDylib state changes must be made via an ExecutionSession to guarantee
 /// that they are synchronized with respect to other JITDylib operations.
-class JITDylib : public ThreadSafeRefCountedBase<JITDylib> {
+class JITDylib : public ThreadSafeRefCountedBase<JITDylib>,
+                 public jitlink::JITLinkDylib {
   friend class AsynchronousSymbolQuery;
   friend class ExecutionSession;
   friend class Platform;

include/llvm/ExecutionEngine/Orc/OrcRPCTargetProcessControl.h

@@ -148,7 +148,8 @@ public:
   OrcRPCTPCJITLinkMemoryManager(OrcRPCTPCImplT &Parent) : Parent(Parent) {}
 
   Expected<std::unique_ptr<Allocation>>
-  allocate(const SegmentsRequestMap &Request) override {
+  allocate(const jitlink::JITLinkDylib *JD,
+           const SegmentsRequestMap &Request) override {
     orcrpctpc::ReserveMemRequest RMR;
     HostAllocMap HostAllocs;
 

include/llvm/ExecutionEngine/Orc/OrcRemoteTargetClient.h

@@ -535,7 +535,8 @@ public:
     }
 
     Expected<std::unique_ptr<Allocation>>
-    allocate(const SegmentsRequestMap &Request) override {
+    allocate(const jitlink::JITLinkDylib *JD,
+             const SegmentsRequestMap &Request) override {
       return RPCMMAlloc::Create(Client, Id, Request);
     }
 

lib/ExecutionEngine/JITLink/JITLinkGeneric.cpp

@@ -266,7 +266,8 @@ Error JITLinkerBase::allocateSegments(const SegmentLayoutMap &Layout) {
   }
   LLVM_DEBUG(dbgs() << " }\n");
 
-  if (auto AllocOrErr = Ctx->getMemoryManager().allocate(Segments))
+  if (auto AllocOrErr =
+          Ctx->getMemoryManager().allocate(Ctx->getJITLinkDylib(), Segments))
     Alloc = std::move(*AllocOrErr);
   else
     return AllocOrErr.takeError();

lib/ExecutionEngine/JITLink/JITLinkMemoryManager.cpp

@@ -17,7 +17,8 @@ JITLinkMemoryManager::~JITLinkMemoryManager() = default;
 JITLinkMemoryManager::Allocation::~Allocation() = default;
 
 Expected<std::unique_ptr<JITLinkMemoryManager::Allocation>>
-InProcessMemoryManager::allocate(const SegmentsRequestMap &Request) {
+InProcessMemoryManager::allocate(const JITLinkDylib *JD,
+                                 const SegmentsRequestMap &Request) {
 
   using AllocationMap = DenseMap<unsigned, sys::MemoryBlock>;
 

lib/ExecutionEngine/Orc/ObjectLinkingLayer.cpp

@@ -28,7 +28,8 @@ public:
       ObjectLinkingLayer &Layer,
       std::unique_ptr<MaterializationResponsibility> MR,
       std::unique_ptr<MemoryBuffer> ObjBuffer)
-      : Layer(Layer), MR(std::move(MR)), ObjBuffer(std::move(ObjBuffer)) {}
+      : JITLinkContext(&MR->getTargetJITDylib()), Layer(Layer),
+        MR(std::move(MR)), ObjBuffer(std::move(ObjBuffer)) {}
 
   ~ObjectLinkingLayerJITLinkContext() {
     // If there is an object buffer return function then use it to

lib/ExecutionEngine/Orc/TPCIndirectionUtils.cpp

@@ -109,7 +109,7 @@ Error TPCTrampolinePool::grow() {
   jitlink::JITLinkMemoryManager::SegmentsRequestMap Request;
   Request[TrampolinePagePermissions] = {PageSize, static_cast<size_t>(PageSize),
                                         0};
-  auto Alloc = TPC.getMemMgr().allocate(Request);
+  auto Alloc = TPC.getMemMgr().allocate(nullptr, Request);
 
   if (!Alloc)
     return Alloc.takeError();
@@ -294,7 +294,7 @@ TPCIndirectionUtils::writeResolverBlock(JITTargetAddress ReentryFnAddr,
   jitlink::JITLinkMemoryManager::SegmentsRequestMap Request;
   Request[ResolverBlockPermissions] = {TPC.getPageSize(),
                                        static_cast<size_t>(ResolverSize), 0};
-  auto Alloc = TPC.getMemMgr().allocate(Request);
+  auto Alloc = TPC.getMemMgr().allocate(nullptr, Request);
   if (!Alloc)
     return Alloc.takeError();
 
@@ -364,7 +364,7 @@ TPCIndirectionUtils::getIndirectStubs(unsigned NumStubs) {
     Request[StubPagePermissions] = {PageSize, static_cast<size_t>(StubBytes),
                                     0};
     Request[PointerPagePermissions] = {PageSize, 0, PointerBytes};
-    auto Alloc = TPC.getMemMgr().allocate(Request);
+    auto Alloc = TPC.getMemMgr().allocate(nullptr, Request);
     if (!Alloc)
       return Alloc.takeError();
 

tools/llvm-jitlink/llvm-jitlink.cpp

@@ -330,7 +330,7 @@ public:
   }
 
   Expected<std::unique_ptr<JITLinkMemoryManager::Allocation>>
-  allocate(const SegmentsRequestMap &Request) override {
+  allocate(const JITLinkDylib *JD, const SegmentsRequestMap &Request) override {
 
     using AllocationMap = DenseMap<unsigned, sys::MemoryBlock>;