[MemProf] Fix bug introduced by restructuring in optional handling

[teresajohnson]

The restructuring of the context pruning patch in PR138792
(764614e) introduced a bug under the
non-default -memprof-keep-all-not-cold-contexts handling.

Added more testing of this mode which would have caught the issue.

While here, fix the newly added function name to match code style.

[llvmbot]

@llvm/pr-subscribers-llvm-analysis

Author: Teresa Johnson (teresajohnson)

Changes

The restructuring of the context pruning patch in PR138792
(764614e) introduced a bug under the
non-default -memprof-keep-all-not-cold-contexts handling.

Added more testing of this mode which would have caught the issue.

While here, fix the newly added function name to match code style.

---

Full diff: https://github.com/llvm/llvm-project/pull/139092.diff

2 Files Affected:

• (modified) llvm/lib/Analysis/MemoryProfileInfo.cpp (+6-4)
• (modified) llvm/unittests/Analysis/MemoryProfileInfoTest.cpp (+73)

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.

[teresajohnson]

@snehasish noticed this newly introduced bug when prepping a follow on change this morning, so I went ahead and merged the fix. Please do post-commit review when you get a chance and I'll address any feedback in a follow up PR

[snehasish]

lgtm

[snehasish]

nit: I don't think it's comon to use this API in unittests to reset the cl::opt, the only other tests which use this are CommandLineTest.cpp and ExecutionTest.cpp (apart from another usage in this file). In other unittests I've seen just the value being set to the default e.g. false in this case.

[teresajohnson]

I use it in one other place in this file to do something similar, and landed on that solution because there is no copy constructor for cl::opt so I couldn't save and restore the default (and was concerned about hardwiring it). But I can just use if conditions to do the save and restore. Will send a follow on change to modify both usages in this file to that approach.

[teresajohnson]

Actually, no if condition needed. I can't copy to another cl::opt, but changing "auto" to "bool" on the save variable works. Sent PR139117 to fix both instances in this file.