diff --git a/llvm/tools/llvm-profgen/ProfileGenerator.cpp b/llvm/tools/llvm-profgen/ProfileGenerator.cpp
index 32641328f649..3d3abda731d9 100644
--- a/llvm/tools/llvm-profgen/ProfileGenerator.cpp
+++ b/llvm/tools/llvm-profgen/ProfileGenerator.cpp
@@ -216,19 +216,21 @@ void ProfileGenerator::findDisjointRanges(RangeSample &DisjointRanges,
   }
 }
 
-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 Ret.first->second;
+  return I->second;
 }
 
 void CSProfileGenerator::generateProfile() {
@@ -543,11 +545,8 @@ void PseudoProbeCSProfileGenerator::populateBodySamplesWithProbes(
         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);
@@ -610,13 +609,11 @@ FunctionSamples &PseudoProbeCSProfileGenerator::getFunctionProfileForLeafProbe(
   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;
 }
diff --git a/llvm/tools/llvm-profgen/ProfileGenerator.h b/llvm/tools/llvm-profgen/ProfileGenerator.h
index 0b065d70e9ba..96b3675cdb86 100644
--- a/llvm/tools/llvm-profgen/ProfileGenerator.h
+++ b/llvm/tools/llvm-profgen/ProfileGenerator.h
@@ -190,8 +190,9 @@ public:
 
 protected:
   // Lookup or create FunctionSamples for the context
-  FunctionSamples &getFunctionProfileForContext(SampleContextFrames ContextId,
-                                                bool WasLeafInlined = false);
+  FunctionSamples &
+  getFunctionProfileForContext(const SampleContextFrameVector &Context,
+                               bool WasLeafInlined = false);
   // Post processing for profiles before writing out, such as mermining
   // and trimming cold profiles, running preinliner on profiles.
   void postProcessProfiles();
@@ -203,6 +204,9 @@ protected:
   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
@@ -249,9 +253,6 @@ private:
   FunctionSamples &
   getFunctionProfileForLeafProbe(SampleContextFrames ContextStack,
                                  const MCDecodedPseudoProbe *LeafProbe);
-
-  // Underlying context table serves for sample profile writer.
-  std::unordered_set<SampleContextFrameVector, SampleContextFrameHash> Contexts;
 };
 
 } // end namespace sampleprof