[perf] Refactor kernel profiler

benchmarks/fill_dense.py

@@ -128,9 +128,9 @@ def fill():
 # ti.cfg.arch = ti.cuda
 # ti.cfg.print_kernel_llvm_ir_optimized = True
 # ti.cfg.print_kernel_llvm_ir = True
-ti.cfg.enable_profiler = True
+ti.cfg.kernel_profiler = True
 # ti.cfg.verbose_kernel_launches = True
 print(benchmark_nested_struct_listgen_8x8())
 # print(benchmark_root_listgen())
-ti.profiler_print()
+ti.kernel_profiler_print()
 '''

benchmarks/fill_sparse.py

@@ -43,7 +43,7 @@ def task():
 
 
 '''
-ti.init(arch=ti.cuda, enable_profiler=True)
+ti.init(arch=ti.cuda, kernel_profiler=True)
 benchmark_nested_struct_fill_and_clear()
-ti.profiler_print()
+ti.kernel_profiler_print()
 '''

examples/mgpcg.py

@@ -2,7 +2,7 @@
 import taichi as ti
 
 real = ti.f32
-ti.init(default_fp=real, arch=ti.x64, enable_profiler=True)
+ti.init(default_fp=real, arch=ti.x64, kernel_profiler=True)
 
 # grid parameters
 N = 128
@@ -207,4 +207,4 @@ def paint():
     gui.set_image(pixels)
     gui.show()
 
-ti.profiler_print()
+ti.kernel_profiler_print()

examples/mgpcg_advanced.py

@@ -3,7 +3,7 @@
 import taichi as ti
 
 real = ti.f32
-ti.init(default_fp=real, arch=ti.x64, enable_profiler=True)
+ti.init(default_fp=real, arch=ti.x64, kernel_profiler=True)
 
 
 @ti.data_oriented
@@ -207,7 +207,7 @@ def run(self):
             gui.set_image(self.pixels)
             gui.show()
 
-        ti.profiler_print()
+        ti.kernel_profiler_print()
 
 
 solver = MGPCG()

misc/benchmark_parallel_compilation.py

@@ -65,5 +65,5 @@ def substep():
 for i in range(32):
     substep()
 
-ti.profiler_print()
+ti.kernel_profiler_print()
 ti.core.print_profile_info()

python/taichi/lang/__init__.py

@@ -35,10 +35,10 @@
 opengl = core.opengl
 gpu = [cuda, metal, opengl]
 cpu = core.host_arch()
-profiler_print = lambda: core.get_current_program().profiler_print()
-profiler_clear = lambda: core.get_current_program().profiler_clear()
-profiler_start = lambda n: core.get_current_program().profiler_start(n)
-profiler_stop = lambda: core.get_current_program().profiler_stop()
+kernel_profiler_print = lambda: core.get_current_program(
+).kernel_profiler_print()
+kernel_profiler_clear = lambda: core.get_current_program(
+).kernel_profiler_clear()
 
 
 class _Extension(object):

python/taichi/misc/gui.py

@@ -41,9 +41,6 @@ def __init__(self, name, res=512, background_color=0x0):
         self.key_pressed = set()
         self.event = None
         self.clear()
-        if ti.core.get_current_program():
-            self.core.set_profiler(
-                ti.core.get_current_program().get_profiler())
 
     def __enter__(self):
         return self

taichi/backends/cpu/codegen_cpu.cpp

@@ -50,7 +50,7 @@ class CodeGenLLVMCPU : public CodeGenLLVM {
     stat.add("codegen_offloaded_tasks");
     using Type = OffloadedStmt::TaskType;
     auto offloaded_task_name = init_offloaded_task_function(stmt);
-    if (prog->config.enable_profiler) {
+    if (prog->config.kernel_profiler && arch_is_cpu(prog->config.arch)) {
       call(
           builder.get(), "LLVMRuntime_profiler_start",
           {get_runtime(), builder->CreateGlobalStringPtr(offloaded_task_name)});
@@ -72,7 +72,7 @@ class CodeGenLLVMCPU : public CodeGenLLVM {
     } else {
       TI_NOT_IMPLEMENTED
     }
-    if (prog->config.enable_profiler) {
+    if (prog->config.kernel_profiler && arch_is_cpu(prog->config.arch)) {
       call(builder.get(), "LLVMRuntime_profiler_stop", {get_runtime()});
     }
     finalize_offloaded_task_function();

taichi/backends/cuda/cuda_context.h

@@ -4,7 +4,7 @@
 #include <unordered_map>
 #include <thread>
 
-#include "taichi/program/profiler.h"
+#include "taichi/program/kernel_profiler.h"
 #include "taichi/backends/cuda/cuda_driver.h"
 
 TLANG_NAMESPACE_BEGIN
@@ -23,7 +23,7 @@ class CUDAContext {
   int dev_count;
   std::string mcpu;
   std::mutex lock;
-  ProfilerBase *profiler;
+  KernelProfilerBase *profiler;
   CUDADriver &driver;
 
   static std::unordered_map<std::thread::id, std::unique_ptr<CUDAContext>>
@@ -45,9 +45,10 @@ class CUDAContext {
               unsigned gridDim,
               unsigned blockDim);
 
-  void set_profiler(ProfilerBase *profiler) {
+  void set_profiler(KernelProfilerBase *profiler) {
     this->profiler = profiler;
   }
+
   std::string get_mcpu() const {
     return mcpu;
   }

taichi/backends/metal/kernel_manager.cpp

@@ -212,7 +212,7 @@ class CompiledTaichiKernel {
     const SNodeDescriptorsMap *snode_descriptors;
     MTLDevice *device;
     MemoryPool *mem_pool;
-    ProfilerBase *profiler;
+    KernelProfilerBase *profiler;
   };
 
   CompiledTaichiKernel(Params params) : ctx_attribs(*params.ctx_attribs) {
@@ -651,7 +651,7 @@ class KernelManager::Impl {
   CompileConfig *const config_;
   const CompiledStructs compiled_structs_;
   MemoryPool *const mem_pool_;
-  ProfilerBase *const profiler_;
+  KernelProfilerBase *const profiler_;
   nsobj_unique_ptr<MTLDevice> device_;
   nsobj_unique_ptr<MTLCommandQueue> command_queue_;
   nsobj_unique_ptr<MTLCommandBuffer> cur_command_buffer_;

taichi/backends/metal/kernel_manager.h

@@ -7,7 +7,7 @@
 
 #include "taichi/backends/metal/kernel_util.h"
 #include "taichi/lang_util.h"
-#include "taichi/program/profiler.h"
+#include "taichi/program/kernel_profiler.h"
 #include "taichi/backends/metal/struct_metal.h"
 #include "taichi/system/memory_pool.h"
 
@@ -27,7 +27,7 @@ class KernelManager {
     CompiledStructs compiled_structs;
     CompileConfig *config;
     MemoryPool *mem_pool;
-    ProfilerBase *profiler;
+    KernelProfilerBase *profiler;
     int root_id;
   };
 

taichi/gui/gui.h

@@ -2,7 +2,7 @@
 
 #include "taichi/math/math.h"
 #include "taichi/system/timer.h"
-#include "taichi/program/profiler.h"
+#include "taichi/program/kernel_profiler.h"
 
 #include <atomic>
 #include <ctime>
@@ -488,7 +488,6 @@ class GUI : public GUIBase {
   Vector2i cursor_pos;
   bool button_status[3];
   int widget_height;
-  lang::ProfilerBase *profiler;
 
   void set_mouse_pos(int x, int y) {
     cursor_pos = Vector2i(x, y);
@@ -886,10 +885,6 @@ class GUI : public GUIBase {
   }
 
   ~GUI();
-
-  void set_profiler(lang::ProfilerBase *profiler) {
-    this->profiler = profiler;
-  }
 };
 
 TI_NAMESPACE_END

taichi/jit/jit_module.h

@@ -6,7 +6,7 @@
 #include "taichi/inc/constants.h"
 #include "taichi/llvm/llvm_fwd.h"
 #include "taichi/lang_util.h"
-#include "taichi/program/profiler.h"
+#include "taichi/program/kernel_profiler.h"
 
 TLANG_NAMESPACE_BEGIN
 

taichi/program/compile_config.cpp

@@ -29,7 +29,7 @@ CompileConfig::CompileConfig() {
   default_fp = DataType::f32;
   default_ip = DataType::i32;
   verbose_kernel_launches = false;
-  enable_profiler = false;
+  kernel_profiler = false;
   default_cpu_block_dim = 0;  // 0 = adaptive
   default_gpu_block_dim = 64;
   verbose = true;

taichi/program/compile_config.h

@@ -32,7 +32,7 @@ struct CompileConfig {
   bool print_kernel_llvm_ir_optimized;
   bool print_kernel_nvptx;
   bool verbose_kernel_launches;
-  bool enable_profiler;
+  bool kernel_profiler;
   bool verbose;
   bool fast_math;
   bool use_unified_memory;

taichi/program/kernel.cpp

@@ -90,13 +90,17 @@ void Kernel::operator()() {
     }
     compiled(program.get_context());
     program.sync = (program.sync && arch_is_cpu(arch));
-    if (program.config.debug && arch_is_cpu(arch)) {
+    // Note that Kernel::arch may be different from program.config.arch
+    if (program.config.debug && arch_is_cpu(arch) &&
+        arch_is_cpu(program.config.arch)) {
       program.check_runtime_error();
     }
   } else {
     program.sync = false;
     program.async_engine->launch(this);
-    if (program.config.debug && arch_is_cpu(arch)) {
+    // Note that Kernel::arch may be different from program.config.arch
+    if (program.config.debug && arch_is_cpu(arch) &&
+        arch_is_cpu(program.config.arch)) {
       program.check_runtime_error();
     }
   }

taichi/program/profiler.cpp → taichi/program/kernel_profiler.cpp

@@ -1,11 +1,11 @@
-#include "profiler.h"
+#include "kernel_profiler.h"
 
 #include "taichi/system/timer.h"
 #include "taichi/backends/cuda/cuda_driver.h"
 
 TLANG_NAMESPACE_BEGIN
 
-void ProfileRecord::insert_sample(double t) {
+void KernelProfileRecord::insert_sample(double t) {
   if (counter == 0) {
     min = t;
     max = t;
@@ -16,21 +16,23 @@ void ProfileRecord::insert_sample(double t) {
   total += t;
 }
 
-void ProfilerBase::profiler_start(ProfilerBase *profiler,
-                                  const char *kernel_name) {
+void KernelProfilerBase::profiler_start(KernelProfilerBase *profiler,
+                                        const char *kernel_name) {
+  TI_ASSERT(profiler);
   profiler->start(std::string(kernel_name));
 }
 
-void ProfilerBase::profiler_stop(ProfilerBase *profiler) {
+void KernelProfilerBase::profiler_stop(KernelProfilerBase *profiler) {
+  TI_ASSERT(profiler);
   profiler->stop();
 }
 
-void ProfilerBase::print() {
+void KernelProfilerBase::print() {
   sync();
   printf("%s\n", title().c_str());
   for (auto &rec : records) {
     printf(
-        "[%6.2f%%] %-40s     min %7.3f ms   avg %7.3f ms    max %7.3f ms   "
+        "[%6.2f%%] %-40s    min %7.3f ms   avg %7.3f ms   max %7.3f ms   "
         "total %7.3f s [%7dx]\n",
         rec.total / total_time * 100.0f, rec.name.c_str(), rec.min,
         rec.total / rec.counter, rec.max, rec.total / 1000.0f, rec.counter);
@@ -39,7 +41,7 @@ void ProfilerBase::print() {
 
 namespace {
 // A simple profiler that uses Time::get_time()
-class DefaultProfiler : public ProfilerBase {
+class DefaultProfiler : public KernelProfilerBase {
  public:
   explicit DefaultProfiler(Arch arch)
       : title_(fmt::format("{} Profiler", arch_name(arch))) {
@@ -60,9 +62,9 @@ class DefaultProfiler : public ProfilerBase {
   void stop() override {
     auto t = Time::get_time() - start_t_;
     auto ms = t * 1000.0;
-    auto it =
-        std::find_if(records.begin(), records.end(),
-                     [&](ProfileRecord &r) { return r.name == event_name_; });
+    auto it = std::find_if(
+        records.begin(), records.end(),
+        [&](KernelProfileRecord &r) { return r.name == event_name_; });
     if (it == records.end()) {
       records.emplace_back(event_name_);
       it = std::prev(records.end());
@@ -78,7 +80,7 @@ class DefaultProfiler : public ProfilerBase {
 };
 
 // A CUDA kernel profiler that uses CUDA timing events
-class CUDAProfiler : public ProfilerBase {
+class KernelProfilerCUDA : public KernelProfilerBase {
  public:
 #if defined(TI_WITH_CUDA)
   void *current_stop;
@@ -123,7 +125,7 @@ class CUDAProfiler : public ProfilerBase {
         CUDADriver::get_instance().event_elapsed_time(&ms, start, stop);
         auto it = std::find_if(
             records.begin(), records.end(),
-            [&](ProfileRecord &r) { return r.name == map_elem.first; });
+            [&](KernelProfileRecord &r) { return r.name == map_elem.first; });
         if (it == records.end()) {
           records.emplace_back(map_elem.first);
           it = std::prev(records.end());
@@ -138,19 +140,19 @@ class CUDAProfiler : public ProfilerBase {
 #endif
   }
 
-  static CUDAProfiler &get_instance() {
-    static CUDAProfiler profiler;
+  static KernelProfilerCUDA &get_instance() {
+    static KernelProfilerCUDA profiler;
     return profiler;
   }
 };
 }  // namespace
 
-std::unique_ptr<ProfilerBase> make_profiler(Arch arch) {
+std::unique_ptr<KernelProfilerBase> make_profiler(Arch arch) {
   if (arch == Arch::x64 || arch == Arch::arm64 || arch == Arch::metal ||
       arch == Arch::opengl) {
     return std::make_unique<DefaultProfiler>(arch);
   } else if (arch == Arch::cuda) {
-    return std::make_unique<CUDAProfiler>();
+    return std::make_unique<KernelProfilerCUDA>();
   } else {
     TI_NOT_IMPLEMENTED;
   }