[PGO] Support PGO annotation of CallBrInst

We currently instrument CallBrInst but do not annotate it with
the branch weight. This patch enables PGO annotation of CallBrInst.

Differential Revision: https://reviews.llvm.org/D133040

diff --git a/llvm/lib/Analysis/BranchProbabilityInfo.cpp b/llvm/lib/Analysis/BranchProbabilityInfo.cpp
index 8918fb9..50fdf57 100644 (file)
--- a/llvm/lib/Analysis/BranchProbabilityInfo.cpp
+++ b/llvm/lib/Analysis/BranchProbabilityInfo.cpp
@@ -380,7 +380,7 @@ 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) ||
-        isa<InvokeInst>(TI)))
+        isa<InvokeInst>(TI) || isa<CallBrInst>(TI)))
     return false;
 
   MDNode *WeightsNode = TI->getMetadata(LLVMContext::MD_prof);

diff --git a/llvm/lib/IR/Verifier.cpp b/llvm/lib/IR/Verifier.cpp
index 3c617d2..fe6ae62 100644 (file)
--- a/llvm/lib/IR/Verifier.cpp
+++ b/llvm/lib/IR/Verifier.cpp
@@ -4502,6 +4502,8 @@ void Verifier::visitProfMetadata(Instruction &I, MDNode *MD) {
         ExpectedNumOperands = IBI->getNumDestinations();
       else if (isa<SelectInst>(&I))
         ExpectedNumOperands = 2;
+      else if (CallBrInst *CI = dyn_cast<CallBrInst>(&I))
+        ExpectedNumOperands = CI->getNumSuccessors();
       else
         CheckFailed("!prof branch_weights are not allowed for this instruction",
                     MD);

diff --git a/llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp b/llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp
index ffa222a..310e83d 100644 (file)
--- a/llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp
+++ b/llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp
@@ -1398,7 +1398,8 @@ void PGOUseFunc::setBranchWeights() {
     if (TI->getNumSuccessors() < 2)
       continue;
     if (!(isa<BranchInst>(TI) || isa<SwitchInst>(TI) ||
-          isa<IndirectBrInst>(TI) || isa<InvokeInst>(TI)))
+          isa<IndirectBrInst>(TI) || isa<InvokeInst>(TI) ||
+          isa<CallBrInst>(TI)))
       continue;
 
     if (getBBInfo(&BB).CountValue == 0)

diff --git a/llvm/test/Transforms/PGOProfile/Inputs/callbr.proftext b/llvm/test/Transforms/PGOProfile/Inputs/callbr.proftext
new file mode 100644 (file)
index 0000000..c2b1bee
--- /dev/null
+++ b/llvm/test/Transforms/PGOProfile/Inputs/callbr.proftext
@@ -0,0 +1,8 @@
+# :ir is the flag to indicate this is IR level profile.
+:ir
+a
+784007059655560962
+2
+100
+80
+

diff --git a/llvm/test/Transforms/PGOProfile/callbr.ll b/llvm/test/Transforms/PGOProfile/callbr.ll
index d171f1a..c7baf5b 100644 (file)
--- a/llvm/test/Transforms/PGOProfile/callbr.ll
+++ b/llvm/test/Transforms/PGOProfile/callbr.ll
@@ -1,4 +1,8 @@
 ; RUN: opt -passes=pgo-instr-gen -S 2>&1 < %s | FileCheck %s
+;
+; RUN: llvm-profdata merge %S/Inputs/callbr.proftext -o %t.profdata
+; RUN: opt < %s -passes=pgo-instr-use -pgo-test-profile-file=%t.profdata -S | FileCheck %s --check-prefix=USE
+
 
 define i32 @a() {
 entry:
@@ -15,3 +19,11 @@ b:
   %0 = load i32, ptr %retval, align 4
   ret i32 %0
 }
+; USE-LABEL: @a
+; USE-SAME: !prof ![[BW_ENTRY:[0-9]+]]
+; USE: callbr void asm sideeffect
+; USE: to label
+; USE-SAME: !prof ![[BW_CALLBR:[0-9]+]]
+; USE ![[BW_ENTRY]] = !{!"function_entry_count", i64 100}
+; USE ![[BW_CALLBR]] = !{!"branch_weights", i32 80, i32 20}
+