Replace is_custom_type() with is_quant()

python/taichi/lang/matrix.py

@@ -1456,10 +1456,7 @@ def _calc_dynamic_index_stride(self):
             self.dynamic_index_stride = 0
             return
         length = len(paths[0])
-        if any(
-                len(path) != length or ti_core.is_custom_type(path[length -
-                                                                   1]._dtype)
-                for path in paths):
+        if any(len(path) != length or ti_core.is_quant(path[length - 1]._dtype) for path in paths):
             return
         for i in range(length):
             if any(path[i] != paths[0][i] for path in paths):

taichi/ir/frontend_ir.cpp

@@ -510,10 +510,8 @@ void AtomicOpExpression::type_check(CompileConfig *) {
   };
   if (!val->ret_type->is<PrimitiveType>())
     error();
-  if (auto cit = dest->ret_type->cast<CustomIntType>()) {
-    ret_type = cit->get_compute_type();
-  } else if (auto cft = dest->ret_type->cast<CustomFloatType>()) {
-    ret_type = cft->get_compute_type();
+  if (is_quant(dest->ret_type)) {
+    ret_type = dest->ret_type->get_compute_type();
   } else if (dest->ret_type->is<PrimitiveType>()) {
     ret_type = dest->ret_type;
   } else {

taichi/ir/type_utils.h

@@ -73,7 +73,7 @@ inline PrimitiveTypeID get_primitive_data_type() {
   }
 }
 
-inline bool is_custom_type(DataType dt) {
+inline bool is_quant(DataType dt) {
   return dt->is<CustomIntType>() || dt->is<CustomFloatType>();
 }
 

taichi/python/export_lang.cpp

@@ -859,7 +859,7 @@ void export_lang(py::module &m) {
 #undef PER_TYPE
 
   m.def("data_type_size", data_type_size);
-  m.def("is_custom_type", is_custom_type);
+  m.def("is_quant", is_quant);
   m.def("is_integral", is_integral);
   m.def("is_signed", is_signed);
   m.def("is_real", is_real);

taichi/transforms/type_check.cpp

@@ -23,7 +23,7 @@ class TypeCheck : public IRVisitor {
                          Stmt *&val,
                          const std::string &stmt_name) {
     auto dst_type = dst->ret_type.ptr_removed();
-    if (dst_type->is<CustomIntType>() || dst_type->is<CustomFloatType>()) {
+    if (is_quant(dst_type)) {
       // We force the value type to be the compute_type of the bit pointer.
       // Casting from compute_type to physical_type is handled in codegen.
       dst_type = dst_type->get_compute_type();