[ORC] WIP Speculative compilation

[pree-jackie]

Summary

Initial Patch to Support Speculative Compilation in ORC v2 APIs. Idea is to replicate the mapping provided by the lazyreexports utility (SymbolAliasMap) to find which Callable Symbol the stubs symbols lazy reexports and where the Callable Symbol lives (Impl JITDylibs).

Instrument IR with calls to the run time function orc_speculate_for at the entry block of Functions which calls other functions. At run time the call to __orc_speculate_for triggers the compilation of functions and update the pointer of the lazy reexports stubs.

This Patch introduces two maps : Imap and SpeculationMap,

Imap - replicates the source symbols and impl JITDylibs for each alias.
Speculation Map - lives inside the Speculator, which keeps track of for each function address what are the likely functions that the function may call at runtime.

Simple flow is at runtime Speculator queries the Imap to find the source symbols and impl JITDylibs.

Things to do:

• Write run time function implementation.
• Make Stub Pointer Updation through Notify Resolved
• Test Cases

Current Implementation finds likely functions one hop at a time, in my opinion is it not a drawback.

[hfinkel]

Can you please update this patch with full context?

I assume that you're posting this for initial feedback, but a patch description explaining the usage model would be helpful.

llvm/include/llvm/ExecutionEngine/Orc/Speculation.h

55

Don't use assert(0 && ..., use llvm_unreachable.

[hfinkel]

I assume that you're posting this for initial feedback, but a patch description explaining the usage model would be helpful.

Ignore this; you just did ;)

[lhames]

This is an excellent start!

llvm/include/llvm/ExecutionEngine/Orc/Speculation.h

46–47

This should be:

assert(It.second && "Source Entities are already tracked for this symbol?");

55

LLVM's coding style does not use braces for single conditional statements, so it should be:

if (Position != Maps.end())
  return Position->getSecond();

rather than

if (Position != Maps.end()) {
  return Position->getSecond();
}

More generally, this function can be reduced to:

ImplPair getImplFor(SymbolStringPtr StubAddr) {
  auto Position = Maps.find(StubAddr);
  assert(Position != Maps.end() &&
         "Source Entities are not tracked for Symbol?");
  return Position->getSecond();
}

78–84

Applying the coding conventions (no braces for single statements, folding assert conditions into the assert):

auto It = GlobalSpecMap.find(FAddr);
assert(It != GlobalSpecMap.end() && "launching compiles for Unexpected function Address?");
for (auto &Pos : It->getSecond())
  auto SourceEntities = AliaseeImplTable.getImplFor(Pos);

"Pos" is not a very descriptive name here. "Callee" ?

Likewise for SourceEntities. "ImplSymbol"?

SourceEntities is currently unused, but all you should need to do is issue an asynchronous lookup for it with a no-op 'OnComplete' callback. OnComplete takes an error, so you'll have to deal with that, but the usual way is just to log it to the ExecutionSession:

auto ImplSymbol = AliaseeImplTable.getImplFor(Callee);
const auto &ImplSymbolName = ImplSymbol.first;
auto *ImplJD = *ImplSymbol.second;
ES.lookup({ImplJD, true}, {ImplSymbolName},
          SymbolState::Ready,
          [](Expected<SymbolMap> Result) {
            if (auto Err = Result.takeError())
            ES.reportError(std::move(Err));
          });

98–99

"Walking" functions describes how you build this set, but not what it is. CalleeSet might be a better name for now?

llvm/lib/ExecutionEngine/Orc/Speculation.cpp

57

What line is this comment in reference to?

83

Variable names should be capitalized, so this should be 'Likely'.

106–108

No braces needed here either, since the body is a single statement.

[lhames]

llvm/lib/ExecutionEngine/Orc/Speculation.cpp

91

If I have understood your patch correctly, you could test Dave Blaikie's proof-of-concept speculation idea by issuing a call to:

SpecMap.speculateFor(RAddr);

You would need to remove the std::move from RAddr on the previous line. JITTargetAddresses are cheap to copy so that's fine to get rid of.

To test that this is really working you could compile a test module:

void a() {}
void b() {}

int main(int argc, char *argv[]) {
  if (argc > 2)
    a();
  else
    b();
  return 0;
}

If your simple speculator is working you should see modules for both a() and b() pass through your speculation layer when main is called, even though only one will ever be executed.

[pree-jackie]

Hi Lang, thank you for the comments.
I'm reading those and I'll address those shortly.

Regarding, testing I have already done it in single and multi-threaded JIT mode with one and more JITDylibs. The expected behavior was achieved in my test cases. Only thing that blocks this patch to support speculative compiles without run time call is issuing a ES.lookup, I will add that here shortly.

[pree-jackie]

llvm/lib/ExecutionEngine/Orc/Speculation.cpp

57

Converting the function definition address to the unsigned int of 64 bits. Which are used as the key in Speculator to lookup likely functions.

91

Not exactly, registering the symbols with speculator doesn't mean in any way that the Impl Symbols are recorded in the Imap. To replicate the lazyreexports symbol mapping we use the same technique as earlier.

Triggering a call at run time is easy, because we can guarantee that when we call the function, all the external references by the function are already stub materialized. So when we call the function at run time through __orc_speculate_for, Imap has the Impl Symbols.

But triggering speculative compiles without the run time function, requires 2 look ups
1 ) Lookup likely Symbols to Stub Materialize and handover the remaining to the linking layer.
2 ) Lookup based on the Imap Entries of the likely functions for speculative compiles.

I have done it, it is working and I have to write tests.

Actually, it work along with the same sequence of actions that ORC v2 do.

[pree-jackie]

Fixed comments,

[pree-jackie]

llvm/lib/ExecutionEngine/Orc/Speculation.cpp

91

This becomes meaning less, since Callback changed from Resolved to Ready

[pree-jackie]

Added __orc_speculate_for , and orc_speculator. Remaining are test-cases, once test case is added I think it is ready to go, without major changes

[pree-jackie]

Additions in this Change Set

1. Functionality for LLVM and Custom Analysis Integration with Speculator.

2. Speculator Class is refactored to let the client with different program representations re-use the common functionality.

3. Template Parameter is Added to IRSpeculator class to identify the top-level analysis the clients want to run.

[pree-jackie]

Ping :)

[hfinkel]

Can you please upload a patch with full context?

[pree-jackie]

In D63378#1574241, @hfinkel wrote:

Can you please upload a patch with full context?

I have updated

[pree-jackie]

lib/ExecutionEngine/Orc/Speculation.cpp

32 ↗

(On Diff #208493)

Resolve warning in Release build

[lhames]

Hi Praveen,

I'm not sure the IRSpeculator is necessary. My instinct would be to just have the Speculator class as the global repository of estimates of the form "execution point -> { candidates to speculatively compile }" (i.e. your GlobalSpecMap). Then the IRSpeculationLayer should take a reference to the Speculator and use a custom FunctionPass to populates the estimates.

I've attached an example of a hand rolled pass pipeline that should give you access to statically calculated block frequency estimates:

block-freq-pass-example.cpp2 KBDownload

• Lang.

[pree-jackie]

In D63378#1575196, @lhames wrote:

Hi Praveen,

I'm not sure the IRSpeculator is necessary. My instinct would be to just have the Speculator class as the global repository of estimates of the form "execution point -> { candidates to speculatively compile }" (i.e. your GlobalSpecMap). Then the IRSpeculationLayer should take a reference to the Speculator and use a custom FunctionPass to populates the estimates.

I've attached an example of a hand rolled pass pipeline that should give you access to statically calculated block frequency estimates:

block-freq-pass-example.cpp2 KBDownload

• Lang.

I think it is really a matter of separation, that is where we analysis the module and retrieve the results that means whether in a layer or a speculator object, If we take over the layer approach we have to make the Pass to be aware of speculator.

[pree-jackie]

In D63378#1575917, @pree-jackie wrote:

In D63378#1575196, @lhames wrote:

Hi Praveen,

I'm not sure the IRSpeculator is necessary. My instinct would be to just have the Speculator class as the global repository of estimates of the form "execution point -> { candidates to speculatively compile }" (i.e. your GlobalSpecMap). Then the IRSpeculationLayer should take a reference to the Speculator and use a custom FunctionPass to populates the estimates.

I've attached an example of a hand rolled pass pipeline that should give you access to statically calculated block frequency estimates:

block-freq-pass-example.cpp2 KBDownload

• Lang.

I think it is really a matter of separation, that is where we analysis the module and retrieve the results that means whether in a layer or a speculator object, If we take over the layer approach we have to make the Pass to be aware of speculator.

Also, we have an extra argument (JITDylib*) which cannot be passed around.. Is there any strong reason we should go with your approach?

In my point of view, we don't need to setup whole Pass Manager Pipeline, the minimalistic requirement is FunctionAnalysisManager.

[pree-jackie]

lhames idea :

Remove Speculator variants for different program representations and handle them via ORC Layers.

[pree-jackie]

In D63378#1576628, @pree-jackie wrote:

lhames idea :

Remove Speculator variants for different program representations and handle them via ORC Layers.

Plus, remove code boilerplate.

[pree-jackie]

lib/ExecutionEngine/Orc/Speculation.cpp

32 ↗

(On Diff #208493)

Addressed.

[lhames]

This is looking good. I think it would be nice to have an example QueryAnalysis implementation to demonstrate how that is supposed to work.

include/llvm/ExecutionEngine/Orc/Speculation.h

144–151 ↗

(On Diff #208755)

One minor observation: I think that adding the ability to configure the pass pipeline here provides more degrees of freedom than necessary. If the aim is to simplify the QueryAnalysis by providing a preconfigured pass pipeline then the IRSpeculationLayer should configure that pipeline (and expose no registration mechanism). If the aim is for the pipeline to remain configurable (which seems like a good idea) then you can remove the PassBuilder and FunctionPassManager from the layer entirely (simplifying the layer implementation) and have the QueryAnalysis capture them when it is created.

Note that this is just polish if you have time. I would not let it block your progress on real experiments.

[pree-jackie]

1. Included a Example to show how to use speculation,

2. With the current implementation, we see speedups but it is not benchmark and quantized yet, the analysis are experimental.

3. Launching lookup's for already compiled symbols are avoided by a simple pre map lookup, but the jump can be avoided entirely by having a global variable and creating the new entry Basic Block, which comparison failure : jump into the JIT compiler memory and on success branch to original "pre entry Basic Block of function".

[lhames]

Minor style nits aside, this looks great. Once you've addressed them (and run clang-format over the patch) feel free to commit this to the mainline and continue development there. :)

llvm/examples/SpeculativeJIT/SpeculativeJIT.cpp

198 ↗

(On Diff #212603)

There should be a newline here.

llvm/include/llvm/ExecutionEngine/Orc/LazyReexports.h

23

Is this FileOutputBuffer include needed?

188

The SrcJDLoc parameter should have a default value of nullptr, so that people who are not using speculation can ignore it.

llvm/include/llvm/ExecutionEngine/Orc/Speculation.h

55

It's best to use "return None" rather than "return {}". I've seen some older compilers struggle with the latter.

72

The LLVM style guide says that variable names should start with a capital letter, so lockit should be Lockit.

81

Lockit should start with a capital letter here too.

115

Speculator doesn't seem to have any virtual methods. Are you intending for this to be subclassed? If not, you can remove the virtual destructor too.

164–171

I think this should be able to be removed now?

llvm/lib/ExecutionEngine/Orc/SpeculateAnalyses.cpp

76 ↗

(On Diff #212603)

topk should start with a capital letter.

llvm/lib/ExecutionEngine/Orc/Speculation.cpp

31

lockit should start with a capital letter.

[pree-jackie]

llvm/include/llvm/ExecutionEngine/Orc/Speculation.h

164–171

I tend to keep this interface, because it helps us to register analysis which uses already available LLVM pass results, without this, it may cumbersome at the client code (SpeculativeJIT).
Let's keep it, and see if it is not worthy we can kill this.

[pree-jackie]

Resolve comments

[pree-jackie]

Adapt to new locking scheme

[pree-jackie]

The change set is committed in rL367756.

llvm/include/llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h

 #include "llvm/ExecutionEngine/Orc/LazyReexports.h"
 #include "llvm/ExecutionEngine/Orc/Legacy.h"
 #include "llvm/ExecutionEngine/Orc/OrcError.h"
+#include "llvm/ExecutionEngine/Orc/Speculation.h"
 #include "llvm/ExecutionEngine/RuntimeDyld.h"
 #include "llvm/IR/Attributes.h"
 #include "llvm/IR/Constant.h"

   /// Sets the partition function.
   void setPartitionFunction(PartitionFunction Partition);
 
+  /// Sets the ImplSymbolMap
+  void setImplMap(ImplSymbolMap *Imp);
   /// Emits the given module. This should not be called by clients: it will be
   /// called by the JIT when a definition added via the add method is requested.
   void emit(MaterializationResponsibility R, ThreadSafeModule TSM) override;

   PerDylibResourcesMap DylibResources;
   PartitionFunction Partition = compileRequested;
   SymbolLinkagePromoter PromoteSymbols;
+  ImplSymbolMap *AliaseeImpls = nullptr;
 };
 
 /// Compile-on-demand layer.

llvm/include/llvm/ExecutionEngine/Orc/LazyReexports.h

 
 #include "llvm/ExecutionEngine/Orc/Core.h"
 #include "llvm/ExecutionEngine/Orc/IndirectionUtils.h"
+#include "llvm/ExecutionEngine/Orc/Speculation.h"
 
 namespace llvm {
 

                                    IndirectStubsManager &ISManager,
                                    JITDylib &SourceJD,
                                    SymbolAliasMap CallableAliases,
-                                   VModuleKey K);
+                                   ImplSymbolMap *SrcJDLoc, VModuleKey K);
 
   StringRef getName() const override;
 

   SymbolAliasMap CallableAliases;
   std::shared_ptr<LazyCallThroughManager::NotifyResolvedFunction>
       NotifyResolved;
+  ImplSymbolMap *AliaseeTable;
 };
 
 /// Define lazy-reexports based on the given SymbolAliasMap. Each lazy re-export

 inline std::unique_ptr<LazyReexportsMaterializationUnit>
 lazyReexports(LazyCallThroughManager &LCTManager,
               IndirectStubsManager &ISManager, JITDylib &SourceJD,
-              SymbolAliasMap CallableAliases, VModuleKey K = VModuleKey()) {
+              SymbolAliasMap CallableAliases, ImplSymbolMap *SrcJDLoc,
+              VModuleKey K = VModuleKey()) {
   return llvm::make_unique<LazyReexportsMaterializationUnit>(
-      LCTManager, ISManager, SourceJD, std::move(CallableAliases),
+      LCTManager, ISManager, SourceJD, std::move(CallableAliases), SrcJDLoc,
       std::move(K));
 }
 

llvm/include/llvm/ExecutionEngine/Orc/Speculation.h

+//===-- Speculation.h - Speculative Compilation --*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Contains the definition to support speculative compilation when laziness is
+// enabled.
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_EXECUTIONENGINE_ORC_SPECULATION_H
+#define LLVM_EXECUTIONENGINE_ORC_SPECULATION_H
+
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ExecutionEngine/Orc/Core.h"
+#include "llvm/ExecutionEngine/Orc/IRCompileLayer.h"
+#include "llvm/IR/PassManager.h"
+#include "llvm/Passes/PassBuilder.h"
+#include <mutex>
+#include <type_traits>
+#include <utility>
+#include <vector>
+
+namespace llvm {
+namespace orc {
+
+class Speculator;
+
+// Track the Impls (JITDylib,Symbols) of Symbols while lazy call through
+// trampolines are created. Operations are guarded by locks tp ensure that Imap
+// stays in consistent state after read/write
+
+class ImplSymbolMap {
+  friend class Speculator;
+
+public:
+  using ImplPair = std::pair<SymbolStringPtr, JITDylib *>;
+  using ImapTy = DenseMap<SymbolStringPtr, ImplPair>;
+
+  void trackImpls(SymbolAliasMap ImplMaps, JITDylib *SrcJD);
+
+private:
+  ImplPair getImplFor(const SymbolStringPtr &StubAddr) {
+    std::lock_guard<std::mutex> Lockit(ConcurrentAccess);
+    auto Position = Maps.find(StubAddr);
+    assert(Position != Maps.end() &&
+           "ImplSymbols are not tracked for this Symbol?");
+    return Position->getSecond();
+  }
+
+  std::mutex ConcurrentAccess;
+  ImapTy Maps;
+};
+
+// Defines Speculator Concept, clients can re-use this common functionality
+// for different program representation
+
+class Speculator {
+public:
+  using TargetFAddr = JITTargetAddress;
+  using TargetAndLikelies = DenseMap<SymbolStringPtr, SymbolNameSet>;
+  using StubAddrLikelies = DenseMap<TargetFAddr, SymbolNameSet>;
+
+private:
+  void registerSymbolsWithAddr(TargetFAddr ImplAddr,
+                               SymbolNameSet likelySymbols) {
+    std::lock_guard<std::mutex> lockit(ConcurrentAccess);
+    GlobalSpecMap.insert({ImplAddr, std::move(likelySymbols)});
+  }
+
+  void launchCompile(JITTargetAddress FAddr) {
+    auto It = GlobalSpecMap.find(FAddr);
+    assert(It != GlobalSpecMap.end() &&
+           "launching speculative compiles for unexpected function address?");
+    for (auto &Callee : It->getSecond()) {
+      auto ImplSymbol = AliaseeImplTable.getImplFor(Callee);
+      const auto &ImplSymbolName = ImplSymbol.first;
+      auto *ImplJD = ImplSymbol.second;
+      ES.lookup(JITDylibSearchList({{ImplJD, true}}),
+                SymbolNameSet({ImplSymbolName}), SymbolState::Ready,
+                [this](Expected<SymbolMap> Result) {
+                  if (auto Err = Result.takeError())
+                    ES.reportError(std::move(Err));
+                },
+                NoDependenciesToRegister);
+    }
+  }
+
+public:
+  Speculator(ImplSymbolMap &Impl, ExecutionSession &ref)
+      : AliaseeImplTable(Impl), ES(ref), GlobalSpecMap(0) {}
+  Speculator(const Speculator &) = delete;
+  Speculator(Speculator &&) = delete;
+  Speculator &operator=(const Speculator &) = delete;
+  Speculator &operator=(Speculator &&) = delete;
+  virtual ~Speculator() {}
+
+  // Speculatively compile likely functions for the given Stub Address.
+  // destination of __orc_speculate_for jump
+  void speculateFor(TargetFAddr StubAddr) { launchCompile(StubAddr); }
+
+  // FIXME : Register with Stub Address, after JITLink Fix.
+  void registerSymbols(TargetAndLikelies Candidates, JITDylib *JD) {
+    for (auto &SymPair : Candidates) {
+      auto Target = SymPair.first;
+      auto Likely = SymPair.second;
+
+      auto OnReadyFixUp = [Likely, Target,
+                           this](Expected<SymbolMap> ReadySymbol) {
+        if (ReadySymbol) {
+          auto RAddr = (*ReadySymbol)[Target].getAddress();
+          registerSymbolsWithAddr(RAddr, std::move(Likely));
+        } else
+          this->getES().reportError(ReadySymbol.takeError());
+      };
+      ES.lookup(JITDylibSearchList({{JD, false}}), SymbolNameSet({Target}),
+                SymbolState::Ready, OnReadyFixUp, NoDependenciesToRegister);
+    }
+  }
+
+  ExecutionSession &getES() { return ES; }
+
+private:
+  std::mutex ConcurrentAccess;
+  ImplSymbolMap &AliaseeImplTable;
+  ExecutionSession &ES;
+  StubAddrLikelies GlobalSpecMap;
+};
+
+// Speculator Interface for LLVM IR representation
+class IRSpeculatorImpl : private Speculator {
+
+public:
+  using CompilationUnitTy = Module;
+  using PassIRUnit = Function;
+  using IRlikiesStrRef = DenseMap<StringRef, DenseSet<StringRef>>;
+
+private:
+  // register all function analysis pass in LLVM
+  void setUpAnalysesPass() { PB.registerFunctionAnalyses(FAM); }
+
+public:
+  IRSpeculatorImpl(ImplSymbolMap &impl, ExecutionSession &ES)
+      : Speculator(impl, ES) {
+    setUpAnalysesPass();
+  }
+
+  template <
+      typename AnalysisTy,
+      typename std::enable_if<
+          std::is_base_of<AnalysisInfoMixin<AnalysisTy>, AnalysisTy>::value,
+          bool>::type = true>
+  void registerAnalysis() {
+    FAM.registerPass([]() { return AnalysisTy(); });
+  }
+
+  void runOn(const CompilationUnitTy &ProgramRepr, JITDylib &LookupDylib) {
+    for (auto &Fn : ProgramRepr.getFunctionList()) {
+      if (!Fn.isDeclaration()) {
+        auto InternedNames = internToJITSymbols(collect(Fn));
+        registerSymbols(InternedNames, &LookupDylib);
+      }
+    }
+  }
+
+  FunctionAnalysisManager &getAnalysisManager() { return FAM; }
+
+private:
+  virtual IRlikiesStrRef collect(const PassIRUnit &) = 0;
+
+  TargetAndLikelies internToJITSymbols(IRlikiesStrRef IRNames) {
+    TargetAndLikelies InternedNames;
+    DenseSet<SymbolStringPtr> TargetJITNames;
+    ExecutionSession &Es = getES();
+    for (auto &NamePair : IRNames) {
+      for (auto &TargetNames : NamePair.second)
+        TargetJITNames.insert(Es.intern(TargetNames));
+
+      InternedNames.insert(
+          {Es.intern(NamePair.first), std::move(TargetJITNames)});
+    }
+    return InternedNames;
+  }
+
+  PassBuilder PB;
+  FunctionAnalysisManager FAM;
+};
+
+// Create Object of this class templated with top-level analysis pass, type is
+// used to query the Function Analysis Manager, Pass a function object to help
+// speculator to retrieve likely next executable based on the Analysis Result
+
+template <
+    typename AnalysisType,
+    typename std::enable_if<
+        std::is_base_of<AnalysisInfoMixin<AnalysisType>, AnalysisType>::value,
+        bool>::type = true>
+class IRSpeculator : public IRSpeculatorImpl {
+public:
+  using EvalResultTy = std::function<IRlikiesStrRef(
+      typename AnalysisType::Result &, const PassIRUnit &)>;
+  IRSpeculator(ImplSymbolMap &Impl, ExecutionSession &ES,
+               EvalResultTy &&Interpreter)
+      : IRSpeculatorImpl(Impl, ES) {
+    ResultResolver = Interpreter;
+  }
+
+private:
+  IRlikiesStrRef collect(const PassIRUnit &Fn) override {
+    auto &V = getAnalysisManager().template getResult<AnalysisType>(
+        const_cast<Function &>(Fn));
+    auto IRNames = ResultResolver(V, Fn);
+    return IRNames;
+  }
+  EvalResultTy ResultResolver;
+};
+
+// Helps to Speculate and Instrument LLVM IR
+class IRSpeculationLayer : public IRLayer {
+public:
+  IRSpeculationLayer(ExecutionSession &ES, IRCompileLayer &BaseLayer,
+                     IRSpeculatorImpl &Spec)
+      : IRLayer(ES), NextLayer(BaseLayer), S(Spec) {}
+
+  void emit(MaterializationResponsibility R, ThreadSafeModule TSM);
+
+private:
+  IRCompileLayer &NextLayer;
+  IRSpeculatorImpl &S;
+};
+
+// Runtime Function Interface
+extern "C" {
+void __orc_speculate_for(Speculator *, uint64_t stub_id);
+}
+
+} // namespace orc
+} // namespace llvm
+#endif
 No newline at end of file

llvm/lib/ExecutionEngine/Orc/CMakeLists.txt

   RPCUtils.cpp
   RTDyldObjectLinkingLayer.cpp
   ThreadSafeModule.cpp
-
+  Speculation.cpp
   ADDITIONAL_HEADER_DIRS
   ${LLVM_MAIN_INCLUDE_DIR}/llvm/ExecutionEngine/Orc
 

   PRIVATE
   LLVMBitReader
   LLVMBitWriter
+  LLVMAnalysis
   )

llvm/lib/ExecutionEngine/Orc/CompileOnDemandLayer.cpp

   this->Partition = std::move(Partition);
 }
 
+void CompileOnDemandLayer::setImplMap(ImplSymbolMap *Imp) {
+  this->AliaseeImpls = Imp;
+}
 void CompileOnDemandLayer::emit(MaterializationResponsibility R,
                                 ThreadSafeModule TSM) {
   assert(TSM.getModule() && "Null module");

 
   R.replace(reexports(PDR.getImplDylib(), std::move(NonCallables), true));
   R.replace(lazyReexports(LCTMgr, PDR.getISManager(), PDR.getImplDylib(),
-                          std::move(Callables)));
+                          std::move(Callables), AliaseeImpls));
 }
 
 CompileOnDemandLayer::PerDylibResources &

llvm/lib/ExecutionEngine/Orc/LazyReexports.cpp

 
 LazyReexportsMaterializationUnit::LazyReexportsMaterializationUnit(
     LazyCallThroughManager &LCTManager, IndirectStubsManager &ISManager,
-    JITDylib &SourceJD, SymbolAliasMap CallableAliases, VModuleKey K)
+    JITDylib &SourceJD, SymbolAliasMap CallableAliases, ImplSymbolMap *SrcJDLoc,
+    VModuleKey K)
     : MaterializationUnit(extractFlags(CallableAliases), std::move(K)),
       LCTManager(LCTManager), ISManager(ISManager), SourceJD(SourceJD),
       CallableAliases(std::move(CallableAliases)),

           [&ISManager](JITDylib &JD, const SymbolStringPtr &SymbolName,
                        JITTargetAddress ResolvedAddr) {
             return ISManager.updatePointer(*SymbolName, ResolvedAddr);
-          })) {}
+          })),
+      AliaseeTable(SrcJDLoc) {}
 
 StringRef LazyReexportsMaterializationUnit::getName() const {
   return "<Lazy Reexports>";

 
   if (!CallableAliases.empty())
     R.replace(lazyReexports(LCTManager, ISManager, SourceJD,
-                            std::move(CallableAliases)));
+                            std::move(CallableAliases), AliaseeTable));
 
   IndirectStubsManager::StubInitsMap StubInits;
   for (auto &Alias : RequestedAliases) {

         std::make_pair(*CallThroughTrampoline, Alias.second.AliasFlags);
   }
 
+  if (AliaseeTable != nullptr)
+    AliaseeTable->trackImpls(RequestedAliases, &SourceJD);
+
   if (auto Err = ISManager.createStubs(StubInits)) {
     SourceJD.getExecutionSession().reportError(std::move(Err));
     R.failMaterialization();

llvm/lib/ExecutionEngine/Orc/Speculation.cpp

+//===---------- speculation.cpp - Utilities for Speculation ----------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/ExecutionEngine/Orc/Speculation.h"
+
+#include "llvm/IR/BasicBlock.h"
+#include "llvm/IR/Function.h"
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/Instruction.h"
+#include "llvm/IR/Instructions.h"
+#include "llvm/IR/LLVMContext.h"
+#include "llvm/IR/Module.h"
+#include "llvm/IR/Type.h"
+#include "llvm/IR/Verifier.h"
+
+#include <vector>
+
+namespace llvm {
+
+namespace orc {
+
+// ImplSymbolMap methods
+void ImplSymbolMap::trackImpls(SymbolAliasMap ImplMaps, JITDylib *SrcJD) {
+  std::lock_guard<std::mutex> lockit(ConcurrentAccess);
+  for (auto &I : ImplMaps) {
+    auto It = Maps.insert({I.first, {I.second.Aliasee, SrcJD}});
+    assert(It.second && "ImplSymbols are already tracked for this Symbol?");
+  }
+}
+
+void IRSpeculationLayer::emit(MaterializationResponsibility R,
+                              ThreadSafeModule TSM) {
+  auto M = TSM.getModule();
+
+  // Do not spawn new thread.
+  S.runOn(*M, R.getTargetJITDylib());
+
+  assert(M && "Speculation Layer received Null Module ?");
+
+  // Instrumentation of runtime calls
+  auto &InContext = M->getContext();
+  auto SpeculatorVTy = StructType::create(InContext, "Class.Speculator");
+  auto RuntimeCallTy = FunctionType::get(
+      Type::getVoidTy(InContext),
+      {SpeculatorVTy->getPointerTo(), Type::getInt64Ty(InContext)}, false);
+
+  auto RuntimeCall =
+      Function::Create(RuntimeCallTy, Function::LinkageTypes::ExternalLinkage,
+                       "__orc_speculate_for", *M);
+
+  auto SpeclAddr = new GlobalVariable(
+      *M, SpeculatorVTy, false, GlobalValue::LinkageTypes::ExternalLinkage,
+      nullptr, "__orc_speculator");
+
+  IRBuilder<> Mutator(InContext);
+  for (auto &Fn : M->getFunctionList()) {
+    if (!Fn.isDeclaration()) {
+      Mutator.SetInsertPoint(&(Fn.getEntryBlock().front()));
+      auto ImplAddrToUint =
+          Mutator.CreatePtrToInt(&Fn, Type::getInt64Ty(InContext));
+      Mutator.CreateCall(RuntimeCallTy, RuntimeCall,
+                         {SpeclAddr, ImplAddrToUint});
+    }
+  }
+
+  assert(!verifyModule(*M) && "Speculation Instrumentation breaks IR?");
+
+  NextLayer.emit(std::move(R), std::move(TSM));
+}
+
+// Runtime Function Implementation
+extern "C" void __orc_speculate_for(Speculator *Ptr, uint64_t StubId) {
+  assert(Ptr && " Null Address Received in orc_speculate_for ");
+  Ptr->speculateFor(StubId);
+}
+
+} // namespace orc
+} // namespace llvm
 No newline at end of file