[Lang] Deprecate SOA-layout for NdarrayMatrix/NdarrayVector

python/taichi/aot/utils.py

@@ -1,6 +1,5 @@
 from taichi.lang._ndarray import ScalarNdarray
 from taichi.lang._texture import Texture
-from taichi.lang.enums import Layout
 from taichi.lang.exception import TaichiCompilationError
 from taichi.lang.matrix import Matrix, MatrixNdarray, MatrixType, VectorNdarray
 from taichi.lang.util import cook_dtype
@@ -77,15 +76,13 @@ def produce_injected_args(kernel, symbolic_args=None):
                 injected_args.append(
                     VectorNdarray(element_shape[0],
                                   dtype=dtype,
-                                  shape=(2, ) * field_dim,
-                                  layout=Layout.AOS))
+                                  shape=(2, ) * field_dim))
             elif element_dim == 2:
                 injected_args.append(
                     MatrixNdarray(element_shape[0],
                                   element_shape[1],
                                   dtype=dtype,
-                                  shape=(2, ) * field_dim,
-                                  layout=Layout.AOS))
+                                  shape=(2, ) * field_dim))
             else:
                 raise RuntimeError('')
         elif isinstance(anno, (TextureType, RWTextureType)):

python/taichi/lang/_ndarray.py

@@ -17,14 +17,14 @@ class Ndarray:
     """
     def __init__(self):
         self.host_accessor = None
-        self.layout = None
         self.shape = None
         self.element_type = None
         self.dtype = None
         self.arr = None
+        self.layout = Layout.AOS
 
     def get_type(self):
-        return NdarrayTypeMetadata(self.element_type, self.shape, self.layout)
+        return NdarrayTypeMetadata(self.element_type, self.shape)
 
     @property
     def element_shape(self):
@@ -91,7 +91,7 @@ def _ndarray_to_numpy(self):
         return arr
 
     @python_scope
-    def _ndarray_matrix_to_numpy(self, layout, as_vector):
+    def _ndarray_matrix_to_numpy(self, as_vector):
         """Converts matrix ndarray to a numpy array.
 
         Returns:
@@ -101,7 +101,7 @@ def _ndarray_matrix_to_numpy(self, layout, as_vector):
                        dtype=to_numpy_type(self.dtype))
         from taichi._kernels import \
             ndarray_matrix_to_ext_arr  # pylint: disable=C0415
-        layout_is_aos = 1 if layout == Layout.AOS else 0
+        layout_is_aos = 1
         ndarray_matrix_to_ext_arr(self, arr, layout_is_aos, as_vector)
         impl.get_runtime().sync()
         return arr
@@ -127,7 +127,7 @@ def _ndarray_from_numpy(self, arr):
         impl.get_runtime().sync()
 
     @python_scope
-    def _ndarray_matrix_from_numpy(self, arr, layout, as_vector):
+    def _ndarray_matrix_from_numpy(self, arr, as_vector):
         """Loads all values from a numpy array.
 
         Args:
@@ -144,7 +144,7 @@ def _ndarray_matrix_from_numpy(self, arr, layout, as_vector):
 
         from taichi._kernels import \
             ext_arr_to_ndarray_matrix  # pylint: disable=C0415
-        layout_is_aos = 1 if layout == Layout.AOS else 0
+        layout_is_aos = 1
         ext_arr_to_ndarray_matrix(arr, self, layout_is_aos, as_vector)
         impl.get_runtime().sync()
 
@@ -301,10 +301,7 @@ class NdarrayHostAccess:
     def __init__(self, arr, indices_first, indices_second):
         self.ndarr = arr
         self.arr = arr.arr
-        if arr.layout == Layout.SOA:
-            self.indices = indices_second + indices_first
-        else:
-            self.indices = indices_first + indices_second
+        self.indices = indices_first + indices_second
 
         def getter():
             self.ndarr._initialize_host_accessor()

python/taichi/lang/any_array.py

@@ -33,8 +33,8 @@ def layout(self):
     def get_type(self):
         return NdarrayTypeMetadata(
             self.ptr.get_ret_type(),
-            None,  # AnyArray can take any shape
-            self.layout())
+            None  # AnyArray can take any shape
+        )
 
     @property
     @taichi_scope

python/taichi/lang/impl.py

@@ -8,7 +8,7 @@
 from taichi.lang._ndarray import ScalarNdarray
 from taichi.lang._ndrange import GroupedNDRange, _Ndrange
 from taichi.lang.any_array import AnyArray, AnyArrayAccess
-from taichi.lang.enums import Layout, SNodeGradType
+from taichi.lang.enums import SNodeGradType
 from taichi.lang.exception import (TaichiRuntimeError, TaichiSyntaxError,
                                    TaichiTypeError)
 from taichi.lang.expr import Expr, make_expr_group
@@ -746,13 +746,12 @@ def field(dtype,
 
 
 @python_scope
-def ndarray(dtype, shape, layout=Layout.NULL):
+def ndarray(dtype, shape):
     """Defines a Taichi ndarray with scalar elements.
 
     Args:
         dtype (Union[DataType, MatrixType]): Data type of each element. This can be either a scalar type like ti.f32 or a compound type like ti.types.vector(3, ti.i32).
         shape (Union[int, tuple[int]]): Shape of the ndarray.
-        layout (Layout, optional): Layout of ndarray, only applicable when element is non-scalar type. Default is Layout.AOS.
 
     Example:
         The code below shows how a Taichi ndarray with scalar elements can be declared and defined::
@@ -761,16 +760,14 @@ def ndarray(dtype, shape, layout=Layout.NULL):
             >>> vec3 = ti.types.vector(3, ti.i32)
             >>> y = ti.ndarray(vec3, shape=(10, 2))  # ndarray of shape (10, 2), each element is a vector of 3 ti.i32 scalars.
             >>> matrix_ty = ti.types.matrix(3, 4, float)
-            >>> z = ti.ndarray(matrix_ty, shape=(4, 5), layout=ti.Layout.SOA)  # ndarray of shape (4, 5), each element is a matrix of (3, 4) ti.float scalars.
+            >>> z = ti.ndarray(matrix_ty, shape=(4, 5))  # ndarray of shape (4, 5), each element is a matrix of (3, 4) ti.float scalars.
     """
     if isinstance(shape, numbers.Number):
         shape = (shape, )
     if dtype in all_types:
-        assert layout == Layout.NULL
         return ScalarNdarray(dtype, shape)
     if isinstance(dtype, MatrixType):
-        layout = Layout.AOS if layout == Layout.NULL else layout
-        return MatrixNdarray(dtype.n, dtype.m, dtype.dtype, shape, layout)
+        return MatrixNdarray(dtype.n, dtype.m, dtype.dtype, shape)
 
     raise TaichiRuntimeError(
         f'{dtype} is not supported as ndarray element type')

python/taichi/lang/kernel_impl.py

@@ -355,24 +355,21 @@ def extract_arg(arg, anno):
                 return arg.dtype, len(arg.shape), (), Layout.AOS
             if isinstance(arg, taichi.lang.matrix.VectorNdarray):
                 anno.check_matched(arg.get_type())
-                return arg.dtype, len(arg.shape) + 1, (arg.n, ), arg.layout
+                return arg.dtype, len(arg.shape) + 1, (arg.n, ), Layout.AOS
             if isinstance(arg, taichi.lang.matrix.MatrixNdarray):
                 anno.check_matched(arg.get_type())
                 return arg.dtype, len(arg.shape) + 2, (arg.n,
-                                                       arg.m), arg.layout
+                                                       arg.m), Layout.AOS
             # external arrays
             element_dim = 0 if anno.element_dim is None else anno.element_dim
-            layout = Layout.AOS if anno.layout is None else anno.layout
             shape = tuple(arg.shape)
             if len(shape) < element_dim:
                 raise ValueError(
                     f"Invalid argument into ti.types.ndarray() - required element_dim={element_dim}, "
                     f"but the argument has only {len(shape)} dimensions")
-            element_shape = (
-            ) if element_dim == 0 else shape[:
-                                             element_dim] if layout == Layout.SOA else shape[
-                                                 -element_dim:]
-            return to_taichi_type(arg.dtype), len(shape), element_shape, layout
+            element_shape = () if element_dim == 0 else shape[-element_dim:]
+            return to_taichi_type(
+                arg.dtype), len(shape), element_shape, Layout.AOS
         if isinstance(anno, sparse_matrix_builder):
             return arg.dtype
         # Use '#' as a placeholder because other kinds of arguments are not involved in template instantiation

python/taichi/lang/matrix.py

@@ -1258,7 +1258,7 @@ def field(cls,
 
     @classmethod
     @python_scope
-    def ndarray(cls, n, m, dtype, shape, layout=Layout.AOS):
+    def ndarray(cls, n, m, dtype, shape):
         """Defines a Taichi ndarray with matrix elements.
         This function must be called in Python scope, and after `ti.init` is called.
 
@@ -1267,7 +1267,6 @@ def ndarray(cls, n, m, dtype, shape, layout=Layout.AOS):
             m (int): Number of columns of the matrix.
             dtype (DataType): Data type of each value.
             shape (Union[int, tuple[int]]): Shape of the ndarray.
-            layout (Layout, optional): Memory layout, AOS by default.
 
         Example::
 
@@ -1278,7 +1277,7 @@ def ndarray(cls, n, m, dtype, shape, layout=Layout.AOS):
         """
         if isinstance(shape, numbers.Number):
             shape = (shape, )
-        return MatrixNdarray(n, m, dtype, shape, layout)
+        return MatrixNdarray(n, m, dtype, shape)
 
     @staticmethod
     def rows(rows):
@@ -1457,7 +1456,7 @@ def field(cls, n, dtype, *args, **kwargs):
 
     @classmethod
     @python_scope
-    def ndarray(cls, n, dtype, shape, layout=Layout.AOS):
+    def ndarray(cls, n, dtype, shape):
         """Defines a Taichi ndarray with vector elements.
 
         Args:
@@ -1473,7 +1472,7 @@ def ndarray(cls, n, dtype, shape, layout=Layout.AOS):
         """
         if isinstance(shape, numbers.Number):
             shape = (shape, )
-        return VectorNdarray(n, dtype, shape, layout)
+        return VectorNdarray(n, dtype, shape)
 
 
 class _IntermediateMatrix(Matrix):
@@ -1902,33 +1901,32 @@ class MatrixNdarray(Ndarray):
         m (int): Number of columns of the matrix.
         dtype (DataType): Data type of each value.
         shape (Union[int, tuple[int]]): Shape of the ndarray.
-        layout (Layout): Memory layout.
 
     Example::
 
-        >>> arr = ti.MatrixNdarray(2, 2, ti.f32, shape=(3, 3), layout=Layout.SOA)
+        >>> arr = ti.MatrixNdarray(2, 2, ti.f32, shape=(3, 3))
     """
-    def __init__(self, n, m, dtype, shape, layout):
+    def __init__(self, n, m, dtype, shape):
         self.n = n
         self.m = m
         super().__init__()
         # TODO(zhanlue): remove self.dtype and migrate its usages to element_type
         self.dtype = cook_dtype(dtype)
 
-        self.layout = layout
+        self.layout = Layout.AOS
         self.shape = tuple(shape)
         self.element_type = TensorType((self.n, self.m), dtype)
         # TODO: we should pass in element_type, shape, layout instead.
         self.arr = impl.get_runtime().prog.create_ndarray(
-            cook_dtype(self.element_type.ptr), shape, layout)
+            cook_dtype(self.element_type.ptr), shape, Layout.AOS)
 
     @property
     def element_shape(self):
         """Returns the shape of each element (a 2D matrix) in this ndarray.
 
         Example::
 
-            >>> arr = ti.MatrixNdarray(2, 2, ti.f32, shape=(3, 3), layout=Layout.SOA)
+            >>> arr = ti.MatrixNdarray(2, 2, ti.f32, shape=(3, 3))
             >>> arr.element_shape
             (2, 2)
         """
@@ -1958,15 +1956,15 @@ def to_numpy(self):
 
         Example::
 
-            >>> arr = ti.MatrixNdarray(2, 2, ti.f32, shape=(2, 1), layout=Layout.SOA)
+            >>> arr = ti.MatrixNdarray(2, 2, ti.f32, shape=(2, 1))
             >>> arr.to_numpy()
             [[[[0. 0.]
                [0. 0.]]]
 
              [[[0. 0.]
                [0. 0.]]]]
         """
-        return self._ndarray_matrix_to_numpy(self.layout, as_vector=0)
+        return self._ndarray_matrix_to_numpy(as_vector=0)
 
     @python_scope
     def from_numpy(self, arr):
@@ -1978,12 +1976,11 @@ def from_numpy(self, arr):
             >>> arr = np.ones((2, 1, 2, 2))
             >>> m.from_numpy(arr)
         """
-        self._ndarray_matrix_from_numpy(arr, self.layout, as_vector=0)
+        self._ndarray_matrix_from_numpy(arr, as_vector=0)
 
     @python_scope
     def __deepcopy__(self, memo=None):
-        ret_arr = MatrixNdarray(self.n, self.m, self.dtype, self.shape,
-                                self.layout)
+        ret_arr = MatrixNdarray(self.n, self.m, self.dtype, self.shape)
         ret_arr.copy_from(self)
         return ret_arr
 
@@ -1995,7 +1992,7 @@ def _fill_by_kernel(self, val):
 
     @python_scope
     def __repr__(self):
-        return f'<{self.n}x{self.m} {self.layout} ti.Matrix.ndarray>'
+        return f'<{self.n}x{self.m} {Layout.AOS} ti.Matrix.ndarray>'
 
 
 class VectorNdarray(Ndarray):
@@ -2009,27 +2006,27 @@ class VectorNdarray(Ndarray):
 
     Example::
 
-        >>> a = ti.VectorNdarray(3, ti.f32, (3, 3), layout=Layout.SOA)
+        >>> a = ti.VectorNdarray(3, ti.f32, (3, 3))
     """
-    def __init__(self, n, dtype, shape, layout):
+    def __init__(self, n, dtype, shape):
         self.n = n
         super().__init__()
         # TODO(zhanlue): remove self.dtype and migrate its usages to element_type
         self.dtype = cook_dtype(dtype)
 
-        self.layout = layout
+        self.layout = Layout.AOS
         self.shape = tuple(shape)
         self.element_type = TensorType((n, ), self.dtype)
         self.arr = impl.get_runtime().prog.create_ndarray(
-            cook_dtype(self.element_type.ptr), shape, layout)
+            cook_dtype(self.element_type.ptr), shape, Layout.AOS)
 
     @property
     def element_shape(self):
         """Gets the dimension of the vector of this ndarray.
 
         Example::
 
-            >>> a = ti.VectorNdarray(3, ti.f32, (3, 3), layout=Layout.SOA)
+            >>> a = ti.VectorNdarray(3, ti.f32, (3, 3))
             >>> a.element_shape
             (3,)
         """
@@ -2055,15 +2052,15 @@ def to_numpy(self):
 
         Example::
 
-            >>> a = ti.VectorNdarray(3, ti.f32, (2, 2), layout=Layout.SOA)
+            >>> a = ti.VectorNdarray(3, ti.f32, (2, 2))
             >>> a.to_numpy()
             array([[[0., 0., 0.],
                     [0., 0., 0.]],
 
                    [[0., 0., 0.],
                     [0., 0., 0.]]], dtype=float32)
         """
-        return self._ndarray_matrix_to_numpy(self.layout, as_vector=1)
+        return self._ndarray_matrix_to_numpy(as_vector=1)
 
     @python_scope
     def from_numpy(self, arr):
@@ -2078,11 +2075,11 @@ def from_numpy(self, arr):
             >>> b = np.ones((2, 2, 3), dtype=np.float32)
             >>> a.from_numpy(b)
         """
-        self._ndarray_matrix_from_numpy(arr, self.layout, as_vector=1)
+        self._ndarray_matrix_from_numpy(arr, as_vector=1)
 
     @python_scope
     def __deepcopy__(self, memo=None):
-        ret_arr = VectorNdarray(self.n, self.dtype, self.shape, self.layout)
+        ret_arr = VectorNdarray(self.n, self.dtype, self.shape)
         ret_arr.copy_from(self)
         return ret_arr
 
@@ -2094,7 +2091,7 @@ def _fill_by_kernel(self, val):
 
     @python_scope
     def __repr__(self):
-        return f'<{self.n} {self.layout} ti.Vector.ndarray>'
+        return f'<{self.n} {Layout.AOS} ti.Vector.ndarray>'
 
 
 __all__ = ["Matrix", "Vector", "MatrixField", "MatrixNdarray", "VectorNdarray"]