diff --git a/llvm/include/llvm/ExecutionEngine/Orc/MapperJITLinkMemoryManager.h b/llvm/include/llvm/ExecutionEngine/Orc/MapperJITLinkMemoryManager.h --- 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 @@ class MapperJITLinkMemoryManager : public jitlink::JITLinkMemoryManager { public: - MapperJITLinkMemoryManager(std::unique_ptr Mapper); + MapperJITLinkMemoryManager(size_t ReservationGranularity, + std::unique_ptr Mapper); template static Expected> - CreateWithMapper(Args &&...A) { + CreateWithMapper(size_t ReservationGranularity, Args &&...A) { auto Mapper = MemoryMapperType::Create(std::forward(A)...); if (!Mapper) return Mapper.takeError(); - return std::make_unique(std::move(*Mapper)); + return std::make_unique(ReservationGranularity, + std::move(*Mapper)); } void allocate(const jitlink::JITLinkDylib *JD, jitlink::LinkGraph &G, @@ -47,6 +48,15 @@ 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 AvailableMemory; + // Ranges that have been reserved in executor and already allocated + DenseMap UsedMemory; + std::unique_ptr Mapper; }; diff --git a/llvm/lib/ExecutionEngine/Orc/MapperJITLinkMemoryManager.cpp b/llvm/lib/ExecutionEngine/Orc/MapperJITLinkMemoryManager.cpp --- 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 @@ }; MapperJITLinkMemoryManager::MapperJITLinkMemoryManager( - std::unique_ptr Mapper) - : Mapper(std::move(Mapper)) {} + size_t ReservationGranularity, std::unique_ptr Mapper) + : ReservationUnits(ReservationGranularity), Mapper(std::move(Mapper)) {} void MapperJITLinkMemoryManager::allocate(const JITLinkDylib *JD, LinkGraph &G, OnAllocatedFunction OnAllocated) { @@ -69,66 +70,94 @@ return; } - // Check if total size fits in address space - if (SegsSizes->total() > std::numeric_limits::max()) { - OnAllocated(make_error( - 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::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 Result) mutable { - if (!Result) { - return OnAllocated(Result.takeError()); - } + auto CompleteAllocation = [this, &G, BL = std::move(BL), + OnAllocated = std::move(OnAllocated)]( + Expected Result) mutable { + if (!Result) { + Mutex.unlock(); + return OnAllocated(Result.takeError()); + } - auto NextSegAddr = Result->Start; + auto NextSegAddr = Result->Start; - std::vector SegInfos; + std::vector 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(*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(*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 Allocs, OnDeallocatedFunction OnDeallocated) { - std::vector Bases; - Bases.reserve(Allocs.size()); + std::lock_guard 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 --- 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 @@ -21,28 +22,107 @@ namespace { +class CounterMapper final : public MemoryMapper { +public: + CounterMapper(std::unique_ptr 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 Allocations, + OnDeinitializedFunction OnDeInitialized) override { + ++DeinitCount; + return Mapper->deinitialize(Allocations, std::move(OnDeInitialized)); + } + + void release(ArrayRef 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 Mapper; +}; + TEST(MapperJITLinkMemoryManagerTest, InProcess) { - auto MemMgr = cantFail( - MapperJITLinkMemoryManager::CreateWithMapper()); + auto Mapper = std::make_unique( + cantFail(InProcessMemoryMapper::Create())); + + auto *Counter = static_cast(Mapper.get()); + + auto MemMgr = std::make_unique(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(); - 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(); + StringRef TargetHello1(TargetMem1, Hello.size()); + EXPECT_EQ(Hello, TargetHello1); + + const char *TargetMem2 = TargetAddr2.toPtr(); + 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