[refactor] Specialized Ndarray Type is (element_type, shape, layout)

ghstack-source-id: 977cd453359b8ccc09deccacc62a915abcd42734
Pull Request resolved: #5065

Author: ailzhang

python/taichi/aot/utils.py

@@ -35,29 +35,33 @@ def produce_injected_args(kernel, symbolic_args=None):
                 raise TaichiCompilationError(
                     f'Expected Ndaray type, got {anno}')
             if symbolic_args is not None:
-                anno.element_shape = tuple(symbolic_args[i].element_shape)
-                anno.element_dim = len(anno.element_shape)
-                anno.dtype = symbolic_args[i].dtype()
+                element_shape = tuple(symbolic_args[i].element_shape)
+                element_dim = len(element_shape)
+                dtype = symbolic_args[i].dtype()
+            else:
+                element_shape = anno.element_shape
+                element_dim = anno.element_dim
+                dtype = anno.dtype
 
-            if anno.element_shape is None or anno.field_dim is None:
+            if element_shape is None or anno.field_dim is None:
                 raise TaichiCompilationError(
                     'Please either specify both `element_shape` and `field_dim` '
                     'in the param annotation, or provide an example '
                     f'ndarray for param={arg.name}')
-            if anno.element_dim == 0:
+            if element_dim is None or element_dim == 0:
                 injected_args.append(
-                    ScalarNdarray(anno.dtype, (2, ) * anno.field_dim))
-            elif anno.element_dim == 1:
+                    ScalarNdarray(dtype, (2, ) * anno.field_dim))
+            elif element_dim == 1:
                 injected_args.append(
-                    VectorNdarray(anno.element_shape[0],
-                                  dtype=anno.dtype,
+                    VectorNdarray(element_shape[0],
+                                  dtype=dtype,
                                   shape=(2, ) * anno.field_dim,
                                   layout=Layout.AOS))
-            elif anno.element_dim == 2:
+            elif element_dim == 2:
                 injected_args.append(
-                    MatrixNdarray(anno.element_shape[0],
-                                  anno.element_shape[1],
-                                  dtype=anno.dtype,
+                    MatrixNdarray(element_shape[0],
+                                  element_shape[1],
+                                  dtype=dtype,
                                   shape=(2, ) * anno.field_dim,
                                   layout=Layout.AOS))
             else:

python/taichi/lang/_ndarray.py

@@ -4,6 +4,7 @@
 from taichi.lang.enums import Layout
 from taichi.lang.util import cook_dtype, python_scope, to_numpy_type
 from taichi.types import primitive_types
+from taichi.types.ndarray_type import SpecializeNdarrayType
 
 
 class Ndarray:
@@ -15,10 +16,17 @@ class Ndarray:
     """
     def __init__(self, dtype, arr_shape):
         self.host_accessor = None
+        self.layout = None
+        self.shape = None
+        self.element_type = None
         self.dtype = cook_dtype(dtype)
         self.arr = impl.get_runtime().prog.create_ndarray(
             cook_dtype(dtype), arr_shape)
 
+    def get_type(self):
+        return SpecializeNdarrayType(self.element_type, self.shape,
+                                     self.layout)
+
     @property
     def element_shape(self):
         """Gets ndarray element shape.
@@ -209,6 +217,7 @@ class ScalarNdarray(Ndarray):
     def __init__(self, dtype, arr_shape):
         super().__init__(dtype, arr_shape)
         self.shape = tuple(self.arr.shape)
+        self.element_type = dtype
 
     @property
     def element_shape(self):

python/taichi/lang/kernel_impl.py

@@ -337,21 +337,13 @@ def extract_arg(arg, anno):
             return arg
         if isinstance(anno, ndarray_type.NdarrayType):
             if isinstance(arg, taichi.lang._ndarray.ScalarNdarray):
-                anno._check_element_dim(arg, 0)
-                anno._check_element_shape(())
-                anno._check_field_dim(len(arg.shape))
+                anno.match(arg.get_type())
                 return arg.dtype, len(arg.shape), (), Layout.AOS
             if isinstance(arg, taichi.lang.matrix.VectorNdarray):
-                anno._check_element_dim(arg, 1)
-                anno._check_element_shape((arg.n, ))
-                anno._check_field_dim(len(arg.shape))
-                anno._check_layout(arg)
+                anno.match(arg.get_type())
                 return arg.dtype, len(arg.shape) + 1, (arg.n, ), arg.layout
             if isinstance(arg, taichi.lang.matrix.MatrixNdarray):
-                anno._check_element_dim(arg, 2)
-                anno._check_element_shape((arg.n, arg.m))
-                anno._check_field_dim(len(arg.shape))
-                anno._check_layout(arg)
+                anno.match(arg.get_type())
                 return arg.dtype, len(arg.shape) + 2, (arg.n,
                                                        arg.m), arg.layout
             # external arrays

python/taichi/lang/matrix.py

@@ -17,7 +17,7 @@
                               taichi_scope, to_numpy_type, to_paddle_type,
                               to_pytorch_type, warning)
 from taichi.types import primitive_types
-from taichi.types.compound_types import CompoundType
+from taichi.types.compound_types import CompoundType, TensorType
 
 
 def _gen_swizzles(cls):
@@ -1688,12 +1688,14 @@ class MatrixNdarray(Ndarray):
         >>> arr = ti.MatrixNdarray(2, 2, ti.f32, shape=(3, 3), layout=Layout.SOA)
     """
     def __init__(self, n, m, dtype, shape, layout):
-        self.layout = layout
-        self.shape = shape
         self.n = n
         self.m = m
+        # TODO: we should pass in element_type, shape, layout instead.
         arr_shape = (n, m) + shape if layout == Layout.SOA else shape + (n, m)
         super().__init__(dtype, arr_shape)
+        self.layout = layout
+        self.shape = shape
+        self.element_type = TensorType((self.n, self.m), dtype)
 
     @property
     def element_shape(self):
@@ -1783,11 +1785,13 @@ class VectorNdarray(Ndarray):
         >>> a = ti.VectorNdarray(3, ti.f32, (3, 3), layout=Layout.SOA)
     """
     def __init__(self, n, dtype, shape, layout):
-        self.layout = layout
-        self.shape = shape
         self.n = n
+        # TODO: pass in element_type, shape, layout directly
         arr_shape = (n, ) + shape if layout == Layout.SOA else shape + (n, )
         super().__init__(dtype, arr_shape)
+        self.layout = layout
+        self.shape = shape
+        self.element_type = TensorType((n, ), dtype)
 
     @property
     def element_shape(self):

python/taichi/types/compound_types.py

@@ -5,6 +5,12 @@ class CompoundType:
     pass
 
 
+class TensorType(CompoundType):
+    def __init__(self, shape, dtype):
+        self.dtype = taichi.lang.util.cook_dtype(dtype)
+        self.shape = shape
+
+
 # TODO: maybe move MatrixType, StructType here to avoid the circular import?
 def matrix(n, m, dtype):
     """Creates a matrix type with given shape and data type.

python/taichi/types/ndarray_type.py

@@ -1,6 +1,13 @@
 from taichi.types.primitive_types import f32
 
 
+class SpecializeNdarrayType:
+    def __init__(self, element_type, shape=None, layout=None):
+        self.element_type = element_type
+        self.shape = shape
+        self.layout = layout
+
+
 class NdarrayType:
     """Type annotation for arbitrary arrays, including external arrays (numpy ndarrays and torch tensors) and Taichi ndarrays.
 
@@ -32,31 +39,31 @@ def __init__(self,
         self.element_shape = element_shape
         self.element_dim = len(
             element_shape) if element_shape is not None else element_dim
+
         self.field_dim = field_dim
         self.layout = layout
 
-    def _check_element_dim(self, arg, arg_dim):
-        if self.element_dim is not None and self.element_dim != arg_dim:
+    def match(self, ndarray_type: SpecializeNdarrayType):
+        if self.element_dim is not None and self.element_dim != len(
+                ndarray_type.element_type.shape):
             raise ValueError(
-                f"Invalid argument into ti.types.ndarray() - required element_dim={self.element_dim}, but {arg} is provided"
+                f"Invalid argument into ti.types.ndarray() - required element_dim={self.element_dim}, but {len(ndarray_type.element_type.shape)} is provided"
             )
 
-    def _check_layout(self, arg):
-        if self.layout is not None and self.layout != arg.layout:
+        if self.element_shape is not None and self.element_shape != ndarray_type.element_type.shape:
             raise ValueError(
-                f"Invalid argument into ti.types.ndarray() - required layout={self.layout}, but {arg} is provided"
+                f"Invalid argument into ti.types.ndarray() - required element_shape={self.element_shape}, but {ndarray_type.element_type.shape} is provided"
             )
 
-    def _check_element_shape(self, shapes):
-        if self.element_shape is not None and shapes != self.element_shape:
+        if self.layout is not None and self.layout != ndarray_type.layout:
             raise ValueError(
-                f"Invalid argument into ti.types.ndarray() - required element_shape={self.element_shape}, but {shapes} is provided"
+                f"Invalid argument into ti.types.ndarray() - required layout={self.layout}, but {ndarray_type.layout} is provided"
             )
 
-    def _check_field_dim(self, field_dim):
-        if self.field_dim is not None and field_dim != self.field_dim:
+        if self.field_dim is not None and self.field_dim != len(
+                ndarray_type.shape):
             raise ValueError(
-                f"Invalid argument into ti.types.ndarray() - required field_dim={self.field_dim}, but {field_dim} is provided"
+                f"Invalid argument into ti.types.ndarray() - required field_dim={self.field_dim}, but {ndarray_type.element_type} is provided"
             )
 
 

tests/python/test_aot.py

@@ -84,15 +84,14 @@ def test_aot_bind_id():
     density1 = ti.ndarray(dtype=ti.f32, shape=(8, 8))
 
     @ti.kernel
-    def init(x: ti.f32, density1: ti.types.ndarray(field_dim=2,
-                                                   element_shape=())):
+    def init(x: ti.f32, density1: ti.types.ndarray(field_dim=2)):
         for i, j in density1:
             density[i, j] = x
             density1[i, j] = x + 1
 
     with tempfile.TemporaryDirectory() as tmpdir:
         m = ti.aot.Module(ti.lang.impl.current_cfg().arch)
-        m.add_kernel(init)
+        m.add_kernel(init, {'density1': density1})
         m.save(tmpdir, '')
         with open(os.path.join(tmpdir, 'metadata.json')) as json_file:
             res = json.load(json_file)

tests/python/test_api.py

@@ -113,9 +113,9 @@ def _get_expected_matrix_apis():
     'to_torch'
 ]
 user_api[ti.MatrixNdarray] = [
-    'copy_from', 'element_shape', 'fill', 'from_numpy', 'to_numpy'
+    'copy_from', 'element_shape', 'fill', 'from_numpy', 'get_type', 'to_numpy'
 ]
-user_api[ti.Ndarray] = ['copy_from', 'element_shape', 'fill']
+user_api[ti.Ndarray] = ['copy_from', 'element_shape', 'fill', 'get_type']
 user_api[ti.SNode] = [
     'bit_array', 'bit_struct', 'bitmasked', 'deactivate_all', 'dense',
     'dynamic', 'lazy_grad', 'parent', 'place', 'pointer', 'shape'
@@ -125,7 +125,7 @@ def _get_expected_matrix_apis():
     'parent', 'shape', 'snode', 'to_numpy', 'to_paddle', 'to_torch'
 ]
 user_api[ti.ScalarNdarray] = [
-    'copy_from', 'element_shape', 'fill', 'from_numpy', 'to_numpy'
+    'copy_from', 'element_shape', 'fill', 'from_numpy', 'get_type', 'to_numpy'
 ]
 user_api[ti.Struct] = ['field', 'fill', 'items', 'keys', 'to_dict']
 user_api[ti.StructField] = [
@@ -134,7 +134,7 @@ def _get_expected_matrix_apis():
     'to_paddle', 'to_torch'
 ]
 user_api[ti.VectorNdarray] = [
-    'copy_from', 'element_shape', 'fill', 'from_numpy', 'to_numpy'
+    'copy_from', 'element_shape', 'fill', 'from_numpy', 'get_type', 'to_numpy'
 ]
 
 

tests/python/test_graph.py

@@ -9,7 +9,7 @@ def test_ndarray_int():
     n = 4
 
     @ti.kernel
-    def test(pos: ti.types.ndarray(field_dim=1, element_shape=())):
+    def test(pos: ti.types.ndarray(field_dim=1)):
         for i in range(n):
             pos[i] = 1