temporaries

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
@@ -217,15 +217,21 @@ class Stmt(Node): ...
 
 
 class Assign(Stmt):
-    target: SymRef
-    expr: Expr  # TODO Program expression
+    target: Expr  # `make_tuple` or SymRef
+    expr: Expr  # only `apply_stencil`
+
+
+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[Sym]
+    declarations: List[Temporary]
     body: List[Stmt]
 
     _NODE_SYMBOLS_: ClassVar[List[Sym]] = [Sym(id=name) for name in BUILTINS]

src/gt4py/next/iterator/transforms/fencil_to_program.py

@@ -34,3 +34,12 @@ def visit_FencilDefinition(self, node: itir.FencilDefinition) -> itir.Program:
             declarations=[],
             body=self.visit(node.closures),
         )
+
+    def visit_FencilWithTemporaries(self, node) -> 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),
+        )

src/gt4py/next/iterator/transforms/global_tmps.py

@@ -17,9 +17,8 @@
 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,26 +53,18 @@
 # 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):
     """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]:
+def pformat_Temporary(printer: PrettyPrinter, node: ir.Temporary, *, prec: int) -> list[str]:
     start, end = [node.id + " = temporary("], [");"]
     args = []
     if node.domain is not None:
@@ -367,7 +358,7 @@ def always_extract_heuristics(_):
             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 +629,8 @@ def convert_type(dtype):
         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
         ],
     )
 

src/gt4py/next/iterator/transforms/pass_manager.py

@@ -204,7 +204,6 @@ def apply_common_transforms(
         ir, opcount_preserving=True, force_inline_lambda_args=force_inline_lambda_args
     )
 
-    assert isinstance(ir, itir.FencilDefinition)
-    prog = FencilToProgram.apply(ir)
+    prog = FencilToProgram.apply(ir)  # type: ignore[arg-type]  # TODO: remove after refactoring
 
     return prog

src/gt4py/next/program_processors/codegens/gtfn/itir_to_gtfn_ir.py

@@ -575,11 +575,11 @@ def visit_Program(self, node: itir.Program, **kwargs: Any) -> FencilDefinition:
             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):
@@ -601,13 +601,6 @@ def dtype_to_cpp(x: int | tuple | str) -> str:
     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]),
-        )
+        raise AssertionError(
+            "Internal error: Fencils are no longer supported."
+        )  # TODO remove after refactoring is complete

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

@@ -20,11 +20,7 @@
 import gt4py.eve as eve
 from gt4py.next import Dimension, DimensionKind, type_inference as next_typing
 from gt4py.next.common import NeighborTable
-from gt4py.next.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
 
@@ -164,7 +160,7 @@ def __init__(
         self,
         param_types: list[ts.TypeSpec],
         offset_provider: dict[str, NeighborTable],
-        tmps: list[itir_transforms.global_tmps.Temporary],
+        tmps: list[itir.Temporary],
         use_field_canonical_representation: bool,
         column_axis: Optional[Dimension] = None,
     ):

src/gt4py/next/program_processors/runners/roundtrip.py

@@ -86,7 +86,7 @@ def visit_FencilWithTemporaries(
             + 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)