From 1d6a90418e4bfc294b4174880e93cb43835ebdf5 Mon Sep 17 00:00:00 2001
From: Stella Laurenzo <stellaraccident@gmail.com>
Date: Fri, 15 Jul 2022 21:03:04 -0700
Subject: [PATCH] Example:

```
// -----// IR Dump Before LowerLinalgMicrokernels (iree-vmvx-lower-linalg-microkernels) //----- //
```

I've been meaning to suggest this for a long time, and I think the only reason we don't have it is because we didn't used to have the `getArgument()` handy when printing these comments. When debugging or putting a pipeline together based on such dumps, I often find myself grepping for the argument name of the pass (which is often related but not universally).
---
 mlir/lib/Pass/IRPrinting.cpp       |  9 ++++---
 mlir/test/Pass/ir-printing.mlir    | 38 +++++++++++++++---------------
 mlir/test/Pass/run-reproducer.mlir |  6 ++---
 3 files changed, 28 insertions(+), 25 deletions(-)

diff --git a/mlir/lib/Pass/IRPrinting.cpp b/mlir/lib/Pass/IRPrinting.cpp
index c17963c51112..c20d9b1e9135 100644
--- a/mlir/lib/Pass/IRPrinting.cpp
+++ b/mlir/lib/Pass/IRPrinting.cpp
@@ -126,7 +126,8 @@ void IRPrinterInstrumentation::runBeforePass(Pass *pass, Operation *op) {
     beforePassFingerPrints.try_emplace(pass, op);
 
   config->printBeforeIfEnabled(pass, op, [&](raw_ostream &out) {
-    out << "// -----// IR Dump Before " << pass->getName();
+    out << "// -----// IR Dump Before " << pass->getName() << " ("
+        << pass->getArgument() << ")";
     printIR(op, config->shouldPrintAtModuleScope(), out,
             config->getOpPrintingFlags());
     out << "\n\n";
@@ -156,7 +157,8 @@ void IRPrinterInstrumentation::runAfterPass(Pass *pass, Operation *op) {
   }
 
   config->printAfterIfEnabled(pass, op, [&](raw_ostream &out) {
-    out << "// -----// IR Dump After " << pass->getName();
+    out << "// -----// IR Dump After " << pass->getName() << " ("
+        << pass->getArgument() << ")";
     printIR(op, config->shouldPrintAtModuleScope(), out,
             config->getOpPrintingFlags());
     out << "\n\n";
@@ -170,7 +172,8 @@ void IRPrinterInstrumentation::runAfterPassFailed(Pass *pass, Operation *op) {
     beforePassFingerPrints.erase(pass);
 
   config->printAfterIfEnabled(pass, op, [&](raw_ostream &out) {
-    out << formatv("// -----// IR Dump After {0} Failed", pass->getName());
+    out << formatv("// -----// IR Dump After {0} Failed ({1})", pass->getName(),
+                   pass->getArgument());
     printIR(op, config->shouldPrintAtModuleScope(), out, OpPrintingFlags());
     out << "\n\n";
   });
diff --git a/mlir/test/Pass/ir-printing.mlir b/mlir/test/Pass/ir-printing.mlir
index 8c9303d5da74..bd506c2c7ebd 100644
--- a/mlir/test/Pass/ir-printing.mlir
+++ b/mlir/test/Pass/ir-printing.mlir
@@ -15,53 +15,53 @@ func.func @bar() {
   return
 }
 
-// BEFORE: // -----// IR Dump Before{{.*}}CSE //----- //
+// BEFORE: // -----// IR Dump Before{{.*}}CSE (cse) //----- //
 // BEFORE-NEXT: func @foo()
-// BEFORE: // -----// IR Dump Before{{.*}}CSE //----- //
+// BEFORE: // -----// IR Dump Before{{.*}}CSE (cse) //----- //
 // BEFORE-NEXT: func @bar()
-// BEFORE-NOT: // -----// IR Dump Before{{.*}}Canonicalizer //----- //
+// BEFORE-NOT: // -----// IR Dump Before{{.*}}Canonicalizer (canonicalize) //----- //
 // BEFORE-NOT: // -----// IR Dump After
 
-// BEFORE_ALL: // -----// IR Dump Before{{.*}}CSE //----- //
+// BEFORE_ALL: // -----// IR Dump Before{{.*}}CSE (cse) //----- //
 // BEFORE_ALL-NEXT: func @foo()
-// BEFORE_ALL: // -----// IR Dump Before{{.*}}Canonicalizer //----- //
+// BEFORE_ALL: // -----// IR Dump Before{{.*}}Canonicalizer (canonicalize) //----- //
 // BEFORE_ALL-NEXT: func @foo()
-// BEFORE_ALL: // -----// IR Dump Before{{.*}}CSE //----- //
+// BEFORE_ALL: // -----// IR Dump Before{{.*}}CSE (cse) //----- //
 // BEFORE_ALL-NEXT: func @bar()
-// BEFORE_ALL: // -----// IR Dump Before{{.*}}Canonicalizer //----- //
+// BEFORE_ALL: // -----// IR Dump Before{{.*}}Canonicalizer (canonicalize) //----- //
 // BEFORE_ALL-NEXT: func @bar()
 // BEFORE_ALL-NOT: // -----// IR Dump After
 
 // AFTER-NOT: // -----// IR Dump Before
-// AFTER: // -----// IR Dump After{{.*}}CSE //----- //
+// AFTER: // -----// IR Dump After{{.*}}CSE (cse) //----- //
 // AFTER-NEXT: func @foo()
-// AFTER: // -----// IR Dump After{{.*}}CSE //----- //
+// AFTER: // -----// IR Dump After{{.*}}CSE (cse) //----- //
 // AFTER-NEXT: func @bar()
-// AFTER-NOT: // -----// IR Dump After{{.*}}Canonicalizer //----- //
+// AFTER-NOT: // -----// IR Dump After{{.*}}Canonicalizer (canonicalize) //----- //
 
 // AFTER_ALL-NOT: // -----// IR Dump Before
-// AFTER_ALL: // -----// IR Dump After{{.*}}CSE //----- //
+// AFTER_ALL: // -----// IR Dump After{{.*}}CSE (cse) //----- //
 // AFTER_ALL-NEXT: func @foo()
-// AFTER_ALL: // -----// IR Dump After{{.*}}Canonicalizer //----- //
+// AFTER_ALL: // -----// IR Dump After{{.*}}Canonicalizer (canonicalize) //----- //
 // AFTER_ALL-NEXT: func @foo()
-// AFTER_ALL: // -----// IR Dump After{{.*}}CSE //----- //
+// AFTER_ALL: // -----// IR Dump After{{.*}}CSE (cse) //----- //
 // AFTER_ALL-NEXT: func @bar()
-// AFTER_ALL: // -----// IR Dump After{{.*}}Canonicalizer //----- //
+// AFTER_ALL: // -----// IR Dump After{{.*}}Canonicalizer (canonicalize) //----- //
 // AFTER_ALL-NEXT: func @bar()
 
-// BEFORE_MODULE: // -----// IR Dump Before{{.*}}CSE ('func.func' operation: @foo) //----- //
+// BEFORE_MODULE: // -----// IR Dump Before{{.*}}CSE (cse) ('func.func' operation: @foo) //----- //
 // BEFORE_MODULE: func @foo()
 // BEFORE_MODULE: func @bar()
-// BEFORE_MODULE: // -----// IR Dump Before{{.*}}CSE ('func.func' operation: @bar) //----- //
+// BEFORE_MODULE: // -----// IR Dump Before{{.*}}CSE (cse) ('func.func' operation: @bar) //----- //
 // BEFORE_MODULE: func @foo()
 // BEFORE_MODULE: func @bar()
 
-// AFTER_ALL_CHANGE: // -----// IR Dump After{{.*}}CSE //----- //
+// AFTER_ALL_CHANGE: // -----// IR Dump After{{.*}}CSE (cse) //----- //
 // AFTER_ALL_CHANGE-NEXT: func @foo()
-// AFTER_ALL_CHANGE-NOT: // -----// IR Dump After{{.*}}CSE //----- //
+// AFTER_ALL_CHANGE-NOT: // -----// IR Dump After{{.*}}CSE (cse) //----- //
 // We expect that only 'foo' changed during CSE, and the second run of CSE did
 // nothing.
 
 // AFTER_FAILURE-NOT: // -----// IR Dump After{{.*}}CSE
-// AFTER_FAILURE: // -----// IR Dump After{{.*}}TestFailurePass Failed //----- //
+// AFTER_FAILURE: // -----// IR Dump After{{.*}}TestFailurePass Failed (test-pass-failure) //----- //
 // AFTER_FAILURE: func @foo()
diff --git a/mlir/test/Pass/run-reproducer.mlir b/mlir/test/Pass/run-reproducer.mlir
index c4a2f06801d5..6627033d8be7 100644
--- a/mlir/test/Pass/run-reproducer.mlir
+++ b/mlir/test/Pass/run-reproducer.mlir
@@ -18,9 +18,9 @@ func.func @bar() {
   }
 #-}
 
-// BEFORE: // -----// IR Dump Before{{.*}}CSE //----- //
+// BEFORE: // -----// IR Dump Before{{.*}}CSE (cse) //----- //
 // BEFORE-NEXT: func @foo()
-// BEFORE: // -----// IR Dump Before{{.*}}CSE //----- //
+// BEFORE: // -----// IR Dump Before{{.*}}CSE (cse) //----- //
 // BEFORE-NEXT: func @bar()
-// BEFORE-NOT: // -----// IR Dump Before{{.*}}Canonicalizer //----- //
+// BEFORE-NOT: // -----// IR Dump Before{{.*}}Canonicalizer (canonicalize) //----- //
 // BEFORE-NOT: // -----// IR Dump After