include/llvm/ExecutionEngine/Orc/Speculation.h

   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", for (auto &I
 104                                 : 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     // for a given symbol, there may be no symbol qualified for speculatively
 111     // compile try to fix this before jumping to this code if possible.
 112     for (auto &LookupPair : SpeculativeLookUpImpls)
 113       ES.lookup(JITDylibSearchList({{LookupPair.first, true}}),
 114                 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(JITDylibSearchList({{JD, true}}), SymbolNameSet({Target}),
 155                 SymbolState::Ready, OnReadyFixUp, NoDependenciesToRegister);
 156     }
 157   }
 158
 159   ExecutionSession &getES() { return ES; }
 160
 161 private:
 162   static void speculateForEntryPoint(Speculator *Ptr, uint64_t StubId);
 163   std::mutex ConcurrentAccess;
 164   ImplSymbolMap &AliaseeImplTable;
 165   ExecutionSession &ES;
 166   StubAddrLikelies GlobalSpecMap;
 167 };
 168
 169 class IRSpeculationLayer : public IRLayer {
 170 public:
 171   using IRlikiesStrRef = Optional<DenseMap<StringRef, DenseSet<StringRef>>>;
 172   using ResultEval = std::function<IRlikiesStrRef(Function &)>;
 173   using TargetAndLikelies = DenseMap<SymbolStringPtr, SymbolNameSet>;
 174
 175   IRSpeculationLayer(ExecutionSession &ES, IRCompileLayer &BaseLayer,
 176                      Speculator &Spec, MangleAndInterner &Mangle,
 177                      ResultEval Interpreter)
 178       : IRLayer(ES), NextLayer(BaseLayer), S(Spec), Mangle(Mangle),
 179         QueryAnalysis(Interpreter) {}
 180
 181   void emit(MaterializationResponsibility R, ThreadSafeModule TSM);
 182
 183 private:
 184   TargetAndLikelies
 185   internToJITSymbols(DenseMap<StringRef, DenseSet<StringRef>> IRNames) {
 186     assert(!IRNames.empty() && "No IRNames received to Intern?");
 187     TargetAndLikelies InternedNames;
 188     DenseSet<SymbolStringPtr> TargetJITNames;
 189     for (auto &NamePair : IRNames) {
 190       for (auto &TargetNames : NamePair.second)
 191         TargetJITNames.insert(Mangle(TargetNames));
 192
 193       InternedNames[Mangle(NamePair.first)] = std::move(TargetJITNames);
 194     }
 195     return InternedNames;
 196   }
 197
 198   IRCompileLayer &NextLayer;
 199   Speculator &S;
 200   MangleAndInterner &Mangle;
 201   ResultEval QueryAnalysis;
 202 };
 203
 204 } // namespace orc
 205 } // namespace llvm
 206
 207 #endif // LLVM_EXECUTIONENGINE_ORC_SPECULATION_H