[llvm][misexpect] Track provenance of branch weights

[paulkirth]

Summary

MisExpect diagnostics need to know if branch_weight metadata originates
from an llvm.expect intrinsic. Without that information, we check branch
weights multiple times in the case if ThinLTO + SampleProfiling, leading
to some inaccuracy in how we report MisExpect related diagnostics to
users.

This patch allows us to track that provenance by adding a new optional
field to the branch weight metadata type. Since we change the format of
MD_prof metadata in a fundamental way, we need to update code handling
branch weights in a number of places.

We also update the lang ref for branch weights to reflect the change.

[tejohnson]

@davidxl @xur for review and thoughts.

I'm a little wary of requiring that both pieces of metadata be carried together, as that seems very brittle to maintain in the compiler. What would happen if the MD_expected didn't get propagated by some pass along with the MD_prof? I think you would get a false negative, which I suppose is better than a false positive. An alternative, that I guess would require more extensive changes, is to add an additional item to the "branch_weights" list (would need to be obviously distinguishable by type from the list of weights since that can be variable though).

Patch needs tests showing uses of the new metadata, and some documentation in LangRef (i.e. near https://llvm.org/docs/LangRef.html#prof-metadata).

[paulkirth]

In D131306#3756009, @tejohnson wrote:

@davidxl @xur for review and thoughts.

I'm a little wary of requiring that both pieces of metadata be carried together, as that seems very brittle to maintain in the compiler. What would happen if the MD_expected didn't get propagated by some pass along with the MD_prof? I think you would get a false negative, which I suppose is better than a false positive. An alternative, that I guess would require more extensive changes, is to add an additional item to the "branch_weights" list (would need to be obviously distinguishable by type from the list of weights since that can be variable though).

Agreed. This isn't my preferred solution, but it seemed far less invasive than changing the format of profiling metadata. Originally, I had looked into adding a provenance field to the metadata. It required changes to every test that has branch weights, and I balked at submitting that. I'm also a little wary of making a heavily used metadata type larger. I guess there isn't a lot of difference, but the external metadata is optional, and we could remove it after checking is complete.

Unfortunately, I just don't think there is a clean solution here. Either we make an invasive change to the metadata format, or we deal w/ updating 2 pieces of metadata everywhere. I'm just very unsure about which is the right tradeoff.

Patch needs tests showing uses of the new metadata, and some documentation in LangRef (i.e. near https://llvm.org/docs/LangRef.html#prof-metadata).

Yes, testing and documentation is something I plan to improve, but I wanted to get some feedback on this approach before investing too heavily.

[tejohnson]

In D131306#3756074, @paulkirth wrote:

In D131306#3756009, @tejohnson wrote:

@davidxl @xur for review and thoughts.

I'm a little wary of requiring that both pieces of metadata be carried together, as that seems very brittle to maintain in the compiler. What would happen if the MD_expected didn't get propagated by some pass along with the MD_prof? I think you would get a false negative, which I suppose is better than a false positive. An alternative, that I guess would require more extensive changes, is to add an additional item to the "branch_weights" list (would need to be obviously distinguishable by type from the list of weights since that can be variable though).

Agreed. This isn't my preferred solution, but it seemed far less invasive than changing the format of profiling metadata. Originally, I had looked into adding a provenance field to the metadata. It required changes to every test that has branch weights, and I balked at submitting that. I'm also a little wary of making a heavily used metadata type larger. I guess there isn't a lot of difference, but the external metadata is optional, and we could remove it after checking is complete.

Well I was thinking the extra field would be optional as well and could be removed. But understood that this requires more changes (although maybe not if it is optional, and after your recent changes to centralize some of the prof metadata handling in the compiler).

Unfortunately, I just don't think there is a clean solution here. Either we make an invasive change to the metadata format, or we deal w/ updating 2 pieces of metadata everywhere. I'm just very unsure about which is the right tradeoff.

Patch needs tests showing uses of the new metadata, and some documentation in LangRef (i.e. near https://llvm.org/docs/LangRef.html#prof-metadata).

Yes, testing and documentation is something I plan to improve, but I wanted to get some feedback on this approach before investing too heavily.

Ok understood. Hopefully others will chime in with feedback.

[paulkirth]

In D131306#3756087, @tejohnson wrote:

Well I was thinking the extra field would be optional as well and could be removed. But understood that this requires more changes (although maybe not if it is optional, and after your recent changes to centralize some of the prof metadata handling in the compiler).

Hmm, I don't think I considered that a field in the metadata could be optional. Do you mean something like this?

!{!"branch_weights", !10, i32 1717986918, i32 429496731}

where !10 is just some optional metadata, and we'd just ensure things that parse the MD_prof data skip it correctly? Given that we've mostly consolidated how branch weights are extracted and manipulated that might only require a limited number of updates to the code and tests.

[tejohnson]

In D131306#3756153, @paulkirth wrote:

In D131306#3756087, @tejohnson wrote:

Well I was thinking the extra field would be optional as well and could be removed. But understood that this requires more changes (although maybe not if it is optional, and after your recent changes to centralize some of the prof metadata handling in the compiler).

Hmm, I don't think I considered that a field in the metadata could be optional. Do you mean something like this?

!{!"branch_weights", !10, i32 1717986918, i32 429496731}

where !10 is just some optional metadata, and we'd just ensure things that parse the MD_prof data skip it correctly? Given that we've mostly consolidated how branch weights are extracted and manipulated that might only require a limited number of updates to the code and tests.

Yep! Or a string. I was thinking it might be easier at the end of the list of fields, but either way.

[paulkirth]

I think that may be a better approach. Let me investigate a bit and see what the exact tradeoffs are.

If anyone else has a suggestion or thought on another way to handle this, I'm all ears.

[xur]

I agreed with Teresa: adding an optional string is better than using a separated metadata here.

[paulkirth]

rebase and change implementation to include provenance information directly in the MD_prof metadata

[paulkirth]

@tejohnson @xur I kind of dropped the ball on these patches, but what are your thoughts on this approach over the old(more invasive) change to the profdata format I had prototyped before? the patch will obviously need to be rebased, but other than that, do we see a downside to handling provenance tracking for branch weights this way?

[tejohnson]

In D131306#4037037, @paulkirth wrote:

@tejohnson @xur I kind of dropped the ball on these patches, but what are your thoughts on this approach over the old(more invasive) change to the profdata format I had prototyped before? the patch will obviously need to be rebased, but other than that, do we see a downside to handling provenance tracking for branch weights this way?

Sorry, it looks like you were waiting on a review of the latest changes from me but I didn't get to it. I don't recall the other changes you prototyped off the top of my head - can you point me to that? But I don't have an issue with this approach as I recall it seemed cleanest at the time.

[paulkirth]

In D131306#4052782, @tejohnson wrote:

In D131306#4037037, @paulkirth wrote:

@tejohnson @xur I kind of dropped the ball on these patches, but what are your thoughts on this approach over the old(more invasive) change to the profdata format I had prototyped before? the patch will obviously need to be rebased, but other than that, do we see a downside to handling provenance tracking for branch weights this way?

Sorry, it looks like you were waiting on a review of the latest changes from me but I didn't get to it. I don't recall the other changes you prototyped off the top of my head - can you point me to that?

It was just the last revision of this patch https://reviews.llvm.org/D131306?vs=on&id=450448#toc. The way I handled it before was to leave the MD_prof layout alone and use a new MD type to track the provenance. It had the benefit of leaving the layout alone, and the downside that every place that did something w/ MD_prof needed to copy that as well.

But I don't have an issue with this approach as I recall it seemed cleanest at the time.

Sounds good. I'll start work updating this then. Thanks!

[tejohnson]

In D131306#4052878, @paulkirth wrote:

In D131306#4052782, @tejohnson wrote:

In D131306#4037037, @paulkirth wrote:

@tejohnson @xur I kind of dropped the ball on these patches, but what are your thoughts on this approach over the old(more invasive) change to the profdata format I had prototyped before? the patch will obviously need to be rebased, but other than that, do we see a downside to handling provenance tracking for branch weights this way?

Sorry, it looks like you were waiting on a review of the latest changes from me but I didn't get to it. I don't recall the other changes you prototyped off the top of my head - can you point me to that?

It was just the last revision of this patch https://reviews.llvm.org/D131306?vs=on&id=450448#toc. The way I handled it before was to leave the MD_prof layout alone and use a new MD type to track the provenance. It had the benefit of leaving the layout alone, and the downside that every place that did something w/ MD_prof needed to copy that as well.

Ok, yeah, I definitely prefer it to not be a separate metadata that we need to keep in sync. I prefer this approach.

But I don't have an issue with this approach as I recall it seemed cleanest at the time.

Sounds good. I'll start work updating this then. Thanks!

[paulkirth]

Rebase.

• fix test cases
• improve code swaping branch weights in instruction.cpp
• clean up some dead metadata types left over from an old version of this patch

[paulkirth]

Rebase

[jdoerfert]

Drive by: There should be a lang ref component to this. Also, changing branch weight metadata might impact downstream users, an RFC seems in order (assuming I didn't miss one).

[paulkirth]

In D131306#4081864, @jdoerfert wrote:

Drive by: There should be a lang ref component to this.

Thanks for pointing out that we need to update the langref. Offhand, do you recall if there are places other than BranchWeightMetadata.rst that would need to be updated?

Also, changing branch weight metadata might impact downstream users, an RFC seems in order (assuming I didn't miss one).

You didn't miss one. I was under the impression that the layout of MD_prof was an internal detail that wouldn't require one. I think that is doubly true since the new field is optional, and existing bitcode/IR can still be consumed.

But you raise a fair point that it would impact anyone trying to process new IR, so maybe one is in order to be a good citizen. I'll try to set aside some time to draft a short RFC for discourse.

[paulkirth]

Rebase and update lang ref

[paulkirth]

Update summary

[tejohnson]

Sorry for the slow response. Looks pretty good, just a few minor suggestions and questions.

llvm/include/llvm/IR/MDBuilder.h

61 ↗

(On Diff #492493)

Update comment to mention new parameter in this and the below methods.

Also, I wonder if it would be better to make the IsExpected non-optional, to force any new callers to think about whether they need that and help avoid issues where it gets dropped incorrectly?

llvm/include/llvm/IR/ProfDataUtils.h

58 ↗

(On Diff #492493)

Maybe note the meaning of the expected field (i.e. had an llvm.expect intrinsic).

llvm/lib/Analysis/BranchProbabilityInfo.cpp

394 ↗

(On Diff #492493)

Isn't it now the first 1 or 2 are a name?

397 ↗

(On Diff #492493)

Should this be an assert?

llvm/lib/IR/Instructions.cpp

732 ↗

(On Diff #492493)

I guess we don't have to worry about pushing the expect provenance because this is a call and it shouldn't have one. Maybe add an assert that the offset is 1?

llvm/lib/IR/Verifier.cpp

4553 ↗

(On Diff #492493)

Related to an earlier comment regarding calls - maybe the verifier should ensure that "expected" only shows up on certain instruction types?

llvm/lib/Transforms/Utils/MisExpect.cpp

164

Why not check it within verifyMisExpect, seems less error prone than requiring callers to check?

[paulkirth]

llvm/include/llvm/IR/MDBuilder.h

61 ↗

(On Diff #492493)

Also, I wonder if it would be better to make the IsExpected non-optional, to force any new callers to think about whether they need that and help avoid issues where it gets dropped incorrectly?

I think that's a good idea. I can see how much that affects this patch after addressing the current set of comments.

llvm/lib/Analysis/BranchProbabilityInfo.cpp

394 ↗

(On Diff #492493)

Actually, this code block doesn' need to be here at all anymore. That check was moved up into getValidBranchWeightMDNode. I must have missed that when I rebased.

397 ↗

(On Diff #492493)

Same as above. This wasn't an assert originally, so I didn't change its behavior.

llvm/lib/IR/Verifier.cpp

4553 ↗

(On Diff #492493)

I think the issue is that we've overload "branch_weights" for lots of things that aren't branch weights, but are similar. IMO it should be something like: !{!"call_count", i32 2000}, but that likely would break a number of things in the codebase, and require lots of special handling for the various different uses. Right now the IR change I've made doesn't really disambiguate the various cases and just treats them all uniformly.

I'll take a look at updating the verifier rules after addressing the other comments.

llvm/lib/Transforms/Utils/MisExpect.cpp

164

Centralizing the checking would be ideal, but because Frontend and Backend checking are not completely symmetric, we cannot rely on the instruction having the Expected flag set in the metadata when verifyMisExpect is called for Frontend based instrumentation.

Maybe I should update the comment to reflect that the "provenance" check ensures we won't run backend checks multiple times, and put a more explicit comment about it in checkBackendInstrumentation?

[paulkirth]

Rebase and address comments.

• update doc comments and describe IsExpected parameter
• add assert when scaling branch weights on calls
• remove redundant check, now done in getValidBranchWeightMDNode()
• update assert comment, and describe rationale for hasExpectedProvenance in checkBackendInstrumentation()

TODO:

• investigate making IsExpected a non-optional parameter.
• investigate more checks on the instruction type when metadata has an !expected field

[paulkirth]

Remove redunant assert

clang/test/CodeGenCXX/attr-likelihood-if-branch-weights.cpp

if (b)
     [[likely]] {
       return A();
     }
   return B();
 }
 
 bool g() {
   // CHECK-LABEL: define{{.*}} zeroext i1 @_Z1gv
-  // CHECK: br {{.*}} !prof !8
+  // CHECK: br {{.*}} !prof !9
   if (b)
     [[unlikely]] {
       return A();
     }
 
   return B();
 }
 
 bool h() {
   // CHECK-LABEL: define{{.*}} zeroext i1 @_Z1hv
-  // CHECK: br {{.*}} !prof !8
+  // CHECK: br {{.*}} !prof !9
   if (b)
     [[unlikely]] return A();
 
   return B();
 }
 
 void NullStmt() {
   // CHECK-LABEL: define{{.*}}NullStmt
-  // CHECK: br {{.*}} !prof !8
+  // CHECK: br {{.*}} !prof !9
   if (b)
     [[unlikely]];
   else {
     // Make sure the branches aren't optimized away.
     b = true;
   }
 }
 
 void IfStmt() {
   // CHECK-LABEL: define{{.*}}IfStmt
-  // CHECK: br {{.*}} !prof !8
+  // CHECK: br {{.*}} !prof !9
   if (b)
     [[unlikely]] if (B()) {}
 
   // CHECK-NOT: br {{.*}} !prof
   // CHECK: br {{.*}} !prof
   if (b) {
     if (B())
       [[unlikely]] { b = false; }
   }
 }
 
 void WhileStmt() {
   // CHECK-LABEL: define{{.*}}WhileStmt
-  // CHECK: br {{.*}} !prof !8
+  // CHECK: br {{.*}} !prof !9
   if (b)
     [[unlikely]] while (B()) {}
 
   // CHECK-NOT: br {{.*}} %if.end{{.*}} !prof
   if (b)
     // CHECK: br {{.*}} !prof !7
     while (B())
       [[unlikely]] { b = false; }
 }
 
 void DoStmt() {
   // CHECK-LABEL: define{{.*}}DoStmt
-  // CHECK: br {{.*}} !prof !8
+  // CHECK: br {{.*}} !prof !9
   if (b)
     [[unlikely]] do {}
     while (B())
       ;
 
   // CHECK-NOT: br {{.*}} %if.end{{.*}} !prof
   if (b)
     do
       [[unlikely]] {}
     while (B());
 }
 
 void ForStmt() {
   // CHECK-LABEL: define{{.*}}ForStmt
-  // CHECK: br {{.*}} !prof !8
+  // CHECK: br {{.*}} !prof !9
   if (b)
     [[unlikely]] for (; B();) {}
 
   // CHECK-NOT: br {{.*}} %if.end{{.*}} !prof
   if (b)
     // CHECK: br {{.*}} !prof !7
     for (; B();)
       [[unlikely]] {}
 }
 
 void GotoStmt() {
   // CHECK-LABEL: define{{.*}}GotoStmt
-  // CHECK: br {{.*}} !prof !8
+  // CHECK: br {{.*}} !prof !9
   if (b)
     [[unlikely]] goto end;
   else {
     // Make sure the branches aren't optimized away.
     b = true;
   }
 end:;
 }
 
 void ReturnStmt() {
   // CHECK-LABEL: define{{.*}}ReturnStmt
-  // CHECK: br {{.*}} !prof !8
+  // CHECK: br {{.*}} !prof !9
   if (b)
     [[unlikely]] return;
   else {
     // Make sure the branches aren't optimized away.
     b = true;
   }
 }
 
 void SwitchStmt() {
   // CHECK-LABEL: define{{.*}}SwitchStmt
-  // CHECK: br {{.*}} !prof !8
+  // CHECK: br {{.*}} !prof !9
   if (b)
     [[unlikely]] switch (i) {}
   else {
     // Make sure the branches aren't optimized away.
     b = true;
   }
   // CHECK-NOT: br {{.*}} %if.end{{.*}} !prof
   if (b)
     switch (i)
       [[unlikely]] {}
   else {
     // Make sure the branches aren't optimized away.
     b = true;
   }
 }
 
 // CHECK: !7 = !{!"branch_weights", i32 [[UNLIKELY]], i32 [[LIKELY]]}
-// CHECK: !8 = !{!"branch_weights", i32 [[LIKELY]], i32 [[UNLIKELY]]}
+// CHECK: !9 = !{!"branch_weights", i32 [[LIKELY]], i32 [[UNLIKELY]]}

clang/test/CodeGenCXX/attr-likelihood-if-vs-builtin-expect.cpp

 // RUN: %clang_cc1 -O1 -disable-llvm-passes -emit-llvm %s -o - -triple=x86_64-linux-gnu | opt --lower-expect -S | FileCheck %s
 
 // Verifies the output of __builtin_expect versus the output of the likelihood
 // attributes. They should generate the same probabilities for the branches.
 
 extern bool a();
 extern bool b();
 extern bool c();
 
 void ab1(int &i) {
   // CHECK-LABEL: define{{.*}}ab1
-  // CHECK: br {{.*}} !prof [[BW_LIKELY:!.+]]
+  // CHECK: br {{.*}} !prof [[BW_LIKELY:!.+]],
   // CHECK: br {{.*}} !prof [[BW_LIKELY]]
   // CHECK: br {{.*}} !prof [[BW_LIKELY]]
   if (__builtin_expect(a() && b() && a(), 1)) {
     ++i;
   } else {
     --i;
   }
 }
void al(int &i) {
   }
 }
 
 void ab0(int &i) {
   // CHECK-LABEL: define{{.*}}ab0
   // CHECK: br {{.*}}end{{$}}
   // CHECK: br {{.*}}end{{$}}
   // CHECK: br {{.*}}end{{$}}
-  // CHECK: br {{.*}} !prof [[BW_UNLIKELY:!.+]]
+  // CHECK: br {{.*}} !prof [[BW_UNLIKELY:!.+]],
   if (__builtin_expect(a() && b() && c(), 0)) {
     ++i;
   } else {
     --i;
   }
 }
 
 void au(int &i) {

llvm/include/llvm/IR/FixedMetadataKinds.def

 LLVM_FIXED_MD_KIND(MD_vcall_visibility, "vcall_visibility", 28)
 LLVM_FIXED_MD_KIND(MD_noundef, "noundef", 29)
 LLVM_FIXED_MD_KIND(MD_annotation, "annotation", 30)
 LLVM_FIXED_MD_KIND(MD_nosanitize, "nosanitize", 31)
 LLVM_FIXED_MD_KIND(MD_func_sanitize, "func_sanitize", 32)
 LLVM_FIXED_MD_KIND(MD_exclude, "exclude", 33)
 LLVM_FIXED_MD_KIND(MD_memprof, "memprof", 34)
 LLVM_FIXED_MD_KIND(MD_callsite, "callsite", 35)
+LLVM_FIXED_MD_KIND(MD_expected, "expected", 36)

llvm/lib/Transforms/IPO/ArgumentPromotion.cpp

if (InvokeInst *II = dyn_cast<InvokeInst>(&CB)) {
       auto *NewCall = CallInst::Create(NF, Args, OpBundles, "", &CB);
       NewCall->setTailCallKind(cast<CallInst>(&CB)->getTailCallKind());
       NewCS = NewCall;
     }
     NewCS->setCallingConv(CB.getCallingConv());
     NewCS->setAttributes(AttributeList::get(F->getContext(),
                                             CallPAL.getFnAttrs(),
                                             CallPAL.getRetAttrs(), ArgAttrVec));
-    NewCS->copyMetadata(CB, {LLVMContext::MD_prof, LLVMContext::MD_dbg});
+    NewCS->copyMetadata(CB, {LLVMContext::MD_prof, LLVMContext::MD_expected,
+                             LLVMContext::MD_dbg});
     Args.clear();
     ArgAttrVec.clear();
 
     AttributeFuncs::updateMinLegalVectorWidthAttr(*CB.getCaller(),
                                                   LargestVectorWidth);
 
     if (!CB.use_empty()) {
       CB.replaceAllUsesWith(NewCS);

llvm/lib/Transforms/IPO/Attributor.cpp

auto CallSiteReplacementCreator = [&](AbstractCallSite ACS) {
       } else {
         auto *NewCI = CallInst::Create(NewFn, NewArgOperands, OperandBundleDefs,
                                        "", OldCB);
         NewCI->setTailCallKind(cast<CallInst>(OldCB)->getTailCallKind());
         NewCB = NewCI;
       }
 
       // Copy over various properties and the new attributes.
-      NewCB->copyMetadata(*OldCB, {LLVMContext::MD_prof, LLVMContext::MD_dbg});
+      NewCB->copyMetadata(*OldCB,
+                          {LLVMContext::MD_prof, LLVMContext::MD_expected,
+                           LLVMContext::MD_dbg});
       NewCB->setCallingConv(OldCB->getCallingConv());
       NewCB->takeName(OldCB);
       NewCB->setAttributes(AttributeList::get(
           Ctx, OldCallAttributeList.getFnAttrs(),
           OldCallAttributeList.getRetAttrs(), NewArgOperandAttributes));
 
       AttributeFuncs::updateMinLegalVectorWidthAttr(*NewCB->getCaller(),
                                                     LargestVectorWidth);

llvm/lib/Transforms/IPO/DeadArgumentElimination.cpp

if (InvokeInst *II = dyn_cast<InvokeInst>(CB)) {
                                  Args, OpBundles, "", CB);
     } else {
       NewCB = CallInst::Create(NF, Args, OpBundles, "", CB);
       cast<CallInst>(NewCB)->setTailCallKind(
           cast<CallInst>(CB)->getTailCallKind());
     }
     NewCB->setCallingConv(CB->getCallingConv());
     NewCB->setAttributes(PAL);
-    NewCB->copyMetadata(*CB, {LLVMContext::MD_prof, LLVMContext::MD_dbg});
+    NewCB->copyMetadata(*CB, {LLVMContext::MD_prof, LLVMContext::MD_expected,
+                              LLVMContext::MD_dbg});
 
     Args.clear();
 
     if (!CB->use_empty())
       CB->replaceAllUsesWith(NewCB);
 
     NewCB->takeName(CB);
 
if (InvokeInst *II = dyn_cast<InvokeInst>(&CB)) {
                                  Args, OpBundles, "", CB.getParent());
     } else {
       NewCB = CallInst::Create(NFTy, NF, Args, OpBundles, "", &CB);
       cast<CallInst>(NewCB)->setTailCallKind(
           cast<CallInst>(&CB)->getTailCallKind());
     }
     NewCB->setCallingConv(CB.getCallingConv());
     NewCB->setAttributes(NewCallPAL);
-    NewCB->copyMetadata(CB, {LLVMContext::MD_prof, LLVMContext::MD_dbg});
+    NewCB->copyMetadata(CB, {LLVMContext::MD_prof, LLVMContext::MD_expected,
+                             LLVMContext::MD_dbg});
     Args.clear();
     ArgAttrVec.clear();
 
     if (!CB.use_empty() || CB.isUsedByMetadata()) {
       if (NewCB->getType() == CB.getType()) {
         // Return type not changed? Just replace users then.
         CB.replaceAllUsesWith(NewCB);
         NewCB->takeName(&CB);

llvm/lib/Transforms/IPO/LowerTypeTests.cpp

Value *LowerTypeTestsModule::lowerTypeTestCall(Metadata *TypeId, CallInst *CI,
   if (CI->hasOneUse())
     if (auto *Br = dyn_cast<BranchInst>(*CI->user_begin()))
       if (CI->getNextNode() == Br) {
         BasicBlock *Then = InitialBB->splitBasicBlock(CI->getIterator());
         BasicBlock *Else = Br->getSuccessor(1);
         BranchInst *NewBr = BranchInst::Create(Then, Else, OffsetInRange);
         NewBr->setMetadata(LLVMContext::MD_prof,
                            Br->getMetadata(LLVMContext::MD_prof));
+        NewBr->setMetadata(LLVMContext::MD_expected,
+                           Br->getMetadata(LLVMContext::MD_expected));
         ReplaceInstWithInst(InitialBB->getTerminator(), NewBr);
 
         // Update phis in Else resulting from InitialBB being split
         for (auto &Phi : Else->phis())
           Phi.addIncoming(Phi.getIncomingValueForBlock(Then), InitialBB);
 
         IRBuilder<> ThenB(CI);
         return createBitSetTest(ThenB, TIL, BitOffset);

llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp

if (InvokeInst *II = dyn_cast<InvokeInst>(Caller)) {
     cast<CallInst>(NewCall)->setTailCallKind(
         cast<CallInst>(Caller)->getTailCallKind());
   }
   NewCall->takeName(Caller);
   NewCall->setCallingConv(Call.getCallingConv());
   NewCall->setAttributes(NewCallerPAL);
 
   // Preserve prof metadata if any.
-  NewCall->copyMetadata(*Caller, {LLVMContext::MD_prof});
+  NewCall->copyMetadata(*Caller,
+                        {LLVMContext::MD_prof, LLVMContext::MD_expected});
 
   // Insert a cast of the return type as necessary.
   Instruction *NC = NewCall;
   Value *NV = NC;
   if (OldRetTy != NV->getType() && !Caller->use_empty()) {
     if (!NV->getType()->isVoidTy()) {
       NV = NC = CastInst::CreateBitOrPointerCast(NC, OldRetTy);
       NC->setDebugLoc(Caller->getDebugLoc());

llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp

for (const auto &MDPair : MD) {
     // never be the case in LLVM, but we explicitly switch over only known
     // metadata to be conservatively correct. If you are adding metadata to
     // LLVM which pertains to stores, you almost certainly want to add it
     // here.
     switch (ID) {
     case LLVMContext::MD_dbg:
     case LLVMContext::MD_tbaa:
     case LLVMContext::MD_prof:
+    case LLVMContext::MD_expected:
     case LLVMContext::MD_fpmath:
     case LLVMContext::MD_tbaa_struct:
     case LLVMContext::MD_alias_scope:
     case LLVMContext::MD_noalias:
     case LLVMContext::MD_nontemporal:
     case LLVMContext::MD_mem_parallel_loop_access:
     case LLVMContext::MD_access_group:
       // All of these directly apply.

llvm/lib/Transforms/Scalar/LowerExpectIntrinsic.cpp

static bool handleSwitchExpect(SwitchInst &SI) {
 
   misexpect::checkExpectAnnotations(SI, Weights, /*IsFrontend=*/true);
 
   SI.setCondition(ArgValue);
 
   SI.setMetadata(LLVMContext::MD_prof,
                  MDBuilder(CI->getContext()).createBranchWeights(Weights));
 
+  SI.setMetadata(LLVMContext::MD_expected, MDNode::get(CI->getContext(), None));
+
   return true;
 }
 
 /// Handler for PHINodes that define the value argument to an
 /// @llvm.expect call.
 ///
 /// If the operand of the phi has a constant value and it 'contradicts'
 /// with the expected value of phi def, then the corresponding incoming
for (unsigned i = 0, e = PhiDef->getNumIncomingValues(); i != e; ++i) {
     if (IsOpndComingFromSuccessor(BI->getSuccessor(1)))
       BI->setMetadata(LLVMContext::MD_prof,
                       MDB.createBranchWeights(LikelyBranchWeightVal,
                                               UnlikelyBranchWeightVal));
     else if (IsOpndComingFromSuccessor(BI->getSuccessor(0)))
       BI->setMetadata(LLVMContext::MD_prof,
                       MDB.createBranchWeights(UnlikelyBranchWeightVal,
                                               LikelyBranchWeightVal));
+    BI->setMetadata(LLVMContext::MD_expected,
+                    MDNode::get(BI->getContext(), None));
   }
 }
 
 // Handle both BranchInst and SelectInst.
 template <class BrSelInst> static bool handleBrSelExpect(BrSelInst &BSI) {
 
   // Handle non-optimized IR code like:
   //   %expval = call i64 @llvm.expect.i64(i64 %conv1, i64 1)
template <class BrSelInst> static bool handleBrSelExpect(BrSelInst &BSI) {
   if (CmpI)
     CmpI->setOperand(0, ArgValue);
   else
     BSI.setCondition(ArgValue);
 
   misexpect::checkFrontendInstrumentation(BSI, ExpectedWeights);
 
   BSI.setMetadata(LLVMContext::MD_prof, Node);
+  BSI.setMetadata(LLVMContext::MD_expected,
+                  MDNode::get(BSI.getContext(), None));
 
   return true;
 }
 
 static bool handleBranchExpect(BranchInst &BI) {
   if (BI.isUnconditional())
     return false;
 

llvm/lib/Transforms/Utils/Local.cpp

if (SI->getNumCases() == 1) {
         ConstantInt *SIDef =
             mdconst::dyn_extract<ConstantInt>(MD->getOperand(1));
         assert(SICase && SIDef);
         // The TrueWeight should be the weight for the single case of SI.
         NewBr->setMetadata(LLVMContext::MD_prof,
                         MDBuilder(BB->getContext()).
                         createBranchWeights(SICase->getValue().getZExtValue(),
                                             SIDef->getValue().getZExtValue()));
+        NewBr->setMetadata(LLVMContext::MD_expected,
+                           SI->getMetadata(LLVMContext::MD_expected));
       }
 
       // Update make.implicit metadata to the newly-created conditional branch.
       MDNode *MakeImplicitMD = SI->getMetadata(LLVMContext::MD_make_implicit);
       if (MakeImplicitMD)
         NewBr->setMetadata(LLVMContext::MD_make_implicit, MakeImplicitMD);
 
       // Delete the old switch.
for (const auto &MDPair : MD) {
     // when the pointer type changes. This should essentially never be the case
     // in LLVM, but we explicitly switch over only known metadata to be
     // conservatively correct. If you are adding metadata to LLVM which pertains
     // to loads, you almost certainly want to add it here.
     switch (ID) {
     case LLVMContext::MD_dbg:
     case LLVMContext::MD_tbaa:
     case LLVMContext::MD_prof:
+    case LLVMContext::MD_expected:
     case LLVMContext::MD_fpmath:
     case LLVMContext::MD_tbaa_struct:
     case LLVMContext::MD_invariant_load:
     case LLVMContext::MD_alias_scope:
     case LLVMContext::MD_noalias:
     case LLVMContext::MD_nontemporal:
     case LLVMContext::MD_mem_parallel_loop_access:
     case LLVMContext::MD_access_group:

llvm/lib/Transforms/Utils/MisExpect.cpp

 
 // Command line option to enable/disable the warning when profile data suggests
 // a mismatch with the use of the llvm.expect intrinsic
 static cl::opt<bool> PGOWarnMisExpect(
     "pgo-warn-misexpect", cl::init(false), cl::Hidden,
     cl::desc("Use this option to turn on/off "
              "warnings about incorrect usage of llvm.expect intrinsics."));
 
+// Command line option for setting the diagnostic tolerance threshold
 static cl::opt<unsigned> MisExpectTolerance(
     "misexpect-tolerance", cl::init(0),
-    cl::desc("Prevents emiting diagnostics when profile counts are "
+    cl::desc("Prevents emitting diagnostics when profile counts are "
              "within N% of the threshold.."));
 
 } // namespace llvm
 
 namespace {
 
 bool isMisExpectDiagEnabled(LLVMContext &Ctx) {
   return PGOWarnMisExpect || Ctx.getMisExpectWarningRequested();
Instruction *getInstCondition(Instruction *I) {
   // TODO: Find a way to resolve condition location for switches
   // Using the condition of the switch seems to often resolve to an earlier
   // point in the program, i.e. the calculation of the switch condition, rather
   // than the switch's location in the source code. Thus, we should use the
   // instruction to get source code locations rather than the condition to
   // improve diagnostic output, such as the caret. If the same problem exists
   // for branch instructions, then we should remove this function and directly
   // use the instruction
-  //
   else if (auto *S = dyn_cast<SwitchInst>(I)) {
     Ret = dyn_cast<Instruction>(S->getCondition());
   }
   return Ret ? Ret : I;
 }
 
 void emitMisexpectDiagnostic(Instruction *I, LLVMContext &Ctx,
                              uint64_t ProfCount, uint64_t TotalCount) {
void verifyMisExpect(Instruction &I, ArrayRef<uint32_t> RealWeights,
   const uint64_t RealWeightsTotal =
       std::accumulate(RealWeights.begin(), RealWeights.end(), (uint64_t)0,
                       std::plus<uint64_t>());
   const uint64_t NumUnlikelyTargets = RealWeights.size() - 1;
 
   uint64_t TotalBranchWeight =
       LikelyBranchWeight + (UnlikelyBranchWeight * NumUnlikelyTargets);
 
-  // FIXME: When we've addressed sample profiling, restore the assertion
-  //
-  // We cannot calculate branch probability if either of these invariants aren't
-  // met. However, MisExpect diagnostics should not prevent code from compiling,
-  // so we simply forgo emitting diagnostics here, and return early.
-  // assert((TotalBranchWeight >= LikelyBranchWeight) && (TotalBranchWeight > 0)
-  //              && "TotalBranchWeight is less than the Likely branch weight");
-  if ((TotalBranchWeight == 0) || (TotalBranchWeight <= LikelyBranchWeight))
-    return;
+  // Failing this assert means that we've either got corrupted metadata, that
+  // we're checking stale or manually inserted branch weights, or that branch
+  // weights are being added multiple times, as is the case for SampleProfiling
+  // under ThinLTO. We gate all known entry paths to verifyMisExpect() by first
+  // checking for the presence of the MD_expected metadata, which is *only*
+  // added in the LowerExpectIntrinsic Pass, avoiding the assert.
+  assert((TotalBranchWeight >= LikelyBranchWeight) && (TotalBranchWeight > 0) &&
+         "TotalBranchWeight is less than the Likely branch weight");
 
   // To determine our threshold value we need to obtain the branch probability
   // for the weights added by llvm.expect and use that proportion to calculate
   // our threshold based on the collected profile data.
   auto LikelyProbablilty = BranchProbability::getBranchProbability(
       LikelyBranchWeight, TotalBranchWeight);
 
   uint64_t ScaledThreshold = LikelyProbablilty.scale(RealWeightsTotal);
void verifyMisExpect(Instruction &I, ArrayRef<uint32_t> RealWeights,
   // When the profile weight is below the threshold, we emit the diagnostic
   if (ProfiledWeight < ScaledThreshold)
     emitMisexpectDiagnostic(&I, I.getContext(), ProfiledWeight,
                             RealWeightsTotal);
 }
 
 void checkBackendInstrumentation(Instruction &I,
                                  const ArrayRef<uint32_t> RealWeights) {
+  if (!I.getMetadata(LLVMContext::MD_expected))
+    return;
+  // TODO consider writing a small pass that will remove all MD_expected
+  // metadata that we can run at after the PGO and SampleProfiling passes
+  //
+  // For now, since we're doing the misexpect check here, drop the MD_expected
+  // metadata
+  I.setMetadata(LLVMContext::MD_expected, nullptr);
   SmallVector<uint32_t> ExpectedWeights;
   if (!extractBranchWeights(I, ExpectedWeights))
     return;
   verifyMisExpect(I, RealWeights, ExpectedWeights);
 }
 
 void checkFrontendInstrumentation(Instruction &I,
                                   const ArrayRef<uint32_t> ExpectedWeights) {

llvm/test/Transforms/LowerExpectIntrinsic/basic.ll

entry:
   %x.addr = alloca i32, align 4
   store i32 %x, i32* %x.addr, align 4
   %tmp = load i32, i32* %x.addr, align 4
   %cmp = icmp slt i32 %tmp, 0
   %conv = zext i1 %cmp to i32
   %conv1 = sext i32 %conv to i64
   %expval = call i64 @llvm.expect.i64(i64 %conv1, i64 0)
   %tobool = icmp ne i64 %expval, 0
-; CHECK: !prof !1
+; CHECK: !prof !2
 ; CHECK-NOT: @llvm.expect
   br i1 %tobool, label %if.then, label %if.end
 
 if.then:                                          ; preds = %entry
   %call = call i32 (...) @f()
   store i32 %call, i32* %retval
   br label %return
 

 define i32 @test6(i32 %x) nounwind uwtable ssp {
 entry:
   %retval = alloca i32, align 4
   %x.addr = alloca i32, align 4
   store i32 %x, i32* %x.addr, align 4
   %tmp = load i32, i32* %x.addr, align 4
   %conv = sext i32 %tmp to i64
   %expval = call i64 @llvm.expect.i64(i64 %conv, i64 2)
-; CHECK: !prof !2
+; CHECK: !prof !3
 ; CHECK-NOT: @llvm.expect
   switch i64 %expval, label %sw.epilog [
     i64 1, label %sw.bb
     i64 2, label %sw.bb
   ]
 
 sw.bb:                                            ; preds = %entry, %entry
   store i32 0, i32* %retval

 define i32 @test7(i32 %x) nounwind uwtable ssp {
 entry:
   %retval = alloca i32, align 4
   %x.addr = alloca i32, align 4
   store i32 %x, i32* %x.addr, align 4
   %tmp = load i32, i32* %x.addr, align 4
   %conv = sext i32 %tmp to i64
   %expval = call i64 @llvm.expect.i64(i64 %conv, i64 1)
-; CHECK: !prof !3
+; CHECK: !prof !4
 ; CHECK-NOT: @llvm.expect
   switch i64 %expval, label %sw.epilog [
     i64 2, label %sw.bb
     i64 3, label %sw.bb
   ]
 
 sw.bb:                                            ; preds = %entry, %entry
   %tmp1 = load i32, i32* %x.addr, align 4
return:                                           ; preds = %if.end, %if.then
   ret i32 %0
 }
 
 ; CHECK-LABEL: @test10(
 define i32 @test10(i64 %t6) {
   %t7 = call i64 @llvm.expect.i64(i64 %t6, i64 0)
   %t8 = icmp ne i64 %t7, 0
   %t9 = select i1 %t8, i32 1, i32 2
-; CHECK: select{{.*}}, !prof !1
+; CHECK: select{{.*}}, !prof !2
   ret i32 %t9
 }
 
 
 declare i1 @llvm.expect.i1(i1, i1) nounwind readnone
 
 ; CHECK: !0 = !{!"branch_weights", i32 2000, i32 1}
-; CHECK: !1 = !{!"branch_weights", i32 1, i32 2000}
-; CHECK: !2 = !{!"branch_weights", i32 1, i32 1, i32 2000}
-; CHECK: !3 = !{!"branch_weights", i32 2000, i32 1, i32 1}
+; CHECK: !2 = !{!"branch_weights", i32 1, i32 2000}
+; CHECK: !3 = !{!"branch_weights", i32 1, i32 1, i32 2000}
+; CHECK: !4 = !{!"branch_weights", i32 2000, i32 1, i32 1}

llvm/test/Transforms/LowerExpectIntrinsic/expect-with-probability.ll

entry:
   %x.addr = alloca i32, align 4
   store i32 %x, i32* %x.addr, align 4
   %tmp = load i32, i32* %x.addr, align 4
   %cmp = icmp slt i32 %tmp, 0
   %conv = zext i1 %cmp to i32
   %conv1 = sext i32 %conv to i64
   %expval = call i64 @llvm.expect.with.probability.i64(i64 %conv1, i64 0, double 8.000000e-01)
   %tobool = icmp ne i64 %expval, 0
-; CHECK: !prof !1
+; CHECK: !prof !2
 ; CHECK-NOT: @llvm.expect.with.probability
   br i1 %tobool, label %if.then, label %if.end
 
 if.then:                                          ; preds = %entry
   %call = call i32 (...) @f()
   store i32 %call, i32* %retval
   br label %return
 

 define i32 @test6(i32 %x) nounwind uwtable ssp {
 entry:
   %retval = alloca i32, align 4
   %x.addr = alloca i32, align 4
   store i32 %x, i32* %x.addr, align 4
   %tmp = load i32, i32* %x.addr, align 4
   %conv = sext i32 %tmp to i64
   %expval = call i64 @llvm.expect.with.probability.i64(i64 %conv, i64 2, double 8.000000e-01)
-; CHECK: !prof !2
+; CHECK: !prof !3
 ; CHECK-NOT: @llvm.expect.with.probability
   switch i64 %expval, label %sw.epilog [
     i64 1, label %sw.bb
     i64 2, label %sw.bb
   ]
 
 sw.bb:                                            ; preds = %entry, %entry
   store i32 0, i32* %retval

 define i32 @test7(i32 %x) nounwind uwtable ssp {
 entry:
   %retval = alloca i32, align 4
   %x.addr = alloca i32, align 4
   store i32 %x, i32* %x.addr, align 4
   %tmp = load i32, i32* %x.addr, align 4
   %conv = sext i32 %tmp to i64
   %expval = call i64 @llvm.expect.with.probability.i64(i64 %conv, i64 1, double 8.000000e-01)
-; CHECK: !prof !3
+; CHECK: !prof !4
 ; CHECK-NOT: @llvm.expect.with.probability
   switch i64 %expval, label %sw.epilog [
     i64 2, label %sw.bb
     i64 3, label %sw.bb
   ]
 
 sw.bb:                                            ; preds = %entry, %entry
   %tmp1 = load i32, i32* %x.addr, align 4
return:                                           ; preds = %if.end, %if.then
   ret i32 %0
 }
 
 ; CHECK-LABEL: @test10(
 define i32 @test10(i64 %t6) {
   %t7 = call i64 @llvm.expect.with.probability.i64(i64 %t6, i64 0, double 8.000000e-01)
   %t8 = icmp ne i64 %t7, 0
   %t9 = select i1 %t8, i32 1, i32 2
-; CHECK: select{{.*}}, !prof !1
+; CHECK: select{{.*}}, !prof !2
   ret i32 %t9
 }
 
 
 declare i1 @llvm.expect.with.probability.i1(i1, i1, double) nounwind readnone
 
 ; CHECK: !0 = !{!"branch_weights", i32 1717986918, i32 429496731}
-; CHECK: !1 = !{!"branch_weights", i32 429496731, i32 1717986918}
-; CHECK: !2 = !{!"branch_weights", i32 214748366, i32 214748366, i32 1717986918}
-; CHECK: !3 = !{!"branch_weights", i32 1717986918, i32 214748366, i32 214748366}
+; CHECK: !2 = !{!"branch_weights", i32 429496731, i32 1717986918}
+; CHECK: !3 = !{!"branch_weights", i32 214748366, i32 214748366, i32 1717986918}
+; CHECK: !4 = !{!"branch_weights", i32 1717986918, i32 214748366, i32 214748366}