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-exception
6 //
7 //===----------------------------------------------------------------------===//
8
9 #include "llvm/ExecutionEngine/Orc/MapperJITLinkMemoryManager.h"
10
11 #include "llvm/ADT/STLExtras.h"
12 #include "llvm/ExecutionEngine/JITLink/JITLink.h"
13 #include "llvm/Support/Process.h"
14
15 using namespace llvm::jitlink;
16
17 namespace llvm {
18 namespace orc {
19
20 class MapperJITLinkMemoryManager::InFlightAlloc
21 : public JITLinkMemoryManager::InFlightAlloc {
22 public:
InFlightAlloc(MapperJITLinkMemoryManager & Parent,LinkGraph & G,ExecutorAddr AllocAddr,std::vector<MemoryMapper::AllocInfo::SegInfo> Segs)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)) {}
27
finalize(OnFinalizedFunction OnFinalize)28 void finalize(OnFinalizedFunction OnFinalize) override {
29 MemoryMapper::AllocInfo AI;
30 AI.MappingBase = AllocAddr;
31
32 std::swap(AI.Segments, Segs);
33 std::swap(AI.Actions, G.allocActions());
34
35 Parent.Mapper->initialize(AI, [OnFinalize = std::move(OnFinalize)](
36 Expected<ExecutorAddr> Result) mutable {
37 if (!Result) {
38 OnFinalize(Result.takeError());
39 return;
40 }
41
42 OnFinalize(FinalizedAlloc(*Result));
43 });
44 }
45
abandon(OnAbandonedFunction OnFinalize)46 void abandon(OnAbandonedFunction OnFinalize) override {
47 Parent.Mapper->release({AllocAddr}, std::move(OnFinalize));
48 }
49
50 private:
51 MapperJITLinkMemoryManager &Parent;
52 LinkGraph &G;
53 ExecutorAddr AllocAddr;
54 std::vector<MemoryMapper::AllocInfo::SegInfo> Segs;
55 };
56
MapperJITLinkMemoryManager(size_t ReservationGranularity,std::unique_ptr<MemoryMapper> Mapper)57 MapperJITLinkMemoryManager::MapperJITLinkMemoryManager(
58 size_t ReservationGranularity, std::unique_ptr<MemoryMapper> Mapper)
59 : ReservationUnits(ReservationGranularity), AvailableMemory(AMAllocator),
60 Mapper(std::move(Mapper)) {}
61
allocate(const JITLinkDylib * JD,LinkGraph & G,OnAllocatedFunction OnAllocated)62 void MapperJITLinkMemoryManager::allocate(const JITLinkDylib *JD, LinkGraph &G,
63 OnAllocatedFunction OnAllocated) {
64 BasicLayout BL(G);
65
66 // find required address space
67 auto SegsSizes = BL.getContiguousPageBasedLayoutSizes(Mapper->getPageSize());
68 if (!SegsSizes) {
69 OnAllocated(SegsSizes.takeError());
70 return;
71 }
72
73 auto TotalSize = SegsSizes->total();
74
75 auto CompleteAllocation = [this, &G, BL = std::move(BL),
76 OnAllocated = std::move(OnAllocated)](
77 Expected<ExecutorAddrRange> Result) mutable {
78 if (!Result) {
79 Mutex.unlock();
80 return OnAllocated(Result.takeError());
81 }
82
83 auto NextSegAddr = Result->Start;
84
85 std::vector<MemoryMapper::AllocInfo::SegInfo> SegInfos;
86
87 for (auto &KV : BL.segments()) {
88 auto &AG = KV.first;
89 auto &Seg = KV.second;
90
91 auto TotalSize = Seg.ContentSize + Seg.ZeroFillSize;
92
93 Seg.Addr = NextSegAddr;
94 Seg.WorkingMem = Mapper->prepare(NextSegAddr, TotalSize);
95
96 NextSegAddr += alignTo(TotalSize, Mapper->getPageSize());
97
98 MemoryMapper::AllocInfo::SegInfo SI;
99 SI.Offset = Seg.Addr - Result->Start;
100 SI.ContentSize = Seg.ContentSize;
101 SI.ZeroFillSize = Seg.ZeroFillSize;
102 SI.AG = AG;
103 SI.WorkingMem = Seg.WorkingMem;
104
105 SegInfos.push_back(SI);
106 }
107
108 UsedMemory.insert({Result->Start, NextSegAddr - Result->Start});
109
110 if (NextSegAddr < Result->End) {
111 // Save the remaining memory for reuse in next allocation(s)
112 AvailableMemory.insert(NextSegAddr, Result->End - 1, true);
113 }
114 Mutex.unlock();
115
116 if (auto Err = BL.apply()) {
117 OnAllocated(std::move(Err));
118 return;
119 }
120
121 OnAllocated(std::make_unique<InFlightAlloc>(*this, G, Result->Start,
122 std::move(SegInfos)));
123 };
124
125 Mutex.lock();
126
127 // find an already reserved range that is large enough
128 ExecutorAddrRange SelectedRange{};
129
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);
134 It.erase();
135 break;
136 }
137 }
138
139 if (SelectedRange.empty()) { // no already reserved range was found
140 auto TotalAllocation = alignTo(TotalSize, ReservationUnits);
141 Mapper->reserve(TotalAllocation, std::move(CompleteAllocation));
142 } else {
143 CompleteAllocation(SelectedRange);
144 }
145 }
146
deallocate(std::vector<FinalizedAlloc> Allocs,OnDeallocatedFunction OnDeallocated)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);
154 }
155
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?
163 if (Err) {
164 for (auto &FA : Allocs)
165 FA.release();
166 OnDeallocated(std::move(Err));
167 return;
168 }
169
170 {
171 std::lock_guard<std::mutex> Lock(Mutex);
172
173 for (auto &FA : Allocs) {
174 ExecutorAddr Addr = FA.getAddress();
175 ExecutorAddrDiff Size = UsedMemory[Addr];
176
177 UsedMemory.erase(Addr);
178 AvailableMemory.insert(Addr, Addr + Size - 1, true);
179
180 FA.release();
181 }
182 }
183
184 OnDeallocated(Error::success());
185 });
186 }
187
188 } // end namespace orc
189 } // end namespace llvm
190