[Inline][PGO] After inline, update InvokeInst profile counts in caller and cloned callee

llvm/include/llvm/IR/Instructions.h

@@ -4370,6 +4370,9 @@ class InvokeInst : public CallBase {
 
   unsigned getNumSuccessors() const { return 2; }
 
+  /// Updates profile metadata by scaling it by \p S / \p T.
+  void updateProfWeight(uint64_t S, uint64_t T);
+
   // Methods for support type inquiry through isa, cast, and dyn_cast:
   static bool classof(const Instruction *I) {
     return (I->getOpcode() == Instruction::Invoke);

llvm/lib/IR/Instructions.cpp

@@ -927,6 +927,18 @@ LandingPadInst *InvokeInst::getLandingPadInst() const {
   return cast<LandingPadInst>(getUnwindDest()->getFirstNonPHI());
 }
 
+void InvokeInst::updateProfWeight(uint64_t S, uint64_t T) {
+  if (T == 0) {
+    LLVM_DEBUG(dbgs() << "Attempting to update profile weights will result in "
+                         "div by 0. Ignoring. Likely the function "
+                      << getParent()->getParent()->getName()
+                      << " has 0 entry count, and contains call instructions "
+                         "with non-zero prof info.");
+    return;
+  }
+  scaleProfData(*this, S, T);
+}
+
 //===----------------------------------------------------------------------===//
 //                        CallBrInst Implementation
 //===----------------------------------------------------------------------===//

llvm/lib/IR/ProfDataUtils.cpp

@@ -46,6 +46,9 @@ constexpr unsigned WeightsIdx = 1;
 // the minimum number of operands for MD_prof nodes with branch weights
 constexpr unsigned MinBWOps = 3;
 
+// the minimum number of operands for MD_prof nodes with value profiles
+constexpr unsigned MinVPOps = 5;
+
 // We may want to add support for other MD_prof types, so provide an abstraction
 // for checking the metadata type.
 bool isTargetMD(const MDNode *ProfData, const char *Name, unsigned MinOps) {
@@ -97,11 +100,25 @@ bool isBranchWeightMD(const MDNode *ProfileData) {
   return isTargetMD(ProfileData, "branch_weights", MinBWOps);
 }
 
+bool isValueProfileMD(const MDNode *ProfileData) {
+  return isTargetMD(ProfileData, "VP", MinVPOps);
+}
+
 bool hasBranchWeightMD(const Instruction &I) {
   auto *ProfileData = I.getMetadata(LLVMContext::MD_prof);
   return isBranchWeightMD(ProfileData);
 }
 
+bool hasCountTypeMD(const Instruction &I) {
+  auto *ProfileData = I.getMetadata(LLVMContext::MD_prof);
+  // Value profiles record count-type information.
+  if (isValueProfileMD(ProfileData))
+    return true;
+  // Conservatively assume non CallBase instruction only get taken/not-taken
+  // branch probability, so not interpret them as count.
+  return isa<CallBase>(I) && !isBranchWeightMD(ProfileData);
+}
+
 bool hasValidBranchWeightMD(const Instruction &I) {
   return getValidBranchWeightMDNode(I);
 }
@@ -212,6 +229,9 @@ void scaleProfData(Instruction &I, uint64_t S, uint64_t T) {
                         ProfDataName->getString() != "VP"))
     return;
 
+  if (!hasCountTypeMD(I))
+    return;
+
   LLVMContext &C = I.getContext();
 
   MDBuilder MDB(C);

llvm/lib/Transforms/Utils/InlineFunction.cpp

@@ -1982,10 +1982,14 @@ void llvm::updateProfileCallee(
   // During inlining ?
   if (VMap) {
     uint64_t CloneEntryCount = PriorEntryCount - NewEntryCount;
-    for (auto Entry : *VMap)
+    for (auto Entry : *VMap) {
       if (isa<CallInst>(Entry.first))
         if (auto *CI = dyn_cast_or_null<CallInst>(Entry.second))
           CI->updateProfWeight(CloneEntryCount, PriorEntryCount);
+      if (isa<InvokeInst>(Entry.first))
+        if (auto *II = dyn_cast_or_null<InvokeInst>(Entry.second))
+          II->updateProfWeight(CloneEntryCount, PriorEntryCount);
+    }
   }
 
   if (EntryDelta) {
@@ -1994,9 +1998,12 @@ void llvm::updateProfileCallee(
     for (BasicBlock &BB : *Callee)
       // No need to update the callsite if it is pruned during inlining.
       if (!VMap || VMap->count(&BB))
-        for (Instruction &I : BB)
+        for (Instruction &I : BB) {
           if (CallInst *CI = dyn_cast<CallInst>(&I))
             CI->updateProfWeight(NewEntryCount, PriorEntryCount);
+          if (InvokeInst *II = dyn_cast<InvokeInst>(&I))
+            II->updateProfWeight(NewEntryCount, PriorEntryCount);
+        }
   }
 }
 

llvm/test/Transforms/Inline/update_invoke_prof.ll

@@ -1,22 +1,31 @@
-; A pre-commit test to show that branch weights and value profiles associated with invoke are not updated.
+; Test that branch weights and value profiles associated with invoke are updated
+; in both caller and callee after inline, but invoke instructions with taken or
+; not taken branch probabilities are not updated.
 ; RUN: opt < %s -passes='require<profile-summary>,cgscc(inline)' -S | FileCheck %s
 
 declare i32 @__gxx_personality_v0(...)
 
 define void @caller(ptr %func) personality ptr @__gxx_personality_v0 !prof !15 {
   call void @callee(ptr %func), !prof !16
+
   ret void
 }
 
-declare void @inner_callee(ptr %func)
+declare void @callee1(ptr %func)
+
+declare void @callee2(ptr %func)
 
 define void @callee(ptr %func) personality ptr @__gxx_personality_v0 !prof !17 {
   invoke void %func()
           to label %next unwind label %lpad, !prof !18
 
 next:
-  invoke void @inner_callee(ptr %func)
-          to label %ret unwind label %lpad, !prof !19
+  invoke void @callee1(ptr %func)
+          to label %cont unwind label %lpad, !prof !19
+
+cont:
+  invoke void @callee2(ptr %func)
+          to label %ret unwind label %lpad, !prof !20
 
 lpad:
   %exn = landingpad {ptr, i32}
@@ -47,18 +56,27 @@ ret:
 !17 = !{!"function_entry_count", i32 1500}
 !18 = !{!"VP", i32 0, i64 1500, i64 123, i64 900, i64 456, i64 600}
 !19 = !{!"branch_weights", i32 1500}
+!20 = !{!"branch_weights", i32 1234, i32 5678}
 
 ; CHECK-LABEL: @caller(
 ; CHECK:  invoke void %func(
 ; CHECK-NEXT: {{.*}} !prof ![[PROF1:[0-9]+]]
-; CHECK:  invoke void @inner_callee(
+; CHECK:  invoke void @callee1(
 ; CHECK-NEXT: {{.*}} !prof ![[PROF2:[0-9]+]]
+; CHECK:  invoke void @callee2(
+; CHECK-NEXT: {{.*}} !prof ![[PROF3:[0-9]+]]
 
 ; CHECK-LABL: @callee(
 ; CHECK:  invoke void %func(
-; CHECK-NEXT: {{.*}} !prof ![[PROF1]] 
-; CHECK:  invoke void @inner_callee(
-; CHECK-NEXT: {{.*}} !prof ![[PROF2]]
+; CHECK-NEXT: {{.*}} !prof ![[PROF4:[0-9]+]]
+; CHECK:  invoke void @callee1(
+; CHECK-NEXT: {{.*}} !prof ![[PROF5:[0-9]+]]
+; CHECK:  invoke void @callee2(
+; CHECK-NEXT: {{.*}} !prof ![[PROF3]]
+
 
-; CHECK: ![[PROF1]] = !{!"VP", i32 0, i64 1500, i64 123, i64 900, i64 456, i64 600}
-; CHECK: ![[PROF2]] = !{!"branch_weights", i32 1500}
+; CHECK: ![[PROF1]] = !{!"VP", i32 0, i64 1000, i64 123, i64 600, i64 456, i64 400}
+; CHECK: ![[PROF2]] = !{!"branch_weights", i32 1000}
+; CHECK: ![[PROF3]] = !{!"branch_weights", i32 1234, i32 5678}
+; CHECK: ![[PROF4]] = !{!"VP", i32 0, i64 500, i64 123, i64 300, i64 456, i64 200}
+; CHECK: ![[PROF5]] = !{!"branch_weights", i32 500}