1 files changed, 26 insertions, 10 deletions

diff --git a/llvm/lib/ProfileData/SampleProf.cpp b/llvm/lib/ProfileData/SampleProf.cpp
index 294f646..addb473 100644
--- a/llvm/lib/ProfileData/SampleProf.cpp
+++ b/llvm/lib/ProfileData/SampleProf.cpp
@@ -236,7 +236,9 @@ LineLocation FunctionSamples::getCallSiteIdentifier(const DILocation *DIL,
 }
 
 const FunctionSamples *FunctionSamples::findFunctionSamples(
-    const DILocation *DIL, SampleProfileReaderItaniumRemapper *Remapper) const {
+    const DILocation *DIL, SampleProfileReaderItaniumRemapper *Remapper,
+    const HashKeyMap<std::unordered_map, FunctionId, FunctionId>
+        *FuncNameToProfNameMap) const {
   assert(DIL);
   SmallVector<std::pair<LineLocation, StringRef>, 10> S;
 
@@ -256,7 +258,8 @@ const FunctionSamples *FunctionSamples::findFunctionSamples(
     return this;
   const FunctionSamples *FS = this;
   for (int i = S.size() - 1; i >= 0 && FS != nullptr; i--) {
-    FS = FS->findFunctionSamplesAt(S[i].first, S[i].second, Remapper);
+    FS = FS->findFunctionSamplesAt(S[i].first, S[i].second, Remapper,
+                                   FuncNameToProfNameMap);
   }
   return FS;
 }
@@ -277,19 +280,32 @@ void FunctionSamples::findAllNames(DenseSet<FunctionId> &NameSet) const {
 
 const FunctionSamples *FunctionSamples::findFunctionSamplesAt(
     const LineLocation &Loc, StringRef CalleeName,
-    SampleProfileReaderItaniumRemapper *Remapper) const {
+    SampleProfileReaderItaniumRemapper *Remapper,
+    const HashKeyMap<std::unordered_map, FunctionId, FunctionId>
+        *FuncNameToProfNameMap) const {
   CalleeName = getCanonicalFnName(CalleeName);
 
-  auto iter = CallsiteSamples.find(mapIRLocToProfileLoc(Loc));
-  if (iter == CallsiteSamples.end())
+  auto I = CallsiteSamples.find(mapIRLocToProfileLoc(Loc));
+  if (I == CallsiteSamples.end())
     return nullptr;
-  auto FS = iter->second.find(getRepInFormat(CalleeName));
-  if (FS != iter->second.end())
+  auto FS = I->second.find(getRepInFormat(CalleeName));
+  if (FS != I->second.end())
     return &FS->second;
+
+  if (FuncNameToProfNameMap && !FuncNameToProfNameMap->empty()) {
+    auto R = FuncNameToProfNameMap->find(FunctionId(CalleeName));
+    if (R != FuncNameToProfNameMap->end()) {
+      CalleeName = R->second.stringRef();
+      auto FS = I->second.find(getRepInFormat(CalleeName));
+      if (FS != I->second.end())
+        return &FS->second;
+    }
+  }
+
   if (Remapper) {
     if (auto NameInProfile = Remapper->lookUpNameInProfile(CalleeName)) {
-      auto FS = iter->second.find(getRepInFormat(*NameInProfile));
-      if (FS != iter->second.end())
+      auto FS = I->second.find(getRepInFormat(*NameInProfile));
+      if (FS != I->second.end())
         return &FS->second;
     }
   }
@@ -300,7 +316,7 @@ const FunctionSamples *FunctionSamples::findFunctionSamplesAt(
     return nullptr;
   uint64_t MaxTotalSamples = 0;
   const FunctionSamples *R = nullptr;
-  for (const auto &NameFS : iter->second)
+  for (const auto &NameFS : I->second)
     if (NameFS.second.getTotalSamples() >= MaxTotalSamples) {
       MaxTotalSamples = NameFS.second.getTotalSamples();
       R = &NameFS.second;