[llvm] 2d10b7b - Reapply "[ORC][llvm-jitlink] Add SimpleLazyReexportsSpeculator..." with fixes.

Lang Hames via llvm-commits llvm-commits at lists.llvm.org
Thu Jan 9 17:18:44 PST 2025


Author: Lang Hames
Date: 2025-01-10T12:16:21+11:00
New Revision: 2d10b7b750f97b42055d5b9b08a88c18ff811cd2

URL: https://github.com/llvm/llvm-project/commit/2d10b7b750f97b42055d5b9b08a88c18ff811cd2
DIFF: https://github.com/llvm/llvm-project/commit/2d10b7b750f97b42055d5b9b08a88c18ff811cd2.diff

LOG: Reapply "[ORC][llvm-jitlink] Add SimpleLazyReexportsSpeculator..." with fixes.

This reapplies 6d72bf47606, which was reverted in 57447d3ddf to investigate
build failures, e.g. https://lab.llvm.org/buildbot/#/builders/3/builds/10114.

The original patch contained an invalid unused friend declaration of
std::make_shared. This has been removed.

Added: 
    

Modified: 
    llvm/include/llvm/ExecutionEngine/Orc/JITLinkReentryTrampolines.h
    llvm/include/llvm/ExecutionEngine/Orc/LazyReexports.h
    llvm/include/llvm/ExecutionEngine/Orc/TaskDispatch.h
    llvm/lib/ExecutionEngine/Orc/JITLinkReentryTrampolines.cpp
    llvm/lib/ExecutionEngine/Orc/LazyReexports.cpp
    llvm/lib/ExecutionEngine/Orc/TaskDispatch.cpp
    llvm/tools/llvm-jitlink/llvm-jitlink.cpp
    llvm/tools/llvm-jitlink/llvm-jitlink.h

Removed: 
    


################################################################################
diff  --git a/llvm/include/llvm/ExecutionEngine/Orc/JITLinkReentryTrampolines.h b/llvm/include/llvm/ExecutionEngine/Orc/JITLinkReentryTrampolines.h
index 673019b748b3b2..94d2ff0717eb4d 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/JITLinkReentryTrampolines.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/JITLinkReentryTrampolines.h
@@ -65,7 +65,8 @@ class JITLinkReentryTrampolines {
 Expected<std::unique_ptr<LazyReexportsManager>>
 createJITLinkLazyReexportsManager(ObjectLinkingLayer &ObjLinkingLayer,
                                   RedirectableSymbolManager &RSMgr,
-                                  JITDylib &PlatformJD);
+                                  JITDylib &PlatformJD,
+                                  LazyReexportsManager::Listener *L = nullptr);
 
 } // namespace llvm::orc
 

diff  --git a/llvm/include/llvm/ExecutionEngine/Orc/LazyReexports.h b/llvm/include/llvm/ExecutionEngine/Orc/LazyReexports.h
index c6b5d08544b1ef..635f2b08367e67 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/LazyReexports.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/LazyReexports.h
@@ -179,6 +179,37 @@ class LazyReexportsManager : public ResourceManager {
   lazyReexports(LazyReexportsManager &, SymbolAliasMap);
 
 public:
+  struct CallThroughInfo {
+    JITDylibSP JD;
+    SymbolStringPtr Name;
+    SymbolStringPtr BodyName;
+  };
+
+  class Listener {
+  public:
+    using CallThroughInfo = LazyReexportsManager::CallThroughInfo;
+
+    virtual ~Listener();
+
+    /// Called under the session lock when new lazy reexports are created.
+    virtual void onLazyReexportsCreated(JITDylib &JD, ResourceKey K,
+                                        const SymbolAliasMap &Reexports) = 0;
+
+    /// Called under the session lock when lazy reexports have their ownership
+    /// transferred to a new ResourceKey.
+    virtual void onLazyReexportsTransfered(JITDylib &JD, ResourceKey DstK,
+                                           ResourceKey SrcK) = 0;
+
+    /// Called under the session lock when lazy reexports are removed.
+    virtual Error onLazyReexportsRemoved(JITDylib &JD, ResourceKey K) = 0;
+
+    /// Called outside the session lock when a lazy reexport is called.
+    /// NOTE: Since this is called outside the session lock there is a chance
+    ///       that the reexport referred to has already been removed. Listeners
+    ///       must be prepared to handle requests for stale reexports.
+    virtual void onLazyReexportCalled(const CallThroughInfo &CTI) = 0;
+  };
+
   using OnTrampolinesReadyFn = unique_function<void(
       Expected<std::vector<ExecutorSymbolDef>> EntryAddrs)>;
   using EmitTrampolinesFn =
@@ -189,7 +220,7 @@ class LazyReexportsManager : public ResourceManager {
   /// This will work both in-process and out-of-process.
   static Expected<std::unique_ptr<LazyReexportsManager>>
   Create(EmitTrampolinesFn EmitTrampolines, RedirectableSymbolManager &RSMgr,
-         JITDylib &PlatformJD);
+         JITDylib &PlatformJD, Listener *L = nullptr);
 
   LazyReexportsManager(LazyReexportsManager &&) = delete;
   LazyReexportsManager &operator=(LazyReexportsManager &&) = delete;
@@ -199,12 +230,6 @@ class LazyReexportsManager : public ResourceManager {
                                ResourceKey SrcK) override;
 
 private:
-  struct CallThroughInfo {
-    SymbolStringPtr Name;
-    SymbolStringPtr BodyName;
-    JITDylibSP JD;
-  };
-
   class MU;
   class Plugin;
 
@@ -213,7 +238,7 @@ class LazyReexportsManager : public ResourceManager {
 
   LazyReexportsManager(EmitTrampolinesFn EmitTrampolines,
                        RedirectableSymbolManager &RSMgr, JITDylib &PlatformJD,
-                       Error &Err);
+                       Listener *L, Error &Err);
 
   std::unique_ptr<MaterializationUnit>
   createLazyReexports(SymbolAliasMap Reexports);
@@ -229,6 +254,7 @@ class LazyReexportsManager : public ResourceManager {
   ExecutionSession &ES;
   EmitTrampolinesFn EmitTrampolines;
   RedirectableSymbolManager &RSMgr;
+  Listener *L;
 
   DenseMap<ResourceKey, std::vector<ExecutorAddr>> KeyToReentryAddrs;
   DenseMap<ExecutorAddr, CallThroughInfo> CallThroughs;
@@ -242,6 +268,64 @@ lazyReexports(LazyReexportsManager &LRM, SymbolAliasMap Reexports) {
   return LRM.createLazyReexports(std::move(Reexports));
 }
 
+class SimpleLazyReexportsSpeculator : public LazyReexportsManager::Listener {
+public:
+  using RecordExecutionFunction =
+      unique_function<void(const CallThroughInfo &CTI)>;
+
+  static std::shared_ptr<SimpleLazyReexportsSpeculator>
+  Create(ExecutionSession &ES, RecordExecutionFunction RecordExec = {}) {
+    class make_shared_helper : public SimpleLazyReexportsSpeculator {
+    public:
+      make_shared_helper(ExecutionSession &ES,
+                         RecordExecutionFunction RecordExec)
+          : SimpleLazyReexportsSpeculator(ES, std::move(RecordExec)) {}
+    };
+
+    auto Instance =
+        std::make_shared<make_shared_helper>(ES, std::move(RecordExec));
+    Instance->WeakThis = Instance;
+    return Instance;
+  }
+
+  SimpleLazyReexportsSpeculator(SimpleLazyReexportsSpeculator &&) = delete;
+  SimpleLazyReexportsSpeculator &
+  operator=(SimpleLazyReexportsSpeculator &&) = delete;
+  ~SimpleLazyReexportsSpeculator() override;
+
+  void onLazyReexportsCreated(JITDylib &JD, ResourceKey K,
+                              const SymbolAliasMap &Reexports) override;
+
+  void onLazyReexportsTransfered(JITDylib &JD, ResourceKey DstK,
+                                 ResourceKey SrcK) override;
+
+  Error onLazyReexportsRemoved(JITDylib &JD, ResourceKey K) override;
+
+  void onLazyReexportCalled(const CallThroughInfo &CTI) override;
+
+  void addSpeculationSuggestions(
+      std::vector<std::pair<std::string, SymbolStringPtr>> NewSuggestions);
+
+private:
+  SimpleLazyReexportsSpeculator(ExecutionSession &ES,
+                                RecordExecutionFunction RecordExec)
+      : ES(ES), RecordExec(std::move(RecordExec)) {}
+
+  bool doNextSpeculativeLookup();
+
+  class SpeculateTask;
+
+  using KeyToFunctionBodiesMap =
+      DenseMap<ResourceKey, std::vector<SymbolStringPtr>>;
+
+  ExecutionSession &ES;
+  RecordExecutionFunction RecordExec;
+  std::weak_ptr<SimpleLazyReexportsSpeculator> WeakThis;
+  DenseMap<JITDylib *, KeyToFunctionBodiesMap> LazyReexports;
+  std::deque<std::pair<std::string, SymbolStringPtr>> SpeculateSuggestions;
+  bool SpeculateTaskActive = false;
+};
+
 } // End namespace orc
 } // End namespace llvm
 

diff  --git a/llvm/include/llvm/ExecutionEngine/Orc/TaskDispatch.h b/llvm/include/llvm/ExecutionEngine/Orc/TaskDispatch.h
index d7939864fd609b..67164679f9e302 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/TaskDispatch.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/TaskDispatch.h
@@ -92,6 +92,16 @@ makeGenericNamedTask(FnT &&Fn, const char *Desc = nullptr) {
                                                      Desc);
 }
 
+/// IdleTask can be used as the basis for low-priority tasks, e.g. speculative
+/// lookup.
+class IdleTask : public RTTIExtends<IdleTask, Task> {
+public:
+  static char ID;
+
+private:
+  void anchor() override;
+};
+
 /// Abstract base for classes that dispatch ORC Tasks.
 class TaskDispatcher {
 public:
@@ -118,9 +128,13 @@ class DynamicThreadPoolTaskDispatcher : public TaskDispatcher {
   DynamicThreadPoolTaskDispatcher(
       std::optional<size_t> MaxMaterializationThreads)
       : MaxMaterializationThreads(MaxMaterializationThreads) {}
+
   void dispatch(std::unique_ptr<Task> T) override;
   void shutdown() override;
 private:
+  bool canRunMaterializationTaskNow();
+  bool canRunIdleTaskNow();
+
   std::mutex DispatchMutex;
   bool Shutdown = false;
   size_t Outstanding = 0;
@@ -129,6 +143,7 @@ class DynamicThreadPoolTaskDispatcher : public TaskDispatcher {
   std::optional<size_t> MaxMaterializationThreads;
   size_t NumMaterializationThreads = 0;
   std::deque<std::unique_ptr<Task>> MaterializationTaskQueue;
+  std::deque<std::unique_ptr<Task>> IdleTaskQueue;
 };
 
 #endif // LLVM_ENABLE_THREADS

diff  --git a/llvm/lib/ExecutionEngine/Orc/JITLinkReentryTrampolines.cpp b/llvm/lib/ExecutionEngine/Orc/JITLinkReentryTrampolines.cpp
index be574ef7279c27..d9f859654d41e9 100644
--- a/llvm/lib/ExecutionEngine/Orc/JITLinkReentryTrampolines.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/JITLinkReentryTrampolines.cpp
@@ -173,7 +173,8 @@ void JITLinkReentryTrampolines::emit(ResourceTrackerSP RT,
 Expected<std::unique_ptr<LazyReexportsManager>>
 createJITLinkLazyReexportsManager(ObjectLinkingLayer &ObjLinkingLayer,
                                   RedirectableSymbolManager &RSMgr,
-                                  JITDylib &PlatformJD) {
+                                  JITDylib &PlatformJD,
+                                  LazyReexportsManager::Listener *L) {
   auto JLT = JITLinkReentryTrampolines::Create(ObjLinkingLayer);
   if (!JLT)
     return JLT.takeError();
@@ -184,7 +185,7 @@ createJITLinkLazyReexportsManager(ObjectLinkingLayer &ObjLinkingLayer,
                                   OnTrampolinesReady) mutable {
         JLT->emit(std::move(RT), NumTrampolines, std::move(OnTrampolinesReady));
       },
-      RSMgr, PlatformJD);
+      RSMgr, PlatformJD, L);
 }
 
 } // namespace llvm::orc

diff  --git a/llvm/lib/ExecutionEngine/Orc/LazyReexports.cpp b/llvm/lib/ExecutionEngine/Orc/LazyReexports.cpp
index df4539ba4329ff..7b38621eba8245 100644
--- a/llvm/lib/ExecutionEngine/Orc/LazyReexports.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/LazyReexports.cpp
@@ -280,31 +280,34 @@ class LazyReexportsManager::Plugin : public ObjectLinkingLayer::Plugin {
   std::mutex M;
 };
 
+LazyReexportsManager::Listener::~Listener() = default;
+
 Expected<std::unique_ptr<LazyReexportsManager>>
 LazyReexportsManager::Create(EmitTrampolinesFn EmitTrampolines,
                              RedirectableSymbolManager &RSMgr,
-                             JITDylib &PlatformJD) {
+                             JITDylib &PlatformJD, Listener *L) {
   Error Err = Error::success();
   std::unique_ptr<LazyReexportsManager> LRM(new LazyReexportsManager(
-      std::move(EmitTrampolines), RSMgr, PlatformJD, Err));
+      std::move(EmitTrampolines), RSMgr, PlatformJD, L, Err));
   if (Err)
     return std::move(Err);
   return std::move(LRM);
 }
 
 Error LazyReexportsManager::handleRemoveResources(JITDylib &JD, ResourceKey K) {
-  JD.getExecutionSession().runSessionLocked([&]() {
+  return JD.getExecutionSession().runSessionLocked([&]() -> Error {
     auto I = KeyToReentryAddrs.find(K);
-    if (I != KeyToReentryAddrs.end()) {
-      auto &ReentryAddrs = I->second;
-      for (auto &ReentryAddr : ReentryAddrs) {
-        assert(CallThroughs.count(ReentryAddr) && "CallTrhough missing");
-        CallThroughs.erase(ReentryAddr);
-      }
-      KeyToReentryAddrs.erase(I);
+    if (I == KeyToReentryAddrs.end())
+      return Error::success();
+
+    auto &ReentryAddrs = I->second;
+    for (auto &ReentryAddr : ReentryAddrs) {
+      assert(CallThroughs.count(ReentryAddr) && "CallTrhough missing");
+      CallThroughs.erase(ReentryAddr);
     }
+    KeyToReentryAddrs.erase(I);
+    return L ? L->onLazyReexportsRemoved(JD, K) : Error::success();
   });
-  return Error::success();
 }
 
 void LazyReexportsManager::handleTransferResources(JITDylib &JD,
@@ -323,14 +326,17 @@ void LazyReexportsManager::handleTransferResources(JITDylib &JD,
       DstAddrs.insert(DstAddrs.end(), SrcAddrs.begin(), SrcAddrs.end());
       KeyToReentryAddrs.erase(I);
     }
+    if (L)
+      L->onLazyReexportsTransfered(JD, DstK, SrcK);
   }
 }
 
 LazyReexportsManager::LazyReexportsManager(EmitTrampolinesFn EmitTrampolines,
                                            RedirectableSymbolManager &RSMgr,
-                                           JITDylib &PlatformJD, Error &Err)
+                                           JITDylib &PlatformJD, Listener *L,
+                                           Error &Err)
     : ES(PlatformJD.getExecutionSession()),
-      EmitTrampolines(std::move(EmitTrampolines)), RSMgr(RSMgr) {
+      EmitTrampolines(std::move(EmitTrampolines)), RSMgr(RSMgr), L(L) {
 
   using namespace shared;
 
@@ -384,17 +390,22 @@ void LazyReexportsManager::emitRedirectableSymbols(
     Redirs[Name] = (*ReentryPoints)[I++];
 
   I = 0;
-  if (auto Err = MR->withResourceKeyDo([&](ResourceKey K) {
-        for (auto &[Name, AI] : Reexports) {
-          const auto &ReentryPoint = (*ReentryPoints)[I++];
-          CallThroughs[ReentryPoint.getAddress()] = {Name, AI.Aliasee,
-                                                     &MR->getTargetJITDylib()};
-          KeyToReentryAddrs[K].push_back(ReentryPoint.getAddress());
-        }
-      })) {
-    MR->getExecutionSession().reportError(std::move(Err));
-    MR->failMaterialization();
-    return;
+  if (!Reexports.empty()) {
+    if (auto Err = MR->withResourceKeyDo([&](ResourceKey K) {
+          auto &JD = MR->getTargetJITDylib();
+          auto &ReentryAddrsForK = KeyToReentryAddrs[K];
+          for (auto &[Name, AI] : Reexports) {
+            const auto &ReentryPoint = (*ReentryPoints)[I++];
+            CallThroughs[ReentryPoint.getAddress()] = {&JD, Name, AI.Aliasee};
+            ReentryAddrsForK.push_back(ReentryPoint.getAddress());
+          }
+          if (L)
+            L->onLazyReexportsCreated(JD, K, Reexports);
+        })) {
+      MR->getExecutionSession().reportError(std::move(Err));
+      MR->failMaterialization();
+      return;
+    }
   }
 
   RSMgr.emitRedirectableSymbols(std::move(MR), std::move(Redirs));
@@ -415,6 +426,9 @@ void LazyReexportsManager::resolve(ResolveSendResultFn SendResult,
     LandingInfo = I->second;
   });
 
+  if (L)
+    L->onLazyReexportCalled(LandingInfo);
+
   SymbolInstance LandingSym(LandingInfo.JD, std::move(LandingInfo.BodyName));
   LandingSym.lookupAsync([this, JD = std::move(LandingInfo.JD),
                           ReentryName = std::move(LandingInfo.Name),
@@ -432,5 +446,168 @@ void LazyReexportsManager::resolve(ResolveSendResultFn SendResult,
   });
 }
 
+class SimpleLazyReexportsSpeculator::SpeculateTask : public IdleTask {
+public:
+  SpeculateTask(std::weak_ptr<SimpleLazyReexportsSpeculator> Speculator)
+      : Speculator(std::move(Speculator)) {}
+
+  void printDescription(raw_ostream &OS) override {
+    OS << "Speculative Lookup Task";
+  }
+
+  void run() override {
+    if (auto S = Speculator.lock())
+      S->doNextSpeculativeLookup();
+  }
+
+private:
+  std::weak_ptr<SimpleLazyReexportsSpeculator> Speculator;
+};
+
+SimpleLazyReexportsSpeculator::~SimpleLazyReexportsSpeculator() {
+  for (auto &[JD, _] : LazyReexports)
+    JITDylibSP(JD)->Release();
+}
+
+void SimpleLazyReexportsSpeculator::onLazyReexportsCreated(
+    JITDylib &JD, ResourceKey K, const SymbolAliasMap &Reexports) {
+  if (!LazyReexports.count(&JD))
+    JD.Retain();
+  auto &BodiesVec = LazyReexports[&JD][K];
+  for (auto &[Name, AI] : Reexports)
+    BodiesVec.push_back(AI.Aliasee);
+  if (!SpeculateTaskActive) {
+    SpeculateTaskActive = true;
+    ES.dispatchTask(std::make_unique<SpeculateTask>(WeakThis));
+  }
+}
+
+void SimpleLazyReexportsSpeculator::onLazyReexportsTransfered(
+    JITDylib &JD, ResourceKey DstK, ResourceKey SrcK) {
+
+  auto I = LazyReexports.find(&JD);
+  if (I == LazyReexports.end())
+    return;
+
+  auto &MapForJD = I->second;
+  auto J = MapForJD.find(SrcK);
+  if (J == MapForJD.end())
+    return;
+
+  // We have something to transfer.
+  auto K = MapForJD.find(DstK);
+  if (K == MapForJD.end()) {
+    auto Tmp = std::move(J->second);
+    MapForJD.erase(J);
+    MapForJD[DstK] = std::move(Tmp);
+  } else {
+    auto &SrcNames = J->second;
+    auto &DstNames = K->second;
+    DstNames.insert(DstNames.end(), SrcNames.begin(), SrcNames.end());
+    MapForJD.erase(J);
+  }
+}
+
+Error SimpleLazyReexportsSpeculator::onLazyReexportsRemoved(JITDylib &JD,
+                                                            ResourceKey K) {
+
+  auto I = LazyReexports.find(&JD);
+  if (I == LazyReexports.end())
+    return Error::success();
+
+  auto &MapForJD = I->second;
+  MapForJD.erase(K);
+
+  if (MapForJD.empty()) {
+    LazyReexports.erase(I);
+    JD.Release();
+  }
+
+  return Error::success();
+}
+
+void SimpleLazyReexportsSpeculator::onLazyReexportCalled(
+    const CallThroughInfo &CTI) {
+  if (RecordExec)
+    RecordExec(CTI);
+}
+
+void SimpleLazyReexportsSpeculator::addSpeculationSuggestions(
+    std::vector<std::pair<std::string, SymbolStringPtr>> NewSuggestions) {
+  ES.runSessionLocked([&]() {
+    for (auto &[JDName, SymbolName] : NewSuggestions)
+      SpeculateSuggestions.push_back(
+          {std::move(JDName), std::move(SymbolName)});
+  });
+}
+
+bool SimpleLazyReexportsSpeculator::doNextSpeculativeLookup() {
+  // Use existing speculation queue if available, otherwise take the next
+  // element from LazyReexports.
+  JITDylibSP SpeculateJD = nullptr;
+  SymbolStringPtr SpeculateFn;
+
+  auto SpeculateAgain = ES.runSessionLocked([&]() {
+    while (!SpeculateSuggestions.empty()) {
+      auto [JDName, SymbolName] = std::move(SpeculateSuggestions.front());
+      SpeculateSuggestions.pop_front();
+
+      if (auto *JD = ES.getJITDylibByName(JDName)) {
+        SpeculateJD = JD;
+        SpeculateFn = std::move(SymbolName);
+        break;
+      }
+    }
+
+    if (!SpeculateJD) {
+      assert(!LazyReexports.empty() && "LazyReexports map is empty");
+      auto LRItr =
+          std::next(LazyReexports.begin(), rand() % LazyReexports.size());
+      auto &[JD, KeyToFnBodies] = *LRItr;
+
+      assert(!KeyToFnBodies.empty() && "Key to function bodies map empty");
+      auto KeyToFnBodiesItr =
+          std::next(KeyToFnBodies.begin(), rand() % KeyToFnBodies.size());
+      auto &[Key, FnBodies] = *KeyToFnBodiesItr;
+
+      assert(!FnBodies.empty() && "Function bodies list empty");
+      auto FnBodyItr = std::next(FnBodies.begin(), rand() % FnBodies.size());
+
+      SpeculateJD = JITDylibSP(JD);
+      SpeculateFn = std::move(*FnBodyItr);
+
+      FnBodies.erase(FnBodyItr);
+      if (FnBodies.empty()) {
+        KeyToFnBodies.erase(KeyToFnBodiesItr);
+        if (KeyToFnBodies.empty()) {
+          LRItr->first->Release();
+          LazyReexports.erase(LRItr);
+        }
+      }
+    }
+
+    SpeculateTaskActive =
+        !SpeculateSuggestions.empty() || !LazyReexports.empty();
+    return SpeculateTaskActive;
+  });
+
+  LLVM_DEBUG({
+    dbgs() << "Issuing speculative lookup for ( " << SpeculateJD->getName()
+           << ", " << SpeculateFn << " )...\n";
+  });
+
+  ES.lookup(
+      LookupKind::Static, makeJITDylibSearchOrder(SpeculateJD.get()),
+      {{std::move(SpeculateFn), SymbolLookupFlags::WeaklyReferencedSymbol}},
+      SymbolState::Ready,
+      [](Expected<SymbolMap> Result) { consumeError(Result.takeError()); },
+      NoDependenciesToRegister);
+
+  if (SpeculateAgain)
+    ES.dispatchTask(std::make_unique<SpeculateTask>(WeakThis));
+
+  return false;
+}
+
 } // End namespace orc.
 } // End namespace llvm.

diff  --git a/llvm/lib/ExecutionEngine/Orc/TaskDispatch.cpp b/llvm/lib/ExecutionEngine/Orc/TaskDispatch.cpp
index 1af17e85220db5..e87a14f3ea7c4f 100644
--- a/llvm/lib/ExecutionEngine/Orc/TaskDispatch.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/TaskDispatch.cpp
@@ -15,9 +15,13 @@ namespace orc {
 
 char Task::ID = 0;
 char GenericNamedTask::ID = 0;
+char IdleTask::ID = 0;
+
 const char *GenericNamedTask::DefaultDescription = "Generic Task";
 
 void Task::anchor() {}
+void IdleTask::anchor() {}
+
 TaskDispatcher::~TaskDispatcher() = default;
 
 void InPlaceTaskDispatcher::dispatch(std::unique_ptr<Task> T) { T->run(); }
@@ -26,7 +30,15 @@ void InPlaceTaskDispatcher::shutdown() {}
 
 #if LLVM_ENABLE_THREADS
 void DynamicThreadPoolTaskDispatcher::dispatch(std::unique_ptr<Task> T) {
-  bool IsMaterializationTask = isa<MaterializationTask>(*T);
+
+  enum { Normal, Materialization, Idle } TaskKind;
+
+  if (isa<MaterializationTask>(*T))
+    TaskKind = Materialization;
+  else if (isa<IdleTask>(*T))
+    TaskKind = Idle;
+  else
+    TaskKind = Normal;
 
   {
     std::lock_guard<std::mutex> Lock(DispatchMutex);
@@ -35,24 +47,24 @@ void DynamicThreadPoolTaskDispatcher::dispatch(std::unique_ptr<Task> T) {
     if (Shutdown)
       return;
 
-    if (IsMaterializationTask) {
+    if (TaskKind == Materialization) {
 
       // If this is a materialization task and there are too many running
       // already then queue this one up and return early.
-      if (MaxMaterializationThreads &&
-          NumMaterializationThreads == *MaxMaterializationThreads) {
-        MaterializationTaskQueue.push_back(std::move(T));
-        return;
-      }
+      if (!canRunMaterializationTaskNow())
+        return MaterializationTaskQueue.push_back(std::move(T));
 
       // Otherwise record that we have a materialization task running.
       ++NumMaterializationThreads;
+    } else if (TaskKind == Idle) {
+      if (!canRunIdleTaskNow())
+        return IdleTaskQueue.push_back(std::move(T));
     }
 
     ++Outstanding;
   }
 
-  std::thread([this, T = std::move(T), IsMaterializationTask]() mutable {
+  std::thread([this, T = std::move(T), TaskKind]() mutable {
     while (true) {
 
       // Run the task.
@@ -67,18 +79,24 @@ void DynamicThreadPoolTaskDispatcher::dispatch(std::unique_ptr<Task> T) {
       // Check the work queue state and either proceed with the next task or
       // end this thread.
       std::lock_guard<std::mutex> Lock(DispatchMutex);
-      if (!MaterializationTaskQueue.empty()) {
+
+      if (TaskKind == Materialization)
+        --NumMaterializationThreads;
+      --Outstanding;
+
+      if (!MaterializationTaskQueue.empty() && canRunMaterializationTaskNow()) {
         // If there are any materialization tasks running then steal that work.
         T = std::move(MaterializationTaskQueue.front());
         MaterializationTaskQueue.pop_front();
-        if (!IsMaterializationTask) {
-          ++NumMaterializationThreads;
-          IsMaterializationTask = true;
-        }
+        TaskKind = Materialization;
+        ++NumMaterializationThreads;
+        ++Outstanding;
+      } else if (!IdleTaskQueue.empty() && canRunIdleTaskNow()) {
+        T = std::move(IdleTaskQueue.front());
+        IdleTaskQueue.pop_front();
+        TaskKind = Idle;
+        ++Outstanding;
       } else {
-        if (IsMaterializationTask)
-          --NumMaterializationThreads;
-        --Outstanding;
         if (Outstanding == 0)
           OutstandingCV.notify_all();
         return;
@@ -92,6 +110,17 @@ void DynamicThreadPoolTaskDispatcher::shutdown() {
   Shutdown = true;
   OutstandingCV.wait(Lock, [this]() { return Outstanding == 0; });
 }
+
+bool DynamicThreadPoolTaskDispatcher::canRunMaterializationTaskNow() {
+  return !MaxMaterializationThreads ||
+         (NumMaterializationThreads < *MaxMaterializationThreads);
+}
+
+bool DynamicThreadPoolTaskDispatcher::canRunIdleTaskNow() {
+  return !MaxMaterializationThreads ||
+         (Outstanding < *MaxMaterializationThreads);
+}
+
 #endif
 
 } // namespace orc

diff  --git a/llvm/tools/llvm-jitlink/llvm-jitlink.cpp b/llvm/tools/llvm-jitlink/llvm-jitlink.cpp
index 963c36322c8ab8..ece8484fe8c7c6 100644
--- a/llvm/tools/llvm-jitlink/llvm-jitlink.cpp
+++ b/llvm/tools/llvm-jitlink/llvm-jitlink.cpp
@@ -83,13 +83,35 @@ using namespace llvm::orc;
 
 static cl::OptionCategory JITLinkCategory("JITLink Options");
 
+static cl::list<std::string> InputFiles(cl::Positional, cl::OneOrMore,
+                                        cl::desc("input files"),
+                                        cl::cat(JITLinkCategory));
+
 static cl::list<bool> LazyLink("lazy",
                                cl::desc("Link the following file lazily"),
                                cl::cat(JITLinkCategory));
 
-static cl::list<std::string> InputFiles(cl::Positional, cl::OneOrMore,
-                                        cl::desc("input files"),
-                                        cl::cat(JITLinkCategory));
+enum class SpeculateKind { None, Simple };
+
+cl::opt<SpeculateKind> Speculate(
+    "speculate", cl::desc("Choose speculation scheme"),
+    cl::init(SpeculateKind::None),
+    cl::values(clEnumValN(SpeculateKind::None, "none", "No speculation"),
+               clEnumValN(SpeculateKind::Simple, "simple",
+                          "Simple speculation")),
+    cl::cat(JITLinkCategory));
+
+cl::opt<std::string> SpeculateOrder(
+    "speculate-order",
+    cl::desc("A CSV file containing (JITDylib, Function) pairs to"
+             "speculatively look up"),
+    cl::cat(JITLinkCategory));
+
+cl::opt<std::string> RecordLazyExecs(
+    "record-lazy-execs",
+    cl::desc("Write lazy-function executions to a CSV file as (JITDylib, "
+             "function) pairs"),
+    cl::cat(JITLinkCategory));
 
 static cl::opt<size_t> MaterializationThreads(
     "num-threads", cl::desc("Number of materialization threads to use"),
@@ -959,17 +981,50 @@ class PhonyExternalsGenerator : public DefinitionGenerator {
 };
 
 Expected<std::unique_ptr<Session::LazyLinkingSupport>>
-createLazyLinkingSupport(ObjectLinkingLayer &OLL, JITDylib &PlatformJD) {
-  auto RSMgr = JITLinkRedirectableSymbolManager::Create(OLL);
+createLazyLinkingSupport(Session &S) {
+  auto RSMgr = JITLinkRedirectableSymbolManager::Create(S.ObjLayer);
   if (!RSMgr)
     return RSMgr.takeError();
 
-  auto LRMgr = createJITLinkLazyReexportsManager(OLL, **RSMgr, PlatformJD);
+  std::shared_ptr<SimpleLazyReexportsSpeculator> Speculator;
+  switch (Speculate) {
+  case SpeculateKind::None:
+    break;
+  case SpeculateKind::Simple:
+    SimpleLazyReexportsSpeculator::RecordExecutionFunction RecordExecs;
+
+    if (!RecordLazyExecs.empty())
+      RecordExecs = [&S](const LazyReexportsManager::CallThroughInfo &CTI) {
+        S.LazyFnExecOrder.push_back({CTI.JD->getName(), CTI.BodyName});
+      };
+
+    Speculator =
+        SimpleLazyReexportsSpeculator::Create(S.ES, std::move(RecordExecs));
+    break;
+  }
+
+  auto LRMgr = createJITLinkLazyReexportsManager(
+      S.ObjLayer, **RSMgr, *S.PlatformJD, Speculator.get());
   if (!LRMgr)
     return LRMgr.takeError();
 
-  return std::make_unique<Session::LazyLinkingSupport>(std::move(*RSMgr),
-                                                       std::move(*LRMgr), OLL);
+  return std::make_unique<Session::LazyLinkingSupport>(
+      std::move(*RSMgr), std::move(Speculator), std::move(*LRMgr), S.ObjLayer);
+}
+
+static Error writeLazyExecOrder(Session &S) {
+  if (RecordLazyExecs.empty())
+    return Error::success();
+
+  std::error_code EC;
+  raw_fd_ostream ExecOrderOut(RecordLazyExecs, EC);
+  if (EC)
+    return createFileError(RecordLazyExecs, EC);
+
+  for (auto &[JDName, FunctionName] : S.LazyFnExecOrder)
+    ExecOrderOut << JDName << ", " << FunctionName << "\n";
+
+  return Error::success();
 }
 
 Expected<std::unique_ptr<Session>> Session::Create(Triple TT,
@@ -1017,8 +1072,7 @@ Expected<std::unique_ptr<Session>> Session::Create(Triple TT,
   S->Features = std::move(Features);
 
   if (lazyLinkingRequested()) {
-    if (auto LazyLinking =
-            createLazyLinkingSupport(S->ObjLayer, *S->PlatformJD))
+    if (auto LazyLinking = createLazyLinkingSupport(*S))
       S->LazyLinking = std::move(*LazyLinking);
     else
       return LazyLinking.takeError();
@@ -1028,6 +1082,9 @@ Expected<std::unique_ptr<Session>> Session::Create(Triple TT,
 }
 
 Session::~Session() {
+  if (auto Err = writeLazyExecOrder(*this))
+    ES.reportError(std::move(Err));
+
   if (auto Err = ES.endSession())
     ES.reportError(std::move(Err));
 }
@@ -1702,6 +1759,23 @@ static Error sanitizeArguments(const Triple &TT, const char *ArgV0) {
       return make_error<StringError>(
           "Lazy linking cannot be used with -harness mode",
           inconvertibleErrorCode());
+  } else if (Speculate != SpeculateKind::None) {
+    errs() << "Warning: -speculate ignored as there are no -lazy inputs\n";
+    Speculate = SpeculateKind::None;
+  }
+
+  if (Speculate == SpeculateKind::None) {
+    if (!SpeculateOrder.empty()) {
+      errs() << "Warning: -speculate-order ignored because speculation is "
+                "disabled\n";
+      SpeculateOrder = "";
+    }
+
+    if (!RecordLazyExecs.empty()) {
+      errs() << "Warning: -record-lazy-execs ignored because speculation is "
+                "disabled\n";
+      RecordLazyExecs = "";
+    }
   }
 
   return Error::success();
@@ -2267,6 +2341,59 @@ static Error addLibraries(Session &S,
   return Error::success();
 }
 
+static Error addSpeculationOrder(Session &S) {
+
+  if (SpeculateOrder.empty())
+    return Error::success();
+
+  assert(S.LazyLinking && "SpeculateOrder set, but lazy linking not enabled");
+  assert(S.LazyLinking->Speculator && "SpeculatoOrder set, but no speculator");
+
+  auto SpecOrderBuffer = getFile(SpeculateOrder);
+  if (!SpecOrderBuffer)
+    return SpecOrderBuffer.takeError();
+
+  StringRef LineStream((*SpecOrderBuffer)->getBuffer());
+  std::vector<std::pair<std::string, SymbolStringPtr>> SpecOrder;
+
+  size_t LineNumber = 0;
+  while (!LineStream.empty()) {
+    ++LineNumber;
+
+    auto MakeSpecOrderErr = [&](StringRef Reason) {
+      return make_error<StringError>("Error in speculation order file \"" +
+                                         SpeculateOrder + "\" on line " +
+                                         Twine(LineNumber) + ": " + Reason,
+                                     inconvertibleErrorCode());
+    };
+
+    StringRef CurLine;
+    std::tie(CurLine, LineStream) = LineStream.split('\n');
+    CurLine = CurLine.trim();
+    if (CurLine.empty())
+      continue;
+
+    auto [JDName, FuncName] = CurLine.split(',');
+
+    if (FuncName.empty())
+      return MakeSpecOrderErr("missing ',' separator");
+
+    JDName = JDName.trim();
+    if (JDName.empty())
+      return MakeSpecOrderErr("no value for column 1 (JIT Dylib name)");
+
+    FuncName = FuncName.trim();
+    if (FuncName.empty())
+      return MakeSpecOrderErr("no value for column 2 (function name)");
+
+    SpecOrder.push_back({JDName.str(), S.ES.intern(FuncName)});
+  }
+
+  S.LazyLinking->Speculator->addSpeculationSuggestions(std::move(SpecOrder));
+
+  return Error::success();
+}
+
 static Error addSessionInputs(Session &S) {
   std::map<unsigned, JITDylib *> IdxToJD;
   DenseSet<unsigned> LazyLinkIdxs;
@@ -2299,6 +2426,9 @@ static Error addSessionInputs(Session &S) {
   if (auto Err = addLibraries(S, IdxToJD, LazyLinkIdxs))
     return Err;
 
+  if (auto Err = addSpeculationOrder(S))
+    return Err;
+
   return Error::success();
 }
 

diff  --git a/llvm/tools/llvm-jitlink/llvm-jitlink.h b/llvm/tools/llvm-jitlink/llvm-jitlink.h
index be3710971729a4..92d667e797b881 100644
--- a/llvm/tools/llvm-jitlink/llvm-jitlink.h
+++ b/llvm/tools/llvm-jitlink/llvm-jitlink.h
@@ -33,13 +33,17 @@ namespace llvm {
 struct Session {
 
   struct LazyLinkingSupport {
-    LazyLinkingSupport(std::unique_ptr<orc::RedirectableSymbolManager> RSMgr,
-                       std::unique_ptr<orc::LazyReexportsManager> LRMgr,
-                       orc::ObjectLinkingLayer &ObjLinkingLayer)
-        : RSMgr(std::move(RSMgr)), LRMgr(std::move(LRMgr)),
+    LazyLinkingSupport(
+        std::unique_ptr<orc::RedirectableSymbolManager> RSMgr,
+        std::shared_ptr<orc::SimpleLazyReexportsSpeculator> Speculator,
+        std::unique_ptr<orc::LazyReexportsManager> LRMgr,
+        orc::ObjectLinkingLayer &ObjLinkingLayer)
+        : RSMgr(std::move(RSMgr)), Speculator(std::move(Speculator)),
+          LRMgr(std::move(LRMgr)),
           LazyObjLinkingLayer(ObjLinkingLayer, *this->LRMgr) {}
 
     std::unique_ptr<orc::RedirectableSymbolManager> RSMgr;
+    std::shared_ptr<orc::SimpleLazyReexportsSpeculator> Speculator;
     std::unique_ptr<orc::LazyReexportsManager> LRMgr;
     orc::LazyObjectLinkingLayer LazyObjLinkingLayer;
   };
@@ -52,6 +56,7 @@ struct Session {
   std::unique_ptr<LazyLinkingSupport> LazyLinking;
   orc::JITDylibSearchOrder JDSearchOrder;
   SubtargetFeatures Features;
+  std::vector<std::pair<std::string, orc::SymbolStringPtr>> LazyFnExecOrder;
 
   ~Session();
 


        


More information about the llvm-commits mailing list