[memprof] Add YAML read/write support to llvm-profdata (llvm#118915)

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.

Author: kazutakahirata

llvm/include/llvm/ProfileData/InstrProfReader.h

@@ -716,6 +716,9 @@ class IndexedMemProfReader {
 
   DenseMap<uint64_t, SmallVector<memprof::CallEdgeTy, 0>>
   getMemProfCallerCalleePairs() const;
+
+  // Return the entire MemProf profile.
+  memprof::AllMemProfData getAllMemProfData() const;
 };
 
 /// Reader for the indexed binary instrprof format.
@@ -823,6 +826,10 @@ class IndexedInstrProfReader : public InstrProfReader {
     return MemProfReader.getMemProfCallerCalleePairs();
   }
 
+  memprof::AllMemProfData getAllMemProfData() const {
+    return MemProfReader.getAllMemProfData();
+  }
+
   /// Fill Counts with the profile data for the given function name.
   Error getFunctionCounts(StringRef FuncName, uint64_t FuncHash,
                           std::vector<uint64_t> &Counts);

llvm/include/llvm/ProfileData/MemProfReader.h

@@ -213,6 +213,20 @@ class RawMemProfReader final : public MemProfReader {
 class YAMLMemProfReader final : public MemProfReader {
 public:
   YAMLMemProfReader() = default;
+
+  // 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.
+  static bool hasFormat(const StringRef Path);
+
+  // 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);
+
   void parse(StringRef YAMLData);
 };
 } // namespace memprof

llvm/lib/ProfileData/InstrProfReader.cpp

@@ -1664,6 +1664,22 @@ IndexedMemProfReader::getMemProfCallerCalleePairs() const {
   return Pairs;
 }
 
+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,
                                                 uint64_t FuncHash,
                                                 std::vector<uint64_t> &Counts) {

llvm/lib/ProfileData/MemProfReader.cpp

@@ -751,6 +751,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) {
+  auto Reader = std::make_unique<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);

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
+---
+...

llvm/test/tools/llvm-profdata/memprof-yaml.test

@@ -0,0 +1,34 @@
+; 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
+; 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 }
+...

llvm/tools/llvm-profdata/llvm-profdata.cpp

@@ -723,6 +723,43 @@ 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 incompatible 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);
+    };
+
+    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;
+  }
+
   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 +3279,36 @@ 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);
+
+  // Show the raw profile in YAML.
+  if (memprof::RawMemProfReader::hasFormat(Filename)) {
+    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());
+  memprof::AllMemProfData Data = Reader->getAllMemProfData();
+  yaml::Output Yout(OS);
+  Yout << Data;
+
   return 0;
 }