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::SPSSimpleExecutorMemoryManagerFinalizeSignature>(
61 WrapperFnAddr: Parent.SAs.Finalize,
62 SendResult: [OnFinalize = std::move(OnFinalize), AllocAddr = this->AllocAddr](
63 Error SerializationErr, Error FinalizeErr) mutable {
64 // FIXME: Release abandoned alloc.
65 if (SerializationErr) {
66 cantFail(Err: std::move(FinalizeErr));
67 OnFinalize(std::move(SerializationErr));
68 } else if (FinalizeErr)
69 OnFinalize(std::move(FinalizeErr));
70 else
71 OnFinalize(FinalizedAlloc(AllocAddr));
72 },
73 Args: Parent.SAs.Allocator, Args: std::move(FR));
74 }
75
76 void abandon(OnAbandonedFunction OnAbandoned) override {
77 // FIXME: Return memory to pool instead.
78 Parent.EPC.callSPSWrapperAsync<
79 rt::SPSSimpleExecutorMemoryManagerDeallocateSignature>(
80 WrapperFnAddr: Parent.SAs.Deallocate,
81 SendResult: [OnAbandoned = std::move(OnAbandoned)](Error SerializationErr,
82 Error DeallocateErr) mutable {
83 if (SerializationErr) {
84 cantFail(Err: std::move(DeallocateErr));
85 OnAbandoned(std::move(SerializationErr));
86 } else
87 OnAbandoned(std::move(DeallocateErr));
88 },
89 Args: Parent.SAs.Allocator, Args: ArrayRef<ExecutorAddr>(AllocAddr));
90 }
91
92private:
93 EPCGenericJITLinkMemoryManager &Parent;
94 LinkGraph &G;
95 ExecutorAddr AllocAddr;
96 SegInfoMap Segs;
97};
98
99void EPCGenericJITLinkMemoryManager::allocate(const JITLinkDylib *JD,
100 LinkGraph &G,
101 OnAllocatedFunction OnAllocated) {
102 BasicLayout BL(G);
103
104 auto Pages = BL.getContiguousPageBasedLayoutSizes(PageSize: EPC.getPageSize());
105 if (!Pages)
106 return OnAllocated(Pages.takeError());
107
108 EPC.callSPSWrapperAsync<rt::SPSSimpleExecutorMemoryManagerReserveSignature>(
109 WrapperFnAddr: SAs.Reserve,
110 SendResult: [this, BL = std::move(BL), OnAllocated = std::move(OnAllocated)](
111 Error SerializationErr, Expected<ExecutorAddr> AllocAddr) mutable {
112 if (SerializationErr) {
113 cantFail(Err: AllocAddr.takeError());
114 return OnAllocated(std::move(SerializationErr));
115 }
116 if (!AllocAddr)
117 return OnAllocated(AllocAddr.takeError());
118
119 completeAllocation(AllocAddr: *AllocAddr, BL: std::move(BL), OnAllocated: std::move(OnAllocated));
120 },
121 Args: SAs.Allocator, Args: Pages->total());
122}
123
124void EPCGenericJITLinkMemoryManager::deallocate(
125 std::vector<FinalizedAlloc> Allocs, OnDeallocatedFunction OnDeallocated) {
126 EPC.callSPSWrapperAsync<
127 rt::SPSSimpleExecutorMemoryManagerDeallocateSignature>(
128 WrapperFnAddr: SAs.Deallocate,
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