| 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 | |