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