[PGO][profcheck] ignore explicitly cold functions (#151778)

mtrofin · web-flow · commit 9a60841dc446 · 2025-08-04T03:53:49.000+02:00
There is a case when branch profile metadata is OK to miss, namely, cold functions. The goal of the RFC (see the referenced issue) is to avoid accidental omission (and, at a later date, corruption) of profile metadata. However, asking cold functions to have all their conditional branches marked with "0" probabilities would be overdoing it. We can just ask cold functions to have an explicit 0 entry count.

This patch:
- injects an entry count for functions, unless they have one (synthetic or not)
- if the entry count is 0, doesn't inject, nor does it verify the rest of the metadata
- at verification, if the entry count is missing, it reports an error

Issue #147390
diff --git a/llvm/lib/Transforms/Utils/ProfileVerify.cpp b/llvm/lib/Transforms/Utils/ProfileVerify.cpp
@@ -20,8 +20,12 @@
 #include "llvm/IR/MDBuilder.h"
 #include "llvm/IR/ProfDataUtils.h"
 #include "llvm/Support/BranchProbability.h"
+#include "llvm/Support/CommandLine.h"
 
 using namespace llvm;
+static cl::opt<int64_t>
+    DefaultFunctionEntryCount("profcheck-default-function-entry-count",
+                              cl::init(1000));
 namespace {
 class ProfileInjector {
   Function &F;
@@ -63,6 +67,19 @@ bool ProfileInjector::inject() {
   // will get the same BPI it does if the injector wasn't running.
   auto &BPI = FAM.getResult<BranchProbabilityAnalysis>(F);
 
+  // Inject a function count if there's none. It's reasonable for a pass to
+  // want to clear the MD_prof of a function with zero entry count. If the
+  // original profile (iFDO or AFDO) is empty for a function, it's simpler to
+  // require assigning it the 0-entry count explicitly than to mark every branch
+  // as cold (we do want some explicit information in the spirit of what this
+  // verifier wants to achieve - make dropping / corrupting MD_prof
+  // unit-testable)
+  if (!F.getEntryCount(/*AllowSynthetic=*/true))
+    F.setEntryCount(DefaultFunctionEntryCount);
+  // If there is an entry count that's 0, then don't bother injecting. We won't
+  // verify these either.
+  if (F.getEntryCount(/*AllowSynthetic=*/true)->getCount() == 0)
+    return false;
   bool Changed = false;
   for (auto &BB : F) {
     auto *Term = getTerminatorBenefitingFromMDProf(BB);
@@ -119,11 +136,20 @@ PreservedAnalyses ProfileInjectorPass::run(Function &F,
 
 PreservedAnalyses ProfileVerifierPass::run(Function &F,
                                            FunctionAnalysisManager &FAM) {
+  const auto EntryCount = F.getEntryCount(/*AllowSynthetic=*/true);
+  if (!EntryCount) {
+    F.getContext().emitError("Profile verification failed: function entry "
+                             "count missing (set to 0 if cold)");
+    return PreservedAnalyses::all();
+  }
+  if (EntryCount->getCount() == 0)
+    return PreservedAnalyses::all();
   for (const auto &BB : F)
     if (const auto *Term =
             ProfileInjector::getTerminatorBenefitingFromMDProf(BB))
       if (!Term->getMetadata(LLVMContext::MD_prof))
-        F.getContext().emitError("Profile verification failed");
+        F.getContext().emitError(
+            "Profile verification failed: branch annotation missing");
 
-  return PreservedAnalyses::none();
+  return PreservedAnalyses::all();
 }
diff --git a/llvm/test/Transforms/PGOProfile/prof-inject-existing.ll b/llvm/test/Transforms/PGOProfile/prof-inject-existing.ll
@@ -0,0 +1,22 @@
+; Test that prof-inject does not modify existing metadata (incl. "unknown")
+
+; RUN: opt -passes=prof-inject %s -S -o - | FileCheck %s
+
+define void @foo(i32 %i) {
+  %c = icmp eq i32 %i, 0
+  br i1 %c, label %yes, label %no, !prof !0
+yes:
+  br i1 %c, label %yes2, label %no, !prof !1
+yes2:
+  ret void
+no:
+  ret void
+}
+
+!0 = !{!"branch_weights", i32 1, i32 2}
+!1 = !{!"unknown"}
+; CHECK: define void @foo(i32 %i) !prof !0
+; CHECK: br i1 %c, label %yes, label %no, !prof !1
+; CHECK: !0 = !{!"function_entry_count", i64 1000}
+; CHECK: !1 = !{!"branch_weights", i32 1, i32 2}
+; CHECK: !2 = !{!"unknown"}
diff --git a/llvm/test/Transforms/PGOProfile/prof-verify-as-needed.ll b/llvm/test/Transforms/PGOProfile/prof-verify-as-needed.ll
@@ -1,6 +1,6 @@
 ; Test that prof-inject only injects missing metadata
 
-; RUN: opt -passes=prof-inject %s -S -o - | FileCheck %s
+; RUN: opt -passes=prof-inject -profcheck-default-function-entry-count=10 %s -S -o - | FileCheck %s
 
 define void @foo(i32 %i) {
   %c = icmp eq i32 %i, 0
@@ -13,8 +13,26 @@ no:
   ret void
 }
 
+define void @cold(i32 %i) !prof !1 {
+  %c = icmp eq i32 %i, 0
+  br i1 %c, label %yes, label %no
+yes:
+  br i1 %c, label %yes2, label %no
+yes2:
+  ret void
+no:
+  ret void
+}
 !0 = !{!"branch_weights", i32 1, i32 2}
-; CHECK: br i1 %c, label %yes, label %no, !prof !0
-; CHECK: br i1 %c, label %yes2, label %no, !prof !1
-; CHECK: !0 = !{!"branch_weights", i32 1, i32 2}
-; CHECK: !1 = !{!"branch_weights", i32 3, i32 5}
+!1 = !{!"function_entry_count", i32 0}
+
+; CHECK-LABEL: @foo
+; CHECK: br i1 %c, label %yes, label %no, !prof !1
+; CHECK: br i1 %c, label %yes2, label %no, !prof !2
+; CHECK-LABEL: @cold
+; CHECK: br i1 %c, label %yes, label %no{{$}}
+; CHECK: br i1 %c, label %yes2, label %no{{$}}
+; CHECK: !0 = !{!"function_entry_count", i64 10}
+; CHECK: !1 = !{!"branch_weights", i32 1, i32 2}
+; CHECK: !2 = !{!"branch_weights", i32 3, i32 5}
+; CHECK: !3 = !{!"function_entry_count", i32 0}
diff --git a/llvm/test/Transforms/PGOProfile/prof-verify-existing.ll b/llvm/test/Transforms/PGOProfile/prof-verify-existing.ll
@@ -1,21 +1,23 @@
 ; Test that prof-inject does not modify existing metadata (incl. "unknown")
 
-; RUN: opt -passes=prof-inject %s -S -o - | FileCheck %s
 ; RUN: opt -passes=prof-verify %s -S --disable-output
 
-define void @foo(i32 %i) {
+define void @foo(i32 %i) !prof !0 {
   %c = icmp eq i32 %i, 0
-  br i1 %c, label %yes, label %no, !prof !0
+  br i1 %c, label %yes, label %no, !prof !1
 yes:
-  br i1 %c, label %yes2, label %no, !prof !1
+  br i1 %c, label %yes2, label %no, !prof !2
 yes2:
   ret void
 no:
   ret void
 }
 
-!0 = !{!"branch_weights", i32 1, i32 2}
-!1 = !{!"unknown"}
-; CHECK: br i1 %c, label %yes, label %no, !prof !0
-; CHECK: !0 = !{!"branch_weights", i32 1, i32 2}
-; CHECK: !1 = !{!"unknown"}
+!0 = !{!"function_entry_count", i32 1}
+!1 = !{!"branch_weights", i32 1, i32 2}
+!2 = !{!"unknown"}
+; CHECK: define void @foo(i32 %i) !prof !0
+; CHECK: br i1 %c, label %yes, label %no, !prof !1
+; CHECK: !0 = !{!"function_entry_count", i64 1}
+; CHECK: !1 = !{!"branch_weights", i32 1, i32 2}
+; CHECK: !2 = !{!"unknown"}
diff --git a/llvm/test/Transforms/PGOProfile/prof-verify-known-cold.ll b/llvm/test/Transforms/PGOProfile/prof-verify-known-cold.ll
@@ -0,0 +1,15 @@
+; Test prof-verify for functions explicitly marked as cold
+
+; RUN: opt -passes=prof-inject,prof-verify %s -o - 2>&1 | FileCheck %s
+
+define void @foo(i32 %i) !prof !0 {
+  %c = icmp eq i32 %i, 0
+  br i1 %c, label %yes, label %no
+yes:
+  ret void
+no:
+  ret void
+}
+!0 = !{!"function_entry_count", i32 0}
+
+; CHECK-NOT: Profile verification failed
diff --git a/llvm/test/Transforms/PGOProfile/prof-verify-no-entrycount.ll b/llvm/test/Transforms/PGOProfile/prof-verify-no-entrycount.ll
@@ -0,0 +1,14 @@
+; Test prof-verify for functions without entry count
+
+; RUN: not opt -passes=prof-verify %s -o - 2>&1 | FileCheck %s
+
+define void @foo(i32 %i) {
+  %c = icmp eq i32 %i, 0
+  br i1 %c, label %yes, label %no
+yes:
+  ret void
+no:
+  ret void
+}
+
+; CHECK: Profile verification failed: function entry count missing (set to 0 if cold)
diff --git a/llvm/test/Transforms/PGOProfile/prof-verify.ll b/llvm/test/Transforms/PGOProfile/prof-verify.ll
@@ -5,16 +5,17 @@
 ; RUN: opt -passes=prof-inject,prof-verify %s --disable-output
 ; RUN: opt -enable-profcheck %s -S -o - | FileCheck %s --check-prefix=INJECT
 
-define void @foo(i32 %i) {
+define void @foo(i32 %i) !prof !0 {
   %c = icmp eq i32 %i, 0
   br i1 %c, label %yes, label %no
 yes:
   ret void
 no:
   ret void
 }
+!0 = !{!"function_entry_count", i32 1}
 
-; INJECT: br i1 %c, label %yes, label %no, !prof !0
-; INJECT: !0 = !{!"branch_weights", i32 3, i32 5}
+; INJECT: br i1 %c, label %yes, label %no, !prof !1
+; INJECT: !1 = !{!"branch_weights", i32 3, i32 5}
 
-; VERIFY: Profile verification failed
+; VERIFY: Profile verification failed: branch annotation missing
