[ORC] Make runAllocActions and runDeallocActions asynchorous.

Part of ongoing work to make core ORC operations asynchronous.

This only affects the runner utilities, not the AllocationActions themselves.
The AllocationActions will be made asynchronous in a future patch.

Author: lhames

llvm/include/llvm/ExecutionEngine/Orc/Shared/AllocationActions.h

@@ -13,6 +13,7 @@
 #ifndef LLVM_EXECUTIONENGINE_ORC_SHARED_ALLOCATIONACTIONS_H
 #define LLVM_EXECUTIONENGINE_ORC_SHARED_ALLOCATIONACTIONS_H
 
+#include "llvm/ADT/FunctionExtras.h"
 #include "llvm/ExecutionEngine/Orc/Shared/ExecutorAddress.h"
 #include "llvm/ExecutionEngine/Orc/Shared/WrapperFunctionUtils.h"
 #include "llvm/Support/Memory.h"
@@ -52,6 +53,9 @@ inline size_t numDeallocActions(const AllocActions &AAs) {
       AAs, [](const AllocActionCallPair &P) { return !!P.Dealloc; });
 }
 
+using OnRunFinalizeActionsCompleteFn =
+    unique_function<void(Expected<std::vector<WrapperFunctionCall>>)>;
+
 /// Run finalize actions.
 ///
 /// If any finalize action fails then the corresponding dealloc actions will be
@@ -62,13 +66,16 @@ inline size_t numDeallocActions(const AllocActions &AAs) {
 /// be returned. The dealloc actions should be run by calling
 /// runDeallocationActions. If this function succeeds then the AA argument will
 /// be cleared before the function returns.
-Expected<std::vector<WrapperFunctionCall>>
-runFinalizeActions(AllocActions &AAs);
+void runFinalizeActions(AllocActions &AAs,
+                        OnRunFinalizeActionsCompleteFn OnComplete);
+
+using OnRunDeallocActionsComeleteFn = unique_function<void(Error)>;
 
 /// Run deallocation actions.
 /// Dealloc actions will be run in reverse order (from last element of DAs to
 /// first).
-Error runDeallocActions(ArrayRef<WrapperFunctionCall> DAs);
+void runDeallocActions(ArrayRef<WrapperFunctionCall> DAs,
+                       OnRunDeallocActionsComeleteFn OnComplete);
 
 using SPSAllocActionCallPair =
     SPSTuple<SPSWrapperFunctionCall, SPSWrapperFunctionCall>;

llvm/lib/ExecutionEngine/JITLink/JITLinkMemoryManager.cpp

@@ -260,17 +260,22 @@ class InProcessMemoryManager::IPInFlightAlloc
     }
 
     // Run finalization actions.
-    auto DeallocActions = runFinalizeActions(G->allocActions());
-    if (!DeallocActions) {
-      OnFinalized(DeallocActions.takeError());
-      return;
-    }
+    using WrapperFunctionCall = orc::shared::WrapperFunctionCall;
+    runFinalizeActions(
+        G->allocActions(),
+        [this, OnFinalized = std::move(OnFinalized)](
+            Expected<std::vector<WrapperFunctionCall>> DeallocActions) mutable {
+          completeFinalization(std::move(OnFinalized),
+                               std::move(DeallocActions));
+        });
+  }
 
-    // Release the finalize segments slab.
-    if (auto EC = sys::Memory::releaseMappedMemory(FinalizationSegments)) {
-      OnFinalized(errorCodeToError(EC));
-      return;
-    }
+  void abandon(OnAbandonedFunction OnAbandoned) override {
+    Error Err = Error::success();
+    if (auto EC = sys::Memory::releaseMappedMemory(FinalizationSegments))
+      Err = joinErrors(std::move(Err), errorCodeToError(EC));
+    if (auto EC = sys::Memory::releaseMappedMemory(StandardSegments))
+      Err = joinErrors(std::move(Err), errorCodeToError(EC));
 
 #ifndef NDEBUG
     // Set 'G' to null to flag that we've been successfully finalized.
@@ -279,17 +284,22 @@ class InProcessMemoryManager::IPInFlightAlloc
     G = nullptr;
 #endif
 
-    // Continue with finalized allocation.
-    OnFinalized(MemMgr.createFinalizedAlloc(std::move(StandardSegments),
-                                            std::move(*DeallocActions)));
+    OnAbandoned(std::move(Err));
   }
 
-  void abandon(OnAbandonedFunction OnAbandoned) override {
-    Error Err = Error::success();
-    if (auto EC = sys::Memory::releaseMappedMemory(FinalizationSegments))
-      Err = joinErrors(std::move(Err), errorCodeToError(EC));
-    if (auto EC = sys::Memory::releaseMappedMemory(StandardSegments))
-      Err = joinErrors(std::move(Err), errorCodeToError(EC));
+private:
+  void completeFinalization(
+      OnFinalizedFunction OnFinalized,
+      Expected<std::vector<orc::shared::WrapperFunctionCall>> DeallocActions) {
+
+    if (!DeallocActions)
+      return OnFinalized(DeallocActions.takeError());
+
+    // Release the finalize segments slab.
+    if (auto EC = sys::Memory::releaseMappedMemory(FinalizationSegments)) {
+      OnFinalized(errorCodeToError(EC));
+      return;
+    }
 
 #ifndef NDEBUG
     // Set 'G' to null to flag that we've been successfully finalized.
@@ -298,10 +308,11 @@ class InProcessMemoryManager::IPInFlightAlloc
     G = nullptr;
 #endif
 
-    OnAbandoned(std::move(Err));
+    // Continue with finalized allocation.
+    OnFinalized(MemMgr.createFinalizedAlloc(std::move(StandardSegments),
+                                            std::move(*DeallocActions)));
   }
 
-private:
   Error applyProtections() {
     for (auto &KV : BL.segments()) {
       const auto &AG = KV.first;

llvm/lib/ExecutionEngine/Orc/MemoryMapper.cpp

@@ -90,17 +90,27 @@ void InProcessMemoryMapper::initialize(MemoryMapper::AllocInfo &AI,
       sys::Memory::InvalidateInstructionCache(Base.toPtr<void *>(), Size);
   }
 
-  auto DeinitializeActions = shared::runFinalizeActions(AI.Actions);
-  if (!DeinitializeActions)
-    return OnInitialized(DeinitializeActions.takeError());
+  std::vector<shared::WrapperFunctionCall> DeinitializeActions;
+  {
+    std::promise<MSVCPExpected<std::vector<shared::WrapperFunctionCall>>> P;
+    auto F = P.get_future();
+    shared::runFinalizeActions(
+        AI.Actions, [&](Expected<std::vector<shared::WrapperFunctionCall>> R) {
+          P.set_value(std::move(R));
+        });
+    if (auto DeinitializeActionsOrErr = F.get())
+      DeinitializeActions = std::move(*DeinitializeActionsOrErr);
+    else
+      return OnInitialized(DeinitializeActionsOrErr.takeError());
+  }
 
   {
     std::lock_guard<std::mutex> Lock(Mutex);
 
     // This is the maximum range whose permission have been possibly modified
     auto &Alloc = Allocations[MinAddr];
     Alloc.Size = MaxAddr - MinAddr;
-    Alloc.DeinitializationActions = std::move(*DeinitializeActions);
+    Alloc.DeinitializationActions = std::move(DeinitializeActions);
     Reservations[AI.MappingBase.toPtr<void *>()].Allocations.push_back(MinAddr);
   }
 
@@ -117,10 +127,10 @@ void InProcessMemoryMapper::deinitialize(
 
     for (auto Base : llvm::reverse(Bases)) {
 
-      if (Error Err = shared::runDeallocActions(
-              Allocations[Base].DeinitializationActions)) {
-        AllErr = joinErrors(std::move(AllErr), std::move(Err));
-      }
+      shared::runDeallocActions(
+          Allocations[Base].DeinitializationActions, [&](Error Err) {
+            AllErr = joinErrors(std::move(AllErr), std::move(Err));
+          });
 
       // Reset protections to read/write so the area can be reused
       if (auto EC = sys::Memory::protectMappedMemory(

llvm/lib/ExecutionEngine/Orc/Shared/AllocationActions.cpp

@@ -12,31 +12,39 @@ namespace llvm {
 namespace orc {
 namespace shared {
 
-Expected<std::vector<WrapperFunctionCall>>
-runFinalizeActions(AllocActions &AAs) {
+void runFinalizeActions(AllocActions &AAs,
+                        OnRunFinalizeActionsCompleteFn OnComplete) {
   std::vector<WrapperFunctionCall> DeallocActions;
   DeallocActions.reserve(numDeallocActions(AAs));
 
   for (auto &AA : AAs) {
     if (AA.Finalize)
-      if (auto Err = AA.Finalize.runWithSPSRetErrorMerged())
-        return joinErrors(std::move(Err), runDeallocActions(DeallocActions));
+
+      if (auto Err = AA.Finalize.runWithSPSRetErrorMerged()) {
+        while (!DeallocActions.empty()) {
+          Err = joinErrors(std::move(Err),
+                           DeallocActions.back().runWithSPSRetErrorMerged());
+          DeallocActions.pop_back();
+        }
+        return OnComplete(std::move(Err));
+      }
 
     if (AA.Dealloc)
       DeallocActions.push_back(std::move(AA.Dealloc));
   }
 
   AAs.clear();
-  return DeallocActions;
+  OnComplete(std::move(DeallocActions));
 }
 
-Error runDeallocActions(ArrayRef<WrapperFunctionCall> DAs) {
+void runDeallocActions(ArrayRef<WrapperFunctionCall> DAs,
+                       OnRunDeallocActionsComeleteFn OnComplete) {
   Error Err = Error::success();
   while (!DAs.empty()) {
     Err = joinErrors(std::move(Err), DAs.back().runWithSPSRetErrorMerged());
     DAs = DAs.drop_back();
   }
-  return Err;
+  OnComplete(std::move(Err));
 }
 
 } // namespace shared

llvm/lib/ExecutionEngine/Orc/TargetProcess/ExecutorSharedMemoryMapperService.cpp

@@ -9,8 +9,10 @@
 #include "llvm/ExecutionEngine/Orc/TargetProcess/ExecutorSharedMemoryMapperService.h"
 #include "llvm/Config/llvm-config.h" // for LLVM_ON_UNIX
 #include "llvm/ExecutionEngine/Orc/Shared/OrcRTBridge.h"
+#include "llvm/Support/MSVCErrorWorkarounds.h"
 #include "llvm/Support/Process.h"
 #include "llvm/Support/WindowsError.h"
+#include <future>
 #include <sstream>
 
 #if defined(LLVM_ON_UNIX)
@@ -181,15 +183,24 @@ Expected<ExecutorAddr> ExecutorSharedMemoryMapperService::initialize(
   }
 
   // Run finalization actions and get deinitlization action list.
-  auto DeinitializeActions = shared::runFinalizeActions(FR.Actions);
-  if (!DeinitializeActions) {
-    return DeinitializeActions.takeError();
+  std::vector<shared::WrapperFunctionCall> DeinitializeActions;
+  {
+    std::promise<MSVCPExpected<std::vector<shared::WrapperFunctionCall>>> P;
+    auto F = P.get_future();
+    shared::runFinalizeActions(
+        FR.Actions, [&](Expected<std::vector<shared::WrapperFunctionCall>> R) {
+          P.set_value(std::move(R));
+        });
+    if (auto DeinitializeActionsOrErr = F.get())
+      DeinitializeActions = std::move(*DeinitializeActionsOrErr);
+    else
+      return DeinitializeActionsOrErr.takeError();
   }
 
   {
     std::lock_guard<std::mutex> Lock(Mutex);
     Allocations[MinAddr].DeinitializationActions =
-        std::move(*DeinitializeActions);
+        std::move(DeinitializeActions);
     Reservations[Reservation.toPtr<void *>()].Allocations.push_back(MinAddr);
   }
 
@@ -210,10 +221,11 @@ Error ExecutorSharedMemoryMapperService::deinitialize(
     std::lock_guard<std::mutex> Lock(Mutex);
 
     for (auto Base : llvm::reverse(Bases)) {
-      if (Error Err = shared::runDeallocActions(
-              Allocations[Base].DeinitializationActions)) {
-        AllErr = joinErrors(std::move(AllErr), std::move(Err));
-      }
+      shared::runDeallocActions(
+          Allocations[Base].DeinitializationActions, [&](Error Err) {
+            if (Err)
+              AllErr = joinErrors(std::move(AllErr), std::move(Err));
+          });
 
       // Remove the allocation from the allocation list of its reservation
       for (auto &Reservation : Reservations) {