[llvm] [profcheck] Annotate `select` instructions (PR #152171)

[Mircea Trofin via llvm-commits]

https://github.com/mtrofin updated https://github.com/llvm/llvm-project/pull/152171

>From 5d25fe96b58417b544f4ab88776e860c13a29322 Mon Sep 17 00:00:00 2001
From: Mircea Trofin <mtrofin at google.com>
Date: Tue, 5 Aug 2025 09:29:14 -0700
Subject: [PATCH] [profcheck] Annotate `select` instructions

---
 llvm/lib/Transforms/Utils/ProfileVerify.cpp   | 29 ++++++++-
 .../Transforms/PGOProfile/profcheck-select.ll | 63 +++++++++++++++++++
 2 files changed, 90 insertions(+), 2 deletions(-)
 create mode 100644 llvm/test/Transforms/PGOProfile/profcheck-select.ll

diff --git a/llvm/lib/Transforms/Utils/ProfileVerify.cpp b/llvm/lib/Transforms/Utils/ProfileVerify.cpp
index d67192f9d44ee..b5197204853b6 100644
--- a/llvm/lib/Transforms/Utils/ProfileVerify.cpp
+++ b/llvm/lib/Transforms/Utils/ProfileVerify.cpp
@@ -26,6 +26,18 @@ using namespace llvm;
 static cl::opt<int64_t>
     DefaultFunctionEntryCount("profcheck-default-function-entry-count",
                               cl::init(1000));
+static cl::opt<bool>
+    AnnotateSelect("profcheck-annotate-select", cl::init(true),
+                   cl::desc("Also inject (if missing) and verify MD_prof for "
+                            "`select` instructions"));
+static cl::opt<uint32_t> SelectTrueWeight(
+    "profcheck-select-true-weight", cl::init(2U),
+    cl::desc("When annotating `select` instructions, this value will be used "
+             "for the first ('true') case."));
+static cl::opt<uint32_t> SelectFalseWeight(
+    "profcheck-select-false-weight", cl::init(3U),
+    cl::desc("When annotating `select` instructions, this value will be used "
+             "for the second ('false') case."));
 namespace {
 class ProfileInjector {
   Function &F;
@@ -82,6 +94,13 @@ bool ProfileInjector::inject() {
     return false;
   bool Changed = false;
   for (auto &BB : F) {
+    if (AnnotateSelect) {
+      for (auto &I : BB) {
+        if (isa<SelectInst>(I) && !I.getMetadata(LLVMContext::MD_prof))
+          setBranchWeights(I, {SelectTrueWeight, SelectFalseWeight},
+                           /*IsExpected=*/false);
+      }
+    }
     auto *Term = getTerminatorBenefitingFromMDProf(BB);
     if (!Term || Term->getMetadata(LLVMContext::MD_prof))
       continue;
@@ -144,12 +163,18 @@ PreservedAnalyses ProfileVerifierPass::run(Function &F,
   }
   if (EntryCount->getCount() == 0)
     return PreservedAnalyses::all();
-  for (const auto &BB : F)
+  for (const auto &BB : F) {
+    if (AnnotateSelect) {
+      for (const auto &I : BB)
+        if (isa<SelectInst>(I) && !I.getMetadata(LLVMContext::MD_prof))
+          F.getContext().emitError(
+              "Profile verification failed: select annotation missing");
+    }
     if (const auto *Term =
             ProfileInjector::getTerminatorBenefitingFromMDProf(BB))
       if (!Term->getMetadata(LLVMContext::MD_prof))
         F.getContext().emitError(
             "Profile verification failed: branch annotation missing");
-
+  }
   return PreservedAnalyses::all();
 }
diff --git a/llvm/test/Transforms/PGOProfile/profcheck-select.ll b/llvm/test/Transforms/PGOProfile/profcheck-select.ll
new file mode 100644
index 0000000000000..00077586dcbb0
--- /dev/null
+++ b/llvm/test/Transforms/PGOProfile/profcheck-select.ll
@@ -0,0 +1,63 @@
+; RUN: split-file %s %t
+
+; RUN: opt -passes=prof-inject %t/inject.ll -S -o - | FileCheck %t/inject.ll
+
+; RUN: opt -passes=prof-inject %t/inject-some.ll \
+; RUN:   -profcheck-select-true-weight=1 -profcheck-select-false-weight=6 \
+; RUN:   -S -o - | FileCheck %t/inject-some.ll
+
+; RUN: opt -passes=prof-verify %t/verify.ll 2>&1 | FileCheck %t/verify.ll
+
+; RUN: not opt -passes=prof-verify %t/verify-missing.ll 2>&1 | FileCheck %t/verify-missing.ll
+
+; verify we can disable it. It's sufficient to see opt not failing. 
+; RUN: opt -passes=prof-verify -profcheck-annotate-select=0 %t/verify-missing.ll
+
+;--- inject.ll
+declare void @foo(i32 %a);
+define void @bar(i1 %c) {
+  %v = select i1 %c, i32 1, i32 2
+  call void @foo(i32 %v)
+  ret void
+}
+; CHECK-LABEL: @bar
+; CHECK: %v = select i1 %c, i32 1, i32 2, !prof !1
+; CHECK: !0 = !{!"function_entry_count", i64 1000}
+; CHECK: !1 = !{!"branch_weights", i32 2, i32 3}
+
+;--- inject-some.ll
+declare void @foo(i32 %a);
+define void @bar(i1 %c) {
+  %e = select i1 %c, i32 1, i32 2, !prof !0
+  %c2 = icmp eq i32 %e, 2
+  %v = select i1 %c2, i32 5, i32 10
+  call void @foo(i32 %v)
+  ret void
+}
+!0 = !{!"branch_weights", i32 2, i32 3}
+; CHECK-LABEL: @bar
+; CHECK: %v = select i1 %c2, i32 5, i32 10, !prof !2
+; CHECK: !0 = !{!"function_entry_count", i64 1000}
+; CHECK: !1 = !{!"branch_weights", i32 2, i32 3}
+; CHECK: !2 = !{!"branch_weights", i32 1, i32 6}
+
+;--- verify.ll
+declare void @foo(i32 %a);
+define void @bar(i1 %c) !prof !0 {
+  %v = select i1 %c, i32 1, i32 2, !prof !1
+  call void @foo(i32 %v)
+  ret void
+}
+!0 = !{!"function_entry_count", i64 1000}
+!1 = !{!"branch_weights", i32 1, i32 7}
+; CHECK-NOT: Profile verification failed: select annotation missing
+
+;--- verify-missing.ll
+declare void @foo(i32 %a);
+define void @bar(i1 %c) !prof !0 {
+  %v = select i1 %c, i32 1, i32 2
+  call void @foo(i32 %v)
+  ret void
+}
+!0 = !{!"function_entry_count", i64 1000}
+; CHECK: Profile verification failed: select annotation missing
\ No newline at end of file