Repository: GridTools/gt4py
Pull request #1664 (2024-09-25 19:20:38)

feat[next]: Add support for IfStmt in ITIR
Adds support for stateful if statements in ITIR for ITIR embedded and GTFN backend.


Patch:
diff --git a/src/gt4py/next/iterator/embedded.py b/src/gt4py/next/iterator/embedded.py
index e1b52043..afe0cec4 100644
--- a/src/gt4py/next/iterator/embedded.py
+++ b/src/gt4py/next/iterator/embedded.py
@@ -1603,6 +1603,26 @@ def set_at(expr: common.Field, domain: common.DomainLike, target: common.Mutable
     operators._tuple_assign_field(target, expr, common.domain(domain))
 
 
+@runtime.if_stmt.register(EMBEDDED)
+def if_stmt(cond: bool, true_branch: Callable[[], None], false_branch: Callable[[], None]) -> None:
+    """
+    (Stateful) if statement.
+
+    The two branches are represented as lambda functions, such that they are not executed eagerly.
+    This is required to avoid out-of-bounds accesses. Note that a dedicated built-in is required,
+    contrary to using a plain python if-stmt, such that tracing / double roundtrip works.
+
+    Arguments:
+        cond: The condition to decide which branch to execute.
+        true_branch: A lambda function to be executed when `cond` is `True`.
+        false_branch: A lambda function to be executed when `cond` is `False`.
+    """
+    if cond:
+        true_branch()
+    else:
+        false_branch()
+
+
 def _compute_at_position(
     sten: Callable,
     ins: Sequence[common.Field],
diff --git a/src/gt4py/next/iterator/ir.py b/src/gt4py/next/iterator/ir.py
index 43e249f5..b2a54950 100644
--- a/src/gt4py/next/iterator/ir.py
+++ b/src/gt4py/next/iterator/ir.py
@@ -210,6 +210,12 @@ class SetAt(Stmt):  # from JAX array.at[...].set()
     target: Expr  # `make_tuple` or SymRef
 
 
+class IfStmt(Stmt):
+    cond: Expr
+    true_branch: list[Stmt]
+    false_branch: list[Stmt]
+
+
 class Temporary(Node):
     id: Coerced[eve.SymbolName]
     domain: Optional[Expr] = None
@@ -242,3 +248,4 @@ StencilClosure.__hash__ = Node.__hash__  # type: ignore[method-assign]
 FencilDefinition.__hash__ = Node.__hash__  # type: ignore[method-assign]
 Program.__hash__ = Node.__hash__  # type: ignore[method-assign]
 SetAt.__hash__ = Node.__hash__  # type: ignore[method-assign]
+IfStmt.__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 6af57327..08459a94 100644
--- a/src/gt4py/next/iterator/pretty_parser.py
+++ b/src/gt4py/next/iterator/pretty_parser.py
@@ -8,7 +8,7 @@
 
 from typing import Union
 
-from lark import lark, lexer as lark_lexer, visitors as lark_visitors
+from lark import lark, lexer as lark_lexer, tree as lark_tree, visitors as lark_visitors
 
 from gt4py.next.iterator import ir
 from gt4py.next.iterator.ir_utils import ir_makers as im
@@ -21,6 +21,7 @@ GRAMMAR = """
         | declaration
         | stencil_closure
         | set_at
+        | if_stmt
         | program
         | prec0
 
@@ -78,13 +79,17 @@ GRAMMAR = """
         | named_range
         | "(" prec0 ")"
 
+    ?stmt: set_at | if_stmt
+    set_at: prec0 "@" prec0 "←" prec1 ";"
+    else_branch_seperator: "else"
+    if_stmt: "if" "(" prec0 ")" "{" ( stmt )* "}" else_branch_seperator "{" ( stmt )* "}"
+
     named_range: AXIS_NAME ":" "[" prec0 "," prec0 ")"
     function_definition: ID_NAME "=" "λ(" ( SYM "," )* SYM? ")" "→" prec0 ";"
     declaration: ID_NAME "=" "temporary(" "domain=" prec0 "," "dtype=" TYPE_LITERAL ")" ";"
     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 )+ "}"
+    program: ID_NAME "(" ( SYM "," )* SYM ")" "{" ( function_definition )* ( declaration )* ( stmt )+ "}"
 
     %import common (CNAME, SIGNED_FLOAT, SIGNED_INT, WS)
     %ignore WS
@@ -215,6 +220,27 @@ class ToIrTransformer(lark_visitors.Transformer):
         output, stencil, *inputs, domain = args
         return ir.StencilClosure(domain=domain, stencil=stencil, output=output, inputs=inputs)
 
+    def if_stmt(self, cond: ir.Expr, *args):
+        found_else_seperator = False
+        true_branch = []
+        false_branch = []
+        for arg in args:
+            if isinstance(arg, lark_tree.Tree):
+                assert arg.data == "else_branch_seperator"
+                found_else_seperator = True
+                continue
+
+            if not found_else_seperator:
+                true_branch.append(arg)
+            else:
+                false_branch.append(arg)
+
+        return ir.IfStmt(
+            cond=cond,
+            true_branch=true_branch,
+            false_branch=false_branch,
+        )
+
     def declaration(self, *args: ir.Expr) -> ir.Temporary:
         tid, domain, dtype = args
         return ir.Temporary(id=tid, domain=domain, dtype=dtype)
@@ -253,7 +279,7 @@ class ToIrTransformer(lark_visitors.Transformer):
             elif isinstance(arg, ir.Temporary):
                 declarations.append(arg)
             else:
-                assert isinstance(arg, ir.SetAt)
+                assert isinstance(arg, ir.Stmt)
                 body.append(arg)
         return ir.Program(
             id=fid,
diff --git a/src/gt4py/next/iterator/pretty_printer.py b/src/gt4py/next/iterator/pretty_printer.py
index d9f62717..99287f8a 100644
--- a/src/gt4py/next/iterator/pretty_printer.py
+++ b/src/gt4py/next/iterator/pretty_printer.py
@@ -299,6 +299,19 @@ class PrettyPrinter(NodeTranslator):
         )
         return self._optimum(h, v)
 
+    def visit_IfStmt(self, node: ir.IfStmt, *, prec: int) -> list[str]:
+        cond = self.visit(node.cond, prec=0)
+        true_branch = self._vmerge(*self.visit(node.true_branch, prec=0))
+        false_branch = self._vmerge(*self.visit(node.false_branch, prec=0))
+
+        hhead = self._hmerge(["if ("], cond, [") {"])
+        vhead = self._vmerge(["if ("], cond, [") {"])
+        head = self._optimum(hhead, vhead)
+
+        return self._vmerge(
+            head, self._indent(true_branch), ["} else {"], self._indent(false_branch), ["}"]
+        )
+
     def visit_FencilDefinition(self, node: ir.FencilDefinition, *, prec: int) -> list[str]:
         assert prec == 0
         function_definitions = self.visit(node.function_definitions, prec=0)
diff --git a/src/gt4py/next/iterator/runtime.py b/src/gt4py/next/iterator/runtime.py
index 5627e5c4..ad85d154 100644
--- a/src/gt4py/next/iterator/runtime.py
+++ b/src/gt4py/next/iterator/runtime.py
@@ -26,7 +26,7 @@ if TYPE_CHECKING:
     # TODO(tehrengruber): remove cirular dependency and import unconditionally
     from gt4py.next import backend as next_backend
 
-__all__ = ["offset", "fundef", "fendef", "closure", "set_at"]
+__all__ = ["offset", "fundef", "fendef", "closure", "set_at", "if_stmt"]
 
 
 @dataclass(frozen=True)
@@ -214,3 +214,8 @@ def closure(*args):  # TODO remove
 @builtin_dispatch
 def set_at(*args):
     return BackendNotSelectedError()
+
+
+@builtin_dispatch
+def if_stmt(*args):
+    return BackendNotSelectedError()
diff --git a/src/gt4py/next/iterator/tracing.py b/src/gt4py/next/iterator/tracing.py
index a0f0b863..6772d4b5 100644
--- a/src/gt4py/next/iterator/tracing.py
+++ b/src/gt4py/next/iterator/tracing.py
@@ -247,6 +247,27 @@ def set_at(expr: itir.Expr, domain: itir.Expr, target: itir.Expr) -> None:
     TracerContext.add_stmt(itir.SetAt(expr=expr, domain=domain, target=target))
 
 
+@iterator.runtime.if_stmt.register(TRACING)
+def if_stmt(
+    cond: itir.Expr, true_branch_f: typing.Callable, false_branch_f: typing.Callable
+) -> None:
+    true_branch: List[itir.Stmt] = []
+    false_branch: List[itir.Stmt] = []
+
+    old_body = TracerContext.body
+    TracerContext.body = true_branch
+    true_branch_f()
+
+    TracerContext.body = false_branch
+    false_branch_f()
+
+    TracerContext.body = old_body
+
+    TracerContext.add_stmt(
+        itir.IfStmt(cond=cond, true_branch=true_branch, false_branch=false_branch)
+    )
+
+
 def _contains_tuple_dtype_field(arg):
     if isinstance(arg, tuple):
         return any(_contains_tuple_dtype_field(el) for el in arg)
diff --git a/src/gt4py/next/iterator/type_system/inference.py b/src/gt4py/next/iterator/type_system/inference.py
index 46566924..bc1095df 100644
--- a/src/gt4py/next/iterator/type_system/inference.py
+++ b/src/gt4py/next/iterator/type_system/inference.py
@@ -499,6 +499,12 @@ class ITIRTypeInference(eve.NodeTranslator):
             lambda dtype: ts.FieldType(dims=domain.dims, dtype=dtype), node.dtype
         )
 
+    def visit_IfStmt(self, node: itir.IfStmt, *, ctx) -> None:
+        cond = self.visit(node.cond, ctx=ctx)
+        assert cond == ts.ScalarType(kind=ts.ScalarKind.BOOL)
+        self.visit(node.true_branch, ctx=ctx)
+        self.visit(node.false_branch, ctx=ctx)
+
     def visit_SetAt(self, node: itir.SetAt, *, ctx) -> None:
         self.visit(node.expr, ctx=ctx)
         self.visit(node.domain, ctx=ctx)
diff --git a/src/gt4py/next/program_processors/codegens/gtfn/codegen.py b/src/gt4py/next/program_processors/codegens/gtfn/codegen.py
index c9107469..92dbcede 100644
--- a/src/gt4py/next/program_processors/codegens/gtfn/codegen.py
+++ b/src/gt4py/next/program_processors/codegens/gtfn/codegen.py
@@ -166,6 +166,16 @@ class GTFNCodegen(codegen.TemplatedGenerator):
         "{backend}.vertical_executor({axis})().{'.'.join('arg(' + a + ')' for a in args)}.{'.'.join(scans)}.execute();"
     )
 
+    IfStmt = as_mako(
+        """
+          if (${cond}) {
+            ${'\\n'.join(true_branch)}
+          } else {
+            ${'\\n'.join(false_branch)}
+          }
+        """
+    )
+
     ScanPassDefinition = as_mako(
         """
         struct ${id} : ${'gtfn::fwd' if _this_node.forward else 'gtfn::bwd'} {
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 825623d6..1995e4de 100644
--- a/src/gt4py/next/program_processors/codegens/gtfn/gtfn_ir.py
+++ b/src/gt4py/next/program_processors/codegens/gtfn/gtfn_ir.py
@@ -96,23 +96,23 @@ class Backend(Node):
     domain: Union[SymRef, CartesianDomain, UnstructuredDomain]
 
 
-def _is_ref_or_tuple_expr_of_ref(expr: Expr) -> bool:
+def _is_ref_literal_or_tuple_expr_of_ref(expr: Expr) -> bool:
     if (
         isinstance(expr, FunCall)
         and isinstance(expr.fun, SymRef)
         and expr.fun.id == "tuple_get"
         and len(expr.args) == 2
-        and _is_ref_or_tuple_expr_of_ref(expr.args[1])
+        and _is_ref_literal_or_tuple_expr_of_ref(expr.args[1])
     ):
         return True
     if (
         isinstance(expr, FunCall)
         and isinstance(expr.fun, SymRef)
         and expr.fun.id == "make_tuple"
-        and all(_is_ref_or_tuple_expr_of_ref(arg) for arg in expr.args)
+        and all(_is_ref_literal_or_tuple_expr_of_ref(arg) for arg in expr.args)
     ):
         return True
-    if isinstance(expr, SymRef):
+    if isinstance(expr, (SymRef, Literal)):
         return True
     return False
 
@@ -125,7 +125,8 @@ class SidComposite(Expr):
         self: datamodels.DataModelTP, attribute: datamodels.Attribute, value: list[Expr]
     ) -> None:
         if not all(
-            isinstance(el, (SidFromScalar, SidComposite)) or _is_ref_or_tuple_expr_of_ref(el)
+            isinstance(el, (SidFromScalar, SidComposite))
+            or _is_ref_literal_or_tuple_expr_of_ref(el)
             for el in value
         ):
             raise ValueError(
@@ -140,11 +141,15 @@ class SidFromScalar(Expr):
     def _arg_validator(
         self: datamodels.DataModelTP, attribute: datamodels.Attribute, value: Expr
     ) -> None:
-        if not _is_ref_or_tuple_expr_of_ref(value):
+        if not _is_ref_literal_or_tuple_expr_of_ref(value):
             raise ValueError("Only 'SymRef' or tuple expr of 'SymRef' allowed.")
 
 
-class StencilExecution(Node):
+class Stmt(Node):
+    pass
+
+
+class StencilExecution(Stmt):
     backend: Backend
     stencil: SymRef
     output: Union[SymRef, SidComposite]
@@ -158,13 +163,19 @@ class Scan(Node):
     init: Expr
 
 
-class ScanExecution(Node):
+class ScanExecution(Stmt):
     backend: Backend
     scans: list[Scan]
     args: list[Expr]
     axis: SymRef
 
 
+class IfStmt(Stmt):
+    cond: Expr
+    true_branch: list[Stmt]
+    false_branch: list[Stmt]
+
+
 class TemporaryAllocation(Node):
     id: SymbolName
     dtype: str
@@ -199,7 +210,7 @@ class Program(Node, ValidatedSymbolTableTrait):
     function_definitions: list[
         Union[FunctionDefinition, ScanPassDefinition, ImperativeFunctionDefinition]
     ]
-    executions: list[Union[StencilExecution, ScanExecution]]
+    executions: list[Stmt]
     offset_definitions: list[TagDefinition]
     grid_type: common.GridType
     temporaries: list[TemporaryAllocation]
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 7fa7bfec..d729a5ba 100644
--- a/src/gt4py/next/program_processors/codegens/gtfn/gtfn_module.py
+++ b/src/gt4py/next/program_processors/codegens/gtfn/gtfn_module.py
@@ -82,7 +82,7 @@ class GTFNTranslationStep(
 
     def _process_regular_arguments(
         self,
-        program: itir.FencilDefinition,
+        program: itir.FencilDefinition | itir.Program,
         arg_types: tuple[ts.TypeSpec, ...],
         offset_provider: dict[str, Connectivity | Dimension],
     ) -> tuple[list[interface.Parameter], list[str]]:
@@ -161,10 +161,10 @@ class GTFNTranslationStep(
 
     def _preprocess_program(
         self,
-        program: itir.FencilDefinition,
+        program: itir.FencilDefinition | itir.Program,
         offset_provider: dict[str, Connectivity | Dimension],
     ) -> itir.Program:
-        if not self.enable_itir_transforms:
+        if isinstance(program, itir.FencilDefinition) and not self.enable_itir_transforms:
             return fencil_to_program.FencilToProgram().apply(
                 program
             )  # FIXME[#1582](tehrengruber): should be removed after refactoring to combined IR
@@ -195,7 +195,7 @@ class GTFNTranslationStep(
 
     def generate_stencil_source(
         self,
-        program: itir.FencilDefinition,
+        program: itir.FencilDefinition | itir.Program,
         offset_provider: dict[str, Connectivity | Dimension],
         column_axis: Optional[common.Dimension],
     ) -> str:
@@ -216,7 +216,6 @@ class GTFNTranslationStep(
     ) -> stages.ProgramSource[languages.NanobindSrcL, languages.LanguageWithHeaderFilesSettings]:
         """Generate GTFN C++ code from the ITIR definition."""
         program: itir.FencilDefinition | itir.Program = inp.data
-        assert isinstance(program, itir.FencilDefinition)
 
         # handle regular parameters and arguments of the program (i.e. what the user defined in
         #  the program)
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 e9a0ad16..3bd96d14 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
@@ -24,6 +24,7 @@ from gt4py.next.program_processors.codegens.gtfn.gtfn_ir import (
     CastExpr,
     FunCall,
     FunctionDefinition,
+    IfStmt,
     IntegralConstant,
     Lambda,
     Literal,
@@ -66,8 +67,11 @@ _vertical_dimension = "gtfn::unstructured::dim::vertical"
 _horizontal_dimension = "gtfn::unstructured::dim::horizontal"
 
 
-def _get_domains(node: Iterable[itir.Stmt]) -> Iterable[itir.FunCall]:
-    return eve_utils.xiter(node).if_isinstance(itir.SetAt).getattr("domain").to_set()
+def _get_domains(nodes: Iterable[itir.Stmt]) -> Iterable[itir.FunCall]:
+    result = set()
+    for node in nodes:
+        result |= node.walk_values().if_isinstance(itir.SetAt).getattr("domain").to_set()
+    return result
 
 
 def _extract_grid_type(domain: itir.FunCall) -> common.GridType:
@@ -573,6 +577,13 @@ class GTFN_lowering(eve.NodeTranslator, eve.VisitorWithSymbolTableTrait):
     def visit_Stmt(self, node: itir.Stmt, **kwargs: Any) -> None:
         raise AssertionError("All Stmts need to be handled explicitly.")
 
+    def visit_IfStmt(self, node: itir.IfStmt, **kwargs: Any) -> IfStmt:
+        return IfStmt(
+            cond=self.visit(node.cond, **kwargs),
+            true_branch=self.visit(node.true_branch, **kwargs),
+            false_branch=self.visit(node.false_branch, **kwargs),
+        )
+
     def visit_SetAt(
         self, node: itir.SetAt, *, extracted_functions: list, **kwargs: Any
     ) -> Union[StencilExecution, ScanExecution]:
diff --git a/src/gt4py/next/program_processors/runners/roundtrip.py b/src/gt4py/next/program_processors/runners/roundtrip.py
index c062e14b..93e6d09c 100644
--- a/src/gt4py/next/program_processors/runners/roundtrip.py
+++ b/src/gt4py/next/program_processors/runners/roundtrip.py
@@ -63,6 +63,10 @@ def ${id}(${','.join(params)}):
     """
     )
     SetAt = as_mako("set_at(${expr}, ${domain}, ${target})")
+    IfStmt = as_mako("""if_stmt(${cond}, 
+        lambda: [${','.join(true_branch)}],
+        lambda: [${','.join(false_branch)}]
+    )""")
 
     def visit_Temporary(self, node: itir.Temporary, **kwargs: Any) -> str:
         assert (


Test patch:
diff --git a/tests/next_tests/definitions.py b/tests/next_tests/definitions.py
index a62d0518..123384a0 100644
--- a/tests/next_tests/definitions.py
+++ b/tests/next_tests/definitions.py
@@ -93,6 +93,7 @@ USES_CONSTANT_FIELDS = "uses_constant_fields"
 USES_DYNAMIC_OFFSETS = "uses_dynamic_offsets"
 USES_FLOORDIV = "uses_floordiv"
 USES_IF_STMTS = "uses_if_stmts"
+USES_IR_IF_STMTS = "uses_ir_if_stmts"
 USES_INDEX_FIELDS = "uses_index_fields"
 USES_LIFT_EXPRESSIONS = "uses_lift_expressions"
 USES_NEGATIVE_MODULO = "uses_negative_modulo"
@@ -135,6 +136,7 @@ COMMON_SKIP_TEST_LIST = [
     (USES_SPARSE_FIELDS_AS_OUTPUT, XFAIL, UNSUPPORTED_MESSAGE),
 ]
 DACE_SKIP_TEST_LIST = COMMON_SKIP_TEST_LIST + [
+    (USES_IR_IF_STMTS, XFAIL, UNSUPPORTED_MESSAGE),
     (USES_SCALAR_IN_DOMAIN_AND_FO, XFAIL, UNSUPPORTED_MESSAGE),
     (USES_INDEX_FIELDS, XFAIL, UNSUPPORTED_MESSAGE),
     (USES_LIFT_EXPRESSIONS, XFAIL, UNSUPPORTED_MESSAGE),
diff --git a/tests/next_tests/integration_tests/multi_feature_tests/iterator_tests/test_if_stmt.py b/tests/next_tests/integration_tests/multi_feature_tests/iterator_tests/test_if_stmt.py
new file mode 100644
index 00000000..2dde7d76
--- /dev/null
+++ b/tests/next_tests/integration_tests/multi_feature_tests/iterator_tests/test_if_stmt.py
@@ -0,0 +1,75 @@
+# GT4Py - GridTools Framework
+#
+# Copyright (c) 2014-2024, ETH Zurich
+# All rights reserved.
+#
+# Please, refer to the LICENSE file in the root directory.
+# SPDX-License-Identifier: BSD-3-Clause
+
+# 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
+
+import numpy as np
+import pytest
+
+import gt4py.next as gtx
+from gt4py.next.iterator.builtins import cartesian_domain, deref, as_fieldop, named_range
+from gt4py.next.iterator.runtime import set_at, if_stmt, fendef, fundef, offset
+from gt4py.next.program_processors.runners import gtfn
+
+from next_tests.unit_tests.conftest import program_processor, run_processor
+
+i = offset("i")
+
+
+@fundef
+def multiply(alpha, inp):
+    return deref(alpha) * deref(inp)
+
+
+IDim = gtx.Dimension("IDim")
+
+
+@pytest.mark.uses_ir_if_stmts
+@pytest.mark.parametrize("cond", [True, False])
+def test_if_stmt(program_processor, cond):
+    program_processor, validate = program_processor
+    size = 10
+
+    @fendef(offset_provider={"i": IDim})
+    def fencil(cond1, inp, out):
+        domain = cartesian_domain(named_range(IDim, 0, size))
+        if_stmt(
+            cond1,
+            lambda: set_at(
+                as_fieldop(multiply, domain)(1.0, inp),
+                domain,
+                out,
+            ),
+            lambda: set_at(
+                as_fieldop(multiply, domain)(2.0, inp),
+                domain,
+                out,
+            ),
+        )
+
+    rng = np.random.default_rng()
+    inp = gtx.as_field([IDim], rng.normal(size=size))
+    out = gtx.as_field([IDim], np.zeros(size))
+    ref = inp if cond else 2.0 * inp
+
+    run_processor(fencil, program_processor, cond, inp, out)
+
+    if validate:
+        assert np.allclose(out.asnumpy(), ref.asnumpy())
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 5bb359df..da4bea88 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
@@ -231,6 +231,37 @@ def test_set_at():
     assert actual == expected
 
 
+def test_if_stmt():
+    testee = """if (cond) { 
+      y @ cartesian_domain() ← x;
+      if (cond) {
+        y @ cartesian_domain() ← x;
+      } else {
+      }
+    } else {
+      y @ cartesian_domain() ← x;
+    }"""
+    stmt = ir.SetAt(
+        expr=im.ref("x"),
+        domain=im.domain("cartesian_domain", {}),
+        target=im.ref("y"),
+    )
+    expected = ir.IfStmt(
+        cond=im.ref("cond"),
+        true_branch=[
+            stmt,
+            ir.IfStmt(
+                cond=im.ref("cond"),
+                true_branch=[stmt],
+                false_branch=[],
+            ),
+        ],
+        false_branch=[stmt],
+    )
+    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}"
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 acfb1d0b..05cd6b68 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
@@ -456,3 +456,25 @@ def test_program_tuple_setat_short_target():
         isinstance(result.body[0].target.type, ts.TupleType)
         and len(result.body[0].target.type.types) == 1
     )
+
+
+def test_if_stmt():
+    cartesian_domain = im.call("cartesian_domain")(
+        im.call("named_range")(itir.AxisLiteral(value="IDim"), 0, 1)
+    )
+
+    testee = itir.IfStmt(
+        cond=im.literal_from_value(True),
+        true_branch=[
+            itir.SetAt(
+                expr=im.as_fieldop("deref", cartesian_domain)(im.ref("inp", float_i_field)),
+                domain=cartesian_domain,
+                target=im.ref("out", float_i_field),
+            )
+        ],
+        false_branch=[],
+    )
+
+    result = itir_type_inference.infer(testee, offset_provider={}, allow_undeclared_symbols=True)
+    assert result.cond.type == bool_type
+    assert result.true_branch[0].expr.type == float_i_field