diff --git a/llvm/lib/Analysis/MemoryProfileInfo.cpp b/llvm/lib/Analysis/MemoryProfileInfo.cpp
index 30b674c320ef6..924325b97eaaa 100644
--- a/llvm/lib/Analysis/MemoryProfileInfo.cpp
+++ b/llvm/lib/Analysis/MemoryProfileInfo.cpp
@@ -54,7 +54,7 @@ cl::opt<bool> MemProfReportHintedSizes(
 // This is useful if we have enabled reporting of hinted sizes, and want to get
 // information from the indexing step for all contexts (especially for testing),
 // or have specified a value less than 100% for -memprof-cloning-cold-threshold.
-static cl::opt<bool> MemProfKeepAllNotColdContexts(
+cl::opt<bool> MemProfKeepAllNotColdContexts(
     "memprof-keep-all-not-cold-contexts", cl::init(false), cl::Hidden,
     cl::desc("Keep all non-cold contexts (increases cloning overheads)"));
 
@@ -244,12 +244,14 @@ void CallStackTrie::convertHotToNotCold(CallStackTrieNode *Node) {
 
 // Copy over some or all of NewMIBNodes to the SavedMIBNodes vector, depending
 // on options that enable filtering out some NotCold contexts.
-static void SaveFilteredNewMIBNodes(std::vector<Metadata *> &NewMIBNodes,
+static void saveFilteredNewMIBNodes(std::vector<Metadata *> &NewMIBNodes,
                                     std::vector<Metadata *> &SavedMIBNodes,
                                     unsigned CallerContextLength) {
   // In the simplest case, with pruning disabled, keep all the new MIB nodes.
-  if (MemProfKeepAllNotColdContexts)
+  if (MemProfKeepAllNotColdContexts) {
     append_range(SavedMIBNodes, NewMIBNodes);
+    return;
+  }
 
   auto EmitMessageForRemovedContexts = [](const MDNode *MIBMD, StringRef Tag,
                                           StringRef Extra) {
@@ -372,7 +374,7 @@ bool CallStackTrie::buildMIBNodes(CallStackTrieNode *Node, LLVMContext &Ctx,
     }
     // Pass in the stack length of the MIB nodes added for the immediate caller,
     // which is the current stack length plus 1.
-    SaveFilteredNewMIBNodes(NewMIBNodes, MIBNodes, MIBCallStack.size() + 1);
+    saveFilteredNewMIBNodes(NewMIBNodes, MIBNodes, MIBCallStack.size() + 1);
     if (AddedMIBNodesForAllCallerContexts)
       return true;
     // We expect that the callers should be forced to add MIBs to disambiguate
diff --git a/llvm/unittests/Analysis/MemoryProfileInfoTest.cpp b/llvm/unittests/Analysis/MemoryProfileInfoTest.cpp
index 637fff93cc1ad..2943631150650 100644
--- a/llvm/unittests/Analysis/MemoryProfileInfoTest.cpp
+++ b/llvm/unittests/Analysis/MemoryProfileInfoTest.cpp
@@ -27,6 +27,7 @@ extern cl::opt<float> MemProfLifetimeAccessDensityColdThreshold;
 extern cl::opt<unsigned> MemProfAveLifetimeColdThreshold;
 extern cl::opt<unsigned> MemProfMinAveLifetimeAccessDensityHotThreshold;
 extern cl::opt<bool> MemProfUseHotHints;
+extern cl::opt<bool> MemProfKeepAllNotColdContexts;
 
 namespace {
 
@@ -530,6 +531,78 @@ declare dso_local noalias noundef i8* @malloc(i64 noundef)
   EXPECT_THAT(MemProfMD, MemprofMetadataEquals(ExpectedVals));
 }
 
+// Test to ensure that we keep optionally keep unneeded NotCold contexts.
+// Same as PruneUnneededNotColdContexts test but with the
+// MemProfKeepAllNotColdContexts set to true.
+TEST_F(MemoryProfileInfoTest, KeepUnneededNotColdContexts) {
+  LLVMContext C;
+  std::unique_ptr<Module> M = makeLLVMModule(C,
+                                             R"IR(
+target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-pc-linux-gnu"
+define i32* @test() {
+entry:
+  %call = call noalias dereferenceable_or_null(40) i8* @malloc(i64 noundef 40)
+  %0 = bitcast i8* %call to i32*
+  ret i32* %0
+}
+declare dso_local noalias noundef i8* @malloc(i64 noundef)
+)IR");
+
+  Function *Func = M->getFunction("test");
+
+  CallStackTrie Trie;
+
+  Trie.addCallStack(AllocationType::Cold, {1, 2, 3, 4});
+  Trie.addCallStack(AllocationType::Cold, {1, 2, 3, 5, 6, 7});
+  // This NotCold context is needed to know where the above two Cold contexts
+  // must be cloned from:
+  Trie.addCallStack(AllocationType::NotCold, {1, 2, 3, 5, 6, 13});
+
+  Trie.addCallStack(AllocationType::Cold, {1, 2, 3, 8, 9, 10});
+  // This NotCold context is needed to know where the above Cold context must be
+  // cloned from:
+  Trie.addCallStack(AllocationType::NotCold, {1, 2, 3, 8, 9, 14});
+  // This NotCold context is not needed since the above is sufficient (we pick
+  // the first in sorted order).
+  Trie.addCallStack(AllocationType::NotCold, {1, 2, 3, 8, 9, 15});
+
+  // None of these NotCold contexts are needed as the Cold contexts they
+  // overlap with are covered by longer overlapping NotCold contexts.
+  Trie.addCallStack(AllocationType::NotCold, {1, 2, 3, 12});
+  Trie.addCallStack(AllocationType::NotCold, {1, 2, 11});
+  Trie.addCallStack(AllocationType::NotCold, {1, 16});
+
+  // We should keep all of the above contexts, even those that are unneeded by
+  // default, because we will set the option to keep them.
+  std::vector<std::pair<AllocationType, std::vector<unsigned>>> ExpectedVals = {
+      {AllocationType::Cold, {1, 2, 3, 4}},
+      {AllocationType::Cold, {1, 2, 3, 5, 6, 7}},
+      {AllocationType::NotCold, {1, 2, 3, 5, 6, 13}},
+      {AllocationType::Cold, {1, 2, 3, 8, 9, 10}},
+      {AllocationType::NotCold, {1, 2, 3, 8, 9, 14}},
+      {AllocationType::NotCold, {1, 2, 3, 8, 9, 15}},
+      {AllocationType::NotCold, {1, 2, 3, 12}},
+      {AllocationType::NotCold, {1, 2, 11}},
+      {AllocationType::NotCold, {1, 16}}};
+
+  CallBase *Call = findCall(*Func, "call");
+  ASSERT_NE(Call, nullptr);
+
+  // Specify that all non-cold contexts should be kept.
+  MemProfKeepAllNotColdContexts = true;
+
+  Trie.buildAndAttachMIBMetadata(Call);
+
+  // Undo the manual set of the MemProfKeepAllNotColdContexts above.
+  cl::ResetAllOptionOccurrences();
+
+  EXPECT_FALSE(Call->hasFnAttr("memprof"));
+  EXPECT_TRUE(Call->hasMetadata(LLVMContext::MD_memprof));
+  MDNode *MemProfMD = Call->getMetadata(LLVMContext::MD_memprof);
+  EXPECT_THAT(MemProfMD, MemprofMetadataEquals(ExpectedVals));
+}
+
 // Test CallStackTrie::addCallStack interface taking memprof MIB metadata.
 // Check that allocations annotated with memprof metadata with a single
 // allocation type get simplified to an attribute.