1 files changed, 19 insertions, 10 deletions

diff --git a/llvm/lib/Transforms/IPO/SampleProfile.cpp b/llvm/lib/Transforms/IPO/SampleProfile.cpp
index 6df17ec400c..46b9914347b 100644
--- a/llvm/lib/Transforms/IPO/SampleProfile.cpp
+++ b/llvm/lib/Transforms/IPO/SampleProfile.cpp
@@ -171,7 +171,7 @@ protected:
   ErrorOr<uint64_t> getBlockWeight(const BasicBlock *BB);
   const FunctionSamples *findCalleeFunctionSamples(const Instruction &I) const;
   std::vector<const FunctionSamples *>
-  findIndirectCallFunctionSamples(const Instruction &I) const;
+  findIndirectCallFunctionSamples(const Instruction &I, uint64_t &Sum) const;
   const FunctionSamples *findFunctionSamples(const Instruction &I) const;
   bool inlineCallInstruction(Instruction *I);
   bool inlineHotFunctions(Function &F,
@@ -625,10 +625,11 @@ SampleProfileLoader::findCalleeFunctionSamples(const Instruction &Inst) const {
 }
 
 /// Returns a vector of FunctionSamples that are the indirect call targets
-/// of \p Inst. The vector is sorted by the total number of samples.
+/// of \p Inst. The vector is sorted by the total number of samples. Stores
+/// the total call count of the indirect call in \p Sum.
 std::vector<const FunctionSamples *>
 SampleProfileLoader::findIndirectCallFunctionSamples(
-    const Instruction &Inst) const {
+    const Instruction &Inst, uint64_t &Sum) const {
   const DILocation *DIL = Inst.getDebugLoc();
   std::vector<const FunctionSamples *> R;
 
@@ -640,16 +641,25 @@ SampleProfileLoader::findIndirectCallFunctionSamples(
   if (FS == nullptr)
     return R;
 
+  uint32_t LineOffset = getOffset(DIL);
+  uint32_t Discriminator = DIL->getBaseDiscriminator();
+
+  auto T = FS->findCallTargetMapAt(LineOffset, Discriminator);
+  Sum = 0;
+  if (T)
+    for (const auto &T_C : T.get())
+      Sum += T_C.second;
   if (const FunctionSamplesMap *M = FS->findFunctionSamplesMapAt(
           LineLocation(getOffset(DIL), DIL->getBaseDiscriminator()))) {
     if (M->size() == 0)
       return R;
     for (const auto &NameFS : *M) {
+      Sum += NameFS.second.getEntrySamples();
       R.push_back(&NameFS.second);
     }
     std::sort(R.begin(), R.end(),
               [](const FunctionSamples *L, const FunctionSamples *R) {
-                return L->getTotalSamples() > R->getTotalSamples();
+                return L->getEntrySamples() > R->getEntrySamples();
               });
   }
   return R;
@@ -764,7 +774,8 @@ bool SampleProfileLoader::inlineHotFunctions(
       if (CallSite(I).isIndirectCall()) {
         if (PromotedInsns.count(I))
           continue;
-        for (const auto *FS : findIndirectCallFunctionSamples(*I)) {
+        uint64_t Sum;
+        for (const auto *FS : findIndirectCallFunctionSamples(*I, Sum)) {
           if (IsThinLTOPreLink) {
             FS->findImportedFunctions(ImportGUIDs, F.getParent(),
                                       Samples->getTotalSamples() *
@@ -786,12 +797,10 @@ bool SampleProfileLoader::inlineHotFunctions(
               !R->getValue()->isDeclaration() &&
               R->getValue()->getSubprogram() &&
               isLegalToPromote(I, R->getValue(), &Reason)) {
-            // The indirect target was promoted and inlined in the profile,
-            // as a result, we do not have profile info for the branch
-            // probability. We set the probability to 80% taken to indicate
-            // that the static call is likely taken.
+            uint64_t C = FS->getEntrySamples();
             Instruction *DI = promoteIndirectCall(
-                I, R->getValue(), 80, 100, false, ORE);
+                I, R->getValue(), C, Sum, false, ORE);
+            Sum -= C;
             PromotedInsns.insert(I);
             // If profile mismatches, we should not attempt to inline DI.
             if ((isa<CallInst>(DI) || isa<InvokeInst>(DI)) &&