[llvm-profgen] Fix a dangling vector reference in CS line number based generator

[wlei]

Summary

It seems we missed one spot to persist SampleContextFrameVector into the global table (CSProfileGenerator::populateFunctionBoundarySamples:340) which causes a crash.

This change tried to fix it in a centralized way i. e. where we generate the FunctionSamples.

[wenlei]

lgtm. so the problem was with profile generation without probe, right?

[hoy]

lgtm, thanks for the fix.

[wlei]

In D110275#3016049, @wenlei wrote:

lgtm. so the problem was with profile generation without probe, right?

Yeah, this is for CS line-number based generation which we didn't test for large service.

llvm/tools/llvm-profgen/ProfileGenerator.cpp

225

Side note:

This might be improved in the future, if it's a new context, it will always do a twice search(find and emplace), I guess context search is heavy.

[wenlei]

llvm/tools/llvm-profgen/ProfileGenerator.cpp

225

I think emplace alone is enough. The return value tells you whether the element exists, and insertion only happens when the element does not exist. Though we would need to create NewContext in Contexts first in order to combine find with emplace.

[wlei]

using const SampleContextFrameVector &Context instead of ArrayRef to avoid obj copy

llvm/tools/llvm-profgen/ProfileGenerator.h

while (I <= MaxDedupSize) {
       I++;
       Context.resize(End);
       MaxDedupSize = std::min(static_cast<uint32_t>(End / 2), MaxDedupSize);
     }
   }
 
 protected:
   // Lookup or create FunctionSamples for the context
-  FunctionSamples &getFunctionProfileForContext(SampleContextFrames ContextId,
-                                                bool WasLeafInlined = false);
+  FunctionSamples &
+  getFunctionProfileForContext(const SampleContextFrameVector &Context,
   // Post processing for profiles before writing out, such as mermining
   // and trimming cold profiles, running preinliner on profiles.
   void postProcessProfiles();
   void computeSummaryAndThreshold();
   void write(std::unique_ptr<SampleProfileWriter> Writer,
              SampleProfileMap &ProfileMap) override;
 
   // Thresholds from profile summary to answer isHotCount/isColdCount queries.
   uint64_t HotCountThreshold;
   uint64_t ColdCountThreshold;
 
+  // Underlying context table serves for sample profile writer.
+  std::unordered_set<SampleContextFrameVector, SampleContextFrameHash> Contexts;
+
 private:
   // Helper function for updating body sample for a leaf location in
   // FunctionProfile
   void updateBodySamplesforFunctionProfile(FunctionSamples &FunctionProfile,
                                            const SampleContextFrame &LeafLoc,
                                            uint64_t Count);
   void populateFunctionBodySamples(FunctionSamples &FunctionProfile,
                                    const RangeSample &RangeCounters);
void populateBodySamplesWithProbes(const RangeSample &RangeCounter,
                                      SampleContextFrames ContextStack);
   // Fill in boundary samples for a call probe
   void populateBoundarySamplesWithProbes(const BranchSample &BranchCounter,
                                          SampleContextFrames ContextStack);
   // Helper function to get FunctionSamples for the leaf probe
   FunctionSamples &
   getFunctionProfileForLeafProbe(SampleContextFrames ContextStack,
                                  const MCDecodedPseudoProbe *LeafProbe);
-
-  // Underlying context table serves for sample profile writer.
-  std::unordered_set<SampleContextFrameVector, SampleContextFrameHash> Contexts;
 };
 
 } // end namespace sampleprof
 } // end namespace llvm
 
 #endif

llvm/tools/llvm-profgen/ProfileGenerator.cpp

if (Point.EndCount) {
       DisjointRanges[{BeginAddress, Address}] = Count;
       assert(Count >= Point.EndCount && "Mismatched live ranges");
       Count -= Point.EndCount;
       BeginAddress = Address + 1;
     }
   }
 }
 
-FunctionSamples &
-CSProfileGenerator::getFunctionProfileForContext(SampleContextFrames Context,
-                                                 bool WasLeafInlined) {
-  SampleContext FContext(Context);
-  auto Ret = ProfileMap.emplace(Context, FunctionSamples());
-  if (Ret.second) {
-    SampleContext FContext(Context, RawContext);
+FunctionSamples &CSProfileGenerator::getFunctionProfileForContext(
+    const SampleContextFrameVector &Context, bool WasLeafInlined) {
+  auto I = ProfileMap.find(SampleContext(Context));
+  if (I == ProfileMap.end()) {
+    // Save the new context for future references.
+    SampleContextFrames NewContext = *Contexts.insert(Context).first;
+    SampleContext FContext(NewContext, RawContext);
+    auto Ret = ProfileMap.emplace(FContext, FunctionSamples());
     if (WasLeafInlined)
       FContext.setAttribute(ContextWasInlined);
     FunctionSamples &FProfile = Ret.first->second;
     FProfile.setContext(FContext);
-  }
     return Ret.first->second;
   }
+  return I->second;
+}
 
 void CSProfileGenerator::generateProfile() {
   FunctionSamples::ProfileIsCS = true;
   for (const auto &CI : SampleCounters) {
     const StringBasedCtxKey *CtxKey =
         dyn_cast<StringBasedCtxKey>(CI.first.getPtr());
     // Get or create function profile for the range
     FunctionSamples &FunctionProfile =
if (Probe->isEntry()) {
         SampleContextFrames CalleeContextId =
             FunctionProfile.getContext().getContextFrames();
         SampleContextFrameVector CallerContextId;
         SampleContextFrame &&CallerLeafFrameLoc =
             getCallerContext(CalleeContextId, CallerContextId);
         uint64_t CallerIndex = CallerLeafFrameLoc.Callsite.LineOffset;
         assert(CallerIndex &&
                "Inferred caller's location index shouldn't be zero!");
-        // Save the new context for future references.
-        SampleContextFrames CallerContext =
-            *Contexts.insert(CallerContextId).first;
         FunctionSamples &CallerProfile =
-            getFunctionProfileForContext(CallerContext);
+            getFunctionProfileForContext(CallerContextId);
         CallerProfile.setFunctionHash(InlinerDesc->FuncHash);
         CallerProfile.addBodySamples(CallerIndex, 0, Count);
         CallerProfile.addTotalSamples(Count);
         CallerProfile.addCalledTargetSamples(
             CallerIndex, 0, FunctionProfile.getContext().getName(), Count);
       }
     }
   }
FunctionSamples &PseudoProbeCSProfileGenerator::getFunctionProfileForLeafProbe(
   // Inlined stack: [foo:1, bar:2], the ContextId will be "foo:1 @ bar"
   auto LeafFrame = NewContextStack.back();
   LeafFrame.Callsite = LineLocation(0, 0);
   NewContextStack.pop_back();
   // Compress the context string except for the leaf frame
   CSProfileGenerator::compressRecursionContext(NewContextStack);
   CSProfileGenerator::trimContext(NewContextStack);
   NewContextStack.push_back(LeafFrame);
-  // Save the new context for future references.
-  SampleContextFrames NewContext = *Contexts.insert(NewContextStack).first;
 
   const auto *FuncDesc = Binary->getFuncDescForGUID(LeafProbe->getGuid());
   bool WasLeafInlined = LeafProbe->getInlineTreeNode()->hasInlineSite();
   FunctionSamples &FunctionProile =
-      getFunctionProfileForContext(NewContext, WasLeafInlined);
+      getFunctionProfileForContext(NewContextStack, WasLeafInlined);
   FunctionProile.setFunctionHash(FuncDesc->FuncHash);
   return FunctionProile;
 }
 
 } // end namespace sampleprof
 } // end namespace llvm