llvm/lib/ExecutionEngine/Orc/EPCGenericJITLinkMemoryManager.cpp

   1 //===---- EPCGenericJITLinkMemoryManager.cpp -- Mem management via EPC ----===//
   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/EPCGenericJITLinkMemoryManager.h"
  10
  11 #include "llvm/ExecutionEngine/JITLink/JITLink.h"
  12 #include "llvm/ExecutionEngine/Orc/LookupAndRecordAddrs.h"
  13 #include "llvm/ExecutionEngine/Orc/Shared/OrcRTBridge.h"
  14
  15 #include <limits>
  16
  17 using namespace llvm::jitlink;
  18
  19 namespace llvm {
  20 namespace orc {
  21
  22 class EPCGenericJITLinkMemoryManager::InFlightAlloc
  23     : public jitlink::JITLinkMemoryManager::InFlightAlloc {
  24 public:
  25
  26   // FIXME: The C++98 initializer is an attempt to work around compile failures
  27   // due to http://www.open-std.org/jtc1/sc22/wg21/docs/cwg_defects.html#1397.
  28   // We should be able to switch this back to member initialization once that
  29   // issue is fixed.
  30   struct SegInfo {
  31     SegInfo() : WorkingMem(nullptr), ContentSize(0), ZeroFillSize(0) {}
  32
  33     char *WorkingMem;
  34     ExecutorAddr Addr;
  35     uint64_t ContentSize;
  36     uint64_t ZeroFillSize;
  37   };
  38
  39   using SegInfoMap = AllocGroupSmallMap<SegInfo>;
  40
  41   InFlightAlloc(EPCGenericJITLinkMemoryManager &Parent, LinkGraph &G,
  42                 ExecutorAddr AllocAddr, SegInfoMap Segs)
  43       : Parent(Parent), G(G), AllocAddr(AllocAddr), Segs(std::move(Segs)) {}
  44
  45   void finalize(OnFinalizedFunction OnFinalize) override {
  46     tpctypes::FinalizeRequest FR;
  47     for (auto &KV : Segs) {
  48       assert(KV.second.ContentSize <= std::numeric_limits<size_t>::max());
  49       FR.Segments.push_back(tpctypes::SegFinalizeRequest{
  50           KV.first,
  51           KV.second.Addr,
  52           alignTo(KV.second.ContentSize + KV.second.ZeroFillSize,
  53                   Parent.EPC.getPageSize()),
  54           {KV.second.WorkingMem, static_cast<size_t>(KV.second.ContentSize)}});
  55     }
  56
  57     // Transfer allocation actions.
  58     std::swap(FR.Actions, G.allocActions());
  59
  60     Parent.EPC.callSPSWrapperAsync<
  61         rt::SPSSimpleExecutorMemoryManagerFinalizeSignature>(
  62         Parent.SAs.Finalize,
  63         [OnFinalize = std::move(OnFinalize), AllocAddr = this->AllocAddr](
  64             Error SerializationErr, Error FinalizeErr) mutable {
  65           // FIXME: Release abandoned alloc.
  66           if (SerializationErr) {
  67             cantFail(std::move(FinalizeErr));
  68             OnFinalize(std::move(SerializationErr));
  69           } else if (FinalizeErr)
  70             OnFinalize(std::move(FinalizeErr));
  71           else
  72             OnFinalize(FinalizedAlloc(AllocAddr));
  73         },
  74         Parent.SAs.Allocator, std::move(FR));
  75   }
  76
  77   void abandon(OnAbandonedFunction OnAbandoned) override {
  78     // FIXME: Return memory to pool instead.
  79     Parent.EPC.callSPSWrapperAsync<
  80         rt::SPSSimpleExecutorMemoryManagerDeallocateSignature>(
  81         Parent.SAs.Deallocate,
  82         [OnAbandoned = std::move(OnAbandoned)](Error SerializationErr,
  83                                                Error DeallocateErr) mutable {
  84           if (SerializationErr) {
  85             cantFail(std::move(DeallocateErr));
  86             OnAbandoned(std::move(SerializationErr));
  87           } else
  88             OnAbandoned(std::move(DeallocateErr));
  89         },
  90         Parent.SAs.Allocator, ArrayRef<ExecutorAddr>(AllocAddr));
  91   }
  92
  93 private:
  94   EPCGenericJITLinkMemoryManager &Parent;
  95   LinkGraph &G;
  96   ExecutorAddr AllocAddr;
  97   SegInfoMap Segs;
  98 };
  99
 100 void EPCGenericJITLinkMemoryManager::allocate(const JITLinkDylib *JD,
 101                                               LinkGraph &G,
 102                                               OnAllocatedFunction OnAllocated) {
 103   BasicLayout BL(G);
 104
 105   auto Pages = BL.getContiguousPageBasedLayoutSizes(EPC.getPageSize());
 106   if (!Pages)
 107     return OnAllocated(Pages.takeError());
 108
 109   EPC.callSPSWrapperAsync<rt::SPSSimpleExecutorMemoryManagerReserveSignature>(
 110       SAs.Reserve,
 111       [this, BL = std::move(BL), OnAllocated = std::move(OnAllocated)](
 112           Error SerializationErr, Expected<ExecutorAddr> AllocAddr) mutable {
 113         if (SerializationErr) {
 114           cantFail(AllocAddr.takeError());
 115           return OnAllocated(std::move(SerializationErr));
 116         }
 117         if (!AllocAddr)
 118           return OnAllocated(AllocAddr.takeError());
 119
 120         completeAllocation(*AllocAddr, std::move(BL), std::move(OnAllocated));
 121       },
 122       SAs.Allocator, Pages->total());
 123 }
 124
 125 void EPCGenericJITLinkMemoryManager::deallocate(
 126     std::vector<FinalizedAlloc> Allocs, OnDeallocatedFunction OnDeallocated) {
 127   EPC.callSPSWrapperAsync<
 128       rt::SPSSimpleExecutorMemoryManagerDeallocateSignature>(
 129       SAs.Deallocate,
 130       [OnDeallocated = std::move(OnDeallocated)](Error SerErr,
 131                                                  Error DeallocErr) mutable {
 132         if (SerErr) {
 133           cantFail(std::move(DeallocErr));
 134           OnDeallocated(std::move(SerErr));
 135         } else
 136           OnDeallocated(std::move(DeallocErr));
 137       },
 138       SAs.Allocator, Allocs);
 139   for (auto &A : Allocs)
 140     A.release();
 141 }
 142
 143 void EPCGenericJITLinkMemoryManager::completeAllocation(
 144     ExecutorAddr AllocAddr, BasicLayout BL, OnAllocatedFunction OnAllocated) {
 145
 146   InFlightAlloc::SegInfoMap SegInfos;
 147
 148   ExecutorAddr NextSegAddr = AllocAddr;
 149   for (auto &KV : BL.segments()) {
 150     const auto &AG = KV.first;
 151     auto &Seg = KV.second;
 152
 153     Seg.Addr = NextSegAddr;
 154     KV.second.WorkingMem = BL.getGraph().allocateBuffer(Seg.ContentSize).data();
 155     NextSegAddr += ExecutorAddrDiff(
 156         alignTo(Seg.ContentSize + Seg.ZeroFillSize, EPC.getPageSize()));
 157
 158     auto &SegInfo = SegInfos[AG];
 159     SegInfo.ContentSize = Seg.ContentSize;
 160     SegInfo.ZeroFillSize = Seg.ZeroFillSize;
 161     SegInfo.Addr = Seg.Addr;
 162     SegInfo.WorkingMem = Seg.WorkingMem;
 163   }
 164
 165   if (auto Err = BL.apply())
 166     return OnAllocated(std::move(Err));
 167
 168   OnAllocated(std::make_unique<InFlightAlloc>(*this, BL.getGraph(), AllocAddr,
 169                                               std::move(SegInfos)));
 170 }
 171
 172 } // end namespace orc
 173 } // end namespace llvm