[ORC][llvm-jitlink] Add SimpleLazyReexportsSpeculator, use in llvm-jitlink.

Also adds a new IdleTask type and updates DynamicThreadPoolTaskDispatcher to
schedule IdleTasks whenever the total number of threads running is less than
the maximum number of MaterializationThreads.

A SimpleLazyReexportsSpeculator instance maintains a list of speculation
suggestions ((JITDylib, Function) pairs) and registered lazy reexports. When
speculation opportunities are available (having been added via
addSpeculationSuggestions or when lazy reexports were created) it schedules
an IdleTask that triggers the next speculative lookup as soon as resources
are available. Speculation suggestions are processed first, followed by
lookups for lazy reexport bodies. A callback can be registered at object
construction time to record lazy reexport executions as they happen, and these
executions can be fed back into the speculator as suggestions on subsequent
executions.

The llvm-jitlink tool is updated to support speculation when lazy linking is
used via three new arguments:

 -speculate=[none|simple] : When the 'simple' value is specified a
                            SimpleLazyReexportsSpeculator instances is used
                            for speculation.

 -speculate-order <path> : Specifies a path to a CSV containing
                           (jit dylib name, function name) triples to use
                           as speculative suggestions in the current run.

 -record-lazy-execs <path> : Specifies a path in which to record lazy function
                             executions as a CSV of (jit dylib name, function
                             name) pairs, suitable for use with
                             -speculate-order.

The same path can be passed to -speculate-order and -record-lazy-execs, in
which case the file will be overwritten at the end of the execution.

No testcase yet: Speculative linking is difficult to test (since by definition
execution behavior should be unaffected by speculation) and this is an new
prototype of the concept*. Tests will be added in the future once the interface
and behavior settle down.

* An earlier implementation of the speculation concept can be found in
  llvm/include/llvm/ExecutionEngine/Orc/Speculation.h. Both systems have the
  same goal (hiding compilation latency) but different mechanisms. This patch
  relies entirely on information available in the controller, where the old
  system could receive additional information from the JIT'd runtime via
  callbacks. I aim to combine the two in the future, but want to gain more
  practical experience with speculation first.

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,66 @@ lazyReexports(LazyReexportsManager &LRM, SymbolAliasMap Reexports) {
   return LRM.createLazyReexports(std::move(Reexports));
 }
 
+class SimpleLazyReexportsSpeculator : public LazyReexportsManager::Listener {
+  friend std::shared_ptr<SimpleLazyReexportsSpeculator> std::make_shared();
+
+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