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

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)