[memprof] Take Schema into account in PortableMemInfoBlock::serializedSize (#89824)

PortableMemInfoBlock::{serialize,deserialize} take Schema into
account, allowing us to serialize/deserialize a subset of the fields.
However, PortableMemInfoBlock::serializedSize does not.  That is, it
assumes that all fields are always serialized and deserialized.  In
other words, if we choose to serialize/deserialize a subset of the
fields, serializedSize would claim more storage than we actually need.

This patch fixes the problem by teaching serializedSize to take Schema
into account.  For now, this patch has no effect on the actual indexed
MemProf profile because we serialize/deserialize all fields, but that
might change in the future.

Aside from check-llvm, I tested this patch by verifying that
llvm-profdata generates bit-wise identical files for each version for
a large raw MemProf file I have.

Author: kazutakahirata

llvm/include/llvm/ProfileData/MemProf.h

@@ -138,11 +138,22 @@ struct PortableMemInfoBlock {
     return !operator==(Other);
   }
 
-  static constexpr size_t serializedSize() {
+  static size_t serializedSize(const MemProfSchema &Schema) {
     size_t Result = 0;
-#define MIBEntryDef(NameTag, Name, Type) Result += sizeof(Type);
+
+    for (const Meta Id : Schema) {
+      switch (Id) {
+#define MIBEntryDef(NameTag, Name, Type)                                       \
+  case Meta::Name: {                                                           \
+    Result += sizeof(Type);                                                    \
+  } break;
 #include "llvm/ProfileData/MIBEntryDef.inc"
 #undef MIBEntryDef
+      default:
+        llvm_unreachable("Unknown meta type id, invalid input?");
+      }
+    }
+
     return Result;
   }
 
@@ -292,7 +303,8 @@ struct IndexedAllocationInfo {
       : CallStack(CS.begin(), CS.end()), CSId(CSId), Info(MB) {}
 
   // Returns the size in bytes when this allocation info struct is serialized.
-  size_t serializedSize(IndexedVersion Version) const;
+  size_t serializedSize(const MemProfSchema &Schema,
+                        IndexedVersion Version) const;
 
   bool operator==(const IndexedAllocationInfo &Other) const {
     if (Other.Info != Info)
@@ -367,7 +379,8 @@ struct IndexedMemProfRecord {
     CallSites.append(Other.CallSites);
   }
 
-  size_t serializedSize(IndexedVersion Version) const;
+  size_t serializedSize(const MemProfSchema &Schema,
+                        IndexedVersion Version) const;
 
   bool operator==(const IndexedMemProfRecord &Other) const {
     if (Other.AllocSites != AllocSites)
@@ -535,7 +548,7 @@ class RecordWriterTrait {
     endian::Writer LE(Out, llvm::endianness::little);
     offset_type N = sizeof(K);
     LE.write<offset_type>(N);
-    offset_type M = V.serializedSize(Version);
+    offset_type M = V.serializedSize(*Schema, Version);
     LE.write<offset_type>(M);
     return std::make_pair(N, M);
   }

llvm/lib/ProfileData/MemProf.cpp

@@ -10,42 +10,46 @@
 
 namespace llvm {
 namespace memprof {
-static size_t serializedSizeV0(const IndexedAllocationInfo &IAI) {
+static size_t serializedSizeV0(const IndexedAllocationInfo &IAI,
+                               const MemProfSchema &Schema) {
   size_t Size = 0;
   // The number of frames to serialize.
   Size += sizeof(uint64_t);
   // The callstack frame ids.
   Size += sizeof(FrameId) * IAI.CallStack.size();
   // The size of the payload.
-  Size += PortableMemInfoBlock::serializedSize();
+  Size += PortableMemInfoBlock::serializedSize(Schema);
   return Size;
 }
 
-static size_t serializedSizeV2(const IndexedAllocationInfo &IAI) {
+static size_t serializedSizeV2(const IndexedAllocationInfo &IAI,
+                               const MemProfSchema &Schema) {
   size_t Size = 0;
   // The CallStackId
   Size += sizeof(CallStackId);
   // The size of the payload.
-  Size += PortableMemInfoBlock::serializedSize();
+  Size += PortableMemInfoBlock::serializedSize(Schema);
   return Size;
 }
 
-size_t IndexedAllocationInfo::serializedSize(IndexedVersion Version) const {
+size_t IndexedAllocationInfo::serializedSize(const MemProfSchema &Schema,
+                                             IndexedVersion Version) const {
   switch (Version) {
   case Version0:
   case Version1:
-    return serializedSizeV0(*this);
+    return serializedSizeV0(*this, Schema);
   case Version2:
-    return serializedSizeV2(*this);
+    return serializedSizeV2(*this, Schema);
   }
   llvm_unreachable("unsupported MemProf version");
 }
 
-static size_t serializedSizeV0(const IndexedMemProfRecord &Record) {
+static size_t serializedSizeV0(const IndexedMemProfRecord &Record,
+                               const MemProfSchema &Schema) {
   // The number of alloc sites to serialize.
   size_t Result = sizeof(uint64_t);
   for (const IndexedAllocationInfo &N : Record.AllocSites)
-    Result += N.serializedSize(Version0);
+    Result += N.serializedSize(Schema, Version0);
 
   // The number of callsites we have information for.
   Result += sizeof(uint64_t);
@@ -57,11 +61,12 @@ static size_t serializedSizeV0(const IndexedMemProfRecord &Record) {
   return Result;
 }
 
-static size_t serializedSizeV2(const IndexedMemProfRecord &Record) {
+static size_t serializedSizeV2(const IndexedMemProfRecord &Record,
+                               const MemProfSchema &Schema) {
   // The number of alloc sites to serialize.
   size_t Result = sizeof(uint64_t);
   for (const IndexedAllocationInfo &N : Record.AllocSites)
-    Result += N.serializedSize(Version2);
+    Result += N.serializedSize(Schema, Version2);
 
   // The number of callsites we have information for.
   Result += sizeof(uint64_t);
@@ -70,13 +75,14 @@ static size_t serializedSizeV2(const IndexedMemProfRecord &Record) {
   return Result;
 }
 
-size_t IndexedMemProfRecord::serializedSize(IndexedVersion Version) const {
+size_t IndexedMemProfRecord::serializedSize(const MemProfSchema &Schema,
+                                            IndexedVersion Version) const {
   switch (Version) {
   case Version0:
   case Version1:
-    return serializedSizeV0(*this);
+    return serializedSizeV0(*this, Schema);
   case Version2:
-    return serializedSizeV2(*this);
+    return serializedSizeV2(*this, Schema);
   }
   llvm_unreachable("unsupported MemProf version");
 }
@@ -156,7 +162,7 @@ static IndexedMemProfRecord deserializeV0(const MemProfSchema &Schema,
     }
     Node.CSId = hashCallStack(Node.CallStack);
     Node.Info.deserialize(Schema, Ptr);
-    Ptr += PortableMemInfoBlock::serializedSize();
+    Ptr += PortableMemInfoBlock::serializedSize(Schema);
     Record.AllocSites.push_back(Node);
   }
 
@@ -193,7 +199,7 @@ static IndexedMemProfRecord deserializeV2(const MemProfSchema &Schema,
     IndexedAllocationInfo Node;
     Node.CSId = endian::readNext<CallStackId, llvm::endianness::little>(Ptr);
     Node.Info.deserialize(Schema, Ptr);
-    Ptr += PortableMemInfoBlock::serializedSize();
+    Ptr += PortableMemInfoBlock::serializedSize(Schema);
     Record.AllocSites.push_back(Node);
   }