1 //=== MapperJITLinkMemoryManager.cpp - Memory management with MemoryMapper ===//
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-exception
7 //===----------------------------------------------------------------------===//
9 #include "llvm/ExecutionEngine/Orc/MapperJITLinkMemoryManager.h"
11 #include "llvm/ADT/STLExtras.h"
12 #include "llvm/ExecutionEngine/JITLink/JITLink.h"
13 #include "llvm/Support/Process.h"
15 using namespace llvm::jitlink;
20 class MapperJITLinkMemoryManager::InFlightAlloc
21 : public JITLinkMemoryManager::InFlightAlloc {
23 InFlightAlloc(MapperJITLinkMemoryManager &Parent, LinkGraph &G,
24 ExecutorAddr AllocAddr,
25 std::vector<MemoryMapper::AllocInfo::SegInfo> Segs)
26 : Parent(Parent), G(G), AllocAddr(AllocAddr), Segs(std::move(Segs)) {}
28 void finalize(OnFinalizedFunction OnFinalize) override {
29 MemoryMapper::AllocInfo AI;
30 AI.MappingBase = AllocAddr;
32 std::swap(AI.Segments, Segs);
33 std::swap(AI.Actions, G.allocActions());
35 Parent.Mapper->initialize(AI, [OnFinalize = std::move(OnFinalize)](
36 Expected<ExecutorAddr> Result) mutable {
38 OnFinalize(Result.takeError());
42 OnFinalize(FinalizedAlloc(*Result));
46 void abandon(OnAbandonedFunction OnFinalize) override {
47 Parent.Mapper->release({AllocAddr}, std::move(OnFinalize));
51 MapperJITLinkMemoryManager &Parent;
53 ExecutorAddr AllocAddr;
54 std::vector<MemoryMapper::AllocInfo::SegInfo> Segs;
57 MapperJITLinkMemoryManager::MapperJITLinkMemoryManager(
58 size_t ReservationGranularity, std::unique_ptr<MemoryMapper> Mapper)
59 : ReservationUnits(ReservationGranularity), AvailableMemory(AMAllocator),
60 Mapper(std::move(Mapper)) {}
62 void MapperJITLinkMemoryManager::allocate(const JITLinkDylib *JD, LinkGraph &G,
63 OnAllocatedFunction OnAllocated) {
66 // find required address space
67 auto SegsSizes = BL.getContiguousPageBasedLayoutSizes(Mapper->getPageSize());
69 OnAllocated(SegsSizes.takeError());
73 auto TotalSize = SegsSizes->total();
75 auto CompleteAllocation = [this, &G, BL = std::move(BL),
76 OnAllocated = std::move(OnAllocated)](
77 Expected<ExecutorAddrRange> Result) mutable {
80 return OnAllocated(Result.takeError());
83 auto NextSegAddr = Result->Start;
85 std::vector<MemoryMapper::AllocInfo::SegInfo> SegInfos;
87 for (auto &KV : BL.segments()) {
89 auto &Seg = KV.second;
91 auto TotalSize = Seg.ContentSize + Seg.ZeroFillSize;
93 Seg.Addr = NextSegAddr;
94 Seg.WorkingMem = Mapper->prepare(NextSegAddr, TotalSize);
96 NextSegAddr += alignTo(TotalSize, Mapper->getPageSize());
98 MemoryMapper::AllocInfo::SegInfo SI;
99 SI.Offset = Seg.Addr - Result->Start;
100 SI.ContentSize = Seg.ContentSize;
101 SI.ZeroFillSize = Seg.ZeroFillSize;
103 SI.WorkingMem = Seg.WorkingMem;
105 SegInfos.push_back(SI);
108 UsedMemory.insert({Result->Start, NextSegAddr - Result->Start});
110 if (NextSegAddr < Result->End) {
111 // Save the remaining memory for reuse in next allocation(s)
112 AvailableMemory.insert(NextSegAddr, Result->End - 1, true);
116 if (auto Err = BL.apply()) {
117 OnAllocated(std::move(Err));
121 OnAllocated(std::make_unique<InFlightAlloc>(*this, G, Result->Start,
122 std::move(SegInfos)));
127 // find an already reserved range that is large enough
128 ExecutorAddrRange SelectedRange{};
130 for (AvailableMemoryMap::iterator It = AvailableMemory.begin();
131 It != AvailableMemory.end(); It++) {
132 if (It.stop() - It.start() + 1 >= TotalSize) {
133 SelectedRange = ExecutorAddrRange(It.start(), It.stop() + 1);
139 if (SelectedRange.empty()) { // no already reserved range was found
140 auto TotalAllocation = alignTo(TotalSize, ReservationUnits);
141 Mapper->reserve(TotalAllocation, std::move(CompleteAllocation));
143 CompleteAllocation(SelectedRange);
147 void MapperJITLinkMemoryManager::deallocate(
148 std::vector<FinalizedAlloc> Allocs, OnDeallocatedFunction OnDeallocated) {
149 std::vector<ExecutorAddr> Bases;
150 Bases.reserve(Allocs.size());
151 for (auto &FA : Allocs) {
152 ExecutorAddr Addr = FA.getAddress();
153 Bases.push_back(Addr);
156 Mapper->deinitialize(Bases, [this, Allocs = std::move(Allocs),
157 OnDeallocated = std::move(OnDeallocated)](
158 llvm::Error Err) mutable {
159 // TODO: How should we treat memory that we fail to deinitialize?
160 // We're currently bailing out and treating it as "burned" -- should we
161 // require that a failure to deinitialize still reset the memory so that
162 // we can reclaim it?
164 for (auto &FA : Allocs)
166 OnDeallocated(std::move(Err));
171 std::lock_guard<std::mutex> Lock(Mutex);
173 for (auto &FA : Allocs) {
174 ExecutorAddr Addr = FA.getAddress();
175 ExecutorAddrDiff Size = UsedMemory[Addr];
177 UsedMemory.erase(Addr);
178 AvailableMemory.insert(Addr, Addr + Size - 1, true);
184 OnDeallocated(Error::success());
188 } // end namespace orc
189 } // end namespace llvm