[mlir] Support drawing control-flow graphs in ViewOpGraph.cpp

* Add new pass option `print-data-flow-edges`, default value `true`.
* Add new pass option `print-control-flow-edges`, default value `false`.
* Remove `PrintCFGPass`. Same functionality now provided by
  `PrintOpPass`.

Differential Revision: https://reviews.llvm.org/D106342

Author: matthias-springer

mlir/include/mlir/Transforms/Passes.h

@@ -17,7 +17,6 @@
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/LocationSnapshot.h"
 #include "mlir/Transforms/ViewOpGraph.h"
-#include "mlir/Transforms/ViewRegionGraph.h"
 #include <limits>
 
 namespace mlir {

mlir/include/mlir/Transforms/Passes.td

@@ -612,11 +612,6 @@ def ParallelLoopCollapsing : Pass<"parallel-loop-collapsing"> {
   ];
 }
 
-def PrintCFG : FunctionPass<"print-cfg-graph"> {
-  let summary = "Print CFG graph per-Region";
-  let constructor = "mlir::createPrintCFGGraphPass()";
-}
-
 def PrintOpStats : Pass<"print-op-stats"> {
   let summary = "Print statistics of operations";
   let constructor = "mlir::createPrintOpStatsPass()";
@@ -689,14 +684,17 @@ def SymbolDCE : Pass<"symbol-dce"> {
 }
 
 def ViewOpGraphPass : Pass<"view-op-graph"> {
-  let summary = "Print Graphviz dataflow visualization of an operation";
+  let summary = "Print Graphviz visualization of an operation";
   let description = [{
-    This pass prints a Graphviz dataflow graph of a module.
+    This pass prints a Graphviz graph of a module.
 
     - Operations are represented as nodes;
-    - Uses as edges;
+    - Uses (data flow) as edges;
+    - Control flow as dashed edges;
     - Regions/blocks as subgraphs.
 
+    By default, only data flow edges are printed.
+
     Note: See https://www.graphviz.org/doc/info/lang.html for more information
     about the Graphviz DOT language.
   }];
@@ -705,6 +703,10 @@ def ViewOpGraphPass : Pass<"view-op-graph"> {
             /*default=*/"20", "Limit attribute/type length to number of chars">,
     Option<"printAttrs", "print-attrs", "bool",
            /*default=*/"true", "Print attributes of operations">,
+    Option<"printControlFlowEdges", "print-control-flow-edges", "bool",
+           /*default=*/"false", "Print control flow edges">,
+    Option<"printDataFlowEdges", "print-data-flow-edges", "bool",
+           /*default=*/"true", "Print data flow edges">,
     Option<"printResultTypes", "print-result-types", "bool",
             /*default=*/"true", "Print result types of operations">
   ];

mlir/include/mlir/Transforms/ViewRegionGraph.h

@@ -21,7 +21,6 @@ add_mlir_library(MLIRTransforms
   StripDebugInfo.cpp
   SymbolDCE.cpp
   ViewOpGraph.cpp
-  ViewRegionGraph.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Transforms

mlir/lib/Transforms/CMakeLists.txt

@@ -17,6 +17,7 @@
 #include "mlir/Transforms/Passes.h"
 #include "mlir/Transforms/Utils.h"
 #include "llvm/ADT/SmallSet.h"
+#include "llvm/Support/Debug.h"
 
 #define DEBUG_TYPE "normalize-memrefs"
 

mlir/lib/Transforms/NormalizeMemRefs.cpp

@@ -12,9 +12,11 @@
 #include "mlir/IR/Operation.h"
 #include "mlir/Support/IndentedOstream.h"
 #include "llvm/Support/Format.h"
+#include "llvm/Support/GraphWriter.h"
 
 using namespace mlir;
 
+static const StringRef kLineStyleControlFlow = "dashed";
 static const StringRef kLineStyleDataFlow = "solid";
 static const StringRef kShapeNode = "ellipse";
 static const StringRef kShapeNone = "plain";
@@ -78,6 +80,13 @@ class PrintOpPass : public ViewOpGraphPassBase<PrintOpPass> {
     });
   }
 
+  /// Create a CFG graph for a region. Used in `Region::viewGraph`.
+  void emitRegionCFG(Region &region) {
+    printControlFlowEdges = true;
+    printDataFlowEdges = false;
+    emitGraph([&]() { processRegion(region); });
+  }
+
 private:
   /// Emit all edges. This function should be called after all nodes have been
   /// emitted.
@@ -151,8 +160,7 @@ class PrintOpPass : public ViewOpGraphPassBase<PrintOpPass> {
 
   /// Append an edge to the list of edges.
   /// Note: Edges are written to the output stream via `emitAllEdgeStmts`.
-  void emitEdgeStmt(Node n1, Node n2, std::string label,
-                    StringRef style = kLineStyleDataFlow) {
+  void emitEdgeStmt(Node n1, Node n2, std::string label, StringRef style) {
     AttributeMap attrs;
     attrs["style"] = style.str();
     // Do not label edges that start/end at a cluster boundary. Such edges are
@@ -233,14 +241,20 @@ class PrintOpPass : public ViewOpGraphPassBase<PrintOpPass> {
         valueToNode[blockArg] = emitNodeStmt(getLabel(blockArg));
 
       // Emit a node for each operation.
-      for (Operation &op : block)
-        processOperation(&op);
+      Optional<Node> prevNode;
+      for (Operation &op : block) {
+        Node nextNode = processOperation(&op);
+        if (printControlFlowEdges && prevNode)
+          emitEdgeStmt(*prevNode, nextNode, /*label=*/"",
+                       kLineStyleControlFlow);
+        prevNode = nextNode;
+      }
     });
   }
 
   /// Process an operation. If the operation has regions, emit a cluster.
   /// Otherwise, emit a node.
-  void processOperation(Operation *op) {
+  Node processOperation(Operation *op) {
     Node node;
     if (op->getNumRegions() > 0) {
       // Emit cluster for op with regions.
@@ -254,14 +268,19 @@ class PrintOpPass : public ViewOpGraphPassBase<PrintOpPass> {
       node = emitNodeStmt(getLabel(op));
     }
 
-    // Insert edges originating from each operand.
-    unsigned numOperands = op->getNumOperands();
-    for (unsigned i = 0; i < numOperands; i++)
-      emitEdgeStmt(valueToNode[op->getOperand(i)], node,
-                   /*label=*/numOperands == 1 ? "" : std::to_string(i));
+    // Insert data flow edges originating from each operand.
+    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);
+    }
 
     for (Value result : op->getResults())
       valueToNode[result] = node;
+
+    return node;
   }
 
   /// Process a region.
@@ -294,3 +313,25 @@ std::unique_ptr<Pass>
 mlir::createPrintOpGraphPass(raw_ostream &os) {
   return std::make_unique<PrintOpPass>(os);
 }
+
+/// Generate a CFG for a region and show it in a window.
+static void llvmViewGraph(Region &region, const Twine &name) {
+  int fd;
+  std::string filename = llvm::createGraphFilename(name.str(), fd);
+  {
+    llvm::raw_fd_ostream os(fd, /*shouldClose=*/true);
+    if (fd == -1) {
+      llvm::errs() << "error opening file '" << filename << "' for writing\n";
+      return;
+    }
+    PrintOpPass pass(os);
+    pass.emitRegionCFG(region);
+  }
+  llvm::DisplayGraph(filename, /*wait=*/false, llvm::GraphProgram::DOT);
+}
+
+void mlir::Region::viewGraph(const Twine &regionName) {
+  llvmViewGraph(*this, regionName);
+}
+
+void mlir::Region::viewGraph() { viewGraph("region"); }

mlir/lib/Transforms/ViewOpGraph.cpp

@@ -1,26 +1,55 @@
-// RUN: mlir-opt -allow-unregistered-dialect -mlir-elide-elementsattrs-if-larger=2 -view-op-graph %s -o %t 2>&1 | FileCheck %s
+// RUN: mlir-opt -allow-unregistered-dialect -mlir-elide-elementsattrs-if-larger=2 -view-op-graph %s -o %t 2>&1 | FileCheck -check-prefix=DFG %s
+// RUN: mlir-opt -allow-unregistered-dialect -mlir-elide-elementsattrs-if-larger=2 -view-op-graph='print-data-flow-edges=false print-control-flow-edges=true' %s -o %t 2>&1 | FileCheck -check-prefix=CFG %s
+
+// DFG-LABEL: digraph G {
+//       DFG:   subgraph {{.*}} {
+//       DFG:     subgraph {{.*}}
+//       DFG:       label = "builtin.func{{.*}}merge_blocks
+//       DFG:       subgraph {{.*}} {
+//       DFG:         v[[ARG0:.*]] [label = "arg0"
+//       DFG:         v[[CONST10:.*]] [label ={{.*}}10 : i32
+//       DFG:         subgraph [[CLUSTER_MERGE_BLOCKS:.*]] {
+//       DFG:           v[[ANCHOR:.*]] [label = " ", shape = plain]
+//       DFG:           label = "test.merge_blocks
+//       DFG:           subgraph {{.*}} {
+//       DFG:             v[[TEST_BR:.*]] [label = "test.br
+//       DFG:           }
+//       DFG:           subgraph {{.*}} {
+//       DFG:           }
+//       DFG:         }
+//       DFG:         v[[TEST_RET:.*]] [label = "test.return
+//       DFG:   v[[ARG0]] -> v[[TEST_BR]]
+//       DFG:   v[[CONST10]] -> v[[TEST_BR]]
+//       DFG:   v[[ANCHOR]] -> v[[TEST_RET]] [{{.*}}, ltail = [[CLUSTER_MERGE_BLOCKS]]]
+//       DFG:   v[[ANCHOR]] -> v[[TEST_RET]] [{{.*}}, ltail = [[CLUSTER_MERGE_BLOCKS]]]
+
+// CFG-LABEL: digraph G {
+//       CFG:   subgraph {{.*}} {
+//       CFG:     subgraph {{.*}}
+//       CFG:       label = "builtin.func{{.*}}merge_blocks
+//       CFG:       subgraph {{.*}} {
+//       CFG:         v[[C1:.*]] [label = "std.constant
+//       CFG:         v[[C2:.*]] [label = "std.constant
+//       CFG:         v[[C3:.*]] [label = "std.constant
+//       CFG:         v[[C4:.*]] [label = "std.constant
+//       CFG:         v[[TEST_FUNC:.*]] [label = "test.func
+//       CFG:         subgraph [[CLUSTER_MERGE_BLOCKS:.*]] {
+//       CFG:           v[[ANCHOR:.*]] [label = " ", shape = plain]
+//       CFG:           label = "test.merge_blocks
+//       CFG:           subgraph {{.*}} {
+//       CFG:             v[[TEST_BR:.*]] [label = "test.br
+//       CFG:           }
+//       CFG:           subgraph {{.*}} {
+//       CFG:           }
+//       CFG:         }
+//       CFG:         v[[TEST_RET:.*]] [label = "test.return
+//       CFG:   v[[C1]] -> v[[C2]]
+//       CFG:   v[[C2]] -> v[[C3]]
+//       CFG:   v[[C3]] -> v[[C4]]
+//       CFG:   v[[C4]] -> v[[TEST_FUNC]]
+//       CFG:   v[[TEST_FUNC]] -> v[[ANCHOR]] [{{.*}}, lhead = [[CLUSTER_MERGE_BLOCKS]]]
+//       CFG:   v[[ANCHOR]] -> v[[TEST_RET]] [{{.*}}, ltail = [[CLUSTER_MERGE_BLOCKS]]]
 
-// CHECK-LABEL: digraph G {
-//       CHECK:   subgraph {{.*}} {
-//       CHECK:     subgraph {{.*}}
-//       CHECK:       label = "builtin.func{{.*}}merge_blocks
-//       CHECK:       subgraph {{.*}} {
-//       CHECK:         v[[ARG0:.*]] [label = "arg0"
-//       CHECK:         v[[CONST10:.*]] [label ={{.*}}10 : i32
-//       CHECK:         subgraph [[CLUSTER_MERGE_BLOCKS:.*]] {
-//       CHECK:           v[[ANCHOR:.*]] [label = " ", shape = plain]
-//       CHECK:           label = "test.merge_blocks
-//       CHECK:           subgraph {{.*}} {
-//       CHECK:             v[[TEST_BR:.*]] [label = "test.br
-//       CHECK:           }
-//       CHECK:           subgraph {{.*}} {
-//       CHECK:           }
-//       CHECK:         }
-//       CHECK:         v[[TEST_RET:.*]] [label = "test.return
-//       CHECK:   v[[ARG0]] -> v[[TEST_BR]]
-//       CHECK:   v[[CONST10]] -> v[[TEST_BR]]
-//       CHECK:   v[[ANCHOR]] -> v[[TEST_RET]] [{{.*}}, ltail = [[CLUSTER_MERGE_BLOCKS]]]
-//       CHECK:   v[[ANCHOR]] -> v[[TEST_RET]] [{{.*}}, ltail = [[CLUSTER_MERGE_BLOCKS]]]
 func @merge_blocks(%arg0: i32, %arg1 : i32) -> () {
   %0 = constant dense<[[0, 1], [2, 3]]> : tensor<2x2xi32>
   %1 = constant dense<1> : tensor<5xi32>