[ORC] Update CompileOnDemandLayer2 to use the new lazyReexports mechanism

for lazy compilation, rather than a callback manager.

The new mechanism does not block compile threads, and does not require
function bodies to be renamed.

Future modifications should allow laziness on a per-module basis to work
without any modification of the input module.

llvm-svn: 343065

7 files changed, 82 insertions, 93 deletions

diff --git a/llvm/include/llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h b/llvm/include/llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h
index 592f14843e9..0c048e4c251 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h
@@ -23,6 +23,7 @@
 #include "llvm/ExecutionEngine/Orc/IndirectionUtils.h"
 #include "llvm/ExecutionEngine/Orc/LambdaResolver.h"
 #include "llvm/ExecutionEngine/Orc/Layer.h"
+#include "llvm/ExecutionEngine/Orc/LazyReexports.h"
 #include "llvm/ExecutionEngine/Orc/Legacy.h"
 #include "llvm/ExecutionEngine/Orc/OrcError.h"
 #include "llvm/ExecutionEngine/RuntimeDyld.h"
@@ -69,7 +70,7 @@ public:
       std::function<std::unique_ptr<IndirectStubsManager>()>;
 
   CompileOnDemandLayer2(ExecutionSession &ES, IRLayer &BaseLayer,
-                        JITCompileCallbackManager &CCMgr,
+                        LazyCallThroughManager &LCTMgr,
                         IndirectStubsManagerBuilder BuildIndirectStubsManager);
 
   Error add(JITDylib &V, VModuleKey K, ThreadSafeModule TSM) override;
@@ -78,10 +79,22 @@ public:
             ThreadSafeModule TSM) override;
 
 private:
-  using StubManagersMap =
-      std::map<const JITDylib *, std::unique_ptr<IndirectStubsManager>>;
+  struct PerDylibResources {
+  public:
+    PerDylibResources(JITDylib &ImplD,
+                      std::unique_ptr<IndirectStubsManager> ISMgr)
+        : ImplD(ImplD), ISMgr(std::move(ISMgr)) {}
+    JITDylib &getImplDylib() { return ImplD; }
+    IndirectStubsManager &getISManager() { return *ISMgr; }
+
+  private:
+    JITDylib &ImplD;
+    std::unique_ptr<IndirectStubsManager> ISMgr;
+  };
+
+  using PerDylibResourcesMap = std::map<const JITDylib *, PerDylibResources>;
 
-  IndirectStubsManager &getStubsManager(const JITDylib &JD);
+  PerDylibResources &getPerDylibResources(JITDylib &TargetD);
 
   void emitExtractedFunctionsModule(MaterializationResponsibility R,
                                     ThreadSafeModule TSM);
@@ -89,9 +102,9 @@ private:
   mutable std::mutex CODLayerMutex;
 
   IRLayer &BaseLayer;
-  JITCompileCallbackManager &CCMgr;
+  LazyCallThroughManager &LCTMgr;
   IndirectStubsManagerBuilder BuildIndirectStubsManager;
-  StubManagersMap StubsMgrs;
+  PerDylibResourcesMap DylibResources;
 };
 
 /// Compile-on-demand layer.
diff --git a/llvm/include/llvm/ExecutionEngine/Orc/Core.h b/llvm/include/llvm/ExecutionEngine/Orc/Core.h
index c7a925ced27..f75705732e2 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/Core.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/Core.h
@@ -745,6 +745,9 @@ public:
     DispatchMaterialization(JD, std::move(MU));
   }
 
+  /// Dump the state of all the JITDylibs in this session.
+  void dump(raw_ostream &OS);
+
 private:
   static void logErrorsToStdErr(Error Err) {
     logAllUnhandledErrors(std::move(Err), errs(), "JIT session error: ");
diff --git a/llvm/include/llvm/ExecutionEngine/Orc/LLJIT.h b/llvm/include/llvm/ExecutionEngine/Orc/LLJIT.h
index 4c840dae69f..61c5cd29de7 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/LLJIT.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/LLJIT.h
@@ -159,17 +159,16 @@ private:
   // Create a single-threaded LLLazyJIT instance.
   LLLazyJIT(std::unique_ptr<ExecutionSession> ES,
             std::unique_ptr<TargetMachine> TM, DataLayout DL,
-            std::unique_ptr<JITCompileCallbackManager> CCMgr,
+            std::unique_ptr<LazyCallThroughManager> LCTMgr,
             std::function<std::unique_ptr<IndirectStubsManager>()> ISMBuilder);
 
   // Create a multi-threaded LLLazyJIT instance.
-  LLLazyJIT(std::unique_ptr<ExecutionSession> ES,
-            JITTargetMachineBuilder JTMB, DataLayout DL,
-            unsigned NumCompileThreads,
-            std::unique_ptr<JITCompileCallbackManager> CCMgr,
+  LLLazyJIT(std::unique_ptr<ExecutionSession> ES, JITTargetMachineBuilder JTMB,
+            DataLayout DL, unsigned NumCompileThreads,
+            std::unique_ptr<LazyCallThroughManager> LCTMgr,
             std::function<std::unique_ptr<IndirectStubsManager>()> ISMBuilder);
 
-  std::unique_ptr<JITCompileCallbackManager> CCMgr;
+  std::unique_ptr<LazyCallThroughManager> LCTMgr;
   std::function<std::unique_ptr<IndirectStubsManager>()> ISMBuilder;
 
   IRTransformLayer2 TransformLayer;
diff --git a/llvm/lib/ExecutionEngine/Orc/CompileOnDemandLayer.cpp b/llvm/lib/ExecutionEngine/Orc/CompileOnDemandLayer.cpp
index a68848f2f30..b0ae6c0a777 100644
--- a/llvm/lib/ExecutionEngine/Orc/CompileOnDemandLayer.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/CompileOnDemandLayer.cpp
@@ -186,9 +186,9 @@ private:
 };
 
 CompileOnDemandLayer2::CompileOnDemandLayer2(
-    ExecutionSession &ES, IRLayer &BaseLayer, JITCompileCallbackManager &CCMgr,
+    ExecutionSession &ES, IRLayer &BaseLayer, LazyCallThroughManager &LCTMgr,
     IndirectStubsManagerBuilder BuildIndirectStubsManager)
-    : IRLayer(ES), BaseLayer(BaseLayer), CCMgr(CCMgr),
+    : IRLayer(ES), BaseLayer(BaseLayer), LCTMgr(LCTMgr),
       BuildIndirectStubsManager(std::move(BuildIndirectStubsManager)) {}
 
 Error CompileOnDemandLayer2::add(JITDylib &V, VModuleKey K,
@@ -212,12 +212,15 @@ void CompileOnDemandLayer2::emit(MaterializationResponsibility R, VModuleKey K,
 
   auto GlobalsModule = extractGlobals(TSM);
 
-  // Delete the bodies of any available externally functions, rename the
-  // rest, and build the compile callbacks.
+  // Delete the bodies of any available externally functions and build the
+  // lazy reexports alias map.
   std::map<SymbolStringPtr, std::pair<JITTargetAddress, JITSymbolFlags>>
       StubCallbacksAndLinkages;
   auto &TargetJD = R.getTargetJITDylib();
+  auto &Resources = getPerDylibResources(TargetJD);
+  auto &ImplD = Resources.getImplDylib();
 
+  SymbolAliasMap LazyReexports;
   for (auto &F : M.functions()) {
     if (F.isDeclaration())
       continue;
@@ -228,81 +231,44 @@ void CompileOnDemandLayer2::emit(MaterializationResponsibility R, VModuleKey K,
       continue;
     }
 
-    assert(F.hasName() && "Function should have a name");
-    std::string StubUnmangledName = F.getName();
-    F.setName(F.getName() + "$body");
-    auto StubDecl = cloneFunctionDecl(*TSM.getModule(), F);
-    StubDecl->setName(StubUnmangledName);
-    StubDecl->setPersonalityFn(nullptr);
-    StubDecl->setLinkage(GlobalValue::ExternalLinkage);
-    F.replaceAllUsesWith(StubDecl);
-
-    auto StubName = Mangle(StubUnmangledName);
-    auto BodyName = Mangle(F.getName());
-    if (auto CallbackAddr = CCMgr.getCompileCallback(
-            [BodyName, &TargetJD, &ES]() -> JITTargetAddress {
-              if (auto Sym = lookup({&TargetJD}, BodyName))
-                return Sym->getAddress();
-              else {
-                ES.reportError(Sym.takeError());
-                return 0;
-              }
-            })) {
-      auto Flags = JITSymbolFlags::fromGlobalValue(F);
-      Flags &= ~JITSymbolFlags::Weak;
-      StubCallbacksAndLinkages[std::move(StubName)] =
-          std::make_pair(*CallbackAddr, Flags);
-    } else {
-      ES.reportError(CallbackAddr.takeError());
-      R.failMaterialization();
-      return;
-    }
-  }
-
-  // Build the stub inits map.
-  IndirectStubsManager::StubInitsMap StubInits;
-  for (auto &KV : StubCallbacksAndLinkages)
-    StubInits[*KV.first] = KV.second;
+    auto Flags = JITSymbolFlags::fromGlobalValue(F);
+    assert(Flags.isCallable() && "Non-callable definition in functions module");
 
-  // Build the function-body-extracting materialization unit.
-  if (auto Err = R.getTargetJITDylib().define(
-          llvm::make_unique<ExtractingIRMaterializationUnit>(ES, *this,
-                                                             std::move(TSM)))) {
-    ES.reportError(std::move(Err));
-    R.failMaterialization();
-    return;
+    auto MangledName = Mangle(F.getName());
+    LazyReexports[MangledName] = SymbolAliasMapEntry(MangledName, Flags);
   }
 
-  // Build the stubs.
-  // FIXME: Remove function bodies materialization unit if stub creation fails.
-  auto &StubsMgr = getStubsManager(TargetJD);
-  if (auto Err = StubsMgr.createStubs(StubInits)) {
+  // Add the functions module to the implementation dylib using an extracting
+  // materialization unit.
+  if (auto Err =
+          ImplD.define(llvm::make_unique<ExtractingIRMaterializationUnit>(
+              ES, *this, std::move(TSM)))) {
     ES.reportError(std::move(Err));
     R.failMaterialization();
     return;
   }
 
-  // Resolve and finalize stubs.
-  SymbolMap ResolvedStubs;
-  for (auto &KV : StubCallbacksAndLinkages) {
-    if (auto Sym = StubsMgr.findStub(*KV.first, false))
-      ResolvedStubs[KV.first] = Sym;
-    else
-      llvm_unreachable("Stub went missing");
-  }
-
-  R.resolve(ResolvedStubs);
+  // Handle responsibility for function symbols by returning lazy reexports.
+  auto &ISMgr = Resources.getISManager();
+  R.replace(lazyReexports(LCTMgr, ISMgr, ImplD, LazyReexports));
 
   BaseLayer.emit(std::move(R), std::move(K), std::move(GlobalsModule));
 }
 
-IndirectStubsManager &
-CompileOnDemandLayer2::getStubsManager(const JITDylib &V) {
-  std::lock_guard<std::mutex> Lock(CODLayerMutex);
-  StubManagersMap::iterator I = StubsMgrs.find(&V);
-  if (I == StubsMgrs.end())
-    I = StubsMgrs.insert(std::make_pair(&V, BuildIndirectStubsManager())).first;
-  return *I->second;
+CompileOnDemandLayer2::PerDylibResources &
+CompileOnDemandLayer2::getPerDylibResources(JITDylib &TargetD) {
+  auto I = DylibResources.find(&TargetD);
+  if (I == DylibResources.end()) {
+    auto &ImplD =
+        getExecutionSession().createJITDylib(TargetD.getName() + ".impl");
+    TargetD.withSearchOrderDo([&](const JITDylibList &TargetSearchOrder) {
+      ImplD.setSearchOrder(TargetSearchOrder, false);
+    });
+    PerDylibResources PDR(ImplD, BuildIndirectStubsManager());
+    I = DylibResources.insert(std::make_pair(&TargetD, std::move(PDR))).first;
+  }
+
+  return I->second;
 }
 
 void CompileOnDemandLayer2::emitExtractedFunctionsModule(
diff --git a/llvm/lib/ExecutionEngine/Orc/Core.cpp b/llvm/lib/ExecutionEngine/Orc/Core.cpp
index b1a82e89c6f..7df4b74aad0 100644
--- a/llvm/lib/ExecutionEngine/Orc/Core.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/Core.cpp
@@ -1702,6 +1702,13 @@ ExecutionSession::lookup(const JITDylibList &JDs, const SymbolNameSet &Symbols,
 #endif
 }
 
+void ExecutionSession::dump(raw_ostream &OS) {
+  runSessionLocked([this, &OS]() {
+    for (auto &JD : JDs)
+      JD->dump(OS);
+  });
+}
+
 void ExecutionSession::runOutstandingMUs() {
   while (1) {
     std::pair<JITDylib *, std::unique_ptr<MaterializationUnit>> JITDylibAndMU;
diff --git a/llvm/lib/ExecutionEngine/Orc/LLJIT.cpp b/llvm/lib/ExecutionEngine/Orc/LLJIT.cpp
index a3da0741e11..0a0faaa681c 100644
--- a/llvm/lib/ExecutionEngine/Orc/LLJIT.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/LLJIT.cpp
@@ -151,9 +151,9 @@ Expected<std::unique_ptr<LLLazyJIT>>
 
   const Triple &TT = JTMB.getTargetTriple();
 
-  auto CCMgr = createLocalCompileCallbackManager(TT, *ES, 0);
-  if (!CCMgr)
-    return CCMgr.takeError();
+  auto LCTMgr = createLocalLazyCallThroughManager(TT, *ES, 0);
+  if (!LCTMgr)
+    return LCTMgr.takeError();
 
   auto ISMBuilder = createLocalIndirectStubsManagerBuilder(TT);
   if (!ISMBuilder)
@@ -167,12 +167,12 @@ Expected<std::unique_ptr<LLLazyJIT>>
       return TM.takeError();
     return std::unique_ptr<LLLazyJIT>(
         new LLLazyJIT(std::move(ES), std::move(*TM), std::move(DL),
-                      std::move(*CCMgr), std::move(ISMBuilder)));
+                      std::move(*LCTMgr), std::move(ISMBuilder)));
   }
 
   return std::unique_ptr<LLLazyJIT>(new LLLazyJIT(
       std::move(ES), std::move(JTMB), std::move(DL), NumCompileThreads,
-      std::move(*CCMgr), std::move(ISMBuilder)));
+      std::move(*LCTMgr), std::move(ISMBuilder)));
 }
 
 Error LLLazyJIT::addLazyIRModule(JITDylib &JD, ThreadSafeModule TSM) {
@@ -191,21 +191,22 @@ Error LLLazyJIT::addLazyIRModule(JITDylib &JD, ThreadSafeModule TSM) {
 
 LLLazyJIT::LLLazyJIT(
     std::unique_ptr<ExecutionSession> ES, std::unique_ptr<TargetMachine> TM,
-    DataLayout DL, std::unique_ptr<JITCompileCallbackManager> CCMgr,
+    DataLayout DL, std::unique_ptr<LazyCallThroughManager> LCTMgr,
     std::function<std::unique_ptr<IndirectStubsManager>()> ISMBuilder)
     : LLJIT(std::move(ES), std::move(TM), std::move(DL)),
-      CCMgr(std::move(CCMgr)), TransformLayer(*this->ES, CompileLayer),
-      CODLayer(*this->ES, TransformLayer, *this->CCMgr, std::move(ISMBuilder)) {
-}
+      LCTMgr(std::move(LCTMgr)), TransformLayer(*this->ES, CompileLayer),
+      CODLayer(*this->ES, TransformLayer, *this->LCTMgr,
+               std::move(ISMBuilder)) {}
 
 LLLazyJIT::LLLazyJIT(
     std::unique_ptr<ExecutionSession> ES, JITTargetMachineBuilder JTMB,
-    DataLayout DL, unsigned NumCompileThreads, std::unique_ptr<JITCompileCallbackManager> CCMgr,
+    DataLayout DL, unsigned NumCompileThreads,
+    std::unique_ptr<LazyCallThroughManager> LCTMgr,
     std::function<std::unique_ptr<IndirectStubsManager>()> ISMBuilder)
     : LLJIT(std::move(ES), std::move(JTMB), std::move(DL), NumCompileThreads),
-      CCMgr(std::move(CCMgr)), TransformLayer(*this->ES, CompileLayer),
-      CODLayer(*this->ES, TransformLayer, *this->CCMgr, std::move(ISMBuilder)) {
-}
+      LCTMgr(std::move(LCTMgr)), TransformLayer(*this->ES, CompileLayer),
+      CODLayer(*this->ES, TransformLayer, *this->LCTMgr,
+               std::move(ISMBuilder)) {}
 
 } // End namespace orc.
 } // End namespace llvm.
diff --git a/llvm/test/ExecutionEngine/OrcLazy/multiple-compile-threads-basic.ll b/llvm/test/ExecutionEngine/OrcLazy/multiple-compile-threads-basic.ll
index a53f23b100c..24379128cc7 100644
--- a/llvm/test/ExecutionEngine/OrcLazy/multiple-compile-threads-basic.ll
+++ b/llvm/test/ExecutionEngine/OrcLazy/multiple-compile-threads-basic.ll
@@ -1,4 +1,4 @@
-; RUN: lli -jit-kind=orc-lazy -compile-threads=5 -thread-entry hello %s | FileCheck %s
+; RUN: lli -jit-kind=orc-lazy -compile-threads=2 -thread-entry hello %s | FileCheck %s
 ;
 ; CHECK: Hello