From 91b6735ecba71230d9be431d8b7a8c8c6004fd50 Mon Sep 17 00:00:00 2001 From: Yi Xu Date: Wed, 23 Nov 2022 13:59:55 +0800 Subject: [PATCH] [Lang] Limit non-first division of an axis on a SNodeTree path to a power of two (#6690) Issue: #6660 Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- cpp_examples/aot_save.cpp | 2 +- cpp_examples/autograd.cpp | 4 +- cpp_examples/run_snode.cpp | 2 +- python/taichi/lang/snode.py | 14 +-- taichi/ir/snode.cpp | 62 +++++++------ taichi/ir/snode.h | 88 +++++++++++++------ taichi/program/snode_expr_utils.cpp | 2 +- taichi/python/export_lang.cpp | 20 +++-- tests/cpp/analysis/bls_analyzer_test.cpp | 3 +- tests/cpp/aot/dx12/aot_save_load_test.cpp | 2 +- tests/cpp/codegen/refine_coordinates_test.cpp | 6 +- tests/cpp/struct/snode_tree_test.cpp | 4 +- .../cpp/transforms/make_block_local_test.cpp | 10 ++- .../scalar_pointer_lowerer_test.cpp | 10 +-- tests/python/test_bitmasked.py | 4 +- tests/python/test_packed_size.py | 4 +- 16 files changed, 140 insertions(+), 97 deletions(-) diff --git a/cpp_examples/aot_save.cpp b/cpp_examples/aot_save.cpp index a0c5f37308ee0..d4611e5c17928 100644 --- a/cpp_examples/aot_save.cpp +++ b/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(root), /*compile_only=*/true); diff --git a/cpp_examples/autograd.cpp b/cpp_examples/autograd.cpp index b25c39485b20a..d02bd6c3008d8 100644 --- a/cpp_examples/autograd.cpp +++ b/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( - &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(); return snode; diff --git a/cpp_examples/run_snode.cpp b/cpp_examples/run_snode.cpp index ef58dd06c935b..57b2198bf8bdf 100644 --- a/cpp_examples/run_snode.cpp +++ b/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(root), /*compile_only=*/false); diff --git a/python/taichi/lang/snode.py b/python/taichi/lang/snode.py index c8f67c3396b15..3d96752356a3a 100644 --- a/python/taichi/lang/snode.py +++ b/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): diff --git a/taichi/ir/snode.cpp b/taichi/ir/snode.cpp index 0cb333fd22039..1f5195a160ee5 100644 --- a/taichi/ir/snode.cpp +++ b/taichi/ir/snode.cpp @@ -37,7 +37,8 @@ SNode &SNode::insert_children(SNodeType t) { SNode &SNode::create_node(std::vector axes, std::vector 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(axes.size(), sizes[0]); @@ -54,30 +55,28 @@ SNode &SNode::create_node(std::vector 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 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 &axes, const std::vector &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; diff --git a/taichi/ir/snode.h b/taichi/ir/snode.h index 947e7e93b2769..7bf655c5f6abe 100644 --- a/taichi/ir/snode.h +++ b/taichi/ir/snode.h @@ -171,83 +171,115 @@ class SNode { SNode &create_node(std::vector axes, std::vector 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 &axes, const std::vector &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 &axes, int sizes, bool packed) { - return create_node(axes, std::vector{sizes}, SNodeType::dense, packed); + SNode &dense(const std::vector &axes, + int sizes, + bool packed, + const std::string &tb) { + return create_node(axes, std::vector{sizes}, SNodeType::dense, packed, + tb); } - SNode &dense(const Axis &axis, int size, bool packed) { - return SNode::dense(std::vector{axis}, size, packed); + SNode &dense(const Axis &axis, int size, bool packed, const std::string &tb) { + return SNode::dense(std::vector{axis}, size, packed, tb); } SNode &pointer(const std::vector &axes, const std::vector &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 &axes, int sizes, bool packed) { + SNode &pointer(const std::vector &axes, + int sizes, + bool packed, + const std::string &tb) { return create_node(axes, std::vector{sizes}, SNodeType::pointer, - packed); + packed, tb); } - SNode &pointer(const Axis &axis, int size, bool packed) { - return SNode::pointer(std::vector{axis}, size, packed); + SNode &pointer(const Axis &axis, + int size, + bool packed, + const std::string &tb) { + return SNode::pointer(std::vector{axis}, size, packed, tb); } SNode &bitmasked(const std::vector &axes, const std::vector &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 &axes, int sizes, bool packed) { + SNode &bitmasked(const std::vector &axes, + int sizes, + bool packed, + const std::string &tb) { return create_node(axes, std::vector{sizes}, SNodeType::bitmasked, - packed); + packed, tb); } - SNode &bitmasked(const Axis &axis, int size, bool packed) { - return SNode::bitmasked(std::vector{axis}, size, packed); + SNode &bitmasked(const Axis &axis, + int size, + bool packed, + const std::string &tb) { + return SNode::bitmasked(std::vector{axis}, size, packed, tb); } SNode &hash(const std::vector &axes, const std::vector &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 &axes, int sizes, bool packed) { - return create_node(axes, std::vector{sizes}, SNodeType::hash, packed); + SNode &hash(const std::vector &axes, + int sizes, + bool packed, + const std::string &tb) { + return create_node(axes, std::vector{sizes}, SNodeType::hash, packed, + tb); } - SNode &hash(const Axis &axis, int size, bool packed) { - return hash(std::vector{axis}, size, packed); + SNode &hash(const Axis &axis, int size, bool packed, const std::string &tb) { + return hash(std::vector{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 &axes, const std::vector &sizes, int bits, - bool packed); + bool packed, + const std::string &tb); void print(); void set_index_offsets(std::vector 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; diff --git a/taichi/program/snode_expr_utils.cpp b/taichi/program/snode_expr_utils.cpp index bd73f820a86a8..90d97d3fbb435 100644 --- a/taichi/program/snode_expr_utils.cpp +++ b/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(), {}, false); + auto &ds = parent->dense(std::vector(), {}, false, ""); place_child(expr_arg, offset, id_in_bit_struct, &ds, snode_to_exprs); } else { TI_ASSERT(expr_arg->is()); diff --git a/taichi/python/export_lang.cpp b/taichi/python/export_lang.cpp index 4e131e2b30837..120909777458b 100644 --- a/taichi/python/export_lang.cpp +++ b/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 &, - const std::vector &, bool))(&SNode::dense), + const std::vector &, bool, + const std::string &))(&SNode::dense), + py::return_value_policy::reference) + .def("pointer", + (SNode & (SNode::*)(const std::vector &, + const std::vector &, bool, + const std::string &))(&SNode::pointer), py::return_value_policy::reference) - .def( - "pointer", - (SNode & (SNode::*)(const std::vector &, - const std::vector &, bool))(&SNode::pointer), - py::return_value_policy::reference) .def("hash", (SNode & (SNode::*)(const std::vector &, - const std::vector &, bool))(&SNode::hash), + const std::vector &, 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 &, - const std::vector &, - bool))(&SNode::bitmasked), + const std::vector &, 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, diff --git a/tests/cpp/analysis/bls_analyzer_test.cpp b/tests/cpp/analysis/bls_analyzer_test.cpp index 205617df3b97e..24aeb9949ef3a 100644 --- a/tests/cpp/analysis/bls_analyzer_test.cpp +++ b/tests/cpp/analysis/bls_analyzer_test.cpp @@ -20,7 +20,8 @@ class BLSAnalyzerTest : public ::testing::Test { void SetUp() override { const std::vector axes = {Axis{0}, Axis{1}}; root_snode_ = std::make_unique(/*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; diff --git a/tests/cpp/aot/dx12/aot_save_load_test.cpp b/tests/cpp/aot/dx12/aot_save_load_test.cpp index bac7a58ebd25d..2b87c00cc07d1 100644 --- a/tests/cpp/aot/dx12/aot_save_load_test.cpp +++ b/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(root), /*compile_only=*/true); diff --git a/tests/cpp/codegen/refine_coordinates_test.cpp b/tests/cpp/codegen/refine_coordinates_test.cpp index f776f444fb819..5b3da54132c62 100644 --- a/tests/cpp/codegen/refine_coordinates_test.cpp +++ b/tests/cpp/codegen/refine_coordinates_test.cpp @@ -114,7 +114,7 @@ struct BitsRange { }; constexpr int kPointerSize = 5; -constexpr int kDenseSize = 7; +constexpr int kDenseSize = 8; class RefineCoordinatesTest : public ::testing::Test { protected: @@ -128,8 +128,8 @@ class RefineCoordinatesTest : public ::testing::Test { root_snode_ = std::make_unique(/*depth=*/0, /*t=*/SNodeType::root); const std::vector axes = {Axis{0}}; - ptr_snode_ = &(root_snode_->pointer(axes, kPointerSize, false)); - dense_snode_ = &(ptr_snode_->dense(axes, kDenseSize, false)); + ptr_snode_ = &(root_snode_->pointer(axes, kPointerSize, false, "")); + dense_snode_ = &(ptr_snode_->dense(axes, kDenseSize, false, "")); // Must end with a `place` SNode. auto &leaf_snode = dense_snode_->insert_children(SNodeType::place); leaf_snode.dt = PrimitiveType::f32; diff --git a/tests/cpp/struct/snode_tree_test.cpp b/tests/cpp/struct/snode_tree_test.cpp index ceec825bd5c32..07dbe9fe8a256 100644 --- a/tests/cpp/struct/snode_tree_test.cpp +++ b/tests/cpp/struct/snode_tree_test.cpp @@ -11,8 +11,8 @@ TEST(SNodeTree, GetSNodeToRootMapping) { const std::vector axes = {Axis{0}}; std::vector all_snode_ids; for (int i = 0; i < 3; ++i) { - auto &ptr_snode = root.pointer(axes, kSNodeSize, kPacked); - auto &dense_snode = ptr_snode.dense(axes, kSNodeSize, kPacked); + auto &ptr_snode = root.pointer(axes, kSNodeSize, kPacked, ""); + auto &dense_snode = ptr_snode.dense(axes, kSNodeSize, kPacked, ""); auto &leaf_snode = dense_snode.insert_children(SNodeType::place); all_snode_ids.push_back(ptr_snode.id); all_snode_ids.push_back(dense_snode.id); diff --git a/tests/cpp/transforms/make_block_local_test.cpp b/tests/cpp/transforms/make_block_local_test.cpp index c38a80f4916c2..59beb27ba4703 100644 --- a/tests/cpp/transforms/make_block_local_test.cpp +++ b/tests/cpp/transforms/make_block_local_test.cpp @@ -41,14 +41,16 @@ class MakeBlockLocalTest : public ::testing::Test { // want to see if the tests can handle the loop index scaling multiplier // (block_size) and infer the BLS size correctly. const std::vector axes = {Axis{0}, Axis{1}}; - pointer_snode_ = &(root_snode_->pointer(axes, pointer_size, false)); + pointer_snode_ = &(root_snode_->pointer(axes, pointer_size, false, "")); - bls_snode_ = &(pointer_snode_->dense(axes, /*sizes=*/block_size, false)); + bls_snode_ = + &(pointer_snode_->dense(axes, /*sizes=*/block_size, false, "")); bls_place_snode_ = &(bls_snode_->insert_children(SNodeType::place)); bls_place_snode_->dt = PrimitiveType::f32; - struct_for_snode_ = &(pointer_snode_->dynamic({Axis{2}}, /*n=*/1024, - /*chunk_size=*/128, false)); + struct_for_snode_ = + &(pointer_snode_->dynamic({Axis{2}}, /*n=*/1024, + /*chunk_size=*/128, false, "")); struct_for_place_snode_ = &(struct_for_snode_->insert_children(SNodeType::place)); struct_for_place_snode_->dt = PrimitiveType::i32; diff --git a/tests/cpp/transforms/scalar_pointer_lowerer_test.cpp b/tests/cpp/transforms/scalar_pointer_lowerer_test.cpp index dbfeabfe0c22c..ea6b20c197e6e 100644 --- a/tests/cpp/transforms/scalar_pointer_lowerer_test.cpp +++ b/tests/cpp/transforms/scalar_pointer_lowerer_test.cpp @@ -35,8 +35,8 @@ class ScalarPointerLowererTest : public ::testing::Test { void SetUp() override { root_snode_ = std::make_unique(/*depth=*/0, /*t=*/SNodeType::root); const std::vector axes = {Axis{0}}; - ptr_snode_ = &(root_snode_->pointer(axes, kPointerSize, false)); - dense_snode_ = &(ptr_snode_->dense(axes, kDenseSize, false)); + ptr_snode_ = &(root_snode_->pointer(axes, kPointerSize, false, "")); + dense_snode_ = &(ptr_snode_->dense(axes, kDenseSize, false, "")); // Must end with a `place` SNode. leaf_snode_ = &(dense_snode_->insert_children(SNodeType::place)); leaf_snode_->dt = PrimitiveType::f32; @@ -109,9 +109,9 @@ TEST(ScalarPointerLowerer, EliminateMod) { VecStatement lowered; Stmt *index = builder.get_int32(2); auto root = std::make_unique(/*depth=*/0, SNodeType::root); - SNode *dense_1 = &(root->dense({Axis{2}, Axis{1}}, /*size=*/7, kPacked)); - SNode *dense_2 = &(root->dense({Axis{1}}, /*size=*/3, kPacked)); - SNode *dense_3 = &(dense_2->dense({Axis{0}}, /*size=*/5, kPacked)); + SNode *dense_1 = &(root->dense({Axis{2}, Axis{1}}, /*size=*/7, kPacked, "")); + SNode *dense_2 = &(root->dense({Axis{1}}, /*size=*/3, kPacked, "")); + SNode *dense_3 = &(dense_2->dense({Axis{0}}, /*size=*/5, kPacked, "")); SNode *leaf_1 = &(dense_1->insert_children(SNodeType::place)); SNode *leaf_2 = &(dense_3->insert_children(SNodeType::place)); LowererImpl lowerer_1{leaf_1, diff --git a/tests/python/test_bitmasked.py b/tests/python/test_bitmasked.py index e31b4eb71f1e1..38babb36aab4a 100644 --- a/tests/python/test_bitmasked.py +++ b/tests/python/test_bitmasked.py @@ -7,7 +7,7 @@ def _test_basic(): c = ti.field(ti.i32) s = ti.field(ti.i32) - bm = ti.root.bitmasked(ti.ij, (3, 6)).bitmasked(ti.i, 5) + bm = ti.root.bitmasked(ti.ij, (3, 6)).bitmasked(ti.i, 8) bm.place(x) ti.root.place(c, s) @@ -183,7 +183,7 @@ def _test_sparsity_changes(): c = ti.field(ti.i32) s = ti.field(ti.i32) - bm = ti.root.bitmasked(ti.i, 5).bitmasked(ti.i, 3) + bm = ti.root.bitmasked(ti.i, 5).bitmasked(ti.i, 4) bm.place(x) ti.root.place(c, s) diff --git a/tests/python/test_packed_size.py b/tests/python/test_packed_size.py index e52c0e5e95dfc..a5802a52aed25 100644 --- a/tests/python/test_packed_size.py +++ b/tests/python/test_packed_size.py @@ -5,6 +5,6 @@ @test_utils.test(require=ti.extension.packed, packed=True) def test_packed_size(): x = ti.field(ti.i32) - ti.root.dense(ti.i, 17).dense(ti.ijk, 129).place(x) - assert x.shape == (17 * 129, 129, 129) + ti.root.dense(ti.l, 3).dense(ti.ijk, 129).place(x) + assert x.shape == (129, 129, 129, 3) assert x.snode.parent().parent()._cell_size_bytes == 4 * 129**3