xref: /llvm-project/llvm/lib/ExecutionEngine/Orc/JITLinkReentryTrampolines.cpp (revision 4eaff6c58ae2f130ac8d63cf2c87bbb483114876)
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         [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(&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(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({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>("JITLinkReentryTrampolines: architecture " +
100 				   TT.getArchName() + " not supported",
101                                    inconvertibleErrorCode());
102   }
103 
104   return std::make_unique<JITLinkReentryTrampolines>(ObjLinkingLayer,
105                                                      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(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(("__orc_reentry_graph_#" + Twine(++ReentryGraphIdx)).str());
129 
130   auto G = std::make_unique<jitlink::LinkGraph>(
131       (*ReentryGraphSym).str(), ES.getSymbolStringPool(), ES.getTargetTriple(),
132       SubtargetFeatures(), jitlink::getGenericEdgeKindName);
133 
134   auto &ReentryFnSym = G->addExternalSymbol(ReentryFnName, 0, false);
135 
136   auto &ReentrySection =
137       G->createSection(ReentrySectionName, 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(FirstBlock, 0, *ReentryGraphSym, FirstBlock.getSize(),
144                       Linkage::Strong, Scope::SideEffectsOnly, true, true);
145 
146   auto TrampolineAddrs = std::make_shared<std::vector<ExecutorSymbolDef>>();
147   TrampolineAddrScraper->registerGraph(*G, TrampolineAddrs);
148 
149   // Add Graph via object linking layer.
150   if (auto Err = ObjLinkingLayer.add(std::move(RT), std::move(G)))
151     return OnTrampolinesReady(std::move(Err));
152 
153   // Trigger graph emission.
154   ES.lookup(
155       LookupKind::Static, {{JD.get(), JITDylibLookupFlags::MatchAllSymbols}},
156       SymbolLookupSet(ReentryGraphSym,
157                       SymbolLookupFlags::WeaklyReferencedSymbol),
158       SymbolState::Ready,
159       [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       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       [JLT = std::move(*JLT)](ResourceTrackerSP RT, size_t NumTrampolines,
181                               LazyReexportsManager::OnTrampolinesReadyFn
182                                   OnTrampolinesReady) mutable {
183         JLT->emit(std::move(RT), NumTrampolines, std::move(OnTrampolinesReady));
184       },
185       RSMgr, PlatformJD, L);
186 }
187 
188 } // namespace llvm::orc
189