Skip to content

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

Closed
havogt wants to merge 6 commits into
GridTools:mainfrom
havogt:offset_type_name
Closed

havogt wants to merge 6 commits into
GridTools:mainfrom
havogt:offset_type_name

Conversation

@havogt

@havogt havogt commented Jul 29, 2026 •

Copy link
Copy Markdown
Contributor

Description

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-96 — dim_name = dim.value → common.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_shift — neighbor_sum(a(RenamedV2E), axis=V2EDim)
  • test_import_renamed_offset_sparse_shift — a(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.

egparedes added a commit that referenced this pull request Sep 18, 2026
`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 <vogt@hey.com>
egparedes added a commit that referenced this pull request Sep 23, 2026
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 —
confirmed by running it on v1.2.2:

```
MyOff = FieldOffset("TAGNAME", ...)
embedded:  {"TAGNAME": conn} OK ; {"MyOff": conn}   -> KeyError 'TAGNAME'
compiled:  {"MyOff": conn}   OK ; {"TAGNAME": conn} -> KeyError 'MyOff'
```

`ts.OffsetType` now carries `tag`, and lowering emits it.

Lowering also matches shift arguments on their **type** instead of their
node shape, so module-qualified offsets work. `a(mod.V2E)`,
`a(mod.E2V[0])` and `a(mod.Koff[1])` used to fail on every compiled
backend with `Unexpected shift arguments!` (a `foast.Attribute` is not a
`foast.Name`), while embedded ran them — the same embedded/compiled
divergence. Only tagged offsets are matched (`tag=str()`), and
`field(Off)` requires two targets, so an untagged `(Dim + 1)[0]` or a
bare Cartesian `a(Koff)` still gets the lowering error rather than an
assertion or a silently wrong `neighbors`.

This supersedes #2730 by @havogt, which carried the same fix under the
field name `name`; the type-driven matching and one of its tests are
taken from there (co-authored). `tag` is kept over `name` because it
matches `common.Tag` / `FieldOffset.value`, and `name` is easily
confused with the Python variable name — which is exactly what the bug
conflated.

### Why `tag` is `Optional`

A Cartesian shift written `Dim + offset` has no tag and needs none — it
lowers to an `itir.CartesianOffset` carrying both dimensions, with no
provider lookup. `type_deduction` builds an `OffsetType` for exactly
that case (`:711`), so a required field would break it. Subscripting
(`Off[1]`) drops the local dimension but *propagates* the tag, which is
the offset's identity.

### ⚠️ Behaviour change

For a declaration whose tag differs from the variable it is bound to,
compiled backends previously required `offset_provider={"MyOff": conn}`
and now require `{"TAGNAME": conn}`. That divergence from embedded
execution is the bug being fixed, but it is user-visible.

**ICON4Py is unaffected**: all 16 `FieldOffset` declarations have tag ==
variable name, there are no `Koff[...]` subscripts in model code, and
`as_offset` never consults the tag.

No `CHANGELOG.md` entry — that file is only ever touched by release PRs.

### Tests

The regression test grows from one cell (`a(Off[1])` on `GTFN_CPU`) to
`{shift, neighbor_sum} × {tag ≠ variable name, tag ≠ local dimension
name}` across the whole backend matrix, plus two lowering unit tests in
`test_foast_to_gtir.py` that assert the emitted `OffsetLiteral`
directly. Both unit tests were verified to **fail** with the fix
reverted.

For the type-driven matching: `test_import_from_mod.py` gains a
module-qualified `neighbor_sum(a(cases.V2E))` (from #2730) and
`a(cases.E2V[0])` across the backend matrix (20 compiled-backend
failures before the fix), and `test_foast_to_gtir.py` gains two
error-path tests (`(TDim + 1)[0]`, bare `inp(TOff)`), each verified to
fail with its guard removed.

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 never uses it.

The remaining failures are marked per backend, from measurement rather
than assumption:

| | roundtrip | roundtrip.gtir | gtfn | embedded | dace |
|---|---|---|---|---|---|
| shift, tag ≠ varname | ✅ | ✅ | ✅ | ✅ | ✅ |
| reduction, tag ≠ varname | ✅ | ✅ | ✅ | ✅ | ✅ |
| shift, tag ≠ local dim | ✅ | ✅ | ✅ | ✅ | xfail |
| reduction, tag ≠ local dim | ✅ | xfail | xfail | xfail | xfail |

Note `roundtrip` passes the reduction case while `roundtrip.gtir` does
not: `roundtrip.default` runs `apply_common_transforms`, so the
reduction unrolls keyed by the *offset tag*, whereas `roundtrip.gtir`
runs only the fieldview transforms and reaches `iterator/embedded.py`
keyed on the *local dimension*.

Both remaining constraints are one underlying issue: those paths resolve
a connectivity through the local dimension's name rather than the
offset's identity. Fixing it needs a back-pointer from the local
dimension to its connectivity, which is a later step in this stack.

GPU and JAX cells are marked by shared-code-path reasoning, not
measurement — they skip locally. `xfail_strict` is on, so if any is
wrong CI fails loudly rather than passing silently.

### Verification

At head (`d151516`): `pytest tests/next_tests -m "not uses_dace"` → 4495
passed / 0 failed; `pytest tests/next_tests/regression_tests -m
uses_dace` → 6 passed / 0 failed; `test_import_from_mod.py` on all CPU
backends incl. DaCe → 36 passed; `mypy src/`, `tach check`, `pre-commit
run` clean. The full `-m uses_dace` run (1477 passed / 0 failed) was
done on the first commit.

### Context

First PR of a stack implementing
[`egparedes/connectivities-as-types`](GridTools/gt4py_knowledge#32),
an alternative to #2844. **This PR stands alone** — it is a bugfix that
is correct regardless of whether the rest of the stack lands.

---------

Co-authored-by: Hannes Vogt <vogt@hey.com>
@egparedes

Copy link
Copy Markdown
Contributor

Close in favor of #2898

@egparedes egparedes closed this Sep 23, 2026
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.

3 participants