[llvm] [memprof] Update YAML traits for writer purposes (PR #118720)

[Kazu Hirata via llvm-commits]

https://github.com/kazutakahirata created https://github.com/llvm/llvm-project/pull/118720

For Frames, we prefer the inline notation for the brevity.

For PortableMemInfoBlock, we go through all member fields and print
out those that are populated.


>From d403516a577a48f23150a42808a5cf6dda1d7702 Mon Sep 17 00:00:00 2001
From: Kazu Hirata <kazu at google.com>
Date: Wed, 4 Dec 2024 10:23:21 -0800
Subject: [PATCH] [memprof] Update YAML traits for writer purposes

For Frames, we prefer the inline notation for the brevity.

For PortableMemInfoBlock, we go through all member fields and print
out those that are populated.
---
 llvm/include/llvm/ProfileData/MemProf.h    | 14 +++++++--
 llvm/unittests/ProfileData/MemProfTest.cpp | 34 ++++++++++++++++++++++
 2 files changed, 46 insertions(+), 2 deletions(-)

diff --git a/llvm/include/llvm/ProfileData/MemProf.h b/llvm/include/llvm/ProfileData/MemProf.h
index fe017913f6de24..88666087007c91 100644
--- a/llvm/include/llvm/ProfileData/MemProf.h
+++ b/llvm/include/llvm/ProfileData/MemProf.h
@@ -1177,6 +1177,11 @@ template <> struct MappingTraits<memprof::Frame> {
     (void)Column;
     (void)IsInlineFrame;
   }
+
+  // Request the inline notation for brevity:
+  //
+  //   { Function: 123, LineOffset: 11, Column: 10; IsInlineFrame: true }
+  static const bool flow = true;
 };
 
 template <> struct CustomMappingTraits<memprof::PortableMemInfoBlock> {
@@ -1201,8 +1206,13 @@ template <> struct CustomMappingTraits<memprof::PortableMemInfoBlock> {
     Io.setError("Key is not a valid validation event");
   }
 
-  static void output(IO &Io, memprof::PortableMemInfoBlock &VI) {
-    llvm_unreachable("To be implemented");
+  static void output(IO &Io, memprof::PortableMemInfoBlock &MIB) {
+    auto Schema = MIB.getSchema();
+#define MIBEntryDef(NameTag, Name, Type)                                       \
+  if (Schema.test(llvm::to_underlying(memprof::Meta::Name)))                   \
+    Io.mapRequired(#Name, MIB.Name);
+#include "llvm/ProfileData/MIBEntryDef.inc"
+#undef MIBEntryDef
   }
 };
 
diff --git a/llvm/unittests/ProfileData/MemProfTest.cpp b/llvm/unittests/ProfileData/MemProfTest.cpp
index 2f8589bbfbb962..0f3af93b096852 100644
--- a/llvm/unittests/ProfileData/MemProfTest.cpp
+++ b/llvm/unittests/ProfileData/MemProfTest.cpp
@@ -807,4 +807,38 @@ TEST(MemProf, YAMLParser) {
   EXPECT_THAT(Record.CallSiteIds,
               ElementsAre(hashCallStack(CS3), hashCallStack(CS4)));
 }
+
+TEST(MemProf, YAMLWriterFrame) {
+  Frame F(12, 34, 56, true);
+
+  std::string Out;
+  llvm::raw_string_ostream OS(Out);
+  llvm::yaml::Output Yout(OS);
+  Yout << F;
+  EXPECT_EQ(Out, R"YAML(---
+{ Function: 12, LineOffset: 34, Column: 56, Inline: true }
+...
+)YAML");
+}
+
+TEST(MemProf, YAMLWriterMIB) {
+  MemInfoBlock MIB;
+  MIB.AllocCount = 111;
+  MIB.TotalSize = 222;
+  MIB.TotalLifetime = 333;
+  MIB.TotalLifetimeAccessDensity = 444;
+  PortableMemInfoBlock PMIB(MIB, llvm::memprof::getHotColdSchema());
+
+  std::string Out;
+  llvm::raw_string_ostream OS(Out);
+  llvm::yaml::Output Yout(OS);
+  Yout << PMIB;
+  EXPECT_EQ(Out, R"YAML(---
+AllocCount:      111
+TotalSize:       222
+TotalLifetime:   333
+TotalLifetimeAccessDensity: 444
+...
+)YAML");
+}
 } // namespace