[llvm] 2299873 - [SamplePGO] Keeping prof metadata for IndirectBrInst

[Wenlei He via llvm-commits]

Author: spupyrev
Date: 2021-03-30T10:44:48-07:00
New Revision: 22998738e8f78b447b04a7fd23efa4deaeda306b

URL: https://github.com/llvm/llvm-project/commit/22998738e8f78b447b04a7fd23efa4deaeda306b
DIFF: https://github.com/llvm/llvm-project/commit/22998738e8f78b447b04a7fd23efa4deaeda306b.diff

LOG: [SamplePGO] Keeping prof metadata for IndirectBrInst

Currently prof metadata with branch counts is added only for BranchInst and SwitchInst, but not for IndirectBrInst. As a result, BPI/BFI make incorrect inferences for indirect branches, which can be very hot.
This diff adds metadata for IndirectBrInst, in addition to BranchInst and SwitchInst.

Reviewed By: wmi, wenlei

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

Added: 
    llvm/test/Transforms/SampleProfile/Inputs/pseudo-probe-profile-metadata.prof
    llvm/test/Transforms/SampleProfile/pseudo-probe-profile-metadata.ll

Modified: 
    llvm/lib/IR/Metadata.cpp
    llvm/lib/Transforms/IPO/SampleProfile.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/lib/IR/Metadata.cpp b/llvm/lib/IR/Metadata.cpp
index 6ac7215645c06..4d6ccfa43a640 100644
--- a/llvm/lib/IR/Metadata.cpp
+++ b/llvm/lib/IR/Metadata.cpp
@@ -1413,12 +1413,12 @@ bool Instruction::extractProfMetadata(uint64_t &TrueVal,
 }
 
 bool Instruction::extractProfTotalWeight(uint64_t &TotalVal) const {
-  assert((getOpcode() == Instruction::Br ||
-          getOpcode() == Instruction::Select ||
-          getOpcode() == Instruction::Call ||
-          getOpcode() == Instruction::Invoke ||
-          getOpcode() == Instruction::Switch) &&
-         "Looking for branch weights on something besides branch");
+  assert(
+      (getOpcode() == Instruction::Br || getOpcode() == Instruction::Select ||
+       getOpcode() == Instruction::Call || getOpcode() == Instruction::Invoke ||
+       getOpcode() == Instruction::IndirectBr ||
+       getOpcode() == Instruction::Switch) &&
+      "Looking for branch weights on something besides branch");
 
   TotalVal = 0;
   auto *ProfileData = getMetadata(LLVMContext::MD_prof);

diff  --git a/llvm/lib/Transforms/IPO/SampleProfile.cpp b/llvm/lib/Transforms/IPO/SampleProfile.cpp
index 1cb0492cb0b43..627b9e3d98906 100644
--- a/llvm/lib/Transforms/IPO/SampleProfile.cpp
+++ b/llvm/lib/Transforms/IPO/SampleProfile.cpp
@@ -1468,7 +1468,8 @@ void SampleProfileLoader::generateMDProfMetadata(Function &F) {
     Instruction *TI = BB->getTerminator();
     if (TI->getNumSuccessors() == 1)
       continue;
-    if (!isa<BranchInst>(TI) && !isa<SwitchInst>(TI))
+    if (!isa<BranchInst>(TI) && !isa<SwitchInst>(TI) &&
+        !isa<IndirectBrInst>(TI))
       continue;
 
     DebugLoc BranchLoc = TI->getDebugLoc();

diff  --git a/llvm/test/Transforms/SampleProfile/Inputs/pseudo-probe-profile-metadata.prof b/llvm/test/Transforms/SampleProfile/Inputs/pseudo-probe-profile-metadata.prof
new file mode 100644
index 0000000000000..3f00007e9ce7a
--- /dev/null
+++ b/llvm/test/Transforms/SampleProfile/Inputs/pseudo-probe-profile-metadata.prof
@@ -0,0 +1,8 @@
+foo:3200:13
+ 1: 6
+ 2: 86
+ 3: 8212
+ 4: 1
+ 5: 17747
+ 6: 5
+ !CFGChecksum: 158517001042

diff  --git a/llvm/test/Transforms/SampleProfile/pseudo-probe-profile-metadata.ll b/llvm/test/Transforms/SampleProfile/pseudo-probe-profile-metadata.ll
new file mode 100644
index 0000000000000..47c874d610bbe
--- /dev/null
+++ b/llvm/test/Transforms/SampleProfile/pseudo-probe-profile-metadata.ll
@@ -0,0 +1,54 @@
+; RUN: opt < %s -passes=pseudo-probe,sample-profile -sample-profile-file=%S/Inputs/pseudo-probe-profile-metadata.prof -S | FileCheck %s
+
+; The test verifies the presence of prof metadata for BranchInst, SwitchInst,
+; and IndirectBrInst
+
+ at yydebug = dso_local global i32 0, align 4
+
+define dso_local i32 @foo() #0 {
+entry:
+  call void @llvm.pseudoprobe(i64 6699318081062747564, i64 1, i32 0, i64 -1)
+  %0 = load i32, i32* @yydebug, align 4
+  %cmp = icmp ne i32 %0, 0
+  br i1 %cmp, label %b1, label %exit
+; CHECK: br i1 %cmp, label %b1, label %exit, !prof ![[ENTRY_PROF:[0-9]+]]
+
+b1:
+  call void @llvm.pseudoprobe(i64 6699318081062747564, i64 2, i32 0, i64 -1)
+  %1 = load i32, i32* @yydebug, align 4
+  switch i32 %1, label %b3 [
+    i32 124, label %indirectgoto
+    i32 92, label %b2
+  ]
+; CHECK: ], !prof ![[SWITCH_PROF:[0-9]+]]
+
+b2:
+  call void @llvm.pseudoprobe(i64 6699318081062747564, i64 3, i32 0, i64 -1)
+  br label %indirectgoto
+
+b3:
+  call void @llvm.pseudoprobe(i64 6699318081062747564, i64 4, i32 0, i64 -1)
+  %2 = load i32, i32* @yydebug, align 4
+  ret i32 %2
+
+indirectgoto:
+  %indirect.goto.dest = alloca i8, align 4
+  call void @llvm.pseudoprobe(i64 6699318081062747564, i64 5, i32 0, i64 -1)
+  indirectbr i8* %indirect.goto.dest, [label %b1, label %b3, label %b2]
+; CHECK: indirectbr i8* %indirect.goto.dest, [label %b1, label %b3, label %b2], !prof ![[GOTO_PROF:[0-9]+]]
+
+exit:
+  call void @llvm.pseudoprobe(i64 6699318081062747564, i64 6, i32 0, i64 -1)
+  %3 = load i32, i32* @yydebug, align 4
+  ret i32 %3
+
+}
+
+attributes #0 = {"use-sample-profile"}
+declare void @llvm.pseudoprobe(i64, i64, i32, i64) #1
+!llvm.pseudo_probe_desc = !{!4496}
+!4496 = !{i64 6699318081062747564, i64 158517001042, !"foo", null}
+
+; CHECK: ![[ENTRY_PROF]] = !{!"branch_weights", i32 10, i32 6}
+; CHECK: ![[SWITCH_PROF]] = !{!"branch_weights", i32 1, i32 9536, i32 1}
+; CHECK: ![[GOTO_PROF]] = !{!"branch_weights", i32 17739, i32 1, i32 1}