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:
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 
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 
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 
57 MapperJITLinkMemoryManager::MapperJITLinkMemoryManager(
58     size_t ReservationGranularity, std::unique_ptr<MemoryMapper> Mapper)
59     : ReservationUnits(ReservationGranularity), AvailableMemory(AMAllocator),
60       Mapper(std::move(Mapper)) {}
61 
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 
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