[ctxprof] Move test serialization to yaml (llvm#122545)

We have a textual representation of contextual profiles for test scenarios, mainly. This patch moves that to YAML instead of JSON. YAML is more succinct and readable (some of the .ll tests should be illustrative). In addition, JSON is parse-able by the YAML reader. A subsequent patch will address deserialization. (thanks, @kazutakahirata, for showing me how to use the llvm YAML reader/writer APIs, which I incorrectly thought to be more low-level than the JSON ones!)
DKLoehr · Jan 17, 2025 · 96cb54e · 96cb54e
1 parent 6dbac7f
commit 96cb54e
Show file tree

Hide file tree

Showing 31 changed files with 195 additions and 310 deletions.
diff --git a/llvm/include/llvm/ProfileData/PGOCtxProfWriter.h b/llvm/include/llvm/ProfileData/PGOCtxProfWriter.h
@@ -81,6 +81,14 @@ class PGOCtxProfileWriter final {
   static constexpr StringRef ContainerMagic = "CTXP";
 };
 
-Error createCtxProfFromJSON(StringRef Profile, raw_ostream &Out);
+/// 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
diff --git a/llvm/lib/ProfileData/PGOCtxProfWriter.cpp b/llvm/lib/ProfileData/PGOCtxProfWriter.cpp
@@ -13,7 +13,11 @@
 #include "llvm/ProfileData/PGOCtxProfWriter.h"
 #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"
 
 using namespace llvm;
 using namespace llvm::ctx_profile;
@@ -85,22 +89,15 @@ void PGOCtxProfileWriter::write(const ContextNode &RootNode) {
 }
 
 namespace {
-// A structural representation of the JSON input.
-struct DeserializableCtx {
-  ctx_profile::GUID Guid = 0;
-  std::vector<uint64_t> Counters;
-  std::vector<std::vector<DeserializableCtx>> Callsites;
-};
-
 ctx_profile::ContextNode *
 createNode(std::vector<std::unique_ptr<char[]>> &Nodes,
-           const std::vector<DeserializableCtx> &DCList);
+           const std::vector<SerializableCtxRepresentation> &DCList);
 
 // Convert a DeserializableCtx into a ContextNode, potentially linking it to
 // its sibling (e.g. callee at same callsite) "Next".
 ctx_profile::ContextNode *
 createNode(std::vector<std::unique_ptr<char[]>> &Nodes,
-           const DeserializableCtx &DC,
+           const SerializableCtxRepresentation &DC,
            ctx_profile::ContextNode *Next = nullptr) {
   auto AllocSize = ctx_profile::ContextNode::getAllocSize(DC.Counters.size(),
                                                           DC.Callsites.size());
@@ -115,38 +112,34 @@ createNode(std::vector<std::unique_ptr<char[]>> &Nodes,
   return Ret;
 }
 
-// Convert a list of DeserializableCtx into a linked list of ContextNodes.
+// Convert a list of SerializableCtxRepresentation into a linked list of
+// ContextNodes.
 ctx_profile::ContextNode *
 createNode(std::vector<std::unique_ptr<char[]>> &Nodes,
-           const std::vector<DeserializableCtx> &DCList) {
+           const std::vector<SerializableCtxRepresentation> &DCList) {
   ctx_profile::ContextNode *List = nullptr;
   for (const auto &DC : DCList)
     List = createNode(Nodes, DC, List);
   return List;
 }
 } // namespace
 
-namespace llvm {
-namespace json {
-bool fromJSON(const Value &E, DeserializableCtx &R, Path P) {
-  json::ObjectMapper Mapper(E, P);
-  return Mapper && Mapper.map("Guid", R.Guid) &&
-         Mapper.map("Counters", R.Counters) &&
-         Mapper.mapOptional("Callsites", R.Callsites);
-}
-} // namespace json
-} // namespace llvm
-
-Error llvm::createCtxProfFromJSON(StringRef Profile, raw_ostream &Out) {
-  auto P = json::parse(Profile);
-  if (!P)
-    return P.takeError();
+LLVM_YAML_IS_SEQUENCE_VECTOR(SerializableCtxRepresentation)
+LLVM_YAML_IS_SEQUENCE_VECTOR(std::vector<SerializableCtxRepresentation>)
+template <> struct yaml::MappingTraits<SerializableCtxRepresentation> {
+  static void mapping(yaml::IO &IO, SerializableCtxRepresentation &SCR) {
+    IO.mapRequired("Guid", SCR.Guid);
+    IO.mapRequired("Counters", SCR.Counters);
+    IO.mapOptional("Callsites", SCR.Callsites);
+  }
+};
 
-  json::Path::Root R("");
-  std::vector<DeserializableCtx> DCList;
-  if (!fromJSON(*P, DCList, R))
-    return R.getError();
-  // Nodes provides memory backing for the ContextualNodes.
+Error llvm::createCtxProfFromYAML(StringRef Profile, raw_ostream &Out) {
+  yaml::Input In(Profile);
+  std::vector<SerializableCtxRepresentation> DCList;
+  In >> DCList;
+  if (In.error())
+    return createStringError(In.error(), "incorrect yaml content");
   std::vector<std::unique_ptr<char[]>> Nodes;
   std::error_code EC;
   if (EC)
@@ -162,4 +155,4 @@ Error llvm::createCtxProfFromJSON(StringRef Profile, raw_ostream &Out) {
   if (EC)
     return createStringError(EC, "failed to write output");
   return Error::success();
-}
+}
diff --git a/llvm/test/Analysis/CtxProfAnalysis/flatten-and-annotate.ll b/llvm/test/Analysis/CtxProfAnalysis/flatten-and-annotate.ll
@@ -2,7 +2,7 @@
 ;
 ; RUN: rm -rf %t
 ; RUN: split-file %s %t
-; RUN: llvm-ctxprof-util fromJSON --input=%t/profile.json --output=%t/profile.ctxprofdata
+; RUN: llvm-ctxprof-util fromYAML --input=%t/profile.yaml --output=%t/profile.ctxprofdata
 ; RUN: opt -module-summary -passes='thinlto-pre-link<O2>' -use-ctx-profile=%t/profile.ctxprofdata \
 ; RUN:   %t/example.ll -S -o %t/prelink.ll
 ; RUN: FileCheck --input-file %t/prelink.ll %s --check-prefix=PRELINK
@@ -58,27 +58,15 @@
 ; CHECK:       ![[AN_ENTRYPOINT_EP]] = !{!"function_entry_count", i64 100}
 ; CHECK:       ![[AN_ENTRYPOINT_BW]] = !{!"branch_weights", i32 40, i32 60} 
 
-;--- profile.json
-[
-  {
-    "Guid": 4909520559318251808,
-    "Counters": [100, 40],
-    "Callsites": [
-      [
-        {
-          "Guid": 11872291593386833696,
-          "Counters": [ 100, 5 ]
-        }
-      ],
-      [
-        {
-          "Guid": 11872291593386833696,
-          "Counters": [ 40, 10 ]
-        }
-      ]
-    ]
-  }
-]
+;--- profile.yaml
+- Guid: 4909520559318251808
+  Counters: [100, 40]
+  Callsites: -
+              - Guid: 11872291593386833696
+                Counters: [ 100, 5 ]
+             -
+              - Guid: 11872291593386833696
+                Counters: [ 40, 10 ]
 ;--- example.ll
 declare void @bar()
 

diff --git a/llvm/test/Analysis/CtxProfAnalysis/flatten-check-path.ll b/llvm/test/Analysis/CtxProfAnalysis/flatten-check-path.ll
@@ -3,9 +3,9 @@
 ; already visited blocks count as taken (i.e. the flow continues through them).
 ;
 ; RUN: split-file %s %t
-; RUN: llvm-ctxprof-util fromJSON --input=%t/profile_ok.json --output=%t/profile_ok.ctxprofdata
-; RUN: llvm-ctxprof-util fromJSON --input=%t/profile_pump.json --output=%t/profile_pump.ctxprofdata
-; RUN: llvm-ctxprof-util fromJSON --input=%t/profile_unreachable.json --output=%t/profile_unreachable.ctxprofdata
+; RUN: llvm-ctxprof-util fromYAML --input=%t/profile_ok.yaml --output=%t/profile_ok.ctxprofdata
+; RUN: llvm-ctxprof-util fromYAML --input=%t/profile_pump.yaml --output=%t/profile_pump.ctxprofdata
+; RUN: llvm-ctxprof-util fromYAML --input=%t/profile_unreachable.yaml --output=%t/profile_unreachable.ctxprofdata
 ;
 ; RUN: opt -passes=ctx-prof-flatten %t/example_ok.ll -use-ctx-profile=%t/profile_ok.ctxprofdata -S -o - | FileCheck %s
 ; RUN: not --crash opt -passes=ctx-prof-flatten %t/message_pump.ll -use-ctx-profile=%t/profile_pump.ctxprofdata -S 2>&1 | FileCheck %s --check-prefix=ASSERTION
@@ -38,8 +38,9 @@ exit:
 }
 !0 = !{i64 1234}
 
-;--- profile_ok.json
-[{"Guid":1234, "Counters":[2, 2, 1, 2]}]
+;--- profile_ok.yaml
+- Guid: 1234 
+  Counters: [2, 2, 1, 2]
 
 ;--- message_pump.ll
 ; This is a message pump: the loop never exits. This should result in an
@@ -59,8 +60,9 @@ exit:
 }
 !0 = !{i64 1234}
 
-;--- profile_pump.json
-[{"Guid":1234, "Counters":[2, 10, 0]}]
+;--- profile_pump.yaml
+- Guid: 1234
+  Counters: [2, 10, 0]
 
 ;--- unreachable.ll
 ; An unreachable block is reached, that's an error
@@ -81,5 +83,6 @@ exit:
 }
 !0 = !{i64 1234}
 
-;--- profile_unreachable.json
-[{"Guid":1234, "Counters":[2, 1, 1, 2]}]
+;--- profile_unreachable.yaml
+- Guid: 1234
+  Counters: [2, 1, 1, 2]
diff --git a/llvm/test/Analysis/CtxProfAnalysis/flatten-icp.ll b/llvm/test/Analysis/CtxProfAnalysis/flatten-icp.ll
@@ -1,5 +1,5 @@
 ; RUN: split-file %s %t
-; RUN: llvm-ctxprof-util fromJSON --input %t/profile.json --output %t/profile.ctxprofdata
+; RUN: llvm-ctxprof-util fromYAML --input %t/profile.yaml --output %t/profile.ctxprofdata
 ;
 ; In the given profile, in one of the contexts the indirect call is taken, the
 ; target we're trying to ICP - GUID:2000 - doesn't appear at all. That should
@@ -45,11 +45,18 @@ attributes #1 = { noinline }
 !1 = !{i64 3000}
 !2 = !{i64 4000}
 
-;--- profile.json
-[ {
-  "Guid": 4000, "Counters":[10], "Callsites": [
-    [{"Guid":3000, "Counters":[10], "Callsites":[[{"Guid":1000, "Counters":[10]}]]}],
-    [{"Guid":3000, "Counters":[10], "Callsites":[[{"Guid":9000, "Counters":[10]}]]}]
-  ]
-}
-]
+;--- profile.yaml
+- Guid: 4000
+  Counters: [10]
+  Callsites:  -
+                - Guid: 3000
+                  Counters: [10]
+                  Callsites: -
+                              - Guid: 1000
+                                Counters: [10]
+              -
+                - Guid: 3000
+                  Counters: [10]
+                  Callsites: -
+                              - Guid: 9000
+                                Counters: [10]
diff --git a/llvm/test/Analysis/CtxProfAnalysis/flatten-zero-path.ll b/llvm/test/Analysis/CtxProfAnalysis/flatten-zero-path.ll
@@ -1,6 +1,6 @@
 ; Check that flattened profile lowering handles cold subgraphs that end in "unreachable"
 ; RUN: split-file %s %t
-; RUN: llvm-ctxprof-util fromJSON --input=%t/profile.json --output=%t/profile.ctxprofdata
+; RUN: llvm-ctxprof-util fromYAML --input=%t/profile.yaml --output=%t/profile.ctxprofdata
 ; RUN: opt -passes=ctx-prof-flatten %t/example.ll -use-ctx-profile=%t/profile.ctxprofdata -S -o - | FileCheck %s
 
 ; CHECK-LABEL: entry:
@@ -51,5 +51,6 @@ exit:
 
 !0 = !{i64 1234}
 
-;--- profile.json
-[{"Guid":1234, "Counters":[6,0,0,0]}]
+;--- profile.yaml
+- Guid: 1234
+  Counters: [6,0,0,0]
diff --git a/llvm/test/Analysis/CtxProfAnalysis/full-cycle.ll b/llvm/test/Analysis/CtxProfAnalysis/full-cycle.ll
@@ -8,7 +8,7 @@
 ; different counter values, and we expect resulting flat profile to be the sum
 ; (of values at the same index).
 ;
-; RUN: llvm-ctxprof-util fromJSON --input=%t/profile.json --output=%t/profile.ctxprofdata
+; RUN: llvm-ctxprof-util fromYAML --input=%t/profile.yaml --output=%t/profile.ctxprofdata
 ;
 ; RUN: opt -module-summary -passes='thinlto-pre-link<O2>' -use-ctx-profile=%t/profile.ctxprofdata -o %t/m1.bc %t/m1.ll
 ; RUN: opt -module-summary -passes='thinlto-pre-link<O2>' -use-ctx-profile=%t/profile.ctxprofdata -o %t/m2.bc %t/m2.ll
@@ -65,55 +65,21 @@ define void @entrypoint() {
   call void @f3()
   ret void
 }
-;--- profile.json
-[
-  {
-    "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
-  }
-]
+;--- profile.yaml
+- 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]
 ;--- expected.txt
 Function Info:
 2072045998141807037 : f1. MaxCounterID: 1. MaxCallsiteID: 1

diff --git a/llvm/test/Analysis/CtxProfAnalysis/handle-select.ll b/llvm/test/Analysis/CtxProfAnalysis/handle-select.ll
@@ -4,7 +4,7 @@
 ; the `select` is elided.
 ;
 ; RUN: split-file %s %t
-; RUN: llvm-ctxprof-util fromJSON --input=%t/profile.json --output=%t/profile.ctxprofdata
+; RUN: llvm-ctxprof-util fromYAML --input=%t/profile.yaml --output=%t/profile.ctxprofdata
 ;
 ; RUN: opt -passes=ctx-instr-gen %t/example.ll -use-ctx-profile=%t/profile.ctxprofdata -S -o - | FileCheck %s --check-prefix=INSTR
 ; RUN: opt -passes=ctx-instr-gen,module-inline %t/example.ll -use-ctx-profile=%t/profile.ctxprofdata -S -o - | FileCheck %s --check-prefix=POST-INL
@@ -72,5 +72,12 @@ define i32 @bar(i32 %t) !guid !1 {
 !0 = !{i64 1234}
 !1 = !{i64 5678}
 
-;--- profile.json
-[{"Guid":1234, "Counters":[10, 4], "Callsites":[[{"Guid": 5678, "Counters":[4,3]}],[{"Guid": 5678, "Counters":[6,6]}]]}]
+;--- profile.yaml
+- Guid: 1234
+  Counters: [10, 4]
+  Callsites:  -
+                - Guid: 5678
+                  Counters: [4,3]
+              - 
+                - Guid: 5678
+                  Counters: [6,6]