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