Skip to content

refactor[next]: carry the connectivity name in OffsetType - #2730

Draft
havogt wants to merge 5 commits into
GridTools:mainfrom
havogt:offset_type_name
Draft

refactor[next]: carry the connectivity name in OffsetType#2730
havogt wants to merge 5 commits into
GridTools:mainfrom
havogt:offset_type_name

Conversation

@havogt

@havogt havogt commented Jul 29, 2026

Copy link
Copy Markdown
Contributor

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 also foast.Attribute.attr), 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 next_tests.integration_tests.cases import V2E as RenamedV2E

@gtx.field_operator
def testee(a: cases.EField) -> cases.VField:
    return neighbor_sum(a(RenamedV2E), axis=V2EDim)
embedded: [1 3]
gtfn:     KeyError: "Offset 'RenamedV2E' not found in offset provider."

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.OffsetType gains name, set from FieldOffset.value in __gt_type__. Both shift arms in _visit_shift now match on the type rather than the node shape — mirroring the Cartesian arm, which already reads the Dimension out of ts.DimensionType. The foast.Attribute special case from #2724 is dropped; it is subsumed.

name is None for the Dim + idx offsets synthesized in _deduce_binop_type. Those are Cartesian, resolved from source/target alone, and never looked up in the offset provider. FieldOffset.__gt_type__ always sets it, so Cartesian FieldOffsets keep a name and the field(Off[idx]) arm is unaffected.

The field(Off) arm moved below the as_offset arm: as_offset(...) propagates its argument's OffsetType, so it would otherwise be caught by the now type-driven pattern. (It also has a single target, so the target=[_, _] guard would reject it — the move makes the ordering intentional rather than incidental.)

This is a step toward the # TODO(havogt): replace by ConnectivityType already on OffsetType, 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-96dim_name = dim.valuecommon.get_offset_type(offset_provider_type, dim_name)
  • runners/dace/lowering/gtir_to_sdfg.py:553 — raises ValueError("The provided local dimension … does not match any offset provider type.")
  • runners/dace/lowering/gtir_dataflow.py:733

That 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_shiftneighbor_sum(a(RenamedV2E), axis=V2EDim)
  • test_import_renamed_offset_sparse_shifta(RenamedE2V[0]), covering the field(Off[idx]) arm

Both fail on #2724 and pass here.

Local runs (full backend matrix incl. GPU, mypy src/ and pre-commit clean):

  • tests/next_tests/unit_tests — 2151 passed, 9 skipped, 15 xfailed
  • tests/next_tests/integration_tests — 5316 passed, 197 skipped, 22 deselected, 612 xfailed

The 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

  • All fixes and/or new features come with corresponding tests.
  • Important design decisions have been documented in the appropriate ADR.

havogt added 2 commits July 29, 2026 11:32
`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
havogt marked this pull request as draft July 29, 2026 11:44
havogt added 3 commits July 29, 2026 15:53
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`).

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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] to ts.OffsetType and set it from FieldOffset.value during __gt_type__.
  • Update _visit_shift lowering to pattern-match on ts.OffsetType(name=...) (type-driven) for both field(Off) and field(Off[idx]), and adjust match ordering around as_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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants