[MemProf] Optionally print or record the profiled sizes of allocations

[teresajohnson]

This is the first step in being able to track the total profiled sizes
of allocations successfully marked as cold.

Under a new option -memprof-report-hinted-sizes:

• For unambiguous (non-context-sensitive) allocations, print the
  profiled size and the allocation coldness, along with a hash of the
  allocation's location (to allow for deduplication across modules or
  inline instances).
• For context sensitive allocations, add the size as a 3rd operand on
  the MIB metadata. A follow on patch will propagate this through to the
  thin link where the sizes will be reported for each context after
  cloning.

[llvmbot]

@llvm/pr-subscribers-pgo
@llvm/pr-subscribers-llvm-ir
@llvm/pr-subscribers-llvm-analysis

@llvm/pr-subscribers-llvm-transforms

Author: Teresa Johnson (teresajohnson)

Changes

This is the first step in being able to track the total profiled sizes
of allocations successfully marked as cold.

Under a new option -memprof-report-hinted-sizes:

• For unambiguous (non-context-sensitive) allocations, print the
  profiled size and the allocation coldness, along with a hash of the
  allocation's location (to allow for deduplication across modules or
  inline instances).
• For context sensitive allocations, add the size as a 3rd operand on
  the MIB metadata. A follow on patch will propagate this through to the
  thin link where the sizes will be reported for each context after
  cloning.

---

Full diff: https://github.com/llvm/llvm-project/pull/98248.diff

6 Files Affected:

• (modified) llvm/include/llvm/Analysis/MemoryProfileInfo.h (+9-3)
• (modified) llvm/lib/Analysis/MemoryProfileInfo.cpp (+34-10)
• (modified) llvm/lib/IR/Verifier.cpp (+7-3)
• (modified) llvm/lib/Transforms/Instrumentation/MemProfiler.cpp (+9-2)
• (modified) llvm/test/Transforms/PGOProfile/memprof.ll (+15)
• (modified) llvm/test/Verifier/memprof-metadata-bad.ll (+2-2)

diff --git a/llvm/include/llvm/Analysis/MemoryProfileInfo.h b/llvm/include/llvm/Analysis/MemoryProfileInfo.h
index 355bff46f6279..75ddf5cacf74d 100644
--- a/llvm/include/llvm/Analysis/MemoryProfileInfo.h
+++ b/llvm/include/llvm/Analysis/MemoryProfileInfo.h
@@ -37,6 +37,10 @@ MDNode *getMIBStackNode(const MDNode *MIB);
 /// Returns the allocation type from an MIB metadata node.
 AllocationType getMIBAllocType(const MDNode *MIB);
 
+/// Returns the total size from an MIB metadata node, or 0 if it was not
+/// recorded.
+uint64_t getMIBTotalSize(const MDNode *MIB);
+
 /// Returns the string to use in attributes with the given type.
 std::string getAllocTypeAttributeString(AllocationType Type);
 
@@ -54,10 +58,11 @@ class CallStackTrie {
     // Allocation types for call context sharing the context prefix at this
     // node.
     uint8_t AllocTypes;
+    uint64_t TotalSize;
     // 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)) {}
+    CallStackTrieNode(AllocationType Type, uint64_t TotalSize)
+        : AllocTypes(static_cast<uint8_t>(Type)), TotalSize(TotalSize) {}
   };
 
   // The node for the allocation at the root.
@@ -90,7 +95,8 @@ class CallStackTrie {
   /// 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);
+  void addCallStack(AllocationType AllocType, ArrayRef<uint64_t> StackIds,
+                    uint64_t TotalSize = 0);
 
   /// Add the call stack context along with its allocation type from the MIB
   /// metadata to the Trie.
diff --git a/llvm/lib/Analysis/MemoryProfileInfo.cpp b/llvm/lib/Analysis/MemoryProfileInfo.cpp
index 5c09ba9462716..c81c65fd4e691 100644
--- a/llvm/lib/Analysis/MemoryProfileInfo.cpp
+++ b/llvm/lib/Analysis/MemoryProfileInfo.cpp
@@ -41,6 +41,10 @@ cl::opt<unsigned> MemProfMinAveLifetimeAccessDensityHotThreshold(
     cl::desc("The minimum TotalLifetimeAccessDensity / AllocCount for an "
              "allocation to be considered hot"));
 
+cl::opt<bool> MemProfReportHintedSizes(
+    "memprof-report-hinted-sizes", cl::init(false), cl::Hidden,
+    cl::desc("Report total allocation sizes of hinted allocations"));
+
 AllocationType llvm::memprof::getAllocType(uint64_t TotalLifetimeAccessDensity,
                                            uint64_t AllocCount,
                                            uint64_t TotalLifetime) {
@@ -74,13 +78,13 @@ MDNode *llvm::memprof::buildCallstackMetadata(ArrayRef<uint64_t> CallStack,
 }
 
 MDNode *llvm::memprof::getMIBStackNode(const MDNode *MIB) {
-  assert(MIB->getNumOperands() == 2);
+  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);
+  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.
@@ -94,6 +98,12 @@ AllocationType llvm::memprof::getMIBAllocType(const MDNode *MIB) {
   return AllocationType::NotCold;
 }
 
+uint64_t llvm::memprof::getMIBTotalSize(const MDNode *MIB) {
+  if (MIB->getNumOperands() < 3)
+    return 0;
+  return mdconst::dyn_extract<ConstantInt>(MIB->getOperand(2))->getZExtValue();
+}
+
 std::string llvm::memprof::getAllocTypeAttributeString(AllocationType Type) {
   switch (Type) {
   case AllocationType::NotCold:
@@ -125,7 +135,8 @@ bool llvm::memprof::hasSingleAllocType(uint8_t AllocTypes) {
 }
 
 void CallStackTrie::addCallStack(AllocationType AllocType,
-                                 ArrayRef<uint64_t> StackIds) {
+                                 ArrayRef<uint64_t> StackIds,
+                                 uint64_t TotalSize) {
   bool First = true;
   CallStackTrieNode *Curr = nullptr;
   for (auto StackId : StackIds) {
@@ -135,9 +146,10 @@ void CallStackTrie::addCallStack(AllocationType AllocType,
       if (Alloc) {
         assert(AllocStackId == StackId);
         Alloc->AllocTypes |= static_cast<uint8_t>(AllocType);
+        Alloc->TotalSize += TotalSize;
       } else {
         AllocStackId = StackId;
-        Alloc = new CallStackTrieNode(AllocType);
+        Alloc = new CallStackTrieNode(AllocType, TotalSize);
       }
       Curr = Alloc;
       continue;
@@ -147,10 +159,11 @@ void CallStackTrie::addCallStack(AllocationType AllocType,
     if (Next != Curr->Callers.end()) {
       Curr = Next->second;
       Curr->AllocTypes |= static_cast<uint8_t>(AllocType);
+      Curr->TotalSize += TotalSize;
       continue;
     }
     // Otherwise add a new caller node.
-    auto *New = new CallStackTrieNode(AllocType);
+    auto *New = new CallStackTrieNode(AllocType, TotalSize);
     Curr->Callers[StackId] = New;
     Curr = New;
   }
@@ -167,16 +180,19 @@ void CallStackTrie::addCallStack(MDNode *MIB) {
     assert(StackId);
     CallStack.push_back(StackId->getZExtValue());
   }
-  addCallStack(getMIBAllocType(MIB), CallStack);
+  addCallStack(getMIBAllocType(MIB), CallStack, getMIBTotalSize(MIB));
 }
 
 static MDNode *createMIBNode(LLVMContext &Ctx,
                              std::vector<uint64_t> &MIBCallStack,
-                             AllocationType AllocType) {
+                             AllocationType AllocType, uint64_t TotalSize) {
   std::vector<Metadata *> MIBPayload(
       {buildCallstackMetadata(MIBCallStack, Ctx)});
   MIBPayload.push_back(
       MDString::get(Ctx, getAllocTypeAttributeString(AllocType)));
+  if (TotalSize)
+    MIBPayload.push_back(ValueAsMetadata::get(
+        ConstantInt::get(Type::getInt64Ty(Ctx), TotalSize)));
   return MDNode::get(Ctx, MIBPayload);
 }
 
@@ -190,8 +206,8 @@ bool CallStackTrie::buildMIBNodes(CallStackTrieNode *Node, LLVMContext &Ctx,
   // 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));
+    MIBNodes.push_back(createMIBNode(
+        Ctx, MIBCallStack, (AllocationType)Node->AllocTypes, Node->TotalSize));
     return true;
   }
 
@@ -227,7 +243,8 @@ bool CallStackTrie::buildMIBNodes(CallStackTrieNode *Node, LLVMContext &Ctx,
   // non-cold allocation type.
   if (!CalleeHasAmbiguousCallerContext)
     return false;
-  MIBNodes.push_back(createMIBNode(Ctx, MIBCallStack, AllocationType::NotCold));
+  MIBNodes.push_back(createMIBNode(Ctx, MIBCallStack, AllocationType::NotCold,
+                                   Node->TotalSize));
   return true;
 }
 
@@ -238,6 +255,13 @@ bool CallStackTrie::buildAndAttachMIBMetadata(CallBase *CI) {
   auto &Ctx = CI->getContext();
   if (hasSingleAllocType(Alloc->AllocTypes)) {
     addAllocTypeAttribute(Ctx, CI, (AllocationType)Alloc->AllocTypes);
+    if (MemProfReportHintedSizes) {
+      assert(Alloc->TotalSize);
+      errs() << "Total size for allocation with location hash " << AllocStackId
+             << " and single alloc type "
+             << getAllocTypeAttributeString((AllocationType)Alloc->AllocTypes)
+             << ": " << Alloc->TotalSize << "\n";
+    }
     return false;
   }
   std::vector<uint64_t> MIBCallStack;
diff --git a/llvm/lib/IR/Verifier.cpp b/llvm/lib/IR/Verifier.cpp
index c98f61d555140..1794dd3cc028b 100644
--- a/llvm/lib/IR/Verifier.cpp
+++ b/llvm/lib/IR/Verifier.cpp
@@ -4934,10 +4934,14 @@ void Verifier::visitMemProfMetadata(Instruction &I, MDNode *MD) {
     MDNode *StackMD = dyn_cast<MDNode>(MIB->getOperand(0));
     visitCallStackMetadata(StackMD);
 
-    // Check that remaining operands are MDString.
-    Check(llvm::all_of(llvm::drop_begin(MIB->operands()),
+    // Check that remaining operands, except possibly the last, are MDString.
+    Check(llvm::all_of(MIB->operands().drop_front().drop_back(),
                        [](const MDOperand &Op) { return isa<MDString>(Op); }),
-          "Not all !memprof MemInfoBlock operands 1 to N are MDString", MIB);
+          "Not all !memprof MemInfoBlock operands 1 to N-1 are MDString", MIB);
+    // The last operand might be the total profiled size so can be an integer.
+    auto &LastOperand = MIB->operands().back();
+    Check(isa<MDString>(LastOperand) || mdconst::hasa<ConstantInt>(LastOperand),
+          "Last !memprof MemInfoBlock operand not MDString or int", MIB);
   }
 }
 
diff --git a/llvm/lib/Transforms/Instrumentation/MemProfiler.cpp b/llvm/lib/Transforms/Instrumentation/MemProfiler.cpp
index 5d5e205d45462..2c5d749d4a67a 100644
--- a/llvm/lib/Transforms/Instrumentation/MemProfiler.cpp
+++ b/llvm/lib/Transforms/Instrumentation/MemProfiler.cpp
@@ -161,6 +161,8 @@ static cl::opt<bool>
                                      "context in this module's profiles"),
                             cl::Hidden, cl::init(false));
 
+extern cl::opt<bool> MemProfReportHintedSizes;
+
 // Instrumentation statistics
 STATISTIC(NumInstrumentedReads, "Number of instrumented reads");
 STATISTIC(NumInstrumentedWrites, "Number of instrumented writes");
@@ -712,7 +714,12 @@ static AllocationType addCallStack(CallStackTrie &AllocTrie,
   auto AllocType = getAllocType(AllocInfo->Info.getTotalLifetimeAccessDensity(),
                                 AllocInfo->Info.getAllocCount(),
                                 AllocInfo->Info.getTotalLifetime());
-  AllocTrie.addCallStack(AllocType, StackIds);
+  uint64_t TotalSize = 0;
+  if (MemProfReportHintedSizes) {
+    TotalSize = AllocInfo->Info.getTotalSize();
+    assert(TotalSize);
+  }
+  AllocTrie.addCallStack(AllocType, StackIds, TotalSize);
   return AllocType;
 }
 
@@ -1055,4 +1062,4 @@ PreservedAnalyses MemProfUsePass::run(Module &M, ModuleAnalysisManager &AM) {
   }
 
   return PreservedAnalyses::none();
-}
\ No newline at end of file
+}
diff --git a/llvm/test/Transforms/PGOProfile/memprof.ll b/llvm/test/Transforms/PGOProfile/memprof.ll
index 4a87f4f9d7449..b6407f7b123d3 100644
--- a/llvm/test/Transforms/PGOProfile/memprof.ll
+++ b/llvm/test/Transforms/PGOProfile/memprof.ll
@@ -63,6 +63,9 @@
 ;; give both memprof and pgo metadata.
 ; RUN: opt < %s -passes='pgo-instr-use,memprof-use<profile-filename=%t.pgomemprofdata>' -pgo-test-profile-file=%t.pgomemprofdata -pgo-warn-missing-function -S 2>&1 | FileCheck %s --check-prefixes=MEMPROF,ALL,PGO
 
+;; Check that the total sizes are reported if requested.
+; RUN: opt < %s -passes='memprof-use<profile-filename=%t.memprofdata>' -pgo-warn-missing-function -S -memprof-report-hinted-sizes 2>&1 | FileCheck %s --check-prefixes=TOTALSIZES
+
 ; MEMPROFMATCHINFO: MemProf notcold context with id 1093248920606587996 has total profiled size 10 is matched
 ; MEMPROFMATCHINFO: MemProf notcold context with id 5725971306423925017 has total profiled size 10 is matched
 ; MEMPROFMATCHINFO: MemProf notcold context with id 6792096022461663180 has total profiled size 10 is matched
@@ -331,6 +334,18 @@ for.end:                                          ; preds = %for.cond
 ; MEMPROF: ![[C10]] = !{i64 2061451396820446691}
 ; MEMPROF: ![[C11]] = !{i64 1544787832369987002}
 
+;; For non-context sensitive allocations that get attributes we emit a message
+;; with the allocation hash, type, and size in bytes.
+; TOTALSIZES: Total size for allocation with location hash 6792096022461663180 and single alloc type notcold: 10
+; TOTALSIZES: Total size for allocation with location hash 15737101490731057601 and single alloc type cold: 10
+;; For context sensitive allocations the size in bytes is included on the MIB
+;; metadata.
+; TOTALSIZES: !"cold", i64 10}
+; TOTALSIZES: !"cold", i64 10}
+; TOTALSIZES: !"notcold", i64 10}
+; TOTALSIZES: !"cold", i64 20}
+; TOTALSIZES: !"notcold", i64 10}
+
 
 ; MEMPROFNOCOLINFO: #[[A1]] = { builtin allocsize(0) "memprof"="notcold" }
 ; MEMPROFNOCOLINFO: #[[A2]] = { builtin allocsize(0) "memprof"="cold" }
diff --git a/llvm/test/Verifier/memprof-metadata-bad.ll b/llvm/test/Verifier/memprof-metadata-bad.ll
index 83a10764d1808..f4f1f6bb0a463 100644
--- a/llvm/test/Verifier/memprof-metadata-bad.ll
+++ b/llvm/test/Verifier/memprof-metadata-bad.ll
@@ -43,8 +43,8 @@ declare dso_local noalias noundef ptr @malloc(i64 noundef)
 !6 = !{i64 0}
 !7 = !{!8}
 ; CHECK: call stack metadata should have at least 1 operand
-; CHECK: Not all !memprof MemInfoBlock operands 1 to N are MDString
-!8 = !{!0, !"default", i64 0}
+; CHECK: Not all !memprof MemInfoBlock operands 1 to N-1 are MDString
+!8 = !{!0, !"default", i64 0, i64 5}
 !9 = !{i64 123}
 ; CHECK: call stack metadata operand should be constant integer
 !10 = !{!"wrongtype"}

[snehasish]

lgtm