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