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 ↗

(On Diff #334212)

[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::IndirectBr ||
           getOpcode() == Instruction::Switch) &&
          "Looking for branch weights on something besides branch");
 
   TotalVal = 0;
   auto *ProfileData = getMetadata(LLVMContext::MD_prof);
   if (!ProfileData)
     return false;
 

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;