[PGO] detect infinite loop and form MST properly

Differential Revision: http://reviews.llvm.org/D40873

llvm-svn: 320104

llvm/include/llvm/Analysis/BlockFrequencyInfo.h

@@ -56,6 +56,7 @@ class BlockFrequencyInfo {
 
   const Function *getFunction() const;
   const BranchProbabilityInfo *getBPI() const;
+  const LoopInfo *getLoopInfo() const;
   void view() const;
 
   /// getblockFreq - Return block frequency. Return 0 if we don't have the

llvm/include/llvm/Analysis/BlockFrequencyInfoImpl.h

@@ -993,6 +993,7 @@ template <class BT> class BlockFrequencyInfoImpl : BlockFrequencyInfoImplBase {
   }
 
   const BranchProbabilityInfoT &getBPI() const { return *BPI; }
+  const LoopInfoT &getLoopInfo() const { return *LI; }
 
   /// \brief Print the frequencies for the current function.
   ///

llvm/lib/Analysis/BlockFrequencyInfo.cpp

@@ -264,6 +264,10 @@ const BranchProbabilityInfo *BlockFrequencyInfo::getBPI() const {
   return BFI ? &BFI->getBPI() : nullptr;
 }
 
+const LoopInfo *BlockFrequencyInfo::getLoopInfo() const {
+  return BFI ? &BFI->getLoopInfo() : nullptr;
+}
+
 raw_ostream &BlockFrequencyInfo::
 printBlockFreq(raw_ostream &OS, const BlockFrequency Freq) const {
   return BFI ? BFI->printBlockFreq(OS, Freq) : OS;

llvm/lib/Transforms/Instrumentation/CFGMST.h

@@ -20,6 +20,7 @@
 #include "llvm/Analysis/BlockFrequencyInfo.h"
 #include "llvm/Analysis/BranchProbabilityInfo.h"
 #include "llvm/Analysis/CFG.h"
+#include "llvm/Analysis/LoopInfo.h"
 #include "llvm/Support/BranchProbability.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
@@ -136,6 +137,21 @@ template <class Edge, class BBInfo> class CFGMST {
                      << " w = " << BBWeight << "\n");
       }
     }
+    // check if there is any infinite loop. If yes, add a fake edge from
+    // the header block to the fake node:
+    for (auto *L : *LI) {
+      SmallVector<BasicBlock *, 2> ExitingBlocks;
+      L->getExitingBlocks(ExitingBlocks);
+      if (!ExitingBlocks.empty())
+        continue;
+      auto *HB = L->getHeader();
+      if (!HB)
+        continue;
+      addEdge(HB, nullptr, UINT64_MAX);
+      DEBUG(dbgs() << "  Edge: from infinite loop header " << HB->getName()
+                   << " to exit"
+                   << " w = " << UINT64_MAX << "\n");
+    }
   }
 
   // Sort CFG edges based on its weight.
@@ -210,13 +226,13 @@ template <class Edge, class BBInfo> class CFGMST {
     return *AllEdges.back();
   }
 
-  BranchProbabilityInfo *BPI;
+  const BranchProbabilityInfo *BPI;
   BlockFrequencyInfo *BFI;
+  const LoopInfo *LI;
 
 public:
-  CFGMST(Function &Func, BranchProbabilityInfo *BPI_ = nullptr,
-         BlockFrequencyInfo *BFI_ = nullptr)
-      : F(Func), BPI(BPI_), BFI(BFI_) {
+  CFGMST(Function &Func, BlockFrequencyInfo *BFI_ = nullptr)
+      : F(Func), BPI(BFI_->getBPI()), BFI(BFI_), LI(BFI_->getLoopInfo()) {
     buildEdges();
     sortEdgesByWeight();
     computeMinimumSpanningTree();

llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp

@@ -425,7 +425,6 @@ char PGOInstrumentationGenLegacyPass::ID = 0;
 INITIALIZE_PASS_BEGIN(PGOInstrumentationGenLegacyPass, "pgo-instr-gen",
                       "PGO instrumentation.", false, false)
 INITIALIZE_PASS_DEPENDENCY(BlockFrequencyInfoWrapperPass)
-INITIALIZE_PASS_DEPENDENCY(BranchProbabilityInfoWrapperPass)
 INITIALIZE_PASS_END(PGOInstrumentationGenLegacyPass, "pgo-instr-gen",
                     "PGO instrumentation.", false, false)
 
@@ -438,7 +437,6 @@ char PGOInstrumentationUseLegacyPass::ID = 0;
 INITIALIZE_PASS_BEGIN(PGOInstrumentationUseLegacyPass, "pgo-instr-use",
                       "Read PGO instrumentation profile.", false, false)
 INITIALIZE_PASS_DEPENDENCY(BlockFrequencyInfoWrapperPass)
-INITIALIZE_PASS_DEPENDENCY(BranchProbabilityInfoWrapperPass)
 INITIALIZE_PASS_END(PGOInstrumentationUseLegacyPass, "pgo-instr-use",
                     "Read PGO instrumentation profile.", false, false)
 
@@ -529,10 +527,9 @@ template <class Edge, class BBInfo> class FuncPGOInstrumentation {
   FuncPGOInstrumentation(
       Function &Func,
       std::unordered_multimap<Comdat *, GlobalValue *> &ComdatMembers,
-      bool CreateGlobalVar = false, BranchProbabilityInfo *BPI = nullptr,
-      BlockFrequencyInfo *BFI = nullptr)
+      bool CreateGlobalVar = false, BlockFrequencyInfo *BFI = nullptr)
       : F(Func), ComdatMembers(ComdatMembers), ValueSites(IPVK_Last + 1),
-        SIVisitor(Func), MIVisitor(Func), MST(F, BPI, BFI) {
+        SIVisitor(Func), MIVisitor(Func), MST(F, BFI) {
     // This should be done before CFG hash computation.
     SIVisitor.countSelects(Func);
     MIVisitor.countMemIntrinsics(Func);
@@ -714,10 +711,9 @@ BasicBlock *FuncPGOInstrumentation<Edge, BBInfo>::getInstrBB(Edge *E) {
 // Visit all edge and instrument the edges not in MST, and do value profiling.
 // Critical edges will be split.
 static void instrumentOneFunc(
-    Function &F, Module *M, BranchProbabilityInfo *BPI, BlockFrequencyInfo *BFI,
+    Function &F, Module *M, BlockFrequencyInfo *BFI,
     std::unordered_multimap<Comdat *, GlobalValue *> &ComdatMembers) {
-  FuncPGOInstrumentation<PGOEdge, BBInfo> FuncInfo(F, ComdatMembers, true, BPI,
-                                                   BFI);
+  FuncPGOInstrumentation<PGOEdge, BBInfo> FuncInfo(F, ComdatMembers, true, BFI);
   unsigned NumCounters = FuncInfo.getNumCounters();
 
   uint32_t I = 0;
@@ -843,11 +839,9 @@ class PGOUseFunc {
 public:
   PGOUseFunc(Function &Func, Module *Modu,
              std::unordered_multimap<Comdat *, GlobalValue *> &ComdatMembers,
-             BranchProbabilityInfo *BPI = nullptr,
              BlockFrequencyInfo *BFIin = nullptr)
       : F(Func), M(Modu), BFI(BFIin),
-        FuncInfo(Func, ComdatMembers, false, BPI, BFIin),
-        FreqAttr(FFA_Normal) {}
+        FuncInfo(Func, ComdatMembers, false, BFIin), FreqAttr(FFA_Normal) {}
 
   // Read counts for the instrumented BB from profile.
   bool readCounters(IndexedInstrProfReader *PGOReader);
@@ -1378,18 +1372,16 @@ static void collectComdatMembers(
 }
 
 static bool InstrumentAllFunctions(
-    Module &M, function_ref<BranchProbabilityInfo *(Function &)> LookupBPI,
-    function_ref<BlockFrequencyInfo *(Function &)> LookupBFI) {
+    Module &M, function_ref<BlockFrequencyInfo *(Function &)> LookupBFI) {
   createIRLevelProfileFlagVariable(M);
   std::unordered_multimap<Comdat *, GlobalValue *> ComdatMembers;
   collectComdatMembers(M, ComdatMembers);
 
   for (auto &F : M) {
     if (F.isDeclaration())
       continue;
-    auto *BPI = LookupBPI(F);
     auto *BFI = LookupBFI(F);
-    instrumentOneFunc(F, &M, BPI, BFI, ComdatMembers);
+    instrumentOneFunc(F, &M, BFI, ComdatMembers);
   }
   return true;
 }
@@ -1398,35 +1390,28 @@ bool PGOInstrumentationGenLegacyPass::runOnModule(Module &M) {
   if (skipModule(M))
     return false;
 
-  auto LookupBPI = [this](Function &F) {
-    return &this->getAnalysis<BranchProbabilityInfoWrapperPass>(F).getBPI();
-  };
   auto LookupBFI = [this](Function &F) {
     return &this->getAnalysis<BlockFrequencyInfoWrapperPass>(F).getBFI();
   };
-  return InstrumentAllFunctions(M, LookupBPI, LookupBFI);
+  return InstrumentAllFunctions(M, LookupBFI);
 }
 
 PreservedAnalyses PGOInstrumentationGen::run(Module &M,
                                              ModuleAnalysisManager &AM) {
   auto &FAM = AM.getResult<FunctionAnalysisManagerModuleProxy>(M).getManager();
-  auto LookupBPI = [&FAM](Function &F) {
-    return &FAM.getResult<BranchProbabilityAnalysis>(F);
-  };
 
   auto LookupBFI = [&FAM](Function &F) {
     return &FAM.getResult<BlockFrequencyAnalysis>(F);
   };
 
-  if (!InstrumentAllFunctions(M, LookupBPI, LookupBFI))
+  if (!InstrumentAllFunctions(M, LookupBFI))
     return PreservedAnalyses::all();
 
   return PreservedAnalyses::none();
 }
 
 static bool annotateAllFunctions(
     Module &M, StringRef ProfileFileName,
-    function_ref<BranchProbabilityInfo *(Function &)> LookupBPI,
     function_ref<BlockFrequencyInfo *(Function &)> LookupBFI) {
   DEBUG(dbgs() << "Read in profile counters: ");
   auto &Ctx = M.getContext();
@@ -1461,9 +1446,8 @@ static bool annotateAllFunctions(
   for (auto &F : M) {
     if (F.isDeclaration())
       continue;
-    auto *BPI = LookupBPI(F);
     auto *BFI = LookupBFI(F);
-    PGOUseFunc Func(F, &M, ComdatMembers, BPI, BFI);
+    PGOUseFunc Func(F, &M, ComdatMembers, BFI);
     if (!Func.readCounters(PGOReader.get()))
       continue;
     Func.populateCounters();
@@ -1531,15 +1515,12 @@ PreservedAnalyses PGOInstrumentationUse::run(Module &M,
                                              ModuleAnalysisManager &AM) {
 
   auto &FAM = AM.getResult<FunctionAnalysisManagerModuleProxy>(M).getManager();
-  auto LookupBPI = [&FAM](Function &F) {
-    return &FAM.getResult<BranchProbabilityAnalysis>(F);
-  };
 
   auto LookupBFI = [&FAM](Function &F) {
     return &FAM.getResult<BlockFrequencyAnalysis>(F);
   };
 
-  if (!annotateAllFunctions(M, ProfileFileName, LookupBPI, LookupBFI))
+  if (!annotateAllFunctions(M, ProfileFileName, LookupBFI))
     return PreservedAnalyses::all();
 
   return PreservedAnalyses::none();
@@ -1549,14 +1530,11 @@ bool PGOInstrumentationUseLegacyPass::runOnModule(Module &M) {
   if (skipModule(M))
     return false;
 
-  auto LookupBPI = [this](Function &F) {
-    return &this->getAnalysis<BranchProbabilityInfoWrapperPass>(F).getBPI();
-  };
   auto LookupBFI = [this](Function &F) {
     return &this->getAnalysis<BlockFrequencyInfoWrapperPass>(F).getBFI();
   };
 
-  return annotateAllFunctions(M, ProfileFileName, LookupBPI, LookupBFI);
+  return annotateAllFunctions(M, ProfileFileName, LookupBFI);
 }
 
 static std::string getSimpleNodeName(const BasicBlock *Node) {

llvm/test/Transforms/PGOProfile/infinite_loop_gen.ll

@@ -0,0 +1,18 @@
+; RUN: opt < %s -pgo-instr-gen -S -o -   | FileCheck %s
+
+define void @foo() {
+entry:
+  br label %while.body
+; CHECK: llvm.instrprof.increment
+
+  while.body:                                       ; preds = %entry, %while.body
+; CHECK: llvm.instrprof.increment
+    call void (...) @bar() #2
+    br label %while.body
+}
+
+declare void @bar(...)
+
+
+attributes #0 = { nounwind }
+