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