Repository: GridTools/gt4py
Pull request #1834 (2025-01-29 12:45:48)

refactor[next][daxce]: cleanup get_sdfg_args() parameter list
Remove unnecessary `kwargs` from parameter list of `get_sdfg_args()`.


Patch:
diff --git a/src/gt4py/next/program_processors/runners/dace/gtir_dataflow.py b/src/gt4py/next/program_processors/runners/dace/gtir_dataflow.py
index e00e363a..a34828af 100644
--- a/src/gt4py/next/program_processors/runners/dace/gtir_dataflow.py
+++ b/src/gt4py/next/program_processors/runners/dace/gtir_dataflow.py
@@ -332,6 +332,7 @@ class LambdaToDataflow(eve.NodeVisitor):
     sdfg: dace.SDFG
     state: dace.SDFGState
     subgraph_builder: gtir_sdfg.DataflowBuilder
+    scan_carry_symbol: Optional[gtir.Sym]
     input_edges: list[DataflowInputEdge] = dataclasses.field(default_factory=lambda: [])
     symbol_map: dict[
         str,
@@ -693,7 +694,7 @@ class LambdaToDataflow(eve.NodeVisitor):
         # visit each branch of the if-statement as if it was a Lambda node
         lambda_node = gtir.Lambda(params=lambda_params, expr=expr)
         input_edges, output_tree = translate_lambda_to_dataflow(
-            if_sdfg, if_branch_state, self.subgraph_builder, lambda_node, args=lambda_args
+            if_sdfg, if_branch_state, self.subgraph_builder, lambda_node, lambda_args
         )
 
         for data_node in if_branch_state.data_nodes():
@@ -736,7 +737,12 @@ class LambdaToDataflow(eve.NodeVisitor):
         """
         Lowers an if-expression with exclusive branch execution into a nested SDFG,
         in which each branch is lowered into a dataflow in a separate state and
-        the if-condition is represented as the inter-state edge condtion.
+        the if-condition is represented as the inter-state edge condition.
+
+        Exclusive branch execution for local if expressions is meant to be used
+        in iterator view. Iterator view is required ONLY inside scan field operators.
+        For regular field operators, the fieldview behavior of if-expressions
+        corresponds to a local select, therefore it should be lowered to a tasklet.
         """
 
         def write_output_of_nested_sdfg_to_temporary(inner_value: ValueExpr) -> ValueExpr:
@@ -1633,13 +1639,87 @@ class LambdaToDataflow(eve.NodeVisitor):
         tuple_fields = self.visit(node.args[1])
         return tuple_fields[index]
 
+    def requires_exclusive_if(self, node: gtir.FunCall) -> bool:
+        """
+        The meaning of `if_` builtin function is unclear in GTIR.
+        In some context, it corresponds to a ternary operator where, depending on
+        the condition result, only one branch or the other should be executed,
+        because one of them is invalid. The typical case is the use of `if_` to
+        decide whether it is possible or not to access a shifted iterator, for
+        example when the condition expression calls `can_deref`.
+        The ternary operator is also used in iterator view, where the field arguments
+        are not necessarily both defined on the entire output domain (this behavior
+        should not appear in field view, because there the user code should use
+        `concat_where` instead of `where` for such cases). It is difficult to catch
+        such behavior, because it would require to know the exact domain of all
+        fields, which is not known at compile time. However, the iterator view
+        behavior should only appear inside scan field operators.
+        A different usage of `if_` expressions is selecting one argument value or
+        the other, where both arguments are defined on the output domain, therefore
+        always valid.
+        In order to simplify the SDFG and facilitate the optimization stage, we
+        try to avoid the ternary operator form when not needed. The reason is that
+        exclusive branch execution is represented in the SDFG as a conditional
+        state transition, which prevents fusion.
+        """
+        assert cpm.is_call_to(node, "if_")
+        assert len(node.args) == 3
+
+        condition_vars = (
+            eve.walk_values(node.args[0])
+            .if_isinstance(gtir.SymRef)
+            .map(lambda node: str(node.id))
+            .filter(lambda x: x in self.symbol_map)
+            .to_set()
+        )
+
+        # first, check if any argument contains shift expressions that depend on the condition variables
+        for arg in node.args[1:3]:
+            shift_nodes = (
+                eve.walk_values(arg).filter(lambda node: cpm.is_applied_shift(node)).to_set()
+            )
+            for shift_node in shift_nodes:
+                shift_vars = (
+                    eve.walk_values(shift_node)
+                    .if_isinstance(gtir.SymRef)
+                    .map(lambda node: str(node.id))
+                    .filter(lambda x: x in self.symbol_map)
+                    .to_set()
+                )
+                # require exclusive branch execution if any shift expression one of
+                # the if branches accesses a variable used in the condition expression
+                depend_vars = condition_vars.intersection(shift_vars)
+                if len(depend_vars) != 0:
+                    return True
+
+        # secondly, check whether the `if_` branches access different sets of fields
+        # and this happens inside a scan field operator
+        if self.scan_carry_symbol is not None:
+            # the `if_` node is inside a scan stencil expression
+            scan_carry_var = str(self.scan_carry_symbol.id)
+            if scan_carry_var in condition_vars:
+                br1_vars, br2_vars = (
+                    eve.walk_values(arg)
+                    .if_isinstance(gtir.SymRef)
+                    .map(lambda node: str(node.id))
+                    .filter(lambda x: isinstance(self.symbol_map.get(x, None), MemletExpr))
+                    .to_set()
+                    for arg in node.args[1:3]
+                )
+                if br1_vars != br2_vars:
+                    # the two branches of the `if_` expression access different sets of fields,
+                    # depending on the scan carry value
+                    return True
+
+        return False
+
     def visit_FunCall(
         self, node: gtir.FunCall
     ) -> IteratorExpr | DataExpr | tuple[IteratorExpr | DataExpr | tuple[Any, ...], ...]:
         if cpm.is_call_to(node, "deref"):
             return self._visit_deref(node)
 
-        elif cpm.is_call_to(node, "if_"):
+        elif cpm.is_call_to(node, "if_") and self.requires_exclusive_if(node):
             return self._visit_if(node)
 
         elif cpm.is_call_to(node, "neighbors"):
@@ -1776,6 +1856,7 @@ def translate_lambda_to_dataflow(
         | ValueExpr
         | tuple[IteratorExpr | MemletExpr | ValueExpr | tuple[Any, ...], ...]
     ],
+    scan_carry_symbol: Optional[gtir.Sym] = None,
 ) -> tuple[
     list[DataflowInputEdge],
     tuple[DataflowOutputEdge | tuple[Any, ...], ...],
@@ -1794,13 +1875,15 @@ def translate_lambda_to_dataflow(
         sdfg_builder: Helper class to build the dataflow inside the given SDFG.
         node: Lambda node to visit.
         args: Arguments passed to lambda node.
+        scan_carry_symbol: When set, the lowering of `if_` expression will consider
+            using the ternary operator form with exclusive branch execution.
 
     Returns:
         A tuple of two elements:
         - List of connections for data inputs to the dataflow.
         - Tree representation of output data connections.
     """
-    taskgen = LambdaToDataflow(sdfg, state, sdfg_builder)
+    taskgen = LambdaToDataflow(sdfg, state, sdfg_builder, scan_carry_symbol)
     lambda_output = taskgen.visit_let(node, args)
 
     if isinstance(lambda_output, DataflowOutputEdge):
diff --git a/src/gt4py/next/program_processors/runners/dace/gtir_scan_translator.py b/src/gt4py/next/program_processors/runners/dace/gtir_scan_translator.py
index d3d8e101..ec88cd8f 100644
--- a/src/gt4py/next/program_processors/runners/dace/gtir_scan_translator.py
+++ b/src/gt4py/next/program_processors/runners/dace/gtir_scan_translator.py
@@ -417,7 +417,12 @@ def _lower_lambda_to_nested_sdfg(
     # stil inside the 'compute' state, generate the dataflow representing the stencil
     # to be applied on the horizontal domain
     lambda_input_edges, lambda_result = gtir_dataflow.translate_lambda_to_dataflow(
-        nsdfg, compute_state, lambda_translator, lambda_node, args=stencil_args
+        nsdfg,
+        compute_state,
+        lambda_translator,
+        lambda_node,
+        stencil_args,
+        scan_carry_symbol=scan_carry_symbol,
     )
     # connect the dataflow input directly to the source data nodes, without passing through a map node;
     # the reason is that the map for horizontal domain is outside the scan loop region
diff --git a/src/gt4py/next/program_processors/runners/dace/sdfg_callable.py b/src/gt4py/next/program_processors/runners/dace/sdfg_callable.py
index ecd23619..7f221a5a 100644
--- a/src/gt4py/next/program_processors/runners/dace/sdfg_callable.py
+++ b/src/gt4py/next/program_processors/runners/dace/sdfg_callable.py
@@ -129,15 +129,23 @@ def get_sdfg_args(
     *args: Any,
     check_args: bool = False,
     on_gpu: bool = False,
-    **kwargs: Any,
 ) -> dict[str, Any]:
     """Extracts the arguments needed to call the SDFG.
 
-    This function can handle the same arguments that are passed to dace runner.
+    This function can handle the arguments that are passed to the dace runner
+    and that end up in the decoration stage of the dace backend workflow.
 
     Args:
         sdfg:               The SDFG for which we want to get the arguments.
-        offset_provider:    Offset provider.
+        offset_provider:    The offset provider.
+        args:               The list of arguments passed to the dace runner.
+        check_args:         If True, return only the arguments that are expected
+                            according to the SDFG signature.
+        on_gpu:             If True, this method ensures that the arrays for the
+                            connectivity tables are allocated in GPU memory.
+
+    Returns:
+        A dictionary of keyword arguments to be passed in the SDFG call.
     """
 
     dace_args = _get_args(sdfg, args)
diff --git a/src/gt4py/next/program_processors/runners/dace/transformations/utils.py b/src/gt4py/next/program_processors/runners/dace/transformations/utils.py
index 87308061..3cc2dadd 100644
--- a/src/gt4py/next/program_processors/runners/dace/transformations/utils.py
+++ b/src/gt4py/next/program_processors/runners/dace/transformations/utils.py
@@ -50,6 +50,7 @@ def gt_make_transients_persistent(
         not_modify_lifetime: set[str] = set()
 
         for state in nsdfg.states():
+            scope_dict = state.scope_dict()
             for dnode in state.data_nodes():
                 if dnode.data in not_modify_lifetime:
                     continue
@@ -70,6 +71,15 @@ def gt_make_transients_persistent(
                     not_modify_lifetime.add(dnode.data)
                     continue
 
+                # If the data is referenced inside a scope, such as a map, it might be possible
+                #  that it is only used inside that scope. If we would make it persistent, then
+                #  it would essentially be allocated outside and be shared among the different
+                #  map iterations. So we can not make it persistent.
+                #  The downside is, that we might have to perform dynamic allocation.
+                if scope_dict[dnode] is not None:
+                    not_modify_lifetime.add(dnode.data)
+                    continue
+
                 try:
                     # The symbols describing the total size must be a subset of the
                     #  free symbols of the SDFG (symbols passed as argument).
diff --git a/src/gt4py/next/program_processors/runners/dace/workflow/decoration.py b/src/gt4py/next/program_processors/runners/dace/workflow/decoration.py
index 2ee99f5f..9648ac9e 100644
--- a/src/gt4py/next/program_processors/runners/dace/workflow/decoration.py
+++ b/src/gt4py/next/program_processors/runners/dace/workflow/decoration.py
@@ -86,7 +86,6 @@ def convert_args(
             *flat_args,
             check_args=False,
             on_gpu=on_gpu,
-            use_field_canonical_representation=use_field_canonical_representation,
         )
 
         with dace.config.temporary_config():


Test patch:
diff --git a/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/transformation_tests/test_make_transients_persistent.py b/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/transformation_tests/test_make_transients_persistent.py
new file mode 100644
index 00000000..d8cf8e33
--- /dev/null
+++ b/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/transformation_tests/test_make_transients_persistent.py
@@ -0,0 +1,74 @@
+# GT4Py - GridTools Framework
+#
+# Copyright (c) 2014-2024, ETH Zurich
+# All rights reserved.
+#
+# Please, refer to the LICENSE file in the root directory.
+# SPDX-License-Identifier: BSD-3-Clause
+
+import pytest
+
+dace = pytest.importorskip("dace")
+from dace.sdfg import nodes as dace_nodes
+
+from gt4py.next.program_processors.runners.dace import (
+    transformations as gtx_transformations,
+)
+
+from . import util
+
+import dace
+
+
+def _make_transients_persistent_inner_access_sdfg() -> tuple[dace.SDFG, dace.SDFGState]:
+    sdfg = dace.SDFG(util.unique_name("transients_persistent_inner_access_sdfg"))
+    state = sdfg.add_state(is_start_block=True)
+
+    for name in "abc":
+        sdfg.add_array(
+            name,
+            shape=(10,),
+            dtype=dace.float64,
+            transient=False,
+        )
+    sdfg.arrays["b"].transient = True
+
+    me: dace_nodes.MapEntry
+    mx: dace_nodes.MapExit
+    me, mx = state.add_map("comp", ndrange={"__i0": "0:10"})
+    a, b, c = (state.add_access(name) for name in "abc")
+    tsklt: dace_nodes.Tasklet = state.add_tasklet(
+        "tsklt",
+        inputs={"__in"},
+        code="__out = __in + 1.0",
+        outputs={"__out"},
+    )
+
+    me.add_in_connector("IN_A")
+    state.add_edge(a, None, me, "IN_A", dace.Memlet("a[0:10]"))
+
+    me.add_out_connector("OUT_A")
+    state.add_edge(me, "OUT_A", b, None, dace.Memlet("a[__i0] -> [__i0]"))
+
+    state.add_edge(b, None, tsklt, "__in", dace.Memlet("b[__i0]"))
+
+    mx.add_in_connector("IN_C")
+    state.add_edge(tsklt, "__out", mx, "IN_C", dace.Memlet("c[__i0]"))
+
+    mx.add_out_connector("OUT_C")
+    state.add_edge(mx, "OUT_C", c, None, dace.Memlet("c[0:10]"))
+    sdfg.validate()
+    return sdfg, state
+
+
+def test_make_transients_persistent_inner_access():
+    sdfg, state = _make_transients_persistent_inner_access_sdfg()
+    assert sdfg.arrays["b"].lifetime is dace.dtypes.AllocationLifetime.Scope
+
+    # Because `b`, the only transient, is used inside a map scope, it is not selected,
+    #  although in this situation it would be possible.
+    change_report: dict[int, set[str]] = gtx_transformations.gt_make_transients_persistent(
+        sdfg, device=dace.DeviceType.CPU
+    )
+    assert len(change_report) == 1
+    assert change_report[sdfg.cfg_id] == set()