feat[next]: Tracer support part 1: tree_map

[SF-N]

As a first step towards tracer support (enabling vector operations), this PR introduces a new tree_map operator for mapping functions over tuples (including nesting). tree_map is unrolled to make_tuple calls in UnrollTreeMap.

[Copilot]

Pull request overview

This PR introduces an iterator-level tree_map builtin as an IR operator to support mapping functions over (nested) tuples, as a first step towards tracer support and future vector operations. It includes type synthesis for tree_map and a transform (UnrollTreeMap) that lowers tree_map(f)(...) into explicit make_tuple / tuple_get IR.

Changes:

• Add tree_map builtin plumbing (builtin dispatch + IR maker helper) and update tuple-where lowering to emit tree_map.
• Add tree_map type synthesizer and a new UnrollTreeMap transform, wired into the iterator pass pipeline.
• Add unit tests for the _unroll helper and adjust existing frontend lowering expectations.

Reviewed changes

Copilot reviewed 8 out of 8 changed files in this pull request and generated 6 comments.

Show a summary per file

File	Description
tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_unroll_tree_map.py	Adds unit tests for _unroll tuple expansion behavior.
tests/next_tests/unit_tests/ffront_tests/test_foast_to_gtir.py	Updates tuple where reference IR to use tree_map.
src/gt4py/next/iterator/type_system/type_synthesizer.py	Registers and implements type synthesis for tree_map.
src/gt4py/next/iterator/transforms/unroll_tree_map.py	New transform to unroll tree_map into tuple primitives.
src/gt4py/next/iterator/transforms/pass_manager.py	Runs UnrollTreeMap and tuple-collapsing before domain inference.
src/gt4py/next/iterator/ir_utils/ir_makers.py	Adds im.tree_map(...) helper for constructing IR.
src/gt4py/next/iterator/builtins.py	Adds tree_map to builtin dispatch and builtin name set.
src/gt4py/next/ffront/foast_to_gtir.py	Lowers tuple where via tree_map instead of explicit tuple construction.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[SF-N]

Without this test_reduction_expression_with_where_and_tuples fails with ValueError: 'target_domain' cannot be 'NEVER' unless "allow_uninferred=True".

[tehrengruber]

Note: probably this is also the test case where the loop is required. I'll take a look if another configuration of the pass helps to avoid the loop.

[Copilot]

Pull request overview

Copilot reviewed 9 out of 9 changed files in this pull request and generated 5 comments.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[SF-N]

I think _visit_astype cannot use tree_map because it needs type-dependent lowering per leaf: fields use _map(cast) while scalars use cast directly. Let's discuss if you have something else in mind.

[SF-N]

As far as I see, _visit_concat_where already has its own expand_tuple_args pass for handling nested tuples, and each branch can have a different domain. Attempting to wrap it in tree_map caused type inference failures. Let's discuss if you have something else in mind.

[tehrengruber]

Other reasons: A single concat_where is better to digest for optimizations. The lowering would actually be complicated if we would emit tree_map in foast_to_gtir.

[tehrengruber]

Suggested change

	def _visit_concat_where(self, node: foast.Call, **kwargs: Any) -> itir.FunCall:
	# TODO(tehrengruber): Use `tree_map_tuple` when the domain inference is able to handle
	# lambda functions (with the results domain depending on the caller / args)

[Copilot]

Pull request overview

Copilot reviewed 21 out of 21 changed files in this pull request and generated 4 comments.

[tehrengruber]

Partial review only.

[tehrengruber]

Contrary to the tuple comprehension case where we can not type the el_expr in el_expr for ... in ... and hence have to disallow cases with mixed element types, this is not the case here. The process_elements approach works for the case above. Let's add a test (if not already done) and revert to the previous strategy. Additionally a comment that we would like to use tree_map_tuple and why this is not possible (or a reference) would be good. If you want to be thorough create an ADR that describes the strategy for all the cases and reference that. https://hackmd.io/PxqTkA4rSGCgCmUAPq15SA could serve as a starting point.

[tehrengruber]

Suggested change

	def _visit_concat_where(self, node: foast.Call, **kwargs: Any) -> itir.FunCall:
	# TODO(tehrengruber): Use `tree_map_tuple` when the domain inference is able to handle
	# lambda functions (with the results domain depending on the caller / args)

[tehrengruber]

Maybe we should make map_tuple support multiple args like tree_map_tuple.

Open question: Does using multiple args in the lowering improve what collapse tuple can do before the tree_map_tuple can do? Then let's use that path even for the ... for tuple in ... in the frontend.

[tehrengruber]

tuple[0]+tuple[1] for tuple in nested_tuple
or
let(tuple, {a, b})(tuple[0]+tuple[1]) for a, b in {nested_tuple[0], nested_tuple[1]}
Another dimension: a+b for a, b in nested_tuple
The interesting parts are what is the advantage / disadvantage in the lowering & unroll pass.

[SF-N]

All examples single-arg map_tuple(op)(n) over eg.g n : ((Field[[IDim], float64], Field[[IDim], float64]),(Field[[IDim], float64],Field[[IDim], float64])), and they collapse to the identical result:

1. t[0] + t[1] for t in n:
   lowered:
   map_tuple(λ(t) → t[0] + t[1])(n)
   before collapse:
   (λ(_utm_0) → {(λ(t) → t[0] + t[1])(_utm_0[0]),(λ(t) → t[0] + t[1])(_utm_0[1])})(n)
   after collapse:
   {n[0][0] + n[0][1], n[1][0] + n[1][1]}

2. let(tup, {a, b})(tup[0] + tup[1]) for a, b in {n[0], n[1]}:
   lowered:
   map_tuple(λ(t) → (λ(a) → (λ(b) → (λ(tup) → tup[0] + tup[1])({a, b}))(t[1]))(t[0]))({n[0], n[1]})
   before collapse:
   (λ(_utm_0) → {(λ(t) → (λ(a) → (λ(b) → (λ(tup) → tup[0] + tup[1])({a, b}))(t[1]))(t[0]))(_utm_0[0]), (λ(t) → (λ(a) → (λ(b) → (λ(tup) → tup[0] + tup[1])({a, b}))(t[1]))(t[0]))(_utm_0[1])})({n[0], n[1]})
   after collapse:
   {n[0][0] + n[0][1], n[1][0] + n[1][1]}

3. a + b for a, b in n:
   lowered:
   map_tuple(λ(t) → (λ(a) → (λ(b) → a + b)(t[1]))(t[0]))(n)
   before collapse:
   (λ(_utm_0) → {(λ(t) → (λ(a) → (λ(b) → a + b)(t[1]))(t[0]))(_utm_0[0]), (λ(t) → (λ(a) → (λ(b) → a + b)(t[1]))(t[0]))(_utm_0[1])})(n)
   after collapse:
   {n[0][0] + n[0][1], n[1][0] + n[1][1]}

All three collapse to the same result, so the choice only affects the intermediate size that the unroll+collapse passes must remove.

There is no use-case for multi-arg map_tuple so far and also not in #PR2487 , whereas for tree_map_tuple the multi-arg use-case is where(cond, a, b).

Nevertheless, I tried implementing multi-arg map_tuple and no changes in the lowering were necessary, only small changes in the type_synthesizer, ir_makers and in the unroll_tuple_maps replacing

def _map_tuple_body(
    f: itir.Expr, tup_exprs: list[itir.Expr], tup_types: list[ts.TupleType]
) -> itir.Expr:
    """Unroll `map_tuple(f)(t)` over top-level elements only (no recursion)."""
    (tup_expr,) = tup_exprs
    (tup_type,) = tup_types
    return im.make_tuple(
        *(im.call(f)(im.tuple_get(i, tup_expr)) for i in range(len(tup_type.types)))
    )

by

def _map_tuple_body(
    f: itir.Expr, tup_exprs: list[itir.Expr], tup_types: list[ts.TupleType]
) -> itir.Expr:
    """Unroll `map_tuple(f)(t1, ..., tN)` over top-level elements only (no recursion)."""
    return im.make_tuple(
        *(
            im.call(f)(*(im.tuple_get(i, tup_expr) for tup_expr in tup_exprs))
            for i in range(len(tup_types[0].types))
        )
    )

The only benefit would be in the case
map_tuple(λa,b → a+b)(g(), h()) vs map_tuple(λt → t[0]+t[1])(zip(g(),h())) where g()/h() would be evaluated 2× in the former and 4x in the latter. But usually a comprehension has a single source, so there is nothing independent to combine.

So, I reverted the changes again and would keep map_tuple single-arg.