[Orc] Enable user-supplied memory managers in the CompileOnDemand layer.

Previously the CompileOnDemand layer was hard-coded to use a new
SectionMemoryManager for each function when it was called.

llvm-svn: 257265

3 files changed, 57 insertions, 39 deletions

diff --git a/llvm/include/llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h b/llvm/include/llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h
index 7dab5d1bc67..0eb8693df42 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h
@@ -19,7 +19,6 @@
 #include "LambdaResolver.h"
 #include "LogicalDylib.h"
 #include "llvm/ADT/STLExtras.h"
-#include "llvm/ExecutionEngine/SectionMemoryManager.h"
 #include "llvm/Transforms/Utils/Cloning.h"
 #include <list>
 #include <memory>
@@ -61,31 +60,36 @@ private:
 
   typedef typename BaseLayerT::ModuleSetHandleT BaseLayerModuleSetHandleT;
 
-  class ModuleOwner {
+  // Provide type-erasure for the Modules and MemoryManagers.
+  template <typename ResourceT>
+  class ResourceOwner {
   public:
-    ModuleOwner() = default;
-    ModuleOwner(const ModuleOwner&) = delete;
-    ModuleOwner& operator=(const ModuleOwner&) = delete;
-    virtual ~ModuleOwner() { }
-    virtual Module& getModule() const = 0;
+    ResourceOwner() = default;
+    ResourceOwner(const ResourceOwner&) = delete;
+    ResourceOwner& operator=(const ResourceOwner&) = delete;
+    virtual ~ResourceOwner() { }
+    virtual ResourceT& getResource() const = 0;
   };
 
-  template <typename ModulePtrT>
-  class ModuleOwnerImpl : public ModuleOwner {
+  template <typename ResourceT, typename ResourcePtrT>
+  class ResourceOwnerImpl : public ResourceOwner<ResourceT> {
   public:
-    ModuleOwnerImpl(ModulePtrT ModulePtr) : ModulePtr(std::move(ModulePtr)) {}
-    Module& getModule() const override { return *ModulePtr; }
+    ResourceOwnerImpl(ResourcePtrT ResourcePtr)
+      : ResourcePtr(std::move(ResourcePtr)) {}
+    ResourceT& getResource() const override { return *ResourcePtr; }
   private:
-    ModulePtrT ModulePtr;
+    ResourcePtrT ResourcePtr;
   };
 
-  template <typename ModulePtrT>
-  std::unique_ptr<ModuleOwner> wrapOwnership(ModulePtrT ModulePtr) {
-    return llvm::make_unique<ModuleOwnerImpl<ModulePtrT>>(std::move(ModulePtr));
+  template <typename ResourceT, typename ResourcePtrT>
+  std::unique_ptr<ResourceOwner<ResourceT>>
+  wrapOwnership(ResourcePtrT ResourcePtr) {
+    typedef ResourceOwnerImpl<ResourceT, ResourcePtrT> RO;
+    return llvm::make_unique<RO>(std::move(ResourcePtr));
   }
 
   struct LogicalModuleResources {
-    std::unique_ptr<ModuleOwner> SourceModuleOwner;
+    std::unique_ptr<ResourceOwner<Module>> SourceModule;
     std::set<const Function*> StubsToClone;
     std::unique_ptr<IndirectStubsMgrT> StubsMgr;
 
@@ -93,13 +97,13 @@ private:
 
     // Explicit move constructor to make MSVC happy.
     LogicalModuleResources(LogicalModuleResources &&Other)
-        : SourceModuleOwner(std::move(Other.SourceModuleOwner)),
+        : SourceModule(std::move(Other.SourceModule)),
           StubsToClone(std::move(Other.StubsToClone)),
           StubsMgr(std::move(Other.StubsMgr)) {}
 
     // Explicit move assignment to make MSVC happy.
     LogicalModuleResources& operator=(LogicalModuleResources &&Other) {
-      SourceModuleOwner = std::move(Other.SourceModuleOwner);
+      SourceModule = std::move(Other.SourceModule);
       StubsToClone = std::move(Other.StubsToClone);
       StubsMgr = std::move(Other.StubsMgr);
     }
@@ -114,12 +118,19 @@ private:
 
   };
 
-
-
   struct LogicalDylibResources {
     typedef std::function<RuntimeDyld::SymbolInfo(const std::string&)>
       SymbolResolverFtor;
+
+    typedef std::function<typename BaseLayerT::ModuleSetHandleT(
+                            BaseLayerT&,
+                            std::unique_ptr<Module>,
+                            std::unique_ptr<RuntimeDyld::SymbolResolver>)>
+      ModuleAdderFtor;
+
     SymbolResolverFtor ExternalSymbolResolver;
+    std::unique_ptr<ResourceOwner<RuntimeDyld::MemoryManager>> MemMgr;
+    ModuleAdderFtor ModuleAdder;
   };
 
   typedef LogicalDylib<BaseLayerT, LogicalModuleResources,
@@ -157,9 +168,6 @@ public:
                                 MemoryManagerPtrT MemMgr,
                                 SymbolResolverPtrT Resolver) {
 
-    assert(MemMgr == nullptr &&
-           "User supplied memory managers not supported with COD yet.");
-
     LogicalDylibs.push_back(CODLogicalDylib(BaseLayer));
     auto &LDResources = LogicalDylibs.back().getDylibResources();
 
@@ -168,6 +176,18 @@ public:
         return Resolver->findSymbol(Name);
       };
 
+    auto &MemMgrRef = *MemMgr;
+    LDResources.MemMgr =
+      wrapOwnership<RuntimeDyld::MemoryManager>(std::move(MemMgr));
+
+    LDResources.ModuleAdder =
+      [&MemMgrRef](BaseLayerT &B, std::unique_ptr<Module> M,
+                   std::unique_ptr<RuntimeDyld::SymbolResolver> R) {
+        std::vector<std::unique_ptr<Module>> Ms;
+        Ms.push_back(std::move(M));
+        return B.addModuleSet(std::move(Ms), &MemMgrRef, std::move(R));
+      };
+
     // Process each of the modules in this module set.
     for (auto &M : Ms)
       addLogicalModule(LogicalDylibs.back(), std::move(M));
@@ -215,9 +235,9 @@ private:
     auto LMH = LD.createLogicalModule();
     auto &LMResources =  LD.getLogicalModuleResources(LMH);
 
-    LMResources.SourceModuleOwner = wrapOwnership(std::move(SrcMPtr));
+    LMResources.SourceModule = wrapOwnership<Module>(std::move(SrcMPtr));
 
-    Module &SrcM = LMResources.SourceModuleOwner->getModule();
+    Module &SrcM = LMResources.SourceModule->getResource();
 
     // Create the GlobalValues module.
     const DataLayout &DL = SrcM.getDataLayout();
@@ -326,12 +346,9 @@ private:
           return RuntimeDyld::SymbolInfo(nullptr);
         });
 
-    std::vector<std::unique_ptr<Module>> GVsMSet;
-    GVsMSet.push_back(std::move(GVsM));
     auto GVsH =
-      BaseLayer.addModuleSet(std::move(GVsMSet),
-                             llvm::make_unique<SectionMemoryManager>(),
-                             std::move(GVsResolver));
+      LD.getDylibResources().ModuleAdder(BaseLayer, std::move(GVsM),
+				         std::move(GVsResolver));
     LD.addToLogicalModule(LMH, GVsH);
   }
 
@@ -348,7 +365,7 @@ private:
                                   LogicalModuleHandle LMH,
                                   Function &F) {
     auto &LMResources = LD.getLogicalModuleResources(LMH);
-    Module &SrcM = LMResources.SourceModuleOwner->getModule();
+    Module &SrcM = LMResources.SourceModule->getResource();
 
     // If F is a declaration we must already have compiled it.
     if (F.isDeclaration())
@@ -386,7 +403,7 @@ private:
                                           LogicalModuleHandle LMH,
                                           const PartitionT &Part) {
     auto &LMResources = LD.getLogicalModuleResources(LMH);
-    Module &SrcM = LMResources.SourceModuleOwner->getModule();
+    Module &SrcM = LMResources.SourceModule->getResource();
 
     // Create the module.
     std::string NewName = SrcM.getName();
@@ -445,7 +462,6 @@ private:
       moveFunctionBody(*F, VMap, &Materializer);
 
     // Create memory manager and symbol resolver.
-    auto MemMgr = llvm::make_unique<SectionMemoryManager>();
     auto Resolver = createLambdaResolver(
         [this, &LD, LMH](const std::string &Name) {
           if (auto Symbol = LD.findSymbolInternally(LMH, Name))
@@ -459,10 +475,9 @@ private:
                                            Symbol.getFlags());
           return RuntimeDyld::SymbolInfo(nullptr);
         });
-    std::vector<std::unique_ptr<Module>> PartMSet;
-    PartMSet.push_back(std::move(M));
-    return BaseLayer.addModuleSet(std::move(PartMSet), std::move(MemMgr),
-                                  std::move(Resolver));
+
+    return LD.getDylibResources().ModuleAdder(BaseLayer, std::move(M),
+					      std::move(Resolver));
   }
 
   BaseLayerT &BaseLayer;
diff --git a/llvm/lib/ExecutionEngine/Orc/OrcCBindingsStack.h b/llvm/lib/ExecutionEngine/Orc/OrcCBindingsStack.h
index 2e17624ff47..aae6a99432b 100644
--- a/llvm/lib/ExecutionEngine/Orc/OrcCBindingsStack.h
+++ b/llvm/lib/ExecutionEngine/Orc/OrcCBindingsStack.h
@@ -221,7 +221,8 @@ public:
   ModuleHandleT addIRModuleLazy(Module* M,
                                 LLVMOrcSymbolResolverFn ExternalResolver,
                                 void *ExternalResolverCtx) {
-    return addIRModule(CODLayer, std::move(M), nullptr,
+    return addIRModule(CODLayer, std::move(M),
+		       llvm::make_unique<SectionMemoryManager>(),
                        std::move(ExternalResolver), ExternalResolverCtx);
   }
 
diff --git a/llvm/tools/lli/OrcLazyJIT.h b/llvm/tools/lli/OrcLazyJIT.h
index bb4da33ea9b..2a5b31d1bfb 100644
--- a/llvm/tools/lli/OrcLazyJIT.h
+++ b/llvm/tools/lli/OrcLazyJIT.h
@@ -105,7 +105,9 @@ public:
     // Add the module to the JIT.
     std::vector<std::unique_ptr<Module>> S;
     S.push_back(std::move(M));
-    auto H = CODLayer.addModuleSet(std::move(S), nullptr, std::move(Resolver));
+    auto H = CODLayer.addModuleSet(std::move(S),
+				   llvm::make_unique<SectionMemoryManager>(),
+				   std::move(Resolver));
 
     // Run the static constructors, and save the static destructor runner for
     // execution when the JIT is torn down.