From faca090f3510e20b49533bb312b580f9778da695 Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Thu, 16 Nov 2023 12:58:56 +0000 Subject: [PATCH 01/29] feat[next] Enable embedded field view in ffront_tests --- pyproject.toml | 4 ++- src/gt4py/next/allocators.py | 2 +- src/gt4py/next/ffront/decorator.py | 34 ++++++++++++++----- tests/next_tests/exclusion_matrices.py | 7 ++++ .../ffront_tests/ffront_test_utils.py | 19 +++++------ .../ffront_tests/test_execution.py | 4 +++ 6 files changed, 48 insertions(+), 22 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index e2d2a7dfe9..9132a347a4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -344,7 +344,9 @@ markers = [ 'uses_strided_neighbor_offset: tests that require backend support for strided neighbor offset', 'uses_tuple_args: tests that require backend support for tuple arguments', 'uses_tuple_returns: tests that require backend support for tuple results', - 'uses_zero_dimensional_fields: tests that require backend support for zero-dimensional fields' + 'uses_zero_dimensional_fields: tests that require backend support for zero-dimensional fields', + 'uses_cartesian_shift: tests that use a Cartesian connectivity', + 'uses_unstructured_shift: tests that use a unstructured connectivity' ] norecursedirs = ['dist', 'build', 'cpp_backend_tests/build*', '_local/*', '.*'] testpaths = 'tests' diff --git a/src/gt4py/next/allocators.py b/src/gt4py/next/allocators.py index 58600d8cda..cbb5da19f0 100644 --- a/src/gt4py/next/allocators.py +++ b/src/gt4py/next/allocators.py @@ -331,7 +331,7 @@ def allocate( """ if device is None and allocator is None: - raise ValueError("No 'device' or 'allocator' specified") + allocator = StandardCPUFieldBufferAllocator() actual_allocator = get_allocator(allocator) if actual_allocator is None: assert device is not None # for mypy diff --git a/src/gt4py/next/ffront/decorator.py b/src/gt4py/next/ffront/decorator.py index 2d12331513..2afb4d9cb5 100644 --- a/src/gt4py/next/ffront/decorator.py +++ b/src/gt4py/next/ffront/decorator.py @@ -32,7 +32,7 @@ from gt4py._core import definitions as core_defs from gt4py.eve import utils as eve_utils from gt4py.eve.extended_typing import Any, Optional -from gt4py.next import allocators as next_allocators +from gt4py.next import allocators as next_allocators, common from gt4py.next.common import Dimension, DimensionKind, GridType from gt4py.next.ffront import ( dialect_ast_enums, @@ -171,7 +171,9 @@ class Program: past_node: past.Program closure_vars: dict[str, Any] definition: Optional[types.FunctionType] = None - backend: Optional[ppi.ProgramExecutor] = None + backend: None | eve_utils.NOTHING | ppi.ProgramExecutor = ( + eve_utils.NOTHING + ) # TODO(havogt): temporary change, remove once `None` is default backend grid_type: Optional[GridType] = None @classmethod @@ -282,21 +284,20 @@ def itir(self) -> itir.FencilDefinition: ) def __call__(self, *args, offset_provider: dict[str, Dimension], **kwargs) -> None: - if ( - self.backend is None and DEFAULT_BACKEND is None - ): # TODO(havogt): for now enable embedded execution by setting DEFAULT_BACKEND to None + if self.backend is None: self.definition(*args, **kwargs) return rewritten_args, size_args, kwargs = self._process_args(args, kwargs) - if not self.backend: + backend = self.backend + if self.backend is eve_utils.NOTHING: warnings.warn( UserWarning( f"Field View Program '{self.itir.id}': Using default ({DEFAULT_BACKEND}) backend." ) ) - backend = self.backend or DEFAULT_BACKEND + backend = DEFAULT_BACKEND ppi.ensure_processor_kind(backend, ppi.ProgramExecutor) if "debug" in kwargs: @@ -687,7 +688,7 @@ def __call__( # if we are reaching this from a program call. if "out" in kwargs: out = kwargs.pop("out") - if "offset_provider" in kwargs: + if self.backend is not None: # "out" and "offset_provider" -> field_operator as program offset_provider = kwargs.pop("offset_provider") args, kwargs = type_info.canonicalize_arguments(self.foast_node.type, args, kwargs) @@ -705,13 +706,28 @@ def __call__( ) else: # "out" -> field_operator called from program in embedded execution - out.ndarray[:] = self.definition(*args, **kwargs).ndarray[:] + + res = self.definition(*args) # TODO put offset_provider in contextvar + _tuple_assign_field(out, res) return else: # field_operator called from other field_operator in embedded execution return self.definition(*args, **kwargs) +def _tuple_assign_field( + tgt: tuple[common.Field | tuple, ...] | common.Field, + src: tuple[common.Field | tuple, ...] | common.Field, +): + if isinstance(tgt, tuple): + if not isinstance(src, tuple): + raise RuntimeError(f"Cannot assign {src} to {tgt}.") + for t, s in zip(tgt, src): + _tuple_assign_field(t, s) + else: + tgt.ndarray[...] = src.ndarray[...] # TODO + + @typing.overload def field_operator( definition: types.FunctionType, *, backend: Optional[ppi.ProgramExecutor] diff --git a/tests/next_tests/exclusion_matrices.py b/tests/next_tests/exclusion_matrices.py index ddea04649f..bbda9e6d26 100644 --- a/tests/next_tests/exclusion_matrices.py +++ b/tests/next_tests/exclusion_matrices.py @@ -98,6 +98,8 @@ class ProgramFormatterId(_PythonObjectIdMixin, str, enum.Enum): USES_TUPLE_ARGS = "uses_tuple_args" USES_TUPLE_RETURNS = "uses_tuple_returns" USES_ZERO_DIMENSIONAL_FIELDS = "uses_zero_dimensional_fields" +USES_CARTESIAN_SHIFT = "uses_cartesian_shift" +USES_UNSTRUCTURED_SHIFT = "uses_unstructured_shift" # Skip messages (available format keys: 'marker', 'backend') UNSUPPORTED_MESSAGE = "'{marker}' tests not supported by '{backend}' backend" @@ -114,10 +116,15 @@ class ProgramFormatterId(_PythonObjectIdMixin, str, enum.Enum): (USES_REDUCTION_WITH_ONLY_SPARSE_FIELDS, XFAIL, REDUCTION_WITH_ONLY_SPARSE_FIELDS_MESSAGE), (USES_SCAN_IN_FIELD_OPERATOR, XFAIL, UNSUPPORTED_MESSAGE), ] +EMBEDDED_SKIP_LIST = [ + (USES_CARTESIAN_SHIFT, XFAIL, UNSUPPORTED_MESSAGE), + (USES_UNSTRUCTURED_SHIFT, XFAIL, UNSUPPORTED_MESSAGE), +] #: Skip matrix, contains for each backend processor a list of tuples with following fields: #: (, ) BACKEND_SKIP_TEST_MATRIX = { + None: EMBEDDED_SKIP_LIST, OptionalProgramBackendId.DACE_CPU: GTFN_SKIP_TEST_LIST + [ (USES_CAN_DEREF, XFAIL, UNSUPPORTED_MESSAGE), diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/ffront_test_utils.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/ffront_test_utils.py index 386e64451d..fb753bf169 100644 --- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/ffront_test_utils.py +++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/ffront_test_utils.py @@ -53,6 +53,7 @@ def no_backend(program: itir.FencilDefinition, *args: Any, **kwargs: Any) -> Non definitions.ProgramBackendId.GTFN_CPU, definitions.ProgramBackendId.GTFN_CPU_IMPERATIVE, definitions.ProgramBackendId.GTFN_CPU_WITH_TEMPORARIES, + None, ] + OPTIONAL_PROCESSORS, ids=lambda p: p.short_id() if p is not None else "None", @@ -65,19 +66,15 @@ def fieldview_backend(request): Check ADR 15 for details on the test-exclusion matrices. """ backend_id = request.param - if backend_id is None: - backend = None - else: - backend = backend_id.load() - - for marker, skip_mark, msg in next_tests.exclusion_matrices.BACKEND_SKIP_TEST_MATRIX.get( - backend_id, [] - ): - if request.node.get_closest_marker(marker): - skip_mark(msg.format(marker=marker, backend=backend_id)) + backend = None if backend_id is None else backend_id.load() - backup_backend = decorator.DEFAULT_BACKEND + for marker, skip_mark, msg in next_tests.exclusion_matrices.BACKEND_SKIP_TEST_MATRIX.get( + backend_id, [] + ): + if request.node.get_closest_marker(marker): + skip_mark(msg.format(marker=marker, backend=backend_id)) + backup_backend = decorator.DEFAULT_BACKEND decorator.DEFAULT_BACKEND = no_backend yield backend decorator.DEFAULT_BACKEND = backup_backend diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_execution.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_execution.py index d381a2242a..8def5d5dfc 100644 --- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_execution.py +++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_execution.py @@ -76,6 +76,7 @@ def testee(a: cases.IJKField, b: cases.IJKField) -> tuple[cases.IJKField, cases. cases.verify_with_default_data(cartesian_case, testee, ref=lambda a, b: (a, b)) +@pytest.mark.uses_cartesian_shift def test_cartesian_shift(cartesian_case): # noqa: F811 # fixtures @gtx.field_operator def testee(a: cases.IJKField) -> cases.IJKField: @@ -87,6 +88,7 @@ def testee(a: cases.IJKField) -> cases.IJKField: cases.verify(cartesian_case, testee, a, out=out, ref=a[1:]) +@pytest.mark.uses_unstructured_shift def test_unstructured_shift(unstructured_case): # noqa: F811 # fixtures @gtx.field_operator def testee(a: cases.VField) -> cases.EField: @@ -99,6 +101,7 @@ def testee(a: cases.VField) -> cases.EField: ) +@pytest.mark.uses_unstructured_shift def test_composed_unstructured_shift(unstructured_case): @gtx.field_operator def composed_shift_unstructured_flat(inp: cases.VField) -> cases.CField: @@ -143,6 +146,7 @@ def composed_shift_unstructured(inp: cases.VField) -> cases.CField: ) +@pytest.mark.uses_cartesian_shift def test_fold_shifts(cartesian_case): # noqa: F811 # fixtures """Shifting the result of an addition should work.""" From 7931cfd543d0e0a46c698c75e314f16c92528500 Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Thu, 16 Nov 2023 13:32:28 +0000 Subject: [PATCH 02/29] broadcast for scalars --- pyproject.toml | 3 +- src/gt4py/next/ffront/fbuiltins.py | 52 ++++++++++--------- tests/next_tests/exclusion_matrices.py | 2 + .../ffront_tests/test_execution.py | 4 ++ .../embedded_tests/test_nd_array_field.py | 2 +- 5 files changed, 37 insertions(+), 26 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 9132a347a4..647fd6243e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -346,7 +346,8 @@ markers = [ 'uses_tuple_returns: tests that require backend support for tuple results', 'uses_zero_dimensional_fields: tests that require backend support for zero-dimensional fields', 'uses_cartesian_shift: tests that use a Cartesian connectivity', - 'uses_unstructured_shift: tests that use a unstructured connectivity' + 'uses_unstructured_shift: tests that use a unstructured connectivity', + 'uses_scan: tests that uses scan' ] norecursedirs = ['dist', 'build', 'cpp_backend_tests/build*', '_local/*', '.*'] testpaths = 'tests' diff --git a/src/gt4py/next/ffront/fbuiltins.py b/src/gt4py/next/ffront/fbuiltins.py index 13c21eb516..0a8cf504ba 100644 --- a/src/gt4py/next/ffront/fbuiltins.py +++ b/src/gt4py/next/ffront/fbuiltins.py @@ -28,10 +28,12 @@ cast, ) +import numpy as np from numpy import float32, float64, int32, int64 -from gt4py._core import definitions as gt4py_defs -from gt4py.next.common import Dimension, DimensionKind, Field +from gt4py._core import definitions as core_defs +from gt4py.next import common +from gt4py.next.common import Dimension, DimensionKind, Field # TODO fix import style from gt4py.next.ffront.experimental import as_offset # noqa F401 from gt4py.next.iterator import runtime from gt4py.next.type_system import type_specifications as ts @@ -58,7 +60,7 @@ def _type_conversion_helper(t: type) -> type[ts.TypeSpec] | tuple[type[ts.TypeSp return ts.FieldType elif t is Dimension: return ts.DimensionType - elif t is gt4py_defs.ScalarT: + elif t is core_defs.ScalarT: return ts.ScalarType elif t is type: return ( @@ -128,12 +130,8 @@ def __gt_type__(self) -> ts.FunctionType: ) -def builtin_function(fun: Callable[_P, _R]) -> BuiltInFunction[_R, _P]: - return BuiltInFunction(fun) - - MaskT = TypeVar("MaskT", bound=Field) -FieldT = TypeVar("FieldT", bound=Union[Field, gt4py_defs.Scalar, Tuple]) +FieldT = TypeVar("FieldT", bound=Union[Field, core_defs.Scalar, Tuple]) class WhereBuiltinFunction( @@ -153,7 +151,7 @@ def __call__(self, mask: MaskT, true_field: FieldT, false_field: FieldT) -> _R: return super().__call__(mask, true_field, false_field) -@builtin_function +@BuiltInFunction def neighbor_sum( field: Field, /, @@ -162,7 +160,7 @@ def neighbor_sum( raise NotImplementedError() -@builtin_function +@BuiltInFunction def max_over( field: Field, /, @@ -171,7 +169,7 @@ def max_over( raise NotImplementedError() -@builtin_function +@BuiltInFunction def min_over( field: Field, /, @@ -180,23 +178,29 @@ def min_over( raise NotImplementedError() -@builtin_function -def broadcast(field: Field | gt4py_defs.ScalarT, dims: Tuple[Dimension, ...], /) -> Field: - raise NotImplementedError() +@BuiltInFunction +def broadcast(field: Field | core_defs.ScalarT, dims: tuple[Dimension, ...], /) -> Field: + assert core_defs.is_scalar_type(field) + return common.field( + np.asarray(field)[ + tuple([np.newaxis] * len(dims)) + ], # TODO(havogt) use FunctionField once available + domain=common.Domain(dims=dims, ranges=tuple([common.UnitRange.infinity()] * len(dims))), + ) @WhereBuiltinFunction def where( mask: Field, - true_field: Field | gt4py_defs.ScalarT | Tuple, - false_field: Field | gt4py_defs.ScalarT | Tuple, + true_field: Field | core_defs.ScalarT | Tuple, + false_field: Field | core_defs.ScalarT | Tuple, /, ) -> Field | Tuple: raise NotImplementedError() -@builtin_function -def astype(field: Field | gt4py_defs.ScalarT, type_: type, /) -> Field: +@BuiltInFunction +def astype(field: Field | core_defs.ScalarT, type_: type, /) -> Field: raise NotImplementedError() @@ -229,11 +233,11 @@ def astype(field: Field | gt4py_defs.ScalarT, type_: type, /) -> Field: def _make_unary_math_builtin(name): - def impl(value: Field | gt4py_defs.ScalarT, /) -> Field | gt4py_defs.ScalarT: + def impl(value: Field | core_defs.ScalarT, /) -> Field | core_defs.ScalarT: raise NotImplementedError() impl.__name__ = name - globals()[name] = builtin_function(impl) + globals()[name] = BuiltInFunction(impl) for f in ( @@ -248,14 +252,14 @@ def impl(value: Field | gt4py_defs.ScalarT, /) -> Field | gt4py_defs.ScalarT: def _make_binary_math_builtin(name): def impl( - lhs: Field | gt4py_defs.ScalarT, - rhs: Field | gt4py_defs.ScalarT, + lhs: Field | core_defs.ScalarT, + rhs: Field | core_defs.ScalarT, /, - ) -> Field | gt4py_defs.ScalarT: + ) -> Field | core_defs.ScalarT: raise NotImplementedError() impl.__name__ = name - globals()[name] = builtin_function(impl) + globals()[name] = BuiltInFunction(impl) for f in BINARY_MATH_NUMBER_BUILTIN_NAMES: diff --git a/tests/next_tests/exclusion_matrices.py b/tests/next_tests/exclusion_matrices.py index bbda9e6d26..b87cae2153 100644 --- a/tests/next_tests/exclusion_matrices.py +++ b/tests/next_tests/exclusion_matrices.py @@ -100,6 +100,7 @@ class ProgramFormatterId(_PythonObjectIdMixin, str, enum.Enum): USES_ZERO_DIMENSIONAL_FIELDS = "uses_zero_dimensional_fields" USES_CARTESIAN_SHIFT = "uses_cartesian_shift" USES_UNSTRUCTURED_SHIFT = "uses_unstructured_shift" +USES_SCAN = "uses_scan" # Skip messages (available format keys: 'marker', 'backend') UNSUPPORTED_MESSAGE = "'{marker}' tests not supported by '{backend}' backend" @@ -119,6 +120,7 @@ class ProgramFormatterId(_PythonObjectIdMixin, str, enum.Enum): EMBEDDED_SKIP_LIST = [ (USES_CARTESIAN_SHIFT, XFAIL, UNSUPPORTED_MESSAGE), (USES_UNSTRUCTURED_SHIFT, XFAIL, UNSUPPORTED_MESSAGE), + (USES_SCAN, XFAIL, UNSUPPORTED_MESSAGE), ] #: Skip matrix, contains for each backend processor a list of tuples with following fields: diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_execution.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_execution.py index 8def5d5dfc..aef43df47d 100644 --- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_execution.py +++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_execution.py @@ -210,6 +210,7 @@ def testee(a: int32) -> cases.VField: @pytest.mark.uses_index_fields +@pytest.mark.uses_cartesian_shift def test_scalar_arg_with_field(cartesian_case): # noqa: F811 # fixtures @gtx.field_operator def testee(a: cases.IJKField, b: int32) -> cases.IJKField: @@ -250,6 +251,7 @@ def testee(size: gtx.IndexType, out: gtx.Field[[IDim], gtx.IndexType]): ) +@pytest.mark.uses_scan def test_scalar_scan(cartesian_case): # noqa: F811 # fixtures @gtx.scan_operator(axis=KDim, forward=True, init=(0.0)) def testee_scan(state: float, qc_in: float, scalar: float) -> float: @@ -268,6 +270,7 @@ def testee(qc: cases.IKFloatField, scalar: float): cases.verify(cartesian_case, testee, qc, scalar, inout=qc, ref=expected) +@pytest.mark.uses_scan @pytest.mark.uses_scan_in_field_operator def test_tuple_scalar_scan(cartesian_case): # noqa: F811 # fixtures @gtx.scan_operator(axis=KDim, forward=True, init=0.0) @@ -289,6 +292,7 @@ def testee_op( cases.verify(cartesian_case, testee_op, qc, tuple_scalar, out=qc, ref=expected) +@pytest.mark.uses_scan @pytest.mark.uses_index_fields def test_scalar_scan_vertical_offset(cartesian_case): # noqa: F811 # fixtures @gtx.scan_operator(axis=KDim, forward=True, init=(0.0)) diff --git a/tests/next_tests/unit_tests/embedded_tests/test_nd_array_field.py b/tests/next_tests/unit_tests/embedded_tests/test_nd_array_field.py index 49aeece87e..00dbf68274 100644 --- a/tests/next_tests/unit_tests/embedded_tests/test_nd_array_field.py +++ b/tests/next_tests/unit_tests/embedded_tests/test_nd_array_field.py @@ -259,7 +259,7 @@ def test_mixed_fields(product_nd_array_implementation): def test_non_dispatched_function(): - @fbuiltins.builtin_function + @fbuiltins.BuiltInFunction def fma(a: common.Field, b: common.Field, c: common.Field, /) -> common.Field: return a * b + c From 4734c84b6ee4c2df07165c63968a811c2f28e60d Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Thu, 16 Nov 2023 13:47:39 +0000 Subject: [PATCH 03/29] implement astype --- src/gt4py/next/embedded/nd_array_field.py | 15 +++++++++++---- src/gt4py/next/ffront/fbuiltins.py | 3 ++- tests/next_tests/exclusion_matrices.py | 1 + .../feature_tests/ffront_tests/test_execution.py | 11 +++++++++++ 4 files changed, 25 insertions(+), 5 deletions(-) diff --git a/src/gt4py/next/embedded/nd_array_field.py b/src/gt4py/next/embedded/nd_array_field.py index ea88948841..1cd08ab40c 100644 --- a/src/gt4py/next/embedded/nd_array_field.py +++ b/src/gt4py/next/embedded/nd_array_field.py @@ -135,16 +135,16 @@ def from_array( /, *, domain: common.DomainLike, - dtype_like: Optional[core_defs.DType] = None, # TODO define DTypeLike + dtype: Optional[core_defs.DTypeLike] = None, ) -> NdArrayField: domain = common.domain(domain) xp = cls.array_ns - xp_dtype = None if dtype_like is None else xp.dtype(core_defs.dtype(dtype_like).scalar_type) + xp_dtype = None if dtype is None else xp.dtype(core_defs.dtype(dtype).scalar_type) array = xp.asarray(data, dtype=xp_dtype) - if dtype_like is not None: - assert array.dtype.type == core_defs.dtype(dtype_like).scalar_type + if dtype is not None: + assert array.dtype.type == core_defs.dtype(dtype).scalar_type assert issubclass(array.dtype.type, core_defs.SCALAR_TYPES) @@ -351,6 +351,13 @@ def _builtins_broadcast( NdArrayField.register_builtin_func(fbuiltins.broadcast, _builtins_broadcast) +def _astype(field: NdArrayField, type_: type) -> NdArrayField: + return field.__class__.from_array(field.ndarray, domain=field.domain, dtype=type_) + + +NdArrayField.register_builtin_func(fbuiltins.astype, _astype) # type: ignore[arg-type] # TODO(havogt) the registry should not be for any Field + + def _get_slices_from_domain_slice( domain: common.Domain, domain_slice: common.Domain | Sequence[common.NamedRange | common.NamedIndex | Any], diff --git a/src/gt4py/next/ffront/fbuiltins.py b/src/gt4py/next/ffront/fbuiltins.py index 0a8cf504ba..65c6a834e8 100644 --- a/src/gt4py/next/ffront/fbuiltins.py +++ b/src/gt4py/next/ffront/fbuiltins.py @@ -201,7 +201,8 @@ def where( @BuiltInFunction def astype(field: Field | core_defs.ScalarT, type_: type, /) -> Field: - raise NotImplementedError() + assert core_defs.is_scalar_type(field) + return type_(field) UNARY_MATH_NUMBER_BUILTIN_NAMES = ["abs"] diff --git a/tests/next_tests/exclusion_matrices.py b/tests/next_tests/exclusion_matrices.py index b87cae2153..d74fb68677 100644 --- a/tests/next_tests/exclusion_matrices.py +++ b/tests/next_tests/exclusion_matrices.py @@ -121,6 +121,7 @@ class ProgramFormatterId(_PythonObjectIdMixin, str, enum.Enum): (USES_CARTESIAN_SHIFT, XFAIL, UNSUPPORTED_MESSAGE), (USES_UNSTRUCTURED_SHIFT, XFAIL, UNSUPPORTED_MESSAGE), (USES_SCAN, XFAIL, UNSUPPORTED_MESSAGE), + (USES_DYNAMIC_OFFSETS, XFAIL, UNSUPPORTED_MESSAGE), ] #: Skip matrix, contains for each backend processor a list of tuples with following fields: diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_execution.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_execution.py index aef43df47d..e6cd683dc3 100644 --- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_execution.py +++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_execution.py @@ -421,6 +421,7 @@ def combine(a: cases.IField, b: cases.IField) -> cases.IField: cases.verify_with_default_data(cartesian_case, combine, ref=lambda a, b: a + a + b) +@pytest.mark.uses_unstructured_shift @pytest.mark.uses_reduction_over_lift_expressions def test_nested_reduction(unstructured_case): @gtx.field_operator @@ -442,6 +443,7 @@ def testee(a: cases.EField) -> cases.EField: ) +@pytest.mark.uses_unstructured_shift @pytest.mark.xfail(reason="Not yet supported in lowering, requires `map_`ing of inner reduce op.") def test_nested_reduction_shift_first(unstructured_case): @gtx.field_operator @@ -462,6 +464,7 @@ def testee(inp: cases.EField) -> cases.EField: ) +@pytest.mark.uses_unstructured_shift @pytest.mark.uses_tuple_returns def test_tuple_return_2(unstructured_case): @gtx.field_operator @@ -481,6 +484,7 @@ def testee(a: cases.EField, b: cases.EField) -> tuple[cases.VField, cases.VField ) +@pytest.mark.uses_unstructured_shift @pytest.mark.uses_constant_fields def test_tuple_with_local_field_in_reduction_shifted(unstructured_case): @gtx.field_operator @@ -510,6 +514,7 @@ def testee(a: tuple[tuple[cases.IField, cases.IField], cases.IField]) -> cases.I ) +@pytest.mark.uses_scan @pytest.mark.parametrize("forward", [True, False]) def test_fieldop_from_scan(cartesian_case, forward): init = 1.0 @@ -530,6 +535,7 @@ def simple_scan_operator(carry: float) -> float: cases.verify(cartesian_case, simple_scan_operator, out=out, ref=expected) +@pytest.mark.uses_scan @pytest.mark.uses_lift_expressions def test_solve_triag(cartesian_case): if cartesian_case.backend in [ @@ -618,6 +624,7 @@ def testee( ) +@pytest.mark.uses_unstructured_shift @pytest.mark.uses_reduction_over_lift_expressions def test_ternary_builtin_neighbor_sum(unstructured_case): @gtx.field_operator @@ -636,6 +643,7 @@ def testee(a: cases.EField, b: cases.EField) -> cases.VField: ) +@pytest.mark.uses_scan def test_ternary_scan(cartesian_case): if cartesian_case.backend in [gtfn.run_gtfn_with_temporaries]: pytest.xfail("Temporary extraction does not work correctly in combination with scans.") @@ -658,6 +666,7 @@ def simple_scan_operator(carry: float, a: float) -> float: @pytest.mark.parametrize("forward", [True, False]) +@pytest.mark.uses_scan @pytest.mark.uses_tuple_returns def test_scan_nested_tuple_output(forward, cartesian_case): if cartesian_case.backend in [gtfn.run_gtfn_with_temporaries]: @@ -690,6 +699,7 @@ def testee(out: tuple[cases.KField, tuple[cases.KField, cases.KField]]): @pytest.mark.uses_tuple_args +@pytest.mark.uses_scan def test_scan_nested_tuple_input(cartesian_case): init = 1.0 k_size = cartesian_case.default_sizes[KDim] @@ -880,6 +890,7 @@ def program_domain_tuple( ) +@pytest.mark.uses_cartesian_shift def test_where_k_offset(cartesian_case): @gtx.field_operator def fieldop_where_k_offset( From e1463d00dea8db6d091268a2a99a55f0b3c61279 Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Thu, 16 Nov 2023 13:59:23 +0000 Subject: [PATCH 04/29] support binary builtins for scalars --- src/gt4py/next/ffront/decorator.py | 6 +++--- src/gt4py/next/ffront/fbuiltins.py | 13 +++++++++++-- 2 files changed, 14 insertions(+), 5 deletions(-) diff --git a/src/gt4py/next/ffront/decorator.py b/src/gt4py/next/ffront/decorator.py index 2afb4d9cb5..8ae46a45d2 100644 --- a/src/gt4py/next/ffront/decorator.py +++ b/src/gt4py/next/ffront/decorator.py @@ -284,12 +284,12 @@ def itir(self) -> itir.FencilDefinition: ) def __call__(self, *args, offset_provider: dict[str, Dimension], **kwargs) -> None: + rewritten_args, size_args, kwargs = self._process_args(args, kwargs) + if self.backend is None: - self.definition(*args, **kwargs) + self.definition(*rewritten_args, **kwargs) return - rewritten_args, size_args, kwargs = self._process_args(args, kwargs) - backend = self.backend if self.backend is eve_utils.NOTHING: warnings.warn( diff --git a/src/gt4py/next/ffront/fbuiltins.py b/src/gt4py/next/ffront/fbuiltins.py index 65c6a834e8..6611d2a228 100644 --- a/src/gt4py/next/ffront/fbuiltins.py +++ b/src/gt4py/next/ffront/fbuiltins.py @@ -14,6 +14,7 @@ import dataclasses import inspect +import math from builtins import bool, float, int, tuple from typing import ( Any, @@ -248,7 +249,13 @@ def impl(value: Field | core_defs.ScalarT, /) -> Field | core_defs.ScalarT: ): _make_unary_math_builtin(f) -BINARY_MATH_NUMBER_BUILTIN_NAMES = ["minimum", "maximum", "fmod", "power"] +BINARY_MATH_NUMBER_BUILTIN_TO_PYTHON_SCALAR_FUNCTION = { + "minimum": min, + "maximum": max, + "fmod": math.fmod, + "power": pow, +} +BINARY_MATH_NUMBER_BUILTIN_NAMES = list(BINARY_MATH_NUMBER_BUILTIN_TO_PYTHON_SCALAR_FUNCTION.keys()) def _make_binary_math_builtin(name): @@ -257,7 +264,9 @@ def impl( rhs: Field | core_defs.ScalarT, /, ) -> Field | core_defs.ScalarT: - raise NotImplementedError() + assert core_defs.is_scalar_type(lhs) + assert core_defs.is_scalar_type(rhs) + return BINARY_MATH_NUMBER_BUILTIN_TO_PYTHON_SCALAR_FUNCTION[name](lhs, rhs) # type: ignore[operator] # Cannot call function of unknown type impl.__name__ = name globals()[name] = BuiltInFunction(impl) From f1047dc9970d4d0f5bf4392eeb2c21abbb1dd206 Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Thu, 16 Nov 2023 14:38:35 +0000 Subject: [PATCH 05/29] support domain --- src/gt4py/next/common.py | 7 +++++ src/gt4py/next/embedded/nd_array_field.py | 5 +--- src/gt4py/next/ffront/decorator.py | 17 +++++++---- .../ffront_tests/test_arg_call_interface.py | 4 +++ .../ffront_tests/test_execution.py | 30 ++++++++++++++----- 5 files changed, 46 insertions(+), 17 deletions(-) diff --git a/src/gt4py/next/common.py b/src/gt4py/next/common.py index ffaa410563..e5d5d7d463 100644 --- a/src/gt4py/next/common.py +++ b/src/gt4py/next/common.py @@ -141,6 +141,13 @@ def __and__(self, other: Set[Any]) -> UnitRange: else: raise NotImplementedError("Can only find the intersection between UnitRange instances.") + def __le__(self, other: Set[Any]): + if isinstance(other, UnitRange): + # required for infinity comparison + return self.start >= other.start and self.stop <= other.stop + else: + return Set.__le__(self, other) + def __str__(self) -> str: return f"({self.start}:{self.stop})" diff --git a/src/gt4py/next/embedded/nd_array_field.py b/src/gt4py/next/embedded/nd_array_field.py index 1cd08ab40c..dd2e0d2c34 100644 --- a/src/gt4py/next/embedded/nd_array_field.py +++ b/src/gt4py/next/embedded/nd_array_field.py @@ -150,10 +150,7 @@ def from_array( assert all(isinstance(d, common.Dimension) for d in domain.dims), domain assert len(domain) == array.ndim - assert all( - len(r) == s or (s == 1 and r == common.UnitRange.infinity()) - for r, s in zip(domain.ranges, array.shape) - ) + assert all(len(r) == s or s == 1 for r, s in zip(domain.ranges, array.shape)) return cls(domain, array) diff --git a/src/gt4py/next/ffront/decorator.py b/src/gt4py/next/ffront/decorator.py index 8ae46a45d2..95b87e0efc 100644 --- a/src/gt4py/next/ffront/decorator.py +++ b/src/gt4py/next/ffront/decorator.py @@ -688,9 +688,9 @@ def __call__( # if we are reaching this from a program call. if "out" in kwargs: out = kwargs.pop("out") + offset_provider = kwargs.pop("offset_provider", None) if self.backend is not None: # "out" and "offset_provider" -> field_operator as program - offset_provider = kwargs.pop("offset_provider") args, kwargs = type_info.canonicalize_arguments(self.foast_node.type, args, kwargs) # TODO(tehrengruber): check all offset providers are given # deduce argument types @@ -706,9 +706,12 @@ def __call__( ) else: # "out" -> field_operator called from program in embedded execution - - res = self.definition(*args) # TODO put offset_provider in contextvar - _tuple_assign_field(out, res) + # TODO put offset_provider in ctxt var + domain = kwargs.pop("domain", None) + res = self.definition(*args, **kwargs) + _tuple_assign_field( + out, res, domain=None if domain is None else common.domain(domain) + ) return else: # field_operator called from other field_operator in embedded execution @@ -718,14 +721,16 @@ def __call__( def _tuple_assign_field( tgt: tuple[common.Field | tuple, ...] | common.Field, src: tuple[common.Field | tuple, ...] | common.Field, + domain: common.Domain, ): if isinstance(tgt, tuple): if not isinstance(src, tuple): raise RuntimeError(f"Cannot assign {src} to {tgt}.") for t, s in zip(tgt, src): - _tuple_assign_field(t, s) + _tuple_assign_field(t, s, domain) else: - tgt.ndarray[...] = src.ndarray[...] # TODO + domain = domain or tgt.domain + tgt[domain] = src[domain] @typing.overload diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_arg_call_interface.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_arg_call_interface.py index deb1382dfb..6957e628bb 100644 --- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_arg_call_interface.py +++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_arg_call_interface.py @@ -158,6 +158,7 @@ def testee( ) +@pytest.mark.uses_scan @pytest.mark.uses_scan_in_field_operator def test_call_scan_operator_from_field_operator(cartesian_case): @scan_operator(axis=KDim, forward=True, init=0.0) @@ -183,6 +184,7 @@ def testee(a: IJKFloatField, b: IJKFloatField) -> IJKFloatField: cases.verify(cartesian_case, testee, a, b, out=out, ref=expected) +@pytest.mark.uses_scan def test_call_scan_operator_from_program(cartesian_case): @scan_operator(axis=KDim, forward=True, init=0.0) def testee_scan(state: float, x: float, y: float) -> float: @@ -222,6 +224,7 @@ def testee( ) +@pytest.mark.uses_scan def test_scan_wrong_return_type(cartesian_case): with pytest.raises( errors.DSLError, @@ -239,6 +242,7 @@ def testee(qc: cases.IKFloatField, param_1: int32, param_2: float, scalar: float testee_scan(qc, param_1, param_2, scalar, out=(qc, param_1, param_2)) +@pytest.mark.uses_scan def test_scan_wrong_state_type(cartesian_case): with pytest.raises( errors.DSLError, diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_execution.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_execution.py index e6cd683dc3..b47c5c4256 100644 --- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_execution.py +++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_execution.py @@ -772,7 +772,10 @@ def program_domain(a: cases.IField, out: cases.IField): a = cases.allocate(cartesian_case, program_domain, "a")() out = cases.allocate(cartesian_case, program_domain, "out")() - cases.verify(cartesian_case, program_domain, a, out, inout=out[1:9], ref=a[1:9] * 2) + ref = out.ndarray.copy() # ensure we are not overwriting out outside of the domain + ref[1:9] = a[1:9] * 2 + + cases.verify(cartesian_case, program_domain, a, out, inout=out, ref=ref) def test_domain_input_bounds(cartesian_case): @@ -803,6 +806,9 @@ def program_domain( inp = cases.allocate(cartesian_case, program_domain, "inp")() out = cases.allocate(cartesian_case, fieldop_domain, cases.RETURN)() + ref = out.ndarray.copy() + ref[lower_i : int(upper_i / 2)] = inp[lower_i : int(upper_i / 2)] * 2 + cases.verify( cartesian_case, program_domain, @@ -810,8 +816,8 @@ def program_domain( out, lower_i, upper_i, - inout=out[lower_i : int(upper_i / 2)], - ref=inp[lower_i : int(upper_i / 2)] * 2, + inout=out, + ref=ref, ) @@ -843,6 +849,11 @@ def program_domain( a = cases.allocate(cartesian_case, program_domain, "a")() out = cases.allocate(cartesian_case, program_domain, "out")() + ref = out.ndarray.copy() + ref[1 * lower_i : upper_i + 0, lower_j - 0 : upper_j] = ( + a[1 * lower_i : upper_i + 0, lower_j - 0 : upper_j] * 2 + ) + cases.verify( cartesian_case, program_domain, @@ -852,8 +863,8 @@ def program_domain( upper_i, lower_j, upper_j, - inout=out[1 * lower_i : upper_i + 0, lower_j - 0 : upper_j], - ref=a[1 * lower_i : upper_i + 0, lower_j - 0 : upper_j] * 2, + inout=out, + ref=ref, ) @@ -878,6 +889,11 @@ def program_domain_tuple( out0 = cases.allocate(cartesian_case, program_domain_tuple, "out0")() out1 = cases.allocate(cartesian_case, program_domain_tuple, "out1")() + ref0 = out0.ndarray.copy() + ref0[1:9, 4:6] = inp0[1:9, 4:6] + inp1[1:9, 4:6] + ref1 = out1.ndarray.copy() + ref1[1:9, 4:6] = inp1[1:9, 4:6] + cases.verify( cartesian_case, program_domain_tuple, @@ -885,8 +901,8 @@ def program_domain_tuple( inp1, out0, out1, - inout=(out0[1:9, 4:6], out1[1:9, 4:6]), - ref=(inp0[1:9, 4:6] + inp1[1:9, 4:6], inp1[1:9, 4:6]), + inout=(out0, out1), + ref=(ref0, ref1), ) From 9ac0ddd13d424e8e2a1cbd54d282a26eca384ffd Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Thu, 16 Nov 2023 16:13:20 +0000 Subject: [PATCH 06/29] add __ne__, __eq__ --- pyproject.toml | 3 ++- src/gt4py/next/common.py | 8 ++++++++ src/gt4py/next/embedded/nd_array_field.py | 4 ++++ src/gt4py/next/ffront/fbuiltins.py | 3 +++ src/gt4py/next/iterator/embedded.py | 12 ++++++++++++ tests/next_tests/exclusion_matrices.py | 2 ++ .../ffront_tests/test_execution.py | 7 +++++++ .../ffront_tests/test_external_local_field.py | 3 +++ .../ffront_tests/test_gt4py_builtins.py | 7 +++++++ .../ffront_tests/test_math_builtin_execution.py | 3 +++ .../ffront_tests/test_math_unary_builtins.py | 17 +++++++---------- .../feature_tests/ffront_tests/test_program.py | 2 ++ .../ffront_tests/test_icon_like_scan.py | 3 +++ .../ffront_tests/test_laplacian.py | 4 ++++ 14 files changed, 67 insertions(+), 11 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 647fd6243e..a200451b2d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -347,7 +347,8 @@ markers = [ 'uses_zero_dimensional_fields: tests that require backend support for zero-dimensional fields', 'uses_cartesian_shift: tests that use a Cartesian connectivity', 'uses_unstructured_shift: tests that use a unstructured connectivity', - 'uses_scan: tests that uses scan' + 'uses_scan: tests that uses scan', + 'checks_specific_error: tests that rely on the backend to produce a specific error message' ] norecursedirs = ['dist', 'build', 'cpp_backend_tests/build*', '_local/*', '.*'] testpaths = 'tests' diff --git a/src/gt4py/next/common.py b/src/gt4py/next/common.py index e5d5d7d463..da6c60be3f 100644 --- a/src/gt4py/next/common.py +++ b/src/gt4py/next/common.py @@ -493,6 +493,14 @@ def __neg__(self) -> Field: def __invert__(self) -> Field: """Only defined for `Field` of value type `bool`.""" + @abc.abstractmethod + def __eq__(self, other: Any) -> Field: + ... + + @abc.abstractmethod + def __neq__(self, other: Any) -> Field: + ... + @abc.abstractmethod def __add__(self, other: Field | core_defs.ScalarT) -> Field: ... diff --git a/src/gt4py/next/embedded/nd_array_field.py b/src/gt4py/next/embedded/nd_array_field.py index dd2e0d2c34..a26b8612c1 100644 --- a/src/gt4py/next/embedded/nd_array_field.py +++ b/src/gt4py/next/embedded/nd_array_field.py @@ -191,6 +191,10 @@ def restrict(self, index: common.AnyIndexSpec) -> common.Field | core_defs.Scala __mod__ = __rmod__ = _make_builtin("mod", "mod") + __ne__ = _make_builtin("not_equal", "not_equal") + + __eq__ = _make_builtin("equal", "equal") + def __and__(self, other: common.Field | core_defs.ScalarT) -> NdArrayField: if self.dtype == core_defs.BoolDType(): return _make_builtin("logical_and", "logical_and")(self, other) diff --git a/src/gt4py/next/ffront/fbuiltins.py b/src/gt4py/next/ffront/fbuiltins.py index 6611d2a228..4a7db4d0dc 100644 --- a/src/gt4py/next/ffront/fbuiltins.py +++ b/src/gt4py/next/ffront/fbuiltins.py @@ -236,6 +236,9 @@ def astype(field: Field | core_defs.ScalarT, type_: type, /) -> Field: def _make_unary_math_builtin(name): def impl(value: Field | core_defs.ScalarT, /) -> Field | core_defs.ScalarT: + # TODO(havogt): enable once we have a failing test (see `test_math_builtin_execution.py`) + # assert core_defs.is_scalar_type(value) # noqa: E800 # commented code + # return getattr(math, name)(value)# noqa: E800 # commented code raise NotImplementedError() impl.__name__ = name diff --git a/src/gt4py/next/iterator/embedded.py b/src/gt4py/next/iterator/embedded.py index 674f99f61c..d7abbbdffc 100644 --- a/src/gt4py/next/iterator/embedded.py +++ b/src/gt4py/next/iterator/embedded.py @@ -1093,6 +1093,12 @@ def __neg__(self) -> common.Field: def __invert__(self) -> common.Field: raise NotImplementedError() + def __eq__(self, other: Any) -> common.Field: + raise NotImplementedError() + + def __neq__(self, other: Any) -> common.Field: + raise NotImplementedError() + def __add__(self, other: common.Field | core_defs.ScalarT) -> common.Field: raise NotImplementedError() @@ -1194,6 +1200,12 @@ def __neg__(self) -> common.Field: def __invert__(self) -> common.Field: raise NotImplementedError() + def __eq__(self, other: Any) -> common.Field: + raise NotImplementedError() + + def __neq__(self, other: Any) -> common.Field: + raise NotImplementedError() + def __add__(self, other: common.Field | core_defs.ScalarT) -> common.Field: raise NotImplementedError() diff --git a/tests/next_tests/exclusion_matrices.py b/tests/next_tests/exclusion_matrices.py index d74fb68677..249e17d358 100644 --- a/tests/next_tests/exclusion_matrices.py +++ b/tests/next_tests/exclusion_matrices.py @@ -101,6 +101,7 @@ class ProgramFormatterId(_PythonObjectIdMixin, str, enum.Enum): USES_CARTESIAN_SHIFT = "uses_cartesian_shift" USES_UNSTRUCTURED_SHIFT = "uses_unstructured_shift" USES_SCAN = "uses_scan" +CHECKS_SPECIFIC_ERROR = "checks_specific_error" # Skip messages (available format keys: 'marker', 'backend') UNSUPPORTED_MESSAGE = "'{marker}' tests not supported by '{backend}' backend" @@ -122,6 +123,7 @@ class ProgramFormatterId(_PythonObjectIdMixin, str, enum.Enum): (USES_UNSTRUCTURED_SHIFT, XFAIL, UNSUPPORTED_MESSAGE), (USES_SCAN, XFAIL, UNSUPPORTED_MESSAGE), (USES_DYNAMIC_OFFSETS, XFAIL, UNSUPPORTED_MESSAGE), + (CHECKS_SPECIFIC_ERROR, XFAIL, UNSUPPORTED_MESSAGE), ] #: Skip matrix, contains for each backend processor a list of tuples with following fields: diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_execution.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_execution.py index b47c5c4256..22154da9a7 100644 --- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_execution.py +++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_execution.py @@ -1044,6 +1044,13 @@ def _invalid_unpack() -> tuple[int32, float64, int32]: def test_constant_closure_vars(cartesian_case): + if cartesian_case.backend is None: + # >>> field = gtx.zeros(domain) + # >>> np.int32(1)*field # steals the buffer from the field + # array([0.]) + + # TODO(havogt): remove `__array__`` from `NdArrayField` + pytest.xfail("Bug: Binary operation between np datatype and Field returns ndarray.") from gt4py.eve.utils import FrozenNamespace constants = FrozenNamespace( diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_external_local_field.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_external_local_field.py index 04b27c6c17..5135b3d47a 100644 --- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_external_local_field.py +++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_external_local_field.py @@ -26,6 +26,9 @@ ) +pytestmark = pytest.mark.uses_unstructured_shift + + def test_external_local_field(unstructured_case): @gtx.field_operator def testee( diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_gt4py_builtins.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_gt4py_builtins.py index 8213f54a45..1eba95e880 100644 --- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_gt4py_builtins.py +++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_gt4py_builtins.py @@ -39,6 +39,7 @@ ) +@pytest.mark.uses_unstructured_shift @pytest.mark.parametrize( "strategy", [cases.UniqueInitializer(1), cases.UniqueInitializer(-100)], @@ -65,6 +66,7 @@ def testee(edge_f: cases.EField) -> cases.VField: cases.verify(unstructured_case, testee, inp, ref=ref, out=out) +@pytest.mark.uses_unstructured_shift def test_minover_execution(unstructured_case): @gtx.field_operator def minover(edge_f: cases.EField) -> cases.VField: @@ -77,6 +79,7 @@ def minover(edge_f: cases.EField) -> cases.VField: ) +@pytest.mark.uses_unstructured_shift def test_reduction_execution(unstructured_case): @gtx.field_operator def reduction(edge_f: cases.EField) -> cases.VField: @@ -93,6 +96,7 @@ def fencil(edge_f: cases.EField, out: cases.VField): ) +@pytest.mark.uses_unstructured_shift @pytest.mark.uses_constant_fields def test_reduction_expression_in_call(unstructured_case): @gtx.field_operator @@ -113,6 +117,7 @@ def fencil(edge_f: cases.EField, out: cases.VField): ) +@pytest.mark.uses_unstructured_shift def test_reduction_with_common_expression(unstructured_case): @gtx.field_operator def testee(flux: cases.EField) -> cases.VField: @@ -191,6 +196,7 @@ def broadcast_two_fields(inp1: cases.IField, inp2: gtx.Field[[JDim], int32]) -> ) +@pytest.mark.uses_cartesian_shift def test_broadcast_shifted(cartesian_case): @gtx.field_operator def simple_broadcast(inp: cases.IField) -> cases.IJField: @@ -249,6 +255,7 @@ def conditional_promotion(a: cases.IFloatField) -> cases.IFloatField: ) +@pytest.mark.uses_cartesian_shift def test_conditional_shifted(cartesian_case): @gtx.field_operator def conditional_shifted( diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_math_builtin_execution.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_math_builtin_execution.py index a5d2b92719..a1839b8e17 100644 --- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_math_builtin_execution.py +++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_math_builtin_execution.py @@ -116,6 +116,9 @@ def make_builtin_field_operator(builtin_name: str): @pytest.mark.parametrize("builtin_name, inputs", math_builtin_test_data()) def test_math_function_builtins_execution(cartesian_case, builtin_name: str, inputs): + if cartesian_case.backend is None: + # TODO(havogt) find a way that works for embedded + pytest.xfail("Test does not have a field view program.") if builtin_name == "gamma": # numpy has no gamma function ref_impl: Callable = np.vectorize(math.gamma) diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_math_unary_builtins.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_math_unary_builtins.py index 5a277f9440..59e11a7de8 100644 --- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_math_unary_builtins.py +++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_math_unary_builtins.py @@ -173,17 +173,14 @@ def tilde_fieldop(inp1: cases.IBoolField) -> cases.IBoolField: def test_unary_not(cartesian_case): - @gtx.field_operator - def not_fieldop(inp1: cases.IBoolField) -> cases.IBoolField: - return not inp1 + pytest.xfail( + "We accidentally supported `not` on fields. This is wrong, we should raise an error." + ) + with pytest.raises: # TODO `not` on a field should be illegal - size = cartesian_case.default_sizes[IDim] - bool_field = np.random.choice(a=[False, True], size=(size)) - inp1 = cases.allocate(cartesian_case, not_fieldop, "inp1").strategy( - cases.ConstInitializer(bool_field) - )() - out = cases.allocate(cartesian_case, not_fieldop, cases.RETURN)() - cases.verify(cartesian_case, not_fieldop, inp1, out=out, ref=~inp1) + @gtx.field_operator + def not_fieldop(inp1: cases.IBoolField) -> cases.IBoolField: + return not inp1 # Trig builtins diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_program.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_program.py index 7a1c827a0d..8f7d3faf81 100644 --- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_program.py +++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_program.py @@ -51,6 +51,7 @@ def test_identity_fo_execution(cartesian_case, identity_def): ) +@pytest.mark.uses_cartesian_shift def test_shift_by_one_execution(cartesian_case): @gtx.field_operator def shift_by_one(in_field: cases.IFloatField) -> cases.IFloatField: @@ -230,6 +231,7 @@ def test_wrong_argument_type(cartesian_case, copy_program_def): assert re.search(msg, exc_info.value.__cause__.args[0]) is not None +@pytest.mark.checks_specific_error def test_dimensions_domain(cartesian_case): @gtx.field_operator def empty_domain_fieldop(a: cases.IJField): diff --git a/tests/next_tests/integration_tests/multi_feature_tests/ffront_tests/test_icon_like_scan.py b/tests/next_tests/integration_tests/multi_feature_tests/ffront_tests/test_icon_like_scan.py index 108ee25862..eaae9a2a3e 100644 --- a/tests/next_tests/integration_tests/multi_feature_tests/ffront_tests/test_icon_like_scan.py +++ b/tests/next_tests/integration_tests/multi_feature_tests/ffront_tests/test_icon_like_scan.py @@ -25,6 +25,9 @@ ) +pytestmark = pytest.mark.uses_unstructured_shift + + Cell = gtx.Dimension("Cell") KDim = gtx.Dimension("KDim", kind=gtx.DimensionKind.VERTICAL) Koff = gtx.FieldOffset("Koff", KDim, (KDim,)) diff --git a/tests/next_tests/integration_tests/multi_feature_tests/ffront_tests/test_laplacian.py b/tests/next_tests/integration_tests/multi_feature_tests/ffront_tests/test_laplacian.py index d275a977dd..9a1e968de0 100644 --- a/tests/next_tests/integration_tests/multi_feature_tests/ffront_tests/test_laplacian.py +++ b/tests/next_tests/integration_tests/multi_feature_tests/ffront_tests/test_laplacian.py @@ -13,6 +13,7 @@ # SPDX-License-Identifier: GPL-3.0-or-later import numpy as np +import pytest import gt4py.next as gtx @@ -23,6 +24,9 @@ ) +pytestmark = pytest.mark.uses_cartesian_shift + + @gtx.field_operator def lap(in_field: gtx.Field[[IDim, JDim], "float"]) -> gtx.Field[[IDim, JDim], "float"]: return ( From f8682ed2d677fbe336287a27a685507ac9408f75 Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Thu, 16 Nov 2023 16:20:56 +0000 Subject: [PATCH 07/29] fix typo --- src/gt4py/next/allocators.py | 2 +- src/gt4py/next/constructors.py | 2 ++ 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/src/gt4py/next/allocators.py b/src/gt4py/next/allocators.py index cbb5da19f0..58600d8cda 100644 --- a/src/gt4py/next/allocators.py +++ b/src/gt4py/next/allocators.py @@ -331,7 +331,7 @@ def allocate( """ if device is None and allocator is None: - allocator = StandardCPUFieldBufferAllocator() + raise ValueError("No 'device' or 'allocator' specified") actual_allocator = get_allocator(allocator) if actual_allocator is None: assert device is not None # for mypy diff --git a/src/gt4py/next/constructors.py b/src/gt4py/next/constructors.py index 30ef8452aa..42b0bcda90 100644 --- a/src/gt4py/next/constructors.py +++ b/src/gt4py/next/constructors.py @@ -82,6 +82,8 @@ def empty( (3, 3) """ dtype = core_defs.dtype(dtype) + if allocator is None and device is None: + device = core_defs.Device(core_defs.DeviceType.CPU, device_id=0) buffer = next_allocators.allocate( domain, dtype, aligned_index=aligned_index, allocator=allocator, device=device ) From 42805f710ccc15454b41693ee81eaf07a931648a Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Thu, 16 Nov 2023 16:21:18 +0000 Subject: [PATCH 08/29] this is the typo, the other was improve alloc --- src/gt4py/next/common.py | 2 +- src/gt4py/next/iterator/embedded.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/gt4py/next/common.py b/src/gt4py/next/common.py index da6c60be3f..e93f970dd5 100644 --- a/src/gt4py/next/common.py +++ b/src/gt4py/next/common.py @@ -498,7 +498,7 @@ def __eq__(self, other: Any) -> Field: ... @abc.abstractmethod - def __neq__(self, other: Any) -> Field: + def __ne__(self, other: Any) -> Field: ... @abc.abstractmethod diff --git a/src/gt4py/next/iterator/embedded.py b/src/gt4py/next/iterator/embedded.py index d7abbbdffc..5ba94dec44 100644 --- a/src/gt4py/next/iterator/embedded.py +++ b/src/gt4py/next/iterator/embedded.py @@ -1096,7 +1096,7 @@ def __invert__(self) -> common.Field: def __eq__(self, other: Any) -> common.Field: raise NotImplementedError() - def __neq__(self, other: Any) -> common.Field: + def __ne__(self, other: Any) -> common.Field: raise NotImplementedError() def __add__(self, other: common.Field | core_defs.ScalarT) -> common.Field: @@ -1203,7 +1203,7 @@ def __invert__(self) -> common.Field: def __eq__(self, other: Any) -> common.Field: raise NotImplementedError() - def __neq__(self, other: Any) -> common.Field: + def __ne__(self, other: Any) -> common.Field: raise NotImplementedError() def __add__(self, other: common.Field | core_defs.ScalarT) -> common.Field: From ec0a0d582180e6e19b11bcb799a5afffc9eec7ac Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Thu, 16 Nov 2023 16:35:07 +0000 Subject: [PATCH 09/29] cleanup import in fbuiltin --- src/gt4py/next/ffront/fbuiltins.py | 67 +++++++++++++++++------------- 1 file changed, 38 insertions(+), 29 deletions(-) diff --git a/src/gt4py/next/ffront/fbuiltins.py b/src/gt4py/next/ffront/fbuiltins.py index 4a7db4d0dc..23d772ae3e 100644 --- a/src/gt4py/next/ffront/fbuiltins.py +++ b/src/gt4py/next/ffront/fbuiltins.py @@ -34,7 +34,7 @@ from gt4py._core import definitions as core_defs from gt4py.next import common -from gt4py.next.common import Dimension, DimensionKind, Field # TODO fix import style +from gt4py.next import Dimension, Field # direct import for TYPE_BUILTINS from gt4py.next.ffront.experimental import as_offset # noqa F401 from gt4py.next.iterator import runtime from gt4py.next.type_system import type_specifications as ts @@ -43,7 +43,14 @@ PYTHON_TYPE_BUILTINS = [bool, int, float, tuple] PYTHON_TYPE_BUILTIN_NAMES = [t.__name__ for t in PYTHON_TYPE_BUILTINS] -TYPE_BUILTINS = [Field, Dimension, int32, int64, float32, float64] + PYTHON_TYPE_BUILTINS +TYPE_BUILTINS = [ + Field, + Dimension, + int32, + int64, + float32, + float64, +] + PYTHON_TYPE_BUILTINS TYPE_BUILTIN_NAMES = [t.__name__ for t in TYPE_BUILTINS] # Be aware: Type aliases are not fully supported in the frontend yet, e.g. `IndexType(1)` will not @@ -57,9 +64,9 @@ def _type_conversion_helper(t: type) -> type[ts.TypeSpec] | tuple[type[ts.TypeSpec], ...]: - if t is Field: + if t is common.Field: return ts.FieldType - elif t is Dimension: + elif t is common.Dimension: return ts.DimensionType elif t is core_defs.ScalarT: return ts.ScalarType @@ -131,8 +138,8 @@ def __gt_type__(self) -> ts.FunctionType: ) -MaskT = TypeVar("MaskT", bound=Field) -FieldT = TypeVar("FieldT", bound=Union[Field, core_defs.Scalar, Tuple]) +MaskT = TypeVar("MaskT", bound=common.Field) +FieldT = TypeVar("FieldT", bound=Union[common.Field, core_defs.Scalar, Tuple]) class WhereBuiltinFunction( @@ -154,33 +161,35 @@ def __call__(self, mask: MaskT, true_field: FieldT, false_field: FieldT) -> _R: @BuiltInFunction def neighbor_sum( - field: Field, + field: common.Field, /, - axis: Dimension, -) -> Field: + axis: common.Dimension, +) -> common.Field: raise NotImplementedError() @BuiltInFunction def max_over( - field: Field, + field: common.Field, /, - axis: Dimension, -) -> Field: + axis: common.Dimension, +) -> common.Field: raise NotImplementedError() @BuiltInFunction def min_over( - field: Field, + field: common.Field, /, - axis: Dimension, -) -> Field: + axis: common.Dimension, +) -> common.Field: raise NotImplementedError() @BuiltInFunction -def broadcast(field: Field | core_defs.ScalarT, dims: tuple[Dimension, ...], /) -> Field: +def broadcast( + field: common.Field | core_defs.ScalarT, dims: tuple[common.Dimension, ...], / +) -> common.Field: assert core_defs.is_scalar_type(field) return common.field( np.asarray(field)[ @@ -192,16 +201,16 @@ def broadcast(field: Field | core_defs.ScalarT, dims: tuple[Dimension, ...], /) @WhereBuiltinFunction def where( - mask: Field, - true_field: Field | core_defs.ScalarT | Tuple, - false_field: Field | core_defs.ScalarT | Tuple, + mask: common.Field, + true_field: common.Field | core_defs.ScalarT | Tuple, + false_field: common.Field | core_defs.ScalarT | Tuple, /, -) -> Field | Tuple: +) -> common.Field | Tuple: raise NotImplementedError() @BuiltInFunction -def astype(field: Field | core_defs.ScalarT, type_: type, /) -> Field: +def astype(field: common.Field | core_defs.ScalarT, type_: type, /) -> common.Field: assert core_defs.is_scalar_type(field) return type_(field) @@ -235,7 +244,7 @@ def astype(field: Field | core_defs.ScalarT, type_: type, /) -> Field: def _make_unary_math_builtin(name): - def impl(value: Field | core_defs.ScalarT, /) -> Field | core_defs.ScalarT: + def impl(value: common.Field | core_defs.ScalarT, /) -> common.Field | core_defs.ScalarT: # TODO(havogt): enable once we have a failing test (see `test_math_builtin_execution.py`) # assert core_defs.is_scalar_type(value) # noqa: E800 # commented code # return getattr(math, name)(value)# noqa: E800 # commented code @@ -263,10 +272,10 @@ def impl(value: Field | core_defs.ScalarT, /) -> Field | core_defs.ScalarT: def _make_binary_math_builtin(name): def impl( - lhs: Field | core_defs.ScalarT, - rhs: Field | core_defs.ScalarT, + lhs: common.Field | core_defs.ScalarT, + rhs: common.Field | core_defs.ScalarT, /, - ) -> Field | core_defs.ScalarT: + ) -> common.Field | core_defs.ScalarT: assert core_defs.is_scalar_type(lhs) assert core_defs.is_scalar_type(rhs) return BINARY_MATH_NUMBER_BUILTIN_TO_PYTHON_SCALAR_FUNCTION[name](lhs, rhs) # type: ignore[operator] # Cannot call function of unknown type @@ -308,13 +317,13 @@ def impl( # guidelines for decision. @dataclasses.dataclass(frozen=True) class FieldOffset(runtime.Offset): - source: Dimension - target: tuple[Dimension] | tuple[Dimension, Dimension] + source: common.Dimension + target: tuple[common.Dimension] | tuple[common.Dimension, common.Dimension] connectivity: Optional[Any] = None # TODO def __post_init__(self): - if len(self.target) == 2 and self.target[1].kind != DimensionKind.LOCAL: - raise ValueError("Second dimension in offset must be a local dimension.") + if len(self.target) == 2 and self.target[1].kind != common.DimensionKind.LOCAL: + raise ValueError("Second Dimension in offset must be a local Dimension.") def __gt_type__(self): return ts.OffsetType(source=self.source, target=self.target) From ac28ea0a2471d9bcd9423c3ea3e6575fd51e82f1 Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Thu, 16 Nov 2023 17:02:48 +0000 Subject: [PATCH 10/29] fix test case --- .../feature_tests/ffront_tests/test_program.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_program.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_program.py index 8f7d3faf81..545abd2825 100644 --- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_program.py +++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_program.py @@ -248,4 +248,4 @@ def empty_domain_program(a: cases.IJField, out_field: cases.IJField): ValueError, match=(r"Dimensions in out field and field domain are not equivalent"), ): - empty_domain_program(a, out_field, offset_provider={}) + cases.run(cartesian_case, empty_domain_program, a, out_field, offset_provider={}) From 89e05ea4136ab001e98924342c3155946d236658 Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Thu, 16 Nov 2023 18:54:23 +0000 Subject: [PATCH 11/29] fix/ignore typing --- src/gt4py/next/common.py | 4 ++-- src/gt4py/next/embedded/nd_array_field.py | 4 ++-- src/gt4py/next/ffront/fbuiltins.py | 6 ++++-- src/gt4py/next/iterator/embedded.py | 8 ++++---- 4 files changed, 12 insertions(+), 10 deletions(-) diff --git a/src/gt4py/next/common.py b/src/gt4py/next/common.py index e93f970dd5..c4d9c64524 100644 --- a/src/gt4py/next/common.py +++ b/src/gt4py/next/common.py @@ -494,11 +494,11 @@ def __invert__(self) -> Field: """Only defined for `Field` of value type `bool`.""" @abc.abstractmethod - def __eq__(self, other: Any) -> Field: + def __eq__(self, other: Any) -> Field: # type: ignore[override] # mypy wants return `bool` ... @abc.abstractmethod - def __ne__(self, other: Any) -> Field: + def __ne__(self, other: Any) -> Field: # type: ignore[override] # mypy wants return `bool` ... @abc.abstractmethod diff --git a/src/gt4py/next/embedded/nd_array_field.py b/src/gt4py/next/embedded/nd_array_field.py index a26b8612c1..e1695686b0 100644 --- a/src/gt4py/next/embedded/nd_array_field.py +++ b/src/gt4py/next/embedded/nd_array_field.py @@ -191,9 +191,9 @@ def restrict(self, index: common.AnyIndexSpec) -> common.Field | core_defs.Scala __mod__ = __rmod__ = _make_builtin("mod", "mod") - __ne__ = _make_builtin("not_equal", "not_equal") + __ne__ = _make_builtin("not_equal", "not_equal") # type: ignore[assignment] # mypy wants return `bool` - __eq__ = _make_builtin("equal", "equal") + __eq__ = _make_builtin("equal", "equal") # type: ignore[assignment] # mypy wants return `bool` def __and__(self, other: common.Field | core_defs.ScalarT) -> NdArrayField: if self.dtype == core_defs.BoolDType(): diff --git a/src/gt4py/next/ffront/fbuiltins.py b/src/gt4py/next/ffront/fbuiltins.py index 23d772ae3e..38f9c7809f 100644 --- a/src/gt4py/next/ffront/fbuiltins.py +++ b/src/gt4py/next/ffront/fbuiltins.py @@ -34,7 +34,7 @@ from gt4py._core import definitions as core_defs from gt4py.next import common -from gt4py.next import Dimension, Field # direct import for TYPE_BUILTINS +from gt4py.next.common import Dimension, Field # direct import for TYPE_BUILTINS from gt4py.next.ffront.experimental import as_offset # noqa F401 from gt4py.next.iterator import runtime from gt4py.next.type_system import type_specifications as ts @@ -188,7 +188,9 @@ def min_over( @BuiltInFunction def broadcast( - field: common.Field | core_defs.ScalarT, dims: tuple[common.Dimension, ...], / + field: common.Field | core_defs.ScalarT, + dims: tuple[common.Dimension, ...], + /, ) -> common.Field: assert core_defs.is_scalar_type(field) return common.field( diff --git a/src/gt4py/next/iterator/embedded.py b/src/gt4py/next/iterator/embedded.py index 5ba94dec44..44294a3a71 100644 --- a/src/gt4py/next/iterator/embedded.py +++ b/src/gt4py/next/iterator/embedded.py @@ -1093,10 +1093,10 @@ def __neg__(self) -> common.Field: def __invert__(self) -> common.Field: raise NotImplementedError() - def __eq__(self, other: Any) -> common.Field: + def __eq__(self, other: Any) -> common.Field: # type: ignore[override] # mypy wants return `bool` raise NotImplementedError() - def __ne__(self, other: Any) -> common.Field: + def __ne__(self, other: Any) -> common.Field: # type: ignore[override] # mypy wants return `bool` raise NotImplementedError() def __add__(self, other: common.Field | core_defs.ScalarT) -> common.Field: @@ -1200,10 +1200,10 @@ def __neg__(self) -> common.Field: def __invert__(self) -> common.Field: raise NotImplementedError() - def __eq__(self, other: Any) -> common.Field: + def __eq__(self, other: Any) -> common.Field: # type: ignore[override] # mypy wants return `bool` raise NotImplementedError() - def __ne__(self, other: Any) -> common.Field: + def __ne__(self, other: Any) -> common.Field: # type: ignore[override] # mypy wants return `bool` raise NotImplementedError() def __add__(self, other: common.Field | core_defs.ScalarT) -> common.Field: From cdbaf0bfed8051dac11964fb2127b07b1468a112 Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Thu, 16 Nov 2023 21:05:11 +0000 Subject: [PATCH 12/29] remove __array__ from field --- src/gt4py/next/common.py | 3 +++ src/gt4py/next/embedded/nd_array_field.py | 4 ++-- tests/next_tests/integration_tests/cases.py | 18 ++++++++++++------ .../ffront_tests/test_arg_call_interface.py | 8 ++++---- .../ffront_tests/test_execution.py | 17 +++++------------ 5 files changed, 26 insertions(+), 24 deletions(-) diff --git a/src/gt4py/next/common.py b/src/gt4py/next/common.py index c4d9c64524..d6fe58cadb 100644 --- a/src/gt4py/next/common.py +++ b/src/gt4py/next/common.py @@ -461,6 +461,9 @@ def dtype(self) -> core_defs.DType[core_defs.ScalarT]: def ndarray(self) -> core_defs.NDArrayObject: ... + def asnumpy(self) -> np.ndarray: + return np.asarray(self.ndarray) + def __str__(self) -> str: return f"⟨{self.domain!s} → {self.dtype}⟩" diff --git a/src/gt4py/next/embedded/nd_array_field.py b/src/gt4py/next/embedded/nd_array_field.py index e1695686b0..2346c1e625 100644 --- a/src/gt4py/next/embedded/nd_array_field.py +++ b/src/gt4py/next/embedded/nd_array_field.py @@ -120,8 +120,8 @@ def __gt_origin__(self) -> tuple[int, ...]: def ndarray(self) -> core_defs.NDArrayObject: return self._ndarray - def __array__(self, dtype: npt.DTypeLike = None) -> np.ndarray: - return np.asarray(self._ndarray, dtype) + def asnumpy(self) -> np.ndarray: + return np.asarray(self._ndarray) @property def dtype(self) -> core_defs.DType[core_defs.ScalarT]: diff --git a/tests/next_tests/integration_tests/cases.py b/tests/next_tests/integration_tests/cases.py index 634d85e64c..e80b174c6d 100644 --- a/tests/next_tests/integration_tests/cases.py +++ b/tests/next_tests/integration_tests/cases.py @@ -377,6 +377,13 @@ def run( fieldview_prog.with_grid_type(case.grid_type).with_backend(case.backend)(*args, **kwargs) +def _asndarray(fields: tuple[common.Field | np.ndarray | tuple]) -> tuple[np.ndarray | tuple]: + if isinstance(fields, tuple): + return tuple(_asndarray(f) for f in fields) + else: + return fields.asnumpy() if hasattr(fields, "asnumpy") else fields + + def verify( case: Case, fieldview_prog: decorator.FieldOperator | decorator.Program, @@ -427,14 +434,13 @@ def verify( run(case, fieldview_prog, *args, offset_provider=offset_provider) out_comp = out or inout - out_comp_str = str(out_comp) assert out_comp is not None - if hasattr(out_comp, "ndarray"): - out_comp_str = str(out_comp.ndarray) - assert comparison(ref, out_comp), ( + out_comp_ndarray = _asndarray(out_comp) + ref_ndarray = _asndarray(ref) + assert comparison(ref_ndarray, out_comp_ndarray), ( f"Verification failed:\n" f"\tcomparison={comparison.__name__}(ref, out)\n" - f"\tref = {ref}\n\tout = {out_comp_str}" + f"\tref = {ref_ndarray}\n\tout = {str(out_comp_ndarray)}" ) @@ -460,7 +466,7 @@ def verify_with_default_data( ``comparison(ref, )`` and should return a boolean. """ inps, kwfields = get_default_data(case, fieldop) - ref_args = tuple(i.ndarray if hasattr(i, "ndarray") else i for i in inps) + ref_args = tuple(i.asnumpy() if hasattr(i, "asnumpy") else i for i in inps) verify( case, fieldop, diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_arg_call_interface.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_arg_call_interface.py index 6957e628bb..2f94301f38 100644 --- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_arg_call_interface.py +++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_arg_call_interface.py @@ -63,9 +63,9 @@ def testee(a: IField, b: IField, c: IField) -> IField: *pos_args, **kw_args, out=out, offset_provider=cartesian_case.offset_provider ) - expected = np.asarray(args["a"]) * 2 * np.asarray(args["b"]) - np.asarray(args["c"]) + expected = args["a"] * 2 * args["b"] - args["c"] - assert np.allclose(out, expected) + assert np.allclose(out.asnumpy(), expected.asnumpy()) @pytest.mark.parametrize("arg_spec", _generate_arg_permutations(("a", "b", "out"))) @@ -89,9 +89,9 @@ def testee(a: IField, b: IField, out: IField): *pos_args, **kw_args, offset_provider=cartesian_case.offset_provider ) - expected = np.asarray(args["a"]) + 2 * np.asarray(args["b"]) + expected = args["a"] + 2 * args["b"] - assert np.allclose(args["out"], expected) + assert np.allclose(args["out"].asnumpy(), expected.asnumpy()) def test_call_field_operator_from_field_operator(cartesian_case): diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_execution.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_execution.py index 22154da9a7..94baafdbaf 100644 --- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_execution.py +++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_execution.py @@ -773,7 +773,7 @@ def program_domain(a: cases.IField, out: cases.IField): out = cases.allocate(cartesian_case, program_domain, "out")() ref = out.ndarray.copy() # ensure we are not overwriting out outside of the domain - ref[1:9] = a[1:9] * 2 + ref[1:9] = a.asnumpy()[1:9] * 2 cases.verify(cartesian_case, program_domain, a, out, inout=out, ref=ref) @@ -807,7 +807,7 @@ def program_domain( out = cases.allocate(cartesian_case, fieldop_domain, cases.RETURN)() ref = out.ndarray.copy() - ref[lower_i : int(upper_i / 2)] = inp[lower_i : int(upper_i / 2)] * 2 + ref[lower_i : int(upper_i / 2)] = inp.asnumpy()[lower_i : int(upper_i / 2)] * 2 cases.verify( cartesian_case, @@ -851,7 +851,7 @@ def program_domain( ref = out.ndarray.copy() ref[1 * lower_i : upper_i + 0, lower_j - 0 : upper_j] = ( - a[1 * lower_i : upper_i + 0, lower_j - 0 : upper_j] * 2 + a.asnumpy()[1 * lower_i : upper_i + 0, lower_j - 0 : upper_j] * 2 ) cases.verify( @@ -890,9 +890,9 @@ def program_domain_tuple( out1 = cases.allocate(cartesian_case, program_domain_tuple, "out1")() ref0 = out0.ndarray.copy() - ref0[1:9, 4:6] = inp0[1:9, 4:6] + inp1[1:9, 4:6] + ref0[1:9, 4:6] = inp0.asnumpy()[1:9, 4:6] + inp1.asnumpy()[1:9, 4:6] ref1 = out1.ndarray.copy() - ref1[1:9, 4:6] = inp1[1:9, 4:6] + ref1[1:9, 4:6] = inp1.asnumpy()[1:9, 4:6] cases.verify( cartesian_case, @@ -1044,13 +1044,6 @@ def _invalid_unpack() -> tuple[int32, float64, int32]: def test_constant_closure_vars(cartesian_case): - if cartesian_case.backend is None: - # >>> field = gtx.zeros(domain) - # >>> np.int32(1)*field # steals the buffer from the field - # array([0.]) - - # TODO(havogt): remove `__array__`` from `NdArrayField` - pytest.xfail("Bug: Binary operation between np datatype and Field returns ndarray.") from gt4py.eve.utils import FrozenNamespace constants = FrozenNamespace( From 43cb189fa0e8ec9597e87e36d5d937bee3dc328d Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Thu, 16 Nov 2023 21:14:59 +0000 Subject: [PATCH 13/29] some more... --- .../ffront_tests/test_gt4py_builtins.py | 14 +++++++++++--- .../feature_tests/test_util_cases.py | 18 +++++++++--------- 2 files changed, 20 insertions(+), 12 deletions(-) diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_gt4py_builtins.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_gt4py_builtins.py index 1eba95e880..529ff45b24 100644 --- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_gt4py_builtins.py +++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_gt4py_builtins.py @@ -157,8 +157,8 @@ def conditional_nested_tuple( b, out=cases.allocate(cartesian_case, conditional_nested_tuple, cases.RETURN)(), ref=np.where( - mask, - ((a, b), (b, a)), + mask.asnumpy(), + ((a.asnumpy(), b.asnumpy()), (b.asnumpy(), a.asnumpy())), ((np.full(size, 5.0), np.full(size, 7.0)), (np.full(size, 7.0), np.full(size, 5.0))), ), ) @@ -224,7 +224,15 @@ def conditional( b = cases.allocate(cartesian_case, conditional, "b")() out = cases.allocate(cartesian_case, conditional, cases.RETURN)() - cases.verify(cartesian_case, conditional, mask, a, b, out=out, ref=np.where(mask, a, b)) + cases.verify( + cartesian_case, + conditional, + mask, + a, + b, + out=out, + ref=np.where(mask.asnumpy(), a.asnumpy(), b.asnumpy()), + ) def test_conditional_promotion(cartesian_case): diff --git a/tests/next_tests/integration_tests/feature_tests/test_util_cases.py b/tests/next_tests/integration_tests/feature_tests/test_util_cases.py index 3f229ef389..579dec11f8 100644 --- a/tests/next_tests/integration_tests/feature_tests/test_util_cases.py +++ b/tests/next_tests/integration_tests/feature_tests/test_util_cases.py @@ -41,30 +41,30 @@ def mixed_args( def test_allocate_default_unique(cartesian_case): # noqa: F811 # fixtures a = cases.allocate(cartesian_case, mixed_args, "a")() - assert np.min(a) == 0 - assert np.max(a) == np.prod(tuple(cartesian_case.default_sizes.values())) - 1 + assert np.min(a.asnumpy()) == 0 + assert np.max(a.asnumpy()) == np.prod(tuple(cartesian_case.default_sizes.values())) - 1 b = cases.allocate(cartesian_case, mixed_args, "b")() - assert b == np.max(a) + 1 + assert b == np.max(a.asnumpy()) + 1 c = cases.allocate(cartesian_case, mixed_args, "c")() - assert np.min(c) == b + 1 - assert np.max(c) == np.prod(tuple(cartesian_case.default_sizes.values())) * 2 + assert np.min(c.asnumpy()) == b + 1 + assert np.max(c.asnumpy()) == np.prod(tuple(cartesian_case.default_sizes.values())) * 2 def test_allocate_return_default_zeros(cartesian_case): # noqa: F811 # fixtures a, (b, c) = cases.allocate(cartesian_case, mixed_args, cases.RETURN)() - assert np.all(np.asarray(a) == 0) - assert np.all(np.asarray(a) == b) - assert np.all(np.asarray(b) == c) + assert np.all(a.asnumpy() == 0) + assert np.all(b.asnumpy() == 0) + assert np.all(c.asnumpy() == 0) def test_allocate_const(cartesian_case): # noqa: F811 # fixtures a = cases.allocate(cartesian_case, mixed_args, "a").strategy(cases.ConstInitializer(42))() - assert np.all(np.asarray(a) == 42) + assert np.all(a.asnumpy() == 42) b = cases.allocate(cartesian_case, mixed_args, "b").strategy(cases.ConstInitializer(42))() assert b == 42.0 From 520da0d4d69e9b942aca4fb27d0d1995a13d6c6e Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Thu, 16 Nov 2023 21:19:28 +0000 Subject: [PATCH 14/29] fix some iterator tests --- .../feature_tests/iterator_tests/test_builtins.py | 12 ++++++------ .../iterator_tests/test_horizontal_indirection.py | 4 ++-- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/tests/next_tests/integration_tests/feature_tests/iterator_tests/test_builtins.py b/tests/next_tests/integration_tests/feature_tests/iterator_tests/test_builtins.py index d5d57c9024..ff153cc812 100644 --- a/tests/next_tests/integration_tests/feature_tests/iterator_tests/test_builtins.py +++ b/tests/next_tests/integration_tests/feature_tests/iterator_tests/test_builtins.py @@ -178,7 +178,7 @@ def test_arithmetic_and_logical_builtins(program_processor, builtin, inputs, exp fencil(builtin, out, *inps, processor=program_processor, as_column=as_column) if validate: - assert np.allclose(np.asarray(out), expected) + assert np.allclose(out.asnumpy(), expected) @pytest.mark.parametrize("builtin, inputs, expected", arithmetic_and_logical_test_data()) @@ -199,7 +199,7 @@ def test_arithmetic_and_logical_functors_gtfn(builtin, inputs, expected): ) # avoid inlining the function fencil(builtin, out, *inps, processor=gtfn_without_transforms) - assert np.allclose(np.asarray(out), expected) + assert np.allclose(out.asnumpy(), expected) @pytest.mark.parametrize("as_column", [False, True]) @@ -228,7 +228,7 @@ def test_math_function_builtins(program_processor, builtin_name, inputs, as_colu ) if validate: - assert np.allclose(np.asarray(out), expected) + assert np.allclose(out.asnumpy(), expected) Neighbor = offset("Neighbor") @@ -269,7 +269,7 @@ def test_can_deref(program_processor, stencil): ) if validate: - assert np.allclose(np.asarray(out), -1.0) + assert np.allclose(out.asnumpy(), -1.0) a_neighbor_tbl = gtx.NeighborTableOffsetProvider(np.array([[0]]), Node, Node, 1) run_processor( @@ -281,7 +281,7 @@ def test_can_deref(program_processor, stencil): ) if validate: - assert np.allclose(np.asarray(out), 1.0) + assert np.allclose(out.asnumpy(), 1.0) # def test_can_deref_lifted(program_processor): @@ -337,7 +337,7 @@ def test_cast(program_processor, as_column, input_value, dtype, np_dtype): def sten_cast(it, casted_valued): return eq(cast_(deref(it), dtype), deref(casted_valued)) - out = field_maker(np.zeros_like(inp, dtype=builtins.bool))[0] + out = field_maker(np.zeros_like(inp.asnumpy(), dtype=builtins.bool))[0] run_processor( sten_cast[{IDim: range(1)}], program_processor, diff --git a/tests/next_tests/integration_tests/feature_tests/iterator_tests/test_horizontal_indirection.py b/tests/next_tests/integration_tests/feature_tests/iterator_tests/test_horizontal_indirection.py index f9bd2cc33b..69f594a2bc 100644 --- a/tests/next_tests/integration_tests/feature_tests/iterator_tests/test_horizontal_indirection.py +++ b/tests/next_tests/integration_tests/feature_tests/iterator_tests/test_horizontal_indirection.py @@ -82,7 +82,7 @@ def test_simple_indirection(program_processor): ) if validate: - assert np.allclose(ref, out) + assert np.allclose(ref, out.asnumpy()) @fundef @@ -113,4 +113,4 @@ def test_direct_offset_for_indirection(program_processor): ) if validate: - assert np.allclose(ref, out) + assert np.allclose(ref, out.asnumpy()) From 1e7b0edc0071fde217fd39feb973acdbefcafe71 Mon Sep 17 00:00:00 2001 From: nfarabullini Date: Fri, 17 Nov 2023 14:47:03 +0100 Subject: [PATCH 15/29] added asnumpy where needed --- .../feature_tests/test_call_interface.py | 4 +-- .../ffront_tests/test_arg_call_interface.py | 6 ++--- .../ffront_tests/test_execution.py | 4 +-- .../ffront_tests/test_gt4py_builtins.py | 5 ++-- .../test_math_builtin_execution.py | 2 +- .../ffront_tests/test_program.py | 10 ++++--- .../ffront_tests/test_scalar_if.py | 8 +++--- .../iterator_tests/test_implicit_fencil.py | 6 ++--- .../feature_tests/iterator_tests/test_scan.py | 2 +- .../iterator_tests/test_trivial.py | 8 +++--- .../iterator_tests/test_tuple.py | 26 +++++++++---------- .../ffront_tests/test_icon_like_scan.py | 12 ++++----- .../iterator_tests/test_anton_toy.py | 2 +- .../iterator_tests/test_column_stencil.py | 6 ++--- .../iterator_tests/test_hdiff.py | 2 +- .../iterator_tests/test_vertical_advection.py | 2 +- .../test_with_toy_connectivity.py | 26 +++++++++---------- .../otf_tests/test_gtfn_workflow.py | 2 +- 18 files changed, 67 insertions(+), 66 deletions(-) diff --git a/tests/cartesian_tests/integration_tests/feature_tests/test_call_interface.py b/tests/cartesian_tests/integration_tests/feature_tests/test_call_interface.py index b3af613fec..ef8a8a08e3 100644 --- a/tests/cartesian_tests/integration_tests/feature_tests/test_call_interface.py +++ b/tests/cartesian_tests/integration_tests/feature_tests/test_call_interface.py @@ -186,7 +186,7 @@ def test_default_arguments(backend): arg3 = gt_storage.ones( backend=backend, dtype=np.float64, shape=(3, 3, 3), aligned_index=(0, 0, 0) ) - tmp = np.asarray(arg3) + tmp = arg3.asnumpy() tmp *= 2 branch_true(arg1, None, arg3, par1=2.0) @@ -208,7 +208,7 @@ def test_default_arguments(backend): arg3 = gt_storage.ones( backend=backend, dtype=np.float64, shape=(3, 3, 3), aligned_index=(0, 0, 0) ) - tmp = np.asarray(arg3) + tmp = arg3.asnumpy() tmp *= 2 branch_true(arg1, arg2=None, par1=2.0, par2=5.0, par3=3.0) diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_arg_call_interface.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_arg_call_interface.py index 2f94301f38..6293ff76bd 100644 --- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_arg_call_interface.py +++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_arg_call_interface.py @@ -177,9 +177,7 @@ def testee(a: IJKFloatField, b: IJKFloatField) -> IJKFloatField: a, b, out = ( cases.allocate(cartesian_case, testee, name)() for name in ("a", "b", cases.RETURN) ) - expected = (1.0 + 3.0 + 5.0 + 7.0) * np.add.accumulate( - np.asarray(a) + 2.0 * np.asarray(b), axis=2 - ) + expected = (1.0 + 3.0 + 5.0 + 7.0) * np.add.accumulate(a.asnumpy() + 2.0 * b.asnumpy(), axis=2) cases.verify(cartesian_case, testee, a, b, out=out, ref=expected) @@ -210,7 +208,7 @@ def testee( for name in ("out1", "out2", "out3", "out4") ) - ref = np.add.accumulate(np.asarray(a) + 2 * np.asarray(b), axis=2) + ref = np.add.accumulate(a.asnumpy() + 2 * b.asnumpy(), axis=2) cases.verify( cartesian_case, diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_execution.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_execution.py index 94baafdbaf..dd3565645d 100644 --- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_execution.py +++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_execution.py @@ -403,7 +403,7 @@ def testee(a: cases.IKField, offset_field: cases.IKField) -> gtx.Field[[IDim, KD comparison=lambda out, ref: np.all(out == ref), ) - assert np.allclose(out, ref) + assert np.allclose(out.asnumpy(), ref) def test_nested_tuple_return(cartesian_case): @@ -920,7 +920,7 @@ def fieldop_where_k_offset( )() out = cases.allocate(cartesian_case, fieldop_where_k_offset, "inp")() - ref = np.where(np.asarray(k_index) > 0, np.roll(inp, 1, axis=1), 2) + ref = np.where(k_index.asnumpy() > 0, np.roll(inp.asnumpy(), 1, axis=1), 2) cases.verify(cartesian_case, fieldop_where_k_offset, inp, k_index, out=out, ref=ref) diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_gt4py_builtins.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_gt4py_builtins.py index 529ff45b24..ec74612c12 100644 --- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_gt4py_builtins.py +++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_gt4py_builtins.py @@ -247,10 +247,9 @@ def conditional_promotion(mask: cases.IBoolField, a: cases.IFloatField) -> cases )() a = cases.allocate(cartesian_case, conditional_promotion, "a")() out = cases.allocate(cartesian_case, conditional_promotion, cases.RETURN)() + ref = np.where(mask.asnumpy(), a.asnumpy(), 10.0) - cases.verify( - cartesian_case, conditional_promotion, mask, a, out=out, ref=np.where(mask, a, 10.0) - ) + cases.verify(cartesian_case, conditional_promotion, mask, a, out=out, ref=ref) def test_conditional_compareop(cartesian_case): diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_math_builtin_execution.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_math_builtin_execution.py index a1839b8e17..f401e86468 100644 --- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_math_builtin_execution.py +++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_math_builtin_execution.py @@ -135,4 +135,4 @@ def test_math_function_builtins_execution(cartesian_case, builtin_name: str, inp builtin_field_op(*inps, out=out, offset_provider={}) - assert np.allclose(np.asarray(out), expected) + assert np.allclose(out.asnumpy(), expected) diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_program.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_program.py index 545abd2825..7b0f3d8f60 100644 --- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_program.py +++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_program.py @@ -152,7 +152,7 @@ def prog( cases.run(cartesian_case, prog, a, b, out_a, out_b, offset_provider={}) - assert np.allclose((a, b), (out_a, out_b)) + assert np.allclose((a.asnumpy(), b.asnumpy()), (out_a.asnumpy(), out_b.asnumpy())) def test_tuple_program_return_constructed_inside_with_slicing(cartesian_case): @@ -178,7 +178,9 @@ def prog( cases.run(cartesian_case, prog, a, b, out_a, out_b, offset_provider={}) - assert np.allclose((a[1:], b[1:]), (out_a[1:], out_b[1:])) + assert np.allclose( + (a[1:].asnumpy(), b[1:].asnumpy()), (out_a[1:].asnumpy(), out_b[1:].asnumpy()) + ) assert out_a[0] == 0 and out_b[0] == 0 @@ -209,7 +211,9 @@ def prog( cases.run(cartesian_case, prog, a, b, c, out_a, out_b, out_c, offset_provider={}) - assert np.allclose((a, b, c), (out_a, out_b, out_c)) + assert np.allclose( + (a.asnumpy(), b.asnumpy(), c.asnumpy()), (out_a.asnumpy(), out_b.asnumpy(), out_c.asnumpy()) + ) def test_wrong_argument_type(cartesian_case, copy_program_def): diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_scalar_if.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_scalar_if.py index e9c3ac8d19..84b480a23d 100644 --- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_scalar_if.py +++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_scalar_if.py @@ -315,10 +315,10 @@ def if_without_else( out = cases.allocate(cartesian_case, if_without_else, cases.RETURN)() ref = { - (True, True): np.asarray(a) + 2, - (True, False): np.asarray(a), - (False, True): np.asarray(b) + 1, - (False, False): np.asarray(b) + 1, + (True, True): a.asnumpy() + 2, + (True, False): a.asnumpy(), + (False, True): b.asnumpy() + 1, + (False, False): b.asnumpy() + 1, } cases.verify( diff --git a/tests/next_tests/integration_tests/feature_tests/iterator_tests/test_implicit_fencil.py b/tests/next_tests/integration_tests/feature_tests/iterator_tests/test_implicit_fencil.py index 2df7691f9e..6f600414db 100644 --- a/tests/next_tests/integration_tests/feature_tests/iterator_tests/test_implicit_fencil.py +++ b/tests/next_tests/integration_tests/feature_tests/iterator_tests/test_implicit_fencil.py @@ -53,7 +53,7 @@ def test_single_argument(program_processor, dom): run_processor(copy_stencil[dom], program_processor, inp, out=out, offset_provider={}) if validate: - assert np.allclose(inp, out) + assert np.allclose(inp.asnumpy(), out.asnumpy()) def test_2_arguments(program_processor, dom): @@ -70,7 +70,7 @@ def fun(inp0, inp1): run_processor(fun[dom], program_processor, inp0, inp1, out=out, offset_provider={}) if validate: - assert np.allclose(inp0 + inp1, out) + assert np.allclose(inp0.asnumpy() + inp1.asnumpy(), out.asnumpy()) def test_lambda_domain(program_processor): @@ -82,4 +82,4 @@ def test_lambda_domain(program_processor): run_processor(copy_stencil[dom], program_processor, inp, out=out, offset_provider={}) if validate: - assert np.allclose(inp, out) + assert np.allclose(inp.asnumpy(), out.asnumpy()) diff --git a/tests/next_tests/integration_tests/feature_tests/iterator_tests/test_scan.py b/tests/next_tests/integration_tests/feature_tests/iterator_tests/test_scan.py index 3af0440c27..fce1aa3960 100644 --- a/tests/next_tests/integration_tests/feature_tests/iterator_tests/test_scan.py +++ b/tests/next_tests/integration_tests/feature_tests/iterator_tests/test_scan.py @@ -60,4 +60,4 @@ def wrapped(inp): ) if validate: - assert np.allclose(out[:, :-1], reference) + assert np.allclose(out[:, :-1].asnumpy(), reference) diff --git a/tests/next_tests/integration_tests/feature_tests/iterator_tests/test_trivial.py b/tests/next_tests/integration_tests/feature_tests/iterator_tests/test_trivial.py index 8c59f994ee..8e12647c1b 100644 --- a/tests/next_tests/integration_tests/feature_tests/iterator_tests/test_trivial.py +++ b/tests/next_tests/integration_tests/feature_tests/iterator_tests/test_trivial.py @@ -65,7 +65,7 @@ def test_trivial(program_processor, lift_mode): ) if validate: - assert np.allclose(out[:, :, 0], out_s) + assert np.allclose(out[:, :, 0], out_s.asnumpy()) @fundef @@ -100,7 +100,7 @@ def test_shifted_arg_to_lift(program_processor, lift_mode): ) if validate: - assert np.allclose(out, out_s) + assert np.allclose(out, out_s.asnumpy()) @fendef @@ -137,7 +137,7 @@ def test_direct_deref(program_processor, lift_mode): ) if validate: - assert np.allclose(out, out_s) + assert np.allclose(out, out_s.asnumpy()) @fundef @@ -167,4 +167,4 @@ def test_vertical_shift_unstructured(program_processor): ) if validate: - assert np.allclose(inp_s[:, 1:], np.asarray(out_s)[:, :-1]) + assert np.allclose(inp_s[:, 1:].asnumpy(), out_s[:, :-1].asnumpy()) diff --git a/tests/next_tests/integration_tests/feature_tests/iterator_tests/test_tuple.py b/tests/next_tests/integration_tests/feature_tests/iterator_tests/test_tuple.py index 97a51508f5..add772e7ef 100644 --- a/tests/next_tests/integration_tests/feature_tests/iterator_tests/test_tuple.py +++ b/tests/next_tests/integration_tests/feature_tests/iterator_tests/test_tuple.py @@ -76,8 +76,8 @@ def test_tuple_output(program_processor, stencil): } run_processor(stencil[dom], program_processor, inp1, inp2, out=out, offset_provider={}) if validate: - assert np.allclose(inp1, out[0]) - assert np.allclose(inp2, out[1]) + assert np.allclose(inp1.asnumpy(), out[0].asnumpy()) + assert np.allclose(inp2.asnumpy(), out[1].asnumpy()) @fundef @@ -144,10 +144,10 @@ def stencil(inp1, inp2, inp3, inp4): offset_provider={}, ) if validate: - assert np.allclose(inp1, out[0][0]) - assert np.allclose(inp2, out[0][1]) - assert np.allclose(inp3, out[1][0]) - assert np.allclose(inp4, out[1][1]) + assert np.allclose(inp1.asnumpy(), out[0][0].asnumpy()) + assert np.allclose(inp2.asnumpy(), out[0][1].asnumpy()) + assert np.allclose(inp3.asnumpy(), out[1][0].asnumpy()) + assert np.allclose(inp4.asnumpy(), out[1][1].asnumpy()) @pytest.mark.parametrize( @@ -197,8 +197,8 @@ def fencil(size0, size1, size2, inp1, inp2, out1, out2): offset_provider={}, ) if validate: - assert np.allclose(inp1, out1) - assert np.allclose(inp2, out2) + assert np.allclose(inp1.asnumpy(), out1.asnumpy()) + assert np.allclose(inp2.asnumpy(), out2.asnumpy()) def test_asymetric_nested_tuple_of_field_output_constructed_inside(program_processor): @@ -255,9 +255,9 @@ def fencil(size0, size1, size2, inp1, inp2, inp3, out1, out2, out3): offset_provider={}, ) if validate: - assert np.allclose(inp1, out1) - assert np.allclose(inp2, out2) - assert np.allclose(inp3, out3) + assert np.allclose(inp1.asnumpy(), out1.asnumpy()) + assert np.allclose(inp2.asnumpy(), out2.asnumpy()) + assert np.allclose(inp3.asnumpy(), out3.asnumpy()) @pytest.mark.xfail(reason="Implement wrapper for extradim as tuple") @@ -323,7 +323,7 @@ def test_tuple_field_input(program_processor): } run_processor(tuple_input[dom], program_processor, (inp1, inp2), out=out, offset_provider={}) if validate: - assert np.allclose(np.asarray(inp1) + np.asarray(inp2), out) + assert np.allclose(inp1.asnumpy() + inp2.asnumpy(), out.asnumpy()) @pytest.mark.xfail(reason="Implement wrapper for extradim as tuple") @@ -389,7 +389,7 @@ def test_tuple_of_tuple_of_field_input(program_processor): ) if validate: assert np.allclose( - (np.asarray(inp1) + np.asarray(inp2) + np.asarray(inp3) + np.asarray(inp4)), out + (inp1.asnumpy() + inp2.asnumpy() + inp3.asnumpy() + inp4.asnumpy()), out.asnumpy() ) diff --git a/tests/next_tests/integration_tests/multi_feature_tests/ffront_tests/test_icon_like_scan.py b/tests/next_tests/integration_tests/multi_feature_tests/ffront_tests/test_icon_like_scan.py index eaae9a2a3e..f09b45ee56 100644 --- a/tests/next_tests/integration_tests/multi_feature_tests/ffront_tests/test_icon_like_scan.py +++ b/tests/next_tests/integration_tests/multi_feature_tests/ffront_tests/test_icon_like_scan.py @@ -228,7 +228,7 @@ def test_solve_nonhydro_stencil_52_like_z_q(test_setup, fieldview_backend): offset_provider={"Koff": KDim}, ) - assert np.allclose(test_setup.z_q_ref[:, 1:], test_setup.z_q_out[:, 1:]) + assert np.allclose(test_setup.z_q_ref[:, 1:], test_setup.z_q_out[:, 1:].asnumpy()) @pytest.mark.uses_tuple_returns @@ -250,7 +250,7 @@ def test_solve_nonhydro_stencil_52_like_z_q_tup(test_setup, fieldview_backend): offset_provider={"Koff": KDim}, ) - assert np.allclose(test_setup.z_q_ref[:, 1:], test_setup.z_q_out[:, 1:]) + assert np.allclose(test_setup.z_q_ref[:, 1:], test_setup.z_q_out[:, 1:].asnumpy()) @pytest.mark.uses_tuple_returns @@ -266,8 +266,8 @@ def test_solve_nonhydro_stencil_52_like(test_setup, fieldview_backend): offset_provider={"Koff": KDim}, ) - assert np.allclose(test_setup.z_q_ref, test_setup.z_q) - assert np.allclose(test_setup.w_ref, test_setup.w) + assert np.allclose(test_setup.z_q_ref, test_setup.z_q.asnumpy()) + assert np.allclose(test_setup.w_ref, test_setup.w.asnumpy()) @pytest.mark.uses_tuple_returns @@ -285,5 +285,5 @@ def test_solve_nonhydro_stencil_52_like_with_gtfn_tuple_merge(test_setup, fieldv offset_provider={"Koff": KDim}, ) - assert np.allclose(test_setup.z_q_ref, test_setup.z_q) - assert np.allclose(test_setup.w_ref, test_setup.w) + assert np.allclose(test_setup.z_q_ref, test_setup.z_q.asnumpy()) + assert np.allclose(test_setup.w_ref, test_setup.w.asnumpy()) diff --git a/tests/next_tests/integration_tests/multi_feature_tests/iterator_tests/test_anton_toy.py b/tests/next_tests/integration_tests/multi_feature_tests/iterator_tests/test_anton_toy.py index 829bc497cb..806ab7eb9a 100644 --- a/tests/next_tests/integration_tests/multi_feature_tests/iterator_tests/test_anton_toy.py +++ b/tests/next_tests/integration_tests/multi_feature_tests/iterator_tests/test_anton_toy.py @@ -103,4 +103,4 @@ def test_anton_toy(program_processor, lift_mode): ) if validate: - assert np.allclose(out, ref) + assert np.allclose(out.asnumpy(), ref) diff --git a/tests/next_tests/integration_tests/multi_feature_tests/iterator_tests/test_column_stencil.py b/tests/next_tests/integration_tests/multi_feature_tests/iterator_tests/test_column_stencil.py index d05b14d73d..51bd039ab3 100644 --- a/tests/next_tests/integration_tests/multi_feature_tests/iterator_tests/test_column_stencil.py +++ b/tests/next_tests/integration_tests/multi_feature_tests/iterator_tests/test_column_stencil.py @@ -102,7 +102,7 @@ def test_basic_column_stencils(program_processor, lift_mode, basic_stencils): ) if validate: - assert np.allclose(ref, out) + assert np.allclose(ref, out.asnumpy()) @fundef @@ -392,7 +392,7 @@ def test_different_vertical_sizes_with_origin(program_processor): inp0 = gtx.as_field([KDim], np.arange(0, k_size)) inp1 = gtx.as_field([KDim], np.arange(0, k_size + 1), origin={KDim: 1}) out = gtx.as_field([KDim], np.zeros(k_size, dtype=np.int64)) - ref = np.asarray(inp0) + np.asarray(inp1)[:-1] + ref = inp0.asnumpy() + inp1.asnumpy()[:-1] run_processor( sum_fencil, @@ -405,7 +405,7 @@ def test_different_vertical_sizes_with_origin(program_processor): ) if validate: - assert np.allclose(ref, out) + assert np.allclose(ref, out.asnumpy()) # TODO(havogt) test tuple_get builtin on a Column diff --git a/tests/next_tests/integration_tests/multi_feature_tests/iterator_tests/test_hdiff.py b/tests/next_tests/integration_tests/multi_feature_tests/iterator_tests/test_hdiff.py index 8aabd18267..9bba1ab89c 100644 --- a/tests/next_tests/integration_tests/multi_feature_tests/iterator_tests/test_hdiff.py +++ b/tests/next_tests/integration_tests/multi_feature_tests/iterator_tests/test_hdiff.py @@ -97,4 +97,4 @@ def test_hdiff(hdiff_reference, program_processor, lift_mode): ) if validate: - assert np.allclose(out[:, :, 0], out_s) + assert np.allclose(out[:, :, 0], out_s.asnumpy()) diff --git a/tests/next_tests/integration_tests/multi_feature_tests/iterator_tests/test_vertical_advection.py b/tests/next_tests/integration_tests/multi_feature_tests/iterator_tests/test_vertical_advection.py index 29c82442ea..f2a6505a7e 100644 --- a/tests/next_tests/integration_tests/multi_feature_tests/iterator_tests/test_vertical_advection.py +++ b/tests/next_tests/integration_tests/multi_feature_tests/iterator_tests/test_vertical_advection.py @@ -158,4 +158,4 @@ def test_tridiag(fencil, tridiag_reference, program_processor, lift_mode): ) if validate: - assert np.allclose(x, x_s) + assert np.allclose(x, x_s.asnumpy()) diff --git a/tests/next_tests/integration_tests/multi_feature_tests/iterator_tests/test_with_toy_connectivity.py b/tests/next_tests/integration_tests/multi_feature_tests/iterator_tests/test_with_toy_connectivity.py index 6354e45451..000d3c4822 100644 --- a/tests/next_tests/integration_tests/multi_feature_tests/iterator_tests/test_with_toy_connectivity.py +++ b/tests/next_tests/integration_tests/multi_feature_tests/iterator_tests/test_with_toy_connectivity.py @@ -99,7 +99,7 @@ def test_sum_edges_to_vertices(program_processor, lift_mode, stencil): lift_mode=lift_mode, ) if validate: - assert np.allclose(out, ref) + assert np.allclose(out.asnumpy(), ref) @fundef @@ -122,7 +122,7 @@ def test_map_neighbors(program_processor, lift_mode): lift_mode=lift_mode, ) if validate: - assert np.allclose(out, ref) + assert np.allclose(out.asnumpy(), ref) @fundef @@ -146,7 +146,7 @@ def test_map_make_const_list(program_processor, lift_mode): lift_mode=lift_mode, ) if validate: - assert np.allclose(out, ref) + assert np.allclose(out.asnumpy(), ref) @fundef @@ -172,7 +172,7 @@ def test_first_vertex_neigh_of_first_edge_neigh_of_cells_fencil(program_processo lift_mode=lift_mode, ) if validate: - assert np.allclose(out, ref) + assert np.allclose(out.asnumpy(), ref) @fundef @@ -200,7 +200,7 @@ def test_sparse_input_field(program_processor, lift_mode): ) if validate: - assert np.allclose(out, ref) + assert np.allclose(out.asnumpy(), ref) def test_sparse_input_field_v2v(program_processor, lift_mode): @@ -226,7 +226,7 @@ def test_sparse_input_field_v2v(program_processor, lift_mode): ) if validate: - assert np.allclose(out, ref) + assert np.allclose(out.asnumpy(), ref) @fundef @@ -254,7 +254,7 @@ def test_slice_sparse(program_processor, lift_mode): ) if validate: - assert np.allclose(out, ref) + assert np.allclose(out.asnumpy(), ref) @fundef @@ -309,7 +309,7 @@ def test_shift_sliced_sparse(program_processor, lift_mode): ) if validate: - assert np.allclose(out, ref) + assert np.allclose(out.asnumpy(), ref) @fundef @@ -337,7 +337,7 @@ def test_slice_shifted_sparse(program_processor, lift_mode): ) if validate: - assert np.allclose(out, ref) + assert np.allclose(out.asnumpy(), ref) @fundef @@ -365,7 +365,7 @@ def test_lift(program_processor, lift_mode): lift_mode=lift_mode, ) if validate: - assert np.allclose(out, ref) + assert np.allclose(out.asnumpy(), ref) @fundef @@ -390,7 +390,7 @@ def test_shift_sparse_input_field(program_processor, lift_mode): ) if validate: - assert np.allclose(out, ref) + assert np.allclose(out.asnumpy(), ref) @fundef @@ -443,7 +443,7 @@ def test_shift_sparse_input_field2(program_processor, lift_mode): ) if validate: - assert np.allclose(out1, out2) + assert np.allclose(out1.asnumpy(), out2.asnumpy()) @fundef @@ -484,4 +484,4 @@ def test_sparse_shifted_stencil_reduce(program_processor, lift_mode): ) if validate: - assert np.allclose(np.asarray(out), ref) + assert np.allclose(out.asnumpy(), ref) diff --git a/tests/next_tests/integration_tests/multi_feature_tests/otf_tests/test_gtfn_workflow.py b/tests/next_tests/integration_tests/multi_feature_tests/otf_tests/test_gtfn_workflow.py index d851c5560a..c91be04999 100644 --- a/tests/next_tests/integration_tests/multi_feature_tests/otf_tests/test_gtfn_workflow.py +++ b/tests/next_tests/integration_tests/multi_feature_tests/otf_tests/test_gtfn_workflow.py @@ -41,4 +41,4 @@ def copy(inp: gtx.Field[[IDim, JDim], gtx.int32]) -> gtx.Field[[IDim, JDim], gtx copy(inp, out=out, offset_provider={}) - assert np.allclose(inp[:out_nx, :out_ny], out) + assert np.allclose(inp[:out_nx, :out_ny].asnumpy(), out.asnumpy()) From 5fbf8268436c50c0598b72278c91c935f2770319 Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Mon, 20 Nov 2023 11:19:09 +0000 Subject: [PATCH 16/29] remove test exclusion --- .../feature_tests/ffront_tests/test_execution.py | 7 ------- 1 file changed, 7 deletions(-) diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_execution.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_execution.py index ce44960604..52687edbb2 100644 --- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_execution.py +++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_execution.py @@ -1119,13 +1119,6 @@ def _invalid_unpack() -> tuple[int32, float64, int32]: def test_constant_closure_vars(cartesian_case): - if cartesian_case.backend is None: - # >>> field = gtx.zeros(domain) - # >>> np.int32(1)*field # steals the buffer from the field - # array([0.]) - - # TODO(havogt): remove `__array__`` from `NdArrayField` - pytest.xfail("Bug: Binary operation between np datatype and Field returns ndarray.") from gt4py.eve.utils import FrozenNamespace constants = FrozenNamespace( From 074265055385fa6900e129e3280318a278d6deeb Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Mon, 20 Nov 2023 11:20:19 +0000 Subject: [PATCH 17/29] undo unintendend change --- src/gt4py/next/ffront/decorator.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/gt4py/next/ffront/decorator.py b/src/gt4py/next/ffront/decorator.py index e7664d3c14..107415eb06 100644 --- a/src/gt4py/next/ffront/decorator.py +++ b/src/gt4py/next/ffront/decorator.py @@ -294,6 +294,10 @@ def __call__(self, *args, offset_provider: dict[str, Dimension], **kwargs) -> No self.definition(*rewritten_args, **kwargs) return + ppi.ensure_processor_kind(self.backend, ppi.ProgramExecutor) + if "debug" in kwargs: + debug(self.itir) + self.backend( self.itir, *rewritten_args, From 3b9e35dfdaa804f82bc53a9c8a091bda8de51f1c Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Mon, 20 Nov 2023 12:03:08 +0000 Subject: [PATCH 18/29] undo unrelated change --- .../integration_tests/feature_tests/test_call_interface.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/cartesian_tests/integration_tests/feature_tests/test_call_interface.py b/tests/cartesian_tests/integration_tests/feature_tests/test_call_interface.py index ef8a8a08e3..b3af613fec 100644 --- a/tests/cartesian_tests/integration_tests/feature_tests/test_call_interface.py +++ b/tests/cartesian_tests/integration_tests/feature_tests/test_call_interface.py @@ -186,7 +186,7 @@ def test_default_arguments(backend): arg3 = gt_storage.ones( backend=backend, dtype=np.float64, shape=(3, 3, 3), aligned_index=(0, 0, 0) ) - tmp = arg3.asnumpy() + tmp = np.asarray(arg3) tmp *= 2 branch_true(arg1, None, arg3, par1=2.0) @@ -208,7 +208,7 @@ def test_default_arguments(backend): arg3 = gt_storage.ones( backend=backend, dtype=np.float64, shape=(3, 3, 3), aligned_index=(0, 0, 0) ) - tmp = arg3.asnumpy() + tmp = np.asarray(arg3) tmp *= 2 branch_true(arg1, arg2=None, par1=2.0, par2=5.0, par3=3.0) From 8a6e6c20ca8154ab408a67b5b967739b5edfe328 Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Mon, 20 Nov 2023 12:06:09 +0000 Subject: [PATCH 19/29] undo change --- src/gt4py/next/ffront/fbuiltins.py | 9 +-------- 1 file changed, 1 insertion(+), 8 deletions(-) diff --git a/src/gt4py/next/ffront/fbuiltins.py b/src/gt4py/next/ffront/fbuiltins.py index e523e94eaa..706b6a4606 100644 --- a/src/gt4py/next/ffront/fbuiltins.py +++ b/src/gt4py/next/ffront/fbuiltins.py @@ -14,7 +14,6 @@ import dataclasses import inspect -import math from builtins import bool, float, int, tuple from typing import ( Any, @@ -272,13 +271,7 @@ def impl(value: common.Field | core_defs.ScalarT, /) -> common.Field | core_defs ): _make_unary_math_builtin(f) -BINARY_MATH_NUMBER_BUILTIN_TO_PYTHON_SCALAR_FUNCTION = { - "minimum": min, - "maximum": max, - "fmod": math.fmod, - "power": pow, -} -BINARY_MATH_NUMBER_BUILTIN_NAMES = list(BINARY_MATH_NUMBER_BUILTIN_TO_PYTHON_SCALAR_FUNCTION.keys()) +BINARY_MATH_NUMBER_BUILTIN_NAMES = ["minimum", "maximum", "fmod", "power"] def _make_binary_math_builtin(name): From 5f52a6c65720b8e1cb4604534e81a45cae751e45 Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Mon, 20 Nov 2023 17:49:16 +0000 Subject: [PATCH 20/29] fix column stencil test --- src/gt4py/next/utils.py | 15 +++++++++++++++ tests/next_tests/integration_tests/cases.py | 13 +++---------- .../iterator_tests/test_column_stencil.py | 13 +++++++------ 3 files changed, 25 insertions(+), 16 deletions(-) diff --git a/src/gt4py/next/utils.py b/src/gt4py/next/utils.py index 006b3057b0..ddbe93d10b 100644 --- a/src/gt4py/next/utils.py +++ b/src/gt4py/next/utils.py @@ -12,8 +12,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +import functools from typing import Any, ClassVar, TypeGuard, TypeVar +from gt4py.next import common + class RecursionGuard: """ @@ -56,3 +59,15 @@ def __exit__(self, *exc): def is_tuple_of(v: Any, t: type[_T]) -> TypeGuard[tuple[_T, ...]]: return isinstance(v, tuple) and all(isinstance(e, t) for e in v) + + +def apply_to_tuple_elems(fun, *args): # TODO type annotations + if isinstance(args[0], tuple): + assert all(isinstance(arg, tuple) for arg in args) + return tuple(apply_to_tuple_elems(fun, *arg) for arg in zip(*args)) + return fun(*args) + + +asnumpy = functools.partial( + apply_to_tuple_elems, lambda f: f.asnumpy() if common.is_field(f) else f +) diff --git a/tests/next_tests/integration_tests/cases.py b/tests/next_tests/integration_tests/cases.py index a157cecdc1..7ef724ee2f 100644 --- a/tests/next_tests/integration_tests/cases.py +++ b/tests/next_tests/integration_tests/cases.py @@ -28,7 +28,7 @@ from gt4py._core import definitions as core_defs from gt4py.eve import extended_typing as xtyping from gt4py.eve.extended_typing import Self -from gt4py.next import common, constructors +from gt4py.next import common, constructors, utils from gt4py.next.ffront import decorator from gt4py.next.program_processors import processor_interface as ppi from gt4py.next.type_system import type_specifications as ts, type_translation @@ -385,13 +385,6 @@ def run( fieldview_prog.with_grid_type(case.grid_type).with_backend(case.backend)(*args, **kwargs) -def _asndarray(fields: tuple[common.Field | np.ndarray | tuple]) -> tuple[np.ndarray | tuple]: - if isinstance(fields, tuple): - return tuple(_asndarray(f) for f in fields) - else: - return fields.asnumpy() if hasattr(fields, "asnumpy") else fields - - def verify( case: Case, fieldview_prog: decorator.FieldOperator | decorator.Program, @@ -443,8 +436,8 @@ def verify( out_comp = out or inout assert out_comp is not None - out_comp_ndarray = _asndarray(out_comp) - ref_ndarray = _asndarray(ref) + out_comp_ndarray = utils.asnumpy(out_comp) + ref_ndarray = utils.asnumpy(ref) assert comparison(ref_ndarray, out_comp_ndarray), ( f"Verification failed:\n" f"\tcomparison={comparison.__name__}(ref, out)\n" diff --git a/tests/next_tests/integration_tests/multi_feature_tests/iterator_tests/test_column_stencil.py b/tests/next_tests/integration_tests/multi_feature_tests/iterator_tests/test_column_stencil.py index 51bd039ab3..fd571514ac 100644 --- a/tests/next_tests/integration_tests/multi_feature_tests/iterator_tests/test_column_stencil.py +++ b/tests/next_tests/integration_tests/multi_feature_tests/iterator_tests/test_column_stencil.py @@ -16,6 +16,7 @@ import pytest import gt4py.next as gtx +from gt4py.next import utils from gt4py.next.iterator.builtins import * from gt4py.next.iterator.runtime import closure, fendef, fundef, offset @@ -89,7 +90,7 @@ def test_basic_column_stencils(program_processor, lift_mode, basic_stencils): ) out = gtx.as_field([IDim, KDim], np.zeros(shape)) - ref = ref_fun(inp) + ref = ref_fun(inp.asnumpy()) run_processor( stencil[{IDim: range(0, shape[0]), KDim: range(0, shape[1])}], @@ -157,7 +158,7 @@ def test_k_level_condition(program_processor, lift_mode, fun, k_level, inp_funct k_size = 5 inp = inp_function(k_size) - ref = ref_function(inp) + ref = ref_function(utils.asnumpy(inp)) out = gtx.as_field([KDim], np.zeros((5,), dtype=np.int32)) @@ -173,7 +174,7 @@ def test_k_level_condition(program_processor, lift_mode, fun, k_level, inp_funct ) if validate: - np.allclose(ref, out) + np.allclose(ref, out.asnumpy()) @fundef @@ -222,7 +223,7 @@ def test_ksum_scan(program_processor, lift_mode, kstart, reference): ) if validate: - assert np.allclose(reference, np.asarray(out)) + assert np.allclose(reference, out.asnumpy()) @fundef @@ -260,7 +261,7 @@ def test_ksum_back_scan(program_processor, lift_mode): ) if validate: - assert np.allclose(ref, np.asarray(out)) + assert np.allclose(ref, out.asnumpy()) @fundef @@ -366,7 +367,7 @@ def test_different_vertical_sizes(program_processor): ) if validate: - assert np.allclose(ref[1:], out[1:]) + assert np.allclose(ref[1:], out.asnumpy()[1:]) @fundef From d2e840bc1c69ab0b20a1aec06fa823e87e16668f Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Mon, 20 Nov 2023 17:49:58 +0000 Subject: [PATCH 21/29] fix ndarray test --- .../unit_tests/embedded_tests/test_nd_array_field.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/next_tests/unit_tests/embedded_tests/test_nd_array_field.py b/tests/next_tests/unit_tests/embedded_tests/test_nd_array_field.py index 00dbf68274..436e672cc5 100644 --- a/tests/next_tests/unit_tests/embedded_tests/test_nd_array_field.py +++ b/tests/next_tests/unit_tests/embedded_tests/test_nd_array_field.py @@ -555,8 +555,8 @@ def test_setitem(index, value): domain=common.Domain(dims=(IDim, JDim), ranges=(UnitRange(0, 10), UnitRange(0, 10))), ) - expected = np.copy(field.ndarray) - expected[index] = value + expected = np.copy(field.asnumpy()) + expected[index] = value.asnumpy() if common.is_field(value) else value field[index] = value From c394ed0cb03a8024c6513c517bf167e43ed8e3df Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Mon, 20 Nov 2023 17:55:09 +0000 Subject: [PATCH 22/29] fix 2 more --- .../feature_tests/ffront_tests/test_gt4py_builtins.py | 2 +- .../iterator_tests/test_strided_offset_provider.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_gt4py_builtins.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_gt4py_builtins.py index d60e396817..e2434d860a 100644 --- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_gt4py_builtins.py +++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_gt4py_builtins.py @@ -286,7 +286,7 @@ def conditional_program( b, out, inout=out, - ref=np.where(mask, a, b)[1:], + ref=np.where(mask.asnumpy(), a.asnumpy(), b.asnumpy())[1:], ) diff --git a/tests/next_tests/integration_tests/feature_tests/iterator_tests/test_strided_offset_provider.py b/tests/next_tests/integration_tests/feature_tests/iterator_tests/test_strided_offset_provider.py index abdfffd74e..dd603fa3be 100644 --- a/tests/next_tests/integration_tests/feature_tests/iterator_tests/test_strided_offset_provider.py +++ b/tests/next_tests/integration_tests/feature_tests/iterator_tests/test_strided_offset_provider.py @@ -63,9 +63,9 @@ def test_strided_offset_provider(program_processor): ), ) out = gtx.as_field([LocA], np.zeros((LocA_size,))) - ref = np.sum(np.asarray(inp).reshape(LocA_size, max_neighbors), axis=-1) + ref = np.sum(inp.asnumpy().reshape(LocA_size, max_neighbors), axis=-1) run_processor(fencil, program_processor, LocA_size, out, inp) if validate: - assert np.allclose(out, ref) + assert np.allclose(out.asnumpy(), ref) From 894e41d4e163bdc92cc16074bf604a435c3f7821 Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Mon, 20 Nov 2023 17:58:29 +0000 Subject: [PATCH 23/29] astype tuple --- .../feature_tests/ffront_tests/test_execution.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_execution.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_execution.py index 52687edbb2..1f3b54d6f0 100644 --- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_execution.py +++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_execution.py @@ -371,8 +371,8 @@ def cast_nested_tuple( a = cases.allocate(cartesian_case, cast_tuple, "a")() b = cases.allocate(cartesian_case, cast_tuple, "b")() - a_asint = cartesian_case.as_field([IDim], np.asarray(a).astype(int32)) - b_asint = cartesian_case.as_field([IDim], np.asarray(b).astype(int32)) + a_asint = cartesian_case.as_field([IDim], a.asnumpy().astype(int32)) + b_asint = cartesian_case.as_field([IDim], b.asnumpy().astype(int32)) out_tuple = cases.allocate(cartesian_case, cast_tuple, cases.RETURN)() out_nested_tuple = cases.allocate(cartesian_case, cast_nested_tuple, cases.RETURN)() @@ -384,7 +384,10 @@ def cast_nested_tuple( a_asint, b_asint, out=out_tuple, - ref=(np.full_like(a, True, dtype=bool), np.full_like(b, True, dtype=bool)), + ref=( + np.full_like(a.asnumpy(), True, dtype=bool), + np.full_like(b.asnumpy(), True, dtype=bool), + ), ) cases.verify( @@ -396,9 +399,9 @@ def cast_nested_tuple( b_asint, out=out_nested_tuple, ref=( - np.full_like(a, True, dtype=bool), - np.full_like(a, True, dtype=bool), - np.full_like(b, True, dtype=bool), + np.full_like(a.asnumpy(), True, dtype=bool), + np.full_like(a.asnumpy(), True, dtype=bool), + np.full_like(b.asnumpy(), True, dtype=bool), ), ) From 89c429d6871a5d386db16776172e9d7a489043d5 Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Mon, 20 Nov 2023 18:57:33 +0000 Subject: [PATCH 24/29] 2 more --- .../feature_tests/iterator_tests/test_conditional.py | 4 ++-- .../feature_tests/iterator_tests/test_constant.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/next_tests/integration_tests/feature_tests/iterator_tests/test_conditional.py b/tests/next_tests/integration_tests/feature_tests/iterator_tests/test_conditional.py index de7ebf2869..8536dbea90 100644 --- a/tests/next_tests/integration_tests/feature_tests/iterator_tests/test_conditional.py +++ b/tests/next_tests/integration_tests/feature_tests/iterator_tests/test_conditional.py @@ -51,5 +51,5 @@ def test_conditional_w_tuple(program_processor): offset_provider={}, ) if validate: - assert np.all(out.ndarray[np.asarray(inp) == 0] == 3.0) - assert np.all(out.ndarray[np.asarray(inp) == 1] == 7.0) + assert np.all(out.asnumpy()[inp.asnumpy() == 0] == 3.0) + assert np.all(out.asnumpy()[inp.asnumpy() == 1] == 7.0) diff --git a/tests/next_tests/integration_tests/feature_tests/iterator_tests/test_constant.py b/tests/next_tests/integration_tests/feature_tests/iterator_tests/test_constant.py index 83a86319b4..faae549086 100644 --- a/tests/next_tests/integration_tests/feature_tests/iterator_tests/test_constant.py +++ b/tests/next_tests/integration_tests/feature_tests/iterator_tests/test_constant.py @@ -31,8 +31,8 @@ def constant_stencil(): # this is traced as a lambda, TODO directly feed iterat return deref(inp) + deref(lift(constant_stencil)()) inp = gtx.as_field([IDim], np.asarray([0, 42], dtype=np.int32)) - res = gtx.as_field([IDim], np.zeros_like(inp)) + res = gtx.as_field([IDim], np.zeros_like(inp.asnumpy())) add_constant[{IDim: range(2)}](inp, out=res, offset_provider={}, backend=roundtrip.executor) - assert np.allclose(res, np.asarray([1, 43])) + assert np.allclose(res.asnumpy(), np.asarray([1, 43])) From 27999c18b302540e2d7652eff35bd298eb4929a6 Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Mon, 20 Nov 2023 19:07:31 +0000 Subject: [PATCH 25/29] fix fvm nabla --- src/gt4py/next/utils.py | 2 +- .../iterator_tests/test_fvm_nabla.py | 40 +++++++++---------- 2 files changed, 21 insertions(+), 21 deletions(-) diff --git a/src/gt4py/next/utils.py b/src/gt4py/next/utils.py index ddbe93d10b..9e964e5b98 100644 --- a/src/gt4py/next/utils.py +++ b/src/gt4py/next/utils.py @@ -63,7 +63,7 @@ def is_tuple_of(v: Any, t: type[_T]) -> TypeGuard[tuple[_T, ...]]: def apply_to_tuple_elems(fun, *args): # TODO type annotations if isinstance(args[0], tuple): - assert all(isinstance(arg, tuple) for arg in args) + assert all(isinstance(arg, tuple) and len(args[0]) == len(arg) for arg in args) return tuple(apply_to_tuple_elems(fun, *arg) for arg in zip(*args)) return fun(*args) diff --git a/tests/next_tests/integration_tests/multi_feature_tests/iterator_tests/test_fvm_nabla.py b/tests/next_tests/integration_tests/multi_feature_tests/iterator_tests/test_fvm_nabla.py index 47867b9a64..e1d959aba9 100644 --- a/tests/next_tests/integration_tests/multi_feature_tests/iterator_tests/test_fvm_nabla.py +++ b/tests/next_tests/integration_tests/multi_feature_tests/iterator_tests/test_fvm_nabla.py @@ -159,8 +159,8 @@ def test_compute_zavgS(program_processor, lift_mode): ) if validate: - assert_close(-199755464.25741270, np.min(zavgS)) - assert_close(388241977.58389181, np.max(zavgS)) + assert_close(-199755464.25741270, np.min(zavgS.asnumpy())) + assert_close(388241977.58389181, np.max(zavgS.asnumpy())) run_processor( compute_zavgS_fencil, @@ -173,8 +173,8 @@ def test_compute_zavgS(program_processor, lift_mode): lift_mode=lift_mode, ) if validate: - assert_close(-1000788897.3202186, np.min(zavgS)) - assert_close(1000788897.3202186, np.max(zavgS)) + assert_close(-1000788897.3202186, np.min(zavgS.asnumpy())) + assert_close(1000788897.3202186, np.max(zavgS.asnumpy())) @fendef @@ -222,11 +222,11 @@ def test_compute_zavgS2(program_processor, lift_mode): ) if validate: - assert_close(-199755464.25741270, np.min(zavgS[0])) - assert_close(388241977.58389181, np.max(zavgS[0])) + assert_close(-199755464.25741270, np.min(zavgS[0].asnumpy())) + assert_close(388241977.58389181, np.max(zavgS[0].asnumpy())) - assert_close(-1000788897.3202186, np.min(zavgS[1])) - assert_close(1000788897.3202186, np.max(zavgS[1])) + assert_close(-1000788897.3202186, np.min(zavgS[1].asnumpy())) + assert_close(1000788897.3202186, np.max(zavgS[1].asnumpy())) @pytest.mark.requires_atlas @@ -266,10 +266,10 @@ def test_nabla(program_processor, lift_mode): ) if validate: - assert_close(-3.5455427772566003e-003, np.min(pnabla_MXX)) - assert_close(3.5455427772565435e-003, np.max(pnabla_MXX)) - assert_close(-3.3540113705465301e-003, np.min(pnabla_MYY)) - assert_close(3.3540113705465301e-003, np.max(pnabla_MYY)) + assert_close(-3.5455427772566003e-003, np.min(pnabla_MXX.asnumpy())) + assert_close(3.5455427772565435e-003, np.max(pnabla_MXX.asnumpy())) + assert_close(-3.3540113705465301e-003, np.min(pnabla_MYY.asnumpy())) + assert_close(3.3540113705465301e-003, np.max(pnabla_MYY.asnumpy())) @fendef @@ -322,10 +322,10 @@ def test_nabla2(program_processor, lift_mode): ) if validate: - assert_close(-3.5455427772566003e-003, np.min(pnabla_MXX)) - assert_close(3.5455427772565435e-003, np.max(pnabla_MXX)) - assert_close(-3.3540113705465301e-003, np.min(pnabla_MYY)) - assert_close(3.3540113705465301e-003, np.max(pnabla_MYY)) + assert_close(-3.5455427772566003e-003, np.min(pnabla_MXX.asnumpy())) + assert_close(3.5455427772565435e-003, np.max(pnabla_MXX.asnumpy())) + assert_close(-3.3540113705465301e-003, np.min(pnabla_MYY.asnumpy())) + assert_close(3.3540113705465301e-003, np.max(pnabla_MYY.asnumpy())) @fundef @@ -407,7 +407,7 @@ def test_nabla_sign(program_processor, lift_mode): ) if validate: - assert_close(-3.5455427772566003e-003, np.min(pnabla_MXX)) - assert_close(3.5455427772565435e-003, np.max(pnabla_MXX)) - assert_close(-3.3540113705465301e-003, np.min(pnabla_MYY)) - assert_close(3.3540113705465301e-003, np.max(pnabla_MYY)) + assert_close(-3.5455427772566003e-003, np.min(pnabla_MXX.asnumpy())) + assert_close(3.5455427772565435e-003, np.max(pnabla_MXX.asnumpy())) + assert_close(-3.3540113705465301e-003, np.min(pnabla_MYY.asnumpy())) + assert_close(3.3540113705465301e-003, np.max(pnabla_MYY.asnumpy())) From fb76d11bb285fa83d9b0d4a58208dc5ffe0cdc7d Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Mon, 20 Nov 2023 19:56:22 +0000 Subject: [PATCH 26/29] refactor asnumpy --- src/gt4py/next/utils.py | 39 ++++++++++++++++++++++++++++++--------- 1 file changed, 30 insertions(+), 9 deletions(-) diff --git a/src/gt4py/next/utils.py b/src/gt4py/next/utils.py index 9e964e5b98..76a42e8448 100644 --- a/src/gt4py/next/utils.py +++ b/src/gt4py/next/utils.py @@ -13,7 +13,9 @@ # SPDX-License-Identifier: GPL-3.0-or-later import functools -from typing import Any, ClassVar, TypeGuard, TypeVar +from typing import Any, Callable, ClassVar, ParamSpec, TypeGuard, TypeVar, cast + +import numpy as np from gt4py.next import common @@ -56,18 +58,37 @@ def __exit__(self, *exc): _T = TypeVar("_T") +_P = ParamSpec("_P") +_R = TypeVar("_R") + def is_tuple_of(v: Any, t: type[_T]) -> TypeGuard[tuple[_T, ...]]: return isinstance(v, tuple) and all(isinstance(e, t) for e in v) -def apply_to_tuple_elems(fun, *args): # TODO type annotations - if isinstance(args[0], tuple): - assert all(isinstance(arg, tuple) and len(args[0]) == len(arg) for arg in args) - return tuple(apply_to_tuple_elems(fun, *arg) for arg in zip(*args)) - return fun(*args) +def apply_to_tuple_elements(fun: Callable[_P, _R]) -> Callable[..., _R | tuple[_R | tuple, ...]]: + """Apply `fun` to each entry of (possibly nested) tuples. + + Examples: + >>> apply_to_tuple_elements(lambda x: x + 1)(((1, 2), 3)) + ((2, 3), 4) + + >>> apply_to_tuple_elements(lambda x, y: x + y)(((1, 2), 3), ((4, 5), 6)) + ((5, 7), 9) + """ + + @functools.wraps(fun) + def impl(*args: Any | tuple[Any | tuple, ...]) -> _R | tuple[_R | tuple, ...]: + if isinstance(args[0], tuple): + assert all(isinstance(arg, tuple) and len(args[0]) == len(arg) for arg in args) + return tuple(impl(*arg) for arg in zip(*args)) + + return fun(*cast(_P.args, args)) + + return impl -asnumpy = functools.partial( - apply_to_tuple_elems, lambda f: f.asnumpy() if common.is_field(f) else f -) +# TODO(havogt): consider moving to module like `field_utils` +@apply_to_tuple_elements +def asnumpy(field: common.Field | np.ndarray) -> np.ndarray: + return field.asnumpy() if common.is_field(field) else field # type: ignore[return-value] # mypy doesn't understand the condition From 753417e689d5e58f9d9068402bc033d767c46f2f Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Tue, 21 Nov 2023 10:29:34 +0000 Subject: [PATCH 27/29] address review comments --- src/gt4py/next/common.py | 7 ++++--- src/gt4py/next/iterator/embedded.py | 6 ++++++ src/gt4py/next/utils.py | 8 ++++---- 3 files changed, 14 insertions(+), 7 deletions(-) diff --git a/src/gt4py/next/common.py b/src/gt4py/next/common.py index 5fb50e1e59..51ad14f22d 100644 --- a/src/gt4py/next/common.py +++ b/src/gt4py/next/common.py @@ -464,12 +464,13 @@ def dtype(self) -> core_defs.DType[core_defs.ScalarT]: def ndarray(self) -> core_defs.NDArrayObject: ... - def asnumpy(self) -> np.ndarray: - return np.asarray(self.ndarray) - def __str__(self) -> str: return f"⟨{self.domain!s} → {self.dtype}⟩" + @abc.abstractmethod + def asnumpy(self) -> np.ndarray: + ... + @abc.abstractmethod def remap(self, index_field: Field) -> Field: ... diff --git a/src/gt4py/next/iterator/embedded.py b/src/gt4py/next/iterator/embedded.py index 44294a3a71..9000b00d8f 100644 --- a/src/gt4py/next/iterator/embedded.py +++ b/src/gt4py/next/iterator/embedded.py @@ -1068,6 +1068,9 @@ def dtype(self) -> core_defs.Int32DType: def ndarray(self) -> core_defs.NDArrayObject: raise AttributeError("Cannot get `ndarray` of an infinite Field.") + def asnumpy(self) -> np.ndarray: + raise NotImplementedError() + def remap(self, index_field: common.Field) -> common.Field: # TODO can be implemented by constructing and ndarray (but do we know of which kind?) raise NotImplementedError() @@ -1180,6 +1183,9 @@ def dtype(self) -> core_defs.DType[core_defs.ScalarT]: def ndarray(self) -> core_defs.NDArrayObject: raise AttributeError("Cannot get `ndarray` of an infinite Field.") + def asnumpy(self) -> np.ndarray: + raise NotImplementedError() + def remap(self, index_field: common.Field) -> common.Field: # TODO can be implemented by constructing and ndarray (but do we know of which kind?) raise NotImplementedError() diff --git a/src/gt4py/next/utils.py b/src/gt4py/next/utils.py index 76a42e8448..f64d1e198a 100644 --- a/src/gt4py/next/utils.py +++ b/src/gt4py/next/utils.py @@ -66,14 +66,14 @@ def is_tuple_of(v: Any, t: type[_T]) -> TypeGuard[tuple[_T, ...]]: return isinstance(v, tuple) and all(isinstance(e, t) for e in v) -def apply_to_tuple_elements(fun: Callable[_P, _R]) -> Callable[..., _R | tuple[_R | tuple, ...]]: +def tree_map(fun: Callable[_P, _R]) -> Callable[..., _R | tuple[_R | tuple, ...]]: """Apply `fun` to each entry of (possibly nested) tuples. Examples: - >>> apply_to_tuple_elements(lambda x: x + 1)(((1, 2), 3)) + >>> tree_map(lambda x: x + 1)(((1, 2), 3)) ((2, 3), 4) - >>> apply_to_tuple_elements(lambda x, y: x + y)(((1, 2), 3), ((4, 5), 6)) + >>> tree_map(lambda x, y: x + y)(((1, 2), 3), ((4, 5), 6)) ((5, 7), 9) """ @@ -89,6 +89,6 @@ def impl(*args: Any | tuple[Any | tuple, ...]) -> _R | tuple[_R | tuple, ...]: # TODO(havogt): consider moving to module like `field_utils` -@apply_to_tuple_elements +@tree_map def asnumpy(field: common.Field | np.ndarray) -> np.ndarray: return field.asnumpy() if common.is_field(field) else field # type: ignore[return-value] # mypy doesn't understand the condition From 117e2c9af664a3ac5fd873354119ebbecbbaeed3 Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Tue, 21 Nov 2023 11:30:26 +0100 Subject: [PATCH 28/29] Update src/gt4py/next/utils.py --- src/gt4py/next/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/gt4py/next/utils.py b/src/gt4py/next/utils.py index f64d1e198a..703e592840 100644 --- a/src/gt4py/next/utils.py +++ b/src/gt4py/next/utils.py @@ -83,7 +83,7 @@ def impl(*args: Any | tuple[Any | tuple, ...]) -> _R | tuple[_R | tuple, ...]: assert all(isinstance(arg, tuple) and len(args[0]) == len(arg) for arg in args) return tuple(impl(*arg) for arg in zip(*args)) - return fun(*cast(_P.args, args)) + return fun(*cast(_P.args, args)) # mypy doesn't understand that `args` at this point is of type `_P.args` return impl From 9671c54b4baef0b20ddb0e1de934f9be65c4c269 Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Tue, 21 Nov 2023 10:31:53 +0000 Subject: [PATCH 29/29] fix formatting --- src/gt4py/next/utils.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/gt4py/next/utils.py b/src/gt4py/next/utils.py index 703e592840..baae8361c5 100644 --- a/src/gt4py/next/utils.py +++ b/src/gt4py/next/utils.py @@ -83,7 +83,9 @@ def impl(*args: Any | tuple[Any | tuple, ...]) -> _R | tuple[_R | tuple, ...]: assert all(isinstance(arg, tuple) and len(args[0]) == len(arg) for arg in args) return tuple(impl(*arg) for arg in zip(*args)) - return fun(*cast(_P.args, args)) # mypy doesn't understand that `args` at this point is of type `_P.args` + return fun( + *cast(_P.args, args) + ) # mypy doesn't understand that `args` at this point is of type `_P.args` return impl