[CSSPGO] Fix an invalid hash table reference issue in the CS preinliner.

[hoy]

Summary

We were using a StringMap object to store all profiles to be emitted. The object is basically an unordered hash table, therefore updating it in the process of trasvering it may cause issue since the underlying bucket array could change.

I'm also moving the csspgo-preinliner switch around so that no context tri will be constructed (by the constructor of CSPreInliner) when the switch is off.

[wlei]

Thanks for the fix!

llvm/lib/ProfileData/SampleProf.cpp

383–384

typo: Context

384

Do we need this?

384

typo: Context

[wenlei]

llvm/lib/ProfileData/SampleProf.cpp

373–374

Good catch - not sure how I come to incorrectly rely on the order of StringMap..

However the intention was to avoid extra copies of FunctionSamples as much as possible - these can large objects for copying.

I think what we could do is to collect a std::map<StringRef, StringRef> (old context to new context) for how we should update the StringMap, then iterate over the (top-down) ordered std::map to move profile to be keyed by correct context. This way we can avoid one extra copy for each FunctionSamples that needs to be moved.

[hoy]

llvm/lib/ProfileData/SampleProf.cpp

373–374

IIUC, we are trying to update the key of a promoted profile in ProfileMap. Since we cannot do that in place, we basically remove the current entry and add a new entry which just differs in the context key. If that's the case, not sure how to avoid copying a profile since the value field of an entry is not reference.

384

It's to make sure Ret is used in the release build where the following assert turns out nothing. Otherwise the ninja complaints Ret is defined by not used.

[wenlei]

llvm/lib/ProfileData/SampleProf.cpp

373–374

Here you're doing two copies, Map[A] -> temp -> Map[B]. With a std::map<StringRef, StringRef> for name mapping in top-down order, we should be able to do in-place copy Map[A] -> Map[B] by iterating over the name map?

[hoy]

llvm/lib/ProfileData/SampleProf.cpp

373–374

A and B are context strings, right? Why does the top-down order matter, say A is main->foo->bar and B is bar?

Yeah, we are using a intermediate temp. Should change ProfilesToBeAdded to StringMap<FunctionSamples&> and do the addition before the deletion.

[wlei]

llvm/lib/ProfileData/SampleProf.cpp

384

I see, how about we use find instead of try_emplace, like assert(ProfilesToBeAdded.find(ContextStr) == ProfilesToBeAdded.end() && "Context conflict during canonicalization");

[wenlei]

llvm/lib/ProfileData/SampleProf.cpp

373–374

This is my thought - two things blocks in-place copy: 1) iterating over StringMap while updating, which is solved by iterating over name mapping std::map<StringRef, StringRef>; 2) not overwriting items yet to be updated, this can be solved by ordering the updates so we make room for destination key of an update before we actually update.

[hoy]

llvm/lib/ProfileData/SampleProf.cpp

373–374

Still trying to follow. Say ProfileMap only has one entry {"main->foo->bar", Profile}. The context in Profile is "bar" and we are trying to update the key to "bar". With an extra std::map<StringRef, StringRef> , we would make something like {"main->foo->bar", "bar"}. Then how do we proceed next? Still figuring out how to avoid copying Profile.

384

That looks better but may slow down the assert compiler. Actually the (void) usage is quite common in the code base.

[wlei]

llvm/lib/ProfileData/SampleProf.cpp

384

Now I got it, thanks for the clarification.

[hoy]

Addressing Wenlei and Lei's comments.

[wenlei]

llvm/lib/ProfileData/SampleProf.cpp

373–374

So the idea is to avoid copying out to a temp map then copying back. Now I see that you changed the temp map to use ref type. In that case, it's equally good as there's going to be only one copy.

400

With try_emplace, we're expecting some of the insertions will fail? In that case, such profile will be lost. Perhaps we still need to order the insertion so such failure will not happen.

[hoy]

Addressing Wenlei's comment.

[hoy]

llvm/lib/ProfileData/SampleProf.cpp

400

Good point! Didn't realize that conflict could happen. Adjusted the insertion and deletion order.

[hoy]

I compared the performance of the new algorithm and the old implementation with extra copying with a decent large profile. It turned out the old implementation was better in speed. This is due to the fewer string hashing operations on long context string involved in the new implementation. While the new implementation could be better with shorter context string and bigger function profiles, I'm inclined to use the old one since it's easy to read and maintain.

[hoy]

Rolling back to first iteration.

[hoy]

Updating D104267: [CSSPGO] Fix an invalid hash table reference issue in the CS preinliner.

[wenlei]

In D104267#2826019, @hoy wrote:

I compared the performance of the new algorithm and the old implementation with extra copying with a decent large profile. It turned out the old implementation was better in speed. This is due to the fewer string hashing operations on long context string involved in the new implementation. While the new implementation could be better with shorter context string and bigger function profiles, I'm inclined to use the old one since it's easy to read and maintain.

Sounds good, thanks for giving it a try. LGTM with a minor comment.

llvm/lib/ProfileData/SampleProf.cpp

372–373

This no longer needs to be a set. I used set because of the erase operation. Now a vector would do.

373–374

May want to add a short comment explaining why an extra map (ProfilesToBeAdded) is needed, it seem not as trivial based on our discussions..

[hoy]

Addresing Wenlei's feedbacks.

llvm/lib/ProfileData/SampleProf.cpp

if (Ret.second) {
       FProfile.setContext(FContext);
       FProfile.setName(FContext.getNameWithoutContext());
     }
     FunctionSamples &OrigProfile = Ret.first->second;
     OrigProfile.merge(I.second);
   }
 }
 
 void SampleContextTrimmer::canonicalizeContextProfiles() {
-  StringSet<> ProfilesToBeRemoved;
-  // Note that StringMap order is guaranteed to be top-down order,
-  // this makes sure we make room for promoted/merged context in the
-  // map, before we move profiles in the map.
+  std::vector<StringRef> ProfilesToBeRemoved;
+  StringMap<FunctionSamples> ProfilesToBeAdded;
   for (auto &I : ProfileMap) {
     FunctionSamples &FProfile = I.second;
     StringRef ContextStr = FProfile.getNameWithContext();
     if (I.first() == ContextStr)
       continue;
 
     // Use the context string from FunctionSamples to update the keys of
     // ProfileMap. They can get out of sync after context profile promotion
     // through pre-inliner.
-    auto Ret = ProfileMap.try_emplace(ContextStr, FProfile);
-    assert(Ret.second && "Conext conflict during canonicalization");
-    FProfile = Ret.first->second;
-
-    // Track the context profile to remove
-    ProfilesToBeRemoved.erase(ContextStr);
-    ProfilesToBeRemoved.insert(I.first());
+    // Duplicate the function profile for later insertion to avoid a conflict
+    // caused by a context both to be add and to be removed. This could happen
+    // when a context is promoted to another context which is also promoted to
+    // the third context. For example, given an original context A @ B @ C that
+    // is promoted to B @ C and the original context B @ C which is promoted to
+    // just C, adding B @ C to the profile map while removing same context (but
+    // with different profiles) from the map can cause a conflict if they are
+    // not handled in a right order. This can be solved by just caching the
+    // profiles to be added.
+    auto Ret = ProfilesToBeAdded.try_emplace(ContextStr, FProfile);
+    (void)Ret;
+    assert(Ret.second && "Context conflict during canonicalization");
+    ProfilesToBeRemoved.push_back(I.first());
   }
 
   for (auto &I : ProfilesToBeRemoved) {
-    ProfileMap.erase(I.first());
+    ProfileMap.erase(I);
+  }
+
+  for (auto &I : ProfilesToBeAdded) {
+    ProfileMap.try_emplace(I.first(), I.second);
   }
 }
 
 std::error_code ProfileSymbolList::write(raw_ostream &OS) {
   // Sort the symbols before output. If doing compression.
   // It will make the compression much more effective.
   std::vector<StringRef> SortedList(Syms.begin(), Syms.end());
   llvm::sort(SortedList);

llvm/tools/llvm-profgen/CSPreInliner.cpp

 #include <cstdint>
 #include <queue>
 
 #define DEBUG_TYPE "cs-preinliner"
 
 using namespace llvm;
 using namespace sampleprof;
 
-static cl::opt<bool> EnableCSPreInliner(
-    "csspgo-preinliner", cl::Hidden, cl::init(false),
-    cl::desc("Run a global pre-inliner to merge context profile based on "
-             "estimated global top-down inline decisions"));
-
 // The switches specify inline thresholds used in SampleProfileLoader inlining.
 // TODO: the actual threshold to be tuned here because the size here is based
 // on machine code not LLVM IR.
 extern cl::opt<int> SampleHotCallSiteThreshold;
 extern cl::opt<int> SampleColdCallSiteThreshold;
 extern cl::opt<int> ProfileInlineGrowthLimit;
 extern cl::opt<int> ProfileInlineLimitMin;
 extern cl::opt<int> ProfileInlineLimitMax;
while (!CQueue.empty()) {
              << " (candidate size:" << Candidate.SizeCost
              << ", call count: " << Candidate.CallsiteCount << ", previously "
              << (WasInlined ? "inlined)\n" : "not inlined)\n");
     }
   });
 }
 
 void CSPreInliner::run() {
-  if (!EnableCSPreInliner)
-    return;
-
 #ifndef NDEBUG
   auto printProfileNames = [](StringMap<FunctionSamples> &Profiles,
                               bool IsInput) {
     dbgs() << (IsInput ? "Input" : "Output") << " context-sensitive profiles ("
            << Profiles.size() << " total):\n";
     for (auto &It : Profiles) {
       const FunctionSamples &Samples = It.second;
       dbgs() << "  [" << Samples.getNameWithContext() << "] "

llvm/tools/llvm-profgen/ProfileGenerator.cpp

static cl::opt<bool> CSProfTrimColdContext(
     cl::desc("If the total count of the profile after all merge is done "
              "is still smaller than threshold, it will be trimmed."));
 
 static cl::opt<uint32_t> CSProfColdContextFrameDepth(
     "csprof-frame-depth-for-cold-context", cl::init(1), cl::ZeroOrMore,
     cl::desc("Keep the last K frames while merging cold profile. 1 means the "
              "context-less base profile"));
 
+static cl::opt<bool> EnableCSPreInliner(
+    "csspgo-preinliner", cl::Hidden, cl::init(false),
+    cl::desc("Run a global pre-inliner to merge context profile based on "
+             "estimated global top-down inline decisions"));
+
 extern cl::opt<int> ProfileSummaryCutoffCold;
 
 using namespace llvm;
 using namespace sampleprof;
 
 namespace llvm {
 namespace sampleprof {
 

 
 void CSProfileGenerator::postProcessProfiles() {
   // Compute hot/cold threshold based on profile. This will be used for cold
   // context profile merging/trimming.
   computeSummaryAndThreshold();
 
   // Run global pre-inliner to adjust/merge context profile based on estimated
   // inline decisions.
+  if (EnableCSPreInliner)
     CSPreInliner(ProfileMap, HotCountThreshold, ColdCountThreshold).run();
 
   // Trim and merge cold context profile using cold threshold above;
   SampleContextTrimmer(ProfileMap)
       .trimAndMergeColdContextProfiles(
           ColdCountThreshold, CSProfTrimColdContext, CSProfMergeColdContext,
           CSProfColdContextFrameDepth);
 }