[Lang] Limit non-first division of an axis on a SNodeTree path to a power of two

cpp_examples/aot_save.cpp

@@ -13,7 +13,7 @@ void aot_save(taichi::Arch arch) {
 
   // program.materialize_runtime();
   auto *root = new SNode(0, SNodeType::root);
-  auto *pointer = &root->dense(Axis(0), n, false);
+  auto *pointer = &root->dense(Axis(0), n, false, "");
   auto *place = &pointer->insert_children(SNodeType::place);
   place->dt = PrimitiveType::i32;
   program.add_snode_tree(std::unique_ptr<SNode>(root), /*compile_only=*/true);

cpp_examples/autograd.cpp

@@ -90,10 +90,10 @@ void autograd() {
     };
 
     auto *snode =
-        &root->dense(Axis(0), n, false).insert_children(SNodeType::place);
+        &root->dense(Axis(0), n, false, "").insert_children(SNodeType::place);
     snode->dt = PrimitiveType::f32;
     snode->grad_info = std::make_unique<GradInfoPrimal>(
-        &root->dense(Axis(0), n, false).insert_children(SNodeType::place));
+        &root->dense(Axis(0), n, false, "").insert_children(SNodeType::place));
     snode->get_adjoint()->dt = PrimitiveType::f32;
     snode->get_adjoint()->grad_info = std::make_unique<GradInfoAdjoint>();
     return snode;

cpp_examples/run_snode.cpp

@@ -47,7 +47,7 @@ void run_snode() {
   int n = 10;
   program.materialize_runtime();
   auto *root = new SNode(0, SNodeType::root);
-  auto *pointer = &root->pointer(Axis(0), n, false);
+  auto *pointer = &root->pointer(Axis(0), n, false, "");
   auto *place = &pointer->insert_children(SNodeType::place);
   place->dt = PrimitiveType::i32;
   program.add_snode_tree(std::unique_ptr<SNode>(root), /*compile_only=*/false);

python/taichi/lang/snode.py

@@ -3,7 +3,7 @@
 from taichi._lib import core as _ti_core
 from taichi.lang import expr, impl, matrix
 from taichi.lang.field import BitpackedFields, Field
-from taichi.lang.util import is_taichi_class
+from taichi.lang.util import get_traceback, is_taichi_class
 
 
 class SNode:
@@ -35,7 +35,7 @@ def dense(self, axes, dimensions):
             dimensions = [dimensions] * len(axes)
         return SNode(
             self.ptr.dense(axes, dimensions,
-                           impl.current_cfg().packed))
+                           impl.current_cfg().packed, get_traceback()))
 
     def pointer(self, axes, dimensions):
         """Adds a pointer SNode as a child component of `self`.
@@ -51,7 +51,7 @@ def pointer(self, axes, dimensions):
             dimensions = [dimensions] * len(axes)
         return SNode(
             self.ptr.pointer(axes, dimensions,
-                             impl.current_cfg().packed))
+                             impl.current_cfg().packed, get_traceback()))
 
     @staticmethod
     def _hash(axes, dimensions):
@@ -79,7 +79,7 @@ def dynamic(self, axis, dimension, chunk_size=None):
             chunk_size = dimension
         return SNode(
             self.ptr.dynamic(axis[0], dimension, chunk_size,
-                             impl.current_cfg().packed))
+                             impl.current_cfg().packed, get_traceback()))
 
     def bitmasked(self, axes, dimensions):
         """Adds a bitmasked SNode as a child component of `self`.
@@ -95,7 +95,7 @@ def bitmasked(self, axes, dimensions):
             dimensions = [dimensions] * len(axes)
         return SNode(
             self.ptr.bitmasked(axes, dimensions,
-                               impl.current_cfg().packed))
+                               impl.current_cfg().packed, get_traceback()))
 
     def quant_array(self, axes, dimensions, max_num_bits):
         """Adds a quant_array SNode as a child component of `self`.
@@ -112,7 +112,7 @@ def quant_array(self, axes, dimensions, max_num_bits):
             dimensions = [dimensions] * len(axes)
         return SNode(
             self.ptr.quant_array(axes, dimensions, max_num_bits,
-                                 impl.current_cfg().packed))
+                                 impl.current_cfg().packed, get_traceback()))
 
     def place(self, *args, offset=None):
         """Places a list of Taichi fields under the `self` container.
@@ -134,7 +134,7 @@ def place(self, *args, offset=None):
                 bit_struct_type = arg.bit_struct_type_builder.build()
                 bit_struct_snode = self.ptr.bit_struct(
                     bit_struct_type,
-                    impl.current_cfg().packed)
+                    impl.current_cfg().packed, get_traceback())
                 for (field, id_in_bit_struct) in arg.fields:
                     bit_struct_snode.place(field, offset, id_in_bit_struct)
             elif isinstance(arg, Field):

taichi/ir/snode.cpp

@@ -37,7 +37,8 @@ SNode &SNode::insert_children(SNodeType t) {
 SNode &SNode::create_node(std::vector<Axis> axes,
                           std::vector<int> sizes,
                           SNodeType type,
-                          bool packed) {
+                          bool packed,
+                          const std::string &tb) {
   TI_ASSERT(axes.size() == sizes.size() || sizes.size() == 1);
   if (sizes.size() == 1) {
     sizes = std::vector<int>(axes.size(), sizes[0]);
@@ -54,30 +55,28 @@ SNode &SNode::create_node(std::vector<Axis> axes,
       throw TaichiRuntimeError(
           "Every dimension of a Taichi field should be positive");
     }
-    auto &ind = axes[i];
-    new_node.extractors[ind.value].activate(
+    int ind = axes[i].value;
+    auto end = new_node.physical_index_position + new_node.num_active_indices;
+    bool is_first_division =
+        std::find(new_node.physical_index_position, end, ind) == end;
+    if (is_first_division) {
+      new_node.physical_index_position[new_node.num_active_indices++] = ind;
+    } else {
+      TI_WARN_IF(
+          packed && !bit::is_power_of_two(sizes[i]),
+          "Non-first division of an axis on a SNodeTree path should be a power "
+          "of two to achieve best performance:\n{} We plan to turn this "
+          "warning into an error at v1.4.0. If you do have a use case that "
+          "needs to violate this rule, please submit an issue to notify us.",
+          tb);
+    }
+    new_node.extractors[ind].activate(
         bit::log2int(bit::least_pot_bound(sizes[i])));
-    new_node.extractors[ind.value].num_elements_from_root *= sizes[i];
+    new_node.extractors[ind].num_elements_from_root *= sizes[i];
     if (packed) {
-      new_node.extractors[ind.value].shape = sizes[i];
+      new_node.extractors[ind].shape = sizes[i];
     } else {  // if not in packed mode, pad shape to POT
-      new_node.extractors[ind.value].shape =
-          1 << new_node.extractors[ind.value].num_bits;
-    }
-  }
-  // infer mappings
-  for (int i = 0; i < taichi_max_num_indices; i++) {
-    bool found = false;
-    for (int k = 0; k < taichi_max_num_indices; k++) {
-      if (new_node.physical_index_position[k] == i) {
-        found = true;
-        break;
-      }
-    }
-    if (found)
-      continue;
-    if (new_node.extractors[i].active) {
-      new_node.physical_index_position[new_node.num_active_indices++] = i;
+      new_node.extractors[ind].shape = 1 << new_node.extractors[ind].num_bits;
     }
   }
   std::sort(new_node.physical_index_position,
@@ -130,14 +129,20 @@ SNode &SNode::create_node(std::vector<Axis> axes,
   return new_node;
 }
 
-SNode &SNode::dynamic(const Axis &expr, int n, int chunk_size, bool packed) {
-  auto &snode = create_node({expr}, {n}, SNodeType::dynamic, packed);
+SNode &SNode::dynamic(const Axis &expr,
+                      int n,
+                      int chunk_size,
+                      bool packed,
+                      const std::string &tb) {
+  auto &snode = create_node({expr}, {n}, SNodeType::dynamic, packed, tb);
   snode.chunk_size = chunk_size;
   return snode;
 }
 
-SNode &SNode::bit_struct(BitStructType *bit_struct_type, bool packed) {
-  auto &snode = create_node({}, {}, SNodeType::bit_struct, packed);
+SNode &SNode::bit_struct(BitStructType *bit_struct_type,
+                         bool packed,
+                         const std::string &tb) {
+  auto &snode = create_node({}, {}, SNodeType::bit_struct, packed, tb);
   snode.dt = bit_struct_type;
   snode.physical_type = bit_struct_type->get_physical_type();
   return snode;
@@ -146,8 +151,9 @@ SNode &SNode::bit_struct(BitStructType *bit_struct_type, bool 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);
+                          bool packed,
+                          const std::string &tb) {
+  auto &snode = create_node(axes, sizes, SNodeType::quant_array, packed, tb);
   snode.physical_type =
       TypeFactory::get_instance().get_primitive_int_type(bits, false);
   return snode;

taichi/ir/snode.h

@@ -171,83 +171,115 @@ class SNode {
   SNode &create_node(std::vector<Axis> axes,
                      std::vector<int> sizes,
                      SNodeType type,
-                     bool packed);
+                     bool packed,
+                     const std::string &tb);
 
   // SNodes maintains how flattened index bits are taken from indices
   SNode &dense(const std::vector<Axis> &axes,
                const std::vector<int> &sizes,
-               bool packed) {
-    return create_node(axes, sizes, SNodeType::dense, packed);
+               bool packed,
+               const std::string &tb) {
+    return create_node(axes, sizes, SNodeType::dense, packed, tb);
   }
 
-  SNode &dense(const std::vector<Axis> &axes, int sizes, bool packed) {
-    return create_node(axes, std::vector<int>{sizes}, SNodeType::dense, packed);
+  SNode &dense(const std::vector<Axis> &axes,
+               int sizes,
+               bool packed,
+               const std::string &tb) {
+    return create_node(axes, std::vector<int>{sizes}, SNodeType::dense, packed,
+                       tb);
   }
 
-  SNode &dense(const Axis &axis, int size, bool packed) {
-    return SNode::dense(std::vector<Axis>{axis}, size, packed);
+  SNode &dense(const Axis &axis, int size, bool packed, const std::string &tb) {
+    return SNode::dense(std::vector<Axis>{axis}, size, packed, tb);
   }
 
   SNode &pointer(const std::vector<Axis> &axes,
                  const std::vector<int> &sizes,
-                 bool packed) {
-    return create_node(axes, sizes, SNodeType::pointer, packed);
+                 bool packed,
+                 const std::string &tb) {
+    return create_node(axes, sizes, SNodeType::pointer, packed, tb);
   }
 
-  SNode &pointer(const std::vector<Axis> &axes, int sizes, bool packed) {
+  SNode &pointer(const std::vector<Axis> &axes,
+                 int sizes,
+                 bool packed,
+                 const std::string &tb) {
     return create_node(axes, std::vector<int>{sizes}, SNodeType::pointer,
-                       packed);
+                       packed, tb);
   }
 
-  SNode &pointer(const Axis &axis, int size, bool packed) {
-    return SNode::pointer(std::vector<Axis>{axis}, size, packed);
+  SNode &pointer(const Axis &axis,
+                 int size,
+                 bool packed,
+                 const std::string &tb) {
+    return SNode::pointer(std::vector<Axis>{axis}, size, packed, tb);
   }
 
   SNode &bitmasked(const std::vector<Axis> &axes,
                    const std::vector<int> &sizes,
-                   bool packed) {
-    return create_node(axes, sizes, SNodeType::bitmasked, packed);
+                   bool packed,
+                   const std::string &tb) {
+    return create_node(axes, sizes, SNodeType::bitmasked, packed, tb);
   }
 
-  SNode &bitmasked(const std::vector<Axis> &axes, int sizes, bool packed) {
+  SNode &bitmasked(const std::vector<Axis> &axes,
+                   int sizes,
+                   bool packed,
+                   const std::string &tb) {
     return create_node(axes, std::vector<int>{sizes}, SNodeType::bitmasked,
-                       packed);
+                       packed, tb);
   }
 
-  SNode &bitmasked(const Axis &axis, int size, bool packed) {
-    return SNode::bitmasked(std::vector<Axis>{axis}, size, packed);
+  SNode &bitmasked(const Axis &axis,
+                   int size,
+                   bool packed,
+                   const std::string &tb) {
+    return SNode::bitmasked(std::vector<Axis>{axis}, size, packed, tb);
   }
 
   SNode &hash(const std::vector<Axis> &axes,
               const std::vector<int> &sizes,
-              bool packed) {
-    return create_node(axes, sizes, SNodeType::hash, packed);
+              bool packed,
+              const std::string &tb) {
+    return create_node(axes, sizes, SNodeType::hash, packed, tb);
   }
 
-  SNode &hash(const std::vector<Axis> &axes, int sizes, bool packed) {
-    return create_node(axes, std::vector<int>{sizes}, SNodeType::hash, packed);
+  SNode &hash(const std::vector<Axis> &axes,
+              int sizes,
+              bool packed,
+              const std::string &tb) {
+    return create_node(axes, std::vector<int>{sizes}, SNodeType::hash, packed,
+                       tb);
   }
 
-  SNode &hash(const Axis &axis, int size, bool packed) {
-    return hash(std::vector<Axis>{axis}, size, packed);
+  SNode &hash(const Axis &axis, int size, bool packed, const std::string &tb) {
+    return hash(std::vector<Axis>{axis}, size, packed, tb);
   }
 
   std::string type_name() {
     return snode_type_name(type);
   }
 
-  SNode &bit_struct(BitStructType *bit_struct_type, bool packed);
+  SNode &bit_struct(BitStructType *bit_struct_type,
+                    bool packed,
+                    const std::string &tb);
 
   SNode &quant_array(const std::vector<Axis> &axes,
                      const std::vector<int> &sizes,
                      int bits,
-                     bool packed);
+                     bool packed,
+                     const std::string &tb);
 
   void print();
 
   void set_index_offsets(std::vector<int> index_offsets);
 
-  SNode &dynamic(const Axis &expr, int n, int chunk_size, bool packed);
+  SNode &dynamic(const Axis &expr,
+                 int n,
+                 int chunk_size,
+                 bool packed,
+                 const std::string &tb);
 
   SNode &morton(bool val = true) {
     _morton = val;

taichi/program/snode_expr_utils.cpp

@@ -56,7 +56,7 @@ void place_child(Expr *expr_arg,
                  SNodeFieldMap *snode_to_exprs) {
   if (parent->type == SNodeType::root) {
     // never directly place to root
-    auto &ds = parent->dense(std::vector<Axis>(), {}, false);
+    auto &ds = parent->dense(std::vector<Axis>(), {}, false, "");
     place_child(expr_arg, offset, id_in_bit_struct, &ds, snode_to_exprs);
   } else {
     TI_ASSERT(expr_arg->is<FieldExpression>());

taichi/python/export_lang.cpp

@@ -499,22 +499,24 @@ void export_lang(py::module &m) {
       .def_readonly("id", &SNode::id)
       .def("dense",
            (SNode & (SNode::*)(const std::vector<Axis> &,
-                               const std::vector<int> &, bool))(&SNode::dense),
+                               const std::vector<int> &, bool,
+                               const std::string &))(&SNode::dense),
+           py::return_value_policy::reference)
+      .def("pointer",
+           (SNode & (SNode::*)(const std::vector<Axis> &,
+                               const std::vector<int> &, bool,
+                               const std::string &))(&SNode::pointer),
            py::return_value_policy::reference)
-      .def(
-          "pointer",
-          (SNode & (SNode::*)(const std::vector<Axis> &,
-                              const std::vector<int> &, bool))(&SNode::pointer),
-          py::return_value_policy::reference)
       .def("hash",
            (SNode & (SNode::*)(const std::vector<Axis> &,
-                               const std::vector<int> &, bool))(&SNode::hash),
+                               const std::vector<int> &, bool,
+                               const std::string &))(&SNode::hash),
            py::return_value_policy::reference)
       .def("dynamic", &SNode::dynamic, py::return_value_policy::reference)
       .def("bitmasked",
            (SNode & (SNode::*)(const std::vector<Axis> &,
-                               const std::vector<int> &,
-                               bool))(&SNode::bitmasked),
+                               const std::vector<int> &, bool,
+                               const std::string &))(&SNode::bitmasked),
            py::return_value_policy::reference)
       .def("bit_struct", &SNode::bit_struct, py::return_value_policy::reference)
       .def("quant_array", &SNode::quant_array,

tests/cpp/analysis/bls_analyzer_test.cpp

@@ -20,7 +20,8 @@ class BLSAnalyzerTest : public ::testing::Test {
   void SetUp() override {
     const std::vector<Axis> axes = {Axis{0}, Axis{1}};
     root_snode_ = std::make_unique<SNode>(/*depth=*/0, /*t=*/SNodeType::root);
-    parent_snode_ = &(root_snode_->dense(axes, /*sizes=*/kBlockSize, false));
+    parent_snode_ =
+        &(root_snode_->dense(axes, /*sizes=*/kBlockSize, false, ""));
     child_snode_ = &(parent_snode_->insert_children(SNodeType::place));
     child_snode_->dt = PrimitiveType::i32;
 

tests/cpp/aot/dx12/aot_save_load_test.cpp

@@ -23,7 +23,7 @@ namespace fs = std::filesystem;
   int n = 10;
 
   auto *root = new SNode(0, SNodeType::root);
-  auto *pointer = &root->dense(Axis(0), n, false);
+  auto *pointer = &root->dense(Axis(0), n, false, "");
   auto *place = &pointer->insert_children(SNodeType::place);
   place->dt = PrimitiveType::i32;
   program.add_snode_tree(std::unique_ptr<SNode>(root), /*compile_only=*/true);