feat[next]: Output argument with non-zero domain start

src/gt4py/next/ffront/past_process_args.py

@@ -83,40 +83,46 @@ def _process_args(
             # TODO(tehrengruber): Previously this function was called with the actual arguments
             #  not their type. The check using the shape here is not functional anymore and
             #  should instead be placed in a proper location.
-            shapes_and_dims = [*_field_constituents_shape_and_dims(args[param_idx], param.type)]
+            ranges_and_dims = [*_field_constituents_range_and_dims(args[param_idx], param.type)]
             # check that all non-scalar like constituents have the same shape and dimension, e.g.
             # for `(scalar, (field1, field2))` the two fields need to have the same shape and
             # dimension
-            if shapes_and_dims:
-                shape, dims = shapes_and_dims[0]
+            if ranges_and_dims:
+                range_, dims = ranges_and_dims[0]
                 if not all(
-                    el_shape == shape and el_dims == dims for (el_shape, el_dims) in shapes_and_dims
+                    el_range == range_ and el_dims == dims
+                    for (el_range, el_dims) in ranges_and_dims
                 ):
                     raise ValueError(
                         "Constituents of composite arguments (e.g. the elements of a"
                         " tuple) need to have the same shape and dimensions."
                     )
+                index_type = ts.ScalarType(kind=ts.ScalarKind.INT32)
                 size_args.extend(
-                    shape if shape else [ts.ScalarType(kind=ts.ScalarKind.INT32)] * len(dims)  # type: ignore[arg-type]  # shape is always empty
+                    range_ if range_ else [ts.TupleType(types=[index_type, index_type])] * len(dims)  # type: ignore[arg-type]  # shape is always empty
                 )
     return tuple(rewritten_args), tuple(size_args), kwargs
 
 
-def _field_constituents_shape_and_dims(
+def _field_constituents_range_and_dims(
     arg: Any,  # TODO(havogt): improve typing
     arg_type: ts.DataType,
-) -> Iterator[tuple[tuple[int, ...], list[common.Dimension]]]:
+) -> Iterator[tuple[tuple[tuple[int, int], ...], list[common.Dimension]]]:
     match arg_type:
         case ts.TupleType():
             for el, el_type in zip(arg, arg_type.types):
-                yield from _field_constituents_shape_and_dims(el, el_type)
+                yield from _field_constituents_range_and_dims(el, el_type)
         case ts.FieldType():
             dims = type_info.extract_dims(arg_type)
             if isinstance(arg, ts.TypeSpec):  # TODO
                 yield (tuple(), dims)
             elif dims:
-                assert hasattr(arg, "shape") and len(arg.shape) == len(dims)
-                yield (arg.shape, dims)
+                assert (
+                    hasattr(arg, "domain")
+                    and isinstance(arg.domain, common.Domain)
+                    and len(arg.domain.dims) == len(dims)
+                )
+                yield (tuple((r.start, r.stop) for r in arg.domain.ranges), dims)
             else:
                 yield from []  # ignore 0-dim fields
         case ts.ScalarType():

src/gt4py/next/ffront/past_to_itir.py

@@ -138,8 +138,8 @@ def _column_axis(all_closure_vars: dict[str, Any]) -> Optional[common.Dimension]
     return iter(scanops_per_axis.keys()).__next__()
 
 
-def _size_arg_from_field(field_name: str, dim: int) -> str:
-    return f"__{field_name}_size_{dim}"
+def _range_arg_from_field(field_name: str, dim: int) -> str:
+    return f"__{field_name}_{dim}_range"
 
 
 def _flatten_tuple_expr(node: past.Expr) -> list[past.Name | past.Subscript]:
@@ -217,13 +217,14 @@ def _gen_size_params_from_program(self, node: past.Program) -> list[itir.Sym]:
             )
             if len(fields_dims) > 0:  # otherwise `param` has no constituent which is of `FieldType`
                 assert all(field_dims == fields_dims[0] for field_dims in fields_dims)
+                index_type = ts.ScalarType(
+                    kind=getattr(ts.ScalarKind, itir.INTEGER_INDEX_BUILTIN.upper())
+                )
                 for dim_idx in range(len(fields_dims[0])):
                     size_params.append(
                         itir.Sym(
-                            id=_size_arg_from_field(param.id, dim_idx),
-                            type=ts.ScalarType(
-                                kind=getattr(ts.ScalarKind, itir.INTEGER_INDEX_BUILTIN.upper())
-                            ),
+                            id=_range_arg_from_field(param.id, dim_idx),
+                            type=ts.TupleType(types=[index_type, index_type]),
                         )
                     )
 
@@ -286,7 +287,8 @@ def _visit_slice_bound(
         self,
         slice_bound: Optional[past.Constant],
         default_value: itir.Expr,
-        dim_size: itir.Expr,
+        start_idx: itir.Expr,
+        stop_idx: itir.Expr,
         **kwargs: Any,
     ) -> itir.Expr:
         if slice_bound is None:
@@ -296,11 +298,9 @@ def _visit_slice_bound(
                 slice_bound.type
             )
             if slice_bound.value < 0:
-                lowered_bound = itir.FunCall(
-                    fun=itir.SymRef(id="plus"), args=[dim_size, self.visit(slice_bound, **kwargs)]
-                )
+                lowered_bound = im.plus(stop_idx, self.visit(slice_bound, **kwargs))
             else:
-                lowered_bound = self.visit(slice_bound, **kwargs)
+                lowered_bound = im.plus(start_idx, self.visit(slice_bound, **kwargs))
         else:
             raise AssertionError("Expected 'None' or 'past.Constant'.")
         if slice_bound:
@@ -348,8 +348,9 @@ def _construct_itir_domain_arg(
         domain_args = []
         domain_args_kind = []
         for dim_i, dim in enumerate(out_dims):
-            # an expression for the size of a dimension
-            dim_size = itir.SymRef(id=_size_arg_from_field(out_field.id, dim_i))
+            # an expression for the range of a dimension
+            dim_range = itir.SymRef(id=_range_arg_from_field(out_field.id, dim_i))
+            dim_start, dim_stop = im.tuple_get(0, dim_range), im.tuple_get(1, dim_range)
             # bounds
             lower: itir.Expr
             upper: itir.Expr
@@ -359,11 +360,12 @@ def _construct_itir_domain_arg(
             else:
                 lower = self._visit_slice_bound(
                     slices[dim_i].lower if slices else None,
-                    im.literal("0", itir.INTEGER_INDEX_BUILTIN),
-                    dim_size,
+                    dim_start,
+                    dim_start,
+                    dim_stop,
                 )
                 upper = self._visit_slice_bound(
-                    slices[dim_i].upper if slices else None, dim_size, dim_size
+                    slices[dim_i].upper if slices else None, dim_stop, dim_start, dim_stop
                 )
 
             if dim.kind == common.DimensionKind.LOCAL:

src/gt4py/next/otf/arguments.py

@@ -122,7 +122,7 @@ def find_first_field(tuple_arg: tuple[Any, ...]) -> Optional[common.Field]:
     return None
 
 
-def iter_size_args(args: tuple[Any, ...]) -> Iterator[int]:
+def iter_size_args(args: tuple[Any, ...]) -> Iterator[tuple[int, int]]:
     """
     Yield the size of each field argument in each dimension.
 
@@ -136,7 +136,9 @@ def iter_size_args(args: tuple[Any, ...]) -> Iterator[int]:
                 if first_field:
                     yield from iter_size_args((first_field,))
             case common.Field():
-                yield from arg.ndarray.shape
+                for range_ in arg.domain.ranges:
+                    assert isinstance(range_, common.UnitRange)
+                    yield (range_.start, range_.stop)
             case _:
                 pass
 
@@ -156,6 +158,7 @@ def iter_size_compile_args(
         )
         if field_constituents:
             # we only need the first field, because all fields in a tuple must have the same dims and sizes
+            index_type = ts.ScalarType(kind=ts.ScalarKind.INT32)
             yield from [
-                ts.ScalarType(kind=ts.ScalarKind.INT32) for dim in field_constituents[0].dims
+                ts.TupleType(types=[index_type, index_type]) for dim in field_constituents[0].dims
             ]

src/gt4py/next/program_processors/runners/dace_common/utility.py

@@ -19,7 +19,7 @@
 
 
 # regex to match the symbols for field shape and strides
-FIELD_SYMBOL_RE: Final[re.Pattern] = re.compile(r"__.+_(size|stride)_\d+")
+FIELD_SYMBOL_RE: Final[re.Pattern] = re.compile(r"^__.+_(range_[01]|((size|stride)_\d+))$")
 
 
 def as_dace_type(type_: ts.ScalarType) -> dace.typeclass:

src/gt4py/next/program_processors/runners/dace_fieldview/gtir_python_codegen.py

@@ -84,6 +84,11 @@ def builtin_if(*args: Any) -> str:
     return f"{true_val} if {cond} else {false_val}"
 
 
+def builtin_tuple_get(*args: Any) -> str:
+    index, tuple_name = args
+    return f"{tuple_name}_{index}"
+
+
 def make_const_list(arg: str) -> str:
     """
     Takes a single scalar argument and broadcasts this value on the local dimension
@@ -97,6 +102,7 @@ def make_const_list(arg: str) -> str:
     "cast_": builtin_cast,
     "if_": builtin_if,
     "make_const_list": make_const_list,
+    "tuple_get": builtin_tuple_get,
 }
 
 

tests/next_tests/definitions.py

@@ -134,6 +134,7 @@ class ProgramFormatterId(_PythonObjectIdMixin, str, enum.Enum):
 ]
 DACE_SKIP_TEST_LIST = DOMAIN_INFERENCE_SKIP_LIST + [
     (USES_NEGATIVE_MODULO, XFAIL, UNSUPPORTED_MESSAGE),
+    (USES_ORIGIN, XFAIL, UNSUPPORTED_MESSAGE),
     (USES_SCAN, XFAIL, UNSUPPORTED_MESSAGE),
     (USES_SPARSE_FIELDS_AS_OUTPUT, XFAIL, UNSUPPORTED_MESSAGE),
 ]

tests/next_tests/integration_tests/feature_tests/ffront_tests/test_execution.py

@@ -994,7 +994,7 @@ def program_domain(a: cases.IField, out: cases.IField):
     a = cases.allocate(cartesian_case, program_domain, "a")()
     out = cases.allocate(cartesian_case, program_domain, "out")()
 
-    ref = out.asnumpy().copy()  # ensure we are not overwriting out outside of the domain
+    ref = out.asnumpy().copy()  # ensure we are not overwriting out outside the domain
     ref[1:9] = a.asnumpy()[1:9] * 2
 
     cases.verify(cartesian_case, program_domain, a, out, inout=out, ref=ref)

tests/next_tests/integration_tests/feature_tests/ffront_tests/test_program.py

@@ -14,7 +14,7 @@
 import pytest
 
 import gt4py.next as gtx
-from gt4py.next import errors
+from gt4py.next import errors, constructors, common
 
 from next_tests.integration_tests import cases
 from next_tests.integration_tests.cases import (
@@ -251,3 +251,42 @@ def empty_domain_program(a: cases.IJField, out_field: cases.IJField):
         ValueError, match=(r"Dimensions in out field and field domain are not equivalent")
     ):
         cases.run(cartesian_case, empty_domain_program, a, out_field, offset_provider={})
+
+
+@pytest.mark.uses_origin
+def test_out_field_arg_with_non_zero_domain_start(cartesian_case, copy_program_def):
+    copy_program = gtx.program(copy_program_def, backend=cartesian_case.backend)
+
+    size = cartesian_case.default_sizes[IDim]
+
+    inp = cases.allocate(cartesian_case, copy_program, "in_field").unique()()
+    out = constructors.empty(
+        common.domain({IDim: (1, size - 2)}),
+        allocator=cartesian_case.allocator,
+    )
+    ref = inp.ndarray[1:-2]
+
+    cases.verify(cartesian_case, copy_program, inp, out=out, ref=ref)
+
+
+@pytest.mark.uses_origin
+def test_in_field_arg_with_non_zero_domain_start(cartesian_case, copy_program_def):
+    @gtx.field_operator
+    def identity(a: cases.IField) -> cases.IField:
+        return a
+
+    @gtx.program
+    def copy_program(a: cases.IField, out: cases.IField):
+        identity(a, out=out, domain={IDim: (1, 9)})
+
+    inp = constructors.empty(
+        common.domain({IDim: (1, 9)}),
+        dtype=np.int32,
+        allocator=cartesian_case.allocator,
+    )
+    inp.ndarray[...] = 42
+    out = cases.allocate(cartesian_case, copy_program, "out", sizes={IDim: 10})()
+    ref = out.asnumpy().copy()  # ensure we are not overwriting `out` outside the domain
+    ref[1:9] = inp.asnumpy()
+
+    cases.verify(cartesian_case, copy_program, inp, out=out, ref=ref)

tests/next_tests/unit_tests/ffront_tests/test_past_to_gtir.py

@@ -58,8 +58,30 @@ def test_copy_lowering(copy_program_def, gtir_identity_fundef):
                     fun=P(itir.SymRef, id=eve.SymbolRef("named_range")),
                     args=[
                         P(itir.AxisLiteral, value="IDim"),
-                        P(itir.Literal, value="0", type=ts.ScalarType(kind=ts.ScalarKind.INT32)),
-                        P(itir.SymRef, id=eve.SymbolRef("__out_size_0")),
+                        P(
+                            itir.FunCall,
+                            fun=P(itir.SymRef, id=eve.SymbolRef("tuple_get")),
+                            args=[
+                                P(
+                                    itir.Literal,
+                                    value="0",
+                                    type=ts.ScalarType(kind=ts.ScalarKind.INT32),
+                                ),
+                                P(itir.SymRef, id=eve.SymbolRef("__out_0_range")),
+                            ],
+                        ),
+                        P(
+                            itir.FunCall,
+                            fun=P(itir.SymRef, id=eve.SymbolRef("tuple_get")),
+                            args=[
+                                P(
+                                    itir.Literal,
+                                    value="1",
+                                    type=ts.ScalarType(kind=ts.ScalarKind.INT32),
+                                ),
+                                P(itir.SymRef, id=eve.SymbolRef("__out_0_range")),
+                            ],
+                        ),
                     ],
                 )
             ],
@@ -77,8 +99,8 @@ def test_copy_lowering(copy_program_def, gtir_identity_fundef):
         params=[
             P(itir.Sym, id=eve.SymbolName("in_field")),
             P(itir.Sym, id=eve.SymbolName("out")),
-            P(itir.Sym, id=eve.SymbolName("__in_field_size_0")),
-            P(itir.Sym, id=eve.SymbolName("__out_size_0")),
+            P(itir.Sym, id=eve.SymbolName("__in_field_0_range")),
+            P(itir.Sym, id=eve.SymbolName("__out_0_range")),
         ],
         body=[set_at_pattern],
     )
@@ -105,18 +127,58 @@ def test_copy_restrict_lowering(copy_restrict_program_def, gtir_identity_fundef)
                     args=[
                         P(itir.AxisLiteral, value="IDim"),
                         P(
-                            itir.Literal,
-                            value="1",
-                            type=ts.ScalarType(
-                                kind=getattr(ts.ScalarKind, itir.INTEGER_INDEX_BUILTIN.upper())
-                            ),
+                            itir.FunCall,
+                            fun=P(itir.SymRef, id=eve.SymbolRef("plus")),
+                            args=[
+                                P(
+                                    itir.FunCall,
+                                    fun=P(itir.SymRef, id=eve.SymbolRef("tuple_get")),
+                                    args=[
+                                        P(
+                                            itir.Literal,
+                                            value="0",
+                                            type=ts.ScalarType(kind=ts.ScalarKind.INT32),
+                                        ),
+                                        P(itir.SymRef, id=eve.SymbolRef("__out_0_range")),
+                                    ],
+                                ),
+                                P(
+                                    itir.Literal,
+                                    value="1",
+                                    type=ts.ScalarType(
+                                        kind=getattr(
+                                            ts.ScalarKind, itir.INTEGER_INDEX_BUILTIN.upper()
+                                        )
+                                    ),
+                                ),
+                            ],
                         ),
                         P(
-                            itir.Literal,
-                            value="2",
-                            type=ts.ScalarType(
-                                kind=getattr(ts.ScalarKind, itir.INTEGER_INDEX_BUILTIN.upper())
-                            ),
+                            itir.FunCall,
+                            fun=P(itir.SymRef, id=eve.SymbolRef("plus")),
+                            args=[
+                                P(
+                                    itir.FunCall,
+                                    fun=P(itir.SymRef, id=eve.SymbolRef("tuple_get")),
+                                    args=[
+                                        P(
+                                            itir.Literal,
+                                            value="0",
+                                            type=ts.ScalarType(kind=ts.ScalarKind.INT32),
+                                        ),
+                                        P(itir.SymRef, id=eve.SymbolRef("__out_0_range")),
+                                    ],
+                                ),
+                                P(
+                                    itir.Literal,
+                                    value="2",
+                                    type=ts.ScalarType(
+                                        kind=getattr(
+                                            ts.ScalarKind, itir.INTEGER_INDEX_BUILTIN.upper()
+                                        )
+                                    ),
+                                ),
+                            ],
                         ),
                     ],
                 )
@@ -129,8 +191,8 @@ def test_copy_restrict_lowering(copy_restrict_program_def, gtir_identity_fundef)
         params=[
             P(itir.Sym, id=eve.SymbolName("in_field")),
             P(itir.Sym, id=eve.SymbolName("out")),
-            P(itir.Sym, id=eve.SymbolName("__in_field_size_0")),
-            P(itir.Sym, id=eve.SymbolName("__out_size_0")),
+            P(itir.Sym, id=eve.SymbolName("__in_field_0_range")),
+            P(itir.Sym, id=eve.SymbolName("__out_0_range")),
         ],
         body=[set_at_pattern],
     )