From d403516a577a48f23150a42808a5cf6dda1d7702 Mon Sep 17 00:00:00 2001
From: Kazu Hirata <kazu@google.com>
Date: Wed, 4 Dec 2024 10:23:21 -0800
Subject: [PATCH 1/2] [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 fe017913f6de2..88666087007c9 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 2f8589bbfbb96..0f3af93b09685 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

From 3abef1b0960c2d7c3e7897d48779463e6c3964cb Mon Sep 17 00:00:00 2001
From: Kazu Hirata <kazu@google.com>
Date: Wed, 4 Dec 2024 16:47:31 -0800
Subject: [PATCH 2/2] Address comments.

---
 llvm/include/llvm/ProfileData/MemProf.h    |  1 -
 llvm/unittests/ProfileData/MemProfTest.cpp | 20 +++++++++++---------
 2 files changed, 11 insertions(+), 10 deletions(-)

diff --git a/llvm/include/llvm/ProfileData/MemProf.h b/llvm/include/llvm/ProfileData/MemProf.h
index 88666087007c9..5af460efb6ddf 100644
--- a/llvm/include/llvm/ProfileData/MemProf.h
+++ b/llvm/include/llvm/ProfileData/MemProf.h
@@ -1179,7 +1179,6 @@ template <> struct MappingTraits<memprof::Frame> {
   }
 
   // Request the inline notation for brevity:
-  //
   //   { Function: 123, LineOffset: 11, Column: 10; IsInlineFrame: true }
   static const bool flow = true;
 };
diff --git a/llvm/unittests/ProfileData/MemProfTest.cpp b/llvm/unittests/ProfileData/MemProfTest.cpp
index 0f3af93b09685..83383eeaa3fe6 100644
--- a/llvm/unittests/ProfileData/MemProfTest.cpp
+++ b/llvm/unittests/ProfileData/MemProfTest.cpp
@@ -808,15 +808,20 @@ TEST(MemProf, YAMLParser) {
               ElementsAre(hashCallStack(CS3), hashCallStack(CS4)));
 }
 
-TEST(MemProf, YAMLWriterFrame) {
-  Frame F(12, 34, 56, true);
-
+template <typename T> std::string serializeInYAML(T &Val) {
   std::string Out;
   llvm::raw_string_ostream OS(Out);
   llvm::yaml::Output Yout(OS);
-  Yout << F;
+  Yout << Val;
+  return Out;
+}
+
+TEST(MemProf, YAMLWriterFrame) {
+  Frame F(11, 22, 33, true);
+
+  std::string Out = serializeInYAML(F);
   EXPECT_EQ(Out, R"YAML(---
-{ Function: 12, LineOffset: 34, Column: 56, Inline: true }
+{ Function: 11, LineOffset: 22, Column: 33, Inline: true }
 ...
 )YAML");
 }
@@ -829,10 +834,7 @@ TEST(MemProf, YAMLWriterMIB) {
   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;
+  std::string Out = serializeInYAML(PMIB);
   EXPECT_EQ(Out, R"YAML(---
 AllocCount:      111
 TotalSize:       222