[BOLT][NFCI] Strip suffix in getLTOCommonName

Also provide a mechanism to override the list of suffixes to consider.
Override LTOSuffixes for getGUID in pseudo probe parsing.

Test Plan: 

Reviewers:
Subscribers:

Tasks:

Tags:

Differential Revision: https://phabricator.intern.facebook.com/D61857819

Pull Request: #106243

Author: aaupov

bolt/include/bolt/Utils/Utils.h

@@ -41,6 +41,11 @@ std::string getEscapedName(const StringRef &Name);
 /// Return the unescaped name
 std::string getUnescapedName(const StringRef &Name);
 
+/// Return a common part for a given \p Name wrt a given \p Suffixes list.
+/// Preserve the suffix if \p KeepSuffix is set, only dropping characters
+/// following it, otherwise drop the suffix as well.
+std::optional<StringRef> getCommonName(const StringRef Name, bool KeepSuffix,
+                                       ArrayRef<StringRef> Suffixes);
 /// LTO-generated function names take a form:
 ///
 ///   <function_name>.lto_priv.<decimal_number>/...

bolt/lib/Rewrite/PseudoProbeRewriter.cpp

@@ -14,6 +14,7 @@
 #include "bolt/Rewrite/MetadataRewriter.h"
 #include "bolt/Rewrite/MetadataRewriters.h"
 #include "bolt/Utils/CommandLineOpts.h"
+#include "bolt/Utils/Utils.h"
 #include "llvm/IR/Function.h"
 #include "llvm/MC/MCPseudoProbe.h"
 #include "llvm/Support/CommandLine.h"
@@ -133,10 +134,16 @@ void PseudoProbeRewriter::parsePseudoProbe() {
 
   MCPseudoProbeDecoder::Uint64Set GuidFilter;
   MCPseudoProbeDecoder::Uint64Map FuncStartAddrs;
+  SmallVector<StringRef, 3> Suffixes({".destroy", ".resume", ".llvm."});
   for (const BinaryFunction *F : BC.getAllBinaryFunctions()) {
     for (const MCSymbol *Sym : F->getSymbols()) {
-      FuncStartAddrs[Function::getGUID(NameResolver::restore(Sym->getName()))] =
-          F->getAddress();
+      StringRef SymName = NameResolver::restore(Sym->getName());
+      if (std::optional<StringRef> CommonName =
+              getCommonName(SymName, false, Suffixes)) {
+        SymName = *CommonName;
+      }
+      uint64_t GUID = Function::getGUID(SymName);
+      FuncStartAddrs[GUID] = F->getAddress();
     }
   }
   Contents = PseudoProbeSection->getContents();
@@ -155,13 +162,25 @@ void PseudoProbeRewriter::parsePseudoProbe() {
     ProbeDecoder.printProbesForAllAddresses(outs());
   }
 
-  for (const auto &FuncDesc : ProbeDecoder.getGUID2FuncDescMap()) {
-    uint64_t GUID = FuncDesc.FuncGUID;
-    if (!FuncStartAddrs.contains(GUID))
-      continue;
-    BinaryFunction *BF = BC.getBinaryFunctionAtAddress(FuncStartAddrs[GUID]);
-    assert(BF);
-    BF->setGUID(GUID);
+  const GUIDProbeFunctionMap &GUID2Func = ProbeDecoder.getGUID2FuncDescMap();
+  // Checks GUID in GUID2Func and returns it if it's present or null otherwise.
+  auto checkGUID = [&](StringRef SymName) {
+    uint64_t GUID = Function::getGUID(SymName);
+    if (GUID2Func.find(GUID) == GUID2Func.end())
+      return 0ull;
+    return GUID;
+  };
+  for (BinaryFunction *F : BC.getAllBinaryFunctions()) {
+    for (const MCSymbol *Sym : F->getSymbols()) {
+      StringRef SymName = NameResolver::restore(Sym->getName());
+      uint64_t GUID = checkGUID(SymName);
+      std::optional<StringRef> CommonName =
+          getCommonName(SymName, false, Suffixes);
+      if (!GUID && CommonName)
+        GUID = checkGUID(*CommonName);
+      if (GUID)
+        F->setGUID(GUID);
+    }
   }
 }
 

bolt/lib/Utils/Utils.cpp

@@ -66,15 +66,21 @@ std::string getUnescapedName(const StringRef &Name) {
   return Output;
 }
 
-std::optional<StringRef> getLTOCommonName(const StringRef Name) {
-  for (StringRef Suffix : {".__uniq.", ".lto_priv.", ".constprop.", ".llvm."}) {
+std::optional<StringRef> getCommonName(const StringRef Name, bool KeepSuffix,
+                                       ArrayRef<StringRef> Suffixes) {
+  for (StringRef Suffix : Suffixes) {
     size_t LTOSuffixPos = Name.find(Suffix);
     if (LTOSuffixPos != StringRef::npos)
-      return Name.substr(0, LTOSuffixPos + Suffix.size());
+      return Name.substr(0, LTOSuffixPos + (KeepSuffix ? Suffix.size() : 0));
   }
   return std::nullopt;
 }
 
+std::optional<StringRef> getLTOCommonName(const StringRef Name) {
+  return getCommonName(Name, true,
+                       {".__uniq.", ".lto_priv.", ".constprop.", ".llvm."});
+}
+
 std::optional<uint8_t> readDWARFExpressionTargetReg(StringRef ExprBytes) {
   uint8_t Opcode = ExprBytes[0];
   if (Opcode == dwarf::DW_CFA_def_cfa_expression)