feat[next][dace]: Use gt4py symbols for field size as dace array shape (

GridTools#1458)

This PR changes the ITIR DaCe backend to reuse the gt4py symbols for field size as dace symbols in array shape. The gt4py symbols are passed as scalar arguments to the program and are used in the definition of the closure domain.

src/gt4py/next/program_processors/runners/dace_iterator/__init__.py

@@ -215,7 +215,7 @@ def offset_invariants(offset):
     return m.hexdigest()
 
 
-def get_sdfg_args(sdfg: dace.SDFG, *args, **kwargs) -> dict[str, Any]:
+def get_sdfg_args(sdfg: dace.SDFG, *args, check_args: bool = False, **kwargs) -> dict[str, Any]:
     """Extracts the arguments needed to call the SDFG.
 
     This function can handle the same arguments that are passed to `run_dace_iterator()`.
@@ -229,7 +229,6 @@ def get_sdfg_args(sdfg: dace.SDFG, *args, **kwargs) -> dict[str, Any]:
     neighbor_tables = filter_neighbor_tables(offset_provider)
     device = dace.DeviceType.GPU if on_gpu else dace.DeviceType.CPU
 
-    sdfg_sig = sdfg.signature_arglist(with_types=False)
     dace_args = get_args(sdfg, args)
     dace_field_args = {n: v for n, v in dace_args.items() if not np.isscalar(v)}
     dace_conn_args = get_connectivity_args(neighbor_tables, device)
@@ -247,9 +246,13 @@ def get_sdfg_args(sdfg: dace.SDFG, *args, **kwargs) -> dict[str, Any]:
         **dace_conn_strides,
         **dace_offsets,
     }
-    expected_args = {key: all_args[key] for key in sdfg_sig}
 
-    return expected_args
+    if check_args:
+        # return only arguments expected in SDFG signature (note hat `signature_arglist` takes time)
+        sdfg_sig = sdfg.signature_arglist(with_types=False)
+        return {key: all_args[key] for key in sdfg_sig}
+
+    return all_args
 
 
 def build_sdfg_from_itir(
@@ -390,12 +393,12 @@ def run_dace_iterator(program: itir.FencilDefinition, *args, **kwargs):
         if build_cache is not None:
             build_cache[cache_id] = sdfg_program
 
-    expected_args = get_sdfg_args(sdfg, *args, **kwargs)
+    sdfg_args = get_sdfg_args(sdfg, *args, **kwargs)
 
     with dace.config.temporary_config():
         dace.config.Config.set("compiler", "allow_view_arguments", value=True)
         dace.config.Config.set("frontend", "check_args", value=True)
-        sdfg_program(**expected_args)
+        sdfg_program(**sdfg_args)
 
 
 def _run_dace_cpu(program: itir.FencilDefinition, *args, **kwargs) -> None:

src/gt4py/next/program_processors/runners/dace_iterator/itir_to_sdfg.py

@@ -104,7 +104,7 @@ def _make_array_shape_and_strides(
     dims: Sequence[Dimension],
     neighbor_tables: Mapping[str, NeighborTable],
     sort_dims: bool,
-) -> tuple[list[dace.symbol], list[dace.symbol]]:
+) -> tuple[list[dace.symbol], list[dace.symbol], list[dace.symbol]]:
     """
     Parse field dimensions and allocate symbols for array shape and strides.
 
@@ -116,18 +116,20 @@ def _make_array_shape_and_strides(
     tuple(shape, strides)
         The output tuple fields are arrays of dace symbolic expressions.
     """
-    dtype = dace.int64
-    sorted_dims = [dim for _, dim in get_sorted_dims(dims)] if sort_dims else dims
+    dtype = dace.int32
+    sorted_dims = get_sorted_dims(dims) if sort_dims else enumerate(dims)
     shape = [
         (
             neighbor_tables[dim.value].max_neighbors
             if dim.kind == DimensionKind.LOCAL
-            else dace.symbol(unique_name(f"{name}_shape{i}"), dtype)
+            # we reuse the same gt4py symbol for field size passed as scalar argument which is used in closure domain
+            else dace.symbol(f"__{name}_size_{i}", dtype)
         )
-        for i, dim in enumerate(sorted_dims)
+        for i, dim in sorted_dims
     ]
-    strides = [dace.symbol(unique_name(f"{name}_stride{i}"), dtype) for i, _ in enumerate(shape)]
-    return shape, strides
+    offset = [dace.symbol(unique_name(f"{name}_offset{i}"), dtype) for i, _ in sorted_dims]
+    strides = [dace.symbol(unique_name(f"{name}_stride{i}"), dtype) for i, _ in sorted_dims]
+    return shape, offset, strides
 
 
 def _check_no_lifts(node: itir.StencilClosure):
@@ -179,19 +181,24 @@ def add_storage(
         sort_dimensions: bool = True,
     ):
         if isinstance(type_, ts.FieldType):
-            shape, strides = _make_array_shape_and_strides(
+            shape, offset, strides = _make_array_shape_and_strides(
                 name, type_.dims, neighbor_tables, sort_dimensions
             )
-            offset = (
-                [dace.symbol(unique_name(f"{name}_offset{i}_")) for i in range(len(type_.dims))]
-                if has_offset
-                else None
-            )
             dtype = as_dace_type(type_.dtype)
-            sdfg.add_array(name, shape=shape, strides=strides, offset=offset, dtype=dtype)
+            sdfg.add_array(
+                name,
+                shape=shape,
+                strides=strides,
+                offset=(offset if has_offset else None),
+                dtype=dtype,
+            )
 
         elif isinstance(type_, ts.ScalarType):
-            sdfg.add_symbol(name, as_dace_type(type_))
+            dtype = as_dace_type(type_)
+            if name in sdfg.symbols:
+                assert sdfg.symbols[name].dtype == dtype
+            else:
+                sdfg.add_symbol(name, dtype)
 
         else:
             raise NotImplementedError()
@@ -429,7 +436,6 @@ def visit_StencilClosure(
         for name, type_ in self.storage_types.items():
             if isinstance(type_, ts.ScalarType):
                 dtype = as_dace_type(type_)
-                closure_sdfg.add_symbol(name, dtype)
                 if name in input_names:
                     out_name = unique_var_name()
                     closure_sdfg.add_scalar(out_name, dtype, transient=True)