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
32namespace llvm {
33namespace jitlink {
34
35class Block;
36class LinkGraph;
37class 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.
43class JITLinkMemoryManager {
44public:
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.
210class BasicLayout {
211public:
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
248private:
249  using SegmentMap = orc::AllocGroupSmallMap<Segment>;
250
251public:
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
282private:
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.
297class SimpleSegmentAlloc {
298public:
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
346private:
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.
358class InProcessMemoryManager : public JITLinkMemoryManager {
359public:
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
380private:
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