[SampleFDO] Fix a crash when the sample profile uses md5 and -sample-profile-merge-inlinee is enabled

[wmi]

Summary

When -sample-profile-merge-inlinee is enabled, new FunctionSamples may be created during profile merge without GUIDToFuncNameMap being initialized. That will occasionally cause compiler crash. The patch fixes it.

[davidxl]

lgtm

llvm/lib/Transforms/IPO/SampleProfile.cpp

238

nit: name it UpdateFunctionSamples

[wenlei]

Just realized this is fixing the same issue as @hoyFB's D84997 - we ran into it immediately after -sample-profile-merge-inlinee is turned on too.

It looks to me that we only need to transfer GUIDToFuncNameMap during merge as the map is shared, so it can be done at merge time like D84997. What do you think? That's the fixed we used internally.

[wmi]

Yes, the fix in D84997 is more concise. Here I use a straightforward fix. It is ok for Hongtao to commit D84997, but just add a testcase (free to use the one here).

[wmi]

Commit the fix for the second problem in D84997 together with the assertion and testcase in this patch, so D84997 can be dedicated to solve the other problem.

llvm/lib/Transforms/IPO/SampleProfile.cpp

private:
   /// 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;
 
   SampleProfileLoader &SPLoader;
 };
 
+/// Set GUIDToFuncNameMap for FunctionSamples in FSToUpdate and all the
+/// callsite FunctionSamples recursively.
+void SetGUIDToFuncNameMapImpl(std::queue<FunctionSamples *> &FSToUpdate,
+                              DenseMap<uint64_t, StringRef> *Map) {
+  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);
+      }
+    }
+  }
+}
+
 class GUIDToFuncNameMapper {
 public:
   GUIDToFuncNameMapper(Module &M, SampleProfileReader &Reader,
                         DenseMap<uint64_t, StringRef> &GUIDToFuncNameMap)
       : CurrentReader(Reader), CurrentModule(M),
       CurrentGUIDToFuncNameMap(GUIDToFuncNameMap) {
     if (!CurrentReader.useMD5())
       return;
public:
   }
 
 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);
-        }
-      }
-    }
+    SetGUIDToFuncNameMapImpl(FSToUpdate, Map);
   }
 
   SampleProfileReader &CurrentReader;
   Module &CurrentModule;
   DenseMap<uint64_t, StringRef> &CurrentGUIDToFuncNameMap;
 };
 
 /// Sample profile pass.
while (true) {
     for (auto &BB : F) {
       bool Hot = false;
       SmallVector<CallBase *, 10> AllCandidates;
       SmallVector<CallBase *, 10> ColdCandidates;
       for (auto &I : BB.getInstList()) {
         const FunctionSamples *FS = nullptr;
         if (auto *CB = dyn_cast<CallBase>(&I)) {
           if (!isa<IntrinsicInst>(I) && (FS = findCalleeFunctionSamples(*CB))) {
+            assert((!FunctionSamples::UseMD5 || FS->GUIDToFuncNameMap) &&
+                   "GUIDToFuncNameMap has to be populated");
             AllCandidates.push_back(CB);
             if (FS->getEntrySamples() > 0)
               localNotInlinedCallSites.try_emplace(CB, FS);
             if (callsiteIsHot(FS, PSI))
               Hot = true;
             else if (shouldInlineColdCallee(*CB))
               ColdCandidates.push_back(CB);
           }
if (ProfileMergeInlinee) {
       assert(FS->getHeadSamples() == 0 && "Expect 0 head sample for inlinee");
       const_cast<FunctionSamples *>(FS)->addHeadSamples(FS->getEntrySamples());
 
       // Note that we have to do the merge right after processing function.
       // This allows OutlineFS's profile to be used for annotation during
       // top-down processing of functions' annotation.
       FunctionSamples *OutlineFS = Reader->getOrCreateSamplesFor(*Callee);
       OutlineFS->merge(*FS);
+
+      // Recursively setup GUIDToFuncNameMap in OutlineFS.
+      std::queue<FunctionSamples *> FSToUpdate;
+      FSToUpdate.push(OutlineFS);
+      SetGUIDToFuncNameMapImpl(FSToUpdate, &GUIDToFuncNameMap);
     } else {
       auto pair =
           notInlinedCallInfo.try_emplace(Callee, NotInlinedProfileInfo{0});
       pair.first->second.entryCount += FS->getEntrySamples();
     }
   }
   return Changed;
 }

llvm/test/Transforms/SampleProfile/inline-mergeprof.ll

 ; Test we lose details of not inlined profile without '-sample-profile-merge-inlinee'
 ; RUN: opt < %s -sample-profile -sample-profile-file=%S/Inputs/inline-mergeprof.prof -sample-profile-merge-inlinee=false -S | FileCheck -check-prefix=SCALE %s
 ; RUN: opt < %s -sample-profile -sample-profile-file=%S/Inputs/inline-mergeprof.prof -sample-profile-merge-inlinee=true -S | FileCheck -check-prefix=SCALE %s
 ; RUN: opt < %s -passes=sample-profile -sample-profile-file=%S/Inputs/inline-mergeprof.prof -sample-profile-merge-inlinee=false -S | FileCheck -check-prefix=SCALE %s
 
-; Test we properly merge not inlined profile properly with '-sample-profile-merge-inlinee'
+; Test we properly merge not inlined profile with '-sample-profile-merge-inlinee'
 ; RUN: opt < %s -passes=sample-profile -sample-profile-file=%S/Inputs/inline-mergeprof.prof -sample-profile-merge-inlinee=true -S | FileCheck -check-prefix=MERGE  %s
 
+; Test we properly merge not inlined profile with '-sample-profile-merge-inlinee'
+; when the profile uses md5.
+; RUN: opt < %s -passes=sample-profile -sample-profile-file=%S/Inputs/inline-mergeprof.md5.prof -sample-profile-merge-inlinee=true -S | FileCheck -check-prefix=MERGE  %s
+
 @.str = private unnamed_addr constant [11 x i8] c"sum is %d\0A\00", align 1
 
 define i32 @main() #0 !dbg !6 {
 entry:
   %retval = alloca i32, align 4
   %s = alloca i32, align 4
   %i = alloca i32, align 4
   %tmp = load i32, i32* %i, align 4, !dbg !8