Repository: GridTools/gt4py
Pull request #1985 (2025-04-15 18:48:59)

fix[next]: Temporary extraction for fields defined on different domains
Splits tuple values that are defined on different domains into separate SetAts.

During temporary extraction we might extract a SetAt that contains a tuple with elements defined on different domains (e.g. Vertex and Edge). SetAts are required to have a single domain, therefore these SetAts are split into one SetAt per domain at the end of temporary extraction.

Note: it seems many of the cases in icon4py are actually cases where the domain is semantically equal, but the domain expression is different (alla `1+1` != `2`). This can result in many unnecessary kernel calls.

havogt: >it seems many of the cases in icon4py are actually cases where the domain is semantically equal, but the domain expression is different (alla 1+1 != 2). This can result in many unnecessary kernel calls.

Is this horizontal or vertical? With the variant precompile the domain expressions should get simpler, also I would like to do another round of removing superfluous concat_wheres on the horizontal. So the problem might disappear.
tehrengruber: The case I saw was horizontal only.

Patch:
diff --git a/src/gt4py/next/iterator/ir_utils/domain_utils.py b/src/gt4py/next/iterator/ir_utils/domain_utils.py
index 5973715d..0655a840 100644
--- a/src/gt4py/next/iterator/ir_utils/domain_utils.py
+++ b/src/gt4py/next/iterator/ir_utils/domain_utils.py
@@ -40,7 +40,7 @@ def _max_domain_sizes_by_location_type(offset_provider: Mapping[str, Any]) -> di
     return sizes
 
 
-@dataclasses.dataclass
+@dataclasses.dataclass(frozen=True)
 class SymbolicRange:
     start: itir.Expr
     stop: itir.Expr
@@ -49,13 +49,16 @@ class SymbolicRange:
         return SymbolicRange(im.plus(self.start, distance), im.plus(self.stop, distance))
 
 
-@dataclasses.dataclass
+@dataclasses.dataclass(frozen=True)
 class SymbolicDomain:
     grid_type: Literal["unstructured_domain", "cartesian_domain"]
     ranges: dict[
         common.Dimension, SymbolicRange
     ]  # TODO(havogt): remove `AxisLiteral` by `Dimension` everywhere
 
+    def __hash__(self):
+        return hash((self.grid_type, frozenset(self.ranges.items())))
+
     @classmethod
     def from_expr(cls, node: itir.Node) -> SymbolicDomain:
         assert isinstance(node, itir.FunCall) and node.fun in [
diff --git a/src/gt4py/next/iterator/ir_utils/ir_makers.py b/src/gt4py/next/iterator/ir_utils/ir_makers.py
index aae98335..f47dfbd1 100644
--- a/src/gt4py/next/iterator/ir_utils/ir_makers.py
+++ b/src/gt4py/next/iterator/ir_utils/ir_makers.py
@@ -7,7 +7,7 @@
 # SPDX-License-Identifier: BSD-3-Clause
 
 import typing
-from typing import Callable, Optional, Union
+from typing import Any, Callable, Optional, Union
 
 from gt4py._core import definitions as core_defs
 from gt4py.next import common
@@ -38,7 +38,9 @@ def sym(sym_or_name: Union[str, itir.Sym], type_: str | ts.TypeSpec | None = Non
 
 
 def ref(
-    ref_or_name: Union[str, itir.SymRef], type_: str | ts.TypeSpec | None = None
+    ref_or_name: Union[str, itir.SymRef],
+    type_: str | ts.TypeSpec | None = None,
+    annex: dict[str, Any] | None = None,
 ) -> itir.SymRef:
     """
     Convert to SymRef if necessary.
@@ -57,8 +59,13 @@ def ref(
     """
     if isinstance(ref_or_name, itir.SymRef):
         assert not type_
+        assert not annex
         return ref_or_name
-    return itir.SymRef(id=ref_or_name, type=ensure_type(type_))
+    ref = itir.SymRef(id=ref_or_name, type=ensure_type(type_))
+    if annex is not None:
+        for key, value in annex.items():
+            setattr(ref.annex, key, value)
+    return ref
 
 
 def ensure_expr(literal_or_expr: Union[str, core_defs.Scalar, itir.Expr]) -> itir.Expr:
diff --git a/src/gt4py/next/iterator/transforms/global_tmps.py b/src/gt4py/next/iterator/transforms/global_tmps.py
index b7b36305..682dcb5e 100644
--- a/src/gt4py/next/iterator/transforms/global_tmps.py
+++ b/src/gt4py/next/iterator/transforms/global_tmps.py
@@ -9,6 +9,7 @@
 from __future__ import annotations
 
 import functools
+from collections.abc import Sequence
 from typing import Callable, Optional
 
 from gt4py.eve import utils as eve_utils
@@ -19,11 +20,94 @@ from gt4py.next.iterator.ir_utils import (
     domain_utils,
     ir_makers as im,
 )
+from gt4py.next.iterator.ir_utils.domain_utils import SymbolicDomain
 from gt4py.next.iterator.transforms import cse, infer_domain, inline_lambdas
 from gt4py.next.iterator.type_system import inference as type_inference
 from gt4py.next.type_system import type_info, type_specifications as ts
 
 
+def select_elems_by_domain(
+    select_domain: SymbolicDomain,
+    target: itir.Expr,
+    args: Sequence[itir.Expr],
+    domains: tuple[SymbolicDomain, ...],
+):
+    """
+    Select all elements of possibly nested tuples for the given domain.
+
+    Returns (non-nested) tuples of the selected elements and the corresponding targets.
+    """
+    new_targets = []
+    new_els = []
+    for i, (el, el_domain) in enumerate(zip(args, domains)):
+        current_target = im.tuple_get(i, target)
+        if isinstance(el_domain, tuple):
+            assert cpm.is_call_to(el, "make_tuple")
+            more_targets, more_els = select_elems_by_domain(
+                select_domain, current_target, el.args, el_domain
+            )
+            new_els.extend(more_els)
+            new_targets.extend(more_targets)
+        else:
+            assert isinstance(el_domain, SymbolicDomain)
+            if el_domain == select_domain:
+                new_targets.append(current_target)
+                new_els.append(el)
+    return new_targets, new_els
+
+
+def _set_at_for_domain(stmt: itir.SetAt, domain: SymbolicDomain) -> itir.SetAt:
+    """Extract all elements with given domain into a new `SetAt` statement."""
+    tuple_expr = stmt.expr
+    assert cpm.is_call_to(tuple_expr, "make_tuple")
+    targets, expr_els = select_elems_by_domain(
+        domain, stmt.target, tuple_expr.args, stmt.expr.annex.domain
+    )
+    new_expr = im.make_tuple(*expr_els)
+    new_expr.annex.domain = domain
+
+    return itir.SetAt(expr=new_expr, domain=domain.as_expr(), target=im.make_tuple(*targets))
+
+
+def _populate_and_homogenize_domains(stmts: list[itir.Stmt]) -> list[itir.Stmt]:
+    """
+    Splits `SetAt` statements with multiple domains into multiple statements.
+
+    `SetAt`s have a single domain therefore the target domains have to be the same.
+    We support nested tuples by flattening them and recombining them into non-nested tuples.
+    """
+    new_stmts: list[itir.Stmt] = []
+    for stmt in stmts:
+        if isinstance(stmt, itir.SetAt):
+            assert hasattr(stmt.expr.annex, "domain")
+            distinct_domains: list[domain_utils.SymbolicDomain] = list(
+                # ordered set for reproducibility
+                dict.fromkeys(next_utils.flatten_nested_tuple(stmt.expr.annex.domain))
+            )
+            if len(distinct_domains) == 1:
+                new_stmts.append(
+                    itir.SetAt(
+                        expr=stmt.expr,
+                        domain=distinct_domains[0].as_expr(),  # insert concrete domain
+                        target=stmt.target,
+                    )
+                )
+            else:
+                new_stmts.extend(_set_at_for_domain(stmt, domain) for domain in distinct_domains)
+        elif isinstance(stmt, itir.IfStmt):
+            new_stmts.append(
+                itir.IfStmt(
+                    cond=stmt.cond,
+                    true_branch=_populate_and_homogenize_domains(stmt.true_branch),
+                    false_branch=_populate_and_homogenize_domains(stmt.false_branch),
+                )
+            )
+        else:
+            raise ValueError("Expected `itir.SetAt` or `itir.IfStmt`.")
+
+    return new_stmts
+
+
 def _transform_if(
     stmt: itir.Stmt, declarations: list[itir.Temporary], uids: eve_utils.UIDGenerator
 ) -> Optional[list[itir.Stmt]]:
@@ -74,24 +158,6 @@ def _transform_by_pattern(
         #    or a tuple thereof)
         #  - one `SetAt` statement that materializes the expression into the temporary
         for tmp_sym, tmp_expr in extracted_fields.items():
-            domain: infer_domain.DomainAccess = tmp_expr.annex.domain
-
-            # TODO(tehrengruber): Implement. This happens when the expression is a combination
-            #  of an `if_` call with a tuple, e.g., `if_(cond, {a, b}, {c, d})`. As long as we are
-            #  able to eliminate all tuples, e.g., by propagating the scalar ifs to the top-level
-            #  of a SetAt, the CollapseTuple pass will eliminate most of this cases.
-            if isinstance(domain, tuple):
-                flattened_domains: tuple[domain_utils.SymbolicDomain] = (
-                    next_utils.flatten_nested_tuple(domain)  # type: ignore[assignment]  # mypy not smart enough
-                )
-                if not all(d == flattened_domains[0] for d in flattened_domains):
-                    raise NotImplementedError(
-                        "Tuple expressions with different domains is not supported yet."
-                    )
-                domain = flattened_domains[0]
-            assert isinstance(domain, domain_utils.SymbolicDomain)
-            domain_expr = domain.as_expr()
-
             assert isinstance(tmp_expr.type, ts.TypeSpec)
             tmp_names: str | tuple[str | tuple, ...] = type_info.apply_to_primitive_constituents(
                 lambda x: uids.sequential_id(),
@@ -103,20 +169,58 @@ def _transform_by_pattern(
             ) = type_info.apply_to_primitive_constituents(
                 type_info.extract_dtype,
                 tmp_expr.type,
-                tuple_constructor=lambda *elements: tuple(elements),
+                tuple_constructor=lambda *elements: elements,
             )
 
-            # allocate temporary for all tuple elements
-            def allocate_temporary(tmp_name: str, dtype: ts.ScalarType):
-                declarations.append(itir.Temporary(id=tmp_name, domain=domain_expr, dtype=dtype))  # noqa: B023 # function only used inside loop
+            tmp_domains: SymbolicDomain | tuple[SymbolicDomain | tuple, ...] = tmp_expr.annex.domain
+
+            if cpm.is_applied_as_fieldop(tmp_expr):
+                # In this case all tuple elements have the same size (or will be `NEVER`).
+                # Create the tuple structure with that domain.
+                domain = list(
+                    set(next_utils.flatten_nested_tuple((tmp_domains,)))
+                    - {infer_domain.DomainAccessDescriptor.NEVER}  # type: ignore[arg-type] # type should always be `SymbolicDomain`
+                )
+                assert len(domain) == 1
+                # this is the domain used as initial value in the tuple construction below
+                tmp_domains = domain[0]
+
+            def get_domain(
+                _, path: tuple[int, ...]
+            ) -> domain_utils.SymbolicDomain:  # function only used inside loop
+                domain = functools.reduce(
+                    lambda var, idx: var[idx] if isinstance(var, tuple) else var,
+                    path,
+                    tmp_domains,  # noqa: B023
+                )
+                assert isinstance(domain, domain_utils.SymbolicDomain)
+                return domain
 
-            next_utils.tree_map(allocate_temporary)(tmp_names, tmp_dtypes)
+            # The following propagates the domains to the tuple structure of `tmp_expr.type`.
+            # `tmp_domains` might not have this structure because domain inference was not able to infer the tuple structure.
+            tmp_domains = type_info.apply_to_primitive_constituents(
+                get_domain,
+                tmp_expr.type,
+                with_path_arg=True,
+                tuple_constructor=lambda *elements: tuple(elements),
+            )
+
+            declarations.extend(
+                itir.Temporary(id=tmp_name, domain=domain.as_expr(), dtype=dtype)
+                for tmp_name, domain, dtype in zip(
+                    next_utils.flatten_nested_tuple((tmp_names,)),
+                    next_utils.flatten_nested_tuple((tmp_domains,)),
+                    next_utils.flatten_nested_tuple((tmp_dtypes,)),
+                    strict=True,
+                )
+            )
 
             # if the expr is a field this just gives a simple `itir.SymRef`, otherwise we generate a
             #  `make_tuple` expression.
             target_expr: itir.Expr = next_utils.tree_map(
-                lambda x: im.ref(x), result_collection_constructor=lambda els: im.make_tuple(*els)
-            )(tmp_names)  # type: ignore[assignment]  # typing of tree_map does not reflect action of `result_collection_constructor` yet
+                lambda name, domain: im.ref(name, annex={"domain": domain}),
+                result_collection_constructor=lambda els: im.make_tuple(*els),
+            )(tmp_names, tmp_domains)  # type: ignore[assignment]  # typing of tree_map does not reflect action of `result_collection_constructor` yet
 
             # note: the let would be removed automatically by the `cse.extract_subexpression`, but
             # we remove it here for readability & debuggability.
@@ -127,7 +231,13 @@ def _transform_by_pattern(
             # TODO(tehrengruber): _transform_stmt not needed if deepest_expr_first=True
             tmp_stmts.extend(
                 _transform_stmt(
-                    itir.SetAt(target=target_expr, domain=domain_expr, expr=tmp_expr),
+                    itir.SetAt(
+                        target=target_expr,
+                        # The domain is populated later in `_populate_and_homogenize_domains`
+                        # at this point the `SetAt` contains possibly expressions on different domains.
+                        domain=im.ref("UNPOPULATED"),
+                        expr=tmp_expr,
+                    ),
                     declarations,
                     uids,
                 )
@@ -171,6 +281,8 @@ def _transform_stmt(
         if not did_transform:
             stmts.append(stmt)
 
+    stmts = _populate_and_homogenize_domains(stmts)
+
     return stmts
 
 


Test patch:
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 52d77e5f..67213a30 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
@@ -11,7 +11,7 @@ from typing import Optional
 from gt4py.next import common
 from gt4py.next.iterator import builtins, ir as itir
 from gt4py.next.iterator.ir_utils import ir_makers as im
-from gt4py.next.iterator.transforms import global_tmps, infer_domain
+from gt4py.next.iterator.transforms import global_tmps, infer_domain, collapse_tuple
 from gt4py.next.iterator.type_system import inference as type_inference
 from gt4py.next.type_system import type_specifications as ts
 
@@ -222,3 +222,206 @@ def test_nested_if():
 
     actual = global_tmps.create_global_tmps(testee, offset_provider)
     assert actual == expected
+
+
+def test_tuple_different_domain():
+    domain01 = im.domain("cartesian_domain", {IDim: (0, 1)})
+    domain12 = im.domain("cartesian_domain", {IDim: (1, 2)})
+    offset_provider = {"I": IDim}
+
+    def add_shifted(domain: itir.FunCall | None = None):
+        return im.as_fieldop(
+            im.lambda_("it1", "it2")(im.plus(im.deref("it1"), im.deref(im.shift("I", 1)("it2")))),
+            domain,
+        )
+
+    params = [
+        im.sym("cond", ts.ScalarType(kind=ts.ScalarKind.BOOL)),
+        im.sym("inp1", i_field_type),
+        im.sym("inp2", i_field_type),
+        im.sym("out", i_field_type),
+    ]
+    testee = program_factory(
+        params=params,
+        body=[
+            itir.SetAt(
+                # val = if_(cond, make_tuple(as_fieldop(...), as_fieldop(...)), make_tuple())
+                # val1 = if_(cond, as_fieldop(...), as_fieldop(...))
+                # val2 = if_(cond, as_fieldop(...), as_fieldop(...))
+                # materialize_into(out, make_tuple(as_fieldop(...), ...))
+                target=im.ref("out"),
+                expr=im.let(
+                    "val",
+                    im.if_(
+                        "cond", im.make_tuple("inp1", "inp2"), im.make_tuple("inp2", "inp1")
+                    ),  # todo: fix, the domain is strange
+                )(add_shifted(None)(im.tuple_get(0, "val"), im.tuple_get(1, "val"))),
+                domain=domain01,
+            )
+        ],
+    )
+    testee = type_inference.infer(testee, offset_provider_type=offset_provider)
+    testee = infer_domain.infer_program(testee, offset_provider=offset_provider)
+
+    expected = program_factory(
+        params=params,
+        declarations=[
+            itir.Temporary(id="__tmp_1", domain=domain01, dtype=float_type),
+            itir.Temporary(id="__tmp_2", domain=domain12, dtype=float_type),
+        ],
+        body=[
+            itir.IfStmt(
+                cond=im.ref("cond"),
+                true_branch=[
+                    itir.SetAt(
+                        target=im.make_tuple(im.ref("__tmp_1")),
+                        expr=im.make_tuple(im.ref("inp1")),
+                        domain=domain01,
+                    ),
+                    itir.SetAt(
+                        target=im.make_tuple(im.ref("__tmp_2")),
+                        expr=im.make_tuple(im.ref("inp2")),
+                        domain=domain12,
+                    ),
+                ],
+                false_branch=[
+                    itir.SetAt(
+                        target=im.make_tuple(im.ref("__tmp_1")),
+                        expr=im.make_tuple(im.ref("inp2")),
+                        domain=domain01,
+                    ),
+                    itir.SetAt(
+                        target=im.make_tuple(im.ref("__tmp_2")),
+                        expr=im.make_tuple(im.ref("inp1")),
+                        domain=domain12,
+                    ),
+                ],
+            ),
+            itir.SetAt(
+                target=im.ref("out"),
+                expr=add_shifted(domain01)(
+                    im.ref("__tmp_1"),
+                    im.ref("__tmp_2"),
+                ),
+                domain=domain01,
+            ),
+        ],
+    )
+
+    actual = global_tmps.create_global_tmps(testee, offset_provider)
+    actual = collapse_tuple.CollapseTuple.apply(actual)
+    assert actual == expected
+
+
+def test_tuple_different_domain_nested():
+    domain01 = im.domain("cartesian_domain", {IDim: (0, 1)})
+    domain12 = im.domain("cartesian_domain", {IDim: (1, 2)})
+    domainm10 = im.domain("cartesian_domain", {IDim: (-1, 0)})
+    offset_provider = {"I": IDim}
+
+    def add_shifted(domain: itir.FunCall | None = None):
+        return im.as_fieldop(
+            im.lambda_("it1", "it2", "it3")(
+                im.plus(
+                    im.deref("it1"),
+                    im.plus(im.deref(im.shift("I", 1)("it2")), im.deref(im.shift("I", -1)("it3"))),
+                )
+            ),
+            domain,
+        )
+
+    params = [
+        im.sym("cond", ts.ScalarType(kind=ts.ScalarKind.BOOL)),
+        im.sym("inp1", i_field_type),
+        im.sym("inp2", i_field_type),
+        im.sym("inp3", i_field_type),
+        im.sym("inp4", i_field_type),
+        im.sym("inp5", i_field_type),
+        im.sym("out", i_field_type),
+    ]
+    testee = program_factory(
+        params=params,
+        body=[
+            itir.SetAt(
+                target=im.ref("out"),
+                expr=im.let(
+                    "val",
+                    im.if_(
+                        "cond",
+                        im.make_tuple(im.make_tuple("inp1", "inp2"), "inp3"),
+                        im.make_tuple(im.make_tuple("inp4", "inp5"), "inp1"),
+                    ),
+                )(
+                    add_shifted(None)(
+                        im.tuple_get(0, im.tuple_get(0, "val")),
+                        im.tuple_get(1, im.tuple_get(0, "val")),
+                        im.tuple_get(1, "val"),
+                    )
+                ),
+                domain=domain01,
+            )
+        ],
+    )
+    testee = type_inference.infer(testee, offset_provider_type=offset_provider)
+    testee = infer_domain.infer_program(testee, offset_provider=offset_provider)
+
+    expected = program_factory(
+        params=params,
+        declarations=[
+            itir.Temporary(id="__tmp_1", domain=domain01, dtype=float_type),
+            itir.Temporary(id="__tmp_2", domain=domain12, dtype=float_type),
+            itir.Temporary(id="__tmp_3", domain=domainm10, dtype=float_type),
+        ],
+        body=[
+            itir.IfStmt(
+                cond=im.ref("cond"),
+                true_branch=[
+                    itir.SetAt(
+                        target=im.make_tuple(im.ref("__tmp_1")),
+                        expr=im.make_tuple(im.ref("inp1")),
+                        domain=domain01,
+                    ),
+                    itir.SetAt(
+                        target=im.make_tuple(im.ref("__tmp_2")),
+                        expr=im.make_tuple(im.ref("inp2")),
+                        domain=domain12,
+                    ),
+                    itir.SetAt(
+                        target=im.make_tuple(im.ref("__tmp_3")),
+                        expr=im.make_tuple(im.ref("inp3")),
+                        domain=domainm10,
+                    ),
+                ],
+                false_branch=[
+                    itir.SetAt(
+                        target=im.make_tuple(im.ref("__tmp_1")),
+                        expr=im.make_tuple(im.ref("inp4")),
+                        domain=domain01,
+                    ),
+                    itir.SetAt(
+                        target=im.make_tuple(im.ref("__tmp_2")),
+                        expr=im.make_tuple(im.ref("inp5")),
+                        domain=domain12,
+                    ),
+                    itir.SetAt(
+                        target=im.make_tuple(im.ref("__tmp_3")),
+                        expr=im.make_tuple(im.ref("inp1")),
+                        domain=domainm10,
+                    ),
+                ],
+            ),
+            itir.SetAt(
+                target=im.ref("out"),
+                expr=add_shifted(domain01)(
+                    im.ref("__tmp_1"),
+                    im.ref("__tmp_2"),
+                    im.ref("__tmp_3"),
+                ),
+                domain=domain01,
+            ),
+        ],
+    )
+
+    actual = global_tmps.create_global_tmps(testee, offset_provider)
+    actual = collapse_tuple.CollapseTuple.apply(actual)
+    assert actual == expected