From f0594b6fd49f0668e8240776a3376fea5fd81267 Mon Sep 17 00:00:00 2001
From: Frederik Gossen <frgossen@google.com>
Date: Fri, 17 Mar 2023 14:47:18 -0400
Subject: [PATCH] [MLIR] Add label to print-ir pass

Differential Revision: https://reviews.llvm.org/D145523
---
 mlir/include/mlir/Transforms/Passes.h  |  7 ++++---
 mlir/include/mlir/Transforms/Passes.td |  3 +++
 mlir/lib/Transforms/PrintIR.cpp        | 18 +++++++++++++-----
 3 files changed, 20 insertions(+), 8 deletions(-)

diff --git a/mlir/include/mlir/Transforms/Passes.h b/mlir/include/mlir/Transforms/Passes.h
index 12986c9..1d95e25 100644
--- a/mlir/include/mlir/Transforms/Passes.h
+++ b/mlir/include/mlir/Transforms/Passes.h
@@ -31,10 +31,11 @@ class GreedyRewriteConfig;
 #define GEN_PASS_DECL_CANONICALIZER
 #define GEN_PASS_DECL_CONTROLFLOWSINK
 #define GEN_PASS_DECL_CSEPASS
+#define GEN_PASS_DECL_INLINER
 #define GEN_PASS_DECL_LOOPINVARIANTCODEMOTION
-#define GEN_PASS_DECL_STRIPDEBUGINFO
+#define GEN_PASS_DECL_PRINTIRPASS
 #define GEN_PASS_DECL_PRINTOPSTATS
-#define GEN_PASS_DECL_INLINER
+#define GEN_PASS_DECL_STRIPDEBUGINFO
 #define GEN_PASS_DECL_SCCP
 #define GEN_PASS_DECL_SYMBOLDCE
 #define GEN_PASS_DECL_SYMBOLPRIVATIZE
@@ -65,7 +66,7 @@ std::unique_ptr<Pass> createControlFlowSinkPass();
 std::unique_ptr<Pass> createCSEPass();
 
 /// Creates a pass to print IR on the debug stream.
-std::unique_ptr<Pass> createPrintIRPass();
+std::unique_ptr<Pass> createPrintIRPass(const PrintIRPassOptions & = {});
 
 /// Creates a pass that generates IR to verify ops at runtime.
 std::unique_ptr<Pass> createGenerateRuntimeVerificationPass();
diff --git a/mlir/include/mlir/Transforms/Passes.td b/mlir/include/mlir/Transforms/Passes.td
index 1ed014a..cecefb3 100644
--- a/mlir/include/mlir/Transforms/Passes.td
+++ b/mlir/include/mlir/Transforms/Passes.td
@@ -92,6 +92,9 @@ def PrintIRPass : Pass<"print-ir"> {
     purposes to inspect the IR at a specific point in the pipeline.
   }];
   let constructor = "mlir::createPrintIRPass()";
+  let options = [
+    Option<"label", "label", "std::string", /*default=*/"", "Label">,
+  ];
 }
 
 def GenerateRuntimeVerification : Pass<"generate-runtime-verification"> {
diff --git a/mlir/lib/Transforms/PrintIR.cpp b/mlir/lib/Transforms/PrintIR.cpp
index 41cb6ef..cc42c7e 100644
--- a/mlir/lib/Transforms/PrintIR.cpp
+++ b/mlir/lib/Transforms/PrintIR.cpp
@@ -7,6 +7,8 @@
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/Passes.h"
+#include "llvm/Support/Debug.h"
 
 namespace mlir {
 namespace {
@@ -15,15 +17,21 @@ namespace {
 #include "mlir/Transforms/Passes.h.inc"
 
 struct PrintIRPass : public impl::PrintIRPassBase<PrintIRPass> {
-  PrintIRPass() = default;
-
-  void runOnOperation() override { getOperation()->dump(); }
+  using impl::PrintIRPassBase<PrintIRPass>::PrintIRPassBase;
+
+  void runOnOperation() override {
+    llvm::dbgs() << "// -----// IR Dump";
+    if (!this->label.empty())
+      llvm::dbgs() << " " << this->label;
+    llvm::dbgs() << " //----- //\n";
+    getOperation()->dump();
+  }
 };
 
 } // namespace
 
-std::unique_ptr<Pass> createPrintIRPass() {
-  return std::make_unique<PrintIRPass>();
+std::unique_ptr<Pass> createPrintIRPass(const PrintIRPassOptions &options) {
+  return std::make_unique<PrintIRPass>(options);
 }
 
 } // namespace mlir
\ No newline at end of file
-- 
2.7.4