[ViewOpGraph] Improve GraphViz output

mlir/lib/Transforms/ViewOpGraph.cpp

@@ -14,6 +14,7 @@
 #include "mlir/IR/Operation.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Support/IndentedOstream.h"
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/Support/Format.h"
 #include "llvm/Support/GraphWriter.h"
 #include <map>
@@ -29,7 +30,7 @@ using namespace mlir;
 
 static const StringRef kLineStyleControlFlow = "dashed";
 static const StringRef kLineStyleDataFlow = "solid";
-static const StringRef kShapeNode = "ellipse";
+static const StringRef kShapeNode = "Mrecord";
 static const StringRef kShapeNone = "plain";
 
 /// Return the size limits for eliding large attributes.
@@ -49,16 +50,25 @@ static std::string strFromOs(function_ref<void(raw_ostream &)> func) {
   return buf;
 }
 
-/// Escape special characters such as '\n' and quotation marks.
-static std::string escapeString(std::string str) {
-  return strFromOs([&](raw_ostream &os) { os.write_escaped(str); });
-}
-
 /// Put quotation marks around a given string.
 static std::string quoteString(const std::string &str) {
   return "\"" + str + "\"";
 }
 
+/// For Graphviz record nodes:
+/// " Braces, vertical bars and angle brackets must be escaped with a backslash
+/// character if you wish them to appear as a literal character "
+std::string escapeLabelString(const std::string &str) {
+  std::string buf;
+  llvm::raw_string_ostream os(buf);
+  for (char c : str) {
+    if (llvm::is_contained({'{', '|', '<', '}', '>', '\n', '"'}, c))
+      os << '\\';
+    os << c;
+  }
+  return buf;
+}
+
 using AttributeMap = std::map<std::string, std::string>;
 
 namespace {
@@ -79,6 +89,12 @@ struct Node {
   std::optional<int> clusterId;
 };
 
+struct DataFlowEdge {
+  Value value;
+  Node node;
+  std::string port;
+};
+
 /// This pass generates a Graphviz dataflow visualization of an MLIR operation.
 /// Note: See https://www.graphviz.org/doc/info/lang.html for more information
 /// about the Graphviz DOT language.
@@ -107,7 +123,7 @@ class PrintOpPass : public impl::ViewOpGraphBase<PrintOpPass> {
 private:
   /// Generate a color mapping that will color every operation with the same
   /// name the same way. It'll interpolate the hue in the HSV color-space,
-  /// attempting to keep the contrast suitable for black text.
+  /// using muted colors that provide good contrast for black text.
   template <typename T>
   void initColorMapping(T &irEntity) {
     backgroundColors.clear();
@@ -120,17 +136,19 @@ class PrintOpPass : public impl::ViewOpGraphBase<PrintOpPass> {
     });
     for (auto indexedOps : llvm::enumerate(ops)) {
       double hue = ((double)indexedOps.index()) / ops.size();
+      // Use lower saturation (0.3) and higher value (0.95) for better
+      // readability
       backgroundColors[indexedOps.value()->getName()].second =
-          std::to_string(hue) + " 1.0 1.0";
+          std::to_string(hue) + " 0.3 0.95";
     }
   }
 
   /// 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 auto &e : dataFlowEdges) {
+        emitEdgeStmt(valueToNode[e.value], e.node, e.port, kLineStyleDataFlow);
       }
     }
 
@@ -147,8 +165,7 @@ class PrintOpPass : public impl::ViewOpGraphBase<PrintOpPass> {
     os.indent();
     // Emit invisible anchor node from/to which arrows can be drawn.
     Node anchorNode = emitNodeStmt(" ", kShapeNone);
-    os << attrStmt("label", quoteString(escapeString(std::move(label))))
-       << ";\n";
+    os << attrStmt("label", quoteString(label)) << ";\n";
     builder();
     os.unindent();
     os << "}\n";
@@ -176,16 +193,17 @@ class PrintOpPass : public impl::ViewOpGraphBase<PrintOpPass> {
 
     // Always emit splat attributes.
     if (isa<SplatElementsAttr>(attr)) {
-      attr.print(os);
+      os << escapeLabelString(
+          strFromOs([&](raw_ostream &os) { attr.print(os); }));
       return;
     }
 
     // Elide "big" elements attributes.
     auto elements = dyn_cast<ElementsAttr>(attr);
     if (elements && elements.getNumElements() > largeAttrLimit) {
       os << std::string(elements.getShapedType().getRank(), '[') << "..."
-         << std::string(elements.getShapedType().getRank(), ']') << " : "
-         << elements.getType();
+         << std::string(elements.getShapedType().getRank(), ']') << " : ";
+      emitMlirType(os, elements.getType());
       return;
     }
 
@@ -199,27 +217,43 @@ class PrintOpPass : public impl::ViewOpGraphBase<PrintOpPass> {
     std::string buf;
     llvm::raw_string_ostream ss(buf);
     attr.print(ss);
-    os << truncateString(buf);
+    os << escapeLabelString(truncateString(buf));
+  }
+
+  // Print a truncated and escaped MLIR type to `os`.
+  void emitMlirType(raw_ostream &os, Type type) {
+    std::string buf;
+    llvm::raw_string_ostream ss(buf);
+    type.print(ss);
+    os << escapeLabelString(truncateString(buf));
+  }
+
+  // Print a truncated and escaped MLIR operand to `os`.
+  void emitMlirOperand(raw_ostream &os, Value operand) {
+    operand.printAsOperand(os, OpPrintingFlags());
   }
 
   /// 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) {
+  void emitEdgeStmt(Node n1, Node n2, std::string port, StringRef style) {
     AttributeMap attrs;
     attrs["style"] = style.str();
-    // Do not label edges that start/end at a cluster boundary. Such edges are
-    // clipped at the boundary, but labels are not. This can lead to labels
-    // floating around without any edge next to them.
-    if (!n1.clusterId && !n2.clusterId)
-      attrs["label"] = quoteString(escapeString(std::move(label)));
     // Use `ltail` and `lhead` to draw edges between clusters.
     if (n1.clusterId)
       attrs["ltail"] = "cluster_" + std::to_string(*n1.clusterId);
     if (n2.clusterId)
       attrs["lhead"] = "cluster_" + std::to_string(*n2.clusterId);
 
     edges.push_back(strFromOs([&](raw_ostream &os) {
-      os << llvm::format("v%i -> v%i ", n1.id, n2.id);
+      os << "v" << n1.id;
+      if (!port.empty() && !n1.clusterId)
+        // Attach edge to south compass point of the result
+        os << ":res" << port << ":s";
+      os << " -> ";
+      os << "v" << n2.id;
+      if (!port.empty() && !n2.clusterId)
+        // Attach edge to north compass point of the operand
+        os << ":arg" << port << ":n";
       emitAttrList(os, attrs);
     }));
   }
@@ -240,20 +274,30 @@ class PrintOpPass : public impl::ViewOpGraphBase<PrintOpPass> {
                     StringRef background = "") {
     int nodeId = ++counter;
     AttributeMap attrs;
-    attrs["label"] = quoteString(escapeString(std::move(label)));
+    attrs["label"] = quoteString(label);
     attrs["shape"] = shape.str();
     if (!background.empty()) {
       attrs["style"] = "filled";
-      attrs["fillcolor"] = ("\"" + background + "\"").str();
+      attrs["fillcolor"] = quoteString(background.str());
     }
     os << llvm::format("v%i ", nodeId);
     emitAttrList(os, attrs);
     os << ";\n";
     return Node(nodeId);
   }
 
-  /// Generate a label for an operation.
-  std::string getLabel(Operation *op) {
+  std::string getValuePortName(Value operand) {
+    // Print value as an operand and omit the leading '%' character.
+    auto str = strFromOs([&](raw_ostream &os) {
+      operand.printAsOperand(os, OpPrintingFlags());
+    });
+    // Replace % and # with _
+    std::replace(str.begin(), str.end(), '%', '_');
+    std::replace(str.begin(), str.end(), '#', '_');
+    return str;
+  }
+
+  std::string getClusterLabel(Operation *op) {
     return strFromOs([&](raw_ostream &os) {
       // Print operation name and type.
       os << op->getName();
@@ -267,18 +311,73 @@ class PrintOpPass : public impl::ViewOpGraphBase<PrintOpPass> {
 
       // Print attributes.
       if (printAttrs) {
-        os << "\n";
+        os << "\\l";
+        for (const NamedAttribute &attr : op->getAttrs()) {
+          os << escapeLabelString(attr.getName().getValue().str()) << ": ";
+          emitMlirAttr(os, attr.getValue());
+          os << "\\l";
+        }
+      }
+    });
+  }
+
+  /// Generate a label for an operation.
+  std::string getRecordLabel(Operation *op) {
+    return strFromOs([&](raw_ostream &os) {
+      os << "{";
+
+      // Print operation inputs.
+      if (op->getNumOperands() > 0) {
+        os << "{";
+        auto operandToPort = [&](Value operand) {
+          os << "<arg" << getValuePortName(operand) << "> ";
+          emitMlirOperand(os, operand);
+        };
+        interleave(op->getOperands(), os, operandToPort, "|");
+        os << "}|";
+      }
+      // Print operation name and type.
+      os << op->getName() << "\\l";
+
+      // Print attributes.
+      if (printAttrs && !op->getAttrs().empty()) {
+        // Extra line break to separate attributes from the operation name.
+        os << "\\l";
         for (const NamedAttribute &attr : op->getAttrs()) {
-          os << '\n' << attr.getName().getValue() << ": ";
+          os << attr.getName().getValue() << ": ";
           emitMlirAttr(os, attr.getValue());
+          os << "\\l";
         }
       }
+
+      if (op->getNumResults() > 0) {
+        os << "|{";
+        auto resultToPort = [&](Value result) {
+          os << "<res" << getValuePortName(result) << "> ";
+          emitMlirOperand(os, result);
+          if (printResultTypes) {
+            os << " ";
+            emitMlirType(os, result.getType());
+          }
+        };
+        interleave(op->getResults(), os, resultToPort, "|");
+        os << "}";
+      }
+
+      os << "}";
     });
   }
 
   /// Generate a label for a block argument.
   std::string getLabel(BlockArgument arg) {
-    return "arg" + std::to_string(arg.getArgNumber());
+    return strFromOs([&](raw_ostream &os) {
+      os << "<res" << getValuePortName(arg) << "> ";
+      arg.printAsOperand(os, OpPrintingFlags());
+      if (printResultTypes) {
+        os << " ";
+        emitMlirType(os, arg.getType());
+      }
+    });
   }
 
   /// Process a block. Emit a cluster and one node per block argument and
@@ -287,14 +386,12 @@ class PrintOpPass : public impl::ViewOpGraphBase<PrintOpPass> {
     emitClusterStmt([&]() {
       for (BlockArgument &blockArg : block.getArguments())
         valueToNode[blockArg] = emitNodeStmt(getLabel(blockArg));
-
       // Emit a node for each operation.
       std::optional<Node> prevNode;
       for (Operation &op : block) {
         Node nextNode = processOperation(&op);
         if (printControlFlowEdges && prevNode)
-          emitEdgeStmt(*prevNode, nextNode, /*label=*/"",
-                       kLineStyleControlFlow);
+          emitEdgeStmt(*prevNode, nextNode, /*port=*/"", kLineStyleControlFlow);
         prevNode = nextNode;
       }
     });
@@ -311,18 +408,19 @@ class PrintOpPass : public impl::ViewOpGraphBase<PrintOpPass> {
             for (Region &region : op->getRegions())
               processRegion(region);
           },
-          getLabel(op));
+          getClusterLabel(op));
     } else {
-      node = emitNodeStmt(getLabel(op), kShapeNode,
+      node = emitNodeStmt(getRecordLabel(op), kShapeNode,
                           backgroundColors[op->getName()].second);
     }
 
     // Insert data flow edges originating from each operand.
     if (printDataFlowEdges) {
       unsigned numOperands = op->getNumOperands();
-      for (unsigned i = 0; i < numOperands; i++)
-        dataFlowEdges.push_back({op->getOperand(i), node,
-                                 numOperands == 1 ? "" : std::to_string(i)});
+      for (unsigned i = 0; i < numOperands; i++) {
+        auto operand = op->getOperand(i);
+        dataFlowEdges.push_back({operand, node, getValuePortName(operand)});
+      }
     }
 
     for (Value result : op->getResults())
@@ -352,7 +450,7 @@ class PrintOpPass : public impl::ViewOpGraphBase<PrintOpPass> {
   /// 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;
+  std::vector<DataFlowEdge> dataFlowEdges;
   /// Counter for generating unique node/subgraph identifiers.
   int counter = 0;
 

mlir/test/Transforms/print-op-graph-back-edges.mlir

@@ -1,21 +1,21 @@
 // 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: }
+//  DFG-NEXT:   compound = true;
+//  DFG-NEXT:   subgraph cluster_1 {
+//  DFG-NEXT:     v2 [label = " ", shape = plain];
+//  DFG-NEXT:     label = "builtin.module : ()\l";
+//  DFG-NEXT:     subgraph cluster_3 {
+//  DFG-NEXT:       v4 [label = " ", shape = plain];
+//  DFG-NEXT:       label = "";
+//  DFG-NEXT:       v5 [fillcolor = "0.000000 0.3 0.95", label = "{{\{\{}}<arg_c0> %c0|<arg_c1> %c1}|arith.addi\l\loverflowFlags: #arith.overflow\<none...\l|{<res_0> %0 index}}", shape = Mrecord, style = filled];
+//  DFG-NEXT:       v6 [fillcolor = "0.333333 0.3 0.95", label = "{arith.constant\l\lvalue: 0 : index\l|{<res_c0> %c0 index}}", shape = Mrecord, style = filled];
+//  DFG-NEXT:       v7 [fillcolor = "0.333333 0.3 0.95", label = "{arith.constant\l\lvalue: 1 : index\l|{<res_c1> %c1 index}}", shape = Mrecord, style = filled];
+//  DFG-NEXT:     }
+//  DFG-NEXT:   }
+//  DFG-NEXT:   v6:res_c0:s -> v5:arg_c0:n[style = solid];
+//  DFG-NEXT:   v7:res_c1:s -> v5:arg_c1:n[style = solid];
+//  DFG-NEXT: }
 
 module {
   %add = arith.addi %c0, %c1 : index