[ORC] Add an asynchronous jit-link function, jitLinkForORC, to RuntimeDyld and

switch RTDyldObjectLinkingLayer2 to use it.

RuntimeDyld::loadObject is currently a blocking operation. This means that any
JIT'd code whose call-graph contains an embedded complete K graph will require
at least K threads to link, which precludes the use of a fixed sized thread
pool for concurrent JITing of arbitrary code (whatever K the thread-pool is set
at, any code with a K+1 complete subgraph will deadlock at JIT-link time).

To address this issue, this commmit introduces a function called jitLinkForORC
that uses continuation-passing style to pass the fix-up and finalization steps
to the asynchronous symbol resolver interface so that linking can be performed
without blocking.

llvm-svn: 343043

5 files changed, 215 insertions, 60 deletions

diff --git a/llvm/include/llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h b/llvm/include/llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h
index 2376f36d38d..0c30520a21b 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h
@@ -103,6 +103,14 @@ public:
   }
 
 private:
+  Error onObjLoad(VModuleKey K, MaterializationResponsibility &R,
+                  object::ObjectFile &Obj,
+                  std::unique_ptr<RuntimeDyld::LoadedObjectInfo> LoadedObjInfo,
+                  std::map<StringRef, JITEvaluatedSymbol> Resolved,
+                  std::set<StringRef> &InternalSymbols);
+
+  void onObjEmit(VModuleKey K, MaterializationResponsibility &R, Error Err);
+
   mutable std::mutex RTDyldLayerMutex;
   GetMemoryManagerFunction GetMemoryManager;
   NotifyLoadedFunction NotifyLoaded;
diff --git a/llvm/include/llvm/ExecutionEngine/RuntimeDyld.h b/llvm/include/llvm/ExecutionEngine/RuntimeDyld.h
index 5dd5add1bb3..e419ee05e56 100644
--- a/llvm/include/llvm/ExecutionEngine/RuntimeDyld.h
+++ b/llvm/include/llvm/ExecutionEngine/RuntimeDyld.h
@@ -250,6 +250,16 @@ public:
   void finalizeWithMemoryManagerLocking();
 
 private:
+  friend void
+  jitLinkForORC(object::ObjectFile &Obj,
+                std::unique_ptr<MemoryBuffer> UnderlyingBuffer,
+                RuntimeDyld::MemoryManager &MemMgr, JITSymbolResolver &Resolver,
+                bool ProcessAllSections,
+                std::function<Error(std::unique_ptr<LoadedObjectInfo>,
+                                    std::map<StringRef, JITEvaluatedSymbol>)>
+                    OnLoaded,
+                std::function<void(Error)> OnEmitted);
+
   // RuntimeDyldImpl is the actual class. RuntimeDyld is just the public
   // interface.
   std::unique_ptr<RuntimeDyldImpl> Dyld;
@@ -259,6 +269,21 @@ private:
   RuntimeDyldCheckerImpl *Checker;
 };
 
+// Asynchronous JIT link for ORC.
+//
+// Warning: This API is experimental and probably should not be used by anyone
+// but ORC's RTDyldObjectLinkingLayer2. Internally it constructs a RuntimeDyld
+// instance and uses continuation passing to perform the fix-up and finalize
+// steps asynchronously.
+void jitLinkForORC(object::ObjectFile &Obj,
+                   std::unique_ptr<MemoryBuffer> UnderlyingBuffer,
+                   RuntimeDyld::MemoryManager &MemMgr,
+                   JITSymbolResolver &Resolver, bool ProcessAllSections,
+                   std::function<Error(std::unique_ptr<LoadedObjectInfo>,
+                                       std::map<StringRef, JITEvaluatedSymbol>)>
+                       OnLoaded,
+                   std::function<void(Error)> OnEmitted);
+
 } // end namespace llvm
 
 #endif // LLVM_EXECUTIONENGINE_RUNTIMEDYLD_H
diff --git a/llvm/lib/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.cpp b/llvm/lib/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.cpp
index d87078f2e94..5d814bd3b9d 100644
--- a/llvm/lib/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.cpp
@@ -88,36 +88,30 @@ void RTDyldObjectLinkingLayer2::emit(MaterializationResponsibility R,
                                      std::unique_ptr<MemoryBuffer> O) {
   assert(O && "Object must not be null");
 
-  auto &ES = getExecutionSession();
-
-  auto ObjFile = object::ObjectFile::createObjectFile(*O);
-  if (!ObjFile) {
-    getExecutionSession().reportError(ObjFile.takeError());
-    R.failMaterialization();
-  }
+  // This method launches an asynchronous link step that will fulfill our
+  // materialization responsibility. We need to switch R to be heap
+  // allocated before that happens so it can live as long as the asynchronous
+  // link needs it to (i.e. it must be able to outlive this method).
+  auto SharedR = std::make_shared<MaterializationResponsibility>(std::move(R));
 
-  auto MemoryManager = GetMemoryManager(K);
-
-  JITDylibSearchOrderResolver Resolver(R);
-  auto RTDyld = llvm::make_unique<RuntimeDyld>(*MemoryManager, Resolver);
-  RTDyld->setProcessAllSections(ProcessAllSections);
+  auto &ES = getExecutionSession();
 
-  {
-    std::lock_guard<std::mutex> Lock(RTDyldLayerMutex);
+  auto Obj = object::ObjectFile::createObjectFile(*O);
 
-    assert(!MemMgrs.count(K) &&
-           "A memory manager already exists for this key?");
-    MemMgrs[K] = std::move(MemoryManager);
+  if (!Obj) {
+    getExecutionSession().reportError(Obj.takeError());
+    SharedR->failMaterialization();
+    return;
   }
 
-  auto Info = RTDyld->loadObject(**ObjFile);
-
+  // Collect the internal symbols from the object file: We will need to
+  // filter these later.
+  auto InternalSymbols = std::make_shared<std::set<StringRef>>();
   {
-    std::set<StringRef> InternalSymbols;
-    for (auto &Sym : (*ObjFile)->symbols()) {
+    for (auto &Sym : (*Obj)->symbols()) {
       if (!(Sym.getFlags() & object::BasicSymbolRef::SF_Global)) {
         if (auto SymName = Sym.getName())
-          InternalSymbols.insert(*SymName);
+          InternalSymbols->insert(*SymName);
         else {
           ES.reportError(SymName.takeError());
           R.failMaterialization();
@@ -125,51 +119,89 @@ void RTDyldObjectLinkingLayer2::emit(MaterializationResponsibility R,
         }
       }
     }
+  }
 
-    SymbolFlagsMap ExtraSymbolsToClaim;
-    SymbolMap Symbols;
-    for (auto &KV : RTDyld->getSymbolTable()) {
-      // Scan the symbols and add them to the Symbols map for resolution.
-
-      // We never claim internal symbols.
-      if (InternalSymbols.count(KV.first))
-        continue;
-
-      auto InternedName = ES.getSymbolStringPool().intern(KV.first);
-      auto Flags = KV.second.getFlags();
+  auto MemoryManager = GetMemoryManager(K);
+  auto &MemMgr = *MemoryManager;
+  {
+    std::lock_guard<std::mutex> Lock(RTDyldLayerMutex);
 
-      // Override object flags and claim responsibility for symbols if
-      // requested.
-      if (OverrideObjectFlags || AutoClaimObjectSymbols) {
-        auto I = R.getSymbols().find(InternedName);
+    assert(!MemMgrs.count(K) &&
+           "A memory manager already exists for this key?");
+    MemMgrs[K] = std::move(MemoryManager);
+  }
 
-        if (OverrideObjectFlags && I != R.getSymbols().end())
-          Flags = JITSymbolFlags::stripTransientFlags(I->second);
-        else if (AutoClaimObjectSymbols && I == R.getSymbols().end())
-          ExtraSymbolsToClaim[InternedName] = Flags;
-      }
+  JITDylibSearchOrderResolver Resolver(*SharedR);
+
+  /* Thoughts on proper cross-dylib weak symbol handling:
+   *
+   * Change selection of canonical defs to be a manually triggered process, and
+   * add a 'canonical' bit to symbol definitions. When canonical def selection
+   * is triggered, sweep the JITDylibs to mark defs as canonical, discard
+   * duplicate defs.
+   */
+  jitLinkForORC(
+      **Obj, std::move(O), MemMgr, Resolver, ProcessAllSections,
+      [this, K, SharedR, &Obj, InternalSymbols](
+          std::unique_ptr<RuntimeDyld::LoadedObjectInfo> LoadedObjInfo,
+          std::map<StringRef, JITEvaluatedSymbol> ResolvedSymbols) {
+        return onObjLoad(K, *SharedR, **Obj, std::move(LoadedObjInfo),
+                         ResolvedSymbols, *InternalSymbols);
+      },
+      [this, K, SharedR](Error Err) {
+        onObjEmit(K, *SharedR, std::move(Err));
+      });
+}
 
-      Symbols[InternedName] = JITEvaluatedSymbol(KV.second.getAddress(), Flags);
+Error RTDyldObjectLinkingLayer2::onObjLoad(
+    VModuleKey K, MaterializationResponsibility &R, object::ObjectFile &Obj,
+    std::unique_ptr<RuntimeDyld::LoadedObjectInfo> LoadedObjInfo,
+    std::map<StringRef, JITEvaluatedSymbol> Resolved,
+    std::set<StringRef> &InternalSymbols) {
+  SymbolFlagsMap ExtraSymbolsToClaim;
+  SymbolMap Symbols;
+  for (auto &KV : Resolved) {
+    // Scan the symbols and add them to the Symbols map for resolution.
+
+    // We never claim internal symbols.
+    if (InternalSymbols.count(KV.first))
+      continue;
+
+    auto InternedName =
+        getExecutionSession().getSymbolStringPool().intern(KV.first);
+    auto Flags = KV.second.getFlags();
+
+    // Override object flags and claim responsibility for symbols if
+    // requested.
+    if (OverrideObjectFlags || AutoClaimObjectSymbols) {
+      auto I = R.getSymbols().find(InternedName);
+
+      if (OverrideObjectFlags && I != R.getSymbols().end())
+        Flags = JITSymbolFlags::stripTransientFlags(I->second);
+      else if (AutoClaimObjectSymbols && I == R.getSymbols().end())
+        ExtraSymbolsToClaim[InternedName] = Flags;
     }
 
-    if (!ExtraSymbolsToClaim.empty())
-      if (auto Err = R.defineMaterializing(ExtraSymbolsToClaim)) {
-        ES.reportError(std::move(Err));
-        R.failMaterialization();
-        return;
-      }
-
-    R.resolve(Symbols);
+    Symbols[InternedName] = JITEvaluatedSymbol(KV.second.getAddress(), Flags);
   }
 
+  if (!ExtraSymbolsToClaim.empty())
+    if (auto Err = R.defineMaterializing(ExtraSymbolsToClaim))
+      return Err;
+
+  R.resolve(Symbols);
+
   if (NotifyLoaded)
-    NotifyLoaded(K, **ObjFile, *Info);
+    NotifyLoaded(K, Obj, *LoadedObjInfo);
 
-  RTDyld->finalizeWithMemoryManagerLocking();
+  return Error::success();
+}
 
-  if (RTDyld->hasError()) {
-    ES.reportError(make_error<StringError>(RTDyld->getErrorString(),
-                                           inconvertibleErrorCode()));
+void RTDyldObjectLinkingLayer2::onObjEmit(VModuleKey K,
+                                          MaterializationResponsibility &R,
+                                          Error Err) {
+  if (Err) {
+    getExecutionSession().reportError(std::move(Err));
     R.failMaterialization();
     return;
   }
diff --git a/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyld.cpp b/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyld.cpp
index 563a9725142..53cb782c55c 100644
--- a/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyld.cpp
+++ b/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyld.cpp
@@ -134,6 +134,14 @@ void RuntimeDyldImpl::resolveRelocations() {
     ErrorStr = toString(std::move(Err));
   }
 
+  resolveLocalRelocations();
+
+  // Print out sections after relocation.
+  LLVM_DEBUG(for (int i = 0, e = Sections.size(); i != e; ++i)
+                 dumpSectionMemory(Sections[i], "after relocations"););
+}
+
+void RuntimeDyldImpl::resolveLocalRelocations() {
   // Iterate over all outstanding relocations
   for (auto it = Relocations.begin(), e = Relocations.end(); it != e; ++it) {
     // The Section here (Sections[i]) refers to the section in which the
@@ -146,10 +154,6 @@ void RuntimeDyldImpl::resolveRelocations() {
     resolveRelocationList(it->second, Addr);
   }
   Relocations.clear();
-
-  // Print out sections after relocation.
-  LLVM_DEBUG(for (int i = 0, e = Sections.size(); i != e; ++i)
-                 dumpSectionMemory(Sections[i], "after relocations"););
 }
 
 void RuntimeDyldImpl::mapSectionAddress(const void *LocalAddress,
@@ -1120,6 +1124,56 @@ Error RuntimeDyldImpl::resolveExternalSymbols() {
   return Error::success();
 }
 
+void RuntimeDyldImpl::finalizeAsync(
+    std::unique_ptr<RuntimeDyldImpl> This, std::function<void(Error)> OnEmitted,
+    std::unique_ptr<MemoryBuffer> UnderlyingBuffer) {
+
+  // FIXME: Move-capture OnRelocsApplied and UnderlyingBuffer once we have
+  // c++14.
+  auto SharedUnderlyingBuffer =
+      std::shared_ptr<MemoryBuffer>(std::move(UnderlyingBuffer));
+  auto SharedThis = std::shared_ptr<RuntimeDyldImpl>(std::move(This));
+  auto PostResolveContinuation =
+      [SharedThis, OnEmitted, SharedUnderlyingBuffer](
+          Expected<JITSymbolResolver::LookupResult> Result) {
+        if (!Result) {
+          OnEmitted(Result.takeError());
+          return;
+        }
+
+        /// Copy the result into a StringMap, where the keys are held by value.
+        StringMap<JITEvaluatedSymbol> Resolved;
+        for (auto &KV : *Result)
+          Resolved[KV.first] = KV.second;
+
+        SharedThis->applyExternalSymbolRelocations(Resolved);
+        SharedThis->resolveLocalRelocations();
+        SharedThis->registerEHFrames();
+        std::string ErrMsg;
+        if (SharedThis->MemMgr.finalizeMemory(&ErrMsg))
+          OnEmitted(make_error<StringError>(std::move(ErrMsg),
+                                            inconvertibleErrorCode()));
+        else
+          OnEmitted(Error::success());
+      };
+
+  JITSymbolResolver::LookupSet Symbols;
+
+  for (auto &RelocKV : SharedThis->ExternalSymbolRelocations) {
+    StringRef Name = RelocKV.first();
+    assert(!Name.empty() && "Symbol has no name?");
+    assert(!SharedThis->GlobalSymbolTable.count(Name) &&
+           "Name already processed. RuntimeDyld instances can not be re-used "
+           "when finalizing with finalizeAsync.");
+    Symbols.insert(Name);
+  }
+
+  if (!Symbols.empty()) {
+    SharedThis->Resolver.lookup(Symbols, PostResolveContinuation);
+  } else
+    PostResolveContinuation(std::map<StringRef, JITEvaluatedSymbol>());
+}
+
 //===----------------------------------------------------------------------===//
 // RuntimeDyld class implementation
 
@@ -1267,5 +1321,35 @@ void RuntimeDyld::deregisterEHFrames() {
   if (Dyld)
     Dyld->deregisterEHFrames();
 }
+// FIXME: Kill this with fire once we have a new JIT linker: this is only here
+// so that we can re-use RuntimeDyld's implementation without twisting the
+// interface any further for ORC's purposes.
+void jitLinkForORC(object::ObjectFile &Obj,
+                   std::unique_ptr<MemoryBuffer> UnderlyingBuffer,
+                   RuntimeDyld::MemoryManager &MemMgr,
+                   JITSymbolResolver &Resolver, bool ProcessAllSections,
+                   std::function<Error(
+                       std::unique_ptr<RuntimeDyld::LoadedObjectInfo> LoadedObj,
+                       std::map<StringRef, JITEvaluatedSymbol>)>
+                       OnLoaded,
+                   std::function<void(Error)> OnEmitted) {
+
+  RuntimeDyld RTDyld(MemMgr, Resolver);
+  RTDyld.setProcessAllSections(ProcessAllSections);
+
+  auto Info = RTDyld.loadObject(Obj);
+
+  if (RTDyld.hasError()) {
+    OnEmitted(make_error<StringError>(RTDyld.getErrorString(),
+                                      inconvertibleErrorCode()));
+    return;
+  }
+
+  if (auto Err = OnLoaded(std::move(Info), RTDyld.getSymbolTable()))
+    OnEmitted(std::move(Err));
+
+  RuntimeDyldImpl::finalizeAsync(std::move(RTDyld.Dyld), std::move(OnEmitted),
+                                 std::move(UnderlyingBuffer));
+}
 
 } // end namespace llvm
diff --git a/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyldImpl.h b/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyldImpl.h
index f01d3103f82..4c650e09ac1 100644
--- a/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyldImpl.h
+++ b/llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyldImpl.h
@@ -539,6 +539,12 @@ public:
 
   void resolveRelocations();
 
+  void resolveLocalRelocations();
+
+  static void finalizeAsync(std::unique_ptr<RuntimeDyldImpl> This,
+                            std::function<void(Error)> OnEmitted,
+                            std::unique_ptr<MemoryBuffer> UnderlyingBuffer);
+
   void reassignSectionAddress(unsigned SectionID, uint64_t Addr);
 
   void mapSectionAddress(const void *LocalAddress, uint64_t TargetAddress);