From 66d641276a4d5a1f63e6010a5facef0b82799cce Mon Sep 17 00:00:00 2001
From: Kazu Hirata <kazu@google.com>
Date: Wed, 4 Dec 2024 10:23:21 -0800
Subject: [PATCH 1/8] [memprof] Add YAML read/write support to llvm-profdata

This patch adds YAML read/write support to llvm-profdata.  The primary
intent is to accommodate MemProf profiles in test cases, thereby
avoiding the binary format.

The read support is via llvm-profdata merge.  This is useful when we
want to verify that the compiler does the right thing on a given .ll
file and a MemProf profile in a test case.  In the test case, we would
convert the MemProf profile in YAML to an indexed profile and invoke
the compiler on the .ll file along with the indexed profile.

The write support is via llvm-profdata show --memory.  This is useful
when we wish to convert an indexed MemProf profile to YAML while
writing tests.  We would compile a test case in C++, run it for an
indexed MemProf profile, and then convert it to the text format.
---
 .../llvm/ProfileData/InstrProfReader.h        |  7 ++
 llvm/include/llvm/ProfileData/MemProfReader.h | 15 ++++
 llvm/lib/ProfileData/InstrProfReader.cpp      |  8 ++
 llvm/lib/ProfileData/MemProfReader.cpp        | 30 +++++++
 .../tools/llvm-profdata/memprof-yaml.test     | 33 +++++++
 llvm/tools/llvm-profdata/llvm-profdata.cpp    | 89 ++++++++++++++++---
 6 files changed, 170 insertions(+), 12 deletions(-)
 create mode 100644 llvm/test/tools/llvm-profdata/memprof-yaml.test

diff --git a/llvm/include/llvm/ProfileData/InstrProfReader.h b/llvm/include/llvm/ProfileData/InstrProfReader.h
index 330cf540c099b..95a8171021b43 100644
--- a/llvm/include/llvm/ProfileData/InstrProfReader.h
+++ b/llvm/include/llvm/ProfileData/InstrProfReader.h
@@ -716,6 +716,9 @@ class IndexedMemProfReader {
 
   DenseMap<uint64_t, SmallVector<memprof::CallEdgeTy, 0>>
   getMemProfCallerCalleePairs() const;
+
+  // Return a vector of all GUIDs that we have corresponding MemProfRecords for.
+  SmallVector<uint64_t, 0> getMemProfRecordKeys() const;
 };
 
 /// Reader for the indexed binary instrprof format.
@@ -823,6 +826,10 @@ class IndexedInstrProfReader : public InstrProfReader {
     return MemProfReader.getMemProfCallerCalleePairs();
   }
 
+  SmallVector<uint64_t, 0> getMemProfRecordKeys() {
+    return MemProfReader.getMemProfRecordKeys();
+  }
+
   /// Fill Counts with the profile data for the given function name.
   Error getFunctionCounts(StringRef FuncName, uint64_t FuncHash,
                           std::vector<uint64_t> &Counts);
diff --git a/llvm/include/llvm/ProfileData/MemProfReader.h b/llvm/include/llvm/ProfileData/MemProfReader.h
index 0529f79460646..421985b55056c 100644
--- a/llvm/include/llvm/ProfileData/MemProfReader.h
+++ b/llvm/include/llvm/ProfileData/MemProfReader.h
@@ -213,6 +213,21 @@ class RawMemProfReader final : public MemProfReader {
 class YAMLMemProfReader final : public MemProfReader {
 public:
   YAMLMemProfReader() = default;
+
+  // Return true if the \p DataBuffer starts with magic bytes indicating it is
+  // a raw binary memprof profile.
+  static bool hasFormat(const MemoryBuffer &DataBuffer);
+  // Return true if the file at \p Path starts with magic bytes indicating it is
+  // a raw binary memprof profile.
+  static bool hasFormat(const StringRef Path);
+
+  // Create a RawMemProfReader after sanity checking the contents of the file at
+  // \p Path or the \p Buffer. The binary from which the profile has been
+  // collected is specified via a path in \p ProfiledBinary.
+  static Expected<std::unique_ptr<YAMLMemProfReader>> create(const Twine &Path);
+  static Expected<std::unique_ptr<YAMLMemProfReader>>
+  create(std::unique_ptr<MemoryBuffer> Buffer);
+
   void parse(StringRef YAMLData);
 };
 } // namespace memprof
diff --git a/llvm/lib/ProfileData/InstrProfReader.cpp b/llvm/lib/ProfileData/InstrProfReader.cpp
index dad79b2c1761e..cd10429cd742b 100644
--- a/llvm/lib/ProfileData/InstrProfReader.cpp
+++ b/llvm/lib/ProfileData/InstrProfReader.cpp
@@ -1664,6 +1664,14 @@ IndexedMemProfReader::getMemProfCallerCalleePairs() const {
   return Pairs;
 }
 
+SmallVector<uint64_t, 0> IndexedMemProfReader::getMemProfRecordKeys() const {
+  SmallVector<uint64_t, 0> Keys;
+  Keys.reserve(MemProfRecordTable->getNumEntries());
+  for (uint64_t Key : MemProfRecordTable->keys())
+    Keys.push_back(Key);
+  return Keys;
+}
+
 Error IndexedInstrProfReader::getFunctionCounts(StringRef FuncName,
                                                 uint64_t FuncHash,
                                                 std::vector<uint64_t> &Counts) {
diff --git a/llvm/lib/ProfileData/MemProfReader.cpp b/llvm/lib/ProfileData/MemProfReader.cpp
index 9dacf29898593..8325fc75aaa85 100644
--- a/llvm/lib/ProfileData/MemProfReader.cpp
+++ b/llvm/lib/ProfileData/MemProfReader.cpp
@@ -755,6 +755,36 @@ Error RawMemProfReader::readNextRecord(
   return MemProfReader::readNextRecord(GuidRecord, IdToFrameCallback);
 }
 
+Expected<std::unique_ptr<YAMLMemProfReader>>
+YAMLMemProfReader::create(const Twine &Path) {
+  auto BufferOr = MemoryBuffer::getFileOrSTDIN(Path);
+  if (std::error_code EC = BufferOr.getError())
+    return report(errorCodeToError(EC), Path.getSingleStringRef());
+
+  std::unique_ptr<MemoryBuffer> Buffer(BufferOr.get().release());
+  return create(std::move(Buffer));
+}
+
+Expected<std::unique_ptr<YAMLMemProfReader>>
+YAMLMemProfReader::create(std::unique_ptr<MemoryBuffer> Buffer) {
+  std::unique_ptr<YAMLMemProfReader> Reader(new YAMLMemProfReader());
+  Reader->parse(Buffer->getBuffer());
+  return std::move(Reader);
+}
+
+bool YAMLMemProfReader::hasFormat(const StringRef Path) {
+  auto BufferOr = MemoryBuffer::getFileOrSTDIN(Path);
+  if (!BufferOr)
+    return false;
+
+  std::unique_ptr<MemoryBuffer> Buffer(BufferOr.get().release());
+  return hasFormat(*Buffer);
+}
+
+bool YAMLMemProfReader::hasFormat(const MemoryBuffer &Buffer) {
+  return Buffer.getBuffer().starts_with("---");
+}
+
 void YAMLMemProfReader::parse(StringRef YAMLData) {
   memprof::AllMemProfData Doc;
   yaml::Input Yin(YAMLData);
diff --git a/llvm/test/tools/llvm-profdata/memprof-yaml.test b/llvm/test/tools/llvm-profdata/memprof-yaml.test
new file mode 100644
index 0000000000000..9875faf355582
--- /dev/null
+++ b/llvm/test/tools/llvm-profdata/memprof-yaml.test
@@ -0,0 +1,33 @@
+; RUN: split-file %s %t
+; RUN: llvm-profdata merge %t/memprof-in.yaml -o %t/memprof-out.indexed
+; RUN: llvm-profdata show --memory %t/memprof-out.indexed > %t/memprof-out.yaml
+; RUN: cmp %t/memprof-in.yaml %t/memprof-out.yaml
+
+; Verify that the YAML output is identical to the YAML input.
+;--- memprof-in.yaml
+---
+HeapProfileRecords:
+  - GUID:            16045690981402826360
+    AllocSites:
+      - Callstack:
+          - { Function: 100, LineOffset: 11, Column: 10, IsInlineFrame: true }
+          - { Function: 200, LineOffset: 22, Column: 20, IsInlineFrame: false }
+        MemInfoBlock:
+          AllocCount:      111
+          TotalSize:       222
+          TotalLifetime:   333
+          TotalLifetimeAccessDensity: 444
+      - Callstack:
+          - { Function: 300, LineOffset: 33, Column: 30, IsInlineFrame: false }
+          - { Function: 400, LineOffset: 44, Column: 40, IsInlineFrame: true }
+        MemInfoBlock:
+          AllocCount:      555
+          TotalSize:       666
+          TotalLifetime:   777
+          TotalLifetimeAccessDensity: 888
+    CallSites:
+      - - { Function: 500, LineOffset: 55, Column: 50, IsInlineFrame: true }
+        - { Function: 600, LineOffset: 66, Column: 60, IsInlineFrame: false }
+      - - { Function: 700, LineOffset: 77, Column: 70, IsInlineFrame: true }
+        - { Function: 800, LineOffset: 88, Column: 80, IsInlineFrame: false }
+...
diff --git a/llvm/tools/llvm-profdata/llvm-profdata.cpp b/llvm/tools/llvm-profdata/llvm-profdata.cpp
index 1d9d7bcf76549..69973db532853 100644
--- a/llvm/tools/llvm-profdata/llvm-profdata.cpp
+++ b/llvm/tools/llvm-profdata/llvm-profdata.cpp
@@ -723,6 +723,35 @@ loadInput(const WeightedFile &Input, SymbolRemapper *Remapper,
     return;
   }
 
+  using ::llvm::memprof::YAMLMemProfReader;
+  if (YAMLMemProfReader::hasFormat(Input.Filename)) {
+    auto ReaderOrErr = YAMLMemProfReader::create(Input.Filename);
+    if (!ReaderOrErr) {
+      exitWithError(ReaderOrErr.takeError(), Input.Filename);
+    }
+    std::unique_ptr<YAMLMemProfReader> Reader = std::move(ReaderOrErr.get());
+    // Check if the profile types can be merged, e.g. clang frontend profiles
+    // should not be merged with memprof profiles.
+    if (Error E = WC->Writer.mergeProfileKind(Reader->getProfileKind())) {
+      consumeError(std::move(E));
+      WC->Errors.emplace_back(
+          make_error<StringError>(
+              "Cannot merge MemProf profile with Clang generated profile.",
+              std::error_code()),
+          Filename);
+      return;
+    }
+
+    auto MemProfError = [&](Error E) {
+      auto [ErrorCode, Msg] = InstrProfError::take(std::move(E));
+      WC->Errors.emplace_back(make_error<InstrProfError>(ErrorCode, Msg),
+                              Filename);
+    };
+
+    WC->Writer.addMemProfData(Reader->takeMemProfData(), MemProfError);
+    return;
+  }
+
   auto FS = vfs::getRealFileSystem();
   // TODO: This only saves the first non-fatal error from InstrProfReader, and
   // then added to WriterContext::Errors. However, this is not extensible, if
@@ -3242,18 +3271,54 @@ static int showSampleProfile(ShowFormat SFormat, raw_fd_ostream &OS) {
 static int showMemProfProfile(ShowFormat SFormat, raw_fd_ostream &OS) {
   if (SFormat == ShowFormat::Json)
     exitWithError("JSON output is not supported for MemProf");
-  auto ReaderOr = llvm::memprof::RawMemProfReader::create(
-      Filename, ProfiledBinary, /*KeepNames=*/true);
-  if (Error E = ReaderOr.takeError())
-    // Since the error can be related to the profile or the binary we do not
-    // pass whence. Instead additional context is provided where necessary in
-    // the error message.
-    exitWithError(std::move(E), /*Whence*/ "");
-
-  std::unique_ptr<llvm::memprof::RawMemProfReader> Reader(
-      ReaderOr.get().release());
-
-  Reader->printYAML(OS);
+
+  // Load the file to check the magic bytes.
+  llvm::ErrorOr<std::unique_ptr<llvm::MemoryBuffer>> BufferOrError =
+      llvm::MemoryBuffer::getFile(Filename);
+  if (auto EC = BufferOrError.getError())
+    exitWithError("Error opening profile file '" + Filename + "'");
+  auto Buffer = std::move(BufferOrError.get());
+
+  // Show the raw profile in YAML.
+  if (memprof::RawMemProfReader::hasFormat(*Buffer)) {
+    auto ReaderOr = llvm::memprof::RawMemProfReader::create(
+        Filename, ProfiledBinary, /*KeepNames=*/true);
+    if (Error E = ReaderOr.takeError())
+      // Since the error can be related to the profile or the binary we do not
+      // pass whence. Instead additional context is provided where necessary in
+      // the error message.
+      exitWithError(std::move(E), /*Whence*/ "");
+
+    std::unique_ptr<llvm::memprof::RawMemProfReader> Reader(
+        ReaderOr.get().release());
+
+    Reader->printYAML(OS);
+    return 0;
+  }
+
+  // Show the indexed MemProf profile in YAML.
+  auto FS = vfs::getRealFileSystem();
+  auto ReaderOrErr = IndexedInstrProfReader::create(Filename, *FS);
+  if (Error E = ReaderOrErr.takeError())
+    exitWithError(std::move(E), Filename);
+
+  auto Reader = std::move(ReaderOrErr.get());
+
+  // Build pairs of GUID and MemProfRecord.
+  memprof::AllMemProfData Data;
+  for (const uint64_t Key : Reader->getMemProfRecordKeys()) {
+    auto Record = Reader->getMemProfRecord(Key);
+    if (Record.takeError())
+      continue;
+    memprof::GUIDMemProfRecordPair Pair;
+    Pair.GUID = Key;
+    Pair.Record = std::move(*Record);
+    Data.HeapProfileRecords.push_back(std::move(Pair));
+  }
+
+  yaml::Output Yout(OS);
+  Yout << Data;
+
   return 0;
 }
 

From d8a8da668fce68efb6371d49a15ca5519c580ee3 Mon Sep 17 00:00:00 2001
From: Kazu Hirata <kazu@google.com>
Date: Thu, 5 Dec 2024 21:24:52 -0800
Subject: [PATCH 2/8] Add "REQUIRES: x86_64-linux".

---
 llvm/test/tools/llvm-profdata/memprof-yaml.test | 1 +
 1 file changed, 1 insertion(+)

diff --git a/llvm/test/tools/llvm-profdata/memprof-yaml.test b/llvm/test/tools/llvm-profdata/memprof-yaml.test
index 9875faf355582..fbbf8a6e7e5b5 100644
--- a/llvm/test/tools/llvm-profdata/memprof-yaml.test
+++ b/llvm/test/tools/llvm-profdata/memprof-yaml.test
@@ -1,3 +1,4 @@
+; REQUIRES: x86_64-linux
 ; RUN: split-file %s %t
 ; RUN: llvm-profdata merge %t/memprof-in.yaml -o %t/memprof-out.indexed
 ; RUN: llvm-profdata show --memory %t/memprof-out.indexed > %t/memprof-out.yaml

From 13e6da8e81421d0b03bf996b6e1059898663f961 Mon Sep 17 00:00:00 2001
From: Kazu Hirata <kazu@google.com>
Date: Fri, 6 Dec 2024 11:26:42 -0800
Subject: [PATCH 3/8] Address comments.

---
 .../llvm/ProfileData/InstrProfReader.h        |  8 ++++----
 llvm/include/llvm/ProfileData/MemProfReader.h |  8 +++-----
 llvm/lib/ProfileData/InstrProfReader.cpp      | 19 +++++++++++++------
 llvm/tools/llvm-profdata/llvm-profdata.cpp    | 19 +++----------------
 4 files changed, 23 insertions(+), 31 deletions(-)

diff --git a/llvm/include/llvm/ProfileData/InstrProfReader.h b/llvm/include/llvm/ProfileData/InstrProfReader.h
index 95a8171021b43..25e13f575296d 100644
--- a/llvm/include/llvm/ProfileData/InstrProfReader.h
+++ b/llvm/include/llvm/ProfileData/InstrProfReader.h
@@ -717,8 +717,8 @@ class IndexedMemProfReader {
   DenseMap<uint64_t, SmallVector<memprof::CallEdgeTy, 0>>
   getMemProfCallerCalleePairs() const;
 
-  // Return a vector of all GUIDs that we have corresponding MemProfRecords for.
-  SmallVector<uint64_t, 0> getMemProfRecordKeys() const;
+  // Return the entire MemProf profile.
+  memprof::AllMemProfData getAllMemProfData() const;
 };
 
 /// Reader for the indexed binary instrprof format.
@@ -826,8 +826,8 @@ class IndexedInstrProfReader : public InstrProfReader {
     return MemProfReader.getMemProfCallerCalleePairs();
   }
 
-  SmallVector<uint64_t, 0> getMemProfRecordKeys() {
-    return MemProfReader.getMemProfRecordKeys();
+  memprof::AllMemProfData getAllMemProfData() const {
+    return MemProfReader.getAllMemProfData();
   }
 
   /// Fill Counts with the profile data for the given function name.
diff --git a/llvm/include/llvm/ProfileData/MemProfReader.h b/llvm/include/llvm/ProfileData/MemProfReader.h
index 421985b55056c..70e48bbea004a 100644
--- a/llvm/include/llvm/ProfileData/MemProfReader.h
+++ b/llvm/include/llvm/ProfileData/MemProfReader.h
@@ -214,16 +214,14 @@ class YAMLMemProfReader final : public MemProfReader {
 public:
   YAMLMemProfReader() = default;
 
-  // Return true if the \p DataBuffer starts with magic bytes indicating it is
-  // a raw binary memprof profile.
+  // Return true if the \p DataBuffer starts "---" indicating it is a YAML file.
   static bool hasFormat(const MemoryBuffer &DataBuffer);
   // Return true if the file at \p Path starts with magic bytes indicating it is
   // a raw binary memprof profile.
   static bool hasFormat(const StringRef Path);
 
-  // Create a RawMemProfReader after sanity checking the contents of the file at
-  // \p Path or the \p Buffer. The binary from which the profile has been
-  // collected is specified via a path in \p ProfiledBinary.
+  // Create a YAMLMemProfReader after sanity checking the contents of the file at
+  // \p Path or the \p Buffer.
   static Expected<std::unique_ptr<YAMLMemProfReader>> create(const Twine &Path);
   static Expected<std::unique_ptr<YAMLMemProfReader>>
   create(std::unique_ptr<MemoryBuffer> Buffer);
diff --git a/llvm/lib/ProfileData/InstrProfReader.cpp b/llvm/lib/ProfileData/InstrProfReader.cpp
index cd10429cd742b..d2ff10bd1d847 100644
--- a/llvm/lib/ProfileData/InstrProfReader.cpp
+++ b/llvm/lib/ProfileData/InstrProfReader.cpp
@@ -1664,12 +1664,19 @@ IndexedMemProfReader::getMemProfCallerCalleePairs() const {
   return Pairs;
 }
 
-SmallVector<uint64_t, 0> IndexedMemProfReader::getMemProfRecordKeys() const {
-  SmallVector<uint64_t, 0> Keys;
-  Keys.reserve(MemProfRecordTable->getNumEntries());
-  for (uint64_t Key : MemProfRecordTable->keys())
-    Keys.push_back(Key);
-  return Keys;
+memprof::AllMemProfData IndexedMemProfReader::getAllMemProfData() const {
+  memprof::AllMemProfData AllMemProfData;
+  AllMemProfData.HeapProfileRecords.reserve(MemProfRecordTable->getNumEntries());
+  for (uint64_t Key : MemProfRecordTable->keys()) {
+    auto Record = getMemProfRecord(Key);
+    if (Record.takeError())
+      continue;
+    memprof::GUIDMemProfRecordPair Pair;
+    Pair.GUID = Key;
+    Pair.Record = std::move(*Record);
+    AllMemProfData.HeapProfileRecords.push_back(std::move(Pair));
+  }
+  return AllMemProfData;
 }
 
 Error IndexedInstrProfReader::getFunctionCounts(StringRef FuncName,
diff --git a/llvm/tools/llvm-profdata/llvm-profdata.cpp b/llvm/tools/llvm-profdata/llvm-profdata.cpp
index 69973db532853..35f8fbe8688ef 100644
--- a/llvm/tools/llvm-profdata/llvm-profdata.cpp
+++ b/llvm/tools/llvm-profdata/llvm-profdata.cpp
@@ -726,9 +726,8 @@ loadInput(const WeightedFile &Input, SymbolRemapper *Remapper,
   using ::llvm::memprof::YAMLMemProfReader;
   if (YAMLMemProfReader::hasFormat(Input.Filename)) {
     auto ReaderOrErr = YAMLMemProfReader::create(Input.Filename);
-    if (!ReaderOrErr) {
+    if (!ReaderOrErr)
       exitWithError(ReaderOrErr.takeError(), Input.Filename);
-    }
     std::unique_ptr<YAMLMemProfReader> Reader = std::move(ReaderOrErr.get());
     // Check if the profile types can be merged, e.g. clang frontend profiles
     // should not be merged with memprof profiles.
@@ -736,7 +735,7 @@ loadInput(const WeightedFile &Input, SymbolRemapper *Remapper,
       consumeError(std::move(E));
       WC->Errors.emplace_back(
           make_error<StringError>(
-              "Cannot merge MemProf profile with Clang generated profile.",
+              "Cannot merge MemProf profile with incompatible profile.",
               std::error_code()),
           Filename);
       return;
@@ -3303,19 +3302,7 @@ static int showMemProfProfile(ShowFormat SFormat, raw_fd_ostream &OS) {
     exitWithError(std::move(E), Filename);
 
   auto Reader = std::move(ReaderOrErr.get());
-
-  // Build pairs of GUID and MemProfRecord.
-  memprof::AllMemProfData Data;
-  for (const uint64_t Key : Reader->getMemProfRecordKeys()) {
-    auto Record = Reader->getMemProfRecord(Key);
-    if (Record.takeError())
-      continue;
-    memprof::GUIDMemProfRecordPair Pair;
-    Pair.GUID = Key;
-    Pair.Record = std::move(*Record);
-    Data.HeapProfileRecords.push_back(std::move(Pair));
-  }
-
+  memprof::AllMemProfData Data = Reader->getAllMemProfData();
   yaml::Output Yout(OS);
   Yout << Data;
 

From ff591bfc35f1b543bcb112fe2d42cdfe2bbc48a5 Mon Sep 17 00:00:00 2001
From: Kazu Hirata <kazu@google.com>
Date: Fri, 6 Dec 2024 11:57:34 -0800
Subject: [PATCH 4/8] Fix formatting

---
 llvm/include/llvm/ProfileData/MemProfReader.h | 4 ++--
 llvm/lib/ProfileData/InstrProfReader.cpp      | 3 ++-
 2 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/llvm/include/llvm/ProfileData/MemProfReader.h b/llvm/include/llvm/ProfileData/MemProfReader.h
index 70e48bbea004a..43a3fe4e3b2d7 100644
--- a/llvm/include/llvm/ProfileData/MemProfReader.h
+++ b/llvm/include/llvm/ProfileData/MemProfReader.h
@@ -220,8 +220,8 @@ class YAMLMemProfReader final : public MemProfReader {
   // a raw binary memprof profile.
   static bool hasFormat(const StringRef Path);
 
-  // Create a YAMLMemProfReader after sanity checking the contents of the file at
-  // \p Path or the \p Buffer.
+  // Create a YAMLMemProfReader after sanity checking the contents of the file
+  // at \p Path or the \p Buffer.
   static Expected<std::unique_ptr<YAMLMemProfReader>> create(const Twine &Path);
   static Expected<std::unique_ptr<YAMLMemProfReader>>
   create(std::unique_ptr<MemoryBuffer> Buffer);
diff --git a/llvm/lib/ProfileData/InstrProfReader.cpp b/llvm/lib/ProfileData/InstrProfReader.cpp
index d2ff10bd1d847..cac1760d3ef80 100644
--- a/llvm/lib/ProfileData/InstrProfReader.cpp
+++ b/llvm/lib/ProfileData/InstrProfReader.cpp
@@ -1666,7 +1666,8 @@ IndexedMemProfReader::getMemProfCallerCalleePairs() const {
 
 memprof::AllMemProfData IndexedMemProfReader::getAllMemProfData() const {
   memprof::AllMemProfData AllMemProfData;
-  AllMemProfData.HeapProfileRecords.reserve(MemProfRecordTable->getNumEntries());
+  AllMemProfData.HeapProfileRecords.reserve(
+      MemProfRecordTable->getNumEntries());
   for (uint64_t Key : MemProfRecordTable->keys()) {
     auto Record = getMemProfRecord(Key);
     if (Record.takeError())

From 40ac8ffdb3559e8e0e9180617606dc792d798cb3 Mon Sep 17 00:00:00 2001
From: Kazu Hirata <kazu@google.com>
Date: Fri, 6 Dec 2024 13:32:04 -0800
Subject: [PATCH 5/8] Address comments.

---
 llvm/include/llvm/ProfileData/MemProfReader.h |  6 +++---
 llvm/lib/ProfileData/MemProfReader.cpp        |  2 +-
 llvm/tools/llvm-profdata/llvm-profdata.cpp    | 12 +++---------
 3 files changed, 7 insertions(+), 13 deletions(-)

diff --git a/llvm/include/llvm/ProfileData/MemProfReader.h b/llvm/include/llvm/ProfileData/MemProfReader.h
index 43a3fe4e3b2d7..c8eec19b83a83 100644
--- a/llvm/include/llvm/ProfileData/MemProfReader.h
+++ b/llvm/include/llvm/ProfileData/MemProfReader.h
@@ -214,10 +214,10 @@ class YAMLMemProfReader final : public MemProfReader {
 public:
   YAMLMemProfReader() = default;
 
-  // Return true if the \p DataBuffer starts "---" indicating it is a YAML file.
+  // Return true if the \p DataBuffer starts with "---" indicating it is a YAML
+  // file.
   static bool hasFormat(const MemoryBuffer &DataBuffer);
-  // Return true if the file at \p Path starts with magic bytes indicating it is
-  // a raw binary memprof profile.
+  // Wrapper around hasFormat above, reading the file instead of the memory buffer.
   static bool hasFormat(const StringRef Path);
 
   // Create a YAMLMemProfReader after sanity checking the contents of the file
diff --git a/llvm/lib/ProfileData/MemProfReader.cpp b/llvm/lib/ProfileData/MemProfReader.cpp
index 8325fc75aaa85..e4082fff41686 100644
--- a/llvm/lib/ProfileData/MemProfReader.cpp
+++ b/llvm/lib/ProfileData/MemProfReader.cpp
@@ -767,7 +767,7 @@ YAMLMemProfReader::create(const Twine &Path) {
 
 Expected<std::unique_ptr<YAMLMemProfReader>>
 YAMLMemProfReader::create(std::unique_ptr<MemoryBuffer> Buffer) {
-  std::unique_ptr<YAMLMemProfReader> Reader(new YAMLMemProfReader());
+  auto Reader = std::make_unique<YAMLMemProfReader>();
   Reader->parse(Buffer->getBuffer());
   return std::move(Reader);
 }
diff --git a/llvm/tools/llvm-profdata/llvm-profdata.cpp b/llvm/tools/llvm-profdata/llvm-profdata.cpp
index 35f8fbe8688ef..2c455f0ae1435 100644
--- a/llvm/tools/llvm-profdata/llvm-profdata.cpp
+++ b/llvm/tools/llvm-profdata/llvm-profdata.cpp
@@ -3271,22 +3271,16 @@ static int showMemProfProfile(ShowFormat SFormat, raw_fd_ostream &OS) {
   if (SFormat == ShowFormat::Json)
     exitWithError("JSON output is not supported for MemProf");
 
-  // Load the file to check the magic bytes.
-  llvm::ErrorOr<std::unique_ptr<llvm::MemoryBuffer>> BufferOrError =
-      llvm::MemoryBuffer::getFile(Filename);
-  if (auto EC = BufferOrError.getError())
-    exitWithError("Error opening profile file '" + Filename + "'");
-  auto Buffer = std::move(BufferOrError.get());
-
   // Show the raw profile in YAML.
-  if (memprof::RawMemProfReader::hasFormat(*Buffer)) {
+  if (memprof::RawMemProfReader::hasFormat(Filename)) {
     auto ReaderOr = llvm::memprof::RawMemProfReader::create(
         Filename, ProfiledBinary, /*KeepNames=*/true);
-    if (Error E = ReaderOr.takeError())
+    if (Error E = ReaderOr.takeError()) {
       // Since the error can be related to the profile or the binary we do not
       // pass whence. Instead additional context is provided where necessary in
       // the error message.
       exitWithError(std::move(E), /*Whence*/ "");
+    }
 
     std::unique_ptr<llvm::memprof::RawMemProfReader> Reader(
         ReaderOr.get().release());

From a395c9189133ba9615ac4e0aa6c95e6d04129d78 Mon Sep 17 00:00:00 2001
From: Kazu Hirata <kazu@google.com>
Date: Fri, 6 Dec 2024 15:02:04 -0800
Subject: [PATCH 6/8] Fix formatting

---
 llvm/include/llvm/ProfileData/MemProfReader.h | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/llvm/include/llvm/ProfileData/MemProfReader.h b/llvm/include/llvm/ProfileData/MemProfReader.h
index c8eec19b83a83..29d9e57cae3e3 100644
--- a/llvm/include/llvm/ProfileData/MemProfReader.h
+++ b/llvm/include/llvm/ProfileData/MemProfReader.h
@@ -217,7 +217,8 @@ class YAMLMemProfReader final : public MemProfReader {
   // Return true if the \p DataBuffer starts with "---" indicating it is a YAML
   // file.
   static bool hasFormat(const MemoryBuffer &DataBuffer);
-  // Wrapper around hasFormat above, reading the file instead of the memory buffer.
+  // Wrapper around hasFormat above, reading the file instead of the memory
+  // buffer.
   static bool hasFormat(const StringRef Path);
 
   // Create a YAMLMemProfReader after sanity checking the contents of the file

From a3865b8ce9dbc82f9936b0a970176849047467a1 Mon Sep 17 00:00:00 2001
From: Kazu Hirata <kazu@google.com>
Date: Fri, 6 Dec 2024 15:25:27 -0800
Subject: [PATCH 7/8] Check for the empty input.

---
 llvm/tools/llvm-profdata/llvm-profdata.cpp | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

diff --git a/llvm/tools/llvm-profdata/llvm-profdata.cpp b/llvm/tools/llvm-profdata/llvm-profdata.cpp
index 2c455f0ae1435..f34d99e10f316 100644
--- a/llvm/tools/llvm-profdata/llvm-profdata.cpp
+++ b/llvm/tools/llvm-profdata/llvm-profdata.cpp
@@ -747,7 +747,16 @@ loadInput(const WeightedFile &Input, SymbolRemapper *Remapper,
                               Filename);
     };
 
-    WC->Writer.addMemProfData(Reader->takeMemProfData(), MemProfError);
+    auto MemProfData = Reader->takeMemProfData();
+
+    // Check for the empty input in case the YAML file is invalid.
+    if (MemProfData.Records.empty()) {
+      WC->Errors.emplace_back(
+          make_error<StringError>("The profile is empty.", std::error_code()),
+          Filename);
+    }
+
+    WC->Writer.addMemProfData(std::move(MemProfData), MemProfError);
     return;
   }
 

From a57f294c587e28473ce3bf25d4e4cdba4fcb9f16 Mon Sep 17 00:00:00 2001
From: Kazu Hirata <kazu@google.com>
Date: Fri, 6 Dec 2024 16:54:41 -0800
Subject: [PATCH 8/8] Add a test for the invalid YAML input.

---
 llvm/test/tools/llvm-profdata/memprof-yaml-invalid.test | 8 ++++++++
 1 file changed, 8 insertions(+)
 create mode 100644 llvm/test/tools/llvm-profdata/memprof-yaml-invalid.test

diff --git a/llvm/test/tools/llvm-profdata/memprof-yaml-invalid.test b/llvm/test/tools/llvm-profdata/memprof-yaml-invalid.test
new file mode 100644
index 0000000000000..a13451cee5116
--- /dev/null
+++ b/llvm/test/tools/llvm-profdata/memprof-yaml-invalid.test
@@ -0,0 +1,8 @@
+; REQUIRES: x86_64-linux
+; RUN: split-file %s %t
+; RUN: not llvm-profdata merge %t/memprof-invalid.yaml -o %t/memprof-invalid.indexed
+
+; Verify that the invalid YAML input results in an error.
+;--- memprof-invalid.yaml
+---
+...