diff --git a/src/gt4py/next/embedded/nd_array_field.py b/src/gt4py/next/embedded/nd_array_field.py index 1760cb17e8..f7adbed89b 100644 --- a/src/gt4py/next/embedded/nd_array_field.py +++ b/src/gt4py/next/embedded/nd_array_field.py @@ -211,8 +211,11 @@ def remap( # then compute the index array xp = self.array_ns new_idx_array = xp.asarray(restricted_connectivity.ndarray) - current_range.start - # finally, take the new array - new_buffer = xp.take(self._ndarray, new_idx_array, axis=dim_idx) + if self._ndarray.ndim > 1 and restricted_connectivity_domain == new_domain: + new_buffer = self._take_mdim(new_idx_array, axis=dim_idx) + else: + # finally, take the new array + new_buffer = xp.take(self._ndarray, new_idx_array, axis=dim_idx) return self.__class__.from_array( new_buffer, @@ -324,6 +327,24 @@ def _slice( assert common.is_relative_index_sequence(slice_) return new_domain, slice_ + def _take_mdim( + self, + restricted_connectivity: core_defs.NDArrayObject, + axis: int, + ) -> core_defs.NDArrayObject: + xp = self.array_ns + dim = self.domain.dims[axis] + offset_abs = [ + restricted_connectivity if d == dim else xp.indices(restricted_connectivity.shape)[d_i] + for d_i, d in enumerate(self.domain.dims) + ] + new_buffer_flat = xp.take( + xp.asarray(self._ndarray).flatten(), + xp.ravel_multi_index(tuple(offset_abs), self._ndarray.shape).flatten(), + ) + new_buffer = new_buffer_flat.reshape(restricted_connectivity.shape) + return new_buffer + @dataclasses.dataclass(frozen=True) class NdArrayConnectivityField( # type: ignore[misc] # for __ne__, __eq__ diff --git a/src/gt4py/next/ffront/experimental.py b/src/gt4py/next/ffront/experimental.py index 39da80a5de..f53a75973a 100644 --- a/src/gt4py/next/ffront/experimental.py +++ b/src/gt4py/next/ffront/experimental.py @@ -12,30 +12,17 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -from dataclasses import dataclass +from gt4py.next import common +from gt4py.next.ffront.fbuiltins import BuiltInFunction, FieldOffset -from gt4py.next.type_system import type_specifications as ts +@BuiltInFunction +def as_offset( + offset_: FieldOffset, + field: common.Field, + /, +) -> common.ConnectivityField: + raise NotImplementedError() -@dataclass -class BuiltInFunction: - __gt_type: ts.FunctionType - def __call__(self, *args, **kwargs): - """Act as an empty place holder for the built in function.""" - - def __gt_type__(self): - return self.__gt_type - - -as_offset = BuiltInFunction( - ts.FunctionType( - pos_only_args=[ - ts.DeferredType(constraint=ts.OffsetType), - ts.DeferredType(constraint=ts.FieldType), - ], - pos_or_kw_args={}, - kw_only_args={}, - returns=ts.DeferredType(constraint=ts.OffsetType), - ) -) +EXPERIMENTAL_FUN_BUILTIN_NAMES = ["as_offset"] diff --git a/src/gt4py/next/ffront/foast_passes/type_deduction.py b/src/gt4py/next/ffront/foast_passes/type_deduction.py index 9bc004e7b8..0e7c430596 100644 --- a/src/gt4py/next/ffront/foast_passes/type_deduction.py +++ b/src/gt4py/next/ffront/foast_passes/type_deduction.py @@ -20,6 +20,7 @@ from gt4py.next.common import DimensionKind from gt4py.next.ffront import ( # noqa dialect_ast_enums, + experimental, fbuiltins, type_info as ti_ffront, type_specifications as ts_ffront, @@ -727,7 +728,8 @@ def visit_Call(self, node: foast.Call, **kwargs) -> foast.Call: isinstance(new_func.type, ts.FunctionType) and not type_info.is_concrete(return_type) and isinstance(new_func, foast.Name) - and new_func.id in fbuiltins.FUN_BUILTIN_NAMES + and new_func.id + in (fbuiltins.FUN_BUILTIN_NAMES + experimental.EXPERIMENTAL_FUN_BUILTIN_NAMES) ): visitor = getattr(self, f"_visit_{new_func.id}") return visitor(new_node, **kwargs) diff --git a/tests/next_tests/definitions.py b/tests/next_tests/definitions.py index a8ba72f366..c6927a01d4 100644 --- a/tests/next_tests/definitions.py +++ b/tests/next_tests/definitions.py @@ -163,7 +163,6 @@ class ProgramFormatterId(_PythonObjectIdMixin, str, enum.Enum): (USES_ZERO_DIMENSIONAL_FIELDS, XFAIL, UNSUPPORTED_MESSAGE), ] EMBEDDED_SKIP_LIST = [ - (USES_DYNAMIC_OFFSETS, XFAIL, UNSUPPORTED_MESSAGE), (CHECKS_SPECIFIC_ERROR, XFAIL, UNSUPPORTED_MESSAGE), ( USES_SCAN_WITHOUT_FIELD_ARGS, diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_as_offset.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_as_offset.py new file mode 100644 index 0000000000..401cabf6ce --- /dev/null +++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_as_offset.py @@ -0,0 +1,210 @@ +# -*- coding: utf-8 -*- +# GT4Py - GridTools Framework +# +# Copyright (c) 2014-2023, ETH Zurich +# All rights reserved. +# +# This file is part of the GT4Py project and the GridTools framework. +# GT4Py is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np +import pytest + +import gt4py.next as gtx +from gt4py.next.ffront.experimental import as_offset + +from next_tests.integration_tests import cases +from next_tests.integration_tests.cases import IDim, Ioff, JDim, KDim, Koff, cartesian_case +from next_tests.integration_tests.feature_tests.ffront_tests.ffront_test_utils import ( + fieldview_backend, + reduction_setup, +) + + +@pytest.mark.uses_dynamic_offsets +def test_offset_field_regular_input_size(cartesian_case): + @gtx.field_operator + def testee( + a: gtx.Field[[IDim, KDim], int], offset_field: gtx.Field[[IDim, KDim], int] + ) -> gtx.Field[[IDim, KDim], int]: + return a(as_offset(Ioff, offset_field)) + + out = cases.allocate(cartesian_case, testee, cases.RETURN)() + a = cases.allocate(cartesian_case, testee, "a")() + offset_field = cases.allocate(cartesian_case, testee, "offset_field").strategy( + cases.ConstInitializer(3) + )() + ref = a[3:] + + cases.verify( + cartesian_case, + testee, + a, + offset_field, + out=out[: len(out.asnumpy()) - 3], + offset_provider={"Ioff": IDim}, + ref=ref, + comparison=lambda out, ref: np.all(out == ref), + ) + + +@pytest.mark.uses_dynamic_offsets +def test_offset_field_domain(cartesian_case): + @gtx.field_operator + def testee_fo( + a: gtx.Field[[IDim, KDim], int], offset_field: gtx.Field[[IDim, KDim], int] + ) -> gtx.Field[[IDim, KDim], int]: + return a(as_offset(Ioff, offset_field)) + + @gtx.program + def testee( + a: gtx.Field[[IDim, KDim], int], + offset_field: gtx.Field[[IDim, KDim], int], + out: gtx.Field[[IDim, KDim], int], + ): + testee_fo(a, offset_field, out=out, domain={IDim: (1, 5), KDim: (0, 10)}) + + out = cases.allocate(cartesian_case, testee, "out")() + a = cases.allocate(cartesian_case, testee, "a").extend({IDim: (0, 3)})() + offset_field = cases.allocate(cartesian_case, testee, "offset_field").strategy( + cases.ConstInitializer(3) + )() + ref = out.asnumpy().copy() + a_shift = a[3:] + ref[1:5] = a_shift.asnumpy()[1:5] + + cases.verify( + cartesian_case, + testee, + a, + offset_field, + out, + inout=out, + offset_provider={"Ioff": IDim}, + ref=ref, + comparison=lambda out, ref: np.all(out == ref), + ) + + +@pytest.mark.uses_dynamic_offsets +def test_offset_field_input_domain(cartesian_case): + i_size = cartesian_case.default_sizes[IDim] + k_size = cartesian_case.default_sizes[KDim] + + @gtx.field_operator + def testee_fo( + a: gtx.Field[[IDim, KDim], int], offset_field: gtx.Field[[IDim, KDim], int] + ) -> gtx.Field[[IDim, KDim], int]: + return a(as_offset(Ioff, offset_field)) + + @gtx.program + def testee( + a: gtx.Field[[IDim, KDim], int], + offset_field: gtx.Field[[IDim, KDim], int], + out: gtx.Field[[IDim, KDim], int], + ): + testee_fo(a, offset_field, out=out) + + a_fo = cases.allocate(cartesian_case, testee, "a").extend({IDim: (0, 3)})().ndarray + out = cartesian_case.as_field( + [IDim, KDim], + np.zeros([i_size, k_size], dtype=int), + origin={IDim: 2, KDim: 0}, + ) + a = cartesian_case.as_field([IDim, KDim], a_fo, origin={IDim: 2, KDim: 0}) + offset_field = cases.allocate(cartesian_case, testee, "offset_field").strategy( + cases.ConstInitializer(3) + )() + ref = a[5:] + cases.verify( + cartesian_case, + testee, + a, + offset_field, + out[2:], + inout=out[2:], + offset_provider={"Ioff": IDim}, + ref=ref, + ) + + +@pytest.mark.parametrize( + "dims, offset, offset_dim, off_dim, other_dim", + [ + ([IDim, KDim], Ioff, IDim, 0, 1), + ([IDim, JDim], Ioff, IDim, 0, 1), + ([IDim, KDim], Koff, KDim, 1, 0), + ], +) +@pytest.mark.uses_dynamic_offsets +def test_offset_field_rand_connectivity( + cartesian_case, dims, offset, offset_dim, off_dim, other_dim +): + @gtx.field_operator + def testee( + a: gtx.Field[[dims[0], dims[1]], int], offset_field: gtx.Field[[dims[0], dims[1]], int] + ) -> gtx.Field[[dims[0], dims[1]], int]: + return a(as_offset(offset, offset_field)) + + out = cases.allocate(cartesian_case, testee, cases.RETURN)() + a = cases.allocate(cartesian_case, testee, "a").extend({offset_dim: (0, 3)})() + offset_field = cartesian_case.as_field( + [dims[0], dims[1]], + np.random.randint( + 0, + 4, + size=(cartesian_case.default_sizes[dims[0]], cartesian_case.default_sizes[dims[1]]), + ), + ) + ref = np.zeros(offset_field.asnumpy().shape, dtype=int) + + for j in range(len(offset_field.asnumpy()[off_dim])): + for i in range(len(offset_field.asnumpy())): + if offset_dim.kind == "vertical": + y = offset_field.asnumpy()[i][j] + np.indices(offset_field.shape)[off_dim][i][j] + x = np.indices(offset_field.shape)[other_dim][i][j] + else: + x = offset_field.asnumpy()[i][j] + np.indices(offset_field.shape)[off_dim][i][j] + y = np.indices(offset_field.shape)[other_dim][i][j] + ref[i][j] = a[x][y] + + cases.verify( + cartesian_case, + testee, + a, + offset_field, + out=out, + offset_provider={"offset": offset_dim, "Ioff": IDim, "Koff": KDim}, + ref=ref, + comparison=lambda out, ref: np.all(out == ref), + ) + + +@pytest.mark.uses_dynamic_offsets +def test_offset_field_3d(cartesian_case): + @gtx.field_operator + def testee(a: cases.IJKField, offset_field: cases.IJKField) -> cases.IJKField: + return a(as_offset(Ioff, offset_field)) + + out = cases.allocate(cartesian_case, testee, cases.RETURN)() + a = cases.allocate(cartesian_case, testee, "a").extend({IDim: (0, 2)})() + offset_field = cases.allocate(cartesian_case, testee, "offset_field").strategy( + cases.ConstInitializer(2) + )() + + cases.verify( + cartesian_case, + testee, + a, + offset_field, + out=out, + offset_provider={"Ioff": IDim, "Koff": KDim}, + ref=a[2:], + comparison=lambda out, ref: np.all(out == ref), + ) 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 5db9886966..7af97c5c3f 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 @@ -32,7 +32,6 @@ neighbor_sum, where, ) -from gt4py.next.ffront.experimental import as_offset from gt4py.next.program_processors.runners import gtfn from next_tests.integration_tests import cases @@ -462,40 +461,6 @@ def testee(a: cases.IFloatField) -> gtx.Field[[IDim], np.float32]: ) -@pytest.mark.uses_dynamic_offsets -def test_offset_field(cartesian_case): - ref = np.full( - (cartesian_case.default_sizes[IDim], cartesian_case.default_sizes[KDim]), True, dtype=bool - ) - - @gtx.field_operator - def testee(a: cases.IKField, offset_field: cases.IKField) -> gtx.Field[[IDim, KDim], bool]: - a_i = a(as_offset(Ioff, offset_field)) - a_i_k = a_i(as_offset(Koff, offset_field)) - b_i = a(Ioff[1]) - b_i_k = b_i(Koff[1]) - return a_i_k == b_i_k - - out = cases.allocate(cartesian_case, testee, cases.RETURN)() - a = cases.allocate(cartesian_case, testee, "a").extend({IDim: (0, 1), KDim: (0, 1)})() - offset_field = cases.allocate(cartesian_case, testee, "offset_field").strategy( - cases.ConstInitializer(1) - )() - - cases.verify( - cartesian_case, - testee, - a, - offset_field, - out=out, - offset_provider={"Ioff": IDim, "Koff": KDim}, - ref=np.full_like(offset_field.asnumpy(), True, dtype=bool), - comparison=lambda out, ref: np.all(out == ref), - ) - - assert np.allclose(out.asnumpy(), ref) - - def test_nested_tuple_return(cartesian_case): @gtx.field_operator def pack_tuple(