feature[next]: Extend cartesian offset syntax

src/gt4py/next/common.py

@@ -70,6 +70,20 @@ def __str__(self) -> str:
         return self.value
 
 
+def dimension_to_implicit_offset(dim: str) -> str:
+    """
+    Return name of offset implicitly defined by a dimension.
+
+    Each dimension implicitly also defines an offset, such that we can allow syntax like::
+
+        field(TDim + 1)
+
+    without having to explicitly define an offset for ``TDim``. This function defines the respective
+    naming convention.
+    """
+    return f"_{dim}Off"
+
+
 @dataclasses.dataclass(frozen=True)
 class Dimension:
     value: str
@@ -81,6 +95,18 @@ def __str__(self) -> str:
     def __call__(self, val: int) -> NamedIndex:
         return NamedIndex(self, val)
 
+    def __add__(self, offset: int) -> ConnectivityField:
+        # TODO(sf-n): just to avoid circular import. Move or refactor the FieldOffset to avoid this.
+        from gt4py.next.ffront import fbuiltins
+
+        assert isinstance(self.value, str)
+        return fbuiltins.FieldOffset(
+            dimension_to_implicit_offset(self.value), source=self, target=(self,)
+        )[offset]
+
+    def __sub__(self, offset: int) -> ConnectivityField:
+        return self + (-offset)
+
 
 class Infinity(enum.Enum):
     """Describes an unbounded `UnitRange`."""
@@ -672,7 +698,11 @@ def as_scalar(self) -> core_defs.ScalarT: ...
 
     # Operators
     @abc.abstractmethod
-    def __call__(self, index_field: ConnectivityField | fbuiltins.FieldOffset) -> Field: ...
+    def __call__(
+        self,
+        index_field: ConnectivityField | fbuiltins.FieldOffset,
+        *args: ConnectivityField | fbuiltins.FieldOffset,
+    ) -> Field: ...
 
     @abc.abstractmethod
     def __getitem__(self, item: AnyIndexSpec) -> Self: ...
@@ -992,7 +1022,11 @@ def inverse_image(self, image_range: UnitRange | NamedRange) -> Sequence[NamedRa
         assert isinstance(image_range, UnitRange)
         return (named_range((self.domain_dim, image_range - self.offset)),)
 
-    def premap(self, index_field: ConnectivityField | fbuiltins.FieldOffset) -> ConnectivityField:
+    def premap(
+        self,
+        index_field: ConnectivityField | fbuiltins.FieldOffset,
+        *args: ConnectivityField | fbuiltins.FieldOffset,
+    ) -> ConnectivityField:
         raise NotImplementedError()
 
     __call__ = premap

src/gt4py/next/embedded/nd_array_field.py

@@ -315,7 +315,16 @@ def premap(
         assert len(conn_fields) == 1
         return _remapping_premap(self, conn_fields[0])
 
-    __call__ = premap  # type: ignore[assignment]
+    def __call__(
+        self,
+        index_field: common.ConnectivityField | fbuiltins.FieldOffset,
+        *args: common.ConnectivityField | fbuiltins.FieldOffset,
+    ) -> common.Field:
+        return functools.reduce(
+            lambda field, current_index_field: field.premap(current_index_field),
+            [index_field, *args],
+            self,
+        )
 
     def restrict(self, index: common.AnyIndexSpec) -> NdArrayField:
         new_domain, buffer_slice = self._slice(index)

src/gt4py/next/ffront/decorator.py

@@ -26,6 +26,7 @@
 from gt4py.next import (
     allocators as next_allocators,
     backend as next_backend,
+    common,
     embedded as next_embedded,
     errors,
 )
@@ -185,7 +186,35 @@ def itir(self) -> itir.FencilDefinition:
             return self.backend.transforms_prog.past_to_itir(no_args_past).program
         return past_to_itir.PastToItirFactory()(no_args_past).program
 
+    @functools.cached_property
+    def _implicit_offset_provider(self) -> dict[common.Tag, common.OffsetProviderElem]:
+        """
+        Add all implicit offset providers.
+
+        Each dimension implicitly defines an offset provider such that we can allow syntax like::
+
+            field(TDim + 1)
+
+        This function adds these implicit offset providers.
+        """
+        # TODO(tehrengruber): We add all dimensions here regardless of whether they are cartesian
+        #  or unstructured. While it is conceptually fine, but somewhat meaningless,
+        #  to do something `Cell+1` the GTFN backend for example doesn't support these. We should
+        #  find a way to avoid adding these dimensions, but since we don't have the grid type here
+        #  and since the dimensions don't this information either, we just add all dimensions here
+        #  and filter them out in the backends that don't support this.
+        implicit_offset_provider = {}
+        params = self.past_stage.past_node.params
+        for param in params:
+            if isinstance(param.type, ts.FieldType):
+                for dim in param.type.dims:
+                    implicit_offset_provider.update(
+                        {common.dimension_to_implicit_offset(dim.value): dim}
+                    )
+        return implicit_offset_provider
+
     def __call__(self, *args, offset_provider: dict[str, Dimension], **kwargs: Any) -> None:
+        offset_provider = offset_provider | self._implicit_offset_provider
         if self.backend is None:
             warnings.warn(
                 UserWarning(

src/gt4py/next/ffront/foast_passes/type_deduction.py

@@ -589,6 +589,18 @@ 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}'."
+            )
+
         logical_ops = {
             dialect_ast_enums.BinaryOperator.BIT_AND,
             dialect_ast_enums.BinaryOperator.BIT_OR,

src/gt4py/next/ffront/foast_to_gtir.py

@@ -13,6 +13,7 @@
 from gt4py import eve
 from gt4py.eve import utils as eve_utils
 from gt4py.eve.extended_typing import Never
+from gt4py.next import common
 from gt4py.next.ffront import (
     dialect_ast_enums,
     experimental as experimental_builtins,
@@ -219,25 +220,62 @@ def visit_Compare(self, node: foast.Compare, **kwargs: Any) -> itir.FunCall:
         return self._map(node.op.value, node.left, node.right)
 
     def _visit_shift(self, node: foast.Call, **kwargs: Any) -> itir.Expr:
-        match node.args[0]:
-            case foast.Subscript(value=foast.Name(id=offset_name), index=int(offset_index)):
-                shift_offset = im.shift(offset_name, offset_index)
-                return im.as_fieldop(im.lambda_("__it")(im.deref(shift_offset("__it"))))(
-                    self.visit(node.func, **kwargs)
-                )
-            case foast.Name(id=offset_name):
-                return im.as_fieldop_neighbors(str(offset_name), self.visit(node.func, **kwargs))
-            case foast.Call(func=foast.Name(id="as_offset")):
-                # TODO(havogt): discuss this representation
-                func_args = node.args[0]
-                offset_dim = func_args.args[0]
-                assert isinstance(offset_dim, foast.Name)
-                shift_offset = im.shift(offset_dim.id, im.deref("__offset"))
-                return im.as_fieldop(
-                    im.lambda_("__it", "__offset")(im.deref(shift_offset("__it")))
-                )(self.visit(node.func, **kwargs), self.visit(func_args.args[1], **kwargs))
-            case _:
-                raise FieldOperatorLoweringError("Unexpected shift arguments!")
+        current_expr = self.visit(node.func, **kwargs)
+
+        for arg in node.args:
+            match arg:
+                # `field(Off[idx])`
+                case foast.Subscript(value=foast.Name(id=offset_name), index=int(offset_index)):
+                    current_expr = im.as_fieldop(
+                        im.lambda_("__it")(im.deref(im.shift(offset_name, offset_index)("__it")))
+                    )(current_expr)
+                # `field(Dim + idx)`
+                case foast.BinOp(
+                    op=dialect_ast_enums.BinaryOperator.ADD
+                    | dialect_ast_enums.BinaryOperator.SUB,
+                    left=foast.Name(id=dimension),  # TODO(tehrengruber): use type of lhs
+                    right=foast.Constant(value=offset_index),
+                ):
+                    if arg.op == dialect_ast_enums.BinaryOperator.SUB:
+                        offset_index *= -1
+                    current_expr = im.as_fieldop(
+                        # TODO(SF-N): we rely on the naming-convention that the cartesian dimensions
+                        #  are passed suffixed with `off`, e.g. the `K` is passed as `Koff` in the
+                        #  offset provider. This is a rather unclean solution and should be
+                        #  improved.
+                        im.lambda_("__it")(
+                            im.deref(
+                                im.shift(
+                                    common.dimension_to_implicit_offset(dimension), offset_index
+                                )("__it")
+                            )
+                        )
+                    )(current_expr)
+                # `field(Off)`
+                case foast.Name(id=offset_name):
+                    # only a single unstructured shift is supported so returning here is fine even though we
+                    # are in a loop.
+                    assert len(node.args) == 1 and len(arg.type.target) > 1  # type: ignore[attr-defined] # ensured by pattern
+                    return im.as_fieldop_neighbors(
+                        str(offset_name), self.visit(node.func, **kwargs)
+                    )
+                # `field(as_offset(Off, offset_field))`
+                case foast.Call(func=foast.Name(id="as_offset")):
+                    func_args = arg
+                    # TODO(tehrengruber): Discuss representation. We could use the type system to
+                    #  deduce the offset dimension instead of (e.g. to allow aliasing).
+                    offset_dim = func_args.args[0]
+                    assert isinstance(offset_dim, foast.Name)
+                    offset_field = self.visit(func_args.args[1], **kwargs)
+                    current_expr = im.as_fieldop(
+                        im.lambda_("__it", "__offset")(
+                            im.deref(im.shift(offset_dim.id, im.deref("__offset"))("__it"))
+                        )
+                    )(current_expr, offset_field)
+                case _:
+                    raise FieldOperatorLoweringError("Unexpected shift arguments!")
+
+        return current_expr
 
     def visit_Call(self, node: foast.Call, **kwargs: Any) -> itir.Expr:
         if type_info.type_class(node.func.type) is ts.FieldType:

src/gt4py/next/ffront/foast_to_itir.py

@@ -14,6 +14,7 @@
 from gt4py.eve import NodeTranslator, PreserveLocationVisitor
 from gt4py.eve.extended_typing import Never
 from gt4py.eve.utils import UIDGenerator
+from gt4py.next import common
 from gt4py.next.ffront import (
     dialect_ast_enums,
     fbuiltins,
@@ -276,23 +277,60 @@ def visit_Compare(self, node: foast.Compare, **kwargs: Any) -> itir.FunCall:
         return self._map(node.op.value, node.left, node.right)
 
     def _visit_shift(self, node: foast.Call, **kwargs: Any) -> itir.Expr:
-        match node.args[0]:
-            case foast.Subscript(value=foast.Name(id=offset_name), index=int(offset_index)):
-                shift_offset = im.shift(offset_name, offset_index)
-            case foast.Name(id=offset_name):
-                return im.lifted_neighbors(str(offset_name), self.visit(node.func, **kwargs))
-            case foast.Call(func=foast.Name(id="as_offset")):
-                func_args = node.args[0]
-                offset_dim = func_args.args[0]
-                assert isinstance(offset_dim, foast.Name)
-                shift_offset = im.shift(
-                    offset_dim.id, im.deref(self.visit(func_args.args[1], **kwargs))
-                )
-            case _:
-                raise FieldOperatorLoweringError("Unexpected shift arguments!")
-        return im.lift(im.lambda_("it")(im.deref(shift_offset("it"))))(
-            self.visit(node.func, **kwargs)
-        )
+        current_expr = self.visit(node.func, **kwargs)
+
+        for arg in node.args:
+            match arg:
+                # `field(Off[idx])`
+                case foast.Subscript(value=foast.Name(id=offset_name), index=int(offset_index)):
+                    current_expr = im.lift(
+                        im.lambda_("it")(im.deref(im.shift(offset_name, offset_index)("it")))
+                    )(current_expr)
+                # `field(Dim + idx)`
+                case foast.BinOp(
+                    op=dialect_ast_enums.BinaryOperator.ADD
+                    | dialect_ast_enums.BinaryOperator.SUB,
+                    left=foast.Name(id=dimension),
+                    right=foast.Constant(value=offset_index),
+                ):
+                    if arg.op == dialect_ast_enums.BinaryOperator.SUB:
+                        offset_index *= -1
+                    current_expr = im.lift(
+                        # TODO(SF-N): we rely on the naming-convention that the cartesian dimensions
+                        #  are passed suffixed with `off`, e.g. the `K` is passed as `Koff` in the
+                        #  offset provider. This is a rather unclean solution and should be
+                        #  improved.
+                        im.lambda_("it")(
+                            im.deref(
+                                im.shift(
+                                    common.dimension_to_implicit_offset(dimension), offset_index
+                                )("it")
+                            )
+                        )
+                    )(current_expr)
+                # `field(Off)`
+                case foast.Name(id=offset_name):
+                    # only a single unstructured shift is supported so returning here is fine even though we
+                    # are in a loop.
+                    assert len(node.args) == 1 and len(arg.type.target) > 1  # type: ignore[attr-defined] # ensured by pattern
+                    return im.lifted_neighbors(str(offset_name), self.visit(node.func, **kwargs))
+                # `field(as_offset(Off, offset_field))`
+                case foast.Call(func=foast.Name(id="as_offset")):
+                    func_args = arg
+                    # TODO(tehrengruber): Use type system to deduce the offset dimension instead of
+                    #  (e.g. to allow aliasing)
+                    offset_dim = func_args.args[0]
+                    assert isinstance(offset_dim, foast.Name)
+                    offset_it = self.visit(func_args.args[1], **kwargs)
+                    current_expr = im.lift(
+                        im.lambda_("it", "offset")(
+                            im.deref(im.shift(offset_dim.id, im.deref("offset"))("it"))
+                        )
+                    )(current_expr, offset_it)
+                case _:
+                    raise FieldOperatorLoweringError("Unexpected shift arguments!")
+
+        return current_expr
 
     def visit_Call(self, node: foast.Call, **kwargs: Any) -> itir.Expr:
         if type_info.type_class(node.func.type) is ts.FieldType:

src/gt4py/next/ffront/past_to_itir.py

@@ -15,7 +15,7 @@
 import factory
 
 from gt4py.eve import NodeTranslator, concepts, traits
-from gt4py.next import common, config
+from gt4py.next import common, config, errors
 from gt4py.next.ffront import (
     fbuiltins,
     gtcallable,
@@ -435,9 +435,10 @@ def _compute_field_slice(node: past.Subscript) -> list[past.Slice]:
         node_dims = cast(ts.FieldType, node.type).dims
         assert isinstance(node_dims, list)
         if isinstance(node.type, ts.FieldType) and len(out_field_slice_) != len(node_dims):
-            raise ValueError(
+            raise errors.DSLError(
+                node.location,
                 f"Too many indices for field '{out_field_name}': field is {len(node_dims)}"
-                f"-dimensional, but {len(out_field_slice_)} were indexed."
+                f"-dimensional, but {len(out_field_slice_)} were indexed.",
             )
         return out_field_slice_
 

src/gt4py/next/iterator/embedded.py

@@ -1153,7 +1153,11 @@ def as_scalar(self) -> core_defs.IntegralScalar:
         assert self._cur_index is not None
         return self._cur_index
 
-    def premap(self, index_field: common.ConnectivityField | fbuiltins.FieldOffset) -> common.Field:
+    def premap(
+        self,
+        index_field: common.ConnectivityField | fbuiltins.FieldOffset,
+        *args: common.ConnectivityField | fbuiltins.FieldOffset,
+    ) -> common.Field:
         # TODO can be implemented by constructing and ndarray (but do we know of which kind?)
         raise NotImplementedError()
 
@@ -1273,7 +1277,11 @@ def ndarray(self) -> core_defs.NDArrayObject:
     def asnumpy(self) -> np.ndarray:
         raise NotImplementedError()
 
-    def premap(self, index_field: common.ConnectivityField | fbuiltins.FieldOffset) -> common.Field:
+    def premap(
+        self,
+        index_field: common.ConnectivityField | fbuiltins.FieldOffset,
+        *args: common.ConnectivityField | fbuiltins.FieldOffset,
+    ) -> common.Field:
         # TODO can be implemented by constructing and ndarray (but do we know of which kind?)
         raise NotImplementedError()
 

src/gt4py/next/program_processors/codegens/gtfn/itir_to_gtfn_ir.py

@@ -154,6 +154,17 @@ def _collect_offset_definitions(
                 )
             else:
                 assert grid_type == common.GridType.UNSTRUCTURED
+                # TODO(tehrengruber): The implicit offset providers added to support syntax like
+                #  `KDim+1` can also include horizontal dimensions. Cartesian shifts in this
+                #  dimension are not supported by the backend and also never occur in user code.
+                #  We just skip these here for now, but this is not a clean solution. Not having
+                #  any unstructured dimensions in here would be preferred.
+                if (
+                    dim.kind == common.DimensionKind.HORIZONTAL
+                    and offset_name == common.dimension_to_implicit_offset(dim.value)
+                ):
+                    continue
+
                 if not dim.kind == common.DimensionKind.VERTICAL:
                     raise ValueError(
                         "Mapping an offset to a horizontal dimension in unstructured is not allowed."

src/gt4py/next/program_processors/runners/dace_iterator/__init__.py

@@ -345,7 +345,9 @@ def __sdfg__(self, *args, **kwargs) -> dace.sdfg.sdfg.SDFG:
             raise ValueError(
                 "[DaCe Orchestration] Connectivities -at compile time- are required to generate the SDFG. Use `with_connectivities` method."
             )
-        offset_provider = self.connectivities  # tables are None at this point
+        offset_provider = (
+            self.connectivities | self._implicit_offset_provider
+        )  # tables are None at this point
 
         sdfg = self.backend.executor.otf_workflow.step.translation.generate_sdfg(  # type: ignore[union-attr]
             self.itir,