[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

1632

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

1648

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

[huangjd]

llvm/lib/Transforms/IPO/SampleProfile.cpp

1632

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.

1648

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

1632

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.

1648

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

1660–1663

Curious why we need to reset the SampleProfileMatcher?

[huangjd]

Updated comment, remove unnecessary reset of SampleProfileMatcher

llvm/lib/Transforms/IPO/SampleProfile.cpp

1660–1663

Removed. This field does not need to be updated

[wlei]

llvm/lib/Transforms/IPO/SampleProfile.cpp

1658

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

1649–1651

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);
}

1657–1658

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.

1659–1660

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

533 ↗

(On Diff #549583)

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.

1627

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

1627

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

1627

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

[wenlei]

llvm/lib/Transforms/IPO/SampleProfile.cpp

1627

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 ↗

(On Diff #550852)

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/lib/Transforms/IPO/SampleProfile.cpp

protected:
   // Inline cold/small functions in addition to hot ones
   bool shouldInlineColdCallee(CallBase &CallInst);
   void emitOptimizationRemarksForInlineCandidates(
       const SmallVectorImpl<CallBase *> &Candidates, const Function &F,
       bool Hot);
   void promoteMergeNotInlinedContextSamples(
       MapVector<CallBase *, const FunctionSamples *> NonInlinedCallSites,
       const Function &F);
+  void updateFunctionSamplesPointers(const Function &F);
   std::vector<Function *> buildFunctionOrder(Module &M, LazyCallGraph &CG);
   std::unique_ptr<ProfiledCallGraph> buildProfiledCallGraph(Module &M);
   void generateMDProfMetadata(Function &F);
 
   /// Map from function name to Function *. Used to find the function from
   /// the function name. If the function name contains suffix, additional
   /// entry is added to map from the stripped name to the function if there
   /// is one-to-one mapping.
bool SampleProfileLoader::inlineHotFunctionsWithPriority(
   if (!FunctionSamples::ProfileIsCS)
     promoteMergeNotInlinedContextSamples(LocalNotInlinedCallSites, F);
   return Changed;
 }
 
 void SampleProfileLoader::promoteMergeNotInlinedContextSamples(
     MapVector<CallBase *, const FunctionSamples *> NonInlinedCallSites,
     const Function &F) {
+  std::vector<std::pair<Function *, FunctionSamples>> NewFunctionSamples;
   // Accumulate not inlined callsite information into notInlinedSamples
   for (const auto &Pair : NonInlinedCallSites) {
     CallBase *I = Pair.first;
     Function *Callee = I->getCalledFunction();
     if (!Callee || Callee->isDeclaration())
       continue;
 
     ORE->emit(
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);
+        // We need to handle create new samples later because it can trigger
+        // a rehash invalidating all pointers to FunctionSamples.
+        FunctionSamples *OutlineFS = Reader->getSamplesFor(*Callee);
+        if (OutlineFS != nullptr) {
           OutlineFS->merge(*FS, 1);
           // Set outlined profile to be synthetic to not bias the inliner.
           OutlineFS->SetContextSynthetic();
+        } else
+          NewFunctionSamples.emplace_back(Callee, *FS);
       }
     } else {
       auto pair =
           notInlinedCallInfo.try_emplace(Callee, NotInlinedProfileInfo{0});
       pair.first->second.entryCount += FS->getHeadSamplesEstimate();
     }
   }
+  for (const auto &Pair : NewFunctionSamples) {
+    FunctionSamples *OutlineFS = Reader->getOrCreateSamplesFor(*Pair.first);
+    OutlineFS->merge(Pair.second, 1);
+    // Set outlined profile to be synthetic to not bias the inliner.
+    OutlineFS->SetContextSynthetic();
+  }
+  // Insertion to the sample profile map may invalidates all pointers to
+  // existing FunctionSamples, so we need to update any fields with them.
+  if (NewFunctionSamples.size() > 0) {
+    updateFunctionSamplesPointers(F);
+  }
+}
+
+/// Update all fields containing a pointer to FunctionSamples, used after the
+/// sample profile map is rehashed because FunctionSamples may be moved.
+void SampleProfileLoader::updateFunctionSamplesPointers(const Function &F) {
+  assert(Reader->profileIsCS() == FunctionSamples::ProfileIsCS);
+  if (Reader->profileIsCS())
+    ContextTracker = std::make_unique<SampleContextTracker>(
+        Reader->getProfiles(), &GUIDToFuncNameMap);
+  if (ReportProfileStaleness || PersistProfileStaleness || SalvageStaleProfile)
+    MatchingManager =
+        std::make_unique<SampleProfileMatcher>(
+          const_cast<Module&>(*F.getParent()), *Reader, ProbeManager.get());
+
+  if (FunctionSamples::ProfileIsCS)
+    Samples = ContextTracker->getBaseSamplesFor(F);
+  else
+    Samples = Reader->getSamplesFor(F);
+  assert(Samples && !Samples->empty());
+  for (auto &DILocation2Sample : DILocation2SampleMap) {
+    auto DIL = DILocation2Sample.first;
+    auto &Sample = DILocation2Sample.second;
+    if (FunctionSamples::ProfileIsCS)
+      Sample = ContextTracker->getContextSamplesFor(DIL);
+    else
+      Sample = Samples->findFunctionSamples(DIL, Reader->getRemapper());
+  }
 }
 
 /// 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(