Repository: GridTools/gt4py
Pull request #1568 (2024-07-03 09:38:05)

GTIR domain inference
 #### Uses TraceShifts to infer the minimal domain of (nested) `as_fieldop`s and SetAt:
Domain inference of let-statements is not implemented yet and will be done in PR #1591
```
expr = as_fieldop(stencil)(as_fieldop(stencil)(inp_field))

domain = cartesian_domain(...)

# new_expr has transformed all as_fieldop(stencil) into as_fieldop(stencil, domain)
infer_as_fieldop(expr, domain) -> new_expr, dict(referenced_field => accessed domain)
# with new_expr = as_fieldop(stencil, cartesian_domain(...))(inp_field, ...)

# multiple SetAt (no composition)
tmp <- as_fieldop(stencil)(inp_field, ...) @ auto_domain()
out <- as_fieldop(stencil)(tmp, ...) @ ...

# domain is inferred bottom up from last SetAt, only domains of SetAt on temporaries are changed
infer_program(program_expr)
```
### PR Description
- New file gt4py.next.iterator.transforms.infer_domain with function `infer_as_fieldop` and `infer_program`
- Several new tests in test_infer_domain.py to test functionality
- New function SymbolicDomain.translate
- New helper-function `domain` and `as_fieldop` in ir.makers
- Uses constant folding of `itir.Literal`s to simplify resulting domains



### Todos
- [x] Domain inference for nested  `as_fieldop`s (cartesian_domain)
- [x] Domain inference for SetAt
- [x] Extend to unstructured grids

### Note
I agree that my contribution in this PR is additionally licensed under the BSD-3-Clause license (SPDX short identifier).

tehrengruber: Note to myself: One function for itir.Program with temporaries, one without (only the simple for loop). When writing temporary pass introduce `ProgramWithUnresolvedTemporaries` that allows `_gtmp_auto_domain` and potentially move the inference (continuing to build upon the primitive function operaton on `as_fieldop` here) there.

Patch:
diff --git a/src/gt4py/next/iterator/ir_utils/ir_makers.py b/src/gt4py/next/iterator/ir_utils/ir_makers.py
index 8d29e49b..166f25d3 100644
--- a/src/gt4py/next/iterator/ir_utils/ir_makers.py
+++ b/src/gt4py/next/iterator/ir_utils/ir_makers.py
@@ -10,6 +10,8 @@ import typing
 from typing import Callable, Iterable, Optional, Union
 
 from gt4py._core import definitions as core_defs
+from gt4py.eve.extended_typing import Dict, Tuple
+from gt4py.next import common
 from gt4py.next.iterator import ir as itir
 from gt4py.next.type_system import type_specifications as ts, type_translation
 
@@ -397,9 +399,46 @@ def map_(op):
     return call(call("map_")(op))
 
 
+def domain(
+    grid_type: Union[common.GridType, str],
+    ranges: Dict[Union[common.Dimension, str], Tuple[itir.Expr, itir.Expr]],
+) -> itir.FunCall:
+    """
+    >>> str(
+    ...     domain(
+    ...         common.GridType.CARTESIAN,
+    ...         {
+    ...             common.Dimension(value="IDim", kind=common.DimensionKind.HORIZONTAL): (0, 10),
+    ...             common.Dimension(value="JDim", kind=common.DimensionKind.HORIZONTAL): (0, 20),
+    ...         },
+    ...     )
+    ... )
+    'c⟨ IDimₕ: [0, 10), JDimₕ: [0, 20) ⟩'
+    >>> str(domain(common.GridType.CARTESIAN, {"IDim": (0, 10), "JDim": (0, 20)}))
+    'c⟨ IDimₕ: [0, 10), JDimₕ: [0, 20) ⟩'
+    >>> str(domain(common.GridType.UNSTRUCTURED, {"IDim": (0, 10), "JDim": (0, 20)}))
+    'u⟨ IDimₕ: [0, 10), JDimₕ: [0, 20) ⟩'
+    """
+    if isinstance(grid_type, common.GridType):
+        grid_type = f"{grid_type!s}_domain"
+    return call(grid_type)(
+        *[
+            call("named_range")(
+                itir.AxisLiteral(value=d.value, kind=d.kind)
+                if isinstance(d, common.Dimension)
+                else itir.AxisLiteral(value=d),
+                r[0],
+                r[1],
+            )
+            for d, r in ranges.items()
+        ]
+    )
+
+
 def as_fieldop(expr: itir.Expr, domain: Optional[itir.FunCall] = None) -> call:
     """
     Create an `as_fieldop` call.
+
     Examples
     --------
     >>> str(as_fieldop(lambda_("it1", "it2")(plus(deref("it1"), deref("it2"))))("field1", "field2"))
diff --git a/src/gt4py/next/iterator/transforms/global_tmps.py b/src/gt4py/next/iterator/transforms/global_tmps.py
index 1fdcdbb1..0c9f68a5 100644
--- a/src/gt4py/next/iterator/transforms/global_tmps.py
+++ b/src/gt4py/next/iterator/transforms/global_tmps.py
@@ -16,6 +16,7 @@ from typing import Any, Callable, Final, Iterable, Literal, Optional, Sequence
 
 import gt4py.next as gtx
 from gt4py.eve import NodeTranslator, PreserveLocationVisitor
+from gt4py.eve.extended_typing import Dict, Tuple
 from gt4py.eve.traits import SymbolTableTrait
 from gt4py.eve.utils import UIDGenerator
 from gt4py.next import common
@@ -446,12 +447,56 @@ class SymbolicDomain:
         return cls(node.fun.id, ranges)  # type: ignore[attr-defined]  # ensure by assert above
 
     def as_expr(self) -> ir.FunCall:
-        return im.call(self.grid_type)(
-            *[
-                im.call("named_range")(ir.AxisLiteral(value=d.value, kind=d.kind), r.start, r.stop)
-                for d, r in self.ranges.items()
-            ]
-        )
+        converted_ranges: dict[common.Dimension | str, tuple[ir.Expr, ir.Expr]] = {
+            key: (value.start, value.stop) for key, value in self.ranges.items()
+        }
+        return im.domain(self.grid_type, converted_ranges)
+
+    def translate(
+        self: SymbolicDomain,
+        shift: Tuple[ir.OffsetLiteral, ...],
+        offset_provider: Dict[str, common.Dimension],
+    ) -> SymbolicDomain:
+        dims = list(self.ranges.keys())
+        new_ranges = {dim: self.ranges[dim] for dim in dims}
+        if len(shift) == 0:
+            return self
+        if len(shift) == 2:
+            off, val = shift
+            assert isinstance(off.value, str) and isinstance(val.value, int)
+            nbt_provider = offset_provider[off.value]
+            if isinstance(nbt_provider, common.Dimension):
+                current_dim = nbt_provider
+                # cartesian offset
+                new_ranges[current_dim] = SymbolicRange.translate(
+                    self.ranges[current_dim], val.value
+                )
+            elif isinstance(nbt_provider, common.Connectivity):
+                # unstructured shift
+                # note: ugly but cheap re-computation, but should disappear
+                horizontal_sizes = _max_domain_sizes_by_location_type(offset_provider)
+
+                old_dim = nbt_provider.origin_axis
+                new_dim = nbt_provider.neighbor_axis
+
+                assert new_dim not in new_ranges or old_dim == new_dim
+
+                # TODO(tehrengruber): Do we need symbolic sizes, e.g., for ICON?
+                new_range = SymbolicRange(
+                    im.literal("0", ir.INTEGER_INDEX_BUILTIN),
+                    im.literal(str(horizontal_sizes[new_dim.value]), ir.INTEGER_INDEX_BUILTIN),
+                )
+                new_ranges = dict(
+                    (dim, range_) if dim != old_dim else (new_dim, new_range)
+                    for dim, range_ in new_ranges.items()
+                )
+            else:
+                raise AssertionError()
+            return SymbolicDomain(self.grid_type, new_ranges)
+        elif len(shift) > 2:
+            return self.translate(shift[0:2], offset_provider).translate(shift[2:], offset_provider)
+        else:
+            raise AssertionError("Number of shifts must be a multiple of 2.")
 
 
 def domain_union(domains: list[SymbolicDomain]) -> SymbolicDomain:
@@ -469,6 +514,7 @@ def domain_union(domains: list[SymbolicDomain]) -> SymbolicDomain:
             [domain.ranges[dim].stop for domain in domains],
         )
         new_domain_ranges[dim] = SymbolicRange(start, stop)
+
     return SymbolicDomain(domains[0].grid_type, new_domain_ranges)
 
 
diff --git a/src/gt4py/next/iterator/transforms/infer_domain.py b/src/gt4py/next/iterator/transforms/infer_domain.py
new file mode 100644
index 00000000..2d465eb3
--- /dev/null
+++ b/src/gt4py/next/iterator/transforms/infer_domain.py
@@ -0,0 +1,205 @@
+# GT4Py - GridTools Framework
+#
+# Copyright (c) 2014-2024, ETH Zurich
+# All rights reserved.
+#
+# Please, refer to the LICENSE file in the root directory.
+# SPDX-License-Identifier: BSD-3-Clause
+
+from gt4py.eve import utils as eve_utils
+from gt4py.eve.extended_typing import Dict, Tuple
+from gt4py.next.common import Dimension
+from gt4py.next.iterator import ir as itir
+from gt4py.next.iterator.ir_utils import common_pattern_matcher as cpm, ir_makers as im
+from gt4py.next.iterator.transforms.global_tmps import AUTO_DOMAIN, SymbolicDomain, domain_union
+from gt4py.next.iterator.transforms.trace_shifts import TraceShifts
+
+
+def _merge_domains(
+    original_domains: Dict[str, SymbolicDomain], additional_domains: Dict[str, SymbolicDomain]
+) -> Dict[str, SymbolicDomain]:
+    new_domains = {**original_domains}
+    for key, value in additional_domains.items():
+        if key in original_domains:
+            new_domains[key] = domain_union([original_domains[key], value])
+        else:
+            new_domains[key] = value
+
+    return new_domains
+
+
+# FIXME[#1582](tehrengruber): Use new TraceShift API when #1592 is merged.
+def trace_shifts(
+    stencil: itir.Expr, input_ids: list[str], domain: itir.Expr
+) -> dict[str, set[tuple[itir.OffsetLiteral, ...]]]:
+    node = itir.StencilClosure(
+        stencil=stencil,
+        inputs=[im.ref(id_) for id_ in input_ids],
+        output=im.ref("__dummy"),
+        domain=domain,
+    )
+    return TraceShifts.apply(node, inputs_only=True)  # type: ignore[return-value]  # ensured by inputs_only=True
+
+
+def extract_shifts_and_translate_domains(
+    stencil: itir.Expr,
+    input_ids: list[str],
+    target_domain: SymbolicDomain,
+    offset_provider: Dict[str, Dimension],
+    accessed_domains: Dict[str, SymbolicDomain],
+):
+    shifts_results = trace_shifts(stencil, input_ids, SymbolicDomain.as_expr(target_domain))
+
+    for in_field_id in input_ids:
+        shifts_list = shifts_results[in_field_id]
+
+        new_domains = [
+            SymbolicDomain.translate(target_domain, shift, offset_provider) for shift in shifts_list
+        ]
+        if new_domains:
+            accessed_domains[in_field_id] = domain_union(new_domains)
+
+
+def infer_as_fieldop(
+    applied_fieldop: itir.FunCall,
+    target_domain: SymbolicDomain | itir.FunCall,
+    offset_provider: Dict[str, Dimension],
+) -> Tuple[itir.FunCall, Dict[str, SymbolicDomain]]:
+    assert isinstance(applied_fieldop, itir.FunCall)
+    assert cpm.is_call_to(applied_fieldop.fun, "as_fieldop")
+
+    # `as_fieldop(stencil)(inputs...)`
+    stencil, inputs = applied_fieldop.fun.args[0], applied_fieldop.args
+
+    # ensure stencil has as many params as arguments
+    assert not isinstance(stencil, itir.Lambda) or len(stencil.params) == len(applied_fieldop.args)
+
+    input_ids: list[str] = []
+    accessed_domains: Dict[str, SymbolicDomain] = {}
+
+    # Assign ids for all inputs to `as_fieldop`. `SymRef`s stay as is, nested `as_fieldop` get a
+    # temporary id.
+    tmp_uid_gen = eve_utils.UIDGenerator(prefix="__dom_inf")
+    for in_field in inputs:
+        if isinstance(in_field, itir.FunCall) or isinstance(in_field, itir.Literal):
+            id_ = tmp_uid_gen.sequential_id()
+        elif isinstance(in_field, itir.SymRef):
+            id_ = in_field.id
+        else:
+            raise ValueError(f"Unsupported type {type(in_field)}")
+        input_ids.append(id_)
+
+    if isinstance(target_domain, itir.FunCall):
+        target_domain = SymbolicDomain.from_expr(target_domain)
+
+    extract_shifts_and_translate_domains(
+        stencil, input_ids, target_domain, offset_provider, accessed_domains
+    )
+
+    # Recursively infer domain of inputs and update domain arg of nested `as_fieldops`
+    transformed_inputs: list[itir.Expr] = []
+    for in_field_id, in_field in zip(input_ids, inputs):
+        if isinstance(in_field, itir.FunCall):
+            transformed_input, accessed_domains_tmp = infer_as_fieldop(
+                in_field, accessed_domains[in_field_id], offset_provider
+            )
+            transformed_inputs.append(transformed_input)
+
+            # Merge accessed_domains and accessed_domains_tmp
+            accessed_domains = _merge_domains(accessed_domains, accessed_domains_tmp)
+        elif isinstance(in_field, itir.SymRef) or isinstance(in_field, itir.Literal):
+            transformed_inputs.append(in_field)
+        else:
+            raise ValueError(f"Unsupported type {type(in_field)}")
+
+    transformed_call = im.as_fieldop(stencil, SymbolicDomain.as_expr(target_domain))(
+        *transformed_inputs
+    )
+
+    accessed_domains_without_tmp = {
+        k: v
+        for k, v in accessed_domains.items()
+        if not k.startswith(tmp_uid_gen.prefix)  # type: ignore[arg-type] # prefix is always str
+    }
+
+    return transformed_call, accessed_domains_without_tmp
+
+
+def _validate_temporary_usage(body: list[itir.Stmt], temporaries: list[str]):
+    assigned_targets = set()
+    for stmt in body:
+        assert isinstance(stmt, itir.SetAt)  # TODO: extend for if-statements when they land
+        assert isinstance(
+            stmt.target, itir.SymRef
+        )  # TODO: stmt.target can be an expr, e.g. make_tuple
+        if stmt.target.id in assigned_targets:
+            raise ValueError("Temporaries can only be used once within a program.")
+        if stmt.target.id in temporaries:
+            assigned_targets.add(stmt.target.id)
+
+
+def infer_program(
+    program: itir.Program,
+    offset_provider: Dict[str, Dimension],
+) -> itir.Program:
+    accessed_domains: dict[str, SymbolicDomain] = {}
+    transformed_set_ats: list[itir.SetAt] = []
+
+    temporaries: list[str] = [tmp.id for tmp in program.declarations]
+
+    _validate_temporary_usage(program.body, temporaries)
+
+    for set_at in reversed(program.body):
+        assert isinstance(set_at, itir.SetAt)
+        if isinstance(set_at.expr, itir.SymRef):
+            transformed_set_ats.insert(0, set_at)
+            continue
+        assert isinstance(set_at.expr, itir.FunCall)
+        assert cpm.is_call_to(set_at.expr.fun, "as_fieldop")
+        assert isinstance(
+            set_at.target, itir.SymRef
+        )  # TODO: stmt.target can be an expr, e.g. make_tuple
+        if set_at.target.id in temporaries:
+            # ignore temporaries as their domain is the `AUTO_DOMAIN` placeholder
+            assert set_at.domain == AUTO_DOMAIN
+        else:
+            accessed_domains[set_at.target.id] = SymbolicDomain.from_expr(set_at.domain)
+
+        transformed_as_fieldop, current_accessed_domains = infer_as_fieldop(
+            set_at.expr, accessed_domains[set_at.target.id], offset_provider
+        )
+        transformed_set_ats.insert(
+            0,
+            itir.SetAt(
+                expr=transformed_as_fieldop,
+                domain=SymbolicDomain.as_expr(accessed_domains[set_at.target.id]),
+                target=set_at.target,
+            ),
+        )
+
+        for field in current_accessed_domains:
+            if field in accessed_domains:
+                # multiple accesses to the same field -> compute union of accessed domains
+                if field in temporaries:
+                    accessed_domains[field] = domain_union(
+                        [accessed_domains[field], current_accessed_domains[field]]
+                    )
+                else:
+                    # TODO(tehrengruber): if domain_ref is an external field the domain must
+                    #  already be larger. This should be checked, but would require additions
+                    #  to the IR.
+                    pass
+            else:
+                accessed_domains[field] = current_accessed_domains[field]
+
+    new_declarations = program.declarations
+    for temporary in new_declarations:
+        temporary.domain = SymbolicDomain.as_expr(accessed_domains[temporary.id])
+
+    return itir.Program(
+        id=program.id,
+        function_definitions=program.function_definitions,
+        params=program.params,
+        declarations=new_declarations,
+        body=transformed_set_ats,
+    )
diff --git a/src/gt4py/next/iterator/transforms/trace_shifts.py b/src/gt4py/next/iterator/transforms/trace_shifts.py
index db87d4a3..d181940b 100644
--- a/src/gt4py/next/iterator/transforms/trace_shifts.py
+++ b/src/gt4py/next/iterator/transforms/trace_shifts.py
@@ -7,7 +7,6 @@
 # SPDX-License-Identifier: BSD-3-Clause
 
 import dataclasses
-import enum
 import sys
 from collections.abc import Callable
 from typing import Any, Final, Iterable, Literal
@@ -45,11 +44,11 @@ def copy_recorded_shifts(from_: ir.Node, to: ir.Node) -> None:
     to.annex.recorded_shifts = from_.annex.recorded_shifts
 
 
-class Sentinel(enum.Enum):
-    VALUE = enum.auto()
-    TYPE = enum.auto()
+class Sentinel(eve.StrEnum):
+    VALUE = "VALUE"
+    TYPE = "TYPE"
 
-    ALL_NEIGHBORS = enum.auto()
+    ALL_NEIGHBORS = "ALL_NEIGHBORS"
 
 
 @dataclasses.dataclass(frozen=True)
@@ -137,6 +136,11 @@ def _can_deref(x):
 
 
 def _shift(*offsets):
+    assert all(
+        isinstance(offset, ir.OffsetLiteral) or offset in [Sentinel.ALL_NEIGHBORS, Sentinel.VALUE]
+        for offset in offsets
+    )
+
     def apply(arg):
         assert isinstance(arg, IteratorTracer)
         return arg.shift(offsets)


Test patch:
diff --git a/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_domain_inference.py b/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_domain_inference.py
new file mode 100644
index 00000000..eb3deb1a
--- /dev/null
+++ b/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_domain_inference.py
@@ -0,0 +1,545 @@
+# GT4Py - GridTools Framework
+#
+# Copyright (c) 2014-2024, ETH Zurich
+# All rights reserved.
+#
+# Please, refer to the LICENSE file in the root directory.
+# SPDX-License-Identifier: BSD-3-Clause
+
+# TODO(SF-N): test scan operator
+
+import numpy as np
+from typing import Iterable
+
+from gt4py import eve
+from gt4py.next.iterator.ir_utils import common_pattern_matcher as cpm, ir_makers as im
+from gt4py.next.iterator import ir as itir
+from gt4py.next.iterator.transforms.infer_domain import infer_as_fieldop, infer_program
+from gt4py.next.iterator.transforms.global_tmps import SymbolicDomain, AUTO_DOMAIN
+import pytest
+from gt4py.eve.extended_typing import Dict
+from gt4py.next.common import Dimension, DimensionKind
+from gt4py.next import common, NeighborTableOffsetProvider
+from gt4py.next.type_system import type_specifications as ts
+from gt4py.next.iterator.transforms.constant_folding import ConstantFolding
+
+float_type = ts.ScalarType(kind=ts.ScalarKind.FLOAT64)
+IDim = common.Dimension(value="IDim", kind=common.DimensionKind.HORIZONTAL)
+JDim = common.Dimension(value="JDim", kind=common.DimensionKind.HORIZONTAL)
+KDim = common.Dimension(value="KDim", kind=common.DimensionKind.VERTICAL)
+Vertex = common.Dimension(value="Vertex", kind=common.DimensionKind.HORIZONTAL)
+Edge = common.Dimension(value="Edge", kind=common.DimensionKind.HORIZONTAL)
+
+
+@pytest.fixture
+def offset_provider():
+    return {
+        "Ioff": IDim,
+        "Joff": JDim,
+        "Koff": KDim,
+    }
+
+
+@pytest.fixture
+def unstructured_offset_provider():
+    return {
+        "E2V": NeighborTableOffsetProvider(
+            np.array([[0, 1]], dtype=np.int32),
+            Edge,
+            Vertex,
+            2,
+        )
+    }
+
+
+def run_test_as_fieldop(
+    stencil: itir.Lambda,
+    domain: itir.FunCall,
+    expected_domain_dict: Dict[str, Dict[str | Dimension, tuple[itir.Expr, itir.Expr]]],
+    offset_provider: Dict[str, Dimension],
+    *,
+    refs: Iterable[itir.SymRef] = None,
+    domain_type: str = common.GridType.CARTESIAN,
+) -> None:
+    if refs is None:
+        refs = [f"in_field{i+1}" for i in range(0, len(stencil.params))]
+
+    testee = im.as_fieldop(stencil)(*refs)
+    expected = im.as_fieldop(stencil, domain)(*refs)
+
+    actual_call, actual_domains = infer_as_fieldop(
+        testee, SymbolicDomain.from_expr(domain), offset_provider
+    )
+
+    folded_domains = constant_fold_accessed_domains(actual_domains)
+    expected_domains = {
+        ref: SymbolicDomain.from_expr(im.domain(domain_type, d))
+        for ref, d in expected_domain_dict.items()
+    }
+
+    assert actual_call == expected
+    assert folded_domains == expected_domains
+
+
+def run_test_program(
+    testee: itir.Program, expected: itir.Program, offset_provider: dict[str, Dimension]
+) -> None:
+    actual_program = infer_program(testee, offset_provider)
+
+    folded_program = constant_fold_domain_exprs(actual_program)
+    assert folded_program == expected
+
+
+class _ConstantFoldDomainsExprs(eve.NodeTranslator):
+    def visit_FunCall(self, node: itir.FunCall):
+        if cpm.is_call_to(node, ("cartesian_domain", "unstructured_domain")):
+            return ConstantFolding.apply(node)
+        return self.generic_visit(node)
+
+
+def constant_fold_domain_exprs(arg: itir.Node) -> itir.Node:
+    return _ConstantFoldDomainsExprs().visit(arg)
+
+
+def constant_fold_accessed_domains(domains: Dict[str, SymbolicDomain]) -> Dict[str, SymbolicDomain]:
+    return {
+        k: SymbolicDomain.from_expr(constant_fold_domain_exprs(v.as_expr()))
+        for k, v in domains.items()
+    }
+
+
+def test_forward_difference_x(offset_provider):
+    stencil = im.lambda_("arg0")(im.minus(im.deref(im.shift("Ioff", 1)("arg0")), im.deref("arg0")))
+    domain = im.domain(common.GridType.CARTESIAN, {IDim: (0, 11)})
+    expected_accessed_domains = {"in_field1": {IDim: (0, 12)}}
+    run_test_as_fieldop(stencil, domain, expected_accessed_domains, offset_provider)
+
+
+def test_multi_length_shift(offset_provider):
+    stencil = im.lambda_("arg0")(
+        im.deref(
+            im.call(
+                im.call("shift")(
+                    im.ensure_offset("Ioff"),
+                    im.ensure_offset(1),
+                    im.ensure_offset("Ioff"),
+                    im.ensure_offset(2),
+                )
+            )("arg0")
+        )
+    )
+    domain = im.domain(common.GridType.CARTESIAN, {IDim: (0, 11)})
+    expected_accessed_domains = {"in_field1": {IDim: (3, 14)}}
+    run_test_as_fieldop(stencil, domain, expected_accessed_domains, offset_provider)
+
+
+def test_unused_input(offset_provider):
+    stencil = im.lambda_("arg0", "arg1")(im.deref("arg0"))
+
+    domain = im.domain(common.GridType.CARTESIAN, {IDim: (0, 11)})
+    expected_accessed_domains = {
+        "in_field1": {IDim: (0, 11)},
+    }
+    run_test_as_fieldop(
+        stencil,
+        domain,
+        expected_accessed_domains,
+        offset_provider,
+    )
+
+
+def test_unstructured_shift(unstructured_offset_provider):
+    stencil = im.lambda_("arg0")(im.deref(im.shift("E2V", 1)("arg0")))
+    domain = im.domain(common.GridType.UNSTRUCTURED, {Edge: (0, 1)})
+    expected_accessed_domains = {"in_field1": {Vertex: (0, 2)}}
+
+    run_test_as_fieldop(
+        stencil,
+        domain,
+        expected_accessed_domains,
+        unstructured_offset_provider,
+        domain_type=common.GridType.UNSTRUCTURED,
+    )
+
+
+def test_laplace(offset_provider):
+    stencil = im.lambda_("arg0")(
+        im.plus(
+            im.plus(
+                im.plus(
+                    im.plus(
+                        im.multiplies_(-4.0, im.deref("arg0")),
+                        im.deref(im.shift("Ioff", 1)("arg0")),
+                    ),
+                    im.deref(im.shift("Joff", 1)("arg0")),
+                ),
+                im.deref(im.shift("Ioff", -1)("arg0")),
+            ),
+            im.deref(im.shift("Joff", -1)("arg0")),
+        )
+    )
+    domain = im.domain(common.GridType.CARTESIAN, {IDim: (0, 11), JDim: (0, 7)})
+    expected_accessed_domains = {"in_field1": {IDim: (-1, 12), JDim: (-1, 8)}}
+
+    run_test_as_fieldop(stencil, domain, expected_accessed_domains, offset_provider)
+
+
+def test_shift_x_y_two_inputs(offset_provider):
+    stencil = im.lambda_("arg0", "arg1")(
+        im.plus(
+            im.deref(im.shift("Ioff", -1)("arg0")),
+            im.deref(im.shift("Joff", 1)("arg1")),
+        )
+    )
+    domain = im.domain(common.GridType.CARTESIAN, {IDim: (0, 11), JDim: (0, 7)})
+    expected_accessed_domains = {
+        "in_field1": {IDim: (-1, 10), JDim: (0, 7)},
+        "in_field2": {IDim: (0, 11), JDim: (1, 8)},
+    }
+    run_test_as_fieldop(
+        stencil,
+        domain,
+        expected_accessed_domains,
+        offset_provider,
+    )
+
+
+def test_shift_x_y_two_inputs_literal(offset_provider):
+    stencil = im.lambda_("arg0", "arg1")(
+        im.plus(
+            im.deref(im.shift("Ioff", -1)("arg0")),
+            im.deref(im.shift("Joff", 1)("arg1")),
+        )
+    )
+    domain = im.domain(common.GridType.CARTESIAN, {IDim: (0, 11), JDim: (0, 7)})
+    expected_accessed_domains = {
+        "in_field1": {IDim: (-1, 10), JDim: (0, 7)},
+    }
+    run_test_as_fieldop(
+        stencil,
+        domain,
+        expected_accessed_domains,
+        offset_provider,
+        refs=(im.ref("in_field1"), 2),
+    )
+
+
+def test_shift_x_y_z_three_inputs(offset_provider):
+    stencil = im.lambda_("arg0", "arg1", "arg2")(
+        im.plus(
+            im.plus(
+                im.deref(im.shift("Ioff", 1)("arg0")),
+                im.deref(im.shift("Joff", 1)("arg1")),
+            ),
+            im.deref(im.shift("Koff", -1)("arg2")),
+        )
+    )
+    domain_dict = {
+        IDim: (0, 11),
+        JDim: (0, 7),
+        KDim: (0, 3),
+    }
+    expected_domain_dict = {
+        "in_field1": {
+            IDim: (1, 12),
+            JDim: (0, 7),
+            KDim: (0, 3),
+        },
+        "in_field2": {
+            IDim: (0, 11),
+            JDim: (1, 8),
+            KDim: (0, 3),
+        },
+        "in_field3": {
+            IDim: (0, 11),
+            JDim: (0, 7),
+            KDim: (-1, 2),
+        },
+    }
+    run_test_as_fieldop(
+        stencil,
+        im.domain(common.GridType.CARTESIAN, domain_dict),
+        expected_domain_dict,
+        offset_provider,
+    )
+
+
+def test_nested_stencils(offset_provider):
+    inner_stencil = im.lambda_("arg0_tmp", "arg1_tmp")(
+        im.plus(
+            im.deref(im.shift("Ioff", 1)("arg0_tmp")),
+            im.deref(im.shift("Joff", -1)("arg1_tmp")),
+        )
+    )
+    stencil = im.lambda_("arg0", "arg1")(
+        im.plus(
+            im.deref(im.shift("Ioff", 1)("arg0")),
+            im.deref(im.shift("Joff", -1)("arg1")),
+        )
+    )
+    tmp = im.as_fieldop(inner_stencil)(im.ref("in_field1"), im.ref("in_field2"))
+    testee = im.as_fieldop(stencil)(im.ref("in_field1"), tmp)
+
+    domain_inner = im.domain(common.GridType.CARTESIAN, {IDim: (0, 11), JDim: (-1, 6)})
+    domain = im.domain(common.GridType.CARTESIAN, {IDim: (0, 11), JDim: (0, 7)})
+
+    expected_inner = im.as_fieldop(inner_stencil, domain_inner)(
+        im.ref("in_field1"), im.ref("in_field2")
+    )
+    expected = im.as_fieldop(stencil, domain)(im.ref("in_field1"), expected_inner)
+
+    expected_domains = {
+        "in_field1": SymbolicDomain.from_expr(
+            im.domain(common.GridType.CARTESIAN, {IDim: (1, 12), JDim: (-1, 7)})
+        ),
+        "in_field2": SymbolicDomain.from_expr(
+            im.domain(common.GridType.CARTESIAN, {IDim: (0, 11), JDim: (-2, 5)})
+        ),
+    }
+    actual_call, actual_domains = infer_as_fieldop(
+        testee, SymbolicDomain.from_expr(domain), offset_provider
+    )
+    folded_domains = constant_fold_accessed_domains(actual_domains)
+    folded_call = constant_fold_domain_exprs(actual_call)
+    assert folded_call == expected
+    assert folded_domains == expected_domains
+
+
+@pytest.mark.parametrize("iterations", [3, 5])
+def test_nested_stencils_n_times(offset_provider, iterations):
+    stencil = im.lambda_("arg0", "arg1")(
+        im.plus(
+            im.deref(im.shift("Ioff", 1)("arg0")),
+            im.deref(im.shift("Joff", 1)("arg1")),
+        )
+    )
+    assert iterations >= 2
+
+    domain = im.domain(
+        common.GridType.CARTESIAN, {IDim: (0, 11), JDim: (iterations - 1, 7 + iterations - 1)}
+    )
+    testee = im.as_fieldop(stencil)(im.ref("in_field1"), im.ref("in_field2"))
+    expected = im.as_fieldop(stencil, domain)(im.ref("in_field1"), im.ref("in_field2"))
+
+    for n in range(1, iterations):
+        domain = im.domain(
+            common.GridType.CARTESIAN,
+            {IDim: (0, 11), JDim: (iterations - 1 - n, 7 + iterations - 1 - n)},
+        )
+        testee = im.as_fieldop(stencil)(im.ref("in_field1"), testee)
+        expected = im.as_fieldop(stencil, domain)(im.ref("in_field1"), expected)
+
+    testee = testee
+
+    expected_domains = {
+        "in_field1": SymbolicDomain.from_expr(
+            im.domain(common.GridType.CARTESIAN, {IDim: (1, 12), JDim: (0, 7 + iterations - 1)})
+        ),
+        "in_field2": SymbolicDomain.from_expr(
+            im.domain(
+                common.GridType.CARTESIAN, {IDim: (0, 11), JDim: (iterations, 7 + iterations)}
+            )
+        ),
+    }
+
+    actual_call, actual_domains = infer_as_fieldop(
+        testee, SymbolicDomain.from_expr(domain), offset_provider
+    )
+
+    folded_domains = constant_fold_accessed_domains(actual_domains)
+    folded_call = constant_fold_domain_exprs(actual_call)
+    assert folded_call == expected
+    assert folded_domains == expected_domains
+
+
+def test_program(offset_provider):
+    stencil = im.lambda_("arg0")(im.minus(im.deref(im.shift("Ioff", 1)("arg0")), im.deref("arg0")))
+
+    applied_as_fieldop_tmp = im.as_fieldop(stencil)(im.ref("in_field"))
+    applied_as_fieldop = im.as_fieldop(stencil)(im.ref("tmp"))
+
+    domain_tmp = im.domain(common.GridType.CARTESIAN, {IDim: (0, 12)})
+    domain = im.domain(common.GridType.CARTESIAN, {IDim: (0, 11)})
+
+    params = [im.sym(name) for name in ["in_field", "out_field", "_gtmp_auto_domain"]]
+
+    testee = itir.Program(
+        id="forward_diff_with_tmp",
+        function_definitions=[],
+        params=params,
+        declarations=[itir.Temporary(id="tmp", domain=AUTO_DOMAIN, dtype=float_type)],
+        body=[
+            itir.SetAt(expr=applied_as_fieldop_tmp, domain=AUTO_DOMAIN, target=im.ref("tmp")),
+            itir.SetAt(expr=applied_as_fieldop, domain=domain, target=im.ref("out_field")),
+        ],
+    )
+
+    expected_expr_tmp = im.as_fieldop(stencil, domain_tmp)(im.ref("in_field"))
+    expected_epxr = im.as_fieldop(stencil, domain)(im.ref("tmp"))
+
+    expected = itir.Program(
+        id="forward_diff_with_tmp",
+        function_definitions=[],
+        params=params,
+        declarations=[itir.Temporary(id="tmp", domain=domain_tmp, dtype=float_type)],
+        body=[
+            itir.SetAt(expr=expected_expr_tmp, domain=domain_tmp, target=im.ref("tmp")),
+            itir.SetAt(expr=expected_epxr, domain=domain, target=im.ref("out_field")),
+        ],
+    )
+
+    run_test_program(testee, expected, offset_provider)
+
+
+def test_program_two_tmps(offset_provider):
+    stencil = im.lambda_("arg0")(im.minus(im.deref(im.shift("Ioff", 1)("arg0")), im.deref("arg0")))
+
+    as_fieldop_tmp1 = im.as_fieldop(stencil)(im.ref("in_field"))
+    as_fieldop_tmp2 = im.as_fieldop(stencil)(im.ref("tmp1"))
+    as_fieldop = im.as_fieldop(stencil)(im.ref("tmp2"))
+
+    domain = im.domain(common.GridType.CARTESIAN, {IDim: (0, 11)})
+    domain_tmp1 = im.domain(common.GridType.CARTESIAN, {IDim: (0, 13)})
+    domain_tmp2 = im.domain(common.GridType.CARTESIAN, {IDim: (0, 12)})
+
+    params = [im.sym(name) for name in ["in_field", "out_field", "_gtmp_auto_domain"]]
+
+    testee = itir.Program(
+        id="forward_diff_with_two_tmps",
+        function_definitions=[],
+        params=params,
+        declarations=[
+            itir.Temporary(id="tmp1", domain=AUTO_DOMAIN, dtype=float_type),
+            itir.Temporary(id="tmp2", domain=AUTO_DOMAIN, dtype=float_type),
+        ],
+        body=[
+            itir.SetAt(expr=as_fieldop_tmp1, domain=AUTO_DOMAIN, target=im.ref("tmp1")),
+            itir.SetAt(expr=as_fieldop_tmp2, domain=AUTO_DOMAIN, target=im.ref("tmp2")),
+            itir.SetAt(expr=as_fieldop, domain=domain, target=im.ref("out_field")),
+        ],
+    )
+
+    expected_expr_tmp1 = im.as_fieldop(stencil, domain_tmp1)(im.ref("in_field"))
+    expected_expr_tmp2 = im.as_fieldop(stencil, domain_tmp2)(im.ref("tmp1"))
+    expected_expr = im.as_fieldop(stencil, domain)(im.ref("tmp2"))
+
+    expected = itir.Program(
+        id="forward_diff_with_two_tmps",
+        function_definitions=[],
+        params=params,
+        declarations=[
+            itir.Temporary(id="tmp1", domain=domain_tmp1, dtype=float_type),
+            itir.Temporary(id="tmp2", domain=domain_tmp2, dtype=float_type),
+        ],
+        body=[
+            itir.SetAt(expr=expected_expr_tmp1, domain=domain_tmp1, target=im.ref("tmp1")),
+            itir.SetAt(expr=expected_expr_tmp2, domain=domain_tmp2, target=im.ref("tmp2")),
+            itir.SetAt(expr=expected_expr, domain=domain, target=im.ref("out_field")),
+        ],
+    )
+
+    run_test_program(testee, expected, offset_provider)
+
+
+@pytest.mark.xfail(raises=ValueError)
+def test_program_ValueError(offset_provider):
+    with pytest.raises(ValueError, match=r"Temporaries can only be used once within a program."):
+        stencil = im.lambda_("arg0")(im.deref("arg0"))
+
+        as_fieldop_tmp = im.as_fieldop(stencil)(im.ref("in_field"))
+        as_fieldop = im.as_fieldop(stencil)(im.ref("tmp"))
+
+        domain = im.domain(common.GridType.CARTESIAN, {IDim: (0, 11)})
+
+        params = [im.sym(name) for name in ["in_field", "out_field", "_gtmp_auto_domain"]]
+
+        infer_program(
+            itir.Program(
+                id="forward_diff_with_tmp",
+                function_definitions=[],
+                params=params,
+                declarations=[itir.Temporary(id="tmp", domain=AUTO_DOMAIN, dtype=float_type)],
+                body=[
+                    # target occurs twice here which is prohibited
+                    itir.SetAt(expr=as_fieldop_tmp, domain=AUTO_DOMAIN, target=im.ref("tmp")),
+                    itir.SetAt(expr=as_fieldop_tmp, domain=AUTO_DOMAIN, target=im.ref("tmp")),
+                    itir.SetAt(expr=as_fieldop, domain=domain, target=im.ref("out_field")),
+                ],
+            ),
+            offset_provider,
+        )
+
+
+def test_program_tree_tmps_two_inputs(offset_provider):
+    stencil = im.lambda_("arg0", "arg1")(
+        im.minus(im.deref(im.shift("Ioff", 1)("arg0")), im.deref("arg1"))
+    )
+    stencil_tmp = im.lambda_("arg0")(
+        im.minus(im.deref(im.shift("Ioff", 1)("arg0")), im.deref("arg0"))
+    )
+    stencil_tmp_minus = im.lambda_("arg0", "arg1")(
+        im.minus(im.deref(im.shift("Ioff", -1)("arg0")), im.deref("arg1"))
+    )
+
+    as_fieldop_tmp1 = im.as_fieldop(stencil)(im.ref("in_field1"), im.ref("in_field2"))
+    as_fieldop_tmp2 = im.as_fieldop(stencil_tmp)(im.ref("tmp1"))
+    as_fieldop_out1 = im.as_fieldop(stencil_tmp)(im.ref("tmp2"))
+    as_fieldop_tmp3 = im.as_fieldop(stencil)(im.ref("tmp1"), im.ref("in_field2"))
+    as_fieldop_out2 = im.as_fieldop(stencil_tmp_minus)(im.ref("tmp2"), im.ref("tmp3"))
+
+    domain_tmp1 = im.domain(common.GridType.CARTESIAN, {IDim: (-1, 13)})
+    domain_tmp2 = im.domain(common.GridType.CARTESIAN, {IDim: (-1, 12)})
+    domain_tmp3 = im.domain(common.GridType.CARTESIAN, {IDim: (0, 11)})
+    domain_out = im.domain(common.GridType.CARTESIAN, {IDim: (0, 11)})
+    params = [
+        im.sym(name)
+        for name in ["in_field1", "in_field2", "out_field1", "out_field2", "_gtmp_auto_domain"]
+    ]
+
+    testee = itir.Program(
+        id="differences_three_tmps_two_inputs",
+        function_definitions=[],
+        params=params,
+        declarations=[
+            itir.Temporary(id="tmp1", domain=AUTO_DOMAIN, dtype=float_type),
+            itir.Temporary(id="tmp2", domain=AUTO_DOMAIN, dtype=float_type),
+            itir.Temporary(id="tmp3", domain=AUTO_DOMAIN, dtype=float_type),
+        ],
+        body=[
+            itir.SetAt(expr=as_fieldop_tmp1, domain=AUTO_DOMAIN, target=im.ref("tmp1")),
+            itir.SetAt(expr=as_fieldop_tmp2, domain=AUTO_DOMAIN, target=im.ref("tmp2")),
+            itir.SetAt(expr=as_fieldop_out1, domain=domain_out, target=im.ref("out_field1")),
+            itir.SetAt(expr=as_fieldop_tmp3, domain=AUTO_DOMAIN, target=im.ref("tmp3")),
+            itir.SetAt(expr=as_fieldop_out2, domain=domain_out, target=im.ref("out_field2")),
+        ],
+    )
+
+    expected_expr_tmp1 = im.as_fieldop(stencil, domain_tmp1)(
+        im.ref("in_field1"), im.ref("in_field2")
+    )
+    expected_expr_tmp2 = im.as_fieldop(stencil_tmp, domain_tmp2)(im.ref("tmp1"))
+    expected_expr_out1 = im.as_fieldop(stencil_tmp, domain_out)(im.ref("tmp2"))
+    expected_expr_tmp3 = im.as_fieldop(stencil, domain_tmp3)(im.ref("tmp1"), im.ref("in_field2"))
+    expected_expr_out2 = im.as_fieldop(stencil_tmp_minus, domain_out)(
+        im.ref("tmp2"), im.ref("tmp3")
+    )
+
+    expected = itir.Program(
+        id="differences_three_tmps_two_inputs",
+        function_definitions=[],
+        params=params,
+        declarations=[
+            itir.Temporary(id="tmp1", domain=domain_tmp1, dtype=float_type),
+            itir.Temporary(id="tmp2", domain=domain_tmp2, dtype=float_type),
+            itir.Temporary(id="tmp3", domain=domain_tmp3, dtype=float_type),
+        ],
+        body=[
+            itir.SetAt(expr=expected_expr_tmp1, domain=domain_tmp1, target=im.ref("tmp1")),
+            itir.SetAt(expr=expected_expr_tmp2, domain=domain_tmp2, target=im.ref("tmp2")),
+            itir.SetAt(expr=expected_expr_out1, domain=domain_out, target=im.ref("out_field1")),
+            itir.SetAt(expr=expected_expr_tmp3, domain=domain_tmp3, target=im.ref("tmp3")),
+            itir.SetAt(expr=expected_expr_out2, domain=domain_out, target=im.ref("out_field2")),
+        ],
+    )
+
+    run_test_program(testee, expected, offset_provider)