1 | //===-- JITLinkMemoryManager.h - JITLink mem manager interface --*- C++ -*-===// |
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 | // Contains the JITLinkMemoryManager interface. |
10 | // |
11 | //===----------------------------------------------------------------------===// |
12 | |
13 | #ifndef LLVM_EXECUTIONENGINE_JITLINK_JITLINKMEMORYMANAGER_H |
14 | #define LLVM_EXECUTIONENGINE_JITLINK_JITLINKMEMORYMANAGER_H |
15 | |
16 | #include "llvm/ADT/FunctionExtras.h" |
17 | #include "llvm/ADT/SmallVector.h" |
18 | #include "llvm/ExecutionEngine/JITLink/JITLinkDylib.h" |
19 | #include "llvm/ExecutionEngine/Orc/Shared/AllocationActions.h" |
20 | #include "llvm/ExecutionEngine/Orc/Shared/ExecutorAddress.h" |
21 | #include "llvm/ExecutionEngine/Orc/Shared/MemoryFlags.h" |
22 | #include "llvm/Support/Allocator.h" |
23 | #include "llvm/Support/Error.h" |
24 | #include "llvm/Support/MSVCErrorWorkarounds.h" |
25 | #include "llvm/Support/Memory.h" |
26 | #include "llvm/Support/RecyclingAllocator.h" |
27 | |
28 | #include <cstdint> |
29 | #include <future> |
30 | #include <mutex> |
31 | |
32 | namespace llvm { |
33 | namespace jitlink { |
34 | |
35 | class Block; |
36 | class LinkGraph; |
37 | class Section; |
38 | |
39 | /// Manages allocations of JIT memory. |
40 | /// |
41 | /// Instances of this class may be accessed concurrently from multiple threads |
42 | /// and their implemetations should include any necessary synchronization. |
43 | class JITLinkMemoryManager { |
44 | public: |
45 | |
46 | /// Represents a finalized allocation. |
47 | /// |
48 | /// Finalized allocations must be passed to the |
49 | /// JITLinkMemoryManager:deallocate method prior to being destroyed. |
50 | /// |
51 | /// The interpretation of the Address associated with the finalized allocation |
52 | /// is up to the memory manager implementation. Common options are using the |
53 | /// base address of the allocation, or the address of a memory management |
54 | /// object that tracks the allocation. |
55 | class FinalizedAlloc { |
56 | friend class JITLinkMemoryManager; |
57 | |
58 | static constexpr auto InvalidAddr = ~uint64_t(0); |
59 | |
60 | public: |
61 | FinalizedAlloc() = default; |
62 | explicit FinalizedAlloc(orc::ExecutorAddr A) : A(A) { |
63 | assert(A.getValue() != InvalidAddr && |
64 | "Explicitly creating an invalid allocation?" ); |
65 | } |
66 | FinalizedAlloc(const FinalizedAlloc &) = delete; |
67 | FinalizedAlloc(FinalizedAlloc &&Other) : A(Other.A) { |
68 | Other.A.setValue(InvalidAddr); |
69 | } |
70 | FinalizedAlloc &operator=(const FinalizedAlloc &) = delete; |
71 | FinalizedAlloc &operator=(FinalizedAlloc &&Other) { |
72 | assert(A.getValue() == InvalidAddr && |
73 | "Cannot overwrite active finalized allocation" ); |
74 | std::swap(a&: A, b&: Other.A); |
75 | return *this; |
76 | } |
77 | ~FinalizedAlloc() { |
78 | assert(A.getValue() == InvalidAddr && |
79 | "Finalized allocation was not deallocated" ); |
80 | } |
81 | |
82 | /// FinalizedAllocs convert to false for default-constructed, and |
83 | /// true otherwise. Default-constructed allocs need not be deallocated. |
84 | explicit operator bool() const { return A.getValue() != InvalidAddr; } |
85 | |
86 | /// Returns the address associated with this finalized allocation. |
87 | /// The allocation is unmodified. |
88 | orc::ExecutorAddr getAddress() const { return A; } |
89 | |
90 | /// Returns the address associated with this finalized allocation and |
91 | /// resets this object to the default state. |
92 | /// This should only be used by allocators when deallocating memory. |
93 | orc::ExecutorAddr release() { |
94 | orc::ExecutorAddr Tmp = A; |
95 | A.setValue(InvalidAddr); |
96 | return Tmp; |
97 | } |
98 | |
99 | private: |
100 | orc::ExecutorAddr A{InvalidAddr}; |
101 | }; |
102 | |
103 | /// Represents an allocation which has not been finalized yet. |
104 | /// |
105 | /// InFlightAllocs manage both executor memory allocations and working |
106 | /// memory allocations. |
107 | /// |
108 | /// On finalization, the InFlightAlloc should transfer the content of |
109 | /// working memory into executor memory, apply memory protections, and |
110 | /// run any finalization functions. |
111 | /// |
112 | /// Working memory should be kept alive at least until one of the following |
113 | /// happens: (1) the InFlightAlloc instance is destroyed, (2) the |
114 | /// InFlightAlloc is abandoned, (3) finalized target memory is destroyed. |
115 | /// |
116 | /// If abandon is called then working memory and executor memory should both |
117 | /// be freed. |
118 | class InFlightAlloc { |
119 | public: |
120 | using OnFinalizedFunction = unique_function<void(Expected<FinalizedAlloc>)>; |
121 | using OnAbandonedFunction = unique_function<void(Error)>; |
122 | |
123 | virtual ~InFlightAlloc(); |
124 | |
125 | /// Called prior to finalization if the allocation should be abandoned. |
126 | virtual void abandon(OnAbandonedFunction OnAbandoned) = 0; |
127 | |
128 | /// Called to transfer working memory to the target and apply finalization. |
129 | virtual void finalize(OnFinalizedFunction OnFinalized) = 0; |
130 | |
131 | /// Synchronous convenience version of finalize. |
132 | Expected<FinalizedAlloc> finalize() { |
133 | std::promise<MSVCPExpected<FinalizedAlloc>> FinalizeResultP; |
134 | auto FinalizeResultF = FinalizeResultP.get_future(); |
135 | finalize(OnFinalized: [&](Expected<FinalizedAlloc> Result) { |
136 | FinalizeResultP.set_value(std::move(Result)); |
137 | }); |
138 | return FinalizeResultF.get(); |
139 | } |
140 | }; |
141 | |
142 | /// Typedef for the argument to be passed to OnAllocatedFunction. |
143 | using AllocResult = Expected<std::unique_ptr<InFlightAlloc>>; |
144 | |
145 | /// Called when allocation has been completed. |
146 | using OnAllocatedFunction = unique_function<void(AllocResult)>; |
147 | |
148 | /// Called when deallocation has completed. |
149 | using OnDeallocatedFunction = unique_function<void(Error)>; |
150 | |
151 | virtual ~JITLinkMemoryManager(); |
152 | |
153 | /// Start the allocation process. |
154 | /// |
155 | /// If the initial allocation is successful then the OnAllocated function will |
156 | /// be called with a std::unique_ptr<InFlightAlloc> value. If the assocation |
157 | /// is unsuccessful then the OnAllocated function will be called with an |
158 | /// Error. |
159 | virtual void allocate(const JITLinkDylib *JD, LinkGraph &G, |
160 | OnAllocatedFunction OnAllocated) = 0; |
161 | |
162 | /// Convenience function for blocking allocation. |
163 | AllocResult allocate(const JITLinkDylib *JD, LinkGraph &G) { |
164 | std::promise<MSVCPExpected<std::unique_ptr<InFlightAlloc>>> AllocResultP; |
165 | auto AllocResultF = AllocResultP.get_future(); |
166 | allocate(JD, G, OnAllocated: [&](AllocResult Alloc) { |
167 | AllocResultP.set_value(std::move(Alloc)); |
168 | }); |
169 | return AllocResultF.get(); |
170 | } |
171 | |
172 | /// Deallocate a list of allocation objects. |
173 | /// |
174 | /// Dealloc actions will be run in reverse order (from the end of the vector |
175 | /// to the start). |
176 | virtual void deallocate(std::vector<FinalizedAlloc> Allocs, |
177 | OnDeallocatedFunction OnDeallocated) = 0; |
178 | |
179 | /// Convenience function for deallocation of a single alloc. |
180 | void deallocate(FinalizedAlloc Alloc, OnDeallocatedFunction OnDeallocated) { |
181 | std::vector<FinalizedAlloc> Allocs; |
182 | Allocs.push_back(x: std::move(Alloc)); |
183 | deallocate(Allocs: std::move(Allocs), OnDeallocated: std::move(OnDeallocated)); |
184 | } |
185 | |
186 | /// Convenience function for blocking deallocation. |
187 | Error deallocate(std::vector<FinalizedAlloc> Allocs) { |
188 | std::promise<MSVCPError> DeallocResultP; |
189 | auto DeallocResultF = DeallocResultP.get_future(); |
190 | deallocate(Allocs: std::move(Allocs), |
191 | OnDeallocated: [&](Error Err) { DeallocResultP.set_value(std::move(Err)); }); |
192 | return DeallocResultF.get(); |
193 | } |
194 | |
195 | /// Convenience function for blocking deallocation of a single alloc. |
196 | Error deallocate(FinalizedAlloc Alloc) { |
197 | std::vector<FinalizedAlloc> Allocs; |
198 | Allocs.push_back(x: std::move(Alloc)); |
199 | return deallocate(Allocs: std::move(Allocs)); |
200 | } |
201 | }; |
202 | |
203 | /// BasicLayout simplifies the implementation of JITLinkMemoryManagers. |
204 | /// |
205 | /// BasicLayout groups Sections into Segments based on their memory protection |
206 | /// and deallocation policies. JITLinkMemoryManagers can construct a BasicLayout |
207 | /// from a Graph, and then assign working memory and addresses to each of the |
208 | /// Segments. These addreses will be mapped back onto the Graph blocks in |
209 | /// the apply method. |
210 | class BasicLayout { |
211 | public: |
212 | /// The Alignment, ContentSize and ZeroFillSize of each segment will be |
213 | /// pre-filled from the Graph. Clients must set the Addr and WorkingMem fields |
214 | /// prior to calling apply. |
215 | // |
216 | // FIXME: The C++98 initializer is an attempt to work around compile failures |
217 | // due to http://www.open-std.org/jtc1/sc22/wg21/docs/cwg_defects.html#1397. |
218 | // We should be able to switch this back to member initialization once that |
219 | // issue is fixed. |
220 | class Segment { |
221 | friend class BasicLayout; |
222 | |
223 | public: |
224 | Segment() |
225 | : ContentSize(0), ZeroFillSize(0), Addr(0), WorkingMem(nullptr), |
226 | NextWorkingMemOffset(0) {} |
227 | Align Alignment; |
228 | size_t ContentSize; |
229 | uint64_t ZeroFillSize; |
230 | orc::ExecutorAddr Addr; |
231 | char *WorkingMem = nullptr; |
232 | |
233 | private: |
234 | size_t NextWorkingMemOffset; |
235 | std::vector<Block *> ContentBlocks, ZeroFillBlocks; |
236 | }; |
237 | |
238 | /// A convenience class that further groups segments based on memory |
239 | /// deallocation policy. This allows clients to make two slab allocations: |
240 | /// one for all standard segments, and one for all finalize segments. |
241 | struct ContiguousPageBasedLayoutSizes { |
242 | uint64_t StandardSegs = 0; |
243 | uint64_t FinalizeSegs = 0; |
244 | |
245 | uint64_t total() const { return StandardSegs + FinalizeSegs; } |
246 | }; |
247 | |
248 | private: |
249 | using SegmentMap = orc::AllocGroupSmallMap<Segment>; |
250 | |
251 | public: |
252 | BasicLayout(LinkGraph &G); |
253 | |
254 | /// Return a reference to the graph this allocation was created from. |
255 | LinkGraph &getGraph() { return G; } |
256 | |
257 | /// Returns the total number of required to allocate all segments (with each |
258 | /// segment padded out to page size) for all standard segments, and all |
259 | /// finalize segments. |
260 | /// |
261 | /// This is a convenience function for the common case where the segments will |
262 | /// be allocated contiguously. |
263 | /// |
264 | /// This function will return an error if any segment has an alignment that |
265 | /// is higher than a page. |
266 | Expected<ContiguousPageBasedLayoutSizes> |
267 | getContiguousPageBasedLayoutSizes(uint64_t PageSize); |
268 | |
269 | /// Returns an iterator over the segments of the layout. |
270 | iterator_range<SegmentMap::iterator> segments() { |
271 | return {Segments.begin(), Segments.end()}; |
272 | } |
273 | |
274 | /// Apply the layout to the graph. |
275 | Error apply(); |
276 | |
277 | /// Returns a reference to the AllocActions in the graph. |
278 | /// This convenience function saves callers from having to #include |
279 | /// LinkGraph.h if all they need are allocation actions. |
280 | orc::shared::AllocActions &graphAllocActions(); |
281 | |
282 | private: |
283 | LinkGraph &G; |
284 | SegmentMap Segments; |
285 | }; |
286 | |
287 | /// A utility class for making simple allocations using JITLinkMemoryManager. |
288 | /// |
289 | /// SimpleSegementAlloc takes a mapping of AllocGroups to Segments and uses |
290 | /// this to create a LinkGraph with one Section (containing one Block) per |
291 | /// Segment. Clients can obtain a pointer to the working memory and executor |
292 | /// address of that block using the Segment's AllocGroup. Once memory has been |
293 | /// populated, clients can call finalize to finalize the memory. |
294 | /// |
295 | /// Note: Segments with MemLifetime::NoAlloc are not permitted, since they would |
296 | /// not be useful, and their presence is likely to indicate a bug. |
297 | class SimpleSegmentAlloc { |
298 | public: |
299 | /// Describes a segment to be allocated. |
300 | struct Segment { |
301 | Segment() = default; |
302 | Segment(size_t ContentSize, Align ContentAlign) |
303 | : ContentSize(ContentSize), ContentAlign(ContentAlign) {} |
304 | |
305 | size_t ContentSize = 0; |
306 | Align ContentAlign; |
307 | }; |
308 | |
309 | /// Describes the segment working memory and executor address. |
310 | struct SegmentInfo { |
311 | orc::ExecutorAddr Addr; |
312 | MutableArrayRef<char> WorkingMem; |
313 | }; |
314 | |
315 | using SegmentMap = orc::AllocGroupSmallMap<Segment>; |
316 | |
317 | using OnCreatedFunction = unique_function<void(Expected<SimpleSegmentAlloc>)>; |
318 | |
319 | using OnFinalizedFunction = |
320 | JITLinkMemoryManager::InFlightAlloc::OnFinalizedFunction; |
321 | |
322 | static void Create(JITLinkMemoryManager &MemMgr, const JITLinkDylib *JD, |
323 | SegmentMap Segments, OnCreatedFunction OnCreated); |
324 | |
325 | static Expected<SimpleSegmentAlloc> Create(JITLinkMemoryManager &MemMgr, |
326 | const JITLinkDylib *JD, |
327 | SegmentMap Segments); |
328 | |
329 | SimpleSegmentAlloc(SimpleSegmentAlloc &&); |
330 | SimpleSegmentAlloc &operator=(SimpleSegmentAlloc &&); |
331 | ~SimpleSegmentAlloc(); |
332 | |
333 | /// Returns the SegmentInfo for the given group. |
334 | SegmentInfo getSegInfo(orc::AllocGroup AG); |
335 | |
336 | /// Finalize all groups (async version). |
337 | void finalize(OnFinalizedFunction OnFinalized) { |
338 | Alloc->finalize(OnFinalized: std::move(OnFinalized)); |
339 | } |
340 | |
341 | /// Finalize all groups. |
342 | Expected<JITLinkMemoryManager::FinalizedAlloc> finalize() { |
343 | return Alloc->finalize(); |
344 | } |
345 | |
346 | private: |
347 | SimpleSegmentAlloc( |
348 | std::unique_ptr<LinkGraph> G, |
349 | orc::AllocGroupSmallMap<Block *> ContentBlocks, |
350 | std::unique_ptr<JITLinkMemoryManager::InFlightAlloc> Alloc); |
351 | |
352 | std::unique_ptr<LinkGraph> G; |
353 | orc::AllocGroupSmallMap<Block *> ContentBlocks; |
354 | std::unique_ptr<JITLinkMemoryManager::InFlightAlloc> Alloc; |
355 | }; |
356 | |
357 | /// A JITLinkMemoryManager that allocates in-process memory. |
358 | class InProcessMemoryManager : public JITLinkMemoryManager { |
359 | public: |
360 | class IPInFlightAlloc; |
361 | |
362 | /// Attempts to auto-detect the host page size. |
363 | static Expected<std::unique_ptr<InProcessMemoryManager>> Create(); |
364 | |
365 | /// Create an instance using the given page size. |
366 | InProcessMemoryManager(uint64_t PageSize) : PageSize(PageSize) {} |
367 | |
368 | void allocate(const JITLinkDylib *JD, LinkGraph &G, |
369 | OnAllocatedFunction OnAllocated) override; |
370 | |
371 | // Use overloads from base class. |
372 | using JITLinkMemoryManager::allocate; |
373 | |
374 | void deallocate(std::vector<FinalizedAlloc> Alloc, |
375 | OnDeallocatedFunction OnDeallocated) override; |
376 | |
377 | // Use overloads from base class. |
378 | using JITLinkMemoryManager::deallocate; |
379 | |
380 | private: |
381 | // FIXME: Use an in-place array instead of a vector for DeallocActions. |
382 | // There shouldn't need to be a heap alloc for this. |
383 | struct FinalizedAllocInfo { |
384 | sys::MemoryBlock StandardSegments; |
385 | std::vector<orc::shared::WrapperFunctionCall> DeallocActions; |
386 | }; |
387 | |
388 | FinalizedAlloc createFinalizedAlloc( |
389 | sys::MemoryBlock StandardSegments, |
390 | std::vector<orc::shared::WrapperFunctionCall> DeallocActions); |
391 | |
392 | uint64_t PageSize; |
393 | std::mutex FinalizedAllocsMutex; |
394 | RecyclingAllocator<BumpPtrAllocator, FinalizedAllocInfo> FinalizedAllocInfos; |
395 | }; |
396 | |
397 | } // end namespace jitlink |
398 | } // end namespace llvm |
399 | |
400 | #endif // LLVM_EXECUTIONENGINE_JITLINK_JITLINKMEMORYMANAGER_H |
401 | |