Repository: GridTools/gt4py
Pull request #2013 (2025-04-30 10:35:18)

fix[next]: Prevent extraction from lambdas in cse pass
The common subexpression elimination pass already has a mechanism to avoid extracting from `if_` calls, such that the evaluation of expressions can be guarded in order to to avoid out-of-bounds accesses or excessive pre-computation. The former protection can be circumvented by placing things inside a lambda outside of the protecting if:
```
if cond then
  let f = λ() → ·guarded_it in
    if cond2 then f() else 0
else
  ·guarded_it
```
This PR prevents that by not extracting from lambdas unless they are a let-form in which case it is clear that the function is unconditionally evaluated.


Patch:
diff --git a/pyproject.toml b/pyproject.toml
index b5925be4..078bb5a6 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -434,7 +434,7 @@ url = 'https://test.pypi.org/simple/'
 
 [tool.uv.sources]
 atlas4py = {index = "test.pypi"}
-dace = {git = "https://github.com/spcl/dace", branch = "main", extra = "dace-next"}
+dace = {git = "https://github.com/GridTools/dace", branch = "gt4py-next-integration", extra = "dace-next"}
 
 # -- versioningit --
 [tool.versioningit]
diff --git a/src/gt4py/eve/visitors.py b/src/gt4py/eve/visitors.py
index 6e559a26..31f9f59b 100644
--- a/src/gt4py/eve/visitors.py
+++ b/src/gt4py/eve/visitors.py
@@ -121,20 +121,6 @@ class NodeVisitor:
         return None
 
 
-def _preserve_annex(
-    node: concepts.Node, new_node: concepts.Node, preserved_annex_attrs: tuple[str, ...]
-) -> None:
-    # access to `new_node.annex` implicitly creates the `__node_annex__` attribute in the property getter
-    new_annex_dict = new_node.annex.__dict__
-    old_annex_dict = node.annex.__dict__
-    for key in (old_annex_dict.keys() & preserved_annex_attrs) - new_annex_dict.keys():
-        # Note: The annex item's value of the new node might not be equal
-        # (in the sense that an equality comparison returns false), to
-        # the old one, but in the context of the pass, they are
-        # equivalent. Therefore, we don't assert equality here.
-        new_annex_dict[key] = old_annex_dict[key]
-
-
 class NodeTranslator(NodeVisitor):
     """Special `NodeVisitor` to translate nodes and trees.
 
@@ -173,7 +159,7 @@ class NodeTranslator(NodeVisitor):
                 }
             )
             if self.PRESERVED_ANNEX_ATTRS and getattr(node, "__node_annex__", None):
-                _preserve_annex(node, new_node, self.PRESERVED_ANNEX_ATTRS)
+                self._preserve_annex(node, new_node)
 
             return new_node
 
@@ -208,6 +194,17 @@ class NodeTranslator(NodeVisitor):
             and self.PRESERVED_ANNEX_ATTRS
             and getattr(node, "__node_annex__", None)
         ):
-            _preserve_annex(node, new_node, self.PRESERVED_ANNEX_ATTRS)
+            self._preserve_annex(node, new_node)
 
         return new_node
+
+    def _preserve_annex(self, node: concepts.Node, new_node: concepts.Node) -> None:
+        # access to `new_node.annex` implicitly creates the `__node_annex__` attribute in the property getter
+        new_annex_dict = new_node.annex.__dict__
+        old_annex_dict = node.annex.__dict__
+        for key in (old_annex_dict.keys() & self.PRESERVED_ANNEX_ATTRS) - new_annex_dict.keys():
+            # Note: The annex item's value of the new node might not be equal
+            # (in the sense that an equality comparison returns false), to
+            # the old one, but in the context of the pass, they are
+            # equivalent. Therefore, we don't assert equality here.
+            new_annex_dict[key] = old_annex_dict[key]
diff --git a/src/gt4py/next/iterator/transforms/collapse_tuple.py b/src/gt4py/next/iterator/transforms/collapse_tuple.py
index 55644c7d..a29870fb 100644
--- a/src/gt4py/next/iterator/transforms/collapse_tuple.py
+++ b/src/gt4py/next/iterator/transforms/collapse_tuple.py
@@ -325,9 +325,10 @@ class CollapseTuple(
             new_args: list[itir.Expr] = []
             for arg in node.args:
                 if cpm.is_call_to(node, "make_tuple") and not _is_trivial_make_tuple_call(node):
-                    el_name = self.uids.sequential_id(prefix="__ct_el")
-                    new_args.append(im.ref(el_name, arg.type))
-                    bound_vars[im.sym(el_name, arg.type)] = arg
+                    new_arg = im.ref(self.uids.sequential_id(prefix="__ct_el"), arg.type)
+                    self._preserve_annex(arg, new_arg)
+                    new_args.append(new_arg)
+                    bound_vars[im.sym(new_arg.id, arg.type)] = arg
                 else:
                     new_args.append(arg)
 
@@ -547,7 +548,11 @@ class CollapseTuple(
         for param, arg in zip(stencil.params, node.args, strict=True):
             if isinstance(arg.type, ts.TupleType):
                 assert isinstance(param.type, it_ts.IteratorType)
-                ref_to_remapped_arg = im.ref(f"__ct_flat_remapped_{len(remapped_args)}", arg.type)
+                ref_to_remapped_arg = im.ref(
+                    f"__ct_flat_remapped_{len(remapped_args)}",
+                    arg.type,
+                )
+                self._preserve_annex(arg, ref_to_remapped_arg)
                 remapped_args[im.sym(ref_to_remapped_arg.id, arg.type)] = arg
                 new_params_inner, lift_params = [], []
                 assert isinstance(param.type.element_type, ts.TupleType)
diff --git a/src/gt4py/next/iterator/transforms/cse.py b/src/gt4py/next/iterator/transforms/cse.py
index cc1ffc3c..2611c721 100644
--- a/src/gt4py/next/iterator/transforms/cse.py
+++ b/src/gt4py/next/iterator/transforms/cse.py
@@ -8,10 +8,9 @@
 
 from __future__ import annotations
 
+import collections
 import dataclasses
-import functools
 import math
-import operator
 from typing import Callable, Iterable, TypeVar, Union, cast
 
 import gt4py.next.iterator.ir_utils.ir_makers as im
@@ -149,8 +148,15 @@ class CollectSubexpressions(PreserveLocationVisitor, VisitorWithSymbolTableTrait
         depth = kwargs.pop("depth")
         return super().generic_visit(node, depth=depth + 1, **kwargs)
 
-    def visit(self, node: itir.Node, **kwargs) -> None:  # type: ignore[override] # supertype accepts any node, but we want to be more specific here.
-        if not isinstance(node, SymbolTableTrait) and not _is_collectable_expr(node):
+    def visit(self, node: itir.Node, is_let_form: bool = False, **kwargs) -> None:  # type: ignore[override] # supertype accepts any node, but we want to be more specific here.
+        # do not descent into un-applied lambda
+        can_collect_children = not isinstance(node, itir.Lambda) or is_let_form
+
+        if (
+            can_collect_children
+            and not isinstance(node, SymbolTableTrait)
+            and not _is_collectable_expr(node)
+        ):
             return super().visit(node, **kwargs)
 
         depth = kwargs["depth"]
@@ -158,48 +164,62 @@ class CollectSubexpressions(PreserveLocationVisitor, VisitorWithSymbolTableTrait
         collected_child_node_ids: set[int] = set()
         used_symbol_ids: set[int] = set()
 
-        # Special handling of `if_(condition, true_branch, false_branch)` like expressions that
-        # avoids extracting subexpressions unless they are used in either the condition or both
-        # branches.
-        if isinstance(node, itir.FunCall) and node.fun == itir.SymRef(id="if_"):
-            assert len(node.args) == 3
-            # collect subexpressions for all arguments to the `if_`
-            arg_states = [self.State() for _ in node.args]
-            for arg, state in zip(node.args, arg_states):
-                self.visit(arg, state=state, **{**kwargs, "depth": depth + 1})
-
-            # remove all subexpressions that are not eligible for collection
-            #  (either they occur in the condition or in both branches)
-            eligible_subexprs = arg_states[0].subexprs.keys() | (
-                arg_states[1].subexprs.keys() & arg_states[2].subexprs.keys()
-            )
-            for arg_state in arg_states:
-                arg_state.remove_subexprs(arg_state.subexprs.keys() - eligible_subexprs)
-
-            # merge the states of the three arguments
-            subexprs: dict[itir.Node, CollectSubexpressions.SubexpressionData] = {}
-            for state in arg_states:
-                for subexpr, data in state.subexprs.items():
-                    merged_data = subexprs.setdefault(subexpr, self.SubexpressionData())
-                    merged_data.subexprs.extend(data.subexprs)
-                    merged_data.max_depth = max(merged_data.max_depth, data.max_depth)
-            collected_child_node_ids = functools.reduce(
-                operator.or_, (state.collected_child_node_ids for state in arg_states)
-            )
-            used_symbol_ids = functools.reduce(
-                operator.or_, (state.used_symbol_ids for state in arg_states)
-            )
-            # propagate collected subexpressions to parent
-            for subexpr, data in subexprs.items():
-                parent_data = parent_state.subexprs.setdefault(subexpr, self.SubexpressionData())
-                parent_data.subexprs.extend(data.subexprs)
-                parent_data.max_depth = max(merged_data.max_depth, data.max_depth)
-        else:
-            super().visit(
-                node,
-                state=self.State(parent_state.subexprs, collected_child_node_ids, used_symbol_ids),
-                **kwargs,
-            )
+        if can_collect_children:
+            # Special handling of `if_(condition, true_branch, false_branch)` like expressions that
+            # avoids extracting subexpressions unless they are used in either the condition or both
+            # branches.
+            if cpm.is_call_to(node, "if_"):
+                assert len(node.args) == 3
+                # collect subexpressions for all arguments to the `if_`
+                arg_states = [self.State() for _ in node.args]
+                for arg, state in zip(node.args, arg_states):
+                    self.visit(arg, state=state, **(kwargs | {"depth": depth + 1}))
+
+                # remove all subexpressions that are not eligible for collection
+                #  (either they occur in the condition or in both branches)
+                eligible_subexprs = arg_states[0].subexprs.keys() | (
+                    arg_states[1].subexprs.keys() & arg_states[2].subexprs.keys()
+                )
+                for arg_state in arg_states:
+                    arg_state.remove_subexprs(arg_state.subexprs.keys() - eligible_subexprs)
+
+                # merge the states of the three arguments
+                subexprs: dict[itir.Node, CollectSubexpressions.SubexpressionData] = (
+                    collections.defaultdict(self.SubexpressionData)
+                )
+                for state in arg_states:
+                    for subexpr, data in state.subexprs.items():
+                        merged_data = subexprs[subexpr]
+                        merged_data.subexprs.extend(data.subexprs)
+                        merged_data.max_depth = max(merged_data.max_depth, data.max_depth)
+                collected_child_node_ids = set.union(
+                    *(state.collected_child_node_ids for state in arg_states)
+                )
+                used_symbol_ids = set.union(*(state.used_symbol_ids for state in arg_states))
+                # propagate collected subexpressions to parent
+                for subexpr, data in subexprs.items():
+                    parent_data = parent_state.subexprs.setdefault(
+                        subexpr, self.SubexpressionData()
+                    )
+                    parent_data.subexprs.extend(data.subexprs)
+                    parent_data.max_depth = max(merged_data.max_depth, data.max_depth)
+            elif cpm.is_let(node):
+                new_state = self.State(
+                    parent_state.subexprs, collected_child_node_ids, used_symbol_ids
+                )
+                self.visit(
+                    node.fun, state=new_state, is_let_form=True, **(kwargs | {"depth": depth + 1})
+                )
+                for arg in node.args:
+                    self.visit(arg, state=new_state, **(kwargs | {"depth": depth + 1}))
+            else:
+                super().visit(
+                    node,
+                    state=self.State(
+                        parent_state.subexprs, collected_child_node_ids, used_symbol_ids
+                    ),
+                    **kwargs,
+                )
 
         if isinstance(node, SymbolTableTrait):
             # remove symbols used in child nodes if they are declared in the current node
diff --git a/src/gt4py/next/iterator/transforms/fixed_point_transformation.py b/src/gt4py/next/iterator/transforms/fixed_point_transformation.py
index f1176b4b..53721534 100644
--- a/src/gt4py/next/iterator/transforms/fixed_point_transformation.py
+++ b/src/gt4py/next/iterator/transforms/fixed_point_transformation.py
@@ -63,5 +63,6 @@ class FixedPointTransformation(eve.NodeTranslator):
                         result is not node
                     ), f"Transformation {transformation.name.lower()} should have returned None, since nothing changed."
                     itir_type_inference.reinfer(result)
+                    self._preserve_annex(node, result)
                     return result
         return None
diff --git a/uv.lock b/uv.lock
index 3748f9b0..f75e6c70 100644
--- a/uv.lock
+++ b/uv.lock
@@ -663,7 +663,7 @@ wheels = [
 [[package]]
 name = "dace"
 version = "1.0.0"
-source = { git = "https://github.com/spcl/dace?branch=main#9c67b4c6fe02a3248fcc25c6686715988e9db015" }
+source = { git = "https://github.com/GridTools/dace?branch=gt4py-next-integration#8b9b674abfcbd0148a32a87fa6d88e240030435d" }
 resolution-markers = [
     "python_full_version >= '3.11'",
     "python_full_version < '3.11'",
@@ -1081,7 +1081,7 @@ dace = [
     { name = "dace", version = "1.0.2", source = { registry = "https://pypi.org/simple" } },
 ]
 dace-next = [
-    { name = "dace", version = "1.0.0", source = { git = "https://github.com/spcl/dace?branch=main#9c67b4c6fe02a3248fcc25c6686715988e9db015" } },
+    { name = "dace", version = "1.0.0", source = { git = "https://github.com/GridTools/dace?branch=gt4py-next-integration#8b9b674abfcbd0148a32a87fa6d88e240030435d" } },
 ]
 formatting = [
     { name = "clang-format" },
@@ -1204,7 +1204,7 @@ requires-dist = [
     { name = "cupy-rocm-5-0", marker = "extra == 'rocm5-0'", specifier = ">=13.3.0" },
     { name = "cytoolz", specifier = ">=0.12.1" },
     { name = "dace", marker = "extra == 'dace'", specifier = ">=1.0.2,<1.1.0" },
-    { name = "dace", marker = "extra == 'dace-next'", git = "https://github.com/spcl/dace?branch=main" },
+    { name = "dace", marker = "extra == 'dace-next'", git = "https://github.com/GridTools/dace?branch=gt4py-next-integration" },
     { name = "deepdiff", specifier = ">=5.6.0" },
     { name = "devtools", specifier = ">=0.6" },
     { name = "diskcache", specifier = ">=5.6.3" },


Test patch:
diff --git a/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_cse.py b/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_cse.py
index 3909c6f2..1a10478a 100644
--- a/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_cse.py
+++ b/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_cse.py
@@ -106,6 +106,12 @@ def test_lambda_redef_same_arg():
 
 
 def test_lambda_redef_same_arg_scope():
+    pytest.xfail(
+        reason="Not implemented. The CSE pass may not extract from a lambda function"
+        "unless it knows the function is unconditionally evaluated. For this"
+        "case this currently beyond scope."
+    )
+
     def common_expr():
         return im.lambda_("a")(im.plus("a", im.plus(1, 1)))
 
@@ -293,3 +299,25 @@ def test_scalar_extraction_inside_as_fieldop():
 
     actual = CSE.apply(testee, within_stencil=False)
     assert actual == expected
+
+
+def test_no_extraction_from_unapplied_lambda():
+    testee = im.if_(
+        "cond",
+        im.let("f", im.lambda_()(im.deref("guarded_it")))(im.if_("cond2", im.call("f")(), 0)),
+        im.deref("guarded_it"),
+    )
+
+    actual = CSE.apply(testee, within_stencil=True)
+    assert actual == testee  # no extraction should happen
+
+
+def test_extraction_from_let_form():
+    common = im.plus(1, 2)
+    testee = im.plus(im.let("a", 1)(im.plus("a", common)), im.let("b", 2)(im.plus("b", common)))
+    expected = im.let("_cs_1", common)(
+        im.plus(im.let("a", 1)(im.plus("a", "_cs_1")), im.let("b", 2)(im.plus("b", "_cs_1")))
+    )
+
+    actual = CSE.apply(testee, within_stencil=True)
+    assert actual == expected