Repository: GridTools/gt4py
Pull request #1524 (2024-04-09 08:59:36)

refactor[next]: Fencil to itir.Program for gtfn
First PR for preparing itir for a combined field view + iterator representation.

Adds new nodes:
- itir.Program (to replace itir.FencilDefinition)
- itir.SetAt (to replace itir.StencilClosure)
and a new builtin `as_fieldop`.

The semantic of `SetAt` is that the `expr` is directly computed into the `target` field. `as_fieldop` aka map, takes an itir stencil (a function from iterators to values) and promotes it to a field_operator, i.e. a function from fields to fields.

The idea for `itir.Program` is to align and ultimately merge with `past.Program`.

In this PR,  the transition from Fencil to Program is implemented only for GTFN with a separate pass before lowering to gtfn.

Additionally the pretty_printer/parser is extended with the new nodes/builtin.


Patch:
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 4f1eafeb..a01cf095 100644
--- a/src/gt4py/next/ffront/decorator.py
+++ b/src/gt4py/next/ffront/decorator.py
@@ -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/iterator/ir.py b/src/gt4py/next/iterator/ir.py
index 56f931f4..5d1907b2 100644
--- a/src/gt4py/next/iterator/ir.py
+++ b/src/gt4py/next/iterator/ir.py
@@ -13,7 +13,7 @@
 # SPDX-License-Identifier: GPL-3.0-or-later
 
 import typing
-from typing import ClassVar, List, Optional, Union
+from typing import Any, ClassVar, List, Optional, Union
 
 import gt4py.eve as eve
 from gt4py.eve import Coerced, SymbolName, SymbolRef, datamodels
@@ -22,6 +22,11 @@ from gt4py.eve.traits import SymbolTableTrait, ValidatedSymbolTableTrait
 from gt4py.eve.utils import noninstantiable
 
 
+# TODO(havogt):
+# After completion of refactoring to GTIR, FencilDefinition and StencilClosure should be removed everywhere.
+# During transition, we lower to FencilDefinitions and apply a transformation to GTIR-style afterwards.
+
+
 @noninstantiable
 class Node(eve.Node):
     location: Optional[SourceLocation] = eve.field(default=None, repr=False, compare=False)
@@ -202,6 +207,13 @@ BUILTINS = {
     *TYPEBUILTINS,
 }
 
+# only used in `Program`` not `FencilDefinition`
+# TODO(havogt): restructure after refactoring to GTIR
+GTIR_BUILTINS = {
+    *BUILTINS,
+    "as_fieldop",  # `as_fieldop(stencil)` creates field_operator from stencil
+}
+
 
 class FencilDefinition(Node, ValidatedSymbolTableTrait):
     id: Coerced[SymbolName]
@@ -212,6 +224,31 @@ class FencilDefinition(Node, ValidatedSymbolTableTrait):
     _NODE_SYMBOLS_: ClassVar[List[Sym]] = [Sym(id=name) for name in BUILTINS]
 
 
+class Stmt(Node): ...
+
+
+class SetAt(Stmt):  # from JAX array.at[...].set()
+    expr: Expr  # only `as_fieldop(stencil)(inp0, ...)` in first refactoring
+    domain: Expr
+    target: Expr  # `make_tuple` or SymRef
+
+
+class Temporary(Node):
+    id: Coerced[eve.SymbolName]
+    domain: Optional[Expr] = None
+    dtype: Optional[Any] = None  # TODO
+
+
+class Program(Node, ValidatedSymbolTableTrait):
+    id: Coerced[SymbolName]
+    function_definitions: List[FunctionDefinition]
+    params: List[Sym]
+    declarations: List[Temporary]
+    body: List[Stmt]
+
+    _NODE_SYMBOLS_: ClassVar[List[Sym]] = [Sym(id=name) for name in GTIR_BUILTINS]
+
+
 # TODO(fthaler): just use hashable types in nodes (tuples instead of lists)
 Sym.__hash__ = Node.__hash__  # type: ignore[method-assign]
 Expr.__hash__ = Node.__hash__  # type: ignore[method-assign]
diff --git a/src/gt4py/next/iterator/pretty_parser.py b/src/gt4py/next/iterator/pretty_parser.py
index 9dd96b07..f6d532ee 100644
--- a/src/gt4py/next/iterator/pretty_parser.py
+++ b/src/gt4py/next/iterator/pretty_parser.py
@@ -23,7 +23,10 @@ from gt4py.next.iterator.ir_utils import ir_makers as im
 GRAMMAR = """
     start: fencil_definition
         | function_definition
+        | declaration
         | stencil_closure
+        | set_at
+        | program
         | prec0
 
     SYM: CNAME
@@ -64,6 +67,7 @@ GRAMMAR = """
         | "·" prec7 -> deref
         | "¬" prec7 -> bool_not
         | "↑" prec7 -> lift
+        | "⇑" prec7 -> as_fieldop
 
     ?prec8: prec9
         | prec8 "[" prec0 "]" -> tuple_get
@@ -80,8 +84,11 @@ GRAMMAR = """
 
     named_range: AXIS_NAME ":" "[" prec0 "," prec0 ")"
     function_definition: ID_NAME "=" "λ(" ( SYM "," )* SYM? ")" "→" prec0 ";"
+    declaration: ID_NAME "=" "temporary(" "domain=" prec0 "," "dtype=" prec0 ")" ";"
     stencil_closure: prec0 "←" "(" prec0 ")" "(" ( SYM_REF ", " )* SYM_REF ")" "@" prec0 ";"
+    set_at: prec0 "@" prec0 "←" prec1 ";"
     fencil_definition: ID_NAME "(" ( SYM "," )* SYM ")" "{" ( function_definition )* ( stencil_closure )+ "}"
+    program: ID_NAME "(" ( SYM "," )* SYM ")" "{" ( function_definition )* ( declaration )* ( set_at )+ "}"
 
     %import common (CNAME, SIGNED_FLOAT, SIGNED_INT, WS)
     %ignore WS
@@ -167,6 +174,9 @@ class ToIrTransformer(lark_visitors.Transformer):
     def lift(self, arg: ir.Expr) -> ir.FunCall:
         return ir.FunCall(fun=ir.SymRef(id="lift"), args=[arg])
 
+    def as_fieldop(self, arg: ir.Expr) -> ir.FunCall:
+        return ir.FunCall(fun=ir.SymRef(id="as_fieldop"), args=[arg])
+
     def astype(self, arg: ir.Expr) -> ir.FunCall:
         return ir.FunCall(fun=ir.SymRef(id="cast_"), args=[arg])
 
@@ -202,6 +212,15 @@ class ToIrTransformer(lark_visitors.Transformer):
         output, stencil, *inputs, domain = args
         return ir.StencilClosure(domain=domain, stencil=stencil, output=output, inputs=inputs)
 
+    def declaration(self, *args: ir.Expr) -> ir.Temporary:
+        tid, domain, dtype = args
+        return ir.Temporary(id=tid, domain=domain, dtype=dtype)
+
+    def set_at(self, *args: ir.Expr) -> ir.SetAt:
+        target, domain, expr = args
+        return ir.SetAt(expr=expr, domain=domain, target=target)
+
+    # TODO(havogt): remove after refactoring.
     def fencil_definition(self, fid: str, *args: ir.Node) -> ir.FencilDefinition:
         params = []
         function_definitions = []
@@ -218,6 +237,29 @@ class ToIrTransformer(lark_visitors.Transformer):
             id=fid, function_definitions=function_definitions, params=params, closures=closures
         )
 
+    def program(self, fid: str, *args: ir.Node) -> ir.Program:
+        params = []
+        function_definitions = []
+        body = []
+        declarations = []
+        for arg in args:
+            if isinstance(arg, ir.Sym):
+                params.append(arg)
+            elif isinstance(arg, ir.FunctionDefinition):
+                function_definitions.append(arg)
+            elif isinstance(arg, ir.Temporary):
+                declarations.append(arg)
+            else:
+                assert isinstance(arg, ir.SetAt)
+                body.append(arg)
+        return ir.Program(
+            id=fid,
+            function_definitions=function_definitions,
+            params=params,
+            body=body,
+            declarations=declarations,
+        )
+
     def start(self, arg: ir.Node) -> ir.Node:
         return arg
 
diff --git a/src/gt4py/next/iterator/pretty_printer.py b/src/gt4py/next/iterator/pretty_printer.py
index 041c12e7..3f224d2e 100644
--- a/src/gt4py/next/iterator/pretty_printer.py
+++ b/src/gt4py/next/iterator/pretty_printer.py
@@ -43,7 +43,7 @@ BINARY_OPS: Final = {
 }
 
 # replacements for builtin unary operations
-UNARY_OPS: Final = {"deref": "·", "lift": "↑", "not_": "¬"}
+UNARY_OPS: Final = {"deref": "·", "lift": "↑", "not_": "¬", "as_fieldop": "⇑"}
 
 # operator precedence
 PRECEDENCE: Final = {
@@ -63,6 +63,7 @@ PRECEDENCE: Final = {
     "deref": 7,
     "not_": 7,
     "lift": 7,
+    "as_fieldop": 7,
     "tuple_get": 8,
     "__call__": 8,
 }
@@ -272,6 +273,35 @@ class PrettyPrinter(NodeTranslator):
         )
         return self._optimum(h, v)
 
+    def visit_Temporary(self, node: ir.Temporary, *, prec: int) -> list[str]:
+        start, end = [node.id + " = temporary("], [");"]
+        args = []
+        if node.domain is not None:
+            args.append(self._hmerge(["domain="], self.visit(node.domain, prec=0)))
+        if node.dtype is not None:
+            args.append(self._hmerge(["dtype="], [str(node.dtype)]))
+        hargs = self._hmerge(*self._hinterleave(args, ", "))
+        vargs = self._vmerge(*self._hinterleave(args, ","))
+        oargs = self._optimum(hargs, vargs)
+        h = self._hmerge(start, oargs, end)
+        v = self._vmerge(start, self._indent(oargs), end)
+        return self._optimum(h, v)
+
+    def visit_SetAt(self, node: ir.SetAt, *, prec: int) -> list[str]:
+        expr = self.visit(node.expr, prec=0)
+        domain = self.visit(node.domain, prec=0)
+        target = self.visit(node.target, prec=0)
+
+        head = self._hmerge(target, [" @ "], domain)
+        foot = self._hmerge([" ← "], expr, [";"])
+
+        h = self._hmerge(head, foot)
+        v = self._vmerge(
+            head,
+            self._indent(self._indent(foot)),
+        )
+        return self._optimum(h, v)
+
     def visit_FencilDefinition(self, node: ir.FencilDefinition, *, prec: int) -> list[str]:
         assert prec == 0
         function_definitions = self.visit(node.function_definitions, prec=0)
@@ -291,6 +321,31 @@ class PrettyPrinter(NodeTranslator):
             params, self._indent(function_definitions), self._indent(closures), ["}"]
         )
 
+    def visit_Program(self, node: ir.Program, *, prec: int) -> list[str]:
+        assert prec == 0
+        function_definitions = self.visit(node.function_definitions, prec=0)
+        body = self.visit(node.body, prec=0)
+        declarations = self.visit(node.declarations, prec=0)
+        params = self.visit(node.params, prec=0)
+
+        hparams = self._hmerge([node.id + "("], *self._hinterleave(params, ", "), [") {"])
+        vparams = self._vmerge(
+            [node.id + "("], *self._hinterleave(params, ",", indent=True), [") {"]
+        )
+        params = self._optimum(hparams, vparams)
+
+        function_definitions = self._vmerge(*function_definitions)
+        declarations = self._vmerge(*declarations)
+        body = self._vmerge(*body)
+
+        return self._vmerge(
+            params,
+            self._indent(function_definitions),
+            self._indent(declarations),
+            self._indent(body),
+            ["}"],
+        )
+
     @classmethod
     def apply(cls, node: ir.Node, indent: int, width: int) -> str:
         return "\n".join(cls(indent=indent, width=width).visit(node, prec=0))
diff --git a/src/gt4py/next/iterator/transforms/fencil_to_program.py b/src/gt4py/next/iterator/transforms/fencil_to_program.py
new file mode 100644
index 00000000..fff12a6c
--- /dev/null
+++ b/src/gt4py/next/iterator/transforms/fencil_to_program.py
@@ -0,0 +1,46 @@
+# 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
+
+from gt4py import eve
+from gt4py.next.iterator import ir as itir
+from gt4py.next.iterator.ir_utils import ir_makers as im
+from gt4py.next.iterator.transforms import global_tmps
+
+
+class FencilToProgram(eve.NodeTranslator):
+    @classmethod
+    def apply(cls, node: itir.FencilDefinition | global_tmps.FencilWithTemporaries) -> itir.Program:
+        return cls().visit(node)
+
+    def visit_StencilClosure(self, node: itir.StencilClosure) -> itir.SetAt:
+        as_fieldop = im.call(im.call("as_fieldop")(node.stencil, node.domain))(*node.inputs)
+        return itir.SetAt(expr=as_fieldop, domain=node.domain, target=node.output)
+
+    def visit_FencilDefinition(self, node: itir.FencilDefinition) -> itir.Program:
+        return itir.Program(
+            id=node.id,
+            function_definitions=node.function_definitions,
+            params=node.params,
+            declarations=[],
+            body=self.visit(node.closures),
+        )
+
+    def visit_FencilWithTemporaries(self, node: global_tmps.FencilWithTemporaries) -> itir.Program:
+        return itir.Program(
+            id=node.fencil.id,
+            function_definitions=node.fencil.function_definitions,
+            params=node.params,
+            declarations=node.tmps,
+            body=self.visit(node.fencil.closures),
+        )
diff --git a/src/gt4py/next/iterator/transforms/global_tmps.py b/src/gt4py/next/iterator/transforms/global_tmps.py
index d099272b..42d68318 100644
--- a/src/gt4py/next/iterator/transforms/global_tmps.py
+++ b/src/gt4py/next/iterator/transforms/global_tmps.py
@@ -17,9 +17,8 @@ import functools
 from collections.abc import Mapping
 from typing import Any, Callable, Final, Iterable, Literal, Optional, Sequence
 
-import gt4py.eve as eve
 import gt4py.next as gtx
-from gt4py.eve import Coerced, NodeTranslator, PreserveLocationVisitor
+from gt4py.eve import NodeTranslator, PreserveLocationVisitor
 from gt4py.eve.traits import SymbolTableTrait
 from gt4py.eve.utils import UIDGenerator
 from gt4py.next import common
@@ -54,40 +53,19 @@ AUTO_DOMAIN: Final = ir.FunCall(fun=ir.SymRef(id="_gtmp_auto_domain"), args=[])
 # Iterator IR extension nodes
 
 
-class Temporary(ir.Node):
-    """Iterator IR extension: declaration of a temporary buffer."""
-
-    id: Coerced[eve.SymbolName]
-    domain: Optional[ir.Expr] = None
-    dtype: Optional[Any] = None
-
-
-class FencilWithTemporaries(ir.Node, SymbolTableTrait):
+class FencilWithTemporaries(
+    ir.Node, SymbolTableTrait
+):  # TODO(havogt): remove and use new `itir.Program` instead.
     """Iterator IR extension: declaration of a fencil with temporary buffers."""
 
     fencil: ir.FencilDefinition
     params: list[ir.Sym]
-    tmps: list[Temporary]
+    tmps: list[ir.Temporary]
 
 
 # Extensions for `PrettyPrinter` for easier debugging
 
 
-def pformat_Temporary(printer: PrettyPrinter, node: Temporary, *, prec: int) -> list[str]:
-    start, end = [node.id + " = temporary("], [");"]
-    args = []
-    if node.domain is not None:
-        args.append(printer._hmerge(["domain="], printer.visit(node.domain, prec=0)))
-    if node.dtype is not None:
-        args.append(printer._hmerge(["dtype="], [str(node.dtype)]))
-    hargs = printer._hmerge(*printer._hinterleave(args, ", "))
-    vargs = printer._vmerge(*printer._hinterleave(args, ","))
-    oargs = printer._optimum(hargs, vargs)
-    h = printer._hmerge(start, oargs, end)
-    v = printer._vmerge(start, printer._indent(oargs), end)
-    return printer._optimum(h, v)
-
-
 def pformat_FencilWithTemporaries(
     printer: PrettyPrinter, node: FencilWithTemporaries, *, prec: int
 ) -> list[str]:
@@ -117,7 +95,6 @@ def pformat_FencilWithTemporaries(
     return printer._vmerge(params, printer._indent(body), ["}"])
 
 
-PrettyPrinter.visit_Temporary = pformat_Temporary  # type: ignore
 PrettyPrinter.visit_FencilWithTemporaries = pformat_FencilWithTemporaries  # type: ignore
 
 
@@ -367,7 +344,7 @@ def split_closures(
             location=node.location,
         ),
         params=node.params,
-        tmps=[Temporary(id=tmp.id) for tmp in tmps],
+        tmps=[ir.Temporary(id=tmp.id) for tmp in tmps],
     )
 
 
@@ -638,7 +615,8 @@ def collect_tmps_info(node: FencilWithTemporaries, *, offset_provider) -> Fencil
         fencil=node.fencil,
         params=node.params,
         tmps=[
-            Temporary(id=tmp.id, domain=domains[tmp.id], dtype=types[tmp.id]) for tmp in node.tmps
+            ir.Temporary(id=tmp.id, domain=domains[tmp.id], dtype=types[tmp.id])
+            for tmp in node.tmps
         ],
     )
 
diff --git a/src/gt4py/next/iterator/transforms/pass_manager.py b/src/gt4py/next/iterator/transforms/pass_manager.py
index 5852ba9a..8a192032 100644
--- a/src/gt4py/next/iterator/transforms/pass_manager.py
+++ b/src/gt4py/next/iterator/transforms/pass_manager.py
@@ -24,7 +24,7 @@ from gt4py.next.iterator.transforms.constant_folding import ConstantFolding
 from gt4py.next.iterator.transforms.cse import CommonSubexpressionElimination
 from gt4py.next.iterator.transforms.eta_reduction import EtaReduction
 from gt4py.next.iterator.transforms.fuse_maps import FuseMaps
-from gt4py.next.iterator.transforms.global_tmps import CreateGlobalTmps
+from gt4py.next.iterator.transforms.global_tmps import CreateGlobalTmps, FencilWithTemporaries
 from gt4py.next.iterator.transforms.inline_center_deref_lift_vars import InlineCenterDerefLiftVars
 from gt4py.next.iterator.transforms.inline_fundefs import InlineFundefs, PruneUnreferencedFundefs
 from gt4py.next.iterator.transforms.inline_into_scan import InlineIntoScan
@@ -88,7 +88,7 @@ def apply_common_transforms(
         Callable[[itir.StencilClosure], Callable[[itir.Expr], bool]]
     ] = None,
     symbolic_domain_sizes: Optional[dict[str, str]] = None,
-):
+) -> itir.FencilDefinition | FencilWithTemporaries:
     icdlv_uids = eve_utils.UIDGenerator()
 
     if lift_mode is None:
@@ -203,4 +203,5 @@ def apply_common_transforms(
         ir, opcount_preserving=True, force_inline_lambda_args=force_inline_lambda_args
     )
 
+    assert isinstance(ir, (itir.FencilDefinition, FencilWithTemporaries))
     return ir
diff --git a/src/gt4py/next/program_processors/codegens/gtfn/codegen.py b/src/gt4py/next/program_processors/codegens/gtfn/codegen.py
index 987fe6ca..cc07fac5 100644
--- a/src/gt4py/next/program_processors/codegens/gtfn/codegen.py
+++ b/src/gt4py/next/program_processors/codegens/gtfn/codegen.py
@@ -234,9 +234,7 @@ class GTFNCodegen(codegen.TemplatedGenerator):
         "auto {id} = gtfn::allocate_global_tmp<{dtype}>(tmp_alloc__, {tmp_sizes});"
     )
 
-    def visit_FencilDefinition(
-        self, node: gtfn_ir.FencilDefinition, **kwargs: Any
-    ) -> Union[str, Collection[str]]:
+    def visit_Program(self, node: gtfn_ir.Program, **kwargs: Any) -> Union[str, Collection[str]]:
         self.is_cartesian = node.grid_type == common.GridType.CARTESIAN
         self.user_defined_function_ids = list(
             str(fundef.id) for fundef in node.function_definitions
@@ -248,7 +246,7 @@ class GTFNCodegen(codegen.TemplatedGenerator):
             **kwargs,
         )
 
-    FencilDefinition = as_mako(
+    Program = as_mako(
         """
     #include <cmath>
     #include <cstdint>
diff --git a/src/gt4py/next/program_processors/codegens/gtfn/gtfn_ir.py b/src/gt4py/next/program_processors/codegens/gtfn/gtfn_ir.py
index b266f577..65ea6c91 100644
--- a/src/gt4py/next/program_processors/codegens/gtfn/gtfn_ir.py
+++ b/src/gt4py/next/program_processors/codegens/gtfn/gtfn_ir.py
@@ -155,7 +155,7 @@ class TagDefinition(Node):
     alias: Optional[Union[str, SymRef]] = None
 
 
-class FencilDefinition(Node, ValidatedSymbolTableTrait):
+class Program(Node, ValidatedSymbolTableTrait):
     id: SymbolName
     params: list[Sym]
     function_definitions: list[
diff --git a/src/gt4py/next/program_processors/codegens/gtfn/gtfn_module.py b/src/gt4py/next/program_processors/codegens/gtfn/gtfn_module.py
index ca293aa2..d28e5130 100644
--- a/src/gt4py/next/program_processors/codegens/gtfn/gtfn_module.py
+++ b/src/gt4py/next/program_processors/codegens/gtfn/gtfn_module.py
@@ -28,7 +28,7 @@ from gt4py.next import common
 from gt4py.next.common import Connectivity, Dimension
 from gt4py.next.ffront import fbuiltins
 from gt4py.next.iterator import ir as itir
-from gt4py.next.iterator.transforms import LiftMode, pass_manager
+from gt4py.next.iterator.transforms import LiftMode, fencil_to_program, global_tmps, pass_manager
 from gt4py.next.otf import languages, stages, step_types, workflow
 from gt4py.next.otf.binding import cpp_interface, interface
 from gt4py.next.program_processors.codegens.gtfn.codegen import GTFNCodegen, GTFNIMCodegen
@@ -183,7 +183,7 @@ class GTFNTranslationStep(
         program: itir.FencilDefinition,
         offset_provider: dict[str, Connectivity | Dimension],
         runtime_lift_mode: Optional[LiftMode],
-    ) -> itir.FencilDefinition:
+    ) -> itir.FencilDefinition | global_tmps.FencilWithTemporaries:
         # 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
@@ -231,8 +231,11 @@ class GTFNTranslationStep(
         runtime_lift_mode: Optional[LiftMode] = None,
     ) -> str:
         new_program = self._preprocess_program(program, offset_provider, runtime_lift_mode)
+        program_itir = fencil_to_program.FencilToProgram().apply(
+            new_program
+        )  # TODO(havogt): should be removed after refactoring to combined IR
         gtfn_ir = GTFN_lowering.apply(
-            new_program, offset_provider=offset_provider, column_axis=column_axis
+            program_itir, offset_provider=offset_provider, column_axis=column_axis
         )
 
         if self.use_imperative_backend:
diff --git a/src/gt4py/next/program_processors/codegens/gtfn/itir_to_gtfn_ir.py b/src/gt4py/next/program_processors/codegens/gtfn/itir_to_gtfn_ir.py
index 4617e54e..bbde7206 100644
--- a/src/gt4py/next/program_processors/codegens/gtfn/itir_to_gtfn_ir.py
+++ b/src/gt4py/next/program_processors/codegens/gtfn/itir_to_gtfn_ir.py
@@ -16,23 +16,22 @@ import dataclasses
 from typing import Any, ClassVar, Iterable, Optional, Type, TypeGuard, Union
 
 import gt4py.eve as eve
+from gt4py.eve import utils as eve_utils
 from gt4py.eve.concepts import SymbolName
-from gt4py.eve.utils import UIDGenerator
 from gt4py.next import common
 from gt4py.next.iterator import ir as itir
-from gt4py.next.iterator.transforms import global_tmps
 from gt4py.next.program_processors.codegens.gtfn.gtfn_ir import (
     Backend,
     BinaryExpr,
     CartesianDomain,
     CastExpr,
-    FencilDefinition,
     FunCall,
     FunctionDefinition,
     IntegralConstant,
     Lambda,
     Literal,
     OffsetLiteral,
+    Program,
     Scan,
     ScanExecution,
     ScanPassDefinition,
@@ -66,8 +65,8 @@ _vertical_dimension = "gtfn::unstructured::dim::vertical"
 _horizontal_dimension = "gtfn::unstructured::dim::horizontal"
 
 
-def _get_domains(closures: Iterable[itir.StencilClosure]) -> Iterable[itir.FunCall]:
-    return (c.domain for c in closures)
+def _get_domains(node: Iterable[itir.Stmt]) -> Iterable[itir.FunCall]:
+    return eve_utils.xiter(node).if_isinstance(itir.SetAt).getattr("domain").to_set()
 
 
 def _extract_grid_type(domain: itir.FunCall) -> common.GridType:
@@ -78,8 +77,8 @@ def _extract_grid_type(domain: itir.FunCall) -> common.GridType:
         return common.GridType.UNSTRUCTURED
 
 
-def _get_gridtype(closures: list[itir.StencilClosure]) -> common.GridType:
-    domains = _get_domains(closures)
+def _get_gridtype(body: list[itir.Stmt]) -> common.GridType:
+    domains = _get_domains(body)
     grid_types = {_extract_grid_type(d) for d in domains}
     if len(grid_types) != 1:
         raise ValueError(
@@ -97,9 +96,9 @@ def _name_from_named_range(named_range_call: itir.FunCall) -> str:
 
 
 def _collect_dimensions_from_domain(
-    closures: Iterable[itir.StencilClosure],
+    body: Iterable[itir.Stmt],
 ) -> dict[str, TagDefinition]:
-    domains = _get_domains(closures)
+    domains = _get_domains(body)
     offset_definitions = {}
     for domain in domains:
         if domain.fun == itir.SymRef(id="cartesian_domain"):
@@ -198,6 +197,14 @@ def _bool_from_literal(node: itir.Node) -> bool:
     return node.value == "True"
 
 
+def _is_applied_as_fieldop(arg: itir.Expr) -> TypeGuard[itir.FunCall]:
+    return (
+        isinstance(arg, itir.FunCall)
+        and isinstance(arg.fun, itir.FunCall)
+        and arg.fun.fun == itir.SymRef(id="as_fieldop")
+    )
+
+
 @dataclasses.dataclass(frozen=True)
 class GTFN_lowering(eve.NodeTranslator, eve.VisitorWithSymbolTableTrait):
     _binary_op_map: ClassVar[dict[str, str]] = {
@@ -224,26 +231,22 @@ class GTFN_lowering(eve.NodeTranslator, eve.VisitorWithSymbolTableTrait):
 
     # we use one UID generator per instance such that the generated ids are
     # stable across multiple runs (required for caching to properly work)
-    uids: UIDGenerator = dataclasses.field(init=False, repr=False, default_factory=UIDGenerator)
+    uids: eve_utils.UIDGenerator = dataclasses.field(
+        init=False, repr=False, default_factory=eve_utils.UIDGenerator
+    )
 
     @classmethod
     def apply(
         cls,
-        node: itir.FencilDefinition | global_tmps.FencilWithTemporaries,
+        node: itir.Program,
         *,
         offset_provider: dict,
         column_axis: Optional[common.Dimension],
-    ) -> FencilDefinition:
-        if isinstance(node, global_tmps.FencilWithTemporaries):
-            fencil_definition = node.fencil
-        elif isinstance(node, itir.FencilDefinition):
-            fencil_definition = node
-        else:
-            raise TypeError(
-                f"Expected a 'FencilDefinition' or 'FencilWithTemporaries', got '{type(node).__name__}'."
-            )
+    ) -> Program:
+        if not isinstance(node, itir.Program):
+            raise TypeError(f"Expected a 'Program', got '{type(node).__name__}'.")
 
-        grid_type = _get_gridtype(fencil_definition.closures)
+        grid_type = _get_gridtype(node.body)
         return cls(
             offset_provider=offset_provider, column_axis=column_axis, grid_type=grid_type
         ).visit(node)
@@ -434,44 +437,6 @@ class GTFN_lowering(eve.NodeTranslator, eve.VisitorWithSymbolTableTrait):
             return SidComposite(values=[self._visit_output_argument(v) for v in node.args])
         raise ValueError("Expected 'SymRef' or 'make_tuple' in output argument.")
 
-    def visit_StencilClosure(
-        self, node: itir.StencilClosure, extracted_functions: list, **kwargs: Any
-    ) -> Union[ScanExecution, StencilExecution]:
-        backend = Backend(domain=self.visit(node.domain, stencil=node.stencil, **kwargs))
-        if _is_scan(node.stencil):
-            scan_id = self.uids.sequential_id(prefix="_scan")
-            scan_lambda = self.visit(node.stencil.args[0], **kwargs)
-            forward = _bool_from_literal(node.stencil.args[1])
-            scan_def = ScanPassDefinition(
-                id=scan_id, params=scan_lambda.params, expr=scan_lambda.expr, forward=forward
-            )
-            extracted_functions.append(scan_def)
-            scan = Scan(
-                function=SymRef(id=scan_id),
-                output=0,
-                inputs=[i + 1 for i, _ in enumerate(node.inputs)],
-                init=self.visit(node.stencil.args[2], **kwargs),
-            )
-            column_axis = self.column_axis
-            assert isinstance(column_axis, common.Dimension)
-            return ScanExecution(
-                backend=backend,
-                scans=[scan],
-                args=[self._visit_output_argument(node.output), *self.visit(node.inputs)],
-                axis=SymRef(id=column_axis.value),
-            )
-        return StencilExecution(
-            stencil=self.visit(
-                node.stencil,
-                force_function_extraction=True,
-                extracted_functions=extracted_functions,
-                **kwargs,
-            ),
-            output=self._visit_output_argument(node.output),
-            inputs=self.visit(node.inputs, **kwargs),
-            backend=backend,
-        )
-
     @staticmethod
     def _merge_scans(
         executions: list[Union[StencilExecution, ScanExecution]],
@@ -517,29 +482,72 @@ class GTFN_lowering(eve.NodeTranslator, eve.VisitorWithSymbolTableTrait):
                 res.append(execution)
         return res
 
-    def visit_FencilDefinition(
-        self, node: itir.FencilDefinition, **kwargs: Any
-    ) -> FencilDefinition:
+    def visit_Stmt(self, node: itir.Stmt, **kwargs: Any) -> None:
+        raise AssertionError("All Stmts need to be handled explicitly.")
+
+    def visit_SetAt(
+        self, node: itir.SetAt, *, extracted_functions: list, **kwargs: Any
+    ) -> Union[StencilExecution, ScanExecution]:
+        assert _is_applied_as_fieldop(node.expr)
+        stencil = node.expr.fun.args[0]  # type: ignore[attr-defined] # checked in assert
+        domain = node.domain
+        inputs = node.expr.args
+        backend = Backend(domain=self.visit(domain, stencil=stencil, **kwargs))
+        if _is_scan(stencil):
+            scan_id = self.uids.sequential_id(prefix="_scan")
+            scan_lambda = self.visit(stencil.args[0], **kwargs)
+            forward = _bool_from_literal(stencil.args[1])
+            scan_def = ScanPassDefinition(
+                id=scan_id, params=scan_lambda.params, expr=scan_lambda.expr, forward=forward
+            )
+            extracted_functions.append(scan_def)
+            scan = Scan(
+                function=SymRef(id=scan_id),
+                output=0,
+                inputs=[i + 1 for i, _ in enumerate(inputs)],
+                init=self.visit(stencil.args[2], **kwargs),
+            )
+            column_axis = self.column_axis
+            assert isinstance(column_axis, common.Dimension)
+            return ScanExecution(
+                backend=backend,
+                scans=[scan],
+                args=[self._visit_output_argument(node.target), *self.visit(inputs)],
+                axis=SymRef(id=column_axis.value),
+            )
+        return StencilExecution(
+            stencil=self.visit(
+                stencil,
+                force_function_extraction=True,
+                extracted_functions=extracted_functions,
+                **kwargs,
+            ),
+            output=self._visit_output_argument(node.target),
+            inputs=self.visit(inputs, **kwargs),
+            backend=backend,
+        )
+
+    def visit_Program(self, node: itir.Program, **kwargs: Any) -> Program:
         extracted_functions: list[Union[FunctionDefinition, ScanPassDefinition]] = []
-        executions = self.visit(node.closures, extracted_functions=extracted_functions)
+        executions = self.visit(node.body, extracted_functions=extracted_functions)
         executions = self._merge_scans(executions)
         function_definitions = self.visit(node.function_definitions) + extracted_functions
         offset_definitions = {
-            **_collect_dimensions_from_domain(node.closures),
+            **_collect_dimensions_from_domain(node.body),
             **_collect_offset_definitions(node, self.grid_type, self.offset_provider),
         }
-        return FencilDefinition(
+        return Program(
             id=SymbolName(node.id),
             params=self.visit(node.params),
             executions=executions,
             grid_type=self.grid_type,
             offset_definitions=list(offset_definitions.values()),
             function_definitions=function_definitions,
-            temporaries=[],
+            temporaries=self.visit(node.declarations, params=[p.id for p in node.params]),
         )
 
     def visit_Temporary(
-        self, node: global_tmps.Temporary, *, params: list, **kwargs: Any
+        self, node: itir.Temporary, *, params: list, **kwargs: Any
     ) -> TemporaryAllocation:
         def dtype_to_cpp(x: int | tuple | str) -> str:
             if isinstance(x, int):
@@ -557,17 +565,3 @@ class GTFN_lowering(eve.NodeTranslator, eve.VisitorWithSymbolTableTrait):
         return TemporaryAllocation(
             id=node.id, dtype=dtype_to_cpp(node.dtype), domain=self.visit(node.domain, **kwargs)
         )
-
-    def visit_FencilWithTemporaries(
-        self, node: global_tmps.FencilWithTemporaries, **kwargs: Any
-    ) -> FencilDefinition:
-        fencil = self.visit(node.fencil, **kwargs)
-        return FencilDefinition(
-            id=fencil.id,
-            params=self.visit(node.params),
-            executions=fencil.executions,
-            grid_type=fencil.grid_type,
-            offset_definitions=fencil.offset_definitions,
-            function_definitions=fencil.function_definitions,
-            temporaries=self.visit(node.tmps, params=[p.id for p in node.params]),
-        )
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 39a5440e..487404b2 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
@@ -20,11 +20,7 @@ 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 Connectivity
-from gt4py.next.iterator import (
-    ir as itir,
-    transforms as itir_transforms,
-    type_inference as itir_typing,
-)
+from gt4py.next.iterator import ir as itir, 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 as tt
 
@@ -164,7 +160,7 @@ class ItirToSDFG(eve.NodeVisitor):
         self,
         param_types: list[ts.TypeSpec],
         offset_provider: dict[str, Connectivity | Dimension],
-        tmps: list[itir_transforms.global_tmps.Temporary],
+        tmps: list[itir.Temporary],
         use_field_canonical_representation: bool,
         column_axis: Optional[Dimension] = None,
     ):
diff --git a/src/gt4py/next/program_processors/runners/roundtrip.py b/src/gt4py/next/program_processors/runners/roundtrip.py
index 38714221..aba2dbf7 100644
--- a/src/gt4py/next/program_processors/runners/roundtrip.py
+++ b/src/gt4py/next/program_processors/runners/roundtrip.py
@@ -86,7 +86,7 @@ def ${id}(${','.join(params)}):
             + f"\n    {node.fencil.id}({args}, **kwargs)\n"
         )
 
-    def visit_Temporary(self, node: gtmps_transform.Temporary, **kwargs: Any) -> str:
+    def visit_Temporary(self, node: itir.Temporary, **kwargs: Any) -> str:
         assert (
             isinstance(node.domain, itir.FunCall)
             and isinstance(node.domain.fun, itir.SymRef)


Test patch:
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(
diff --git a/tests/next_tests/unit_tests/iterator_tests/test_pretty_parser.py b/tests/next_tests/unit_tests/iterator_tests/test_pretty_parser.py
index b3a9ba80..b02c610a 100644
--- a/tests/next_tests/unit_tests/iterator_tests/test_pretty_parser.py
+++ b/tests/next_tests/unit_tests/iterator_tests/test_pretty_parser.py
@@ -69,6 +69,13 @@ def test_lift():
     assert actual == expected
 
 
+def test_as_fieldop():
+    testee = "⇑x"
+    expected = ir.FunCall(fun=ir.SymRef(id="as_fieldop"), args=[ir.SymRef(id="x")])
+    actual = pparse(testee)
+    assert actual == expected
+
+
 def test_bool_arithmetic():
     testee = "¬(¬a ∨ b ∧ (c ∨ d))"
     expected = ir.FunCall(
@@ -183,6 +190,13 @@ def test_function_definition():
     assert actual == expected
 
 
+def test_temporary():
+    testee = "t = temporary(domain=domain, dtype=float64);"
+    expected = ir.Temporary(id="t", domain=ir.SymRef(id="domain"), dtype=ir.SymRef(id="float64"))
+    actual = pparse(testee)
+    assert actual == expected
+
+
 def test_stencil_closure():
     testee = "y ← (deref)(x) @ cartesian_domain();"
     expected = ir.StencilClosure(
@@ -195,6 +209,18 @@ def test_stencil_closure():
     assert actual == expected
 
 
+def test_set_at():
+    testee = "y @ cartesian_domain() ← x;"
+    expected = ir.SetAt(
+        expr=ir.SymRef(id="x"),
+        domain=ir.FunCall(fun=ir.SymRef(id="cartesian_domain"), args=[]),
+        target=ir.SymRef(id="y"),
+    )
+    actual = pparse(testee)
+    assert actual == expected
+
+
+# TODO(havogt): remove after refactoring to GTIR
 def test_fencil_definition():
     testee = "f(d, x, y) {\n  g = λ(x) → x;\n  y ← (deref)(x) @ cartesian_domain();\n}"
     expected = ir.FencilDefinition(
@@ -214,3 +240,30 @@ def test_fencil_definition():
     )
     actual = pparse(testee)
     assert actual == expected
+
+
+def test_program():
+    testee = "f(d, x, y) {\n  g = λ(x) → x;\n  tmp = temporary(domain=cartesian_domain(), dtype=float64);\n  y @ cartesian_domain() ← x;\n}"
+    expected = ir.Program(
+        id="f",
+        function_definitions=[
+            ir.FunctionDefinition(id="g", params=[ir.Sym(id="x")], expr=ir.SymRef(id="x"))
+        ],
+        params=[ir.Sym(id="d"), ir.Sym(id="x"), ir.Sym(id="y")],
+        declarations=[
+            ir.Temporary(
+                id="tmp",
+                domain=ir.FunCall(fun=ir.SymRef(id="cartesian_domain"), args=[]),
+                dtype=ir.SymRef(id="float64"),
+            ),
+        ],
+        body=[
+            ir.SetAt(
+                expr=ir.SymRef(id="x"),
+                domain=ir.FunCall(fun=ir.SymRef(id="cartesian_domain"), args=[]),
+                target=ir.SymRef(id="y"),
+            )
+        ],
+    )
+    actual = pparse(testee)
+    assert actual == expected
diff --git a/tests/next_tests/unit_tests/iterator_tests/test_pretty_printer.py b/tests/next_tests/unit_tests/iterator_tests/test_pretty_printer.py
index 844c905e..bc1372ce 100644
--- a/tests/next_tests/unit_tests/iterator_tests/test_pretty_printer.py
+++ b/tests/next_tests/unit_tests/iterator_tests/test_pretty_printer.py
@@ -159,6 +159,13 @@ def test_lift():
     assert actual == expected
 
 
+def test_as_fieldop():
+    testee = ir.FunCall(fun=ir.SymRef(id="as_fieldop"), args=[ir.SymRef(id="x")])
+    expected = "⇑x"
+    actual = pformat(testee)
+    assert actual == expected
+
+
 def test_bool_arithmetic():
     testee = ir.FunCall(
         fun=ir.SymRef(id="not_"),
@@ -287,6 +294,13 @@ def test_function_definition():
     assert actual == expected
 
 
+def test_temporary():
+    testee = ir.Temporary(id="t", domain=ir.SymRef(id="domain"), dtype="float64")
+    expected = "t = temporary(domain=domain, dtype=float64);"
+    actual = pformat(testee)
+    assert actual == expected
+
+
 def test_stencil_closure():
     testee = ir.StencilClosure(
         domain=ir.FunCall(fun=ir.SymRef(id="cartesian_domain"), args=[]),
@@ -299,6 +313,18 @@ def test_stencil_closure():
     assert actual == expected
 
 
+def test_set_at():
+    testee = ir.SetAt(
+        expr=ir.SymRef(id="x"),
+        domain=ir.FunCall(fun=ir.SymRef(id="cartesian_domain"), args=[]),
+        target=ir.SymRef(id="y"),
+    )
+    expected = "y @ cartesian_domain() ← x;"
+    actual = pformat(testee)
+    assert actual == expected
+
+
+# TODO(havogt): remove after refactoring.
 def test_fencil_definition():
     testee = ir.FencilDefinition(
         id="f",
@@ -318,3 +344,30 @@ def test_fencil_definition():
     actual = pformat(testee)
     expected = "f(d, x, y) {\n  g = λ(x) → x;\n  y ← (deref)(x) @ cartesian_domain();\n}"
     assert actual == expected
+
+
+def test_program():
+    testee = ir.Program(
+        id="f",
+        function_definitions=[
+            ir.FunctionDefinition(id="g", params=[ir.Sym(id="x")], expr=ir.SymRef(id="x"))
+        ],
+        params=[ir.Sym(id="d"), ir.Sym(id="x"), ir.Sym(id="y")],
+        declarations=[
+            ir.Temporary(
+                id="tmp",
+                domain=ir.FunCall(fun=ir.SymRef(id="cartesian_domain"), args=[]),
+                dtype="float64",
+            ),
+        ],
+        body=[
+            ir.SetAt(
+                expr=ir.SymRef(id="x"),
+                domain=ir.FunCall(fun=ir.SymRef(id="cartesian_domain"), args=[]),
+                target=ir.SymRef(id="y"),
+            )
+        ],
+    )
+    actual = pformat(testee)
+    expected = "f(d, x, y) {\n  g = λ(x) → x;\n  tmp = temporary(domain=cartesian_domain(), dtype=float64);\n  y @ cartesian_domain() ← x;\n}"
+    assert actual == expected
diff --git a/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_global_tmps.py b/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_global_tmps.py
index 0521b041..442f71bc 100644
--- a/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_global_tmps.py
+++ b/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_global_tmps.py
@@ -21,7 +21,6 @@ from gt4py.next.iterator.transforms.global_tmps import (
     AUTO_DOMAIN,
     FencilWithTemporaries,
     SimpleTemporaryExtractionHeuristics,
-    Temporary,
     collect_tmps_info,
     split_closures,
     update_domains,
@@ -87,7 +86,7 @@ def test_split_closures():
         ],
     )
     actual = split_closures(testee, offset_provider={})
-    assert actual.tmps == [Temporary(id="_tmp_1"), Temporary(id="_tmp_2")]
+    assert actual.tmps == [ir.Temporary(id="_tmp_1"), ir.Temporary(id="_tmp_2")]
     assert actual.fencil == expected
 
 
@@ -141,7 +140,7 @@ def test_split_closures_simple_heuristics():
     actual = split_closures(
         testee, extraction_heuristics=SimpleTemporaryExtractionHeuristics, offset_provider={}
     )
-    assert actual.tmps == [Temporary(id="_tmp_1")]
+    assert actual.tmps == [ir.Temporary(id="_tmp_1")]
     assert actual.fencil == expected
 
 
@@ -211,7 +210,7 @@ def test_split_closures_lifted_scan():
     )
 
     actual = split_closures(testee, offset_provider={})
-    assert actual.tmps == [Temporary(id="_tmp_1")]
+    assert actual.tmps == [ir.Temporary(id="_tmp_1")]
     assert actual.fencil == expected
 
 
@@ -255,7 +254,7 @@ def test_update_cartesian_domains():
             ],
         ),
         params=[im.sym("i"), im.sym("j"), im.sym("k"), im.sym("inp"), im.sym("out")],
-        tmps=[Temporary(id="_gtmp_0"), Temporary(id="_gtmp_1")],
+        tmps=[ir.Temporary(id="_gtmp_0"), ir.Temporary(id="_gtmp_1")],
     )
     expected = copy.deepcopy(testee)
     assert expected.fencil.params.pop() == im.sym("_gtmp_auto_domain")
@@ -413,14 +412,14 @@ def test_collect_tmps_info():
             ],
         ),
         params=[ir.Sym(id="i"), ir.Sym(id="j"), ir.Sym(id="k"), ir.Sym(id="inp"), ir.Sym(id="out")],
-        tmps=[Temporary(id="_gtmp_0"), Temporary(id="_gtmp_1")],
+        tmps=[ir.Temporary(id="_gtmp_0"), ir.Temporary(id="_gtmp_1")],
     )
     expected = FencilWithTemporaries(
         fencil=testee.fencil,
         params=testee.params,
         tmps=[
-            Temporary(id="_gtmp_0", domain=tmp_domain, dtype="float64"),
-            Temporary(id="_gtmp_1", domain=tmp_domain, dtype="float64"),
+            ir.Temporary(id="_gtmp_0", domain=tmp_domain, dtype="float64"),
+            ir.Temporary(id="_gtmp_1", domain=tmp_domain, dtype="float64"),
         ],
     )
     actual = collect_tmps_info(testee, offset_provider={})
diff --git a/tests/next_tests/unit_tests/program_processor_tests/codegens_tests/gtfn_tests/test_itir_to_gtfn_ir.py b/tests/next_tests/unit_tests/program_processor_tests/codegens_tests/gtfn_tests/test_itir_to_gtfn_ir.py
index 43a0b45c..bfb51175 100644
--- a/tests/next_tests/unit_tests/program_processor_tests/codegens_tests/gtfn_tests/test_itir_to_gtfn_ir.py
+++ b/tests/next_tests/unit_tests/program_processor_tests/codegens_tests/gtfn_tests/test_itir_to_gtfn_ir.py
@@ -14,6 +14,7 @@
 
 import gt4py.next as gtx
 from gt4py.next.iterator import ir as itir
+from gt4py.next.iterator.ir_utils import ir_makers as im
 from gt4py.next.program_processors.codegens.gtfn import gtfn_ir, itir_to_gtfn_ir as it2gtfn
 
 
@@ -41,3 +42,23 @@ def test_unapplied_funcall_to_function_object():
     ).visit(testee)
 
     assert expected == actual
+
+
+def test_get_domains():
+    domain = im.call("cartesian_domain")(im.call("named_range")(itir.AxisLiteral(value="D"), 1, 2))
+    testee = itir.Program(
+        id="foo",
+        function_definitions=[],
+        params=[itir.Sym(id="bar")],
+        declarations=[],
+        body=[
+            itir.SetAt(
+                expr=im.call(im.call("as_fieldop")("deref"))(),
+                domain=domain,
+                target=itir.SymRef(id="bar"),
+            )
+        ],
+    )
+
+    result = list(it2gtfn._get_domains(testee.body))
+    assert result == [domain]