[ORC] Generalize materialization dispatch to task dispatch.

Generalizing this API allows work to be distributed more evenly. In particular,
query callbacks can now be dispatched (rather than running immediately on the
thread that satisfied the query). This avoids the pathalogical case where an
operation on one thread satisfies many queries simultaneously, causing large
amounts of work to be run on that thread while other threads potentially sit
idle.

include/llvm/ExecutionEngine/Orc/Core.h

@@ -22,6 +22,7 @@
 #include "llvm/ExecutionEngine/Orc/SymbolStringPool.h"
 #include "llvm/ExecutionEngine/OrcV1Deprecation.h"
 #include "llvm/Support/Debug.h"
+#include "llvm/Support/ExtensibleRTTI.h"
 
 #include <atomic>
 #include <memory>
@@ -637,6 +638,8 @@ class MaterializationUnit {
   friend class JITDylib;
 
 public:
+  static char ID;
+
   MaterializationUnit(SymbolFlagsMap InitalSymbolFlags,
                       SymbolStringPtr InitSymbol)
       : SymbolFlags(std::move(InitalSymbolFlags)),
@@ -1217,6 +1220,37 @@ public:
                     const DenseMap<JITDylib *, SymbolLookupSet> &InitSyms);
 };
 
+/// Represents an abstract task for ORC to run.
+class Task : public RTTIExtends<Task, RTTIRoot> {
+public:
+  static char ID;
+
+  /// Description of the task to be performed. Used for logging.
+  virtual void printDescription(raw_ostream &OS) = 0;
+
+  /// Run the task.
+  virtual void run() = 0;
+
+private:
+  void anchor() override;
+};
+
+/// A materialization task.
+class MaterializationTask : public RTTIExtends<MaterializationTask, Task> {
+public:
+  static char ID;
+
+  MaterializationTask(std::unique_ptr<MaterializationUnit> MU,
+                      std::unique_ptr<MaterializationResponsibility> MR)
+      : MU(std::move(MU)), MR(std::move(MR)) {}
+  void printDescription(raw_ostream &OS) override;
+  void run() override;
+
+private:
+  std::unique_ptr<MaterializationUnit> MU;
+  std::unique_ptr<MaterializationResponsibility> MR;
+};
+
 /// An ExecutionSession represents a running JIT program.
 class ExecutionSession {
   friend class InProgressLookupFlagsState;
@@ -1230,10 +1264,8 @@ public:
   /// For reporting errors.
   using ErrorReporter = std::function<void(Error)>;
 
-  /// For dispatching MaterializationUnit::materialize calls.
-  using DispatchMaterializationFunction =
-      std::function<void(std::unique_ptr<MaterializationUnit> MU,
-                         std::unique_ptr<MaterializationResponsibility> MR)>;
+  /// For dispatching ORC tasks (typically materialization tasks).
+  using DispatchTaskFunction = std::function<void(std::unique_ptr<Task> T)>;
 
   /// Construct an ExecutionSession.
   ///
@@ -1306,10 +1338,9 @@ public:
   /// Unhandled errors can be sent here to log them.
   void reportError(Error Err) { ReportError(std::move(Err)); }
 
-  /// Set the materialization dispatch function.
-  ExecutionSession &setDispatchMaterialization(
-      DispatchMaterializationFunction DispatchMaterialization) {
-    this->DispatchMaterialization = std::move(DispatchMaterialization);
+  /// Set the task dispatch function.
+  ExecutionSession &setDispatchTask(DispatchTaskFunction DispatchTask) {
+    this->DispatchTask = std::move(DispatchTask);
     return *this;
   }
 
@@ -1384,12 +1415,10 @@ public:
          SymbolState RequiredState = SymbolState::Ready);
 
   /// Materialize the given unit.
-  void
-  dispatchMaterialization(std::unique_ptr<MaterializationUnit> MU,
-                          std::unique_ptr<MaterializationResponsibility> MR) {
-    assert(MU && "MU must be non-null");
-    DEBUG_WITH_TYPE("orc", dumpDispatchInfo(MR->getTargetJITDylib(), *MU));
-    DispatchMaterialization(std::move(MU), std::move(MR));
+  void dispatchTask(std::unique_ptr<Task> T) {
+    assert(T && "T must be non-null");
+    DEBUG_WITH_TYPE("orc", dumpDispatchInfo(*T));
+    DispatchTask(std::move(T));
   }
 
   /// Dump the state of all the JITDylibs in this session.
@@ -1400,11 +1429,7 @@ private:
     logAllUnhandledErrors(std::move(Err), errs(), "JIT session error: ");
   }
 
-  static void materializeOnCurrentThread(
-      std::unique_ptr<MaterializationUnit> MU,
-      std::unique_ptr<MaterializationResponsibility> MR) {
-    MU->materialize(std::move(MR));
-  }
+  static void runOnCurrentThread(std::unique_ptr<Task> T) { T->run(); }
 
   void dispatchOutstandingMUs();
 
@@ -1474,7 +1499,7 @@ private:
                                 const SymbolDependenceMap &Dependencies);
 
 #ifndef NDEBUG
-  void dumpDispatchInfo(JITDylib &JD, MaterializationUnit &MU);
+  void dumpDispatchInfo(Task &T);
 #endif // NDEBUG
 
   mutable std::recursive_mutex SessionMutex;
@@ -1482,8 +1507,7 @@ private:
   std::shared_ptr<SymbolStringPool> SSP;
   std::unique_ptr<Platform> P;
   ErrorReporter ReportError = logErrorsToStdErr;
-  DispatchMaterializationFunction DispatchMaterialization =
-      materializeOnCurrentThread;
+  DispatchTaskFunction DispatchTask = runOnCurrentThread;
 
   std::vector<ResourceManager *> ResourceManagers;
 

lib/ExecutionEngine/Orc/Core.cpp

@@ -29,6 +29,8 @@ char SymbolsNotFound::ID = 0;
 char SymbolsCouldNotBeRemoved::ID = 0;
 char MissingSymbolDefinitions::ID = 0;
 char UnexpectedSymbolDefinitions::ID = 0;
+char Task::ID = 0;
+char MaterializationTask::ID = 0;
 
 RegisterDependenciesFunction NoDependenciesToRegister =
     RegisterDependenciesFunction();
@@ -750,7 +752,8 @@ Error JITDylib::replace(MaterializationResponsibility &FromMR,
 
   if (MustRunMU) {
     assert(MustRunMR && "MustRunMU set implies MustRunMR set");
-    ES.dispatchMaterialization(std::move(MustRunMU), std::move(MustRunMR));
+    ES.dispatchTask(std::make_unique<MaterializationTask>(
+        std::move(MustRunMU), std::move(MustRunMR)));
   } else {
     assert(!MustRunMR && "MustRunMU unset implies MustRunMR unset");
   }
@@ -1730,6 +1733,15 @@ Expected<DenseMap<JITDylib *, SymbolMap>> Platform::lookupInitSymbols(
   return std::move(CompoundResult);
 }
 
+void Task::anchor() {}
+
+void MaterializationTask::printDescription(raw_ostream &OS) {
+  OS << "Materialization task: " << MU->getName() << " in "
+     << MR->getTargetJITDylib().getName() << "\n";
+}
+
+void MaterializationTask::run() { MU->materialize(std::move(MR)); }
+
 ExecutionSession::ExecutionSession(std::shared_ptr<SymbolStringPool> SSP)
     : SSP(SSP ? std::move(SSP) : std::make_shared<SymbolStringPool>()) {}
 
@@ -2003,7 +2015,8 @@ void ExecutionSession::dispatchOutstandingMUs() {
 
     assert(JMU->first && "No MU?");
     LLVM_DEBUG(dbgs() << "  Dispatching \"" << JMU->first->getName() << "\"\n");
-    dispatchMaterialization(std::move(JMU->first), std::move(JMU->second));
+    dispatchTask(std::make_unique<MaterializationTask>(std::move(JMU->first),
+                                                       std::move(JMU->second)));
   }
   LLVM_DEBUG(dbgs() << "Done dispatching MaterializationUnits.\n");
 }
@@ -2776,9 +2789,10 @@ void ExecutionSession::OL_addDependenciesForAll(
 }
 
 #ifndef NDEBUG
-void ExecutionSession::dumpDispatchInfo(JITDylib &JD, MaterializationUnit &MU) {
+void ExecutionSession::dumpDispatchInfo(Task &T) {
   runSessionLocked([&]() {
-    dbgs() << "Dispatching " << MU << " for " << JD.getName() << "\n";
+    dbgs() << "Dispatching: ";
+    T.printDescription(dbgs());
   });
 }
 #endif // NDEBUG

lib/ExecutionEngine/Orc/LLJIT.cpp

@@ -1126,20 +1126,16 @@ LLJIT::LLJIT(LLJITBuilderState &S, Error &Err)
     InitHelperTransformLayer->setCloneToNewContextOnEmit(true);
     CompileThreads =
         std::make_unique<ThreadPool>(hardware_concurrency(S.NumCompileThreads));
-    ES->setDispatchMaterialization(
-        [this](std::unique_ptr<MaterializationUnit> MU,
-               std::unique_ptr<MaterializationResponsibility> MR) {
-          // FIXME: We should be able to use move-capture here, but ThreadPool's
-          // AsyncTaskTys are std::functions rather than unique_functions
-          // (because MSVC's std::packaged_tasks don't support move-only types).
-          // Fix this when all the above gets sorted out.
-          CompileThreads->async(
-              [UnownedMU = MU.release(), UnownedMR = MR.release()]() mutable {
-                std::unique_ptr<MaterializationUnit> MU(UnownedMU);
-                std::unique_ptr<MaterializationResponsibility> MR(UnownedMR);
-                MU->materialize(std::move(MR));
-              });
-        });
+    ES->setDispatchTask([this](std::unique_ptr<Task> T) {
+      // FIXME: We should be able to use move-capture here, but ThreadPool's
+      // AsyncTaskTys are std::functions rather than unique_functions
+      // (because MSVC's std::packaged_tasks don't support move-only types).
+      // Fix this when all the above gets sorted out.
+      CompileThreads->async([UnownedT = T.release()]() mutable {
+        std::unique_ptr<Task> T(UnownedT);
+        T->run();
+      });
+    });
   }
 
   if (S.SetUpPlatform)

unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp

@@ -1019,12 +1019,11 @@ TEST_F(CoreAPIsStandardTest, TestBasicWeakSymbolMaterialization) {
 
 TEST_F(CoreAPIsStandardTest, DefineMaterializingSymbol) {
   bool ExpectNoMoreMaterialization = false;
-  ES.setDispatchMaterialization(
-      [&](std::unique_ptr<MaterializationUnit> MU,
-          std::unique_ptr<MaterializationResponsibility> MR) {
+  ES.setDispatchTask(
+      [&](std::unique_ptr<Task> T) {
         if (ExpectNoMoreMaterialization)
           ADD_FAILURE() << "Unexpected materialization";
-        MU->materialize(std::move(MR));
+        T->run();
       });
 
   auto MU = std::make_unique<SimpleMaterializationUnit>(
@@ -1252,12 +1251,11 @@ TEST_F(CoreAPIsStandardTest, TestLookupWithThreadedMaterialization) {
 #if LLVM_ENABLE_THREADS
 
   std::thread MaterializationThread;
-  ES.setDispatchMaterialization(
-      [&](std::unique_ptr<MaterializationUnit> MU,
-          std::unique_ptr<MaterializationResponsibility> MR) {
+  ES.setDispatchTask(
+      [&](std::unique_ptr<Task> T) {
         MaterializationThread =
-            std::thread([MU = std::move(MU), MR = std::move(MR)]() mutable {
-              MU->materialize(std::move(MR));
+            std::thread([T = std::move(T)]() mutable {
+              T->run();
             });
       });