[ORC] Add finalization & deallocation actions, SimpleExecutorMemoryManager class
authorLang Hames <lhames@gmail.com>
Tue, 14 Sep 2021 01:47:08 +0000 (11:47 +1000)
committerLang Hames <lhames@gmail.com>
Thu, 16 Sep 2021 23:55:45 +0000 (09:55 +1000)
Finalization and deallocation actions are a key part of the upcoming
JITLinkMemoryManager redesign: They generalize the existing finalization and
deallocate concepts (basically "copy-and-mprotect", and "munmap") to include
support for arbitrary registration and deregistration of parts of JIT linked
code. This allows us to register and deregister eh-frames, TLV sections,
language metadata, etc. using regular memory management calls with no additional
IPC/RPC overhead, which should both improve JIT performance and simplify
interactions between ORC and the ORC runtime.

The SimpleExecutorMemoryManager class provides executor-side support for memory
management operations, including finalization and deallocation actions.

This support is being added in advance of the rest of the memory manager
redesign as it will simplify the introduction of an EPC based
RuntimeDyld::MemoryManager (since eh-frame registration/deregistration will be
expressible as actions). The new RuntimeDyld::MemoryManager will in turn allow
us to remove older remote allocators that are blocking the rest of the memory
manager changes.

17 files changed:
llvm/include/llvm/ExecutionEngine/Orc/EPCGenericJITLinkMemoryManager.h
llvm/include/llvm/ExecutionEngine/Orc/Shared/OrcRTBridge.h
llvm/include/llvm/ExecutionEngine/Orc/Shared/TargetProcessControlTypes.h
llvm/include/llvm/ExecutionEngine/Orc/TargetProcess/ExecutorBootstrapService.h [new file with mode: 0644]
llvm/include/llvm/ExecutionEngine/Orc/TargetProcess/SimpleExecutorMemoryManager.h [new file with mode: 0644]
llvm/include/llvm/ExecutionEngine/Orc/TargetProcess/SimpleRemoteEPCServer.h
llvm/lib/ExecutionEngine/Orc/EPCGenericJITLinkMemoryManager.cpp
llvm/lib/ExecutionEngine/Orc/Shared/OrcRTBridge.cpp
llvm/lib/ExecutionEngine/Orc/SimpleRemoteEPC.cpp
llvm/lib/ExecutionEngine/Orc/TargetProcess/CMakeLists.txt
llvm/lib/ExecutionEngine/Orc/TargetProcess/OrcRTBootstrap.cpp
llvm/lib/ExecutionEngine/Orc/TargetProcess/SimpleExecutorMemoryManager.cpp [new file with mode: 0644]
llvm/lib/ExecutionEngine/Orc/TargetProcess/SimpleRemoteEPCServer.cpp
llvm/tools/llvm-jitlink/llvm-jitlink-executor/llvm-jitlink-executor.cpp
llvm/unittests/ExecutionEngine/Orc/CMakeLists.txt
llvm/unittests/ExecutionEngine/Orc/EPCGenericJITLinkMemoryManagerTest.cpp
llvm/unittests/ExecutionEngine/Orc/SimpleExecutorMemoryManagerTest.cpp [new file with mode: 0644]

index 4800712..96b9522 100644 (file)
@@ -27,7 +27,8 @@ namespace orc {
 class EPCGenericJITLinkMemoryManager : public jitlink::JITLinkMemoryManager {
 public:
   /// Function addresses for memory access.
-  struct FuncAddrs {
+  struct SymbolAddrs {
+    ExecutorAddress Allocator;
     ExecutorAddress Reserve;
     ExecutorAddress Finalize;
     ExecutorAddress Deallocate;
@@ -35,8 +36,8 @@ public:
 
   /// Create an EPCGenericJITLinkMemoryManager instance from a given set of
   /// function addrs.
-  EPCGenericJITLinkMemoryManager(ExecutorProcessControl &EPC, FuncAddrs FAs)
-      : EPC(EPC), FAs(FAs) {}
+  EPCGenericJITLinkMemoryManager(ExecutorProcessControl &EPC, SymbolAddrs SAs)
+      : EPC(EPC), SAs(SAs) {}
 
   Expected<std::unique_ptr<Allocation>>
   allocate(const jitlink::JITLinkDylib *JD,
@@ -46,7 +47,7 @@ private:
   class Alloc;
 
   ExecutorProcessControl &EPC;
-  FuncAddrs FAs;
+  SymbolAddrs SAs;
 };
 
 } // end namespace orc
index 688439f..54db1f7 100644 (file)
@@ -21,9 +21,11 @@ namespace llvm {
 namespace orc {
 namespace rt {
 
-extern const char *MemoryReserveWrapperName;
-extern const char *MemoryFinalizeWrapperName;
-extern const char *MemoryDeallocateWrapperName;
+extern const char *SimpleExecutorMemoryManagerInstanceName;
+extern const char *SimpleExecutorMemoryManagerReserveWrapperName;
+extern const char *SimpleExecutorMemoryManagerFinalizeWrapperName;
+extern const char *SimpleExecutorMemoryManagerDeallocateWrapperName;
+
 extern const char *MemoryWriteUInt8sWrapperName;
 extern const char *MemoryWriteUInt16sWrapperName;
 extern const char *MemoryWriteUInt32sWrapperName;
@@ -31,11 +33,15 @@ extern const char *MemoryWriteUInt64sWrapperName;
 extern const char *MemoryWriteBuffersWrapperName;
 extern const char *RunAsMainWrapperName;
 
-using SPSMemoryReserveSignature =
-    shared::SPSExpected<shared::SPSExecutorAddress>(uint64_t);
-using SPSMemoryFinalizeSignature = shared::SPSError(shared::SPSFinalizeRequest);
-using SPSMemoryDeallocateSignature =
-    shared::SPSError(shared::SPSExecutorAddress, uint64_t);
+using SPSSimpleExecutorMemoryManagerReserveSignature =
+    shared::SPSExpected<shared::SPSExecutorAddress>(shared::SPSExecutorAddress,
+                                                    uint64_t);
+using SPSSimpleExecutorMemoryManagerFinalizeSignature =
+    shared::SPSError(shared::SPSExecutorAddress, shared::SPSFinalizeRequest);
+using SPSSimpleExecutorMemoryManagerDeallocateSignature =
+    shared::SPSError(shared::SPSExecutorAddress,
+                     shared::SPSSequence<shared::SPSExecutorAddress>);
+
 using SPSRunAsMainSignature = int64_t(shared::SPSExecutorAddress,
                                       shared::SPSSequence<shared::SPSString>);
 
index f4fd619..84b1558 100644 (file)
@@ -19,6 +19,7 @@
 #include "llvm/ExecutionEngine/JITSymbol.h"
 #include "llvm/ExecutionEngine/Orc/Shared/ExecutorAddress.h"
 #include "llvm/ExecutionEngine/Orc/Shared/SimplePackedSerialization.h"
+#include "llvm/ExecutionEngine/Orc/Shared/WrapperFunctionUtils.h"
 #include "llvm/Support/Memory.h"
 
 #include <vector>
@@ -68,6 +69,31 @@ inline std::string getWireProtectionFlagsStr(WireProtectionFlags WPF) {
   return Result;
 }
 
+struct SupportFunctionCall {
+  using FnTy = shared::detail::CWrapperFunctionResult(const char *ArgData,
+                                                      size_t ArgSize);
+  ExecutorAddress Func;
+  ExecutorAddress ArgData;
+  uint64_t ArgSize;
+
+  Error run() {
+    shared::WrapperFunctionResult WFR(
+        Func.toPtr<FnTy *>()(ArgData.toPtr<const char *>(), ArgSize));
+    if (const char *ErrMsg = WFR.getOutOfBandError())
+      return make_error<StringError>(ErrMsg, inconvertibleErrorCode());
+    if (!WFR.empty())
+      return make_error<StringError>("Unexpected result bytes from "
+                                     "support function call",
+                                     inconvertibleErrorCode());
+    return Error::success();
+  }
+};
+
+struct AllocationActionsPair {
+  SupportFunctionCall Finalize;
+  SupportFunctionCall Deallocate;
+};
+
 struct SegFinalizeRequest {
   WireProtectionFlags Prot;
   ExecutorAddress Addr;
@@ -75,7 +101,10 @@ struct SegFinalizeRequest {
   ArrayRef<char> Content;
 };
 
-using FinalizeRequest = std::vector<SegFinalizeRequest>;
+struct FinalizeRequest {
+  std::vector<SegFinalizeRequest> Segments;
+  std::vector<AllocationActionsPair> Actions;
+};
 
 template <typename T> struct UIntWrite {
   UIntWrite() = default;
@@ -120,11 +149,18 @@ namespace shared {
 
 class SPSMemoryProtectionFlags {};
 
+using SPSSupportFunctionCall =
+    SPSTuple<SPSExecutorAddress, SPSExecutorAddress, uint64_t>;
+
 using SPSSegFinalizeRequest =
     SPSTuple<SPSMemoryProtectionFlags, SPSExecutorAddress, uint64_t,
              SPSSequence<char>>;
 
-using SPSFinalizeRequest = SPSSequence<SPSSegFinalizeRequest>;
+using SPSAllocationActionsPair =
+    SPSTuple<SPSSupportFunctionCall, SPSSupportFunctionCall>;
+
+using SPSFinalizeRequest = SPSTuple<SPSSequence<SPSSegFinalizeRequest>,
+                                    SPSSequence<SPSAllocationActionsPair>>;
 
 template <typename T>
 using SPSMemoryAccessUIntWrite = SPSTuple<SPSExecutorAddress, T>;
@@ -161,6 +197,48 @@ public:
 };
 
 template <>
+class SPSSerializationTraits<SPSSupportFunctionCall,
+                             tpctypes::SupportFunctionCall> {
+  using AL = SPSSupportFunctionCall::AsArgList;
+
+public:
+  static size_t size(const tpctypes::SupportFunctionCall &SFC) {
+    return AL::size(SFC.Func, SFC.ArgData, SFC.ArgSize);
+  }
+
+  static bool serialize(SPSOutputBuffer &OB,
+                        const tpctypes::SupportFunctionCall &SFC) {
+    return AL::serialize(OB, SFC.Func, SFC.ArgData, SFC.ArgSize);
+  }
+
+  static bool deserialize(SPSInputBuffer &IB,
+                          tpctypes::SupportFunctionCall &SFC) {
+    return AL::deserialize(IB, SFC.Func, SFC.ArgData, SFC.ArgSize);
+  }
+};
+
+template <>
+class SPSSerializationTraits<SPSAllocationActionsPair,
+                             tpctypes::AllocationActionsPair> {
+  using AL = SPSAllocationActionsPair::AsArgList;
+
+public:
+  static size_t size(const tpctypes::AllocationActionsPair &AAP) {
+    return AL::size(AAP.Finalize, AAP.Deallocate);
+  }
+
+  static bool serialize(SPSOutputBuffer &OB,
+                        const tpctypes::AllocationActionsPair &AAP) {
+    return AL::serialize(OB, AAP.Finalize, AAP.Deallocate);
+  }
+
+  static bool deserialize(SPSInputBuffer &IB,
+                          tpctypes::AllocationActionsPair &AAP) {
+    return AL::deserialize(IB, AAP.Finalize, AAP.Deallocate);
+  }
+};
+
+template <>
 class SPSSerializationTraits<SPSSegFinalizeRequest,
                              tpctypes::SegFinalizeRequest> {
   using SFRAL = SPSSegFinalizeRequest::AsArgList;
@@ -181,6 +259,25 @@ public:
   }
 };
 
+template <>
+class SPSSerializationTraits<SPSFinalizeRequest, tpctypes::FinalizeRequest> {
+  using FRAL = SPSFinalizeRequest::AsArgList;
+
+public:
+  static size_t size(const tpctypes::FinalizeRequest &FR) {
+    return FRAL::size(FR.Segments, FR.Actions);
+  }
+
+  static bool serialize(SPSOutputBuffer &OB,
+                        const tpctypes::FinalizeRequest &FR) {
+    return FRAL::serialize(OB, FR.Segments, FR.Actions);
+  }
+
+  static bool deserialize(SPSInputBuffer &IB, tpctypes::FinalizeRequest &FR) {
+    return FRAL::deserialize(IB, FR.Segments, FR.Actions);
+  }
+};
+
 template <typename T>
 class SPSSerializationTraits<SPSMemoryAccessUIntWrite<T>,
                              tpctypes::UIntWrite<T>> {
diff --git a/llvm/include/llvm/ExecutionEngine/Orc/TargetProcess/ExecutorBootstrapService.h b/llvm/include/llvm/ExecutionEngine/Orc/TargetProcess/ExecutorBootstrapService.h
new file mode 100644 (file)
index 0000000..206d46c
--- /dev/null
@@ -0,0 +1,36 @@
+//===- ExecutorService.h - Provide bootstrap symbols to session -*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Provides a service by supplying some set of bootstrap symbols.
+//
+// FIXME: The functionality in this file should be moved to the ORC runtime.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_EXECUTIONENGINE_ORC_TARGETPROCESS_EXECUTORBOOTSTRAPSERVICE_H
+#define LLVM_EXECUTIONENGINE_ORC_TARGETPROCESS_EXECUTORBOOTSTRAPSERVICE_H
+
+#include "llvm/ADT/StringMap.h"
+#include "llvm/ExecutionEngine/Orc/Shared/ExecutorAddress.h"
+
+namespace llvm {
+namespace orc {
+
+class ExecutorBootstrapService {
+public:
+  virtual ~ExecutorBootstrapService();
+
+  virtual void
+  addBootstrapSymbols(StringMap<ExecutorAddress> &BootstrapSymbols) = 0;
+  virtual Error shutdown() = 0;
+};
+
+} // end namespace orc
+} // end namespace llvm
+
+#endif // LLVM_EXECUTIONENGINE_ORC_TARGETPROCESS_EXECUTORBOOTSTRAPSERVICE_H
diff --git a/llvm/include/llvm/ExecutionEngine/Orc/TargetProcess/SimpleExecutorMemoryManager.h b/llvm/include/llvm/ExecutionEngine/Orc/TargetProcess/SimpleExecutorMemoryManager.h
new file mode 100644 (file)
index 0000000..b2a3ddf
--- /dev/null
@@ -0,0 +1,70 @@
+//===---------------- SimpleExecutorMemoryManager.h -------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// A simple allocator class suitable for basic remote-JIT use.
+//
+// FIXME: The functionality in this file should be moved to the ORC runtime.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_EXECUTIONENGINE_ORC_TARGETPROCESS_SIMPLEEXECUTORMEMORYMANAGER_H
+#define LLVM_EXECUTIONENGINE_ORC_TARGETPROCESS_SIMPLEEXECUTORMEMORYMANAGER_H
+
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ExecutionEngine/Orc/Shared/ExecutorAddress.h"
+#include "llvm/ExecutionEngine/Orc/Shared/TargetProcessControlTypes.h"
+#include "llvm/ExecutionEngine/Orc/Shared/WrapperFunctionUtils.h"
+#include "llvm/ExecutionEngine/Orc/TargetProcess/ExecutorBootstrapService.h"
+#include "llvm/Support/Error.h"
+
+#include <mutex>
+
+namespace llvm {
+namespace orc {
+namespace rt_bootstrap {
+
+/// Simple page-based allocator.
+class SimpleExecutorMemoryManager : public ExecutorBootstrapService {
+public:
+  virtual ~SimpleExecutorMemoryManager();
+
+  Expected<ExecutorAddress> allocate(uint64_t Size);
+  Error finalize(tpctypes::FinalizeRequest &FR);
+  Error deallocate(const std::vector<ExecutorAddress> &Bases);
+
+  Error shutdown() override;
+  void addBootstrapSymbols(StringMap<ExecutorAddress> &M) override;
+
+private:
+  struct Allocation {
+    size_t Size = 0;
+    std::vector<tpctypes::SupportFunctionCall> DeallocationActions;
+  };
+
+  using AllocationsMap = DenseMap<void *, Allocation>;
+
+  Error deallocateImpl(void *Base, Allocation &A);
+
+  static llvm::orc::shared::detail::CWrapperFunctionResult
+  reserveWrapper(const char *ArgData, size_t ArgSize);
+
+  static llvm::orc::shared::detail::CWrapperFunctionResult
+  finalizeWrapper(const char *ArgData, size_t ArgSize);
+
+  static llvm::orc::shared::detail::CWrapperFunctionResult
+  deallocateWrapper(const char *ArgData, size_t ArgSize);
+
+  std::mutex M;
+  AllocationsMap Allocations;
+};
+
+} // end namespace rt_bootstrap
+} // end namespace orc
+} // end namespace llvm
+
+#endif // LLVM_EXECUTIONENGINE_ORC_TARGETPROCESS_SIMPLEEXECUTORMEMORYMANAGER_H
index 121d8d4..083ed64 100644 (file)
@@ -19,6 +19,7 @@
 #include "llvm/ExecutionEngine/Orc/Shared/SimpleRemoteEPCUtils.h"
 #include "llvm/ExecutionEngine/Orc/Shared/TargetProcessControlTypes.h"
 #include "llvm/ExecutionEngine/Orc/Shared/WrapperFunctionUtils.h"
+#include "llvm/ExecutionEngine/Orc/TargetProcess/ExecutorBootstrapService.h"
 #include "llvm/Support/DynamicLibrary.h"
 #include "llvm/Support/Error.h"
 
@@ -35,6 +36,7 @@ class SimpleRemoteEPCServer : public SimpleRemoteEPCTransportClient {
 public:
   using ReportErrorFunction = unique_function<void(Error)>;
 
+  /// Dispatches calls to runWrapper.
   class Dispatcher {
   public:
     virtual ~Dispatcher();
@@ -56,26 +58,60 @@ public:
   };
 #endif
 
+  class Setup {
+    friend class SimpleRemoteEPCServer;
+
+  public:
+    SimpleRemoteEPCServer &server();
+    StringMap<ExecutorAddress> &bootstrapSymbols() { return BootstrapSymbols; }
+    std::vector<std::unique_ptr<ExecutorBootstrapService>> &services() {
+      return Services;
+    }
+    void setDispatcher(std::unique_ptr<Dispatcher> D) { S.D = std::move(D); }
+    void setErrorReporter(unique_function<void(Error)> ReportError) {
+      S.ReportError = std::move(ReportError);
+    }
+
+  private:
+    Setup(SimpleRemoteEPCServer &S) : S(S) {}
+    SimpleRemoteEPCServer &S;
+    StringMap<ExecutorAddress> BootstrapSymbols;
+    std::vector<std::unique_ptr<ExecutorBootstrapService>> Services;
+  };
+
   static StringMap<ExecutorAddress> defaultBootstrapSymbols();
 
   template <typename TransportT, typename... TransportTCtorArgTs>
   static Expected<std::unique_ptr<SimpleRemoteEPCServer>>
-  Create(std::unique_ptr<Dispatcher> D,
-         StringMap<ExecutorAddress> BootstrapSymbols,
+  Create(unique_function<Error(Setup &S)> SetupFunction,
          TransportTCtorArgTs &&...TransportTCtorArgs) {
-    auto SREPCServer = std::make_unique<SimpleRemoteEPCServer>();
-    SREPCServer->D = std::move(D);
-    SREPCServer->ReportError = [](Error Err) {
-      logAllUnhandledErrors(std::move(Err), errs(), "SimpleRemoteEPCServer ");
-    };
+    auto Server = std::make_unique<SimpleRemoteEPCServer>();
+    Setup S(*Server);
+    if (auto Err = SetupFunction(S))
+      return std::move(Err);
+
+    // Set ReportError up-front so that it can be used if construction
+    // process fails.
+    if (!Server->ReportError)
+      Server->ReportError = [](Error Err) {
+        logAllUnhandledErrors(std::move(Err), errs(), "SimpleRemoteEPCServer ");
+      };
+
+    // Attempt to create transport.
     auto T = TransportT::Create(
-        *SREPCServer, std::forward<TransportTCtorArgTs>(TransportTCtorArgs)...);
+        *Server, std::forward<TransportTCtorArgTs>(TransportTCtorArgs)...);
     if (!T)
       return T.takeError();
-    SREPCServer->T = std::move(*T);
-    if (auto Err = SREPCServer->sendSetupMessage(std::move(BootstrapSymbols)))
+    Server->T = std::move(*T);
+
+    // If transport creation succeeds then start up services.
+    Server->Services = std::move(S.services());
+    for (auto &Service : Server->Services)
+      Service->addBootstrapSymbols(S.bootstrapSymbols());
+
+    if (auto Err = Server->sendSetupMessage(std::move(S.BootstrapSymbols)))
       return std::move(Err);
-    return std::move(SREPCServer);
+    return std::move(Server);
   }
 
   /// Set an error reporter for this server.
@@ -136,6 +172,7 @@ private:
   Error ShutdownErr = Error::success();
   std::unique_ptr<SimpleRemoteEPCTransport> T;
   std::unique_ptr<Dispatcher> D;
+  std::vector<std::unique_ptr<ExecutorBootstrapService>> Services;
   ReportErrorFunction ReportError;
 
   uint64_t NextSeqNo = 0;
index 0cdb0aa..5308f8c 100644 (file)
@@ -27,9 +27,8 @@ public:
   using SegInfoMap = DenseMap<unsigned, SegInfo>;
 
   Alloc(EPCGenericJITLinkMemoryManager &Parent, ExecutorAddress TargetAddr,
-        uint64_t TargetSize, std::unique_ptr<char[]> WorkingBuffer,
-        SegInfoMap Segs)
-      : Parent(Parent), TargetAddr(TargetAddr), TargetSize(TargetSize),
+        std::unique_ptr<char[]> WorkingBuffer, SegInfoMap Segs)
+      : Parent(Parent), TargetAddr(TargetAddr),
         WorkingBuffer(std::move(WorkingBuffer)), Segs(std::move(Segs)) {}
 
   MutableArrayRef<char> getWorkingMemory(ProtectionFlags Seg) override {
@@ -50,7 +49,7 @@ public:
     tpctypes::FinalizeRequest FR;
     for (auto &KV : Segs) {
       assert(KV.second.ContentSize <= std::numeric_limits<size_t>::max());
-      FR.push_back(tpctypes::SegFinalizeRequest{
+      FR.Segments.push_back(tpctypes::SegFinalizeRequest{
           tpctypes::toWireProtectionFlags(
               static_cast<sys::Memory::ProtectionFlags>(KV.first)),
           KV.second.TargetAddr,
@@ -59,7 +58,8 @@ public:
           {WorkingMem, static_cast<size_t>(KV.second.ContentSize)}});
       WorkingMem += KV.second.ContentSize;
     }
-    Parent.EPC.callSPSWrapperAsync<rt::SPSMemoryFinalizeSignature>(
+    Parent.EPC.callSPSWrapperAsync<
+        rt::SPSSimpleExecutorMemoryManagerFinalizeSignature>(
         [OnFinalize = std::move(OnFinalize)](Error SerializationErr,
                                              Error FinalizeErr) {
           if (SerializationErr)
@@ -67,13 +67,15 @@ public:
           else
             OnFinalize(std::move(FinalizeErr));
         },
-        Parent.FAs.Finalize.getValue(), std::move(FR));
+        Parent.SAs.Finalize.getValue(), Parent.SAs.Allocator, std::move(FR));
   }
 
   Error deallocate() override {
     Error Err = Error::success();
-    if (auto E2 = Parent.EPC.callSPSWrapper<rt::SPSMemoryDeallocateSignature>(
-            Parent.FAs.Deallocate.getValue(), Err, TargetAddr, TargetSize))
+    if (auto E2 = Parent.EPC.callSPSWrapper<
+                  rt::SPSSimpleExecutorMemoryManagerDeallocateSignature>(
+            Parent.SAs.Deallocate.getValue(), Err, Parent.SAs.Allocator,
+            ArrayRef<ExecutorAddress>(TargetAddr)))
       return E2;
     return Err;
   }
@@ -81,7 +83,6 @@ public:
 private:
   EPCGenericJITLinkMemoryManager &Parent;
   ExecutorAddress TargetAddr;
-  uint64_t TargetSize;
   std::unique_ptr<char[]> WorkingBuffer;
   SegInfoMap Segs;
 };
@@ -111,8 +112,9 @@ EPCGenericJITLinkMemoryManager::allocate(const jitlink::JITLinkDylib *JD,
   if (WorkingSize > 0)
     WorkingBuffer = std::make_unique<char[]>(WorkingSize);
   Expected<ExecutorAddress> TargetAllocAddr((ExecutorAddress()));
-  if (auto Err = EPC.callSPSWrapper<rt::SPSMemoryReserveSignature>(
-          FAs.Reserve.getValue(), TargetAllocAddr, AllocSize))
+  if (auto Err = EPC.callSPSWrapper<
+                 rt::SPSSimpleExecutorMemoryManagerReserveSignature>(
+          SAs.Reserve.getValue(), TargetAllocAddr, SAs.Allocator, AllocSize))
     return std::move(Err);
   if (!TargetAllocAddr)
     return TargetAllocAddr.takeError();
@@ -127,7 +129,7 @@ EPCGenericJITLinkMemoryManager::allocate(const jitlink::JITLinkDylib *JD,
     WorkingMem += Seg.ContentSize;
   }
 
-  return std::make_unique<Alloc>(*this, *TargetAllocAddr, AllocSize,
+  return std::make_unique<Alloc>(*this, *TargetAllocAddr,
                                  std::move(WorkingBuffer), std::move(Segs));
 }
 
index 16fee1f..0927fce 100644 (file)
@@ -12,12 +12,14 @@ namespace llvm {
 namespace orc {
 namespace rt {
 
-const char *MemoryReserveWrapperName =
-    "__llvm_orc_bootstrap_mem_reserve_wrapper";
-const char *MemoryFinalizeWrapperName =
-    "__llvm_orc_bootstrap_mem_finalize_wrapper";
-const char *MemoryDeallocateWrapperName =
-    "__llvm_orc_bootstrap_mem_deallocate_wrapper";
+const char *SimpleExecutorMemoryManagerInstanceName =
+    "__llvm_orc_SimpleExecutorMemoryManager_Instance";
+const char *SimpleExecutorMemoryManagerReserveWrapperName =
+    "__llvm_orc_SimpleExecutorMemoryManager_reserve_wrapper";
+const char *SimpleExecutorMemoryManagerFinalizeWrapperName =
+    "__llvm_orc_SimpleExecutorMemoryManager_finalize_wrapper";
+const char *SimpleExecutorMemoryManagerDeallocateWrapperName =
+    "__llvm_orc_SimpleExecutorMemoryManager_deallocate_wrapper";
 const char *MemoryWriteUInt8sWrapperName =
     "__llvm_orc_bootstrap_mem_write_uint8s_wrapper";
 const char *MemoryWriteUInt16sWrapperName =
index c9d7db2..c720f9d 100644 (file)
@@ -170,14 +170,16 @@ void SimpleRemoteEPC::handleDisconnect(Error Err) {
 
 Expected<std::unique_ptr<jitlink::JITLinkMemoryManager>>
 SimpleRemoteEPC::createMemoryManager() {
-  EPCGenericJITLinkMemoryManager::FuncAddrs FAs;
+  EPCGenericJITLinkMemoryManager::SymbolAddrs SAs;
   if (auto Err = getBootstrapSymbols(
-          {{FAs.Reserve, rt::MemoryReserveWrapperName},
-           {FAs.Finalize, rt::MemoryFinalizeWrapperName},
-           {FAs.Deallocate, rt::MemoryDeallocateWrapperName}}))
+          {{SAs.Allocator, rt::SimpleExecutorMemoryManagerInstanceName},
+           {SAs.Reserve, rt::SimpleExecutorMemoryManagerReserveWrapperName},
+           {SAs.Finalize, rt::SimpleExecutorMemoryManagerFinalizeWrapperName},
+           {SAs.Deallocate,
+            rt::SimpleExecutorMemoryManagerDeallocateWrapperName}}))
     return std::move(Err);
 
-  return std::make_unique<EPCGenericJITLinkMemoryManager>(*this, FAs);
+  return std::make_unique<EPCGenericJITLinkMemoryManager>(*this, SAs);
 }
 
 Expected<std::unique_ptr<ExecutorProcessControl::MemoryAccess>>
index 1ae02c3..6bed786 100644 (file)
@@ -2,6 +2,7 @@ add_llvm_component_library(LLVMOrcTargetProcess
   JITLoaderGDB.cpp
   OrcRTBootstrap.cpp
   RegisterEHFrames.cpp
+  SimpleExecutorMemoryManager.cpp
   SimpleRemoteEPCServer.cpp
   TargetExecutionUtils.cpp
 
index a95bcdb..eb80eb3 100644 (file)
@@ -20,57 +20,6 @@ namespace llvm {
 namespace orc {
 namespace rt_bootstrap {
 
-static llvm::orc::shared::detail::CWrapperFunctionResult
-reserveWrapper(const char *ArgData, size_t ArgSize) {
-  return WrapperFunction<rt::SPSMemoryReserveSignature>::handle(
-             ArgData, ArgSize,
-             [](uint64_t Size) -> Expected<ExecutorAddress> {
-               std::error_code EC;
-               auto MB = sys::Memory::allocateMappedMemory(
-                   Size, 0, sys::Memory::MF_READ | sys::Memory::MF_WRITE, EC);
-               if (EC)
-                 return errorCodeToError(EC);
-               return ExecutorAddress::fromPtr(MB.base());
-             })
-      .release();
-}
-
-static llvm::orc::shared::detail::CWrapperFunctionResult
-finalizeWrapper(const char *ArgData, size_t ArgSize) {
-  return WrapperFunction<rt::SPSMemoryFinalizeSignature>::handle(
-             ArgData, ArgSize,
-             [](const tpctypes::FinalizeRequest &FR) -> Error {
-               for (auto &Seg : FR) {
-                 char *Mem = Seg.Addr.toPtr<char *>();
-                 memcpy(Mem, Seg.Content.data(), Seg.Content.size());
-                 memset(Mem + Seg.Content.size(), 0,
-                        Seg.Size - Seg.Content.size());
-                 assert(Seg.Size <= std::numeric_limits<size_t>::max());
-                 if (auto EC = sys::Memory::protectMappedMemory(
-                         {Mem, static_cast<size_t>(Seg.Size)},
-                         tpctypes::fromWireProtectionFlags(Seg.Prot)))
-                   return errorCodeToError(EC);
-                 if (Seg.Prot & tpctypes::WPF_Exec)
-                   sys::Memory::InvalidateInstructionCache(Mem, Seg.Size);
-               }
-               return Error::success();
-             })
-      .release();
-}
-
-static llvm::orc::shared::detail::CWrapperFunctionResult
-deallocateWrapper(const char *ArgData, size_t ArgSize) {
-  return WrapperFunction<rt::SPSMemoryDeallocateSignature>::handle(
-             ArgData, ArgSize,
-             [](ExecutorAddress Base, uint64_t Size) -> Error {
-               sys::MemoryBlock MB(Base.toPtr<void *>(), Size);
-               if (auto EC = sys::Memory::releaseMappedMemory(MB))
-                 return errorCodeToError(EC);
-               return Error::success();
-             })
-      .release();
-}
-
 template <typename WriteT, typename SPSWriteT>
 static llvm::orc::shared::detail::CWrapperFunctionResult
 writeUIntsWrapper(const char *ArgData, size_t ArgSize) {
@@ -108,10 +57,6 @@ runAsMainWrapper(const char *ArgData, size_t ArgSize) {
 }
 
 void addTo(StringMap<ExecutorAddress> &M) {
-  M[rt::MemoryReserveWrapperName] = ExecutorAddress::fromPtr(&reserveWrapper);
-  M[rt::MemoryFinalizeWrapperName] = ExecutorAddress::fromPtr(&finalizeWrapper);
-  M[rt::MemoryDeallocateWrapperName] =
-      ExecutorAddress::fromPtr(&deallocateWrapper);
   M[rt::MemoryWriteUInt8sWrapperName] = ExecutorAddress::fromPtr(
       &writeUIntsWrapper<tpctypes::UInt8Write,
                          shared::SPSMemoryAccessUInt8Write>);
diff --git a/llvm/lib/ExecutionEngine/Orc/TargetProcess/SimpleExecutorMemoryManager.cpp b/llvm/lib/ExecutionEngine/Orc/TargetProcess/SimpleExecutorMemoryManager.cpp
new file mode 100644 (file)
index 0000000..ea839f9
--- /dev/null
@@ -0,0 +1,251 @@
+//===------- ResourceManager.cpp - Abstract executor resource mgmt --------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/ExecutionEngine/Orc/TargetProcess/SimpleExecutorMemoryManager.h"
+
+#include "llvm/ExecutionEngine/Orc/Shared/OrcRTBridge.h"
+#include "llvm/Support/FormatVariadic.h"
+
+#define DEBUG_TYPE "orc"
+
+namespace llvm {
+namespace orc {
+namespace rt_bootstrap {
+
+SimpleExecutorMemoryManager::~SimpleExecutorMemoryManager() {
+  assert(Allocations.empty() && "shutdown not called?");
+}
+
+Expected<ExecutorAddress> SimpleExecutorMemoryManager::allocate(uint64_t Size) {
+  std::error_code EC;
+  auto MB = sys::Memory::allocateMappedMemory(
+      Size, 0, sys::Memory::MF_READ | sys::Memory::MF_WRITE, EC);
+  if (EC)
+    return errorCodeToError(EC);
+  std::lock_guard<std::mutex> Lock(M);
+  assert(!Allocations.count(MB.base()) && "Duplicate allocation addr");
+  Allocations[MB.base()].Size = Size;
+  return ExecutorAddress::fromPtr(MB.base());
+}
+
+Error SimpleExecutorMemoryManager::finalize(tpctypes::FinalizeRequest &FR) {
+  ExecutorAddress Base(~0ULL);
+  std::vector<tpctypes::SupportFunctionCall> DeallocationActions;
+  size_t SuccessfulFinalizationActions = 0;
+
+  for (auto &Seg : FR.Segments)
+    Base = std::min(Base, Seg.Addr);
+
+  for (auto &ActPair : FR.Actions)
+    if (ActPair.Deallocate.Func)
+      DeallocationActions.push_back(ActPair.Deallocate);
+
+  // Get the Allocation for this finalization.
+  size_t AllocSize = 0;
+  {
+    std::lock_guard<std::mutex> Lock(M);
+    auto I = Allocations.find(Base.toPtr<void *>());
+    if (I == Allocations.end())
+      return make_error<StringError>("Attempt to finalize unrecognized "
+                                     "allocation " +
+                                         formatv("{0:x}", Base.getValue()),
+                                     inconvertibleErrorCode());
+    AllocSize = I->second.Size;
+    I->second.DeallocationActions = std::move(DeallocationActions);
+  }
+  ExecutorAddress AllocEnd = Base + ExecutorAddrDiff(AllocSize);
+
+  // Bail-out function: this will run deallocation actions corresponding to any
+  // completed finalization actions, then deallocate memory.
+  auto BailOut = [&](Error Err) {
+    std::pair<void *, Allocation> AllocToDestroy;
+
+    // Get allocation to destory.
+    {
+      std::lock_guard<std::mutex> Lock(M);
+      auto I = Allocations.find(Base.toPtr<void *>());
+
+      // Check for missing allocation (effective a double free).
+      if (I == Allocations.end())
+        return joinErrors(
+            std::move(Err),
+            make_error<StringError>("No allocation entry found "
+                                    "for " +
+                                        formatv("{0:x}", Base.getValue()),
+                                    inconvertibleErrorCode()));
+      AllocToDestroy = std::move(*I);
+      Allocations.erase(I);
+    }
+
+    // Run deallocation actions for all completed finalization actions.
+    while (SuccessfulFinalizationActions)
+      Err = joinErrors(
+          std::move(Err),
+          FR.Actions[--SuccessfulFinalizationActions].Deallocate.run());
+
+    // Deallocate memory.
+    sys::MemoryBlock MB(AllocToDestroy.first, AllocToDestroy.second.Size);
+    if (auto EC = sys::Memory::releaseMappedMemory(MB))
+      Err = joinErrors(std::move(Err), errorCodeToError(EC));
+
+    return Err;
+  };
+
+  // Copy content and apply permissions.
+  for (auto &Seg : FR.Segments) {
+
+    // Check segment ranges.
+    if (LLVM_UNLIKELY(Seg.Size < Seg.Content.size()))
+      return BailOut(make_error<StringError>(
+          formatv("Segment {0:x} content size ({1:x} bytes) "
+                  "exceeds segment size ({2:x} bytes)",
+                  Seg.Addr.getValue(), Seg.Content.size(), Seg.Size),
+          inconvertibleErrorCode()));
+    ExecutorAddress SegEnd = Seg.Addr + ExecutorAddrDiff(Seg.Size);
+    if (LLVM_UNLIKELY(Seg.Addr < Base || SegEnd > AllocEnd))
+      return BailOut(make_error<StringError>(
+          formatv("Segment {0:x} -- {1:x} crosses boundary of "
+                  "allocation {2:x} -- {3:x}",
+                  Seg.Addr.getValue(), SegEnd.getValue(), Base.getValue(),
+                  AllocEnd.getValue()),
+          inconvertibleErrorCode()));
+
+    char *Mem = Seg.Addr.toPtr<char *>();
+    memcpy(Mem, Seg.Content.data(), Seg.Content.size());
+    memset(Mem + Seg.Content.size(), 0, Seg.Size - Seg.Content.size());
+    assert(Seg.Size <= std::numeric_limits<size_t>::max());
+    if (auto EC = sys::Memory::protectMappedMemory(
+            {Mem, static_cast<size_t>(Seg.Size)},
+            tpctypes::fromWireProtectionFlags(Seg.Prot)))
+      return BailOut(errorCodeToError(EC));
+    if (Seg.Prot & tpctypes::WPF_Exec)
+      sys::Memory::InvalidateInstructionCache(Mem, Seg.Size);
+  }
+
+  // Run finalization actions.
+  for (auto &ActPair : FR.Actions) {
+    if (auto Err = ActPair.Finalize.run())
+      return BailOut(std::move(Err));
+    ++SuccessfulFinalizationActions;
+  }
+
+  return Error::success();
+}
+
+Error SimpleExecutorMemoryManager::deallocate(
+    const std::vector<ExecutorAddress> &Bases) {
+  std::vector<std::pair<void *, Allocation>> AllocPairs;
+  AllocPairs.reserve(Bases.size());
+
+  // Get allocation to destory.
+  Error Err = Error::success();
+  {
+    std::lock_guard<std::mutex> Lock(M);
+    for (auto &Base : Bases) {
+      auto I = Allocations.find(Base.toPtr<void *>());
+
+      // Check for missing allocation (effective a double free).
+      if (I != Allocations.end()) {
+        AllocPairs.push_back(std::move(*I));
+        Allocations.erase(I);
+      } else
+        Err = joinErrors(
+            std::move(Err),
+            make_error<StringError>("No allocation entry found "
+                                    "for " +
+                                        formatv("{0:x}", Base.getValue()),
+                                    inconvertibleErrorCode()));
+    }
+  }
+
+  while (!AllocPairs.empty()) {
+    auto &P = AllocPairs.back();
+    Err = joinErrors(std::move(Err), deallocateImpl(P.first, P.second));
+    AllocPairs.pop_back();
+  }
+
+  return Err;
+}
+
+Error SimpleExecutorMemoryManager::shutdown() {
+
+  AllocationsMap AM;
+  {
+    std::lock_guard<std::mutex> Lock(M);
+    AM = std::move(Allocations);
+  }
+
+  Error Err = Error::success();
+  for (auto &KV : AM)
+    Err = joinErrors(std::move(Err), deallocateImpl(KV.first, KV.second));
+  return Err;
+}
+
+void SimpleExecutorMemoryManager::addBootstrapSymbols(
+    StringMap<ExecutorAddress> &M) {
+  M[rt::SimpleExecutorMemoryManagerInstanceName] =
+      ExecutorAddress::fromPtr(this);
+  M[rt::SimpleExecutorMemoryManagerReserveWrapperName] =
+      ExecutorAddress::fromPtr(&reserveWrapper);
+  M[rt::SimpleExecutorMemoryManagerFinalizeWrapperName] =
+      ExecutorAddress::fromPtr(&finalizeWrapper);
+  M[rt::SimpleExecutorMemoryManagerDeallocateWrapperName] =
+      ExecutorAddress::fromPtr(&deallocateWrapper);
+}
+
+Error SimpleExecutorMemoryManager::deallocateImpl(void *Base, Allocation &A) {
+  Error Err = Error::success();
+
+  while (!A.DeallocationActions.empty()) {
+    Err = joinErrors(std::move(Err), A.DeallocationActions.back().run());
+    A.DeallocationActions.pop_back();
+  }
+
+  sys::MemoryBlock MB(Base, A.Size);
+  if (auto EC = sys::Memory::releaseMappedMemory(MB))
+    Err = joinErrors(std::move(Err), errorCodeToError(EC));
+
+  return Err;
+}
+
+llvm::orc::shared::detail::CWrapperFunctionResult
+SimpleExecutorMemoryManager::reserveWrapper(const char *ArgData,
+                                            size_t ArgSize) {
+  return shared::WrapperFunction<
+             rt::SPSSimpleExecutorMemoryManagerReserveSignature>::
+      handle(ArgData, ArgSize,
+             shared::makeMethodWrapperHandler(
+                 &SimpleExecutorMemoryManager::allocate))
+          .release();
+}
+
+llvm::orc::shared::detail::CWrapperFunctionResult
+SimpleExecutorMemoryManager::finalizeWrapper(const char *ArgData,
+                                             size_t ArgSize) {
+  return shared::WrapperFunction<
+             rt::SPSSimpleExecutorMemoryManagerFinalizeSignature>::
+      handle(ArgData, ArgSize,
+             shared::makeMethodWrapperHandler(
+                 &SimpleExecutorMemoryManager::finalize))
+          .release();
+}
+
+llvm::orc::shared::detail::CWrapperFunctionResult
+SimpleExecutorMemoryManager::deallocateWrapper(const char *ArgData,
+                                               size_t ArgSize) {
+  return shared::WrapperFunction<
+             rt::SPSSimpleExecutorMemoryManagerDeallocateSignature>::
+      handle(ArgData, ArgSize,
+             shared::makeMethodWrapperHandler(
+                 &SimpleExecutorMemoryManager::deallocate))
+          .release();
+}
+
+} // namespace rt_bootstrap
+} // end namespace orc
+} // end namespace llvm
index e95c6cb..e49cce5 100644 (file)
@@ -22,6 +22,8 @@ using namespace llvm::orc::shared;
 namespace llvm {
 namespace orc {
 
+ExecutorBootstrapService::~ExecutorBootstrapService() {}
+
 SimpleRemoteEPCServer::Dispatcher::~Dispatcher() {}
 
 #if LLVM_ENABLE_THREADS
index 2fb4cb3..e79639b 100644 (file)
@@ -14,6 +14,7 @@
 #include "llvm/ExecutionEngine/Orc/Shared/FDRawByteChannel.h"
 #include "llvm/ExecutionEngine/Orc/TargetProcess/JITLoaderGDB.h"
 #include "llvm/ExecutionEngine/Orc/TargetProcess/RegisterEHFrames.h"
+#include "llvm/ExecutionEngine/Orc/TargetProcess/SimpleExecutorMemoryManager.h"
 #include "llvm/ExecutionEngine/Orc/TargetProcess/SimpleRemoteEPCServer.h"
 #include "llvm/Support/DynamicLibrary.h"
 #include "llvm/Support/Error.h"
@@ -138,8 +139,16 @@ int main(int argc, char *argv[]) {
 
   auto Server =
       ExitOnErr(SimpleRemoteEPCServer::Create<FDSimpleRemoteEPCTransport>(
-          std::make_unique<SimpleRemoteEPCServer::ThreadDispatcher>(),
-          SimpleRemoteEPCServer::defaultBootstrapSymbols(), InFD, OutFD));
+          [](SimpleRemoteEPCServer::Setup &S) -> Error {
+            S.setDispatcher(
+                std::make_unique<SimpleRemoteEPCServer::ThreadDispatcher>());
+            S.bootstrapSymbols() =
+                SimpleRemoteEPCServer::defaultBootstrapSymbols();
+            S.services().push_back(
+                std::make_unique<rt_bootstrap::SimpleExecutorMemoryManager>());
+            return Error::success();
+          },
+          InFD, OutFD));
 
   ExitOnErr(Server->waitForDisconnect());
   return 0;
index 25a0e1f..7f6796b 100644 (file)
@@ -30,6 +30,7 @@ add_llvm_unittest(OrcJITTests
   ResourceTrackerTest.cpp
   RPCUtilsTest.cpp
   RTDyldObjectLinkingLayerTest.cpp
+  SimpleExecutorMemoryManagerTest.cpp
   SimplePackedSerializationTest.cpp
   SymbolStringPoolTest.cpp
   ThreadSafeModuleTest.cpp
index 9a0481c..2f34811 100644 (file)
@@ -9,12 +9,16 @@
 #include "OrcTestCommon.h"
 
 #include "llvm/ExecutionEngine/Orc/EPCGenericJITLinkMemoryManager.h"
+
+#include "llvm/ADT/DenseMap.h"
 #include "llvm/ExecutionEngine/Orc/Shared/OrcRTBridge.h"
 #include "llvm/ExecutionEngine/Orc/Shared/TargetProcessControlTypes.h"
+#include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/Memory.h"
 #include "llvm/Testing/Support/Error.h"
 
 #include <limits>
+#include <vector>
 
 using namespace llvm;
 using namespace llvm::orc;
@@ -22,66 +26,95 @@ using namespace llvm::orc::shared;
 
 namespace {
 
+class SimpleAllocator {
+public:
+  Expected<ExecutorAddress> reserve(uint64_t Size) {
+    std::error_code EC;
+    auto MB = sys::Memory::allocateMappedMemory(
+        Size, 0, sys::Memory::MF_READ | sys::Memory::MF_WRITE, EC);
+    if (EC)
+      return errorCodeToError(EC);
+    Blocks[MB.base()] = sys::OwningMemoryBlock(std::move(MB));
+    return ExecutorAddress::fromPtr(MB.base());
+  }
+
+  Error finalize(tpctypes::FinalizeRequest FR) {
+    for (auto &Seg : FR.Segments) {
+      char *Mem = Seg.Addr.toPtr<char *>();
+      memcpy(Mem, Seg.Content.data(), Seg.Content.size());
+      memset(Mem + Seg.Content.size(), 0, Seg.Size - Seg.Content.size());
+      assert(Seg.Size <= std::numeric_limits<size_t>::max());
+      if (auto EC = sys::Memory::protectMappedMemory(
+              {Mem, static_cast<size_t>(Seg.Size)},
+              tpctypes::fromWireProtectionFlags(Seg.Prot)))
+        return errorCodeToError(EC);
+      if (Seg.Prot & tpctypes::WPF_Exec)
+        sys::Memory::InvalidateInstructionCache(Mem, Seg.Size);
+    }
+    return Error::success();
+  }
+
+  Error deallocate(std::vector<ExecutorAddress> &Bases) {
+    Error Err = Error::success();
+    for (auto &Base : Bases) {
+      auto I = Blocks.find(Base.toPtr<void *>());
+      if (I == Blocks.end()) {
+        Err = joinErrors(
+            std::move(Err),
+            make_error<StringError>("No allocation for " +
+                                        formatv("{0:x}", Base.getValue()),
+                                    inconvertibleErrorCode()));
+        continue;
+      }
+      auto MB = std::move(I->second);
+      Blocks.erase(I);
+      auto MBToRelease = MB.getMemoryBlock();
+      if (auto EC = sys::Memory::releaseMappedMemory(MBToRelease))
+        Err = joinErrors(std::move(Err), errorCodeToError(EC));
+    }
+    return Err;
+  }
+
+private:
+  DenseMap<void *, sys::OwningMemoryBlock> Blocks;
+};
+
 llvm::orc::shared::detail::CWrapperFunctionResult
 testReserve(const char *ArgData, size_t ArgSize) {
-  return WrapperFunction<rt::SPSMemoryReserveSignature>::handle(
-             ArgData, ArgSize,
-             [](uint64_t Size) -> Expected<ExecutorAddress> {
-               std::error_code EC;
-               auto MB = sys::Memory::allocateMappedMemory(
-                   Size, 0, sys::Memory::MF_READ | sys::Memory::MF_WRITE, EC);
-               if (EC)
-                 return errorCodeToError(EC);
-               return ExecutorAddress::fromPtr(MB.base());
-             })
-      .release();
+  return WrapperFunction<rt::SPSSimpleExecutorMemoryManagerReserveSignature>::
+      handle(ArgData, ArgSize,
+             makeMethodWrapperHandler(&SimpleAllocator::reserve))
+          .release();
 }
 
 llvm::orc::shared::detail::CWrapperFunctionResult
 testFinalize(const char *ArgData, size_t ArgSize) {
-  return WrapperFunction<rt::SPSMemoryFinalizeSignature>::handle(
-             ArgData, ArgSize,
-             [](const tpctypes::FinalizeRequest &FR) -> Error {
-               for (auto &Seg : FR) {
-                 char *Mem = Seg.Addr.toPtr<char *>();
-                 memcpy(Mem, Seg.Content.data(), Seg.Content.size());
-                 memset(Mem + Seg.Content.size(), 0,
-                        Seg.Size - Seg.Content.size());
-                 assert(Seg.Size <= std::numeric_limits<size_t>::max());
-                 if (auto EC = sys::Memory::protectMappedMemory(
-                         {Mem, static_cast<size_t>(Seg.Size)},
-                         tpctypes::fromWireProtectionFlags(Seg.Prot)))
-                   return errorCodeToError(EC);
-                 if (Seg.Prot & tpctypes::WPF_Exec)
-                   sys::Memory::InvalidateInstructionCache(Mem, Seg.Size);
-               }
-               return Error::success();
-             })
-      .release();
+  return WrapperFunction<rt::SPSSimpleExecutorMemoryManagerFinalizeSignature>::
+      handle(ArgData, ArgSize,
+             makeMethodWrapperHandler(&SimpleAllocator::finalize))
+          .release();
 }
 
 llvm::orc::shared::detail::CWrapperFunctionResult
 testDeallocate(const char *ArgData, size_t ArgSize) {
-  return WrapperFunction<rt::SPSMemoryDeallocateSignature>::handle(
-             ArgData, ArgSize,
-             [](ExecutorAddress Base, uint64_t Size) -> Error {
-               sys::MemoryBlock MB(Base.toPtr<void *>(), Size);
-               if (auto EC = sys::Memory::releaseMappedMemory(MB))
-                 return errorCodeToError(EC);
-               return Error::success();
-             })
-      .release();
+  return WrapperFunction<
+             rt::SPSSimpleExecutorMemoryManagerDeallocateSignature>::
+      handle(ArgData, ArgSize,
+             makeMethodWrapperHandler(&SimpleAllocator::deallocate))
+          .release();
 }
 
 TEST(EPCGenericJITLinkMemoryManagerTest, AllocFinalizeFree) {
   auto SelfEPC = cantFail(SelfExecutorProcessControl::Create());
+  SimpleAllocator SA;
 
-  EPCGenericJITLinkMemoryManager::FuncAddrs FAs;
-  FAs.Reserve = ExecutorAddress::fromPtr(&testReserve);
-  FAs.Finalize = ExecutorAddress::fromPtr(&testFinalize);
-  FAs.Deallocate = ExecutorAddress::fromPtr(&testDeallocate);
+  EPCGenericJITLinkMemoryManager::SymbolAddrs SAs;
+  SAs.Allocator = ExecutorAddress::fromPtr(&SA);
+  SAs.Reserve = ExecutorAddress::fromPtr(&testReserve);
+  SAs.Finalize = ExecutorAddress::fromPtr(&testFinalize);
+  SAs.Deallocate = ExecutorAddress::fromPtr(&testDeallocate);
 
-  auto MemMgr = std::make_unique<EPCGenericJITLinkMemoryManager>(*SelfEPC, FAs);
+  auto MemMgr = std::make_unique<EPCGenericJITLinkMemoryManager>(*SelfEPC, SAs);
 
   jitlink::JITLinkMemoryManager::SegmentsRequestMap SRM;
   StringRef Hello = "hello";
diff --git a/llvm/unittests/ExecutionEngine/Orc/SimpleExecutorMemoryManagerTest.cpp b/llvm/unittests/ExecutionEngine/Orc/SimpleExecutorMemoryManagerTest.cpp
new file mode 100644 (file)
index 0000000..5bd600a
--- /dev/null
@@ -0,0 +1,86 @@
+//===---------------- SimpleExecutorMemoryManagerTest.cpp -----------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/ExecutionEngine/Orc/TargetProcess/SimpleExecutorMemoryManager.h"
+#include "llvm/Testing/Support/Error.h"
+#include "gtest/gtest.h"
+
+#include <limits>
+#include <vector>
+
+using namespace llvm;
+using namespace llvm::orc;
+using namespace llvm::orc::shared;
+using namespace llvm::orc::rt_bootstrap;
+
+namespace {
+
+orc::shared::detail::CWrapperFunctionResult
+incrementWrapper(const char *ArgData, size_t ArgSize) {
+  return WrapperFunction<void(SPSExecutorAddress)>::handle(
+             ArgData, ArgSize,
+             [](ExecutorAddress A) { *A.toPtr<int *>() += 1; })
+      .release();
+}
+
+TEST(SimpleExecutorMemoryManagerTest, AllocFinalizeFree) {
+  SimpleExecutorMemoryManager MemMgr;
+
+  constexpr unsigned AllocSize = 16384;
+  auto Mem = MemMgr.allocate(AllocSize);
+  EXPECT_THAT_ERROR(Mem.takeError(), Succeeded());
+
+  std::string HW = "Hello, world!";
+
+  int FinalizeCounter = 0;
+  auto FinalizeCounterAddrArgBuffer =
+      orc::shared::detail::serializeViaSPSToWrapperFunctionResult<
+          SPSArgList<SPSExecutorAddress>>(
+          ExecutorAddress::fromPtr(&FinalizeCounter));
+
+  int DeallocateCounter = 0;
+  auto DeallocateCounterAddrArgBuffer =
+      orc::shared::detail::serializeViaSPSToWrapperFunctionResult<
+          SPSArgList<SPSExecutorAddress>>(
+          ExecutorAddress::fromPtr(&DeallocateCounter));
+
+  tpctypes::FinalizeRequest FR;
+  FR.Segments.push_back(
+      tpctypes::SegFinalizeRequest{tpctypes::WPF_Read | tpctypes::WPF_Write,
+                                   *Mem,
+                                   AllocSize,
+                                   {HW.data(), HW.size() + 1}});
+  FR.Actions.push_back(
+      {/* Finalize: */
+       {ExecutorAddress::fromPtr(incrementWrapper),
+        ExecutorAddress::fromPtr(FinalizeCounterAddrArgBuffer.data()),
+        FinalizeCounterAddrArgBuffer.size()},
+       /*  Deallocate: */
+       {ExecutorAddress::fromPtr(incrementWrapper),
+        ExecutorAddress::fromPtr(DeallocateCounterAddrArgBuffer.data()),
+        DeallocateCounterAddrArgBuffer.size()}});
+
+  EXPECT_EQ(FinalizeCounter, 0);
+  EXPECT_EQ(DeallocateCounter, 0);
+
+  auto FinalizeErr = MemMgr.finalize(FR);
+  EXPECT_THAT_ERROR(std::move(FinalizeErr), Succeeded());
+
+  EXPECT_EQ(FinalizeCounter, 1);
+  EXPECT_EQ(DeallocateCounter, 0);
+
+  EXPECT_EQ(HW, std::string(Mem->toPtr<const char *>()));
+
+  auto DeallocateErr = MemMgr.deallocate({*Mem});
+  EXPECT_THAT_ERROR(std::move(DeallocateErr), Succeeded());
+
+  EXPECT_EQ(FinalizeCounter, 1);
+  EXPECT_EQ(DeallocateCounter, 1);
+}
+
+} // namespace