1 | //===----- JITLinkReentryTrampolines.cpp -- JITLink-based trampoline- -----===// |
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/JITLinkReentryTrampolines.h" |
10 | |
11 | #include "llvm/ExecutionEngine/JITLink/aarch64.h" |
12 | #include "llvm/ExecutionEngine/JITLink/x86_64.h" |
13 | #include "llvm/ExecutionEngine/Orc/ObjectLinkingLayer.h" |
14 | |
15 | #include <memory> |
16 | |
17 | #define DEBUG_TYPE "orc" |
18 | |
19 | using namespace llvm; |
20 | using namespace llvm::jitlink; |
21 | |
22 | namespace { |
23 | constexpr StringRef ReentryFnName = "__orc_rt_reenter" ; |
24 | constexpr StringRef ReentrySectionName = "__orc_stubs" ; |
25 | } // namespace |
26 | |
27 | namespace llvm::orc { |
28 | |
29 | class JITLinkReentryTrampolines::TrampolineAddrScraperPlugin |
30 | : public ObjectLinkingLayer::Plugin { |
31 | public: |
32 | void modifyPassConfig(MaterializationResponsibility &MR, |
33 | jitlink::LinkGraph &G, |
34 | jitlink::PassConfiguration &Config) override { |
35 | Config.PreFixupPasses.push_back( |
36 | x: [this](LinkGraph &G) { return recordTrampolineAddrs(G); }); |
37 | } |
38 | |
39 | Error notifyFailed(MaterializationResponsibility &MR) override { |
40 | return Error::success(); |
41 | } |
42 | |
43 | Error notifyRemovingResources(JITDylib &JD, ResourceKey K) override { |
44 | return Error::success(); |
45 | } |
46 | |
47 | void notifyTransferringResources(JITDylib &JD, ResourceKey DstKey, |
48 | ResourceKey SrcKey) override {} |
49 | |
50 | void registerGraph(LinkGraph &G, |
51 | std::shared_ptr<std::vector<ExecutorSymbolDef>> Addrs) { |
52 | std::lock_guard<std::mutex> Lock(M); |
53 | assert(!PendingAddrs.count(&G) && "Duplicate registration" ); |
54 | PendingAddrs[&G] = std::move(Addrs); |
55 | } |
56 | |
57 | Error recordTrampolineAddrs(LinkGraph &G) { |
58 | std::shared_ptr<std::vector<ExecutorSymbolDef>> Addrs; |
59 | { |
60 | std::lock_guard<std::mutex> Lock(M); |
61 | auto I = PendingAddrs.find(Val: &G); |
62 | if (I == PendingAddrs.end()) |
63 | return Error::success(); |
64 | Addrs = std::move(I->second); |
65 | PendingAddrs.erase(I); |
66 | } |
67 | |
68 | auto *Sec = G.findSectionByName(Name: ReentrySectionName); |
69 | assert(Sec && "Reentry graph missing reentry section" ); |
70 | assert(!Sec->empty() && "Reentry graph is empty" ); |
71 | |
72 | for (auto *Sym : Sec->symbols()) |
73 | if (!Sym->hasName()) |
74 | Addrs->push_back(x: {Sym->getAddress(), JITSymbolFlags()}); |
75 | |
76 | return Error::success(); |
77 | } |
78 | |
79 | private: |
80 | std::mutex M; |
81 | DenseMap<LinkGraph *, std::shared_ptr<std::vector<ExecutorSymbolDef>>> |
82 | PendingAddrs; |
83 | }; |
84 | |
85 | Expected<std::unique_ptr<JITLinkReentryTrampolines>> |
86 | JITLinkReentryTrampolines::Create(ObjectLinkingLayer &ObjLinkingLayer) { |
87 | |
88 | EmitTrampolineFn EmitTrampoline; |
89 | |
90 | const auto &TT = ObjLinkingLayer.getExecutionSession().getTargetTriple(); |
91 | switch (TT.getArch()) { |
92 | case Triple::aarch64: |
93 | EmitTrampoline = aarch64::createAnonymousReentryTrampoline; |
94 | break; |
95 | case Triple::x86_64: |
96 | EmitTrampoline = x86_64::createAnonymousReentryTrampoline; |
97 | break; |
98 | default: |
99 | return make_error<StringError>(Args: "JITLinkReentryTrampolines: architecture " + |
100 | TT.getArchName() + " not supported" , |
101 | Args: inconvertibleErrorCode()); |
102 | } |
103 | |
104 | return std::make_unique<JITLinkReentryTrampolines>(args&: ObjLinkingLayer, |
105 | args: std::move(EmitTrampoline)); |
106 | } |
107 | |
108 | JITLinkReentryTrampolines::JITLinkReentryTrampolines( |
109 | ObjectLinkingLayer &ObjLinkingLayer, EmitTrampolineFn EmitTrampoline) |
110 | : ObjLinkingLayer(ObjLinkingLayer), |
111 | EmitTrampoline(std::move(EmitTrampoline)) { |
112 | auto TAS = std::make_shared<TrampolineAddrScraperPlugin>(); |
113 | TrampolineAddrScraper = TAS.get(); |
114 | ObjLinkingLayer.addPlugin(P: std::move(TAS)); |
115 | } |
116 | |
117 | void JITLinkReentryTrampolines::emit(ResourceTrackerSP RT, |
118 | size_t NumTrampolines, |
119 | OnTrampolinesReadyFn OnTrampolinesReady) { |
120 | |
121 | if (NumTrampolines == 0) |
122 | return OnTrampolinesReady(std::vector<ExecutorSymbolDef>()); |
123 | |
124 | JITDylibSP JD(&RT->getJITDylib()); |
125 | auto &ES = ObjLinkingLayer.getExecutionSession(); |
126 | |
127 | auto ReentryGraphSym = |
128 | ES.intern(SymName: ("__orc_reentry_graph_#" + Twine(++ReentryGraphIdx)).str()); |
129 | |
130 | auto G = std::make_unique<jitlink::LinkGraph>( |
131 | args: (*ReentryGraphSym).str(), args: ES.getSymbolStringPool(), args: ES.getTargetTriple(), |
132 | args: SubtargetFeatures(), args&: jitlink::getGenericEdgeKindName); |
133 | |
134 | auto &ReentryFnSym = G->addExternalSymbol(Name: ReentryFnName, Size: 0, IsWeaklyReferenced: false); |
135 | |
136 | auto &ReentrySection = |
137 | G->createSection(Name: ReentrySectionName, Prot: MemProt::Exec | MemProt::Read); |
138 | |
139 | for (size_t I = 0; I != NumTrampolines; ++I) |
140 | EmitTrampoline(*G, ReentrySection, ReentryFnSym).setLive(true); |
141 | |
142 | auto &FirstBlock = **ReentrySection.blocks().begin(); |
143 | G->addDefinedSymbol(Content&: FirstBlock, Offset: 0, Name: *ReentryGraphSym, Size: FirstBlock.getSize(), |
144 | L: Linkage::Strong, S: Scope::SideEffectsOnly, IsCallable: true, IsLive: true); |
145 | |
146 | auto TrampolineAddrs = std::make_shared<std::vector<ExecutorSymbolDef>>(); |
147 | TrampolineAddrScraper->registerGraph(G&: *G, Addrs: TrampolineAddrs); |
148 | |
149 | // Add Graph via object linking layer. |
150 | if (auto Err = ObjLinkingLayer.add(RT: std::move(RT), G: std::move(G))) |
151 | return OnTrampolinesReady(std::move(Err)); |
152 | |
153 | // Trigger graph emission. |
154 | ES.lookup( |
155 | K: LookupKind::Static, SearchOrder: {{JD.get(), JITDylibLookupFlags::MatchAllSymbols}}, |
156 | Symbols: SymbolLookupSet(ReentryGraphSym, |
157 | SymbolLookupFlags::WeaklyReferencedSymbol), |
158 | RequiredState: SymbolState::Ready, |
159 | NotifyComplete: [OnTrampolinesReady = std::move(OnTrampolinesReady), |
160 | TrampolineAddrs = |
161 | std::move(TrampolineAddrs)](Expected<SymbolMap> Result) mutable { |
162 | if (Result) |
163 | OnTrampolinesReady(std::move(*TrampolineAddrs)); |
164 | else |
165 | OnTrampolinesReady(Result.takeError()); |
166 | }, |
167 | RegisterDependencies: NoDependenciesToRegister); |
168 | } |
169 | |
170 | Expected<std::unique_ptr<LazyReexportsManager>> |
171 | createJITLinkLazyReexportsManager(ObjectLinkingLayer &ObjLinkingLayer, |
172 | RedirectableSymbolManager &RSMgr, |
173 | JITDylib &PlatformJD, |
174 | LazyReexportsManager::Listener *L) { |
175 | auto JLT = JITLinkReentryTrampolines::Create(ObjLinkingLayer); |
176 | if (!JLT) |
177 | return JLT.takeError(); |
178 | |
179 | return LazyReexportsManager::Create( |
180 | EmitTrampolines: [JLT = std::move(*JLT)](ResourceTrackerSP RT, size_t NumTrampolines, |
181 | LazyReexportsManager::OnTrampolinesReadyFn |
182 | OnTrampolinesReady) mutable { |
183 | JLT->emit(RT: std::move(RT), NumTrampolines, OnTrampolinesReady: std::move(OnTrampolinesReady)); |
184 | }, |
185 | RSMgr, PlatformJD, L); |
186 | } |
187 | |
188 | } // namespace llvm::orc |
189 | |