[metal] Fix listgen when iterating over children of a bitmasked SNode

taichi/backends/metal/shaders/runtime_kernels.metal.h

@@ -91,14 +91,16 @@ STR(
         const int child_idx = (ii % num_slots);
         const auto parent_elem = parent_list.get<ListgenElement>(parent_idx);
         ListgenElement child_elem;
-        child_elem.root_mem_offset = parent_elem.root_mem_offset +
-                                     child_idx * child_stride +
-                                     child_meta.mem_offset_in_parent;
+        child_elem.root_mem_offset =
+            parent_elem.root_mem_offset + child_idx * child_stride;
+        // We are testing the activation of *parent* cells, do not add the
+        // offset of the child within its parent cell here.
         if (is_active(root_addr + child_elem.root_mem_offset, parent_meta,
                       child_idx)) {
           refine_coordinates(parent_elem,
                              runtime->snode_extractors[parent_snode_id],
                              child_idx, &child_elem);
+          child_elem.root_mem_offset += child_meta.mem_offset_in_parent;
           child_list.append(child_elem);
         }
       }

tests/python/test_bitmasked.py

@@ -203,3 +203,38 @@ def run():
     run()
     assert c[None] == 4
     assert s[None] == 42
+
+
+@archs_support_bitmasked
+def test_bitmasked_offset_child():
+    x = ti.var(ti.i32)
+    y = ti.var(ti.i32)
+    z = ti.var(ti.i32)
+    s = ti.var(ti.i32, shape=())
+
+    n = 16
+    # Offset children:
+    # * In |bm|'s cell: |bm2| has a non-zero offset
+    # * In |bm2|'s cell: |z| has a non-zero offset
+    # * We iterate over |z| to test the listgen handles offsets correctly
+    bm = ti.root.bitmasked(ti.i, n)
+    bm.dense(ti.i, 2).place(x)
+    bm2 = bm.bitmasked(ti.i, 4)
+    bm2.dense(ti.i, 2).place(y)
+    bm2.bitmasked(ti.i, 4).place(z)
+
+    @ti.kernel
+    def func():
+        for _ in z:
+            s[None] += 1
+
+    z[0] = 1
+    z[7] = 1
+    z[42] = 1
+    z[53] = 1
+    z[88] = 1
+    z[101] = 1
+    z[233] = 1
+
+    func()
+    assert s[None] == 7