[Bug] Fix ndarray memory leak

python/gstaichi/lang/_ndarray.py

@@ -38,6 +38,12 @@ def __init__(self):
         # we register with runtime, in order to enable reset to work later
         impl.get_runtime().ndarrays.add(self)
 
+    def __del__(self):
+        if impl is not None and impl.get_runtime is not None and impl.get_runtime() is not None:
+            prog = impl.get_runtime()._prog
+            if prog is not None:
+                prog.delete_ndarray(self.arr)
+
     def to_dlpack(self):
         return impl.get_runtime().prog.ndarray_to_dlpack(self, self.arr)
 
@@ -270,12 +276,6 @@ def __init__(self, dtype, arr_shape):
         self.shape = tuple(self.arr.shape)
         self.element_type = dtype
 
-    def __del__(self):
-        if impl is not None and impl.get_runtime is not None and impl.get_runtime() is not None:
-            prog = impl.get_runtime()._prog
-            if prog is not None:
-                prog.delete_ndarray(self.arr)
-
     @property
     def element_shape(self):
         return ()

tests/python/test_ndarray.py

@@ -1199,3 +1199,16 @@ def my_kernel({args}) -> None:
         my_kernel(*arg_objs_l)
     for i in range(num_args):
         assert arg_objs_l[i][0] == i + 1
+
+
+@pytest.mark.parametrize("dtype", [ti.i32, ti.types.vector(3, ti.f32), ti.types.matrix(2, 2, ti.f32)])
+@test_utils.test()
+def test_ndarray_del(dtype) -> None:
+    def foo():
+        nd = ti.ndarray(ti.i32, (1000,))
+        assert nd.arr.shape != []
+        return nd.arr
+
+    arr = foo()
+    # after deletion, the shape is now empty, so this shows the arr was deleted
+    assert arr.shape == []