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 
24 BasicLayout::BasicLayout(LinkGraph &G) : G(G) {
25 
26   for (auto &Sec : G.sections()) {
27     // Skip empty sections.
28     if (empty(Sec.blocks()))
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>
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() == MemDeallocPolicy::Standard)
93       SegsSizes.StandardSegs += SegSize;
94     else
95       SegsSizes.FinalizeSegs += SegSize;
96   }
97 
98   return SegsSizes;
99 }
100 
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 
141 orc::shared::AllocActions &BasicLayout::graphAllocActions() {
142   return G.allocActions();
143 }
144 
145 void SimpleSegmentAlloc::Create(JITLinkMemoryManager &MemMgr,
146                                 const JITLinkDylib *JD, SegmentMap Segments,
147                                 OnCreatedFunction OnCreated) {
148 
149   static_assert(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   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>
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 SimpleSegmentAlloc::getSegInfo(AllocGroup AG) {
217   auto I = ContentBlocks.find(AG);
218   if (I != ContentBlocks.end()) {
219     auto &B = *I->second;
220     return {B.getAddress(), B.getAlreadyMutableContent()};
221   }
222   return {};
223 }
224 
225 SimpleSegmentAlloc::SimpleSegmentAlloc(
226     std::unique_ptr<LinkGraph> G, AllocGroupSmallMap<Block *> ContentBlocks,
227     std::unique_ptr<JITLinkMemoryManager::InFlightAlloc> Alloc)
228     : G(std::move(G)), ContentBlocks(std::move(ContentBlocks)),
229       Alloc(std::move(Alloc)) {}
230 
231 class InProcessMemoryManager::IPInFlightAlloc
232     : public JITLinkMemoryManager::InFlightAlloc {
233 public:
234   IPInFlightAlloc(InProcessMemoryManager &MemMgr, LinkGraph &G, BasicLayout BL,
235                   sys::MemoryBlock StandardSegments,
236                   sys::MemoryBlock FinalizationSegments)
237       : MemMgr(MemMgr), G(G), BL(std::move(BL)),
238         StandardSegments(std::move(StandardSegments)),
239         FinalizationSegments(std::move(FinalizationSegments)) {}
240 
241   void finalize(OnFinalizedFunction OnFinalized) override {
242 
243     // Apply memory protections to all segments.
244     if (auto Err = applyProtections()) {
245       OnFinalized(std::move(Err));
246       return;
247     }
248 
249     // Run finalization actions.
250     auto DeallocActions = runFinalizeActions(G.allocActions());
251     if (!DeallocActions) {
252       OnFinalized(DeallocActions.takeError());
253       return;
254     }
255 
256     // Release the finalize segments slab.
257     if (auto EC = sys::Memory::releaseMappedMemory(FinalizationSegments)) {
258       OnFinalized(errorCodeToError(EC));
259       return;
260     }
261 
262     // Continue with finalized allocation.
263     OnFinalized(MemMgr.createFinalizedAlloc(std::move(StandardSegments),
264                                             std::move(*DeallocActions)));
265   }
266 
267   void abandon(OnAbandonedFunction OnAbandoned) override {
268     Error Err = Error::success();
269     if (auto EC = sys::Memory::releaseMappedMemory(FinalizationSegments))
270       Err = joinErrors(std::move(Err), errorCodeToError(EC));
271     if (auto EC = sys::Memory::releaseMappedMemory(StandardSegments))
272       Err = joinErrors(std::move(Err), errorCodeToError(EC));
273     OnAbandoned(std::move(Err));
274   }
275 
276 private:
277   Error applyProtections() {
278     for (auto &KV : BL.segments()) {
279       const auto &AG = KV.first;
280       auto &Seg = KV.second;
281 
282       auto Prot = toSysMemoryProtectionFlags(AG.getMemProt());
283 
284       uint64_t SegSize =
285           alignTo(Seg.ContentSize + Seg.ZeroFillSize, MemMgr.PageSize);
286       sys::MemoryBlock MB(Seg.WorkingMem, SegSize);
287       if (auto EC = sys::Memory::protectMappedMemory(MB, Prot))
288         return errorCodeToError(EC);
289       if (Prot & sys::Memory::MF_EXEC)
290         sys::Memory::InvalidateInstructionCache(MB.base(), MB.allocatedSize());
291     }
292     return Error::success();
293   }
294 
295   InProcessMemoryManager &MemMgr;
296   LinkGraph &G;
297   BasicLayout BL;
298   sys::MemoryBlock StandardSegments;
299   sys::MemoryBlock FinalizationSegments;
300 };
301 
302 Expected<std::unique_ptr<InProcessMemoryManager>>
303 InProcessMemoryManager::Create() {
304   if (auto PageSize = sys::Process::getPageSize())
305     return std::make_unique<InProcessMemoryManager>(*PageSize);
306   else
307     return PageSize.takeError();
308 }
309 
310 void InProcessMemoryManager::allocate(const JITLinkDylib *JD, LinkGraph &G,
311                                       OnAllocatedFunction OnAllocated) {
312 
313   // FIXME: Just check this once on startup.
314   if (!isPowerOf2_64((uint64_t)PageSize)) {
315     OnAllocated(make_error<StringError>("Page size is not a power of 2",
316                                         inconvertibleErrorCode()));
317     return;
318   }
319 
320   BasicLayout BL(G);
321 
322   /// Scan the request and calculate the group and total sizes.
323   /// Check that segment size is no larger than a page.
324   auto SegsSizes = BL.getContiguousPageBasedLayoutSizes(PageSize);
325   if (!SegsSizes) {
326     OnAllocated(SegsSizes.takeError());
327     return;
328   }
329 
330   /// Check that the total size requested (including zero fill) is not larger
331   /// than a size_t.
332   if (SegsSizes->total() > std::numeric_limits<size_t>::max()) {
333     OnAllocated(make_error<JITLinkError>(
334         "Total requested size " + formatv("{0:x}", SegsSizes->total()) +
335         " for graph " + G.getName() + " exceeds address space"));
336     return;
337   }
338 
339   // Allocate one slab for the whole thing (to make sure everything is
340   // in-range), then partition into standard and finalization blocks.
341   //
342   // FIXME: Make two separate allocations in the future to reduce
343   // fragmentation: finalization segments will usually be a single page, and
344   // standard segments are likely to be more than one page. Where multiple
345   // allocations are in-flight at once (likely) the current approach will leave
346   // a lot of single-page holes.
347   sys::MemoryBlock Slab;
348   sys::MemoryBlock StandardSegsMem;
349   sys::MemoryBlock FinalizeSegsMem;
350   {
351     const sys::Memory::ProtectionFlags ReadWrite =
352         static_cast<sys::Memory::ProtectionFlags>(sys::Memory::MF_READ |
353                                                   sys::Memory::MF_WRITE);
354 
355     std::error_code EC;
356     Slab = sys::Memory::allocateMappedMemory(SegsSizes->total(), nullptr,
357                                              ReadWrite, EC);
358 
359     if (EC) {
360       OnAllocated(errorCodeToError(EC));
361       return;
362     }
363 
364     // Zero-fill the whole slab up-front.
365     memset(Slab.base(), 0, Slab.allocatedSize());
366 
367     StandardSegsMem = {Slab.base(),
368                        static_cast<size_t>(SegsSizes->StandardSegs)};
369     FinalizeSegsMem = {(void *)((char *)Slab.base() + SegsSizes->StandardSegs),
370                        static_cast<size_t>(SegsSizes->FinalizeSegs)};
371   }
372 
373   auto NextStandardSegAddr = orc::ExecutorAddr::fromPtr(StandardSegsMem.base());
374   auto NextFinalizeSegAddr = orc::ExecutorAddr::fromPtr(FinalizeSegsMem.base());
375 
376   LLVM_DEBUG({
377     dbgs() << "InProcessMemoryManager allocated:\n";
378     if (SegsSizes->StandardSegs)
379       dbgs() << formatv("  [ {0:x16} -- {1:x16} ]", NextStandardSegAddr,
380                         NextStandardSegAddr + StandardSegsMem.allocatedSize())
381              << " to stardard segs\n";
382     else
383       dbgs() << "  no standard segs\n";
384     if (SegsSizes->FinalizeSegs)
385       dbgs() << formatv("  [ {0:x16} -- {1:x16} ]", NextFinalizeSegAddr,
386                         NextFinalizeSegAddr + FinalizeSegsMem.allocatedSize())
387              << " to finalize segs\n";
388     else
389       dbgs() << "  no finalize segs\n";
390   });
391 
392   // Build ProtMap, assign addresses.
393   for (auto &KV : BL.segments()) {
394     auto &AG = KV.first;
395     auto &Seg = KV.second;
396 
397     auto &SegAddr = (AG.getMemDeallocPolicy() == MemDeallocPolicy::Standard)
398                         ? NextStandardSegAddr
399                         : NextFinalizeSegAddr;
400 
401     Seg.WorkingMem = SegAddr.toPtr<char *>();
402     Seg.Addr = SegAddr;
403 
404     SegAddr += alignTo(Seg.ContentSize + Seg.ZeroFillSize, PageSize);
405   }
406 
407   if (auto Err = BL.apply()) {
408     OnAllocated(std::move(Err));
409     return;
410   }
411 
412   OnAllocated(std::make_unique<IPInFlightAlloc>(*this, G, std::move(BL),
413                                                 std::move(StandardSegsMem),
414                                                 std::move(FinalizeSegsMem)));
415 }
416 
417 void InProcessMemoryManager::deallocate(std::vector<FinalizedAlloc> Allocs,
418                                         OnDeallocatedFunction OnDeallocated) {
419   std::vector<sys::MemoryBlock> StandardSegmentsList;
420   std::vector<std::vector<orc::shared::WrapperFunctionCall>> DeallocActionsList;
421 
422   {
423     std::lock_guard<std::mutex> Lock(FinalizedAllocsMutex);
424     for (auto &Alloc : Allocs) {
425       auto *FA = Alloc.release().toPtr<FinalizedAllocInfo *>();
426       StandardSegmentsList.push_back(std::move(FA->StandardSegments));
427       if (!FA->DeallocActions.empty())
428         DeallocActionsList.push_back(std::move(FA->DeallocActions));
429       FA->~FinalizedAllocInfo();
430       FinalizedAllocInfos.Deallocate(FA);
431     }
432   }
433 
434   Error DeallocErr = Error::success();
435 
436   while (!DeallocActionsList.empty()) {
437     auto &DeallocActions = DeallocActionsList.back();
438     auto &StandardSegments = StandardSegmentsList.back();
439 
440     /// Run any deallocate calls.
441     while (!DeallocActions.empty()) {
442       if (auto Err = DeallocActions.back().runWithSPSRetErrorMerged())
443         DeallocErr = joinErrors(std::move(DeallocErr), std::move(Err));
444       DeallocActions.pop_back();
445     }
446 
447     /// Release the standard segments slab.
448     if (auto EC = sys::Memory::releaseMappedMemory(StandardSegments))
449       DeallocErr = joinErrors(std::move(DeallocErr), errorCodeToError(EC));
450 
451     DeallocActionsList.pop_back();
452     StandardSegmentsList.pop_back();
453   }
454 
455   OnDeallocated(std::move(DeallocErr));
456 }
457 
458 JITLinkMemoryManager::FinalizedAlloc
459 InProcessMemoryManager::createFinalizedAlloc(
460     sys::MemoryBlock StandardSegments,
461     std::vector<orc::shared::WrapperFunctionCall> DeallocActions) {
462   std::lock_guard<std::mutex> Lock(FinalizedAllocsMutex);
463   auto *FA = FinalizedAllocInfos.Allocate<FinalizedAllocInfo>();
464   new (FA) FinalizedAllocInfo(
465       {std::move(StandardSegments), std::move(DeallocActions)});
466   return FinalizedAlloc(orc::ExecutorAddr::fromPtr(FA));
467 }
468 
469 } // end namespace jitlink
470 } // end namespace llvm
471