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