[llvm] [ORC] Make dynamic library search generation async (PR #81205)

Ben Langmuir via llvm-commits llvm-commits at lists.llvm.org
Thu Feb 8 15:57:28 PST 2024


https://github.com/benlangmuir created https://github.com/llvm/llvm-project/pull/81205

There are three stacked changes here that ultimately allow search generation to be async and stop blocking a task. The commits break up the change by layer; let me know if you'd prefer separate reviews.

In places where the existing sync API is something I expect is called directly I kept a sync version of the API using futures on top of the async one. In places where the API is likely only used internally in ORC (e.g. tryToGenerate) I dropped the synchronous version.

### [ORC] Make DefinitionGenerator::tryToGenerate async

Support unblocking the execution session task pool while running
definition generators. This allows reentrant generators (e.g. triggering
JIT lookups) with single-threaded JITs, and may improve performance when
there are slow generators.

### [ORC] Add ExecutorProcessControl::lookupSymbolsAsync

Switch the primary implementation of EPC lookupSymbols to be async,
keeping a synchronous wrapper for compatibility. Use the new async
implementation inside EPCDynamicLibrarySearchGenerator to continue
working towards a fully async search generator (remainining TODO is the
EPCGenericDylibManager::lookup).

### [ORC] Add EPCGenericDylibManager::lookupAsync and use from SimpleRemoteEPC

Provide an asynchronous lookup API for EPCGenericDylibManager and adopt
that from the SimpleRemoteEPC. This enables an end-to-end async
EPCDynamicLibrarySearchGenerator. Note: currently we keep the current
per-dlhandle lookup model, but a future improvement could do a single
async call for a given lookup operation.

>From 657a5002e4b64e20d9639522a66c65c5229c1cc7 Mon Sep 17 00:00:00 2001
From: Ben Langmuir <blangmuir at apple.com>
Date: Fri, 12 Jan 2024 15:15:33 -0800
Subject: [PATCH 1/3] [ORC] Make DefinitionGenerator::tryToGenerate async

Support unblocking the execution session task pool while running
definition generators. This allows reentrant generators (e.g. triggering
JIT lookups) with single-threaded JITs, and may improve performance when
there are slow generators.
---
 llvm/include/llvm/ExecutionEngine/Orc/Core.h  | 16 ++--
 .../Orc/EPCDynamicLibrarySearchGenerator.h    |  7 +-
 .../llvm/ExecutionEngine/Orc/ExecutionUtils.h | 21 +++--
 llvm/lib/ExecutionEngine/Orc/Core.cpp         | 92 +++++++++----------
 .../Orc/EPCDynamicLibrarySearchGenerator.cpp  | 21 +++--
 .../ExecutionEngine/Orc/ExecutionUtils.cpp    | 49 +++++-----
 .../ExecutionEngine/Orc/OrcV2CBindings.cpp    |  9 +-
 llvm/tools/llvm-jitlink/llvm-jitlink.cpp      | 10 +-
 .../ExecutionEngine/Orc/CoreAPIsTest.cpp      | 34 ++++---
 9 files changed, 141 insertions(+), 118 deletions(-)

diff --git a/llvm/include/llvm/ExecutionEngine/Orc/Core.h b/llvm/include/llvm/ExecutionEngine/Orc/Core.h
index 45bf9adcf2a4ec..e4a3dc6ea6de7b 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/Core.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/Core.h
@@ -949,15 +949,18 @@ class DefinitionGenerator {
 public:
   virtual ~DefinitionGenerator();
 
+  using NotifyCompleteFn = unique_function<void(LookupState LS, Error)>;
+
   /// DefinitionGenerators should override this method to insert new
   /// definitions into the parent JITDylib. K specifies the kind of this
   /// lookup. JD specifies the target JITDylib being searched, and
   /// JDLookupFlags specifies whether the search should match against
   /// hidden symbols. Finally, Symbols describes the set of unresolved
   /// symbols and their associated lookup flags.
-  virtual Error tryToGenerate(LookupState &LS, LookupKind K, JITDylib &JD,
-                              JITDylibLookupFlags JDLookupFlags,
-                              const SymbolLookupSet &LookupSet) = 0;
+  virtual void tryToGenerate(LookupState LS, LookupKind K, JITDylib &JD,
+                             JITDylibLookupFlags JDLookupFlags,
+                             const SymbolLookupSet &LookupSet,
+                             NotifyCompleteFn NotifyComplete) = 0;
 
 private:
   std::mutex M;
@@ -2006,9 +2009,10 @@ class ReexportsGenerator : public DefinitionGenerator {
                      JITDylibLookupFlags SourceJDLookupFlags,
                      SymbolPredicate Allow = SymbolPredicate());
 
-  Error tryToGenerate(LookupState &LS, LookupKind K, JITDylib &JD,
-                      JITDylibLookupFlags JDLookupFlags,
-                      const SymbolLookupSet &LookupSet) override;
+  void tryToGenerate(LookupState LS, LookupKind K, JITDylib &JD,
+                     JITDylibLookupFlags JDLookupFlags,
+                     const SymbolLookupSet &LookupSet,
+                     NotifyCompleteFn NotifyComplete) override;
 
 private:
   JITDylib &SourceJD;
diff --git a/llvm/include/llvm/ExecutionEngine/Orc/EPCDynamicLibrarySearchGenerator.h b/llvm/include/llvm/ExecutionEngine/Orc/EPCDynamicLibrarySearchGenerator.h
index e56afe4fe656b3..2db8510a200686 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/EPCDynamicLibrarySearchGenerator.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/EPCDynamicLibrarySearchGenerator.h
@@ -60,9 +60,10 @@ class EPCDynamicLibrarySearchGenerator : public DefinitionGenerator {
     return Load(ES, nullptr, std::move(Allow), std::move(AddAbsoluteSymbols));
   }
 
-  Error tryToGenerate(LookupState &LS, LookupKind K, JITDylib &JD,
-                      JITDylibLookupFlags JDLookupFlags,
-                      const SymbolLookupSet &Symbols) override;
+  void tryToGenerate(LookupState LS, LookupKind K, JITDylib &JD,
+                     JITDylibLookupFlags JDLookupFlags,
+                     const SymbolLookupSet &Symbols,
+                     NotifyCompleteFn NotifyComplete) override;
 
 private:
   ExecutorProcessControl &EPC;
diff --git a/llvm/include/llvm/ExecutionEngine/Orc/ExecutionUtils.h b/llvm/include/llvm/ExecutionEngine/Orc/ExecutionUtils.h
index f7c286bec778b0..973654731423e2 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/ExecutionUtils.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/ExecutionUtils.h
@@ -250,9 +250,10 @@ class DynamicLibrarySearchGenerator : public DefinitionGenerator {
                 std::move(AddAbsoluteSymbols));
   }
 
-  Error tryToGenerate(LookupState &LS, LookupKind K, JITDylib &JD,
-                      JITDylibLookupFlags JDLookupFlags,
-                      const SymbolLookupSet &Symbols) override;
+  void tryToGenerate(LookupState LS, LookupKind K, JITDylib &JD,
+                     JITDylibLookupFlags JDLookupFlags,
+                     const SymbolLookupSet &Symbols,
+                     NotifyCompleteFn NotifyComplete) override;
 
 private:
   sys::DynamicLibrary Dylib;
@@ -307,9 +308,10 @@ class StaticLibraryDefinitionGenerator : public DefinitionGenerator {
     return ImportedDynamicLibraries;
   }
 
-  Error tryToGenerate(LookupState &LS, LookupKind K, JITDylib &JD,
-                      JITDylibLookupFlags JDLookupFlags,
-                      const SymbolLookupSet &Symbols) override;
+  void tryToGenerate(LookupState LS, LookupKind K, JITDylib &JD,
+                     JITDylibLookupFlags JDLookupFlags,
+                     const SymbolLookupSet &Symbols,
+                     NotifyCompleteFn NotifyComplete) override;
 
 private:
   StaticLibraryDefinitionGenerator(ObjectLayer &L,
@@ -343,9 +345,10 @@ class DLLImportDefinitionGenerator : public DefinitionGenerator {
   static std::unique_ptr<DLLImportDefinitionGenerator>
   Create(ExecutionSession &ES, ObjectLinkingLayer &L);
 
-  Error tryToGenerate(LookupState &LS, LookupKind K, JITDylib &JD,
-                      JITDylibLookupFlags JDLookupFlags,
-                      const SymbolLookupSet &Symbols) override;
+  void tryToGenerate(LookupState LS, LookupKind K, JITDylib &JD,
+                     JITDylibLookupFlags JDLookupFlags,
+                     const SymbolLookupSet &Symbols,
+                     NotifyCompleteFn NotifyComplete) override;
 
 private:
   DLLImportDefinitionGenerator(ExecutionSession &ES, ObjectLinkingLayer &L)
diff --git a/llvm/lib/ExecutionEngine/Orc/Core.cpp b/llvm/lib/ExecutionEngine/Orc/Core.cpp
index 37c32f86e8d873..d696306863c80a 100644
--- a/llvm/lib/ExecutionEngine/Orc/Core.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/Core.cpp
@@ -614,17 +614,18 @@ ReexportsGenerator::ReexportsGenerator(JITDylib &SourceJD,
     : SourceJD(SourceJD), SourceJDLookupFlags(SourceJDLookupFlags),
       Allow(std::move(Allow)) {}
 
-Error ReexportsGenerator::tryToGenerate(LookupState &LS, LookupKind K,
-                                        JITDylib &JD,
-                                        JITDylibLookupFlags JDLookupFlags,
-                                        const SymbolLookupSet &LookupSet) {
+void ReexportsGenerator::tryToGenerate(LookupState LS, LookupKind K,
+                                       JITDylib &JD,
+                                       JITDylibLookupFlags JDLookupFlags,
+                                       const SymbolLookupSet &LookupSet,
+                                       NotifyCompleteFn NotifyComplete) {
   assert(&JD != &SourceJD && "Cannot re-export from the same dylib");
 
   // Use lookupFlags to find the subset of symbols that match our lookup.
   auto Flags = JD.getExecutionSession().lookupFlags(
       K, {{&SourceJD, JDLookupFlags}}, LookupSet);
   if (!Flags)
-    return Flags.takeError();
+    return NotifyComplete(std::move(LS), Flags.takeError());
 
   // Create an alias map.
   orc::SymbolAliasMap AliasMap;
@@ -633,10 +634,11 @@ Error ReexportsGenerator::tryToGenerate(LookupState &LS, LookupKind K,
       AliasMap[KV.first] = SymbolAliasMapEntry(KV.first, KV.second);
 
   if (AliasMap.empty())
-    return Error::success();
+    return NotifyComplete(std::move(LS), Error::success());
 
   // Define the re-exports.
-  return JD.define(reexports(SourceJD, AliasMap, SourceJDLookupFlags));
+  auto Err = JD.define(reexports(SourceJD, AliasMap, SourceJDLookupFlags));
+  NotifyComplete(std::move(LS), std::move(Err));
 }
 
 LookupState::LookupState(std::unique_ptr<InProgressLookupState> IPLS)
@@ -2232,8 +2234,8 @@ void ExecutionSession::OL_applyQueryPhase1(
                << " remaining generators for "
                << IPLS->DefGeneratorCandidates.size() << " candidates\n";
     });
-    while (!IPLS->CurDefGeneratorStack.empty() &&
-           !IPLS->DefGeneratorCandidates.empty()) {
+    if (!IPLS->CurDefGeneratorStack.empty() &&
+        !IPLS->DefGeneratorCandidates.empty()) {
       auto DG = IPLS->CurDefGeneratorStack.back().lock();
 
       if (!DG)
@@ -2266,49 +2268,41 @@ void ExecutionSession::OL_applyQueryPhase1(
 
       // Run the generator. If the generator takes ownership of QA then this
       // will break the loop.
-      {
-        LLVM_DEBUG(dbgs() << "  Attempting to generate " << LookupSet << "\n");
-        LookupState LS(std::move(IPLS));
-        Err = DG->tryToGenerate(LS, K, JD, JDLookupFlags, LookupSet);
-        IPLS = std::move(LS.IPLS);
-      }
-
-      // If the lookup returned then pop the generator stack and unblock the
-      // next lookup on this generator (if any).
-      if (IPLS)
-        OL_resumeLookupAfterGeneration(*IPLS);
-
-      // If there was an error then fail the query.
-      if (Err) {
-        LLVM_DEBUG({
-          dbgs() << "  Error attempting to generate " << LookupSet << "\n";
-        });
-        assert(IPLS && "LS cannot be retained if error is returned");
-        return IPLS->fail(std::move(Err));
-      }
+      LLVM_DEBUG(dbgs() << "  Attempting to generate " << LookupSet << "\n");
+      DG->tryToGenerate(
+          std::move(IPLS), K, JD, JDLookupFlags, LookupSet,
+          [this](LookupState LS, Error Err) {
+            // If the lookup returned then pop the generator stack and unblock
+            // the next lookup on this generator (if any).
+            if (LS.IPLS)
+              OL_resumeLookupAfterGeneration(*LS.IPLS);
+
+            // If there was an error then fail the query.
+            if (Err) {
+              assert(LS.IPLS && "LS cannot be retained if error is returned");
+              LLVM_DEBUG({
+                auto &LookupSet = LS.IPLS->DefGeneratorCandidates;
+                dbgs() << "  Error attempting to generate " << LookupSet
+                       << "\n";
+              });
+              return LS.IPLS->fail(std::move(Err));
+            }
 
-      // Otherwise if QA was captured then break the loop.
-      if (!IPLS) {
-        LLVM_DEBUG(
-            { dbgs() << "  LookupState captured. Exiting phase1 for now.\n"; });
-        return;
-      }
+            // Otherwise if QA was captured then break the loop.
+            if (!LS.IPLS) {
+              LLVM_DEBUG({
+                dbgs() << "  LookupState captured. Exiting phase1 for now.\n";
+              });
+              return;
+            }
 
-      // Otherwise if we're continuing around the loop then update candidates
-      // for the next round.
-      runSessionLocked([&] {
-        LLVM_DEBUG(dbgs() << "  Updating candidate set post-generation\n");
-        Err = IL_updateCandidatesFor(
-            JD, JDLookupFlags, IPLS->DefGeneratorCandidates,
-            JD.DefGenerators.empty() ? nullptr
-                                     : &IPLS->DefGeneratorNonCandidates);
-      });
+            // Otherwise if we're continuing around the loop, reenter
+            // OL_applyQueryPhase1. Updating candidates is handled there.
+            dispatchTask(std::make_unique<LookupTask>(std::move(LS)));
+          });
 
-      // If updating candidates failed then fail the query.
-      if (Err) {
-        LLVM_DEBUG(dbgs() << "  Error encountered while updating candidates\n");
-        return IPLS->fail(std::move(Err));
-      }
+      // Continuation will resume query as necessary.
+      return;
     }
 
     if (IPLS->DefGeneratorCandidates.empty() &&
diff --git a/llvm/lib/ExecutionEngine/Orc/EPCDynamicLibrarySearchGenerator.cpp b/llvm/lib/ExecutionEngine/Orc/EPCDynamicLibrarySearchGenerator.cpp
index 460f4e1c448e67..bc062e199aa5e2 100644
--- a/llvm/lib/ExecutionEngine/Orc/EPCDynamicLibrarySearchGenerator.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/EPCDynamicLibrarySearchGenerator.cpp
@@ -23,12 +23,13 @@ EPCDynamicLibrarySearchGenerator::Load(
       ES, *Handle, std::move(Allow), std::move(AddAbsoluteSymbols));
 }
 
-Error EPCDynamicLibrarySearchGenerator::tryToGenerate(
-    LookupState &LS, LookupKind K, JITDylib &JD,
-    JITDylibLookupFlags JDLookupFlags, const SymbolLookupSet &Symbols) {
+void EPCDynamicLibrarySearchGenerator::tryToGenerate(
+    LookupState LS, LookupKind K, JITDylib &JD,
+    JITDylibLookupFlags JDLookupFlags, const SymbolLookupSet &Symbols,
+    NotifyCompleteFn NotifyComplete) {
 
   if (Symbols.empty())
-    return Error::success();
+    return NotifyComplete(std::move(LS), Error::success());
 
   SymbolLookupSet LookupSymbols;
 
@@ -44,7 +45,7 @@ Error EPCDynamicLibrarySearchGenerator::tryToGenerate(
   ExecutorProcessControl::LookupRequest Request(H, LookupSymbols);
   auto Result = EPC.lookupSymbols(Request);
   if (!Result)
-    return Result.takeError();
+    return NotifyComplete(std::move(LS), Result.takeError());
 
   assert(Result->size() == 1 && "Results for more than one library returned");
   assert(Result->front().size() == LookupSymbols.size() &&
@@ -59,12 +60,14 @@ Error EPCDynamicLibrarySearchGenerator::tryToGenerate(
 
   // If there were no resolved symbols bail out.
   if (NewSymbols.empty())
-    return Error::success();
+    return NotifyComplete(std::move(LS), Error::success());
 
   // Define resolved symbols.
-  if (AddAbsoluteSymbols)
-    return AddAbsoluteSymbols(JD, std::move(NewSymbols));
-  return JD.define(absoluteSymbols(std::move(NewSymbols)));
+  Error Err = AddAbsoluteSymbols
+                  ? AddAbsoluteSymbols(JD, std::move(NewSymbols))
+                  : JD.define(absoluteSymbols(std::move(NewSymbols)));
+
+  NotifyComplete(std::move(LS), std::move(Err));
 }
 
 } // end namespace orc
diff --git a/llvm/lib/ExecutionEngine/Orc/ExecutionUtils.cpp b/llvm/lib/ExecutionEngine/Orc/ExecutionUtils.cpp
index 3952445bb1aa4b..155784d82f558f 100644
--- a/llvm/lib/ExecutionEngine/Orc/ExecutionUtils.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/ExecutionUtils.cpp
@@ -16,6 +16,7 @@
 #include "llvm/IR/Module.h"
 #include "llvm/MC/TargetRegistry.h"
 #include "llvm/Object/MachOUniversal.h"
+#include "llvm/Support/Error.h"
 #include "llvm/Support/FormatVariadic.h"
 #include "llvm/Target/TargetMachine.h"
 #include <string>
@@ -237,9 +238,10 @@ DynamicLibrarySearchGenerator::Load(const char *FileName, char GlobalPrefix,
       std::move(AddAbsoluteSymbols));
 }
 
-Error DynamicLibrarySearchGenerator::tryToGenerate(
-    LookupState &LS, LookupKind K, JITDylib &JD,
-    JITDylibLookupFlags JDLookupFlags, const SymbolLookupSet &Symbols) {
+void DynamicLibrarySearchGenerator::tryToGenerate(
+    LookupState LS, LookupKind K, JITDylib &JD,
+    JITDylibLookupFlags JDLookupFlags, const SymbolLookupSet &Symbols,
+    NotifyCompleteFn NotifyComplete) {
   orc::SymbolMap NewSymbols;
 
   bool HasGlobalPrefix = (GlobalPrefix != '\0');
@@ -263,11 +265,13 @@ Error DynamicLibrarySearchGenerator::tryToGenerate(
   }
 
   if (NewSymbols.empty())
-    return Error::success();
+    return NotifyComplete(std::move(LS), Error::success());
+
+  Error Err = AddAbsoluteSymbols
+                  ? AddAbsoluteSymbols(JD, std::move(NewSymbols))
+                  : JD.define(absoluteSymbols(std::move(NewSymbols)));
 
-  if (AddAbsoluteSymbols)
-    return AddAbsoluteSymbols(JD, std::move(NewSymbols));
-  return JD.define(absoluteSymbols(std::move(NewSymbols)));
+  NotifyComplete(std::move(LS), std::move(Err));
 }
 
 Expected<std::unique_ptr<StaticLibraryDefinitionGenerator>>
@@ -379,17 +383,18 @@ StaticLibraryDefinitionGenerator::Create(
                                  inconvertibleErrorCode());
 }
 
-Error StaticLibraryDefinitionGenerator::tryToGenerate(
-    LookupState &LS, LookupKind K, JITDylib &JD,
-    JITDylibLookupFlags JDLookupFlags, const SymbolLookupSet &Symbols) {
+void StaticLibraryDefinitionGenerator::tryToGenerate(
+    LookupState LS, LookupKind K, JITDylib &JD,
+    JITDylibLookupFlags JDLookupFlags, const SymbolLookupSet &Symbols,
+    NotifyCompleteFn NotifyComplete) {
   // Don't materialize symbols from static archives unless this is a static
   // lookup.
   if (K != LookupKind::Static)
-    return Error::success();
+    return NotifyComplete(std::move(LS), Error::success());
 
   // Bail out early if we've already freed the archive.
   if (!Archive)
-    return Error::success();
+    return NotifyComplete(std::move(LS), Error::success());
 
   DenseSet<std::pair<StringRef, StringRef>> ChildBufferInfos;
 
@@ -408,14 +413,14 @@ Error StaticLibraryDefinitionGenerator::tryToGenerate(
 
     auto I = GetObjFileInterface(L.getExecutionSession(), ChildBufferRef);
     if (!I)
-      return I.takeError();
+      return NotifyComplete(std::move(LS), I.takeError());
 
     if (auto Err = L.add(JD, MemoryBuffer::getMemBuffer(ChildBufferRef, false),
                          std::move(*I)))
-      return Err;
+      return NotifyComplete(std::move(LS), std::move(Err));
   }
 
-  return Error::success();
+  NotifyComplete(std::move(LS), Error::success());
 }
 
 Error StaticLibraryDefinitionGenerator::buildObjectFilesMap() {
@@ -489,9 +494,10 @@ DLLImportDefinitionGenerator::Create(ExecutionSession &ES,
       new DLLImportDefinitionGenerator(ES, L));
 }
 
-Error DLLImportDefinitionGenerator::tryToGenerate(
-    LookupState &LS, LookupKind K, JITDylib &JD,
-    JITDylibLookupFlags JDLookupFlags, const SymbolLookupSet &Symbols) {
+void DLLImportDefinitionGenerator::tryToGenerate(
+    LookupState LS, LookupKind K, JITDylib &JD,
+    JITDylibLookupFlags JDLookupFlags, const SymbolLookupSet &Symbols,
+    NotifyCompleteFn NotifyComplete) {
   JITDylibSearchOrder LinkOrder;
   JD.withLinkOrderDo([&](const JITDylibSearchOrder &LO) {
     LinkOrder.reserve(LO.size());
@@ -524,12 +530,13 @@ Error DLLImportDefinitionGenerator::tryToGenerate(
   auto Resolved =
       ES.lookup(LinkOrder, LookupSet, LookupKind::DLSym, SymbolState::Resolved);
   if (!Resolved)
-    return Resolved.takeError();
+    return NotifyComplete(std::move(LS), Resolved.takeError());
 
   auto G = createStubsGraph(*Resolved);
   if (!G)
-    return G.takeError();
-  return L.add(JD, std::move(*G));
+    return NotifyComplete(std::move(LS), G.takeError());
+
+  NotifyComplete(std::move(LS), L.add(JD, std::move(*G)));
 }
 
 Expected<unsigned>
diff --git a/llvm/lib/ExecutionEngine/Orc/OrcV2CBindings.cpp b/llvm/lib/ExecutionEngine/Orc/OrcV2CBindings.cpp
index 453b8f86868adb..e5c1da4a36e9ca 100644
--- a/llvm/lib/ExecutionEngine/Orc/OrcV2CBindings.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/OrcV2CBindings.cpp
@@ -268,9 +268,10 @@ class CAPIDefinitionGenerator final : public DefinitionGenerator {
       Dispose(Ctx);
   }
 
-  Error tryToGenerate(LookupState &LS, LookupKind K, JITDylib &JD,
-                      JITDylibLookupFlags JDLookupFlags,
-                      const SymbolLookupSet &LookupSet) override {
+  void tryToGenerate(LookupState LS, LookupKind K, JITDylib &JD,
+                     JITDylibLookupFlags JDLookupFlags,
+                     const SymbolLookupSet &LookupSet,
+                     NotifyCompleteFn NotifyComplete) override {
 
     // Take the lookup state.
     LLVMOrcLookupStateRef LSR = ::wrap(OrcV2CAPIHelper::extractLookupState(LS));
@@ -300,7 +301,7 @@ class CAPIDefinitionGenerator final : public DefinitionGenerator {
     // Restore the lookup state.
     OrcV2CAPIHelper::resetLookupState(LS, ::unwrap(LSR));
 
-    return Err;
+    NotifyComplete(std::move(LS), std::move(Err));
   }
 
 private:
diff --git a/llvm/tools/llvm-jitlink/llvm-jitlink.cpp b/llvm/tools/llvm-jitlink/llvm-jitlink.cpp
index 769ed17ac4cbd5..8d548bbcd69fce 100644
--- a/llvm/tools/llvm-jitlink/llvm-jitlink.cpp
+++ b/llvm/tools/llvm-jitlink/llvm-jitlink.cpp
@@ -888,13 +888,15 @@ static Expected<std::unique_ptr<ExecutorProcessControl>> connectToExecutor() {
 
 class PhonyExternalsGenerator : public DefinitionGenerator {
 public:
-  Error tryToGenerate(LookupState &LS, LookupKind K, JITDylib &JD,
-                      JITDylibLookupFlags JDLookupFlags,
-                      const SymbolLookupSet &LookupSet) override {
+  void tryToGenerate(LookupState LS, LookupKind K, JITDylib &JD,
+                     JITDylibLookupFlags JDLookupFlags,
+                     const SymbolLookupSet &LookupSet,
+                     NotifyCompleteFn NotifyComplete) override {
     SymbolMap PhonySymbols;
     for (auto &KV : LookupSet)
       PhonySymbols[KV.first] = {ExecutorAddr(), JITSymbolFlags::Exported};
-    return JD.define(absoluteSymbols(std::move(PhonySymbols)));
+    auto Err = JD.define(absoluteSymbols(std::move(PhonySymbols)));
+    NotifyComplete(std::move(LS), std::move(Err));
   }
 };
 
diff --git a/llvm/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp b/llvm/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp
index 5e2b5f35bcf471..e87b7887951cba 100644
--- a/llvm/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp
+++ b/llvm/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp
@@ -11,6 +11,7 @@
 #include "llvm/Config/llvm-config.h"
 #include "llvm/ExecutionEngine/Orc/Core.h"
 #include "llvm/ExecutionEngine/Orc/Shared/OrcError.h"
+#include "llvm/Support/Error.h"
 #include "llvm/Testing/Support/Error.h"
 
 #include <deque>
@@ -339,9 +340,12 @@ TEST_F(CoreAPIsStandardTest, LookupWithGeneratorFailure) {
 
   class BadGenerator : public DefinitionGenerator {
   public:
-    Error tryToGenerate(LookupState &LS, LookupKind K, JITDylib &,
-                        JITDylibLookupFlags, const SymbolLookupSet &) override {
-      return make_error<StringError>("BadGenerator", inconvertibleErrorCode());
+    void tryToGenerate(LookupState LS, LookupKind K, JITDylib &,
+                       JITDylibLookupFlags, const SymbolLookupSet &,
+                       NotifyCompleteFn NotifyComplete) override {
+      NotifyComplete(
+          std::move(LS),
+          make_error<StringError>("BadGenerator", inconvertibleErrorCode()));
     }
   };
 
@@ -1040,9 +1044,10 @@ TEST_F(CoreAPIsStandardTest, GeneratorTest) {
   class TestGenerator : public DefinitionGenerator {
   public:
     TestGenerator(SymbolMap Symbols) : Symbols(std::move(Symbols)) {}
-    Error tryToGenerate(LookupState &LS, LookupKind K, JITDylib &JD,
-                        JITDylibLookupFlags JDLookupFlags,
-                        const SymbolLookupSet &Names) override {
+    void tryToGenerate(LookupState LS, LookupKind K, JITDylib &JD,
+                       JITDylibLookupFlags JDLookupFlags,
+                       const SymbolLookupSet &Names,
+                       NotifyCompleteFn NotifyComplete) override {
       SymbolMap NewDefs;
 
       for (const auto &KV : Names) {
@@ -1052,7 +1057,7 @@ TEST_F(CoreAPIsStandardTest, GeneratorTest) {
       }
 
       cantFail(JD.define(absoluteSymbols(std::move(NewDefs))));
-      return Error::success();
+      NotifyComplete(std::move(LS), Error::success());
     };
 
   private:
@@ -1088,13 +1093,16 @@ class SimpleAsyncGenerator : public DefinitionGenerator {
                         JITDylibLookupFlags, const SymbolLookupSet &)>
       TryToGenerateOverride;
 
-  Error tryToGenerate(LookupState &LS, LookupKind K, JITDylib &JD,
-                      JITDylibLookupFlags JDLookupFlags,
-                      const SymbolLookupSet &Names) override {
-    if (TryToGenerateOverride)
-      return TryToGenerateOverride(LS, K, JD, JDLookupFlags, Names);
+  void tryToGenerate(LookupState LS, LookupKind K, JITDylib &JD,
+                     JITDylibLookupFlags JDLookupFlags,
+                     const SymbolLookupSet &Names,
+                     NotifyCompleteFn NotifyComplete) override {
+    if (TryToGenerateOverride) {
+      auto Err = TryToGenerateOverride(LS, K, JD, JDLookupFlags, Names);
+      return NotifyComplete(std::move(LS), std::move(Err));
+    }
     Lookup = SuspendedLookupInfo{std::move(LS), K, &JD, JDLookupFlags, Names};
-    return Error::success();
+    NotifyComplete(std::move(LS), Error::success());
   }
 
   SuspendedLookupInfo takeLookup() {

>From 968ad419db6999629f10457ba16cd36f73c160ff Mon Sep 17 00:00:00 2001
From: Ben Langmuir <blangmuir at apple.com>
Date: Wed, 17 Jan 2024 13:58:37 -0800
Subject: [PATCH 2/3] [ORC] Add ExecutorProcessControl::lookupSymbolsAsync

Switch the primary implementation of EPC lookupSymbols to be async,
keeping a synchronous wrapper for compatibility. Use the new async
implementation inside EPCDynamicLibrarySearchGenerator to continue
working towards a fully async search generator (remainining TODO is the
EPCGenericDylibManager::lookup).
---
 .../Orc/ExecutorProcessControl.h              | 30 ++++++++--
 .../ExecutionEngine/Orc/SimpleRemoteEPC.h     |  4 +-
 .../Orc/EPCDynamicLibrarySearchGenerator.cpp  | 57 ++++++++++---------
 .../Orc/ExecutorProcessControl.cpp            | 10 ++--
 .../ExecutionEngine/Orc/SimpleRemoteEPC.cpp   |  8 +--
 .../Orc/ObjectLinkingLayerTest.cpp            |  6 +-
 6 files changed, 69 insertions(+), 46 deletions(-)

diff --git a/llvm/include/llvm/ExecutionEngine/Orc/ExecutorProcessControl.h b/llvm/include/llvm/ExecutionEngine/Orc/ExecutorProcessControl.h
index 9e42d6dd615dfd..51064cf2b7acd6 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/ExecutorProcessControl.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/ExecutorProcessControl.h
@@ -290,8 +290,26 @@ class ExecutorProcessControl {
   /// that correspond to the lookup order. If a required symbol is not
   /// found then this method will return an error. If a weakly referenced
   /// symbol is not found then it be assigned a '0' value.
-  virtual Expected<std::vector<tpctypes::LookupResult>>
-  lookupSymbols(ArrayRef<LookupRequest> Request) = 0;
+  Expected<std::vector<tpctypes::LookupResult>>
+  lookupSymbols(ArrayRef<LookupRequest> Request) {
+    std::promise<Expected<std::vector<tpctypes::LookupResult>>> RP;
+    auto RF = RP.get_future();
+    lookupSymbolsAsync(Request,
+                       [&RP](auto Result) { RP.set_value(std::move(Result)); });
+    return RF.get();
+  }
+
+  using SymbolLookupCompleteFn =
+      unique_function<void(Expected<std::vector<tpctypes::LookupResult>>)>;
+
+  /// Search for symbols in the target process.
+  ///
+  /// The result of the lookup is a 2-dimensional array of target addresses
+  /// that correspond to the lookup order. If a required symbol is not
+  /// found then this method will return an error. If a weakly referenced
+  /// symbol is not found then it be assigned a '0' value.
+  virtual void lookupSymbolsAsync(ArrayRef<LookupRequest> Request,
+                                  SymbolLookupCompleteFn F) = 0;
 
   /// Run function with a main-like signature.
   virtual Expected<int32_t> runAsMain(ExecutorAddr MainFnAddr,
@@ -462,8 +480,8 @@ class UnsupportedExecutorProcessControl : public ExecutorProcessControl,
     llvm_unreachable("Unsupported");
   }
 
-  Expected<std::vector<tpctypes::LookupResult>>
-  lookupSymbols(ArrayRef<LookupRequest> Request) override {
+  void lookupSymbolsAsync(ArrayRef<LookupRequest> Request,
+                          SymbolLookupCompleteFn F) override {
     llvm_unreachable("Unsupported");
   }
 
@@ -510,8 +528,8 @@ class SelfExecutorProcessControl : public ExecutorProcessControl,
 
   Expected<tpctypes::DylibHandle> loadDylib(const char *DylibPath) override;
 
-  Expected<std::vector<tpctypes::LookupResult>>
-  lookupSymbols(ArrayRef<LookupRequest> Request) override;
+  void lookupSymbolsAsync(ArrayRef<LookupRequest> Request,
+                          SymbolLookupCompleteFn F) override;
 
   Expected<int32_t> runAsMain(ExecutorAddr MainFnAddr,
                               ArrayRef<std::string> Args) override;
diff --git a/llvm/include/llvm/ExecutionEngine/Orc/SimpleRemoteEPC.h b/llvm/include/llvm/ExecutionEngine/Orc/SimpleRemoteEPC.h
index 25b79be48810c6..c10b8df01cc0a4 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/SimpleRemoteEPC.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/SimpleRemoteEPC.h
@@ -71,8 +71,8 @@ class SimpleRemoteEPC : public ExecutorProcessControl,
 
   Expected<tpctypes::DylibHandle> loadDylib(const char *DylibPath) override;
 
-  Expected<std::vector<tpctypes::LookupResult>>
-  lookupSymbols(ArrayRef<LookupRequest> Request) override;
+  void lookupSymbolsAsync(ArrayRef<LookupRequest> Request,
+                          SymbolLookupCompleteFn F) override;
 
   Expected<int32_t> runAsMain(ExecutorAddr MainFnAddr,
                               ArrayRef<std::string> Args) override;
diff --git a/llvm/lib/ExecutionEngine/Orc/EPCDynamicLibrarySearchGenerator.cpp b/llvm/lib/ExecutionEngine/Orc/EPCDynamicLibrarySearchGenerator.cpp
index bc062e199aa5e2..94566df130f955 100644
--- a/llvm/lib/ExecutionEngine/Orc/EPCDynamicLibrarySearchGenerator.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/EPCDynamicLibrarySearchGenerator.cpp
@@ -40,34 +40,37 @@ void EPCDynamicLibrarySearchGenerator::tryToGenerate(
     LookupSymbols.add(KV.first, SymbolLookupFlags::WeaklyReferencedSymbol);
   }
 
-  SymbolMap NewSymbols;
-
   ExecutorProcessControl::LookupRequest Request(H, LookupSymbols);
-  auto Result = EPC.lookupSymbols(Request);
-  if (!Result)
-    return NotifyComplete(std::move(LS), Result.takeError());
-
-  assert(Result->size() == 1 && "Results for more than one library returned");
-  assert(Result->front().size() == LookupSymbols.size() &&
-         "Result has incorrect number of elements");
-
-  auto ResultI = Result->front().begin();
-  for (auto &KV : LookupSymbols) {
-    if (ResultI->getAddress())
-      NewSymbols[KV.first] = *ResultI;
-    ++ResultI;
-  }
-
-  // If there were no resolved symbols bail out.
-  if (NewSymbols.empty())
-    return NotifyComplete(std::move(LS), Error::success());
-
-  // Define resolved symbols.
-  Error Err = AddAbsoluteSymbols
-                  ? AddAbsoluteSymbols(JD, std::move(NewSymbols))
-                  : JD.define(absoluteSymbols(std::move(NewSymbols)));
-
-  NotifyComplete(std::move(LS), std::move(Err));
+  // Copy-capture LookupSymbols, since LookupRequest keeps a reference.
+  EPC.lookupSymbolsAsync(Request, [this, &JD, LS = std::move(LS),
+                                   NotifyComplete = std::move(NotifyComplete),
+                                   LookupSymbols](auto Result) mutable {
+    if (!Result)
+      return NotifyComplete(std::move(LS), Result.takeError());
+
+    assert(Result->size() == 1 && "Results for more than one library returned");
+    assert(Result->front().size() == LookupSymbols.size() &&
+           "Result has incorrect number of elements");
+
+    SymbolMap NewSymbols;
+    auto ResultI = Result->front().begin();
+    for (auto &KV : LookupSymbols) {
+      if (ResultI->getAddress())
+        NewSymbols[KV.first] = *ResultI;
+      ++ResultI;
+    }
+
+    // If there were no resolved symbols bail out.
+    if (NewSymbols.empty())
+      return NotifyComplete(std::move(LS), Error::success());
+
+    // Define resolved symbols.
+    Error Err = AddAbsoluteSymbols
+                    ? AddAbsoluteSymbols(JD, std::move(NewSymbols))
+                    : JD.define(absoluteSymbols(std::move(NewSymbols)));
+
+    NotifyComplete(std::move(LS), std::move(Err));
+  });
 }
 
 } // end namespace orc
diff --git a/llvm/lib/ExecutionEngine/Orc/ExecutorProcessControl.cpp b/llvm/lib/ExecutionEngine/Orc/ExecutorProcessControl.cpp
index f0c551cd778047..efafca949e61ef 100644
--- a/llvm/lib/ExecutionEngine/Orc/ExecutorProcessControl.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/ExecutorProcessControl.cpp
@@ -89,8 +89,9 @@ SelfExecutorProcessControl::loadDylib(const char *DylibPath) {
   return ExecutorAddr::fromPtr(Dylib.getOSSpecificHandle());
 }
 
-Expected<std::vector<tpctypes::LookupResult>>
-SelfExecutorProcessControl::lookupSymbols(ArrayRef<LookupRequest> Request) {
+void SelfExecutorProcessControl::lookupSymbolsAsync(
+    ArrayRef<LookupRequest> Request,
+    ExecutorProcessControl::SymbolLookupCompleteFn Complete) {
   std::vector<tpctypes::LookupResult> R;
 
   for (auto &Elem : Request) {
@@ -105,7 +106,8 @@ SelfExecutorProcessControl::lookupSymbols(ArrayRef<LookupRequest> Request) {
         // FIXME: Collect all failing symbols before erroring out.
         SymbolNameVector MissingSymbols;
         MissingSymbols.push_back(Sym);
-        return make_error<SymbolsNotFound>(SSP, std::move(MissingSymbols));
+        return Complete(
+            make_error<SymbolsNotFound>(SSP, std::move(MissingSymbols)));
       }
       // FIXME: determine accurate JITSymbolFlags.
       R.back().push_back(
@@ -113,7 +115,7 @@ SelfExecutorProcessControl::lookupSymbols(ArrayRef<LookupRequest> Request) {
     }
   }
 
-  return R;
+  Complete(std::move(R));
 }
 
 Expected<int32_t>
diff --git a/llvm/lib/ExecutionEngine/Orc/SimpleRemoteEPC.cpp b/llvm/lib/ExecutionEngine/Orc/SimpleRemoteEPC.cpp
index 3d3ca891d88108..c681d8fc8c5fb1 100644
--- a/llvm/lib/ExecutionEngine/Orc/SimpleRemoteEPC.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/SimpleRemoteEPC.cpp
@@ -29,8 +29,8 @@ SimpleRemoteEPC::loadDylib(const char *DylibPath) {
   return DylibMgr->open(DylibPath, 0);
 }
 
-Expected<std::vector<tpctypes::LookupResult>>
-SimpleRemoteEPC::lookupSymbols(ArrayRef<LookupRequest> Request) {
+void SimpleRemoteEPC::lookupSymbolsAsync(ArrayRef<LookupRequest> Request,
+                                         SymbolLookupCompleteFn Complete) {
   std::vector<tpctypes::LookupResult> Result;
 
   for (auto &Element : Request) {
@@ -40,9 +40,9 @@ SimpleRemoteEPC::lookupSymbols(ArrayRef<LookupRequest> Request) {
       for (auto Addr : *R)
         Result.back().push_back(Addr);
     } else
-      return R.takeError();
+      return Complete(R.takeError());
   }
-  return std::move(Result);
+  Complete(std::move(Result));
 }
 
 Expected<int32_t> SimpleRemoteEPC::runAsMain(ExecutorAddr MainFnAddr,
diff --git a/llvm/unittests/ExecutionEngine/Orc/ObjectLinkingLayerTest.cpp b/llvm/unittests/ExecutionEngine/Orc/ObjectLinkingLayerTest.cpp
index edd12ebb62e1a4..7ab3e40df7459d 100644
--- a/llvm/unittests/ExecutionEngine/Orc/ObjectLinkingLayerTest.cpp
+++ b/llvm/unittests/ExecutionEngine/Orc/ObjectLinkingLayerTest.cpp
@@ -189,8 +189,8 @@ TEST(ObjectLinkingLayerSearchGeneratorTest, AbsoluteSymbolsObjectLayer) {
       return ExecutorAddr::fromPtr((void *)nullptr);
     }
 
-    Expected<std::vector<tpctypes::LookupResult>>
-    lookupSymbols(ArrayRef<LookupRequest> Request) override {
+    void lookupSymbolsAsync(ArrayRef<LookupRequest> Request,
+                            SymbolLookupCompleteFn Complete) override {
       std::vector<ExecutorSymbolDef> Result;
       EXPECT_EQ(Request.size(), 1u);
       for (auto &LR : Request) {
@@ -205,7 +205,7 @@ TEST(ObjectLinkingLayerSearchGeneratorTest, AbsoluteSymbolsObjectLayer) {
           }
         }
       }
-      return std::vector<tpctypes::LookupResult>{1, Result};
+      Complete(std::vector<tpctypes::LookupResult>{1, Result});
     }
   };
 

>From f49853b67522d8a8ecbe42f0d8e3b5556461fb89 Mon Sep 17 00:00:00 2001
From: Ben Langmuir <blangmuir at apple.com>
Date: Thu, 8 Feb 2024 14:20:15 -0800
Subject: [PATCH 3/3] [ORC] Add EPCGenericDylibManager::lookupAsync and use
 from SimpleRemoteEPC

Provide an asynchronous lookup API for EPCGenericDylibManager and adopt
that from the SimpleRemoteEPC. This enables an end-to-end async
EPCDynamicLibrarySearchGenerator. Note: currently we keep the current
per-dlhandle lookup model, but a future improvement could do a single
async call for a given lookup operation.
---
 .../Orc/EPCGenericDylibManager.h              | 25 +++++++++-
 .../Orc/EPCGenericDylibManager.cpp            | 50 +++++++++++--------
 .../ExecutionEngine/Orc/SimpleRemoteEPC.cpp   | 41 ++++++++++-----
 3 files changed, 82 insertions(+), 34 deletions(-)

diff --git a/llvm/include/llvm/ExecutionEngine/Orc/EPCGenericDylibManager.h b/llvm/include/llvm/ExecutionEngine/Orc/EPCGenericDylibManager.h
index 6ee2deef04d091..e0bce5d246129d 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/EPCGenericDylibManager.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/EPCGenericDylibManager.h
@@ -51,11 +51,32 @@ class EPCGenericDylibManager {
 
   /// Looks up symbols within the given dylib.
   Expected<std::vector<ExecutorSymbolDef>>
-  lookup(tpctypes::DylibHandle H, const SymbolLookupSet &Lookup);
+  lookup(tpctypes::DylibHandle H, const SymbolLookupSet &Lookup) {
+    std::promise<Expected<std::vector<ExecutorSymbolDef>>> RP;
+    auto RF = RP.get_future();
+    lookupAsync(H, Lookup, [&RP](auto R) { RP.set_value(std::move(R)); });
+    return RF.get();
+  }
 
   /// Looks up symbols within the given dylib.
   Expected<std::vector<ExecutorSymbolDef>>
-  lookup(tpctypes::DylibHandle H, const RemoteSymbolLookupSet &Lookup);
+  lookup(tpctypes::DylibHandle H, const RemoteSymbolLookupSet &Lookup) {
+    std::promise<Expected<std::vector<ExecutorSymbolDef>>> RP;
+    auto RF = RP.get_future();
+    lookupAsync(H, Lookup, [&RP](auto R) { RP.set_value(std::move(R)); });
+    return RF.get();
+  }
+
+  using SymbolLookupCompleteFn =
+      unique_function<void(Expected<std::vector<ExecutorSymbolDef>>)>;
+
+  /// Looks up symbols within the given dylib.
+  void lookupAsync(tpctypes::DylibHandle H, const SymbolLookupSet &Lookup,
+                   SymbolLookupCompleteFn Complete);
+
+  /// Looks up symbols within the given dylib.
+  void lookupAsync(tpctypes::DylibHandle H, const RemoteSymbolLookupSet &Lookup,
+                   SymbolLookupCompleteFn Complete);
 
 private:
   ExecutorProcessControl &EPC;
diff --git a/llvm/lib/ExecutionEngine/Orc/EPCGenericDylibManager.cpp b/llvm/lib/ExecutionEngine/Orc/EPCGenericDylibManager.cpp
index da185c80c6c7d3..6a7cab4a55100b 100644
--- a/llvm/lib/ExecutionEngine/Orc/EPCGenericDylibManager.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/EPCGenericDylibManager.cpp
@@ -81,28 +81,38 @@ Expected<tpctypes::DylibHandle> EPCGenericDylibManager::open(StringRef Path,
   return H;
 }
 
-Expected<std::vector<ExecutorSymbolDef>>
-EPCGenericDylibManager::lookup(tpctypes::DylibHandle H,
-                               const SymbolLookupSet &Lookup) {
-  Expected<std::vector<ExecutorSymbolDef>> Result(
-      (std::vector<ExecutorSymbolDef>()));
-  if (auto Err =
-          EPC.callSPSWrapper<rt::SPSSimpleExecutorDylibManagerLookupSignature>(
-              SAs.Lookup, Result, SAs.Instance, H, Lookup))
-    return std::move(Err);
-  return Result;
+void EPCGenericDylibManager::lookupAsync(tpctypes::DylibHandle H,
+                                         const SymbolLookupSet &Lookup,
+                                         SymbolLookupCompleteFn Complete) {
+  EPC.callSPSWrapperAsync<rt::SPSSimpleExecutorDylibManagerLookupSignature>(
+      SAs.Lookup,
+      [Complete = std::move(Complete)](
+          Error SerializationErr,
+          Expected<std::vector<ExecutorSymbolDef>> Result) mutable {
+        if (SerializationErr) {
+          cantFail(Result.takeError());
+          Complete(std::move(SerializationErr));
+        }
+        Complete(std::move(Result));
+      },
+      SAs.Instance, H, Lookup);
 }
 
-Expected<std::vector<ExecutorSymbolDef>>
-EPCGenericDylibManager::lookup(tpctypes::DylibHandle H,
-                               const RemoteSymbolLookupSet &Lookup) {
-  Expected<std::vector<ExecutorSymbolDef>> Result(
-      (std::vector<ExecutorSymbolDef>()));
-  if (auto Err =
-          EPC.callSPSWrapper<rt::SPSSimpleExecutorDylibManagerLookupSignature>(
-              SAs.Lookup, Result, SAs.Instance, H, Lookup))
-    return std::move(Err);
-  return Result;
+void EPCGenericDylibManager::lookupAsync(tpctypes::DylibHandle H,
+                                         const RemoteSymbolLookupSet &Lookup,
+                                         SymbolLookupCompleteFn Complete) {
+  EPC.callSPSWrapperAsync<rt::SPSSimpleExecutorDylibManagerLookupSignature>(
+      SAs.Lookup,
+      [Complete = std::move(Complete)](
+          Error SerializationErr,
+          Expected<std::vector<ExecutorSymbolDef>> Result) mutable {
+        if (SerializationErr) {
+          cantFail(Result.takeError());
+          Complete(std::move(SerializationErr));
+        }
+        Complete(std::move(Result));
+      },
+      SAs.Instance, H, Lookup);
 }
 
 } // end namespace orc
diff --git a/llvm/lib/ExecutionEngine/Orc/SimpleRemoteEPC.cpp b/llvm/lib/ExecutionEngine/Orc/SimpleRemoteEPC.cpp
index c681d8fc8c5fb1..a81019cb1dabb0 100644
--- a/llvm/lib/ExecutionEngine/Orc/SimpleRemoteEPC.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/SimpleRemoteEPC.cpp
@@ -29,20 +29,37 @@ SimpleRemoteEPC::loadDylib(const char *DylibPath) {
   return DylibMgr->open(DylibPath, 0);
 }
 
+/// Async helper to chain together calls to DylibMgr::lookupAsync to fulfill all
+/// all the requests.
+/// FIXME: The dylib manager should support multiple LookupRequests natively.
+static void
+lookupSymbolsAsyncHelper(EPCGenericDylibManager &DylibMgr,
+                         ArrayRef<SimpleRemoteEPC::LookupRequest> Request,
+                         std::vector<tpctypes::LookupResult> Result,
+                         SimpleRemoteEPC::SymbolLookupCompleteFn Complete) {
+  if (Request.empty())
+    return Complete(std::move(Result));
+
+  auto &Element = Request.front();
+  DylibMgr.lookupAsync(Element.Handle, Element.Symbols,
+                       [&DylibMgr, Request, Complete = std::move(Complete),
+                        Result = std::move(Result)](auto R) mutable {
+                         if (!R)
+                           return Complete(R.takeError());
+                         Result.push_back({});
+                         Result.back().reserve(R->size());
+                         for (auto Addr : *R)
+                           Result.back().push_back(Addr);
+
+                         lookupSymbolsAsyncHelper(
+                             DylibMgr, Request.drop_front(), std::move(Result),
+                             std::move(Complete));
+                       });
+}
+
 void SimpleRemoteEPC::lookupSymbolsAsync(ArrayRef<LookupRequest> Request,
                                          SymbolLookupCompleteFn Complete) {
-  std::vector<tpctypes::LookupResult> Result;
-
-  for (auto &Element : Request) {
-    if (auto R = DylibMgr->lookup(Element.Handle, Element.Symbols)) {
-      Result.push_back({});
-      Result.back().reserve(R->size());
-      for (auto Addr : *R)
-        Result.back().push_back(Addr);
-    } else
-      return Complete(R.takeError());
-  }
-  Complete(std::move(Result));
+  lookupSymbolsAsyncHelper(*DylibMgr, Request, {}, std::move(Complete));
 }
 
 Expected<int32_t> SimpleRemoteEPC::runAsMain(ExecutorAddr MainFnAddr,



More information about the llvm-commits mailing list