xref: /llvm-project/llvm/lib/ExecutionEngine/JITLink/JITLinkMemoryManager.cpp (revision c5965a411c635106a47738b8d2e24db822b7416f) 
1 //===--- JITLinkMemoryManager.cpp - JITLinkMemoryManager implementation ---===//
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/JITLink/JITLinkMemoryManager.h"
10 #include "llvm/ExecutionEngine/JITLink/JITLink.h"
11 #include "llvm/Support/FormatVariadic.h"
12 #include "llvm/Support/Process.h"
13 
14 #define DEBUG_TYPE "jitlink"
15 
16 using namespace llvm;
17 
18 namespace {
19 
20 // FIXME: Remove this copy of CWrapperFunctionResult as soon as JITLink can
21 // depend on shared utils from Orc.
22 
23 // Must be kept in-sync with compiler-rt/lib/orc/c-api.h.
24 union CWrapperFunctionResultDataUnion {
25   char *ValuePtr;
26   char Value[sizeof(ValuePtr)];
27 };
28 
29 // Must be kept in-sync with compiler-rt/lib/orc/c-api.h.
30 typedef struct {
31   CWrapperFunctionResultDataUnion Data;
32   size_t Size;
33 } CWrapperFunctionResult;
34 
35 Error toError(CWrapperFunctionResult R) {
36   bool HasError = false;
37   std::string ErrMsg;
38   if (R.Size) {
39     bool Large = R.Size > sizeof(CWrapperFunctionResultDataUnion);
40     char *Content = Large ? R.Data.ValuePtr : R.Data.Value;
41     if (Content[0]) {
42       HasError = true;
43       constexpr unsigned StrStart = 1 + sizeof(uint64_t);
44       ErrMsg.resize(R.Size - StrStart);
45       memcpy(&ErrMsg[0], Content + StrStart, R.Size - StrStart);
46     }
47     if (Large)
48       free(R.Data.ValuePtr);
49   } else if (R.Data.ValuePtr) {
50     HasError = true;
51     ErrMsg = R.Data.ValuePtr;
52     free(R.Data.ValuePtr);
53   }
54 
55   if (HasError)
56     return make_error<StringError>(std::move(ErrMsg), inconvertibleErrorCode());
57   return Error::success();
58 }
59 } // namespace
60 
61 namespace llvm {
62 namespace jitlink {
63 
64 JITLinkMemoryManager::~JITLinkMemoryManager() = default;
65 JITLinkMemoryManager::InFlightAlloc::~InFlightAlloc() = default;
66 
67 static Error runAllocAction(AllocActionCall &C) {
68   using WrapperFnTy = CWrapperFunctionResult (*)(const void *, size_t);
69   auto *Fn = jitTargetAddressToPointer<WrapperFnTy>(C.FnAddr);
70 
71   return toError(Fn(jitTargetAddressToPointer<const void *>(C.CtxAddr),
72                     static_cast<size_t>(C.CtxSize)));
73 }
74 
75 BasicLayout::BasicLayout(LinkGraph &G) : G(G) {
76 
77   for (auto &Sec : G.sections()) {
78     // Skip empty sections.
79     if (empty(Sec.blocks()))
80       continue;
81 
82     auto &Seg = Segments[{Sec.getMemProt(), Sec.getMemDeallocPolicy()}];
83     for (auto *B : Sec.blocks())
84       if (LLVM_LIKELY(!B->isZeroFill()))
85         Seg.ContentBlocks.push_back(B);
86       else
87         Seg.ZeroFillBlocks.push_back(B);
88   }
89 
90   // Build Segments map.
91   auto CompareBlocks = [](const Block *LHS, const Block *RHS) {
92     // Sort by section, address and size
93     if (LHS->getSection().getOrdinal() != RHS->getSection().getOrdinal())
94       return LHS->getSection().getOrdinal() < RHS->getSection().getOrdinal();
95     if (LHS->getAddress() != RHS->getAddress())
96       return LHS->getAddress() < RHS->getAddress();
97     return LHS->getSize() < RHS->getSize();
98   };
99 
100   LLVM_DEBUG(dbgs() << "Generated BasicLayout for " << G.getName() << ":\n");
101   for (auto &KV : Segments) {
102     auto &Seg = KV.second;
103 
104     llvm::sort(Seg.ContentBlocks, CompareBlocks);
105     llvm::sort(Seg.ZeroFillBlocks, CompareBlocks);
106 
107     for (auto *B : Seg.ContentBlocks) {
108       Seg.ContentSize = alignToBlock(Seg.ContentSize, *B);
109       Seg.ContentSize += B->getSize();
110       Seg.Alignment = std::max(Seg.Alignment, Align(B->getAlignment()));
111     }
112 
113     uint64_t SegEndOffset = Seg.ContentSize;
114     for (auto *B : Seg.ZeroFillBlocks) {
115       SegEndOffset = alignToBlock(SegEndOffset, *B);
116       SegEndOffset += B->getSize();
117       Seg.Alignment = std::max(Seg.Alignment, Align(B->getAlignment()));
118     }
119     Seg.ZeroFillSize = SegEndOffset - Seg.ContentSize;
120 
121     LLVM_DEBUG({
122       dbgs() << "  Seg " << KV.first
123              << ": content-size=" << formatv("{0:x}", Seg.ContentSize)
124              << ", zero-fill-size=" << formatv("{0:x}", Seg.ZeroFillSize)
125              << ", align=" << formatv("{0:x}", Seg.Alignment.value()) << "\n";
126     });
127   }
128 }
129 
130 Expected<BasicLayout::ContiguousPageBasedLayoutSizes>
131 BasicLayout::getContiguousPageBasedLayoutSizes(uint64_t PageSize) {
132   ContiguousPageBasedLayoutSizes SegsSizes;
133 
134   for (auto &KV : segments()) {
135     auto &AG = KV.first;
136     auto &Seg = KV.second;
137 
138     if (Seg.Alignment > PageSize)
139       return make_error<StringError>("Segment alignment greater than page size",
140                                      inconvertibleErrorCode());
141 
142     uint64_t SegSize = alignTo(Seg.ContentSize + Seg.ZeroFillSize, PageSize);
143     if (AG.getMemDeallocPolicy() == MemDeallocPolicy::Standard)
144       SegsSizes.StandardSegs += SegSize;
145     else
146       SegsSizes.FinalizeSegs += SegSize;
147   }
148 
149   return SegsSizes;
150 }
151 
152 Error BasicLayout::apply() {
153   for (auto &KV : Segments) {
154     auto &Seg = KV.second;
155 
156     assert(!(Seg.ContentBlocks.empty() && Seg.ZeroFillBlocks.empty()) &&
157            "Empty section recorded?");
158 
159     for (auto *B : Seg.ContentBlocks) {
160       // Align addr and working-mem-offset.
161       Seg.Addr = alignToBlock(Seg.Addr, *B);
162       Seg.NextWorkingMemOffset = alignToBlock(Seg.NextWorkingMemOffset, *B);
163 
164       // Update block addr.
165       B->setAddress(Seg.Addr);
166       Seg.Addr += B->getSize();
167 
168       // Copy content to working memory, then update content to point at working
169       // memory.
170       memcpy(Seg.WorkingMem + Seg.NextWorkingMemOffset, B->getContent().data(),
171              B->getSize());
172       B->setMutableContent(
173           {Seg.WorkingMem + Seg.NextWorkingMemOffset, B->getSize()});
174       Seg.NextWorkingMemOffset += B->getSize();
175     }
176 
177     for (auto *B : Seg.ZeroFillBlocks) {
178       // Align addr.
179       Seg.Addr = alignToBlock(Seg.Addr, *B);
180       // Update block addr.
181       B->setAddress(Seg.Addr);
182       Seg.Addr += B->getSize();
183     }
184 
185     Seg.ContentBlocks.clear();
186     Seg.ZeroFillBlocks.clear();
187   }
188 
189   return Error::success();
190 }
191 
192 AllocActions &BasicLayout::graphAllocActions() { return G.allocActions(); }
193 
194 void SimpleSegmentAlloc::Create(JITLinkMemoryManager &MemMgr,
195                                 const JITLinkDylib *JD, SegmentMap Segments,
196                                 OnCreatedFunction OnCreated) {
197 
198   static_assert(AllocGroup::NumGroups == 16,
199                 "AllocGroup has changed. Section names below must be updated");
200   StringRef AGSectionNames[] = {
201       "__---.standard", "__R--.standard", "__-W-.standard", "__RW-.standard",
202       "__--X.standard", "__R-X.standard", "__-WX.standard", "__RWX.standard",
203       "__---.finalize", "__R--.finalize", "__-W-.finalize", "__RW-.finalize",
204       "__--X.finalize", "__R-X.finalize", "__-WX.finalize", "__RWX.finalize"};
205 
206   auto G =
207       std::make_unique<LinkGraph>("", Triple(), 0, support::native, nullptr);
208   AllocGroupSmallMap<Block *> ContentBlocks;
209 
210   JITTargetAddress NextAddr = 0x100000;
211   for (auto &KV : Segments) {
212     auto &AG = KV.first;
213     auto &Seg = KV.second;
214 
215     auto AGSectionName =
216         AGSectionNames[static_cast<unsigned>(AG.getMemProt()) |
217                        static_cast<bool>(AG.getMemDeallocPolicy()) << 3];
218 
219     auto &Sec = G->createSection(AGSectionName, AG.getMemProt());
220     Sec.setMemDeallocPolicy(AG.getMemDeallocPolicy());
221 
222     if (Seg.ContentSize != 0) {
223       NextAddr = alignTo(NextAddr, Seg.ContentAlign);
224       auto &B =
225           G->createMutableContentBlock(Sec, G->allocateBuffer(Seg.ContentSize),
226                                        NextAddr, Seg.ContentAlign.value(), 0);
227       ContentBlocks[AG] = &B;
228       NextAddr += Seg.ContentSize;
229     }
230   }
231 
232   // GRef declared separately since order-of-argument-eval isn't specified.
233   auto &GRef = *G;
234   MemMgr.allocate(JD, GRef,
235                   [G = std::move(G), ContentBlocks = std::move(ContentBlocks),
236                    OnCreated = std::move(OnCreated)](
237                       JITLinkMemoryManager::AllocResult Alloc) mutable {
238                     if (!Alloc)
239                       OnCreated(Alloc.takeError());
240                     else
241                       OnCreated(SimpleSegmentAlloc(std::move(G),
242                                                    std::move(ContentBlocks),
243                                                    std::move(*Alloc)));
244                   });
245 }
246 
247 Expected<SimpleSegmentAlloc>
248 SimpleSegmentAlloc::Create(JITLinkMemoryManager &MemMgr, const JITLinkDylib *JD,
249                            SegmentMap Segments) {
250   std::promise<MSVCPExpected<SimpleSegmentAlloc>> AllocP;
251   auto AllocF = AllocP.get_future();
252   Create(MemMgr, JD, std::move(Segments),
253          [&](Expected<SimpleSegmentAlloc> Result) {
254            AllocP.set_value(std::move(Result));
255          });
256   return AllocF.get();
257 }
258 
259 SimpleSegmentAlloc::SimpleSegmentAlloc(SimpleSegmentAlloc &&) = default;
260 SimpleSegmentAlloc &
261 SimpleSegmentAlloc::operator=(SimpleSegmentAlloc &&) = default;
262 SimpleSegmentAlloc::~SimpleSegmentAlloc() {}
263 
264 SimpleSegmentAlloc::SegmentInfo SimpleSegmentAlloc::getSegInfo(AllocGroup AG) {
265   auto I = ContentBlocks.find(AG);
266   if (I != ContentBlocks.end()) {
267     auto &B = *I->second;
268     return {B.getAddress(), B.getAlreadyMutableContent()};
269   }
270   return {};
271 }
272 
273 SimpleSegmentAlloc::SimpleSegmentAlloc(
274     std::unique_ptr<LinkGraph> G, AllocGroupSmallMap<Block *> ContentBlocks,
275     std::unique_ptr<JITLinkMemoryManager::InFlightAlloc> Alloc)
276     : G(std::move(G)), ContentBlocks(std::move(ContentBlocks)),
277       Alloc(std::move(Alloc)) {}
278 
279 class InProcessMemoryManager::IPInFlightAlloc
280     : public JITLinkMemoryManager::InFlightAlloc {
281 public:
282   IPInFlightAlloc(InProcessMemoryManager &MemMgr, LinkGraph &G, BasicLayout BL,
283                   sys::MemoryBlock StandardSegments,
284                   sys::MemoryBlock FinalizationSegments)
285       : MemMgr(MemMgr), G(G), BL(std::move(BL)),
286         StandardSegments(std::move(StandardSegments)),
287         FinalizationSegments(std::move(FinalizationSegments)) {}
288 
289   void finalize(OnFinalizedFunction OnFinalized) override {
290 
291     // Apply memory protections to all segments.
292     if (auto Err = applyProtections()) {
293       OnFinalized(std::move(Err));
294       return;
295     }
296 
297     // Run finalization actions.
298     // FIXME: Roll back previous successful actions on failure.
299     std::vector<AllocActionCall> DeallocActions;
300     DeallocActions.reserve(G.allocActions().size());
301     for (auto &ActPair : G.allocActions()) {
302       if (ActPair.Finalize.FnAddr)
303         if (auto Err = runAllocAction(ActPair.Finalize)) {
304           OnFinalized(std::move(Err));
305           return;
306         }
307       if (ActPair.Dealloc.FnAddr)
308         DeallocActions.push_back(ActPair.Dealloc);
309     }
310     G.allocActions().clear();
311 
312     // Release the finalize segments slab.
313     if (auto EC = sys::Memory::releaseMappedMemory(FinalizationSegments)) {
314       OnFinalized(errorCodeToError(EC));
315       return;
316     }
317 
318     // Continue with finalized allocation.
319     OnFinalized(MemMgr.createFinalizedAlloc(std::move(StandardSegments),
320                                             std::move(DeallocActions)));
321   }
322 
323   void abandon(OnAbandonedFunction OnAbandoned) override {
324     Error Err = Error::success();
325     if (auto EC = sys::Memory::releaseMappedMemory(FinalizationSegments))
326       Err = joinErrors(std::move(Err), errorCodeToError(EC));
327     if (auto EC = sys::Memory::releaseMappedMemory(StandardSegments))
328       Err = joinErrors(std::move(Err), errorCodeToError(EC));
329     OnAbandoned(std::move(Err));
330   }
331 
332 private:
333   Error applyProtections() {
334     for (auto &KV : BL.segments()) {
335       const auto &AG = KV.first;
336       auto &Seg = KV.second;
337 
338       auto Prot = toSysMemoryProtectionFlags(AG.getMemProt());
339 
340       uint64_t SegSize =
341           alignTo(Seg.ContentSize + Seg.ZeroFillSize, MemMgr.PageSize);
342       sys::MemoryBlock MB(Seg.WorkingMem, SegSize);
343       if (auto EC = sys::Memory::protectMappedMemory(MB, Prot))
344         return errorCodeToError(EC);
345       if (Prot & sys::Memory::MF_EXEC)
346         sys::Memory::InvalidateInstructionCache(MB.base(), MB.allocatedSize());
347     }
348     return Error::success();
349   }
350 
351   InProcessMemoryManager &MemMgr;
352   LinkGraph &G;
353   BasicLayout BL;
354   sys::MemoryBlock StandardSegments;
355   sys::MemoryBlock FinalizationSegments;
356 };
357 
358 Expected<std::unique_ptr<InProcessMemoryManager>>
359 InProcessMemoryManager::Create() {
360   if (auto PageSize = sys::Process::getPageSize())
361     return std::make_unique<InProcessMemoryManager>(*PageSize);
362   else
363     return PageSize.takeError();
364 }
365 
366 void InProcessMemoryManager::allocate(const JITLinkDylib *JD, LinkGraph &G,
367                                       OnAllocatedFunction OnAllocated) {
368 
369   // FIXME: Just check this once on startup.
370   if (!isPowerOf2_64((uint64_t)PageSize)) {
371     OnAllocated(make_error<StringError>("Page size is not a power of 2",
372                                         inconvertibleErrorCode()));
373     return;
374   }
375 
376   BasicLayout BL(G);
377 
378   /// Scan the request and calculate the group and total sizes.
379   /// Check that segment size is no larger than a page.
380   auto SegsSizes = BL.getContiguousPageBasedLayoutSizes(PageSize);
381   if (!SegsSizes) {
382     OnAllocated(SegsSizes.takeError());
383     return;
384   }
385 
386   /// Check that the total size requested (including zero fill) is not larger
387   /// than a size_t.
388   if (SegsSizes->total() > std::numeric_limits<size_t>::max()) {
389     OnAllocated(make_error<JITLinkError>(
390         "Total requested size " + formatv("{0:x}", SegsSizes->total()) +
391         " for graph " + G.getName() + " exceeds address space"));
392     return;
393   }
394 
395   // Allocate one slab for the whole thing (to make sure everything is
396   // in-range), then partition into standard and finalization blocks.
397   //
398   // FIXME: Make two separate allocations in the future to reduce
399   // fragmentation: finalization segments will usually be a single page, and
400   // standard segments are likely to be more than one page. Where multiple
401   // allocations are in-flight at once (likely) the current approach will leave
402   // a lot of single-page holes.
403   sys::MemoryBlock Slab;
404   sys::MemoryBlock StandardSegsMem;
405   sys::MemoryBlock FinalizeSegsMem;
406   {
407     const sys::Memory::ProtectionFlags ReadWrite =
408         static_cast<sys::Memory::ProtectionFlags>(sys::Memory::MF_READ |
409                                                   sys::Memory::MF_WRITE);
410 
411     std::error_code EC;
412     Slab = sys::Memory::allocateMappedMemory(SegsSizes->total(), nullptr,
413                                              ReadWrite, EC);
414 
415     if (EC) {
416       OnAllocated(errorCodeToError(EC));
417       return;
418     }
419 
420     // Zero-fill the whole slab up-front.
421     memset(Slab.base(), 0, Slab.allocatedSize());
422 
423     StandardSegsMem = {Slab.base(),
424                        static_cast<size_t>(SegsSizes->StandardSegs)};
425     FinalizeSegsMem = {(void *)((char *)Slab.base() + SegsSizes->StandardSegs),
426                        static_cast<size_t>(SegsSizes->FinalizeSegs)};
427   }
428 
429   auto NextStandardSegAddr = pointerToJITTargetAddress(StandardSegsMem.base());
430   auto NextFinalizeSegAddr = pointerToJITTargetAddress(FinalizeSegsMem.base());
431 
432   LLVM_DEBUG({
433     dbgs() << "InProcessMemoryManager allocated:\n";
434     if (SegsSizes->StandardSegs)
435       dbgs() << formatv("  [ {0:x16} -- {1:x16} ]", NextStandardSegAddr,
436                         NextStandardSegAddr + StandardSegsMem.allocatedSize())
437              << " to stardard segs\n";
438     else
439       dbgs() << "  no standard segs\n";
440     if (SegsSizes->FinalizeSegs)
441       dbgs() << formatv("  [ {0:x16} -- {1:x16} ]", NextFinalizeSegAddr,
442                         NextFinalizeSegAddr + FinalizeSegsMem.allocatedSize())
443              << " to finalize segs\n";
444     else
445       dbgs() << "  no finalize segs\n";
446   });
447 
448   // Build ProtMap, assign addresses.
449   for (auto &KV : BL.segments()) {
450     auto &AG = KV.first;
451     auto &Seg = KV.second;
452 
453     auto &SegAddr = (AG.getMemDeallocPolicy() == MemDeallocPolicy::Standard)
454                         ? NextStandardSegAddr
455                         : NextFinalizeSegAddr;
456 
457     Seg.WorkingMem = jitTargetAddressToPointer<char *>(SegAddr);
458     Seg.Addr = SegAddr;
459 
460     SegAddr += alignTo(Seg.ContentSize + Seg.ZeroFillSize, PageSize);
461   }
462 
463   if (auto Err = BL.apply()) {
464     OnAllocated(std::move(Err));
465     return;
466   }
467 
468   OnAllocated(std::make_unique<IPInFlightAlloc>(*this, G, std::move(BL),
469                                                 std::move(StandardSegsMem),
470                                                 std::move(FinalizeSegsMem)));
471 }
472 
473 void InProcessMemoryManager::deallocate(std::vector<FinalizedAlloc> Allocs,
474                                         OnDeallocatedFunction OnDeallocated) {
475   std::vector<sys::MemoryBlock> StandardSegmentsList;
476   std::vector<std::vector<AllocActionCall>> DeallocActionsList;
477 
478   {
479     std::lock_guard<std::mutex> Lock(FinalizedAllocsMutex);
480     for (auto &Alloc : Allocs) {
481       auto *FA =
482           jitTargetAddressToPointer<FinalizedAllocInfo *>(Alloc.release());
483       StandardSegmentsList.push_back(std::move(FA->StandardSegments));
484       if (!FA->DeallocActions.empty())
485         DeallocActionsList.push_back(std::move(FA->DeallocActions));
486       FA->~FinalizedAllocInfo();
487       FinalizedAllocInfos.Deallocate(FA);
488     }
489   }
490 
491   Error DeallocErr = Error::success();
492 
493   while (!DeallocActionsList.empty()) {
494     auto &DeallocActions = DeallocActionsList.back();
495     auto &StandardSegments = StandardSegmentsList.back();
496 
497     /// Run any deallocate calls.
498     while (!DeallocActions.empty()) {
499       if (auto Err = runAllocAction(DeallocActions.back()))
500         DeallocErr = joinErrors(std::move(DeallocErr), std::move(Err));
501       DeallocActions.pop_back();
502     }
503 
504     /// Release the standard segments slab.
505     if (auto EC = sys::Memory::releaseMappedMemory(StandardSegments))
506       DeallocErr = joinErrors(std::move(DeallocErr), errorCodeToError(EC));
507 
508     DeallocActionsList.pop_back();
509     StandardSegmentsList.pop_back();
510   }
511 
512   OnDeallocated(std::move(DeallocErr));
513 }
514 
515 JITLinkMemoryManager::FinalizedAlloc
516 InProcessMemoryManager::createFinalizedAlloc(
517     sys::MemoryBlock StandardSegments,
518     std::vector<AllocActionCall> DeallocActions) {
519   std::lock_guard<std::mutex> Lock(FinalizedAllocsMutex);
520   auto *FA = FinalizedAllocInfos.Allocate<FinalizedAllocInfo>();
521   new (FA) FinalizedAllocInfo(
522       {std::move(StandardSegments), std::move(DeallocActions)});
523   return FinalizedAlloc(pointerToJITTargetAddress(FA));
524 }
525 
526 } // end namespace jitlink
527 } // end namespace llvm
528