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

[Kazu Hirata via llvm-commits]

https://github.com/kazutakahirata updated https://github.com/llvm/llvm-project/pull/116769

>From d8e152c90af80425b9b846ed868ecee64e9b343c Mon Sep 17 00:00:00 2001
From: Kazu Hirata <kazu at google.com>
Date: Mon, 18 Nov 2024 11:58:05 -0800
Subject: [PATCH] [memprof] Add MemProfReader::takeMemProfData

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);
  }
---
 llvm/include/llvm/ProfileData/MemProfReader.h | 18 ++++++++++++
 llvm/lib/ProfileData/InstrProfWriter.cpp      |  3 +-
 llvm/tools/llvm-profdata/llvm-profdata.cpp    | 28 +------------------
 3 files changed, 21 insertions(+), 28 deletions(-)

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 8a42e430fb54e8..d7ba73233c021a 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;
   }