[Lang] Add parameter 'order' to specify layout for scalar, vector, ma…

…trix fields (#5617)

* [Lang] Add parameter 'order' to specify layout for scalar, vector, matrix fields

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix format

* Fix lowercase

* Nit

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

python/taichi/lang/impl.py

@@ -1,6 +1,6 @@
 import numbers
 from types import FunctionType, MethodType
-from typing import Iterable
+from typing import Iterable, Sequence
 
 import numpy as np
 from taichi._lib import core as _ti_core
@@ -9,7 +9,8 @@
 from taichi.lang._ndrange import GroupedNDRange, _Ndrange
 from taichi.lang.any_array import AnyArray, AnyArrayAccess
 from taichi.lang.enums import Layout
-from taichi.lang.exception import TaichiRuntimeError, TaichiTypeError
+from taichi.lang.exception import (TaichiRuntimeError, TaichiSyntaxError,
+                                   TaichiTypeError)
 from taichi.lang.expr import Expr, make_expr_group
 from taichi.lang.field import Field, ScalarField
 from taichi.lang.kernel_arguments import SparseMatrixProxy
@@ -526,6 +527,19 @@ def __repr__(self):
 """
 
 
+def _create_snode(axis_seq: Sequence[int], shape_seq: Sequence[numbers.Number],
+                  same_level: bool):
+    dim = len(axis_seq)
+    assert dim == len(shape_seq)
+    snode = root
+    if same_level:
+        snode = snode.dense(axes(*axis_seq), shape_seq)
+    else:
+        for i in range(dim):
+            snode = snode.dense(axes(axis_seq[i]), (shape_seq[i], ))
+    return snode
+
+
 @python_scope
 def create_field_member(dtype, name, needs_grad, needs_dual):
     dtype = cook_dtype(dtype)
@@ -576,6 +590,7 @@ def create_field_member(dtype, name, needs_grad, needs_dual):
 @python_scope
 def field(dtype,
           shape=None,
+          order=None,
           name="",
           offset=None,
           needs_grad=False,
@@ -592,6 +607,7 @@ def field(dtype,
     Args:
         dtype (DataType): data type of the field.
         shape (Union[int, tuple[int]], optional): shape of the field.
+        order (str, optional): order of the shape laid out in memory.
         name (str, optional): name of the field.
         offset (Union[int, tuple[int]], optional): offset of the field domain.
         needs_grad (bool, optional): whether this field participates in autodiff (reverse mode)
@@ -604,41 +620,64 @@ def field(dtype,
         The code below shows how a Taichi field can be declared and defined::
 
             >>> x1 = ti.field(ti.f32, shape=(16, 8))
-            >>>
             >>> # Equivalently
             >>> x2 = ti.field(ti.f32)
             >>> ti.root.dense(ti.ij, shape=(16, 8)).place(x2)
-    """
-
-    if isinstance(shape, numbers.Number):
-        shape = (shape, )
-
-    if isinstance(offset, numbers.Number):
-        offset = (offset, )
-
-    if shape is not None and offset is not None:
-        assert len(shape) == len(
-            offset
-        ), f'The dimensionality of shape and offset must be the same  ({len(shape)} != {len(offset)})'
-
-    assert (offset is None or shape
-            is not None), 'The shape cannot be None when offset is being set'
+            >>>
+            >>> x3 = ti.field(ti.f32, shape=(16, 8), order='ji')
+            >>> # Equivalently
+            >>> x4 = ti.field(ti.f32)
+            >>> ti.root.dense(ti.j, shape=8).dense(ti.i, shape=16).place(x4)
 
+    """
     x, x_grad, x_dual = create_field_member(dtype, name, needs_grad,
                                             needs_dual)
     x, x_grad, x_dual = ScalarField(x), ScalarField(x_grad), ScalarField(
         x_dual)
-
     x._set_grad(x_grad)
     x._set_dual(x_dual)
 
-    if shape is not None:
+    if shape is None:
+        if offset is not None:
+            raise TaichiSyntaxError('shape cannot be None when offset is set')
+        if order is not None:
+            raise TaichiSyntaxError('shape cannot be None when order is set')
+    else:
+        if isinstance(shape, numbers.Number):
+            shape = (shape, )
+        if isinstance(offset, numbers.Number):
+            offset = (offset, )
         dim = len(shape)
-        root.dense(index_nd(dim), shape).place(x, offset=offset)
+        if offset is not None and dim != len(offset):
+            raise TaichiSyntaxError(
+                f'The dimensionality of shape and offset must be the same ({dim} != {len(offset)})'
+            )
+        axis_seq = []
+        shape_seq = []
+        if order is not None:
+            if dim != len(order):
+                raise TaichiSyntaxError(
+                    f'The dimensionality of shape and order must be the same ({dim} != {len(order)})'
+                )
+            if dim != len(set(order)):
+                raise TaichiSyntaxError('The axes in order must be different')
+            for ch in order:
+                axis = ord(ch) - ord('i')
+                if axis < 0 or axis >= dim:
+                    raise TaichiSyntaxError(f'Invalid axis {ch}')
+                axis_seq.append(axis)
+                shape_seq.append(shape[axis])
+        else:
+            axis_seq = list(range(dim))
+            shape_seq = list(shape)
+        same_level = order is None
+        _create_snode(axis_seq, shape_seq, same_level).place(x, offset=offset)
         if needs_grad:
-            root.dense(index_nd(dim), shape).place(x_grad)
+            _create_snode(axis_seq, shape_seq, same_level).place(x_grad,
+                                                                 offset=offset)
         if needs_dual:
-            root.dense(index_nd(dim), shape).place(x_dual)
+            _create_snode(axis_seq, shape_seq, same_level).place(x_dual,
+                                                                 offset=offset)
     return x
 
 

python/taichi/lang/matrix.py

@@ -1069,6 +1069,7 @@ def field(cls,
               m,
               dtype,
               shape=None,
+              order=None,
               name="",
               offset=None,
               needs_grad=False,
@@ -1081,6 +1082,7 @@ def field(cls,
             m (int): The desired number of columns of the Matrix.
             dtype (DataType, optional): The desired data type of the Matrix.
             shape (Union[int, tuple of int], optional): The desired shape of the Matrix.
+            order (str, optional): order of the shape laid out in memory.
             name (string, optional): The custom name of the field.
             offset (Union[int, tuple of int], optional): The coordinate offset
                 of all elements in a field.
@@ -1136,43 +1138,68 @@ def field(cls,
         impl.get_runtime().matrix_fields.append(entries)
 
         if shape is None:
-            assert offset is None, "shape cannot be None when offset is being set"
-
-        if shape is not None:
+            if offset is not None:
+                raise TaichiSyntaxError(
+                    'shape cannot be None when offset is set')
+            if order is not None:
+                raise TaichiSyntaxError(
+                    'shape cannot be None when order is set')
+        else:
             if isinstance(shape, numbers.Number):
                 shape = (shape, )
             if isinstance(offset, numbers.Number):
                 offset = (offset, )
-
-            if offset is not None:
-                assert len(shape) == len(
-                    offset
-                ), f'The dimensionality of shape and offset must be the same  ({len(shape)} != {len(offset)})'
-
             dim = len(shape)
+            if offset is not None and dim != len(offset):
+                raise TaichiSyntaxError(
+                    f'The dimensionality of shape and offset must be the same ({dim} != {len(offset)})'
+                )
+            axis_seq = []
+            shape_seq = []
+            if order is not None:
+                if dim != len(order):
+                    raise TaichiSyntaxError(
+                        f'The dimensionality of shape and order must be the same ({dim} != {len(order)})'
+                    )
+                if dim != len(set(order)):
+                    raise TaichiSyntaxError(
+                        'The axes in order must be different')
+                for ch in order:
+                    axis = ord(ch) - ord('i')
+                    if axis < 0 or axis >= dim:
+                        raise TaichiSyntaxError(f'Invalid axis {ch}')
+                    axis_seq.append(axis)
+                    shape_seq.append(shape[axis])
+            else:
+                axis_seq = list(range(dim))
+                shape_seq = list(shape)
+            same_level = order is None
             if layout == Layout.SOA:
                 for e in entries._get_field_members():
-                    impl.root.dense(impl.index_nd(dim),
-                                    shape).place(ScalarField(e), offset=offset)
+                    impl._create_snode(axis_seq, shape_seq,
+                                       same_level).place(ScalarField(e),
+                                                         offset=offset)
                 if needs_grad:
                     for e in entries_grad._get_field_members():
-                        impl.root.dense(impl.index_nd(dim),
-                                        shape).place(ScalarField(e),
-                                                     offset=offset)
+                        impl._create_snode(axis_seq, shape_seq,
+                                           same_level).place(ScalarField(e),
+                                                             offset=offset)
                 if needs_dual:
                     for e in entries_dual._get_field_members():
-                        impl.root.dense(impl.index_nd(dim),
-                                        shape).place(ScalarField(e),
-                                                     offset=offset)
+                        impl._create_snode(axis_seq, shape_seq,
+                                           same_level).place(ScalarField(e),
+                                                             offset=offset)
             else:
-                impl.root.dense(impl.index_nd(dim), shape).place(entries,
-                                                                 offset=offset)
+                impl._create_snode(axis_seq, shape_seq,
+                                   same_level).place(entries, offset=offset)
                 if needs_grad:
-                    impl.root.dense(impl.index_nd(dim),
-                                    shape).place(entries_grad, offset=offset)
+                    impl._create_snode(axis_seq, shape_seq,
+                                       same_level).place(entries_grad,
+                                                         offset=offset)
                 if needs_dual:
-                    impl.root.dense(impl.index_nd(dim),
-                                    shape).place(entries_dual, offset=offset)
+                    impl._create_snode(axis_seq, shape_seq,
+                                       same_level).place(entries_dual,
+                                                         offset=offset)
         return entries
 
     @classmethod

tests/python/test_offset.py

@@ -1,4 +1,5 @@
 import pytest
+from taichi.lang.misc import get_host_arch_list
 
 import taichi as ti
 from tests import test_utils
@@ -116,26 +117,38 @@ def test():
             assert a[i, j][0, 0] == i + j
 
 
-@test_utils.test()
-def test_offset_must_throw_var():
-    with pytest.raises(AssertionError):
-        a = ti.field(dtype=ti.float32, shape=3, offset=(3, 4))
-        b = ti.field(dtype=ti.float32, shape=None, offset=(3, 4))
+@test_utils.test(arch=get_host_arch_list())
+def test_offset_must_throw_scalar():
+    with pytest.raises(
+            ti.TaichiCompilationError,
+            match='The dimensionality of shape and offset must be the same'):
+        a = ti.field(dtype=ti.f32, shape=3, offset=(3, 4))
+    with pytest.raises(ti.TaichiCompilationError,
+                       match='shape cannot be None when offset is set'):
+        b = ti.field(dtype=ti.f32, shape=None, offset=(3, 4))
 
 
-@test_utils.test()
+@test_utils.test(arch=get_host_arch_list())
 def test_offset_must_throw_vector():
-    with pytest.raises(AssertionError):
-        a = ti.Vector.field(3, dtype=ti.float32, shape=3, offset=(3, 4))
-        b = ti.Vector.field(3, dtype=ti.float32, shape=None, offset=(3, ))
+    with pytest.raises(
+            ti.TaichiCompilationError,
+            match='The dimensionality of shape and offset must be the same'):
+        a = ti.Vector.field(3, dtype=ti.f32, shape=3, offset=(3, 4))
+    with pytest.raises(ti.TaichiCompilationError,
+                       match='shape cannot be None when offset is set'):
+        b = ti.Vector.field(3, dtype=ti.f32, shape=None, offset=(3, ))
 
 
-@test_utils.test()
+@test_utils.test(arch=get_host_arch_list())
 def test_offset_must_throw_matrix():
-    with pytest.raises(AssertionError):
-        c = ti.Matrix.field(3,
+    with pytest.raises(
+            ti.TaichiCompilationError,
+            match='The dimensionality of shape and offset must be the same'):
+        a = ti.Matrix.field(3,
                             3,
                             dtype=ti.i32,
                             shape=(32, 16, 8),
                             offset=(32, 16))
-        d = ti.Matrix.field(3, 3, dtype=ti.i32, shape=None, offset=(32, 16))
+    with pytest.raises(ti.TaichiCompilationError,
+                       match='shape cannot be None when offset is set'):
+        b = ti.Matrix.field(3, 3, dtype=ti.i32, shape=None, offset=(32, 16))