[nfc][instr] Encapsulate CFGMST (#72207)

Very little of it needs to be public.

Author: mtrofin

llvm/include/llvm/Transforms/Instrumentation/CFGMST.h

@@ -35,7 +35,6 @@ namespace llvm {
 /// Implements a Union-find algorithm to compute Minimum Spanning Tree
 /// for a given CFG.
 template <class Edge, class BBInfo> class CFGMST {
-public:
   Function &F;
 
   // Store all the edges in CFG. It may contain some stale edges
@@ -49,6 +48,12 @@ template <class Edge, class BBInfo> class CFGMST {
   // (For function with an infinite loop, this block may be absent)
   bool ExitBlockFound = false;
 
+  BranchProbabilityInfo *const BPI;
+  BlockFrequencyInfo *const BFI;
+
+  // If function entry will be always instrumented.
+  const bool InstrumentFuncEntry;
+
   // Find the root group of the G and compress the path from G to the root.
   BBInfo *findAndCompressGroup(BBInfo *G) {
     if (G->Group != G)
@@ -77,21 +82,6 @@ template <class Edge, class BBInfo> class CFGMST {
     return true;
   }
 
-  // Give BB, return the auxiliary information.
-  BBInfo &getBBInfo(const BasicBlock *BB) const {
-    auto It = BBInfos.find(BB);
-    assert(It->second.get() != nullptr);
-    return *It->second.get();
-  }
-
-  // Give BB, return the auxiliary information if it's available.
-  BBInfo *findBBInfo(const BasicBlock *BB) const {
-    auto It = BBInfos.find(BB);
-    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.
@@ -236,6 +226,7 @@ template <class Edge, class BBInfo> class CFGMST {
     }
   }
 
+public:
   // Dump the Debug information about the instrumentation.
   void dumpEdges(raw_ostream &OS, const Twine &Message) const {
     if (!Message.str().empty())
@@ -274,25 +265,42 @@ template <class Edge, class BBInfo> class CFGMST {
     return *AllEdges.back();
   }
 
-  BranchProbabilityInfo *BPI;
-  BlockFrequencyInfo *BFI;
-
-  // If function entry will be always instrumented.
-  bool InstrumentFuncEntry;
-
-public:
-  CFGMST(Function &Func, bool InstrumentFuncEntry_,
-         BranchProbabilityInfo *BPI_ = nullptr,
-         BlockFrequencyInfo *BFI_ = nullptr)
-      : F(Func), BPI(BPI_), BFI(BFI_),
-        InstrumentFuncEntry(InstrumentFuncEntry_) {
+  CFGMST(Function &Func, bool InstrumentFuncEntry,
+         BranchProbabilityInfo *BPI = nullptr,
+         BlockFrequencyInfo *BFI = nullptr)
+      : F(Func), BPI(BPI), BFI(BFI), InstrumentFuncEntry(InstrumentFuncEntry) {
     buildEdges();
     sortEdgesByWeight();
     computeMinimumSpanningTree();
     if (AllEdges.size() > 1 && InstrumentFuncEntry)
       std::iter_swap(std::move(AllEdges.begin()),
                      std::move(AllEdges.begin() + AllEdges.size() - 1));
   }
+
+  const std::vector<std::unique_ptr<Edge>> &allEdges() const {
+    return AllEdges;
+  }
+
+  std::vector<std::unique_ptr<Edge>> &allEdges() { return AllEdges; }
+
+  size_t numEdges() const { return AllEdges.size(); }
+
+  size_t bbInfoSize() const { return BBInfos.size(); }
+
+  // Give BB, return the auxiliary information.
+  BBInfo &getBBInfo(const BasicBlock *BB) const {
+    auto It = BBInfos.find(BB);
+    assert(It->second.get() != nullptr);
+    return *It->second.get();
+  }
+
+  // Give BB, return the auxiliary information if it's available.
+  BBInfo *findBBInfo(const BasicBlock *BB) const {
+    auto It = BBInfos.find(BB);
+    if (It == BBInfos.end())
+      return nullptr;
+    return It->second.get();
+  }
 };
 
 } // end namespace llvm

llvm/lib/Transforms/Instrumentation/GCOVProfiling.cpp

@@ -750,7 +750,7 @@ static BasicBlock *getInstrBB(CFGMST<Edge, BBInfo> &MST, Edge &E,
 #ifndef NDEBUG
 static void dumpEdges(CFGMST<Edge, BBInfo> &MST, GCOVFunction &GF) {
   size_t ID = 0;
-  for (auto &E : make_pointee_range(MST.AllEdges)) {
+  for (const auto &E : make_pointee_range(MST.allEdges())) {
     GCOVBlock &Src = E.SrcBB ? GF.getBlock(E.SrcBB) : GF.getEntryBlock();
     GCOVBlock &Dst = E.DestBB ? GF.getBlock(E.DestBB) : GF.getReturnBlock();
     dbgs() << "  Edge " << ID++ << ": " << Src.Number << "->" << Dst.Number
@@ -820,8 +820,8 @@ bool GCOVProfiler::emitProfileNotes(
       CFGMST<Edge, BBInfo> MST(F, /*InstrumentFuncEntry_=*/false, BPI, BFI);
 
       // getInstrBB can split basic blocks and push elements to AllEdges.
-      for (size_t I : llvm::seq<size_t>(0, MST.AllEdges.size())) {
-        auto &E = *MST.AllEdges[I];
+      for (size_t I : llvm::seq<size_t>(0, MST.numEdges())) {
+        auto &E = *MST.allEdges()[I];
         // For now, disable spanning tree optimization when fork or exec* is
         // used.
         if (HasExecOrFork)
@@ -836,16 +836,16 @@ bool GCOVProfiler::emitProfileNotes(
 
       // Some non-tree edges are IndirectBr which cannot be split. Ignore them
       // as well.
-      llvm::erase_if(MST.AllEdges, [](std::unique_ptr<Edge> &E) {
+      llvm::erase_if(MST.allEdges(), [](std::unique_ptr<Edge> &E) {
         return E->Removed || (!E->InMST && !E->Place);
       });
       const size_t Measured =
           std::stable_partition(
-              MST.AllEdges.begin(), MST.AllEdges.end(),
+              MST.allEdges().begin(), MST.allEdges().end(),
               [](std::unique_ptr<Edge> &E) { return E->Place; }) -
-          MST.AllEdges.begin();
+          MST.allEdges().begin();
       for (size_t I : llvm::seq<size_t>(0, Measured)) {
-        Edge &E = *MST.AllEdges[I];
+        Edge &E = *MST.allEdges()[I];
         GCOVBlock &Src =
             E.SrcBB ? Func.getBlock(E.SrcBB) : Func.getEntryBlock();
         GCOVBlock &Dst =
@@ -854,13 +854,13 @@ bool GCOVProfiler::emitProfileNotes(
         E.DstNumber = Dst.Number;
       }
       std::stable_sort(
-          MST.AllEdges.begin(), MST.AllEdges.begin() + Measured,
+          MST.allEdges().begin(), MST.allEdges().begin() + Measured,
           [](const std::unique_ptr<Edge> &L, const std::unique_ptr<Edge> &R) {
             return L->SrcNumber != R->SrcNumber ? L->SrcNumber < R->SrcNumber
                                                 : L->DstNumber < R->DstNumber;
           });
 
-      for (const Edge &E : make_pointee_range(MST.AllEdges)) {
+      for (const Edge &E : make_pointee_range(MST.allEdges())) {
         GCOVBlock &Src =
             E.SrcBB ? Func.getBlock(E.SrcBB) : Func.getEntryBlock();
         GCOVBlock &Dst =
@@ -917,7 +917,7 @@ bool GCOVProfiler::emitProfileNotes(
         CountersBySP.emplace_back(Counters, SP);
 
         for (size_t I : llvm::seq<size_t>(0, Measured)) {
-          const Edge &E = *MST.AllEdges[I];
+          const Edge &E = *MST.allEdges()[I];
           IRBuilder<> Builder(E.Place, E.Place->getFirstInsertionPt());
           Value *V = Builder.CreateConstInBoundsGEP2_64(
               Counters->getValueType(), Counters, 0, I);

llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp

@@ -582,12 +582,12 @@ template <class Edge, class BBInfo> class FuncPGOInstrumentation {
     if (!IsCS) {
       NumOfPGOSelectInsts += SIVisitor.getNumOfSelectInsts();
       NumOfPGOMemIntrinsics += ValueSites[IPVK_MemOPSize].size();
-      NumOfPGOBB += MST.BBInfos.size();
+      NumOfPGOBB += MST.bbInfoSize();
       ValueSites[IPVK_IndirectCallTarget] = VPC.get(IPVK_IndirectCallTarget);
     } else {
       NumOfCSPGOSelectInsts += SIVisitor.getNumOfSelectInsts();
       NumOfCSPGOMemIntrinsics += ValueSites[IPVK_MemOPSize].size();
-      NumOfCSPGOBB += MST.BBInfos.size();
+      NumOfCSPGOBB += MST.bbInfoSize();
     }
 
     FuncName = getIRPGOFuncName(F);
@@ -597,7 +597,7 @@ template <class Edge, class BBInfo> class FuncPGOInstrumentation {
       renameComdatFunction();
     LLVM_DEBUG(dumpInfo("after CFGMST"));
 
-    for (auto &E : MST.AllEdges) {
+    for (const auto &E : MST.allEdges()) {
       if (E->Removed)
         continue;
       IsCS ? NumOfCSPGOEdge++ : NumOfPGOEdge++;
@@ -640,7 +640,7 @@ void FuncPGOInstrumentation<Edge, BBInfo>::computeCFGHash() {
     FunctionHash = (uint64_t)SIVisitor.getNumOfSelectInsts() << 56 |
                    (uint64_t)ValueSites[IPVK_IndirectCallTarget].size() << 48 |
                    //(uint64_t)ValueSites[IPVK_MemOPSize].size() << 40 |
-                   (uint64_t)MST.AllEdges.size() << 32 | JC.getCRC();
+                   (uint64_t)MST.numEdges() << 32 | JC.getCRC();
   } else {
     // The higher 32 bits.
     auto updateJCH = [&JCH](uint64_t Num) {
@@ -654,7 +654,7 @@ void FuncPGOInstrumentation<Edge, BBInfo>::computeCFGHash() {
     if (BCI) {
       updateJCH(BCI->getInstrumentedBlocksHash());
     } else {
-      updateJCH((uint64_t)MST.AllEdges.size());
+      updateJCH((uint64_t)MST.numEdges());
     }
 
     // Hash format for context sensitive profile. Reserve 4 bits for other
@@ -669,7 +669,7 @@ void FuncPGOInstrumentation<Edge, BBInfo>::computeCFGHash() {
   LLVM_DEBUG(dbgs() << "Function Hash Computation for " << F.getName() << ":\n"
                     << " CRC = " << JC.getCRC()
                     << ", Selects = " << SIVisitor.getNumOfSelectInsts()
-                    << ", Edges = " << MST.AllEdges.size() << ", ICSites = "
+                    << ", Edges = " << MST.numEdges() << ", ICSites = "
                     << ValueSites[IPVK_IndirectCallTarget].size());
   if (!PGOOldCFGHashing) {
     LLVM_DEBUG(dbgs() << ", Memops = " << ValueSites[IPVK_MemOPSize].size()
@@ -757,8 +757,8 @@ void FuncPGOInstrumentation<Edge, BBInfo>::getInstrumentBBs(
 
   // Use a worklist as we will update the vector during the iteration.
   std::vector<Edge *> EdgeList;
-  EdgeList.reserve(MST.AllEdges.size());
-  for (auto &E : MST.AllEdges)
+  EdgeList.reserve(MST.numEdges());
+  for (const auto &E : MST.allEdges())
     EdgeList.push_back(E.get());
 
   for (auto &E : EdgeList) {
@@ -1163,12 +1163,12 @@ class PGOUseFunc {
 } // end anonymous namespace
 
 /// Set up InEdges/OutEdges for all BBs in the MST.
-static void
-setupBBInfoEdges(FuncPGOInstrumentation<PGOUseEdge, PGOUseBBInfo> &FuncInfo) {
+static void setupBBInfoEdges(
+    const FuncPGOInstrumentation<PGOUseEdge, PGOUseBBInfo> &FuncInfo) {
   // This is not required when there is block coverage inference.
   if (FuncInfo.BCI)
     return;
-  for (auto &E : FuncInfo.MST.AllEdges) {
+  for (const auto &E : FuncInfo.MST.allEdges()) {
     if (E->Removed)
       continue;
     const BasicBlock *SrcBB = E->SrcBB;
@@ -1224,7 +1224,7 @@ bool PGOUseFunc::setInstrumentedCounts(
   // Set the profile count the Instrumented edges. There are BBs that not in
   // MST but not instrumented. Need to set the edge count value so that we can
   // populate the profile counts later.
-  for (auto &E : FuncInfo.MST.AllEdges) {
+  for (const auto &E : FuncInfo.MST.allEdges()) {
     if (E->Removed || E->InMST)
       continue;
     const BasicBlock *SrcBB = E->SrcBB;