Extend InvokeInst !prof branch_weights metadata to unwind branches

[yrouban]

Summary

Allow InvokeInst to have the second optional prof branch weight for its unwind branch.
InvokeInst is a terminator with two successors. It might have its unwind branch taken many times. If so the BranchProbabilityInfo unwind branch heuristic can be inaccurate. This patch allows a higher accuracy calculated with both branch weights set.

Changes:

• A new section about InvokeInst is added to the BranchWeightMetadata page. It states the old information that missed in the doc and adds new about the second branch weight.
• Verifier is changed to allow either 1 or 2 branch weights for InvokeInst.
• A new test is written for BranchProbabilityInfo to demonstrate the main improvement of the simple fix in calcMetadataWeights().
• Several new testcases are created for Inliner. Those check that both weights are accounted for invoke instruction weight calculation.

[yrouban]

llvm/lib/IR/Instruction.cpp

779

return true;

[yrouban]

fixed return statement in Instruction::copyProfWeights().

[davidxl]

llvm/lib/IR/Instruction.cpp

760

naming nit: copyProfData or copyProfMetadata because it copies branchweights and VP data

[yrouban]

Renamed method copyProfWeighs() to copyProfMetadata() and improved its description.

[hjyamauchi]

llvm/include/llvm/IR/Instruction.h

337

Can this function move to the CallBase class as it seems it's already specific to CallBase and all the existing call sites call on CallBase.

llvm/lib/IR/Verifier.cpp

4158

"Wrong number of InvokeInst branch_weights operands" to be clearer?

[xur]

Thanks for the fix. We also noticed that the branch weight for invoke instructions were not set properly, even thought we had the profile information in PGO.
I have a check that disables setting the branch weight metadata for IRPGO. Once this patch is in. I can remove that check. Of you can integrate that change into this patch.
It's in PGOUseFunc::setBranchWeights(). Just white-list InvokeInst.

[reames]

Mostly minor comments on style to help drive review forward. This is a drive by comment, don't wait on me please.

I generally think this patch is a really good idea. Longer term, I think we should extend this to calls as well. There's nothing preventing a call from being hot throw w/o a local exception handler and the current static heuristic doesn't consider that case. I specifically do not want this rolled into the current patch. That would be scope creep. :)

llvm/docs/BranchWeightMetadata.rst

98

I'd suggest calling this INVOKE_NORMAL_WEIGHT.

llvm/lib/IR/Instruction.cpp

763–778

Why not just have the argument be of CallBase type then?

773

Note sure you need these validation checks here since the source can be presumed valid.

llvm/lib/IR/Verifier.cpp

4139

Please rearrange this code as:
if (is invoke) {

do invoke check;

} else {

do current code block;

}

llvm/lib/Transforms/IPO/ArgumentPromotion.cpp

340

It looks like copyProfMetadata could be pulled out into a separate NFC. Please do so.

[davidxl]

lgtm with Philip and Rong's comments addressed.

[yrouban]

In D80618#2063604, @xur wrote:

Thanks for the fix. We also noticed that the branch weight for invoke instructions were not set properly, even thought we had the profile information in PGO.
I have a check that disables setting the branch weight metadata for IRPGO. Once this patch is in. I can remove that check. Of you can integrate that change into this patch.
It's in PGOUseFunc::setBranchWeights(). Just white-list InvokeInst.

done

llvm/lib/IR/Instruction.cpp

773

With this comment and having in mind the comments about type of this and SrcCB I decided to get rid of the method at all and just call NewCB->copyMetadata(OldCB, {LLVMContext::MD_prof}).

[yrouban]

Addressed all comments.
Extracted setProfWeights() related changes to D80987.
I will land this patch once D80987 is reviewed and landed.

[yrouban]

no changes but added wider context for ease of review

[w2yehia]

Hi @yrouban
this commit breaks the following testcase. If you agree this is a regression, then can you please revert your commit until this is addressed.

; RUN: opt -simplifycfg -S %s
;
target datalayout = "e-m:e-i64:64-n32:64"
target triple = "powerpc64le-unknown-linux-gnu"

%cM = type { i32 }
%cE = type { i32 }

@use = dso_local unnamed_addr alias void (%cM*), void (%cM*)* @foo

define dso_local void @foo(%cM* %this) unnamed_addr align 2 personality i8* bitcast (i32 (...)* @__gxx_personality_v0 to i8*) !prof !0 {
entry:
  %0 = load void (%cE, %cM*)*, void (%cE, %cM*)** undef, align 8
  invoke void %0(%cE undef, %cM* %this)
          to label %cond.end unwind label %lpad

cond.end:                                         ; preds = %entry
  invoke void @bar()
          to label %for.cond unwind label %lpad22.loopexit.split-lp, !prof !1

for.cond:                             ; preds = %cond.end
  unreachable

lpad:                                             ; preds = %entry
  %1 = landingpad { i8*, i32 }
          catch i8* null
  unreachable

lpad22.loopexit.split-lp:                         ; preds = %cond.end
  %lpad.loopexit.split-lp = landingpad { i8*, i32 }
          catch i8* null
  unreachable
}

declare i32 @__gxx_personality_v0(...)

declare dso_local void @bar() local_unnamed_addr align 2

!0 = !{!"function_entry_count", i64 9}
!1 = !{!"branch_weights", i32 9, i32 0}

The verifier throws:

Wrong number of operands
!1 = !{!"branch_weights", i32 9, i32 0}

The issue is that SimplifyCFG tries to replace an invoke with a call, and copies over the MD as is.
One quick fix is

--- a/llvm/lib/IR/Instruction.cpp
+++ b/llvm/lib/IR/Instruction.cpp
@@ -749,9 +749,14 @@ void Instruction::copyMetadata(const Instruction &SrcInst,
   // new one.
   SmallVector<std::pair<unsigned, MDNode *>, 4> TheMDs;
   SrcInst.getAllMetadataOtherThanDebugLoc(TheMDs);
+  bool fixupBranchWeight = isa<CallInst>(this) && isa<InvokeInst>(SrcInst);
   for (const auto &MD : TheMDs) {
-    if (WL.empty() || WLS.count(MD.first))
-      setMetadata(MD.first, MD.second);
+    if (!WL.empty() && !WLS.count(MD.first))
+      continue;
+    MDNode *MDN = MD.second;
+    if (fixupBranchWeight && MDN->getNumOperands() == 3)
+      MDN = MDNode::get(MDN->getContext(), {MDN->getOperand(0), MDN->getOperand(1)});
+    setMetadata(MD.first, MDN);
   }
   if (WL.empty() || WLS.count(LLVMContext::MD_dbg))
     setDebugLoc(SrcInst.getDebugLoc());

But i'm not sure if it's the best one so I'll leave to you to decide.

Thank you for your attention.

[yrouban]

In D80618#2099753, @w2yehia wrote:

Thank you for reporting the case.
The bug is fixed with https://reviews.llvm.org/D81996 by Hans Wennborg.

llvm/docs/BranchWeightMetadata.rst

 
 .. code-block:: none
 
   !0 = metadata !{
     metadata !"branch_weights",
     i32 <CALL_BRANCH_WEIGHT>
   }
 
+``InvokeInst``
+^^^^^^^^^^^^^^^^^^
+
+Invoke instruction may have branch weight metadata with one or two weights.
+The second weight is optional and corresponds to the unwind branch.
+If only one weight is set then it contains the execution count of the call
+and used in SamplePGO mode only as described for the call instruction. If both
+weights are specified then the second weight contains count of unwind branch
+taken and the first weights contains the execution count of the call minus
+the count of unwind branch taken. Both weights specified are used to calculate
+BranchProbability as for BranchInst and for SamplePGO the sum of both weights
+is used.
+
+.. code-block:: none
+
+  !0 = metadata !{
+    metadata !"branch_weights",
+    i32 <INVOKE_BRANCH_WEIGHT>
+    [ , i32 <INVOKE_UNWIND_WEIGHT> ]
+  }
+
 Other
 ^^^^^
 
 Other terminator instructions are not allowed to contain Branch Weight Metadata.
 
 .. _\__builtin_expect:
 
 Built-in ``expect`` Instructions

llvm/include/llvm/IR/Instruction.h

public:
   /// Returns false if no metadata or invalid metadata was found.
   bool extractProfMetadata(uint64_t &TrueVal, uint64_t &FalseVal) const;
 
   /// Retrieve total raw weight values of a branch.
   /// Returns true on success with profile total weights filled in.
   /// Returns false if no metadata was found.
   bool extractProfTotalWeight(uint64_t &TotalVal) const;
 
-  /// Sets the branch_weights metadata to \p W for CallInst.
-  void setProfWeight(uint64_t W);
+  /// Copies the branch_weights metadata if any from source CallBase to this.
+  bool copyProfWeights(const Instruction &SrcCB);
 
   /// Set the debug location information for this instruction.
   void setDebugLoc(DebugLoc Loc) { DbgLoc = std::move(Loc); }
 
   /// Return the debug location for this node as a DebugLoc.
   const DebugLoc &getDebugLoc() const { return DbgLoc; }
 
   /// Set or clear the nuw flag on this instruction, which must be an operator
   /// which supports this flag. See LangRef.html for the meaning of this flag.

llvm/lib/Analysis/BranchProbabilityInfo.cpp

 
 // Propagate existing explicit probabilities from either profile data or
 // 'expect' intrinsic processing. Examine metadata against unreachable
 // heuristic. The probability of the edge coming to unreachable block is
 // set to min of metadata and unreachable heuristic.
 bool BranchProbabilityInfo::calcMetadataWeights(const BasicBlock *BB) {
   const Instruction *TI = BB->getTerminator();
   assert(TI->getNumSuccessors() > 1 && "expected more than one successor!");
-  if (!(isa<BranchInst>(TI) || isa<SwitchInst>(TI) || isa<IndirectBrInst>(TI)))
+  if (!(isa<BranchInst>(TI) || isa<SwitchInst>(TI) || isa<IndirectBrInst>(TI) ||
+        isa<InvokeInst>(TI)))
     return false;
 
   MDNode *WeightsNode = TI->getMetadata(LLVMContext::MD_prof);
   if (!WeightsNode)
     return false;
 
   // Check that the number of successors is manageable.
   assert(TI->getNumSuccessors() < UINT32_MAX && "Too many successors");

llvm/lib/IR/Instruction.cpp

 #undef HANDLE_INST
   }
 
   New->SubclassOptionalData = SubclassOptionalData;
   New->copyMetadata(*this);
   return New;
 }
 
-void Instruction::setProfWeight(uint64_t W) {
+bool Instruction::copyProfWeights(const Instruction &SrcCB) {
   assert(isa<CallBase>(this) &&
          "Can only set weights for call like instructions");
-  SmallVector<uint32_t, 1> Weights;
-  Weights.push_back(W);
-  MDBuilder MDB(getContext());
-  setMetadata(LLVMContext::MD_prof, MDB.createBranchWeights(Weights));
+  assert(isa<CallBase>(&SrcCB) &&
+         "Can only set weights for call like instructions");
+  auto *ProfileData = SrcCB.getMetadata(LLVMContext::MD_prof);
+  if (!ProfileData)
+    return false;
+
+  auto *ProfDataName = dyn_cast<MDString>(ProfileData->getOperand(0));
+  if (!ProfDataName)
+    return false;
+
+  if (!ProfDataName->getString().equals("branch_weights") &&
+      !(ProfDataName->getString().equals("VP") &&
+        ProfileData->getNumOperands() > 3))
+    return false;
+
+  setMetadata(LLVMContext::MD_prof, ProfileData);
 }

llvm/lib/IR/Verifier.cpp

void Verifier::visitProfMetadata(Instruction &I, MDNode *MD) {
   Assert(MD->getOperand(0) != nullptr, "first operand should not be null", MD);
   Assert(isa<MDString>(MD->getOperand(0)),
          "expected string with name of the !prof annotation", MD);
   MDString *MDS = cast<MDString>(MD->getOperand(0));
   StringRef ProfName = MDS->getString();
 
   // Check consistency of !prof branch_weights metadata.
   if (ProfName.equals("branch_weights")) {
     unsigned ExpectedNumOperands = 0;
     if (BranchInst *BI = dyn_cast<BranchInst>(&I))
       ExpectedNumOperands = BI->getNumSuccessors();
     else if (SwitchInst *SI = dyn_cast<SwitchInst>(&I))
       ExpectedNumOperands = SI->getNumSuccessors();
-    else if (isa<CallInst>(&I) || isa<InvokeInst>(&I))
+    else if (isa<CallInst>(&I))
       ExpectedNumOperands = 1;
+    else if (isa<InvokeInst>(&I))
+      ; // ExpectedNumOperands will not be needed.
     else if (IndirectBrInst *IBI = dyn_cast<IndirectBrInst>(&I))
       ExpectedNumOperands = IBI->getNumDestinations();
     else if (isa<SelectInst>(&I))
       ExpectedNumOperands = 2;
     else
       CheckFailed("!prof branch_weights are not allowed for this instruction",
                   MD);
 
+    if (isa<InvokeInst>(&I))
+      Assert(MD->getNumOperands() == 2 || MD->getNumOperands() == 3,
+             "Wrong number of InvokeInst operands", MD);
+    else
       Assert(MD->getNumOperands() == 1 + ExpectedNumOperands,
              "Wrong number of operands", MD);
     for (unsigned i = 1; i < MD->getNumOperands(); ++i) {
       auto &MDO = MD->getOperand(i);
       Assert(MDO, "second operand should not be null", MD);
       Assert(mdconst::dyn_extract<ConstantInt>(MDO),
              "!prof brunch_weights operand is not a const int");
     }
   }
 }

llvm/lib/Transforms/IPO/ArgumentPromotion.cpp

if (InvokeInst *II = dyn_cast<InvokeInst>(&CB)) {
       NewCall->setTailCallKind(cast<CallInst>(&CB)->getTailCallKind());
       NewCS = NewCall;
     }
     NewCS->setCallingConv(CB.getCallingConv());
     NewCS->setAttributes(
         AttributeList::get(F->getContext(), CallPAL.getFnAttributes(),
                            CallPAL.getRetAttributes(), ArgAttrVec));
     NewCS->setDebugLoc(CB.getDebugLoc());
-    uint64_t W;
-    if (CB.extractProfTotalWeight(W))
-      NewCS->setProfWeight(W);
+    NewCS->copyProfWeights(CB);
     Args.clear();
     ArgAttrVec.clear();
 
     // Update the callgraph to know that the callsite has been transformed.
     if (ReplaceCallSite)
       (*ReplaceCallSite)(CB, *NewCS);
 
     if (!CB.use_empty()) {

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.
-      uint64_t W;
-      if (OldCB->extractProfTotalWeight(W))
-        NewCB->setProfWeight(W);
+      NewCB->copyProfWeights(*OldCB);
       NewCB->setCallingConv(OldCB->getCallingConv());
       NewCB->setDebugLoc(OldCB->getDebugLoc());
       NewCB->takeName(OldCB);
       NewCB->setAttributes(AttributeList::get(
           Ctx, OldCallAttributeList.getFnAttributes(),
           OldCallAttributeList.getRetAttributes(), NewArgOperandAttributes));
 
       CallSitePairs.push_back({OldCB, NewCB});

llvm/lib/Transforms/IPO/DeadArgumentElimination.cpp

for (Value::user_iterator I = Fn.user_begin(), E = Fn.user_end(); I != E; ) {
     } else {
       NewCB = CallInst::Create(NF, Args, OpBundles, "", CB);
       cast<CallInst>(NewCB)->setTailCallKind(
           cast<CallInst>(CB)->getTailCallKind());
     }
     NewCB->setCallingConv(CB->getCallingConv());
     NewCB->setAttributes(PAL);
     NewCB->setDebugLoc(CB->getDebugLoc());
-    uint64_t W;
-    if (CB->extractProfTotalWeight(W))
-      NewCB->setProfWeight(W);
+    NewCB->copyProfWeights(*CB);
 
     Args.clear();
 
     if (!CB->use_empty())
       CB->replaceAllUsesWith(NewCB);
 
     NewCB->takeName(CB);
 
while (!F->use_empty()) {
     } 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->setDebugLoc(CB.getDebugLoc());
-    uint64_t W;
-    if (CB.extractProfTotalWeight(W))
-      NewCB->setProfWeight(W);
+    NewCB->copyProfWeights(CB);
     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/InstCombine/InstCombineCalls.cpp

cast<CallInst>(NewCall)->setTailCallKind(
         cast<CallInst>(Caller)->getTailCallKind());
   }
   NewCall->takeName(Caller);
   NewCall->setCallingConv(Call.getCallingConv());
   NewCall->setAttributes(NewCallerPAL);
 
   // Preserve the weight metadata for the new call instruction. The metadata
   // is used by SamplePGO to check callsite's hotness.
-  uint64_t W;
-  if (Caller->extractProfTotalWeight(W))
-    NewCall->setProfWeight(W);
+  NewCall->copyProfWeights(*Caller);
 
   // 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/test/Analysis/BranchProbabilityInfo/basic.ll

%ll = landingpad { i8*, i32 }
           cleanup
   call void @ColdFunc() #0
   br label %if.end
 
 if.end:
   ret i32 0
 }
 
+; CHECK-LABEL: test_invoke_code_profiled
+define void @test_invoke_code_profiled(i1 %c) personality i8* bitcast (i32 (...)* @__gxx_personality_v0 to i8*) {
+entry:
+; CHECK: edge entry -> invoke.to0 probability is 0x7ffff800 / 0x80000000 = 100.00% [HOT edge]
+; CHECK: edge entry -> lpad probability is 0x00000800 / 0x80000000 = 0.00%
+  invoke i32 @InvokeCall() to label %invoke.to0 unwind label %lpad
+
+invoke.to0:
+; CHECK: edge invoke.to0 -> invoke.to1 probability is 0x7ffff800 / 0x80000000 = 100.00% [HOT edge]
+; CHECK: edge invoke.to0 -> lpad probability is 0x00000800 / 0x80000000 = 0.00%
+  invoke i32 @InvokeCall() to label %invoke.to1 unwind label %lpad,
+     !prof !{!"branch_weights", i32 444}
+
+invoke.to1:
+; CHECK: invoke.to1 -> invoke.to2 probability is 0x55555555 / 0x80000000 = 66.67%
+; CHECK: invoke.to1 -> lpad probability is 0x2aaaaaab / 0x80000000 = 33.33%
+  invoke i32 @InvokeCall() to label %invoke.to2 unwind label %lpad,
+     !prof !{!"branch_weights", i32 222, i32 111}
+  ret void
+
+invoke.to2:
+  ret void
+
+lpad:
+  %ll = landingpad { i8*, i32 }
+          cleanup
+  ret void
+}
+
 declare i32 @__gxx_personality_v0(...)
 declare void  @ColdFunc()
 declare i32 @InvokeCall()
 
 attributes #0 = { cold }
 
 
 define i32 @zero1(i32 %i, i32 %a, i32 %b) {

llvm/test/Transforms/Inline/inline-hot-callsite.ll

 ; CHECK: call i32 @callee2
 ; CHECK-NOT: call i32 @callee1
 ; CHECK: ret i32 %x3.i
   %y2 = call i32 @callee2(i32 %y1), !prof !22
   %y3 = call i32 @callee1(i32 %y2), !prof !21
   ret i32 %y3
 }
 
+declare i32 @__gxx_personality_v0(...)
+
+define i32 @invoker2(i32 %y1) personality i8* bitcast (i32 (...)* @__gxx_personality_v0 to i8*) {
+; CHECK-LABEL: @invoker2(
+; CHECK: invoke i32 @callee2
+; CHECK-NOT: invoke i32 @callee1
+; CHECK: ret i32
+  %y2 = invoke i32 @callee2(i32 %y1) to label %next unwind label %lpad, !prof !22
+
+next:
+  %y3 = invoke i32 @callee1(i32 %y2) to label %exit unwind label %lpad, !prof !21
+
+exit:
+  ret i32 1
+
+lpad:
+  %ll = landingpad { i8*, i32 } cleanup
+  ret i32 1
+}
+
+define i32 @invoker3(i32 %y1) personality i8* bitcast (i32 (...)* @__gxx_personality_v0 to i8*) {
+; CHECK-LABEL: @invoker3(
+; CHECK: invoke i32 @callee2
+; CHECK-NOT: invoke i32 @callee1
+; CHECK: ret i32
+  %y2 = invoke i32 @callee2(i32 %y1) to label %next unwind label %lpad,
+          !prof !{!"branch_weights", i64 1, i64 0}
+
+next:
+  %y3 = invoke i32 @callee1(i32 %y2) to label %exit unwind label %lpad,
+          !prof !{!"branch_weights", i64 300, i64 1}
+
+exit:
+  ret i32 1
+
+lpad:
+  %ll = landingpad { i8*, i32 } cleanup
+  ret i32 1
+}
+
+define i32 @invoker4(i32 %y1) personality i8* bitcast (i32 (...)* @__gxx_personality_v0 to i8*) {
+; CHECK-LABEL: @invoker4(
+; CHECK: invoke i32 @callee2
+; CHECK-NOT: invoke i32 @callee1
+; CHECK: ret i32
+  %y2 = invoke i32 @callee2(i32 %y1) to label %next unwind label %lpad,
+          !prof !{!"branch_weights", i64 1, i64 0}
+
+next:
+  %y3 = invoke i32 @callee1(i32 %y2) to label %exit unwind label %lpad,
+          !prof !{!"branch_weights", i64 0, i64 300}
+
+exit:
+  ret i32 1
+
+lpad:
+  %ll = landingpad { i8*, i32 } cleanup
+  ret i32 1
+}
+
 declare void @extern()
 
 !llvm.module.flags = !{!1}
 !21 = !{!"branch_weights", i64 300}
 !22 = !{!"branch_weights", i64 1}
 
 !1 = !{i32 1, !"ProfileSummary", !2}
 !2 = !{!3, !4, !5, !6, !7, !8, !9, !10}

llvm/test/Transforms/InstCombine/cast-call-combine-prof.ll

 lpad:
   %lp = landingpad { i8*, i32 }
           filter [0 x i8*] zeroinitializer
   %ehptr = extractvalue { i8*, i32 } %lp, 0
   tail call void @__cxa_call_unexpected(i8* %ehptr) noreturn nounwind
   unreachable
 }
 
-; CHECK: ![[$PROF]] = !{!"branch_weights", i32 2000}
+; CHECK: ![[$PROF]] = !{!"VP", i32 0, i64 2000, i64 -3913987384944532146, i64 2000}
 !0 = !{!"VP", i32 0, i64 2000, i64 -3913987384944532146, i64 2000}
 
 !llvm.module.flags = !{!1}
 
 !1 = !{i32 1, !"ProfileSummary", !2}
 !2 = !{!3, !4, !5, !6, !7, !8, !9, !10}
 !3 = !{!"ProfileFormat", !"InstrProf"}
 !4 = !{!"TotalCount", i64 10000}