[Orc][JITLink] Slab based memory allocator to reduce RPC calls
authorAnubhab Ghosh <anubhabghosh.me@gmail.com>
Fri, 22 Jul 2022 19:14:03 +0000 (00:44 +0530)
committerAnubhab Ghosh <anubhabghosh.me@gmail.com>
Mon, 8 Aug 2022 09:43:41 +0000 (15:13 +0530)
Calling reserve() used to require an RPC call. This commit allows large
ranges of executor address space to be reserved. Subsequent calls to
reserve() will return subranges of already reserved address space while
there is still space available.

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

llvm/include/llvm/ExecutionEngine/Orc/MapperJITLinkMemoryManager.h
llvm/lib/ExecutionEngine/Orc/MapperJITLinkMemoryManager.cpp
llvm/unittests/ExecutionEngine/Orc/MapperJITLinkMemoryManagerTest.cpp

index 37d75bfff546fa26a92420401d7f0b030f3c10bc..4017b120790da9d78c2b9e807a5b04ee0ebbed4e 100644 (file)
@@ -14,7 +14,6 @@
 #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 {
@@ -22,16 +21,18 @@ namespace orc {
 
 class MapperJITLinkMemoryManager : public jitlink::JITLinkMemoryManager {
 public:
-  MapperJITLinkMemoryManager(std::unique_ptr<MemoryMapper> Mapper);
+  MapperJITLinkMemoryManager(size_t ReservationGranularity,
+                             std::unique_ptr<MemoryMapper> Mapper);
 
   template <class MemoryMapperType, class... Args>
   static Expected<std::unique_ptr<MapperJITLinkMemoryManager>>
-  CreateWithMapper(Args &&...A) {
+  CreateWithMapper(size_t ReservationGranularity, Args &&...A) {
     auto Mapper = MemoryMapperType::Create(std::forward<Args>(A)...);
     if (!Mapper)
       return Mapper.takeError();
 
-    return std::make_unique<MapperJITLinkMemoryManager>(std::move(*Mapper));
+    return std::make_unique<MapperJITLinkMemoryManager>(ReservationGranularity,
+                                                        std::move(*Mapper));
   }
 
   void allocate(const jitlink::JITLinkDylib *JD, jitlink::LinkGraph &G,
@@ -47,6 +48,15 @@ public:
 private:
   class InFlightAlloc;
 
+  std::mutex Mutex;
+
+  // We reserve multiples of this from the executor address space
+  size_t ReservationUnits;
+  // Ranges that have been reserved in executor but not yet allocated
+  std::vector<ExecutorAddrRange> AvailableMemory;
+  // Ranges that have been reserved in executor and already allocated
+  DenseMap<ExecutorAddr, ExecutorAddrDiff> UsedMemory;
+
   std::unique_ptr<MemoryMapper> Mapper;
 };
 
index c2e7baabb994d2ceca608cc05d38b15cdf40f149..b9328f0a16edf557a96692a530b9abde3b75f2cb 100644 (file)
@@ -8,6 +8,7 @@
 
 #include "llvm/ExecutionEngine/Orc/MapperJITLinkMemoryManager.h"
 
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/ExecutionEngine/JITLink/JITLink.h"
 #include "llvm/Support/Process.h"
 
@@ -55,8 +56,8 @@ private:
 };
 
 MapperJITLinkMemoryManager::MapperJITLinkMemoryManager(
-    std::unique_ptr<MemoryMapper> Mapper)
-    : Mapper(std::move(Mapper)) {}
+    size_t ReservationGranularity, std::unique_ptr<MemoryMapper> Mapper)
+    : ReservationUnits(ReservationGranularity), Mapper(std::move(Mapper)) {}
 
 void MapperJITLinkMemoryManager::allocate(const JITLinkDylib *JD, LinkGraph &G,
                                           OnAllocatedFunction OnAllocated) {
@@ -69,66 +70,94 @@ void MapperJITLinkMemoryManager::allocate(const JITLinkDylib *JD, LinkGraph &G,
     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;
+  auto TotalSize = SegsSizes->total();
+
+  Mutex.lock();
+
+  // find an already reserved range that is large enough
+  ExecutorAddrRange SelectedRange{};
+  std::vector<ExecutorAddrRange>::iterator SelectedRangeIt;
+  SelectedRangeIt =
+      llvm::find_if(AvailableMemory, [TotalSize](ExecutorAddrRange Range) {
+        return TotalSize < Range.size();
+      });
+  if (SelectedRangeIt != AvailableMemory.end()) {
+    SelectedRange = *SelectedRangeIt;
+    AvailableMemory.erase(SelectedRangeIt);
   }
 
-  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 CompleteAllocation = [this, &G, BL = std::move(BL),
+                             OnAllocated = std::move(OnAllocated)](
+                                Expected<ExecutorAddrRange> Result) mutable {
+    if (!Result) {
+      Mutex.unlock();
+      return OnAllocated(Result.takeError());
+    }
 
-        auto NextSegAddr = Result->Start;
+    auto NextSegAddr = Result->Start;
 
-        std::vector<MemoryMapper::AllocInfo::SegInfo> SegInfos;
+    std::vector<MemoryMapper::AllocInfo::SegInfo> SegInfos;
 
-        for (auto &KV : BL.segments()) {
-          auto &AG = KV.first;
-          auto &Seg = KV.second;
+    for (auto &KV : BL.segments()) {
+      auto &AG = KV.first;
+      auto &Seg = KV.second;
 
-          auto TotalSize = Seg.ContentSize + Seg.ZeroFillSize;
+      auto TotalSize = Seg.ContentSize + Seg.ZeroFillSize;
 
-          Seg.Addr = NextSegAddr;
-          Seg.WorkingMem = Mapper->prepare(NextSegAddr, TotalSize);
+      Seg.Addr = NextSegAddr;
+      Seg.WorkingMem = Mapper->prepare(NextSegAddr, TotalSize);
 
-          NextSegAddr += alignTo(TotalSize, Mapper->getPageSize());
+      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;
+      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);
-        }
+      SegInfos.push_back(SI);
+    }
 
-        if (auto Err = BL.apply()) {
-          OnAllocated(std::move(Err));
-          return;
-        }
+    UsedMemory.insert({Result->Start, NextSegAddr - Result->Start});
 
-        OnAllocated(std::make_unique<InFlightAlloc>(*this, G, Result->Start,
-                                                    std::move(SegInfos)));
-      });
+    if (NextSegAddr < Result->End) {
+      // Save the remaining memory for reuse in next allocation(s)
+      auto RemainingRange = ExecutorAddrRange(NextSegAddr, Result->End);
+      AvailableMemory.push_back(RemainingRange);
+    }
+    Mutex.unlock();
+
+    if (auto Err = BL.apply()) {
+      OnAllocated(std::move(Err));
+      return;
+    }
+
+    OnAllocated(std::make_unique<InFlightAlloc>(*this, G, Result->Start,
+                                                std::move(SegInfos)));
+  };
+
+  if (SelectedRange.empty()) { // no already reserved range was found
+    auto TotalAllocation = alignTo(TotalSize, ReservationUnits);
+    Mapper->reserve(TotalAllocation, std::move(CompleteAllocation));
+  } else {
+    CompleteAllocation(SelectedRange);
+  }
 }
 
 void MapperJITLinkMemoryManager::deallocate(
     std::vector<FinalizedAlloc> Allocs, OnDeallocatedFunction OnDeallocated) {
-  std::vector<ExecutorAddr> Bases;
-  Bases.reserve(Allocs.size());
+  std::lock_guard<std::mutex> Lock(Mutex);
+
   for (auto &FA : Allocs) {
-    Bases.push_back(FA.getAddress());
+    ExecutorAddr Addr = FA.getAddress();
+    ExecutorAddrDiff Size = UsedMemory[Addr];
+    UsedMemory.erase(Addr);
+
+    AvailableMemory.push_back({Addr, Addr + Size});
     FA.release();
   }
-  Mapper->release(Bases, std::move(OnDeallocated));
+  OnDeallocated(Error::success());
 }
 
 } // end namespace orc
index e6dc303b6b296dfaaf1d59b1902af020e018f29c..74d83877013591174d3b81e0ba6cb5b2e04be64c 100644 (file)
@@ -10,6 +10,7 @@
 
 #include "llvm/ExecutionEngine/Orc/MapperJITLinkMemoryManager.h"
 
+#include "llvm/ExecutionEngine/Orc/MemoryMapper.h"
 #include "llvm/Testing/Support/Error.h"
 
 #include <vector>
@@ -21,28 +22,107 @@ using namespace llvm::orc::shared;
 
 namespace {
 
+class CounterMapper final : public MemoryMapper {
+public:
+  CounterMapper(std::unique_ptr<MemoryMapper> Mapper)
+      : Mapper(std::move(Mapper)) {}
+
+  unsigned int getPageSize() override { return Mapper->getPageSize(); }
+
+  void reserve(size_t NumBytes, OnReservedFunction OnReserved) override {
+    ++ReserveCount;
+    return Mapper->reserve(NumBytes, std::move(OnReserved));
+  }
+
+  void initialize(AllocInfo &AI, OnInitializedFunction OnInitialized) override {
+    ++InitCount;
+    return Mapper->initialize(AI, std::move(OnInitialized));
+  }
+
+  char *prepare(ExecutorAddr Addr, size_t ContentSize) override {
+    return Mapper->prepare(Addr, ContentSize);
+  }
+
+  void deinitialize(ArrayRef<ExecutorAddr> Allocations,
+                    OnDeinitializedFunction OnDeInitialized) override {
+    ++DeinitCount;
+    return Mapper->deinitialize(Allocations, std::move(OnDeInitialized));
+  }
+
+  void release(ArrayRef<ExecutorAddr> Reservations,
+               OnReleasedFunction OnRelease) override {
+    ++ReleaseCount;
+
+    return Mapper->release(Reservations, std::move(OnRelease));
+  }
+
+  int ReserveCount = 0, InitCount = 0, DeinitCount = 0, ReleaseCount = 0;
+
+private:
+  std::unique_ptr<MemoryMapper> Mapper;
+};
+
 TEST(MapperJITLinkMemoryManagerTest, InProcess) {
-  auto MemMgr = cantFail(
-      MapperJITLinkMemoryManager::CreateWithMapper<InProcessMemoryMapper>());
+  auto Mapper = std::make_unique<CounterMapper>(
+      cantFail(InProcessMemoryMapper::Create()));
+
+  auto *Counter = static_cast<CounterMapper *>(Mapper.get());
+
+  auto MemMgr = std::make_unique<MapperJITLinkMemoryManager>(16 * 1024 * 1024,
+                                                             std::move(Mapper));
+
+  EXPECT_EQ(Counter->ReserveCount, 0);
+  EXPECT_EQ(Counter->InitCount, 0);
 
   StringRef Hello = "hello";
-  auto SSA = jitlink::SimpleSegmentAlloc::Create(
+  auto SSA1 = 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());
+  EXPECT_THAT_EXPECTED(SSA1, Succeeded());
+
+  EXPECT_EQ(Counter->ReserveCount, 1);
+  EXPECT_EQ(Counter->InitCount, 0);
+
+  auto SegInfo1 = SSA1->getSegInfo(jitlink::MemProt::Read);
+  memcpy(SegInfo1.WorkingMem.data(), Hello.data(), Hello.size());
 
-  auto FA = SSA->finalize();
-  EXPECT_THAT_EXPECTED(FA, Succeeded());
+  auto FA1 = SSA1->finalize();
+  EXPECT_THAT_EXPECTED(FA1, Succeeded());
+
+  EXPECT_EQ(Counter->ReserveCount, 1);
+  EXPECT_EQ(Counter->InitCount, 1);
+
+  auto SSA2 = jitlink::SimpleSegmentAlloc::Create(
+      *MemMgr, nullptr, {{jitlink::MemProt::Read, {Hello.size(), Align(1)}}});
+  EXPECT_THAT_EXPECTED(SSA2, Succeeded());
 
-  ExecutorAddr TargetAddr(SegInfo.Addr);
+  // last reservation should be reused
+  EXPECT_EQ(Counter->ReserveCount, 1);
+  EXPECT_EQ(Counter->InitCount, 1);
 
-  const char *TargetMem = TargetAddr.toPtr<const char *>();
-  StringRef TargetHello(TargetMem, Hello.size());
-  EXPECT_EQ(Hello, TargetHello);
+  auto SegInfo2 = SSA2->getSegInfo(jitlink::MemProt::Read);
+  memcpy(SegInfo2.WorkingMem.data(), Hello.data(), Hello.size());
+  auto FA2 = SSA2->finalize();
+  EXPECT_THAT_EXPECTED(FA2, Succeeded());
 
-  auto Err2 = MemMgr->deallocate(std::move(*FA));
+  EXPECT_EQ(Counter->ReserveCount, 1);
+  EXPECT_EQ(Counter->InitCount, 2);
+
+  ExecutorAddr TargetAddr1(SegInfo1.Addr);
+  ExecutorAddr TargetAddr2(SegInfo2.Addr);
+
+  const char *TargetMem1 = TargetAddr1.toPtr<const char *>();
+  StringRef TargetHello1(TargetMem1, Hello.size());
+  EXPECT_EQ(Hello, TargetHello1);
+
+  const char *TargetMem2 = TargetAddr2.toPtr<const char *>();
+  StringRef TargetHello2(TargetMem2, Hello.size());
+  EXPECT_EQ(Hello, TargetHello2);
+
+  auto Err2 = MemMgr->deallocate(std::move(*FA1));
   EXPECT_THAT_ERROR(std::move(Err2), Succeeded());
+
+  auto Err3 = MemMgr->deallocate(std::move(*FA2));
+  EXPECT_THAT_ERROR(std::move(Err3), Succeeded());
 }
 
 } // namespace