feat[next-dace]: remove canonical sorting based on node and connector labels

.pre-commit-config.yaml

@@ -117,3 +117,10 @@ repos:
         entry: uv run --group dev --frozen --isolated validate-pyproject pyproject.toml
         files: ^pyproject\.toml$
         pass_filenames: false
+
+      - <<: *uv-managed-hook
+        id: no-sorted-in-dace-backend
+        name: no sorted() in dace backend
+        entry: uv run --group dev --frozen --isolated python scripts/python/check_no_sorted_graph_traversal.py
+        files: ^src/gt4py/next/program_processors/runners/dace/.*\.py$
+        types_or: [python]

scripts/python/check_no_sorted_graph_traversal.py

@@ -0,0 +1,112 @@
+#!/usr/bin/env -S uv run -q --frozen --isolated --python 3.12 --group scripts python3
+#
+# 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
+
+"""Check that ``sorted()`` is not used in gt4py's DaCe backend.
+
+Using ``sorted()`` to obtain a "canonical" ordering is usually a sign that the
+underlying data structure has non-deterministic iteration order. DaCe graph
+methods such as ``in_edges()``, ``out_edges()`` and ``nodes()`` already return
+lists, connector names are already stored in dicts, and sets should be replaced
+by ``OrderedSet`` when order matters. Wrapping any of these in ``sorted()`` is
+therefore unnecessary and can hide non-determinism when the sort key
+accidentally depends on object identity, e.g. ``id()`` or ``repr()``.
+
+The allowed exceptions are:
+- sorting ``free_symbols``, because the set of symbols is a genuine mathematical
+  set and sorting it yields a stable order;
+- ``splitting_tools.py``, where ``sorted(split_description, key=str)`` is used
+  to obtain a deterministic processing order that cannot be derived from the
+  existing data structures.
+
+This script is intended as a pre-commit / CI check for the DaCe backend code.
+"""
+
+from __future__ import annotations
+
+import ast
+import sys
+from pathlib import Path
+
+
+# ``sorted()`` may only be used on these expressions.
+ALLOWED_SORTED_ARGUMENTS = frozenset({"free_symbols"})
+
+# Some files are allowed to use ``sorted()`` because their specific ordering
+# requirement cannot be satisfied by the existing data structures.
+EXEMPT_FILE_NAMES = frozenset({"splitting_tools.py"})
+
+
+def _is_allowed_sorted_argument(node: ast.AST) -> bool:
+    """Return True if ``sorted(node)`` is an allowed use case."""
+    arg_src = ast.unparse(node)
+    return any(allowed in arg_src for allowed in ALLOWED_SORTED_ARGUMENTS)
+
+
+def check_file(path: Path) -> list[str]:
+    """Return a list of violation messages for ``path``."""
+    if path.name in EXEMPT_FILE_NAMES:
+        return []
+
+    try:
+        source = path.read_text()
+    except (OSError, UnicodeDecodeError) as e:
+        return [f"{path}: error reading file: {e}"]
+
+    try:
+        tree = ast.parse(source)
+    except SyntaxError as e:
+        return [f"{path}:{e.lineno}: syntax error: {e.msg}"]
+
+    violations: list[str] = []
+    for node in ast.walk(tree):
+        if not (
+            isinstance(node, ast.Call)
+            and isinstance(node.func, ast.Name)
+            and node.func.id == "sorted"
+        ):
+            continue
+        if node.args and not _is_allowed_sorted_argument(node.args[0]):
+            arg_src = ast.unparse(node.args[0])
+            violations.append(f"{path}:{node.lineno}: sorted({arg_src})")
+
+    return violations
+
+
+def main(argv: list[str] | None = None) -> int:
+    if argv is None:
+        argv = sys.argv[1:]
+
+    if not argv:
+        print("Usage: check_no_sorted_graph_traversal.py <file>...", file=sys.stderr)
+        return 2
+
+    all_violations: list[str] = []
+    for arg in argv:
+        path = Path(arg)
+        if not path.is_file():
+            continue
+        all_violations.extend(check_file(path))
+
+    if all_violations:
+        print(
+            "Found sorted() calls in the DaCe backend. "
+            "Use OrderedSet or rely on the existing list/dict order instead; "
+            "only sorting free_symbols (and the exempt splitting_tools.py) is allowed:",
+            file=sys.stderr,
+        )
+        for violation in all_violations:
+            print(violation, file=sys.stderr)
+        return 1
+
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

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

@@ -29,6 +29,7 @@
 import dace
 from dace import nodes as dace_nodes, subsets as dace_subsets
 from dace.libraries import standard as dace_stdlib
+from ordered_set import OrderedSet
 
 from gt4py import eve
 from gt4py.eve.extended_typing import MaybeNestedInTuple, NestedTuple
@@ -461,8 +462,8 @@ def _add_map(
     def _add_tasklet(
         self,
         name: str,
-        inputs: set[str] | Mapping[str, dace.dtypes.typeclass | None],
-        outputs: set[str] | Mapping[str, dace.dtypes.typeclass | None],
+        inputs: OrderedSet[str] | Mapping[str, dace.dtypes.typeclass | None],
+        outputs: OrderedSet[str] | Mapping[str, dace.dtypes.typeclass | None],
         code: str,
         **kwargs: Any,
     ) -> tuple[dace_nodes.Tasklet, dict[str, str]]:
@@ -631,8 +632,8 @@ def _visit_deref(self, node: gtir.FunCall) -> DataExpr:
         )
         deref_node, connector_mapping = self._add_tasklet(
             name="deref",
-            inputs={"field"} | set(index_connectors),
-            outputs={"val"},
+            inputs={"field": None} | {conn: None for conn in index_connectors},
+            outputs={"val": None},
             code=f"val = field[{index_internals}]",
         )
         # add new termination point for the field parameter
@@ -989,7 +990,7 @@ def write_output_of_nested_sdfg_to_temporary(inner_value: ValueExpr) -> ValueExp
         else:
             result = outer_value
 
-        outputs = {outval.dc_node.data for outval in gtx_utils.flatten_nested_tuple((result,))}
+        outputs = [outval.dc_node.data for outval in gtx_utils.flatten_nested_tuple((result,))]
 
         # map the connectivities that were used inside the nested SDFG
         used_connectivities = [
@@ -1004,8 +1005,8 @@ def write_output_of_nested_sdfg_to_temporary(inner_value: ValueExpr) -> ValueExp
             nsdfg_symbols_mapping["__cond"] = condition_value.value
         nsdfg_node = self.state.add_nested_sdfg(
             nsdfg,
-            inputs={key: None for key in sorted(used_connectivities | input_memlets.keys())},
-            outputs={key: None for key in sorted(outputs)},
+            inputs={key: None for key in [*used_connectivities, *input_memlets.keys()]},
+            outputs={key: None for key in outputs},
             symbol_mapping=nsdfg_symbols_mapping,
         )
 
@@ -1207,8 +1208,8 @@ def _visit_list_get(self, node: gtir.FunCall) -> ValueExpr:
         elif isinstance(index_arg, ValueExpr):
             tasklet_node, connector_mapping = self._add_tasklet(
                 name="list_get",
-                inputs={"index", "data"},
-                outputs={"val"},
+                inputs={"index": None, "data": None},
+                outputs={"val": None},
                 code="val = data[index]",
             )
             self._add_edge(
@@ -1516,22 +1517,22 @@ def _make_cartesian_shift(
             if isinstance(index_expr, SymbolExpr):
                 dynamic_offset_tasklet, connector_mapping = self._add_tasklet(
                     name="dynamic_offset",
-                    inputs={"offset"},
-                    outputs={new_index_connector},
+                    inputs={"offset": None},
+                    outputs={new_index_connector: None},
                     code=f"{new_index_connector} = {index_expr.value} + offset",
                 )
             elif isinstance(offset_expr, SymbolExpr):
                 dynamic_offset_tasklet, connector_mapping = self._add_tasklet(
                     name="dynamic_offset",
-                    inputs={"index"},
-                    outputs={new_index_connector},
+                    inputs={"index": None},
+                    outputs={new_index_connector: None},
                     code=f"{new_index_connector} = index + {offset_expr}",
                 )
             else:
                 dynamic_offset_tasklet, connector_mapping = self._add_tasklet(
                     name="dynamic_offset",
-                    inputs={"index", "offset"},
-                    outputs={new_index_connector},
+                    inputs={"index": None, "offset": None},
+                    outputs={new_index_connector: None},
                     code=f"{new_index_connector} = index + offset",
                 )
             for input_expr, input_connector in [
@@ -1587,8 +1588,8 @@ def _make_dynamic_neighbor_offset(
         new_index_connector = "neighbor_index"
         tasklet_node, connector_mapping = self._add_tasklet(
             "dynamic_neighbor_offset",
-            {"table", "offset"},
-            {new_index_connector},
+            {"table": None, "offset": None},
+            {new_index_connector: None},
             f"{new_index_connector} = table[{origin_index.value}, offset]",
         )
         self._add_input_data_edge(
@@ -1723,8 +1724,8 @@ def _visit_generic_builtin(self, node: gtir.FunCall) -> ValueExpr:
         out_connector = "result"
         tasklet_node, connector_mapping = self._add_tasklet(
             name=builtin_name,
-            inputs=set(node_connections.keys()),
-            outputs={out_connector},
+            inputs={conn: None for conn in node_connections.keys()},
+            outputs={out_connector: None},
             code="{} = {}".format(out_connector, code),
         )
 
@@ -1841,7 +1842,10 @@ def _visit_Lambda_impl(
                 # special case where the field operator is simply copying data from source to destination node
                 output_dtype = output_expr.dc_node.desc(self.sdfg).dtype
                 tasklet_node, connector_mapping = self._add_tasklet(
-                    name="copy", inputs={"inp"}, outputs={"out"}, code="out = inp"
+                    name="copy",
+                    inputs={"inp": None},
+                    outputs={"out": None},
+                    code="out = inp",
                 )
                 self._add_input_data_edge(
                     output_expr.dc_node,
@@ -1853,7 +1857,10 @@ def _visit_Lambda_impl(
                 # even simpler case, where a constant value is written to destination node
                 output_dtype = output_expr.dc_dtype
                 tasklet_node, connector_mapping = self._add_tasklet(
-                    name="write", inputs={}, outputs={"out"}, code=f"out = {output_expr.value}"
+                    name="write",
+                    inputs={},
+                    outputs={"out": None},
+                    code=f"out = {output_expr.value}",
                 )
 
             output_expr = self._construct_tasklet_result(

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

@@ -21,6 +21,7 @@
 import dace
 from dace import nodes as dace_nodes, subsets as dace_subsets
 from dace.frontend.python import astutils as dace_astutils
+from ordered_set import OrderedSet
 
 from gt4py import eve
 from gt4py.eve import concepts
@@ -116,8 +117,8 @@ def add_tasklet(
         name: str,
         sdfg: dace.SDFG,
         state: dace.SDFGState,
-        inputs: set[str] | Mapping[str, dace.dtypes.typeclass | None],
-        outputs: set[str] | Mapping[str, dace.dtypes.typeclass | None],
+        inputs: OrderedSet[str] | Mapping[str, dace.dtypes.typeclass | None],
+        outputs: OrderedSet[str] | Mapping[str, dace.dtypes.typeclass | None],
         code: str,
         language: dace.dtypes.Language = dace.dtypes.Language.Python,
         **kwargs: Any,
@@ -132,10 +133,10 @@ def add_tasklet(
             The created tasklet node and the mapping from original connector names to
             modified connector names.
         """
-        if isinstance(inputs, set):
-            inputs = {k: None for k in sorted(inputs)}
-        if isinstance(outputs, set):
-            outputs = {k: None for k in sorted(outputs)}
+        if not isinstance(inputs, Mapping):
+            inputs = {k: None for k in inputs}
+        if not isinstance(outputs, Mapping):
+            outputs = {k: None for k in outputs}
         assert inputs.keys().isdisjoint(outputs.keys())
 
         connector_mapping = {
@@ -627,7 +628,7 @@ def setup_nested_context(
 
         # Sorting the parameter list in alphabetical order to improve determinism.
         input_params = [
-            gtir.Sym(id=name, type=lambda_symbols[name]) for name in sorted(lambda_symbols.keys())
+            gtir.Sym(id=name, type=lambda_symbols[name]) for name in lambda_symbols.keys()
         ]
 
         nsdfg = dace.SDFG(name=self.unique_nsdfg_name(sdfg_name))
@@ -684,11 +685,11 @@ def add_nested_sdfg(
         )
         # The output connectors only need to be setup for the actual result of the
         # internal dataflow that writes to some sink data nodes of the nested SDFG.
-        lambda_outputs = {
+        lambda_outputs = [
             dataname
             for output in lambda_output_data
             if output is not None and (dataname := output.dc_node.data) not in data_args
-        }
+        ]
 
         connectivity_arrays = {
             gtx_dace_args.connectivity_identifier(offset)
@@ -766,8 +767,8 @@ def add_nested_sdfg(
 
         nsdfg_node = outer_ctx.state.add_nested_sdfg(
             inner_ctx.sdfg,
-            inputs={key: None for key in sorted(input_memlets.keys())},
-            outputs={key: None for key in sorted(lambda_outputs)},
+            inputs={key: None for key in input_memlets.keys()},
+            outputs={key: None for key in lambda_outputs},
             symbol_mapping=nsdfg_symbols_mapping,
             debuginfo=gtir_to_sdfg_utils.debug_info(node, default=outer_ctx.sdfg.debuginfo),
         )

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

@@ -481,7 +481,7 @@ def translate_index(
         sdfg=ctx.sdfg,
         state=ctx.state,
         inputs={},
-        outputs={"val"},
+        outputs={"val": None},
         code=f"val = {dim_index}",
     )
     ctx.state.add_edge(
@@ -547,7 +547,7 @@ def _get_symbolic_value(
         sdfg=sdfg,
         state=state,
         inputs={},
-        outputs={"out"},
+        outputs={"out": None},
         code=f"out = {symbolic_expr}",
     )
     temp_name, _ = sdfg.add_scalar(
@@ -682,8 +682,8 @@ def translate_scalar_expr(
         name="scalar_expr",
         sdfg=ctx.sdfg,
         state=ctx.state,
-        inputs=set(connectors),
-        outputs={"out"},
+        inputs={connector: None for connector in connectors},
+        outputs={"out": None},
         code=f"out = {python_code}",
     )
     # create edges for the input data connectors

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

@@ -569,9 +569,9 @@ def _handle_dataflow_result_of_nested_sdfg(
         inner_ctx.sdfg.arrays[inner_dataname] = inner_desc
         # We write the result to the output field after the loop region has ended.
         if len(inner_ctx.sdfg.states()) == 3:
-            assert sorted(st.label for st in inner_ctx.sdfg.states()) == [
-                "scan_compute",
+            assert [st.label for st in inner_ctx.sdfg.states()] == [
                 "scan_entry",
+                "scan_compute",
                 "scan_update",
             ]
             scan_loop = next(
@@ -581,11 +581,11 @@ def _handle_dataflow_result_of_nested_sdfg(
             )
             last_level_state = inner_ctx.sdfg.add_state_after(scan_loop, "scan_last_level")
         else:
-            assert sorted(st.label for st in inner_ctx.sdfg.states()) == [
-                "scan_compute",
+            assert [st.label for st in inner_ctx.sdfg.states()] == [
                 "scan_entry",
-                "scan_last_level",
+                "scan_compute",
                 "scan_update",
+                "scan_last_level",
             ]
             last_level_state = next(
                 s for s in inner_ctx.sdfg.states() if s.label == "scan_last_level"