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
13#include <limits>
14
15using namespace llvm::jitlink;
16
17namespace llvm {
18namespace orc {
19
20class EPCGenericJITLinkMemoryManager::InFlightAlloc
21 : public jitlink::JITLinkMemoryManager::InFlightAlloc {
22public:
23
24 // FIXME: The C++98 initializer is an attempt to work around compile failures
25 // due to http://www.open-std.org/jtc1/sc22/wg21/docs/cwg_defects.html#1397.
26 // We should be able to switch this back to member initialization once that
27 // issue is fixed.
28 struct SegInfo {
29 SegInfo() : WorkingMem(nullptr), ContentSize(0), ZeroFillSize(0) {}
30
31 char *WorkingMem;
32 ExecutorAddr Addr;
33 uint64_t ContentSize;
34 uint64_t ZeroFillSize;
35 };
36
37 using SegInfoMap = AllocGroupSmallMap<SegInfo>;
38
39 InFlightAlloc(EPCGenericJITLinkMemoryManager &Parent, LinkGraph &G,
40 ExecutorAddr AllocAddr, SegInfoMap Segs)
41 : Parent(Parent), G(G), AllocAddr(AllocAddr), Segs(std::move(Segs)) {}
42
43 void finalize(OnFinalizedFunction OnFinalize) override {
44 tpctypes::FinalizeRequest FR;
45 for (auto &KV : Segs) {
46 assert(KV.second.ContentSize <= std::numeric_limits<size_t>::max());
47 FR.Segments.push_back(x: tpctypes::SegFinalizeRequest{
48 .RAG: KV.first,
49 .Addr: KV.second.Addr,
50 .Size: alignTo(Value: KV.second.ContentSize + KV.second.ZeroFillSize,
51 Align: Parent.ES.getExecutorProcessControl().getPageSize()),
52 .Content: {KV.second.WorkingMem, static_cast<size_t>(KV.second.ContentSize)}});
53 }
54
55 // Transfer allocation actions.
56 std::swap(x&: FR.Actions, y&: G.allocActions());
57
58 Parent.B.Initialize(
59 [OnFinalize = std::move(OnFinalize), AllocAddr = this->AllocAddr](
60 Expected<ExecutorAddr> InitializeKey) mutable {
61 // FIXME: Release abandoned alloc.
62 if (!InitializeKey)
63 return OnFinalize(InitializeKey.takeError());
64 OnFinalize(FinalizedAlloc(AllocAddr));
65 },
66 Parent.ES, Parent.B.Instance, FR);
67 }
68
69 void abandon(OnAbandonedFunction OnAbandoned) override {
70 // FIXME: Return memory to pool instead.
71 Parent.B.Release(std::move(OnAbandoned), Parent.ES, Parent.B.Instance,
72 ArrayRef<ExecutorAddr>(AllocAddr));
73 }
74
75private:
76 EPCGenericJITLinkMemoryManager &Parent;
77 LinkGraph &G;
78 ExecutorAddr AllocAddr;
79 SegInfoMap Segs;
80};
81
82void EPCGenericJITLinkMemoryManager::allocate(const JITLinkDylib *JD,
83 LinkGraph &G,
84 OnAllocatedFunction OnAllocated) {
85 BasicLayout BL(G);
86
87 auto Pages = BL.getContiguousPageBasedLayoutSizes(
88 PageSize: ES.getExecutorProcessControl().getPageSize());
89 if (!Pages)
90 return OnAllocated(Pages.takeError());
91
92 B.Reserve(
93 [this, BL = std::move(BL), OnAllocated = std::move(OnAllocated)](
94 Expected<ExecutorAddr> AllocAddr) mutable {
95 if (!AllocAddr)
96 return OnAllocated(AllocAddr.takeError());
97 completeAllocation(AllocAddr: *AllocAddr, BL: std::move(BL), OnAllocated: std::move(OnAllocated));
98 },
99 ES, B.Instance, Pages->total());
100}
101
102void EPCGenericJITLinkMemoryManager::deallocate(
103 std::vector<FinalizedAlloc> Allocs, OnDeallocatedFunction OnDeallocated) {
104 std::vector<ExecutorAddr> Bases;
105 Bases.reserve(n: Allocs.size());
106 for (auto &A : Allocs)
107 Bases.push_back(x: ExecutorAddr(A.getAddress()));
108 B.Release(std::move(OnDeallocated), ES, B.Instance, Bases);
109 for (auto &A : Allocs)
110 A.release();
111}
112
113void EPCGenericJITLinkMemoryManager::completeAllocation(
114 ExecutorAddr AllocAddr, BasicLayout BL, OnAllocatedFunction OnAllocated) {
115
116 InFlightAlloc::SegInfoMap SegInfos;
117
118 ExecutorAddr NextSegAddr = AllocAddr;
119 for (auto &KV : BL.segments()) {
120 const auto &AG = KV.first;
121 auto &Seg = KV.second;
122
123 Seg.Addr = NextSegAddr;
124 KV.second.WorkingMem = BL.getGraph().allocateBuffer(Size: Seg.ContentSize).data();
125 NextSegAddr +=
126 ExecutorAddrDiff(alignTo(Value: Seg.ContentSize + Seg.ZeroFillSize,
127 Align: ES.getExecutorProcessControl().getPageSize()));
128
129 auto &SegInfo = SegInfos[AG];
130 SegInfo.ContentSize = Seg.ContentSize;
131 SegInfo.ZeroFillSize = Seg.ZeroFillSize;
132 SegInfo.Addr = Seg.Addr;
133 SegInfo.WorkingMem = Seg.WorkingMem;
134 }
135
136 if (auto Err = BL.apply())
137 return OnAllocated(std::move(Err));
138
139 OnAllocated(std::make_unique<InFlightAlloc>(args&: *this, args&: BL.getGraph(), args&: AllocAddr,
140 args: std::move(SegInfos)));
141}
142
143} // end namespace orc
144} // end namespace llvm
145