Repository: GridTools/gt4py
Pull request #1773 (2024-12-05 20:59:48)

refactor[next]: remove all FencilDefinitions from tests


Patch:
diff --git a/src/gt4py/next/iterator/ir.py b/src/gt4py/next/iterator/ir.py
index 7098e9fa..6efee293 100644
--- a/src/gt4py/next/iterator/ir.py
+++ b/src/gt4py/next/iterator/ir.py
@@ -189,17 +189,11 @@ BUILTINS = {
     "scan",
     "if_",
     "index",  # `index(dim)` creates a dim-field that has the current index at each point
+    "as_fieldop",  # `as_fieldop(stencil, domain)` creates field_operator from stencil (domain is optional, but for now required for embedded execution)
     *ARITHMETIC_BUILTINS,
     *TYPEBUILTINS,
 }
 
-# only used in `Program`` not `FencilDefinition`
-# TODO(havogt): restructure after refactoring to GTIR
-GTIR_BUILTINS = {
-    *BUILTINS,
-    "as_fieldop",  # `as_fieldop(stencil, domain)` creates field_operator from stencil (domain is optional, but for now required for embedded execution)
-}
-
 
 class FencilDefinition(Node, ValidatedSymbolTableTrait):
     id: Coerced[SymbolName]
@@ -243,7 +237,7 @@ class Program(Node, ValidatedSymbolTableTrait):
     implicit_domain: bool = False
 
     _NODE_SYMBOLS_: ClassVar[List[Sym]] = [
-        Sym(id=name) for name in sorted(GTIR_BUILTINS)
+        Sym(id=name) for name in sorted(BUILTINS)
     ]  # sorted for serialization stability
 
 
diff --git a/src/gt4py/next/iterator/transforms/symbol_ref_utils.py b/src/gt4py/next/iterator/transforms/symbol_ref_utils.py
index 05163a36..1765259a 100644
--- a/src/gt4py/next/iterator/transforms/symbol_ref_utils.py
+++ b/src/gt4py/next/iterator/transforms/symbol_ref_utils.py
@@ -140,6 +140,4 @@ def collect_symbol_refs(
 
 
 def get_user_defined_symbols(symtable: dict[eve.SymbolName, itir.Sym]) -> set[str]:
-    return {str(sym) for sym in symtable.keys()} - {
-        str(n.id) for n in itir.FencilDefinition._NODE_SYMBOLS_
-    }
+    return {str(sym) for sym in symtable.keys()} - {str(n.id) for n in itir.Program._NODE_SYMBOLS_}


Test patch:
diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_decorator.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_decorator.py
index 47419c27..45bf7428 100644
--- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_decorator.py
+++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_decorator.py
@@ -30,10 +30,8 @@ def test_program_itir_regression(cartesian_case):
     def testee(a: cases.IField, out: cases.IField):
         testee_op(a, out=out)
 
-    assert isinstance(testee.itir, (itir.Program, itir.FencilDefinition))
-    assert isinstance(
-        testee.with_backend(cartesian_case.backend).itir, (itir.Program, itir.FencilDefinition)
-    )
+    assert isinstance(testee.itir, itir.Program)
+    assert isinstance(testee.with_backend(cartesian_case.backend).itir, itir.Program)
 
 
 def test_frozen(cartesian_case):
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 da4bea88..bf47f997 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
@@ -7,8 +7,8 @@
 # SPDX-License-Identifier: BSD-3-Clause
 
 from gt4py.next.iterator import ir
-from gt4py.next.iterator.pretty_parser import pparse
 from gt4py.next.iterator.ir_utils import ir_makers as im
+from gt4py.next.iterator.pretty_parser import pparse
 from gt4py.next.type_system import type_specifications as ts
 
 
@@ -208,18 +208,6 @@ def test_temporary():
     assert actual == expected
 
 
-def test_stencil_closure():
-    testee = "y ← (deref)(x) @ cartesian_domain();"
-    expected = ir.StencilClosure(
-        domain=ir.FunCall(fun=ir.SymRef(id="cartesian_domain"), args=[]),
-        stencil=ir.SymRef(id="deref"),
-        output=ir.SymRef(id="y"),
-        inputs=[ir.SymRef(id="x")],
-    )
-    actual = pparse(testee)
-    assert actual == expected
-
-
 def test_set_at():
     testee = "y @ cartesian_domain() ← x;"
     expected = ir.SetAt(
@@ -262,28 +250,6 @@ def test_if_stmt():
     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(
-        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")],
-        closures=[
-            ir.StencilClosure(
-                domain=ir.FunCall(fun=ir.SymRef(id="cartesian_domain"), args=[]),
-                stencil=ir.SymRef(id="deref"),
-                output=ir.SymRef(id="y"),
-                inputs=[ir.SymRef(id="x")],
-            )
-        ],
-    )
-    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(
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 69a45cf1..11f50dbf 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
@@ -7,8 +7,8 @@
 # SPDX-License-Identifier: BSD-3-Clause
 
 from gt4py.next.iterator import ir
-from gt4py.next.iterator.pretty_printer import PrettyPrinter, pformat
 from gt4py.next.iterator.ir_utils import ir_makers as im
+from gt4py.next.iterator.pretty_printer import PrettyPrinter, pformat
 from gt4py.next.type_system import type_specifications as ts
 
 
@@ -313,18 +313,6 @@ def test_temporary():
     assert actual == expected
 
 
-def test_stencil_closure():
-    testee = ir.StencilClosure(
-        domain=ir.FunCall(fun=ir.SymRef(id="cartesian_domain"), args=[]),
-        stencil=ir.SymRef(id="deref"),
-        output=ir.SymRef(id="y"),
-        inputs=[ir.SymRef(id="x")],
-    )
-    expected = "y ← (deref)(x) @ cartesian_domain();"
-    actual = pformat(testee)
-    assert actual == expected
-
-
 def test_set_at():
     testee = ir.SetAt(
         expr=ir.SymRef(id="x"),
@@ -336,28 +324,6 @@ def test_set_at():
     assert actual == expected
 
 
-# TODO(havogt): remove after refactoring.
-def test_fencil_definition():
-    testee = ir.FencilDefinition(
-        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")],
-        closures=[
-            ir.StencilClosure(
-                domain=ir.FunCall(fun=ir.SymRef(id="cartesian_domain"), args=[]),
-                stencil=ir.SymRef(id="deref"),
-                output=ir.SymRef(id="y"),
-                inputs=[ir.SymRef(id="x")],
-            )
-        ],
-    )
-    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",
diff --git a/tests/next_tests/unit_tests/iterator_tests/test_type_inference.py b/tests/next_tests/unit_tests/iterator_tests/test_type_inference.py
index 65a5b588..7eb4e86a 100644
--- a/tests/next_tests/unit_tests/iterator_tests/test_type_inference.py
+++ b/tests/next_tests/unit_tests/iterator_tests/test_type_inference.py
@@ -23,13 +23,12 @@ from gt4py.next.iterator.type_system import (
 )
 from gt4py.next.type_system import type_specifications as ts
 
-from next_tests.integration_tests.feature_tests.ffront_tests.ffront_test_utils import simple_mesh
-
 from next_tests.integration_tests.cases import (
     C2E,
     E2V,
     V2E,
     E2VDim,
+    Edge,
     IDim,
     Ioff,
     JDim,
@@ -37,11 +36,12 @@ from next_tests.integration_tests.cases import (
     Koff,
     V2EDim,
     Vertex,
-    Edge,
-    mesh_descriptor,
     exec_alloc_descriptor,
+    mesh_descriptor,
     unstructured_case,
 )
+from next_tests.integration_tests.feature_tests.ffront_tests.ffront_test_utils import simple_mesh
+
 
 bool_type = ts.ScalarType(kind=ts.ScalarKind.BOOL)
 int_type = ts.ScalarType(kind=ts.ScalarKind.INT32)
@@ -275,48 +275,35 @@ def test_cast_first_arg_inference():
     assert result.type == float64_type
 
 
-# TODO(tehrengruber): Rewrite tests to use itir.Program
 def test_cartesian_fencil_definition():
     cartesian_domain = im.call("cartesian_domain")(
         im.call("named_range")(itir.AxisLiteral(value="IDim"), 0, 1)
     )
 
-    testee = itir.FencilDefinition(
+    testee = itir.Program(
         id="f",
         function_definitions=[],
         params=[im.sym("inp", float_i_field), im.sym("out", float_i_field)],
-        closures=[
-            itir.StencilClosure(
+        declarations=[],
+        body=[
+            itir.SetAt(
+                expr=im.call(im.call("as_fieldop")(im.ref("deref"), cartesian_domain))(
+                    im.ref("inp")
+                ),
                 domain=cartesian_domain,
-                stencil=im.ref("deref"),
-                output=im.ref("out"),
-                inputs=[im.ref("inp")],
+                target=im.ref("out"),
             ),
         ],
     )
 
     result = itir_type_inference.infer(testee, offset_provider_type={"Ioff": IDim})
 
-    closure_type = it_ts.StencilClosureType(
-        domain=it_ts.DomainType(dims=[IDim]),
-        stencil=ts.FunctionType(
-            pos_only_args=[
-                it_ts.IteratorType(
-                    position_dims=[IDim], defined_dims=[IDim], element_type=float64_type
-                )
-            ],
-            pos_or_kw_args={},
-            kw_only_args={},
-            returns=float64_type,
-        ),
-        output=float_i_field,
-        inputs=[float_i_field],
-    )
-    fencil_type = it_ts.FencilType(
-        params={"inp": float_i_field, "out": float_i_field}, closures=[closure_type]
-    )
-    assert result.type == fencil_type
-    assert result.closures[0].type == closure_type
+    program_type = it_ts.ProgramType(params={"inp": float_i_field, "out": float_i_field})
+    assert result.type == program_type
+    domain_type = it_ts.DomainType(dims=[IDim])
+    assert result.body[0].domain.type == domain_type
+    assert result.body[0].expr.type == float_i_field
+    assert result.body[0].target.type == float_i_field
 
 
 def test_unstructured_fencil_definition():
@@ -326,44 +313,34 @@ def test_unstructured_fencil_definition():
         im.call("named_range")(itir.AxisLiteral(value="KDim"), 0, 1),
     )
 
-    testee = itir.FencilDefinition(
+    testee = itir.Program(
         id="f",
         function_definitions=[],
         params=[im.sym("inp", float_edge_k_field), im.sym("out", float_vertex_k_field)],
-        closures=[
-            itir.StencilClosure(
+        declarations=[],
+        body=[
+            itir.SetAt(
+                expr=im.call(
+                    im.call("as_fieldop")(
+                        im.lambda_("it")(im.deref(im.shift("V2E", 0)("it"))), unstructured_domain
+                    )
+                )(im.ref("inp")),
                 domain=unstructured_domain,
-                stencil=im.lambda_("it")(im.deref(im.shift("V2E", 0)("it"))),
-                output=im.ref("out"),
-                inputs=[im.ref("inp")],
+                target=im.ref("out"),
             ),
         ],
     )
 
     result = itir_type_inference.infer(testee, offset_provider_type=mesh.offset_provider_type)
 
-    closure_type = it_ts.StencilClosureType(
-        domain=it_ts.DomainType(dims=[Vertex, KDim]),
-        stencil=ts.FunctionType(
-            pos_only_args=[
-                it_ts.IteratorType(
-                    position_dims=[Vertex, KDim],
-                    defined_dims=[Edge, KDim],
-                    element_type=float64_type,
-                )
-            ],
-            pos_or_kw_args={},
-            kw_only_args={},
-            returns=float64_type,
-        ),
-        output=float_vertex_k_field,
-        inputs=[float_edge_k_field],
-    )
-    fencil_type = it_ts.FencilType(
-        params={"inp": float_edge_k_field, "out": float_vertex_k_field}, closures=[closure_type]
+    program_type = it_ts.ProgramType(
+        params={"inp": float_edge_k_field, "out": float_vertex_k_field}
     )
-    assert result.type == fencil_type
-    assert result.closures[0].type == closure_type
+    assert result.type == program_type
+    domain_type = it_ts.DomainType(dims=[Vertex, KDim])
+    assert result.body[0].domain.type == domain_type
+    assert result.body[0].expr.type == float_vertex_k_field
+    assert result.body[0].target.type == float_vertex_k_field
 
 
 def test_function_definition():
@@ -371,45 +348,29 @@ def test_function_definition():
         im.call("named_range")(itir.AxisLiteral(value="IDim"), 0, 1)
     )
 
-    testee = itir.FencilDefinition(
+    testee = itir.Program(
         id="f",
         function_definitions=[
             itir.FunctionDefinition(id="foo", params=[im.sym("it")], expr=im.deref("it")),
             itir.FunctionDefinition(id="bar", params=[im.sym("it")], expr=im.call("foo")("it")),
         ],
         params=[im.sym("inp", float_i_field), im.sym("out", float_i_field)],
-        closures=[
-            itir.StencilClosure(
+        declarations=[],
+        body=[
+            itir.SetAt(
                 domain=cartesian_domain,
-                stencil=im.ref("bar"),
-                output=im.ref("out"),
-                inputs=[im.ref("inp")],
+                expr=im.call(im.call("as_fieldop")(im.ref("bar"), cartesian_domain))(im.ref("inp")),
+                target=im.ref("out"),
             ),
         ],
     )
 
     result = itir_type_inference.infer(testee, offset_provider_type={"Ioff": IDim})
 
-    closure_type = it_ts.StencilClosureType(
-        domain=it_ts.DomainType(dims=[IDim]),
-        stencil=ts.FunctionType(
-            pos_only_args=[
-                it_ts.IteratorType(
-                    position_dims=[IDim], defined_dims=[IDim], element_type=float64_type
-                )
-            ],
-            pos_or_kw_args={},
-            kw_only_args={},
-            returns=float64_type,
-        ),
-        output=float_i_field,
-        inputs=[float_i_field],
-    )
-    fencil_type = it_ts.FencilType(
-        params={"inp": float_i_field, "out": float_i_field}, closures=[closure_type]
-    )
-    assert result.type == fencil_type
-    assert result.closures[0].type == closure_type
+    program_type = it_ts.ProgramType(params={"inp": float_i_field, "out": float_i_field})
+    assert result.type == program_type
+    assert result.body[0].expr.type == float_i_field
+    assert result.body[0].target.type == float_i_field
 
 
 def test_fencil_with_nb_field_input():
@@ -419,24 +380,30 @@ def test_fencil_with_nb_field_input():
         im.call("named_range")(itir.AxisLiteral(value="KDim"), 0, 1),
     )
 
-    testee = itir.FencilDefinition(
+    testee = itir.Program(
         id="f",
         function_definitions=[],
         params=[im.sym("inp", float_vertex_v2e_field), im.sym("out", float_vertex_k_field)],
-        closures=[
-            itir.StencilClosure(
+        declarations=[],
+        body=[
+            itir.SetAt(
                 domain=unstructured_domain,
-                stencil=im.lambda_("it")(im.call(im.call("reduce")("plus", 0.0))(im.deref("it"))),
-                output=im.ref("out"),
-                inputs=[im.ref("inp")],
+                expr=im.call(
+                    im.call("as_fieldop")(
+                        im.lambda_("it")(im.call(im.call("reduce")("plus", 0.0))(im.deref("it"))),
+                        unstructured_domain,
+                    )
+                )(im.ref("inp")),
+                target=im.ref("out"),
             ),
         ],
     )
 
     result = itir_type_inference.infer(testee, offset_provider_type=mesh.offset_provider_type)
 
-    assert result.closures[0].stencil.expr.args[0].type == float64_list_type
-    assert result.closures[0].stencil.type.returns == float64_type
+    stencil = result.body[0].expr.fun.args[0]
+    assert stencil.expr.args[0].type == float64_list_type
+    assert stencil.type.returns == float64_type
 
 
 def test_program_tuple_setat_short_target():
diff --git a/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_symbol_ref_utils.py b/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_symbol_ref_utils.py
index 0c118ff6..c162860c 100644
--- a/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_symbol_ref_utils.py
+++ b/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_symbol_ref_utils.py
@@ -6,28 +6,23 @@
 # Please, refer to the LICENSE file in the root directory.
 # SPDX-License-Identifier: BSD-3-Clause
 
-from dataclasses import dataclass
-from typing import Optional
 
-from gt4py import eve
 from gt4py.next.iterator import ir as itir
-from gt4py.next.iterator.transforms.symbol_ref_utils import (
-    collect_symbol_refs,
-    get_user_defined_symbols,
-)
+from gt4py.next.iterator.transforms.symbol_ref_utils import get_user_defined_symbols
 
 
 def test_get_user_defined_symbols():
-    ir = itir.FencilDefinition(
+    domain = itir.FunCall(fun=itir.SymRef(id="cartesian_domain"), args=[])
+    ir = itir.Program(
         id="foo",
         function_definitions=[],
         params=[itir.Sym(id="target_symbol")],
-        closures=[
-            itir.StencilClosure(
-                domain=itir.FunCall(fun=itir.SymRef(id="cartesian_domain"), args=[]),
-                stencil=itir.SymRef(id="deref"),
-                output=itir.SymRef(id="target_symbol"),
-                inputs=[],
+        declarations=[],
+        body=[
+            itir.SetAt(
+                expr=itir.Lambda(params=[itir.Sym(id="foo")], expr=itir.SymRef(id="foo")),
+                domain=domain,
+                target=itir.SymRef(id="target_symbol"),
             )
         ],
     )
diff --git a/tests/next_tests/unit_tests/program_processor_tests/codegens_tests/gtfn_tests/test_gtfn_module.py b/tests/next_tests/unit_tests/program_processor_tests/codegens_tests/gtfn_tests/test_gtfn_module.py
index e64bd8a5..0586d487 100644
--- a/tests/next_tests/unit_tests/program_processor_tests/codegens_tests/gtfn_tests/test_gtfn_module.py
+++ b/tests/next_tests/unit_tests/program_processor_tests/codegens_tests/gtfn_tests/test_gtfn_module.py
@@ -6,11 +6,11 @@
 # Please, refer to the LICENSE file in the root directory.
 # SPDX-License-Identifier: BSD-3-Clause
 
-import numpy as np
-import pytest
 import copy
-import diskcache
 
+import diskcache
+import numpy as np
+import pytest
 
 import gt4py.next as gtx
 from gt4py.next.iterator import ir as itir
@@ -19,18 +19,17 @@ from gt4py.next.otf import arguments, languages, stages
 from gt4py.next.program_processors.codegens.gtfn import gtfn_module
 from gt4py.next.program_processors.runners import gtfn
 from gt4py.next.type_system import type_translation
-from next_tests.integration_tests import cases
-from next_tests.integration_tests.feature_tests.ffront_tests.ffront_test_utils import KDim
 
+from next_tests.integration_tests import cases
 from next_tests.integration_tests.cases import cartesian_case
-
 from next_tests.integration_tests.feature_tests.ffront_tests.ffront_test_utils import (
+    KDim,
     exec_alloc_descriptor,
 )
 
 
 @pytest.fixture
-def fencil_example():
+def program_example():
     IDim = gtx.Dimension("I")
     params = [gtx.as_field([IDim], np.empty((1,), dtype=np.float32)), np.float32(3.14)]
     param_types = [type_translation.from_value(param) for param in params]
@@ -48,7 +47,7 @@ def fencil_example():
             )
         ],
     )
-    fencil = itir.FencilDefinition(
+    program = itir.Program(
         id="example",
         params=[im.sym(name, type_) for name, type_ in zip(("buf", "sc"), param_types)],
         function_definitions=[
@@ -58,20 +57,22 @@ def fencil_example():
                 expr=im.literal("1", "float32"),
             )
         ],
-        closures=[
-            itir.StencilClosure(
+        declarations=[],
+        body=[
+            itir.SetAt(
+                expr=im.call(im.call("as_fieldop")(itir.SymRef(id="stencil"), domain))(
+                    itir.SymRef(id="buf"), itir.SymRef(id="sc")
+                ),
                 domain=domain,
-                stencil=itir.SymRef(id="stencil"),
-                output=itir.SymRef(id="buf"),
-                inputs=[itir.SymRef(id="buf"), itir.SymRef(id="sc")],
+                target=itir.SymRef(id="buf"),
             )
         ],
     )
-    return fencil, params
+    return program, params
 
 
-def test_codegen(fencil_example):
-    fencil, parameters = fencil_example
+def test_codegen(program_example):
+    fencil, parameters = program_example
     module = gtfn_module.translate_program_cpu(
         stages.CompilableProgram(
             data=fencil,
@@ -85,8 +86,8 @@ def test_codegen(fencil_example):
     assert module.language is languages.CPP
 
 
-def test_hash_and_diskcache(fencil_example, tmp_path):
-    fencil, parameters = fencil_example
+def test_hash_and_diskcache(program_example, tmp_path):
+    fencil, parameters = program_example
     compilable_program = stages.CompilableProgram(
         data=fencil,
         args=arguments.CompileTimeArgs.from_concrete_no_size(
@@ -129,8 +130,8 @@ def test_hash_and_diskcache(fencil_example, tmp_path):
     ) != gtfn.fingerprint_compilable_program(altered_program_column_axis)
 
 
-def test_gtfn_file_cache(fencil_example):
-    fencil, parameters = fencil_example
+def test_gtfn_file_cache(program_example):
+    fencil, parameters = program_example
     compilable_program = stages.CompilableProgram(
         data=fencil,
         args=arguments.CompileTimeArgs.from_concrete_no_size(