[Mlir-commits] [mlir] [ViewOpGraph] Re-use `AsmState` to significantly speed-up printing. (PR #208240)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Philipp Schilk (schilkp)

<details>
<summary>Changes</summary>

I was trying to print some rather large functions (5000+ ops) to a graphivz/dot file, and noticed that the `ViewOpGraph` pass rather slow.

`--view-op-graph` previously always called `Operation::printAsOperand` with the default `OpPrintingFlags` causing a new `AsmState` to be constructed every time, which is extremely inefficient.

This patch instead creates a top-level `AsmState` once for each pass run, and re-uses it for all printing.

This makes printing graphviz/dot files for large graphs *significantly* faster: A single `func.func` with 5000 operations previously took 33 seconds to print, and can now be printed (essentially) instantly :)

---

quick benchmark:

gen.bash:
```bash
#!/bin/bash
# Usage: ./gen.bash <num_ops> <output_file>
num_ops=$1
output_file=$2
{
    echo "func.func @<!-- -->core(%arg0: i32) -> i32 {"
    echo "  %c1 = arith.constant 1 : i32"
    echo "  %0 = arith.addi %arg0, %c1 : i32"
    for ((i=1; i<num_ops; i++)); do
        echo "  %$i = arith.addi %$((i-1)), %c1 : i32"
    done
    echo "  return %$((num_ops-1)) : i32"
    echo "}"
} > "$output_file"
```


```bash
> gen.bash 5000 test_5000.mlir
> # Before:
> time ./build/bin/mlir-opt --view-op-graph ./test_5000.mlir 2> /dev/null -o /dev/null
./build/bin/mlir-opt --view-op-graph ./test_5000.mlir -o out.mlir 2> out.dot  35.15s user 0.06s system 99% cpu 35.279 total
> # After:
> time ./build/bin/mlir-opt --view-op-graph ./test_5000.mlir 2> out.dot -o out.mlir
./build/bin/mlir-opt --view-op-graph ./test_5000.mlir -o out.mlir 2> out.dot  0.03s user 0.04s system 99% cpu 0.073 total
```


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


1 Files Affected:

- (modified) mlir/lib/Transforms/ViewOpGraph.cpp (+12-5) 


``````````diff
diff --git a/mlir/lib/Transforms/ViewOpGraph.cpp b/mlir/lib/Transforms/ViewOpGraph.cpp
index 2d7e40d18efca..26923984ccdf3 100644
--- a/mlir/lib/Transforms/ViewOpGraph.cpp
+++ b/mlir/lib/Transforms/ViewOpGraph.cpp
@@ -8,6 +8,7 @@
 
 #include "mlir/Transforms/ViewOpGraph.h"
 
+#include "mlir/IR/AsmState.h"
 #include "mlir/IR/Block.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Operation.h"
@@ -108,10 +109,12 @@ class PrintOpPass : public impl::ViewOpGraphPassBase<PrintOpPass> {
 
   void runOnOperation() override {
     initColorMapping(*getOperation());
+    asmState.emplace(getOperation(), OpPrintingFlags());
     emitGraph([&]() {
       processOperation(getOperation());
       emitAllEdgeStmts();
     });
+    asmState.reset();
     markAllAnalysesPreserved();
   }
 
@@ -120,7 +123,9 @@ class PrintOpPass : public impl::ViewOpGraphPassBase<PrintOpPass> {
     printControlFlowEdges = true;
     printDataFlowEdges = false;
     initColorMapping(region);
+    asmState.emplace(region.getParentOp(), OpPrintingFlags());
     emitGraph([&]() { processRegion(region); });
+    asmState.reset();
   }
 
 private:
@@ -234,7 +239,7 @@ class PrintOpPass : public impl::ViewOpGraphPassBase<PrintOpPass> {
 
   // Print a truncated and escaped MLIR operand to `os`.
   void emitMlirOperand(raw_ostream &os, Value operand) {
-    operand.printAsOperand(os, OpPrintingFlags());
+    operand.printAsOperand(os, *asmState);
   }
 
   /// Append an edge to the list of edges.
@@ -292,9 +297,8 @@ class PrintOpPass : public impl::ViewOpGraphPassBase<PrintOpPass> {
 
   std::string getValuePortName(Value operand) {
     // Print value as an operand and omit the leading '%' character.
-    auto str = strFromOs([&](raw_ostream &os) {
-      operand.printAsOperand(os, OpPrintingFlags());
-    });
+    auto str = strFromOs(
+        [&](raw_ostream &os) { operand.printAsOperand(os, *asmState); });
     // Replace % and # with _
     llvm::replace(str, '%', '_');
     llvm::replace(str, '#', '_');
@@ -376,7 +380,7 @@ class PrintOpPass : public impl::ViewOpGraphPassBase<PrintOpPass> {
   std::string getLabel(BlockArgument arg) {
     return strFromOs([&](raw_ostream &os) {
       os << "<res" << getValuePortName(arg) << "> ";
-      arg.printAsOperand(os, OpPrintingFlags());
+      arg.printAsOperand(os, *asmState);
       if (printResultTypes) {
         os << " ";
         emitMlirType(os, arg.getType());
@@ -448,6 +452,9 @@ class PrintOpPass : public impl::ViewOpGraphPassBase<PrintOpPass> {
 
   /// Output stream to write DOT file to.
   raw_indented_ostream os;
+  /// Re-usable assembly printer state for efficient printing. Initialized
+  /// on each pass run.
+  std::optional<AsmState> asmState;
   /// A list of edges. For simplicity, should be emitted after all nodes were
   /// emitted.
   std::vector<std::string> edges;

``````````

</details>


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