[PGO] More fix to infinite loop profiling

[davidxl]

Summary

The IR instrumentation uses MST to decide instrumentation points. It relies adding a fake edge from the function exit block (with no successors) to the fake node. Otherwise edges needed to be instrumented may end up in MST (and therefore not instrumented).

With infinite loop, the exit block may not exist thus the instrumentation will be wrong. Without the fix, the runtime test case added in r319463 will have zero entry count for main (or any parent function of an infinite loop).

After this is fix, there will be more fixes needed in BFI to avoid the body count from being dropped.

[vsk]

This fix lgtm, but I've never touched this code, so it'd be worth getting another +1.

lib/Transforms/Instrumentation/CFGMST.h

144

Nit: it might be nice to add a hasExitingBlocks() predicate to LoopInfo, since it wouldn't do lower-bound(# ExitingBlocks) work.

[xur]

Looks good to me.

lib/Transforms/Instrumentation/CFGMST.h

 #ifndef LLVM_LIB_TRANSFORMS_INSTRUMENTATION_CFGMST_H
 #define LLVM_LIB_TRANSFORMS_INSTRUMENTATION_CFGMST_H
 
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/STLExtras.h"
 #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"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
 #include <utility>
 #include <vector>
 
 #define DEBUG_TYPE "cfgmst"
for (Function::iterator BB = F.begin(), E = F.end(); BB != E; ++BB) {
                        << TargetBB->getName() << "  w=" << Weight << "\n");
         }
       } else {
         addEdge(&*BB, nullptr, BBWeight);
         DEBUG(dbgs() << "  Edge: from " << BB->getName() << " to exit"
                      << " 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.
   void sortEdgesByWeight() {
     std::stable_sort(AllEdges.begin(), AllEdges.end(),
                      [](const std::unique_ptr<Edge> &Edge1,
                         const std::unique_ptr<Edge> &Edge2) {
                        return Edge1->Weight > Edge2->Weight;
if (Inserted)
       // Newly inserted, update the real info.
       Iter->second = std::move(llvm::make_unique<BBInfo>(Index));
     AllEdges.emplace_back(new Edge(Src, Dest, W));
     return *AllEdges.back();
   }
 
   BranchProbabilityInfo *BPI;
   BlockFrequencyInfo *BFI;
+  LoopInfo *LI;
 
 public:
   CFGMST(Function &Func, BranchProbabilityInfo *BPI_ = nullptr,
-         BlockFrequencyInfo *BFI_ = nullptr)
-      : F(Func), BPI(BPI_), BFI(BFI_) {
+         BlockFrequencyInfo *BFI_ = nullptr,
+         LoopInfo *LI_ = nullptr)
+      : F(Func), BPI(BPI_), BFI(BFI_), LI(LI_) {
     buildEdges();
     sortEdgesByWeight();
     computeMinimumSpanningTree();
   }
 };
 
 } // end namespace llvm
 
 #undef DEBUG_TYPE // "cfgmst"
 
 #endif // LLVM_LIB_TRANSFORMS_INSTRUMENTATION_CFGMST_H

lib/Transforms/Instrumentation/PGOInstrumentation.cpp

 
   StringRef getPassName() const override { return "PGOInstrumentationGenPass"; }
 
 private:
   bool runOnModule(Module &M) override;
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.addRequired<BlockFrequencyInfoWrapperPass>();
+    AU.addRequired<LoopInfoWrapperPass>();
   }
 };
 
 class PGOInstrumentationUseLegacyPass : public ModulePass {
 public:
   static char ID;
 
   // Provide the profile filename as the parameter.

 
 private:
   std::string ProfileFileName;
 
   bool runOnModule(Module &M) override;
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.addRequired<BlockFrequencyInfoWrapperPass>();
+    AU.addRequired<LoopInfoWrapperPass>();
   }
 };
 
 } // end anonymous namespace
 
 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_DEPENDENCY(LoopInfoWrapperPass)
 INITIALIZE_PASS_END(PGOInstrumentationGenLegacyPass, "pgo-instr-gen",
                     "PGO instrumentation.", false, false)
 
 ModulePass *llvm::createPGOInstrumentationGenLegacyPass() {
   return new PGOInstrumentationGenLegacyPass();
 }
 
 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_DEPENDENCY(LoopInfoWrapperPass)
 INITIALIZE_PASS_END(PGOInstrumentationUseLegacyPass, "pgo-instr-use",
                     "Read PGO instrumentation profile.", false, false)
 
 ModulePass *llvm::createPGOInstrumentationUseLegacyPass(StringRef Filename) {
   return new PGOInstrumentationUseLegacyPass(Filename.str());
 }
 
 namespace {

     MST.dumpEdges(dbgs(), Twine("Dump Function ") + FuncName + " Hash: " +
                               Twine(FunctionHash) + "\t" + Str);
   }
 
   FuncPGOInstrumentation(
       Function &Func,
       std::unordered_multimap<Comdat *, GlobalValue *> &ComdatMembers,
       bool CreateGlobalVar = false, BranchProbabilityInfo *BPI = nullptr,
-      BlockFrequencyInfo *BFI = nullptr)
+      BlockFrequencyInfo *BFI = nullptr, LoopInfo *LI = nullptr)
       : F(Func), ComdatMembers(ComdatMembers), ValueSites(IPVK_Last + 1),
-        SIVisitor(Func), MIVisitor(Func), MST(F, BPI, BFI) {
+        SIVisitor(Func), MIVisitor(Func), MST(F, BPI, BFI, LI) {
     // This should be done before CFG hash computation.
     SIVisitor.countSelects(Func);
     MIVisitor.countMemIntrinsics(Func);
     NumOfPGOSelectInsts += SIVisitor.getNumOfSelectInsts();
     NumOfPGOMemIntrinsics += MIVisitor.getNumOfMemIntrinsics();
     ValueSites[IPVK_IndirectCallTarget] = findIndirectCallSites(Func);
     ValueSites[IPVK_MemOPSize] = MIVisitor.findMemIntrinsics(Func);
 

   E->Removed = true;
   return InstrBB;
 }
 
 // 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,
+    LoopInfo *LI,
     std::unordered_multimap<Comdat *, GlobalValue *> &ComdatMembers) {
   FuncPGOInstrumentation<PGOEdge, BBInfo> FuncInfo(F, ComdatMembers, true, BPI,
-                                                   BFI);
+                                                   BFI, LI);
   unsigned NumCounters = FuncInfo.getNumCounters();
 
   uint32_t I = 0;
   Type *I8PtrTy = Type::getInt8PtrTy(M->getContext());
   for (auto &E : FuncInfo.MST.AllEdges) {
     BasicBlock *InstrBB = FuncInfo.getInstrBB(E.get());
     if (!InstrBB)
       continue;

 
 namespace {
 
 class PGOUseFunc {
 public:
   PGOUseFunc(Function &Func, Module *Modu,
              std::unordered_multimap<Comdat *, GlobalValue *> &ComdatMembers,
              BranchProbabilityInfo *BPI = nullptr,
-             BlockFrequencyInfo *BFIin = nullptr)
+             BlockFrequencyInfo *BFIin = nullptr,
+             LoopInfo *LI = nullptr)
       : F(Func), M(Modu), BFI(BFIin),
-        FuncInfo(Func, ComdatMembers, false, BPI, BFIin),
+        FuncInfo(Func, ComdatMembers, false, BPI, BFIin, LI),
         FreqAttr(FFA_Normal) {}
 
   // Read counts for the instrumented BB from profile.
   bool readCounters(IndexedInstrProfReader *PGOReader);
 
   // Populate the counts for all BBs.
   void populateCounters();
 

       ComdatMembers.insert(std::make_pair(C, &GV));
   for (GlobalAlias &GA : M.aliases())
     if (Comdat *C = GA.getComdat())
       ComdatMembers.insert(std::make_pair(C, &GA));
 }
 
 static bool InstrumentAllFunctions(
     Module &M, function_ref<BranchProbabilityInfo *(Function &)> LookupBPI,
-    function_ref<BlockFrequencyInfo *(Function &)> LookupBFI) {
+    function_ref<BlockFrequencyInfo *(Function &)> LookupBFI,
+    function_ref<LoopInfo *(Function &)> LookupLI) {
   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);
+    auto *LI = LookupLI(F);
+    instrumentOneFunc(F, &M, BPI, BFI, LI, ComdatMembers);
   }
   return true;
 }
 
 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);
+  auto LookupLI = [this](Function &F) {
+    return &this->getAnalysis<LoopInfoWrapperPass>(F).getLoopInfo();
+  };
+  return InstrumentAllFunctions(M, LookupBPI, LookupBFI, LookupLI);
 }
 
 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))
+  auto LookupLI = [&FAM](Function &F) {
+    return &FAM.getResult<LoopAnalysis>(F);
+  };
+
+  if (!InstrumentAllFunctions(M, LookupBPI, LookupBFI, LookupLI))
     return PreservedAnalyses::all();
 
   return PreservedAnalyses::none();
 }
 
 static bool annotateAllFunctions(
     Module &M, StringRef ProfileFileName,
     function_ref<BranchProbabilityInfo *(Function &)> LookupBPI,
-    function_ref<BlockFrequencyInfo *(Function &)> LookupBFI) {
+    function_ref<BlockFrequencyInfo *(Function &)> LookupBFI,
+    function_ref<LoopInfo *(Function &)> LookupLI) {
   DEBUG(dbgs() << "Read in profile counters: ");
   auto &Ctx = M.getContext();
   // Read the counter array from file.
   auto ReaderOrErr = IndexedInstrProfReader::create(ProfileFileName);
   if (Error E = ReaderOrErr.takeError()) {
     handleAllErrors(std::move(E), [&](const ErrorInfoBase &EI) {
       Ctx.diagnose(
           DiagnosticInfoPGOProfile(ProfileFileName.data(), EI.message()));

   collectComdatMembers(M, ComdatMembers);
   std::vector<Function *> HotFunctions;
   std::vector<Function *> ColdFunctions;
   for (auto &F : M) {
     if (F.isDeclaration())
       continue;
     auto *BPI = LookupBPI(F);
     auto *BFI = LookupBFI(F);
-    PGOUseFunc Func(F, &M, ComdatMembers, BPI, BFI);
+    auto *LI = LookupLI(F);
+    PGOUseFunc Func(F, &M, ComdatMembers, BPI, BFI, LI);
     if (!Func.readCounters(PGOReader.get()))
       continue;
     Func.populateCounters();
     Func.setBranchWeights();
     Func.annotateValueSites();
     Func.annotateIrrLoopHeaderWeights();
     PGOUseFunc::FuncFreqAttr FreqAttr = Func.getFuncFreqAttr();
     if (FreqAttr == PGOUseFunc::FFA_Cold)

   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))
+  auto LookupLI = [&FAM](Function &F) {
+    return &FAM.getResult<LoopAnalysis>(F);
+  };
+
+  if (!annotateAllFunctions(M, ProfileFileName, LookupBPI, LookupBFI, LookupLI))
     return PreservedAnalyses::all();
 
   return PreservedAnalyses::none();
 }
 
 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();
   };
+  auto LookupLI = [this](Function &F) {
+    return &this->getAnalysis<LoopInfoWrapperPass>(F).getLoopInfo();
+  };
 
-  return annotateAllFunctions(M, ProfileFileName, LookupBPI, LookupBFI);
+  return annotateAllFunctions(M, ProfileFileName, LookupBPI, LookupBFI, LookupLI);
 }
 
 static std::string getSimpleNodeName(const BasicBlock *Node) {
   if (!Node->getName().empty())
     return Node->getName();
 
   std::string SimpleNodeName;
   raw_string_ostream OS(SimpleNodeName);

test/Transforms/PGOProfile/infinite_loop_gen.ll

+; 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 }
+