xref: /freebsd/contrib/llvm-project/llvm/lib/ExecutionEngine/Orc/EPCGenericJITLinkMemoryManager.cpp (revision 5e801ac66d24704442eba426ed13c3effb8a34e7)
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 
17 using namespace llvm::jitlink;
18 
19 namespace llvm {
20 namespace orc {
21 
22 class EPCGenericJITLinkMemoryManager::InFlightAlloc
23     : public jitlink::JITLinkMemoryManager::InFlightAlloc {
24 public:
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(tpctypes::SegFinalizeRequest{
50           tpctypes::toWireProtectionFlags(
51               toSysMemoryProtectionFlags(KV.first.getMemProt())),
52           KV.second.Addr,
53           alignTo(KV.second.ContentSize + KV.second.ZeroFillSize,
54                   Parent.EPC.getPageSize()),
55           {KV.second.WorkingMem, static_cast<size_t>(KV.second.ContentSize)}});
56     }
57 
58     // Transfer allocation actions.
59     // FIXME: Merge JITLink and ORC SupportFunctionCall and Action list types,
60     //        turn this into a std::swap.
61     FR.Actions.reserve(G.allocActions().size());
62     for (auto &ActPair : G.allocActions())
63       FR.Actions.push_back({{ExecutorAddr(ActPair.Finalize.FnAddr),
64                              {ExecutorAddr(ActPair.Finalize.CtxAddr),
65                               ExecutorAddrDiff(ActPair.Finalize.CtxSize)}},
66                             {ExecutorAddr(ActPair.Dealloc.FnAddr),
67                              {ExecutorAddr(ActPair.Dealloc.CtxAddr),
68                               ExecutorAddrDiff(ActPair.Dealloc.CtxSize)}}});
69     G.allocActions().clear();
70 
71     Parent.EPC.callSPSWrapperAsync<
72         rt::SPSSimpleExecutorMemoryManagerFinalizeSignature>(
73         Parent.SAs.Finalize,
74         [OnFinalize = std::move(OnFinalize), AllocAddr = this->AllocAddr](
75             Error SerializationErr, Error FinalizeErr) mutable {
76           // FIXME: Release abandoned alloc.
77           if (SerializationErr) {
78             cantFail(std::move(FinalizeErr));
79             OnFinalize(std::move(SerializationErr));
80           } else if (FinalizeErr)
81             OnFinalize(std::move(FinalizeErr));
82           else
83             OnFinalize(FinalizedAlloc(AllocAddr.getValue()));
84         },
85         Parent.SAs.Allocator, std::move(FR));
86   }
87 
88   void abandon(OnAbandonedFunction OnAbandoned) override {
89     // FIXME: Return memory to pool instead.
90     Parent.EPC.callSPSWrapperAsync<
91         rt::SPSSimpleExecutorMemoryManagerDeallocateSignature>(
92         Parent.SAs.Deallocate,
93         [OnAbandoned = std::move(OnAbandoned)](Error SerializationErr,
94                                                Error DeallocateErr) mutable {
95           if (SerializationErr) {
96             cantFail(std::move(DeallocateErr));
97             OnAbandoned(std::move(SerializationErr));
98           } else
99             OnAbandoned(std::move(DeallocateErr));
100         },
101         Parent.SAs.Allocator, ArrayRef<ExecutorAddr>(AllocAddr));
102   }
103 
104 private:
105   EPCGenericJITLinkMemoryManager &Parent;
106   LinkGraph &G;
107   ExecutorAddr AllocAddr;
108   SegInfoMap Segs;
109 };
110 
111 void EPCGenericJITLinkMemoryManager::allocate(const JITLinkDylib *JD,
112                                               LinkGraph &G,
113                                               OnAllocatedFunction OnAllocated) {
114   BasicLayout BL(G);
115 
116   auto Pages = BL.getContiguousPageBasedLayoutSizes(EPC.getPageSize());
117   if (!Pages)
118     return OnAllocated(Pages.takeError());
119 
120   EPC.callSPSWrapperAsync<rt::SPSSimpleExecutorMemoryManagerReserveSignature>(
121       SAs.Reserve,
122       [this, BL = std::move(BL), OnAllocated = std::move(OnAllocated)](
123           Error SerializationErr, Expected<ExecutorAddr> AllocAddr) mutable {
124         if (SerializationErr) {
125           cantFail(AllocAddr.takeError());
126           return OnAllocated(std::move(SerializationErr));
127         }
128         if (!AllocAddr)
129           return OnAllocated(AllocAddr.takeError());
130 
131         completeAllocation(*AllocAddr, std::move(BL), std::move(OnAllocated));
132       },
133       SAs.Allocator, Pages->total());
134 }
135 
136 void EPCGenericJITLinkMemoryManager::deallocate(
137     std::vector<FinalizedAlloc> Allocs, OnDeallocatedFunction OnDeallocated) {
138   EPC.callSPSWrapperAsync<
139       rt::SPSSimpleExecutorMemoryManagerDeallocateSignature>(
140       SAs.Deallocate,
141       [OnDeallocated = std::move(OnDeallocated)](Error SerErr,
142                                                  Error DeallocErr) mutable {
143         if (SerErr) {
144           cantFail(std::move(DeallocErr));
145           OnDeallocated(std::move(SerErr));
146         } else
147           OnDeallocated(std::move(DeallocErr));
148       },
149       SAs.Allocator, Allocs);
150   for (auto &A : Allocs)
151     A.release();
152 }
153 
154 void EPCGenericJITLinkMemoryManager::completeAllocation(
155     ExecutorAddr AllocAddr, BasicLayout BL, OnAllocatedFunction OnAllocated) {
156 
157   InFlightAlloc::SegInfoMap SegInfos;
158 
159   ExecutorAddr NextSegAddr = AllocAddr;
160   for (auto &KV : BL.segments()) {
161     const auto &AG = KV.first;
162     auto &Seg = KV.second;
163 
164     Seg.Addr = NextSegAddr.getValue();
165     KV.second.WorkingMem = BL.getGraph().allocateBuffer(Seg.ContentSize).data();
166     NextSegAddr += ExecutorAddrDiff(
167         alignTo(Seg.ContentSize + Seg.ZeroFillSize, EPC.getPageSize()));
168 
169     auto &SegInfo = SegInfos[AG];
170     SegInfo.ContentSize = Seg.ContentSize;
171     SegInfo.ZeroFillSize = Seg.ZeroFillSize;
172     SegInfo.Addr = ExecutorAddr(Seg.Addr);
173     SegInfo.WorkingMem = Seg.WorkingMem;
174   }
175 
176   if (auto Err = BL.apply())
177     return OnAllocated(std::move(Err));
178 
179   OnAllocated(std::make_unique<InFlightAlloc>(*this, BL.getGraph(), AllocAddr,
180                                               std::move(SegInfos)));
181 }
182 
183 } // end namespace orc
184 } // end namespace llvm
185