feat[next]: as_offset implementation in embedded

src/gt4py/next/embedded/nd_array_field.py

@@ -211,8 +211,11 @@ def remap(
             # then compute the index array
             xp = self.array_ns
             new_idx_array = xp.asarray(restricted_connectivity.ndarray) - current_range.start
-            # finally, take the new array
-            new_buffer = xp.take(self._ndarray, new_idx_array, axis=dim_idx)
+            if self._ndarray.ndim > 1 and restricted_connectivity_domain == new_domain:
+                new_buffer = self._take_mdim(new_idx_array, axis=dim_idx)
+            else:
+                # finally, take the new array
+                new_buffer = xp.take(self._ndarray, new_idx_array, axis=dim_idx)
 
         return self.__class__.from_array(
             new_buffer,
@@ -324,6 +327,24 @@ def _slice(
         assert common.is_relative_index_sequence(slice_)
         return new_domain, slice_
 
+    def _take_mdim(
+        self,
+        restricted_connectivity: core_defs.NDArrayObject,
+        axis: int,
+    ) -> core_defs.NDArrayObject:
+        xp = self.array_ns
+        dim = self.domain.dims[axis]
+        offset_abs = [
+            restricted_connectivity if d == dim else xp.indices(restricted_connectivity.shape)[d_i]
+            for d_i, d in enumerate(self.domain.dims)
+        ]
+        new_buffer_flat = xp.take(
+            xp.asarray(self._ndarray).flatten(),
+            xp.ravel_multi_index(tuple(offset_abs), self._ndarray.shape).flatten(),
+        )
+        new_buffer = new_buffer_flat.reshape(restricted_connectivity.shape)
+        return new_buffer
+
 
 @dataclasses.dataclass(frozen=True)
 class NdArrayConnectivityField(  # type: ignore[misc] # for __ne__, __eq__

src/gt4py/next/ffront/experimental.py

@@ -12,30 +12,17 @@
 #
 # SPDX-License-Identifier: GPL-3.0-or-later
 
-from dataclasses import dataclass
+from gt4py.next import common
+from gt4py.next.ffront.fbuiltins import BuiltInFunction, FieldOffset
 
-from gt4py.next.type_system import type_specifications as ts
 
+@BuiltInFunction
+def as_offset(
+    offset_: FieldOffset,
+    field: common.Field,
+    /,
+) -> common.ConnectivityField:
+    raise NotImplementedError()
 
-@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(
-        pos_only_args=[
-            ts.DeferredType(constraint=ts.OffsetType),
-            ts.DeferredType(constraint=ts.FieldType),
-        ],
-        pos_or_kw_args={},
-        kw_only_args={},
-        returns=ts.DeferredType(constraint=ts.OffsetType),
-    )
-)
+EXPERIMENTAL_FUN_BUILTIN_NAMES = ["as_offset"]

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

@@ -20,6 +20,7 @@
 from gt4py.next.common import DimensionKind
 from gt4py.next.ffront import (  # noqa
     dialect_ast_enums,
+    experimental,
     fbuiltins,
     type_info as ti_ffront,
     type_specifications as ts_ffront,
@@ -727,7 +728,8 @@ def visit_Call(self, node: foast.Call, **kwargs) -> foast.Call:
             isinstance(new_func.type, ts.FunctionType)
             and not type_info.is_concrete(return_type)
             and isinstance(new_func, foast.Name)
-            and new_func.id in fbuiltins.FUN_BUILTIN_NAMES
+            and new_func.id
+            in (fbuiltins.FUN_BUILTIN_NAMES + experimental.EXPERIMENTAL_FUN_BUILTIN_NAMES)
         ):
             visitor = getattr(self, f"_visit_{new_func.id}")
             return visitor(new_node, **kwargs)

tests/next_tests/definitions.py

@@ -163,7 +163,6 @@ class ProgramFormatterId(_PythonObjectIdMixin, str, enum.Enum):
     (USES_ZERO_DIMENSIONAL_FIELDS, XFAIL, UNSUPPORTED_MESSAGE),
 ]
 EMBEDDED_SKIP_LIST = [
-    (USES_DYNAMIC_OFFSETS, XFAIL, UNSUPPORTED_MESSAGE),
     (CHECKS_SPECIFIC_ERROR, XFAIL, UNSUPPORTED_MESSAGE),
     (
         USES_SCAN_WITHOUT_FIELD_ARGS,

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

@@ -0,0 +1,210 @@
+# -*- coding: utf-8 -*-
+# 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
+
+import numpy as np
+import pytest
+
+import gt4py.next as gtx
+from gt4py.next.ffront.experimental import as_offset
+
+from next_tests.integration_tests import cases
+from next_tests.integration_tests.cases import IDim, Ioff, JDim, KDim, Koff, cartesian_case
+from next_tests.integration_tests.feature_tests.ffront_tests.ffront_test_utils import (
+    fieldview_backend,
+    reduction_setup,
+)
+
+
+@pytest.mark.uses_dynamic_offsets
+def test_offset_field_regular_input_size(cartesian_case):
+    @gtx.field_operator
+    def testee(
+        a: gtx.Field[[IDim, KDim], int], offset_field: gtx.Field[[IDim, KDim], int]
+    ) -> gtx.Field[[IDim, KDim], int]:
+        return a(as_offset(Ioff, offset_field))
+
+    out = cases.allocate(cartesian_case, testee, cases.RETURN)()
+    a = cases.allocate(cartesian_case, testee, "a")()
+    offset_field = cases.allocate(cartesian_case, testee, "offset_field").strategy(
+        cases.ConstInitializer(3)
+    )()
+    ref = a[3:]
+
+    cases.verify(
+        cartesian_case,
+        testee,
+        a,
+        offset_field,
+        out=out[: len(out.asnumpy()) - 3],
+        offset_provider={"Ioff": IDim},
+        ref=ref,
+        comparison=lambda out, ref: np.all(out == ref),
+    )
+
+
+@pytest.mark.uses_dynamic_offsets
+def test_offset_field_domain(cartesian_case):
+    @gtx.field_operator
+    def testee_fo(
+        a: gtx.Field[[IDim, KDim], int], offset_field: gtx.Field[[IDim, KDim], int]
+    ) -> gtx.Field[[IDim, KDim], int]:
+        return a(as_offset(Ioff, offset_field))
+
+    @gtx.program
+    def testee(
+        a: gtx.Field[[IDim, KDim], int],
+        offset_field: gtx.Field[[IDim, KDim], int],
+        out: gtx.Field[[IDim, KDim], int],
+    ):
+        testee_fo(a, offset_field, out=out, domain={IDim: (1, 5), KDim: (0, 10)})
+
+    out = cases.allocate(cartesian_case, testee, "out")()
+    a = cases.allocate(cartesian_case, testee, "a").extend({IDim: (0, 3)})()
+    offset_field = cases.allocate(cartesian_case, testee, "offset_field").strategy(
+        cases.ConstInitializer(3)
+    )()
+    ref = out.asnumpy().copy()
+    a_shift = a[3:]
+    ref[1:5] = a_shift.asnumpy()[1:5]
+
+    cases.verify(
+        cartesian_case,
+        testee,
+        a,
+        offset_field,
+        out,
+        inout=out,
+        offset_provider={"Ioff": IDim},
+        ref=ref,
+        comparison=lambda out, ref: np.all(out == ref),
+    )
+
+
+@pytest.mark.uses_dynamic_offsets
+def test_offset_field_input_domain(cartesian_case):
+    i_size = cartesian_case.default_sizes[IDim]
+    k_size = cartesian_case.default_sizes[KDim]
+
+    @gtx.field_operator
+    def testee_fo(
+        a: gtx.Field[[IDim, KDim], int], offset_field: gtx.Field[[IDim, KDim], int]
+    ) -> gtx.Field[[IDim, KDim], int]:
+        return a(as_offset(Ioff, offset_field))
+
+    @gtx.program
+    def testee(
+        a: gtx.Field[[IDim, KDim], int],
+        offset_field: gtx.Field[[IDim, KDim], int],
+        out: gtx.Field[[IDim, KDim], int],
+    ):
+        testee_fo(a, offset_field, out=out)
+
+    a_fo = cases.allocate(cartesian_case, testee, "a").extend({IDim: (0, 3)})().ndarray
+    out = cartesian_case.as_field(
+        [IDim, KDim],
+        np.zeros([i_size, k_size], dtype=int),
+        origin={IDim: 2, KDim: 0},
+    )
+    a = cartesian_case.as_field([IDim, KDim], a_fo, origin={IDim: 2, KDim: 0})
+    offset_field = cases.allocate(cartesian_case, testee, "offset_field").strategy(
+        cases.ConstInitializer(3)
+    )()
+    ref = a[5:]
+    cases.verify(
+        cartesian_case,
+        testee,
+        a,
+        offset_field,
+        out[2:],
+        inout=out[2:],
+        offset_provider={"Ioff": IDim},
+        ref=ref,
+    )
+
+
+@pytest.mark.parametrize(
+    "dims, offset, offset_dim, off_dim, other_dim",
+    [
+        ([IDim, KDim], Ioff, IDim, 0, 1),
+        ([IDim, JDim], Ioff, IDim, 0, 1),
+        ([IDim, KDim], Koff, KDim, 1, 0),
+    ],
+)
+@pytest.mark.uses_dynamic_offsets
+def test_offset_field_rand_connectivity(
+    cartesian_case, dims, offset, offset_dim, off_dim, other_dim
+):
+    @gtx.field_operator
+    def testee(
+        a: gtx.Field[[dims[0], dims[1]], int], offset_field: gtx.Field[[dims[0], dims[1]], int]
+    ) -> gtx.Field[[dims[0], dims[1]], int]:
+        return a(as_offset(offset, offset_field))
+
+    out = cases.allocate(cartesian_case, testee, cases.RETURN)()
+    a = cases.allocate(cartesian_case, testee, "a").extend({offset_dim: (0, 3)})()
+    offset_field = cartesian_case.as_field(
+        [dims[0], dims[1]],
+        np.random.randint(
+            0,
+            4,
+            size=(cartesian_case.default_sizes[dims[0]], cartesian_case.default_sizes[dims[1]]),
+        ),
+    )
+    ref = np.zeros(offset_field.asnumpy().shape, dtype=int)
+
+    for j in range(len(offset_field.asnumpy()[off_dim])):
+        for i in range(len(offset_field.asnumpy())):
+            if offset_dim.kind == "vertical":
+                y = offset_field.asnumpy()[i][j] + np.indices(offset_field.shape)[off_dim][i][j]
+                x = np.indices(offset_field.shape)[other_dim][i][j]
+            else:
+                x = offset_field.asnumpy()[i][j] + np.indices(offset_field.shape)[off_dim][i][j]
+                y = np.indices(offset_field.shape)[other_dim][i][j]
+            ref[i][j] = a[x][y]
+
+    cases.verify(
+        cartesian_case,
+        testee,
+        a,
+        offset_field,
+        out=out,
+        offset_provider={"offset": offset_dim, "Ioff": IDim, "Koff": KDim},
+        ref=ref,
+        comparison=lambda out, ref: np.all(out == ref),
+    )
+
+
+@pytest.mark.uses_dynamic_offsets
+def test_offset_field_3d(cartesian_case):
+    @gtx.field_operator
+    def testee(a: cases.IJKField, offset_field: cases.IJKField) -> cases.IJKField:
+        return a(as_offset(Ioff, offset_field))
+
+    out = cases.allocate(cartesian_case, testee, cases.RETURN)()
+    a = cases.allocate(cartesian_case, testee, "a").extend({IDim: (0, 2)})()
+    offset_field = cases.allocate(cartesian_case, testee, "offset_field").strategy(
+        cases.ConstInitializer(2)
+    )()
+
+    cases.verify(
+        cartesian_case,
+        testee,
+        a,
+        offset_field,
+        out=out,
+        offset_provider={"Ioff": IDim, "Koff": KDim},
+        ref=a[2:],
+        comparison=lambda out, ref: np.all(out == ref),
+    )

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

@@ -32,7 +32,6 @@
     neighbor_sum,
     where,
 )
-from gt4py.next.ffront.experimental import as_offset
 from gt4py.next.program_processors.runners import gtfn
 
 from next_tests.integration_tests import cases
@@ -462,40 +461,6 @@ def testee(a: cases.IFloatField) -> gtx.Field[[IDim], np.float32]:
     )
 
 
-@pytest.mark.uses_dynamic_offsets
-def test_offset_field(cartesian_case):
-    ref = np.full(
-        (cartesian_case.default_sizes[IDim], cartesian_case.default_sizes[KDim]), True, dtype=bool
-    )
-
-    @gtx.field_operator
-    def testee(a: cases.IKField, offset_field: cases.IKField) -> gtx.Field[[IDim, KDim], bool]:
-        a_i = a(as_offset(Ioff, offset_field))
-        a_i_k = a_i(as_offset(Koff, offset_field))
-        b_i = a(Ioff[1])
-        b_i_k = b_i(Koff[1])
-        return a_i_k == b_i_k
-
-    out = cases.allocate(cartesian_case, testee, cases.RETURN)()
-    a = cases.allocate(cartesian_case, testee, "a").extend({IDim: (0, 1), KDim: (0, 1)})()
-    offset_field = cases.allocate(cartesian_case, testee, "offset_field").strategy(
-        cases.ConstInitializer(1)
-    )()
-
-    cases.verify(
-        cartesian_case,
-        testee,
-        a,
-        offset_field,
-        out=out,
-        offset_provider={"Ioff": IDim, "Koff": KDim},
-        ref=np.full_like(offset_field.asnumpy(), True, dtype=bool),
-        comparison=lambda out, ref: np.all(out == ref),
-    )
-
-    assert np.allclose(out.asnumpy(), ref)
-
-
 def test_nested_tuple_return(cartesian_case):
     @gtx.field_operator
     def pack_tuple(