[ORC] Add a new MemoryMapper-based JITLinkMemoryManager implementation.

MapperJITLinkMemoryManager supports executor memory management using any
implementation of MemoryMapper to do the transfer such as InProcessMapper or
SharedMemoryMapper.

Reviewed By: lhames

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

llvm/include/llvm/ExecutionEngine/Orc/MapperJITLinkMemoryManager.h

@@ -0,0 +1,56 @@
+//===--------------- MapperJITLinkMemoryManager.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
+//
+//===----------------------------------------------------------------------===//
+//
+// Implements JITLinkMemoryManager using MemoryMapper
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_EXECUTIONENGINE_ORC_MAPPERJITLINKMEMORYMANAGER_H
+#define LLVM_EXECUTIONENGINE_ORC_MAPPERJITLINKMEMORYMANAGER_H
+
+#include "llvm/ExecutionEngine/JITLink/JITLinkMemoryManager.h"
+#include "llvm/ExecutionEngine/Orc/Core.h"
+#include "llvm/ExecutionEngine/Orc/MemoryMapper.h"
+
+namespace llvm {
+namespace orc {
+
+class MapperJITLinkMemoryManager : public jitlink::JITLinkMemoryManager {
+public:
+  MapperJITLinkMemoryManager(std::unique_ptr<MemoryMapper> Mapper);
+
+  template <class MemoryMapperType, class... Args>
+  static Expected<std::unique_ptr<MapperJITLinkMemoryManager>>
+  CreateWithMapper(Args &&...A) {
+    auto Mapper = MemoryMapperType::Create(std::forward<Args>(A)...);
+    if (!Mapper)
+      return Mapper.takeError();
+
+    return std::make_unique<MapperJITLinkMemoryManager>(std::move(*Mapper));
+  }
+
+  void allocate(const jitlink::JITLinkDylib *JD, jitlink::LinkGraph &G,
+                OnAllocatedFunction OnAllocated) override;
+  // synchronous overload
+  using JITLinkMemoryManager::allocate;
+
+  void deallocate(std::vector<FinalizedAlloc> Allocs,
+                  OnDeallocatedFunction OnDeallocated) override;
+  // synchronous overload
+  using JITLinkMemoryManager::deallocate;
+
+private:
+  class InFlightAlloc;
+
+  std::unique_ptr<MemoryMapper> Mapper;
+};
+
+} // end namespace orc
+} // end namespace llvm
+
+#endif // LLVM_EXECUTIONENGINE_ORC_MAPPERJITLINKMEMORYMANAGER_H

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

@@ -14,6 +14,7 @@
 #define LLVM_EXECUTIONENGINE_ORC_MEMORYMAPPER_H
 
 #include "llvm/ExecutionEngine/Orc/Core.h"
+#include "llvm/Support/Process.h"
 
 #include <mutex>
 
@@ -41,6 +42,9 @@ class MemoryMapper {
 
   using OnReservedFunction = unique_function<void(Expected<ExecutorAddrRange>)>;
 
+  // Page size of the target process
+  virtual unsigned int getPageSize() = 0;
+
   /// Reserves address space in executor process
   virtual void reserve(size_t NumBytes, OnReservedFunction OnReserved) = 0;
 
@@ -76,7 +80,11 @@ class MemoryMapper {
 
 class InProcessMemoryMapper final : public MemoryMapper {
 public:
-  InProcessMemoryMapper() {}
+  InProcessMemoryMapper(size_t PageSize);
+
+  static Expected<std::unique_ptr<InProcessMemoryMapper>> Create();
+
+  unsigned int getPageSize() override { return PageSize; }
 
   void reserve(size_t NumBytes, OnReservedFunction OnReserved) override;
 
@@ -107,6 +115,8 @@ class InProcessMemoryMapper final : public MemoryMapper {
   std::mutex Mutex;
   ReservationMap Reservations;
   AllocationMap Allocations;
+
+  size_t PageSize;
 };
 
 class SharedMemoryMapper final : public MemoryMapper {
@@ -119,8 +129,13 @@ class SharedMemoryMapper final : public MemoryMapper {
     ExecutorAddr Release;
   };
 
-  SharedMemoryMapper(ExecutorProcessControl &EPC, SymbolAddrs SAs)
-      : EPC(EPC), SAs(SAs) {}
+  SharedMemoryMapper(ExecutorProcessControl &EPC, SymbolAddrs SAs,
+                     size_t PageSize);
+
+  static Expected<std::unique_ptr<SharedMemoryMapper>>
+  Create(ExecutorProcessControl &EPC, SymbolAddrs SAs);
+
+  unsigned int getPageSize() override { return PageSize; }
 
   void reserve(size_t NumBytes, OnReservedFunction OnReserved) override;
 
@@ -148,6 +163,8 @@ class SharedMemoryMapper final : public MemoryMapper {
   std::mutex Mutex;
 
   std::map<ExecutorAddr, Reservation> Reservations;
+
+  size_t PageSize;
 };
 
 } // namespace orc

llvm/lib/ExecutionEngine/Orc/CMakeLists.txt

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

llvm/lib/ExecutionEngine/Orc/MapperJITLinkMemoryManager.cpp

@@ -0,0 +1,135 @@
+//=== MapperJITLinkMemoryManager.cpp - Memory management with MemoryMapper ===//
+//
+// 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/MapperJITLinkMemoryManager.h"
+
+#include "llvm/ExecutionEngine/JITLink/JITLink.h"
+#include "llvm/Support/Process.h"
+
+#include <limits>
+
+using namespace llvm::jitlink;
+
+namespace llvm {
+namespace orc {
+
+class MapperJITLinkMemoryManager::InFlightAlloc
+    : public JITLinkMemoryManager::InFlightAlloc {
+public:
+  InFlightAlloc(MapperJITLinkMemoryManager &Parent, LinkGraph &G,
+                ExecutorAddr AllocAddr,
+                std::vector<MemoryMapper::AllocInfo::SegInfo> Segs)
+      : Parent(Parent), G(G), AllocAddr(AllocAddr), Segs(std::move(Segs)) {}
+
+  void finalize(OnFinalizedFunction OnFinalize) override {
+    MemoryMapper::AllocInfo AI;
+    AI.MappingBase = AllocAddr;
+
+    std::swap(AI.Segments, Segs);
+    std::swap(AI.Actions, G.allocActions());
+
+    Parent.Mapper->initialize(AI, [&](Expected<ExecutorAddr> Result) {
+      if (!Result) {
+        OnFinalize(Result.takeError());
+        return;
+      }
+
+      OnFinalize(FinalizedAlloc(*Result));
+    });
+  }
+
+  void abandon(OnAbandonedFunction OnFinalize) override {
+    Parent.Mapper->release({AllocAddr}, std::move(OnFinalize));
+  }
+
+private:
+  MapperJITLinkMemoryManager &Parent;
+  LinkGraph &G;
+  ExecutorAddr AllocAddr;
+  std::vector<MemoryMapper::AllocInfo::SegInfo> Segs;
+};
+
+MapperJITLinkMemoryManager::MapperJITLinkMemoryManager(
+    std::unique_ptr<MemoryMapper> Mapper)
+    : Mapper(std::move(Mapper)) {}
+
+void MapperJITLinkMemoryManager::allocate(const JITLinkDylib *JD, LinkGraph &G,
+                                          OnAllocatedFunction OnAllocated) {
+  BasicLayout BL(G);
+
+  // find required address space
+  auto SegsSizes = BL.getContiguousPageBasedLayoutSizes(Mapper->getPageSize());
+  if (!SegsSizes) {
+    OnAllocated(SegsSizes.takeError());
+    return;
+  }
+
+  // Check if total size fits in address space
+  if (SegsSizes->total() > std::numeric_limits<size_t>::max()) {
+    OnAllocated(make_error<JITLinkError>(
+        formatv("Total requested size {:x} for graph {} exceeds address space",
+                SegsSizes->total(), G.getName())));
+    return;
+  }
+
+  Mapper->reserve(
+      SegsSizes->total(),
+      [this, &G, BL = std::move(BL), OnAllocated = std::move(OnAllocated)](
+          Expected<ExecutorAddrRange> Result) mutable {
+        if (!Result) {
+          return OnAllocated(Result.takeError());
+        }
+
+        auto NextSegAddr = Result->Start;
+
+        std::vector<MemoryMapper::AllocInfo::SegInfo> SegInfos;
+
+        for (auto &KV : BL.segments()) {
+          auto &AG = KV.first;
+          auto &Seg = KV.second;
+
+          auto TotalSize = Seg.ContentSize + Seg.ZeroFillSize;
+
+          Seg.Addr = NextSegAddr;
+          Seg.WorkingMem = Mapper->prepare(NextSegAddr, TotalSize);
+
+          NextSegAddr += alignTo(TotalSize, Mapper->getPageSize());
+
+          MemoryMapper::AllocInfo::SegInfo SI;
+          SI.Offset = Seg.Addr - Result->Start;
+          SI.ContentSize = Seg.ContentSize;
+          SI.ZeroFillSize = Seg.ZeroFillSize;
+          SI.Prot = (toSysMemoryProtectionFlags(AG.getMemProt()));
+          SI.WorkingMem = Seg.WorkingMem;
+
+          SegInfos.push_back(SI);
+        }
+
+        if (auto Err = BL.apply()) {
+          OnAllocated(std::move(Err));
+          return;
+        }
+
+        OnAllocated(std::make_unique<InFlightAlloc>(*this, G, Result->Start,
+                                                    std::move(SegInfos)));
+      });
+}
+
+void MapperJITLinkMemoryManager::deallocate(
+    std::vector<FinalizedAlloc> Allocs, OnDeallocatedFunction OnDeallocated) {
+  std::vector<ExecutorAddr> Bases;
+  Bases.reserve(Allocs.size());
+  for (auto &FA : Allocs) {
+    Bases.push_back(FA.getAddress());
+    FA.release();
+  }
+  Mapper->release(Bases, std::move(OnDeallocated));
+}
+
+} // end namespace orc
+} // end namespace llvm

llvm/lib/ExecutionEngine/Orc/MemoryMapper.cpp

@@ -24,6 +24,17 @@ namespace orc {
 
 MemoryMapper::~MemoryMapper() {}
 
+InProcessMemoryMapper::InProcessMemoryMapper(size_t PageSize)
+    : PageSize(PageSize) {}
+
+Expected<std::unique_ptr<InProcessMemoryMapper>>
+InProcessMemoryMapper::Create() {
+  auto PageSize = sys::Process::getPageSize();
+  if (!PageSize)
+    return PageSize.takeError();
+  return std::make_unique<InProcessMemoryMapper>(*PageSize);
+}
+
 void InProcessMemoryMapper::reserve(size_t NumBytes,
                                     OnReservedFunction OnReserved) {
   std::error_code EC;
@@ -160,6 +171,19 @@ InProcessMemoryMapper::~InProcessMemoryMapper() {
 
 // SharedMemoryMapper
 
+SharedMemoryMapper::SharedMemoryMapper(ExecutorProcessControl &EPC,
+                                       SymbolAddrs SAs, size_t PageSize)
+    : EPC(EPC), SAs(SAs), PageSize(PageSize) {}
+
+Expected<std::unique_ptr<SharedMemoryMapper>>
+SharedMemoryMapper::Create(ExecutorProcessControl &EPC, SymbolAddrs SAs) {
+  auto PageSize = sys::Process::getPageSize();
+  if (!PageSize)
+    return PageSize.takeError();
+
+  return std::make_unique<SharedMemoryMapper>(EPC, SAs, *PageSize);
+}
+
 void SharedMemoryMapper::reserve(size_t NumBytes,
                                  OnReservedFunction OnReserved) {
 #if defined(LLVM_ON_UNIX) || defined(_WIN32)

llvm/unittests/ExecutionEngine/Orc/CMakeLists.txt

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

llvm/unittests/ExecutionEngine/Orc/MapperJITLinkMemoryManagerTest.cpp

@@ -0,0 +1,48 @@
+//===-------------- MapperJITLinkMemoryManagerTest.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 "OrcTestCommon.h"
+
+#include "llvm/ExecutionEngine/Orc/MapperJITLinkMemoryManager.h"
+
+#include "llvm/Testing/Support/Error.h"
+
+#include <vector>
+
+using namespace llvm;
+using namespace llvm::jitlink;
+using namespace llvm::orc;
+using namespace llvm::orc::shared;
+
+namespace {
+
+TEST(MapperJITLinkMemoryManagerTest, InProcess) {
+  auto MemMgr = cantFail(
+      MapperJITLinkMemoryManager::CreateWithMapper<InProcessMemoryMapper>());
+
+  StringRef Hello = "hello";
+  auto SSA = jitlink::SimpleSegmentAlloc::Create(
+      *MemMgr, nullptr, {{jitlink::MemProt::Read, {Hello.size(), Align(1)}}});
+  EXPECT_THAT_EXPECTED(SSA, Succeeded());
+  auto SegInfo = SSA->getSegInfo(jitlink::MemProt::Read);
+  memcpy(SegInfo.WorkingMem.data(), Hello.data(), Hello.size());
+
+  auto FA = SSA->finalize();
+  EXPECT_THAT_EXPECTED(FA, Succeeded());
+
+  ExecutorAddr TargetAddr(SegInfo.Addr);
+
+  const char *TargetMem = TargetAddr.toPtr<const char *>();
+  StringRef TargetHello(TargetMem, Hello.size());
+  EXPECT_EQ(Hello, TargetHello);
+
+  auto Err2 = MemMgr->deallocate(std::move(*FA));
+  EXPECT_THAT_ERROR(std::move(Err2), Succeeded());
+}
+
+} // namespace

llvm/unittests/ExecutionEngine/Orc/MemoryMapperTest.cpp

@@ -66,10 +66,10 @@ TEST(MemoryMapperTest, InitializeDeinitialize) {
   int DeinitializeCounter = 0;
   {
     std::unique_ptr<MemoryMapper> Mapper =
-        std::make_unique<InProcessMemoryMapper>();
+        cantFail(InProcessMemoryMapper::Create());
 
     // We will do two separate allocations
-    auto PageSize = cantFail(sys::Process::getPageSize());
+    auto PageSize = Mapper->getPageSize();
     auto TotalSize = PageSize * 2;
 
     // Reserve address space

llvm/unittests/ExecutionEngine/Orc/SharedMemoryMapperTest.cpp

@@ -62,11 +62,11 @@ TEST(SharedMemoryMapperTest, MemReserveInitializeDeinitializeRelease) {
   auto F = P.get_future();
 
   {
-    auto PageSize = cantFail(sys::Process::getPageSize());
-    size_t ReqSize = PageSize;
-
     std::unique_ptr<MemoryMapper> Mapper =
-        std::make_unique<SharedMemoryMapper>(*SelfEPC, SAs);
+        cantFail(SharedMemoryMapper::Create(*SelfEPC, SAs));
+
+    auto PageSize = Mapper->getPageSize();
+    size_t ReqSize = PageSize;
 
     Mapper->reserve(ReqSize, [&](Expected<ExecutorAddrRange> Result) {
       EXPECT_THAT_ERROR(Result.takeError(), Succeeded());