diff --git a/src/gt4py/next/ffront/foast_to_itir.py b/src/gt4py/next/ffront/foast_to_itir.py index b5d84d4bd6..e2fdf165ce 100644 --- a/src/gt4py/next/ffront/foast_to_itir.py +++ b/src/gt4py/next/ffront/foast_to_itir.py @@ -233,6 +233,10 @@ def visit_Name(self, node: foast.Name, **kwargs: Any) -> itir.SymRef: return im.ref(node.id) def visit_Subscript(self, node: foast.Subscript, **kwargs: Any) -> itir.Expr: + if isinstance(node.type, ts.OffsetType): + # TODO(havogt): problem is that we need to resolve `E2V[0]` -> `E2V, 0` (not as tuple) as it needs to be inserted as `X` in `shift(X)` + raise NotImplementedError("Offsets can only be used inlined in the premap operation.") + return im.tuple_get(node.index, self.visit(node.value, **kwargs)) def visit_TupleExpr(self, node: foast.TupleExpr, **kwargs: Any) -> itir.Expr: diff --git a/src/gt4py/next/ffront/func_to_foast.py b/src/gt4py/next/ffront/func_to_foast.py index 120068220f..e16010e867 100644 --- a/src/gt4py/next/ffront/func_to_foast.py +++ b/src/gt4py/next/ffront/func_to_foast.py @@ -300,7 +300,7 @@ def visit_Assign( if not isinstance(target, ast.Name): raise errors.DSLError(self.get_location(node), "Can only assign to names.") new_value = self.visit(node.value) - constraint_type: Type[ts.DataType] = ts.DataType + constraint_type: Type[ts.TypeSpec] = ts.TypeSpec if isinstance(new_value, foast.TupleExpr): constraint_type = ts.TupleType elif ( @@ -308,6 +308,8 @@ def visit_Assign( and type_info.type_class(new_value.type) is ts.ScalarType ): constraint_type = ts.ScalarType + else: + pass return foast.Assign( target=foast.DataSymbol( id=target.id, 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 26825f9e4c..8a3324b6e8 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 @@ -16,9 +16,7 @@ astype, broadcast, common, - constructors, errors, - field_utils, float32, float64, int32, @@ -28,14 +26,11 @@ ) from gt4py.next.ffront.experimental import as_offset from gt4py.next.program_processors.runners import gtfn -from gt4py.next.type_system import type_specifications as ts from next_tests.integration_tests import cases from next_tests.integration_tests.cases import ( - C2E, E2V, V2E, - E2VDim, IDim, Ioff, JDim, @@ -43,7 +38,6 @@ Koff, V2EDim, Vertex, - Edge, cartesian_case, unstructured_case, ) @@ -73,92 +67,6 @@ 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): - @gtx.field_operator - def testee(a: cases.IJKField) -> cases.IJKField: - return a(Ioff[1]) - - a = cases.allocate(cartesian_case, testee, "a").extend({IDim: (0, 1)})() - out = cases.allocate(cartesian_case, testee, cases.RETURN)() - - cases.verify(cartesian_case, testee, a, out=out, ref=a[1:]) - - -@pytest.mark.uses_unstructured_shift -def test_unstructured_shift(unstructured_case): - @gtx.field_operator - def testee(a: cases.VField) -> cases.EField: - return a(E2V[0]) - - cases.verify_with_default_data( - unstructured_case, - testee, - ref=lambda a: a[unstructured_case.offset_provider["E2V"].table[:, 0]], - ) - - -@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: - return inp(E2V[0])(C2E[0]) - - @gtx.field_operator - def composed_shift_unstructured_intermediate_result(inp: cases.VField) -> cases.CField: - tmp = inp(E2V[0]) - return tmp(C2E[0]) - - @gtx.field_operator - def shift_e2v(inp: cases.VField) -> cases.EField: - return inp(E2V[0]) - - @gtx.field_operator - def composed_shift_unstructured(inp: cases.VField) -> cases.CField: - return shift_e2v(inp)(C2E[0]) - - cases.verify_with_default_data( - unstructured_case, - composed_shift_unstructured_flat, - ref=lambda inp: inp[unstructured_case.offset_provider["E2V"].table[:, 0]][ - unstructured_case.offset_provider["C2E"].table[:, 0] - ], - ) - - cases.verify_with_default_data( - unstructured_case, - composed_shift_unstructured_intermediate_result, - ref=lambda inp: inp[unstructured_case.offset_provider["E2V"].table[:, 0]][ - unstructured_case.offset_provider["C2E"].table[:, 0] - ], - comparison=lambda inp, tmp: np.all(inp == tmp), - ) - - cases.verify_with_default_data( - unstructured_case, - composed_shift_unstructured, - ref=lambda inp: inp[unstructured_case.offset_provider["E2V"].table[:, 0]][ - unstructured_case.offset_provider["C2E"].table[:, 0] - ], - ) - - -@pytest.mark.uses_cartesian_shift -def test_fold_shifts(cartesian_case): - """Shifting the result of an addition should work.""" - - @gtx.field_operator - def testee(a: cases.IJKField, b: cases.IJKField) -> cases.IJKField: - tmp = a + b(Ioff[1]) - return tmp(Ioff[1]) - - a = cases.allocate(cartesian_case, testee, "a").extend({cases.IDim: (0, 1)})() - b = cases.allocate(cartesian_case, testee, "b").extend({cases.IDim: (0, 2)})() - out = cases.allocate(cartesian_case, testee, cases.RETURN)() - - cases.verify(cartesian_case, testee, a, b, out=out, ref=a.ndarray[1:] + b.ndarray[2:]) - - def test_tuples(cartesian_case): @gtx.field_operator def testee(a: cases.IJKFloatField, b: cases.IJKFloatField) -> cases.IJKFloatField: @@ -483,50 +391,6 @@ 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 - def testee(a: cases.VField) -> cases.VField: - tmp = neighbor_sum(a(E2V), axis=E2VDim) - tmp_2 = neighbor_sum(tmp(V2E), axis=V2EDim) - return tmp_2 - - cases.verify_with_default_data( - unstructured_case, - testee, - ref=lambda a: np.sum( - np.sum(a[unstructured_case.offset_provider["E2V"].table], axis=1, initial=0)[ - unstructured_case.offset_provider["V2E"].table - ], - axis=1, - where=unstructured_case.offset_provider["V2E"].table != common._DEFAULT_SKIP_VALUE, - ), - comparison=lambda a, tmp_2: np.all(a == tmp_2), - ) - - -@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 - def testee(inp: cases.EField) -> cases.EField: - tmp = inp(V2E) - tmp2 = tmp(E2V) - return neighbor_sum(neighbor_sum(tmp2, axis=V2EDim), axis=E2VDim) - - cases.verify_with_default_data( - unstructured_case, - testee, - ref=lambda inp: np.sum( - np.sum(inp[unstructured_case.offset_provider["V2E"].table], axis=1)[ - unstructured_case.offset_provider["E2V"].table - ], - axis=1, - ), - ) - - @pytest.mark.uses_unstructured_shift @pytest.mark.uses_tuple_returns def test_tuple_return_2(unstructured_case): diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_premap.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_premap.py new file mode 100644 index 0000000000..12613064d9 --- /dev/null +++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_premap.py @@ -0,0 +1,174 @@ +# GT4Py - GridTools Framework +# +# Copyright (c) 2014-2024, ETH Zurich +# All rights reserved. +# +# Please, refer to the LICENSE file in the root directory. +# SPDX-License-Identifier: BSD-3-Clause + +import numpy as np +import pytest + +from gt4py import next as gtx +from gt4py.next import common, neighbor_sum + +from next_tests.integration_tests import cases +from next_tests.integration_tests.cases import ( + C2E, + E2V, + V2E, + E2VDim, + IDim, + Ioff, + V2EDim, + cartesian_case, + unstructured_case, +) +from next_tests.integration_tests.feature_tests.ffront_tests.ffront_test_utils import ( + exec_alloc_descriptor, + mesh_descriptor, +) + + +@pytest.mark.uses_cartesian_shift +def test_cartesian_shift(cartesian_case): + @gtx.field_operator + def testee(a: cases.IJKField) -> cases.IJKField: + return a(Ioff[1]) + + a = cases.allocate(cartesian_case, testee, "a").extend({IDim: (0, 1)})() + out = cases.allocate(cartesian_case, testee, cases.RETURN)() + + cases.verify(cartesian_case, testee, a, out=out, ref=a[1:]) + + +@pytest.mark.uses_unstructured_shift +def test_unstructured_shift(unstructured_case): + @gtx.field_operator + def testee(a: cases.VField) -> cases.EField: + return a(E2V[0]) + + cases.verify_with_default_data( + unstructured_case, + testee, + ref=lambda a: a[unstructured_case.offset_provider["E2V"].table[:, 0]], + ) + + +@pytest.mark.uses_unstructured_shift +def test_unstructured_shift_offset_symbol(unstructured_case): + @gtx.field_operator + def testee(a: cases.VField) -> cases.EField: + e2v0 = E2V[0] + return a(e2v0) + + cases.verify_with_default_data( + unstructured_case, + testee, + ref=lambda a: a[unstructured_case.offset_provider["E2V"].table[:, 0]], + ) + + +@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: + return inp(E2V[0])(C2E[0]) + + @gtx.field_operator + def composed_shift_unstructured_intermediate_result(inp: cases.VField) -> cases.CField: + tmp = inp(E2V[0]) + return tmp(C2E[0]) + + @gtx.field_operator + def shift_e2v(inp: cases.VField) -> cases.EField: + return inp(E2V[0]) + + @gtx.field_operator + def composed_shift_unstructured(inp: cases.VField) -> cases.CField: + return shift_e2v(inp)(C2E[0]) + + cases.verify_with_default_data( + unstructured_case, + composed_shift_unstructured_flat, + ref=lambda inp: inp[unstructured_case.offset_provider["E2V"].table[:, 0]][ + unstructured_case.offset_provider["C2E"].table[:, 0] + ], + ) + + cases.verify_with_default_data( + unstructured_case, + composed_shift_unstructured_intermediate_result, + ref=lambda inp: inp[unstructured_case.offset_provider["E2V"].table[:, 0]][ + unstructured_case.offset_provider["C2E"].table[:, 0] + ], + comparison=lambda inp, tmp: np.all(inp == tmp), + ) + + cases.verify_with_default_data( + unstructured_case, + composed_shift_unstructured, + ref=lambda inp: inp[unstructured_case.offset_provider["E2V"].table[:, 0]][ + unstructured_case.offset_provider["C2E"].table[:, 0] + ], + ) + + +@pytest.mark.uses_cartesian_shift +def test_fold_shifts(cartesian_case): + """Shifting the result of an addition should work.""" + + @gtx.field_operator + def testee(a: cases.IJKField, b: cases.IJKField) -> cases.IJKField: + tmp = a + b(Ioff[1]) + return tmp(Ioff[1]) + + a = cases.allocate(cartesian_case, testee, "a").extend({cases.IDim: (0, 1)})() + b = cases.allocate(cartesian_case, testee, "b").extend({cases.IDim: (0, 2)})() + out = cases.allocate(cartesian_case, testee, cases.RETURN)() + + cases.verify(cartesian_case, testee, a, b, out=out, ref=a.ndarray[1:] + b.ndarray[2:]) + + +@pytest.mark.uses_unstructured_shift +@pytest.mark.uses_reduction_over_lift_expressions +def test_nested_reduction(unstructured_case): + @gtx.field_operator + def testee(a: cases.VField) -> cases.VField: + tmp = neighbor_sum(a(E2V), axis=E2VDim) + tmp_2 = neighbor_sum(tmp(V2E), axis=V2EDim) + return tmp_2 + + cases.verify_with_default_data( + unstructured_case, + testee, + ref=lambda a: np.sum( + np.sum(a[unstructured_case.offset_provider["E2V"].table], axis=1, initial=0)[ + unstructured_case.offset_provider["V2E"].table + ], + axis=1, + where=unstructured_case.offset_provider["V2E"].table != common._DEFAULT_SKIP_VALUE, + ), + comparison=lambda a, tmp_2: np.all(a == tmp_2), + ) + + +@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 + def testee(inp: cases.EField) -> cases.EField: + tmp = inp(V2E) + tmp2 = tmp(E2V) + return neighbor_sum(neighbor_sum(tmp2, axis=V2EDim), axis=E2VDim) + + cases.verify_with_default_data( + unstructured_case, + testee, + ref=lambda inp: np.sum( + np.sum(inp[unstructured_case.offset_provider["V2E"].table], axis=1)[ + unstructured_case.offset_provider["E2V"].table + ], + axis=1, + ), + )