[llvm] 0caa4a9 - [PGO] Support PGO annotation of CallBrInst

[Rong Xu via llvm-commits]

Author: Rong Xu
Date: 2022-09-01T14:13:50-07:00
New Revision: 0caa4a9559edcbb4195ce140de0658e46a90b52a

URL: https://github.com/llvm/llvm-project/commit/0caa4a9559edcbb4195ce140de0658e46a90b52a
DIFF: https://github.com/llvm/llvm-project/commit/0caa4a9559edcbb4195ce140de0658e46a90b52a.diff

LOG: [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

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

Modified: 
    llvm/lib/Analysis/BranchProbabilityInfo.cpp
    llvm/lib/IR/Verifier.cpp
    llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp
    llvm/test/Transforms/PGOProfile/callbr.ll

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Analysis/BranchProbabilityInfo.cpp b/llvm/lib/Analysis/BranchProbabilityInfo.cpp
index 8918fb967594..50fdf57e01db 100644
--- 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 3c617d2548c7..fe6ae62813b9 100644
--- 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 ffa222a37547..310e83df34df 100644
--- 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
index 000000000000..c2b1bee27b98
--- /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 d171f1a58758..c7baf5bfa684 100644
--- 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}
+