Repository: GridTools/gt4py
Pull request #1521 (2024-04-04 10:58:35)

refactor[next]: itir embedded: cleaner closure run
This enables to run the closure also outside of the fendef context, e.g. from field view embedded.


Patch:
diff --git a/constraints.txt b/constraints.txt
index 9488949c..fe4f777a 100644
--- a/constraints.txt
+++ b/constraints.txt
@@ -103,7 +103,7 @@ mypy==1.9.0               # via -r requirements-dev.in
 mypy-extensions==1.0.0    # via black, mypy
 nanobind==1.9.2           # via gt4py (pyproject.toml)
 nbclient==0.6.8           # via nbmake
-nbformat==5.10.3          # via jupytext, nbclient, nbmake
+nbformat==5.10.4          # via jupytext, nbclient, nbmake
 nbmake==1.5.3             # via -r requirements-dev.in
 nest-asyncio==1.6.0       # via ipykernel, nbclient
 networkx==3.1             # via dace
diff --git a/pyproject.toml b/pyproject.toml
index 0d5a2973..3ee66c94 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -48,7 +48,7 @@ dependencies = [
   'pybind11>=2.10.1',
   'setuptools>=65.5.0',
   'tabulate>=0.8.10',
-  'typing-extensions>=4.6.0',
+  'typing-extensions>=4.10.0',
   'xxhash>=1.4.4,<3.1.0'
 ]
 description = 'Python library for generating high-performance implementations of stencil kernels for weather and climate modeling from a domain-specific language (DSL)'
diff --git a/requirements-dev.txt b/requirements-dev.txt
index df79f905..71a41525 100644
--- a/requirements-dev.txt
+++ b/requirements-dev.txt
@@ -103,7 +103,7 @@ mypy==1.9.0               # via -c constraints.txt, -r requirements-dev.in
 mypy-extensions==1.0.0    # via -c constraints.txt, black, mypy
 nanobind==1.9.2           # via -c constraints.txt, gt4py (pyproject.toml)
 nbclient==0.6.8           # via -c constraints.txt, nbmake
-nbformat==5.10.3          # via -c constraints.txt, jupytext, nbclient, nbmake
+nbformat==5.10.4          # via -c constraints.txt, jupytext, nbclient, nbmake
 nbmake==1.5.3             # via -c constraints.txt, -r requirements-dev.in
 nest-asyncio==1.6.0       # via -c constraints.txt, ipykernel, nbclient
 networkx==3.1             # via -c constraints.txt, dace
diff --git a/src/gt4py/next/embedded/context.py b/src/gt4py/next/embedded/context.py
index 672dc9c6..be5ca39e 100644
--- a/src/gt4py/next/embedded/context.py
+++ b/src/gt4py/next/embedded/context.py
@@ -27,12 +27,8 @@ import gt4py.next.common as common
 #: closure execution context.
 closure_column_range: cvars.ContextVar[common.NamedRange] = cvars.ContextVar("column_range")
 
-_undefined_offset_provider: common.OffsetProvider = {}
-
 #: Offset provider dict in the current embedded execution context.
-offset_provider: cvars.ContextVar[common.OffsetProvider] = cvars.ContextVar(
-    "offset_provider", default=_undefined_offset_provider
-)
+offset_provider: cvars.ContextVar[common.OffsetProvider] = cvars.ContextVar("offset_provider")
 
 
 @contextlib.contextmanager
@@ -41,6 +37,8 @@ def new_context(
     closure_column_range: common.NamedRange | eve.NothingType = eve.NOTHING,
     offset_provider: common.OffsetProvider | eve.NothingType = eve.NOTHING,
 ) -> Generator[cvars.Context, None, None]:
+    """Create a new context, updating the provided values."""
+
     import gt4py.next.embedded.context as this_module
 
     updates: list[tuple[cvars.ContextVar[Any], Any]] = []
@@ -61,5 +59,5 @@ def new_context(
     yield ctx
 
 
-def within_context() -> bool:
-    return offset_provider.get() is not _undefined_offset_provider
+def within_valid_context() -> bool:
+    return offset_provider.get(eve.NOTHING) is not eve.NOTHING
diff --git a/src/gt4py/next/embedded/operators.py b/src/gt4py/next/embedded/operators.py
index be1768b4..dd197ac2 100644
--- a/src/gt4py/next/embedded/operators.py
+++ b/src/gt4py/next/embedded/operators.py
@@ -98,7 +98,7 @@ def field_operator_call(op: EmbeddedOperator[_R, _P], args: Any, kwargs: Any) ->
     if "out" in kwargs:
         # called from program or direct field_operator as program
         new_context_kwargs = {}
-        if embedded_context.within_context():
+        if embedded_context.within_valid_context():
             # called from program
             assert "offset_provider" not in kwargs
         else:
diff --git a/src/gt4py/next/ffront/decorator.py b/src/gt4py/next/ffront/decorator.py
index f55e8297..a01cf095 100644
--- a/src/gt4py/next/ffront/decorator.py
+++ b/src/gt4py/next/ffront/decorator.py
@@ -208,7 +208,7 @@ class Program:
             kwargs={},
         )
         if self.backend is not None and self.backend.transformer is not None:
-            return self.backend.transformer.past_to_itir(no_args_past)
+            return self.backend.transformer.past_to_itir(no_args_past).program
         return past_to_itir.PastToItirFactory()(no_args_past).program
 
     def __call__(self, *args, offset_provider: dict[str, Dimension], **kwargs: Any) -> None:
@@ -545,7 +545,7 @@ class FieldOperator(GTCallable, Generic[OperatorNodeT]):
         return self._program_cache[hash_]
 
     def __call__(self, *args, **kwargs) -> None:
-        if not next_embedded.context.within_context() and self.backend is not None:
+        if not next_embedded.context.within_valid_context() and self.backend is not None:
             # non embedded execution
             if "offset_provider" not in kwargs:
                 raise errors.MissingArgumentError(None, "offset_provider", True)
diff --git a/src/gt4py/next/iterator/embedded.py b/src/gt4py/next/iterator/embedded.py
index 5c6976ea..f5d4c6e5 100644
--- a/src/gt4py/next/iterator/embedded.py
+++ b/src/gt4py/next/iterator/embedded.py
@@ -17,7 +17,6 @@
 from __future__ import annotations
 
 import abc
-import contextvars as cvars
 import copy
 import dataclasses
 import itertools
@@ -28,6 +27,7 @@ import warnings
 import numpy as np
 import numpy.typing as npt
 
+from gt4py import eve
 from gt4py._core import definitions as core_defs
 from gt4py.eve import extended_typing as xtyping
 from gt4py.eve.extended_typing import (
@@ -52,8 +52,8 @@ from gt4py.eve.extended_typing import (
     overload,
     runtime_checkable,
 )
-from gt4py.next import common, embedded as next_embedded
-from gt4py.next.embedded import exceptions as embedded_exceptions
+from gt4py.next import common
+from gt4py.next.embedded import context as embedded_context, exceptions as embedded_exceptions
 from gt4py.next.ffront import fbuiltins
 from gt4py.next.iterator import builtins, runtime
 
@@ -191,12 +191,6 @@ class MutableLocatedField(LocatedField, Protocol):
     def field_setitem(self, indices: NamedFieldIndices, value: Any) -> None: ...
 
 
-#: Column range used in column mode (`column_axis != None`) in the current closure execution context.
-column_range_cvar: cvars.ContextVar[common.NamedRange] = next_embedded.context.closure_column_range
-#: Offset provider dict in the current closure execution context.
-offset_provider_cvar: cvars.ContextVar[OffsetProvider] = next_embedded.context.offset_provider
-
-
 class Column(np.lib.mixins.NDArrayOperatorsMixin):
     """Represents a column when executed in column mode (`column_axis != None`).
 
@@ -207,7 +201,7 @@ class Column(np.lib.mixins.NDArrayOperatorsMixin):
     def __init__(self, kstart: int, data: np.ndarray | Scalar) -> None:
         self.kstart = kstart
         assert isinstance(data, (np.ndarray, Scalar))  # type: ignore # mypy bug #11673
-        column_range: common.NamedRange = column_range_cvar.get()
+        column_range: common.NamedRange = embedded_context.closure_column_range.get()
         self.data = (
             data if isinstance(data, np.ndarray) else np.full(len(column_range.unit_range), data)
         )
@@ -751,7 +745,7 @@ def _make_tuple(
             except embedded_exceptions.IndexOutOfBounds:
                 return _UNDEFINED
     else:
-        column_range = column_range_cvar.get().unit_range
+        column_range = embedded_context.closure_column_range.get().unit_range
         assert column_range is not None
 
         col: list[
@@ -796,7 +790,7 @@ class MDIterator:
 
     def shift(self, *offsets: OffsetPart) -> MDIterator:
         complete_offsets = group_offsets(*offsets)
-        offset_provider = offset_provider_cvar.get()
+        offset_provider = embedded_context.offset_provider.get()
         assert offset_provider is not None
         return MDIterator(
             self.field,
@@ -821,8 +815,8 @@ class MDIterator:
             if not all(axis.value in shifted_pos.keys() for axis in axes if axis is not None):
                 raise IndexError("Iterator position doesn't point to valid location for its field.")
         slice_column = dict[Tag, range]()
-        column_range = column_range_cvar.get()
         if self.column_axis is not None:
+            column_range = embedded_context.closure_column_range.get()
             assert column_range is not None
             k_pos = shifted_pos.pop(self.column_axis)
             assert isinstance(k_pos, int)
@@ -862,7 +856,7 @@ def make_in_iterator(
         init = [None] * sparse_dimensions.count(sparse_dim)
         new_pos[sparse_dim] = init  # type: ignore[assignment] # looks like mypy is confused
     if column_axis is not None:
-        column_range = column_range_cvar.get().unit_range
+        column_range = embedded_context.closure_column_range.get().unit_range
         # if we deal with column stencil the column position is just an offset by which the whole column needs to be shifted
         assert column_range is not None
         new_pos[column_axis] = column_range.start
@@ -1303,7 +1297,7 @@ class _ConstList(Generic[DT]):
 def neighbors(offset: runtime.Offset, it: ItIterator) -> _List:
     offset_str = offset.value if isinstance(offset, runtime.Offset) else offset
     assert isinstance(offset_str, str)
-    offset_provider = offset_provider_cvar.get()
+    offset_provider = embedded_context.offset_provider.get()
     assert offset_provider is not None
     connectivity = offset_provider[offset_str]
     assert isinstance(connectivity, common.Connectivity)
@@ -1359,7 +1353,7 @@ class SparseListIterator:
     offsets: Sequence[OffsetPart] = dataclasses.field(default_factory=list, kw_only=True)
 
     def deref(self) -> Any:
-        offset_provider = offset_provider_cvar.get()
+        offset_provider = embedded_context.offset_provider.get()
         assert offset_provider is not None
         connectivity = offset_provider[self.list_offset]
         assert isinstance(connectivity, common.Connectivity)
@@ -1376,12 +1370,6 @@ class SparseListIterator:
         return SparseListIterator(self.it, self.list_offset, offsets=[*offsets, *self.offsets])
 
 
-@dataclasses.dataclass(frozen=True)
-class ColumnDescriptor:
-    axis: str
-    col_range: range  # TODO(havogt) introduce range type that doesn't have step
-
-
 @dataclasses.dataclass(frozen=True)
 class ScanArgIterator:
     wrapped_iter: ItIterator
@@ -1480,7 +1468,7 @@ def _column_dtype(elem: Any) -> np.dtype:
 @builtins.scan.register(EMBEDDED)
 def scan(scan_pass, is_forward: bool, init):
     def impl(*iters: ItIterator):
-        column_range = column_range_cvar.get().unit_range
+        column_range = embedded_context.closure_column_range.get().unit_range
         if column_range is None:
             raise RuntimeError("Column range is not defined, cannot scan.")
 
@@ -1508,64 +1496,78 @@ def _validate_domain(domain: Domain, offset_provider: OffsetProvider) -> None:
             )
 
 
-def fendef_embedded(fun: Callable[..., None], *args: Any, **kwargs: Any):
-    if "offset_provider" not in kwargs:
-        raise RuntimeError("'offset_provider' not provided.")
-
-    offset_provider = kwargs["offset_provider"]
-
-    @runtime.closure.register(EMBEDDED)
-    def closure(
-        domain_: Domain,
-        sten: Callable[..., Any],
-        out,  #: MutableLocatedField,
-        ins: list[common.Field],
-    ) -> None:
-        _validate_domain(domain_, kwargs["offset_provider"])
-        domain: dict[Tag, range] = _dimension_to_tag(domain_)
-        if not (isinstance(out, common.Field) or is_tuple_of_field(out)):
-            raise TypeError("'Out' needs to be a located field.")
-
-        column_range = None
-        column: Optional[ColumnDescriptor] = None
-        if kwargs.get("column_axis") and kwargs["column_axis"].value in domain:
-            column_axis = kwargs["column_axis"]
-            column = ColumnDescriptor(column_axis.value, domain[column_axis.value])
-            del domain[column_axis.value]
-
+@runtime.closure.register(EMBEDDED)
+def closure(
+    domain_: Domain,
+    sten: Callable[..., Any],
+    out,  #: MutableLocatedField,
+    ins: list[common.Field],
+) -> None:
+    assert embedded_context.within_valid_context()
+    offset_provider = embedded_context.offset_provider.get()
+    _validate_domain(domain_, offset_provider)
+    domain: dict[Tag, range] = _dimension_to_tag(domain_)
+    if not (isinstance(out, common.Field) or is_tuple_of_field(out)):
+        raise TypeError("'Out' needs to be a located field.")
+
+    column_range: common.NamedRange | eve.NothingType = eve.NOTHING
+    if (col_range_placeholder := embedded_context.closure_column_range.get(None)) is not None:
+        assert (
+            col_range_placeholder.unit_range.is_empty()
+        )  # check it's just the placeholder with empty range
+        column_axis = col_range_placeholder.dim
+        if column_axis is not None and column_axis.value in domain:
             column_range = common.NamedRange(
-                column_axis, common.UnitRange(column.col_range.start, column.col_range.stop)
+                column_axis,
+                common.UnitRange(domain[column_axis.value].start, domain[column_axis.value].stop),
             )
+            del domain[column_axis.value]
 
-        out = as_tuple_field(out) if is_tuple_of_field(out) else _wrap_field(out)
+    out = as_tuple_field(out) if is_tuple_of_field(out) else _wrap_field(out)
 
-        def _closure_runner():
-            # Set context variables before executing the closure
-            column_range_cvar.set(column_range)
-            offset_provider_cvar.set(offset_provider)
+    with embedded_context.new_context(closure_column_range=column_range) as ctx:
 
+        def _iterate():
             for pos in _domain_iterator(domain):
                 promoted_ins = [promote_scalars(inp) for inp in ins]
                 ins_iters = list(
-                    make_in_iterator(inp, pos, column_axis=column.axis if column else None)
+                    make_in_iterator(
+                        inp,
+                        pos,
+                        column_axis=column_range.dim.value
+                        if column_range is not eve.NOTHING
+                        else None,
+                    )
                     for inp in promoted_ins
                 )
                 res = sten(*ins_iters)
 
-                if column is None:
+                if column_range is eve.NOTHING:
                     assert _is_concrete_position(pos)
                     out.field_setitem(pos, res)
                 else:
                     col_pos = pos.copy()
-                    for k in column.col_range:
-                        col_pos[column.axis] = k
+                    for k in column_range.unit_range:
+                        col_pos[column_range.dim.value] = k
                         assert _is_concrete_position(col_pos)
                         out.field_setitem(col_pos, res[k])
 
-        ctx = cvars.copy_context()
-        ctx.run(_closure_runner)
+        ctx.run(_iterate)
+
+
+def fendef_embedded(fun: Callable[..., None], *args: Any, **kwargs: Any):
+    if "offset_provider" not in kwargs:
+        raise RuntimeError("'offset_provider' not provided.")
+
+    context_vars = {"offset_provider": kwargs["offset_provider"]}
+    if "column_axis" in kwargs:
+        context_vars["closure_column_range"] = common.NamedRange(
+            kwargs["column_axis"],
+            common.UnitRange(0, 0),  # empty: indicates column operation, will update later
+        )
 
-    fun(*args)
+    with embedded_context.new_context(**context_vars) as ctx:
+        ctx.run(fun, *args)
 
 
 runtime.fendef_embedded = fendef_embedded
diff --git a/src/gt4py/next/program_processors/runners/dace.py b/src/gt4py/next/program_processors/runners/dace.py
index 2920b3d8..266d7b35 100644
--- a/src/gt4py/next/program_processors/runners/dace.py
+++ b/src/gt4py/next/program_processors/runners/dace.py
@@ -83,5 +83,11 @@ class DaCeBackendFactory(GTFNBackendFactory):
 
 
 run_dace_cpu = DaCeBackendFactory(cached=True, auto_optimize=True)
+run_dace_cpu_with_temporaries = DaCeBackendFactory(
+    cached=True, auto_optimize=True, use_temporaries=True
+)
 
 run_dace_gpu = DaCeBackendFactory(gpu=True, cached=True, auto_optimize=True)
+run_dace_gpu_with_temporaries = DaCeBackendFactory(
+    gpu=True, cached=True, auto_optimize=True, use_temporaries=True
+)
diff --git a/src/gt4py/next/program_processors/runners/dace_iterator/__init__.py b/src/gt4py/next/program_processors/runners/dace_iterator/__init__.py
index 7456aa1c..2b56dc04 100644
--- a/src/gt4py/next/program_processors/runners/dace_iterator/__init__.py
+++ b/src/gt4py/next/program_processors/runners/dace_iterator/__init__.py
@@ -14,7 +14,7 @@
 import warnings
 from inspect import currentframe, getframeinfo
 from pathlib import Path
-from typing import Any, Mapping, Optional, Sequence
+from typing import Any, Callable, Mapping, Optional, Sequence
 
 import dace
 import numpy as np
@@ -24,10 +24,10 @@ from dace.transformation.auto import auto_optimize as autoopt
 import gt4py.next.iterator.ir as itir
 from gt4py.next import common
 from gt4py.next.iterator import transforms as itir_transforms
-from gt4py.next.type_system import type_translation
+from gt4py.next.type_system import type_specifications as ts
 
 from .itir_to_sdfg import ItirToSDFG
-from .utility import connectivity_identifier, filter_neighbor_tables, get_sorted_dims
+from .utility import connectivity_identifier, filter_connectivities, get_sorted_dims
 
 
 try:
@@ -67,6 +67,10 @@ def preprocess_program(
     program: itir.FencilDefinition,
     offset_provider: Mapping[str, Any],
     lift_mode: itir_transforms.LiftMode,
+    symbolic_domain_sizes: Optional[dict[str, str]] = None,
+    temporary_extraction_heuristics: Optional[
+        Callable[[itir.StencilClosure], Callable[[itir.Expr], bool]]
+    ] = None,
     unroll_reduce: bool = False,
 ):
     node = itir_transforms.apply_common_transforms(
@@ -75,6 +79,8 @@ def preprocess_program(
         force_inline_lambda_args=True,
         lift_mode=lift_mode,
         offset_provider=offset_provider,
+        symbolic_domain_sizes=symbolic_domain_sizes,
+        temporary_extraction_heuristics=temporary_extraction_heuristics,
         unroll_reduce=unroll_reduce,
     )
 
@@ -180,7 +186,10 @@ def get_sdfg_args(
     """
     offset_provider = kwargs["offset_provider"]
 
-    neighbor_tables = filter_neighbor_tables(offset_provider)
+    neighbor_tables: dict[str, common.NeighborTable] = {}
+    for offset, connectivity in filter_connectivities(offset_provider).items():
+        assert isinstance(connectivity, common.NeighborTable)
+        neighbor_tables[offset] = connectivity
     device = dace.DeviceType.GPU if on_gpu else dace.DeviceType.CPU
 
     dace_args = get_args(sdfg, args, use_field_canonical_representation)
@@ -211,12 +220,16 @@ def get_sdfg_args(
 
 def build_sdfg_from_itir(
     program: itir.FencilDefinition,
-    *args,
+    arg_types: list[ts.TypeSpec],
     offset_provider: dict[str, Any],
     auto_optimize: bool = False,
     on_gpu: bool = False,
     column_axis: Optional[common.Dimension] = None,
     lift_mode: itir_transforms.LiftMode = itir_transforms.LiftMode.FORCE_INLINE,
+    symbolic_domain_sizes: Optional[dict[str, str]] = None,
+    temporary_extraction_heuristics: Optional[
+        Callable[[itir.StencilClosure], Callable[[itir.Expr], bool]]
+    ] = None,
     load_sdfg_from_file: bool = False,
     save_sdfg: bool = True,
     use_field_canonical_representation: bool = True,
@@ -225,12 +238,13 @@ def build_sdfg_from_itir(
 
     Args:
         program:             The Fencil that should be translated.
-        *args:               Arguments for which the fencil should be called.
+        arg_types:           Types of the arguments passed to the fencil.
         offset_provider:     The set of offset providers that should be used.
         auto_optimize:       Apply DaCe's `auto_optimize` heuristic.
         on_gpu:              Performs the translation for GPU, defaults to `False`.
         column_axis:         The column axis to be used, defaults to `None`.
         lift_mode:           Which lift mode should be used, defaults `FORCE_INLINE`.
+        symbolic_domain_sizes: Used for generation of liskov bindings when temporaries are enabled.
         load_sdfg_from_file: Allows to read the SDFG from file, instead of generating it, for debug only.
         save_sdfg:           If `True`, the default the SDFG is stored as a file and can be loaded, this allows to skip the lowering step, requires `load_sdfg_from_file` set to `True`.
         use_field_canonical_representation: If `True`,  assume that the fields dimensions are sorted alphabetically.
@@ -245,10 +259,10 @@ def build_sdfg_from_itir(
         sdfg.validate()
         return sdfg
 
-    arg_types = [type_translation.from_value(arg) for arg in args]
-
     # visit ITIR and generate SDFG
-    program, tmps = preprocess_program(program, offset_provider, lift_mode)
+    program, tmps = preprocess_program(
+        program, offset_provider, lift_mode, symbolic_domain_sizes, temporary_extraction_heuristics
+    )
     sdfg_genenerator = ItirToSDFG(
         arg_types, offset_provider, tmps, use_field_canonical_representation, column_axis
     )
diff --git a/src/gt4py/next/program_processors/runners/dace_iterator/itir_to_sdfg.py b/src/gt4py/next/program_processors/runners/dace_iterator/itir_to_sdfg.py
index f0cfad5f..39a5440e 100644
--- a/src/gt4py/next/program_processors/runners/dace_iterator/itir_to_sdfg.py
+++ b/src/gt4py/next/program_processors/runners/dace_iterator/itir_to_sdfg.py
@@ -19,14 +19,14 @@ from dace.sdfg.state import LoopRegion
 
 import gt4py.eve as eve
 from gt4py.next import Dimension, DimensionKind, type_inference as next_typing
-from gt4py.next.common import NeighborTable
+from gt4py.next.common import Connectivity
 from gt4py.next.iterator import (
     ir as itir,
     transforms as itir_transforms,
     type_inference as itir_typing,
 )
 from gt4py.next.iterator.ir import Expr, FunCall, Literal, Sym, SymRef
-from gt4py.next.type_system import type_specifications as ts, type_translation
+from gt4py.next.type_system import type_specifications as ts, type_translation as tt
 
 from .itir_to_tasklet import (
     Context,
@@ -44,10 +44,10 @@ from .utility import (
     as_scalar_type,
     connectivity_identifier,
     dace_debuginfo,
-    filter_neighbor_tables,
+    filter_connectivities,
     flatten_list,
     get_sorted_dims,
-    get_used_neighbor_tables,
+    get_used_connectivities,
     map_nested_sdfg_symbols,
     new_array_symbols,
     unique_name,
@@ -119,7 +119,7 @@ def _make_array_shape_and_strides(
     """
     dtype = dace.int32
     sorted_dims = get_sorted_dims(dims) if sort_dims else list(enumerate(dims))
-    neighbor_tables = filter_neighbor_tables(offset_provider)
+    neighbor_tables = filter_connectivities(offset_provider)
     shape = [
         (
             neighbor_tables[dim.value].max_neighbors
@@ -163,7 +163,7 @@ class ItirToSDFG(eve.NodeVisitor):
     def __init__(
         self,
         param_types: list[ts.TypeSpec],
-        offset_provider: dict[str, NeighborTable],
+        offset_provider: dict[str, Connectivity | Dimension],
         tmps: list[itir_transforms.global_tmps.Temporary],
         use_field_canonical_representation: bool,
         column_axis: Optional[Dimension] = None,
@@ -292,7 +292,7 @@ class ItirToSDFG(eve.NodeVisitor):
         self.node_types = itir_typing.infer_all(node)
 
         # Filter neighbor tables from offset providers.
-        neighbor_tables = get_used_neighbor_tables(node, self.offset_provider)
+        neighbor_tables = get_used_connectivities(node, self.offset_provider)
 
         # Add program parameters as SDFG storages.
         for param, type_ in zip(node.params, self.param_types):
@@ -312,7 +312,7 @@ class ItirToSDFG(eve.NodeVisitor):
 
         # Add connectivities as SDFG storages.
         for offset, offset_provider in neighbor_tables.items():
-            scalar_kind = type_translation.get_scalar_kind(offset_provider.table.dtype)
+            scalar_kind = tt.get_scalar_kind(offset_provider.index_type)
             local_dim = Dimension(offset, kind=DimensionKind.LOCAL)
             type_ = ts.FieldType(
                 [offset_provider.origin_axis, local_dim], ts.ScalarType(scalar_kind)
@@ -382,7 +382,7 @@ class ItirToSDFG(eve.NodeVisitor):
         closure_init_state = closure_sdfg.add_state_before(closure_state, "closure_init", True)
 
         input_names = [str(inp.id) for inp in node.inputs]
-        neighbor_tables = get_used_neighbor_tables(node, self.offset_provider)
+        neighbor_tables = get_used_connectivities(node, self.offset_provider)
         connectivity_names = [connectivity_identifier(offset) for offset in neighbor_tables.keys()]
 
         output_nodes = self.get_output_nodes(node, closure_sdfg, closure_state)
@@ -574,7 +574,7 @@ class ItirToSDFG(eve.NodeVisitor):
         )
 
         assert isinstance(node.output, SymRef)
-        neighbor_tables = get_used_neighbor_tables(node, self.offset_provider)
+        neighbor_tables = get_used_connectivities(node, self.offset_provider)
         input_names = [str(inp.id) for inp in node.inputs]
         connectivity_names = [connectivity_identifier(offset) for offset in neighbor_tables.keys()]
 
@@ -737,7 +737,7 @@ class ItirToSDFG(eve.NodeVisitor):
             tuple[str, tuple[ValueExpr | SymbolExpr, ValueExpr | SymbolExpr]], ...
         ],
     ) -> tuple[dace.SDFG, dict[str, str | dace.subsets.Subset], list[str]]:
-        neighbor_tables = get_used_neighbor_tables(node, self.offset_provider)
+        neighbor_tables = get_used_connectivities(node, self.offset_provider)
         input_names = [str(inp.id) for inp in node.inputs]
         connectivity_names = [connectivity_identifier(offset) for offset in neighbor_tables.keys()]
 
diff --git a/src/gt4py/next/program_processors/runners/dace_iterator/itir_to_tasklet.py b/src/gt4py/next/program_processors/runners/dace_iterator/itir_to_tasklet.py
index f4005936..778846b2 100644
--- a/src/gt4py/next/program_processors/runners/dace_iterator/itir_to_tasklet.py
+++ b/src/gt4py/next/program_processors/runners/dace_iterator/itir_to_tasklet.py
@@ -22,10 +22,9 @@ import numpy as np
 
 import gt4py.eve.codegen
 from gt4py import eve
-from gt4py.next import Dimension, StridedNeighborOffsetProvider, type_inference as next_typing
-from gt4py.next.common import _DEFAULT_SKIP_VALUE as neighbor_skip_value
+from gt4py.next import Dimension, type_inference as next_typing
+from gt4py.next.common import _DEFAULT_SKIP_VALUE as neighbor_skip_value, Connectivity
 from gt4py.next.iterator import ir as itir, type_inference as itir_typing
-from gt4py.next.iterator.embedded import NeighborTableOffsetProvider
 from gt4py.next.iterator.ir import FunCall, Lambda
 from gt4py.next.iterator.type_inference import Val
 from gt4py.next.type_system import type_specifications as ts
@@ -36,7 +35,7 @@ from .utility import (
     connectivity_identifier,
     dace_debuginfo,
     flatten_list,
-    get_used_neighbor_tables,
+    get_used_connectivities,
     map_nested_sdfg_symbols,
     new_array_symbols,
     unique_name,
@@ -184,7 +183,7 @@ def _visit_lift_in_neighbors_reduction(
     transformer: "PythonTaskletCodegen",
     node: itir.FunCall,
     node_args: Sequence[IteratorExpr | list[ValueExpr]],
-    offset_provider: NeighborTableOffsetProvider,
+    offset_provider: Connectivity,
     map_entry: dace.nodes.MapEntry,
     map_exit: dace.nodes.MapExit,
     neighbor_index_node: dace.nodes.AccessNode,
@@ -229,7 +228,7 @@ def _visit_lift_in_neighbors_reduction(
             assert isinstance(y, ValueExpr)
             input_nodes[x] = y.value
 
-    neighbor_tables = get_used_neighbor_tables(node.args[0], transformer.offset_provider)
+    neighbor_tables = get_used_connectivities(node.args[0], transformer.offset_provider)
     connectivity_names = [connectivity_identifier(offset) for offset in neighbor_tables.keys()]
 
     parent_sdfg = transformer.context.body
@@ -328,7 +327,7 @@ def builtin_neighbors(
     offset_dim = offset_literal.value
     assert isinstance(offset_dim, str)
     offset_provider = transformer.offset_provider[offset_dim]
-    if not isinstance(offset_provider, NeighborTableOffsetProvider):
+    if not isinstance(offset_provider, Connectivity):
         raise NotImplementedError(
             "Neighbor reduction only implemented for connectivity based on neighbor tables."
         )
@@ -917,7 +916,7 @@ class PythonTaskletCodegen(gt4py.eve.codegen.TemplatedGenerator):
     ]:
         func_name = f"lambda_{abs(hash(node)):x}"
         neighbor_tables = (
-            get_used_neighbor_tables(node, self.offset_provider) if use_neighbor_tables else {}
+            get_used_connectivities(node, self.offset_provider) if use_neighbor_tables else {}
         )
         connectivity_names = [connectivity_identifier(offset) for offset in neighbor_tables.keys()]
 
@@ -1070,7 +1069,7 @@ class PythonTaskletCodegen(gt4py.eve.codegen.TemplatedGenerator):
                         store, self.context.body.arrays[store]
                     )
 
-        neighbor_tables = get_used_neighbor_tables(node.fun, self.offset_provider)
+        neighbor_tables = get_used_connectivities(node.fun, self.offset_provider)
         for offset in neighbor_tables.keys():
             var = connectivity_identifier(offset)
             nsdfg_inputs[var] = dace.Memlet.from_array(var, self.context.body.arrays[var])
@@ -1216,7 +1215,7 @@ class PythonTaskletCodegen(gt4py.eve.codegen.TemplatedGenerator):
         offset_node = self.visit(tail[1])[0]
         assert offset_node.dtype in dace.dtypes.INTEGER_TYPES
 
-        if isinstance(self.offset_provider[offset_dim], NeighborTableOffsetProvider):
+        if isinstance(self.offset_provider[offset_dim], Connectivity):
             offset_provider = self.offset_provider[offset_dim]
             connectivity = self.context.state.add_access(
                 connectivity_identifier(offset_dim), debuginfo=di
@@ -1225,20 +1224,12 @@ class PythonTaskletCodegen(gt4py.eve.codegen.TemplatedGenerator):
             shifted_dim = offset_provider.origin_axis.value
             target_dim = offset_provider.neighbor_axis.value
             args = [
-                ValueExpr(connectivity, offset_provider.table.dtype),
+                ValueExpr(connectivity, _INDEX_DTYPE),
                 ValueExpr(iterator.indices[shifted_dim], offset_node.dtype),
                 offset_node,
             ]
             internals = [f"{arg.value.data}_v" for arg in args]
             expr = f"{internals[0]}[{internals[1]}, {internals[2]}]"
-        elif isinstance(self.offset_provider[offset_dim], StridedNeighborOffsetProvider):
-            offset_provider = self.offset_provider[offset_dim]
-
-            shifted_dim = offset_provider.origin_axis.value
-            target_dim = offset_provider.neighbor_axis.value
-            args = [ValueExpr(iterator.indices[shifted_dim], offset_node.dtype), offset_node]
-            internals = [f"{arg.value.data}_v" for arg in args]
-            expr = f"{internals[0]} * {offset_provider.max_neighbors} + {internals[1]}"
         else:
             assert isinstance(self.offset_provider[offset_dim], Dimension)
 
diff --git a/src/gt4py/next/program_processors/runners/dace_iterator/utility.py b/src/gt4py/next/program_processors/runners/dace_iterator/utility.py
index fd679f04..5f852b28 100644
--- a/src/gt4py/next/program_processors/runners/dace_iterator/utility.py
+++ b/src/gt4py/next/program_processors/runners/dace_iterator/utility.py
@@ -19,7 +19,7 @@ import dace
 import gt4py.next.iterator.ir as itir
 from gt4py import eve
 from gt4py.next import Dimension
-from gt4py.next.common import NeighborTable
+from gt4py.next.common import Connectivity
 from gt4py.next.type_system import type_specifications as ts
 
 
@@ -60,20 +60,20 @@ def as_scalar_type(typestr: str) -> ts.ScalarType:
     return ts.ScalarType(kind)
 
 
-def filter_neighbor_tables(offset_provider: Mapping[str, Any]) -> dict[str, NeighborTable]:
+def filter_connectivities(offset_provider: Mapping[str, Any]) -> dict[str, Connectivity]:
     return {
         offset: table
         for offset, table in offset_provider.items()
-        if isinstance(table, NeighborTable)
+        if isinstance(table, Connectivity)
     }
 
 
-def get_used_neighbor_tables(
+def get_used_connectivities(
     node: itir.Node, offset_provider: Mapping[str, Any]
-) -> dict[str, NeighborTable]:
-    neighbor_tables = filter_neighbor_tables(offset_provider)
+) -> dict[str, Connectivity]:
+    connectivities = filter_connectivities(offset_provider)
     offset_dims = set(eve.walk_values(node).if_isinstance(itir.OffsetLiteral).getattr("value"))
-    return {offset: neighbor_tables[offset] for offset in offset_dims if offset in neighbor_tables}
+    return {offset: connectivities[offset] for offset in offset_dims if offset in connectivities}
 
 
 def connectivity_identifier(name: str) -> str:
diff --git a/src/gt4py/next/program_processors/runners/dace_iterator/workflow.py b/src/gt4py/next/program_processors/runners/dace_iterator/workflow.py
index 537bda79..96a40e74 100644
--- a/src/gt4py/next/program_processors/runners/dace_iterator/workflow.py
+++ b/src/gt4py/next/program_processors/runners/dace_iterator/workflow.py
@@ -24,14 +24,13 @@ from dace.codegen.compiled_sdfg import CompiledSDFG
 
 from gt4py._core import definitions as core_defs
 from gt4py.next import common, config
-from gt4py.next.common import Dimension
 from gt4py.next.iterator import ir as itir
 from gt4py.next.iterator.transforms import LiftMode
 from gt4py.next.otf import languages, stages, step_types, workflow
 from gt4py.next.otf.binding import interface
 from gt4py.next.otf.compilation import cache
 from gt4py.next.otf.languages import LanguageSettings
-from gt4py.next.type_system import type_translation as tt
+from gt4py.next.type_system import type_specifications as ts, type_translation as tt
 
 from . import build_sdfg_from_itir, get_sdfg_args
 
@@ -46,6 +45,7 @@ class DaCeTranslator(
     auto_optimize: bool = False
     lift_mode: LiftMode = LiftMode.FORCE_INLINE
     device_type: core_defs.DeviceType = core_defs.DeviceType.CPU
+    symbolic_domain_sizes: Optional[dict[str, str]] = None
     temporary_extraction_heuristics: Optional[
         Callable[[itir.StencilClosure], Callable[[itir.Expr], bool]]
     ] = None
@@ -56,20 +56,16 @@ class DaCeTranslator(
             formatter_key="", formatter_style="", file_extension="sdfg"
         )
 
-    def __call__(
-        self, inp: stages.ProgramCall
-    ) -> stages.ProgramSource[languages.SDFG, LanguageSettings]:
-        """Generate DaCe SDFG file from the ITIR definition."""
-        program: itir.FencilDefinition = inp.program
+    def generate_sdfg(
+        self,
+        program: itir.FencilDefinition,
+        arg_types: list[ts.TypeSpec],
+        offset_provider: dict[str, common.Dimension | common.Connectivity],
+        column_axis: Optional[common.Dimension],
+        runtime_lift_mode: Optional[LiftMode] = None,
+    ) -> dace.SDFG:
         on_gpu = True if self.device_type == core_defs.DeviceType.CUDA else False
 
-        # ITIR parameters
-        column_axis: Optional[Dimension] = inp.kwargs.get("column_axis", None)
-        offset_provider: dict[str, common.Dimension | common.Connectivity] = inp.kwargs[
-            "offset_provider"
-        ]
-        runtime_lift_mode: Optional[LiftMode] = inp.kwargs.get("lift_mode", None)
-
         # TODO(tehrengruber): Remove `lift_mode` from call interface. It has been implicitly added
         #  to the interface of all (or at least all of concern) backends, but instead should be
         #  configured in the backend itself (like it is here), until then we respect the argument
@@ -82,27 +78,44 @@ class DaCeTranslator(
                 stacklevel=2,
             )
 
-        sdfg = build_sdfg_from_itir(
+        return build_sdfg_from_itir(
             program,
-            *inp.args,
+            arg_types,
             offset_provider=offset_provider,
             auto_optimize=self.auto_optimize,
             on_gpu=on_gpu,
             column_axis=column_axis,
             lift_mode=lift_mode,
+            symbolic_domain_sizes=self.symbolic_domain_sizes,
+            temporary_extraction_heuristics=self.temporary_extraction_heuristics,
             load_sdfg_from_file=False,
             save_sdfg=False,
             use_field_canonical_representation=self.use_field_canonical_representation,
         )
 
-        arg_types = tuple(
+    def __call__(
+        self, inp: stages.ProgramCall
+    ) -> stages.ProgramSource[languages.SDFG, LanguageSettings]:
+        """Generate DaCe SDFG file from the ITIR definition."""
+        program: itir.FencilDefinition = inp.program
+        arg_types = [tt.from_value(arg) for arg in inp.args]
+
+        sdfg = self.generate_sdfg(
+            program,
+            arg_types,
+            inp.kwargs["offset_provider"],
+            inp.kwargs.get("column_axis", None),
+            inp.kwargs.get("lift_mode", None),
+        )
+
+        param_types = tuple(
             interface.Parameter(param, tt.from_value(arg))
             for param, arg in zip(sdfg.arg_names, inp.args)
         )
 
         module: stages.ProgramSource[languages.SDFG, languages.LanguageSettings] = (
             stages.ProgramSource(
-                entry_point=interface.Function(program.id, arg_types),
+                entry_point=interface.Function(program.id, param_types),
                 source_code=sdfg.to_json(),
                 library_deps=tuple(),
                 language=languages.SDFG,


Test patch:
diff --git a/min-extra-requirements-test.txt b/min-extra-requirements-test.txt
index d3871db2..d6689803 100644
--- a/min-extra-requirements-test.txt
+++ b/min-extra-requirements-test.txt
@@ -113,6 +113,6 @@ tabulate==0.8.10
 tomli==2.0.1; python_version < "3.11"
 tox==3.2.0
 types-all==1.0.0
-typing-extensions==4.6.0
+typing-extensions==4.10.0
 xxhash==1.4.4
 ##[[[end]]]
diff --git a/min-requirements-test.txt b/min-requirements-test.txt
index e985b08a..eeb633dd 100644
--- a/min-requirements-test.txt
+++ b/min-requirements-test.txt
@@ -106,6 +106,6 @@ tabulate==0.8.10
 tomli==2.0.1; python_version < "3.11"
 tox==3.2.0
 types-all==1.0.0
-typing-extensions==4.6.0
+typing-extensions==4.10.0
 xxhash==1.4.4
 ##[[[end]]]
diff --git a/tests/next_tests/definitions.py b/tests/next_tests/definitions.py
index 44989e78..86eac697 100644
--- a/tests/next_tests/definitions.py
+++ b/tests/next_tests/definitions.py
@@ -146,6 +146,7 @@ DACE_SKIP_TEST_LIST = COMMON_SKIP_TEST_LIST + [
     (USES_INDEX_FIELDS, XFAIL, UNSUPPORTED_MESSAGE),
     (USES_LIFT_EXPRESSIONS, XFAIL, UNSUPPORTED_MESSAGE),
     (USES_ORIGIN, XFAIL, UNSUPPORTED_MESSAGE),
+    (USES_STRIDED_NEIGHBOR_OFFSET, XFAIL, BINDINGS_UNSUPPORTED_MESSAGE),
     (USES_TUPLE_ARGS, XFAIL, UNSUPPORTED_MESSAGE),
     (USES_TUPLE_RETURNS, XFAIL, UNSUPPORTED_MESSAGE),
     (USES_ZERO_DIMENSIONAL_FIELDS, XFAIL, UNSUPPORTED_MESSAGE),
diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_decorator.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_decorator.py
new file mode 100644
index 00000000..a6d00ab7
--- /dev/null
+++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_decorator.py
@@ -0,0 +1,35 @@
+# 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
+# TODO(dropd): Remove as soon as `gt4py.next.ffront.decorator` is type checked.
+from gt4py import next as gtx
+from gt4py.next.iterator import ir as itir
+
+from next_tests.integration_tests import cases
+from next_tests.integration_tests.cases import cartesian_case
+from next_tests.integration_tests.feature_tests.ffront_tests.ffront_test_utils import (
+    exec_alloc_descriptor,
+)
+
+
+def test_program_itir_regression(cartesian_case):
+    @gtx.field_operator(backend=None)
+    def testee_op(a: cases.IField) -> cases.IField:
+        return a
+
+    @gtx.program(backend=None)
+    def testee(a: cases.IField, out: cases.IField):
+        testee_op(a, out=out)
+
+    assert isinstance(testee.itir, itir.FencilDefinition)
+    assert isinstance(testee.with_backend(cartesian_case.executor).itir, itir.FencilDefinition)
diff --git a/tests/next_tests/unit_tests/iterator_tests/test_embedded_internals.py b/tests/next_tests/unit_tests/iterator_tests/test_embedded_internals.py
index ec6e6135..4ce8394b 100644
--- a/tests/next_tests/unit_tests/iterator_tests/test_embedded_internals.py
+++ b/tests/next_tests/unit_tests/iterator_tests/test_embedded_internals.py
@@ -13,13 +13,13 @@
 # SPDX-License-Identifier: GPL-3.0-or-later
 
 import contextvars as cvars
-import threading
 from typing import Any, Callable, Optional
 
 import numpy as np
 import pytest
 
 from gt4py.next import common
+from gt4py.next.embedded import context as embedded_context
 from gt4py.next.iterator import embedded
 
 
@@ -30,8 +30,8 @@ def _run_within_context(
     offset_provider: Optional[embedded.OffsetProvider] = None,
 ) -> Any:
     def wrapped_func():
-        embedded.column_range_cvar.set(column_range)
-        embedded.offset_provider_cvar.set(offset_provider)
+        embedded_context.closure_column_range.set(column_range)
+        embedded_context.offset_provider.set(offset_provider)
         func()
 
     cvars.copy_context().run(wrapped_func)
@@ -59,7 +59,7 @@ def test_column_ufunc():
         assert res.kstart == 1
 
     # Setting an invalid column_range here shouldn't affect other contexts
-    embedded.column_range_cvar.set(range(2, 999))
+    embedded_context.closure_column_range.set(range(2, 999))
     _run_within_context(
         lambda: test_func(2, 3),
         column_range=common.NamedRange(