[llvm-profgen] Loading binary functions from .symtab when DWARF info is incomplete (#163654)

Indexing in .debug_str section could lead to integer overflow when in
DWARF32 format.

https://github.com/llvm/llvm-project/blob/e61e6251b692ffe71910bad22b82e41313f003cf/llvm/lib/DWP/DWP.cpp#L35C30-L35C47

This can lead to missing symbols from the DWARF info, and hurts profile
quality. As a workaround, we may use information from the symbol table
(.symtab), and recover the missing symbols with addresses and ranges.

Output:
```
# Before
...
warning: 6.64%(338916009/5106344252) of samples are from ranges that do not belong to any functions.

# After
...
warning: 0.07%(3501587/4906133148) of samples are from ranges that do not belong to any functions.
warning: 5.71%(280266919/4906133148) of samples are from ranges that belong to functions recovered from symbol table.
```

We see 0.4% - 1.35% performance improvements on our internal services
where profiles are generated with binaries that have .debug_str section
over 4GB.

Author: HighW4y2H3ll

llvm/include/llvm/MC/MCPseudoProbe.h

@@ -328,6 +328,7 @@ class MCDecodedPseudoProbeInlineTree
 
   // Return false if it's a dummy inline site
   bool hasInlineSite() const { return !isRoot() && !Parent->isRoot(); }
+  bool isTopLevelFunc() const { return !isRoot() && Parent->isRoot(); }
   InlineSite getInlineSite() const { return InlineSite(Guid, ProbeId); }
   void setProbes(MutableArrayRef<MCDecodedPseudoProbe> ProbesRef) {
     Probes = ProbesRef.data();

llvm/include/llvm/ProfileData/SampleProf.h

@@ -1214,13 +1214,19 @@ class FunctionSamples {
     // Note the sequence of the suffixes in the knownSuffixes array matters.
     // If suffix "A" is appended after the suffix "B", "A" should be in front
     // of "B" in knownSuffixes.
-    const char *KnownSuffixes[] = {LLVMSuffix, PartSuffix, UniqSuffix};
+    const SmallVector<StringRef> KnownSuffixes{LLVMSuffix, PartSuffix,
+                                               UniqSuffix};
+    return getCanonicalFnName(FnName, KnownSuffixes, Attr);
+  }
+
+  static StringRef getCanonicalFnName(StringRef FnName,
+                                      ArrayRef<StringRef> Suffixes,
+                                      StringRef Attr = "selected") {
     if (Attr == "" || Attr == "all")
       return FnName.split('.').first;
     if (Attr == "selected") {
       StringRef Cand(FnName);
-      for (const auto &Suf : KnownSuffixes) {
-        StringRef Suffix(Suf);
+      for (const auto Suffix : Suffixes) {
         // If the profile contains ".__uniq." suffix, don't strip the
         // suffix for names in the IR.
         if (Suffix == UniqSuffix && FunctionSamples::HasUniqSuffix)
@@ -1229,7 +1235,7 @@ class FunctionSamples {
         if (It == StringRef::npos)
           continue;
         auto Dit = Cand.rfind('.');
-        if (Dit == It + Suffix.size() - 1)
+        if (Dit == It || Dit == It + Suffix.size() - 1)
           Cand = Cand.substr(0, It);
       }
       return Cand;

llvm/test/tools/llvm-profgen/Inputs/missing-dwarf.exe

@@ -0,0 +1,37 @@
+; RUN: rm -rf %t
+; RUN: mkdir -p %t
+; RUN: cd %t
+
+; RUN: echo -e "1\n401120-40113b:1\n1\n40112f->401110:1" > %t.prof
+
+; Test --load-function-from-symbol=0
+; RUN: llvm-profgen --format=text --unsymbolized-profile=%t.prof --binary=%S/Inputs/missing-dwarf.exe --output=%t1 --fill-zero-for-all-funcs --show-detailed-warning --use-offset=0 --load-function-from-symbol=0 2>&1 | FileCheck %s --check-prefix=CHECK-NO-LOAD-SYMTAB
+
+; CHECK-NO-LOAD-SYMTAB:      warning: Loading of DWARF info completed, but no binary functions have been retrieved.
+
+; Test --load-function-from-symbol=1
+; RUN: llvm-profgen --format=text --unsymbolized-profile=%t.prof --binary=%S/Inputs/missing-dwarf.exe --output=%t2 --fill-zero-for-all-funcs --show-detailed-warning --use-offset=0 --load-function-from-symbol=1
+; RUN: FileCheck %s --input-file %t2 --check-prefix=CHECK-LOAD-SYMTAB
+
+; CHECK-LOAD-SYMTAB:      main:2:1
+; CHECK-LOAD-SYMTAB-NEXT:  1: 1
+; CHECK-LOAD-SYMTAB-NEXT:  2: 1 foo:1
+; CHECK-LOAD-SYMTAB-NEXT:  !CFGChecksum: 281479271677951
+; CHECK-LOAD-SYMTAB-NEXT: foo:0:0
+; CHECK-LOAD-SYMTAB-NEXT:  1: 0
+; CHECK-LOAD-SYMTAB-NEXT:  !CFGChecksum: 4294967295
+
+; Build instructions:
+; missing-dwarf.o:       clang -gsplit-dwarf=split -fdebug-compilation-dir=.  test.c   -fdebug-info-for-profiling  -fpseudo-probe-for-profiling  -O0 -g -o missing-dwarf.o -c
+; missing-dwarf.exe:     clang -fdebug-compilation-dir=.  missing-dwarf.o -o missing-dwarf.exe  -fdebug-info-for-profiling  -fpseudo-probe-for-profiling  -O0 -g
+
+; Source code:
+
+int foo() {
+  return 1;
+}
+
+int main() {
+  foo();
+  return 0;
+}

llvm/test/tools/llvm-profgen/missing-dwarf.test

@@ -22,6 +22,7 @@ extern cl::opt<bool> ShowDetailedWarning;
 extern cl::opt<bool> InferMissingFrames;
 extern cl::opt<bool> EnableCSPreInliner;
 extern cl::opt<bool> UseContextCostForPreInliner;
+extern cl::opt<bool> LoadFunctionFromSymbol;
 
 } // end namespace llvm
 

llvm/tools/llvm-profgen/Options.h

@@ -1284,6 +1284,7 @@ void PerfScriptReader::warnInvalidRange() {
   uint64_t TotalRangeNum = 0;
   uint64_t InstNotBoundary = 0;
   uint64_t UnmatchedRange = 0;
+  uint64_t RecoveredRange = 0;
   uint64_t RangeCrossFunc = 0;
   uint64_t BogusRange = 0;
 
@@ -1309,6 +1310,9 @@ void PerfScriptReader::warnInvalidRange() {
       continue;
     }
 
+    if (FRange->Func->NameStatus != DwarfNameStatus::Matched)
+      RecoveredRange += I.second;
+
     if (EndAddress >= FRange->EndAddress) {
       RangeCrossFunc += I.second;
       WarnInvalidRange(StartAddress, EndAddress, RangeCrossFuncMsg);
@@ -1328,6 +1332,9 @@ void PerfScriptReader::warnInvalidRange() {
   emitWarningSummary(
       UnmatchedRange, TotalRangeNum,
       "of samples are from ranges that do not belong to any functions.");
+  emitWarningSummary(RecoveredRange, TotalRangeNum,
+                     "of samples are from ranges that belong to functions "
+                     "recovered from symbol table.");
   emitWarningSummary(
       RangeCrossFunc, TotalRangeNum,
       "of samples are from ranges that do cross function boundaries.");

llvm/tools/llvm-profgen/PerfReader.cpp

@@ -503,8 +503,11 @@ ProfileGenerator::getTopLevelFunctionProfile(FunctionId FuncName) {
 void ProfileGenerator::generateProfile() {
   collectProfiledFunctions();
 
-  if (Binary->usePseudoProbes())
+  if (Binary->usePseudoProbes()) {
     Binary->decodePseudoProbe();
+    if (LoadFunctionFromSymbol)
+      Binary->loadSymbolsFromPseudoProbe();
+  }
 
   if (SampleCounters) {
     if (Binary->usePseudoProbes()) {
@@ -732,6 +735,14 @@ ProfileGeneratorBase::getCalleeNameForAddress(uint64_t TargetAddress) {
   if (!FRange || !FRange->IsFuncEntry)
     return StringRef();
 
+  // DWARF and symbol table may have mismatching function names. Instead, we'll
+  // try to use its pseudo probe name first.
+  if (Binary->usePseudoProbes()) {
+    auto FuncName = Binary->findPseudoProbeName(FRange->Func);
+    if (FuncName.size())
+      return FunctionSamples::getCanonicalFnName(FuncName);
+  }
+
   return FunctionSamples::getCanonicalFnName(FRange->getFuncName());
 }
 
@@ -919,6 +930,8 @@ void CSProfileGenerator::generateProfile() {
     Binary->decodePseudoProbe();
     if (InferMissingFrames)
       initializeMissingFrameInferrer();
+    if (LoadFunctionFromSymbol)
+      Binary->loadSymbolsFromPseudoProbe();
   }
 
   if (SampleCounters) {

llvm/tools/llvm-profgen/ProfileGenerator.cpp

@@ -37,6 +37,14 @@ cl::opt<bool> ShowSourceLocations("show-source-locations",
                                   cl::desc("Print source locations."),
                                   cl::cat(ProfGenCategory));
 
+cl::opt<bool> LoadFunctionFromSymbol(
+    "load-function-from-symbol", cl::init(true),
+    cl::desc(
+        "Gather additional binary function info from symbols (e.g. .symtab) in "
+        "case dwarf info is incomplete. Only support binaries in ELF format "
+        "with pseudo probe, for other formats, this flag will be a no-op."),
+    cl::cat(ProfGenCategory));
+
 static cl::opt<bool>
     ShowCanonicalFnName("show-canonical-fname",
                         cl::desc("Print canonical function name."),
@@ -257,6 +265,9 @@ void ProfiledBinary::load() {
   if (ShowDisassemblyOnly)
     decodePseudoProbe(Obj);
 
+  if (LoadFunctionFromSymbol && UsePseudoProbes)
+    loadSymbolsFromSymtab(Obj);
+
   // Disassemble the text sections.
   disassemble(Obj);
 
@@ -461,6 +472,13 @@ void ProfiledBinary::decodePseudoProbe(const ObjectFile *Obj) {
   } else {
     for (auto *F : ProfiledFunctions) {
       GuidFilter.insert(Function::getGUIDAssumingExternalLinkage(F->FuncName));
+      // DWARF name might be broken when a DWARF32 .debug_str.dwo section
+      // execeeds 4GB. We expect symbol table to contain the correct function
+      // names which matches the pseudo probe. Adding back all the GUIDs if
+      // possible.
+      auto AltGUIDs = AlternativeFunctionGUIDs.equal_range(F);
+      for (const auto &[_, Func] : make_range(AltGUIDs))
+        GuidFilter.insert(Func);
       for (auto &Range : F->Ranges) {
         auto GUIDs = StartAddrToSymMap.equal_range(Range.first);
         for (const auto &[StartAddr, Func] : make_range(GUIDs))
@@ -522,7 +540,9 @@ void ProfiledBinary::setIsFuncEntry(FuncRange *FuncRange,
   // Set IsFuncEntry to ture if there is only one range in the function or the
   // RangeSymName from ELF is equal to its DWARF-based function name.
   if (FuncRange->Func->Ranges.size() == 1 ||
-      (!FuncRange->IsFuncEntry && FuncRange->getFuncName() == RangeSymName))
+      (!FuncRange->IsFuncEntry &&
+       (FuncRange->getFuncName() == RangeSymName ||
+        FuncRange->Func->NameStatus != DwarfNameStatus::Matched)))
     FuncRange->IsFuncEntry = true;
 }
 
@@ -604,13 +624,13 @@ bool ProfiledBinary::dissassembleSymbol(std::size_t SI, ArrayRef<uint8_t> Bytes,
       // Record potential call targets for tail frame inference later-on.
       if (InferMissingFrames && FRange) {
         uint64_t Target = 0;
-        MIA->evaluateBranch(Inst, Address, Size, Target);
+        bool Err = MIA->evaluateBranch(Inst, Address, Size, Target);
         if (MCDesc.isCall()) {
           // Indirect call targets are unknown at this point. Recording the
           // unknown target (zero) for further LBR-based refinement.
           MissingContextInferrer->CallEdges[Address].insert(Target);
         } else if (MCDesc.isUnconditionalBranch()) {
-          assert(Target &&
+          assert(Err &&
                  "target should be known for unconditional direct branch");
           // Any inter-function unconditional jump is considered tail call at
           // this point. This is not 100% accurate and could further be
@@ -820,6 +840,100 @@ void ProfiledBinary::populateSymbolAddressList(const ObjectFile *Obj) {
   }
 }
 
+void ProfiledBinary::loadSymbolsFromSymtab(const ObjectFile *Obj) {
+  // Load binary functions from symbol table when Debug info is incomplete.
+  // Strip the internal suffixes which are not reflected in the DWARF info.
+  const SmallVector<StringRef, 10> Suffixes(
+      {// Internal suffixes from CoroSplit pass
+       ".cleanup", ".destroy", ".resume",
+       // Internal suffixes from Bolt
+       ".cold", ".warm",
+       // Compiler/LTO internal
+       ".llvm.", ".part.", ".isra.", ".constprop.", ".lto_priv."});
+  StringRef FileName = Obj->getFileName();
+  // Only apply this to ELF binary. e.g. COFF file format doesn't have `size`
+  // field in the symbol table.
+  bool IsELFObject = isa<ELFObjectFileBase>(Obj);
+  if (!IsELFObject)
+    return;
+  for (const SymbolRef &Symbol : Obj->symbols()) {
+    const SymbolRef::Type Type = unwrapOrError(Symbol.getType(), FileName);
+    const uint64_t StartAddr = unwrapOrError(Symbol.getAddress(), FileName);
+    const StringRef Name = unwrapOrError(Symbol.getName(), FileName);
+    uint64_t Size = 0;
+    if (LLVM_LIKELY(IsELFObject)) {
+      ELFSymbolRef ElfSymbol(Symbol);
+      Size = ElfSymbol.getSize();
+    }
+
+    if (Size == 0 || Type != SymbolRef::ST_Function)
+      continue;
+
+    const uint64_t EndAddr = StartAddr + Size;
+    const StringRef SymName =
+        FunctionSamples::getCanonicalFnName(Name, Suffixes);
+    assert(StartAddr < EndAddr && StartAddr >= getPreferredBaseAddress() &&
+           "Function range is invalid.");
+
+    auto Range = findFuncRange(StartAddr);
+    if (!Range) {
+      assert(findFuncRange(EndAddr - 1) == nullptr &&
+             "Function range overlaps with existing functions.");
+      // Function from symbol table not found previously in DWARF, store ranges.
+      auto Ret = BinaryFunctions.emplace(SymName, BinaryFunction());
+      auto &Func = Ret.first->second;
+      if (Ret.second) {
+        Func.FuncName = Ret.first->first;
+        HashBinaryFunctions[Function::getGUIDAssumingExternalLinkage(SymName)] =
+            &Func;
+      }
+
+      Func.NameStatus = DwarfNameStatus::Missing;
+      Func.Ranges.emplace_back(StartAddr, EndAddr);
+
+      auto R = StartAddrToFuncRangeMap.emplace(StartAddr, FuncRange());
+      FuncRange &FRange = R.first->second;
+
+      FRange.Func = &Func;
+      FRange.StartAddress = StartAddr;
+      FRange.EndAddress = EndAddr;
+
+    } else if (SymName != Range->getFuncName()) {
+      // Function range already found from DWARF, but the symbol name from
+      // symbol table is inconsistent with debug info. Log this discrepancy and
+      // the alternative function GUID.
+      if (ShowDetailedWarning)
+        WithColor::warning()
+            << "Conflicting name for symbol " << Name << " with range ("
+            << format("%8" PRIx64, StartAddr) << ", "
+            << format("%8" PRIx64, EndAddr) << ")"
+            << ", but the DWARF symbol " << Range->getFuncName()
+            << " indicates an overlapping range ("
+            << format("%8" PRIx64, Range->StartAddress) << ", "
+            << format("%8" PRIx64, Range->EndAddress) << ")\n";
+
+      assert(StartAddr == Range->StartAddress && EndAddr == Range->EndAddress &&
+             "Mismatched function range");
+
+      Range->Func->NameStatus = DwarfNameStatus::Mismatch;
+      AlternativeFunctionGUIDs.emplace(
+          Range->Func, Function::getGUIDAssumingExternalLinkage(SymName));
+
+    } else if (StartAddr != Range->StartAddress &&
+               EndAddr != Range->EndAddress) {
+      // Function already found in DWARF, but the address range from symbol
+      // table conflicts/overlaps with the debug info.
+      WithColor::warning() << "Conflicting range for symbol " << Name
+                           << " with range (" << format("%8" PRIx64, StartAddr)
+                           << ", " << format("%8" PRIx64, EndAddr) << ")"
+                           << ", but the DWARF symbol " << Range->getFuncName()
+                           << " indicates another range ("
+                           << format("%8" PRIx64, Range->StartAddress) << ", "
+                           << format("%8" PRIx64, Range->EndAddress) << ")\n";
+    }
+  }
+}
+
 void ProfiledBinary::loadSymbolsFromDWARFUnit(DWARFUnit &CompilationUnit) {
   for (const auto &DieInfo : CompilationUnit.dies()) {
     llvm::DWARFDie Die(&CompilationUnit, &DieInfo);
@@ -1034,6 +1148,58 @@ void ProfiledBinary::computeInlinedContextSizeForFunc(
   }
 }
 
+void ProfiledBinary::loadSymbolsFromPseudoProbe() {
+  if (!UsePseudoProbes)
+    return;
+
+  const AddressProbesMap &Address2ProbesMap = getAddress2ProbesMap();
+  for (auto *Func : ProfiledFunctions) {
+    if (Func->NameStatus != DwarfNameStatus::Mismatch)
+      continue;
+    for (auto &[StartAddr, EndAddr] : Func->Ranges) {
+      auto Range = findFuncRangeForStartAddr(StartAddr);
+      if (!Range->IsFuncEntry)
+        continue;
+      const auto &Probe = Address2ProbesMap.find(StartAddr, EndAddr);
+      if (Probe.begin() != Probe.end()) {
+        const MCDecodedPseudoProbeInlineTree *InlineTreeNode =
+            Probe.begin()->get().getInlineTreeNode();
+        while (!InlineTreeNode->isTopLevelFunc())
+          InlineTreeNode = static_cast<MCDecodedPseudoProbeInlineTree *>(
+              InlineTreeNode->Parent);
+
+        auto TopLevelProbes = InlineTreeNode->getProbes();
+        auto TopProbe = TopLevelProbes.begin();
+        assert(TopProbe != TopLevelProbes.end() &&
+               TopProbe->getAddress() >= StartAddr &&
+               TopProbe->getAddress() < EndAddr &&
+               "Top level pseudo probe does not match function range");
+
+        const auto *ProbeDesc = getFuncDescForGUID(InlineTreeNode->Guid);
+        auto Ret = PseudoProbeNames.emplace(Func, ProbeDesc->FuncName);
+        if (!Ret.second && Ret.first->second != ProbeDesc->FuncName &&
+            ShowDetailedWarning)
+          WithColor::warning()
+              << "Mismatched pseudo probe names in function " << Func->FuncName
+              << " at range: (" << format("%8" PRIx64, StartAddr) << ", "
+              << format("%8" PRIx64, EndAddr) << "). "
+              << "The previously found pseudo probe name is "
+              << Ret.first->second << " but it conflicts with name "
+              << ProbeDesc->FuncName
+              << " This likely indicates a DWARF error that produces "
+                 "conflicting symbols at the same starting address.\n";
+      }
+    }
+  }
+}
+
+StringRef ProfiledBinary::findPseudoProbeName(const BinaryFunction *Func) {
+  auto ProbeName = PseudoProbeNames.find(Func);
+  if (ProbeName == PseudoProbeNames.end())
+    return StringRef();
+  return ProbeName->second;
+}
+
 void ProfiledBinary::inferMissingFrames(
     const SmallVectorImpl<uint64_t> &Context,
     SmallVectorImpl<uint64_t> &NewContext) {