Program::get_total_compilation_time

iqeqsun · Nov 10, 2019 · 21f0b8f · 21f0b8f
1 parent 7821f7c
commit 21f0b8f
Show file tree

Hide file tree

Showing 4 changed files with 27 additions and 28 deletions.
diff --git a/examples/difftaichi/smoke_taichi_gpu.py b/examples/difftaichi/smoke_taichi_gpu.py
@@ -8,7 +8,7 @@
 real = ti.f32
 ti.set_default_fp(real)
 
-num_iterations = 50
+num_iterations = 150
 n_grid = 110
 dx = 1.0 / n_grid
 num_iterations_gauss_seidel = 6
@@ -191,12 +191,13 @@ def main():
   for opt in range(num_iterations):
     t = time.time()
     with ti.Tape(loss):
-      output = "test" if opt % 10 == 9 else None
+      output = "test" if opt % 10 == -1 else None
       forward(output)
     print('total time', (time.time() - t) * 1000, 'ms')
 
     print('Iter', opt, ' Loss =', loss[None])
     apply_grad()
+    print("Compilation time:", ti.get_runtime().prog.get_total_compilation_time())
 
   forward("output")
 

diff --git a/taichi/program.cpp b/taichi/program.cpp
@@ -20,6 +20,7 @@ std::atomic<int> Program::num_instances;
 SNode root;
 
 FunctionType Program::compile(Kernel &kernel) {
+  auto start_t = Time::get_time();
   TI_AUTO_PROF;
   FunctionType ret = nullptr;
   if (kernel.arch == Arch::x86_64) {
@@ -32,6 +33,7 @@ FunctionType Program::compile(Kernel &kernel) {
     TC_NOT_IMPLEMENTED;
   }
   TC_ASSERT(ret);
+  total_compilation_time += Time::get_time() - start_t;
   return ret;
 }
 
@@ -158,6 +160,7 @@ Program::Program(Arch arch) {
   }
 #endif
   TC_ASSERT_INFO(num_instances == 0, "Only one instance at a time");
+  total_compilation_time = 0;
   num_instances += 1;
   SNode::counter = 0;
   // llvm_context_device is initialized before kernel compilation

diff --git a/taichi/program.h b/taichi/program.h
@@ -2,15 +2,15 @@
 
 #pragma once
 
-#include <taichi/context.h>
-#include <taichi/unified_allocator.h>
-#include <taichi/profiler.h>
-#include <atomic>
-#include "tlang_util.h"
-#include "snode.h"
 #include "ir.h"
-#include "taichi_llvm_context.h"
 #include "kernel.h"
+#include "snode.h"
+#include "taichi_llvm_context.h"
+#include "tlang_util.h"
+#include <atomic>
+#include <taichi/context.h>
+#include <taichi/profiler.h>
+#include <taichi/unified_allocator.h>
 #if defined(TC_PLATFORM_UNIX)
 #include <dlfcn.h>
 #endif
@@ -20,12 +20,10 @@ TLANG_NAMESPACE_BEGIN
 extern Program *current_program;
 extern SNode root;
 
-TC_FORCE_INLINE Program &get_current_program() {
-  return *current_program;
-}
+TC_FORCE_INLINE Program &get_current_program() { return *current_program; }
 
 class Program {
- public:
+public:
   using Kernel = taichi::Tlang::Kernel;
   // Should be copiable
   std::vector<void *> loaded_dlls;
@@ -40,9 +38,10 @@ class Program {
   CPUProfiler cpu_profiler;
   Context context;
   std::unique_ptr<TaichiLLVMContext> llvm_context_host, llvm_context_device;
-  bool sync;  // device/host synchronized?
+  bool sync; // device/host synchronized?
   bool clear_all_gradients_initialized;
   bool finalized;
+  float64 total_compilation_time;
   static std::atomic<int> num_instances;
 
   std::vector<std::unique_ptr<Kernel>> functions;
@@ -76,11 +75,10 @@ class Program {
     return context;
   }
 
-  Program() : Program(default_compile_config.arch) {
-  }
+  Program() : Program(default_compile_config.arch) {}
 
   Program(const Program &){
-      TC_NOT_IMPLEMENTED  // for pybind11..
+      TC_NOT_IMPLEMENTED // for pybind11..
   }
 
   Program(Arch arch);
@@ -95,7 +93,7 @@ class Program {
 #if defined(TC_PLATFORM_UNIX)
       dlclose(dll);
 #else
-	TC_NOT_IMPLEMENTED
+      TC_NOT_IMPLEMENTED
 #endif
     }
     UnifiedAllocator::free();
@@ -136,8 +134,7 @@ class Program {
   }
 
   Kernel &kernel(const std::function<void()> &body,
-                 const std::string &name = "",
-                 bool grad = false) {
+                 const std::string &name = "", bool grad = false) {
     // Expr::set_allow_store(true);
     auto func = std::make_unique<Kernel>(*this, body, name, grad);
     // Expr::set_allow_store(false);
@@ -146,12 +143,9 @@ class Program {
     return *functions.back();
   }
 
-  void start_function_definition(Kernel *func) {
-    current_kernel = func;
-  }
+  void start_function_definition(Kernel *func) { current_kernel = func; }
 
-  void end_function_definition() {
-  }
+  void end_function_definition() {}
 
   FunctionType compile(Kernel &kernel);
 
@@ -180,9 +174,9 @@ class Program {
 
   Kernel &get_snode_writer(SNode *snode);
 
-  Arch get_host_arch() {
-    return Arch::x86_64;
-  }
+  Arch get_host_arch() { return Arch::x86_64; }
+
+  float64 get_total_compilation_time() { return total_compilation_time; }
 };
 
 TLANG_NAMESPACE_END
diff --git a/taichi/python_bindings.cpp b/taichi/python_bindings.cpp
@@ -72,6 +72,7 @@ void export_lang(py::module &m) {
       .def("profiler_print", &Program::profiler_clear)
       .def("finalize", &Program::finalize)
       .def("get_snode_writer", &Program::get_snode_writer)
+      .def("get_total_compilation_time", &Program::get_total_compilation_time)
       .def("synchronize", &Program::synchronize);
 
   m.def("get_current_program", get_current_program,