diff --git a/taichi/ir/snode.cpp b/taichi/ir/snode.cpp index 6a4fe1f3890f0..0cb333fd22039 100644 --- a/taichi/ir/snode.cpp +++ b/taichi/ir/snode.cpp @@ -91,8 +91,8 @@ SNode &SNode::create_node(std::vector axes, } if (acc_shape > std::numeric_limits::max()) { TI_WARN( - "Snode index might be out of int32 boundary but int64 indexing is not " - "supported yet."); + "SNode index might be out of int32 boundary but int64 indexing is not " + "supported yet. Struct fors might not work either."); } new_node.num_cells_per_container = acc_shape; // infer extractors (only for POT) diff --git a/taichi/transforms/demote_dense_struct_fors.cpp b/taichi/transforms/demote_dense_struct_fors.cpp index 55ba8c780c7f7..192aa7cf52742 100644 --- a/taichi/transforms/demote_dense_struct_fors.cpp +++ b/taichi/transforms/demote_dense_struct_fors.cpp @@ -26,7 +26,6 @@ void convert_to_range_for(OffloadedStmt *offloaded, bool packed) { snode = snode->parent; } std::reverse(snodes.begin(), snodes.end()); - TI_ASSERT(total_bits <= 30); // general shape calculation - no dependence on POT int64 total_n = 1; @@ -38,6 +37,7 @@ void convert_to_range_for(OffloadedStmt *offloaded, bool packed) { } total_n *= s->num_cells_per_container; } + TI_ASSERT(total_n <= std::numeric_limits::max()); offloaded->const_begin = true; offloaded->const_end = true; diff --git a/tests/python/test_struct_for_non_pot.py b/tests/python/test_struct_for_non_pot.py index c8eefbe3defb5..dd35434dfe517 100644 --- a/tests/python/test_struct_for_non_pot.py +++ b/tests/python/test_struct_for_non_pot.py @@ -66,3 +66,18 @@ def test_2d(): @test_utils.test(require=ti.extension.packed, packed=True) def test_2d_packed(): _test_2d() + + +@test_utils.test(require=ti.extension.packed, packed=True) +def test_2d_overflow_if_not_packed(): + n, m, p = 2**9 + 1, 2**9 + 1, 2**10 + 1 + arr = ti.field(ti.u8, (n, m, p)) + + @ti.kernel + def count() -> ti.i32: + res = 0 + for _ in ti.grouped(arr): + res += 1 + return res + + assert count() == n * m * p