Merge pull request #88 from NervanaSystems/bob/graphviz

add cmake finder for graphviz dot program and then conditionally comp…

cmake/Modules/FindGraphviz.cmake

+find_program(GRAPHVIZ_EXECUTABLE dot)
+
+# Handle REQUIRED and QUIET arguments
+# this will also set GRAPHVIZ_FOUND to true if GRAPHVIZ_EXECUTABLE exists
+include(FindPackageHandleStandardArgs)
+find_package_handle_standard_args(Graphviz
+                     "Failed to locate graphviz executable"
+                     GRAPHVIZ_EXECUTABLE)

src/CMakeLists.txt

@@ -38,15 +38,22 @@ set(NGRAPH_INCLUDE_PATH
     ${CMAKE_CURRENT_SOURCE_DIR}/ngraph
 )
 
+# find_program (GRAPHVIZ dot)
+# message (STATUS "graphviz '${GRAPHVIZ}'")
+find_package(Graphviz)
+if (GRAPHVIZ_FOUND)
+    set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DGRAPHVIZ_FOUND")
+endif()
+
 include_directories("${NGRAPH_INCLUDE_PATH}")
 
 add_library(ngraph SHARED ${SRC})
 target_include_directories(ngraph PUBLIC "${NGRAPH_INCLUDE_PATH}")
 
-if ( APPLE )
-    set_property( TARGET ngraph PROPERTY PREFIX "lib" )
-    set_property( TARGET ngraph PROPERTY OUTPUT_NAME "ngraph.so" )
-    set_property( TARGET ngraph PROPERTY SUFFIX "" )
+if (APPLE)
+    set_property(TARGET ngraph PROPERTY PREFIX "lib")
+    set_property(TARGET ngraph PROPERTY OUTPUT_NAME "ngraph.so")
+    set_property(TARGET ngraph PROPERTY SUFFIX "")
 endif()
 
 #-----------------------------------------------------------------------------------------------
@@ -65,7 +72,7 @@ set(CMAKE_INSTALL_PREFIX "$ENV{HOME}/ngraph_dist" CACHE PATH "Install directory"
 message (STATUS "Installation directory: ${CMAKE_INSTALL_PREFIX}")
 message (STATUS "To Override use: cmake -DCMAKE_INSTALL_PREFIX=/foo -P cmake_install.cmake")
 
-install(TARGETS ngraph DESTINATION ${CMAKE_INSTALL_PREFIX} )
+install(TARGETS ngraph DESTINATION ${CMAKE_INSTALL_PREFIX})
 install(DIRECTORY 
     ${CMAKE_CURRENT_SOURCE_DIR}/ngraph 
     DESTINATION ${CMAKE_INSTALL_PREFIX} 

src/ngraph/visualize.cpp

@@ -34,13 +34,46 @@ void Visualize::add(node_ptr p)
     traverse_nodes(p, [&](node_ptr node) {
         for (auto arg : node->get_arguments())
         {
-            m_ss << "    " << arg->get_node_id() << " -> " << node->get_node_id() << ";\n";
+            m_ss << add_attributes(arg);
+            m_ss << add_attributes(node);
+            m_ss << "    " << arg->get_node_id() << " -> " << node->get_node_id();
+            m_ss << ";\n";
         }
     });
 }
 
+std::string Visualize::add_attributes(node_ptr node)
+{
+    string rc;
+    if (!contains(m_nodes_with_attributes, node))
+    {
+        m_nodes_with_attributes.insert(node);
+        rc = get_attributes(node);
+    }
+    return rc;
+}
+
+std::string Visualize::get_attributes(node_ptr node)
+{
+    stringstream ss;
+    if (node->is_parameter())
+    {
+        ss << "    " << node->get_node_id() << " [shape=box color=blue]\n";
+    }
+    else if (node->is_op())
+    {
+        ss << "    " << node->get_node_id() << " [shape=ellipse color=black]\n";
+    }
+    else
+    {
+        ss << "    " << node->get_node_id() << " [shape=diamond color=red]\n";
+    }
+    return ss.str();
+}
+
 void Visualize::save_dot(const string& path) const
 {
+#if GRAPHVIZ_FOUND
     auto     tmp_file = path + ".tmp";
     ofstream out(tmp_file);
     if (out)
@@ -56,6 +89,8 @@ void Visualize::save_dot(const string& path) const
         auto stream = popen(cmd.c_str(), "r");
         pclose(stream);
 
-        // remove(tmp_file.c_str());
+        remove(tmp_file.c_str());
     }
+#else
+#endif
 }

src/ngraph/visualize.hpp

@@ -36,6 +36,10 @@ public:
     void save_dot(const std::string& path) const;
 
 private:
-    std::stringstream m_ss;
-    std::string       m_name;
+    std::string add_attributes(node_ptr node);
+    std::string get_attributes(node_ptr node);
+
+    std::stringstream  m_ss;
+    std::string        m_name;
+    std::set<node_ptr> m_nodes_with_attributes;
 };