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 173