Repository: GridTools/gt4py
Pull request #1273 (2023-06-19 11:33:10)

fix[next]: Fix CSE extraction for if statements
The criterion for extracting common subexpressions out of `if_` calls was wrong. The correct criterion as implemented in this PR is that all subexpressions that either occur in the condition or both branches are eligible for extraction.


Patch:
diff --git a/src/gt4py/next/iterator/transforms/cse.py b/src/gt4py/next/iterator/transforms/cse.py
index e3338676..d2e8d594 100644
--- a/src/gt4py/next/iterator/transforms/cse.py
+++ b/src/gt4py/next/iterator/transforms/cse.py
@@ -114,8 +114,8 @@ class CollectSubexpressions(VisitorWithSymbolTableTrait, NodeVisitor):
         used_symbol_ids: set[int] = set()
 
         # Special handling of `if_(condition, true_branch, false_branch)` like expressions that
-        # avoids extracting subexpressions unless they are used in at least two of the three
-        # arguments.
+        # avoids extracting subexpressions unless they are used in either the condition or both
+        # branches.
         if isinstance(node, ir.FunCall) and node.fun == ir.SymRef(id="if_"):
             assert len(node.args) == 3
             # collect subexpressions for all arguments to the `if_`
@@ -123,15 +123,11 @@ class CollectSubexpressions(VisitorWithSymbolTableTrait, NodeVisitor):
             for arg, state in zip(node.args, arg_states):
                 self.visit(arg, state=state, **kwargs)
 
-            # for each subexpression find in how many of the three arguments they occur
-            subexpr_count: dict[ir.Node, int] = {}
-            for arg_state in arg_states:
-                for subexpr in arg_state.subexprs.keys():
-                    subexpr_count.setdefault(subexpr, 0)
-                    subexpr_count[subexpr] += 1
-
             # remove all subexpressions that are not eligible for collection
-            eligible_subexprs = {subexpr for subexpr, count in subexpr_count.items() if count >= 2}
+            #  (either they occur in the condition or in both branches)
+            eligible_subexprs = arg_states[0].subexprs.keys() | (
+                arg_states[1].subexprs.keys() & arg_states[2].subexprs.keys()
+            )
             for arg_state in arg_states:
                 arg_state.remove_subexprs(arg_state.subexprs.keys() - eligible_subexprs)
 


Test patch:
diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_gt4py_builtins.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_gt4py_builtins.py
index 92bd5b99..8f5a5f69 100644
--- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_gt4py_builtins.py
+++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_gt4py_builtins.py
@@ -32,7 +32,6 @@ from next_tests.integration_tests.feature_tests.cases import (
     V2EDim,
     Vertex,
     cartesian_case,
-    fieldview_backend,
     unstructured_case,
 )
 from next_tests.integration_tests.feature_tests.ffront_tests.ffront_test_utils import (
diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_math_builtin_execution.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_math_builtin_execution.py
index 94b80103..81ca50d7 100644
--- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_math_builtin_execution.py
+++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_math_builtin_execution.py
@@ -18,13 +18,18 @@ from typing import Callable
 import numpy as np
 import pytest
 
-from gt4py.next import Field, float64, np_as_located_field
+from gt4py.next import np_as_located_field
 from gt4py.next.ffront import dialect_ast_enums, fbuiltins, field_operator_ast as foast
 from gt4py.next.ffront.decorator import FieldOperator
 from gt4py.next.ffront.foast_passes.type_deduction import FieldOperatorTypeDeduction
 from gt4py.next.type_system import type_translation
 
+from next_tests.integration_tests.feature_tests import cases
+from next_tests.integration_tests.feature_tests.cases import IDim, cartesian_case, unstructured_case
 from next_tests.integration_tests.feature_tests.ffront_tests.ffront_test_utils import *
+from next_tests.integration_tests.feature_tests.ffront_tests.ffront_test_utils import (
+    fieldview_backend,
+)
 from next_tests.integration_tests.feature_tests.math_builtin_test_data import math_builtin_test_data
 
 
@@ -41,16 +46,16 @@ def make_builtin_field_operator(builtin_name: str):
     # construct annotation dictionary containing the input argument and return
     #  types
     if builtin_name in fbuiltins.UNARY_MATH_NUMBER_BUILTIN_NAMES:
-        annotations = {"inp": Field[[IDim], float64], "return": Field[[IDim], float64]}
+        annotations = {"inp": cases.IFloatField, "return": cases.IFloatField}
     elif builtin_name in fbuiltins.UNARY_MATH_FP_BUILTIN_NAMES:
-        annotations = {"inp": Field[[IDim], float64], "return": Field[[IDim], float64]}
+        annotations = {"inp": cases.IFloatField, "return": cases.IFloatField}
     elif builtin_name in fbuiltins.UNARY_MATH_FP_PREDICATE_BUILTIN_NAMES:
-        annotations = {"inp": Field[[IDim], float64], "return": Field[[IDim], bool]}
+        annotations = {"inp": cases.IFloatField, "return": cases.IBoolField}
     elif builtin_name in fbuiltins.BINARY_MATH_NUMBER_BUILTIN_NAMES:
         annotations = {
-            "inp1": Field[[IDim], float64],
-            "inp2": Field[[IDim], float64],
-            "return": Field[[IDim], float64],
+            "inp1": cases.IFloatField,
+            "inp2": cases.IFloatField,
+            "return": cases.IFloatField,
         }
     else:
         raise AssertionError(f"Unknown builtin `{builtin_name}`")
@@ -111,7 +116,7 @@ def make_builtin_field_operator(builtin_name: str):
 
 
 @pytest.mark.parametrize("builtin_name, inputs", math_builtin_test_data())
-def test_math_function_builtins_execution(fieldview_backend, builtin_name: str, inputs):
+def test_math_function_builtins_execution(cartesian_case, builtin_name: str, inputs):
     if builtin_name == "gamma":
         # numpy has no gamma function
         ref_impl: Callable = np.vectorize(math.gamma)
@@ -122,7 +127,9 @@ def test_math_function_builtins_execution(fieldview_backend, builtin_name: str,
     expected = ref_impl(*inputs)
     out = np_as_located_field(IDim)(np.zeros_like(expected))
 
-    builtin_field_op = make_builtin_field_operator(builtin_name).with_backend(fieldview_backend)
+    builtin_field_op = make_builtin_field_operator(builtin_name).with_backend(
+        cartesian_case.backend
+    )
 
     builtin_field_op(*inps, out=out, offset_provider={})
 
diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_program.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_program.py
index dddc5ff2..bc6642c5 100644
--- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_program.py
+++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_program.py
@@ -19,15 +19,19 @@ import re
 import numpy as np
 import pytest
 
-from gt4py.next import Field, field_operator, np_as_located_field, program
+import gt4py.next as gtx
 from gt4py.next.common import GTTypeError
 from gt4py.next.ffront.past_passes.type_deduction import ProgramTypeError
-from gt4py.next.program_processors.runners import gtfn_cpu, roundtrip
 
-from next_tests.past_common_fixtures import (
+from next_tests.integration_tests.feature_tests import cases
+from next_tests.integration_tests.feature_tests.cases import (
     IDim,
     Ioff,
     JDim,
+    cartesian_case,
+    fieldview_backend,
+)
+from next_tests.past_common_fixtures import (
     copy_program_def,
     copy_restrict_program_def,
     double_copy_program_def,
@@ -36,205 +40,190 @@ from next_tests.past_common_fixtures import (
 )
 
 
-@pytest.fixture(params=[roundtrip.executor, gtfn_cpu.run_gtfn])
-def fieldview_backend(request):
-    yield request.param
-
-
-def test_identity_fo_execution(fieldview_backend, identity_def):
-    size = 10
-    in_field = np_as_located_field(IDim)(np.ones((size)))
-    out_field = np_as_located_field(IDim)(np.zeros((size)))
-    identity = field_operator(identity_def, backend=fieldview_backend)
-
-    identity(in_field, out=out_field, offset_provider={})
-
-    assert np.allclose(in_field, out_field)
+def test_identity_fo_execution(cartesian_case, identity_def):
+    identity = gtx.field_operator(identity_def, backend=cartesian_case.backend)
 
+    in_field = cases.allocate(cartesian_case, identity, "in_field").strategy(
+        cases.ConstInitializer(1)
+    )()
+    out_field = cases.allocate(cartesian_case, identity, "in_field").strategy(
+        cases.ConstInitializer(0)
+    )()
 
-def test_shift_by_one_execution(fieldview_backend):
-    size = 10
-    in_field = np_as_located_field(IDim)(np.arange(0, size, 1, dtype=np.float64))
-    out_field = np_as_located_field(IDim)(np.zeros((size)))
-    out_field_ref = np_as_located_field(IDim)(
-        np.array([i + 1.0 if i in range(0, size - 1) else 0 for i in range(0, size)])
+    cases.verify(
+        cartesian_case,
+        identity,
+        in_field,
+        out=out_field,
+        ref=np.ones((cartesian_case.default_sizes[IDim])),
     )
 
-    @field_operator
-    def shift_by_one(in_field: Field[[IDim], float64]) -> Field[[IDim], float64]:
+
+def test_shift_by_one_execution(cartesian_case):
+    @gtx.field_operator
+    def shift_by_one(in_field: cases.IFloatField) -> cases.IFloatField:
         return in_field(Ioff[1])
 
     # direct call to field operator
     # TODO(tehrengruber): slicing located fields not supported currently
     # shift_by_one(in_field, out=out_field[:-1], offset_provider={"Ioff": IDim})
 
-    @program
-    def shift_by_one_program(in_field: Field[[IDim], float64], out_field: Field[[IDim], float64]):
+    @gtx.program
+    def shift_by_one_program(in_field: cases.IFloatField, out_field: cases.IFloatField):
         shift_by_one(in_field, out=out_field[:-1])
 
-    shift_by_one_program.with_backend(fieldview_backend)(
-        in_field, out_field, offset_provider={"Ioff": IDim}
+    in_field = cases.allocate(cartesian_case, shift_by_one_program, "in_field").extend(
+        {IDim: (0, 1)}
+    )()
+    out_field = cases.allocate(cartesian_case, shift_by_one_program, "out_field")()
+
+    cases.verify(
+        cartesian_case,
+        shift_by_one_program,
+        in_field,
+        out_field,
+        inout=out_field.array()[:-1],
+        ref=in_field.array()[1:-1],
     )
 
-    assert np.allclose(out_field, out_field_ref)
-
-
-def test_copy_execution(fieldview_backend, copy_program_def):
-    size = 10
-    in_field = np_as_located_field(IDim)(np.ones((size)))
-    out_field = np_as_located_field(IDim)(np.zeros((size)))
-    copy_program = program(copy_program_def, backend=fieldview_backend)
-
-    copy_program(in_field, out_field, offset_provider={})
-
-    assert np.allclose(in_field, out_field)
-
 
-def test_double_copy_execution(fieldview_backend, double_copy_program_def):
-    size = 10
-    in_field = np_as_located_field(IDim)(np.ones((size)))
-    intermediate_field = np_as_located_field(IDim)(np.zeros((size)))
-    out_field = np_as_located_field(IDim)(np.zeros((size)))
-    double_copy_program = program(double_copy_program_def, backend=fieldview_backend)
+def test_copy_execution(cartesian_case, copy_program_def):
+    copy_program = gtx.program(copy_program_def, backend=cartesian_case.backend)
 
-    double_copy_program(in_field, intermediate_field, out_field, offset_provider={})
+    cases.verify_with_default_data(cartesian_case, copy_program, ref=lambda in_field: in_field)
 
-    assert np.allclose(in_field, out_field)
 
+def test_double_copy_execution(cartesian_case, double_copy_program_def):
+    double_copy_program = gtx.program(double_copy_program_def, backend=cartesian_case.backend)
 
-def test_copy_restricted_execution(fieldview_backend, copy_restrict_program_def):
-    size = 10
-    in_field = np_as_located_field(IDim)(np.ones((size)))
-    out_field = np_as_located_field(IDim)(np.zeros((size)))
-    out_field_ref = np_as_located_field(IDim)(
-        np.array([1 if i in range(1, 2) else 0 for i in range(0, size)])
+    cases.verify_with_default_data(
+        cartesian_case, double_copy_program, ref=lambda in_field, intermediate_field: in_field
     )
-    copy_restrict_program = program(copy_restrict_program_def, backend=fieldview_backend)
 
-    copy_restrict_program(in_field, out_field, offset_provider={})
 
-    assert np.allclose(out_field_ref, out_field)
+def test_copy_restricted_execution(cartesian_case, copy_restrict_program_def):
+    copy_restrict_program = gtx.program(copy_restrict_program_def, backend=cartesian_case.backend)
 
+    cases.verify_with_default_data(
+        cartesian_case,
+        copy_restrict_program,
+        ref=lambda in_field: np.array(
+            [
+                in_field[i] if i in range(1, 2) else 0
+                for i in range(0, cartesian_case.default_sizes[IDim])
+            ]
+        ),
+    )
 
-def test_calling_fo_from_fo_execution(fieldview_backend, identity_def):
-    size = 10
-    in_field = np_as_located_field(IDim)(2 * np.ones((size)))
-    out_field = np_as_located_field(IDim)(np.zeros((size)))
-    out_field_ref = np_as_located_field(IDim)(2 * 2 * 2 * np.ones((size)))
 
-    @field_operator
-    def pow_two(field: Field[[IDim], "float64"]) -> Field[[IDim], "float64"]:
+def test_calling_fo_from_fo_execution(cartesian_case):
+    @gtx.field_operator
+    def pow_two(field: cases.IFloatField) -> cases.IFloatField:
         return field * field
 
-    @field_operator
-    def pow_three(field: Field[[IDim], "float64"]) -> Field[[IDim], "float64"]:
+    @gtx.field_operator
+    def pow_three(field: cases.IFloatField) -> cases.IFloatField:
         return field * pow_two(field)
 
-    @program(backend=fieldview_backend)
-    def fo_from_fo_program(in_field: Field[[IDim], "float64"], out_field: Field[[IDim], "float64"]):
-        pow_three(in_field, out=out_field)
-
-    fo_from_fo_program(in_field, out_field, offset_provider={})
-
-    assert np.allclose(out_field, out_field_ref)
+    @gtx.program
+    def fo_from_fo_program(in_field: cases.IFloatField, out: cases.IFloatField):
+        pow_three(in_field, out=out)
 
+    cases.verify_with_default_data(
+        cartesian_case,
+        fo_from_fo_program,
+        ref=lambda in_field: in_field**3,
+    )
 
-def test_tuple_program_return_constructed_inside(fieldview_backend):
-    size = 10
-    a = np_as_located_field(IDim)(np.ones((size,)))
-    b = np_as_located_field(IDim)(2 * np.ones((size,)))
-    out_a = np_as_located_field(IDim)(np.zeros((size,)))
-    out_b = np_as_located_field(IDim)(np.zeros((size,)))
 
-    @field_operator
+def test_tuple_program_return_constructed_inside(cartesian_case):
+    @gtx.field_operator
     def pack_tuple(
-        a: Field[[IDim], float64], b: Field[[IDim], float64]
-    ) -> tuple[Field[[IDim], float64], Field[[IDim], float64]]:
+        a: cases.IFloatField, b: cases.IFloatField
+    ) -> tuple[cases.IFloatField, cases.IFloatField]:
         return (a, b)
 
-    @program(backend=fieldview_backend)
+    @gtx.program
     def prog(
-        a: Field[[IDim], float64],
-        b: Field[[IDim], float64],
-        out_a: Field[[IDim], float64],
-        out_b: Field[[IDim], float64],
+        a: cases.IFloatField,
+        b: cases.IFloatField,
+        out_a: cases.IFloatField,
+        out_b: cases.IFloatField,
     ):
         pack_tuple(a, b, out=(out_a, out_b))
 
-    prog(a, b, out_a, out_b, offset_provider={})
+    a = cases.allocate(cartesian_case, prog, "a")()
+    b = cases.allocate(cartesian_case, prog, "b")()
+    out_a = cases.allocate(cartesian_case, prog, "out_a")()
+    out_b = cases.allocate(cartesian_case, prog, "out_b")()
 
-    assert np.allclose(a, out_a)
-    assert np.allclose(b, out_b)
+    cases.run(cartesian_case, prog, a, b, out_a, out_b, offset_provider={})
 
+    assert np.allclose((a, b), (out_a, out_b))
 
-def test_tuple_program_return_constructed_inside_with_slicing(fieldview_backend):
-    size = 10
-    a = np_as_located_field(IDim)(np.ones((size,)))
-    b = np_as_located_field(IDim)(2 * np.ones((size,)))
-    out_a = np_as_located_field(IDim)(np.zeros((size,)))
-    out_b = np_as_located_field(IDim)(np.zeros((size,)))
 
-    @field_operator
+def test_tuple_program_return_constructed_inside_with_slicing(cartesian_case):
+    @gtx.field_operator
     def pack_tuple(
-        a: Field[[IDim], float64], b: Field[[IDim], float64]
-    ) -> tuple[Field[[IDim], float64], Field[[IDim], float64]]:
+        a: cases.IFloatField, b: cases.IFloatField
+    ) -> tuple[cases.IFloatField, cases.IFloatField]:
         return (a, b)
 
-    @program(backend=fieldview_backend)
+    @gtx.program
     def prog(
-        a: Field[[IDim], float64],
-        b: Field[[IDim], float64],
-        out_a: Field[[IDim], float64],
-        out_b: Field[[IDim], float64],
+        a: cases.IFloatField,
+        b: cases.IFloatField,
+        out_a: cases.IFloatField,
+        out_b: cases.IFloatField,
     ):
         pack_tuple(a, b, out=(out_a[1:], out_b[1:]))
 
-    prog(a, b, out_a, out_b, offset_provider={})
+    a = cases.allocate(cartesian_case, prog, "a").strategy(cases.ConstInitializer(1))()
+    b = cases.allocate(cartesian_case, prog, "b").strategy(cases.ConstInitializer(2))()
+    out_a = cases.allocate(cartesian_case, prog, "out_a").strategy(cases.ConstInitializer(0))()
+    out_b = cases.allocate(cartesian_case, prog, "out_b").strategy(cases.ConstInitializer(0))()
 
-    assert np.allclose(a[1:], out_a[1:])
-    assert out_a[0] == 0.0
-    assert np.allclose(b[1:], out_b[1:])
-    assert out_b[0] == 0.0
+    cases.run(cartesian_case, prog, a, b, out_a, out_b, offset_provider={})
 
+    assert np.allclose((a.array()[1:], b.array()[1:]), (out_a.array()[1:], out_b.array()[1:]))
+    assert out_a[0] == 0 and out_b[0] == 0
 
-def test_tuple_program_return_constructed_inside_nested(fieldview_backend):
-    size = 10
-    a = np_as_located_field(IDim)(np.ones((size,)))
-    b = np_as_located_field(IDim)(2 * np.ones((size,)))
-    c = np_as_located_field(IDim)(3 * np.ones((size,)))
-    out_a = np_as_located_field(IDim)(np.zeros((size,)))
-    out_b = np_as_located_field(IDim)(np.zeros((size,)))
-    out_c = np_as_located_field(IDim)(np.zeros((size,)))
 
-    @field_operator
+def test_tuple_program_return_constructed_inside_nested(cartesian_case):
+    @gtx.field_operator
     def pack_tuple(
-        a: Field[[IDim], float64], b: Field[[IDim], float64], c: Field[[IDim], float64]
-    ) -> tuple[tuple[Field[[IDim], float64], Field[[IDim], float64]], Field[[IDim], float64]]:
+        a: cases.IFloatField, b: cases.IFloatField, c: cases.IFloatField
+    ) -> tuple[tuple[cases.IFloatField, cases.IFloatField], cases.IFloatField]:
         return ((a, b), c)
 
-    @program(backend=fieldview_backend)
+    @gtx.program
     def prog(
-        a: Field[[IDim], float64],
-        b: Field[[IDim], float64],
-        c: Field[[IDim], float64],
-        out_a: Field[[IDim], float64],
-        out_b: Field[[IDim], float64],
-        out_c: Field[[IDim], float64],
+        a: cases.IFloatField,
+        b: cases.IFloatField,
+        c: cases.IFloatField,
+        out_a: cases.IFloatField,
+        out_b: cases.IFloatField,
+        out_c: cases.IFloatField,
     ):
         pack_tuple(a, b, c, out=((out_a, out_b), out_c))
 
-    prog(a, b, c, out_a, out_b, out_c, offset_provider={})
+    a = cases.allocate(cartesian_case, prog, "a").strategy(cases.ConstInitializer(1))()
+    b = cases.allocate(cartesian_case, prog, "b").strategy(cases.ConstInitializer(2))()
+    c = cases.allocate(cartesian_case, prog, "b").strategy(cases.ConstInitializer(3))()
+    out_a = cases.allocate(cartesian_case, prog, "out_a").strategy(cases.ConstInitializer(0))()
+    out_b = cases.allocate(cartesian_case, prog, "out_b").strategy(cases.ConstInitializer(0))()
+    out_c = cases.allocate(cartesian_case, prog, "out_c").strategy(cases.ConstInitializer(0))()
 
-    assert np.allclose(a, out_a)
-    assert np.allclose(b, out_b)
+    cases.run(cartesian_case, prog, a, b, c, out_a, out_b, out_c, offset_provider={})
 
+    assert np.allclose((a, b, c), (out_a, out_b, out_c))
 
-def test_wrong_argument_type(fieldview_backend, copy_program_def):
-    size = 10
-    inp = np_as_located_field(JDim)(np.ones((size,)))
-    out = np_as_located_field(IDim)(np.zeros((size,)))
 
-    copy_program = program(copy_program_def, backend=fieldview_backend)
+def test_wrong_argument_type(cartesian_case, copy_program_def):
+    copy_program = gtx.program(copy_program_def, backend=cartesian_case.backend)
+
+    inp = gtx.np_as_located_field(JDim)(np.ones((cartesian_case.default_sizes[JDim],)))
+    out = cases.allocate(cartesian_case, copy_program, "out").strategy(cases.ConstInitializer(1))()
 
     with pytest.raises(ProgramTypeError) as exc_info:
         # program is defined on Field[[IDim], ...], but we call with
@@ -249,21 +238,18 @@ def test_wrong_argument_type(fieldview_backend, copy_program_def):
         assert re.search(msg, exc_info.value.__cause__.args[0]) is not None
 
 
-def test_dimensions_domain():
-    size = 10
-    a = np_as_located_field(IDim, JDim)(np.ones((size, size)))
-    out_field = np_as_located_field(IDim, JDim)(np.ones((size, size)))
-
-    @field_operator()
-    def empty_domain_fieldop(a: Field[[IDim, JDim], float64]):
+def test_dimensions_domain(cartesian_case):
+    @gtx.field_operator
+    def empty_domain_fieldop(a: cases.IJField):
         return a
 
-    @program
-    def empty_domain_program(
-        a: Field[[IDim, JDim], float64], out_field: Field[[IDim, JDim], float64]
-    ):
+    @gtx.program
+    def empty_domain_program(a: cases.IJField, out_field: cases.IJField):
         empty_domain_fieldop(a, out=out_field, domain={JDim: (0, 1), IDim: (0, 1)})
 
+    a = cases.allocate(cartesian_case, empty_domain_program, "a")()
+    out_field = cases.allocate(cartesian_case, empty_domain_program, "out_field")()
+
     with pytest.raises(
         GTTypeError,
         match=(r"Dimensions in out field and field domain are not equivalent"),
@@ -271,40 +257,42 @@ def test_dimensions_domain():
         empty_domain_program(a, out_field, offset_provider={})
 
 
-def test_input_kwargs(fieldview_backend):
-    size = 10
-    input_1 = np_as_located_field(IDim, JDim)(np.ones((size, size)))
-    input_2 = np_as_located_field(IDim, JDim)(np.ones((size, size)) * 2)
-    input_3 = np_as_located_field(IDim, JDim)(np.ones((size, size)) * 3)
-
-    expected = np.asarray(input_3) * np.asarray(input_1) - np.asarray(input_2)
-
-    @field_operator(backend=fieldview_backend)
+def test_input_kwargs(cartesian_case):
+    @gtx.field_operator(backend=cartesian_case.backend)
     def fieldop_input_kwargs(
-        a: Field[[IDim, JDim], float64],
-        b: Field[[IDim, JDim], float64],
-        c: Field[[IDim, JDim], float64],
-    ) -> Field[[IDim, JDim], float64]:
+        a: cases.IJField,
+        b: cases.IJField,
+        c: cases.IJField,
+    ) -> cases.IJField:
         return c * a - b
 
-    out = np_as_located_field(IDim, JDim)(np.zeros((size, size)))
+    input_1 = cases.allocate(cartesian_case, fieldop_input_kwargs, "a")()
+    input_2 = cases.allocate(cartesian_case, fieldop_input_kwargs, "b")()
+    input_3 = cases.allocate(cartesian_case, fieldop_input_kwargs, "c")()
+    expected = np.asarray(input_3) * np.asarray(input_1) - np.asarray(input_2)
+    out = cases.allocate(cartesian_case, fieldop_input_kwargs, cases.RETURN)()
+
     fieldop_input_kwargs(input_1, b=input_2, c=input_3, out=out, offset_provider={})
     assert np.allclose(expected, out)
 
-    @program(backend=fieldview_backend)
+    @gtx.program(backend=cartesian_case.backend)
     def program_input_kwargs(
-        a: Field[[IDim, JDim], float64],
-        b: Field[[IDim, JDim], float64],
-        c: Field[[IDim, JDim], float64],
-        out: Field[[IDim, JDim], float64],
+        a: cases.IJField,
+        b: cases.IJField,
+        c: cases.IJField,
+        out: cases.IJField,
     ):
         fieldop_input_kwargs(a, b, c, out=out)
 
-    out = np_as_located_field(IDim, JDim)(np.zeros((size, size)))
+    out = cases.allocate(cartesian_case, program_input_kwargs, "out").strategy(
+        cases.ConstInitializer(0)
+    )()
     program_input_kwargs(input_1, b=input_2, c=input_3, out=out, offset_provider={})
     assert np.allclose(expected, out)
 
-    out = np_as_located_field(IDim, JDim)(np.zeros((size, size)))
+    out = cases.allocate(cartesian_case, program_input_kwargs, "out").strategy(
+        cases.ConstInitializer(0)
+    )()
     program_input_kwargs(a=input_1, b=input_2, c=input_3, out=out, offset_provider={})
     assert np.allclose(expected, out)
 
diff --git a/tests/next_tests/past_common_fixtures.py b/tests/next_tests/past_common_fixtures.py
index fbd1b306..3ac931f3 100644
--- a/tests/next_tests/past_common_fixtures.py
+++ b/tests/next_tests/past_common_fixtures.py
@@ -55,10 +55,8 @@ def make_tuple_op():
 def copy_program_def(identity_def):
     identity = gtx.field_operator(identity_def)
 
-    def copy_program(
-        in_field: gtx.Field[[IDim], "float64"], out_field: gtx.Field[[IDim], "float64"]
-    ):
-        identity(in_field, out=out_field)
+    def copy_program(in_field: gtx.Field[[IDim], "float64"], out: gtx.Field[[IDim], "float64"]):
+        identity(in_field, out=out)
 
     return copy_program
 
@@ -70,10 +68,10 @@ def double_copy_program_def(identity_def):
     def double_copy_program(
         in_field: gtx.Field[[IDim], "float64"],
         intermediate_field: gtx.Field[[IDim], "float64"],
-        out_field: gtx.Field[[IDim], "float64"],
+        out: gtx.Field[[IDim], "float64"],
     ):
         identity(in_field, out=intermediate_field)
-        identity(intermediate_field, out=out_field)
+        identity(intermediate_field, out=out)
 
     return double_copy_program
 
@@ -83,9 +81,9 @@ def copy_restrict_program_def(identity_def):
     identity = gtx.field_operator(identity_def)
 
     def copy_restrict_program(
-        in_field: gtx.Field[[IDim], "float64"], out_field: gtx.Field[[IDim], "float64"]
+        in_field: gtx.Field[[IDim], "float64"], out: gtx.Field[[IDim], "float64"]
     ):
-        identity(in_field, out=out_field[1:2])
+        identity(in_field, out=out[1:2])
 
     return copy_restrict_program
 
diff --git a/tests/next_tests/unit_tests/ffront_tests/test_func_to_past.py b/tests/next_tests/unit_tests/ffront_tests/test_func_to_past.py
index 12285d10..f1070c93 100644
--- a/tests/next_tests/unit_tests/ffront_tests/test_func_to_past.py
+++ b/tests/next_tests/unit_tests/ffront_tests/test_func_to_past.py
@@ -59,14 +59,14 @@ def test_copy_parsing(copy_program_def):
         id=eve.SymbolName("copy_program"),
         params=[
             P(past.Symbol, id=eve.SymbolName("in_field"), type=field_type),
-            P(past.Symbol, id=eve.SymbolName("out_field"), type=field_type),
+            P(past.Symbol, id=eve.SymbolName("out"), type=field_type),
         ],
         body=[
             P(
                 past.Call,
                 func=P(past.Name, id=past.SymbolRef("identity")),
                 args=[P(past.Name, id=past.SymbolRef("in_field"))],
-                kwargs={"out": P(past.Name, id=past.SymbolRef("out_field"))},
+                kwargs={"out": P(past.Name, id=past.SymbolRef("out"))},
             )
         ],
         location=P(past.SourceLocation),
@@ -87,7 +87,7 @@ def test_double_copy_parsing(double_copy_program_def):
         params=[
             P(past.Symbol, id=eve.SymbolName("in_field"), type=field_type),
             P(past.Symbol, id=eve.SymbolName("intermediate_field"), type=field_type),
-            P(past.Symbol, id=eve.SymbolName("out_field"), type=field_type),
+            P(past.Symbol, id=eve.SymbolName("out"), type=field_type),
         ],
         body=[
             P(
@@ -100,7 +100,7 @@ def test_double_copy_parsing(double_copy_program_def):
                 past.Call,
                 func=P(past.Name, id=past.SymbolRef("identity")),
                 args=[P(past.Name, id=past.SymbolRef("intermediate_field"))],
-                kwargs={"out": P(past.Name, id=past.SymbolRef("out_field"))},
+                kwargs={"out": P(past.Name, id=past.SymbolRef("out"))},
             ),
         ],
     )
@@ -135,7 +135,7 @@ def test_copy_restrict_parsing(copy_restrict_program_def):
         id=eve.SymbolName("copy_restrict_program"),
         params=[
             P(past.Symbol, id=eve.SymbolName("in_field"), type=field_type),
-            P(past.Symbol, id=eve.SymbolName("out_field"), type=field_type),
+            P(past.Symbol, id=eve.SymbolName("out"), type=field_type),
         ],
         body=[
             P(
@@ -145,7 +145,7 @@ def test_copy_restrict_parsing(copy_restrict_program_def):
                 kwargs={
                     "out": P(
                         past.Subscript,
-                        value=P(past.Name, id=past.SymbolRef("out_field")),
+                        value=P(past.Name, id=past.SymbolRef("out")),
                         slice_=slice_pattern_node,
                     )
                 },
diff --git a/tests/next_tests/unit_tests/ffront_tests/test_past_to_itir.py b/tests/next_tests/unit_tests/ffront_tests/test_past_to_itir.py
index d6ae5318..ff3abcf2 100644
--- a/tests/next_tests/unit_tests/ffront_tests/test_past_to_itir.py
+++ b/tests/next_tests/unit_tests/ffront_tests/test_past_to_itir.py
@@ -59,24 +59,24 @@ def test_copy_lowering(copy_program_def, itir_identity_fundef):
                     fun=P(itir.SymRef, id=eve.SymbolRef("named_range")),
                     args=[
                         P(itir.AxisLiteral, value="IDim"),
-                        P(itir.Literal, value="0", type=itir.INTEGER_INDEX_BUILTIN),
-                        P(itir.SymRef, id=eve.SymbolRef("__out_field_size_0")),
+                        P(itir.Literal, value="0", type="int32"),
+                        P(itir.SymRef, id=eve.SymbolRef("__out_size_0")),
                     ],
                 )
             ],
         ),
         stencil=P(itir.SymRef, id=eve.SymbolRef("identity")),
         inputs=[P(itir.SymRef, id=eve.SymbolRef("in_field"))],
-        output=P(itir.SymRef, id=eve.SymbolRef("out_field")),
+        output=P(itir.SymRef, id=eve.SymbolRef("out")),
     )
     fencil_pattern = P(
         itir.FencilDefinition,
         id=eve.SymbolName("copy_program"),
         params=[
             P(itir.Sym, id=eve.SymbolName("in_field")),
-            P(itir.Sym, id=eve.SymbolName("out_field")),
+            P(itir.Sym, id=eve.SymbolName("out")),
             P(itir.Sym, id=eve.SymbolName("__in_field_size_0")),
-            P(itir.Sym, id=eve.SymbolName("__out_field_size_0")),
+            P(itir.Sym, id=eve.SymbolName("__out_size_0")),
         ],
         closures=[closure_pattern],
     )
@@ -112,9 +112,9 @@ def test_copy_restrict_lowering(copy_restrict_program_def, itir_identity_fundef)
         id=eve.SymbolName("copy_restrict_program"),
         params=[
             P(itir.Sym, id=eve.SymbolName("in_field")),
-            P(itir.Sym, id=eve.SymbolName("out_field")),
+            P(itir.Sym, id=eve.SymbolName("out")),
             P(itir.Sym, id=eve.SymbolName("__in_field_size_0")),
-            P(itir.Sym, id=eve.SymbolName("__out_field_size_0")),
+            P(itir.Sym, id=eve.SymbolName("__out_size_0")),
         ],
         closures=[closure_pattern],
     )
diff --git a/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_cse.py b/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_cse.py
index 3ce17bc8..a53667d9 100644
--- a/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_cse.py
+++ b/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_cse.py
@@ -177,3 +177,20 @@ def test_if_can_deref_eligible_extraction():
 
     actual = CSE().visit(testee)
     assert actual == expected
+
+
+def test_if_eligible_extraction():
+    # Test that a subexpression only occurring in the condition of an `if_` is moved outside the
+    # if statement.
+
+    # if ((a ∧ b) ∧ (a ∧ b)) then c else d
+    testee = im.if_(
+        im.and_(im.and_("a", "b"), im.and_("a", "b")),
+        "c",
+        "d",
+    )
+    # (λ(_cs_1) → if _cs_1 ∧ _cs_1 then c else d)(a ∧ b)
+    expected = im.let("_cs_1", im.and_("a", "b"))(im.if_(im.and_("_cs_1", "_cs_1"), "c", "d"))
+
+    actual = CSE().visit(testee)
+    assert actual == expected