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

[huangjd]

Summary

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.

[davidxl]

llvm/lib/Transforms/IPO/SampleProfile.cpp

1638

is it enough to just use a flag to indicate insertion happens (instead of using a vector)? The getOrCreateSamples can set the flag.

1646

insertion does not mean rehashing happens. Is there a better way to avoid the unncessary update?

[huangjd]

llvm/lib/Transforms/IPO/SampleProfile.cpp

1638

The vector also stores a deep copy of the caller's FunctionSamples that's to be merged with newly created outlined function's FunctionSamples. Note that NonInlinedCallSites uses pointers to existing FunctionSamples, which will be invalidated if relocation happens, so we need a copy instead.

1646

C++ unordered_map or llvm::DenseMap does not provide an interface to check if data are relocated. There's ugly way to check it such as comparing the address to the first element. However, rehashing is likely to happen if insertion happens, because the profile reader reserves the exact number of FunctionSamples in the map. Meanwhile insertion itself is unlikely to happen as this function is called very few times in a large production project, so it wouldn't matter for performance.

[davidxl]

llvm/lib/Transforms/IPO/SampleProfile.cpp

1638

Ok. Can you modify the comments to reflect this: avoid creating new function samples (for outline function) in the the loop, otherwise the pointers used in the loop iteration can be invalidated.

1646

make sense. Can you also update the comment saying this is rare event (i.e., missing function samples for outline function).

[wlei]

llvm/lib/Transforms/IPO/SampleProfile.cpp

1658–1661

Curious why we need to reset the SampleProfileMatcher?

[huangjd]

Updated comment, remove unnecessary reset of SampleProfileMatcher

llvm/lib/Transforms/IPO/SampleProfile.cpp

1658–1661

Removed. This field does not need to be updated

[wlei]

llvm/lib/Transforms/IPO/SampleProfile.cpp

1656

It looks like promoteMergeNotInlinedContextSamples are not for CS profile, see the comments:

// For CS profile, profile for not inlined context will be merged when
// base profile is being retrieved.
if (!FunctionSamples::ProfileIsCS)
  promoteMergeNotInlinedContextSamples(LocalNotInlinedCallSites, F);

probably just assert on !Reader->profileIsCS()

[wenlei]

llvm/lib/Transforms/IPO/SampleProfile.cpp

1647–1649

Patching all pointers seems a bit hacky. Have you considered maintaining reference stability by using a separate container for new function samples, so there will be no insertion to SampleProfileReader::Profiles?

Strictly speaking SampleProfileReader::Profiles should only contain profile directly read out of input, and those newly created profiles are synthesized which doesn't belong to reader.

If we use a separate container (say SynthesizedFuncSamples), we would only need to change the code below to look for profiles from the addition container as well for non-CS profile.

if (FunctionSamples::ProfileIsCS)
  Samples = ContextTracker->getBaseSamplesFor(F);
else
  Samples = Reader->getSamplesFor(F);

-->

if (FunctionSamples::ProfileIsCS)
  Samples = ContextTracker->getBaseSamplesFor(F);
else {
  Samples = Reader->getSamplesFor(F);
  if (!Samples)
    Samples = findSynthesizedProfile(F);
}

1655–1656

Yeah we should assert this is non CS profile. CS profile doesn't need promoteMergeNotInlinedContextSamples as all adjustments are done in SampleContextTracker without creating new FunctionSamples.

1657–1658

This is super heavy hammer to rebuild all context trie, which would have been a deal breaker. But again, we shouldn't need any of this for CS profile.

[huangjd]

Rewrite patch, use a much cleaner approach by separating the sample profile map

[wenlei]

llvm/include/llvm/ProfileData/SampleProfReader.h

510

Can this be owned by SampleProfileLoader instead? Reader should just present what's being read out of an input profile, and should not couple with profile loader in other ways.

[huangjd]

Update to make SampleProfileMatcher own the outline profiles

[wenlei]

llvm/lib/Transforms/IPO/SampleProfile.cpp

1622

We can actually remove the getOrCreateSamplesFor API from reader now to guarantee it's immutable after profile is read.

1626

We could also be inserting into OutlineFunctionSamples which can cause rehash while using a FunctionSamples from OutlineFunctionSamples?

I'm wondering if OutlineFunctionSamples should use a container that guarantees reference stability, e.g. std::map instead of std::unordered_map.

[huangjd]

llvm/lib/Transforms/IPO/SampleProfile.cpp

1626

That's OK, there is no pointer reference OutlineFunctionSamples, so there's no concern of segfault. Performance is trivial because very few inlined functions are actually promoted to outline function (say < 10 in a production project )

[huangjd]

Remove unused getOrCreateSamples function

llvm/lib/Transforms/IPO/SampleProfile.cpp

1626

The pointer here is used immediately and won't be used outside of the function

[wenlei]

llvm/lib/Transforms/IPO/SampleProfile.cpp

1626

That's OK, there is no pointer reference OutlineFunctionSamples, so there's no concern of segfault.

Imagine this case: we are processing a function (inside SampleProfileLoader::runOnFunction) whose own profile is retrieved from OutlineFunctionSamples, now its callee profile got inserted into OutlineFunctionSamples, causing a rehash, so SampleProfileLoader::Samples now becomes a dangling reference,

Performance is trivial because very few inlined functions are actually promoted to outline function (say < 10 in a production project )

Right, so we can afford to use std::map.

[huangjd]

use map instead of unordered_map

[wenlei]

thanks for the fix and following up with the comments, lgtm now.

llvm/include/llvm/Transforms/Utils/SampleProfileLoaderBaseImpl.h

270

Add a comment to explain why we use std::map here instead of existing SampleProfileMap, in case others stomped on this code later and tempted to change it to SampleProfileMap for consistency.

[huangjd]

Added comment

[huangjd]

rebase

[ayermolo]

This comment was removed by ayermolo.

llvm/include/llvm/ProfileData/SampleProfReader.h

public:
   FunctionSamples *getSamplesFor(const Function &F) {
     // The function name may have been updated by adding suffix. Call
     // a helper to (optionally) strip off suffixes so that we can
     // match against the original function name in the profile.
     StringRef CanonName = FunctionSamples::getCanonicalFnName(F);
     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;
     Fname = getRepInFormat(Fname, useMD5(), FGUID);
     auto It = Profiles.find(Fname);
     if (It != Profiles.end())
       return &It->second;
 
protected:
   /// The profile of every function executed at runtime is collected
   /// in the structure FunctionSamples. This maps function objects
   /// to their corresponding profiles.
   SampleProfileMap Profiles;
 
   /// LLVM context used to emit diagnostics.
   LLVMContext &Ctx;
 
   /// Memory buffer holding the profile file.
   std::unique_ptr<MemoryBuffer> Buffer;
 
   /// Extra name buffer holding names created on demand.
   /// This should only be needed for md5 profiles.
   std::unordered_set<std::string> MD5NameBuffer;
 
   /// Profile summary information.
   std::unique_ptr<ProfileSummary> Summary;

llvm/include/llvm/Transforms/Utils/SampleProfileLoaderBaseImpl.h

protected:
   BlockEdgeMap Successors;
 
   /// Profile coverage tracker.
   SampleCoverageTracker CoverageTracker;
 
   /// 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;
 
   /// Samples collected for the body of this function.
   FunctionSamples *Samples = nullptr;
 
   /// Name of the profile file to load.
   std::string Filename;

llvm/lib/Transforms/IPO/SampleProfile.cpp

if (ProfileMergeInlinee) {
         // Use entry samples as head samples during the merge, as inlinees
         // don't have head samples.
         const_cast<FunctionSamples *>(FS)->addHeadSamples(
             FS->getHeadSamplesEstimate());
 
         // 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();
       }
     } else {
       auto pair =
           notInlinedCallInfo.try_emplace(Callee, NotInlinedProfileInfo{0});
       pair.first->second.entryCount += FS->getHeadSamplesEstimate();
     }
   }
 }
 
 /// Returns the sorted CallTargetMap \p M by count in descending order.
 static SmallVector<InstrProfValueData, 2>
 GetSortedValueDataFromCallTargets(const SampleRecord::CallTargetMap &M) {
   SmallVector<InstrProfValueData, 2> R;
   for (const auto &I : SampleRecord::SortCallTargets(M)) {
     R.emplace_back(
         InstrProfValueData{FunctionSamples::getGUID(I.first), I.second});
   }
   return R;
 }
 
 // Generate MD_prof metadata for every branch instruction using the
 // edge weights computed during propagation.
 void SampleProfileLoader::generateMDProfMetadata(Function &F) {
   // Generate MD_prof metadata for every branch instruction using the
   // edge weights computed during propagation.
   LLVM_DEBUG(dbgs() << "\nPropagation complete. Setting branch weights\n");
   LLVMContext &Ctx = F.getContext();
   MDBuilder MDB(Ctx);
   for (auto &BI : F) {
     BasicBlock *BB = &BI;
 
     if (BlockWeights[BB]) {
       for (auto &I : *BB) {
         if (!isa<CallInst>(I) && !isa<InvokeInst>(I))
           continue;
         if (!cast<CallBase>(I).getCalledFunction()) {
           const DebugLoc &DLoc = I.getDebugLoc();
           if (!DLoc)
             continue;
if (AM) {
     ORE = &FAM.getResult<OptimizationRemarkEmitterAnalysis>(F);
   } else {
     OwnedORE = std::make_unique<OptimizationRemarkEmitter>(&F);
     ORE = OwnedORE.get();
   }
 
   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);
   return false;
 }
 SampleProfileLoaderPass::SampleProfileLoaderPass(
     std::string File, std::string RemappingFile, ThinOrFullLTOPhase LTOPhase,
     IntrusiveRefCntPtr<vfs::FileSystem> FS)