more design enhancements, docs

Author: vtjnash

llvm/include/llvm/ExecutionEngine/Orc/TaskDispatch.h

@@ -14,6 +14,7 @@
 #define LLVM_EXECUTIONENGINE_ORC_TASKDISPATCH_H
 
 #include "llvm/Config/llvm-config.h"
+#include "llvm/ADT/BitmaskEnum.h"
 #include "llvm/Support/Compiler.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/ErrorHandling.h"
@@ -115,7 +116,7 @@ class LLVM_ABI TaskDispatcher {
 public:
   virtual ~TaskDispatcher();
 
-  /// Run the given task.
+  /// Schedule the given task to run.
   virtual void dispatch(std::unique_ptr<Task> T) = 0;
 
   /// Called by ExecutionSession. Waits until all tasks have completed.
@@ -125,15 +126,21 @@ class LLVM_ABI TaskDispatcher {
   virtual void work_until(future_base &F) = 0;
 };
 
-/// Runs all tasks on the current thread.
+/// Runs all tasks on the current thread, at the next work_until yield point.
 class LLVM_ABI InPlaceTaskDispatcher : public TaskDispatcher {
 public:
   void dispatch(std::unique_ptr<Task> T) override;
   void shutdown() override;
   void work_until(future_base &F) override;
 
 private:
-  thread_local static SmallVector<std::unique_ptr<Task>> TaskQueue;
+  /// C++ does not support non-static thread_local variables, so this needs to
+  /// store both the task and the associated dispatcher queue so that shutdown
+  /// can wait for the correct tasks to finish.
+#if LLVM_ENABLE_THREADS
+  thread_local static
+#endif
+    SmallVector<std::pair<std::unique_ptr<Task>, InPlaceTaskDispatcher*>> TaskQueue;
 
 #if LLVM_ENABLE_THREADS
   std::mutex DispatchMutex;
@@ -172,32 +179,55 @@ 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
+/// @name ORC Promise/Future Classes
+/// 
+/// ORC-aware promise/future implementation that integrates with the
+/// TaskDispatcher system to allow efficient cooperative multitasking while
+/// waiting for results (with certain limitations on what can be awaited).
+/// Together they provide building blocks for a full async/await-like runtime
+/// for llvm that supports multiple threads.
+/// 
+/// Unlike std::promise/std::future alone, these classes can help dispatch other
+/// tasks while waiting, preventing deadlocks and improving overall system
+/// throughput. They have a similar API, though with some important differences
+/// and some features simply not currently implemented.
+///
+/// @{
 
 /// Status for future/promise state
-enum class FutureStatus : uint8_t { NotReady = 0, Ready = 1, NotValid = 2 };
+enum class FutureStatus : uint8_t { NotReady = 0, Ready = 1 };
+
+/// Status for promise state tracking
+enum class PromiseStatus : uint8_t { 
+  Initial = 0, 
+  FutureCreated = 1, 
+  ValueSet = 2,
+  LLVM_MARK_AS_BITMASK_ENUM(ValueSet)
+};
 
-/// Type-erased base class for futures
+/// Type-erased base class for futures, generally for scheduler use to avoid
+/// needing virtual dispatches
 class future_base {
 public:
+  /// Check if the future is now ready with a value (precondition: should be valid)
   bool ready() const {
     return state_->status_.load(std::memory_order_acquire) !=
            FutureStatus::NotReady;
   }
 
   /// Check if the future is in a valid state (not moved-from and not consumed)
   bool valid() const {
-    return state_ && state_->status_.load(std::memory_order_acquire) !=
-                         FutureStatus::NotValid;
+    return state_ != nullptr;
   }
 
   /// 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 (!ready())
+    if (!ready()) {
       D.work_until(*this);
-    assert(ready());
+      if (state_->status_.load(std::memory_order_relaxed) != FutureStatus::Ready)
+        report_fatal_error("work_until() returned without this future being ready");
+    }
   }
 
 protected:
@@ -207,19 +237,22 @@ class future_base {
 
   future_base(state_base *state) : state_(state) {}
   future_base() = default;
+
+  /// Only allow deleting the future once it is invalid
   ~future_base() {
     if (valid())
       report_fatal_error("get() must be called before future destruction");
-    delete state_;
+    // state_ is already nullptr if get() was called, otherwise we have an error above
   }
 
   // Move constructor and assignment
   future_base(future_base &&other) noexcept : state_(other.state_) {
     other.state_ = nullptr;
   }
+
   future_base &operator=(future_base &&other) noexcept {
     if (this != &other) {
-      delete state_;
+      this->~future_base();
       state_ = other.state_;
       other.state_ = nullptr;
     }
@@ -229,12 +262,33 @@ class future_base {
   state_base *state_;
 };
 
-/// ORC-aware future class that can help with task dispatch while waiting
+/// TaskDispatcher-aware future class for cooperative await.
+///
+/// @tparam T The type of value this future will provide. Use void for futures that
+///           signal completion without providing a value.
+///
+/// This future implementation is similar to `std::future`, so most code can
+/// transition to it easily`. However, it differs from `std::future` in a few key
+/// ways to be aware of:
+/// - No exception support (or the overhead for it).
+/// - Waiting operations (get(&D), wait(&D)) help dispatch other tasks while
+///   blocked, requiring an additional argument of which TaskDispatcher object of
+///   where all associated work will be scheduled.
+/// - Must call get() exactly once before destruction (enforced with
+///   `report_fatal_error`). Internal state is freed when `get` returns.
+///
+/// Other notable features, in common with `std::future`:
+/// - Supports both value types and void specialization through the same interface.
+/// - Thread-safe through atomic operations.
+/// - Provides acquire-release ordering with `std::promise::set_value()`.
+/// - Concurrent access to any method (including to `ready`) on multiple threads
+///   is not allowed.
 
 template <typename T> class future;
 template <typename T> class promise;
 template <typename T> class future : public future_base {
 public:
+  // Template the state struct so that future<void> has no wasted overhead for the value
   struct state : public future_base::state_base {
     template <typename U> struct value_storage {
       U value_;
@@ -254,15 +308,11 @@ template <typename T> class future : public future_base {
   future &operator=(future &&) = default;
 
   /// Get the value, helping with task dispatch while waiting.
-  /// This will destroy the underlying value, so this must only be called once.
+  /// This will destroy the underlying value, so this must be called exactly once.
   T get(TaskDispatcher &D) {
     if (!valid())
       report_fatal_error("get() must only be called once");
     wait(D);
-    auto old_status = state_->status_.exchange(FutureStatus::NotValid,
-                                               std::memory_order_release);
-    if (old_status != FutureStatus::Ready)
-      report_fatal_error("get() must only be called once");
     return take_value();
   }
 
@@ -271,56 +321,93 @@ template <typename T> class future : public future_base {
 
   template <typename U = T>
   typename std::enable_if<!std::is_void<U>::value, U>::type take_value() {
-    return std::move(
+    T result = std::move(
         static_cast<typename future<T>::state *>(state_)->storage.value_);
+    delete state_;
+    state_ = nullptr;
+    return result;
   }
 
   template <typename U = T>
-  typename std::enable_if<std::is_void<U>::value, U>::type take_value() {}
+  typename std::enable_if<std::is_void<U>::value, U>::type take_value() {
+    delete state_;
+    state_ = nullptr;
+  }
 
   explicit future(state *state) : future_base(state) {}
 };
 
-/// ORC-aware promise class that works with ORC future
+/// TaskDispatcher-aware promise class that provides values to associated futures.
+///
+/// @tparam T The type of value this promise will provide. Use void for promises that
+///           signal completion without providing a value.
+///
+/// This promise implementation provides the value-setting side of the promise/future
+/// pair and integrates with the ORC TaskDispatcher system. Key characteristics:
+/// - Must call get_future() no more than once to create the associated future.
+/// - Must call set_value() exactly once to provide the result.
+/// - Automatically cleans up state if get_future() is never called.
+/// - Thread-safe from set_value to get.
+/// - Move-only semantics to prevent accidental copying.
+///
+/// @par Error Handling:
+/// The promise/future system uses report_fatal_error() for misuse:
+/// - Calling get_future() more than once.
+/// - Calling set_value() more than once.
+/// - Destroying a future without calling get().
+/// - Calling get() more than once on a future.
+///
+/// @par Thread Safety:
+/// - Each promise/future must only be accessed by one thread, as concurrent
+///   calls to the API functions may result in crashes.
+/// - Multiple threads can safely access different promise/future pairs.
+/// - set_value() and get() operations are atomic and thread-safe.
+/// - Move operations should only be performed by a single thread.
 template <typename T> class promise {
   friend class future<T>;
 
 public:
-  promise() : state_(new typename future<T>::state()), future_created_(false) {}
+  promise() : state_(new typename future<T>::state()), status_(PromiseStatus::Initial) {}
 
   ~promise() {
+    // Assert proper promise lifecycle: ensure set_value was called if
+    // get_future() was called. This can catch deadlocks where get_future() is
+    // called but set_value() is never called, though only if the promise is
+    // moved from instead of borrowed from the frame with the future.
+    assert((bool)(status_ & PromiseStatus::ValueSet) == (bool)(status_ & PromiseStatus::FutureCreated) && 
+           "Promise destroyed with mismatched ValueSet and FutureCreated state");
     // Delete state only if get_future() was never called
-    if (!future_created_) {
+    if (!(status_ & PromiseStatus::FutureCreated)) {
       delete state_;
+    } else {
+      assert((bool)(status_ & PromiseStatus::ValueSet));
     }
   }
 
   promise(const promise &) = delete;
   promise &operator=(const promise &) = delete;
 
   promise(promise &&other) noexcept
-      : state_(other.state_), future_created_(other.future_created_) {
+      : state_(other.state_), status_(other.status_) {
     other.state_ = nullptr;
-    other.future_created_ = false;
+    other.status_ = PromiseStatus::Initial;
   }
 
   promise &operator=(promise &&other) noexcept {
     if (this != &other) {
-      if (!future_created_) {
-        delete state_;
-      }
+      this->~promise();
       state_ = other.state_;
-      future_created_ = other.future_created_;
+      status_ = other.status_;
       other.state_ = nullptr;
-      other.future_created_ = false;
+      other.status_ = PromiseStatus::Initial;
     }
     return *this;
   }
 
   /// Get the associated future (must only be called once)
   future<T> get_future() {
-    assert(!future_created_ && "get_future() can only be called once");
-    future_created_ = true;
+    assert(!(status_ & PromiseStatus::FutureCreated) && "get_future() can only be called once");
+    status_ |= PromiseStatus::FutureCreated;
     return future<T>(state_);
   }
 
@@ -334,15 +421,19 @@ template <typename T> class promise {
   set_value(const typename std::conditional<std::is_void<T>::value,
                                             std::nullopt_t, T>::type &value) {
     assert(state_ && "Invalid promise state");
+    assert(!(status_ & PromiseStatus::ValueSet) && "set_value() can only be called once");
     state_->storage.value_ = value;
+    status_ |= PromiseStatus::ValueSet;
     state_->status_.store(FutureStatus::Ready, std::memory_order_release);
   }
 
   template <typename U = T>
   void set_value(typename std::conditional<std::is_void<T>::value,
                                            std::nullopt_t, T>::type &&value) {
     assert(state_ && "Invalid promise state");
+    assert(!(status_ & PromiseStatus::ValueSet) && "set_value() can only be called once");
     state_->storage.value_ = std::move(value);
+    status_ |= PromiseStatus::ValueSet;
     state_->status_.store(FutureStatus::Ready, std::memory_order_release);
   }
 
@@ -357,21 +448,25 @@ template <typename T> class promise {
   template <typename U = T>
   typename std::enable_if<std::is_void<U>::value, void>::type set_value() {
     assert(state_ && "Invalid promise state");
+    assert(!(status_ & PromiseStatus::ValueSet) && "set_value() can only be called once");
+    status_ |= PromiseStatus::ValueSet;
     state_->status_.store(FutureStatus::Ready, std::memory_order_release);
   }
 
   /// Swap with another promise
   void swap(promise &other) noexcept {
     using std::swap;
     swap(state_, other.state_);
-    swap(future_created_, other.future_created_);
+    swap(status_, other.status_);
   }
 
 private:
   typename future<T>::state *state_;
-  bool future_created_;
+  PromiseStatus status_;
 };
 
+/// @}
+
 } // End namespace orc
 } // End namespace llvm
 

llvm/lib/ExecutionEngine/Orc/Core.cpp

@@ -2735,7 +2735,8 @@ void ExecutionSession::OL_completeLookup(
           auto MR = createMaterializationResponsibility(
               *UMI->RT, std::move(UMI->MU->SymbolFlags),
               std::move(UMI->MU->InitSymbol));
-          LLVM_DEBUG(dbgs() << "  Dispatching \"" << UMI->MU->getName() << "\"\n");
+          LLVM_DEBUG(dbgs()
+                     << "  Dispatching \"" << UMI->MU->getName() << "\"\n");
           dispatchTask(std::make_unique<MaterializationTask>(std::move(UMI->MU),
                                                              std::move(MR)));
         }