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(); }