[NFC] Refactor MBB hotness/coldness into templated PSI functions

[shenhan]

Summary

[NFC] 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 @wenlei mentioned in his commet, 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.

[wenlei]

Thanks for the follow up. It doesn't feel like this belongs to MBFI either. BFI for IR doesn't contain any of the hotness query - instead these are all part of PSI.

I'm wondering if we can make some template functions in PSI that cover both IR and MIR hotness queries, since BFI and MBFI have very similar APIs. We have precedence for PGO code sharing between IR and MIR (SampleProfileLoaderBaseImpl<BT>).

[shenhan]

In D152758#4416070, @wenlei wrote:

Thanks for the follow up. It doesn't feel like this belongs to MBFI either. BFI for IR doesn't contain any of the hotness query - instead these are all part of PSI.

I'm wondering if we can make some template functions in PSI that cover both IR and MIR hotness queries, since BFI and MBFI have very similar APIs. We have precedence for PGO code sharing between IR and MIR (SampleProfileLoaderBaseImpl<BT>).

Thanks. That makes sense. I've reverted all changes to MBFI. And moved the code to PSI with templates to accommodate Function/MachineFunction, BasicBlock/MachineBasicBlock, BFI/MBFI. Also moved some function implementations from the cpp file to the header, as they are templated. (I have not chosen to template the whole PSI class, because currently I see limited usage for such changes. (And the scope of template-PSI would require a serials of CLs.)).

[wenlei]

+@davidxl

llvm/include/llvm/Analysis/ProfileSummaryInfo.h

136 ↗

(On Diff #531140)

This is better down through template specialization. Check how SampleProfileInference<BT>::findUnlikelyJumps handles Invoke in IR.

You can extract only this part into a small helper, and specialize that helper based on template type (i.e. check call/invoke for IR, no-op for MIR)

170 ↗

(On Diff #531140)

same here

311 ↗

(On Diff #531140)

same here

362 ↗

(On Diff #531140)

same here

llvm/lib/Analysis/ProfileSummaryInfo.cpp

98–100 ↗

(On Diff #531140)

nit: this comment seem lost?

[davidxl]

llvm/include/llvm/Analysis/ProfileSummaryInfo.h

136 ↗

(On Diff #531140)

Agree. If we plan to have this for MBFI in the future, we can just extract this part out as a helper function too.

171 ↗

(On Diff #531140)

Can this part share code with the Hot checking counter part?

[shenhan]

llvm/include/llvm/Analysis/ProfileSummaryInfo.h

136 ↗

(On Diff #531140)

Done.

171 ↗

(On Diff #531140)

Done by extracting part of this into templated function "findTotalCallCount", leaving hot/cold checking part unchanged. This keeps code simple and also easy to understand.

[shenhan]

Build failure is caused by one time out case:
Timed Out Tests (1):

MLIR :: Examples/standalone/test.toy

Unlikely to be caused by this change. Rerunning locally and will retry the bot later.

[wenlei]

llvm/include/llvm/Analysis/ProfileSummaryInfo.h

277 ↗

(On Diff #531175)

Call and Invoke are specific to IR version, so checking that in generic implementation is a bit weird. I'd suggest let generic implementation return std::nullopt, and have IR version check on Call and Invoke.

277 ↗

(On Diff #531175)

also nit: findTotalCallCount->getTotalCallCount to be consistent with getProfileCount.

[shenhan]

llvm/include/llvm/Analysis/ProfileSummaryInfo.h

277 ↗

(On Diff #531175)

Good point. Done.

[wenlei]

llvm/lib/CodeGen/MachineSizeOpts.cpp

28

Now with PSI directly supporting IR and MIR, do we still need MachineBasicBlockBFIAdapter?

This function below (main use of adapter) can now interface with templated PSI APIs (isHotBlockNthPercentile/isColdBlockNthPercentile/isColdBlock) directly without going through adapter.

template<typename AdapterT, typename BlockTOrBlockFreq, typename BFIT>
bool shouldOptimizeForSizeImpl(BlockTOrBlockFreq BBOrBlockFreq, ProfileSummaryInfo *PSI,
                               BFIT *BFI, PGSOQueryType QueryType)

[shenhan]

llvm/lib/CodeGen/MachineSizeOpts.cpp

28

Thanks. Good hint.

Removed "AdapterT" template-parameter from "shouldOptimizeForSizeImpl" and "shouldFuncOptimizeForSizeImpl".

This also leads to the removal of "MachineBasicBlockBFIAdapter" class. Now MachineSizeOpts.cpp looks neat.

[shenhan]

llvm/include/llvm/Analysis/ProfileSummaryInfo.h

128 ↗

(On Diff #533403)

Also renamed all occurrences of template type name "FreqInfoType" -> "BFIT", so this is in accordance with existing code.

[shenhan]

Previous build bot failure was caused by 31af18bccea95fe1ae8aa2c51cf7c8e92a1c208e, which was later reverted by c9a8a0e8a9b267fae0352af352e9c735fef019f7. Synced this patch past problematic CLs.

[wenlei]

Thanks for the refactoring and cleanup. LGTM.

llvm/include/llvm/Analysis/ProfileSummaryInfo.h

127 ↗

(On Diff #534376)

nit: FunctionType->FuncT to be consistent with BFIT and other template definitions.

[shenhan]

Thanks @wenlei for the review.

llvm/include/llvm/Analysis/ProfileSummaryInfo.h

127 ↗

(On Diff #534376)

Done for all occurrences.

[aidengrossman]

This seems to have broken quite a few builders. Do you mind reverting if you don't have a quick fix available?

llvm/include/llvm/CodeGen/MachineBlockFrequencyInfo.h

 
 namespace llvm {
 
 template <class BlockT> class BlockFrequencyInfoImpl;
 class MachineBasicBlock;
 class MachineBranchProbabilityInfo;
 class MachineFunction;
 class MachineLoopInfo;
+class ProfileSummaryInfo;
 class raw_ostream;
 
 /// MachineBlockFrequencyInfo pass uses BlockFrequencyInfoImpl implementation
 /// to estimate machine basic block frequencies.
 class MachineBlockFrequencyInfo : public MachineFunctionPass {
   using ImplType = BlockFrequencyInfoImpl<MachineBasicBlock>;
   std::unique_ptr<ImplType> MBFI;
 
public:
   // Convenience method that attempts to look up the frequency associated with
   // BB and print it to OS.
   raw_ostream &printBlockFreq(raw_ostream &OS,
                               const MachineBasicBlock *MBB) const;
 
   /// Divide a block's BlockFrequency::getFrequency() value by this value to
   /// obtain the entry block - relative frequency of said block.
   uint64_t getEntryFreq() const;
-};
 
+  /// Below utility functions are like ProfileSummaryInfo, except for
+  /// MachineBasicBlock.
+  bool isFunctionHotInCallGraphNthPercentile(int PercentileCutoff,
+                                             const MachineFunction *MF,
+                                             ProfileSummaryInfo *PSI) const;
+
+  bool isFunctionColdInCallGraphNthPercentile(int PercentileCutoff,
+                                              const MachineFunction *MF,
+                                              ProfileSummaryInfo *PSI) const;
+
+  bool isFunctionHotInCallGraph(const MachineFunction *MF,
+                                ProfileSummaryInfo *PSI) const;
+
+  bool isFunctionColdInCallGraph(const MachineFunction *MF,
+                                 ProfileSummaryInfo *PSI) const;
+
+  bool isHotBlock(const MachineBasicBlock *MBB, ProfileSummaryInfo *PSI) const;
+
+  bool isHotBlock(BlockFrequency BlockFreq, ProfileSummaryInfo *PSI) const;
+
+  bool isColdBlock(const MachineBasicBlock *MBB, ProfileSummaryInfo *PSI) const;
+
+  bool isColdBlock(BlockFrequency BlockFreq, ProfileSummaryInfo *PSI) const;
+
+  bool isHotBlockNthPercentile(int PercentileCutoff,
+                               const MachineBasicBlock *MBB,
+                               ProfileSummaryInfo *PSI) const;
+
+  bool isHotBlockNthPercentile(int PercentileCutoff, BlockFrequency BlockFreq,
+                               ProfileSummaryInfo *PSI) const;
+
+  bool isColdBlockNthPercentile(int PercentileCutoff,
+                                const MachineBasicBlock *MBB,
+                                ProfileSummaryInfo *PSI) const;
+
+  bool isColdBlockNthPercentile(int PercentileCutoff, BlockFrequency BlockFreq,
+                                ProfileSummaryInfo *PSI) const;
+};
 } // end namespace llvm
 
 #endif // LLVM_CODEGEN_MACHINEBLOCKFREQUENCYINFO_H

llvm/lib/CodeGen/MachineBlockFrequencyInfo.cpp

 //===- MachineBlockFrequencyInfo.cpp - MBB Frequency Analysis -------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // Loops should be simplified before this analysis.
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/CodeGen/MachineBlockFrequencyInfo.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/iterator.h"
 #include "llvm/Analysis/BlockFrequencyInfoImpl.h"
+#include "llvm/Analysis/ProfileSummaryInfo.h"
 #include "llvm/CodeGen/MachineBasicBlock.h"
 #include "llvm/CodeGen/MachineBranchProbabilityInfo.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineLoopInfo.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/GraphWriter.h"

 MachineBlockFrequencyInfo::printBlockFreq(raw_ostream &OS,
                                           const MachineBasicBlock *MBB) const {
   return MBFI ? MBFI->printBlockFreq(OS, MBB) : OS;
 }
 
 uint64_t MachineBlockFrequencyInfo::getEntryFreq() const {
   return MBFI ? MBFI->getEntryFreq() : 0;
 }
+
+bool MachineBlockFrequencyInfo::isHotBlock(const MachineBasicBlock *MBB,
+                                           ProfileSummaryInfo *PSI) const {
+  auto Count = getBlockProfileCount(MBB);
+  return Count && PSI->isHotCount(*Count);
+}
+
+bool MachineBlockFrequencyInfo::isHotBlock(BlockFrequency BlockFreq,
+                                           ProfileSummaryInfo *PSI) const {
+  auto Count = getProfileCountFromFreq(BlockFreq.getFrequency());
+  return Count && PSI->isHotCount(*Count);
+}
+
+bool MachineBlockFrequencyInfo::isColdBlock(const MachineBasicBlock *MBB,
+                                            ProfileSummaryInfo *PSI) const {
+  auto Count = getBlockProfileCount(MBB);
+  return Count && PSI->isColdCount(*Count);
+}
+
+bool MachineBlockFrequencyInfo::isColdBlock(BlockFrequency BlockFreq,
+                                            ProfileSummaryInfo *PSI) const {
+  auto Count = getProfileCountFromFreq(BlockFreq.getFrequency());
+  return Count && PSI->isColdCount(*Count);
+}
+
+bool MachineBlockFrequencyInfo::isFunctionHotInCallGraphNthPercentile(
+    int PercentileCutoff, const MachineFunction *MF,
+    ProfileSummaryInfo *PSI) const {
+  if (auto FunctionCount = MF->getFunction().getEntryCount())
+    if (PSI->isHotCountNthPercentile(PercentileCutoff,
+                                     FunctionCount->getCount()))
+      return true;
+  for (const auto &MBB : *MF)
+    if (isHotBlockNthPercentile(PercentileCutoff, &MBB, PSI))
+      return true;
+  return false;
+}
+
+bool MachineBlockFrequencyInfo::isFunctionColdInCallGraphNthPercentile(
+    int PercentileCutoff, const MachineFunction *MF,
+    ProfileSummaryInfo *PSI) const {
+  if (auto FunctionCount = MF->getFunction().getEntryCount())
+    if (!PSI->isColdCountNthPercentile(PercentileCutoff,
+                                       FunctionCount->getCount()))
+      return false;
+  for (const auto &MBB : *MF)
+    if (!isColdBlockNthPercentile(PercentileCutoff, &MBB, PSI))
+      return false;
+  return true;
+}
+
+bool MachineBlockFrequencyInfo::isFunctionHotInCallGraph(
+    const MachineFunction *MF, ProfileSummaryInfo *PSI) const {
+  if (auto FunctionCount = MF->getFunction().getEntryCount())
+    if (PSI->isHotCount(FunctionCount->getCount()))
+      return true;
+  for (const auto &MBB : *MF)
+    if (isHotBlock(&MBB, PSI))
+      return true;
+  return false;
+}
+
+bool MachineBlockFrequencyInfo::isFunctionColdInCallGraph(
+    const MachineFunction *MF, ProfileSummaryInfo *PSI) const {
+  if (auto FunctionCount = MF->getFunction().getEntryCount())
+    if (!PSI->isColdCount(FunctionCount->getCount()))
+      return false;
+  for (const auto &MBB : *MF)
+    if (!isColdBlock(&MBB, PSI))
+      return false;
+  return true;
+}
+
+bool MachineBlockFrequencyInfo::isHotBlockNthPercentile(
+    int PercentileCutoff, const MachineBasicBlock *MBB,
+    ProfileSummaryInfo *PSI) const {
+  auto Count = getBlockProfileCount(MBB);
+  return Count && PSI->isHotCountNthPercentile(PercentileCutoff, *Count);
+}
+
+bool MachineBlockFrequencyInfo::isHotBlockNthPercentile(
+    int PercentileCutoff, BlockFrequency BlockFreq,
+    ProfileSummaryInfo *PSI) const {
+  auto Count = getProfileCountFromFreq(BlockFreq.getFrequency());
+  return Count && PSI->isHotCountNthPercentile(PercentileCutoff, *Count);
+}
+
+bool MachineBlockFrequencyInfo::isColdBlockNthPercentile(
+    int PercentileCutoff, const MachineBasicBlock *MBB,
+    ProfileSummaryInfo *PSI) const {
+  auto Count = getBlockProfileCount(MBB);
+  return Count && PSI->isColdCountNthPercentile(PercentileCutoff, *Count);
+}
+
+bool MachineBlockFrequencyInfo::isColdBlockNthPercentile(
+    int PercentileCutoff, BlockFrequency BlockFreq,
+    ProfileSummaryInfo *PSI) const {
+  auto Count = getProfileCountFromFreq(BlockFreq.getFrequency());
+  return Count && PSI->isColdCountNthPercentile(PercentileCutoff, *Count);
+}

llvm/lib/CodeGen/MachineSizeOpts.cpp

 
 extern cl::opt<bool> EnablePGSO;
 extern cl::opt<bool> PGSOLargeWorkingSetSizeOnly;
 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.
-bool isColdBlock(const MachineBasicBlock *MBB,
-                 ProfileSummaryInfo *PSI,
-                 const MachineBlockFrequencyInfo *MBFI) {
-  auto Count = MBFI->getBlockProfileCount(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,
-                                    ProfileSummaryInfo *PSI,
-                                    const MachineBlockFrequencyInfo *MBFI) {
-  auto Count = MBFI->getBlockProfileCount(MBB);
-  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(
-    const MachineFunction *MF,
-    ProfileSummaryInfo *PSI,
-    const MachineBlockFrequencyInfo &MBFI) {
-  if (auto FunctionCount = MF->getFunction().getEntryCount())
-    if (!PSI->isColdCount(FunctionCount->getCount()))
-      return false;
-  for (const auto &MBB : *MF)
-    if (!isColdBlock(&MBB, PSI, &MBFI))
-      return false;
-  return true;
-}
-
-/// Like ProfileSummaryInfo::isFunctionHotInCallGraphNthPercentile but for
-/// MachineFunction.
-bool isFunctionHotInCallGraphNthPercentile(
-    int PercentileCutoff,
-    const MachineFunction *MF,
-    ProfileSummaryInfo *PSI,
-    const MachineBlockFrequencyInfo &MBFI) {
-  if (auto FunctionCount = MF->getFunction().getEntryCount())
-    if (PSI->isHotCountNthPercentile(PercentileCutoff,
-                                     FunctionCount->getCount()))
-      return true;
-  for (const auto &MBB : *MF)
-    if (isHotBlockNthPercentile(PercentileCutoff, &MBB, PSI, &MBFI))
-      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
-
 struct MachineBasicBlockBFIAdapter {
   static bool isFunctionColdInCallGraph(const MachineFunction *MF,
                                         ProfileSummaryInfo *PSI,
                                         const MachineBlockFrequencyInfo &MBFI) {
-    return machine_size_opts_detail::isFunctionColdInCallGraph(MF, PSI, MBFI);
+    return MBFI.isFunctionColdInCallGraph(MF, PSI);
   }
   static bool isFunctionHotInCallGraphNthPercentile(
       int CutOff,
       const MachineFunction *MF,
       ProfileSummaryInfo *PSI,
       const MachineBlockFrequencyInfo &MBFI) {
-    return machine_size_opts_detail::isFunctionHotInCallGraphNthPercentile(
-        CutOff, MF, PSI, MBFI);
+    return MBFI.isFunctionHotInCallGraphNthPercentile(CutOff, MF, PSI);
   }
   static bool isFunctionColdInCallGraphNthPercentile(
       int CutOff, const MachineFunction *MF, ProfileSummaryInfo *PSI,
       const MachineBlockFrequencyInfo &MBFI) {
-    return machine_size_opts_detail::isFunctionColdInCallGraphNthPercentile(
-        CutOff, MF, PSI, MBFI);
+    return MBFI.isFunctionColdInCallGraphNthPercentile(CutOff, MF, PSI);
   }
   static bool isColdBlock(const MachineBasicBlock *MBB,
                           ProfileSummaryInfo *PSI,
                           const MachineBlockFrequencyInfo *MBFI) {
-    return machine_size_opts_detail::isColdBlock(MBB, PSI, MBFI);
+    return MBFI->isColdBlock(MBB, PSI);
   }
   static bool isColdBlock(BlockFrequency BlockFreq,
                           ProfileSummaryInfo *PSI,
                           const MachineBlockFrequencyInfo *MBFI) {
-    return machine_size_opts_detail::isColdBlock(BlockFreq, PSI, MBFI);
+    return MBFI->isColdBlock(BlockFreq, PSI);
   }
   static bool isHotBlockNthPercentile(int CutOff,
                                       const MachineBasicBlock *MBB,
                                       ProfileSummaryInfo *PSI,
                                       const MachineBlockFrequencyInfo *MBFI) {
-    return machine_size_opts_detail::isHotBlockNthPercentile(
-        CutOff, MBB, PSI, MBFI);
+    return MBFI->isHotBlockNthPercentile(CutOff, MBB, PSI);
   }
   static bool isHotBlockNthPercentile(int CutOff,
                                       BlockFrequency BlockFreq,
                                       ProfileSummaryInfo *PSI,
                                       const MachineBlockFrequencyInfo *MBFI) {
-    return machine_size_opts_detail::isHotBlockNthPercentile(
-        CutOff, BlockFreq, PSI, MBFI);
+    return MBFI->isHotBlockNthPercentile(CutOff, BlockFreq, PSI);
   }
   static bool isColdBlockNthPercentile(int CutOff, const MachineBasicBlock *MBB,
                                        ProfileSummaryInfo *PSI,
                                        const MachineBlockFrequencyInfo *MBFI) {
-    return machine_size_opts_detail::isColdBlockNthPercentile(CutOff, MBB, PSI,
-                                                              MBFI);
+    return MBFI->isColdBlockNthPercentile(CutOff, MBB, PSI);
   }
   static bool isColdBlockNthPercentile(int CutOff, BlockFrequency BlockFreq,
                                        ProfileSummaryInfo *PSI,
                                        const MachineBlockFrequencyInfo *MBFI) {
-    return machine_size_opts_detail::isColdBlockNthPercentile(CutOff, BlockFreq,
-                                                              PSI, MBFI);
+    return MBFI->isColdBlockNthPercentile(CutOff, BlockFreq, PSI);
   }
 };
 } // end anonymous namespace
 
 bool llvm::shouldOptimizeForSize(const MachineFunction *MF,
                                  ProfileSummaryInfo *PSI,
                                  const MachineBlockFrequencyInfo *MBFI,
                                  PGSOQueryType QueryType) {