1//=== MapperJITLinkMemoryManager.cpp - Memory management with MemoryMapper ===//
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/MapperJITLinkMemoryManager.h"
10
11#include "llvm/ADT/STLExtras.h"
12#include "llvm/ExecutionEngine/JITLink/JITLink.h"
13#include "llvm/Support/Process.h"
14
15using namespace llvm::jitlink;
16
17namespace llvm {
18namespace orc {
19
20class MapperJITLinkMemoryManager::InFlightAlloc
21 : public JITLinkMemoryManager::InFlightAlloc {
22public:
23 InFlightAlloc(MapperJITLinkMemoryManager &Parent, LinkGraph &G,
24 ExecutorAddr AllocAddr,
25 std::vector<MemoryMapper::AllocInfo::SegInfo> Segs)
26 : Parent(Parent), G(G), AllocAddr(AllocAddr), Segs(std::move(Segs)) {}
27
28 void finalize(OnFinalizedFunction OnFinalize) override {
29 MemoryMapper::AllocInfo AI;
30 AI.MappingBase = AllocAddr;
31
32 std::swap(x&: AI.Segments, y&: Segs);
33 std::swap(x&: AI.Actions, y&: G.allocActions());
34
35 Parent.Mapper->initialize(AI, OnInitialized: [OnFinalize = std::move(OnFinalize)](
36 Expected<ExecutorAddr> Result) mutable {
37 if (!Result) {
38 OnFinalize(Result.takeError());
39 return;
40 }
41
42 OnFinalize(FinalizedAlloc(*Result));
43 });
44 }
45
46 void abandon(OnAbandonedFunction OnFinalize) override {
47 Parent.Mapper->release(Reservations: {AllocAddr}, OnRelease: std::move(OnFinalize));
48 }
49
50private:
51 MapperJITLinkMemoryManager &Parent;
52 LinkGraph &G;
53 ExecutorAddr AllocAddr;
54 std::vector<MemoryMapper::AllocInfo::SegInfo> Segs;
55};
56
57MapperJITLinkMemoryManager::MapperJITLinkMemoryManager(
58 size_t ReservationGranularity, std::unique_ptr<MemoryMapper> Mapper)
59 : ReservationUnits(ReservationGranularity), AvailableMemory(AMAllocator),
60 Mapper(std::move(Mapper)) {}
61
62void MapperJITLinkMemoryManager::allocate(const JITLinkDylib *JD, LinkGraph &G,
63 OnAllocatedFunction OnAllocated) {
64 BasicLayout BL(G);
65
66 // find required address space
67 auto SegsSizes = BL.getContiguousPageBasedLayoutSizes(PageSize: Mapper->getPageSize());
68 if (!SegsSizes) {
69 OnAllocated(SegsSizes.takeError());
70 return;
71 }
72
73 auto TotalSize = SegsSizes->total();
74
75 auto CompleteAllocation = [this, &G, BL = std::move(BL),
76 OnAllocated = std::move(OnAllocated)](
77 Expected<ExecutorAddrRange> Result) mutable {
78 if (!Result) {
79 Mutex.unlock();
80 return OnAllocated(Result.takeError());
81 }
82
83 auto NextSegAddr = Result->Start;
84
85 std::vector<MemoryMapper::AllocInfo::SegInfo> SegInfos;
86
87 for (auto &KV : BL.segments()) {
88 auto &AG = KV.first;
89 auto &Seg = KV.second;
90
91 auto TotalSize = Seg.ContentSize + Seg.ZeroFillSize;
92
93 Seg.Addr = NextSegAddr;
94 Seg.WorkingMem = Mapper->prepare(Addr: NextSegAddr, ContentSize: TotalSize);
95
96 NextSegAddr += alignTo(Value: TotalSize, Align: Mapper->getPageSize());
97
98 MemoryMapper::AllocInfo::SegInfo SI;
99 SI.Offset = Seg.Addr - Result->Start;
100 SI.ContentSize = Seg.ContentSize;
101 SI.ZeroFillSize = Seg.ZeroFillSize;
102 SI.AG = AG;
103 SI.WorkingMem = Seg.WorkingMem;
104
105 SegInfos.push_back(x: SI);
106 }
107
108 UsedMemory.insert(KV: {Result->Start, NextSegAddr - Result->Start});
109
110 if (NextSegAddr < Result->End) {
111 // Save the remaining memory for reuse in next allocation(s)
112 AvailableMemory.insert(a: NextSegAddr, b: Result->End - 1, y: true);
113 }
114 Mutex.unlock();
115
116 if (auto Err = BL.apply()) {
117 OnAllocated(std::move(Err));
118 return;
119 }
120
121 OnAllocated(std::make_unique<InFlightAlloc>(args&: *this, args&: G, args&: Result->Start,
122 args: std::move(SegInfos)));
123 };
124
125 Mutex.lock();
126
127 // find an already reserved range that is large enough
128 ExecutorAddrRange SelectedRange{};
129
130 for (AvailableMemoryMap::iterator It = AvailableMemory.begin();
131 It != AvailableMemory.end(); It++) {
132 if (It.stop() - It.start() + 1 >= TotalSize) {
133 SelectedRange = ExecutorAddrRange(It.start(), It.stop() + 1);
134 It.erase();
135 break;
136 }
137 }
138
139 if (SelectedRange.empty()) { // no already reserved range was found
140 auto TotalAllocation = alignTo(Value: TotalSize, Align: ReservationUnits);
141 Mapper->reserve(NumBytes: TotalAllocation, OnReserved: std::move(CompleteAllocation));
142 } else {
143 CompleteAllocation(SelectedRange);
144 }
145}
146
147void MapperJITLinkMemoryManager::deallocate(
148 std::vector<FinalizedAlloc> Allocs, OnDeallocatedFunction OnDeallocated) {
149 std::vector<ExecutorAddr> Bases;
150 Bases.reserve(n: Allocs.size());
151 for (auto &FA : Allocs) {
152 ExecutorAddr Addr = FA.getAddress();
153 Bases.push_back(x: Addr);
154 }
155
156 Mapper->deinitialize(Allocations: Bases, OnDeInitialized: [this, Allocs = std::move(Allocs),
157 OnDeallocated = std::move(OnDeallocated)](
158 llvm::Error Err) mutable {
159 // TODO: How should we treat memory that we fail to deinitialize?
160 // We're currently bailing out and treating it as "burned" -- should we
161 // require that a failure to deinitialize still reset the memory so that
162 // we can reclaim it?
163 if (Err) {
164 for (auto &FA : Allocs)
165 FA.release();
166 OnDeallocated(std::move(Err));
167 return;
168 }
169
170 {
171 std::lock_guard<std::mutex> Lock(Mutex);
172
173 for (auto &FA : Allocs) {
174 ExecutorAddr Addr = FA.getAddress();
175 ExecutorAddrDiff Size = UsedMemory[Addr];
176
177 UsedMemory.erase(Val: Addr);
178 AvailableMemory.insert(a: Addr, b: Addr + Size - 1, y: true);
179
180 FA.release();
181 }
182 }
183
184 OnDeallocated(Error::success());
185 });
186}
187
188} // end namespace orc
189} // end namespace llvm
190