Repository: GridTools/gt4py
Pull request #1695 (2024-10-20 19:17:47)

feat[next][dace]: Cleanup scalar args
This PR simplifies the representation of program scalar arguments in the SDFG: instead of promoting them to symbols, we represent them as scalar data. Scalar to symbol promotion is a transformation pass available in DaCe that should be applied after lowering to SDFG, eventually.

Additional changes:
- Fixes a problem in naming of input connectors for scalar expression tasklets: the connector name cannot match the argument name, otherwise we cannot pass the same value to two connectors in expressions like `out = tmp * tmp`.
- Only propagate to lambda scope the symbols that are referenced.

edopao: cscs-ci run default

Patch:
diff --git a/src/gt4py/next/program_processors/runners/dace_fieldview/gtir_builtin_translators.py b/src/gt4py/next/program_processors/runners/dace_fieldview/gtir_builtin_translators.py
index a8ae1cc0..3cd2b17b 100644
--- a/src/gt4py/next/program_processors/runners/dace_fieldview/gtir_builtin_translators.py
+++ b/src/gt4py/next/program_processors/runners/dace_fieldview/gtir_builtin_translators.py
@@ -481,13 +481,12 @@ def _get_data_nodes(
         sym_node = state.add_access(sym_name)
         return Field(sym_node, sym_type)
     elif isinstance(sym_type, ts.ScalarType):
-        if sym_name in sdfg.arrays:
-            # access the existing scalar container
-            sym_node = state.add_access(sym_name)
-        else:
+        if sym_name in sdfg.symbols:
             sym_node = _get_symbolic_value(
                 sdfg, state, sdfg_builder, sym_name, sym_type, temp_name=f"__{sym_name}"
             )
+        else:
+            sym_node = state.add_access(sym_name)
         return Field(sym_node, sym_type)
     elif isinstance(sym_type, ts.TupleType):
         tuple_fields = dace_gtir_utils.get_tuple_fields(sym_name, sym_type)
@@ -612,7 +611,7 @@ def translate_scalar_expr(
     connectors = []
     scalar_expr_args = []
 
-    for arg_expr in node.args:
+    for i, arg_expr in enumerate(node.args):
         visit_expr = True
         if isinstance(arg_expr, gtir.SymRef):
             try:
@@ -636,7 +635,7 @@ def translate_scalar_expr(
             )
             if not (isinstance(arg, Field) and isinstance(arg.data_type, ts.ScalarType)):
                 raise ValueError(f"Invalid argument to scalar expression {arg_expr}.")
-            param = f"__in_{arg.data_node.data}"
+            param = f"__arg{i}"
             args.append(arg.data_node)
             connectors.append(param)
             scalar_expr_args.append(gtir.SymRef(id=param))
diff --git a/src/gt4py/next/program_processors/runners/dace_fieldview/gtir_dataflow.py b/src/gt4py/next/program_processors/runners/dace_fieldview/gtir_dataflow.py
index 0e571fc1..fdd92b57 100644
--- a/src/gt4py/next/program_processors/runners/dace_fieldview/gtir_dataflow.py
+++ b/src/gt4py/next/program_processors/runners/dace_fieldview/gtir_dataflow.py
@@ -939,7 +939,7 @@ class LambdaToDataflow(eve.NodeVisitor):
             arg_expr = self.visit(arg)
             if isinstance(arg_expr, MemletExpr | DataExpr):
                 # the argument value is the result of a tasklet node or direct field access
-                connector = f"__inp_{i}"
+                connector = f"__arg{i}"
                 node_connections[connector] = arg_expr
                 node_internals.append(connector)
             else:
diff --git a/src/gt4py/next/program_processors/runners/dace_fieldview/gtir_sdfg.py b/src/gt4py/next/program_processors/runners/dace_fieldview/gtir_sdfg.py
index 3697609d..31e561a1 100644
--- a/src/gt4py/next/program_processors/runners/dace_fieldview/gtir_sdfg.py
+++ b/src/gt4py/next/program_processors/runners/dace_fieldview/gtir_sdfg.py
@@ -17,6 +17,7 @@ from __future__ import annotations
 import abc
 import dataclasses
 import itertools
+import operator
 from typing import Any, Dict, Iterable, List, Optional, Protocol, Sequence, Set, Tuple, Union
 
 import dace
@@ -26,7 +27,7 @@ from gt4py.eve import concepts
 from gt4py.next import common as gtx_common, utils as gtx_utils
 from gt4py.next.iterator import ir as gtir
 from gt4py.next.iterator.ir_utils import common_pattern_matcher as cpm
-from gt4py.next.iterator.transforms import prune_casts as ir_prune_casts
+from gt4py.next.iterator.transforms import prune_casts as ir_prune_casts, symbol_ref_utils
 from gt4py.next.iterator.type_system import inference as gtir_type_inference
 from gt4py.next.program_processors.runners.dace_common import utility as dace_utils
 from gt4py.next.program_processors.runners.dace_fieldview import (
@@ -106,6 +107,38 @@ class SDFGBuilder(DataflowBuilder, Protocol):
         ...
 
 
+def _collect_symbols_in_domain_expressions(
+    ir: gtir.Node, ir_params: Sequence[gtir.Sym]
+) -> set[str]:
+    """
+    Collect symbols accessed in domain expressions that also appear in the paremeter list.
+
+    This function is used to identify all parameters that are accessed in domain
+    expressions. They have to be passed to the SDFG call as DaCe symbols (instead
+    of scalars) such that they can be used as bounds in map ranges.
+
+    Args:
+        ir: GTIR node to be traversed and where to search for domain expressions.
+        ir_params: List of parameters to search for in domain expressions.
+
+    Returns:
+        A set of names corresponding to the parameters found in domain expressions.
+    """
+    params = {str(sym.id) for sym in ir_params}
+    return set(
+        eve.walk_values(ir)
+        .filter(lambda node: cpm.is_call_to(node, ("cartesian_domain", "unstructured_domain")))
+        .map(
+            lambda domain: eve.walk_values(domain)
+            .if_isinstance(gtir.SymRef)
+            .map(lambda symref: str(symref.id))
+            .filter(lambda sym: sym in params)
+            .to_list()
+        )
+        .reduce(operator.add, init=[])
+    )
+
+
 @dataclasses.dataclass(frozen=True)
 class GTIRToSDFG(eve.NodeVisitor, SDFGBuilder):
     """Provides translation capability from a GTIR program to a DaCe SDFG.
@@ -172,60 +205,79 @@ class GTIRToSDFG(eve.NodeVisitor, SDFGBuilder):
     def _add_storage(
         self,
         sdfg: dace.SDFG,
+        symbolic_arguments: set[str],
         name: str,
-        symbol_type: ts.DataType,
+        gt_type: ts.DataType,
         transient: bool = True,
-        is_tuple_member: bool = False,
     ) -> list[tuple[str, ts.DataType]]:
         """
-        Add storage for data containers used in the SDFG. For fields, it allocates dace arrays,
-        while scalars are stored as SDFG symbols.
+        Add storage in the SDFG for a given GT4Py data symbol.
 
-        The fields used as temporary arrays, when `transient = True`, are allocated and exist
-        only within the SDFG; when `transient = False`, the fields have to be allocated outside
-        and have to be passed as array arguments to the SDFG.
+        GT4Py fields are allocated as DaCe arrays. GT4Py scalars are represented
+        as DaCe scalar objects in the SDFG; the exception are the symbols passed as
+        `symbolic_arguments`, e.g. symbols used in domain expressions, and those used
+        for symbolic array shape and strides.
+
+        The fields used as temporary arrays, when `transient = True`, are allocated
+        and exist only within the SDFG; when `transient = False`, the fields have
+        to be allocated outside and have to be passed as arguments to the SDFG call.
+
+        Args:
+            sdfg: The SDFG where storage needs to be allocated.
+            symbolic_arguments: Set of GT4Py scalars that must be represented as SDFG symbols.
+            name: Symbol Name to be allocated.
+            gt_type: GT4Py symbol type.
+            transient: True when the data symbol has to be allocated as internal storage.
 
         Returns:
-            List of data containers or symbols allocated as storage. This is a list, not a single value,
-            because in case of tuples we flat the tuple fields (eventually nested) and allocate storage
-            for each tuple element.
+            List of tuples '(data_name, gt_type)' where 'data_name' is the name of
+            the data container used as storage in the SDFG and 'gt_type' is the
+            corresponding GT4Py type. In case the storage has to be allocated for
+            a tuple symbol the list contains a flattened version of the tuple,
+            otherwise the list will contain a single entry.
         """
-        if isinstance(symbol_type, ts.TupleType):
+        if isinstance(gt_type, ts.TupleType):
             tuple_fields = []
             for tname, tsymbol_type in dace_gtir_utils.get_tuple_fields(
-                name, symbol_type, flatten=True
+                name, gt_type, flatten=True
             ):
                 tuple_fields.extend(
-                    self._add_storage(sdfg, tname, tsymbol_type, transient, is_tuple_member=True)
+                    self._add_storage(sdfg, symbolic_arguments, tname, tsymbol_type, transient)
                 )
             return tuple_fields
 
-        elif isinstance(symbol_type, ts.FieldType):
-            dtype = dace_utils.as_dace_type(symbol_type.dtype)
+        elif isinstance(gt_type, ts.FieldType):
+            dtype = dace_utils.as_dace_type(gt_type.dtype)
             # use symbolic shape, which allows to invoke the program with fields of different size;
             # and symbolic strides, which enables decoupling the memory layout from generated code.
-            sym_shape, sym_strides = self._make_array_shape_and_strides(name, symbol_type.dims)
+            sym_shape, sym_strides = self._make_array_shape_and_strides(name, gt_type.dims)
             sdfg.add_array(name, sym_shape, dtype, strides=sym_strides, transient=transient)
 
-            return [(name, symbol_type)]
+            return [(name, gt_type)]
 
-        elif isinstance(symbol_type, ts.ScalarType):
-            dtype = dace_utils.as_dace_type(symbol_type)
-            # Scalar arguments passed to the program are represented as symbols in DaCe SDFG;
-            # the exception are members of tuple arguments, that are represented as scalar containers.
-            # The field size is sometimes passed as scalar argument to the program, so we have to
-            # check if the shape symbol was already allocated by `_make_array_shape_and_strides`.
-            # We assume that the scalar argument for field size always follows the field argument.
-            if is_tuple_member:
-                sdfg.add_scalar(name, dtype, transient=transient)
-            elif name in sdfg.symbols:
-                assert sdfg.symbols[name].dtype == dtype
-            else:
+        elif isinstance(gt_type, ts.ScalarType):
+            dtype = dace_utils.as_dace_type(gt_type)
+            if name in symbolic_arguments:
                 sdfg.add_symbol(name, dtype)
+            elif dace_utils.is_field_symbol(name):
+                # Sometimes, when the field domain is implicitly derived from the
+                # field domain, the gt4py lowering adds the field size as a scalar
+                # argument to the program IR. Suppose a field '__sym', then gt4py
+                # will add '__sym_size_0'.
+                # Therefore, here we check whether the shape symbol was already
+                # created by `_make_array_shape_and_strides`, when allocating
+                # storage for field arguments. We assume that the scalar argument
+                # for field size, if present, always follows the field argument.
+                if name in sdfg.symbols:
+                    assert sdfg.symbols[name].dtype == dtype
+                else:
+                    sdfg.add_symbol(name, dtype)
+            else:
+                sdfg.add_scalar(name, dtype, transient=transient)
 
-            return [(name, symbol_type)]
+            return [(name, gt_type)]
 
-        raise RuntimeError(f"Data type '{type(symbol_type)}' not supported.")
+        raise RuntimeError(f"Data type '{type(gt_type)}' not supported.")
 
     def _add_storage_for_temporary(self, temp_decl: gtir.Temporary) -> dict[str, str]:
         """
@@ -246,10 +298,6 @@ class GTIRToSDFG(eve.NodeVisitor, SDFGBuilder):
 
         Returns:
             A list of array nodes containing the result fields.
-
-        TODO: Do we need to return the GT4Py `FieldType`/`ScalarType`? It is needed
-        in case the transient arrays containing the expression result are not guaranteed
-        to have the same memory layout as the target array.
         """
         result = self.visit(node, sdfg=sdfg, head_state=head_state, reduce_identity=None)
 
@@ -274,14 +322,28 @@ class GTIRToSDFG(eve.NodeVisitor, SDFGBuilder):
         temp_result = gtx_utils.tree_map(make_temps)(result)
         return list(gtx_utils.flatten_nested_tuple((temp_result,)))
 
-    def _add_sdfg_params(self, sdfg: dace.SDFG, node_params: Sequence[gtir.Sym]) -> list[str]:
-        """Helper function to add storage for node parameters and connectivity tables."""
+    def _add_sdfg_params(
+        self,
+        sdfg: dace.SDFG,
+        node_params: Sequence[gtir.Sym],
+        symbolic_arguments: set[str],
+    ) -> list[str]:
+        """
+        Helper function to add storage for node parameters and connectivity tables.
+
+        GT4Py field arguments will be translated to `dace.data.Array` objects.
+        GT4Py scalar arguments will be translated to `dace.data.Scalar` objects,
+        except when they are listed in 'symbolic_arguments', in which case they
+        will be represented in the SDFG as DaCe symbols.
+        """
         # add non-transient arrays and/or SDFG symbols for the program arguments
         sdfg_args = []
         for param in node_params:
             pname = str(param.id)
             assert isinstance(param.type, (ts.DataType))
-            sdfg_args += self._add_storage(sdfg, pname, param.type, transient=False)
+            sdfg_args += self._add_storage(
+                sdfg, symbolic_arguments, pname, param.type, transient=False
+            )
             self.global_symbols[pname] = param.type
 
         # add SDFG storage for connectivity tables
@@ -290,7 +352,7 @@ class GTIRToSDFG(eve.NodeVisitor, SDFGBuilder):
         ).items():
             scalar_kind = tt.get_scalar_kind(offset_provider.index_type)
             local_dim = gtx_common.Dimension(offset, kind=gtx_common.DimensionKind.LOCAL)
-            type_ = ts.FieldType(
+            gt_type = ts.FieldType(
                 [offset_provider.origin_axis, local_dim], ts.ScalarType(scalar_kind)
             )
             # We store all connectivity tables as transient arrays here; later, while building
@@ -298,7 +360,9 @@ class GTIRToSDFG(eve.NodeVisitor, SDFGBuilder):
             # the tables that are actually used. This way, we avoid adding SDFG arguments for
             # the connectivity tables that are not used. The remaining unused transient arrays
             # are removed by the dace simplify pass.
-            self._add_storage(sdfg, dace_utils.connectivity_identifier(offset), type_)
+            self._add_storage(
+                sdfg, symbolic_arguments, dace_utils.connectivity_identifier(offset), gt_type
+            )
 
         # the list of all sdfg arguments (aka non-transient arrays) which include tuple-element fields
         return [arg_name for arg_name, _ in sdfg_args]
@@ -330,7 +394,8 @@ class GTIRToSDFG(eve.NodeVisitor, SDFGBuilder):
         else:
             head_state = entry_state
 
-        sdfg_arg_names = self._add_sdfg_params(sdfg, node.params)
+        domain_symbols = _collect_symbols_in_domain_expressions(node, node.params)
+        sdfg_arg_names = self._add_sdfg_params(sdfg, node.params, symbolic_arguments=domain_symbols)
 
         # visit one statement at a time and expand the SDFG from the current head state
         for i, stmt in enumerate(node.body):
@@ -490,7 +555,10 @@ class GTIRToSDFG(eve.NodeVisitor, SDFGBuilder):
         ]
 
         # inherit symbols from parent scope but eventually override with local symbols
-        lambda_symbols = self.global_symbols | {
+        lambda_symbols = {
+            sym: self.global_symbols[sym]
+            for sym in symbol_ref_utils.collect_symbol_refs(node.expr, self.global_symbols.keys())
+        } | {
             pname: dace_gtir_utils.get_tuple_type(arg) if isinstance(arg, tuple) else arg.data_type
             for pname, arg in lambda_args_mapping
         }
@@ -501,11 +569,12 @@ class GTIRToSDFG(eve.NodeVisitor, SDFGBuilder):
         nstate = nsdfg.add_state("lambda")
 
         # add sdfg storage for the symbols that need to be passed as input parameters
+        lambda_params = [
+            gtir.Sym(id=p_name, type=p_type) for p_name, p_type in lambda_symbols.items()
+        ]
+        lambda_domain_symbols = _collect_symbols_in_domain_expressions(node.expr, lambda_params)
         lambda_translator._add_sdfg_params(
-            nsdfg,
-            node_params=[
-                gtir.Sym(id=p_name, type=p_type) for p_name, p_type in lambda_symbols.items()
-            ],
+            nsdfg, node_params=lambda_params, symbolic_arguments=lambda_domain_symbols
         )
 
         lambda_result = lambda_translator.visit(


Test patch:
diff --git a/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_gtir_to_sdfg.py b/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_gtir_to_sdfg.py
index 728b4b02..230ff695 100644
--- a/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_gtir_to_sdfg.py
+++ b/tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/test_gtir_to_sdfg.py
@@ -819,9 +819,7 @@ def test_gtir_cartesian_shift_left():
 
         sdfg = dace_backend.build_sdfg_from_gtir(testee, CARTESIAN_OFFSETS)
 
-        FSYMBOLS_tmp = FSYMBOLS.copy()
-        FSYMBOLS_tmp["__x_offset_stride_0"] = 1
-        sdfg(a, a_offset, b, **FSYMBOLS_tmp)
+        sdfg(a, a_offset, b, **FSYMBOLS, __x_offset_size_0=N, __x_offset_stride_0=1)
         assert np.allclose(a[OFFSET:] + DELTA, b[:-OFFSET])
 
 
@@ -914,7 +912,7 @@ def test_gtir_cartesian_shift_right():
 
         sdfg = dace_backend.build_sdfg_from_gtir(testee, CARTESIAN_OFFSETS)
 
-        sdfg(a, a_offset, b, **FSYMBOLS, __x_offset_stride_0=1)
+        sdfg(a, a_offset, b, **FSYMBOLS, __x_offset_size_0=N, __x_offset_stride_0=1)
         assert np.allclose(a[:-OFFSET] + DELTA, b[OFFSET:])
 
 
@@ -1072,7 +1070,9 @@ def test_gtir_connectivity_shift():
             __ev_field_size_1=SIMPLE_MESH.num_vertices,
             __ev_field_stride_0=SIMPLE_MESH.num_vertices,
             __ev_field_stride_1=1,
+            __c2e_offset_size_0=SIMPLE_MESH.num_cells,
             __c2e_offset_stride_0=1,
+            __e2v_offset_size_0=SIMPLE_MESH.num_edges,
             __e2v_offset_stride_0=1,
         )
         assert np.allclose(ce, ref)
@@ -1592,6 +1592,41 @@ def test_gtir_let_lambda():
     assert np.allclose(b, ref)
 
 
+def test_gtir_let_lambda_scalar_expression():
+    domain = im.domain(gtx_common.GridType.CARTESIAN, ranges={IDim: (0, "size")})
+    testee = gtir.Program(
+        id="let_lambda_scalar_expression",
+        function_definitions=[],
+        params=[
+            gtir.Sym(id="a", type=IFTYPE.dtype),
+            gtir.Sym(id="b", type=IFTYPE.dtype),
+            gtir.Sym(id="x", type=IFTYPE),
+            gtir.Sym(id="y", type=IFTYPE),
+            gtir.Sym(id="size", type=SIZE_TYPE),
+        ],
+        declarations=[],
+        body=[
+            gtir.SetAt(
+                expr=im.let("tmp", im.multiplies_("a", "b"))(
+                    im.op_as_fieldop("multiplies", domain)("x", im.multiplies_("tmp", "tmp"))
+                ),
+                domain=domain,
+                target=gtir.SymRef(id="y"),
+            )
+        ],
+    )
+
+    a = np.random.rand()
+    b = np.random.rand()
+    c = np.random.rand(N)
+    d = np.empty_like(c)
+
+    sdfg = dace_backend.build_sdfg_from_gtir(testee, {})
+
+    sdfg(a, b, c, d, **FSYMBOLS)
+    assert np.allclose(d, (a * a * b * b * c))
+
+
 def test_gtir_let_lambda_with_connectivity():
     C2E_neighbor_idx = 1
     C2V_neighbor_idx = 2
@@ -1757,9 +1792,9 @@ def test_gtir_if_scalars():
         body=[
             gtir.SetAt(
                 expr=im.let("f", im.tuple_get(0, "x"))(
-                    im.let("y", im.tuple_get(1, "x"))(
-                        im.let("y_0", im.tuple_get(0, "y"))(
-                            im.let("y_1", im.tuple_get(1, "y"))(
+                    im.let("g", im.tuple_get(1, "x"))(
+                        im.let("y_0", im.tuple_get(0, "g"))(
+                            im.let("y_1", im.tuple_get(1, "g"))(
                                 im.op_as_fieldop("plus", domain)(
                                     "f",
                                     im.if_(