[mlir] Handle cycles and back edges in --view-op-graph (#82002)

Fixes #62128.

3 files changed, 86 insertions, 3 deletions

diff --git a/mlir/lib/Transforms/ViewOpGraph.cpp b/mlir/lib/Transforms/ViewOpGraph.cpp
index 3d2723839957..c2eb2b893cea 100644
--- a/mlir/lib/Transforms/ViewOpGraph.cpp
+++ b/mlir/lib/Transforms/ViewOpGraph.cpp
@@ -13,6 +13,7 @@
 #include "mlir/IR/Operation.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Support/IndentedOstream.h"
+#include "mlir/Transforms/TopologicalSortUtils.h"
 #include "llvm/Support/Format.h"
 #include "llvm/Support/GraphWriter.h"
 #include <map>
@@ -126,6 +127,12 @@ private:
   /// Emit all edges. This function should be called after all nodes have been
   /// emitted.
   void emitAllEdgeStmts() {
+    if (printDataFlowEdges) {
+      for (const auto &[value, node, label] : dataFlowEdges) {
+        emitEdgeStmt(valueToNode[value], node, label, kLineStyleDataFlow);
+      }
+    }
+
     for (const std::string &edge : edges)
       os << edge << ";\n";
     edges.clear();
@@ -313,9 +320,8 @@ private:
     if (printDataFlowEdges) {
       unsigned numOperands = op->getNumOperands();
       for (unsigned i = 0; i < numOperands; i++)
-        emitEdgeStmt(valueToNode[op->getOperand(i)], node,
-                     /*label=*/numOperands == 1 ? "" : std::to_string(i),
-                     kLineStyleDataFlow);
+        dataFlowEdges.push_back({op->getOperand(i), node,
+                                 numOperands == 1 ? "" : std::to_string(i)});
     }
 
     for (Value result : op->getResults())
@@ -344,6 +350,8 @@ private:
   std::vector<std::string> edges;
   /// Mapping of SSA values to Graphviz nodes/clusters.
   DenseMap<Value, Node> valueToNode;
+  /// Output for data flow edges is delayed until the end to handle cycles
+  std::vector<std::tuple<Value, Node, std::string>> dataFlowEdges;
   /// Counter for generating unique node/subgraph identifiers.
   int counter = 0;
 
diff --git a/mlir/test/Transforms/print-op-graph-back-edges.mlir b/mlir/test/Transforms/print-op-graph-back-edges.mlir
new file mode 100644
index 000000000000..ed922dd7cb13
--- /dev/null
+++ b/mlir/test/Transforms/print-op-graph-back-edges.mlir
@@ -0,0 +1,24 @@
+// RUN: mlir-opt -view-op-graph %s -o %t 2>&1 | FileCheck -check-prefix=DFG %s
+
+// DFG-LABEL: digraph G {
+//       DFG:   compound = true;
+//       DFG:   subgraph cluster_1 {
+//       DFG:     v2 [label = " ", shape = plain];
+//       DFG:     label = "builtin.module : ()\n";
+//       DFG:     subgraph cluster_3 {
+//       DFG:       v4 [label = " ", shape = plain];
+//       DFG:       label = "";
+//       DFG:       v5 [fillcolor = "0.000000 1.0 1.0", label = "arith.addi : (index)\n\noverflowFlags: #arith.overflow<none...", shape = ellipse, style = filled];
+//       DFG:       v6 [fillcolor = "0.333333 1.0 1.0", label = "arith.constant : (index)\n\nvalue: 0 : index", shape = ellipse, style = filled];
+//       DFG:       v7 [fillcolor = "0.333333 1.0 1.0", label = "arith.constant : (index)\n\nvalue: 1 : index", shape = ellipse, style = filled];
+//       DFG:     }
+//       DFG:   }
+//       DFG:   v6 -> v5 [label = "0", style = solid];
+//       DFG:   v7 -> v5 [label = "1", style = solid];
+//       DFG: }
+
+module {
+  %add = arith.addi %c0, %c1 : index
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+}
diff --git a/mlir/test/Transforms/print-op-graph-cycles.mlir b/mlir/test/Transforms/print-op-graph-cycles.mlir
new file mode 100644
index 000000000000..7e4eb5616a28
--- /dev/null
+++ b/mlir/test/Transforms/print-op-graph-cycles.mlir
@@ -0,0 +1,51 @@
+// RUN: mlir-opt -view-op-graph -allow-unregistered-dialect %s -o %t 2>&1 | FileCheck -check-prefix=DFG %s
+
+// DFG-LABEL: digraph G {
+//       DFG:   compound = true;
+//       DFG:   subgraph cluster_1 {
+//       DFG:     v2 [label = " ", shape = plain];
+//       DFG:     label = "builtin.module : ()\n";
+//       DFG:     subgraph cluster_3 {
+//       DFG:       v4 [label = " ", shape = plain];
+//       DFG:       label = "";
+//       DFG:       subgraph cluster_5 {
+//       DFG:         v6 [label = " ", shape = plain];
+//       DFG:         label = "test.graph_region : ()\n";
+//       DFG:         subgraph cluster_7 {
+//       DFG:           v8 [label = " ", shape = plain];
+//       DFG:           label = "";
+//       DFG:           v9 [fillcolor = "0.000000 1.0 1.0", label = "op1 : (i32)\n", shape = ellipse, style = filled];
+//       DFG:           subgraph cluster_10 {
+//       DFG:             v11 [label = " ", shape = plain];
+//       DFG:             label = "test.ssacfg_region : (i32)\n";
+//       DFG:             subgraph cluster_12 {
+//       DFG:               v13 [label = " ", shape = plain];
+//       DFG:               label = "";
+//       DFG:               v14 [fillcolor = "0.166667 1.0 1.0", label = "op2 : (i32)\n", shape = ellipse, style = filled];
+//       DFG:             }
+//       DFG:           }
+//       DFG:           v15 [fillcolor = "0.166667 1.0 1.0", label = "op2 : (i32)\n", shape = ellipse, style = filled];
+//       DFG:           v16 [fillcolor = "0.500000 1.0 1.0", label = "op3 : (i32)\n", shape = ellipse, style = filled];
+//       DFG:         }
+//       DFG:       }
+//       DFG:     }
+//       DFG:   }
+//       DFG:   v9 -> v9 [label = "0", style = solid];
+//       DFG:   v15 -> v9 [label = "1", style = solid];
+//       DFG:   v9 -> v14 [label = "0", style = solid];
+//       DFG:   v11 -> v14 [ltail = cluster_10, style = solid];
+//       DFG:   v15 -> v14 [label = "2", style = solid];
+//       DFG:   v16 -> v14 [label = "3", style = solid];
+//       DFG:   v9 -> v15 [label = "0", style = solid];
+//       DFG:   v16 -> v15 [label = "1", style = solid];
+//       DFG:   v9 -> v16 [label = "", style = solid];
+//       DFG: }
+
+"test.graph_region"() ({ // A Graph region
+  %1 = "op1"(%1, %3) : (i32, i32) -> (i32)  // OK: %1, %3 allowed here
+  %2 = "test.ssacfg_region"() ({
+     %5 = "op2"(%1, %2, %3, %4) : (i32, i32, i32, i32) -> (i32) // OK: %1, %2, %3, %4 all defined in the containing region
+  }) : () -> (i32)
+  %3 = "op2"(%1, %4) : (i32, i32) -> (i32)  // OK: %4 allowed here
+  %4 = "op3"(%1) : (i32) -> (i32)
+}) : () -> ()