[Orc][JITLink] JITLinkMemoryManager implementation using MemoryMapper

[argentite]

Summary

MapperJITLinkMemoryManager supports executor memory management using any
implementation of MemoryMapper to do the transfer such as in process or
shared memory based.

[lhames]

Otherwise LGTM. Nice work!

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

46

SharedMemoryMapper will need a definition of this now too.

llvm/lib/ExecutionEngine/Orc/MemoryMapper.cpp

31

getPageSize can fail -- this needs to return an error. The usual idiom is a named constructor:

class InProcessMemoryMapper ... {
public:
  Expected<InProcessMemoryMapper> Create() {
    if (auto PageSize = sys::Process::getPageSize())
      return InProcessMemoryMapper(*PageSize);
    else
      return PageSize.takeError();
  }

  InProcessMemoryMapper(size_t PageSize) : PageSize(PageSize) {}

  ...
}

[argentite]

Added support for new SharedMemoryMapper

Improved error handling with create() methods

[lhames]

Otherwise LGTM.

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

85

Create methods _usually_ buck the method naming convention and use an uppercase 'C' to emphasize that they're named constructors (though honestly there's plenty of instances of both styles in the codebase). Since ORC more consistently uses upper-case 'C', we should stick with that for now.

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

+//===--------------- 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

 //===- 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 "llvm/Support/Process.h"
 
 #include <mutex>
 
 namespace llvm {
 namespace orc {
 
 /// Manages mapping, content transfer and protections for JIT memory
 class MemoryMapper {
struct AllocInfo {
 
     ExecutorAddr MappingBase;
     std::vector<SegInfo> Segments;
     shared::AllocActions Actions;
   };
 
   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;
 
   /// Provides working memory
   virtual char *prepare(ExecutorAddr Addr, size_t ContentSize) = 0;
 
   using OnInitializedFunction = unique_function<void(Expected<ExecutorAddr>)>;
 
public:
   virtual void release(ArrayRef<ExecutorAddr> Reservations,
                        OnReleasedFunction OnRelease) = 0;
 
   virtual ~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;
 
   void initialize(AllocInfo &AI, OnInitializedFunction OnInitialized) override;
 
   char *prepare(ExecutorAddr Addr, size_t ContentSize) override;
 
   void deinitialize(ArrayRef<ExecutorAddr> Allocations,
struct Reservation {
     size_t Size;
     std::vector<ExecutorAddr> Allocations;
   };
   using ReservationMap = DenseMap<void *, Reservation>;
 
   std::mutex Mutex;
   ReservationMap Reservations;
   AllocationMap Allocations;
+
+  size_t PageSize;
 };
 
 class SharedMemoryMapper final : public MemoryMapper {
 public:
   struct SymbolAddrs {
     ExecutorAddr Instance;
     ExecutorAddr Reserve;
     ExecutorAddr Initialize;
     ExecutorAddr Deinitialize;
     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;
 
   char *prepare(ExecutorAddr Addr, size_t ContentSize) override;
 
   void initialize(AllocInfo &AI, OnInitializedFunction OnInitialized) override;
 
   void deinitialize(ArrayRef<ExecutorAddr> Allocations,
private:
   };
 
   ExecutorProcessControl &EPC;
   SymbolAddrs SAs;
 
   std::mutex Mutex;
 
   std::map<ExecutorAddr, Reservation> Reservations;
+
+  size_t PageSize;
 };
 
 } // namespace orc
 } // end namespace llvm
 
 #endif // LLVM_EXECUTIONENGINE_ORC_MEMORYMAPPER_H

llvm/lib/ExecutionEngine/Orc/CMakeLists.txt

add_llvm_component_library(LLVMOrcJIT
   IRCompileLayer.cpp
   IRTransformLayer.cpp
   JITTargetMachineBuilder.cpp
   LazyReexports.cpp
   Layer.cpp
   LookupAndRecordAddrs.cpp
   LLJIT.cpp
   MachOPlatform.cpp
+  MapperJITLinkMemoryManager.cpp
   MemoryMapper.cpp
   ELFNixPlatform.cpp
   Mangling.cpp
   ObjectLinkingLayer.cpp
   ObjectTransformLayer.cpp
   OrcABISupport.cpp
   OrcV2CBindings.cpp
   RTDyldObjectLinkingLayer.cpp

llvm/lib/ExecutionEngine/Orc/MapperJITLinkMemoryManager.cpp

+//=== 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

 #include <windows.h>
 #endif
 
 namespace llvm {
 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;
   auto MB = sys::Memory::allocateMappedMemory(
       NumBytes, nullptr, sys::Memory::MF_READ | sys::Memory::MF_WRITE, EC);
 
   if (EC)
     return OnReserved(errorCodeToError(EC));
InProcessMemoryMapper::~InProcessMemoryMapper() {
   std::promise<MSVCPError> P;
   auto F = P.get_future();
   release(ReservationAddrs, [&](Error Err) { P.set_value(std::move(Err)); });
   cantFail(F.get());
 }
 
 // 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)
 
   EPC.callSPSWrapperAsync<
       rt::SPSExecutorSharedMemoryMapperServiceReserveSignature>(
       SAs.Reserve,
       [this, NumBytes, OnReserved = std::move(OnReserved)](

llvm/unittests/ExecutionEngine/Orc/CMakeLists.txt

add_llvm_unittest(OrcJITTests
   ExecutorAddressTest.cpp
   ExecutionSessionWrapperFunctionCallsTest.cpp
   EPCGenericJITLinkMemoryManagerTest.cpp
   EPCGenericMemoryAccessTest.cpp
   IndirectionUtilsTest.cpp
   JITTargetMachineBuilderTest.cpp
   LazyCallThroughAndReexportsTest.cpp
   LookupAndRecordAddrsTest.cpp
+  MapperJITLinkMemoryManagerTest.cpp
   MemoryMapperTest.cpp
   ObjectLinkingLayerTest.cpp
   OrcCAPITest.cpp
   OrcTestCommon.cpp
   ResourceTrackerTest.cpp
   RTDyldObjectLinkingLayerTest.cpp
   SharedMemoryMapperTest.cpp
   SimpleExecutorMemoryManagerTest.cpp

llvm/unittests/ExecutionEngine/Orc/MapperJITLinkMemoryManagerTest.cpp

+//===-------------- 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

 
 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>();
+        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
     auto Mem1 = reserve(*Mapper, TotalSize);
     EXPECT_THAT_ERROR(Mem1.takeError(), Succeeded());
 
     // Test string for memory transfer
     std::string HW = "Hello, world!";

llvm/unittests/ExecutionEngine/Orc/SharedMemoryMapperTest.cpp

TEST(SharedMemoryMapperTest, MemReserveInitializeDeinitializeRelease) {
 
   std::string TestString = "Hello, World!";
 
   // barrier
   std::promise<void> P;
   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());
       auto Reservation = std::move(*Result);
       {
         char *Addr = Mapper->prepare(Reservation.Start, TestString.size() + 1);
         std::strcpy(Addr, TestString.c_str());
       }