[PGO] Support PGO annotation of CallBrInst

[xur]

Summary

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

[davidxl]

lgtm

[wenlei]

cc @hoy

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) ||
-        isa<InvokeInst>(TI)))
+        isa<InvokeInst>(TI) || isa<CallBrInst>(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/Verifier.cpp

if (isa<InvokeInst>(&I)) {
       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 if (CallBrInst *CI = dyn_cast<CallBrInst>(&I))
+        ExpectedNumOperands = CI->getNumSuccessors();
       else
         CheckFailed("!prof branch_weights are not allowed for this instruction",
                     MD);
 
       Check(MD->getNumOperands() == 1 + ExpectedNumOperands,
             "Wrong number of operands", MD);
     }
     for (unsigned i = 1; i < MD->getNumOperands(); ++i) {

llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp

void PGOUseFunc::setBranchWeights() {
   // Generate MD_prof metadata for every branch instruction.
   LLVM_DEBUG(dbgs() << "\nSetting branch weights for func " << F.getName()
                     << " IsCS=" << IsCS << "\n");
   for (auto &BB : F) {
     Instruction *TI = BB.getTerminator();
     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)
       continue;
 
     // We have a non-zero Branch BB.
     const UseBBInfo &BBCountInfo = getBBInfo(&BB);
     unsigned Size = BBCountInfo.OutEdges.size();

llvm/test/Transforms/PGOProfile/Inputs/callbr.proftext

+# :ir is the flag to indicate this is IR level profile.
+:ir
+a
+784007059655560962
+2
+100
+80
+

llvm/test/Transforms/PGOProfile/callbr.ll

 ; 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:
 ; CHECK-NOT: ptrtoint ptr asm sideeffect
 ; CHECK: callbr void asm sideeffect
   %retval = alloca i32, align 4
   callbr void asm sideeffect "", "!i,~{dirflag},~{fpsr},~{flags}"() #1
           to label %asm.fallthrough [label %b]
 
 asm.fallthrough:
   br label %b
 
 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}
+