feat[next]: as_offset implementation in embedded

src/gt4py/next/common.py

@@ -499,8 +499,12 @@ def replace(self, index: int | Dimension, *named_ranges: NamedRange) -> Domain:
         if index < 0:
             index += len(self.dims)
         new_dims, new_ranges = zip(*named_ranges) if len(named_ranges) > 0 else ((), ())
-        dims = self.dims[:index] + new_dims + self.dims[index + 1 :]
-        ranges = self.ranges[:index] + new_ranges + self.ranges[index + 1 :]
+        if new_dims == self.dims:
+            dims = new_dims
+            ranges = new_ranges
+        else:
+            dims = self.dims[:index] + new_dims + self.dims[index + 1 :]
+            ranges = self.ranges[:index] + new_ranges + self.ranges[index + 1 :]
 
         return Domain(dims=dims, ranges=ranges)
 

src/gt4py/next/embedded/nd_array_field.py

@@ -27,7 +27,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:
@@ -74,6 +74,24 @@ def _builtin_op(*fields: common.Field | core_defs.Scalar) -> NdArrayField:
     return _builtin_op
 
 
+def _take_mdim(
+    input_arr: core_defs.NDArrayObject,
+    restricted_connectivity: core_defs.NDArrayObject,
+    new_domain: common.Domain,
+    dim: common.Dimension,
+) -> core_defs.NDArrayObject:
+    offset_abs = [
+        restricted_connectivity if d == dim else np.indices(restricted_connectivity.shape)[d_i]
+        for d_i, d in enumerate(new_domain.dims)
+    ]
+    new_buffer_flat = np.take(
+        np.asarray(input_arr).flatten(),
+        np.ravel_multi_index(tuple(offset_abs), input_arr.shape).flatten(),
+    )
+    new_buffer = new_buffer_flat.reshape(restricted_connectivity.shape)
+    return new_buffer
+
+
 _Value: TypeAlias = common.Field | core_defs.ScalarT
 _P = ParamSpec("_P")
 _R = TypeVar("_R", _Value, tuple[_Value, ...])
@@ -198,8 +216,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 = _take_mdim(self._ndarray, new_idx_array, new_domain, dim)
+            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, domain=new_domain, dtype=self.dtype)
 
@@ -592,6 +613,24 @@ def _astype(field: common.Field | core_defs.ScalarT | tuple, type_: type) -> NdA
 NdArrayField.register_builtin_func(fbuiltins.astype, _astype)
 
 
+def _as_offset(offset_: fbuiltins.FieldOffset, field: common.Field) -> NdArrayConnectivityField:
+    if isinstance(field, NdArrayField):
+        # change field.ndarray from relative to absolute
+        offset_dim = np.squeeze(
+            np.where(list(map(lambda x: x == offset_.source, field.domain.dims)))
+        ).item()
+        new_connectivity = np.indices(field.ndarray.shape)[offset_dim] + field.ndarray
+        return NumPyArrayConnectivityField.from_array(
+            new_connectivity, codomain=offset_.source, domain=field.domain
+        )
+    raise AssertionError(
+        "This is the NdArrayConnectivityField implementation of `experimental.as_offset`."
+    )
+
+
+NdArrayField.register_builtin_func(experimental.as_offset, _as_offset)  # type: ignore[has-type] #type specified in experimental
+
+
 def _get_slices_from_domain_slice(
     domain: common.Domain,
     domain_slice: common.Domain | Sequence[common.NamedRange | common.NamedIndex | Any],

src/gt4py/next/ffront/experimental.py

@@ -12,30 +12,23 @@
 #
 # SPDX-License-Identifier: GPL-3.0-or-later
 
-from dataclasses import dataclass
+import numpy as np
 
-from gt4py.next.type_system import type_specifications as ts
+from gt4py.next import common
+from gt4py.next.ffront.fbuiltins import BuiltInFunction, FieldOffset
 
 
-@dataclass
-class BuiltInFunction:
-    __gt_type: ts.FunctionType
+@BuiltInFunction
+def as_offset(
+    offset_: FieldOffset,
+    field: common.Field,
+    /,
+) -> common.ConnectivityField:
+    offset_dim = np.squeeze(
+        np.where(list(map(lambda x: x == offset_.source, field.domain.dims)))
+    ).item()
+    new_connectivity = np.indices(field.ndarray.shape)[offset_dim] + field.ndarray
+    return common.connectivity(new_connectivity, codomain=offset_.source, domain=field.domain)
 
-    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/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:
@@ -300,7 +303,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

@@ -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,
@@ -717,7 +718,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)

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

@@ -167,24 +167,29 @@ def get_cache_id(
     column_axis: Optional[common.Dimension],
     offset_provider: Mapping[str, Any],
 ) -> str:
-    max_neighbors = [
-        (k, v.max_neighbors)
-        for k, v in offset_provider.items()
+    def offset_invariants(offset):
         if isinstance(
-            v,
+            offset,
             (
                 itir_embedded.NeighborTableOffsetProvider,
                 itir_embedded.StridedNeighborOffsetProvider,
             ),
-        )
+        ):
+            return offset.origin_axis, offset.neighbor_axis, offset.max_neighbors
+        if isinstance(offset, common.Dimension):
+            return (offset,)
+        return tuple()
+
+    offset_cache_keys = [
+        (name, offset_invariants(offset)) for name, offset in offset_provider.items()
     ]
     cache_id_args = [
         str(arg)
         for arg in (
             program,
             *arg_types,
             column_axis,
-            *max_neighbors,
+            *offset_cache_keys,
         )
     ]
     m = hashlib.sha256()

tests/next_tests/exclusion_matrices.py

@@ -132,7 +132,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),
 ]
 GTFN_SKIP_TEST_LIST = COMMON_SKIP_TEST_LIST + [