[Mlir-commits] [mlir] [mlir] Add IR printing flag to skip regions (PR #77726)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-core

Author: Jakub Kuderski (kuhar)

<details>
<summary>Changes</summary>

The new flag, `--mlir-print-ir-skip-regions`, sets the op printing option that disables region printing. This results in the usual `--mlir-print-ir-*` debug options printing only the names of the executed passes and the signatures of the ops.

Example:
```mlir
// -----// IR Dump Before CSE (cse) //----- //
func.func @<!-- -->bar(%arg0: f32, %arg1: f32) -> f32 {...}

// -----// IR Dump Before Canonicalizer (canonicalize) //----- //
func.func @<!-- -->bar(%arg0: f32, %arg1: f32) -> f32 {...}
```

The main use-case is to be triage compilation issues (crashes, slowness) on very deep pass pipelines and with very large IR files, where printing IR is prohibitively slow otherwise.

Unlike the pass statistics mechanism, the new flag can be used as an 'interactive' progress indicator.

---
Full diff: https://github.com/llvm/llvm-project/pull/77726.diff


3 Files Affected:

- (modified) mlir/docs/PassManagement.md (+21) 
- (modified) mlir/lib/Pass/PassManagerOptions.cpp (+11-1) 
- (modified) mlir/test/Pass/ir-printing.mlir (+11) 


``````````diff
diff --git a/mlir/docs/PassManagement.md b/mlir/docs/PassManagement.md
index 95a38207b7f854..a303e350d28b6b 100644
--- a/mlir/docs/PassManagement.md
+++ b/mlir/docs/PassManagement.md
@@ -1302,6 +1302,27 @@ func.func @simple_constant() -> (i32, i32) {
 }
 ```
 
+*   `mlir-print-ir-skip-regions`
+    *   Do not print op regions.
+    *   This option is useful to when we are interested in the names of the
+        executed passes but do not want to inspect the IR, e.g., when the
+        module is very large and printing is prohibitively slow.
+
+```shell
+$ mlir-opt foo.mlir -mlir-disable-threading -pass-pipeline='func.func(canonicalize,cse)' -mlir-print-ir-after-all -mlir-print-ir-skip-regions
+// -----// IR Dump Before CSE (cse) //----- //
+func.func @bar(%arg0: f32, %arg1: f32) -> f32 {...}
+
+// -----// IR Dump Before Canonicalizer (canonicalize) //----- //
+func.func @bar(%arg0: f32, %arg1: f32) -> f32 {...}
+
+// -----// IR Dump Before CSE (cse) //----- //
+func.func @simple_constant() -> (i32, i32) {...}
+
+// -----// IR Dump Before Canonicalizer (canonicalize) //----- //
+func.func @simple_constant() -> (i32, i32) {...}
+```
+
 ## Crash and Failure Reproduction
 
 The [pass manager](#pass-manager) in MLIR contains a builtin mechanism to
diff --git a/mlir/lib/Pass/PassManagerOptions.cpp b/mlir/lib/Pass/PassManagerOptions.cpp
index ffc53b7e3ed023..da7f27fb75a29c 100644
--- a/mlir/lib/Pass/PassManagerOptions.cpp
+++ b/mlir/lib/Pass/PassManagerOptions.cpp
@@ -58,6 +58,12 @@ struct PassManagerOptions {
       llvm::cl::desc("When printing IR for print-ir-[before|after]{-all} "
                      "always print the top-level operation"),
       llvm::cl::init(false)};
+  llvm::cl::opt<bool> printSkipRegions{
+      "mlir-print-ir-skip-regions",
+      llvm::cl::desc(
+          "When printing IR for print-ir-[beforef|after]{-all} do not output "
+          "regions. This is useful to indicate pass pipeline progress."),
+      llvm::cl::init(false)};
 
   /// Add an IR printing instrumentation if enabled by any 'print-ir' flags.
   void addPrinterInstrumentation(PassManager &pm);
@@ -119,10 +125,14 @@ void PassManagerOptions::addPrinterInstrumentation(PassManager &pm) {
   if (!shouldPrintBeforePass && !shouldPrintAfterPass)
     return;
 
+  OpPrintingFlags opPrintingFlags;
+  if (printSkipRegions)
+    opPrintingFlags.skipRegions();
+
   // Otherwise, add the IR printing instrumentation.
   pm.enableIRPrinting(shouldPrintBeforePass, shouldPrintAfterPass,
                       printModuleScope, printAfterChange, printAfterFailure,
-                      llvm::errs());
+                      llvm::errs(), opPrintingFlags);
 }
 
 void mlir::registerPassManagerCLOptions() {
diff --git a/mlir/test/Pass/ir-printing.mlir b/mlir/test/Pass/ir-printing.mlir
index 048b721ba6d530..cc2f277d76ee52 100644
--- a/mlir/test/Pass/ir-printing.mlir
+++ b/mlir/test/Pass/ir-printing.mlir
@@ -3,6 +3,7 @@
 // RUN: mlir-opt %s -mlir-disable-threading=true -pass-pipeline='builtin.module(func.func(cse,canonicalize))' -mlir-print-ir-after=cse -o /dev/null 2>&1 | FileCheck -check-prefix=AFTER %s
 // RUN: mlir-opt %s -mlir-disable-threading=true -pass-pipeline='builtin.module(func.func(cse,canonicalize))' -mlir-print-ir-after-all -o /dev/null 2>&1 | FileCheck -check-prefix=AFTER_ALL %s
 // RUN: mlir-opt %s -mlir-disable-threading=true -pass-pipeline='builtin.module(func.func(cse,canonicalize))' -mlir-print-ir-before=cse -mlir-print-ir-module-scope -o /dev/null 2>&1 | FileCheck -check-prefix=BEFORE_MODULE %s
+// RUN: mlir-opt %s -mlir-disable-threading=true -pass-pipeline='builtin.module(func.func(cse,canonicalize))' -mlir-print-ir-before-all -mlir-print-ir-skip-regions -o /dev/null 2>&1 | FileCheck -check-prefix=BEFORE_ALL_SKIP_REGIONS %s
 // RUN: mlir-opt %s -mlir-disable-threading=true -pass-pipeline='builtin.module(func.func(cse,cse))' -mlir-print-ir-after-all -mlir-print-ir-after-change -o /dev/null 2>&1 | FileCheck -check-prefix=AFTER_ALL_CHANGE %s
 // RUN: not mlir-opt %s -mlir-disable-threading=true -pass-pipeline='builtin.module(func.func(cse,test-pass-failure))' -mlir-print-ir-after-failure -o /dev/null 2>&1 | FileCheck -check-prefix=AFTER_FAILURE %s
 
@@ -56,6 +57,16 @@ func.func @bar() {
 // BEFORE_MODULE: func @foo()
 // BEFORE_MODULE: func @bar()
 
+// BEFORE_ALL_SKIP_REGIONS: // -----// IR Dump Before{{.*}}CSE (cse) //----- //
+// BEFORE_ALL_SKIP_REGIONS-NEXT: func @foo() {...}
+// BEFORE_ALL_SKIP_REGIONS: // -----// IR Dump Before{{.*}}Canonicalizer (canonicalize) //----- //
+// BEFORE_ALL_SKIP_REGIONS-NEXT: func @foo() {...}
+
+// BEFORE_ALL_SKIP_REGIONS: // -----// IR Dump Before{{.*}}CSE (cse) //----- //
+// BEFORE_ALL_SKIP_REGIONS-NEXT: func @bar() {...}
+// BEFORE_ALL_SKIP_REGIONS: // -----// IR Dump Before{{.*}}Canonicalizer (canonicalize) //----- //
+// BEFORE_ALL_SKIP_REGIONS-NEXT: func @bar() {...}
+
 // AFTER_ALL_CHANGE: // -----// IR Dump After{{.*}}CSE (cse) //----- //
 // AFTER_ALL_CHANGE-NEXT: func @foo()
 // AFTER_ALL_CHANGE-NOT: // -----// IR Dump After{{.*}}CSE (cse) //----- //

``````````

</details>


https://github.com/llvm/llvm-project/pull/77726