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
22class EPCGenericJITLinkMemoryManager::InFlightAlloc
23 : public jitlink::JITLinkMemoryManager::InFlightAlloc {
24public:
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(x: tpctypes::SegFinalizeRequest{
50 .RAG: KV.first,
51 .Addr: KV.second.Addr,
52 .Size: alignTo(Value: KV.second.ContentSize + KV.second.ZeroFillSize,
53 Align: Parent.EPC.getPageSize()),
54 .Content: {KV.second.WorkingMem, static_cast<size_t>(KV.second.ContentSize)}});
55 }
56
57 // Transfer allocation actions.
58 std::swap(x&: FR.Actions, y&: G.allocActions());
59
60 Parent.EPC.callSPSWrapperAsync<
61 rt::SPSSimpleExecutorMemoryManagerInitializeSignature>(
62 WrapperFnAddr: Parent.SAs.Initialize,
63 SendResult: [OnFinalize = std::move(OnFinalize), AllocAddr = this->AllocAddr](
64 Error SerializationErr,
65 Expected<ExecutorAddr> InitializeKey) mutable {
66 // FIXME: Release abandoned alloc.
67 if (SerializationErr) {
68 cantFail(Err: InitializeKey.takeError());
69 OnFinalize(std::move(SerializationErr));
70 } else if (!InitializeKey)
71 OnFinalize(InitializeKey.takeError());
72 else
73 OnFinalize(FinalizedAlloc(AllocAddr));
74 },
75 Args: Parent.SAs.Allocator, Args: std::move(FR));
76 }
77
78 void abandon(OnAbandonedFunction OnAbandoned) override {
79 // FIXME: Return memory to pool instead.
80 Parent.EPC.callSPSWrapperAsync<
81 rt::SPSSimpleExecutorMemoryManagerReleaseSignature>(
82 WrapperFnAddr: Parent.SAs.Release,
83 SendResult: [OnAbandoned = std::move(OnAbandoned)](Error SerializationErr,
84 Error DeallocateErr) mutable {
85 if (SerializationErr) {
86 cantFail(Err: std::move(DeallocateErr));
87 OnAbandoned(std::move(SerializationErr));
88 } else
89 OnAbandoned(std::move(DeallocateErr));
90 },
91 Args: Parent.SAs.Allocator, Args: ArrayRef<ExecutorAddr>(AllocAddr));
92 }
93
94private:
95 EPCGenericJITLinkMemoryManager &Parent;
96 LinkGraph &G;
97 ExecutorAddr AllocAddr;
98 SegInfoMap Segs;
99};
100
101Expected<std::unique_ptr<EPCGenericJITLinkMemoryManager>>
102EPCGenericJITLinkMemoryManager::Create(
103 JITDylib &JD, rt::SimpleExecutorMemoryManagerSymbolNames SNs) {
104 auto &ES = JD.getExecutionSession();
105 SymbolAddrs SAs;
106 if (auto Err = lookupAndRecordAddrs(
107 ES, K: LookupKind::Static, SearchOrder: makeJITDylibSearchOrder(JDs: {&JD}),
108 Pairs: {
109 {ES.intern(SymName: SNs.AllocatorName), &SAs.Allocator},
110 {ES.intern(SymName: SNs.ReserveName), &SAs.Reserve},
111 {ES.intern(SymName: SNs.InitializeName), &SAs.Initialize},
112 {ES.intern(SymName: SNs.DeinitializeName), &SAs.Deinitialize},
113 {ES.intern(SymName: SNs.ReleaseName), &SAs.Release},
114 }))
115 return Err;
116 return std::make_unique<EPCGenericJITLinkMemoryManager>(
117 args&: ES.getExecutorProcessControl(), args&: SAs);
118}
119
120Expected<std::unique_ptr<EPCGenericJITLinkMemoryManager>>
121EPCGenericJITLinkMemoryManager::Create(
122 ExecutionSession &ES, rt::SimpleExecutorMemoryManagerSymbolNames SNs) {
123 return Create(JD&: ES.getBootstrapJITDylib(), SNs: std::move(SNs));
124}
125
126void EPCGenericJITLinkMemoryManager::allocate(const JITLinkDylib *JD,
127 LinkGraph &G,
128 OnAllocatedFunction OnAllocated) {
129 BasicLayout BL(G);
130
131 auto Pages = BL.getContiguousPageBasedLayoutSizes(PageSize: EPC.getPageSize());
132 if (!Pages)
133 return OnAllocated(Pages.takeError());
134
135 EPC.callSPSWrapperAsync<rt::SPSSimpleExecutorMemoryManagerReserveSignature>(
136 WrapperFnAddr: SAs.Reserve,
137 SendResult: [this, BL = std::move(BL), OnAllocated = std::move(OnAllocated)](
138 Error SerializationErr, Expected<ExecutorAddr> AllocAddr) mutable {
139 if (SerializationErr) {
140 cantFail(Err: AllocAddr.takeError());
141 return OnAllocated(std::move(SerializationErr));
142 }
143 if (!AllocAddr)
144 return OnAllocated(AllocAddr.takeError());
145
146 completeAllocation(AllocAddr: *AllocAddr, BL: std::move(BL), OnAllocated: std::move(OnAllocated));
147 },
148 Args: SAs.Allocator, Args: Pages->total());
149}
150
151void EPCGenericJITLinkMemoryManager::deallocate(
152 std::vector<FinalizedAlloc> Allocs, OnDeallocatedFunction OnDeallocated) {
153 EPC.callSPSWrapperAsync<rt::SPSSimpleExecutorMemoryManagerReleaseSignature>(
154 WrapperFnAddr: SAs.Release,
155 SendResult: [OnDeallocated = std::move(OnDeallocated)](Error SerErr,
156 Error DeallocErr) mutable {
157 if (SerErr) {
158 cantFail(Err: std::move(DeallocErr));
159 OnDeallocated(std::move(SerErr));
160 } else
161 OnDeallocated(std::move(DeallocErr));
162 },
163 Args: SAs.Allocator, Args: Allocs);
164 for (auto &A : Allocs)
165 A.release();
166}
167
168void EPCGenericJITLinkMemoryManager::completeAllocation(
169 ExecutorAddr AllocAddr, BasicLayout BL, OnAllocatedFunction OnAllocated) {
170
171 InFlightAlloc::SegInfoMap SegInfos;
172
173 ExecutorAddr NextSegAddr = AllocAddr;
174 for (auto &KV : BL.segments()) {
175 const auto &AG = KV.first;
176 auto &Seg = KV.second;
177
178 Seg.Addr = NextSegAddr;
179 KV.second.WorkingMem = BL.getGraph().allocateBuffer(Size: Seg.ContentSize).data();
180 NextSegAddr += ExecutorAddrDiff(
181 alignTo(Value: Seg.ContentSize + Seg.ZeroFillSize, Align: EPC.getPageSize()));
182
183 auto &SegInfo = SegInfos[AG];
184 SegInfo.ContentSize = Seg.ContentSize;
185 SegInfo.ZeroFillSize = Seg.ZeroFillSize;
186 SegInfo.Addr = Seg.Addr;
187 SegInfo.WorkingMem = Seg.WorkingMem;
188 }
189
190 if (auto Err = BL.apply())
191 return OnAllocated(std::move(Err));
192
193 OnAllocated(std::make_unique<InFlightAlloc>(args&: *this, args&: BL.getGraph(), args&: AllocAddr,
194 args: std::move(SegInfos)));
195}
196
197} // end namespace orc
198} // end namespace llvm
199