[refactor] Move OffloadedStmt::TaskType to a separate file (#1946)

taichi/analysis/clone.cpp

@@ -1,5 +1,6 @@
 #include "taichi/ir/ir.h"
 #include "taichi/ir/analysis.h"
+#include "taichi/ir/statements.h"
 #include "taichi/ir/transforms.h"
 #include "taichi/ir/visitors.h"
 #include "taichi/program/program.h"

taichi/backends/cc/codegen_cc.cpp

@@ -2,6 +2,7 @@
 #include "cc_kernel.h"
 #include "cc_layout.h"
 #include "taichi/ir/ir.h"
+#include "taichi/ir/statements.h"
 #include "taichi/ir/transforms.h"
 #include "taichi/util/line_appender.h"
 #include "taichi/util/str.h"

taichi/backends/cpu/codegen_cpu.cpp

@@ -6,6 +6,7 @@
 #include "taichi/lang_util.h"
 #include "taichi/program/program.h"
 #include "taichi/ir/ir.h"
+#include "taichi/ir/statements.h"
 #include "taichi/util/statistics.h"
 
 TLANG_NAMESPACE_BEGIN

taichi/backends/cuda/codegen_cuda.cpp

@@ -7,6 +7,7 @@
 #include "taichi/util/io.h"
 #include "taichi/util/statistics.h"
 #include "taichi/ir/ir.h"
+#include "taichi/ir/statements.h"
 #include "taichi/program/program.h"
 #include "taichi/lang_util.h"
 #include "taichi/backends/cuda/cuda_driver.h"

taichi/backends/metal/codegen_metal.cpp

@@ -8,6 +8,7 @@
 #include "taichi/backends/metal/env_config.h"
 #include "taichi/backends/metal/features.h"
 #include "taichi/ir/ir.h"
+#include "taichi/ir/statements.h"
 #include "taichi/ir/transforms.h"
 #include "taichi/math/arithmetic.h"
 #include "taichi/util/line_appender.h"

taichi/backends/metal/kernel_manager.cpp

@@ -34,7 +34,7 @@ namespace shaders {
 #include "taichi/backends/metal/shaders/runtime_utils.metal.h"
 }  // namespace shaders
 
-using KernelTaskType = OffloadedStmt::TaskType;
+using KernelTaskType = OffloadedTaskType;
 using BufferEnum = KernelAttributes::Buffers;
 
 inline int infer_msl_version(const TaichiKernelAttributes::UsedFeatures &f) {

taichi/backends/metal/kernel_manager.h

@@ -6,9 +6,10 @@
 #include <vector>
 
 #include "taichi/backends/metal/kernel_util.h"
+#include "taichi/backends/metal/struct_metal.h"
 #include "taichi/lang_util.h"
+#include "taichi/program/compile_config.h"
 #include "taichi/program/kernel_profiler.h"
-#include "taichi/backends/metal/struct_metal.h"
 #include "taichi/system/memory_pool.h"
 
 TLANG_NAMESPACE_BEGIN

taichi/backends/metal/kernel_util.cpp

@@ -42,13 +42,13 @@ std::string KernelAttributes::debug_string() const {
   std::string result;
   result += fmt::format(
       "<KernelAttributes name={} num_threads={} task_type={} buffers=[ ", name,
-      num_threads, OffloadedStmt::task_type_name(task_type));
+      num_threads, offloaded_task_type_name(task_type));
   for (auto b : buffers) {
     result += buffers_name(b) + " ";
   }
   result += "]";  // closes |buffers|
   // TODO(k-ye): show range_for
-  if (task_type == OffloadedStmt::TaskType::listgen) {
+  if (task_type == OffloadedTaskType::listgen) {
     result += fmt::format(" snode={}", runtime_list_op_attribs->snode->id);
   }
   result += ">";

taichi/backends/metal/kernel_util.h

@@ -4,7 +4,7 @@
 #include <string>
 #include <vector>
 
-#include "taichi/ir/statements.h"
+#include "taichi/ir/offloaded_task_type.h"
 #include "taichi/backends/metal/data_types.h"
 
 // Data structures defined in this file may overlap with some of the Taichi data
@@ -43,7 +43,7 @@ struct KernelAttributes {
   };
   std::string name;
   int num_threads;
-  OffloadedStmt::TaskType task_type;
+  OffloadedTaskType task_type;
 
   struct RangeForAttributes {
     // |begin| has differen meanings depending on |const_begin|:

taichi/backends/opengl/codegen_opengl.cpp

@@ -7,6 +7,7 @@
 #include "taichi/backends/opengl/opengl_data_types.h"
 #include "taichi/backends/opengl/opengl_kernel_util.h"
 #include "taichi/ir/ir.h"
+#include "taichi/ir/statements.h"
 #include "taichi/ir/transforms.h"
 #include "taichi/util/line_appender.h"
 #include "taichi/util/macros.h"

taichi/backends/opengl/opengl_api.h

@@ -6,8 +6,11 @@
 #include <vector>
 #include <optional>
 
-#include "opengl_kernel_util.h"
-#include "opengl_kernel_launcher.h"
+#include "taichi/backends/opengl/opengl_kernel_util.h"
+#include "taichi/backends/opengl/opengl_kernel_launcher.h"
+#define TI_RUNTIME_HOST
+#include "taichi/program/context.h"
+#undef TI_RUNTIME_HOST
 
 TLANG_NAMESPACE_BEGIN
 

taichi/backends/opengl/opengl_kernel_util.h

@@ -4,7 +4,7 @@
 #include <vector>
 #include <unordered_map>
 
-#include "taichi/ir/statements.h"
+#include "taichi/ir/snode.h"
 
 TLANG_NAMESPACE_BEGIN
 

taichi/backends/opengl/struct_opengl.cpp

@@ -1,5 +1,7 @@
 #include "struct_opengl.h"
 
+#include "taichi/ir/snode.h"
+
 TLANG_NAMESPACE_BEGIN
 namespace opengl {
 

taichi/codegen/codegen_llvm.cpp

@@ -1,5 +1,6 @@
 #include "taichi/codegen/codegen_llvm.h"
 
+#include "taichi/ir/statements.h"
 #include "taichi/struct/struct_llvm.h"
 #include "taichi/util/file_sequence_writer.h"
 

taichi/inc/offloaded_task_type.inc.h

@@ -0,0 +1,5 @@
+PER_TASK_TYPE(serial)
+PER_TASK_TYPE(range_for)
+PER_TASK_TYPE(struct_for)
+PER_TASK_TYPE(listgen)
+PER_TASK_TYPE(gc)

taichi/ir/frontend.cpp

@@ -1,6 +1,7 @@
 // Frontend constructs
 
-#include "frontend.h"
+#include "taichi/ir/frontend.h"
+#include "taichi/ir/statements.h"
 #include "taichi/program/program.h"
 
 TLANG_NAMESPACE_BEGIN

taichi/ir/offloaded_task_type.cpp

@@ -0,0 +1,15 @@
+#include "taichi/ir/offloaded_task_type.h"
+
+TLANG_NAMESPACE_BEGIN
+
+std::string offloaded_task_type_name(OffloadedTaskType tt) {
+  if (false) {
+  }
+#define PER_TASK_TYPE(x) else if (tt == OffloadedTaskType::x) return #x;
+#include "taichi/inc/offloaded_task_type.inc.h"
+#undef PER_TASK_TYPE
+  else
+    TI_NOT_IMPLEMENTED
+}
+
+TLANG_NAMESPACE_END

taichi/ir/offloaded_task_type.h

@@ -0,0 +1,17 @@
+#pragma once
+
+#include "taichi/common/core.h"
+
+#include <string>
+
+TLANG_NAMESPACE_BEGIN
+
+enum class OffloadedTaskType : int {
+#define PER_TASK_TYPE(x) x,
+#include "taichi/inc/offloaded_task_type.inc.h"
+#undef PER_TASK_TYPE
+};
+
+std::string offloaded_task_type_name(OffloadedTaskType tt);
+
+TLANG_NAMESPACE_END

taichi/ir/snode.cpp

@@ -2,6 +2,7 @@
 
 #include "taichi/ir/ir.h"
 #include "taichi/ir/frontend.h"
+#include "taichi/ir/statements.h"
 #include "taichi/backends/cuda/cuda_driver.h"
 
 TLANG_NAMESPACE_BEGIN
@@ -18,7 +19,7 @@ void set_kernel_args(const std::vector<int> &I,
 
 }  // namespace
 
-std::atomic<int> SNode::counter = 0;
+std::atomic<int> SNode::counter{0};
 
 SNode &SNode::insert_children(SNodeType t) {
   TI_ASSERT(t != SNodeType::root);

taichi/ir/statements.cpp

@@ -317,15 +317,7 @@ std::string OffloadedStmt::task_name() const {
 
 // static
 std::string OffloadedStmt::task_type_name(TaskType tt) {
-#define REGISTER_NAME(x) \
-  { TaskType::x, #x }
-  const static std::unordered_map<TaskType, std::string> m = {
-      REGISTER_NAME(serial),     REGISTER_NAME(range_for),
-      REGISTER_NAME(struct_for), REGISTER_NAME(listgen),
-      REGISTER_NAME(gc),
-  };
-#undef REGISTER_NAME
-  return m.find(tt)->second;
+  return offloaded_task_type_name(tt);
 }
 
 std::unique_ptr<Stmt> OffloadedStmt::clone() const {

taichi/ir/statements.h

@@ -1,6 +1,7 @@
 #pragma once
 
 #include "taichi/ir/ir.h"
+#include "taichi/ir/offloaded_task_type.h"
 #include "taichi/ir/scratch_pad.h"
 
 TLANG_NAMESPACE_BEGIN
@@ -789,13 +790,7 @@ class GetChStmt : public Stmt {
 
 class OffloadedStmt : public Stmt {
  public:
-  enum TaskType : int {
-    serial,
-    range_for,
-    struct_for,
-    listgen,
-    gc,
-  };
+  using TaskType = OffloadedTaskType;
 
   TaskType task_type;
   SNode *snode;
@@ -830,7 +825,7 @@ class OffloadedStmt : public Stmt {
   static std::string task_type_name(TaskType tt);
 
   bool has_body() const {
-    return task_type != listgen && task_type != gc;
+    return task_type != TaskType::listgen && task_type != TaskType::gc;
   }
 
   bool is_container_statement() const override {

taichi/program/async_engine.cpp

@@ -7,8 +7,9 @@
 #include "taichi/backends/cpu/codegen_cpu.h"
 #include "taichi/util/testing.h"
 #include "taichi/util/statistics.h"
-#include "taichi/ir/transforms.h"
 #include "taichi/ir/analysis.h"
+#include "taichi/ir/statements.h"
+#include "taichi/ir/transforms.h"
 #include "taichi/program/extension.h"
 
 TLANG_NAMESPACE_BEGIN

taichi/program/async_engine.h

@@ -8,7 +8,6 @@
 #include <unordered_map>
 
 #include "taichi/ir/ir.h"
-#include "taichi/ir/statements.h"
 #include "taichi/lang_util.h"
 #define TI_RUNTIME_HOST
 #include "taichi/program/context.h"

taichi/program/async_utils.cpp

@@ -38,7 +38,7 @@ bool TaskLaunchRecord::empty() const {
 
 void TaskMeta::print() const {
   fmt::print("TaskMeta\n  name {}\n", name);
-  fmt::print("  type {}\n", OffloadedStmt::task_type_name(type));
+  fmt::print("  type {}\n", offloaded_task_type_name(type));
   if (snode != nullptr) {
     fmt::print("  snode {}\n", snode->get_node_type_name_hinted());
   } else {
@@ -99,8 +99,8 @@ TaskMeta *get_task_meta(IRBank *ir_bank, const TaskLaunchRecord &t) {
   TaskMeta meta;
   // TODO: this is an abuse since it gathers nothing...
   auto *root_stmt = t.stmt();
-  meta.name = t.kernel->name + "_" +
-              OffloadedStmt::task_type_name(root_stmt->task_type);
+  meta.name =
+      t.kernel->name + "_" + offloaded_task_type_name(root_stmt->task_type);
   meta.type = root_stmt->task_type;
   get_meta_input_value_states(root_stmt, &meta);
   gather_statements(root_stmt, [&](Stmt *stmt) {
@@ -160,14 +160,14 @@ TaskMeta *get_task_meta(IRBank *ir_bank, const TaskLaunchRecord &t) {
     }
   }
 
-  if (root_stmt->task_type == OffloadedStmt::listgen) {
+  if (root_stmt->task_type == OffloadedTaskType::listgen) {
     TI_ASSERT(root_stmt->snode->parent);
     meta.snode = root_stmt->snode;
     meta.input_states.emplace(root_stmt->snode->parent, AsyncState::Type::list);
     meta.input_states.emplace(root_stmt->snode, AsyncState::Type::list);
     meta.input_states.emplace(root_stmt->snode, AsyncState::Type::mask);
     meta.output_states.emplace(root_stmt->snode, AsyncState::Type::list);
-  } else if (root_stmt->task_type == OffloadedStmt::struct_for) {
+  } else if (root_stmt->task_type == OffloadedTaskType::struct_for) {
     meta.snode = root_stmt->snode;
     meta.input_states.emplace(root_stmt->snode, AsyncState::Type::list);
   }
@@ -197,10 +197,10 @@ TaskFusionMeta get_task_fusion_meta(IRBank *bank, const TaskLaunchRecord &t) {
 
   auto *task = t.stmt();
   meta.type = task->task_type;
-  if (task->task_type == OffloadedStmt::struct_for) {
+  if (task->task_type == OffloadedTaskType::struct_for) {
     meta.snode = task->snode;
     meta.block_dim = task->block_dim;
-  } else if (task->task_type == OffloadedStmt::range_for) {
+  } else if (task->task_type == OffloadedTaskType::range_for) {
     // TODO: a few problems with the range-for test condition:
     // 1. This could incorrectly fuse two range-for kernels that have
     // different sizes, but then the loop ranges get padded to the same
@@ -215,7 +215,7 @@ TaskFusionMeta get_task_fusion_meta(IRBank *bank, const TaskLaunchRecord &t) {
     }
     meta.begin_value = task->begin_value;
     meta.end_value = task->end_value;
-  } else if (task->task_type != OffloadedStmt::serial) {
+  } else if (task->task_type != OffloadedTaskType::serial) {
     // Do not fuse gc/listgen tasks.
     return fusion_meta_bank[t.ir_handle] = TaskFusionMeta();
   }

taichi/program/async_utils.h

@@ -5,7 +5,7 @@
 #include <atomic>
 
 #include "taichi/ir/snode.h"
-#include "taichi/ir/statements.h"
+#include "taichi/ir/offloaded_task_type.h"
 #define TI_RUNTIME_HOST
 #include "taichi/program/context.h"
 #undef TI_RUNTIME_HOST
@@ -14,6 +14,9 @@ TLANG_NAMESPACE_BEGIN
 
 struct TaskMeta;
 
+class IRNode;
+class OffloadedStmt;
+
 class IRHandle {
  public:
   IRHandle() : ir_(nullptr), hash_(0) {
@@ -110,7 +113,7 @@ struct AsyncState {
 
 struct TaskFusionMeta {
   // meta for task fusion
-  OffloadedStmt::TaskType type{OffloadedStmt::TaskType::serial};
+  OffloadedTaskType type{OffloadedTaskType::serial};
   SNode *snode{nullptr};  // struct-for only
   int block_dim{0};       // struct-for only
   int32 begin_value{0};   // range-for only
@@ -172,7 +175,8 @@ struct hash<taichi::lang::AsyncState> {
 template <>
 struct hash<taichi::lang::TaskFusionMeta> {
   std::size_t operator()(const taichi::lang::TaskFusionMeta &t) const noexcept {
-    std::size_t result = (t.type << 1) ^ t.fusible ^ (std::size_t)t.kernel;
+    std::size_t result =
+        ((std::size_t)t.type << 1) ^ t.fusible ^ (std::size_t)t.kernel;
     result ^= (std::size_t)t.block_dim * 100000007UL + (std::size_t)t.snode;
     result ^= ((std::size_t)t.begin_value << 32) ^ t.end_value;
     return result;
@@ -185,7 +189,7 @@ TLANG_NAMESPACE_BEGIN
 
 struct TaskMeta {
   std::string name;
-  OffloadedStmt::TaskType type{OffloadedStmt::TaskType::serial};
+  OffloadedTaskType type{OffloadedTaskType::serial};
   SNode *snode{nullptr};  // struct-for and listgen only
   std::unordered_set<AsyncState> input_states;
   std::unordered_set<AsyncState> output_states;

taichi/program/ir_bank.cpp

@@ -1,7 +1,8 @@
 #include "taichi/program/ir_bank.h"
 
-#include "taichi/ir/transforms.h"
 #include "taichi/ir/analysis.h"
+#include "taichi/ir/statements.h"
+#include "taichi/ir/transforms.h"
 #include "taichi/program/kernel.h"
 
 TLANG_NAMESPACE_BEGIN

taichi/program/kernel.cpp

@@ -6,6 +6,7 @@
 #include "taichi/program/async_engine.h"
 #include "taichi/codegen/codegen.h"
 #include "taichi/backends/cuda/cuda_driver.h"
+#include "taichi/ir/statements.h"
 #include "taichi/ir/transforms.h"
 #include "taichi/util/action_recorder.h"
 #include "taichi/program/extension.h"