diff --git a/ci/cscs-ci-dace-determinism.yml b/ci/cscs-ci-dace-determinism.yml index 289b738b60..7e80e03e70 100644 --- a/ci/cscs-ci-dace-determinism.yml +++ b/ci/cscs-ci-dace-determinism.yml @@ -64,7 +64,6 @@ build_cscs_amd_rocm: USE_MPI: 0 # TODO(havogt): to workaround the libfabric hook injecting incompatible libraries SLURM_JOB_NUM_NODES: 1 SLURM_TIMELIMIT: 60 - allow_failure: true artifacts: when: always paths: diff --git a/docs/development/dace_codegen_reproducability.md b/docs/development/dace_codegen_reproducability.md new file mode 100644 index 0000000000..a335a694c1 --- /dev/null +++ b/docs/development/dace_codegen_reproducability.md @@ -0,0 +1,95 @@ +# Debugging indeterministic behavior of dace transformations + +- Enable printing each transformation step, e.g. using + ``` + dace.Config.set("progress", value=True) + ``` + TODO: introduce new config var that prints the hash instead of hard-coding it. +- Execute the program in question twice and compare the output. +- Set a conditonal breakpoint in beginning of the `apply` method of the first pass where the SDFG + hash changes with condition `sdfg.hash_sdfg() == `. + Note: In case running the previous passes takes a long time it makes sense to serialize the SDFG + to json (`sdfg.to_json("sdfg(1|2).json")`) and loading it again (see debug script below) to + ease debugging. In rare cases the serializing and deserializing the sdfg changes the hash. In such + cases this trick doesn't work and the first location where the hash changes might not be the exact + location where the indeterministic behavior is. It helps to use a different hash, e.g. + `content_hash`, but this should be solved in general. + Note: It makes sense to also place a breakpoint after `DaceTranslator.generate_sdfg` to recognize + when all executions finished. +- When the location is found it is usually easy to spot the origin of the indeterminism. Often + there is a set operation or a symbol is named in an indeterministic way. Use ordered sets and + deterministic symbol names. + +## Appendix + +__Debugging sdfg autooptimize__ + +Usage `python debug_auto_optimize_sdfg.py sdfg1.json` + +```python +import pickle +import sys + +import dace +import json + +from dace import SDFG + +from gt4py.next.program_processors.runners.dace import ( + lowering as gtx_dace_lowering, + sdfg_args as gtx_dace_args, + transformations as gtx_transformations, +) +from dace.utils import print_sdfg_hash + +file = sys.argv[1] + +with open(file) as f: + data = json.load(f) + sdfg = dace.SDFG.from_json(data) + print_sdfg_hash(sdfg) + + gtx_transformations.gt_auto_optimize( + sdfg, + gpu=False, + constant_symbols={}, + unit_strides_kind=None, + ) +``` + +__Debugging single sdfg transform__ + +Usage `python debug_single_sdfg_transform.py sdfg1.json` + +```python +import pickle +import sys + +import dace +import json + +from dace import SDFG + +from gt4py.next.program_processors.runners.dace import ( + lowering as gtx_dace_lowering, + sdfg_args as gtx_dace_args, + transformations as gtx_transformations, +) +from dace.utils import print_sdfg_hash + +transformation = gtx_transformations.MoveDataflowIntoIfBody +file = sys.argv[1] + +with open(file) as f: + data = json.load(f) + sdfg = dace.SDFG.from_json(data) + print_sdfg_hash(sdfg) + + sdfg.apply_transformations_repeated( + transformation( + ignore_upstream_blocks=False, + ), + validate=False, + validate_all=True, + ) +``` diff --git a/pyproject.toml b/pyproject.toml index eb27848426..6c8c23f688 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -122,7 +122,8 @@ dependencies = [ 'toolz>=0.12.1', 'typing-extensions>=4.12.0', 'versioningit>=3.1.1', - 'xxhash>=3.5.0' + 'xxhash>=3.5.0', + 'ordered-set>=4.0.0' ] description = 'Python library for generating high-performance implementations of stencil kernels for weather and climate modeling from a domain-specific language (DSL)' dynamic = ['version'] @@ -479,6 +480,7 @@ url = 'https://gridtools.github.io/pypi/' # dace = {index = "gridtools"} [tool.uv.sources] atlas4py = {index = "test.pypi"} +dace = {git = "https://github.com/GridTools/dace", branch = "dace_toolchain_deterministic"} # -- versioningit -- [tool.versioningit] diff --git a/src/gt4py/next/program_processors/runners/dace/transformations/auto_optimize.py b/src/gt4py/next/program_processors/runners/dace/transformations/auto_optimize.py index a7b34f6bbf..30e4740bdd 100644 --- a/src/gt4py/next/program_processors/runners/dace/transformations/auto_optimize.py +++ b/src/gt4py/next/program_processors/runners/dace/transformations/auto_optimize.py @@ -18,6 +18,7 @@ from dace.transformation import dataflow as dace_dataflow from dace.transformation.auto import auto_optimize as dace_aoptimize from dace.transformation.passes import analysis as dace_analysis +from dace.utils import print_sdfg_hash from gt4py.next import common as gtx_common, utils as gtx_utils from gt4py.next.program_processors.runners.dace import ( @@ -254,12 +255,14 @@ def gt_auto_optimize( # Initial Cleanup # NOTE: The initial simplification stage must be synchronized with the one that # `gt_substitute_compiletime_symbols()` performs! + print_sdfg_hash(sdfg) gtx_transformations.gt_simplify( sdfg=sdfg, validate=False, skip=gtx_transformations.constants._GT_AUTO_OPT_INITIAL_STEP_SIMPLIFY_SKIP_LIST, validate_all=validate_all, ) + print_sdfg_hash(sdfg) if constant_symbols: gtx_transformations.gt_substitute_compiletime_symbols( @@ -271,6 +274,7 @@ def gt_auto_optimize( validate=False, validate_all=validate_all, ) + print_sdfg_hash(sdfg) # Demote the fields. # Actually they should probably be at the very start of this function, however, @@ -304,10 +308,12 @@ def gt_auto_optimize( skip=gtx_transformations.constants._GT_AUTO_OPT_INITIAL_STEP_SIMPLIFY_SKIP_LIST, validate_all=validate_all, ) + print_sdfg_hash(sdfg) gtx_transformations.gt_reduce_distributed_buffering( sdfg, validate=False, validate_all=validate_all ) + print_sdfg_hash(sdfg) # Process top level Maps sdfg = _gt_auto_process_top_level_maps( @@ -317,12 +323,14 @@ def gt_auto_optimize( optimization_hooks=optimization_hooks, validate_all=validate_all, ) + print_sdfg_hash(sdfg) # We now ensure that point wise computations are properly double buffered, # this ensures that rule 3 of ADR-18 is maintained. # TODO(phimuell): Figuring out if it is important to do it before the inner # Map optimization. I think it is, especially when we apply `LoopBlocking`. gtx_transformations.gt_create_local_double_buffering(sdfg) + print_sdfg_hash(sdfg) # Optimize the interior of the Maps: sdfg = _gt_auto_process_dataflow_inside_maps( @@ -336,6 +344,7 @@ def gt_auto_optimize( validate_all=validate_all, uids=uids, ) + print_sdfg_hash(sdfg) # Configure the Maps: # Will also perform the GPU transformation. @@ -395,6 +404,7 @@ def gt_auto_optimize( gpu_block_size_spec=gpu_block_size_spec if gpu_block_size_spec else None, validate_all=validate_all, ) + print_sdfg_hash(sdfg) # Transients sdfg = _gt_auto_post_processing( @@ -409,6 +419,7 @@ def gt_auto_optimize( gpu_memory_pool=gpu_memory_pool, validate_all=validate_all, ) + print_sdfg_hash(sdfg) # Canonicalize the SDFG. This ensures that the code generator will see SDFGs # that conform to the historical expected version. 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 8742f1dc6a..f9be88dee5 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 @@ -594,7 +594,7 @@ def _check_for_data_and_symbol_conflicts( self, sdfg: dace.SDFG, state: dace.SDFGState, - relocatable_dataflow: set[dace_nodes.Node], + relocatable_dataflow: OrderedSet[dace_nodes.Node], if_block: dace_nodes.NestedSDFG, enclosing_map: dace_nodes.MapEntry, ) -> bool: @@ -724,7 +724,7 @@ def _filter_relocatable_dataflow( non_relocatable_dataflow: dict[str, OrderedSet[dace_nodes.Node]], connector_usage_location: dict[str, tuple[dace.SDFGState, dace_nodes.AccessNode]], enclosing_map: dace_nodes.MapEntry, - ) -> set[dace_nodes.Node]: + ) -> OrderedSet[dace_nodes.Node]: """Compute the final set of the relocatable nodes. The function expects the dataflow that is upstream of every connector @@ -749,8 +749,8 @@ def _filter_relocatable_dataflow( """ # These are the nodes that can not be relocated anyway. - all_non_relocatable_dataflow: set[dace_nodes.Node] = functools.reduce( - lambda s1, s2: s1.union(s2), non_relocatable_dataflow.values(), set() + all_non_relocatable_dataflow: OrderedSet[dace_nodes.Node] = functools.reduce( + lambda s1, s2: s1.union(s2), non_relocatable_dataflow.values(), OrderedSet() ) # While we can relocate nodes that are needed by multiple connectors, we can @@ -782,8 +782,8 @@ def _filter_relocatable_dataflow( # process all of them together. We do this because a node can be associated to # multiple connectors and as such data dependencies can show up. We will, # after the filtering distribute them back. - nodes_proposed_for_reloc: set[dace_nodes.Node] = functools.reduce( - lambda s1, s2: s1.union(s2), raw_relocatable_dataflow.values(), set() + nodes_proposed_for_reloc: OrderedSet[dace_nodes.Node] = functools.reduce( + lambda s1, s2: s1.union(s2), raw_relocatable_dataflow.values(), OrderedSet() ) # Filtering out all nodes that can not be relocated anyway. @@ -874,8 +874,8 @@ def _partition_if_block( if len(if_block.out_connectors.keys()) == 0: return None - input_names: set[str] = set(if_block.in_connectors.keys()) - output_names: set[str] = set(if_block.out_connectors.keys()) + input_names: OrderedSet[str] = OrderedSet(if_block.in_connectors.keys()) + output_names: OrderedSet[str] = OrderedSet(if_block.out_connectors.keys()) # If data is used as input and output we ignore it. # TODO(phimuell): Think if this case can be handled. @@ -895,7 +895,7 @@ def _partition_if_block( connector_usage_location: dict[str, tuple[dace.SDFGState, dace_nodes.AccessNode]] = {} # This is the dataflow that can not be relocated. - non_relocatable_connectors: set[str] = set() + non_relocatable_connectors: OrderedSet[str] = OrderedSet() # Now inspect all states. for _, if_branch in inner_if_block.branches: @@ -943,7 +943,7 @@ def _partition_if_block( # In addition to the non relocatable connectors that were found above, we also # mark all connectors that were not found as non relocatable. non_relocatable_connectors.update( - conn for conn in input_names if conn not in connector_usage_location + [conn for conn in input_names if conn not in connector_usage_location] ) # We require that at least one non relocatable dataflow is there, this is for diff --git a/src/gt4py/next/program_processors/runners/dace/transformations/simplify.py b/src/gt4py/next/program_processors/runners/dace/transformations/simplify.py index 4361c5f1b3..d2c95f69a4 100644 --- a/src/gt4py/next/program_processors/runners/dace/transformations/simplify.py +++ b/src/gt4py/next/program_processors/runners/dace/transformations/simplify.py @@ -1036,7 +1036,9 @@ def apply( # Now we will reroute the edges went through the inner map, through the # inner access node instead. - for old_inner_edge in list( + for ( + old_inner_edge + ) in list( # TODO(tehrengruber): Why all these list comprehensions everywhere? graph.out_edges_by_connector(map_entry, "OUT_" + connector_name) ): # We now modify the downstream data. This is because we no longer refer diff --git a/src/gt4py/next/program_processors/runners/dace/transformations/splitting_tools.py b/src/gt4py/next/program_processors/runners/dace/transformations/splitting_tools.py index 167678f5f2..94a26ac616 100644 --- a/src/gt4py/next/program_processors/runners/dace/transformations/splitting_tools.py +++ b/src/gt4py/next/program_processors/runners/dace/transformations/splitting_tools.py @@ -266,6 +266,7 @@ def split_node( # NOTE: Turning them into a string is the best solution is probably the only way # to achieve some stability. The only downside is that the order now depends # on the specialization level that is used, i.e. if we have numbers or symbols. + # TODO(tehrengruber): Is this still needed? split_description = sorted(split_description, key=lambda split: str(split)) desc_to_split = node_to_split.desc(sdfg) diff --git a/src/gt4py/next/program_processors/runners/dace/workflow/common.py b/src/gt4py/next/program_processors/runners/dace/workflow/common.py index 312f418c11..c065a55a77 100644 --- a/src/gt4py/next/program_processors/runners/dace/workflow/common.py +++ b/src/gt4py/next/program_processors/runners/dace/workflow/common.py @@ -57,6 +57,7 @@ def set_dace_config( # NOTE: Each thread maintains its own set of configuration, i.e. `dace.Config` is # a thread local variable. This means it is safe to set values that are different # for each thread. + dace.Config.set("progress", value=True) # We rely on dace cache to avoid recompiling the SDFG. # Note that the workflow step with the persistent `FileCache` store diff --git a/src/gt4py/next/program_processors/runners/dace/workflow/translation.py b/src/gt4py/next/program_processors/runners/dace/workflow/translation.py index cca4114fc1..37e438de6a 100644 --- a/src/gt4py/next/program_processors/runners/dace/workflow/translation.py +++ b/src/gt4py/next/program_processors/runners/dace/workflow/translation.py @@ -9,6 +9,7 @@ from __future__ import annotations import dataclasses +import json from typing import Any, Optional import dace @@ -29,6 +30,35 @@ from gt4py.next.type_system import type_specifications as ts +def remove_guid(data: Any) -> Any: + """ + Recursively traverse a dict and remove all keys named 'guid'. + + Args: + data: A dictionary, list, or other data structure to traverse + + Returns: + The data structure with all 'guid' keys removed + """ + if isinstance(data, dict): + return {key: remove_guid(value) for key, value in data.items() if key != "guid"} + elif isinstance(data, list): + return [remove_guid(item) for item in data] + elif isinstance(data, tuple): + return tuple(remove_guid(item) for item in data) + elif isinstance(data, (str, int, float, bool)) or data is None: + return data + else: + raise RuntimeError("Unsupported data type") + + +def remove_guid_and_save_to_file(sdfg: dace.SDFG, filename: str) -> None: + cleaned_data = remove_guid(sdfg.to_json()) + + with open(filename, "w") as f: + json.dump(cleaned_data, f, indent=2) + + def find_constant_symbols( ir: itir.Program, sdfg: dace.SDFG, @@ -364,7 +394,10 @@ def generate_sdfg( *args: Any, **kwargs: Any, ) -> dace.SDFG: - with gtx_wfdcommon.dace_context(device_type=self.device_type): + with ( + gtx_wfdcommon.dace_context(device_type=self.device_type), + dace.sdfg.nodes.reset_node_id_counter(), + ): return self._generate_sdfg_without_configuring_dace(*args, **kwargs) def _generate_sdfg_without_configuring_dace( diff --git a/uv.lock b/uv.lock index f853aee738..bca76bf77a 100644 --- a/uv.lock +++ b/uv.lock @@ -1207,7 +1207,7 @@ wheels = [ [[package]] name = "dace" version = "2.0.0a4" -source = { registry = "https://pypi.org/simple" } +source = { git = "https://github.com/GridTools/dace?branch=dace_toolchain_deterministic#fbf235c565fdef94ba6af879c008bcc34ef11c0b" } dependencies = [ { name = "astunparse" }, { name = "dill" }, @@ -1216,6 +1216,7 @@ dependencies = [ { name = "networkx", version = "3.5", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' or (extra == 'extra-5-gt4py-cuda12' and extra == 'extra-5-gt4py-jax-cuda13') or (extra == 'extra-5-gt4py-cuda12' and extra == 'extra-5-gt4py-rocm6') or (extra == 'extra-5-gt4py-cuda12' and extra == 'extra-5-gt4py-rocm7') or (extra == 'extra-5-gt4py-jax-cuda13' and extra == 'extra-5-gt4py-rocm6') or (extra == 'extra-5-gt4py-jax-cuda13' and extra == 'extra-5-gt4py-rocm7') or (extra == 'extra-5-gt4py-rocm6' and extra == 'extra-5-gt4py-rocm7')" }, { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11' or (extra == 'extra-5-gt4py-cuda12' and extra == 'extra-5-gt4py-jax-cuda13') or (extra == 'extra-5-gt4py-cuda12' and extra == 'extra-5-gt4py-rocm6') or (extra == 'extra-5-gt4py-cuda12' and extra == 'extra-5-gt4py-rocm7') or (extra == 'extra-5-gt4py-jax-cuda13' and extra == 'extra-5-gt4py-rocm6') or (extra == 'extra-5-gt4py-jax-cuda13' and extra == 'extra-5-gt4py-rocm7') or (extra == 'extra-5-gt4py-rocm6' and extra == 'extra-5-gt4py-rocm7')" }, { name = "numpy", version = "2.4.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' or (extra == 'extra-5-gt4py-cuda12' and extra == 'extra-5-gt4py-jax-cuda13') or (extra == 'extra-5-gt4py-cuda12' and extra == 'extra-5-gt4py-rocm6') or (extra == 'extra-5-gt4py-cuda12' and extra == 'extra-5-gt4py-rocm7') or (extra == 'extra-5-gt4py-jax-cuda13' and extra == 'extra-5-gt4py-rocm6') or (extra == 'extra-5-gt4py-jax-cuda13' and extra == 'extra-5-gt4py-rocm7') or (extra == 'extra-5-gt4py-rocm6' and extra == 'extra-5-gt4py-rocm7')" }, + { name = "ordered-set" }, { name = "packaging" }, { name = "ply" }, { name = "pyreadline", marker = "sys_platform == 'win32' or (extra == 'extra-5-gt4py-cuda12' and extra == 'extra-5-gt4py-jax-cuda13') or (extra == 'extra-5-gt4py-cuda12' and extra == 'extra-5-gt4py-rocm6') or (extra == 'extra-5-gt4py-cuda12' and extra == 'extra-5-gt4py-rocm7') or (extra == 'extra-5-gt4py-jax-cuda13' and extra == 'extra-5-gt4py-rocm6') or (extra == 'extra-5-gt4py-jax-cuda13' and extra == 'extra-5-gt4py-rocm7') or (extra == 'extra-5-gt4py-rocm6' and extra == 'extra-5-gt4py-rocm7')" }, @@ -1223,7 +1224,6 @@ dependencies = [ { name = "sympy" }, { name = "typing-extensions" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/3a/07/a0410822c1aaef052c0015e7a8399a7d7da76a11e0a97939027c15d125ee/dace-2.0.0a4.tar.gz", hash = "sha256:f45f84b59a34846e0049347559416c1fac39c739aa6cbc791e27e1973d58b1ad", size = 6026693, upload-time = "2026-06-25T14:09:32.009Z" } [[package]] name = "debugpy" @@ -1859,7 +1859,7 @@ requires-dist = [ { name = "cupy-cuda13x", marker = "extra == 'cuda13'", specifier = ">=14.0" }, { name = "cupy-rocm-7-0", marker = "extra == 'rocm7'", specifier = ">=14.0" }, { name = "cytoolz", specifier = ">=1.0.1" }, - { name = "dace", specifier = ">=2.0.0a4" }, + { name = "dace", git = "https://github.com/GridTools/dace?branch=dace_toolchain_deterministic" }, { name = "deepdiff", specifier = ">=8.1.0" }, { name = "devtools", specifier = ">=0.6" }, { name = "factory-boy", specifier = ">=3.3.3" }, @@ -1882,6 +1882,7 @@ requires-dist = [ { name = "ninja", specifier = ">=1.11" }, { name = "numpy", marker = "python_full_version < '3.14'", specifier = ">=2.0.0" }, { name = "numpy", marker = "python_full_version >= '3.14'", specifier = ">=2.3.2" }, + { name = "ordered-set", specifier = ">=4.0.0" }, { name = "ordered-set", specifier = ">=4.1.0" }, { name = "packaging", specifier = ">=20.0" }, { name = "pybind11", specifier = ">=3.0.3" },