[llvm-profdata] Support JSON as as an output-only format

[kazu]

Summary

This patch teaches llvm-profdata to output the sample profile in the
JSON format. The new option is intended to be used for research and
development purposes. For example, one can write a Python script to
take a JSON file and analyze how similar different inline instances of
a given function are to each other.

I've chosen JSON because Python can parse it reasonably fast, and it
just takes a couple of lines to read the whole data:

import json
with open ('profile.json') as f:
  profile = json.load(f)

[hoy]

Perhaps make it explicit in the title that this is json an output-only format so far?

llvm/lib/ProfileData/SampleProfWriter.cpp

539

Wondering if this is necessary since S.getBodySamples() returns a BodySampleMap object which is already ordered.

Similarly with getCallsiteSamples below.

llvm/test/tools/llvm-profdata/sample-profile-json.test

3

nit: use JSON-next for this line and below

[wenlei]

Do you absolutely need this to be a supported format and generated natively in llvm tools?

For a profile format, we'd expect full reader/writer support, in both llvm-profdata and llvm-profgen. Having a partially supported output only format seems less than ideal.

For research purpose, would it be possible to generate json from a text profile through offline scripts?

Another option is to leverage llvm-profdata show to have an extra flag to output in json form. show output doesn't need to be a valid profile format and it doesn't need to be loadable by profile reader.

[davidxl]

Re-reply in Phab: +1 on the show command direction. If we have a reader in the future, consider promote this into a supported format.

[kazu]

Updated to move the JSON format support to "llvm-profdata show".

[kazu]

In D130944#3695666, @wenlei wrote:

Do you absolutely need this to be a supported format and generated natively in llvm tools?

For a profile format, we'd expect full reader/writer support, in both llvm-profdata and llvm-profgen. Having a partially supported output only format seems less than ideal.

For research purpose, would it be possible to generate json from a text profile through offline scripts?

Another option is to leverage llvm-profdata show to have an extra flag to output in json form. show output doesn't need to be a valid profile format and it doesn't need to be loadable by profile reader.

The main reason for adding the JSON support in an llvm tool is that I can leverage the parser.

llvm-profdata show makes sense.

I wrote another "dumper" just for the JSON format in SampleProfileReader. It's harder than I thought to wire SampleProfileWriter* to llvm-profdata show. Specifically, I had trouble casting raw_fd_ostream to raw_ostream for json::OStream purposes.

[kazu]

I've manually updated the patch description as that's not part of the diff in Phabricator.

[davidxl]

llvm/lib/ProfileData/SampleProfReader.cpp

87 ↗

(On Diff #451251)

too many nesting levels. Perhaps split out the lamda?

115 ↗

(On Diff #451251)

split out the lamda?

[kazu]

Broke up the big lambda into smaller pieces.

[kazu]

I've broken up the big lambda into smaller pieces. It should be easier to see that the output has two major parts, namely`BodySamples` and CallsiteSamples, with the implementation details factored out elsewhere.

PTAL. Thanks!

[davidxl]

lgtm

llvm/include/llvm/ProfileData/SampleProf.h

 namespace sampleprof {
 
 enum SampleProfileFormat {
   SPF_None = 0,
   SPF_Text = 0x1,
   SPF_Compact_Binary = 0x2,
   SPF_GCC = 0x3,
   SPF_Ext_Binary = 0x4,
+  SPF_JSON = 0x5,  // Output only
   SPF_Binary = 0xff
 };
 
 static inline uint64_t SPMagic(SampleProfileFormat Format = SPF_Binary) {
   return uint64_t('S') << (64 - 8) | uint64_t('P') << (64 - 16) |
          uint64_t('R') << (64 - 24) | uint64_t('O') << (64 - 32) |
          uint64_t('F') << (64 - 40) | uint64_t('4') << (64 - 48) |
          uint64_t('2') << (64 - 56) | uint64_t(Format);

llvm/include/llvm/ProfileData/SampleProfWriter.h

 #ifndef LLVM_PROFILEDATA_SAMPLEPROFWRITER_H
 #define LLVM_PROFILEDATA_SAMPLEPROFWRITER_H
 
 #include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/IR/ProfileSummary.h"
 #include "llvm/ProfileData/SampleProf.h"
 #include "llvm/Support/ErrorOr.h"
+#include "llvm/Support/JSON.h"
 #include "llvm/Support/raw_ostream.h"
 #include <cstdint>
 #include <memory>
 #include <set>
 #include <system_error>
 
 namespace llvm {
 namespace sampleprof {

 
 protected:
   SampleProfileWriter(std::unique_ptr<raw_ostream> &OS)
       : OutputStream(std::move(OS)) {}
 
   /// Write a file header for the profile file.
   virtual std::error_code writeHeader(const SampleProfileMap &ProfileMap) = 0;
 
+  /// Write a file footer for the profile file.
+  virtual std::error_code writeFooter(const SampleProfileMap &ProfileMap) {
+    return sampleprof_error::success;
+  }
+
   // Write function profiles to the profile file.
   virtual std::error_code writeFuncProfiles(const SampleProfileMap &ProfileMap);
 
   /// Output stream where to emit the profile to.
   std::unique_ptr<raw_ostream> OutputStream;
 
   /// Profile summary.
   std::unique_ptr<ProfileSummary> Summary;
private:
   /// This is used when printing inlined callees.
   unsigned Indent;
 
   friend ErrorOr<std::unique_ptr<SampleProfileWriter>>
   SampleProfileWriter::create(std::unique_ptr<raw_ostream> &OS,
                               SampleProfileFormat Format);
 };
 
+/// Sample-based profile writer (JSON format).
+class SampleProfileWriterJSON : public SampleProfileWriter {
+public:
+  std::error_code writeSample(const FunctionSamples &S) override;
+
+protected:
+  SampleProfileWriterJSON(std::unique_ptr<raw_ostream> &OS)
+      : SampleProfileWriter(OS),
+        JOS(std::make_unique<json::OStream>(*OutputStream, 2)) {}
+
+  std::error_code writeHeader(const SampleProfileMap &ProfileMap) override {
+    JOS->arrayBegin();
+    return sampleprof_error::success;
+  }
+
+  std::error_code writeFooter(const SampleProfileMap &ProfileMap) override {
+    JOS->arrayEnd();
+    return sampleprof_error::success;
+  }
+
+private:
+  std::unique_ptr<json::OStream> JOS;
+
+  void writeSampleImpl(const FunctionSamples &S, bool TopLevel = false);
+
+  friend ErrorOr<std::unique_ptr<SampleProfileWriter>>
+  SampleProfileWriter::create(std::unique_ptr<raw_ostream> &OS,
+                              SampleProfileFormat Format);
+};
+
 /// Sample-based profile writer (binary format).
 class SampleProfileWriterBinary : public SampleProfileWriter {
 public:
   SampleProfileWriterBinary(std::unique_ptr<raw_ostream> &OS)
       : SampleProfileWriter(OS) {}
 
   std::error_code writeSample(const FunctionSamples &S) override;
 

llvm/lib/ProfileData/SampleProfWriter.cpp

 
 std::error_code SampleProfileWriter::write(const SampleProfileMap &ProfileMap) {
   if (std::error_code EC = writeHeader(ProfileMap))
     return EC;
 
   if (std::error_code EC = writeFuncProfiles(ProfileMap))
     return EC;
 
+  if (std::error_code EC = writeFooter(ProfileMap))
+    return EC;
+
   return sampleprof_error::success;
 }
 
 /// Return the current position and prepare to use it as the start
 /// position of a section given the section type \p Type and its position
 /// \p LayoutIdx in SectionHdrLayout.
 uint64_t
 SampleProfileWriterExtBinaryBase::markSectionStart(SecType Type,
std::error_code SampleProfileWriterText::writeSample(const FunctionSamples &S) {
   if (S.getContext().getAllAttributes()) {
     OS.indent(Indent + 1);
     OS << "!Attributes: " << S.getContext().getAllAttributes() << "\n";
   }
 
   return sampleprof_error::success;
 }
 
+void SampleProfileWriterJSON::writeSampleImpl(const FunctionSamples &S,
+                                              bool TopLevel) {
+  JOS->object([&] {
+    JOS->attribute("name", S.getName());
+    JOS->attribute("total", S.getTotalSamples());
+    if (TopLevel)
+      JOS->attribute("head", S.getHeadSamples());
+
+    SampleSorter<LineLocation, SampleRecord> SortedSamples(
+        S.getBodySamples());
+    if (!SortedSamples.get().empty()) {
+      JOS->attributeArray("body", [&] {
+        for (const auto &I : SortedSamples.get()) {
+          LineLocation Loc = I->first;
+          const SampleRecord &Sample = I->second;
+          JOS->object([&] {
+            JOS->attribute("line", Loc.LineOffset);
+            if (Loc.Discriminator)
+              JOS->attribute("discriminator", Loc.Discriminator);
+            JOS->attribute("samples", Sample.getSamples());
+
+            auto CallTargets = Sample.getSortedCallTargets();
+            if (!CallTargets.empty()) {
+              JOS->attributeArray("calls", [&] {
+                for (const auto &J : CallTargets) {
+                  JOS->object([&] {
+                    JOS->attribute("function", J.first);
+                    JOS->attribute("samples", J.second);
+                  });
+                }
+              });
+            }
+          });
+        }
+      });
+    }
+
+    SampleSorter<LineLocation, FunctionSamplesMap> SortedCallsiteSamples(
+        S.getCallsiteSamples());
+    auto CallsiteSamples = SortedCallsiteSamples.get();
+    if (!CallsiteSamples.empty()) {
+      JOS->attributeArray("callsites", [&] {
+        for (const auto &I : CallsiteSamples)
+          for (const auto &FS : I->second) {
+            LineLocation Loc = I->first;
+            const FunctionSamples &CalleeSamples = FS.second;
+            JOS->object([&] {
+              JOS->attribute("line", Loc.LineOffset);
+              if (Loc.Discriminator)
+                JOS->attribute("discriminator", Loc.Discriminator);
+              JOS->attributeArray("samples",
+                                  [&] { writeSampleImpl(CalleeSamples); });
+            });
+          }
+      });
+    }
+  });
+}
+
+std::error_code SampleProfileWriterJSON::writeSample(const FunctionSamples &S) {
+  writeSampleImpl(S, true);
+  return sampleprof_error::success;
+}
+
 std::error_code
 SampleProfileWriterBinary::writeContextIdx(const SampleContext &Context) {
   assert(!Context.hasContext() && "cs profile is not supported");
   return writeNameIdx(Context.getName());
 }
 
 std::error_code SampleProfileWriterBinary::writeNameIdx(StringRef FName) {
   auto &NTable = getNameTable();
SampleProfileWriter::create(std::unique_ptr<raw_ostream> &OS,
   if (Format == SPF_Binary)
     Writer.reset(new SampleProfileWriterRawBinary(OS));
   else if (Format == SPF_Ext_Binary)
     Writer.reset(new SampleProfileWriterExtBinary(OS));
   else if (Format == SPF_Compact_Binary)
     Writer.reset(new SampleProfileWriterCompactBinary(OS));
   else if (Format == SPF_Text)
     Writer.reset(new SampleProfileWriterText(OS));
+  else if (Format == SPF_JSON)
+    Writer.reset(new SampleProfileWriterJSON(OS));
   else if (Format == SPF_GCC)
     EC = sampleprof_error::unsupported_writing_format;
   else
     EC = sampleprof_error::unrecognized_format;
 
   if (EC)
     return EC;
 
   Writer->Format = Format;
   return std::move(Writer);
 }
 
 void SampleProfileWriter::computeSummary(const SampleProfileMap &ProfileMap) {
   SampleProfileSummaryBuilder Builder(ProfileSummaryBuilder::DefaultCutoffs);
   Summary = Builder.computeSummaryForProfiles(ProfileMap);
 }

llvm/test/tools/llvm-profdata/sample-profile-json.test

+RUN: llvm-profdata merge --sample --json %p/Inputs/sample-profile.proftext | FileCheck %s --check-prefix=JSON
+JSON: [
+JSON:   {
+JSON:     "name": "main",
+JSON:     "total": 184019,
+JSON:     "head": 0,
+JSON:     "body": [
+JSON:       {
+JSON:         "line": 4,
+JSON:         "samples": 534
+JSON:       },
+JSON:       {
+JSON:         "line": 4,
+JSON:         "discriminator": 2,
+JSON:         "samples": 534
+JSON:       },
+JSON:       {
+JSON:         "line": 5,
+JSON:         "samples": 1075
+JSON:       },
+JSON:       {
+JSON:         "line": 5,
+JSON:         "discriminator": 1,
+JSON:         "samples": 1075
+JSON:       },
+JSON:       {
+JSON:         "line": 6,
+JSON:         "samples": 2080
+JSON:       },
+JSON:       {
+JSON:         "line": 7,
+JSON:         "samples": 534
+JSON:       },
+JSON:       {
+JSON:         "line": 9,
+JSON:         "samples": 2064,
+JSON:         "calls": [
+JSON:           {
+JSON:             "function": "_Z3bari",
+JSON:             "samples": 1471
+JSON:           },
+JSON:           {
+JSON:             "function": "_Z3fooi",
+JSON:             "samples": 631
+JSON:           }
+JSON:         ]
+JSON:       }
+JSON:     ],
+JSON:     "callsites": [
+JSON:       {
+JSON:         "line": 10,
+JSON:         "samples": [
+JSON:           {
+JSON:             "name": "inline1",
+JSON:             "total": 1000,
+JSON:             "body": [
+JSON:               {
+JSON:                 "line": 1,
+JSON:                 "samples": 1000
+JSON:               }
+JSON:             ]
+JSON:           }
+JSON:         ]
+JSON:       },
+JSON:       {
+JSON:         "line": 10,
+JSON:         "samples": [
+JSON:           {
+JSON:             "name": "inline2",
+JSON:             "total": 2000,
+JSON:             "body": [
+JSON:               {
+JSON:                 "line": 1,
+JSON:                 "samples": 2000
+JSON:               }
+JSON:             ]
+JSON:           }
+JSON:         ]
+JSON:       }
+JSON:     ]
+JSON:   },
+JSON:   {
+JSON:     "name": "_Z3bari",
+JSON:     "total": 20301,
+JSON:     "head": 1437,
+JSON:     "body": [
+JSON:       {
+JSON:         "line": 1,
+JSON:         "samples": 1437
+JSON:       }
+JSON:     ]
+JSON:   },
+JSON:   {
+JSON:     "name": "_Z3fooi",
+JSON:     "total": 7711,
+JSON:     "head": 610,
+JSON:     "body": [
+JSON:       {
+JSON:         "line": 1,
+JSON:         "samples": 610
+JSON:       }
+JSON:     ]
+JSON:   }
+JSON: ]

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

 #include <algorithm>
 #include <queue>
 
 using namespace llvm;
 
 enum ProfileFormat {
   PF_None = 0,
   PF_Text,
+  PF_JSON,
   PF_Compact_Binary,
   PF_Ext_Binary,
   PF_GCC,
   PF_Binary
 };
 
 static void warn(Twine Message, std::string Whence = "",
                  std::string Hint = "") {
for (const auto &CallsiteSamples : Samples.getCallsiteSamples()) {
     }
   }
   return Result;
 }
 
 static sampleprof::SampleProfileFormat FormatMap[] = {
     sampleprof::SPF_None,
     sampleprof::SPF_Text,
+    sampleprof::SPF_JSON,
     sampleprof::SPF_Compact_Binary,
     sampleprof::SPF_Ext_Binary,
     sampleprof::SPF_GCC,
     sampleprof::SPF_Binary};
 
 static std::unique_ptr<MemoryBuffer>
 getInputFileBuf(const StringRef &InputFile) {
   if (InputFile == "")
static int merge_main(int argc, const char *argv[]) {
   cl::opt<ProfileFormat> OutputFormat(
       cl::desc("Format of output profile"), cl::init(PF_Binary),
       cl::values(
           clEnumValN(PF_Binary, "binary", "Binary encoding (default)"),
           clEnumValN(PF_Compact_Binary, "compbinary",
                      "Compact binary encoding"),
           clEnumValN(PF_Ext_Binary, "extbinary", "Extensible binary encoding"),
           clEnumValN(PF_Text, "text", "Text encoding"),
+          clEnumValN(PF_JSON, "json", "JSON"),
           clEnumValN(PF_GCC, "gcc",
                      "GCC encoding (only meaningful for -sample)")));
   cl::opt<FailureMode> FailureMode(
       "failure-mode", cl::init(failIfAnyAreInvalid), cl::desc("Failure mode:"),
       cl::values(clEnumValN(failIfAnyAreInvalid, "any",
                             "Fail if any profile is invalid."),
                  clEnumValN(failIfAllAreInvalid, "all",
                             "Fail only if all profiles are invalid.")));