[PGO] handle infinite loop properly in pgo instrumention

[davidxl]

Summary

The original fix id D40702 introduces a memory leak -- a surprising side effect of calling 'getAnalysis<>' method. Calling this has the effect of releasing the memory of current BFI and BPI.

The right way is to get the cached BPI and LI from the current BFI.

[vsk]

I don't have my head wrapped around the cause of the leak yet. Was the issue that getAnalysis<LoopInfo...>() recomputed BFI, thereby leaking the existing analysis? If so, would it be sensible for the pass manager to provide a cached BFI instead of allowing the leak? (I'm not suggesting changing the PM, just curious about how these things fit together.)

At a more basic level: Say CFGMST depends on BFI, which in turn depends on {BPI, LI}. Is it possible for CFGMST to be scheduled after a competing pass "DestroysBPI", which preserves BFI but invalidates BPI? If so, how do we avoid that?

[davidxl]

What happens is that the call of getAnalysis<LoopInfoWrapperPass>() is invoked after getAnalysis<BlockFrequencyInfoWrapperPass>(). The legacy pass manager in this case will do the analysis run to compute LoopInfo on the fly. Before running the analysis, it will also invoke the 'release' memory method of the analysis passes contained in, in this case is BPI and BFI. After that we basically have use-after free problem. The leaking is probably just a side effect.

The fix simply fetches the cached LI and BPI from BFI. This is a common pattern used (for getting BPI from BFI). This is only an issue with the old PM -- as it will go away soon, we don't need to worry about fixing that.

[vsk]

In D40873#946971, @davidxl wrote:

What happens is that the call of getAnalysis<LoopInfoWrapperPass>() is invoked after getAnalysis<BlockFrequencyInfoWrapperPass>(). The legacy pass manager in this case will do the analysis run to compute LoopInfo on the fly. Before running the analysis, it will also invoke the 'release' memory method of the analysis passes contained in, in this case is BPI and BFI. After that we basically have use-after free problem. The leaking is probably just a side effect.

Ah ok, this makes a lot more sense to me as a use-after-free. Thanks, lgtm.

llvm/trunk/include/llvm/Analysis/BlockFrequencyInfo.h

public:
   ~BlockFrequencyInfo();
 
   /// Handle invalidation explicitly.
   bool invalidate(Function &F, const PreservedAnalyses &PA,
                   FunctionAnalysisManager::Invalidator &);
 
   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
   /// information. Please note that initial frequency is equal to ENTRY_FREQ. It
   /// means that we should not rely on the value itself, but only on the
   /// comparison to the other block frequencies. We do this to avoid using of
   /// floating points.
   BlockFrequency getBlockFreq(const BasicBlock *BB) const;

llvm/trunk/include/llvm/Analysis/BlockFrequencyInfoImpl.h

public:
 
   void setBlockFreq(const BlockT *BB, uint64_t Freq);
 
   Scaled64 getFloatingBlockFreq(const BlockT *BB) const {
     return BlockFrequencyInfoImplBase::getFloatingBlockFreq(getNode(BB));
   }
 
   const BranchProbabilityInfoT &getBPI() const { return *BPI; }
+  const LoopInfoT &getLoopInfo() const { return *LI; }
 
   /// \brief Print the frequencies for the current function.
   ///
   /// Prints the frequencies for the blocks in the current function.
   ///
   /// Blocks are printed in the natural iteration order of the function, rather
   /// than reverse post-order.  This provides two advantages:  writing -analyze
   /// tests is easier (since blocks come out in source order), and even

llvm/trunk/lib/Analysis/BlockFrequencyInfo.cpp

 const Function *BlockFrequencyInfo::getFunction() const {
   return BFI ? BFI->getFunction() : nullptr;
 }
 
 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;
 }
 
 raw_ostream &
 BlockFrequencyInfo::printBlockFreq(raw_ostream &OS,
                                    const BasicBlock *BB) const {

llvm/trunk/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;
Edge &addEdge(const BasicBlock *Src, const BasicBlock *Dest, uint64_t W) {
     std::tie(Iter, Inserted) = BBInfos.insert(std::make_pair(Dest, nullptr));
     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;
+  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();
   }
 };
 
 } // end namespace llvm
 
 #undef DEBUG_TYPE // "cfgmst"
 
 #endif // LLVM_LIB_TRANSFORMS_INSTRUMENTATION_CFGMST_H

llvm/trunk/lib/Transforms/Instrumentation/PGOInstrumentation.cpp

 
 } // 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_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_END(PGOInstrumentationUseLegacyPass, "pgo-instr-use",
                     "Read PGO instrumentation profile.", false, false)
 
 ModulePass *llvm::createPGOInstrumentationUseLegacyPass(StringRef Filename) {
   return new PGOInstrumentationUseLegacyPass(Filename.str());
 }
 
 namespace {
public:
   void dumpInfo(std::string Str = "") const {
     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)
+      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);
     NumOfPGOSelectInsts += SIVisitor.getNumOfSelectInsts();
     NumOfPGOMemIntrinsics += MIVisitor.getNumOfMemIntrinsics();
     ValueSites[IPVK_IndirectCallTarget] = findIndirectCallSites(Func);
     ValueSites[IPVK_MemOPSize] = MIVisitor.findMemIntrinsics(Func);
 
BasicBlock *FuncPGOInstrumentation<Edge, BBInfo>::getInstrBB(Edge *E) {
 
   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,
+    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;
   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)
       : 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);
 
   // Populate the counts for all BBs.
   void populateCounters();
 
   // Set the branch weights based on the count values.
for (GlobalVariable &GV : M.globals())
     if (Comdat *C = GV.getComdat())
       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) {
+    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;
 }
 
 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();
   // 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(
static bool annotateAllFunctions(
 
   std::unordered_multimap<Comdat *, GlobalValue *> ComdatMembers;
   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);
+    PGOUseFunc Func(F, &M, ComdatMembers, BFI);
     if (!Func.readCounters(PGOReader.get()))
       continue;
     Func.populateCounters();
     Func.setBranchWeights();
     Func.annotateValueSites();
     Func.annotateIrrLoopHeaderWeights();
     PGOUseFunc::FuncFreqAttr FreqAttr = Func.getFuncFreqAttr();
     if (FreqAttr == PGOUseFunc::FFA_Cold)
PGOInstrumentationUse::PGOInstrumentationUse(std::string Filename)
   if (!PGOTestProfileFile.empty())
     ProfileFileName = PGOTestProfileFile;
 }
 
 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();
 }
 
 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) {
   if (!Node->getName().empty())
     return Node->getName();
 
   std::string SimpleNodeName;
   raw_string_ostream OS(SimpleNodeName);

llvm/trunk/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 }
+