[SampleFDO] Support enabling -funique-internal-linkage-name.

now -funique-internal-linkage-name flag is available, and we want to flip
it on by default since it is beneficial to have separate sample profiles
for different internal symbols with the same name. As a preparation, we
want to avoid regression caused by the flip.

When we flip -funique-internal-linkage-name on, the profile is collected
from binary built without -funique-internal-linkage-name so it has no uniq
suffix, but the IR in the optimized build contains the suffix. This kind of
mismatch may introduce transient regression.

To avoid such mismatch, we introduce a NameTable section flag indicating
whether there is any name in the profile containing uniq suffix. Compiler
will decide whether to keep uniq suffix during name canonicalization
depending on the NameTable section flag. The flag is only available for
extbinary format. For other formats, by default compiler will keep uniq
suffix so they will only experience transient regression when
-funique-internal-linkage-name is just flipped.

Another type of regression is caused by places where we miss to call
getCanonicalFnName. Those places are fixed.

Differential Revision: https://reviews.llvm.org/D96932

llvm/include/llvm/ProfileData/SampleProf.h

@@ -177,7 +177,10 @@ enum class SecNameTableFlags : uint32_t {
   SecFlagMD5Name = (1 << 0),
   // Store MD5 in fixed length instead of ULEB128 so NameTable can be
   // accessed like an array.
-  SecFlagFixedLengthMD5 = (1 << 1)
+  SecFlagFixedLengthMD5 = (1 << 1),
+  // Profile contains ".__uniq." suffix name. Compiler shouldn't strip
+  // the suffix when doing profile matching when seeing the flag.
+  SecFlagUniqSuffix = (1 << 2)
 };
 enum class SecProfSummaryFlags : uint32_t {
   SecFlagInValid = 0,
@@ -728,13 +731,14 @@ class FunctionSamples {
   /// GUID to \p S. Also traverse the BodySamples to add hot CallTarget's GUID
   /// to \p S.
   void findInlinedFunctions(DenseSet<GlobalValue::GUID> &S, const Module *M,
+                            const StringMap<Function *> &SymbolMap,
                             uint64_t Threshold) const {
     if (TotalSamples <= Threshold)
       return;
     auto isDeclaration = [](const Function *F) {
       return !F || F->isDeclaration();
     };
-    if (isDeclaration(M->getFunction(getFuncName()))) {
+    if (isDeclaration(SymbolMap.lookup(getFuncName()))) {
       // Add to the import list only when it's defined out of module.
       S.insert(getGUID(Name));
     }
@@ -743,13 +747,13 @@ class FunctionSamples {
     for (const auto &BS : BodySamples)
       for (const auto &TS : BS.second.getCallTargets())
         if (TS.getValue() > Threshold) {
-          const Function *Callee = M->getFunction(getFuncName(TS.getKey()));
+          const Function *Callee = SymbolMap.lookup(getFuncName(TS.getKey()));
           if (isDeclaration(Callee))
             S.insert(getGUID(TS.getKey()));
         }
     for (const auto &CS : CallsiteSamples)
       for (const auto &NameFS : CS.second)
-        NameFS.second.findInlinedFunctions(S, M, Threshold);
+        NameFS.second.findInlinedFunctions(S, M, SymbolMap, Threshold);
   }
 
   /// Set the name of the function.
@@ -780,17 +784,31 @@ class FunctionSamples {
     return getCanonicalFnName(F.getName(), Attr);
   }
 
-  static StringRef getCanonicalFnName(StringRef FnName, StringRef Attr = "") {
-    static const char *knownSuffixes[] = { ".llvm.", ".part." };
+  /// Name suffixes which canonicalization should handle to avoid
+  /// profile mismatch.
+  static constexpr const char *LLVMSuffix = ".llvm.";
+  static constexpr const char *PartSuffix = ".part.";
+  static constexpr const char *UniqSuffix = ".__uniq.";
+
+  static StringRef getCanonicalFnName(StringRef FnName,
+                                      StringRef Attr = "selected") {
+    // 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};
     if (Attr == "" || Attr == "all") {
       return FnName.split('.').first;
     } else if (Attr == "selected") {
       StringRef Cand(FnName);
       for (const auto &Suf : knownSuffixes) {
         StringRef Suffix(Suf);
+        // If the profile contains ".__uniq." suffix, don't strip the
+        // suffix for names in the IR.
+        if (Suffix == UniqSuffix && FunctionSamples::HasUniqSuffix)
+          continue;
         auto It = Cand.rfind(Suffix);
         if (It == StringRef::npos)
-          return Cand;
+          continue;
         auto Dit = Cand.rfind('.');
         if (Dit == It + Suffix.size() - 1)
           Cand = Cand.substr(0, It);
@@ -861,6 +879,9 @@ class FunctionSamples {
   /// Whether the profile uses MD5 to represent string.
   static bool UseMD5;
 
+  /// Whether the profile contains any ".__uniq." suffix in a name.
+  static bool HasUniqSuffix;
+
   /// GUIDToFuncNameMap saves the mapping from GUID to the symbol name, for
   /// all the function symbols defined or declared in current module.
   DenseMap<uint64_t, StringRef> *GUIDToFuncNameMap = nullptr;

llvm/include/llvm/ProfileData/SampleProfReader.h

@@ -363,7 +363,11 @@ class SampleProfileReader {
   /// Print the profile for \p FName on stream \p OS.
   void dumpFunctionProfile(StringRef FName, raw_ostream &OS = dbgs());
 
-  virtual void collectFuncsFrom(const Module &M) {}
+  /// Collect functions with definitions in Module M. For reader which
+  /// support loading function profiles on demand, return true when the
+  /// reader has been given a module. Always return false for reader
+  /// which doesn't support loading function profiles on demand.
+  virtual bool collectFuncsFromModule() { return false; }
 
   /// Print all the profiles on stream \p OS.
   void dump(raw_ostream &OS = dbgs());
@@ -454,9 +458,13 @@ class SampleProfileReader {
   /// Don't read profile without context if the flag is set. This is only meaningful
   /// for ExtBinary format.
   virtual void setSkipFlatProf(bool Skip) {}
+  /// Return whether any name in the profile contains ".__uniq." suffix.
+  virtual bool hasUniqSuffix() { return false; }
 
   SampleProfileReaderItaniumRemapper *getRemapper() { return Remapper.get(); }
 
+  void setModule(const Module *Mod) { M = Mod; }
+
 protected:
   /// Map every function to its associated profile.
   ///
@@ -496,6 +504,11 @@ class SampleProfileReader {
 
   /// \brief The format of sample.
   SampleProfileFormat Format = SPF_None;
+
+  /// \brief The current module being compiled if SampleProfileReader
+  /// is used by compiler. If SampleProfileReader is used by other
+  /// tools which are not compiler, M is usually nullptr.
+  const Module *M = nullptr;
 };
 
 class SampleProfileReaderText : public SampleProfileReader {
@@ -656,8 +669,6 @@ class SampleProfileReaderExtBinaryBase : public SampleProfileReaderBinary {
   DenseMap<StringRef, uint64_t> FuncOffsetTable;
   /// The set containing the functions to use when compiling a module.
   DenseSet<StringRef> FuncsToUse;
-  /// Use all functions from the input profile.
-  bool UseAllFuncs = true;
 
   /// Use fixed length MD5 instead of ULEB128 encoding so NameTable doesn't
   /// need to be read in up front and can be directly accessed using index.
@@ -692,8 +703,9 @@ class SampleProfileReaderExtBinaryBase : public SampleProfileReaderBinary {
   uint64_t getFileSize();
   virtual bool dumpSectionInfo(raw_ostream &OS = dbgs()) override;
 
-  /// Collect functions with definitions in Module \p M.
-  void collectFuncsFrom(const Module &M) override;
+  /// Collect functions with definitions in Module M. Return true if
+  /// the reader has been given a module.
+  bool collectFuncsFromModule() override;
 
   /// Return whether names in the profile are all MD5 numbers.
   virtual bool useMD5() override { return MD5StringBuf.get(); }
@@ -731,8 +743,6 @@ class SampleProfileReaderCompactBinary : public SampleProfileReaderBinary {
   DenseMap<StringRef, uint64_t> FuncOffsetTable;
   /// The set containing the functions to use when compiling a module.
   DenseSet<StringRef> FuncsToUse;
-  /// Use all functions from the input profile.
-  bool UseAllFuncs = true;
   virtual std::error_code verifySPMagic(uint64_t Magic) override;
   virtual std::error_code readNameTable() override;
   /// Read a string indirectly via the name table.
@@ -751,8 +761,9 @@ class SampleProfileReaderCompactBinary : public SampleProfileReaderBinary {
   /// Read samples only for functions to use.
   std::error_code readImpl() override;
 
-  /// Collect functions to be used when compiling Module \p M.
-  void collectFuncsFrom(const Module &M) override;
+  /// Collect functions with definitions in Module M. Return true if
+  /// the reader has been given a module.
+  bool collectFuncsFromModule() override;
 
   /// Return whether names in the profile are all MD5 numbers.
   virtual bool useMD5() override { return true; }

llvm/lib/ProfileData/SampleProf.cpp

@@ -40,7 +40,8 @@ namespace sampleprof {
 SampleProfileFormat FunctionSamples::Format;
 bool FunctionSamples::ProfileIsProbeBased = false;
 bool FunctionSamples::ProfileIsCS = false;
-bool FunctionSamples::UseMD5;
+bool FunctionSamples::UseMD5 = false;
+bool FunctionSamples::HasUniqSuffix = true;
 } // namespace sampleprof
 } // namespace llvm
 
@@ -262,6 +263,8 @@ void FunctionSamples::findAllNames(DenseSet<StringRef> &NameSet) const {
 const FunctionSamples *FunctionSamples::findFunctionSamplesAt(
     const LineLocation &Loc, StringRef CalleeName,
     SampleProfileReaderItaniumRemapper *Remapper) const {
+  CalleeName = getCanonicalFnName(CalleeName);
+
   std::string CalleeGUID;
   CalleeName = getRepInFormat(CalleeName, UseMD5, CalleeGUID);
 

llvm/lib/ProfileData/SampleProfReader.cpp

@@ -584,6 +584,8 @@ std::error_code SampleProfileReaderExtBinaryBase::readOneSection(
     bool UseMD5 = hasSecFlag(Entry, SecNameTableFlags::SecFlagMD5Name);
     assert((!FixedLengthMD5 || UseMD5) &&
            "If FixedLengthMD5 is true, UseMD5 has to be true");
+    FunctionSamples::HasUniqSuffix =
+        hasSecFlag(Entry, SecNameTableFlags::SecFlagUniqSuffix);
     if (std::error_code EC = readNameTableSec(UseMD5))
       return EC;
     break;
@@ -615,11 +617,13 @@ std::error_code SampleProfileReaderExtBinaryBase::readOneSection(
   return sampleprof_error::success;
 }
 
-void SampleProfileReaderExtBinaryBase::collectFuncsFrom(const Module &M) {
-  UseAllFuncs = false;
+bool SampleProfileReaderExtBinaryBase::collectFuncsFromModule() {
+  if (!M)
+    return false;
   FuncsToUse.clear();
-  for (auto &F : M)
+  for (auto &F : *M)
     FuncsToUse.insert(FunctionSamples::getCanonicalFnName(F));
+  return true;
 }
 
 std::error_code SampleProfileReaderExtBinaryBase::readFuncOffsetTable() {
@@ -648,14 +652,24 @@ std::error_code SampleProfileReaderExtBinaryBase::readFuncOffsetTable() {
 }
 
 std::error_code SampleProfileReaderExtBinaryBase::readFuncProfiles() {
+  // Collect functions used by current module if the Reader has been
+  // given a module.
+  // collectFuncsFromModule uses FunctionSamples::getCanonicalFnName
+  // which will query FunctionSamples::HasUniqSuffix, so it has to be
+  // called after FunctionSamples::HasUniqSuffix is set, i.e. after
+  // NameTable section is read.
+  bool LoadFuncsToBeUsed = collectFuncsFromModule();
+
+  // When LoadFuncsToBeUsed is false, load all the function profiles.
   const uint8_t *Start = Data;
-  if (UseAllFuncs) {
+  if (!LoadFuncsToBeUsed) {
     while (Data < End) {
       if (std::error_code EC = readFuncProfile(Data))
         return EC;
     }
     assert(Data == End && "More data is read than expected");
   } else {
+    // Load function profiles on demand.
     if (Remapper) {
       for (auto Name : FuncsToUse) {
         Remapper->insert(Name);
@@ -688,7 +702,6 @@ std::error_code SampleProfileReaderExtBinaryBase::readFuncProfiles() {
     }
     Data = End;
   }
-
   assert((CSProfileCount == 0 || CSProfileCount == Profiles.size()) &&
          "Cannot have both context-sensitive and regular profile");
   ProfileIsCS = (CSProfileCount > 0);
@@ -783,13 +796,18 @@ std::error_code SampleProfileReaderExtBinaryBase::readImpl() {
 }
 
 std::error_code SampleProfileReaderCompactBinary::readImpl() {
+  // Collect functions used by current module if the Reader has been
+  // given a module.
+  bool LoadFuncsToBeUsed = collectFuncsFromModule();
+
   std::vector<uint64_t> OffsetsToUse;
-  if (UseAllFuncs) {
+  if (!LoadFuncsToBeUsed) {
+    // load all the function profiles.
     for (auto FuncEntry : FuncOffsetTable) {
       OffsetsToUse.push_back(FuncEntry.second);
     }
-  }
-  else {
+  } else {
+    // load function profiles on demand.
     for (auto Name : FuncsToUse) {
       auto GUID = std::to_string(MD5Hash(Name));
       auto iter = FuncOffsetTable.find(StringRef(GUID));
@@ -1010,6 +1028,8 @@ static std::string getSecFlagsStr(const SecHdrTableEntry &Entry) {
       Flags.append("fixlenmd5,");
     else if (hasSecFlag(Entry, SecNameTableFlags::SecFlagMD5Name))
       Flags.append("md5,");
+    if (hasSecFlag(Entry, SecNameTableFlags::SecFlagUniqSuffix))
+      Flags.append("uniq,");
     break;
   case SecProfSummary:
     if (hasSecFlag(Entry, SecProfSummaryFlags::SecFlagPartial))
@@ -1117,11 +1137,13 @@ std::error_code SampleProfileReaderCompactBinary::readFuncOffsetTable() {
   return sampleprof_error::success;
 }
 
-void SampleProfileReaderCompactBinary::collectFuncsFrom(const Module &M) {
-  UseAllFuncs = false;
+bool SampleProfileReaderCompactBinary::collectFuncsFromModule() {
+  if (!M)
+    return false;
   FuncsToUse.clear();
-  for (auto &F : M)
+  for (auto &F : *M)
     FuncsToUse.insert(FunctionSamples::getCanonicalFnName(F));
+  return true;
 }
 
 std::error_code SampleProfileReaderBinary::readSummaryEntry(

llvm/lib/ProfileData/SampleProfWriter.cpp

@@ -204,6 +204,17 @@ std::error_code SampleProfileWriterExtBinaryBase::writeNameTableSection(
     addName(I.first());
     addNames(I.second);
   }
+
+  // If NameTable contains ".__uniq." suffix, set SecFlagUniqSuffix flag
+  // so compiler won't strip the suffix during profile matching after
+  // seeing the flag in the profile.
+  for (const auto &I : NameTable) {
+    if (I.first.find(FunctionSamples::UniqSuffix) != StringRef::npos) {
+      addSectionFlag(SecNameTable, SecNameTableFlags::SecFlagUniqSuffix);
+      break;
+    }
+  }
+
   if (auto EC = writeNameTable())
     return EC;
   return sampleprof_error::success;

llvm/lib/Transforms/IPO/SampleContextTracker.cpp

@@ -199,6 +199,8 @@ SampleContextTracker::getCalleeContextSamplesFor(const CallBase &Inst,
   if (!DIL)
     return nullptr;
 
+  CalleeName = FunctionSamples::getCanonicalFnName(CalleeName);
+
   // For indirect call, CalleeName will be empty, in which case the context
   // profile for callee with largest total samples will be returned.
   ContextTrieNode *CalleeContext = getCalleeContextFor(DIL, CalleeName);