[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

7 ↗

(On Diff #112410)

Are these needed?

[xur]

test/Transforms/PGOProfile/indirectbr.ll

7 ↗

(On Diff #112410)

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

[davidxl]

lgtm

llvm/trunk/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");

llvm/trunk/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);

llvm/trunk/test/Transforms/PGOProfile/Inputs/indirectbr.proftext

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

llvm/trunk/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
+; RUN: opt < %s -passes=pgo-instr-use -pgo-test-profile-file=%t.profdata -S | opt -S -analyze -branch-prob | FileCheck %s --check-prefix=BRANCHPROB
+
+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
+}
+
+; BRANCHPROB: Printing analysis 'Branch Probability Analysis' for function 'foo':
+; BRANCHPROB:---- Branch Probabilities ----
+; BRANCHPROB:  edge entry -> if.then probability is 0x37c32b17 / 0x80000000 = 43.56%
+; BRANCHPROB:  edge entry -> return probability is 0x483cd4e9 / 0x80000000 = 56.44%
+; BRANCHPROB:  edge if.then -> return probability is 0x5ba2e8ba / 0x80000000 = 71.59%
+; BRANCHPROB:  edge if.then -> label2 probability is 0x1d1745d1 / 0x80000000 = 22.73%
+; BRANCHPROB:  edge if.then -> label3 probability is 0x0745d174 / 0x80000000 = 5.68%
+; BRANCHPROB:  edge label2 -> return probability is 0x80000000 / 0x80000000 = 100.00% [HOT edge]
+; BRANCHPROB:  edge label3 -> return probability is 0x80000000 / 0x80000000 = 100.00% [HOT edge]
+
+
+