[llvm] [MemProf] Fix when CallStackTrie is a single chain with multi alloc type (PR #79433)

[via llvm-commits]

https://github.com/lifengxiang1025 updated https://github.com/llvm/llvm-project/pull/79433

>From 9884c20326fe2b00f9b9ed0d6fae0c5e8f23a2d0 Mon Sep 17 00:00:00 2001
From: lifengxiang <lifengxiang.1025 at bytedance.com>
Date: Thu, 25 Jan 2024 20:15:02 +0800
Subject: [PATCH] [MemProf] Fix when CallStackTrie is a single chain with multi
 alloc type

---
 llvm/lib/Analysis/MemoryProfileInfo.cpp | 15 ++++++++++-----
 1 file changed, 10 insertions(+), 5 deletions(-)

diff --git a/llvm/lib/Analysis/MemoryProfileInfo.cpp b/llvm/lib/Analysis/MemoryProfileInfo.cpp
index 7fbcffc6489dcfd..274316f8038f04f 100644
--- a/llvm/lib/Analysis/MemoryProfileInfo.cpp
+++ b/llvm/lib/Analysis/MemoryProfileInfo.cpp
@@ -244,11 +244,16 @@ bool CallStackTrie::buildAndAttachMIBMetadata(CallBase *CI) {
   MIBCallStack.push_back(AllocStackId);
   std::vector<Metadata *> MIBNodes;
   assert(!Alloc->Callers.empty() && "addCallStack has not been called yet");
-  buildMIBNodes(Alloc, Ctx, MIBCallStack, MIBNodes,
-                /*CalleeHasAmbiguousCallerContext=*/true);
-  assert(MIBCallStack.size() == 1 &&
-         "Should only be left with Alloc's location in stack");
-  CI->setMetadata(LLVMContext::MD_memprof, MDNode::get(Ctx, MIBNodes));
+  if (buildMIBNodes(Alloc, Ctx, MIBCallStack, MIBNodes,
+                    Alloc->Callers.size() > 1)) {
+    assert(MIBCallStack.size() == 1 &&
+           "Should only be left with Alloc's location in stack");
+    CI->setMetadata(LLVMContext::MD_memprof, MDNode::get(Ctx, MIBNodes));
+  } else
+    // If there exists corner case that CallStackTrie is one chain and all node
+    // with multi alloc type, Conservatively give it non-cold allocation type.
+    // FIXME: Avoid this case before memory profile created.
+    addAllocTypeAttribute(Ctx, CI, AllocationType::NotCold);
   return true;
 }