Fixup GuidFilter && Pseudo probe callee mismatch

Author: HighW4y2H3ll

llvm/tools/llvm-profgen/ProfileGenerator.cpp

@@ -723,7 +723,7 @@ void ProfileGenerator::populateBodySamplesForAllFunctions(
 }
 
 StringRef
-ProfileGeneratorBase::getCalleeNameForAddress(uint64_t TargetAddress) {
+ProfileGeneratorBase::getCalleeNameForAddress(uint64_t TargetAddress, bool RestoreSymbolName) {
   // Get the function range by branch target if it's a call branch.
   auto *FRange = Binary->findFuncRangeForStartAddr(TargetAddress);
 
@@ -732,6 +732,16 @@ ProfileGeneratorBase::getCalleeNameForAddress(uint64_t TargetAddress) {
   if (!FRange || !FRange->IsFuncEntry)
     return StringRef();
 
+  if (RestoreSymbolName && FRange->Func->FromSymtab) {
+    const AddressProbesMap &Address2ProbesMap =
+        Binary->getAddress2ProbesMap();
+    for (const MCDecodedPseudoProbe &Probe :
+         Address2ProbesMap.find(TargetAddress)) {
+      if (const auto *ProbeDesc = Binary->getFuncDescForGUID(Probe.getGuid()))
+        return FunctionSamples::getCanonicalFnName(ProbeDesc->FuncName);
+    }
+  }
+
   return FunctionSamples::getCanonicalFnName(FRange->getFuncName());
 }
 
@@ -1352,7 +1362,7 @@ void CSProfileGenerator::populateBoundarySamplesWithProbes(
         getFunctionProfileForLeafProbe(CtxKey, CallProbe);
     FunctionProfile.addBodySamples(CallProbe->getIndex(), 0, Count);
     FunctionProfile.addTotalSamples(Count);
-    StringRef CalleeName = getCalleeNameForAddress(TargetAddress);
+    StringRef CalleeName = getCalleeNameForAddress(TargetAddress, true);
     if (CalleeName.size() == 0)
       continue;
     FunctionProfile.addCalledTargetSamples(CallProbe->getIndex(),

llvm/tools/llvm-profgen/ProfileGenerator.h

@@ -109,7 +109,7 @@ class ProfileGeneratorBase {
 
   bool filterAmbiguousProfile(FunctionSamples &FS);
 
-  StringRef getCalleeNameForAddress(uint64_t TargetAddress);
+  StringRef getCalleeNameForAddress(uint64_t TargetAddress, bool RestoreSymbolName = false);
 
   void computeSummaryAndThreshold(SampleProfileMap &ProfileMap);
 

llvm/tools/llvm-profgen/ProfiledBinary.cpp

@@ -473,8 +473,12 @@ void ProfiledBinary::decodePseudoProbe(const ObjectFile *Obj) {
       GuidFilter.insert(Function::getGUIDAssumingExternalLinkage(F->FuncName));
       for (auto &Range : F->Ranges) {
         auto GUIDs = StartAddrToSymMap.equal_range(Range.first);
-        for (const auto &[StartAddr, Func] : make_range(GUIDs))
+        for (const auto &[StartAddr, Func] : make_range(GUIDs)) {
           FuncStartAddresses[Func] = StartAddr;
+          // Function name may be changed when symbol table is loaded. Adding
+          // back the original GUID if possible
+          GuidFilter.insert(Func);
+        }
       }
     }
   }
@@ -865,9 +869,11 @@ void ProfiledBinary::loadSymbolsFromSymtab(const ObjectFile *Obj) {
     const uint64_t EndAddr = StartAddr + Size;
     const StringRef SymName =
         FunctionSamples::getCanonicalFnName(Name, Suffixes);
+    assert(StartAddr < EndAddr && StartAddr >= getPreferredBaseAddress());
 
     auto Range = findFuncRange(StartAddr);
     if (!Range) {
+      assert(findFuncRange(EndAddr - 1) == nullptr);
       // Function from symbol table not found previously in DWARF, store ranges.
       auto Ret = BinaryFunctions.emplace(SymName, BinaryFunction());
       auto &Func = Ret.first->second;