[PGO] Set edge weights for indirectbr instruction with profile counts

[xur]

Summary

Current PGO only annotates the edge weight for branch and switch instructions with profile counts. We should also annotate the indirectbr instruction as all the information is there.
This patch enables the annotating for indirectbr instructions. Also uses this annotation in branch probability analysis.

[davidxl]

test/Transforms/PGOProfile/indirectbr.ll

8

Are these needed?

[xur]

test/Transforms/PGOProfile/indirectbr.ll

8

probably not, but not of the .ll tests in that test directory have this triple.

[davidxl]

lgtm

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 TerminatorInst *TI = BB->getTerminator();
   assert(TI->getNumSuccessors() > 1 && "expected more than one successor!");
-  if (!isa<BranchInst>(TI) && !isa<SwitchInst>(TI))
+  if (!(isa<BranchInst>(TI) || isa<SwitchInst>(TI) || isa<IndirectBrInst>(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");

lib/Transforms/Instrumentation/PGOInstrumentation.cpp

 // Assign the scaled count values to the BB with multiple out edges.
 void PGOUseFunc::setBranchWeights() {
   // Generate MD_prof metadata for every branch instruction.
   DEBUG(dbgs() << "\nSetting branch weights.\n");
   for (auto &BB : F) {
     TerminatorInst *TI = BB.getTerminator();
     if (TI->getNumSuccessors() < 2)
       continue;
-    if (!isa<BranchInst>(TI) && !isa<SwitchInst>(TI))
+    if (!(isa<BranchInst>(TI) || isa<SwitchInst>(TI) ||
+          isa<IndirectBrInst>(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();
     SmallVector<uint64_t, 2> EdgeCounts(Size, 0);

test/Transforms/PGOProfile/Inputs/indirectbr.proftext

+# IR level Instrumentation Flag
+:ir
+foo
+# Func Hash:
[phone]
+# Num Counters:
+4
+# Counter Values:
+202
+88
+20
+5
+

test/Transforms/PGOProfile/indirectbr.ll

+; RUN: llvm-profdata merge %S/Inputs/indirectbr.proftext -o %t.profdata
+; RUN: opt < %s -pgo-instr-use -pgo-test-profile-file=%t.profdata -S | FileCheck %s --check-prefix=USE
+; New PM
+; RUN: opt < %s -passes=pgo-instr-use -pgo-test-profile-file=%t.profdata -S | FileCheck %s --check-prefix=USE
+
+target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
[e-mail] = internal unnamed_addr constant [3 x i8*] [i8* blockaddress(@foo, %return), i8* blockaddress(@foo, %label2), i8* blockaddress(@foo, %label3)], align 16
+
+define i32 @foo(i32 %i) {
+entry:
+  %cmp = icmp ult i32 %i, 3
+  br i1 %cmp, label %if.then, label %return
+
+if.then:
+  %idxprom = zext i32 %i to i64
+  %arrayidx = getelementptr inbounds [3 x i8*], [3 x i8*]* @foo.table, i64 0, i64 %idxprom
+  %0 = load i8*, i8** %arrayidx, align 8
+  indirectbr i8* %0, [label %return, label %label2, label %label3]
+; USE:  indirectbr i8* %0, [label %return, label %label2, label %label3]
+; USE-SAME: !prof ![[BW_INDBR:[0-9]+]]
+; USE: ![[BW_INDBR]] = !{!"branch_weights", i32 63, i32 20, i32 5}
+
+label2:
+  br label %return
+
+label3:
+  br label %return
+
+return:
+  %retval.0 = phi i32 [ 3, %label3 ], [ 2, %label2 ], [ 0, %entry ], [ 1, %if.then ]
+  ret i32 %retval.0
+}
+
+