[llvm] 6ede652 - [JITLink][Orc] Add MemoryMapper interface with InProcess implementation

Stefan Gränitz via llvm-commits llvm-commits at lists.llvm.org
Tue Jun 21 04:45:26 PDT 2022


Author: Anubhab Ghosh
Date: 2022-06-21T13:44:17+02:00
New Revision: 6ede65205073d3cf6b1ed4d101e66eae3e0fc8e6

URL: https://github.com/llvm/llvm-project/commit/6ede65205073d3cf6b1ed4d101e66eae3e0fc8e6
DIFF: https://github.com/llvm/llvm-project/commit/6ede65205073d3cf6b1ed4d101e66eae3e0fc8e6.diff

LOG: [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).

Reviewed By: sgraenitz, lhames

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

Added: 
    llvm/include/llvm/ExecutionEngine/Orc/MemoryMapper.h
    llvm/lib/ExecutionEngine/Orc/MemoryMapper.cpp
    llvm/unittests/ExecutionEngine/Orc/MemoryMapperTest.cpp

Modified: 
    llvm/lib/ExecutionEngine/Orc/CMakeLists.txt
    llvm/unittests/ExecutionEngine/Orc/CMakeLists.txt

Removed: 
    


################################################################################
diff  --git a/llvm/include/llvm/ExecutionEngine/Orc/MemoryMapper.h b/llvm/include/llvm/ExecutionEngine/Orc/MemoryMapper.h
new file mode 100644
index 0000000000000..d023bfbdb5b6f
--- /dev/null
+++ b/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

diff  --git a/llvm/lib/ExecutionEngine/Orc/CMakeLists.txt b/llvm/lib/ExecutionEngine/Orc/CMakeLists.txt
index cf8e0c85840c9..1c9b9e8a1b31e 100644
--- a/llvm/lib/ExecutionEngine/Orc/CMakeLists.txt
+++ b/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

diff  --git a/llvm/lib/ExecutionEngine/Orc/MemoryMapper.cpp b/llvm/lib/ExecutionEngine/Orc/MemoryMapper.cpp
new file mode 100644
index 0000000000000..f00cac3e2c3e9
--- /dev/null
+++ b/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<Error> 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<Error> P;
+  auto F = P.get_future();
+  release(ReservationAddrs, [&](Error Err) { P.set_value(std::move(Err)); });
+  cantFail(F.get());
+}
+
+} // namespace orc
+
+} // namespace llvm

diff  --git a/llvm/unittests/ExecutionEngine/Orc/CMakeLists.txt b/llvm/unittests/ExecutionEngine/Orc/CMakeLists.txt
index 9d45957266b8a..dd4028e5302ca 100644
--- a/llvm/unittests/ExecutionEngine/Orc/CMakeLists.txt
+++ b/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

diff  --git a/llvm/unittests/ExecutionEngine/Orc/MemoryMapperTest.cpp b/llvm/unittests/ExecutionEngine/Orc/MemoryMapperTest.cpp
new file mode 100644
index 0000000000000..f352715429c02
--- /dev/null
+++ b/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<Expected<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<Expected<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<Error> 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<Error> 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


        


More information about the llvm-commits mailing list