[type] Initial TypeFactory class

taichi/ir/type_factory.cpp

@@ -0,0 +1,15 @@
+#include "taichi/ir/type_factory.h"
+
+TLANG_NAMESPACE_BEGIN
+
+Type *TypeFactory::get_primitive_type(PrimitiveType::primitive_type id) {
+  std::lock_guard<std::mutex> _(mut_);
+
+  if (primitive_types_.find(id) == primitive_types_.end()) {
+    primitive_types_[id] = std::make_unique<PrimitiveType>(id);
+  }
+
+  return primitive_types_[id].get();
+}
+
+TLANG_NAMESPACE_END

taichi/ir/type_factory.h

@@ -0,0 +1,18 @@
+#include "taichi/lang_util.h"
+
+#include <mutex>
+
+TLANG_NAMESPACE_BEGIN
+
+class TypeFactory {
+ public:
+  Type *get_primitive_type(PrimitiveType::primitive_type id);
+
+ private:
+  std::unordered_map<PrimitiveType::primitive_type, std::unique_ptr<Type>>
+      primitive_types_;
+
+  std::mutex mut_;
+};
+
+TLANG_NAMESPACE_END

taichi/lang_util.cpp

@@ -4,6 +4,7 @@
 
 #include "taichi/math/linalg.h"
 #include "taichi/program/arch.h"
+#include "taichi/program/program.h"
 #include "taichi/program/compile_config.h"
 #include "taichi/system/timer.h"
 
@@ -31,10 +32,14 @@ real default_measurement_time = 1;
 
 // Note: these primitive types should never be freed. They are supposed to live
 // together with the process. This is a temporary solution. Later we should
-// manage its ownership more systematically
-#define PER_TYPE(x)           \
-  DataType PrimitiveType::x = \
-      DataType(new PrimitiveType(PrimitiveType::primitive_type::x));
+// manage its ownership more systematically.
+
+// This part doesn't look good, but we will remove it soon anyway.
+#define PER_TYPE(x)                                            \
+  DataType PrimitiveType::x =                                  \
+      DataType(Program::get_type_factory().get_primitive_type( \
+          PrimitiveType::primitive_type::x));
+
 #include "taichi/inc/data_type.inc.h"
 #undef PER_TYPE
 

taichi/program/program.cpp

@@ -194,6 +194,12 @@ Program::Program(Arch desired_arch) {
            arch_name(arch));
 }
 
+TypeFactory &Program::get_type_factory() {
+  // type_factory should never be destroyed, hence the raw new operator.
+  static TypeFactory *type_factory = new TypeFactory;
+  return *type_factory;
+}
+
 FunctionType Program::compile(Kernel &kernel) {
   auto start_t = Time::get_time();
   TI_AUTO_PROF;

taichi/program/program.h

@@ -8,6 +8,7 @@
 
 #define TI_RUNTIME_HOST
 #include "taichi/ir/ir.h"
+#include "taichi/ir/type_factory.h"
 #include "taichi/ir/snode.h"
 #include "taichi/lang_util.h"
 #include "taichi/llvm/llvm_context.h"
@@ -105,6 +106,10 @@ class Program {
       jit_evaluator_cache;
   std::mutex jit_evaluator_cache_mut;
 
+  // Note: for now we let all Programs share a single TypeFactory for smooth
+  // migration. In the future each program should have its own copy.
+  static TypeFactory &get_type_factory();
+
   Program() : Program(default_compile_config.arch) {
   }