keeping prof metadata for IndirectBrInst

[spupyrev]

Summary

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.

[wmi]

Thanks for the change. Could you add a test for it?

[spupyrev]

• added a test to verify the presense of metadata
• lint suggestions

[wmi]

LGTM.

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

3–6

[wenlei]

Thanks for the fix, Sergey! Looking forward to more improvements on weight propagation and BFI inference. :) LGTM as well and I will push this one on your behalf.

[xbolva00]

Does llvm docs/langref need to be updated?

[wenlei]

In D99550#2659341, @xbolva00 wrote:

Does llvm docs/langref need to be updated?

prof metadata for indirected branch is supported according to the doc https://llvm.org/docs/BranchWeightMetadata.html, the change fixed an oversight in keeping such metadata. So doc should be good.

llvm/lib/IR/Metadata.cpp

bool Instruction::extractProfMetadata(uint64_t &TrueVal,
 
   TrueVal = CITrue->getValue().getZExtValue();
   FalseVal = CIFalse->getValue().getZExtValue();
 
   return true;
 }
 
 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 ||
 
   TotalVal = 0;
   auto *ProfileData = getMetadata(LLVMContext::MD_prof);
   if (!ProfileData)
     return false;
 
   auto *ProfDataName = dyn_cast<MDString>(ProfileData->getOperand(0));
   if (!ProfDataName)

llvm/lib/Transforms/IPO/SampleProfile.cpp

if (BlockWeights[BB]) {
                         MDB.createBranchWeights(
                             {static_cast<uint32_t>(BlockWeights[BB])}));
         }
       }
     }
     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();
     LLVM_DEBUG(dbgs() << "\nGetting weights for branch at line "
                       << ((BranchLoc) ? Twine(BranchLoc.getLine())
                                       : Twine("<UNKNOWN LOCATION>"))
                       << ".\n");
     SmallVector<uint32_t, 4> Weights;

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

+foo:3200:13
+ 1: 6
+ 2: 86
+ 3: 8212
+ 4: 1
+ 5: 17747
+ 6: 5
+ !CFGChecksum: 158517001042

llvm/test/Transforms/SampleProfile/pseudo-probe-profile-metadata.ll

+; 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
+
+@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}