diff --git a/mlir/include/mlir/Transforms/Passes.td b/mlir/include/mlir/Transforms/Passes.td --- a/mlir/include/mlir/Transforms/Passes.td +++ b/mlir/include/mlir/Transforms/Passes.td @@ -695,6 +695,7 @@ - Operations are represented as nodes; - Uses as edges; + - Control flow as dashed edges; - Regions/blocks as subgraphs. Note: See https://www.graphviz.org/doc/info/lang.html for more information @@ -705,6 +706,10 @@ /*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"> ]; diff --git a/mlir/lib/Transforms/ViewOpGraph.cpp b/mlir/lib/Transforms/ViewOpGraph.cpp --- a/mlir/lib/Transforms/ViewOpGraph.cpp +++ b/mlir/lib/Transforms/ViewOpGraph.cpp @@ -15,6 +15,7 @@ using namespace mlir; +static const StringRef kLineStyleControlFlow = "dashed"; static const StringRef kLineStyleDataFlow = "solid"; static const StringRef kShapeNode = "ellipse"; static const StringRef kShapeNone = "plain"; @@ -151,8 +152,7 @@ /// 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 +233,20 @@ valueToNode[blockArg] = emitNodeStmt(getLabel(blockArg)); // Emit a node for each operation. - for (Operation &op : block) - processOperation(&op); + Optional 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 +260,19 @@ 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. diff --git a/mlir/test/Transforms/print-op-graph.mlir b/mlir/test/Transforms/print-op-graph.mlir --- a/mlir/test/Transforms/print-op-graph.mlir +++ b/mlir/test/Transforms/print-op-graph.mlir @@ -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>