[llvm-branch-commits] [llvm] [pgo] add means to specify "unknown" MD_prof (PR #145578)

[Mircea Trofin via llvm-branch-commits]

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

>From 6c6ac88c73fbffa56983bf8a0cf269e0bc59cb14 Mon Sep 17 00:00:00 2001
From: Mircea Trofin <mtrofin at google.com>
Date: Tue, 24 Jun 2025 09:50:40 -0700
Subject: [PATCH] [pgo] add means to specify "unknown" MD_prof

---
 llvm/include/llvm/IR/ProfDataUtils.h       | 12 +++++++++
 llvm/lib/IR/ProfDataUtils.cpp              | 22 ++++++++++++++++
 llvm/lib/IR/Verifier.cpp                   |  3 +++
 llvm/test/Bitcode/branch-weight-unknown.ll | 30 ++++++++++++++++++++++
 4 files changed, 67 insertions(+)
 create mode 100644 llvm/test/Bitcode/branch-weight-unknown.ll

diff --git a/llvm/include/llvm/IR/ProfDataUtils.h b/llvm/include/llvm/IR/ProfDataUtils.h
index 8e8d069b836f1..89fa7f735f5d4 100644
--- a/llvm/include/llvm/IR/ProfDataUtils.h
+++ b/llvm/include/llvm/IR/ProfDataUtils.h
@@ -133,6 +133,18 @@ LLVM_ABI bool extractProfTotalWeight(const Instruction &I,
 LLVM_ABI void setBranchWeights(Instruction &I, ArrayRef<uint32_t> Weights,
                                bool IsExpected);
 
+/// Specify that the branch weights for this terminator cannot be known at
+/// compile time. This should only be called by passes, and never as a default
+/// behavior in e.g. MDBuilder. The goal is to use this info to validate passes
+/// do not accidentally drop profile info, and this API is called in cases where
+/// the pass explicitly cannot provide that info. Defaulting it in would hide
+/// bugs where the pass forgets to transfer over or otherwise specify profile
+/// info.
+LLVM_ABI void setExplicitlyUnknownBranchWeights(Instruction &I);
+
+LLVM_ABI bool isExplicitlyUnknownBranchWeightsMetadata(const MDNode &MD);
+LLVM_ABI bool hasExplicitlyUnknownBranchWeights(const Instruction &I);
+
 /// Scaling the profile data attached to 'I' using the ratio of S/T.
 LLVM_ABI void scaleProfData(Instruction &I, uint64_t S, uint64_t T);
 
diff --git a/llvm/lib/IR/ProfDataUtils.cpp b/llvm/lib/IR/ProfDataUtils.cpp
index 21524eb840539..1585771c0d0ae 100644
--- a/llvm/lib/IR/ProfDataUtils.cpp
+++ b/llvm/lib/IR/ProfDataUtils.cpp
@@ -44,6 +44,8 @@ constexpr unsigned MinBWOps = 3;
 // the minimum number of operands for MD_prof nodes with value profiles
 constexpr unsigned MinVPOps = 5;
 
+const char *UnknownBranchWeightsMarker = "unknown";
+
 // We may want to add support for other MD_prof types, so provide an abstraction
 // for checking the metadata type.
 bool isTargetMD(const MDNode *ProfData, const char *Name, unsigned MinOps) {
@@ -232,6 +234,26 @@ bool extractProfTotalWeight(const Instruction &I, uint64_t &TotalVal) {
   return extractProfTotalWeight(I.getMetadata(LLVMContext::MD_prof), TotalVal);
 }
 
+void setExplicitlyUnknownBranchWeights(Instruction &I) {
+  MDBuilder MDB(I.getContext());
+  I.setMetadata(LLVMContext::MD_prof,
+                MDNode::get(I.getContext(),
+                            MDB.createString(UnknownBranchWeightsMarker)));
+}
+
+bool isExplicitlyUnknownBranchWeightsMetadata(const MDNode &MD) {
+  if (MD.getNumOperands() != 1)
+    return false;
+  return MD.getOperand(0).equalsStr(UnknownBranchWeightsMarker);
+}
+
+bool hasExplicitlyUnknownBranchWeights(const Instruction &I) {
+  auto *MD = I.getMetadata(LLVMContext::MD_prof);
+  if (!MD)
+    return false;
+  return isExplicitlyUnknownBranchWeightsMetadata(*MD);
+}
+
 void setBranchWeights(Instruction &I, ArrayRef<uint32_t> Weights,
                       bool IsExpected) {
   MDBuilder MDB(I.getContext());
diff --git a/llvm/lib/IR/Verifier.cpp b/llvm/lib/IR/Verifier.cpp
index ae95e3e2bff8d..0ffe4ac257da5 100644
--- a/llvm/lib/IR/Verifier.cpp
+++ b/llvm/lib/IR/Verifier.cpp
@@ -4964,6 +4964,9 @@ void Verifier::visitDereferenceableMetadata(Instruction& I, MDNode* MD) {
 }
 
 void Verifier::visitProfMetadata(Instruction &I, MDNode *MD) {
+  if (isExplicitlyUnknownBranchWeightsMetadata(*MD))
+    return;
+
   Check(MD->getNumOperands() >= 2,
         "!prof annotations should have no less than 2 operands", MD);
 
diff --git a/llvm/test/Bitcode/branch-weight-unknown.ll b/llvm/test/Bitcode/branch-weight-unknown.ll
new file mode 100644
index 0000000000000..921be1ff5da97
--- /dev/null
+++ b/llvm/test/Bitcode/branch-weight-unknown.ll
@@ -0,0 +1,30 @@
+; Test branch weight unknown validation
+
+; RUN: split-file %s %t
+; RUN: opt -passes=verify %t/correct.ll --disable-output
+; RUN: not opt -passes=verify %t/incorrect.ll --disable-output
+; RUN: not opt -passes=verify %t/on_function.ll --disable-output
+
+;--- correct.ll
+define void @correct(i32 %a) {
+  %c = icmp eq i32 %a, 0
+  br i1 %c, label %yes, label %no, !prof !0
+yes:
+  ret void
+no:
+  ret void
+}
+
+!0 = !{!"unknown"}
+
+;--- incorrect.ll
+define void @correct(i32 %a) {
+  %c = icmp eq i32 %a, 0
+  br i1 %c, label %yes, label %no, !prof !0
+yes:
+  ret void
+no:
+  ret void
+}
+
+!0 = !{!"unknown", i32 12, i32 67}