remove dispatchOutstandingMUs and other cleanups to scheduling

vtjnash · vtjnash · commit c4b5d9006904 · 2025-08-19T14:58:15.000Z
diff --git a/llvm/include/llvm/ExecutionEngine/Orc/Core.h b/llvm/include/llvm/ExecutionEngine/Orc/Core.h
@@ -1720,8 +1720,6 @@ class ExecutionSession {
     logAllUnhandledErrors(std::move(Err), errs(), "JIT session error: ");
   }
 
-  void dispatchOutstandingMUs();
-
   static std::unique_ptr<MaterializationResponsibility>
   createMaterializationResponsibility(ResourceTracker &RT,
                                       SymbolFlagsMap Symbols,
@@ -1831,13 +1829,6 @@ class ExecutionSession {
 
   std::vector<JITDylibSP> JDs;
 
-  // FIXME: Remove this (and runOutstandingMUs) once the linking layer works
-  //        with callbacks from asynchronous queries.
-  mutable std::recursive_mutex OutstandingMUsMutex;
-  std::vector<std::pair<std::unique_ptr<MaterializationUnit>,
-                        std::unique_ptr<MaterializationResponsibility>>>
-      OutstandingMUs;
-
   mutable std::mutex JITDispatchHandlersMutex;
   DenseMap<ExecutorAddr, std::shared_ptr<JITDispatchHandlerFunction>>
       JITDispatchHandlers;
diff --git a/llvm/include/llvm/ExecutionEngine/Orc/TaskDispatch.h b/llvm/include/llvm/ExecutionEngine/Orc/TaskDispatch.h
@@ -162,6 +162,7 @@ class LLVM_ABI DynamicThreadPoolTaskDispatcher : public TaskDispatcher {
   bool Shutdown = false;
   size_t Outstanding = 0;
   std::condition_variable OutstandingCV;
+  SmallVector<future_base *> WaitingFutures;
 
   std::optional<size_t> MaxMaterializationThreads;
   size_t NumMaterializationThreads = 0;
@@ -171,13 +172,16 @@ class LLVM_ABI DynamicThreadPoolTaskDispatcher : public TaskDispatcher {
 
 #endif // LLVM_ENABLE_THREADS
 
+/// ORC-TaskDispatch aware promise/future class that can help with task dispatch while waiting
+// TODO: docs
+
 /// Status for future/promise state
 enum class FutureStatus : uint8_t { NotReady = 0, Ready = 1, NotValid = 2 };
 
 /// Type-erased base class for futures
 class future_base {
 public:
-  bool is_ready() const {
+  bool ready() const {
     return state_->status_.load(std::memory_order_acquire) !=
            FutureStatus::NotReady;
   }
@@ -191,9 +195,9 @@ class future_base {
   /// Wait for the future to be ready, helping with task dispatch
   void wait(TaskDispatcher &D) {
     // Keep helping with task dispatch until our future is ready
-    if (!is_ready())
+    if (!ready())
       D.work_until(*this);
-    assert(is_ready());
+    assert(ready());
   }
 
 protected:
diff --git a/llvm/lib/ExecutionEngine/Orc/Core.cpp b/llvm/lib/ExecutionEngine/Orc/Core.cpp
@@ -1816,11 +1816,6 @@ void ExecutionSession::lookup(
     });
   });
 
-  // lookup can be re-entered recursively if running on a single thread. Run any
-  // outstanding MUs in case this query depends on them, otherwise this lookup
-  // will starve waiting for a result from an MU that is stuck in the queue.
-  dispatchOutstandingMUs();
-
   auto Unresolved = std::move(Symbols);
   auto Q = std::make_shared<AsynchronousSymbolQuery>(Unresolved, RequiredState,
                                                      std::move(NotifyComplete));
@@ -2156,32 +2151,6 @@ bool ExecutionSession::verifySessionState(Twine Phase) {
 }
 #endif // EXPENSIVE_CHECKS
 
-void ExecutionSession::dispatchOutstandingMUs() {
-  LLVM_DEBUG(dbgs() << "Dispatching MaterializationUnits...\n");
-  while (true) {
-    std::optional<std::pair<std::unique_ptr<MaterializationUnit>,
-                            std::unique_ptr<MaterializationResponsibility>>>
-        JMU;
-
-    {
-      std::lock_guard<std::recursive_mutex> Lock(OutstandingMUsMutex);
-      if (!OutstandingMUs.empty()) {
-        JMU.emplace(std::move(OutstandingMUs.back()));
-        OutstandingMUs.pop_back();
-      }
-    }
-
-    if (!JMU)
-      break;
-
-    assert(JMU->first && "No MU?");
-    LLVM_DEBUG(dbgs() << "  Dispatching \"" << JMU->first->getName() << "\"\n");
-    dispatchTask(std::make_unique<MaterializationTask>(std::move(JMU->first),
-                                                       std::move(JMU->second)));
-  }
-  LLVM_DEBUG(dbgs() << "Done dispatching MaterializationUnits.\n");
-}
-
 Error ExecutionSession::removeResourceTracker(ResourceTracker &RT) {
   LLVM_DEBUG({
     dbgs() << "In " << RT.getJITDylib().getName() << " removing tracker "
@@ -2755,8 +2724,6 @@ void ExecutionSession::OL_completeLookup(
 
     // Move the collected MUs to the OutstandingMUs list.
     if (!CollectedUMIs.empty()) {
-      std::lock_guard<std::recursive_mutex> Lock(OutstandingMUsMutex);
-
       LLVM_DEBUG(dbgs() << "Adding MUs to dispatch:\n");
       for (auto &KV : CollectedUMIs) {
         LLVM_DEBUG({
@@ -2768,8 +2735,9 @@ void ExecutionSession::OL_completeLookup(
           auto MR = createMaterializationResponsibility(
               *UMI->RT, std::move(UMI->MU->SymbolFlags),
               std::move(UMI->MU->InitSymbol));
-          OutstandingMUs.push_back(
-              std::make_pair(std::move(UMI->MU), std::move(MR)));
+          LLVM_DEBUG(dbgs() << "  Dispatching \"" << UMI->MU->getName() << "\"\n");
+          dispatchTask(std::make_unique<MaterializationTask>(std::move(UMI->MU),
+                                                             std::move(MR)));
         }
       }
     } else
@@ -2795,8 +2763,6 @@ void ExecutionSession::OL_completeLookup(
     LLVM_DEBUG(dbgs() << "Completing query\n");
     Q->handleComplete(*this);
   }
-
-  dispatchOutstandingMUs();
 }
 
 void ExecutionSession::OL_completeLookupFlags(
diff --git a/llvm/lib/ExecutionEngine/Orc/TaskDispatch.cpp b/llvm/lib/ExecutionEngine/Orc/TaskDispatch.cpp
@@ -40,7 +40,7 @@ void InPlaceTaskDispatcher::shutdown() {
 }
 
 void InPlaceTaskDispatcher::work_until(future_base &F) {
-  while (!F.is_ready()) {
+  while (!F.ready()) {
     // First, process any tasks in our local queue
     // Process in LIFO order (most recently added first) to avoid deadlocks
     // when tasks have dependencies on each other
@@ -54,7 +54,7 @@ void InPlaceTaskDispatcher::work_until(future_base &F) {
       {
         std::lock_guard<std::mutex> Lock(DispatchMutex);
         bool ShouldNotify = llvm::any_of(
-            WaitingFutures, [](future_base *F) { return F->is_ready(); });
+            WaitingFutures, [](future_base *F) { return F->ready(); });
         if (ShouldNotify) {
           WaitingFutures.clear();
           WorkFinishedCV.notify_all();
@@ -63,23 +63,23 @@ void InPlaceTaskDispatcher::work_until(future_base &F) {
 #endif
 
       // Check if our future is now ready
-      if (F.is_ready())
+      if (F.ready())
         return;
     }
 
     // If we get here, our queue is empty but the future isn't ready
-    // We need to wait for other threads to finish work that might complete our
+    // We need to wait for other threads to finish work that should complete our
     // future
 #if LLVM_ENABLE_THREADS
     {
       std::unique_lock<std::mutex> Lock(DispatchMutex);
       WaitingFutures.push_back(&F);
-      WorkFinishedCV.wait(Lock, [&F]() { return F.is_ready(); });
+      WorkFinishedCV.wait(Lock, [&F]() { return F.ready(); });
     }
 #else
     // Without threading, if our queue is empty and future isn't ready,
-    // we can't make progress
-    return;
+    // the library must have forgotten to schedule it, causing deadlock here
+    report_fatal_error("waiting for future that was never dispatched");
 #endif
   }
 }
@@ -140,6 +140,13 @@ void DynamicThreadPoolTaskDispatcher::dispatch(std::unique_ptr<Task> T) {
         --NumMaterializationThreads;
       --Outstanding;
 
+      bool ShouldNotify = Outstanding == 0 || llvm::any_of(
+          WaitingFutures, [](future_base *F) { return F->ready(); });
+      if (ShouldNotify) {
+        WaitingFutures.clear();
+        OutstandingCV.notify_all();
+      }
+
       if (!MaterializationTaskQueue.empty() && canRunMaterializationTaskNow()) {
         // If there are any materialization tasks running then steal that work.
         T = std::move(MaterializationTaskQueue.front());
@@ -153,8 +160,6 @@ void DynamicThreadPoolTaskDispatcher::dispatch(std::unique_ptr<Task> T) {
         TaskKind = Idle;
         ++Outstanding;
       } else {
-        if (Outstanding == 0)
-          OutstandingCV.notify_all();
         return;
       }
     }
@@ -178,9 +183,9 @@ bool DynamicThreadPoolTaskDispatcher::canRunIdleTaskNow() {
 }
 
 void DynamicThreadPoolTaskDispatcher::work_until(future_base &F) {
-  // TODO: Implement efficient work_until for DynamicThreadPoolTaskDispatcher
   std::unique_lock<std::mutex> Lock(DispatchMutex);
-  OutstandingCV.wait(Lock, [this]() { return Outstanding == 0; });
+  WaitingFutures.push_back(&F);
+  OutstandingCV.wait(Lock, [&F]() { return F.ready(); });
 }
 
 #endif
