[llvm] 1dad624 - [MemProf] Add memprof metadata related analysis utilities

Teresa Johnson via llvm-commits llvm-commits at lists.llvm.org
Thu Jul 21 13:46:24 PDT 2022


Author: Teresa Johnson
Date: 2022-07-21T13:46:01-07:00
New Revision: 1dad6247d275c5787cbc3131174d9c4aa543e653

URL: https://github.com/llvm/llvm-project/commit/1dad6247d275c5787cbc3131174d9c4aa543e653
DIFF: https://github.com/llvm/llvm-project/commit/1dad6247d275c5787cbc3131174d9c4aa543e653.diff

LOG: [MemProf] Add memprof metadata related analysis utilities

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.

Differential Revision: https://reviews.llvm.org/D128854

Added: 
    llvm/include/llvm/Analysis/MemoryProfileInfo.h
    llvm/lib/Analysis/MemoryProfileInfo.cpp
    llvm/unittests/Analysis/MemoryProfileInfoTest.cpp

Modified: 
    llvm/lib/Analysis/CMakeLists.txt
    llvm/unittests/Analysis/CMakeLists.txt

Removed: 
    


################################################################################
diff  --git a/llvm/include/llvm/Analysis/MemoryProfileInfo.h b/llvm/include/llvm/Analysis/MemoryProfileInfo.h
new file mode 100644
index 0000000000000..1b12e78eaeba3
--- /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 
diff erent behavior with 
diff erent 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
index 821d1938b2163..e59725c990170 100644
--- a/llvm/lib/Analysis/CMakeLists.txt
+++ b/llvm/lib/Analysis/CMakeLists.txt
@@ -98,6 +98,7 @@ add_llvm_component_library(LLVMAnalysis
   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
index 0000000000000..3d11cb81226e3
--- /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 
diff erent 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
index ed2cb81055b38..af45304fdaf10 100644
--- a/llvm/unittests/Analysis/CMakeLists.txt
+++ b/llvm/unittests/Analysis/CMakeLists.txt
@@ -38,6 +38,7 @@ add_llvm_unittest_with_input_files(AnalysisTests
   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
index 0000000000000..a3549aca81e51
--- /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 
diff erent 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


        


More information about the llvm-commits mailing list