[type] [refactor] Add compute_type for CustomIntType

taichi/codegen/codegen_llvm.cpp

@@ -328,7 +328,8 @@ void CodeGenLLVM::visit(UnaryOpStmt *stmt) {
       auto from_size = 0;
       if (from->is<CustomIntType>()) {
         // TODO: replace 32 with a customizable type
-        from_size = 32;
+        from_size =
+            data_type_size(from->cast<CustomIntType>()->get_compute_type());
       } else {
         from_size = data_type_size(from);
       }
@@ -618,6 +619,31 @@ llvm::Type *CodeGenLLVM::llvm_type(DataType dt) {
   return nullptr;
 }
 
+llvm::Type *CodeGenLLVM::llvm_ptr_type(DataType dt) {
+  if (dt->is_primitive(PrimitiveTypeID::i8) ||
+      dt->is_primitive(PrimitiveTypeID::u8)) {
+    return llvm::Type::getInt8PtrTy(*llvm_context);
+  } else if (dt->is_primitive(PrimitiveTypeID::i16) ||
+             dt->is_primitive(PrimitiveTypeID::u16)) {
+    return llvm::Type::getInt16PtrTy(*llvm_context);
+  } else if (dt->is_primitive(PrimitiveTypeID::i32) ||
+             dt->is_primitive(PrimitiveTypeID::u32)) {
+    return llvm::Type::getInt32PtrTy(*llvm_context);
+  } else if (dt->is_primitive(PrimitiveTypeID::i64) ||
+             dt->is_primitive(PrimitiveTypeID::u64)) {
+    return llvm::Type::getInt64PtrTy(*llvm_context);
+  } else if (dt->is_primitive(PrimitiveTypeID::u1)) {
+    return llvm::Type::getInt1PtrTy(*llvm_context);
+  } else if (dt->is_primitive(PrimitiveTypeID::f32)) {
+    return llvm::Type::getFloatPtrTy(*llvm_context);
+  } else if (dt->is_primitive(PrimitiveTypeID::f64)) {
+    return llvm::Type::getDoublePtrTy(*llvm_context);
+  } else {
+    TI_NOT_IMPLEMENTED;
+  }
+  return nullptr;
+}
+
 void CodeGenLLVM::visit(TernaryOpStmt *stmt) {
   TI_ASSERT(stmt->op_type == TernaryOpType::select);
   llvm_val[stmt] = builder->CreateSelect(
@@ -907,8 +933,8 @@ void CodeGenLLVM::visit(KernelReturnStmt *stmt) {
     TI_NOT_IMPLEMENTED
   } else {
     auto intermediate_bits = 0;
-    if (stmt->value->ret_type->is<CustomIntType>()) {
-      intermediate_bits = 32;
+    if (auto cit = stmt->value->ret_type->cast<CustomIntType>()) {
+      intermediate_bits = data_type_bits(cit->get_compute_type());
     } else {
       intermediate_bits =
           tlctx->get_data_type(stmt->value->ret_type)->getPrimitiveSizeInBits();
@@ -1103,14 +1129,16 @@ void CodeGenLLVM::visit(GlobalStoreStmt *stmt) {
   auto ptr_type = stmt->ptr->ret_type->as<PointerType>();
   if (ptr_type->is_bit_pointer()) {
     auto cit = ptr_type->get_pointee_type()->as<CustomIntType>();
-    llvm::Value *byte_ptr = nullptr, *bit_offset = nullptr;
-    read_bit_pointer(llvm_val[stmt->ptr], byte_ptr, bit_offset);
+    llvm::Value *byte_ptr = nullptr, *bit_offset = nullptr,
+                *physical_type_size = nullptr;
+    read_bit_pointer(llvm_val[stmt->ptr], byte_ptr, bit_offset,
+                     physical_type_size);
     builder->CreateCall(
-        get_runtime_function("set_partial_bits_b32"),
+        get_runtime_function("set_partial_bits"),
         {builder->CreateBitCast(byte_ptr,
                                 llvm::Type::getInt32PtrTy(*llvm_context)),
          bit_offset, tlctx->get_constant(cit->get_num_bits()),
-         llvm_val[stmt->data]});
+         llvm_val[stmt->data], physical_type_size});
   } else {
     builder->CreateStore(llvm_val[stmt->data], llvm_val[stmt->ptr]);
   }
@@ -1122,25 +1150,16 @@ void CodeGenLLVM::visit(GlobalLoadStmt *stmt) {
   if (stmt->ptr->ret_type->as<PointerType>()->is_bit_pointer()) {
     auto cit = stmt->ret_type->as<CustomIntType>();
     // 1. load bit pointer
-    llvm::Value *byte_ptr, *bit_offset;
-    read_bit_pointer(llvm_val[stmt->ptr], byte_ptr, bit_offset);
-    auto bit_level_container = builder->CreateLoad(builder->CreateBitCast(
-        byte_ptr, llvm::Type::getInt32PtrTy(*llvm_context)));
-    // 2. bit shifting
-    //    first left shift `32 - (offset + num_bits)`
-    //    then right shift `32 - num_bits`
-    auto bit_end = builder->CreateAdd(bit_offset,
-                                      tlctx->get_constant(cit->get_num_bits()));
-    auto left = builder->CreateSub(tlctx->get_constant(32), bit_end);
-    auto right = builder->CreateAdd(tlctx->get_constant(32),
-                                    tlctx->get_constant(-cit->get_num_bits()));
-    auto step1 = builder->CreateShl(bit_level_container, left);
-    llvm::Value *step2 = nullptr;
-    if (cit->get_is_signed())
-      step2 = builder->CreateAShr(step1, right);
-    else
-      step2 = builder->CreateLShr(step1, right);
-    llvm_val[stmt] = step2;
+    llvm::Value *byte_ptr, *bit_offset, *physical_type_size;
+    read_bit_pointer(llvm_val[stmt->ptr], byte_ptr, bit_offset,
+                     physical_type_size);
+    auto tmp = builder->CreateCall(
+        get_runtime_function("load_partial_bits"),
+        {byte_ptr, bit_offset, tlctx->get_constant(cit->get_num_bits()),
+         physical_type_size,
+         tlctx->get_constant((uint32)cit->get_is_signed())});
+    llvm_val[stmt] = builder->CreateIntCast(
+        tmp, llvm_type(cit->get_compute_type()), cit->get_is_signed());
   } else {
     llvm_val[stmt] = builder->CreateLoad(tlctx->get_data_type(stmt->ret_type),
                                          llvm_val[stmt->ptr]);
@@ -1236,14 +1255,17 @@ void CodeGenLLVM::visit(LinearizeStmt *stmt) {
 void CodeGenLLVM::visit(IntegerOffsetStmt *stmt){TI_NOT_IMPLEMENTED}
 
 llvm::Value *CodeGenLLVM::create_bit_ptr_struct(llvm::Value *byte_ptr_base,
-                                                llvm::Value *bit_offset) {
+                                                llvm::Value *bit_offset,
+                                                int num_bits) {
   // 1. create a bit pointer struct
   // struct bit_pointer {
   //    i8* byte_ptr;
   //    i32 offset;
+  //    i32 physical_type_size;
   // };
   auto struct_type = llvm::StructType::get(
       *llvm_context, {llvm::Type::getInt8PtrTy(*llvm_context),
+                      llvm::Type::getInt32Ty(*llvm_context),
                       llvm::Type::getInt32Ty(*llvm_context)});
   // 2. alloca the bit pointer struct
   auto bit_ptr_struct = create_entry_block_alloca(struct_type);
@@ -1257,6 +1279,11 @@ llvm::Value *CodeGenLLVM::create_bit_ptr_struct(llvm::Value *byte_ptr_base,
   builder->CreateStore(
       bit_offset, builder->CreateGEP(bit_ptr_struct, {tlctx->get_constant(0),
                                                       tlctx->get_constant(1)}));
+  // 5. store `physical_type` in `bit_ptr_struct`
+  builder->CreateStore(
+      tlctx->get_constant(num_bits),
+      builder->CreateGEP(bit_ptr_struct,
+                         {tlctx->get_constant(0), tlctx->get_constant(2)}));
   return bit_ptr_struct;
 }
 
@@ -1284,7 +1311,10 @@ void CodeGenLLVM::visit(SNodeLookupStmt *stmt) {
         snode->dt.get_ptr()->as<BitArrayType>()->get_element_num_bits();
     auto offset = tlctx->get_constant(element_num_bits);
     offset = builder->CreateMul(offset, llvm_val[stmt->input_index]);
-    llvm_val[stmt] = create_bit_ptr_struct(llvm_val[stmt->input_snode], offset);
+    llvm_val[stmt] = create_bit_ptr_struct(
+        llvm_val[stmt->input_snode], offset,
+        data_type_bits(
+            snode->dt.get_ptr()->as<BitArrayType>()->get_physical_type()));
   } else {
     TI_INFO(snode_type_name(snode->type));
     TI_NOT_IMPLEMENTED
@@ -1299,7 +1329,9 @@ void CodeGenLLVM::visit(GetChStmt *stmt) {
     auto bit_offset = bit_struct->get_member_bit_offset(
         stmt->input_snode->child_id(stmt->output_snode));
     auto offset = tlctx->get_constant(bit_offset);
-    llvm_val[stmt] = create_bit_ptr_struct(llvm_val[stmt->input_ptr], offset);
+    llvm_val[stmt] =
+        create_bit_ptr_struct(llvm_val[stmt->input_ptr], offset,
+                              data_type_bits(bit_struct->get_physical_type()));
   } else {
     auto ch = create_call(stmt->output_snode->get_ch_from_parent_func_name(),
                           {builder->CreateBitCast(

taichi/codegen/codegen_llvm.h

@@ -148,6 +148,8 @@ class CodeGenLLVM : public IRVisitor, public LLVMModuleBuilder {
 
   llvm::Type *llvm_type(DataType dt);
 
+  llvm::Type *llvm_ptr_type(DataType dt);
+
   void visit(Block *stmt_list) override;
 
   void visit(AllocaStmt *stmt) override;
@@ -207,7 +209,8 @@ class CodeGenLLVM : public IRVisitor, public LLVMModuleBuilder {
   void visit(IntegerOffsetStmt *stmt) override;
 
   llvm::Value *create_bit_ptr_struct(llvm::Value *byte_ptr_base,
-                                     llvm::Value *bit_offset);
+                                     llvm::Value *bit_offset,
+                                     int num_bits = 32);
 
   void visit(SNodeLookupStmt *stmt) override;
 

taichi/ir/type.cpp

@@ -101,6 +101,58 @@ std::string CustomIntType::to_string() const {
   return fmt::format("c{}{}", is_signed_ ? 'i' : 'u', num_bits_);
 }
 
+CustomIntType::CustomIntType(int num_bits, bool is_signed)
+    : compute_type(nullptr), num_bits_(num_bits), is_signed_(is_signed) {
+  // TODO(type): support customizable compute_type
+  //  and should we expose it to users?
+  TI_ASSERT(num_bits <= 32);
+  if (is_signed) {
+    compute_type =
+        TypeFactory::get_instance().get_primitive_type(PrimitiveTypeID::i32);
+  } else {
+    compute_type =
+        TypeFactory::get_instance().get_primitive_type(PrimitiveTypeID::u32);
+  }
+}
+
+#define SET_COMPUTE_TYPE(n, N)         \
+  else if (n == N) {                   \
+    if (is_signed)                     \
+      type_id = PrimitiveTypeID::i##N; \
+    else                               \
+      type_id = PrimitiveTypeID::u##N; \
+  }
+
+CustomIntType::CustomIntType(int compute_type_bits,
+                             int num_bits,
+                             bool is_signed)
+    : compute_type(nullptr), num_bits_(num_bits), is_signed_(is_signed) {
+  auto type_id = PrimitiveTypeID::unknown;
+  if (false) {
+  }
+  SET_COMPUTE_TYPE(compute_type_bits, 64)
+  SET_COMPUTE_TYPE(compute_type_bits, 32)
+  SET_COMPUTE_TYPE(compute_type_bits, 16)
+  SET_COMPUTE_TYPE(compute_type_bits, 8)
+  else {TI_NOT_IMPLEMENTED} compute_type =
+      TypeFactory::get_instance().get_primitive_type(type_id);
+}
+
+BitStructType::BitStructType(PrimitiveType *physical_type,
+                             std::vector<Type *> member_types,
+                             std::vector<int> member_bit_offsets)
+    : physical_type_(physical_type),
+      member_types_(member_types),
+      member_bit_offsets_(member_bit_offsets) {
+  TI_ASSERT(member_types_.size() == member_bit_offsets_.size());
+  int physical_type_bits = data_type_bits(physical_type);
+  for (auto i = 0; i < member_types_.size(); ++i) {
+    auto bits_end = member_types_[i]->as<CustomIntType>()->get_num_bits() +
+                    member_bit_offsets_[i];
+    TI_ASSERT(physical_type_bits >= bits_end)
+  }
+}
+
 std::string BitStructType::to_string() const {
   std::string str = "bs(";
   int num_members = (int)member_bit_offsets_.size();

taichi/ir/type.h

@@ -166,12 +166,20 @@ class VectorType : public Type {
 
 class CustomIntType : public Type {
  public:
-  CustomIntType(int num_bits, bool is_signed)
-      : num_bits_(num_bits), is_signed_(is_signed) {
+  CustomIntType(int num_bits, bool is_signed);
+
+  CustomIntType(int compute_type_bits, int numBits, bool isSigned);
+
+  ~CustomIntType() override {
+    delete compute_type;
   }
 
   std::string to_string() const override;
 
+  Type *get_compute_type() {
+    return compute_type;
+  }
+
   int get_num_bits() const {
     return num_bits_;
   }
@@ -183,20 +191,26 @@ class CustomIntType : public Type {
  private:
   // TODO(type): for now we can uniformly use i32 as the "compute_type". It may
   // be a good idea to make "compute_type" also customizable.
-  int num_bits_;
-  bool is_signed_;
+  Type *compute_type{nullptr};
+  int num_bits_{32};
+  bool is_signed_{true};
 };
 
 class BitStructType : public Type {
  public:
   BitStructType(PrimitiveType *physical_type,
                 std::vector<Type *> member_types,
-                std::vector<int> member_bit_offsets)
-      : physical_type_(physical_type),
-        member_types_(member_types),
-        member_bit_offsets_(member_bit_offsets) {
-    TI_ASSERT(member_types_.size() == member_bit_offsets_.size());
-  }
+                std::vector<int> member_bit_offsets);
+  //      : physical_type_(physical_type),
+  //        member_types_(member_types),
+  //        member_bit_offsets_(member_bit_offsets) {
+  //    TI_ASSERT(member_types_.size() == member_bit_offsets_.size());
+  //    int physical_type_bits = data_type_bits(physical_type);
+  //    for (auto i = 0; i < member_types_.size(); ++i) {
+  //      auto bits_end = member_types_[i]->as<CustomIntType>()->get_num_bits()
+  //      + member_bit_offsets_[i]; TI_ASSERT(physical_type_bits >= bits_end)
+  //    }
+  //  }
 
   std::string to_string() const override;
 

taichi/ir/type_factory.cpp

@@ -46,6 +46,18 @@ Type *TypeFactory::get_custom_int_type(int num_bits, bool is_signed) {
   return custom_int_types_[key].get();
 }
 
+Type *TypeFactory::_get_custom_int_type(int compute_type_bits,
+                                        int num_bits,
+                                        bool is_signed) {
+  auto key = std::make_pair(num_bits, is_signed);
+  if (custom_int_types_with_compute_types_.find(key) ==
+      custom_int_types_with_compute_types_.end()) {
+    custom_int_types_with_compute_types_[key] =
+        std::make_unique<CustomIntType>(compute_type_bits, num_bits, is_signed);
+  }
+  return custom_int_types_with_compute_types_[key].get();
+}
+
 Type *TypeFactory::get_bit_struct_type(PrimitiveType *physical_type,
                                        std::vector<Type *> member_types,
                                        std::vector<int> member_bit_offsets) {

taichi/ir/type_factory.h

@@ -23,6 +23,10 @@ class TypeFactory {
 
   Type *get_custom_int_type(int num_bits, bool is_signed);
 
+  Type *_get_custom_int_type(int compute_type_bits,
+                             int num_bits,
+                             bool is_signed);
+
   Type *get_bit_struct_type(PrimitiveType *physical_type,
                             std::vector<Type *> member_types,
                             std::vector<int> member_bit_offsets);
@@ -49,6 +53,9 @@ class TypeFactory {
   // TODO: use unordered map
   std::map<std::pair<int, bool>, std::unique_ptr<Type>> custom_int_types_;
 
+  std::map<std::pair<int, bool>, std::unique_ptr<Type>>
+      custom_int_types_with_compute_types_;
+
   // TODO: avoid duplication
   std::vector<std::unique_ptr<Type>> bit_struct_types_;
 

taichi/lang_util.cpp

@@ -125,6 +125,10 @@ int data_type_size(DataType t) {
   }
 }
 
+int data_type_bits(DataType t) {
+  return data_type_size(t) * 8;
+}
+
 std::string data_type_short_name(DataType t) {
   if (!t->is<PrimitiveType>()) {
     return t->to_string();

taichi/lang_util.h

@@ -345,6 +345,7 @@ std::string make_list(const std::vector<T> &data,
 }
 
 int data_type_size(DataType t);
+int data_type_bits(DataType t);
 DataType promoted_type(DataType a, DataType b);
 
 extern std::string compiled_lib_dir;

taichi/llvm/llvm_codegen_utils.h

@@ -128,7 +128,8 @@ class LLVMModuleBuilder {
 
   void read_bit_pointer(llvm::Value *ptr,
                         llvm::Value *&byte_ptr,
-                        llvm::Value *&bit_offset) {
+                        llvm::Value *&bit_offset,
+                        llvm::Value *&physical_type) {
     // 1. load byte pointer
     auto byte_ptr_in_bit_struct = builder->CreateGEP(
         ptr, {tlctx->get_constant(0), tlctx->get_constant(0)});
@@ -140,6 +141,11 @@ class LLVMModuleBuilder {
         ptr, {tlctx->get_constant(0), tlctx->get_constant(1)});
     bit_offset = builder->CreateLoad(bit_offset_in_bit_struct);
     TI_ASSERT(bit_offset->getType()->isIntegerTy(32));
+
+    auto physical_type_bit_struct = builder->CreateGEP(
+        ptr, {tlctx->get_constant(0), tlctx->get_constant(2)});
+    physical_type = builder->CreateLoad(physical_type_bit_struct);
+    TI_ASSERT(bit_offset->getType()->isIntegerTy(32));
   }
 };
 

taichi/python/export_lang.cpp

@@ -728,6 +728,8 @@ void export_lang(py::module &m) {
   // TypeFactory on Python-scope pointer destruction.
   py::class_<TypeFactory>(m, "TypeFactory")
       .def("get_custom_int_type", &TypeFactory::get_custom_int_type,
+           py::return_value_policy::reference)
+      .def("_get_custom_int_type", &TypeFactory::_get_custom_int_type,
            py::return_value_policy::reference);
 
   m.def("get_type_factory_instance", TypeFactory::get_instance,