173 lines · 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-exception6//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 16using namespace llvm::jitlink;17 18namespace llvm {19namespace orc {20 21class EPCGenericJITLinkMemoryManager::InFlightAlloc22 : public jitlink::JITLinkMemoryManager::InFlightAlloc {23public:24 25 // FIXME: The C++98 initializer is an attempt to work around compile failures26 // 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 that28 // 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::SPSSimpleExecutorMemoryManagerInitializeSignature>(61 Parent.SAs.Initialize,62 [OnFinalize = std::move(OnFinalize), AllocAddr = this->AllocAddr](63 Error SerializationErr,64 Expected<ExecutorAddr> InitializeKey) mutable {65 // FIXME: Release abandoned alloc.66 if (SerializationErr) {67 cantFail(InitializeKey.takeError());68 OnFinalize(std::move(SerializationErr));69 } else if (!InitializeKey)70 OnFinalize(InitializeKey.takeError());71 else72 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::SPSSimpleExecutorMemoryManagerReleaseSignature>(81 Parent.SAs.Release,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 } else88 OnAbandoned(std::move(DeallocateErr));89 },90 Parent.SAs.Allocator, ArrayRef<ExecutorAddr>(AllocAddr));91 }92 93private:94 EPCGenericJITLinkMemoryManager &Parent;95 LinkGraph &G;96 ExecutorAddr AllocAddr;97 SegInfoMap Segs;98};99 100void 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 125void EPCGenericJITLinkMemoryManager::deallocate(126 std::vector<FinalizedAlloc> Allocs, OnDeallocatedFunction OnDeallocated) {127 EPC.callSPSWrapperAsync<rt::SPSSimpleExecutorMemoryManagerReleaseSignature>(128 SAs.Release,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 } else135 OnDeallocated(std::move(DeallocErr));136 },137 SAs.Allocator, Allocs);138 for (auto &A : Allocs)139 A.release();140}141 142void 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 orc172} // end namespace llvm173