[llvm] 1eee6de - [Orc][JITLink] Slab based memory allocator to reduce RPC calls
Anubhab Ghosh via llvm-commits
llvm-commits at lists.llvm.org
Mon Aug 8 02:43:56 PDT 2022
Author: Anubhab Ghosh
Date: 2022-08-08T15:13:41+05:30
New Revision: 1eee6de873974f55538df976bf7802f019eac70a
URL: https://github.com/llvm/llvm-project/commit/1eee6de873974f55538df976bf7802f019eac70a
DIFF: https://github.com/llvm/llvm-project/commit/1eee6de873974f55538df976bf7802f019eac70a.diff
LOG: [Orc][JITLink] Slab based memory allocator to reduce RPC calls
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
Added:
Modified:
llvm/include/llvm/ExecutionEngine/Orc/MapperJITLinkMemoryManager.h
llvm/lib/ExecutionEngine/Orc/MapperJITLinkMemoryManager.cpp
llvm/unittests/ExecutionEngine/Orc/MapperJITLinkMemoryManagerTest.cpp
Removed:
################################################################################
diff --git a/llvm/include/llvm/ExecutionEngine/Orc/MapperJITLinkMemoryManager.h b/llvm/include/llvm/ExecutionEngine/Orc/MapperJITLinkMemoryManager.h
index 37d75bfff546f..4017b120790da 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/MapperJITLinkMemoryManager.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/MapperJITLinkMemoryManager.h
@@ -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 @@ class MapperJITLinkMemoryManager : public jitlink::JITLinkMemoryManager {
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;
};
diff --git a/llvm/lib/ExecutionEngine/Orc/MapperJITLinkMemoryManager.cpp b/llvm/lib/ExecutionEngine/Orc/MapperJITLinkMemoryManager.cpp
index c2e7baabb994d..b9328f0a16edf 100644
--- a/llvm/lib/ExecutionEngine/Orc/MapperJITLinkMemoryManager.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/MapperJITLinkMemoryManager.cpp
@@ -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 @@ class MapperJITLinkMemoryManager::InFlightAlloc
};
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
diff --git a/llvm/unittests/ExecutionEngine/Orc/MapperJITLinkMemoryManagerTest.cpp b/llvm/unittests/ExecutionEngine/Orc/MapperJITLinkMemoryManagerTest.cpp
index e6dc303b6b296..74d8387701359 100644
--- a/llvm/unittests/ExecutionEngine/Orc/MapperJITLinkMemoryManagerTest.cpp
+++ b/llvm/unittests/ExecutionEngine/Orc/MapperJITLinkMemoryManagerTest.cpp
@@ -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
More information about the llvm-commits
mailing list