From 12eacd16f38def70719bf8f31e4e6885413e4320 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Enrique=20Gonz=C3=A1lez=20Paredes?= Date: Fri, 18 Sep 2026 03:47:35 +0200 Subject: [PATCH 1/4] fix[next]: lower unstructured shifts with the offset's own tag Lowering emitted the *Python variable name* an offset was bound to as the IR shift tag, because `ts.OffsetType` did not carry the tag. Embedded execution keys on `FieldOffset.value`, so the same program needed a different offset provider depending on how it was run: MyOff = FieldOffset("TAGNAME", ...) embedded: {"TAGNAME": conn} OK ; {"MyOff": conn} -> KeyError 'TAGNAME' compiled: {"MyOff": conn} OK ; {"TAGNAME": conn} -> KeyError 'MyOff' `ts.OffsetType` now carries `tag`. It is `Optional`, not required: a Cartesian shift built from `Dim + offset` has no tag and needs none, since it lowers to a `CartesianOffset` carrying both dimensions with no provider lookup. Subscripting (`Off[1]`) drops the local dimension but propagates the tag, which is the offset's identity. The regression test grows from one cell -- `a(Off[1])` on gtfn -- to the cross product of {shift, neighbor_sum} x {tag != variable name, tag != local dimension name} over the whole backend matrix. Each `Case` holds exactly one connectivity on purpose: DaCe walks every offset-provider entry while building the SDFG and looks a connectivity up by its *local dimension's* name, so a second, non-conforming entry fails a program that does not use it. The cells that still fail are marked, per backend, from measurement: * `uses_offset_tag_differing_from_local_dim` -- DaCe only; the gtfn shift path was fixed in #1789. * `uses_offset_tag_differing_from_local_dim_in_reduction` -- embedded, gtfn, DaCe and the lower-level `iterator/embedded.py` execution. Notably *not* the roundtrip backend, which passes; only `roundtrip.gtir` fails. Both remaining constraints are the same one: those paths resolve a connectivity through the local dimension's name rather than the offset's identity. Fixing that needs a back-pointer from the local dimension to its connectivity, which is a separate change. --- pyproject.toml | 2 + src/gt4py/next/ffront/fbuiltins.py | 6 +- .../ffront/foast_passes/type_deduction.py | 6 +- src/gt4py/next/ffront/foast_to_gtir.py | 29 +++- .../next/type_system/type_specifications.py | 13 +- tests/next_tests/definitions.py | 43 +++++- .../test_offset_dimensions_names.py | 141 +++++++++++++++--- .../ffront_tests/test_foast_to_gtir.py | 32 ++++ 8 files changed, 244 insertions(+), 28 deletions(-) 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..6e592a7d08 100644 --- a/src/gt4py/next/ffront/fbuiltins.py +++ b/src/gt4py/next/ffront/fbuiltins.py @@ -480,6 +480,10 @@ def impl( # guidelines for decision. @dataclasses.dataclass(frozen=True) class FieldOffset(runtime.Offset): + #: Narrows `runtime.Offset.value`, which is `int | str`. A `FieldOffset`'s value is its + #: tag: the key its connectivity has in the offset provider, and what lowering emits as + #: the shift. An integer tag has no meaning here, and every consumer treats it as a name. + value: str source: common.Dimension target: tuple[common.Dimension] | tuple[common.Dimension, common.Dimension] @@ -492,7 +496,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..36db949bb8 100644 --- a/src/gt4py/next/ffront/foast_passes/type_deduction.py +++ b/src/gt4py/next/ffront/foast_passes/type_deduction.py @@ -456,12 +456,14 @@ 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,)) + # NOTE: `tag` is propagated: subscripting drops the local dimension but not + # the identity of the offset, which lowering needs to emit the shift. + new_type = ts.OffsetType(source=source, target=(target1,), tag=tag) case ts.OffsetType(source=source, target=(target,)): # for cartesian axes (e.g. I, J) the index of the subscript only # signifies the displacement in the respective dimension, diff --git a/src/gt4py/next/ffront/foast_to_gtir.py b/src/gt4py/next/ffront/foast_to_gtir.py index 68c8267b62..aa369fac1b 100644 --- a/src/gt4py/next/ffront/foast_to_gtir.py +++ b/src/gt4py/next/ffront/foast_to_gtir.py @@ -66,6 +66,26 @@ def promote_to_list(node_type: ts.TypeSpec) -> Callable[[itir.Expr], itir.Expr]: return lambda x: x +def _offset_tag(offset_type: ts.OffsetType) -> str: + """ + Return the tag to emit for a `FieldOffset`-based shift, i.e. its offset-provider key. + + Covers both the unstructured form and a Cartesian `FieldOffset` subscript + (`a(Koff[1])`), which reaches the same branch. + + This used to be the name of the Python variable the `FieldOffset` was bound to, which + is not the offset's identity: a declaration `Off = FieldOffset("Tag", ...)` made + compiled backends look up `'Off'` while embedded execution looked up `'Tag'`, so the + same program needed a different offset provider depending on how it was run. + """ + assert offset_type.tag is not None, ( + f"Offset '{offset_type}' has no tag. Only a Cartesian shift written as" + " 'Dim + offset' is untagged, and that lowers to a 'CartesianOffset' carrying both" + " dimensions, without reaching this function." + ) + return offset_type.tag + + @dataclasses.dataclass class FieldOperatorLowering(eve.PreserveLocationVisitor, eve.NodeTranslator): """ @@ -302,7 +322,9 @@ def _visit_shift(self, node: foast.Call, **kwargs: Any) -> itir.Expr: 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(offset_name.type), new_index)("__it")) + ) )(current_expr) # `field(Dim + idx)` (where `idx` is integer or half integer) case foast.BinOp( @@ -323,12 +345,13 @@ def _visit_shift(self, node: foast.Call, **kwargs: Any) -> itir.Expr: ) )(current_expr) # `field(Off)` - case foast.Name(id=offset_name): + case foast.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 + assert isinstance(arg.type, ts.OffsetType) return im.as_fieldop_neighbors( - str(offset_name), self.visit(node.func, **kwargs) + _offset_tag(arg.type), self.visit(node.func, **kwargs) ) # `field(as_offset(Off, offset_field))` case foast.Call(func=foast.Name(id="as_offset")): diff --git a/src/gt4py/next/type_system/type_specifications.py b/src/gt4py/next/type_system/type_specifications.py index 59ac40f0f3..21efe3efe3 100644 --- a/src/gt4py/next/type_system/type_specifications.py +++ b/src/gt4py/next/type_system/type_specifications.py @@ -74,9 +74,20 @@ class OffsetType(TypeSpec): # TODO(havogt): replace by ConnectivityType source: common.Dimension target: tuple[common.Dimension] | tuple[common.Dimension, common.Dimension] + #: The offset's own tag, i.e. the key its connectivity has in the offset provider. + #: + #: `None` for a Cartesian shift built from `Dim + offset`, which has no tag and needs + #: none: it lowers to a `CartesianOffset` carrying the two dimensions, with no lookup. + #: Set for every offset that comes from a `FieldOffset` declaration, so that lowering + #: does not have to fall back to the name of the Python variable the declaration + #: happens to be bound to. + tag: Optional[common.Tag] = None def __str__(self) -> str: - return f"Offset[{self.source}, {self.target}]" + # NOTE: the tag is part of the identity, so two offsets over the same dimensions + # must not render identically in diagnostics. + 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/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_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) From 36e76a6000ba4208865519b2fe54ee0f641fec58 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Enrique=20Gonz=C3=A1lez=20Paredes?= Date: Fri, 18 Sep 2026 15:14:29 +0200 Subject: [PATCH 2/4] fix[next]: lower shifts by module-qualified offsets `field(mod.Off)` and `field(mod.Off[idx])` failed lowering with "Unexpected shift arguments!": both shift arms of `_visit_shift` matched a `foast.Name`, and a module-qualified offset is a `foast.Attribute`. Embedded execution was unaffected, so again the execution paths disagreed. Now that `ts.OffsetType` carries the tag, both arms match on the type instead of the node shape, like the `Dim + idx` arm already does. Matching on `tag=str()` rather than asserting it keeps an untagged `Dim + idx` offset out: `(Dim + 1)[0]` still gets the lowering error instead of an assertion. The `field(Off)` arm now requires two targets in its pattern, so a bare Cartesian `FieldOffset` (`field(Koff)`) also gets the lowering error instead of an assertion. That arm moves below `as_offset(...)`, whose type is its first argument's `OffsetType`. Taken from #2730, which carried the same tag under the name `name`. Co-authored-by: Hannes Vogt --- src/gt4py/next/ffront/foast_to_gtir.py | 53 +++++++------------ .../ffront_tests/test_import_from_mod.py | 31 ++++++++++- .../ffront_tests/test_foast_to_gtir.py | 26 ++++++++- 3 files changed, 73 insertions(+), 37 deletions(-) diff --git a/src/gt4py/next/ffront/foast_to_gtir.py b/src/gt4py/next/ffront/foast_to_gtir.py index aa369fac1b..b726f2f45c 100644 --- a/src/gt4py/next/ffront/foast_to_gtir.py +++ b/src/gt4py/next/ffront/foast_to_gtir.py @@ -66,26 +66,6 @@ def promote_to_list(node_type: ts.TypeSpec) -> Callable[[itir.Expr], itir.Expr]: return lambda x: x -def _offset_tag(offset_type: ts.OffsetType) -> str: - """ - Return the tag to emit for a `FieldOffset`-based shift, i.e. its offset-provider key. - - Covers both the unstructured form and a Cartesian `FieldOffset` subscript - (`a(Koff[1])`), which reaches the same branch. - - This used to be the name of the Python variable the `FieldOffset` was bound to, which - is not the offset's identity: a declaration `Off = FieldOffset("Tag", ...)` made - compiled backends look up `'Off'` while embedded execution looked up `'Tag'`, so the - same program needed a different offset provider depending on how it was run. - """ - assert offset_type.tag is not None, ( - f"Offset '{offset_type}' has no tag. Only a Cartesian shift written as" - " 'Dim + offset' is untagged, and that lowers to a 'CartesianOffset' carrying both" - " dimensions, without reaching this function." - ) - return offset_type.tag - - @dataclasses.dataclass class FieldOperatorLowering(eve.PreserveLocationVisitor, eve.NodeTranslator): """ @@ -315,16 +295,22 @@ 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): + # NOTE: matched on the type rather than on the node shape, so that a + # module-qualified offset (`field(mod.Off[idx])`) is lowered as well. The + # emitted shift is the offset's tag, i.e. its offset-provider key, and not the + # name of the Python variable the `FieldOffset` happens to be bound to. This + # also covers a Cartesian `FieldOffset` (`field(Koff[idx])`); only the + # untagged `Dim + idx` offsets fall through, to the error below. + 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_tag(offset_name.type), 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( @@ -344,15 +330,6 @@ def _visit_shift(self, node: foast.Call, **kwargs: Any) -> itir.Expr: ) ) )(current_expr) - # `field(Off)` - case foast.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 - assert isinstance(arg.type, ts.OffsetType) - return im.as_fieldop_neighbors( - _offset_tag(arg.type), self.visit(node.func, **kwargs) - ) # `field(as_offset(Off, offset_field))` case foast.Call(func=foast.Name(id="as_offset")): func_args = arg @@ -367,6 +344,14 @@ def _visit_shift(self, node: foast.Call, **kwargs: Any) -> itir.Expr: ) ) )(current_expr, offset_field) + # `field(Off)` + # NOTE: matched on the type, like `field(Off[idx])` above, and placed after + # `as_offset(...)`, whose type is the `OffsetType` of its first argument. + 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/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/unit_tests/ffront_tests/test_foast_to_gtir.py b/tests/next_tests/unit_tests/ffront_tests/test_foast_to_gtir.py index eb1927a776..183077d583 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 @@ -35,7 +35,7 @@ from gt4py.next.ffront.ast_passes import single_static_assign as ssa from gt4py.next.ffront.experimental import as_offset from gt4py.next.ffront.fbuiltins import exp, minimum -from gt4py.next.ffront.foast_to_gtir import FieldOperatorLowering +from gt4py.next.ffront.foast_to_gtir import FieldOperatorLowering, FieldOperatorLoweringError from gt4py.next.ffront.func_to_foast import FieldOperatorParser from gt4py.next.iterator.ir_utils import ir_makers as im from gt4py.next.iterator.transforms import inline_lambdas @@ -814,6 +814,30 @@ def foo(edge_f: gtx.Field[[Edge], float64]): assert lowered.expr == reference +def test_subscripted_untagged_offset_is_a_lowering_error(): + """Only a tagged offset lowers as `Off[idx]`; the untagged `Dim + idx` one has no key.""" + + def foo(inp: gtx.Field[[TDim], float64]): + return inp((TDim + 1)[0]) + + parsed = FieldOperatorParser.apply_to_function(foo) + + with pytest.raises(FieldOperatorLoweringError, match="Unexpected shift arguments"): + FieldOperatorLowering.apply(parsed) + + +def test_bare_cartesian_offset_is_a_lowering_error(): + """Only an offset with a local dimension lowers as `field(Off)`, i.e. to `neighbors`.""" + + def foo(inp: gtx.Field[[TDim], float64]): + return inp(TOff) + + parsed = FieldOperatorParser.apply_to_function(foo) + + with pytest.raises(FieldOperatorLoweringError, match="Unexpected shift arguments"): + FieldOperatorLowering.apply(parsed) + + def test_reduction_lowering_neighbor_sum(): def foo(edge_f: gtx.Field[[Edge], float64]): return neighbor_sum(edge_f(V2E), axis=V2EDim) From ef327f9e9aea39c2c845cc058bbf640ba4a955af Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Enrique=20Gonz=C3=A1lez=20Paredes?= Date: Fri, 18 Sep 2026 18:16:05 +0200 Subject: [PATCH 3/4] style[next]: trim comments on the offset tag --- src/gt4py/next/ffront/fbuiltins.py | 4 +--- src/gt4py/next/ffront/foast_passes/type_deduction.py | 2 -- src/gt4py/next/ffront/foast_to_gtir.py | 9 +-------- src/gt4py/next/type_system/type_specifications.py | 10 +--------- 4 files changed, 3 insertions(+), 22 deletions(-) diff --git a/src/gt4py/next/ffront/fbuiltins.py b/src/gt4py/next/ffront/fbuiltins.py index 6e592a7d08..2fd072a483 100644 --- a/src/gt4py/next/ffront/fbuiltins.py +++ b/src/gt4py/next/ffront/fbuiltins.py @@ -480,9 +480,7 @@ def impl( # guidelines for decision. @dataclasses.dataclass(frozen=True) class FieldOffset(runtime.Offset): - #: Narrows `runtime.Offset.value`, which is `int | str`. A `FieldOffset`'s value is its - #: tag: the key its connectivity has in the offset provider, and what lowering emits as - #: the shift. An integer tag has no meaning here, and every consumer treats it as a name. + #: The tag, i.e. the offset-provider key. value: str source: common.Dimension target: tuple[common.Dimension] | tuple[common.Dimension, common.Dimension] diff --git a/src/gt4py/next/ffront/foast_passes/type_deduction.py b/src/gt4py/next/ffront/foast_passes/type_deduction.py index 36db949bb8..2996cd6b2a 100644 --- a/src/gt4py/next/ffront/foast_passes/type_deduction.py +++ b/src/gt4py/next/ffront/foast_passes/type_deduction.py @@ -461,8 +461,6 @@ def visit_Subscript(self, node: foast.Subscript, **kwargs: Any) -> foast.Subscri raise errors.DSLError( new_value.location, "Second dimension in offset must be a local dimension." ) - # NOTE: `tag` is propagated: subscripting drops the local dimension but not - # the identity of the offset, which lowering needs to emit the shift. new_type = ts.OffsetType(source=source, target=(target1,), tag=tag) case ts.OffsetType(source=source, target=(target,)): # for cartesian axes (e.g. I, J) the index of the subscript only diff --git a/src/gt4py/next/ffront/foast_to_gtir.py b/src/gt4py/next/ffront/foast_to_gtir.py index b726f2f45c..9b3c15d7eb 100644 --- a/src/gt4py/next/ffront/foast_to_gtir.py +++ b/src/gt4py/next/ffront/foast_to_gtir.py @@ -295,12 +295,7 @@ def _visit_shift(self, node: foast.Call, **kwargs: Any) -> itir.Expr: for arg in node.args: match arg: # `field(Off[idx])` - # NOTE: matched on the type rather than on the node shape, so that a - # module-qualified offset (`field(mod.Off[idx])`) is lowered as well. The - # emitted shift is the offset's tag, i.e. its offset-provider key, and not the - # name of the Python variable the `FieldOffset` happens to be bound to. This - # also covers a Cartesian `FieldOffset` (`field(Koff[idx])`); only the - # untagged `Dim + idx` offsets fall through, to the error below. + # (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, @@ -345,8 +340,6 @@ def _visit_shift(self, node: foast.Call, **kwargs: Any) -> itir.Expr: ) )(current_expr, offset_field) # `field(Off)` - # NOTE: matched on the type, like `field(Off[idx])` above, and placed after - # `as_offset(...)`, whose type is the `OffsetType` of its first argument. 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. diff --git a/src/gt4py/next/type_system/type_specifications.py b/src/gt4py/next/type_system/type_specifications.py index 21efe3efe3..eb130bd994 100644 --- a/src/gt4py/next/type_system/type_specifications.py +++ b/src/gt4py/next/type_system/type_specifications.py @@ -74,18 +74,10 @@ class OffsetType(TypeSpec): # TODO(havogt): replace by ConnectivityType source: common.Dimension target: tuple[common.Dimension] | tuple[common.Dimension, common.Dimension] - #: The offset's own tag, i.e. the key its connectivity has in the offset provider. - #: - #: `None` for a Cartesian shift built from `Dim + offset`, which has no tag and needs - #: none: it lowers to a `CartesianOffset` carrying the two dimensions, with no lookup. - #: Set for every offset that comes from a `FieldOffset` declaration, so that lowering - #: does not have to fall back to the name of the Python variable the declaration - #: happens to be bound to. + #: The offset-provider key; `None` for the untagged Cartesian `Dim + offset`. tag: Optional[common.Tag] = None def __str__(self) -> str: - # NOTE: the tag is part of the identity, so two offsets over the same dimensions - # must not render identically in diagnostics. tag = "" if self.tag is None else f"{self.tag}: " return f"Offset[{tag}{self.source}, {self.target}]" From 25f7a32c4d3cbcf49a4d6f8b6152fd9bcc328b6d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Enrique=20Gonz=C3=A1lez=20Paredes?= Date: Wed, 23 Sep 2026 11:29:46 +0200 Subject: [PATCH 4/4] fix[next]: report invalid shift arguments as DSL errors `a(Koff)` (a Cartesian `FieldOffset` without an index) and `a((Dim + 1)[0])` (an indexed dimension shift) are programs a user can write, but they reached lowering and died there with `FieldOperatorLoweringError: Unexpected shift arguments!`, which carries no source location and no explanation. Reject both in type deduction with a located `DSLError` and a hint. The Cartesian-offset message quotes the offset as written (`Koff`, or `mod.Koff`), not its tag, since they can differ and the user's spelling is the one in their file. The lowering error stays as the internal guard it is. Note `as_offset(Off, f)` has the same type as a bare `Off` -- tagged, one target -- so the check matches a bare offset reference, not the type alone. Reported by @havogt in review of #2898. --- .../ffront/foast_passes/type_deduction.py | 30 +++++++++++++++++-- .../ffront_tests/test_diagnostic_messages.py | 27 +++++++++++++++++ .../ffront_tests/test_foast_to_gtir.py | 26 +--------------- 3 files changed, 56 insertions(+), 27 deletions(-) diff --git a/src/gt4py/next/ffront/foast_passes/type_deduction.py b/src/gt4py/next/ffront/foast_passes/type_deduction.py index 2996cd6b2a..738b105a61 100644 --- a/src/gt4py/next/ffront/foast_passes/type_deduction.py +++ b/src/gt4py/next/ffront/foast_passes/type_deduction.py @@ -462,7 +462,7 @@ def visit_Subscript(self, node: foast.Subscript, **kwargs: Any) -> foast.Subscri new_value.location, "Second dimension in offset must be a local dimension." ) new_type = ts.OffsetType(source=source, target=(target1,), tag=tag) - case ts.OffsetType(source=source, target=(target,)): + 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/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 183077d583..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 @@ -35,7 +35,7 @@ from gt4py.next.ffront.ast_passes import single_static_assign as ssa from gt4py.next.ffront.experimental import as_offset from gt4py.next.ffront.fbuiltins import exp, minimum -from gt4py.next.ffront.foast_to_gtir import FieldOperatorLowering, FieldOperatorLoweringError +from gt4py.next.ffront.foast_to_gtir import FieldOperatorLowering from gt4py.next.ffront.func_to_foast import FieldOperatorParser from gt4py.next.iterator.ir_utils import ir_makers as im from gt4py.next.iterator.transforms import inline_lambdas @@ -814,30 +814,6 @@ def foo(edge_f: gtx.Field[[Edge], float64]): assert lowered.expr == reference -def test_subscripted_untagged_offset_is_a_lowering_error(): - """Only a tagged offset lowers as `Off[idx]`; the untagged `Dim + idx` one has no key.""" - - def foo(inp: gtx.Field[[TDim], float64]): - return inp((TDim + 1)[0]) - - parsed = FieldOperatorParser.apply_to_function(foo) - - with pytest.raises(FieldOperatorLoweringError, match="Unexpected shift arguments"): - FieldOperatorLowering.apply(parsed) - - -def test_bare_cartesian_offset_is_a_lowering_error(): - """Only an offset with a local dimension lowers as `field(Off)`, i.e. to `neighbors`.""" - - def foo(inp: gtx.Field[[TDim], float64]): - return inp(TOff) - - parsed = FieldOperatorParser.apply_to_function(foo) - - with pytest.raises(FieldOperatorLoweringError, match="Unexpected shift arguments"): - FieldOperatorLowering.apply(parsed) - - def test_reduction_lowering_neighbor_sum(): def foo(edge_f: gtx.Field[[Edge], float64]): return neighbor_sum(edge_f(V2E), axis=V2EDim)