Repository: taichi-dev/taichi (language: python)

Issue:
Certain pointer-based SNode tree behaves incorrect in packed mode
<!-- We've collected some common issue solutions in https://docs.taichi-lang.org/docs/install. Make sure you've check them out first. Hopefully they could address your problem. -->

**Describe the bug**

When `packed=True` in global config and the SNode Tree is `pointer(3, 3)-dense(8, 8)`, struct-for loop will omit certain index even it has been activated:

**To Reproduce**

``` Python3 
import taichi as ti
ti.init(arch=ti.cpu, packed=True)
​
block_size, leaf_size = 3, 8
​
x = ti.field(ti.i32)
block = ti.root.pointer(ti.ij, (block_size, block_size))
block.dense(ti.ij, (leaf_size, leaf_size)).place(x)
​
print(x.to_numpy().shape)
​
@ti.kernel
def activate():
  x[7, 7] = 1
​
activate()
​
@ti.kernel
def test():
  for I in ti.grouped(x):
    if I[0] == 7 and I[1] == 7: print(I)
​
test()
​
# Expected Output:
# (24, 24)
# [7, 7]
​
# Actual Output (iff block_size == 3 and x[7, 7] = 1):
# (24, 24)
```


Discussion / hints:

Fix patch:
diff --git a/taichi/codegen/llvm/codegen_llvm.cpp b/taichi/codegen/llvm/codegen_llvm.cpp
index 57210e300ae76..2c3d0a5096e25 100644
--- a/taichi/codegen/llvm/codegen_llvm.cpp
+++ b/taichi/codegen/llvm/codegen_llvm.cpp
@@ -2261,7 +2261,8 @@ void TaskCodeGenLLVM::create_offload_struct_for(OffloadedStmt *stmt,
 
   int list_element_size = std::min(leaf_block->max_num_elements(),
                                    (int64)taichi_listgen_max_element_size);
-  int num_splits = std::max(1, list_element_size / stmt->block_dim);
+  int num_splits = std::max(1, list_element_size / stmt->block_dim +
+                                   (list_element_size % stmt->block_dim != 0));
 
   auto struct_for_func = get_runtime_function("parallel_struct_for");
 


Test patch:
diff --git a/tests/python/test_struct_for_non_pot.py b/tests/python/test_struct_for_non_pot.py
index 9632fde3383e8..3840e82797997 100644
--- a/tests/python/test_struct_for_non_pot.py
+++ b/tests/python/test_struct_for_non_pot.py
@@ -66,3 +66,40 @@ def test_2d():
 @test_utils.test(packed=True)
 def test_2d_packed():
     _test_2d()
+
+
+def _test_2d_pointer():
+    block_size, leaf_size = 3, 8
+    x = ti.field(ti.i32)
+    block = ti.root.pointer(ti.ij, (block_size, block_size))
+    block.dense(ti.ij, (leaf_size, leaf_size)).place(x)
+
+    @ti.kernel
+    def activate():
+        x[7, 7] = 1
+
+    activate()
+
+    @ti.kernel
+    def test() -> ti.i32:
+        res = 0
+        for I in ti.grouped(x):
+            res += I[0] + I[1] * 2
+        return res
+
+    ans = 0
+    for i in range(leaf_size):
+        for j in range(leaf_size):
+            ans += i + j * 2
+
+    assert ans == test()
+
+
+@test_utils.test(require=ti.extension.sparse, packed=False)
+def test_2d_pointer():
+    _test_2d_pointer()
+
+
+@test_utils.test(require=ti.extension.sparse, packed=True)
+def test_2d_pointer_packed():
+    _test_2d_pointer()