Repository: GridTools/gt4py
Pull request #2146 (2025-07-16 06:13:45)

fix[dace][next]: Removed Special Case For `CopyChainRemover`.
`CopyChainRemover` removes patterns that are generated by `concat_where` expressions.
It checks if the source is not used anywhere else and if it is _fully_ copied into the destination.
Before there was an exception to that rule for cases when it was not possible to verify.
However, in `8_to_13_corrector` we encountered a case where this lead to an error.
In fact the error was probably present in more stencils, but it only surfaced there, the issue was that something was generated that was not used and then some memory location was written twice.
The correctness then dependent on the order in which the code was generated.

This PR removes this special rule. It will not impact SL2, but might impact SL1. 


Patch:
diff --git a/src/gt4py/next/iterator/transforms/concat_where/expand_tuple_args.py b/src/gt4py/next/iterator/transforms/concat_where/expand_tuple_args.py
index ea408697..40d956fc 100644
--- a/src/gt4py/next/iterator/transforms/concat_where/expand_tuple_args.py
+++ b/src/gt4py/next/iterator/transforms/concat_where/expand_tuple_args.py
@@ -5,16 +5,18 @@
 #
 # Please, refer to the LICENSE file in the root directory.
 # SPDX-License-Identifier: BSD-3-Clause
+from typing import Optional
 
-from gt4py.eve import NodeTranslator, PreserveLocationVisitor
+from gt4py.eve import PreserveLocationVisitor
 from gt4py.next import common
 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.transforms.fixed_point_transformation import FixedPointTransformation
 from gt4py.next.iterator.type_system import inference as type_inference
 from gt4py.next.type_system import type_specifications as ts
 
 
-class _ExpandTupleArgs(PreserveLocationVisitor, NodeTranslator):
+class _ExpandTupleArgs(PreserveLocationVisitor, FixedPointTransformation):
     PRESERVED_ANNEX_ATTRS = (
         "type",
         "domain",
@@ -35,28 +37,31 @@ class _ExpandTupleArgs(PreserveLocationVisitor, NodeTranslator):
         )
         return cls().visit(node)
 
-    def visit_FunCall(self, node: itir.FunCall) -> itir.FunCall:
-        node = self.generic_visit(node)
-
+    def transform(self, node: itir.Node, **kwargs) -> Optional[itir.Node]:
         # `concat_where(cond, {a, b}, {c, d})`
         # -> `{concat_where(cond, a, c), concat_where(cond, a, c)}`
-        if cpm.is_call_to(node, "concat_where") and isinstance(
+        if not cpm.is_call_to(node, "concat_where") or not isinstance(
             type_inference.reinfer(node.args[1]).type, ts.TupleType
         ):
-            cond, true_branch, false_branch = node.args
-            new_els = []
-            assert isinstance(true_branch.type, ts.TupleType)
-            for i in range(len(true_branch.type.types)):
-                new_els.append(
-                    im.concat_where(cond, im.tuple_get(i, "__tb"), im.tuple_get(i, "__fb"))
-                )
+            return None
 
-            new_node = im.let(("__tb", true_branch), ("__fb", false_branch))(
-                im.make_tuple(*new_els)
+        cond, true_branch, false_branch = node.args
+        new_els = []
+        assert isinstance(true_branch.type, ts.TupleType)
+        for i in range(len(true_branch.type.types)):
+            new_els.append(
+                self.fp_transform(
+                    im.concat_where(
+                        cond,
+                        im.tuple_get(i, im.ref("__tb", true_branch.type)),
+                        im.tuple_get(i, im.ref("__fb", false_branch.type)),
+                    ),
+                    **kwargs,
+                )
             )
-            return new_node
 
-        return node
+        new_node = im.let(("__tb", true_branch), ("__fb", false_branch))(im.make_tuple(*new_els))
+        return new_node
 
 
 expand_tuple_args = _ExpandTupleArgs.apply
diff --git a/src/gt4py/next/program_processors/runners/dace/transformations/redundant_array_removers.py b/src/gt4py/next/program_processors/runners/dace/transformations/redundant_array_removers.py
index 87947c9d..ebe43453 100644
--- a/src/gt4py/next/program_processors/runners/dace/transformations/redundant_array_removers.py
+++ b/src/gt4py/next/program_processors/runners/dace/transformations/redundant_array_removers.py
@@ -249,17 +249,11 @@ class CopyChainRemover(dace_transformation.SingleStateTransformation):
     - `A1` is a transient and must have the same dimensionality than `A2`.
     - `A1` is fully read by `A2`.
 
-    In certain cases the last rule can not be verified, an alternative formulation,
-    which is a consequence of the lowering and domain interference, will be checked:
-    - `A1` is read from the beginning, i.e. all subsets starts at literal `0`.
-    - `A1` has only one going edge.
-    - `A2` is global memory.
-
     Notes:
-        - The transformation assumes that the domain inference adjusted the ranges of
-            the maps such that, in case they write into a transient, the full shape of the transient array is written.
-            has the same size, i.e. there is not padding, or data that is not written
-            to.
+        The transformation assumes that the domain inference adjusted the ranges of
+        the maps such that, in case they write into a transient, the full shape of
+        the transient array is written, has the same size, i.e. there is not padding,
+        or data that is not written to.
 
     Args:
         single_use_data: List of data containers that are used only at one place.
@@ -362,24 +356,15 @@ class CopyChainRemover(dace_transformation.SingleStateTransformation):
             return False
 
         # Checking if the whole array is read.
-        #  As described in the description of the class there are two different
-        #  formulation of this rule, either the simple one or one that proves this
-        #  indirectly.
-        # NOTE: The main benefit of requiring that the whole array is read is
-        #  that we do not have to adjust maps.
+        # NOTE: The main benefit of requiring that the whole array is read is that we
+        #   do not have to adjust maps.
+        # NOTE: In previous versions there was an ad hoc rule, to bypass the "full
+        #   read rule". However, it caused problems, so it was removed.
+        # TODO: We have to improve this test, because sometimes the expressions are
+        #   so complex that without information about relations, such as
+        #   `vertical_start <= vertical_end` it is not possible to prove this check.
         a1_range = dace_sbs.Range.from_array(a1_desc)
-        if src_subset.covers(a1_range):
-            pass
-        elif all(ss.is_constant() for ss in src_subset.size()):
-            # If the subset is fully known then we require that it is covered.
-            return False
-        elif (
-            (not a2_desc.transient)
-            and (graph.out_degree(a1) == 1)
-            and all(ss_start == 0 for ss_start in src_subset.min_element())
-        ):
-            pass
-        else:
+        if not src_subset.covers(a1_range):
             return False
 
         # We have to ensure that no cycle is created through the removal of `a1`.


Test patch:
diff --git a/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_concat_where_expand_tuple_args.py b/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_concat_where_expand_tuple_args.py
index 42ad2920..8aa5df10 100644
--- a/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_concat_where_expand_tuple_args.py
+++ b/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_concat_where_expand_tuple_args.py
@@ -51,3 +51,43 @@ def test_trivial():
     )
 
     assert actual == expected
+
+
+def test_nested():
+    cond = im.domain(common.GridType.CARTESIAN, {IDim: (itir.InfinityLiteral.NEGATIVE, 1)})
+    domain = im.domain(common.GridType.CARTESIAN, {IDim: (0, 2)})
+    symbolic_domain = domain_utils.SymbolicDomain.from_expr(domain)
+
+    testee = im.concat_where(
+        cond,
+        im.make_tuple(
+            im.ref("a", field_type), im.make_tuple(im.ref("c", field_type), im.ref("e", field_type))
+        ),
+        im.make_tuple(
+            im.ref("b", field_type), im.make_tuple(im.ref("d", field_type), im.ref("f", field_type))
+        ),
+    )
+    testee, _ = infer_domain.infer_expr(
+        testee,
+        (symbolic_domain, symbolic_domain),
+        keep_existing_domains=True,
+        offset_provider={},
+    )
+
+    expected = im.make_tuple(
+        im.concat_where(cond, "a", "b"),
+        im.make_tuple(
+            im.concat_where(cond, "c", "d"),
+            im.concat_where(cond, "e", "f"),
+        ),
+    )
+
+    actual = concat_where.expand_tuple_args(
+        testee, offset_provider_type={}, allow_undeclared_symbols=True
+    )
+
+    actual = collapse_tuple.CollapseTuple.apply(
+        actual, allow_undeclared_symbols=True, within_stencil=False
+    )
+
+    assert actual == expected