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