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

[Artem Tyurin]

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

>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 1/2] [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
+}

>From 96abe87d71acff42c34b43b4bf86c4d849b7cb2e Mon Sep 17 00:00:00 2001
From: Artem Tyurin <artem.tyurin at gmail.com>
Date: Sat, 17 Feb 2024 00:24:24 +0100
Subject: [PATCH 2/2] Do not modify existing blocks

---
 mlir/lib/Transforms/ViewOpGraph.cpp | 12 ++++++++----
 1 file changed, 8 insertions(+), 4 deletions(-)

diff --git a/mlir/lib/Transforms/ViewOpGraph.cpp b/mlir/lib/Transforms/ViewOpGraph.cpp
index 485f93aeb9c3c4..8ebd8e631bb57f 100644
--- a/mlir/lib/Transforms/ViewOpGraph.cpp
+++ b/mlir/lib/Transforms/ViewOpGraph.cpp
@@ -277,16 +277,20 @@ 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));
 
+      SmallVector<Operation*> sortedOperations;
+      for (Operation &op : block) {
+        sortedOperations.push_back(&op);
+      }
+      computeTopologicalSorting(sortedOperations);
+
       // Emit a node for each operation.
       std::optional<Node> prevNode;
-      for (Operation &op : block) {
-        Node nextNode = processOperation(&op);
+      for (Operation *op : sortedOperations) {
+        Node nextNode = processOperation(op);
         if (printControlFlowEdges && prevNode)
           emitEdgeStmt(*prevNode, nextNode, /*label=*/"",
                        kLineStyleControlFlow);