diff options
Diffstat (limited to 'llvm/lib/ExecutionEngine/Orc/MapperJITLinkMemoryManager.cpp')
| -rw-r--r-- | llvm/lib/ExecutionEngine/Orc/MapperJITLinkMemoryManager.cpp | 146 |
1 files changed, 100 insertions, 46 deletions
diff --git a/llvm/lib/ExecutionEngine/Orc/MapperJITLinkMemoryManager.cpp b/llvm/lib/ExecutionEngine/Orc/MapperJITLinkMemoryManager.cpp index c2e7baabb994..d099a251232e 100644 --- a/llvm/lib/ExecutionEngine/Orc/MapperJITLinkMemoryManager.cpp +++ b/llvm/lib/ExecutionEngine/Orc/MapperJITLinkMemoryManager.cpp @@ -8,11 +8,10 @@ #include "llvm/ExecutionEngine/Orc/MapperJITLinkMemoryManager.h" +#include "llvm/ADT/STLExtras.h" #include "llvm/ExecutionEngine/JITLink/JITLink.h" #include "llvm/Support/Process.h" -#include <limits> - using namespace llvm::jitlink; namespace llvm { @@ -33,7 +32,8 @@ public: std::swap(AI.Segments, Segs); std::swap(AI.Actions, G.allocActions()); - Parent.Mapper->initialize(AI, [&](Expected<ExecutorAddr> Result) { + Parent.Mapper->initialize(AI, [OnFinalize = std::move(OnFinalize)]( + Expected<ExecutorAddr> Result) mutable { if (!Result) { OnFinalize(Result.takeError()); return; @@ -55,8 +55,9 @@ private: }; MapperJITLinkMemoryManager::MapperJITLinkMemoryManager( - std::unique_ptr<MemoryMapper> Mapper) - : Mapper(std::move(Mapper)) {} + size_t ReservationGranularity, std::unique_ptr<MemoryMapper> Mapper) + : ReservationUnits(ReservationGranularity), AvailableMemory(AMAllocator), + Mapper(std::move(Mapper)) {} void MapperJITLinkMemoryManager::allocate(const JITLinkDylib *JD, LinkGraph &G, OnAllocatedFunction OnAllocated) { @@ -69,55 +70,78 @@ 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(); + + 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; + + std::vector<MemoryMapper::AllocInfo::SegInfo> SegInfos; + + for (auto &KV : BL.segments()) { + auto &AG = KV.first; + auto &Seg = KV.second; - 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 TotalSize = Seg.ContentSize + Seg.ZeroFillSize; - auto NextSegAddr = Result->Start; + Seg.Addr = NextSegAddr; + Seg.WorkingMem = Mapper->prepare(NextSegAddr, TotalSize); - std::vector<MemoryMapper::AllocInfo::SegInfo> SegInfos; + NextSegAddr += alignTo(TotalSize, Mapper->getPageSize()); - for (auto &KV : BL.segments()) { - auto &AG = KV.first; - auto &Seg = KV.second; + MemoryMapper::AllocInfo::SegInfo SI; + SI.Offset = Seg.Addr - Result->Start; + SI.ContentSize = Seg.ContentSize; + SI.ZeroFillSize = Seg.ZeroFillSize; + SI.AG = AG; + SI.WorkingMem = Seg.WorkingMem; - auto TotalSize = Seg.ContentSize + Seg.ZeroFillSize; + SegInfos.push_back(SI); + } - Seg.Addr = NextSegAddr; - Seg.WorkingMem = Mapper->prepare(NextSegAddr, TotalSize); + UsedMemory.insert({Result->Start, NextSegAddr - Result->Start}); - NextSegAddr += alignTo(TotalSize, Mapper->getPageSize()); + if (NextSegAddr < Result->End) { + // Save the remaining memory for reuse in next allocation(s) + AvailableMemory.insert(NextSegAddr, Result->End - 1, true); + } + Mutex.unlock(); - 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; + if (auto Err = BL.apply()) { + OnAllocated(std::move(Err)); + return; + } - SegInfos.push_back(SI); - } + OnAllocated(std::make_unique<InFlightAlloc>(*this, G, Result->Start, + std::move(SegInfos))); + }; - if (auto Err = BL.apply()) { - OnAllocated(std::move(Err)); - return; - } + Mutex.lock(); - OnAllocated(std::make_unique<InFlightAlloc>(*this, G, Result->Start, - std::move(SegInfos))); - }); + // 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)); + } else { + CompleteAllocation(SelectedRange); + } } void MapperJITLinkMemoryManager::deallocate( @@ -125,10 +149,40 @@ void MapperJITLinkMemoryManager::deallocate( std::vector<ExecutorAddr> Bases; Bases.reserve(Allocs.size()); for (auto &FA : Allocs) { - Bases.push_back(FA.getAddress()); - FA.release(); + ExecutorAddr Addr = FA.getAddress(); + Bases.push_back(Addr); } - Mapper->release(Bases, std::move(OnDeallocated)); + + Mapper->deinitialize(Bases, [this, Allocs = std::move(Allocs), + OnDeallocated = std::move(OnDeallocated)]( + llvm::Error Err) mutable { + // TODO: How should we treat memory that we fail to deinitialize? + // We're currently bailing out and treating it as "burned" -- should we + // require that a failure to deinitialize still reset the memory so that + // we can reclaim it? + if (Err) { + for (auto &FA : Allocs) + FA.release(); + OnDeallocated(std::move(Err)); + return; + } + + { + std::lock_guard<std::mutex> Lock(Mutex); + + for (auto &FA : Allocs) { + ExecutorAddr Addr = FA.getAddress(); + ExecutorAddrDiff Size = UsedMemory[Addr]; + + UsedMemory.erase(Addr); + AvailableMemory.insert(Addr, Addr + Size - 1, true); + + FA.release(); + } + } + + OnDeallocated(Error::success()); + }); } } // end namespace orc |