Skip to content

fix[next]: lower unstructured shifts with the offset's own tag - #2896

Closed
egparedes wants to merge 1 commit into
GridTools:mainfrom
egparedes:connectivities-as-types-1-shift-tag
Closed

egparedes wants to merge 1 commit into
GridTools:mainfrom
egparedes:connectivities-as-types-1-shift-tag

Conversation

@egparedes

Copy link
Copy Markdown
Contributor

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.

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.

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

nox-equivalent local runs: pytest tests/next_tests -m "not uses_dace" → 4477 passed / 0 failed; -m uses_dace → 1477 passed / 0 failed; mypy src/ clean; tach check clean; pre-commit run clean.

Context

First PR of a stack implementing egparedes/connectivities-as-types, an alternative to #2844. This PR stands alone — it is a bugfix that is correct regardless of whether the rest of the stack lands.

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 GridTools#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.
@egparedes

Copy link
Copy Markdown
Contributor Author

Superseded by #2898. This one was opened from a fork branch, so #2897 could not use it as a base — GitHub requires a stacked PR's base to be a branch in the target repository. Reopened from a GridTools/gt4py branch instead, matching #2844/#2845, so the stack is configured properly.

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

1 participant