record data access profile payload length

Author: mingmingl-llvm

llvm/lib/ProfileData/IndexedMemProfData.cpp

@@ -228,8 +228,10 @@ static Error writeMemProfRadixTreeBased(
   OS.write(0ULL); // Reserve space for the memprof call stack payload offset.
   OS.write(0ULL); // Reserve space for the memprof record payload offset.
   OS.write(0ULL); // Reserve space for the memprof record table offset.
-  if (Version >= memprof::Version4)
+  if (Version >= memprof::Version4) {
     OS.write(0ULL); // Reserve space for the data access profile offset.
+    OS.write(0ULL); // Reserve space for the size of data access profiles.
+  }
 
   auto Schema = memprof::getHotColdSchema();
   if (MemProfFullSchema)
@@ -257,12 +259,14 @@ static Error writeMemProfRadixTreeBased(
       OS, MemProfData.Records, &Schema, Version, &MemProfCallStackIndexes);
 
   uint64_t DataAccessProfOffset = 0;
+  uint64_t DataAccessProfLength = 0;
   if (DataAccessProfileData.has_value()) {
     assert(Version >= memprof::Version4 &&
            "Data access profiles are added starting from v4");
     DataAccessProfOffset = OS.tell();
     if (Error E = (*DataAccessProfileData).get().serialize(OS))
       return E;
+    DataAccessProfLength = OS.tell() - DataAccessProfOffset;
   }
 
   // Verify that the computation for the number of elements in the call stack
@@ -271,13 +275,15 @@ static Error writeMemProfRadixTreeBased(
              NumElements * sizeof(memprof::LinearFrameId) ==
          RecordPayloadOffset);
 
-  SmallVector<uint64_t, 4> Header = {
+  SmallVector<uint64_t, 8> Header = {
       CallStackPayloadOffset,
       RecordPayloadOffset,
       RecordTableOffset,
   };
-  if (Version >= memprof::Version4)
+  if (Version >= memprof::Version4) {
     Header.push_back(DataAccessProfOffset);
+    Header.push_back(DataAccessProfLength);
+  }
   OS.patch({{HeaderUpdatePos, Header}});
 
   return Error::success();
@@ -394,9 +400,13 @@ Error IndexedMemProfReader::deserializeRadixTreeBased(
       support::endian::readNext<uint64_t, llvm::endianness::little>(Ptr);
 
   uint64_t DataAccessProfOffset = 0;
-  if (Version == memprof::Version4)
+  uint64_t DataAccessProfLength = 0;
+  if (Version == memprof::Version4) {
     DataAccessProfOffset =
         support::endian::readNext<uint64_t, llvm::endianness::little>(Ptr);
+    DataAccessProfLength =
+        support::endian::readNext<uint64_t, llvm::endianness::little>(Ptr);
+  }
 
   // Read the schema.
   auto SchemaOr = memprof::readMemProfSchema(Ptr);
@@ -419,7 +429,7 @@ Error IndexedMemProfReader::deserializeRadixTreeBased(
       /*Payload=*/Start + RecordPayloadOffset,
       /*Base=*/Start, memprof::RecordLookupTrait(Version, Schema)));
 
-  if (DataAccessProfOffset > RecordTableOffset) {
+  if (DataAccessProfLength > 0) {
     DataAccessProfileData =
         std::make_unique<data_access_prof::DataAccessProfData>();
     const unsigned char *DAPPtr = Start + DataAccessProfOffset;