#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 {
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,
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;
};
#include "llvm/ExecutionEngine/Orc/MapperJITLinkMemoryManager.h"
+#include "llvm/ADT/STLExtras.h"
#include "llvm/ExecutionEngine/JITLink/JITLink.h"
#include "llvm/Support/Process.h"
};
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) {
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
#include "llvm/ExecutionEngine/Orc/MapperJITLinkMemoryManager.h"
+#include "llvm/ExecutionEngine/Orc/MemoryMapper.h"
#include "llvm/Testing/Support/Error.h"
#include <vector>
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