Skip to content

refactor[next]: Use cartesian shift syntax in all tests#2411

Open
tehrengruber wants to merge 14 commits into
GridTools:mainfrom
tehrengruber:plus_minus_cart_shift
Open

refactor[next]: Use cartesian shift syntax in all tests#2411
tehrengruber wants to merge 14 commits into
GridTools:mainfrom
tehrengruber:plus_minus_cart_shift

Conversation

@tehrengruber
Copy link
Copy Markdown
Contributor

@tehrengruber tehrengruber commented Dec 4, 2025

Prerequisite for staggered fields (#2339), split out so that PR reduces to the staggering feature itself.

Changes

  • New CartesianOffset IR node. field(Dim ± n) now lowers to shift(cartesian_offset(Dim, Dim), n) — a node carrying the offset's source/target dimensions with their DimensionKind — replacing the implicit string offset (_OffDim), which dropped the kind and broke vertical KDim ± n on the gtfn/dace/roundtrip pipelines. Handled in type inference/synthesis, domain inference, trace_shifts, pretty-print/parse, and the gtfn/dace/roundtrip backends (roundtrip cartesian-offset support was previously missing). The shift no longer needs an offset-provider entry — the node is self-describing.
  • Test migration. All tests use the field(Dim ± n) syntax; the Ioff/Joff/Koff FieldOffsets and their offset-provider entries are removed, kept only where still required (dynamic as_offset, iterator-level string offsets).

Why the node carries domain and codomain

CartesianOffset encodes a source→target dimension mapping. In this PR every cartesian shift is a translation (domain == codomain, e.g. IDim → IDim), enforced by explicit assert domain == codomain in the domain-inference and backend handlers. The second dimension exists for #2339's relocation shifts (domain != codomain, e.g. IDim → IHalfDim for half-cell ±0.5 staggered shifts): keeping the node shape here lets #2339 add relocation by dropping those asserts and implementing the domain != codomain case — without reshaping the node or re-touching every handler.

Out of scope (stays in #2339)

No OffsetProviderElem retype (bare Dimension is still valid in offset providers), no Connectivity Protocol→concrete refactor, and no staggering (_Staggered dims / ±0.5 half shifts).

@tehrengruber tehrengruber marked this pull request as draft December 4, 2025 15:57
Copy link
Copy Markdown

Copilot AI left a comment

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 introduces a dedicated CartesianOffset iterator-IR node to represent cartesian shifts (field(Dim ± n)) while preserving dimension kind (notably vertical KDim), and migrates the test suite to consistently use the cartesian shift syntax instead of Ioff/Joff/Koff FieldOffsets.

Changes:

  • Add itir.CartesianOffset plus type inference/synthesis, domain inference/translation, pretty-printing, and backend support (gtfn, DaCe, roundtrip).
  • Update ffront lowering so field(Dim ± n) lowers to shift(cartesian_offset(Dim, Dim), n) instead of implicit string offsets.
  • Migrate tests to field(Dim ± n) and remove most Ioff/Joff/Koff usage and offset-provider entries where no longer needed.

Reviewed changes

Copilot reviewed 27 out of 27 changed files in this pull request and generated 4 comments.

Show a summary per file
File Description
tests/next_tests/unit_tests/iterator_tests/test_type_inference.py Drop Ioff/Koff fixture imports where tests now rely on cartesian shift encoding.
tests/next_tests/unit_tests/ffront_tests/test_foast_to_gtir.py Update reference IR to expect CartesianOffset in lowered shift.
tests/next_tests/unit_tests/embedded_tests/test_basic_program.py Replace FieldOffset-based shift with IDim + 1 and remove offset_provider from call site.
tests/next_tests/integration_tests/multi_feature_tests/ffront_tests/test_laplacian.py Remove Joff import (tests now use cartesian shift syntax).
tests/next_tests/integration_tests/multi_feature_tests/ffront_tests/test_icon_like_scan.py Replace Koff[...] shifts with KDim ± n; stop providing Koff in offset_provider.
tests/next_tests/integration_tests/feature_tests/iterator_tests/test_program.py Keep iterator-level string offset usage (shift("Ioff", ...)) but remove FieldOffset declaration.
tests/next_tests/integration_tests/feature_tests/ffront_tests/test_where.py Replace Koff[-1] usage with KDim - 1.
tests/next_tests/integration_tests/feature_tests/ffront_tests/test_program.py Replace Ioff[1] with IDim + 1.
tests/next_tests/integration_tests/feature_tests/ffront_tests/test_gt4py_builtins.py Migrate shifts from Ioff/Joff/Koff to Dim ± n.
tests/next_tests/integration_tests/feature_tests/ffront_tests/test_execution.py Migrate shift tests to Dim ± n; keep explicit FieldOffsets only for dynamic as_offset paths.
tests/next_tests/integration_tests/feature_tests/ffront_tests/test_decorator.py Remove simple_cartesian_grid dependency; adjust offset-provider cache test input.
tests/next_tests/integration_tests/feature_tests/ffront_tests/ffront_test_utils.py Remove Ioff/Joff/Koff exports/definitions; make simple_cartesian_grid() offset_provider empty.
tests/next_tests/integration_tests/cases.py Remove Ioff/Joff/Koff imports from shared case definitions.
tests/next_tests/fixtures/past_common.py Remove Ioff/Joff fixture definitions as tests switch to cartesian shift syntax.
src/gt4py/next/program_processors/runners/roundtrip.py Add CartesianOffset rendering and ensure implicit offsets are declared in generated module.
src/gt4py/next/program_processors/runners/dace/lowering/gtir_dataflow.py Teach DaCe lowering to handle shift(CartesianOffset, ...) as cartesian shift with kind-correct dimension.
src/gt4py/next/program_processors/codegens/gtfn/itir_to_gtfn_ir.py Define tags for dimensions referenced by CartesianOffset; render CartesianOffset as dimension tag literal.
src/gt4py/next/iterator/type_system/type_synthesizer.py Extend dimension resolution and shift type synthesis to accept CartesianOffset.
src/gt4py/next/iterator/type_system/type_specifications.py Add CartesianOffsetType for type inference/synthesis.
src/gt4py/next/iterator/type_system/inference.py Add inference for itir.CartesianOffset nodes.
src/gt4py/next/iterator/transforms/trace_shifts.py Allow tracing shift(...) sequences that include CartesianOffset.
src/gt4py/next/iterator/pretty_printer.py Add pretty-printer support for CartesianOffset.
src/gt4py/next/iterator/ir.py Introduce CartesianOffset IR node and hashing support.
src/gt4py/next/iterator/ir_utils/misc.py Add helper to reconstruct common.Dimension (incl. kind) from AxisLiteral.
src/gt4py/next/iterator/ir_utils/ir_makers.py Add im.cartesian_offset(domain, codomain) helper.
src/gt4py/next/iterator/ir_utils/domain_utils.py Add domain translation support for cartesian shifts encoded via CartesianOffset.
src/gt4py/next/ffront/foast_to_gtir.py Lower field(Dim ± n) to shift(CartesianOffset(...), n) using dimension type info.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Comment thread src/gt4py/next/ffront/foast_to_gtir.py
Comment thread src/gt4py/next/program_processors/runners/roundtrip.py Outdated
Comment thread src/gt4py/next/program_processors/codegens/gtfn/itir_to_gtfn_ir.py
Comment thread tests/next_tests/unit_tests/iterator_tests/test_pretty_printer.py Outdated
Copy link
Copy Markdown

Copilot AI left a comment

Choose a reason for hiding this comment

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

Pull request overview

Copilot reviewed 32 out of 32 changed files in this pull request and generated 1 comment.

Comment thread src/gt4py/next/iterator/embedded.py
Copy link
Copy Markdown

Copilot AI left a comment

Choose a reason for hiding this comment

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

Pull request overview

Copilot reviewed 33 out of 33 changed files in this pull request and generated 1 comment.

Comment thread src/gt4py/next/iterator/embedded.py
@havogt havogt marked this pull request as ready for review May 22, 2026 14:57
Copy link
Copy Markdown

Copilot AI left a comment

Choose a reason for hiding this comment

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

Pull request overview

Copilot reviewed 33 out of 33 changed files in this pull request and generated 3 comments.

Comment thread src/gt4py/next/iterator/type_system/type_synthesizer.py Outdated
Comment thread src/gt4py/next/iterator/type_system/type_synthesizer.py Outdated
Comment thread src/gt4py/next/iterator/transforms/trace_shifts.py
Copy link
Copy Markdown

Copilot AI left a comment

Choose a reason for hiding this comment

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

Pull request overview

Copilot reviewed 34 out of 34 changed files in this pull request and generated 1 comment.

Comment thread src/gt4py/next/iterator/ir_utils/domain_utils.py
Copy link
Copy Markdown
Contributor Author

@tehrengruber tehrengruber left a comment

Choose a reason for hiding this comment

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

Note for the discussion: The ITIR test changes are missing, feels like they should be in this PR.

Comment thread src/gt4py/next/iterator/ir_utils/domain_utils.py
Comment thread src/gt4py/next/iterator/transforms/trace_shifts.py


class CartesianOffset(Expr):
domain: AxisLiteral
Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Is this actually a good name?

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

You mean it should be something like domain_dim, codomain_dim?

Comment thread src/gt4py/next/iterator/type_system/type_synthesizer.py
Comment thread src/gt4py/next/iterator/embedded.py Outdated
# A cartesian shift can be encoded by a `common.Dimension` tag (carrying its kind), shifting the
# position along that dimension directly; named offsets use a string `Tag` resolved via the
# offset provider.
OffsetPart: TypeAlias = Tag | common.Dimension | common.IntIndex
Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

In the IR we have CartesianOffset, but here we have Dimension. That feels inconsistent. Eventually this should be CartesianConnectivity. Maybe this is a hint that this change also belongs in the PR here. Otherwise it is not easy to asses if disallowing shift(ref_to_cart_conn, distance) is a good idea.

I did not review the rest of this file.

Comment thread src/gt4py/next/iterator/embedded.py
Copy link
Copy Markdown

Copilot AI left a comment

Choose a reason for hiding this comment

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

Pull request overview

Copilot reviewed 42 out of 42 changed files in this pull request and generated no new comments.

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