[llvm] da2855c - [SampleProfile] Potential use after move in SampleProfileLoader::promoteMergeNotInlinedContextSamples

[William Huang via llvm-commits]

Author: William Huang
Date: 2023-08-16T20:32:15Z
New Revision: da2855c0bad09a53c122797ecc18bdd02505e1c0

URL: https://github.com/llvm/llvm-project/commit/da2855c0bad09a53c122797ecc18bdd02505e1c0
DIFF: https://github.com/llvm/llvm-project/commit/da2855c0bad09a53c122797ecc18bdd02505e1c0.diff

LOG: [SampleProfile] Potential use after move in SampleProfileLoader::promoteMergeNotInlinedContextSamples

SampleProfileLoader::promoteMergeNotInlinedContextSample adds
certain uninlined functions to the sample profile map (unordered_map, which is
previously read from a profile file). This action may cause the map to
be rehashed, invalidating all pointers to FunctionSamples used by many
members of SampleProfileLoader, while the existing code did nothing to
guard against that. This bug is theoretical since adding a few
new functions to a large profile usually won't trigger a rehash, or even
if there's a rehash std::unordered_map tries its best to expand its
capacity in-place.

This bug will trigger if the container type of sample profile map is
changed to llvm::DenseMap or other implementation, such as in D147740,
for SampleProfReader's performance reason.

Reviewed By: wenlei

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

Added: 
    

Modified: 
    llvm/include/llvm/ProfileData/SampleProfReader.h
    llvm/include/llvm/Transforms/Utils/SampleProfileLoaderBaseImpl.h
    llvm/lib/Transforms/IPO/SampleProfile.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/include/llvm/ProfileData/SampleProfReader.h b/llvm/include/llvm/ProfileData/SampleProfReader.h
index e14b0bfc7912aa..cd21da3863309e 100644
--- a/llvm/include/llvm/ProfileData/SampleProfReader.h
+++ b/llvm/include/llvm/ProfileData/SampleProfReader.h
@@ -407,22 +407,6 @@ class SampleProfileReader {
     return getSamplesFor(CanonName);
   }
 
-  /// Return the samples collected for function \p F, create empty
-  /// FunctionSamples if it doesn't exist.
-  FunctionSamples *getOrCreateSamplesFor(const Function &F) {
-    std::string FGUID;
-    StringRef CanonName = FunctionSamples::getCanonicalFnName(F);
-    CanonName = getRepInFormat(CanonName, useMD5(), FGUID);
-    auto It = Profiles.find(CanonName);
-    if (It != Profiles.end())
-      return &It->second;
-    if (!FGUID.empty()) {
-      assert(useMD5() && "New name should only be generated for md5 profile");
-      CanonName = *MD5NameBuffer.insert(FGUID).first;
-    }
-    return &Profiles[CanonName];
-  }
-
   /// Return the samples collected for function \p F.
   virtual FunctionSamples *getSamplesFor(StringRef Fname) {
     std::string FGUID;

diff  --git a/llvm/include/llvm/Transforms/Utils/SampleProfileLoaderBaseImpl.h b/llvm/include/llvm/Transforms/Utils/SampleProfileLoaderBaseImpl.h
index 1c6ba530e3df4c..cd5f8ecf1543e1 100644
--- a/llvm/include/llvm/Transforms/Utils/SampleProfileLoaderBaseImpl.h
+++ b/llvm/include/llvm/Transforms/Utils/SampleProfileLoaderBaseImpl.h
@@ -265,6 +265,11 @@ template <typename FT> class SampleProfileLoaderBaseImpl {
   /// Profile reader object.
   std::unique_ptr<SampleProfileReader> Reader;
 
+  /// Synthetic samples created by duplicating the samples of inlined functions
+  /// from the original profile as if they were top level sample profiles.
+  /// Use std::map because insertion may happen while its content is referenced.
+  std::map<SampleContext, FunctionSamples> OutlineFunctionSamples;
+
   // A pseudo probe helper to correlate the imported sample counts.
   std::unique_ptr<PseudoProbeManager> ProbeManager;
 

diff  --git a/llvm/lib/Transforms/IPO/SampleProfile.cpp b/llvm/lib/Transforms/IPO/SampleProfile.cpp
index a53baecd4776db..85431fb32b4b4d 100644
--- a/llvm/lib/Transforms/IPO/SampleProfile.cpp
+++ b/llvm/lib/Transforms/IPO/SampleProfile.cpp
@@ -1619,7 +1619,12 @@ void SampleProfileLoader::promoteMergeNotInlinedContextSamples(
         // 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);
+        FunctionSamples *OutlineFS = Reader->getSamplesFor(*Callee);
+        // If outlined function does not exist in the profile, add it to a
+        // separate map so that it does not rehash the original profile.
+        if (!OutlineFS)
+          OutlineFS = &OutlineFunctionSamples[
+              FunctionSamples::getCanonicalFnName(Callee->getName())];
         OutlineFS->merge(*FS, 1);
         // Set outlined profile to be synthetic to not bias the inliner.
         OutlineFS->SetContextSynthetic();
@@ -2571,8 +2576,24 @@ bool SampleProfileLoader::runOnFunction(Function &F, ModuleAnalysisManager *AM)
 
   if (FunctionSamples::ProfileIsCS)
     Samples = ContextTracker->getBaseSamplesFor(F);
-  else
+  else {
     Samples = Reader->getSamplesFor(F);
+    // Try search in previously inlined functions that were split or duplicated
+    // into base.
+    if (!Samples) {
+      StringRef CanonName = FunctionSamples::getCanonicalFnName(F);
+      auto It = OutlineFunctionSamples.find(CanonName);
+      if (It != OutlineFunctionSamples.end()) {
+        Samples = &It->second;
+      } else if (auto Remapper = Reader->getRemapper()) {
+        if (auto RemppedName = Remapper->lookUpNameInProfile(CanonName)) {
+          It = OutlineFunctionSamples.find(*RemppedName);
+          if (It != OutlineFunctionSamples.end())
+            Samples = &It->second;
+        }
+      }
+    }
+  }
 
   if (Samples && !Samples->empty())
     return emitAnnotations(F);