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::SPSSimpleExecutorMemoryManagerFinalizeSignature>(
62 WrapperFnAddr: Parent.SAs.Finalize,
63 SendResult: [OnFinalize = std::move(OnFinalize), AllocAddr = this->AllocAddr](
64 Error SerializationErr, Error FinalizeErr) mutable {
65 // FIXME: Release abandoned alloc.
66 if (SerializationErr) {
67 cantFail(Err: 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 Args: Parent.SAs.Allocator, Args: 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 WrapperFnAddr: Parent.SAs.Deallocate,
82 SendResult: [OnAbandoned = std::move(OnAbandoned)](Error SerializationErr,
83 Error DeallocateErr) mutable {
84 if (SerializationErr) {
85 cantFail(Err: std::move(DeallocateErr));
86 OnAbandoned(std::move(SerializationErr));
87 } else
88 OnAbandoned(std::move(DeallocateErr));
89 },
90 Args: Parent.SAs.Allocator, Args: 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(PageSize: EPC.getPageSize());
106 if (!Pages)
107 return OnAllocated(Pages.takeError());
108
109 EPC.callSPSWrapperAsync<rt::SPSSimpleExecutorMemoryManagerReserveSignature>(
110 WrapperFnAddr: SAs.Reserve,
111 SendResult: [this, BL = std::move(BL), OnAllocated = std::move(OnAllocated)](
112 Error SerializationErr, Expected<ExecutorAddr> AllocAddr) mutable {
113 if (SerializationErr) {
114 cantFail(Err: AllocAddr.takeError());
115 return OnAllocated(std::move(SerializationErr));
116 }
117 if (!AllocAddr)
118 return OnAllocated(AllocAddr.takeError());
119
120 completeAllocation(AllocAddr: *AllocAddr, BL: std::move(BL), OnAllocated: std::move(OnAllocated));
121 },
122 Args: SAs.Allocator, Args: Pages->total());
123}
124
125void EPCGenericJITLinkMemoryManager::deallocate(
126 std::vector<FinalizedAlloc> Allocs, OnDeallocatedFunction OnDeallocated) {
127 EPC.callSPSWrapperAsync<
128 rt::SPSSimpleExecutorMemoryManagerDeallocateSignature>(
129 WrapperFnAddr: SAs.Deallocate,
130 SendResult: [OnDeallocated = std::move(OnDeallocated)](Error SerErr,
131 Error DeallocErr) mutable {
132 if (SerErr) {
133 cantFail(Err: std::move(DeallocErr));
134 OnDeallocated(std::move(SerErr));
135 } else
136 OnDeallocated(std::move(DeallocErr));
137 },
138 Args: SAs.Allocator, Args: Allocs);
139 for (auto &A : Allocs)
140 A.release();
141}
142
143void 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(Size: Seg.ContentSize).data();
155 NextSegAddr += ExecutorAddrDiff(
156 alignTo(Value: Seg.ContentSize + Seg.ZeroFillSize, Align: EPC.getPageSize()));
157
158 auto &SegInfo = SegInfos[AG];
159 SegInfo.ContentSize = Seg.ContentSize;
160 SegInfo.ZeroFillSize = Seg.ZeroFillSize;
161 SegInfo.Addr = 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>(args&: *this, args&: BL.getGraph(), args&: AllocAddr,
169 args: std::move(SegInfos)));
170}
171
172} // end namespace orc
173} // end namespace llvm
174