as_offset implementation in embedded

src/gt4py/next/embedded/nd_array_field.py

@@ -28,7 +28,7 @@
 from gt4py.eve.extended_typing import Any, Never, Optional, ParamSpec, TypeAlias, TypeVar
 from gt4py.next import common
 from gt4py.next.embedded import common as embedded_common
-from gt4py.next.ffront import fbuiltins
+from gt4py.next.ffront import experimental, fbuiltins
 
 
 try:
@@ -173,13 +173,19 @@ def remap(
         assert common.is_connectivity_field(connectivity)
 
         # Compute the new domain
-        dim = connectivity.codomain
+        dim = (
+            connectivity.codomain
+            if isinstance(connectivity.codomain, common.Dimension)
+            else connectivity.codomain.source
+        )
         dim_idx = self.domain.dim_index(dim)
         if dim_idx is None:
             raise ValueError(f"Incompatible index field, expected a field with dimension {dim}.")
 
         current_range: common.UnitRange = self.domain[dim_idx][1]
         new_ranges = connectivity.inverse_image(current_range)
+        if isinstance(connectivity.codomain, fbuiltins.FieldOffset):
+            new_ranges = [new_ranges[dim_idx]]
         new_domain = self.domain.replace(dim_idx, *new_ranges)
 
         # perform contramap
@@ -201,6 +207,14 @@ def remap(
             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 len(new_idx_array.shape) > 1 and isinstance(
+                connectivity.codomain, fbuiltins.FieldOffset
+            ):
+                new_buffer = (
+                    np.diagonal(new_buffer).T
+                    if dim.kind == "horizontal"
+                    else np.diagonal(new_buffer.T)
+                )
 
         return self.__class__.from_array(new_buffer, domain=new_domain, dtype=self.dtype)
 
@@ -343,7 +357,7 @@ def from_array(  # type: ignore[override]
         cls,
         data: npt.ArrayLike | core_defs.NDArrayObject,
         /,
-        codomain: common.DimT,
+        codomain: common.DimT | fbuiltins.FieldOffset,
         *,
         domain: common.DomainLike,
         dtype: Optional[core_defs.DTypeLike] = None,
@@ -363,7 +377,7 @@ def from_array(  # type: ignore[override]
         assert len(domain) == array.ndim
         assert all(len(r) == s or s == 1 for r, s in zip(domain.ranges, array.shape))
 
-        assert isinstance(codomain, common.Dimension)
+        assert isinstance(codomain, (common.Dimension, fbuiltins.FieldOffset))
 
         return cls(domain, array, codomain)
 
@@ -586,7 +600,20 @@ def _astype(field: common.Field | core_defs.ScalarT | tuple, type_: type) -> NdA
     raise AssertionError("This is the NdArrayField implementation of `fbuiltins.astype`.")
 
 
+def _as_offset(offset_: fbuiltins.FieldOffset, field: common.Field) -> NdArrayConnectivityField:
+    if isinstance(field, NdArrayField):
+        # change field.ndarray from local to global
+        global_index_arr = np.arange(field.shape[0]) + field.ndarray
+        return NumPyArrayConnectivityField.from_array(
+            global_index_arr, codomain=offset_, domain=field.domain
+        )
+    raise AssertionError(
+        "This is the NdArrayConnectivityField implementation of `experimental.as_offset`."
+    )
+
+
 NdArrayField.register_builtin_func(fbuiltins.astype, _astype)
+NdArrayField.register_builtin_func(experimental.as_offset, _as_offset)
 
 
 def _get_slices_from_domain_slice(

src/gt4py/next/ffront/experimental.py

@@ -12,30 +12,14 @@
 #
 # 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
 
-
-@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),
-    )
-)
+@BuiltInFunction
+def as_offset(
+    offset_: FieldOffset,
+    field: common.Field,
+    /,
+) -> common.ConnectivityField:
+    raise NotImplementedError()

src/gt4py/next/ffront/fbuiltins.py

@@ -25,7 +25,6 @@
 from gt4py._core import definitions as core_defs
 from gt4py.next import common, embedded
 from gt4py.next.common import Dimension, Field  # noqa: F401  # direct import for TYPE_BUILTINS
-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
 
@@ -58,6 +57,10 @@ def _type_conversion_helper(t: type) -> type[ts.TypeSpec] | tuple[type[ts.TypeSp
         return ts.FieldType
     elif t is common.Dimension:
         return ts.DimensionType
+    elif t is FieldOffset:
+        return ts.OffsetType
+    elif t is common.ConnectivityField:
+        return ts.OffsetType
     elif t is core_defs.ScalarT:
         return ts.ScalarType
     elif t is type:
@@ -297,7 +300,6 @@ def impl(
     "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

@@ -716,7 +716,7 @@ 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 + ["as_offset"])
         ):
             visitor = getattr(self, f"_visit_{new_func.id}")
             return visitor(new_node, **kwargs)

tests/next_tests/exclusion_matrices.py

@@ -134,7 +134,6 @@ class ProgramFormatterId(_PythonObjectIdMixin, str, enum.Enum):
 ]
 EMBEDDED_SKIP_LIST = [
     (USES_SCAN, XFAIL, UNSUPPORTED_MESSAGE),
-    (USES_DYNAMIC_OFFSETS, XFAIL, UNSUPPORTED_MESSAGE),
     (CHECKS_SPECIFIC_ERROR, XFAIL, UNSUPPORTED_MESSAGE),
 ]
 GTFN_SKIP_TEST_LIST = COMMON_SKIP_TEST_LIST + [