[llvm] [MemProf] Prune unneeded non-cold contexts (PR #124823)

[Snehasish Kumar via llvm-commits]

================
@@ -428,14 +469,65 @@ declare dso_local noalias noundef i8* @malloc(i64 noundef)
       EXPECT_EQ(getMIBAllocType(MIB), AllocationType::Cold);
     else if (StackId->getZExtValue() == 5u)
       EXPECT_EQ(getMIBAllocType(MIB), AllocationType::NotCold);
-    else {
-      ASSERT_EQ(StackId->getZExtValue(), 8u);
-      // Hot contexts are converted to NotCold when building the metadata.
-      EXPECT_EQ(getMIBAllocType(MIB), AllocationType::NotCold);
-    }
   }
 }
 
+// Test to ensure that we prune NotCold contexts that are unneeded for
+// determining where Cold contexts need to be cloned to enable correct hinting.
+TEST_F(MemoryProfileInfoTest, PruneUnneededNotColdContexts) {
+  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});
+
+  // Neither of these two 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});
+
+  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}}};
+
+  CallBase *Call = findCall(*Func, "call");
+  Trie.buildAndAttachMIBMetadata(Call);
----------------
snehasish wrote:

Should we assert Call is not a nullptr before we pass it to the function?

https://github.com/llvm/llvm-project/pull/124823