nomnigraph - support subgraph visualization (pytorch#13795)

Summary:
Pull Request resolved: pytorch#13795

Add ability for dot string generation for a single subgraph and python bindings (which is pretty useful for model exploration in Python)
Restructure DotGenerator class a bit to make it easy to implement this feature

Reviewed By: bwasti

Differential Revision: D13010512

fbshipit-source-id: 825665438394b7e6968ab6da167b477af82a7b62

Author: duc0

caffe2/core/nomnigraph/include/nomnigraph/Converters/Dot.h

@@ -1,6 +1,7 @@
 #ifndef NOM_CONVERTERS_DOT_H
 #define NOM_CONVERTERS_DOT_H
 
+#include "nomnigraph/Graph/Algorithms.h"
 #include "nomnigraph/Graph/Graph.h"
 #include "nomnigraph/Support/Casting.h"
 
@@ -10,57 +11,40 @@
 #include <queue>
 #include <sstream>
 #include <unordered_map>
+#include <vector>
 
 namespace nom {
 namespace converters {
 
-template <typename T, typename... U>
+template <typename GraphT>
 class DotGenerator {
  public:
   using NodePrinter = std::function<std::map<std::string, std::string>(
-      typename nom::Graph<T, U...>::NodeRef)>;
+      typename GraphT::NodeRef)>;
   using EdgePrinter = std::function<std::map<std::string, std::string>(
-      typename nom::Graph<T, U...>::EdgeRef)>;
-  using NodeRef = typename nom::Graph<T, U...>::NodeRef;
+      typename GraphT::EdgeRef)>;
 
   static std::map<std::string, std::string> defaultEdgePrinter(
-      typename nom::Graph<T, U...>::EdgeRef e) {
+      typename GraphT::EdgeRef) {
     std::map<std::string, std::string> labelMap;
     return labelMap;
   }
 
-  DotGenerator(typename nom::Graph<T, U...>* g) : g_(g) {}
+  DotGenerator(NodePrinter nodePrinter, EdgePrinter edgePrinter)
+      : nodePrinter_(nodePrinter), edgePrinter_(edgePrinter) {}
 
-  ~DotGenerator() {}
-
-  /**
-   * Converts given graph into DOT string
-   * @param  nodePrinter node attribute extractor
-   * @param  edgePrinter edge attribute extractor
-   * @return             DOT string representation of graph
-   */
-  std::string convert(NodePrinter nodePrinter, EdgePrinter edgePrinter) {
+  // Convert a graph (with optional subgraphs cluster) to dot.
+  std::string convert(
+      const typename GraphT::SubgraphType& sg,
+      const std::vector<typename GraphT::SubgraphType*>& subgraphs) const {
     std::ostringstream output;
     output << "digraph G {\n\
       ";
-    for (const auto& node : g_->getMutableNodes()) {
-      output << (uint64_t)node; // dot doesn't like hex
-      output << "[";
-      for (const auto& attrib : nodePrinter(node)) {
-        output << attrib.first << "=\"" << attrib.second << "\",";
-      }
-      output << "];\n";
-      for (const auto& edge : node->getOutEdges()) {
-        output << (uint64_t)edge->tail() << " -> " << (uint64_t)edge->head();
-        output << "[";
-        for (const auto& attrib : edgePrinter(edge)) {
-          output << attrib.first << "=\"" << attrib.second << "\",";
-        }
-        output << "];\n";
-      }
+    for (const auto& node : sg.getNodes()) {
+      generateNode(node, sg, output);
     }
-    for (auto i = 0; i < subgraphs_.size(); ++i) {
-      const auto& subgraph = subgraphs_[i];
+    for (auto i = 0; i < subgraphs.size(); ++i) {
+      const auto& subgraph = subgraphs[i];
       output << "subgraph cluster" << i << " {\n";
       output << "style=dotted;\n";
       for (const auto& node : subgraph->getNodes()) {
@@ -73,6 +57,18 @@ class DotGenerator {
     return output.str();
   }
 
+  // Convert a subgraph to dot.
+  std::string convert(const typename GraphT::SubgraphType& sg) const {
+    std::ostringstream output;
+    output << "digraph G {\n\
+      ";
+    for (const auto& node : sg.getNodes()) {
+      generateNode(node, sg, output);
+    }
+    output << "}";
+    return output.str();
+  }
+
   /**
    * NOTE No subgraph support
    * Converts given graph into DOT string w/operator input-order preserved
@@ -84,19 +80,20 @@ class DotGenerator {
    *     - Node: op_ptr[shape=record, label="{{<i0>*|<i1>*|...}|{op}|{<o0>*}"]
    *     - Edge: <parent_node_ptr>:<ref>:s -> <this_node_ptr>:<ref>:n
    */
-  std::string convertStruct(NodePrinter nodePrinter, EdgePrinter edgePrinter) {
+  std::string convertStruct(const typename GraphT::SubgraphType& sg) const {
     std::ostringstream output;
     output << "digraph G {\n";
 
     // Get input nodes (nodes w/o parents)
-    std::unordered_map<NodeRef, int> nodeDepthMap; // Touched nodes for BFS
-    std::queue<NodeRef> workList; // Init w/parentless nodes
-    for (const auto& node : g_->getMutableNodes()) {
+    std::unordered_map<typename GraphT::NodeRef, int>
+        nodeDepthMap; // Touched nodes for BFS
+    std::queue<typename GraphT::NodeRef> workList; // Init w/parentless nodes
+    for (const auto& node : sg.getNodes()) {
       if (node->getInEdges().size() == 0 && node->getOutEdges().size() > 0) {
         // Add input node to dot string
         output << (uint64_t)node << "[shape=record, label=\"{{Data In}|{<"
                << (uint64_t)node << ">";
-        for (const auto& attr : nodePrinter(node)) {
+        for (const auto& attr : nodePrinter_(node)) {
           output << attr.second;
         }
         output << "}}\"]\n";
@@ -108,7 +105,7 @@ class DotGenerator {
     }
 
     // BFS to get operator nodes
-    std::vector<NodeRef> ops;
+    std::vector<typename GraphT::NodeRef> ops;
     while (workList.size() > 0) {
       const auto& node = workList.front();
       for (const auto& edge : node->getOutEdges()) {
@@ -127,25 +124,21 @@ class DotGenerator {
     }
 
     // Finalize output
-    output << getOperatorSubtreeDotString(ops, nodePrinter) << "}\n";
+    output << getOperatorSubtreeDotString(ops) << "}\n";
     return output.str();
   }
 
-  void addSubgraph(const nom::Subgraph<T, U...>* s) {
-    subgraphs_.emplace_back(s);
-  }
-
  private:
-  typename nom::Graph<T, U...>* g_;
-  typename std::vector<const nom::Subgraph<T, U...>*> subgraphs_;
+  NodePrinter nodePrinter_;
+  EdgePrinter edgePrinter_;
 
   /**
    * Get DOT string record of given operator and DOT string of its input edges
    * @param  op          operator to parse
    * @param  nodePrinter node attribute extractor
    * @return             '\n' sep string of operator & input edges
    */
-  std::string getOperatorDotString(NodeRef op, NodePrinter nodePrinter) {
+  std::string getOperatorDotString(typename GraphT::NodeRef op) const {
     std::ostringstream output;
     std::ostringstream record; // Operator node record
     record << (uint64_t)op << "[shape=record, label=\"{{";
@@ -167,15 +160,15 @@ class DotGenerator {
 
       // Add input to operator record
       record << sep << "<" << (uint64_t)input << ">";
-      for (const auto& attr : nodePrinter(input)) {
+      for (const auto& attr : nodePrinter_(input)) {
         record << attr.second;
       }
       sep = "|";
     }
 
     // Extract operator name
     record << "}|{";
-    for (const auto& attr : nodePrinter(op)) {
+    for (const auto& attr : nodePrinter_(op)) {
       record << attr.second;
     }
     record << "}|{";
@@ -185,7 +178,7 @@ class DotGenerator {
     for (const auto& edge : op->getOutEdges()) {
       const auto& child = edge->head();
       record << sep << "<" << (uint64_t)child << ">";
-      for (const auto& attr : nodePrinter(child)) {
+      for (const auto& attr : nodePrinter_(child)) {
         record << attr.second;
       }
       sep = "|";
@@ -203,48 +196,81 @@ class DotGenerator {
    * @return             DOT string that renders operators subgraph
    */
   std::string getOperatorSubtreeDotString(
-      std::vector<NodeRef> ops,
-      NodePrinter nodePrinter) {
+      std::vector<typename GraphT::NodeRef> ops) const {
     std::ostringstream output;
     for (const auto& op : ops) {
-      output << getOperatorDotString(op, nodePrinter);
+      output << getOperatorDotString(op);
     }
     return output.str();
   }
+
+  // Generate dot string for a node.
+  void generateNode(
+      typename GraphT::NodeRef node,
+      const typename GraphT::SubgraphType& sg,
+      std::ostringstream& output) const {
+    output << (uint64_t)node; // dot doesn't like hex
+    output << "[";
+    for (const auto& attrib : nodePrinter_(node)) {
+      output << attrib.first << "=\"" << attrib.second << "\",";
+    }
+    output << "];\n";
+    for (const auto& edge : node->getOutEdges()) {
+      if (!sg.hasEdge(edge)) {
+        continue;
+      }
+      output << (uint64_t)edge->tail() << " -> " << (uint64_t)edge->head();
+      output << "[";
+      for (const auto& attrib : edgePrinter_(edge)) {
+        output << attrib.first << "=\"" << attrib.second << "\",";
+      }
+      output << "];\n";
+    }
+  }
 };
 
-template <typename T, typename... U>
+// Convert a graph to dot string.
+template <typename GraphT>
 std::string convertToDotString(
-    nom::Graph<T, U...>* g,
-    typename DotGenerator<T, U...>::NodePrinter nodePrinter,
-    typename DotGenerator<T, U...>::EdgePrinter edgePrinter =
-        DotGenerator<T, U...>::defaultEdgePrinter) {
-  auto d = DotGenerator<T, U...>(g);
-  return d.convert(nodePrinter, edgePrinter);
+    GraphT* g,
+    typename DotGenerator<GraphT>::NodePrinter nodePrinter,
+    typename DotGenerator<GraphT>::EdgePrinter edgePrinter =
+        DotGenerator<GraphT>::defaultEdgePrinter) {
+  auto d = DotGenerator<GraphT>(nodePrinter, edgePrinter);
+  return d.convert(algorithm::createSubgraph(g), {});
 }
 
-template <typename T, typename... U>
+// Convert a graph to dot string and annotate subgraph clusters.
+template <typename GraphT>
 std::string convertToDotString(
-    nom::Graph<T, U...>* g,
-    const std::vector<nom::Subgraph<T, U...>>& subgraphs,
-    typename DotGenerator<T, U...>::NodePrinter nodePrinter,
-    typename DotGenerator<T, U...>::EdgePrinter edgePrinter =
-        DotGenerator<T, U...>::defaultEdgePrinter) {
-  auto d = DotGenerator<T, U...>(g);
-  for (const auto& subgraph : subgraphs) {
-    d.addSubgraph(&subgraph);
-  }
-  return d.convert(nodePrinter, edgePrinter);
+    GraphT* g,
+    const std::vector<typename GraphT::SubgraphType*>& subgraphs,
+    typename DotGenerator<GraphT>::NodePrinter nodePrinter,
+    typename DotGenerator<GraphT>::EdgePrinter edgePrinter =
+        DotGenerator<GraphT>::defaultEdgePrinter) {
+  auto d = DotGenerator<GraphT>(nodePrinter, edgePrinter);
+  return d.convert(algorithm::createSubgraph(g), subgraphs);
+}
+
+// Convert a subgraph to dot string.
+template <typename GraphT>
+std::string convertToDotString(
+    const typename GraphT::SubgraphType& sg,
+    typename DotGenerator<GraphT>::NodePrinter nodePrinter,
+    typename DotGenerator<GraphT>::EdgePrinter edgePrinter =
+        DotGenerator<GraphT>::defaultEdgePrinter) {
+  auto d = DotGenerator<GraphT>(nodePrinter, edgePrinter);
+  return d.convert(sg);
 }
 
-template <typename T, typename... U>
+template <typename GraphT>
 std::string convertToDotRecordString(
-    nom::Graph<T, U...>* g,
-    typename DotGenerator<T, U...>::NodePrinter nodePrinter,
-    typename DotGenerator<T, U...>::EdgePrinter edgePrinter =
-        DotGenerator<T, U...>::defaultEdgePrinter) {
-  auto d = DotGenerator<T, U...>(g);
-  return d.convertStruct(nodePrinter, edgePrinter);
+    GraphT* g,
+    typename DotGenerator<GraphT>::NodePrinter nodePrinter,
+    typename DotGenerator<GraphT>::EdgePrinter edgePrinter =
+        DotGenerator<GraphT>::defaultEdgePrinter) {
+  auto d = DotGenerator<GraphT>(nodePrinter, edgePrinter);
+  return d.convertStruct(algorithm::createSubgraph(g));
 }
 
 } // namespace converters

caffe2/core/nomnigraph/include/nomnigraph/Graph/Algorithms.h

@@ -196,6 +196,17 @@ void induceEdges(SubgraphType* sg) {
   }
 }
 
+/// \brief Create subgraph object from graph.
+template <typename GraphType>
+typename GraphType::SubgraphType createSubgraph(GraphType* g) {
+  typename GraphType::SubgraphType subgraph;
+  for (auto& node : g->getMutableNodes()) {
+    subgraph.addNode(node);
+  }
+  induceEdges(&subgraph);
+  return subgraph;
+}
+
 } // namespace algorithm
 } // namespace nom
 

caffe2/core/nomnigraph/tests/GraphTest.cc

@@ -176,3 +176,44 @@ TEST(Basic, MoveSubgraph) {
   EXPECT_EQ(g.getMutableEdges().size(), 0);
   EXPECT_EQ(g2.getMutableEdges().size(), 1);
 }
+
+TEST(Basic, DotGenerator) {
+  TestGraph g;
+  auto n1 = createTestNode(g);
+  auto n2 = createTestNode(g);
+  auto n3 = createTestNode(g);
+  auto e12 = g.createEdge(n1, n2);
+  g.createEdge(n1, n3);
+
+  std::string dot = nom::converters::convertToDotString(&g, TestNodePrinter);
+
+  // sanity check
+  std::string prefix = "digraph G";
+  // Full string comparison of the output is not stable because the dot
+  // string includes node pointer address as node id. We should switch to
+  // comparing full output once dot generator no longer uses addresses.
+  EXPECT_TRUE(dot.compare(0, prefix.length(), prefix) == 0);
+
+  TestGraph::SubgraphType sg;
+  sg.addNode(n1);
+  sg.addNode(n2);
+  sg.addEdge(e12);
+
+  // Convert to dot with subgraph clusters.
+  dot = nom::converters::convertToDotString<TestGraph>(
+      &g, {&sg}, TestNodePrinter);
+
+  // sanity check
+  EXPECT_TRUE(dot.compare(0, prefix.length(), prefix) == 0);
+
+  // Convert a single subgraph to dot.
+  dot = nom::converters::convertToDotString<TestGraph>(sg, TestNodePrinter);
+
+  // sanity check
+  EXPECT_TRUE(dot.compare(0, prefix.length(), prefix) == 0);
+
+  dot =
+      nom::converters::convertToDotRecordString<TestGraph>(&g, TestNodePrinter);
+  // sanity check
+  EXPECT_TRUE(dot.compare(0, prefix.length(), prefix) == 0);
+}

caffe2/core/nomnigraph/tests/test_util.cc

@@ -120,3 +120,10 @@ std::map<std::string, std::string> NNPrinter(typename nom::repr::NNGraph::NodeRe
 nom::Graph<TestClass>::NodeRef createTestNode(nom::Graph<TestClass>& g) {
   return g.createNode(TestClass());
 }
+
+std::map<std::string, std::string> TestNodePrinter(
+    nom::Graph<TestClass>::NodeRef /* unused */) {
+  std::map<std::string, std::string> labelMap;
+  labelMap["label"] = "Node";
+  return labelMap;
+}

caffe2/core/nomnigraph/tests/test_util.h

@@ -114,4 +114,7 @@ std::map<std::string, std::string> NNPrinter(typename nom::repr::NNGraph::NodeRe
 
 CAFFE2_API nom::Graph<TestClass>::NodeRef createTestNode(
     nom::Graph<TestClass>& g);
+
+CAFFE2_API std::map<std::string, std::string> TestNodePrinter(
+    nom::Graph<TestClass>::NodeRef node);
 #endif // NOM_TESTS_TEST_UTIL_H