diff --git a/python/taichi/lang/impl.py b/python/taichi/lang/impl.py index b4c714c46eaa8..3cd2f20738184 100644 --- a/python/taichi/lang/impl.py +++ b/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 diff --git a/python/taichi/lang/matrix.py b/python/taichi/lang/matrix.py index 704d926f1fe9d..629751e3412fb 100644 --- a/python/taichi/lang/matrix.py +++ b/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 diff --git a/tests/python/test_offset.py b/tests/python/test_offset.py index fd41b732a0f12..4152b65bbf4ab 100644 --- a/tests/python/test_offset.py +++ b/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)) diff --git a/tests/python/test_order.py b/tests/python/test_order.py new file mode 100644 index 0000000000000..940749659627d --- /dev/null +++ b/tests/python/test_order.py @@ -0,0 +1,153 @@ +import pytest +from taichi.lang.misc import get_host_arch_list + +import taichi as ti +from tests import test_utils + + +@test_utils.test(arch=get_host_arch_list()) +def test_order_scalar(): + X = 16 + Y = 8 + Z = 4 + S = 4 + + a = ti.field(ti.i32, shape=(X, Y, Z), order='ijk') + b = ti.field(ti.i32, shape=(X, Y, Z), order='ikj') + c = ti.field(ti.i32, shape=(X, Y, Z), order='jik') + d = ti.field(ti.i32, shape=(X, Y, Z), order='jki') + e = ti.field(ti.i32, shape=(X, Y, Z), order='kij') + f = ti.field(ti.i32, shape=(X, Y, Z), order='kji') + + @ti.kernel + def fill(): + for i, j, k in b: + a[i, j, k] = i * j * k + b[i, j, k] = i * j * k + c[i, j, k] = i * j * k + d[i, j, k] = i * j * k + e[i, j, k] = i * j * k + f[i, j, k] = i * j * k + + @ti.kernel + def get_field_addr(a: ti.template(), i: ti.i32, j: ti.i32, + k: ti.i32) -> ti.u64: + return ti.get_addr(a, [i, j, k]) + + fill() + + a_addr = get_field_addr(a, 0, 0, 0) + b_addr = get_field_addr(b, 0, 0, 0) + c_addr = get_field_addr(c, 0, 0, 0) + d_addr = get_field_addr(d, 0, 0, 0) + e_addr = get_field_addr(e, 0, 0, 0) + f_addr = get_field_addr(f, 0, 0, 0) + for i in range(X): + for j in range(Y): + for k in range(Z): + assert a[i, j, k] == b[i, j, k] == c[i, j, k] == i * j * k + assert d[i, j, k] == e[i, j, k] == f[i, j, k] == i * j * k + assert a_addr + (i * (Y * Z) + j * Z + k) * S == \ + get_field_addr(a, i, j, k) + assert b_addr + (i * (Z * Y) + k * Y + j) * S == \ + get_field_addr(b, i, j, k) + assert c_addr + (j * (X * Z) + i * Z + k) * S == \ + get_field_addr(c, i, j, k) + assert d_addr + (j * (Z * X) + k * X + i) * S == \ + get_field_addr(d, i, j, k) + assert e_addr + (k * (X * Y) + i * Y + j) * S == \ + get_field_addr(e, i, j, k) + assert f_addr + (k * (Y * X) + j * X + i) * S == \ + get_field_addr(f, i, j, k) + + +@test_utils.test(arch=get_host_arch_list()) +def test_order_vector(): + X = 4 + Y = 2 + Z = 2 + S = 4 + + a = ti.Vector.field(Z, + ti.i32, + shape=(X, Y), + order='ij', + layout=ti.Layout.AOS) + b = ti.Vector.field(Z, + ti.i32, + shape=(X, Y), + order='ji', + layout=ti.Layout.AOS) + c = ti.Vector.field(Z, + ti.i32, + shape=(X, Y), + order='ij', + layout=ti.Layout.SOA) + d = ti.Vector.field(Z, + ti.i32, + shape=(X, Y), + order='ji', + layout=ti.Layout.SOA) + + @ti.kernel + def fill(): + for i, j in b: + a[i, j] = [i, j] + b[i, j] = [i, j] + c[i, j] = [i, j] + d[i, j] = [i, j] + + @ti.kernel + def get_field_addr(a: ti.template(), i: ti.i32, j: ti.i32) -> ti.u64: + return ti.get_addr(a, [i, j]) + + fill() + + a_addr = get_field_addr(a, 0, 0) + b_addr = get_field_addr(b, 0, 0) + c_addr = get_field_addr(c, 0, 0) + d_addr = get_field_addr(d, 0, 0) + for i in range(X): + for j in range(Y): + assert a[i, j] == b[i, j] == c[i, j] == d[i, j] == [i, j] + for k in range(Z): + assert a_addr + (i * (Y * Z) + j * Z + k) * S == \ + get_field_addr(a.get_scalar_field(k), i, j) + assert b_addr + (j * (X * Z) + i * Z + k) * S == \ + get_field_addr(b.get_scalar_field(k), i, j) + assert c_addr + (k * (X * Y) + i * Y + j) * S == \ + get_field_addr(c.get_scalar_field(k), i, j) + assert d_addr + (k * (Y * X) + j * X + i) * S == \ + get_field_addr(d.get_scalar_field(k), i, j) + + +@test_utils.test(arch=get_host_arch_list()) +def test_order_must_throw_scalar(): + with pytest.raises( + ti.TaichiCompilationError, + match='The dimensionality of shape and order must be the same'): + a = ti.field(dtype=ti.f32, shape=3, order='ij') + with pytest.raises(ti.TaichiCompilationError, + match='shape cannot be None when order is set'): + b = ti.field(dtype=ti.f32, shape=None, order='i') + with pytest.raises(ti.TaichiCompilationError, + match='The axes in order must be different'): + c = ti.field(dtype=ti.f32, shape=(3, 4, 3), order='iji') + with pytest.raises(ti.TaichiCompilationError, match='Invalid axis'): + d = ti.field(dtype=ti.f32, shape=(3, 4, 3), order='ijl') + + +@test_utils.test(arch=get_host_arch_list()) +def test_order_must_throw_vector(): + with pytest.raises( + ti.TaichiCompilationError, + match='The dimensionality of shape and order must be the same'): + a = ti.Vector.field(3, dtype=ti.f32, shape=3, order='ij') + with pytest.raises(ti.TaichiCompilationError, + match='shape cannot be None when order is set'): + b = ti.Vector.field(3, dtype=ti.f32, shape=None, order='i') + with pytest.raises(ti.TaichiCompilationError, + match='The axes in order must be different'): + c = ti.Vector.field(3, dtype=ti.f32, shape=(3, 4, 3), order='iii') + with pytest.raises(ti.TaichiCompilationError, match='Invalid axis'): + d = ti.Vector.field(3, dtype=ti.f32, shape=(3, 4, 3), order='ihj')