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(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 = JITTargetAddress;
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({ImplAddr, std::move(likelySymbols)});
71   }
72 
73   void launchCompile(JITTargetAddress 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(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(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(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           LookupKind::Static,
111           makeJITDylibSearchOrder(LookupPair.first,
112                                   JITDylibLookupFlags::MatchAllSymbols),
113           SymbolLookupSet(LookupPair.second), SymbolState::Ready,
114           [this](Expected<SymbolMap> Result) {
115             if (auto Err = Result.takeError())
116               ES.reportError(std::move(Err));
117           },
118           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(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 RAddr = (*ReadySymbol)[Target].getAddress();
148           registerSymbolsWithAddr(RAddr, std::move(Likely));
149         } else
150           this->getES().reportError(ReadySymbol.takeError());
151       };
152       // Include non-exported symbols also.
153       ES.lookup(
154           LookupKind::Static,
155           makeJITDylibSearchOrder(JD, JITDylibLookupFlags::MatchAllSymbols),
156           SymbolLookupSet(Target, SymbolLookupFlags::WeaklyReferencedSymbol),
157           SymbolState::Ready, OnReadyFixUp, 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(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