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 @@ -13,6 +13,7 @@ #ifndef LLVM_EXECUTIONENGINE_ORC_MAPPERJITLINKMEMORYMANAGER_H #define LLVM_EXECUTIONENGINE_ORC_MAPPERJITLINKMEMORYMANAGER_H +#include "llvm/ADT/IntervalMap.h" #include "llvm/ExecutionEngine/JITLink/JITLinkMemoryManager.h" #include "llvm/ExecutionEngine/Orc/MemoryMapper.h" @@ -52,8 +53,12 @@ // 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; + using AvailableMemoryMap = IntervalMap; + AvailableMemoryMap::Allocator AMAllocator; + IntervalMap AvailableMemory; + // Ranges that have been reserved in executor and already allocated DenseMap UsedMemory; 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 @@ -57,7 +57,8 @@ MapperJITLinkMemoryManager::MapperJITLinkMemoryManager( size_t ReservationGranularity, std::unique_ptr Mapper) - : ReservationUnits(ReservationGranularity), Mapper(std::move(Mapper)) {} + : ReservationUnits(ReservationGranularity), AvailableMemory(AMAllocator), + Mapper(std::move(Mapper)) {} void MapperJITLinkMemoryManager::allocate(const JITLinkDylib *JD, LinkGraph &G, OnAllocatedFunction OnAllocated) { @@ -72,20 +73,6 @@ 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); - } - auto CompleteAllocation = [this, &G, BL = std::move(BL), OnAllocated = std::move(OnAllocated)]( Expected Result) mutable { @@ -123,8 +110,7 @@ if (NextSegAddr < Result->End) { // Save the remaining memory for reuse in next allocation(s) - auto RemainingRange = ExecutorAddrRange(NextSegAddr, Result->End); - AvailableMemory.push_back(RemainingRange); + AvailableMemory.insert(NextSegAddr, Result->End - 1, true); } Mutex.unlock(); @@ -137,6 +123,20 @@ std::move(SegInfos))); }; + Mutex.lock(); + + // find an already reserved range that is large enough + ExecutorAddrRange SelectedRange{}; + + for (AvailableMemoryMap::iterator It = AvailableMemory.begin(); + It != AvailableMemory.end(); It++) { + if (It.stop() - It.start() + 1 >= TotalSize) { + SelectedRange = ExecutorAddrRange(It.start(), It.stop() + 1); + It.erase(); + break; + } + } + if (SelectedRange.empty()) { // no already reserved range was found auto TotalAllocation = alignTo(TotalSize, ReservationUnits); Mapper->reserve(TotalAllocation, std::move(CompleteAllocation)); @@ -168,7 +168,7 @@ ExecutorAddrDiff Size = UsedMemory[Addr]; UsedMemory.erase(Addr); - AvailableMemory.push_back({Addr, Addr + Size}); + AvailableMemory.insert(Addr, Addr + Size - 1, true); FA.release(); } 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 @@ -1,4 +1,4 @@ -//===-------------- 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. @@ -131,4 +131,46 @@ EXPECT_EQ(Counter->DeinitCount, 2); } +TEST(MapperJITLinkMemoryManagerTest, Coalescing) { + auto Mapper = cantFail(InProcessMemoryMapper::Create()); + auto MemMgr = std::make_unique(16 * 1024 * 1024, + std::move(Mapper)); + + auto SSA1 = jitlink::SimpleSegmentAlloc::Create( + *MemMgr, nullptr, {{jitlink::MemProt::Read, {1024, Align(1)}}}); + EXPECT_THAT_EXPECTED(SSA1, Succeeded()); + auto SegInfo1 = SSA1->getSegInfo(jitlink::MemProt::Read); + ExecutorAddr TargetAddr1(SegInfo1.Addr); + auto FA1 = SSA1->finalize(); + EXPECT_THAT_EXPECTED(FA1, Succeeded()); + + auto SSA2 = jitlink::SimpleSegmentAlloc::Create( + *MemMgr, nullptr, {{jitlink::MemProt::Read, {1024, Align(1)}}}); + EXPECT_THAT_EXPECTED(SSA2, Succeeded()); + auto FA2 = SSA2->finalize(); + EXPECT_THAT_EXPECTED(FA2, Succeeded()); + + 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()); + + auto SSA3 = jitlink::SimpleSegmentAlloc::Create( + *MemMgr, nullptr, {{jitlink::MemProt::Read, {2048, Align(1)}}}); + EXPECT_THAT_EXPECTED(SSA3, Succeeded()); + + auto SegInfo3 = SSA3->getSegInfo(jitlink::MemProt::Read); + ExecutorAddr TargetAddr3(SegInfo3.Addr); + + auto FA3 = SSA3->finalize(); + EXPECT_THAT_EXPECTED(FA3, Succeeded()); + + // previous two freed 1024 blocks should be fused to form a 2048 block + EXPECT_EQ(TargetAddr1, TargetAddr3); + + auto Err4 = MemMgr->deallocate(std::move(*FA3)); + EXPECT_THAT_ERROR(std::move(Err4), Succeeded()); +} + } // namespace