[ctxprof] dump profiles using yaml (for testing) (#123108)

This is a follow-up from PR #122545, which enabled converting yaml to contextual profiles.

This change uses the lower level yaml APIs because:
- the mapping APIs `llvm::yaml` offers don't work with `const` values, because they (the APIs) want to enable both serialization and deserialization
- building a helper data structure would be an alternative, but it'd be either memory-consuming or overly-complex design, given the recursive nature of the contextual profiles.

Author: mtrofin

llvm/include/llvm/Analysis/CtxProfAnalysis.h

@@ -140,7 +140,7 @@ class CtxProfAnalysis : public AnalysisInfoMixin<CtxProfAnalysis> {
 class CtxProfAnalysisPrinterPass
     : public PassInfoMixin<CtxProfAnalysisPrinterPass> {
 public:
-  enum class PrintMode { Everything, JSON };
+  enum class PrintMode { Everything, YAML };
   explicit CtxProfAnalysisPrinterPass(raw_ostream &OS);
 
   PreservedAnalyses run(Module &M, ModuleAnalysisManager &MAM);

llvm/include/llvm/ProfileData/PGOCtxProfReader.h

@@ -183,5 +183,8 @@ class PGOCtxProfileReader final {
 
   Expected<std::map<GlobalValue::GUID, PGOCtxProfContext>> loadContexts();
 };
+
+void convertCtxProfToYaml(raw_ostream &OS,
+                          const PGOCtxProfContext::CallTargetMapTy &);
 } // namespace llvm
 #endif

llvm/include/llvm/ProfileData/PGOCtxProfWriter.h

@@ -81,14 +81,6 @@ class PGOCtxProfileWriter final {
   static constexpr StringRef ContainerMagic = "CTXP";
 };
 
-/// Representation of the context node suitable for yaml / json serialization /
-/// deserialization.
-struct SerializableCtxRepresentation {
-  ctx_profile::GUID Guid = 0;
-  std::vector<uint64_t> Counters;
-  std::vector<std::vector<SerializableCtxRepresentation>> Callsites;
-};
-
 Error createCtxProfFromYAML(StringRef Profile, raw_ostream &Out);
 } // namespace llvm
 #endif

llvm/lib/Analysis/CtxProfAnalysis.cpp

@@ -19,7 +19,6 @@
 #include "llvm/IR/PassManager.h"
 #include "llvm/ProfileData/PGOCtxProfReader.h"
 #include "llvm/Support/CommandLine.h"
-#include "llvm/Support/JSON.h"
 #include "llvm/Support/MemoryBuffer.h"
 
 #define DEBUG_TYPE "ctx_prof"
@@ -31,49 +30,13 @@ cl::opt<std::string>
 
 static cl::opt<CtxProfAnalysisPrinterPass::PrintMode> PrintLevel(
     "ctx-profile-printer-level",
-    cl::init(CtxProfAnalysisPrinterPass::PrintMode::JSON), cl::Hidden,
+    cl::init(CtxProfAnalysisPrinterPass::PrintMode::YAML), cl::Hidden,
     cl::values(clEnumValN(CtxProfAnalysisPrinterPass::PrintMode::Everything,
                           "everything", "print everything - most verbose"),
-               clEnumValN(CtxProfAnalysisPrinterPass::PrintMode::JSON, "json",
-                          "just the json representation of the profile")),
+               clEnumValN(CtxProfAnalysisPrinterPass::PrintMode::YAML, "yaml",
+                          "just the yaml representation of the profile")),
     cl::desc("Verbosity level of the contextual profile printer pass."));
 
-namespace llvm {
-namespace json {
-Value toJSON(const PGOCtxProfContext &P) {
-  Object Ret;
-  Ret["Guid"] = P.guid();
-  Ret["Counters"] = Array(P.counters());
-  if (P.callsites().empty())
-    return Ret;
-  auto AllCS =
-      ::llvm::map_range(P.callsites(), [](const auto &P) { return P.first; });
-  auto MaxIt = ::llvm::max_element(AllCS);
-  assert(MaxIt != AllCS.end() && "We should have a max value because the "
-                                 "callsites collection is not empty.");
-  Array CSites;
-  // Iterate to, and including, the maximum index.
-  for (auto I = 0U, Max = *MaxIt; I <= Max; ++I) {
-    CSites.push_back(Array());
-    Array &Targets = *CSites.back().getAsArray();
-    if (P.hasCallsite(I))
-      for (const auto &[_, Ctx] : P.callsite(I))
-        Targets.push_back(toJSON(Ctx));
-  }
-  Ret["Callsites"] = std::move(CSites);
-
-  return Ret;
-}
-
-Value toJSON(const PGOCtxProfContext::CallTargetMapTy &P) {
-  Array Ret;
-  for (const auto &[_, Ctx] : P)
-    Ret.push_back(toJSON(Ctx));
-  return Ret;
-}
-} // namespace json
-} // namespace llvm
-
 const char *AssignGUIDPass::GUIDMetadataName = "guid";
 
 PreservedAnalyses AssignGUIDPass::run(Module &M, ModuleAnalysisManager &MAM) {
@@ -214,15 +177,13 @@ PreservedAnalyses CtxProfAnalysisPrinterPass::run(Module &M,
          << ". MaxCallsiteID: " << FuncInfo.NextCallsiteIndex << "\n";
   }
 
-  const auto JSONed = ::llvm::json::toJSON(C.profiles());
-
   if (Mode == PrintMode::Everything)
     OS << "\nCurrent Profile:\n";
-  OS << formatv("{0:2}", JSONed);
-  if (Mode == PrintMode::JSON)
+  convertCtxProfToYaml(OS, C.profiles());
+  OS << "\n";
+  if (Mode == PrintMode::YAML)
     return PreservedAnalyses::all();
 
-  OS << "\n";
   OS << "\nFlat Profile:\n";
   auto Flat = C.flatten();
   for (const auto &[Guid, Counters] : Flat) {

llvm/lib/ProfileData/PGOCtxProfReader.cpp

@@ -17,6 +17,10 @@
 #include "llvm/ProfileData/InstrProf.h"
 #include "llvm/ProfileData/PGOCtxProfWriter.h"
 #include "llvm/Support/Error.h"
+#include "llvm/Support/ErrorHandling.h"
+#include "llvm/Support/YAMLTraits.h"
+#include <iterator>
+#include <utility>
 
 using namespace llvm;
 
@@ -176,3 +180,86 @@ PGOCtxProfileReader::loadContexts() {
   }
   return std::move(Ret);
 }
+
+namespace {
+// We want to pass `const` values PGOCtxProfContext references to the yaml
+// converter, and the regular yaml mapping APIs are designed to handle both
+// serialization and deserialization, which prevents using const for
+// serialization. Using an intermediate datastructure is overkill, both
+// space-wise and design complexity-wise. Instead, we use the lower-level APIs.
+void toYaml(yaml::Output &Out, const PGOCtxProfContext &Ctx);
+
+void toYaml(yaml::Output &Out,
+            const PGOCtxProfContext::CallTargetMapTy &CallTargets) {
+  Out.beginSequence();
+  size_t Index = 0;
+  void *SaveData = nullptr;
+  for (const auto &[_, Ctx] : CallTargets) {
+    Out.preflightElement(Index++, SaveData);
+    toYaml(Out, Ctx);
+    Out.postflightElement(nullptr);
+  }
+  Out.endSequence();
+}
+
+void toYaml(yaml::Output &Out,
+            const PGOCtxProfContext::CallsiteMapTy &Callsites) {
+  auto AllCS = ::llvm::make_first_range(Callsites);
+  auto MaxIt = ::llvm::max_element(AllCS);
+  assert(MaxIt != AllCS.end() && "We should have a max value because the "
+                                 "callsites collection is not empty.");
+  void *SaveData = nullptr;
+  Out.beginSequence();
+  for (auto I = 0U; I <= *MaxIt; ++I) {
+    Out.preflightElement(I, SaveData);
+    auto It = Callsites.find(I);
+    if (It == Callsites.end()) {
+      // This will produce a `[ ]` sequence, which is what we want here.
+      Out.beginFlowSequence();
+      Out.endFlowSequence();
+    } else {
+      toYaml(Out, It->second);
+    }
+    Out.postflightElement(nullptr);
+  }
+  Out.endSequence();
+}
+
+void toYaml(yaml::Output &Out, const PGOCtxProfContext &Ctx) {
+  yaml::EmptyContext Empty;
+  Out.beginMapping();
+  void *SaveInfo = nullptr;
+  bool UseDefault = false;
+  {
+    Out.preflightKey("Guid", /*Required=*/true, /*SameAsDefault=*/false,
+                     UseDefault, SaveInfo);
+    auto Guid = Ctx.guid();
+    yaml::yamlize(Out, Guid, true, Empty);
+    Out.postflightKey(nullptr);
+  }
+  {
+    Out.preflightKey("Counters", true, false, UseDefault, SaveInfo);
+    Out.beginFlowSequence();
+    for (size_t I = 0U, E = Ctx.counters().size(); I < E; ++I) {
+      Out.preflightFlowElement(I, SaveInfo);
+      uint64_t V = Ctx.counters()[I];
+      yaml::yamlize(Out, V, true, Empty);
+      Out.postflightFlowElement(SaveInfo);
+    }
+    Out.endFlowSequence();
+    Out.postflightKey(nullptr);
+  }
+  if (!Ctx.callsites().empty()) {
+    Out.preflightKey("Callsites", true, false, UseDefault, SaveInfo);
+    toYaml(Out, Ctx.callsites());
+    Out.postflightKey(nullptr);
+  }
+  Out.endMapping();
+}
+} // namespace
+
+void llvm::convertCtxProfToYaml(
+    raw_ostream &OS, const PGOCtxProfContext::CallTargetMapTy &Profiles) {
+  yaml::Output Out(OS);
+  toYaml(Out, Profiles);
+}

llvm/lib/ProfileData/PGOCtxProfWriter.cpp

@@ -14,7 +14,6 @@
 #include "llvm/Bitstream/BitCodeEnums.h"
 #include "llvm/ProfileData/CtxInstrContextNode.h"
 #include "llvm/Support/Error.h"
-#include "llvm/Support/JSON.h"
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/YAMLTraits.h"
 #include "llvm/Support/raw_ostream.h"
@@ -89,6 +88,15 @@ void PGOCtxProfileWriter::write(const ContextNode &RootNode) {
 }
 
 namespace {
+
+/// Representation of the context node suitable for yaml serialization /
+/// deserialization.
+struct SerializableCtxRepresentation {
+  ctx_profile::GUID Guid = 0;
+  std::vector<uint64_t> Counters;
+  std::vector<std::vector<SerializableCtxRepresentation>> Callsites;
+};
+
 ctx_profile::ContextNode *
 createNode(std::vector<std::unique_ptr<char[]>> &Nodes,
            const std::vector<SerializableCtxRepresentation> &DCList);

llvm/test/Analysis/CtxProfAnalysis/full-cycle.ll

@@ -88,54 +88,20 @@ Function Info:
 10507721908651011566 : entrypoint. MaxCounterID: 1. MaxCallsiteID: 2
 
 Current Profile:
-[
-  {
-    "Callsites": [
-      [
-        {
-          "Callsites": [
-            [
-              {
-                "Counters": [
-                  10,
-                  7
-                ],
-                "Guid": 3087265239403591524
-              }
-            ]
-          ],
-          "Counters": [
-            7
-          ],
-          "Guid": 2072045998141807037
-        }
-      ],
-      [
-        {
-          "Callsites": [
-            [
-              {
-                "Counters": [
-                  1,
-                  2
-                ],
-                "Guid": 3087265239403591524
-              }
-            ]
-          ],
-          "Counters": [
-            2
-          ],
-          "Guid": 4197650231481825559
-        }
-      ]
-    ],
-    "Counters": [
-      1
-    ],
-    "Guid": 10507721908651011566
-  }
-]
+
+- Guid:            10507721908651011566
+  Counters:        [ 1 ]
+  Callsites:
+    - - Guid:            2072045998141807037
+        Counters:        [ 7 ]
+        Callsites:
+          - - Guid:            3087265239403591524
+              Counters:        [ 10, 7 ]
+    - - Guid:            4197650231481825559
+        Counters:        [ 2 ]
+        Callsites:
+          - - Guid:            3087265239403591524
+              Counters:        [ 1, 2 ]
 
 Flat Profile:
 2072045998141807037 : 7 

llvm/test/Analysis/CtxProfAnalysis/inline.ll

@@ -1,11 +1,12 @@
+; REQUIRES: x86_64-linux
 ; RUN: rm -rf %t
 ; RUN: split-file %s %t
 ; RUN: llvm-ctxprof-util fromYAML --input=%t/profile.yaml --output=%t/profile.ctxprofdata
 
 ; RUN: opt -passes='module-inline,print<ctx-prof-analysis>' -ctx-profile-printer-level=everything %t/module.ll -S \
-; RUN:   -use-ctx-profile=%t/profile.ctxprofdata -ctx-profile-printer-level=json \
-; RUN:   -o - 2> %t/profile-final.txt | FileCheck %s
-; RUN: %python %S/json_equals.py %t/profile-final.txt %t/expected.json
+; RUN:   -use-ctx-profile=%t/profile.ctxprofdata -ctx-profile-printer-level=yaml \
+; RUN:   -o - 2> %t/profile-final.yaml | FileCheck %s
+; RUN: diff %t/profile-final.yaml %t/expected.yaml
 
 ; There are 2 calls to @a from @entrypoint. We only inline the one callsite
 ; marked as alwaysinline, the rest are blocked (marked noinline). After the inline,
@@ -109,17 +110,16 @@ define i32 @b() !guid !2 {
                   Callsites:  -
                                 - Guid: 1002
                                   Counters: [500]
-;--- expected.json
-[
-  { "Guid": 1000,
-    "Counters": [10, 2, 8, 100],
-    "Callsites": [
-      [],
-      [ { "Guid": 1001,
-          "Counters": [8, 500],
-          "Callsites": [[{"Guid": 1002, "Counters": [500]}]]}
-      ],
-      [{ "Guid": 1002, "Counters": [100]}]
-    ]
-  }
-]
+;--- expected.yaml
+
+- Guid:            1000
+  Counters:        [ 10, 2, 8, 100 ]
+  Callsites:
+    - [  ]
+    - - Guid:            1001
+        Counters:        [ 8, 500 ]
+        Callsites:
+          - - Guid:            1002
+              Counters:        [ 500 ]
+    - - Guid:            1002
+        Counters:        [ 100 ]