1 | //===-- Speculation.h - Speculative Compilation --*- C++ -*-===// |
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 | // Contains the definition to support speculative compilation when laziness is |
10 | // enabled. |
11 | //===----------------------------------------------------------------------===// |
12 | |
13 | #ifndef LLVM_EXECUTIONENGINE_ORC_SPECULATION_H |
14 | #define LLVM_EXECUTIONENGINE_ORC_SPECULATION_H |
15 | |
16 | #include "llvm/ADT/DenseMap.h" |
17 | #include "llvm/ExecutionEngine/Orc/Core.h" |
18 | #include "llvm/ExecutionEngine/Orc/DebugUtils.h" |
19 | #include "llvm/ExecutionEngine/Orc/IRCompileLayer.h" |
20 | #include "llvm/Support/Debug.h" |
21 | #include <mutex> |
22 | #include <type_traits> |
23 | #include <utility> |
24 | |
25 | namespace llvm { |
26 | namespace orc { |
27 | |
28 | class Speculator; |
29 | |
30 | // Track the Impls (JITDylib,Symbols) of Symbols while lazy call through |
31 | // trampolines are created. Operations are guarded by locks tp ensure that Imap |
32 | // stays in consistent state after read/write |
33 | |
34 | class ImplSymbolMap { |
35 | friend class Speculator; |
36 | |
37 | public: |
38 | using AliaseeDetails = std::pair<SymbolStringPtr, JITDylib *>; |
39 | using Alias = SymbolStringPtr; |
40 | using ImapTy = DenseMap<Alias, AliaseeDetails>; |
41 | void trackImpls(SymbolAliasMap ImplMaps, JITDylib *SrcJD); |
42 | |
43 | private: |
44 | // FIX ME: find a right way to distinguish the pre-compile Symbols, and update |
45 | // the callsite |
46 | std::optional<AliaseeDetails> getImplFor(const SymbolStringPtr &StubSymbol) { |
47 | std::lock_guard<std::mutex> Lockit(ConcurrentAccess); |
48 | auto Position = Maps.find(Val: StubSymbol); |
49 | if (Position != Maps.end()) |
50 | return Position->getSecond(); |
51 | else |
52 | return std::nullopt; |
53 | } |
54 | |
55 | std::mutex ConcurrentAccess; |
56 | ImapTy Maps; |
57 | }; |
58 | |
59 | // Defines Speculator Concept, |
60 | class Speculator { |
61 | public: |
62 | using TargetFAddr = ExecutorAddr; |
63 | using FunctionCandidatesMap = DenseMap<SymbolStringPtr, SymbolNameSet>; |
64 | using StubAddrLikelies = DenseMap<TargetFAddr, SymbolNameSet>; |
65 | |
66 | private: |
67 | void registerSymbolsWithAddr(TargetFAddr ImplAddr, |
68 | SymbolNameSet likelySymbols) { |
69 | std::lock_guard<std::mutex> Lockit(ConcurrentAccess); |
70 | GlobalSpecMap.insert(KV: {ImplAddr, std::move(likelySymbols)}); |
71 | } |
72 | |
73 | void launchCompile(ExecutorAddr FAddr) { |
74 | SymbolNameSet CandidateSet; |
75 | // Copy CandidateSet is necessary, to avoid unsynchronized access to |
76 | // the datastructure. |
77 | { |
78 | std::lock_guard<std::mutex> Lockit(ConcurrentAccess); |
79 | auto It = GlobalSpecMap.find(Val: FAddr); |
80 | if (It == GlobalSpecMap.end()) |
81 | return; |
82 | CandidateSet = It->getSecond(); |
83 | } |
84 | |
85 | SymbolDependenceMap SpeculativeLookUpImpls; |
86 | |
87 | for (auto &Callee : CandidateSet) { |
88 | auto ImplSymbol = AliaseeImplTable.getImplFor(StubSymbol: Callee); |
89 | // try to distinguish already compiled & library symbols |
90 | if (!ImplSymbol) |
91 | continue; |
92 | const auto &ImplSymbolName = ImplSymbol->first; |
93 | JITDylib *ImplJD = ImplSymbol->second; |
94 | auto &SymbolsInJD = SpeculativeLookUpImpls[ImplJD]; |
95 | SymbolsInJD.insert(V: ImplSymbolName); |
96 | } |
97 | |
98 | DEBUG_WITH_TYPE("orc" , { |
99 | for (auto &I : SpeculativeLookUpImpls) { |
100 | llvm::dbgs() << "\n In " << I.first->getName() << " JITDylib " ; |
101 | for (auto &N : I.second) |
102 | llvm::dbgs() << "\n Likely Symbol : " << N; |
103 | } |
104 | }); |
105 | |
106 | // for a given symbol, there may be no symbol qualified for speculatively |
107 | // compile try to fix this before jumping to this code if possible. |
108 | for (auto &LookupPair : SpeculativeLookUpImpls) |
109 | ES.lookup( |
110 | K: LookupKind::Static, |
111 | SearchOrder: makeJITDylibSearchOrder(JDs: LookupPair.first, |
112 | Flags: JITDylibLookupFlags::MatchAllSymbols), |
113 | Symbols: SymbolLookupSet(LookupPair.second), RequiredState: SymbolState::Ready, |
114 | NotifyComplete: [this](Expected<SymbolMap> Result) { |
115 | if (auto Err = Result.takeError()) |
116 | ES.reportError(Err: std::move(Err)); |
117 | }, |
118 | RegisterDependencies: NoDependenciesToRegister); |
119 | } |
120 | |
121 | public: |
122 | Speculator(ImplSymbolMap &Impl, ExecutionSession &ref) |
123 | : AliaseeImplTable(Impl), ES(ref), GlobalSpecMap(0) {} |
124 | Speculator(const Speculator &) = delete; |
125 | Speculator(Speculator &&) = delete; |
126 | Speculator &operator=(const Speculator &) = delete; |
127 | Speculator &operator=(Speculator &&) = delete; |
128 | |
129 | /// Define symbols for this Speculator object (__orc_speculator) and the |
130 | /// speculation runtime entry point symbol (__orc_speculate_for) in the |
131 | /// given JITDylib. |
132 | Error addSpeculationRuntime(JITDylib &JD, MangleAndInterner &Mangle); |
133 | |
134 | // Speculatively compile likely functions for the given Stub Address. |
135 | // destination of __orc_speculate_for jump |
136 | void speculateFor(TargetFAddr StubAddr) { launchCompile(FAddr: StubAddr); } |
137 | |
138 | // FIXME : Register with Stub Address, after JITLink Fix. |
139 | void registerSymbols(FunctionCandidatesMap Candidates, JITDylib *JD) { |
140 | for (auto &SymPair : Candidates) { |
141 | auto Target = SymPair.first; |
142 | auto Likely = SymPair.second; |
143 | |
144 | auto OnReadyFixUp = [Likely, Target, |
145 | this](Expected<SymbolMap> ReadySymbol) { |
146 | if (ReadySymbol) { |
147 | auto RDef = (*ReadySymbol)[Target]; |
148 | registerSymbolsWithAddr(ImplAddr: RDef.getAddress(), likelySymbols: std::move(Likely)); |
149 | } else |
150 | this->getES().reportError(Err: ReadySymbol.takeError()); |
151 | }; |
152 | // Include non-exported symbols also. |
153 | ES.lookup( |
154 | K: LookupKind::Static, |
155 | SearchOrder: makeJITDylibSearchOrder(JDs: JD, Flags: JITDylibLookupFlags::MatchAllSymbols), |
156 | Symbols: SymbolLookupSet(Target, SymbolLookupFlags::WeaklyReferencedSymbol), |
157 | RequiredState: SymbolState::Ready, NotifyComplete: OnReadyFixUp, RegisterDependencies: NoDependenciesToRegister); |
158 | } |
159 | } |
160 | |
161 | ExecutionSession &getES() { return ES; } |
162 | |
163 | private: |
164 | static void speculateForEntryPoint(Speculator *Ptr, uint64_t StubId); |
165 | std::mutex ConcurrentAccess; |
166 | ImplSymbolMap &AliaseeImplTable; |
167 | ExecutionSession &ES; |
168 | StubAddrLikelies GlobalSpecMap; |
169 | }; |
170 | |
171 | class IRSpeculationLayer : public IRLayer { |
172 | public: |
173 | using IRlikiesStrRef = |
174 | std::optional<DenseMap<StringRef, DenseSet<StringRef>>>; |
175 | using ResultEval = std::function<IRlikiesStrRef(Function &)>; |
176 | using TargetAndLikelies = DenseMap<SymbolStringPtr, SymbolNameSet>; |
177 | |
178 | IRSpeculationLayer(ExecutionSession &ES, IRLayer &BaseLayer, Speculator &Spec, |
179 | MangleAndInterner &Mangle, ResultEval Interpreter) |
180 | : IRLayer(ES, BaseLayer.getManglingOptions()), NextLayer(BaseLayer), |
181 | S(Spec), Mangle(Mangle), QueryAnalysis(Interpreter) {} |
182 | |
183 | void emit(std::unique_ptr<MaterializationResponsibility> R, |
184 | ThreadSafeModule TSM) override; |
185 | |
186 | private: |
187 | TargetAndLikelies |
188 | internToJITSymbols(DenseMap<StringRef, DenseSet<StringRef>> IRNames) { |
189 | assert(!IRNames.empty() && "No IRNames received to Intern?" ); |
190 | TargetAndLikelies InternedNames; |
191 | for (auto &NamePair : IRNames) { |
192 | DenseSet<SymbolStringPtr> TargetJITNames; |
193 | for (auto &TargetNames : NamePair.second) |
194 | TargetJITNames.insert(V: Mangle(TargetNames)); |
195 | InternedNames[Mangle(NamePair.first)] = std::move(TargetJITNames); |
196 | } |
197 | return InternedNames; |
198 | } |
199 | |
200 | IRLayer &NextLayer; |
201 | Speculator &S; |
202 | MangleAndInterner &Mangle; |
203 | ResultEval QueryAnalysis; |
204 | }; |
205 | |
206 | } // namespace orc |
207 | } // namespace llvm |
208 | |
209 | #endif // LLVM_EXECUTIONENGINE_ORC_SPECULATION_H |
210 | |