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
17using namespace llvm::jitlink;
18
19namespace llvm {
20namespace orc {
21
22const EPCGenericJITLinkMemoryManager::SymbolNames
23 EPCGenericJITLinkMemoryManager::orc_rt_SimpleNativeMemoryMapSPSSymbols = {
24 .AllocatorName: "orc_rt_SimpleNativeMemoryMap_Instance",
25 .ReserveName: "orc_rt_SimpleNativeMemoryMap_reserve_sps_wrapper",
26 .InitializeName: "orc_rt_SimpleNativeMemoryMap_initialize_sps_wrapper",
27 .DeinitializeName: "orc_rt_SimpleNativeMemoryMap_deinitializeMultiple_sps_wrapper",
28 .ReleaseName: "orc_rt_SimpleNativeMemoryMap_releaseMultiple_sps_wrapper",
29};
30
31class EPCGenericJITLinkMemoryManager::InFlightAlloc
32 : public jitlink::JITLinkMemoryManager::InFlightAlloc {
33public:
34
35 // FIXME: The C++98 initializer is an attempt to work around compile failures
36 // due to http://www.open-std.org/jtc1/sc22/wg21/docs/cwg_defects.html#1397.
37 // We should be able to switch this back to member initialization once that
38 // issue is fixed.
39 struct SegInfo {
40 SegInfo() : WorkingMem(nullptr), ContentSize(0), ZeroFillSize(0) {}
41
42 char *WorkingMem;
43 ExecutorAddr Addr;
44 uint64_t ContentSize;
45 uint64_t ZeroFillSize;
46 };
47
48 using SegInfoMap = AllocGroupSmallMap<SegInfo>;
49
50 InFlightAlloc(EPCGenericJITLinkMemoryManager &Parent, LinkGraph &G,
51 ExecutorAddr AllocAddr, SegInfoMap Segs)
52 : Parent(Parent), G(G), AllocAddr(AllocAddr), Segs(std::move(Segs)) {}
53
54 void finalize(OnFinalizedFunction OnFinalize) override {
55 tpctypes::FinalizeRequest FR;
56 for (auto &KV : Segs) {
57 assert(KV.second.ContentSize <= std::numeric_limits<size_t>::max());
58 FR.Segments.push_back(x: tpctypes::SegFinalizeRequest{
59 .RAG: KV.first,
60 .Addr: KV.second.Addr,
61 .Size: alignTo(Value: KV.second.ContentSize + KV.second.ZeroFillSize,
62 Align: Parent.EPC.getPageSize()),
63 .Content: {KV.second.WorkingMem, static_cast<size_t>(KV.second.ContentSize)}});
64 }
65
66 // Transfer allocation actions.
67 std::swap(x&: FR.Actions, y&: G.allocActions());
68
69 Parent.EPC.callSPSWrapperAsync<
70 rt::SPSSimpleExecutorMemoryManagerInitializeSignature>(
71 WrapperFnAddr: Parent.SAs.Initialize,
72 SendResult: [OnFinalize = std::move(OnFinalize), AllocAddr = this->AllocAddr](
73 Error SerializationErr,
74 Expected<ExecutorAddr> InitializeKey) mutable {
75 // FIXME: Release abandoned alloc.
76 if (SerializationErr) {
77 cantFail(Err: InitializeKey.takeError());
78 OnFinalize(std::move(SerializationErr));
79 } else if (!InitializeKey)
80 OnFinalize(InitializeKey.takeError());
81 else
82 OnFinalize(FinalizedAlloc(AllocAddr));
83 },
84 Args: Parent.SAs.Allocator, Args: std::move(FR));
85 }
86
87 void abandon(OnAbandonedFunction OnAbandoned) override {
88 // FIXME: Return memory to pool instead.
89 Parent.EPC.callSPSWrapperAsync<
90 rt::SPSSimpleExecutorMemoryManagerReleaseSignature>(
91 WrapperFnAddr: Parent.SAs.Release,
92 SendResult: [OnAbandoned = std::move(OnAbandoned)](Error SerializationErr,
93 Error DeallocateErr) mutable {
94 if (SerializationErr) {
95 cantFail(Err: std::move(DeallocateErr));
96 OnAbandoned(std::move(SerializationErr));
97 } else
98 OnAbandoned(std::move(DeallocateErr));
99 },
100 Args: Parent.SAs.Allocator, Args: ArrayRef<ExecutorAddr>(AllocAddr));
101 }
102
103private:
104 EPCGenericJITLinkMemoryManager &Parent;
105 LinkGraph &G;
106 ExecutorAddr AllocAddr;
107 SegInfoMap Segs;
108};
109
110Expected<std::unique_ptr<EPCGenericJITLinkMemoryManager>>
111EPCGenericJITLinkMemoryManager::Create(JITDylib &JD, SymbolNames SNs) {
112 auto &ES = JD.getExecutionSession();
113 SymbolAddrs SAs;
114 if (auto Err = lookupAndRecordAddrs(
115 ES, K: LookupKind::Static, SearchOrder: makeJITDylibSearchOrder(JDs: {&JD}),
116 Pairs: {
117 {ES.intern(SymName: SNs.AllocatorName), &SAs.Allocator},
118 {ES.intern(SymName: SNs.ReserveName), &SAs.Reserve},
119 {ES.intern(SymName: SNs.InitializeName), &SAs.Initialize},
120 {ES.intern(SymName: SNs.DeinitializeName), &SAs.Deinitialize},
121 {ES.intern(SymName: SNs.ReleaseName), &SAs.Release},
122 }))
123 return Err;
124 return std::make_unique<EPCGenericJITLinkMemoryManager>(
125 args&: ES.getExecutorProcessControl(), args&: SAs);
126}
127
128Expected<std::unique_ptr<EPCGenericJITLinkMemoryManager>>
129EPCGenericJITLinkMemoryManager::Create(ExecutionSession &ES, SymbolNames SNs) {
130 return Create(JD&: ES.getBootstrapJITDylib(), SNs: std::move(SNs));
131}
132
133void EPCGenericJITLinkMemoryManager::allocate(const JITLinkDylib *JD,
134 LinkGraph &G,
135 OnAllocatedFunction OnAllocated) {
136 BasicLayout BL(G);
137
138 auto Pages = BL.getContiguousPageBasedLayoutSizes(PageSize: EPC.getPageSize());
139 if (!Pages)
140 return OnAllocated(Pages.takeError());
141
142 EPC.callSPSWrapperAsync<rt::SPSSimpleExecutorMemoryManagerReserveSignature>(
143 WrapperFnAddr: SAs.Reserve,
144 SendResult: [this, BL = std::move(BL), OnAllocated = std::move(OnAllocated)](
145 Error SerializationErr, Expected<ExecutorAddr> AllocAddr) mutable {
146 if (SerializationErr) {
147 cantFail(Err: AllocAddr.takeError());
148 return OnAllocated(std::move(SerializationErr));
149 }
150 if (!AllocAddr)
151 return OnAllocated(AllocAddr.takeError());
152
153 completeAllocation(AllocAddr: *AllocAddr, BL: std::move(BL), OnAllocated: std::move(OnAllocated));
154 },
155 Args: SAs.Allocator, Args: Pages->total());
156}
157
158void EPCGenericJITLinkMemoryManager::deallocate(
159 std::vector<FinalizedAlloc> Allocs, OnDeallocatedFunction OnDeallocated) {
160 EPC.callSPSWrapperAsync<rt::SPSSimpleExecutorMemoryManagerReleaseSignature>(
161 WrapperFnAddr: SAs.Release,
162 SendResult: [OnDeallocated = std::move(OnDeallocated)](Error SerErr,
163 Error DeallocErr) mutable {
164 if (SerErr) {
165 cantFail(Err: std::move(DeallocErr));
166 OnDeallocated(std::move(SerErr));
167 } else
168 OnDeallocated(std::move(DeallocErr));
169 },
170 Args: SAs.Allocator, Args: Allocs);
171 for (auto &A : Allocs)
172 A.release();
173}
174
175void EPCGenericJITLinkMemoryManager::completeAllocation(
176 ExecutorAddr AllocAddr, BasicLayout BL, OnAllocatedFunction OnAllocated) {
177
178 InFlightAlloc::SegInfoMap SegInfos;
179
180 ExecutorAddr NextSegAddr = AllocAddr;
181 for (auto &KV : BL.segments()) {
182 const auto &AG = KV.first;
183 auto &Seg = KV.second;
184
185 Seg.Addr = NextSegAddr;
186 KV.second.WorkingMem = BL.getGraph().allocateBuffer(Size: Seg.ContentSize).data();
187 NextSegAddr += ExecutorAddrDiff(
188 alignTo(Value: Seg.ContentSize + Seg.ZeroFillSize, Align: EPC.getPageSize()));
189
190 auto &SegInfo = SegInfos[AG];
191 SegInfo.ContentSize = Seg.ContentSize;
192 SegInfo.ZeroFillSize = Seg.ZeroFillSize;
193 SegInfo.Addr = Seg.Addr;
194 SegInfo.WorkingMem = Seg.WorkingMem;
195 }
196
197 if (auto Err = BL.apply())
198 return OnAllocated(std::move(Err));
199
200 OnAllocated(std::make_unique<InFlightAlloc>(args&: *this, args&: BL.getGraph(), args&: AllocAddr,
201 args: std::move(SegInfos)));
202}
203
204} // end namespace orc
205} // end namespace llvm
206