cowait

llvm/include/llvm/ExecutionEngine/JITLink/JITLinkMemoryManager.h

@@ -20,6 +20,7 @@
 #include "llvm/ExecutionEngine/Orc/Shared/ExecutorAddress.h"
 #include "llvm/ExecutionEngine/Orc/Shared/MemoryFlags.h"
 #include "llvm/ExecutionEngine/Orc/SymbolStringPool.h"
+#include "llvm/ExecutionEngine/Orc/TaskDispatch.h"
 #include "llvm/Support/Allocator.h"
 #include "llvm/Support/Compiler.h"
 #include "llvm/Support/Error.h"
@@ -34,6 +35,11 @@
 #include <mutex>
 
 namespace llvm {
+
+namespace orc {
+class ExecutorProcessControl;
+}
+
 namespace jitlink {
 
 class Block;
@@ -132,11 +138,11 @@ class LLVM_ABI JITLinkMemoryManager {
     /// Called to transfer working memory to the target and apply finalization.
     virtual void finalize(OnFinalizedFunction OnFinalized) = 0;
 
-    /// Synchronous convenience version of finalize.
-    Expected<FinalizedAlloc> finalize() {
-      std::promise<MSVCPExpected<FinalizedAlloc>> FinalizeResultP;
-      auto FinalizeResultF = FinalizeResultP.get_future();
-      finalize([&](Expected<FinalizedAlloc> Result) {
+    /// Co-synchronous convenience version of finalize.
+    Expected<FinalizedAlloc> finalize(orc::TaskDispatcher &D) {
+      orc::future<MSVCPExpected<FinalizedAlloc>> FinalizeResultF;
+      finalize([FinalizeResultP = FinalizeResultF.get_promise(D)](
+                   Expected<FinalizedAlloc> Result) mutable {
         FinalizeResultP.set_value(std::move(Result));
       });
       return FinalizeResultF.get();
@@ -163,13 +169,14 @@ class LLVM_ABI JITLinkMemoryManager {
   virtual void allocate(const JITLinkDylib *JD, LinkGraph &G,
                         OnAllocatedFunction OnAllocated) = 0;
 
-  /// Convenience function for blocking allocation.
-  AllocResult allocate(const JITLinkDylib *JD, LinkGraph &G) {
-    std::promise<MSVCPExpected<std::unique_ptr<InFlightAlloc>>> AllocResultP;
-    auto AllocResultF = AllocResultP.get_future();
-    allocate(JD, G, [&](AllocResult Alloc) {
-      AllocResultP.set_value(std::move(Alloc));
-    });
+  /// Convenience function for co-blocking allocation.
+  AllocResult allocate(const JITLinkDylib *JD, LinkGraph &G,
+                       orc::TaskDispatcher &D) {
+    orc::future<MSVCPExpected<std::unique_ptr<InFlightAlloc>>> AllocResultF;
+    allocate(JD, G,
+             [AllocResultP = AllocResultF.get_promise(D)](AllocResult Alloc) {
+               AllocResultP.set_value(std::move(Alloc));
+             });
     return AllocResultF.get();
   }
 
@@ -187,20 +194,21 @@ class LLVM_ABI JITLinkMemoryManager {
     deallocate(std::move(Allocs), std::move(OnDeallocated));
   }
 
-  /// Convenience function for blocking deallocation.
-  Error deallocate(std::vector<FinalizedAlloc> Allocs) {
-    std::promise<MSVCPError> DeallocResultP;
-    auto DeallocResultF = DeallocResultP.get_future();
+  /// Convenience function for co-blocking deallocation.
+  Error deallocate(std::vector<FinalizedAlloc> Allocs, orc::TaskDispatcher &D) {
+    orc::future<MSVCPError> DeallocResultF;
     deallocate(std::move(Allocs),
-               [&](Error Err) { DeallocResultP.set_value(std::move(Err)); });
+               [DeallocResultP = DeallocResultF.get_promise(D)](Error Err) {
+                 DeallocResultP.set_value(std::move(Err));
+               });
     return DeallocResultF.get();
   }
 
-  /// Convenience function for blocking deallocation of a single alloc.
-  Error deallocate(FinalizedAlloc Alloc) {
+  /// Convenience function for co-blocking deallocation of a single alloc.
+  Error deallocate(FinalizedAlloc Alloc, orc::TaskDispatcher &D) {
     std::vector<FinalizedAlloc> Allocs;
     Allocs.push_back(std::move(Alloc));
-    return deallocate(std::move(Allocs));
+    return deallocate(std::move(Allocs), D);
   }
 };
 
@@ -326,10 +334,20 @@ class SimpleSegmentAlloc {
                               Triple TT, const JITLinkDylib *JD,
                               SegmentMap Segments, OnCreatedFunction OnCreated);
 
+  // The blocking version of this should be deprecated, and requires an
+  // TaskDispatcher for co-async correctness.
   LLVM_ABI static Expected<SimpleSegmentAlloc>
   Create(JITLinkMemoryManager &MemMgr,
          std::shared_ptr<orc::SymbolStringPool> SSP, Triple TT,
-         const JITLinkDylib *JD, SegmentMap Segments);
+         const JITLinkDylib *JD, SegmentMap Segments, orc::TaskDispatcher &D) {
+    orc::future<MSVCPExpected<SimpleSegmentAlloc>> AllocF;
+    Create(
+        MemMgr, std::move(SSP), std::move(TT), JD, std::move(Segments),
+        [AllocP = AllocF.get_promise(D)](Expected<SimpleSegmentAlloc> Result) {
+          AllocP.set_value(std::move(Result));
+        });
+    return AllocF.get();
+  }
 
   LLVM_ABI SimpleSegmentAlloc(SimpleSegmentAlloc &&);
   LLVM_ABI SimpleSegmentAlloc &operator=(SimpleSegmentAlloc &&);
@@ -343,9 +361,10 @@ class SimpleSegmentAlloc {
     Alloc->finalize(std::move(OnFinalized));
   }
 
-  /// Finalize all groups.
-  Expected<JITLinkMemoryManager::FinalizedAlloc> finalize() {
-    return Alloc->finalize();
+  /// Finalize all groups (deprecated co-blocking version).
+  Expected<JITLinkMemoryManager::FinalizedAlloc>
+  finalize(orc::TaskDispatcher &D) {
+    return Alloc->finalize(D);
   }
 
 private:

llvm/include/llvm/ExecutionEngine/Orc/COFFPlatform.h

@@ -22,7 +22,7 @@
 #include "llvm/ExecutionEngine/Orc/Shared/ExecutorAddress.h"
 #include "llvm/Support/Compiler.h"
 
-#include <future>
+#include "llvm/ExecutionEngine/Orc/TaskDispatch.h"
 #include <list>
 #include <memory>
 #include <thread>

llvm/include/llvm/ExecutionEngine/Orc/COFFVCRuntimeSupport.h

@@ -20,7 +20,7 @@
 #include "llvm/ExecutionEngine/Orc/Shared/ExecutorAddress.h"
 #include "llvm/Support/Compiler.h"
 
-#include <future>
+#include "llvm/ExecutionEngine/Orc/TaskDispatch.h"
 #include <memory>
 #include <thread>
 #include <vector>

llvm/include/llvm/ExecutionEngine/Orc/Core.h

@@ -31,9 +31,9 @@
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/ExtensibleRTTI.h"
 
+#include "llvm/ExecutionEngine/Orc/TaskDispatch.h"
 #include <atomic>
 #include <deque>
-#include <future>
 #include <memory>
 #include <vector>
 
@@ -1707,8 +1707,6 @@ class ExecutionSession {
     logAllUnhandledErrors(std::move(Err), errs(), "JIT session error: ");
   }
 
-  void dispatchOutstandingMUs();
-
   static std::unique_ptr<MaterializationResponsibility>
   createMaterializationResponsibility(ResourceTracker &RT,
                                       SymbolFlagsMap Symbols,
@@ -1815,13 +1813,6 @@ class ExecutionSession {
   std::vector<JITDylibSP> JDs;
   WaitingOnGraph G;
 
-  // 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;

llvm/include/llvm/ExecutionEngine/Orc/DylibManager.h

@@ -13,21 +13,26 @@
 #ifndef LLVM_EXECUTIONENGINE_ORC_DYLIBMANAGER_H
 #define LLVM_EXECUTIONENGINE_ORC_DYLIBMANAGER_H
 
+#include "llvm/ExecutionEngine/Orc/ExecutorProcessControl.h"
+#include "llvm/ExecutionEngine/Orc/InProcessMemoryAccess.h"
 #include "llvm/ExecutionEngine/Orc/Shared/TargetProcessControlTypes.h"
+#include "llvm/ExecutionEngine/Orc/TaskDispatch.h"
 #include "llvm/Support/Compiler.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/MSVCErrorWorkarounds.h"
-
-#include <future>
 #include <mutex>
 #include <vector>
 
 namespace llvm::orc {
 
 class SymbolLookupSet;
 
-class LLVM_ABI DylibManager {
+class LLVM_ABI DylibManager : public ExecutorProcessControl {
 public:
+  DylibManager(std::shared_ptr<SymbolStringPool> SSP,
+               std::unique_ptr<TaskDispatcher> D)
+      : ExecutorProcessControl(std::move(SSP), std::move(D)) {}
+
   /// A pair of a dylib and a set of symbols to be looked up.
   struct LookupRequest {
     LookupRequest(tpctypes::DylibHandle Handle, const SymbolLookupSet &Symbols)
@@ -51,10 +56,11 @@ class LLVM_ABI DylibManager {
   /// symbol is not found then it be assigned a '0' value.
   Expected<std::vector<tpctypes::LookupResult>>
   lookupSymbols(ArrayRef<LookupRequest> Request) {
-    std::promise<MSVCPExpected<std::vector<tpctypes::LookupResult>>> RP;
-    auto RF = RP.get_future();
+    orc::future<MSVCPExpected<std::vector<tpctypes::LookupResult>>> RF;
     lookupSymbolsAsync(Request,
-                       [&RP](auto Result) { RP.set_value(std::move(Result)); });
+                       [RP = RF.get_promise(getDispatcher())](auto Result) {
+                         RP.set_value(std::move(Result));
+                       });
     return RF.get();
   }
 

llvm/include/llvm/ExecutionEngine/Orc/ELFNixPlatform.h

@@ -20,7 +20,7 @@
 #include "llvm/ExecutionEngine/Orc/Shared/ExecutorAddress.h"
 #include "llvm/Support/Compiler.h"
 
-#include <future>
+#include "llvm/ExecutionEngine/Orc/TaskDispatch.h"
 #include <thread>
 #include <unordered_map>
 #include <vector>

llvm/include/llvm/ExecutionEngine/Orc/EPCGenericDylibManager.h

@@ -21,6 +21,7 @@
 #include "llvm/ExecutionEngine/Orc/ExecutorProcessControl.h"
 #include "llvm/ExecutionEngine/Orc/Shared/ExecutorSymbolDef.h"
 #include "llvm/ExecutionEngine/Orc/Shared/SimpleRemoteEPCUtils.h"
+#include "llvm/ExecutionEngine/Orc/TaskDispatch.h"
 #include "llvm/Support/Compiler.h"
 
 namespace llvm {
@@ -50,21 +51,23 @@ class EPCGenericDylibManager {
   /// Loads the dylib with the given name.
   LLVM_ABI Expected<tpctypes::DylibHandle> open(StringRef Path, uint64_t Mode);
 
-  /// Looks up symbols within the given dylib.
+  /// Blocking lookup of symbols within the given dylib.
   Expected<tpctypes::LookupResult> lookup(tpctypes::DylibHandle H,
                                           const SymbolLookupSet &Lookup) {
-    std::promise<MSVCPExpected<tpctypes::LookupResult>> RP;
-    auto RF = RP.get_future();
-    lookupAsync(H, Lookup, [&RP](auto R) { RP.set_value(std::move(R)); });
+    orc::future<Expected<tpctypes::LookupResult>> RF;
+    lookupAsync(H, Lookup, [RP = RF.get_promise(EPC.getDispatcher())](auto R) {
+      RP.set_value(std::move(R));
+    });
     return RF.get();
   }
 
-  /// Looks up symbols within the given dylib.
+  /// Blocking lookup of symbols within the given dylib.
   Expected<tpctypes::LookupResult> lookup(tpctypes::DylibHandle H,
                                           const RemoteSymbolLookupSet &Lookup) {
-    std::promise<MSVCPExpected<tpctypes::LookupResult>> RP;
-    auto RF = RP.get_future();
-    lookupAsync(H, Lookup, [&RP](auto R) { RP.set_value(std::move(R)); });
+    orc::future<Expected<tpctypes::LookupResult>> RF;
+    lookupAsync(H, Lookup, [RP = RF.get_promise(EPC.getDispatcher())](auto R) {
+      RP.set_value(std::move(R));
+    });
     return RF.get();
   }
 

llvm/include/llvm/ExecutionEngine/Orc/EPCGenericMemoryAccess.h

@@ -6,7 +6,7 @@
 //
 //===----------------------------------------------------------------------===//
 //
-// Implements ExecutorProcessControl::MemoryAccess by making calls to
+// Implements MemoryAccess by making calls to
 // ExecutorProcessControl::callWrapperAsync.
 //
 // This simplifies the implementaton of new ExecutorProcessControl instances,
@@ -19,6 +19,7 @@
 #define LLVM_EXECUTIONENGINE_ORC_EPCGENERICMEMORYACCESS_H
 
 #include "llvm/ExecutionEngine/Orc/Core.h"
+#include "llvm/ExecutionEngine/Orc/MemoryAccess.h"
 
 namespace llvm {
 namespace orc {
@@ -45,7 +46,7 @@ class EPCGenericMemoryAccess : public MemoryAccess {
   /// Create an EPCGenericMemoryAccess instance from a given set of
   /// function addrs.
   EPCGenericMemoryAccess(ExecutorProcessControl &EPC, FuncAddrs FAs)
-      : EPC(EPC), FAs(FAs) {}
+      : MemoryAccess(EPC), FAs(FAs) {}
 
   void writeUInt8sAsync(ArrayRef<tpctypes::UInt8Write> Ws,
                         WriteResultFn OnWriteComplete) override {
@@ -202,7 +203,6 @@ class EPCGenericMemoryAccess : public MemoryAccess {
   }
 
 private:
-  ExecutorProcessControl &EPC;
   FuncAddrs FAs;
 };
 

llvm/include/llvm/ExecutionEngine/Orc/ExecutorProcessControl.h

@@ -15,8 +15,6 @@
 
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ExecutionEngine/JITLink/JITLinkMemoryManager.h"
-#include "llvm/ExecutionEngine/Orc/DylibManager.h"
-#include "llvm/ExecutionEngine/Orc/MemoryAccess.h"
 #include "llvm/ExecutionEngine/Orc/Shared/ExecutorAddress.h"
 #include "llvm/ExecutionEngine/Orc/Shared/TargetProcessControlTypes.h"
 #include "llvm/ExecutionEngine/Orc/Shared/WrapperFunctionUtils.h"
@@ -26,25 +24,28 @@
 #include "llvm/Support/Compiler.h"
 #include "llvm/TargetParser/Triple.h"
 
-#include <future>
 #include <mutex>
 #include <vector>
 
-namespace llvm::orc {
+namespace llvm {
+namespace orc {
 
 class ExecutionSession;
+class DylibManager;
+class MemoryAccess;
+class SymbolLookupSet;
 
 /// ExecutorProcessControl supports interaction with a JIT target process.
 class LLVM_ABI ExecutorProcessControl {
   friend class ExecutionSession;
-public:
 
+public:
   /// A handler or incoming WrapperFunctionResults -- either return values from
   /// callWrapper* calls, or incoming JIT-dispatch requests.
   ///
   /// IncomingWFRHandlers are constructible from
   /// unique_function<void(shared::WrapperFunctionResult)>s using the
-  /// runInPlace function or a RunWithDispatch object.
+  /// RunInPlace function or a RunAsTask object.
   class IncomingWFRHandler {
     friend class ExecutorProcessControl;
   public:
@@ -83,15 +84,16 @@ class LLVM_ABI ExecutorProcessControl {
 
     template <typename FnT>
     IncomingWFRHandler operator()(FnT &&Fn) {
-      return IncomingWFRHandler(
-          [&D = this->D, Fn = std::move(Fn)]
-          (shared::WrapperFunctionResult WFR) mutable {
-              D.dispatch(
-                makeGenericNamedTask(
-                    [Fn = std::move(Fn), WFR = std::move(WFR)]() mutable {
-                      Fn(std::move(WFR));
-                    }, "WFR handler task"));
+      orc::future<shared::WrapperFunctionResult> F;
+      auto H = IncomingWFRHandler(
+          [P = F.get_promise(D)](shared::WrapperFunctionResult WFR) {
+            P.set_value(std::move(WFR));
           });
+      std::move(F).then(
+          [Fn = std::move(Fn)](shared::WrapperFunctionResult &&WFR) mutable {
+            Fn(std::move(WFR));
+          });
+      return H;
     }
   private:
     TaskDispatcher &D;
@@ -251,13 +253,13 @@ class LLVM_ABI ExecutorProcessControl {
   /// \endcode{.cpp}
   shared::WrapperFunctionResult callWrapper(ExecutorAddr WrapperFnAddr,
                                             ArrayRef<char> ArgBuffer) {
-    std::promise<shared::WrapperFunctionResult> RP;
-    auto RF = RP.get_future();
+    orc::future<shared::WrapperFunctionResult> RF;
     callWrapperAsync(
         RunInPlace(), WrapperFnAddr,
-        [&](shared::WrapperFunctionResult R) {
+        [RP = RF.get_promise(*D)](shared::WrapperFunctionResult R) {
           RP.set_value(std::move(R));
-        }, ArgBuffer);
+        },
+        ArgBuffer);
     return RF.get();
   }
 
@@ -322,6 +324,7 @@ class LLVM_ABI ExecutorProcessControl {
   StringMap<ExecutorAddr> BootstrapSymbols;
 };
 
-} // namespace llvm::orc
+} // end namespace orc
+} // end namespace llvm
 
 #endif // LLVM_EXECUTIONENGINE_ORC_EXECUTORPROCESSCONTROL_H