add option to nbench to generate only a dot file (#2781)

src/ngraph/pass/visualize_tree.cpp

@@ -65,9 +65,10 @@ bool pass::VisualizeTree::run_on_module(vector<shared_ptr<Function>>& functions)
     return false;
 }
 
-pass::VisualizeTree::VisualizeTree(const string& file_name, node_modifiers_t nm)
+pass::VisualizeTree::VisualizeTree(const string& file_name, node_modifiers_t nm, bool dot_only)
     : m_name{file_name}
     , m_node_modifiers{nm}
+    , m_dot_only(dot_only)
 {
 }
 
@@ -177,15 +178,18 @@ void pass::VisualizeTree::render() const
         out << "}\n";
         out.close();
 
-#ifndef _WIN32
-        stringstream ss;
-        ss << "dot -T" << get_file_ext() << " " << dot_file << " -o " << m_name;
-        auto cmd = ss.str();
-        auto stream = popen(cmd.c_str(), "r");
-        if (stream)
+        if (!m_dot_only)
         {
-            pclose(stream);
-        }
+#ifndef _WIN32
+            stringstream ss;
+            ss << "dot -T" << get_file_ext() << " " << dot_file << " -o " << m_name;
+            auto cmd = ss.str();
+            auto stream = popen(cmd.c_str(), "r");
+            if (stream)
+            {
+                pclose(stream);
+            }
 #endif
+        }
     }
 }

src/ngraph/pass/visualize_tree.hpp

@@ -41,7 +41,9 @@ class ngraph::pass::VisualizeTree : public ModulePass
 public:
     using node_modifiers_t =
         std::function<void(const Node& node, std::vector<std::string>& attributes)>;
-    VisualizeTree(const std::string& file_name, node_modifiers_t nm = nullptr);
+    VisualizeTree(const std::string& file_name,
+                  node_modifiers_t nm = nullptr,
+                  bool dot_only = false);
     bool run_on_module(std::vector<std::shared_ptr<ngraph::Function>>&) override;
 
     static std::string get_file_ext();
@@ -57,4 +59,5 @@ private:
     std::unordered_map<std::type_index, std::function<void(const Node&, std::ostream& ss)>>
         m_ops_to_details;
     node_modifiers_t m_node_modifiers = nullptr;
+    bool m_dot_only;
 };

src/tools/nbench/nbench.cpp

@@ -197,6 +197,7 @@ int main(int argc, char** argv)
     bool visualize = false;
     int warmup_iterations = 1;
     bool copy_data = true;
+    bool dot_file = false;
 
     for (size_t i = 1; i < argc; i++)
     {
@@ -237,6 +238,10 @@ int main(int argc, char** argv)
         {
             visualize = true;
         }
+        else if (arg == "--dot")
+        {
+            dot_file = true;
+        }
         else if (arg == "-d" || arg == "--directory")
         {
             directory = argv[++i];
@@ -294,6 +299,7 @@ OPTIONS
         --timing_detail           Gather detailed timing
         -w|--warmup_iterations    Number of warm-up iterations
         --no_copy_data            Disable copy of input/result data every iteration
+        --dot                     Generate graphviz dot file
 )###";
         return 1;
     }
@@ -339,10 +345,10 @@ OPTIONS
             {
                 shared_ptr<Function> f = deserialize(model);
                 auto model_file_name = ngraph::file_util::get_file_name(model) + std::string(".") +
-                                       pass::VisualizeTree::get_file_ext();
+                                       (dot_file ? "dot" : pass::VisualizeTree::get_file_ext());
 
                 pass::Manager pass_manager;
-                pass_manager.register_pass<pass::VisualizeTree>(model_file_name);
+                pass_manager.register_pass<pass::VisualizeTree>(model_file_name, nullptr, true);
                 pass_manager.run_passes(f);
             }