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

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 make GUIDToFuncNameMap a member of SampleProfileLoader instead of a file static data.

Reviewers: wmi, davidxl, danielcdh

Subscribers: mehdi_amini, inglorion, hiraditya, dexonsmith, llvm-commits

Tags: #llvm

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

llvm-svn: 368596

llvm/include/llvm/ProfileData/SampleProf.h

@@ -447,11 +447,10 @@ public:
   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)
+    assert(GUIDToFuncNameMap && "GUIDToFuncNameMap needs to be popluated first");
-      llvm_unreachable("Input Module should be the same as CurrentModule");
+    auto iter = GUIDToFuncNameMap->find(std::stoull(Name.data()));
-    auto iter = GUIDToFuncNameMap.find(std::stoull(Name.data()));
+    if (iter == GUIDToFuncNameMap->end())
-    if (iter == GUIDToFuncNameMap.end())
       return StringRef();
     return iter->second;
   }
@@ -472,42 +471,10 @@ public:
   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.
+  /// all the function symbols defined or declared in current module.
-  static DenseMap<uint64_t, StringRef> GUIDToFuncNameMap;
+  DenseMap<uint64_t, StringRef> *GUIDToFuncNameMap = nullptr;
-  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;
-    }
-  };
 
   // 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

llvm/lib/ProfileData/SampleProf.cpp

@@ -28,8 +28,6 @@ using namespace sampleprof;
 namespace llvm {
 namespace sampleprof {
 SampleProfileFormat FunctionSamples::Format;
-DenseMap<uint64_t, StringRef> FunctionSamples::GUIDToFuncNameMap;
-Module *FunctionSamples::CurrentModule;
 } // namespace sampleprof
 } // namespace llvm
 

llvm/lib/Transforms/IPO/SampleProfile.cpp

@@ -79,6 +79,7 @@
 #include <limits>
 #include <map>
 #include <memory>
+#include <queue>
 #include <string>
 #include <system_error>
 #include <utility>
@@ -187,6 +188,74 @@ private:
   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
@@ -326,6 +395,10 @@ protected:
     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 {
@@ -1594,7 +1667,7 @@ ModulePass *llvm::createSampleProfileLoaderPass(StringRef Name) {
 
 bool SampleProfileLoader::runOnModule(Module &M, ModuleAnalysisManager *AM,
                                       ProfileSummaryInfo *_PSI) {
-  FunctionSamples::GUIDToFuncNameMapper Mapper(M);
+  GUIDToFuncNameMapper Mapper(M, *Reader, GUIDToFuncNameMap);
   if (!ProfileIsValid)
     return false;