Extend InvokeInst !prof branch_weights metadata to unwind branches

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.
 - PGOUseFunc::setBranchWeights() is fixed to be applicable to
   InvokeInst.

Reviewers: davidxl, reames, xur, yamauchi
Tags: #llvm
Differential Revision: https://reviews.llvm.org/D80618

docs/BranchWeightMetadata.rst

@@ -78,6 +78,27 @@ block and entry counts which may not be accurate with sampling.
     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
 ^^^^^
 

lib/Analysis/BranchProbabilityInfo.cpp

@@ -282,7 +282,8 @@ bool BranchProbabilityInfo::calcUnreachableHeuristics(const BasicBlock *BB) {
 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);

lib/IR/Verifier.cpp

@@ -4173,23 +4173,28 @@ void Verifier::visitProfMetadata(Instruction &I, MDNode *MD) {
 
   // 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);
+      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);

lib/Transforms/Instrumentation/PGOInstrumentation.cpp

@@ -1343,7 +1343,7 @@ void PGOUseFunc::setBranchWeights() {
     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)

test/Analysis/BranchProbabilityInfo/basic.ll

@@ -317,6 +317,35 @@ 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()

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

@@ -39,6 +39,66 @@ define i32 @caller2(i32 %y1) {
   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}