[Mlir-commits] [mlir] [MLIR] Expose API for setting output format enum on DefaultTimingManger (PR #104548)

[Shaurya Sharma]

https://github.com/shaurya0 created https://github.com/llvm/llvm-project/pull/104548

This change allows for selecting the output format of the `DefaultTimingManager` through an API which was previously only accessible when running e.g. `mlir-opt --mlir-timing --mlir-output-format=<value>` at the command line. The use case is to be able to enable MLIR pass timing in situations when setting CL options is not viable. The `DefaultTimingManager` already allows for setting the `DisplayMode` in a similar manner.

>From 79d33e277d77885e6bd90c01db92153ce357914c Mon Sep 17 00:00:00 2001
From: Shaurya Sharma <shaurya at modular.com>
Date: Fri, 16 Aug 2024 08:20:41 +0200
Subject: [PATCH] [MLIR] Expose API for setting output format enum on
 DefaultTimingManger

---
 mlir/include/mlir/Support/Timing.h | 3 +++
 mlir/lib/Support/Timing.cpp        | 8 ++++++++
 2 files changed, 11 insertions(+)

diff --git a/mlir/include/mlir/Support/Timing.h b/mlir/include/mlir/Support/Timing.h
index a8a4bfd1c6cf16..133456e8499c6f 100644
--- a/mlir/include/mlir/Support/Timing.h
+++ b/mlir/include/mlir/Support/Timing.h
@@ -430,6 +430,9 @@ class DefaultTimingManager : public TimingManager {
   /// Change the stream where the output will be printed to.
   void setOutput(std::unique_ptr<OutputStrategy> output);
 
+  /// Change the output format.
+  void setOutputFormat(OutputFormat format);
+
   /// Print and clear the timing results. Only call this when there are no more
   /// references to nested timers around, as printing post-processes and clears
   /// the timers.
diff --git a/mlir/lib/Support/Timing.cpp b/mlir/lib/Support/Timing.cpp
index ac16eb7d224c9a..9cde0ef6503e15 100644
--- a/mlir/lib/Support/Timing.cpp
+++ b/mlir/lib/Support/Timing.cpp
@@ -527,6 +527,14 @@ void DefaultTimingManager::setOutput(std::unique_ptr<OutputStrategy> output) {
   out = std::move(output);
 }
 
+/// Change the output format.
+void DefaultTimingManager::setOutputFormat(OutputFormat outputFormat) {
+  if (outputFormat == OutputFormat::Text)
+    setOutput(std::make_unique<OutputTextStrategy>(llvm::errs()));
+  else if (outputFormat == OutputFormat::Json)
+    setOutput(std::make_unique<OutputJsonStrategy>(llvm::errs()));
+}
+
 /// Print and clear the timing results.
 void DefaultTimingManager::print() {
   if (impl->enabled) {