[MCJIT][Orc] Refactor RTDyldMemoryManager, weave RuntimeDyld::SymbolInfo through

MCJIT.

This patch decouples the two responsibilities of the RTDyldMemoryManager class,
memory management and symbol resolution, into two new classes:
RuntimeDyld::MemoryManager and RuntimeDyld::SymbolResolver.

The symbol resolution interface is modified slightly, from:

  uint64_t getSymbolAddress(const std::string &Name);

to:

  RuntimeDyld::SymbolInfo findSymbol(const std::string &Name);

The latter passes symbol flags along with symbol addresses, allowing RuntimeDyld
and others to reason about non-strong/non-exported symbols.


The memory management interface removes the following method:

  void notifyObjectLoaded(ExecutionEngine *EE,
                          const object::ObjectFile &) {}

as it is not related to memory management. (Note: Backwards compatibility *is*
maintained for this method in MCJIT and OrcMCJITReplacement, see below).


The RTDyldMemoryManager class remains in-tree for backwards compatibility.
It inherits directly from RuntimeDyld::SymbolResolver, and indirectly from
RuntimeDyld::MemoryManager via the new MCJITMemoryManager class, which
just subclasses RuntimeDyld::MemoryManager and reintroduces the
notifyObjectLoaded method for backwards compatibility).

The EngineBuilder class retains the existing method:

  EngineBuilder&
  setMCJITMemoryManager(std::unique_ptr<RTDyldMemoryManager> mcjmm);

and includes two new methods:

  EngineBuilder&
  setMemoryManager(std::unique_ptr<MCJITMemoryManager> MM);

  EngineBuilder&
  setSymbolResolver(std::unique_ptr<RuntimeDyld::SymbolResolver> SR);

Clients should use EITHER:

A single call to setMCJITMemoryManager with an RTDyldMemoryManager.

OR (exclusive)

One call each to each of setMemoryManager and setSymbolResolver.

This patch should be fully compatible with existing uses of RTDyldMemoryManager.
If it is not it should be considered a bug, and the patch either fixed or
reverted.

If clients find the new API to be an improvement the goal will be to deprecate
and eventually remove the RTDyldMemoryManager class in favor of the new classes.

llvm-svn: 233509

4 files changed, 66 insertions, 46 deletions

diff --git a/llvm/examples/Kaleidoscope/Orc/fully_lazy/toy.cpp b/llvm/examples/Kaleidoscope/Orc/fully_lazy/toy.cpp
index 3b5dcd0b3e0..e3359dd0e47 100644
--- a/llvm/examples/Kaleidoscope/Orc/fully_lazy/toy.cpp
+++ b/llvm/examples/Kaleidoscope/Orc/fully_lazy/toy.cpp
@@ -1,6 +1,7 @@
 #include "llvm/Analysis/Passes.h"
 #include "llvm/ExecutionEngine/Orc/CompileUtils.h"
 #include "llvm/ExecutionEngine/Orc/IRCompileLayer.h"
+#include "llvm/ExecutionEngine/Orc/LambdaResolver.h"
 #include "llvm/ExecutionEngine/Orc/LazyEmittingLayer.h"
 #include "llvm/ExecutionEngine/Orc/ObjectLinkingLayer.h"
 #include "llvm/ExecutionEngine/Orc/OrcTargetSupport.h"
@@ -1168,11 +1169,9 @@ public:
   KaleidoscopeJIT(SessionContext &Session)
     : Session(Session),
       Mang(Session.getTarget().getDataLayout()),
-      ObjectLayer(
-        [](){ return llvm::make_unique<SectionMemoryManager>(); }),
       CompileLayer(ObjectLayer, SimpleCompiler(Session.getTarget())),
       LazyEmitLayer(CompileLayer),
-      CompileCallbacks(LazyEmitLayer, Session.getLLVMContext(),
+      CompileCallbacks(LazyEmitLayer, CCMgrMemMgr, Session.getLLVMContext(),
                        reinterpret_cast<uintptr_t>(EarthShatteringKaboom),
                        64) {}
 
@@ -1194,20 +1193,22 @@ public:
     // We need a memory manager to allocate memory and resolve symbols for this
     // new module. Create one that resolves symbols by looking back into the
     // JIT.
-    auto MM = createLookasideRTDyldMM<SectionMemoryManager>(
-                [&](const std::string &Name) {
-                  // First try to find 'Name' within the JIT.
-                  if (auto Symbol = findSymbol(Name))
-                    return Symbol.getAddress();
-
-                  // If we don't already have a definition of 'Name' then search
-                  // the ASTs.
-                  return searchFunctionASTs(Name);
-                },
-                [](const std::string &S) { return 0; } );
+    auto Resolver = createLambdaResolver(
+                      [&](const std::string &Name) {
+                        // First try to find 'Name' within the JIT.
+                        if (auto Symbol = findSymbol(Name))
+                          return RuntimeDyld::SymbolInfo(Symbol.getAddress(),
+                                                         Symbol.getFlags());
+
+                        // If we don't already have a definition of 'Name' then search
+                        // the ASTs.
+                        return searchFunctionASTs(Name);
+                      },
+                      [](const std::string &S) { return nullptr; } );
 
     return LazyEmitLayer.addModuleSet(singletonSet(std::move(M)),
-                                      std::move(MM));
+                                      make_unique<SectionMemoryManager>(),
+                                      std::move(Resolver));
   }
 
   void removeModule(ModuleHandleT H) { LazyEmitLayer.removeModuleSet(H); }
@@ -1232,7 +1233,7 @@ private:
 
   // This method searches the FunctionDefs map for a definition of 'Name'. If it
   // finds one it generates a stub for it and returns the address of the stub.
-  TargetAddress searchFunctionASTs(const std::string &Name) {
+  RuntimeDyld::SymbolInfo searchFunctionASTs(const std::string &Name) {
     auto DefI = FunctionDefs.find(Name);
     if (DefI == FunctionDefs.end())
       return 0;
@@ -1244,7 +1245,8 @@ private:
 
     // IRGen the AST, add it to the JIT, and return the address for it.
     auto H = irGenStub(std::move(FnAST));
-    return findSymbolIn(H, Name).getAddress();
+    auto Sym = findSymbolIn(H, Name);
+    return RuntimeDyld::SymbolInfo(Sym.getAddress(), Sym.getFlags());
   }
 
   // This method will take the AST for a function definition and IR-gen a stub
@@ -1261,14 +1263,16 @@ private:
     //         compile and update actions for the callback, and get a pointer to
     //         the jit trampoline that we need to call to trigger those actions.
     auto CallbackInfo =
-      CompileCallbacks.getCompileCallback(*F->getFunctionType());
+      CompileCallbacks.getCompileCallback(F->getContext());
 
     // Step 3) Create a stub that will indirectly call the body of this
     //         function once it is compiled. Initially, set the function
     //         pointer for the indirection to point at the trampoline.
     std::string BodyPtrName = (F->getName() + "$address").str();
     GlobalVariable *FunctionBodyPointer =
-      createImplPointer(*F, BodyPtrName, CallbackInfo.getAddress());
+      createImplPointer(*F, BodyPtrName,
+                        createIRTypedAddress(*F->getFunctionType(),
+                                             CallbackInfo.getAddress()));
     makeStub(*F, *FunctionBodyPointer);
 
     // Step 4) Add the module containing the stub to the JIT.
@@ -1297,6 +1301,7 @@ private:
 
   SessionContext &Session;
   Mangler Mang;
+  SectionMemoryManager CCMgrMemMgr;
   ObjLayerT ObjectLayer;
   CompileLayerT CompileLayer;
   LazyEmitLayerT LazyEmitLayer;
diff --git a/llvm/examples/Kaleidoscope/Orc/initial/toy.cpp b/llvm/examples/Kaleidoscope/Orc/initial/toy.cpp
index cc6fb8e0a01..bf43f2952c7 100644
--- a/llvm/examples/Kaleidoscope/Orc/initial/toy.cpp
+++ b/llvm/examples/Kaleidoscope/Orc/initial/toy.cpp
@@ -1,6 +1,7 @@
 #include "llvm/Analysis/Passes.h"
 #include "llvm/ExecutionEngine/Orc/CompileUtils.h"
 #include "llvm/ExecutionEngine/Orc/IRCompileLayer.h"
+#include "llvm/ExecutionEngine/Orc/LambdaResolver.h"
 #include "llvm/ExecutionEngine/Orc/LazyEmittingLayer.h"
 #include "llvm/ExecutionEngine/Orc/ObjectLinkingLayer.h"
 #include "llvm/IR/DataLayout.h"
@@ -1175,13 +1176,18 @@ public:
     // We need a memory manager to allocate memory and resolve symbols for this
     // new module. Create one that resolves symbols by looking back into the
     // JIT.
-    auto MM = createLookasideRTDyldMM<SectionMemoryManager>(
-                [&](const std::string &Name) {
-                  return findSymbol(Name).getAddress();
-                },
-                [](const std::string &S) { return 0; } );
-
-    return CompileLayer.addModuleSet(singletonSet(std::move(M)), std::move(MM));
+    auto Resolver = createLambdaResolver(
+                      [&](const std::string &Name) {
+                        if (auto Sym = findSymbol(Name))
+                          return RuntimeDyld::SymbolInfo(Sym.getAddress(),
+                                                         Sym.getFlags());
+                        return RuntimeDyld::SymbolInfo(nullptr);
+                      },
+                      [](const std::string &S) { return nullptr; }
+                    );
+    return CompileLayer.addModuleSet(singletonSet(std::move(M)),
+                                     make_unique<SectionMemoryManager>(),
+                                     std::move(Resolver));
   }
 
   void removeModule(ModuleHandleT H) { CompileLayer.removeModuleSet(H); }
diff --git a/llvm/examples/Kaleidoscope/Orc/lazy_codegen/toy.cpp b/llvm/examples/Kaleidoscope/Orc/lazy_codegen/toy.cpp
index 6e2ec2723a4..1369ba6f5ee 100644
--- a/llvm/examples/Kaleidoscope/Orc/lazy_codegen/toy.cpp
+++ b/llvm/examples/Kaleidoscope/Orc/lazy_codegen/toy.cpp
@@ -1,6 +1,7 @@
 #include "llvm/Analysis/Passes.h"
 #include "llvm/ExecutionEngine/Orc/CompileUtils.h"
 #include "llvm/ExecutionEngine/Orc/IRCompileLayer.h"
+#include "llvm/ExecutionEngine/Orc/LambdaResolver.h"
 #include "llvm/ExecutionEngine/Orc/LazyEmittingLayer.h"
 #include "llvm/ExecutionEngine/Orc/ObjectLinkingLayer.h"
 #include "llvm/IR/DataLayout.h"
@@ -1178,14 +1179,18 @@ public:
     // We need a memory manager to allocate memory and resolve symbols for this
     // new module. Create one that resolves symbols by looking back into the
     // JIT.
-    auto MM = createLookasideRTDyldMM<SectionMemoryManager>(
-                [&](const std::string &Name) {
-                  return findSymbol(Name).getAddress();
-                },
-                [](const std::string &S) { return 0; } );
+    auto Resolver = createLambdaResolver(
+                      [&](const std::string &Name) {
+                        if (auto Sym = findSymbol(Name))
+                          return RuntimeDyld::SymbolInfo(Sym.getAddress(),
+                                                         Sym.getFlags());
+                        return RuntimeDyld::SymbolInfo(nullptr);
+                      },
+                      [](const std::string &S) { return nullptr; } );
 
     return LazyEmitLayer.addModuleSet(singletonSet(std::move(M)),
-                                      std::move(MM));
+                                      make_unique<SectionMemoryManager>(),
+                                      std::move(Resolver));
   }
 
   void removeModule(ModuleHandleT H) { LazyEmitLayer.removeModuleSet(H); }
diff --git a/llvm/examples/Kaleidoscope/Orc/lazy_irgen/toy.cpp b/llvm/examples/Kaleidoscope/Orc/lazy_irgen/toy.cpp
index 19801e12030..c489a450d79 100644
--- a/llvm/examples/Kaleidoscope/Orc/lazy_irgen/toy.cpp
+++ b/llvm/examples/Kaleidoscope/Orc/lazy_irgen/toy.cpp
@@ -1,6 +1,7 @@
 #include "llvm/Analysis/Passes.h"
 #include "llvm/ExecutionEngine/Orc/CompileUtils.h"
 #include "llvm/ExecutionEngine/Orc/IRCompileLayer.h"
+#include "llvm/ExecutionEngine/Orc/LambdaResolver.h"
 #include "llvm/ExecutionEngine/Orc/LazyEmittingLayer.h"
 #include "llvm/ExecutionEngine/Orc/ObjectLinkingLayer.h"
 #include "llvm/IR/DataLayout.h"
@@ -1183,20 +1184,22 @@ public:
     // We need a memory manager to allocate memory and resolve symbols for this
     // new module. Create one that resolves symbols by looking back into the
     // JIT.
-    auto MM = createLookasideRTDyldMM<SectionMemoryManager>(
-                [&](const std::string &Name) {
-                  // First try to find 'Name' within the JIT.
-                  if (auto Symbol = findSymbol(Name))
-                    return Symbol.getAddress();
-
-                  // If we don't already have a definition of 'Name' then search
-                  // the ASTs.
-                  return searchFunctionASTs(Name);
-                },
-                [](const std::string &S) { return 0; } );
+    auto Resolver = createLambdaResolver(
+                      [&](const std::string &Name) {
+                        // First try to find 'Name' within the JIT.
+                        if (auto Symbol = findSymbol(Name))
+                          return RuntimeDyld::SymbolInfo(Symbol.getAddress(),
+                                                         Symbol.getFlags());
+
+                        // If we don't already have a definition of 'Name' then search
+                        // the ASTs.
+                        return searchFunctionASTs(Name);
+                      },
+                      [](const std::string &S) { return nullptr; } );
 
     return LazyEmitLayer.addModuleSet(singletonSet(std::move(M)),
-                                      std::move(MM));
+                                      make_unique<SectionMemoryManager>(),
+                                      std::move(Resolver));
   }
 
   void removeModule(ModuleHandleT H) { LazyEmitLayer.removeModuleSet(H); }
@@ -1217,7 +1220,7 @@ private:
 
   // This method searches the FunctionDefs map for a definition of 'Name'. If it
   // finds one it generates a stub for it and returns the address of the stub.
-  TargetAddress searchFunctionASTs(const std::string &Name) {
+  RuntimeDyld::SymbolInfo searchFunctionASTs(const std::string &Name) {
     auto DefI = FunctionDefs.find(Name);
     if (DefI == FunctionDefs.end())
       return 0;
@@ -1228,7 +1231,8 @@ private:
 
     // IRGen the AST, add it to the JIT, and return the address for it.
     auto H = addModule(IRGen(Session, *FnAST));
-    return findSymbolIn(H, Name).getAddress();
+    auto Sym = findSymbolIn(H, Name);
+    return RuntimeDyld::SymbolInfo(Sym.getAddress(), Sym.getFlags());
   }
 
   SessionContext &Session;