[llvm] [MemProf] Optionally discard small non-cold contexts (PR #139113)

[Teresa Johnson via llvm-commits]

================
@@ -208,21 +222,32 @@ void CallStackTrie::addCallStack(MDNode *MIB) {
 
 static MDNode *createMIBNode(LLVMContext &Ctx, ArrayRef<uint64_t> MIBCallStack,
                              AllocationType AllocType,
-                             ArrayRef<ContextTotalSize> ContextSizeInfo) {
+                             ArrayRef<ContextTotalSize> ContextSizeInfo,
+                             uint64_t &TotalBytes, uint64_t &ColdBytes) {
   SmallVector<Metadata *> MIBPayload(
       {buildCallstackMetadata(MIBCallStack, Ctx)});
   MIBPayload.push_back(
       MDString::get(Ctx, getAllocTypeAttributeString(AllocType)));
   if (!ContextSizeInfo.empty()) {
     for (const auto &[FullStackId, TotalSize] : ContextSizeInfo) {
-      auto *FullStackIdMD = ValueAsMetadata::get(
-          ConstantInt::get(Type::getInt64Ty(Ctx), FullStackId));
-      auto *TotalSizeMD = ValueAsMetadata::get(
-          ConstantInt::get(Type::getInt64Ty(Ctx), TotalSize));
-      auto *ContextSizeMD = MDNode::get(Ctx, {FullStackIdMD, TotalSizeMD});
-      MIBPayload.push_back(ContextSizeMD);
+      TotalBytes += TotalSize;
+      if (AllocType == AllocationType::Cold)
+        ColdBytes += TotalSize;
+      // Only add the context size info as metadata if we need it in the thin
+      // link (currently if reporting of hinted sizes is enabled or we have
+      // specified a threshold for marking allocations cold after cloning).
+      if (MemProfReportHintedSizes || MinClonedColdBytePercent < 100) {
+        auto *FullStackIdMD = ValueAsMetadata::get(
+            ConstantInt::get(Type::getInt64Ty(Ctx), FullStackId));
+        auto *TotalSizeMD = ValueAsMetadata::get(
+            ConstantInt::get(Type::getInt64Ty(Ctx), TotalSize));
+        auto *ContextSizeMD = MDNode::get(Ctx, {FullStackIdMD, TotalSizeMD});
+        MIBPayload.push_back(ContextSizeMD);
+      }
     }
   }
+  assert(MinCallsiteColdBytePercent >= 100 ||
+         (!ContextSizeInfo.empty() && TotalBytes > 0));
----------------
teresajohnson wrote:

Done on the early return. For your second question, I debated checking == 100, but decided to more gracefully handle a user that provided a larger value than 100 (in lieu of checking option values somewhere) - the other handling checks if it is < 100, so that will all work fine. Added a comment.

https://github.com/llvm/llvm-project/pull/139113