[ORC] Add finalization & deallocation actions, SimpleExecutorMemoryManager class

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.

Author: lhames

llvm/include/llvm/ExecutionEngine/Orc/EPCGenericJITLinkMemoryManager.h

@@ -27,16 +27,17 @@ 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;
   };
 
   /// 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 @@ class EPCGenericJITLinkMemoryManager : public jitlink::JITLinkMemoryManager {
   class Alloc;
 
   ExecutorProcessControl &EPC;
-  FuncAddrs FAs;
+  SymbolAddrs SAs;
 };
 
 } // end namespace orc

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

@@ -21,21 +21,27 @@ 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;
 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>);
 

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

@@ -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,14 +69,42 @@ 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;
   uint64_t Size;
   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>;
@@ -160,6 +196,48 @@ class SPSSerializationTraits<SPSMemoryProtectionFlags,
   }
 };
 
+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> {
@@ -181,6 +259,25 @@ class SPSSerializationTraits<SPSSegFinalizeRequest,
   }
 };
 
+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>> {

llvm/include/llvm/ExecutionEngine/Orc/TargetProcess/ExecutorBootstrapService.h

@@ -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

llvm/include/llvm/ExecutionEngine/Orc/TargetProcess/SimpleExecutorMemoryManager.h

@@ -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