feat[next]: Tracer support part 1: tree_map

src/gt4py/next/ffront/foast_to_gtir.py

@@ -412,23 +412,15 @@ def _visit_where(self, node: foast.Call, **kwargs: Any) -> itir.FunCall:
             return self._lower_and_map("if_", *node.args)
 
         cond_ = self.visit(node.args[0])
+        true_ = self.visit(node.args[1])
+        false_ = self.visit(node.args[2])
         cond_symref_name = f"__cond_{cond_.fingerprint()}"
 
-        def create_if(
-            true_: itir.Expr, false_: itir.Expr, arg_types: tuple[ts.TypeSpec, ts.TypeSpec]
-        ) -> itir.FunCall:
-            return _map(
-                "if_",
-                (im.ref(cond_symref_name), true_, false_),
-                (node.args[0].type, *arg_types),
+        result = im.tree_map(
+            im.lambda_("__a", "__b")(
+                im.op_as_fieldop("if_")(im.ref(cond_symref_name), im.ref("__a"), im.ref("__b"))
             )
-
-        result = lowering_utils.process_elements(
-            create_if,
-            (self.visit(node.args[1]), self.visit(node.args[2])),
-            node.type,
-            arg_types=(node.args[1].type, node.args[2].type),
-        )
+        )(true_, false_)
 
         return im.let(cond_symref_name, cond_)(result)
 

src/gt4py/next/iterator/builtins.py

@@ -57,6 +57,11 @@ def map_(*args):
     raise BackendNotSelectedError()
 
 
+@builtin_dispatch
+def tree_map(*args):
+    raise BackendNotSelectedError()
+
+
 @builtin_dispatch
 def make_const_list(*args):
     raise BackendNotSelectedError()
@@ -498,7 +503,8 @@ def get_domain_range(*args):
     "lift",
     "make_const_list",
     "make_tuple",
-    "map_",
+    "tree_map",
+    "map_",  # TODO: rename to map_list
     "named_range",
     "neighbors",
     "reduce",

src/gt4py/next/iterator/ir_utils/ir_makers.py

@@ -629,6 +629,11 @@ def map_(op):
     return call(call("map_")(op))
 
 
+def tree_map(op):
+    """Create a `tree_map` call: tree_map(op)(tup1, tup2, ...)."""
+    return call(call("tree_map")(op))
+
+
 def reduce(op, expr):
     """Create a `reduce` call."""
     return call(call("reduce")(op, expr))

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

@@ -23,6 +23,7 @@
     prune_empty_concat_where,
     remove_broadcast,
     symbol_ref_utils,
+    unroll_tree_map,
 )
 from gt4py.next.iterator.transforms.collapse_list_get import CollapseListGet
 from gt4py.next.iterator.transforms.collapse_tuple import CollapseTuple
@@ -176,6 +177,26 @@ def apply_common_transforms(
     )  # domain inference does not support dynamic offsets yet
     ir = infer_domain_ops.InferDomainOps.apply(ir)
     ir = concat_where.canonicalize_domain_argument(ir)
+    ir = unroll_tree_map.UnrollTreeMap.apply(ir, uids=uids)
+
+    # After UnrollTreeMap, collapse `tuple_get(i, let(...)(make_tuple(...)))` patterns so that
+    # domain inference does not encounter `as_fieldop` nodes inside dead tuple elements
+    # (which would receive NEVER domain). Do multiple iterations for nested `let`s.
+    for _ in range(10):
+        collapsed = ir
+        ir = CollapseTuple.apply(
+            ir,
+            enabled_transformations=(
+                CollapseTuple.Transformation.PROPAGATE_TUPLE_GET
+                | CollapseTuple.Transformation.COLLAPSE_TUPLE_GET_MAKE_TUPLE
+            ),
+            uids=uids,
+            offset_provider_type=offset_provider_type,
+        )  # type: ignore[assignment]  # always an itir.Program
+        if ir == collapsed:
+            break
+    else:
+        raise RuntimeError("'CollapseTuple' did not converge after `UnrollTreeMap`.")
 
     ir = infer_domain.infer_program(
         ir,
@@ -290,6 +311,23 @@ def apply_fieldview_transforms(
 
     ir = infer_domain_ops.InferDomainOps.apply(ir)
     ir = concat_where.canonicalize_domain_argument(ir)
+    ir = unroll_tree_map.UnrollTreeMap.apply(ir, uids=uids)
+    for _ in range(10):
+        prev = ir
+        ir = CollapseTuple.apply(
+            ir,
+            enabled_transformations=(
+                CollapseTuple.Transformation.PROPAGATE_TUPLE_GET
+                | CollapseTuple.Transformation.COLLAPSE_TUPLE_GET_MAKE_TUPLE
+            ),
+            uids=uids,
+            offset_provider_type=offset_provider_type,
+        )  # type: ignore[assignment]  # always an itir.Program
+        if ir == prev:
+            break
+    else:
+        raise RuntimeError("'CollapseTuple' did not converge after `UnrollTreeMap`.")
+
     ir = ConstantFolding.apply(ir)  # type: ignore[assignment]  # always an itir.Program
 
     ir = infer_domain.infer_program(

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

@@ -0,0 +1,85 @@
+# 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
+import dataclasses
+
+from gt4py import eve
+from gt4py.next import utils
+from gt4py.next.iterator import ir as itir
+from gt4py.next.iterator.ir_utils import common_pattern_matcher as cpm, ir_makers as im
+from gt4py.next.iterator.type_system import inference as itir_inference
+from gt4py.next.type_system import type_specifications as ts
+
+
+def _unroll(
+    f: itir.Expr,
+    tup_types: list[ts.TupleType],
+    tup_exprs: list[itir.Expr],
+) -> itir.Expr:
+    """Recursively expand ``tree_map(f)(tup0, tup1, ...)`` into ``make_tuple`` / ``tuple_get``."""
+    assert tup_types, "tree_map requires at least one tuple argument."
+    n = len(tup_types[0].types)
+    if any(len(t.types) != n for t in tup_types[1:]):
+        raise ValueError(
+            f"All tree_map arguments must have the same tuple structure at each level, "
+            f"got {[len(t.types) for t in tup_types]}."
+        )
+
+    elements: list[itir.Expr] = []
+    for i in range(n):
+        child_types = [t.types[i] for t in tup_types]
+        child_exprs = [im.tuple_get(i, e) for e in tup_exprs]
+
+        all_tuples = all(isinstance(ct, ts.TupleType) for ct in child_types)
+        all_leaves = all(not isinstance(ct, ts.TupleType) for ct in child_types)
+        if all_tuples:
+            nested_types = [ct for ct in child_types if isinstance(ct, ts.TupleType)]
+            elements.append(_unroll(f, nested_types, child_exprs))
+        elif all_leaves:
+            elements.append(im.call(f)(*child_exprs))
+        else:
+            raise ValueError(
+                "All tree_map arguments must have the same tree structure "
+                "(all leaves must be reached simultaneously)."
+            )
+
+    return im.make_tuple(*elements)
+
+
+@dataclasses.dataclass
+class UnrollTreeMap(eve.NodeTranslator):
+    PRESERVED_ANNEX_ATTRS = ("domain",)
+
+    uids: utils.IDGeneratorPool
+
+    @classmethod
+    def apply(cls, program: itir.Program, *, uids: utils.IDGeneratorPool):
+        return cls(uids=uids).visit(program)
+
+    def visit_FunCall(self, node: itir.FunCall):
+        node = self.generic_visit(node)
+
+        if not cpm.is_call_to(node.fun, "tree_map"):
+            return node
+
+        f = node.fun.args[0]
+        tup_args = node.args
+        tup_types: list[ts.TupleType] = []
+        for tup in tup_args:
+            itir_inference.reinfer(tup)
+            assert isinstance(tup.type, ts.TupleType)
+            tup_types.append(tup.type)
+
+        tup_refs = [next(self.uids["_utm"]) for _ in tup_args]
+        body = _unroll(f, tup_types, [im.ref(r) for r in tup_refs])
+
+        result = body
+        for ref_name, tup in reversed(list(zip(tup_refs, tup_args))):
+            result = im.let(ref_name, tup)(result)
+
+        itir_inference.reinfer(result)
+        return result

src/gt4py/next/iterator/type_system/type_synthesizer.py

@@ -633,6 +633,42 @@ def applied_map(
     return applied_map
 
 
+@_register_builtin_type_synthesizer(fun_names=["tree_map"])
+def _tree_map(op: TypeSynthesizer) -> TypeSynthesizer:
+    @type_synthesizer
+    def applied_map(
+        *args: ts.TupleType, offset_provider_type: common.OffsetProviderType
+    ) -> ts.TupleType:
+        if not args:
+            raise TypeError("tree_map requires at least one argument.")
+
+        def _recurse(*arg_types: ts.TypeSpec) -> ts.TypeSpec:
+            all_tuples = all(isinstance(a, ts.TupleType) for a in arg_types)
+            all_leaves = all(not isinstance(a, ts.TupleType) for a in arg_types)
+            if all_tuples:
+                tup_types = [a for a in arg_types if isinstance(a, ts.TupleType)]
+                n = len(tup_types[0].types)
+                if any(len(t.types) != n for t in tup_types[1:]):
+                    raise TypeError(
+                        f"All tree_map arguments must have the same tuple structure at each level, "
+                        f"got {[len(t.types) for t in tup_types]}."
+                    )
+                return ts.TupleType(
+                    types=[_recurse(*(a.types[i] for a in tup_types)) for i in range(n)]
+                )
+            elif all_leaves:
+                return op(*arg_types, offset_provider_type=offset_provider_type)  # type: ignore[return-value]
+            else:
+                raise TypeError(
+                    "All tree_map arguments must have the same tree structure "
+                    "(all leaves must be reached simultaneously)."
+                )
+
+        return _recurse(*args)  # type: ignore[return-value]
+
+    return applied_map
+
+
 @_register_builtin_type_synthesizer
 def reduce(op: TypeSynthesizer, init: ts.TypeSpec) -> TypeSynthesizer:
     @type_synthesizer

tests/next_tests/unit_tests/ffront_tests/test_foast_to_gtir.py

@@ -207,10 +207,9 @@ def foo(
         lowered
     )  # we generate a let for the condition which is removed by inlining for easier testing
 
-    reference = im.make_tuple(
-        im.op_as_fieldop("if_")("a", im.tuple_get(0, "b"), im.tuple_get(0, "c")),
-        im.op_as_fieldop("if_")("a", im.tuple_get(1, "b"), im.tuple_get(1, "c")),
-    )
+    reference = im.tree_map(
+        im.lambda_("__a", "__b")(im.op_as_fieldop("if_")("a", im.ref("__a"), im.ref("__b")))
+    )("b", "c")
 
     assert lowered_inlined.expr == reference
 

tests/next_tests/unit_tests/iterator_tests/test_type_inference.py

@@ -155,6 +155,23 @@ def expression_test_cases():
             ),
             ts.ListType(element_type=int_type, offset_type=V2EDim),
         ),
+        # tree_map
+        (
+            im.tree_map(im.ref("plus"))(
+                im.ref("t1", ts.TupleType(types=[int_type, int_type])),
+                im.ref("t2", ts.TupleType(types=[int_type, int_type])),
+            ),
+            ts.TupleType(types=[int_type, int_type]),
+        ),
+        (
+            im.tree_map(im.ref("not_"))(
+                im.ref(
+                    "t",
+                    ts.TupleType(types=[bool_type, ts.TupleType(types=[bool_type, bool_type])]),
+                ),
+            ),
+            ts.TupleType(types=[bool_type, ts.TupleType(types=[bool_type, bool_type])]),
+        ),
         # reduce
         (im.reduce("plus", 0)(im.ref("l", int_list_type)), int_type),
         (

tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_unroll_tree_map.py

@@ -0,0 +1,92 @@
+# 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
+
+from gt4py.next import common, utils
+from gt4py.next.iterator import ir as itir
+from gt4py.next.iterator.ir_utils import ir_makers as im
+from gt4py.next.iterator.transforms.unroll_tree_map import UnrollTreeMap
+from gt4py.next.type_system import type_specifications as ts
+
+IDim = common.Dimension("IDim")
+T = ts.ScalarType(kind=ts.ScalarKind.FLOAT64)
+i_field = ts.FieldType(dims=[IDim], dtype=T)
+i_tuple_field = ts.TupleType(types=[i_field, i_field])
+i_nested_tuple_field = ts.TupleType(types=[i_tuple_field, i_field])
+
+i_domain = im.call("cartesian_domain")(im.named_range(itir.AxisLiteral(value="IDim"), 0, 1))
+
+
+def _make_program(params: list[itir.Sym], expr: itir.Expr) -> itir.Program:
+    return itir.Program(
+        id="testee",
+        function_definitions=[],
+        params=[*params, im.sym("out", i_field)],
+        declarations=[],
+        body=[
+            itir.SetAt(
+                expr=expr,
+                domain=i_domain,
+                target=im.ref("out"),
+            )
+        ],
+    )
+
+
+def _neg():
+    return im.lambda_("__a")(im.op_as_fieldop("neg")("__a"))
+
+
+def _plus():
+    return im.lambda_("__a", "__b")(im.op_as_fieldop("plus")("__a", "__b"))
+
+
+def test_multi_arg():
+    uids = utils.IDGeneratorPool()
+    program = _make_program(
+        [im.sym("a", i_tuple_field), im.sym("b", i_tuple_field)],
+        im.call(im.call("tree_map")(_plus()))(
+            im.ref("a", i_tuple_field), im.ref("b", i_tuple_field)
+        ),
+    )
+    result = UnrollTreeMap.apply(program, uids=uids)
+
+    expected = _make_program(
+        [im.sym("a", i_tuple_field), im.sym("b", i_tuple_field)],
+        im.let("_utm_0", "a")(
+            im.let("_utm_1", "b")(
+                im.make_tuple(
+                    im.call(_plus())(im.tuple_get(0, "_utm_0"), im.tuple_get(0, "_utm_1")),
+                    im.call(_plus())(im.tuple_get(1, "_utm_0"), im.tuple_get(1, "_utm_1")),
+                )
+            )
+        ),
+    )
+    assert result == expected
+
+
+def test_nested():
+    uids = utils.IDGeneratorPool()
+    program = _make_program(
+        [im.sym("t", i_nested_tuple_field)],
+        im.call(im.call("tree_map")(_neg()))(im.ref("t", i_nested_tuple_field)),
+    )
+    result = UnrollTreeMap.apply(program, uids=uids)
+
+    expected = _make_program(
+        [im.sym("t", i_nested_tuple_field)],
+        im.let("_utm_0", "t")(
+            im.make_tuple(
+                im.make_tuple(
+                    im.call(_neg())(im.tuple_get(0, im.tuple_get(0, "_utm_0"))),
+                    im.call(_neg())(im.tuple_get(1, im.tuple_get(0, "_utm_0"))),
+                ),
+                im.call(_neg())(im.tuple_get(1, "_utm_0")),
+            )
+        ),
+    )
+    assert result == expected