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