189 lines · cpp
1//=== MapperJITLinkMemoryManager.cpp - Memory management with MemoryMapper ===//2//3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.4// See https://llvm.org/LICENSE.txt for license information.5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception6//7//===----------------------------------------------------------------------===//8 9#include "llvm/ExecutionEngine/Orc/MapperJITLinkMemoryManager.h"10 11#include "llvm/ExecutionEngine/JITLink/JITLink.h"12#include "llvm/Support/Process.h"13 14using namespace llvm::jitlink;15 16namespace llvm {17namespace orc {18 19class MapperJITLinkMemoryManager::InFlightAlloc20 : public JITLinkMemoryManager::InFlightAlloc {21public:22 InFlightAlloc(MapperJITLinkMemoryManager &Parent, LinkGraph &G,23 ExecutorAddr AllocAddr,24 std::vector<MemoryMapper::AllocInfo::SegInfo> Segs)25 : Parent(Parent), G(G), AllocAddr(AllocAddr), Segs(std::move(Segs)) {}26 27 void finalize(OnFinalizedFunction OnFinalize) override {28 MemoryMapper::AllocInfo AI;29 AI.MappingBase = AllocAddr;30 31 std::swap(AI.Segments, Segs);32 std::swap(AI.Actions, G.allocActions());33 34 Parent.Mapper->initialize(AI, [OnFinalize = std::move(OnFinalize)](35 Expected<ExecutorAddr> Result) mutable {36 if (!Result) {37 OnFinalize(Result.takeError());38 return;39 }40 41 OnFinalize(FinalizedAlloc(*Result));42 });43 }44 45 void abandon(OnAbandonedFunction OnFinalize) override {46 Parent.Mapper->deinitialize({AllocAddr}, std::move(OnFinalize));47 }48 49private:50 MapperJITLinkMemoryManager &Parent;51 LinkGraph &G;52 ExecutorAddr AllocAddr;53 std::vector<MemoryMapper::AllocInfo::SegInfo> Segs;54};55 56MapperJITLinkMemoryManager::MapperJITLinkMemoryManager(57 size_t ReservationGranularity, std::unique_ptr<MemoryMapper> Mapper)58 : ReservationUnits(ReservationGranularity), AvailableMemory(AMAllocator),59 Mapper(std::move(Mapper)) {}60 61void MapperJITLinkMemoryManager::allocate(const JITLinkDylib *JD, LinkGraph &G,62 OnAllocatedFunction OnAllocated) {63 BasicLayout BL(G);64 65 // find required address space66 auto SegsSizes = BL.getContiguousPageBasedLayoutSizes(Mapper->getPageSize());67 if (!SegsSizes) {68 OnAllocated(SegsSizes.takeError());69 return;70 }71 72 auto TotalSize = SegsSizes->total();73 74 auto CompleteAllocation = [this, &G, BL = std::move(BL),75 OnAllocated = std::move(OnAllocated)](76 Expected<ExecutorAddrRange> Result) mutable {77 if (!Result) {78 Mutex.unlock();79 return OnAllocated(Result.takeError());80 }81 82 auto NextSegAddr = Result->Start;83 84 std::vector<MemoryMapper::AllocInfo::SegInfo> SegInfos;85 86 for (auto &KV : BL.segments()) {87 auto &AG = KV.first;88 auto &Seg = KV.second;89 90 auto TotalSize = Seg.ContentSize + Seg.ZeroFillSize;91 92 Seg.Addr = NextSegAddr;93 Seg.WorkingMem = Mapper->prepare(G, NextSegAddr, TotalSize);94 95 NextSegAddr += alignTo(TotalSize, Mapper->getPageSize());96 97 MemoryMapper::AllocInfo::SegInfo SI;98 SI.Offset = Seg.Addr - Result->Start;99 SI.ContentSize = Seg.ContentSize;100 SI.ZeroFillSize = Seg.ZeroFillSize;101 SI.AG = AG;102 SI.WorkingMem = Seg.WorkingMem;103 104 SegInfos.push_back(SI);105 }106 107 UsedMemory.insert({Result->Start, NextSegAddr - Result->Start});108 109 if (NextSegAddr < Result->End) {110 // Save the remaining memory for reuse in next allocation(s)111 AvailableMemory.insert(NextSegAddr, Result->End - 1, true);112 }113 Mutex.unlock();114 115 if (auto Err = BL.apply()) {116 OnAllocated(std::move(Err));117 return;118 }119 120 OnAllocated(std::make_unique<InFlightAlloc>(*this, G, Result->Start,121 std::move(SegInfos)));122 };123 124 Mutex.lock();125 126 // find an already reserved range that is large enough127 ExecutorAddrRange SelectedRange{};128 129 for (AvailableMemoryMap::iterator It = AvailableMemory.begin();130 It != AvailableMemory.end(); It++) {131 if (It.stop() - It.start() + 1 >= TotalSize) {132 SelectedRange = ExecutorAddrRange(It.start(), It.stop() + 1);133 It.erase();134 break;135 }136 }137 138 if (SelectedRange.empty()) { // no already reserved range was found139 auto TotalAllocation = alignTo(TotalSize, ReservationUnits);140 Mapper->reserve(TotalAllocation, std::move(CompleteAllocation));141 } else {142 CompleteAllocation(SelectedRange);143 }144}145 146void MapperJITLinkMemoryManager::deallocate(147 std::vector<FinalizedAlloc> Allocs, OnDeallocatedFunction OnDeallocated) {148 std::vector<ExecutorAddr> Bases;149 Bases.reserve(Allocs.size());150 for (auto &FA : Allocs) {151 ExecutorAddr Addr = FA.getAddress();152 Bases.push_back(Addr);153 }154 155 Mapper->deinitialize(Bases, [this, Allocs = std::move(Allocs),156 OnDeallocated = std::move(OnDeallocated)](157 llvm::Error Err) mutable {158 // TODO: How should we treat memory that we fail to deinitialize?159 // We're currently bailing out and treating it as "burned" -- should we160 // require that a failure to deinitialize still reset the memory so that161 // we can reclaim it?162 if (Err) {163 for (auto &FA : Allocs)164 FA.release();165 OnDeallocated(std::move(Err));166 return;167 }168 169 {170 std::lock_guard<std::mutex> Lock(Mutex);171 172 for (auto &FA : Allocs) {173 ExecutorAddr Addr = FA.getAddress();174 ExecutorAddrDiff Size = UsedMemory[Addr];175 176 UsedMemory.erase(Addr);177 AvailableMemory.insert(Addr, Addr + Size - 1, true);178 179 FA.release();180 }181 }182 183 OnDeallocated(Error::success());184 });185}186 187} // end namespace orc188} // end namespace llvm189