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 = ExecutorAddr(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 174