[ThinLTO][AutoFDO] Fix memory corruption due to race condition from thin backends

[wenlei]

Summary

This commit fixed a race condition from multi-threaded thinLTO backends that causes non-deterministic memory corruption for a data structure used only by AutoFDO with compact binary profile.
GUIDToFuncNameMap, a static data member of type DenseMap in FunctionSamples is used as a per-module mapping from function name MD5 to name string when input AutoFDO profile is in compact binary format. However with ThinLTO, we can have parallel backends modifying and accessing the class static map concurrently. The fix is to mark GUIDToFuncNameMap as thread_local.

[wmi]

Thanks for working on the issue.

From reading the comment of LLVM_THREAD_LOCAL in include/llvm/Support/Compiler.h, it seems we should only use TLS for POD if we want it to be functioning as we expect on all platforms.

How about change GUIDToFuncNameMap and CurrentModule as member fields in class GUIDToFuncNameMapper so we can create them separately for each module? Add a unique_ptr of GUIDToFuncNameMapper in SampleProfileLoader and a pointer field of type "GUIDToFuncNameMapper *" for each FunctionSamples object, so FunctionSamples object can access GUIDToFuncNameMap and CurrentModule via the pointer of GUIDToFuncNameMapper. The initialization of GUIDToFuncNameMapper pointer of each FunctionSamples object can be done in GUIDToFuncNameMapper's constructor (need to pass Reader as a param to get all the FunctionSamples).

[wenlei]

In D65848#1621227, @wmi wrote:

Thanks for working on the issue.

From reading the comment of LLVM_THREAD_LOCAL in include/llvm/Support/Compiler.h, it seems we should only use TLS for POD if we want it to be functioning as we expect on all platforms.

How about change GUIDToFuncNameMap and CurrentModule as member fields in class GUIDToFuncNameMapper so we can create them separately for each module? Add a unique_ptr of GUIDToFuncNameMapper in SampleProfileLoader and a pointer field of type "GUIDToFuncNameMapper *" for each FunctionSamples object, so FunctionSamples object can access GUIDToFuncNameMap and CurrentModule via the pointer of GUIDToFuncNameMapper. The initialization of GUIDToFuncNameMapper pointer of each FunctionSamples object can be done in GUIDToFuncNameMapper's constructor (need to pass Reader as a param to get all the FunctionSamples).

Thanks for review. Using thread_local is the "minimum" fix, I actually implemented something similar to what you suggested (under testing), except I let GUIDToFuncNameMap be a member of SampleProfileLoader, and keep a reference in GUIDToFuncNameMapper. That way we can avoid TLS overhead. I will update once testing is done.

[wenlei]

avoid TLS and move GUIDToFuncNameMapper to class member

[wmi]

Thanks for the change.

llvm/include/llvm/ProfileData/SampleProf.h

383

If we can avoid adding the param, that will be great. GUIDToFuncNameMap is only used when the format is compact. When the interface of findInlinedFunctions is used, I hope the user doesn't have to worry about how to create a GUIDToFuncNameMap. That is why I think we may pass the handle of GUIDToFuncNameMap to FunctionSamples objects in GUIDToFuncNameMapper constructor.

[wenlei]

let GUIDToFuncNameMapper populate GUIDToFuncNameMap field of FunctionSamples.

[wenlei]

llvm/include/llvm/ProfileData/SampleProf.h

383

Agreed it's better to avoid this extra param. Changed to populate GUIDToFuncNameMap for FunctionSamples in GUIDToFuncNameMapper's ctor now. We have to traverse all inline trees though to cover all FunctionSamples, which is a bit heavy and why I didn't do this in previous version. But this is cleaner for sure - thanks! I also moved GUIDToFuncNameMapper out of FunctionSamples.

[wmi]

LGTM.

llvm/include/llvm/ProfileData/SampleProf.h

451

Please add some assertion message.

[wenlei]

add message for assertion.

llvm/include/llvm/ProfileData/SampleProf.h

sampleprof_error merge(const FunctionSamples &Other, uint64_t Weight = 1) {
     return Result;
   }
 
   /// Recursively traverses all children, if the total sample count of the
   /// corresponding function is no less than \p Threshold, add its corresponding
   /// 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,
                             uint64_t Threshold) const {
     if (TotalSamples <= Threshold)
       return;
     S.insert(getGUID(Name));
     // Import hot CallTargets, which may not be available in IR because full
     // profile annotation cannot be done until backend compilation in ThinLTO.
     for (const auto &BS : BodySamples)
       for (const auto &TS : BS.second.getCallTargets())
         if (TS.getValue() > Threshold) {
public:
   /// When the Format is not SPF_Compact_Binary, \p Name needs no translation.
   /// When the Format is SPF_Compact_Binary, \p Name in current FunctionSamples
   /// is actually GUID of the original function name. getNameInModule will
   /// translate \p Name in current FunctionSamples into its original name.
   /// If the original name doesn't exist in \p M, return empty StringRef.
   StringRef getNameInModule(StringRef Name, const Module *M) const {
     if (Format != SPF_Compact_Binary)
       return Name;
-    // Expect CurrentModule to be initialized by GUIDToFuncNameMapper.
-    if (M != CurrentModule)
-      llvm_unreachable("Input Module should be the same as CurrentModule");
-    auto iter = GUIDToFuncNameMap.find(std::stoull(Name.data()));
-    if (iter == GUIDToFuncNameMap.end())
+
+    assert(GUIDToFuncNameMap && "GUIDToFuncNameMap needs to be popluated first");
+    auto iter = GUIDToFuncNameMap->find(std::stoull(Name.data()));
+    if (iter == GUIDToFuncNameMap->end())
       return StringRef();
     return iter->second;
   }
 
   /// Returns the line offset to the start line of the subprogram.
   /// We assume that a single function will not exceed 65535 LOC.
   static unsigned getOffset(const DILocation *DIL);
 
   /// Get the FunctionSamples of the inline instance where DIL originates
   /// from.
   ///
   /// The FunctionSamples of the instruction (Machine or IR) associated to
   /// \p DIL is the inlined instance in which that instruction is coming from.
   /// We traverse the inline stack of that instruction, and match it with the
   /// tree nodes in the profile.
   ///
   /// \returns the FunctionSamples pointer to the inlined instance.
   const FunctionSamples *findFunctionSamples(const DILocation *DIL) const;
 
   static SampleProfileFormat Format;
-  /// GUIDToFuncNameMap saves the mapping from GUID to the symbol name, for
-  /// all the function symbols defined or declared in CurrentModule.
-  static DenseMap<uint64_t, StringRef> GUIDToFuncNameMap;
-  static Module *CurrentModule;
-
-  class GUIDToFuncNameMapper {
-  public:
-    GUIDToFuncNameMapper(Module &M) {
-      if (Format != SPF_Compact_Binary)
-        return;
 
-      for (const auto &F : M) {
-        StringRef OrigName = F.getName();
-        GUIDToFuncNameMap.insert({Function::getGUID(OrigName), OrigName});
-        /// Local to global var promotion used by optimization like thinlto
-        /// will rename the var and add suffix like ".llvm.xxx" to the
-        /// original local name. In sample profile, the suffixes of function
-        /// names are all stripped. Since it is possible that the mapper is
-        /// built in post-thin-link phase and var promotion has been done,
-        /// we need to add the substring of function name without the suffix
-        /// into the GUIDToFuncNameMap.
-        StringRef CanonName = getCanonicalFnName(F);
-        if (CanonName != OrigName)
-          GUIDToFuncNameMap.insert({Function::getGUID(CanonName), CanonName});
-      }
-      CurrentModule = &M;
-    }
-
-    ~GUIDToFuncNameMapper() {
-      if (Format != SPF_Compact_Binary)
-        return;
-
-      GUIDToFuncNameMap.clear();
-      CurrentModule = nullptr;
-    }
-  };
+  /// 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;
 
   // Assume the input \p Name is a name coming from FunctionSamples itself.
   // If the format is SPF_Compact_Binary, the name is already a GUID and we
   // don't want to return the GUID of GUID.
   static uint64_t getGUID(StringRef Name) {
     return (Format == SPF_Compact_Binary) ? std::stoull(Name.data())
                                           : Function::getGUID(Name);
   }

llvm/lib/ProfileData/SampleProf.cpp

 #include <system_error>
 
 using namespace llvm;
 using namespace sampleprof;
 
 namespace llvm {
 namespace sampleprof {
 SampleProfileFormat FunctionSamples::Format;
-DenseMap<uint64_t, StringRef> FunctionSamples::GUIDToFuncNameMap;
-Module *FunctionSamples::CurrentModule;
 } // namespace sampleprof
 } // namespace llvm
 
 namespace {
 
 // FIXME: This class is only here to support the transition to llvm::Error. It
 // will be removed once this transition is complete. Clients should prefer to
 // deal with the Error value directly, rather than converting to error_code.

llvm/lib/Transforms/IPO/SampleProfile.cpp

 #include "llvm/Transforms/Utils/Cloning.h"
 #include <algorithm>
 #include <cassert>
 #include <cstdint>
 #include <functional>
 #include <limits>
 #include <map>
 #include <memory>
+#include <queue>
 #include <string>
 #include <system_error>
 #include <utility>
 #include <vector>
 
 using namespace llvm;
 using namespace sampleprof;
 using ProfileCount = Function::ProfileCount;
private:
   ///
   /// Note that this accumulator tracks samples used from a single function
   /// and all the inlined callsites. Strictly, we should have a map of counters
   /// keyed by FunctionSamples pointers, but these stats are cleared after
   /// every function, so we just need to keep a single counter.
   uint64_t TotalUsedSamples = 0;
 };
 
+class GUIDToFuncNameMapper {
+public:
+  GUIDToFuncNameMapper(Module &M, SampleProfileReader &Reader,
+                        DenseMap<uint64_t, StringRef> &GUIDToFuncNameMap)
+      : CurrentReader(Reader), CurrentModule(M), 
+      CurrentGUIDToFuncNameMap(GUIDToFuncNameMap) {
+    if (CurrentReader.getFormat() != SPF_Compact_Binary)
+      return;
+
+    for (const auto &F : CurrentModule) {
+      StringRef OrigName = F.getName();
+      CurrentGUIDToFuncNameMap.insert(
+          {Function::getGUID(OrigName), OrigName});
+
+      // Local to global var promotion used by optimization like thinlto
+      // will rename the var and add suffix like ".llvm.xxx" to the
+      // original local name. In sample profile, the suffixes of function
+      // names are all stripped. Since it is possible that the mapper is
+      // built in post-thin-link phase and var promotion has been done,
+      // we need to add the substring of function name without the suffix
+      // into the GUIDToFuncNameMap.
+      StringRef CanonName = FunctionSamples::getCanonicalFnName(F);
+      if (CanonName != OrigName)
+        CurrentGUIDToFuncNameMap.insert(
+            {Function::getGUID(CanonName), CanonName});
+    }
+
+    // Update GUIDToFuncNameMap for each function including inlinees.
+    SetGUIDToFuncNameMapForAll(&CurrentGUIDToFuncNameMap);
+  }
+
+  ~GUIDToFuncNameMapper() {
+    if (CurrentReader.getFormat() != SPF_Compact_Binary)
+      return;
+
+    CurrentGUIDToFuncNameMap.clear();
+
+    // Reset GUIDToFuncNameMap for of each function as they're no
+    // longer valid at this point.
+    SetGUIDToFuncNameMapForAll(nullptr);
+  }
+
+private:
+  void SetGUIDToFuncNameMapForAll(DenseMap<uint64_t, StringRef> *Map) {
+    std::queue<FunctionSamples *> FSToUpdate;
+    for (auto &IFS : CurrentReader.getProfiles()) {
+      FSToUpdate.push(&IFS.second);
+    }
+
+    while (!FSToUpdate.empty()) {
+      FunctionSamples *FS = FSToUpdate.front();
+      FSToUpdate.pop();
+      FS->GUIDToFuncNameMap = Map;
+      for (const auto &ICS : FS->getCallsiteSamples()) {
+        const FunctionSamplesMap &FSMap = ICS.second;
+        for (auto &IFS : FSMap) {
+          FunctionSamples &FS = const_cast<FunctionSamples &>(IFS.second);
+          FSToUpdate.push(&FS);
+        }
+      }
+    }
+  }
+
+  SampleProfileReader &CurrentReader;
+  Module &CurrentModule;
+  DenseMap<uint64_t, StringRef> &CurrentGUIDToFuncNameMap;
+};
+
 /// Sample profile pass.
 ///
 /// This pass reads profile data from the file specified by
 /// -sample-profile-file and annotates every affected function with the
 /// profile information found in that file.
 class SampleProfileLoader {
 public:
   SampleProfileLoader(
protected:
 
   // Information recorded when we declined to inline a call site
   // because we have determined it is too cold is accumulated for
   // each callee function. Initially this is just the entry count.
   struct NotInlinedProfileInfo {
     uint64_t entryCount;
   };
   DenseMap<Function *, NotInlinedProfileInfo> notInlinedCallInfo;
+
+  // 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;
 };
 
 class SampleProfileLoaderLegacyPass : public ModulePass {
 public:
   // Class identification, replacement for typeinfo
   static char ID;
 
   SampleProfileLoaderLegacyPass(StringRef Name = SampleProfileFile,

 }
 
 ModulePass *llvm::createSampleProfileLoaderPass(StringRef Name) {
   return new SampleProfileLoaderLegacyPass(Name);
 }
 
 bool SampleProfileLoader::runOnModule(Module &M, ModuleAnalysisManager *AM,
                                       ProfileSummaryInfo *_PSI) {
-  FunctionSamples::GUIDToFuncNameMapper Mapper(M);
+  GUIDToFuncNameMapper Mapper(M, *Reader, GUIDToFuncNameMap);
   if (!ProfileIsValid)
     return false;
 
   PSI = _PSI;
   if (M.getProfileSummary(/* IsCS */ false) == nullptr)
     M.setProfileSummary(Reader->getSummary().getMD(M.getContext()),
                         ProfileSummary::PSK_Sample);