[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();
@@ -1105,12 +1131,16 @@ void CodeGenLLVM::visit(GlobalStoreStmt *stmt) {
     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);
+    auto runtime_func_name = fmt::format(
+        "set_partial_bits_b{}", data_type_bits(cit->get_compute_type()));
     builder->CreateCall(
-        get_runtime_function("set_partial_bits_b32"),
+        get_runtime_function(runtime_func_name),
         {builder->CreateBitCast(byte_ptr,
-                                llvm::Type::getInt32PtrTy(*llvm_context)),
+                                llvm_ptr_type(cit->get_compute_type())),
          bit_offset, tlctx->get_constant(cit->get_num_bits()),
-         llvm_val[stmt->data]});
+         builder->CreateIntCast(llvm_val[stmt->data],
+                                llvm_type(cit->get_compute_type()),
+                                cit->get_is_signed())});
   } else {
     builder->CreateStore(llvm_val[stmt->data], llvm_val[stmt->ptr]);
   }
@@ -1125,15 +1155,21 @@ void CodeGenLLVM::visit(GlobalLoadStmt *stmt) {
     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)));
+        byte_ptr, llvm_ptr_type(cit->get_compute_type())));
     // 2. bit shifting
-    //    first left shift `32 - (offset + num_bits)`
-    //    then right shift `32 - num_bits`
+    //    first left shift `compute_type_size(like 32, 64, ...) - (offset +
+    //    num_bits)` then right shift `compute_type_size - num_bits`
+    auto compute_type_size = data_type_bits(cit->get_compute_type());
     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 left =
+        builder->CreateSub(tlctx->get_constant(compute_type_size), bit_end);
+    auto right = builder->CreateSub(tlctx->get_constant(compute_type_size),
+                                    tlctx->get_constant(cit->get_num_bits()));
+    left = builder->CreateIntCast(left, llvm_type(cit->get_compute_type()),
+                                  cit->get_is_signed());
+    right = builder->CreateIntCast(right, llvm_type(cit->get_compute_type()),
+                                   cit->get_is_signed());
     auto step1 = builder->CreateShl(bit_level_container, left);
     llvm::Value *step2 = nullptr;
     if (cit->get_is_signed())

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;

taichi/ir/type.cpp

@@ -101,6 +101,20 @@ 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);
+  }
+}
+
 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,18 @@ 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() 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,8 +189,9 @@ 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 {

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/runtime/llvm/runtime.cpp

@@ -81,9 +81,11 @@ using float32 = float;
 using float64 = double;
 
 using i8 = int8;
+using i16 = int16;
 using i32 = int32;
 using i64 = int64;
 using u8 = uint8;
+using u16 = uint16;
 using u32 = uint32;
 using u64 = uint64;
 using f32 = float32;
@@ -1551,17 +1553,24 @@ void stack_push(Ptr stack, size_t max_num_elements, std::size_t element_size) {
 
 #include "internal_functions.h"
 
-void set_partial_bits_b32(u32 *ptr, u32 offset, u32 bits, u32 value) {
-  u32 mask = ((((u32)1 << bits) - 1) << offset);
-  u32 new_value = 0;
-  u32 old_value = *ptr;
-  do {
-    old_value = *ptr;
-    new_value = (old_value & (~mask)) | (value << offset);
-  } while (!__atomic_compare_exchange(ptr, &old_value, &new_value, true,
-                                      std::memory_order::memory_order_seq_cst,
-                                      std::memory_order::memory_order_seq_cst));
-}
+#define DEFINE_SET_PARTIAL_BITS(N)                                            \
+  void set_partial_bits_b##N(u##N *ptr, u32 offset, u32 bits, u##N value) {   \
+    u##N mask = ((((u##N)1 << bits) - 1) << offset);                          \
+    u##N new_value = 0;                                                       \
+    u##N old_value = *ptr;                                                    \
+    do {                                                                      \
+      old_value = *ptr;                                                       \
+      new_value = (old_value & (~mask)) | (u##N)(value << offset);            \
+    } while (                                                                 \
+        !__atomic_compare_exchange(ptr, &old_value, &new_value, true,         \
+                                   std::memory_order::memory_order_seq_cst,   \
+                                   std::memory_order::memory_order_seq_cst)); \
+  }
+
+DEFINE_SET_PARTIAL_BITS(8);
+DEFINE_SET_PARTIAL_BITS(16);
+DEFINE_SET_PARTIAL_BITS(32);
+DEFINE_SET_PARTIAL_BITS(64);
 }
 
 #endif

taichi/transforms/type_check.cpp

@@ -268,11 +268,8 @@ class TypeCheck : public IRVisitor {
 
     if (stmt->lhs->ret_type != stmt->rhs->ret_type) {
       auto promote_custom_int_type = [&](Stmt *stmt, Stmt *hs) {
-        if (hs->ret_type->is<CustomIntType>()) {
-          if (hs->ret_type->cast<CustomIntType>()->get_is_signed())
-            return insert_type_cast_before(stmt, hs, get_data_type<int32>());
-          else
-            return insert_type_cast_before(stmt, hs, get_data_type<uint32>());
+        if (auto cit = hs->ret_type->cast<CustomIntType>()) {
+          return insert_type_cast_before(stmt, hs, cit->get_compute_type());
         }
         return hs;
       };