edit

Author: edopao

src/gt4py/next/program_processors/runners/dace/lowering/gtir_dataflow.py

@@ -269,6 +269,7 @@ def connect(
         map_exit: Optional[dace_nodes.MapExit],
         dest: dace_nodes.AccessNode,
         dest_subset: dace_subsets.Range,
+        allow_removal_of_last_node: bool,
     ) -> bool:
         """Create a connection to the `dest` node, writing the given `dest_subset`.
 
@@ -277,35 +278,34 @@ def connect(
         outside data container is removed, the caller is responsible to propagate
         the strides of the destination array to the array inside the nested SDFG.
         """
-        # We don't allow removing the last node of the dataflow inside a map scope,
-        # because the same reults could be written to multiple destiation nodes ouside
-        # the map scope, in case of field operators returning tuples.
-        allow_removal_of_last_node: Final[bool] = False
         dest_desc = self.result.dc_node.desc(self.state)
         write_edge = self.state.in_edges(self.result.dc_node)[0]
 
-        if self.state.out_degree(self.result.dc_node) != 0:
-            remove_last_node = False
-        elif isinstance(write_edge.src, dace_nodes.Tasklet):
-            # The temporary data written by a tasklet can be safely deleted.
-            remove_last_node = True
-        elif isinstance(write_edge.src, dace_nodes.NestedSDFG):
-            if isinstance(dest_desc, dace.data.Scalar):
-                # We keep scalar temporary storage, as a general rule, since it
-                # does not affect performance of the generated code. This scalar
-                # is only required in some cases, e.g. for nested SDFGs implementing
-                # reduction, which use a WCR memlet for the reduction operation.
+        if allow_removal_of_last_node:
+            if self.state.out_degree(self.result.dc_node) != 0:
                 remove_last_node = False
-            else:
-                # We remove the transient array on the output connection of a nested
-                # SDFG and write directly to the destination node.
-                # The caller is responsible to propagate the strides of the destination
-                # array to the array inside the nested SDFG.
+            elif isinstance(write_edge.src, dace_nodes.Tasklet):
+                # The temporary data written by a tasklet can be safely deleted.
                 remove_last_node = True
+            elif isinstance(write_edge.src, dace_nodes.NestedSDFG):
+                if isinstance(dest_desc, dace.data.Scalar):
+                    # We keep scalar temporary storage, as a general rule, since it
+                    # does not affect performance of the generated code. This scalar
+                    # is only required in some cases, e.g. for nested SDFGs implementing
+                    # reduction, which use a WCR memlet for the reduction operation.
+                    remove_last_node = False
+                else:
+                    # We remove the transient array on the output connection of a nested
+                    # SDFG and write directly to the destination node.
+                    # The caller is responsible to propagate the strides of the destination
+                    # array to the array inside the nested SDFG.
+                    remove_last_node = True
+            else:
+                remove_last_node = False
         else:
             remove_last_node = False
 
-        if allow_removal_of_last_node and remove_last_node:
+        if remove_last_node:
             src_node = write_edge.src
             src_node_connector = write_edge.src_conn
             src_subset = write_edge.data.src_subset

src/gt4py/next/program_processors/runners/dace/lowering/gtir_to_sdfg_primitives.py

@@ -9,6 +9,7 @@
 from __future__ import annotations
 
 import abc
+from collections import Counter
 from typing import TYPE_CHECKING, Iterable, Optional, Protocol
 
 import dace
@@ -95,6 +96,7 @@ def _create_field_operator_impl(
     output_edge: gtir_dataflow.DataflowOutputEdge,
     output_type: ts.FieldType,
     map_exit: dace_nodes.MapExit,
+    output_consumer_count: dict[dace_nodes.AccessNode, int],
 ) -> gtir_to_sdfg_types.FieldopData:
     """
     Helper method to allocate a temporary array that stores one field computed
@@ -166,8 +168,11 @@ def _create_field_operator_impl(
         )
     field_node = ctx.state.add_access(field_name)
 
-    # and here the edge writing the dataflow result data through the map exit node
-    output_edge.connect(map_exit, field_node, field_subset)
+    # and here the edge writing the dataflow result data through the map exit node.
+    # Note that we cannot remove the output data access node only if this is used
+    #  for mutiple fields in a return tuple.
+    allow_removal_of_last_node = output_consumer_count[output_edge.result.dc_node] == 1
+    output_edge.connect(map_exit, field_node, field_subset, allow_removal_of_last_node)
 
     return gtir_to_sdfg_types.FieldopData(
         field_node, ts.FieldType(field_dims, output_edge.result.gt_dtype), tuple(field_origin)
@@ -217,17 +222,24 @@ def _create_field_operator(
     for edge in input_edges:
         edge.connect(map_entry)
 
+    # The same output node could be used for multiple fields in case of tuple return.
+    # In this case, the output access node cannot be removed.
+    consumer_count = Counter(
+        oedge.result.dc_node
+        for oedge in gtx_utils.flatten_nested_tuple((output_tree,))
+        if oedge is not None
+    )
     if isinstance(node_type, ts.FieldType):
         assert isinstance(output_tree, gtir_dataflow.DataflowOutputEdge)
         return _create_field_operator_impl(
-            ctx, sdfg_builder, domain, output_tree, node_type, map_exit
+            ctx, sdfg_builder, domain, output_tree, node_type, map_exit, consumer_count
         )
     else:
         # handle tuples of fields
         output_symbol_tree = gtir_to_sdfg_utils.make_symbol_tree("x", node_type)
         return gtx_utils.tree_map(
             lambda output_edge, output_sym: _create_field_operator_impl(
-                ctx, sdfg_builder, domain, output_edge, output_sym.type, map_exit
+                ctx, sdfg_builder, domain, output_edge, output_sym.type, map_exit, consumer_count
             )
         )(output_tree, output_symbol_tree)
 

src/gt4py/next/program_processors/runners/dace/lowering/gtir_to_sdfg_scan.py

@@ -22,6 +22,7 @@
 
 from __future__ import annotations
 
+from collections import Counter
 from typing import Iterable, Sequence
 
 import dace
@@ -73,6 +74,7 @@ def _create_scan_field_operator_impl(
     output_domain: infer_domain.NonTupleDomainAccess,
     output_type: ts.FieldType,
     map_exit: dace_nodes.MapExit | None,
+    output_consumer_count: dict[dace_nodes.AccessNode, int],
 ) -> gtir_to_sdfg_types.FieldopData | None:
     """
     Helper method to allocate a temporary array that stores one field computed
@@ -146,7 +148,12 @@ def _create_scan_field_operator_impl(
     #  Up to now the nested SDFG is writing into a transient data container that
     #  has the size to hold one column. The function below, that does the connection,
     #  will remove that transient and write directly to the result field.
-    inner_map_output_temporary_removed = output_edge.connect(map_exit, field_node, field_subset)
+    #  Note that we cannot remove the output data access node only if this is used
+    #  for mutiple fields in a return tuple.
+    allow_removal_of_last_node = output_consumer_count[output_edge.result.dc_node] == 1
+    inner_map_output_temporary_removed = output_edge.connect(
+        map_exit, field_node, field_subset, allow_removal_of_last_node
+    )
     if not inner_map_output_temporary_removed:
         raise ValueError("The scan nested SDFG is expected to write directly to the result field.")
 
@@ -244,6 +251,14 @@ def _create_scan_field_operator(
         else im.sym("__gtir_unused_dummy_var", node_type)
     )
 
+    # The same output node could be used for multiple fields in case of tuple return.
+    # In this case, the output access node cannot be removed.
+    consumer_count = Counter(
+        oedge.result.dc_node
+        for oedge in gtx_utils.flatten_nested_tuple((output,))
+        if oedge is not None
+    )
+
     return gtx_utils.tree_map(
         lambda edge, domain, sym: _create_scan_field_operator_impl(
             ctx,
@@ -252,6 +267,7 @@ def _create_scan_field_operator(
             domain,
             sym.type,
             map_exit,
+            consumer_count,
         )
     )(output, output_domain, dummy_output_symbol)