[pgo] add means to specify "unknown" MD_prof

[mtrofin]

This PR is part of https://discourse.llvm.org/t/rfc-profile-information-propagation-unittesting/73595

In a slight departure from the RFC, instead of a brand-new MD_prof_unknown kind, this adds a first operand to MD_prof metadata. This makes it easy to replace with valid metadata (only one MD_prof), otherwise sites inserting valid MD_prof would also have to check to remove the unknown one.

The patch just introduces the notion and fixes the verifier accordingly. Existing APIs working (esp. reading) MD_prof will be updated subsequently.

[mtrofin]

• [pgo] add means to specify "unknown" MD_prof #145578 👈 (View in Graphite)
• [IR][PGO] Verify invalid MD_prof metadata on instructions #145576 : 2 other dependent PRs (#145584 , #145721 )
• main

This stack of pull requests is managed by Graphite. Learn more about stacking.

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.

[llvmbot]

@llvm/pr-subscribers-llvm-ir

Author: Mircea Trofin (mtrofin)

Changes

This PR is part of https://discourse.llvm.org/t/rfc-profile-information-propagation-unittesting/73595

In a slight departure from the RFC, instead of a brand-new MD_prof_unknown kind, this adds a first operand to MD_prof metadata. This makes it easy to replace with valid metadata (only one MD_prof), otherwise sites inserting valid MD_prof would also have to check to remove the unknown one.

The patch just introduces the notion and fixes the verifier accordingly. Existing APIs working (esp. reading) MD_prof will be updated subsequently.

---

Full diff: https://github.com/llvm/llvm-project/pull/145578.diff

4 Files Affected:

• (modified) llvm/include/llvm/IR/ProfDataUtils.h (+12)
• (modified) llvm/lib/IR/ProfDataUtils.cpp (+22)
• (modified) llvm/lib/IR/Verifier.cpp (+3)
• (added) llvm/test/Bitcode/branch-weight-unknown.ll (+30)

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}

[nikic]

Can you please check the error message?

[mtrofin]

yes, and I'll move these to llvm/test/Verifier/branch-weight.ll from the last patch

[nikic]

Not seeing an on_function test case in here?

[mtrofin]

added in llvm/test/Verifier/branch-weight.ll

[teresajohnson]

What is the error message in this case?

[mtrofin]

The CHECK one, i.e. "expected either branch_weights or VP profile name".

hmm.. maybe we want something like "unknown should have no operands"? My thinking was that if you had the operands then maybe you messed up the name.... idk. wdyt?

[teresajohnson]

The CHECK one, i.e. "expected either branch_weights or VP profile name".

hmm.. maybe we want something like "unknown should have no operands"? My thinking was that if you had the operands then maybe you messed up the name.... idk. wdyt?

I think if someone specified unknown it would be clearer to give an error about the operands, if that isn't hard to do