[llvm-branch-commits] [llvm] [ctxprof] add conversion to yaml to llvm-ctxprof-utils (PR #123131)

[Mircea Trofin via llvm-branch-commits]

https://github.com/mtrofin created https://github.com/llvm/llvm-project/pull/123131

None

>From c8e600b9969239c2491fd97f011172f4032524bb Mon Sep 17 00:00:00 2001
From: Mircea Trofin <mtrofin at google.com>
Date: Wed, 15 Jan 2025 14:10:19 -0800
Subject: [PATCH] [ctxprof] add conversion to yaml to llvm-ctxprof-utils

---
 .../tools/llvm-ctxprof-util/Inputs/valid.yaml | 26 ++++++------
 .../llvm-ctxprof-util/llvm-ctxprof-util.test  |  4 +-
 .../llvm-ctxprof-util/llvm-ctxprof-util.cpp   | 41 ++++++++++++++++---
 3 files changed, 51 insertions(+), 20 deletions(-)

diff --git a/llvm/test/tools/llvm-ctxprof-util/Inputs/valid.yaml b/llvm/test/tools/llvm-ctxprof-util/Inputs/valid.yaml
index 6080c2414d64a0..9bbf82d59c9134 100644
--- a/llvm/test/tools/llvm-ctxprof-util/Inputs/valid.yaml
+++ b/llvm/test/tools/llvm-ctxprof-util/Inputs/valid.yaml
@@ -1,13 +1,13 @@
-- Guid: 1000
-  Counters: [1, 2, 3]
-  Callsites:  - []
-              - 
-                - Guid: 2000
-                  Counters: [4, 5]
-                - Guid: 18446744073709551613
-                  Counters: [6, 7, 8]
-              - 
-                - Guid: 3000
-                  Counters: [40, 50]
-- Guid: 18446744073709551612
-  Counters: [5, 9, 10]
+
+- Guid:            1000
+  Counters:        [ 1, 2, 3 ]
+  Callsites:
+    - [  ]
+    - - Guid:            2000
+        Counters:        [ 4, 5 ]
+      - Guid:            18446744073709551613
+        Counters:        [ 6, 7, 8 ]
+    - - Guid:            3000
+        Counters:        [ 40, 50 ]
+- Guid:            18446744073709551612
+  Counters:        [ 5, 9, 10 ]
diff --git a/llvm/test/tools/llvm-ctxprof-util/llvm-ctxprof-util.test b/llvm/test/tools/llvm-ctxprof-util/llvm-ctxprof-util.test
index 91ebd1de59bb54..30bc8bce054101 100644
--- a/llvm/test/tools/llvm-ctxprof-util/llvm-ctxprof-util.test
+++ b/llvm/test/tools/llvm-ctxprof-util/llvm-ctxprof-util.test
@@ -4,7 +4,9 @@
 ; RUN: llvm-ctxprof-util fromYAML --input %S/Inputs/empty.yaml -output %t/empty.bitstream
 ; RUN: llvm-bcanalyzer --dump %t/empty.bitstream | FileCheck %s --check-prefix=EMPTY
 
-; RUN: llvm-ctxprof-util fromYAML --input %S/Inputs/valid.yaml -output %t/valid.bitstream
+; RUN: llvm-ctxprof-util fromYAML -input %S/Inputs/valid.yaml -output %t/valid.bitstream
+; RUN: llvm-ctxprof-util toYAML -input %t/valid.bitstream -output %t/valid2.yaml
+; RUN: diff %t/valid2.yaml %S/Inputs/valid.yaml
 
 ; For the valid case, check against a reference output.
 ; Note that uint64_t are printed as signed values by llvm-bcanalyzer:
diff --git a/llvm/tools/llvm-ctxprof-util/llvm-ctxprof-util.cpp b/llvm/tools/llvm-ctxprof-util/llvm-ctxprof-util.cpp
index cfa14b22c14698..314144ac6624c9 100644
--- a/llvm/tools/llvm-ctxprof-util/llvm-ctxprof-util.cpp
+++ b/llvm/tools/llvm-ctxprof-util/llvm-ctxprof-util.cpp
@@ -13,6 +13,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/IR/GlobalValue.h"
+#include "llvm/ProfileData/PGOCtxProfReader.h"
 #include "llvm/ProfileData/PGOCtxProfWriter.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Error.h"
@@ -23,6 +24,7 @@
 using namespace llvm;
 
 static cl::SubCommand FromYAML("fromYAML", "Convert from yaml");
+static cl::SubCommand ToYAML("toYAML", "Convert to yaml");
 
 static cl::opt<std::string> InputFilename(
     "input", cl::value_desc("input"), cl::init("-"),
@@ -35,15 +37,16 @@ static cl::opt<std::string> InputFilename(
         "'Contexts', optional. An array containing arrays of contexts. The "
         "context array at a position 'i' is the set of callees at that "
         "callsite index. Use an empty array to indicate no callees."),
-    cl::sub(FromYAML));
+    cl::sub(FromYAML), cl::sub(ToYAML));
 
 static cl::opt<std::string> OutputFilename("output", cl::value_desc("output"),
                                            cl::init("-"),
                                            cl::desc("Output file"),
-                                           cl::sub(FromYAML));
+                                           cl::sub(FromYAML), cl::sub(ToYAML));
 
+namespace {
 // Save the bitstream profile from the JSON representation.
-Error convertFromYAML() {
+Error convertFromYaml() {
   auto BufOrError =
       MemoryBuffer::getFileOrSTDIN(InputFilename, /*IsText=*/true);
   if (!BufOrError)
@@ -61,11 +64,30 @@ Error convertFromYAML() {
   return llvm::createCtxProfFromYAML(BufOrError.get()->getBuffer(), Out);
 }
 
+Error convertToYaml() {
+  auto BufOrError = MemoryBuffer::getFileOrSTDIN(InputFilename);
+  if (!BufOrError)
+    return createFileError(InputFilename, BufOrError.getError());
+
+  std::error_code EC;
+  raw_fd_ostream Out(OutputFilename, EC);
+  if (EC)
+    return createStringError(EC, "failed to open output");
+  PGOCtxProfileReader Reader(BufOrError.get()->getBuffer());
+  auto Prof = Reader.loadContexts();
+  if (!Prof)
+    return Prof.takeError();
+  llvm::convertCtxProfToYaml(Out, *Prof);
+  Out << "\n";
+  return Error::success();
+}
+} // namespace
+
 int main(int argc, const char **argv) {
   cl::ParseCommandLineOptions(argc, argv, "LLVM Contextual Profile Utils\n");
   ExitOnError ExitOnErr("llvm-ctxprof-util: ");
-  if (FromYAML) {
-    if (auto E = convertFromYAML()) {
+  auto HandleErr = [&](Error E) -> int {
+    if (E) {
       handleAllErrors(std::move(E), [&](const ErrorInfoBase &E) {
         E.log(errs());
         errs() << "\n";
@@ -73,7 +95,14 @@ int main(int argc, const char **argv) {
       return 1;
     }
     return 0;
-  }
+  };
+
+  if (FromYAML)
+    return HandleErr(convertFromYaml());
+
+  if (ToYAML)
+    return HandleErr(convertToYaml());
+
   cl::PrintHelpMessage();
   return 1;
 }