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