[PGO] fix PGO use ICE when there are unreachable BBs

[xur]

Summary

For -O0 there might be unreachable BBs, which breaks the assumption that all the BBs have an auxiliary data structure. In this patch, I change the interface for getBBInfo() so that a nullptr can be returned for the unreachable BBs (and the callers can ignore those BBs).

This fixes the bug reported
https://llvm.org/bugs/show_bug.cgi?id=31209

[davidxl]

I suggest adding a new interface 'findBBInfo' returns pointer (which can be null) and use it in context that can do null check. Other users of the getBBInfo won't change.

[xur]

Integrated David's suggestions.
Actually there are two more places need to be changed.

-Rong

[xur]

to clarify my previous comments: I meant two more places needed to use the new interface. they were included in patch set 2 that just uploaded.

[davidxl]

lgtm

lib/Transforms/Instrumentation/CFGMST.h

BBInfo *findAndCompressGroup(BBInfo *G) {
     if (G->Group != G)
       G->Group = findAndCompressGroup(static_cast<BBInfo *>(G->Group));
     return static_cast<BBInfo *>(G->Group);
   }
 
   // Union BB1 and BB2 into the same group and return true.
   // Returns false if BB1 and BB2 are already in the same group.
   bool unionGroups(const BasicBlock *BB1, const BasicBlock *BB2) {
-    BBInfo *BB1G = findAndCompressGroup(&getBBInfo(BB1));
-    BBInfo *BB2G = findAndCompressGroup(&getBBInfo(BB2));
+    BBInfo *BB1G = findAndCompressGroup(getBBInfo(BB1));
+    BBInfo *BB2G = findAndCompressGroup(getBBInfo(BB2));
 
     if (BB1G == BB2G)
       return false;
 
     // Make the smaller rank tree a direct child or the root of high rank tree.
     if (BB1G->Rank < BB2G->Rank)
       BB1G->Group = BB2G;
     else {
       BB2G->Group = BB1G;
       // If the ranks are the same, increment root of one tree by one.
       if (BB1G->Rank == BB2G->Rank)
         BB1G->Rank++;
     }
     return true;
   }
 
   // Give BB, return the auxiliary information.
-  BBInfo &getBBInfo(const BasicBlock *BB) const {
+  BBInfo *getBBInfo(const BasicBlock *BB) const {
     auto It = BBInfos.find(BB);
-    assert(It->second.get() != nullptr);
-    return *It->second.get();
+    if (It == BBInfos.end())
+      return nullptr;
+    return It->second.get();
   }
 
   // Traverse the CFG using a stack. Find all the edges and assign the weight.
   // Edges with large weight will be put into MST first so they are less likely
   // to be instrumented.
   void buildEdges() {
     DEBUG(dbgs() << "Build Edge on " << F.getName() << "\n");
 
for (auto &BI : BBInfos) {
       OS << "  BB: " << (BB == nullptr ? "FakeNode" : BB->getName()) << "  "
          << BI.second->infoString() << "\n";
     }
 
     OS << "  Number of Edges: " << AllEdges.size()
        << " (*: Instrument, C: CriticalEdge, -: Removed)\n";
     uint32_t Count = 0;
     for (auto &EI : AllEdges)
-      OS << "  Edge " << Count++ << ": " << getBBInfo(EI->SrcBB).Index << "-->"
-         << getBBInfo(EI->DestBB).Index << EI->infoString() << "\n";
+      OS << "  Edge " << Count++ << ": " << getBBInfo(EI->SrcBB)->Index << "-->"
+         << getBBInfo(EI->DestBB)->Index << EI->infoString() << "\n";
   }
 
   // Add an edge to AllEdges with weight W.
   Edge &addEdge(const BasicBlock *Src, const BasicBlock *Dest, uint64_t W) {
     uint32_t Index = BBInfos.size();
     auto Iter = BBInfos.end();
     bool Inserted;
     std::tie(Iter, Inserted) = BBInfos.insert(std::make_pair(Src, nullptr));

lib/Transforms/Instrumentation/PGOInstrumentation.cpp

public:
   // The Minimum Spanning Tree of function CFG.
   CFGMST<Edge, BBInfo> MST;
 
   // Give an edge, find the BB that will be instrumented.
   // Return nullptr if there is no BB to be instrumented.
   BasicBlock *getInstrBB(Edge *E);
 
   // Return the auxiliary BB information.
-  BBInfo &getBBInfo(const BasicBlock *BB) const { return MST.getBBInfo(BB); }
+  BBInfo *getBBInfo(const BasicBlock *BB) const { return MST.getBBInfo(BB); }
 
   // Dump edges and BB information.
   void dumpInfo(std::string Str = "") const {
     MST.dumpEdges(dbgs(), Twine("Dump Function ") + FuncName + " Hash: " +
                               Twine(FunctionHash) + "\t" + Str);
   }
 
   FuncPGOInstrumentation(

 template <class Edge, class BBInfo>
 void FuncPGOInstrumentation<Edge, BBInfo>::computeCFGHash() {
   std::vector<char> Indexes;
   JamCRC JC;
   for (auto &BB : F) {
     const TerminatorInst *TI = BB.getTerminator();
     for (unsigned I = 0, E = TI->getNumSuccessors(); I != E; ++I) {
       BasicBlock *Succ = TI->getSuccessor(I);
-      uint32_t Index = getBBInfo(Succ).Index;
+      uint32_t Index = getBBInfo(Succ)->Index;
       for (int J = 0; J < 4; J++)
         Indexes.push_back((char)(Index >> (J * 8)));
     }
   }
   JC.update(Indexes);
   FunctionHash = (uint64_t)SIVisitor.getNumOfSelectInsts() << 56 |
                  (uint64_t)IndirectCallSites.size() << 48 |
                  (uint64_t)MST.AllEdges.size() << 32 | JC.getCRC();
BasicBlock *FuncPGOInstrumentation<Edge, BBInfo>::getInstrBB(Edge *E) {
   if (TI->getNumSuccessors() <= 1)
     return SrcBB;
   if (!E->IsCritical)
     return DestBB;
 
   // For a critical edge, we have to split. Instrument the newly
   // created BB.
   NumOfPGOSplit++;
-  DEBUG(dbgs() << "Split critical edge: " << getBBInfo(SrcBB).Index << " --> "
-               << getBBInfo(DestBB).Index << "\n");
+  DEBUG(dbgs() << "Split critical edge: " << getBBInfo(SrcBB)->Index << " --> "
+               << getBBInfo(DestBB)->Index << "\n");
   unsigned SuccNum = GetSuccessorNumber(SrcBB, DestBB);
   BasicBlock *InstrBB = SplitCriticalEdge(TI, SuccNum);
   assert(InstrBB && "Critical edge is not split");
 
   E->Removed = true;
   return InstrBB;
 }
 
public:
   FuncFreqAttr getFuncFreqAttr() const { return FreqAttr; }
 
   // Return the function hash.
   uint64_t getFuncHash() const { return FuncInfo.FunctionHash; }
   // Return the profile record for this function;
   InstrProfRecord &getProfileRecord() { return ProfileRecord; }
 
   // Return the auxiliary BB information.
-  UseBBInfo &getBBInfo(const BasicBlock *BB) const {
+  UseBBInfo *getBBInfo(const BasicBlock *BB) const {
     return FuncInfo.getBBInfo(BB);
   }
 
 private:
   Function &F;
   Module *M;
   // This member stores the shared information with class PGOGenFunc.
   FuncPGOInstrumentation<PGOUseEdge, UseBBInfo> FuncInfo;
void PGOUseFunc::setInstrumentedCounts(
 
   uint32_t I = 0;
   for (auto &E : WorkList) {
     BasicBlock *InstrBB = FuncInfo.getInstrBB(E);
     if (!InstrBB)
       continue;
     uint64_t CountValue = CountFromProfile[I++];
     if (!E->Removed) {
-      getBBInfo(InstrBB).setBBInfoCount(CountValue);
+      getBBInfo(InstrBB)->setBBInfoCount(CountValue);
       E->setEdgeCount(CountValue);
       continue;
     }
 
     // Need to add two new edges.
     BasicBlock *SrcBB = const_cast<BasicBlock *>(E->SrcBB);
     BasicBlock *DestBB = const_cast<BasicBlock *>(E->DestBB);
     // Add new edge of SrcBB->InstrBB.
     PGOUseEdge &NewEdge = FuncInfo.MST.addEdge(SrcBB, InstrBB, 0);
     NewEdge.setEdgeCount(CountValue);
     // Add new edge of InstrBB->DestBB.
     PGOUseEdge &NewEdge1 = FuncInfo.MST.addEdge(InstrBB, DestBB, 0);
     NewEdge1.setEdgeCount(CountValue);
     NewEdge1.InMST = true;
-    getBBInfo(InstrBB).setBBInfoCount(CountValue);
+    getBBInfo(InstrBB)->setBBInfoCount(CountValue);
   }
   ProfileCountSize =  CountFromProfile.size();
   CountPosition = I;
 }
 
 // Set the count value for the unknown edge. There should be one and only one
 // unknown edge in Edges vector.
 void PGOUseFunc::setEdgeCount(DirectEdges &Edges, uint64_t Value) {
   for (auto &E : Edges) {
     if (E->CountValid)
       continue;
     E->setEdgeCount(Value);
 
-    getBBInfo(E->SrcBB).UnknownCountOutEdge--;
-    getBBInfo(E->DestBB).UnknownCountInEdge--;
+    getBBInfo(E->SrcBB)->UnknownCountOutEdge--;
+    getBBInfo(E->DestBB)->UnknownCountInEdge--;
     return;
   }
   llvm_unreachable("Cannot find the unknown count edge");
 }
 
 // Read the profile from ProfileFileName and assign the value to the
 // instrumented BB and the edges. This function also updates ProgramMaxCount.
 // Return true if the profile are successfully read, and false on errors.
bool PGOUseFunc::readCounters(IndexedInstrProfReader *PGOReader) {
   uint64_t ValueSum = 0;
   for (unsigned I = 0, S = CountFromProfile.size(); I < S; I++) {
     DEBUG(dbgs() << "  " << I << ": " << CountFromProfile[I] << "\n");
     ValueSum += CountFromProfile[I];
   }
 
   DEBUG(dbgs() << "SUM =  " << ValueSum << "\n");
 
-  getBBInfo(nullptr).UnknownCountOutEdge = 2;
-  getBBInfo(nullptr).UnknownCountInEdge = 2;
+  getBBInfo(nullptr)->UnknownCountOutEdge = 2;
+  getBBInfo(nullptr)->UnknownCountInEdge = 2;
 
   setInstrumentedCounts(CountFromProfile);
   ProgramMaxCount = PGOReader->getMaximumFunctionCount();
   return true;
 }
 
 // Populate the counters from instrumented BBs to all BBs.
 // In the end of this operation, all BBs should have a valid count value.
 void PGOUseFunc::populateCounters() {
   // First set up Count variable for all BBs.
   for (auto &E : FuncInfo.MST.AllEdges) {
     if (E->Removed)
       continue;
 
     const BasicBlock *SrcBB = E->SrcBB;
     const BasicBlock *DestBB = E->DestBB;
-    UseBBInfo &SrcInfo = getBBInfo(SrcBB);
-    UseBBInfo &DestInfo = getBBInfo(DestBB);
-    SrcInfo.OutEdges.push_back(E.get());
-    DestInfo.InEdges.push_back(E.get());
-    SrcInfo.UnknownCountOutEdge++;
-    DestInfo.UnknownCountInEdge++;
+    UseBBInfo *SrcInfo = getBBInfo(SrcBB);
+    UseBBInfo *DestInfo = getBBInfo(DestBB);
+    SrcInfo->OutEdges.push_back(E.get());
+    DestInfo->InEdges.push_back(E.get());
+    SrcInfo->UnknownCountOutEdge++;
+    DestInfo->UnknownCountInEdge++;
 
     if (!E->CountValid)
       continue;
-    DestInfo.UnknownCountInEdge--;
-    SrcInfo.UnknownCountOutEdge--;
+    DestInfo->UnknownCountInEdge--;
+    SrcInfo->UnknownCountOutEdge--;
   }
 
   bool Changes = true;
   unsigned NumPasses = 0;
   while (Changes) {
     NumPasses++;
     Changes = false;
 
     // For efficient traversal, it's better to start from the end as most
     // of the instrumented edges are at the end.
     for (auto &BB : reverse(F)) {
-      UseBBInfo &Count = getBBInfo(&BB);
-      if (!Count.CountValid) {
-        if (Count.UnknownCountOutEdge == 0) {
-          Count.CountValue = sumEdgeCount(Count.OutEdges);
-          Count.CountValid = true;
+      UseBBInfo *Count = getBBInfo(&BB);
+      if (Count == nullptr)
+        continue;
+      if (!Count->CountValid) {
+        if (Count->UnknownCountOutEdge == 0) {
+          Count->CountValue = sumEdgeCount(Count->OutEdges);
+          Count->CountValid = true;
           Changes = true;
-        } else if (Count.UnknownCountInEdge == 0) {
-          Count.CountValue = sumEdgeCount(Count.InEdges);
-          Count.CountValid = true;
+        } else if (Count->UnknownCountInEdge == 0) {
+          Count->CountValue = sumEdgeCount(Count->InEdges);
+          Count->CountValid = true;
           Changes = true;
         }
       }
-      if (Count.CountValid) {
-        if (Count.UnknownCountOutEdge == 1) {
-          uint64_t Total = Count.CountValue - sumEdgeCount(Count.OutEdges);
-          setEdgeCount(Count.OutEdges, Total);
+      if (Count->CountValid) {
+        if (Count->UnknownCountOutEdge == 1) {
+          uint64_t Total = Count->CountValue - sumEdgeCount(Count->OutEdges);
+          setEdgeCount(Count->OutEdges, Total);
           Changes = true;
         }
-        if (Count.UnknownCountInEdge == 1) {
-          uint64_t Total = Count.CountValue - sumEdgeCount(Count.InEdges);
-          setEdgeCount(Count.InEdges, Total);
+        if (Count->UnknownCountInEdge == 1) {
+          uint64_t Total = Count->CountValue - sumEdgeCount(Count->InEdges);
+          setEdgeCount(Count->InEdges, Total);
           Changes = true;
         }
       }
     }
   }
 
   DEBUG(dbgs() << "Populate counts in " << NumPasses << " passes.\n");
 #ifndef NDEBUG
   // Assert every BB has a valid counter.
-  for (auto &BB : F)
-    assert(getBBInfo(&BB).CountValid && "BB count is not valid");
+  for (auto &BB : F) {
+    auto BI = getBBInfo(&BB);
+    if (BI != nullptr)
+      assert(BI->CountValid && "BB count is not valid");
+  }
 #endif
-  uint64_t FuncEntryCount = getBBInfo(&*F.begin()).CountValue;
+  uint64_t FuncEntryCount = getBBInfo(&*F.begin())->CountValue;
   F.setEntryCount(FuncEntryCount);
   uint64_t FuncMaxCount = FuncEntryCount;
-  for (auto &BB : F)
-    FuncMaxCount = std::max(FuncMaxCount, getBBInfo(&BB).CountValue);
+  for (auto &BB : F) {
+    auto BI = getBBInfo(&BB);
+    if (BI == nullptr)
+      continue;
+    FuncMaxCount = std::max(FuncMaxCount, BI->CountValue);
+  }
   markFunctionAttributes(FuncEntryCount, FuncMaxCount);
 
   // Now annotate select instructions
   FuncInfo.SIVisitor.annotateSelects(F, this, &CountPosition);
   assert(CountPosition == ProfileCountSize);
 
   DEBUG(FuncInfo.dumpInfo("after reading profile."));
 }
void PGOUseFunc::setBranchWeights() {
   // Generate MD_prof metadata for every branch instruction.
   DEBUG(dbgs() << "\nSetting branch weights.\n");
   for (auto &BB : F) {
     TerminatorInst *TI = BB.getTerminator();
     if (TI->getNumSuccessors() < 2)
       continue;
     if (!isa<BranchInst>(TI) && !isa<SwitchInst>(TI))
       continue;
-    if (getBBInfo(&BB).CountValue == 0)
+    if (getBBInfo(&BB)->CountValue == 0)
       continue;
 
     // We have a non-zero Branch BB.
-    const UseBBInfo &BBCountInfo = getBBInfo(&BB);
-    unsigned Size = BBCountInfo.OutEdges.size();
+    const UseBBInfo *BBCountInfo = getBBInfo(&BB);
+    unsigned Size = BBCountInfo->OutEdges.size();
     SmallVector<uint64_t, 2> EdgeCounts(Size, 0);
     uint64_t MaxCount = 0;
     for (unsigned s = 0; s < Size; s++) {
-      const PGOUseEdge *E = BBCountInfo.OutEdges[s];
+      const PGOUseEdge *E = BBCountInfo->OutEdges[s];
       const BasicBlock *SrcBB = E->SrcBB;
       const BasicBlock *DestBB = E->DestBB;
       if (DestBB == nullptr)
         continue;
       unsigned SuccNum = GetSuccessorNumber(SrcBB, DestBB);
       uint64_t EdgeCount = E->CountValue;
       if (EdgeCount > MaxCount)
         MaxCount = EdgeCount;

 
 void SelectInstVisitor::annotateOneSelectInst(SelectInst &SI) {
   std::vector<uint64_t> &CountFromProfile = UseFunc->getProfileRecord().Counts;
   assert(*CurCtrIdx < CountFromProfile.size() &&
          "Out of bound access of counters");
   uint64_t SCounts[2];
   SCounts[0] = CountFromProfile[*CurCtrIdx]; // True count
   ++(*CurCtrIdx);
-  uint64_t TotalCount = UseFunc->getBBInfo(SI.getParent()).CountValue;
+  uint64_t TotalCount = UseFunc->getBBInfo(SI.getParent())->CountValue;
   // False Count
   SCounts[1] = (TotalCount > SCounts[0] ? TotalCount - SCounts[0] : 0);
   uint64_t MaxCount = std::max(SCounts[0], SCounts[1]);
   if (MaxCount)
     setProfMetadata(F.getParent(), &SI, SCounts, MaxCount);
 }
 
 void SelectInstVisitor::visitSelectInst(SelectInst &SI) {

test/Transforms/PGOProfile/Inputs/unreachable_bb.proftext

+# IR level Instrumentation Flag
+:ir
+foo
+# Func Hash:
[phone]
+# Num Counters:
+1
+# Counter Values:
+0

test/Transforms/PGOProfile/unreachable_bb.ll

+; RUN: llvm-profdata merge %S/Inputs/unreachable_bb.proftext -o %t.profdata
+; RUN: opt < %s -pgo-instr-use -pgo-test-profile-file=%t.profdata -S | FileCheck %s --check-prefix=USE
+; RUN: opt < %s -passes=pgo-instr-use -pgo-test-profile-file=%t.profdata -S | FileCheck %s --check-prefix=USE
+
+target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+define void @foo() {
+entry:
+  call void @bar()
+  unreachable
+return:
+  ret void
+}
+
+declare void @bar()
+
+;USE: !0 = !{i32 1, !"ProfileSummary", !1}
+;USE: !1 = !{!2, !3, !4, !5, !6, !7, !8, !9}
+;USE: !2 = !{!"ProfileFormat", !"InstrProf"}
+;USE: !3 = !{!"TotalCount", i64 0}
+
+