[memprof] Group MemProf data structures into a struct (NFC) (#92360)

kazutakahirata · web-flow · commit 0dc80e4b2640 · 2024-05-16T10:35:45.000-07:00
This patch groups the three Memprof data structures into a struct
named IndexedMemProfData and teaches InstrProfWriter to use it.  This
way, we can pass IndexedMemProfData to writeMemProf and its helpers
instead of individual data structures.

As a follow-up, we can use the new struct in MemProfReader also.  That
in turn allows loadInput in llvm-profdata to move the MemProf data
into the writer context, saving a few seconds for a large MemProf
profile.
diff --git a/llvm/include/llvm/ProfileData/InstrProfWriter.h b/llvm/include/llvm/ProfileData/InstrProfWriter.h
@@ -51,18 +51,8 @@ class InstrProfWriter {
   SmallVector<TemporalProfTraceTy> TemporalProfTraces;
   std::mt19937 RNG;
 
-  // A map to hold memprof data per function. The lower 64 bits obtained from
-  // the md5 hash of the function name is used to index into the map.
-  llvm::MapVector<GlobalValue::GUID, memprof::IndexedMemProfRecord>
-      MemProfRecordData;
-  // A map to hold frame id to frame mappings. The mappings are used to
-  // convert IndexedMemProfRecord to MemProfRecords with frame information
-  // inline.
-  llvm::MapVector<memprof::FrameId, memprof::Frame> MemProfFrameData;
-
-  // A map to hold call stack id to call stacks.
-  llvm::MapVector<memprof::CallStackId, llvm::SmallVector<memprof::FrameId>>
-      MemProfCallStackData;
+  // The MemProf data.
+  memprof::IndexedMemProfData MemProfData;
 
   // List of binary ids.
   std::vector<llvm::object::BuildID> BinaryIds;
diff --git a/llvm/include/llvm/ProfileData/MemProf.h b/llvm/include/llvm/ProfileData/MemProf.h
@@ -833,6 +833,20 @@ template <typename MapTy> struct CallStackIdConverter {
   }
 };
 
+struct IndexedMemProfData {
+  // A map to hold memprof data per function. The lower 64 bits obtained from
+  // the md5 hash of the function name is used to index into the map.
+  llvm::MapVector<GlobalValue::GUID, IndexedMemProfRecord> RecordData;
+
+  // A map to hold frame id to frame mappings. The mappings are used to
+  // convert IndexedMemProfRecord to MemProfRecords with frame information
+  // inline.
+  llvm::MapVector<FrameId, Frame> FrameData;
+
+  // A map to hold call stack id to call stacks.
+  llvm::MapVector<CallStackId, llvm::SmallVector<FrameId>> CallStackData;
+};
+
 // Verify that each CallStackId is computed with hashCallStack.  This function
 // is intended to help transition from CallStack to CSId in
 // IndexedAllocationInfo.
diff --git a/llvm/lib/ProfileData/InstrProfWriter.cpp b/llvm/lib/ProfileData/InstrProfWriter.cpp
@@ -273,7 +273,7 @@ void InstrProfWriter::addRecord(StringRef Name, uint64_t Hash,
 
 void InstrProfWriter::addMemProfRecord(
     const Function::GUID Id, const memprof::IndexedMemProfRecord &Record) {
-  auto [Iter, Inserted] = MemProfRecordData.insert({Id, Record});
+  auto [Iter, Inserted] = MemProfData.RecordData.insert({Id, Record});
   // If we inserted a new record then we are done.
   if (Inserted) {
     return;
@@ -285,7 +285,7 @@ void InstrProfWriter::addMemProfRecord(
 bool InstrProfWriter::addMemProfFrame(const memprof::FrameId Id,
                                       const memprof::Frame &Frame,
                                       function_ref<void(Error)> Warn) {
-  auto [Iter, Inserted] = MemProfFrameData.insert({Id, Frame});
+  auto [Iter, Inserted] = MemProfData.FrameData.insert({Id, Frame});
   // If a mapping already exists for the current frame id and it does not
   // match the new mapping provided then reset the existing contents and bail
   // out. We don't support the merging of memprof data whose Frame -> Id
@@ -302,7 +302,7 @@ bool InstrProfWriter::addMemProfCallStack(
     const memprof::CallStackId CSId,
     const llvm::SmallVector<memprof::FrameId> &CallStack,
     function_ref<void(Error)> Warn) {
-  auto [Iter, Inserted] = MemProfCallStackData.insert({CSId, CallStack});
+  auto [Iter, Inserted] = MemProfData.CallStackData.insert({CSId, CallStack});
   // If a mapping already exists for the current call stack id and it does not
   // match the new mapping provided then reset the existing contents and bail
   // out. We don't support the merging of memprof data whose CallStack -> Id
@@ -389,22 +389,22 @@ void InstrProfWriter::mergeRecordsFromWriter(InstrProfWriter &&IPW,
   addTemporalProfileTraces(IPW.TemporalProfTraces,
                            IPW.TemporalProfTraceStreamSize);
 
-  MemProfFrameData.reserve(IPW.MemProfFrameData.size());
-  for (auto &[FrameId, Frame] : IPW.MemProfFrameData) {
+  MemProfData.FrameData.reserve(IPW.MemProfData.FrameData.size());
+  for (auto &[FrameId, Frame] : IPW.MemProfData.FrameData) {
     // If we weren't able to add the frame mappings then it doesn't make sense
     // to try to merge the records from this profile.
     if (!addMemProfFrame(FrameId, Frame, Warn))
       return;
   }
 
-  MemProfCallStackData.reserve(IPW.MemProfCallStackData.size());
-  for (auto &[CSId, CallStack] : IPW.MemProfCallStackData) {
+  MemProfData.CallStackData.reserve(IPW.MemProfData.CallStackData.size());
+  for (auto &[CSId, CallStack] : IPW.MemProfData.CallStackData) {
     if (!addMemProfCallStack(CSId, CallStack, Warn))
       return;
   }
 
-  MemProfRecordData.reserve(IPW.MemProfRecordData.size());
-  for (auto &[GUID, Record] : IPW.MemProfRecordData) {
+  MemProfData.RecordData.reserve(IPW.MemProfData.RecordData.size());
+  for (auto &[GUID, Record] : IPW.MemProfData.RecordData) {
     addMemProfRecord(GUID, Record);
   }
 }
@@ -499,11 +499,8 @@ static uint64_t writeMemProfCallStacks(
   return CallStackTableGenerator.Emit(OS.OS);
 }
 
-static Error writeMemProfV0(
-    ProfOStream &OS,
-    llvm::MapVector<GlobalValue::GUID, memprof::IndexedMemProfRecord>
-        &MemProfRecordData,
-    llvm::MapVector<memprof::FrameId, memprof::Frame> &MemProfFrameData) {
+static Error writeMemProfV0(ProfOStream &OS,
+                            memprof::IndexedMemProfData &MemProfData) {
   uint64_t HeaderUpdatePos = OS.tell();
   OS.write(0ULL); // Reserve space for the memprof record table offset.
   OS.write(0ULL); // Reserve space for the memprof frame payload offset.
@@ -512,23 +509,20 @@ static Error writeMemProfV0(
   auto Schema = memprof::getFullSchema();
   writeMemProfSchema(OS, Schema);
 
-  uint64_t RecordTableOffset =
-      writeMemProfRecords(OS, MemProfRecordData, &Schema, memprof::Version0);
+  uint64_t RecordTableOffset = writeMemProfRecords(OS, MemProfData.RecordData,
+                                                   &Schema, memprof::Version0);
 
   uint64_t FramePayloadOffset = OS.tell();
-  uint64_t FrameTableOffset = writeMemProfFrames(OS, MemProfFrameData);
+  uint64_t FrameTableOffset = writeMemProfFrames(OS, MemProfData.FrameData);
 
   uint64_t Header[] = {RecordTableOffset, FramePayloadOffset, FrameTableOffset};
   OS.patch({{HeaderUpdatePos, Header, std::size(Header)}});
 
   return Error::success();
 }
 
-static Error writeMemProfV1(
-    ProfOStream &OS,
-    llvm::MapVector<GlobalValue::GUID, memprof::IndexedMemProfRecord>
-        &MemProfRecordData,
-    llvm::MapVector<memprof::FrameId, memprof::Frame> &MemProfFrameData) {
+static Error writeMemProfV1(ProfOStream &OS,
+                            memprof::IndexedMemProfData &MemProfData) {
   OS.write(memprof::Version1);
   uint64_t HeaderUpdatePos = OS.tell();
   OS.write(0ULL); // Reserve space for the memprof record table offset.
@@ -538,26 +532,21 @@ static Error writeMemProfV1(
   auto Schema = memprof::getFullSchema();
   writeMemProfSchema(OS, Schema);
 
-  uint64_t RecordTableOffset =
-      writeMemProfRecords(OS, MemProfRecordData, &Schema, memprof::Version1);
+  uint64_t RecordTableOffset = writeMemProfRecords(OS, MemProfData.RecordData,
+                                                   &Schema, memprof::Version1);
 
   uint64_t FramePayloadOffset = OS.tell();
-  uint64_t FrameTableOffset = writeMemProfFrames(OS, MemProfFrameData);
+  uint64_t FrameTableOffset = writeMemProfFrames(OS, MemProfData.FrameData);
 
   uint64_t Header[] = {RecordTableOffset, FramePayloadOffset, FrameTableOffset};
   OS.patch({{HeaderUpdatePos, Header, std::size(Header)}});
 
   return Error::success();
 }
 
-static Error writeMemProfV2(
-    ProfOStream &OS,
-    llvm::MapVector<GlobalValue::GUID, memprof::IndexedMemProfRecord>
-        &MemProfRecordData,
-    llvm::MapVector<memprof::FrameId, memprof::Frame> &MemProfFrameData,
-    llvm::MapVector<memprof::CallStackId, llvm::SmallVector<memprof::FrameId>>
-        &MemProfCallStackData,
-    bool MemProfFullSchema) {
+static Error writeMemProfV2(ProfOStream &OS,
+                            memprof::IndexedMemProfData &MemProfData,
+                            bool MemProfFullSchema) {
   OS.write(memprof::Version2);
   uint64_t HeaderUpdatePos = OS.tell();
   OS.write(0ULL); // Reserve space for the memprof record table offset.
@@ -571,15 +560,15 @@ static Error writeMemProfV2(
     Schema = memprof::getFullSchema();
   writeMemProfSchema(OS, Schema);
 
-  uint64_t RecordTableOffset =
-      writeMemProfRecords(OS, MemProfRecordData, &Schema, memprof::Version2);
+  uint64_t RecordTableOffset = writeMemProfRecords(OS, MemProfData.RecordData,
+                                                   &Schema, memprof::Version2);
 
   uint64_t FramePayloadOffset = OS.tell();
-  uint64_t FrameTableOffset = writeMemProfFrames(OS, MemProfFrameData);
+  uint64_t FrameTableOffset = writeMemProfFrames(OS, MemProfData.FrameData);
 
   uint64_t CallStackPayloadOffset = OS.tell();
   uint64_t CallStackTableOffset =
-      writeMemProfCallStacks(OS, MemProfCallStackData);
+      writeMemProfCallStacks(OS, MemProfData.CallStackData);
 
   uint64_t Header[] = {
       RecordTableOffset,      FramePayloadOffset,   FrameTableOffset,
@@ -603,23 +592,17 @@ static Error writeMemProfV2(
 // uint64_t Schema entry N - 1
 // OnDiskChainedHashTable MemProfRecordData
 // OnDiskChainedHashTable MemProfFrameData
-static Error writeMemProf(
-    ProfOStream &OS,
-    llvm::MapVector<GlobalValue::GUID, memprof::IndexedMemProfRecord>
-        &MemProfRecordData,
-    llvm::MapVector<memprof::FrameId, memprof::Frame> &MemProfFrameData,
-    llvm::MapVector<memprof::CallStackId, llvm::SmallVector<memprof::FrameId>>
-        &MemProfCallStackData,
-    memprof::IndexedVersion MemProfVersionRequested, bool MemProfFullSchema) {
-
+static Error writeMemProf(ProfOStream &OS,
+                          memprof::IndexedMemProfData &MemProfData,
+                          memprof::IndexedVersion MemProfVersionRequested,
+                          bool MemProfFullSchema) {
   switch (MemProfVersionRequested) {
   case memprof::Version0:
-    return writeMemProfV0(OS, MemProfRecordData, MemProfFrameData);
+    return writeMemProfV0(OS, MemProfData);
   case memprof::Version1:
-    return writeMemProfV1(OS, MemProfRecordData, MemProfFrameData);
+    return writeMemProfV1(OS, MemProfData);
   case memprof::Version2:
-    return writeMemProfV2(OS, MemProfRecordData, MemProfFrameData,
-                          MemProfCallStackData, MemProfFullSchema);
+    return writeMemProfV2(OS, MemProfData, MemProfFullSchema);
   }
 
   return make_error<InstrProfError>(
@@ -737,8 +720,7 @@ Error InstrProfWriter::writeImpl(ProfOStream &OS) {
   uint64_t MemProfSectionStart = 0;
   if (static_cast<bool>(ProfileKind & InstrProfKind::MemProf)) {
     MemProfSectionStart = OS.tell();
-    if (auto E = writeMemProf(OS, MemProfRecordData, MemProfFrameData,
-                              MemProfCallStackData, MemProfVersionRequested,
+    if (auto E = writeMemProf(OS, MemProfData, MemProfVersionRequested,
                               MemProfFullSchema))
       return E;
   }
