Improve MLIR "view-op-graph" to color operations according to their name

[jpienaar]

I think reasonable step.

mlir/lib/Transforms/ViewOpGraph.cpp

119

Later we could enable a pre populate or some such too - TF side I often only had a few ops I cared about, and as discussed in chat there was a lambda that could be populated from python side and so I could for example only focus on some nodes.

mlir/lib/Transforms/ViewOpGraph.cpp

 /// Note: See https://www.graphviz.org/doc/info/lang.html for more information
 /// about the Graphviz DOT language.
 class PrintOpPass : public impl::ViewOpGraphBase<PrintOpPass> {
 public:
   PrintOpPass(raw_ostream &os) : os(os) {}
   PrintOpPass(const PrintOpPass &o) : PrintOpPass(o.os.getOStream()) {}
 
   void runOnOperation() override {
+    initColorMapping(*getOperation());
     emitGraph([&]() {
       processOperation(getOperation());
       emitAllEdgeStmts();
     });
   }
 
   /// Create a CFG graph for a region. Used in `Region::viewGraph`.
   void emitRegionCFG(Region &region) {
     printControlFlowEdges = true;
     printDataFlowEdges = false;
+    initColorMapping(region);
     emitGraph([&]() { processRegion(region); });
   }
 
 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.
+  template <typename T>
+  void initColorMapping(T &irEntity) {
+    backgroundColors.clear();
+    SmallVector<Operation *> ops;
+    irEntity.walk([&](Operation *op) {
+      auto &entry = backgroundColors[op->getName()];
+      if (entry.first == 0)
+        ops.push_back(op);
+      ++entry.first;
+    });
+    for (auto indexedOps : llvm::enumerate(ops)) {
+      double hue = ((double)indexedOps.index()) / ops.size();
+      backgroundColors[indexedOps.value()->getName()].second =
+          std::to_string(hue) + " 1.0 1.0";
+    }
+  }
+
   /// Emit all edges. This function should be called after all nodes have been
   /// emitted.
   void emitAllEdgeStmts() {
     for (const std::string &edge : edges)
       os << edge << ";\n";
     edges.clear();
   }
 
void emitGraph(function_ref<void()> builder) {
     // Edges between clusters are allowed only in compound mode.
     os << attrStmt("compound", "true") << ";\n";
     builder();
     os.unindent();
     os << "}\n";
   }
 
   /// Emit a node statement.
-  Node emitNodeStmt(std::string label, StringRef shape = kShapeNode) {
+  Node emitNodeStmt(std::string label, StringRef shape = kShapeNode,
+                    StringRef background = "") {
     int nodeId = ++counter;
     AttributeMap attrs;
     attrs["label"] = quoteString(escapeString(std::move(label)));
     attrs["shape"] = shape.str();
+    if (!background.empty()) {
+      attrs["style"] = "filled";
+      attrs["fillcolor"] = ("\"" + 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) {
if (op->getNumRegions() > 0) {
       // Emit cluster for op with regions.
       node = emitClusterStmt(
           [&]() {
             for (Region &region : op->getRegions())
               processRegion(region);
           },
           getLabel(op));
     } else {
-      node = emitNodeStmt(getLabel(op));
+      node = emitNodeStmt(getLabel(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++)
         emitEdgeStmt(valueToNode[op->getOperand(i)], node,
                      /*label=*/numOperands == 1 ? "" : std::to_string(i),
private:
   raw_indented_ostream os;
   /// A list of edges. For simplicity, should be emitted after all nodes were
   /// emitted.
   std::vector<std::string> edges;
   /// Mapping of SSA values to Graphviz nodes/clusters.
   DenseMap<Value, Node> valueToNode;
   /// Counter for generating unique node/subgraph identifiers.
   int counter = 0;
+
+  DenseMap<OperationName, std::pair<int, std::string>> backgroundColors;
 };
 
 } // namespace
 
 std::unique_ptr<Pass> mlir::createPrintOpGraphPass(raw_ostream &os) {
   return std::make_unique<PrintOpPass>(os);
 }