diff --git a/docs/development/ADRs/next/0026-Staggered_Dimensions.md b/docs/development/ADRs/next/0026-Staggered_Dimensions.md new file mode 100644 index 0000000000..77a8cff00d --- /dev/null +++ b/docs/development/ADRs/next/0026-Staggered_Dimensions.md @@ -0,0 +1,101 @@ +--- +tags: [] +--- + +# Staggered Dimensions + +- **Status**: valid +- **Authors**: Till Ehrengruber (@tehrengruber) +- **Created**: 2026-07-08 +- **Updated**: 2026-07-09 + +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 by staying close to the mathematical formulation. + +## Context + +Without a distinction between an entity and its staggered counterpart, fields +living on, e.g., cells and on edges share the same dimension and can be combined +silently even though they are geometrically different. We want the type system to +keep them apart, while still letting users move between the two with an intuitive, +mathematically-flavoured syntax. + +## Indexing convention + +A shift by a **half-integer** offset moves between a dimension and its staggered +counterpart; an integer offset shifts within the (non-staggered) dimension as +usual. + +Let `i_field` be a 1D field of cell values defined on `I`, and `IHalf` the +corresponding staggered dimension of the 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½ +``` + +Symmetrically, let `ihalf_field` be defined on `I − ½` (i.e. on `IHalf`). 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. This +index arithmetic is encoded in `common.connectivity_for_cartesian_shift`. + +## Encoding + +A staggered dimension is encoded as its base dimension's name with the internal +`_Staggered` prefix (`common._STAGGERED_PREFIX`), rather than as a new attribute +on `Dimension`. The helpers `is_staggered`, `flip_staggered` and +`as_non_staggered` operate purely on that prefix. + +Dimensions are identified by their **name** and appear throughout the toolchain +in more than one form: as a `common.Dimension` instance, but also as an +`itir.AxisLiteral` (and as plain strings in generated backend code). Adding a +"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 the frontend and IR pass +it through unchanged, without any of those formats having to learn about +staggering. The leading underscore marks the prefix as internal — user code goes +through the `is_staggered` / `flip_staggered` / `as_non_staggered` helpers +(exported from `gt4py.next`) rather than manipulating the name directly. + +## Dimension Constraints + +**A dimension and its staggered counterpart must not appear together** in the +same field or domain. This follows directly from the geometric meaning of +staggered dimension and is enforced by `check_dims` (`common.py`), which +raises otherwise. Other functions may rely on this, e.g `order_dims` just orders +by the `as_non_staggered` name. + +## gtfn backend + +In gtfn a staggered dimension is emitted as a C++ `using` alias of its base +dimension's tag (`_add_staggered_aliases` in `itir_to_gtfn_ir.py`): e.g. +`_StaggeredIDim_t` becomes an alias of `IDim_t`, and `visit_AxisLiteral` +correspondingly emits the base dimension name. + +The reason is that gtfn lowers a shift to an integer offset along a SID axis, and +a half-integer shift produces exactly such an offset (see +`connectivity_for_cartesian_shift`). But a staggered shift also *relocates* the +field onto the staggered dimension, a different axis identity. If the staggered +dimension had its own tag, reproducing that would require not just shifting the +SID but also *renaming* its axis from the base tag to the staggered tag — which +the gtfn shift primitive cannot do. Aliasing the two tags makes the base and its +staggered counterpart a single axis, so the relocation collapses into the plain +offset shift and no axis has to be renamed. diff --git a/docs/development/ADRs/next/README.md b/docs/development/ADRs/next/README.md index 9fe0695f24..a6d06da00e 100644 --- a/docs/development/ADRs/next/README.md +++ b/docs/development/ADRs/next/README.md @@ -18,6 +18,7 @@ Writing a new ADR is simple: - [0005 - Extending Iterator IR](0005-Extending_Iterator_IR.md) - [0023 - Fingerprinting](0023-Fingerprinting.md) +- [0026 - Staggered Dimensions](0024-Staggered_Dimensions.md) ### Frontend and Parsing #frontend diff --git a/src/gt4py/next/__init__.py b/src/gt4py/next/__init__.py index f93ef105b1..806265a1d1 100644 --- a/src/gt4py/next/__init__.py +++ b/src/gt4py/next/__init__.py @@ -31,7 +31,10 @@ Field, GridType, UnitRange, + as_non_staggered, domain, + flip_staggered, + is_staggered, unit_range, ) from .constructors import FieldConstructor, as_connectivity, as_field, empty, full, ones, zeros @@ -122,6 +125,9 @@ "Domain", "unit_range", "UnitRange", + "is_staggered", + "flip_staggered", + "as_non_staggered", # from constructors "FieldConstructor", "empty", diff --git a/src/gt4py/next/common.py b/src/gt4py/next/common.py index d71f0e7f4f..65b564c14f 100644 --- a/src/gt4py/next/common.py +++ b/src/gt4py/next/common.py @@ -86,10 +86,10 @@ def __str__(self) -> str: def __call__(self, val: int) -> NamedIndex: return NamedIndex(self, val) - def __add__(self, offset: int) -> Connectivity: - return CartesianConnectivity(self, offset) + def __add__(self, offset: int | float) -> Connectivity: + return connectivity_for_cartesian_shift(self, offset) - def __sub__(self, offset: int) -> Connectivity: + def __sub__(self, offset: int | float) -> Connectivity: return self + (-offset) def __gt__(self, value: core_defs.IntegralScalar) -> Domain: @@ -1336,11 +1336,29 @@ def order_dimensions(dims: Iterable[Dimension]) -> list[Dimension]: """Find the canonical ordering of the dimensions in `dims`.""" if sum(1 for dim in dims if dim.kind == DimensionKind.LOCAL) > 1: raise ValueError("There are more than one dimension with DimensionKind 'LOCAL'.") - return sorted(dims, key=lambda dim: (_DIM_KIND_ORDER[dim.kind], dim.value)) + return sorted( + dims, + key=lambda dim: ( + _DIM_KIND_ORDER[dim.kind], + as_non_staggered(dim).value, + ), + ) def check_dims(dims: Sequence[Dimension]) -> None: - if dims != order_dimensions(dims): + # A dimension and its staggered counterpart (i.e. sharing the same non-staggered base dimension) + # denote different grid locations and must not appear together in the same field/domain: mixing + # them is ambiguous (it makes `order_dimensions` non-total and produces duplicate backend tags). + seen: dict[Dimension, Dimension] = {} + for dim in dims: + base = as_non_staggered(dim) + if base in seen: + raise ValueError( + f"Dimensions '{seen[base]}' and '{dim}' cannot be combined: a dimension and its " + f"staggered counterpart must not appear together in the same field or domain." + ) + seen[base] = dim + if list(dims) != order_dimensions(dims): raise ValueError( f"Dimensions '{', '.join(map(str, dims))}' are not ordered correctly, expected '{', '.join(map(str, order_dimensions(dims)))}'." ) @@ -1424,3 +1442,48 @@ def __gt_builtin_func__(cls, /, func: fbuiltins.BuiltInFunction[_R, _P]) -> Call #: Equivalent to the `_FillValue` attribute in the UGRID Conventions #: (see: http://ugrid-conventions.github.io/ugrid-conventions/). _DEFAULT_SKIP_VALUE: Final[int] = -1 +_STAGGERED_PREFIX = "_Staggered" + + +def is_staggered(dim: Dimension) -> bool: + """Return whether `dim` is a staggered dimension.""" + return dim.value.startswith(_STAGGERED_PREFIX) + + +def flip_staggered(dim: Dimension) -> Dimension: + """Return the staggered counterpart of `dim`.""" + if is_staggered(dim): + return Dimension(dim.value[len(_STAGGERED_PREFIX) :], dim.kind) + else: + return Dimension(f"{_STAGGERED_PREFIX}{dim.value}", dim.kind) + + +def as_non_staggered(dim: Dimension) -> Dimension: + """Return the non-staggered base dimension of `dim` (`dim` itself if already non-staggered).""" + if is_staggered(dim): + return flip_staggered(dim) + return dim + + +def connectivity_for_cartesian_shift(dim: Dimension, offset: int | float) -> CartesianConnectivity: + """ + Build the connectivity that shifts `dim` by `offset`. + + An integer `offset` shifts within `dim` (the codomain stays `dim`). A half-integer `offset` + (fractional part `0.5`) shifts to the staggered counterpart of `dim` (the codomain becomes + `flip_staggered(dim)`). + + The half-integer case encodes the convention that a staggered index sits half a cell *below* + its base index (see ADR 0024): `IHalf(0)` is the edge below `I(0)`. Because of this asymmetry, + shifting out of a non-staggered dimension needs a `+1` index correction that shifting out of a + staggered dimension does not, e.g. `I + 0.5` maps `I(i)` to `IHalf(i+1)` (position `i+½`) while + `IHalf + 0.5` maps `IHalf(i)` to `I(i)`. + """ + integral_offset, staggered_offset = divmod(offset, 1) + if staggered_offset == 0.5: + if not is_staggered(dim): + integral_offset += 1 + return CartesianConnectivity(dim, int(integral_offset), codomain=flip_staggered(dim)) + else: + assert staggered_offset == 0 + return CartesianConnectivity(dim, int(integral_offset), codomain=dim) diff --git a/src/gt4py/next/ffront/foast_passes/type_deduction.py b/src/gt4py/next/ffront/foast_passes/type_deduction.py index 1b090489ff..a96533f36a 100644 --- a/src/gt4py/next/ffront/foast_passes/type_deduction.py +++ b/src/gt4py/next/ffront/foast_passes/type_deduction.py @@ -12,7 +12,7 @@ import gt4py.next.ffront.field_operator_ast as foast from gt4py import eve from gt4py.eve import NodeTranslator, NodeVisitor, traits -from gt4py.next import errors +from gt4py.next import common, errors from gt4py.next.common import Dimension, DimensionKind, promote_dims from gt4py.next.ffront import ( dialect_ast_enums, @@ -607,13 +607,6 @@ def _deduce_compare_type( def _deduce_binop_type( self, node: foast.BinOp, *, left: foast.Expr, right: foast.Expr, **kwargs: Any ) -> Optional[ts.TypeSpec]: - # e.g. `IDim+1` - if ( - isinstance(left.type, ts.DimensionType) - and isinstance(right.type, ts.ScalarType) - and type_info.is_integral(right.type) - ): - return ts.OffsetType(source=left.type.dim, target=(left.type.dim,)) if isinstance(left.type, ts.OffsetType): raise errors.DSLError( node.location, f"Type '{left.type}' can not be used in operator '{node.op}'." @@ -679,6 +672,34 @@ def _deduce_binop_type( f"must be one of {', '.join((str(op) for op in logical_ops))}.", ) return ts.DomainType(dims=promote_dims(left.type.dims, right.type.dims)) + elif ( + node.op in (dialect_ast_enums.BinaryOperator.ADD, dialect_ast_enums.BinaryOperator.SUB) + and isinstance(left.type, ts.DimensionType) + and isinstance(right.type, ts.ScalarType) + and type_info.is_arithmetic(right.type) + ): + # e.g. `IDim+1` or `IDim+0.5` + if not isinstance(right, foast.Constant): + raise errors.DSLError( + right.location, + "Cartesian offsets are only supported with a literal right-hand side, " + "e.g. 'IDim + 1', but not 'IDim + expr'.", + ) + offset_index = right.value + if node.op == dialect_ast_enums.BinaryOperator.SUB: + offset_index *= -1 + if not isinstance(offset_index, (int, float)) or offset_index % 1 not in (0, 0.5): + raise errors.DSLError( + right.location, + f"Invalid offset '{right.value}' for a Cartesian shift of dimension " + f"'{left.type.dim.value}'.", + hints=[ + "Use an integer offset to shift within the dimension, or a half-integer " + "offset such as '0.5' or '-1.5' to shift to the staggered dimension." + ], + ) + conn = common.connectivity_for_cartesian_shift(left.type.dim, offset_index) + return ts.OffsetType(source=conn.codomain, target=(conn.domain_dim,)) else: raise errors.DSLError(node.location, err_msg) diff --git a/src/gt4py/next/ffront/foast_to_gtir.py b/src/gt4py/next/ffront/foast_to_gtir.py index 3d5b4d0dbd..4df2714fe0 100644 --- a/src/gt4py/next/ffront/foast_to_gtir.py +++ b/src/gt4py/next/ffront/foast_to_gtir.py @@ -12,7 +12,7 @@ from gt4py import eve from gt4py.eve.extended_typing import Never, cast -from gt4py.next import utils +from gt4py.next import common, utils from gt4py.next.ffront import ( dialect_ast_enums, experimental as experimental_builtins, @@ -304,19 +304,22 @@ def _visit_shift(self, node: foast.Call, **kwargs: Any) -> itir.Expr: current_expr = im.as_fieldop( im.lambda_("__it")(im.deref(im.shift(offset_name.id, new_index)("__it"))) )(current_expr) - # `field(Dim + idx)` + # `field(Dim + idx)` (where `idx` is integer or half integer) 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)), right=foast.Constant(value=offset_index), ): if arg.op == dialect_ast_enums.BinaryOperator.SUB: offset_index *= -1 - assert isinstance(dim_name.type, ts.DimensionType) - dim = dim_name.type.dim + conn = common.connectivity_for_cartesian_shift(dim, offset_index) current_expr = im.as_fieldop( im.lambda_("__it")( - im.deref(im.shift(im.cartesian_offset(dim), offset_index)("__it")) + im.deref( + im.shift( + im.cartesian_offset(conn.domain_dim, conn.codomain), conn.offset + )("__it") + ) ) )(current_expr) # `field(Off)` diff --git a/src/gt4py/next/iterator/embedded.py b/src/gt4py/next/iterator/embedded.py index e01ecc181e..76d25c1770 100644 --- a/src/gt4py/next/iterator/embedded.py +++ b/src/gt4py/next/iterator/embedded.py @@ -587,16 +587,11 @@ def execute_shift( # the assertions above confirm pos is incomplete casting here to avoid duplicating work in a type guard return cast(IncompletePosition, pos) | {tag: new_entry} - # a `CartesianConnectivity` tag is a self-describing cartesian shift (e.g. from a - # `CartesianOffset` IR node); named offsets are resolved through the offset provider if isinstance(tag, common.CartesianConnectivity): - assert tag.domain_dim == tag.codomain # relocation (staggering) is not supported here new_pos = copy.copy(pos) - key = tag.domain_dim.value - if common.is_int_index(value := new_pos[key]): - new_pos[key] = value + index + tag.offset - else: - raise AssertionError() + value = new_pos.pop(tag.domain_dim.value) + assert common.is_int_index(value) + new_pos[tag.codomain.value] = value + index + tag.offset return new_pos offset_implementation = common.get_offset(offset_provider, tag) if common.is_neighbor_table(offset_implementation): diff --git a/src/gt4py/next/iterator/ir.py b/src/gt4py/next/iterator/ir.py index 5a3fc1713f..3deff7618c 100644 --- a/src/gt4py/next/iterator/ir.py +++ b/src/gt4py/next/iterator/ir.py @@ -88,6 +88,7 @@ def __str__(self): InfinityLiteral.POSITIVE = InfinityLiteral(name="POSITIVE") +# TODO(tehrengruber): allow int only and create OffsetRef for str instead class OffsetLiteral(Expr): value: Union[int, str] diff --git a/src/gt4py/next/iterator/ir_utils/domain_utils.py b/src/gt4py/next/iterator/ir_utils/domain_utils.py index 305e69c1eb..983f9d0d33 100644 --- a/src/gt4py/next/iterator/ir_utils/domain_utils.py +++ b/src/gt4py/next/iterator/ir_utils/domain_utils.py @@ -65,14 +65,12 @@ def _unstructured_translate_range_statically( tag: str, val: itir.OffsetLiteral | Literal[trace_shifts.Sentinel.VALUE, trace_shifts.Sentinel.ALL_NEIGHBORS], - offset_provider: common.OffsetProvider, + connectivity: common.NeighborTable, expr: itir.Expr | None = None, ) -> SymbolicRange: """ Translate `range_` using static connectivity information from `offset_provider`. """ - assert common.is_offset_provider(offset_provider) - connectivity = offset_provider[tag] assert common.is_neighbor_table(connectivity) skip_value = connectivity.skip_value @@ -181,36 +179,48 @@ def translate( #: func:`gt4py.next.iterator.transforms.infer_domain.infer_expr` for more details. symbolic_domain_sizes: Optional[dict[str, itir.Expr]] = None, ) -> SymbolicDomain: - offset_provider_type = common.offset_provider_to_type(offset_provider) - dims = list(self.ranges.keys()) new_ranges = {dim: self.ranges[dim] for dim in dims} if len(shift) == 0: return self if len(shift) == 2: off, val = shift + if isinstance(off, itir.CartesianOffset): if val is trace_shifts.Sentinel.VALUE: - raise NotImplementedError("Dynamic offsets not supported.") + raise NotImplementedError("Dynamic cartesian offsets not supported.") assert isinstance(val, itir.OffsetLiteral) and isinstance(val.value, int) - dom = misc.dim_from_axis_literal(off.domain) - cod = misc.dim_from_axis_literal(off.codomain) - assert dom == cod # relocation (staggering) is not supported here - new_ranges[dom] = SymbolicRange.translate(self.ranges[dom], val.value) - return SymbolicDomain(self.grid_type, new_ranges) - assert isinstance(off, itir.OffsetLiteral) and isinstance(off.value, str) - connectivity_type = common.get_offset_type(offset_provider_type, off.value) - - if isinstance(connectivity_type, common.NeighborConnectivityType): - # unstructured shift + + old_dim = misc.dim_from_axis_literal(off.domain) + new_dim = misc.dim_from_axis_literal(off.codomain) + + assert new_dim not in new_ranges or old_dim == new_dim + + new_range = SymbolicRange.translate(self.ranges[old_dim], val.value) + new_ranges = dict( + (dim, range_) if dim != old_dim else (new_dim, new_range) + for dim, range_ in new_ranges.items() + ) + else: + assert isinstance(off, itir.OffsetLiteral) and isinstance(off.value, str) assert ( isinstance(val, itir.OffsetLiteral) and isinstance(val.value, int) ) or val in [ trace_shifts.Sentinel.ALL_NEIGHBORS, trace_shifts.Sentinel.VALUE, ] - old_dim = connectivity_type.source_dim - new_dim = connectivity_type.codomain + + connectivity: common.NeighborTable | common.NeighborConnectivityType + if common.is_offset_provider(offset_provider): + connectivity = common.get_offset(offset_provider, off.value) + old_dim = connectivity.domain.dims[0] + new_dim = connectivity.codomain + else: + assert common.is_offset_provider_type(offset_provider) + connectivity = common.get_offset_type(offset_provider, off.value) + old_dim = connectivity.domain[0] + new_dim = connectivity.codomain + assert new_dim not in new_ranges or old_dim == new_dim if symbolic_domain_sizes is not None and new_dim.value in symbolic_domain_sizes: new_range = SymbolicRange( @@ -218,18 +228,20 @@ def translate( im.ensure_expr(symbolic_domain_sizes[new_dim.value]), ) else: - assert common.is_offset_provider(offset_provider) - assert not isinstance(val, itir.CartesianOffset) # offset value, never a node + assert common.is_neighbor_table(connectivity) new_range = _unstructured_translate_range_statically( - new_ranges[old_dim], off.value, val, offset_provider, self.as_expr() + new_ranges[old_dim], + off.value, + val, # type: ignore[arg-type] # mypy not smart enough + connectivity, + self.as_expr(), ) new_ranges = dict( (dim, range_) if dim != old_dim else (new_dim, new_range) for dim, range_ in new_ranges.items() ) - else: - raise AssertionError() + return SymbolicDomain(self.grid_type, new_ranges) elif len(shift) > 2: return self.translate(shift[0:2], offset_provider, symbolic_domain_sizes).translate( diff --git a/src/gt4py/next/iterator/type_system/type_synthesizer.py b/src/gt4py/next/iterator/type_system/type_synthesizer.py index e204377bc0..6587eb437e 100644 --- a/src/gt4py/next/iterator/type_system/type_synthesizer.py +++ b/src/gt4py/next/iterator/type_system/type_synthesizer.py @@ -16,8 +16,8 @@ from gt4py.eve import utils as eve_utils from gt4py.eve.extended_typing import Callable, Iterable, Optional, Union from gt4py.next import common, utils -from gt4py.next.ffront import fbuiltins from gt4py.next.iterator import builtins, ir as itir +from gt4py.next.iterator.ir_utils import misc as ir_misc from gt4py.next.iterator.type_system import type_specifications as it_ts from gt4py.next.type_system import type_info, type_specifications as ts from gt4py.next.utils import tree_map @@ -482,14 +482,25 @@ def _resolve_dimensions( >>> Edge = common.Dimension(value="Edge") >>> Vertex = common.Dimension(value="Vertex") + >>> Cell = common.Dimension(value="Cell") >>> K = common.Dimension(value="K", kind=common.DimensionKind.VERTICAL) >>> V2E = common.Dimension(value="V2E") + >>> C2V = common.Dimension(value="C2V") >>> input_dims = [Edge, K] >>> shift_tuple = ( + ... itir.OffsetLiteral(value="C2V"), + ... itir.OffsetLiteral(value=0), ... itir.OffsetLiteral(value="V2E"), ... itir.OffsetLiteral(value=0), ... ) >>> offset_provider_type = { + ... "C2V": common.NeighborConnectivityType( + ... domain=(Cell, C2V), + ... codomain=Vertex, + ... skip_value=None, + ... dtype=None, + ... max_neighbors=3, + ... ), ... "V2E": common.NeighborConnectivityType( ... domain=(Vertex, V2E), ... codomain=Edge, @@ -497,26 +508,54 @@ def _resolve_dimensions( ... dtype=None, ... max_neighbors=4, ... ), - ... "KOff": K, ... } >>> _resolve_dimensions(input_dims, shift_tuple, offset_provider_type) - [Dimension(value='Vertex', kind=), Dimension(value='K', kind=)] + [Dimension(value='Cell', kind=), Dimension(value='K', kind=)] + >>> from gt4py.next.iterator.ir_utils import ir_makers as im + >>> IDim = common.Dimension(value="IDim") + >>> IHalfDim = common.flip_staggered(IDim) + >>> JDim = common.Dimension(value="JDim") + >>> JHalfDim = common.flip_staggered(JDim) + >>> input_dims = [IDim, JDim] + >>> shift_tuple = ( + ... itir.CartesianOffset( + ... domain=im.axis_literal(IDim), codomain=im.axis_literal(IHalfDim) + ... ), + ... itir.OffsetLiteral(value=0), + ... itir.CartesianOffset(domain=im.axis_literal(JDim), codomain=im.axis_literal(IDim)), + ... itir.OffsetLiteral(value=0), + ... itir.CartesianOffset( + ... domain=im.axis_literal(IHalfDim), codomain=im.axis_literal(JDim) + ... ), + ... itir.OffsetLiteral(value=0), + ... ) + >>> _resolve_dimensions(input_dims, shift_tuple, offset_provider_type) + [Dimension(value='JDim', kind=), Dimension(value='IDim', kind=)] + """ resolved_dims = [] for input_dim in input_dims: - for off_literal in reversed( - shift_tuple[::2] - ): # Only OffsetLiterals are processed, located at even indices in shift_tuple. Shifts are applied in reverse order: the last shift in the tuple is applied first. + resolved_dim = input_dim + for off_literal in reversed(shift_tuple[::2]): + # Only OffsetLiterals/CartesianOffsets are processed, located at even indices in + # shift_tuple. Shifts are applied in reverse order: the last shift in the tuple is + # applied first. if isinstance(off_literal, itir.CartesianOffset): - if off_literal.domain != off_literal.codomain: - raise NotImplementedError("Relocation (staggering) is not supported.") - continue # translation does not change the dimension - assert isinstance(off_literal.value, str) - offset_type = common.get_offset_type(offset_provider_type, off_literal.value) - if isinstance(offset_type, (fbuiltins.FieldOffset, common.NeighborConnectivityType)): - if input_dim == offset_type.codomain: # Check if input fits to offset - input_dim = offset_type.domain[0] # Update input_dim for next iteration - resolved_dims.append(input_dim) + if resolved_dim == ir_misc.dim_from_axis_literal(off_literal.codomain): + resolved_dim = ir_misc.dim_from_axis_literal(off_literal.domain) + else: + assert isinstance(off_literal, itir.OffsetLiteral) and isinstance( + off_literal.value, str + ) + offset_type = common.get_offset_type(offset_provider_type, off_literal.value) + if isinstance(offset_type, common.NeighborConnectivityType): + if resolved_dim == offset_type.codomain: # Check if input fits to offset + resolved_dim = offset_type.domain[0] # Update input_dim for next iteration + else: + raise NotImplementedError( + f"'{offset_type}' is not a supported connectivity type." + ) + resolved_dims.append(resolved_dim) return resolved_dims @@ -660,24 +699,24 @@ def apply_shift( new_position_dims = [*it.position_dims] assert len(offset_literals) % 2 == 0 for offset_axis, _ in zip(offset_literals[:-1:2], offset_literals[1::2], strict=True): + source_dim: common.Dimension + target_dim: common.Dimension if isinstance(offset_axis, it_ts.CartesianOffsetType): - if offset_axis.domain != offset_axis.codomain: - raise NotImplementedError("Relocation (staggering) is not supported.") - continue # translation leaves position dims unchanged - assert isinstance(offset_axis, it_ts.OffsetLiteralType) and isinstance( - offset_axis.value, str - ) - type_ = common.get_offset_type(offset_provider_type, offset_axis.value) - if isinstance(type_, common.NeighborConnectivityType): - found = False - for i, dim in enumerate(new_position_dims): - if dim.value == type_.source_dim.value: - assert not found - new_position_dims[i] = type_.codomain - found = True - assert found + source_dim, target_dim = offset_axis.domain, offset_axis.codomain else: - raise NotImplementedError(f"{type_} is not a supported Connectivity type.") + assert isinstance(offset_axis, it_ts.OffsetLiteralType) + assert isinstance(offset_axis.value, str) + type_ = common.get_offset_type(offset_provider_type, offset_axis.value) + assert isinstance(type_, common.NeighborConnectivityType) + source_dim, target_dim = type_.domain[0], type_.codomain + + found = False + for i, dim in enumerate(new_position_dims): + if dim == source_dim: + assert not found + new_position_dims[i] = target_dim + found = True + assert found else: # during re-inference we don't have an offset provider type new_position_dims = "unknown" diff --git a/src/gt4py/next/program_processors/codegens/gtfn/codegen.py b/src/gt4py/next/program_processors/codegens/gtfn/codegen.py index 5f1cfcc516..13348b4555 100644 --- a/src/gt4py/next/program_processors/codegens/gtfn/codegen.py +++ b/src/gt4py/next/program_processors/codegens/gtfn/codegen.py @@ -315,14 +315,13 @@ def visit_Program(self, node: gtfn_ir.Program, **kwargs: Any) -> Union[str, Coll def _block_sizes(self, offset_definitions: list[gtfn_ir.TagDefinition]) -> str: if self.is_cartesian: - block_dims = [] - block_sizes = [32, 8] + [1] * (len(offset_definitions) - 2) - for i, tag in enumerate(offset_definitions): - if tag.alias is None: - block_dims.append( - f"gridtools::meta::list<{tag.name.id}_t, " - f"gridtools::integral_constant>" - ) + dims = [tag for tag in offset_definitions if tag.alias is None] + block_sizes = [32, 8] + [1] * (len(dims) - 2) + block_dims = [ + f"gridtools::meta::list<{tag.name.id}_t, " + f"gridtools::integral_constant>" + for tag, block_size in zip(dims, block_sizes) + ] sizes_str = ",\n".join(block_dims) return f"using block_sizes_t = gridtools::meta::list<{sizes_str}>;" else: diff --git a/src/gt4py/next/program_processors/codegens/gtfn/itir_to_gtfn_ir.py b/src/gt4py/next/program_processors/codegens/gtfn/itir_to_gtfn_ir.py index cc9dbd1d12..01b403fa18 100644 --- a/src/gt4py/next/program_processors/codegens/gtfn/itir_to_gtfn_ir.py +++ b/src/gt4py/next/program_processors/codegens/gtfn/itir_to_gtfn_ir.py @@ -151,36 +151,28 @@ def _collect_offset_definitions( grid_type: common.GridType, offset_provider_type: common.OffsetProviderType, ) -> dict[str, TagDefinition]: - used_offset_tags: set[str] = ( - node.walk_values() - .if_isinstance(itir.OffsetLiteral) - .filter(lambda offset_literal: isinstance(offset_literal.value, str)) - .getattr("value") - ).to_set() - # implicit offsets don't occur in the `offset_provider_type`, get them from the used offset tags - offset_provider_type = { - offset_name: common.get_offset_type(offset_provider_type, offset_name) - for offset_name in used_offset_tags - } | {**offset_provider_type} offset_definitions = {} + offset_provider_type = {**offset_provider_type} - # cartesian shifts (`field(Dim + n)`) are encoded as `CartesianOffset` nodes and don't - # occur in the `offset_provider_type`; define a tag for each of their dimensions cartesian_offsets: set[itir.CartesianOffset] = ( node.walk_values().if_isinstance(itir.CartesianOffset) ).to_set() for cart_offset in cartesian_offsets: - for axis in (cart_offset.domain, cart_offset.codomain): + dims = [ + ir_utils_misc.dim_from_axis_literal(v) + for v in (cart_offset.domain, cart_offset.codomain) + ] + for dim in dims: if grid_type == common.GridType.CARTESIAN: - offset_definitions[axis.value] = TagDefinition(name=Sym(id=axis.value)) + offset_definitions[dim.value] = TagDefinition(name=Sym(id=dim.value)) else: assert grid_type == common.GridType.UNSTRUCTURED - if axis.kind != common.DimensionKind.VERTICAL: + if dim.kind != common.DimensionKind.VERTICAL: raise ValueError( "Mapping an offset to a horizontal dimension in unstructured is not allowed." ) - offset_definitions[axis.value] = TagDefinition( - name=Sym(id=axis.value), alias=_vertical_dimension + offset_definitions[dim.value] = TagDefinition( + name=Sym(id=dim.value), alias=_vertical_dimension ) for offset_name, connectivity_type in offset_provider_type.items(): @@ -205,6 +197,23 @@ def _collect_offset_definitions( return offset_definitions +def _add_staggered_aliases( + offset_definitions: dict[str, TagDefinition], +) -> dict[str, TagDefinition]: + """Turn every staggered dimension tag into a alias of its base dimension.""" + result: dict[str, TagDefinition] = {} + aliases: dict[str, TagDefinition] = {} + for name, tag_def in offset_definitions.items(): + if tag_def.alias is None and common.is_staggered(common.Dimension(value=name)): + base_name = common.as_non_staggered(common.Dimension(value=name)).value + # ensure the base tag exists (as alias target and loop dimension) in this position + result.setdefault(base_name, TagDefinition(name=Sym(id=base_name))) + aliases[name] = TagDefinition(name=Sym(id=name), alias=SymRef(id=base_name)) + else: + result[name] = tag_def + return {**result, **aliases} + + def _literal_as_integral_constant(node: itir.Literal) -> IntegralConstant: assert type_info.is_integral_scalar(node.type) return IntegralConstant(value=int(node.value)) @@ -387,12 +396,11 @@ def visit_OffsetLiteral(self, node: itir.OffsetLiteral, **kwargs: Any) -> Offset return OffsetLiteral(value=node.value) def visit_CartesianOffset(self, node: itir.CartesianOffset, **kwargs: Any) -> Literal: - # render as the (shared) dimension tag - assert node.domain == node.codomain, "relocation (staggering) is not supported" return self.visit(node.codomain, **kwargs) def visit_AxisLiteral(self, node: itir.AxisLiteral, **kwargs: Any) -> Literal: - return Literal(value=node.value, type="axis_literal") + assert isinstance(node.type, ts.DimensionType) + return Literal(value=node.type.dim.value, type="axis_literal") def _make_domain(self, node: itir.FunCall) -> tuple[TaggedValues, TaggedValues]: tags = [] @@ -689,6 +697,7 @@ def visit_Program(self, node: itir.Program, **kwargs: Any) -> Program: **_collect_dimensions_from_domain(node.body), **_collect_offset_definitions(node, self.grid_type, self.offset_provider_type), } + offset_definitions = _add_staggered_aliases(offset_definitions) return Program( id=SymbolName(node.id), params=self.visit(node.params), diff --git a/src/gt4py/next/program_processors/runners/dace/lowering/gtir_dataflow.py b/src/gt4py/next/program_processors/runners/dace/lowering/gtir_dataflow.py index 0b755d9f03..bdb7bb20a2 100644 --- a/src/gt4py/next/program_processors/runners/dace/lowering/gtir_dataflow.py +++ b/src/gt4py/next/program_processors/runners/dace/lowering/gtir_dataflow.py @@ -1497,12 +1497,13 @@ def _visit_shift_multidim( return offset_provider_arg, offset_value_arg, it def _make_cartesian_shift( - self, it: IteratorExpr, offset_dim: gtx_common.Dimension, offset_expr: DataExpr + self, it: IteratorExpr, offset: gtir.CartesianOffset, offset_expr: DataExpr ) -> IteratorExpr: """Implements cartesian shift along one dimension.""" - assert any(dim == offset_dim for dim, _ in it.field_domain) + old_dim = itir_misc.dim_from_axis_literal(offset.domain) + new_dim = itir_misc.dim_from_axis_literal(offset.codomain) new_index: SymbolExpr | ValueExpr - index_expr = it.indices[offset_dim] + index_expr = it.indices[old_dim] if isinstance(index_expr, SymbolExpr) and isinstance(offset_expr, SymbolExpr): # purely symbolic expression which can be interpreted at compile time new_index = SymbolExpr( @@ -1565,9 +1566,10 @@ def _make_cartesian_shift( ) # a new iterator with a shifted index along one dimension - shifted_indices = { - dim: (new_index if dim == offset_dim else index) for dim, index in it.indices.items() - } + shifted_indices = dict( + (new_dim, new_index) if dim == old_dim else (dim, index) + for dim, index in it.indices.items() + ) return IteratorExpr(it.field, it.gt_dtype, it.field_domain, shifted_indices) def _make_dynamic_neighbor_offset( @@ -1673,31 +1675,26 @@ def _visit_shift(self, node: gtir.FunCall) -> IteratorExpr: ) if isinstance(offset_provider_arg, gtir.CartesianOffset): - # cartesian shift; the dimension (incl. kind) is encoded in the node - assert offset_provider_arg.domain == offset_provider_arg.codomain, ( - "relocation (staggering) is not supported" + return self._make_cartesian_shift(it, offset_provider_arg, offset_expr) + else: + assert isinstance(offset_provider_arg, gtir.OffsetLiteral) + assert isinstance(offset_provider_arg.value, str) + offset_provider_type = self.subgraph_builder.get_offset_provider_type( + offset_provider_arg.value + ) + assert isinstance(offset_provider_type, gtx_common.NeighborConnectivityType) + # a named offset → unstructured shift; the offset value may be a static + # `OffsetLiteral` or a dynamic offset (handled by `_make_unstructured_shift`). + # initially, the storage for the connectivity tables is created as transient; + # when the tables are used, the storage is changed to non-transient, + # so the corresponding arrays are supposed to be allocated by the SDFG caller + offset_table = gtx_dace_args.connectivity_identifier(offset_provider_arg.value) + self.sdfg.arrays[offset_table].transient = False + offset_table_node = self.state.add_access(offset_table) + + return self._make_unstructured_shift( + it, offset_provider_type, offset_table_node, offset_expr ) - offset_dim = itir_misc.dim_from_axis_literal(offset_provider_arg.codomain) - return self._make_cartesian_shift(it, offset_dim, offset_expr) - - # first argument of the shift node is the offset provider - assert isinstance(offset_provider_arg, gtir.OffsetLiteral) - offset = offset_provider_arg.value - assert isinstance(offset, str) - offset_provider_type = self.subgraph_builder.get_offset_provider_type(offset) - - # reaching here means a named offset → unstructured shift (cartesian shifts took the - # CartesianOffset branch above) - # initially, the storage for the connectivity tables is created as transient; - # when the tables are used, the storage is changed to non-transient, - # so the corresponding arrays are supposed to be allocated by the SDFG caller - offset_table = gtx_dace_args.connectivity_identifier(offset) - self.sdfg.arrays[offset_table].transient = False - offset_table_node = self.state.add_access(offset_table) - - return self._make_unstructured_shift( - it, offset_provider_type, offset_table_node, offset_expr - ) def _visit_generic_builtin(self, node: gtir.FunCall) -> ValueExpr: """ diff --git a/src/gt4py/next/program_processors/runners/gtfn.py b/src/gt4py/next/program_processors/runners/gtfn.py index 1ee1860c5d..88547c0627 100644 --- a/src/gt4py/next/program_processors/runners/gtfn.py +++ b/src/gt4py/next/program_processors/runners/gtfn.py @@ -90,16 +90,16 @@ def extract_connectivity_args( # This is currently true only because when hashing offset provider dicts, # the keys' order is taken into account. Any modification to the hashing # of offset providers may break this assumption here. - args: list[tuple[core_defs.NDArrayObject, tuple[int, ...]]] = [ - (ndarray, zero_origin) - for conn in offset_provider.values() - if (ndarray := getattr(conn, "ndarray", None)) is not None - ] assert all( common.is_neighbor_table(conn) and field_utils.verify_device_field_type(conn, device) for conn in offset_provider.values() if hasattr(conn, "ndarray") ) + args: list[tuple[core_defs.NDArrayObject, tuple[int, ...]]] = [ + (conn.ndarray, zero_origin) + for conn in offset_provider.values() + if common.is_neighbor_table(conn) + ] return args diff --git a/tests/next_tests/integration_tests/cases.py b/tests/next_tests/integration_tests/cases.py index d2efee1fc5..08fb817856 100644 --- a/tests/next_tests/integration_tests/cases.py +++ b/tests/next_tests/integration_tests/cases.py @@ -58,7 +58,9 @@ E2VDim, Edge, IDim, + IHalfDim, JDim, + JHalfDim, KDim, KHalfDim, V2EDim, @@ -74,6 +76,7 @@ # mypy does not accept [IDim, ...] as a type IField: TypeAlias = gtx.Field[[IDim], np.int32] # type: ignore [valid-type] +IHalfField: TypeAlias = gtx.Field[[IHalfDim], np.int32] # type: ignore [valid-type] JField: TypeAlias = gtx.Field[[JDim], np.int32] # type: ignore [valid-type] IFloatField: TypeAlias = gtx.Field[[IDim], np.float64] # type: ignore [valid-type] IBoolField: TypeAlias = gtx.Field[[IDim], bool] # type: ignore [valid-type] @@ -501,6 +504,7 @@ def verify_with_default_data( case: Case, fieldop: decorator.FieldOperator, ref: Callable, + offset_provider: Optional[OffsetProvider] = None, comparison: Callable[[Any, Any], bool] = tree_mapped_np_allclose, ) -> None: """ @@ -515,6 +519,8 @@ def verify_with_default_data( fieldview_prog: The field operator or program to be verified. ref: A callable which will be called with all the input arguments of the fieldview code, after applying ``.ndarray`` on the fields. + offset_provider: An override for the test case's offset_provider. + Use with care! comparison: A comparison function, which will be called as ``comparison(ref, )`` and should return a boolean. """ @@ -528,7 +534,7 @@ def verify_with_default_data( *inps, **kwfields, ref=ref(*ref_args), - offset_provider=case.offset_provider, + offset_provider=offset_provider, comparison=comparison, ) @@ -731,7 +737,9 @@ def from_cartesian_grid_descriptor( IDim: grid_descriptor.sizes[0], JDim: grid_descriptor.sizes[1], KDim: grid_descriptor.sizes[2], - KHalfDim: grid_descriptor.sizes[3], + IHalfDim: grid_descriptor.sizes[0] - 1, + JHalfDim: grid_descriptor.sizes[1] - 1, + KHalfDim: grid_descriptor.sizes[2] - 1, }, grid_type=common.GridType.CARTESIAN, allocator=allocator, diff --git a/tests/next_tests/integration_tests/cases_utils.py b/tests/next_tests/integration_tests/cases_utils.py index 3f7b979def..d29a0b6d90 100644 --- a/tests/next_tests/integration_tests/cases_utils.py +++ b/tests/next_tests/integration_tests/cases_utils.py @@ -154,9 +154,11 @@ def debug_itir(tree): DType = TypeVar("DType") IDim = gtx.Dimension("IDim") +IHalfDim = common.flip_staggered(IDim) JDim = gtx.Dimension("JDim") +JHalfDim = common.flip_staggered(JDim) KDim = gtx.Dimension("KDim", kind=gtx.DimensionKind.VERTICAL) -KHalfDim = gtx.Dimension("KHalf", kind=gtx.DimensionKind.VERTICAL) +KHalfDim = common.flip_staggered(KDim) Vertex = gtx.Dimension("Vertex") Edge = gtx.Dimension("Edge") @@ -190,11 +192,11 @@ def offset_provider_type(self) -> common.OffsetProviderType: ... def simple_cartesian_grid( - sizes: int | tuple[int, int, int, int] = (5, 7, 9, 11), + sizes: int | tuple[int, int, int, int] = (5, 7, 9), ) -> CartesianGridDescriptor: if isinstance(sizes, int): - sizes = (sizes,) * 4 - assert len(sizes) == 4, "sizes must be a tuple of four integers" + sizes = (sizes,) * 3 + assert len(sizes) == 3, "sizes must be a tuple of three integers" offset_provider = {} diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_staggered.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_staggered.py new file mode 100644 index 0000000000..d79bf4adda --- /dev/null +++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_staggered.py @@ -0,0 +1,142 @@ +# GT4Py - GridTools Framework +# +# Copyright (c) 2014-2024, ETH Zurich +# All rights reserved. +# +# Please, refer to the LICENSE file in the root directory. +# SPDX-License-Identifier: BSD-3-Clause +import numpy as np +import pytest + +import gt4py.next as gtx + +from next_tests.integration_tests import cases +from next_tests.integration_tests.cases import ( + IDim, + IHalfDim, + JDim, + KDim, + KHalfDim, + cartesian_case, +) +from next_tests.integration_tests.cases_utils import exec_alloc_descriptor + + +@pytest.mark.uses_cartesian_shift +def test_copy_half_field(cartesian_case): + @gtx.field_operator + def testee(a: cases.IHalfField) -> cases.IHalfField: + field_tuple = (a, a) + field_0 = field_tuple[0] + field_1 = field_tuple[1] + return field_0 + + cases.verify_with_default_data(cartesian_case, testee, ref=lambda a: a, offset_provider={}) + + +@pytest.mark.uses_cartesian_shift +def test_cartesian_shift_plus(cartesian_case): + @gtx.field_operator + def testee(a: cases.IField) -> cases.IField: + return a(IDim + 1) # always pass an I-index to an IField + + size = cartesian_case.default_sizes[IDim] + a = cases.allocate(cartesian_case, testee, "a", domain={IDim: (1, size + 1)})() + out = cases.allocate(cartesian_case, testee, cases.RETURN, domain={IDim: (0, size)})() + + cases.verify(cartesian_case, testee, a, out=out, ref=a[:], offset_provider={}) + + +@pytest.mark.uses_cartesian_shift +def test_cartesian_half_shift_plus(cartesian_case): + @gtx.field_operator + def testee(a: cases.IField) -> cases.IHalfField: + return a(IHalfDim + 0.5) # always pass an I-index to an IField + + size = cartesian_case.default_sizes[IDim] + a = cases.allocate(cartesian_case, testee, "a", sizes={IDim: size})() + out = cases.allocate(cartesian_case, testee, cases.RETURN, sizes={IHalfDim: size})() + + cases.verify(cartesian_case, testee, a, out=out, ref=a, offset_provider={}) + + +@pytest.mark.uses_cartesian_shift +def test_cartesian_half_shift_back(cartesian_case): + @gtx.field_operator + def testee(a: cases.IHalfField) -> cases.IHalfField: + return a(IDim + 0.5)(IHalfDim - 0.5) # always pass an I-index to an IField + + a = cases.allocate(cartesian_case, testee, "a")() + out = cases.allocate(cartesian_case, testee, cases.RETURN)() + + cases.verify(cartesian_case, testee, a, out=out, ref=a, offset_provider={}) + + +@pytest.mark.uses_cartesian_shift +def test_cartesian_half_shift_plus1(cartesian_case): + @gtx.field_operator + def testee(a: cases.IHalfField) -> cases.IHalfField: + return a(IHalfDim + 1) # always pass an IHalf-index to an IHalfField + + size = cartesian_case.default_sizes[IDim] + a = cases.allocate(cartesian_case, testee, "a", domain={IHalfDim: (1, size + 1)})() + out = cases.allocate(cartesian_case, testee, cases.RETURN, domain={IHalfDim: (0, size)})() + + cases.verify(cartesian_case, testee, a, out=out[:], ref=a[:], offset_provider={}) + + +@pytest.mark.uses_cartesian_shift +def test_cartesian_half_shift_minus(cartesian_case): + @gtx.field_operator + def testee(a: cases.IField) -> cases.IHalfField: + return a(IHalfDim - 0.5) # always pass an I-index to an IField + + size = cartesian_case.default_sizes[IDim] + a = cases.allocate(cartesian_case, testee, "a", domain={IDim: (-1, size - 1)})() + out = cases.allocate(cartesian_case, testee, cases.RETURN, domain={IHalfDim: (0, size)})() + + cases.verify(cartesian_case, testee, a, out=out, ref=a[:], offset_provider={}) + + +@pytest.mark.uses_cartesian_shift +def test_cartesian_half_shift_half2center(cartesian_case): + @gtx.field_operator + def testee(a: cases.IHalfField) -> cases.IField: + return 2 * a(IDim + 0.5) # always pass an IHalf-index to an IHalfField + + size = cartesian_case.default_sizes[IDim] + a = cases.allocate(cartesian_case, testee, "a", domain={IHalfDim: (1, size + 1)})() + out = cases.allocate(cartesian_case, testee, cases.RETURN, sizes={IDim: size})() + + cases.verify(cartesian_case, testee, a, out=out, ref=2 * a[:], offset_provider={}) + + +@pytest.mark.uses_cartesian_shift +def test_cartesian_half_shift_vertical(cartesian_case): + # vertical (K) staggering: identical mechanism, different dimension kind. + @gtx.field_operator + def testee(a: cases.KField) -> gtx.Field[[KHalfDim], np.int32]: + return a(KHalfDim + 0.5) + + size = cartesian_case.default_sizes[KDim] + a = cases.allocate(cartesian_case, testee, "a", sizes={KDim: size})() + out = cases.allocate(cartesian_case, testee, cases.RETURN, sizes={KHalfDim: size})() + + cases.verify(cartesian_case, testee, a, out=out, ref=a, offset_provider={}) + + +@pytest.mark.uses_cartesian_shift +def test_cartesian_half_shift_multi_dim(cartesian_case): + # staggering one axis of a multi-dimensional field leaves the other axis untouched. + @gtx.field_operator + def testee(a: cases.IJField) -> gtx.Field[[IHalfDim, JDim], np.int32]: + return a(IHalfDim + 0.5) + + isize = cartesian_case.default_sizes[IDim] + jsize = cartesian_case.default_sizes[JDim] + a = cases.allocate(cartesian_case, testee, "a", sizes={IDim: isize, JDim: jsize})() + out = cases.allocate( + cartesian_case, testee, cases.RETURN, sizes={IHalfDim: isize, JDim: jsize} + )() + + cases.verify(cartesian_case, testee, a, out=out, ref=a, offset_provider={}) diff --git a/tests/next_tests/integration_tests/multi_feature_tests/ffront_tests/test_multiple_output_domains.py b/tests/next_tests/integration_tests/multi_feature_tests/ffront_tests/test_multiple_output_domains.py index d4b27e78ee..0ddb42a521 100644 --- a/tests/next_tests/integration_tests/multi_feature_tests/ffront_tests/test_multiple_output_domains.py +++ b/tests/next_tests/integration_tests/multi_feature_tests/ffront_tests/test_multiple_output_domains.py @@ -17,6 +17,7 @@ IDim, JDim, KDim, + KHalfDim, C2E, E2V, V2E, @@ -37,7 +38,7 @@ mesh_descriptor, ) -KHalfDim = gtx.Dimension("KHalf", kind=gtx.DimensionKind.VERTICAL) +pytestmark = pytest.mark.uses_cartesian_shift @gtx.field_operator diff --git a/tests/next_tests/unit_tests/ffront_tests/test_diagnostic_messages.py b/tests/next_tests/unit_tests/ffront_tests/test_diagnostic_messages.py index 70573d280e..d0bd52f360 100644 --- a/tests/next_tests/unit_tests/ffront_tests/test_diagnostic_messages.py +++ b/tests/next_tests/unit_tests/ffront_tests/test_diagnostic_messages.py @@ -127,6 +127,19 @@ def with_and(a: gtx.Field[[IDim], bool], b: gtx.Field[[IDim], bool]) -> gtx.Fiel assert any("'&' and '|'" in hint for hint in err.hints) +def test_invalid_cartesian_offset_suggests_valid_offsets(): + def foo(a: gtx.Field[[IDim], float64]) -> gtx.Field[[IDim], float64]: + return a(IDim + 0.25) + + err = parse_error(foo) + + assert err.message == "Invalid offset '0.25' for a Cartesian shift of dimension 'IDim'." + assert any("half-integer offset" in hint for hint in err.hints) + rendered = str(err) + assert "return a(IDim + 0.25)" in rendered + assert "Hint:" in rendered + + def test_add_note_uses_pep678_notes(): # 'add_note' uses the standard PEP 678 mechanism ('__notes__'); the # structured 'notes' field is reserved for content authored at the raise diff --git a/tests/next_tests/unit_tests/ffront_tests/test_foast_to_gtir.py b/tests/next_tests/unit_tests/ffront_tests/test_foast_to_gtir.py index e354d878a1..0d1805fd6f 100644 --- a/tests/next_tests/unit_tests/ffront_tests/test_foast_to_gtir.py +++ b/tests/next_tests/unit_tests/ffront_tests/test_foast_to_gtir.py @@ -21,6 +21,7 @@ from gt4py.next import ( astype, broadcast, + errors, float32, float64, int32, @@ -110,18 +111,38 @@ def foo(inp1: gtx.Field[[TDim], float64], inp2: gtx.Field[[TDim], float64]): assert lowered.expr == reference +def test_premap_cartesian_invalid_fractional_offset(): + # only integer or half-integer offsets are valid; expect a located DSLError. + def foo(inp: gtx.Field[[TDim], float64]): + return inp(TDim + 0.25) + + with pytest.raises(errors.DSLError, match="Invalid offset"): + FieldOperatorParser.apply_to_function(foo) + + +def test_premap_cartesian_non_literal_offset(): + def foo(inp: gtx.Field[[TDim], float64], i: int): + return inp(TDim + i) + + with pytest.raises(errors.DSLError, match="literal right-hand side"): + FieldOperatorParser.apply_to_function(foo) + + def test_premap_cartesian_syntax(): + # an integer offset shifts within the dimension; '1.0' is equivalent to '1'. def foo(inp: gtx.Field[[TDim], float64]): return inp(TDim + 1) - parsed = FieldOperatorParser.apply_to_function(foo) - lowered = FieldOperatorLowering.apply(parsed) + def foo_float(inp: gtx.Field[[TDim], float64]): + return inp(TDim + 1.0) reference = im.as_fieldop( im.lambda_("__it")(im.deref(im.shift(im.cartesian_offset(TDim, TDim), 1)("__it"))) )("inp") - assert lowered.expr == reference + for f in (foo, foo_float): + lowered = FieldOperatorLowering.apply(FieldOperatorParser.apply_to_function(f)) + assert lowered.expr == reference def test_as_offset(): diff --git a/tests/next_tests/unit_tests/iterator_tests/ir_utils_tests/test_domain_utils.py b/tests/next_tests/unit_tests/iterator_tests/ir_utils_tests/test_domain_utils.py index cf56c6593e..4aac031406 100644 --- a/tests/next_tests/unit_tests/iterator_tests/ir_utils_tests/test_domain_utils.py +++ b/tests/next_tests/unit_tests/iterator_tests/ir_utils_tests/test_domain_utils.py @@ -15,6 +15,7 @@ from gt4py.next import common, constructors I = common.Dimension("I") +IHalf = common.flip_staggered(I) J = common.Dimension("J") K = common.Dimension("J", kind=common.DimensionKind.VERTICAL) Vertex = common.Dimension("Vertex") @@ -292,6 +293,47 @@ def test_unstructured_translate(shift_chain, expected_end_domain): assert end_domain == expected_end_domain +@pytest.mark.parametrize("as_type", [False, True]) +def test_unstructured_translate_with_symbolic_domain_sizes(as_type): + # With `symbolic_domain_sizes` the translated range is taken from the provided size + # expression instead of the connectivity table. This makes `translate` work for a type-only + # `OffsetProviderType` (which has no table) as well as a runtime `OffsetProvider`. + offset_provider = { + "V2E": constructors.as_connectivity( + domain={Vertex: (0, 4), V2EDim: 1}, + codomain=Edge, + data=np.asarray([0, 1, 2, 3], dtype=fbuiltins.IndexType).reshape((4, 1)), + ) + } + if as_type: + offset_provider = common.offset_provider_to_type(offset_provider) + + domain = domain_utils.SymbolicDomain.from_expr( + im.domain(common.GridType.UNSTRUCTURED, {Vertex: (0, 4)}) + ) + translated = domain.translate( + [itir.OffsetLiteral(value="V2E"), itir.OffsetLiteral(value=0)], + offset_provider, + symbolic_domain_sizes={"Edge": im.ref("num_edges")}, + ) + + expected = im.domain(common.GridType.UNSTRUCTURED, {Edge: (0, im.ref("num_edges"))}) + assert translated.as_expr() == expected + + +def test_translate_staggered_cartesian_offset(): + shift = (im.cartesian_offset(I, IHalf), im.ensure_offset(1)) + + domain = domain_utils.SymbolicDomain.from_expr( + im.domain(common.GridType.CARTESIAN, {I: (0, 10)}) + ) + translated = domain.translate(shift, {}) + + assert translated == domain_utils.SymbolicDomain.from_expr( + im.domain(common.GridType.CARTESIAN, {IHalf: (im.plus(0, 1), im.plus(10, 1))}) + ) + + def test_non_contiguous_domain_warning(monkeypatch): monkeypatch.setattr(domain_utils, "_NON_CONTIGUOUS_DOMAIN_WARNING_SKIPPED_OFFSET_TAGS", set()) diff --git a/tests/next_tests/unit_tests/test_common.py b/tests/next_tests/unit_tests/test_common.py index cc1c5505d9..fddc9dbd2e 100644 --- a/tests/next_tests/unit_tests/test_common.py +++ b/tests/next_tests/unit_tests/test_common.py @@ -37,6 +37,7 @@ IDim = Dimension("IDim") JDim = Dimension("JDim") KDim = Dimension("KDim", kind=DimensionKind.VERTICAL) +IHalfDim = common.flip_staggered(IDim) @pytest.fixture @@ -588,6 +589,12 @@ def dimension_promotion_cases() -> list[ "There are more than one dimension with DimensionKind 'LOCAL'.", ), ([[JDim, V2E], [IDim, KDim]], [IDim, JDim, V2E, KDim], None), + # a dimension and its staggered counterpart must not be promoted into the same field + ( + [[IDim], [common.flip_staggered(IDim)]], + None, + "staggered counterpart must not appear together", + ), ] return [ ([[el for el in arg] for arg in args], [el for el in result] if result else result, msg) @@ -677,6 +684,25 @@ def test_for_relocation(self): assert result.codomain == I assert result.offset == 0 + @pytest.mark.parametrize( + "dim, offset, exp_codomain, exp_offset", + [ + (IDim, 3, IDim, 3), + (IDim, 2.0, IDim, 2), + (IDim, 0.5, IHalfDim, 1), + (IDim, -0.5, IHalfDim, 0), + (IDim, 1.5, IHalfDim, 2), + (IDim, -1.5, IHalfDim, -1), + (IHalfDim, 0.5, IDim, 0), + (IHalfDim, -0.5, IDim, -1), + ], + ) + def test_connectivity_for_cartesian_shift(self, dim, offset, exp_codomain, exp_offset): + conn = common.connectivity_for_cartesian_shift(dim, offset) + assert conn.domain_dim == dim + assert conn.codomain == exp_codomain + assert conn.offset == exp_offset + class TestDimensionComparisonOperators: """Test Dimension comparison operators return correct Domain objects."""