refactor[next]: Use cartesian shift syntax in all tests#2411
refactor[next]: Use cartesian shift syntax in all tests#2411tehrengruber wants to merge 14 commits into
Conversation
There was a problem hiding this comment.
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.CartesianOffsetplus type inference/synthesis, domain inference/translation, pretty-printing, and backend support (gtfn, DaCe, roundtrip). - Update ffront lowering so
field(Dim ± n)lowers toshift(cartesian_offset(Dim, Dim), n)instead of implicit string offsets. - Migrate tests to
field(Dim ± n)and remove mostIoff/Joff/Koffusage 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.
tehrengruber
left a comment
There was a problem hiding this comment.
Note for the discussion: The ITIR test changes are missing, feels like they should be in this PR.
|
|
||
|
|
||
| class CartesianOffset(Expr): | ||
| domain: AxisLiteral |
There was a problem hiding this comment.
Is this actually a good name?
There was a problem hiding this comment.
You mean it should be something like domain_dim, codomain_dim?
| # 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 |
There was a problem hiding this comment.
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.
Prerequisite for staggered fields (#2339), split out so that PR reduces to the staggering feature itself.
Changes
CartesianOffsetIR node.field(Dim ± n)now lowers toshift(cartesian_offset(Dim, Dim), n)— a node carrying the offset's source/target dimensions with theirDimensionKind— replacing the implicit string offset (_OffDim), which dropped the kind and broke verticalKDim ± non 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.field(Dim ± n)syntax; theIoff/Joff/KoffFieldOffsets and their offset-provider entries are removed, kept only where still required (dynamicas_offset, iterator-level string offsets).Why the node carries
domainandcodomainCartesianOffsetencodes a source→target dimension mapping. In this PR every cartesian shift is a translation (domain == codomain, e.g.IDim → IDim), enforced by explicitassert domain == codomainin the domain-inference and backend handlers. The second dimension exists for #2339's relocation shifts (domain != codomain, e.g.IDim → IHalfDimfor half-cell ±0.5 staggered shifts): keeping the node shape here lets #2339 add relocation by dropping those asserts and implementing thedomain != codomaincase — without reshaping the node or re-touching every handler.Out of scope (stays in #2339)
No
OffsetProviderElemretype (bareDimensionis still valid in offset providers), noConnectivityProtocol→concrete refactor, and no staggering (_Staggereddims / ±0.5 half shifts).