[bug] Enable field-related checks in materialize() not only in first …

…call (taichi-dev#3906)

python/taichi/examples/algorithm/mciso_advanced.py

@@ -294,13 +294,6 @@ def __init__(self, N=64, dim=3, blk_size=None):
         self.use_sparse = blk_size is not None
         self.blk_size = blk_size
 
-        et = [self.et2, self.et3][dim - 2]
-        self.et = ti.Vector.field(dim, int, et.shape[:2])
-
-        @ti.materialize_callback
-        def init_et():
-            self.et.from_numpy(et)
-
         self.m = ti.field(float)  # field to sample
         self.g = ti.Vector.field(self.dim, float)  # normalized gradient
         indices = [ti.ij, ti.ijk][dim - 2]
@@ -316,6 +309,10 @@ def init_et():
             dim, float,
             (self.N**self.dim, self.dim))  # result buffer, TODO: optimize this
 
+        et = [self.et2, self.et3][dim - 2]
+        self.et = ti.Vector.field(dim, int, et.shape[:2])
+        self.et.from_numpy(et)
+
     @ti.kernel
     def compute_grad(self):
         for I in ti.grouped(self.g):

python/taichi/lang/__init__.py

@@ -27,9 +27,8 @@
                               begin_frontend_struct_for, call_internal,
                               current_cfg, expr_init, expr_init_func,
                               expr_init_list, field, get_runtime, grouped,
-                              insert_expr_stmt_if_ti_func,
-                              materialize_callback, ndarray, one, root, static,
-                              static_assert, static_print, stop_grad,
+                              insert_expr_stmt_if_ti_func, ndarray, one, root,
+                              static, static_assert, static_print, stop_grad,
                               subscript, ti_assert, ti_float, ti_format,
                               ti_int, ti_print, zero)
 from taichi.lang.kernel_arguments import SparseMatrixProxy

python/taichi/lang/impl.py

@@ -255,7 +255,6 @@ class PyTaichi:
     def __init__(self, kernels=None):
         self.materialized = False
         self.prog = None
-        self.materialize_callbacks = []
         self.compiled_functions = {}
         self.compiled_grad_functions = {}
         self.scope_stack = []
@@ -346,20 +345,12 @@ def _calc_matrix_field_dynamic_index_stride(self):
 
     def materialize(self):
         self.materialize_root_fb(not self.materialized)
-
-        if self.materialized:
-            return
-
         self.materialized = True
 
         self._check_field_not_placed()
         self._check_matrix_field_member_shape()
         self._calc_matrix_field_dynamic_index_stride()
 
-        for callback in self.materialize_callbacks:
-            callback()
-        self.materialize_callbacks = []
-
     def _register_signal_handlers(self):
         if self._signal_handler_registry is None:
             self._signal_handler_registry = _ti_core.HackedSignalRegister()
@@ -386,10 +377,6 @@ def get_runtime():
     return pytaichi
 
 
-def materialize_callback(foo):
-    get_runtime().materialize_callbacks.append(foo)
-
-
 def _clamp_unsigned_to_range(npty, val):
     # npty: np.int32 or np.int64
     iif = np.iinfo(npty)

python/taichi/lang/snode.py

@@ -284,18 +284,17 @@ def get_children(self):
     @property
     def num_dynamically_allocated(self):
         runtime = impl.get_runtime()
-        runtime.materialize()
+        runtime.materialize_root_fb(False)
         return runtime.prog.get_snode_num_dynamically_allocated(self.ptr)
 
     @property
     def cell_size_bytes(self):
-        runtime = impl.get_runtime()
-        runtime.materialize()
+        impl.get_runtime().materialize_root_fb(False)
         return self.ptr.cell_size_bytes
 
     @property
     def offset_bytes_in_parent_cell(self):
-        impl.get_runtime().materialize()
+        impl.get_runtime().materialize_root_fb(False)
         return self.ptr.offset_bytes_in_parent_cell
 
     def deactivate_all(self):

tests/python/test_matrix.py

@@ -486,3 +486,25 @@ def test_matrix_field_dynamic_index_different_stride():
 
     ti.get_runtime().materialize()
     assert v.dynamic_index_stride is None
+
+
+@ti.test(arch=[ti.cpu, ti.cuda], dynamic_index=True)
+def test_matrix_field_dynamic_index_multiple_materialize():
+    @ti.kernel
+    def empty():
+        pass
+
+    empty()
+
+    n = 5
+    a = ti.Vector.field(3, dtype=ti.i32, shape=n)
+
+    @ti.kernel
+    def func():
+        for i in a:
+            a[i][i % 3] = i
+
+    func()
+    for i in range(n):
+        for j in range(3):
+            assert a[i][j] == (i if j == i % 3 else 0)

tests/python/test_runtime.py

@@ -126,23 +126,6 @@ def test_init_bad_arg():
         ti.init(_test_mode=True, debug=True, foo_bar=233)
 
 
-@ti.test(arch=ti.cpu)
-def test_materialize_callback():
-    x = ti.field(ti.f32, (3, 4))
-
-    @ti.materialize_callback
-    @ti.kernel
-    def init_x():
-        for i in range(3):
-            for j in range(4):
-                x[i, j] = i + j + 1
-
-    # x will be initialized on first invocation
-    for i in range(3):
-        for j in range(4):
-            assert x[i, j] == i + j + 1
-
-
 @pytest.mark.parametrize('level', ti.supported_log_levels)
 @ti.test()
 def test_supported_log_levels(level):

tests/python/test_tensor_reflection.py

@@ -59,7 +59,7 @@ def test_unordered():
     assert val.snode in blk3.get_children()
     assert blk3 in blk2.get_children()
     assert blk2 in blk1.get_children()
-    ti.get_runtime().materialize()
+    ti.get_runtime().materialize_root_fb(False)
     assert blk1 in ti.FieldsBuilder.finalized_roots()[0].get_children()
 
     expected_str = f'ti.root => dense {[n]} => dense {[m, n]}' \