[ORCv2] - New Speculate Query Implementation

[pree-jackie]

Summary

This patch introduces, SequenceBBQuery - new heuristic to find likely next callable functions it tries to find the blocks with calls in order of execution sequence of Blocks.

It still uses BlockFrequencyAnalysis to find high frequency blocks. For a handful of hottest blocks (plan to customize), the algorithm traverse and discovered the caller blocks along the way to Entry Basic Block and Exit Basic Block. It uses Block Hint, to stop traversing the already visited blocks in both direction. It implicitly assumes that once the block is visited during discovering entry or exit nodes, revisiting them again does not add much. It also branch probability info (cached result) to traverse only hot edges (planned to customize) from hot blocks. Without BPI, the algorithm mostly return's all the blocks in the CFG with calls.

It also changes the heuristic queries, so they don't maintain states. Hence it is safe to call from multiple threads.

It also implements, new instrumentation to avoid jumping into JIT on every call to the function with the help _orc_speculate.decision.block and _orc_speculate.block.

"Speculator Registration Mechanism is also changed" - kudos to @lhames

Open to review, mostly looking to change implementation of SequeceBBQuery heuristics with good data structure choices.

[mgrang]

llvm/lib/ExecutionEngine/Orc/SpeculateAnalyses.cpp

262

You can use the range-based version of llvm::sort here:

llvm::sort(BBFreqs);

[pree-jackie]

llvm/lib/ExecutionEngine/Orc/SpeculateAnalyses.cpp

262

Thanks for the catch.

[lhames]

llvm/lib/ExecutionEngine/Orc/SpeculateAnalyses.cpp

168–173

I think this can be simplified to:

BlockHint.CallerBlock = llvm::find(CallerBlocks, AtBB) != CallerBlocks.end();
VisitedBlocks.insert(std::make_pair(AtBB, BlockHint));

244–245

Can you describe more about how SequencedBBQuery differs from BlockFreqQuery?

In particular: (1) What is the output of this algorithm, queryCFG, and (2) what benefit do we get from scanning to the exit blocks?

[dblaikie]

llvm/include/llvm/ExecutionEngine/Orc/SpeculateAnalyses.h

38

You can remove this 'private' since it's implicit

51

no need for 'private' here, when a class members are private by default

53

no need for 'public' in a struct where it's public by default

57

Perhaps use non-static data member initializers instead of a ctor here? (add initializers "= true", etc to the member declarations directly, and remove this ctor)

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

86–87

Remove else-after-return ( https://llvm.org/docs/CodingStandards.html#don-t-use-else-after-a-return )

129–133

Does this dtor need to be explicitly declared? Or could this line be deleted?

llvm/lib/ExecutionEngine/Orc/SpeculateAnalyses.cpp

70–71

Does this capture any case that isn't captured by the all_of below? Perhaps this special case could be removed?

72

You can drop the "-> bool" part - it's implicit & probably clear enough from the implementation that the return type is bool.

73

"x ? true : false" seems redundant. Perhaps "return BB.getSingleSuccessor() != nullptr;" would be OK/better?

131–132

Remove else-after-return

142

Can probably use llvm::find(BBList, &Block) rather than having to pass begin/end here

145

*FItr == &Block here, right? Perhaps it'd be more clear to use &Block here, then - so it's clear there's nothing being accessed in BBList - it's only used for a lookup test?

(& maybe there's a neater way to write that - ah, llvm::is_contained?)

206–207

Remove else-after-return

218–220

Perhaps these two calls to VisitedBlocks could be coalesced?

if (llvm::is_contained(CallerBlocks, AtBB))
  BlockHint.CallerBlock = true;
VisitedBlocks.insert(std::make_pair(AtBB, BlockHint));

231

Should this use a range-based for loop?

235

Could this use a range-based for loop too?

265

Skip the "? true : false" here, just "return BBF.second > BBS.second;" - that's already a boolean expression.

277–278

Possibly using range-based for loop here (maybe using an arrayref sliced down to NHotBlocks)

[pree-jackie]

llvm/lib/ExecutionEngine/Orc/SpeculateAnalyses.cpp

244–245

The blockfreqquery, only selects the blocks with highest execution frequencies. Based on the observation, often the blocks within the loop nests have high frequency. The idea here is to, traverse back towards the entry node and exit node to collect all the predecessor and successor blocks with call instructions respectively.
Here the standard hot branch probability (4/5) is used to guide traversal to entry and exit nodes but we
can customize the branch probability to define own hotness instead of 80%, to have better coverage.

Algorithm produces a set of basic blocks N, where N >= NHotBlocks,

that is when N = NHotBlocks, there are no predecessor & successor blocks with call inst.

when N > NHotBlocks, x = (N - NHotBlocks)
there are x number of predecessor & successor blocks with call inst which we miss in block freq query.

[pree-jackie]

llvm/lib/ExecutionEngine/Orc/SpeculateAnalyses.cpp

244–245

This heuristic outputs a sequence of basic blocks,
for example : entry -> bb1.call -> bb2.call.
The last operation of the heuristic, rearranges the collected set of basic blocks by taking the distance from the entry node into account. In the above example, bb2.call may have high frequency than bb1.call, but based on distance from entry node they are swapped.

244–245

Plus, this sequence will be very helpful in future that when we are introducing custom speculate.points inside the function instead of just having them in top of CFG.
In essence, we can decide the appropriate position of speculate.points by knowing the sequence of blocks with call,

Plus, this is experimental try.

[pree-jackie]

llvm/lib/ExecutionEngine/Orc/SpeculateAnalyses.cpp

235

No, I guess.

[pree-jackie]

I resolved N-1 comments. Plus, this change set implements the launching bulk speculative queries and locking the module before instrumenting.
I kindly request you to review the else-after-return resolved part of code, to check whether I understand it correctly.

Thanks.

[lhames]

LGTM, pending else-after-return fixes.

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

86–88

This is not quite right.

"else after return" refers to the following idiom:

if (C) {
  // do stuff
  return;
} else {
  // do other stuff.
}

LLVM's style guide recommends that this be rewritten as:

if (C) {
  // do stuff
  return;
}
// do other stuff.

In this case the right fix was to replace:

if (It == GlobalSpecMap.end() || Iv.second == false)
  return;
else
  CandidateSet = It->getSecond();

with

if (It == GlobalSpecMap.end() || Iv.second == false)
  return;

CandidateSet = It->getSecond();

llvm/lib/ExecutionEngine/Orc/SpeculateAnalyses.cpp

130–134

To remove the else-after-return here, this should be rewritten from:

if (TotalBlocks == 1)
  return TotalBlocks;
else
  return TotalBlocks / 2;

to:

if (TotalBlocks == 1)
  return TotalBlocks;
return TotalBlocks / 2;

201–204

This can be rewritten as:

if (!Itr->second.Downward)
  return;
Itr->second.Downward = false;

[pree-jackie]

else-after-return fixes. Going to commit

[pree-jackie]

Committing tomorrow as it is difficult to watch build bots continuously now.

llvm/examples/SpeculativeJIT/SpeculativeJIT.cpp

SpeculativeJIT(
     this->ES->setDispatchMaterialization(
 
         [this](JITDylib &JD, std::unique_ptr<MaterializationUnit> MU) {
           // FIXME: Switch to move capture once we have c  14.
           auto SharedMU = std::shared_ptr<MaterializationUnit>(std::move(MU));
           auto Work = [SharedMU, &JD]() { SharedMU->doMaterialize(JD); };
           CompileThreads.async(std::move(Work));
         });
-    JITEvaluatedSymbol SpeculatorSymbol(JITTargetAddress(&S),
-                                        JITSymbolFlags::Exported);
-    ExitOnErr(this->ES->getMainJITDylib().define(
-        absoluteSymbols({{Mangle("__orc_speculator"), SpeculatorSymbol}})));
+    ExitOnErr(S.addSpeculationRuntime(this->ES->getMainJITDylib(), Mangle));
     LocalCXXRuntimeOverrides CXXRuntimeoverrides;
     ExitOnErr(CXXRuntimeoverrides.enable(this->ES->getMainJITDylib(), Mangle));
   }
 
   static std::unique_ptr<SectionMemoryManager> createMemMgr() {
     return std::make_unique<SectionMemoryManager>();
   }
 

llvm/include/llvm/ExecutionEngine/Orc/SpeculateAnalyses.h

 //===-- SpeculateAnalyses.h  --*- 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
 //
 //===----------------------------------------------------------------------===//
 // \file
 /// Contains the Analyses and Result Interpretation to select likely functions
-/// to Speculatively compile before they are called. [Experimentation]
+/// to Speculatively compile before they are called. [Purely Experimentation]
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_EXECUTIONENGINE_ORC_SPECULATEANALYSES_H
 #define LLVM_EXECUTIONENGINE_ORC_SPECULATEANALYSES_H
 
+#include "llvm/Analysis/BranchProbabilityInfo.h"
 #include "llvm/ExecutionEngine/Orc/Core.h"
 #include "llvm/ExecutionEngine/Orc/Speculation.h"
 
 #include <vector>
 
 namespace llvm {
 
 namespace orc {
 
-// Direct calls in high frequency basic blocks are extracted.
-class BlockFreqQuery {
-private:
+// Provides common code.
+class SpeculateQuery {
+protected:
   void findCalles(const BasicBlock *, DenseSet<StringRef> &);
-  size_t numBBToGet(size_t);
+  bool isStraightLine(const Function &F);
 
 public:
   using ResultTy = Optional<DenseMap<StringRef, DenseSet<StringRef>>>;
+};
 
+// Direct calls in high frequency basic blocks are extracted.
+class BlockFreqQuery : public SpeculateQuery {
+  size_t numBBToGet(size_t);
+
+public:
   // Find likely next executables based on IR Block Frequency
-  ResultTy operator()(Function &F, FunctionAnalysisManager &FAM);
+  ResultTy operator()(Function &F);
+};
+
+// This Query generates a sequence of basic blocks which follows the order of
+// execution.
+// A handful of BB with higher block frequencies are taken, then path to entry
+// and end BB are discovered by traversing up & down the CFG.
+class SequenceBBQuery : public SpeculateQuery {
+  struct WalkDirection {
+    bool Upward = true, Downward = true;
+    // the block associated contain a call
+    bool CallerBlock = false;
+  };
+
+public:
+  using VisitedBlocksInfoTy = DenseMap<const BasicBlock *, WalkDirection>;
+  using BlockListTy = SmallVector<const BasicBlock *, 8>;
+  using BackEdgesInfoTy =
+      SmallVector<std::pair<const BasicBlock *, const BasicBlock *>, 8>;
+  using BlockFreqInfoTy =
+      SmallVector<std::pair<const BasicBlock *, uint64_t>, 8>;
+
+private:
+  std::size_t getHottestBlocks(std::size_t TotalBlocks);
+  BlockListTy rearrangeBB(const Function &, const BlockListTy &);
+  BlockListTy queryCFG(Function &, const BlockListTy &);
+  void traverseToEntryBlock(const BasicBlock *, const BlockListTy &,
+                            const BackEdgesInfoTy &,
+                            const BranchProbabilityInfo *,
+                            VisitedBlocksInfoTy &);
+  void traverseToExitBlock(const BasicBlock *, const BlockListTy &,
+                           const BackEdgesInfoTy &,
+                           const BranchProbabilityInfo *,
+                           VisitedBlocksInfoTy &);
+
+public:
+  ResultTy operator()(Function &F);
 };
 
 } // namespace orc
 } // namespace llvm
 
 #endif // LLVM_EXECUTIONENGINE_ORC_SPECULATEANALYSES_H

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

 
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/Optional.h"
 #include "llvm/ExecutionEngine/Orc/Core.h"
 #include "llvm/ExecutionEngine/Orc/IRCompileLayer.h"
 #include "llvm/IR/PassManager.h"
 #include "llvm/Passes/PassBuilder.h"
+#include "llvm/Support/Debug.h"
 
 #include <mutex>
 #include <type_traits>
 #include <utility>
 #include <vector>
 
 namespace llvm {
 namespace orc {
private:
 
   void launchCompile(JITTargetAddress FAddr) {
     SymbolNameSet CandidateSet;
     // Copy CandidateSet is necessary, to avoid unsynchronized access to
     // the datastructure.
     {
       std::lock_guard<std::mutex> Lockit(ConcurrentAccess);
       auto It = GlobalSpecMap.find(FAddr);
-      // Kill this when jump on first call instrumentation is in place;
-      auto Iv = AlreadyExecuted.insert(FAddr);
-      if (It == GlobalSpecMap.end() || Iv.second == false)
+      if (It == GlobalSpecMap.end())
         return;
-      else
-        CandidateSet = It->getSecond();
     }
 
-    // Try to distinguish pre-compiled symbols!
+    SymbolDependenceMap SpeculativeLookUpImpls;
+
     for (auto &Callee : CandidateSet) {
       auto ImplSymbol = AliaseeImplTable.getImplFor(Callee);
+      // try to distinguish already compiled & library symbols
       if (!ImplSymbol.hasValue())
         continue;
       const auto &ImplSymbolName = ImplSymbol.getPointer()->first;
-      auto *ImplJD = ImplSymbol.getPointer()->second;
-      ES.lookup(JITDylibSearchList({{ImplJD, true}}),
-                SymbolNameSet({ImplSymbolName}), SymbolState::Ready,
+      JITDylib *ImplJD = ImplSymbol.getPointer()->second;
+      auto &SymbolsInJD = SpeculativeLookUpImpls[ImplJD];
+      SymbolsInJD.insert(ImplSymbolName);
+    }
+
+    DEBUG_WITH_TYPE("orc", for (auto &I
+                                : SpeculativeLookUpImpls) {
+      llvm::dbgs() << "\n In " << I.first->getName() << " JITDylib ";
+      for (auto &N : I.second)
+        llvm::dbgs() << "\n Likely Symbol : " << N;
+    });
+
+    // for a given symbol, there may be no symbol qualified for speculatively
+    // compile try to fix this before jumping to this code if possible.
+    for (auto &LookupPair : SpeculativeLookUpImpls)
+      ES.lookup(JITDylibSearchList({{LookupPair.first, true}}),
+                LookupPair.second, 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;
-  ~Speculator() {}
+
+  /// Define symbols for this Speculator object (__orc_speculator) and the
+  /// speculation runtime entry point symbol (__orc_speculate_for) in the
+  /// given JITDylib.
+  Error addSpeculationRuntime(JITDylib &JD, MangleAndInterner &Mangle);
 
   // 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(FunctionCandidatesMap Candidates, JITDylib *JD) {
     for (auto &SymPair : Candidates) {
for (auto &SymPair : Candidates) {
       ES.lookup(JITDylibSearchList({{JD, true}}), SymbolNameSet({Target}),
                 SymbolState::Ready, OnReadyFixUp, NoDependenciesToRegister);
     }
   }
 
   ExecutionSession &getES() { return ES; }
 
 private:
+  static void speculateForEntryPoint(Speculator *Ptr, uint64_t StubId);
   std::mutex ConcurrentAccess;
   ImplSymbolMap &AliaseeImplTable;
   ExecutionSession &ES;
-  DenseSet<TargetFAddr> AlreadyExecuted;
   StubAddrLikelies GlobalSpecMap;
 };
-// replace DenseMap with Pair
+
 class IRSpeculationLayer : public IRLayer {
 public:
   using IRlikiesStrRef = Optional<DenseMap<StringRef, DenseSet<StringRef>>>;
-  using ResultEval =
-      std::function<IRlikiesStrRef(Function &, FunctionAnalysisManager &)>;
+  using ResultEval = std::function<IRlikiesStrRef(Function &)>;
   using TargetAndLikelies = DenseMap<SymbolStringPtr, SymbolNameSet>;
 
   IRSpeculationLayer(ExecutionSession &ES, IRCompileLayer &BaseLayer,
                      Speculator &Spec, MangleAndInterner &Mangle,
                      ResultEval Interpreter)
       : IRLayer(ES), NextLayer(BaseLayer), S(Spec), Mangle(Mangle),
-        QueryAnalysis(Interpreter) {
-    PB.registerFunctionAnalyses(FAM);
-  }
+        QueryAnalysis(Interpreter) {}
 
   void emit(MaterializationResponsibility R, ThreadSafeModule TSM);
 
 private:
   TargetAndLikelies
   internToJITSymbols(DenseMap<StringRef, DenseSet<StringRef>> IRNames) {
     assert(!IRNames.empty() && "No IRNames received to Intern?");
     TargetAndLikelies InternedNames;
     DenseSet<SymbolStringPtr> TargetJITNames;
     for (auto &NamePair : IRNames) {
       for (auto &TargetNames : NamePair.second)
         TargetJITNames.insert(Mangle(TargetNames));
 
       InternedNames[Mangle(NamePair.first)] = std::move(TargetJITNames);
     }
     return InternedNames;
   }
 
   IRCompileLayer &NextLayer;
   Speculator &S;
   MangleAndInterner &Mangle;
-  PassBuilder PB;
-  FunctionAnalysisManager FAM;
   ResultEval QueryAnalysis;
 };
 
-// Runtime Function Interface
-extern "C" {
-void __orc_speculate_for(Speculator *, uint64_t stub_id);
-}
-
 } // namespace orc
 } // namespace llvm
 
 #endif // LLVM_EXECUTIONENGINE_ORC_SPECULATION_H

llvm/lib/ExecutionEngine/Orc/SpeculateAnalyses.cpp

 //===-- SpeculateAnalyses.cpp  --*- 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/ExecutionEngine/Orc/SpeculateAnalyses.h"
+#include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Analysis/BlockFrequencyInfo.h"
+#include "llvm/Analysis/BranchProbabilityInfo.h"
+#include "llvm/Analysis/CFG.h"
+#include "llvm/IR/PassManager.h"
+#include "llvm/Passes/PassBuilder.h"
+#include "llvm/Support/ErrorHandling.h"
+
+#include <algorithm>
 
 namespace {
 using namespace llvm;
-std::vector<const BasicBlock *> findBBwithCalls(const Function &F,
+SmallVector<const BasicBlock *, 8> findBBwithCalls(const Function &F,
                                                    bool IndirectCall = false) {
-  std::vector<const BasicBlock *> BBs;
+  SmallVector<const BasicBlock *, 8> BBs;
 
   auto findCallInst = [&IndirectCall](const Instruction &I) {
-    if (auto Call = dyn_cast<CallBase>(&I)) {
-      if (Call->isIndirectCall())
-        return IndirectCall;
-      else
-        return true;
-    } else
+    if (auto Call = dyn_cast<CallBase>(&I))
+      return Call->isIndirectCall() ? IndirectCall : true;
       return false;
   };
   for (auto &BB : F)
     if (findCallInst(*BB.getTerminator()) ||
         llvm::any_of(BB.instructionsWithoutDebug(), findCallInst))
       BBs.emplace_back(&BB);
 
   return BBs;
 }
 } // namespace
 
 // Implementations of Queries shouldn't need to lock the resources
 // such as LLVMContext, each argument (function) has a non-shared LLVMContext
+// Plus, if Queries contain states necessary locking scheme should be provided.
 namespace llvm {
 namespace orc {
 
 // Collect direct calls only
-void BlockFreqQuery::findCalles(const BasicBlock *BB,
+void SpeculateQuery::findCalles(const BasicBlock *BB,
                                 DenseSet<StringRef> &CallesNames) {
   assert(BB != nullptr && "Traversing Null BB to find calls?");
 
   auto getCalledFunction = [&CallesNames](const CallBase *Call) {
     auto CalledValue = Call->getCalledOperand()->stripPointerCasts();
     if (auto DirectCall = dyn_cast<Function>(CalledValue))
       CallesNames.insert(DirectCall->getName());
   };
   for (auto &I : BB->instructionsWithoutDebug())
     if (auto CI = dyn_cast<CallInst>(&I))
       getCalledFunction(CI);
 
   if (auto II = dyn_cast<InvokeInst>(BB->getTerminator()))
     getCalledFunction(II);
 }
 
-// blind calculation
+bool SpeculateQuery::isStraightLine(const Function &F) {
+  return llvm::all_of(F.getBasicBlockList(), [](const BasicBlock &BB) {
+    return BB.getSingleSuccessor() != nullptr;
+  });
+}
+
+// BlockFreqQuery Implementations
+
 size_t BlockFreqQuery::numBBToGet(size_t numBB) {
   // small CFG
   if (numBB < 4)
     return numBB;
   // mid-size CFG
   else if (numBB < 20)
     return (numBB / 2);
   else
     return (numBB / 2) + (numBB / 4);
 }
 
-BlockFreqQuery::ResultTy BlockFreqQuery::
-operator()(Function &F, FunctionAnalysisManager &FAM) {
+BlockFreqQuery::ResultTy BlockFreqQuery::operator()(Function &F) {
   DenseMap<StringRef, DenseSet<StringRef>> CallerAndCalles;
   DenseSet<StringRef> Calles;
   SmallVector<std::pair<const BasicBlock *, uint64_t>, 8> BBFreqs;
 
+  PassBuilder PB;
+  FunctionAnalysisManager FAM;
+  PB.registerFunctionAnalyses(FAM);
+
   auto IBBs = findBBwithCalls(F);
 
   if (IBBs.empty())
     return None;
 
   auto &BFI = FAM.getResult<BlockFrequencyAnalysis>(F);
 
   for (const auto I : IBBs)
for (size_t i = 0; i < Topk; i++)
     findCalles(BBFreqs[i].first, Calles);
 
   assert(!Calles.empty() && "Running Analysis on Function with no calls?");
 
   CallerAndCalles.insert({F.getName(), std::move(Calles)});
 
   return CallerAndCalles;
 }
+
+// SequenceBBQuery Implementation
+std::size_t SequenceBBQuery::getHottestBlocks(std::size_t TotalBlocks) {
+  if (TotalBlocks == 1)
+    return TotalBlocks;
+  return TotalBlocks / 2;
+}
+
+// FIXME : find good implementation.
+SequenceBBQuery::BlockListTy
+SequenceBBQuery::rearrangeBB(const Function &F, const BlockListTy &BBList) {
+  BlockListTy RearrangedBBSet;
+
+  for (auto &Block : F.getBasicBlockList())
+    if (llvm::is_contained(BBList, &Block))
+      RearrangedBBSet.push_back(&Block);
+
+  assert(RearrangedBBSet.size() == BBList.size() &&
+         "BasicBlock missing while rearranging?");
+  return RearrangedBBSet;
+}
+
+void SequenceBBQuery::traverseToEntryBlock(const BasicBlock *AtBB,
+                                           const BlockListTy &CallerBlocks,
+                                           const BackEdgesInfoTy &BackEdgesInfo,
+                                           const BranchProbabilityInfo *BPI,
+                                           VisitedBlocksInfoTy &VisitedBlocks) {
+  auto Itr = VisitedBlocks.find(AtBB);
+  if (Itr != VisitedBlocks.end()) { // already visited.
+    if (!Itr->second.Upward)
+      return;
+    Itr->second.Upward = false;
+  } else {
+    // Create hint for newly discoverd blocks.
+    WalkDirection BlockHint;
+    BlockHint.Upward = false;
+    // FIXME: Expensive Check
+    if (llvm::is_contained(CallerBlocks, AtBB))
+      BlockHint.CallerBlock = true;
+    VisitedBlocks.insert(std::make_pair(AtBB, BlockHint));
+  }
+
+  const_pred_iterator PIt = pred_begin(AtBB), EIt = pred_end(AtBB);
+  // Move this check to top, when we have code setup to launch speculative
+  // compiles for function in entry BB, this triggers the speculative compiles
+  // before running the program.
+  if (PIt == EIt) // No Preds.
+    return;
+
+  DenseSet<const BasicBlock *> PredSkipNodes;
+
+  // Since we are checking for predecessor's backedges, this Block
+  // occurs in second position.
+  for (auto &I : BackEdgesInfo)
+    if (I.second == AtBB)
+      PredSkipNodes.insert(I.first);
+
+  // Skip predecessors which source of back-edges.
+  for (; PIt != EIt; ++PIt)
+    // checking EdgeHotness is cheaper
+    if (BPI->isEdgeHot(*PIt, AtBB) && !PredSkipNodes.count(*PIt))
+      traverseToEntryBlock(*PIt, CallerBlocks, BackEdgesInfo, BPI,
+                           VisitedBlocks);
+}
+
+void SequenceBBQuery::traverseToExitBlock(const BasicBlock *AtBB,
+                                          const BlockListTy &CallerBlocks,
+                                          const BackEdgesInfoTy &BackEdgesInfo,
+                                          const BranchProbabilityInfo *BPI,
+                                          VisitedBlocksInfoTy &VisitedBlocks) {
+  auto Itr = VisitedBlocks.find(AtBB);
+  if (Itr != VisitedBlocks.end()) { // already visited.
+    if (!Itr->second.Downward)
+      return;
+    Itr->second.Downward = false;
+  } else {
+    // Create hint for newly discoverd blocks.
+    WalkDirection BlockHint;
+    BlockHint.Downward = false;
+    // FIXME: Expensive Check
+    if (llvm::is_contained(CallerBlocks, AtBB))
+      BlockHint.CallerBlock = true;
+    VisitedBlocks.insert(std::make_pair(AtBB, BlockHint));
+  }
+
+  succ_const_iterator PIt = succ_begin(AtBB), EIt = succ_end(AtBB);
+  if (PIt == EIt) // No succs.
+    return;
+
+  // If there are hot edges, then compute SuccSkipNodes.
+  DenseSet<const BasicBlock *> SuccSkipNodes;
+
+  // Since we are checking for successor's backedges, this Block
+  // occurs in first position.
+  for (auto &I : BackEdgesInfo)
+    if (I.first == AtBB)
+      SuccSkipNodes.insert(I.second);
+
+  for (; PIt != EIt; ++PIt)
+    if (BPI->isEdgeHot(AtBB, *PIt) && !SuccSkipNodes.count(*PIt))
+      traverseToExitBlock(*PIt, CallerBlocks, BackEdgesInfo, BPI,
+                          VisitedBlocks);
+}
+
+// Get Block frequencies for blocks and take most frquently executed block,
+// walk towards the entry block from those blocks and discover the basic blocks
+// with call.
+SequenceBBQuery::BlockListTy
+SequenceBBQuery::queryCFG(Function &F, const BlockListTy &CallerBlocks) {
+
+  BlockFreqInfoTy BBFreqs;
+  VisitedBlocksInfoTy VisitedBlocks;
+  BackEdgesInfoTy BackEdgesInfo;
+
+  PassBuilder PB;
+  FunctionAnalysisManager FAM;
+  PB.registerFunctionAnalyses(FAM);
+
+  auto &BFI = FAM.getResult<BlockFrequencyAnalysis>(F);
+
+  llvm::FindFunctionBackedges(F, BackEdgesInfo);
+
+  for (const auto I : CallerBlocks)
+    BBFreqs.push_back({I, BFI.getBlockFreq(I).getFrequency()});
+
+  llvm::sort(BBFreqs, [](decltype(BBFreqs)::const_reference Bbf,
+                         decltype(BBFreqs)::const_reference Bbs) {
+    return Bbf.second > Bbs.second;
+  });
+
+  ArrayRef<std::pair<const BasicBlock *, uint64_t>> HotBlocksRef(BBFreqs);
+  HotBlocksRef =
+      HotBlocksRef.drop_back(BBFreqs.size() - getHottestBlocks(BBFreqs.size()));
+
+  BranchProbabilityInfo *BPI =
+      FAM.getCachedResult<BranchProbabilityAnalysis>(F);
+
+  // visit NHotBlocks,
+  // traverse upwards to entry
+  // traverse downwards to end.
+
+  for (auto I : HotBlocksRef) {
+    traverseToEntryBlock(I.first, CallerBlocks, BackEdgesInfo, BPI,
+                         VisitedBlocks);
+    traverseToExitBlock(I.first, CallerBlocks, BackEdgesInfo, BPI,
+                        VisitedBlocks);
+  }
+
+  BlockListTy MinCallerBlocks;
+  for (auto &I : VisitedBlocks)
+    if (I.second.CallerBlock)
+      MinCallerBlocks.push_back(std::move(I.first));
+
+  return rearrangeBB(F, MinCallerBlocks);
+}
+
+SpeculateQuery::ResultTy SequenceBBQuery::operator()(Function &F) {
+  // reduce the number of lists!
+  DenseMap<StringRef, DenseSet<StringRef>> CallerAndCalles;
+  DenseSet<StringRef> Calles;
+  BlockListTy SequencedBlocks;
+  BlockListTy CallerBlocks;
+
+  CallerBlocks = findBBwithCalls(F);
+  if (CallerBlocks.empty())
+    return None;
+
+  if (isStraightLine(F))
+    SequencedBlocks = rearrangeBB(F, CallerBlocks);
+  else
+    SequencedBlocks = queryCFG(F, CallerBlocks);
+
+  for (auto BB : SequencedBlocks)
+    findCalles(BB, Calles);
+
+  CallerAndCalles.insert({F.getName(), std::move(Calles)});
+  return CallerAndCalles;
+}
+
 } // namespace orc
 } // namespace llvm

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 "llvm/Support/Debug.h"
 
 #include <vector>
 
 namespace llvm {
 
 namespace orc {
 
 // ImplSymbolMap methods
 void ImplSymbolMap::trackImpls(SymbolAliasMap ImplMaps, JITDylib *SrcJD) {
   assert(SrcJD && "Tracking on Null Source .impl dylib");
   std::lock_guard<std::mutex> Lockit(ConcurrentAccess);
   for (auto &I : ImplMaps) {
     auto It = Maps.insert({I.first, {I.second.Aliasee, SrcJD}});
     // check rationale when independent dylibs have same symbol name?
     assert(It.second && "ImplSymbols are already tracked for this Symbol?");
     (void)(It);
   }
 }
 
+// Trigger Speculative Compiles.
+void Speculator::speculateForEntryPoint(Speculator *Ptr, uint64_t StubId) {
+  assert(Ptr && " Null Address Received in orc_speculate_for ");
+  Ptr->speculateFor(StubId);
+}
+
+Error Speculator::addSpeculationRuntime(JITDylib &JD,
+                                        MangleAndInterner &Mangle) {
+  JITEvaluatedSymbol ThisPtr(pointerToJITTargetAddress(this),
+                             JITSymbolFlags::Exported);
+  JITEvaluatedSymbol SpeculateForEntryPtr(
+      pointerToJITTargetAddress(&speculateForEntryPoint),
+      JITSymbolFlags::Exported);
+  return JD.define(absoluteSymbols({
+      {Mangle("__orc_speculator"), ThisPtr},                // Data Symbol
+      {Mangle("__orc_speculate_for"), SpeculateForEntryPtr} // Callable Symbol
+  }));
+}
+
 // If two modules, share the same LLVMContext, different threads must
-// not access those modules concurrently, doing so leave the
-// LLVMContext in in-consistent state.
-// But here since each TSM has a unique Context associated with it,
-// on locking is necessary!
+// not access them concurrently without locking the associated LLVMContext
+// this implementation follows this contract.
 void IRSpeculationLayer::emit(MaterializationResponsibility R,
                               ThreadSafeModule TSM) {
 
   assert(TSM && "Speculation Layer received Null Module ?");
   assert(TSM.getContext().getContext() != nullptr &&
          "Module with null LLVMContext?");
 
-  // Instrumentation of runtime calls
-  auto &InContext = *TSM.getContext().getContext();
-  auto SpeculatorVTy = StructType::create(InContext, "Class.Speculator");
+  // Instrumentation of runtime calls, lock the Module
+  TSM.withModuleDo([this, &R](Module &M) {
+    auto &MContext = M.getContext();
+    auto SpeculatorVTy = StructType::create(MContext, "Class.Speculator");
     auto RuntimeCallTy = FunctionType::get(
-      Type::getVoidTy(InContext),
-      {SpeculatorVTy->getPointerTo(), Type::getInt64Ty(InContext)}, false);
+        Type::getVoidTy(MContext),
+        {SpeculatorVTy->getPointerTo(), Type::getInt64Ty(MContext)}, false);
     auto RuntimeCall =
         Function::Create(RuntimeCallTy, Function::LinkageTypes::ExternalLinkage,
-                       "__orc_speculate_for", TSM.getModuleUnlocked());
+                         "__orc_speculate_for", &M);
     auto SpeclAddr = new GlobalVariable(
-      *TSM.getModuleUnlocked(), SpeculatorVTy, false,
-      GlobalValue::LinkageTypes::ExternalLinkage, nullptr, "__orc_speculator");
+        M, SpeculatorVTy, false, GlobalValue::LinkageTypes::ExternalLinkage,
+        nullptr, "__orc_speculator");
 
-  IRBuilder<> Mutator(InContext);
+    IRBuilder<> Mutator(MContext);
 
     // QueryAnalysis allowed to transform the IR source, one such example is
     // Simplify CFG helps the static branch prediction heuristics!
-  for (auto &Fn : TSM.getModuleUnlocked()->getFunctionList()) {
+    for (auto &Fn : M.getFunctionList()) {
       if (!Fn.isDeclaration()) {
-      auto IRNames = QueryAnalysis(Fn, FAM);
+
+        auto IRNames = QueryAnalysis(Fn);
         // Instrument and register if Query has result
         if (IRNames.hasValue()) {
-        Mutator.SetInsertPoint(&(Fn.getEntryBlock().front()));
+
+          // Emit globals for each function.
+          auto LoadValueTy = Type::getInt8Ty(MContext);
+          auto SpeculatorGuard = new GlobalVariable(
+              M, LoadValueTy, false, GlobalValue::LinkageTypes::InternalLinkage,
+              ConstantInt::get(LoadValueTy, 0),
+              "__orc_speculate.guard.for." + Fn.getName());
+          SpeculatorGuard->setAlignment(1);
+          SpeculatorGuard->setUnnamedAddr(GlobalValue::UnnamedAddr::Local);
+
+          BasicBlock &ProgramEntry = Fn.getEntryBlock();
+          // Create BasicBlocks before the program's entry basicblock
+          BasicBlock *SpeculateBlock = BasicBlock::Create(
+              MContext, "__orc_speculate.block", &Fn, &ProgramEntry);
+          BasicBlock *SpeculateDecisionBlock = BasicBlock::Create(
+              MContext, "__orc_speculate.decision.block", &Fn, SpeculateBlock);
+
+          assert(SpeculateDecisionBlock == &Fn.getEntryBlock() &&
+                 "SpeculateDecisionBlock not updated?");
+          Mutator.SetInsertPoint(SpeculateDecisionBlock);
+
+          auto LoadGuard =
+              Mutator.CreateLoad(LoadValueTy, SpeculatorGuard, "guard.value");
+          // if just loaded value equal to 0,return true.
+          auto CanSpeculate =
+              Mutator.CreateICmpEQ(LoadGuard, ConstantInt::get(LoadValueTy, 0),
+                                   "compare.to.speculate");
+          Mutator.CreateCondBr(CanSpeculate, SpeculateBlock, &ProgramEntry);
+
+          Mutator.SetInsertPoint(SpeculateBlock);
           auto ImplAddrToUint =
-            Mutator.CreatePtrToInt(&Fn, Type::getInt64Ty(InContext));
+              Mutator.CreatePtrToInt(&Fn, Type::getInt64Ty(MContext));
           Mutator.CreateCall(RuntimeCallTy, RuntimeCall,
                              {SpeclAddr, ImplAddrToUint});
+          Mutator.CreateStore(ConstantInt::get(LoadValueTy, 1),
+                              SpeculatorGuard);
+          Mutator.CreateBr(&ProgramEntry);
+
+          assert(Mutator.GetInsertBlock()->getParent() == &Fn &&
+                 "IR builder association mismatch?");
           S.registerSymbols(internToJITSymbols(IRNames.getValue()),
                             &R.getTargetJITDylib());
         }
       }
     }
-  // No locking needed read only operation.
-  assert(!(verifyModule(*TSM.getModuleUnlocked())) &&
+  });
+
+  assert(!TSM.withModuleDo([](const Module &M) { return 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