Re-apply "[JITLink][Orc] Add MemoryMapper interface with InProcess implementation"

[JITLink][Orc] Add MemoryMapper interface with InProcess implementation

MemoryMapper class takes care of cross-process and in-process address space
reservation, mapping, transferring content and applying protections.

Implementations of this class can support different ways to do this such
as using shared memory, transferring memory contents over EPC or just
mapping memory in the same process (InProcessMemoryMapper).

The original patch landed with commit 6ede652050
It was reverted temporarily in commit 6a4056ab2a

Reviewed By: sgraenitz, lhames

Differential Revision: https://reviews.llvm.org/D127491

llvm/include/llvm/ExecutionEngine/Orc/MemoryMapper.h

@@ -0,0 +1,115 @@
+//===- MemoryMapper.h - Cross-process memory mapper -------------*- 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Cross-process (and in-process) memory mapping and transfer
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_EXECUTIONENGINE_ORC_MEMORYMAPPER_H
+#define LLVM_EXECUTIONENGINE_ORC_MEMORYMAPPER_H
+
+#include "llvm/ExecutionEngine/Orc/Core.h"
+
+#include <mutex>
+
+namespace llvm {
+namespace orc {
+
+/// Manages mapping, content transfer and protections for JIT memory
+class MemoryMapper {
+public:
+  /// Represents a single allocation containing multiple segments and
+  /// initialization and deinitialization actions
+  struct AllocInfo {
+    struct SegInfo {
+      ExecutorAddrDiff Offset;
+      const char *WorkingMem;
+      size_t ContentSize;
+      size_t ZeroFillSize;
+      unsigned Prot;
+    };
+
+    ExecutorAddr MappingBase;
+    std::vector<SegInfo> Segments;
+    shared::AllocActions Actions;
+  };
+
+  using OnReservedFunction = unique_function<void(Expected<ExecutorAddrRange>)>;
+
+  /// Reserves address space in executor process
+  virtual void reserve(size_t NumBytes, OnReservedFunction OnReserved) = 0;
+
+  /// Provides working memory
+  virtual char *prepare(ExecutorAddr Addr, size_t ContentSize) = 0;
+
+  using OnInitializedFunction = unique_function<void(Expected<ExecutorAddr>)>;
+
+  /// Ensures executor memory is synchronized with working copy memory, sends
+  /// functions to be called after initilization and before deinitialization and
+  /// applies memory protections
+  /// Returns a unique address identifying the allocation. This address should
+  /// be passed to deinitialize to run deallocation actions (and reset
+  /// permissions where possible).
+  virtual void initialize(AllocInfo &AI,
+                          OnInitializedFunction OnInitialized) = 0;
+
+  using OnDeinitializedFunction = unique_function<void(Error)>;
+
+  /// Runs previously specified deinitialization actions
+  /// Executor addresses returned by initialize should be passed
+  virtual void deinitialize(ArrayRef<ExecutorAddr> Allocations,
+                            OnDeinitializedFunction OnDeInitialized) = 0;
+
+  using OnReleasedFunction = unique_function<void(Error)>;
+
+  /// Release address space acquired through reserve()
+  virtual void release(ArrayRef<ExecutorAddr> Reservations,
+                       OnReleasedFunction OnRelease) = 0;
+
+  virtual ~MemoryMapper();
+};
+
+class InProcessMemoryMapper final : public MemoryMapper {
+public:
+  InProcessMemoryMapper() {}
+
+  void reserve(size_t NumBytes, OnReservedFunction OnReserved) override;
+
+  void initialize(AllocInfo &AI, OnInitializedFunction OnInitialized) override;
+
+  char *prepare(ExecutorAddr Addr, size_t ContentSize) override;
+
+  void deinitialize(ArrayRef<ExecutorAddr> Allocations,
+                    OnDeinitializedFunction OnDeInitialized) override;
+
+  void release(ArrayRef<ExecutorAddr> Reservations,
+               OnReleasedFunction OnRelease) override;
+
+  ~InProcessMemoryMapper() override;
+
+private:
+  struct Allocation {
+    std::vector<shared::WrapperFunctionCall> DeinitializationActions;
+  };
+  using AllocationMap = DenseMap<ExecutorAddr, Allocation>;
+
+  struct Reservation {
+    size_t Size;
+    std::vector<ExecutorAddr> Allocations;
+  };
+  using ReservationMap = DenseMap<void *, Reservation>;
+
+  std::mutex Mutex;
+  ReservationMap Reservations;
+  AllocationMap Allocations;
+};
+
+} // namespace orc
+} // end namespace llvm
+
+#endif // LLVM_EXECUTIONENGINE_ORC_MEMORYMAPPER_H

llvm/lib/ExecutionEngine/Orc/CMakeLists.txt

@@ -23,6 +23,7 @@ add_llvm_component_library(LLVMOrcJIT
   LookupAndRecordAddrs.cpp
   LLJIT.cpp
   MachOPlatform.cpp
+  MemoryMapper.cpp
   ELFNixPlatform.cpp
   Mangling.cpp
   ObjectLinkingLayer.cpp

llvm/lib/ExecutionEngine/Orc/MemoryMapper.cpp

@@ -0,0 +1,152 @@
+//===- MemoryMapper.cpp - Cross-process memory mapper ------------*- 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/ExecutionEngine/Orc/MemoryMapper.h"
+
+namespace llvm {
+namespace orc {
+
+MemoryMapper::~MemoryMapper() {}
+
+void InProcessMemoryMapper::reserve(size_t NumBytes,
+                                    OnReservedFunction OnReserved) {
+  std::error_code EC;
+  auto MB = sys::Memory::allocateMappedMemory(
+      NumBytes, nullptr, sys::Memory::MF_READ | sys::Memory::MF_WRITE, EC);
+
+  if (EC)
+    return OnReserved(errorCodeToError(EC));
+
+  {
+    std::lock_guard<std::mutex> Lock(Mutex);
+    Reservations[MB.base()].Size = MB.allocatedSize();
+  }
+
+  OnReserved(
+      ExecutorAddrRange(ExecutorAddr::fromPtr(MB.base()), MB.allocatedSize()));
+}
+
+char *InProcessMemoryMapper::prepare(ExecutorAddr Addr, size_t ContentSize) {
+  return Addr.toPtr<char *>();
+}
+
+void InProcessMemoryMapper::initialize(MemoryMapper::AllocInfo &AI,
+                                       OnInitializedFunction OnInitialized) {
+  ExecutorAddr MinAddr(~0ULL);
+
+  for (auto &Segment : AI.Segments) {
+    auto Base = AI.MappingBase + Segment.Offset;
+    auto Size = Segment.ContentSize + Segment.ZeroFillSize;
+
+    if (Base < MinAddr)
+      MinAddr = Base;
+
+    std::memset((Base + Segment.ContentSize).toPtr<void *>(), 0,
+                Segment.ZeroFillSize);
+
+    if (auto EC = sys::Memory::protectMappedMemory({Base.toPtr<void *>(), Size},
+                                                   Segment.Prot)) {
+      return OnInitialized(errorCodeToError(EC));
+    }
+    if (Segment.Prot & sys::Memory::MF_EXEC)
+      sys::Memory::InvalidateInstructionCache(Base.toPtr<void *>(), Size);
+  }
+
+  auto DeinitializeActions = shared::runFinalizeActions(AI.Actions);
+  if (!DeinitializeActions)
+    return OnInitialized(DeinitializeActions.takeError());
+
+  {
+    std::lock_guard<std::mutex> Lock(Mutex);
+    Allocations[MinAddr].DeinitializationActions =
+        std::move(*DeinitializeActions);
+    Reservations[AI.MappingBase.toPtr<void *>()].Allocations.push_back(MinAddr);
+  }
+
+  OnInitialized(MinAddr);
+}
+
+void InProcessMemoryMapper::deinitialize(
+    ArrayRef<ExecutorAddr> Bases,
+    MemoryMapper::OnDeinitializedFunction OnDeinitialized) {
+  Error AllErr = Error::success();
+
+  {
+    std::lock_guard<std::mutex> Lock(Mutex);
+
+    for (auto Base : Bases) {
+
+      if (Error Err = shared::runDeallocActions(
+              Allocations[Base].DeinitializationActions)) {
+        AllErr = joinErrors(std::move(AllErr), std::move(Err));
+      }
+
+      Allocations.erase(Base);
+    }
+  }
+
+  OnDeinitialized(std::move(AllErr));
+}
+
+void InProcessMemoryMapper::release(ArrayRef<ExecutorAddr> Bases,
+                                    OnReleasedFunction OnReleased) {
+  Error Err = Error::success();
+
+  for (auto Base : Bases) {
+    std::vector<ExecutorAddr> AllocAddrs;
+    size_t Size;
+    {
+      std::lock_guard<std::mutex> Lock(Mutex);
+      auto &R = Reservations[Base.toPtr<void *>()];
+      Size = R.Size;
+      AllocAddrs.swap(R.Allocations);
+    }
+
+    // deinitialize sub allocations
+    std::promise<MSVCPError> P;
+    auto F = P.get_future();
+    deinitialize(AllocAddrs, [&](Error Err) { P.set_value(std::move(Err)); });
+    if (Error E = F.get()) {
+      Err = joinErrors(std::move(Err), std::move(E));
+    }
+
+    // free the memory
+    auto MB = sys::MemoryBlock(Base.toPtr<void *>(), Size);
+
+    auto EC = sys::Memory::releaseMappedMemory(MB);
+    if (EC) {
+      Err = joinErrors(std::move(Err), errorCodeToError(EC));
+    }
+
+    std::lock_guard<std::mutex> Lock(Mutex);
+    Reservations.erase(Base.toPtr<void *>());
+  }
+
+  OnReleased(std::move(Err));
+}
+
+InProcessMemoryMapper::~InProcessMemoryMapper() {
+  std::vector<ExecutorAddr> ReservationAddrs;
+  {
+    std::lock_guard<std::mutex> Lock(Mutex);
+
+    ReservationAddrs.reserve(Reservations.size());
+    for (const auto &R : Reservations) {
+      ReservationAddrs.push_back(ExecutorAddr::fromPtr(R.getFirst()));
+    }
+  }
+
+  std::promise<MSVCPError> P;
+  auto F = P.get_future();
+  release(ReservationAddrs, [&](Error Err) { P.set_value(std::move(Err)); });
+  cantFail(F.get());
+}
+
+} // namespace orc
+
+} // namespace llvm

llvm/unittests/ExecutionEngine/Orc/CMakeLists.txt

@@ -24,6 +24,7 @@ add_llvm_unittest(OrcJITTests
   JITTargetMachineBuilderTest.cpp
   LazyCallThroughAndReexportsTest.cpp
   LookupAndRecordAddrsTest.cpp
+  MemoryMapperTest.cpp
   ObjectLinkingLayerTest.cpp
   OrcCAPITest.cpp
   OrcTestCommon.cpp

llvm/unittests/ExecutionEngine/Orc/MemoryMapperTest.cpp

@@ -0,0 +1,203 @@
+//===------------------------ MemoryMapperTest.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/MemoryMapper.h"
+#include "llvm/Support/Process.h"
+#include "llvm/Testing/Support/Error.h"
+#include "gtest/gtest.h"
+
+using namespace llvm;
+using namespace llvm::orc;
+using namespace llvm::orc::shared;
+
+namespace {
+
+Expected<ExecutorAddrRange> reserve(MemoryMapper &M, size_t NumBytes) {
+  std::promise<MSVCPExpected<ExecutorAddrRange>> P;
+  auto F = P.get_future();
+  M.reserve(NumBytes, [&](auto R) { P.set_value(std::move(R)); });
+  return F.get();
+}
+
+Expected<ExecutorAddr> initialize(MemoryMapper &M,
+                                  MemoryMapper::AllocInfo &AI) {
+  std::promise<MSVCPExpected<ExecutorAddr>> P;
+  auto F = P.get_future();
+  M.initialize(AI, [&](auto R) { P.set_value(std::move(R)); });
+  return F.get();
+}
+
+Error deinitialize(MemoryMapper &M,
+                   const std::vector<ExecutorAddr> &Allocations) {
+  std::promise<MSVCPError> P;
+  auto F = P.get_future();
+  M.deinitialize(Allocations, [&](auto R) { P.set_value(std::move(R)); });
+  return F.get();
+}
+
+Error release(MemoryMapper &M, const std::vector<ExecutorAddr> &Reservations) {
+  std::promise<MSVCPError> P;
+  auto F = P.get_future();
+  M.release(Reservations, [&](auto R) { P.set_value(std::move(R)); });
+  return F.get();
+}
+
+// A basic function to be used as both initializer/deinitializer
+orc::shared::CWrapperFunctionResult incrementWrapper(const char *ArgData,
+                                                     size_t ArgSize) {
+  return WrapperFunction<SPSError(SPSExecutorAddr)>::handle(
+             ArgData, ArgSize,
+             [](ExecutorAddr A) -> Error {
+               *A.toPtr<int *>() += 1;
+               return Error::success();
+             })
+      .release();
+}
+
+TEST(MemoryMapperTest, InitializeDeinitialize) {
+  // These counters are used to track how many times the initializer and
+  // deinitializer functions are called
+  int InitializeCounter = 0;
+  int DeinitializeCounter = 0;
+  {
+    std::unique_ptr<MemoryMapper> Mapper =
+        std::make_unique<InProcessMemoryMapper>();
+
+    // We will do two separate allocations
+    auto PageSize = cantFail(sys::Process::getPageSize());
+    auto TotalSize = PageSize * 2;
+
+    // Reserve address space
+    auto Mem1 = reserve(*Mapper, TotalSize);
+    EXPECT_THAT_ERROR(Mem1.takeError(), Succeeded());
+
+    // Test string for memory transfer
+    std::string HW = "Hello, world!";
+
+    {
+      // Provide working memory
+      char *WA1 = Mapper->prepare(Mem1->Start, HW.size() + 1);
+      std::strcpy(static_cast<char *>(WA1), HW.c_str());
+    }
+
+    // A structure to be passed to initialize
+    MemoryMapper::AllocInfo Alloc1;
+    {
+      MemoryMapper::AllocInfo::SegInfo Seg1;
+      Seg1.Offset = 0;
+      Seg1.ContentSize = HW.size();
+      Seg1.ZeroFillSize = PageSize - Seg1.ContentSize;
+      Seg1.Prot = sys::Memory::MF_READ | sys::Memory::MF_WRITE;
+
+      Alloc1.MappingBase = Mem1->Start;
+      Alloc1.Segments.push_back(Seg1);
+      Alloc1.Actions.push_back(
+          {cantFail(WrapperFunctionCall::Create<SPSArgList<SPSExecutorAddr>>(
+               ExecutorAddr::fromPtr(incrementWrapper),
+               ExecutorAddr::fromPtr(&InitializeCounter))),
+           cantFail(WrapperFunctionCall::Create<SPSArgList<SPSExecutorAddr>>(
+               ExecutorAddr::fromPtr(incrementWrapper),
+               ExecutorAddr::fromPtr(&DeinitializeCounter)))});
+    }
+
+    {
+      char *WA2 = Mapper->prepare(Mem1->Start + PageSize, HW.size() + 1);
+      std::strcpy(static_cast<char *>(WA2), HW.c_str());
+    }
+
+    MemoryMapper::AllocInfo Alloc2;
+    {
+      MemoryMapper::AllocInfo::SegInfo Seg2;
+      Seg2.Offset = PageSize;
+      Seg2.ContentSize = HW.size();
+      Seg2.ZeroFillSize = PageSize - Seg2.ContentSize;
+      Seg2.Prot = sys::Memory::MF_READ | sys::Memory::MF_WRITE;
+
+      Alloc2.MappingBase = Mem1->Start;
+      Alloc2.Segments.push_back(Seg2);
+      Alloc2.Actions.push_back(
+          {cantFail(WrapperFunctionCall::Create<SPSArgList<SPSExecutorAddr>>(
+               ExecutorAddr::fromPtr(incrementWrapper),
+               ExecutorAddr::fromPtr(&InitializeCounter))),
+           cantFail(WrapperFunctionCall::Create<SPSArgList<SPSExecutorAddr>>(
+               ExecutorAddr::fromPtr(incrementWrapper),
+               ExecutorAddr::fromPtr(&DeinitializeCounter)))});
+    }
+
+    EXPECT_EQ(InitializeCounter, 0);
+    EXPECT_EQ(DeinitializeCounter, 0);
+
+    // Set memory protections and run initializers
+    auto Init1 = initialize(*Mapper, Alloc1);
+    EXPECT_THAT_ERROR(Init1.takeError(), Succeeded());
+    EXPECT_EQ(HW, std::string(static_cast<char *>(Init1->toPtr<char *>())));
+
+    EXPECT_EQ(InitializeCounter, 1);
+    EXPECT_EQ(DeinitializeCounter, 0);
+
+    auto Init2 = initialize(*Mapper, Alloc2);
+    EXPECT_THAT_ERROR(Init2.takeError(), Succeeded());
+    EXPECT_EQ(HW, std::string(static_cast<char *>(Init2->toPtr<char *>())));
+
+    EXPECT_EQ(InitializeCounter, 2);
+    EXPECT_EQ(DeinitializeCounter, 0);
+
+    // Explicit deinitialization of first allocation
+    std::vector<ExecutorAddr> DeinitAddr = {*Init1};
+    EXPECT_THAT_ERROR(deinitialize(*Mapper, DeinitAddr), Succeeded());
+
+    EXPECT_EQ(InitializeCounter, 2);
+    EXPECT_EQ(DeinitializeCounter, 1);
+
+    // Test explicit release
+    {
+      auto Mem2 = reserve(*Mapper, PageSize);
+      EXPECT_THAT_ERROR(Mem2.takeError(), Succeeded());
+
+      char *WA = Mapper->prepare(Mem2->Start, HW.size() + 1);
+      std::strcpy(static_cast<char *>(WA), HW.c_str());
+
+      MemoryMapper::AllocInfo Alloc3;
+      {
+        MemoryMapper::AllocInfo::SegInfo Seg3;
+        Seg3.Offset = 0;
+        Seg3.ContentSize = HW.size();
+        Seg3.ZeroFillSize = PageSize - Seg3.ContentSize;
+        Seg3.Prot = sys::Memory::MF_READ | sys::Memory::MF_WRITE;
+
+        Alloc3.MappingBase = Mem2->Start;
+        Alloc3.Segments.push_back(Seg3);
+        Alloc3.Actions.push_back(
+            {cantFail(WrapperFunctionCall::Create<SPSArgList<SPSExecutorAddr>>(
+                 ExecutorAddr::fromPtr(incrementWrapper),
+                 ExecutorAddr::fromPtr(&InitializeCounter))),
+             cantFail(WrapperFunctionCall::Create<SPSArgList<SPSExecutorAddr>>(
+                 ExecutorAddr::fromPtr(incrementWrapper),
+                 ExecutorAddr::fromPtr(&DeinitializeCounter)))});
+      }
+      auto Init3 = initialize(*Mapper, Alloc3);
+      EXPECT_THAT_ERROR(Init3.takeError(), Succeeded());
+      EXPECT_EQ(HW, std::string(static_cast<char *>(Init3->toPtr<char *>())));
+
+      EXPECT_EQ(InitializeCounter, 3);
+      EXPECT_EQ(DeinitializeCounter, 1);
+
+      std::vector<ExecutorAddr> ReleaseAddrs = {Mem2->Start};
+      EXPECT_THAT_ERROR(release(*Mapper, ReleaseAddrs), Succeeded());
+
+      EXPECT_EQ(InitializeCounter, 3);
+      EXPECT_EQ(DeinitializeCounter, 2);
+    }
+  }
+
+  // Implicit deinitialization by the destructor
+  EXPECT_EQ(InitializeCounter, 3);
+  EXPECT_EQ(DeinitializeCounter, 3);
+}
+
+} // namespace