1 files changed, 89 insertions, 1 deletions
diff --git a/llvm/lib/CodeGen/MachineSizeOpts.cpp b/llvm/lib/CodeGen/MachineSizeOpts.cpp
index aff67f9cfd55..584d43b42004 100644
--- a/llvm/lib/CodeGen/MachineSizeOpts.cpp
+++ b/llvm/lib/CodeGen/MachineSizeOpts.cpp
@@ -12,6 +12,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/CodeGen/MachineSizeOpts.h"
+#include "llvm/CodeGen/MBFIWrapper.h"
 #include "llvm/Analysis/ProfileSummaryInfo.h"
 #include "llvm/CodeGen/MachineBlockFrequencyInfo.h"
 
@@ -23,6 +24,7 @@ extern cl::opt<bool> ForcePGSO;
 extern cl::opt<int> PgsoCutoffInstrProf;
 extern cl::opt<int> PgsoCutoffSampleProf;
 
+namespace {
 namespace machine_size_opts_detail {
 
 /// Like ProfileSummaryInfo::isColdBlock but for MachineBasicBlock.
@@ -33,6 +35,13 @@ bool isColdBlock(const MachineBasicBlock *MBB,
   return Count && PSI->isColdCount(*Count);
 }
 
+bool isColdBlock(BlockFrequency BlockFreq,
+                 ProfileSummaryInfo *PSI,
+                 const MachineBlockFrequencyInfo *MBFI) {
+  auto Count = MBFI->getProfileCountFromFreq(BlockFreq.getFrequency());
+  return Count && PSI->isColdCount(*Count);
+}
+
 /// Like ProfileSummaryInfo::isHotBlockNthPercentile but for MachineBasicBlock.
 static bool isHotBlockNthPercentile(int PercentileCutoff,
                                     const MachineBasicBlock *MBB,
@@ -42,6 +51,30 @@ static bool isHotBlockNthPercentile(int PercentileCutoff,
   return Count && PSI->isHotCountNthPercentile(PercentileCutoff, *Count);
 }
 
+static bool isHotBlockNthPercentile(int PercentileCutoff,
+                                    BlockFrequency BlockFreq,
+                                    ProfileSummaryInfo *PSI,
+                                    const MachineBlockFrequencyInfo *MBFI) {
+  auto Count = MBFI->getProfileCountFromFreq(BlockFreq.getFrequency());
+  return Count && PSI->isHotCountNthPercentile(PercentileCutoff, *Count);
+}
+
+static bool isColdBlockNthPercentile(int PercentileCutoff,
+                                     const MachineBasicBlock *MBB,
+                                     ProfileSummaryInfo *PSI,
+                                     const MachineBlockFrequencyInfo *MBFI) {
+  auto Count = MBFI->getBlockProfileCount(MBB);
+  return Count && PSI->isColdCountNthPercentile(PercentileCutoff, *Count);
+}
+
+static bool isColdBlockNthPercentile(int PercentileCutoff,
+                                     BlockFrequency BlockFreq,
+                                     ProfileSummaryInfo *PSI,
+                                     const MachineBlockFrequencyInfo *MBFI) {
+  auto Count = MBFI->getProfileCountFromFreq(BlockFreq.getFrequency());
+  return Count && PSI->isColdCountNthPercentile(PercentileCutoff, *Count);
+}
+
 /// Like ProfileSummaryInfo::isFunctionColdInCallGraph but for
 /// MachineFunction.
 bool isFunctionColdInCallGraph(
@@ -73,9 +106,21 @@ bool isFunctionHotInCallGraphNthPercentile(
       return true;
   return false;
 }
+
+bool isFunctionColdInCallGraphNthPercentile(
+    int PercentileCutoff, const MachineFunction *MF, ProfileSummaryInfo *PSI,
+    const MachineBlockFrequencyInfo &MBFI) {
+  if (auto FunctionCount = MF->getFunction().getEntryCount())
+    if (!PSI->isColdCountNthPercentile(PercentileCutoff,
+                                       FunctionCount.getCount()))
+      return false;
+  for (const auto &MBB : *MF)
+    if (!isColdBlockNthPercentile(PercentileCutoff, &MBB, PSI, &MBFI))
+      return false;
+  return true;
+}
 } // namespace machine_size_opts_detail
 
-namespace {
 struct MachineBasicBlockBFIAdapter {
   static bool isFunctionColdInCallGraph(const MachineFunction *MF,
                                         ProfileSummaryInfo *PSI,
@@ -90,11 +135,22 @@ struct MachineBasicBlockBFIAdapter {
     return machine_size_opts_detail::isFunctionHotInCallGraphNthPercentile(
         CutOff, MF, PSI, MBFI);
   }
+  static bool isFunctionColdInCallGraphNthPercentile(
+      int CutOff, const MachineFunction *MF, ProfileSummaryInfo *PSI,
+      const MachineBlockFrequencyInfo &MBFI) {
+    return machine_size_opts_detail::isFunctionColdInCallGraphNthPercentile(
+        CutOff, MF, PSI, MBFI);
+  }
   static bool isColdBlock(const MachineBasicBlock *MBB,
                           ProfileSummaryInfo *PSI,
                           const MachineBlockFrequencyInfo *MBFI) {
     return machine_size_opts_detail::isColdBlock(MBB, PSI, MBFI);
   }
+  static bool isColdBlock(BlockFrequency BlockFreq,
+                          ProfileSummaryInfo *PSI,
+                          const MachineBlockFrequencyInfo *MBFI) {
+    return machine_size_opts_detail::isColdBlock(BlockFreq, PSI, MBFI);
+  }
   static bool isHotBlockNthPercentile(int CutOff,
                                       const MachineBasicBlock *MBB,
                                       ProfileSummaryInfo *PSI,
@@ -102,6 +158,25 @@ struct MachineBasicBlockBFIAdapter {
     return machine_size_opts_detail::isHotBlockNthPercentile(
         CutOff, MBB, PSI, MBFI);
   }
+  static bool isHotBlockNthPercentile(int CutOff,
+                                      BlockFrequency BlockFreq,
+                                      ProfileSummaryInfo *PSI,
+                                      const MachineBlockFrequencyInfo *MBFI) {
+    return machine_size_opts_detail::isHotBlockNthPercentile(
+        CutOff, BlockFreq, PSI, MBFI);
+  }
+  static bool isColdBlockNthPercentile(int CutOff, const MachineBasicBlock *MBB,
+                                       ProfileSummaryInfo *PSI,
+                                       const MachineBlockFrequencyInfo *MBFI) {
+    return machine_size_opts_detail::isColdBlockNthPercentile(CutOff, MBB, PSI,
+                                                              MBFI);
+  }
+  static bool isColdBlockNthPercentile(int CutOff, BlockFrequency BlockFreq,
+                                       ProfileSummaryInfo *PSI,
+                                       const MachineBlockFrequencyInfo *MBFI) {
+    return machine_size_opts_detail::isColdBlockNthPercentile(CutOff, BlockFreq,
+                                                              PSI, MBFI);
+  }
 };
 } // end anonymous namespace
 
@@ -117,6 +192,19 @@ bool llvm::shouldOptimizeForSize(const MachineBasicBlock *MBB,
                                  ProfileSummaryInfo *PSI,
                                  const MachineBlockFrequencyInfo *MBFI,
                                  PGSOQueryType QueryType) {
+  assert(MBB);
   return shouldOptimizeForSizeImpl<MachineBasicBlockBFIAdapter>(
       MBB, PSI, MBFI, QueryType);
 }
+
+bool llvm::shouldOptimizeForSize(const MachineBasicBlock *MBB,
+                                 ProfileSummaryInfo *PSI,
+                                 MBFIWrapper *MBFIW,
+                                 PGSOQueryType QueryType) {
+  assert(MBB);
+  if (!PSI || !MBFIW)
+    return false;
+  BlockFrequency BlockFreq = MBFIW->getBlockFreq(MBB);
+  return shouldOptimizeForSizeImpl<MachineBasicBlockBFIAdapter>(
+      BlockFreq, PSI, &MBFIW->getMBFI(), QueryType);
+}