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

[William Junda Huang via Phabricator via llvm-commits]

huangjd created this revision.
huangjd added reviewers: davidxl, snehasish, MatzeB, ro, wenlei.
Herald added subscribers: wlei, ormris, hiraditya.
Herald added a project: All.
huangjd requested review of this revision.
Herald added a project: LLVM.
Herald added a subscriber: llvm-commits.

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 <https://reviews.llvm.org/D147740>,
for SampleProfReader's performance reason.


Repository:
  rG LLVM Github Monorepo

https://reviews.llvm.org/D157061

Files:
  llvm/lib/Transforms/IPO/SampleProfile.cpp


Index: llvm/lib/Transforms/IPO/SampleProfile.cpp
===================================================================
--- llvm/lib/Transforms/IPO/SampleProfile.cpp
+++ llvm/lib/Transforms/IPO/SampleProfile.cpp
@@ -565,6 +565,7 @@
   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);
@@ -1581,6 +1582,7 @@
 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;
@@ -1619,10 +1621,15 @@
         // 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, 1);
-        // Set outlined profile to be synthetic to not bias the inliner.
-        OutlineFS->SetContextSynthetic();
+        // 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 =
@@ -1630,6 +1637,44 @@
       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.


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D157061.547068.patch
Type: text/x-patch
Size: 3926 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20230804/b17582df/attachment.bin>