[PGO] Annotate branch_weight for invoke instruction

[mingmingl]

Summary

Invoke instruction can have both branch_weights and indirect-call VP count.
We currently only have indirect-call VP count for it.
This patch enables branch_weight metadata for invoke instruction.

[davidxl]

Can you simplify the test case (e.g. stripping the debug info)?

[hoy]

llvm/include/llvm/IR/Instructions.h

3969

nit: metadat -> metadata

3973

I'm wondering if this function can be made more generic and not restrict to invoke instruction only. I feel it's likely in the future that more types of metadata will be allowed for an instruction. How about returning a something like a hash map with the keys being the type of metadata?

4011

nit: the return nullptr isn't necessary as BranchWeightMD will be set to null if the branch metadata isn't there? Similarly below.

llvm/lib/IR/Metadata.cpp

1512

This seems to change existing behavior for invokes from legacy IR that have only VP metadata.

[wenlei]

Why is invoke unique here? Can regular call never have both VP and branch weights? Ideally we'd avoid special case for a narrow type of instruction.

[xur]

llvm/include/llvm/IR/Instructions.h

3973

As of now, the branch weight metadata is only attached to terminator instructions in FDO.
In sampleFDO, a branch weight can also be attached to a direct call.
VP metadata (more specifically, indirect call metadata) can only be with indirect-call or invoke instruction.
So in current code, we only have this situation for invoke instruction. This function can be written in a more generic way. But there is no use case for now.
I prefer keep this simple form and add a TODO (and expand when there is a use case). What do you think?

4011

that's correct. Fixed.

[xur]

In D133121#3765650, @wenlei wrote:

Why is invoke unique here? Can regular call never have both VP and branch weights? Ideally we'd avoid special case for a narrow type of instruction.

In FDO, branch weights only attached to terminal instruction (and call instruction is not, while invoke is). Indirect-call VP only attached to indirect-call or invoke instruction.
In SampleFDO, we also attach a branch weight to direct-call. But again indirect-call only only for indirect-call and invoke.
So only invoke instruction can have both profile metadata.

[xur]

Integrated with review comments.

[xur]

In D133121#3764425, @davidxl wrote:

Can you simplify the test case (e.g. stripping the debug info)?

Test case simplified in the new patch set.

[hoy]

llvm/include/llvm/IR/Instructions.h

3973

Sounds good for now.

llvm/lib/IR/Metadata.cpp

1512

How about this? Should both branch_weight and vp metadata be checked here to be downwards compatible?

[xur]

llvm/lib/IR/Metadata.cpp

1512

Sorry that I missed this comment.

Yes. you are right on this. Currently we return the VP counts as the total count.
This does create a backward compatibility issue if we use an older profile.

I will check both VP and branch_weights.

[xur]

Integrated Hongtao's suggestion.

[hoy]

Thanks for the change. Wondering if you saw performance improvement with it.

The corresponding change on the sample profile side isn't included, do you have any plan for it?

llvm/lib/ProfileData/InstrProf.cpp

1077

Can we use get instead of getDistinct so that invokes with same metadata can potentially share the same node?

llvm/lib/Transforms/Utils/Local.cpp

2168

Should the check be if (ProfileData)? If the original invoke doesn't have !VP, the new call should end up with no metadata based on the existing behavior.

[xur]

llvm/lib/ProfileData/InstrProf.cpp

1077

Yes. get() seems to be better suited here.

llvm/lib/Transforms/Utils/Local.cpp

2168

If the original invoke does not have VP, it can still have branch_weights.

Maybe my words in patch description is confusing: this patch tries to fix the case where invoke calls to a indirect function.
If the argument of invoke is a direct function -- it will only have a branch_weight and we are doing the right annotations.

For this transformation, mostly like the invoke is a direct function, in which case, we will annotate a total counts -- FDO will not use
this metadata. AutoFDO might use this in the inliner.

If here the newcall is an indirect function, we have have two annotations in invoke and can only have one in the new call. I choose to keep the indirect metadata. This changes the semantic but I think it's the right thing to do.

[xur]

In D133121#3772956, @hoy wrote:

Thanks for the change. Wondering if you saw performance improvement with it.

The corresponding change on the sample profile side isn't included, do you have any plan for it?

Yes. For AutoFDO, we still need some change for the invoke in sampleloader
The situation is slightly different in AutoFDO. In FDO, we do branch_weight annotation first, then do the VP.

In AutoFDO, we do VP annotation first and then branch_weight annotation. We need to set the branch_weight metadata where there is already a VP metadata (which is the reverse of FDO).
I was planning to have a separated patch for AutoFDO, but I can also include it in this patch.

[hoy]

In D133121#3774657, @xur wrote:

In D133121#3772956, @hoy wrote:

Thanks for the change. Wondering if you saw performance improvement with it.

The corresponding change on the sample profile side isn't included, do you have any plan for it?

Yes. For AutoFDO, we still need some change for the invoke in sampleloader
The situation is slightly different in AutoFDO. In FDO, we do branch_weight annotation first, then do the VP.

In AutoFDO, we do VP annotation first and then branch_weight annotation. We need to set the branch_weight metadata where there is already a VP metadata (which is the reverse of FDO).
I was planning to have a separated patch for AutoFDO, but I can also include it in this patch.

A separate patch sounds good, thanks.

llvm/lib/Transforms/Utils/Local.cpp

2168

I see, thanks for the explanation.

[xur]

In D133121#3774764, @hoy wrote:

In D133121#3774657, @xur wrote:

In D133121#3772956, @hoy wrote:

Thanks for the change. Wondering if you saw performance improvement with it.

The corresponding change on the sample profile side isn't included, do you have any plan for it?

Yes. For AutoFDO, we still need some change for the invoke in sampleloader
The situation is slightly different in AutoFDO. In FDO, we do branch_weight annotation first, then do the VP.

In AutoFDO, we do VP annotation first and then branch_weight annotation. We need to set the branch_weight metadata where there is already a VP metadata (which is the reverse of FDO).
I was planning to have a separated patch for AutoFDO, but I can also include it in this patch.

A separate patch sounds good, thanks.

AutoFDO is less problematic as branch_weights metadata overwrites VP in the end. (and many VP metadata is already consumed by the prompting in preinliner).

In FDO, VP metadata overwrites branch_weights -- This is bad as we lost more information in BPI/BFI.
(branch_weights is more important for downstream opt).

[hoy]

In D133121#3774799, @xur wrote:

In D133121#3774764, @hoy wrote:

In D133121#3774657, @xur wrote:

In D133121#3772956, @hoy wrote:

Thanks for the change. Wondering if you saw performance improvement with it.

The corresponding change on the sample profile side isn't included, do you have any plan for it?

Yes. For AutoFDO, we still need some change for the invoke in sampleloader
The situation is slightly different in AutoFDO. In FDO, we do branch_weight annotation first, then do the VP.

In AutoFDO, we do VP annotation first and then branch_weight annotation. We need to set the branch_weight metadata where there is already a VP metadata (which is the reverse of FDO).
I was planning to have a separated patch for AutoFDO, but I can also include it in this patch.

A separate patch sounds good, thanks.

AutoFDO is less problematic as branch_weights metadata overwrites VP in the end. (and many VP metadata is already consumed by the prompting in preinliner).

In FDO, VP metadata overwrites branch_weights -- This is bad as we lost more information in BPI/BFI.
(branch_weights is more important for downstream opt).

It looks like same happens to autofdo, i.e, no branch_weights assigned to indirect invokes? https://github.com/llvm/llvm-project/blob/7b81a81d5f9cccb1b091cfc5264bc483b0acc83a/llvm/lib/Transforms/IPO/SampleProfile.cpp#L1628

Oh, I guess branch_weight is attached to invokes that are promoted from an indirect invoke, so it is better handled than the FDO case.

llvm/lib/ProfileData/InstrProf.cpp

1077

LGTM after this is fixed.

[wenlei]

llvm/include/llvm/IR/Instructions.h

3973

A slight preference for this to be generic. The problem you're dealing with is the unordered nature of MD. Previously it works because there's the assumption that only one profile MD exists, which is not true for invoke.

It's cleaner and more robust if you simply make all APIs not depending on MD order. You don't need a hashmap, but how about just always search through MD when looking for branch_weights or VP. Specifically, that's changing the general implementation of extractBranchWeights and extractProfTotalWeight to be order independent regardless of the opcode.

[xur]

llvm/include/llvm/IR/Instructions.h

3973

OK. This makes sense. I'll update the patch with this suggestion.
Thanks!

[xur]

In D133121#3774815, @hoy wrote:

In D133121#3774799, @xur wrote:

In D133121#3774764, @hoy wrote:

In D133121#3774657, @xur wrote:

In D133121#3772956, @hoy wrote:

Thanks for the change. Wondering if you saw performance improvement with it.

The corresponding change on the sample profile side isn't included, do you have any plan for it?

Yes. For AutoFDO, we still need some change for the invoke in sampleloader
The situation is slightly different in AutoFDO. In FDO, we do branch_weight annotation first, then do the VP.

In AutoFDO, we do VP annotation first and then branch_weight annotation. We need to set the branch_weight metadata where there is already a VP metadata (which is the reverse of FDO).
I was planning to have a separated patch for AutoFDO, but I can also include it in this patch.

A separate patch sounds good, thanks.

AutoFDO is less problematic as branch_weights metadata overwrites VP in the end. (and many VP metadata is already consumed by the prompting in preinliner).

In FDO, VP metadata overwrites branch_weights -- This is bad as we lost more information in BPI/BFI.
(branch_weights is more important for downstream opt).

It looks like same happens to autofdo, i.e, no branch_weights assigned to indirect invokes? https://github.com/llvm/llvm-project/blob/7b81a81d5f9cccb1b091cfc5264bc483b0acc83a/llvm/lib/Transforms/IPO/SampleProfile.cpp#L1628

Oh, I guess branch_weight is attached to invokes that are promoted from an indirect invoke, so it is better handled than the FDO case.

the branch weight for invoke is set in line 1724 (https://github.com/llvm/llvm-project/blob/7b81a81d5f9cccb1b091cfc5264bc483b0acc83a/llvm/lib/Transforms/IPO/SampleProfile.cpp#L1724)

line 1628 sets the call's branch-weight (which is a single number, call counts) -- this should be overwritten by line 1724.

[hoy]

In D133121#3775043, @xur wrote:

In D133121#3774815, @hoy wrote:

In D133121#3774799, @xur wrote:

In D133121#3774764, @hoy wrote:

In D133121#3774657, @xur wrote:

In D133121#3772956, @hoy wrote:

Thanks for the change. Wondering if you saw performance improvement with it.

The corresponding change on the sample profile side isn't included, do you have any plan for it?

Yes. For AutoFDO, we still need some change for the invoke in sampleloader
The situation is slightly different in AutoFDO. In FDO, we do branch_weight annotation first, then do the VP.

In AutoFDO, we do VP annotation first and then branch_weight annotation. We need to set the branch_weight metadata where there is already a VP metadata (which is the reverse of FDO).
I was planning to have a separated patch for AutoFDO, but I can also include it in this patch.

A separate patch sounds good, thanks.

AutoFDO is less problematic as branch_weights metadata overwrites VP in the end. (and many VP metadata is already consumed by the prompting in preinliner).

In FDO, VP metadata overwrites branch_weights -- This is bad as we lost more information in BPI/BFI.
(branch_weights is more important for downstream opt).

It looks like same happens to autofdo, i.e, no branch_weights assigned to indirect invokes? https://github.com/llvm/llvm-project/blob/7b81a81d5f9cccb1b091cfc5264bc483b0acc83a/llvm/lib/Transforms/IPO/SampleProfile.cpp#L1628

Oh, I guess branch_weight is attached to invokes that are promoted from an indirect invoke, so it is better handled than the FDO case.

the branch weight for invoke is set in line 1724 (https://github.com/llvm/llvm-project/blob/7b81a81d5f9cccb1b091cfc5264bc483b0acc83a/llvm/lib/Transforms/IPO/SampleProfile.cpp#L1724)

line 1628 sets the call's branch-weight (which is a single number, call counts) -- this should be overwritten by line 1724.

Line 1724 is guarded by !TI->extractProfTotalWeight(TempWeight) which is non-zero for invokes with !VP?

[xur]

In D133121#3775072, @hoy wrote:

In D133121#3775043, @xur wrote:

In D133121#3774815, @hoy wrote:

In D133121#3774799, @xur wrote:

In D133121#3774764, @hoy wrote:

In D133121#3774657, @xur wrote:

In D133121#3772956, @hoy wrote:

Thanks for the change. Wondering if you saw performance improvement with it.

The corresponding change on the sample profile side isn't included, do you have any plan for it?

Yes. For AutoFDO, we still need some change for the invoke in sampleloader
The situation is slightly different in AutoFDO. In FDO, we do branch_weight annotation first, then do the VP.

In AutoFDO, we do VP annotation first and then branch_weight annotation. We need to set the branch_weight metadata where there is already a VP metadata (which is the reverse of FDO).
I was planning to have a separated patch for AutoFDO, but I can also include it in this patch.

A separate patch sounds good, thanks.

AutoFDO is less problematic as branch_weights metadata overwrites VP in the end. (and many VP metadata is already consumed by the prompting in preinliner).

In FDO, VP metadata overwrites branch_weights -- This is bad as we lost more information in BPI/BFI.
(branch_weights is more important for downstream opt).

It looks like same happens to autofdo, i.e, no branch_weights assigned to indirect invokes? https://github.com/llvm/llvm-project/blob/7b81a81d5f9cccb1b091cfc5264bc483b0acc83a/llvm/lib/Transforms/IPO/SampleProfile.cpp#L1628

Oh, I guess branch_weight is attached to invokes that are promoted from an indirect invoke, so it is better handled than the FDO case.

the branch weight for invoke is set in line 1724 (https://github.com/llvm/llvm-project/blob/7b81a81d5f9cccb1b091cfc5264bc483b0acc83a/llvm/lib/Transforms/IPO/SampleProfile.cpp#L1724)

line 1628 sets the call's branch-weight (which is a single number, call counts) -- this should be overwritten by line 1724.

Line 1724 is guarded by !TI->extractProfTotalWeight(TempWeight) which is non-zero for invokes with !VP?

Yes. This is the place that need to be changed.
Here TI is the invoke. It can have either call-count branch weight metadata (set in 1628), or VP metadata set in line 878.

We need to set the right branch_weight here -- if there is a VP metadada, we create a compound one. If this is single callcount weight, we replace it with real branch_weight.

[mingmingl]

Ohhh no!

I just want to do suggested edit for a comment (around IndirectCallPromotion.cpp line 316) but obviously this is not the correct way to do it.. I suppose there is a way you could get it back?

[mingmingl]

In D133121#4285524, @mingmingl wrote:

Ohhh no!

I just want to do suggested edit for a comment (around IndirectCallPromotion.cpp line 316) but obviously this is not the correct way to do it.. I suppose there is a way you could get it back?

Looks like you could commandeer it back (https://stackoverflow.com/questions/28296729/how-to-undo-commandeer-revision-in-phabricator) My apologies about it :(

llvm/include/llvm/IR/Instructions.h

public:
   // Methods for support type inquiry through isa, cast, and dyn_cast:
   static bool classof(const Instruction *I) {
     return (I->getOpcode() == Instruction::Invoke);
   }
   static bool classof(const Value *V) {
     return isa<Instruction>(V) && classof(cast<Instruction>(V));
   }
 
+  // Find the "branch_weights" metadata and indirect-call metadata.
+  // Return true if either metadata is found.
+  // Return false if none of the metadtata is found.
+  // This method sets the reference argument to nullptr if the metadata
+  // is not found.
+  // TODO: make this function more generic to return a hash_map of
+  // the metdata and its key.
+  static bool getProfMetaData(MDNode *MD, MDNode *&BranchWeightMD,
+                              MDNode *&IndirectCallMD) {
+    BranchWeightMD = IndirectCallMD = nullptr;
+
+    auto isMDwithName = [](MDNode *MD, const char *Str) {
+      auto *MDName = dyn_cast<MDString>(MD->getOperand(0));
+      if (!MDName)
+        return false;
+      return (MDName->getString().equals(Str));
+    };
+
+    if (isMDwithName(MD, "branch_weights")) {
+      BranchWeightMD = MD;
+      return true;
+    }
+    if (isMDwithName(MD, "VP")) {
+      IndirectCallMD = MD;
+      return true;
+    }
+    if (MD->getNumOperands() != 3)
+      return false;
+    for (int I = 1; I < 3; I++) {
+      auto *MDT = cast<MDNode>(MD->getOperand(I));
+      if (!MDT)
+        continue;
+      if (isMDwithName(MDT, "branch_weights"))
+        BranchWeightMD = MDT;
+      else if (isMDwithName(MDT, "VP"))
+        IndirectCallMD = MDT;
+    }
+    return (BranchWeightMD || IndirectCallMD);
+  }
+
+  // A utility method that return "branch_weights" metadata. The return
+  // value will be nullptr if the metadata is not found.
+  static MDNode *getBranchWeightProfData(MDNode *MD) {
+    MDNode *BranchWeightMD, *IndirectCallMD;
+    getProfMetaData(MD, BranchWeightMD, IndirectCallMD);
+    return BranchWeightMD;
+  }
+
+  // A utility method that return indirect-call metadata. The return
+  // value will be nullptr if the metadata is not found.
+  static MDNode *getIndirectCallProfData(MDNode *MD) {
+    MDNode *BranchWeightMD, *IndirectCallMD;
+    getProfMetaData(MD, BranchWeightMD, IndirectCallMD);
+    return IndirectCallMD;
+  }
+
 private:
   // Shadow Instruction::setInstructionSubclassData with a private forwarding
   // method so that subclasses cannot accidentally use it.
   template <typename Bitfield>
   void setSubclassData(typename Bitfield::Type Value) {
     Instruction::setSubclassData<Bitfield>(Value);
   }
 };

llvm/lib/Analysis/BranchProbabilityInfo.cpp

bool BranchProbabilityInfo::calcMetadataWeights(const BasicBlock *BB) {
   // Build up the final weights that will be used in a temporary buffer.
   // Compute the sum of all weights to later decide whether they need to
   // be scaled to fit in 32 bits.
   uint64_t WeightSum = 0;
   SmallVector<uint32_t, 2> Weights;
   SmallVector<unsigned, 2> UnreachableIdxs;
   SmallVector<unsigned, 2> ReachableIdxs;
 
-  extractBranchWeights(*TI, Weights);
+  if (!extractBranchWeights(*TI, Weights))
+    return false;
+
   for (unsigned I = 0, E = Weights.size(); I != E; ++I) {
     WeightSum += Weights[I];
     const LoopBlock SrcLoopBB = getLoopBlock(BB);
     const LoopBlock DstLoopBB = getLoopBlock(TI->getSuccessor(I));
     auto EstimatedWeight = getEstimatedEdgeWeight({SrcLoopBB, DstLoopBB});
     if (EstimatedWeight &&
         *EstimatedWeight <= static_cast<uint32_t>(BlockExecWeight::UNREACHABLE))
       UnreachableIdxs.push_back(I);

llvm/lib/IR/Metadata.cpp

 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");
 
-  return ::extractProfTotalWeight(getMetadata(LLVMContext::MD_prof), TotalVal);
+  TotalVal = 0;
+  auto *ProfileData = getMetadata(LLVMContext::MD_prof);
+  if (!ProfileData)
+    return false;
+  if (getOpcode() == Instruction::Invoke) {
+    ProfileData = InvokeInst::getBranchWeightProfData(ProfileData);
+    if (!ProfileData)
+      ProfileData = InvokeInst::getIndirectCallProfData(ProfileData);
+  }
+  return ::extractProfTotalWeight(ProfileData, TotalVal);
 }
 
 void GlobalObject::copyMetadata(const GlobalObject *Other, unsigned Offset) {
   SmallVector<std::pair<unsigned, MDNode *>, 8> MDs;
   Other->getAllMetadata(MDs);
   for (auto &MD : MDs) {
     // We need to adjust the type metadata offset.
     if (Offset != 0 && MD.first == LLVMContext::MD_type) {

llvm/lib/IR/ProfDataUtils.cpp

bool extractBranchWeights(const MDNode *ProfileData,
   if (!isBranchWeightMD(ProfileData))
     return false;
   return extractWeights(ProfileData, Weights);
 }
 
 bool extractBranchWeights(const Instruction &I,
                           SmallVectorImpl<uint32_t> &Weights) {
   auto *ProfileData = I.getMetadata(LLVMContext::MD_prof);
+  if (!ProfileData)
+    return false;
+  if (I.getOpcode() == Instruction::Invoke)
+    ProfileData = InvokeInst::getBranchWeightProfData(ProfileData);
   return extractBranchWeights(ProfileData, Weights);
 }
 
 bool extractBranchWeights(const Instruction &I, uint64_t &TrueVal,
                           uint64_t &FalseVal) {
   assert((I.getOpcode() == Instruction::Br ||
-          I.getOpcode() == Instruction::Select) &&
+          I.getOpcode() == Instruction::Select ||
+          I.getOpcode() == Instruction::Invoke) &&
          "Looking for branch weights on something besides branch or select");
 
   SmallVector<uint32_t, 2> Weights;
   auto *ProfileData = I.getMetadata(LLVMContext::MD_prof);
+  if (!ProfileData)
+    return false;
+
+  if (I.getOpcode() == Instruction::Invoke)
+    ProfileData = InvokeInst::getBranchWeightProfData(ProfileData);
+
   if (!extractBranchWeights(ProfileData, Weights))
     return false;
 
   if (Weights.size() > 2)
     return false;
 
   TrueVal = Weights[0];
   FalseVal = Weights[1];

llvm/lib/IR/Verifier.cpp

void Verifier::visitDereferenceableMetadata(Instruction& I, MDNode* MD) {
   ConstantInt *CI = mdconst::dyn_extract<ConstantInt>(MD->getOperand(0));
   Check(CI && CI->getType()->isIntegerTy(64),
         "dereferenceable, "
         "dereferenceable_or_null metadata value must be an i64!",
         &I);
 }
 
 void Verifier::visitProfMetadata(Instruction &I, MDNode *MD) {
+  auto checkInst = [this](Instruction &I, MDNode *MD) {
     Check(MD->getNumOperands() >= 2,
           "!prof annotations should have no less than 2 operands", MD);
 
     // Check first operand.
     Check(MD->getOperand(0) != nullptr, "first operand should not be null", MD);
     Check(isa<MDString>(MD->getOperand(0)),
           "expected string with name of the !prof annotation", MD);
     MDString *MDS = cast<MDString>(MD->getOperand(0));
     StringRef ProfName = MDS->getString();
 
     // Check consistency of !prof branch_weights metadata.
     if (ProfName.equals("branch_weights")) {
       if (isa<InvokeInst>(&I)) {
         Check(MD->getNumOperands() == 2 || MD->getNumOperands() == 3,
               "Wrong number of InvokeInst branch_weights operands", MD);
       } else {
         unsigned ExpectedNumOperands = 0;
         if (BranchInst *BI = dyn_cast<BranchInst>(&I))
           ExpectedNumOperands = BI->getNumSuccessors();
         else if (SwitchInst *SI = dyn_cast<SwitchInst>(&I))
           ExpectedNumOperands = SI->getNumSuccessors();
         else if (isa<CallInst>(&I))
           ExpectedNumOperands = 1;
         else if (IndirectBrInst *IBI = dyn_cast<IndirectBrInst>(&I))
           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);
+          CheckFailed(
+              "!prof branch_weights are not allowed for this instruction", MD);
 
         Check(MD->getNumOperands() == 1 + ExpectedNumOperands,
               "Wrong number of operands", MD);
       }
       for (unsigned i = 1; i < MD->getNumOperands(); ++i) {
         auto &MDO = MD->getOperand(i);
         Check(MDO, "second operand should not be null", MD);
         Check(mdconst::dyn_extract<ConstantInt>(MDO),
               "!prof brunch_weights operand is not a const int");
       }
     }
+  };
+
+  // Check MD for non-invoke instructions or InvokeInst only with branch_weight.
+  if (!isa<InvokeInst>(&I) || isa<MDString>(MD->getOperand(0))) {
+    checkInst(I, MD);
+    return;
+  }
+
+  // For an InvokeInst that has both branch_weight and VP profiles.
+  checkInst(I, cast<MDNode>(MD->getOperand(1)));
+  checkInst(I, cast<MDNode>(MD->getOperand(2)));
 }
 
 void Verifier::visitCallStackMetadata(MDNode *MD) {
   // Call stack metadata should consist of a list of at least 1 constant int
   // (representing a hash of the location).
   Check(MD->getNumOperands() >= 1,
         "call stack metadata should have at least 1 operand", MD);
 

llvm/lib/ProfileData/InstrProf.cpp

 #include "llvm/ADT/Triple.h"
 #include "llvm/Config/config.h"
 #include "llvm/IR/Constant.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/GlobalValue.h"
 #include "llvm/IR/GlobalVariable.h"
 #include "llvm/IR/Instruction.h"
+#include "llvm/IR/Instructions.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/MDBuilder.h"
 #include "llvm/IR/Metadata.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/Type.h"
 #include "llvm/ProfileData/InstrProfReader.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/CommandLine.h"
void annotateValueSite(Module &M, Instruction &Inst,
   for (auto &VD : VDs) {
     Vals.push_back(MDHelper.createConstant(
         ConstantInt::get(Type::getInt64Ty(Ctx), VD.Value)));
     Vals.push_back(MDHelper.createConstant(
         ConstantInt::get(Type::getInt64Ty(Ctx), VD.Count)));
     if (--MDCount == 0)
       break;
   }
+  if (!isa<InvokeInst>(&Inst)) {
     Inst.setMetadata(LLVMContext::MD_prof, MDNode::get(Ctx, Vals));
+    return;
+  }
+
+  // For an invoke instruction, if we have VP counts, we should already have a
+  // "branch_weights" MD_prof.
+  MDNode *BranchWeight = Inst.getMetadata(LLVMContext::MD_prof);
+  assert(BranchWeight != nullptr &&
+         dyn_cast<MDString>(BranchWeight->getOperand(0))
+             ->getString()
+             .equals("branch_weights") &&
+         "Expect non-null branch_weight metadata");
+  SmallVector<Metadata *, 3> InvokeVals;
+  InvokeVals.push_back(nullptr);
+  InvokeVals.push_back(BranchWeight);
+  InvokeVals.push_back(MDNode::get(Ctx, Vals));
+  MDNode *NewM = MDNode::getDistinct(Ctx, InvokeVals);
+  NewM->replaceOperandWith(0, NewM);
+  Inst.setMetadata(LLVMContext::MD_prof, NewM);
 }
 
 bool getValueProfDataFromInst(const Instruction &Inst,
                               InstrProfValueKind ValueKind,
                               uint32_t MaxNumValueData,
                               InstrProfValueData ValueData[],
                               uint32_t &ActualNumValueData, uint64_t &TotalC,
                               bool GetNoICPValue) {
   MDNode *MD = Inst.getMetadata(LLVMContext::MD_prof);
   if (!MD)
     return false;
 
+  if (isa<InvokeInst>(&Inst)) {
+    MD = cast<InvokeInst>(&Inst)->getIndirectCallProfData(MD);
+    if (!MD)
+      return false;
+  }
+
   unsigned NOps = MD->getNumOperands();
 
   if (NOps < 5)
     return false;
 
   // Operand 0 is a string tag "VP":
   MDString *Tag = cast<MDString>(MD->getOperand(0));
   if (!Tag)

llvm/lib/Transforms/Instrumentation/IndirectCallPromotion.cpp

if (!NumCandidates ||
       continue;
     auto PromotionCandidates = getPromotionCandidatesForCallSite(
         *CB, ICallProfDataRef, TotalCount, NumCandidates);
     uint32_t NumPromoted = tryToPromote(*CB, PromotionCandidates, TotalCount);
     if (NumPromoted == 0)
       continue;
 
     Changed = true;
-    // Adjust the MD.prof metadata. First delete the old one.
+    // Adjust the MD.prof metadata. First delete the old one. xur
+    if (isa<InvokeInst>(CB)) {
+      MDNode *MD = CB->getMetadata(LLVMContext::MD_prof);
+      MD = InvokeInst::getBranchWeightProfData(MD);
+      CB->setMetadata(LLVMContext::MD_prof, MD);
+    } else
       CB->setMetadata(LLVMContext::MD_prof, nullptr);
     // If all promoted, we don't need the MD.prof metadata.
     if (TotalCount == 0 || NumPromoted == NumVals)
       continue;
     // Otherwise we need update with the un-promoted records back.
     annotateValueSite(*M, *CB, ICallProfDataRef.slice(NumPromoted), TotalCount,
                       IPVK_IndirectCallTarget, NumCandidates);
   }
   return Changed;

llvm/lib/Transforms/Utils/Local.cpp

CallInst *llvm::createCallMatchingInvoke(InvokeInst *II) {
   II->getOperandBundlesAsDefs(OpBundles);
   CallInst *NewCall = CallInst::Create(II->getFunctionType(),
                                        II->getCalledOperand(), Args, OpBundles);
   NewCall->setCallingConv(II->getCallingConv());
   NewCall->setAttributes(II->getAttributes());
   NewCall->setDebugLoc(II->getDebugLoc());
   NewCall->copyMetadata(*II);
 
-  // If the invoke had profile metadata, try converting them for CallInst.
+  // Move MD_prof data to the new instruction. We can only have one.
+  // First to keep indirect-call profile metadata.
+  // Otherwise, if try converting them for total call profile data.
+  auto *ProfileData = II->getMetadata(LLVMContext::MD_prof);
+  if (!ProfileData)
+    return NewCall;
+  ProfileData = InvokeInst::getIndirectCallProfData(ProfileData);
+  if (!ProfileData) {
     uint64_t TotalWeight;
-  if (NewCall->extractProfTotalWeight(TotalWeight)) {
+    if (II->extractProfTotalWeight(TotalWeight)) {
       // Set the total weight if it fits into i32, otherwise reset.
       MDBuilder MDB(NewCall->getContext());
-    auto NewWeights = uint32_t(TotalWeight) != TotalWeight
-                          ? nullptr
-                          : MDB.createBranchWeights({uint32_t(TotalWeight)});
-    NewCall->setMetadata(LLVMContext::MD_prof, NewWeights);
+      ProfileData = uint32_t(TotalWeight) != TotalWeight
     }
-
+  }
+  // Set the MD_prof metadata.
+  NewCall->setMetadata(LLVMContext::MD_prof, ProfileData);
   return NewCall;
 }
 
 // changeToCall - Convert the specified invoke into a normal call.
 CallInst *llvm::changeToCall(InvokeInst *II, DomTreeUpdater *DTU) {
   CallInst *NewCall = createCallMatchingInvoke(II);
   NewCall->takeName(II);
   NewCall->insertBefore(II);

llvm/test/Transforms/PGOProfile/Inputs/invoke_icall.proftext

+# IR level Instrumentation Flag
+:ir
+_Z3bari
+# Func Hash:
+784007058953177093
+# Num Counters:
+2
+# Counter Values:
+10000
+7500
+
+_Z3fooi
+# Func Hash:
+264495437628610820
+# Num Counters:
+4
+# Counter Values:
+10000
+0
+7500
+0
+# Num Value Kinds:
+1
+# ValueKind = IPVK_IndirectCallTarget:
+0
+# NumValueSites:
+1
+1
+_Z3bari:10000
+

llvm/test/Transforms/PGOProfile/invoke_icall.ll

+; RUN: llvm-profdata merge %S/Inputs/invoke_icall.proftext -o %t.profdata
+; RUN: opt < %s -passes=pgo-instr-use -pgo-instrument-entry=false -pgo-test-profile-file=%t.profdata -S | FileCheck %s
+
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+%"class.std::basic_ostream" = type { i32 (...)**, %"class.std::basic_ios" }
+%"class.std::basic_ios" = type { %"class.std::ios_base", %"class.std::basic_ostream"*, i8, i8 }
+%"class.std::ios_base" = type { i32 (...)**, i64, i64, i32, i32, i32, %"struct.std::ios_base::_Callback_list"* }
+%"struct.std::ios_base::_Callback_list" = type { %"struct.std::ios_base::_Callback_list"*, void (i32, %"class.std::ios_base"*, i32)*, i32, i32 }
+
+@_ZTIi = external dso_local constant i8*
+@goo = dso_local global i32 (i32)* null, align 8
+@_ZSt4cerr = external dso_local global %"class.std::basic_ostream", align 8
+@.str = private unnamed_addr constant [19 x i8] c"got an exception.\0A\00", align 1
+
+define dso_local i32 @_Z3bari(i32 %a) {
+entry:
+  %0 = and i32 %a, 3
+  %tobool.not = icmp eq i32 %0, 0
+  br i1 %tobool.not, label %if.end, label %if.then
+
+if.then:
+  %exception = call i8* @__cxa_allocate_exception(i64 4)
+  %1 = bitcast i8* %exception to i32*
+  store i32 20, i32* %1, align 16
+  call void @__cxa_throw(i8* %exception, i8* bitcast (i8** @_ZTIi to i8*), i8* null)
+  unreachable
+
+if.end:
+  %mul = mul nsw i32 %a, %a
+  ret i32 %mul
+}
+
+declare dso_local i8* @__cxa_allocate_exception(i64)
+
+declare dso_local void @__cxa_throw(i8*, i8*, i8*)
+
+define dso_local i32 @_Z3fooi(i32 %n) personality i8* bitcast (i32 (...)* @__gxx_personality_v0 to i8*) {
+entry:
+  %0 = load i32 (i32)*, i32 (i32)** @goo, align 8
+  %call = invoke i32 %0(i32 %n)
+          to label %cleanup unwind label %lpad
+; CHECK-LABEL: @_Z3fooi
+; CHECK-SAME: !prof ![[BW_ENTRY:[0-9]+]]
+; CHECK: %call = invoke
+; CHECK: to label
+; CHECK-SAME: !prof ![[MD_INVOKE:[0-9]+]]
+; CHECK-DAG: ![[BW_ENTRY]] = !{!"function_entry_count", i64 10000}
+; CHECK-DAG: ![[BW_INVOKE:[0-9]+]] = !{!"branch_weights", i32 2500, i32 7500}
+; CHECK-DAG: ![[MD_INVOKE]] = distinct !{![[MD_INVOKE]], ![[BW_INVOKE]], ![[VP_INVOKE:[0-9]+]]}
+; CHECK-DAG: ![[VP_INVOKE]] = !{!"VP", i32 0, i64 10000, i64 3727899762981752933, i64 10000}
+
+lpad:
+  %1 = landingpad { i8*, i32 }
+          cleanup
+          catch i8* bitcast (i8** @_ZTIi to i8*)
+  %2 = extractvalue { i8*, i32 } %1, 1
+  %3 = call i32 @llvm.eh.typeid.for(i8* bitcast (i8** @_ZTIi to i8*))
+  %matches = icmp eq i32 %2, %3
+  br i1 %matches, label %catch, label %ehcleanup
+
+catch:
+  %4 = extractvalue { i8*, i32 } %1, 0
+  %5 = call i8* @__cxa_begin_catch(i8* %4)
+  %call1.i12 = invoke nonnull align 8 dereferenceable(8) %"class.std::basic_ostream"* @_ZSt16__ostream_insertIcSt11char_traitsIcEERSt13basic_ostreamIT_T0_ES6_PKS3_l(%"class.std::basic_ostream"* nonnull align 8 dereferenceable(8) @_ZSt4cerr, i8* nonnull getelementptr inbounds ([19 x i8], [19 x i8]* @.str, i64 0, i64 0), i64 18)
+          to label %invoke.cont2 unwind label %lpad1
+
+invoke.cont2:
+  call void @__cxa_end_catch()
+  br label %cleanup
+
+lpad1:
+  %6 = landingpad { i8*, i32 }
+          cleanup
+  call void @__cxa_end_catch()
+  br label %ehcleanup
+
+cleanup:
+  %retval.0 = phi i32 [ -1, %invoke.cont2 ], [ %call, %entry ]
+  ret i32 %retval.0
+
+ehcleanup:
+  %lpad.val6.merged = phi { i8*, i32 } [ %6, %lpad1 ], [ %1, %lpad ]
+  resume { i8*, i32 } %lpad.val6.merged
+}
+
+declare dso_local i32 @__gxx_personality_v0(...)
+
+declare i32 @llvm.eh.typeid.for(i8*)
+
+declare dso_local i8* @__cxa_begin_catch(i8*)
+
+declare dso_local void @__cxa_end_catch()
+
+declare dso_local nonnull align 8 dereferenceable(8) %"class.std::basic_ostream"* @_ZSt16__ostream_insertIcSt11char_traitsIcEERSt13basic_ostreamIT_T0_ES6_PKS3_l(%"class.std::basic_ostream"* nonnull align 8 dereferenceable(8), i8*, i64)
+