ExecutionEngine/Orc/Speculation.h

0001 //===-- Speculation.h - Speculative Compilation --*- C++ -*-===//
0002 //
0003 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
0004 // See https://llvm.org/LICENSE.txt for license information.
0005 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
0006 //
0007 //===----------------------------------------------------------------------===//
0008 //
0009 // Contains the definition to support speculative compilation when laziness is
0010 // enabled.
0011 //===----------------------------------------------------------------------===//
0012
0013 #ifndef LLVM_EXECUTIONENGINE_ORC_SPECULATION_H
0014 #define LLVM_EXECUTIONENGINE_ORC_SPECULATION_H
0015
0016 #include "llvm/ADT/DenseMap.h"
0017 #include "llvm/ExecutionEngine/Orc/Core.h"
0018 #include "llvm/ExecutionEngine/Orc/DebugUtils.h"
0019 #include "llvm/ExecutionEngine/Orc/IRCompileLayer.h"
0020 #include "llvm/Support/Debug.h"
0021 #include <mutex>
0022 #include <type_traits>
0023 #include <utility>
0024
0025 namespace llvm {
0026 namespace orc {
0027
0028 class Speculator;
0029
0030 // Track the Impls (JITDylib,Symbols) of Symbols while lazy call through
0031 // trampolines are created. Operations are guarded by locks tp ensure that Imap
0032 // stays in consistent state after read/write
0033
0034 class ImplSymbolMap {
0035   friend class Speculator;
0036
0037 public:
0038   using AliaseeDetails = std::pair<SymbolStringPtr, JITDylib *>;
0039   using Alias = SymbolStringPtr;
0040   using ImapTy = DenseMap<Alias, AliaseeDetails>;
0041   void trackImpls(SymbolAliasMap ImplMaps, JITDylib *SrcJD);
0042
0043 private:
0044   // FIX ME: find a right way to distinguish the pre-compile Symbols, and update
0045   // the callsite
0046   std::optional<AliaseeDetails> getImplFor(const SymbolStringPtr &StubSymbol) {
0047     std::lock_guard<std::mutex> Lockit(ConcurrentAccess);
0048     auto Position = Maps.find(StubSymbol);
0049     if (Position != Maps.end())
0050       return Position->getSecond();
0051     else
0052       return std::nullopt;
0053   }
0054
0055   std::mutex ConcurrentAccess;
0056   ImapTy Maps;
0057 };
0058
0059 // Defines Speculator Concept,
0060 class Speculator {
0061 public:
0062   using TargetFAddr = ExecutorAddr;
0063   using FunctionCandidatesMap = DenseMap<SymbolStringPtr, SymbolNameSet>;
0064   using StubAddrLikelies = DenseMap<TargetFAddr, SymbolNameSet>;
0065
0066 private:
0067   void registerSymbolsWithAddr(TargetFAddr ImplAddr,
0068                                SymbolNameSet likelySymbols) {
0069     std::lock_guard<std::mutex> Lockit(ConcurrentAccess);
0070     GlobalSpecMap.insert({ImplAddr, std::move(likelySymbols)});
0071   }
0072
0073   void launchCompile(ExecutorAddr FAddr) {
0074     SymbolNameSet CandidateSet;
0075     // Copy CandidateSet is necessary, to avoid unsynchronized access to
0076     // the datastructure.
0077     {
0078       std::lock_guard<std::mutex> Lockit(ConcurrentAccess);
0079       auto It = GlobalSpecMap.find(FAddr);
0080       if (It == GlobalSpecMap.end())
0081         return;
0082       CandidateSet = It->getSecond();
0083     }
0084
0085     SymbolDependenceMap SpeculativeLookUpImpls;
0086
0087     for (auto &Callee : CandidateSet) {
0088       auto ImplSymbol = AliaseeImplTable.getImplFor(Callee);
0089       // try to distinguish already compiled & library symbols
0090       if (!ImplSymbol)
0091         continue;
0092       const auto &ImplSymbolName = ImplSymbol->first;
0093       JITDylib *ImplJD = ImplSymbol->second;
0094       auto &SymbolsInJD = SpeculativeLookUpImpls[ImplJD];
0095       SymbolsInJD.insert(ImplSymbolName);
0096     }
0097
0098     DEBUG_WITH_TYPE("orc", {
0099       for (auto &I : SpeculativeLookUpImpls) {
0100         llvm::dbgs() << "\n In " << I.first->getName() << " JITDylib ";
0101         for (auto &N : I.second)
0102           llvm::dbgs() << "\n Likely Symbol : " << N;
0103       }
0104     });
0105
0106     // for a given symbol, there may be no symbol qualified for speculatively
0107     // compile try to fix this before jumping to this code if possible.
0108     for (auto &LookupPair : SpeculativeLookUpImpls)
0109       ES.lookup(
0110           LookupKind::Static,
0111           makeJITDylibSearchOrder(LookupPair.first,
0112                                   JITDylibLookupFlags::MatchAllSymbols),
0113           SymbolLookupSet(LookupPair.second), SymbolState::Ready,
0114           [this](Expected<SymbolMap> Result) {
0115             if (auto Err = Result.takeError())
0116               ES.reportError(std::move(Err));
0117           },
0118           NoDependenciesToRegister);
0119   }
0120
0121 public:
0122   Speculator(ImplSymbolMap &Impl, ExecutionSession &ref)
0123       : AliaseeImplTable(Impl), ES(ref), GlobalSpecMap(0) {}
0124   Speculator(const Speculator &) = delete;
0125   Speculator(Speculator &&) = delete;
0126   Speculator &operator=(const Speculator &) = delete;
0127   Speculator &operator=(Speculator &&) = delete;
0128
0129   /// Define symbols for this Speculator object (__orc_speculator) and the
0130   /// speculation runtime entry point symbol (__orc_speculate_for) in the
0131   /// given JITDylib.
0132   Error addSpeculationRuntime(JITDylib &JD, MangleAndInterner &Mangle);
0133
0134   // Speculatively compile likely functions for the given Stub Address.
0135   // destination of __orc_speculate_for jump
0136   void speculateFor(TargetFAddr StubAddr) { launchCompile(StubAddr); }
0137
0138   // FIXME : Register with Stub Address, after JITLink Fix.
0139   void registerSymbols(FunctionCandidatesMap Candidates, JITDylib *JD) {
0140     for (auto &SymPair : Candidates) {
0141       auto Target = SymPair.first;
0142       auto Likely = SymPair.second;
0143
0144       auto OnReadyFixUp = [Likely, Target,
0145                            this](Expected<SymbolMap> ReadySymbol) {
0146         if (ReadySymbol) {
0147           auto RDef = (*ReadySymbol)[Target];
0148           registerSymbolsWithAddr(RDef.getAddress(), std::move(Likely));
0149         } else
0150           this->getES().reportError(ReadySymbol.takeError());
0151       };
0152       // Include non-exported symbols also.
0153       ES.lookup(
0154           LookupKind::Static,
0155           makeJITDylibSearchOrder(JD, JITDylibLookupFlags::MatchAllSymbols),
0156           SymbolLookupSet(Target, SymbolLookupFlags::WeaklyReferencedSymbol),
0157           SymbolState::Ready, OnReadyFixUp, NoDependenciesToRegister);
0158     }
0159   }
0160
0161   ExecutionSession &getES() { return ES; }
0162
0163 private:
0164   static void speculateForEntryPoint(Speculator *Ptr, uint64_t StubId);
0165   std::mutex ConcurrentAccess;
0166   ImplSymbolMap &AliaseeImplTable;
0167   ExecutionSession &ES;
0168   StubAddrLikelies GlobalSpecMap;
0169 };
0170
0171 class IRSpeculationLayer : public IRLayer {
0172 public:
0173   using IRlikiesStrRef =
0174       std::optional<DenseMap<StringRef, DenseSet<StringRef>>>;
0175   using ResultEval = std::function<IRlikiesStrRef(Function &)>;
0176   using TargetAndLikelies = DenseMap<SymbolStringPtr, SymbolNameSet>;
0177
0178   IRSpeculationLayer(ExecutionSession &ES, IRLayer &BaseLayer, Speculator &Spec,
0179                      MangleAndInterner &Mangle, ResultEval Interpreter)
0180       : IRLayer(ES, BaseLayer.getManglingOptions()), NextLayer(BaseLayer),
0181         S(Spec), Mangle(Mangle), QueryAnalysis(Interpreter) {}
0182
0183   void emit(std::unique_ptr<MaterializationResponsibility> R,
0184             ThreadSafeModule TSM) override;
0185
0186 private:
0187   TargetAndLikelies
0188   internToJITSymbols(DenseMap<StringRef, DenseSet<StringRef>> IRNames) {
0189     assert(!IRNames.empty() && "No IRNames received to Intern?");
0190     TargetAndLikelies InternedNames;
0191     for (auto &NamePair : IRNames) {
0192       DenseSet<SymbolStringPtr> TargetJITNames;
0193       for (auto &TargetNames : NamePair.second)
0194         TargetJITNames.insert(Mangle(TargetNames));
0195       InternedNames[Mangle(NamePair.first)] = std::move(TargetJITNames);
0196     }
0197     return InternedNames;
0198   }
0199
0200   IRLayer &NextLayer;
0201   Speculator &S;
0202   MangleAndInterner &Mangle;
0203   ResultEval QueryAnalysis;
0204 };
0205
0206 } // namespace orc
0207 } // namespace llvm
0208
0209 #endif // LLVM_EXECUTIONENGINE_ORC_SPECULATION_H