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

This reapplies 6d72bf4, which was reverted in 57447d3 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.

Author: lhames

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
 

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
 

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

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