Repository: GridTools/gt4py
Pull request #2018 (2025-05-07 07:38:21)

fix[next]: Fix cse extraction of lambda functions
Fix a small bug introduced in #2013, where `f` in the following expression
```
let val = 1 in
  let f = λ() → val in
     opaque_fun(f, g)
  end
end
```
was erroneously extracted beyond the definition of `val`.


Patch:
diff --git a/src/gt4py/next/iterator/transforms/cse.py b/src/gt4py/next/iterator/transforms/cse.py
index 2611c721..14fcbead 100644
--- a/src/gt4py/next/iterator/transforms/cse.py
+++ b/src/gt4py/next/iterator/transforms/cse.py
@@ -220,6 +220,9 @@ class CollectSubexpressions(PreserveLocationVisitor, VisitorWithSymbolTableTrait
                     ),
                     **kwargs,
                 )
+        else:
+            # we don't want to collect any expression, but still need to track all symbol refs
+            super().visit(node, state=self.State({}, set(), 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/program_processors/runners/dace/transformations/gpu_utils.py b/src/gt4py/next/program_processors/runners/dace/transformations/gpu_utils.py
index 003da07e..d121e460 100644
--- a/src/gt4py/next/program_processors/runners/dace/transformations/gpu_utils.py
+++ b/src/gt4py/next/program_processors/runners/dace/transformations/gpu_utils.py
@@ -258,7 +258,9 @@ def _gt_expand_non_standard_memlets_sdfg(
 ) -> set[dace_nodes.MapEntry]:
     """Implementation of `_gt_expand_non_standard_memlets()` that process a single SDFG."""
     new_maps: set[dace_nodes.MapEntry] = set()
-    # The implementation is based on DaCe's code generator.
+    # The implementation is based on DaCe's code generator, see `dace/codegen/targets/cuda.py`
+    #  in the function `preprocess()`
+    # NOTE: This implementation needs a DaCe version that includes https://github.com/spcl/dace/pull/1976
     for state in sdfg.states():
         for e in state.edges():
             # We are only interested in edges that connects two access nodes of GPU memory.
@@ -279,16 +281,9 @@ def _gt_expand_non_standard_memlets_sdfg(
             if dims == 1:
                 continue
             elif dims == 2:
-                if src_strides[-1] != 1 or dst_strides[-1] != 1:
-                    try:
-                        is_src_cont = src_strides[0] / src_strides[1] == copy_shape[1]
-                        is_dst_cont = dst_strides[0] / dst_strides[1] == copy_shape[1]
-                    except (TypeError, ValueError):
-                        is_src_cont = False
-                        is_dst_cont = False
-                    if is_src_cont and is_dst_cont:
-                        continue
-                else:
+                is_fortran_order = src_strides[0] == 1 and dst_strides[0] == 1
+                is_c_order = src_strides[-1] == 1 and dst_strides[-1] == 1
+                if is_c_order or is_fortran_order:
                     continue
             elif dims > 2:
                 if not (src_strides[-1] != 1 or dst_strides[-1] != 1):
diff --git a/src/gt4py/next/program_processors/runners/dace/transformations/move_dataflow_into_if_body.py b/src/gt4py/next/program_processors/runners/dace/transformations/move_dataflow_into_if_body.py
index 46f785bb..0a70b8d2 100644
--- a/src/gt4py/next/program_processors/runners/dace/transformations/move_dataflow_into_if_body.py
+++ b/src/gt4py/next/program_processors/runners/dace/transformations/move_dataflow_into_if_body.py
@@ -17,7 +17,7 @@ from dace import (
     subsets as dace_sbs,
     transformation as dace_transformation,
 )
-from dace.sdfg import nodes as dace_nodes, propagation as dace_propagation
+from dace.sdfg import nodes as dace_nodes, propagation as dace_propagation, utils as dace_sutils
 
 from gt4py.next.program_processors.runners.dace import transformations as gtx_transformations
 
@@ -212,17 +212,13 @@ class MoveDataflowIntoIfBody(dace_transformation.SingleStateTransformation):
                 connector=conn_name,
             )
 
-        # Now fix the symbol mapping.
         self._update_symbol_mapping(if_block)
 
-        # Now remove the nodes that have been relocated from the SDFG and also
-        #  clean up the registry.
-        for nodes_to_move in relocatable_dataflow.values():
-            for node_to_remove in nodes_to_move:
-                if isinstance(node_to_remove, dace_nodes.AccessNode):
-                    assert node_to_remove.desc(sdfg).transient
-                    sdfg.remove_data(node_to_remove.data, validate=False)
-                graph.remove_node(node_to_remove)
+        self._remove_outside_dataflow(
+            sdfg=sdfg,
+            state=graph,
+            relocatable_dataflow=relocatable_dataflow,
+        )
 
         # Because we relocate some node it seems that DaCe gets a bit confused.
         #  So we have to reset the list. Without it the `test_if_mover_chain`
@@ -352,7 +348,13 @@ class MoveDataflowIntoIfBody(dace_transformation.SingleStateTransformation):
                 else:
                     # The data is defined somewhere in the Map scope itself.
                     outer_data = iedge.src
-                assert isinstance(outer_data, dace_nodes.AccessNode)
+                # TODO(phimuell): It is possible that this does not lead to an
+                #   AccessNode on the outside, but to something inside the Map scope
+                #   such as the MapExit of an inner map. To handle such a case we need
+                #   to construct the set of nodes to move differently, i.e.
+                #   considering this case already there.
+                if not isinstance(outer_data, dace_nodes.AccessNode):
+                    raise NotImplementedError()
                 assert not gtx_transformations.utils.is_view(outer_data, sdfg)
 
                 # If the data is not yet available in the inner SDFG made
@@ -374,6 +376,13 @@ class MoveDataflowIntoIfBody(dace_transformation.SingleStateTransformation):
                         ),
                     )
                     if_block.add_in_connector(outer_data.data)
+                else:
+                    # This is the case that we found a node, that refers to data that
+                    #  was already patched into the `if_block`. We would have to remove
+                    #  this, but since this function just replicates the dataflow,
+                    #  it will not do that. Instead we postpone this to the cleanup
+                    #  phase, see `_remove_outside_dataflow()`.
+                    pass
 
                 if outer_data not in new_nodes:
                     assert all(
@@ -398,6 +407,36 @@ class MoveDataflowIntoIfBody(dace_transformation.SingleStateTransformation):
         inner_sdfg.arrays[connector].transient = True
         if_block.remove_in_connector(connector)
 
+    def _remove_outside_dataflow(
+        self,
+        sdfg: dace.SDFG,
+        state: dace.SDFGState,
+        relocatable_dataflow: dict[str, set[dace_nodes.Node]],
+    ) -> None:
+        """Removes the original dataflow, that has been relocated.
+
+        The function will also remove data containers that are no longer in use.
+        """
+        all_relocatable_dataflow: set[dace_nodes.Node] = functools.reduce(
+            lambda s1, s2: s1.union(s2), relocatable_dataflow.values(), set()
+        )
+
+        # Before we can clean the original nodes, we must clean the dataflow. If a
+        #  node, that was relocated, has incoming connections we must remove them
+        #  and the parent dataflow.
+        for node_to_remove in all_relocatable_dataflow:
+            for iedge in list(state.in_edges(node_to_remove)):
+                if iedge.src in all_relocatable_dataflow:
+                    continue
+                dace_sutils.remove_edge_and_dangling_path(state, iedge)
+
+            if isinstance(node_to_remove, dace_nodes.AccessNode):
+                assert node_to_remove.desc(sdfg).transient
+                sdfg.remove_data(node_to_remove.data, validate=False)
+
+        # Remove the original nodes (data descriptors were deleted in the loop above).
+        state.remove_nodes_from(all_relocatable_dataflow)
+
     def _update_symbol_mapping(
         self,
         if_block: dace_nodes.NestedSDFG,
diff --git a/uv.lock b/uv.lock
index f75e6c70..16d3d3c9 100644
--- a/uv.lock
+++ b/uv.lock
@@ -663,7 +663,7 @@ wheels = [
 [[package]]
 name = "dace"
 version = "1.0.0"
-source = { git = "https://github.com/GridTools/dace?branch=gt4py-next-integration#8b9b674abfcbd0148a32a87fa6d88e240030435d" }
+source = { git = "https://github.com/GridTools/dace?branch=gt4py-next-integration#268fc18b99b4e76bdd4f9e1bdc289086fe8c808b" }
 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/GridTools/dace?branch=gt4py-next-integration#8b9b674abfcbd0148a32a87fa6d88e240030435d" } },
+    { name = "dace", version = "1.0.0", source = { git = "https://github.com/GridTools/dace?branch=gt4py-next-integration#268fc18b99b4e76bdd4f9e1bdc289086fe8c808b" } },
 ]
 formatting = [
     { name = "clang-format" },


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 1a10478a..88ae203f 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
@@ -25,6 +25,22 @@ def offset_provider_type(request):
     return {"I": common.Dimension("I", kind=common.DimensionKind.HORIZONTAL)}
 
 
+@pytest.fixture
+def opaque_fun(request):
+    return im.ref(
+        "opaque_fun",
+        ts.FunctionType(
+            pos_only_args=[],
+            pos_or_kw_args={
+                "f": ts.DeferredType(constraint=None),
+                "g": ts.DeferredType(constraint=None),
+            },
+            kw_only_args={},
+            returns=ts.DeferredType(constraint=None),
+        ),
+    )
+
+
 def test_trivial():
     common = im.plus("x", "y")
     testee = im.plus(common, common)
@@ -321,3 +337,21 @@ def test_extraction_from_let_form():
 
     actual = CSE.apply(testee, within_stencil=True)
     assert actual == expected
+
+
+def test_sym_ref_collection_from_lambda(opaque_fun):
+    # This is more a regression than a unit test. When `f` & `g` are collected, we don't want to
+    # collect subexpression from their body (as the pass does not recognize them to be evaluated),
+    # but still need to respect that `val` is used inside. Hence, when extracting `f` and `g,` the
+    # extracted lambda needs to be placed after the definition of `val`.
+    testee = im.let("val", 1)(
+        im.let(("f", im.lambda_()("val")), ("g", im.lambda_()("val")))(
+            im.call(opaque_fun)("f", "g")
+        )
+    )
+    expected = im.let("val", 1)(
+        im.let("_cs_1", im.lambda_()("val"))(im.call(opaque_fun)("_cs_1", "_cs_1"))
+    )
+
+    actual = CSE.apply(testee, within_stencil=True)
+    assert actual == expected  # no extraction should happen
diff --git a/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/transformation_tests/test_move_dataflow_into_if_body.py b/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/transformation_tests/test_move_dataflow_into_if_body.py
index 6f3871f1..971772c8 100644
--- a/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/transformation_tests/test_move_dataflow_into_if_body.py
+++ b/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/transformation_tests/test_move_dataflow_into_if_body.py
@@ -16,6 +16,7 @@ import gc
 
 dace = pytest.importorskip("dace")
 from dace.sdfg import nodes as dace_nodes, propagation as dace_propagation
+from dace import data as dace_data
 from dace.transformation import dataflow as dace_dataflow
 
 from gt4py.next.program_processors.runners.dace import (
@@ -535,12 +536,7 @@ def test_if_mover_dependent_branch_2():
     state.add_edge(tasklet_a1, "__out", a1, None, dace.Memlet("a1[0]"))
 
     # Computing of `b1`:
-    tasklet_b1 = state.add_tasklet(
-        "tasklet_b1", inputs={"__in"}, outputs={"__out"}, code="__out = math.exp(__in)"
-    )
-
-    state.add_edge(me, "OUT_b", tasklet_b1, "__in", dace.Memlet("b[__i]"))
-    state.add_edge(tasklet_b1, "__out", b1, None, dace.Memlet("b1[0]"))
+    state.add_edge(me, "OUT_b", b1, None, dace.Memlet("b[__i] -> [0]"))
 
     # Now the condition.
     tasklet_cond = state.add_tasklet(
@@ -597,6 +593,99 @@ def test_if_mover_dependent_branch_2():
     assert if_block.sdfg.arrays.keys() == expected_data
 
 
+def test_if_mover_dependent_branch_3():
+    """
+    Very similar test to `test_if_mover_dependent_branch_1()`, but the common data
+    is an AccessNode outside the Map.
+    """
+    sdfg = dace.SDFG(util.unique_name("if_mover_dependent_branches"))
+    state = sdfg.add_state(is_start_block=True)
+
+    gnames = ["a", "b", "c", "d", "cond"]
+    for name in gnames:
+        sdfg.add_array(
+            name,
+            shape=(10,),
+            dtype=(dace.bool_ if name == "cond" else dace.float64),
+            transient=False,
+        )
+    a, b, c, d, cond = (state.add_access(name) for name in gnames)
+
+    temp_names = ["t1", "t2"]
+    for name in temp_names:
+        sdfg.add_scalar(
+            name,
+            dtype=dace.float64,
+            transient=True,
+        )
+    t1, t2 = (state.add_access(name) for name in temp_names)
+
+    me, mx = state.add_map("map", ndrange={"__i": "0:10"})
+    if_block = _make_if_block(state=state, outer_sdfg=sdfg)
+    tlet1 = state.add_tasklet(
+        "tlet1", inputs={"__in1", "__in2"}, outputs={"__out"}, code="__out = __in1 + __in2"
+    )
+
+    # Note we could also use a Tasklet here, but using a Map makes handling a bit more
+    #  difficult, because it is not a proper end of a Memlet path.
+    ime2, imx2 = state.add_map("think_of_me_as_a_tasklet", ndrange={"__UNUSED": "0"})
+    tlet2 = state.add_tasklet(
+        "tlet2", inputs={"__in1", "__in2"}, outputs={"__out"}, code="__out = __in1 + __in2"
+    )
+
+    state.add_edge(a, None, me, "IN_a", dace.Memlet("a[0:10]"))
+    state.add_edge(b, None, me, "IN_b1", dace.Memlet("b[0:10]"))
+    state.add_edge(me, "OUT_a", tlet1, "__in1", dace.Memlet("a[__i]"))
+    state.add_edge(me, "OUT_b1", tlet1, "__in2", dace.Memlet("b[__i]"))
+    state.add_edge(tlet1, "__out", t1, None, dace.Memlet("t1[0]"))
+    me.add_scope_connectors("a")
+    me.add_scope_connectors("b1")
+
+    state.add_edge(c, None, me, "IN_c", dace.Memlet("c[0:10]"))
+    state.add_edge(b, None, me, "IN_b2", dace.Memlet("b[0:10]"))
+    state.add_edge(me, "OUT_b2", ime2, "IN_b", dace.Memlet("b[__i]"))
+    state.add_edge(ime2, "OUT_b", tlet2, "__in1", dace.Memlet("b[__i]"))
+    state.add_edge(me, "OUT_c", ime2, "IN_c", dace.Memlet("c[__i]"))
+    state.add_edge(ime2, "OUT_c", tlet2, "__in2", dace.Memlet("c[__i]"))
+    state.add_edge(tlet2, "__out", imx2, "IN_t2", dace.Memlet("t2[0]"))
+    state.add_edge(imx2, "OUT_t2", t2, None, dace.Memlet("t2[0]"))
+    me.add_scope_connectors("b2")
+    me.add_scope_connectors("c")
+    ime2.add_scope_connectors("b")
+    ime2.add_scope_connectors("c")
+    imx2.add_scope_connectors("t2")
+
+    state.add_edge(cond, None, me, "IN_cond", dace.Memlet("cond[0:10]"))
+    state.add_edge(me, "OUT_cond", if_block, "__cond", dace.Memlet("cond[__i]"))
+    me.add_scope_connectors("cond")
+
+    state.add_edge(t1, None, if_block, "__arg1", dace.Memlet("t1[0]"))
+    state.add_edge(t2, None, if_block, "__arg2", dace.Memlet("t2[0]"))
+
+    state.add_edge(if_block, "__output", mx, "IN_d", dace.Memlet("d[__i]"))
+    state.add_edge(mx, "OUT_d", d, None, dace.Memlet("d[0:10]"))
+    mx.add_scope_connectors("d")
+
+    sdfg.validate()
+
+    assert state.out_degree(b) == 2
+    assert len(me.in_connectors) == 5
+    assert util.count_nodes(state, dace_nodes.Tasklet) == 2
+    assert util.count_nodes(state, dace_nodes.MapEntry) == 2
+    assert util.count_nodes(state, dace_nodes.AccessNode) == 7
+
+    _perform_test(sdfg, explected_applies=1)
+
+    # It is unspecific if `IN_b1` or `IN_b2` remains, but `b` should have only one connector.
+    assert any(iconn in me.in_connectors for iconn in ["IN_b1", "IN_b2"])
+    assert state.out_degree(b) == 1
+    assert len(me.in_connectors) == 4
+    assert util.count_nodes(state, dace_nodes.Tasklet) == 0
+    assert util.count_nodes(state, dace_nodes.MapEntry) == 1
+    assert util.count_nodes(state, dace_nodes.AccessNode) == 5
+    assert set(gnames) == sdfg.arrays.keys()
+
+
 def test_if_mover_no_ops():
     """
     Essentially tests the following situation: