[Mlir-commits] [mlir] [mlir] Handle backedges in --view-op-graph (PR #82002)

[Artem Tyurin]

https://github.com/agentcooper created https://github.com/llvm/llvm-project/pull/82002

Fixes #62128.

Before:

![before](https://github.com/llvm/llvm-project/assets/794591/531b47f2-0638-47d1-864f-9dbed06e176b)

After:

![after](https://github.com/llvm/llvm-project/assets/794591/8a8f438a-7c8b-4b9c-a7af-40971bda8bfd)


>From 8dc6364a6e68d032d161bf2f9316edf1adf03e41 Mon Sep 17 00:00:00 2001
From: Artem Tyurin <artem.tyurin at gmail.com>
Date: Fri, 16 Feb 2024 16:42:42 +0100
Subject: [PATCH] [mlir] Handle backedges in --view-op-graph

---
 mlir/lib/Transforms/ViewOpGraph.cpp           |  3 +++
 .../Transforms/print-op-graph-backedges.mlir  | 24 +++++++++++++++++++
 2 files changed, 27 insertions(+)
 create mode 100644 mlir/test/Transforms/print-op-graph-backedges.mlir

diff --git a/mlir/lib/Transforms/ViewOpGraph.cpp b/mlir/lib/Transforms/ViewOpGraph.cpp
index 3d2723839957c4..485f93aeb9c3c4 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>
@@ -276,6 +277,8 @@ class PrintOpPass : public impl::ViewOpGraphBase<PrintOpPass> {
   /// Process a block. Emit a cluster and one node per block argument and
   /// operation inside the cluster.
   void processBlock(Block &block) {
+    sortTopologically(&block);
+
     emitClusterStmt([&]() {
       for (BlockArgument &blockArg : block.getArguments())
         valueToNode[blockArg] = emitNodeStmt(getLabel(blockArg));
diff --git a/mlir/test/Transforms/print-op-graph-backedges.mlir b/mlir/test/Transforms/print-op-graph-backedges.mlir
new file mode 100644
index 00000000000000..cd03bd0c1c298f
--- /dev/null
+++ b/mlir/test/Transforms/print-op-graph-backedges.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.333333 1.0 1.0", label = "arith.constant : (index)\n\nvalue: 0 : index", shape = ellipse, style = filled];
+//       DFG:       v6 [fillcolor = "0.333333 1.0 1.0", label = "arith.constant : (index)\n\nvalue: 1 : index", shape = ellipse, style = filled];
+//       DFG:       v7 [fillcolor = "0.000000 1.0 1.0", label = "arith.addi : (index)\n\noverflowFlags: #arith.overflow<none...", shape = ellipse, style = filled];
+//       DFG:     }
+//       DFG:   }
+//       DFG:   v5 -> v7 [label = "0", style = solid];
+//       DFG:   v6 -> v7 [label = "1", style = solid];
+//       DFG: }
+
+module {
+  %add = arith.addi %c0, %c1 : index
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+}