Resubmit with fix: [NFC] Refactor MBB hotness/coldness into templated PSI functions.

In D152399, we calculate BPI->BFI in MachineFunctionSplit pass just to
use PSI->isFunctionHotInCallGraph, which is expensive. Instead, we can
implement this directly with MBFI.

Reviewer mentioned in the comment, that machine_size_opts already has
isFunctionColdInCallGraph, isFunctionHotInCallGraphNthPercentile, etc
implemented. These can be refactored and reused across MFS and machine
size opts.

This CL does this - it refactors out those internal static functions
into PSI as templated functions, so they can be accessed easily.

Differential Revision: https://reviews.llvm.org/D153927

Author: shenhanc78

llvm/include/llvm/Analysis/ProfileSummaryInfo.h

@@ -15,6 +15,10 @@
 #define LLVM_ANALYSIS_PROFILESUMMARYINFO_H
 
 #include "llvm/ADT/DenseMap.h"
+#include "llvm/Analysis/BlockFrequencyInfo.h"
+#include "llvm/CodeGen/MachineFunction.h"
+#include "llvm/IR/Function.h"
+#include "llvm/IR/Instructions.h"
 #include "llvm/IR/PassManager.h"
 #include "llvm/IR/ProfileSummary.h"
 #include "llvm/Pass.h"
@@ -23,9 +27,7 @@
 
 namespace llvm {
 class BasicBlock;
-class BlockFrequencyInfo;
 class CallBase;
-class Function;
 
 /// Analysis providing profile information.
 ///
@@ -107,28 +109,77 @@ class ProfileSummaryInfo {
   bool hasHugeWorkingSetSize() const;
   /// Returns true if the working set size of the code is considered large.
   bool hasLargeWorkingSetSize() const;
-  /// Returns true if \p F has hot function entry.
-  bool isFunctionEntryHot(const Function *F) const;
+  /// Returns true if \p F has hot function entry. If it returns false, it
+  /// either means it is not hot or it is unknown whether it is hot or not (for
+  /// example, no profile data is available).
+  template <typename FuncT> bool isFunctionEntryHot(const FuncT *F) const {
+    if (!F || !hasProfileSummary())
+      return false;
+    std::optional<Function::ProfileCount> FunctionCount = getEntryCount(F);
+    // FIXME: The heuristic used below for determining hotness is based on
+    // preliminary SPEC tuning for inliner. This will eventually be a
+    // convenience method that calls isHotCount.
+    return FunctionCount && isHotCount(FunctionCount->getCount());
+  }
+
   /// Returns true if \p F contains hot code.
-  bool isFunctionHotInCallGraph(const Function *F,
-                                BlockFrequencyInfo &BFI) const;
+  template <typename FuncT, typename BFIT>
+  bool isFunctionHotInCallGraph(const FuncT *F, BFIT &BFI) const {
+    if (!F || !hasProfileSummary())
+      return false;
+    if (auto FunctionCount = getEntryCount(F))
+      if (isHotCount(FunctionCount->getCount()))
+        return true;
+
+    if (auto TotalCallCount = getTotalCallCount(F)) {
+      if (isHotCount(*TotalCallCount))
+        return true;
+    }
+
+    for (const auto &BB : *F)
+      if (isHotBlock(&BB, &BFI))
+        return true;
+    return false;
+  }
   /// Returns true if \p F has cold function entry.
   bool isFunctionEntryCold(const Function *F) const;
   /// Returns true if \p F contains only cold code.
-  bool isFunctionColdInCallGraph(const Function *F,
-                                 BlockFrequencyInfo &BFI) const;
+  template <typename FuncT, typename BFIT>
+  bool isFunctionColdInCallGraph(const FuncT *F, BFIT &BFI) const {
+    if (!F || !hasProfileSummary())
+      return false;
+    if (auto FunctionCount = getEntryCount(F))
+      if (!isColdCount(FunctionCount->getCount()))
+        return false;
+
+    if (auto TotalCallCount = getTotalCallCount(F)) {
+      if (!isColdCount(*TotalCallCount))
+        return false;
+    }
+
+    for (const auto &BB : *F)
+      if (!isColdBlock(&BB, &BFI))
+        return false;
+    return true;
+  }
   /// Returns true if the hotness of \p F is unknown.
   bool isFunctionHotnessUnknown(const Function &F) const;
   /// Returns true if \p F contains hot code with regard to a given hot
   /// percentile cutoff value.
+  template <typename FuncT, typename BFIT>
   bool isFunctionHotInCallGraphNthPercentile(int PercentileCutoff,
-                                             const Function *F,
-                                             BlockFrequencyInfo &BFI) const;
+                                             const FuncT *F, BFIT &BFI) const {
+    return isFunctionHotOrColdInCallGraphNthPercentile<true, FuncT, BFIT>(
+        PercentileCutoff, F, BFI);
+  }
   /// Returns true if \p F contains cold code with regard to a given cold
   /// percentile cutoff value.
+  template <typename FuncT, typename BFIT>
   bool isFunctionColdInCallGraphNthPercentile(int PercentileCutoff,
-                                              const Function *F,
-                                              BlockFrequencyInfo &BFI) const;
+                                              const FuncT *F, BFIT &BFI) const {
+    return isFunctionHotOrColdInCallGraphNthPercentile<false, FuncT, BFIT>(
+        PercentileCutoff, F, BFI);
+  }
   /// Returns true if count \p C is considered hot.
   bool isHotCount(uint64_t C) const;
   /// Returns true if count \p C is considered cold.
@@ -143,22 +194,57 @@ class ProfileSummaryInfo {
   /// PercentileCutoff is encoded as a 6 digit decimal fixed point number, where
   /// the first two digits are the whole part. E.g. 995000 for 99.5 percentile.
   bool isColdCountNthPercentile(int PercentileCutoff, uint64_t C) const;
+
   /// Returns true if BasicBlock \p BB is considered hot.
-  bool isHotBlock(const BasicBlock *BB, BlockFrequencyInfo *BFI) const;
+  template <typename BBType, typename BFIT>
+  bool isHotBlock(const BBType *BB, BFIT *BFI) const {
+    auto Count = BFI->getBlockProfileCount(BB);
+    return Count && isHotCount(*Count);
+  }
+
   /// Returns true if BasicBlock \p BB is considered cold.
-  bool isColdBlock(const BasicBlock *BB, BlockFrequencyInfo *BFI) const;
-  /// Returns true if BasicBlock \p BB is considered hot with regard to a given
-  /// hot percentile cutoff value.
-  /// PercentileCutoff is encoded as a 6 digit decimal fixed point number, where
-  /// the first two digits are the whole part. E.g. 995000 for 99.5 percentile.
-  bool isHotBlockNthPercentile(int PercentileCutoff, const BasicBlock *BB,
-                               BlockFrequencyInfo *BFI) const;
+  template <typename BBType, typename BFIT>
+  bool isColdBlock(const BBType *BB, BFIT *BFI) const {
+    auto Count = BFI->getBlockProfileCount(BB);
+    return Count && isColdCount(*Count);
+  }
+
+  template <typename BFIT>
+  bool isColdBlock(BlockFrequency BlockFreq, const BFIT *BFI) const {
+    auto Count = BFI->getProfileCountFromFreq(BlockFreq.getFrequency());
+    return Count && isColdCount(*Count);
+  }
+
+  template <typename BBType, typename BFIT>
+  bool isHotBlockNthPercentile(int PercentileCutoff, const BBType *BB,
+                               BFIT *BFI) const {
+    return isHotOrColdBlockNthPercentile<true, BBType, BFIT>(PercentileCutoff,
+                                                             BB, BFI);
+  }
+
+  template <typename BFIT>
+  bool isHotBlockNthPercentile(int PercentileCutoff, BlockFrequency BlockFreq,
+                               BFIT *BFI) const {
+    return isHotOrColdBlockNthPercentile<true, BFIT>(PercentileCutoff,
+                                                     BlockFreq, BFI);
+  }
+
   /// Returns true if BasicBlock \p BB is considered cold with regard to a given
   /// cold percentile cutoff value.
   /// PercentileCutoff is encoded as a 6 digit decimal fixed point number, where
   /// the first two digits are the whole part. E.g. 995000 for 99.5 percentile.
-  bool isColdBlockNthPercentile(int PercentileCutoff, const BasicBlock *BB,
-                                BlockFrequencyInfo *BFI) const;
+  template <typename BBType, typename BFIT>
+  bool isColdBlockNthPercentile(int PercentileCutoff, const BBType *BB,
+                                BFIT *BFI) const {
+    return isHotOrColdBlockNthPercentile<false, BBType, BFIT>(PercentileCutoff,
+                                                              BB, BFI);
+  }
+  template <typename BFIT>
+  bool isColdBlockNthPercentile(int PercentileCutoff, BlockFrequency BlockFreq,
+                                BFIT *BFI) const {
+    return isHotOrColdBlockNthPercentile<false, BFIT>(PercentileCutoff,
+                                                      BlockFreq, BFI);
+  }
   /// Returns true if the call site \p CB is considered hot.
   bool isHotCallSite(const CallBase &CB, BlockFrequencyInfo *BFI) const;
   /// Returns true if call site \p CB is considered cold.
@@ -178,18 +264,97 @@ class ProfileSummaryInfo {
     return ColdCountThreshold.value_or(0);
   }
 
- private:
-   template <bool isHot>
-   bool isFunctionHotOrColdInCallGraphNthPercentile(
-       int PercentileCutoff, const Function *F, BlockFrequencyInfo &BFI) const;
-   template <bool isHot>
-   bool isHotOrColdCountNthPercentile(int PercentileCutoff, uint64_t C) const;
-   template <bool isHot>
-   bool isHotOrColdBlockNthPercentile(int PercentileCutoff,
-                                      const BasicBlock *BB,
-                                      BlockFrequencyInfo *BFI) const;
+private:
+  template <typename FuncT>
+  std::optional<uint64_t> getTotalCallCount(const FuncT *F) const {
+    return std::nullopt;
+  }
+
+  template <bool isHot, typename FuncT, typename BFIT>
+  bool isFunctionHotOrColdInCallGraphNthPercentile(int PercentileCutoff,
+                                                   const FuncT *F,
+                                                   BFIT &FI) const {
+    if (!F || !hasProfileSummary())
+      return false;
+    if (auto FunctionCount = getEntryCount(F)) {
+      if (isHot &&
+          isHotCountNthPercentile(PercentileCutoff, FunctionCount->getCount()))
+        return true;
+      if (!isHot && !isColdCountNthPercentile(PercentileCutoff,
+                                              FunctionCount->getCount()))
+        return false;
+    }
+    if (auto TotalCallCount = getTotalCallCount(F)) {
+      if (isHot && isHotCountNthPercentile(PercentileCutoff, *TotalCallCount))
+        return true;
+      if (!isHot &&
+          !isColdCountNthPercentile(PercentileCutoff, *TotalCallCount))
+        return false;
+    }
+    for (const auto &BB : *F) {
+      if (isHot && isHotBlockNthPercentile(PercentileCutoff, &BB, &FI))
+        return true;
+      if (!isHot && !isColdBlockNthPercentile(PercentileCutoff, &BB, &FI))
+        return false;
+    }
+    return !isHot;
+  }
+
+  template <bool isHot>
+  bool isHotOrColdCountNthPercentile(int PercentileCutoff, uint64_t C) const;
+
+  template <bool isHot, typename BBType, typename BFIT>
+  bool isHotOrColdBlockNthPercentile(int PercentileCutoff, const BBType *BB,
+                                     BFIT *BFI) const {
+    auto Count = BFI->getBlockProfileCount(BB);
+    if (isHot)
+      return Count && isHotCountNthPercentile(PercentileCutoff, *Count);
+    else
+      return Count && isColdCountNthPercentile(PercentileCutoff, *Count);
+  }
+
+  template <bool isHot, typename BFIT>
+  bool isHotOrColdBlockNthPercentile(int PercentileCutoff,
+                                     BlockFrequency BlockFreq,
+                                     BFIT *BFI) const {
+    auto Count = BFI->getProfileCountFromFreq(BlockFreq.getFrequency());
+    if (isHot)
+      return Count && isHotCountNthPercentile(PercentileCutoff, *Count);
+    else
+      return Count && isColdCountNthPercentile(PercentileCutoff, *Count);
+  }
+
+  template <typename FuncT>
+  std::optional<Function::ProfileCount> getEntryCount(const FuncT *F) const {
+    return F->getEntryCount();
+  }
 };
 
+template <>
+inline std::optional<uint64_t>
+ProfileSummaryInfo::getTotalCallCount<Function>(const Function *F) const {
+  if (!hasSampleProfile())
+    return std::nullopt;
+  uint64_t TotalCallCount = 0;
+  for (const auto &BB : *F) {
+    for (const auto &I : BB) {
+      if (isa<CallInst>(I) || isa<InvokeInst>(I)) {
+        if (auto CallCount = getProfileCount(cast<CallBase>(I), nullptr)) {
+          TotalCallCount += *CallCount;
+        }
+      }
+    }
+  }
+  return TotalCallCount;
+}
+
+template <>
+inline std::optional<Function::ProfileCount>
+ProfileSummaryInfo::getEntryCount<MachineFunction>(
+    const MachineFunction *F) const {
+  return F->getFunction().getEntryCount();
+}
+
 /// An analysis pass based on legacy pass manager to deliver ProfileSummaryInfo.
 class ProfileSummaryInfoWrapperPass : public ImmutablePass {
   std::unique_ptr<ProfileSummaryInfo> PSI;

llvm/include/llvm/Transforms/Utils/SizeOpts.h

@@ -47,7 +47,7 @@ static inline bool isPGSOColdCodeOnly(ProfileSummaryInfo *PSI) {
          (PGSOLargeWorkingSetSizeOnly && !PSI->hasLargeWorkingSetSize());
 }
 
-template<typename AdapterT, typename FuncT, typename BFIT>
+template <typename FuncT, typename BFIT>
 bool shouldFuncOptimizeForSizeImpl(const FuncT *F, ProfileSummaryInfo *PSI,
                                    BFIT *BFI, PGSOQueryType QueryType) {
   assert(F);
@@ -58,34 +58,34 @@ bool shouldFuncOptimizeForSizeImpl(const FuncT *F, ProfileSummaryInfo *PSI,
   if (!EnablePGSO)
     return false;
   if (isPGSOColdCodeOnly(PSI))
-    return AdapterT::isFunctionColdInCallGraph(F, PSI, *BFI);
+    return PSI->isFunctionColdInCallGraph(F, *BFI);
   if (PSI->hasSampleProfile())
     // The "isCold" check seems to work better for Sample PGO as it could have
     // many profile-unannotated functions.
-    return AdapterT::isFunctionColdInCallGraphNthPercentile(
-        PgsoCutoffSampleProf, F, PSI, *BFI);
-  return !AdapterT::isFunctionHotInCallGraphNthPercentile(PgsoCutoffInstrProf,
-                                                          F, PSI, *BFI);
+    return PSI->isFunctionColdInCallGraphNthPercentile(PgsoCutoffSampleProf, F,
+                                                       *BFI);
+  return !PSI->isFunctionHotInCallGraphNthPercentile(PgsoCutoffInstrProf, F,
+                                                     *BFI);
 }
 
-template<typename AdapterT, typename BlockTOrBlockFreq, typename BFIT>
-bool shouldOptimizeForSizeImpl(BlockTOrBlockFreq BBOrBlockFreq, ProfileSummaryInfo *PSI,
-                               BFIT *BFI, PGSOQueryType QueryType) {
+template <typename BlockTOrBlockFreq, typename BFIT>
+bool shouldOptimizeForSizeImpl(BlockTOrBlockFreq BBOrBlockFreq,
+                               ProfileSummaryInfo *PSI, BFIT *BFI,
+                               PGSOQueryType QueryType) {
   if (!PSI || !BFI || !PSI->hasProfileSummary())
     return false;
   if (ForcePGSO)
     return true;
   if (!EnablePGSO)
     return false;
   if (isPGSOColdCodeOnly(PSI))
-    return AdapterT::isColdBlock(BBOrBlockFreq, PSI, BFI);
+    return PSI->isColdBlock(BBOrBlockFreq, BFI);
   if (PSI->hasSampleProfile())
     // The "isCold" check seems to work better for Sample PGO as it could have
     // many profile-unannotated functions.
-    return AdapterT::isColdBlockNthPercentile(PgsoCutoffSampleProf,
-                                              BBOrBlockFreq, PSI, BFI);
-  return !AdapterT::isHotBlockNthPercentile(PgsoCutoffInstrProf, BBOrBlockFreq,
-                                            PSI, BFI);
+    return PSI->isColdBlockNthPercentile(PgsoCutoffSampleProf, BBOrBlockFreq,
+                                         BFI);
+  return !PSI->isHotBlockNthPercentile(PgsoCutoffInstrProf, BBOrBlockFreq, BFI);
 }
 
 /// Returns true if function \p F is suggested to be size-optimized based on the