llvm/lib/ExecutionEngine/Orc/MapperJITLinkMemoryManager.cpp

   1 //=== MapperJITLinkMemoryManager.cpp - Memory management with MemoryMapper ===//
   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/Orc/MapperJITLinkMemoryManager.h"
  10
  11 #include "llvm/ADT/STLExtras.h"
  12 #include "llvm/ExecutionEngine/JITLink/JITLink.h"
  13 #include "llvm/Support/Process.h"
  14
  15 using namespace llvm::jitlink;
  16
  17 namespace llvm {
  18 namespace orc {
  19
  20 class MapperJITLinkMemoryManager::InFlightAlloc
  21     : public JITLinkMemoryManager::InFlightAlloc {
  22 public:
  23   InFlightAlloc(MapperJITLinkMemoryManager &Parent, LinkGraph &G,
  24                 ExecutorAddr AllocAddr,
  25                 std::vector<MemoryMapper::AllocInfo::SegInfo> Segs)
  26       : Parent(Parent), G(G), AllocAddr(AllocAddr), Segs(std::move(Segs)) {}
  27
  28   void finalize(OnFinalizedFunction OnFinalize) override {
  29     MemoryMapper::AllocInfo AI;
  30     AI.MappingBase = AllocAddr;
  31
  32     std::swap(AI.Segments, Segs);
  33     std::swap(AI.Actions, G.allocActions());
  34
  35     Parent.Mapper->initialize(AI, [OnFinalize = std::move(OnFinalize)](
  36                                       Expected<ExecutorAddr> Result) mutable {
  37       if (!Result) {
  38         OnFinalize(Result.takeError());
  39         return;
  40       }
  41
  42       OnFinalize(FinalizedAlloc(*Result));
  43     });
  44   }
  45
  46   void abandon(OnAbandonedFunction OnFinalize) override {
  47     Parent.Mapper->release({AllocAddr}, std::move(OnFinalize));
  48   }
  49
  50 private:
  51   MapperJITLinkMemoryManager &Parent;
  52   LinkGraph &G;
  53   ExecutorAddr AllocAddr;
  54   std::vector<MemoryMapper::AllocInfo::SegInfo> Segs;
  55 };
  56
  57 MapperJITLinkMemoryManager::MapperJITLinkMemoryManager(
  58     size_t ReservationGranularity, std::unique_ptr<MemoryMapper> Mapper)
  59     : ReservationUnits(ReservationGranularity), AvailableMemory(AMAllocator),
  60       Mapper(std::move(Mapper)) {}
  61
  62 void MapperJITLinkMemoryManager::allocate(const JITLinkDylib *JD, LinkGraph &G,
  63                                           OnAllocatedFunction OnAllocated) {
  64   BasicLayout BL(G);
  65
  66   // find required address space
  67   auto SegsSizes = BL.getContiguousPageBasedLayoutSizes(Mapper->getPageSize());
  68   if (!SegsSizes) {
  69     OnAllocated(SegsSizes.takeError());
  70     return;
  71   }
  72
  73   auto TotalSize = SegsSizes->total();
  74
  75   auto CompleteAllocation = [this, &G, BL = std::move(BL),
  76                              OnAllocated = std::move(OnAllocated)](
  77                                 Expected<ExecutorAddrRange> Result) mutable {
  78     if (!Result) {
  79       Mutex.unlock();
  80       return OnAllocated(Result.takeError());
  81     }
  82
  83     auto NextSegAddr = Result->Start;
  84
  85     std::vector<MemoryMapper::AllocInfo::SegInfo> SegInfos;
  86
  87     for (auto &KV : BL.segments()) {
  88       auto &AG = KV.first;
  89       auto &Seg = KV.second;
  90
  91       auto TotalSize = Seg.ContentSize + Seg.ZeroFillSize;
  92
  93       Seg.Addr = NextSegAddr;
  94       Seg.WorkingMem = Mapper->prepare(NextSegAddr, TotalSize);
  95
  96       NextSegAddr += alignTo(TotalSize, Mapper->getPageSize());
  97
  98       MemoryMapper::AllocInfo::SegInfo SI;
  99       SI.Offset = Seg.Addr - Result->Start;
 100       SI.ContentSize = Seg.ContentSize;
 101       SI.ZeroFillSize = Seg.ZeroFillSize;
 102       SI.AG = AG;
 103       SI.WorkingMem = Seg.WorkingMem;
 104
 105       SegInfos.push_back(SI);
 106     }
 107
 108     UsedMemory.insert({Result->Start, NextSegAddr - Result->Start});
 109
 110     if (NextSegAddr < Result->End) {
 111       // Save the remaining memory for reuse in next allocation(s)
 112       AvailableMemory.insert(NextSegAddr, Result->End - 1, true);
 113     }
 114     Mutex.unlock();
 115
 116     if (auto Err = BL.apply()) {
 117       OnAllocated(std::move(Err));
 118       return;
 119     }
 120
 121     OnAllocated(std::make_unique<InFlightAlloc>(*this, G, Result->Start,
 122                                                 std::move(SegInfos)));
 123   };
 124
 125   Mutex.lock();
 126
 127   // find an already reserved range that is large enough
 128   ExecutorAddrRange SelectedRange{};
 129
 130   for (AvailableMemoryMap::iterator It = AvailableMemory.begin();
 131        It != AvailableMemory.end(); It++) {
 132     if (It.stop() - It.start() + 1 >= TotalSize) {
 133       SelectedRange = ExecutorAddrRange(It.start(), It.stop() + 1);
 134       It.erase();
 135       break;
 136     }
 137   }
 138
 139   if (SelectedRange.empty()) { // no already reserved range was found
 140     auto TotalAllocation = alignTo(TotalSize, ReservationUnits);
 141     Mapper->reserve(TotalAllocation, std::move(CompleteAllocation));
 142   } else {
 143     CompleteAllocation(SelectedRange);
 144   }
 145 }
 146
 147 void MapperJITLinkMemoryManager::deallocate(
 148     std::vector<FinalizedAlloc> Allocs, OnDeallocatedFunction OnDeallocated) {
 149   std::vector<ExecutorAddr> Bases;
 150   Bases.reserve(Allocs.size());
 151   for (auto &FA : Allocs) {
 152     ExecutorAddr Addr = FA.getAddress();
 153     Bases.push_back(Addr);
 154   }
 155
 156   Mapper->deinitialize(Bases, [this, Allocs = std::move(Allocs),
 157                                OnDeallocated = std::move(OnDeallocated)](
 158                                   llvm::Error Err) mutable {
 159     // TODO: How should we treat memory that we fail to deinitialize?
 160     // We're currently bailing out and treating it as "burned" -- should we
 161     // require that a failure to deinitialize still reset the memory so that
 162     // we can reclaim it?
 163     if (Err) {
 164       for (auto &FA : Allocs)
 165         FA.release();
 166       OnDeallocated(std::move(Err));
 167       return;
 168     }
 169
 170     {
 171       std::lock_guard<std::mutex> Lock(Mutex);
 172
 173       for (auto &FA : Allocs) {
 174         ExecutorAddr Addr = FA.getAddress();
 175         ExecutorAddrDiff Size = UsedMemory[Addr];
 176
 177         UsedMemory.erase(Addr);
 178         AvailableMemory.insert(Addr, Addr + Size - 1, true);
 179
 180         FA.release();
 181       }
 182     }
 183
 184     OnDeallocated(Error::success());
 185   });
 186 }
 187
 188 } // end namespace orc
 189 } // end namespace llvm