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,24 @@ 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
 
     ir = infer_domain.infer_program(
         ir,
@@ -290,6 +309,20 @@ 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
     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,72 @@
+# 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``."""
+    n = len(tup_types[0].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]
+
+        if all(isinstance(ct, ts.TupleType) for ct in child_types):
+            nested_types = [ct for ct in child_types if isinstance(ct, ts.TupleType)]
+            elements.append(_unroll(f, nested_types, child_exprs))
+        else:
+            elements.append(im.call(f)(*child_exprs))
+
+    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,28 @@ 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:
+        def _recurse(*arg_types: ts.TypeSpec) -> ts.TypeSpec:
+            if isinstance(arg_types[0], ts.TupleType):
+                tup_types = [a for a in arg_types if isinstance(a, ts.TupleType)]
+                return ts.TupleType(
+                    types=[
+                        _recurse(*(a.types[i] for a in tup_types))
+                        for i in range(len(arg_types[0].types))
+                    ]
+                )
+            return op(*arg_types, offset_provider_type=offset_provider_type)  # type: ignore[return-value]
+
+        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/transforms_tests/test_unroll_tree_map.py

@@ -0,0 +1,43 @@
+# 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.iterator.ir_utils import ir_makers as im
+from gt4py.next.iterator.transforms.unroll_tree_map import _unroll
+from gt4py.next.type_system import type_specifications as ts
+
+
+T = ts.ScalarType(kind=ts.ScalarKind.FLOAT64)
+TT = ts.TupleType(types=[T, T])
+
+
+def test_single_arg():
+    result = _unroll(im.ref("f"), [TT], [im.ref("t")])
+    expected = im.make_tuple(im.call("f")(im.tuple_get(0, "t")), im.call("f")(im.tuple_get(1, "t")))
+    assert result == expected
+
+
+def test_multi_arg():
+    result = _unroll(im.ref("f"), [TT, TT], [im.ref("a"), im.ref("b")])
+    expected = im.make_tuple(
+        im.call("f")(im.tuple_get(0, "a"), im.tuple_get(0, "b")),
+        im.call("f")(im.tuple_get(1, "a"), im.tuple_get(1, "b")),
+    )
+    assert result == expected
+
+
+def test_nested():
+    outer = ts.TupleType(types=[TT, T])
+    result = _unroll(im.ref("f"), [outer], [im.ref("t")])
+    expected = im.make_tuple(
+        im.make_tuple(
+            im.call("f")(im.tuple_get(0, im.tuple_get(0, "t"))),
+            im.call("f")(im.tuple_get(1, im.tuple_get(0, "t"))),
+        ),
+        im.call("f")(im.tuple_get(1, "t")),
+    )
+    assert result == expected