Skip to content

feat[next]: staggered fields#2667

Open
tehrengruber wants to merge 7 commits into
GridTools:mainfrom
tehrengruber:staggered_fields
Open

feat[next]: staggered fields#2667
tehrengruber wants to merge 7 commits into
GridTools:mainfrom
tehrengruber:staggered_fields

Conversation

@tehrengruber

@tehrengruber tehrengruber commented Jun 18, 2026

Copy link
Copy Markdown
Contributor

A staggered dimension is a dimension sitting at the half-integer positions of a base dimension. For example in a cell centered 2d cartesian grid with cells located at I, J the edges are located at I - ½, J and I, J - ½. Differentiating between staggered and non-staggered dimensions avoids accidental errors when fields defined on different entities are combined and improves readability for users by staying close to the mathematical formulations.

Let i_field for simplicity be a 1d field of cell values defined on I and IHalf the corresponding staggered dimension of edges between those cells then

  • i_field(IHalf + 0.5) → maps edges to the cell value above (i−½)+½ = i
  • i_field(IHalf − 0.5) → maps edges to the cell value below (i−½)−½ = i−1
           ┊      ●      ┊       ●      ┊      ●      ┊
 I               -1              0             1
 IHalf    -1½           -½             +½            +1½

Similary, let ihalf_field be defined on I - ½ then

  • ihalf_field(I + 0.5) → maps cells to the edge above i+½
  • ihalf_field(I − 0.5) → maps cells to the edge below i−½

Storage

ihalf_field.ndarray[idx] holds the value for logical index i = domain.start + idx, which is interpreted as sitting at i − ½. In the example above, IHalf(0) therefore sits at 0 − ½ = −½, i.e. it is the index of the edge just below the cell I(0). The memory layout is identical to any normal field; the "half" is purely how the index is interpreted geometrically.

Replaces #2339

@tehrengruber
tehrengruber marked this pull request as draft June 18, 2026 03:50
Comment thread src/gt4py/next/ffront/foast_to_gtir.py Outdated
Comment thread src/gt4py/next/iterator/ir_utils/domain_utils.py Outdated
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/type_system/type_synthesizer.py Outdated
Comment thread src/gt4py/next/iterator/embedded.py Outdated
Comment thread src/gt4py/next/program_processors/codegens/gtfn/gtfn_module.py Outdated
Comment thread src/gt4py/next/program_processors/codegens/gtfn/itir_to_gtfn_ir.py Outdated
Comment thread src/gt4py/next/common.py
Comment thread tests/next_tests/integration_tests/feature_tests/ffront_tests/test_basic.py Outdated
@tehrengruber
tehrengruber force-pushed the staggered_fields branch 3 times, most recently from 38714de to 8dd276b Compare June 22, 2026 12:11
Comment thread src/gt4py/next/iterator/ir_utils/domain_utils.py Outdated
@tehrengruber
tehrengruber requested a review from havogt July 1, 2026 14:37
@tehrengruber
tehrengruber marked this pull request as ready for review July 2, 2026 08:50
Comment thread src/gt4py/next/ffront/foast_passes/type_deduction.py Outdated
Comment thread src/gt4py/next/ffront/foast_to_gtir.py Outdated
Comment thread src/gt4py/next/ffront/foast_to_gtir.py Outdated
Comment thread src/gt4py/next/iterator/ir_utils/misc.py Outdated
Comment thread src/gt4py/next/iterator/ir_utils/domain_utils.py Outdated
Comment thread src/gt4py/next/program_processors/runners/gtfn.py
Comment thread src/gt4py/next/iterator/type_system/type_synthesizer.py
Comment thread src/gt4py/next/program_processors/codegens/gtfn/itir_to_gtfn_ir.py Outdated
Comment thread src/gt4py/next/program_processors/runners/gtfn.py
Comment thread src/gt4py/next/common.py
@tehrengruber
tehrengruber changed the base branch from main to tehrengruber-remove_cartesian_offset_subscript July 8, 2026 09:11
@tehrengruber
tehrengruber force-pushed the tehrengruber-remove_cartesian_offset_subscript branch from b6b4771 to 674d4fa Compare July 8, 2026 09:17
@tehrengruber
tehrengruber force-pushed the tehrengruber-remove_cartesian_offset_subscript branch from 674d4fa to 874f468 Compare July 8, 2026 09:21
Comment thread docs/development/ADRs/next/0024-Staggered_Dimension_Encoding.md Outdated
Comment thread docs/development/ADRs/next/0024-Staggered_Dimension_Encoding.md Outdated
@tehrengruber
tehrengruber requested a review from havogt July 8, 2026 13:47
@tehrengruber
tehrengruber changed the base branch from tehrengruber-remove_cartesian_offset_subscript to main July 8, 2026 13:51

@havogt havogt left a comment

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.

GLM 5.2: Review of the staggered-fields feature, run against the current head (df11778e2). Most of the earlier round (mixing restriction, _collect_dimensions_from_params block sizes, extract_connectivity_args device guard, field(Off[idx]) removal) is resolved in the rebased tree — thanks.

One item that was reported earlier and marked resolved (connectivity_for_cartesian_shift's bare assert half == 0.5, tehrengruber-ai "Done (c7ad10812)") does not appear at this head, so the fix looks to have been lost during the rebase. Details in the inline comment at common.py:1468.

Remaining items are mostly ADR completeness, error-message formatting, and a few test gaps for the new code paths. No other correctness issues found — the half-integer offset arithmetic is correct for all sign/direction combinations I checked, and the dace/gtfn/embedded backends handle relocation consistently.

Comment thread src/gt4py/next/common.py Outdated
Comment thread src/gt4py/next/common.py
Comment thread src/gt4py/next/ffront/foast_passes/type_deduction.py Outdated
Comment thread src/gt4py/next/ffront/foast_passes/type_deduction.py Outdated
Comment thread src/gt4py/next/iterator/ir_utils/domain_utils.py
Comment thread src/gt4py/next/program_processors/codegens/gtfn/itir_to_gtfn_ir.py Outdated
Comment thread tests/next_tests/integration_tests/feature_tests/ffront_tests/test_staggered.py Outdated
Comment thread docs/development/ADRs/next/0024-Staggered_Dimensions.md
"staggered" flag to `Dimension` would have required threading it through every one
of these representations and every place that constructs, compares, or lowers a
dimension — a large, invasive change touching the frontend, the IR, and all
backends. Carrying the marker in the name instead means it round-trips through all

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.

GLM 5.2: The "round-trips through all representations for free" claim isn't quite accurate — backends actively handle the prefix:

  • gtfn (itir_to_gtfn_ir.py:_add_staggered_aliases, ~line 200) turns every staggered tag into a C++ alias of its base dimension's TagDefinition, and visit_AxisLiteral (~line 401) emits node.type.dim.value instead of node.value.
  • nanobind (nanobind.py:212) uses common.as_non_staggered(dim).value so the generated binding references the base dimension name.

A contributor adding a new backend needs to know the prefix must be handled at code-generation boundaries. Suggest adding a short "Backend handling" subsection.

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.

@tehrengruber-ai Let's remove the common.as_non_staggered(dim) in the bindings unless there is a reason this is needed.

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

Addressed in 544232d: softened the 'round-trips for free' wording and added a 'gtfn backend' section describing the alias handling.

@tehrengruber
tehrengruber requested a review from havogt July 10, 2026 13:04

@havogt havogt left a comment

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.

Please check my questions. No need to do anything if you think current state makes sense.

@@ -0,0 +1,101 @@
---

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.

Index 24 and 25 is taken in the meantime.


- [0005 - Extending Iterator IR](0005-Extending_Iterator_IR.md)
- [0023 - Fingerprinting](0023-Fingerprinting.md)
- [0024 - Staggered Dimensions](0024-Staggered_Dimensions.md)

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.

Suggested change
- [0024 - Staggered Dimensions](0024-Staggered_Dimensions.md)
- [0026 - Staggered Dimensions](0026-Staggered_Dimensions.md)

case foast.BinOp(
op=dialect_ast_enums.BinaryOperator.ADD | dialect_ast_enums.BinaryOperator.SUB,
left=foast.Name() as dim_name,
left=foast.LocatedNode(type=ts.DimensionType(dim=common.Dimension() as dim)),

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.

LocatedNode makes sense?

"Domain",
"unit_range",
"UnitRange",
"is_staggered",

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.

Should is_staggered and as_non_staggered be part of the public API? Do you have a use-case?

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