Use the Propeller CFG profile in the PGO analysis map if it is available. (#163252)

This PR implements the emitting of the post-link CFG information in PGO
analysis map, as explained in the
[RFC](https://discourse.llvm.org/t/rfc-extending-the-pgo-analysis-map-with-propeller-cfg-frequencies/88617).
This is enabled by a flag `pgo-analysis-map-emit-bb-sections-cfg`.

This PR bumps the SHT_LLVM_BB_ADDR_MAP version to 5.
Also includes some refactoring changes related to storing the CFG in the
Basic block sections profile reader.

Author: rlavaee

llvm/docs/Extensions.rst

@@ -416,7 +416,13 @@ as offsets relative to prior addresses.
 The following versioning schemes are currently supported (newer versions support
 features of the older versions).
 
-Version 4 (newest): Capable of encoding basic block hashes. This feature is
+Version 5 (newest): Capable of encoding Post-Link CFG information, which
+provides basic block and edge frequencies obtained from a post-link tool like
+Propeller, reflecting the final binary layout. This feature is enabled by the 8th
+bit of the feature entry.
+The feature data will be two bytes long to accommodate future extensions.
+
+Version 4: Capable of encoding basic block hashes. This feature is
 enabled by the 7th bit of the feature byte.
 
 Example:
@@ -526,6 +532,13 @@ those bits are:
    defined in ``llvm/Support/BranchProbability.h``. It indicates the probability
    that the block is followed by a given successor block during execution.
 
+#. Post-Link CFG - When enabled, the PGO Analysis Map will include CFG
+   information obtained from a post-link tool, such as Propeller. This feature
+   is enabled with the ``-pgo-analysis-map-emit-bb-sections-cfg`` flag. When
+   this option is active, the map will contain basic block and edge frequencies
+   from the basic block sections profile. This provides more accurate profiling
+   information that reflects the final binary layout.
+
 This extra data requires version 2 or above. This is necessary since successors
 of basic blocks won't know their index but will know their BB ID.
 

llvm/include/llvm/CodeGen/BasicBlockSectionsProfileReader.h

@@ -42,22 +42,47 @@ struct BBClusterInfo {
   unsigned PositionInCluster;
 };
 
-// This represents the raw input profile for one function.
-struct FunctionPathAndClusterInfo {
-  // BB Cluster information specified by `UniqueBBID`s.
-  SmallVector<BBClusterInfo> ClusterInfo;
-  // Paths to clone. A path a -> b -> c -> d implies cloning b, c, and d along
-  // the edge a -> b (a is not cloned). The index of the path in this vector
-  // determines the `UniqueBBID::CloneID` of the cloned blocks in that path.
-  SmallVector<SmallVector<unsigned>> ClonePaths;
+// This represents the CFG profile data for a function.
+struct CFGProfile {
   // Node counts for each basic block.
   DenseMap<UniqueBBID, uint64_t> NodeCounts;
   // Edge counts for each edge, stored as a nested map.
   DenseMap<UniqueBBID, DenseMap<UniqueBBID, uint64_t>> EdgeCounts;
+
   // Hash for each basic block. The Hashes are stored for every original block
   // (not cloned blocks), hence the map key being unsigned instead of
   // UniqueBBID.
   DenseMap<unsigned, uint64_t> BBHashes;
+
+  // Returns the profile count for the given basic block or zero if it does not
+  // exist.
+  uint64_t getBlockCount(const UniqueBBID &BBID) const {
+    return NodeCounts.lookup(BBID);
+  }
+
+  // Returns the profile count for the edge from `SrcBBID` to `SinkBBID` or
+  // zero if it does not exist.
+  uint64_t getEdgeCount(const UniqueBBID &SrcBBID,
+                        const UniqueBBID &SinkBBID) const {
+    auto It = EdgeCounts.find(SrcBBID);
+    if (It == EdgeCounts.end())
+      return 0;
+    return It->second.lookup(SinkBBID);
+  }
+};
+
+// This struct represents the raw optimization profile for a function,
+// including CFG data (block and edge counts) and layout directives (clustering
+// and cloning paths).
+struct FunctionOptimizationProfile {
+  // BB Cluster information specified by `UniqueBBID`s.
+  SmallVector<BBClusterInfo> ClusterInfo;
+  // Paths to clone. A path a -> b -> c -> d implies cloning b, c, and d along
+  // the edge a -> b (a is not cloned). The index of the path in this vector
+  // determines the `UniqueBBID::CloneID` of the cloned blocks in that path.
+  SmallVector<SmallVector<unsigned>> ClonePaths;
+  // Cfg profile data (block and edge frequencies).
+  CFGProfile CFG;
 };
 
 class BasicBlockSectionsProfileReader {
@@ -81,14 +106,17 @@ class BasicBlockSectionsProfileReader {
   SmallVector<SmallVector<unsigned>>
   getClonePathsForFunction(StringRef FuncName) const;
 
-  // Returns the profile count for the edge from `SrcBBID` to `SinkBBID` in
-  // function `FuncName` or zero if it does not exist.
   uint64_t getEdgeCount(StringRef FuncName, const UniqueBBID &SrcBBID,
-                        const UniqueBBID &SinkBBID) const;
+                        const UniqueBBID &DestBBID) const;
 
-  // Return the complete function path and cluster info for the given function.
-  std::pair<bool, FunctionPathAndClusterInfo>
-  getFunctionPathAndClusterInfo(StringRef FuncName) const;
+  // Returns a pointer to the CFGProfile for the function \p FuncName.
+  // Returns nullptr if no profile data is available for the function.
+  const CFGProfile *getFunctionCFGProfile(StringRef FuncName) const {
+    auto It = ProgramOptimizationProfile.find(getAliasName(FuncName));
+    if (It == ProgramOptimizationProfile.end())
+      return nullptr;
+    return &It->second.CFG;
+  }
 
 private:
   StringRef getAliasName(StringRef FuncName) const {
@@ -130,16 +158,14 @@ class BasicBlockSectionsProfileReader {
   // empty string if no debug info is available.
   StringMap<SmallString<128>> FunctionNameToDIFilename;
 
-  // This contains the BB cluster information for the whole program.
-  //
-  // For every function name, it contains the cloning and cluster information
-  // for (all or some of) its basic blocks. The cluster information for every
-  // basic block includes its cluster ID along with the position of the basic
-  // block in that cluster.
-  StringMap<FunctionPathAndClusterInfo> ProgramPathAndClusterInfo;
+  // This map contains the optimization profile for each function in the
+  // program. A function's optimization profile consists of CFG data (node and
+  // edge counts) and layout directives such as basic block clustering and
+  // cloning paths.
+  StringMap<FunctionOptimizationProfile> ProgramOptimizationProfile;
 
   // Some functions have alias names. We use this map to find the main alias
-  // name which appears in ProgramPathAndClusterInfo as a key.
+  // name which appears in ProgramOptimizationProfile as a key.
   StringMap<StringRef> FuncAliasMap;
 };
 
@@ -196,12 +222,11 @@ class BasicBlockSectionsProfileReaderWrapperPass : public ImmutablePass {
   SmallVector<SmallVector<unsigned>>
   getClonePathsForFunction(StringRef FuncName) const;
 
+  const CFGProfile *getFunctionCFGProfile(StringRef FuncName) const;
+
   uint64_t getEdgeCount(StringRef FuncName, const UniqueBBID &SrcBBID,
                         const UniqueBBID &DestBBID) const;
 
-  std::pair<bool, FunctionPathAndClusterInfo>
-  getFunctionPathAndClusterInfo(StringRef FuncName) const;
-
   // Initializes the FunctionNameToDIFilename map for the current module and
   // then reads the profile for the matching functions.
   bool doInitialization(Module &M) override;

llvm/include/llvm/MC/MCContext.h

@@ -175,7 +175,7 @@ class MCContext {
   unsigned GetInstance(unsigned LocalLabelVal);
 
   /// SHT_LLVM_BB_ADDR_MAP version to emit.
-  uint8_t BBAddrMapVersion = 4;
+  uint8_t BBAddrMapVersion = 5;
 
   /// The file name of the log file from the environment variable
   /// AS_SECURE_LOG_FILE.  Which must be set before the .secure_log_unique

llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp

@@ -20,7 +20,6 @@
 #include "WinException.h"
 #include "llvm/ADT/APFloat.h"
 #include "llvm/ADT/APInt.h"
-#include "llvm/ADT/BitmaskEnum.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallPtrSet.h"
@@ -37,6 +36,7 @@
 #include "llvm/BinaryFormat/COFF.h"
 #include "llvm/BinaryFormat/Dwarf.h"
 #include "llvm/BinaryFormat/ELF.h"
+#include "llvm/CodeGen/BasicBlockSectionsProfileReader.h"
 #include "llvm/CodeGen/GCMetadata.h"
 #include "llvm/CodeGen/GCMetadataPrinter.h"
 #include "llvm/CodeGen/LazyMachineBlockFrequencyInfo.h"
@@ -150,6 +150,7 @@ enum class PGOMapFeaturesEnum {
   FuncEntryCount,
   BBFreq,
   BrProb,
+  PropellerCFG,
   All,
 };
 static cl::bits<PGOMapFeaturesEnum> PgoAnalysisMapFeatures(
@@ -166,6 +167,12 @@ static cl::bits<PGOMapFeaturesEnum> PgoAnalysisMapFeatures(
         "Enable extended information within the SHT_LLVM_BB_ADDR_MAP that is "
         "extracted from PGO related analysis."));
 
+static cl::opt<bool> PgoAnalysisMapEmitBBSectionsCfg(
+    "pgo-analysis-map-emit-bb-sections-cfg",
+    cl::desc("Enable the post-link cfg information from the basic block "
+             "sections profile in the PGO analysis map"),
+    cl::Hidden, cl::init(false));
+
 static cl::opt<bool> BBAddrMapSkipEmitBBEntries(
     "basic-block-address-map-skip-bb-entries",
     cl::desc("Skip emitting basic block entries in the SHT_LLVM_BB_ADDR_MAP "
@@ -479,6 +486,7 @@ void AsmPrinter::getAnalysisUsage(AnalysisUsage &AU) const {
   AU.addRequired<MachineBranchProbabilityInfoWrapperPass>();
   if (EmitBBHash)
     AU.addRequired<MachineBlockHashInfo>();
+  AU.addUsedIfAvailable<BasicBlockSectionsProfileReaderWrapperPass>();
 }
 
 bool AsmPrinter::doInitialization(Module &M) {
@@ -1409,7 +1417,7 @@ static uint32_t getBBAddrMapMetadata(const MachineBasicBlock &MBB) {
 
 static llvm::object::BBAddrMap::Features
 getBBAddrMapFeature(const MachineFunction &MF, int NumMBBSectionRanges,
-                    bool HasCalls) {
+                    bool HasCalls, const CFGProfile *FuncCFGProfile) {
   // Ensure that the user has not passed in additional options while also
   // specifying all or none.
   if ((PgoAnalysisMapFeatures.isSet(PGOMapFeaturesEnum::None) ||
@@ -1431,17 +1439,17 @@ getBBAddrMapFeature(const MachineFunction &MF, int NumMBBSectionRanges,
   bool BrProbEnabled =
       AllFeatures ||
       (!NoFeatures && PgoAnalysisMapFeatures.isSet(PGOMapFeaturesEnum::BrProb));
+  bool PostLinkCfgEnabled = FuncCFGProfile && PgoAnalysisMapEmitBBSectionsCfg;
 
   if ((BBFreqEnabled || BrProbEnabled) && BBAddrMapSkipEmitBBEntries) {
     MF.getFunction().getContext().emitError(
-        "BB entries info is required for BBFreq and BrProb "
-        "features");
+        "BB entries info is required for BBFreq and BrProb features");
   }
   return {FuncEntryCountEnabled, BBFreqEnabled, BrProbEnabled,
           MF.hasBBSections() && NumMBBSectionRanges > 1,
           // Use static_cast to avoid breakage of tests on windows.
           static_cast<bool>(BBAddrMapSkipEmitBBEntries), HasCalls,
-          static_cast<bool>(EmitBBHash), false};
+          static_cast<bool>(EmitBBHash), PostLinkCfgEnabled};
 }
 
 void AsmPrinter::emitBBAddrMapSection(const MachineFunction &MF) {
@@ -1450,6 +1458,14 @@ void AsmPrinter::emitBBAddrMapSection(const MachineFunction &MF) {
   assert(BBAddrMapSection && ".llvm_bb_addr_map section is not initialized.");
   bool HasCalls = !CurrentFnCallsiteEndSymbols.empty();
 
+  const BasicBlockSectionsProfileReader *BBSPR = nullptr;
+  if (auto *BBSPRPass =
+          getAnalysisIfAvailable<BasicBlockSectionsProfileReaderWrapperPass>())
+    BBSPR = &BBSPRPass->getBBSPR();
+  const CFGProfile *FuncCFGProfile = nullptr;
+  if (BBSPR)
+    FuncCFGProfile = BBSPR->getFunctionCFGProfile(MF.getFunction().getName());
+
   const MCSymbol *FunctionSymbol = getFunctionBegin();
 
   OutStreamer->pushSection();
@@ -1458,8 +1474,9 @@ void AsmPrinter::emitBBAddrMapSection(const MachineFunction &MF) {
   uint8_t BBAddrMapVersion = OutStreamer->getContext().getBBAddrMapVersion();
   OutStreamer->emitInt8(BBAddrMapVersion);
   OutStreamer->AddComment("feature");
-  auto Features = getBBAddrMapFeature(MF, MBBSectionRanges.size(), HasCalls);
-  OutStreamer->emitInt8(Features.encode());
+  auto Features = getBBAddrMapFeature(MF, MBBSectionRanges.size(), HasCalls,
+                                      FuncCFGProfile);
+  OutStreamer->emitInt16(Features.encode());
   // Emit BB Information for each basic block in the function.
   if (Features.MultiBBRange) {
     OutStreamer->AddComment("number of basic block ranges");
@@ -1563,6 +1580,11 @@ void AsmPrinter::emitBBAddrMapSection(const MachineFunction &MF) {
           OutStreamer->AddComment("basic block frequency");
           OutStreamer->emitULEB128IntValue(
               MBFI->getBlockFreq(&MBB).getFrequency());
+          if (Features.PostLinkCfg) {
+            OutStreamer->AddComment("basic block frequency (propeller)");
+            OutStreamer->emitULEB128IntValue(
+                FuncCFGProfile->getBlockCount(*MBB.getBBID()));
+          }
         }
         if (Features.BrProb) {
           unsigned SuccCount = MBB.succ_size();
@@ -1574,6 +1596,11 @@ void AsmPrinter::emitBBAddrMapSection(const MachineFunction &MF) {
             OutStreamer->AddComment("successor branch probability");
             OutStreamer->emitULEB128IntValue(
                 MBPI->getEdgeProbability(&MBB, SuccMBB).getNumerator());
+            if (Features.PostLinkCfg) {
+              OutStreamer->AddComment("successor branch frequency (propeller)");
+              OutStreamer->emitULEB128IntValue(FuncCFGProfile->getEdgeCount(
+                  *MBB.getBBID(), *SuccMBB->getBBID()));
+            }
           }
         }
       }

llvm/lib/CodeGen/BasicBlockMatchingAndInference.cpp

@@ -119,33 +119,32 @@ BasicBlockMatchingAndInference::initWeightInfoByMatching(MachineFunction &MF) {
   StaleMatcher Matcher;
   Matcher.init(Blocks, Hashes);
   BasicBlockMatchingAndInference::WeightInfo MatchWeight;
-  auto [IsValid, PathAndClusterInfo] =
-      BSPR->getFunctionPathAndClusterInfo(MF.getName());
-  if (!IsValid)
+  const CFGProfile *CFG = BSPR->getFunctionCFGProfile(MF.getName());
+  if (CFG == nullptr)
     return MatchWeight;
-  for (auto &BlockCount : PathAndClusterInfo.NodeCounts) {
-    if (PathAndClusterInfo.BBHashes.count(BlockCount.first.BaseID)) {
-      auto Hash = PathAndClusterInfo.BBHashes[BlockCount.first.BaseID];
+  for (auto &BlockCount : CFG->NodeCounts) {
+    if (CFG->BBHashes.count(BlockCount.first.BaseID)) {
+      auto Hash = CFG->BBHashes.lookup(BlockCount.first.BaseID);
       MachineBasicBlock *Block = Matcher.matchBlock(BlendedBlockHash(Hash));
       // When a basic block has clone copies, sum their counts.
       if (Block != nullptr)
         MatchWeight.BlockWeights[Block] += BlockCount.second;
     }
   }
-  for (auto &PredItem : PathAndClusterInfo.EdgeCounts) {
+  for (auto &PredItem : CFG->EdgeCounts) {
     auto PredID = PredItem.first.BaseID;
-    if (!PathAndClusterInfo.BBHashes.count(PredID))
+    if (!CFG->BBHashes.count(PredID))
       continue;
-    auto PredHash = PathAndClusterInfo.BBHashes[PredID];
+    auto PredHash = CFG->BBHashes.lookup(PredID);
     MachineBasicBlock *PredBlock =
         Matcher.matchBlock(BlendedBlockHash(PredHash));
     if (PredBlock == nullptr)
       continue;
     for (auto &SuccItem : PredItem.second) {
       auto SuccID = SuccItem.first.BaseID;
       auto EdgeWeight = SuccItem.second;
-      if (PathAndClusterInfo.BBHashes.count(SuccID)) {
-        auto SuccHash = PathAndClusterInfo.BBHashes[SuccID];
+      if (CFG->BBHashes.count(SuccID)) {
+        auto SuccHash = CFG->BBHashes.lookup(SuccID);
         MachineBasicBlock *SuccBlock =
             Matcher.matchBlock(BlendedBlockHash(SuccHash));
         // When an edge has clone copies, sum their counts.