[Orc][JITLink] Slab based memory allocator to reduce RPC calls

[argentite]

Summary

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.

[lhames]

llvm/lib/ExecutionEngine/Orc/MapperJITLinkMemoryManager.cpp

62

It looks like the algorithm here is:

1. lock
2. search for range
3. unlock
4. check if range found
5. if so then CompleteAllocation, otherwise reserve and CompleteAllocation in the callback.

The problem is that if we start off with two concurrent allocate calls we might allocate two different slabs, which we don't want to do. I think the allocator should do a blocking call to reserve under the lock instead. That'll hit the first caller hard, but it guarantees that we don't allocate more than one slab until/unless the first one fills up.

(We could also make this configurable too -- we want the slab allocator to be able to guarantee memory range constraints to enable use of arbitrary small code model code, but it's also useful as a performance optimization, so two arguments "slab size" and "allow concurrent reservations" would make sense. I think that should be a follow-up patch though).

[argentite]

When enough memory is not already reserved, lock the mutex while doing the RPC to prevent subsequent allocations to request new memory ranges

[argentite]

Updated title and commit message

[lhames]

LGTM!

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

55–58

// Ranges that have been reserved on both side
DenseMap<ExecutorAddr, ExecutorAddrDiff> UsedMemory;

I didn't understand this comment? This is just the used memory list, right? Whether or not there's any memory associated with it on the controller side should be up to the mapper, rather than the allocator.

llvm/lib/ExecutionEngine/Orc/MapperJITLinkMemoryManager.cpp

72–87

I think this check should probably be dropped: We support JITing on a 32-bit controller for a 64-bit target, so you could request, e.g., an 8Gb zero-fill data section. That would overflow a size_t on the controller, but is still perfectly legal.

On the other hand if you overflow the address scape for the executor then you'll definitely overflow the available space too, so you'll error out below anyway.

157

There's no coalescing of the AvailableMemory list yet, so you'll get fragmentation in the allocator. I think that should be addressed in a follow-up patch though.

[argentite]

Remove address space size check and improve comments

[argentite]

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

55–58

I have tried to make the language clearer.

llvm/lib/ExecutionEngine/Orc/MapperJITLinkMemoryManager.cpp

157

I plan to do a more efficient implementation as a follow up.

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);
+  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,
                 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::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;
 };
 
 } // end namespace orc
 } // end namespace llvm
 
 #endif // LLVM_EXECUTIONENGINE_ORC_MAPPERJITLINKMEMORYMANAGER_H

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/ADT/STLExtras.h"
 #include "llvm/ExecutionEngine/JITLink/JITLink.h"
 #include "llvm/Support/Process.h"
 
 #include <limits>
 
 using namespace llvm::jitlink;
 
 namespace llvm {

 private:
   MapperJITLinkMemoryManager &Parent;
   LinkGraph &G;
   ExecutorAddr AllocAddr;
   std::vector<MemoryMapper::AllocInfo::SegInfo> Segs;
 };
 
 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) {
   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;
+  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)](
+  auto CompleteAllocation = [this, &G, BL = std::move(BL),
+                             OnAllocated = std::move(OnAllocated)](
                                 Expected<ExecutorAddrRange> Result) mutable {
-        if (!Result) {
-          return OnAllocated(Result.takeError());
-        }
+      Mutex.unlock();
 
-        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;
+      SI.Prot = toSysMemoryProtectionFlags(AG.getMemProt());
 
-          SegInfos.push_back(SI);
-        }
 
-        if (auto Err = BL.apply()) {
-          OnAllocated(std::move(Err));
-          return;
-        }
+    UsedMemory.insert({Result->Start, NextSegAddr - Result->Start});
+
+    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();
+
 
-        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
 } // end namespace llvm

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/ExecutionEngine/Orc/MemoryMapper.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 {
 
+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 FA = SSA->finalize();
-  EXPECT_THAT_EXPECTED(FA, Succeeded());
+  auto SegInfo1 = SSA1->getSegInfo(jitlink::MemProt::Read);
+  memcpy(SegInfo1.WorkingMem.data(), Hello.data(), Hello.size());
 
-  ExecutorAddr TargetAddr(SegInfo.Addr);
+  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());
 
-  const char *TargetMem = TargetAddr.toPtr<const char *>();
-  StringRef TargetHello(TargetMem, Hello.size());
-  EXPECT_EQ(Hello, TargetHello);
+  // last reservation should be reused
+  EXPECT_EQ(Counter->ReserveCount, 1);
+  EXPECT_EQ(Counter->InitCount, 1);
 
-  auto Err2 = MemMgr->deallocate(std::move(*FA));
+  auto SegInfo2 = SSA2->getSegInfo(jitlink::MemProt::Read);
+  memcpy(SegInfo2.WorkingMem.data(), Hello.data(), Hello.size());
+  auto FA2 = SSA2->finalize();
+  EXPECT_THAT_EXPECTED(FA2, Succeeded());
+
+  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