Implementations for as_offset

src/gt4py/next/ffront/experimental.py

@@ -0,0 +1,40 @@
+# GT4Py - GridTools Framework
+#
+# Copyright (c) 2014-2023, ETH Zurich
+# All rights reserved.
+#
+# This file is part of the GT4Py project and the GridTools framework.
+# GT4Py is free software: you can redistribute it and/or modify it under
+# the terms of the GNU General Public License as published by the
+# Free Software Foundation, either version 3 of the License, or any later
+# version. See the LICENSE.txt file at the top-level directory of this
+# distribution for a copy of the license or check <https://www.gnu.org/licenses/>.
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+
+from dataclasses import dataclass
+
+from gt4py.next.type_system import type_specifications as ts
+
+
+@dataclass
+class BuiltInFunction:
+    __gt_type: ts.FunctionType
+
+    def __call__(self, *args, **kwargs):
+        """Act as an empty place holder for the built in function."""
+
+    def __gt_type__(self):
+        return self.__gt_type
+
+
+as_offset = BuiltInFunction(
+    ts.FunctionType(
+        args=[
+            ts.DeferredType(constraint=ts.OffsetType),
+            ts.DeferredType(constraint=ts.FieldType),
+        ],
+        kwargs={},
+        returns=ts.DeferredType(constraint=ts.OffsetType),
+    )
+)

src/gt4py/next/ffront/fbuiltins.py

@@ -18,6 +18,7 @@
 from numpy import float32, float64, int32, int64
 
 from gt4py.next.common import Dimension, DimensionKind, Field
+from gt4py.next.ffront.experimental import as_offset  # noqa F401
 from gt4py.next.iterator import runtime
 from gt4py.next.type_system import type_specifications as ts
 
@@ -192,6 +193,7 @@ def __gt_type__(self):
     "broadcast",
     "where",
     "astype",
+    "as_offset",
 ] + MATH_BUILTIN_NAMES
 
 BUILTIN_NAMES = TYPE_BUILTIN_NAMES + FUN_BUILTIN_NAMES

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

@@ -718,6 +718,35 @@ def _visit_astype(self, node: foast.Call, **kwargs) -> foast.Call:
             location=node.location,
         )
 
+    def _visit_as_offset(self, node: foast.Call, **kwargs) -> foast.Call:
+        arg_0 = node.args[0].type
+        arg_1 = node.args[1].type
+        assert isinstance(arg_0, ts.OffsetType)
+        assert isinstance(arg_1, ts.FieldType)
+        if not type_info.is_integral(arg_1):
+            raise FieldOperatorTypeDeductionError.from_foast_node(
+                node,
+                msg=f"Incompatible argument in call to `{node.func.id}`. "
+                f"Excepted integer for offset field dtype, but got {arg_1.dtype}"
+                f"{node.location}",
+            )
+
+        if arg_0.source not in arg_1.dims:
+            raise FieldOperatorTypeDeductionError.from_foast_node(
+                node,
+                msg=f"Incompatible argument in call to `{node.func.id}`. "
+                f"{arg_0.source} not in list of offset field dimensions {arg_1.dims}. "
+                f"{node.location}",
+            )
+
+        return foast.Call(
+            func=node.func,
+            args=node.args,
+            kwargs=node.kwargs,
+            type=arg_0,
+            location=node.location,
+        )
+
     def _visit_where(self, node: foast.Call, **kwargs) -> foast.Call:
         mask_type = cast(ts.FieldType, node.args[0].type)
         true_branch_type = node.args[1].type

src/gt4py/next/ffront/foast_to_itir.py

@@ -357,10 +357,19 @@ def visit_Compare(self, node: foast.Compare, **kwargs) -> itir.FunCall:
     def _visit_shift(self, node: foast.Call, **kwargs) -> itir.FunCall:
         match node.args[0]:
             case foast.Subscript(value=foast.Name(id=offset_name), index=int(offset_index)):
-                return im.shift_(offset_name, offset_index)(self.visit(node.func, **kwargs))
+                shift_offset = im.shift_(offset_name, offset_index)
             case foast.Name(id=offset_name):
-                return im.shift_(offset_name)(self.visit(node.func, **kwargs))
-        raise FieldOperatorLoweringError("Unexpected shift arguments!")
+                shift_offset = im.shift_(offset_name)
+            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 shift_offset(self.visit(node.func, **kwargs))
 
     def _make_reduction_expr(
         self,

tests/next_tests/ffront_tests/test_execution.py

@@ -20,22 +20,21 @@
 import pytest as pytest
 
 from gt4py.next.ffront.decorator import field_operator, program, scan_operator
+from gt4py.next.ffront.experimental import as_offset
 from gt4py.next.ffront.fbuiltins import (
     Dimension,
     Field,
     FieldOffset,
     astype,
     broadcast,
+    float32,
     float64,
     int32,
     int64,
-    max_over,
-    min_over,
     neighbor_sum,
     where,
 )
 from gt4py.next.ffront.foast_passes.type_deduction import FieldOperatorTypeDeductionError
-from gt4py.next.iterator.builtins import float32
 from gt4py.next.iterator.embedded import (
     NeighborTableOffsetProvider,
     index_field,
@@ -336,6 +335,50 @@ def astype_fieldop_float(b: Field[[IDim], int64]) -> Field[[IDim], int32]:
     assert np.allclose(c_int32.array(), out_int_32)
 
 
+def test_offset_field(fieldview_backend):
+    a_I_arr = np.random.randn(size, size).astype("float64")
+    a_I_float = np_as_located_field(IDim, KDim)(a_I_arr)
+    a_I_float_1 = np_as_located_field(IDim, KDim)(
+        np.append(np.insert(a_I_arr, size, 0, axis=1), [np.array([0] * (size + 1))], axis=0)
+    )
+    offset_field_arr = np.asarray(np.ones((size - 1, size - 1)), dtype=int64)
+    offset_field_comp = np.append(
+        np.insert(offset_field_arr, size - 1, 0, axis=1), [np.array([0] * size)], axis=0
+    )
+    offset_field = np_as_located_field(IDim, KDim)(offset_field_comp)
+    out_I_float = np_as_located_field(IDim, KDim)(np.zeros((size, size), dtype=float64))
+    out_I_float_1 = np_as_located_field(IDim, KDim)(np.zeros((size, size), dtype=float64))
+
+    @field_operator(backend=fieldview_backend)
+    def offset_index_field_fo(
+        a: Field[[IDim, KDim], float64],
+        offset_field: Field[[IDim, KDim], int64],
+    ) -> Field[[IDim, KDim], float64]:
+        a_i = a(as_offset(Ioff, offset_field))
+        a_i_k = a_i(as_offset(Koff, offset_field))
+        return a_i_k
+
+    offset_index_field_fo(
+        a_I_float,
+        offset_field,
+        out=out_I_float,
+        offset_provider={"Ioff": IDim, "Koff": KDim},
+    )
+
+    @field_operator(backend=fieldview_backend)
+    def offset_index_int_fo(a: Field[[IDim, KDim], float64]) -> Field[[IDim, KDim], float64]:
+        a_i = a(Ioff[1])
+        a_i_k = a_i(Koff[1])
+        return a_i_k
+
+    offset_index_int_fo(
+        a_I_float_1, out=out_I_float_1, offset_provider={"Ioff": IDim, "Koff": KDim}
+    )
+    assert np.allclose(
+        out_I_float.array()[: size - 1, : size - 1], out_I_float_1.array()[: size - 1, : size - 1]
+    )
+
+
 def test_nested_tuple_return(fieldview_backend):
     a_I_float = np_as_located_field(IDim)(np.random.randn(size).astype("float64"))
     b_I_float = np_as_located_field(IDim)(np.random.randn(size).astype("float64"))

tests/next_tests/ffront_tests/test_type_deduction.py

@@ -18,6 +18,7 @@
 
 import gt4py.next.ffront.type_specifications
 from gt4py.next.common import DimensionKind, GTTypeError
+from gt4py.next.ffront.experimental import as_offset
 from gt4py.next.ffront.fbuiltins import (
     Dimension,
     Field,
@@ -684,3 +685,33 @@ def return_undefined():
 
     with pytest.raises(FieldOperatorTypeDeductionError, match="Undeclared symbol"):
         _ = FieldOperatorParser.apply_to_function(return_undefined)
+
+
+def test_as_offset_dim():
+    ADim = Dimension("ADim")
+    BDim = Dimension("BDim")
+    Boff = FieldOffset("Boff", source=BDim, target=(BDim,))
+
+    def as_offset_dim(a: Field[[ADim, BDim], float], b: Field[[ADim], int]):
+        return a(as_offset(Boff, b))
+
+    with pytest.raises(
+        FieldOperatorTypeDeductionError,
+        match=f"not in list of offset field dimensions",
+    ):
+        _ = FieldOperatorParser.apply_to_function(as_offset_dim)
+
+
+def test_as_offset_dtype():
+    ADim = Dimension("ADim")
+    BDim = Dimension("BDim")
+    Boff = FieldOffset("Boff", source=BDim, target=(BDim,))
+
+    def as_offset_dtype(a: Field[[ADim, BDim], float], b: Field[[BDim], float]):
+        return a(as_offset(Boff, b))
+
+    with pytest.raises(
+        FieldOperatorTypeDeductionError,
+        match=f"Excepted integer for offset field dtype",
+    ):
+        _ = FieldOperatorParser.apply_to_function(as_offset_dtype)