diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 521e577d58..764552b047 100644 --- a/.pre-commit-config.yaml +++ b/.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] diff --git a/scripts/python/check_no_sorted_graph_traversal.py b/scripts/python/check_no_sorted_graph_traversal.py new file mode 100644 index 0000000000..c0fc7a91ad --- /dev/null +++ b/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=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()) diff --git a/src/gt4py/next/program_processors/runners/dace/lowering/gtir_dataflow.py b/src/gt4py/next/program_processors/runners/dace/lowering/gtir_dataflow.py index 36a48ae119..0e73f77942 100644 --- a/src/gt4py/next/program_processors/runners/dace/lowering/gtir_dataflow.py +++ b/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( diff --git a/src/gt4py/next/program_processors/runners/dace/lowering/gtir_to_sdfg.py b/src/gt4py/next/program_processors/runners/dace/lowering/gtir_to_sdfg.py index 8a8d42f9ba..8617f7bf3d 100644 --- a/src/gt4py/next/program_processors/runners/dace/lowering/gtir_to_sdfg.py +++ b/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), ) diff --git a/src/gt4py/next/program_processors/runners/dace/lowering/gtir_to_sdfg_primitives.py b/src/gt4py/next/program_processors/runners/dace/lowering/gtir_to_sdfg_primitives.py index a0d09cab1f..d15d39de6b 100644 --- a/src/gt4py/next/program_processors/runners/dace/lowering/gtir_to_sdfg_primitives.py +++ b/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 diff --git a/src/gt4py/next/program_processors/runners/dace/lowering/gtir_to_sdfg_scan.py b/src/gt4py/next/program_processors/runners/dace/lowering/gtir_to_sdfg_scan.py index fddcb080b8..7ff584b51d 100644 --- a/src/gt4py/next/program_processors/runners/dace/lowering/gtir_to_sdfg_scan.py +++ b/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" diff --git a/src/gt4py/next/program_processors/runners/dace/lowering/gtir_to_sdfg_utils.py b/src/gt4py/next/program_processors/runners/dace/lowering/gtir_to_sdfg_utils.py index 2b55baf5b6..b36f7784eb 100644 --- a/src/gt4py/next/program_processors/runners/dace/lowering/gtir_to_sdfg_utils.py +++ b/src/gt4py/next/program_processors/runners/dace/lowering/gtir_to_sdfg_utils.py @@ -11,6 +11,7 @@ from typing import Dict, Final, Mapping, Optional, TypeVar import dace +from ordered_set import OrderedSet from gt4py import eve from gt4py.eve.extended_typing import NestedTuple @@ -136,14 +137,14 @@ def visit_SymRef(self, node: gtir.SymRef, *, symtable: Dict[str, str]) -> gtir.S f"Unexpectd symbol with prefix '{_TASKLET_CONNECTOR_PREFIX}' in program parameters." ) - ir_sym_ids = {str(sym.id) for sym in eve.walk_values(ir).if_isinstance(gtir.Sym).to_set()} + ir_sym_ids = OrderedSet( + eve.walk_values(ir).if_isinstance(gtir.Sym).map(lambda x: str(x.id)).to_list() + ) ir_ssa_uuid = eve.utils.SequentialIDGenerator(prefix="gtir_var") # note: traverse in alphabetical order to generate UIDs in deterministic way invalid_symbols_mapping = { - sym_id: next(ir_ssa_uuid) - for sym_id in sorted(ir_sym_ids) - if not dace.dtypes.validate_name(sym_id) + sym_id: next(ir_ssa_uuid) for sym_id in ir_sym_ids if not dace.dtypes.validate_name(sym_id) } if len(invalid_symbols_mapping) == 0: return ir diff --git a/src/gt4py/next/program_processors/runners/dace/transformations/concat_where_mapper.py b/src/gt4py/next/program_processors/runners/dace/transformations/concat_where_mapper.py index c8a1a120f3..f70e12928c 100644 --- a/src/gt4py/next/program_processors/runners/dace/transformations/concat_where_mapper.py +++ b/src/gt4py/next/program_processors/runners/dace/transformations/concat_where_mapper.py @@ -293,7 +293,6 @@ def gt_apply_concat_where_replacement_on_sdfg( found_nsdfgs.append((state, node)) if len(suitable_concat_nodes) > 0: - suitable_concat_nodes = sorted(suitable_concat_nodes, key=lambda x: (repr(x[0]), x[1].data)) for state, concat_node in suitable_concat_nodes: nb_applies += gt_replace_concat_where_node( sdfg=sdfg, @@ -305,7 +304,6 @@ def gt_apply_concat_where_replacement_on_sdfg( sdfg.validate() if len(found_nsdfgs) > 0: - found_nsdfgs = sorted(found_nsdfgs, key=lambda x: (repr(x[0]), str(x[1]))) for _, nsdfg in found_nsdfgs: nb_applies += gt_apply_concat_where_replacement_on_sdfg( sdfg=nsdfg.sdfg, @@ -482,7 +480,7 @@ def _setup_initial_producer_description_on_top_level( ) assert all(str(fs) in sdfg.symbols for fs in initial_producer_specs[-1].free_symbols) - return sorted(initial_producer_specs) + return initial_producer_specs def _process_descending_points_of_state( @@ -867,10 +865,9 @@ def _map_data_into_nested_scopes( # its ancestors scopes too. We will bring them in a deterministic order before # process them, handling missing parent scopes on the fly, see # `_map_data_into_nested_scopes_impl()`. - scopes_containing_consumers: list[_ScopeLocation] = sorted( - {scope_dict[consumer_spec.consumer] for consumer_spec in consumer_specs}, - key=lambda scope: "NONE" if scope is None else str(scope), - ) + scopes_containing_consumers: list[_ScopeLocation] = [ + scope_dict[consumer_spec.consumer] for consumer_spec in consumer_specs + ] for scope in scopes_containing_consumers: _map_data_into_nested_scopes_impl( state=state, @@ -1327,7 +1324,7 @@ def _find_consumer_specs_single_source_single_level( if for_check: return stairs_to_deeper_levels - return sorted(consumer_specs), sorted(stairs_to_deeper_levels) + return consumer_specs, stairs_to_deeper_levels def _handle_special_case_of_gt4py_scan_point( diff --git a/src/gt4py/next/program_processors/runners/dace/transformations/dead_dataflow_elimination.py b/src/gt4py/next/program_processors/runners/dace/transformations/dead_dataflow_elimination.py index 39dd321644..a8e99343fc 100644 --- a/src/gt4py/next/program_processors/runners/dace/transformations/dead_dataflow_elimination.py +++ b/src/gt4py/next/program_processors/runners/dace/transformations/dead_dataflow_elimination.py @@ -12,6 +12,7 @@ from dace import properties as dace_properties, transformation as dace_transformation from dace.sdfg import nodes as dace_nodes from dace.transformation.passes import analysis as dace_analysis +from ordered_set import OrderedSet from gt4py.next.program_processors.runners.dace import transformations as gtx_transformations @@ -189,9 +190,10 @@ def gt_remove_map( # Now find all notes that are producer or consumer of the Map, after we removed # the nodes of the Maps we need to check if they have become isolated. - # NOTE: Needs to be a `set` to handle multiple connections with the MapEntry node. - adjacent_nodes: set[dace_nodes.AccessNode] = {iedge.src for iedge in state.in_edges(map_entry)} - adjacent_nodes.update(oedge.dst for oedge in state.out_edges(map_exit)) + # NOTE: Use an OrderedSet to handle multiple connections with the MapEntry node + # while preserving deterministic iteration order. + adjacent_nodes = OrderedSet(iedge.src for iedge in state.in_edges(map_entry)) + adjacent_nodes.update([oedge.dst for oedge in state.out_edges(map_exit)]) if not all(isinstance(ac, dace_nodes.AccessNode) for ac in adjacent_nodes): raise ValueError( @@ -208,8 +210,7 @@ def gt_remove_map( state.remove_nodes_from(map_scope.nodes()) # Now check for potentially isolated nodes. - # Process them in deterministic order. - for adjacent_node in sorted(adjacent_nodes, key=lambda ac: ac.data): + for adjacent_node in adjacent_nodes: if state.degree(adjacent_node) == 0: map_scope_datas.add(adjacent_node.data) state.remove_node(adjacent_node) diff --git a/src/gt4py/next/program_processors/runners/dace/transformations/map_fusion_utils.py b/src/gt4py/next/program_processors/runners/dace/transformations/map_fusion_utils.py index 9cd4368415..5a74d95a68 100644 --- a/src/gt4py/next/program_processors/runners/dace/transformations/map_fusion_utils.py +++ b/src/gt4py/next/program_processors/runners/dace/transformations/map_fusion_utils.py @@ -78,8 +78,8 @@ def _new_name(old_name: str) -> str: elif isinstance(node, dace_nodes.NestedSDFG): node_ = graph.add_nested_sdfg( sdfg=copy.deepcopy(node.sdfg), - inputs={k: None for k in sorted(node.in_connectors.keys())}, - outputs={k: None for k in sorted(node.out_connectors.keys())}, + inputs={k: None for k in node.in_connectors.keys()}, + outputs={k: None for k in node.out_connectors.keys()}, symbol_mapping=node.symbol_mapping.copy(), debuginfo=copy.copy(node.debuginfo), ) @@ -202,19 +202,18 @@ def split_overlapping_map_range( Two lists, each containing the ranges corresponding to the splitted range for the first and the second map, respectively. """ - first_map_params = set(first_map.params) - second_map_params = set(second_map.params) - if first_map_params != second_map_params: + if set(first_map.params) != set(second_map.params): return None first_map_dict = dict(zip(first_map.params, first_map.range.ranges, strict=True)) second_map_dict = dict(zip(second_map.params, second_map.range.ranges, strict=True)) + # Follow order of parameters as they appear in the first map. first_map_sorted_range = dace_subsets.Range( - [first_map_dict[param] for param in sorted(first_map_params)] + [first_map_dict[param] for param in first_map.params] ) second_map_sorted_range = dace_subsets.Range( - [second_map_dict[param] for param in sorted(second_map_params)] + [second_map_dict[param] for param in first_map.params] ) if gtx_dace_split.never_intersecting(first_map_sorted_range, second_map_sorted_range): diff --git a/src/gt4py/next/program_processors/runners/dace/transformations/move_dataflow_into_if_body.py b/src/gt4py/next/program_processors/runners/dace/transformations/move_dataflow_into_if_body.py index 130a55d245..39a94f6e17 100644 --- a/src/gt4py/next/program_processors/runners/dace/transformations/move_dataflow_into_if_body.py +++ b/src/gt4py/next/program_processors/runners/dace/transformations/move_dataflow_into_if_body.py @@ -304,7 +304,7 @@ def _replicate_dataflow_into_branch( # aliasing, i.e. multiple inner names refer to the same outer name. # TODO(phimuell): Investigate if it would be better if we handle partially # mapped in data such by fully map it in and perform the slicing outside. - fully_mapped_in_data: dict[str, set[str]] = collections.defaultdict(set) + fully_mapped_in_data: dict[str, OrderedSet[str]] = collections.defaultdict(OrderedSet) for if_iedge in state.in_edges(if_block): if if_iedge.data.is_empty(): continue @@ -385,18 +385,13 @@ def _replicate_dataflow_into_branch( # `branch_state`. We first look if the state contains an AccessNode # referring to that data. outer_aliases = fully_mapped_in_data[outer_data] - candidate_nodes: list[dace_nodes.AccessNode] = sorted( - ( - dnode - for dnode in branch_state.data_nodes() - if dnode.data in outer_aliases - ), - key=lambda dnode: dnode.data, - ) + candidate_nodes: list[dace_nodes.AccessNode] = [ + dnode for dnode in branch_state.data_nodes() if dnode.data in outer_aliases + ] if len(candidate_nodes) == 0: # There is no AccessNode in the state so we have to create one. - inner_data = sorted(outer_aliases)[0] + inner_data = next(iter(outer_aliases)) inner_node = branch_state.add_access(inner_data) else: @@ -410,7 +405,8 @@ def _replicate_dataflow_into_branch( if len(candidate_source_nodes) != len(candidate_nodes): raise NotImplementedError() - # We take the first node, since they are sorted it is deterministic. + # We take the first node; `data_nodes()` returns a list, + # so the order is already deterministic. inner_node = candidate_source_nodes[0] # A different AccessNode object for the same outer data may have @@ -550,10 +546,9 @@ def _update_symbol_mapping( are available in the parent SDFG. """ symbol_mapping = if_block.symbol_mapping - missing_symbols = sorted( - (ms for ms in if_block.sdfg.free_symbols if ms not in symbol_mapping), - key=lambda sym: str(sym), - ) + missing_symbols = [ + ms for ms in sorted(if_block.sdfg.free_symbols) if ms not in symbol_mapping + ] symbol_mapping.update({s: s for s in missing_symbols}) if_block.symbol_mapping = symbol_mapping # Performs conversion. diff --git a/src/gt4py/next/program_processors/runners/dace/transformations/split_access_nodes.py b/src/gt4py/next/program_processors/runners/dace/transformations/split_access_nodes.py index e55ff4d1b7..475f0fd908 100644 --- a/src/gt4py/next/program_processors/runners/dace/transformations/split_access_nodes.py +++ b/src/gt4py/next/program_processors/runners/dace/transformations/split_access_nodes.py @@ -102,14 +102,11 @@ def _apply_split_access_node_non_recursive( # We can only split single use data that is also transient. Because GT4Py uses # an SSA style we know that there is only one AccessNode that refers to that # data. Thus, all AccessNodes that are stored refer to different data - access_nodes_to_process = sorted( - ( - dnode - for dnode in state.data_nodes() - if dnode.data in single_use_data and scope_dict[dnode] is None - ), - key=lambda dnode: dnode.data, - ) + access_nodes_to_process = [ + dnode + for dnode in state.data_nodes() + if dnode.data in single_use_data and scope_dict[dnode] is None + ] assert len(access_nodes_to_process) == len( set(map(lambda ac: ac.data, access_nodes_to_process)) ) 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 f80809bec5..a09ccd07cf 100644 --- a/src/gt4py/next/program_processors/runners/dace/transformations/utils.py +++ b/src/gt4py/next/program_processors/runners/dace/transformations/utils.py @@ -802,7 +802,7 @@ def gt_data_descriptor_mapping( nsdfg: The nested SDFG node we want to process. only_fully_mapped: Only look at the fully mapped data. only_inputs: Only consider the data that are used as inputs. - only_inputs: Only consider the data that are used as outputs. + only_outputs: Only consider the data that are used as outputs. """ assert not (only_inputs and only_outputs) name_mapping: dict[str, str] = {} @@ -811,11 +811,9 @@ def gt_data_descriptor_mapping( # When we have to return both, we start with the input such that the output # descriptors are dominant. if not only_outputs: - iedges = sorted( - (iedge for iedge in state.in_edges(nsdfg) if not iedge.data.is_empty()), - key=lambda iedge: iedge.dst_conn, - ) - for iedge in iedges: + for iedge in state.in_edges(nsdfg): + if iedge.data.is_empty(): + continue data_outside = iedge.data.data data_inside = iedge.dst_conn if only_fully_mapped and ( @@ -827,11 +825,9 @@ def gt_data_descriptor_mapping( if only_inputs: return name_mapping - oedges = sorted( - (oedge for oedge in state.out_edges(nsdfg) if not oedge.data.is_empty()), - key=lambda oedge: iedge.src_conn, - ) - for oedge in oedges: + for oedge in state.out_edges(nsdfg): + if oedge.data.is_empty(): + continue data_outside = oedge.data.data data_inside = oedge.src_conn if only_fully_mapped and (