[CSSPGO] Refactoring SampleProfileMatcher::runOnFunction

[wlei]

Summary

• rename IRLocation --> IRAnchors, ProfileLocation --> ProfileAnchors
• reorganize runOnFunction, fact out the finding IR anchors code into findIRAnchors
• introduce a new function findProfileAnchors to populate the profile related anchors, the result is saved into ProfileAnchors, it's later used for both mismatch report and matching, this can avoid to parse the getBodySamples and getCallsiteSamples for multiple times.
• move the MatchedCallsiteLocs stuffs from findIRAnchors to countProfileMismatches so that all the staleness metrics report are computed in one function.
• move all matching related into runStaleProfileMatching, and move all mismatching report into countProfileMismatches

[wlei]

llvm/test/Transforms/SampleProfile/profile-mismatch.ll

9

This is due to I changed to use the totalSamples instead of headerSamples for the mismatched inlined samples, since we switched to use the original nested profile, using totalSamples should makes more sense.

[wenlei]

llvm/lib/Transforms/IPO/SampleProfile.cpp

464

this can actually be a valid c mangle name.. we need something more distinctive.. *UnkownIndirectCallee?

2158

nit: the CalleeName -> CalleeName.

2160

nit: avoid nums of false positive ->

reduce num of false positives

or

avoid false positive

2238

This is passed by value, missing &?

2365–2389

Trying to organize the code a bit and keep high level flow clean, how about this:

// ..
std::map<LineLocation, StringRef> IRAnchors;
findIRAnchors(FS, IRAnchors);

// ..
std::map<LineLocation, StringSet<>> ProfileAnchors;
findProfileAnchors(FS, ProfileLocations);

// Detect profile mismatch for profile staleness metrics report.
if (ReportProfileStaleness || PersistProfileStaleness) {
  countProfileMismatches(FS, IRAnchors, ProfileAnchors, ...)
}

if (IsFuncHashMismatch && SalvageStaleProfile) {
   runStaleProfileMatching(IRAnchors, ProfileAnchors, getIRToProfileLocationMap(F));
}

Accordingly, we would need to move populateProfileCallsites in runStaleProfileMatching. Maybe even inline expand runStaleProfileMatching there without a function call since the function is now simplified? Also move all statistics into countProfileMismatches.

Inside runStaleProfileMatching, we'd also rename ProfileAnchors to MatchedAnchors.

llvm/test/Transforms/SampleProfile/profile-mismatch.ll

27

what causes this change?

[wlei]

addressing feedback

[wenlei]

looks cleaner and more structured now, thanks for update!

llvm/lib/Transforms/IPO/SampleProfile.cpp

2377–2384

nit: move this into runStaleProfileMatching as well?

[wlei]

llvm/lib/Transforms/IPO/SampleProfile.cpp

464

Good catch, changed use "dot" as the separate char, that should not be used as c/c++ function name.

2158

done.

2160

done

2238

Good catch, removed populateProfileCallsites

2365–2389

That would looks much cleaner, thanks for the suggestion! done accordingly.

Only one things is that the name IRAnchors, it is not only the Anchors, we also need the non-call locations whose callee name is empty. but I understand separating the non-call locations and the anchors doesn't look neat. I don't have strong option on this, maybe this is straightforward to understand.

llvm/test/Transforms/SampleProfile/profile-mismatch.ll

27

this is the same reason as the line 9.

"I changed to use the totalSamples instead of headerSamples for the mismatched inlined samples, since we switched to use the original nested profile, using totalSamples should makes more sense."

[wenlei]

llvm/lib/Transforms/IPO/SampleProfile.cpp

2365–2389

Yeah, I think the term Anchor is abstract enough that including non-call is also okay.

[wlei]

move dbg print into runStaleProfileMatching

[wenlei]

lgtm, thanks.

[wlei]

fix typo

[hoy]

llvm/lib/Transforms/IPO/SampleProfile.cpp

2128

extractProbe can also be used to extract a probe from a callsite. Can this code be unified with the callsite handling code below to avoid the overwrite?

[wlei]

llvm/lib/Transforms/IPO/SampleProfile.cpp

2128

Try to understand this, do you mean we can just use the LineLocation(Probe->Id, 0) as the callsite location, it then can replace the LineLocation IRCallsite = FunctionSamples::getCallSiteIdentifier(DLoc);?

If so, that's a good idea for pseudo-probe.

However, currently callsite handling part of the code are still used for AutoFDO, we use the calliste related staleness metrics for AutoFDO. Then we probably still need the callsite handling code.

[hoy]

llvm/lib/Transforms/IPO/SampleProfile.cpp

2128

I think you can just do this

if (FunctionSamples::ProfileIsProbeBased) {
     if (auto Probe = extractProbe(I)) {
         StringRef CalleeName;
         if (const auto *CB = dyn_cast<CallBase>(&I)) {
             if (Function *Callee = I.getCalledFunction()){
                 CalleeName = FunctionSamples::getCanonicalFnName(Callee->getName());
            } else {
               CalleeName = UnknownIndirectCallee;
            }
         }
         IRAnchors.emplace(LineLocation(Probe->Id, 0), CalleeName);
    }
}

Let me know if it works. As for AutoFDO, we could probably leave it a TODO just as what you do now and redesign the whole logic later.

[wlei]

llvm/lib/Transforms/IPO/SampleProfile.cpp

2128

Thanks for the suggestion. That looks cleaner, let me give it a try.

I guess we even don't need the FunctionSamples::ProfileIsProbeBased.

Also do you mind I try it in a new patch, on top of the stack, because this patch is the bottom of the patch stack, then I need to more rebasing and then need to request for review for the accepted patch.

[hoy]

llvm/lib/Transforms/IPO/SampleProfile.cpp

2128

I guess we even don't need the FunctionSamples::ProfileIsProbeBased.

It should still be useful. In theory, we can still use AutoFDO profile for a probed function. Actually our dual-profile mode is just a use case.

Sounds good to make it a new patch.

llvm/lib/Transforms/IPO/SampleProfile.cpp

class SampleProfileMatcher {
   uint64_t NumMismatchedCallsites = 0;
   uint64_t MismatchedCallsiteSamples = 0;
   uint64_t TotalCallsiteSamples = 0;
   uint64_t TotalProfiledFunc = 0;
   uint64_t NumMismatchedFuncHash = 0;
   uint64_t MismatchedFuncHashSamples = 0;
   uint64_t TotalFuncHashSamples = 0;
 
+  // A dummy name for unknown indirect callee, used to differentiate from a
+  // non-call instruction that also has an empty callee name.
+  static constexpr const char *UnknownIndirectCallee =
+      "unknown.indirect.callee";
+
 public:
   SampleProfileMatcher(Module &M, SampleProfileReader &Reader,
                        const PseudoProbeManager *ProbeManager)
       : M(M), Reader(Reader), ProbeManager(ProbeManager) {
     if (FlattenProfileForMatching) {
       ProfileConverter::flattenProfile(Reader.getProfiles(), FlattenedProfiles,
                                        FunctionSamples::ProfileIsCS);
     }
   }
   void runOnModule();
 
 private:
   FunctionSamples *getFlattenedSamplesFor(const Function &F) {
     StringRef CanonFName = FunctionSamples::getCanonicalFnName(F);
     auto It = FlattenedProfiles.find(CanonFName);
     if (It != FlattenedProfiles.end())
       return &It->second;
     return nullptr;
   }
   void runOnFunction(const Function &F, const FunctionSamples &FS);
+  void findIRAnchors(const Function &F,
+                     std::map<LineLocation, StringRef> &IRAnchors);
+  void findProfileAnchors(const FunctionSamples &FS,
+                          std::map<LineLocation, StringSet<>> &ProfileAnchors);
   void countProfileMismatches(
+      const Function &F, const FunctionSamples &FS,
+      const std::map<LineLocation, StringRef> &IRAnchors,
+      const std::map<LineLocation, StringSet<>> &ProfileAnchors);
+  void countProfileCallsiteMismatches(
       const FunctionSamples &FS,
-      const std::unordered_set<LineLocation, LineLocationHash>
-          &MatchedCallsiteLocs,
-      uint64_t &FuncMismatchedCallsites, uint64_t &FuncProfiledCallsites);
+      const std::map<LineLocation, StringRef> &IRAnchors,
+      const std::map<LineLocation, StringSet<>> &ProfileAnchors,
 
+      uint64_t &FuncMismatchedCallsites, uint64_t &FuncProfiledCallsites);
   LocToLocMap &getIRToProfileLocationMap(const Function &F) {
     auto Ret = FuncMappings.try_emplace(
         FunctionSamples::getCanonicalFnName(F.getName()), LocToLocMap());
     return Ret.first->second;
   }
   void distributeIRToProfileLocationMap();
   void distributeIRToProfileLocationMap(FunctionSamples &FS);
-  void populateProfileCallsites(
-      const FunctionSamples &FS,
-      StringMap<std::set<LineLocation>> &CalleeToCallsitesMap);
   void runStaleProfileMatching(
-      const std::map<LineLocation, StringRef> &IRLocations,
-      StringMap<std::set<LineLocation>> &CalleeToCallsitesMap,
+      const Function &F, const std::map<LineLocation, StringRef> &IRAnchors,
+      const std::map<LineLocation, StringSet<>> &ProfileAnchors,
       LocToLocMap &IRToProfileLocationMap);
 };
 
 /// Sample profile pass.
 ///
 /// This pass reads profile data from the file specified by
 /// -sample-profile-file and annotates every affected function with the
 /// profile information found in that file.
if (ReportProfileStaleness || PersistProfileStaleness ||
       SalvageStaleProfile) {
     MatchingManager =
         std::make_unique<SampleProfileMatcher>(M, *Reader, ProbeManager.get());
   }
 
   return true;
 }
 
-void SampleProfileMatcher::countProfileMismatches(
-    const FunctionSamples &FS,
-    const std::unordered_set<LineLocation, LineLocationHash>
-        &MatchedCallsiteLocs,
-    uint64_t &FuncMismatchedCallsites, uint64_t &FuncProfiledCallsites) {
-
-  auto isInvalidLineOffset = [](uint32_t LineOffset) {
-    return LineOffset & 0x8000;
-  };
+void SampleProfileMatcher::findIRAnchors(
+    const Function &F, std::map<LineLocation, StringRef> &IRAnchors) {
+  // Extract profile matching anchors in the IR.
+  for (auto &BB : F) {
+    for (auto &I : BB) {
+      // TODO: Support line-number based location(AutoFDO).
+      if (FunctionSamples::ProfileIsProbeBased && isa<PseudoProbeInst>(&I)) {
+        if (std::optional<PseudoProbe> Probe = extractProbe(I))
+          IRAnchors.emplace(LineLocation(Probe->Id, 0), StringRef());
+      }
 
-  // Check if there are any callsites in the profile that does not match to any
-  // IR callsites, those callsite samples will be discarded.
-  for (auto &I : FS.getBodySamples()) {
-    const LineLocation &Loc = I.first;
-    if (isInvalidLineOffset(Loc.LineOffset))
+      if (!isa<CallBase>(&I) || isa<IntrinsicInst>(&I))
         continue;
 
-    uint64_t Count = I.second.getSamples();
-    if (!I.second.getCallTargets().empty()) {
-      TotalCallsiteSamples += Count;
-      FuncProfiledCallsites++;
-      if (!MatchedCallsiteLocs.count(Loc)) {
-        MismatchedCallsiteSamples += Count;
-        FuncMismatchedCallsites++;
-      }
-    }
-  }
+      const auto *CB = dyn_cast<CallBase>(&I);
+      if (auto &DLoc = I.getDebugLoc()) {
+        LineLocation IRCallsite = FunctionSamples::getCallSiteIdentifier(DLoc);
+        StringRef CalleeName = UnknownIndirectCallee;
+        if (Function *Callee = CB->getCalledFunction())
+          CalleeName = FunctionSamples::getCanonicalFnName(Callee->getName());
+
+        // Force to overwrite the callee name in case any non-call location was
+        // written before.
+        auto R = IRAnchors.emplace(IRCallsite, CalleeName);
+        R.first->second = CalleeName;
+        assert((!FunctionSamples::ProfileIsProbeBased || R.second ||
+                R.first->second == CalleeName) &&
+               "Overwrite non-call or different callee name location for "
+               "pseudo probe callsite");
+      }
 
-  for (auto &I : FS.getCallsiteSamples()) {
-    const LineLocation &Loc = I.first;
-    if (isInvalidLineOffset(Loc.LineOffset))
-      continue;
+void SampleProfileMatcher::countProfileMismatches(
+    const Function &F, const FunctionSamples &FS,
+    const std::map<LineLocation, StringRef> &IRAnchors,
+    const std::map<LineLocation, StringSet<>> &ProfileAnchors) {
+  bool IsFuncHashMismatch = false;
+  if (FunctionSamples::ProfileIsProbeBased) {
+    uint64_t Count = FS.getTotalSamples();
+    TotalFuncHashSamples += Count;
+    TotalProfiledFunc++;
+    if (!ProbeManager->profileIsValid(F, FS)) {
+      MismatchedFuncHashSamples += Count;
+      NumMismatchedFuncHash++;
+      IsFuncHashMismatch = true;
+    }
+  }
 
-    uint64_t Count = 0;
-    for (auto &FM : I.second) {
-      Count += FM.second.getHeadSamplesEstimate();
-    }
-    TotalCallsiteSamples += Count;
+  uint64_t FuncMismatchedCallsites = 0;
+  uint64_t FuncProfiledCallsites = 0;
+  countProfileCallsiteMismatches(FS, IRAnchors, ProfileAnchors,
+                                 FuncMismatchedCallsites,
+                                 FuncProfiledCallsites);
+  TotalProfiledCallsites += FuncProfiledCallsites;
+  NumMismatchedCallsites += FuncMismatchedCallsites;
+  LLVM_DEBUG({
+    if (FunctionSamples::ProfileIsProbeBased && !IsFuncHashMismatch &&
+        FuncMismatchedCallsites)
+      dbgs() << "Function checksum is matched but there are "
+             << FuncMismatchedCallsites << "/" << FuncProfiledCallsites
+             << " mismatched callsites.\n";
+  });
+
+void SampleProfileMatcher::countProfileCallsiteMismatches(
+    const FunctionSamples &FS,
+    const std::map<LineLocation, StringRef> &IRAnchors,
+    const std::map<LineLocation, StringSet<>> &ProfileAnchors,
+    uint64_t &FuncMismatchedCallsites, uint64_t &FuncProfiledCallsites) {
+
+  // Check if there are any callsites in the profile that does not match to any
+  // IR callsites, those callsite samples will be discarded.
+  for (const auto &I : ProfileAnchors) {
+    const auto &Loc = I.first;
+    const auto &Callees = I.second;
+    assert(!Callees.empty() && "Callees should not be empty");
+
+    StringRef IRCalleeName;
+    const auto &IR = IRAnchors.find(Loc);
+    if (IR != IRAnchors.end())
+      IRCalleeName = IR->second;
+
+    // Compute number of samples in the original profile.
+    uint64_t CallsiteSamples = 0;
+    auto CTM = FS.findCallTargetMapAt(Loc);
+    if (CTM) {
+      for (const auto &I : CTM.get())
+        CallsiteSamples += I.second;
+    }
+    const auto *FSMap = FS.findFunctionSamplesMapAt(Loc);
+    if (FSMap) {
+      for (const auto &I : *FSMap)
+        CallsiteSamples += I.second.getTotalSamples();
+    }
+
+    bool CallsiteIsMatched = false;
+    // Since indirect call does not have CalleeName, check conservatively if
+    // callsite in the profile is a callsite location. This is to reduce num of
+    // false positive since otherwise all the indirect call samples will be
+    // reported as mismatching.
+    if (IRCalleeName == UnknownIndirectCallee)
+      CallsiteIsMatched = true;
+    else if (Callees.size() == 1 && Callees.count(IRCalleeName))
+      CallsiteIsMatched = true;
+
     FuncProfiledCallsites++;
-    if (!MatchedCallsiteLocs.count(Loc)) {
-      MismatchedCallsiteSamples += Count;
+    TotalCallsiteSamples += CallsiteSamples;
+    if (!CallsiteIsMatched) {
       FuncMismatchedCallsites++;
+      MismatchedCallsiteSamples += CallsiteSamples;
     }
   }
 }
 
-// Populate the anchors(direct callee name) from profile.
-void SampleProfileMatcher::populateProfileCallsites(
+void SampleProfileMatcher::findProfileAnchors(
     const FunctionSamples &FS,
-    StringMap<std::set<LineLocation>> &CalleeToCallsitesMap) {
+    std::map<LineLocation, StringSet<>> &ProfileAnchors) {
+  auto isInvalidLineOffset = [](uint32_t LineOffset) {
+    return LineOffset & 0x8000;
+  };
+
   for (const auto &I : FS.getBodySamples()) {
-    const auto &Loc = I.first;
-    const auto &CTM = I.second.getCallTargets();
-    // Filter out possible indirect calls, use direct callee name as anchor.
-    if (CTM.size() == 1) {
-      StringRef CalleeName = CTM.begin()->first();
-      const auto &Candidates = CalleeToCallsitesMap.try_emplace(
-          CalleeName, std::set<LineLocation>());
-      Candidates.first->second.insert(Loc);
+    const LineLocation &Loc = I.first;
+    if (isInvalidLineOffset(Loc.LineOffset))
+      continue;
+    for (const auto &I : I.second.getCallTargets()) {
+      auto Ret = ProfileAnchors.try_emplace(Loc, StringSet<>());
+      Ret.first->second.insert(I.first());
     }
   }
 
   for (const auto &I : FS.getCallsiteSamples()) {
     const LineLocation &Loc = I.first;
+    if (isInvalidLineOffset(Loc.LineOffset))
+      continue;
     const auto &CalleeMap = I.second;
-    // Filter out possible indirect calls, use direct callee name as anchor.
-    if (CalleeMap.size() == 1) {
-      StringRef CalleeName = CalleeMap.begin()->first;
-      const auto &Candidates = CalleeToCallsitesMap.try_emplace(
-          CalleeName, std::set<LineLocation>());
-      Candidates.first->second.insert(Loc);
+    for (const auto &I : CalleeMap) {
+      auto Ret = ProfileAnchors.try_emplace(Loc, StringSet<>());
+      Ret.first->second.insert(I.first);
     }
   }
 }
 
 // Call target name anchor based profile fuzzy matching.
 // Input:
 // For IR locations, the anchor is the callee name of direct callsite; For
 // profile locations, it's the call target name for BodySamples or inlinee's
 // profile name for CallsiteSamples.
 // Matching heuristic:
 // First match all the anchors in lexical order, then split the non-anchor
 // locations between the two anchors evenly, first half are matched based on the
 // start anchor, second half are matched based on the end anchor.
 // For example, given:
 // IR locations:      [1, 2(foo), 3, 5, 6(bar), 7]
 // Profile locations: [1, 2, 3(foo), 4, 7, 8(bar), 9]
 // The matching gives:
 //   [1,    2(foo), 3,  5,  6(bar), 7]
 //    |     |       |   |     |     |
 //   [1, 2, 3(foo), 4,  7,  8(bar), 9]
 // The output mapping: [2->3, 3->4, 5->7, 6->8, 7->9].
 void SampleProfileMatcher::runStaleProfileMatching(
-    const std::map<LineLocation, StringRef> &IRLocations,
-    StringMap<std::set<LineLocation>> &CalleeToCallsitesMap,
+    const Function &F, const std::map<LineLocation, StringRef> &IRAnchors,
+    const std::map<LineLocation, StringSet<>> &ProfileAnchors,
     LocToLocMap &IRToProfileLocationMap) {
+  LLVM_DEBUG(dbgs() << "Run stale profile matching for " << F.getName()
+                    << "\n");
   assert(IRToProfileLocationMap.empty() &&
          "Run stale profile matching only once per function");
 
+  StringMap<std::set<LineLocation>> CalleeToCallsitesMap;
+  for (const auto &I : ProfileAnchors) {
+    const auto &Loc = I.first;
+    const auto &Callees = I.second;
+    // Filter out possible indirect calls, use direct callee name as anchor.
+    if (Callees.size() == 1) {
+      StringRef CalleeName = Callees.begin()->first();
+      const auto &Candidates = CalleeToCallsitesMap.try_emplace(
+          CalleeName, std::set<LineLocation>());
+      Candidates.first->second.insert(Loc);
+    }
+  }
+
   auto InsertMatching = [&](const LineLocation &From, const LineLocation &To) {
     // Skip the unchanged location mapping to save memory.
     if (From != To)
       IRToProfileLocationMap.insert({From, To});
   };
 
   // Use function's beginning location as the initial anchor.
   int32_t LocationDelta = 0;
   SmallVector<LineLocation> LastMatchedNonAnchors;
 
-  for (const auto &IR : IRLocations) {
+  for (const auto &IR : IRAnchors) {
     const auto &Loc = IR.first;
     StringRef CalleeName = IR.second;
     bool IsMatchedAnchor = false;
     // Match the anchor location in lexical order.
     if (!CalleeName.empty()) {
-      auto ProfileAnchors = CalleeToCallsitesMap.find(CalleeName);
-      if (ProfileAnchors != CalleeToCallsitesMap.end() &&
-          !ProfileAnchors->second.empty()) {
-        auto CI = ProfileAnchors->second.begin();
+      auto CandidateAnchors = CalleeToCallsitesMap.find(CalleeName);
+      if (CandidateAnchors != CalleeToCallsitesMap.end() &&
+          !CandidateAnchors->second.empty()) {
+        auto CI = CandidateAnchors->second.begin();
         const auto Candidate = *CI;
-        ProfileAnchors->second.erase(CI);
+        CandidateAnchors->second.erase(CI);
         InsertMatching(Loc, Candidate);
         LLVM_DEBUG(dbgs() << "Callsite with callee:" << CalleeName
                           << " is matched from " << Loc << " to " << Candidate
                           << "\n");
         LocationDelta = Candidate.LineOffset - Loc.LineOffset;
 
         // Match backwards for non-anchor locations.
         // The locations in LastMatchedNonAnchors have been matched forwards
if (!IsMatchedAnchor) {
       LLVM_DEBUG(dbgs() << "Location is matched from " << Loc << " to "
                         << Candidate << "\n");
       LastMatchedNonAnchors.emplace_back(Loc);
     }
   }
 }
 
 void SampleProfileMatcher::runOnFunction(const Function &F,
                                          const FunctionSamples &FS) {
-  bool IsFuncHashMismatch = false;
-  if (FunctionSamples::ProfileIsProbeBased) {
-    uint64_t Count = FS.getTotalSamples();
-    TotalFuncHashSamples += Count;
-    TotalProfiledFunc++;
-    if (!ProbeManager->profileIsValid(F, FS)) {
-      MismatchedFuncHashSamples += Count;
-      NumMismatchedFuncHash++;
-      IsFuncHashMismatch = true;
-    }
-  }
-
-  std::unordered_set<LineLocation, LineLocationHash> MatchedCallsiteLocs;
-  // The value of the map is the name of direct callsite and use empty StringRef
-  // for non-direct-call site.
-  std::map<LineLocation, StringRef> IRLocations;
-
-  // Extract profile matching anchors and profile mismatch metrics in the IR.
-  for (auto &BB : F) {
-    for (auto &I : BB) {
-      // TODO: Support line-number based location(AutoFDO).
-      if (FunctionSamples::ProfileIsProbeBased && isa<PseudoProbeInst>(&I)) {
-        if (std::optional<PseudoProbe> Probe = extractProbe(I))
-          IRLocations.emplace(LineLocation(Probe->Id, 0), StringRef());
-      }
-
-      if (!isa<CallBase>(&I) || isa<IntrinsicInst>(&I))
-        continue;
-
-      const auto *CB = dyn_cast<CallBase>(&I);
-      if (auto &DLoc = I.getDebugLoc()) {
-        LineLocation IRCallsite = FunctionSamples::getCallSiteIdentifier(DLoc);
-
-        StringRef CalleeName;
-        if (Function *Callee = CB->getCalledFunction())
-          CalleeName = FunctionSamples::getCanonicalFnName(Callee->getName());
-
-        // Force to overwrite the callee name in case any non-call location was
-        // written before.
-        auto R = IRLocations.emplace(IRCallsite, CalleeName);
-        R.first->second = CalleeName;
-        assert((!FunctionSamples::ProfileIsProbeBased || R.second ||
-                R.first->second == CalleeName) &&
-               "Overwrite non-call or different callee name location for "
-               "pseudo probe callsite");
-
-        // Go through all the callsites on the IR and flag the callsite if the
-        // target name is the same as the one in the profile.
-        const auto CTM = FS.findCallTargetMapAt(IRCallsite);
-        const auto CallsiteFS = FS.findFunctionSamplesMapAt(IRCallsite);
-
-        // Indirect call case.
-        if (CalleeName.empty()) {
-          // Since indirect call does not have the CalleeName, check
-          // conservatively if callsite in the profile is a callsite location.
-          // This is to avoid nums of false positive since otherwise all the
-          // indirect call samples will be reported as mismatching.
-          if ((CTM && !CTM->empty()) || (CallsiteFS && !CallsiteFS->empty()))
-            MatchedCallsiteLocs.insert(IRCallsite);
-        } else {
-          // Check if the call target name is matched for direct call case.
-          if ((CTM && CTM->count(CalleeName)) ||
-              (CallsiteFS && CallsiteFS->count(CalleeName)))
-            MatchedCallsiteLocs.insert(IRCallsite);
-        }
-      }
-    }
-  }
+  // Anchors for IR. It's a map from IR location to callee name, callee name is
+  // empty for non-call instruction and use a dummy name(UnknownIndirectCallee)
+  // for unknown indrect callee name.
+  std::map<LineLocation, StringRef> IRAnchors;
+  findIRAnchors(F, IRAnchors);
+  // Anchors for profile. It's a map from callsite location to a set of callee
+  // name.
+  std::map<LineLocation, StringSet<>> ProfileAnchors;
+  findProfileAnchors(FS, ProfileAnchors);
 
   // Detect profile mismatch for profile staleness metrics report.
   if (ReportProfileStaleness || PersistProfileStaleness) {
-    uint64_t FuncMismatchedCallsites = 0;
-    uint64_t FuncProfiledCallsites = 0;
-    countProfileMismatches(FS, MatchedCallsiteLocs, FuncMismatchedCallsites,
-                           FuncProfiledCallsites);
-    TotalProfiledCallsites += FuncProfiledCallsites;
-    NumMismatchedCallsites += FuncMismatchedCallsites;
-    LLVM_DEBUG({
-      if (FunctionSamples::ProfileIsProbeBased && !IsFuncHashMismatch &&
-          FuncMismatchedCallsites)
-        dbgs() << "Function checksum is matched but there are "
-               << FuncMismatchedCallsites << "/" << FuncProfiledCallsites
-               << " mismatched callsites.\n";
-    });
+    countProfileMismatches(F, FS, IRAnchors, ProfileAnchors);
   }
 
-  if (IsFuncHashMismatch && SalvageStaleProfile) {
-    LLVM_DEBUG(dbgs() << "Run stale profile matching for " << F.getName()
-                      << "\n");
-
-    StringMap<std::set<LineLocation>> CalleeToCallsitesMap;
-    populateProfileCallsites(FS, CalleeToCallsitesMap);
-
+  // Run profile matching for checksum mismatched profile, currently only
+  // support for pseudo-probe.
+  if (SalvageStaleProfile && FunctionSamples::ProfileIsProbeBased &&
+      !ProbeManager->profileIsValid(F, FS)) {
     // The matching result will be saved to IRToProfileLocationMap, create a new
     // map for each function.
-    auto &IRToProfileLocationMap = getIRToProfileLocationMap(F);
-
-    runStaleProfileMatching(IRLocations, CalleeToCallsitesMap,
-                            IRToProfileLocationMap);
+    runStaleProfileMatching(F, IRAnchors, ProfileAnchors,
+                            getIRToProfileLocationMap(F));
   }
 }
 
 void SampleProfileMatcher::runOnModule() {
   for (auto &F : M) {
     if (F.isDeclaration() || !F.hasFnAttribute("use-sample-profile"))
       continue;
     FunctionSamples *FS = nullptr;
     if (FlattenProfileForMatching)

llvm/test/Transforms/SampleProfile/profile-mismatch.ll

 ; REQUIRES: x86_64-linux
 ; RUN: opt < %s -passes=sample-profile -sample-profile-file=%S/Inputs/profile-mismatch.prof -report-profile-staleness -persist-profile-staleness -flatten-profile-for-matching=0 -S 2>%t -o %t.ll
 ; RUN: FileCheck %s --input-file %t
 ; RUN: FileCheck %s --input-file %t.ll -check-prefix=CHECK-MD
 ; RUN: llc < %t.ll -filetype=obj -o %t.obj
 ; RUN: llvm-objdump --section-headers %t.obj | FileCheck %s --check-prefix=CHECK-OBJ
 ; RUN: llc < %t.ll -filetype=asm -o - | FileCheck %s --check-prefix=CHECK-ASM
 
-; CHECK: (2/3) of callsites' profile are invalid and (15/25) of samples are discarded due to callsite location mismatch.
+; CHECK: (2/3) of callsites' profile are invalid and (25/35) of samples are discarded due to callsite location mismatch.
 
-; CHECK-MD: ![[#]] = !{!"NumMismatchedCallsites", i64 2, !"TotalProfiledCallsites", i64 3, !"MismatchedCallsiteSamples", i64 15, !"TotalCallsiteSamples", i64 25}
+; CHECK-MD: ![[#]] = !{!"NumMismatchedCallsites", i64 2, !"TotalProfiledCallsites", i64 3, !"MismatchedCallsiteSamples", i64 25, !"TotalCallsiteSamples", i64 35}
 
 ; CHECK-OBJ: .llvm_stats
 
 ; CHECK-ASM: .section  .llvm_stats,"",@progbits
 ; CHECK-ASM: .byte 22
 ; CHECK-ASM: .ascii  "NumMismatchedCallsites"
 ; CHECK-ASM: .byte 4
 ; CHECK-ASM: .ascii  "Mg=="
 ; CHECK-ASM: .byte 22
 ; CHECK-ASM: .ascii  "TotalProfiledCallsites"
 ; CHECK-ASM: .byte 4
 ; CHECK-ASM: .ascii  "Mw=="
 ; CHECK-ASM: .byte 25
 ; CHECK-ASM: .ascii  "MismatchedCallsiteSamples"
 ; CHECK-ASM: .byte 4
-; CHECK-ASM: .ascii  "MTU="
+; CHECK-ASM: .ascii  "MjU="
 ; CHECK-ASM: .byte 20
 ; CHECK-ASM: .ascii  "TotalCallsiteSamples"
 ; CHECK-ASM: .byte 4
-; CHECK-ASM: .ascii  "MjU="
+; CHECK-ASM: .ascii  "MzU="
 
 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"
 
 @x = dso_local global i32 0, align 4, !dbg !0
 
 ; Function Attrs: nounwind uwtable
 define dso_local i32 @foo(i32 noundef %x) #0 !dbg !12 {