From 24850c07acbf18f881a1960df703f9394a5a29e6 Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Wed, 29 Jul 2026 11:32:22 +0200 Subject: [PATCH 1/5] fix[next]: support module-qualified offsets in unstructured shifts `field(mod.Off)` raised `FieldOperatorLoweringError: Unexpected shift arguments!` because the `field(Off)` arm of `_visit_shift` matched `foast.Name` only. Extend it to `foast.Attribute`. `ts.OffsetType` carries no name, so the lowering takes the offset tag from the source-level identifier. The Cartesian arm does not have this problem: it matches on `ts.DimensionType`, which carries the `Dimension` itself, so it is node-shape agnostic already. Also add tests for module-qualified Cartesian and staggered shifts, which work but were uncovered. Known remaining gap: an offset renamed on import (`from mod import V2E as W2E`) still lowers to the wrong tag, since `attr`/`id` is the source-level name rather than `FieldOffset.value`. Embedded resolves via `FieldOffset.value` and disagrees. Fixing that means putting the name in `OffsetType`/`ConnectivityType`. --- src/gt4py/next/ffront/foast_to_gtir.py | 4 +- .../ffront_tests/test_import_from_mod.py | 50 ++++++++++++++++++- 2 files changed, 50 insertions(+), 4 deletions(-) diff --git a/src/gt4py/next/ffront/foast_to_gtir.py b/src/gt4py/next/ffront/foast_to_gtir.py index 10bc754526..0ae08053e8 100644 --- a/src/gt4py/next/ffront/foast_to_gtir.py +++ b/src/gt4py/next/ffront/foast_to_gtir.py @@ -322,8 +322,8 @@ def _visit_shift(self, node: foast.Call, **kwargs: Any) -> itir.Expr: ) ) )(current_expr) - # `field(Off)` - case foast.Name(id=offset_name): + # `field(Off)` or `field(mod.Off)` + case foast.Name(id=offset_name) | foast.Attribute(attr=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 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..666f53777c 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,17 @@ 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, + IHalfDim, + KDim, +) from next_tests.integration_tests.cases_utils import ( exec_alloc_descriptor, @@ -54,6 +60,46 @@ 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_cartesian_shift +def test_import_dims_module_cartesian_shift(cartesian_case): + @gtx.field_operator + def testee(a: cases.IField) -> cases.IField: + return a(cases.IDim + 1) + + size = cartesian_case.default_sizes[IDim] + a = cases.allocate(cartesian_case, testee, "a", domain={IDim: (1, size + 1)})() + out = cases.allocate(cartesian_case, testee, cases.RETURN, domain={IDim: (0, size)})() + + cases.verify(cartesian_case, testee, a, out=out, ref=a[:], offset_provider={}) + + +@pytest.mark.uses_cartesian_shift +def test_import_dims_module_staggered_shift(cartesian_case): + @gtx.field_operator + def testee(a: cases.IField) -> cases.IHalfField: + return a(cases.IHalfDim + 0.5) + + size = cartesian_case.default_sizes[IDim] + a = cases.allocate(cartesian_case, testee, "a", sizes={IDim: size})() + out = cases.allocate(cartesian_case, testee, cases.RETURN, sizes={IHalfDim: size})() + + cases.verify(cartesian_case, testee, a, out=out, ref=a, offset_provider={}) + + +@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), + ) + + # 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 From 8427bb87ad3081fbe9bb7f628b482614ce475b53 Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Wed, 29 Jul 2026 13:38:00 +0200 Subject: [PATCH 2/5] refactor[next]: carry the connectivity name in `OffsetType` The unstructured shift lowering took the offset tag from the source-level identifier, so `field(Off)` had to pattern-match the FOAST node shape and resolved to the wrong tag whenever the identifier was not the connectivity name: from mod import V2E as W2E neighbor_sum(a(W2E), axis=V2EDim) # compiled: KeyError 'W2E' Embedded resolves via `FieldOffset.value` and was unaffected, so the two execution paths disagreed. Add `name` to `ts.OffsetType`, set from `FieldOffset.value`. Both shift arms now match on the type instead of the node, mirroring the Cartesian arm, which already reads the `Dimension` from `ts.DimensionType`. `name` is `None` for the `Dim + idx` offsets synthesized in type deduction; those are resolved structurally and never looked up in the offset provider. --- src/gt4py/next/ffront/fbuiltins.py | 3 +- .../ffront/foast_passes/type_deduction.py | 4 +-- src/gt4py/next/ffront/foast_to_gtir.py | 24 ++++++++------- .../next/type_system/type_specifications.py | 3 ++ .../ffront_tests/test_import_from_mod.py | 29 +++++++++++++++++++ 5 files changed, 49 insertions(+), 14 deletions(-) diff --git a/src/gt4py/next/ffront/fbuiltins.py b/src/gt4py/next/ffront/fbuiltins.py index 37ddf9183a..93290b95c6 100644 --- a/src/gt4py/next/ffront/fbuiltins.py +++ b/src/gt4py/next/ffront/fbuiltins.py @@ -476,7 +476,8 @@ 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) + assert isinstance(self.value, str) + return ts.OffsetType(source=self.source, target=self.target, name=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 a96533f36a..c561eafd65 100644 --- a/src/gt4py/next/ffront/foast_passes/type_deduction.py +++ b/src/gt4py/next/ffront/foast_passes/type_deduction.py @@ -453,12 +453,12 @@ 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), name=name): 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,)) + new_type = ts.OffsetType(source=source, target=(target1,), name=name) 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 0ae08053e8..a74e876055 100644 --- a/src/gt4py/next/ffront/foast_to_gtir.py +++ b/src/gt4py/next/ffront/foast_to_gtir.py @@ -295,14 +295,16 @@ 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): + case foast.Subscript( + value=foast.LocatedNode(type=ts.OffsetType(name=str() as offset_name)), + 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_name, new_index)("__it"))) )(current_expr) # `field(Dim + idx)` (where `idx` is integer or half integer) case foast.BinOp( @@ -322,14 +324,6 @@ def _visit_shift(self, node: foast.Call, **kwargs: Any) -> itir.Expr: ) ) )(current_expr) - # `field(Off)` or `field(mod.Off)` - case foast.Name(id=offset_name) | foast.Attribute(attr=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 +338,14 @@ def _visit_shift(self, node: foast.Call, **kwargs: Any) -> itir.Expr: ) ) )(current_expr, offset_field) + # `field(Off)` + case foast.LocatedNode( + type=ts.OffsetType(name=str() as offset_name, 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_name, 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..32793e53ef 100644 --- a/src/gt4py/next/type_system/type_specifications.py +++ b/src/gt4py/next/type_system/type_specifications.py @@ -74,6 +74,9 @@ class OffsetType(TypeSpec): # TODO(havogt): replace by ConnectivityType source: common.Dimension target: tuple[common.Dimension] | tuple[common.Dimension, common.Dimension] + # Key of the connectivity in the offset provider. `None` for Cartesian shifts, which are + # resolved from `source`/`target` alone and never looked up in the offset provider. + name: Optional[str] = None def __str__(self) -> str: return f"Offset[{self.source}, {self.target}]" 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 666f53777c..67dcf93f8e 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 @@ -20,7 +20,9 @@ IDim, IHalfDim, KDim, + V2EDim, ) +from next_tests.integration_tests.cases import E2V as RenamedE2V, V2E as RenamedV2E from next_tests.integration_tests.cases_utils import ( exec_alloc_descriptor, @@ -100,6 +102,33 @@ def testee(a: cases.EField) -> cases.VField: ) +@pytest.mark.uses_unstructured_shift +def test_import_renamed_offset_unstructured_shift(unstructured_case): + @gtx.field_operator + def testee(a: cases.EField) -> cases.VField: + return neighbor_sum(a(RenamedV2E), axis=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_renamed_offset_sparse_shift(unstructured_case): + @gtx.field_operator + def testee(a: cases.VField) -> cases.EField: + return a(RenamedE2V[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 From 490afc0871697fd670d5ed5e6e86783a2880280d Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Wed, 29 Jul 2026 15:53:41 +0200 Subject: [PATCH 3/5] docs[next]: correct the `OffsetType.name` comment The comment claimed `name` is `None` for Cartesian shifts. A Cartesian `FieldOffset` goes through `__gt_type__` and does carry a name; the nameless case is the `OffsetType` synthesized in `_deduce_binop_type` for `Dim + idx`, regardless of it being Cartesian. Annotate as `common.Tag`, the key type of `OffsetProvider`, so the type states what the string has to match. --- src/gt4py/next/type_system/type_specifications.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/gt4py/next/type_system/type_specifications.py b/src/gt4py/next/type_system/type_specifications.py index 32793e53ef..4afe697f0d 100644 --- a/src/gt4py/next/type_system/type_specifications.py +++ b/src/gt4py/next/type_system/type_specifications.py @@ -74,9 +74,9 @@ class OffsetType(TypeSpec): # TODO(havogt): replace by ConnectivityType source: common.Dimension target: tuple[common.Dimension] | tuple[common.Dimension, common.Dimension] - # Key of the connectivity in the offset provider. `None` for Cartesian shifts, which are - # resolved from `source`/`target` alone and never looked up in the offset provider. - name: Optional[str] = None + #: `None` for the offsets synthesized from `Dim + idx` dimension arithmetic, which are + #: resolved from `source`/`target` alone rather than looked up in the offset provider. + name: Optional[common.Tag] = None def __str__(self) -> str: return f"Offset[{self.source}, {self.target}]" From 73bca3190d9da82a692bc87d434db26d5617d058 Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Wed, 29 Jul 2026 15:58:06 +0200 Subject: [PATCH 4/5] docs[next]: record what blocks the `OffsetType` -> `ConnectivityType` TODO `source`/`target` already map onto `ConnectivityType.codomain`/`domain`. What is missing is `skip_value`/`dtype`/`max_neighbors`, which are unknown until the offset provider is supplied at call time, and a home for `name` (the same gap as the TODO on `NeighborConnectivityType`). --- src/gt4py/next/type_system/type_specifications.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/gt4py/next/type_system/type_specifications.py b/src/gt4py/next/type_system/type_specifications.py index 4afe697f0d..58ad93d90f 100644 --- a/src/gt4py/next/type_system/type_specifications.py +++ b/src/gt4py/next/type_system/type_specifications.py @@ -71,7 +71,8 @@ def __str__(self) -> str: class OffsetType(TypeSpec): - # TODO(havogt): replace by ConnectivityType + # TODO(havogt): replace by ConnectivityType, once it can be built without an offset provider + # (`skip_value`/`dtype`/`max_neighbors`) and has a home for `name`. source: common.Dimension target: tuple[common.Dimension] | tuple[common.Dimension, common.Dimension] #: `None` for the offsets synthesized from `Dim + idx` dimension arithmetic, which are From b4bd07f0a7df04d4fd20c57cebba2d453f1d9c04 Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Wed, 29 Jul 2026 16:00:22 +0200 Subject: [PATCH 5/5] Revert "docs[next]: record what blocks the `OffsetType` -> `ConnectivityType` TODO" This reverts commit 73bca3190. --- src/gt4py/next/type_system/type_specifications.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/gt4py/next/type_system/type_specifications.py b/src/gt4py/next/type_system/type_specifications.py index 58ad93d90f..4afe697f0d 100644 --- a/src/gt4py/next/type_system/type_specifications.py +++ b/src/gt4py/next/type_system/type_specifications.py @@ -71,8 +71,7 @@ def __str__(self) -> str: class OffsetType(TypeSpec): - # TODO(havogt): replace by ConnectivityType, once it can be built without an offset provider - # (`skip_value`/`dtype`/`max_neighbors`) and has a home for `name`. + # TODO(havogt): replace by ConnectivityType source: common.Dimension target: tuple[common.Dimension] | tuple[common.Dimension, common.Dimension] #: `None` for the offsets synthesized from `Dim + idx` dimension arithmetic, which are