[Lang] [type] Refine SNode with quant 5/n: Rename bit_array to quant_array

python/taichi/_snode/fields_builder.py

@@ -104,12 +104,12 @@ def bit_struct(self, num_bits: int):
         self.empty = False
         return self.root.bit_struct(num_bits)
 
-    def bit_array(self, indices: Union[Sequence[_Axis], _Axis],
-                  dimensions: Union[Sequence[int], int], num_bits: int):
-        """Same as :func:`taichi.lang.snode.SNode.bit_array`"""
+    def quant_array(self, indices: Union[Sequence[_Axis], _Axis],
+                    dimensions: Union[Sequence[int], int], num_bits: int):
+        """Same as :func:`taichi.lang.snode.SNode.quant_array`"""
         self._check_not_finalized()
         self.empty = False
-        return self.root.bit_array(indices, dimensions, num_bits)
+        return self.root.quant_array(indices, dimensions, num_bits)
 
     def place(self,
               *args: Any,

python/taichi/lang/misc.py

@@ -596,7 +596,7 @@ def _block_dim_adaptive(block_dim_adaptive):
 
 
 def _bit_vectorize():
-    """Enable bit vectorization of struct fors on bit_arrays.
+    """Enable bit vectorization of struct fors on quant_arrays.
     """
     get_runtime().prog.current_ast_builder().bit_vectorize()
 
@@ -614,7 +614,7 @@ def loop_config(*,
         serialize (bool): Whether to let the for loop execute serially, `serialize=True` equals to `parallelize=1`
         parallelize (int): The number of threads to use on CPU
         block_dim_adaptive (bool): Whether to allow backends set block_dim adaptively, enabled by default
-        bit_vectorize (bool): Whether to enable bit vectorization of struct fors on bit_arrays.
+        bit_vectorize (bool): Whether to enable bit vectorization of struct fors on quant_arrays.
 
     Examples::
 
@@ -644,8 +644,8 @@ def fill():
         x = ti.field(dtype=u1)
         y = ti.field(dtype=u1)
         cell = ti.root.dense(ti.ij, (128, 4))
-        cell.bit_array(ti.j, 32).place(x)
-        cell.bit_array(ti.j, 32).place(y)
+        cell.quant_array(ti.j, 32).place(x)
+        cell.quant_array(ti.j, 32).place(y)
         @ti.kernel
         def copy():
             ti.loop_config(bit_vectorize=True)

python/taichi/lang/snode.py

@@ -107,8 +107,8 @@ def bit_struct(self, num_bits: int):
         """
         return SNode(self.ptr.bit_struct(num_bits, impl.current_cfg().packed))
 
-    def bit_array(self, axes, dimensions, num_bits):
-        """Adds a bit_array SNode as a child component of `self`.
+    def quant_array(self, axes, dimensions, num_bits):
+        """Adds a quant_array SNode as a child component of `self`.
 
         Args:
             axes (List[Axis]): Axes to activate.
@@ -121,8 +121,8 @@ def bit_array(self, axes, dimensions, num_bits):
         if isinstance(dimensions, int):
             dimensions = [dimensions] * len(axes)
         return SNode(
-            self.ptr.bit_array(axes, dimensions, num_bits,
-                               impl.current_cfg().packed))
+            self.ptr.quant_array(axes, dimensions, num_bits,
+                                 impl.current_cfg().packed))
 
     def place(self, *args, offset=None, shared_exponent=False):
         """Places a list of Taichi fields under the `self` container.

taichi/codegen/llvm/codegen_llvm.cpp

@@ -241,7 +241,7 @@ std::unique_ptr<RuntimeObject> CodeGenLLVM::emit_struct_meta_object(
     meta =
         std::make_unique<RuntimeObject>("BitmaskedMeta", this, builder.get());
     emit_struct_meta_base("Bitmasked", meta->ptr, snode);
-  } else if (snode->type == SNodeType::bit_array) {
+  } else if (snode->type == SNodeType::quant_array) {
     meta = std::make_unique<RuntimeObject>("DenseMeta", this, builder.get());
     emit_struct_meta_base("Dense", meta->ptr, snode);
   } else {
@@ -1351,7 +1351,7 @@ void CodeGenLLVM::visit(GlobalStoreStmt *stmt) {
             "handled by BitStructStoreStmt.",
             pointee_type->to_string());
       } else {
-        TI_ERROR("Bit array only supports quant int type.");
+        TI_ERROR("Quant array only supports quant int type.");
       }
     }
     store_quant_int(llvm_val[stmt->dest], pointee_type->as<QuantIntType>(),
@@ -1414,8 +1414,8 @@ std::string CodeGenLLVM::get_runtime_snode_name(SNode *snode) {
     return "Bitmasked";
   } else if (snode->type == SNodeType::bit_struct) {
     return "BitStruct";
-  } else if (snode->type == SNodeType::bit_array) {
-    return "BitArray";
+  } else if (snode->type == SNodeType::quant_array) {
+    return "QuantArray";
   } else {
     TI_P(snode_type_name(snode->type));
     TI_NOT_IMPLEMENTED
@@ -1536,9 +1536,9 @@ void CodeGenLLVM::visit(SNodeLookupStmt *stmt) {
                           {llvm_val[stmt->input_index]});
   } else if (snode->type == SNodeType::bit_struct) {
     llvm_val[stmt] = parent;
-  } else if (snode->type == SNodeType::bit_array) {
+  } else if (snode->type == SNodeType::quant_array) {
     auto element_num_bits =
-        snode->dt->as<BitArrayType>()->get_element_num_bits();
+        snode->dt->as<QuantArrayType>()->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(llvm_val[stmt->input_snode], offset);
@@ -1549,7 +1549,7 @@ void CodeGenLLVM::visit(SNodeLookupStmt *stmt) {
 }
 
 void CodeGenLLVM::visit(GetChStmt *stmt) {
-  if (stmt->input_snode->type == SNodeType::bit_array) {
+  if (stmt->input_snode->type == SNodeType::quant_array) {
     llvm_val[stmt] = llvm_val[stmt->input_ptr];
   } else if (stmt->ret_type->as<PointerType>()->is_bit_pointer()) {
     auto bit_struct = stmt->input_snode->dt->cast<BitStructType>();
@@ -1728,16 +1728,16 @@ void CodeGenLLVM::create_offload_struct_for(OffloadedStmt *stmt, bool spmd) {
   llvm::Function *body = nullptr;
   auto leaf_block = stmt->snode;
 
-  // For a bit-vectorized loop over a bit array, we generate struct for on its
+  // For a bit-vectorized loop over a quant array, we generate struct for on its
   // parent node (must be "dense") instead of itself for higher performance.
   if (stmt->is_bit_vectorized) {
-    if (leaf_block->type == SNodeType::bit_array &&
+    if (leaf_block->type == SNodeType::quant_array &&
         leaf_block->parent->type == SNodeType::dense) {
       leaf_block = leaf_block->parent;
     } else {
       TI_ERROR(
-          "A bit-vectorized struct-for must loop over a bit array with a dense "
-          "parent");
+          "A bit-vectorized struct-for must loop over a quant array with a "
+          "dense parent");
     }
   }
 
@@ -1869,7 +1869,7 @@ void CodeGenLLVM::create_offload_struct_for(OffloadedStmt *stmt, bool spmd) {
     create_call(refine, {parent_coordinates, new_coordinates,
                          builder->CreateLoad(loop_index)});
 
-    // For a bit-vectorized loop over a bit array, one more refine step is
+    // For a bit-vectorized loop over a quant array, one more refine step is
     // needed to make final coordinates non-consecutive, since each thread will
     // process multiple coordinates via vectorization
     if (stmt->is_bit_vectorized) {

taichi/codegen/llvm/struct_llvm.cpp

@@ -87,17 +87,17 @@ void StructCompilerLLVM::generate_types(SNode &snode) {
     }
     snode.dt = snode.bit_struct_type_builder->build();
     body_type = tlctx_->get_data_type(snode.physical_type);
-  } else if (type == SNodeType::bit_array) {
-    // A bit array SNode should have only one child
+  } else if (type == SNodeType::quant_array) {
+    // A quant array SNode should have only one child
     TI_ASSERT(snode.ch.size() == 1);
     auto &ch = snode.ch[0];
     Type *ch_type = ch->dt;
     if (!arch_is_cpu(arch_)) {
       TI_ERROR_IF(data_type_bits(snode.physical_type) <= 16,
-                  "bit_array physical type must be at least 32 bits on "
+                  "quant_array physical type must be at least 32 bits on "
                   "non-CPU backends.");
     }
-    snode.dt = TypeFactory::get_instance().get_bit_array_type(
+    snode.dt = TypeFactory::get_instance().get_quant_array_type(
         snode.physical_type, ch_type, snode.num_cells_per_container);
 
     DataType container_primitive_type(snode.physical_type);

taichi/codegen/metal/struct_metal.cpp

@@ -67,7 +67,7 @@ class StructCompiler {
         if (ty == SNodeType::place) {
           // do nothing
         }
-        CHECK_UNSUPPORTED_TYPE(bit_array)
+        CHECK_UNSUPPORTED_TYPE(quant_array)
         CHECK_UNSUPPORTED_TYPE(hash)
         else {
           max_snodes_ = std::max(max_snodes_, sn->id);

taichi/inc/snodes.inc.h

@@ -6,5 +6,5 @@ PER_SNODE(bitmasked)
 PER_SNODE(hash)
 PER_SNODE(place)
 PER_SNODE(bit_struct)
-PER_SNODE(bit_array)
+PER_SNODE(quant_array)
 PER_SNODE(undefined)

taichi/ir/snode.cpp

@@ -25,7 +25,7 @@ SNode &SNode::insert_children(SNodeType t) {
   std::memcpy(new_ch->physical_index_position, physical_index_position,
               sizeof(physical_index_position));
   new_ch->num_active_indices = num_active_indices;
-  if (type == SNodeType::bit_struct || type == SNodeType::bit_array) {
+  if (type == SNodeType::bit_struct || type == SNodeType::quant_array) {
     new_ch->is_bit_level = true;
   } else {
     new_ch->is_bit_level = is_bit_level;
@@ -145,11 +145,11 @@ SNode &SNode::bit_struct(int num_bits, bool packed) {
   return snode;
 }
 
-SNode &SNode::bit_array(const std::vector<Axis> &axes,
-                        const std::vector<int> &sizes,
-                        int bits,
-                        bool packed) {
-  auto &snode = create_node(axes, sizes, SNodeType::bit_array, packed);
+SNode &SNode::quant_array(const std::vector<Axis> &axes,
+                          const std::vector<int> &sizes,
+                          int bits,
+                          bool packed) {
+  auto &snode = create_node(axes, sizes, SNodeType::quant_array, packed);
   snode.physical_type =
       TypeFactory::get_instance().get_primitive_int_type(bits, false);
   return snode;

taichi/ir/snode.h

@@ -128,7 +128,7 @@ class SNode {
   int chunk_size{0};
   std::size_t cell_size_bytes{0};
   std::size_t offset_bytes_in_parent_cell{0};
-  PrimitiveType *physical_type{nullptr};  // for bit_struct and bit_array only
+  PrimitiveType *physical_type{nullptr};  // for bit_struct and quant_array only
   DataType dt;
   bool has_ambient{false};
   TypedConstant ambient_val;
@@ -147,7 +147,7 @@ class SNode {
 
   // is_bit_level=false: the SNode is not bitpacked
   // is_bit_level=true: the SNode is bitpacked (i.e., strictly inside bit_struct
-  // or bit_array)
+  // or quant_array)
   bool is_bit_level{false};
 
   // Whether the path from root to |this| contains only `dense` SNodes.
@@ -247,10 +247,10 @@ class SNode {
 
   SNode &bit_struct(int bits, bool packed);
 
-  SNode &bit_array(const std::vector<Axis> &axes,
-                   const std::vector<int> &sizes,
-                   int bits,
-                   bool packed);
+  SNode &quant_array(const std::vector<Axis> &axes,
+                     const std::vector<int> &sizes,
+                     int bits,
+                     bool packed);
 
   void print();
 

taichi/ir/type.cpp

@@ -203,8 +203,8 @@ std::string BitStructType::to_string() const {
   return str + ")";
 }
 
-std::string BitArrayType::to_string() const {
-  return fmt::format("ba({}x{})", element_type_->to_string(), num_elements_);
+std::string QuantArrayType::to_string() const {
+  return fmt::format("qa({}x{})", element_type_->to_string(), num_elements_);
 }
 
 std::string TypedConstant::stringify() const {

taichi/ir/type.h

@@ -298,11 +298,11 @@ class BitStructType : public Type {
   std::vector<int> member_bit_offsets_;
 };
 
-class BitArrayType : public Type {
+class QuantArrayType : public Type {
  public:
-  BitArrayType(PrimitiveType *physical_type,
-               Type *element_type_,
-               int num_elements_)
+  QuantArrayType(PrimitiveType *physical_type,
+                 Type *element_type_,
+                 int num_elements_)
       : physical_type_(physical_type),
         element_type_(element_type_),
         num_elements_(num_elements_) {

taichi/ir/type_factory.cpp

@@ -86,12 +86,12 @@ Type *TypeFactory::get_bit_struct_type(PrimitiveType *physical_type,
   return bit_struct_types_.back().get();
 }
 
-Type *TypeFactory::get_bit_array_type(PrimitiveType *physical_type,
-                                      Type *element_type,
-                                      int num_elements) {
-  bit_array_types_.push_back(std::make_unique<BitArrayType>(
+Type *TypeFactory::get_quant_array_type(PrimitiveType *physical_type,
+                                        Type *element_type,
+                                        int num_elements) {
+  quant_array_types_.push_back(std::make_unique<QuantArrayType>(
       physical_type, element_type, num_elements));
-  return bit_array_types_.back().get();
+  return quant_array_types_.back().get();
 }
 
 PrimitiveType *TypeFactory::get_primitive_int_type(int bits, bool is_signed) {

taichi/ir/type_factory.h

@@ -37,9 +37,9 @@ class TypeFactory {
                             std::vector<Type *> member_types,
                             std::vector<int> member_bit_offsets);
 
-  Type *get_bit_array_type(PrimitiveType *physical_type,
-                           Type *element_type,
-                           int num_elements);
+  Type *get_quant_array_type(PrimitiveType *physical_type,
+                             Type *element_type,
+                             int num_elements);
 
   static DataType create_vector_or_scalar_type(int width,
                                                DataType element,
@@ -77,7 +77,7 @@ class TypeFactory {
   std::vector<std::unique_ptr<Type>> bit_struct_types_;
 
   // TODO: avoid duplication
-  std::vector<std::unique_ptr<Type>> bit_array_types_;
+  std::vector<std::unique_ptr<Type>> quant_array_types_;
 
   std::mutex mut_;
 };

taichi/python/export_lang.cpp

@@ -521,7 +521,8 @@ void export_lang(py::module &m) {
                                bool))(&SNode::bitmasked),
            py::return_value_policy::reference)
       .def("bit_struct", &SNode::bit_struct, py::return_value_policy::reference)
-      .def("bit_array", &SNode::bit_array, py::return_value_policy::reference)
+      .def("quant_array", &SNode::quant_array,
+           py::return_value_policy::reference)
       .def("place", &SNode::place)
       .def("data_type", [](SNode *snode) { return snode->dt; })
       .def("name", [](SNode *snode) { return snode->name; })