xref: /freebsd/contrib/llvm-project/llvm/include/llvm/ExecutionEngine/JITLink/JITLinkMemoryManager.h (revision 5f757f3ff9144b609b3c433dfd370cc6bdc191ad)
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, 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([&](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, [&](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(std::move(Alloc));
183     deallocate(std::move(Allocs), 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(std::move(Allocs),
191                [&](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(std::move(Alloc));
199     return deallocate(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(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