[PGO][PGSO] Handle MBFIWrapper

[hjyamauchi]

Summary

Some code gen passes use MBFIWrapper to keep track of the frequency of new
blocks. This was not taken into account and could lead to incorrect frequencies
as MBFI silently returns zero frequency for unknown/new blocks.

Add a variant for MBFIWrapper in the PGSO query interface.

Depends on D73494.

[davidxl]

The MBFIWrapper change seems NFC, can it be extracted out first?

[hjyamauchi]

Split the movement of MBFIWrapper into D73494.

[hjyamauchi]

Rebased past D73494. PTAL.

[davidxl]

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

92

Is this function necessary? It is mostly the same as the previous impl.

[hjyamauchi]

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

92

This version takes a BlockFrequency rather than a block (BlockT). Since only MBFIWrapper, rather than BFI, has the true frequency for the block, we'd need to either (1) get the frequency out of MBFIWrapper and then pass it here, or (2) pass MBFIWrapper down here. The current approach is (1). Either way, we'd need to take the frequency out of MBFIWrapper somwhere along the way, and it won't likely dedup the code as they both require a new param type that doesn't seem to be amenable to additional templatization on top of this already templatized code.

Do you see a better idea?

[davidxl]

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

92

How about changing BlockT template parameter into BlockTOrBlockFreq, so that it can be instantiated with BlockFrequency too?

[hjyamauchi]

Addressed comment. PTAL.

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

92

That just worked out, except for the assert on the BB pointer non-nullness, which I resolved by moving it to the callers. Done.

[davidxl]

lgtm

[hjyamauchi]

This was submitted as https://github.com/llvm/llvm-project/commit/ac8da31a0f9895685e34e661d44075f46b559fbe

llvm/include/llvm/CodeGen/MBFIWrapper.h

+//===- llvm/CodeGen/MBFIWrapper.h -------------------------------*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This class keeps track of branch frequencies of newly created blocks and
+// transformed blocks. Used by the TailDuplication and MachineBlockPlacement.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_CODEGEN_MBFIWRAPPER_H
+#define LLVM_CODEGEN_MBFIWRAPPER_H
+
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/Support/BlockFrequency.h"
+
+namespace llvm {
+
+class MachineBasicBlock;
+class MachineBlockFrequencyInfo;
+
+class MBFIWrapper {
+ public:
+  MBFIWrapper(const MachineBlockFrequencyInfo &I) : MBFI(I) {}
+
+  BlockFrequency getBlockFreq(const MachineBasicBlock *MBB) const;
+  void setBlockFreq(const MachineBasicBlock *MBB, BlockFrequency F);
+  raw_ostream &printBlockFreq(raw_ostream &OS,
+                              const MachineBasicBlock *MBB) const;
+  raw_ostream &printBlockFreq(raw_ostream &OS,
+                              const BlockFrequency Freq) const;
+  void view(const Twine &Name, bool isSimple = true);
+  uint64_t getEntryFreq() const;
+  const MachineBlockFrequencyInfo &getMBFI() { return MBFI; }
+
+ private:
+  const MachineBlockFrequencyInfo &MBFI;
+  DenseMap<const MachineBasicBlock *, BlockFrequency> MergedBBFreq;
+};
+
+} // end namespace llvm
+
+#endif // LLVM_CODEGEN_MBFIWRAPPER_H

llvm/include/llvm/CodeGen/MachineSizeOpts.h

 #include "llvm/Transforms/Utils/SizeOpts.h"
 
 namespace llvm {
 
 class ProfileSummaryInfo;
 class MachineBasicBlock;
 class MachineBlockFrequencyInfo;
 class MachineFunction;
+class MBFIWrapper;
 
 /// Returns true if machine function \p MF is suggested to be size-optimized
 /// based on the profile.
 bool shouldOptimizeForSize(const MachineFunction *MF, ProfileSummaryInfo *PSI,
                            const MachineBlockFrequencyInfo *BFI,
                            PGSOQueryType QueryType = PGSOQueryType::Other);
 /// Returns true if machine basic block \p MBB is suggested to be size-optimized
 /// based on the profile.
 bool shouldOptimizeForSize(const MachineBasicBlock *MBB,
                            ProfileSummaryInfo *PSI,
                            const MachineBlockFrequencyInfo *MBFI,
                            PGSOQueryType QueryType = PGSOQueryType::Other);
+/// Returns true if machine basic block \p MBB is suggested to be size-optimized
+/// based on the profile.
+bool shouldOptimizeForSize(const MachineBasicBlock *MBB,
+                           ProfileSummaryInfo *PSI,
+                           MBFIWrapper *MBFIWrapper,
+                           PGSOQueryType QueryType = PGSOQueryType::Other);
 
 } // end namespace llvm
 
 #endif // LLVM_CODEGEN_MACHINE_SIZEOPTS_H

llvm/include/llvm/CodeGen/TailDuplicator.h

 
 #ifndef LLVM_CODEGEN_TAILDUPLICATOR_H
 #define LLVM_CODEGEN_TAILDUPLICATOR_H
 
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallVector.h"
+#include "llvm/CodeGen/MBFIWrapper.h"
 #include "llvm/CodeGen/TargetInstrInfo.h"
 #include <utility>
 #include <vector>
 
 namespace llvm {
 
 class MachineBasicBlock;
 class MachineBlockFrequencyInfo;
 class MachineBranchProbabilityInfo;
 class MachineFunction;
 class MachineInstr;
 class MachineModuleInfo;
 class MachineRegisterInfo;
 class ProfileSummaryInfo;
 class TargetRegisterInfo;
 
 /// Utility class to perform tail duplication.
 class TailDuplicator {
   const TargetInstrInfo *TII;
   const TargetRegisterInfo *TRI;
   const MachineBranchProbabilityInfo *MBPI;
   const MachineModuleInfo *MMI;
   MachineRegisterInfo *MRI;
   MachineFunction *MF;
-  const MachineBlockFrequencyInfo *MBFI;
+  MBFIWrapper *MBFI;
   ProfileSummaryInfo *PSI;
   bool PreRegAlloc;
   bool LayoutMode;
   unsigned TailDupSize;
 
   // A list of virtual registers for which to update SSA form.
   SmallVector<unsigned, 16> SSAUpdateVRs;
 
public:
   /// @param MBPI - Branch Probability Info. Used to propagate correct
   ///     probabilities when modifying the CFG.
   /// @param LayoutMode - When true, don't use the existing layout to make
   ///     decisions.
   /// @param TailDupSize - Maxmimum size of blocks to tail-duplicate. Zero
   ///     default implies using the command line value TailDupSize.
   void initMF(MachineFunction &MF, bool PreRegAlloc,
               const MachineBranchProbabilityInfo *MBPI,
-              const MachineBlockFrequencyInfo *MBFI,
+              MBFIWrapper *MBFI,
               ProfileSummaryInfo *PSI,
               bool LayoutMode, unsigned TailDupSize = 0);
 
   bool tailDuplicateBlocks();
   static bool isSimpleBB(MachineBasicBlock *TailBB);
   bool shouldTailDuplicate(bool IsSimple, MachineBasicBlock &TailBB);
 
   /// Returns true if TailBB can successfully be duplicated into PredBB

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

if (PGSOColdCodeOnly ||
     // Even if the working set size isn't large, size-optimize cold code.
     return AdapterT::isColdBlock(BB, PSI, BFI);
   }
   return !AdapterT::isHotBlockNthPercentile(
       PSI->hasSampleProfile() ? PgsoCutoffSampleProf : PgsoCutoffInstrProf,
       BB, PSI, BFI);
 }
 
+template<typename AdapterT, typename BFIT>
+bool shouldOptimizeForSizeImpl(BlockFrequency BlockFreq,
+                               ProfileSummaryInfo *PSI, BFIT *BFI,
+                               PGSOQueryType QueryType) {
+  if (!PSI || !BFI || !PSI->hasProfileSummary())
+    return false;
+  if (ForcePGSO)
+    return true;
+  if (!EnablePGSO)
+    return false;
+  // Temporarily enable size optimizations only for the IR pass or test query
+  // sites for gradual commit/rollout. This is to be removed later.
+  if (PGSOIRPassOrTestOnly && !(QueryType == PGSOQueryType::IRPass ||
+                                QueryType == PGSOQueryType::Test))
+    return false;
+  if (PGSOColdCodeOnly ||
+      (PGSOLargeWorkingSetSizeOnly && !PSI->hasLargeWorkingSetSize())) {
+    // Even if the working set size isn't large, size-optimize cold code.
+    return AdapterT::isColdBlock(BlockFreq, PSI, BFI);
+  }
+  return !AdapterT::isHotBlockNthPercentile(
+      PSI->hasSampleProfile() ? PgsoCutoffSampleProf : PgsoCutoffInstrProf,
+      BlockFreq, PSI, BFI);
+}
+
 /// Returns true if function \p F is suggested to be size-optimized based on the
 /// profile.
 bool shouldOptimizeForSize(const Function *F, ProfileSummaryInfo *PSI,
                            BlockFrequencyInfo *BFI,
                            PGSOQueryType QueryType = PGSOQueryType::Other);
 
 /// Returns true if basic block \p BB is suggested to be size-optimized based on
 /// the profile.
 bool shouldOptimizeForSize(const BasicBlock *BB, ProfileSummaryInfo *PSI,
                            BlockFrequencyInfo *BFI,
                            PGSOQueryType QueryType = PGSOQueryType::Other);
 
 } // end namespace llvm
 
 #endif // LLVM_TRANSFORMS_UTILS_SIZEOPTS_H

llvm/lib/CodeGen/BranchFolding.h

 //===- BranchFolding.h - Fold machine code branch instructions --*- C++ -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_LIB_CODEGEN_BRANCHFOLDING_H
 #define LLVM_LIB_CODEGEN_BRANCHFOLDING_H
 
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/CodeGen/LivePhysRegs.h"
 #include "llvm/CodeGen/MachineBasicBlock.h"
+#include "llvm/CodeGen/MBFIWrapper.h"
 #include "llvm/Support/BlockFrequency.h"
 #include "llvm/Support/Compiler.h"
 #include <cstdint>
 #include <vector>
 
 namespace llvm {
 
 class BasicBlock;
 class MachineBlockFrequencyInfo;
 class MachineBranchProbabilityInfo;
 class MachineFunction;
 class MachineLoopInfo;
 class MachineModuleInfo;
 class MachineRegisterInfo;
 class ProfileSummaryInfo;
 class raw_ostream;
 class TargetInstrInfo;
 class TargetRegisterInfo;
 
   class LLVM_LIBRARY_VISIBILITY BranchFolder {
   public:
-    class MBFIWrapper;
-
     explicit BranchFolder(bool defaultEnableTailMerge,
                           bool CommonHoist,
                           MBFIWrapper &FreqInfo,
                           const MachineBranchProbabilityInfo &ProbInfo,
                           ProfileSummaryInfo *PSI,
                           // Min tail length to merge. Defaults to commandline
                           // flag. Ignored for optsize.
                           unsigned MinTailLength = 0);
private:
     unsigned MinCommonTailLength;
     const TargetInstrInfo *TII;
     const MachineRegisterInfo *MRI;
     const TargetRegisterInfo *TRI;
     MachineModuleInfo *MMI;
     MachineLoopInfo *MLI;
     LivePhysRegs LiveRegs;
 
-  public:
-    /// This class keeps track of branch frequencies of newly created
-    /// blocks and tail-merged blocks.
-    class MBFIWrapper {
-    public:
-      MBFIWrapper(const MachineBlockFrequencyInfo &I) : MBFI(I) {}
-
-      BlockFrequency getBlockFreq(const MachineBasicBlock *MBB) const;
-      void setBlockFreq(const MachineBasicBlock *MBB, BlockFrequency F);
-      raw_ostream &printBlockFreq(raw_ostream &OS,
-                                  const MachineBasicBlock *MBB) const;
-      raw_ostream &printBlockFreq(raw_ostream &OS,
-                                  const BlockFrequency Freq) const;
-      void view(const Twine &Name, bool isSimple = true);
-      uint64_t getEntryFreq() const;
-      const MachineBlockFrequencyInfo &getMBFI() { return MBFI; }
-
-    private:
-      const MachineBlockFrequencyInfo &MBFI;
-      DenseMap<const MachineBasicBlock *, BlockFrequency> MergedBBFreq;
-    };
-
   private:
     MBFIWrapper &MBBFreqInfo;
     const MachineBranchProbabilityInfo &MBPI;
     ProfileSummaryInfo *PSI;
 
     bool TailMergeBlocks(MachineFunction &MF);
     bool TryTailMergeBlocks(MachineBasicBlock* SuccBB,
                        MachineBasicBlock* PredBB,

llvm/lib/CodeGen/BranchFolding.cpp

bool BranchFolderPass::runOnMachineFunction(MachineFunction &MF) {
   if (skipFunction(MF.getFunction()))
     return false;
 
   TargetPassConfig *PassConfig = &getAnalysis<TargetPassConfig>();
   // TailMerge can create jump into if branches that make CFG irreducible for
   // HW that requires structurized CFG.
   bool EnableTailMerge = !MF.getTarget().requiresStructuredCFG() &&
                          PassConfig->getEnableTailMerge();
-  BranchFolder::MBFIWrapper MBBFreqInfo(
+  MBFIWrapper MBBFreqInfo(
       getAnalysis<MachineBlockFrequencyInfo>());
   BranchFolder Folder(EnableTailMerge, /*CommonHoist=*/true, MBBFreqInfo,
                       getAnalysis<MachineBranchProbabilityInfo>(),
                       &getAnalysis<ProfileSummaryInfoWrapperPass>().getPSI());
   auto *MMIWP = getAnalysisIfAvailable<MachineModuleInfoWrapperPass>();
   return Folder.OptimizeFunction(
       MF, MF.getSubtarget().getInstrInfo(), MF.getSubtarget().getRegisterInfo(),
       MMIWP ? &MMIWP->getMMI() : nullptr);
BranchFolder::MergePotentialsElt::operator<(const MergePotentialsElt &o) const {
   // an object with itself.
 #ifndef _GLIBCXX_DEBUG
   llvm_unreachable("Predecessor appears twice");
 #else
   return false;
 #endif
 }
 
-BlockFrequency
-BranchFolder::MBFIWrapper::getBlockFreq(const MachineBasicBlock *MBB) const {
-  auto I = MergedBBFreq.find(MBB);
-
-  if (I != MergedBBFreq.end())
-    return I->second;
-
-  return MBFI.getBlockFreq(MBB);
-}
-
-void BranchFolder::MBFIWrapper::setBlockFreq(const MachineBasicBlock *MBB,
-                                             BlockFrequency F) {
-  MergedBBFreq[MBB] = F;
-}
-
-raw_ostream &
-BranchFolder::MBFIWrapper::printBlockFreq(raw_ostream &OS,
-                                          const MachineBasicBlock *MBB) const {
-  return MBFI.printBlockFreq(OS, getBlockFreq(MBB));
-}
-
-raw_ostream &
-BranchFolder::MBFIWrapper::printBlockFreq(raw_ostream &OS,
-                                          const BlockFrequency Freq) const {
-  return MBFI.printBlockFreq(OS, Freq);
-}
-
-void BranchFolder::MBFIWrapper::view(const Twine &Name, bool isSimple) {
-  MBFI.view(Name, isSimple);
-}
-
-uint64_t
-BranchFolder::MBFIWrapper::getEntryFreq() const {
-  return MBFI.getEntryFreq();
-}
-
 /// CountTerminators - Count the number of terminators in the given
 /// block and set I to the position of the first non-terminator, if there
 /// is one, or MBB->end() otherwise.
 static unsigned CountTerminators(MachineBasicBlock *MBB,
                                  MachineBasicBlock::iterator &I) {
   I = MBB->end();
   unsigned NumTerms = 0;
   while (true) {

 static bool
 ProfitableToMerge(MachineBasicBlock *MBB1, MachineBasicBlock *MBB2,
                   unsigned MinCommonTailLength, unsigned &CommonTailLen,
                   MachineBasicBlock::iterator &I1,
                   MachineBasicBlock::iterator &I2, MachineBasicBlock *SuccBB,
                   MachineBasicBlock *PredBB,
                   DenseMap<const MachineBasicBlock *, int> &EHScopeMembership,
                   bool AfterPlacement,
-                  BranchFolder::MBFIWrapper &MBBFreqInfo,
+                  MBFIWrapper &MBBFreqInfo,
                   ProfileSummaryInfo *PSI) {
   // It is never profitable to tail-merge blocks from two different EH scopes.
   if (!EHScopeMembership.empty()) {
     auto EHScope1 = EHScopeMembership.find(MBB1);
     assert(EHScope1 != EHScopeMembership.end());
     auto EHScope2 = EHScopeMembership.find(MBB2);
     assert(EHScope2 != EHScopeMembership.end());
     if (EHScope1->second != EHScope2->second)
ProfitableToMerge(MachineBasicBlock *MBB1, MachineBasicBlock *MBB2,
 
   // If we are optimizing for code size, 2 instructions in common is enough if
   // we don't have to split a block.  At worst we will be introducing 1 new
   // branch instruction, which is likely to be smaller than the 2
   // instructions that would be deleted in the merge.
   MachineFunction *MF = MBB1->getParent();
   bool OptForSize =
       MF->getFunction().hasOptSize() ||
-      (llvm::shouldOptimizeForSize(MBB1, PSI, &MBBFreqInfo.getMBFI()) &&
-       llvm::shouldOptimizeForSize(MBB2, PSI, &MBBFreqInfo.getMBFI()));
+      (llvm::shouldOptimizeForSize(MBB1, PSI, &MBBFreqInfo) &&
+       llvm::shouldOptimizeForSize(MBB2, PSI, &MBBFreqInfo));
   return EffectiveTailLen >= 2 && OptForSize &&
          (FullBlockTail1 || FullBlockTail2);
 }
 
 unsigned BranchFolder::ComputeSameTails(unsigned CurHash,
                                         unsigned MinCommonTailLength,
                                         MachineBasicBlock *SuccBB,
                                         MachineBasicBlock *PredBB) {
if (MBB->succ_empty() && !PriorCond.empty() && !PriorFBB &&
           return MadeChange;
         }
       }
     }
   }
 
   bool OptForSize =
       MF.getFunction().hasOptSize() ||
-      llvm::shouldOptimizeForSize(MBB, PSI, &MBBFreqInfo.getMBFI());
+      llvm::shouldOptimizeForSize(MBB, PSI, &MBBFreqInfo);
   if (!IsEmptyBlock(MBB) && MBB->pred_size() == 1 && OptForSize) {
     // Changing "Jcc foo; foo: jmp bar;" into "Jcc bar;" might change the branch
     // direction, thereby defeating careful block placement and regressing
     // performance. Therefore, only consider this for optsize functions.
     MachineInstr &TailCall = *MBB->getFirstNonDebugInstr();
     if (TII->isUnconditionalTailCall(TailCall)) {
       MachineBasicBlock *Pred = *MBB->pred_begin();
       MachineBasicBlock *PredTBB = nullptr, *PredFBB = nullptr;

llvm/lib/CodeGen/CMakeLists.txt

add_llvm_component_library(LLVMCodeGen
   MachineScheduler.cpp
   MachineSink.cpp
   MachineSizeOpts.cpp
   MachineSSAUpdater.cpp
   MachineTraceMetrics.cpp
   MachineVerifier.cpp
   ModuloSchedule.cpp
   PatchableFunction.cpp
+  MBFIWrapper.cpp
   MIRPrinter.cpp
   MIRPrintingPass.cpp
   MacroFusion.cpp
   NonRelocatableStringpool.cpp
   OptimizePHIs.cpp
   ParallelCG.cpp
   PeepholeOptimizer.cpp
   PHIElimination.cpp

llvm/lib/CodeGen/IfConversion.cpp

 bool IfConverter::runOnMachineFunction(MachineFunction &MF) {
   if (skipFunction(MF.getFunction()) || (PredicateFtor && !PredicateFtor(MF)))
     return false;
 
   const TargetSubtargetInfo &ST = MF.getSubtarget();
   TLI = ST.getTargetLowering();
   TII = ST.getInstrInfo();
   TRI = ST.getRegisterInfo();
-  BranchFolder::MBFIWrapper MBFI(getAnalysis<MachineBlockFrequencyInfo>());
+  MBFIWrapper MBFI(getAnalysis<MachineBlockFrequencyInfo>());
   MBPI = &getAnalysis<MachineBranchProbabilityInfo>();
   ProfileSummaryInfo *PSI =
       &getAnalysis<ProfileSummaryInfoWrapperPass>().getPSI();
   MRI = &MF.getRegInfo();
   SchedModel.init(&ST);
 
   if (!TII) return false;
 

llvm/lib/CodeGen/MBFIWrapper.cpp

+//===- MBFIWrapper.cpp - MachineBlockFrequencyInfo wrapper ----------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This class keeps track of branch frequencies of newly created blocks and
+// transformed blocks. Used by the TailDuplication and MachineBlockPlacement.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/CodeGen/MBFIWrapper.h"
+#include "llvm/CodeGen/MachineBlockFrequencyInfo.h"
+
+using namespace llvm;
+
+BlockFrequency MBFIWrapper::getBlockFreq(const MachineBasicBlock *MBB) const {
+  auto I = MergedBBFreq.find(MBB);
+
+  if (I != MergedBBFreq.end())
+    return I->second;
+
+  return MBFI.getBlockFreq(MBB);
+}
+
+void MBFIWrapper::setBlockFreq(const MachineBasicBlock *MBB,
+                               BlockFrequency F) {
+  MergedBBFreq[MBB] = F;
+}
+
+raw_ostream & MBFIWrapper::printBlockFreq(raw_ostream &OS,
+                                          const MachineBasicBlock *MBB) const {
+  return MBFI.printBlockFreq(OS, getBlockFreq(MBB));
+}
+
+raw_ostream & MBFIWrapper::printBlockFreq(raw_ostream &OS,
+                                          const BlockFrequency Freq) const {
+  return MBFI.printBlockFreq(OS, Freq);
+}
+
+void MBFIWrapper::view(const Twine &Name, bool isSimple) {
+  MBFI.view(Name, isSimple);
+}
+
+uint64_t MBFIWrapper::getEntryFreq() const {
+  return MBFI.getEntryFreq();
+}

llvm/lib/CodeGen/MachineBlockPlacement.cpp

class MachineBlockPlacement : public MachineFunctionPass {
 
   /// Machine Function
   MachineFunction *F;
 
   /// A handle to the branch probability pass.
   const MachineBranchProbabilityInfo *MBPI;
 
   /// A handle to the function-wide block frequency pass.
-  std::unique_ptr<BranchFolder::MBFIWrapper> MBFI;
+  std::unique_ptr<MBFIWrapper> MBFI;
 
   /// A handle to the loop info.
   MachineLoopInfo *MLI;
 
   /// Preferred loop exit.
   /// Member variable for convenience. It may be removed by duplication deep
   /// in the call stack.
   MachineBasicBlock *PreferredLoopExit;
MachineBlockPlacement::findBestLoopTop(const MachineLoop &L,
   // Placing the latch block before the header may introduce an extra branch
   // that skips this block the first time the loop is executed, which we want
   // to avoid when optimising for size.
   // FIXME: in theory there is a case that does not introduce a new branch,
   // i.e. when the layout predecessor does not fallthrough to the loop header.
   // In practice this never happens though: there always seems to be a preheader
   // that can fallthrough and that is also placed before the header.
   bool OptForSize = F->getFunction().hasOptSize() ||
-                    llvm::shouldOptimizeForSize(L.getHeader(), PSI,
-                                                &MBFI->getMBFI());
+                    llvm::shouldOptimizeForSize(L.getHeader(), PSI, MBFI.get());
   if (OptForSize)
     return L.getHeader();
 
   MachineBasicBlock *OldTop = nullptr;
   MachineBasicBlock *NewTop = L.getHeader();
   while (NewTop != OldTop) {
     OldTop = NewTop;
     NewTop = findBestLoopTopHelper(OldTop, L, LoopBlockSet);
for (MachineBasicBlock *ChainBB : FunctionChain) {
     // If the block is cold relative to its loop header, don't align it
     // regardless of what edges into the block exist.
     MachineBasicBlock *LoopHeader = L->getHeader();
     BlockFrequency LoopHeaderFreq = MBFI->getBlockFreq(LoopHeader);
     if (Freq < (LoopHeaderFreq * ColdProb))
       continue;
 
     // If the global profiles indicates so, don't align it.
-    if (llvm::shouldOptimizeForSize(ChainBB, PSI, &MBFI->getMBFI()) &&
+    if (llvm::shouldOptimizeForSize(ChainBB, PSI, MBFI.get()) &&
         !TLI->alignLoopsWithOptSize())
       continue;
 
     // Check for the existence of a non-layout predecessor which would benefit
     // from aligning this block.
     MachineBasicBlock *LayoutPred =
         &*std::prev(MachineFunction::iterator(ChainBB));
 
if (skipFunction(MF.getFunction()))
     return false;
 
   // Check for single-block functions and skip them.
   if (std::next(MF.begin()) == MF.end())
     return false;
 
   F = &MF;
   MBPI = &getAnalysis<MachineBranchProbabilityInfo>();
-  MBFI = std::make_unique<BranchFolder::MBFIWrapper>(
+  MBFI = std::make_unique<MBFIWrapper>(
       getAnalysis<MachineBlockFrequencyInfo>());
   MLI = &getAnalysis<MachineLoopInfo>();
   TII = MF.getSubtarget().getInstrInfo();
   TLI = MF.getSubtarget().getTargetLowering();
   MPDT = nullptr;
   PSI = &getAnalysis<ProfileSummaryInfoWrapperPass>().getPSI();
 
   // Initialize PreferredLoopExit to nullptr here since it may never be set if
bool MachineBlockPlacement::runOnMachineFunction(MachineFunction &MF) {
 
   if (allowTailDupPlacement()) {
     MPDT = &getAnalysis<MachinePostDominatorTree>();
     bool OptForSize = MF.getFunction().hasOptSize() ||
                       llvm::shouldOptimizeForSize(&MF, PSI, &MBFI->getMBFI());
     if (OptForSize)
       TailDupSize = 1;
     bool PreRegAlloc = false;
-    TailDup.initMF(MF, PreRegAlloc, MBPI, &MBFI->getMBFI(), PSI,
+    TailDup.initMF(MF, PreRegAlloc, MBPI, MBFI.get(), PSI,
                    /* LayoutMode */ true, TailDupSize);
     precomputeTriangleChains();
   }
 
   buildCFGChains();
 
   // Changing the layout can create new tail merging opportunities.
   // TailMerge can create jump into if branches that make CFG irreducible for

llvm/lib/CodeGen/MachineSizeOpts.cpp

 //===- MachineSizeOpts.cpp - code size optimization related code ----------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file contains some shared machine IR code size optimization related
 // code.
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/CodeGen/MachineSizeOpts.h"
+#include "llvm/CodeGen/MBFIWrapper.h"
 #include "llvm/Analysis/ProfileSummaryInfo.h"
 #include "llvm/CodeGen/MachineBlockFrequencyInfo.h"
 
 using namespace llvm;
 
 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 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);
+}
+
 /// 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()))
static bool isFunctionHotInCallGraphNthPercentile(
     return machine_size_opts_detail::isFunctionHotInCallGraphNthPercentile(
         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,
                                       const MachineBlockFrequencyInfo *MBFI) {
     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);
+  }
 };
 } // end anonymous namespace
 
 bool llvm::shouldOptimizeForSize(const MachineFunction *MF,
                                  ProfileSummaryInfo *PSI,
                                  const MachineBlockFrequencyInfo *MBFI,
                                  PGSOQueryType QueryType) {
   return shouldFuncOptimizeForSizeImpl<MachineBasicBlockBFIAdapter>(
       MF, PSI, MBFI, QueryType);
 }
 
 bool llvm::shouldOptimizeForSize(const MachineBasicBlock *MBB,
                                  ProfileSummaryInfo *PSI,
                                  const MachineBlockFrequencyInfo *MBFI,
                                  PGSOQueryType QueryType) {
   return shouldOptimizeForSizeImpl<MachineBasicBlockBFIAdapter>(
       MBB, PSI, MBFI, QueryType);
 }
+
+bool llvm::shouldOptimizeForSize(const MachineBasicBlock *MBB,
+                                 ProfileSummaryInfo *PSI,
+                                 MBFIWrapper *MBFIW,
+                                 PGSOQueryType QueryType) {
+  if (!MBB || !PSI || !MBFIW)
+    return false;
+  BlockFrequency BlockFreq = MBFIW->getBlockFreq(MBB);
+  return shouldOptimizeForSizeImpl<MachineBasicBlockBFIAdapter>(
+      BlockFreq, PSI, &MBFIW->getMBFI(), QueryType);
+}

llvm/lib/CodeGen/TailDuplication.cpp

 using namespace llvm;
 
 #define DEBUG_TYPE "tailduplication"
 
 namespace {
 
 class TailDuplicateBase : public MachineFunctionPass {
   TailDuplicator Duplicator;
+  std::unique_ptr<MBFIWrapper> MBFIW;
   bool PreRegAlloc;
 public:
   TailDuplicateBase(char &PassID, bool PreRegAlloc)
     : MachineFunctionPass(PassID), PreRegAlloc(PreRegAlloc) {}
 
   bool runOnMachineFunction(MachineFunction &MF) override;
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
bool TailDuplicateBase::runOnMachineFunction(MachineFunction &MF) {
   if (skipFunction(MF.getFunction()))
     return false;
 
   auto MBPI = &getAnalysis<MachineBranchProbabilityInfo>();
   auto *PSI = &getAnalysis<ProfileSummaryInfoWrapperPass>().getPSI();
   auto *MBFI = (PSI && PSI->hasProfileSummary()) ?
                &getAnalysis<LazyMachineBlockFrequencyInfoPass>().getBFI() :
                nullptr;
-  Duplicator.initMF(MF, PreRegAlloc, MBPI, MBFI, PSI, /*LayoutMode=*/false);
+  if (MBFI)
+    MBFIW = std::make_unique<MBFIWrapper>(*MBFI);
+  Duplicator.initMF(MF, PreRegAlloc, MBPI, MBFI ? MBFIW.get() : nullptr, PSI,
+                    /*LayoutMode=*/false);
 
   bool MadeChange = false;
   while (Duplicator.tailDuplicateBlocks())
     MadeChange = true;
 
   return MadeChange;
 }

llvm/lib/CodeGen/TailDuplicator.cpp

TailDupVerify("tail-dup-verify",
                   cl::desc("Verify sanity of PHI instructions during taildup"),
                   cl::init(false), cl::Hidden);
 
 static cl::opt<unsigned> TailDupLimit("tail-dup-limit", cl::init(~0U),
                                       cl::Hidden);
 
 void TailDuplicator::initMF(MachineFunction &MFin, bool PreRegAlloc,
                             const MachineBranchProbabilityInfo *MBPIin,
-                            const MachineBlockFrequencyInfo *MBFIin,
+                            MBFIWrapper *MBFIin,
                             ProfileSummaryInfo *PSIin,
                             bool LayoutModeIn, unsigned TailDupSizeIn) {
   MF = &MFin;
   TII = MF->getSubtarget().getInstrInfo();
   TRI = MF->getSubtarget().getRegisterInfo();
   MRI = &MF->getRegInfo();
   MMI = &MF->getMMI();
   MBPI = MBPIin;

llvm/test/CodeGen/X86/tail-opts.ll

cont3:
   br i1 %c4, label %cont4, label %abort4
 abort4:
   call void @alt_abort()
   unreachable
 cont4:
   ret void
 }
 
+; This triggers a situation where a new block (bb4 is split) is created and then
+; would be passed to the PGSO interface llvm::shouldOptimizeForSize().
+@GV = global i32 0
+define void @bfi_new_block_pgso(i32 %c) nounwind {
+; CHECK-LABEL: bfi_new_block_pgso:
+; CHECK:       # %bb.0: # %entry
+; CHECK-NEXT:    testl %edi, %edi
+; CHECK-NEXT:    je .LBB14_4
+; CHECK-NEXT:  # %bb.1: # %bb1
+; CHECK-NEXT:    pushq %rax
+; CHECK-NEXT:    cmpl $16, %edi
+; CHECK-NEXT:    je .LBB14_6
+; CHECK-NEXT:  # %bb.2: # %bb1
+; CHECK-NEXT:    cmpl $17, %edi
+; CHECK-NEXT:    je .LBB14_7
+; CHECK-NEXT:  # %bb.3: # %bb4
+; CHECK-NEXT:    popq %rax
+; CHECK-NEXT:    jmp tail_call_me # TAILCALL
+; CHECK-NEXT:  .LBB14_4: # %bb5
+; CHECK-NEXT:    cmpl $128, %edi
+; CHECK-NEXT:    jne .LBB14_8
+; CHECK-NEXT:  # %bb.5: # %return
+; CHECK-NEXT:    retq
+; CHECK-NEXT:  .LBB14_6: # %bb3
+; CHECK-NEXT:    movl $0, {{.*}}(%rip)
+; CHECK-NEXT:  .LBB14_7: # %bb4
+; CHECK-NEXT:    callq func
+; CHECK-NEXT:    popq %rax
+; CHECK-NEXT:  .LBB14_8: # %bb6
+; CHECK-NEXT:    jmp tail_call_me # TAILCALL
+entry:
+  %0 = icmp eq i32 %c, 0
+  br i1 %0, label %bb5, label %bb1
+
+bb1:
+  switch i32 %c, label %bb4 [
+    i32 16, label %bb3
+    i32 17, label %bb2
+  ]
+
+bb2:
+  call void @func()
+  br label %bb4
+
+bb3:
+  store i32 0, i32* @GV
+  call void @func()
+  br label %bb4
+
+bb4:
+  tail call void @tail_call_me()
+  br label %return
+
+bb5:
+  switch i32 %c, label %bb6 [
+    i32 128, label %return
+  ]
+
+bb6:
+  tail call void @tail_call_me()
+  br label %return
+
+return:
+  ret void
+}
+
 !llvm.module.flags = !{!0}
 !0 = !{i32 1, !"ProfileSummary", !1}
 !1 = !{!2, !3, !4, !5, !6, !7, !8, !9}
 !2 = !{!"ProfileFormat", !"InstrProf"}
 !3 = !{!"TotalCount", i64 10000}
 !4 = !{!"MaxCount", i64 10}
 !5 = !{!"MaxInternalCount", i64 1}
 !6 = !{!"MaxFunctionCount", i64 1000}
 !7 = !{!"NumCounts", i64 3}
 !8 = !{!"NumFunctions", i64 3}
 !9 = !{!"DetailedSummary", !10}
 !10 = !{!11, !12, !13}
 !11 = !{i32 10000, i64 100, i32 1}
 !12 = !{i32 999000, i64 100, i32 1}
 !13 = !{i32 999999, i64 1, i32 2}
 !14 = !{!"function_entry_count", i64 0}