[ORC] Enable use of TargetProcessControl::getMemMgr with ObjectLinkin…

…gLayer.

This patch makes ownership of the JITLinkMemoryManager by ObjectLinkingLayer
optional: the layer can still own the memory manager but no longer has to.

Evevntually we want to move to a state where ObjectLinkingLayer never owns its
memory manager. For now allowing optional ownership makes it easier to develop
classes that can dynamically use either RTDyldObjectLinkingLayer, which owns
its memory managers, or ObjectLinkingLayer (e.g. LLJIT).

llvm/examples/OrcV2Examples/LLJITWithTargetProcessControl/LLJITWithTargetProcessControl.cpp

@@ -134,7 +134,8 @@ int main(int argc, char *argv[]) {
   ExitOnErr.setBanner(std::string(argv[0]) + ": ");
 
   // (1) Create LLJIT instance.
-  auto J = ExitOnErr(LLJITBuilder().create());
+  auto TPC = ExitOnErr(SelfTargetProcessControl::Create());
+  auto J = ExitOnErr(LLJITBuilder().setTargetProcessControl(*TPC).create());
 
   // (2) Install transform to print modules as they are compiled:
   J->getIRTransformLayer().setTransform(
@@ -145,8 +146,6 @@ int main(int argc, char *argv[]) {
       });
 
   // (3) Create stubs and call-through managers:
-
-  auto TPC = ExitOnErr(SelfTargetProcessControl::Create());
   auto TPCIU = ExitOnErr(TPCIndirectionUtils::Create(*TPC));
   ExitOnErr(TPCIU->writeResolverBlock(pointerToJITTargetAddress(&reenter),
                                       pointerToJITTargetAddress(TPCIU.get())));

llvm/include/llvm/ExecutionEngine/Orc/LLJIT.h

@@ -29,6 +29,7 @@ namespace orc {
 
 class LLJITBuilderState;
 class LLLazyJITBuilderState;
+class TargetProcessControl;
 
 /// A pre-fabricated ORC JIT stack that can serve as an alternative to MCJIT.
 ///
@@ -272,6 +273,7 @@ class LLJITBuilderState {
   CompileFunctionCreator CreateCompileFunction;
   PlatformSetupFunction SetUpPlatform;
   unsigned NumCompileThreads = 0;
+  TargetProcessControl *TPC = nullptr;
 
   /// Called prior to JIT class construcion to fix up defaults.
   Error prepareForConstruction();
@@ -354,6 +356,17 @@ class LLJITBuilderSetters {
     return impl();
   }
 
+  /// Set a TargetProcessControl object.
+  ///
+  /// If the platform uses ObjectLinkingLayer by default and no
+  /// ObjectLinkingLayerCreator has been set then the TargetProcessControl
+  /// object will be used to supply the memory manager for the
+  /// ObjectLinkingLayer.
+  SetterImpl &setTargetProcessControl(TargetProcessControl &TPC) {
+    impl().TPC = &TPC;
+    return impl();
+  }
+
   /// Create an instance of the JIT.
   Expected<std::unique_ptr<JITType>> create() {
     if (auto Err = impl().prepareForConstruction())

llvm/include/llvm/ExecutionEngine/Orc/ObjectLinkingLayer.h

@@ -90,8 +90,14 @@ class ObjectLinkingLayer : public ObjectLayer {
   using ReturnObjectBufferFunction =
       std::function<void(std::unique_ptr<MemoryBuffer>)>;
 
-  /// Construct an ObjectLinkingLayer with the given NotifyLoaded,
-  /// and NotifyEmitted functors.
+  /// Construct an ObjectLinkingLayer.
+  ObjectLinkingLayer(ExecutionSession &ES,
+                     jitlink::JITLinkMemoryManager &MemMgr);
+
+  /// Construct an ObjectLinkingLayer. Takes ownership of the given
+  /// JITLinkMemoryManager. This method is a temporary hack to simplify
+  /// co-existence with RTDyldObjectLinkingLayer (which also owns its
+  /// allocators).
   ObjectLinkingLayer(ExecutionSession &ES,
                      std::unique_ptr<jitlink::JITLinkMemoryManager> MemMgr);
 
@@ -159,7 +165,8 @@ class ObjectLinkingLayer : public ObjectLayer {
   Error removeAllModules();
 
   mutable std::mutex LayerMutex;
-  std::unique_ptr<jitlink::JITLinkMemoryManager> MemMgr;
+  jitlink::JITLinkMemoryManager &MemMgr;
+  std::unique_ptr<jitlink::JITLinkMemoryManager> MemMgrOwnership;
   bool OverrideObjectFlags = false;
   bool AutoClaimObjectSymbols = false;
   ReturnObjectBufferFunction ReturnObjectBuffer;

llvm/lib/ExecutionEngine/Orc/LLJIT.cpp

@@ -13,6 +13,7 @@
 #include "llvm/ExecutionEngine/Orc/ObjectLinkingLayer.h"
 #include "llvm/ExecutionEngine/Orc/OrcError.h"
 #include "llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h"
+#include "llvm/ExecutionEngine/Orc/TargetProcessControl.h"
 #include "llvm/ExecutionEngine/SectionMemoryManager.h"
 #include "llvm/IR/GlobalVariable.h"
 #include "llvm/IR/IRBuilder.h"
@@ -973,10 +974,15 @@ Error LLJITBuilderState::prepareForConstruction() {
       JTMB->setRelocationModel(Reloc::PIC_);
       JTMB->setCodeModel(CodeModel::Small);
       CreateObjectLinkingLayer =
-          [](ExecutionSession &ES,
-             const Triple &) -> std::unique_ptr<ObjectLayer> {
-        auto ObjLinkingLayer = std::make_unique<ObjectLinkingLayer>(
-            ES, std::make_unique<jitlink::InProcessMemoryManager>());
+          [TPC = this->TPC](ExecutionSession &ES,
+                            const Triple &) -> std::unique_ptr<ObjectLayer> {
+        std::unique_ptr<ObjectLinkingLayer> ObjLinkingLayer;
+        if (TPC)
+          ObjLinkingLayer =
+              std::make_unique<ObjectLinkingLayer>(ES, TPC->getMemMgr());
+        else
+          ObjLinkingLayer = std::make_unique<ObjectLinkingLayer>(
+              ES, std::make_unique<jitlink::InProcessMemoryManager>());
         ObjLinkingLayer->addPlugin(std::make_unique<EHFrameRegistrationPlugin>(
             jitlink::InProcessEHFrameRegistrar::getInstance()));
         return std::move(ObjLinkingLayer);

llvm/lib/ExecutionEngine/Orc/ObjectLinkingLayer.cpp

@@ -36,7 +36,7 @@ class ObjectLinkingLayerJITLinkContext final : public JITLinkContext {
       Layer.ReturnObjectBuffer(std::move(ObjBuffer));
   }
 
-  JITLinkMemoryManager &getMemoryManager() override { return *Layer.MemMgr; }
+  JITLinkMemoryManager &getMemoryManager() override { return Layer.MemMgr; }
 
   MemoryBufferRef getObjectBuffer() const override {
     return ObjBuffer->getMemBufferRef();
@@ -447,9 +447,13 @@ class ObjectLinkingLayerJITLinkContext final : public JITLinkContext {
 
 ObjectLinkingLayer::Plugin::~Plugin() {}
 
+ObjectLinkingLayer::ObjectLinkingLayer(ExecutionSession &ES,
+                                       JITLinkMemoryManager &MemMgr)
+    : ObjectLayer(ES), MemMgr(MemMgr) {}
+
 ObjectLinkingLayer::ObjectLinkingLayer(
     ExecutionSession &ES, std::unique_ptr<JITLinkMemoryManager> MemMgr)
-    : ObjectLayer(ES), MemMgr(std::move(MemMgr)) {}
+    : ObjectLayer(ES), MemMgr(*MemMgr), MemMgrOwnership(std::move(MemMgr)) {}
 
 ObjectLinkingLayer::~ObjectLinkingLayer() {
   if (auto Err = removeAllModules())