[MemProf] Add memprof metadata related analysis utilities

[tejohnson]

Summary

Adds a number of utilities that are used to help create and update
memprof related metadata. These will be used during profile matching
and annotation, as well as by the inliner when updating the metadata.
Also adds unit tests for the utilities.

See also related RFCs:
RFC: Sanitizer-based Heap Profiler [1]
RFC: A binary serialization format for MemProf [2]
RFC: IR metadata format for MemProf [3]
(Note that the IR metadata format has changed from the RFC during
implementation, as described in the preceeding patch adding the basic
metadata and verification support.)

Depends on D128141.

[snehasish]

Thanks for splitting out the util change and adding unit tests. It made it much easier to review!

llvm/include/llvm/Analysis/MemoryProfileInfo.h

42

I think defining static inline functions in a header will create an instance for each module this file is included in. I think this is a discouraged pattern. Should we move this and the other helper below to the cpp instead?

99

This function is unused in this patch. Perhaps you meant to use it at L100? It would flip the if condition though.

106

I think this function is public for testing only and callers should prefer using the interface addCallStack(MDNode *MIB);? If so can you document this in the comment.

113

Can you elaborate on the rationale for function attributes for single allocation types? It does look like its simpler to access in the code. Is it also because it reduces overhead? If so can you document the rationale here.

llvm/lib/Analysis/MemoryProfileInfo.cpp

78

I think this switch can be replaced with popcount which is a little easier to follow.

const unsigned PopCount = countPopulation(AllocTypes);
assert(PopCount != 0);
return PopCount > 1;

https://github.com/llvm/llvm-project/blob/655bea4226b401a11164f99c6344e38d8742b8e4/llvm/include/llvm/Support/MathExtras.h#L567

129

Can we just inline the getMIBAllocType call into L135 and avoid the single use AllocType var?

130

Prefer a smallvector or maybe add CallStack.reserve since we know the number of operands?

132

nit: A descriptive name like StackHash would make it easier to follow.

166

!Node->Callers.empty()

212

Add a comment that addCallstack has not been called yet?
Also !Alloc->Callers.empty() since we just want to assert that Callers has some entries.

216

Consider inlining the comment into the assert, e.g.
assert(MIBCallStack.size() == 1 && "Comment");

llvm/unittests/Analysis/MemoryProfileInfoTest.cpp

175

Calling this StackId would make the test easier to follow. Also applies for the other tests.

176

This should probably be ASSERT since if the first stackid isn't 1 then it doesn't make much sense to continue.

263

I guess having a repeated stack id is feasible with recursive calls during profiling. It doesn't matter much since its adding a function attribute in this case.

300

I think this test is redundant since at this point we have already tested

• the ability to trim contexts in TrimmedMIBContext
• the ability to extract the metadata node and call the underlying addCallStack() correctly in SimplifyMIBToAttribute.

Up to you if you want to keep it but I'm slightly in favour of removing this test if it doesn't cover any new logic.

[tejohnson]

Address comments

llvm/include/llvm/Analysis/MemoryProfileInfo.h

99

It's used in a follow on. I think it is more natural to use the Alloc variable directly there on L100. So instead I added a couple calls, before and after calling addCallStack on a new Trie, to invoke and test it.

106

No, it is used during matching in D128142.

113

Yes it is lighter weight and simpler to communicate to the lib call handling. It is the way the cloning code I've been testing tags allocation calls once we create allocation call clones with single allocation types. Added comment.

llvm/lib/Analysis/MemoryProfileInfo.cpp

130

Decided to use reserve since we know the number of entries we'll need.

llvm/unittests/Analysis/MemoryProfileInfoTest.cpp

263

I think this is a typo in the test. In the follow on matching patch we collapse recursion. I've removed it here for simplicity.

300

I wanted to have a test that put these together since is functionality that will be used to update after inlining. I agree the components are tested separately already, but I kind of like having this for completeness. If you prefer though I can remove it.

[snehasish]

lgtm, will go through D128142 soon. Thanks!

llvm/unittests/Analysis/MemoryProfileInfoTest.cpp

300

Keeping as is for completeness sounds reasonable to me.

[tejohnson]

Fix bug with use of popcount results that was breaking a few tests.

llvm/include/llvm/Analysis/MemoryProfileInfo.h

+//===- 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

llvm/lib/Analysis/CMakeLists.txt

add_llvm_component_library(LLVMAnalysis
   LoopInfo.cpp
   LoopPass.cpp
   MLInlineAdvisor.cpp
   MemDepPrinter.cpp
   MemDerefPrinter.cpp
   MemoryBuiltins.cpp
   MemoryDependenceAnalysis.cpp
   MemoryLocation.cpp
+  MemoryProfileInfo.cpp
   MemorySSA.cpp
   MemorySSAUpdater.cpp
   ModelUnderTrainingRunner.cpp
   ModuleDebugInfoPrinter.cpp
   ModuleSummaryAnalysis.cpp
   MustExecute.cpp
   NoInferenceModelRunner.cpp
   ObjCARCAliasAnalysis.cpp

llvm/lib/Analysis/MemoryProfileInfo.cpp

+//===-- 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;
+}

llvm/unittests/Analysis/CMakeLists.txt

add_llvm_unittest_with_input_files(AnalysisTests
   InlineCostTest.cpp
   IRSimilarityIdentifierTest.cpp
   IVDescriptorsTest.cpp
   LazyCallGraphTest.cpp
   LoadsTest.cpp
   LoopInfoTest.cpp
   LoopNestTest.cpp
   MemoryBuiltinsTest.cpp
+  MemoryProfileInfoTest.cpp
   MemorySSATest.cpp
   MLModelRunnerTest.cpp
   PhiValuesTest.cpp
   ProfileSummaryInfoTest.cpp
   ScalarEvolutionTest.cpp
   VectorFunctionABITest.cpp
   SparsePropagation.cpp
   TargetLibraryInfoTest.cpp

llvm/unittests/Analysis/MemoryProfileInfoTest.cpp

+//===- 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