[llvm] f97c610 - [memprof] Add MemProfReader::takeMemProfData (#116769)

[via llvm-commits]

Author: Kazu Hirata
Date: 2024-11-19T19:33:26-08:00
New Revision: f97c610d1f824bcd3e078560c836aaaffaaf69b0

URL: https://github.com/llvm/llvm-project/commit/f97c610d1f824bcd3e078560c836aaaffaaf69b0
DIFF: https://github.com/llvm/llvm-project/commit/f97c610d1f824bcd3e078560c836aaaffaaf69b0.diff

LOG: [memprof] Add MemProfReader::takeMemProfData (#116769)

This patch adds MemProfReader::takeMemProfData, a function to return
the complete MemProf profile from the reader.  We can directly pass
its return value to InstrProfWriter::addMemProfData without having to
deal with the indivual components of the MemProf profile.  The new
function is named "take", but it doesn't do std::move yet because of
type differences (DenseMap v.s. MapVector).

The end state I'm trying to get to is roughly as follows:

- MemProfReader accepts IndexedMemProfData as a parameter as opposed
  to the three individual components (frames, call stacks, and
  records).

- MemProfReader keeps IndexedMemProfData as a class member without
  decomposing it into its individual components.

- MemProfReader returns IndexedMemProfData like:

  IndexedMemProfData takeMemProfData() {
    return std::move(MemProfData);
  }

Added: 
    

Modified: 
    llvm/include/llvm/ProfileData/MemProfReader.h
    llvm/lib/ProfileData/InstrProfWriter.cpp
    llvm/tools/llvm-profdata/llvm-profdata.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/include/llvm/ProfileData/MemProfReader.h b/llvm/include/llvm/ProfileData/MemProfReader.h
index da2f14b276ffb0..de2167f97b0dc8 100644
--- a/llvm/include/llvm/ProfileData/MemProfReader.h
+++ b/llvm/include/llvm/ProfileData/MemProfReader.h
@@ -63,6 +63,24 @@ class MemProfReader {
     return FunctionProfileData;
   }
 
+  // Take the complete profile data.
+  IndexedMemProfData takeMemProfData() {
+    // TODO: Once we replace the three member variables, namely IdToFrame,
+    // CSIdToCallStack, and FunctionProfileData, with MemProfData, replace the
+    // following code with just "return std::move(MemProfData);".
+    IndexedMemProfData MemProfData;
+    // Copy key-value pairs because IdToFrame uses DenseMap, whereas
+    // IndexedMemProfData::Frames uses MapVector.
+    for (const auto &[FrameId, F] : IdToFrame)
+      MemProfData.Frames.try_emplace(FrameId, F);
+    // Copy key-value pairs because CSIdToCallStack uses DenseMap, whereas
+    // IndexedMemProfData::CallStacks uses MapVector.
+    for (const auto &[CSId, CS] : CSIdToCallStack)
+      MemProfData.CallStacks.try_emplace(CSId, CS);
+    MemProfData.Records = FunctionProfileData;
+    return MemProfData;
+  }
+
   virtual Error
   readNextRecord(GuidMemProfRecordPair &GuidRecord,
                  std::function<const Frame(const FrameId)> Callback = nullptr) {

diff  --git a/llvm/lib/ProfileData/InstrProfWriter.cpp b/llvm/lib/ProfileData/InstrProfWriter.cpp
index 87a538f35c7865..a1341d0fc973a1 100644
--- a/llvm/lib/ProfileData/InstrProfWriter.cpp
+++ b/llvm/lib/ProfileData/InstrProfWriter.cpp
@@ -370,7 +370,8 @@ bool InstrProfWriter::addMemProfData(memprof::IndexedMemProfData Incoming,
       if (addMemProfCallStack(CSId, CS, Warn))
         return false;
 
-  if (MemProfData.Records.empty())
+  // Add one record at a time if randomization is requested.
+  if (MemProfData.Records.empty() && !MemprofGenerateRandomHotness)
     MemProfData.Records = std::move(Incoming.Records);
   else
     for (const auto &[GUID, Record] : Incoming.Records)

diff  --git a/llvm/tools/llvm-profdata/llvm-profdata.cpp b/llvm/tools/llvm-profdata/llvm-profdata.cpp
index dcc8609f8df942..7641a80129de35 100644
--- a/llvm/tools/llvm-profdata/llvm-profdata.cpp
+++ b/llvm/tools/llvm-profdata/llvm-profdata.cpp
@@ -720,33 +720,7 @@ loadInput(const WeightedFile &Input, SymbolRemapper *Remapper,
                               Filename);
     };
 
-    // Add the frame mappings into the writer context.
-    const auto &IdToFrame = Reader->getFrameMapping();
-    for (const auto &I : IdToFrame) {
-      bool Succeeded = WC->Writer.addMemProfFrame(
-          /*Id=*/I.first, /*Frame=*/I.getSecond(), MemProfError);
-      // If we weren't able to add the frame mappings then it doesn't make sense
-      // to try to add the records from this profile.
-      if (!Succeeded)
-        return;
-    }
-
-    // Add the call stacks into the writer context.
-    const auto &CSIdToCallStacks = Reader->getCallStacks();
-    for (const auto &I : CSIdToCallStacks) {
-      bool Succeeded = WC->Writer.addMemProfCallStack(
-          /*Id=*/I.first, /*Frame=*/I.getSecond(), MemProfError);
-      // If we weren't able to add the call stacks then it doesn't make sense
-      // to try to add the records from this profile.
-      if (!Succeeded)
-        return;
-    }
-
-    const auto &FunctionProfileData = Reader->getProfileData();
-    // Add the memprof records into the writer context.
-    for (const auto &[GUID, Record] : FunctionProfileData) {
-      WC->Writer.addMemProfRecord(GUID, Record);
-    }
+    WC->Writer.addMemProfData(Reader->takeMemProfData(), MemProfError);
     return;
   }