diff --git a/pyproject.toml b/pyproject.toml index ddc792fd5f..64f52fdb8b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -262,6 +262,8 @@ markers = [ 'uses_ir_if_stmts', 'uses_lift: tests that require backend support for lift builtin function', 'uses_negative_modulo: tests that require backend support for modulo on negative numbers', + 'uses_offset_tag_differing_from_local_dim: tests that shift with an offset tag that differs from the name of its local dimension', + 'uses_offset_tag_differing_from_local_dim_in_reduction: tests that reduce over an offset whose tag differs from the name of its local dimension', 'uses_origin: tests that require backend support for domain origin', 'uses_reduce_with_lambda: tests that use lambdas as reduce functions', 'uses_reduction_with_only_sparse_fields: tests that require backend support for with sparse fields', diff --git a/src/gt4py/next/ffront/fbuiltins.py b/src/gt4py/next/ffront/fbuiltins.py index c4f93b13cf..2fd072a483 100644 --- a/src/gt4py/next/ffront/fbuiltins.py +++ b/src/gt4py/next/ffront/fbuiltins.py @@ -480,6 +480,8 @@ def impl( # guidelines for decision. @dataclasses.dataclass(frozen=True) class FieldOffset(runtime.Offset): + #: The tag, i.e. the offset-provider key. + value: str source: common.Dimension target: tuple[common.Dimension] | tuple[common.Dimension, common.Dimension] @@ -492,7 +494,7 @@ def __post_init__(self) -> None: raise ValueError("Second dimension in offset must be a local dimension.") def __gt_type__(self) -> ts.OffsetType: - return ts.OffsetType(source=self.source, target=self.target) + return ts.OffsetType(source=self.source, target=self.target, tag=self.value) def __getitem__(self, offset: int) -> common.Connectivity: """Serve as a connectivity factory.""" diff --git a/src/gt4py/next/ffront/foast_passes/type_deduction.py b/src/gt4py/next/ffront/foast_passes/type_deduction.py index f8f739aa62..738b105a61 100644 --- a/src/gt4py/next/ffront/foast_passes/type_deduction.py +++ b/src/gt4py/next/ffront/foast_passes/type_deduction.py @@ -456,13 +456,13 @@ def visit_Subscript(self, node: foast.Subscript, **kwargs: Any) -> foast.Subscri f"Tuples need to be indexed with literal integers, got '{node.index}'.", ) from ex new_type = types[index] - case ts.OffsetType(source=source, target=(target1, target2)): + case ts.OffsetType(source=source, target=(target1, target2), tag=tag): if not target2.kind == DimensionKind.LOCAL: raise errors.DSLError( new_value.location, "Second dimension in offset must be a local dimension." ) - new_type = ts.OffsetType(source=source, target=(target1,)) - case ts.OffsetType(source=source, target=(target,)): + new_type = ts.OffsetType(source=source, target=(target1,), tag=tag) + case ts.OffsetType(source=source, target=(target,), tag=tag): # for cartesian axes (e.g. I, J) the index of the subscript only # signifies the displacement in the respective dimension, # but does not change the target type. @@ -471,6 +471,19 @@ def visit_Subscript(self, node: foast.Subscript, **kwargs: Any) -> foast.Subscri new_value.location, "Source and target must be equal for offsets with a single target.", ) + if tag is None: + raise errors.DSLError( + new_value.location, + "Cannot index a dimension shift.", + notes=[ + ( + "A shift written as 'Dim + offset' already contains its" + " displacement, unlike a 'FieldOffset', which is indexed to" + " choose one." + ) + ], + hints=[f"Write the displacement directly, e.g. '{source.value} + 1'."], + ) new_type = new_value.type case ts.FieldType(dims=dims, dtype=dtype): # e.g. `field[LocalDim(42)]` @@ -769,7 +782,20 @@ def visit_Call(self, node: foast.Call, **kwargs: Any) -> foast.Call: ): raise errors.DSLError(node.location, "Functions can only be called directly.") elif isinstance(new_func.type, ts.FieldType): - pass + for arg in new_args: + # A Cartesian `FieldOffset` shifts by the index it is subscripted with, so it + # carries no displacement on its own. Only an offset with a local dimension is + # meaningful unsubscripted, as the neighbor access `field(Off)`. + if ( + isinstance(arg, (foast.Name, foast.Attribute)) + and isinstance(arg.type, ts.OffsetType) + and len(arg.type.target) == 1 + ): + raise errors.DSLError( + arg.location, + f"Cannot shift by the Cartesian offset '{arg!s}' without an index.", + hints=[f"Give the displacement, e.g. '{arg!s}[1]'."], + ) elif isinstance(new_func.type, ts.DimensionType): assert new_func.type.dim.kind == DimensionKind.LOCAL return foast.Call( diff --git a/src/gt4py/next/ffront/foast_to_gtir.py b/src/gt4py/next/ffront/foast_to_gtir.py index 68c8267b62..9b3c15d7eb 100644 --- a/src/gt4py/next/ffront/foast_to_gtir.py +++ b/src/gt4py/next/ffront/foast_to_gtir.py @@ -295,14 +295,17 @@ def _visit_shift(self, node: foast.Call, **kwargs: Any) -> itir.Expr: for arg in node.args: match arg: # `field(Off[idx])` - case foast.Subscript(value=foast.Name() as offset_name, index=index): + # (matched on the type, not the node, to also accept `mod.Off[idx]`) + case foast.Subscript( + value=foast.LocatedNode(type=ts.OffsetType(tag=str() as offset_tag)), + index=index, + ): # Constant folding to a `Literal` ensures that `index` becomes an `OffsetLiteral`, # which can be generated as compile-time value backend code. new_index = constant_folding.ConstantFolding.apply(self.visit(index, **kwargs)) assert isinstance(new_index, itir.Literal) - assert isinstance(offset_name.type, ts.OffsetType) current_expr = im.as_fieldop( - im.lambda_("__it")(im.deref(im.shift(offset_name.id, new_index)("__it"))) + im.lambda_("__it")(im.deref(im.shift(offset_tag, new_index)("__it"))) )(current_expr) # `field(Dim + idx)` (where `idx` is integer or half integer) case foast.BinOp( @@ -322,14 +325,6 @@ def _visit_shift(self, node: foast.Call, **kwargs: Any) -> itir.Expr: ) ) )(current_expr) - # `field(Off)` - case foast.Name(id=offset_name): - # only a single unstructured shift is supported so returning here is fine even though we - # are in a loop. - assert len(node.args) == 1 and len(arg.type.target) > 1 # type: ignore[attr-defined] # ensured by pattern - return im.as_fieldop_neighbors( - str(offset_name), self.visit(node.func, **kwargs) - ) # `field(as_offset(Off, offset_field))` case foast.Call(func=foast.Name(id="as_offset")): func_args = arg @@ -344,6 +339,12 @@ def _visit_shift(self, node: foast.Call, **kwargs: Any) -> itir.Expr: ) ) )(current_expr, offset_field) + # `field(Off)` + case foast.LocatedNode(type=ts.OffsetType(tag=str() as offset_tag, target=(_, _))): + # only a single unstructured shift is supported so returning here is fine even though we + # are in a loop. + assert len(node.args) == 1 + return im.as_fieldop_neighbors(offset_tag, self.visit(node.func, **kwargs)) case _: raise FieldOperatorLoweringError("Unexpected shift arguments!") diff --git a/src/gt4py/next/type_system/type_specifications.py b/src/gt4py/next/type_system/type_specifications.py index 59ac40f0f3..eb130bd994 100644 --- a/src/gt4py/next/type_system/type_specifications.py +++ b/src/gt4py/next/type_system/type_specifications.py @@ -74,9 +74,12 @@ class OffsetType(TypeSpec): # TODO(havogt): replace by ConnectivityType source: common.Dimension target: tuple[common.Dimension] | tuple[common.Dimension, common.Dimension] + #: The offset-provider key; `None` for the untagged Cartesian `Dim + offset`. + tag: Optional[common.Tag] = None def __str__(self) -> str: - return f"Offset[{self.source}, {self.target}]" + tag = "" if self.tag is None else f"{self.tag}: " + return f"Offset[{tag}{self.source}, {self.target}]" class ScalarKind(eve_types.IntEnum): diff --git a/tests/next_tests/definitions.py b/tests/next_tests/definitions.py index a75f3e08ff..86a4229860 100644 --- a/tests/next_tests/definitions.py +++ b/tests/next_tests/definitions.py @@ -98,6 +98,10 @@ class ProgramFormatterId(_PythonObjectIdMixin, str, enum.Enum): USES_INDEX_FIELDS = "uses_index_fields" USES_LIFT = "uses_lift" USES_NEGATIVE_MODULO = "uses_negative_modulo" +USES_OFFSET_TAG_DIFFERING_FROM_LOCAL_DIM = "uses_offset_tag_differing_from_local_dim" +USES_OFFSET_TAG_DIFFERING_FROM_LOCAL_DIM_IN_REDUCTION = ( + "uses_offset_tag_differing_from_local_dim_in_reduction" +) USES_ORIGIN = "uses_origin" USES_REDUCE_WITH_LAMBDA = "uses_reduce_with_lambda" USES_SCAN = "uses_scan" @@ -134,6 +138,14 @@ class ProgramFormatterId(_PythonObjectIdMixin, str, enum.Enum): REDUCTION_WITH_ONLY_SPARSE_FIELDS_MESSAGE = ( "We cannot unroll a reduction on a sparse field only (not clear if it is legal ITIR)" ) +#: An offset and its local dimension must currently share a name on most backends, because +#: the connectivity is looked up in the offset provider by the *local dimension's* name. +#: Lifted for the gtfn shift path by #1789; see +#: `regression_tests/ffront_tests/test_offset_dimensions_names.py`. +OFFSET_TAG_DIFFERING_FROM_LOCAL_DIM_MESSAGE = ( + "'{marker}': '{backend}' looks the connectivity up by the local dimension's name," + " so it must equal the offset tag" +) # Index-only vs. consequential markers: # A `uses_*` marker only affects execution if it appears in one of the skip lists below (and thus # in `BACKEND_SKIP_TEST_MATRIX`); such a marker is "consequential" -- it applies the listed @@ -169,6 +181,16 @@ class ProgramFormatterId(_PythonObjectIdMixin, str, enum.Enum): (USES_SCAN_IN_STENCIL, XFAIL, BINDINGS_UNSUPPORTED_MESSAGE), (USES_SPARSE_FIELDS, XFAIL, UNSUPPORTED_MESSAGE), (USES_TUPLE_ITERATOR, XFAIL, UNSUPPORTED_MESSAGE), + ( + USES_OFFSET_TAG_DIFFERING_FROM_LOCAL_DIM, + XFAIL, + OFFSET_TAG_DIFFERING_FROM_LOCAL_DIM_MESSAGE, + ), + ( + USES_OFFSET_TAG_DIFFERING_FROM_LOCAL_DIM_IN_REDUCTION, + XFAIL, + OFFSET_TAG_DIFFERING_FROM_LOCAL_DIM_MESSAGE, + ), ] ) EMBEDDED_SKIP_LIST = [ @@ -180,6 +202,11 @@ class ProgramFormatterId(_PythonObjectIdMixin, str, enum.Enum): ), # we can't extract the field type from scan args (EMBEDDED_CONCAT_WHERE_INFINITE_DOMAIN, XFAIL, UNSUPPORTED_MESSAGE), (EMBEDDED_CONCAT_WHERE_NON_CONTIGUOUS_DOMAIN, XFAIL, UNSUPPORTED_MESSAGE), + ( + USES_OFFSET_TAG_DIFFERING_FROM_LOCAL_DIM_IN_REDUCTION, + XFAIL, + OFFSET_TAG_DIFFERING_FROM_LOCAL_DIM_MESSAGE, + ), ] JAX_EMBEDDED_SKIP_LIST = EMBEDDED_SKIP_LIST + [ (USES_PROGRAM_WITH_SLICED_OUT_ARGUMENTS, XFAIL, UNSUPPORTED_MESSAGE), @@ -190,7 +217,15 @@ class ProgramFormatterId(_PythonObjectIdMixin, str, enum.Enum): (USES_TUPLES_ARGS_WITH_DIFFERENT_BUT_PROMOTABLE_DIMS, XFAIL, UNSUPPORTED_MESSAGE), (USES_CONCAT_WHERE, XFAIL, UNSUPPORTED_MESSAGE), ] -GTIR_EMBEDDED_SKIP_LIST = ROUNDTRIP_SKIP_LIST + [] +GTIR_EMBEDDED_SKIP_LIST = ROUNDTRIP_SKIP_LIST + [ + # NOTE: not in `ROUNDTRIP_SKIP_LIST`: the roundtrip backend passes this, only the + # lower-level `iterator/embedded.py` execution keys on the local dimension's name. + ( + USES_OFFSET_TAG_DIFFERING_FROM_LOCAL_DIM_IN_REDUCTION, + XFAIL, + OFFSET_TAG_DIFFERING_FROM_LOCAL_DIM_MESSAGE, + ), +] GTFN_SKIP_TEST_LIST = ( COMMON_SKIP_TEST_LIST + DOMAIN_INFERENCE_SKIP_LIST @@ -201,6 +236,12 @@ class ProgramFormatterId(_PythonObjectIdMixin, str, enum.Enum): (USES_STRIDED_NEIGHBOR_OFFSET, XFAIL, BINDINGS_UNSUPPORTED_MESSAGE), # max_over broken, see https://github.com/GridTools/gt4py/issues/1289 (USES_MAX_OVER, XFAIL, UNSUPPORTED_MESSAGE), + # NOTE: only the reduction; #1789 lifted this for the shift path. + ( + USES_OFFSET_TAG_DIFFERING_FROM_LOCAL_DIM_IN_REDUCTION, + XFAIL, + OFFSET_TAG_DIFFERING_FROM_LOCAL_DIM_MESSAGE, + ), ] ) diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_import_from_mod.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_import_from_mod.py index 899197c43f..b41f3a9c89 100644 --- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_import_from_mod.py +++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_import_from_mod.py @@ -10,11 +10,11 @@ import numpy as np import gt4py.next as gtx -from gt4py.next import broadcast, astype, int32 +from gt4py.next import broadcast, astype, common, int32, neighbor_sum from next_tests import integration_tests from next_tests.integration_tests import cases -from next_tests.integration_tests.cases import cartesian_case, IDim, KDim +from next_tests.integration_tests.cases import cartesian_case, unstructured_case, IDim, KDim from next_tests.integration_tests.cases_utils import ( exec_alloc_descriptor, @@ -54,6 +54,33 @@ def mod_prog(f: cases.IField, isize: int32, ksize: int32, out: cases.IKField): cases.verify(cartesian_case, mod_prog, f, isize, ksize, out=out, ref=expected) +@pytest.mark.uses_unstructured_shift +def test_import_offset_module_unstructured_shift(unstructured_case): + @gtx.field_operator + def testee(a: cases.EField) -> cases.VField: + return neighbor_sum(a(cases.V2E), axis=cases.V2EDim) + + v2e_table = unstructured_case.offset_provider["V2E"].asnumpy() + cases.verify_with_default_data( + unstructured_case, + testee, + ref=lambda a: np.sum(a[v2e_table], axis=1, where=v2e_table != common._DEFAULT_SKIP_VALUE), + ) + + +@pytest.mark.uses_unstructured_shift +def test_import_offset_module_sparse_shift(unstructured_case): + @gtx.field_operator + def testee(a: cases.VField) -> cases.EField: + return a(cases.E2V[0]) + + cases.verify_with_default_data( + unstructured_case, + testee, + ref=lambda a: a[unstructured_case.offset_provider["E2V"].asnumpy()[:, 0]], + ) + + # TODO: these set of features should be allowed as module imports in a later PR def test_import_module_errors_future_allowed(cartesian_case): from ....artifacts.dummy_package import dummy_module diff --git a/tests/next_tests/regression_tests/ffront_tests/test_offset_dimensions_names.py b/tests/next_tests/regression_tests/ffront_tests/test_offset_dimensions_names.py index 62ddd289eb..89a72b428f 100644 --- a/tests/next_tests/regression_tests/ffront_tests/test_offset_dimensions_names.py +++ b/tests/next_tests/regression_tests/ffront_tests/test_offset_dimensions_names.py @@ -6,59 +6,160 @@ # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause +""" +Regression tests for the four independently authored names of one connectivity. + +Using a single connectivity requires four strings to agree, none of which is +checked against the others at declaration time: + + N1 the `FieldOffset` tag `FieldOffset("V2E", ...)` + N2 the Python variable it is bound to `V2E = FieldOffset(...)` + N3 the local dimension's name `Dimension("V2E", kind=LOCAL)` + N4 the offset-provider key `offset_provider={"V2E": ...}` + +The `V2EDim = Dimension("V2E")` convention makes all four equal, which hides +which one each execution path actually uses. These tests break the convention +deliberately, one name at a time, so the real requirement is visible. +""" + +import numpy as np import pytest from gt4py import next as gtx -from gt4py.next import Dims, Field, common, constructors +from gt4py.next import Dims, Field, common, constructors, neighbor_sum from next_tests import definitions as test_defs from next_tests.integration_tests import cases from next_tests.integration_tests import cases_utils +from next_tests.integration_tests.cases_utils import ( # noqa: F401 [unused-import] # fixture + exec_alloc_descriptor, +) V = gtx.Dimension("V") E = gtx.Dimension("E") + +#: N1 == N3 == N4, but N2 differs: the tag is `TaggedOff`, the variable is `off_a`. +TaggedOffDim = gtx.Dimension("TaggedOff", kind=common.DimensionKind.LOCAL) +off_a = gtx.FieldOffset("TaggedOff", source=E, target=(V, TaggedOffDim)) + +#: N1 == N2 == N4, but N3 differs: the local dimension is `Neigh`, the tag is `OffB`. Neigh = gtx.Dimension("Neigh", kind=common.DimensionKind.LOCAL) -Off = gtx.FieldOffset("Off", source=E, target=(V, Neigh)) +OffB = gtx.FieldOffset("OffB", source=E, target=(V, Neigh)) -@pytest.fixture -def case(): - exec_alloc_descriptor = test_defs.ProgramBackendId.GTFN_CPU.load() +def _case(exec_alloc_descriptor, tag: str, local_dim: common.Dimension) -> cases.Case: + """ + A `Case` whose offset provider holds exactly one connectivity, keyed on its tag. + + One entry per `Case` on purpose: DaCe walks every provider entry while building the + SDFG, and looks a connectivity up by its *local dimension's* name + (`gtir_to_sdfg.py`, constraint A4). A second, non-conforming entry would therefore + fail a program that does not even use it, and the cell under test would be measuring + the wrong thing. + """ mesh = cases_utils.simple_mesh(exec_alloc_descriptor.allocator) - v2e_arr = mesh.offset_provider["V2E"].ndarray + # NOTE: `.asnumpy()`, not `.ndarray`: under a GPU allocator the latter is a device + # array, and `simple_mesh` builds the table from NumPy anyway. + v2e_arr = mesh.offset_provider["V2E"].asnumpy() return cases.Case( - exec_alloc_descriptor, + ( + None + if isinstance(exec_alloc_descriptor, test_defs.EmbeddedDummyBackend) + else exec_alloc_descriptor + ), offset_provider={ - "Off": constructors.as_connectivity( - domain={V: v2e_arr.shape[0], Neigh: 4}, + tag: constructors.as_connectivity( + domain={V: v2e_arr.shape[0], local_dim: v2e_arr.shape[1]}, codomain=E, data=v2e_arr, skip_value=None, allocator=exec_alloc_descriptor.allocator, - ), - }, - default_sizes={ - V: mesh.num_vertices, - E: mesh.num_edges, + ) }, + default_sizes={V: mesh.num_vertices, E: mesh.num_edges}, grid_type=common.GridType.UNSTRUCTURED, allocator=exec_alloc_descriptor.allocator, ) -def test_offset_dimension_name_differ(case): +@pytest.fixture +def case_tag_vs_variable_name(exec_alloc_descriptor): + return _case(exec_alloc_descriptor, "TaggedOff", TaggedOffDim) + + +@pytest.fixture +def case_tag_vs_local_dim(exec_alloc_descriptor): + return _case(exec_alloc_descriptor, "OffB", Neigh) + + +def _neighbor_table(case: cases.Case, tag: str) -> np.ndarray: + return case.offset_provider[tag].asnumpy() + + +# --- N2: the tag differs from the Python variable name ---------------------------- +# Lowering used to emit the *variable* name as the IR shift tag, so embedded and +# compiled execution of the same program needed different provider keys. + + +def test_shift_tag_differs_from_variable_name(case_tag_vs_variable_name): + @gtx.field_operator + def foo(a: Field[Dims[E], float]) -> Field[Dims[V], float]: + return a(off_a[1]) + + cases.verify_with_default_data( + case_tag_vs_variable_name, + foo, + lambda a: a[_neighbor_table(case_tag_vs_variable_name, "TaggedOff")[:, 1]], + ) + + +def test_reduction_tag_differs_from_variable_name(case_tag_vs_variable_name): + @gtx.field_operator + def foo(a: Field[Dims[E], float]) -> Field[Dims[V], float]: + return neighbor_sum(a(off_a), axis=TaggedOffDim) + + cases.verify_with_default_data( + case_tag_vs_variable_name, + foo, + lambda a: np.sum(a[_neighbor_table(case_tag_vs_variable_name, "TaggedOff")], axis=1), + ) + + +# --- N3: the tag differs from the local dimension's name -------------------------- +# Lifted for the gtfn shift path by #1789; still required elsewhere, which is what +# the markers below record. + + +@pytest.mark.uses_offset_tag_differing_from_local_dim +def test_shift_tag_differs_from_local_dim_name(case_tag_vs_local_dim): """ - Ensure that gtfn works with offset name that differs from the name of the local dimension. + Ensure a shift works with an offset tag that differs from the local dimension's name. - If the value of the `NeighborConnectivityType.neighbor_dim` did not match the `FieldOffset` value, - gtfn would silently ignore the neighbor index, see https://github.com/GridTools/gridtools/pull/1814. + If `NeighborConnectivityType.neighbor_dim` did not match the `FieldOffset` value, + gtfn would silently ignore the neighbor index, see + https://github.com/GridTools/gridtools/pull/1814. """ @gtx.field_operator def foo(a: Field[Dims[E], float]) -> Field[Dims[V], float]: - return a(Off[1]) + return a(OffB[1]) + + cases.verify_with_default_data( + case_tag_vs_local_dim, + foo, + lambda a: a[_neighbor_table(case_tag_vs_local_dim, "OffB")[:, 1]], + ) + + +@pytest.mark.uses_offset_tag_differing_from_local_dim_in_reduction +def test_reduction_tag_differs_from_local_dim_name(case_tag_vs_local_dim): + @gtx.field_operator + def foo(a: Field[Dims[E], float]) -> Field[Dims[V], float]: + return neighbor_sum(a(OffB), axis=Neigh) cases.verify_with_default_data( - case, foo, lambda a: a[case.offset_provider["Off"].ndarray[:, 1]] + case_tag_vs_local_dim, + foo, + lambda a: np.sum(a[_neighbor_table(case_tag_vs_local_dim, "OffB")], axis=1), ) diff --git a/tests/next_tests/unit_tests/ffront_tests/test_diagnostic_messages.py b/tests/next_tests/unit_tests/ffront_tests/test_diagnostic_messages.py index 62072cab18..7e2d4fed6e 100644 --- a/tests/next_tests/unit_tests/ffront_tests/test_diagnostic_messages.py +++ b/tests/next_tests/unit_tests/ffront_tests/test_diagnostic_messages.py @@ -28,6 +28,7 @@ IDim = gtx.Dimension("IDim") +IOff = gtx.FieldOffset("Ioff", source=IDim, target=(IDim,)) # A PEP 695 alias whose value raises when it is evaluated, standing in for the # common case of a typo'd dtype ('np.foat64') inside an alias definition. @@ -356,3 +357,29 @@ def broken(a: BrokenFieldAlias) -> gtx.Field[[IDim], float64]: assert err.location.line == err.location.end_line assert err.location.end_column - err.location.column == len("a: BrokenFieldAlias") assert re.search(r"\| +\^{19}(?!\^)", str(err)), str(err) + + +def test_unindexed_cartesian_offset_names_the_offset_as_written(): + # The tag of 'IOff' is 'Ioff'; the message has to quote what the user wrote. + def unindexed(a: gtx.Field[[IDim], float64]) -> gtx.Field[[IDim], float64]: + return a(IOff) + + err = parse_error(unindexed) + + assert err.message == "Cannot shift by the Cartesian offset 'IOff' without an index." + assert err.hints == ["Give the displacement, e.g. 'IOff[1]'."] + rendered = str(err) + assert "return a(IOff)" in rendered + assert re.search(r"\| +\^{4}(?!\^)", rendered), rendered + + +def test_indexed_dimension_shift_is_rejected_with_a_hint(): + def indexed_shift(a: gtx.Field[[IDim], float64]) -> gtx.Field[[IDim], float64]: + return a((IDim + 1)[0]) + + err = parse_error(indexed_shift) + + assert err.message == "Cannot index a dimension shift." + assert err.hints == ["Write the displacement directly, e.g. 'IDim + 1'."] + assert any("already contains its displacement" in note for note in err.notes) + assert "return a((IDim + 1)[0])" in str(err) diff --git a/tests/next_tests/unit_tests/ffront_tests/test_foast_to_gtir.py b/tests/next_tests/unit_tests/ffront_tests/test_foast_to_gtir.py index 407088e4f8..eb1927a776 100644 --- a/tests/next_tests/unit_tests/ffront_tests/test_foast_to_gtir.py +++ b/tests/next_tests/unit_tests/ffront_tests/test_foast_to_gtir.py @@ -50,6 +50,10 @@ TDim = gtx.Dimension("TDim") TOff = gtx.FieldOffset("TDim", source=TDim, target=(TDim,)) +#: An offset whose tag differs from the name of the Python variable it is bound to, and +#: from the name of its local dimension. Lowering must emit the *tag*. +RenamedV2EDim = gtx.Dimension("RenamedLocal", gtx.DimensionKind.LOCAL) +renamed_v2e = gtx.FieldOffset("RenamedTag", source=Edge, target=(Vertex, RenamedV2EDim)) UDim = gtx.Dimension("UDim") @@ -782,6 +786,34 @@ def foo(edge_f: gtx.Field[gtx.Dims[Edge], float64]): assert lowered.expr == reference +def test_unstructured_shift_lowering_emits_offset_tag_not_variable_name(): + """The IR shift tag is the offset's tag, not the variable the offset is bound to.""" + + def foo(edge_f: gtx.Field[[Edge], float64]): + return edge_f(renamed_v2e[1]) + + parsed = FieldOperatorParser.apply_to_function(foo) + lowered = FieldOperatorLowering.apply(parsed) + + reference = im.as_fieldop(im.lambda_("__it")(im.deref(im.shift("RenamedTag", 1)("__it"))))( + "edge_f" + ) + + assert lowered.expr == reference + + +def test_unstructured_neighbors_lowering_emits_offset_tag_not_variable_name(): + def foo(edge_f: gtx.Field[[Edge], float64]): + return edge_f(renamed_v2e) + + parsed = FieldOperatorParser.apply_to_function(foo) + lowered = FieldOperatorLowering.apply(parsed) + + reference = im.as_fieldop_neighbors("RenamedTag", "edge_f") + + assert lowered.expr == reference + + def test_reduction_lowering_neighbor_sum(): def foo(edge_f: gtx.Field[[Edge], float64]): return neighbor_sum(edge_f(V2E), axis=V2EDim)