diff --git a/llvm/include/llvm/Analysis/MemoryProfileInfo.h b/llvm/include/llvm/Analysis/MemoryProfileInfo.h
new file mode 100644
--- /dev/null
+++ b/llvm/include/llvm/Analysis/MemoryProfileInfo.h
@@ -0,0 +1,112 @@
+//===- llvm/Analysis/MemoryProfileInfo.h - memory profile info ---*- C++ -*-==//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file contains utilities to analyze memory profile information.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_ANALYSIS_MEMORYPROFILEINFO_H
+#define LLVM_ANALYSIS_MEMORYPROFILEINFO_H
+
+#include "llvm/IR/Constants.h"
+#include "llvm/IR/InstrTypes.h"
+#include "llvm/IR/Metadata.h"
+#include "llvm/IR/Module.h"
+#include <map>
+
+namespace llvm {
+namespace memprof {
+
+// Allocation type assigned to an allocation reached by a given context.
+// More can be added but initially this is just noncold and cold.
+// Values should be powers of two so that they can be ORed, in particular to
+// track allocations that have different behavior with different calling
+// contexts.
+enum class AllocationType : uint8_t { None = 0, NotCold = 1, Cold = 2 };
+
+/// Return the allocation type for a given set of memory profile values.
+AllocationType getAllocType(uint64_t MaxAccessCount, uint64_t MinSize,
+                            uint64_t MinLifetime);
+
+/// Build callstack metadata from the provided list of call stack ids. Returns
+/// the resulting metadata node.
+MDNode *buildCallstackMetadata(ArrayRef<uint64_t> CallStack, LLVMContext &Ctx);
+
+/// Returns the stack node from an MIB metadata node.
+MDNode *getMIBStackNode(const MDNode *MIB);
+
+/// Returns the allocation type from an MIB metadata node.
+AllocationType getMIBAllocType(const MDNode *MIB);
+
+/// Class to build a trie of call stack contexts for a particular profiled
+/// allocation call, along with their associated allocation types.
+/// The allocation will be at the root of the trie, which is then used to
+/// compute the minimum lists of context ids needed to associate a call context
+/// with a single allocation type.
+class CallStackTrie {
+private:
+  struct CallStackTrieNode {
+    // Allocation types for call context sharing the context prefix at this
+    // node.
+    uint8_t AllocTypes;
+    // Map of caller stack id to the corresponding child Trie node.
+    std::map<uint64_t, CallStackTrieNode *> Callers;
+    CallStackTrieNode(AllocationType Type)
+        : AllocTypes(static_cast<uint8_t>(Type)) {}
+  };
+
+  // The node for the allocation at the root.
+  CallStackTrieNode *Alloc;
+  // The allocation's leaf stack id.
+  uint64_t AllocStackId;
+
+  void deleteTrieNode(CallStackTrieNode *Node) {
+    if (!Node)
+      return;
+    for (auto C : Node->Callers)
+      deleteTrieNode(C.second);
+    delete Node;
+  }
+
+  // Recursive helper to trim contexts and create metadata nodes.
+  bool buildMIBNodes(CallStackTrieNode *Node, LLVMContext &Ctx,
+                     std::vector<uint64_t> &MIBCallStack,
+                     std::vector<Metadata *> &MIBNodes,
+                     bool CalleeHasAmbiguousCallerContext);
+
+public:
+  CallStackTrie() : Alloc(nullptr), AllocStackId(0) {}
+  ~CallStackTrie() { deleteTrieNode(Alloc); }
+
+  bool empty() const { return Alloc == nullptr; }
+
+  /// Add a call stack context with the given allocation type to the Trie.
+  /// The context is represented by the list of stack ids (computed during
+  /// matching via a debug location hash), expected to be in order from the
+  /// allocation call down to the bottom of the call stack (i.e. callee to
+  /// caller order).
+  void addCallStack(AllocationType AllocType, ArrayRef<uint64_t> StackIds);
+
+  /// Add the call stack context along with its allocation type from the MIB
+  /// metadata to the Trie.
+  void addCallStack(MDNode *MIB);
+
+  /// Build and attach the minimal necessary MIB metadata. If the alloc has a
+  /// single allocation type, add a function attribute instead. The reason for
+  /// adding an attribute in this case is that it matches how the behavior for
+  /// allocation calls will be communicated to lib call simplification after
+  /// cloning or another optimization to distinguish the allocation types,
+  /// which is lower overhead and more direct than maintaining this metadata.
+  /// Returns true if memprof metadata attached, false if not (attribute added).
+  bool buildAndAttachMIBMetadata(CallBase *CI);
+};
+
+} // end namespace memprof
+} // end namespace llvm
+
+#endif
diff --git a/llvm/lib/Analysis/CMakeLists.txt b/llvm/lib/Analysis/CMakeLists.txt
--- a/llvm/lib/Analysis/CMakeLists.txt
+++ b/llvm/lib/Analysis/CMakeLists.txt
@@ -98,6 +98,7 @@
   MemoryBuiltins.cpp
   MemoryDependenceAnalysis.cpp
   MemoryLocation.cpp
+  MemoryProfileInfo.cpp
   MemorySSA.cpp
   MemorySSAUpdater.cpp
   ModelUnderTrainingRunner.cpp
diff --git a/llvm/lib/Analysis/MemoryProfileInfo.cpp b/llvm/lib/Analysis/MemoryProfileInfo.cpp
new file mode 100644
--- /dev/null
+++ b/llvm/lib/Analysis/MemoryProfileInfo.cpp
@@ -0,0 +1,226 @@
+//===-- MemoryProfileInfo.cpp - memory profile info ------------------------==//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file contains utilities to analyze memory profile information.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/Analysis/MemoryProfileInfo.h"
+#include "llvm/Support/CommandLine.h"
+
+using namespace llvm;
+using namespace llvm::memprof;
+
+#define DEBUG_TYPE "memory-profile-info"
+
+// Upper bound on accesses per byte for marking an allocation cold.
+cl::opt<float> MemProfAccessesPerByteColdThreshold(
+    "memprof-accesses-per-byte-cold-threshold", cl::init(10.0), cl::Hidden,
+    cl::desc("The threshold the accesses per byte must be under to consider "
+             "an allocation cold"));
+
+// Lower bound on lifetime to mark an allocation cold (in addition to accesses
+// per byte above). This is to avoid pessimizing short lived objects.
+cl::opt<unsigned> MemProfMinLifetimeColdThreshold(
+    "memprof-min-lifetime-cold-threshold", cl::init(200), cl::Hidden,
+    cl::desc("The minimum lifetime (s) for an allocation to be considered "
+             "cold"));
+
+AllocationType llvm::memprof::getAllocType(uint64_t MaxAccessCount,
+                                           uint64_t MinSize,
+                                           uint64_t MinLifetime) {
+  if (((float)MaxAccessCount) / MinSize < MemProfAccessesPerByteColdThreshold &&
+      // MinLifetime is expected to be in ms, so convert the threshold to ms.
+      MinLifetime >= MemProfMinLifetimeColdThreshold * 1000)
+    return AllocationType::Cold;
+  return AllocationType::NotCold;
+}
+
+MDNode *llvm::memprof::buildCallstackMetadata(ArrayRef<uint64_t> CallStack,
+                                              LLVMContext &Ctx) {
+  std::vector<Metadata *> StackVals;
+  for (auto Id : CallStack) {
+    auto *StackValMD =
+        ValueAsMetadata::get(ConstantInt::get(Type::getInt64Ty(Ctx), Id));
+    StackVals.push_back(StackValMD);
+  }
+  return MDNode::get(Ctx, StackVals);
+}
+
+MDNode *llvm::memprof::getMIBStackNode(const MDNode *MIB) {
+  assert(MIB->getNumOperands() == 2);
+  // The stack metadata is the first operand of each memprof MIB metadata.
+  return cast<MDNode>(MIB->getOperand(0));
+}
+
+AllocationType llvm::memprof::getMIBAllocType(const MDNode *MIB) {
+  assert(MIB->getNumOperands() == 2);
+  // The allocation type is currently the second operand of each memprof
+  // MIB metadata. This will need to change as we add additional allocation
+  // types that can be applied based on the allocation profile data.
+  auto *MDS = dyn_cast<MDString>(MIB->getOperand(1));
+  assert(MDS);
+  if (MDS->getString().equals("cold"))
+    return AllocationType::Cold;
+  return AllocationType::NotCold;
+}
+
+static std::string getAllocTypeAttributeString(AllocationType Type) {
+  switch (Type) {
+  case AllocationType::NotCold:
+    return "notcold";
+    break;
+  case AllocationType::Cold:
+    return "cold";
+    break;
+  default:
+    assert(false && "Unexpected alloc type");
+  }
+  llvm_unreachable("invalid alloc type");
+}
+
+static void addAllocTypeAttribute(LLVMContext &Ctx, CallBase *CI,
+                                  AllocationType AllocType) {
+  auto AllocTypeString = getAllocTypeAttributeString(AllocType);
+  auto A = llvm::Attribute::get(Ctx, "memprof", AllocTypeString);
+  CI->addFnAttr(A);
+}
+
+static bool hasSingleAllocType(uint8_t AllocTypes) {
+  const unsigned NumAllocTypes = countPopulation(AllocTypes);
+  assert(NumAllocTypes != 0);
+  return NumAllocTypes == 1;
+}
+
+void CallStackTrie::addCallStack(AllocationType AllocType,
+                                 ArrayRef<uint64_t> StackIds) {
+  bool First = true;
+  CallStackTrieNode *Curr = nullptr;
+  for (auto StackId : StackIds) {
+    // If this is the first stack frame, add or update alloc node.
+    if (First) {
+      First = false;
+      if (Alloc) {
+        assert(AllocStackId == StackId);
+        Alloc->AllocTypes |= static_cast<uint8_t>(AllocType);
+      } else {
+        AllocStackId = StackId;
+        Alloc = new CallStackTrieNode(AllocType);
+      }
+      Curr = Alloc;
+      continue;
+    }
+    // Update existing caller node if it exists.
+    auto Next = Curr->Callers.find(StackId);
+    if (Next != Curr->Callers.end()) {
+      Curr = Next->second;
+      Curr->AllocTypes |= static_cast<uint8_t>(AllocType);
+      continue;
+    }
+    // Otherwise add a new caller node.
+    auto *New = new CallStackTrieNode(AllocType);
+    Curr->Callers[StackId] = New;
+    Curr = New;
+  }
+  assert(Curr);
+}
+
+void CallStackTrie::addCallStack(MDNode *MIB) {
+  MDNode *StackMD = getMIBStackNode(MIB);
+  assert(StackMD);
+  std::vector<uint64_t> CallStack;
+  CallStack.reserve(StackMD->getNumOperands());
+  for (auto &MIBStackIter : StackMD->operands()) {
+    auto *StackId = mdconst::dyn_extract<ConstantInt>(MIBStackIter);
+    assert(StackId);
+    CallStack.push_back(StackId->getZExtValue());
+  }
+  addCallStack(getMIBAllocType(MIB), CallStack);
+}
+
+static MDNode *createMIBNode(LLVMContext &Ctx,
+                             std::vector<uint64_t> &MIBCallStack,
+                             AllocationType AllocType) {
+  std::vector<Metadata *> MIBPayload(
+      {buildCallstackMetadata(MIBCallStack, Ctx)});
+  MIBPayload.push_back(
+      MDString::get(Ctx, getAllocTypeAttributeString(AllocType)));
+  return MDNode::get(Ctx, MIBPayload);
+}
+
+// Recursive helper to trim contexts and create metadata nodes.
+// Caller should have pushed Node's loc to MIBCallStack. Doing this in the
+// caller makes it simpler to handle the many early returns in this method.
+bool CallStackTrie::buildMIBNodes(CallStackTrieNode *Node, LLVMContext &Ctx,
+                                  std::vector<uint64_t> &MIBCallStack,
+                                  std::vector<Metadata *> &MIBNodes,
+                                  bool CalleeHasAmbiguousCallerContext) {
+  // Trim context below the first node in a prefix with a single alloc type.
+  // Add an MIB record for the current call stack prefix.
+  if (hasSingleAllocType(Node->AllocTypes)) {
+    MIBNodes.push_back(
+        createMIBNode(Ctx, MIBCallStack, (AllocationType)Node->AllocTypes));
+    return true;
+  }
+
+  // We don't have a single allocation for all the contexts sharing this prefix,
+  // so recursively descend into callers in trie.
+  if (!Node->Callers.empty()) {
+    bool NodeHasAmbiguousCallerContext = Node->Callers.size() > 1;
+    bool AddedMIBNodesForAllCallerContexts = true;
+    for (auto &Caller : Node->Callers) {
+      MIBCallStack.push_back(Caller.first);
+      AddedMIBNodesForAllCallerContexts &=
+          buildMIBNodes(Caller.second, Ctx, MIBCallStack, MIBNodes,
+                        NodeHasAmbiguousCallerContext);
+      // Remove Caller.
+      MIBCallStack.pop_back();
+    }
+    if (AddedMIBNodesForAllCallerContexts)
+      return true;
+    // We expect that the callers should be forced to add MIBs to disambiguate
+    // the context in this case (see below).
+    assert(!NodeHasAmbiguousCallerContext);
+  }
+
+  // If we reached here, then this node does not have a single allocation type,
+  // and we didn't add metadata for a longer call stack prefix including any of
+  // Node's callers. That means we never hit a single allocation type along all
+  // call stacks with this prefix. This can happen due to recursion collapsing
+  // or the stack being deeper than tracked by the profiler runtime, leading to
+  // contexts with different allocation types being merged. In that case, we
+  // trim the context just below the deepest context split, which is this
+  // node if the callee has an ambiguous caller context (multiple callers),
+  // since the recursive calls above returned false. Conservatively give it
+  // non-cold allocation type.
+  if (!CalleeHasAmbiguousCallerContext)
+    return false;
+  MIBNodes.push_back(createMIBNode(Ctx, MIBCallStack, AllocationType::NotCold));
+  return true;
+}
+
+// Build and attach the minimal necessary MIB metadata. If the alloc has a
+// single allocation type, add a function attribute instead. Returns true if
+// memprof metadata attached, false if not (attribute added).
+bool CallStackTrie::buildAndAttachMIBMetadata(CallBase *CI) {
+  auto &Ctx = CI->getContext();
+  if (hasSingleAllocType(Alloc->AllocTypes)) {
+    addAllocTypeAttribute(Ctx, CI, (AllocationType)Alloc->AllocTypes);
+    return false;
+  }
+  std::vector<uint64_t> MIBCallStack;
+  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));
+  return true;
+}
diff --git a/llvm/unittests/Analysis/CMakeLists.txt b/llvm/unittests/Analysis/CMakeLists.txt
--- a/llvm/unittests/Analysis/CMakeLists.txt
+++ b/llvm/unittests/Analysis/CMakeLists.txt
@@ -38,6 +38,7 @@
   LoopInfoTest.cpp
   LoopNestTest.cpp
   MemoryBuiltinsTest.cpp
+  MemoryProfileInfoTest.cpp
   MemorySSATest.cpp
   MLModelRunnerTest.cpp
   PhiValuesTest.cpp
diff --git a/llvm/unittests/Analysis/MemoryProfileInfoTest.cpp b/llvm/unittests/Analysis/MemoryProfileInfoTest.cpp
new file mode 100644
--- /dev/null
+++ b/llvm/unittests/Analysis/MemoryProfileInfoTest.cpp
@@ -0,0 +1,362 @@
+//===- MemoryProfileInfoTest.cpp - Memory Profile Info Unit Tests-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/Analysis/MemoryProfileInfo.h"
+#include "llvm/AsmParser/Parser.h"
+#include "llvm/IR/Instructions.h"
+#include "llvm/IR/LLVMContext.h"
+#include "llvm/IR/Module.h"
+#include "llvm/Support/CommandLine.h"
+#include "llvm/Support/SourceMgr.h"
+#include "gtest/gtest.h"
+#include <cstring>
+
+using namespace llvm;
+using namespace llvm::memprof;
+
+extern cl::opt<float> MemProfAccessesPerByteColdThreshold;
+extern cl::opt<unsigned> MemProfMinLifetimeColdThreshold;
+
+namespace {
+
+class MemoryProfileInfoTest : public testing::Test {
+protected:
+  std::unique_ptr<Module> makeLLVMModule(LLVMContext &C, const char *IR) {
+    SMDiagnostic Err;
+    std::unique_ptr<Module> Mod = parseAssemblyString(IR, Err, C);
+    if (!Mod)
+      Err.print("MemoryProfileInfoTest", errs());
+    return Mod;
+  }
+
+  // This looks for a call that has the given value name, which
+  // is the name of the value being assigned the call return value.
+  CallBase *findCall(Function &F, const char *Name = nullptr) {
+    for (auto &BB : F)
+      for (auto &I : BB)
+        if (auto *CB = dyn_cast<CallBase>(&I))
+          if (!Name || CB->getName() == Name)
+            return CB;
+    return nullptr;
+  }
+};
+
+// Test getAllocType helper.
+// Basic checks on the allocation type for values just above and below
+// the thresholds.
+TEST_F(MemoryProfileInfoTest, GetAllocType) {
+  // Long lived with more accesses per byte than threshold is not cold.
+  EXPECT_EQ(
+      getAllocType(/*MaxAccessCount=*/MemProfAccessesPerByteColdThreshold + 1,
+                   /*MinSize=*/1,
+                   /*MinLifetime=*/MemProfMinLifetimeColdThreshold * 1000 + 1),
+      AllocationType::NotCold);
+  // Long lived with less accesses per byte than threshold is cold.
+  EXPECT_EQ(
+      getAllocType(/*MaxAccessCount=*/MemProfAccessesPerByteColdThreshold - 1,
+                   /*MinSize=*/1,
+                   /*MinLifetime=*/MemProfMinLifetimeColdThreshold * 1000 + 1),
+      AllocationType::Cold);
+  // Short lived with more accesses per byte than threshold is not cold.
+  EXPECT_EQ(
+      getAllocType(/*MaxAccessCount=*/MemProfAccessesPerByteColdThreshold + 1,
+                   /*MinSize=*/1,
+                   /*MinLifetime=*/MemProfMinLifetimeColdThreshold * 1000 - 1),
+      AllocationType::NotCold);
+  // Short lived with less accesses per byte than threshold is not cold.
+  EXPECT_EQ(
+      getAllocType(/*MaxAccessCount=*/MemProfAccessesPerByteColdThreshold - 1,
+                   /*MinSize=*/1,
+                   /*MinLifetime=*/MemProfMinLifetimeColdThreshold * 1000 - 1),
+      AllocationType::NotCold);
+}
+
+// Test buildCallstackMetadata helper.
+TEST_F(MemoryProfileInfoTest, BuildCallStackMD) {
+  LLVMContext C;
+  MDNode *CallStack = buildCallstackMetadata({1, 2, 3}, C);
+  ASSERT_EQ(CallStack->getNumOperands(), 3u);
+  unsigned ExpectedId = 1;
+  for (auto &Op : CallStack->operands()) {
+    auto *StackId = mdconst::dyn_extract<ConstantInt>(Op);
+    EXPECT_EQ(StackId->getZExtValue(), ExpectedId++);
+  }
+}
+
+// Test CallStackTrie::addCallStack interface taking allocation type and list of
+// call stack ids.
+// Check that allocations with a single allocation type along all call stacks
+// get an attribute instead of memprof metadata.
+TEST_F(MemoryProfileInfoTest, Attribute) {
+  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:
+  %call1 = call noalias dereferenceable_or_null(40) i8* @malloc(i64 noundef 40)
+  %0 = bitcast i8* %call1 to i32*
+  %call2 = call noalias dereferenceable_or_null(40) i8* @malloc(i64 noundef 40)
+  %1 = bitcast i8* %call2 to i32*
+  ret i32* %1
+}
+declare dso_local noalias noundef i8* @malloc(i64 noundef)
+)IR");
+
+  Function *Func = M->getFunction("test");
+
+  // First call has all cold contexts.
+  CallStackTrie Trie1;
+  Trie1.addCallStack(AllocationType::Cold, {1, 2});
+  Trie1.addCallStack(AllocationType::Cold, {1, 3, 4});
+  CallBase *Call1 = findCall(*Func, "call1");
+  Trie1.buildAndAttachMIBMetadata(Call1);
+
+  EXPECT_FALSE(Call1->hasMetadata(LLVMContext::MD_memprof));
+  EXPECT_TRUE(Call1->hasFnAttr("memprof"));
+  EXPECT_EQ(Call1->getFnAttr("memprof").getValueAsString(), "cold");
+
+  // Second call has all non-cold contexts.
+  CallStackTrie Trie2;
+  Trie2.addCallStack(AllocationType::NotCold, {5, 6});
+  Trie2.addCallStack(AllocationType::NotCold, {5, 7, 8});
+  CallBase *Call2 = findCall(*Func, "call2");
+  Trie2.buildAndAttachMIBMetadata(Call2);
+
+  EXPECT_FALSE(Call2->hasMetadata(LLVMContext::MD_memprof));
+  EXPECT_TRUE(Call2->hasFnAttr("memprof"));
+  EXPECT_EQ(Call2->getFnAttr("memprof").getValueAsString(), "notcold");
+}
+
+// Test CallStackTrie::addCallStack interface taking allocation type and list of
+// call stack ids.
+// Test that an allocation call reached by both cold and non cold call stacks
+// gets memprof metadata representing the different allocation type contexts.
+TEST_F(MemoryProfileInfoTest, ColdAndNotColdMIB) {
+  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});
+  Trie.addCallStack(AllocationType::NotCold, {1, 3});
+
+  CallBase *Call = findCall(*Func, "call");
+  Trie.buildAndAttachMIBMetadata(Call);
+
+  EXPECT_FALSE(Call->hasFnAttr("memprof"));
+  EXPECT_TRUE(Call->hasMetadata(LLVMContext::MD_memprof));
+  MDNode *MemProfMD = Call->getMetadata(LLVMContext::MD_memprof);
+  ASSERT_EQ(MemProfMD->getNumOperands(), 2u);
+  for (auto &MIBOp : MemProfMD->operands()) {
+    MDNode *MIB = dyn_cast<MDNode>(MIBOp);
+    MDNode *StackMD = getMIBStackNode(MIB);
+    ASSERT_NE(StackMD, nullptr);
+    ASSERT_EQ(StackMD->getNumOperands(), 2u);
+    auto *StackId = mdconst::dyn_extract<ConstantInt>(StackMD->getOperand(0));
+    ASSERT_EQ(StackId->getZExtValue(), 1u);
+    StackId = mdconst::dyn_extract<ConstantInt>(StackMD->getOperand(1));
+    if (StackId->getZExtValue() == 2u)
+      EXPECT_EQ(getMIBAllocType(MIB), AllocationType::Cold);
+    else {
+      ASSERT_EQ(StackId->getZExtValue(), 3u);
+      EXPECT_EQ(getMIBAllocType(MIB), AllocationType::NotCold);
+    }
+  }
+}
+
+// Test CallStackTrie::addCallStack interface taking allocation type and list of
+// call stack ids.
+// Test that an allocation call reached by multiple call stacks has memprof
+// metadata with the contexts trimmed to the minimum context required to
+// identify the allocation type.
+TEST_F(MemoryProfileInfoTest, TrimmedMIBContext) {
+  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;
+  // We should be able to trim the following two and combine into a single MIB
+  // with the cold context {1, 2}.
+  Trie.addCallStack(AllocationType::Cold, {1, 2, 3});
+  Trie.addCallStack(AllocationType::Cold, {1, 2, 4});
+  // We should be able to trim the following two and combine into a single MIB
+  // with the non-cold context {1, 5}.
+  Trie.addCallStack(AllocationType::NotCold, {1, 5, 6});
+  Trie.addCallStack(AllocationType::NotCold, {1, 5, 7});
+
+  CallBase *Call = findCall(*Func, "call");
+  Trie.buildAndAttachMIBMetadata(Call);
+
+  EXPECT_FALSE(Call->hasFnAttr("memprof"));
+  EXPECT_TRUE(Call->hasMetadata(LLVMContext::MD_memprof));
+  MDNode *MemProfMD = Call->getMetadata(LLVMContext::MD_memprof);
+  ASSERT_EQ(MemProfMD->getNumOperands(), 2u);
+  for (auto &MIBOp : MemProfMD->operands()) {
+    MDNode *MIB = dyn_cast<MDNode>(MIBOp);
+    MDNode *StackMD = getMIBStackNode(MIB);
+    ASSERT_NE(StackMD, nullptr);
+    ASSERT_EQ(StackMD->getNumOperands(), 2u);
+    auto *StackId = mdconst::dyn_extract<ConstantInt>(StackMD->getOperand(0));
+    EXPECT_EQ(StackId->getZExtValue(), 1u);
+    StackId = mdconst::dyn_extract<ConstantInt>(StackMD->getOperand(1));
+    if (StackId->getZExtValue() == 2u)
+      EXPECT_EQ(getMIBAllocType(MIB), AllocationType::Cold);
+    else {
+      ASSERT_EQ(StackId->getZExtValue(), 5u);
+      EXPECT_EQ(getMIBAllocType(MIB), AllocationType::NotCold);
+    }
+  }
+}
+
+// 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.
+TEST_F(MemoryProfileInfoTest, SimplifyMIBToAttribute) {
+  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:
+  %call1 = call noalias dereferenceable_or_null(40) i8* @malloc(i64 noundef 40), !memprof !0
+  %0 = bitcast i8* %call1 to i32*
+  %call2 = call noalias dereferenceable_or_null(40) i8* @malloc(i64 noundef 40), !memprof !3
+  %1 = bitcast i8* %call2 to i32*
+  ret i32* %1
+}
+declare dso_local noalias noundef i8* @malloc(i64 noundef)
+!0 = !{!1}
+!1 = !{!2, !"cold"}
+!2 = !{i64 1, i64 2, i64 3}
+!3 = !{!4}
+!4 = !{!5, !"notcold"}
+!5 = !{i64 4, i64 5, i64 6, i64 7}
+)IR");
+
+  Function *Func = M->getFunction("test");
+
+  // First call has all cold contexts.
+  CallStackTrie Trie1;
+  CallBase *Call1 = findCall(*Func, "call1");
+  MDNode *MemProfMD1 = Call1->getMetadata(LLVMContext::MD_memprof);
+  ASSERT_EQ(MemProfMD1->getNumOperands(), 1u);
+  MDNode *MIB1 = dyn_cast<MDNode>(MemProfMD1->getOperand(0));
+  Trie1.addCallStack(MIB1);
+  Trie1.buildAndAttachMIBMetadata(Call1);
+
+  EXPECT_TRUE(Call1->hasFnAttr("memprof"));
+  EXPECT_EQ(Call1->getFnAttr("memprof").getValueAsString(), "cold");
+
+  // Second call has all non-cold contexts.
+  CallStackTrie Trie2;
+  CallBase *Call2 = findCall(*Func, "call2");
+  MDNode *MemProfMD2 = Call2->getMetadata(LLVMContext::MD_memprof);
+  ASSERT_EQ(MemProfMD2->getNumOperands(), 1u);
+  MDNode *MIB2 = dyn_cast<MDNode>(MemProfMD2->getOperand(0));
+  Trie2.addCallStack(MIB2);
+  Trie2.buildAndAttachMIBMetadata(Call2);
+
+  EXPECT_TRUE(Call2->hasFnAttr("memprof"));
+  EXPECT_EQ(Call2->getFnAttr("memprof").getValueAsString(), "notcold");
+}
+
+// Test CallStackTrie::addCallStack interface taking memprof MIB metadata.
+// Test that allocations annotated with memprof metadata with multiple call
+// stacks gets new memprof metadata with the contexts trimmed to the minimum
+// context required to identify the allocation type.
+TEST_F(MemoryProfileInfoTest, ReTrimMIBContext) {
+  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), !memprof !0
+  %0 = bitcast i8* %call to i32*
+  ret i32* %0
+}
+declare dso_local noalias noundef i8* @malloc(i64 noundef)
+!0 = !{!1, !3, !5, !7}
+!1 = !{!2, !"cold"}
+!2 = !{i64 1, i64 2, i64 3}
+!3 = !{!4, !"cold"}
+!4 = !{i64 1, i64 2, i64 4}
+!5 = !{!6, !"notcold"}
+!6 = !{i64 1, i64 5, i64 6}
+!7 = !{!8, !"notcold"}
+!8 = !{i64 1, i64 5, i64 7}
+)IR");
+
+  Function *Func = M->getFunction("test");
+
+  CallStackTrie Trie;
+  ASSERT_TRUE(Trie.empty());
+  CallBase *Call = findCall(*Func, "call");
+  MDNode *MemProfMD = Call->getMetadata(LLVMContext::MD_memprof);
+  for (auto &MIBOp : MemProfMD->operands()) {
+    MDNode *MIB = dyn_cast<MDNode>(MIBOp);
+    Trie.addCallStack(MIB);
+  }
+  ASSERT_FALSE(Trie.empty());
+  Trie.buildAndAttachMIBMetadata(Call);
+
+  // We should be able to trim the first two and combine into a single MIB
+  // with the cold context {1, 2}.
+  // We should be able to trim the second two and combine into a single MIB
+  // with the non-cold context {1, 5}.
+
+  EXPECT_FALSE(Call->hasFnAttr("memprof"));
+  EXPECT_TRUE(Call->hasMetadata(LLVMContext::MD_memprof));
+  MemProfMD = Call->getMetadata(LLVMContext::MD_memprof);
+  ASSERT_EQ(MemProfMD->getNumOperands(), 2u);
+  for (auto &MIBOp : MemProfMD->operands()) {
+    MDNode *MIB = dyn_cast<MDNode>(MIBOp);
+    MDNode *StackMD = getMIBStackNode(MIB);
+    ASSERT_NE(StackMD, nullptr);
+    ASSERT_EQ(StackMD->getNumOperands(), 2u);
+    auto *StackId = mdconst::dyn_extract<ConstantInt>(StackMD->getOperand(0));
+    EXPECT_EQ(StackId->getZExtValue(), 1u);
+    StackId = mdconst::dyn_extract<ConstantInt>(StackMD->getOperand(1));
+    if (StackId->getZExtValue() == 2u)
+      EXPECT_EQ(getMIBAllocType(MIB), AllocationType::Cold);
+    else {
+      ASSERT_EQ(StackId->getZExtValue(), 5u);
+      EXPECT_EQ(getMIBAllocType(MIB), AllocationType::NotCold);
+    }
+  }
+}
+
+} // end anonymous namespace