llvm/lib/ExecutionEngine/JITLink/JITLinkMemoryManager.cpp

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