[llvm] [BasicBlockSections] Propose a staleness detection by checking the func cfg node num (PR #133244)

[Jinjie Huang via llvm-commits]

https://github.com/Jinjie-Huang created https://github.com/llvm/llvm-project/pull/133244

This patch introduces ‘light-weight’ method to detect the propeller profile staleness. 
To achieve this, the CFG node count of a function can be emitted into the profile (create_llvm_prof tool can obtain this information from the .llvm_bb_addr_map). This allows for  the detection of changes in the CFG during optimization.

If 'node_count' is not provided in the profile, the check will not be enabled.


>From d574f68d14f7f7ef9335e270b40afa9955987b24 Mon Sep 17 00:00:00 2001
From: huangjinjie <huangjinjie at bytedance.com>
Date: Thu, 27 Mar 2025 20:22:53 +0800
Subject: [PATCH] propose a check for func cfg node num

---
 .../CodeGen/BasicBlockSectionsProfileReader.h |  7 ++++++
 llvm/lib/CodeGen/BasicBlockSections.cpp       | 12 ++++++++++
 .../BasicBlockSectionsProfileReader.cpp       | 22 +++++++++++++++++++
 3 files changed, 41 insertions(+)

diff --git a/llvm/include/llvm/CodeGen/BasicBlockSectionsProfileReader.h b/llvm/include/llvm/CodeGen/BasicBlockSectionsProfileReader.h
index 08e6a0e3ef629..51000efd7f219 100644
--- a/llvm/include/llvm/CodeGen/BasicBlockSectionsProfileReader.h
+++ b/llvm/include/llvm/CodeGen/BasicBlockSectionsProfileReader.h
@@ -50,6 +50,8 @@ struct FunctionPathAndClusterInfo {
   // 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;
+  // Node count of the func CFG, record for staleness check
+  unsigned NodeCount;
 };
 
 // Provides DenseMapInfo for UniqueBBID.
@@ -94,6 +96,9 @@ class BasicBlockSectionsProfileReader {
   std::pair<bool, SmallVector<BBClusterInfo>>
   getClusterInfoForFunction(StringRef FuncName) const;
 
+  // Returns the cfg node count of the CFG for the given function.
+  unsigned getCfgNodeNumForFunction(StringRef FuncName) const;
+
   // Returns the path clonings for the given function.
   SmallVector<SmallVector<unsigned>>
   getClonePathsForFunction(StringRef FuncName) const;
@@ -201,6 +206,8 @@ class BasicBlockSectionsProfileReaderWrapperPass : public ImmutablePass {
   std::pair<bool, SmallVector<BBClusterInfo>>
   getClusterInfoForFunction(StringRef FuncName) const;
 
+  unsigned getCfgNodeNumForFunction(StringRef FuncName) const;
+
   SmallVector<SmallVector<unsigned>>
   getClonePathsForFunction(StringRef FuncName) const;
 
diff --git a/llvm/lib/CodeGen/BasicBlockSections.cpp b/llvm/lib/CodeGen/BasicBlockSections.cpp
index 1eedfc4b25912..ad2e86ea16c58 100644
--- a/llvm/lib/CodeGen/BasicBlockSections.cpp
+++ b/llvm/lib/CodeGen/BasicBlockSections.cpp
@@ -80,6 +80,7 @@
 #include "llvm/CodeGen/TargetInstrInfo.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/Target/TargetMachine.h"
+#include "llvm/Support/WithColor.h"
 #include <optional>
 
 using namespace llvm;
@@ -311,6 +312,7 @@ bool BasicBlockSections::handleBBSections(MachineFunction &MF) {
   // original layout positions and finding the original fallthroughs.
   MF.RenumberBlocks();
 
+  unsigned NodeCount = 0;
   DenseMap<UniqueBBID, BBClusterInfo> FuncClusterInfo;
   if (BBSectionsType == BasicBlockSection::List) {
     auto [HasProfile, ClusterInfo] =
@@ -318,6 +320,16 @@ bool BasicBlockSections::handleBBSections(MachineFunction &MF) {
             .getClusterInfoForFunction(MF.getName());
     if (!HasProfile)
       return false;
+    
+    NodeCount = getAnalysis<BasicBlockSectionsProfileReaderWrapperPass>()
+                    .getCfgNodeNumForFunction(MF.getName());
+    if ((NodeCount != 0) && (NodeCount != MF.size())) {
+      WithColor::warning() << "MF " << MF.getName() << ": node count mismatch "
+                        << "(profile=" << NodeCount
+                        << " actual=" << MF.size() << ")\n";
+      return false;
+    }
+
     for (auto &BBClusterInfo : ClusterInfo) {
       FuncClusterInfo.try_emplace(BBClusterInfo.BBID, BBClusterInfo);
     }
diff --git a/llvm/lib/CodeGen/BasicBlockSectionsProfileReader.cpp b/llvm/lib/CodeGen/BasicBlockSectionsProfileReader.cpp
index fa54640265162..06e5034522bd4 100644
--- a/llvm/lib/CodeGen/BasicBlockSectionsProfileReader.cpp
+++ b/llvm/lib/CodeGen/BasicBlockSectionsProfileReader.cpp
@@ -69,6 +69,13 @@ BasicBlockSectionsProfileReader::getClusterInfoForFunction(
              : std::pair(false, SmallVector<BBClusterInfo>());
 }
 
+unsigned BasicBlockSectionsProfileReader::getCfgNodeNumForFunction(
+    StringRef FuncName) const {
+  auto R = ProgramPathAndClusterInfo.find(getAliasName(FuncName));
+  return R != ProgramPathAndClusterInfo.end()
+             ? R->second.NodeCount : 0;
+}
+
 SmallVector<SmallVector<unsigned>>
 BasicBlockSectionsProfileReader::getClonePathsForFunction(
     StringRef FuncName) const {
@@ -263,6 +270,16 @@ Error BasicBlockSectionsProfileReader::ReadV0Profile() {
     StringRef S(*LineIt);
     if (S[0] == '@')
       continue;
+
+    // Record the function cfg node num for staleness check
+    if (S.consume_front("$node_count")) {
+      unsigned long long NodeCount;
+      if (getAsUnsignedInteger(S.trim(), 10, NodeCount))
+        return invalidProfileError("Invalid node count value.");
+      if (FI != ProgramBBClusterInfo.end())
+        FI->second.NodeCount = (unsigned)NodeCount;
+      continue;
+    }
     // Check for the leading "!"
     if (!S.consume_front("!") || S.empty())
       break;
@@ -433,6 +450,11 @@ BasicBlockSectionsProfileReaderWrapperPass::getClusterInfoForFunction(
   return BBSPR.getClusterInfoForFunction(FuncName);
 }
 
+unsigned BasicBlockSectionsProfileReaderWrapperPass::getCfgNodeNumForFunction(
+    StringRef FuncName) const {
+  return BBSPR.getCfgNodeNumForFunction(FuncName);
+}
+
 SmallVector<SmallVector<unsigned>>
 BasicBlockSectionsProfileReaderWrapperPass::getClonePathsForFunction(
     StringRef FuncName) const {