refactor[next]: carry the connectivity name in OffsetType - #2730
Draft
havogt wants to merge 5 commits into
Draft
Conversation
`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`.
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.
havogt
marked this pull request as draft
July 29, 2026 11:44
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.
… 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`).
…ityType` TODO" This reverts commit 73bca31.
There was a problem hiding this comment.
Pull request overview
This PR refactors the gt4py.next frontend type system so unstructured shift lowering uses the connectivity’s declared name carried in ts.OffsetType, rather than relying on the source-level identifier shape (e.g., foast.Name vs foast.Attribute). This resolves mismatches when offsets are module-qualified or renamed on import, aligning compiled lowering behavior with the embedded path.
Changes:
- Add
name: Optional[common.Tag]tots.OffsetTypeand set it fromFieldOffset.valueduring__gt_type__. - Update
_visit_shiftlowering to pattern-match onts.OffsetType(name=...)(type-driven) for bothfield(Off)andfield(Off[idx]), and adjust match ordering aroundas_offset(...). - Extend integration test coverage to include renamed offsets for both unstructured neighbor shifts and sparse shifts across the backend matrix.
Reviewed changes
Copilot reviewed 5 out of 5 changed files in this pull request and generated no comments.
Show a summary per file
| File | Description |
|---|---|
| tests/next_tests/integration_tests/feature_tests/ffront_tests/test_import_from_mod.py | Adds regression tests for renamed/module-imported offsets in unstructured and sparse shift lowering. |
| src/gt4py/next/type_system/type_specifications.py | Extends OffsetType with an optional name to carry the connectivity tag used for offset-provider lookup. |
| src/gt4py/next/ffront/foast_to_gtir.py | Switches unstructured shift lowering to type-driven matching using OffsetType.name, removing node-shape coupling. |
| src/gt4py/next/ffront/foast_passes/type_deduction.py | Preserves OffsetType.name when slicing unstructured offsets (e.g., Off[idx]). |
| src/gt4py/next/ffront/fbuiltins.py | Ensures FieldOffset.__gt_type__ populates OffsetType.name from FieldOffset.value. |
2 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Stacked on #2724 — the first commit here is #2724's. Please review/merge that one first; this PR's own change is the second commit (
refactor[next]: carry the connectivity name in OffsetType).Description
Follow-up to #2724, which patched the symptom. This removes the cause.
The unstructured shift lowering took the ITIR offset tag from the source-level identifier (
foast.Name.id, and after #2724 alsofoast.Attribute.attr), sofield(Off)had to pattern-match the FOAST node shape, and resolved to the wrong tag whenever the identifier was not the connectivity name:Embedded resolves through
FieldOffset.value(fbuiltins.py:485-488) and was unaffected, so the two execution paths disagreed. This dates back to the first FOAST shift lowering (#625, 2022-01-26) and became observable when embedded remap landed (#1309).Change
ts.OffsetTypegainsname, set fromFieldOffset.valuein__gt_type__. Both shift arms in_visit_shiftnow match on the type rather than the node shape — mirroring the Cartesian arm, which already reads theDimensionout ofts.DimensionType. Thefoast.Attributespecial case from #2724 is dropped; it is subsumed.nameisNonefor theDim + idxoffsets synthesized in_deduce_binop_type. Those are Cartesian, resolved fromsource/targetalone, and never looked up in the offset provider.FieldOffset.__gt_type__always sets it, so CartesianFieldOffsets keep a name and thefield(Off[idx])arm is unaffected.The
field(Off)arm moved below theas_offsetarm:as_offset(...)propagates its argument'sOffsetType, so it would otherwise be caught by the now type-driven pattern. (It also has a single target, so thetarget=[_, _]guard would reject it — the move makes the ordering intentional rather than incidental.)This is a step toward the
# TODO(havogt): replace by ConnectivityTypealready onOffsetType, not the whole thing — I did not add an ADR since it does not change the public API shape.Not addressed
A separate name coupling is still live in the backends: a sparse field's LOCAL dimension is looked up in the offset provider by its
value, so the local dimension must be named exactly like its offset.codegens/gtfn/gtfn_module.py:93-96—dim_name = dim.value→common.get_offset_type(offset_provider_type, dim_name)runners/dace/lowering/gtir_to_sdfg.py:553— raisesValueError("The provided local dimension … does not match any offset provider type.")runners/dace/lowering/gtir_dataflow.py:733That is the direction of
common.py:977-979(TODO(havogt): refactor towards encoding this information in the local dimensions of the ConnectivityType.domain) and is untouched here.Tests
Added to
test_import_from_mod.py, across the backend matrix:test_import_renamed_offset_unstructured_shift—neighbor_sum(a(RenamedV2E), axis=V2EDim)test_import_renamed_offset_sparse_shift—a(RenamedE2V[0]), covering thefield(Off[idx])armBoth fail on #2724 and pass here.
Local runs (full backend matrix incl. GPU,
mypy src/andpre-commitclean):tests/next_tests/unit_tests— 2151 passed, 9 skipped, 15 xfailedtests/next_tests/integration_tests— 5316 passed, 197 skipped, 22 deselected, 612 xfailedThe 22 deselected are
test_orchestration.py::test_sdfgConvertible_connectivities[*-dace.run_dace_gpu]and its module siblings: that test fails identically with and without this change on my machine (TypeError: Illegal copy! (from gt_conn_E2V to tlet_0_deref)), verified by stashing and re-running on the base commit. Pre-existing, unrelated.Requirements