[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?

2214

nit: the CalleeName -> CalleeName.

2216

nit: avoid nums of false positive ->

reduce num of false positives

or

avoid false positive

2263

This is passed by value, missing &?

2379–2416

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

2404–2409

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.

2214

done.

2216

done

2263

Good catch, removed populateProfileCallsites

2379–2416

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

2379–2416

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

2193

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

2193

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

2193

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

2193

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

2193

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 the call target name of an indirect call, used to
+  // differentiate from a non-call instruction that also has an empty call
+  // target.
+  static constexpr const char *IndirectCalleeName = "_indirect_call_";
+
 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 countProfileMismatches(
       const FunctionSamples &FS,
-      const std::unordered_set<LineLocation, LineLocationHash>
-          &MatchedCallsiteLocs,
+      const std::map<LineLocation, StringSet<>> &ProfileLocations,
+      const std::map<LineLocation, StringRef> &IRLocations,
       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 populateIRLocations(
-      const Function &F, const FunctionSamples &FS,
-      std::unordered_set<LineLocation, LineLocationHash> &MatchedCallsiteLocs,
+  void populateIRLocations(const Function &F,
                            std::map<LineLocation, StringRef> &IRLocations);
-  void populateProfileCallsites(
+  void populateProfileLocations(
       const FunctionSamples &FS,
+      std::map<LineLocation, StringSet<>> &ProfileLocations);
+  void populateProfileCallsites(
+      const std::map<LineLocation, StringSet<>> ProfileLocations,
       StringMap<std::set<LineLocation>> &CalleeToCallsitesMap);
   void runStaleProfileMatching(
       const std::map<LineLocation, StringRef> &IRLocations,
       StringMap<std::set<LineLocation>> &CalleeToCallsitesMap,
       LocToLocMap &IRToProfileLocationMap);
 };
 
 /// Sample profile pass.
if (ReportProfileStaleness || PersistProfileStaleness ||
     MatchingManager =
         std::make_unique<SampleProfileMatcher>(M, *Reader, ProbeManager.get());
   }
 
   return true;
 }
 
 void SampleProfileMatcher::populateIRLocations(
-    const Function &F, const FunctionSamples &FS,
-    std::unordered_set<LineLocation, LineLocationHash> &MatchedCallsiteLocs,
-    std::map<LineLocation, StringRef> &IRLocations) {
+    const Function &F, std::map<LineLocation, StringRef> &IRLocations) {
   for (auto &BB : F) {
     for (auto &I : BB) {
       // TODO: To support AutoFDO, we need to parse all the regular instructions
       // to extract the line-number based locations. For pseudo-probe,
       // extracting probe instruction is enough.
       if (isa<PseudoProbeInst>(&I)) {
         std::optional<PseudoProbe> Probe = extractProbe(I);
         assert(Probe &&
for (auto &I : BB) {
       if (auto &DLoc = I.getDebugLoc()) {
         // Skip the inlined callsite. For pseudo probe mode, extracting inline
         // info from the probe inst is enough.
         if (DLoc.getInlinedAt())
           continue;
 
         LineLocation IRCallsite = FunctionSamples::getCallSiteIdentifier(DLoc);
 
-        StringRef CalleeName;
+        StringRef CalleeName = IndirectCalleeName;
         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);
+void SampleProfileMatcher::countProfileMismatches(
+    const FunctionSamples &FS,
+    const std::map<LineLocation, StringSet<>> &ProfileLocations,
+    const std::map<LineLocation, StringRef> &IRLocations,
+    uint64_t &FuncMismatchedCallsites, uint64_t &FuncProfiledCallsites) {
 
-        // 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);
-        }
-      }
+  // 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 : ProfileLocations) {
+    const auto &Loc = I.first;
+    const auto &Callees = I.second;
+    assert(!Callees.empty() && "Callees should not be empty");
+
+    StringRef IRCalleeName;
+    const auto &IR = IRLocations.find(Loc);
+    if (IR != IRLocations.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;
+    if (IRCalleeName == IndirectCalleeName)
+      CallsiteIsMatched = true;
+    else if (Callees.size() == 1 && Callees.count(IRCalleeName))
+      CallsiteIsMatched = true;
+
+    FuncProfiledCallsites++;
+    TotalCallsiteSamples += CallsiteSamples;
+    if (!CallsiteIsMatched) {
+      FuncMismatchedCallsites++;
+      MismatchedCallsiteSamples += CallsiteSamples;
     }
   }
 }
 
-void SampleProfileMatcher::countProfileMismatches(
+void SampleProfileMatcher::populateProfileLocations(
     const FunctionSamples &FS,
-    const std::unordered_set<LineLocation, LineLocationHash>
-        &MatchedCallsiteLocs,
-    uint64_t &FuncMismatchedCallsites, uint64_t &FuncProfiledCallsites) {
-
+    std::map<LineLocation, StringSet<>> &ProfileLocations) {
   auto isInvalidLineOffset = [](uint32_t LineOffset) {
     return LineOffset & 0x8000;
   };
 
-  // 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()) {
+  for (const auto &I : FS.getBodySamples()) {
     const LineLocation &Loc = I.first;
     if (isInvalidLineOffset(Loc.LineOffset))
       continue;
-
-    uint64_t Count = I.second.getSamples();
-    if (!I.second.getCallTargets().empty()) {
-      TotalCallsiteSamples += Count;
-      FuncProfiledCallsites++;
-      if (!MatchedCallsiteLocs.count(Loc)) {
-        MismatchedCallsiteSamples += Count;
-        FuncMismatchedCallsites++;
-      }
+    for (const auto &I : I.second.getCallTargets()) {
+      auto Ret = ProfileLocations.try_emplace(Loc, StringSet<>());
+      Ret.first->second.insert(I.first());
     }
   }
 
-  for (auto &I : FS.getCallsiteSamples()) {
+  for (const auto &I : FS.getCallsiteSamples()) {
     const LineLocation &Loc = I.first;
     if (isInvalidLineOffset(Loc.LineOffset))
       continue;
-
-    uint64_t Count = 0;
-    for (auto &FM : I.second) {
-      Count += FM.second.getHeadSamplesEstimate();
-    }
-    TotalCallsiteSamples += Count;
-    FuncProfiledCallsites++;
-    if (!MatchedCallsiteLocs.count(Loc)) {
-      MismatchedCallsiteSamples += Count;
-      FuncMismatchedCallsites++;
+    const auto &CalleeMap = I.second;
+    for (const auto &I : CalleeMap) {
+      auto Ret = ProfileLocations.try_emplace(Loc, StringSet<>());
+      Ret.first->second.insert(I.first);
     }
   }
 }
 
 // Populate the anchors(direct callee name) from profile.
 void SampleProfileMatcher::populateProfileCallsites(
-    const FunctionSamples &FS,
+    const std::map<LineLocation, StringSet<>> ProfileLocations,
     StringMap<std::set<LineLocation>> &CalleeToCallsitesMap) {
-  for (const auto &I : FS.getBodySamples()) {
+  for (const auto &I : ProfileLocations) {
     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);
-    }
-  }
-
-  for (const auto &I : FS.getCallsiteSamples()) {
-    const LineLocation &Loc = I.first;
-    const auto &CalleeMap = I.second;
+    const auto &Calees = I.second;
     // Filter out possible indirect calls, use direct callee name as anchor.
-    if (CalleeMap.size() == 1) {
-      StringRef CalleeName = CalleeMap.begin()->first;
+    if (Calees.size() == 1) {
+      StringRef CalleeName = Calees.begin()->first();
       const auto &Candidates = CalleeToCallsitesMap.try_emplace(
           CalleeName, std::set<LineLocation>());
       Candidates.first->second.insert(Loc);
     }
   }
 }
 
 // Call target name anchor based profile fuzzy matching.
if (FunctionSamples::ProfileIsProbeBased) {
     const auto *FuncDesc = ProbeManager->getDesc(F);
     if (FuncDesc && ProbeManager->profileIsHashMismatched(*FuncDesc, 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.
+  // Anchors for IR. It's a map from IR location to the name of direct call
+  // target. Also use empty StringRef for non-call instruction and use a dummy
+  // name:IndirectCalleeName("_indirect_call_") for indirect callsite.
   std::map<LineLocation, StringRef> IRLocations;
-  // Extract profile matching anchors and profile mismatch metrics in the IR.
-  populateIRLocations(F, FS, MatchedCallsiteLocs, IRLocations);
+  populateIRLocations(F, IRLocations);
+  // Anchors for profile. it's a map from callsite location to set of callee
+  // name.
+  std::map<LineLocation, StringSet<>> ProfileLocations;
+  populateProfileLocations(FS, ProfileLocations);
 
   // Detect profile mismatch for profile staleness metrics report.
   if (ReportProfileStaleness || PersistProfileStaleness) {
     uint64_t FuncMismatchedCallsites = 0;
     uint64_t FuncProfiledCallsites = 0;
-    countProfileMismatches(FS, MatchedCallsiteLocs, FuncMismatchedCallsites,
-                           FuncProfiledCallsites);
+    countProfileMismatches(FS, ProfileLocations, IRLocations,
+                           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";
     });
   }
 
   if (IsFuncHashMismatch && SalvageStaleProfile) {
     LLVM_DEBUG(dbgs() << "Run stale profile matching for " << F.getName()
                       << "\n");
 
     StringMap<std::set<LineLocation>> CalleeToCallsitesMap;
-    populateProfileCallsites(FS, CalleeToCallsitesMap);
+    populateProfileCallsites(ProfileLocations, CalleeToCallsitesMap);
 
     // The matching result will be saved to IRToProfileLocationMap, create a new
     // map for each function.
     auto &IRToProfileLocationMap = getIRToProfileLocationMap(F);
 
     runStaleProfileMatching(IRLocations, CalleeToCallsitesMap,
                             IRToProfileLocationMap);
   }
 }
 
 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 {