[misc] Expose "StateFlowGraph::benchmark_rebuild_graph()" (#2090)

taichi-dev · Dec 7, 2020 · 4168a97 · 4168a97
1 parent fbea0ee
commit 4168a97
Show file tree

Hide file tree

Showing 4 changed files with 7 additions and 5 deletions.
diff --git a/misc/benchmark_rebuild_graph.py b/misc/benchmark_rebuild_graph.py
@@ -23,4 +23,4 @@ def foo():
 for i in range(1000):
     foo()
 
-ti.sync()
+ti.get_runtime().prog.benchmark_rebuild_graph()
diff --git a/taichi/program/async_engine.h b/taichi/program/async_engine.h
@@ -148,8 +148,6 @@ class AsyncEngine {
 
   void debug_sfg(const std::string &suffix);
 
-  void benchmark_rebuild_graph();
-
  private:
   IRBank ir_bank_;
 

diff --git a/taichi/program/state_flow_graph.cpp b/taichi/program/state_flow_graph.cpp
@@ -1495,8 +1495,8 @@ void StateFlowGraph::benchmark_rebuild_graph() {
     for (int i = 0; i < 100; i++)
       rebuild_graph(/*sort=*/false);
     auto rebuild_t = Time::get_time() - t;
-    TI_INFO("nodes = {} total time {:.4f} ns; per_node {:.4f} ns",
-            nodes_.size(), rebuild_t * 1e7, 1e7 * rebuild_t / nodes_.size());
+    TI_INFO("nodes = {} total time {:.4f} ms; per_node {:.4f} ns",
+            nodes_.size(), rebuild_t * 1e4, 1e7 * rebuild_t / nodes_.size());
   }
 }
 

diff --git a/taichi/python/export_lang.cpp b/taichi/python/export_lang.cpp
@@ -196,6 +196,10 @@ void export_lang(py::module &m) {
       .def("print_snode_tree", &Program::print_snode_tree)
       .def("get_snode_num_dynamically_allocated",
            &Program::get_snode_num_dynamically_allocated)
+      .def("benchmark_rebuild_graph",
+           [](Program *program) {
+             program->async_engine->sfg->benchmark_rebuild_graph();
+           })
       .def("synchronize", &Program::synchronize);
 
   m.def("get_current_program", get_current_program,