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 ↗

(On Diff #266466)

return true;

[yrouban]

fixed return statement in Instruction::copyProfWeights().

[davidxl]

llvm/lib/IR/Instruction.cpp

760 ↗

(On Diff #266756)

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 ↗

(On Diff #267150)

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

4173

"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

781 ↗

(On Diff #267150)

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

791 ↗

(On Diff #267150)

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

llvm/lib/IR/Verifier.cpp

4154

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

do invoke check;

} else {

do current code block;

}

llvm/lib/Transforms/IPO/ArgumentPromotion.cpp

340 ↗

(On Diff #267150)

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

791 ↗

(On Diff #267150)

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

     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_NORMAL_WEIGHT>
+    [ , i32 <INVOKE_UNWIND_WEIGHT> ]
+  }
+
 Other
 ^^^^^
 

llvm/lib/Analysis/BranchProbabilityInfo.cpp

 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);

llvm/lib/IR/Verifier.cpp

 
   // 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))
-      ExpectedNumOperands = 1;
-    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 branch_weights operands", MD);
+    } else {
+      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))
+        ExpectedNumOperands = 1;
+      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);
 
       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);

llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp

     if (TI->getNumSuccessors() < 2)
       continue;
     if (!(isa<BranchInst>(TI) || isa<SwitchInst>(TI) ||
-          isa<IndirectBrInst>(TI)))
+          isa<IndirectBrInst>(TI) || isa<InvokeInst>(TI)))
       continue;
 
     if (getBBInfo(&BB).CountValue == 0)

llvm/test/Analysis/BranchProbabilityInfo/basic.ll

   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()

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

   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}