[memprof] Add Version2 of IndexedMemProfRecord serialization

llvm/include/llvm/ProfileData/MemProf.h

@@ -22,6 +22,8 @@ enum IndexedVersion : uint64_t {
   Version0 = 0,
   // Version 1: Added a version field to the header.
   Version1 = 1,
+  // Version 2: Added a call stack table.  Under development.
+  Version2 = 2,
 };
 
 constexpr uint64_t MinimumSupportedVersion = Version0;
@@ -289,10 +291,20 @@ 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() const {
-    return sizeof(uint64_t) + // The number of frames to serialize.
-           sizeof(FrameId) * CallStack.size() +    // The callstack frame ids.
-           PortableMemInfoBlock::serializedSize(); // The size of the payload.
+  size_t serializedSize(IndexedVersion Version = Version0) const {
+    size_t Size = 0;
+    if (Version <= Version1) {
+      // The number of frames to serialize.
+      Size += sizeof(uint64_t);
+      // The callstack frame ids.
+      Size += sizeof(FrameId) * CallStack.size();
+    } else {
+      // The CallStackId
+      Size += sizeof(CallStackId);
+    }
+    // The size of the payload.
+    Size += PortableMemInfoBlock::serializedSize();
+    return Size;
   }
 
   bool operator==(const IndexedAllocationInfo &Other) const {
@@ -306,6 +318,9 @@ struct IndexedAllocationInfo {
       if (Other.CallStack[J] != CallStack[J])
         return false;
     }
+
+    if (Other.CSId != CSId)
+      return false;
     return true;
   }
 
@@ -357,6 +372,9 @@ struct IndexedMemProfRecord {
   // inline location list may include additional entries, users should pick
   // the last entry in the list with the same function GUID.
   llvm::SmallVector<llvm::SmallVector<FrameId>> CallSites;
+  // Conceptually the same as above.  We are going to keep both CallSites and
+  // CallSiteIds while we are transitioning from CallSites to CallSitesIds.
+  llvm::SmallVector<CallStackId> CallSiteIds;
 
   void clear() {
     AllocSites.clear();
@@ -370,17 +388,22 @@ struct IndexedMemProfRecord {
     CallSites.append(Other.CallSites);
   }
 
-  size_t serializedSize() const {
+  size_t serializedSize(IndexedVersion Version = Version0) const {
     size_t Result = sizeof(GlobalValue::GUID);
     for (const IndexedAllocationInfo &N : AllocSites)
-      Result += N.serializedSize();
+      Result += N.serializedSize(Version);
 
     // The number of callsites we have information for.
     Result += sizeof(uint64_t);
-    for (const auto &Frames : CallSites) {
-      // The number of frame ids to serialize.
-      Result += sizeof(uint64_t);
-      Result += Frames.size() * sizeof(FrameId);
+    if (Version <= Version1) {
+      for (const auto &Frames : CallSites) {
+        // The number of frame ids to serialize.
+        Result += sizeof(uint64_t);
+        Result += Frames.size() * sizeof(FrameId);
+      }
+    } else {
+      // The CallStackId
+      Result += CallSiteIds.size() * sizeof(CallStackId);
     }
     return Result;
   }
@@ -401,16 +424,21 @@ struct IndexedMemProfRecord {
       if (CallSites[I] != Other.CallSites[I])
         return false;
     }
+
+    if (Other.CallSiteIds != CallSiteIds)
+      return false;
     return true;
   }
 
   // Serializes the memprof records in \p Records to the ostream \p OS based
   // on the schema provided in \p Schema.
-  void serialize(const MemProfSchema &Schema, raw_ostream &OS);
+  void serialize(const MemProfSchema &Schema, raw_ostream &OS,
+                 IndexedVersion Version = Version0);
 
   // Deserializes memprof records from the Buffer.
   static IndexedMemProfRecord deserialize(const MemProfSchema &Schema,
-                                          const unsigned char *Buffer);
+                                          const unsigned char *Buffer,
+                                          IndexedVersion Version = Version0);
 
   // Returns the GUID for the function name after canonicalization. For
   // memprof, we remove any .llvm suffix added by LTO. MemProfRecords are

llvm/lib/ProfileData/MemProf.cpp

@@ -12,31 +12,42 @@ namespace llvm {
 namespace memprof {
 
 void IndexedMemProfRecord::serialize(const MemProfSchema &Schema,
-                                     raw_ostream &OS) {
+                                     raw_ostream &OS, IndexedVersion Version) {
   using namespace support;
 
   endian::Writer LE(OS, llvm::endianness::little);
 
   LE.write<uint64_t>(AllocSites.size());
   for (const IndexedAllocationInfo &N : AllocSites) {
-    LE.write<uint64_t>(N.CallStack.size());
-    for (const FrameId &Id : N.CallStack)
-      LE.write<FrameId>(Id);
+    if (Version <= Version1) {
+      LE.write<uint64_t>(N.CallStack.size());
+      for (const FrameId &Id : N.CallStack)
+        LE.write<FrameId>(Id);
+    } else {
+      LE.write<CallStackId>(N.CSId);
+    }
     N.Info.serialize(Schema, OS);
   }
 
   // Related contexts.
-  LE.write<uint64_t>(CallSites.size());
-  for (const auto &Frames : CallSites) {
-    LE.write<uint64_t>(Frames.size());
-    for (const FrameId &Id : Frames)
-      LE.write<FrameId>(Id);
+  if (Version <= Version1) {
+    LE.write<uint64_t>(CallSites.size());
+    for (const auto &Frames : CallSites) {
+      LE.write<uint64_t>(Frames.size());
+      for (const FrameId &Id : Frames)
+        LE.write<FrameId>(Id);
+    }
+  } else {
+    LE.write<uint64_t>(CallSiteIds.size());
+    for (const auto &CSId : CallSiteIds)
+      LE.write<CallStackId>(CSId);
   }
 }
 
 IndexedMemProfRecord
 IndexedMemProfRecord::deserialize(const MemProfSchema &Schema,
-                                  const unsigned char *Ptr) {
+                                  const unsigned char *Ptr,
+                                  IndexedVersion Version) {
   using namespace support;
 
   IndexedMemProfRecord Record;
@@ -46,14 +57,20 @@ IndexedMemProfRecord::deserialize(const MemProfSchema &Schema,
       endian::readNext<uint64_t, llvm::endianness::little, unaligned>(Ptr);
   for (uint64_t I = 0; I < NumNodes; I++) {
     IndexedAllocationInfo Node;
-    const uint64_t NumFrames =
-        endian::readNext<uint64_t, llvm::endianness::little, unaligned>(Ptr);
-    for (uint64_t J = 0; J < NumFrames; J++) {
-      const FrameId Id =
-          endian::readNext<FrameId, llvm::endianness::little, unaligned>(Ptr);
-      Node.CallStack.push_back(Id);
+    if (Version <= Version1) {
+      const uint64_t NumFrames =
+          endian::readNext<uint64_t, llvm::endianness::little, unaligned>(Ptr);
+      for (uint64_t J = 0; J < NumFrames; J++) {
+        const FrameId Id =
+            endian::readNext<FrameId, llvm::endianness::little, unaligned>(Ptr);
+        Node.CallStack.push_back(Id);
+      }
+      Node.CSId = hashCallStack(Node.CallStack);
+    } else {
+      Node.CSId =
+          endian::readNext<CallStackId, llvm::endianness::little, unaligned>(
+              Ptr);
     }
-    Node.CSId = hashCallStack(Node.CallStack);
     Node.Info.deserialize(Schema, Ptr);
     Ptr += PortableMemInfoBlock::serializedSize();
     Record.AllocSites.push_back(Node);
@@ -63,16 +80,24 @@ IndexedMemProfRecord::deserialize(const MemProfSchema &Schema,
   const uint64_t NumCtxs =
       endian::readNext<uint64_t, llvm::endianness::little, unaligned>(Ptr);
   for (uint64_t J = 0; J < NumCtxs; J++) {
-    const uint64_t NumFrames =
-        endian::readNext<uint64_t, llvm::endianness::little, unaligned>(Ptr);
-    llvm::SmallVector<FrameId> Frames;
-    Frames.reserve(NumFrames);
-    for (uint64_t K = 0; K < NumFrames; K++) {
-      const FrameId Id =
-          endian::readNext<FrameId, llvm::endianness::little, unaligned>(Ptr);
-      Frames.push_back(Id);
+    if (Version <= Version1) {
+      const uint64_t NumFrames =
+          endian::readNext<uint64_t, llvm::endianness::little, unaligned>(Ptr);
+      llvm::SmallVector<FrameId> Frames;
+      Frames.reserve(NumFrames);
+      for (uint64_t K = 0; K < NumFrames; K++) {
+        const FrameId Id =
+            endian::readNext<FrameId, llvm::endianness::little, unaligned>(Ptr);
+        Frames.push_back(Id);
+      }
+      Record.CallSites.push_back(Frames);
+      Record.CallSiteIds.push_back(hashCallStack(Frames));
+    } else {
+      CallStackId CSId =
+          endian::readNext<CallStackId, llvm::endianness::little, unaligned>(
+              Ptr);
+      Record.CallSiteIds.push_back(CSId);
     }
-    Record.CallSites.push_back(Frames);
   }
 
   return Record;

llvm/unittests/ProfileData/MemProfTest.cpp

@@ -265,7 +265,7 @@ TEST(MemProf, PortableWrapper) {
   EXPECT_EQ(3UL, ReadBlock.getAllocCpuId());
 }
 
-TEST(MemProf, RecordSerializationRoundTrip) {
+TEST(MemProf, RecordSerializationRoundTripVersion0) {
   const MemProfSchema Schema = getFullSchema();
 
   MemInfoBlock Info(/*size=*/16, /*access_count=*/7, /*alloc_timestamp=*/1000,
@@ -284,14 +284,47 @@ TEST(MemProf, RecordSerializationRoundTrip) {
                                    Info);
   }
   Record.CallSites.assign(CallSites);
+  for (const auto &CS : CallSites)
+    Record.CallSiteIds.push_back(llvm::memprof::hashCallStack(CS));
 
   std::string Buffer;
   llvm::raw_string_ostream OS(Buffer);
-  Record.serialize(Schema, OS);
+  Record.serialize(Schema, OS, llvm::memprof::Version0);
   OS.flush();
 
   const IndexedMemProfRecord GotRecord = IndexedMemProfRecord::deserialize(
-      Schema, reinterpret_cast<const unsigned char *>(Buffer.data()));
+      Schema, reinterpret_cast<const unsigned char *>(Buffer.data()),
+      llvm::memprof::Version0);
+
+  EXPECT_EQ(Record, GotRecord);
+}
+
+TEST(MemProf, RecordSerializationRoundTripVerion2) {
+  const MemProfSchema Schema = getFullSchema();
+
+  MemInfoBlock Info(/*size=*/16, /*access_count=*/7, /*alloc_timestamp=*/1000,
+                    /*dealloc_timestamp=*/2000, /*alloc_cpu=*/3,
+                    /*dealloc_cpu=*/4);
+
+  llvm::SmallVector<llvm::memprof::CallStackId> CallStackIds = {0x123, 0x456};
+
+  llvm::SmallVector<llvm::memprof::CallStackId> CallSiteIds = {0x333, 0x444};
+
+  IndexedMemProfRecord Record;
+  for (const auto &CSId : CallStackIds) {
+    // Use the same info block for both allocation sites.
+    Record.AllocSites.emplace_back(llvm::SmallVector<FrameId>(), CSId, Info);
+  }
+  Record.CallSiteIds.assign(CallSiteIds);
+
+  std::string Buffer;
+  llvm::raw_string_ostream OS(Buffer);
+  Record.serialize(Schema, OS, llvm::memprof::Version2);
+  OS.flush();
+
+  const IndexedMemProfRecord GotRecord = IndexedMemProfRecord::deserialize(
+      Schema, reinterpret_cast<const unsigned char *>(Buffer.data()),
+      llvm::memprof::Version2);
 
   EXPECT_EQ(Record, GotRecord);
 }