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
16using namespace llvm::jitlink;
17
18namespace llvm {
19namespace orc {
20
21class EPCGenericJITLinkMemoryManager::InFlightAlloc
22 : public jitlink::JITLinkMemoryManager::InFlightAlloc {
23public:
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(x: tpctypes::SegFinalizeRequest{
49 .RAG: KV.first,
50 .Addr: KV.second.Addr,
51 .Size: alignTo(Value: KV.second.ContentSize + KV.second.ZeroFillSize,
52 Align: Parent.EPC.getPageSize()),
53 .Content: {KV.second.WorkingMem, static_cast<size_t>(KV.second.ContentSize)}});
54 }
55
56 // Transfer allocation actions.
57 std::swap(x&: FR.Actions, y&: G.allocActions());
58
59 Parent.EPC.callSPSWrapperAsync<
60 rt::SPSSimpleExecutorMemoryManagerInitializeSignature>(
61 WrapperFnAddr: Parent.SAs.Initialize,
62 SendResult: [OnFinalize = std::move(OnFinalize), AllocAddr = this->AllocAddr](
63 Error SerializationErr,
64 Expected<ExecutorAddr> InitializeKey) mutable {
65 // FIXME: Release abandoned alloc.
66 if (SerializationErr) {
67 cantFail(Err: InitializeKey.takeError());
68 OnFinalize(std::move(SerializationErr));
69 } else if (!InitializeKey)
70 OnFinalize(InitializeKey.takeError());
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::SPSSimpleExecutorMemoryManagerReleaseSignature>(
81 WrapperFnAddr: Parent.SAs.Release,
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<rt::SPSSimpleExecutorMemoryManagerReleaseSignature>(
128 WrapperFnAddr: SAs.Release,
129 SendResult: [OnDeallocated = std::move(OnDeallocated)](Error SerErr,
130 Error DeallocErr) mutable {
131 if (SerErr) {
132 cantFail(Err: std::move(DeallocErr));
133 OnDeallocated(std::move(SerErr));
134 } else
135 OnDeallocated(std::move(DeallocErr));
136 },
137 Args: SAs.Allocator, Args: Allocs);
138 for (auto &A : Allocs)
139 A.release();
140}
141
142void 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(Size: Seg.ContentSize).data();
154 NextSegAddr += ExecutorAddrDiff(
155 alignTo(Value: Seg.ContentSize + Seg.ZeroFillSize, Align: 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>(args&: *this, args&: BL.getGraph(), args&: AllocAddr,
168 args: std::move(SegInfos)));
169}
170
171} // end namespace orc
172} // end namespace llvm
173