diff --git a/.github/workflows/copilot-setup-steps.yml b/.github/workflows/copilot-setup-steps.yml index b97fbe0869..4579ac2a38 100644 --- a/.github/workflows/copilot-setup-steps.yml +++ b/.github/workflows/copilot-setup-steps.yml @@ -8,7 +8,7 @@ jobs: contents: read steps: - name: Checkout code - uses: actions/checkout@v6 + uses: actions/checkout@v7 with: submodules: recursive diff --git a/.github/workflows/general-ci.yml b/.github/workflows/general-ci.yml index 7535d082cc..8d49f2a70f 100644 --- a/.github/workflows/general-ci.yml +++ b/.github/workflows/general-ci.yml @@ -22,7 +22,7 @@ jobs: simplify: [0,1,autoopt] steps: - - uses: actions/checkout@v6 + - uses: actions/checkout@v7 with: submodules: 'recursive' - name: Set up Python ${{ matrix.python-version }} @@ -106,7 +106,7 @@ jobs: ./tests/xform_test.sh coverage combine .; coverage report; coverage xml - - uses: codecov/codecov-action@v5 + - uses: codecov/codecov-action@v6 with: token: ${{ secrets.CODECOV_TOKEN }} verbose: true diff --git a/.github/workflows/gpu-ci.yml b/.github/workflows/gpu-ci.yml index af276228df..5f8ac948ef 100644 --- a/.github/workflows/gpu-ci.yml +++ b/.github/workflows/gpu-ci.yml @@ -23,7 +23,7 @@ jobs: if: "!contains(github.event.pull_request.labels.*.name, 'no-ci')" runs-on: [self-hosted, gpu] steps: - - uses: actions/checkout@v6 + - uses: actions/checkout@v7 with: submodules: 'recursive' - name: Install dependencies @@ -36,6 +36,7 @@ jobs: pip install flake8 pytest-xdist coverage pip install mpi4py pip install cupy + pip install cutensor-cu12 pip uninstall -y dace pip install -e ".[testing,ml]" curl -Os https://uploader.codecov.io/latest/linux/codecov diff --git a/.github/workflows/heterogeneous-ci.yml b/.github/workflows/heterogeneous-ci.yml index d95667e127..82907599d7 100644 --- a/.github/workflows/heterogeneous-ci.yml +++ b/.github/workflows/heterogeneous-ci.yml @@ -23,7 +23,7 @@ jobs: if: "!contains(github.event.pull_request.labels.*.name, 'no-ci')" runs-on: [self-hosted, linux] steps: - - uses: actions/checkout@v6 + - uses: actions/checkout@v7 with: submodules: 'recursive' - name: Install dependencies diff --git a/.github/workflows/linting.yml b/.github/workflows/linting.yml index c156e9d730..cc50b78241 100644 --- a/.github/workflows/linting.yml +++ b/.github/workflows/linting.yml @@ -14,7 +14,7 @@ jobs: steps: - name: Check repository - uses: actions/checkout@v6 + uses: actions/checkout@v7 - name: Setup Python 3.10 uses: actions/setup-python@v6 @@ -27,10 +27,4 @@ jobs: - name: Run linting tools id: lint - continue-on-error: true - run: pre-commit run --all-files - - - name: Show git diff - if: steps.lint.outcome == 'failure' - run: | - ./.github/workflows/scripts/show-git-diff.sh + run: pre-commit run --show-diff-on-failure --color=always --all-files diff --git a/.github/workflows/ml-ci.yml b/.github/workflows/ml-ci.yml index 94a6180ce7..c32a2121e7 100644 --- a/.github/workflows/ml-ci.yml +++ b/.github/workflows/ml-ci.yml @@ -22,7 +22,7 @@ jobs: simplify: [0,1,autoopt] steps: - - uses: actions/checkout@v6 + - uses: actions/checkout@v7 with: submodules: 'recursive' - name: Set up Python ${{ matrix.python-version }} @@ -56,7 +56,7 @@ jobs: pytest --cov-report=xml --cov=dace --tb=short --timeout_method thread --timeout=600 -v -m "(torch or onnx or autodiff) and not gpu" ./codecov - - uses: codecov/codecov-action@v5 + - uses: codecov/codecov-action@v6 with: token: ${{ secrets.CODECOV_TOKEN }} verbose: true diff --git a/.github/workflows/pyFV3-ci.yml b/.github/workflows/pyFV3-ci.yml index 5543c8bc2d..95e5d5de49 100644 --- a/.github/workflows/pyFV3-ci.yml +++ b/.github/workflows/pyFV3-ci.yml @@ -1,17 +1,12 @@ name: NASA/NOAA pyFV3 repository build test -# Temporarily disabled for main, and instead applied to a specific DaCe v1 maintenance branch (v1/maintenance). Once -# the FV3 bridge has been adapted to DaCe v1, this will need to be reverted back to apply to main. on: push: - #branches: [ main, ci-fix ] - branches: [ v1/maintenance, ci-fix ] + branches: [ main, ci-fix, romanc/stree-v2 ] pull_request: - #branches: [ main, ci-fix ] - branches: [ v1/maintenance, ci-fix ] + branches: [ main, ci-fix, romanc/stree-v2 ] merge_group: - #branches: [ main, ci-fix ] - branches: [ v1/maintenance, ci-fix ] + branches: [ main, ci-fix, romanc/stree-v2 ] defaults: run: @@ -22,84 +17,80 @@ concurrency: cancel-in-progress: true jobs: - build_and_validate_pyFV3: - if: "!contains(github.event.pull_request.labels.*.name, 'no-ci')" - runs-on: ubuntu-latest - strategy: - matrix: - python-version: [3.11.7] + configure: + runs-on: ubuntu-latest + if: "!contains(github.event.pull_request.labels.*.name, 'no-ci')" + name: "Read test configuration" + outputs: + matrix: ${{ steps.set-matrix.outputs.matrix }} + steps: + - name: Checkout to repository + uses: actions/checkout@v7 + - name: Set matrix data + id: set-matrix + run: echo "matrix={\"include\":$(jq -c . < .github/workflows/pyFV3-regression-matrix.json)}" >> $GITHUB_OUTPUT - steps: - - uses: actions/checkout@v6 - with: - repository: 'NOAA-GFDL/PyFV3' - ref: 'ci/DaCe' - submodules: 'recursive' - path: 'pyFV3' - - uses: actions/checkout@v6 - with: - path: 'dace' - submodules: 'recursive' - - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v6 - with: - python-version: ${{ matrix.python-version }} - - name: Install library dependencies - run: | - sudo apt-get update - sudo apt-get install -y libopenmpi-dev libboost-all-dev - gcc --version - # Because Github doesn't allow us to do a git checkout in code - # we use a trick to checkout DaCe first (not using the external submodule) - # install the full suite via requirements_dev, then re-install the correct DaCe - - name: Install Python packages - run: | - python -m pip install --upgrade pip wheel setuptools - pip install -e ./pyFV3[develop] - pip install -e ./dace - - name: Download data - run: | - cd pyFV3 - mkdir -p test_data - cd test_data - wget --retry-connrefused https://portal.nccs.nasa.gov/datashare/astg/smt/pace-regression-data/8.1.3_c12_6ranks_standard.D_SW.tar.gz - tar -xzvf 8.1.3_c12_6ranks_standard.D_SW.tar.gz - wget --retry-connrefused https://portal.nccs.nasa.gov/datashare/astg/smt/pace-regression-data/8.1.3_c12_6ranks_standard.RiemSolver3.tar.gz - tar -xzvf 8.1.3_c12_6ranks_standard.RiemSolver3.tar.gz - wget --retry-connrefused https://portal.nccs.nasa.gov/datashare/astg/smt/pace-regression-data/8.1.3_c12_6ranks_standard.Remapping.tar.gz - tar -xzvf 8.1.3_c12_6ranks_standard.Remapping.tar.gz - cd ../.. - # Clean up caches between run for stale un-expanded SDFG to trip the build system (NDSL side issue) - - name: "Regression test: Riemman Solver on D-grid (RiemSolver3)" - env: - FV3_DACEMODE: BuildAndRun - PACE_CONSTANTS: GFS - PACE_LOGLEVEL: Debug - run: | - pytest -v -s --data_path=./pyFV3/test_data/8.1.3/c12_6ranks_standard/dycore \ - --backend=dace:cpu --which_modules=Riem_Solver3 \ - --threshold_overrides_file=./pyFV3/tests/savepoint/translate/overrides/standard.yaml \ - ./pyFV3/tests/savepoint - rm -r ./.gt_cache_FV3_A - - name: "Regression test: Shallow water lagrangian dynamics on D-grid (D_SW) (on rank 0 only)" - env: - FV3_DACEMODE: BuildAndRun - PACE_CONSTANTS: GFS - PACE_LOGLEVEL: Debug - run: | - pytest -v -s --data_path=./pyFV3/test_data/8.1.3/c12_6ranks_standard/dycore \ - --backend=dace:cpu --which_modules=D_SW --which_rank=0 \ - --threshold_overrides_file=./pyFV3/tests/savepoint/translate/overrides/standard.yaml \ - ./pyFV3/tests/savepoint - rm -r ./.gt_cache_FV3_A - - name: "Regression test: Remapping (on rank 0 only)" - env: - FV3_DACEMODE: BuildAndRun - PACE_CONSTANTS: GFS - PACE_LOGLEVEL: Debug - run: | - pytest -v -s --data_path=./pyFV3/test_data/8.1.3/c12_6ranks_standard/dycore \ - --backend=dace:cpu --which_modules=Remapping --which_rank=0 \ - --threshold_overrides_file=./pyFV3/tests/savepoint/translate/overrides/standard.yaml \ - ./pyFV3/tests/savepoint - rm -r ./.gt_cache_FV3_A + build_and_validate_pyFV3: + runs-on: ubuntu-latest + needs: configure + strategy: + fail-fast: false + matrix: ${{ fromJson(needs.configure.outputs.matrix) }} + name: "${{ matrix.module }} regression test" + + steps: + - name: Checkout pyFV3 + uses: actions/checkout@v7 + with: + repository: 'NOAA-GFDL/pyFV3' + path: 'pyFV3' + + - name: Checkout dace + uses: actions/checkout@v7 + with: + path: 'dace' + submodules: 'recursive' + + - name: Set up Python 3.12 + uses: actions/setup-python@v6 + with: + python-version: 3.12 + + - name: Install OpenMPI + run: pip install openmpi + + # We install `dace` after `pyFV3` to build pyFV3 with the current version of DaCe + - name: Install pyFV3 and dace + run: | + pip install ./pyFV3[ndsl,test] + pip install ./dace + + - name: Print versions + run: | + gcc --version + python --version + pip --version + pip list + + - name: Download data + # Skip "Riemann Solver on D-grid" (currently failing) + if: ${{ !startsWith(matrix.module, 'Riem_Solver3') }} + run: | + mkdir -p pyFV3/test_data + cd pyFV3/test_data + wget --retry-connrefused https://portal.nccs.nasa.gov/datashare/astg/smt/pace-regression-data/${{ matrix.archive }} + tar -xzvf ${{ matrix.archive }} + cd ../.. + + - name: "Regression test: ${{ matrix.name }}" + # Skip "Riemann Solver on D-grid" (currently failing) + if: ${{ !startsWith(matrix.module, 'Riem_Solver3') }} + env: + FV3_DACEMODE: BuildAndRun + NDSL_LOGLEVEL: Debug + run: | + pytest -v -s --data_path=./pyFV3/test_data/8.1.3/c12_6ranks_standard/dycore \ + --backend=orch:dace:cpu:KIJ --which_modules=${{ matrix.module }} \ + ${{ matrix.extra_args }} \ + --threshold_overrides_file=./pyFV3/tests/savepoint/translate/overrides/standard.yaml \ + ./pyFV3/tests/savepoint diff --git a/.github/workflows/pyFV3-regression-matrix.json b/.github/workflows/pyFV3-regression-matrix.json new file mode 100644 index 0000000000..ce876c877c --- /dev/null +++ b/.github/workflows/pyFV3-regression-matrix.json @@ -0,0 +1,20 @@ +[ + { + "name": "Shallow water lagrangian dynamics on D-grid (D_SW) (on rank 0 only)", + "archive": "8.1.3_c12_6ranks_standard.D_SW.tar.gz", + "module": "D_SW", + "extra_args": "--which_rank=0" + }, + { + "name":"Riemann Solver on D-grid (RiemSolver3)", + "archive":"8.1.3_c12_6ranks_standard.RiemSolver3.tar.gz", + "module":"Riem_Solver3", + "extra_args":"--which_rank=0" + }, + { + "name": "Remapping (on rank 0 only)", + "archive": "8.1.3_c12_6ranks_standard.Remapping.tar.gz", + "module": "Remapping", + "extra_args": "--which_rank=0" + } +] diff --git a/.github/workflows/release.sh b/.github/workflows/release.sh index 7b8ff5f5e4..12504fc04d 100755 --- a/.github/workflows/release.sh +++ b/.github/workflows/release.sh @@ -3,7 +3,7 @@ set -e # Install dependencies -pip install --upgrade twine +pip install --upgrade twine build # Synchronize submodules git submodule update --init --recursive diff --git a/.github/workflows/scripts/show-git-diff.sh b/.github/workflows/scripts/show-git-diff.sh deleted file mode 100755 index a811c01672..0000000000 --- a/.github/workflows/scripts/show-git-diff.sh +++ /dev/null @@ -1,21 +0,0 @@ -#!/bin/bash - -# Check for uncommitted changes in the working tree -if [ -n "$(git status --porcelain)" ]; then - echo "Linting tools found the following changes are needed to comply" - echo "with our automatic styling." - echo "" - echo "Please run \"pre-commit run --all-files\" locally to fix these." - echo "See also https://github.com/spcl/dace/blob/main/CONTRIBUTING.md" - echo "" - echo "git status" - echo "----------" - git status - echo "" - echo "git diff" - echo "--------" - git --no-pager diff - echo "" - - exit 1 -fi diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index bfedf3a96d..06faf41167 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -2,6 +2,11 @@ default_language_version: python: python3 repos: +- repo: https://github.com/astral-sh/ruff-pre-commit + rev: v0.15.5 + hooks: + - id: ruff-check + args: [--fix] - repo: https://github.com/pre-commit/pre-commit-hooks rev: v5.0.0 hooks: diff --git a/.readthedocs.yaml b/.readthedocs.yaml index 4cea907235..afe605d31f 100644 --- a/.readthedocs.yaml +++ b/.readthedocs.yaml @@ -6,9 +6,9 @@ version: 2 build: - os: ubuntu-20.04 + os: ubuntu-24.04 tools: - python: "3.10" + python: "3.14" # Build documentation in the docs/ directory with Sphinx sphinx: diff --git a/ci/cscs_gpu.yml b/ci/cscs_gpu.yml index 0299dbaead..0763876534 100644 --- a/ci/cscs_gpu.yml +++ b/ci/cscs_gpu.yml @@ -54,6 +54,12 @@ test_cscs_gh200: - cd "${WORKDIR}/dace" - git fetch --depth 1 origin "${CI_COMMIT_SHA}" - git checkout "${CI_COMMIT_SHA}" + - uv pip install cutensor-cu12 + - CUTENSOR_DIR="${VIRTUAL_ENV}/lib/python${PY_VERSION}/site-packages/cutensor" + - ln -sf libcutensor.so.2 "${CUTENSOR_DIR}/lib/libcutensor.so" + - export CPATH="${CUTENSOR_DIR}/include:${CPATH:-}" + - export LIBRARY_PATH="${CUTENSOR_DIR}/lib:${LIBRARY_PATH:-}" + - export LD_LIBRARY_PATH="${CUTENSOR_DIR}/lib:${LD_LIBRARY_PATH:-}" - uv pip install -e ".[testing]" - export DACE_cache=unique - pytest --cov-report=xml --cov=dace --tb=short --timeout_method thread --timeout=300 -n 32 -m "${TEST_VARIANTS}" diff --git a/dace/autodiff/backward_pass_generator.py b/dace/autodiff/backward_pass_generator.py index 11daf5f3c1..9040cf09e2 100644 --- a/dace/autodiff/backward_pass_generator.py +++ b/dace/autodiff/backward_pass_generator.py @@ -740,11 +740,11 @@ def _zero_out_gradient(self, forward_state: SDFGState, forward_node: nodes.Acces # Create a new memlet that copies what memlet is writing to to the tmp new_memlet_subset = memlet.subset if memlet.data == forward_node.data else memlet.other_subset original_to_tmp_memlet = dace.Memlet(data=backward_node.data, - subset=new_memlet_subset, - other_subset=new_memlet_subset) + subset=copy.deepcopy(new_memlet_subset), + other_subset=copy.deepcopy(new_memlet_subset)) # Remove the src_subset of the new memlet and replace the memlet in the edge - map_exit_memlet.subset = memlet.subset if memlet.data == forward_node.data else memlet.other_subset + map_exit_memlet.subset = copy.deepcopy(new_memlet_subset) map_exit_memlet.other_subset = None edge.data = map_exit_memlet @@ -1761,14 +1761,15 @@ def _connect_given_gradients(self, forward_state: SDFGState, backward_state: SDF # Special case for when the two access nodes are the same if forward_node.data == dest_node.data and fwd_memlet.other_subset is not None: new_memlet = dace.Memlet(data=self.reverse_map[forward_node].data, - subset=fwd_memlet.other_subset, - other_subset=fwd_memlet.subset) + subset=copy.deepcopy(fwd_memlet.other_subset), + other_subset=copy.deepcopy(fwd_memlet.subset)) else: - new_memlet = dace.Memlet(data=self.reverse_map[forward_node].data, - subset=fwd_memlet.subset - if fwd_memlet.data == forward_node.data else fwd_memlet.other_subset, - other_subset=fwd_memlet.other_subset - if fwd_memlet.data == forward_node.data else fwd_memlet.subset) + new_memlet = dace.Memlet( + data=self.reverse_map[forward_node].data, + subset=copy.deepcopy(fwd_memlet.subset if fwd_memlet.data == + forward_node.data else fwd_memlet.other_subset), + other_subset=copy.deepcopy(fwd_memlet.other_subset if fwd_memlet.data == + forward_node.data else fwd_memlet.subset)) memlet = new_memlet if input_conn not in self.result_map[dest_node].required_grad_names: continue diff --git a/dace/autodiff/data_forwarding/recompute.py b/dace/autodiff/data_forwarding/recompute.py index ce6330dbff..3be76f307b 100644 --- a/dace/autodiff/data_forwarding/recompute.py +++ b/dace/autodiff/data_forwarding/recompute.py @@ -5,7 +5,7 @@ # DaCe imports import dace import dace.sdfg.nodes as nodes -from dace.sdfg import SDFG, SDFGState, graph as dgraph, state as dstate +from dace.sdfg import SDFG, SDFGState, state as dstate from dace.sdfg.state import LoopRegion # Autodiff imports @@ -67,7 +67,7 @@ def _connect_recomputation_nsdfg(bwd_generator: 'BackwardPassGenerator', forward # Write the data to a new access node in the backward state # Add a new AccessNode and array to the forward pass - # First, check if a recomputated array with this name already exists + # First, check if a recomputed array with this name already exists if "recomputed_" + target_an.data not in bwd_generator.backward_sdfg.arrays: new_recomp_node_name = "recomputed_" + target_an.data else: @@ -84,12 +84,13 @@ def _connect_recomputation_nsdfg(bwd_generator: 'BackwardPassGenerator', forward # Add the array descriptor and AccessNode to the forward state original_desc = target_an.desc(forward_state) - new_recomp_node = backward_state.add_array( + backward_state.sdfg.add_array( name=new_recomp_node_name, shape=shape, dtype=original_desc.dtype, transient=True, ) + new_recomp_node = backward_state.add_access(new_recomp_node_name) new_recomp_node.setzero = True # Create a memlet passing all the data to the nested-SDFG diff --git a/dace/autodiff/data_forwarding/store.py b/dace/autodiff/data_forwarding/store.py index a744884e36..b37b203b14 100644 --- a/dace/autodiff/data_forwarding/store.py +++ b/dace/autodiff/data_forwarding/store.py @@ -167,12 +167,13 @@ def _store_data(bwd_generator: 'BackwardPassGenerator', forward_state: SDFGState # We want to build the memlet as if this was not in a a loop nb_enclosing_loops = 0 - new_store_node = forward_state.add_array( + forward_state.sdfg.add_array( name=new_store_node_name, shape=shape, dtype=original_desc.dtype, transient=True, ) + new_store_node = forward_state.add_access(new_store_node_name) # Connect the edge source and connector to the new access node # We will save the memlets we create and return them @@ -186,7 +187,7 @@ def _store_data(bwd_generator: 'BackwardPassGenerator', forward_state: SDFGState # In the other cases, we need to route the storing through maps all_edges = ad_utils.get_all_path_edges(forward_state, forward_an, edge) - # Get the map nest memlet informtation + # Get the map nest memlet information start_range, param_list, shape_list, param_dict = ad_utils.get_map_nest_information(all_edges) # The parameters to add for the current memlet in the loop diff --git a/dace/autodiff/implementations/onnx_ops.py b/dace/autodiff/implementations/onnx_ops.py index e04f86c728..740c4fda27 100644 --- a/dace/autodiff/implementations/onnx_ops.py +++ b/dace/autodiff/implementations/onnx_ops.py @@ -28,7 +28,6 @@ # ONNX-specific imports import dace.libraries.onnx as donnx -from dace.libraries.onnx.converters import clean_onnx_name from dace.libraries.onnx.op_implementations.linalg_ops import PureEinsum from dace.transformation.onnx.replacement import onnx_constant_or_none diff --git a/dace/autodiff/library/library.py b/dace/autodiff/library/library.py index b5e2a60e98..e7ed6aec85 100644 --- a/dace/autodiff/library/library.py +++ b/dace/autodiff/library/library.py @@ -4,12 +4,11 @@ Includes the BackwardPass library node, and the replacements for the python frontend """ -from typing import Dict, Set, Optional -import copy +from typing import Dict, Set import dace import dace.library -from dace import data, properties +from dace import properties from dace.transformation import transformation as pm from dace.sdfg import SDFG, SDFGState, graph, nodes @@ -19,7 +18,7 @@ from dace.transformation.passes.analysis import AccessSets # Import ParameterArray from the data package for backward compatibility -from dace.data.ml import ParameterArray +from dace.data.ml import ParameterArray # noqa: F401 @dace.library.expansion diff --git a/dace/cli/daceprof.py b/dace/cli/daceprof.py index b201d40661..86ab6fee02 100644 --- a/dace/cli/daceprof.py +++ b/dace/cli/daceprof.py @@ -9,7 +9,7 @@ import sys import os import shutil -from typing import Callable, List, Optional, Tuple, Union +from typing import List, Optional, Tuple, Union import warnings import dace diff --git a/dace/cli/external_transformation_registry.py b/dace/cli/external_transformation_registry.py index 2d4a5145e3..bf1e2a2cd8 100755 --- a/dace/cli/external_transformation_registry.py +++ b/dace/cli/external_transformation_registry.py @@ -15,7 +15,6 @@ from typing import Dict, Optional from urllib.parse import urlparse import dace -import shutil class TransformationRepoManager: diff --git a/dace/cli/fcdc.py b/dace/cli/fcdc.py index 7b0fc5c02c..3a2f78531d 100644 --- a/dace/cli/fcdc.py +++ b/dace/cli/fcdc.py @@ -1,9 +1,7 @@ # Copyright 2019-2023 ETH Zurich and the DaCe authors. All rights reserved. """ Simple Fortran SDFG command-line compiler. """ -import dace import os -import sys import argparse import shutil from dace.frontend.fortran import fortran_parser diff --git a/dace/cli/sdfgcc.py b/dace/cli/sdfgcc.py index 0d04950be7..eeff31cb9d 100644 --- a/dace/cli/sdfgcc.py +++ b/dace/cli/sdfgcc.py @@ -3,7 +3,6 @@ import dace import os -import sys import argparse import shutil diff --git a/dace/codegen/CMakeLists.txt b/dace/codegen/CMakeLists.txt index abc06ef4ca..614f92a029 100644 --- a/dace/codegen/CMakeLists.txt +++ b/dace/codegen/CMakeLists.txt @@ -235,18 +235,15 @@ if(DACE_ENABLE_HIP) else() if(_arch_res EQUAL 0) set(LOCAL_HIP_ARCHITECTURES "" CACHE STRING "Detected local AMD GPUs for compilation") + else() + set(LOCAL_HIP_ARCHITECTURES "${DACE_HIP_ARCHITECTURES_DEFAULT}" CACHE STRING "Detected local GPUs for compilation") + message(STATUS "No local HIP-capable GPUs found. Using default: ${DACE_HIP_ARCHITECTURES_DEFAULT}") endif() - message(STATUS "No local HIP-capable GPUs found") endif() endif() # Add flags to compile for local AMD architectures - foreach(var ${LOCAL_HIP_ARCHITECTURES}) - list(APPEND HIP_HIPCC_FLAGS --offload-arch=${var}) - endforeach() - - # Add flags from dace config - list(APPEND HIP_HIPCC_FLAGS ${EXTRA_HIP_FLAGS}) + set(CMAKE_HIP_ARCHITECTURES "${LOCAL_HIP_ARCHITECTURES}") # Add include directories for other files set(DACE_LIBS ${DACE_LIBS} hip::host) diff --git a/dace/codegen/codegen.py b/dace/codegen/codegen.py index d83b498de9..fc6791599f 100644 --- a/dace/codegen/codegen.py +++ b/dace/codegen/codegen.py @@ -1,5 +1,6 @@ # Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. import functools +import json from typing import List import dace @@ -178,8 +179,10 @@ def generate_code(sdfg: SDFG, validate=True) -> List[CodeObject]: if not filecmp.cmp(f'{tmp_dir}/test.sdfg', f'{tmp_dir}/test2.sdfg'): with open(f'{tmp_dir}/test.sdfg', 'r') as f1: with open(f'{tmp_dir}/test2.sdfg', 'r') as f2: - diff = difflib.unified_diff(f1.readlines(), - f2.readlines(), + data1 = json.dumps(json.load(f1), indent=2).splitlines(keepends=True) + data2 = json.dumps(json.load(f2), indent=2).splitlines(keepends=True) + diff = difflib.unified_diff(data1, + data2, fromfile='test.sdfg (first save)', tofile='test2.sdfg (after roundtrip)') diff = ''.join(diff) diff --git a/dace/codegen/codeobject.py b/dace/codegen/codeobject.py index aab5cfdb37..aec0a2a610 100644 --- a/dace/codegen/codeobject.py +++ b/dace/codegen/codeobject.py @@ -1,5 +1,6 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. import re +import dace from dace import sourcemap from dace.properties import (Property, DictProperty, SetProperty, make_properties) @@ -45,9 +46,15 @@ def __init__(self, self.linkable = linkable self.environments = environments or set() - if language == 'cpp' and title == 'Frame' and sdfg: - sourcemap.create_maps(sdfg, code, self.target.target_name) + # NOTE: In an earlier version, the source maps were generated here. However, + # this had the side effect that the build folder was generated containing + # the source maps, i.e. the `map/` subfolder. However, no code was actually + # dumped to disc. Another effect is that the generation of code would actually + # overwrite the source map, thus it was moved to `generate_program_folder()`. @property def clean_code(self): return re.sub(r'[ \t]*////__(DACE:|CODEGEN;)[^\n]*', '', self.code) + + def create_source_map(self, sdfg: 'dace.SDFG') -> None: + sourcemap.create_maps(sdfg, self.code, self.target.target_name) diff --git a/dace/codegen/common.py b/dace/codegen/common.py index f5bbf445a2..7e0279022d 100644 --- a/dace/codegen/common.py +++ b/dace/codegen/common.py @@ -37,7 +37,7 @@ def find_outgoing_edges(node, dfg): return list(dfg.out_edges(node)) -@lru_cache(maxsize=16384) +@lru_cache(maxsize=16384, typed=True) def _sym2cpp(s, arrayexprs): return cppunparse.pyexpr2cpp(symbolic.symstr(s, arrayexprs, cpp_mode=True)) diff --git a/dace/codegen/compiled_sdfg.py b/dace/codegen/compiled_sdfg.py index a0ecaf41c4..a188ad9b73 100644 --- a/dace/codegen/compiled_sdfg.py +++ b/dace/codegen/compiled_sdfg.py @@ -13,12 +13,10 @@ import sys import numpy as np -import sympy as sp from dace import data as dt, dtypes, hooks, symbolic from dace.codegen import exceptions as cgx from dace.config import Config -from dace.frontend import operations class ReloadableDLL(object): @@ -27,17 +25,26 @@ class ReloadableDLL(object): bypasses Python's dynamic library reloading issues. """ - def __init__(self, library_filename, program_name): - """ - Creates a new reloadable shared object. + def __init__(self, library_filename, **kwargs): + """Creates a new reloadable shared object. + + The path to the library must be given. The path of the stub library is inferred + from it. This means it is expected that it is located in the same folder, + see `_get_stub_library_path()` for more information. :param library_filename: Path to library file. - :param program_name: Name of the DaCe program (for use in finding - the stub library loader). + :param libstub_path: Optional path to the stub library. """ - self._stub_filename = os.path.join(os.path.dirname(os.path.realpath(library_filename)), - f'libdacestub_{program_name}.{Config.get("compiler", "library_extension")}') - self._library_filename = os.path.realpath(library_filename) + from dace.codegen.compiler import _get_stub_library_path + + self._library_filename = str(pathlib.Path(library_filename).resolve()) + + if "libstub_path" in kwargs: + self._stub_filename = str(pathlib.Path(kwargs.pop("libstub_path")).resolve()) + else: + self._stub_filename = str(_get_stub_library_path(self._library_filename)) + assert len(kwargs) == 0 + self._stub = None self._lib = None @@ -183,11 +190,13 @@ class CompiledSDFG(object): directly, instead, it is generated by utilities such as ``SDFG.compile()``. The class performs the following tasks: + - It ensures that the SDFG object is properly initialized, either by a direct call to ``initialize()`` or the first time it is called. Furthermore, it will also take care of the finalization if it goes out of scope. - It marshalls Python arguments into C arguments. + There are two ways in which a compiled SDFG can be called. The first is using ``__call__()``, i.e. as a normal function. However, this always processes the arguments and performs type checks, which introduces overhead (especially with @@ -478,12 +487,14 @@ def safe_call(self, *args, **kwargs): # Pickle the SDFG and arguments with tempfile.NamedTemporaryFile(mode='wb', delete=False) as f: - pickle.dump({ - 'library_path': self._lib._library_filename, - "sdfg": self.sdfg, - 'args': args, - 'kwargs': kwargs - }, f) + pickle.dump( + { + 'library_path': self._lib._library_filename, + 'stublibrary_path': self._lib._stub_filename, + "sdfg": self.sdfg, + 'args': args, + 'kwargs': kwargs + }, f) temp_path = f.name # Call the SDFG in a separate process @@ -498,9 +509,10 @@ def safe_call(self, *args, **kwargs): with open(r"{temp_path}", "rb") as f: data = pickle.load(f) library_path = data['library_path'] +libstub_path = data['stublibrary_path'] sdfg = data['sdfg'] -lib = csd.ReloadableDLL(library_path, sdfg.name) +lib = csd.ReloadableDLL(library_filename=library_path, libstub_path=libstub_path) obj = csd.CompiledSDFG(sdfg, lib, sdfg.arg_names) obj(*data['args'], **data['kwargs']) @@ -580,7 +592,8 @@ def __del__(self): if self._initialized is True: self.finalize() self._initialized = False - self._libhandle = ctypes.c_void_p(0) + if self._libhandle is not None and ctypes is not None: + self._libhandle = ctypes.c_void_p(0) self._lib.unload() def construct_arguments(self, *args: Any, **kwargs: Any) -> Tuple[Tuple[Any], Tuple[Any]]: diff --git a/dace/codegen/compiler.py b/dace/codegen/compiler.py index 177988f15b..6b25033e72 100644 --- a/dace/codegen/compiler.py +++ b/dace/codegen/compiler.py @@ -4,13 +4,14 @@ returns the corresponding CompiledSDFG object. """ import collections +import io import os -import six +import pathlib +import re import shutil import shlex import subprocess -import re -from typing import Any, Callable, Dict, List, Set, Tuple, TypeVar, Union +from typing import Callable, List, Literal, Set, Tuple, TypeVar, Union, Optional, overload import warnings import dace @@ -24,16 +25,48 @@ T = TypeVar('T') -def generate_program_folder(sdfg, code_objects: List[CodeObject], out_path: str, config=None): - """ Writes all files required to configure and compile the DaCe program - into the specified folder. +def generate_program_folder( + sdfg, + code_objects: List[CodeObject], + out_path: str, + config=None, + folder_mode: Optional[str] = None, +) -> str: + """Writes all files required to configure and compile the DaCe program into the specified folder. + + This function respects the ``compiler.build_folder_mode`` configuration variable, + thus depending on its value the content might be different. However, in any case + the source files are always generated. - :param sdfg: The SDFG to generate the program folder for. - :param code_objects: List of generated code objects. - :param out_path: The folder in which the build files should be written. - :return: Path to the program folder. + :param sdfg: The SDFG to generate the program folder for. + :param code_objects: List of generated code objects. + :param out_path: The folder in which the build files should be written. + :param folder_mode: Select which files should be saved in the program build folder; + if not given, ``compiler.build_folder_mode`` is used. + :return: Path to the program folder. + + :note: The ``config`` argument is retained for compatibility and should not be used. """ + # NOTE: In older version the argument `config` could be a used to pass a custom + # "configuration" (probably a `dict`) object, that would then be written to + # `dace.conf` inside the folder. If nothing was provided the content of the + # global `dace.Config` would be used. However, since _everything_ is consulting + # `dace.Config` for advice, an external configuration, i.e. settings different + # from `dace.Config` can not take effect and storing it is wrong. Thus this + # feature was dropped. + if config is not None: + warnings.warn( + 'Passed a not `None` `config` argument to `generate_program_folder()`.' + ' This has no effect and will be ignored. Instead `dace.Config` will' + ' be used.', + category=UserWarning, + stacklevel=2, + ) + + if folder_mode is None: + folder_mode = Config.get('compiler', 'build_folder_mode') + src_path = os.path.join(out_path, "src") filelist = list() @@ -77,7 +110,13 @@ def generate_program_folder(sdfg, code_objects: List[CodeObject], out_path: str, if code_object.linkable == True: filelist.append("{},{},{}".format(target_name, target_type, basename)) + # Generate the source map. + if sdfg and (folder_mode in ["development"]): + if code_object.language == 'cpp' and code_object.title == 'Frame': + code_object.create_source_map(sdfg) + # Write list of files + # Needed to communicate with `configure_and_compile()`, deleted in production mode. with open(os.path.join(out_path, "dace_files.csv"), "w") as filelist_file: filelist_file.write("\n".join(filelist)) @@ -87,39 +126,65 @@ def generate_program_folder(sdfg, code_objects: List[CodeObject], out_path: str, environments |= obj.environments # Write list of environments + # Needed to communicate with `configure_and_compile()`, deleted in production mode. with open(os.path.join(out_path, "dace_environments.csv"), "w") as env_file: env_file.write("\n".join(environments)) - # Copy a full snapshot of configuration script - if config is not None: - config.save(os.path.join(out_path, "dace.conf"), all=True) - else: - Config.save(os.path.join(out_path, "dace.conf"), all=True) - + # Save the SDFG itself and its hash if sdfg is not None: - # Save the SDFG itself and its hash - hash = sdfg.save(os.path.join(out_path, "program.sdfg"), hash=True) + hash = sdfg.save(os.path.join(out_path, "program.sdfgz"), hash=True, compress=True) filepath = os.path.join(out_path, 'include', 'hash.h') contents = f'#define __HASH_{sdfg.name} "{hash}"\n' if not identical_file_exists(filepath, contents): with open(filepath, 'w') as hfile: hfile.write(contents) + # Write cachedir tag + cachedir_tag = os.path.join(out_path, "CACHEDIR.TAG") + if not os.path.exists(cachedir_tag): + with open(cachedir_tag, "w") as f: + f.write("\n".join([ + "Signature: 8a477f597d28d172789f06886806bc55", + "# This file is a cache directory tag created by DaCe.", + "# For information about cache directory tags, see:", + "# http://www.brynosaurus.com/cachedir/", + ])) + + # Generate the parts of the folder that are exclusive to the development folder mode. + if folder_mode in ["development"]: + # Copy a full snapshot of configuration script + Config.save(os.path.join(out_path, "dace.conf"), all=True) + + # The folder mode file is always generated. In case it is missing we assume the old version. + with open(os.path.join(out_path, "FOLDER_MODE"), "w") as version_file: + version_file.write(folder_mode) + return out_path -def configure_and_compile(program_folder, program_name=None, output_stream=None): - """ Configures and compiles a DaCe program in the specified folder into a - shared library file. +def configure_and_compile( + program_folder, + program_name=None, + output_stream=None, + folder_mode: Optional[str] = None, +) -> pathlib.Path: + """ + Configures and compiles a DaCe program in the specified folder into a shared library file. + + This function respects the ``compiler.build_folder_mode`` configuration variable, + thus depending on its value the content might be different. - :param program_folder: Folder containing all files necessary to build, - equivalent to what was passed to - `generate_program_folder`. - :param output_stream: Additional output stream to write to (used for - other clients such as the vscode extension). - :return: Path to the compiled shared library file. + :param program_folder: Folder containing all files necessary to build, equivalent to + what was passed to `generate_program_folder`. + :param output_stream: Additional output stream to write to (used for other clients + such as the vscode extension). + :return: Path to the compiled shared library file. """ + if folder_mode is None: + folder_mode = Config.get('compiler.build_folder_mode') + assert folder_mode in ["development", "production"] + if program_name is None: program_name = os.path.basename(program_folder) program_folder = os.path.abspath(program_folder) @@ -129,8 +194,9 @@ def configure_and_compile(program_folder, program_name=None, output_stream=None) build_folder = os.path.join(program_folder, "build") os.makedirs(build_folder, exist_ok=True) - # Prepare performance report folder - os.makedirs(os.path.join(program_folder, "perf"), exist_ok=True) + # Prepare performance report folder if requested. + if folder_mode == "development": + os.makedirs(os.path.join(program_folder, "perf"), exist_ok=True) # Read list of DaCe files to compile. # We do this instead of iterating over source files in the directory to @@ -218,6 +284,7 @@ def configure_and_compile(program_folder, program_name=None, output_stream=None) print(f'Running CMake: {cmake_command}') cmake_filename = os.path.join(build_folder, 'cmake_configure.sh') + ############################################## # Configure try: @@ -254,10 +321,200 @@ def configure_and_compile(program_folder, program_name=None, output_stream=None) else: raise cgx.CompilationError('Compiler failure:\n' + ex.output) - shared_library_path = os.path.join(build_folder, "lib{}.{}".format(program_name, - Config.get('compiler', 'library_extension'))) + # Get the names of the library files that were generated. + # Currently we are still in the `development` folder mode. + lib_path = get_binary_name(object_folder=program_folder, sdfg_name=program_name, folder_mode="development") + libstub_path = _get_stub_library_path(lib_path) + + # In production mode, we are now deleting what we need and relocating it. + if folder_mode == "production": + lib_path = pathlib.Path(shutil.move(src=lib_path, dst=program_folder)) + libstub_path = pathlib.Path(shutil.move(src=libstub_path, dst=program_folder)) + program_folder = pathlib.Path(program_folder) + # TODO: Find out where `sample/` are generated and suppress their generation. + for to_delete in ["include", "src", "build", "sample", "dace_environments.csv", "dace_files.csv"]: + if (program_folder / to_delete).is_dir(): + shutil.rmtree(os.path.join(program_folder, to_delete)) + else: + (program_folder / to_delete).unlink() + + return lib_path + + +def get_program_handle( + library_path: Union[pathlib.Path, str], + sdfg: 'dace.SDFG', + stub_library_path: Union[pathlib.Path, str, None] = None, +) -> csd.CompiledSDFG: + """Construct a ``CompiledSDFG`` form a precompiled library directly. + + This function is similar to the (preferred) ``load_precompiled_sdfg()``. However, + instead of passing the build folder of the SDFG to the function, the path to the + compiled library is passed directly. - return shared_library_path + :param library_path: Path to the compiled library representing ``sdfg``. + :param sdfg: The SDFG, will be referenced by the returned ``CompiledSDFG``. + :param stub_library_path: The path to the stub library. + """ + library_path = pathlib.Path(library_path) + if not library_path.is_file(): + raise FileNotFoundError('Compiled SDFG library not found: ' + library_path) + libstub_path = _get_stub_library_path(library_path) if stub_library_path is None else pathlib.Path( + stub_library_path).resolve() + assert libstub_path.is_file() + + lib = csd.ReloadableDLL(library_filename=library_path, libstub_path=libstub_path) + return csd.CompiledSDFG(sdfg, lib, sdfg.arg_names) + + +def load_from_file(sdfg, binary_filename): + warnings.warn( + 'Used deprecated ``load_from_file()`` function, use ``get_program_handle()`` instead.', + category=DeprecationWarning, + stacklevel=2, + ) + return get_program_handle(library_path=binary_filename, sdfg=sdfg) + + +@overload +def get_folder_mode(object_folder: Union[pathlib.Path, str], probe: Literal[False] = False) -> str: + ... + + +@overload +def get_folder_mode(object_folder: Union[pathlib.Path, str], probe: Literal[True]) -> Optional[str]: + ... + + +@overload +def get_folder_mode(object_folder: Union[pathlib.Path, str], probe: bool) -> Optional[str]: + ... + + +def get_folder_mode(object_folder: Union[pathlib.Path, str], probe: bool = False) -> Optional[str]: + """Inspect `object_folder` and determine which save mode the folder has. + + If the function finds the ``FOLDER_MODE`` file it will examine it to get the save mode. + If the folder mode file is absent the function assumes that it is the ``development`` + format, however, some sanity checks are performed. + + The function also has the optional argument ``probe`` if given and the folder + save mode could not be inferred the function will return ``None`` instead of + generating an error. + """ + object_folder = pathlib.Path(object_folder) + + if not object_folder.is_dir(): + if probe: + return None + raise NotADirectoryError("The build folder does not exists.") + + if (object_folder / 'FOLDER_MODE').exists(): + with open(object_folder / 'FOLDER_MODE', 'rt') as F: + folder_mode = F.readline().strip() + return folder_mode + else: + # This is to check an old style folder, i.e. a cache folder that was generated before + # the `FOLDER_MODE` file was introduced. We do some small sanity checks. + # TODO: Phase out this feature, after there are no old style caches. + found_sub_folder = False + for sub_folder in ["build", "map", "src", "include", "sample"]: + if (object_folder / sub_folder).is_dir(): + found_sub_folder = True + elif found_sub_folder: + raise NotADirectoryError(f'Expected that folder ``{object_folder}`` contains ``{sub_folder}``') + + if found_sub_folder: + # All expected folders where found, so expect that this is a 'development' format folder. + return "development" + elif probe: + # None of the files where found. Thus this is probably an empty folder that just exist. + return None + else: + # Up for discussion what to do here. + raise NotADirectoryError(f'``{object_folder}`` does not appear to be a valid build folder.') + + +def get_binary_name( + object_folder: Union[pathlib.Path, str], + sdfg_name: str, + lib_extension: Optional[str] = None, + folder_mode: Optional[str] = None, +) -> pathlib.Path: + """Returns the supposed location of the compiled library given the boundary conditions. + + :param object_folder: The build folder of the SDFG, i.e. `sdfg.build_folder`. + :param sdfg_name: The name of the SDFG, i.e. `sdfg.name`. + :param lib_extension: The extension of the library, i.e. file extension. + If not given the config option `compiler.library_extension` is used. + :param folder_mode: The save mode for the build folder. If not given the config + option `compiler.build_folder_mode` is used. + """ + if lib_extension is None: + lib_extension = Config.get('compiler', 'library_extension') + if folder_mode is None: + folder_mode = Config.get('compiler', 'build_folder_mode') + + folder_hirarchy = [object_folder] + if folder_mode == 'development': + folder_hirarchy.append('build') + elif folder_mode == 'production': + # Nothing to add, they are on the top. + pass + else: + raise ValueError(f"Unknown folder mode '{folder_mode}' found.") + + return pathlib.Path(os.path.join(*folder_hirarchy, f'lib{sdfg_name}.{lib_extension}')) + + +def _get_stub_library_path(sdfg_lib_path: Union[pathlib.Path, str]) -> pathlib.Path: + """Returns the supposed location of the compiled stub library given the path of the compiled library. + """ + sdfg_lib_path = pathlib.Path(sdfg_lib_path) + parent = sdfg_lib_path.parent + lib_name = sdfg_lib_path.name + assert lib_name.startswith('lib') and len(lib_name) > 3 + + return sdfg_lib_path.parent / ('libdacestub_' + lib_name[3:]) + + +def load_precompiled_sdfg( + folder: Union[pathlib.Path, str], + sdfg: Optional['dace.SDFG'] = None, +) -> csd.CompiledSDFG: + """Loads a precompiled SDFG from ``folder``. + + If ``sdfg`` is not given then the function expects to find the ``program.sdfg(z)`` + dump file inside ``folder``. If the folder does not contain a ``FOLDER_MODE`` file + it assumes that it is an old style ``development`` folder otherwise, the information + from ``FOLDER_MODE`` is consulted. + + :param folder: Path to SDFG output folder, i.e. its build folder. + :param sdfg: If given then ``program.sdfg(z)`` does not need to be present. + :return: A callable CompiledSDFG object. + + :note: If ``sdfg`` is given then it is referenced by the returned ``CompiledSDFG``. + """ + folder = pathlib.Path(folder) + + if not folder.is_dir(): + raise NotADirectoryError(f'Can not load the SDFG from folder ``{folder}``.') + + folder_mode = get_folder_mode(folder) + + # Try to find the sdfg from disc, if not given. + if sdfg is not None: + assert isinstance(sdfg, dace.SDFG) + else: + for name in ['program.sdfgz', 'program.sdfg']: + if (folder / name).exists(): + sdfg = dace.SDFG.from_file(folder / name) + break + else: + raise ValueError(f"Could not locate the SDFG for `{folder}`.") + + return get_program_handle(library_path=get_binary_name(folder, sdfg_name=sdfg.name, folder_mode=folder_mode), + sdfg=sdfg) def _get_or_eval(value_or_function: Union[T, Callable[[], T]]) -> T: @@ -381,32 +638,9 @@ def identical_file_exists(filename: str, file_contents: str): return True -def get_program_handle(library_path, sdfg): - lib = csd.ReloadableDLL(library_path, sdfg.name) - # Load and return the compiled function - return csd.CompiledSDFG(sdfg, lib, sdfg.arg_names) - - -def load_from_file(sdfg, binary_filename): - if not os.path.isfile(binary_filename): - raise FileNotFoundError('File not found: ' + binary_filename) - - # Load the generated library - lib = csd.ReloadableDLL(binary_filename, sdfg.name) - - # Load and return the compiled function - return csd.CompiledSDFG(sdfg, lib, sdfg.arg_names) - - -def get_binary_name(object_folder, object_name, lib_extension=Config.get('compiler', 'library_extension')): - name = None - name = os.path.join(object_folder, "build", 'lib%s.%s' % (object_name, lib_extension)) - return name - - def _run_liveoutput(command, output_stream=None, **kwargs): process = subprocess.Popen(command, stderr=subprocess.STDOUT, stdout=subprocess.PIPE, **kwargs) - output = six.StringIO() + output = io.StringIO() while True: line = process.stdout.readline().rstrip() if not line: diff --git a/dace/codegen/control_flow.py b/dace/codegen/control_flow.py index 1b5384428e..063ca6f542 100644 --- a/dace/codegen/control_flow.py +++ b/dace/codegen/control_flow.py @@ -93,50 +93,54 @@ def _loop_region_to_code(region: LoopRegion, dispatch_state: Callable[[SDFGState expr = '' - if loop.update_statement and loop.init_statement and loop.loop_variable: - lsyms = {} - lsyms.update(symbols) - if codegen.dispatcher.defined_vars.has(loop.loop_variable) and not loop.loop_variable in lsyms: - lsyms[loop.loop_variable] = codegen.dispatcher.defined_vars.get(loop.loop_variable)[1] + lsyms = {} + lsyms.update(symbols) + if (loop.loop_variable and codegen.dispatcher.defined_vars.has(loop.loop_variable) + and not loop.loop_variable in lsyms): + lsyms[loop.loop_variable] = codegen.dispatcher.defined_vars.get(loop.loop_variable)[1] + + if loop.init_statement: init = unparse_interstate_edge(loop.init_statement.code[0], sdfg, codegen=codegen, symbols=lsyms) init = init.strip(';') + else: + init = '' + if loop.update_statement: update = unparse_interstate_edge(loop.update_statement.code[0], sdfg, codegen=codegen, symbols=lsyms) update = update.strip(';') + else: + update = '' - if loop.inverted: - if loop.update_before_condition: + if loop.inverted: + if loop.update_before_condition: + if init: expr += f'{init};\n' - expr += 'do {\n' - expr += _clean_loop_body(control_flow_region_to_code(loop, dispatch_state, codegen, symbols)) + expr += 'do {\n' + expr += _clean_loop_body(control_flow_region_to_code(loop, dispatch_state, codegen, symbols)) + if update: expr += f'{update};\n' - expr += f'}} while({cond});\n' else: - expr += f'{init};\n' - expr += 'while (1) {\n' - expr += _clean_loop_body(control_flow_region_to_code(loop, dispatch_state, codegen, symbols)) - expr += f'if (!({cond}))\n' - expr += 'break;\n' - expr += f'{update};\n' - expr += '}\n' + expr += '\n' + expr += f'}} while({cond});\n' else: - if loop.unroll: - if loop.unroll_factor >= 1: - expr += f'#pragma unroll {loop.unroll_factor}\n' - else: - expr += f'#pragma unroll\n' - expr += f'for ({init}; {cond}; {update}) {{\n' + if init: + expr += f'{init};\n' + expr += 'while (1) {\n' expr += _clean_loop_body(control_flow_region_to_code(loop, dispatch_state, codegen, symbols)) - expr += '\n}\n' + expr += f'if (!({cond}))\n' + expr += 'break;\n' + if update: + expr += f'{update};\n' + expr += '}\n' else: - if loop.inverted: - expr += 'do {\n' - expr += _clean_loop_body(control_flow_region_to_code(loop, dispatch_state, codegen, symbols)) - expr += f'\n}} while({cond});\n' - else: - expr += f'while ({cond}) {{\n' - expr += _clean_loop_body(control_flow_region_to_code(loop, dispatch_state, codegen, symbols)) - expr += '\n}\n' + if loop.unroll: + if loop.unroll_factor >= 1: + expr += f'#pragma unroll {loop.unroll_factor}\n' + else: + expr += f'#pragma unroll\n' + expr += f'for ({init}; {cond}; {update}) {{\n' + expr += _clean_loop_body(control_flow_region_to_code(loop, dispatch_state, codegen, symbols)) + expr += '\n}\n' return expr diff --git a/dace/codegen/cppunparse.py b/dace/codegen/cppunparse.py index af79017760..f80a3f841c 100644 --- a/dace/codegen/cppunparse.py +++ b/dace/codegen/cppunparse.py @@ -71,7 +71,6 @@ from functools import lru_cache import inspect -import six import sys import ast import numpy as np @@ -80,7 +79,7 @@ import sympy import dace from numbers import Number -from six import StringIO +from io import StringIO from dace import dtypes from dace.sdfg import type_inference @@ -202,7 +201,7 @@ def fill(self, text=""): def write(self, text): """Append a piece of text to the current line""" - self.f.write(six.text_type(text)) + self.f.write(str(text)) def enter(self): """Print '{', and increase the indentation.""" @@ -1163,6 +1162,6 @@ def py2cpp(code, expr_semicolon=True, defined_symbols=None): raise NotImplementedError('Unsupported type for py2cpp') -@lru_cache(maxsize=16384) +@lru_cache(maxsize=16384, typed=True) def pyexpr2cpp(expr): return py2cpp(expr, expr_semicolon=False) diff --git a/dace/codegen/instrumentation/gpu_tx_markers.py b/dace/codegen/instrumentation/gpu_tx_markers.py index 02c3454e89..7377fd042e 100644 --- a/dace/codegen/instrumentation/gpu_tx_markers.py +++ b/dace/codegen/instrumentation/gpu_tx_markers.py @@ -9,7 +9,6 @@ from dace.memlet import Memlet from dace.sdfg import nodes, SDFG from dace.sdfg.graph import MultiConnectorEdge -from dace.sdfg.nodes import NestedSDFG from dace.sdfg.scope import is_devicelevel_gpu_kernel from dace.sdfg.sdfg import SDFG from dace.sdfg.state import ControlFlowRegion, SDFGState diff --git a/dace/codegen/instrumentation/papi.py b/dace/codegen/instrumentation/papi.py index 116559039a..661730b924 100644 --- a/dace/codegen/instrumentation/papi.py +++ b/dace/codegen/instrumentation/papi.py @@ -451,7 +451,7 @@ class PAPIUtils(object): def available_counters() -> Dict[str, int]: """ Returns the available PAPI counters on this machine. Only works on - *nix based systems with ``grep`` and ``papi-tools`` installed. + posix based systems with ``grep`` and ``papi-tools`` installed. :return: A set of available PAPI counters in the form of a dictionary mapping from counter name to the number of native hardware diff --git a/dace/codegen/instrumentation/report.py b/dace/codegen/instrumentation/report.py index 1658012d90..42a9ccb8f8 100644 --- a/dace/codegen/instrumentation/report.py +++ b/dace/codegen/instrumentation/report.py @@ -443,4 +443,4 @@ def save(self, filename: str) -> None: report_json['sdfgHash'] = self.sdfg_hash report_json['traceEvents'] = [ev.save() for ev in self.events] with open(filename, 'w') as fp: - json.dump(report_json, fp) + json.dump(report_json, fp, indent=None, separators=(',', ':')) diff --git a/dace/codegen/prettycode.py b/dace/codegen/prettycode.py index d41a22d64c..3894e9727e 100644 --- a/dace/codegen/prettycode.py +++ b/dace/codegen/prettycode.py @@ -3,10 +3,10 @@ nodes. """ import inspect -from six import StringIO +from io import StringIO from dace.config import Config from dace.sdfg.graph import NodeNotFoundError -from dace.sdfg.state import ControlFlowRegion, SDFGState +from dace.sdfg.state import ControlFlowRegion class CodeIOStream(StringIO): diff --git a/dace/codegen/target.py b/dace/codegen/target.py index 57c60b5922..18397c08e8 100644 --- a/dace/codegen/target.py +++ b/dace/codegen/target.py @@ -4,7 +4,7 @@ from typing import List, TYPE_CHECKING import warnings -from dace import memlet as mm, data as dt, dtypes, subsets +from dace import memlet as mm, data as dt, dtypes from dace.sdfg import nodes, SDFG, SDFGState, ScopeSubgraphView, graph as gr from dace.registry import make_registry from dace.codegen.prettycode import CodeIOStream diff --git a/dace/codegen/targets/cpp.py b/dace/codegen/targets/cpp.py index c82e1514ae..1fcd55302b 100644 --- a/dace/codegen/targets/cpp.py +++ b/dace/codegen/targets/cpp.py @@ -9,12 +9,11 @@ import itertools import math import numbers -import sys import warnings import sympy as sp -from six import StringIO -from typing import IO, TYPE_CHECKING, List, Optional, Tuple, Union +from io import StringIO +from typing import IO, TYPE_CHECKING, Optional, Tuple, Union import dace from dace import data, subsets, symbolic, dtypes, memlet as mmlt, nodes @@ -23,10 +22,9 @@ from dace.codegen.dispatcher import DefinedType from dace.codegen.prettycode import CodeIOStream from dace.config import Config -from dace.frontend import operations from dace.frontend.python import astutils from dace.frontend.python.astutils import ExtNodeTransformer, rname, unparse -from dace.sdfg import nodes, graph as gr, utils, propagation +from dace.sdfg import nodes, graph as gr, propagation from dace.properties import LambdaProperty from dace.sdfg import SDFG, is_devicelevel_gpu, SDFGState from dace.sdfg.state import ControlFlowRegion, StateSubgraphView @@ -249,6 +247,14 @@ def ptr(name: str, desc: data.Data, sdfg: SDFG = None, framecode: 'DaCeCodeGener if root in sdfg.arrays and isinstance(sdfg.arrays[root], data.Structure): name = name.replace('.', '->') + if isinstance(desc, data.DistributedDescriptor): # Skip CFG ID mangling for distributed descriptors + name = desc.name + if (desc.transient + and desc.lifetime in (dtypes.AllocationLifetime.Persistent, dtypes.AllocationLifetime.External) + and not is_cuda_codegen_in_device(framecode)): + return f'__state->{name}' + return name + # Special case: If memory is persistent and defined in this SDFG, add state # struct to name if (desc.transient and desc.lifetime in (dtypes.AllocationLifetime.Persistent, dtypes.AllocationLifetime.External)): diff --git a/dace/codegen/targets/cpu.py b/dace/codegen/targets/cpu.py index 23e230f9d0..ce0851c351 100644 --- a/dace/codegen/targets/cpu.py +++ b/dace/codegen/targets/cpu.py @@ -34,7 +34,6 @@ class CPUCodeGen(TargetCodeGenerator): language = "cpp" def _define_sdfg_arguments(self, sdfg, arglist): - # NOTE: Multi-nesting with container arrays must be further investigated. def _visit_structure(struct: data.Structure, args: dict, prefix: str = ''): for k, v in struct.members.items(): @@ -1860,6 +1859,16 @@ def _generate_MapEntry( if node.map.schedule in (dtypes.ScheduleType.CPU_Multicore, dtypes.ScheduleType.CPU_Persistent): raise ValueError("An OpenMP map cannot be unrolled (" + node.map.label + ")") + # A symbolic step whose sign is not statically known is guarded by + # an assert() (a statically-negative step is already rejected by + # SDFG validation). Placed before the OpenMP pragma, which must be + # immediately followed by its loop. + for _, _, _skip in node.map.range: + if (_skip > 0) != True: + result.write( + 'assert((%s) > 0 && "Map %s requires a positive step");\n' % (cpp.sym2cpp(_skip), node.map.label), + cfg, state_id, node) + result.write(map_header, cfg, state_id, node) if node.map.schedule == dtypes.ScheduleType.CPU_Persistent: diff --git a/dace/codegen/targets/cuda.py b/dace/codegen/targets/cuda.py index ed99234de7..3b7a3556ed 100644 --- a/dace/codegen/targets/cuda.py +++ b/dace/codegen/targets/cuda.py @@ -6,7 +6,7 @@ import networkx as nx import sympy -from six import StringIO +from io import StringIO import dace from dace import data as dt, Memlet @@ -549,10 +549,8 @@ def cmake_options(): hip_arch = [ha for ha in hip_arch if ha is not None and len(ha) > 0] flags = Config.get("compiler", "cuda", "hip_args") - flags += ' ' + ' '.join( - '--offload-arch={arch}'.format(arch=arch if arch.startswith("gfx") else "gfx" + arch) - for arch in hip_arch) - options.append("-DEXTRA_HIP_FLAGS=\"{}\"".format(flags)) + options.append(f'-DDACE_HIP_ARCHITECTURES_DEFAULT="{";".join(hip_arch)}"') + options.append("-DCMAKE_HIP_FLAGS=\"{}\"".format(flags)) if Config.get('compiler', 'cpu', 'executable'): host_compiler = make_absolute(Config.get("compiler", "cpu", "executable")) @@ -1482,7 +1480,7 @@ def generate_devicelevel_state(self, sdfg: SDFG, cfg: ControlFlowRegion, state: # otherwise streams do not behave as expected becasue they are # allocated on host side streams_to_reset = [ - node for node in state.data_nodes() if isinstance(node.desc(sdfg), dace.nodes.data.Stream) + node for node in state.data_nodes() if isinstance(node.desc(sdfg), dace.data.Stream) and node.desc(sdfg).lifetime == dtypes.AllocationLifetime.Scope ] for stream in streams_to_reset: @@ -1733,6 +1731,9 @@ def generate_scope(self, sdfg: SDFG, cfg: ControlFlowRegion, dfg_scope: StateSub assert node.gpu_maxnreg is not None and node.gpu_maxnreg >= 0 if node.gpu_maxnreg == 0: maxnreg_str = '' + gpu_min_warps_per_eu = '' + if node.gpu_min_warps_per_eu is not None and node.gpu_min_warps_per_eu > 0: + gpu_min_warps_per_eu = f',{node.gpu_min_warps_per_eu}' # Set kernel launch bounds if node.gpu_launch_bounds == "-1": launch_bounds = '' @@ -1740,9 +1741,9 @@ def generate_scope(self, sdfg: SDFG, cfg: ControlFlowRegion, dfg_scope: StateSub if any(symbolic.issymbolic(b) for b in block_dims): launch_bounds = '' else: - launch_bounds = f'__launch_bounds__({_topy(prod(block_dims))})' + launch_bounds = f'__launch_bounds__({_topy(prod(block_dims))}{gpu_min_warps_per_eu})' else: - launch_bounds = f'__launch_bounds__({node.gpu_launch_bounds})' + launch_bounds = f'__launch_bounds__({node.gpu_launch_bounds}{gpu_min_warps_per_eu})' else: maxnreg_str = f'__maxnreg__({node.gpu_maxnreg})' launch_bounds = '' diff --git a/dace/codegen/targets/framecode.py b/dace/codegen/targets/framecode.py index 75fcd9a8ee..573af7cf21 100644 --- a/dace/codegen/targets/framecode.py +++ b/dace/codegen/targets/framecode.py @@ -592,6 +592,11 @@ def determine_allocation_lifetime(self, top_sdfg: SDFG): access_instances[sdfg.cfg_id] = instances for sdfg, name, desc in top_sdfg.arrays_recursive(include_nested_data=True): + if isinstance(desc, data.DistributedDescriptor): + self._dispatcher.defined_vars.add_global(f'__state->{name}', disp.DefinedType.Scalar, + desc.state_field_dtype.ctype) + self.where_allocated[(sdfg, name)] = top_sdfg + continue # NOTE: Assuming here that all Structure members share transient/storage/lifetime properties. # TODO: Study what is needed in the DaCe stack to ensure this assumption is correct. top_desc = sdfg.arrays[name.split('.')[0]] @@ -873,16 +878,17 @@ def generate_code(self, sdfg: SDFG, schedule: Optional[dtypes.ScheduleType], cfg_id: str = "") -> Tuple[str, str, Set[TargetCodeGenerator], Set[str]]: - """ Generate frame code for a given SDFG, calling registered targets' - code generation callbacks for them to generate their own code. - - :param sdfg: The SDFG to generate code for. - :param schedule: The schedule the SDFG is currently located, or - None if the SDFG is top-level. - :param cfg_id An optional string id given to the SDFG label - :return: A tuple of the generated global frame code, local frame - code, and a set of targets that have been used in the - generation of this SDFG. + """ + Generate frame code for a given SDFG, calling registered targets' + code generation callbacks for them to generate their own code. + + :param sdfg: The SDFG to generate code for. + :param schedule: The schedule the SDFG is currently located, or + None if the SDFG is top-level. + :param cfg_id: An optional string id given to the SDFG label + :return: A tuple of the generated global frame code, local frame + code, and a set of targets that have been used in the + generation of this SDFG. """ if len(cfg_id) == 0 and sdfg.cfg_id != 0: cfg_id = '_%d' % sdfg.cfg_id diff --git a/dace/codegen/targets/mpi.py b/dace/codegen/targets/mpi.py index dfff046191..ebdf133711 100644 --- a/dace/codegen/targets/mpi.py +++ b/dace/codegen/targets/mpi.py @@ -1,6 +1,5 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. from typing import TYPE_CHECKING -import dace from dace import registry, symbolic, dtypes from dace.codegen.prettycode import CodeIOStream from dace.codegen.codeobject import CodeObject diff --git a/dace/codegen/targets/sve/codegen.py b/dace/codegen/targets/sve/codegen.py index 14d2c88e5b..02ffd943cd 100644 --- a/dace/codegen/targets/sve/codegen.py +++ b/dace/codegen/targets/sve/codegen.py @@ -11,20 +11,18 @@ from dace.codegen.prettycode import CodeIOStream from dace.codegen.targets.cpp import sym2cpp from dace import dtypes, memlet as mm -from dace.sdfg import graph, state, find_input_arraynode, find_output_arraynode +from dace.sdfg import graph, state from dace.sdfg.scope import is_in_scope -import itertools from dace.codegen.targets.sve import util as util -from typing import List, Optional +from typing import Optional import copy -from six import StringIO +from io import StringIO import dace.codegen.targets.sve.unparse -from dace import registry, symbolic, dtypes +from dace import dtypes from dace.codegen.targets import cpp as cpp from dace.frontend.operations import detect_reduction_type import dace.symbolic from dace.codegen.targets.cpp import sym2cpp -from dace.sdfg import utils as sdutil from dace.codegen.dispatcher import DefinedType import copy import numpy as np diff --git a/dace/codegen/targets/sve/infer.py b/dace/codegen/targets/sve/infer.py index 97f8329c5f..3feb50c0ac 100644 --- a/dace/codegen/targets/sve/infer.py +++ b/dace/codegen/targets/sve/infer.py @@ -2,13 +2,8 @@ """ Inference: This module patches certain dispatchers in the `type_inference.py`, to better suit SVE. """ -import dace.codegen.targets.sve.util -import numpy as np -import ast from dace import dtypes from dace.sdfg import type_inference -from dace.symbolic import SymExpr -import sympy import sys diff --git a/dace/codegen/targets/sve/type_compatibility.py b/dace/codegen/targets/sve/type_compatibility.py index ef0491361a..0209339315 100644 --- a/dace/codegen/targets/sve/type_compatibility.py +++ b/dace/codegen/targets/sve/type_compatibility.py @@ -10,7 +10,6 @@ - Pointers are mixed with vectors or scalars """ -import dace import dace.dtypes as dtypes from dace.codegen.targets.sve import util as util import collections diff --git a/dace/codegen/targets/sve/unparse.py b/dace/codegen/targets/sve/unparse.py index 06c56ff70d..d842dd5945 100644 --- a/dace/codegen/targets/sve/unparse.py +++ b/dace/codegen/targets/sve/unparse.py @@ -3,13 +3,12 @@ AST to SVE: This module is responsible for converting an AST into SVE code. """ -from dace.frontend.python.wrappers import stream import dace import ast from dace.codegen import cppunparse -from dace.sdfg import nodes, SDFG, SDFGState, ScopeSubgraphView, graph as gr -from typing import IO, Tuple, Union -from dace import registry, symbolic, dtypes +from dace.sdfg import nodes, SDFG +from typing import IO +from dace import dtypes from dace.codegen.targets.sve import preprocess as preprocess from dace.codegen.targets.sve import util as util import dace.frontend.python.astutils as astutils diff --git a/dace/codegen/tools/get_hip_arch.cpp b/dace/codegen/tools/get_hip_arch.cpp index 8de47aa0b1..17bfef963e 100644 --- a/dace/codegen/tools/get_hip_arch.cpp +++ b/dace/codegen/tools/get_hip_arch.cpp @@ -1,6 +1,7 @@ // Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. #include +#include #include #include #include diff --git a/dace/config_schema.yml b/dace/config_schema.yml index 2b05d45232..7cd8979d7a 100644 --- a/dace/config_schema.yml +++ b/dace/config_schema.yml @@ -115,6 +115,16 @@ required: If enabled, does not recompile code generated from SDFGs if shared library (.so/.dll) file is present. + build_folder_mode: + type: str + default: development + title: Save mode for the build folder + description: > + Selects which content should be saved in the build folder. + Two modes are currently supported: `development`, that includes + everything; and `production`, that saves only the compiled library + and the folder mode file. + library_prefix: type: str default: "" @@ -201,7 +211,7 @@ required: description: > Specify the default data types to use in generating code. If "Python", Python's semantics will be followed (i.e., `float` and `int` - are represented using 64 bits). If the property is set to "C", C's semantcs will be + are represented using 64 bits). If the property is set to "C", C's semantics will be used (`float` and `int` are represented using 32bits). unique_functions: @@ -251,6 +261,18 @@ required: If set, specifies additional arguments to the initial invocation of ``cmake``. + lineinfo: + type: str + default: "inspect" + title: Add line info + description: > + Wether or not to add line info from the parsed code in the + generated SDFG. Valid options are `inspect` and `none`. + "inspect": During parsing, inspect the python call stack and + automatically add line info from the parsed source code in the + resulting SDFG. + "none": Do not save any line info in the resulting SDFG. + ############################################# # CPU compiler cpu: diff --git a/dace/data/__init__.py b/dace/data/__init__.py index 4620474f01..b87607c9fc 100644 --- a/dace/data/__init__.py +++ b/dace/data/__init__.py @@ -54,6 +54,9 @@ # ML-related data descriptors from dace.data.ml import ParameterArray +# Distributed data descriptors +from dace.data.distributed import DistributedDescriptor, ProcessGrid, SubArray, RedistrArray + # Descriptor creation and array creation from descriptors from dace.data.creation import ( create_datadescriptor, @@ -101,6 +104,11 @@ 'Offset', # ML descriptors 'ParameterArray', + # Distributed descriptors + 'DistributedDescriptor', + 'ProcessGrid', + 'SubArray', + 'RedistrArray', # Functions 'create_datadescriptor', 'make_array_from_descriptor', diff --git a/dace/data/core.py b/dace/data/core.py index 4312f7ebc5..c19a221b2c 100644 --- a/dace/data/core.py +++ b/dace/data/core.py @@ -773,7 +773,7 @@ class Stream(Data): """ Stream (or stream array) data descriptor. """ # Properties - offset = ListProperty(element_type=symbolic.pystr_to_symbolic) + offset = ListProperty(element_type=sp.Basic) buffer_size = SymbolicProperty(desc="Size of internal buffer.", default=0) def __init__(self, diff --git a/dace/data/ctypes_interop.py b/dace/data/ctypes_interop.py index d3f2f370cc..4cd38a7238 100644 --- a/dace/data/ctypes_interop.py +++ b/dace/data/ctypes_interop.py @@ -7,13 +7,16 @@ import ctypes import warnings -from typing import Any, Dict, List, Optional +from typing import Any, Dict, List, Optional, TYPE_CHECKING import numpy as np import sympy as sp from dace import config, dtypes, symbolic +if TYPE_CHECKING: + from dace.data import Data + def make_ctypes_argument(arg: Any, argtype: 'Data', diff --git a/dace/distr_types.py b/dace/data/distributed.py similarity index 86% rename from dace/distr_types.py rename to dace/data/distributed.py index 735a31dbd0..ee4dc27fc6 100644 --- a/dace/distr_types.py +++ b/dace/data/distributed.py @@ -1,25 +1,60 @@ # Copyright 2019-2022 ETH Zurich and the DaCe authors. All rights reserved. """ A module that contains type definitions for distributed SDFGs. """ +import copy from numbers import Integral -from typing import Sequence, Union +from typing import Sequence, Set, Union import dace.dtypes as dtypes from dace import symbolic, serialize -from dace.properties import (Property, make_properties, ShapeProperty, SymbolicProperty, TypeClassProperty, - ListProperty) +from dace.data.core import Data +from dace.properties import Property, make_properties, ShapeProperty, SymbolicProperty, ListProperty ShapeType = Sequence[Union[Integral, str, symbolic.symbol, symbolic.SymExpr, symbolic.sympy.Basic]] RankType = Union[Integral, str, symbolic.symbol, symbolic.SymExpr, symbolic.sympy.Basic] +class DistributedDescriptor(Data): + """Base class for distributed communication descriptors stored in an SDFG.""" + + @property + def state_field_dtype(self): + return self.dtype + + def clone(self): + return copy.deepcopy(self) + + def is_equivalent(self, other): + return type(self) is type(other) and self == other + + def as_arg(self, with_types=True, for_call=False, name=None): + raise TypeError(f'{type(self).__name__} descriptors are not SDFG call arguments') + + def as_python_arg(self, with_types=True, for_call=False, name=None): + raise TypeError(f'{type(self).__name__} descriptors are not Python call arguments') + + @property + def offset(self): + shp = getattr(self, 'shape', []) + return [0] * len(shp) + + +def _symbols_from_shape(shape) -> Set[symbolic.SymbolicType]: + result: Set[symbolic.SymbolicType] = set() + for s in shape: + if isinstance(s, symbolic.sympy.Basic): + result |= set(s.free_symbols) + return result + + @make_properties -class ProcessGrid(object): +class ProcessGrid(DistributedDescriptor): """ - Process-grids implement cartesian topologies similarly to cartesian communicators created with [MPI_Cart_create](https://www.mpich.org/static/docs/latest/www3/MPI_Cart_create.html) + Process-grids implement cartesian topologies similarly to cartesian communicators created with + [MPI_Cart_create](https://www.mpich.org/static/docs/latest/www3/MPI_Cart_create.html) and [MPI_Cart_sub](https://www.mpich.org/static/docs/v3.2/www3/MPI_Cart_sub.html). The boolean property `is_subgrid` provides a switch between "parent" process-grids (equivalent to communicators - create with `MPI_Cart_create`) and sub-grids (equivalent to communicators created with `MPI_Cart_sub`). + created with `MPI_Cart_create`) and sub-grids (equivalent to communicators created with `MPI_Cart_sub`). If `is_subgrid` is false, a "parent" process-grid is created. The `shape` property is equivalent to the `dims` parameter of `MPI_Cart_create`. The other properties are ignored. All "parent" process-grids spawn out of @@ -38,8 +73,7 @@ class ProcessGrid(object): """ name = Property(dtype=str, desc="The process-grid's name.") - is_subgrid = Property(dtype=bool, default=False, desc="If true, spanws sub-grids out of the parent process-grid.") - shape = ShapeProperty(default=[], desc="The process-grid's shape.") + is_subgrid = Property(dtype=bool, default=False, desc="If true, spawns sub-grids out of the parent process-grid.") parent_grid = Property(dtype=str, allow_none=True, default=None, @@ -68,14 +102,21 @@ def __init__(self, self.name = name self.is_subgrid = is_subgrid if is_subgrid: - self.parent_grid = parent_grid.name + self.parent_grid = parent_grid.name if isinstance(parent_grid, ProcessGrid) else parent_grid self.color = color self.exact_grid = exact_grid - self.shape = [parent_grid.shape[i] for i, remain in enumerate(color) if remain] + if isinstance(parent_grid, ProcessGrid): + shape = [parent_grid.shape[i] for i, remain in enumerate(color) if remain] + else: + shape = shape or [] else: - self.shape = shape + self.parent_grid = None + self.color = None + self.exact_grid = None + shape = shape or [] self.root = root - self._validate() + super().__init__(dtypes.opaque('MPI_Comm'), shape, True, dtypes.StorageType.Default, None, + dtypes.AllocationLifetime.Persistent, None) def validate(self): """ Validate the correctness of this object. @@ -96,10 +137,6 @@ def _validate(self): raise ValueError('Color must have only logical true (1) or false (0) values.') return True - @property - def dtype(self): - return type(self) - def to_json(self): attrs = serialize.all_properties_to_json(self) retdict = {"type": type(self).__name__, "attributes": attrs} @@ -110,6 +147,9 @@ def from_json(cls, json_obj, context=None): # Create dummy object ret = cls('tmp', False, []) serialize.set_properties_from_json(ret, json_obj, context=context) + ret.dtype = dtypes.opaque('MPI_Comm') + ret.transient = True + ret.lifetime = dtypes.AllocationLifetime.Persistent # Check validity now ret.validate() return ret @@ -117,7 +157,7 @@ def from_json(cls, json_obj, context=None): def init_code(self): """ Outputs MPI allocation/initialization code for the process-grid. It is assumed that the following variables exist in the SDFG program's state: - - MPI_Comm {self.name}_comm + - MPI_Comm {self.name} - MPI_Group {self.name}_group - int {self.name}_rank - int {self.name}_size @@ -138,11 +178,11 @@ def init_code(self): __state->{self.name}_valid = false; if (__state->{self.parent_grid}_valid) {{ int {self.name}_remain[{len(self.color)}] = {{{', '.join(['1' if c else '0' for c in self.color])}}}; - MPI_Cart_sub(__state->{self.parent_grid}_comm, {self.name}_remain, &__state->{self.name}_comm); - MPI_Comm_group(__state->{self.name}_comm, &__state->{self.name}_group); - MPI_Comm_rank(__state->{self.name}_comm, &__state->{self.name}_rank); - MPI_Comm_size(__state->{self.name}_comm, &__state->{self.name}_size); - MPI_Cart_coords(__state->{self.name}_comm, __state->{self.name}_rank, {len(self.shape)}, __state->{self.name}_coords); + MPI_Cart_sub(__state->{self.parent_grid}, {self.name}_remain, &__state->{self.name}); + MPI_Comm_group(__state->{self.name}, &__state->{self.name}_group); + MPI_Comm_rank(__state->{self.name}, &__state->{self.name}_rank); + MPI_Comm_size(__state->{self.name}, &__state->{self.name}_size); + MPI_Cart_coords(__state->{self.name}, __state->{self.name}_rank, {len(self.shape)}, __state->{self.name}_coords); """ if self.exact_grid is not None: tmp += f""" @@ -164,12 +204,12 @@ def init_code(self): tmp += f"__state->{self.name}_dims[{i}] = {s};\n" tmp += f""" int {self.name}_periods[{len(self.shape)}] = {{0}}; - MPI_Cart_create(MPI_COMM_WORLD, {len(self.shape)}, __state->{self.name}_dims, {self.name}_periods, 0, &__state->{self.name}_comm); - if (__state->{self.name}_comm != MPI_COMM_NULL) {{ - MPI_Comm_group(__state->{self.name}_comm, &__state->{self.name}_group); - MPI_Comm_rank(__state->{self.name}_comm, &__state->{self.name}_rank); - MPI_Comm_size(__state->{self.name}_comm, &__state->{self.name}_size); - MPI_Cart_coords(__state->{self.name}_comm, __state->{self.name}_rank, {len(self.shape)}, __state->{self.name}_coords); + MPI_Cart_create(MPI_COMM_WORLD, {len(self.shape)}, __state->{self.name}_dims, {self.name}_periods, 0, &__state->{self.name}); + if (__state->{self.name} != MPI_COMM_NULL) {{ + MPI_Comm_group(__state->{self.name}, &__state->{self.name}_group); + MPI_Comm_rank(__state->{self.name}, &__state->{self.name}_rank); + MPI_Comm_size(__state->{self.name}, &__state->{self.name}_size); + MPI_Cart_coords(__state->{self.name}, __state->{self.name}_rank, {len(self.shape)}, __state->{self.name}_coords); __state->{self.name}_valid = true; }} else {{ __state->{self.name}_group = MPI_GROUP_NULL; @@ -185,13 +225,13 @@ def exit_code(self): return f""" if (__state->{self.name}_valid) {{ MPI_Group_free(&__state->{self.name}_group); - MPI_Comm_free(&__state->{self.name}_comm); + MPI_Comm_free(&__state->{self.name}); }} """ @make_properties -class SubArray(object): +class SubArray(DistributedDescriptor): """ Sub-arrays describe subsets of Arrays (see `dace::data::Array`) for purposes of distributed communication. They are implemented with [MPI_Type_create_subarray](https://www.mpich.org/static/docs/v3.2/www3/MPI_Type_create_subarray.html). @@ -208,8 +248,6 @@ class SubArray(object): """ name = Property(dtype=str, desc="The type's name.") - dtype = TypeClassProperty(default=dtypes.int32) - shape = ShapeProperty(default=[], desc="The array's shape.") subshape = ShapeProperty(default=[], desc="The sub-array's shape.") pgrid = Property(dtype=str, allow_none=True, @@ -229,12 +267,15 @@ def __init__(self, pgrid: str = None, correspondence: Sequence[Integral] = None): self.name = name - self.dtype = dtype - self.shape = shape self.subshape = subshape self.pgrid = pgrid self.correspondence = correspondence or list(range(len(shape))) - self._validate() + super().__init__(dtype, shape, True, dtypes.StorageType.Default, None, dtypes.AllocationLifetime.Persistent, + None) + + @property + def state_field_dtype(self): + return dtypes.opaque('MPI_Datatype') def validate(self): """ Validate the correctness of this object. @@ -269,10 +310,18 @@ def from_json(cls, json_obj, context=None): # Create dummy object ret = cls('tmp', dtypes.int8, [], [], 'tmp', []) serialize.set_properties_from_json(ret, json_obj, context=context) + ret.transient = True + ret.lifetime = dtypes.AllocationLifetime.Persistent # Check validity now ret.validate() return ret + def used_symbols(self, all_symbols: bool) -> Set[symbolic.SymbolicType]: + result = super().used_symbols(all_symbols) + if self.transient or all_symbols: + result |= _symbols_from_shape(self.subshape) + return result + def init_code(self): """ Outputs MPI allocation/initialization code for the sub-array MPI datatype ONLY if the process-grid is set. It is assumed that the following variables exist in the SDFG program's state: @@ -353,7 +402,7 @@ def exit_code(self): @make_properties -class RedistrArray(object): +class RedistrArray(DistributedDescriptor): """ Describes the redistribution of an Array from one process-grid and sub-array descriptor (`array_a`) to another (`array_b`). The redistribution is implemented with MPI datatypes (see [MPI_Type_create_subarray](https://www.mpich.org/static/docs/v3.2/www3/MPI_Type_create_subarray.html) @@ -369,7 +418,12 @@ def __init__(self, name: str, array_a: str, array_b: str): self.name = name self.array_a = array_a self.array_b = array_b - self._validate() + super().__init__(dtypes.opaque('dace::comm::RedistrArray'), [], True, dtypes.StorageType.Default, None, + dtypes.AllocationLifetime.Persistent, None) + + @property + def state_field_dtype(self): + return dtypes.opaque('MPI_Datatype') def validate(self): """ Validate the correctness of this object. @@ -392,6 +446,9 @@ def from_json(cls, json_obj, context=None): # Create dummy object ret = cls('tmp', 'tmp', 'tmp') serialize.set_properties_from_json(ret, json_obj, context=context) + ret.dtype = dtypes.opaque('dace::comm::RedistrArray') + ret.transient = True + ret.lifetime = dtypes.AllocationLifetime.Persistent # Check validity now ret.validate() return ret diff --git a/dace/data/tensor.py b/dace/data/tensor.py index a09602757d..ec59bfe1b6 100644 --- a/dace/data/tensor.py +++ b/dace/data/tensor.py @@ -662,11 +662,12 @@ def __init__(self, self.indices, self.index_ordering = list(indices), list(index_ordering) num_dims = len(tensor_shape) - dimension_order = [idx for idx in self.index_ordering if isinstance(idx, int)] + # Only concrete Python/SymPy integers refer to tensor dimensions; symbolic values denote unrelated levels. + dimension_order = [int(idx) for idx in index_ordering if isinstance(idx, (int, symbolic.sympy.Integer))] - # all tensor dimensions must occure exactly once in indices + # all tensor dimensions must occur exactly once in indices if not sorted(dimension_order) == list(range(num_dims)): - raise TypeError((f"All tensor dimensions must be refferenced exactly once in " + raise TypeError((f"All tensor dimensions must be referenced exactly once in " f"tensor indices. (referenced dimensions: {dimension_order}; " f"tensor dimensions: {list(range(num_dims))})")) diff --git a/dace/dtypes.py b/dace/dtypes.py index 1ecdc90a82..6d192efba5 100644 --- a/dace/dtypes.py +++ b/dace/dtypes.py @@ -15,6 +15,7 @@ from enum import auto, Enum from dace.attr_enum import ExtensibleAttributeEnum from dace.registry import undefined_safe_enum +from dace.version import __version__ @undefined_safe_enum @@ -379,10 +380,10 @@ def __call__(self, *args, **kwargs): return self.type(*args, **kwargs) def __eq__(self, other): - return other is not None and self.ctype == other.ctype + return other is not None and self.ctype == getattr(other, 'ctype', False) def __ne__(self, other): - return other is not None and self.ctype != other.ctype + return other is not None and self.ctype != getattr(other, 'ctype', False) def __getitem__(self, s): """ This is syntactic sugar that allows us to define an array type @@ -721,8 +722,6 @@ def from_json(json_obj, context=None): if json_obj['type'] != "struct": raise TypeError("Invalid type for struct") - import dace.serialize # Avoid import loop - ret = struct(json_obj['name']) ret._data = {k: json_to_typeclass(v, context) for k, v in json_obj['data']} ret._length = {k: v for k, v in json_obj['length']} @@ -757,7 +756,7 @@ def _parse_field_and_types(self, **fields_and_types): def as_ctypes(self): """ Returns the ctypes version of the typeclass. """ - self_as_json = json.dumps(self.to_json(), sort_keys=True) + self_as_json = json.dumps(self.to_json(), sort_keys=True, separators=(',', ':')) if self_as_json in struct.STRUCT_CTYPES: return struct.STRUCT_CTYPES[self_as_json] # Populate the ctype fields for the struct class. @@ -872,7 +871,7 @@ def __init__(self, return_types, *variadic_args): elif isinstance(arg, data.Data): pass elif isinstance(arg, str): - arg = json_to_typeclass(arg) + arg = json_to_typeclass(arg, {'version': __version__}) else: raise TypeError("Cannot resolve type from: {}".format(arg)) self.input_types.append(arg) @@ -1103,7 +1102,7 @@ def from_json(json_obj, context=None): import dace.serialize # Avoid import loop - return callback([json_to_typeclass(rettype) if rettype else None for rettype in rettypes], + return callback([json_to_typeclass(rettype, context) if rettype else None for rettype in rettypes], *(dace.serialize.from_json(arg, context) for arg in json_obj['arguments'])) def __str__(self): @@ -1263,6 +1262,36 @@ def dtype_to_typeclass(dtype=None): complex128: "dace::complex128" } +TYPECLASS_TO_LITERAL_SUFFIX = { + int8: 'i8', + int16: 'i16', + int32: 'i32', + int64: 'i64', + uint8: 'u8', + uint16: 'u16', + uint32: 'u32', + uint64: 'u64', + float16: 'f16', + float32: 'f32', + float64: 'f64', +} + +LITERAL_SUFFIX_TO_TYPECLASS = {v: k for k, v in TYPECLASS_TO_LITERAL_SUFFIX.items()} + +# Inverse of _CTYPES for numeric types, to parse C++ cast fallbacks. +CTYPE_TO_TYPECLASS = { + ctype: dtype_to_typeclass(nptype) + for nptype, ctype in _CTYPES.items() + if nptype is not None and (numpy.issubdtype(nptype, numpy.integer) or numpy.issubdtype(nptype, numpy.floating)) +} + +TYPECLASS_TO_CPP_LITERAL_SUFFIX = { + float32: 'f', + uint32: 'U', + int64: 'LL', + uint64: 'ULL', +} + TYPECLASS_STRINGS = [ "int", "float", "complex", "bool", "bool_", "int8", "int16", "int32", "int64", "uint8", "uint16", "uint32", "uint64", "float16", "float32", "float64", "complex64", "complex128" @@ -1270,6 +1299,31 @@ def dtype_to_typeclass(dtype=None): INTEGER_TYPES = [bool, bool_, int8, int16, int32, int64, uint8, uint16, uint32, uint64] + +def typeclass_to_literal_suffix(dtype): + return TYPECLASS_TO_LITERAL_SUFFIX[dtype] + + +def literal_suffix_to_typeclass(suffix): + return LITERAL_SUFFIX_TO_TYPECLASS[suffix] + + +def cpp_typed_literal(value, dtype): + """Format a typed constant as a native C++ literal when possible. + + :param value: Constant value to emit. + :param dtype: DaCe typeclass of the constant. + :return: C++ literal string such as ``'1ULL'`` or ``'1.5f'``, or + ``None`` if the dtype should be emitted via an explicit cast. + """ + suffix = TYPECLASS_TO_CPP_LITERAL_SUFFIX.get(dtype, '') + if not suffix: + return None + if dtype == float32: + return f'{float(value)}{suffix}' + return f'{int(value)}{suffix}' + + ####################################################### # Allowed types @@ -1327,28 +1381,42 @@ class DebugInfo: """ Source code location identifier of a node/edge in an SDFG. Used for IDE and debugging purposes. """ - def __init__(self, start_line, start_column=0, end_line=-1, end_column=0, filename=None): + def __init__(self, start_line, start_column=0, end_line=-1, end_column=0, filename=None, file_index=None): self.start_line = start_line self.end_line = end_line if end_line >= 0 else start_line self.start_column = start_column self.end_column = end_column + + if filename is not None and file_index is not None: + raise ValueError("Cannot specify both filename and file_index in DebugInfo") + self.filename = filename + self.file_index = file_index # NOTE: Manually marking as serializable to avoid an import loop # The data structure is a property on its own (pointing to a range of code), # so it is serialized as a dictionary directly. def to_json(self): - return dict(type='DebugInfo', - start_line=self.start_line, - end_line=self.end_line, - start_column=self.start_column, - end_column=self.end_column, - filename=self.filename) + result = {'type': 'DebugInfo'} + if self.start_line is not None: + result['start_line'] = self.start_line + if self.end_line is not None and self.end_line != self.start_line: + result['end_line'] = self.end_line + if self.start_column: + result['start_column'] = self.start_column + if self.end_column and self.end_column != self.start_column: + result['end_column'] = self.end_column + if self.file_index is not None: + result['file_index'] = self.file_index + elif self.filename: + result['filename'] = self.filename + + return result @staticmethod - def from_json(json_obj, context=None): - return DebugInfo(json_obj['start_line'], json_obj['start_column'], json_obj['end_line'], json_obj['end_column'], - json_obj['filename']) + def from_json(json_obj: dict[str, int | str], context=None): + return DebugInfo(json_obj.get('start_line'), json_obj.get('start_column', 0), json_obj.get('end_line', -1), + json_obj.get('end_column', 0), json_obj.get('filename'), json_obj.get('file_index')) def __deepcopy__(self, memo) -> 'DebugInfo': """Performs a `deepcopy` of `self`. diff --git a/dace/frontend/common/einsum.py b/dace/frontend/common/einsum.py index d33b0150f3..4629e3c5b6 100644 --- a/dace/frontend/common/einsum.py +++ b/dace/frontend/common/einsum.py @@ -1,6 +1,5 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. """ Classes to handle Einstein-notation sums (einsum) as a library node. """ -from functools import reduce from itertools import chain from string import ascii_letters from typing import Dict, List, Optional @@ -11,10 +10,8 @@ from dace import dtypes, subsets, symbolic from dace.utils import prod from dace.sdfg.nodes import AccessNode -from dace.sdfg import SDFG, SDFGState, InterstateEdge +from dace.sdfg import SDFG, SDFGState from dace.memlet import Memlet -from dace.frontend.common import op_repository as oprepo -from dace.frontend.python.common import StringLiteral def _is_sequential(index_list): diff --git a/dace/frontend/common/op_repository.py b/dace/frontend/common/op_repository.py index a9102d122a..74924cba7b 100644 --- a/dace/frontend/common/op_repository.py +++ b/dace/frontend/common/op_repository.py @@ -1,7 +1,7 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. import itertools from typing import Any, Callable, Dict, List, Optional, Tuple, Type, Union -from dace.dtypes import paramdec, deduplicate +from dace.dtypes import paramdec MethodType = Callable[..., Tuple[str]] diff --git a/dace/frontend/fortran/ast_components.py b/dace/frontend/fortran/ast_components.py index abdd179a2f..9097460b46 100644 --- a/dace/frontend/fortran/ast_components.py +++ b/dace/frontend/fortran/ast_components.py @@ -1,12 +1,12 @@ # Copyright 2019-2023 ETH Zurich and the DaCe authors. All rights reserved. -from fparser.two import Fortran2008 as f08 +# NOTE: Fortran2003 needs to be imported before Fortran2008 (circular import otherwise). from fparser.two import Fortran2003 as f03 +from fparser.two import Fortran2008 as f08 from fparser.two import symbol_table -import copy from dace.frontend.fortran import ast_internal_classes -from dace.frontend.fortran.ast_internal_classes import FNode, Name_Node -from typing import Any, List, Tuple, Type, TypeVar, Union, overload, TYPE_CHECKING +from dace.frontend.fortran.ast_internal_classes import Name_Node +from typing import Any, List, Type, TypeVar, Union, overload, TYPE_CHECKING if TYPE_CHECKING: from dace.frontend.fortran.intrinsics import FortranIntrinsics diff --git a/dace/frontend/fortran/ast_internal_classes.py b/dace/frontend/fortran/ast_internal_classes.py index 276959337c..0d5e9427f8 100644 --- a/dace/frontend/fortran/ast_internal_classes.py +++ b/dace/frontend/fortran/ast_internal_classes.py @@ -1,5 +1,5 @@ # Copyright 2019-2023 ETH Zurich and the DaCe authors. All rights reserved. -from typing import Any, List, Optional, Tuple, Type, TypeVar, Union, overload +from typing import List, Optional, Union # The node class is the base class for all nodes in the AST. It provides attributes including the line number and fields. # Attributes are not used when walking the tree, but are useful for debugging and for code generation. diff --git a/dace/frontend/fortran/ast_transforms.py b/dace/frontend/fortran/ast_transforms.py index 6b7801a4a8..b8405a2135 100644 --- a/dace/frontend/fortran/ast_transforms.py +++ b/dace/frontend/fortran/ast_transforms.py @@ -1,7 +1,7 @@ # Copyright 2023 ETH Zurich and the DaCe authors. All rights reserved. -from dace.frontend.fortran import ast_components, ast_internal_classes -from typing import Dict, List, Optional, Tuple, Set +from dace.frontend.fortran import ast_internal_classes +from typing import Dict, List, Optional, Tuple import copy diff --git a/dace/frontend/fortran/fortran_parser.py b/dace/frontend/fortran/fortran_parser.py index c89b179a83..4f74bf3417 100644 --- a/dace/frontend/fortran/fortran_parser.py +++ b/dace/frontend/fortran/fortran_parser.py @@ -756,10 +756,10 @@ def subroutine2sdfg(self, node: ast_internal_classes.Subroutine_Subprogram_Node, if not found: if local_name.name in write_names: ast_utils.add_memlet_write(substate, mapped_name, internal_sdfg, - self.name_mapping[new_sdfg][local_name.name], memlet) + self.name_mapping[new_sdfg][local_name.name], dpcp(memlet)) if local_name.name in read_names: ast_utils.add_memlet_read(substate, mapped_name, internal_sdfg, - self.name_mapping[new_sdfg][local_name.name], memlet) + self.name_mapping[new_sdfg][local_name.name], dpcp(memlet)) for i in addedmemlets: diff --git a/dace/frontend/ml/onnx/importer.py b/dace/frontend/ml/onnx/importer.py index 226b3e821f..513a460262 100644 --- a/dace/frontend/ml/onnx/importer.py +++ b/dace/frontend/ml/onnx/importer.py @@ -50,7 +50,7 @@ import copy import tempfile from itertools import chain, repeat -from typing import Any, Callable, Dict, List, Optional, OrderedDict, Tuple, Union +from typing import Any, Dict, List, Optional, OrderedDict, Tuple, Union import numpy as np diff --git a/dace/frontend/ml/tensorflow/tensorflow.py b/dace/frontend/ml/tensorflow/tensorflow.py index ef6cfdb409..cae1c6b9a2 100644 --- a/dace/frontend/ml/tensorflow/tensorflow.py +++ b/dace/frontend/ml/tensorflow/tensorflow.py @@ -4,11 +4,8 @@ from functools import partial import numpy as np -import os -import pickle import re -import math -from typing import Any, List +from typing import Any import warnings import dace @@ -16,9 +13,8 @@ from dace import SDFG, SDFGState, dtypes from dace.data import Scalar from dace.sdfg.nodes import Tasklet, NestedSDFG -from dace.symbolic import symstr, SymExpr +from dace.symbolic import symstr from .winograd import winograd_convolution -from .transformations.redundant_array import TensorflowRedundantArray try: import tensorflow as tf @@ -364,7 +360,6 @@ def compile(self, nodes, gpu, name=None, patterns=[], validate=False, permissive :return: A function that receives a feed_dict, evaluates the nodes, and returns a tuple of values in the same order as nodes. """ - from dace.config import Config # Create a unique name for this session if name is None: @@ -787,7 +782,8 @@ def visit_Const(self, node): # If node already present set it non transient, otherwise add node if not nodeArray: dtype = dace.typeclass(_tensortype(node)) - state.add_array(label, shape, dtype, lifetime=dtypes.AllocationLifetime.SDFG) + self.graph.add_array(label, shape, dtype, lifetime=dtypes.AllocationLifetime.SDFG) + state.add_access(label) else: nodeArray[0].desc(self.graph).transient = False @@ -814,7 +810,8 @@ def visit_VariableV2(self, node): outputNode.desc(self.graph).transient = False except (LookupError): dtype = dace.typeclass(_tensortype(node)) - state.add_array(label, shape, dtype) + self.graph.add_array(label, shape, dtype) + state.add_access(label) # If not already added to the varDict, add a placeholder # zero-initialized array to it so a value error is not triggered. @@ -834,10 +831,8 @@ def visit_Assign(self, node): except (LookupError): dtype = dace.typeclass(_tensortype(node.inputs[1])) shape = dace.properties.ShapeProperty.from_string(str(_tensorshape(node.inputs[1]))) - fillNode = state.add_transient(name=label, - shape=shape, - dtype=dtype, - lifetime=dtypes.AllocationLifetime.SDFG) + self.graph.add_transient(name=label, shape=shape, dtype=dtype, lifetime=dtypes.AllocationLifetime.SDFG) + fillNode = state.add_access(label) label = string_builder(node.inputs[0].name) try: @@ -847,10 +842,8 @@ def visit_Assign(self, node): shape = dace.properties.ShapeProperty.from_string(str(_tensorshape(node.inputs[1]))) assert dtype is not None assert shape is not None - emptyNode = state.add_transient(name=label, - shape=shape, - dtype=dtype, - lifetime=dtypes.AllocationLifetime.SDFG) + self.graph.add_transient(name=label, shape=shape, dtype=dtype, lifetime=dtypes.AllocationLifetime.SDFG) + emptyNode = state.add_access(label) dims = self.get_default_dims(node.inputs[1]) memlet = Memlet.simple(emptyNode, ",".join(dims)) state.add_edge(fillNode, None, emptyNode, None, memlet) @@ -899,7 +892,8 @@ def visit_Placeholder(self, node): # create node for the training examples shape = dace.properties.ShapeProperty.from_string(",".join(inputShape)) dtype = _tensortype(node) - inputNode = state.add_array(name=label + "_Inp", shape=shape, dtype=dace.typeclass(dtype)) + self.graph.add_array(name=label + "_Inp", shape=shape, dtype=dace.typeclass(dtype)) + inputNode = state.add_access(label + "_Inp") # create and add map mapDict = dict(zip(inputParams, inputDims)) @@ -1120,16 +1114,18 @@ def visit_FusedBatchNorm(self, node): for i in inpTensorNode.desc(self.graph).shape[:-1]: normalisationScalar *= i normalisationScalar = float(normalisationScalar) - sumInputs = self.state.add_transient( + self.graph.add_transient( "sigma_x_" + local_ctr, _tensorshape(node.outputs[2]), _tensortype(node.outputs[2]), ) - sumSquareInputs = self.state.add_transient( + sumInputs = self.state.add_access("sigma_x_" + local_ctr) + self.graph.add_transient( "sigma_x2_" + local_ctr, _tensorshape(node.outputs[2]), _tensortype(node.outputs[2]), ) + sumSquareInputs = self.state.add_access("sigma_x2_" + local_ctr) ######## Maps ################### nhwcMapBounds = dict(zip(inpTensorParams, inpTensorDims)) cMapBounds = dict(zip([inpTensorParams[0]], [str(inpTensorDims[-1])])) @@ -1239,16 +1235,18 @@ def visit_FusedBatchNormGrad(self, node): gammaGrads = outputList[1] betaGrads = outputList[2] ############################TRANSIENTS########################################## - gammaPrime = self.state.add_transient( + self.graph.add_transient( "gamma_prime" + local_ctr, _tensorshape(node.outputs[1]), _tensortype(node.outputs[1]), ) - betaPrime = self.state.add_transient( + gammaPrime = self.state.add_access("gamma_prime" + local_ctr) + self.graph.add_transient( "beta_prime" + local_ctr, _tensorshape(node.outputs[2]), _tensortype(node.outputs[2]), ) + betaPrime = self.state.add_access("beta_prime" + local_ctr) ###############################MAPS############################################## # channelMapLabel = string_builder(node.type) + "_outer" # channelMapEntry, channelMapExit = self.state.add_map( @@ -1546,7 +1544,8 @@ def visit_ReadVariableOp(self, node): except (LookupError): dtype = dace.typeclass(_tensortype(node.outputs[0])) shape = dace.properties.ShapeProperty.from_string(str(_tensorshape(node.outputs[0]))) - inputNode = state.add_transient(name=label, shape=shape, dtype=dtype) + self.graph.add_transient(name=label, shape=shape, dtype=dtype) + inputNode = state.add_access(label) outputNode = self.create_and_add_output_node(node)[0] outputDims = self.get_default_dims(node.outputs[0]) @@ -1673,7 +1672,8 @@ def visit_ResourceApplyGradientDescent(self, node): except (LookupError): dtype = dace.typeclass(_tensortype(node.inputs[2])) shape = dace.properties.ShapeProperty.from_string(str(_tensorshape(node.inputs[2]))) - inputNode = state.add_transient(name=label, shape=shape, dtype=dtype) + self.graph.add_transient(name=label, shape=shape, dtype=dtype) + inputNode = state.add_access(label) inputNodes.append(inputNode) inputParams.append(self.get_default_params(node.inputs[2])) inputDims.append(self.get_default_dims(node.inputs[2])) @@ -1854,12 +1854,13 @@ def visit_Conv2D(self, node): outputList = self.create_and_add_output_node(node) ndims = len(outputList[0].desc(self.graph).shape) mapLabel = string_builder(node.type) - reduce_node = self.state.add_transient( + self.graph.add_transient( mapLabel + "_wcr_avoid" + local_ctr, [1], outputList[0].desc(self.graph).dtype, storage=dace.StorageType.Register, ) + reduce_node = self.state.add_access(mapLabel + "_wcr_avoid" + local_ctr) reduce_node.setzero = True mapParams = [] @@ -2109,12 +2110,13 @@ def visit_AvgPool(self, node): # + normalisationScalar # + ");", ) - temp_node = self.state.add_scalar( + self.graph.add_scalar( "scratch_node" + local_count, dace.typeclass(_tensortype(node.outputs[0])), transient=True, storage=dace.StorageType.Register, ) + temp_node = self.state.add_access("scratch_node" + local_count) temp_node.setzero = True memletTempNode = Memlet.simple(str(temp_node), "0", wcr_str="lambda a, b: a+b") memletTempNode_nocr = Memlet.simple(str(temp_node), "0") @@ -2249,12 +2251,13 @@ def visit_ShapeN(self, node): try: outpNode = _find_node(self.state, label) except (LookupError): - outpNode = self.state.add_array( + self.graph.add_array( label, _tensorshape(outputTensor), _tensortype(outputTensor), lifetime=dtypes.AllocationLifetime.SDFG, ) + outpNode = self.state.add_access(label) outpNode.desc(self.graph).transient = False def visit_Reshape(self, node): @@ -2526,7 +2529,7 @@ def visit_Conv2DBackpropInput(self, node): paddedOutputDims = outputDims[0].copy() paddedOutputDims[1] += "+" + str(padding) paddedOutputDims[2] += "+" + str(padding) - paddedOutput = state.add_transient( + self.graph.add_transient( string_builder(node.outputs[0].name) + "_padded", [ paddedOutputDims[0][2:], @@ -2536,6 +2539,7 @@ def visit_Conv2DBackpropInput(self, node): ], _tensortype(node.outputs[0]), ) + paddedOutput = state.add_access(string_builder(node.outputs[0].name) + "_padded") if strides > 1: # Dilate and pad the incoming gradients @@ -2548,11 +2552,12 @@ def visit_Conv2DBackpropInput(self, node): if newShape[1] - ksize + 1 < node.outputs[0].shape[1]: newShape[1] = node.outputs[0].shape[1] + ksize - 1 newShape[2] = node.outputs[0].shape[2] + ksize - 1 - expandedGrads = state.add_transient( + self.graph.add_transient( string_builder(node.inputs[2].name) + "_bigger_strided", newShape, _tensortype(node.inputs[2]), ) + expandedGrads = state.add_access(string_builder(node.inputs[2].name) + "_bigger_strided") expandedGrads.setzero = True mapParams = self.get_default_params(node.inputs[2]) mapRange = self.get_default_dims(node.inputs[2]) @@ -2587,11 +2592,12 @@ def visit_Conv2DBackpropInput(self, node): if newShape[1] - ksize + 1 < node.outputs[0].shape[1]: newShape[1] = node.outputs[0].shape[1] + ksize - 1 newShape[2] = node.outputs[0].shape[2] + ksize - 1 - expandedGrads = state.add_transient( + self.graph.add_transient( string_builder(node.inputs[2].name) + "_bigger", newShape, _tensortype(node.inputs[2]), ) + expandedGrads = state.add_access(string_builder(node.inputs[2].name) + "_bigger") expandedGrads.setzero = True expanderMemlet = Memlet.simple( inputNodes[1], @@ -2624,12 +2630,13 @@ def visit_Conv2DBackpropInput(self, node): tasklet = state.add_tasklet(mapLabel, {"j0", "j1"}, {"out"}, "out = j0 * j1") - reduce_node = self.state.add_transient( + self.graph.add_transient( mapLabel + "_wcr_avoid", [1], outputList[0].desc(self.graph).dtype, storage=dace.StorageType.Register, ) + reduce_node = self.state.add_access(mapLabel + "_wcr_avoid") reduce_node.setzero = True if padding > 0: @@ -2733,11 +2740,12 @@ def visit_Conv2DBackpropFilter(self, node): node.inputs[2].shape[2] + (node.inputs[2].shape[2] - 1) * (strides - 1), node.inputs[2].shape[3], ] - expandedGrads = state.add_transient( + self.graph.add_transient( string_builder(node.inputs[2].name) + "_bigger", newShape, _tensortype(node.inputs[2]), ) + expandedGrads = state.add_access(string_builder(node.inputs[2].name) + "_bigger") expandedGrads.setzero = True mapParams = self.get_default_params(node.inputs[2]) mapRange = self.get_default_dims(node.inputs[2]) @@ -2779,12 +2787,13 @@ def visit_Conv2DBackpropFilter(self, node): tasklet = state.add_tasklet(mapLabel, {"j0", "j1"}, {"out"}, "out = j0*j1") - reduce_node = self.state.add_transient( + self.graph.add_transient( mapLabel + "_wcr_avoid", [1], outputList[0].desc(self.graph).dtype, storage=dace.StorageType.Register, ) + reduce_node = self.state.add_access(mapLabel + "_wcr_avoid") reduce_node.setzero = True self.add_out_memlets( @@ -2837,7 +2846,8 @@ def visit_SparseSoftmaxCrossEntropyWithLogits(self, node): except (LookupError): dtype = dace.typeclass(_tensortype(node)) shape = dace.properties.ShapeProperty.from_string(str(_tensorshape(out))) - outputNode = state.add_transient(label, shape, dtype, lifetime=dtypes.AllocationLifetime.SDFG) + self.graph.add_transient(label, shape, dtype, lifetime=dtypes.AllocationLifetime.SDFG) + outputNode = state.add_access(label) outputList.append(outputNode) mapLabel = string_builder(node.type) @@ -2848,7 +2858,8 @@ def visit_SparseSoftmaxCrossEntropyWithLogits(self, node): dtype = dace.typeclass(_tensortype(node)) shape = dace.properties.ShapeProperty.from_string(str(inputList[1].shape)) - temp1Node = state.add_transient(mapLabel + "_max_tmp", shape, dtype, lifetime=dtypes.AllocationLifetime.SDFG) + self.graph.add_transient(mapLabel + "_max_tmp", shape, dtype, lifetime=dtypes.AllocationLifetime.SDFG) + temp1Node = state.add_access(mapLabel + "_max_tmp") mapEntry, mapExit = state.add_map( mapLabel + "_max", dict(zip(mapParams, mapRange)), @@ -2868,10 +2879,8 @@ def visit_SparseSoftmaxCrossEntropyWithLogits(self, node): ) # 2nd map, calculate the denominator sum - temp2Node = state.add_transient(mapLabel + "_denominator_tmp", - shape, - dtype, - lifetime=dtypes.AllocationLifetime.SDFG) + self.graph.add_transient(mapLabel + "_denominator_tmp", shape, dtype, lifetime=dtypes.AllocationLifetime.SDFG) + temp2Node = state.add_access(mapLabel + "_denominator_tmp") mapEntry, mapExit = state.add_map( mapLabel + "_denominator", dict(zip(mapParams, mapRange)), @@ -2895,10 +2904,8 @@ def visit_SparseSoftmaxCrossEntropyWithLogits(self, node): # 3rd map, calculate the sofmax shape = dace.properties.ShapeProperty.from_string(str(inputList[0].shape)) - temp3Node = state.add_transient(mapLabel + "_softmax_tmp", - shape, - dtype, - lifetime=dtypes.AllocationLifetime.SDFG) + self.graph.add_transient(mapLabel + "_softmax_tmp", shape, dtype, lifetime=dtypes.AllocationLifetime.SDFG) + temp3Node = state.add_access(mapLabel + "_softmax_tmp") mapEntry, mapExit = state.add_map(mapLabel + "_softmax", dict(zip(mapParams, mapRange))) tasklet = state.add_tasklet(mapLabel + "_softmax", {"j0", "j1", "j2"}, {"out"}, "out = (dace::math::exp(j0-j1))/j2;", @@ -3007,21 +3014,23 @@ def visit_LRNGrad(self, node): paddedShape = [] paddedShape += shortAccesses paddedShape[-1] += "+" + depth_radius - paddedInput = state.add_transient( + self.graph.add_transient( label + "_paddedInput", paddedShape, dace.typeclass(_tensortype(node)), lifetime=dtypes.AllocationLifetime.SDFG, ) + paddedInput = state.add_access(label + "_paddedInput") mapEntry, mapExit = state.add_map(label + "_padding", dict(zip(shortParams, shortDims))) tasklet = state.add_tasklet(label + "_padding", {"j0"}, {"out"}, "out=j0") self.add_in_memlets([inputNodes[2]], mapEntry, tasklet, [shortDims], [shortParams]) self.add_out_memlets([paddedInput], mapExit, tasklet, [paddedDims], [copyParams]) - sqrsum = state.add_transient(label + "_Sqrsum", - shortAccesses, - _tensortype(node), - lifetime=dtypes.AllocationLifetime.SDFG) + self.graph.add_transient(label + "_Sqrsum", + shortAccesses, + _tensortype(node), + lifetime=dtypes.AllocationLifetime.SDFG) + sqrsum = state.add_access(label + "_Sqrsum") mapEntry, mapExit = state.add_map(label + "_sqrsum", dict(zip(longParams, longDims))) tasklet = state.add_tasklet(label + "_sqrsum", {"j0"}, {"out"}, "out=j0*j0") self.reinitCR(sqrsum, [shortParams], [shortDims], "0") @@ -3029,19 +3038,21 @@ def visit_LRNGrad(self, node): self.add_out_memlets([sqrsum], mapExit, tasklet, [shortDims], [shortParams], "lambda a,b: a+b", 0) label = string_builder(node.name) - norm = state.add_transient(label + "_Norm", - shortAccesses, - _tensortype(node), - lifetime=dtypes.AllocationLifetime.SDFG) + sel.graph.add_transient(label + "_Norm", + shortAccesses, + _tensortype(node), + lifetime=dtypes.AllocationLifetime.SDFG) + norm = state.add_access(label + "_Norm") mapEntry, mapExit = state.add_map(label + "_norm", dict(zip(shortParams, shortDims))) tasklet = state.add_tasklet(label + "_norm", {"j0"}, {"out"}, "out=" + alpha + "*j0+" + bias) self.add_in_memlets([sqrsum], mapEntry, tasklet, [shortDims], [shortParams]) self.add_out_memlets([norm], mapExit, tasklet, [shortDims], [shortParams]) - preOut = state.add_transient(label + "_preOut", - shortAccesses, - _tensortype(node), - lifetime=dtypes.AllocationLifetime.SDFG) + self.graph.add_transient(label + "_preOut", + shortAccesses, + _tensortype(node), + lifetime=dtypes.AllocationLifetime.SDFG) + preOut = state.add_access(label + "_preOut") mapEntry, mapExit = state.add_map(label, dict(zip(longParams, longDims))) taskletCode = ("if (i4==" + depth_radius + "){\n out = pow(j2," + beta + ")-2*" + alpha + "*" + beta + "*j1*j0/j2;}\n else{\n out = -2*" + alpha + "*" + beta + "*j1*j0/j2;}") @@ -3107,21 +3118,23 @@ def visit_LRN(self, node): paddedShape = [] paddedShape += shortAccesses paddedShape[-1] += "+" + depth_radius - paddedInput = state.add_transient( + self.graph.add_transient( label + "_paddedInput", paddedShape, dace.typeclass(_tensortype(node)), lifetime=dtypes.AllocationLifetime.SDFG, ) + paddedInput = state.add_access(label + "_paddedInput") mapEntry, mapExit = state.add_map(label + "_padding", dict(zip(shortParams, shortDims))) tasklet = state.add_tasklet(label + "_padding", {"j0"}, {"out"}, "out=j0") self.add_in_memlets([inputNodes[0]], mapEntry, tasklet, [shortDims], [shortParams]) self.add_out_memlets([paddedInput], mapExit, tasklet, [paddedDims], [copyParams]) - sqrsum = state.add_transient(label + "_Sqrsum", - shortAccesses, - _tensortype(node), - lifetime=dtypes.AllocationLifetime.SDFG) + self.graph.add_transient(label + "_Sqrsum", + shortAccesses, + _tensortype(node), + lifetime=dtypes.AllocationLifetime.SDFG) + sqrsum = state.add_access(label + "_Sqrsum") mapEntry, mapExit = state.add_map(label + "_sqrsum", dict(zip(longParams, longDims))) tasklet = state.add_tasklet(label + "_sqrsum", {"j0"}, {"out"}, "out=j0*j0") self.reinitCR(sqrsum, [shortParams], [shortDims], "0") @@ -3172,7 +3185,8 @@ def visit_ArgMax(self, node): mapEntry, mapExit = state.add_map(mapLabel + "_max", dict(zip(inputParams[0], inputDims[0]))) dtype = dace.typeclass(_tensortype(node)) shape = dace.properties.ShapeProperty.from_string(",".join(inputAccesses[1])) - temp1Node = state.add_transient(mapLabel + "_max_tmp", shape, dtype, lifetime=dtypes.AllocationLifetime.SDFG) + self.graph.add_transient(mapLabel + "_max_tmp", shape, dtype, lifetime=dtypes.AllocationLifetime.SDFG) + temp1Node = state.add_access(mapLabel + "_max_tmp") tasklet = state.add_tasklet(mapLabel + "_max", {"j0"}, {"out"}, "out = j0") self.reinitCR(temp1Node, [inputParams[1]], [inputDims[1]], "-999999999999") @@ -3309,7 +3323,8 @@ def visit_Softmax(self, node): # 1st map, get maximum in each batchsize dimension dtype = dace.typeclass(_tensortype(node)) shape = dace.properties.ShapeProperty.from_string(str(node.inputs[0].shape.dims[0])) - temp1Node = state.add_transient(mapLabel + "_max_tmp", shape, dtype, lifetime=dtypes.AllocationLifetime.SDFG) + self.graph.add_transient(mapLabel + "_max_tmp", shape, dtype, lifetime=dtypes.AllocationLifetime.SDFG) + temp1Node = state.add_access(mapLabel + "_max_tmp") mapEntry, mapExit = state.add_map(mapLabel + "_max", dict(zip(mapParams, mapRange))) tasklet = state.add_tasklet(mapLabel + "_max", {"j0"}, {"out"}, "out = j0") self.reinitCR(temp1Node, [inputParams[1]], [inputDims[1]], "-999999999999") @@ -3325,10 +3340,8 @@ def visit_Softmax(self, node): ) # 2nd map, calculate the denominator sum - temp2Node = state.add_transient(mapLabel + "_denominator_tmp", - shape, - dtype, - lifetime=dtypes.AllocationLifetime.SDFG) + self.graph.add_transient(mapLabel + "_denominator_tmp", shape, dtype, lifetime=dtypes.AllocationLifetime.SDFG) + temp2Node = state.add_access(mapLabel + "_denominator_tmp") mapEntry, mapExit = state.add_map(mapLabel + "_denominator", dict(zip(mapParams, mapRange))) tasklet = state.add_tasklet(mapLabel + "_denominator", {"j0", "j1"}, {"out"}, "out = dace::math::exp(j0-j1);", @@ -3521,10 +3534,8 @@ def create_and_add_input_node(self, inp): shape = dace.properties.ShapeProperty.from_string(str(_tensorshape(inp))) # Create and add array, default is transient allocated in the # beginning of the SDFG - inputNode = state.add_transient(name=label, - shape=shape, - dtype=dtype, - lifetime=dtypes.AllocationLifetime.SDFG) + self.graph.add_transient(name=label, shape=shape, dtype=dtype, lifetime=dtypes.AllocationLifetime.SDFG) + inputNode = state.add_access(name=label) params = self.get_default_params(inp) dims = self.get_default_dims(inp) @@ -3554,7 +3565,8 @@ def create_and_add_output_node(self, node): # Get type and shape of the tensor dtype = dace.typeclass(_tensortype(out)) shape = dace.properties.ShapeProperty.from_string(str(_tensorshape(out))) - outputNode = state.add_transient(label, shape, dtype, lifetime=dtypes.AllocationLifetime.SDFG) + self.graph.add_transient(label, shape, dtype, lifetime=dtypes.AllocationLifetime.SDFG) + outputNode = state.add_access(label) outputList.append(outputNode) return outputList @@ -3642,10 +3654,11 @@ def inputPadding(self, node, inpnode, inp, outputSize, kernelSize, strides, inpu # Add the padded input to the graph, set it to zero, and add the map. shape = dace.properties.ShapeProperty.from_string(",".join(outputAccesses)) - output = state.add_transient(label + "_padded", - shape=shape, - dtype=inp.dtype, - lifetime=dtypes.AllocationLifetime.SDFG) + self.graph.add_transient(label + "_padded", + shape=shape, + dtype=inp.dtype, + lifetime=dtypes.AllocationLifetime.SDFG) + output = state.add_access(label + "_padded") output.setzero = True # mapParams = inputParams diff --git a/dace/frontend/ml/tensorflow/transformations/redundant_array.py b/dace/frontend/ml/tensorflow/transformations/redundant_array.py index f42df1a847..fded541265 100644 --- a/dace/frontend/ml/tensorflow/transformations/redundant_array.py +++ b/dace/frontend/ml/tensorflow/transformations/redundant_array.py @@ -5,7 +5,6 @@ from dace.sdfg import nodes from dace.sdfg import utils as sdutil from dace.transformation import transformation as pm -from dace.config import Config class TensorflowRedundantArray(pm.SingleStateTransformation): diff --git a/dace/frontend/ml/torch/module.py b/dace/frontend/ml/torch/module.py index ba820faf87..796cae9bc3 100644 --- a/dace/frontend/ml/torch/module.py +++ b/dace/frontend/ml/torch/module.py @@ -1,12 +1,11 @@ # Copyright 2019-2025 ETH Zurich and the DaCe authors. All rights reserved. """ DaCe Python parsing functionality and entry point to Python frontend. """ -from dataclasses import dataclass import collections import itertools import tempfile import copy import os -from typing import Any, Callable, Dict, OrderedDict, List, Optional, Set, Sequence, Tuple, Union +from typing import Any, Callable, Dict, OrderedDict, List, Optional, Set, Tuple, Union # Try importing ML dependencies try: @@ -29,7 +28,6 @@ onnx = None ONNX_AVAILABLE = False -import dace from dace import config, data from dace.codegen import compiled_sdfg from dace.sdfg import SDFG, nodes diff --git a/dace/frontend/octave/parse.py b/dace/frontend/octave/parse.py index 221794a48f..41aef98d17 100644 --- a/dace/frontend/octave/parse.py +++ b/dace/frontend/octave/parse.py @@ -5,12 +5,12 @@ import copy import dace -from .ast_node import AST_Node, AST_Statements +from .ast_node import AST_Statements from .ast_values import AST_Ident, AST_Constant from .ast_expression import AST_BinExpression, AST_UnaryExpression from .ast_matrix import AST_Matrix_Row, AST_Matrix, AST_Transpose from .ast_assign import AST_Assign -from .ast_function import AST_Argument, AST_BuiltInFunCall, AST_FunCall, AST_Function, AST_EndFunc +from .ast_function import AST_FunCall, AST_Function, AST_EndFunc from .ast_range import AST_RangeExpression from .ast_loop import AST_ForLoop from .ast_nullstmt import AST_NullStmt, AST_Comment, AST_EndStmt diff --git a/dace/frontend/operations.py b/dace/frontend/operations.py index 57062f4e3f..41e5d420ad 100644 --- a/dace/frontend/operations.py +++ b/dace/frontend/operations.py @@ -1,5 +1,4 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. -from functools import partial from itertools import chain, repeat from contextlib import contextmanager diff --git a/dace/frontend/python/cached_program.py b/dace/frontend/python/cached_program.py index cdada4b417..d3004528b4 100644 --- a/dace/frontend/python/cached_program.py +++ b/dace/frontend/python/cached_program.py @@ -3,7 +3,7 @@ from collections import OrderedDict from dataclasses import dataclass -from typing import Any, Callable, Dict, Optional, Set, Tuple +from typing import Any, Callable, Dict, Optional, Set import dace from dace import config diff --git a/dace/frontend/python/interface.py b/dace/frontend/python/interface.py index 5581f727bd..169df4680b 100644 --- a/dace/frontend/python/interface.py +++ b/dace/frontend/python/interface.py @@ -2,9 +2,8 @@ """ Python interface for DaCe functions. """ import inspect -from functools import wraps -from typing import (Any, Callable, Deque, Dict, Generator, Optional, Tuple, TypeVar, Union, overload, TYPE_CHECKING, - Generic, Iterable, Iterator) +from typing import (Any, Callable, Deque, Dict, Generator, Optional, TypeVar, Union, overload, Generic, Iterable, + Iterator) from typing_extensions import Self from dace import dtypes @@ -57,7 +56,7 @@ def program(f: F, on the generated DaCe program during compilation. :param device: Transform the function to run on the target device. :param recreate_sdfg: Whether to recreate the SDFG from the Python code. If False, the SDFG will be loaded from the - cache (``//program.sdfg``) if it exists. + cache (``//program.sdfgz``) if it exists. Use this if you want to modify the SDFG after the first call to the function. :param regenerate_code: Whether to regenerate the code from the SDFG. If False, the code in the build folder will be used if it exists. Use this if you want to modify the generated code without DaCe overriding @@ -126,7 +125,7 @@ def method(f: F, on the generated DaCe program during compilation. :param device: Transform the function to run on the target device. :param recreate_sdfg: Whether to recreate the SDFG from the Python code. If False, the SDFG will be loaded from the - cache (``//program.sdfg``) if it exists. + cache (``//program.sdfgz``) if it exists. Use this if you want to modify the SDFG after the first call to the function. :param regenerate_code: Whether to regenerate the code from the SDFG. If False, the code in the build folder will be used if it exists. Use this if you want to modify the generated code without DaCe overriding diff --git a/dace/frontend/python/newast.py b/dace/frontend/python/newast.py index 6f382388eb..30fb537247 100644 --- a/dace/frontend/python/newast.py +++ b/dace/frontend/python/newast.py @@ -2029,23 +2029,27 @@ def _parse_map_inputs(self, name: str, params: List[Tuple[str, str]], self.sdfg.add_symbol(atomstr, self.defined[candidate].dtype) for expr in symbolic.swalk(symval): + # An array access in a bound (legacy ``arr(i)`` or ``Subscript(arr, i)``) + # is hoisted to a dynamic-map-input symbol reading the array. if symbolic.is_sympy_userfunction(expr): - # If function contains a function - if any(symbolic.contains_sympy_functions(a) for a in expr.args): - raise DaceSyntaxError(self, node, 'Indirect accesses not supported in map ranges') - arr = expr.func.__name__ - newvar = '__%s_%s%d' % (name, vid, ctr) - repldict[arr] = newvar - # Create memlet - args = ','.join([str(a) for a in expr.args]) - if arr in self.variables: - arr = self.variables[arr] - if not isinstance(arr, str) or arr not in self.sdfg.arrays: - rng = subsets.Range.from_string(args) - args = str(rng) - map_inputs[newvar] = Memlet.simple(arr, args) - # ','.join([str(a) for a in expr.args])) - ctr += 1 + arr, indices = expr.func.__name__, expr.args + elif isinstance(expr, symbolic.Subscript): + arr, indices = str(expr.args[0]), expr.args[1:] + else: + continue + if any(symbolic.contains_sympy_functions(a) for a in indices): + raise DaceSyntaxError(self, node, 'Indirect accesses not supported in map ranges') + newvar = '__%s_%s%d' % (name, vid, ctr) + repldict[arr] = newvar + # Create memlet + args = ','.join([str(a) for a in indices]) + if arr in self.variables: + arr = self.variables[arr] + if not isinstance(arr, str) or arr not in self.sdfg.arrays: + rng = subsets.Range.from_string(args) + args = str(rng) + map_inputs[newvar] = Memlet.simple(arr, args) + ctr += 1 # Replace functions with new variables for find, replace in repldict.items(): val = re.sub(r"%s\(.*?\)" % find, val, replace) @@ -3597,11 +3601,8 @@ def _visit_assign(self, node, node_target, op, dtype=None, is_return=False): new_data, rng = None, None dtype_keys = tuple(dtypes.dtype_to_typeclass().keys()) - if not ( - result in self.sdfg.symbols or symbolic.issymbolic(result) or isinstance(result, dtype_keys) or - (isinstance(result, str) and any( - result in x - for x in [self.sdfg.arrays, self.sdfg._pgrids, self.sdfg._subarrays, self.sdfg._rdistrarrays]))): + if not (result in self.sdfg.symbols or symbolic.issymbolic(result) or isinstance(result, dtype_keys) or + (isinstance(result, str) and result in self.sdfg.arrays)): raise DaceSyntaxError( self, node, "In assignments, the rhs may only be " "data, numerical/boolean constants " @@ -3625,11 +3626,8 @@ def _visit_assign(self, node, node_target, op, dtype=None, is_return=False): true_name, new_data = self.sdfg.add_scalar(true_name, ttype, transient=True, find_new_name=True) self.variables[name] = true_name defined_vars[name] = true_name - if any(result in x for x in [self.sdfg._pgrids, self.sdfg._rdistrarrays, self.sdfg._subarrays]): - # NOTE: In previous versions some `pgrid` and subgrid related replacement function, - # see `dace/frontend/common/distr.py`, created dummy variables with the same name - # as the entities, such as process grids, they created. Thus the frontend was - # finding them. Since this is now disallowed, we have to explicitly handle this case. + if (isinstance(result, str) and result in self.sdfg.arrays + and isinstance(self.sdfg.arrays[result], data.DistributedDescriptor)): self.variables[name] = result defined_vars[name] = result continue @@ -5510,7 +5508,10 @@ def visit_Subscript(self, node: ast.Subscript, inference: bool = False): except (TypeError, ValueError): pass # Passthrough to exception - raise DaceSyntaxError(self, node.value, 'Subscripted object cannot be a tuple') + raise DaceSyntaxError( + self, node.value, f'Subscript {[t[1] for t in node_parsed]} unsupported. ' + 'DaCe supports sequence of dace.data or an attribute of a dace.data. ' + 'Try `dace.unroll` for more complex types.') array, arrtype = node_parsed[0] if arrtype == 'str' or arrtype in dtypes._CTYPES: raise DaceSyntaxError(self, node, 'Type "%s" cannot be sliced' % arrtype) @@ -5604,8 +5605,12 @@ def make_slice(self, arrname: str, rng: subsets.Range): arrobj.storage, find_new_name=True) wnode = self.current_state.add_write(tmp, debuginfo=self.current_lineinfo) + # ``Memlet.simple`` keeps the subset by reference, but ``rng`` may be a + # cached Range shared by sibling slice reads, so give this edge its own copy. self.current_state.add_nedge( - rnode, wnode, Memlet.simple(array, rng, num_accesses=rng.num_elements(), other_subset_str=other_subset)) + rnode, wnode, + Memlet.simple(array, copy.deepcopy(rng), num_accesses=rng.num_elements(), + other_subset_str=other_subset)) return tmp, other_subset def _compute_output_shape_from_advanced_indexing(self, aname: str, expr: MemletExpr) -> List[symbolic.SymbolicType]: diff --git a/dace/frontend/python/parser.py b/dace/frontend/python/parser.py index 9d9c993fe5..3b9325c2f6 100644 --- a/dace/frontend/python/parser.py +++ b/dace/frontend/python/parser.py @@ -1,9 +1,7 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. """ DaCe Python parsing functionality and entry point to Python frontend. """ import ast -from dataclasses import dataclass import inspect -import itertools import copy import os import sympy @@ -171,7 +169,6 @@ def __init__(self, method: bool = False, use_explicit_cf: bool = True, ignore_type_hints: bool = False): - from dace.codegen import compiled_sdfg # Avoid import loops self.f = f self.dec_args = args @@ -500,7 +497,7 @@ def _parse(self, args, kwargs, simplify=None, save=False, validate=False) -> SDF :param simplify: Whether to apply simplification pass or not (None uses configuration-defined value). :param save: If True, saves the generated SDFG to - ``_dacegraphs/program.sdfg`` after parsing. + ``_dacegraphs/program.sdfgz`` after parsing. :param validate: If True, validates the resulting SDFG after creation. :return: The generated SDFG object. """ @@ -523,7 +520,7 @@ def _parse(self, args, kwargs, simplify=None, save=False, validate=False) -> SDF # Save the SDFG. Skip this step if running from a cached SDFG, as # it might overwrite the cached SDFG. if not cached and not Config.get_bool('compiler', 'use_cache') and save: - sdfg.save(os.path.join('_dacegraphs', 'program.sdfg')) + sdfg.save(os.path.join('_dacegraphs', 'program.sdfgz'), compress=True) # Validate SDFG if validate: @@ -802,13 +799,13 @@ def load_precompiled_sdfg(self, path: str, *args, **kwargs) -> None: function is called. :param path: Path to SDFG build folder (e.g., ".dacecache/program"). - Path has to include ``program.sdfg`` and the binary shared + Path has to include ``program.sdfgz`` and the binary shared object under the ``build`` folder. :param args: Optional compile-time arguments. :param kwargs: Optional compile-time keyword arguments. """ - from dace.sdfg import utils as sdutil # Avoid import loop - csdfg = sdutil.load_precompiled_sdfg(path, self.argnames) + from dace.codegen import compiler as sdfg_compiler # Avoid import loop + csdfg = sdfg_compiler.load_precompiled_sdfg(path) _, cachekey = self._load_sdfg(None, *args, **kwargs) # Update SDFG cache with the SDFG and compiled version @@ -909,7 +906,12 @@ def _generate_pdp(self, # If recreate flag is False, check and load from cache if not self.recreate_sdfg: build_folder = SDFG(self.name).build_folder - sdfg, _ = self.load_sdfg(os.path.join(build_folder, 'program.sdfg'), *args, **kwargs) + sdfg, _ = self.load_sdfg(os.path.join(build_folder, 'program.sdfgz'), *args, **kwargs) + + if sdfg is None: + # attempt to load uncompressed sdfg (backwards compatibility) + sdfg, _ = self.load_sdfg(os.path.join(build_folder, 'program.sdfg'), *args, **kwargs) + if sdfg is not None: return sdfg, True diff --git a/dace/frontend/python/preprocessing.py b/dace/frontend/python/preprocessing.py index dbc9536790..5f6a387b6b 100644 --- a/dace/frontend/python/preprocessing.py +++ b/dace/frontend/python/preprocessing.py @@ -12,7 +12,7 @@ from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union import dace -from dace import data, dtypes, symbolic, sdfg +from dace import data, dtypes, symbolic from dace.config import Config from dace.sdfg import SDFG from dace.frontend.python import astutils diff --git a/dace/frontend/python/replacements/array_creation.py b/dace/frontend/python/replacements/array_creation.py index 71e42b965f..66bcd57163 100644 --- a/dace/frontend/python/replacements/array_creation.py +++ b/dace/frontend/python/replacements/array_creation.py @@ -2,6 +2,7 @@ """ Contains replacements for array-filling methods (zeros, ones, etc.) """ +import dace # noqa from dace.frontend.common import op_repository as oprepo from dace.frontend.python.common import DaceSyntaxError from dace.frontend.python.replacements.utils import ProgramVisitor, Shape, sym_type, broadcast_together diff --git a/dace/frontend/python/replacements/array_manipulation.py b/dace/frontend/python/replacements/array_manipulation.py index 379d9c9fb8..f414d239f3 100644 --- a/dace/frontend/python/replacements/array_manipulation.py +++ b/dace/frontend/python/replacements/array_manipulation.py @@ -1,9 +1,10 @@ -# Copyright 2019-2025 ETH Zurich and the DaCe authors. All rights reserved. +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. """ Contains replacements for N-dimensional array transformations. """ +import dace # noqa from dace.frontend.common import op_repository as oprepo -from dace.frontend.python.common import DaceSyntaxError, StringLiteral +from dace.frontend.python.common import StringLiteral from dace.frontend.python.replacements.utils import ProgramVisitor, UfuncInput, UfuncOutput import dace.frontend.python.memlet_parser as mem_parser from dace import data, dtypes, subsets, symbolic @@ -172,8 +173,8 @@ def _transpose(pv: ProgramVisitor, if axes == (1, 0): # Special case for 2D transposition acc1 = state.add_read(inpname) acc2 = state.add_write(outname) - import dace.libraries.standard # Avoid import loop - tasklet = dace.libraries.standard.Transpose('_Transpose_', restype) + import dace.libraries.linalg # Avoid import loop + tasklet = dace.libraries.linalg.Transpose('_Transpose_', restype) state.add_node(tasklet) state.add_edge(acc1, None, tasklet, '_inp', Memlet.from_array(inpname, arr1)) state.add_edge(tasklet, '_out', acc2, None, Memlet.from_array(outname, arr2)) @@ -190,7 +191,7 @@ def _transpose(pv: ProgramVisitor, read = state.add_read(inpname) write = state.add_write(outname) - from dace.libraries.standard import TensorTranspose + from dace.libraries.linalg import TensorTranspose # Avoid import loop tasklet = TensorTranspose('_TensorTranspose', axes or list(range(len(arr1.shape)))) state.add_node(tasklet) state.add_edge(read, None, tasklet, '_inp_tensor', Memlet.from_array(inpname, arr1)) diff --git a/dace/frontend/python/replacements/array_metadata.py b/dace/frontend/python/replacements/array_metadata.py index fd3d3a2dbe..387c1c9195 100644 --- a/dace/frontend/python/replacements/array_metadata.py +++ b/dace/frontend/python/replacements/array_metadata.py @@ -2,6 +2,7 @@ """ Contains replacements for array metadata (shape, strides, etc.). """ +import dace # noqa from dace.frontend.common import op_repository as oprepo from dace.frontend.python.replacements.utils import ProgramVisitor, Size from dace import data, SDFG, SDFGState diff --git a/dace/frontend/python/replacements/linalg.py b/dace/frontend/python/replacements/linalg.py index 3c10d95d23..f5b054be25 100644 --- a/dace/frontend/python/replacements/linalg.py +++ b/dace/frontend/python/replacements/linalg.py @@ -2,6 +2,7 @@ """ Contains linear algebra function and operator replacements. """ +import dace # noqa from dace.frontend.common import op_repository as oprepo from dace.frontend.python.common import StringLiteral from dace.frontend.python.replacements.utils import ProgramVisitor diff --git a/dace/frontend/python/replacements/misc.py b/dace/frontend/python/replacements/misc.py index 9715edb968..7bfc15e8a4 100644 --- a/dace/frontend/python/replacements/misc.py +++ b/dace/frontend/python/replacements/misc.py @@ -3,6 +3,7 @@ Contains replacements for miscellaneous functions and convenience utility functions, such as a function that calls element-wise operations on data containers. """ +import dace # noqa from dace.frontend.common import op_repository as oprepo from dace.frontend.python import astutils from dace.frontend.python.common import StringLiteral @@ -20,14 +21,14 @@ def _slice(pv: ProgramVisitor, sdfg: SDFG, state: SDFGState, *args, **kwargs): @oprepo.replaces_operator('Array', 'MatMult', otherclass='StorageType') -def _cast_storage(visitor: 'ProgramVisitor', sdfg: SDFG, state: SDFGState, arr: str, stype: dtypes.StorageType) -> str: +def _cast_storage(visitor: ProgramVisitor, sdfg: SDFG, state: SDFGState, arr: str, stype: dtypes.StorageType) -> str: desc = sdfg.arrays[arr] desc.storage = stype return arr @oprepo.replaces('dace.elementwise') -def elementwise(pv: 'ProgramVisitor', +def elementwise(pv: ProgramVisitor, sdfg: SDFG, state: SDFGState, func: Union[StringLiteral, str], diff --git a/dace/frontend/python/replacements/mpi.py b/dace/frontend/python/replacements/mpi.py index 5539833447..a3653ae2ea 100644 --- a/dace/frontend/python/replacements/mpi.py +++ b/dace/frontend/python/replacements/mpi.py @@ -9,7 +9,7 @@ from dace.memlet import Memlet from dace.sdfg import SDFG, SDFGState from numbers import Integral, Number -from typing import Sequence, Tuple, Union +from typing import Optional, Sequence, Tuple, Union ShapeType = Sequence[Union[Integral, str, symbolic.symbol, symbolic.SymExpr, symbolic.sympy.Basic]] RankType = Union[Integral, str, symbolic.symbol, symbolic.SymExpr, symbolic.sympy.Basic] @@ -25,12 +25,13 @@ def _cart_create(pv: ProgramVisitor, sdfg: SDFG, state: SDFGState, dims: ShapeTy :param dims: Shape of the process-grid (see `dims` parameter of `MPI_Cart_create`), e.g., [2, 3, 3]. :return: Name of the new process-grid descriptor. """ - pgrid_name = sdfg.add_pgrid(dims) + name = pv.get_target_name() + pgrid_name = sdfg.add_pgrid(dims, name=name) # Dummy tasklet adds MPI variables to the program's state. from dace.libraries.mpi import Dummy tasklet = Dummy(pgrid_name, [ - f'MPI_Comm {pgrid_name}_comm;', + f'MPI_Comm {pgrid_name};', f'MPI_Group {pgrid_name}_group;', f'int {pgrid_name}_coords[{len(dims)}];', f'int {pgrid_name}_dims[{len(dims)}];', @@ -41,10 +42,8 @@ def _cart_create(pv: ProgramVisitor, sdfg: SDFG, state: SDFGState, dims: ShapeTy state.add_node(tasklet) - # Pseudo-writing to a dummy variable to avoid removal of Dummy node by transformations. - scal_name, scal = sdfg.add_scalar(pgrid_name, dace.int32, transient=True, find_new_name=True) - wnode = state.add_write(scal_name) - state.add_edge(tasklet, '__out', wnode, None, Memlet.from_array(scal_name, scal)) + wnode = state.add_write(pgrid_name) + state.add_edge(tasklet, None, wnode, None, Memlet()) return pgrid_name @@ -77,7 +76,8 @@ def _cart_sub(pv: ProgramVisitor, :param exact_grid: [DEVELOPER] If set then, out of all the sub-grids created, only the one that contains the rank with id `exact_grid` will be utilized for collective communication. :return: Name of the new sub-grid descriptor. """ - pgrid_name = sdfg.add_pgrid(parent_grid=parent_grid, color=color, exact_grid=exact_grid) + name = pv.get_target_name() + pgrid_name = sdfg.add_pgrid(parent_grid=parent_grid, color=color, exact_grid=exact_grid, name=name) # Count sub-grid dimensions. pgrid_ndims = sum([bool(c) for c in color]) @@ -85,7 +85,7 @@ def _cart_sub(pv: ProgramVisitor, # Dummy tasklet adds MPI variables to the program's state. from dace.libraries.mpi import Dummy tasklet = Dummy(pgrid_name, [ - f'MPI_Comm {pgrid_name}_comm;', + f'MPI_Comm {pgrid_name};', f'MPI_Group {pgrid_name}_group;', f'int {pgrid_name}_coords[{pgrid_ndims}];', f'int {pgrid_name}_dims[{pgrid_ndims}];', @@ -96,10 +96,12 @@ def _cart_sub(pv: ProgramVisitor, state.add_node(tasklet) - # Pseudo-writing to a dummy variable to avoid removal of Dummy node by transformations. - scal_name, scal = sdfg.add_scalar(pgrid_name, dace.int32, transient=True, find_new_name=True) - wnode = state.add_write(scal_name) - state.add_edge(tasklet, '__out', wnode, None, Memlet.from_array(scal_name, scal)) + wnode = state.add_write(pgrid_name) + state.add_edge(tasklet, None, wnode, None, Memlet()) + + tasklet.add_in_connector('_parent_grid') + rnode = state.add_read(parent_grid) + state.add_edge(rnode, None, tasklet, '_parent_grid', Memlet(data=parent_grid)) return pgrid_name @@ -163,7 +165,7 @@ def _bcast(pv: ProgramVisitor, from dace.libraries.mpi.nodes.bcast import Bcast from dace.frontend.python.replacements.array_creation_dace import _define_local_scalar - libnode = Bcast('_Bcast_', grid, fcomm) + libnode = Bcast('_Bcast_', fcomm) desc = sdfg.arrays[buffer] in_buffer = state.add_read(buffer) out_buffer = state.add_write(buffer) @@ -175,6 +177,9 @@ def _bcast(pv: ProgramVisitor, root_node = state.add_access(root_name) root_tasklet = state.add_tasklet('_set_root_', {}, {'__out'}, '__out = {}'.format(root)) state.add_edge(root_tasklet, '__out', root_node, None, Memlet.simple(root_name, '0')) + if grid: + libnode.add_in_connector('_grid') + state.add_edge(state.add_read(grid), None, libnode, '_grid', Memlet(data=grid)) state.add_edge(in_buffer, None, libnode, '_inbuffer', Memlet.from_array(buffer, desc)) state.add_edge(root_node, None, libnode, '_root', Memlet.simple(root_node.data, '0')) state.add_edge(libnode, '_outbuffer', out_buffer, None, Memlet.from_array(buffer, desc)) @@ -236,7 +241,7 @@ def _Reduce(pv: ProgramVisitor, from dace.libraries.mpi.nodes.reduce import Reduce from dace.frontend.python.replacements.array_creation_dace import _define_local_scalar - libnode = Reduce('_Reduce_', op, grid) + libnode = Reduce('_Reduce_', op) desc = sdfg.arrays[buffer] in_buffer = state.add_read(buffer) out_buffer = state.add_write(buffer) @@ -248,6 +253,9 @@ def _Reduce(pv: ProgramVisitor, root_node = state.add_access(root_name) root_tasklet = state.add_tasklet('_set_root_', {}, {'__out'}, '__out = {}'.format(root)) state.add_edge(root_tasklet, '__out', root_node, None, Memlet.simple(root_name, '0')) + if grid: + libnode.add_in_connector('_grid') + state.add_edge(state.add_read(grid), None, libnode, '_grid', Memlet(data=grid)) state.add_edge(in_buffer, None, libnode, '_inbuffer', Memlet.from_array(buffer, desc)) state.add_edge(root_node, None, libnode, '_root', Memlet.simple(root_node.data, '0')) state.add_edge(libnode, '_outbuffer', out_buffer, None, Memlet.from_array(buffer, desc)) @@ -261,11 +269,14 @@ def _alltoall(pv: ProgramVisitor, sdfg: SDFG, state: SDFGState, inbuffer: str, o from dace.libraries.mpi.nodes.alltoall import Alltoall - libnode = Alltoall('_Alltoall_', grid) + libnode = Alltoall('_Alltoall_') in_desc = sdfg.arrays[inbuffer] in_buffer = state.add_read(inbuffer) out_desc = sdfg.arrays[outbuffer] out_buffer = state.add_write(outbuffer) + if grid: + libnode.add_in_connector('_grid') + state.add_edge(state.add_read(grid), None, libnode, '_grid', Memlet(data=grid)) state.add_edge(in_buffer, None, libnode, '_inbuffer', Memlet.from_array(in_buffer, in_desc)) state.add_edge(libnode, '_outbuffer', out_buffer, None, Memlet.from_array(out_buffer, out_desc)) @@ -287,7 +298,6 @@ def _intracomm_alltoall(pv: ProgramVisitor, sdfg: SDFG, state: SDFGState, icomm: def _pgrid_alltoall(pv: ProgramVisitor, sdfg: SDFG, state: SDFGState, pgrid: str, inp_buffer: str, out_buffer: str): """ Equivalent to `dace.comm.Alltoall(inp_buffer, out_buffer, grid=pgrid)`. """ - from mpi4py import MPI return _alltoall(pv, sdfg, state, inp_buffer, out_buffer, grid=pgrid) @@ -308,10 +318,13 @@ def _allreduce(pv: ProgramVisitor, op = _mpi4py_to_MPI(MPI, op) if inp_buffer != MPI.IN_PLACE: raise ValueError('DaCe currently supports in-place Allreduce only.') - libnode = Allreduce('_Allreduce_', op, grid) + libnode = Allreduce('_Allreduce_', op) desc = sdfg.arrays[buffer] in_buffer = state.add_read(buffer) out_buffer = state.add_write(buffer) + if grid: + libnode.add_in_connector('_grid') + state.add_edge(state.add_read(grid), None, libnode, '_grid', Memlet(data=grid)) state.add_edge(in_buffer, None, libnode, '_inbuffer', Memlet.from_array(buffer, desc)) state.add_edge(libnode, '_outbuffer', out_buffer, None, Memlet.from_array(buffer, desc)) @@ -519,7 +532,7 @@ def _isend(pv: ProgramVisitor, transient=True, find_new_name=True) - libnode = Isend('_Isend_', grid=grid) + libnode = Isend('_Isend_') buf_range = None if isinstance(buffer, tuple): @@ -589,6 +602,10 @@ def _isend(pv: ProgramVisitor, else: tag_mem = Memlet.simple(tag_name, '0') + if grid: + libnode.add_in_connector('_grid') + state.add_edge(state.add_read(grid), None, libnode, '_grid', Memlet(data=grid)) + state.add_edge(buf_node, None, libnode, '_buffer', buf_mem) state.add_edge(dst_node, None, libnode, '_dest', dst_mem) state.add_edge(tag_node, None, libnode, '_tag', tag_mem) @@ -618,7 +635,6 @@ def _pgrid_isend(pv: ProgramVisitor, sdfg: SDFG, state: SDFGState, pgrid: str, b dst: Union[str, sp.Expr, Number], tag: Union[str, sp.Expr, Number]): """ Equivalent to `dace.comm.Isend(buffer, dst, tag, req, grid=pgrid)`. """ - from mpi4py import MPI req, _ = sdfg.add_array("isend_req", [1], dace.dtypes.opaque("MPI_Request"), transient=True, find_new_name=True) _isend(pv, sdfg, state, buffer, dst, tag, req, grid=pgrid) return req @@ -740,7 +756,7 @@ def _irecv(pv: ProgramVisitor, transient=True, find_new_name=True) - libnode = Irecv('_Irecv_', grid=grid) + libnode = Irecv('_Irecv_') buf_range = None if isinstance(buffer, tuple): @@ -808,6 +824,10 @@ def _irecv(pv: ProgramVisitor, else: tag_mem = Memlet.simple(tag_name, '0') + if grid: + libnode.add_in_connector('_grid') + state.add_edge(state.add_read(grid), None, libnode, '_grid', Memlet(data=grid)) + state.add_edge(libnode, '_buffer', buf_node, None, buf_mem) state.add_edge(src_node, None, libnode, '_src', src_mem) state.add_edge(tag_node, None, libnode, '_tag', tag_mem) @@ -837,7 +857,6 @@ def _pgrid_irecv(pv: ProgramVisitor, sdfg: SDFG, state: SDFGState, pgrid: str, b src: Union[str, sp.Expr, Number], tag: Union[str, sp.Expr, Number]): """ Equivalent to `dace.comm.Isend(buffer, dst, tag, req, grid=pgrid)`. """ - from mpi4py import MPI req, _ = sdfg.add_array("irecv_req", [1], dace.dtypes.opaque("MPI_Request"), transient=True, find_new_name=True) _irecv(pv, sdfg, state, buffer, src, tag, req, grid=pgrid) return req @@ -903,15 +922,14 @@ def _wait(pv: ProgramVisitor, sdfg: SDFG, state: SDFGState, request: str): return None -@oprepo.replaces('dace.comm.Subarray') -def _subarray(pv: ProgramVisitor, - sdfg: SDFG, - state: SDFGState, - array: Union[str, ShapeType], - subarray: Union[str, ShapeType], - dtype: dtypes.typeclass = None, - process_grid: str = None, - correspondence: Sequence[Integral] = None): +def _define_subarray(pv: ProgramVisitor, + sdfg: SDFG, + state: SDFGState, + array: Union[str, ShapeType], + subarray: Union[str, ShapeType], + dtype: dtypes.typeclass = None, + process_grid: str = None, + correspondence: Sequence[Integral] = None) -> Tuple[str, Optional[dace.sdfg.nodes.AccessNode]]: """ Adds a sub-array descriptor to the DaCe Program. Sub-arrays are implemented (when `process_grid` is set) with [MPI_Type_create_subarray](https://www.mpich.org/static/docs/v3.2/www3/MPI_Type_create_subarray.html). @@ -937,7 +955,8 @@ def _subarray(pv: ProgramVisitor, sub_dtype = None dtype = dtype or arr_dtype or sub_dtype - subarray_name = sdfg.add_subarray(dtype, shape, subshape, process_grid, correspondence) + name = pv.get_target_name() + subarray_name = sdfg.add_subarray(dtype, shape, subshape, process_grid, correspondence, name=name) # Generate subgraph only if process-grid is set, i.e., the sub-array will be used for collective scatter/gather ops. if process_grid: @@ -949,11 +968,23 @@ def _subarray(pv: ProgramVisitor, state.add_node(tasklet) - # Pseudo-writing to a dummy variable to avoid removal of Dummy node by transformations. - scal_name, scal = sdfg.add_scalar(subarray_name, dace.int32, transient=True, find_new_name=True) - wnode = state.add_write(scal_name) - state.add_edge(tasklet, '__out', wnode, None, Memlet.from_array(scal_name, scal)) + wnode = state.add_write(subarray_name) + state.add_edge(tasklet, None, wnode, None, Memlet()) + return subarray_name, wnode + + return subarray_name, None + +@oprepo.replaces('dace.comm.Subarray') +def _subarray(pv: ProgramVisitor, + sdfg: SDFG, + state: SDFGState, + array: Union[str, ShapeType], + subarray: Union[str, ShapeType], + dtype: dtypes.typeclass = None, + process_grid: str = None, + correspondence: Sequence[Integral] = None): + subarray_name, _ = _define_subarray(pv, sdfg, state, array, subarray, dtype, process_grid, correspondence) return subarray_name @@ -982,16 +1013,16 @@ def _block_scatter(pv: ProgramVisitor, if in_desc.dtype != out_desc.dtype: raise ValueError("Input/output buffer datatypes must match!") - subarray_name = _subarray(pv, - sdfg, - state, - in_buffer, - out_buffer, - process_grid=scatter_grid, - correspondence=correspondence) + subarray_name, subarray_node = _define_subarray(pv, + sdfg, + state, + in_buffer, + out_buffer, + process_grid=scatter_grid, + correspondence=correspondence) from dace.libraries.mpi import BlockScatter - libnode = BlockScatter('_BlockScatter_', subarray_name, scatter_grid, bcast_grid) + libnode = BlockScatter('_BlockScatter_') inbuf_name = in_buffer in_desc = sdfg.arrays[inbuf_name] @@ -1003,6 +1034,20 @@ def _block_scatter(pv: ProgramVisitor, outbuf_node = state.add_write(outbuf_name) outbuf_mem = Memlet.from_array(outbuf_name, out_desc) + if subarray_node is not None: + libnode.add_in_connector('_subarray', sdfg.arrays[subarray_name].state_field_dtype) + state.add_edge(subarray_node, None, libnode, '_subarray', Memlet(data=subarray_name)) + elif subarray_name: + libnode.add_in_connector('_subarray', sdfg.arrays[subarray_name].state_field_dtype) + state.add_edge(state.add_read(subarray_name), None, libnode, '_subarray', Memlet(data=subarray_name)) + + if scatter_grid: + libnode.add_in_connector('_scatter_grid', sdfg.arrays[scatter_grid].dtype) + state.add_edge(state.add_read(scatter_grid), None, libnode, '_scatter_grid', Memlet(data=scatter_grid)) + if bcast_grid: + libnode.add_in_connector('_bcast_grid', sdfg.arrays[bcast_grid].dtype) + state.add_edge(state.add_read(bcast_grid), None, libnode, '_bcast_grid', Memlet(data=bcast_grid)) + state.add_edge(inbuf_node, None, libnode, '_inp_buffer', inbuf_mem) state.add_edge(libnode, '_out_buffer', outbuf_node, None, outbuf_mem) @@ -1034,16 +1079,16 @@ def _block_gather(pv: ProgramVisitor, if in_desc.dtype != out_desc.dtype: raise ValueError("Input/output buffer datatypes must match!") - subarray_name = _subarray(pv, - sdfg, - state, - out_buffer, - in_buffer, - process_grid=gather_grid, - correspondence=correspondence) + subarray_name, subarray_node = _define_subarray(pv, + sdfg, + state, + out_buffer, + in_buffer, + process_grid=gather_grid, + correspondence=correspondence) from dace.libraries.mpi import BlockGather - libnode = BlockGather('_BlockGather_', subarray_name, gather_grid, reduce_grid) + libnode = BlockGather('_BlockGather_') inbuf_name = in_buffer in_desc = sdfg.arrays[inbuf_name] @@ -1055,6 +1100,20 @@ def _block_gather(pv: ProgramVisitor, outbuf_node = state.add_write(outbuf_name) outbuf_mem = Memlet.from_array(outbuf_name, out_desc) + if subarray_node is not None: + libnode.add_in_connector('_subarray', sdfg.arrays[subarray_name].state_field_dtype) + state.add_edge(subarray_node, None, libnode, '_subarray', Memlet(data=subarray_name)) + elif subarray_name: + libnode.add_in_connector('_subarray', sdfg.arrays[subarray_name].state_field_dtype) + state.add_edge(state.add_read(subarray_name), None, libnode, '_subarray', Memlet(data=subarray_name)) + + if gather_grid: + libnode.add_in_connector('_gather_grid', sdfg.arrays[gather_grid].dtype) + state.add_edge(state.add_read(gather_grid), None, libnode, '_gather_grid', Memlet(data=gather_grid)) + if reduce_grid: + libnode.add_in_connector('_reduce_grid', sdfg.arrays[reduce_grid].dtype) + state.add_edge(state.add_read(reduce_grid), None, libnode, '_reduce_grid', Memlet(data=reduce_grid)) + state.add_edge(inbuf_node, None, libnode, '_inp_buffer', inbuf_mem) state.add_edge(libnode, '_out_buffer', outbuf_node, None, outbuf_mem) @@ -1075,7 +1134,8 @@ def _redistribute(pv: ProgramVisitor, sdfg: SDFG, state: SDFGState, in_buffer: s in_desc = sdfg.arrays[in_buffer] out_desc = sdfg.arrays[out_buffer] - rdistrarray_name = sdfg.add_rdistrarray(in_subarray, out_subarray) + name = pv.get_target_name() + rdistrarray_name = sdfg.add_rdistrarray(in_subarray, out_subarray, name=name) from dace.libraries.mpi import Dummy, Redistribute tasklet = Dummy(rdistrarray_name, [ @@ -1087,9 +1147,8 @@ def _redistribute(pv: ProgramVisitor, sdfg: SDFG, state: SDFGState, in_buffer: s f'int* {rdistrarray_name}_self_size;' ]) state.add_node(tasklet) - scal_name, scal = sdfg.add_scalar(rdistrarray_name, dace.int32, transient=True, find_new_name=True) - wnode = state.add_write(scal_name) - state.add_edge(tasklet, '__out', wnode, None, Memlet.from_array(scal_name, scal)) + wnode = state.add_write(rdistrarray_name) + state.add_edge(tasklet, None, wnode, None, Memlet()) libnode = Redistribute('_Redistribute_', rdistrarray_name) @@ -1118,6 +1177,8 @@ def _redistribute(pv: ProgramVisitor, sdfg: SDFG, state: SDFGState, in_buffer: s else: outbuf_mem = Memlet.from_array(outbuf_name, out_desc) + libnode.add_in_connector('_rdistrarray', sdfg.arrays[rdistrarray_name].state_field_dtype) + state.add_edge(wnode, None, libnode, '_rdistrarray', Memlet(data=rdistrarray_name)) state.add_edge(inbuf_node, None, libnode, '_inp_buffer', inbuf_mem) state.add_edge(libnode, '_out_buffer', outbuf_node, None, outbuf_mem) diff --git a/dace/frontend/python/replacements/operators.py b/dace/frontend/python/replacements/operators.py index a996b2edc6..9fce562074 100644 --- a/dace/frontend/python/replacements/operators.py +++ b/dace/frontend/python/replacements/operators.py @@ -15,7 +15,7 @@ import numpy as np import sympy as sp -import dace # For evaluation of data types +import dace # noqa: F401 (used during evaluation of data types, e.g. casting in replaced op) numpy_version = np.lib.NumpyVersion(np.__version__) diff --git a/dace/frontend/python/replacements/reduction.py b/dace/frontend/python/replacements/reduction.py index ef7f165033..cd7938ece4 100644 --- a/dace/frontend/python/replacements/reduction.py +++ b/dace/frontend/python/replacements/reduction.py @@ -3,6 +3,7 @@ Contains replacements of reduction operations, which cover both NumPy's Mathematical Functions (e.g., ``numpy.sum``) and Sorting, Searching, and Counting Functions (e.g., ``numpy.argmax``). """ +import dace # noqa from dace.frontend.common import op_repository as oprepo from dace.frontend.python.nested_call import NestedCall from dace.frontend.python.replacements.utils import ProgramVisitor, normalize_axes diff --git a/dace/frontend/python/replacements/torch_autodiff.py b/dace/frontend/python/replacements/torch_autodiff.py index d110046be9..2433ae567b 100644 --- a/dace/frontend/python/replacements/torch_autodiff.py +++ b/dace/frontend/python/replacements/torch_autodiff.py @@ -3,17 +3,11 @@ Integration with the dace python frontend """ -from contextlib import contextmanager from typing import Optional, Union, Sequence import itertools -import warnings -import torch -import torch.autograd - -from dace import SDFG, SDFGState, config, data +from dace import SDFG, SDFGState, data import dace.sdfg.sdfg -from dace.transformation import optimizer from dace.frontend.python import common from dace.frontend.common import op_repository from dace.frontend.python import newast diff --git a/dace/frontend/python/replacements/ufunc.py b/dace/frontend/python/replacements/ufunc.py index 317439b021..f924007571 100644 --- a/dace/frontend/python/replacements/ufunc.py +++ b/dace/frontend/python/replacements/ufunc.py @@ -2,6 +2,7 @@ """ Contains replacements for NumPy ufuncs. """ +import dace # noqa from dace.frontend.common import op_repository as oprepo from dace.frontend.python import astutils from dace.frontend.python.nested_call import NestedCall diff --git a/dace/frontend/python/replacements/utils.py b/dace/frontend/python/replacements/utils.py index 0d04591bb6..670a95148e 100644 --- a/dace/frontend/python/replacements/utils.py +++ b/dace/frontend/python/replacements/utils.py @@ -13,7 +13,6 @@ import numpy as np import sympy as sp -import dace # For evaluation of data types ######################################################################## # Type hint definitions diff --git a/dace/frontend/python/tasklet_runner.py b/dace/frontend/python/tasklet_runner.py index b967b13094..01210463d3 100644 --- a/dace/frontend/python/tasklet_runner.py +++ b/dace/frontend/python/tasklet_runner.py @@ -7,15 +7,10 @@ import ast import copy import inspect -import sys from typing import Any, Dict, List, Optional, Tuple -from dace import data, symbolic -from dace.config import Config -from dace.frontend.python import ndloop, wrappers from dace.frontend.python import astutils -from dace.frontend.python.astutils import unparse, rname -from dace.frontend.python.parser import DaceProgram +from dace.frontend.python.astutils import rname def get_tasklet_ast(stack_depth=2, frame=None) -> ast.With: diff --git a/dace/frontend/python/wrappers.py b/dace/frontend/python/wrappers.py index 49f7b32c78..bfd8c5aa8a 100644 --- a/dace/frontend/python/wrappers.py +++ b/dace/frontend/python/wrappers.py @@ -5,7 +5,7 @@ from collections import deque from typing import Deque, Generic, Type, TypeVar -from dace import dtypes, symbolic +from dace import dtypes T = TypeVar('T') diff --git a/dace/libraries/blas/blas_helpers.py b/dace/libraries/blas/blas_helpers.py index 42a5c3287b..10a5052756 100644 --- a/dace/libraries/blas/blas_helpers.py +++ b/dace/libraries/blas/blas_helpers.py @@ -1,7 +1,6 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. import numpy as np from dace import dtypes, data -from dace.data import Array from typing import Any, Dict, Tuple diff --git a/dace/libraries/blas/nodes/axpy.py b/dace/libraries/blas/nodes/axpy.py index 4c6a02906f..c82eabf9f7 100644 --- a/dace/libraries/blas/nodes/axpy.py +++ b/dace/libraries/blas/nodes/axpy.py @@ -3,8 +3,7 @@ import dace.properties import dace.sdfg.nodes from dace.transformation.transformation import ExpandTransformation -from dace.libraries.blas import environments -from dace import config, data as dt, dtypes, memlet as mm, SDFG, SDFGState, symbolic +from dace import data as dt, memlet as mm, SDFG, SDFGState from dace.frontend.common import op_repository as oprepo diff --git a/dace/libraries/blas/nodes/dot.py b/dace/libraries/blas/nodes/dot.py index 1b78f8334f..42ce0c0fa8 100644 --- a/dace/libraries/blas/nodes/dot.py +++ b/dace/libraries/blas/nodes/dot.py @@ -7,7 +7,7 @@ from dace.transformation.transformation import ExpandTransformation from dace.libraries.blas import blas_helpers from .. import environments -from dace import data as dt, dtypes, memlet as mm, SDFG, SDFGState, symbolic +from dace import dtypes, memlet as mm, SDFG, SDFGState from dace.frontend.common import op_repository as oprepo diff --git a/dace/libraries/blas/nodes/gemv.py b/dace/libraries/blas/nodes/gemv.py index 0cc0ff147f..9ca6368b45 100644 --- a/dace/libraries/blas/nodes/gemv.py +++ b/dace/libraries/blas/nodes/gemv.py @@ -4,7 +4,7 @@ import dace.library import dace.sdfg.nodes from dace.sdfg import SDFG, SDFGState -from dace import memlet as mm, data as dt +from dace import memlet as mm from dace.transformation.transformation import ExpandTransformation from dace.libraries.blas.nodes.matmul import _get_matmul_operands from dace.libraries.blas import blas_helpers diff --git a/dace/libraries/fft/nodes/fft.py b/dace/libraries/fft/nodes/fft.py index 2724fee8a6..06abd0e62b 100644 --- a/dace/libraries/fft/nodes/fft.py +++ b/dace/libraries/fft/nodes/fft.py @@ -2,7 +2,6 @@ """ Implements Forward and Inverse Fast Fourier Transform (FFT) library nodes """ -import warnings from dace import data, dtypes, SDFG, SDFGState, symbolic, library, nodes, properties from dace import transformation as xf diff --git a/dace/libraries/lapack/nodes/getri.py b/dace/libraries/lapack/nodes/getri.py index a2de84fb6b..4c8ebfb8f2 100644 --- a/dace/libraries/lapack/nodes/getri.py +++ b/dace/libraries/lapack/nodes/getri.py @@ -4,9 +4,6 @@ import dace.properties import dace.sdfg.nodes from dace.transformation.transformation import ExpandTransformation -from .. import environments -from dace import data as dt, dtypes, memlet as mm, SDFG, SDFGState, symbolic -from dace.frontend.common import op_repository as oprepo from dace.libraries.blas import environments as blas_environments from dace.libraries.blas import blas_helpers diff --git a/dace/libraries/lapack/nodes/getrs.py b/dace/libraries/lapack/nodes/getrs.py index e6080644a6..5dcfb154df 100644 --- a/dace/libraries/lapack/nodes/getrs.py +++ b/dace/libraries/lapack/nodes/getrs.py @@ -5,8 +5,7 @@ import dace.sdfg.nodes from dace.transformation.transformation import ExpandTransformation from .. import environments -from dace import data as dt, dtypes, memlet as mm, SDFG, SDFGState, symbolic -from dace.frontend.common import op_repository as oprepo +from dace import dtypes from dace.libraries.blas import environments as blas_environments from dace.libraries.blas import blas_helpers diff --git a/dace/libraries/linalg/__init__.py b/dace/libraries/linalg/__init__.py index 2673dc66cb..785811db62 100644 --- a/dace/libraries/linalg/__init__.py +++ b/dace/libraries/linalg/__init__.py @@ -1,5 +1,6 @@ -# Copyright 2019-2023 ETH Zurich and the DaCe authors. All rights reserved. +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. from dace.library import register_library +from .environments import * from .nodes import * register_library(__name__, "linalg") diff --git a/dace/libraries/linalg/environments/__init__.py b/dace/libraries/linalg/environments/__init__.py index f8678e6b8e..c736edbd39 100644 --- a/dace/libraries/linalg/environments/__init__.py +++ b/dace/libraries/linalg/environments/__init__.py @@ -1,2 +1,3 @@ -# Copyright 2019-2022 ETH Zurich and the DaCe authors. All rights reserved. -from .cutensor import * +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. +from .hptt import HPTT +from .cutensor import cuTensor diff --git a/dace/libraries/linalg/environments/cutensor.py b/dace/libraries/linalg/environments/cutensor.py index 0022ec1f57..4a03061e57 100644 --- a/dace/libraries/linalg/environments/cutensor.py +++ b/dace/libraries/linalg/environments/cutensor.py @@ -1,9 +1,11 @@ -# Copyright 2019-2023 ETH Zurich and the DaCe authors. All rights reserved. +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. +"""DaCe library environment for the NVIDIA cuTENSOR backend.""" import dace.library @dace.library.environment class cuTensor: + """Build/link configuration and per-node setup code for cuTENSOR-backed library nodes.""" cmake_minimum_version = None cmake_packages = ["CUDA"] @@ -14,12 +16,27 @@ class cuTensor: cmake_link_flags = ["-L -lcutensor"] cmake_files = [] - headers = {'frame': ["../include/dace_cutensor.h"], 'cuda': ["../include/dace_cutensor.h"]} + headers = {'frame': ["dace/dace_cutensor.h"], 'cuda': ["dace/dace_cutensor.h"]} state_fields = ["dace::linalg::CuTensorHandle cutensor_handle;"] init_code = "" finalize_code = "" dependencies = [] + # cuTENSOR v2 type mapping: dtype -> (tensor data type, compute descriptor, + # C scalar type for alpha/beta). The scalar type follows the compute + # descriptor (real even when the tensors are complex). Only float/complex + # are listed: integer descriptors are accepted by the API but crash at + # execution, and reinterpreting ints as floats corrupts negatives (their + # bit patterns are NaN, which the GPU canonicalizes on multiply). Callers + # fall back to the pure expansion for unsupported dtypes. + TYPE_MAP = { + dace.float16: ('CUTENSOR_R_16F', 'CUTENSOR_COMPUTE_DESC_16F', '__half'), + dace.float32: ('CUTENSOR_R_32F', 'CUTENSOR_COMPUTE_DESC_32F', 'float'), + dace.float64: ('CUTENSOR_R_64F', 'CUTENSOR_COMPUTE_DESC_64F', 'double'), + dace.complex64: ('CUTENSOR_C_32F', 'CUTENSOR_COMPUTE_DESC_32F', 'float'), + dace.complex128: ('CUTENSOR_C_64F', 'CUTENSOR_COMPUTE_DESC_64F', 'double'), + } + @staticmethod def handle_setup_code(node): location = node.location diff --git a/dace/libraries/standard/environments/hptt.py b/dace/libraries/linalg/environments/hptt.py similarity index 93% rename from dace/libraries/standard/environments/hptt.py rename to dace/libraries/linalg/environments/hptt.py index 9d65da7111..2927b9e9ee 100644 --- a/dace/libraries/standard/environments/hptt.py +++ b/dace/libraries/linalg/environments/hptt.py @@ -1,4 +1,4 @@ -# Copyright 2019-2023 ETH Zurich and the DaCe authors. All rights reserved. +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. import os from dace import config, library diff --git a/dace/libraries/linalg/nodes/__init__.py b/dace/libraries/linalg/nodes/__init__.py index 5df44ba4a5..793b540dac 100644 --- a/dace/libraries/linalg/nodes/__init__.py +++ b/dace/libraries/linalg/nodes/__init__.py @@ -1,4 +1,6 @@ -# Copyright 2019-2023 ETH Zurich and the DaCe authors. All rights reserved. +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. +from .transpose import Transpose +from .ttranspose import TensorTranspose from .inv import Inv from .solve import Solve from .cholesky import Cholesky diff --git a/dace/libraries/linalg/nodes/cholesky.py b/dace/libraries/linalg/nodes/cholesky.py index ecf80f9364..0aaa468ca9 100644 --- a/dace/libraries/linalg/nodes/cholesky.py +++ b/dace/libraries/linalg/nodes/cholesky.py @@ -1,4 +1,4 @@ -# Copyright 2019-2023 ETH Zurich and the DaCe authors. All rights reserved. +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. import copy import dace.library import dace.properties @@ -6,7 +6,7 @@ from dace import Memlet from dace.libraries.lapack import Potrf -from dace.libraries.standard import Transpose +from dace.libraries.linalg.nodes.transpose import Transpose from dace.transformation.transformation import ExpandTransformation from dace.libraries.lapack import environments from dace.libraries.blas import environments as blas_environments diff --git a/dace/libraries/linalg/nodes/solve.py b/dace/libraries/linalg/nodes/solve.py index 06ec595399..92a94f941f 100644 --- a/dace/libraries/linalg/nodes/solve.py +++ b/dace/libraries/linalg/nodes/solve.py @@ -1,4 +1,4 @@ -# Copyright 2019-2023 ETH Zurich and the DaCe authors. All rights reserved. +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. import copy import dace import dace.library @@ -9,7 +9,7 @@ from dace import Memlet, SDFG, SDFGState from dace import symbolic from dace.libraries.lapack import Getrf, Getrs -from dace.libraries.standard import Transpose +from dace.libraries.linalg.nodes.transpose import Transpose from dace.transformation.transformation import ExpandTransformation from dace.libraries.lapack import environments from dace.libraries.blas import environments as blas_environments diff --git a/dace/libraries/linalg/nodes/tensordot.py b/dace/libraries/linalg/nodes/tensordot.py index b975e6046d..23224d6a1c 100644 --- a/dace/libraries/linalg/nodes/tensordot.py +++ b/dace/libraries/linalg/nodes/tensordot.py @@ -1,11 +1,11 @@ -# Copyright 2019-2023 ETH Zurich and the DaCe authors. All rights reserved. +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. +"""TensorDot library node and its pure / TTGT / cuTENSOR expansions.""" import collections import dace -import dace.libraries.linalg.environments as environments +from dace.libraries.linalg import environments from dace import library, nodes, properties from dace.utils import prod as _prod -from dace.libraries.blas import blas_helpers from dace.symbolic import symstr from dace.transformation.transformation import ExpandTransformation @@ -155,7 +155,7 @@ def expansion(node, parent_state, parent_sdfg): right_tt = _transpose(None, sdfg, state, "_right_tensor", right_axes, outname="ttgt_right_transposed") right_tt_arr = sdfg.arrays[right_tt] - from dace.libraries.blas import Gemm + from dace.libraries.blas import Gemm # Avoid import loop prv_state = state state = sdfg.add_state(f"{node.label}_gemm_state") sdfg.add_edge(prv_state, state, dace.InterstateEdge()) @@ -234,7 +234,8 @@ def expansion(node, parent_state, parent_sdfg): state.add_edge(tasklet, '_c', dot_vnode, None, dace.Memlet.from_array(dot_vname, dot_view)) state.add_edge(dot_vnode, 'views', dot_anode, None, dace.Memlet.from_array(dot_name, dot_arr)) out_node = state.add_write('_out_tensor') - from dace.libraries.standard import TensorTranspose + # Avoid import loop: TensorTranspose is a sibling node in dace.libraries.linalg + from dace.libraries.linalg import TensorTranspose tasklet = TensorTranspose('_TensorTranspose', node.permutation) state.add_edge(dot_anode, None, tasklet, '_inp_tensor', dace.Memlet.from_array(dot_name, dot_arr)) state.add_edge(tasklet, '_out_tensor', out_node, None, dace.Memlet.from_array('_out_tensor', out_arr)) @@ -257,8 +258,13 @@ def expansion(node, parent_state, parent_sdfg): @library.expansion class ExpandCuTensor(ExpandTransformation): """ - Implements the TensorDot library node using cuTENSOR for CUDA-compatible GPUs. - For more information, see https://developer.nvidia.com/cutensor. + Implements the TensorDot library node using cuTENSOR v2 (cutensorContract) + for CUDA-compatible GPUs. Requires cuTENSOR >= 2.0. + + The contraction expresses: + D_{modesD} = alpha * A_{modesA} * B_{modesB} + beta * C_{modesC} + where in this kernel C and D share the same buffer and modes + (i.e. D = alpha * A * B with beta = 0). """ environments = [environments.cuTensor] @@ -268,16 +274,16 @@ def expansion(node, parent_state, parent_sdfg): left_tensor, right_tensor, out_tensor = node.validate(parent_sdfg, parent_state) dtype = out_tensor.dtype.base_type - func, cuda_type, _ = blas_helpers.cublas_type_metadata(dtype) - cuda_dtype = blas_helpers.dtype_to_cudadatatype(dtype) - compute_type = f"CUTENSOR_COMPUTE{cuda_dtype[cuda_dtype.rfind('_'):]}" - func = func + 'getrf' + if dtype not in environments.cuTensor.TYPE_MAP: + raise NotImplementedError(f"cuTENSOR TensorDot does not support dtype {dtype}; supported: " + f"{sorted(str(t) for t in environments.cuTensor.TYPE_MAP)}") + cutensor_dtype, compute_desc, scalar_type = environments.cuTensor.TYPE_MAP[dtype] - alpha = f"({cuda_type})1.0" - beta = f"({cuda_type})0.0" + alpha = f"({scalar_type})1.0" + beta = f"({scalar_type})0.0" abtext = f""" - {cuda_type} alpha = {alpha}; - {cuda_type} beta = {beta}; + {scalar_type} alpha = {alpha}; + {scalar_type} beta = {beta}; """ left_modes = list(range(len(left_tensor.shape))) @@ -291,9 +297,9 @@ def expansion(node, parent_state, parent_sdfg): out_modes = [out_modes[i] for i in node.permutation] modes = f""" - std::vector modeA{{{','.join(str(m) for m in left_modes)}}}; - std::vector modeB{{{','.join(str(m) for m in right_modes)}}}; - std::vector modeC{{{','.join(str(m) for m in out_modes)}}}; + std::vector modeA{{{','.join(str(m) for m in left_modes)}}}; + std::vector modeB{{{','.join(str(m) for m in right_modes)}}}; + std::vector modeC{{{','.join(str(m) for m in out_modes)}}}; """ extents = "std::unordered_map extent;\n" @@ -318,58 +324,65 @@ def expansion(node, parent_state, parent_sdfg): std::vector stridesC{{{','.join(str(s) for s in out_tensor.strides)}}}; """ + # cuTENSOR v2: descriptors take an alignment hint (bytes) instead of + # a per-pointer query; 256 is safe for all CUDA allocations. tdesc = f""" cutensorTensorDescriptor_t descA, descB, descC; - dace::linalg::CheckCuTensorError(cutensorInitTensorDescriptor( - &__dace_cutensor_handle, &descA, modeA.size(), extentA.data(), stridesA.data(), {cuda_dtype}, CUTENSOR_OP_IDENTITY)); - dace::linalg::CheckCuTensorError(cutensorInitTensorDescriptor( - &__dace_cutensor_handle, &descB, modeB.size(), extentB.data(), stridesB.data(), {cuda_dtype}, CUTENSOR_OP_IDENTITY)); - dace::linalg::CheckCuTensorError(cutensorInitTensorDescriptor( - &__dace_cutensor_handle, &descC, modeC.size(), extentC.data(), stridesC.data(), {cuda_dtype}, CUTENSOR_OP_IDENTITY)); - // printf("Tensor descriptors created!\\n"); + dace::linalg::CheckCuTensorError(cutensorCreateTensorDescriptor( + __dace_cutensor_handle, &descA, modeA.size(), + extentA.data(), stridesA.data(), {cutensor_dtype}, 256)); + dace::linalg::CheckCuTensorError(cutensorCreateTensorDescriptor( + __dace_cutensor_handle, &descB, modeB.size(), + extentB.data(), stridesB.data(), {cutensor_dtype}, 256)); + dace::linalg::CheckCuTensorError(cutensorCreateTensorDescriptor( + __dace_cutensor_handle, &descC, modeC.size(), + extentC.data(), stridesC.data(), {cutensor_dtype}, 256)); """ + # Contraction descriptor: D = alpha * A * B + beta * C; here D == C. cdesc = f""" - uint32_t alignmentRequirementA, alignmentRequirementB, alignmentRequirementC; - dace::linalg::CheckCuTensorError(cutensorGetAlignmentRequirement(&__dace_cutensor_handle, _left_tensor, &descA, &alignmentRequirementA)); - dace::linalg::CheckCuTensorError(cutensorGetAlignmentRequirement(&__dace_cutensor_handle, _right_tensor, &descB, &alignmentRequirementB)); - dace::linalg::CheckCuTensorError(cutensorGetAlignmentRequirement(&__dace_cutensor_handle, _out_tensor, &descC, &alignmentRequirementC)); - cutensorContractionDescriptor_t desc; - dace::linalg::CheckCuTensorError(cutensorInitContractionDescriptor( - &__dace_cutensor_handle, &desc, - &descA, modeA.data(), alignmentRequirementA, - &descB, modeB.data(), alignmentRequirementB, - &descC, modeC.data(), alignmentRequirementC, - &descC, modeC.data(), alignmentRequirementC, - {compute_type})); - // printf("Memory alignment and coontraction descriptor created!\\n"); + cutensorOperationDescriptor_t opDesc; + dace::linalg::CheckCuTensorError(cutensorCreateContraction( + __dace_cutensor_handle, &opDesc, + descA, modeA.data(), CUTENSOR_OP_IDENTITY, + descB, modeB.data(), CUTENSOR_OP_IDENTITY, + descC, modeC.data(), CUTENSOR_OP_IDENTITY, + descC, modeC.data(), + {compute_desc})); """ workspace = """ - cutensorContractionFind_t find; - dace::linalg::CheckCuTensorError(cutensorInitContractionFind(&__dace_cutensor_handle, &find, CUTENSOR_ALGO_DEFAULT)); - size_t worksize = 0; - dace::linalg::CheckCuTensorError(cutensorContractionGetWorkspace( - &__dace_cutensor_handle, &desc, &find, CUTENSOR_WORKSPACE_RECOMMENDED, &worksize)); + cutensorPlanPreference_t planPref; + dace::linalg::CheckCuTensorError(cutensorCreatePlanPreference( + __dace_cutensor_handle, &planPref, + CUTENSOR_ALGO_DEFAULT, CUTENSOR_JIT_MODE_DEFAULT)); + uint64_t worksize = 0; + dace::linalg::CheckCuTensorError(cutensorEstimateWorkspaceSize( + __dace_cutensor_handle, opDesc, planPref, + CUTENSOR_WORKSPACE_DEFAULT, &worksize)); void *work = nullptr; if (worksize > 0) cudaMalloc(&work, worksize); - // printf("Workspace created!\\n"); """ execute = """ - cutensorContractionPlan_t plan; - dace::linalg::CheckCuTensorError(cutensorInitContractionPlan(&__dace_cutensor_handle, &plan, &desc, &find, worksize)); - cutensorStatus_t err; - err = cutensorContraction( - &__dace_cutensor_handle, &plan, - (void*)&alpha, _left_tensor, _right_tensor, (void*)&beta, _out_tensor, _out_tensor, + cutensorPlan_t plan; + dace::linalg::CheckCuTensorError(cutensorCreatePlan( + __dace_cutensor_handle, &plan, opDesc, planPref, worksize)); + cutensorStatus_t err = cutensorContract( + __dace_cutensor_handle, plan, + (const void*)&alpha, _left_tensor, _right_tensor, + (const void*)&beta, _out_tensor, _out_tensor, work, worksize, __dace_current_stream); - cudaStreamSynchronize(__dace_current_stream); - if(err != CUTENSOR_STATUS_SUCCESS) { + if (err != CUTENSOR_STATUS_SUCCESS) { printf("ERROR: %s\\n", cutensorGetErrorString(err)); } + cutensorDestroyPlan(plan); + cutensorDestroyPlanPreference(planPref); + cutensorDestroyOperationDescriptor(opDesc); + cutensorDestroyTensorDescriptor(descC); + cutensorDestroyTensorDescriptor(descB); + cutensorDestroyTensorDescriptor(descA); if (work) cudaFree(work); - // printf("Contraction executed!\\n"); """ code = f"{environments.cuTensor.handle_setup_code(node)}{abtext}{modes}{extents}{tdesc}{cdesc}{workspace}{execute}" diff --git a/dace/libraries/standard/nodes/transpose.py b/dace/libraries/linalg/nodes/transpose.py similarity index 99% rename from dace/libraries/standard/nodes/transpose.py rename to dace/libraries/linalg/nodes/transpose.py index 73a01a25d7..7676f09a4a 100644 --- a/dace/libraries/standard/nodes/transpose.py +++ b/dace/libraries/linalg/nodes/transpose.py @@ -1,4 +1,4 @@ -# Copyright 2019-2023 ETH Zurich and the DaCe authors. All rights reserved. +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. import functools from copy import deepcopy as dc import dace.library diff --git a/dace/libraries/standard/nodes/ttranspose.py b/dace/libraries/linalg/nodes/ttranspose.py similarity index 51% rename from dace/libraries/standard/nodes/ttranspose.py rename to dace/libraries/linalg/nodes/ttranspose.py index 6d142db81f..3cebcb221d 100644 --- a/dace/libraries/standard/nodes/ttranspose.py +++ b/dace/libraries/linalg/nodes/ttranspose.py @@ -1,12 +1,13 @@ -# Copyright 2019-2023 ETH Zurich and the DaCe authors. All rights reserved. +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. +"""TensorTranspose library node and its pure / HPTT / cuTENSOR expansions.""" import dace import multiprocessing from dace import library, nodes, properties -from dace.libraries.blas import blas_helpers -from dace.symbolic import symstr from dace.transformation.transformation import ExpandTransformation +from dace.libraries.blas import blas_helpers from numbers import Number -from .. import environments +from dace.libraries.linalg import environments +import warnings @library.expansion @@ -61,14 +62,16 @@ class ExpandHPTT(ExpandTransformation): @staticmethod def expansion(node, parent_state, parent_sdfg): + from dace.codegen.common import sym2cpp # Avoid import loop + inp_tensor, out_tensor = node.validate(parent_sdfg, parent_state) - axes = ','.join([symstr(a) for a in node.axes]) - shape = ','.join([symstr(s) for s in inp_tensor.shape]) + axes = ','.join([sym2cpp(a) for a in node.axes]) + shape = ','.join([sym2cpp(s) for s in inp_tensor.shape]) dchar = blas_helpers.to_blastype(inp_tensor.dtype.type).lower() if dchar not in ('s', 'd', 'c', 'z'): raise TypeError("HPTT supports only single and double (and corresponding complex) FP datatypes") - alpha = symstr(node.alpha) - beta = symstr(node.beta) + alpha = sym2cpp(node.alpha) + beta = sym2cpp(node.beta) code = f""" int perm[{len(inp_tensor.shape)}] = {{{axes}}}; int size[{len(inp_tensor.shape)}] = {{{shape}}}; @@ -84,11 +87,135 @@ def expansion(node, parent_state, parent_sdfg): return tasklet +@library.expansion +class ExpandCuTensor(ExpandTransformation): + """ + Implements the TensorTranspose library node using the cuTENSOR v2 API + (cutensorPermute). Requires cuTENSOR >= 2.0. + + The permutation is expressed as: + C_{modesC} = alpha * A_{modesA} + where modesA is the identity [0, 1, ..., n-1] and modesC encodes the + axis permutation. + + NOTE: beta != 0 is not supported by cutensorPermute (its signature is + ``(handle, plan, alpha, A, B, stream)`` -- out-of-place ``B = alpha*op(A)``, + no beta term). For C = alpha * perm(A) + beta * C, use ExpandPure or + implement via cutensorElementwiseBinary. + See https://docs.nvidia.com/cuda/cutensor/latest/api/cutensor.html#cutensorpermute + """ + + environments = [environments.cuTensor] + + @staticmethod + def expansion(node, parent_state, parent_sdfg): + from dace.codegen.common import sym2cpp # Avoid import loop + + inp_tensor, out_tensor = node.validate(parent_sdfg, parent_state) + + if node.beta != 0: + raise NotImplementedError("cuTENSOR v2 cutensorPermute does not support beta != 0. " + "Use the 'pure' expansion or implement via cutensorElementwiseBinary.") + + ndim = len(inp_tensor.shape) + dtype = inp_tensor.dtype.base_type + + if dtype not in environments.cuTensor.TYPE_MAP: + # Fall back to pure expansion for unsupported types (integers, etc.). + # The pure expansion generates a GPU map when data is GPU_Global, + # so integer transposes still execute on the GPU. + warnings.warn("CuTensor does not support integer tensors, falling back to pure implementation") + return ExpandPure.expansion(node, parent_state, parent_sdfg) + + cutensor_dtype, compute_desc, alpha_type = environments.cuTensor.TYPE_MAP[dtype] + alpha_val = f"({alpha_type}){node.alpha}" + + # Input modes: identity mapping [0, 1, ..., n-1] + modes_a = list(range(ndim)) + # Output modes: the permutation [axes[0], axes[1], ...] + modes_c = list(node.axes) + + modes_a_str = ', '.join(str(m) for m in modes_a) + modes_c_str = ', '.join(str(m) for m in modes_c) + extent_a_str = ', '.join(sym2cpp(s) for s in inp_tensor.shape) + extent_c_str = ', '.join(sym2cpp(s) for s in out_tensor.shape) + stride_a_str = ', '.join(sym2cpp(s) for s in inp_tensor.strides) + stride_c_str = ', '.join(sym2cpp(s) for s in out_tensor.strides) + + code = f"""\ +{environments.cuTensor.handle_setup_code(node)} +{{ + // cuTENSOR v2 permutation + const uint32_t kNdim = {ndim}; + + int32_t modesA[{ndim}] = {{{modes_a_str}}}; + int32_t modesC[{ndim}] = {{{modes_c_str}}}; + int64_t extentA[{ndim}] = {{{extent_a_str}}}; + int64_t extentC[{ndim}] = {{{extent_c_str}}}; + int64_t stridesA[{ndim}] = {{{stride_a_str}}}; + int64_t stridesC[{ndim}] = {{{stride_c_str}}}; + + {alpha_type} alpha = {alpha_val}; + + // tensor descriptors (v2: alignment hint in bytes, 256 is safe) + cutensorTensorDescriptor_t descA, descC; + dace::linalg::CheckCuTensorError(cutensorCreateTensorDescriptor( + __dace_cutensor_handle, &descA, kNdim, + extentA, stridesA, {cutensor_dtype}, 256)); + dace::linalg::CheckCuTensorError(cutensorCreateTensorDescriptor( + __dace_cutensor_handle, &descC, kNdim, + extentC, stridesC, {cutensor_dtype}, 256)); + + // operation descriptor (permutation) + cutensorOperationDescriptor_t opDesc; + dace::linalg::CheckCuTensorError(cutensorCreatePermutation( + __dace_cutensor_handle, &opDesc, + descA, modesA, CUTENSOR_OP_IDENTITY, + descC, modesC, + {compute_desc})); + + // plan preference & plan + cutensorPlanPreference_t planPref; + dace::linalg::CheckCuTensorError(cutensorCreatePlanPreference( + __dace_cutensor_handle, &planPref, + CUTENSOR_ALGO_DEFAULT, CUTENSOR_JIT_MODE_DEFAULT)); + + cutensorPlan_t plan; + dace::linalg::CheckCuTensorError(cutensorCreatePlan( + __dace_cutensor_handle, &plan, opDesc, planPref, 0)); + + // execute + dace::linalg::CheckCuTensorError(cutensorPermute( + __dace_cutensor_handle, plan, + (const void*)&alpha, _inp_tensor, _out_tensor, + __dace_current_stream)); + + // cleanup + cutensorDestroyPlan(plan); + cutensorDestroyPlanPreference(planPref); + cutensorDestroyOperationDescriptor(opDesc); + cutensorDestroyTensorDescriptor(descC); + cutensorDestroyTensorDescriptor(descA); +}} +""" + + tasklet = nodes.Tasklet(node.name, + node.in_connectors, + node.out_connectors, + code, + language=dace.dtypes.Language.CPP) + return tasklet + + @library.node class TensorTranspose(nodes.LibraryNode): """ Implements out-of-place tensor transpositions. """ - implementations = {"pure": ExpandPure, "HPTT": ExpandHPTT} + implementations = { + "pure": ExpandPure, + "HPTT": ExpandHPTT, + "cuTENSOR": ExpandCuTensor, + } default_implementation = 'pure' axes = properties.ListProperty(element_type=int, default=[], desc="Permutation of input tensor's modes") diff --git a/dace/libraries/mpi/nodes/allreduce.py b/dace/libraries/mpi/nodes/allreduce.py index 05171e95b3..ae1280a03e 100644 --- a/dace/libraries/mpi/nodes/allreduce.py +++ b/dace/libraries/mpi/nodes/allreduce.py @@ -4,7 +4,7 @@ import dace.sdfg.nodes from dace.transformation.transformation import ExpandTransformation from .. import environments -from dace.libraries.mpi.nodes.node import MPINode +from dace.libraries.mpi.nodes.node import MPINode, expanded_input_connectors, input_descriptor_name @dace.library.expansion @@ -20,8 +20,9 @@ def expansion(node, parent_state, parent_sdfg, n=None, **kwargs): raise (NotImplementedError) comm = "MPI_COMM_WORLD" - if node.grid: - comm = f"__state->{node.grid}_comm" + grid = input_descriptor_name(node, parent_state, '_grid') + if grid: + comm = "_grid" buffer = '_inbuffer' if in_place: @@ -32,7 +33,7 @@ def expansion(node, parent_state, parent_sdfg, n=None, **kwargs): {node.op}, {comm}); """ tasklet = dace.sdfg.nodes.Tasklet(node.name, - node.in_connectors, + expanded_input_connectors(node, parent_state), node.out_connectors, code, language=dace.dtypes.Language.CPP) @@ -49,12 +50,10 @@ class Allreduce(MPINode): default_implementation = "MPI" op = dace.properties.Property(dtype=str, default='MPI_SUM') - grid = dace.properties.Property(dtype=str, allow_none=True, default=None) - def __init__(self, name, op='MPI_SUM', grid=None, *args, **kwargs): + def __init__(self, name, op='MPI_SUM', *args, **kwargs): super().__init__(name, *args, inputs={"_inbuffer"}, outputs={"_outbuffer"}, **kwargs) self.op = op - self.grid = grid def validate(self, sdfg, state): """ diff --git a/dace/libraries/mpi/nodes/alltoall.py b/dace/libraries/mpi/nodes/alltoall.py index 0089efb533..59d7c6ae4a 100644 --- a/dace/libraries/mpi/nodes/alltoall.py +++ b/dace/libraries/mpi/nodes/alltoall.py @@ -1,10 +1,9 @@ # Copyright 2019-2023 ETH Zurich and the DaCe authors. All rights reserved. import dace.library -import dace.properties import dace.sdfg.nodes from dace.transformation.transformation import ExpandTransformation from .. import environments -from dace.libraries.mpi.nodes.node import MPINode +from dace.libraries.mpi.nodes.node import MPINode, expanded_input_connectors, input_descriptor_name @dace.library.expansion @@ -22,8 +21,9 @@ def expansion(node, parent_state, parent_sdfg, n=None, **kwargs): raise (NotImplementedError) comm = "MPI_COMM_WORLD" - if node.grid: - comm = f"__state->{node.grid}_comm" + grid = input_descriptor_name(node, parent_state, '_grid') + if grid: + comm = "_grid" code = f""" int size; @@ -34,7 +34,7 @@ def expansion(node, parent_state, parent_sdfg, n=None, **kwargs): {comm}); """ tasklet = dace.sdfg.nodes.Tasklet(node.name, - node.in_connectors, + expanded_input_connectors(node, parent_state), node.out_connectors, code, language=dace.dtypes.Language.CPP) @@ -50,11 +50,8 @@ class Alltoall(MPINode): } default_implementation = "MPI" - grid = dace.properties.Property(dtype=str, allow_none=True, default=None) - - def __init__(self, name, grid=None, *args, **kwargs): + def __init__(self, name, *args, **kwargs): super().__init__(name, *args, inputs={"_inbuffer"}, outputs={"_outbuffer"}, **kwargs) - self.grid = grid def validate(self, sdfg, state): """ diff --git a/dace/libraries/mpi/nodes/bcast.py b/dace/libraries/mpi/nodes/bcast.py index bf3729ff38..60149739b5 100644 --- a/dace/libraries/mpi/nodes/bcast.py +++ b/dace/libraries/mpi/nodes/bcast.py @@ -6,7 +6,7 @@ from dace.symbolic import symstr from dace.transformation.transformation import ExpandTransformation from .. import environments -from dace.libraries.mpi.nodes.node import MPINode +from dace.libraries.mpi.nodes.node import MPINode, expanded_input_connectors, input_descriptor_name @dace.library.expansion @@ -44,8 +44,9 @@ def expansion(node, parent_state, parent_sdfg, n=None, **kwargs): init = "" comm = "MPI_COMM_WORLD" - if node.grid: - comm = f"__state->{node.grid}_comm" + grid = input_descriptor_name(node, parent_state, '_grid') + if grid: + comm = "_grid" elif node.fcomm: init = f"MPI_Comm __comm = MPI_Comm_f2c({node.fcomm});" comm = "__comm" @@ -55,7 +56,7 @@ def expansion(node, parent_state, parent_sdfg, n=None, **kwargs): MPI_Bcast({ref}_inbuffer, {count_str}, {mpi_dtype_str}, _root, {comm}); _outbuffer = _inbuffer;""" tasklet = dace.sdfg.nodes.Tasklet(node.name, - node.in_connectors, + expanded_input_connectors(node, parent_state), node.out_connectors, code, language=dtypes.Language.CPP) @@ -71,12 +72,10 @@ class Bcast(MPINode): } default_implementation = "MPI" - grid = dace.properties.Property(dtype=str, allow_none=True, default=None) fcomm = dace.properties.Property(dtype=str, allow_none=True, default=None) - def __init__(self, name, grid=None, fcomm=None, *args, **kwargs): + def __init__(self, name, fcomm=None, *args, **kwargs): super().__init__(name, *args, inputs={"_inbuffer", "_root"}, outputs={"_outbuffer"}, **kwargs) - self.grid = grid self.fcomm = fcomm def validate(self, sdfg, state): diff --git a/dace/libraries/mpi/nodes/dummy.py b/dace/libraries/mpi/nodes/dummy.py index d486fa11e8..b59a7c4cc8 100644 --- a/dace/libraries/mpi/nodes/dummy.py +++ b/dace/libraries/mpi/nodes/dummy.py @@ -1,10 +1,8 @@ # Copyright 2019-2022 ETH Zurich and the DaCe authors. All rights reserved. -from dace.libraries.mpi.nodes.node import MPINode +from dace.libraries.mpi.nodes.node import MPINode, expanded_input_connectors import dace.library import dace.properties import dace.sdfg.nodes -from dace.properties import ListProperty -from dace.symbolic import symstr from dace.transformation.transformation import ExpandTransformation from .. import environments @@ -17,8 +15,8 @@ class ExpandDummyMPI(ExpandTransformation): @staticmethod def expansion(node, parent_state, parent_sdfg, *args, **kwargs): tasklet = dace.sdfg.nodes.Tasklet(node.name, - inputs={}, - outputs={'__out'}, + inputs=expanded_input_connectors(node, parent_state), + outputs=node.out_connectors, code='', state_fields=node.fields, side_effects=True) @@ -29,7 +27,7 @@ def expansion(node, parent_state, parent_sdfg, *args, **kwargs): class Dummy(MPINode): """ The Dummy library node is used to add fields to the DaCe program's state (see code generation) that are strictly - related to MPI communication (see distributed descriptors `dace/distr_types`: ProcessGrid, Subarray, RedistrArray). + related to MPI communication (see distributed descriptors `dace.data.distributed`: ProcessGrid, Subarray, RedistrArray). NOTE: This is a temporary solution until the SDFG API has a method that adds fields to the DaCe program's state directly, without a Tasklet or CodeNode. """ @@ -43,7 +41,7 @@ class Dummy(MPINode): fields = dace.properties.ListProperty(default=[], element_type=str) def __init__(self, name, fields=[], *args, **kwargs): - super().__init__(name, *args, outputs={'__out'}, **kwargs) + super().__init__(name, *args, outputs={}, **kwargs) self.fields = fields def validate(self, sdfg, state): diff --git a/dace/libraries/mpi/nodes/gather.py b/dace/libraries/mpi/nodes/gather.py index b231ff7cee..a56b1de094 100644 --- a/dace/libraries/mpi/nodes/gather.py +++ b/dace/libraries/mpi/nodes/gather.py @@ -1,12 +1,12 @@ # Copyright 2019-2022 ETH Zurich and the DaCe authors. All rights reserved. -from dace import dtypes, library, properties +from dace import dtypes, library from dace.utils import prod as _prod from dace.libraries.mpi import utils from dace.sdfg import nodes from dace.symbolic import symstr from dace.transformation.transformation import ExpandTransformation from .. import environments -from dace.libraries.mpi.nodes.node import MPINode +from dace.libraries.mpi.nodes.node import MPINode, expanded_input_connectors, input_descriptor_name @library.expansion @@ -94,23 +94,35 @@ def expansion(node, parent_state, parent_sdfg): if out_buffer.dtype.veclen > 1: raise NotImplementedError - if node.reduce_grid: + subarray_type = input_descriptor_name(node, parent_state, '_subarray') + gather_grid = input_descriptor_name(node, parent_state, '_gather_grid') + reduce_grid = input_descriptor_name(node, parent_state, '_reduce_grid') + if subarray_type is None: + raise ValueError('BlockGather requires an incoming _subarray connector') + if gather_grid is None: + raise ValueError('BlockGather requires an incoming _gather_grid connector') + + if reduce_grid: code = f""" - if (__state->{node.gather_grid}_valid) {{ - MPI_Reduce(MPI_IN_PLACE, _inp_buffer, {symstr(_prod(inp_buffer.shape))}, {mpi_dtype_str}, MPI_SUM, __state->{node.reduce_grid}_rank, __state->{node.reduce_grid}_comm); - MPI_Gatherv(_inp_buffer, {symstr(_prod(inp_buffer.shape))}, {mpi_dtype_str}, _out_buffer, __state->{node.subarray_type}_counts, __state->{node.subarray_type}_displs, __state->{node.subarray_type}, 0, __state->{node.gather_grid}_comm); - }} else if (__state->{node.reduce_grid}_valid) {{ - MPI_Reduce(_inp_buffer, _inp_buffer, {symstr(_prod(inp_buffer.shape))}, {mpi_dtype_str}, MPI_SUM, 0, __state->{node.reduce_grid}_comm); + if (__state->{gather_grid}_valid) {{ + MPI_Reduce(MPI_IN_PLACE, _inp_buffer, {symstr(_prod(inp_buffer.shape))}, {mpi_dtype_str}, MPI_SUM, __state->{reduce_grid}_rank, _reduce_grid); + MPI_Gatherv(_inp_buffer, {symstr(_prod(inp_buffer.shape))}, {mpi_dtype_str}, _out_buffer, __state->{subarray_type}_counts, __state->{subarray_type}_displs, __state->{subarray_type}, 0, _gather_grid); + }} else if (__state->{reduce_grid}_valid) {{ + MPI_Reduce(_inp_buffer, _inp_buffer, {symstr(_prod(inp_buffer.shape))}, {mpi_dtype_str}, MPI_SUM, 0, _reduce_grid); }} """ else: code = f""" - if (__state->{node.gather_grid}_valid) {{ - MPI_Gatherv(_inp_buffer, {symstr(_prod(inp_buffer.shape))}, {mpi_dtype_str}, _out_buffer, __state->{node.subarray_type}_counts, __state->{node.subarray_type}_displs, __state->{node.subarray_type}, 0, __state->{node.gather_grid}_comm); + if (__state->{gather_grid}_valid) {{ + MPI_Gatherv(_inp_buffer, {symstr(_prod(inp_buffer.shape))}, {mpi_dtype_str}, _out_buffer, __state->{subarray_type}_counts, __state->{subarray_type}_displs, __state->{subarray_type}, 0, _gather_grid); }} """ - tasklet = nodes.Tasklet(node.name, node.in_connectors, node.out_connectors, code, language=dtypes.Language.CPP) + tasklet = nodes.Tasklet(node.name, + expanded_input_connectors(node, parent_state), + node.out_connectors, + code, + language=dtypes.Language.CPP) return tasklet @@ -123,15 +135,8 @@ class BlockGather(MPINode): } default_implementation = "MPI" - subarray_type = properties.Property(dtype=str, default='tmp') - gather_grid = properties.Property(dtype=str, default='tmp') - reduce_grid = properties.Property(dtype=str, allow_none=True, default=None) - - def __init__(self, name, subarray_type='tmp', gather_grid='tmp', reduce_grid=None, *args, **kwargs): + def __init__(self, name, *args, **kwargs): super().__init__(name, *args, inputs={"_inp_buffer"}, outputs={"_out_buffer"}, **kwargs) - self.subarray_type = subarray_type - self.gather_grid = gather_grid - self.reduce_grid = reduce_grid def validate(self, sdfg, state): """ diff --git a/dace/libraries/mpi/nodes/irecv.py b/dace/libraries/mpi/nodes/irecv.py index 845655ca25..ac5c686b66 100644 --- a/dace/libraries/mpi/nodes/irecv.py +++ b/dace/libraries/mpi/nodes/irecv.py @@ -1,11 +1,11 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. import dace.library -import dace.properties import dace.sdfg.nodes from dace.transformation.transformation import ExpandTransformation from .. import environments from dace import dtypes -from dace.libraries.mpi.nodes.node import MPINode +from dace.libraries.mpi.nodes.node import (MPINode, expanded_input_connectors, input_descriptor_name, + validate_integer_descriptor) @dace.library.expansion @@ -22,8 +22,9 @@ def expansion(node, parent_state, parent_sdfg, n=None, **kwargs): raise NotImplementedError comm = "MPI_COMM_WORLD" - if node.grid: - comm = f"__state->{node.grid}_comm" + grid = input_descriptor_name(node, parent_state, '_grid') + if grid: + comm = "_grid" code = "" if ddt is not None: @@ -43,7 +44,7 @@ def expansion(node, parent_state, parent_sdfg, n=None, **kwargs): code += f"""// MPI_Type_free(&newtype); """ tasklet = dace.sdfg.nodes.Tasklet(node.name, - node.in_connectors, + expanded_input_connectors(node, parent_state), node.out_connectors, code, language=dace.dtypes.Language.CPP) @@ -63,11 +64,8 @@ class Irecv(MPINode): } default_implementation = "MPI" - grid = dace.properties.Property(dtype=str, allow_none=True, default=None) - - def __init__(self, name, grid=None, *args, **kwargs): + def __init__(self, name, *args, **kwargs): super().__init__(name, *args, inputs={"_src", "_tag"}, outputs={"_buffer", "_request"}, **kwargs) - self.grid = grid def validate(self, sdfg, state): """ @@ -85,6 +83,9 @@ def validate(self, sdfg, state): if e.dst_conn == "_tag": tag = sdfg.arrays[e.data.data] + validate_integer_descriptor(src, 'Source') + validate_integer_descriptor(tag, 'Tag') + count_str = "XXX" for _, src_conn, _, _, data in state.out_edges(self): if src_conn == '_buffer': diff --git a/dace/libraries/mpi/nodes/isend.py b/dace/libraries/mpi/nodes/isend.py index 8d92b16f57..41c43dae7b 100644 --- a/dace/libraries/mpi/nodes/isend.py +++ b/dace/libraries/mpi/nodes/isend.py @@ -5,7 +5,8 @@ from dace.transformation.transformation import ExpandTransformation from .. import environments from dace import dtypes -from dace.libraries.mpi.nodes.node import MPINode +from dace.libraries.mpi.nodes.node import (MPINode, expanded_input_connectors, input_descriptor_name, + validate_integer_descriptor) @dace.library.expansion @@ -22,8 +23,9 @@ def expansion(node, parent_state, parent_sdfg, n=None, **kwargs): raise NotImplementedError comm = "MPI_COMM_WORLD" - if node.grid: - comm = f"__state->{node.grid}_comm" + grid = input_descriptor_name(node, parent_state, '_grid') + if grid: + comm = "_grid" code = "" @@ -50,7 +52,7 @@ def expansion(node, parent_state, parent_sdfg, n=None, **kwargs): """ tasklet = dace.sdfg.nodes.Tasklet(node.name, - node.in_connectors, + expanded_input_connectors(node, parent_state), node.out_connectors, code, language=dace.dtypes.Language.CPP, @@ -73,12 +75,10 @@ class Isend(MPINode): } default_implementation = "MPI" - grid = dace.properties.Property(dtype=str, allow_none=True, default=None) nosync = dace.properties.Property(dtype=bool, default=False, desc="Do not sync if memory is on GPU") - def __init__(self, name, grid=None, *args, **kwargs): + def __init__(self, name, *args, **kwargs): super().__init__(name, *args, inputs={"_buffer", "_dest", "_tag"}, outputs={"_request"}, **kwargs) - self.grid = grid def validate(self, sdfg, state): """ @@ -97,8 +97,8 @@ def validate(self, sdfg, state): if e.src_conn == "_request": req = sdfg.arrays[e.data.data] - # TODO: Should we expect any integer type for dst/tag and cast to int32 later?. - # TODO: Investigate further in the future. + validate_integer_descriptor(dest, 'Destination') + validate_integer_descriptor(tag, 'Tag') count_str = "XXX" for _, _, _, dst_conn, data in state.in_edges(self): diff --git a/dace/libraries/mpi/nodes/node.py b/dace/libraries/mpi/nodes/node.py index 562874e64c..0777b53488 100644 --- a/dace/libraries/mpi/nodes/node.py +++ b/dace/libraries/mpi/nodes/node.py @@ -1,8 +1,41 @@ # Copyright 2019-2022 ETH Zurich and the DaCe authors. All rights reserved. +from typing import Any, Dict, Optional + +from dace import dtypes from dace.sdfg import nodes +def input_descriptor_name(node: nodes.Node, state: Any, connector: str) -> Optional[str]: + edges = list(state.in_edges_by_connector(node, connector)) + if not edges: + return None + if len(edges) > 1: + raise ValueError(f"Expected at most one input edge for MPI connector {connector}, got {len(edges)}") + + edge = edges[0] + if edge.data is not None and edge.data.data is not None: + return edge.data.data + if isinstance(edge.src, nodes.AccessNode): + return edge.src.data + return None + + +def expanded_input_connectors(node: nodes.Node, state: Any) -> Dict[str, Any]: + connectors = dict(node.in_connectors) + for edge in state.in_edges(node): + if edge.dst_conn is not None: + connectors.setdefault(edge.dst_conn, None) + return connectors + + +def validate_integer_descriptor(desc: Any, name: str) -> None: + if desc is None: + raise ValueError(f"{name} connector is missing") + if desc.dtype.base_type not in dtypes.INTEGER_TYPES: + raise ValueError(f"{name} must be an integer") + + class MPINode(nodes.LibraryNode): """ Abstract class representing an MPI library node. diff --git a/dace/libraries/mpi/nodes/recv.py b/dace/libraries/mpi/nodes/recv.py index b027525fd1..9987ba09c6 100644 --- a/dace/libraries/mpi/nodes/recv.py +++ b/dace/libraries/mpi/nodes/recv.py @@ -4,7 +4,7 @@ import dace.sdfg.nodes from dace.transformation.transformation import ExpandTransformation from .. import environments -from dace.libraries.mpi.nodes.node import MPINode +from dace.libraries.mpi.nodes.node import MPINode, validate_integer_descriptor @dace.library.expansion @@ -32,7 +32,7 @@ def expansion(node, parent_state, parent_sdfg, n=None, **kwargs): mpi_dtype_str = "newtype" count_str = '1' buffer_offset = 0 # this is here because the frontend already changes the ptr - code += f"MPI_Recv(_buffer, {count_str}, {mpi_dtype_str}, _src, _tag, MPI_COMM_WORLD, MPI_STATUS_IGNORE);" + code += f"MPI_Recv(_buffer, {count_str}, {mpi_dtype_str}, int(_src), int(_tag), MPI_COMM_WORLD, MPI_STATUS_IGNORE);" if ddt is not None: code += f"""// MPI_Type_free(&newtype); """ @@ -72,6 +72,9 @@ def validate(self, sdfg, state): if e.dst_conn == "_tag": tag = sdfg.arrays[e.data.data] + validate_integer_descriptor(src, 'Source') + validate_integer_descriptor(tag, 'Tag') + count_str = "XXX" for _, src_conn, _, _, data in state.out_edges(self): if src_conn == '_buffer': diff --git a/dace/libraries/mpi/nodes/redistribute.py b/dace/libraries/mpi/nodes/redistribute.py index fb9f8415c7..73bf92dec1 100644 --- a/dace/libraries/mpi/nodes/redistribute.py +++ b/dace/libraries/mpi/nodes/redistribute.py @@ -1,7 +1,5 @@ # Copyright 2019-2022 ETH Zurich and the DaCe authors. All rights reserved. from dace import dtypes, library, properties, subsets, symbolic -from dace.utils import prod as _prod -from dace.libraries.mpi import utils from dace.sdfg import nodes from dace.transformation.transformation import ExpandTransformation from .. import environments @@ -93,7 +91,7 @@ class Redistribute(MPINode): } default_implementation = "MPI" - redistr = properties.Property(dtype=str, default='tmp') + redistr = properties.DataProperty(default='tmp') def __init__(self, name, redistr='tmp', *args, **kwargs): super().__init__(name, *args, inputs={"_inp_buffer"}, outputs={"_out_buffer"}, **kwargs) diff --git a/dace/libraries/mpi/nodes/reduce.py b/dace/libraries/mpi/nodes/reduce.py index 0b56482637..330e3d264c 100644 --- a/dace/libraries/mpi/nodes/reduce.py +++ b/dace/libraries/mpi/nodes/reduce.py @@ -4,7 +4,7 @@ import dace.sdfg.nodes from dace.transformation.transformation import ExpandTransformation from .. import environments -from dace.libraries.mpi.nodes.node import MPINode +from dace.libraries.mpi.nodes.node import MPINode, expanded_input_connectors, input_descriptor_name @dace.library.expansion @@ -22,21 +22,17 @@ def expansion(node, parent_state, parent_sdfg, n=None, **kwargs): raise ValueError("Reduce root must be an integer!") comm = "MPI_COMM_WORLD" - if node.grid: - comm = f"__state->{node.grid}_comm" + grid = input_descriptor_name(node, parent_state, '_grid') + if grid: + comm = "_grid" code = "" if in_place: - if comm == "MPI_COMM_WORLD": - code += """ - int __world_rank; - MPI_Comm_rank(&__world_rank, MPI_COMM_WORLD); - if (__world_rank == _root) {{ - """ - else: - code += f""" - if (__state->{node.grid}_rank == _root) {{ - """ + code += f""" + int __comm_rank; + MPI_Comm_rank({comm}, &__comm_rank); + if (__comm_rank == _root) {{ + """ code += f""" MPI_Reduce(MPI_IN_PLACE, _outbuffer, {count_str}, {mpi_dtype_str}, {node.op}, _root, {comm}); }} else {{ @@ -45,7 +41,7 @@ def expansion(node, parent_state, parent_sdfg, n=None, **kwargs): if in_place: code += "}" tasklet = dace.sdfg.nodes.Tasklet(node.name, - node.in_connectors, + expanded_input_connectors(node, parent_state), node.out_connectors, code, language=dace.dtypes.Language.CPP) @@ -62,12 +58,10 @@ class Reduce(MPINode): default_implementation = "MPI" op = dace.properties.Property(dtype=str, default='MPI_SUM') - grid = dace.properties.Property(dtype=str, allow_none=True, default=None) - def __init__(self, name, op='MPI_SUM', grid=None, *args, **kwargs): + def __init__(self, name, op='MPI_SUM', *args, **kwargs): super().__init__(name, *args, inputs={"_inbuffer", "_root"}, outputs={"_outbuffer"}, **kwargs) self.op = op - self.grid = grid def validate(self, sdfg, state): """ diff --git a/dace/libraries/mpi/nodes/scatter.py b/dace/libraries/mpi/nodes/scatter.py index 59cc54f3a7..8ef42a7d69 100644 --- a/dace/libraries/mpi/nodes/scatter.py +++ b/dace/libraries/mpi/nodes/scatter.py @@ -1,12 +1,12 @@ # Copyright 2019-2022 ETH Zurich and the DaCe authors. All rights reserved. -from dace import dtypes, library, properties +from dace import dtypes, library from dace.utils import prod as _prod from dace.libraries.mpi import utils from dace.sdfg import nodes from dace.symbolic import symstr from dace.transformation.transformation import ExpandTransformation from .. import environments -from dace.libraries.mpi.nodes.node import MPINode +from dace.libraries.mpi.nodes.node import MPINode, expanded_input_connectors, input_descriptor_name @library.expansion @@ -94,19 +94,31 @@ def expansion(node, parent_state, parent_sdfg): if inp_buffer.dtype.veclen > 1: raise NotImplementedError + subarray_type = input_descriptor_name(node, parent_state, '_subarray') + scatter_grid = input_descriptor_name(node, parent_state, '_scatter_grid') + bcast_grid = input_descriptor_name(node, parent_state, '_bcast_grid') + if subarray_type is None: + raise ValueError('BlockScatter requires an incoming _subarray connector') + if scatter_grid is None: + raise ValueError('BlockScatter requires an incoming _scatter_grid connector') + code = f""" - if (__state->{node.scatter_grid}_valid) {{ - MPI_Scatterv(_inp_buffer, __state->{node.subarray_type}_counts, __state->{node.subarray_type}_displs, __state->{node.subarray_type}, _out_buffer, {symstr(_prod(out_buffer.shape))}, {mpi_dtype_str}, 0, __state->{node.scatter_grid}_comm); + if (__state->{scatter_grid}_valid) {{ + MPI_Scatterv(_inp_buffer, __state->{subarray_type}_counts, __state->{subarray_type}_displs, __state->{subarray_type}, _out_buffer, {symstr(_prod(out_buffer.shape))}, {mpi_dtype_str}, 0, _scatter_grid); }} """ - if node.bcast_grid: + if bcast_grid: code += f""" - if (__state->{node.bcast_grid}_valid) {{ - MPI_Bcast(_out_buffer, {symstr(_prod(out_buffer.shape))}, {mpi_dtype_str}, 0, __state->{node._bcast_grid}_comm); + if (__state->{bcast_grid}_valid) {{ + MPI_Bcast(_out_buffer, {symstr(_prod(out_buffer.shape))}, {mpi_dtype_str}, 0, _bcast_grid); }} """ - tasklet = nodes.Tasklet(node.name, node.in_connectors, node.out_connectors, code, language=dtypes.Language.CPP) + tasklet = nodes.Tasklet(node.name, + expanded_input_connectors(node, parent_state), + node.out_connectors, + code, + language=dtypes.Language.CPP) return tasklet @@ -119,15 +131,8 @@ class BlockScatter(MPINode): } default_implementation = "MPI" - subarray_type = properties.Property(dtype=str, default='tmp') - scatter_grid = properties.Property(dtype=str, default='tmp') - bcast_grid = properties.Property(dtype=str, allow_none=True, default=None) - - def __init__(self, name, subarray_type='tmp', scatter_grid='tmp', bcast_grid=None, *args, **kwargs): + def __init__(self, name, *args, **kwargs): super().__init__(name, *args, inputs={"_inp_buffer"}, outputs={"_out_buffer"}, **kwargs) - self.subarray_type = subarray_type - self.scatter_grid = scatter_grid - self.bcast_grid = bcast_grid def validate(self, sdfg, state): """ diff --git a/dace/libraries/mpi/nodes/send.py b/dace/libraries/mpi/nodes/send.py index 9b2382d551..087e57569d 100644 --- a/dace/libraries/mpi/nodes/send.py +++ b/dace/libraries/mpi/nodes/send.py @@ -4,7 +4,7 @@ import dace.sdfg.nodes from dace.transformation.transformation import ExpandTransformation from .. import environments -from dace.libraries.mpi.nodes.node import MPINode +from dace.libraries.mpi.nodes.node import MPINode, validate_integer_descriptor @dace.library.expansion @@ -33,7 +33,7 @@ def expansion(node, parent_state, parent_sdfg, n=None, **kwargs): count_str = "1" buffer_offset = 0 code += f""" - MPI_Send(&(_buffer[{buffer_offset}]), {count_str}, {mpi_dtype_str}, _dest, _tag, MPI_COMM_WORLD); + MPI_Send(&(_buffer[{buffer_offset}]), {count_str}, {mpi_dtype_str}, int(_dest), int(_tag), MPI_COMM_WORLD); """ if ddt is not None: code += f"""// MPI_Type_free(&newtype); @@ -81,10 +81,8 @@ def validate(self, sdfg, state): if e.dst_conn == "_tag": tag = sdfg.arrays[e.data.data] - if dest.dtype.base_type != dace.dtypes.int32: - raise ValueError("Source must be an integer!") - if tag.dtype.base_type != dace.dtypes.int32: - raise ValueError("Tag must be an integer!") + validate_integer_descriptor(dest, 'Destination') + validate_integer_descriptor(tag, 'Tag') count_str = "XXX" for _, _, _, dst_conn, data in state.in_edges(self): diff --git a/dace/libraries/onnx/forward_implementation_abc.py b/dace/libraries/onnx/forward_implementation_abc.py index b43d7fc4b3..7a15e3e0e6 100644 --- a/dace/libraries/onnx/forward_implementation_abc.py +++ b/dace/libraries/onnx/forward_implementation_abc.py @@ -102,4 +102,4 @@ def registered_implementations(cls, op_name: str) -> typing.List[typing.Tuple[st # Import op_implementations to trigger registration of all implementations -import dace.libraries.onnx.op_implementations +import dace.libraries.onnx.op_implementations # noqa: F401 diff --git a/dace/libraries/onnx/nodes/onnx_op_registry.py b/dace/libraries/onnx/nodes/onnx_op_registry.py index 6b145a0859..91d7ab8982 100644 --- a/dace/libraries/onnx/nodes/onnx_op_registry.py +++ b/dace/libraries/onnx/nodes/onnx_op_registry.py @@ -1,15 +1,11 @@ # Copyright 2019-2025 ETH Zurich and the DaCe authors. All rights reserved. import collections -from typing import Iterator, Tuple, List, Dict, Type +from typing import Type -import dace -import dace.library -import dace.sdfg.nodes as nd import dace.frontend.common.op_repository as dace_op_repo from dace.frontend.python.newast import ProgramVisitor -from dace import config, SDFG, SDFGState, dtypes, data -from dace.properties import Property, ListProperty, make_properties -from dace.sdfg.graph import MultiConnectorEdge +from dace import config, SDFG, SDFGState +from dace.properties import Property, ListProperty from dace.transformation.transformation import ExpandTransformation from dace.libraries.onnx.nodes.node_utils import parse_variadic_param diff --git a/dace/libraries/onnx/op_implementations/elementwise_ops.py b/dace/libraries/onnx/op_implementations/elementwise_ops.py index bf66a868d5..1891ada245 100644 --- a/dace/libraries/onnx/op_implementations/elementwise_ops.py +++ b/dace/libraries/onnx/op_implementations/elementwise_ops.py @@ -20,9 +20,9 @@ from dace.libraries.onnx.forward_implementation_abc import ONNXForward from dace.libraries.onnx.nodes import onnx_op -from dace.libraries.onnx.op_implementations.utils import (op_implementation, out_desc_with_name, program_for_node, +from dace.libraries.onnx.op_implementations.utils import (op_implementation, program_for_node, python_pure_op_implementation) -from dace.sdfg.utils import in_desc_with_name, in_edge_with_name, out_desc_with_name +from dace.sdfg.utils import in_desc_with_name, in_edge_with_name from dace.transformation.onnx.replacement import onnx_constant_or_none # ============================================================================ diff --git a/dace/libraries/onnx/op_implementations/img_op_implementations.py b/dace/libraries/onnx/op_implementations/img_op_implementations.py index 1c3727b07c..daa9fd350e 100644 --- a/dace/libraries/onnx/op_implementations/img_op_implementations.py +++ b/dace/libraries/onnx/op_implementations/img_op_implementations.py @@ -7,8 +7,7 @@ import dace from dace import SDFGState, SDFG, dtypes -from dace.sdfg import nodes, propagation -from dace.transformation.dataflow import MapExpansion, MapCollapse +from dace.sdfg import nodes from dace.sdfg.nodes import Node from dace.libraries.onnx.forward_implementation_abc import ONNXForward from dace.libraries.onnx.nodes.onnx_op import ONNXOp diff --git a/dace/libraries/onnx/op_implementations/reduction_ops.py b/dace/libraries/onnx/op_implementations/reduction_ops.py index cd9a361b8e..ecb17b100e 100644 --- a/dace/libraries/onnx/op_implementations/reduction_ops.py +++ b/dace/libraries/onnx/op_implementations/reduction_ops.py @@ -21,8 +21,8 @@ from dace.libraries.onnx.forward_implementation_abc import ONNXForward from dace.libraries.onnx.nodes import onnx_op from dace.libraries.onnx.op_implementations.common import iterables_equal -from dace.libraries.onnx.op_implementations.utils import (empty_sdfg_for_node, in_desc_with_name, op_implementation, - out_desc_with_name, program_for_node) +from dace.libraries.onnx.op_implementations.utils import (in_desc_with_name, op_implementation, out_desc_with_name, + program_for_node) # ============================================================================ # Cumulative Sum diff --git a/dace/libraries/onnx/op_implementations/utils.py b/dace/libraries/onnx/op_implementations/utils.py index 6da253ed86..eaec455b3d 100644 --- a/dace/libraries/onnx/op_implementations/utils.py +++ b/dace/libraries/onnx/op_implementations/utils.py @@ -2,7 +2,6 @@ import inspect import copy from typing import Dict, Tuple, Optional, Callable, Union, Any -import functools import textwrap import dace diff --git a/dace/libraries/onnx/schema.py b/dace/libraries/onnx/schema.py index 86d6f6fbb9..8f61bc03c1 100644 --- a/dace/libraries/onnx/schema.py +++ b/dace/libraries/onnx/schema.py @@ -114,7 +114,7 @@ def from_onnx_proto(cls, onnx_proto): def from_json(cls, json, context=None): constructor_args = { - name: prop.from_json(json[name] if name in json else prop.default) + name: prop.from_json(json[name] if name in json else prop.default, context=context) for name, prop in cls.__properties__.items() } return cls(**constructor_args) diff --git a/dace/libraries/pblas/environments/intel_mkl_mpich.py b/dace/libraries/pblas/environments/intel_mkl_mpich.py index ab753eae7d..d29af2d758 100644 --- a/dace/libraries/pblas/environments/intel_mkl_mpich.py +++ b/dace/libraries/pblas/environments/intel_mkl_mpich.py @@ -4,7 +4,6 @@ import dace.library import ctypes.util import warnings -from typing import Union @dace.library.environment diff --git a/dace/libraries/standard/environments/__init__.py b/dace/libraries/standard/environments/__init__.py index d8f585ebd9..a47c7755f7 100644 --- a/dace/libraries/standard/environments/__init__.py +++ b/dace/libraries/standard/environments/__init__.py @@ -1,3 +1,2 @@ -# Copyright 2019-2023 ETH Zurich and the DaCe authors. All rights reserved. +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. from .cuda import CUDA -from .hptt import HPTT diff --git a/dace/libraries/standard/nodes/__init__.py b/dace/libraries/standard/nodes/__init__.py index c367a4a109..762e77760c 100644 --- a/dace/libraries/standard/nodes/__init__.py +++ b/dace/libraries/standard/nodes/__init__.py @@ -1,6 +1,4 @@ -# Copyright 2019-2023 ETH Zurich and the DaCe authors. All rights reserved. +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. from .code import CodeLibraryNode from .gearbox import Gearbox from .reduce import Reduce -from .transpose import Transpose -from .ttranspose import TensorTranspose diff --git a/dace/libraries/standard/nodes/gearbox.py b/dace/libraries/standard/nodes/gearbox.py index 32f1383996..52367c73dc 100644 --- a/dace/libraries/standard/nodes/gearbox.py +++ b/dace/libraries/standard/nodes/gearbox.py @@ -1,5 +1,4 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. -import collections import copy import dace diff --git a/dace/libraries/standard/nodes/reduce.py b/dace/libraries/standard/nodes/reduce.py index 3143773f1e..584f5a996f 100644 --- a/dace/libraries/standard/nodes/reduce.py +++ b/dace/libraries/standard/nodes/reduce.py @@ -404,7 +404,7 @@ class ExpandReduceCUDADevice(pm.ExpandTransformation): @staticmethod def expansion(node: 'Reduce', state: SDFGState, sdfg: SDFG): from dace.codegen.prettycode import CodeIOStream - from dace.codegen.targets.cpp import unparse_cr_split, cpp_array_expr + from dace.codegen.targets.cpp import unparse_cr_split node.validate(sdfg, state) input_edge: graph.MultiConnectorEdge = state.in_edges(node)[0] diff --git a/dace/libraries/stencil/cpu.py b/dace/libraries/stencil/cpu.py index 9c4eb25a5b..b8a963a582 100644 --- a/dace/libraries/stencil/cpu.py +++ b/dace/libraries/stencil/cpu.py @@ -1,7 +1,5 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. import collections -import copy -import numpy as np import dace diff --git a/dace/libraries/stencil/stencil.py b/dace/libraries/stencil/stencil.py index bff5d89e06..07ba3f1e29 100644 --- a/dace/libraries/stencil/stencil.py +++ b/dace/libraries/stencil/stencil.py @@ -1,6 +1,6 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. import collections -from typing import Dict, List, Tuple +from typing import Dict, Tuple import dace import dace.library diff --git a/dace/libraries/torch/dispatchers/cpp_torch_extension.py b/dace/libraries/torch/dispatchers/cpp_torch_extension.py index f7449df2f4..e479c3f221 100644 --- a/dace/libraries/torch/dispatchers/cpp_torch_extension.py +++ b/dace/libraries/torch/dispatchers/cpp_torch_extension.py @@ -1,14 +1,11 @@ # Copyright 2019-2025 ETH Zurich and the DaCe authors. All rights reserved. """Code generation for PyTorch C++ dispatched operators.""" import copy -import dataclasses -from distutils import sysconfig import hashlib import itertools import operator import os -import sys -from typing import List, Tuple, Callable, Optional, Dict, Union +from typing import List, Tuple, Optional, Dict import dace.library import numpy as np diff --git a/dace/memlet.py b/dace/memlet.py index c9b12f8bca..bcb854f781 100644 --- a/dace/memlet.py +++ b/dace/memlet.py @@ -4,7 +4,6 @@ from functools import reduce import operator from typing import TYPE_CHECKING, List, Optional, Set, Union -import warnings import dace from dace.sdfg.graph import generate_element_id @@ -48,7 +47,7 @@ class Memlet(object): 'value, and returns the value after resolution') # Code generation and validation hints - debuginfo = DebugInfoProperty(desc='Line information to track source and generated code') + debuginfo = DebugInfoProperty(desc='Line information to track source and generated code', allow_none=True) wcr_nonatomic = Property(dtype=bool, default=False, desc='If True, always generates non-conflicting ' @@ -176,7 +175,7 @@ def to_json(self): attrs['is_data_src'] = self._is_data_src # Fill in legacy (DEPRECATED) values for backwards compatibility - attrs['num_accesses'] = str(self.volume) if not self.dynamic else -1 + attrs['num_accesses'] = symbolic.serialize_symbolic(self.volume) if not self.dynamic else -1 return {"type": "Memlet", "attributes": attrs} diff --git a/dace/optimization/cutout_tuner.py b/dace/optimization/cutout_tuner.py index af8d6e964e..6dc232c084 100644 --- a/dace/optimization/cutout_tuner.py +++ b/dace/optimization/cutout_tuner.py @@ -1,6 +1,5 @@ # Copyright 2019-2022 ETH Zurich and the DaCe authors. All rights reserved. import os -import tempfile import math import dace import json diff --git a/dace/optimization/data_layout_tuner.py b/dace/optimization/data_layout_tuner.py index f45325f313..d9a8411dd2 100644 --- a/dace/optimization/data_layout_tuner.py +++ b/dace/optimization/data_layout_tuner.py @@ -9,7 +9,6 @@ from dace import data as dt, SDFG, dtypes from dace.optimization import cutout_tuner -from dace.sdfg.state import SDFGState from dace.transformation import helpers as xfh from dace.sdfg.analysis.cutout import SDFGCutout from dace.codegen.instrumentation.data import data_report diff --git a/dace/optimization/distributed_cutout_tuner.py b/dace/optimization/distributed_cutout_tuner.py index 39be0b7ae4..b17c94574f 100644 --- a/dace/optimization/distributed_cutout_tuner.py +++ b/dace/optimization/distributed_cutout_tuner.py @@ -1,6 +1,5 @@ # Copyright 2019-2022 ETH Zurich and the DaCe authors. All rights reserved. from collections import OrderedDict -import dace import json import itertools diff --git a/dace/optimization/map_permutation_tuner.py b/dace/optimization/map_permutation_tuner.py index 1ac288d964..1d5f8b7d3d 100644 --- a/dace/optimization/map_permutation_tuner.py +++ b/dace/optimization/map_permutation_tuner.py @@ -8,7 +8,6 @@ from dace.optimization import cutout_tuner from dace.transformation import helpers as xfh from dace.sdfg.analysis.cutout import SDFGCutout -from dace.codegen.instrumentation.data import data_report try: from tqdm import tqdm diff --git a/dace/optimization/on_the_fly_map_fusion_tuner.py b/dace/optimization/on_the_fly_map_fusion_tuner.py index 89baf0587a..8e400c5ed9 100644 --- a/dace/optimization/on_the_fly_map_fusion_tuner.py +++ b/dace/optimization/on_the_fly_map_fusion_tuner.py @@ -1,9 +1,7 @@ # Copyright 2019-2022 ETH Zurich and the DaCe authors. All rights reserved. import dace -import tempfile import math import copy -import numpy as np from typing import Generator, Dict, List, Tuple from collections import Counter diff --git a/dace/optimization/subgraph_fusion_tuner.py b/dace/optimization/subgraph_fusion_tuner.py index 01637d62da..7ae2b4db7f 100644 --- a/dace/optimization/subgraph_fusion_tuner.py +++ b/dace/optimization/subgraph_fusion_tuner.py @@ -1,6 +1,5 @@ # Copyright 2019-2022 ETH Zurich and the DaCe authors. All rights reserved. import dace -import pickle import math import copy diff --git a/dace/optimization/utils.py b/dace/optimization/utils.py index dc1e236bd1..8538c6708a 100644 --- a/dace/optimization/utils.py +++ b/dace/optimization/utils.py @@ -1,7 +1,4 @@ import os -import json -import pickle -import tempfile import math import dace import itertools @@ -9,8 +6,6 @@ from typing import Dict -from dace.codegen.instrumentation.data import data_report - def measure(sdfg, dreport=None, repetitions=30, print_report: bool = False): arguments = {} diff --git a/dace/properties.py b/dace/properties.py index 5137abc621..c8f4e2e2c6 100644 --- a/dace/properties.py +++ b/dace/properties.py @@ -12,9 +12,11 @@ import dace.subsets as sbs import dace import dace.serialize +from packaging.version import parse as parse_version +from dace import symbolic from dace.symbolic import pystr_to_symbolic from dace.dtypes import DebugInfo, typeclass -from numbers import Integral, Number +from numbers import Number from typing import List, Set, Type, Union, TypeVar, Generic, TYPE_CHECKING if TYPE_CHECKING: @@ -22,6 +24,38 @@ T = TypeVar('T') + +def _is_symbolic_converter(converter) -> bool: + """Return True for symbolic converter functions or for the SymExpr class.""" + return converter in (pystr_to_symbolic, symbolic.deserialize_symbolic, symbolic.SymExpr) + + +def _is_symbolic_type(tp) -> bool: + if tp is symbolic.SymExpr: + return True + return isinstance(tp, type) and issubclass(tp, sp.Basic) + + +def _coerce_symbolic_property_value(value): + if isinstance(value, (symbolic.SymExpr, sp.Basic)): + return value + return pystr_to_symbolic(value, simplify=False) + + +def _symbolic_deserializer(value: str, context=None) -> symbolic.SymbolicType: + """ + A backwards compatibility deserializer for symbolic properties. If the version of the + context is less than ``2.0.0a4``, it will use the old ``pystr_to_symbolic`` deserializer. + Otherwise, it will use the new ``symbolic.deserialize_symbolic`` function. + """ + version = (context or {}).get("version", None) + if version is None: + raise TypeError("Context must contain version information for symbolic deserialization") + if version is None or parse_version(version) < parse_version("2.0.0a4"): + return pystr_to_symbolic(value, simplify=False) + return symbolic.deserialize_symbolic(value) + + ############################################################################### # External interface to guarantee correct usage ############################################################################### @@ -455,7 +489,14 @@ def __init__(self, element_type: T, *args, **kwargs): self.element_type = element_type def __set__(self, obj, val): - if isinstance(val, str): + if _is_symbolic_type(self.element_type): + if isinstance(val, str): + val = self.from_string(val) + elif isinstance(val, tuple): + val = [_coerce_symbolic_property_value(v) for v in val] + elif isinstance(val, list): + val = [_coerce_symbolic_property_value(v) for v in val] + elif isinstance(val, str): val = list(map(self.element_type, list(val))) elif isinstance(val, tuple): val = list(map(self.element_type, val)) @@ -468,6 +509,10 @@ def to_string(l): def to_json(self, l): if l is None: return None + if _is_symbolic_type(self.element_type): + return [symbolic.serialize_symbolic(_coerce_symbolic_property_value(elem)) for elem in l] + if _is_symbolic_converter(self.element_type): + return [symbolic.serialize_symbolic(elem) for elem in l] # If element knows how to convert itself, let it if hasattr(self.element_type, "to_json"): return [elem.to_json() for elem in l] @@ -479,18 +524,24 @@ def to_json(self, l): def from_string(self, s): if s.startswith('[') and s.endswith(']'): + if _is_symbolic_type(self.element_type): + return [_coerce_symbolic_property_value(d.strip()) for d in s[1:-1].split(',')] return [self.element_type(d.strip()) for d in s[1:-1].split(',')] else: return list(s) - def from_json(self, data, sdfg=None): + def from_json(self, data, context=None): if data is None: return data if not isinstance(data, list): raise TypeError('ListProperty expects a list input, got %s' % data) + if _is_symbolic_type(self.element_type): + return [_symbolic_deserializer(elem, context=context) for elem in data] + if _is_symbolic_converter(self.element_type): + return [_symbolic_deserializer(elem, context=context) for elem in data] # If element knows how to convert itself, let it if hasattr(self.element_type, "from_json"): - return [self.element_type.from_json(elem) for elem in data] + return [self.element_type.from_json(elem, context) for elem in data] # Type-checks (casts) to the element type return list(map(self.element_type, data)) @@ -518,12 +569,12 @@ def to_json(self, hist): return None return [elem.to_json() if elem is not None else None for elem in hist] - def from_json(self, data, sdfg=None): + def from_json(self, data, context=None): if data is None: return data if not isinstance(data, list): raise TypeError('TransformationHistProperty expects a list input, got %s' % data) - return [dace.serialize.from_json(elem) for elem in data] + return [dace.serialize.from_json(elem, context=context) for elem in data] class DictProperty(Property): @@ -569,7 +620,9 @@ def __set__(self, obj, val): val = {k[0]: k[1] for k in val} elif isinstance(val, dict): val = { - (k if self.is_key(k) else self.key_type(k)): (v if self.is_value(v) else self.value_type(v)) + (k if self.is_key(k) else (_coerce_symbolic_property_value(k) if _is_symbolic_type(self.key_type) else self.key_type(k))): + (v if self.is_value(v) else + (_coerce_symbolic_property_value(v) if _is_symbolic_type(self.value_type) else self.value_type(v))) for k, v in val.items() } super(DictProperty, self).__set__(obj, val) @@ -584,14 +637,28 @@ def to_json(self, d): saved_dictionary = d # If key knows how to convert itself, let it - if hasattr(self.key_type, "to_json"): + if _is_symbolic_type(self.key_type): + saved_dictionary = { + symbolic.serialize_symbolic(_coerce_symbolic_property_value(k)): v + for k, v in saved_dictionary.items() + } + elif _is_symbolic_converter(self.key_type): + saved_dictionary = {symbolic.serialize_symbolic(k): v for k, v in saved_dictionary.items()} + elif hasattr(self.key_type, "to_json"): saved_dictionary = {k.to_json(): v for k, v in saved_dictionary.items()} # Otherwise, if the keys are not a native JSON type, convert to strings elif self.key_type not in (int, float, list, tuple, dict, str): saved_dictionary = {str(k): v for k, v in saved_dictionary.items()} # Same as above, but for values - if hasattr(self.value_type, "to_json"): + if _is_symbolic_type(self.value_type): + saved_dictionary = { + k: symbolic.serialize_symbolic(_coerce_symbolic_property_value(v)) + for k, v in saved_dictionary.items() + } + elif _is_symbolic_converter(self.value_type): + saved_dictionary = {k: symbolic.serialize_symbolic(v) for k, v in saved_dictionary.items()} + elif hasattr(self.value_type, "to_json"): saved_dictionary = {k: v.to_json() for k, v in saved_dictionary.items()} elif self.value_type not in (int, float, list, tuple, dict, str): saved_dictionary = {k: str(v) for k, v in saved_dictionary.items()} @@ -605,7 +672,7 @@ def to_json(self, d): def from_string(s): return dict(s) - def from_json(self, data, sdfg=None): + def from_json(self, data, context=None): if data is None: return data if not isinstance(data, dict): @@ -615,11 +682,21 @@ def from_json(self, data, sdfg=None): key_json = hasattr(self.key_type, "from_json") value_json = hasattr(self.value_type, "from_json") - return { - self.key_type.from_json(k, sdfg) if key_json else self.key_type(k): - self.value_type.from_json(v, sdfg) if value_json else self.value_type(v) - for k, v in data.items() - } + def _convert_key(key): + if _is_symbolic_type(self.key_type): + return _symbolic_deserializer(key, context) + if _is_symbolic_converter(self.key_type): + return _symbolic_deserializer(key, context) + return self.key_type.from_json(key, context) if key_json else self.key_type(key) + + def _convert_value(value): + if _is_symbolic_type(self.value_type): + return _symbolic_deserializer(value, context) + if _is_symbolic_converter(self.value_type): + return _symbolic_deserializer(value, context) + return self.value_type.from_json(value, context) if value_json else self.value_type(value) + + return {_convert_key(k): _convert_value(v) for k, v in data.items()} ############################################################################### @@ -719,7 +796,10 @@ class DebugInfoProperty(Property): def __init__(self, **kwargs): if 'default' not in kwargs: - kwargs['default'] = DebugInfo(0, 0, 0, 0) + if kwargs.get('allow_none', False): + kwargs['default'] = None + else: + kwargs['default'] = DebugInfo(None) super().__init__(dtype=DebugInfo, **kwargs) @property @@ -988,7 +1068,7 @@ def to_json(self): return ret @staticmethod - def from_json(tmp, sdfg=None): + def from_json(tmp, context=None): if tmp is None: return None if isinstance(tmp, CodeBlock): @@ -1139,8 +1219,8 @@ def to_json(self, val): except AttributeError: return SubsetProperty.to_string(val) - def from_json(self, val, sdfg=None): - return dace.serialize.from_json(val) + def from_json(self, val, context=None): + return dace.serialize.from_json(val, context) class SymbolicProperty(Property): @@ -1168,6 +1248,16 @@ def to_string(obj): # Go through sympy once to reorder factors return str(pystr_to_symbolic(str(obj), simplify=False)) + def to_json(self, val): + if val is None: + return None + return symbolic.serialize_symbolic(val) + + def from_json(self, val, context=None): + if val is None: + return None + return _symbolic_deserializer(val, context=context) + class DataProperty(Property): """ Custom Property type that represents a link to a data descriptor. @@ -1265,12 +1355,12 @@ def to_string(obj): def to_json(self, obj): if obj is None: return None - return list(map(str, obj)) + return [symbolic.serialize_symbolic(o) for o in obj] - def from_json(self, d, sdfg=None): + def from_json(self, d, context=None): if d is None: return None - return tuple([dace.symbolic.pystr_to_symbolic(m) for m in d]) + return tuple([_symbolic_deserializer(m, context=context) for m in d]) def __set__(self, obj, val): if isinstance(val, list): @@ -1344,7 +1434,7 @@ def from_json(obj, context=None): return TypeClassProperty.from_string(obj) elif isinstance(obj, dict): # Let the deserializer handle this - return dace.serialize.from_json(obj) + return dace.serialize.from_json(obj, context=context) else: raise TypeError("Cannot parse type from: {}".format(obj)) @@ -1385,7 +1475,7 @@ def from_json(obj, context=None): return NestedDataClassProperty.from_string(obj) elif isinstance(obj, dict): # Let the deserializer handle this - return dace.serialize.from_json(obj) + return dace.serialize.from_json(obj, context=context) else: raise TypeError("Cannot parse type from: {}".format(obj)) @@ -1418,7 +1508,7 @@ def to_json(self, obj): return None return obj.dict() - def from_json(self, d, sdfg=None): + def from_json(self, d, context=None): if d is None: return None return self.dtype.parse_obj(d) diff --git a/dace/registry.py b/dace/registry.py index 07079be865..96761c1c95 100644 --- a/dace/registry.py +++ b/dace/registry.py @@ -3,8 +3,7 @@ subclasses and values can be registered externally. """ from dace import attr_enum -from enum import Enum -from typing import Dict, Type, TypeVar, TYPE_CHECKING +from typing import Dict, Type, TypeVar T = TypeVar('T') E = TypeVar('E', bound=attr_enum.ExtensibleAttributeEnum) diff --git a/dace/libraries/linalg/include/dace_cutensor.h b/dace/runtime/include/dace/dace_cutensor.h similarity index 57% rename from dace/libraries/linalg/include/dace_cutensor.h rename to dace/runtime/include/dace/dace_cutensor.h index ddad2feaa3..ef4f2cb5d8 100644 --- a/dace/libraries/linalg/include/dace_cutensor.h +++ b/dace/runtime/include/dace/dace_cutensor.h @@ -1,21 +1,21 @@ -// Copyright 2019-2023 ETH Zurich and the DaCe authors. All rights reserved. +// Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. #pragma once #include #include -#include // size_t -#include // std::runtime_error -#include // std::to_string +#include +#include +#include #include namespace dace { - namespace linalg { static void CheckCuTensorError(cutensorStatus_t const& status) { if (status != CUTENSOR_STATUS_SUCCESS) { - throw std::runtime_error("cuTENSOR failed with error code: " + std::string(cutensorGetErrorString(status))); + throw std::runtime_error("cuTENSOR failed with error code: " + + std::string(cutensorGetErrorString(status))); } } @@ -26,15 +26,10 @@ static cutensorHandle_t CreateCuTensorHandle(int device) { } } cutensorHandle_t handle; - CheckCuTensorError(cutensorInit(&handle)); + CheckCuTensorError(cutensorCreate(&handle)); return handle; } -/** - * cuTENSOR wrapper class for DaCe. Once constructed, the class can be used to - * get or create a cuTENSOR library handle (cutensorHandle_t) for a given - * GPU ID. The class is constructed when the cuTENSOR DaCe library is used. - **/ class CuTensorHandle { public: CuTensorHandle() = default; @@ -43,8 +38,8 @@ class CuTensorHandle { cutensorHandle_t& Get(int device) { auto f = handles_.find(device); if (f == handles_.end()) { - // Lazily construct new cuSolverDn handle if the specified key does not - // yet exist + // Lazily construct new cuTENSOR handle if the specified key does + // not yet exist auto handle = CreateCuTensorHandle(device); f = handles_.emplace(device, handle).first; } @@ -52,10 +47,9 @@ class CuTensorHandle { } ~CuTensorHandle() { - // NOTE: It seems that the cuTENSOR API is missing a method of destroying a cuTENSOR handle - // for (auto& h : handles_) { - // CheckCuTensorError(cutensorDestroy(h.second)); - // } + for (auto& h : handles_) { + cutensorDestroy(h.second); + } } CuTensorHandle& operator=(CuTensorHandle const&) = delete; @@ -64,5 +58,4 @@ class CuTensorHandle { }; } // namespace linalg - } // namespace dace diff --git a/dace/sdfg/analysis/cutout.py b/dace/sdfg/analysis/cutout.py index db2cddee87..6079b9ec2c 100644 --- a/dace/sdfg/analysis/cutout.py +++ b/dace/sdfg/analysis/cutout.py @@ -109,8 +109,8 @@ def translate_transformation_into(self, transformation: Union[PatternTransformat pass transformation.subgraph = new_subgraph - def to_json(self, hash=False): - cutout_json = super().to_json(hash) + def to_json(self, hash=False, include_transformation_history=False): + cutout_json = super().to_json(hash, include_transformation_history) cutout_json['type'] = SDFG.__name__ return cutout_json diff --git a/dace/sdfg/analysis/schedule_tree/sdfg_to_tree.py b/dace/sdfg/analysis/schedule_tree/sdfg_to_tree.py index 46eb37cdb2..e1c02aece0 100644 --- a/dace/sdfg/analysis/schedule_tree/sdfg_to_tree.py +++ b/dace/sdfg/analysis/schedule_tree/sdfg_to_tree.py @@ -1,4 +1,4 @@ -# Copyright 2019-2025 ETH Zurich and the DaCe authors. All rights reserved. +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. from collections import defaultdict import copy from typing import Dict, List, Set @@ -26,7 +26,7 @@ } -def dealias_sdfg(sdfg: SDFG): +def _dealias_sdfg(sdfg: SDFG) -> None: """ Renames all data containers in an SDFG tree (i.e., nested SDFGs) to use the same data descriptors as the top-level SDFG. This function takes care of offsetting memlets and internal @@ -43,7 +43,8 @@ def dealias_sdfg(sdfg: SDFG): replacements: Dict[str, str] = {} inv_replacements: Dict[str, List[str]] = {} - parent_edges: Dict[str, Memlet] = {} + parent_edges_inputs: Dict[str, gr.MultiConnectorEdge[Memlet]] = {} + parent_edges_outputs: Dict[str, gr.MultiConnectorEdge[Memlet]] = {} to_unsqueeze: Set[str] = set() parent_sdfg = nsdfg.parent_sdfg @@ -53,19 +54,42 @@ def dealias_sdfg(sdfg: SDFG): for name, desc in nsdfg.arrays.items(): if desc.transient: continue - for edge in parent_state.edges_by_connector(parent_node, name): + for edge in parent_state.in_edges_by_connector(parent_node, name): parent_name = edge.data.data assert parent_name in parent_sdfg.arrays - if parent_name != name: + if name != parent_name: + parent_edges_inputs[name] = edge replacements[name] = parent_name - parent_edges[name] = edge if parent_name in inv_replacements: inv_replacements[parent_name].append(name) to_unsqueeze.add(parent_name) else: inv_replacements[parent_name] = [name] + # We found an incoming edge for name and we don't expect a second one. break + for edge in parent_state.out_edges_by_connector(parent_node, name): + parent_name = edge.data.data + assert parent_name in parent_sdfg.arrays + if name != parent_name: + parent_edges_outputs[name] = edge + + if replacements.get(name, None) is not None: + # There's an incoming and an outgoing connector with the same name. + # Make sure both map to the same memory in the parent sdfg. + assert replacements[name] == parent_name + assert name in inv_replacements[parent_name] + break + else: + replacements[name] = parent_name + if parent_name in inv_replacements: + inv_replacements[parent_name].append(name) + to_unsqueeze.add(parent_name) + else: + inv_replacements[parent_name] = [name] + # We found an outgoing edge for name and we don't expect a second one. + break + if to_unsqueeze: for parent_name in to_unsqueeze: parent_arr = parent_sdfg.arrays[parent_name] @@ -94,14 +118,18 @@ def dealias_sdfg(sdfg: SDFG): # destination subset if isinstance(src, nd.AccessNode) and src.data in child_names: src_data = src.data - new_src_memlet = unsqueeze_memlet(e.data, parent_edges[src.data].data, use_src_subset=True) + new_src_memlet = unsqueeze_memlet(e.data, + parent_edges_inputs[src.data].data, + use_src_subset=True) else: src_data = None new_src_memlet = None # We need to take directionality of the memlet into account if isinstance(dst, nd.AccessNode) and dst.data in child_names: dst_data = dst.data - new_dst_memlet = unsqueeze_memlet(e.data, parent_edges[dst.data].data, use_dst_subset=True) + new_dst_memlet = unsqueeze_memlet(e.data, + parent_edges_outputs[dst.data].data, + use_dst_subset=True) else: dst_data = None new_dst_memlet = None @@ -120,23 +148,26 @@ def dealias_sdfg(sdfg: SDFG): syms = e.data.read_symbols() for memlet in e.data.get_read_memlets(nsdfg.arrays): if memlet.data in child_names: - repl_dict[str(memlet)] = unsqueeze_memlet(memlet, parent_edges[memlet.data].data) + repl_dict[str(memlet)] = unsqueeze_memlet(memlet, parent_edges_inputs[memlet.data].data) if memlet.data in syms: syms.remove(memlet.data) for s in syms: - if s in parent_edges: + if s in parent_edges_inputs: if s in nsdfg.arrays: - repl_dict[s] = parent_edges[s].data.data + repl_dict[s] = parent_edges_inputs[s].data.data else: - repl_dict[s] = str(parent_edges[s].data) + repl_dict[s] = str(parent_edges_inputs[s].data) e.data.replace_dict(repl_dict) for name in child_names: - edge = parent_edges[name] - for e in parent_state.memlet_tree(edge): - if e.data.data == parent_name: - e.data.subset = subsets.Range.from_array(parent_arr) - else: - e.data.other_subset = subsets.Range.from_array(parent_arr) + for edge in [parent_edges_inputs.get(name, None), parent_edges_outputs.get(name, None)]: + if edge is None: + continue + + for e in parent_state.memlet_tree(edge): + if e.data.data == parent_name: + e.data.subset = subsets.Range.from_array(parent_arr) + else: + e.data.other_subset = subsets.Range.from_array(parent_arr) if replacements: struct_outside_replacements: Dict[str, str] = {} @@ -166,7 +197,7 @@ def dealias_sdfg(sdfg: SDFG): e._dst_conn = replacements[e.dst_conn] -def normalize_memlet(sdfg: SDFG, state: SDFGState, original: gr.MultiConnectorEdge[Memlet], data: str) -> Memlet: +def _normalize_memlet(sdfg: SDFG, state: SDFGState, original: gr.MultiConnectorEdge[Memlet], data: str) -> Memlet: """ Normalizes a memlet to a given data descriptor. @@ -199,9 +230,10 @@ def normalize_memlet(sdfg: SDFG, state: SDFGState, original: gr.MultiConnectorEd return memlet -def replace_memlets(sdfg: SDFG, input_mapping: Dict[str, Memlet], output_mapping: Dict[str, Memlet]): +def _replace_memlets(sdfg: SDFG, input_mapping: Dict[str, Memlet], output_mapping: Dict[str, Memlet]): """ Replaces all uses of data containers in memlets and interstate edges in an SDFG. + :param sdfg: The SDFG. :param input_mapping: A mapping from internal data descriptor names to external input memlets. :param output_mapping: A mapping from internal data descriptor names to external output memlets. @@ -277,11 +309,13 @@ def replace_memlets(sdfg: SDFG, input_mapping: Dict[str, Memlet], output_mapping e.data.condition.as_string = condstr.replace(find, replace) -def remove_name_collisions(sdfg: SDFG): +def _remove_name_collisions(sdfg: SDFG) -> None: """ Removes name collisions in nested SDFGs by renaming states, data containers, and symbols. - :param sdfg: The SDFG. + This function works in-place on the given SDFG. + + :param sdfg: The SDFG to operate on. """ state_names_seen = set() identifiers_seen = set() @@ -343,7 +377,7 @@ def _make_view_node(state: SDFGState, edge: gr.MultiConnectorEdge[Memlet], view_ Helper function to create a view schedule tree node from a memlet edge. """ sdfg = state.parent - normalized = normalize_memlet(sdfg, state, edge, viewed_name) + normalized = _normalize_memlet(sdfg, state, edge, viewed_name) view_node = tn.ViewNode(target=view_name, source=viewed_name, memlet=normalized, @@ -352,7 +386,7 @@ def _make_view_node(state: SDFGState, edge: gr.MultiConnectorEdge[Memlet], view_ return view_node -def replace_symbols_until_set(nsdfg: dace.nodes.NestedSDFG): +def _replace_symbols_until_set(nsdfg: dace.nodes.NestedSDFG) -> None: """ Replaces symbol values in a nested SDFG until their value has been reset. This is used for matching symbol namespaces between an SDFG and a nested SDFG. @@ -377,7 +411,10 @@ def replace_symbols_until_set(nsdfg: dace.nodes.NestedSDFG): symbolic.safe_replace(per_state_mapping, lambda d: e.data.replace_dict(d, replace_keys=False)) -def prepare_schedule_tree_edges(state: SDFGState) -> Dict[gr.MultiConnectorEdge[Memlet], tn.ScheduleTreeNode]: +def _prepare_schedule_tree_edges( + state: SDFGState +) -> tuple[Dict[gr.MultiConnectorEdge[Memlet], tn.ScheduleTreeNode], Dict[nd.EntryNode, + List[gr.MultiConnectorEdge[Memlet]]]]: """ Creates a dictionary mapping edges to their corresponding schedule tree nodes, if relevant. This handles view edges, reference sets, and dynamic map inputs. @@ -472,7 +509,7 @@ def prepare_schedule_tree_edges(state: SDFGState) -> Dict[gr.MultiConnectorEdge[ result[e] = tn.DynScopeCopyNode(target=e.dst_conn, memlet=e.data) else: target_name = innermost_node.data - new_memlet = normalize_memlet(sdfg, state, e, outermost_node.data) + new_memlet = _normalize_memlet(sdfg, state, e, outermost_node.data) result[e] = tn.CopyNode(target=target_name, memlet=new_memlet) scope = state.entry_node(e.dst if mtree.downwards else e.src) @@ -493,7 +530,7 @@ def _state_schedule_tree(state: SDFGState) -> List[tn.ScheduleTreeNode]: edge_to_stree: Dict[gr.MultiConnectorEdge[Memlet], tn.ScheduleTreeNode] scope_to_edges: Dict[nd.EntryNode, List[gr.MultiConnectorEdge[Memlet]]] - edge_to_stree, scope_to_edges = prepare_schedule_tree_edges(state) + edge_to_stree, scope_to_edges = _prepare_schedule_tree_edges(state) edges_to_ignore = set() # Handle all unscoped edges to generate output views @@ -523,15 +560,19 @@ def _state_schedule_tree(state: SDFGState) -> List[tn.ScheduleTreeNode]: result = subnodes elif isinstance(node, dace.nodes.ExitNode): result = scopes.pop() + parent = result[-1] + assert isinstance(parent, tn.ScheduleTreeScope) + for child in parent.children: + child.parent = parent elif isinstance(node, dace.nodes.NestedSDFG): nested_array_mapping_input = {} nested_array_mapping_output = {} generated_nviews = set() # Replace symbols and memlets in nested SDFGs to match the namespace of the parent SDFG - replace_symbols_until_set(node) + _replace_symbols_until_set(node) - # Create memlets for nested SDFG mapping, or nview schedule nodes if slice cannot be determined + # Create memlets for nested SDFG mapping, or NView schedule nodes if slice cannot be determined. for e in state.all_edges(node): conn = e.dst_conn if e.dst is node else e.src_conn if e.data.is_empty() or not conn: @@ -550,7 +591,7 @@ def _state_schedule_tree(state: SDFGState) -> List[tn.ScheduleTreeNode]: else: nested_array_mapping_output[conn] = e.data - if no_mapping: # Must use view (nview = nested SDFG view) + if no_mapping: # Must use view (NView = nested SDFG view) if conn not in generated_nviews: nview_node = tn.NView(target=conn, source=e.data.data, @@ -560,11 +601,17 @@ def _state_schedule_tree(state: SDFGState) -> List[tn.ScheduleTreeNode]: result.append(nview_node) generated_nviews.add(conn) - replace_memlets(node.sdfg, nested_array_mapping_input, nested_array_mapping_output) + _replace_memlets(node.sdfg, nested_array_mapping_input, nested_array_mapping_output) # Insert the nested SDFG flattened nested_stree = as_schedule_tree(node.sdfg, in_place=True, toplevel=False) result.extend(nested_stree.children) + + if generated_nviews: + # Insert matching NViewEnd nodes to define the scope of NView nodes. + for target in generated_nviews: + result.append(tn.NViewEnd(target=target)) + elif isinstance(node, dace.nodes.Tasklet): in_memlets = {e.dst_conn: e.data for e in state.in_edges(node) if e.dst_conn} out_memlets = {e.src_conn: e.data for e in state.out_edges(node) if e.src_conn} @@ -599,8 +646,7 @@ def _state_schedule_tree(state: SDFGState) -> List[tn.ScheduleTreeNode]: return result -def _isedge_schedule_tree(graph: ControlFlowRegion, - edge: gr.Edge[InterstateEdge], +def _isedge_schedule_tree(edge: gr.Edge[InterstateEdge], emit_goto_for_successors: bool = False) -> List[tn.ScheduleTreeNode]: result: List[tn.ScheduleTreeNode] = [] for aname, aval in edge.data.assignments.items(): @@ -647,7 +693,7 @@ def _block_schedule_tree(block: ControlFlowBlock) -> List[tn.ScheduleTreeNode]: subnodes.extend(_block_schedule_tree(n)) for oe in block.out_edges(n): if oe not in processed_edges: - subnodes.extend(_isedge_schedule_tree(block, oe, emit_goto_for_successors=True)) + subnodes.extend(_isedge_schedule_tree(oe, emit_goto_for_successors=True)) processed_edges.add(oe) gblock = tn.GBlock(children=subnodes) children = [gblock] @@ -659,16 +705,28 @@ def _block_schedule_tree(block: ControlFlowBlock) -> List[tn.ScheduleTreeNode]: oedges = block.out_edges(pivot) if len(oedges) == 1: pivot = oedges[0].dst - children.extend(_isedge_schedule_tree(block, oedges[0], emit_goto_for_successors=False)) + children.extend(_isedge_schedule_tree(oedges[0], emit_goto_for_successors=False)) else: pivot = None if isinstance(block, LoopRegion): - # If this is a loop region, wrap everything in a LoopScope node. - loop_node = tn.LoopScope(loop=block, children=children) - return [loop_node] + # If this is a loop region, wrap everything in a loop scope node. + variant = tn.loop_variant(block) + if variant == "for": + return [tn.ForScope(loop=block, children=children)] + + if variant == "while": + return [tn.WhileScope(loop=block, children=children)] + + if variant == "do-while": + return [tn.DoWhileScope(loop=block, children=children)] + + # If we end up here, we don't need more granularity and just use a general loop scope. + return [tn.LoopScope(loop=block, children=children)] + return children - elif isinstance(block, ConditionalBlock): + + if isinstance(block, ConditionalBlock): result: List[tn.ScheduleTreeNode] = [] if_node = tn.IfScope(condition=block.branches[0][0], children=_block_schedule_tree(block.branches[0][1])) result.append(if_node) @@ -680,9 +738,11 @@ def _block_schedule_tree(block: ControlFlowBlock) -> List[tn.ScheduleTreeNode]: else_node = tn.ElseScope(children=_block_schedule_tree(branch_body)) result.append(else_node) return result - elif isinstance(block, SDFGState): + + if isinstance(block, SDFGState): return _state_schedule_tree(block) - elif isinstance(block, ReturnBlock): + + if isinstance(block, ReturnBlock): # For return blocks, add a goto node to the end of the schedule tree. # NOTE: Return blocks currently always exit the entire SDFG context they are contained in, meaning that the exit # goto has target=None. However, in the future we want to adapt Return blocks to be able to return only a @@ -690,8 +750,8 @@ def _block_schedule_tree(block: ControlFlowBlock) -> List[tn.ScheduleTreeNode]: # entire SDFG. goto_node = tn.GotoNode(target=None) return [goto_node] - else: - raise tn.UnsupportedScopeException(type(block).__name__) + + raise tn.UnsupportedScopeException(type(block).__name__) def _generate_views_in_scope( @@ -725,7 +785,47 @@ def _generate_views_in_scope( return result -def as_schedule_tree(sdfg: SDFG, in_place: bool = False, toplevel: bool = True) -> tn.ScheduleTreeScope: +def _prepare_sdfg_for_conversion(sdfg: SDFG, *, toplevel: bool) -> None: + from dace.transformation import helpers as xfh # Avoid import loop + + # Split edges with assignments and conditions + xfh.split_interstate_edges(sdfg) + + # Replace code->code edges with data<->code edges + xfh.replace_code_to_code_edges(sdfg) + + if toplevel: # Top-level SDFG preparation (only perform once) + # Handle name collisions (in arrays, state labels, symbols) + _remove_name_collisions(sdfg) + + # Ensure no arrays alias in SDFG tree + _dealias_sdfg(sdfg) + + +def _create_unified_descriptor_repository(sdfg: SDFG, stree: tn.ScheduleTreeRoot) -> None: + """ + Creates a single descriptor repository from an SDFG and all nested SDFGs. This includes + data containers, symbols, constants, etc. + + :param sdfg: The top-level SDFG to create the repository from. + :param stree: The tree root in which to make the unified descriptor repository. + """ + stree.containers = sdfg.arrays + stree.symbols = sdfg.symbols + stree.constants = sdfg.constants_prop + + # Since the SDFG is assumed to be de-aliased and contain unique names, we union the contents of + # the nested SDFGs' descriptor repositories + for nsdfg in sdfg.all_sdfgs_recursive(): + transients = {k: v for k, v in nsdfg.arrays.items() if v.transient} + symbols = {k: v for k, v in nsdfg.symbols.items() if k not in stree.symbols} + constants = {k: v for k, v in nsdfg.constants_prop.items() if k not in stree.constants} + stree.containers.update(transients) + stree.symbols.update(symbols) + stree.constants.update(constants) + + +def as_schedule_tree(sdfg: SDFG, *, in_place: bool = False, toplevel: bool = True) -> tn.ScheduleTreeRoot: """ Converts an SDFG into a schedule tree. The schedule tree is a tree of nodes that represent the execution order of the SDFG. @@ -741,30 +841,21 @@ def as_schedule_tree(sdfg: SDFG, in_place: bool = False, toplevel: bool = True) usable after the conversion if ``in_place`` is True! :return: A schedule tree representing the given SDFG. """ - from dace.transformation import helpers as xfh # Avoid import loop if not in_place: sdfg = copy.deepcopy(sdfg) - # Prepare SDFG for conversion - ############################# - - # Split edges with assignments and conditions - xfh.split_interstate_edges(sdfg) - - # Replace code->code edges with data<->code edges - xfh.replace_code_to_code_edges(sdfg) + _prepare_sdfg_for_conversion(sdfg, toplevel=toplevel) - if toplevel: # Top-level SDFG preparation (only perform once) - # Handle name collisions (in arrays, state labels, symbols) - remove_name_collisions(sdfg) - # Ensure no arrays alias in SDFG tree - dealias_sdfg(sdfg) - - ############################# + if toplevel: + result = tn.ScheduleTreeRoot(name=sdfg.name, children=[]) + _create_unified_descriptor_repository(sdfg, result) + result.add_children(_block_schedule_tree(sdfg)) + else: + result = tn.ScheduleTreeScope(children=_block_schedule_tree(sdfg)) - result = tn.ScheduleTreeScope(children=_block_schedule_tree(sdfg)) tn.validate_has_no_other_node_types(result) + tn.validate_children_and_parents_align(result, root=toplevel) # Clean up tree stpasses.remove_unused_and_duplicate_labels(result) diff --git a/dace/sdfg/analysis/schedule_tree/tree_to_sdfg.py b/dace/sdfg/analysis/schedule_tree/tree_to_sdfg.py new file mode 100644 index 0000000000..f8bfab9dde --- /dev/null +++ b/dace/sdfg/analysis/schedule_tree/tree_to_sdfg.py @@ -0,0 +1,1052 @@ +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. +import copy + +from collections import defaultdict +from dataclasses import dataclass +from enum import Enum, auto +from types import TracebackType +from typing import Final + +from dace import symbolic +from dace.memlet import Memlet +from dace.sdfg import nodes, memlet_utils as mmu +from dace.sdfg.sdfg import SDFG, ControlFlowRegion, InterstateEdge +from dace.sdfg.state import ConditionalBlock, ControlFlowBlock, SDFGState, LoopRegion +from dace.sdfg.analysis.schedule_tree import treenodes as tn +from dace.sdfg import propagation + + +class StateBoundaryBehavior(Enum): + STATE_TRANSITION = auto() #: Creates multiple states with a state transition + EMPTY_MEMLET = auto() #: Happens-before empty memlet edges in the same state + + +PREFIX_PASSTHROUGH_IN: Final[str] = "IN_" +PREFIX_PASSTHROUGH_OUT: Final[str] = "OUT_" + + +@dataclass +class _Context: + """Context information for transforming a schedule tree into an SDFG.""" + + root: tn.ScheduleTreeRoot + current_scope: tn.ScheduleTreeScope | None + + access_cache: dict[tuple[SDFGState, int], dict[str, nodes.AccessNode]] + """Per scope (hashed by id(scope_node) access_cache.""" + + +class _TreeScope: + """Automatically set the current scope on the context to the given node.""" + + def __init__(self, node: tn.ScheduleTreeScope, ctx: _Context, state: SDFGState) -> None: + if ctx.current_scope is None and not isinstance(node, tn.ScheduleTreeRoot): + raise ValueError("ctx.current_scope is only allowed to be 'None' when node it tree root.") + + self._ctx = ctx + self._parent_scope = ctx.current_scope + self._node = node + self._state = state + + cache_key = (state, id(node)) + assert cache_key not in self._ctx.access_cache + self._ctx.access_cache[cache_key] = {} + + def __enter__(self) -> None: + assert not self._ctx.access_cache[(self._state, id( + self._node))], "Expecting an empty access_cache when entering the context." + + self._ctx.current_scope = self._node + + def __exit__(self, exc_type: type[BaseException] | None, exc_val: BaseException | None, + exc_tb: TracebackType | None) -> None: + cache_key = (self._state, id(self._node)) + assert cache_key in self._ctx.access_cache + + self._ctx.current_scope = self._parent_scope + + +class _StreeToSDFG(tn.ScheduleNodeVisitor): + + def __init__( + self, + *, + start_state: SDFGState | None = None, + boundary_behavior: StateBoundaryBehavior = StateBoundaryBehavior.STATE_TRANSITION, + max_nested_sdfg: int = 1000, + ) -> None: + if boundary_behavior != StateBoundaryBehavior.STATE_TRANSITION: + raise NotImplementedError("Only STATE_TRANSITION is currently supported as StateBoundaryBehavior.") + + self._ctx: _Context + """Context information like tree root and current scope.""" + + self._current_state = start_state + """Current SDFGState in the SDFG that we are building.""" + + self._current_nestedSDFG: int | None = None + """Id of the current nested SDFG if we are inside one.""" + + self._interstate_symbols: list[tn.AssignNode] = [] + """Interstate symbol assignments. Will be assigned with the next state transition.""" + + self._nviews_free: list[tn.NView] = [] + """Keep track of NView (nested SDFG view) nodes that are "free" to be used.""" + + self._nviews_bound_per_scope: dict[int, list[tn.NView]] = {} + """Mapping of id(SDFG) -> list of active NView nodes in that SDFG.""" + + self._nviews_deferred_removal: dict[int, list[tn.NView]] = {} + """"Mapping of id(SDFG) -> list of NView nodes to be removed once we exit this nested SDFG.""" + + # state management + self._state_stack: list[SDFGState] = [] + + # dataflow scopes + # list[ (MapEntryNode, ToConnect) | (SDFG, {"inputs": set(), "outputs": set()}) ] + self._dataflow_stack: list[tuple[nodes.EntryNode, dict[str, tuple[nodes.AccessNode, Memlet]]] + | tuple[SDFG, dict[str, set[str]]]] = [] + + self._max_nested_sdfg = max_nested_sdfg + + def _apply_nview_array_override(self, array_name: str, sdfg: SDFG) -> bool: + """ + Apply an NView override if applicable. Returns true if the NView was applied. + + See `visit_NView()` for how we keep track of nested SDFG view nodes. + """ + length = len(self._nviews_free) + for index, nview in enumerate(reversed(self._nviews_free), start=1): + if nview.target == array_name and nview not in self._nviews_deferred_removal[id(sdfg)]: + # Add the "override" data descriptor + sdfg.add_datadesc(nview.target, nview.view_desc.clone()) + if nview.src_desc.transient: + sdfg.arrays[nview.target].transient = False + + # Keep track of used NViews per scope (to "free" them again once the scope ends) + self._nviews_bound_per_scope[id(sdfg)].append(nview) + + # This NView is in use now, remove it from the free NViews. + del self._nviews_free[length - index] + return True + + return False + + def _parent_sdfg_with_array(self, name: str, sdfg: SDFG) -> SDFG: + """Find the closest parent SDFG containing an array with the given name.""" + parent_sdfg = sdfg.parent.sdfg + sdfg_counter = 1 + while name not in parent_sdfg.arrays and sdfg_counter < self._max_nested_sdfg: + parent_sdfg = parent_sdfg.parent.sdfg + assert isinstance(parent_sdfg, SDFG) + sdfg_counter += 1 + assert sdfg_counter < self._max_nested_sdfg, f"Array '{name}' not found in any parent of SDFG '{sdfg.name}'." + return parent_sdfg + + def _pop_state(self, label: str | None = None) -> SDFGState: + """Pops the last state from the state stack. + + :param str, optional label: Ensures the popped state's label starts with the given string. + + :return: The popped state. + """ + if not self._state_stack: + raise ValueError("Can't pop state from empty stack.") + + popped = self._state_stack.pop() + if label is not None: + assert popped.label.startswith(label) + + return popped + + def visit_ScheduleTreeRoot(self, node: tn.ScheduleTreeRoot, sdfg: SDFG) -> None: + assert self._current_state is None, "Expected no 'current_state' at root." + assert not self._state_stack, "Expected empty state stack at root." + assert not self._dataflow_stack, "Expected empty dataflow stack at root." + assert not self._interstate_symbols, "Expected empty list of symbols at root." + + self._current_state = sdfg.add_state(label="tree_root", is_start_block=True) + self._ctx = _Context(root=node, access_cache={}, current_scope=None) + with _TreeScope(node, self._ctx, self._current_state): + self.visit(node.children, sdfg=sdfg) + + assert not self._state_stack, "Expected empty state stack." + assert not self._dataflow_stack, "Expected empty dataflow stack." + assert not self._interstate_symbols, "Expected empty list of symbols to add." + + def visit_GBlock(self, node: tn.GBlock, sdfg: SDFG) -> None: + raise NotImplementedError(f"Support for {type(node)} not yet implemented.") + + def visit_StateLabel(self, node: tn.StateLabel, sdfg: SDFG) -> None: + raise NotImplementedError(f"Support for {type(node)} not yet implemented.") + + def visit_GotoNode(self, node: tn.GotoNode, sdfg: SDFG) -> None: + raise NotImplementedError(f"Support for{type(node)} not yet implemented.") + + def visit_AssignNode(self, node: tn.AssignNode, sdfg: SDFG) -> None: + # We just collect them here. They'll be added when state boundaries are added, + # see visitors below. + self._interstate_symbols.append(node) + + # If AssignNode depends on arrays, e.g. `my_sym = my_array[__k] > 0`, make sure array accesses can be resolved. + input_memlets = node.input_memlets() + if not input_memlets: + return + + for entry in reversed(self._dataflow_stack): + scope_node, to_connect = entry + if isinstance(scope_node, SDFG): + # In case we are inside a nested SDFG, make sure memlet data can be + # resolved by explicitly adding inputs. + for memlet in input_memlets: + # Copy data descriptor from parent SDFG and add input connector + if memlet.data not in sdfg.arrays: + parent_sdfg = self._parent_sdfg_with_array(memlet.data, sdfg) + + # Support for NView nodes + use_nview = self._apply_nview_array_override(memlet.data, sdfg) + if not use_nview: + sdfg.add_datadesc(memlet.data, parent_sdfg.arrays[memlet.data].clone()) + + # Transients passed into a nested SDFG become non-transient inside that nested SDFG + if parent_sdfg.arrays[memlet.data].transient: + sdfg.arrays[memlet.data].transient = False + + # Dev note: nview.target and memlet.data are identical + assert memlet.data not in to_connect["inputs"] + to_connect["inputs"].add(memlet.data) + return + + for memlet in input_memlets: + # If we aren't inside a nested SDFG, make sure all memlets can be resolved. + # Imo, this should always be the case. It not, raise an error. + if memlet.data not in sdfg.arrays: + raise ValueError(f"Parsing AssignNode {node} failed. Can't find {memlet.data} in {sdfg}.") + + def visit_ForScope(self, node: tn.ForScope, sdfg: SDFG) -> None: + current_state = self._current_state + assert current_state is not None + cf_region = current_state.parent_graph + + loop_region = LoopRegion(label=node.loop.label, + condition_expr=node.loop.loop_condition, + loop_var=node.loop.loop_variable, + initialize_expr=node.loop.init_statement, + update_expr=node.loop.update_statement, + unroll=node.loop.unroll, + unroll_factor=node.loop.unroll_factor) + cf_region.add_node(loop_region) + loop_state = loop_region.add_state(f"for_loop_state_{id(node)}", is_start_block=True) + + _insert_and_split_assignments(current_state, loop_region) + + self._current_state = loop_state + self.visit(node.children, sdfg=sdfg) + + after_state = _insert_and_split_assignments(loop_region, label="loop_after") + self._current_state = after_state + + def visit_WhileScope(self, node: tn.WhileScope, sdfg: SDFG) -> None: + current_state = self._current_state + assert current_state is not None + cf_region = current_state.parent_graph + + loop_region = node.loop + cf_region.add_node(loop_region) + loop_state = loop_region.add_state(f"while_loop_state_{id(node)}", is_start_block=True) + + _insert_and_split_assignments(current_state, loop_region) + + self._current_state = loop_state + self.visit(node.children, sdfg=sdfg) + + after_state = _insert_and_split_assignments(loop_region, label="loop_after") + self._current_state = after_state + + def visit_DoWhileScope(self, node: tn.DoWhileScope, sdfg: SDFG) -> None: + raise NotImplementedError(f"Support for {type(node)} not yet implemented.") + + def visit_LoopScope(self, node: tn.LoopScope, sdfg: SDFG) -> None: + raise NotImplementedError(f"Support for {type(node)} not yet implemented.") + + def visit_IfScope(self, node: tn.IfScope, sdfg: SDFG) -> None: + before_state = self._current_state + assert before_state is not None + cf_region = before_state.parent_graph + + conditional_block = ConditionalBlock(f"if_scope_{id(node)}") + cf_region.add_node(conditional_block) + _insert_and_split_assignments( + before_state, + conditional_block, + assignments=self._pending_interstate_assignments(), + ) + + if_body = ControlFlowRegion("if_body", sdfg=sdfg) + conditional_block.add_branch(node.condition, if_body) + + if_state = if_body.add_state("if_state", is_start_block=True) + self._current_state = if_state + + # visit children of that branch + self.visit(node.children, sdfg=sdfg) + + self._current_state = conditional_block + + # add merge_state + merge_state = _insert_and_split_assignments( + conditional_block, + label="merge_state", + assignments=self._pending_interstate_assignments(), + ) + + # Check if there's an `ElseScope` following this node (in the parent's children). + # Filter StateBoundaryNodes, which we inserted earlier, for this analysis. + filtered = [n for n in node.parent.children if not isinstance(n, tn.StateBoundaryNode)] + if_index = _list_index(filtered, node) + has_else_branch = len(filtered) > if_index + 1 and isinstance(filtered[if_index + 1], tn.ElseScope) + + if has_else_branch: + # push merge_state on the stack for later usage in `visit_ElseScope` + self._state_stack.append(merge_state) + # push condition_block on the stack for later usage in `visit_ElseScope` + self._state_stack.append(conditional_block) + else: + self._current_state = merge_state + + def visit_StateIfScope(self, node: tn.StateIfScope, sdfg: SDFG) -> None: + raise NotImplementedError(f"Support for {type(node)} not yet implemented.") + + def visit_BreakNode(self, node: tn.BreakNode, sdfg: SDFG) -> None: + raise NotImplementedError(f"Support for {type(node)} not yet implemented.") + + def visit_ContinueNode(self, node: tn.ContinueNode, sdfg: SDFG) -> None: + raise NotImplementedError(f"Support for {type(node)} not yet implemented.") + + def visit_ElifScope(self, node: tn.ElifScope, sdfg: SDFG) -> None: + raise NotImplementedError(f"Support for {type(node)} not yet implemented.") + + def visit_ElseScope(self, node: tn.ElseScope, sdfg: SDFG) -> None: + # get ConditionalBlock from stack + conditional_block: ConditionalBlock = self._pop_state("if_scope") + + else_body = ControlFlowRegion("else_body", sdfg=sdfg) + conditional_block.add_branch(None, else_body) + + else_state = else_body.add_state("else_state", is_start_block=True) + self._current_state = else_state + + # visit children inside the else branch + self.visit(node.children, sdfg=sdfg) + + # merge false-branch into merge_state + merge_state = self._pop_state("merge_state") + self._current_state = merge_state + + if self._pending_interstate_assignments(): + raise NotImplementedError("TODO: update edge with new assignments") + + def _insert_nestedSDFG_in_MapScope(self, node: tn.MapScope, sdfg: SDFG) -> None: + dataflow_stack_size = len(self._dataflow_stack) + state_stack_size = len(self._state_stack) + outer_nestedSDFG = self._current_nestedSDFG + + # prepare inner SDFG + inner_sdfg = SDFG("nested_sdfg", parent=self._current_state) + start_state = inner_sdfg.add_state("nested_root", is_start_block=True) + + # update stacks and current state + old_state_label = self._current_state.label + self._state_stack.append(self._current_state) + self._dataflow_stack.append((inner_sdfg, {"inputs": set(), "outputs": set()})) + self._nviews_bound_per_scope[id(inner_sdfg)] = [] + self._nviews_deferred_removal[id(inner_sdfg)] = [] + self._current_nestedSDFG = id(inner_sdfg) + self._current_state = start_state + + # visit children + with _TreeScope(node, self._ctx, self._current_state): + self.visit(node.children, sdfg=inner_sdfg) + + # restore current state and stacks + self._current_state = self._pop_state(old_state_label) + assert len(self._state_stack) == state_stack_size + _, connectors = self._dataflow_stack.pop() + assert len(self._dataflow_stack) == dataflow_stack_size + + # insert nested SDFG + nsdfg = self._current_state.add_nested_sdfg( + sdfg=inner_sdfg, + inputs=connectors["inputs"], + outputs=connectors["outputs"], + ) + # connect nested SDFG to surrounding map scope + assert self._dataflow_stack + map_entry, to_connect = self._dataflow_stack[-1] + + # connect nsdfg input memlets (to be propagated upon completion of the SDFG) + for name in nsdfg.in_connectors: + out_connector = f"{PREFIX_PASSTHROUGH_OUT}{name}" + new_in_connector = map_entry.add_in_connector(f"{PREFIX_PASSTHROUGH_IN}{name}") + new_out_connector = map_entry.add_out_connector(out_connector) + assert new_in_connector == True + assert new_in_connector == new_out_connector + + # Add Memlet for NView node (if applicable) + edge_added = False + for nview in self._nviews_bound_per_scope[id(inner_sdfg)]: + if name == nview.target: + self._current_state.add_edge(map_entry, out_connector, nsdfg, name, + Memlet.from_memlet(nview.memlet)) + edge_added = True + break + + if not edge_added: + self._current_state.add_edge(map_entry, out_connector, nsdfg, name, + Memlet.from_array(name, nsdfg.sdfg.arrays[name])) + + # Add empty memlet if we didn't add any in the loop above + if self._current_state.out_degree(map_entry) < 1: + self._current_state.add_nedge(map_entry, nsdfg, Memlet()) + + # connect nsdfg output memlets (to be propagated) + for name in nsdfg.out_connectors: + # Add memlets for NView node (if applicable) + edge_added = False + for nview in self._nviews_bound_per_scope[id(inner_sdfg)]: + if name == nview.target: + to_connect[name] = (nsdfg, Memlet.from_memlet(nview.memlet)) + edge_added = True + break + + if not edge_added: + to_connect[name] = (nsdfg, Memlet.from_array(name, nsdfg.sdfg.arrays[name])) + + # Move NViews back to "free" NViews for usage in a sibling scope. + for nview in self._nviews_bound_per_scope[id(inner_sdfg)]: + # If this NView ended in the current nested SDFG, don't add it back to the + # "free NView" nodes. We need to keep it alive until here to make sure that + # we can add the memlets above. + if nview in self._nviews_deferred_removal[id(inner_sdfg)]: + continue + self._nviews_free.append(nview) + + del self._nviews_bound_per_scope[id(inner_sdfg)] + del self._nviews_deferred_removal[id(inner_sdfg)] + + # Restore current nested SDFG + self._current_nestedSDFG = outer_nestedSDFG + + def visit_MapScope(self, node: tn.MapScope, sdfg: SDFG) -> None: + dataflow_stack_size = len(self._dataflow_stack) + cache_state = self._current_state + + # map entry + # --------- + map_entry = nodes.MapEntry(node.node.map) + self._current_state.add_node(map_entry) + self._dataflow_stack.append((map_entry, dict())) + + # visit children inside the map + type_of_children = [type(child) for child in node.children] + last_child_is_MapScope = type_of_children[-1] == tn.MapScope + all_others_are_Boundaries = type_of_children.count(tn.StateBoundaryNode) == len(type_of_children) - 1 + if last_child_is_MapScope and all_others_are_Boundaries: + # skip weirdly added StateBoundaryNode + # tmp: use this - for now - to "backprop-insert" extra state boundaries for nested SDFGs + with _TreeScope(node, self._ctx, self._current_state): + self.visit(node.children[-1], sdfg=sdfg) + elif any([isinstance(child, tn.StateBoundaryNode) for child in node.children]): + self._insert_nestedSDFG_in_MapScope(node, sdfg) + else: + with _TreeScope(node, self._ctx, self._current_state): + self.visit(node.children, sdfg=sdfg) + + cache_key = (cache_state, id(self._ctx.current_scope)) + if cache_key not in self._ctx.access_cache: + self._ctx.access_cache[cache_key] = {} + access_cache = self._ctx.access_cache[cache_key] + + # dataflow stack management + _, to_connect = self._dataflow_stack.pop() + assert len(self._dataflow_stack) == dataflow_stack_size + outer_map_entry, outer_to_connect = self._dataflow_stack[-1] if dataflow_stack_size else (None, None) + + # connect potential input connectors on map_entry + for connector in map_entry.in_connectors: + memlet_data = connector.removeprefix(PREFIX_PASSTHROUGH_IN) + + # connect to local access node (if available) + if memlet_data in access_cache: + cached_access = access_cache[memlet_data] + self._current_state.add_memlet_path(cached_access, + map_entry, + dst_conn=connector, + memlet=Memlet.from_array(memlet_data, sdfg.arrays[memlet_data])) + continue + + if isinstance(outer_map_entry, nodes.EntryNode): + + # get it from outside the map + connector_name = f"{PREFIX_PASSTHROUGH_OUT}{memlet_data}" + if connector_name not in outer_map_entry.out_connectors: + new_in_connector = outer_map_entry.add_in_connector(connector) + new_out_connector = outer_map_entry.add_out_connector(connector_name) + assert new_in_connector == True + assert new_in_connector == new_out_connector + + self._current_state.add_edge(outer_map_entry, connector_name, map_entry, connector, + Memlet.from_array(memlet_data, sdfg.arrays[memlet_data])) + else: + if isinstance(outer_map_entry, SDFG): + # Copy data descriptor from parent SDFG and add input connector + if memlet_data not in sdfg.arrays: + parent_sdfg: SDFG = self._parent_sdfg_with_array(memlet_data, sdfg) + + # Add support for NView nodes + use_nview = self._apply_nview_array_override(memlet_data, sdfg) + if not use_nview: + sdfg.add_datadesc(memlet_data, parent_sdfg.arrays[memlet_data].clone()) + + # Transients passed into a nested SDFG become non-transient inside that nested SDFG + if parent_sdfg.arrays[memlet_data].transient: + sdfg.arrays[memlet_data].transient = False + + # Dev note: nview.target and memlet_data are identical + assert memlet_data not in outer_to_connect["inputs"] + outer_to_connect["inputs"].add(memlet_data) + else: + assert outer_map_entry is None + + # cache local read access + assert memlet_data not in access_cache + access_cache[memlet_data] = self._current_state.add_read(memlet_data) + cached_access = access_cache[memlet_data] + self._current_state.add_memlet_path(cached_access, + map_entry, + dst_conn=connector, + memlet=Memlet.from_array(memlet_data, sdfg.arrays[memlet_data])) + + if isinstance(outer_map_entry, nodes.EntryNode) and self._current_state.out_degree(outer_map_entry) < 1: + self._current_state.add_nedge(outer_map_entry, map_entry, Memlet()) + + # map_exit + # -------- + map_exit = nodes.MapExit(node.node.map) + self._current_state.add_node(map_exit) + + # connect writes to map_exit node + for name in to_connect: + in_connector_name = f"{PREFIX_PASSTHROUGH_IN}{name}" + out_connector_name = f"{PREFIX_PASSTHROUGH_OUT}{name}" + new_in_connector = map_exit.add_in_connector(in_connector_name) + new_out_connector = map_exit.add_out_connector(out_connector_name) + assert new_in_connector == new_out_connector + + # connect "inside the map" + access_node, memlet = to_connect[name] + if isinstance(access_node, nodes.NestedSDFG): + self._current_state.add_edge(access_node, name, map_exit, in_connector_name, memlet) + else: + assert isinstance(access_node, nodes.AccessNode) + if self._current_state.out_degree(access_node) == 0 and self._current_state.in_degree(access_node) == 1: + # this access_node is not used for anything else. + # let's remove it and add a direct connection instead + edges = [edge for edge in self._current_state.edges() if edge.dst == access_node] + assert len(edges) == 1 + self._current_state.add_memlet_path(edges[0].src, + map_exit, + src_conn=edges[0].src_conn, + dst_conn=in_connector_name, + memlet=edges[0].data) + self._current_state.remove_node(access_node) # edge is remove automatically + else: + self._current_state.add_memlet_path(access_node, + map_exit, + dst_conn=in_connector_name, + memlet=memlet) + + if isinstance(outer_map_entry, SDFG): + if name not in sdfg.arrays: + parent_sdfg = self._parent_sdfg_with_array(name, sdfg) + + # Support for NView nodes + use_nview = self._apply_nview_array_override(name, sdfg) + if not use_nview: + sdfg.add_datadesc(name, parent_sdfg.arrays[name].clone()) + + # Transients passed into a nested SDFG become non-transient inside that nested SDFG + if parent_sdfg.arrays[name].transient: + sdfg.arrays[name].transient = False + + # Add out_connector in any case if not yet present, e.g. write after read + # Dev not: name and nview.target are identical + outer_to_connect["outputs"].add(name) + + # connect "outside the map" + # only re-use cached write-only nodes, e.g. don't create a cycle for + # map i=0:20: + # A[i] = tasklet(A[i]) + if name not in access_cache or self._current_state.out_degree(access_cache[name]) > 0: + # cache write access into access_cache + write_access_node = self._current_state.add_write(name) + access_cache[name] = write_access_node + + access_node = access_cache[name] + self._current_state.add_memlet_path(map_exit, + access_node, + src_conn=out_connector_name, + memlet=Memlet.from_array(name, sdfg.arrays[name])) + + if isinstance(outer_map_entry, nodes.EntryNode): + outer_to_connect[name] = (access_node, Memlet.from_array(name, sdfg.arrays[name])) + else: + assert isinstance(outer_map_entry, SDFG) or outer_map_entry is None + + # TODO If nothing is connected at this point, figure out what's the last thing that + # we should connect to. Then, add an empty memlet from that last thing to this + # map_exit. + assert len(self._current_state.in_edges(map_exit)) > 0 + + def visit_ConsumeScope(self, node: tn.ConsumeScope, sdfg: SDFG) -> None: + raise NotImplementedError(f"Support for {type(node)} not yet implemented.") + + def visit_TaskletNode(self, node: tn.TaskletNode, sdfg: SDFG) -> None: + # Add Tasklet to current state + tasklet = node.node + self._current_state.add_node(tasklet) + + cache_key = (self._current_state, id(self._ctx.current_scope)) + if cache_key not in self._ctx.access_cache: + self._ctx.access_cache[cache_key] = {} + cache = self._ctx.access_cache[cache_key] + scope_node, to_connect = self._dataflow_stack[-1] if self._dataflow_stack else (None, None) + + # Connect input memlets + for name, memlet in node.in_memlets.items(): + # connect to local access node if possible + if memlet.data in cache: + cached_access = cache[memlet.data] + self._current_state.add_memlet_path(cached_access, tasklet, dst_conn=name, memlet=memlet) + continue + + if isinstance(scope_node, nodes.MapEntry): + # get it from outside the map + connector_name = f"{PREFIX_PASSTHROUGH_OUT}{memlet.data}" + if connector_name not in scope_node.out_connectors: + new_in_connector = scope_node.add_in_connector(f"{PREFIX_PASSTHROUGH_IN}{memlet.data}") + new_out_connector = scope_node.add_out_connector(connector_name) + assert new_in_connector == True + assert new_in_connector == new_out_connector + + self._current_state.add_edge(scope_node, connector_name, tasklet, name, memlet) + continue + + if isinstance(scope_node, SDFG): + # Copy data descriptor from parent SDFG and add input connector + if memlet.data not in sdfg.arrays: + parent_sdfg = self._parent_sdfg_with_array(memlet.data, sdfg) + + # Support for NView nodes + use_nview = self._apply_nview_array_override(memlet.data, sdfg) + if not use_nview: + sdfg.add_datadesc(memlet.data, parent_sdfg.arrays[memlet.data].clone()) + + # Transients passed into a nested SDFG become non-transient inside that nested SDFG + if parent_sdfg.arrays[memlet.data].transient: + sdfg.arrays[memlet.data].transient = False + + # Dev note: memlet.data and nview.target are identical + assert memlet.data not in to_connect["inputs"] + to_connect["inputs"].add(memlet.data) + else: + assert scope_node is None + + # cache local read access + assert memlet.data not in cache + cache[memlet.data] = self._current_state.add_read(memlet.data) + cached_access = cache[memlet.data] + self._current_state.add_memlet_path(cached_access, tasklet, dst_conn=name, memlet=memlet) + + # Add empty memlet if map_entry has no out_connectors to connect to + if isinstance(scope_node, nodes.MapEntry) and self._current_state.out_degree(scope_node) < 1: + self._current_state.add_nedge(scope_node, tasklet, Memlet()) + + # Connect output memlets + for name, memlet in node.out_memlets.items(): + # only re-use cached write-only nodes, e.g. don't create a cycle for + # A[1] = tasklet(A[1]) + if memlet.data not in cache or self._current_state.out_degree(cache[memlet.data]) > 0: + # cache write access node + write_access_node = self._current_state.add_write(memlet.data) + cache[memlet.data] = write_access_node + + access_node = cache[memlet.data] + self._current_state.add_memlet_path(tasklet, access_node, src_conn=name, memlet=memlet) + + if isinstance(scope_node, nodes.MapEntry): + # copy the memlet since we already used it in the memlet path above + to_connect[memlet.data] = (access_node, copy.deepcopy(memlet)) + continue + + if isinstance(scope_node, SDFG): + if memlet.data not in sdfg.arrays: + parent_sdfg: SDFG = self._parent_sdfg_with_array(memlet.data, sdfg) + + # Support for NView nodes + use_nview = self._apply_nview_array_override(memlet.data, sdfg) + if not use_nview: + sdfg.add_datadesc(memlet.data, parent_sdfg.arrays[memlet.data].clone()) + + # Transients passed into a nested SDFG become non-transient inside that nested SDFG + if parent_sdfg.arrays[memlet.data].transient: + sdfg.arrays[memlet.data].transient = False + + # Add out_connector in any case if not yet present, e.g. write after read + # Dev note: memlet.data and nview.target are identical + to_connect["outputs"].add(memlet.data) + + else: + assert scope_node is None + + def visit_LibraryCall(self, node: tn.LibraryCall, sdfg: SDFG) -> None: + raise NotImplementedError(f"Support for {type(node)} not yet implemented.") + + def visit_CopyNode(self, node: tn.CopyNode, sdfg: SDFG) -> None: + # ensure we have an access_cache and fetch it + cache_key = (self._current_state, id(self._ctx.current_scope)) + if cache_key not in self._ctx.access_cache: + self._ctx.access_cache[cache_key] = {} + access_cache = self._ctx.access_cache[cache_key] + + # assumption source access may or may not yet exist (in this state) + src_name = node.memlet.data + source = access_cache[src_name] if src_name in access_cache else self._current_state.add_read(src_name) + + # assumption: target access node doesn't exist yet + assert node.target not in access_cache + target = self._current_state.add_write(node.target) + + self._current_state.add_memlet_path(source, target, memlet=node.memlet) + + def visit_DynScopeCopyNode(self, node: tn.DynScopeCopyNode, sdfg: SDFG) -> None: + raise NotImplementedError(f"Support for {type(node)} not yet implemented.") + + def visit_ViewNode(self, node: tn.ViewNode, sdfg: SDFG) -> None: + raise NotImplementedError(f"Support for {type(node)} not yet implemented.") + + def visit_NView(self, node: tn.NView, sdfg: SDFG) -> None: + # Basic working principle: + # + # - NView and (artificial) NViewEnd nodes are added in parallel to mark the region where the view applies. + # - Keep a stack of NView nodes (per name) that is pushed/popped when NView and NViewEnd nodes are visited. + # - In between, when going "down into" a NestedSDFG, use the current NView (if it applies) + # - In between, when "coming back up" from a NestedSDFG, pop the NView from the stack. + # - AccessNodes will automatically pick up the right name (from the NestedSDFG's array list) + self._nviews_free.append(node) + + def visit_NViewEnd(self, node: tn.NViewEnd, sdfg: SDFG) -> None: + # If bound to the current nested SDFG, defer cleanup + if self._current_nestedSDFG is not None: + currently_bound = self._nviews_bound_per_scope[self._current_nestedSDFG] + for index, nview in enumerate(reversed(currently_bound)): + if node.target == nview.target: + # Bound to current nested SDFG. Slate for deferred removal once we exit that nested SDFG. + self._nviews_deferred_removal[self._current_nestedSDFG].append(nview) + return + + length = len(self._nviews_free) + for index, nview in enumerate(reversed(self._nviews_free), start=1): + if node.target == nview.target: + # Stack semantics: remove from the back of the list + del self._nviews_free[length - index] + return + + raise RuntimeError(f"No matching NView found for target {node.target} in {self._nviews_free}.") + + def visit_RefSetNode(self, node: tn.RefSetNode, sdfg: SDFG) -> None: + raise NotImplementedError(f"Support for {type(node)} not yet implemented.") + + def visit_StateBoundaryNode(self, node: tn.StateBoundaryNode, sdfg: SDFG) -> None: + # When creating a state boundary, include all inter-state assignments that precede it. + pending = self._pending_interstate_assignments() + + self._current_state = _create_state_boundary( + node, + self._current_state, + assignments=pending, + ) + + def _pending_interstate_assignments(self) -> dict[str, str]: + """ + Return currently pending interstate assignments. Clears the cache. + """ + assignments = {} + + for symbol in self._interstate_symbols: + assignments[symbol.name] = symbol.value.as_string + self._interstate_symbols.clear() + + return assignments + + +def from_schedule_tree( + stree: tn.ScheduleTreeRoot, + state_boundary_behavior: StateBoundaryBehavior = StateBoundaryBehavior.STATE_TRANSITION, + max_nested_sdfgs: int = 1000, +) -> SDFG: + """ + Converts a schedule tree into an SDFG. + + :param stree: The schedule tree root to convert. + :param state_boundary_behavior: Sets the behavior upon encountering a state boundary (e.g., write-after-write). + See the ``StateBoundaryBehavior`` enumeration for more details. + :return: An SDFG representing the schedule tree. + """ + # Setup SDFG descriptor repository + result = SDFG(stree.name, propagate=False) + result.arg_names = copy.deepcopy(stree.arg_names) + for key, container in stree.containers.items(): + result._arrays[key] = copy.deepcopy(container) + result.constants_prop = copy.deepcopy(stree.constants) + result.symbols = copy.deepcopy(stree.symbols) + + # Insert artificial state boundaries after WAW, before label, etc. + stree = _insert_state_boundaries_to_tree(stree) + + # Traverse tree and incrementally build SDFG, finally propagate memlets + _StreeToSDFG(boundary_behavior=state_boundary_behavior, max_nested_sdfg=max_nested_sdfgs).visit(stree, sdfg=result) + propagation.propagate_memlets_sdfg(result) + + return result + + +def _insert_state_boundaries_to_tree(stree: tn.ScheduleTreeRoot) -> tn.ScheduleTreeRoot: + """ + Inserts StateBoundaryNode objects into a schedule tree where more than one SDFG state would be necessary. + Operates in-place on the given schedule tree. + + This happens when there is a: + * write-after-write dependency; + * write-after-read dependency that cannot be fulfilled via memlets; + * control flow block (for/if); or + * otherwise before a state label (which means a state transition could occur, e.g., in a gblock) + + :param stree: The schedule tree to operate on. + """ + + # Simple boundary node inserter for control flow blocks and state labels + class SimpleStateBoundaryInserter(tn.ScheduleNodeTransformer): + + def visit_scope(self, scope: tn.ScheduleTreeScope): + if isinstance(scope, tn.ControlFlowScope) and not isinstance(scope, (tn.ElifScope, tn.ElseScope)): + return [tn.StateBoundaryNode(True), self.generic_visit(scope)] + return self.generic_visit(scope) + + def visit_StateLabel(self, node: tn.StateLabel): + return [tn.StateBoundaryNode(True), self.generic_visit(node)] + + # First, insert boundaries around labels and control flow + stree = SimpleStateBoundaryInserter().visit(stree) + + # Then, insert boundaries after unmet memory dependencies or potential data races + _insert_memory_dependency_state_boundaries(stree) + + # Insert a state boundary after every symbol assignment to ensure symbols are assigned before usage + class SymbolAssignmentBoundaryInserter(tn.ScheduleNodeTransformer): + + def visit_AssignNode(self, node: tn.AssignNode): + # We can assume that assignment nodes are at least contained in the root scope. + assert node.parent, "Expected assignment nodes live a parent scope." + + # Find this node in the parent's children. + node_index = _list_index(node.parent.children, node) + + # Don't add boundary if there's already one or for immediately following assignment nodes. + if node_index < len(node.parent.children) - 1 and isinstance(node.parent.children[node_index + 1], + (tn.StateBoundaryNode, tn.AssignNode)): + return self.generic_visit(node) + + return [self.generic_visit(node), tn.StateBoundaryNode()] + + stree = SymbolAssignmentBoundaryInserter().visit(stree) + + # Hack: "backprop-insert" state boundaries from nested SDFGs + class NestedSDFGStateBoundaryInserter(tn.ScheduleNodeTransformer): + + def visit_MapScope(self, scope: tn.MapScope): + visited = self.generic_visit(scope) + if any([isinstance(child, tn.StateBoundaryNode) for child in scope.children]): + # We can assume that map nodes are at least contained in the root scope. + assert scope.parent is not None + + # Find this scope in its parent's children + node_index = _list_index(scope.parent.children, scope) + + # If there's already a state boundary before the map, don't add another one + if node_index > 0 and isinstance(scope.parent.children[node_index - 1], tn.StateBoundaryNode): + return visited + + return [tn.StateBoundaryNode(), visited] + return visited + + stree = NestedSDFGStateBoundaryInserter().visit(stree) + + return stree + + +def _insert_memory_dependency_state_boundaries(scope: tn.ScheduleTreeScope): + """ + Helper function that inserts boundaries after unmet memory dependencies. + """ + reads: mmu.MemletDict[list[tn.ScheduleTreeNode]] = mmu.MemletDict() + writes: mmu.MemletDict[list[tn.ScheduleTreeNode]] = mmu.MemletDict() + parents: dict[int, set[int]] = defaultdict(set) + boundaries_to_insert: list[int] = [] + + for i, n in enumerate(scope.children): + if isinstance(n, (tn.StateBoundaryNode, tn.ControlFlowScope)): # Clear state + reads.clear() + writes.clear() + parents.clear() + if isinstance(n, tn.ControlFlowScope): # Insert memory boundaries recursively + _insert_memory_dependency_state_boundaries(n) + continue + + # If dataflow scope, insert state boundaries recursively and as a node + if isinstance(n, tn.DataflowScope): + _insert_memory_dependency_state_boundaries(n) + + inputs = n.input_memlets() + outputs = n.output_memlets() + + # Register reads + for inp in inputs: + if inp not in reads: + reads[inp] = [n] + else: + reads[inp].append(n) + + # Transitively add parents + if inp in writes: + for parent in writes[inp]: + parents[id(n)].add(id(parent)) + parents[id(n)].update(parents[id(parent)]) + + # Inter-state assignment nodes with reads necessitate a state transition if they were written to. + if isinstance(n, tn.AssignNode) and any(inp in writes for inp in inputs): + boundaries_to_insert.append(i) + reads.clear() + writes.clear() + parents.clear() + continue + + # Write after write or potential write/write data race, insert state boundary + if any(o in writes and (o not in reads or any(id(r) not in parents for r in reads[o])) for o in outputs): + boundaries_to_insert.append(i) + reads.clear() + writes.clear() + parents.clear() + continue + + # Potential read/write data race: if any read is not in the parents of this node, it might + # be performed in parallel + if any(o in reads and any(id(r) not in parents for r in reads[o]) for o in outputs): + boundaries_to_insert.append(i) + reads.clear() + writes.clear() + parents.clear() + continue + + # Register writes after all hazards have been tested for + for out in outputs: + if out not in writes: + writes[out] = [n] + else: + writes[out].append(n) + + # Insert memory dependency state boundaries in reverse in order to keep indices intact + for i in reversed(boundaries_to_insert): + scope.children.insert(i, tn.StateBoundaryNode()) + + +############################################################################# +# SDFG content creation functions + + +def _create_state_boundary( + boundary_node: tn.StateBoundaryNode, + state: SDFGState, + assignments: dict[str, str] | None = None, +) -> SDFGState: + """ + Creates a boundary between two states + + :param boundary_node: The state boundary node to generate. + :param state: The last state prior to this boundary. + :return: The newly created state. + """ + label = "cf_state_boundary" if boundary_node.due_to_control_flow else "state_boundary" + assignments = assignments if assignments is not None else {} + return _insert_and_split_assignments(state, label=label, assignments=assignments) + + +def _insert_and_split_assignments( + before_state: ControlFlowBlock, + after_state: ControlFlowBlock | None = None, + *, + label: str | None = None, + assignments: dict[str, str] | None = None, +) -> ControlFlowBlock: + """ + Insert given assignments splitting them in case of potential race conditions. + + The semantics of the SDFG dictates that we can not assume any order in the application + of inter-state edge assignments. The only order is that conditions precede assignments. + + Since we just collect all inter-state assignments while parsing the schedule tree, we + need to make sure to split problematic assignments over multiple state transitions. + """ + assignments = assignments if assignments is not None else {} + cf_region = before_state.parent_graph + if after_state is not None and after_state.parent_graph != cf_region: + raise ValueError("Expected before_state and after_state to be in the same control flow region.") + + has_potential_race = False + for key, value in assignments.items(): + syms = symbolic.free_symbols_and_functions(value) + also_assigned = (syms & assignments.keys()) - {key} + if also_assigned: + has_potential_race = True + break + + if not has_potential_race: + if after_state is not None: + cf_region.add_edge(before_state, after_state, InterstateEdge(assignments=assignments)) + return after_state + + return cf_region.add_state_after(before_state, label=label, assignments=assignments) + + last_state = before_state + for index, assignment in enumerate(assignments.items()): + key, value = assignment + is_last_state = index == len(assignments) - 1 + if is_last_state and after_state is not None: + cf_region.add_edge(last_state, after_state, InterstateEdge(assignments={key: value})) + last_state = after_state + else: + last_state = cf_region.add_state_after(last_state, label=label, assignments={key: value}) + + return last_state + + +def _list_index(list: list[tn.ScheduleTreeNode], node: tn.ScheduleTreeNode) -> int: + """Check if node is in list with "is" operator.""" + index = 0 + for element in list: + # compare with "is" to get memory comparison. ".index()" uses value comparison + if element is node: + return index + index += 1 + + raise StopIteration diff --git a/dace/sdfg/analysis/schedule_tree/treenodes.py b/dace/sdfg/analysis/schedule_tree/treenodes.py index 0d709ee0fd..70297e6db9 100644 --- a/dace/sdfg/analysis/schedule_tree/treenodes.py +++ b/dace/sdfg/analysis/schedule_tree/treenodes.py @@ -1,13 +1,17 @@ # Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. +from collections.abc import Mapping from dataclasses import dataclass, field +import sympy -from dace import nodes, data, subsets +from dace import nodes, data, subsets, dtypes, symbolic from dace.properties import CodeBlock from dace.sdfg import InterstateEdge -from dace.sdfg.state import ConditionalBlock, LoopRegion, SDFGState -from dace.symbolic import symbol +from dace.sdfg.memlet_utils import MemletSet +from dace.sdfg.propagation import propagate_subset +from dace.sdfg.sdfg import InterstateEdge, SDFG, memlets_in_ast +from dace.sdfg.state import LoopRegion, SDFGState from dace.memlet import Memlet -from typing import TYPE_CHECKING, Dict, Iterator, List, Literal, Optional, Set, Union +from typing import TYPE_CHECKING, Any, Iterable, Iterator, Literal, Optional if TYPE_CHECKING: from dace import SDFG @@ -21,9 +25,10 @@ class UnsupportedScopeException(Exception): @dataclass class ScheduleTreeNode: + """Base class for nodes in the schedule tree.""" parent: Optional['ScheduleTreeScope'] = field(default=None, init=False, repr=False) - def as_string(self, indent: int = 0): + def as_string(self, indent: int = 0) -> str: return indent * INDENTATION + 'UNSUPPORTED' def preorder_traversal(self) -> Iterator['ScheduleTreeNode']: @@ -32,20 +37,55 @@ def preorder_traversal(self) -> Iterator['ScheduleTreeNode']: """ yield self + def get_root(self) -> 'ScheduleTreeRoot': + if self.parent is None: + raise ValueError('Non-root schedule tree node has no parent.') + return self.parent.get_root() + + def input_memlets(self, root: Optional['ScheduleTreeRoot'] = None, **kwargs: dict[str, Any]) -> MemletSet: + """ + Returns a set of inputs for this node. For scopes, returns the union of its contents. + + :param root: An optional argument specifying the schedule tree's root. If not given, + the value is computed from the current tree node. + :return: A set of memlets representing the inputs of this node. + """ + raise NotImplementedError + + def output_memlets(self, root: Optional['ScheduleTreeRoot'] = None, **kwargs: dict[str, Any]) -> MemletSet: + """ + Returns a set of outputs for this node. For scopes, returns the union of its contents. + + :param root: An optional argument specifying the schedule tree's root. If not given, + the value is computed from the current tree node. + :return: A set of memlets representing the inputs of this node. + """ + raise NotImplementedError + @dataclass class ScheduleTreeScope(ScheduleTreeNode): - children: List['ScheduleTreeNode'] - containers: Optional[Dict[str, data.Data]] = field(default_factory=dict, init=False) - symbols: Optional[Dict[str, symbol]] = field(default_factory=dict, init=False) - - def __init__(self, children: Optional[List['ScheduleTreeNode']] = None): - self.children = children or [] - if self.children: - for child in children: - child.parent = self - self.containers = {} - self.symbols = {} + """ + A `ScheduleTreeScope` is the base class for grouping `ScheduleTreeNode`s hierarchically. + + Each scope holds a list of children and a reference to the parent. + """ + children: list[ScheduleTreeNode] + + def __init__(self, *, children: list[ScheduleTreeNode], parent: Optional['ScheduleTreeScope'] = None) -> None: + for child in children: + child.parent = self + + self.children = children + self.parent = parent + + def add_children(self, children: Iterable[ScheduleTreeNode]) -> None: + for child in children: + child.parent = self + self.children.append(child) + + def add_child(self, child: ScheduleTreeNode) -> None: + self.add_children([child]) def as_string(self, indent: int = 0): if not self.children: @@ -60,18 +100,195 @@ def preorder_traversal(self) -> Iterator['ScheduleTreeNode']: for child in self.children: yield from child.preorder_traversal() - # TODO: Helper function that gets input/output memlets of the scope + def _gather_memlets_in_scope(self, inputs: bool, root: Optional['ScheduleTreeRoot'], keep_locals: bool, + propagate: dict[str, + subsets.Range], disallow_propagation: set[str], **kwargs) -> MemletSet: + gather = (lambda n, root: n.input_memlets(root, **kwargs)) if inputs else ( + lambda n, root: n.output_memlets(root, **kwargs)) + + # Fast path, no propagation necessary + if keep_locals: + return MemletSet().union(*(gather(c, root) for c in self.children)) + + root = root if root is not None else self.get_root() + + if propagate: + to_propagate = list(propagate.items()) + propagate_keys = [a[0] for a in to_propagate] + propagate_values = subsets.Range([a[1] for a in to_propagate]) + + current_locals = set() + current_locals |= disallow_propagation + result = MemletSet() + + # Loop over children in order, if any new symbol is defined within this scope (e.g., symbol assignment, + # dynamic map range), consider it as a new local + for c in self.children: + # Add new locals + if isinstance(c, AssignNode): + current_locals.add(c.name) + elif isinstance(c, DynScopeCopyNode): + current_locals.add(c.target) + + internal_memlets: MemletSet = gather(c, root) + if propagate: + for memlet in internal_memlets: + result.add( + propagate_subset([memlet], + root.containers[memlet.data], + propagate_keys, + propagate_values, + undefined_variables=current_locals, + use_dst=not inputs)) + + return result + + def input_memlets(self, + root: Optional['ScheduleTreeRoot'] = None, + keep_locals: bool = False, + propagate: dict[str, subsets.Range] | None = None, + disallow_propagation: set[str] | None = None, + **kwargs) -> MemletSet: + """ + Returns a union of the set of inputs for this scope. Propagates the memlets used in the scope if ``keep_locals`` + is set to False. + + :param root: An optional argument specifying the schedule tree's root. If not given, + the value is computed from the current tree node. + :param keep_locals: If True, keeps the local symbols defined within the scope as part of the resulting memlets. + Otherwise, performs memlet propagation (see ``propagate`` and ``disallow_propagation``) or + assumes the entire container is used. + :param propagate: An optional dictionary mapping symbols to their corresponding ranges outside of this scope. + For example, the range of values a for-loop may take. + If ``keep_locals`` is False, this dictionary will be used to create projection memlets over + the ranges. See :ref:`memprop` in the documentation for more information. + :param disallow_propagation: If ``keep_locals`` is False, this optional set of strings will be considered + as additional locals. + :return: A set of memlets representing the inputs of this scope. + """ + return self._gather_memlets_in_scope(True, root, keep_locals, propagate or {}, disallow_propagation or set(), + **kwargs) + + def output_memlets(self, + root: Optional['ScheduleTreeRoot'] = None, + keep_locals: bool = False, + propagate: dict[str, subsets.Range] | None = None, + disallow_propagation: set[str] | None = None, + **kwargs) -> MemletSet: + """ + Returns a union of the set of outputs for this scope. Propagates the memlets used in the scope if + ``keep_locals`` is set to False. + + :param root: An optional argument specifying the schedule tree's root. If not given, + the value is computed from the current tree node. + :param keep_locals: If True, keeps the local symbols defined within the scope as part of the resulting memlets. + Otherwise, performs memlet propagation (see ``propagate`` and ``disallow_propagation``) or + assumes the entire container is used. + :param propagate: An optional dictionary mapping symbols to their corresponding ranges outside of this scope. + For example, the range of values a for-loop may take. + If ``keep_locals`` is False, this dictionary will be used to create projection memlets over + the ranges. See :ref:`memprop` in the documentation for more information. + :param disallow_propagation: If ``keep_locals`` is False, this optional set of strings will be considered + as additional locals. + :return: A set of memlets representing the inputs of this scope. + """ + return self._gather_memlets_in_scope(False, root, keep_locals, propagate or {}, disallow_propagation or set(), + **kwargs) + + +@dataclass +class ScheduleTreeRoot(ScheduleTreeScope): + """ + The root of a schedule tree. This is a `ScheduleTreeScope` with additional information on + the available descriptors, symbol types, and constants of the tree, aka the descriptor repository. + + Each schedule tree has only one `ScheduleTreeRoot`. The `ScheduleTreeRoot` is the only `ScheduleTreeScope` + without a parent (because it is the root node of the tree). + """ + name: str + containers: dict[str, data.Data] + symbols: Mapping[str, dtypes.typeclass | symbolic.symbol] + constants: dict[str, tuple[data.Data, Any]] + callback_mapping: dict[str, str] + arg_names: list[str] + + def __init__( + self, + *, + name: str, + children: list[ScheduleTreeNode], + containers: dict[str, data.Data] | None = None, + symbols: Mapping[str, dtypes.typeclass | symbolic.symbol] | None = None, + constants: dict[str, tuple[data.Data, Any]] | None = None, + callback_mapping: dict[str, str] | None = None, + arg_names: list[str] | None = None, + ) -> None: + super().__init__(children=children, parent=None) + + self.name = name + self.containers = containers if containers is not None else dict() + self.symbols = symbols if symbols is not None else dict() + self.constants = constants if constants is not None else dict() + self.callback_mapping = callback_mapping if callback_mapping is not None else dict() + self.arg_names = arg_names if arg_names is not None else list() + + def as_sdfg(self, + validate: bool = True, + simplify: bool = True, + validate_all: bool = False, + skip: set[str] | None = None, + verbose: bool = False) -> SDFG: + """ + Convert this schedule tree representation (back) into an SDFG. + + :param validate: If true, validate generated SDFG. + :param simplify: If true, simplify generated SDFG. The conversion might insert things like extra + empty states that can be cleaned up automatically. The value of `validate` is + passed on to `simplify()`. + :param validate_all: When simplifying, validate all intermediate SDFGs. Unused if simplify is False. + :param skip: Set of names of simplify passes to skip. Unused if simplify is False. + :param verbose: Turn on verbose logging of simplify. Unused if simplify is False. + + :return: SDFG version of this schedule tree. + """ + from dace.sdfg.analysis.schedule_tree import tree_to_sdfg as t2s # Avoid import loop + sdfg = t2s.from_schedule_tree(self) + + if validate: + sdfg.validate() + + if simplify: + skip = set() if skip is None else skip + sdfg.simplify(validate=validate, validate_all=validate_all, verbose=verbose, skip=skip) + + return sdfg + + def get_root(self) -> 'ScheduleTreeRoot': + return self @dataclass class ControlFlowScope(ScheduleTreeScope): - pass + + def __init__(self, *, children: list[ScheduleTreeNode], parent: ScheduleTreeScope | None = None) -> None: + super().__init__(children=children, parent=parent) @dataclass class DataflowScope(ScheduleTreeScope): node: nodes.EntryNode - state: Optional[SDFGState] = None + state: SDFGState | None = None + + def __init__(self, + *, + node: nodes.EntryNode, + children: list[ScheduleTreeNode], + parent: ScheduleTreeScope | None = None, + state: SDFGState | None = None) -> None: + super().__init__(children=children, parent=parent) + + self.node = node + self.state = state @dataclass @@ -82,6 +299,9 @@ class GBlock(ControlFlowScope): Normally contains irreducible control flow. """ + def __init__(self, *, children: list[ScheduleTreeNode], parent: ScheduleTreeScope | None = None) -> None: + super().__init__(children=children, parent=parent) + def as_string(self, indent: int = 0): result = indent * INDENTATION + 'gblock:\n' return result + super().as_string(indent) @@ -94,15 +314,27 @@ class StateLabel(ScheduleTreeNode): def as_string(self, indent: int = 0): return indent * INDENTATION + f'label {self.state.name}:' + def input_memlets(self, root: ScheduleTreeRoot | None = None, **kwargs) -> MemletSet: + return MemletSet() + + def output_memlets(self, root: ScheduleTreeRoot | None = None, **kwargs) -> MemletSet: + return MemletSet() + @dataclass class GotoNode(ScheduleTreeNode): - target: Optional[str] = None #: If None, equivalent to "goto exit" or "return" + target: str | None = None #: If None, equivalent to "goto exit" or "return" def as_string(self, indent: int = 0): name = self.target or 'exit' return indent * INDENTATION + f'goto {name}' + def input_memlets(self, root: ScheduleTreeRoot | None = None, **kwargs) -> MemletSet: + return MemletSet() + + def output_memlets(self, root: ScheduleTreeRoot | None = None, **kwargs) -> MemletSet: + return MemletSet() + @dataclass class AssignNode(ScheduleTreeNode): @@ -116,6 +348,13 @@ class AssignNode(ScheduleTreeNode): def as_string(self, indent: int = 0): return indent * INDENTATION + f'assign {self.name} = {self.value.as_string}' + def input_memlets(self, root: ScheduleTreeRoot | None = None, **kwargs) -> MemletSet: + root = root if root is not None else self.get_root() + return MemletSet(self.edge.get_read_memlets(root.containers)) + + def output_memlets(self, root: ScheduleTreeRoot | None = None, **kwargs) -> MemletSet: + return MemletSet() + @dataclass class LoopScope(ControlFlowScope): @@ -124,51 +363,153 @@ class LoopScope(ControlFlowScope): """ loop: LoopRegion - def _check_loop_variant( - self - ) -> Union[Literal['for'], Literal['while'], Literal['do-while'], Literal['do-for-uncond-increment'], - Literal['do-for']]: - if self.loop.update_statement and self.loop.init_statement and self.loop.loop_variable: - if self.loop.inverted: - if self.loop.update_before_condition: - return 'do-for-uncond-increment' - else: - return 'do-for' - else: - return 'for' - else: - if self.loop.inverted: - return 'do-while' - else: - return 'while' + def __init__(self, + *, + loop: LoopRegion, + children: list[ScheduleTreeNode], + parent: ScheduleTreeScope | None = None) -> None: + super().__init__(children=children, parent=parent) + + self.loop = loop def as_string(self, indent: int = 0): loop = self.loop - loop_variant = self._check_loop_variant() - if loop_variant == 'do-for-uncond-increment': + variant = loop_variant(loop) + if variant == 'do-for-uncond-increment': pre_header = indent * INDENTATION + f'{loop.init_statement.as_string}\n' header = indent * INDENTATION + 'do:\n' pre_footer = (indent + 1) * INDENTATION + f'{loop.update_statement.as_string}\n' footer = indent * INDENTATION + f'while {loop.loop_condition.as_string}' return pre_header + header + super().as_string(indent) + '\n' + pre_footer + footer - elif loop_variant == 'do-for': + + if variant == 'do-for': pre_header = indent * INDENTATION + f'{loop.init_statement.as_string}\n' header = indent * INDENTATION + 'while True:\n' pre_footer = (indent + 1) * INDENTATION + f'if (not {loop.loop_condition.as_string}):\n' pre_footer += (indent + 2) * INDENTATION + 'break\n' footer = (indent + 1) * INDENTATION + f'{loop.update_statement.as_string}\n' return pre_header + header + super().as_string(indent) + '\n' + pre_footer + footer - elif loop_variant == 'for': - result = (indent * INDENTATION + f'for {loop.init_statement.as_string}; ' + - f'{loop.loop_condition.as_string}; ' + f'{loop.update_statement.as_string}:\n') - return result + super().as_string(indent) - elif loop_variant == 'while': - result = indent * INDENTATION + f'while {loop.loop_condition.as_string}:\n' - return result + super().as_string(indent) - else: # 'do-while' - header = indent * INDENTATION + 'do:\n' - footer = indent * INDENTATION + f'while {loop.loop_condition.as_string}' - return header + super().as_string(indent) + '\n' + footer + + if variant in ["for", "while", "do-while"]: + return super().as_string(indent) + + raise NotImplementedError(f"Unknown LoopRegion variant '{variant}.") + + +@dataclass +class ForScope(LoopScope): + """Specialized LoopScope for for-loops.""" + + def __init__(self, + *, + loop: LoopRegion, + children: list[ScheduleTreeNode], + parent: ScheduleTreeScope | None = None) -> None: + super().__init__(loop=loop, children=children, parent=parent) + + def as_string(self, indent: int = 0) -> str: + init_statement = self.loop.init_statement.as_string + condition = self.loop.loop_condition.as_string + update_statement = self.loop.update_statement.as_string + result = indent * INDENTATION + f"for {init_statement}; {condition}; {update_statement}:\n" + return result + super().as_string(indent) + + def input_memlets(self, + root: ScheduleTreeRoot | None = None, + keep_locals: bool = False, + propagate: dict[str, subsets.Range] | None = None, + disallow_propagation: set[str] | None = None, + **kwargs) -> MemletSet: + root = root if root is not None else self.get_root() + + result = MemletSet() + result.update(self.loop.get_meta_read_memlets(arrays=root.containers)) + + # If loop range is well-formed, use it in propagation + range = _loop_range(self.loop) + if range is not None: + propagate = {self.loop.loop_variable: range} + else: + propagate = None + + result.update(super().input_memlets(root, propagate=propagate, **kwargs)) + return result + + def output_memlets(self, + root: ScheduleTreeRoot | None = None, + keep_locals: bool = False, + propagate: dict[str, subsets.Range] | None = None, + disallow_propagation: set[str] | None = None, + **kwargs) -> MemletSet: + + # If loop range is well-formed, use it in propagation + range = _loop_range(self.loop) + if range is not None: + propagate = {self.loop.loop_variable: range} + else: + propagate = None + + return super().output_memlets(root, propagate=propagate, **kwargs) + + +@dataclass +class WhileScope(LoopScope): + """Specialized LoopScope for while-loops.""" + + def __init__(self, + *, + loop: LoopRegion, + children: list[ScheduleTreeNode], + parent: ScheduleTreeScope | None = None) -> None: + super().__init__(loop=loop, children=children, parent=parent) + + def as_string(self, indent: int = 0) -> str: + condition = self.loop.loop_condition.as_string + result = indent * INDENTATION + f'while {condition}:\n' + return result + super().as_string(indent) + + def input_memlets(self, + root: ScheduleTreeRoot | None = None, + keep_locals: bool = False, + propagate: dict[str, subsets.Range] | None = None, + disallow_propagation: set[str] | None = None, + **kwargs) -> MemletSet: + root = root if root is not None else self.get_root() + + result = MemletSet() + result.update(self.loop.get_meta_read_memlets(arrays=root.containers)) + result.update(super().input_memlets(root, **kwargs)) + return result + + +@dataclass +class DoWhileScope(LoopScope): + """Specialized LoopScope for do-while-loops""" + + def __init__(self, + *, + loop: LoopRegion, + children: list[ScheduleTreeNode], + parent: ScheduleTreeScope | None = None) -> None: + super().__init__(loop=loop, children=children, parent=parent) + + def as_string(self, indent: int = 0) -> str: + header = indent * INDENTATION + 'do:\n' + footer = indent * INDENTATION + f'while {self.loop.loop_condition.as_string}' + return header + super().as_string(indent) + '\n' + footer + + def input_memlets(self, + root: ScheduleTreeRoot | None = None, + keep_locals: bool = False, + propagate: dict[str, subsets.Range] | None = None, + disallow_propagation: set[str] | None = None, + **kwargs) -> MemletSet: + root = root if root is not None else self.get_root() + + result = MemletSet() + result.update(self.loop.get_meta_read_memlets(arrays=root.containers)) + result.update(super().input_memlets(root, **kwargs)) + return result @dataclass @@ -178,10 +519,31 @@ class IfScope(ControlFlowScope): """ condition: CodeBlock + def __init__(self, + *, + condition: CodeBlock, + children: list[ScheduleTreeNode], + parent: ScheduleTreeScope | None = None) -> None: + super().__init__(children=children, parent=parent) + + self.condition = condition + def as_string(self, indent: int = 0): result = indent * INDENTATION + f'if {self.condition.as_string}:\n' return result + super().as_string(indent) + def input_memlets(self, + root: ScheduleTreeRoot | None = None, + keep_locals: bool = False, + propagate: dict[str, subsets.Range] | None = None, + disallow_propagation: set[str] | None = None, + **kwargs) -> MemletSet: + root = root if root is not None else self.get_root() + result = MemletSet() + result.update(memlets_in_ast(self.condition.code[0], root.containers)) + result.update(super().input_memlets(root, **kwargs)) + return result + @dataclass class StateIfScope(IfScope): @@ -189,6 +551,13 @@ class StateIfScope(IfScope): A special class of an if scope in general blocks for if statements that are part of a state transition. """ + def __init__(self, + *, + condition: CodeBlock, + children: list[ScheduleTreeNode], + parent: ScheduleTreeScope | None = None) -> None: + super().__init__(condition=condition, children=children, parent=parent) + def as_string(self, indent: int = 0): result = indent * INDENTATION + f'stateif {self.condition.as_string}:\n' return result + super(IfScope, self).as_string(indent) @@ -203,6 +572,12 @@ class BreakNode(ScheduleTreeNode): def as_string(self, indent: int = 0): return indent * INDENTATION + 'break' + def input_memlets(self, root: ScheduleTreeRoot | None = None, **kwargs) -> MemletSet: + return MemletSet() + + def output_memlets(self, root: ScheduleTreeRoot | None = None, **kwargs) -> MemletSet: + return MemletSet() + @dataclass class ContinueNode(ScheduleTreeNode): @@ -213,6 +588,12 @@ class ContinueNode(ScheduleTreeNode): def as_string(self, indent: int = 0): return indent * INDENTATION + 'continue' + def input_memlets(self, root: ScheduleTreeRoot | None = None, **kwargs) -> MemletSet: + return MemletSet() + + def output_memlets(self, root: ScheduleTreeRoot | None = None, **kwargs) -> MemletSet: + return MemletSet() + @dataclass class ElifScope(ControlFlowScope): @@ -221,10 +602,31 @@ class ElifScope(ControlFlowScope): """ condition: CodeBlock + def __init__(self, + *, + condition: CodeBlock, + children: list[ScheduleTreeNode], + parent: ScheduleTreeScope | None = None) -> None: + super().__init__(children=children, parent=parent) + + self.condition = condition + def as_string(self, indent: int = 0): result = indent * INDENTATION + f'elif {self.condition.as_string}:\n' return result + super().as_string(indent) + def input_memlets(self, + root: ScheduleTreeRoot | None = None, + keep_locals: bool = False, + propagate: dict[str, subsets.Range] | None = None, + disallow_propagation: set[str] | None = None, + **kwargs) -> MemletSet: + root = root if root is not None else self.get_root() + result = MemletSet() + result.update(memlets_in_ast(self.condition.code[0], root.containers)) + result.update(super().input_memlets(root, **kwargs)) + return result + @dataclass class ElseScope(ControlFlowScope): @@ -232,6 +634,9 @@ class ElseScope(ControlFlowScope): Else branch scope. """ + def __init__(self, *, children: list[ScheduleTreeNode], parent: ScheduleTreeScope | None = None) -> None: + super().__init__(children=children, parent=parent) + def as_string(self, indent: int = 0): result = indent * INDENTATION + 'else:\n' return result + super().as_string(indent) @@ -242,18 +647,62 @@ class MapScope(DataflowScope): """ Map scope. """ + node: nodes.MapEntry + + def __init__(self, + *, + node: nodes.MapEntry, + children: list[ScheduleTreeNode], + parent: ScheduleTreeScope | None = None, + state: SDFGState | None = None) -> None: + super().__init__(node=node, state=state, children=children, parent=parent) def as_string(self, indent: int = 0): rangestr = ', '.join(subsets.Range.dim_to_string(d) for d in self.node.map.range) result = indent * INDENTATION + f'map {", ".join(self.node.map.params)} in [{rangestr}]:\n' return result + super().as_string(indent) + def input_memlets(self, + root: ScheduleTreeRoot | None = None, + keep_locals: bool = False, + propagate: dict[str, subsets.Range] | None = None, + disallow_propagation: set[str] | None = None, + **kwargs) -> MemletSet: + return super().input_memlets(root, + propagate={ + k: v + for k, v in zip(self.node.map.params, self.node.map.range) + }, + **kwargs) + + def output_memlets(self, + root: ScheduleTreeRoot | None = None, + keep_locals: bool = False, + propagate: dict[str, subsets.Range] | None = None, + disallow_propagation: set[str] | None = None, + **kwargs) -> MemletSet: + return super().output_memlets(root, + propagate={ + k: v + for k, v in zip(self.node.map.params, self.node.map.range) + }, + **kwargs) + @dataclass class ConsumeScope(DataflowScope): """ Consume scope. """ + node: nodes.ConsumeEntry + + def __init__(self, + *, + node: nodes.ConsumeEntry, + children: list[ScheduleTreeNode], + parent: ScheduleTreeScope | None = None, + state: SDFGState | None = None) -> None: + super().__init__(node=node, state=state, children=children, parent=parent) def as_string(self, indent: int = 0): node: nodes.ConsumeEntry = self.node @@ -265,8 +714,8 @@ def as_string(self, indent: int = 0): @dataclass class TaskletNode(ScheduleTreeNode): node: nodes.Tasklet - in_memlets: Dict[str, Memlet] - out_memlets: Dict[str, Memlet] + in_memlets: dict[str, Memlet] + out_memlets: dict[str, Memlet] def as_string(self, indent: int = 0): in_memlets = ', '.join(f'{v}' for v in self.in_memlets.values()) @@ -275,12 +724,18 @@ def as_string(self, indent: int = 0): return indent * INDENTATION + f'tasklet({in_memlets})' return indent * INDENTATION + f'{out_memlets} = tasklet({in_memlets})' + def input_memlets(self, root: ScheduleTreeRoot | None = None, **kwargs) -> MemletSet: + return MemletSet(self.in_memlets.values()) + + def output_memlets(self, root: ScheduleTreeRoot | None = None, **kwargs) -> MemletSet: + return MemletSet(self.out_memlets.values()) + @dataclass class LibraryCall(ScheduleTreeNode): node: nodes.LibraryNode - in_memlets: Union[Dict[str, Memlet], Set[Memlet]] - out_memlets: Union[Dict[str, Memlet], Set[Memlet]] + in_memlets: dict[str, Memlet] | MemletSet + out_memlets: dict[str, Memlet] | MemletSet def as_string(self, indent: int = 0): if isinstance(self.in_memlets, set): @@ -297,6 +752,16 @@ def as_string(self, indent: int = 0): if v.owner not in {nodes.Node, nodes.CodeNode, nodes.LibraryNode}) return indent * INDENTATION + f'{out_memlets} = library {libname}[{own_properties}]({in_memlets})' + def input_memlets(self, root: ScheduleTreeRoot | None = None, **kwargs) -> MemletSet: + if isinstance(self.in_memlets, set): + return MemletSet(self.in_memlets) + return MemletSet(self.in_memlets.values()) + + def output_memlets(self, root: ScheduleTreeRoot | None = None, **kwargs) -> MemletSet: + if isinstance(self.out_memlets, set): + return MemletSet(self.out_memlets) + return MemletSet(self.out_memlets.values()) + @dataclass class CopyNode(ScheduleTreeNode): @@ -315,6 +780,16 @@ def as_string(self, indent: int = 0): return indent * INDENTATION + f'{self.target}{offset} = copy {self.memlet.data}[{self.memlet.subset}]{wcr}' + def input_memlets(self, root: ScheduleTreeRoot | None = None, **kwargs) -> MemletSet: + return MemletSet({self.memlet}) + + def output_memlets(self, root: ScheduleTreeRoot | None = None, **kwargs) -> MemletSet: + root = root if root is not None else self.get_root() + if self.memlet.other_subset is not None: + return MemletSet({Memlet(data=self.target, subset=self.memlet.other_subset, wcr=self.memlet.wcr)}) + + return MemletSet({Memlet.from_array(self.target, root.containers[self.target], self.memlet.wcr)}) + @dataclass class DynScopeCopyNode(ScheduleTreeNode): @@ -327,6 +802,12 @@ class DynScopeCopyNode(ScheduleTreeNode): def as_string(self, indent: int = 0): return indent * INDENTATION + f'{self.target} = dscopy {self.memlet.data}[{self.memlet.subset}]' + def input_memlets(self, root: ScheduleTreeRoot | None = None, **kwargs) -> MemletSet: + return MemletSet({self.memlet}) + + def output_memlets(self, root: ScheduleTreeRoot | None = None, **kwargs) -> MemletSet: + return MemletSet() + @dataclass class ViewNode(ScheduleTreeNode): @@ -339,6 +820,12 @@ class ViewNode(ScheduleTreeNode): def as_string(self, indent: int = 0): return indent * INDENTATION + f'{self.target} = view {self.memlet} as {self.view_desc.shape}' + def input_memlets(self, root: ScheduleTreeRoot | None = None, **kwargs) -> MemletSet: + return MemletSet({self.memlet}) + + def output_memlets(self, root: ScheduleTreeRoot | None = None, **kwargs) -> MemletSet: + return MemletSet({Memlet.from_array(self.target, self.view_desc)}) + @dataclass class NView(ViewNode): @@ -350,6 +837,25 @@ def as_string(self, indent: int = 0): return indent * INDENTATION + f'{self.target} = nview {self.memlet} as {self.view_desc.shape}' +@dataclass +class NViewEnd(ScheduleTreeNode): + """ + Artificial node to denote the scope end of the associated Nested SDFG view node. + """ + + target: str + """Target name of the associated NView container.""" + + def as_string(self, indent: int = 0): + return indent * INDENTATION + f"end nview {self.target}" + + def input_memlets(self, root: ScheduleTreeRoot | None = None, **kwargs) -> MemletSet: + return MemletSet() + + def output_memlets(self, root: ScheduleTreeRoot | None = None, **kwargs) -> MemletSet: + return MemletSet() + + @dataclass class RefSetNode(ScheduleTreeNode): """ @@ -357,7 +863,7 @@ class RefSetNode(ScheduleTreeNode): """ target: str memlet: Memlet - src_desc: Union[data.Data, nodes.CodeNode] + src_desc: data.Data | nodes.CodeNode ref_desc: data.Data def as_string(self, indent: int = 0): @@ -365,46 +871,71 @@ def as_string(self, indent: int = 0): return indent * INDENTATION + f'{self.target} = refset from {type(self.src_desc).__name__.lower()}' return indent * INDENTATION + f'{self.target} = refset to {self.memlet}' + def input_memlets(self, root: ScheduleTreeRoot | None = None, **kwargs) -> MemletSet: + return MemletSet({self.memlet}) + + def output_memlets(self, root: ScheduleTreeRoot | None = None, **kwargs) -> MemletSet: + return MemletSet({Memlet.from_array(self.target, self.ref_desc)}) + + +@dataclass +class StateBoundaryNode(ScheduleTreeNode): + """ + A node that represents a state boundary (e.g., when a write-after-write is encountered). This node + is used only during conversion from a schedule tree to an SDFG. + """ + due_to_control_flow: bool = False + + def as_string(self, indent: int = 0): + return indent * INDENTATION + 'state boundary' + + def input_memlets(self, root: ScheduleTreeRoot | None = None, **kwargs) -> MemletSet: + return MemletSet() + + def output_memlets(self, root: ScheduleTreeRoot | None = None, **kwargs) -> MemletSet: + return MemletSet() + # Classes based on Python's AST NodeVisitor/NodeTransformer for schedule tree nodes class ScheduleNodeVisitor: - def visit(self, node: ScheduleTreeNode): + def visit(self, node: ScheduleTreeNode | list[ScheduleTreeNode], **kwargs: Any): """Visit a node.""" if isinstance(node, list): - return [self.visit(snode) for snode in node] + return [self.visit(snode, **kwargs) for snode in node] + if isinstance(node, ScheduleTreeScope) and hasattr(self, 'visit_scope'): - return self.visit_scope(node) + return self.visit_scope(node, **kwargs) # type: ignore method = 'visit_' + node.__class__.__name__ visitor = getattr(self, method, self.generic_visit) - return visitor(node) + return visitor(node, **kwargs) - def generic_visit(self, node: ScheduleTreeNode): + def generic_visit(self, node: ScheduleTreeNode, **kwargs: Any): if isinstance(node, ScheduleTreeScope): for child in node.children: - self.visit(child) + self.visit(child, **kwargs) class ScheduleNodeTransformer(ScheduleNodeVisitor): - def visit(self, node: ScheduleTreeNode): + def visit(self, node: ScheduleTreeNode | list[ScheduleTreeNode], **kwargs: Any): if isinstance(node, list): result = [] for snode in node: - new_node = self.visit(snode) + new_node = self.visit(snode, **kwargs) if new_node is not None: result.append(new_node) return result - return super().visit(node) + return super().visit(node, **kwargs) - def generic_visit(self, node: ScheduleTreeNode): + def generic_visit(self, node: ScheduleTreeNode, **kwargs: Any): new_values = [] if isinstance(node, ScheduleTreeScope): for value in node.children: if isinstance(value, ScheduleTreeNode): - value = self.visit(value) + value = self.visit(value, **kwargs) if value is None: continue elif not isinstance(value, ScheduleTreeNode): @@ -427,3 +958,96 @@ def validate_has_no_other_node_types(stree: ScheduleTreeScope) -> None: raise RuntimeError(f'Unsupported node type: {type(child).__name__}') if isinstance(child, ScheduleTreeScope): validate_has_no_other_node_types(child) + + +def validate_children_and_parents_align(stree: ScheduleTreeScope, *, root: bool = False) -> None: + """ + Validates the child/parent information of schedule tree scopes are consistent. + + Walks through all children of a scope and raises if the children's parent isn't + the scope. If `root` is true, we additionally check that the top-most scope is + of type `ScheduleTreeRoot`. + + :param stree: Schedule tree scope to be analyzed + :param root: If true, we raise if the top-most scope isn't of type `ScheduleTreeRoot`. + """ + if root and not isinstance(stree, ScheduleTreeRoot): + raise RuntimeError("Expected schedule tree root.") + + for child in stree.children: + if child.parent is not stree: + raise RuntimeError(f"Inconsistent parent/child relationship. child: {child}, parent: {stree}") + if isinstance(child, ScheduleTreeScope): + validate_children_and_parents_align(child) + + +def loop_variant( + loop: LoopRegion +) -> Literal['for'] | Literal['while'] | Literal['do-while'] | Literal['do-for-uncond-increment'] | Literal['do-for']: + if loop.update_statement and loop.init_statement and loop.loop_variable: + if loop.inverted: + if loop.update_before_condition: + return 'do-for-uncond-increment' + return 'do-for' + return 'for' + + if loop.inverted: + return 'do-while' + return 'while' + + +def _loop_range(loop: LoopRegion) -> tuple[symbolic.SymbolicType, symbolic.SymbolicType, symbolic.SymbolicType] | None: + """ + Derive loop range for well-formed `for`-loops. + + :param: loop The loop to be analyzed. + :return: If well formed, `(start, end, step)` where `end` is inclusive, otherwise `None`. + """ + + if loop_variant(loop) != "for" or loop.loop_variable is None: + # Loop range is only defined in for-loops + # and we need to know the loop variable. + return None + + # Avoid cyclic import + from dace.transformation.passes.analysis import loop_analysis + + # If loop information cannot be determined, we cannot derive loop range + start = loop_analysis.get_init_assignment(loop) + step = loop_analysis.get_loop_stride(loop) + end = _match_loop_condition(loop) + if start is None or step is None or end is None: + return None + + return (start, end, step) + + +def _match_loop_condition(loop: LoopRegion) -> symbolic.SymbolicType | None: + """ + Try to find the end of a for-loop by symbolically matching the loop condition. + + :return: loop end (inclusive) or `None` if matching failed. + """ + + condition = symbolic.pystr_to_symbolic(loop.loop_condition.as_string) + loop_symbol = symbolic.pystr_to_symbolic(loop.loop_variable) + a = sympy.Wild('a') + + match = condition.match(loop_symbol < a) + if match is not None: + return match[a] - 1 + + match = condition.match(loop_symbol <= a) + if match is not None: + return match[a] + + match = condition.match(loop_symbol >= a) + if match is not None: + return match[a] + + match = condition.match(loop_symbol > a) + if match is not None: + return match[a] + 1 + + # Matching failed - we can't derive end of loop + return None diff --git a/dace/sdfg/analysis/vector_inference.py b/dace/sdfg/analysis/vector_inference.py index 6f42fb7729..c70135859e 100644 --- a/dace/sdfg/analysis/vector_inference.py +++ b/dace/sdfg/analysis/vector_inference.py @@ -1,18 +1,16 @@ -from enum import Enum, Flag +from enum import Flag from networkx import DiGraph from dace.memlet import Memlet -from collections import defaultdict from dace.sdfg.utils import dfs_topological_sort -from dace.sdfg.graph import MultiConnectorEdge, SubgraphView +from dace.sdfg.graph import MultiConnectorEdge import dace from dace import SDFG, SDFGState import dace.sdfg.nodes as nodes -from collections import defaultdict import dace.transformation.dataflow.sve.infer_types as infer_types import dace.dtypes as dtypes import dace.data as data -from typing import * import dace.symbolic as symbolic +from typing import Dict, Set, Tuple, Union, DefaultDict class VectorInferenceFlags(Flag): diff --git a/dace/sdfg/graph.py b/dace/sdfg/graph.py index f107e2caac..6993fffada 100644 --- a/dace/sdfg/graph.py +++ b/dace/sdfg/graph.py @@ -8,6 +8,7 @@ from dace.dtypes import deduplicate import dace.serialize from typing import Any, Callable, Generic, Iterable, List, Optional, Sequence, TypeVar, Union +from ordered_set import OrderedSet class NodeNotFoundError(Exception): @@ -31,6 +32,9 @@ def __init__(self, src, dst, data: T): self._dst = dst self._data: T = data + def id(self): + return (self._src.id, self._dst.id) + @property def src(self): return self._src @@ -215,7 +219,7 @@ def __getitem__(self, node: NodeT) -> Iterable[NodeT]: def all_edges(self, *nodes: NodeT) -> Iterable[Edge[EdgeT]]: """Returns an iterable to incoming and outgoing Edge objects.""" - result = set() + result = OrderedSet() for node in nodes: result.update(self.in_edges(node)) result.update(self.out_edges(node)) diff --git a/dace/sdfg/infer_types.py b/dace/sdfg/infer_types.py index cead874a20..02f5ae87d6 100644 --- a/dace/sdfg/infer_types.py +++ b/dace/sdfg/infer_types.py @@ -26,7 +26,7 @@ def infer_out_connector_type(sdfg: SDFG, state: SDFGState, node: nodes.CodeNode, e = next(state.out_edges_by_connector(node, cname)) if cname is None: return None - scalar = (e.data.subset and e.data.subset.num_elements() == 1 + scalar = (bool(e.data.subset) and e.data.subset.num_elements() == 1 and (not e.data.dynamic or (e.data.dynamic and e.data.wcr is not None))) if e.data.data is not None: allocated_as_scalar = (sdfg.arrays[e.data.data].storage is not dtypes.StorageType.GPU_Global) @@ -64,7 +64,7 @@ def infer_connector_types(sdfg: SDFG): cname = e.dst_conn if cname is None: continue - scalar = (e.data.subset and e.data.subset.num_elements() == 1) + scalar = bool(e.data.subset) and e.data.subset.num_elements() == 1 if e.data.data is not None: allocated_as_scalar = (sdfg.arrays[e.data.data].storage is not dtypes.StorageType.GPU_Global) else: @@ -143,6 +143,7 @@ def set_default_schedule_and_storage_types(scope: Union[SDFG, SDFGState, nodes.E ``dtypes.DEFAULT_TOPLEVEL_SCHEDULE``. May raise ``InvalidSDFGNodeError`` if a default scope is ambiguous based on surrounding storage types. + :param scope: The SDFG, state, or scope to infer. :param parent_schedules: A list of ScheduleType elements representing an ordered list of schedules, from the global schedule diff --git a/dace/sdfg/memlet_utils.py b/dace/sdfg/memlet_utils.py index 2c91650314..5bbfbdc12d 100644 --- a/dace/sdfg/memlet_utils.py +++ b/dace/sdfg/memlet_utils.py @@ -1,13 +1,15 @@ -# Copyright 2019-2023 ETH Zurich and the DaCe authors. All rights reserved. - +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. import ast -import itertools +from collections import defaultdict +import copy from dace import data, Memlet, subsets, symbolic, dtypes +from dace.frontend.python import memlet_parser from dace.sdfg import SDFGState, SDFG, nodes, utils as sdutil from dace.sdfg.scope import is_devicelevel_gpu from dace.sdfg.graph import MultiConnectorEdge from dace.frontend.python import memlet_parser -from typing import Callable, Dict, Optional, Set, Union, Tuple +import itertools +from typing import Callable, Dict, Iterable, Optional, Set, TypeVar, Tuple, Union, Generator, Any class MemletReplacer(ast.NodeTransformer): @@ -83,6 +85,188 @@ def visit_Subscript(self, node: ast.Subscript): return self.generic_visit(node) +class MemletSet(Set[Memlet]): + """ + Implements a set of memlets that considers subsets that intersect or are covered by its other memlets. + Set updates and unions also perform unions on the contained memlet subsets. + """ + + def __init__(self, iterable: Optional[Iterable[Memlet]] = None, *, intersection_is_contained: bool = True) -> None: + """ + Initializes a memlet set. + + :param iterable: An optional iterable of memlets to initialize the set with. + :param intersection_is_contained: Whether the check ``m in memlet_set`` should return True if the memlet + only intersects with the contents of the set. If False, only completely + covered subsets would return True. + """ + self.internal_set: Dict[str, Set[Memlet]] = {} + self.intersection_is_contained = intersection_is_contained + if iterable is not None: + self.update(iterable) + + def __iter__(self) -> Generator[Memlet, Any, None]: + for subset in self.internal_set.values(): + yield from subset + + def __len__(self) -> int: + return len(self.internal_set) + + def update(self, *iterable: Iterable[Memlet]) -> None: + """ + Updates set of memlets via union of existing ranges. + """ + if len(iterable) == 0: + return + if len(iterable) > 1: + for i in iterable: + self.update(i) + return + + to_update, = iterable + for elem in to_update: + self.add(elem) + + def add(self, elem: Memlet) -> None: + """ + Adds a memlet to the set, potentially performing a union of existing ranges. + """ + if elem.data not in self.internal_set: + self.internal_set[elem.data] = {elem} + return + + # Memlet is in set, either perform a union (if possible) or add to internal set + # TODO(later): Consider other_subset as well + for existing_memlet in self.internal_set[elem.data]: + try: + if subsets.intersects(existing_memlet.subset, elem.subset) == True: # Definitely intersects + if existing_memlet.subset.covers(elem.subset): + break # Nothing to do + + # Create a new union memlet + self.internal_set[elem.data].remove(existing_memlet) + new_memlet = copy.deepcopy(existing_memlet) + new_memlet.subset = subsets.union(existing_memlet.subset, elem.subset) + self.internal_set[elem.data].add(new_memlet) + break + except TypeError: # Indeterminate + pass + else: # all intersections were False or indeterminate (may or does not intersect with existing memlets) + self.internal_set[elem.data].add(elem) + + def __contains__(self, elem: Memlet) -> bool: + """ + Returns True iff the memlet or a range superset thereof exists in this set. + """ + if elem.data not in self.internal_set: + return False + for existing_memlet in self.internal_set[elem.data]: + if existing_memlet.subset.covers(elem.subset): + return True + if self.intersection_is_contained: + try: + if subsets.intersects(existing_memlet.subset, elem.subset) == False: + continue + else: # May intersect or indeterminate + return True + except TypeError: + return True + + return False + + def union(self, *s: Iterable[Memlet]) -> 'MemletSet': + """ + Performs a set-union (with memlet union) over the given sets of memlets. + + :return: New memlet set containing the union of this set and the inputs. + """ + newset = MemletSet(self) + newset.update(s) + return newset + + +T = TypeVar('T') + + +class MemletDict(Dict[Memlet, T]): + """ + Implements a dictionary with memlet keys that considers subsets that intersect + or are covered by its other memlets. + """ + + def __init__(self, **kwargs) -> None: + self.internal_dict: Dict[str, Dict[Memlet, T]] = defaultdict(dict) + self.covers_cache: Dict[Tuple, bool] = defaultdict() + + if kwargs: + self.update(kwargs) + + def __len__(self) -> int: + return len(self.internal_dict) + + def _getkey(self, elem: Memlet) -> Optional[Memlet]: + """ + Returns the corresponding key (exact, covered, intersecting, or indeterminately intersecting memlet) + if it exists in the dictionary, or None if it does not. + """ + if elem.data not in self.internal_dict: + return None + + for existing_memlet in self.internal_dict[elem.data]: + key = (existing_memlet.subset, elem.subset) + is_covered = self.covers_cache.get(key, None) + if is_covered is None: + is_covered = existing_memlet.subset.covers(elem.subset) + self.covers_cache[key] = is_covered + if is_covered: + return existing_memlet + + try: + if subsets.intersects(existing_memlet.subset, elem.subset) == False: # Definitely does not intersect + continue + except TypeError: + pass + + # May or will intersect + return existing_memlet + + return None + + def _setkey(self, key: Memlet, value: T) -> None: + self.internal_dict[key.data][key] = value + + def clear(self) -> None: + self.internal_dict.clear() + + def update(self, mapping: Dict[Memlet, T]) -> None: + for key, value in mapping.items(): + actual_key = self._getkey(key) + if actual_key is None: + self._setkey(key, value) + else: + self._setkey(actual_key, value) + + def __contains__(self, elem: Memlet) -> bool: + """ + Returns True iff the memlet or a range superset thereof exists in this dictionary. + """ + return self._getkey(elem) is not None + + def __getitem__(self, key: Memlet) -> T: + actual_key = self._getkey(key) + if actual_key is None: + raise KeyError(key) + + return self.internal_dict[key.data][actual_key] + + def __setitem__(self, key: Memlet, value: T) -> None: + actual_key = self._getkey(key) + if actual_key is None: + self._setkey(key, value) + else: + self._setkey(actual_key, value) + + def memlet_to_map( edge: MultiConnectorEdge, state: SDFGState, diff --git a/dace/sdfg/nodes.py b/dace/sdfg/nodes.py index 1382ce273d..4a19130dd8 100644 --- a/dace/sdfg/nodes.py +++ b/dace/sdfg/nodes.py @@ -5,20 +5,20 @@ import ast from copy import deepcopy as dcpy from collections.abc import KeysView +import contextvars import dace import itertools import dace.serialize +import sympy as sp from typing import Any, Dict, Optional, Set, Union from dace.config import Config from dace.sdfg import graph -from dace.frontend.python.astutils import unparse, rname -from dace.properties import (EnumProperty, Property, CodeProperty, LambdaProperty, RangeProperty, DebugInfoProperty, - SetProperty, make_properties, indirect_properties, DataProperty, SymbolicProperty, - ListProperty, SDFGReferenceProperty, DictProperty, LibraryImplementationProperty, - CodeBlock) -from dace.frontend.operations import detect_reduction_type +from dace.frontend.python.astutils import rname +from dace.properties import (EnumProperty, Property, CodeProperty, RangeProperty, DebugInfoProperty, SetProperty, + make_properties, indirect_properties, DataProperty, SymbolicProperty, ListProperty, + SDFGReferenceProperty, DictProperty, LibraryImplementationProperty, CodeBlock) from dace.symbolic import issymbolic, pystr_to_symbolic -from dace import data, subsets as sbs, dtypes +from dace import subsets as sbs, dtypes from dace.sdfg import tasklet_validation as tval from dace.sdfg.type_inference import infer_types, infer_expr_type import pydoc @@ -26,6 +26,35 @@ # ----------------------------------------------------------------------------- +# Global context variable to store the node ID counter +_node_id_counter: contextvars.ContextVar[int] = contextvars.ContextVar('_node_id_counter', default=0) + + +def _get_next_node_id() -> int: + """Get the next node ID and increment the counter.""" + current = _node_id_counter.get() + _node_id_counter.set(current + 1) + return current + + +# TODO: check if the edges need an id. If source and dest, are equal they should be interchangable right? + + +class reset_node_id_counter: + """Context manager that resets the node ID counter to zero and restores the old value afterwards.""" + + def __init__(self): + self._old_value = None + + def __enter__(self): + self._old_value = _node_id_counter.get() + _node_id_counter.set(0) + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + _node_id_counter.set(self._old_value) + return False + @make_properties class Node(object): @@ -38,6 +67,7 @@ class Node(object): value_type=dtypes.typeclass, desc="A set of output connectors for this node.") guid = Property(dtype=str, allow_none=False) + id = Property(dtype=int, allow_none=False) def __init__(self, in_connectors=None, out_connectors=None): # Convert connectors to typed connectors with autodetect type @@ -50,6 +80,7 @@ def __init__(self, in_connectors=None, out_connectors=None): self.out_connectors = out_connectors or {} self.guid = graph.generate_element_id(self) + self.id = _get_next_node_id() def __str__(self): if hasattr(self, 'label'): @@ -62,9 +93,14 @@ def __deepcopy__(self, memo): result = cls.__new__(cls) memo[id(self)] = result for k, v in self.__dict__.items(): - if k == 'guid': # Skip ID + if k == 'guid': # Skip GUID + continue + if k == '_id': # Skip ID, will be assigned new one continue setattr(result, k, dcpy(v, memo)) + # Assign new GUID and ID + result.guid = graph.generate_element_id(result) + result.id = _get_next_node_id() return result def validate(self, sdfg, state): @@ -101,6 +137,7 @@ def to_json(self, parent): "type": typestr, "label": labelstr, "attributes": dace.serialize.all_properties_to_json(self), + # TODO(tehrengruber): This id looks very similar to the ID I introduced. "id": parent.node_id(self), "scope_entry": scope_entry_node, "scope_exit": scope_exit_node @@ -277,7 +314,7 @@ class AccessNode(Node): """ A node that accesses data in the SDFG. Denoted by a circular shape. """ setzero = Property(dtype=bool, desc="Initialize to zero", default=False) - debuginfo = DebugInfoProperty() + debuginfo = DebugInfoProperty(allow_none=True) data = DataProperty(desc="Data (array, stream, scalar) to access") instrument = EnumProperty(dtype=dtypes.DataInstrumentationType, @@ -312,6 +349,7 @@ def __deepcopy__(self, memo): node._debuginfo = dcpy(self._debuginfo, memo=memo) node._guid = graph.generate_element_id(node) + node._id = _get_next_node_id() return node @@ -399,7 +437,7 @@ class Tasklet(CodeNode): default=CodeBlock("", dtypes.Language.CPP)) code_exit = CodeProperty(desc="Extra code that is called on DaCe runtime cleanup", default=CodeBlock("", dtypes.Language.CPP)) - debuginfo = DebugInfoProperty() + debuginfo = DebugInfoProperty(allow_none=True) instrument = EnumProperty(dtype=dtypes.InstrumentationType, desc="Measure execution statistics with given method", @@ -598,10 +636,10 @@ class NestedSDFG(CodeNode): allow_none=True, desc='Path to a file containing the SDFG for this nested SDFG') symbol_mapping = DictProperty(key_type=str, - value_type=dace.symbolic.pystr_to_symbolic, + value_type=sp.Basic, desc="Mapping between internal symbols and their values, expressed as " "symbolic expressions") - debuginfo = DebugInfoProperty() + debuginfo = DebugInfoProperty(allow_none=True) is_collapsed = Property(dtype=bool, desc="Show this node/scope/state as collapsed", default=False) instrument = EnumProperty(dtype=dtypes.InstrumentationType, @@ -644,10 +682,13 @@ def __deepcopy__(self, memo): result = cls.__new__(cls) memo[id(self)] = result for k, v in self.__dict__.items(): - # Skip GUID. - if k in ('guid', ): + # Skip GUID and ID. + if k in ('guid', '_id'): continue setattr(result, k, dcpy(v, memo)) + # Assign new GUID and ID + result.guid = graph.generate_element_id(result) + result.id = _get_next_node_id() if result._sdfg is not None: result._sdfg.parent_nsdfg_node = result return result @@ -1008,7 +1049,7 @@ class Map(object): desc="How much iterations should be unrolled." " To prevent unrolling, set this value to 1.") collapse = Property(dtype=int, default=1, desc="How many dimensions to collapse into the parallel range") - debuginfo = DebugInfoProperty() + debuginfo = DebugInfoProperty(allow_none=True) is_collapsed = Property(dtype=bool, desc="Show this node/scope/state as collapsed", default=False) instrument = EnumProperty(dtype=dtypes.InstrumentationType, @@ -1041,6 +1082,11 @@ class Map(object): "(including tuples) sets it explicitly.", serialize_if=lambda m: m.schedule in dtypes.GPU_SCHEDULES) + gpu_min_warps_per_eu = Property(dtype=int, + default=0, + desc="Minimum number of warps per execution unit for GPU kernel", + serialize_if=lambda m: m.schedule in dtypes.GPU_SCHEDULES) + gpu_maxnreg = Property(dtype=int, default=0, desc="Maximum number of registers per thread for GPU kernel", @@ -1094,10 +1140,17 @@ def validate(self, sdfg, state, node): warnings.warn(f'The iteration range of Map {self.label} is {self.range}, which contains a zero' ' or negative sized range, which is allowed but not recommended.' ' The Map will be turned into a no-ops.') - if any((inc <= 0) == True for (_, _, inc) in self.range): - # Should this be turned into an error? - warnings.warn(f'An increment of Map {self.label} was negative, which is allowerd' - ' but probably not useful.') + # A Map is an unordered, ascending iteration domain; the backends + # emit an ascending loop, so a negative step has no valid lowering. + # An empty positive-step range iterates zero times and is left to + # the size warning above. A symbolic step is checked at runtime by + # an assertion in the generated (debug) code. + for (_, _, inc) in self.range: + if (inc < 0) == True: + raise ValueError(f'Map {self.label} has a negative step ({inc}) in range ' + f'{self.range}. Maps must use a positive step; rewrite the ' + 'iteration ascending (canonicalization normalizes loops to a ' + 'positive unit step).') def get_param_num(self): """ Returns the number of map dimension parameters/symbols. """ @@ -1260,7 +1313,7 @@ class Consume(object): condition = CodeProperty(desc="Quiescence condition", allow_none=True, default=None) schedule = EnumProperty(dtype=dtypes.ScheduleType, desc="Consume schedule", default=dtypes.ScheduleType.Default) chunksize = Property(dtype=int, desc="Maximal size of elements to consume at a time", default=1) - debuginfo = DebugInfoProperty() + debuginfo = DebugInfoProperty(allow_none=True) is_collapsed = Property(dtype=bool, desc="Show this node/scope/state as collapsed", default=False) instrument = EnumProperty(dtype=dtypes.InstrumentationType, @@ -1330,7 +1383,7 @@ class LibraryNode(CodeNode): desc="If set, determines the default device mapping of " "the node upon expansion, if expanded to a nested SDFG.", default=dtypes.ScheduleType.Default) - debuginfo = DebugInfoProperty() + debuginfo = DebugInfoProperty(allow_none=True) def __init__(self, name, *args, schedule=None, **kwargs): super().__init__(*args, **kwargs) @@ -1377,15 +1430,17 @@ def expand(self, state_or_sdfg, state_or_impl=None, **kwargs) -> str: node. This method supports two interfaces: - 1. New interface: expand(state, implementation=None, **kwargs) - 2. Old interface: expand(sdfg, state, **kwargs) [for backward compatibility] + + 1. New interface: ``expand(state, implementation=None, **kwargs)`` + 2. Old interface: ``expand(sdfg, state, **kwargs)`` [for backward compatibility] :param state_or_sdfg: Either a ControlFlowBlock (new interface) or SDFG (old interface) :param state_or_impl: Either implementation name (new interface) or SDFGState (old interface) :param kwargs: Additional expansion arguments :return: the name of the expanded implementation - Examples: + Examples:: + # New interface (recommended): result = node.expand(state, 'pure') diff --git a/dace/sdfg/performance_evaluation/work_depth.py b/dace/sdfg/performance_evaluation/work_depth.py index ecb447ba9d..35b60701c7 100644 --- a/dace/sdfg/performance_evaluation/work_depth.py +++ b/dace/sdfg/performance_evaluation/work_depth.py @@ -1,4 +1,4 @@ -# Copyright 2019-2023 ETH Zurich and the DaCe authors. All rights reserved. +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. """ Work depth analysis for any input SDFG. Can be used with the DaCe VS Code extension or from command line as a Python script. """ @@ -12,7 +12,8 @@ import sympy as sp from copy import deepcopy from dace.libraries.blas import MatMul -from dace.libraries.standard import Reduce, Transpose +from dace.libraries.standard import Reduce +from dace.libraries.linalg import Transpose from dace.symbolic import pystr_to_symbolic import ast import astunparse diff --git a/dace/sdfg/propagation.py b/dace/sdfg/propagation.py index a88ef0a367..eb0c5a99f4 100644 --- a/dace/sdfg/propagation.py +++ b/dace/sdfg/propagation.py @@ -1,4 +1,4 @@ -# Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. """ Functionality relating to Memlet propagation (deducing external memlets from internal memory accesses and scope ranges). @@ -963,6 +963,285 @@ def propagate_states(sdfg: 'SDFG', concretize_dynamic_unbounded: bool = False) - sdfg.remove_node(temp_exit_state) +def _make_border_memlets(connectors, as_lists: bool = False): + """ + Create the connector map used while aggregating nested border memlets. + + The returned structure mirrors the nested-SDFG connector layout, + containing separate input and output connector dictionaries. + + :param connectors: A mapping with ``in`` and ``out`` entries whose + values are the connector collections to initialize. + :param as_lists: If True, initializes each connector entry with an empty + list of candidate memlets. Otherwise, initializes each + entry with ``None`` for incremental merged results. + :return: A dictionary of the form ``{'in': {...}, 'out': {...}}`` with + initialized connector entries. + """ + border_memlets = { + 'in': {}, + 'out': {}, + } + for direction in border_memlets: + for connector in connectors[direction]: + border_memlets[direction][connector] = [] if as_lists else None + return border_memlets + + +def _collect_state_border_memlet_candidates(state: 'SDFGState', border_memlets) -> None: + """ + Collect state-local memlets attached to nested border (i.e., not internal) access nodes. + + For nested SDFG inputs this inspects outgoing edges of the corresponding + access node, while for outputs it inspects incoming edges. The collected + memlets are left unpropagated so that the caller can later apply the + appropriate control-flow semantics. + + :param state: The state whose access-node edges should be inspected. + :param border_memlets: A candidate memlet mapping, typically created with + ``_make_border_memlets(..., as_lists=True)``, which + is updated in place. + :note: ``border_memlets`` mapping is updated in-place. + """ + for node in state.data_nodes(): + for direction in border_memlets: + if node.label not in border_memlets[direction]: + continue + + edges = state.out_edges(node) if direction == 'in' else state.in_edges(node) + border_memlets[direction][node.label].extend(edge.data for edge in edges) + + +def _append_border_memlet_candidates(border_memlets, propagated_memlets) -> None: + """ + Append propagated child-region memlets to a candidate memlet map. + + :param border_memlets: The candidate memlet mapping to append to. Each + connector entry is expected to hold a list. + :param propagated_memlets: A propagated connector mapping whose non-None + memlets should be appended as candidates. + :note: ``border_memlets`` mapping is updated in-place. + """ + for direction in border_memlets: + for connector in border_memlets[direction]: + memlet = propagated_memlets[direction][connector] + if memlet is not None: + border_memlets[direction][connector].append(memlet) + + +def _merge_border_memlet(existing: Memlet, incoming: Memlet, array: data.Data) -> Memlet: + """ + Merge two border memlets using union aggregation semantics. + + This merge rule sums volumes, unions subsets, and preserves dynamic + unbounded behavior. It is the default aggregation mode for states and + control-flow regions where all collected contributions conservatively count + toward the same border memlet. + + :param existing: The accumulated border memlet so far, or None if no + memlet has been accumulated yet. + :param incoming: The newly propagated memlet to merge into the result. + :param array: The array descriptor used to fall back to a full-array subset + when subset unioning cannot preserve a more precise result. + :return: The merged memlet. + """ + if incoming is None: + return existing + if existing is None: + return copy.deepcopy(incoming) + + result = copy.deepcopy(existing) + + if incoming.wcr is not None: + if result.wcr is not None and result.wcr != incoming.wcr: + warnings.warn('Memlet appears with more than one type of write-conflict resolution.') + result.wcr = incoming.wcr + + if result.dynamic and result.volume == 0: + pass + elif incoming.dynamic and incoming.volume == 0: + result.dynamic = True + result.volume = 0 + else: + result.volume = simplify(result.volume + incoming.volume) + result.dynamic = (result.dynamic or incoming.dynamic) + + if incoming.subset is not None: + if result.subset is not None: + if result.subset.dims() != incoming.subset.dims(): + raise ValueError('Cannot merge subset ranges of unequal dimension!') + result.subset = subsets.union(result.subset, incoming.subset) + if result.subset is None: + result.subset = subsets.Range.from_array(array) + else: + result.subset = copy.deepcopy(incoming.subset) + + return result + + +def _merge_border_memlet_upper_bound(existing: Memlet, incoming: Memlet, array: data.Data) -> Memlet: + """ + Merge two border memlets using upper-bound conditional semantics. + + Conditional regions may execute only one branch per traversal. Their + propagated subset is therefore the union of all branch subsets, while the + propagated volume is the maximum branch volume instead of the sum. + + :param existing: The accumulated border memlet so far, or None if no + memlet has been accumulated yet. + :param incoming: The newly propagated branch memlet to merge into the + result. + :param array: The array descriptor used to fall back to a full-array subset + when subset unioning cannot preserve a more precise result. + :return: The merged memlet that conservatively upper-bounds the branches. + """ + if incoming is None: + return existing + if existing is None: + return copy.deepcopy(incoming) + + result = copy.deepcopy(existing) + + if incoming.wcr is not None: + if result.wcr is not None and result.wcr != incoming.wcr: + warnings.warn('Memlet appears with more than one type of write-conflict resolution.') + result.wcr = incoming.wcr + + if result.dynamic and result.volume == 0: + pass + elif incoming.dynamic and incoming.volume == 0: + result.dynamic = True + result.volume = 0 + else: + result.volume = simplify(sympy.Max(result.volume, incoming.volume)) + result.dynamic = (result.dynamic or incoming.dynamic) + + if incoming.subset is not None: + if result.subset is not None: + if result.subset.dims() != incoming.subset.dims(): + raise ValueError('Cannot merge subset ranges of unequal dimension!') + result.subset = subsets.union(result.subset, incoming.subset) + if result.subset is None: + result.subset = subsets.Range.from_array(array) + else: + result.subset = copy.deepcopy(incoming.subset) + + return result + + +def _propagate_border_memlet_candidates(candidates, + arrays, + direction: str, + connector: str, + params=None, + rng=None, + scale_by_range: bool = False) -> Memlet: + """ + Propagate candidate memlets through a symbolic iteration range. + + The helper collects all candidate memlets for a single connector, applies + subset propagation through the supplied symbolic range, and optionally keeps + or removes the usual range-size volume scaling. + + :param candidates: A candidate memlet mapping whose connector entries hold + lists of memlets. + :param arrays: The array descriptor mapping of the containing SDFG. + :param direction: Either ``in`` or ``out``, indicating which connector + direction should be propagated. + :param connector: The connector name whose candidate memlets should be + propagated. + :param params: Optional iteration variable names for the propagation range. + If omitted together with ``rng``, a dummy singleton range is + used. + :param rng: Optional symbolic iteration range to propagate through. + :param scale_by_range: If True, preserves the default ``propagate_subset`` + volume scaling by the size of ``rng``. If False, + keeps only the sum of the candidate memlet volumes. + :return: The propagated memlet for the connector, or ``None`` if no candidate + memlets exist. + """ + memlets = candidates[direction][connector] + if len(memlets) == 0: + return None + + if params is None or rng is None: + params = ['__dace_dummy'] + rng = subsets.Range([(0, 0, 1)]) + + array_name = next((memlet.data for memlet in memlets if memlet.data is not None), connector) + array = arrays[array_name] + propagated = propagate_subset(memlets, array, params, rng, use_dst=(direction == 'out')) + + if any(memlet.dynamic and memlet.volume == 0 for memlet in memlets): + propagated.dynamic = True + propagated.volume = 0 + elif not scale_by_range: + propagated.volume = simplify(sum(memlet.volume for memlet in memlets)) + propagated.dynamic = any(memlet.dynamic for memlet in memlets) + + propagated.other_subset = None + return propagated + + +def _propagate_state_border_memlets(state: 'SDFGState', border_memlets, arrays) -> None: + """ + Propagate all connector-adjacent memlets contributed by one state. + + State-local subsets are propagated through the state's symbolic ranges, and + the resulting memlets are then scaled according to the state's execution + metadata before being merged into the nested-SDFG border memlets. + + :param state: The state whose connector-adjacent memlets should be + propagated. + :param border_memlets: The accumulated border memlet mapping to update. + :param arrays: The array descriptor mapping used for propagation and merge + fallback behavior. + :note: The ``border_memlets`` mapping is updated in-place. + """ + candidates = _make_border_memlets(border_memlets, as_lists=True) + _collect_state_border_memlet_candidates(state, candidates) + + params = [] + ranges = [] + for symbol in state.ranges: + params.append(symbol) + ranges.append(state.ranges[symbol][0]) + + if len(params) == 0 or len(ranges) == 0: + params = None + rng = None + else: + rng = subsets.Range(ranges) + + for direction in border_memlets: + for connector in border_memlets[direction]: + propagated = _propagate_border_memlet_candidates( + candidates, + arrays, + direction, + connector, + params=params, + rng=rng, + scale_by_range=False, + ) + if propagated is None: + continue + + if propagated.dynamic and propagated.volume == 0: + pass + elif state.dynamic_executions and state.executions == 0: + propagated.dynamic = True + propagated.volume = 0 + else: + propagated.volume = simplify( + sum(memlet.volume for memlet in candidates[direction][connector]) * state.executions) + propagated.dynamic = propagated.dynamic or state.dynamic_executions + + array_name = propagated.data if propagated.data is not None else connector + border_memlets[direction][connector] = _merge_border_memlet(border_memlets[direction][connector], + propagated, arrays[array_name]) + + def propagate_memlets_nested_sdfg(parent_sdfg: 'SDFG', parent_state: 'SDFGState', nsdfg_node: nodes.NestedSDFG): """ Propagate memlets out of a nested sdfg. @@ -974,109 +1253,27 @@ def propagate_memlets_nested_sdfg(parent_sdfg: 'SDFG', parent_state: 'SDFGState' """ # We import late to avoid cyclic imports here. from dace.transformation.helpers import unsqueeze_memlet + from dace.sdfg.state import AbstractControlFlowRegion, SDFGState # Build a map of connectors to associated 'border' memlets inside # the nested SDFG. This map will be populated with memlets once they # get propagated in the SDFG. - border_memlets = { - 'in': {}, - 'out': {}, - } - for connector in nsdfg_node.in_connectors: - border_memlets['in'][connector] = None - for connector in nsdfg_node.out_connectors: - border_memlets['out'][connector] = None + border_memlets = _make_border_memlets({ + 'in': nsdfg_node.in_connectors, + 'out': nsdfg_node.out_connectors, + }) sdfg = nsdfg_node.sdfg outer_symbols = parent_state.symbols_defined_at(nsdfg_node) - # For each state, go through all access nodes corresponding to any in- or - # out-connectors to and from this SDFG. Given those access nodes, collect - # the corresponding memlets and use them to calculate the memlet volume and - # subset corresponding to the outside memlet attached to that connector. - # This is passed out via `border_memlets` and propagated along from there. - for state in sdfg.states(): - for node in state.data_nodes(): - for direction in border_memlets: - if (node.label not in border_memlets[direction]): - continue - - memlet = border_memlets[direction][node.label] - - # Collect the edges to/from this access node, depending on the - # direction the connector leads in. - edges = [] - if direction == 'in': - edges = state.out_edges(node) - elif direction == 'out': - edges = state.in_edges(node) - - # Collect all memlets belonging to this access node, and - # accumulate the total volume between them. - memlets = [] - for edge in edges: - inside_memlet = edge.data - memlets.append(inside_memlet) - - if memlet is None: - # Use the first encountered memlet as a 'border' memlet - # and accumulate the sum on it. - memlet = Memlet(data=inside_memlet.data, volume=0) - memlet._is_data_src = True - border_memlets[direction][node.label] = memlet - - if inside_memlet.wcr is not None: - if (memlet.wcr is not None and memlet.wcr != inside_memlet.wcr): - warnings.warn('Memlet appears with more than one type of write-conflict resolution.') - memlet.wcr = inside_memlet.wcr - - if memlet.dynamic and memlet.volume == 0: - # Dynamic unbounded - this won't change. - continue - elif ((inside_memlet.dynamic and inside_memlet.volume == 0) - or (state.dynamic_executions and state.executions == 0)): - # At least one dynamic unbounded memlet means the sum - # must be dynamic unbounded. - memlet.dynamic = True - memlet.volume = 0 - else: - memlet.volume += (inside_memlet.volume * state.executions) - memlet.dynamic = (memlet.dynamic or inside_memlet.dynamic or state.dynamic_executions) - - # Given all of this access nodes' memlets, propagate the subset - # according to the state's variable ranges. - if len(memlets) > 0: - params = [] - ranges = [] - for symbol in state.ranges: - params.append(symbol) - ranges.append(state.ranges[symbol][0]) - - if len(params) == 0 or len(ranges) == 0: - params = ['__dace_dummy'] - ranges = [(0, 0, 1)] - - # Propagate the subset based on the direction this memlet is - # pointing. If we're accessing from an incoming connector, - # propagate the source subset, if we're going to an outgoing - # connector, propagate the destination subset. - use_dst = False - if direction == 'out': - use_dst = True - array = sdfg.arrays[node.label] - subset = propagate_subset(memlets, array, params, subsets.Range(ranges), use_dst=use_dst).subset - - # If the border memlet already has a set range, compute the - # union of the ranges to merge the subsets. - if memlet.subset is not None: - if memlet.subset.dims() != subset.dims(): - raise ValueError('Cannot merge subset ranges of unequal dimension!') - else: - memlet.subset = subsets.union(memlet.subset, subset) - if memlet.subset is None: - memlet.subset = subsets.Range.from_array(array) - else: - memlet.subset = subset + # Collect contributions from top-level CFG blocks. Plain states contribute + # directly, while control-flow regions aggregate their child blocks via + # their own propagate_memlets implementations. + for block in sdfg.nodes(): + if isinstance(block, SDFGState): + _propagate_state_border_memlets(block, border_memlets, sdfg.arrays) + elif isinstance(block, AbstractControlFlowRegion): + block.propagate_memlets(border_memlets) # Make sure any potential NSDFG symbol mapping is correctly reversed # when propagating out. @@ -1427,7 +1624,12 @@ def propagate_memlet(dfg_state, # Propagate subset if isinstance(entry_node, nodes.MapEntry): mapnode = entry_node.map - return propagate_subset(aggdata, arr, mapnode.params, mapnode.range, defined_vars, use_dst=use_dst) + return propagate_subset(aggdata, + arr, + mapnode.params, + mapnode.range, + defined_variables=defined_vars, + use_dst=use_dst) elif isinstance(entry_node, nodes.ConsumeEntry): # Nothing to analyze/propagate in consume @@ -1446,7 +1648,9 @@ def propagate_subset(memlets: List[Memlet], arr: data.Data, params: List[str], rng: subsets.Subset, + *, defined_variables: Set[symbolic.SymbolicType] = None, + undefined_variables: Set[symbolic.SymbolicType] = None, use_dst: bool = False) -> Memlet: """ Tries to propagate a list of memlets through a range (computes the image of the memlet function applied on an integer set of, e.g., a @@ -1459,8 +1663,12 @@ def propagate_subset(memlets: List[Memlet], range to propagate with. :param defined_variables: A set of symbols defined that will remain the same throughout propagation. If None, assumes - that all symbols outside of `params` have been - defined. + that all symbols outside of ``params``, except + for ``undefined_variables``, have been defined. + :param undefined_variables: A set of symbols that are explicitly considered + as not defined throughout propagation, such as + locals. Their existence will trigger propagating + the entire memlet. :param use_dst: Whether to propagate the memlets' dst subset or use the src instead, depending on propagation direction. :return: Memlet with propagated subset and volume. @@ -1474,6 +1682,13 @@ def propagate_subset(memlets: List[Memlet], defined_variables |= memlet.free_symbols defined_variables -= set(params) defined_variables = set(symbolic.pystr_to_symbolic(p) for p in defined_variables) + else: + defined_variables = set(defined_variables) + + if undefined_variables is not None: + defined_variables = defined_variables - set(symbolic.pystr_to_symbolic(p) for p in undefined_variables) + else: + undefined_variables = set() # Propagate subset variable_context = [defined_variables, [symbolic.pystr_to_symbolic(p) for p in params]] @@ -1504,18 +1719,33 @@ def propagate_subset(memlets: List[Memlet], tmp_subset = pattern.propagate(arr, [subset], rng) break else: - # No patterns found. Emit a warning and propagate the entire - # array whenever symbols are used + # No patterns found. Propagate the entire array whenever symbols are used. if Config.get_bool('debugprint'): print(f'Cannot find appropriate memlet pattern to propagate {subset} through {rng}') entire_array = subsets.Range.from_array(arr) paramset = set(map(str, params)) # Fill in the entire array only if one of the parameters appears in the - # free symbols list of the subset dimension - tmp_subset = subsets.Range([ - ea if any(set(map(str, _freesyms(sd))) & paramset for sd in s) else s - for s, ea in zip(subset, entire_array) - ]) + # free symbols list of the subset dimension or is undefined outside. + tmp_subset_rng = [] + for s, ea in zip(subset, entire_array): + if isinstance(subset, subsets.Indices): + fsyms = _freesyms(s) + fsyms_str = set(map(str, fsyms)) + contains_params = len(fsyms_str & paramset) != 0 + contains_undefs = len(fsyms & undefined_variables) != 0 + else: + contains_params = False + contains_undefs = False + for sdim in s: + fsyms = _freesyms(sdim) + fsyms_str = set(map(str, fsyms)) + contains_params |= len(fsyms_str & paramset) != 0 + contains_undefs |= len(fsyms & undefined_variables) != 0 + if contains_params or contains_undefs: + tmp_subset_rng.append(ea) + else: + tmp_subset_rng.append(s) + tmp_subset = subsets.Range(tmp_subset_rng) # Union edges as necessary if new_subset is None: diff --git a/dace/sdfg/replace.py b/dace/sdfg/replace.py index 04ceafb261..530ee2b7cc 100644 --- a/dace/sdfg/replace.py +++ b/dace/sdfg/replace.py @@ -1,4 +1,4 @@ -# Copyright 2019-2024 ETH Zurich and the DaCe authors. All rights reserved. +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. """ Contains functionality to perform find-and-replace of symbols in SDFGs. """ import re @@ -81,7 +81,7 @@ def replace_dict(subgraph: 'StateSubgraphView', desc = node.desc(state) # In case the AccessNode name was replaced in the sdfg.arrays but not in the SDFG itself # then we have to look for the replaced value in the sdfg.arrays - elif repl[node.data] in state.sdfg.arrays: + elif node.data in repl and repl[node.data] in state.sdfg.arrays: desc = state.sdfg.arrays[repl[node.data]] else: continue diff --git a/dace/sdfg/scope.py b/dace/sdfg/scope.py index cd139aaa17..f5a397d96b 100644 --- a/dace/sdfg/scope.py +++ b/dace/sdfg/scope.py @@ -8,6 +8,7 @@ from dace.config import Config from dace.sdfg import nodes as nd from dace.sdfg.state import StateSubgraphView +from ordered_set import OrderedSet ScopeDictType = Dict[nd.Node, List[nd.Node]] @@ -62,16 +63,16 @@ def _scope_subgraph(graph, entry_node, include_entry, include_exit) -> ScopeSubg raise TypeError("Received {}: should be dace.nodes.EntryNode".format(type(entry_node).__name__)) node_to_children = graph.scope_children() if include_exit: - children_nodes = set(node_to_children[entry_node]) + children_nodes = OrderedSet(node_to_children[entry_node]) else: - children_nodes = set(n for n in node_to_children[entry_node] if not isinstance(n, nd.ExitNode)) + children_nodes = OrderedSet(n for n in node_to_children[entry_node] if not isinstance(n, nd.ExitNode)) map_nodes = [node for node in children_nodes if isinstance(node, nd.EntryNode)] while len(map_nodes) > 0: next_map_nodes = [] # Traverse children map nodes for map_node in map_nodes: # Get child map subgraph (1 level) - more_nodes = set(node_to_children[map_node]) + more_nodes = OrderedSet(node_to_children[map_node]) # Unionize children_nodes with new nodes children_nodes |= more_nodes # Add nodes of the next level to next_map_nodes diff --git a/dace/sdfg/sdfg.py b/dace/sdfg/sdfg.py index d88effd594..e54a2fb366 100644 --- a/dace/sdfg/sdfg.py +++ b/dace/sdfg/sdfg.py @@ -13,12 +13,9 @@ import sys from typing import Any, AnyStr, Dict, List, Optional, Sequence, Set, Tuple, Type, TYPE_CHECKING, Union import warnings -import subprocess -import tempfile -import pickle import dace -from dace.sdfg.graph import generate_element_id +from dace.sdfg.graph import generate_element_id, SubgraphView import dace.serialize from dace import (data as dt, hooks, memlet as mm, subsets as sbs, dtypes, symbolic) from dace.sdfg.replace import replace_properties_dict @@ -26,9 +23,9 @@ from dace.config import Config from dace.frontend.python import astutils from dace.sdfg import nodes as nd -from dace.sdfg.state import ConditionalBlock, ControlFlowBlock, SDFGState, ControlFlowRegion +from dace.sdfg.state import ConditionalBlock, ControlFlowBlock, SDFGState, ControlFlowRegion, LoopRegion from dace.sdfg.type_inference import infer_expr_type -from dace.distr_types import ProcessGrid, SubArray, RedistrArray +from dace.data.distributed import ProcessGrid, SubArray, RedistrArray from dace.dtypes import validate_name from dace.properties import (DebugInfoProperty, EnumProperty, ListProperty, make_properties, Property, CodeProperty, TransformationHistProperty, OptionalSDFGReferenceProperty, DictProperty, CodeBlock) @@ -42,7 +39,7 @@ from dace.codegen.instrumentation.report import InstrumentationReport from dace.codegen.instrumentation.data.data_report import InstrumentedDataReport from dace.codegen.compiled_sdfg import CompiledSDFG - from dace.sdfg.analysis.schedule_tree.treenodes import ScheduleTreeScope + from dace.sdfg.analysis.schedule_tree.treenodes import ScheduleTreeRoot class NestedDict(dict): @@ -179,13 +176,27 @@ class InterstateEdge(object): loop iterates). """ - assignments = Property(dtype=dict, desc="Assignments to perform upon transition (e.g., 'x=x+1; y = 0')") + assignments = DictProperty( + key_type=str, + value_type=str, + desc="Assignments to perform upon transition (e.g., 'x=x+1; y = 0')", + # NOTE: We serialize assignments as symbolic expressions but store them as strings of CodeBlocks (mostly with + # language=Python). In a future version, we will modify the value type to sympy.Basic and store the + # assignments as symbolic expressions without specialized to/from_json functions. + to_json=lambda d: { + k: symbolic.symstr(symbolic.pystr_to_symbolic(v)) + for k, v in d.items() + }, + from_json=(lambda d, *args, context=None, **kwargs: { + k: symbolic.symstr(symbolic.pystr_to_symbolic(v)) + for k, v in d.items() + })) condition = CodeProperty(desc="Transition condition", default=CodeBlock("1")) guid = Property(dtype=str, allow_none=False) def __init__(self, condition: Optional[Union[CodeBlock, str, ast.AST, list]] = None, - assignments: Optional[Dict] = None): + assignments: Optional[Dict[str, str | ast.AST]] = None): if condition is None: condition = CodeBlock("1") @@ -249,7 +260,7 @@ def read_symbols(self) -> Set[str]: # Symbols in conditions and assignments result = set(map(str, dace.symbolic.symbols_in_ast(self.condition.code[0]))) for assign in self.assignments.values(): - result |= symbolic.free_symbols_and_functions(assign) + result |= symbolic.free_symbols_and_functions(assign) | symbolic.arrays(assign) return result @@ -402,7 +413,6 @@ def to_json(self, parent=None): return { 'type': type(self).__name__, 'attributes': dace.serialize.all_properties_to_json(self), - 'label': self.label } @staticmethod @@ -413,24 +423,6 @@ def from_json(json_obj, context=None): return ret - @property - def label(self): - assignments = ','.join(['%s=%s' % (k, v) for k, v in self.assignments.items()]) - - # Edge with assigment only (no condition) - if self.condition.as_string == '1': - # Edge without conditions or assignments - if len(self.assignments) == 0: - return '' - return assignments - - # Edge with condition only (no assignment) - if len(self.assignments) == 0: - return self.condition.as_string - - # Edges with assigments and conditions - return self.condition.as_string + '; ' + assignments - @make_properties class SDFG(ControlFlowRegion): @@ -480,22 +472,6 @@ class SDFG(ControlFlowRegion): debuginfo = DebugInfoProperty(allow_none=True) - _pgrids = DictProperty(str, - ProcessGrid, - desc="Process-grid descriptors for this SDFG", - to_json=_arrays_to_json, - from_json=_arrays_from_json) - _subarrays = DictProperty(str, - SubArray, - desc="Sub-array descriptors for this SDFG", - to_json=_arrays_to_json, - from_json=_arrays_from_json) - _rdistrarrays = DictProperty(str, - RedistrArray, - desc="Sub-array redistribution descriptors for this SDFG", - to_json=_arrays_to_json, - from_json=_arrays_from_json) - callback_mapping = DictProperty(str, str, desc='Mapping between callback name and its original callback ' @@ -553,11 +529,6 @@ def __init__(self, self._regenerate_code = True self._recompile = True - # Grid-distribution-related fields - self._pgrids = {} - self._subarrays = {} - self._rdistrarrays = {} - # Counter to resolve name conflicts self._orig_name = name self._num = 0 @@ -608,7 +579,61 @@ def sdfg_id(self): """ return self.cfg_id - def to_json(self, hash=False): + def compute_debuginfo_files(self) -> list[str]: + """ + Computes the set of files that are referenced in the SDFG's debug information, and adds them to the SDFG's + debug information properties. + + :return: A list of files that are referenced in the SDFG's debug information, in the order they were indexed. + """ + files: list[str] = [] + + def _replace_file_with_index(element: Any): + if getattr(element, 'debuginfo', False): + debuginfo: dtypes.DebugInfo = element.debuginfo + if debuginfo.filename: + # Replace filename with number + try: + fileindex = files.index(debuginfo.filename) + except ValueError: + files.append(debuginfo.filename) + fileindex = len(files) - 1 + debuginfo.file_index = fileindex + + _replace_file_with_index(self) + for _, _, arr in self.arrays_recursive(): + _replace_file_with_index(arr) + for node, _ in self.all_nodes_recursive(): + _replace_file_with_index(node) + for edge, _ in self.all_edges_recursive(): + _replace_file_with_index(edge) + + return files + + def rematerialize_debuginfo_files(self, files: list[str]) -> None: + """ + Replaces file indices in the SDFG's debug information with the actual file names, using the provided list of + files. + + :param files: A list of files that are referenced in the SDFG's debug information, in the order they were indexed. + """ + + def _replace_index_with_file(element: Any): + if getattr(element, 'debuginfo', False): + debuginfo: dtypes.DebugInfo = element.debuginfo + if debuginfo.file_index is not None: + debuginfo.filename = files[debuginfo.file_index] + debuginfo.file_index = None + + _replace_index_with_file(self) + for _, _, arr in self.arrays_recursive(): + _replace_index_with_file(arr) + for node, _ in self.all_nodes_recursive(): + _replace_index_with_file(node) + for edge, _ in self.all_edges_recursive(): + _replace_index_with_file(edge) + + def to_json(self, hash=False, include_transformation_history=False): """ Serializes this object to JSON format. :return: A string representing the JSON-serialized SDFG. @@ -617,13 +642,40 @@ def to_json(self, hash=False): is_root = self.parent_sdfg is None if is_root: self.reset_cfg_list() - - tmp = super().to_json() + source_files = self.compute_debuginfo_files() + + # Serialize the control-flow graph (states and interstate edges) under this + # SDFG's declared symbols, so symbolic expressions outside any dataflow scope + # (e.g. interstate-edge conditions/assignments) emit a deterministic dtype. + # Each nested SDFG re-pushes its own symbols, and the previous authority is + # restored on exit. + with symbolic.serialization_symbol_dtypes(self.symbols): + tmp = super().to_json() + if is_root: + tmp['source_files'] = source_files # Ensure properties are serialized correctly if 'constants_prop' in tmp['attributes']: tmp['attributes']['constants_prop'] = json.loads(dace.serialize.dumps(tmp['attributes']['constants_prop'])) + if is_root and not include_transformation_history: + # Strip transformation history from JSON recursively + def _strip_transformation_history(json_obj: Any): + if isinstance(json_obj, dict): + if 'type' in json_obj and json_obj['type'] == 'SDFG' and 'attributes' in json_obj: + if 'transformation_hist' in json_obj['attributes']: + del json_obj['attributes']['transformation_hist'] + if 'orig_sdfg' in json_obj['attributes']: + del json_obj['attributes']['orig_sdfg'] + + for v in json_obj.values(): + _strip_transformation_history(v) + elif isinstance(json_obj, (list, tuple)): + for v in json_obj: + _strip_transformation_history(v) + + _strip_transformation_history(tmp) + tmp['attributes']['name'] = self.name if hash: tmp['attributes']['hash'] = self.hash_sdfg(tmp) @@ -636,22 +688,28 @@ def to_json(self, hash=False): @classmethod def from_json(cls, json_obj, context=None): context = context or {'sdfg': None} + context['version'] = json_obj.get('dace_version', context.get('version')) _type = json_obj['type'] if _type != cls.__name__: raise TypeError("Class type mismatch") - attrs = json_obj['attributes'] + attrs = dict(json_obj['attributes']) + json_obj = dict(json_obj) + json_obj['attributes'] = attrs nodes = json_obj['nodes'] edges = json_obj['edges'] if 'constants_prop' in attrs: - constants_prop = dace.serialize.loads(dace.serialize.dumps(attrs['constants_prop'])) + constants_prop = dace.serialize.loads(dace.serialize.dumps(attrs['constants_prop']), context=context) else: constants_prop = None ret = SDFG(name=attrs['name'], constants=constants_prop, parent=context['sdfg']) - dace.serialize.set_properties_from_json(ret, json_obj, ignore_properties={'constants_prop', 'name', 'hash'}) + dace.serialize.set_properties_from_json(ret, + json_obj, + context=context, + ignore_properties={'constants_prop', 'name', 'hash'}) nodelist = [] for n in nodes: @@ -663,12 +721,15 @@ def from_json(cls, json_obj, context=None): nodelist.append(block) for e in edges: - e = dace.serialize.from_json(e) + e = dace.serialize.from_json(e, context=context) ret.add_edge(nodelist[int(e.src)], nodelist[int(e.dst)], e.data) if 'start_block' in json_obj: ret._start_block = json_obj['start_block'] + if 'source_files' in json_obj: # This will only happen on the root SDFG, once deserialization is complete + ret.rematerialize_debuginfo_files(json_obj['source_files']) + return ret def hash_sdfg(self, jsondict: Optional[Dict[str, Any]] = None) -> str: @@ -726,17 +787,17 @@ def arrays(self): @property def process_grids(self): """ Returns a dictionary of process-grid descriptors (`ProcessGrid` objects) used in this SDFG. """ - return self._pgrids + return {name: desc for name, desc in self._arrays.items() if isinstance(desc, ProcessGrid)} @property def subarrays(self): """ Returns a dictionary of sub-array descriptors (`SubArray` objects) used in this SDFG. """ - return self._subarrays + return {name: desc for name, desc in self._arrays.items() if isinstance(desc, SubArray)} @property def rdistrarrays(self): """ Returns a dictionary of sub-array redistribution descriptors (`RedistrArray` objects) used in this SDFG. """ - return self._rdistrarrays + return {name: desc for name, desc in self._arrays.items() if isinstance(desc, RedistrArray)} def data(self, dataname: str): """ Looks up a data descriptor from its name, which can be an array, stream, or scalar symbol. """ @@ -818,12 +879,6 @@ def add_symbol(self, name, stype, find_new_name: bool = False): raise FileExistsError(f'Symbol "{name}" already exists in SDFG') if name in self.arrays: raise FileExistsError(f'Cannot create symbol "{name}", the name is used by a data descriptor.') - if name in self._subarrays: - raise FileExistsError(f'Cannot create symbol "{name}", the name is used by a subarray.') - if name in self._rdistrarrays: - raise FileExistsError(f'Cannot create symbol "{name}", the name is used by a RedistrArray.') - if name in self._pgrids: - raise FileExistsError(f'Cannot create symbol "{name}", the name is used by a ProcessGrid.') if not isinstance(stype, dtypes.typeclass): stype = dtypes.dtype_to_typeclass(stype) self.symbols[name] = stype @@ -1128,7 +1183,7 @@ def call_with_instrumented_data(self, dreport: 'InstrumentedDataReport', *args, ########################################## - def as_schedule_tree(self, in_place: bool = False) -> 'ScheduleTreeScope': + def as_schedule_tree(self, in_place: bool = False) -> 'ScheduleTreeRoot': """ Creates a schedule tree from this SDFG and all nested SDFGs. The schedule tree is a tree of nodes that represent the execution order of the SDFG. @@ -1136,7 +1191,8 @@ def as_schedule_tree(self, in_place: bool = False) -> 'ScheduleTreeScope': etc.) or a ``ScheduleTreeScope`` block (map, for-loop, etc.) that contains other nodes. It can be used to generate code from an SDFG, or to perform schedule transformations on the SDFG. For example, - erasing an empty if branch, or merging two consecutive for-loops. + erasing an empty if branch, or merging two consecutive for-loops. The SDFG can then be reconstructed via the + ``as_sdfg`` method or the ``from_schedule_tree`` function in ``dace.sdfg.analysis.schedule_tree.tree_to_sdfg``. :param in_place: If True, the SDFG is modified in-place. Otherwise, a copy is made. Note that the SDFG might not be usable after the conversion if ``in_place`` is True! @@ -1275,12 +1331,6 @@ def add_constant(self, name: str, value: Any, dtype: dt.Data = None): :param value: The constant value. :param dtype: Optional data type of the symbol, or None to deduce automatically. """ - if name in self._subarrays: - raise FileExistsError(f'Can not create constant "{name}", the name is used by a subarray.') - if name in self._rdistrarrays: - raise FileExistsError(f'Can not create constant "{name}", the name is used by a RedistrArray.') - if name in self._pgrids: - raise FileExistsError(f'Can not create constant "{name}", the name is used by a ProcessGrid.') self.constants_prop[name] = (dtype or dt.create_datadescriptor(value), value) @property @@ -1373,12 +1423,29 @@ def _used_symbols_internal(self, for code in self.exit_code.values(): free_syms |= symbolic.symbols_in_code(code.as_string, self.symbols.keys()) - return super()._used_symbols_internal(all_symbols=all_symbols, - keep_defined_in_mapping=keep_defined_in_mapping, - defined_syms=defined_syms, - free_syms=free_syms, - used_before_assignment=used_before_assignment, - with_contents=with_contents) + # Snapshot the set so super() can filter array names out of + # ``free_syms``; we re-extract them from internal references below. + result = super()._used_symbols_internal(all_symbols=all_symbols, + keep_defined_in_mapping=keep_defined_in_mapping, + defined_syms=defined_syms, + free_syms=free_syms, + used_before_assignment=used_before_assignment, + with_contents=with_contents) + # Expand array-descriptor stride/shape/offset symbols into the free + # set. Without this, a ``ConditionalBlock`` guard or memlet subset + # referencing ``A[i, j]`` leaves the symbols used in ``A`` 's strides + # out of the computed free-symbol set, causing + # ``generate_nsdfg_header`` to emit a nested function signature + # missing those symbols, ceating an invalid SDFG. + res_free, res_defined, res_before = result + if with_contents: + for desc in self.arrays.values(): + res_free |= {str(s) for s in desc.used_symbols(all_symbols)} + # Don't drag in symbols that are genuinely defined inside this + # SDFG (e.g., LoopRegion loop variables); keep only the ones + # outside ``defined_syms``. + res_free -= res_defined + return res_free, res_defined, res_before def get_all_toplevel_symbols(self) -> Set[str]: """ @@ -1646,7 +1713,14 @@ def shared_transients(self, check_toplevel: bool = True, include_nested_data: bo return dtypes.deduplicate(shared) - def save(self, filename: str, use_pickle=False, hash=None, exception=None, compress=False) -> Optional[str]: + def save(self, + filename: str, + use_pickle=False, + hash=None, + exception=None, + compress=False, + readable=False, + include_transformation_history=False) -> Optional[str]: """ Save this SDFG to a file. :param filename: File name to save to. @@ -1657,7 +1731,9 @@ def save(self, filename: str, use_pickle=False, hash=None, exception=None, compr :param exception: If not None, stores error information along with SDFG. :param compress: If True, uses gzip to compress the file upon saving. - :return: The hash of the SDFG, or None if failed/not requested. + :param readable: If True, saves the JSON in a human-readable format. + :param include_transformation_history: If True, includes the transformation history in the saved SDFG. + :return: The hash of the SDFG, or None if not requested. """ filename = os.path.expanduser(filename) @@ -1679,10 +1755,10 @@ def save(self, filename: str, use_pickle=False, hash=None, exception=None, compr else: hash = True if hash is None else hash with fileopen(filename, "w") as fp: - json_output = self.to_json(hash=hash) + json_output = self.to_json(hash=hash, include_transformation_history=include_transformation_history) if exception: json_output['error'] = exception.to_json() - dace.serialize.dump(json_output, fp) + dace.serialize.dump(json_output, fp, readable=readable) if hash and 'hash' in json_output['attributes']: return json_output['attributes']['hash'] @@ -1696,6 +1772,7 @@ def view(self, filename=None, verbose=False): :param verbose: Be verbose, `False` by default. """ from dace.cli.sdfv import view + view(self, filename=filename, verbose=verbose) @staticmethod @@ -1730,12 +1807,11 @@ def from_file(filename: str) -> 'SDFG': # Dynamic SDFG creation API ############################## - + # TODO(tehrengruber): This could also be done using the counter. def _find_new_name(self, name: str): """ Tries to find a new name by adding an underscore and a number. """ - names = (self._arrays.keys() | self.constants_prop.keys() | self._pgrids.keys() | self._subarrays.keys() - | self._rdistrarrays.keys() | self.symbols.keys()) + names = (self._arrays.keys() | self.constants_prop.keys() | self.symbols.keys()) return dt.find_new_name(name, names) def is_name_used(self, name: str) -> bool: @@ -1746,12 +1822,6 @@ def is_name_used(self, name: str) -> bool: return True if name in self.constants_prop: return True - if name in self._pgrids: - return True - if name in self._subarrays: - return True - if name in self._rdistrarrays: - return True return False def is_name_free(self, name: str) -> bool: @@ -2101,12 +2171,6 @@ def add_datadesc(self, name: str, datadesc: dt.Data, find_new_name=False) -> str raise FileExistsError(f'Data descriptor "{name}" already exists in SDFG') if name in self.symbols: raise FileExistsError(f'Can not create data descriptor "{name}", the name is used by a symbol.') - if name in self._subarrays: - raise FileExistsError(f'Can not create data descriptor "{name}", the name is used by a subarray.') - if name in self._rdistrarrays: - raise FileExistsError(f'Can not create data descriptor "{name}", the name is used by a RedistrArray.') - if name in self._pgrids: - raise FileExistsError(f'Can not create data descriptor "{name}", the name is used by a ProcessGrid.') def _add_symbols(sdfg: SDFG, desc: dt.Data): if isinstance(desc, dt.Structure): @@ -2152,7 +2216,8 @@ def add_pgrid(self, parent_grid: str = None, color: Sequence[Union[Integral, bool]] = None, exact_grid: RankType = None, - root: RankType = 0): + root: RankType = 0, + name: Optional[str] = None) -> str: """ Adds a process-grid to the process-grid descriptor store. For more details on process-grids, please read the documentation of the ProcessGrid class. @@ -2161,6 +2226,7 @@ def add_pgrid(self, :param color: The i-th entry specifies whether the i-th dimension is kept in the sub-grid or is dropped (see `remain_dims` input of [MPI_Cart_sub](https://www.mpich.org/static/docs/v3.2/www3/MPI_Cart_sub.html)). :param exact_grid: If set then, out of all the sub-grids created, only the one that contains the rank with id `exact_grid` will be utilized for collective communication. :param root: Root rank (used for collective communication). + :param name: Name hint of the new process-grid descriptor. If None, a name will be automatically generated. :return: Name of the new process-grid descriptor. """ @@ -2177,15 +2243,16 @@ def add_pgrid(self, newshape.append(dace.symbolic.pystr_to_symbolic(s)) shape = newshape - grid_name = self._find_new_name('__pgrid') + grid_name = self._find_new_name(name or '__pgrid') is_subgrid = (parent_grid is not None) if parent_grid and isinstance(parent_grid, str): - parent_grid = self._pgrids[parent_grid] + parent_grid = self.process_grids[parent_grid] - self._pgrids[grid_name] = ProcessGrid(grid_name, is_subgrid, shape, parent_grid, color, exact_grid, root) + pgrid = ProcessGrid(grid_name, is_subgrid, shape, parent_grid, color, exact_grid, root) + self.add_datadesc(grid_name, pgrid) - self.append_init_code(self._pgrids[grid_name].init_code()) - self.append_exit_code(self._pgrids[grid_name].exit_code()) + self.append_init_code(pgrid.init_code()) + self.append_exit_code(pgrid.exit_code()) return grid_name @@ -2194,7 +2261,8 @@ def add_subarray(self, shape: ShapeType, subshape: ShapeType, pgrid: str = None, - correspondence: Sequence[Integral] = None): + correspondence: Sequence[Integral] = None, + name: Optional[str] = None): """ Adds a sub-array to the sub-array descriptor store. For more details on sub-arrays, please read the documentation of the SubArray class. @@ -2203,6 +2271,7 @@ def add_subarray(self, :param subshape: Sub-shape of the sub-array (see `array_of_subsizes` parameter of [MPI_Type_create_subarray](https://www.mpich.org/static/docs/v3.2/www3/MPI_Type_create_subarray.html)). :param pgrid: Process-grid used for collective scatter/gather operations. :param correspondence: Matching among array dimensions and process-grid dimensions. + :param name: Name hint of the new sub-array descriptor. If None, a name will be automatically generated. :return: Name of the new sub-array descriptor. """ @@ -2225,44 +2294,125 @@ def add_subarray(self, subshape = newshape # No need to ensure unique test. - subarray_name = self._find_new_name('__subarray') + subarray_name = self._find_new_name(name or '__subarray') - self._subarrays[subarray_name] = SubArray(subarray_name, dtype, shape, subshape, pgrid, correspondence) - self.append_init_code(self._subarrays[subarray_name].init_code()) - self.append_exit_code(self._subarrays[subarray_name].exit_code()) + subarray = SubArray(subarray_name, dtype, shape, subshape, pgrid, correspondence) + self.add_datadesc(subarray_name, subarray) + self.append_init_code(subarray.init_code()) + self.append_exit_code(subarray.exit_code()) return subarray_name - def add_rdistrarray(self, array_a: str, array_b: str): + def add_rdistrarray(self, array_a: str, array_b: str, name: Optional[str] = None) -> str: """ Adds a sub-array redistribution to the sub-array redistribution descriptor store. For more details on redistributions, please read the documentation of the RedistrArray class. :param array_a: Input sub-array descriptor. :param array_b: Output sub-array descriptor. + :param name: Name hint of the new redistribution descriptor. If None, a name will be automatically generated. :return: Name of the new redistribution descriptor. """ # No need to ensure unique test. - name = self._find_new_name('__rdistrarray') + name = self._find_new_name(name or '__rdistrarray') - self._rdistrarrays[name] = RedistrArray(name, array_a, array_b) - self.append_init_code(self._rdistrarrays[name].init_code(self)) - self.append_exit_code(self._rdistrarrays[name].exit_code(self)) + rdistrarray = RedistrArray(name, array_a, array_b) + self.add_datadesc(name, rdistrarray) + self.append_init_code(rdistrarray.init_code(self)) + self.append_exit_code(rdistrarray.exit_code(self)) return name - def add_loop( + def add_loop(self, + before_block: ControlFlowBlock, + loop_start_block: ControlFlowBlock, + after_block: ControlFlowBlock, + loop_var: str, + initialize_expr: str, + condition_expr: str, + increment_expr: str, + loop_end_block: Optional[ControlFlowBlock] = None, + label: Optional[str] = None) -> LoopRegion: + """ + Helper function that adds a looping control flow block around a + given block/state (or sequence of blocks, if ``loop_end_block`` is provided). + + :param before_block: The block after which the loop should + begin, or None if the loop is the first + block (creates an empty block). + :param loop_start_block: The block that begins the loop. See also + ``loop_end_block`` if the loop is multi-block. + :param after_block: The block that should be invoked after + the loop ends, or None if the program + should terminate (creates an empty block). + :param loop_var: A name of a symbol to use for the loop variable. + :param initialize_expr: A string expression that is assigned + to ``loop_var`` before the loop begins. + If None, does not define an expression. + :param condition_expr: A string condition that occurs every + loop iteration. If None, loops forever. + :param increment_expr: A string expression that is assigned to + ``loop_var`` after every loop iteration. + :param loop_end_block: If the loop wraps multiple blocks, the block + where the loop iteration ends. If None, sets + the end block to ``loop_start_block`` as well. + :param label: An optional label for the loop region. + :return: The generated LoopRegion block. + """ + label = self._ensure_unique_block_name(label or "loop") + loop_region = LoopRegion(label, condition_expr, loop_var, + f'{loop_var} = {initialize_expr}' if initialize_expr else None, + f'{loop_var} = {increment_expr}' if increment_expr else None) + + # Capture subgraphview of loop body + if loop_start_block in self.nodes() or loop_start_block in self.states(): + # Find all reachable blocks in loop body + if loop_end_block is None: + blocks = {loop_start_block} + else: + blocks = set(self.all_nodes_between(loop_start_block, loop_end_block)) + blocks.add(loop_start_block) + blocks.add(loop_end_block) + + subgraph = SubgraphView(self, blocks) + remove_subgraph = True + else: + subgraph = SubgraphView(self, [loop_start_block, loop_end_block] if loop_end_block else [loop_start_block]) + remove_subgraph = False + + # Readd subgraph to main SDFG with loop_region as parent + loop_region.add_node(loop_start_block, is_start_block=True) + loop_region.add_nodes_from(set(subgraph.nodes()) - {loop_start_block}) + for e in subgraph.edges(): + loop_region.add_edge(e.src, e.dst, e.data) + + # Remove subgraph from main SDFG, if necessary + if remove_subgraph: + self.remove_nodes_from(blocks) + + # Connect to graph + self.add_node(loop_region, is_start_block=(before_block is None)) + if before_block is not None: + self.add_edge(before_block, loop_region, InterstateEdge()) + if after_block is not None: + self.add_edge(loop_region, after_block, InterstateEdge()) + + return loop_region + + def add_loop_state_machine( self, - before_state, - loop_state, - after_state, + before_state: SDFGState, + loop_state: SDFGState, + after_state: SDFGState, loop_var: str, initialize_expr: str, condition_expr: str, increment_expr: str, - loop_end_state=None, + loop_end_state: Optional[SDFGState] = None, ): """ Helper function that adds a looping state machine around a - given state (or sequence of states). + given state (or sequence of states). It is recommended to use + ``add_loop`` instead of this function, unless creating a loop + state machine is explicitly requested. :param before_state: The state after which the loop should begin, or None if the loop is the first @@ -2293,9 +2443,6 @@ def add_loop( """ from dace.frontend.python.astutils import negate_expr # Avoid import loops - warnings.warn("SDFG.add_loop is deprecated and will be removed in a future release. Use LoopRegions instead.", - DeprecationWarning) - # Argument checks if loop_var is None and (initialize_expr or increment_expr): raise ValueError("Cannot initalize or increment an empty loop variable") @@ -2359,16 +2506,26 @@ def specialize(self, symbols: Dict[str, Any]): for k, v in symbols.items(): self.add_constant(str(k), v) - def is_loaded(self) -> bool: + def is_loaded(self, folder_mode: Optional[str] = None) -> bool: """ - Returns True if the SDFG binary is already loaded in the current - process. + Returns True if the SDFG binary is already loaded in the current process. + Returns `False` if the file does not exist. """ # Avoid import loops from dace.codegen import compiled_sdfg as cs, compiler - binary_filename = compiler.get_binary_name(self.build_folder, self.name) - dll = cs.ReloadableDLL(binary_filename, self.name) + build_folder = self.build_folder + if folder_mode is None: + folder_mode = compiler.get_folder_mode(build_folder, probe=True) + if folder_mode is None: + folder_mode = Config.get('compiler', 'build_folder_mode') + + # Note here is kind of a "leak". The issue is that while the library is not + # loaded the stub library is loaded. However, it is never unloaded, because + # the `unload()` function is not called. This is technically not an issue + # as `ctypes` will do that at some later point. + binary_filename = compiler.get_binary_name(self.build_folder, self.name, folder_mode=folder_mode) + dll = cs.ReloadableDLL(binary_filename) return dll.is_loaded() def compile(self, output_file=None, validate=True, return_program_handle=True) -> 'CompiledSDFG': @@ -2388,11 +2545,24 @@ def compile(self, output_file=None, validate=True, return_program_handle=True) - # Compute build folder path before running codegen build_folder = self.build_folder + # Get the folder mode, but if the folder already exists, then use the `FOLDER_MODE` file. + folder_mode = compiler.get_folder_mode(build_folder, probe=True) + if folder_mode is None: + folder_mode = Config.get('compiler', 'build_folder_mode') + if not self._recompile or Config.get_bool('compiler', 'use_cache'): # Try to see if a cached version of the binary exists - binary_filename = compiler.get_binary_name(build_folder, self.name) - if os.path.isfile(binary_filename): - return compiler.load_from_file(self, binary_filename) + lib_path = compiler.get_binary_name(object_folder=build_folder, + sdfg_name=self.name, + folder_mode=folder_mode) + if lib_path.is_file(): + if return_program_handle: + # NOTE: We should not pass `self` as `sdfg` argument, but instead deepcopy it. + # The reason is that if code is generated the `CompiledSDFG.sdfg` attribute + # of the returned handle is deepcopied. This means that changes to `self` + # will not affect the attribute. But currently this is the case. + return compiler.load_precompiled_sdfg(folder=build_folder, sdfg=self) + return ############################ # DaCe Compilation Process # @@ -2410,7 +2580,7 @@ def compile(self, output_file=None, validate=True, return_program_handle=True) - # Rename SDFG to avoid runtime issues with clashing names index = 0 - while sdfg.is_loaded(): + while sdfg.is_loaded(folder_mode=folder_mode): sdfg.name = f'{self.name}_{index}' index += 1 if self.name != sdfg.name and Config.get_bool('debugprint'): @@ -2422,6 +2592,7 @@ def compile(self, output_file=None, validate=True, return_program_handle=True) - sdfg.fill_scope_connectors() # Generate code for the program by traversing the SDFG state by state + # It is not yet written to disc. program_objects = codegen.generate_code(sdfg, validate=validate) except Exception: fpath = os.path.join('_dacegraphs', 'failing.sdfgz') @@ -2430,14 +2601,18 @@ def compile(self, output_file=None, validate=True, return_program_handle=True) - raise # Generate the program folder and write the source files - program_folder = compiler.generate_program_folder(sdfg, program_objects, build_folder) + program_folder = compiler.generate_program_folder(sdfg, + program_objects, + build_folder, + folder_mode=folder_mode) else: # The code was already generated, just load the program folder program_folder = build_folder + # NOTE: See the note above about deepcopying. sdfg = self # Compile the code and get the shared library path - shared_library = compiler.configure_and_compile(program_folder, sdfg.name) + shared_library = compiler.configure_and_compile(program_folder, sdfg.name, folder_mode=folder_mode) # If provided, save output to path or filename if output_file is not None: @@ -2447,7 +2622,7 @@ def compile(self, output_file=None, validate=True, return_program_handle=True) - # Get the function handle if return_program_handle: - return compiler.get_program_handle(shared_library, sdfg) + return compiler.load_precompiled_sdfg(folder=build_folder, sdfg=sdfg) def argument_typecheck(self, args, kwargs, types_only=False): """ Checks if arguments and keyword arguments match the SDFG diff --git a/dace/sdfg/state.py b/dace/sdfg/state.py index adc77bf439..68374dbdc5 100644 --- a/dace/sdfg/state.py +++ b/dace/sdfg/state.py @@ -5,12 +5,14 @@ import abc import collections import copy +import hashlib import inspect import itertools +import json import warnings import sympy -from typing import (TYPE_CHECKING, Any, AnyStr, Callable, Dict, Iterable, Iterator, List, Optional, Set, Tuple, Type, - Union, overload) +from typing import (TYPE_CHECKING, Any, AnyStr, Callable, Dict, Iterable, Iterator, List, Optional, Set, Tuple, Union, + overload) import dace from dace.frontend.python import astutils @@ -27,7 +29,7 @@ from dace.sdfg import nodes as nd from dace.sdfg.graph import (MultiConnectorEdge, NodeNotFoundError, OrderedMultiDiConnectorGraph, SubgraphView, OrderedDiGraph, Edge, generate_element_id) -from dace.sdfg.propagation import propagate_memlet +from dace.sdfg import propagation as sdprop from dace.sdfg.type_inference import infer_expr_type from dace.sdfg.validation import validate_state from dace.subsets import Range, Subset @@ -41,18 +43,24 @@ GraphT = Union['ControlFlowRegion', 'SDFGState'] -def _getdebuginfo(old_dinfo=None) -> dtypes.DebugInfo: - """ Returns a DebugInfo object for the position that called this function. +def _get_debug_info(explicit_lineinfo: dtypes.DebugInfo | None) -> dtypes.DebugInfo | None: + """Returns a DebugInfo from the stack, if configured. - :param old_dinfo: Another DebugInfo object that will override the - return value of this function - :return: DebugInfo containing line number and calling file. + If lineinfo is configured in the config, this function inspects the python + stacktrace and returns a DebugInfo object for the position that called this + function. `explicit_lineinfo` has precedence, if given. """ - if old_dinfo is not None: - return old_dinfo + if dace.Config.get("compiler", "lineinfo") == "none": + return None + + if explicit_lineinfo is not None: + return explicit_lineinfo + + if dace.Config.get("compiler", "lineinfo") == "inspect": + caller = inspect.getframeinfo(inspect.stack()[2][0], context=0) + return dtypes.DebugInfo(caller.lineno, 0, caller.lineno, 0, caller.filename) - caller = inspect.getframeinfo(inspect.stack()[2][0], context=0) - return dtypes.DebugInfo(caller.lineno, 0, caller.lineno, 0, caller.filename) + return None def _make_iterators(ndrange): @@ -124,8 +132,8 @@ def all_nodes_recursive(self, (i.e., an SDFG or a scope block). :param predicate: An optional predicate function that decides on whether the traversal should recurse or not. - If the predicate returns False, traversal is not recursed any further into the graph found under NodeT for - a given [NodeT, GraphT] pair. + If the predicate returns False, traversal is not recursed any further into the graph found + under NodeT for a given [NodeT, GraphT] pair. """ return [] @@ -235,6 +243,7 @@ def used_symbols(self, """ return set() + # TODO(tehrengruber): should we make this an ordered set? @property def free_symbols(self) -> Set[str]: """ @@ -1257,7 +1266,7 @@ def edges(self): def sub_regions(self) -> List['AbstractControlFlowRegion']: return [] - def set_default_lineinfo(self, lineinfo: dace.dtypes.DebugInfo): + def set_default_lineinfo(self, lineinfo: Optional[dace.dtypes.DebugInfo]) -> None: """ Sets the default source line information to be lineinfo, or None to revert to default mode. @@ -1288,7 +1297,7 @@ def from_json(cls, json_obj, context=None): ret = cls(label=json_obj['label'], sdfg=context['sdfg']) - dace.serialize.set_properties_from_json(ret, json_obj) + dace.serialize.set_properties_from_json(ret, json_obj, context=context) return ret @@ -1366,7 +1375,7 @@ class SDFGState(OrderedMultiDiConnectorGraph[nd.Node, mm.Memlet], ControlFlowBlo default=CodeBlock("1", language=dtypes.Language.CPP)) location = DictProperty(key_type=str, - value_type=symbolic.pystr_to_symbolic, + value_type=sympy.Basic, desc='Full storage location identifier (e.g., rank, GPU ID)') def __repr__(self) -> str: @@ -1390,6 +1399,52 @@ def __init__(self, label=None, sdfg=None, debuginfo=None, location=None): self.location = location if location is not None else {} self._default_lineinfo = None + # TODO(tehrengruber): unify with sdfg.hash() + def hash_state(self, jsondict: Optional[Dict[str, Any]] = None) -> str: + """ + Returns a hash of the current SDFG, without considering IDs and attribute names. + + :param jsondict: If not None, uses given JSON dictionary as input. + :return: The hash (in SHA-256 format). + """ + + def keyword_remover(json_obj: Any, last_keyword=""): + # Makes non-unique in SDFG hierarchy v2 + # Recursively remove attributes from the SDFG which are not used in + # uniquely representing the SDFG. This, among other things, includes + # the hash, name, transformation history, and meta attributes. + if isinstance(json_obj, dict): + if 'cfg_list_id' in json_obj: + del json_obj['cfg_list_id'] + + keys_to_delete = [] + kv_to_recurse = [] + for key, value in json_obj.items(): + if (isinstance(key, str) + and (key.startswith('_meta_') + or key in ['name', 'hash', 'orig_sdfg', 'transformation_hist', 'instrument', 'guid'])): + keys_to_delete.append(key) + else: + kv_to_recurse.append((key, value)) + + for key in keys_to_delete: + del json_obj[key] + + for key, value in kv_to_recurse: + keyword_remover(value, last_keyword=key) + elif isinstance(json_obj, (list, tuple)): + for value in json_obj: + keyword_remover(value) + + # Clean SDFG of nonstandard objects + jsondict = (json.loads(json.dumps(jsondict)) if jsondict is not None else self.to_json()) + + keyword_remover(jsondict) # Make non-unique in SDFG hierarchy + + string_representation = json.dumps(jsondict) # dict->str + hsh = hashlib.sha256(string_representation.encode('utf-8')) + return hsh.hexdigest() + @property def parent(self): """ Returns the parent SDFG of this state. """ @@ -1479,29 +1534,65 @@ def to_json(self, parent=None): for edge in self.edges(): edge.data.try_initialize(self.sdfg, self, edge) + # Resolve the authoritative dtype of every symbol an element may reference, so + # serialization emits a deterministic dtype that depends only on the enclosing + # scopes -- never on a symbol instance's own (cache-stale) dtype. The authority + # is rebuilt fresh on every serialization (never stored): opening a scope (entry + # node) augments the running map with that scope's ``new_symbols`` (its map + # iterators and dynamic-input connector symbols), and the parent level is + # restored when the recursion returns. + authority_by_node: Dict[nd.Node, Dict[str, dtypes.typeclass]] = {} + try: + scope_children = self.scope_children() + + def _open_scope(scope_entry: Optional[nd.Node], authority: Dict[str, dtypes.typeclass]): + for child in scope_children.get(scope_entry, []): + if isinstance(child, nd.EntryNode): + inner = {**authority, **child.new_symbols(self.sdfg, self, authority)} + authority_by_node[child] = inner + _open_scope(child, inner) + else: + authority_by_node[child] = authority + + _open_scope(None, dict(self.sdfg.symbols)) + except (RuntimeError, ValueError, KeyError): + authority_by_node = {} + + nodes_json = [] + for n in self.nodes(): + with symbolic.serialization_symbol_dtypes(authority_by_node.get(n, self.sdfg.symbols)): + nodes_json.append(n.to_json(self)) + + # An edge's memlet sees the symbols of both endpoints' scopes (the richer + # one is sometimes the source, e.g. MapExit -> AccessNode), so merge them. + edges_json = [] + for e in sorted(self.edges(), key=lambda e: (e.src_conn or '', e.dst_conn or '')): + authority = {**authority_by_node.get(e.src, {}), **authority_by_node.get(e.dst, {})} + with symbolic.serialization_symbol_dtypes(authority): + edges_json.append(e.to_json(self)) + ret = { 'type': type(self).__name__, 'label': self.name, 'id': parent.node_id(self) if parent is not None else None, 'collapsed': self.is_collapsed, 'scope_dict': scope_dict, - 'nodes': [n.to_json(self) for n in self.nodes()], - 'edges': - [e.to_json(self) for e in sorted(self.edges(), key=lambda e: (e.src_conn or '', e.dst_conn or ''))], + 'nodes': nodes_json, + 'edges': edges_json, 'attributes': serialize.all_properties_to_json(self), } return ret @classmethod - def from_json(cls, json_obj, context={'sdfg': None}, pre_ret=None): + def from_json(cls, json_obj, context=None, pre_ret=None): """ Loads the node properties, label and type into a dict. :param json_obj: The object containing information about this node. NOTE: This may not be a string! :return: An SDFGState instance constructed from the passed data """ - + context = context or {'sdfg': None} _type = json_obj['type'] if _type != cls.__name__: raise Exception("Class type mismatch") @@ -1516,7 +1607,8 @@ def from_json(cls, json_obj, context={'sdfg': None}, pre_ret=None): rec_ci = { 'sdfg': context['sdfg'], 'sdfg_state': ret, - 'callback': context['callback'] if 'callback' in context else None + 'callback': context.get('callback'), + 'version': context.get('version'), } serialize.set_properties_from_json(ret, json_obj, rec_ci) @@ -1626,7 +1718,6 @@ def add_read(self, array_or_stream_name: str, debuginfo: Optional[dtypes.DebugIn :return: An array access node. :see: add_access """ - debuginfo = _getdebuginfo(debuginfo or self._default_lineinfo) return self.add_access(array_or_stream_name, debuginfo=debuginfo) def add_write(self, array_or_stream_name: str, debuginfo: Optional[dtypes.DebugInfo] = None) -> nd.AccessNode: @@ -1638,7 +1729,6 @@ def add_write(self, array_or_stream_name: str, debuginfo: Optional[dtypes.DebugI :return: An array access node. :see: add_access """ - debuginfo = _getdebuginfo(debuginfo or self._default_lineinfo) return self.add_access(array_or_stream_name, debuginfo=debuginfo) def add_access(self, array_or_stream_name: str, debuginfo: Optional[dtypes.DebugInfo] = None) -> nd.AccessNode: @@ -1648,7 +1738,7 @@ def add_access(self, array_or_stream_name: str, debuginfo: Optional[dtypes.Debug :param debuginfo: Source line information for this access node. :return: An array access node. """ - debuginfo = _getdebuginfo(debuginfo or self._default_lineinfo) + debuginfo = _get_debug_info(debuginfo or self._default_lineinfo) node = nd.AccessNode(array_or_stream_name, debuginfo=debuginfo) self.add_node(node) return node @@ -1666,10 +1756,10 @@ def add_tasklet( code_exit: str = "", location: dict = None, side_effects: Optional[bool] = None, - debuginfo=None, + debuginfo: Optional[dtypes.DebugInfo] = None, ): """ Adds a tasklet to the SDFG state. """ - debuginfo = _getdebuginfo(debuginfo or self._default_lineinfo) + debuginfo = _get_debug_info(debuginfo or self._default_lineinfo) # Make dictionary of autodetect connector types from set if isinstance(inputs, (set, collections.abc.KeysView)): @@ -1738,7 +1828,7 @@ def add_nested_sdfg( """ if name is None: name = sdfg.label - debuginfo = _getdebuginfo(debuginfo or self._default_lineinfo) + debuginfo = _get_debug_info(debuginfo or self._default_lineinfo) if sdfg is None and external_path is None: raise ValueError('Neither an SDFG nor an external SDFG path has been provided') @@ -1750,9 +1840,13 @@ def add_nested_sdfg( sdfg.update_cfg_list([]) # Make dictionary of autodetect connector types from set - if isinstance(inputs, (set, collections.abc.KeysView)): + # TODO(tehrengruber): Using sets here leads to a situation where self._nodes has a different + # ordering, but to_json from_json restores the order again. Investigate. + if isinstance(inputs, set) or isinstance(outputs, set): + warnings.warn("Using sets as inputs is discouraged as it leads to indeterministic behavior.") + if isinstance(inputs, (set, collections.abc.KeysView, collections.abc.Set)): inputs = {k: None for k in inputs} - if isinstance(outputs, (set, collections.abc.KeysView)): + if isinstance(outputs, (set, collections.abc.KeysView, collections.abc.Set)): outputs = {k: None for k in outputs} s = nd.NestedSDFG( @@ -1802,7 +1896,7 @@ def add_map( ndrange: Union[Dict[str, Union[str, sbs.Subset]], List[Tuple[str, Union[str, sbs.Subset]]]], schedule=dtypes.ScheduleType.Default, unroll=False, - debuginfo=None, + debuginfo: Optional[dtypes.DebugInfo] = None, ) -> Tuple[nd.MapEntry, nd.MapExit]: """ Adds a map entry and map exit. @@ -1814,21 +1908,23 @@ def add_map( :return: (map_entry, map_exit) node 2-tuple """ - debuginfo = _getdebuginfo(debuginfo or self._default_lineinfo) + debuginfo = _get_debug_info(debuginfo or self._default_lineinfo) map = nd.Map(name, *_make_iterators(ndrange), schedule=schedule, unroll=unroll, debuginfo=debuginfo) map_entry = nd.MapEntry(map) map_exit = nd.MapExit(map) self.add_nodes_from([map_entry, map_exit]) return map_entry, map_exit - def add_consume(self, - name, - elements: Tuple[str, str], - condition: str = None, - schedule=dtypes.ScheduleType.Default, - chunksize=1, - debuginfo=None, - language=dtypes.Language.Python) -> Tuple[nd.ConsumeEntry, nd.ConsumeExit]: + def add_consume( + self, + name, + elements: Tuple[str, str], + condition: str = None, + schedule=dtypes.ScheduleType.Default, + chunksize=1, + debuginfo: Optional[dtypes.DebugInfo] = None, + language=dtypes.Language.Python, + ) -> Tuple[nd.ConsumeEntry, nd.ConsumeExit]: """ Adds consume entry and consume exit nodes. :param name: Label @@ -1850,7 +1946,7 @@ def add_consume(self, "(PE_index, num_PEs)") pe_tuple = (elements[0], SymbolicProperty.from_string(elements[1])) - debuginfo = _getdebuginfo(debuginfo or self._default_lineinfo) + debuginfo = _get_debug_info(debuginfo or self._default_lineinfo) if condition is not None: condition = CodeBlock(condition, language) consume = nd.Consume(name, pe_tuple, condition, schedule, chunksize, debuginfo=debuginfo) @@ -1860,24 +1956,23 @@ def add_consume(self, self.add_nodes_from([entry, exit]) return entry, exit - def add_mapped_tasklet(self, - name: str, - map_ranges: Union[Dict[str, Union[str, sbs.Subset]], List[Tuple[str, Union[str, - sbs.Subset]]]], - inputs: Dict[str, mm.Memlet], - code: str, - outputs: Dict[str, mm.Memlet], - schedule=dtypes.ScheduleType.Default, - unroll_map=False, - location=None, - language=dtypes.Language.Python, - debuginfo=None, - external_edges=False, - input_nodes: Optional[Union[Dict[str, nd.AccessNode], List[nd.AccessNode], - Set[nd.AccessNode]]] = None, - output_nodes: Optional[Union[Dict[str, nd.AccessNode], List[nd.AccessNode], - Set[nd.AccessNode]]] = None, - propagate=True) -> Tuple[nd.Tasklet, nd.MapEntry, nd.MapExit]: + def add_mapped_tasklet( + self, + name: str, + map_ranges: Union[Dict[str, Union[str, sbs.Subset]], List[Tuple[str, Union[str, sbs.Subset]]]], + inputs: Dict[str, mm.Memlet], + code: str, + outputs: Dict[str, mm.Memlet], + schedule=dtypes.ScheduleType.Default, + unroll_map=False, + location=None, + language=dtypes.Language.Python, + debuginfo: Optional[dtypes.DebugInfo] = None, + external_edges=False, + input_nodes: Optional[Union[Dict[str, nd.AccessNode], List[nd.AccessNode], Set[nd.AccessNode]]] = None, + output_nodes: Optional[Union[Dict[str, nd.AccessNode], List[nd.AccessNode], Set[nd.AccessNode]]] = None, + propagate=True, + ) -> Tuple[nd.Tasklet, nd.MapEntry, nd.MapExit]: """ Convenience function that adds a map entry, tasklet, map exit, and the respective edges to external arrays. @@ -1910,7 +2005,7 @@ def add_mapped_tasklet(self, :return: tuple of (tasklet, map_entry, map_exit) """ map_name = name + "_map" - debuginfo = _getdebuginfo(debuginfo or self._default_lineinfo) + debuginfo = _get_debug_info(debuginfo or self._default_lineinfo) # Create appropriate dictionaries from inputs tinputs = {k: None for k, v in inputs.items()} @@ -1973,7 +2068,7 @@ def add_mapped_tasklet(self, for inp, inpnode in sorted(inpdict.items()): # Add external edge if propagate: - outer_memlet = propagate_memlet(self, tomemlet[inp], map_entry, True) + outer_memlet = sdprop.propagate_memlet(self, tomemlet[inp], map_entry, True) else: outer_memlet = tomemlet[inp] edges.append(self.add_edge(inpnode, None, map_entry, "IN_" + inp, outer_memlet)) @@ -2004,7 +2099,7 @@ def add_mapped_tasklet(self, for out, outnode in sorted(outdict.items()): # Add external edge if propagate: - outer_memlet = propagate_memlet(self, tomemlet[out], map_exit, True) + outer_memlet = sdprop.propagate_memlet(self, tomemlet[out], map_exit, True) else: outer_memlet = tomemlet[out] edges.append(self.add_edge(map_exit, "OUT_" + out, outnode, None, outer_memlet)) @@ -2030,7 +2125,7 @@ def add_reduce( axes, identity=None, schedule=dtypes.ScheduleType.Default, - debuginfo=None, + debuginfo: Optional[dtypes.DebugInfo] = None, ) -> 'dace.libraries.standard.Reduce': """ Adds a reduction node. @@ -2044,7 +2139,7 @@ def add_reduce( :return: A Reduce node """ import dace.libraries.standard as stdlib # Avoid import loop - debuginfo = _getdebuginfo(debuginfo or self._default_lineinfo) + debuginfo = _get_debug_info(debuginfo or self._default_lineinfo) result = stdlib.Reduce('Reduce', wcr, axes, identity, schedule=schedule, debuginfo=debuginfo) self.add_node(result) return result @@ -2132,7 +2227,7 @@ def add_edge_pair( # Add external edge if external_memlet is None: # If undefined, propagate - external_memlet = propagate_memlet(self, internal_memlet, scope_node, True) + external_memlet = sdprop.propagate_memlet(self, internal_memlet, scope_node, True) if isinstance(scope_node, nd.EntryNode): eedge = self.add_edge( @@ -2253,7 +2348,7 @@ def add_memlet_path(self, *path_nodes, memlet=None, src_conn=None, dst_conn=None snode = edge.dst if propagate_forward else edge.src if not cur_memlet.is_empty(): if propagate: - cur_memlet = propagate_memlet(self, cur_memlet, snode, True) + cur_memlet = sdprop.propagate_memlet(self, cur_memlet, snode, True) # Try to initialize memlets for edge in edges: edge.data.try_initialize(self.sdfg, self, edge) @@ -2599,13 +2694,14 @@ class AbstractControlFlowRegion(OrderedDiGraph[ControlFlowBlock, 'dace.sdfg.Inte ControlFlowBlock, abc.ABC): """ Abstract superclass to represent all kinds of control flow regions in an SDFG. - This is consequently one of the three main classes of control flow graph nodes, which include ``ControlFlowBlock``s, - ``SDFGState``s, and nested ``AbstractControlFlowRegion``s. An ``AbstractControlFlowRegion`` can further be either a - region that directly contains a control flow graph (``ControlFlowRegion``s and subclasses thereof), or something + + This is consequently one of the three main classes of control flow graph nodes, which include ``ControlFlowBlock`` s, + ``SDFGState`` s, and nested ``AbstractControlFlowRegion`` s. An ``AbstractControlFlowRegion`` can further be either a + region that directly contains a control flow graph (``ControlFlowRegion`` s and subclasses thereof), or something that acts like and has the same utilities as a control flow region, including the same API, but is itself not directly a single graph. An example of this is the ``ConditionalBlock``, which acts as a single control flow region to the outside, but contains multiple actual graphs (one per branch). As such, there are very few but important - differences between the subclasses of ``AbstractControlFlowRegion``s, such as how traversals are performed, how many + differences between the subclasses of ``AbstractControlFlowRegion`` s, such as how traversals are performed, how many start blocks there are, etc. """ @@ -2646,6 +2742,44 @@ def replace_meta_accesses(self, replacements: Dict[str, str]) -> None: """ pass + def propagate_memlets(self, border_memlets: Dict[str, Dict[str, Optional[mm.Memlet]]]) -> None: + """ + Propagate child-block memlets to this region boundary. + + The default implementation treats the region conservatively as + straight-line control flow: it collects memlets from child states and + nested regions, then merges them as if every collected contribution may + reach the region boundary. + + :param border_memlets: A mapping from connector direction and name to + the accumulated border memlet. The mapping is + updated in-place. + :note: ``border_memlets`` mapping is updated in-place. + """ + from dace.sdfg import propagation as sdprop + + candidates = sdprop._make_border_memlets(border_memlets, as_lists=True) + + for block in self.nodes(): + if isinstance(block, SDFGState): + sdprop._collect_state_border_memlet_candidates(block, candidates) + elif isinstance(block, AbstractControlFlowRegion): + nested_memlets = sdprop._make_border_memlets(border_memlets) + block.propagate_memlets(nested_memlets) + sdprop._append_border_memlet_candidates(candidates, nested_memlets) + + for direction in border_memlets: + for connector in border_memlets[direction]: + propagated = sdprop._propagate_border_memlet_candidates(candidates, self.sdfg.arrays, direction, + connector) + if propagated is None: + continue + + array_name = propagated.data if propagated.data is not None else connector + border_memlets[direction][connector] = sdprop._merge_border_memlet(border_memlets[direction][connector], + propagated, + self.sdfg.arrays[array_name]) + @property def root_sdfg(self) -> 'SDFG': from dace.sdfg.sdfg import SDFG # Avoid import loop @@ -3046,7 +3180,7 @@ def from_json(cls, json_obj, context=None): ret = cls(label=json_obj['label'], sdfg=context['sdfg']) - dace.serialize.set_properties_from_json(ret, json_obj) + dace.serialize.set_properties_from_json(ret, json_obj, context=context) nodelist = [] for n in nodes: @@ -3058,7 +3192,7 @@ def from_json(cls, json_obj, context=None): nodelist.append(block) for e in edges: - e = dace.serialize.from_json(e) + e = dace.serialize.from_json(e, context=context) ret.add_edge(nodelist[int(e.src)], nodelist[int(e.dst)], e.data) if 'start_block' in json_obj: @@ -3121,6 +3255,7 @@ def start_block(self, block_id): class ControlFlowRegion(AbstractControlFlowRegion): """ A ``ControlFlowRegion`` represents a control flow graph node that itself contains a control flow graph. + This can be an arbitrary control flow graph, but may also be a specific type of control flow region with additional semantics, such as a loop or a function call. """ @@ -3185,34 +3320,27 @@ def __init__(self, update_expr: Optional[Union[str, CodeBlock]] = None, inverted: bool = False, sdfg: Optional['SDFG'] = None, - update_before_condition=True, + update_before_condition: bool = True, unroll: bool = False, unroll_factor: int = 0): super(LoopRegion, self).__init__(label, sdfg) - if initialize_expr is not None: - if isinstance(initialize_expr, CodeBlock): - self.init_statement = initialize_expr - else: - self.init_statement = CodeBlock(initialize_expr) + if initialize_expr is None or isinstance(initialize_expr, CodeBlock): + self.init_statement = initialize_expr else: - self.init_statement = None + self.init_statement = CodeBlock(initialize_expr) - if condition_expr: - if isinstance(condition_expr, CodeBlock): - self.loop_condition = condition_expr - else: - self.loop_condition = CodeBlock(condition_expr) - else: + if condition_expr is None: self.loop_condition = CodeBlock('True') + elif isinstance(condition_expr, CodeBlock): + self.loop_condition = condition_expr + else: + self.loop_condition = CodeBlock(condition_expr) - if update_expr is not None: - if isinstance(update_expr, CodeBlock): - self.update_statement = update_expr - else: - self.update_statement = CodeBlock(update_expr) + if update_expr is None or isinstance(update_expr, CodeBlock): + self.update_statement = update_expr else: - self.update_statement = None + self.update_statement = CodeBlock(update_expr) self.loop_variable = loop_var or '' self.inverted = inverted @@ -3506,14 +3634,23 @@ def get_meta_codeblocks(self): codes.append(self.update_statement) return codes - def get_meta_read_memlets(self) -> List[mm.Memlet]: + def get_meta_read_memlets(self, arrays: Optional[Dict[str, dt.Data]] = None) -> List[mm.Memlet]: + """ + Get a list of all (read) memlets in meta codeblocks. + + :param arrays: An optional dictionary mapping array names to their data descriptors. + If not not given defaults to ``self.sdfg.arrays``. + """ # Avoid cyclic imports. from dace.sdfg.sdfg import memlets_in_ast - read_memlets = memlets_in_ast(self.loop_condition.code[0], self.sdfg.arrays) + + arrays = arrays if arrays is not None else self.sdfg.arrays + + read_memlets = memlets_in_ast(self.loop_condition.code[0], arrays) if self.init_statement: - read_memlets.extend(memlets_in_ast(self.init_statement.code[0], self.sdfg.arrays)) + read_memlets.extend(memlets_in_ast(self.init_statement.code[0], arrays)) if self.update_statement: - read_memlets.extend(memlets_in_ast(self.update_statement.code[0], self.sdfg.arrays)) + read_memlets.extend(memlets_in_ast(self.update_statement.code[0], arrays)) return read_memlets def replace_meta_accesses(self, replacements): @@ -3525,6 +3662,106 @@ def replace_meta_accesses(self, replacements): if self.update_statement: replace_in_codeblock(self.update_statement, replacements) + def propagate_memlets(self, border_memlets: Dict[str, Dict[str, Optional[mm.Memlet]]]) -> None: + """ + Propagate memlets across a loop region boundary. + + When the loop trip count can be derived from the loop metadata, the + region propagates its candidate memlets through that symbolic iteration + range. Inverted loops are split into their unconditional first + iteration and the remaining iterations. Loops that cannot be analyzed + precisely fall back to the conservative base-region behavior. + + :param border_memlets: A mapping from connector direction and name to + the accumulated border memlet. The mapping is + updated in-place. + :note: ``border_memlets`` mapping is updated in-place. + """ + # Avoid cyclic import + from dace.transformation.passes.analysis import loop_analysis + + if self.has_break: + super().propagate_memlets(border_memlets) + return + + init = loop_analysis.get_init_assignment(self) + end = loop_analysis.get_loop_end(self) + stride = loop_analysis.get_loop_stride(self) + if not self.loop_variable or init is None or end is None or stride in (None, 0): + super().propagate_memlets(border_memlets) + return + + candidates = sdprop._make_border_memlets(border_memlets, as_lists=True) + + for block in self.nodes(): + if isinstance(block, SDFGState): + sdprop._collect_state_border_memlet_candidates(block, candidates) + elif isinstance(block, AbstractControlFlowRegion): + nested_memlets = sdprop._make_border_memlets(border_memlets) + block.propagate_memlets(nested_memlets) + sdprop._append_border_memlet_candidates(candidates, nested_memlets) + + def _range_is_definitely_empty(start, stop) -> bool: + """Returns True only when the remaining inverted-loop range is provably empty.""" + simplified_start = symbolic.simplify(start) + simplified_stop = symbolic.simplify(stop) + simplified_stride = symbolic.simplify(stride) + if any( + getattr(expr, 'free_symbols', set()) + for expr in (simplified_start, simplified_stop, simplified_stride)): + return False + if simplified_stride > 0: + return (simplified_start > simplified_stop) == True + return (simplified_start < simplified_stop) == True + + def _propagate_range(start, stop, target_memlets: Dict[str, Dict[str, Optional[mm.Memlet]]]) -> None: + """Propagates all candidate memlets through one loop-iteration range.""" + loop_range = Range([(start, stop, stride)]) + for direction in target_memlets: + for connector in target_memlets[direction]: + propagated = sdprop._propagate_border_memlet_candidates( + candidates, + self.sdfg.arrays, + direction, + connector, + params=[self.loop_variable], + rng=loop_range, + scale_by_range=True, + ) + if propagated is None: + continue + + array_name = propagated.data if propagated.data is not None else connector + target_memlets[direction][connector] = sdprop._merge_border_memlet( + target_memlets[direction][connector], propagated, self.sdfg.arrays[array_name]) + + if not self.inverted: + _propagate_range(init, end, border_memlets) + return + + first_iteration_memlets = sdprop._make_border_memlets(border_memlets) + _propagate_range(init, init, first_iteration_memlets) + + remaining_iterations_memlets = sdprop._make_border_memlets(border_memlets) + remaining_start = symbolic.simplify(init + stride) + # For inverted loops with update-before-condition disabled, the body can + # execute once at ``end + stride`` before the termination condition is + # observed, so the propagated range must include that final iteration. + remaining_end = end if self.update_before_condition else symbolic.simplify(end + stride) + if not _range_is_definitely_empty(remaining_start, remaining_end): + _propagate_range(remaining_start, remaining_end, remaining_iterations_memlets) + + for direction in border_memlets: + for connector in border_memlets[direction]: + for propagated in (first_iteration_memlets[direction][connector], + remaining_iterations_memlets[direction][connector]): + if propagated is None: + continue + + array_name = propagated.data if propagated.data is not None else connector + border_memlets[direction][connector] = sdprop._merge_border_memlet( + border_memlets[direction][connector], propagated, self.sdfg.arrays[array_name]) + def _used_symbols_internal(self, all_symbols: bool, defined_syms: Optional[Set] = None, @@ -3536,8 +3773,10 @@ def _used_symbols_internal(self, free_syms = set() if free_syms is None else free_syms used_before_assignment = set() if used_before_assignment is None else used_before_assignment - defined_syms.add(self.loop_variable) if self.init_statement is not None: + # Loops with no initialization statement do not redefine the loop variable + defined_syms.add(self.loop_variable) + free_syms |= self.init_statement.get_free_symbols() if self.update_statement is not None: free_syms |= self.update_statement.get_free_symbols() @@ -3678,6 +3917,42 @@ def get_meta_read_memlets(self) -> List[mm.Memlet]: read_memlets.extend(memlets_in_ast(c.code[0], self.sdfg.arrays)) return read_memlets + def propagate_memlets(self, border_memlets: Dict[str, Dict[str, Optional[mm.Memlet]]]) -> None: + """ + Propagate memlets across a conditional region boundary. + + Each branch is propagated independently and then merged into a single + boundary memlet. Because at most one branch executes per traversal, the + merged subset is the union of branch subsets while the merged volume is + an upper bound over branch volumes instead of their sum. + + :param border_memlets: A mapping from connector direction and name to + the accumulated border memlet. The mapping is + updated in place. + :note: ``border_memlets`` mapping is updated in-place. + """ + from dace.sdfg import propagation as sdprop + has_condition = False + + for condition, region in self._branches: + has_condition = has_condition or condition is not None + branch_memlets = sdprop._make_border_memlets(border_memlets) + region.propagate_memlets(branch_memlets) + + for direction in border_memlets: + for connector in border_memlets[direction]: + propagated = branch_memlets[direction][connector] + if propagated is None: + continue + + if has_condition: + propagated = copy.deepcopy(propagated) + propagated.dynamic = True + + array_name = propagated.data if propagated.data is not None else connector + border_memlets[direction][connector] = sdprop._merge_border_memlet_upper_bound( + border_memlets[direction][connector], propagated, self.sdfg.arrays[array_name]) + def _used_symbols_internal(self, all_symbols: bool, defined_syms: Optional[Set] = None, @@ -3722,6 +3997,8 @@ def replace_dict(self, def to_json(self, parent=None): json = super().to_json(parent) + del json['nodes'] + del json['edges'] json['branches'] = [(condition.to_json() if condition is not None else None, cfg.to_json()) for condition, cfg in self._branches] return json @@ -3735,11 +4012,11 @@ def from_json(cls, json_obj, context=None): ret = cls(label=json_obj['label'], sdfg=context['sdfg']) - dace.serialize.set_properties_from_json(ret, json_obj) + dace.serialize.set_properties_from_json(ret, json_obj, context=context) for condition, region in json_obj['branches']: if condition is not None: - ret.add_branch(CodeBlock.from_json(condition), ControlFlowRegion.from_json(region, context)) + ret.add_branch(CodeBlock.from_json(condition, context), ControlFlowRegion.from_json(region, context)) else: ret.add_branch(None, ControlFlowRegion.from_json(region, context)) return ret @@ -3853,7 +4130,7 @@ def __repr__(self): @make_properties class NamedRegion(ControlFlowRegion): - debuginfo = DebugInfoProperty() + debuginfo = DebugInfoProperty(allow_none=True) def __init__(self, label: str, sdfg: Optional['SDFG'] = None, debuginfo: Optional[dtypes.DebugInfo] = None): super().__init__(label, sdfg) diff --git a/dace/sdfg/type_inference.py b/dace/sdfg/type_inference.py index b252181dc5..b278b5448e 100644 --- a/dace/sdfg/type_inference.py +++ b/dace/sdfg/type_inference.py @@ -15,7 +15,6 @@ import sympy import sys import dace.frontend.python.astutils -import inspect from typing import Callable, Union # Additional function names that can be used to infer types @@ -109,10 +108,21 @@ def infer_expr_type(code, symbols=None): if code.name in symbols: return symbols[code.name] return code.dtype + elif isinstance(code, symbolic.TypedConstant): + return code.dtype elif isinstance(code, sympy.Basic): - parsed_ast = ast.parse(sympy.printing.pycode(code, allow_unknown_functions=True)) + has_typed_constant = any(isinstance(node, symbolic.TypedConstant) for node in sympy.preorder_traversal(code)) + if has_typed_constant: + parsed_ast = ast.parse(symstr(code)) + else: + parsed_ast = ast.parse(sympy.printing.pycode(code, allow_unknown_functions=True)) elif isinstance(code, SymExpr): - parsed_ast = ast.parse(sympy.printing.pycode(code.expr, allow_unknown_functions=True)) + has_typed_constant = any( + isinstance(node, symbolic.TypedConstant) for node in sympy.preorder_traversal(code.expr)) + if has_typed_constant: + parsed_ast = ast.parse(symstr(code.expr)) + else: + parsed_ast = ast.parse(sympy.printing.pycode(code.expr, allow_unknown_functions=True)) else: raise TypeError(f"Cannot convert type {type(code)} to a Python AST.") diff --git a/dace/sdfg/utils.py b/dace/sdfg/utils.py index 4699c3a5f2..782e98d40d 100644 --- a/dace/sdfg/utils.py +++ b/dace/sdfg/utils.py @@ -1,15 +1,14 @@ -# Copyright 2019-2025 ETH Zurich and the DaCe authors. All rights reserved. +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. """ Various utility functions to create, traverse, and modify SDFGs. """ import collections import copy -import os import warnings import networkx as nx import time import dace.sdfg.nodes -from dace.codegen import compiled_sdfg as csdfg +from dace.codegen import compiled_sdfg as csdfg, compiler as sdfg_compiler from dace.sdfg.graph import MultiConnectorEdge from dace.sdfg.sdfg import SDFG, InterstateEdge from dace.sdfg.nodes import Node, NestedSDFG @@ -1612,7 +1611,7 @@ def inline_sdfgs(sdfg: SDFG, permissive: bool = False, progress: bool = None, mu inaccurate, requires ``tqdm``). If None, prints out progress if over 5 seconds have passed. If False, never shows progress bar. - :param multistate: Include + :param multistate: If True, include multi-state inlining. :return: The total number of SDFGs inlined. """ # Avoid import loops @@ -1662,33 +1661,27 @@ def inline_sdfgs(sdfg: SDFG, permissive: bool = False, progress: bool = None, mu return counter -def load_precompiled_sdfg(folder: str, argnames: Optional[List[str]] = None) -> csdfg.CompiledSDFG: - """ - Loads a pre-compiled SDFG from an output folder (e.g. ".dacecache/program"). - Folder must contain a file called "program.sdfg" and a subfolder called - "build" with the shared object. +def load_precompiled_sdfg(*args, **kwargs) -> csdfg.CompiledSDFG: - :param folder: Path to SDFG output folder. - :param argnames: Names of arguments of the compiled SDFG. - :return: A callable CompiledSDFG object. - """ - sdfg = SDFG.from_file(os.path.join(folder, 'program.sdfg')) - suffix = config.Config.get('compiler', 'library_extension') - return csdfg.CompiledSDFG( - sdfg, - csdfg.ReloadableDLL(os.path.join(folder, 'build', f'lib{sdfg.name}.{suffix}'), sdfg.name), - argnames, + warnings.warn( + 'Used deprecated ``dace.sdfg.utils.load_precompiled_sdfg()`` function, use the one from ``dace.codegen.compiler`` instead.', + category=DeprecationWarning, + stacklevel=2, ) + return sdfg_compiler.load_precompiled_sdfg(*args, **kwargs) + -def distributed_compile(sdfg: SDFG, comm, validate: bool = True) -> csdfg.CompiledSDFG: +def distributed_compile(sdfg: SDFG, comm, *, validate: bool = True) -> csdfg.CompiledSDFG: """ Compiles an SDFG in rank 0 of MPI communicator ``comm``. Then, the compiled SDFG is loaded in all other ranks. :param sdfg: SDFG to be compiled. :param comm: MPI communicator. ``Intracomm`` is the base mpi4py communicator class. + :param validate: If True, validates the SDFG prior to generating code. :return: Compiled SDFG. :note: This method can be used only if the module mpi4py is installed. + :todo: Relocate this function to `dace.codegen.compiler`. """ rank = comm.Get_rank() @@ -1705,7 +1698,7 @@ def distributed_compile(sdfg: SDFG, comm, validate: bool = True) -> csdfg.Compil # Loads compiled SDFG. if rank > 0: - func = load_precompiled_sdfg(folder) + func = sdfg_compiler.load_precompiled_sdfg(folder) comm.Barrier() @@ -1938,6 +1931,7 @@ def map_view_to_array(vdesc: dt.View, adesc: dt.Array, out of a view, we need to compose the subset with the new view dimensions and new subset. The precondition to this method is that the array has unique strides (if not, the process fails). The process works in three steps, as follows: + * First, The degenerate (shape=1) dimensions are removed from both the array and the view for consideration. * The mapping between non-degenerate dimensions is done from the view to the array based on the strides. Note that in a slice, the strides can be expanded or squeezed, but never reordered. This fact is used @@ -1948,9 +1942,11 @@ def map_view_to_array(vdesc: dt.View, adesc: dt.Array, * Third, the remainder of the dimensions of the original (non-view) data descriptor are considered "squeezed". + For example, a scalar view ``A[i, j] -> v`` would return ``({}, [], [0, 1])``. Example 2: ``A[0:2, 3:5, i, j, 0:N] -> V[0:2, 0, 0:2, 0, 0:N, 0]`` would return ``({0: 0, 2: 1, 3: 2, 4: 4}, [1, 5], [3])``. + :param vdesc: The data descriptor of the view. :param adesc: The data descriptor of the viewed data container. :return: A tuple of (mapping of view->array, expanded, squeezed) dimensions, or None if the process failed. diff --git a/dace/sdfg/validation.py b/dace/sdfg/validation.py index c1110ec39b..888c7e77c9 100644 --- a/dace/sdfg/validation.py +++ b/dace/sdfg/validation.py @@ -9,7 +9,7 @@ import networkx as nx -from dace import dtypes, subsets, symbolic, data +from dace import dtypes, subsets, symbolic from dace.dtypes import DebugInfo if TYPE_CHECKING: @@ -255,11 +255,11 @@ def validate_sdfg(sdfg: 'dace.sdfg.SDFG', references: Set[int] = None, **context # Check the names of data descriptors and co. seen_names: Set[str] = set() - for obj_names in [sdfg.arrays.keys(), sdfg.symbols.keys(), sdfg._rdistrarrays.keys(), sdfg._subarrays.keys()]: + for obj_names in [sdfg.arrays.keys(), sdfg.symbols.keys()]: if not seen_names.isdisjoint(obj_names): raise InvalidSDFGError( f'Found duplicated names: "{seen_names.intersection(obj_names)}". Please ensure ' - 'that the names of symbols, data descriptors, subarrays and rdistarrays are unique.', sdfg, None) + 'that the names of symbols and data descriptors are unique.', sdfg, None) seen_names.update(obj_names) # Ensure that there is a mentioning of constants in either the array or symbol. @@ -413,7 +413,6 @@ def validate_state(state: 'dace.sdfg.SDFGState', from dace import data as dt from dace import subsets as sbs from dace.config import Config - from dace.sdfg import SDFG from dace.sdfg import nodes as nd from dace.sdfg import utils as sdutil from dace.sdfg.scope import is_devicelevel_gpu, scope_contains_scope @@ -469,7 +468,7 @@ def validate_state(state: 'dace.sdfg.SDFGState', except InvalidSDFGError: raise except Exception as ex: - raise InvalidSDFGNodeError("Node validation failed: " + str(ex), sdfg, state_id, nid) from ex + raise InvalidSDFGNodeError(f"Node validation failed: {ex}", sdfg, state_id, nid) from ex # Isolated nodes ######################################## @@ -675,6 +674,15 @@ def validate_state(state: 'dace.sdfg.SDFGState', 'rather than using multiple references to the same one', sdfg, state_id, eid) references.add(id(e.data)) + for subset in (e.data.subset, e.data.other_subset): + if subset is None: + continue + if id(subset) in references: + raise InvalidSDFGEdgeError( + f'Duplicate subset detected in memlet "{e.data}". Please copy objects ' + 'rather than using multiple references to the same one', sdfg, state_id, eid) + references.add(id(subset)) + # Edge validation try: e.data.validate(sdfg, state) @@ -1015,8 +1023,7 @@ def to_json(self): def __str__(self): if self.edge_id is not None: e = self.sdfg.edges()[self.edge_id] - edgestr = ' (at edge "%s" (%s -> %s)' % ( - e.data.label, + edgestr = ' (at edge %s -> %s)' % ( str(e.src), str(e.dst), ) diff --git a/dace/serialize.py b/dace/serialize.py index a96f776881..6378b8a760 100644 --- a/dace/serialize.py +++ b/dace/serialize.py @@ -1,8 +1,10 @@ -# Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. +import copy from dataclasses import is_dataclass import enum import json import numpy as np +import sympy import warnings from dace import attr_enum import dace.dtypes @@ -109,6 +111,12 @@ def to_json(obj): elif isinstance(obj, np.ndarray): # Special case for external structures (numpy arrays) return NumpySerializer.to_json(obj) + elif type(obj).__name__ == 'SymExpr': + from dace import symbolic + return symbolic.serialize_symbolic(obj) + elif isinstance(obj, sympy.Basic): + from dace import symbolic + return symbolic.serialize_symbolic(obj) elif is_dataclass(obj): # Serialize dataclass as a dictionary retval = {"type": type(obj).__name__} @@ -201,8 +209,10 @@ def loads(*args, context=None, **kwargs): return from_json(loaded, context) -def dumps(*args, **kwargs): - return json.dumps(*args, default=to_json, indent=2, **kwargs) +def dumps(*args, readable=False, **kwargs): + if readable: + return json.dumps(*args, default=to_json, indent=2, **kwargs) + return json.dumps(*args, default=to_json, indent=None, separators=(',', ':'), **kwargs) def load(*args, context=None, **kwargs): @@ -210,16 +220,22 @@ def load(*args, context=None, **kwargs): return from_json(loaded, context) -def dump(*args, **kwargs): - return json.dump(*args, default=to_json, indent=2, **kwargs) +def dump(*args, readable=False, **kwargs): + if readable: + return json.dump(*args, default=to_json, indent=2, **kwargs) + return json.dump(*args, default=to_json, indent=None, separators=(',', ':'), **kwargs) def all_properties_to_json(object_with_properties): save_all_fields = config.Config.get_bool('testing', 'serialize_all_fields') retdict = {} for x, v in object_with_properties.properties(): - if not save_all_fields and v == x.default: # Skip default fields - continue + if not save_all_fields: + is_default = (v == x.default) + if isinstance(is_default, np.ndarray): + is_default = np.all(is_default) + if is_default: # Skip default fields + continue if not x.serialize_if(object_with_properties): continue retdict[x.attr_name] = x.to_json(v) @@ -241,34 +257,38 @@ def set_properties_from_json(object_with_properties, json_obj, context=None, ign if prop_name in ignore_properties: continue + missing_prop = False try: val = attrs[prop_name] # Make sure we use all properties source_properties.remove(prop_name) except KeyError: + missing_prop = True # Allow a property to not be set if it has a default value # TODO: is this really the job of serialize? if prop.default is not None: - val = prop.default + # Prevent shared mutable defaults from being aliased across deserialized objects. + val = copy.deepcopy(prop.default) elif prop.allow_none: val = None else: raise KeyError("Missing property for object of type " + type(object_with_properties).__name__ + ": " + prop_name) - if isinstance(val, dict): - val = prop.from_json(val, context) - else: - try: + if not missing_prop: + if isinstance(val, dict): val = prop.from_json(val, context) - except TypeError as err: - # TODO: This seems to be called both from places where the - # dictionary has been fully deserialized, and on raw json - # objects. In the interest of time, we're not failing here, but - # should untangle this eventually - warnings.warn("Failed to parse object {}" - " for property {} of type {}. Error was: {}".format(val, prop_name, prop, err)) - raise + else: + try: + val = prop.from_json(val, context) + except TypeError as err: + # TODO: This seems to be called both from places where the + # dictionary has been fully deserialized, and on raw json + # objects. In the interest of time, we're not failing here, but + # should untangle this eventually + warnings.warn("Failed to parse object {}" + " for property {} of type {}. Error was: {}".format(val, prop_name, prop, err)) + raise setattr(object_with_properties, prop_name, val) diff --git a/dace/sourcemap.py b/dace/sourcemap.py index 1b5eb4cb9e..334835f85d 100644 --- a/dace/sourcemap.py +++ b/dace/sourcemap.py @@ -4,15 +4,22 @@ import json import os import socket -from typing import Optional +from typing import Optional, TypedDict from dace import Config, dtypes from dace.sdfg import state from dace.sdfg import nodes +class NodeInfo(TypedDict): + debuginfo: dict[str, int | str] + cfg_id: int + state_id: int + node_id: int + + class SdfgLocation: - def __init__(self, cfg_id, state_id, node_ids): + def __init__(self, cfg_id: int, state_id: int, node_ids: list[int]): self.cfg_id = cfg_id self.state_id = state_id self.node_ids = node_ids @@ -64,7 +71,7 @@ def send(data: json): HOST = socket.gethostname() PORT = os.environ["DACE_port"] - data_bytes = bytes(json.dumps(data), "utf-8") + data_bytes = bytes(json.dumps(data, indent=None, separators=(',', ':')), "utf-8") with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: s.connect((HOST, int(PORT))) s.sendall(data_bytes) @@ -85,7 +92,7 @@ def save(language: str, name: str, map: dict, build_folder: str) -> str: path = os.path.abspath(os.path.join(folder, 'map', 'map_' + language + '.json')) with open(path, "w") as json_file: - json.dump(map, json_file, indent=4) + json.dump(map, json_file, indent=None, separators=(',', ':')) return os.path.abspath(folder) @@ -310,10 +317,10 @@ def divide(self): """ divided = [] for dbinfo in self.debuginfo: - source = dbinfo['debuginfo']['filename'] + source = dbinfo['debuginfo'].get('filename') exists = False for src_infos in divided: - if len(src_infos) > 0 and src_infos[0]['debuginfo']['filename'] == source: + if len(src_infos) > 0 and src_infos[0]['debuginfo'].get('filename') == source: src_infos.append(dbinfo) exists = True break @@ -328,11 +335,12 @@ def sorter(self): for dbinfo_source in self.debuginfo: db_sorted.append( sorted(dbinfo_source, - key=lambda n: (n['debuginfo']['start_line'], n['debuginfo']['start_column'], n['debuginfo'][ - 'end_line'], n['debuginfo']['end_column']))) + key=lambda n: + (n['debuginfo'].get('start_line', 0), n['debuginfo'].get('start_column', 0), n['debuginfo'].get( + 'end_line', 0), n['debuginfo'].get('end_column', 0)))) return db_sorted - def make_info(self, debuginfo, node_id: int, state_id: int, cfg_id: int) -> dict: + def make_info(self, debuginfo, node_id: int, state_id: int, cfg_id: int) -> NodeInfo: """ Creates an object for the current node with the most important information @@ -345,7 +353,7 @@ def make_info(self, debuginfo, node_id: int, state_id: int, cfg_id: int) -> dict return {"debuginfo": debuginfo, "cfg_id": cfg_id, "state_id": state_id, "node_id": node_id} def sdfg_debuginfo(self, graph, cfg_id: int = 0, state_id: int = 0): - """ Recursively retracts all debuginfo from the nodes + """ Recursively extracts all debuginfo from the nodes :param graph: An SDFG or SDFGState to check for nodes :param cfg_id: Id of the current SDFG/NestedSDFG @@ -387,10 +395,10 @@ def create_mapping(self, range_dict=None): """ for file_dbinfo in self.debuginfo: for node in file_dbinfo: - src_file = node["debuginfo"]["filename"] + src_file = node["debuginfo"].get("filename") if not src_file in self.map: self.map[src_file] = {} - for line in range(node["debuginfo"]["start_line"], node["debuginfo"]["end_line"] + 1): + for line in range(node["debuginfo"].get("start_line", 0), node["debuginfo"].get("end_line", 0) + 1): # Maps a python line to a list of nodes # The nodes have been sorted by priority if not str(line) in self.map[src_file]: diff --git a/dace/subsets.py b/dace/subsets.py index c61521a049..3951d91491 100644 --- a/dace/subsets.py +++ b/dace/subsets.py @@ -17,21 +17,21 @@ def nng(expr): def bounding_box_cover_exact(subset_a, subset_b, approximation=False) -> bool: - """Test if `subset_a` covers `subset_b`. + """Test if ``subset_a`` covers ``subset_b``. - The function uses a bounding box to test if `subset_a` covers `subset_b`, - i.e. that `subset_a` is at least as big as `subset_b`. By default the - box is constructed using `{min, max}_element()` or if `approximation` is - `True` `{min, max}_element_approx()`. The most important difference compared - to `bounding_box_cover_exact()` is that this function does not assume + The function uses a bounding box to test if ``subset_a`` covers ``subset_b``, + i.e. that ``subset_a`` is at least as big as ``subset_b``. By default the + box is constructed using ``{min, max}_element()`` or if ``approximation`` is + ``True`` ``{min, max}_element_approx()``. The most important difference compared + to ``bounding_box_cover_exact()`` is that this function does not assume that the symbols are positive. - The function returns `True` if it _can be shown_ that `subset_a` covers `subset_b` - and `False` otherwise. + The function returns ``True`` if it can be shown that ``subset_a`` covers ``subset_b`` + and ``False`` otherwise. :param subset_a: The first subset, the one that should cover. :param subset_b: The second subset, the one that should be covered. - :param approximation: If `True` then use the approximated bounds. + :param approximation: If ``True`` then use the approximated bounds. """ min_elements_a = subset_a.min_element_approx() if approximation else subset_a.min_element() max_elements_a = subset_a.max_element_approx() if approximation else subset_a.max_element() @@ -43,7 +43,7 @@ def bounding_box_cover_exact(subset_a, subset_b, approximation=False) -> bool: return ValueError(f"A bounding box of dimensionality {len(min_elements_a)} cannot" f" test covering a bounding box of dimensionality {len(min_elements_b)}.") - # NOTE: The original implementation always called `nng()`. However, it was decided that + # NOTE: The original implementation always called ``nng()``. However, it was decided that # this is an error and the call was removed in PR#2093. simplify = lambda expr: symbolic.simplify_ext(expr) no_simplify = lambda expr: expr @@ -58,25 +58,25 @@ def bounding_box_cover_exact(subset_a, subset_b, approximation=False) -> bool: def bounding_box_symbolic_positive(subset_a, subset_b, approximation=False) -> bool: - """Checks if `subset_a` covers `subset_b` using positivity assumption. + """Checks if ``subset_a`` covers ``subset_b`` using positivity assumption. - The function uses a bounding box to test if `subset_a` covers `subset_b`, - i.e. that `subset_a` is at least as big as `subset_b`. By default the - box is constructed using `{min, max}_element()` or if `approximation` is - `True` `{min, max}_element_approx()`. The function will perform the + The function uses a bounding box to test if ``subset_a`` covers ``subset_b``, + i.e. that ``subset_a`` is at least as big as ``subset_b``. By default the + box is constructed using ``{min, max}_element()`` or if ``approximation`` is + ``True`` ``{min, max}_element_approx()``. The function will perform the covering check under the assumption that all symbols are positive, - which is the main difference to `bounding_box_cover_exact()`. + which is the main difference to ``bounding_box_cover_exact()``. - The function returns `True` if it _can be shown_ that `subset_a` covers `subset_b` - and `False` otherwise. + The function returns ``True`` if it can be shown that ``subset_a`` covers ``subset_b`` + and ``False`` otherwise. :param subset_a: The first subset, the one that should cover. :param subset_b: The second subset, the one that should be covered. - :param approximation: If `True` then use the approximated bounds. + :param approximation: If ``True`` then use the approximated bounds. - :note: In previous versions this function raised `TypeError` in some cases + :note: In previous versions this function raised ``TypeError`` in some cases when a truth value could not be determined. This behaviour was removed, - since the `bounding_box_cover_exact()` does not show this behaviour. + since the ``bounding_box_cover_exact()`` does not show this behaviour. """ min_elements_a = subset_a.min_element_approx() if approximation else subset_a.min_element() max_elements_a = subset_a.max_element_approx() if approximation else subset_a.max_element() @@ -88,17 +88,17 @@ def bounding_box_symbolic_positive(subset_a, subset_b, approximation=False) -> b return ValueError(f"A bounding box of dimensionality {len(min_elements_a)} cannot" f" test covering a bounding box of dimensionality {len(min_elements_b)}.") - # NOTE: `nng()` is applied inside the loop. + # NOTE: ``nng()`` is applied inside the loop. simplify = lambda expr: symbolic.simplify_ext(expr) no_simplify = lambda expr: expr for rb, re, orb, ore in zip(min_elements_a, max_elements_a, min_elements_b, max_elements_b): # NOTE: Applying simplify takes a lot of time, thus we try to avoid it and try to do the test # first with the symbols we get and if we are unable to figuring out something, we run - # simplify. Furthermore, we also try to postpone `nng()` as long as we can. - # NOTE: We use first `==` in the hope that it is much faster than `<=`. - # NOTE: We have to use the `== True` test because of SymPy's behaviour. Otherwise we would - # get an expression resulting in a `TypeError`. + # simplify. Furthermore, we also try to postpone ``nng()`` as long as we can. + # NOTE: We use first ``==`` in the hope that it is much faster than ``<=``. + # NOTE: We have to use the ``== True`` test because of SymPy's behaviour. Otherwise we would + # get an expression resulting in a ``TypeError``. # lower bound: first check whether symbolic positive condition applies if not (len(rb.free_symbols) == 0 and len(orb.free_symbols) == 1): @@ -172,8 +172,8 @@ def covers_precise(self, other): elif not bounding_box_symbolic_positive(self, other): return False - # NOTE: The original implementation always called `nng()`. However, it was decided that - # and the application was made conditional on `symbolic_positive`, in PR#2093. + # NOTE: The original implementation always called ``nng()``. However, it was decided that + # and the application was made conditional on ``symbolic_positive``, in PR#2093. simplify = (lambda expr: symbolic.simplify_ext(nng(expr))) if symbolic_positive else ( lambda expr: symbolic.simplify_ext(expr)) no_simplify = lambda expr: expr @@ -219,10 +219,10 @@ def covers_precise(self, other): (simp_start % simp_fun(step) == simp_ostart % simp_fun(ostep)) == True)): return False except TypeError: - # If a `TypeError happens during the "no simplify" phase, we immediately + # If a ``TypeError`` happens during the "no simplify" phase, we immediately # go to the simplify phase, in the hope that it might be possible to # simplify the expression more. If we are already using simplify, then - # we return `False`. + # we return ``False``. if i == 0: continue return False @@ -265,7 +265,7 @@ def coord_at(self, i): :param i: A tuple of the same dimensionality as subset.dims() or subset.data_dims(). :return: Absolute coordinates for index i (length equal to - `data_dims()`, may be larger than `dims()`). + ``data_dims()``, may be larger than ``dims()``). """ raise NotImplementedError @@ -331,10 +331,7 @@ def to_json(self): ret = [] def a2s(obj): - if isinstance(obj, symbolic.SymExpr): - return {'main': str(obj.expr), 'approx': str(obj.approx)} - else: - return _simplified_str(obj) + return symbolic.serialize_symbolic(obj) for (start, end, step), tile in zip(self.ranges, self.tile_sizes): ret.append({'start': a2s(start), 'end': a2s(end), 'step': a2s(step), 'tile': a2s(tile)}) @@ -343,6 +340,7 @@ def a2s(obj): @staticmethod def from_json(obj, context=None): + from dace.properties import _symbolic_deserializer # Avoid circular import if not isinstance(obj, dict): raise TypeError("Expected dict, got {}".format(type(obj))) if obj['type'] != 'Range': @@ -352,15 +350,9 @@ def from_json(obj, context=None): ranges = obj['ranges'] tuples = [] - def p2s(x): - pts = symbolic.pystr_to_symbolic - if isinstance(x, str): - return pts(x) - else: - return symbolic.SymExpr(pts(x['main']), pts(x['approx'])) - for r in ranges: - tuples.append((p2s(r['start']), p2s(r['end']), p2s(r['step']), p2s(r['tile']))) + tuples.append((_symbolic_deserializer(r['start'], context), _symbolic_deserializer(r['end'], context), + _symbolic_deserializer(r['step'], context), _symbolic_deserializer(r['tile'], context))) return Range(tuples) @@ -368,7 +360,7 @@ def p2s(x): def from_array(array: 'dace.data.Data'): """ Constructs a range that covers the full array given as input. """ result = Range([(0, s - 1, 1) for s in array.shape]) - if any(o != 0 for o in array.offset): + if hasattr(array, 'offset') and any(o != 0 for o in array.offset): result.offset(array.offset, True) return result @@ -380,7 +372,7 @@ def __add__(self, other): ((*ranges, tile) for ranges, tile in zip(self.ranges + other.ranges, self.tile_sizes + other.tile_sizes))) def __deepcopy__(self, memo) -> 'Range': - """Performs a deepcopy of `self`. + """Performs a deepcopy of ``self``. For performance reasons only the mutable parts are copied. """ @@ -457,7 +449,7 @@ def coord_at(self, i): :param i: A tuple of the same dimensionality as subset.dims() or subset.data_dims(). :return: Absolute coordinates for index i (length equal to - `data_dims()`, may be larger than `dims()`). + ``data_dims()``, may be larger than ``dims()``). """ tiles = sum(1 if ts != 1 else 0 for ts in self.tile_sizes) if len(i) != len(self.ranges) and len(i) != len(self.ranges) + tiles: @@ -533,8 +525,8 @@ def dims(self): def absolute_strides(self, global_shape): """ Returns a list of strides for advancing one element in each - dimension. Size of the list is equal to `data_dims()`, which may - be larger than `dims()` depending on tile sizes. """ + dimension. Size of the list is equal to ``data_dims()``, which may + be larger than ``dims()`` depending on tile sizes. """ # ..., stride2*size1*size0, stride1*size0, stride0, ..., tile strides return [rs * global_shape[i] for i, (_, _, rs) in enumerate(self.ranges) ] + [global_shape[i] for i, ts in enumerate(self.tile_sizes) if ts != 1] @@ -573,14 +565,16 @@ def get_free_symbols_by_indices(self, indices: List[int]) -> Set[str]: result |= symbolic.symlist(d).keys() return result - def reorder(self, order): + def reorder(self, order: Sequence[int]) -> None: """ Re-orders the dimensions in-place according to a permutation list. :param order: List or tuple of integers from 0 to self.dims() - 1, indicating the desired order of the dimensions. """ new_ranges = [self.ranges[o] for o in order] + new_tile_sizes = [self.tile_sizes[o] for o in order] self.ranges = new_ranges + self.tile_sizes = new_tile_sizes @staticmethod def dim_to_string(d, t=1): @@ -954,6 +948,29 @@ def is_contiguous_subset(self, array: 'dace.data.Array') -> bool: if any(s != 1 for (_, _, s) in self): return False + # Special case for 1D subsets. + if self.dims() == 1: + return (array.strides[0] == 1) == True + + # Special case for 1D slicing, e.g. `a[idx0, b:e, idx2]`, i.e. there is only one dimensions + # in which the size of the subset is larger than 1 and in all others it is 1. The subset + # is continuous if that dimension has stride 1. + is_1d_slice = False + for copy_elem, stride in zip(self.size(), array.strides): + if (copy_elem == 1) == True: + continue + if (copy_elem == 0) == True: + return False # Strange case, does it happen? + if is_1d_slice: + is_1d_slice = False # There are multiple >1 size dimensions. + break + if (stride == 1) == False: + break # The >1 size dimension is not associated to stride 1. + is_1d_slice = True + else: + if is_1d_slice: + return True + # Determine array layout and calculate expression lengths accordingly if array.is_packed_fortran_strides(): # Fortran layout: first dimension varies fastest diff --git a/dace/symbolic.py b/dace/symbolic.py index b9d9d138fe..f13522a8f5 100644 --- a/dace/symbolic.py +++ b/dace/symbolic.py @@ -1,11 +1,14 @@ # Copyright 2019-2025 ETH Zurich and the DaCe authors. All rights reserved. import ast +import contextlib from collections import Counter from functools import lru_cache import sympy +import threading import pickle import re -from typing import Any, Callable, Dict, Iterable, Optional, Set, Tuple, Union, TYPE_CHECKING +import types +from typing import Any, Callable, Dict, FrozenSet, Iterable, Optional, Set, Tuple, Union, TYPE_CHECKING, List import numpy import sympy.abc @@ -16,8 +19,104 @@ from dace import dtypes DEFAULT_SYMBOL_TYPE = dtypes.int32 + + +class _SymbolDTypeContext(threading.local): + + def __init__(self): + + # The lowest level in the stack is reserved for "no stack active". + self.ctx_stack: List[types.MappingProxyType[str, 'dtypes.typeclass']] = [types.MappingProxyType({})] + + def push(self, authority: Dict[str, 'dtypes.typeclass']) -> types.MappingProxyType[str, 'dtypes.typeclass']: + """ + Adds a new level of authoritative dtype to the context. + + :param authority: Mapping from symbol name to its authoritative dtype. + """ + new_stack_level = types.MappingProxyType({ + n: dt + for n, dt in authority.items() if self._is_scalar_symbol_dtype(dt) + }) + self.ctx_stack.append(new_stack_level) + return self.ctx_stack[-1] + + def pop(self) -> "_SymbolDTypeContext": + """Remove the current active level of authoritative dtype.""" + if len(self.ctx_stack) == 1: + raise IndexError("Tried to `pop()` from an empty symbol type stack.") + self.ctx_stack.pop() + return self + + def get(self) -> types.MappingProxyType: + """Get the current active set of authoritative dtype.""" + if len(self.ctx_stack) == 0: + raise IndexError("Symbol type stack is empty.") + return self.ctx_stack[-1] + + @staticmethod + def _is_scalar_symbol_dtype(dtype: 'dtypes.typeclass') -> bool: + """ + Whether a dtype is a concrete scalar that can override a symbol's serialized + dtype: a plain :class:`~dace.dtypes.typeclass` with a real numpy scalar type. + Subclasses such as ``pointer``/``callback``/``vector`` are excluded even + though their ``.type`` may be a numpy scalar (a pointer's is its target's), + as is the typeless ``void`` of an untyped dynamic map-range connector. + """ + return type(dtype) is dtypes.typeclass and dtype.type is not None + + +# Authoritative dtype of the symbols an enclosing scope declares while an SDFG +# element is being serialized (name -> typeclass). ``DaceSympySerializer._print_Symbol`` +# consults this so a scoped symbol's emitted dtype is a deterministic function of the +# SDFG scope rather than of the symbol instance's own dtype, which SymPy's expression +# cache can leave stale (it conflates same-named symbols of different dtypes). The map +# only ever *overrides*: a name it does not declare keeps the symbol's own dtype, so a +# bare ``serialize_symbolic`` call (empty map) behaves exactly as before. +_SERIALIZATION_SYMBOL_DTYPES = _SymbolDTypeContext() + + +@contextlib.contextmanager +def serialization_symbol_dtypes(authority: Dict[str, 'dtypes.typeclass']): + """ + Temporarily override, while serializing symbolic expressions, the dtype used for + each scope-declared symbol, restoring the previous mapping on exit. Only concrete + scalar dtypes are kept; any other (e.g. a ``void`` dynamic-connector type) is left + out so that symbol keeps its own dtype. + + :param authority: Mapping from symbol name to its authoritative dtype. + """ + _SERIALIZATION_SYMBOL_DTYPES.push(authority) + try: + yield + finally: + _SERIALIZATION_SYMBOL_DTYPES.pop() + + _NAME_TOKENS = re.compile(r'[a-zA-Z_][a-zA-Z_0-9]*') -_FUNCTION_CALL = re.compile(r'(\w+)\[([^\[\]]+)\]') +_SERIALIZED_SYMBOL_PREFIX = '__DACE_SERIALIZED_SYMBOL_' +_SERIALIZED_UNDEFINED_SYMBOL = '__DACE_SERIALIZED_UNDEFINED_SYMBOL' +_SERIALIZED_SYMBOL = re.compile(r'\$(?P[a-zA-Z_][a-zA-Z_0-9]*)') +_SERIALIZED_TYPED_CONSTANT_VALUE = r'(?:\d+(?:\.\d*)?|\.\d+)(?:[eE][+-]?\d+)?|\d+' +_SERIALIZED_TYPED_CONSTANT_SUFFIX = r'|'.join( + sorted((re.escape(suffix) for suffix in dtypes.LITERAL_SUFFIX_TO_TYPECLASS), key=len, reverse=True)) +_SERIALIZED_TYPED_CONSTANT = re.compile( + rf'(?{_SERIALIZED_TYPED_CONSTANT_VALUE})(?P{_SERIALIZED_TYPED_CONSTANT_SUFFIX})\b' +) +_SERIALIZED_COMPLEX_FLOAT = r'-?(?:\d+(?:\.\d*)?|\.\d+)(?:[eE][+-]?\d+)?' +_SERIALIZED_TYPED_COMPLEX_CONSTANT = re.compile(rf'\(\s*(?:(?P{_SERIALIZED_COMPLEX_FLOAT})\s*(?P[+-])\s*)?' + rf'(?P{_SERIALIZED_COMPLEX_FLOAT})j\s*\)(?Pc64|c128)\b') +# Tokens in a Python expression string that require AST-based rewriting before +# being handed to SymPy (boolean ops, comparisons, bitwise ops, attribute/subscript, etc.). +# The ``.`` matches attribute access only (not a numeric decimal point): routing a float +# literal through ``ast.parse`` would round a near-max value like HUGE up to ``inf``. +_NEEDS_AST_REWRITE = re.compile( + r'\bnot\b|\band\b|\bor\b|\bNone\b|==|!=|\bis\b|\bif\b|[&]|[|]|[\^]|[~]|[<<]|[>>]|[//]|\.(?![0-9])|[\[]|[\]]') + + +def _is_sympy_number(expr) -> bool: + return bool(getattr(expr, 'is_Number', False) or getattr(expr, 'is_number', False)) + # NOTE: Up to (including) version 1.8, sympy.abc._clash is a dictionary of the # form {'N': sympy.abc.N, 'I': sympy.abc.I, 'pi': sympy.abc.pi} @@ -45,7 +144,9 @@ class symbol(sympy.Symbol): s_currentsymbol = 0 - def __new__(cls, name=None, dtype=DEFAULT_SYMBOL_TYPE, **assumptions): + def __new__(cls, name=None, dtype=None, **assumptions): + if dtype is None: + dtype = DEFAULT_SYMBOL_TYPE if name is None: # Set name dynamically name = "sym_" + str(symbol.s_currentsymbol) @@ -60,12 +161,15 @@ def __new__(cls, name=None, dtype=DEFAULT_SYMBOL_TYPE, **assumptions): dkeys = [k for k, v in dtypes.dtype_to_typeclass().items() if v == dtype] is_integer = [issubclass(k, int) or issubclass(k, numpy.integer) for k in dkeys] - if 'integer' in assumptions or not numpy.any(is_integer): - # Using __xnew__ as the regular __new__ is cached, which leads - # to modifying different references of symbols with the same name. - self = sympy.Symbol.__xnew__(cls, name, **assumptions) - else: - self = sympy.Symbol.__xnew__(cls, name, integer=True, **assumptions) + + # Don't pass `commutative` explicitly (SymPy defaults it to True anyway): keeping it out + # of `_assumptions_orig` avoids srepr/serialization order mismatches across build paths. + assumptions = {k: v for k, v in assumptions.items() if k != 'commutative'} + if 'integer' not in assumptions and numpy.any(is_integer): + assumptions['integer'] = True + # Using __xnew__ as the regular __new__ is cached, which leads + # to modifying different references of symbols with the same name. + self = sympy.Symbol.__xnew__(cls, name, **assumptions) self.dtype = dtype self._constraints = [] @@ -168,7 +272,9 @@ class UndefinedSymbol(symbol): >>> expr = N * undefined + 5 """ - def __new__(cls, dtype=DEFAULT_SYMBOL_TYPE, **assumptions): + def __new__(cls, dtype=None, **assumptions): + if dtype is None: + dtype = DEFAULT_SYMBOL_TYPE # Bypass the name validation self = sympy.Symbol.__xnew__(cls, "?", **assumptions) self.dtype = dtype @@ -237,6 +343,79 @@ def __hash__(self): return hash(self.name) +class TypedConstant(sympy.AtomicExpr): + """A typed constant value that participates in symbolic expressions. + + The value is a real ``Integer``/``Float``, or ``re + im*I`` when the + ``dtype`` is complex. The ``dtype`` is a type tag; complex constants must + be numeric (symbolic real/imag parts are unsupported). + + Examples + -------- + >>> from dace import symbolic, int16 + >>> symbolic.serialize_symbolic(symbolic.TypedConstant(2, int16)) + '2i16' + """ + + is_number = True + is_commutative = True + + def __new__(cls, value, dtype=None): + if isinstance(value, TypedConstant): + dtype = dtype or value.dtype + value = value.value + dtype = dtype or _infer_typed_constant_dtype(value) + if not isinstance(dtype, dtypes.typeclass): + raise TypeError(f'dtype must be a DaCe type, got {dtype!r}') + + self = sympy.AtomicExpr.__new__(cls) + self.value = _sympy_constant_value(value) + self.dtype = dtype + return self + + def _hashable_content(self): + return (self.value, self.dtype) + + def _sympystr(self, printer): + if hasattr(printer, '_print_TypedConstant'): + return printer._print_TypedConstant(self) + return printer.doprint(self.value) + + def _eval_evalf(self, prec): + return self.value.evalf(prec) + + def _eval_is_integer(self): + return self.value.is_integer + + def _eval_is_real(self): + return self.value.is_real + + def __int__(self): + return int(self.value) + + def __float__(self): + return float(self.value) + + @property + def is_Integer(self): + return self.value.is_Integer + + @property + def is_Float(self): + return self.value.is_Float + + @property + def _prec(self): + return self.value._prec + + @property + def is_integer(self): + return self.value.is_integer + + def cast_value(self): + return self.dtype(self.value) + + class SymExpr(object): """ Symbolic expressions with support for an overapproximation expression. """ @@ -375,6 +554,101 @@ def __gt__(self, other): return self > pystr_to_symbolic(other) +def _sympy_constant_value(value): + if isinstance(value, TypedConstant): + return value.value + if isinstance(value, sympy.Integer): + return value + if isinstance(value, sympy.Float): + return value + if isinstance(value, numpy.integer): + return sympy.Integer(int(value)) + if isinstance(value, numpy.floating): + return sympy.Float(float(value)) + if isinstance(value, int) and not isinstance(value, bool): + return sympy.Integer(value) + if isinstance(value, float): + return sympy.Float(value) + if isinstance(value, (complex, numpy.complexfloating)): + return sympy.Float(value.real) + sympy.Float(value.imag) * sympy.I + if isinstance(value, sympy.Expr) and value.is_number and value.is_real is False: + return value + raise TypeError(f'Unsupported typed constant value {value!r}') + + +def _infer_typed_constant_dtype(value) -> dtypes.typeclass: + if isinstance(value, TypedConstant): + return value.dtype + if isinstance(value, numpy.generic): + return dtypes.dtype_to_typeclass(type(value)) + if isinstance(value, sympy.Integer): + return DEFAULT_SYMBOL_TYPE + if isinstance(value, sympy.Float): + return dtypes.float64 + if isinstance(value, int) and not isinstance(value, bool): + return DEFAULT_SYMBOL_TYPE + if isinstance(value, float): + return dtypes.float64 + if isinstance(value, (complex, numpy.complexfloating)): + return dtypes.complex128 + raise TypeError(f'Cannot infer a DaCe dtype for {value!r}') + + +def _typed_constant_suffix(dtype: dtypes.typeclass) -> str: + try: + return dtypes.typeclass_to_literal_suffix(dtype) + except KeyError as ex: + raise TypeError(f'Unsupported typed constant dtype {dtype.to_string()}') from ex + + +def _format_float(value: float) -> str: + # Shortest round-trip form, keeping one fractional digit (5.0, not 5 or 5.000...). + s = f'{float(value):.15g}' + if 'e' in s or 'E' in s: + return s + if '.' not in s: + return s + '.0' + int_part, frac_part = s.split('.') + return f'{int_part}.{frac_part.rstrip("0") or "0"}' + + +def _typed_constant_to_string(expr: TypedConstant) -> str: + if expr.dtype in (dtypes.complex64, dtypes.complex128): + re = float(sympy.re(expr.value)) + im = float(sympy.im(expr.value)) + suffix = 'c64' if expr.dtype == dtypes.complex64 else 'c128' + if re == 0.0 and im != 0.0: + sign = '-' if im < 0 else '' + return f'({sign}{_format_float(abs(im))}j){suffix}' + op = '+' if im >= 0 else '-' + return f'({_format_float(re)} {op} {_format_float(abs(im))}j){suffix}' + if isinstance(expr.value, sympy.Float): + value = _format_float(float(expr.value)) + else: + value = sympy.printing.str.sstr(expr.value) + if expr.dtype in dtypes.TYPECLASS_TO_LITERAL_SUFFIX: + return f'{value}{_typed_constant_suffix(expr.dtype)}' + return f'dace.{expr.dtype.to_string()}({value})' + + +def _symbol_default_assumptions(expr: symbol) -> Dict[str, Any]: + return symbol(expr.name, dtype=expr.dtype).assumptions0 + + +def _symbol_serializer_kwargs(expr: symbol, dtype: 'dtypes.typeclass') -> Dict[str, Any]: + kwargs = {} + if dtype != DEFAULT_SYMBOL_TYPE: + kwargs['dtype'] = f'dace.{dtype.to_string()}' + + default_assumptions = _symbol_default_assumptions(symbol(expr.name, dtype=dtype)) + for key, value in sorted(expr.assumptions0.items()): + if key == 'commutative' or key.startswith('extended_'): + continue + if value is True and default_assumptions.get(key) != value: + kwargs[key] = value + return kwargs + + # Type hint for symbolic expressions SymbolicType = Union[sympy.Basic, SymExpr] @@ -421,6 +695,12 @@ def symlist(values): # Skip attributes skip.add(atom.args[1]) continue + if isinstance(atom, Subscript): + # Only the indices are free symbols (consistent with + # ``Subscript.free_symbols``). Skip the whole head subtree so a compound + # head like ``a.b`` in ``a.b[i]`` does not leak ``a``/``b``. + skip.update(sympy.preorder_traversal(atom.args[0])) + continue if isinstance(atom, symbol): result[atom.name] = atom return result @@ -450,6 +730,8 @@ def evaluate(expr: Union[sympy.Basic, int, float], symbols: Dict[Union[symbol, s raise TypeError(f'Symbolic expression "{expr}" cannot be evaluated to a constant') if isinstance(expr, (int, float, numpy.number)): return expr + if isinstance(expr, TypedConstant): + return expr.cast_value() # Evaluate all symbols syms = {(sname if isinstance(sname, sympy.Symbol) else symbol(sname)): sval for sname, sval in symbols.items()} @@ -494,7 +776,7 @@ def overapproximate(expr): return _overapproximate(expr) -@lru_cache(maxsize=2048) +@lru_cache(maxsize=2048, typed=True) def _overapproximate(expr): if isinstance(expr, SymExpr): if expr.expr != expr.approx: @@ -503,7 +785,7 @@ def _overapproximate(expr): return overapproximate(expr.expr) if not isinstance(expr, sympy.Basic): return expr - if isinstance(expr, sympy.Number): + if isinstance(expr, (sympy.Number, TypedConstant)): return expr a = sympy.Wild('a') @@ -636,6 +918,9 @@ def swalk(expr, enter_functions=False): def contains_sympy_functions(expr): """ Returns True if expression contains Sympy functions. """ + if isinstance(expr, Subscript): + # A subscript is a data-container access. + return True if is_sympy_userfunction(expr): if str(expr.func) in _builtin_userfunctions: return False @@ -649,6 +934,16 @@ def contains_sympy_functions(expr): def free_symbols_and_functions(expr: Union[SymbolicType, str]) -> Set[str]: + """ + Return the names of the free symbols and (non-builtin) functions in an expression. + + Data containers accessed via ``Subscript`` (e.g. ``A`` in ``A[i]``) are NOT + reported here, as the container is a data access rather than a free symbol; use + :func:`arrays` to obtain those. + + :param expr: The expression (or its string form) to inspect. + :return: The set of free-symbol and function names. + """ if isinstance(expr, str): if dtypes.validate_name(expr): return {expr} @@ -663,6 +958,43 @@ def free_symbols_and_functions(expr: Union[SymbolicType, str]) -> Set[str]: return result +def arrays(expr: Union[SymbolicType, str]) -> Set[str]: + """ + Return the names of the containers accessed via ``Subscript`` (e.g. ``A`` in ``A[i]``). + + Only subscripted accesses are reported; a rank-0 scalar referenced by name is + indistinguishable from a free symbol here, so it is reported by :func:`scalars`. + + :param expr: The expression (or its string form) to inspect. + :return: The set of subscripted container names. + """ + if isinstance(expr, str): + expr = pystr_to_symbolic(expr) + if not isinstance(expr, sympy.Basic): + return set() + return {str(node.args[0]) for node in expr.atoms(Subscript)} + + +def scalars(expr: Union[SymbolicType, str], descriptors: Dict[str, Any]) -> Set[str]: + """ + Return the names of the rank-0 scalar containers referenced by name. + + A bare name is indistinguishable from a free symbol without the data descriptors, so + ``descriptors`` (e.g. ``sdfg.arrays``) is required. ``arrays(e) | scalars(e, descriptors)`` + yields all referenced data containers. + + :param expr: The expression (or its string form) to inspect. + :param descriptors: Name-to-descriptor mapping used to identify scalar containers. + :return: The set of referenced scalar names. + """ + if isinstance(expr, str): + expr = pystr_to_symbolic(expr) + if not isinstance(expr, sympy.Basic): + return set() + from dace import data # avoid import loop + return {str(s) for s in expr.free_symbols if isinstance(descriptors.get(str(s)), data.Scalar)} + + def is_undefined(expr: Union[SymbolicType, str]) -> bool: """ Checks if a symbolic expression contains any UndefinedSymbol atoms. @@ -736,6 +1068,13 @@ def _eval_is_integer(self): return True +class __int_floor(int_floor): + """ Operator-derived variant of ``int_floor``: the Python ``//`` parses to this so + it round-trips to ``//`` (Python) / ``/`` (C++), while an explicit + ``int_floor(a, b)`` keeps its function spelling. """ + pass + + class int_ceil(sympy.Function): @classmethod @@ -855,11 +1194,12 @@ def eval(cls, x, y): :param x: Value to shift. :param y: Value to shift by. - :return: Return value (literal or symbolic). + :return: The shifted literal if both operands are concrete, else ``None``. """ + # Keep symbolic shifts unevaluated so they round-trip to ``x << y`` rather + # than collapsing to the equivalent ``x * 2**y``. if x.is_Number and y.is_Number: return x << y - return x * (2**y) class right_shift(sympy.Function): @@ -871,11 +1211,40 @@ def eval(cls, x, y): :param x: Value to shift. :param y: Value to shift by. - :return: Return value (literal or symbolic). + :return: The shifted literal if both operands are concrete, else ``None``. """ + # Keep symbolic shifts unevaluated so they round-trip to ``x >> y`` rather + # than collapsing to ``int_floor(x, 2**y)``. if x.is_Number and y.is_Number: return x >> y - return int_floor(x, (2**y)) + + +# Internal variants for the Python operators: ``a | b`` parses to ``__bitwise_or``, etc. +# ``symstr`` prints these as the operator, while the bare ``bitwise_or`` etc. print as +# ``func(a, b)`` (both lower to the operator in C++). They subclass the bare classes so +# ``isinstance`` checks still match. +class __bitwise_and(bitwise_and): + pass + + +class __bitwise_or(bitwise_or): + pass + + +class __bitwise_xor(bitwise_xor): + pass + + +class __bitwise_invert(bitwise_invert): + pass + + +class __left_shift(left_shift): + pass + + +class __right_shift(right_shift): + pass class ROUND(sympy.Function): @@ -910,12 +1279,10 @@ class Attr(sympy.Function): @property def free_symbols(self): - # NOTE: The following handles the case where the attribute is an array access, e.g., "indptr[i]" - if isinstance(self.args[1], sympy.Function): - attribute = str(self.args[1].func) - else: - attribute = str(self.args[1]) - return {sympy.Symbol(f"{self.args[0]}.{attribute}")} + # Report the member as a single ``a.b`` symbol, not the free symbols of its + # parts. An array-access member (``a.b[i]``) is a ``Subscript`` wrapping the + # ``Attr``, so ``args[1]`` is always the plain attribute name. + return {sympy.Symbol(f"{self.args[0]}.{self.args[1]}")} def __str__(self): return f'{self.args[0]}.{self.args[1]}' @@ -924,6 +1291,28 @@ def _subs(self, *args, **kwargs): return Attr(self.args[0].subs(*args, **kwargs), self.args[1].subs(*args, **kwargs)) +class Subscript(sympy.Function): + """ + Represents a subscript expression, equivalent to ``a[i, j, ...]`` in Python. + + The first argument is the subscripted expression; the remaining arguments are + the (single-point) indices. + """ + + @property + def free_symbols(self): + # The subscripted container is a data access, not a free symbol; only the + # indices contribute. The container is reported by :func:`arrays` instead. + return set().union(*(a.free_symbols for a in self.args[1:])) if len(self.args) > 1 else set() + + def __str__(self): + indices = ', '.join(str(a) for a in self.args[1:]) + return f'{self.args[0]}[{indices}]' + + def _subs(self, *args, **kwargs): + return Subscript(*(a.subs(*args, **kwargs) for a in self.args)) + + def sympy_intdiv_fix(expr): """ Fix for SymPy printing out reciprocal values when they should be integral in "ceiling/floor" sympy functions. @@ -1160,13 +1549,13 @@ class PythonOpToSympyConverter(ast.NodeTransformer): } _ast_to_sympy_functions = { - ast.BitAnd: 'bitwise_and', - ast.BitOr: 'bitwise_or', - ast.BitXor: 'bitwise_xor', - ast.Invert: 'bitwise_invert', - ast.LShift: 'left_shift', - ast.RShift: 'right_shift', - ast.FloorDiv: 'int_floor', + ast.BitAnd: '__bitwise_and', + ast.BitOr: '__bitwise_or', + ast.BitXor: '__bitwise_xor', + ast.Invert: '__bitwise_invert', + ast.LShift: '__left_shift', + ast.RShift: '__right_shift', + ast.FloorDiv: '__int_floor', } def visit_UnaryOp(self, node): @@ -1254,13 +1643,28 @@ def visit_IfExp(self, node): return ast.copy_location(new_node, node) def visit_Subscript(self, node): - if isinstance(node.value, ast.Attribute): - attr = ast.Subscript(value=ast.Name(id=node.value.attr, ctx=ast.Load()), slice=node.slice, ctx=ast.Load()) - new_node = ast.Call(func=ast.Name(id='Attr', ctx=ast.Load), - args=[self.visit(node.value.value), self.visit(attr)], - keywords=[]) - return ast.copy_location(new_node, node) - return self.generic_visit(node) + # Collect the index expressions. A tuple slice (``a[i, j]``) becomes + # multiple arguments to ``Subscript``; a single index becomes one argument. + slice_node = node.slice + if isinstance(slice_node, ast.Tuple): + indices = list(slice_node.elts) + else: + indices = [slice_node] + + # Range/slice expressions cannot be represented as symbolic expressions. + for idx in indices: + if isinstance(idx, ast.Slice): + raise SyntaxError(f'Range/slice expressions are not supported in symbolic expressions ' + f'(got slice in subscript of "{ast.unparse(node)}")') + + # Recursively visit the subscripted value (handles attributes and nested + # subscripts via visit_Attribute / visit_Subscript). + value = self.visit(node.value) + + new_node = ast.Call(func=ast.Name(id='Subscript', ctx=ast.Load), + args=[value] + [self.visit(idx) for idx in indices], + keywords=[]) + return ast.copy_location(new_node, node) def visit_Attribute(self, node): new_node = ast.Call(func=ast.Name(id='Attr', ctx=ast.Load), @@ -1269,48 +1673,166 @@ def visit_Attribute(self, node): return ast.copy_location(new_node, node) -@lru_cache(maxsize=16384) -def pystr_to_symbolic(expr, symbol_map=None, simplify=None) -> sympy.Basic: - """ Takes a Python string and converts it into a symbolic expression. """ - from dace.frontend.python.astutils import unparse # Avoid import loops +class _SerializedSymbolicParser(ast.NodeVisitor): + """ + Parser for the deterministic expression strings produced by + :func:`serialize_symbolic`. + """ - if isinstance(expr, (SymExpr, sympy.Basic)): - return expr - if isinstance(expr, str): - try: - return sympy.Integer(int(expr)) - except ValueError: - pass - try: - return sympy.Float(float(expr)) - except ValueError: - pass - if "?" in expr: # Note that this will convert expressions like "a ? b : c" or "some_func(?)" to UndefinedSymbol - return UndefinedSymbol() - if dtypes.validate_name(expr): - return symbol(expr) + @staticmethod + def _python_bool(value): + if value is sympy.true: + return True + if value is sympy.false: + return False + return value - symbol_map = symbol_map or {} - locals = { + @staticmethod + def _flatten_args(expr_type, *args): + result = [] + for arg in args: + if isinstance(arg, expr_type): + result.extend(arg.args) + else: + result.append(arg) + return result + + @staticmethod + def _add(*args): + return sympy.Add._from_args(tuple(args)) + + @staticmethod + def _mul(*args): + return sympy.Mul._from_args(tuple(args)) + + @staticmethod + def _pow(a, b): + a = sympy.sympify(a) + b = sympy.sympify(b) + # Evaluate pure-numeric powers (8**-1 -> 1/8) so coefficients stay reduced Rationals; an + # unevaluated Pow reorders surrounding Mul terms on round-trip. TypedConstants keep dtype. + if (_is_sympy_number(a) and _is_sympy_number(b) and not isinstance(a, TypedConstant) + and not isinstance(b, TypedConstant)): + return a**b + if hasattr(sympy.Pow, '_from_args'): + return sympy.Pow._from_args((a, b)) + if hasattr(sympy.Pow, '__xnew__'): + return sympy.Pow.__xnew__(sympy.Pow, a, b) + return sympy.Pow.__new__.__wrapped__(sympy.Pow, a, b, evaluate=False) + + @staticmethod + def _binop_add(a, b): + flat = _SerializedSymbolicParser._flatten_args(sympy.Add, a, b) + flat.sort(key=sympy.default_sort_key) + return _SerializedSymbolicParser._add(*flat) + + @staticmethod + def _binop_mul(a, b): + args = _SerializedSymbolicParser._flatten_args(sympy.Mul, a, b) + args.sort(key=sympy.default_sort_key) + + if len(args) > 1: + args = [arg for arg in args if not _is_sympy_number(arg) or not equal_valued(arg, 1)] + + coeff = sympy.S.One + nonnumeric_args = [] + for arg in args: + if not isinstance(arg, TypedConstant) and _is_sympy_number(arg): + coeff *= arg + else: + nonnumeric_args.append(arg) + + if coeff != 1 or not nonnumeric_args: + args = [coeff] + nonnumeric_args + else: + args = nonnumeric_args + + if not args: + return sympy.Integer(1) + return _SerializedSymbolicParser._mul(*args) + + @staticmethod + def _negate(a): + if isinstance(a, TypedConstant): + return TypedConstant(-a.value, dtype=a.dtype) + if isinstance(a, sympy.Number): + return -a + return _SerializedSymbolicParser._binop_mul(sympy.S.NegativeOne, a) + + @staticmethod + def _binop_sub(a, b): + return _SerializedSymbolicParser._add( + *_SerializedSymbolicParser._flatten_args(sympy.Add, a, _SerializedSymbolicParser._negate(b))) + + @staticmethod + def _binop_div(a, b): + # Fold the numeric quotient into one reduced factor (1/3 -> Rational(1, 3)) + # instead of leaving a stray identity 1*(1/3). + return _SerializedSymbolicParser._binop_mul(a, _SerializedSymbolicParser._pow(b, -1)) + + @staticmethod + def _binop_pow(a, b): + return _SerializedSymbolicParser._pow(a, b) + + @staticmethod + def _binop_mod(a, b): + return sympy.Mod(a, b, evaluate=False) + + @staticmethod + def _unary_minus(a): + return _SerializedSymbolicParser._negate(a) + + _binops = { + ast.Add: _binop_add, + ast.Sub: _binop_sub, + ast.Mult: _binop_mul, + ast.Div: _binop_div, + ast.Pow: _binop_pow, + ast.Mod: _binop_mod, + ast.FloorDiv: lambda a, b: int_floor(a, b), + } + _unaryops = { + ast.UAdd: lambda a: +a, + ast.USub: _unary_minus, + ast.Not: lambda a: sympy.Not(a), + ast.Invert: lambda a: bitwise_invert(a), + } + _comparators = { + ast.Eq: sympy.Eq, + ast.NotEq: sympy.Ne, + ast.Gt: sympy.Gt, + ast.GtE: sympy.Ge, + ast.Lt: sympy.Lt, + ast.LtE: sympy.Le, + ast.Is: Is, + ast.IsNot: IsNot, + } + _functions = { 'abs': sympy.Abs, + 'Abs': sympy.Abs, 'min': sympy.Min, + 'Min': sympy.Min, 'max': sympy.Max, - 'True': sympy.true, - 'False': sympy.false, - 'GtE': sympy.Ge, - 'LtE': sympy.Le, - 'NotEq': sympy.Ne, + 'Max': sympy.Max, 'floor': sympy.floor, 'ceil': sympy.ceiling, + 'ceiling': sympy.ceiling, + 'sqrt': sympy.sqrt, 'round': ROUND, - # Convert and/or to special sympy functions to avoid boolean evaluation 'And': AND, 'Or': OR, - 'var': sympy.Symbol('var'), - 'root': sympy.Symbol('root'), - 'arg': sympy.Symbol('arg'), - 'Is': Is, - 'IsNot': IsNot, + 'Not': sympy.Not, + 'Eq': sympy.Eq, + 'Ne': sympy.Ne, + 'Gt': sympy.Gt, + 'Ge': sympy.Ge, + 'Lt': sympy.Lt, + 'Le': sympy.Le, + 'int_floor': int_floor, + 'int_ceil': int_ceil, + 'IfExpr': IfExpr, + 'Mod': sympy.Mod, + 'Attr': Attr, 'BitwiseAnd': bitwise_and, 'BitwiseOr': bitwise_or, 'BitwiseXor': bitwise_xor, @@ -1320,47 +1842,441 @@ def pystr_to_symbolic(expr, symbol_map=None, simplify=None) -> sympy.Basic: 'bitwise_xor': bitwise_xor, 'bitwise_invert': bitwise_invert, 'LeftShift': left_shift, - 'left_shift': left_shift, 'RightShift': right_shift, + 'left_shift': left_shift, 'right_shift': right_shift, - 'int_floor': int_floor, - 'int_ceil': int_ceil, - 'IfExpr': IfExpr, - 'Mod': sympy.Mod, - 'Attr': Attr, - 'id': sympy.Symbol('id'), - 'diag': sympy.Symbol('diag'), - 'jn': sympy.Symbol('jn'), } - # _clash1 enables all one-letter variables like N as symbols - # _clash also allows pi, beta, zeta and other common greek letters - locals.update(_sympy_clash) + _constants = { + 'True': sympy.true, + 'False': sympy.false, + 'None': symbol('NoneSymbol'), + 'NoneSymbol': symbol('NoneSymbol'), + 'pi': sympy.pi, + 'E': sympy.E, + 'I': sympy.I, + 'oo': sympy.oo, + 'nan': sympy.nan, + } + + def visit_Expression(self, node): + return self.visit(node.body) + + def visit_Name(self, node): + if node.id in self._constants: + return self._constants[node.id] + if node.id == _SERIALIZED_UNDEFINED_SYMBOL: + return UndefinedSymbol() + if node.id.startswith(_SERIALIZED_SYMBOL_PREFIX): + return symbol(node.id[len(_SERIALIZED_SYMBOL_PREFIX):]) + try: + return getattr(dtypes, node.id) + except AttributeError as ex: + raise TypeError(f'Unknown symbolic identifier "{node.id}"') from ex + + def visit_Constant(self, node): + if isinstance(node.value, bool): + return sympy.true if node.value else sympy.false + if node.value is None: + return symbol('NoneSymbol') + if isinstance(node.value, int): + return sympy.Integer(node.value) + if isinstance(node.value, float): + return sympy.Float(node.value) + if isinstance(node.value, complex): + return TypedConstant(node.value, dtypes.complex128) + raise TypeError(f'Unsupported literal {node.value!r}') + + def visit_UnaryOp(self, node): + return self._unaryops[type(node.op)](self.visit(node.operand)) + + def visit_BinOp(self, node): + return self._binops[type(node.op)](self.visit(node.left), self.visit(node.right)) + + def visit_BoolOp(self, node): + values = [self.visit(v) for v in node.values] + if isinstance(node.op, ast.And): + result = values[0] + for value in values[1:]: + result = AND(result, value) + return result + result = values[0] + for value in values[1:]: + result = OR(result, value) + return result + + def visit_Compare(self, node): + if len(node.ops) != 1 or len(node.comparators) != 1: + raise NotImplementedError('Chained comparisons are not supported in symbolic deserialization') + return self._comparators[type(node.ops[0])](self.visit(node.left), self.visit(node.comparators[0])) + + def visit_IfExp(self, node): + return IfExpr(self.visit(node.test), self.visit(node.body), self.visit(node.orelse)) + + def visit_Call(self, node): + if isinstance(node.func, ast.Name) and node.func.id == '__dace_typed_const__': + value = ast.literal_eval(node.args[0]) + suffix = ast.literal_eval(node.args[1]) + try: + dtype = dtypes.literal_suffix_to_typeclass(suffix) + except KeyError as ex: + raise TypeError(f'Invalid type suffix "{suffix}" in typed constant') from ex + # Keep ``7f64`` integer-valued and ``7.0f64`` float-valued so each + # form round-trips back to itself. + looks_like_float = any(c in value for c in '.eE') + if dtype in (dtypes.float16, dtypes.float32, dtypes.float64) and looks_like_float: + value = float(value) + else: + value = int(value) + return TypedConstant(value, dtype) + + if isinstance(node.func, ast.Name) and node.func.id == 'symbol': + if len(node.args) != 1 or not isinstance(node.args[0], ast.Name): + raise TypeError('symbol(...) expects its first argument to be a serialized symbol name') + symname = node.args[0].id + if not symname.startswith(_SERIALIZED_SYMBOL_PREFIX): + raise TypeError('symbol(...) expects its first argument to be a serialized symbol name') + symname = symname[len(_SERIALIZED_SYMBOL_PREFIX):] + kwargs = {kw.arg: self._python_bool(self.visit(kw.value)) for kw in node.keywords} + dtype = kwargs.pop('dtype', DEFAULT_SYMBOL_TYPE) + return symbol(symname, dtype=dtype, **kwargs) + + if isinstance(node.func, ast.Name) and node.func.id == 'SymExpr': + args = [self.visit(arg) for arg in node.args] + return SymExpr(*args) + + if (isinstance(node.func, ast.Name) and node.func.id in dtypes.CTYPE_TO_TYPECLASS and len(node.args) == 1): + return _cast_symbolic_value(self.visit(node.args[0]), dtypes.CTYPE_TO_TYPECLASS[node.func.id]) + + if isinstance(node.func, ast.Name) and node.func.id == 'complex' and len(node.args) == 2: + re = self.visit(node.args[0]) + im = self.visit(node.args[1]) + return TypedConstant(complex(float(re), float(im)), dtypes.complex128) + + func = self._resolve_function(node.func) + args = [self.visit(arg) for arg in node.args] + if isinstance(func, dtypes.typeclass): + if len(args) != 1: + raise TypeError(f'Type conversion {func} expects one argument') + return _cast_symbolic_value(args[0], func) + + kwargs = {kw.arg: self.visit(kw.value) for kw in node.keywords} + if kwargs: + return func(*args, **kwargs) + return func(*args) + + def visit_Attribute(self, node): + if isinstance(node.value, ast.Name) and node.value.id == 'dace': + try: + return getattr(dtypes, node.attr) + except AttributeError as ex: + raise TypeError(f'Unknown DaCe dtype "{node.attr}"') from ex + return Attr(self.visit(node.value), symbol(node.attr)) + + def generic_visit(self, node): + raise TypeError(f'Unsupported node in symbolic deserialization: {type(node).__name__}') + + def _resolve_function(self, node): + if isinstance(node, ast.Name): + if node.id in self._functions: + return self._functions[node.id] + return sympy.Function(node.id) + if isinstance(node, ast.Attribute): + resolved = self.visit(node) + if isinstance(resolved, dtypes.typeclass): + return resolved + raise TypeError(f'Unsupported symbolic function {ast.dump(node)}') + + +def _cast_symbolic_value(value, dtype: dtypes.typeclass): + if isinstance(value, SymExpr): + return SymExpr(_cast_symbolic_value(value.expr, dtype), _cast_symbolic_value(value.approx, dtype)) + if isinstance(value, symbol): + return symbol(value.name, dtype=dtype, **value.assumptions0) + if isinstance(value, TypedConstant): + return TypedConstant(value.value, dtype=dtype) + if isinstance(value, (sympy.Integer, sympy.Float, int, float, numpy.generic)): + return TypedConstant(value, dtype=dtype) + # Non-constant composite expressions are preserved as explicit casts so they + # can round-trip even when no constant/symbol dtype rewrite is possible. + return sympy.Function(f'dace.{dtype.to_string()}')(value) + + +class DaceSympySerializer(sympy.printing.str.StrPrinter): + + def _print_Symbol(self, expr): + if expr.name == '?': + return '$?' + if expr.name == 'NoneSymbol': + return 'None' + if isinstance(expr, symbol): + # Prefer the dtype the enclosing scope declares for this name over the + # instance's own (possibly cache-stale) dtype; a name the scope does not + # declare keeps the instance dtype (the authority only overrides). + dtype = _SERIALIZATION_SYMBOL_DTYPES.get().get(expr.name, expr.dtype) + kwargs = _symbol_serializer_kwargs(expr, dtype) + if not kwargs: + return f'${expr.name}' + kwlist = ', '.join(f'{key}={value}' for key, value in kwargs.items()) + return f'symbol(${expr.name}, {kwlist})' + return f'${expr.name}' + + def _print_TypedConstant(self, expr): + return _typed_constant_to_string(expr) + + def _print_Integer(self, expr): + return super()._print_Integer(expr) + + def _print_Float(self, expr): + return _format_float(float(sympy_numeric_fix(expr))) + + def _print_Add(self, expr): + # Sort arguments deterministically using SymPy's default_sort_key + sorted_args = sorted(expr.args, key=sympy.core.sorting.default_sort_key) + + # Flatten the sorted arguments + flat_args = [] + + def _flatten(args): + for arg in args: + if isinstance(arg, sympy.Add): + _flatten(sorted(arg.args, key=sympy.core.sorting.default_sort_key)) + else: + flat_args.append(arg) + + _flatten(sorted_args) + + parts = [] + for i, arg in enumerate(flat_args): + # Extract sign for printing + negative = False + if isinstance(arg, (sympy.Number, TypedConstant)) and arg < 0: + negative = True + arg = -arg + elif isinstance(arg, sympy.Basic) and arg.could_extract_minus_sign(): + negative = True + arg = -arg + + rendered = self._print(arg) + if i == 0: + parts.append(f'-{rendered}' if negative else rendered) + else: + parts.append(f' {"-" if negative else "+"} {rendered}') + return ''.join(parts) + + def _print_Mul(self, expr): + sorted_args = sorted(expr.args, key=sympy.default_sort_key) + + flat_args = [] + + def _flatten(args): + for arg in args: + if isinstance(arg, sympy.Mul): + _flatten(sorted(arg.args, key=sympy.default_sort_key)) + else: + flat_args.append(arg) + + _flatten(sorted_args) + + if len(flat_args) > 1: + flat_args = [arg for arg in flat_args if not _is_sympy_number(arg) or not equal_valued(arg, 1)] + + coeff = sympy.S.One + nonnumeric_args = [] + for arg in flat_args: + # Fold only genuine numbers (Integer/Rational/Float). `_is_sympy_number` also matches + # `is_number` composites (pi, sqrt(2), I, (-1)**(1/3)); folding one into `coeff` rebuilds + # this same Mul, and `self._print(coeff)` below would then recurse forever. + if isinstance(arg, sympy.Number) and not isinstance(arg, TypedConstant): + coeff *= arg + else: + nonnumeric_args.append(arg) + + if coeff != 1 or not nonnumeric_args: + flat_args = [coeff] + nonnumeric_args + else: + flat_args = nonnumeric_args + + parts = [] + for arg in flat_args: + rendered = self._print(arg) + if isinstance(arg, sympy.Add): + rendered = f'({rendered})' + parts.append(rendered) + return '*'.join(parts) if parts else '1' + + +def _serialize_symbolic_uncached(expr: Union[SymbolicType, int, float, numpy.number]) -> str: + if isinstance(expr, SymExpr): + return f'SymExpr({serialize_symbolic(expr.expr)}, {serialize_symbolic(expr.approx)})' + if isinstance(expr, TypedConstant): + return _typed_constant_to_string(expr) + if isinstance(expr, numpy.generic): + return _typed_constant_to_string(TypedConstant(expr, dtypes.dtype_to_typeclass(type(expr)))) + if isinstance(expr, int) and not isinstance(expr, bool): + return str(expr) + if isinstance(expr, float): + return sympy.printing.str.sstr(expr) + if isinstance(expr, sympy.Basic): + return DaceSympySerializer().doprint(expr) + return str(expr) + + +def serialize_symbolic(expr): + return _serialize_symbolic_uncached(expr) + + +@lru_cache(maxsize=16384, typed=True) +def deserialize_symbolic(expr) -> SymbolicType: + """Deserialize a string produced by :func:`serialize_symbolic`.""" + if isinstance(expr, (SymExpr, sympy.Basic)): + return expr + if isinstance(expr, numpy.generic): + return TypedConstant(expr, dtypes.dtype_to_typeclass(type(expr))) + if isinstance(expr, int) and not isinstance(expr, bool): + return sympy.Integer(expr) + if isinstance(expr, float): + return sympy.Float(expr) + if not isinstance(expr, str): + return expr + if expr in ('?', '$?'): + return UndefinedSymbol() + + # Symbols are rewritten first because typed-literal replacement introduces + # helper calls, and keeping the substitutions one-way avoids reparsing the + # generated helper identifiers. + expr = expr.replace('$?', _SERIALIZED_UNDEFINED_SYMBOL) + expr = _SERIALIZED_SYMBOL.sub(lambda m: f'{_SERIALIZED_SYMBOL_PREFIX}{m.group("name")}', expr) + expr = _SERIALIZED_TYPED_COMPLEX_CONSTANT.sub(_rewrite_typed_complex, expr) + expr = _SERIALIZED_TYPED_CONSTANT.sub( + lambda m: f'__dace_typed_const__({m.group("value")!r}, {m.group("suffix")!r})', expr) + return _SerializedSymbolicParser().visit(ast.parse(expr, mode='eval')) + + +def _rewrite_typed_complex(m: re.Match) -> str: + dtype = 'complex64' if m.group('csuffix') == 'c64' else 'complex128' + re_val = m.group('re') or '0' + sign = '-' if m.group('op') == '-' else '' + return f'dace.{dtype}(complex({re_val}, {sign}{m.group("im")}))' + + +_PYSTR2SYM_locals = { + 'abs': sympy.Abs, + 'min': sympy.Min, + 'max': sympy.Max, + 'True': sympy.true, + 'False': sympy.false, + 'GtE': sympy.Ge, + 'LtE': sympy.Le, + 'NotEq': sympy.Ne, + 'floor': sympy.floor, + 'ceil': sympy.ceiling, + 'round': ROUND, + # Resolve infinity/NaN names so they survive inside expressions (e.g. ``a + inf``). + 'inf': sympy.oo, + 'nan': sympy.nan, + # Convert and/or to special sympy functions to avoid boolean evaluation + 'And': AND, + 'Or': OR, + 'var': sympy.Symbol('var'), + 'root': sympy.Symbol('root'), + 'arg': sympy.Symbol('arg'), + 'Is': Is, + 'IsNot': IsNot, + 'BitwiseAnd': bitwise_and, + 'BitwiseOr': bitwise_or, + 'BitwiseXor': bitwise_xor, + 'BitwiseNot': bitwise_invert, + 'bitwise_and': bitwise_and, + 'bitwise_or': bitwise_or, + 'bitwise_xor': bitwise_xor, + 'bitwise_invert': bitwise_invert, + '__bitwise_and': __bitwise_and, + '__bitwise_or': __bitwise_or, + '__bitwise_xor': __bitwise_xor, + '__bitwise_invert': __bitwise_invert, + 'LeftShift': left_shift, + 'left_shift': left_shift, + 'RightShift': right_shift, + 'right_shift': right_shift, + '__left_shift': __left_shift, + '__right_shift': __right_shift, + 'int_floor': int_floor, + '__int_floor': __int_floor, + 'int_ceil': int_ceil, + 'IfExpr': IfExpr, + 'Mod': sympy.Mod, + 'Attr': Attr, + 'Subscript': Subscript, + 'id': sympy.Symbol('id'), + 'diag': sympy.Symbol('diag'), + 'jn': sympy.Symbol('jn'), +} +# _clash1 enables all one-letter variables like N as symbols +# _clash also allows pi, beta, zeta and other common greek letters +_PYSTR2SYM_locals.update(_sympy_clash) + + +def pystr_to_symbolic(expr, symbol_map=None, simplify=None) -> sympy.Basic: + """ + The visitor reconstructs symbolic expressions with non-evaluating SymPy + constructors so that serialization round-trips preserve tree structure and + typed-symbol metadata instead of depending on SymPy's automatic + simplification and constructor caches. + """ + # Do not cache SymPy objects: SymPy equality/hash ignores DaCe symbol dtype metadata. + # Keep SymExpr intact even when simplify=True, as it carries exact and approximate forms. + if isinstance(expr, SymExpr): + return expr + if isinstance(expr, sympy.Basic): + return sympy.simplify(expr) if simplify is True else expr + # Symbol maps may contain unhashable or mutable caller-specific replacements, so only cache plain parsing. + if symbol_map is None: + return _pystr_to_symbolic_cached(expr, simplify) + return _pystr_to_symbolic_uncached(expr, symbol_map, simplify) + + +@lru_cache(maxsize=16384, typed=True) +def _pystr_to_symbolic_cached(expr, simplify=None) -> sympy.Basic: + return _pystr_to_symbolic_uncached(expr, None, simplify) + + +def _pystr_to_symbolic_uncached(expr, symbol_map=None, simplify=None) -> sympy.Basic: + from dace.frontend.python.astutils import unparse # Avoid import loops + + if isinstance(expr, str): + try: + return sympy.Integer(int(expr)) + except ValueError: + pass + try: + # Parse the literal directly (SymPy is arbitrary-precision); going via + # Python ``float()`` would round a near-max value (e.g. HUGE) up to inf. + return sympy.Float(expr) + except (ValueError, TypeError): + pass + if "?" in expr: # Note that this will convert expressions like "a ? b : c" or "some_func(?)" to UndefinedSymbol + return UndefinedSymbol() + if dtypes.validate_name(expr): + return symbol(expr) + + symbol_map = symbol_map or {} if isinstance(expr, str): # Sympy processes "not/and/or" as direct evaluation. Replace with And/Or(x, y), Not(x) # Also replaces bitwise operations with user-functions since SymPy does not support bitwise operations. - if re.search(r'\bnot\b|\band\b|\bor\b|\bNone\b|==|!=|\bis\b|\bif\b|[&]|[|]|[\^]|[~]|[<<]|[>>]|[//]|[\.]', expr): + if _NEEDS_AST_REWRITE.search(expr): expr = unparse(PythonOpToSympyConverter().visit(ast.parse(expr).body[0])) # TODO: support SymExpr over-approximated expressions - try: - return sympy_to_dace(sympy.sympify(expr, locals, evaluate=simplify), symbol_map) - except (TypeError, sympy.SympifyError): # Symbol object is not subscriptable - # Replace subscript expressions with function calls - orig_expr = expr - expr = expr.replace('[', '(') - expr = expr.replace(']', ')') - try: - return sympy_to_dace(sympy.sympify(expr, locals, evaluate=simplify), symbol_map) - except TypeError: # Symbol object is not subscriptable - # Replace instances of "xxx[yyy]" with subscript(xxx, yyy) - expr = orig_expr - while _FUNCTION_CALL.search(expr): - expr = _FUNCTION_CALL.sub(r'subscript(\1, \2)', expr) - return sympy_to_dace(sympy.sympify(expr, locals, evaluate=simplify), symbol_map) + result = sympy.sympify(expr, _PYSTR2SYM_locals, evaluate=simplify) + if isinstance(result, bool): + # SymPy parses the literals ``True``/``False`` to Python bools; keep them as + # SymPy booleans so they stay distinct from the integers ``1``/``0``. + result = sympy.true if result else sympy.false + return sympy_to_dace(result, symbol_map) -@lru_cache(maxsize=2048) +@lru_cache(maxsize=2048, typed=True) def simplify(expr: SymbolicType) -> SymbolicType: return sympy.simplify(expr) @@ -1373,6 +2289,9 @@ def __init__(self, arrays, cpp_mode=False, *args, **kwargs): super().__init__(*args, **kwargs) self.arrays = arrays or set() self.cpp_mode = cpp_mode + # Print floats at the shortest precision that round-trips (``3.14`` instead of + # ``3.14000000000000``); higher-precision values (e.g. HUGE) still print in full. + self._settings['full_prec'] = False def _print_Float(self, expr): nf = sympy_numeric_fix(expr) @@ -1380,12 +2299,19 @@ def _print_Float(self, expr): return self._print(nf) return super()._print_Float(expr) + def _print_TypedConstant(self, expr): + value = self._print(expr.value) + if self.cpp_mode: + literal = dtypes.cpp_typed_literal(expr.value, expr.dtype) + if literal is not None: + return literal + return f'{expr.dtype.ctype}({value})' + return f'dace.{expr.dtype.to_string()}({value})' + def _print_Function(self, expr): if str(expr.func) in self.arrays: indices = ", ".join(self._print(arg) for arg in expr.args) return f'{expr.func}[{indices}]' - if self.cpp_mode and str(expr.func) == 'int_floor': - return '((%s) / (%s))' % (self._print(expr.args[0]), self._print(expr.args[1])) if str(expr.func) == 'AND': return f'(({self._print(expr.args[0])}) and ({self._print(expr.args[1])}))' if str(expr.func) == 'OR': @@ -1403,6 +2329,30 @@ def _print_Function(self, expr): attribute = self._print(expr.args[1]) return f'{self._print(expr.args[0])}{sep}{attribute}' # return f'{self._print(expr.args[0])}.{self._print(expr.args[1])}' + if str(expr.func) == 'Subscript': + indices = ', '.join(self._print(a) for a in expr.args[1:]) + return f'{self._print(expr.args[0])}[{indices}]' + # Operator-backed functions: the ``__``-prefixed variants (from the operators) + # always print as the operator; the bare names do so only in C++ (Python keeps + # the ``func(a, b)`` spelling so it round-trips). + name = str(expr.func) + base = name[2:] if name.startswith('__') else name + as_operator = name.startswith('__') or self.cpp_mode + if base == 'bitwise_invert' and as_operator: + return '(~(%s))' % self._print(expr.args[0]) + binop = {'bitwise_and': '&', 'bitwise_or': '|', 'bitwise_xor': '^', 'left_shift': '<<', 'right_shift': '>>'} + if base in binop and as_operator: + return '((%s) %s (%s))' % (self._print(expr.args[0]), binop[base], self._print(expr.args[1])) + # ``int_floor`` divides with ``//`` in Python and ``/`` in C++ (the bare name + # keeps its ``int_floor(a, b)`` spelling in Python so it round-trips). + if base == 'int_floor' and as_operator: + op = '/' if self.cpp_mode else '//' + return '((%s) %s (%s))' % (self._print(expr.args[0]), op, self._print(expr.args[1])) + if str(expr.func) == 'IfExpr': + cond, tval, fval = (self._print(a) for a in expr.args) + if self.cpp_mode: + return '((%s) ? (%s) : (%s))' % (cond, tval, fval) + return '((%s) if (%s) else (%s))' % (tval, cond, fval) return super()._print_Function(expr) def _print_Mod(self, expr): @@ -1418,14 +2368,20 @@ def _print_Not(self, expr): return '(not (%s))' % self._print(expr.args[0]) def _print_Infinity(self, expr): - if self.cpp_mode: - return 'INFINITY' - return super()._print_Infinity(expr) + # Print as ``inf`` so it round-trips back to ``oo`` via ``pystr_to_symbolic``. + return 'INFINITY' if self.cpp_mode else 'inf' def _print_NegativeInfinity(self, expr): - if self.cpp_mode: - return '-INFINITY' - return super()._print_NegativeInfinity(expr) + return '-INFINITY' if self.cpp_mode else '-inf' + + def _print_NaN(self, expr): + return 'NAN' if self.cpp_mode else 'nan' + + def _print_BooleanTrue(self, expr): + return 'true' if self.cpp_mode else 'True' + + def _print_BooleanFalse(self, expr): + return 'false' if self.cpp_mode else 'False' def _print_Symbol(self, expr): if expr.name == 'NoneSymbol': @@ -1466,8 +2422,8 @@ def _print_Pow(self, expr): return f'({self._print(expr.args[0])}) ** ({self._print(expr.args[1])})' -@lru_cache(maxsize=16384) -def symstr(sym, arrayexprs: Optional[Set[str]] = None, cpp_mode=False) -> str: +@lru_cache(maxsize=16384, typed=True) +def symstr(sym, arrayexprs: Optional[FrozenSet[str]] = None, cpp_mode=False) -> str: """ Convert a symbolic expression to a compilable expression. @@ -1482,6 +2438,13 @@ def symstr(sym, arrayexprs: Optional[Set[str]] = None, cpp_mode=False) -> str: if isinstance(sym, SymExpr): return symstr(sym.expr, arrayexprs, cpp_mode=cpp_mode) + # Infinity, NaN and booleans are atomic constants: ``sympy_numeric_fix`` rejects + # the first two and booleans are not ``Number``s, so the generic path below would + # wrap them in parentheses that no longer round-trip. Print them bare instead. + if isinstance(sym, (sympy.core.numbers.Infinity, sympy.core.numbers.NegativeInfinity, sympy.core.numbers.NaN, + sympy.logic.boolalg.BooleanAtom)): + return DaceSympyPrinter(arrayexprs, cpp_mode).doprint(sym) + try: sym = sympy_numeric_fix(sym) sym = sympy_intdiv_fix(sym) @@ -1490,7 +2453,7 @@ def symstr(sym, arrayexprs: Optional[Set[str]] = None, cpp_mode=False) -> str: sstr = DaceSympyPrinter(arrayexprs, cpp_mode).doprint(sym) if isinstance(sym, symbol) or isinstance(sym, sympy.Symbol) or isinstance( - sym, sympy.Number) or dtypes.isconstant(sym): + sym, (sympy.Number, TypedConstant)) or dtypes.isconstant(sym): return sstr else: return '(' + sstr + ')' @@ -1557,7 +2520,7 @@ def safe_replace(mapping: Dict[Union[SymbolicType, str], Union[SymbolicType, str replace_callback(invrepl) -@lru_cache(16384) +@lru_cache(maxsize=16384, typed=True) def _spickle(obj): return str(obj) diff --git a/dace/transformation/auto/auto_optimize.py b/dace/transformation/auto/auto_optimize.py index 8dea17ff71..eb3347429e 100644 --- a/dace/transformation/auto/auto_optimize.py +++ b/dace/transformation/auto/auto_optimize.py @@ -6,12 +6,10 @@ from dace.sdfg import infer_types from dace.sdfg.state import SDFGState, ControlFlowRegion from dace.sdfg.graph import SubgraphView -from dace.sdfg.propagation import propagate_states from dace.sdfg.scope import is_devicelevel_gpu_kernel from dace import config, data as dt, dtypes, Memlet, symbolic from dace.sdfg import SDFG, nodes, graph as gr -from typing import Set, Tuple, Union, List, Iterable, Dict, Callable -import warnings +from typing import Set, Tuple, Union, List, Dict, Callable # Transformations from dace.transformation.passes import FullMapFusion diff --git a/dace/transformation/dataflow/add_threadblock_map.py b/dace/transformation/dataflow/add_threadblock_map.py index febdb12861..f859171e5b 100644 --- a/dace/transformation/dataflow/add_threadblock_map.py +++ b/dace/transformation/dataflow/add_threadblock_map.py @@ -7,8 +7,7 @@ import sympy -import dace -from dace import Config, dtypes, symbolic +from dace import Config, dtypes from dace.properties import make_properties from dace.sdfg import SDFG, SDFGState, nodes, utils as sdutil from dace.transformation import helpers, transformation @@ -196,6 +195,7 @@ def apply(self, state: SDFGState, sdfg: SDFG): new_kernel_entry, *_ = helpers.get_parent_map(state, kernel_map_entry) new_kernel_entry.map.gpu_block_size = gpu_block_size new_kernel_entry.map.gpu_launch_bounds = kernel_map_entry.map.gpu_launch_bounds + new_kernel_entry.map.gpu_min_warps_per_eu = kernel_map_entry.map.gpu_min_warps_per_eu new_kernel_entry.map.gpu_maxnreg = kernel_map_entry.map.gpu_maxnreg def preprocess_default_dims(self): diff --git a/dace/transformation/dataflow/copy_to_map.py b/dace/transformation/dataflow/copy_to_map.py index c52b1da41b..4fe4ec32fd 100644 --- a/dace/transformation/dataflow/copy_to_map.py +++ b/dace/transformation/dataflow/copy_to_map.py @@ -1,6 +1,6 @@ # Copyright 2019-2022 ETH Zurich and the DaCe authors. All rights reserved. -from dace import dtypes, symbolic, data, subsets, Memlet, properties +from dace import properties from dace.transformation import transformation as xf from dace.sdfg import SDFGState, SDFG, nodes, utils as sdutil, memlet_utils as mutils diff --git a/dace/transformation/dataflow/gpu_grid_stride_tiling.py b/dace/transformation/dataflow/gpu_grid_stride_tiling.py index e9475168c2..b9ac5a2e26 100644 --- a/dace/transformation/dataflow/gpu_grid_stride_tiling.py +++ b/dace/transformation/dataflow/gpu_grid_stride_tiling.py @@ -2,7 +2,6 @@ """ This module contains classes and functions that implement the grid-strided map tiling transformation.""" -from typing import Dict import dace from copy import deepcopy as dcpy from dace import dtypes, subsets, symbolic @@ -13,7 +12,6 @@ from dace.transformation import transformation from dace.transformation.dataflow import MapInterchange from dace.transformation.dataflow.strip_mining import calc_set_image, calc_set_union -from dace.sdfg.propagation import propagate_memlet import sympy diff --git a/dace/transformation/dataflow/gpu_transform_local_storage.py b/dace/transformation/dataflow/gpu_transform_local_storage.py index 77a1603dc5..fa54f8ca4c 100644 --- a/dace/transformation/dataflow/gpu_transform_local_storage.py +++ b/dace/transformation/dataflow/gpu_transform_local_storage.py @@ -5,12 +5,11 @@ import copy import collections -from dace import data, dtypes, registry, sdfg as sd, subsets as sbs, symbolic +from dace import data, dtypes, sdfg as sd, subsets as sbs, symbolic from dace.sdfg import nodes, SDFGState from dace.sdfg import utils as sdutil from dace.transformation import transformation from dace.properties import Property, make_properties -from dace.config import Config def in_scope(graph, node, parent): @@ -412,7 +411,7 @@ def apply(self, graph: SDFGState, sdfg): if self.fullcopy: edge.data.subset = sbs.Range.from_array(node.desc(sdfg)) - edge.data.other_subset = newmemlet.subset + edge.data.other_subset = copy.deepcopy(newmemlet.subset) graph.add_edge(edge.src, edge.src_conn, node, None, edge.data) graph.remove_edge(edge) @@ -491,7 +490,7 @@ def apply(self, graph: SDFGState, sdfg): edge.data.wcr = None if self.fullcopy: edge.data.subset = sbs.Range.from_array(node.desc(sdfg)) - edge.data.other_subset = newmemlet.subset + edge.data.other_subset = copy.deepcopy(newmemlet.subset) graph.add_edge(node, None, edge.dst, edge.dst_conn, edge.data) graph.remove_edge(edge) diff --git a/dace/transformation/dataflow/map_distribution.py b/dace/transformation/dataflow/map_distribution.py index eaa8d07f2a..ac064c52e7 100644 --- a/dace/transformation/dataflow/map_distribution.py +++ b/dace/transformation/dataflow/map_distribution.py @@ -3,13 +3,11 @@ from copy import deepcopy from numbers import Number -from typing import Dict, List import dace import sympy -from dace import data, dtypes, subsets, symbolic +from dace import data, subsets from dace.sdfg import nodes from dace.sdfg import utils as sdutil -from dace.sdfg.graph import OrderedMultiDiConnectorGraph from dace.transformation import transformation as pm from dace.transformation.subgraph.helpers import subgraph_from_maps from functools import reduce diff --git a/dace/transformation/dataflow/map_expansion.py b/dace/transformation/dataflow/map_expansion.py index 26341ac378..a985d3b3f1 100644 --- a/dace/transformation/dataflow/map_expansion.py +++ b/dace/transformation/dataflow/map_expansion.py @@ -2,14 +2,13 @@ """ Contains classes that implement the map-expansion transformation. """ from dace.sdfg.utils import consolidate_edges -from typing import Dict, List +from typing import List import copy import dace from dace import dtypes, subsets, symbolic from dace.properties import EnumProperty, make_properties, Property from dace.sdfg import nodes from dace.sdfg import utils as sdutil -from dace.sdfg.graph import OrderedMultiDiConnectorGraph from dace.transformation import transformation as pm from dace.sdfg.propagation import propagate_memlets_scope diff --git a/dace/transformation/dataflow/map_fission.py b/dace/transformation/dataflow/map_fission.py index 7a107c2ac9..33be99fb01 100644 --- a/dace/transformation/dataflow/map_fission.py +++ b/dace/transformation/dataflow/map_fission.py @@ -1,10 +1,10 @@ -# Copyright 2019-2024 ETH Zurich and the DaCe authors. All rights reserved. +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. """ Map Fission transformation. """ from copy import deepcopy as dcpy from collections import defaultdict +from functools import reduce from dace import sdfg as sd, memlet as mm, subsets, data as dt -from dace.codegen import control_flow as cf from dace.properties import CodeBlock from dace.sdfg import nodes, graph as gr from dace.sdfg import utils as sdutil @@ -125,6 +125,14 @@ def can_be_applied(self, graph, expr_index, sdfg, permissive=False): # Get NestedSDFG control flow components nsdfg_node.sdfg.reset_cfg_list() + + # Fissioning a component across a conditional needs the branch + # condition replicated into each fissioned map, currently not supported. + if any( + isinstance(cfg, ConditionalBlock) + for cfg in nsdfg_node.sdfg.all_control_flow_regions(recursive=True)): + return False + if len(nsdfg_node.sdfg.nodes()) == 1: child = nsdfg_node.sdfg.nodes()[0] conditions: List[CodeBlock] = [] @@ -143,6 +151,29 @@ def can_be_applied(self, graph, expr_index, sdfg, permissive=False): return False if any(p in cond.get_free_symbols() for p in map_node.map.params): return False + # Reject if any interstate edge inside the nested SDFG has an + # assignment that depends on the map iterator, either directly or + # through a nested-SDFG input connector whose incoming memlet + # subset references a map parameter. Such assignments cannot be + # safely hoisted out of the fissioned maps. + map_params = set(map_node.map.params) + inputs_dep_on_map = set() + for e in graph.out_edges(map_node): + if e.dst is self.nested_sdfg and e.dst_conn is not None and e.data.subset is not None: + if any(str(s) in map_params for s in e.data.subset.free_symbols): + inputs_dep_on_map.add(e.dst_conn) + for ise in nsdfg_node.sdfg.all_interstate_edges(): + assign_free = set() + for expr in ise.data.assignments.values(): + try: + assign_free.update(str(s) for s in pystr_to_symbolic(expr).free_symbols) + except Exception: + pass + if assign_free & map_params: + return False + if assign_free & inputs_dep_on_map: + return False + helpers.nest_sdfg_control_flow(nsdfg_node.sdfg) subgraphs = list(nsdfg_node.sdfg.nodes()) @@ -207,10 +238,17 @@ def apply(self, graph: sd.SDFGState, sdfg: sd.SDFG): parent = nsdfg_node.sdfg parent_sdfg = parent.parent_sdfg modified_arrays = set() + scalar_like_arrays = set() # Get map information outer_map: nodes.Map = map_entry.map + # Border-transient extent equals the iteration count per dimension. + # Memlets that index border transients are normalized to + # `(p - iMin) / step` so the squeezed array remains in-bounds for + # strided maps. Symbolic steps are assumed non-negative. mapsize = outer_map.range.size() + squeezed_idx = [(pystr_to_symbolic(p) - iMin) / step + for p, (iMin, _iMax, step) in zip(outer_map.params, outer_map.range.ranges)] # Add new symbols from outer map to nested SDFG # Add new symbols also from the adjacent edge subsets and the data descriptors they carry. @@ -260,24 +298,37 @@ def apply(self, graph: sd.SDFGState, sdfg: sd.SDFG): if (isinstance(e.dst, nodes.AccessNode) and not nsdfg_node.sdfg.arrays[e.dst.data].transient) ] - # Map external edges to outer memlets + # Enclosing scope of the fissioned map (``None`` if top-level). + # The new per-component fission maps live in this scope, so + # boundary edges are rewired through it instead of through the about-to-be-removed map. + parent_entry = state.entry_node(map_entry) if self.expr_index == 0 else None + parent_exit = state.exit_node(parent_entry) if parent_entry is not None else None + + # For each boundary edge, record the edge just outside the + # original map so the new per-component maps can reconnect there. + # For a data edge that is its memlet-path neighbour + # (``path[eindex - 1]`` at entry, ``path[eindex + 1]`` at exit). + # An empty dependency edge has a one-element path and thus no + # outside edge: store ``None``. + edge_to_outer = {} edge_to_outer = {} for edge in external_edges_entry: if self.expr_index == 0: # Subgraphs use the corresponding outer map edges path = state.memlet_path(edge) eindex = path.index(edge) - edge_to_outer[edge] = path[eindex - 1] + edge_to_outer[edge] = path[eindex - 1] if eindex > 0 else None else: - # Nested SDFGs use the internal map edges of the node - outer_edge = next(e for e in graph.in_edges(nsdfg_node) if e.dst_conn == edge.src.data) + outer_edge = next((e for e in graph.in_edges(nsdfg_node) if e.dst_conn == edge.src.data), None) + if outer_edge is None: + outer_edge = next(e for e in graph.out_edges(nsdfg_node) if e.src_conn == edge.src.data) edge_to_outer[edge] = outer_edge for edge in external_edges_exit: if self.expr_index == 0: path = state.memlet_path(edge) eindex = path.index(edge) - edge_to_outer[edge] = path[eindex + 1] + edge_to_outer[edge] = path[eindex + 1] if eindex + 1 < len(path) else None else: # Nested SDFGs use the internal map edges of the node outer_edge = next(e for e in graph.out_edges(nsdfg_node) if e.src_conn == edge.dst.data) @@ -307,9 +358,7 @@ def apply(self, graph: sd.SDFGState, sdfg: sd.SDFG): # Add extra nodes in component boundaries for edge in edges: anode = state.add_access(name) - sbs = subsets.Range.from_string(','.join(outer_map.params)) - # Offset memlet by map range begin (to fit the transient) - sbs.offset([r[0] for r in outer_map.range], True) + sbs = subsets.Range([(idx, idx, 1) for idx in squeezed_idx]) state.add_edge(edge.src, edge.src_conn, anode, None, mm.Memlet.simple(name, sbs, num_accesses=outer_map.range.num_elements())) state.add_edge(anode, None, edge.dst, edge.dst_conn, @@ -321,8 +370,7 @@ def apply(self, graph: sd.SDFGState, sdfg: sd.SDFG): for component_in, component_out in components: me, mx = state.add_map(outer_map.label + '_fission', [(p, '0:1') for p in outer_map.params], outer_map.schedule, - unroll=outer_map.unroll, - debuginfo=outer_map.debuginfo) + unroll=outer_map.unroll) # Add dynamic input connectors for conn in map_entry.in_connectors: @@ -347,8 +395,15 @@ def apply(self, graph: sd.SDFGState, sdfg: sd.SDFG): state.add_edge(me, out_conn, e.dst, e.dst_conn, dcpy(e.data)) # Reconnect inner edges at source directly to external nodes if self.expr_index == 0 and e in external_edges_entry: - state.add_edge(edge_to_outer[e].src, edge_to_outer[e].src_conn, me, in_conn, - dcpy(edge_to_outer[e].data)) + outer = edge_to_outer[e] + if outer is None: + # Empty dependency edge: keep the + # new map inside the enclosing scope via an empty + # edge instead of aliasing the removed map_entry. + if parent_entry is not None: + state.add_edge(parent_entry, None, me, None, mm.Memlet()) + else: + state.add_edge(outer.src, outer.src_conn, me, in_conn, dcpy(outer.data)) else: state.add_edge(e.src, e.src_conn, me, in_conn, dcpy(e.data)) state.remove_edge(e) @@ -370,8 +425,15 @@ def apply(self, graph: sd.SDFGState, sdfg: sd.SDFG): state.add_edge(e.src, e.src_conn, mx, in_conn, dcpy(e.data)) # Reconnect inner edges at sink directly to external nodes if self.expr_index == 0 and e in external_edges_exit: - state.add_edge(mx, out_conn, edge_to_outer[e].dst, edge_to_outer[e].dst_conn, - dcpy(edge_to_outer[e].data)) + outer = edge_to_outer[e] + if outer is None: + # Empty dependency edge: keep + # the new map inside the enclosing scope via an + # empty edge instead of aliasing the removed map_exit. + if parent_exit is not None: + state.add_edge(mx, None, parent_exit, None, mm.Memlet()) + else: + state.add_edge(mx, out_conn, outer.dst, outer.dst_conn, dcpy(outer.data)) else: state.add_edge(mx, out_conn, e.dst, e.dst_conn, dcpy(e.data)) state.remove_edge(e) @@ -384,7 +446,9 @@ def apply(self, graph: sd.SDFGState, sdfg: sd.SDFG): for node in sources: if isinstance(node, nodes.AccessNode): for edge in state.in_edges(node): - outer_edge = edge_to_outer[edge] + outer_edge = edge_to_outer.get(edge) + if outer_edge is None: # No outer feeder: nothing to rewire. + continue memlet = dcpy(edge.data) memlet.subset = subsets.Range(outer_map.range.ranges + memlet.subset.ranges) state.add_edge(outer_edge.src, outer_edge.src_conn, edge.dst, edge.dst_conn, memlet) @@ -392,7 +456,9 @@ def apply(self, graph: sd.SDFGState, sdfg: sd.SDFG): for node in sinks: if isinstance(node, nodes.AccessNode): for edge in state.out_edges(node): - outer_edge = edge_to_outer[edge] + outer_edge = edge_to_outer.get(edge) + if outer_edge is None: # No outer consumer: nothing to rewire. + continue state.add_edge(edge.src, edge.src_conn, outer_edge.dst, outer_edge.dst_conn, dcpy(outer_edge.data)) @@ -401,17 +467,33 @@ def apply(self, graph: sd.SDFGState, sdfg: sd.SDFG): if array in modified_arrays: continue desc = parent.arrays[array] - if isinstance(desc, dt.Scalar): # Scalar needs to be augmented to an array + # Treat scalars and length-1 arrays as "scalar-like": their single + # degenerate dimension is replaced by the map dims rather than + # prepended, so the result is shape [extent] rather than + # [extent, 1] (which produces zero-stride aliasing). + scalar_like = (isinstance(desc, dt.Scalar) + or (isinstance(desc, dt.Array) and len(desc.shape) == 1 and desc.shape[0] == 1)) + if isinstance(desc, dt.Scalar): desc = dt.Array(desc.dtype, desc.shape, desc.transient, desc.allow_conflicts, desc.storage, desc.location, desc.strides, desc.offset, False, desc.lifetime, 0, desc.debuginfo, desc.total_size, desc.start_offset) parent.arrays[array] = desc - for sz in reversed(mapsize): - desc.strides = [desc.total_size] + list(desc.strides) - desc.total_size = desc.total_size * sz - desc.shape = mapsize + list(desc.shape) - desc.offset = [0] * len(mapsize) + list(desc.offset) + if scalar_like: + desc.shape = list(mapsize) + strides = [1] * len(mapsize) + for i in range(len(mapsize) - 2, -1, -1): + strides[i] = strides[i + 1] * mapsize[i + 1] + desc.strides = strides + desc.total_size = reduce(lambda a, b: a * b, mapsize, 1) + desc.offset = [0] * len(mapsize) + scalar_like_arrays.add(array) + else: + for sz in reversed(mapsize): + desc.strides = [desc.total_size] + list(desc.strides) + desc.total_size = desc.total_size * sz + desc.shape = list(mapsize) + list(desc.shape) + desc.offset = [0] * len(mapsize) + list(desc.offset) modified_arrays.add(array) # Fill scope connectors so that memlets can be tracked below @@ -473,24 +555,26 @@ def apply(self, graph: sd.SDFGState, sdfg: sd.SDFG): # NOTE: Memlet propagation should run to correct the outer edges for node in subgraph.nodes(): if isinstance(node, nodes.AccessNode) and node.data in arrays: + is_scalar_like = node.data in scalar_like_arrays for edge in state.all_edges(node): for e in state.memlet_tree(edge): # Prepend map dimensions to memlet # NOTE: Do this only for the subset corresponding to `node.data`. If the edge is copying # to/from another AccessNode, the other data may not need extra dimensions. For example, see # `test.transformations.mapfission_test.MapFissionTest.test_array_copy_outside_scope`. + map_ranges = [(idx, idx, 1) for idx in squeezed_idx] if e.data.data == node.data: if e.data.subset: - e.data.subset = subsets.Range([(pystr_to_symbolic(d) - r[0], - pystr_to_symbolic(d) - r[0], 1) - for d, r in zip(outer_map.params, outer_map.range)] + - e.data.subset.ranges) + if is_scalar_like: + e.data.subset = subsets.Range(map_ranges) + else: + e.data.subset = subsets.Range(map_ranges + e.data.subset.ranges) else: if e.data.other_subset: - e.data.other_subset = subsets.Range( - [(pystr_to_symbolic(d) - r[0], pystr_to_symbolic(d) - r[0], 1) - for d, r in zip(outer_map.params, outer_map.range)] + - e.data.other_subset.ranges) + if is_scalar_like: + e.data.other_subset = subsets.Range(map_ranges) + else: + e.data.other_subset = subsets.Range(map_ranges + e.data.other_subset.ranges) # If nested SDFG, reconnect nodes around map and modify memlets if self.expr_index == 1: diff --git a/dace/transformation/dataflow/map_fusion_helper.py b/dace/transformation/dataflow/map_fusion_helper.py index 5191c96b7d..ac962b5489 100644 --- a/dace/transformation/dataflow/map_fusion_helper.py +++ b/dace/transformation/dataflow/map_fusion_helper.py @@ -4,7 +4,7 @@ import dace from dace import symbolic -from dace.sdfg import graph, nodes as nodes, propagation, validation +from dace.sdfg import graph, nodes as nodes, validation from dace.transformation import helpers diff --git a/dace/transformation/dataflow/map_fusion_vertical.py b/dace/transformation/dataflow/map_fusion_vertical.py index 89883425cb..92b0b109d5 100644 --- a/dace/transformation/dataflow/map_fusion_vertical.py +++ b/dace/transformation/dataflow/map_fusion_vertical.py @@ -8,6 +8,7 @@ from dace.sdfg import SDFG, SDFGState, graph, nodes, propagation from dace.transformation.dataflow import map_fusion_helper as mfhelper from dace.sdfg.type_inference import infer_expr_type +from ordered_set import OrderedSet @properties.make_properties @@ -405,9 +406,9 @@ def partition_first_outputs( param_repl: Dict[str, str], ) -> Union[ Tuple[ - Set[graph.MultiConnectorEdge[dace.Memlet]], - Set[graph.MultiConnectorEdge[dace.Memlet]], - Set[graph.MultiConnectorEdge[dace.Memlet]], + OrderedSet[graph.MultiConnectorEdge[dace.Memlet]], + OrderedSet[graph.MultiConnectorEdge[dace.Memlet]], + OrderedSet[graph.MultiConnectorEdge[dace.Memlet]], ], None, ]: @@ -447,9 +448,9 @@ def partition_first_outputs( `require_all_intermediates` and by `self.require_exclusive_intermediates`. """ # The three outputs set. - pure_outputs: Set[graph.MultiConnectorEdge[dace.Memlet]] = set() - exclusive_outputs: Set[graph.MultiConnectorEdge[dace.Memlet]] = set() - shared_outputs: Set[graph.MultiConnectorEdge[dace.Memlet]] = set() + pure_outputs: Set[graph.MultiConnectorEdge[dace.Memlet]] = OrderedSet() + exclusive_outputs: Set[graph.MultiConnectorEdge[dace.Memlet]] = OrderedSet() + shared_outputs: Set[graph.MultiConnectorEdge[dace.Memlet]] = OrderedSet() # Set of intermediate nodes that we have already processed. processed_inter_nodes: Set[nodes.Node] = set() @@ -703,7 +704,7 @@ def partition_first_outputs( def handle_intermediate_set( self, - intermediate_outputs: Set[graph.MultiConnectorEdge[dace.Memlet]], + intermediate_outputs: OrderedSet[graph.MultiConnectorEdge[dace.Memlet]], state: dace.SDFGState, sdfg: SDFG, first_map_exit: nodes.MapExit, @@ -870,7 +871,7 @@ def handle_intermediate_set( # the input connectors on the MapEntry, such that we know where we # have to reroute inside the Map. # NOTE: Assumes that Map (if connected is the direct neighbour). - conn_names: Set[str] = set() + conn_names: OrderedSet[str] = OrderedSet() for inter_node_out_edge in state.out_edges(inter_node): if inter_node_out_edge.dst == second_map_entry: assert inter_node_out_edge.dst_conn.startswith("IN_") diff --git a/dace/transformation/dataflow/map_interchange.py b/dace/transformation/dataflow/map_interchange.py index 971f3172fe..499b8a291c 100644 --- a/dace/transformation/dataflow/map_interchange.py +++ b/dace/transformation/dataflow/map_interchange.py @@ -7,7 +7,6 @@ from dace.symbolic import symlist from dace.transformation import transformation from dace.sdfg.propagation import propagate_memlet -from dace.properties import make_properties class MapInterchange(transformation.SingleStateTransformation): diff --git a/dace/transformation/dataflow/map_unroll.py b/dace/transformation/dataflow/map_unroll.py index 575a4ac93b..b040d0ceab 100644 --- a/dace/transformation/dataflow/map_unroll.py +++ b/dace/transformation/dataflow/map_unroll.py @@ -1,5 +1,5 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. -from dace import data as dt, symbolic, SDFG +from dace import data as dt, symbolic, SDFG, __version__ from dace.sdfg import nodes, utils as sdutil from dace.sdfg.state import SDFGState from dace.transformation import transformation @@ -85,6 +85,7 @@ def apply(self, state: SDFGState, sdfg: SDFG): nsdfg_json = nested_sdfgs[nsdfg] name = nsdfg_json["attributes"]["name"] nsdfg_json["attributes"]["name"] += suffix + nsdfg_json["dace_version"] = __version__ unrolled_nsdfg = SDFG.from_json(nsdfg_json) nsdfg_json["attributes"]["name"] = name # Reinstate # Set all the references diff --git a/dace/transformation/dataflow/mapreduce.py b/dace/transformation/dataflow/mapreduce.py index b10245524b..91174c581d 100644 --- a/dace/transformation/dataflow/mapreduce.py +++ b/dace/transformation/dataflow/mapreduce.py @@ -8,7 +8,6 @@ from dace.properties import Property, make_properties from dace.sdfg import SDFG from dace.sdfg import utils as sdutil -from dace.symbolic import symstr from dace.transformation import transformation as pm from dace.transformation.dataflow.map_collapse import MapCollapse diff --git a/dace/transformation/dataflow/matrix_product_transpose.py b/dace/transformation/dataflow/matrix_product_transpose.py index 5d8bf04fb6..cf83a1022a 100644 --- a/dace/transformation/dataflow/matrix_product_transpose.py +++ b/dace/transformation/dataflow/matrix_product_transpose.py @@ -1,4 +1,4 @@ -# Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. """ Implements the matrix-matrix product transpose transformation. """ from copy import deepcopy as dcpy @@ -17,11 +17,11 @@ class MatrixProductTranspose(transformation.SingleStateTransformation): T(A) @ T(B) = T(B @ A) """ import dace.libraries.blas as blas # Avoid slow imports - import dace.libraries.standard as std # Avoid slow imports + import dace.libraries.linalg as linalg # Avoid slow imports - transpose_a = transformation.PatternNode(std.Transpose) + transpose_a = transformation.PatternNode(linalg.Transpose) at = transformation.PatternNode(nodes.AccessNode) - transpose_b = transformation.PatternNode(std.Transpose) + transpose_b = transformation.PatternNode(linalg.Transpose) bt = transformation.PatternNode(nodes.AccessNode) a_times_b = transformation.PatternNode(blas.MatMul) @@ -58,7 +58,7 @@ def match_to_str(self, graph): return f"{transpose_a.name} -> {a_times_b.name} <- {transpose_b.name}" def apply(self, graph: SDFGState, sdfg: SDFG): - import dace.libraries.standard as std + import dace.libraries.linalg as linalg # Avoid slow imports transpose_a = self.transpose_a _at = self.at @@ -83,7 +83,7 @@ def apply(self, graph: SDFGState, sdfg: SDFG): break tmp_name, tmp_arr = sdfg.add_temp_transient(shape, a_times_b.dtype) tmp_acc = graph.add_access(tmp_name) - transpose_c = std.Transpose('_Transpose_', a_times_b.dtype) + transpose_c = linalg.Transpose('_Transpose_', a_times_b.dtype) for edge in graph.out_edges(a_times_b): _, _, dst, dst_conn, memlet = edge graph.remove_edge(edge) diff --git a/dace/transformation/dataflow/merge_arrays.py b/dace/transformation/dataflow/merge_arrays.py index 2d3aba1299..78911582ce 100644 --- a/dace/transformation/dataflow/merge_arrays.py +++ b/dace/transformation/dataflow/merge_arrays.py @@ -1,9 +1,7 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. from dace.transformation import transformation -from dace import memlet from dace.sdfg.graph import OrderedDiGraph -from dace import memlet -from dace.sdfg import nodes, utils, graph as gr +from dace.sdfg import nodes, utils from dace.sdfg import SDFGState from dace.sdfg.propagation import propagate_memlet diff --git a/dace/transformation/dataflow/otf_map_fusion.py b/dace/transformation/dataflow/otf_map_fusion.py index 0c9a147959..cc2c08687e 100644 --- a/dace/transformation/dataflow/otf_map_fusion.py +++ b/dace/transformation/dataflow/otf_map_fusion.py @@ -13,10 +13,9 @@ from dace.memlet import Memlet from dace.sdfg import utils as sdutil from dace.transformation import transformation -from dace import data as dt from dace import dtypes from dace import symbolic, nodes -from dace.properties import SymbolicProperty, make_properties, Property +from dace.properties import SymbolicProperty, make_properties from dace.transformation.dataflow.stream_transient import AccumulateTransient from dace.transformation.dataflow.local_storage import OutLocalStorage, InLocalStorage diff --git a/dace/transformation/dataflow/prune_connectors.py b/dace/transformation/dataflow/prune_connectors.py index b59cf77cee..3ef6cf0b45 100644 --- a/dace/transformation/dataflow/prune_connectors.py +++ b/dace/transformation/dataflow/prune_connectors.py @@ -2,11 +2,10 @@ from typing import Set, Tuple import re -from dace import dtypes, SDFG, SDFGState, symbolic, properties, data as dt +from dace import dtypes, SDFG, SDFGState, symbolic, properties from dace.transformation import transformation as pm, helpers from dace.sdfg import nodes, utils from dace.sdfg.analysis import cfg -from dace.sdfg.state import StateSubgraphView @properties.make_properties @@ -142,7 +141,7 @@ def _candidates(nsdfg: nodes.NestedSDFG) -> Set[str]: candidates -= (set(map(str, symbolic.symbols_in_ast(e.data.condition.code[0]))) - ignore) for assign in e.data.assignments.values(): - candidates -= (symbolic.free_symbols_and_functions(assign) - ignore) + candidates -= ((symbolic.free_symbols_and_functions(assign) | symbolic.arrays(assign)) - ignore) if local_ignore is None: local_ignore = set(e.data.assignments.keys()) diff --git a/dace/transformation/dataflow/redundant_array.py b/dace/transformation/dataflow/redundant_array.py index e624b9b273..4aacea151d 100644 --- a/dace/transformation/dataflow/redundant_array.py +++ b/dace/transformation/dataflow/redundant_array.py @@ -12,7 +12,6 @@ from dace import data, dtypes from dace import memlet as mm from dace import subsets, symbolic -from dace.config import Config from dace.sdfg import SDFG, SDFGState, graph, nodes from dace.sdfg import utils as sdutil from dace.transformation import helpers @@ -123,35 +122,42 @@ def find_dims_to_pop( a_size: Sequence[symbolic.SymbolicType], b_size: Sequence[symbolic.SymbolicType], ) -> List[int]: - """Determine how the first subset has to be squeezed to get to the dimension of the second subset. + """ + Determine how the first subset has to be squeezed to get to the dimension of the second subset. - Essentially the function determines which dimensions from the subset `A` have to be removed - to get down to the dimensionality of subset `B`. In case they have the same dimensionality + Essentially the function determines which dimensions from the subset ``A`` have to be removed + to get down to the dimensionality of subset ``B``. In case they have the same dimensionality an empty list is returned. It is important that this function does not operates on the actual subsets but on their sizes. Returns: - The function will return the list of dimensions that have to be removed in subset `A` - to bring it to the same dimensionality as subset `B`. + The function will return the list of dimensions that have to be removed in subset ``A`` + to bring it to the same dimensionality as subset ``B``. Note: This function is wrong. However, every attempt of fixing it resulted in breaking other things in strange ways. The only solution would be to redo all transformations. - For example, let's consider the following cases: - - `a_size := [3, 3], b_size := [9]` in this case the function will return `[0, 1]`, which - might be correct in some cases but not in other, depending on the strides of `A` + For example, consider the following cases: + + - ``a_size := [3, 3], b_size := [9]`` in this case the function will return ``[0, 1]``, which + might be correct in some cases but not in other, depending on the strides of ``A`` - While the example above can make sense the following example, - `a_size := [4, 4], b_size := [9]` which also results in `[0, 1]` does not make sense + ``a_size := [4, 4], b_size := [9]`` which also results in ``[0, 1]`` does not make sense and is clearly wrong. - - If a size is listed multiple times, for example `a_size := [1, 1, 5], b := [1, 5]` then - the function will return `[]` although it should return `[0]` or `[1]`. + - If a size is listed multiple times, for example ``a_size := [1, 1, 5], b := [1, 5]`` then + the function will return ``[]`` although it should return ``[0]`` or ``[1]``. - Furthermore, if not all dimensions could be matched then some unspecific value is - returned, for example `a_size := [1, 3, 5], b_size := [1, 8]` will return `[1, 2]`. + returned, for example ``a_size := [1, 3, 5], b_size := [1, 8]`` will return ``[1, 2]``. Note that this in particular is not a real case scenario as it is clearly an error in the subsets. - There is an improved version, `find_dims_to_pop2()` which has a better behaviour. + There is an improved version, ``find_dims_to_pop2()`` which has a better behaviour. + + :param a_size: The size of the first subset. + :param b_size: The size of the second subset. + :return: The function will return the list of dimensions that have to be removed in subset ``A`` + to bring it to the same dimensionality as subset ``B``. """ dims_to_pop = [] for i, sz in enumerate(reversed(a_size)): @@ -163,30 +169,38 @@ def find_dims_to_pop( def find_dims_to_pop2( a_size: Sequence[symbolic.SymbolicType], b_size: Sequence[symbolic.SymbolicType], -) -> List[int]: - """Determine how the first subset has to be squeezed to get to the dimension of the second subset. +) -> Optional[List[int]]: + """ + Determine how the first subset has to be squeezed to get to the dimension of the second subset. - Essentially the function determines which dimensions from the subset `A` have to be removed - to get down to the dimensionality of subset `B`. In case they have the same dimensionality + Essentially the function determines which dimensions from the subset ``A`` have to be removed + to get down to the dimensionality of subset ``B``. In case they have the same dimensionality an empty list is returned. A base assumption is the function is that a View is not used to modify the relative order of dimensions, i.e. only dummy dimensions are inserted at certain locations. It is important that this function does not operates on the actual subsets but on their sizes. - This is an improved version of `find_dims_to_pop()`. + This is an improved version of ``find_dims_to_pop()``. Returns: - The function will return the list of dimensions that have to be removed in subset `A` - to bring it to the same dimensionality as subset `B`. In case the function failed to - associate all `B` dimensions to a corresponding `A` dimension the function returns + The function will return the list of dimensions that have to be removed in subset ``A`` + to bring it to the same dimensionality as subset ``B``. In case the function failed to + associate all ``B`` dimensions to a corresponding ``A`` dimension the function returns `None`. Note: - This function behaves differently than `find_dims_to_pop()` in some cases. In most cases + This function behaves differently than ``find_dims_to_pop()`` in some cases. In most cases the function ensures that a matching has been found. Keep in mind that the function assumes that the relative ordering of dimensions is the same on the source and the destination. + + :param a_size: The size of the first subset. + :param b_size: The size of the second subset. + :return: The function will return the list of dimensions that have to be removed in subset ``A`` + to bring it to the same dimensionality as subset ``B``. In case the function failed to + associate all ``B`` dimensions to a corresponding ``A`` dimension the function returns + ``None``. """ if len(a_size) < len(b_size): raise ValueError( diff --git a/dace/transformation/dataflow/redundant_array_copying.py b/dace/transformation/dataflow/redundant_array_copying.py index 8e4985c626..cda12f04cc 100644 --- a/dace/transformation/dataflow/redundant_array_copying.py +++ b/dace/transformation/dataflow/redundant_array_copying.py @@ -1,13 +1,11 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. """ Contains redundant array removal transformations. """ -from dace import registry from dace.sdfg import nodes from dace.sdfg import utils as sdutil from dace.sdfg.sdfg import SDFG from dace.sdfg.state import SDFGState from dace.transformation import transformation as pm -from dace.config import Config class RedundantArrayCopyingIn(pm.SingleStateTransformation): diff --git a/dace/transformation/dataflow/stream_transient.py b/dace/transformation/dataflow/stream_transient.py index cc928435be..8dfb618d5b 100644 --- a/dace/transformation/dataflow/stream_transient.py +++ b/dace/transformation/dataflow/stream_transient.py @@ -6,8 +6,7 @@ from dace.symbolic import symstr import warnings -from dace import data, dtypes, registry, symbolic, subsets -from dace.frontend.operations import detect_reduction_type +from dace import data, symbolic, subsets from dace.properties import SymbolicProperty, make_properties, Property from dace.sdfg import nodes from dace.sdfg import SDFG diff --git a/dace/transformation/dataflow/streaming_memory.py b/dace/transformation/dataflow/streaming_memory.py index 49d1dc8cb8..800f1819b5 100644 --- a/dace/transformation/dataflow/streaming_memory.py +++ b/dace/transformation/dataflow/streaming_memory.py @@ -8,7 +8,7 @@ import sympy from dace.transformation import transformation as xf -from dace import (data, dtypes, nodes, properties, registry, memlet as mm, subsets, symbolic, symbol, Memlet) +from dace import (data, dtypes, nodes, properties, memlet as mm, subsets, symbolic, symbol, Memlet) from dace.sdfg import SDFG, SDFGState, utils as sdutil, graph as gr from dace.libraries.standard import Gearbox diff --git a/dace/transformation/dataflow/strip_mining.py b/dace/transformation/dataflow/strip_mining.py index 7282a12521..a024489228 100644 --- a/dace/transformation/dataflow/strip_mining.py +++ b/dace/transformation/dataflow/strip_mining.py @@ -4,7 +4,7 @@ import dace from copy import deepcopy as dcpy -from dace import dtypes, registry, subsets, symbolic +from dace import dtypes, subsets, symbolic from dace.sdfg import SDFG, SDFGState from dace.properties import EnumProperty, make_properties, Property, SymbolicProperty from dace.sdfg import nodes diff --git a/dace/transformation/dataflow/sve/infer_types.py b/dace/transformation/dataflow/sve/infer_types.py index cc1c8dc431..67df0feb48 100644 --- a/dace/transformation/dataflow/sve/infer_types.py +++ b/dace/transformation/dataflow/sve/infer_types.py @@ -2,22 +2,15 @@ """ SVE Infer Types: This module is responsible for inferring connector types in the SDFG. """ -from typing import * -from dace.sdfg.graph import MultiConnectorEdge, Graph, SubgraphView +from dace.sdfg.graph import SubgraphView from dace.sdfg.state import SDFGState from dace.sdfg import nodes, SDFG, SDFGState -from dace.sdfg.nodes import Node, Tasklet -import dace.dtypes -import dace.sdfg.infer_types -import dace.transformation.dataflow -import dace.transformation.helpers -import dace.codegen.targets.sve.util -import dace.frontend.operations +from dace.sdfg.nodes import Tasklet import dace.data as data import dace.dtypes as dtypes -from collections import defaultdict from dace.sdfg.utils import dfs_topological_sort from dace.sdfg.type_inference import infer_types +from typing import Tuple, DefaultDict class TypeInferenceDict(DefaultDict[Tuple[Tasklet, str, bool], dtypes.typeclass]): diff --git a/dace/transformation/dataflow/tiling.py b/dace/transformation/dataflow/tiling.py index a89e9813da..415107d37d 100644 --- a/dace/transformation/dataflow/tiling.py +++ b/dace/transformation/dataflow/tiling.py @@ -2,7 +2,7 @@ """ This module contains classes and functions that implement the orthogonal tiling transformation. """ -from dace import registry, symbolic +from dace import symbolic from dace.properties import make_properties, Property, ShapeProperty from dace.sdfg import nodes from dace.sdfg import utils as sdutil diff --git a/dace/transformation/dataflow/trivial_tasklet_elimination.py b/dace/transformation/dataflow/trivial_tasklet_elimination.py index 241d010a2a..21eb50f796 100644 --- a/dace/transformation/dataflow/trivial_tasklet_elimination.py +++ b/dace/transformation/dataflow/trivial_tasklet_elimination.py @@ -1,7 +1,7 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. """ Contains classes that implement the trivial-tasklet-elimination transformation. """ -from dace import data, registry +from dace import data from dace.sdfg import nodes from dace.sdfg import utils as sdutil from dace.transformation import transformation diff --git a/dace/transformation/dataflow/vectorization.py b/dace/transformation/dataflow/vectorization.py index 1f4ad08165..2217e40161 100644 --- a/dace/transformation/dataflow/vectorization.py +++ b/dace/transformation/dataflow/vectorization.py @@ -1,13 +1,12 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. """ Contains classes that implement the vectorization transformation. """ -from dace import data, dtypes, registry, symbolic, subsets +from dace import data, dtypes, symbolic, subsets from dace.sdfg import nodes, SDFG, SDFGState, propagation from dace.sdfg import utils as sdutil from dace.sdfg.scope import ScopeSubgraphView from dace.transformation import transformation from dace.transformation.helpers import replicate_scope from dace.properties import Property, make_properties -import itertools @make_properties diff --git a/dace/transformation/dataflow/warp_tiling.py b/dace/transformation/dataflow/warp_tiling.py index 056e93741a..e6f5462063 100644 --- a/dace/transformation/dataflow/warp_tiling.py +++ b/dace/transformation/dataflow/warp_tiling.py @@ -1,7 +1,7 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. import copy from dace.sdfg.graph import SubgraphView -from dace import registry, properties, nodes, dtypes, subsets, symbolic +from dace import properties, nodes, dtypes, subsets, symbolic from dace import Memlet, SDFG, SDFGState from dace.frontend.operations import detect_reduction_type from dace.transformation import transformation as xf, helpers as xfh diff --git a/dace/transformation/dataflow/wcr_conversion.py b/dace/transformation/dataflow/wcr_conversion.py index c35c3e4545..937d51288f 100644 --- a/dace/transformation/dataflow/wcr_conversion.py +++ b/dace/transformation/dataflow/wcr_conversion.py @@ -9,7 +9,6 @@ from dace.transformation import transformation from dace.sdfg import utils as sdutil from dace import Memlet, SDFG, SDFGState -from dace.transformation import helpers from dace.sdfg.propagation import propagate_memlets_state diff --git a/dace/transformation/estimator/enumeration/brute_force_enumerator.py b/dace/transformation/estimator/enumeration/brute_force_enumerator.py index 1b4ea608e3..91d73e7161 100644 --- a/dace/transformation/estimator/enumeration/brute_force_enumerator.py +++ b/dace/transformation/estimator/enumeration/brute_force_enumerator.py @@ -3,15 +3,12 @@ from dace.transformation.estimator.enumeration import MapScoringEnumerator -from dace.transformation.subgraph import SubgraphFusion, helpers -from dace.properties import make_properties, Property +from dace.transformation.subgraph import helpers +from dace.properties import make_properties from dace.sdfg import SDFG, SDFGState from dace.sdfg.graph import SubgraphView -import dace.sdfg.nodes as nodes - -from collections import deque, defaultdict, ChainMap -from typing import Set, Union, List, Callable +from typing import Callable import itertools diff --git a/dace/transformation/estimator/enumeration/connected_enumerator.py b/dace/transformation/estimator/enumeration/connected_enumerator.py index 8bfb6d4963..ac09b73931 100644 --- a/dace/transformation/estimator/enumeration/connected_enumerator.py +++ b/dace/transformation/estimator/enumeration/connected_enumerator.py @@ -3,16 +3,12 @@ from dace.transformation.estimator.enumeration import MapScoringEnumerator -from dace.transformation.subgraph import SubgraphFusion, helpers +from dace.transformation.subgraph import helpers from dace.properties import make_properties, Property from dace.sdfg import SDFG, SDFGState from dace.sdfg.graph import SubgraphView -import dace.sdfg.nodes as nodes - -from collections import deque, defaultdict, ChainMap -from typing import Set, Union, List, Callable -import itertools +from typing import Set, List, Callable @make_properties diff --git a/dace/transformation/estimator/enumeration/enumerator.py b/dace/transformation/estimator/enumeration/enumerator.py index 912f5c155d..c27c5d1e18 100644 --- a/dace/transformation/estimator/enumeration/enumerator.py +++ b/dace/transformation/estimator/enumeration/enumerator.py @@ -1,17 +1,16 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. """ This file implements the Enuerator class """ -from dace.transformation.subgraph import SubgraphFusion, helpers +from dace.transformation.subgraph import helpers from dace.properties import make_properties, Property from dace.sdfg import SDFG, SDFGState from dace.sdfg.graph import SubgraphView import dace.sdfg.nodes as nodes -from collections import deque, defaultdict, ChainMap -from typing import Set, Union, List, Callable +from collections import defaultdict +from typing import Callable import itertools -import warnings @make_properties diff --git a/dace/transformation/estimator/enumeration/greedy_enumerator.py b/dace/transformation/estimator/enumeration/greedy_enumerator.py index 70cd1a8e34..f12a80941d 100644 --- a/dace/transformation/estimator/enumeration/greedy_enumerator.py +++ b/dace/transformation/estimator/enumeration/greedy_enumerator.py @@ -3,17 +3,13 @@ from dace.transformation.estimator.enumeration import Enumerator -from dace.transformation.subgraph import SubgraphFusion, helpers +from dace.transformation.subgraph import helpers from dace.transformation.subgraph.composite import CompositeFusion from dace.properties import make_properties, Property from dace.sdfg import SDFG, SDFGState from dace.sdfg.graph import SubgraphView -import dace.sdfg.nodes as nodes - -from collections import deque, defaultdict, ChainMap -from typing import Set, Union, List, Callable -import itertools +from typing import Callable import heapq diff --git a/dace/transformation/estimator/enumeration/scoring_enumerator.py b/dace/transformation/estimator/enumeration/scoring_enumerator.py index 0b5603712f..9597868fd9 100644 --- a/dace/transformation/estimator/enumeration/scoring_enumerator.py +++ b/dace/transformation/estimator/enumeration/scoring_enumerator.py @@ -1,18 +1,9 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. """ This file implements the Enuerator class """ -from dace.transformation.subgraph import SubgraphFusion, helpers from dace.properties import make_properties, Property -from dace.sdfg import SDFG, SDFGState -from dace.sdfg.graph import SubgraphView from dace.transformation.estimator.enumeration import Enumerator -import dace.sdfg.nodes as nodes - -from collections import deque, defaultdict, ChainMap -from typing import Set, Union, List, Callable -import itertools -import warnings @make_properties diff --git a/dace/transformation/helpers.py b/dace/transformation/helpers.py index 4a46c3ab19..2b2eed5aa5 100644 --- a/dace/transformation/helpers.py +++ b/dace/transformation/helpers.py @@ -1,5 +1,6 @@ -# Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. """ Transformation helper API. """ +import ast import copy import itertools import warnings @@ -786,13 +787,14 @@ def isolate_nested_sdfg( ) -> Union[Tuple[SDFGState, SDFGState, SDFGState], bool]: """Isolate the nested SDFG. - The function will split `state` into three states: - - Pre State: Contains the data flow that is needed to compute the dependency - of the nested SDFG. - - Middle State: This state contains the nested SDFG and the nodes that are needed - as input or output to the nested SDFG. - - Post State: Contains the data flow that uses the data that was computed - by the nested SDFG. + The function will split ``state`` into three states: + + - Pre State: Contains the data flow that is needed to compute the dependency + of the nested SDFG. + - Middle State: This state contains the nested SDFG and the nodes that are needed + as input or output to the nested SDFG. + - Post State: Contains the data flow that uses the data that was computed + by the nested SDFG. The important aspect of this function is, that it will not increase the number of AccessNodes that are written to, it might, however, add new AccessNodes @@ -800,8 +802,11 @@ def isolate_nested_sdfg( :param state: The state on which we operate. :param nested_sdfg: The nested SDFG node that should be isolated. - :param test_if_applicable: If `True` then do not perform the splitting, only verify - that it can be performed by returning either `True` or `False`. + :param test_if_applicable: If ``True`` then do not perform the splitting, only verify + that it can be performed by returning either ``True`` or ``False``. + :return: If ``test_if_applicable`` is ``False`` then the function will return a tuple of the three states + (pre, middle, post). If ``test_if_applicable`` is ``True``, then the function will return ``True`` + if the split can be performed and ``False`` otherwise. """ # We can only isolate the nested SDFG if it is on global scope, for example @@ -830,6 +835,7 @@ def isolate_nested_sdfg( visited.add(node_to_process) pre_nodes.add(node_to_process) to_visit.extend(iedge.src for iedge in state.in_edges(node_to_process)) + pre_nodes = list(sorted(pre_nodes, key=lambda n: n.id)) # These are the nodes of the middle state. Which are all access nodes that serves # as input to the nested SDFG and the nested SDFG itself. @@ -856,14 +862,14 @@ def isolate_nested_sdfg( "Can only split if the out to the nested SDFG are AccessNodes to non view data and the AccessNodes are only connected to the nested SDFG." ) middle_nodes.add(oedge.dst) + middle_nodes = list(sorted(middle_nodes, key=lambda n: n.id)) # These are the nodes that belongs to the Post State. There are two reasons why a # node belongs to the set of post nodes. # The first is that the node does not belong to any other set. - post_nodes: Set[nodes.Node] = { - node - for node in state.nodes() if (node not in pre_nodes) and (node not in middle_nodes) - } + post_nodes: list[nodes.Node] = [ + node for node in state.nodes() if (node not in pre_nodes) and (node not in middle_nodes) + ] # The second reason, are read dependencies, for this we have to look at the incoming # edges and add any node that we need. @@ -876,7 +882,7 @@ def isolate_nested_sdfg( if test_if_applicable: return False raise ValueError("Can not replicate non non-View AccessNodes into the post state.") - post_nodes.add(node) + post_nodes.append(node) if test_if_applicable: return True @@ -1727,20 +1733,46 @@ def make_map_internal_write_external(sdfg: SDFG, state: SDFGState, map_exit: nod def all_isedges_between(src: ControlFlowBlock, dst: ControlFlowBlock) -> Iterable[Edge[InterstateEdge]]: """ - Helper function that generates an iterable of all edges potentially encountered between two control flow blocks. + Collects all interstate edges that may lie on a path between two blocks (e.g., states, loop regions). + + The result is conservative: when paths cross nested control-flow regions, + interstate edges inside those regions that may contribute to the path are + included as well. + + :param src: The control-flow block where the path starts. + :param dst: The control-flow block where the path ends. + :return: An iterable of unique interstate edges that may be encountered on + any path from ``src`` to ``dst``. """ + + def _parent_chain(block: ControlFlowBlock) -> List[ControlFlowRegion]: + result = [block.parent_graph] + while not isinstance(result[-1], SDFG): + result.append(result[-1].parent_graph) + return result + + def _edge_key(edge: Edge[InterstateEdge]) -> Tuple[int, int, Optional[str], Tuple[Tuple[str, str], ...]]: + condition = edge.data.condition.as_string if edge.data.condition else None + assignments = tuple(sorted((str(k), str(v)) for k, v in edge.data.assignments.items())) + return (id(edge.src), id(edge.dst), condition, assignments) + + def _add_edge(edges: Dict[Tuple[int, int, Optional[str], Tuple[Tuple[str, str], ...]], Edge[InterstateEdge]], + edge: Edge[InterstateEdge]) -> None: + edges[_edge_key(edge)] = edge + if src.sdfg is not dst.sdfg: raise RuntimeError('Blocks reside in different SDFGs') if src.parent_graph is dst.parent_graph: # Simple case where both blocks reside in the same graph: - edges = set() + edges: Dict[Tuple[int, int, Optional[str], Tuple[Tuple[str, str], ...]], Edge[InterstateEdge]] = {} for p in src.parent_graph.all_simple_paths(src, dst, as_edges=True): for e in p: - edges.add(e) + _add_edge(edges, e) if isinstance(e.dst, ControlFlowRegion): - edges.update(e.dst.all_interstate_edges()) - return edges + for nested_edge in e.dst.all_interstate_edges(): + _add_edge(edges, nested_edge) + return list(edges.values()) else: # In the case where the two blocks are not in the same graph, we follow this procedure: # 1. Collect the list of control flow regions on the direct path between the source and destination: @@ -1755,40 +1787,26 @@ def all_isedges_between(src: ControlFlowBlock, dst: ControlFlowBlock) -> Iterabl # the source and destination. # Note that for each edge, if the destination is a control flow region, any edges inside of it may also # be on the path and consequently also need to be added. - edges = set() + edges: Dict[Tuple[int, int, Optional[str], Tuple[Tuple[str, str], ...]], Edge[InterstateEdge]] = {} # Step 1.a): Find the lowest common parent region. - common_regions = set() - pivot_graph = src.parent_graph - all_parent_regions_src = [pivot_graph] - while not isinstance(pivot_graph, SDFG): - pivot_graph = pivot_graph.parent_graph - all_parent_regions_src.append(pivot_graph) - pivot_graph = dst.parent_graph - all_parent_regions_dst = [pivot_graph] - while not isinstance(pivot_graph, SDFG): - pivot_graph = pivot_graph.parent_graph - all_parent_regions_dst.append(pivot_graph) - if pivot_graph in all_parent_regions_src: - common_regions.add(pivot_graph) + all_parent_regions_src = _parent_chain(src) + all_parent_regions_dst = _parent_chain(dst) + common_parent = next((region for region in all_parent_regions_src if region in all_parent_regions_dst), None) + if common_parent is None: + raise RuntimeError('No common parent found') # Step 1.b) and 1.c): Determine the list of parents involved in the path for the source and destination. involved_src: List[ControlFlowRegion] = [] involved_dst: List[ControlFlowRegion] = [] - common_parent: ControlFlowRegion = None for r in all_parent_regions_src: - if r not in common_regions: - involved_src.append(r) - else: - common_parent = r + if r is common_parent: break + involved_src.append(r) for r in all_parent_regions_dst: - if r not in common_regions: - involved_dst.append(r) - else: - if r is not common_parent: - raise RuntimeError('No common parent found') + if r is common_parent: break + involved_dst.append(r) # Step 2 src_pivot = src @@ -1799,9 +1817,10 @@ def all_isedges_between(src: ControlFlowBlock, dst: ControlFlowBlock) -> Iterabl for sink in r.sink_nodes(): for p in r.all_simple_paths(src_pivot, sink, as_edges=True): for e in p: - edges.add(e) + _add_edge(edges, e) if isinstance(e.dst, ControlFlowRegion): - edges.update(e.dst.all_interstate_edges()) + for nested_edge in e.dst.all_interstate_edges(): + _add_edge(edges, nested_edge) src_pivot = r # Step 3 dst_pivot = dst @@ -1811,19 +1830,109 @@ def all_isedges_between(src: ControlFlowBlock, dst: ControlFlowBlock) -> Iterabl continue for p in r.all_simple_paths(r.start_block, dst_pivot, as_edges=True): for e in p: - edges.add(e) + _add_edge(edges, e) if isinstance(e.dst, ControlFlowRegion): - edges.update(e.dst.all_interstate_edges()) + for nested_edge in e.dst.all_interstate_edges(): + _add_edge(edges, nested_edge) dst_pivot = r # Step 4 - for p in common_parent.all_simple_paths(src_pivot, dst_pivot, as_edges=True): + paths = common_parent.all_simple_paths(src_pivot, dst_pivot, as_edges=True) + for p in paths: for e in p: - edges.add(e) - if isinstance(e.dst, ControlFlowRegion) and not e.dst is dst_pivot: - edges.update(e.dst.all_interstate_edges()) + _add_edge(edges, e) + if isinstance(e.src, ControlFlowRegion) and e.src is not src_pivot: + for nested_edge in e.src.all_interstate_edges(): + _add_edge(edges, nested_edge) + if isinstance(e.dst, ControlFlowRegion) and e.dst is not dst_pivot: + for nested_edge in e.dst.all_interstate_edges(): + _add_edge(edges, nested_edge) + + return list(edges.values()) - return edges + +def _assigned_names_in_ast(node: ast.AST) -> Set[str]: + """ + Returns symbol names assigned by an assignment-target AST node (left-hand side of an assignment). + + :param node: The AST node representing an assignment target. + :return: A set of assigned symbol names extracted from the target. + """ + if isinstance(node, ast.Name): + return {node.id} + if isinstance(node, (ast.Tuple, ast.List)): + result = set() + for elt in node.elts: + result.update(_assigned_names_in_ast(elt)) + return result + return set() + + +def _symbols_assigned_in_codeblock(codeblock: Optional[CodeBlock]) -> Set[str]: + """ + Collects symbols assigned within a Python code block. + + Only assignment-like statements represented as Python AST nodes are + inspected. + + :param codeblock: The code block to inspect. + :return: A set of symbol names assigned within the code block. + """ + if codeblock is None or not isinstance(codeblock.code, list): + return set() + + result = set() + for stmt in codeblock.code: + if isinstance(stmt, ast.Assign): + for target in stmt.targets: + result.update(_assigned_names_in_ast(target)) + elif isinstance(stmt, ast.AnnAssign): + result.update(_assigned_names_in_ast(stmt.target)) + elif isinstance(stmt, ast.AugAssign): + result.update(_assigned_names_in_ast(stmt.target)) + return result + + +def _modified_symbols_in_region(region: ControlFlowRegion) -> Set[str]: + """ + Collects symbols modified directly by control flow region metadata. + + This currently includes loop variables and symbols assigned in loop + initializer and update code blocks. + + :param region: The control-flow region to inspect. + :return: A set of symbol names modified by the region metadata itself. + """ + result = set() + if isinstance(region, LoopRegion): + if region.loop_variable: + result.add(region.loop_variable) + result.update(_symbols_assigned_in_codeblock(region.init_statement)) + result.update(_symbols_assigned_in_codeblock(region.update_statement)) + return result + + +def modified_symbols_between(src: ControlFlowBlock, dst: ControlFlowBlock) -> Set[str]: + """ + Returns symbols that may be modified on any path between two blocks. + + In addition to assignments on interstate edges, this includes symbol + mutations expressed in loop-region metadata, such as loop initializers, + update statements, and loop variables. + + :param src: The control-flow block where the path starts. + :param dst: The control-flow block where the path ends. + :return: A conservative set of symbol names that may be modified on any + path between ``src`` and ``dst``. + """ + result = set() + for edge in all_isedges_between(src, dst): + result.update(edge.data.assignments.keys()) + if isinstance(edge.src, ControlFlowRegion): + result.update(_modified_symbols_in_region(edge.src)) + if isinstance(edge.dst, ControlFlowRegion): + result.update(_modified_symbols_in_region(edge.dst)) + return result def replace_sdfg_dtypes( @@ -1958,3 +2067,49 @@ def _is_pointer(obj) -> bool: def _is_structure_view(obj) -> bool: """Check if object is a StructureView.""" return isinstance(obj, data.StructureView) + + +def move_branch_cfg_up_discard_conditions(if_block: ConditionalBlock, body_to_take: ControlFlowRegion): + """ + Move a branch of a ``ConditionalBlock`` up into its parent CFG, replacing the + conditional with the selected branch body and discarding the condition check + and every other branch. Incoming and outgoing edges of the conditional are + rewired to the spliced branch's start and end nodes. + + :param if_block: The conditional block to dissolve. + :param body_to_take: The branch body whose nodes and edges are spliced into + ``if_block.parent_graph`` in place of ``if_block``. + """ + bodies = {b for _, b in if_block.branches} + assert body_to_take in bodies + assert isinstance(if_block, ConditionalBlock) + + graph = if_block.parent_graph + + node_map = dict() + new_start_block = None + new_end_block = None + + for node in body_to_take.nodes(): + copynode = copy.deepcopy(node) + node_map[node] = copynode + start_block_case = (body_to_take.start_block == node) and (graph.start_block == if_block) + if body_to_take.start_block == node: + assert new_start_block is None + new_start_block = copynode + if body_to_take.out_degree(node) == 0: + assert new_end_block is None + new_end_block = copynode + graph.add_node(copynode, is_start_block=start_block_case) + + for edge in body_to_take.edges(): + src = node_map[edge.src] + dst = node_map[edge.dst] + graph.add_edge(src, dst, copy.deepcopy(edge.data)) + + for ie in graph.in_edges(if_block): + graph.add_edge(ie.src, new_start_block, copy.deepcopy(ie.data)) + for oe in graph.out_edges(if_block): + graph.add_edge(new_end_block, oe.dst, copy.deepcopy(oe.data)) + + graph.remove_node(if_block) diff --git a/dace/transformation/interstate/condition_fusion.py b/dace/transformation/interstate/condition_fusion.py index e4a8d5deec..0b595dabaa 100644 --- a/dace/transformation/interstate/condition_fusion.py +++ b/dace/transformation/interstate/condition_fusion.py @@ -4,7 +4,7 @@ from dace import sdfg as sd, properties from dace.properties import CodeBlock from dace.sdfg import utils as sdutil -from dace.sdfg.state import ControlFlowRegion, ConditionalBlock, SDFGState +from dace.sdfg.state import ControlFlowRegion, ConditionalBlock from dace.transformation import transformation as xf diff --git a/dace/transformation/interstate/gpu_transform_sdfg.py b/dace/transformation/interstate/gpu_transform_sdfg.py index caaf05b483..766899319e 100644 --- a/dace/transformation/interstate/gpu_transform_sdfg.py +++ b/dace/transformation/interstate/gpu_transform_sdfg.py @@ -4,8 +4,7 @@ from dace import data, memlet, dtypes, sdfg as sd, subsets as sbs, propagate_memlets_sdfg from dace.sdfg import nodes, scope from dace.sdfg import utils as sdutil -from dace.sdfg.replace import replace_in_codeblock -from dace.sdfg.state import AbstractControlFlowRegion, ConditionalBlock, LoopRegion, SDFGState +from dace.sdfg.state import AbstractControlFlowRegion, SDFGState from dace.transformation import transformation, helpers as xfh from dace.properties import ListProperty, Property, make_properties from collections import defaultdict diff --git a/dace/transformation/interstate/loop_to_map.py b/dace/transformation/interstate/loop_to_map.py index d0815ff23c..7ee9584843 100644 --- a/dace/transformation/interstate/loop_to_map.py +++ b/dace/transformation/interstate/loop_to_map.py @@ -157,8 +157,16 @@ def can_be_applied(self, graph, expr_index, sdfg, permissive=False): if dn.data in write_set: for e in state.in_edges(dn): if e.data.dynamic and e.data.wcr is None: - # If pointers are involved, give up - return False + # A dynamic write (no WCR) is still safe across + # outer-loop iterations if its destination subset + # pins an axis to the iteration variable (same + # ``a*i+b`` pattern enforced below for non-dynamic + # writes). Each iteration then writes to a disjoint + # slab -- whether a given lane fires or not + # cannot race with another iteration's write. + dst_subset = e.data.get_dst_subset(e, state) + if not (dst_subset and _check_range(dst_subset, a, itersym, b, step)): + return False if e.data is None: continue @@ -484,6 +492,30 @@ def apply(self, graph: ControlFlowRegion, sdfg: sd.SDFG): for sym, dtype in nsymbols.items(): nsdfg.symbols[sym] = dtype + # Symbols that the nested SDFG assigns on its own interstate edges + # are internal -- they must not be surfaced onto the NestedSDFG + # node's ``symbol_mapping``. Doing so would make the outer SDFG + # appear to need them as free symbols; a later pruning pass that + # removes them from ``sdfg.symbols`` would then desync the mapping + # (codegen reads ``sdfg.symbols[sym]`` and raises ``KeyError``). + internally_defined = set() + for e in nsdfg.all_interstate_edges(): + internally_defined.update(e.data.assignments.keys()) + + # Propagate free symbols referenced by nested array shapes/strides/offsets: + # ``copy.deepcopy`` of the descriptor carries the symbols, but they + # must be added to the NestedSDFG's symbol mapping. + for desc in nsdfg.arrays.values(): + for sym in desc.free_symbols: + sym_name = str(sym) + if sym_name in internally_defined: + continue + if sym_name in sdfg.symbols: + if sym_name not in nsdfg.symbols: + nsdfg.symbols[sym_name] = sdfg.symbols[sym_name] + if sym_name not in cnode.symbol_mapping: + cnode.symbol_mapping[sym_name] = symbolic.pystr_to_symbolic(sym_name) + # Propagate symbols, where types cannot be inferred alltypes = copy.deepcopy(nsdfg.symbols) alltypes.update({k: v.dtype for k, v in nsdfg.arrays.items()}) @@ -656,8 +688,12 @@ def apply(self, graph: ControlFlowRegion, sdfg: sd.SDFG): # Delete the loop and connected edges. graph.remove_node(self.loop) - # If this had made a variable a free symbol, we can remove it from the SDFG symbols + # If this had made a variable a free symbol, we can remove it from the SDFG symbols. + # Guard both branches with ``in sdfg.symbols`` -- the array-descriptor-symbol + # propagation above may have already cleared entries that were also free. for var in sdfg.free_symbols - fsymbols: + if var not in sdfg.symbols: + continue if sdfg.parent_nsdfg_node: if var not in sdfg.parent_nsdfg_node.symbol_mapping: sdfg.remove_symbol(var) diff --git a/dace/transformation/interstate/loop_unroll.py b/dace/transformation/interstate/loop_unroll.py index ebf9c911e6..af5358d0fb 100644 --- a/dace/transformation/interstate/loop_unroll.py +++ b/dace/transformation/interstate/loop_unroll.py @@ -1,11 +1,11 @@ -# Copyright 2019-2025 ETH Zurich and the DaCe authors. All rights reserved. +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. """ Loop unroll transformation """ import ast import copy from typing import List, Optional, Union -from dace import sdfg as sd, symbolic, serialize +from dace import sdfg as sd, symbolic, serialize, version from dace.properties import Property, make_properties from dace.sdfg import InterstateEdge, utils as sdutil from dace.sdfg.nodes import NestedSDFG @@ -79,6 +79,8 @@ def apply(self, graph: ControlFlowRegion, sdfg): is_symbolic |= symbolic.issymbolic(current_index) iteration_region = self.instantiate_loop_iteration(graph, self.loop, current_index, str(i) if is_symbolic else None) + iteration_region.replace_dict({self.loop.loop_variable: current_index}, replace_keys=True) + iteration_region.replace_meta_accesses({self.loop.loop_variable: str(current_index)}) # Connect iterations with unconditional edges if len(unrolled_iterations) > 0: @@ -103,9 +105,9 @@ def apply(self, graph: ControlFlowRegion, sdfg): # If we remove start block we need to update the new start-block was_start_block = graph.start_block == self.loop + oes = graph.out_edges(self.loop) graph.remove_node(self.loop) if was_start_block: - oes = graph.out_edges(self.loop) if len(oes) > 0: oe = oes[0] oes2 = graph.out_edges(oe.dst) @@ -138,11 +140,16 @@ def instantiate_loop_iteration(self, block_map = {} for block in loop.nodes(): - # Using to/from JSON copies faster than deepcopy. - new_block = serialize.from_json(serialize.to_json(block), context={'sdfg': graph.sdfg}) + # Using to/from JSON is faster for copying blocks than deep copying. + new_block = serialize.from_json(serialize.to_json(block), + context={ + 'sdfg': graph.sdfg, + 'version': version.__version__, + }) assert block not in block_map block_map[block] = new_block - new_block.replace(loop.loop_variable, value) + # The JSON copy is created with SDFG context, so replacement can run before insertion. + new_block.replace_dict({loop.loop_variable: value}) iteration_region.add_node(new_block, is_start_block=(block is loop.start_block)) for edge in loop.edges(): diff --git a/dace/transformation/interstate/move_assignment_outside_if.py b/dace/transformation/interstate/move_assignment_outside_if.py index 10c6da573e..e2d95e350b 100644 --- a/dace/transformation/interstate/move_assignment_outside_if.py +++ b/dace/transformation/interstate/move_assignment_outside_if.py @@ -9,7 +9,7 @@ import sympy as sp from dace import sdfg as sd -from dace.sdfg import graph as gr, utils as sdutil, nodes as nd +from dace.sdfg import utils as sdutil, nodes as nd from dace.sdfg.state import ConditionalBlock, ControlFlowRegion from dace.symbolic import pystr_to_symbolic from dace.transformation import transformation diff --git a/dace/transformation/interstate/multistate_inline.py b/dace/transformation/interstate/multistate_inline.py index c6401dbb55..c8d86ef368 100644 --- a/dace/transformation/interstate/multistate_inline.py +++ b/dace/transformation/interstate/multistate_inline.py @@ -5,16 +5,16 @@ import itertools from typing import Dict, List -from dace import Memlet, symbolic, dtypes, subsets +from dace import Memlet, symbolic, subsets from dace.sdfg import nodes from dace.sdfg.graph import MultiConnectorEdge from dace.sdfg import InterstateEdge, SDFG, SDFGState -from dace.sdfg import utils as sdutil, infer_types +from dace.sdfg import utils as sdutil from dace.sdfg.replace import replace_datadesc_names, replace_properties_dict from dace.transformation import transformation, helpers from dace.properties import make_properties from dace import data -from dace.sdfg.state import LoopRegion, ReturnBlock, StateSubgraphView +from dace.sdfg.state import LoopRegion, ReturnBlock @make_properties diff --git a/dace/transformation/interstate/sdfg_nesting.py b/dace/transformation/interstate/sdfg_nesting.py index 495e7ff4e0..b2dd8bbbd8 100644 --- a/dace/transformation/interstate/sdfg_nesting.py +++ b/dace/transformation/interstate/sdfg_nesting.py @@ -15,7 +15,7 @@ from dace.sdfg import nodes, propagation, utils from dace.sdfg.graph import MultiConnectorEdge, SubgraphView from dace.sdfg import SDFG, SDFGState -from dace.sdfg import utils as sdutil, infer_types, propagation +from dace.sdfg import utils as sdutil, propagation from dace.sdfg.state import LoopRegion from dace.transformation import transformation, helpers from dace.properties import make_properties, Property @@ -154,12 +154,31 @@ def can_be_applied(self, graph: SDFGState, expr_index, sdfg, permissive=False): rem_inpconns = dc(in_connectors) rem_outconns = dc(out_connectors) nstate: SDFGState = nested_sdfg.sdfg.nodes()[0] + # Verify that all connectors have a corresponding access node that is valid + valid_inpconns = set() + valid_outconns = set() for node in nstate.nodes(): if isinstance(node, nodes.AccessNode): + if node.data not in valid_inpconns and node.data in in_connectors: + if nstate.entry_node(node) is not None: + pass # Invalid matching node, in scope + elif nstate.in_degree(node) > 0: + pass # Invalid matching node, has incoming edges + else: + valid_inpconns.add(node.data) + if node.data not in valid_outconns and node.data in out_connectors: + if nstate.entry_node(node) is not None: + pass # Invalid matching node, in scope + elif nstate.out_degree(node) > 0: + pass # Invalid matching node, has outgoing edges + else: + valid_outconns.add(node.data) if node.data in rem_inpconns: rem_inpconns.remove(node.data) if node.data in rem_outconns: rem_outconns.remove(node.data) + if len(valid_inpconns) != len(in_connectors) or len(valid_outconns) != len(out_connectors): + return False if len(rem_outconns) > 0: # Check if remaining outputs would disconnect anything or can be pruned for conn in rem_outconns: diff --git a/dace/transformation/interstate/state_elimination.py b/dace/transformation/interstate/state_elimination.py index 94619576bf..e41abb4af0 100644 --- a/dace/transformation/interstate/state_elimination.py +++ b/dace/transformation/interstate/state_elimination.py @@ -100,7 +100,7 @@ def can_be_applied(self, graph, expr_index, sdfg, permissive=False): return False # Assignments that make descriptors into symbols cannot be eliminated for assign in edge.data.assignments.values(): - if graph.arrays.keys() & symbolic.free_symbols_and_functions(assign): + if graph.arrays.keys() & (symbolic.free_symbols_and_functions(assign) | symbolic.arrays(assign)): return False return True diff --git a/dace/transformation/interstate/state_fusion.py b/dace/transformation/interstate/state_fusion.py index 56e1e64449..930f12aef8 100644 --- a/dace/transformation/interstate/state_fusion.py +++ b/dace/transformation/interstate/state_fusion.py @@ -1,11 +1,11 @@ -# Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. """ State fusion transformation """ from typing import Dict, List, Set import networkx as nx -from dace import data as dt, dtypes, registry, sdfg, subsets +from dace import data as dt, sdfg, subsets from dace.config import Config from dace.sdfg import nodes from dace.sdfg import utils as sdutil @@ -160,7 +160,7 @@ def _check_paths(self, first_state: SDFGState, second_state: SDFGState, match_no break # We keep looking for a potential match with a path that fail to find # a path to the second state to make sure we test memlet_intersections - # independant of the order of the access nodes in the lists + # independent of the order of the access nodes in the lists if fail: break diff --git a/dace/transformation/interstate/state_fusion_with_happens_before.py b/dace/transformation/interstate/state_fusion_with_happens_before.py index 17f9c287ee..3b4b0f5a3d 100644 --- a/dace/transformation/interstate/state_fusion_with_happens_before.py +++ b/dace/transformation/interstate/state_fusion_with_happens_before.py @@ -1,4 +1,4 @@ -# Copyright 2019-2024 ETH Zurich and the DaCe authors. All rights reserved. +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. """ State fusion transformation """ from typing import Dict, List, Set @@ -44,7 +44,6 @@ class StateFusionExtended(transformation.MultiStateTransformation): and fuses them into one state. If permissive, also applies if potential memory access hazards are created. """ - connections_to_make = [] first_state = transformation.PatternNode(sdfg.SDFGState) second_state = transformation.PatternNode(sdfg.SDFGState) @@ -56,6 +55,10 @@ def annotates_memlets(): def expressions(cls): return [sdutil.node_path_graph(cls.first_state, cls.second_state)] + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + self.connections_to_make = [] + @staticmethod def find_fused_components(first_cc_input, first_cc_output, second_cc_input, second_cc_output) -> List[CCDesc]: # Make a bipartite graph out of the first and second components @@ -176,6 +179,8 @@ def _check_paths(self, first_state: SDFGState, second_state: SDFGState, match_no def can_be_applied(self, graph, expr_index, sdfg, permissive=False): first_state: SDFGState = self.first_state second_state: SDFGState = self.second_state + # We keep it alive, such that `apply()` can use it later. + self.connections_to_make.clear() out_edges = graph.out_edges(first_state) in_edges = graph.in_edges(first_state) @@ -376,6 +381,17 @@ def can_be_applied(self, graph, expr_index, sdfg, permissive=False): nodes_second, True, False): return False + # Same-cc write-after-read: first reads ``d``, + # second writes it, ``d`` not a first output. + # Safe fusion would need a dependency edge to + # every first-state sink reading ``d``; refuse + # instead and keep the two states separate. + if d in fused_cc.first_inputs: + nodes_first_read = [n for n in first_input if n.data == d] + if StateFusionExtended.memlets_intersect(first_state, nodes_first_read, True, + second_state, nodes_second, False): + return False + continue # If an input/output of a connected component in the first # state is an output of another connected component in the @@ -385,13 +401,17 @@ def can_be_applied(self, graph, expr_index, sdfg, permissive=False): if d in other_cc.second_outputs: # Check for intersection (if None, fusion is ok) nodes_second = [n for n in second_output if n.data == d] - # Read-Write race + # Read-Write race (write-after-read / anti- + # dependency): the first state reads ``d`` and the + # second writes it. Safely fusing would require a + # dependency edge to every first-state sink reading + # ``d`` (arbitrarily fanned out); refuse instead and + # keep the two states separate. if d in fused_cc.first_inputs: nodes_first = [n for n in first_input if n.data == d] if StateFusionExtended.memlets_intersect(first_state, nodes_first, True, second_state, nodes_second, False): - self.connections_to_make.append([nodes_first, nodes_second]) - #return False + return False # Write-Write race if d in fused_cc.first_outputs: nodes_first = [n for n in first_output if n.data == d] @@ -461,6 +481,9 @@ def apply(self, graph: ControlFlowRegion, sdfg): first_state: SDFGState = self.first_state second_state: SDFGState = self.second_state + # This will populate `self.connections_to_make`. + self.can_be_applied(graph, 0, sdfg) + # Remove interstate edge(s) edges = graph.edges_between(first_state, second_state) for edge in edges: @@ -584,3 +607,6 @@ def apply(self, graph: ControlFlowRegion, sdfg): graph.remove_node(second_state) if graph.start_block == second_state: graph.start_block = graph.node_id(first_state) + + # Technically unneeded, but better to keep track. + self.connections_to_make.clear() diff --git a/dace/transformation/onnx/replacement.py b/dace/transformation/onnx/replacement.py index ce51fdd14c..e6aa74bdd2 100644 --- a/dace/transformation/onnx/replacement.py +++ b/dace/transformation/onnx/replacement.py @@ -2,11 +2,9 @@ """ General class for pattern replacement transformations. """ import abc import dace -from dace import registry, nodes, data as dt +from dace import nodes, data as dt from dace.transformation import transformation, helpers as xfh from typing import Any, Dict, List, Optional, Tuple, Union -from dace.sdfg import utils as sdutil -from dace.libraries.onnx import nodes as onnx_op from dace.sdfg import graph as gr diff --git a/dace/transformation/optimizer.py b/dace/transformation/optimizer.py index 1e77971664..5614df60d6 100644 --- a/dace/transformation/optimizer.py +++ b/dace/transformation/optimizer.py @@ -5,7 +5,6 @@ import copy import os import re -import time from typing import Any, Dict, Iterator, List, Optional, Type import dace @@ -16,7 +15,6 @@ from dace.transformation.transformation import PatternTransformation # This import is necessary since it registers all the patterns -from dace.transformation import dataflow, interstate, subgraph class Optimizer(object): diff --git a/dace/transformation/pass_pipeline.py b/dace/transformation/pass_pipeline.py index f0d602c42c..415dba2a57 100644 --- a/dace/transformation/pass_pipeline.py +++ b/dace/transformation/pass_pipeline.py @@ -2,17 +2,30 @@ """ API for SDFG analysis and manipulation Passes, as well as Pipelines that contain multiple dependent passes. """ -import warnings + from dace import properties, serialize from dace.sdfg import SDFG, SDFGState, graph as gr, nodes, utils as sdutil from enum import Flag, auto -from typing import Any, Dict, Iterator, List, Optional, Set, Type, Union +from typing import Any, Dict, Iterable, Iterator, List, Optional, Set, Type, Union from dataclasses import dataclass from dace.sdfg.state import ConditionalBlock, ControlFlowRegion +def unique_dependencies(passes: Iterable['Pass']) -> List[Union[Type['Pass'], 'Pass']]: + """ + Collects the dependencies of the given passes, preserving their listed order and removing duplicates. + + :param passes: An iterable of passes whose ``depends_on`` results should be combined. + :return: A list of the combined dependencies, in order and without duplicates. + """ + deps = [] + for p in passes: + deps.extend(p.depends_on()) + return list(dict.fromkeys(deps)) # Remove duplicates while preserving order + + class Modifies(Flag): """ Specifies which elements of an SDFG have been modified by a Pass/Transformation. @@ -56,13 +69,13 @@ class Pass: CATEGORY: str = 'Helper' - def depends_on(self) -> Set[Union[Type['Pass'], 'Pass']]: + def depends_on(self) -> List[Union[Type['Pass'], 'Pass']]: """ If in the context of a ``Pipeline``, which other Passes need to run first. - :return: A set of Pass subclasses or objects that need to run prior to this Pass. + :return: A list of Pass subclasses or objects that need to run prior to this Pass. """ - return set() + return [] def modifies(self) -> Modifies: """ @@ -413,7 +426,7 @@ class Pipeline(Pass): def __init__(self, passes: List[Pass]): self.passes = [] - self._pass_names = set(type(p).__name__ for p in passes) + self._pass_names = set(type(p).__name__ for p in passes) # todo sort this? self.passes.extend(passes) # Add missing Pass dependencies @@ -483,11 +496,8 @@ def modifies(self) -> Modifies: def should_reapply(self, modified: Modifies) -> bool: return any(p.should_reapply(modified) for p in self.passes) - def depends_on(self) -> Set[Type[Pass]]: - result = set() - for p in self.passes: - result.update(p.depends_on()) - return result + def depends_on(self) -> List[Union[Type[Pass], Pass]]: + return unique_dependencies(self.passes) def _make_dependency_graph(self) -> gr.OrderedDiGraph: """ diff --git a/dace/transformation/passes/analysis/analysis.py b/dace/transformation/passes/analysis/analysis.py index e9827f7a4d..6518673591 100644 --- a/dace/transformation/passes/analysis/analysis.py +++ b/dace/transformation/passes/analysis/analysis.py @@ -16,6 +16,7 @@ from typing import Dict, Iterable, List, Set, Tuple, Any, Optional, Union import networkx as nx from networkx.algorithms import shortest_paths as nxsp +from ordered_set import OrderedSet from dace.transformation.passes.analysis import loop_analysis @@ -41,9 +42,9 @@ def should_reapply(self, modified: ppl.Modifies) -> bool: return modified & ppl.Modifies.CFG def depends_on(self): - return {ControlFlowBlockReachability} + return [ControlFlowBlockReachability] - def apply_pass(self, top_sdfg: SDFG, pipeline_res: Dict) -> Dict[int, Dict[SDFGState, Set[SDFGState]]]: + def apply_pass(self, top_sdfg: SDFG, pipeline_res: Dict) -> Dict[int, Dict[SDFGState, OrderedSet[SDFGState]]]: """ :return: A dictionary mapping each state to its other reachable states. """ @@ -52,9 +53,9 @@ def apply_pass(self, top_sdfg: SDFG, pipeline_res: Dict) -> Dict[int, Dict[SDFGS cf_block_reach_dict = ControlFlowBlockReachability().apply_pass(top_sdfg, {}) else: cf_block_reach_dict = pipeline_res[ControlFlowBlockReachability.__name__] - reachable: Dict[int, Dict[SDFGState, Set[SDFGState]]] = {} + reachable: Dict[int, Dict[SDFGState, OrderedSet[SDFGState]]] = {} for sdfg in top_sdfg.all_sdfgs_recursive(): - result: Dict[SDFGState, Set[SDFGState]] = defaultdict(set) + result: Dict[SDFGState, OrderedSet[SDFGState]] = defaultdict(OrderedSet) for state in sdfg.states(): for reached in cf_block_reach_dict[state.parent_graph.cfg_id][state]: if isinstance(reached, SDFGState): @@ -85,9 +86,13 @@ def modifies(self) -> ppl.Modifies: def should_reapply(self, modified: ppl.Modifies) -> bool: return modified & ppl.Modifies.CFG - def _region_closure(self, region: ControlFlowRegion, - block_reach: Dict[int, Dict[ControlFlowBlock, Set[ControlFlowBlock]]]) -> Set[SDFGState]: - closure: Set[SDFGState] = set() + def _region_closure( + self, + region: ControlFlowRegion, + block_reach: Dict[int, Dict[ControlFlowBlock, OrderedSet[ControlFlowBlock]]], + cached_closures: dict[int, OrderedSet[ControlFlowBlock]], + ) -> Set[ControlFlowBlock]: + closure: Set[SDFGState] = OrderedSet() if isinstance(region, LoopRegion): # Any point inside the loop may reach any other point inside the loop again. # TODO(later): This is an overapproximation. A branch terminating in a break is excluded from this. @@ -103,23 +108,28 @@ def _region_closure(self, region: ControlFlowRegion, # Walk up the parent tree. pivot = region.parent_graph while pivot and not isinstance(pivot, SDFG): - closure.update(self._region_closure(pivot, block_reach)) + graph_id = id(pivot) + if graph_id not in cached_closures: + cached_closures[graph_id] = self._region_closure(pivot, block_reach, cached_closures) + closure.update(cached_closures[graph_id]) pivot = pivot.parent_graph return closure - def apply_pass(self, top_sdfg: SDFG, _) -> Dict[int, Dict[ControlFlowBlock, Set[ControlFlowBlock]]]: + def apply_pass(self, top_sdfg: SDFG, _) -> Dict[int, Dict[ControlFlowBlock, OrderedSet[ControlFlowBlock]]]: """ :return: For each control flow region, a dictionary mapping each control flow block to its other reachable control flow blocks. """ + top_sdfg.reset_cfg_list() + single_level_reachable: Dict[int, Dict[ControlFlowBlock, - Set[ControlFlowBlock]]] = defaultdict(lambda: defaultdict(set)) + OrderedSet[ControlFlowBlock]]] = defaultdict(lambda: defaultdict(set)) for cfg in top_sdfg.all_control_flow_regions(recursive=True): # In networkx this is currently implemented naively for directed graphs. # The implementation below is faster # tc: nx.DiGraph = nx.transitive_closure(sdfg.nx) for n, v in reachable_nodes(cfg.nx): - reach = set() + reach = OrderedSet() for nd in v: reach.add(nd) if isinstance(nd, AbstractControlFlowRegion): @@ -131,17 +141,22 @@ def apply_pass(self, top_sdfg: SDFG, _) -> Dict[int, Dict[ControlFlowBlock, Set[ if self.contain_to_single_level: return single_level_reachable - reachable: Dict[int, Dict[ControlFlowBlock, Set[ControlFlowBlock]]] = {} + reachable: Dict[int, Dict[ControlFlowBlock, OrderedSet[ControlFlowBlock]]] = {} + cached_closures: dict[int, OrderedSet[ControlFlowBlock]] = {} for sdfg in top_sdfg.all_sdfgs_recursive(): for cfg in sdfg.all_control_flow_regions(): - result: Dict[ControlFlowBlock, Set[ControlFlowBlock]] = defaultdict(set) + result: Dict[ControlFlowBlock, OrderedSet[ControlFlowBlock]] = defaultdict(OrderedSet) for block in cfg.nodes(): for reached in single_level_reachable[block.parent_graph.cfg_id][block]: if isinstance(reached, AbstractControlFlowRegion): result[block].update(reached.all_control_flow_blocks()) result[block].add(reached) if block.parent_graph is not sdfg: - result[block].update(self._region_closure(block.parent_graph, single_level_reachable)) + graph_id = id(block.parent_graph) + if graph_id not in cached_closures: + cached_closures[graph_id] = self._region_closure(block.parent_graph, single_level_reachable, + cached_closures) + result[block].update(cached_closures[graph_id]) reachable[cfg.cfg_id] = result return reachable @@ -165,11 +180,11 @@ def _single_shortest_path_length_no_self(adj, source): seen = {} # level (number of hops) when seen in BFS level = 0 # the current level - nextlevel = set(firstlevel) # set of nodes to check at next level + nextlevel = OrderedSet(firstlevel) # set of nodes to check at next level n = len(adj) while nextlevel: thislevel = nextlevel # advance to next level - nextlevel = set() # and start a new set (fringe) + nextlevel = OrderedSet() # and start a new set (fringe) found = [] for v in thislevel: if v not in seen: @@ -211,9 +226,9 @@ def should_reapply(self, modified: ppl.Modifies) -> bool: return modified & ppl.Modifies.States | ppl.Modifies.Edges | ppl.Modifies.Symbols | ppl.Modifies.Nodes def apply(self, region: ControlFlowRegion, - _) -> Dict[Union[ControlFlowBlock, Edge[InterstateEdge]], Tuple[Set[str], Set[str]]]: + _) -> Dict[Union[ControlFlowBlock, Edge[InterstateEdge]], Tuple[OrderedSet[str], OrderedSet[str]]]: adesc = set(region.sdfg.arrays.keys()) - result: Dict[ControlFlowBlock, Tuple[Set[str], Set[str]]] = {} + result: Dict[ControlFlowBlock, Tuple[OrderedSet[str], OrderedSet[str]]] = {} for block in region.nodes(): # No symbols may be written to inside blocks. result[block] = (block.free_symbols, set()) @@ -240,7 +255,7 @@ def should_reapply(self, modified: ppl.Modifies) -> bool: # If access nodes were modified, reapply return modified & ppl.Modifies.AccessNodes - def _get_loop_region_readset(self, loop: LoopRegion, arrays: Set[str]) -> Set[str]: + def _get_loop_region_readset(self, loop: LoopRegion, arrays: OrderedSet[str]) -> OrderedSet[str]: readset = set() exprs = {loop.loop_condition.as_string} update_stmt = loop_analysis.get_update_assignment(loop) @@ -250,18 +265,18 @@ def _get_loop_region_readset(self, loop: LoopRegion, arrays: Set[str]) -> Set[st if init_stmt: exprs.add(init_stmt) for expr in exprs: - readset |= symbolic.free_symbols_and_functions(expr) & arrays + readset |= (symbolic.free_symbols_and_functions(expr) | symbolic.arrays(expr)) & arrays return readset - def apply_pass(self, top_sdfg: SDFG, _) -> Dict[ControlFlowBlock, Tuple[Set[str], Set[str]]]: + def apply_pass(self, top_sdfg: SDFG, _) -> Dict[ControlFlowBlock, Tuple[OrderedSet[str], OrderedSet[str]]]: """ :return: A dictionary mapping each control flow block to a tuple of its (read, written) data descriptors. """ - result: Dict[ControlFlowBlock, Tuple[Set[str], Set[str]]] = {} + result: Dict[ControlFlowBlock, Tuple[OrderedSet[str], OrderedSet[str]]] = {} for sdfg in top_sdfg.all_sdfgs_recursive(): - arrays: Set[str] = set(sdfg.arrays.keys()) + arrays: OrderedSet[str] = OrderedSet(sdfg.arrays.keys()) for block in sdfg.all_control_flow_blocks(): - readset, writeset = set(), set() + readset, writeset = OrderedSet(), OrderedSet() if isinstance(block, SDFGState): for anode in block.data_nodes(): if block.in_degree(anode) > 0: @@ -280,14 +295,15 @@ def apply_pass(self, top_sdfg: SDFG, _) -> Dict[ControlFlowBlock, Tuple[Set[str] elif isinstance(block, ConditionalBlock): for cond, _ in block.branches: if cond is not None: - readset |= symbolic.free_symbols_and_functions(cond.as_string) & arrays + readset |= (symbolic.free_symbols_and_functions(cond.as_string) + | symbolic.arrays(cond.as_string)) & arrays result[block] = (readset, writeset) # Edges that read from arrays add to both ends' access sets anames = sdfg.arrays.keys() for e in sdfg.all_interstate_edges(): - fsyms = e.data.free_symbols & anames + fsyms = sorted(e.data.free_symbols & anames) if fsyms: result[e.src][0].update(fsyms) result[e.dst][0].update(fsyms) @@ -310,14 +326,14 @@ def should_reapply(self, modified: ppl.Modifies) -> bool: # If anything was modified, reapply return modified & ppl.Modifies.AccessNodes - def apply_pass(self, top_sdfg: SDFG, _) -> Dict[int, Dict[str, Set[SDFGState]]]: + def apply_pass(self, top_sdfg: SDFG, _) -> Dict[int, Dict[str, OrderedSet[SDFGState]]]: """ :return: A dictionary mapping each data descriptor name to states where it can be found in. """ - top_result: Dict[int, Dict[str, Set[SDFGState]]] = {} + top_result: Dict[int, Dict[str, OrderedSet[SDFGState]]] = {} for sdfg in top_sdfg.all_sdfgs_recursive(): - result: Dict[str, Set[SDFGState]] = defaultdict(set) + result: Dict[str, OrderedSet[SDFGState]] = defaultdict(OrderedSet) for state in sdfg.states(): for anode in state.data_nodes(): result[anode.data].add(state) @@ -325,9 +341,9 @@ def apply_pass(self, top_sdfg: SDFG, _) -> Dict[int, Dict[str, Set[SDFGState]]]: # Edges that read from arrays add to both ends' access sets anames = sdfg.arrays.keys() for e in sdfg.all_interstate_edges(): - fsyms = e.data.free_symbols & anames + fsyms = sorted(e.data.free_symbols & anames) for access in fsyms: - result[access].update({e.src, e.dst}) + result[access].update((e.src, e.dst)) top_result[sdfg.cfg_id] = result return top_result @@ -359,18 +375,18 @@ def should_reapply(self, modified: ppl.Modifies) -> bool: # If anything was modified, reapply return modified & ppl.Modifies.AccessNodes & ppl.Modifies.CFG - def apply_pass(self, sdfg: SDFG, _) -> Dict[SDFG, Set[str]]: + def apply_pass(self, sdfg: SDFG, _) -> Dict[SDFG, OrderedSet[str]]: """ :return: A dictionary mapping SDFGs to a `set` of strings containing the name of the data descriptors that are only used once. """ # TODO(pschaad): Should we index on cfg or the SDFG itself. - exclusive_data: Dict[SDFG, Set[str]] = {} + exclusive_data: Dict[SDFG, OrderedSet[str]] = {} for nsdfg in sdfg.all_sdfgs_recursive(): exclusive_data[nsdfg] = self._find_single_use_data_in_sdfg(nsdfg) return exclusive_data - def _find_single_use_data_in_sdfg(self, sdfg: SDFG) -> Set[str]: + def _find_single_use_data_in_sdfg(self, sdfg: SDFG) -> OrderedSet[str]: """Scans an SDFG and computes the data that is only used once in the SDFG. The rules used to classify data descriptors are outlined above. The function @@ -381,8 +397,8 @@ def _find_single_use_data_in_sdfg(self, sdfg: SDFG) -> Set[str]: # If we encounter a data descriptor for the first time we immediately # classify it as single use. We will undo this decision as soon as # learn that it is used somewhere else. - single_use_data: Set[str] = set() - previously_seen: Set[str] = set() + single_use_data: OrderedSet[str] = OrderedSet() + previously_seen: OrderedSet[str] = OrderedSet() for state in sdfg.states(): for dnode in state.data_nodes(): @@ -422,17 +438,19 @@ def modifies(self) -> ppl.Modifies: def should_reapply(self, modified: ppl.Modifies) -> bool: return modified & ppl.Modifies.AccessNodes - def apply_pass(self, top_sdfg: SDFG, - _) -> Dict[int, Dict[str, Dict[SDFGState, Tuple[Set[nd.AccessNode], Set[nd.AccessNode]]]]]: + def apply_pass( + self, top_sdfg: SDFG, + _) -> Dict[int, Dict[str, Dict[SDFGState, Tuple[OrderedSet[nd.AccessNode], OrderedSet[nd.AccessNode]]]]]: """ :return: A dictionary mapping each data descriptor name to a dictionary keyed by states with all access nodes that use that data descriptor. """ - top_result: Dict[int, Dict[str, Set[nd.AccessNode]]] = dict() + top_result: Dict[int, Dict[str, OrderedSet[nd.AccessNode]]] = dict() for sdfg in top_sdfg.all_sdfgs_recursive(): - result: Dict[str, Dict[SDFGState, Tuple[Set[nd.AccessNode], Set[nd.AccessNode]]]] = defaultdict( - lambda: defaultdict(lambda: [set(), set()])) + result: Dict[str, Dict[SDFGState, + Tuple[OrderedSet[nd.AccessNode], OrderedSet[nd.AccessNode]]]] = defaultdict( + lambda: defaultdict(lambda: [OrderedSet(), OrderedSet()])) for state in sdfg.states(): for anode in state.data_nodes(): if state.in_degree(anode) > 0: @@ -460,7 +478,7 @@ def should_reapply(self, modified: ppl.Modifies) -> bool: return modified & ppl.Modifies.Symbols | ppl.Modifies.CFG | ppl.Modifies.Edges | ppl.Modifies.Nodes def depends_on(self): - return {SymbolAccessSets, ControlFlowBlockReachability} + return [SymbolAccessSets, ControlFlowBlockReachability] def _find_dominating_write(self, sym: str, read: Union[ControlFlowBlock, Edge[InterstateEdge]], block_idom: Dict[ControlFlowBlock, ControlFlowBlock]) -> Optional[Edge[InterstateEdge]]: @@ -495,15 +513,16 @@ def _find_dominating_write(self, sym: str, read: Union[ControlFlowBlock, Edge[In return write_isedge def apply(self, region, pipeline_results) -> SymbolScopeDict: - result: SymbolScopeDict = defaultdict(lambda: defaultdict(lambda: set())) + result: SymbolScopeDict = defaultdict(lambda: defaultdict(lambda: OrderedSet())) idom = nx.immediate_dominators(region.nx, region.start_block) all_doms = cfg_analysis.all_dominators(region, idom) - b_reach: Dict[ControlFlowBlock, - Set[ControlFlowBlock]] = pipeline_results[ControlFlowBlockReachability.__name__][region.cfg_id] + b_reach: Dict[ControlFlowBlock, OrderedSet[ControlFlowBlock]] = pipeline_results[ + ControlFlowBlockReachability.__name__][region.cfg_id] symbol_access_sets: Dict[Union[ControlFlowBlock, Edge[InterstateEdge]], - Tuple[Set[str], Set[str]]] = pipeline_results[SymbolAccessSets.__name__][region.cfg_id] + Tuple[OrderedSet[str], + OrderedSet[str]]] = pipeline_results[SymbolAccessSets.__name__][region.cfg_id] for read_loc, (reads, _) in symbol_access_sets.items(): for sym in reads: @@ -537,7 +556,7 @@ def apply(self, region, pipeline_results) -> SymbolScopeDict: other_accesses.update(accesses) other_accesses.add(write) to_remove.add((sym, write)) - result[sym][write] = set() + result[sym][write] = OrderedSet() for sym, write in to_remove: del result[sym][write] @@ -562,15 +581,16 @@ def should_reapply(self, modified: ppl.Modifies) -> bool: return modified & ppl.Modifies.States def depends_on(self): - return {AccessSets, FindAccessNodes, ControlFlowBlockReachability} + return [AccessSets, FindAccessNodes, ControlFlowBlockReachability] def _find_dominating_write(self, desc: str, block: ControlFlowBlock, read: Union[nd.AccessNode, InterstateEdge], - access_nodes: Dict[SDFGState, Tuple[Set[nd.AccessNode], Set[nd.AccessNode]]], + access_nodes: Dict[SDFGState, Tuple[OrderedSet[nd.AccessNode], + OrderedSet[nd.AccessNode]]], idom_dict: Dict[ControlFlowRegion, Dict[ControlFlowBlock, ControlFlowBlock]], - access_sets: Dict[ControlFlowBlock, Tuple[Set[str], Set[str]]], + access_sets: Dict[ControlFlowBlock, Tuple[OrderedSet[str], OrderedSet[str]]], no_self_shadowing: bool = False) -> Optional[Tuple[SDFGState, nd.AccessNode]]: if isinstance(read, nd.AccessNode): state: SDFGState = block @@ -637,16 +657,18 @@ def apply_pass(self, top_sdfg: SDFG, pipeline_results: Dict[str, Any]) -> Dict[i """ top_result: Dict[int, WriteScopeDict] = dict() - access_sets: Dict[ControlFlowBlock, Tuple[Set[str], Set[str]]] = pipeline_results[AccessSets.__name__] + access_sets: Dict[ControlFlowBlock, Tuple[OrderedSet[str], + OrderedSet[str]]] = pipeline_results[AccessSets.__name__] for sdfg in top_sdfg.all_sdfgs_recursive(): - result: WriteScopeDict = defaultdict(lambda: defaultdict(lambda: set())) + result: WriteScopeDict = defaultdict(lambda: defaultdict(lambda: OrderedSet())) idom_dict: Dict[ControlFlowRegion, Dict[ControlFlowBlock, ControlFlowBlock]] = {} - all_doms_transitive: Dict[ControlFlowBlock, Set[ControlFlowBlock]] = defaultdict(lambda: set()) + all_doms_transitive: Dict[ControlFlowBlock, + OrderedSet[ControlFlowBlock]] = defaultdict(lambda: OrderedSet()) for cfg in sdfg.all_control_flow_regions(): if isinstance(cfg, ConditionalBlock): idom_dict[cfg] = {b: b for _, b in cfg.branches} - all_doms = {b: set([b]) for _, b in cfg.branches} + all_doms = {b: OrderedSet([b]) for _, b in cfg.branches} else: idom_dict[cfg] = nx.immediate_dominators(cfg.nx, cfg.start_block) all_doms = cfg_analysis.all_dominators(cfg, idom_dict[cfg]) @@ -658,15 +680,16 @@ def apply_pass(self, top_sdfg: SDFG, pipeline_results: Dict[str, Any]) -> Dict[i all_doms_transitive[k].add(cfg) all_doms_transitive[k].update(all_doms_transitive[cfg]) - access_nodes: Dict[str, Dict[SDFGState, Tuple[Set[nd.AccessNode], Set[nd.AccessNode]]]] = pipeline_results[ - FindAccessNodes.__name__][sdfg.cfg_id] + access_nodes: Dict[str, Dict[SDFGState, Tuple[OrderedSet[nd.AccessNode], + OrderedSet[nd.AccessNode]]]] = pipeline_results[ + FindAccessNodes.__name__][sdfg.cfg_id] block_reach: Dict[ControlFlowBlock, - Set[ControlFlowBlock]] = pipeline_results[ControlFlowBlockReachability.__name__] + OrderedSet[ControlFlowBlock]] = pipeline_results[ControlFlowBlockReachability.__name__] anames = sdfg.arrays.keys() for desc in sdfg.arrays: - desc_states_with_nodes = set(access_nodes[desc].keys()) + desc_states_with_nodes = OrderedSet(access_nodes[desc].keys()) for state in desc_states_with_nodes: for read_node in access_nodes[desc][state][0]: write = self._find_dominating_write(desc, state, read_node, access_nodes, idom_dict, @@ -698,7 +721,7 @@ def apply_pass(self, top_sdfg: SDFG, pipeline_results: Dict[str, Any]) -> Dict[i # If any write A is dominated by another write B and any reads in B's scope are also reachable by A, # then merge A and its scope into B's scope. - to_remove = set() + to_remove = OrderedSet() for write, accesses in result[desc].items(): if write is None: continue @@ -715,7 +738,7 @@ def apply_pass(self, top_sdfg: SDFG, pipeline_results: Dict[str, Any]) -> Dict[i other_accesses.update(accesses) other_accesses.add(write) to_remove.add(write) - result[desc][write] = set() + result[desc][write] = OrderedSet() for write in to_remove: del result[desc][write] top_result[sdfg.cfg_id] = result @@ -737,14 +760,14 @@ def modifies(self) -> ppl.Modifies: def should_reapply(self, modified: ppl.Modifies) -> bool: return modified & ppl.Modifies.Memlets - def apply_pass(self, top_sdfg: SDFG, _) -> Dict[int, Dict[str, Set[Memlet]]]: + def apply_pass(self, top_sdfg: SDFG, _) -> Dict[int, Dict[str, OrderedSet[Memlet]]]: """ :return: A dictionary mapping each data descriptor name to a set of memlets. """ - top_result: Dict[int, Dict[str, Set[Memlet]]] = dict() + top_result: Dict[int, Dict[str, OrderedSet[Memlet]]] = dict() for sdfg in top_sdfg.all_sdfgs_recursive(): - result: Dict[str, Set[Memlet]] = defaultdict(set) + result: Dict[str, OrderedSet[Memlet]] = defaultdict(OrderedSet) for state in sdfg.states(): for anode in state.data_nodes(): for e in state.all_edges(anode): @@ -779,17 +802,17 @@ def modifies(self) -> ppl.Modifies: def should_reapply(self, modified: ppl.Modifies) -> bool: return modified & ppl.Modifies.Memlets - def apply_pass(self, top_sdfg: SDFG, _) -> Dict[int, Dict[str, Set[Union[Memlet, nd.CodeNode]]]]: + def apply_pass(self, top_sdfg: SDFG, _) -> Dict[int, Dict[str, OrderedSet[Union[Memlet, nd.CodeNode]]]]: """ :return: A dictionary mapping each data descriptor name to a set of memlets. """ - top_result: Dict[int, Dict[str, Set[Union[Memlet, nd.CodeNode]]]] = dict() + top_result: Dict[int, Dict[str, OrderedSet[Union[Memlet, nd.CodeNode]]]] = dict() for sdfg in top_sdfg.all_sdfgs_recursive(): - result: Dict[str, Set[Memlet]] = defaultdict(set) - reference_descs = set(k for k, v in sdfg.arrays.items() if isinstance(v, dt.Reference)) + result: Dict[str, OrderedSet[Memlet]] = defaultdict(OrderedSet) + reference_descs = OrderedSet(k for k, v in sdfg.arrays.items() if isinstance(v, dt.Reference)) for state in sdfg.states(): - code_sources: Dict[str, Set[nd.CodeNode]] = defaultdict(set) + code_sources: Dict[str, OrderedSet[nd.CodeNode]] = defaultdict(OrderedSet) for anode in state.data_nodes(): if anode.data not in reference_descs: continue @@ -868,21 +891,22 @@ def should_reapply(self, modified: ppl.Modifies) -> bool: # If anything was modified, reapply return modified & ppl.Modifies.Everything - def _derive_parameter_datasize_constraints(self, sdfg: SDFG, invariants: Dict[str, Set[str]]) -> None: - handled = set() + def _derive_parameter_datasize_constraints(self, sdfg: SDFG, invariants: Dict[str, OrderedSet[str]]) -> None: + handled = OrderedSet() for arr in sdfg.arrays.values(): for dim in arr.shape: if isinstance(dim, symbolic.symbol) and not dim in handled: ds = str(dim) if ds not in invariants: - invariants[ds] = set() + invariants[ds] = OrderedSet() invariants[ds].add(f'{ds} > 0') if self.assume_max_data_size is not None: invariants[ds].add(f'{ds} <= {self.assume_max_data_size}') handled.add(ds) - def apply_pass(self, sdfg: SDFG, _) -> Tuple[Dict[str, Set[str]], Dict[str, Set[str]], Dict[str, Set[str]]]: - invariants: Dict[str, Set[str]] = {} + def apply_pass(self, sdfg: SDFG, + _) -> Tuple[Dict[str, OrderedSet[str]], Dict[str, OrderedSet[str]], Dict[str, OrderedSet[str]]]: + invariants: Dict[str, OrderedSet[str]] = {} self._derive_parameter_datasize_constraints(sdfg, invariants) return {}, invariants, {} @@ -910,11 +934,11 @@ def __init__(self): self.apply_to_conditionals = True def depends_on(self): - return {ControlFlowBlockReachability} + return [ControlFlowBlockReachability] - def _propagate_in_cfg(self, cfg: ControlFlowRegion, reachable: Dict[ControlFlowBlock, Set[ControlFlowBlock]], + def _propagate_in_cfg(self, cfg: ControlFlowRegion, reachable: Dict[ControlFlowBlock, OrderedSet[ControlFlowBlock]], starting_executions: int, starting_dynamic_executions: bool): - visited_blocks: Set[ControlFlowBlock] = set() + visited_blocks: OrderedSet[ControlFlowBlock] = OrderedSet() traversal_q: deque[Tuple[ControlFlowBlock, int, bool, List[str]]] = deque() traversal_q.append((cfg.start_block, starting_executions, starting_dynamic_executions, [])) while traversal_q: @@ -1064,13 +1088,13 @@ def should_reapply(self, modified: ppl.Modifies) -> bool: return modified & ppl.Modifies.CFG def depends_on(self): - return {} + return [] - def apply_pass(self, top_sdfg: SDFG, pipeline_res: Dict) -> Set[nd.AccessNode]: + def apply_pass(self, top_sdfg: SDFG, pipeline_res: Dict) -> OrderedSet[nd.AccessNode]: """ :return: A set of access nodes, which are unique writes in conditional blocks. """ - cond_unique = set() + cond_unique = OrderedSet() for cfb in top_sdfg.all_control_flow_blocks(recursive=True): if not isinstance(cfb, ConditionalBlock): continue @@ -1086,12 +1110,15 @@ def apply_pass(self, top_sdfg: SDFG, pipeline_res: Dict) -> Set[nd.AccessNode]: for st in br.all_states(): for an in st.data_nodes(): array_name = an.data - write_subsets = set(e.data.dst_subset for e in st.in_edges(an)) + write_subsets = OrderedSet(e.data.dst_subset for e in st.in_edges(an)) wss = str(write_subsets) if array_name not in access_write_branch: access_write_branch[array_name] = {} if wss not in access_write_branch[array_name]: - access_write_branch[array_name][wss] = {"branches": set(), "access_nodes": set()} + access_write_branch[array_name][wss] = { + "branches": OrderedSet(), + "access_nodes": OrderedSet() + } access_write_branch[array_name][wss]["branches"].add(br) access_write_branch[array_name][wss]["access_nodes"].add(an) diff --git a/dace/transformation/passes/analysis/scope_data_and_symbol_analysis.py b/dace/transformation/passes/analysis/scope_data_and_symbol_analysis.py index 16b3b7285b..6824776bf8 100644 --- a/dace/transformation/passes/analysis/scope_data_and_symbol_analysis.py +++ b/dace/transformation/passes/analysis/scope_data_and_symbol_analysis.py @@ -52,7 +52,7 @@ def should_reapply(self, modified: ppl.Modifies) -> bool: return modified & (ppl.Modifies.CFG | ppl.Modifies.SDFG | ppl.Modifies.Nodes) def depends_on(self): - return {} + return [] def apply_pass(self, sdfg: dace.SDFG, pipeline_res: Dict) -> Dict[str, ScopeAnalysis]: """ diff --git a/dace/transformation/passes/array_elimination.py b/dace/transformation/passes/array_elimination.py index 9d6d52cf48..1aa6ddc41b 100644 --- a/dace/transformation/passes/array_elimination.py +++ b/dace/transformation/passes/array_elimination.py @@ -13,6 +13,7 @@ SqueezeViewRemove, UnsqueezeViewRemove, RemoveSliceView) from dace.transformation.passes import analysis as ap from dace.transformation.transformation import SingleStateTransformation +from ordered_set import OrderedSet @properties.make_properties @@ -32,7 +33,7 @@ def should_reapply(self, modified: ppl.Modifies) -> bool: return modified & ppl.Modifies.AccessNodes def depends_on(self): - return {ap.StateReachability, ap.FindAccessStates} + return [ap.StateReachability, ap.FindAccessStates] def apply_pass(self, sdfg: SDFG, pipeline_results: Dict[str, Any]) -> Optional[Set[str]]: """ @@ -56,7 +57,7 @@ def apply_pass(self, sdfg: SDFG, pipeline_results: Dict[str, Any]) -> Optional[S return None for state in reversed(state_order): # Find all data descriptors that will no longer be used after this state - removable_data: Set[str] = set( + removable_data: OrderedSet[str] = OrderedSet( s for s in access_sets if state in access_sets[s] and not (access_sets[s] & reachable[state]) - {state}) # Find duplicate access nodes as an ordered list @@ -67,6 +68,7 @@ def apply_pass(self, sdfg: SDFG, pipeline_results: Dict[str, Any]) -> Optional[S # Merge source and sink access nodes removed_nodes = self.merge_access_nodes(state, access_nodes, lambda n: state.in_degree(n) == 0) + removed_nodes |= self.merge_access_nodes(state, access_nodes, lambda n: state.out_degree(n) == 0) # Remove redundant views @@ -85,6 +87,8 @@ def apply_pass(self, sdfg: SDFG, pipeline_results: Dict[str, Any]) -> Optional[S # If node is completely removed from graph, erase data descriptor for aname, desc in list(sdfg.arrays.items()): + if isinstance(desc, data.DistributedDescriptor): + continue if not desc.transient or isinstance(desc, data.Scalar): continue if aname not in access_sets or not access_sets[aname]: diff --git a/dace/transformation/passes/constant_propagation.py b/dace/transformation/passes/constant_propagation.py index f3a87db377..f06ae726dd 100644 --- a/dace/transformation/passes/constant_propagation.py +++ b/dace/transformation/passes/constant_propagation.py @@ -513,6 +513,6 @@ def _data_independent_assignments(self, edge: InterstateEdge, arrays: Set[str]) Return symbol assignments that only depend on other symbols and constants, rather than data descriptors. """ return { - k: v if (not (symbolic.free_symbols_and_functions(v) & arrays)) else _UnknownValue + k: v if (not ((symbolic.free_symbols_and_functions(v) | symbolic.arrays(v)) & arrays)) else _UnknownValue for k, v in edge.assignments.items() } diff --git a/dace/transformation/passes/dead_dataflow_elimination.py b/dace/transformation/passes/dead_dataflow_elimination.py index bc7598bbc8..cd0149d37b 100644 --- a/dace/transformation/passes/dead_dataflow_elimination.py +++ b/dace/transformation/passes/dead_dataflow_elimination.py @@ -3,7 +3,7 @@ import ast from collections import defaultdict from dataclasses import dataclass -from typing import Any, Dict, List, Optional, Set, Tuple, Type +from typing import Dict, List, Set, Tuple, Type from dace import SDFG, Memlet, SDFGState, data, dtypes, properties from dace.frontend.python import astutils @@ -44,8 +44,8 @@ def should_reapply(self, modified: ppl.Modifies) -> bool: # If dataflow or states changed, new dead code may be exposed return modified & (ppl.Modifies.Nodes | ppl.Modifies.Edges | ppl.Modifies.CFG) - def depends_on(self) -> Set[Type[ppl.Pass]]: - return {ap.ControlFlowBlockReachability, ap.AccessSets} + def depends_on(self) -> List[Type[ppl.Pass]]: + return [ap.ControlFlowBlockReachability, ap.AccessSets] def apply(self, region, pipeline_results): """ diff --git a/dace/transformation/passes/full_map_fusion.py b/dace/transformation/passes/full_map_fusion.py index ad9d8e39cf..749a726f36 100644 --- a/dace/transformation/passes/full_map_fusion.py +++ b/dace/transformation/passes/full_map_fusion.py @@ -1,9 +1,7 @@ # Copyright 2019-2025 ETH Zurich and the DaCe authors. All rights reserved. -from typing import Any, Dict, Optional, Set +from typing import Any, Dict, Optional -import warnings - -from dace import SDFG, SDFGState, properties, transformation +from dace import SDFG, properties, transformation from dace.transformation import pass_pipeline as ppl, dataflow as dftrans from dace.transformation.passes import analysis as ap, pattern_matching as pmp @@ -164,7 +162,7 @@ def should_reapply(self, modified: ppl.Modifies) -> bool: return modified & (ppl.Modifies.Scopes | ppl.Modifies.AccessNodes | ppl.Modifies.Memlets | ppl.Modifies.States) def depends_on(self): - return {ap.FindSingleUseData} + return [ap.FindSingleUseData] def apply_pass(self, sdfg: SDFG, pipeline_results: Dict[str, Any]) -> Optional[int]: """ diff --git a/dace/transformation/passes/lift_struct_views.py b/dace/transformation/passes/lift_struct_views.py index c6c38058f1..8c912aa6eb 100644 --- a/dace/transformation/passes/lift_struct_views.py +++ b/dace/transformation/passes/lift_struct_views.py @@ -14,7 +14,6 @@ from dace import data as dt from dace import dtypes -import sys from typing import Literal dirtype = Literal['in', 'out'] @@ -356,7 +355,7 @@ def should_reapply(self, modified: ppl.Modifies) -> bool: return modified & ppl.Modifies.AccessNodes & ppl.Modifies.Tasklets & ppl.Modifies.Memlets def depends_on(self): - return {} + return [] def _lift_control_flow_region_access(self, cfg: ControlFlowRegion, result: Dict[str, Set[str]]) -> bool: lifted_something = False diff --git a/dace/transformation/passes/lift_trivial_if.py b/dace/transformation/passes/lift_trivial_if.py new file mode 100644 index 0000000000..da39724db3 --- /dev/null +++ b/dace/transformation/passes/lift_trivial_if.py @@ -0,0 +1,160 @@ +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. +"""Simplification pass that removes ``ConditionalBlock`` nodes whose condition is provably constant.""" +import re +import dace +from typing import Any, Dict, Optional, Set, Union +from dace import SDFG, ControlFlowRegion +from dace import symbolic +from dace.properties import CodeBlock +from dace.sdfg.sdfg import ConditionalBlock +from dace.sdfg.state import ControlFlowBlock +from dace.transformation.helpers import move_branch_cfg_up_discard_conditions +from dace.transformation import pass_pipeline as ppl, transformation +import dace.sdfg.utils as sdutil +import sympy +from sympy import pycode + + +@transformation.explicit_cf_compatible +class LiftTrivialIf(ppl.Pass): + """Remove ``ConditionalBlock`` nodes whose condition statically evaluates to a literal. + + Handles two shapes: a single-branch ``if`` (drop or replace with an empty + branch depending on the truth value) and an ``if/else`` pair (drop the side + that is unreachable). Runs to a fixed point at each region level and recurses + into nested SDFGs. + """ + + def modifies(self) -> ppl.Modifies: + return ppl.Modifies.CFG | ppl.Modifies.States + + def should_reapply(self, modified: ppl.Modifies) -> bool: + return modified & ppl.Modifies.CFG + + def depends_on(self): + return [] + + def _make_unique_names(self, sdfg: dace.SDFG): + all_blocks = { + n + for n, _ in sdfg.all_nodes_recursive() + if isinstance(n, dace.SDFGState) or isinstance(n, ControlFlowRegion) or isinstance(n, ControlFlowBlock) + } + all_labels: Set[str] = set() + for n in all_blocks: + new_label = dace.utils.find_new_name(n.label, all_labels) + all_labels.add(new_label) + n.label = new_label + + def _trivial_cond_check(self, code: CodeBlock, val: bool): + if code.language != dace.dtypes.Language.Python: + return False + + # Primary: pystr_to_symbolic already handles Python and/or/not and + # comparison operators, We require a concrete literal back -- bool() of + # an unevaluated sympy expression (e.g. ``A[0]`` -> Function(0)) is + # truthy, which would mis-classify dynamic conditions as trivial. + try: + expr = symbolic.pystr_to_symbolic(code.as_string) + result = symbolic.evaluate(expr, symbols={}) + if isinstance(result, (bool, int, sympy.Integer)) or result in (sympy.true, sympy.false): + return bool(result) is val + except Exception: + pass + + # Fallback: Some SDFGs (e.g. Fortran frontend) produce nested comparisons like + # ``(a == 1) == 0`` that sympy refuses to compare bool against an int. + # Try as best effort to rewrite boolean ops/literals to arithmetic over 0/1 and let + # SymExpr.simplify reduce it. + try: + tokens = re.split(r'(\s+|[()\[\]])', code.as_string) + replacements = {"True": "1", "False": "0", "and": "*", "or": "+"} + rewritten = " ".join(replacements.get(t.strip(), t.strip()) for t in tokens).strip() + simplified = dace.symbolic.SymExpr(rewritten).simplify() + result = symbolic.evaluate(dace.symbolic.SymExpr(pycode(simplified)), symbols={}) + if isinstance(result, (bool, int, sympy.Integer)) or result in (sympy.true, sympy.false): + return bool(result) is val + except Exception: + pass + return False + + def _trivially_true(self, code: CodeBlock): + return self._trivial_cond_check(code, True) + + def _trivially_false(self, code: CodeBlock): + return self._trivial_cond_check(code, False) + + def _detect_and_remove_top_level_trivial_ifs(self, graph: Union[ControlFlowRegion, SDFG]): + cfb_to_rm_cfg_to_keep = set() + rmed_count = 0 + for cfb in graph.nodes(): + if isinstance(cfb, ConditionalBlock): + # Supported variants: + # 1. if (cond) where cond is always true + # 2. if (cond) else () + # 2.1 where cond is always true + # 2.2 cond is always false + conditions_and_cfgs = cfb.branches + if len(conditions_and_cfgs) == 1: + cond, cfg = conditions_and_cfgs[0] + if self._trivially_true(cond): + cfb_to_rm_cfg_to_keep.add((cfb, cfg)) + elif self._trivially_false(cond): + _cfg = ControlFlowRegion(label=f"empty_cfg_of_{cfb.label}", sdfg=cfb.sdfg, parent=cfb) + _cfg.add_state(label="empty_placholder", is_start_block=True) + cfb.add_branch(condition=None, branch=_cfg) + cfb_to_rm_cfg_to_keep.add((cfb, _cfg)) + elif len(conditions_and_cfgs) == 2: + cond1, cfg1 = conditions_and_cfgs[0] + cond2, cfg2 = conditions_and_cfgs[1] + # Either one of them must be none + if cond1 is not None and cond2 is not None: + continue + (not_none_cond, not_none_cfg), (none_cond, none_cfg) = (((cond1, cfg1), + (cond2, cfg2)) if cond1 is not None else + ((cond2, cfg2), (cond1, cfg1))) + + if self._trivially_true(not_none_cond): #2.1 + cfb_to_rm_cfg_to_keep.add((cfb, not_none_cfg)) + elif self._trivially_false(not_none_cond): #2.2 + cfb_to_rm_cfg_to_keep.add((cfb, none_cfg)) + + for cfb, cfg in cfb_to_rm_cfg_to_keep: + move_branch_cfg_up_discard_conditions(cfb, cfg) + assert cfb not in graph.nodes() + rmed_count += 1 + + sdutil.set_nested_sdfg_parent_references(graph.sdfg) + graph.sdfg.reset_cfg_list() + + return rmed_count + + def _detect_trivial_ifs_and_rm_cfg(self, graph: Union[ControlFlowRegion, SDFG]): + # We might now have trivial control flow blocks at top level, apply in fixpoint + rmed_count = self._detect_and_remove_top_level_trivial_ifs(graph) + local_rmed_count = rmed_count + while local_rmed_count > 0: + local_rmed_count = self._detect_and_remove_top_level_trivial_ifs(graph) + rmed_count += local_rmed_count + + # Now go one one more level in the node list + for node in graph.all_control_flow_blocks(): + local_rmed_count = self._detect_and_remove_top_level_trivial_ifs(node) + rmed_count += local_rmed_count + + # Recurse in to nSDFGs + for state in graph.all_states(): + for node in state.nodes(): + if isinstance(node, dace.nodes.NestedSDFG): + rmed_count += self._detect_trivial_ifs_and_rm_cfg(node.sdfg) + + return rmed_count + + def apply_pass(self, sdfg: SDFG, pipeline_results: Dict[str, Any]) -> Optional[Dict[str, Set[str]]]: + # Start with top level nodes and continue further to ensure a trivial if within another trivial if + # can be processed correctly + self._make_unique_names(sdfg) + sdfg.reset_cfg_list() + self._detect_trivial_ifs_and_rm_cfg(sdfg) + sdfg.reset_cfg_list() + return None diff --git a/dace/transformation/passes/loop_local_memory_reduction.py b/dace/transformation/passes/loop_local_memory_reduction.py index 233477b6c9..ee8f57e2c9 100644 --- a/dace/transformation/passes/loop_local_memory_reduction.py +++ b/dace/transformation/passes/loop_local_memory_reduction.py @@ -3,8 +3,7 @@ import sympy as sp from dace import sdfg as sd, symbolic, properties from dace import data as dt -from dace.sdfg import SDFGState -from dace.sdfg.state import LoopRegion, ConditionalBlock +from dace.sdfg.state import LoopRegion from dace.data import Scalar from dace.transformation import transformation as xf from dace.transformation import pass_pipeline as ppl @@ -115,7 +114,7 @@ def should_reapply(self, modified: ppl.Modifies) -> bool: return modified != ppl.Modifies.Nothing def depends_on(self): - return {StateReachability, FindAccessStates, ConditionUniqueWrites} + return [StateReachability, FindAccessStates, ConditionUniqueWrites] def apply_pass(self, sdfg: sd.SDFG, pipeline_results: Dict[str, Any]) -> Optional[Set[str]]: self.num_applications = 0 diff --git a/dace/transformation/passes/loop_to_reduce.py b/dace/transformation/passes/loop_to_reduce.py new file mode 100644 index 0000000000..26f2876d42 --- /dev/null +++ b/dace/transformation/passes/loop_to_reduce.py @@ -0,0 +1,512 @@ +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. +"""Detect scalar accumulator loops and replace them with ``Reduce`` nodes. + +Three loop shapes are recognised (``identity=None`` on the emitted +``Reduce`` so the pre-loop accumulator seeds the fold): + +- **Tasklet**: a single-state containing one two-input tasklet that + writes to the accumulator. +- **Interstate edge**: body = 2 empty states joined by one interstate + edge with assignment ``{sym: sym arr[]}``. +- **Conditional interstate edge**: body = a single ``ConditionalBlock`` + with one branch guarded by ``sym arr[]`` (``cmp`` in + ``>``/``>=``/``<``/``<=``) whose body is the 2-empty-states + edge + shape above with assignment ``{sym: arr[]}``. ``>``/``>=`` lift + to ``max``, ``<``/``<=`` lift to ``min``. + +Accumulator forms accepted: a ``Scalar``, a length-1 ``Array``, a single +loop-invariant slice of a multi-element ``Array`` (``C[k]``). +""" +import ast +import copy as _copy +from typing import Dict, NamedTuple, Optional + +import sympy + +from dace import SDFG, SDFGState, data, dtypes, memlet as mm, nodes, properties, subsets, symbolic +from dace.sdfg.state import ConditionalBlock, ControlFlowRegion, LoopRegion +from dace.symbolic import AND, OR, bitwise_and, bitwise_or, Subscript +from dace.transformation import pass_pipeline as ppl +from dace.transformation import transformation as xf +from dace.transformation.passes.analysis import loop_analysis + +# Ops in these tables are commutative by construction, so we skip calling +# ``dace.frontend.operations.is_op_commutative`` (which returns ``None`` for +# ``max`` / ``min`` because Python's builtins choke on symbolic arguments). +_BINOP_TO_WCR: Dict[type, str] = { + ast.Add: "lambda a, b: a + b", + ast.Mult: "lambda a, b: a * b", + ast.BitAnd: "lambda a, b: a & b", + ast.BitOr: "lambda a, b: a | b", + ast.BitXor: "lambda a, b: a ^ b", +} +_BOOLOP_TO_WCR: Dict[type, str] = { + ast.Or: "lambda a, b: a | b", + ast.And: "lambda a, b: a & b", +} +_CALL_TO_WCR: Dict[str, str] = { + "max": "lambda a, b: max(a, b)", + "min": "lambda a, b: min(a, b)", +} +# For a guard `lhs rhs` where the assignment inside writes `sym = arr[i]`, +# the reduction is max iff the condition fires when arr is larger than sym. +_CMP_GT = (ast.Gt, ast.GtE) +_CMP_LT = (ast.Lt, ast.LtE) + + +class _Reduction(NamedTuple): + wcr: str + accum: str # data-descriptor name, or DaCe symbol + accum_subset: subsets.Subset + array: str + array_subset: subsets.Subset + + +@properties.make_properties +@xf.explicit_cf_compatible +class LoopToReduce(ppl.Pass): + """Lift scalar-accumulator loops to Reduction library nodes.""" + + permissive = properties.Property( + dtype=bool, + default=False, + desc="Enable extractors that make semantic assumptions about input " + "data (e.g. the ``any``/``all`` conditional-const-assign pattern " + "which assumes the guard array is 0/1-valued).", + ) + + def __init__(self, permissive: bool = False): + super().__init__() + self.permissive = permissive + + def modifies(self) -> ppl.Modifies: + return ppl.Modifies.CFG | ppl.Modifies.Nodes | ppl.Modifies.Memlets + + def should_reapply(self, modified: ppl.Modifies) -> bool: + return bool(modified & ppl.Modifies.CFG) + + def apply_pass(self, sdfg: SDFG, _) -> Optional[int]: + count = 0 + for node, parent in list(sdfg.all_nodes_recursive()): + if not isinstance(node, LoopRegion): + continue + info = _extract(node, sdfg, permissive=self.permissive) + if info is None: + continue + _lift(parent, node, info) + count += 1 + return count or None + + +def _one_elem(subset) -> Optional[int]: + """Integer number of elements in ``subset``, or ``None`` if non-constant.""" + if subset is None: + return None + try: + s = symbolic.simplify(subset.num_elements()) + except Exception: + return None + return int(s) if s.is_Integer else None + + +def _uses(subset: subsets.Subset, sym: sympy.Symbol) -> bool: + return subset is not None and any(symbolic.pystr_to_symbolic(str(e)) == sym for e in subset.free_symbols) + + +def _scalar_equiv(sdfg: SDFG, a: str, b: str) -> bool: + """Same descriptor, or two distinct dtype-compatible scalar-equivalents.""" + if a == b: + return True + da, db = sdfg.arrays.get(a), sdfg.arrays.get(b) + if da is None or db is None or da.dtype != db.dtype: + return False + + def scalar_like(d) -> bool: + return isinstance(d, data.Scalar) or (isinstance(d, data.Array) and all(s == 1 for s in d.shape)) + + return scalar_like(da) and scalar_like(db) + + +def _expand_over_loop(subset: subsets.Subset, loop_var: sympy.Symbol, start, end) -> Optional[subsets.Range]: + """Widen ``subset`` -- which uses ``loop_var`` linearly -- over the + iteration range ``[start, end]``.""" + if not isinstance(subset, subsets.Range): + return None + ranges = [] + for rb, re_, rs in subset.ndrange(): + if rb != re_ or rs != 1: + return None + offset = symbolic.simplify(rb - loop_var) + if offset.has(loop_var): + return None + ranges.append((symbolic.simplify(start + offset), symbolic.simplify(end + offset), 1)) + return subsets.Range(ranges) + + +def _cmp_to_wcr(cond, target: str, array: str) -> Optional[str]: + """Map a ``sym arr[...]`` (or reversed) guard to a max/min WCR.""" + try: + tree = ast.parse(cond.as_string, mode="eval").body + except (SyntaxError, TypeError, ValueError): + return None + if not isinstance(tree, ast.Compare) or len(tree.ops) != 1: + return None + op_type = type(tree.ops[0]) + if op_type not in _CMP_GT and op_type not in _CMP_LT: + return None + + def _is_target(n): + return isinstance(n, ast.Name) and n.id == target + + def _is_array(n): + return (isinstance(n, ast.Subscript) and isinstance(n.value, ast.Name) and n.value.id == array) + + left, right = tree.left, tree.comparators[0] + if _is_target(left) and _is_array(right): + array_on_left = False + elif _is_array(left) and _is_target(right): + array_on_left = True + else: + return None + is_gt = op_type in _CMP_GT + arr_is_larger = array_on_left == is_gt + return "lambda a, b: max(a, b)" if arr_is_larger else "lambda a, b: min(a, b)" + + +def _extract_any_pattern(cond, const_rhs: int, target: str, sdfg: SDFG, loop_var_sym, start, + end) -> Optional["_Reduction"]: + """Match ``{sym: const}`` conditional-interstate-edge "any"/"all". + + Body = ``ConditionalBlock`` with one branch, guard ``arr[] C`` + (C integer), branch = 2 empty states + interstate edge with assignment + ``{sym: }`` where ``const_rhs`` is 0 or 1. + + The guard array is assumed to be 0/1-valued, so ``any(arr[...] == 1)`` + over the iteration range is equivalent to the bitwise-OR of ``arr[...]`` + -- no predicate synthesis needed, a plain ``Reduce(|)`` over the array + slice suffices. ``const_rhs == 1`` lifts to OR; ``const_rhs == 0`` lifts + to AND. + """ + if const_rhs == 1: + wcr = "lambda a, b: a | b" + elif const_rhs == 0: + wcr = "lambda a, b: a & b" + else: + return None + + try: + tree = ast.parse(cond.as_string, mode="eval").body + except (SyntaxError, TypeError, ValueError): + return None + if not isinstance(tree, ast.Compare) or len(tree.ops) != 1: + return None + left, right = tree.left, tree.comparators[0] + + def _is_subscript_on_array(n): + return (isinstance(n, ast.Subscript) and isinstance(n.value, ast.Name) and n.value.id in sdfg.arrays) + + def _is_int_const(n): + return isinstance(n, ast.Constant) and isinstance(n.value, int) + + if _is_subscript_on_array(left) and _is_int_const(right): + sub = left + elif _is_subscript_on_array(right) and _is_int_const(left): + sub = right + else: + return None + + array = sub.value.id + slice_node = sub.slice + args_ast = slice_node.elts if isinstance(slice_node, ast.Tuple) else [slice_node] + + try: + sym_args = [symbolic.pystr_to_symbolic(ast.unparse(a)) for a in args_ast] + except Exception: + return None + + if len(sym_args) != len(sdfg.arrays[array].shape): + return None + + # Exactly one axis must depend on the loop variable (linearly, offset ∉ sym). + axis_for_iter = None + offset = None + for i, a in enumerate(sym_args): + if a.has(loop_var_sym): + if axis_for_iter is not None: + return None + axis_for_iter = i + try: + off = symbolic.simplify(a - loop_var_sym) + except Exception: + return None + if off.has(loop_var_sym): + return None + offset = off + if axis_for_iter is None: + return None + + ranges = [] + for i, a in enumerate(sym_args): + if i == axis_for_iter: + ranges.append((symbolic.simplify(start + offset), symbolic.simplify(end + offset), 1)) + else: + ranges.append((a, a, 1)) + return _Reduction( + wcr=wcr, + accum=target, + accum_subset=subsets.Range([(0, 0, 1)]), + array=array, + array_subset=subsets.Range(ranges), + ) + + +def _extract(loop: LoopRegion, sdfg: SDFG, permissive: bool = False) -> Optional[_Reduction]: + if not loop.loop_variable: + return None + start = loop_analysis.get_init_assignment(loop) + end = loop_analysis.get_loop_end(loop) + stride = loop_analysis.get_loop_stride(loop) + if start is None or end is None or stride is None or stride != 1: + return None + + blocks = loop.nodes() + loop_var = loop.loop_variable + loop_var_sym = symbolic.pystr_to_symbolic(loop_var) + + # Tasklet pattern: single state with exactly one tasklet. + if len(blocks) == 1 and isinstance(blocks[0], SDFGState): + state = blocks[0] + tasklet = None + for n in state.nodes(): + if isinstance(n, nodes.Tasklet): + if tasklet is not None: + return None + tasklet = n + elif not isinstance(n, nodes.AccessNode): + return None + if tasklet is None: + return None + if tasklet.code.language != dtypes.Language.Python: + return None + + # Classify the tasklet's single-assignment body. + try: + tree = ast.parse((tasklet.code.as_string or "").strip()) + except SyntaxError: + return None + if len(tree.body) != 1 or not isinstance(tree.body[0], ast.Assign): + return None + rhs = tree.body[0].value + if isinstance(rhs, ast.BinOp): + wcr = _BINOP_TO_WCR.get(type(rhs.op)) + elif isinstance(rhs, ast.BoolOp) and len(rhs.values) == 2: + wcr = _BOOLOP_TO_WCR.get(type(rhs.op)) + elif (isinstance(rhs, ast.Call) and isinstance(rhs.func, ast.Name) and len(rhs.args) == 2): + wcr = _CALL_TO_WCR.get(rhs.func.id) + else: + wcr = None + if wcr is None: + return None + + # Tasklet must have exactly 2 data inputs and 1 data output. + def _has_data(e): + return e.data is not None and not e.data.is_empty() + + in_edges = [e for e in state.in_edges(tasklet) if _has_data(e)] + out_edges = [e for e in state.out_edges(tasklet) if _has_data(e)] + if len(in_edges) != 2 or len(out_edges) != 1: + return None + + write_edge = out_edges[0] + if not isinstance(write_edge.dst, nodes.AccessNode): + return None + accum = write_edge.dst.data + if accum not in sdfg.arrays: + return None + write_subset = write_edge.data.subset + if _one_elem(write_subset) != 1 or _uses(write_subset, loop_var_sym): + return None + + # Resolve each tasklet input. + resolved = [] + for e in in_edges: + src = e.src + if not isinstance(src, nodes.AccessNode): + return None + desc = sdfg.arrays.get(src.data) + if desc is None: + return None + if desc.transient and len(state.in_edges(src)) == 1 and len(state.out_edges(src)) == 1: + pred = state.in_edges(src)[0] + if (not isinstance(pred.src, nodes.AccessNode) or pred.data is None or pred.data.subset is None + or _one_elem(e.data.subset) != _one_elem(pred.data.subset)): + return None + resolved.append((pred.src.data, _copy.deepcopy(pred.data.subset))) + else: + resolved.append((src.data, e.data.subset)) + + accum_ok = False + array, arr_subset = None, None + for name, sub in resolved: + if _uses(sub, loop_var_sym): + if array is not None: + return None + array, arr_subset = name, sub + elif _one_elem(sub) == 1 and ((name == accum and sub == write_subset) or + (name != accum and _scalar_equiv(sdfg, name, accum))): + accum_ok = True + if not accum_ok or array is None or array == accum: + return None + + expanded = _expand_over_loop(arr_subset, loop_var_sym, start, end) + if expanded is None: + return None + return _Reduction(wcr, accum, write_subset, array, expanded) + + # Interstate-edge pattern: 2 empty states + 1 edge with 1 assignment, + # either at loop level or inside a single-branch ConditionalBlock whose + # guard is a >/>=/}`` with an array-predicate + # guard; lifts to OR / AND over the (0/1-valued) guard array. + # Gated on ``permissive`` -- the lift is only semantically correct + # if the guard array happens to hold only 0/1 values, which the + # pass cannot verify statically. + if permissive and isinstance(expr, sympy.Integer) and int(expr) in (0, 1): + return _extract_any_pattern(cond, int(expr), target, sdfg, loop_var_sym, start, end) + # Pure copy ``sym = arr[f(i)]`` gated by a max/min comparison. + if not (isinstance(expr, Subscript) and symbolic.arrays(expr) & sdfg.arrays.keys()): + return None + arr_call = expr + wcr = _cmp_to_wcr(cond, target, str(arr_call.args[0])) + if wcr is None: + return None + + array = str(arr_call.args[0]) + # ``Subscript`` carries the head plus indices, so a 1-D access has two args. + if len(sdfg.arrays[array].shape) != 1 or len(arr_call.args) != 2: + return None + offset = symbolic.simplify(arr_call.args[1] - loop_var_sym) + if offset.has(loop_var_sym): + return None + + return _Reduction( + wcr=wcr, + accum=target, + accum_subset=subsets.Range([(0, 0, 1)]), + array=array, + array_subset=subsets.Range([(symbolic.simplify(start + offset), symbolic.simplify(end + offset), 1)]), + ) + + return None + + +def _lift(parent: ControlFlowRegion, loop: LoopRegion, info: _Reduction): + """Replace ``loop`` with a ``Reduce``. If the accumulator is a data + descriptor we write to it directly; if it's a symbol we synthesize a + transient scalar, seed it from the symbol, and assign back on exit.""" + import dace + root = parent + while not isinstance(root, SDFG): + root = root.parent_graph + + was_start = parent.start_block is loop + in_edges = list(parent.in_edges(loop)) + out_edges = list(parent.out_edges(loop)) + extra_assignments: Dict[str, str] = {} + + if info.accum in root.arrays: + red_state = parent.add_state(loop.label + "_reduce", is_start_block=was_start) + entry = red_state + dest_name, dest_subset = info.accum, info.accum_subset + else: + tmp_name, _ = root.add_scalar(f"_red_tmp_{info.accum}", + dtype=root.symbols[info.accum], + transient=True, + find_new_name=True) + init_state = parent.add_state(loop.label + "_init", is_start_block=was_start) + red_state = parent.add_state(loop.label + "_reduce") + parent.add_edge(init_state, red_state, dace.InterstateEdge()) + seed = init_state.add_tasklet("seed", set(), {"_out"}, f"_out = {info.accum}") + init_state.add_edge(seed, "_out", init_state.add_write(tmp_name), None, + mm.Memlet(data=tmp_name, subset=subsets.Range([(0, 0, 1)]))) + entry = init_state + dest_name = tmp_name + dest_subset = subsets.Range([(0, 0, 1)]) + extra_assignments[info.accum] = tmp_name + + for e in in_edges: + parent.add_edge(e.src, entry, e.data) + for e in out_edges: + assigns = dict(e.data.assignments or {}) + assigns.update(extra_assignments) + cond = e.data.condition.as_string if e.data.condition is not None else "1" + parent.add_edge(red_state, e.dst, dace.InterstateEdge(condition=cond, assignments=assigns)) + parent.remove_node(loop) + + arr = red_state.add_read(info.array) + dst = red_state.add_write(dest_name) + red = red_state.add_reduce(info.wcr, axes=list(range(len(info.array_subset))), identity=None) + red_state.add_edge(arr, None, red, None, mm.Memlet(data=info.array, subset=info.array_subset)) + red_state.add_edge(red, None, dst, None, mm.Memlet(data=dest_name, subset=dest_subset)) diff --git a/dace/transformation/passes/pattern_matching.py b/dace/transformation/passes/pattern_matching.py index 9f557527f0..e3ac9f9187 100644 --- a/dace/transformation/passes/pattern_matching.py +++ b/dace/transformation/passes/pattern_matching.py @@ -13,9 +13,10 @@ from dace.sdfg.state import ControlFlowRegion import networkx as nx from networkx.algorithms import isomorphism as iso -from typing import Any, Callable, Dict, Iterable, Iterator, List, Optional, Set, Tuple, Type, Union +from typing import Any, Callable, Dict, Iterable, Iterator, List, Optional, Tuple, Type, Union from dace.sdfg.validation import InvalidSDFGError from dace.transformation import transformation as xf, pass_pipeline as ppl +from dace.utils import print_sdfg_hash @dataclass @@ -79,11 +80,8 @@ def __init__(self, self.print_report = print_report self.progress = progress - def depends_on(self) -> Set[Type[ppl.Pass]]: - result = set() - for p in self.transformations: - result.update(p.depends_on()) - return result + def depends_on(self) -> List[Union[Type[ppl.Pass], ppl.Pass]]: + return ppl.unique_dependencies(self.transformations) def modifies(self) -> ppl.Modifies: result = ppl.Modifies.Nothing @@ -188,6 +186,8 @@ def _apply_and_validate(self, match: xf.PatternTransformation, sdfg: SDFG, start print('Applied {}.\r'.format(', '.join(['%d %s' % (len(v), k) for k, v in applied_transformations.items()])), end='') + # TODO(tehrengruber): make this configurable + print_sdfg_hash(sdfg, frame_index=5, info=type(match).__name__) if self.validate_all: try: sdfg.validate() diff --git a/dace/transformation/passes/reference_reduction.py b/dace/transformation/passes/reference_reduction.py index 801ed21551..60fc07189b 100644 --- a/dace/transformation/passes/reference_reduction.py +++ b/dace/transformation/passes/reference_reduction.py @@ -6,7 +6,7 @@ from dace import SDFG, SDFGState, data, properties, Memlet from dace.sdfg import nodes from dace.transformation import pass_pipeline as ppl, transformation -from dace.transformation.helpers import all_isedges_between +from dace.transformation.helpers import modified_symbols_between from dace.transformation.passes import analysis as ap @@ -26,7 +26,7 @@ def should_reapply(self, modified: ppl.Modifies) -> bool: return modified & ppl.Modifies.AccessNodes def depends_on(self): - return {ap.StateReachability, ap.FindAccessStates, ap.FindReferenceSources} + return [ap.FindAccessStates, ap.FindReferenceSources] def apply_pass(self, sdfg: SDFG, pipeline_results: Dict[str, Any]) -> Optional[Set[str]]: """ @@ -38,7 +38,6 @@ def apply_pass(self, sdfg: SDFG, pipeline_results: Dict[str, Any]) -> Optional[S pipeline, an empty dictionary is expected. :return: A set of removed data descriptor names, or None if nothing changed. """ - reachable: Dict[SDFGState, Set[SDFGState]] = pipeline_results[ap.StateReachability.__name__][sdfg.cfg_id] access_states: Dict[str, Set[SDFGState]] = pipeline_results[ap.FindAccessStates.__name__][sdfg.cfg_id] reference_sources: Dict[str, Set[Memlet]] = pipeline_results[ap.FindReferenceSources.__name__][sdfg.cfg_id] @@ -52,7 +51,7 @@ def apply_pass(self, sdfg: SDFG, pipeline_results: Dict[str, Any]) -> Optional[S refsets = self.find_refsets(candidates, access_states) - result: Set[str] = self.find_candidates(sdfg, reference_sources, refsets, access_states, reachable) + result: Set[str] = self.find_candidates(sdfg, reference_sources, refsets, access_states) if not result: return None @@ -98,7 +97,6 @@ def find_candidates( reference_sources: Dict[str, Set[Memlet]], refsets: Dict[str, List[Tuple[SDFGState, nodes.AccessNode]]], access_states: Dict[str, Set[SDFGState]], - reachable_states: Dict[SDFGState, Set[SDFGState]], ) -> Set[str]: """ Returns a set of candidates for conversion to views. @@ -129,18 +127,13 @@ def find_candidates( if cand not in result: break - # Otherwise, they are only inter-state or free symbols. Test all paths to uses in different states - # NOTE: This is an expensive check! + # Otherwise, they are only inter-state or free symbols. Check whether any such symbol may be + # reassigned on a path to another use, even across sibling control-flow regions. for other_state in access_states[cand]: - # Filter self and unreachable states - if other_state is state or other_state not in reachable_states[state]: + if other_state is state: continue - for e in all_isedges_between(state, other_state): - # The symbol was modified/reassigned in one of the paths, skip - if fsyms & e.data.assignments.keys(): - result.remove(cand) - break - if cand not in result: + if fsyms & modified_symbols_between(state, other_state): + result.remove(cand) break return result diff --git a/dace/transformation/passes/scalar_fission.py b/dace/transformation/passes/scalar_fission.py index 8d88f2752b..76d2324f77 100644 --- a/dace/transformation/passes/scalar_fission.py +++ b/dace/transformation/passes/scalar_fission.py @@ -23,7 +23,7 @@ def should_reapply(self, modified: ppl.Modifies) -> bool: return modified & ppl.Modifies.AccessNodes def depends_on(self): - return {ap.ScalarWriteShadowScopes} + return [ap.ScalarWriteShadowScopes] def apply_pass(self, sdfg: SDFG, pipeline_results: Dict[str, Any]) -> Optional[Dict[str, Set[str]]]: """ diff --git a/dace/transformation/passes/scalar_to_symbol.py b/dace/transformation/passes/scalar_to_symbol.py index 693a4a7777..90594c1a86 100644 --- a/dace/transformation/passes/scalar_to_symbol.py +++ b/dace/transformation/passes/scalar_to_symbol.py @@ -21,7 +21,6 @@ from dace.sdfg import utils as sdutils from dace.sdfg.replace import replace_properties_dict from dace.sdfg.sdfg import InterstateEdge -from dace.sdfg.state import ConditionalBlock, LoopRegion from dace.transformation import helpers as xfh from dace.transformation import pass_pipeline as passes from dace.transformation.transformation import explicit_cf_compatible @@ -87,6 +86,8 @@ def find_promotable_scalars(sdfg: sd.SDFG, transients_only: bool = True, integer # General array checks for aname, desc in sdfg.arrays.items(): + if isinstance(desc, dt.DistributedDescriptor): + continue if isinstance(desc, (dt.View, dt.StructureView)): continue if (transients_only and not desc.transient) or isinstance(desc, dt.Stream): @@ -109,10 +110,6 @@ def find_promotable_scalars(sdfg: sd.SDFG, transients_only: bool = True, integer if candidate not in candidates: continue - # If candidate is read-only, continue normally - if state.in_degree(node) == 0: - continue - # If candidate is read by a library node, skip removed = False for oe in state.out_edges(node): @@ -127,6 +124,10 @@ def find_promotable_scalars(sdfg: sd.SDFG, transients_only: bool = True, integer continue # End of read check + # If candidate is read-only, continue normally + if state.in_degree(node) == 0: + continue + # Candidate may only be accessed in a top-level scope if state.entry_node(node) is not None: candidates.remove(candidate) diff --git a/dace/transformation/passes/simplification/continue_to_condition.py b/dace/transformation/passes/simplification/continue_to_condition.py index 62f08ba296..d5c2a4ac19 100644 --- a/dace/transformation/passes/simplification/continue_to_condition.py +++ b/dace/transformation/passes/simplification/continue_to_condition.py @@ -4,7 +4,7 @@ from dace import sdfg as sd, properties from dace.sdfg import utils as sdutil -from dace.sdfg.state import ContinueBlock, ControlFlowRegion, ConditionalBlock, LoopRegion +from dace.sdfg.state import ContinueBlock, ConditionalBlock, LoopRegion from dace.transformation import transformation from dace.transformation import pass_pipeline as ppl from dace.sdfg.sdfg import SDFG diff --git a/dace/transformation/passes/simplify.py b/dace/transformation/passes/simplify.py index 173647e2a8..0bbdab5a6b 100644 --- a/dace/transformation/passes/simplify.py +++ b/dace/transformation/passes/simplify.py @@ -1,4 +1,4 @@ -# Copyright 2019-2024 ETH Zurich and the DaCe authors. All rights reserved. +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. from dataclasses import dataclass from typing import Any, Dict, List, Optional, Set import warnings @@ -12,6 +12,7 @@ from dace.transformation.passes.dead_dataflow_elimination import DeadDataflowElimination from dace.transformation.passes.dead_state_elimination import DeadStateElimination from dace.transformation.passes.fusion_inline import FuseStates, InlineControlFlowRegions, InlineSDFGs +from dace.transformation.passes.lift_trivial_if import LiftTrivialIf from dace.transformation.passes.optional_arrays import OptionalArrayInference from dace.transformation.passes.scalar_to_symbol import ScalarToSymbolPromotion from dace.transformation.passes.prune_symbols import RemoveUnusedSymbols @@ -32,6 +33,7 @@ DeadDataflowElimination, DeadStateElimination, PruneEmptyConditionalBranches, + LiftTrivialIf, RemoveUnusedSymbols, ReferenceToView, ArrayElimination, diff --git a/dace/transformation/passes/split_tasklets.py b/dace/transformation/passes/split_tasklets.py index 14ea5df28e..dcf008fcba 100644 --- a/dace/transformation/passes/split_tasklets.py +++ b/dace/transformation/passes/split_tasklets.py @@ -9,8 +9,6 @@ import ast from dace.sdfg.nodes import CodeBlock -import re - class ASTSplitter: @@ -107,7 +105,7 @@ def should_reapply(self, modified: ppl.Modifies) -> bool: return modified & ppl.Modifies.Tasklets def depends_on(self): - return {} + return [] tmp_access_identifier = "_split_" diff --git a/dace/transformation/passes/symbol_propagation.py b/dace/transformation/passes/symbol_propagation.py index 694875ab7d..5a4856707b 100644 --- a/dace/transformation/passes/symbol_propagation.py +++ b/dace/transformation/passes/symbol_propagation.py @@ -12,7 +12,7 @@ from dace import SDFG, properties, SDFGState from typing import Any, Dict, Set, Optional from dace import data as dt -from dace.symbolic import pystr_to_symbolic +from dace.symbolic import pystr_to_symbolic, scalars @dataclass(unsafe_hash=True) @@ -98,14 +98,8 @@ def _get_in_syms( ]) } - # Also skip scalars - sym_table = { - k: v - for k, v in sym_table.items() if v is None or not any([ - str(s) in sdfg.arrays and isinstance(sdfg.arrays[str(s)], dt.Scalar) - for s in pystr_to_symbolic(v).free_symbols - ]) - } + # Also skip assignments that read a scalar (scalars cannot be propagated as symbols) + sym_table = {k: v for k, v in sym_table.items() if v is None or not scalars(v, sdfg.arrays)} # Combine the symbols if i == 0: diff --git a/dace/transformation/passes/symbol_ssa.py b/dace/transformation/passes/symbol_ssa.py index da0d1cdbb1..ce30981563 100644 --- a/dace/transformation/passes/symbol_ssa.py +++ b/dace/transformation/passes/symbol_ssa.py @@ -23,7 +23,7 @@ def should_reapply(self, modified: ppl.Modifies) -> bool: return modified & ppl.Modifies.Symbols | ppl.Modifies.Edges | ppl.Modifies.Nodes | ppl.Modifies.States def depends_on(self): - return {ap.SymbolWriteScopes} + return [ap.SymbolWriteScopes] def apply(self, region, pipeline_results) -> Optional[Dict[str, Set[str]]]: """ diff --git a/dace/transformation/subgraph/expansion.py b/dace/transformation/subgraph/expansion.py index 73445dba9a..9b946c92cb 100644 --- a/dace/transformation/subgraph/expansion.py +++ b/dace/transformation/subgraph/expansion.py @@ -5,7 +5,6 @@ from dace import dtypes, symbolic, subsets from dace.sdfg import nodes from dace.sdfg import replace, SDFG, dynamic_map_inputs -from dace.sdfg.graph import SubgraphView from dace.sdfg.state import SDFGState, StateSubgraphView from dace.transformation import transformation from dace.properties import make_properties, Property diff --git a/dace/transformation/subgraph/gpu_persistent_fusion.py b/dace/transformation/subgraph/gpu_persistent_fusion.py index e12ae94504..33758e4eea 100644 --- a/dace/transformation/subgraph/gpu_persistent_fusion.py +++ b/dace/transformation/subgraph/gpu_persistent_fusion.py @@ -5,7 +5,7 @@ from dace.sdfg import SDFG, SDFGState, InterstateEdge from dace.dtypes import StorageType, ScheduleType from dace.properties import Property, make_properties -from dace.sdfg.state import AbstractControlFlowRegion, LoopRegion +from dace.sdfg.state import AbstractControlFlowRegion from dace.sdfg.graph import SubgraphView from dace.transformation.transformation import SubgraphTransformation diff --git a/dace/transformation/subgraph/helpers.py b/dace/transformation/subgraph/helpers.py index 07e08226c0..59da7f4d11 100644 --- a/dace/transformation/subgraph/helpers.py +++ b/dace/transformation/subgraph/helpers.py @@ -5,7 +5,7 @@ from dace.sdfg.graph import SubgraphView import copy -from typing import List, Dict, Set +from typing import List, Dict # **************** # Helper functions diff --git a/dace/transformation/subgraph/stencil_tiling.py b/dace/transformation/subgraph/stencil_tiling.py index caaa367fdb..a1a57615f8 100644 --- a/dace/transformation/subgraph/stencil_tiling.py +++ b/dace/transformation/subgraph/stencil_tiling.py @@ -14,7 +14,6 @@ from dace.transformation.dataflow.map_expansion import MapExpansion from dace.transformation.dataflow.map_collapse import MapCollapse from dace.transformation.dataflow.strip_mining import StripMining -from dace.transformation.interstate.loop_detection import DetectLoop from dace.transformation.subgraph import SubgraphFusion from copy import deepcopy as dcpy diff --git a/dace/transformation/subgraph/subgraph_fusion.py b/dace/transformation/subgraph/subgraph_fusion.py index f360b51f1c..12d31fa515 100644 --- a/dace/transformation/subgraph/subgraph_fusion.py +++ b/dace/transformation/subgraph/subgraph_fusion.py @@ -18,8 +18,6 @@ from typing import List, Tuple import warnings -import dace.libraries.standard as stdlib - from collections import defaultdict from itertools import chain diff --git a/dace/transformation/testing.py b/dace/transformation/testing.py index ba88c8111c..d5c968322c 100644 --- a/dace/transformation/testing.py +++ b/dace/transformation/testing.py @@ -1,7 +1,6 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. import copy from io import StringIO -import os import sys import traceback diff --git a/dace/utils.py b/dace/utils.py index d56c50b5e6..d9a7828188 100644 --- a/dace/utils.py +++ b/dace/utils.py @@ -8,6 +8,7 @@ from typing import Iterable, Sequence, Union import sympy +import inspect # Type alias for numeric or symbolic values NumericType = Union[int, float, sympy.Basic] @@ -62,3 +63,15 @@ def until(val, substr): if substr not in val: return val return val[:val.find(substr)] + + +def json_data_to_file(data, filename): + import json + + with open(filename, "w") as f: + json.dump(data, f, indent=2) + + +def print_sdfg_hash(sdfg: "dace.SDFG", frame_index=1, info="") -> None: + frame = inspect.stack()[frame_index] + print(f"[SDFG Hash] location: {frame.filename}:{frame.lineno} hash: {sdfg.hash_sdfg()}, info: {info}", flush=True) diff --git a/dace/version.py b/dace/version.py index 309569bb14..748ceaf222 100644 --- a/dace/version.py +++ b/dace/version.py @@ -1 +1 @@ -__version__ = '43!2026.02.12' +__version__ = '2.0.0a4' diff --git a/doc/conf.py b/doc/conf.py index 067cea3d17..689d0e85ac 100644 --- a/doc/conf.py +++ b/doc/conf.py @@ -21,7 +21,7 @@ # -- Project information ----------------------------------------------------- project = 'DaCe' -copyright = '2019-2022, Scalable Parallel Computing Laboratory, ETH Zurich' +copyright = '2019-2026, Scalable Parallel Computing Laboratory, ETH Zurich' author = 'Scalable Parallel Computing Laboratory, ETH Zurich and the DaCe authors' # The full version, including alpha/beta/rc tags @@ -34,6 +34,37 @@ # ones. extensions = ['sphinx.ext.autodoc', 'sphinx_autodoc_typehints', 'sphinx.ext.mathjax', 'sphinx_rtd_theme'] +# sphinx_autodoc_typehints emits ``forward_reference`` warnings whenever it +# cannot resolve a forward reference (string annotation) in the *defining* +# module of an object. Sympy >= 1.12 carries ``ClassVar[...]`` annotations as +# strings on classes like ``sympy.Function`` / ``sympy.Number`` without +# importing ``ClassVar`` into those modules' globals. Every DaCe class that +# subclasses sympy classes (``int_floor``, ``int_ceil``, ``OR``, ``AND``, +# ``IfExpr``, ``dace.dtypes.Float``, ...) inherits the broken annotation and +# triggers the warning. The signature ``name 'ClassVar' is not defined`` is +# unique to this upstream sympy issue, so filter only those records and let +# any genuine DaCe-side forward-reference warnings through. +import logging as _logging + + +class _SympyClassVarForwardRefFilter(_logging.Filter): + _signature = "name 'ClassVar' is not defined" + + def filter(self, record: _logging.LogRecord) -> bool: + try: + msg = record.getMessage() + except Exception: + return True + return self._signature not in msg + + +# Attach the filter to both the sphinx logger tree (where Sphinx actually +# emits the formatted warning) and the extension's own logger (belt and +# braces, since sphinx_autodoc_typehints may log through either). +_filter = _SympyClassVarForwardRefFilter() +for _name in ('sphinx', 'sphinx_autodoc_typehints'): + _logging.getLogger(_name).addFilter(_filter) + # Add any paths that contain templates here, relative to this directory. templates_path = ['_templates'] diff --git a/doc/extensions/backend.rst b/doc/extensions/backend.rst new file mode 100644 index 0000000000..e24ee11ae3 --- /dev/null +++ b/doc/extensions/backend.rst @@ -0,0 +1,218 @@ +.. _new-backend: + +New Backends +============ + +A *backend* in DaCe is a code generator that lowers some subset of the SDFG IR +to source code in a specific language and emits the build artifacts needed to +compile that source. The built-in backends live under +:mod:`dace.codegen.targets` (CPU, CUDA, MPI, SVE, MLIR, +Snitch, ...). Adding a new backend means plugging into three points of the +codegen pipeline: the user-facing enumerations, the dispatcher, and the +build/link environment. + +This page covers the plumbing. For a worked example, see the +`Extending the Code Generator +`_ +tutorial, which builds a small custom target end to end. The +``dace/codegen/targets/cpu.py`` and ``dace/codegen/targets/cuda.py`` files are +the canonical references for any non-trivial backend. + +.. _enums: + +Extensible Enumerations +----------------------- + +Most DaCe enumerations - :class:`~dace.dtypes.DeviceType`, +:class:`~dace.dtypes.ScheduleType`, :class:`~dace.dtypes.StorageType`, +:class:`~dace.dtypes.InstrumentationType`, +:class:`~dace.dtypes.DataInstrumentationType`, and +:class:`~dace.dtypes.Language`, among others - are *user-extensible*. +They subclass :class:`dace.attr_enum.ExtensibleAttributeEnum`, which is a +drop-in extension of the standard library's :class:`enum.Enum` that adds +two capabilities a backend frequently needs: + +1. **New members can be registered at runtime**, with either a plain + value or a frozen dataclass acting as a *template* for attributed + instances. +2. **Members can carry attributes** (via dataclass templates), which is + what lets ``ScheduleType.GPU_ThreadBlock`` and friends remain symbolic + while richer entries (e.g., a vendor-specific schedule with tunable + parameters) carry the parameters on the enum value itself. + +Adding plain entries +~~~~~~~~~~~~~~~~~~~~ + +The simplest case is adding a new constant. Use the +:meth:`~dace.attr_enum.ExtensibleAttributeEnum.register` classmethod from +your backend's setup code (typically the package's ``__init__.py``): + +.. code-block:: python + + from dace import dtypes + + # Auto-assigned value (analogous to `auto()`): + dtypes.DeviceType.register('MyAccelerator') + dtypes.ScheduleType.register('MyAccelerator_Device') + dtypes.ScheduleType.register('MyAccelerator_ThreadBlock') + dtypes.StorageType.register('MyAccelerator_Global') + + # Explicit value (pass it as the second argument): + dtypes.StorageType.register('MyAccelerator_Shared', 301) + +After registration, the new entries are reachable as +``dtypes.ScheduleType.MyAccelerator_Device`` and serialize cleanly in +saved SDFGs - the SDFG (de)serializer round-trips any +``ExtensibleAttributeEnum`` member through the registry. + +Adding attributed entries (templates) +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +When a backend needs to attach data to an enum member - for example, the +parameters of a vendor-specific schedule - register a *template*. A +template is just a frozen dataclass; calling the template produces an +instance of the enum that carries the supplied field values: + +.. code-block:: python + + from dataclasses import dataclass + from dace import dtypes + + @dataclass(frozen=True) + class MyAcceleratorKernel: + block_size: int + smem_kb: int + + dtypes.ScheduleType.register_template('MyAccelerator_Kernel', + MyAcceleratorKernel) + + # Use the template like any other enum member: + sched = dtypes.ScheduleType.MyAccelerator_Kernel(block_size=128, + smem_kb=48) + sched.block_size # -> 128 + +Two design points worth noting (mirrored in +:mod:`tests.utils.attrenum_test`): + +* The *template* compares equal to any of its instances, which makes + attributed members ergonomic in ``match`` / ``case`` blocks + (``case dtypes.ScheduleType.MyAccelerator_Kernel:``). +* Two instances created with the same arguments are interned to the + same object (``sched is dtypes.ScheduleType.MyAccelerator_Kernel(block_size=128, smem_kb=48)``), + which keeps SDFG hashing and identity comparisons predictable. + +Templates can also be declared inline at class-definition time when you +own the enum class, by writing ``@dataclass(frozen=True)`` classes +inside the enum body; this is the form used in the test suite. From a +backend extension's perspective, ``register_template`` is the equivalent +runtime form. + +Optional: ``undefined_safe_enum`` +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +The decorator :func:`dace.registry.undefined_safe_enum` is an *optional* +helper that registers a single extra member named ``Undefined`` on an +``ExtensibleAttributeEnum``. The built-in DaCe enums use it so that +deserialization, partially-specified IR, and forward-compatibility paths +have a well-known sentinel to fall back to instead of raising on an +unknown enum value: + +.. code-block:: python + + from dace.attr_enum import ExtensibleAttributeEnum + from dace.registry import undefined_safe_enum + from enum import auto + + @undefined_safe_enum + class MyEnum(ExtensibleAttributeEnum): + First = auto() + Second = auto() + + MyEnum.Undefined # added by the decorator + +Using the decorator is purely a convenience - if you do not need an +``Undefined`` sentinel for your own enum, omit it. Custom enums that +extend the built-in ones (via ``register`` / ``register_template``) +inherit the sentinel from their parent class and do not need to apply +the decorator themselves. + +Implementing a Target Code Generator +------------------------------------ + +A backend is a subclass of +:class:`~dace.codegen.targets.target.TargetCodeGenerator` registered through +:func:`@dace.registry.autoregister_params`. The class is instantiated once +per SDFG by the framecode driver and is responsible for: + +* declaring which schedules, storage types, and node patterns it can handle; +* emitting the source files to be compiled (returning them as + ``CodeObject`` instances from :meth:`get_generated_codeobjects`); +* providing the per-node, per-state, per-copy, and per-allocation + generators that the dispatcher will call back into. + +Inside the constructor, the backend registers itself with the +:class:`~dace.codegen.dispatcher.TargetDispatcher` for the events it wants +to handle: + +.. code-block:: python + + class MyTargetCodeGen(target.TargetCodeGenerator): + def __init__(self, frame, sdfg): + self._frame = frame + self._dispatcher = frame.dispatcher + + # Maps with these schedules will go through `generate_node`/_state. + self._dispatcher.register_map_dispatcher( + [dtypes.ScheduleType.MyAccelerator_Device], self) + + # Arrays in this storage will go through this target's + # allocation/deallocation hooks. + self._dispatcher.register_array_dispatcher( + dtypes.StorageType.MyAccelerator_Global, self) + + # Copies between these storages will be lowered by this target. + self._dispatcher.register_copy_dispatcher( + dtypes.StorageType.CPU_Heap, + dtypes.StorageType.MyAccelerator_Global, + None, + self) + +The full list of registration entry points is documented on +:class:`~dace.codegen.dispatcher.TargetDispatcher`: + +* :meth:`~dace.codegen.dispatcher.TargetDispatcher.register_state_dispatcher` +* :meth:`~dace.codegen.dispatcher.TargetDispatcher.register_node_dispatcher` +* :meth:`~dace.codegen.dispatcher.TargetDispatcher.register_map_dispatcher` +* :meth:`~dace.codegen.dispatcher.TargetDispatcher.register_array_dispatcher` +* :meth:`~dace.codegen.dispatcher.TargetDispatcher.register_copy_dispatcher` + +Each ``register_*`` call accepts an optional predicate so that a target can +opt in or out based on richer criteria than just an enum value. + +The target then implements the corresponding callbacks +(``generate_state``, ``generate_node``, ``allocate_array``, +``deallocate_array``, ``copy_memory``, ...). Each callback receives the +relevant SDFG, state, node(s), and a pair of :class:`CodeIOStream` objects +into which to emit code at the local (current scope) and global +(translation-unit-level) positions. + +Build environments and linked libraries +--------------------------------------- + +A backend almost always needs to add headers, runtime support files, and +link-time dependencies to the build. The recommended way to declare these is +via :ref:`environments `: define a ``@dace.library.environment`` +that lists the ``cmake_packages``, ``headers``, ``cmake_libraries``, and +runtime ``state_fields`` your generated code requires, and add it to the +``environments`` list of the relevant code objects. + +For more invasive integration (e.g., additional source +files generated alongside the SDFG), the ``CodeObject`` returned by the +target can declare extra ``additional_files`` and target-specific +``additional_compiler_flags``. The CMake driver in +:mod:`dace.codegen.compiler` consumes those fields when building the shared +library. + +Once registered, the backend is selected automatically whenever an SDFG +contains nodes/arrays/maps with the schedules, storage types, or node +classes the target claimed. No changes to user-facing APIs are required. diff --git a/doc/extensions/extensions.rst b/doc/extensions/extensions.rst index 6a51da3f72..f0feef98c2 100644 --- a/doc/extensions/extensions.rst +++ b/doc/extensions/extensions.rst @@ -4,7 +4,7 @@ Extending DaCe The DaCe framework is designed to be user-extensible. Developers can externally extend DaCe by adding new types of nodes, frontend languages, transformations, and hardware architecture backends. -The three key mechanisms of extensibility are class inheritance, :ref:`replacements `, and :ref:`enums`. +The three key mechanisms of extensibility are class inheritance, replacements, and enumerations. * Class inheritance is used for new :ref:`library nodes `, :ref:`data types `, :ref:`transformations `, :ref:`passes `, :ref:`code generator targets `, @@ -22,11 +22,12 @@ For more examples of how to extend DaCe, see the following resources: * Code generators: `Extending the Code Generator `_ * Frontend extensions (enumerations and replacements): `Tensor Core code sample `_ -.. .. toctree -.. :maxdepth: 1 +.. toctree:: + :maxdepth: 1 -.. .. symbolic -.. .. libraries -.. .. frontend -.. .. sdfgconvertible -.. .. backend + symbolic + libraries + frontend + sdfgconvertible + backend + instrumentation diff --git a/doc/extensions/frontend.rst b/doc/extensions/frontend.rst new file mode 100644 index 0000000000..6e93fc8b8f --- /dev/null +++ b/doc/extensions/frontend.rst @@ -0,0 +1,129 @@ +.. _new-frontend: + +Writing a New Frontend +====================== + +A frontend in DaCe is the component that translates a source program (Python, +Fortran, MLIR, ...) into an SDFG. Existing frontends are organized under +:mod:`dace.frontend` and serve as references when adding a new one. The +goal of this page is to capture the design conventions that have made those +frontends robust enough to ingest large real-world codebases. + +The general advice that runs through every section below is: **stay +data-centric**. The SDFG IR is good at expressing parallelism, locality, and +data movement; it is not designed to express pointer arithmetic, exceptional +control flow, or escaping closures. A frontend's job is to massage the input +language until it "speaks" in those terms - without prematurely picking a +schedule - that is what :ref:`transformations ` are for. + +Trust the transformations +------------------------- + +Resist the urge to optimize while parsing. Producing a *correct, naive* SDFG +that places statements in distinct states, materializes intermediate buffers, +and uses straightforward maps is almost always the right thing to do; the +:ref:`simplification pipeline ` will fuse states, eliminate +redundant copies, and inline nested SDFGs. Premature fusion or reordering at +parse time tends to interact badly with later passes and makes the frontend +harder to maintain. + +When constructing IR objects, populate ``debuginfo`` with the original source +location. This is much cheaper than reconstructing it later and lets users +correlate generated nodes back to their source. The Python frontend uses +Python AST location information for this purpose; new frontends are encouraged +to follow the same pattern. + +Make things data-centric +------------------------ + +When you encounter a construct that does not map directly to an SDFG node: + +* **Preprocess the input AST** before lowering. The Python frontend uses an + AST-to-AST rewrite pipeline (see :doc:`/frontend/preprocessing`) to desugar + syntactic elements, inline context managers, unroll metaprogramming loops, + and fold compile-time constants. New frontends should prefer this strategy + over teaching the IR builder about every possible source-language quirk. +* **Expose memory access patterns** so that DaCe can reason about them. + Concrete index expressions (``A[i, j]``) are far more useful than opaque + pointer arithmetic. If the source language has loop bounds or array shapes + parameterized by symbols, propagate those symbols through to the SDFG + rather than substituting them with concrete values. +* **Avoid raw pointers** where possible. The recommended progression is: + arrays first, then :class:`~dace.data.View` for slices and aliases, then + (only as a last resort) :class:`~dace.data.Reference` for restricted + pointer semantics. + +Encapsulation, callbacks, and closures +-------------------------------------- + +Most languages have constructs that the SDFG cannot express directly: I/O, +random number generation, calls into the host runtime, exceptions, etc. The +established pattern is to expose these as **Python callbacks** - opaque +function invocations whose body lives in user-provided Python code - and to +thread a per-program ``__state`` value through the call chain so that +internal state is not lost across calls. + +Some practical guidelines: + +* Use a single ``__state`` value per stateful resource (a library, a random + generator, an MPI communicator, ...) so that the SDFG framework can track + data dependencies on it. Without it, code-generator passes are free to + reorder calls in ways that violate the source language's semantics. +* If you can prove that two stateful subsystems do not interact, model them + with separate state values. This unlocks more parallelism downstream. +* Do not lose the closure! When a source-language function captures state + from its lexical environment, the frontend must arrange for that state to + flow into the SDFG (typically as additional arguments or as scalar/array + members of a closure object). The Python frontend's + :class:`~dace.frontend.python.parser.SDFGClosure` is a useful reference + for what such an object looks like. + +Know your language +------------------ + +A robust frontend accounts for **every** AST node it might encounter, even +if only by emitting an explicit error. Silent fall-through to "treat this as +a no-op" is the source of most frontend bugs. Maintain an explicit list of +disallowed constructs (see ``DISALLOWED_STMTS`` in +:mod:`dace.frontend.python.newast` for the Python frontend's variant) and +fail fast when one shows up. + +When the source language has well-known operators or library calls with a +matching DaCe library node, lower directly to that library node rather than +synthesizing the operation manually. For example, the Python frontend lowers +``@`` and ``numpy.matmul`` to the BLAS ``MatMul`` library node, which gives +downstream optimization a hook to swap in vendor-tuned implementations. + +Assumptions +----------- + +Two assumptions are particularly important to surface explicitly: + +* **Aliasing.** If you have whole-program analysis, you can compute precise + no-alias information per call. Otherwise, mark potentially-aliasing + arrays with the :attr:`~dace.data.Data.may_alias` property; the codegen + honors this when emitting array declarations and the optimizer treats it + as a barrier. +* **By-reference vs. by-value.** Source languages that pass aggregates by + reference should be lowered using views; languages that pass them by + value should produce explicit copies into transients. References are + available for cases where neither view nor copy is appropriate, but they + disable some optimizations. + +Limit the scope +--------------- + +Try to build the frontend as a per-scope (per-function, per-translation-unit) +parser. Each scope produces its own SDFG, and inter-scope calls become +nested SDFGs. This mirrors the Python frontend's structure: one +:class:`~dace.frontend.python.newast.ProgramVisitor` per ``@dace.program``, +parented by a single closure. It has two benefits: + +* the IR for each scope stays small and easy to debug; and +* nested SDFGs are first-class citizens of the IR and can be inlined, + duplicated, or specialized by transformations. + +For function calls within a scope, prefer creating a nested SDFG over +inlining the body verbatim. The :class:`~dace.transformation.interstate.sdfg_nesting.InlineSDFG` +transformation will inline at the user's request, but reversing an +inlining is much harder. diff --git a/doc/extensions/instrumentation.rst b/doc/extensions/instrumentation.rst new file mode 100644 index 0000000000..34c4d5ba06 --- /dev/null +++ b/doc/extensions/instrumentation.rst @@ -0,0 +1,113 @@ +.. _new-instrumentation: + +Writing a Custom Instrumentation Provider +========================================= + +Instrumentation providers are the mechanism through which DaCe's code +generator inserts measurement code (timers, hardware counters, NVTX +ranges, ...) around SDFG elements at compile time. The built-in providers +live under :mod:`dace.codegen.instrumentation`; users can add new ones +without modifying DaCe by subclassing +:class:`~dace.codegen.instrumentation.provider.InstrumentationProvider` and +registering the class with :class:`~dace.dtypes.InstrumentationType`. + +When does a provider run? +------------------------- + +Every SDFG element has an ``instrument`` (or ``instrument_condition``) +property that takes an :class:`~dace.dtypes.InstrumentationType`. During +code generation, the framecode driver looks up the provider associated with +that enum value and calls back into it at well-defined points around the +element being lowered. A single provider may handle many elements - the +:class:`~dace.codegen.instrumentation.timer.TimerProvider`, for example, +emits start/stop calls around states, scopes, and full SDFGs. + +Skeleton of a provider +---------------------- + +A minimal provider extends one or more of the ``on_*`` callbacks defined on +the base class. Every callback receives the relevant SDFG/state/node along +with a *local* and *global* :class:`~dace.codegen.prettycode.CodeIOStream` +into which to emit code. The local stream injects code at the current +position in the generated function; the global stream injects code at the +top of the translation unit (useful for ``#include`` directives, helper +functions, and global state). + +.. code-block:: python + + from dace import dtypes, registry + from dace.codegen.instrumentation.provider import InstrumentationProvider + + + @registry.autoregister_params(type=dtypes.InstrumentationType.MyTimer) + class MyTimerProvider(InstrumentationProvider): + """Example provider that brackets every state with a custom timer.""" + + def on_sdfg_begin(self, sdfg, local_stream, global_stream, codegen): + global_stream.write('#include "my_timer.h"', sdfg) + + def on_state_begin(self, sdfg, cfg, state, local_stream, global_stream): + local_stream.write(f'my_timer_start("{state.label}");', sdfg, state.block_id) + + def on_state_end(self, sdfg, cfg, state, local_stream, global_stream): + local_stream.write(f'my_timer_stop("{state.label}");', sdfg, state.block_id) + +The decorator :func:`@registry.autoregister_params ` +adds ``MyTimer`` to :class:`~dace.dtypes.InstrumentationType` so that any +SDFG element can opt in via ``element.instrument = dtypes.InstrumentationType.MyTimer`` +or by setting the ``instrument`` property in the SDFG editor. + +Available hooks +--------------- + +The full set of callbacks is documented on +:class:`~dace.codegen.instrumentation.provider.InstrumentationProvider`. +The most commonly overridden ones are: + +* ``on_sdfg_begin`` / ``on_sdfg_end`` - bracket the entire program. Use + the global stream to emit ``#include``\ s and helper definitions. +* ``on_sdfg_init_begin`` / ``on_sdfg_init_end`` - run inside the + generated ``__dace_init`` function (state allocation, handle creation). +* ``on_sdfg_exit_begin`` / ``on_sdfg_exit_end`` - run inside + ``__dace_exit`` (state teardown). +* ``on_state_begin`` / ``on_state_end`` - bracket each state. +* ``on_scope_entry`` / ``on_scope_exit`` - bracket a map or consume scope. +* ``on_node_begin`` / ``on_node_end`` - bracket an individual node. +* ``on_copy_begin`` / ``on_copy_end`` - bracket a memory copy. +* ``on_allocation_begin`` / ``on_allocation_end`` - around an allocation + performed by the codegen. +* ``on_deallocation_begin`` / ``on_deallocation_end`` - around a + deallocation. + +Each callback's default implementation is a no-op, so a provider only +needs to override the events it cares about. + +Producing instrumentation reports +--------------------------------- + +The :class:`~dace.codegen.instrumentation.report.InstrumentationReport` +class is the standard container for runtime measurements: it deserializes +JSON written by the runtime, and the SDFV viewer can render it as overlays +on the SDFG. If your provider writes its results in the same JSON format +(see :func:`~dace.codegen.instrumentation.report.InstrumentationReport.from_file` +for the schema), the resulting reports integrate transparently into the +existing tooling. Providers that need a custom report layout can return an +``InstrumentationReport`` subclass; the only requirement is that +``durations`` and ``counters`` are accessible per SDFG/state/node so that +overlays render correctly. + +Built-in providers +------------------ + +The built-in providers are concise and instructive examples of the +technique: + +* :mod:`dace.codegen.instrumentation.timer` - host-side timers using + ``std::chrono``. Demonstrates state and scope-level brackets. +* :mod:`dace.codegen.instrumentation.gpu_events` - CUDA/HIP event-based + timers. Demonstrates how to interact with target-specific runtime APIs. +* :mod:`dace.codegen.instrumentation.papi` - PAPI hardware-counter + integration. Demonstrates how to thread per-thread state through the + generated runtime struct via ``state_fields``. +* :mod:`dace.codegen.instrumentation.likwid` - LIKWID region markers, + another good cross-platform reference. diff --git a/doc/extensions/libraries.rst b/doc/extensions/libraries.rst new file mode 100644 index 0000000000..bf00184c74 --- /dev/null +++ b/doc/extensions/libraries.rst @@ -0,0 +1,393 @@ +.. _libraries: + +Libraries and Environments +========================== + +A *library* in DaCe is a collection of :ref:`library node types `, their +expansions (implementations), library-specific transformations, and *environments* +that describe how to compile and link the resulting code against external +dependencies (e.g., BLAS, MPI, cuBLAS). Libraries are the primary mechanism +through which DaCe integrates high-performance third-party code into generated +SDFGs while keeping the IR portable and target-independent. + +Libraries live under :mod:`dace.libraries` and are registered with DaCe at import +time. Each library is a regular Python package whose ``__init__.py`` calls +:func:`dace.library.register_library`. External libraries can be developed in +their own Python packages and registered the same way - they do not need to +live inside the DaCe source tree. + +Defining a Library +------------------ + +The minimal skeleton of a library package looks like this:: + + mylib/ + __init__.py + nodes/ + __init__.py + my_node.py + environments/ + __init__.py + my_env.py + +The package's ``__init__.py`` imports the library nodes and environments and +then registers itself with DaCe: + +.. code-block:: python + + # mylib/__init__.py + from dace.library import register_library + from .nodes import * + from .environments import * + + register_library(__name__, "mylib") + + # Optional: set a default expansion implementation for all nodes in this + # library. Users can override this at runtime by assigning to + # ``mylib.default_implementation``. + default_implementation = "pure" + +Once registered, the library can be discovered via +:func:`dace.library.get_library` and any of its :ref:`library nodes ` +can be expanded using one of its registered implementations. + +Library Nodes and Expansions +~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Library nodes are defined by subclassing :class:`~dace.sdfg.nodes.LibraryNode` +and decorating the class with :func:`@dace.library.node `. +Each node declares an ``implementations`` dictionary that maps an +implementation name (string) to an +:class:`~dace.transformation.transformation.ExpandTransformation` subclass, and +a ``default_implementation`` attribute. Expansions are decorated with +:func:`@dace.library.expansion ` (or registered with +:func:`@dace.library.register_expansion `), +and must declare an ``environments`` list that names the environments the +expansion's generated code depends on. + +See :ref:`libnodes` for a full example (the ``Einsum`` library node) and a +description of how expansion is performed during compilation. + +The ``"pure"`` Implementation +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +By convention, most built-in library nodes provide an expansion named +``"pure"``. A *pure* implementation is a native SDFG expansion that lowers the +library node into ordinary DaCe constructs (maps, tasklets, nested SDFGs) and +declares no external environments. Because it relies only on the DaCe runtime +and the target's standard compiler, the ``"pure"`` implementation is the most +portable - it works on every backend, requires no third-party libraries, and +can be transformed and analyzed like any other SDFG subgraph - but it is +typically slower than vendor-tuned alternatives such as ``"MKL"``, +``"cuBLAS"``, or ``"OpenBLAS"``. When defining a new library node, providing a +``"pure"`` expansion is recommended so that the node always has a working +fallback regardless of the user's environment. + +.. _replacements: + +Registering a Library Call with the Python Frontend +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Defining a library node makes the node available in the SDFG IR, but it does +not by itself make the node callable from a ``@dace.program``. To bridge a +Python-level call site (e.g., ``numpy.dot(a, b)``) to a library node, the +Python frontend uses a *replacement registry* implemented in +:mod:`dace.frontend.common.op_repository`. Each replacement is a Python +function that receives an SDFG and a state and is responsible for adding +nodes (typically a library node, but it can be any subgraph) into that state +and returning the names of the data containers that hold the result. The +frontend invokes these functions whenever it encounters a matching call, +attribute access, operator, or NumPy ufunc in the user's program. + +Replacements live under :mod:`dace.frontend.python.replacements` (one module +per category - ``pymath``, ``linalg``, ``mpi``, ``ufunc``, ``operators``, +etc.) and are registered with one of the decorators exported from +``dace.frontend.common.op_repository``: + +* :func:`@oprepo.replaces(qualname) ` - + replaces a free-standing function call. ``qualname`` is the + fully-qualified, pydoc-compliant name of the function being intercepted + (e.g., ``"numpy.dot"``, ``"math.exp"``, ``"dace.comm.Bcast"``). A single + replacement may be decorated multiple times to cover several aliases. +* :func:`@oprepo.replaces_method(classname, method_name) ` - + replaces a method invocation on an object whose class name matches + ``classname`` (e.g., ``("Intracomm", "Bcast")`` for ``mpi4py``'s + ``Intracomm.Bcast``). +* :func:`@oprepo.replaces_attribute(classname, attr_name) ` - + replaces an attribute access (e.g., ``ndarray.shape``). +* :func:`@oprepo.replaces_operator(classname, optype, otherclass=None) ` - + replaces a binary or unary operator (``optype`` is the corresponding + ``ast`` node name, such as ``"Add"``, ``"Mul"``, or ``"Eq"``) between + two DaCe data classes. +* :func:`@oprepo.replaces_ufunc(name) ` - + replaces a NumPy universal function or one of its methods (``"ufunc"``, + ``"reduce"``, ``"accumulate"``, ...). + +Every replacement function has the same general signature:: + + def my_replacement(pv: ProgramVisitor, sdfg: SDFG, state: SDFGState, + *args, **kwargs) -> Union[str, Tuple[str, ...]]: + ... + +* ``pv`` is the active :class:`~dace.frontend.python.newast.ProgramVisitor`, + which can be queried for the current scope, the user-provided globals, and + utilities such as ``pv.get_target_name()`` for picking unique container + names. +* ``sdfg`` and ``state`` are the SDFG and state that the call site is being + parsed into. +* ``*args`` / ``**kwargs`` mirror the call-site arguments. Array arguments + are passed as the **string name** of an existing data container in + ``sdfg.arrays``; constants and shapes are passed as Python values or + symbolic expressions. +* The return value is the name (or tuple of names) of the data container(s) + that hold the result. The frontend takes care of binding these to the + user's left-hand-side variable. + +A typical replacement does three things: validate inputs, allocate output +containers, and wire a library node into the state. The example below is a +condensed version of the ``numpy.dot`` replacement for the 1-D case +(`see full implementation +`_), +which dispatches to the :class:`~dace.libraries.blas.nodes.dot.Dot` library +node: + +.. code-block:: python + + from dace.frontend.common import op_repository as oprepo + from dace.frontend.python.replacements.utils import ProgramVisitor + from dace import SDFG, SDFGState, Memlet + from dace.libraries.blas.nodes.dot import Dot + + @oprepo.replaces('dace.dot') + @oprepo.replaces('numpy.dot') + def dot(pv: ProgramVisitor, sdfg: SDFG, state: SDFGState, + op_a: str, op_b: str, op_out=None): + arr_a = sdfg.arrays[op_a] + arr_b = sdfg.arrays[op_b] + + # Allocate a transient scalar to hold the result if not provided. + if op_out is None: + op_out, _ = sdfg.add_scalar(pv.get_target_name(), arr_a.dtype, + transient=True, find_new_name=True) + arr_out = sdfg.arrays[op_out] + + # Insert the BLAS Dot library node and connect it to the inputs/outputs. + node = Dot('_Dot_') + state.add_node(node) + state.add_edge(state.add_read(op_a), None, node, '_x', + Memlet.from_array(op_a, arr_a)) + state.add_edge(state.add_read(op_b), None, node, '_y', + Memlet.from_array(op_b, arr_b)) + state.add_edge(node, '_result', state.add_write(op_out), None, + Memlet.from_array(op_out, arr_out)) + + return op_out + +With this registration in place, both ``dace.dot(a, b)`` and ``numpy.dot(a, b)`` +inside a ``@dace.program`` body are lowered into a ``Dot`` library node. The +choice of expansion (``"pure"``, ``"MKL"``, ``"cuBLAS"``, ...) is then made +later, during SDFG simplification or compilation, exactly as described in the +sections above. + +For a replacement to be picked up, its module must be imported during +``dace`` initialization. The built-in replacement modules are wired up via +``dace/frontend/python/replacements/__init__.py``; external replacements only +need to be imported once (e.g., from your library's top-level ``__init__.py``) +before the first ``@dace.program`` is parsed. See the +:mod:`dace.frontend.python.replacements` package for many more concrete +examples covering array creation, reductions, FFTs, MPI collectives, and +linear algebra. + +Environments +------------ + +An *environment* is a declarative description of an external dependency that +generated code may rely on. Environments tell the DaCe code generator and the +CMake build system which packages to find, which headers to include, which +libraries to link against, and which initialization or finalization code to +emit. They isolate platform- and library-specific build logic from the +expansion code, so the same expansion can target multiple backends simply by +changing its ``environments`` list. + +An environment is a Python class decorated with +:func:`@dace.library.environment `. The decorator +verifies that the class implements every required field listed below and +registers it in DaCe's global environment registry, where it can be looked up +by fully qualified class path with :func:`dace.library.get_environment`. + +Each environment **must** define every one of the following attributes (use an +empty list, empty dict, or empty string when a field does not apply): + +* ``cmake_minimum_version`` - minimum required CMake version, or ``None``. +* ``cmake_packages`` - list of CMake packages to ``find_package`` + (e.g., ``["BLAS"]``, ``["MPI"]``). +* ``cmake_variables`` - dict of CMake variables to set before package lookup + (e.g., ``{"BLA_VENDOR": "Intel10_64lp"}``). +* ``cmake_includes`` - list of additional include directories. +* ``cmake_libraries`` - list of libraries to link against. +* ``cmake_compile_flags`` - list of extra compile flags. +* ``cmake_link_flags`` - list of extra linker flags. +* ``cmake_files`` - list of additional CMake files (without the ``.cmake`` + suffix) to be included by the generated build script. See + :ref:`custom-cmake` below. +* ``headers`` - list of C/C++ headers that expansions of this environment + emit ``#include`` directives for. May be a flat list, or a dict keyed by + code-generator target (e.g., ``{'frame': [...], 'cuda': [...]}``) when + different headers are needed per backend. +* ``state_fields`` - list of C/C++ field declarations to add to the generated + SDFG runtime state struct (e.g., ``["cublasHandle_t cublas_handle;"]``). +* ``init_code`` - C/C++ snippet emitted during program initialization + (typically used to populate ``state_fields``). +* ``finalize_code`` - C/C++ snippet emitted during program teardown. +* ``dependencies`` - list of other environment classes that this environment + depends on; DaCe topologically sorts these when resolving the full + dependency graph for a compiled SDFG. + +Any of the ``cmake_*`` and ``headers`` attributes can also be defined as +``@staticmethod``\ s instead of plain class attributes. This is useful when +their values must be computed at runtime, for example to locate a vendor SDK +via environment variables. The Intel MKL environment uses this pattern to +locate ``mkl.h`` and ``libmkl_rt`` from ``MKLROOT`` or a Conda prefix: + +.. code-block:: python + + import os + import dace.library + + @dace.library.environment + class IntelMKL: + """Intel Math Kernel Library environment.""" + + cmake_minimum_version = None + cmake_packages = [] + cmake_variables = {"BLA_VENDOR": "Intel10_64lp"} + cmake_compile_flags = [] + cmake_link_flags = [] + cmake_files = [] + + headers = ["mkl.h", "../include/dace_blas.h"] + state_fields = [] + init_code = "" + finalize_code = "" + dependencies = [] + + @staticmethod + def cmake_includes(): + if 'MKLROOT' in os.environ: + return [os.path.join(os.environ['MKLROOT'], 'include')] + return [] + + @staticmethod + def cmake_libraries(): + ... + +An expansion opts into an environment by listing it on its +``environments`` attribute: + +.. code-block:: python + + from dace.libraries.blas.environments import IntelMKL + + @dace.library.register_expansion(Gemm, "MKL") + class ExpandGemmMKL(xf.ExpandTransformation): + environments = [IntelMKL] + + @staticmethod + def expansion(node, parent_state, parent_sdfg): + ... + +When an SDFG that uses this expansion is compiled, DaCe collects every +environment referenced by every expanded library node, resolves their +``dependencies`` graph, and merges all CMake fragments into the generated +build script. + +.. _custom-cmake: + +Custom CMake Files +------------------ + +For dependencies that cannot be expressed by simply calling ``find_package`` +and listing libraries, an environment may ship one or more custom CMake files +and reference them through the ``cmake_files`` attribute. This mechanism is +useful when integrating a vendor SDK that requires non-trivial discovery +logic, custom compile/link rules, or auto-generated source files. + +The procedure is: + +1. Place one or more ``.cmake`` files next to the environment's Python + module (or anywhere within the library package). +2. List the files in ``cmake_files``, **without** the ``.cmake`` extension and + **without** a directory prefix - DaCe locates them relative to the file + that defined the environment (the path is captured by the + :func:`@dace.library.environment ` decorator + from ``inspect.getmodule(env).__file__``). +3. Inside the custom CMake file, define any variables, targets, or commands + required, and use the standard CMake ``include_directories``, + ``target_link_libraries``, etc., as needed. The generated SDFG build + script will ``include()`` your file before linking the program. + +For example, given the layout:: + + mylib/ + environments/ + my_env.py + FindMyDep.cmake + +the environment can be declared as: + +.. code-block:: python + + @dace.library.environment + class MyDep: + cmake_minimum_version = "3.15" + cmake_packages = ["MyDep"] + cmake_variables = {} + cmake_includes = [] + cmake_libraries = ["MyDep::MyDep"] + cmake_compile_flags = [] + cmake_link_flags = [] + cmake_files = ["FindMyDep"] # resolves to FindMyDep.cmake next to my_env.py + + headers = ["mydep.h"] + state_fields = [] + init_code = "" + finalize_code = "" + dependencies = [] + +If your CMake module exposes a target (e.g., ``MyDep::MyDep``), it is usually +sufficient to list it in ``cmake_libraries`` and let the standard target +resolution machinery propagate include directories and compile flags. + +Built-in Libraries +------------------ + +DaCe ships a number of libraries under :mod:`dace.libraries`. They are loaded +automatically when ``dace`` is imported, and their library nodes are emitted +by the :ref:`Python frontend ` whenever a matching NumPy or +SciPy call is encountered. + +* :mod:`dace.libraries.blas` - Basic Linear Algebra Subprograms (GEMM, GEMV, + AXPY, DOT, GER, ...) with reference, OpenBLAS, Intel MKL, cuBLAS, and + rocBLAS expansions. See :ref:`blas`. +* :mod:`dace.libraries.lapack` - Dense linear algebra solvers (Cholesky, LU, + triangular solves) backed by LAPACK and cuSOLVER. +* :mod:`dace.libraries.linalg` - Higher-level linear algebra routines + (e.g., matrix inverse, tensor contractions) layered on top of BLAS, LAPACK, + and cuTENSOR. +* :mod:`dace.libraries.fft` - Fast Fourier transforms with pure and cuFFT + expansions. +* :mod:`dace.libraries.sparse` - Sparse linear algebra primitives (e.g., + SpMM, SpMV) with pure, Intel MKL, and cuSPARSE expansions. +* :mod:`dace.libraries.mpi` - Point-to-point and collective MPI operations + (Send, Recv, Bcast, Reduce, Allreduce, ...) for distributed SDFGs. +* :mod:`dace.libraries.pblas` - Distributed dense linear algebra via + ScaLAPACK/PBLAS, with Intel MKL and reference variants over OpenMPI/MPICH. +* :mod:`dace.libraries.standard` - Common building blocks (reductions, + transposes, code-region nodes) that are useful across other libraries and + user code. +* :mod:`dace.libraries.stencil` - Stencil computation library node with + CPU and GPU expansions. +* :mod:`dace.libraries.onnx` - Library nodes for ONNX operators, enabling + import and execution of ONNX models inside SDFGs. +* :mod:`dace.libraries.torch` - Interoperability with PyTorch tensors and + modules, including environments for libtorch headers and linkage. diff --git a/doc/extensions/sdfgconvertible.rst b/doc/extensions/sdfgconvertible.rst new file mode 100644 index 0000000000..dc4afea55f --- /dev/null +++ b/doc/extensions/sdfgconvertible.rst @@ -0,0 +1,111 @@ +.. _dsl: + +Python DSL Support via SDFG-Convertible Objects +=============================================== + +Many domain-specific languages (DSLs) embedded in Python are themselves +data-centric and can be lowered to an SDFG. DaCe exposes a small protocol, +*SDFG-convertible objects*, which lets such DSLs participate transparently in +``@dace.program`` parsing without having to be rewritten as DaCe replacements. +Examples of frameworks that integrate with DaCe through this protocol include: + +* the built-in ``@dace.program`` and ``@dace.method`` decorators (a + ``@dace.program`` is itself convertible, which is why one DaCe program + can call another); +* ``dace.ml.torch.module.DaceModule``, which exposes PyTorch + ``nn.Module`` instances as SDFG-convertible objects so they can be + invoked from any ``@dace.program``; +* `GT4Py `_ stencils, whose backends + produce SDFGs and register stencil objects as convertibles. + +The ``SDFGConvertible`` protocol +-------------------------------- + +The protocol is defined in :mod:`dace.frontend.python.common` as the abstract +class :class:`~dace.frontend.python.common.SDFGConvertible`. Any object that +appears in the closure of a ``@dace.program`` and implements (some of) the +methods below is treated as a callable SDFG by the Python frontend. + +.. list-table:: + :header-rows: 1 + :widths: 25 75 + + * - Method + - Purpose + * - ``__sdfg__(*args, **kwargs)`` + - Return the SDFG that should be invoked at the call site. The + arguments are the same Python values that the caller passed in; + implementations typically use them to specialize shapes, + data types, or compile-time options before generating the SDFG. + This is the minimum a convertible must implement. + * - ``__sdfg_signature__()`` + - Return ``(arg_names, constant_args)``, describing the positional + parameter names of the produced SDFG and the names whose values + must be constant at parse time. The frontend uses this to bind the + caller's arguments to the convertible's SDFG arguments. + * - ``__sdfg_closure__(reevaluate=None)`` + - Return a dictionary of additional values (arrays, scalars, callbacks, + nested convertibles) that should be merged into the parent program's + closure. ``reevaluate`` is a list of names that the caller wants + refreshed - implementations that cache arrays should re-read them + in that case. + * - ``closure_resolver(constant_args, given_args, parent_closure=None)`` + - Optional. Return an :class:`~dace.frontend.python.parser.SDFGClosure` + built from the convertible's own captured state. Implementations + that maintain a stateful closure (e.g., a neural network's weights) + use this hook to wire those values into the SDFG. + +Only ``__sdfg__`` and ``__sdfg_signature__`` are required for a basic +convertible. The other methods exist to expose stateful captures cleanly to +the parent program's closure. + +Minimal example +~~~~~~~~~~~~~~~ + +The following stub shows a callable Python class that participates in +``@dace.program`` parsing as if it were itself a DaCe program: + +.. code-block:: python + + import dace + + class MyOperator: + def __init__(self, scale: float): + self.scale = scale + + def __sdfg__(self, A): + sdfg = dace.SDFG('myop') + sdfg.add_array('A', A.shape, A.dtype) + state = sdfg.add_state() + ... + return sdfg + + def __sdfg_signature__(self): + return (['A'], []) + + op = MyOperator(scale=2.0) + + @dace.program + def use(A: dace.float32[16]): + op(A) # parsed as a nested SDFG inside `use` + +When ``use`` is parsed, the Python frontend recognizes ``op`` in the closure +as an :class:`~dace.frontend.python.common.SDFGConvertible`, calls +``op.__sdfg__(A)`` to obtain the operator's SDFG, and inlines it as a nested +SDFG inside ``use``. + +Caveats and recommendations +--------------------------- + +* ``__sdfg__`` is called every time the parent program is parsed. If + generating the SDFG is expensive, cache it on the object and key the cache + on the relevant compile-time arguments. +* The returned SDFG must be self-contained - any state it depends on at + runtime must either be passed through arguments or exposed via + ``__sdfg_closure__`` / ``closure_resolver``. +* For DSLs that produce many SDFGs, consider returning a small SDFG that + delegates to a library node; this keeps the parent program's IR + navigable and allows the DSL to ship its own expansions. +* If your DSL already has a dedicated frontend (e.g., it parses its own + AST), see :doc:`frontend` for guidelines on writing a separate frontend + pipeline rather than extending the Python frontend. diff --git a/doc/extensions/symbolic.rst b/doc/extensions/symbolic.rst new file mode 100644 index 0000000000..d0943ca9d7 --- /dev/null +++ b/doc/extensions/symbolic.rst @@ -0,0 +1,148 @@ +.. _symbolic: + +Working with Symbolic Expressions +================================= + +DaCe relies on `SymPy `_ to reason about array +shapes, memlet ranges, schedules, and any other quantity that may depend +on a runtime parameter. Almost every numeric attribute on the SDFG IR - +shape entries, ranges, parameter bounds, transient sizes - is either a +Python integer or a SymPy expression. This page is a tour of the helpers +in :mod:`dace.symbolic` that the rest of DaCe (and most extensions) use +when they need to manipulate those expressions. + +Symbols +------- + +DaCe symbols are SymPy symbols with an optional fixed type +and attached integer assumption (based on the type): + +.. code-block:: python + + import dace + N = dace.symbol('N', positive=True) + M = dace.symbol('M', dtype=dace.int64) + +The class :class:`dace.symbolic.symbol` extends :class:`sympy.Symbol` so +that any DaCe symbol can be used as-is in SymPy expressions, while still +carrying the metadata DaCe needs (a name, a data type, and a set of SymPy +assumptions). The convenience type alias +:data:`~dace.symbolic.SymbolicType` stands for ``Union[sympy.Basic, +SymExpr]`` and is the type hint to use whenever a function accepts a +symbol or a symbolic expression. + +The richer container :class:`~dace.symbolic.SymExpr` carries both an +*exact* expression and an *approximate* expression. Most code only needs +the exact one; over-approximations come into play when the IR has to +guarantee an upper or lower bound (see below). + +Indeterminate comparisons +~~~~~~~~~~~~~~~~~~~~~~~~~ + +Comparing symbolic expressions returns one of three answers: + +* ``(N > 0) == True`` - the inequality is implied by the assumptions on + the symbols. +* ``(N > 0) == False`` - the inequality is provably false. +* ``(N > 0)`` is an unevaluated SymPy expression - SymPy could neither + prove nor refute the inequality. This is a frequent source of subtle + bugs in transformations: never use such an expression in a Python + ``if``. Use :func:`~dace.symbolic.simplify` (or supply more + assumptions) before branching on it. + +Integer arithmetic caveats +~~~~~~~~~~~~~~~~~~~~~~~~~~ + +SymPy's default rational arithmetic does not match the C/C++ semantics +that the code generators eventually emit. DaCe ships two SymPy functions, +:class:`~dace.symbolic.int_floor` and ``int_ceil``, that correspond to +``a // b`` and ``ceil(a / b)`` for positive integers respectively. Use +them whenever you need to keep the symbolic result in ``int`` arithmetic; +the simplifier and the code generator both know about them. + +Analysis +-------- + +* :func:`~dace.symbolic.issymbolic` checks whether a value is a SymPy + expression that depends on at least one symbol (treating literal + ``Integer``/``Float`` as non-symbolic). +* :func:`~dace.symbolic.free_symbols_and_functions` returns the names of + every free symbol *and* every undefined function appearing in the + expression. This is the right helper to use when computing the symbol + set that must be present in an SDFG before an expression can be + evaluated. +* :func:`~dace.symbolic.swalk` is a small visitor that yields every + sub-expression in pre-order traversal, optionally descending into + function arguments. Use it to look for specific patterns or to gather + all occurrences of a kind of node. +* For sub-expression matching, SymPy's ``expr.match(pattern)`` and + ``expr.find(pattern)`` are usually sufficient; the SymPy documentation + has examples. + +Conversion +---------- + +* :func:`~dace.symbolic.pystr_to_symbolic` parses a Python-style string + (``"N + 2*M"``) into a SymPy expression while honoring DaCe's + conventions (e.g., ``int_floor`` for ``//``). +* :func:`~dace.symbolic.symstr` renders a SymPy expression back to a + Python-style string for generated code and Python-facing utilities. +* :func:`dace.codegen.common.sym2cpp` emits a C/C++-friendly string from + the same expression. Code generators should use ``sym2cpp`` instead of + ``str(expr)`` so that integer division, ``min``/``max``, and the DaCe + helper functions produce valid C++. + +Serialized symbolic form +~~~~~~~~~~~~~~~~~~~~~~~~ + +JSON serialization uses +:func:`~dace.symbolic.serialize_symbolic` and +:func:`~dace.symbolic.deserialize_symbolic` rather than +``pystr_to_symbolic``. The serialized strings are type-accurate and +bijective: + +* symbols are emitted as ``$name``; +* symbols with non-default type/assumptions use + ``symbol($name, dtype=dace.uint64, nonnegative=True)``; +* constants carry an explicit suffix such as ``2i16`` or ``8.0f64``; +* :class:`~dace.symbolic.SymExpr` values serialize as + ``SymExpr(expr, overapprox)``. + +Mutation and simplification +--------------------------- + +* :func:`~dace.symbolic.simplify` is the recommended simplifier + throughout the DaCe codebase. Unlike ``sympy.simplify`` it preserves + integer semantics and runs efficiently on the kinds of expressions + that show up in memlet ranges. +* :func:`~dace.symbolic.safe_replace` performs a substitution that is + safe under aliasing - replacing ``a -> b`` and ``b -> c`` simultaneously + produces ``a -> b, b -> c`` rather than ``a -> c``. Use it whenever + you build a substitution dictionary from a mapping that could overlap. +* :func:`~dace.symbolic.overapproximate` returns a syntactic + over-approximation of an expression (for instance, replacing a + data-dependent ``Min`` with one of its arguments). Together with the + ``approx`` field of :class:`~dace.symbolic.SymExpr` this is what allows + memlet propagation to compute conservative ranges when the exact range + is data-dependent. + +.. _symbolic-when: + +Symbolic types vs. scalars +-------------------------- + +When extending DaCe, a recurring decision is whether a quantity should be +modeled as a *symbol* or as a *scalar transient*. The rule of thumb is: + +* Use a **symbol** for quantities that are constant over the lifetime of + the SDFG (typically loop bounds and array shapes provided by the + caller), or a state (e.g., indices used in memlets). Symbols participate + in the symbolic propagation system. +* Use a **scalar transient** for quantities that may change within + states (counters, accumulators, intermediate results). They live in + :attr:`~dace.sdfg.sdfg.SDFG.arrays` and are written by tasklets like + any other data. + +When in doubt, prefer a symbol if the value is set once or consumed in +ranges or schedules; prefer a scalar otherwise. See the FAQ entry on this +question for a longer discussion. diff --git a/doc/frontend/daceprograms.rst b/doc/frontend/daceprograms.rst index f0617c65c6..198e9cc767 100644 --- a/doc/frontend/daceprograms.rst +++ b/doc/frontend/daceprograms.rst @@ -426,7 +426,7 @@ Calling SDFGs Directly ---------------------- If you want to call SDFGs from a ``@dace.program`` you can do so directly. This is useful when you have a custom implementation -or when you want to use another frontend (for example, in a :ref:`domain-specific language `). Example: +or when you want to use another frontend (for example, in a domain-specific language). Example: .. code-block:: python diff --git a/doc/frontend/parsing.rst b/doc/frontend/parsing.rst index 9868c52e6a..ca7513894b 100644 --- a/doc/frontend/parsing.rst +++ b/doc/frontend/parsing.rst @@ -168,9 +168,9 @@ Example: :align: center :alt: Generated SDFG for-loop for the above Data-Centric Python program -If the :class:`~dace.frontend.python.parser.DaceProgram`'s -:attr:`~dace.frontend.python.parser.DaceProgram.use_explicit_control_flow` attribute is set to true, this will utilize -:class:`~dace.sdfg.state.LoopRegion`s instead of the explicit state machine depicted above. + +If the :class:`~dace.frontend.python.parser.DaceProgram` :attr:`~dace.frontend.python.parser.DaceProgram.use_explicit_control_flow` attribute is set to true, this will utilize +:class:`~dace.sdfg.state.LoopRegion` s instead of the explicit state machine depicted above. :func:`~dace.frontend.python.newast.ProgramVisitor.visit_While` ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ @@ -190,9 +190,7 @@ Parses `while `_ statement :align: center :alt: Generated SDFG while-loop for the above Data-Centric Python program -If the :class:`~dace.frontend.python.parser.DaceProgram`'s -:attr:`~dace.frontend.python.parser.DaceProgram.use_explicit_control_flow` attribute is set to true, this will utilize -:class:`~dace.sdfg.state.LoopRegion`s instead of the explicit state machine depicted above. +If the :class:`~dace.frontend.python.parser.DaceProgram` :attr:`~dace.frontend.python.parser.DaceProgram.use_explicit_control_flow` attribute is set to true, this will utilize :class:`~dace.sdfg.state.LoopRegion` s instead of the explicit state machine depicted above. :func:`~dace.frontend.python.newast.ProgramVisitor.visit_Break` ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ diff --git a/doc/frontend/preprocessing.rst b/doc/frontend/preprocessing.rst new file mode 100644 index 0000000000..368edb02bc --- /dev/null +++ b/doc/frontend/preprocessing.rst @@ -0,0 +1,129 @@ +.. _preprocessing: + +Preprocessing Python AST +======================== + +Before the Python frontend lowers a ``@dace.program`` into an SDFG, the function's +abstract syntax tree (AST) is rewritten by a small pipeline of *preprocessing* +passes. Preprocessing has three goals: + +1. Resolve every name in the function body to either an SDFG symbol, a data + container, a closure constant, or an external Python callback. +2. Specialize Python constructs that the SDFG IR does not represent natively + (for example, list/tuple unpacking, context managers, augmented assignments, + or f-strings) into a form the parser can handle. +3. Reject programs that use language features which DaCe explicitly does not + support (see :doc:`pysupport`). + +The pipeline is implemented in :mod:`dace.frontend.python.preprocessing` and +is entered through :func:`~dace.frontend.python.preprocessing.preprocess_dace_program`. +The same entry point is reused recursively when one ``@dace.program`` calls +another, which keeps the closure of every nested SDFG self-consistent. + +Pipeline overview +----------------- + +The passes run in the following order. Most of them are +:class:`ast.NodeTransformer` subclasses; a few are :class:`ast.NodeVisitor` +checks that only validate the program. + +.. list-table:: + :header-rows: 1 + :widths: 30 70 + + * - Pass + - Role + * - :class:`~dace.frontend.python.preprocessing.StructTransformer` + - Rewrites ``dace.struct`` instantiations into the corresponding type + constructor calls so that downstream passes can see them as plain + function calls. + * - :class:`~dace.frontend.python.preprocessing.ModuleResolver` + - Replaces aliased module references (``import numpy as np``) with their + canonical names so that later replacement lookups (e.g., + ``numpy.dot``) succeed regardless of how the user imported the module. + * - :class:`~dace.frontend.python.preprocessing.MPIResolver` + - Recognizes ``mpi4py`` communicator method calls and rewrites them into + the corresponding ``dace.comm`` library calls. Skipped silently if + ``mpi4py`` is not installed. + * - :class:`~dace.frontend.python.preprocessing.ModuloConverter` + - Normalizes Python's modulo operator to DaCe's C-style semantics (see + :doc:`pysupport`, section 6.7). + * - :func:`~dace.frontend.python.preprocessing.find_disallowed_statements` + - Walks the AST and raises a ``TypeError`` if any statement listed in + :data:`~dace.frontend.python.newast.DISALLOWED_STMTS` is present + (``import``, ``try``, ``yield``, ``del``, ...). When disallowed + statements are found, the function is converted into a Python callback + instead of being parsed as an SDFG. + * - :class:`~dace.frontend.python.preprocessing.GlobalResolver` + - The largest pass. Resolves every free name to one of: a closure + constant (substituted in-place), a captured array (registered in the + :class:`~dace.frontend.python.parser.SDFGClosure`), a DaCe symbol, or + a callable to be inlined or invoked as a callback. ``visit_Assert`` + and ``visit_Raise`` live here as well; ``assert`` is folded statically + and ``raise`` is replaced by a warning. + * - :class:`~dace.frontend.python.preprocessing.DisallowedAssignmentChecker` + - Rejects assignments that would mutate compile-time constants or + walrus-bound names that are visible from the closure. + * - :class:`~dace.frontend.python.preprocessing.LoopUnroller` + - Statically unrolls ``for`` loops whose iterator is a literal sequence + (``range`` over compile-time bounds, captured tuples, ``enumerate``, + ``zip``, ``dace.unroll``). Unrolling is what enables many + metaprogramming patterns that the SDFG IR cannot express directly. + * - :class:`~dace.frontend.python.preprocessing.ExpressionInliner` + - Inlines call expressions whose callee is a pure Python function + captured from the closure (including lambdas), as long as the result + can be expressed as a single AST expression. + * - :class:`~dace.frontend.python.preprocessing.ContextManagerInliner` + - Replaces ``with`` blocks (other than ``with dace.tasklet``) with the + inlined ``__enter__`` / ``__exit__`` calls of the underlying context + manager, including handling of early ``return`` / ``break`` / + ``continue``. + * - :class:`~dace.frontend.python.preprocessing.ConditionalCodeResolver` + - Folds ``if`` / ``elif`` branches whose condition is a compile-time + constant. + * - :class:`~dace.frontend.python.preprocessing.DeadCodeEliminator` + - Removes unreachable code after the previous folding step (e.g., + statements after a ``return`` in a fully-folded branch). + * - :class:`~dace.frontend.python.preprocessing.AugAssignExpander` + - Expands ``a += b`` to ``a = a + b`` for cases where the left-hand + side is not a simple data access. + * - :class:`~dace.frontend.python.preprocessing.CallTreeResolver` + - Discovers all transitively-called ``@dace.program`` and + ``SDFGConvertible`` objects so that they can be parsed and added to + the closure. + +The closure resolver, conditional resolver, dead-code eliminator, and +expression inliner run *together* in a fixed-point loop: as long as one of +them rewrites the AST, the loop is run again. The maximum number of passes +is controlled by the ``frontend.preprocessing_passes`` configuration entry +(``-1`` means "run until quiescent"). + +The closure +----------- + +The output of preprocessing is two-fold: + +* a :class:`~dace.frontend.python.parser.PreprocessedAST` containing the + rewritten AST, the source file, and the resolved global namespace; and +* a :class:`~dace.frontend.python.parser.SDFGClosure` recording every + external object that the program needs at call time - captured arrays, + constants, nested ``@dace.program`` callees, and Python callbacks. + +The closure is what allows ``@dace.program`` to behave like a regular +Python function from the caller's point of view while still producing a +fully self-contained SDFG: when the compiled SDFG is invoked, the closure +is used to bind the captured external state to the program's arguments. + +Disabling or inspecting preprocessing +------------------------------------- + +* Set the configuration option ``frontend.verbose_errors = true`` to see the + exact pass that raised an exception during preprocessing. +* Set ``frontend.preprocessing_passes`` to a positive integer to cap the + number of fixed-point iterations (useful when debugging an infinite loop + in a custom replacement that produces new AST every pass). +* Use :func:`dace.frontend.python.astutils.unparse` to dump the + intermediate AST after any pass during development. + +For background on what comes after preprocessing, see :doc:`parsing` and +:ref:`python-frontend`. diff --git a/doc/frontend/pysupport.rst b/doc/frontend/pysupport.rst index 8993df1988..5bc839c5ef 100644 --- a/doc/frontend/pysupport.rst +++ b/doc/frontend/pysupport.rst @@ -40,16 +40,20 @@ The following keywords are recognized (for at least a subset of their Python fun - or, and, not - if, elif, else - for, while, break, continue, in -- def, return, lambda, with +- def, return, lambda, with, pass +- assert, raise (handled by the preprocessor; see 7.3 and 7.8 below) +- class (only when used together with ``@dace.method``; see 8.7) -The following keywords are NOT accepted: +The following keywords are NOT accepted in the body of a ``@dace.program``: - global, nonlocal -- class - try, except, finally -- raise, yield, pass +- yield - import, from, as -- assert, async, await, del +- async, await, del + +The authoritative list lives in ``DISALLOWED_STMTS`` and ``_DISALLOWED_STMTS`` in +:mod:`dace.frontend.python.newast`. 2.3.2 Reserved classes of identifiers: @@ -178,7 +182,10 @@ Supported 6.14 Lambdas ^^^^^^^^^^^^ -Supported only for defining WCR/reduction operators +Lambdas used as the WCR (write-conflict-resolution) argument of +``dace.reduce`` and similar reductions are supported as-is. Lambdas used +as regular callables are inlined at their call site by the preprocessor. +Returning a lambda from a ``@dace.program`` is not supported. 6.15 Expression lists ^^^^^^^^^^^^^^^^^^^^^ @@ -216,12 +223,22 @@ Supported with the same constraints for targets as in assignment statements. 7.2.2 Annotated assignment statements: -Unsupported +Supported. Annotations on local variables are accepted by ``visit_AnnAssign`` +(see :mod:`~dace.frontend.python.newast`); the annotation is honored when +creating the underlying data container. + +Note that annotated assignments to *compile-time constants* in the closure are +rejected by ``DisallowedAssignmentChecker`` during preprocessing. 7.3 The assert statement ^^^^^^^^^^^^^^^^^^^^^^^^ -Unsupported +The Python preprocessor evaluates assertions statically against the program's +closure (``globals`` plus any captured constants). If the condition is +statically true, the ``assert`` is dropped; if it is statically false, the +DaCe program fails to compile with an ``AssertionError``. ``assert`` +statements that depend on runtime values cannot be checked and are skipped +with a warning. 7.4 The pass statement ^^^^^^^^^^^^^^^^^^^^^^ @@ -245,7 +262,10 @@ Unsupported 7.8 The raise statement ^^^^^^^^^^^^^^^^^^^^^^^ -Unsupported + +Not supported at runtime. ``raise`` statements found in a ``@dace.program`` +are reported with a warning during preprocessing and stripped from the +generated code; control will continue past the would-be raise site. 7.9 The break statement ^^^^^^^^^^^^^^^^^^^^^^^ @@ -262,7 +282,9 @@ SDFG-level as the for/while statement. 7.11 The import statement ^^^^^^^^^^^^^^^^^^^^^^^^^ -Unsupported, including 7.11.1 Future statements +Unsupported inside the body of a ``@dace.program``. ``from __future__`` +imports placed at module scope are tolerated (the preprocessor strips them +before parsing). 7.12 The global statement ^^^^^^^^^^^^^^^^^^^^^^^^^ @@ -316,18 +338,38 @@ Unsupported 8.5 The with statement ^^^^^^^^^^^^^^^^^^^^^^ -Only supported ``with dace.tasklet`` +``with dace.tasklet`` is the canonical form for explicit-dataflow tasklets +(see :ref:`explicit-dataflow-mode`). General ``with`` statements are supported through +the preprocessing-time :class:`~dace.frontend.python.preprocessing.ContextManagerInliner`, +which inlines the ``__enter__`` / ``__exit__`` calls of the context manager +at the appropriate program points. ``with`` blocks that depend on raised +exceptions for control flow are not supported. 8.6 Function definitions ^^^^^^^^^^^^^^^^^^^^^^^^ -Supported only with the ``dace.program`` decorator. Function arguments must be -type-annotated. Nested ``dace.program`` definitions are not supported. +The top-level callable must be decorated with ``@dace.program`` (or +``@dace.method`` when defined inside a class - see 8.7). Function arguments +must be type-annotated. + +Calls from one ``@dace.program`` to another are supported and lower to +nested SDFGs. *Defining* a ``@dace.program`` inside the body of another +``@dace.program`` is not supported; nested ``def`` blocks without the +decorator are inlined by the Python preprocessor when called. + +Lambdas defined as part of a reduction (``dace.reduce``) or other library +nodes that take a WCR are accepted as-is; other lambdas are inlined at their +call site by the preprocessor (see 6.14). 8.7 Class definitions ^^^^^^^^^^^^^^^^^^^^^ -Supported through the Preprocessing component of the Python frontend. +``class`` blocks themselves cannot appear inside a ``@dace.program`` +(``ClassDef`` is in :data:`~dace.frontend.python.newast.DISALLOWED_STMTS`). +Methods of regular Python classes can however be turned into DaCe programs +by decorating them with ``@dace.method``: the preprocessor binds ``self`` as +part of the program's closure so that attribute accesses on ``self`` resolve +to data containers, scalars, or constants of the enclosing instance. 8.8 Coroutines ^^^^^^^^^^^^^^ diff --git a/doc/frontend/python.rst b/doc/frontend/python.rst index 91dc56eebd..ed4465d527 100644 --- a/doc/frontend/python.rst +++ b/doc/frontend/python.rst @@ -25,8 +25,6 @@ capabilities DaCe provides by default. :maxdepth: 1 parsing + preprocessing pysupport npsupport - - -.. preprocessing diff --git a/doc/general/faq.rst b/doc/general/faq.rst new file mode 100644 index 0000000000..624c9f44f2 --- /dev/null +++ b/doc/general/faq.rst @@ -0,0 +1,51 @@ +.. _faq: + +Frequently Asked Questions +========================== + +How do I apply a transformation to a specific node or set of nodes? +------------------------------------------------------------------- + +Use the ``apply_to`` class method of the transformation. Each transformation +class declares a fixed set of pattern variables (e.g., ``MapFusion`` exposes +``first_map_exit``, ``array``, and ``second_map_entry``); pass the matched +nodes as keyword arguments: + +.. code-block:: python + + from dace.transformation.dataflow import MapFusion + MapFusion.apply_to(sdfg, first_map_exit=exit1, array=array_node, + second_map_entry=entry2) + +For pattern-based application across the whole SDFG, use +:meth:`~dace.sdfg.sdfg.SDFG.apply_transformations` and +:meth:`~dace.sdfg.sdfg.SDFG.apply_transformations_repeated`. For +whole-program optimizations, run a :ref:`pass ` instead. + +When should I use a symbol and when should I use a scalar? +---------------------------------------------------------- + +See the discussion in :ref:`symbolic-when` for the general guidelines. + + +How do I perform dynamic memory allocation? +------------------------------------------- + +Array shapes in DaCe are always associated with symbolic expressions. +When the size is known only at runtime, declare the array with a symbolic +shape and assign the actual size to that symbol on a state's +*inter-state edge* before the array is used: + +.. code-block:: python + + sdfg = dace.SDFG('dyn') + N = dace.symbol('N') + sdfg.add_array('A', (N,), dace.float32, transient=True) + + init = sdfg.add_state('init') + body = sdfg.add_state('body') + sdfg.add_edge(init, body, dace.InterstateEdge(assignments={'N': 'compute_size()'})) + +When the size depends on data within the SDFG, place the allocation inside +a nested SDFG so that its symbol can be assigned from a scalar of the +parent. diff --git a/doc/general/glossary.rst b/doc/general/glossary.rst index 17783d1697..28f106d471 100644 --- a/doc/general/glossary.rst +++ b/doc/general/glossary.rst @@ -10,7 +10,7 @@ Below you can find any concept used in the DaCe framework and where it is descri * :ref:`Environment ` * :ref:`Inter-State Edge ` * :ref:`Library ` -* Library Node: :ref:`Definition `, :ref:`Implementation ` +* :ref:`Library Node ` * :ref:`Map scope `, :ref:`sdfg-map` * Memlet: :ref:`Element `, :ref:`Anatomy ` * :ref:`Memlet Propagation ` diff --git a/doc/general/structure.rst b/doc/general/structure.rst index 04d4586651..dc07493a5f 100644 --- a/doc/general/structure.rst +++ b/doc/general/structure.rst @@ -20,7 +20,7 @@ DaCe is divided into the following subfolders: * ``data.py``: Definitions for usage and creation of :ref:`descriptors` * ``memlet.py``: Definition of a :ref:`Memlet ` * ``subsets.py``: Subset types used in memlets (e.g., :class:`~dace.subsets.Range`) - * ``symbolic.py``: Symbolic types, expressions, conversion, and analysis functions. See :ref:`sdfg-symbol` and :ref:`symbolic` + * ``symbolic.py``: Symbolic types, expressions, conversion, and analysis functions. See :ref:`sdfg-symbol` * Optimization and transformation: * ``transformation``: Transformation classes and helpers. See :ref:`transformations` diff --git a/doc/index.rst b/doc/index.rst index 616dca24e7..e58ff4ae45 100644 --- a/doc/index.rst +++ b/doc/index.rst @@ -53,8 +53,7 @@ If you use DaCe, cite us: setup/integration ide/cli general/glossary - -.. general/faq + general/faq .. toctree:: :caption: Developer Guide @@ -64,9 +63,16 @@ If you use DaCe, cite us: general/debugging extensions/properties sdfg/transforming + sdfg/schedule_tree frontend/python codegen/codegen extensions/extensions + extensions/libraries + extensions/symbolic + extensions/frontend + extensions/sdfgconvertible + extensions/backend + extensions/instrumentation general/errors .. toctree:: diff --git a/doc/optimization/guidelines.rst b/doc/optimization/guidelines.rst new file mode 100644 index 0000000000..1b2574e371 --- /dev/null +++ b/doc/optimization/guidelines.rst @@ -0,0 +1,84 @@ +.. _opt_guidelines: + +Optimization Guidelines +======================= + +This page collects the high-level principles that underlie effective DaCe +optimization. The other pages in :ref:`optimization` show *how* to apply +specific transformations; this one is about *what* to aim for. + +Minimize data movement +---------------------- + +Data movement - between memory hierarchy levels, between devices, and +across processes - dominates the cost of most modern workloads. The single +most impactful optimization in DaCe is therefore to remove or shorten +memory transfers. In practice this means: + +* **Fuse states and maps** so that producers and consumers of the same data + share a scope. Map fusion (and the broader subgraph fusion family of + transformations) eliminates the round-trip through main memory between + fused nodes. +* **Stage data in fast memory** with + :class:`~dace.transformation.dataflow.local_storage.InLocalStorage`, + :class:`~dace.transformation.dataflow.local_storage.AccumulateTransient`, + and tiling transformations. Once a tile is in shared/registers, reuse it + for as many computations as possible before evicting. +* **Avoid redundant transients** by running the + :class:`~dace.transformation.passes.transient_reuse.TransientReuse` pass. + Transients that are not reused inflate working sets without speeding up + the program. +* **Specialize storage location**. Set the + :attr:`~dace.data.Data.storage` attribute of arrays to match where they + are produced and consumed; the codegen will emit the right allocation + and copy logic. + +Expose parallelism, then specialize the schedule +------------------------------------------------ + +DaCe's IR separates *what* is parallel (maps, library nodes) from *how* it +is executed (schedules, storage, library implementations). The recommended +order of operations is: + +1. **Expose** all the parallelism in the program. Use the Python frontend's + :func:`dace.map` syntax or transformations like + :class:`~dace.transformation.interstate.loop_to_map.LoopToMap` + and :class:`~dace.transformation.dataflow.map_collapse.MapCollapse` to + produce maximally-parallel loop nests. +2. **Specialize** the schedule to the target platform - use transformations + such as :class:`~dace.transformation.interstate.gpu_transform_sdfg.GPUTransformSDFG` + or choose a schedule manually (:attr:`~dace.dtypes.ScheduleType.GPU_Device` / + :attr:`~dace.dtypes.ScheduleType.GPU_ThreadBlock` schedules for CUDA, + :attr:`~dace.dtypes.ScheduleType.CPU_Multicore` for OpenMP, etc.) +3. **Pick library implementations** for any matrix multiplications, FFTs, + reductions, or collectives. Library nodes default to a portable + ``"pure"`` expansion; switching to ``"MKL"``, ``"cuBLAS"``, + ``"NCCL"``, or another vendor implementation often delivers an + order-of-magnitude speedup. See :ref:`blas`. + +Profile, do not guess +--------------------- + +DaCe makes performance measurement cheap. Use it before deciding where to +optimize: + +* :ref:`profiling` shows how to enable instrumentation on selected SDFG + elements and how to interpret the resulting reports. +* :ref:`optimization_vscode` overlays runtime measurements onto the SDFG + view so that hot regions are visually obvious. +* :ref:`opt_sdfgapi` describes how to drive the optimization loop from a + notebook, applying transformations between profiling runs. + +A good rule of thumb: the earliest and largest speedups typically come from +repeatedly looking for the "hottest" region and applying one or +two transformations to improve its scheduling. + +When in doubt, start with the auto-optimizer +-------------------------------------------- + +The :ref:`automatic heuristics ` are a strong baseline on most +workloads. Calling ``auto_optimize`` once before doing manual optimization +saves time and makes it easier to tell whether a manual change is actually +helping. If the baseline is already close to peak (verifiable through +profiling and a roofline-style calculation), it is often more productive +to switch to algorithmic improvements than to keep tuning the schedule. diff --git a/doc/optimization/interactive.rst b/doc/optimization/interactive.rst new file mode 100644 index 0000000000..0a952b15b2 --- /dev/null +++ b/doc/optimization/interactive.rst @@ -0,0 +1,112 @@ +.. _opt_sdfgapi: + +Interactive Optimization with the SDFG API +========================================== + +Frontends produce a *baseline* SDFG that mirrors the structure of the user's +source program. The Python SDFG API exposes every node, edge, and property +of that graph as a regular Python object, so an SDFG can be optimized +incrementally from a notebook or REPL session without recompiling the +program after every change. This page collects the entry points most often +needed when driving that loop by hand. + +Loading and inspecting a baseline +--------------------------------- + +Any SDFG produced by a frontend or saved to disk can be loaded interactively: + +.. code-block:: python + + import dace + sdfg = dace.SDFG.from_file('matmul.sdfg') # JSON or .sdfgz + sdfg.view() # opens VS Code or Jupyter + sdfg.validate() # raise on inconsistencies + +The :meth:`~dace.sdfg.sdfg.SDFG.view` call renders the graph through the +:ref:`SDFV viewer `. ``validate()`` will raise descriptive +exceptions on common issues such as missing memlets, dangling connectors, or +undefined symbols. + +Operating on the graph +---------------------- + +The SDFG is just a stateful graph object. The most common operations are: + +* Iterating: ``sdfg.states()``, ``sdfg.all_nodes_recursive()``, + ``sdfg.arrays.items()``. +* Locating nodes by predicate, e.g. + ``[n for n, _ in sdfg.all_nodes_recursive() if isinstance(n, dace.nodes.MapEntry)]``. +* Editing properties in place: changing + :attr:`~dace.sdfg.nodes.Map.schedule`, :attr:`~dace.data.Data.storage`, + or :attr:`~dace.sdfg.nodes.AccessNode.data`. +* Saving and reloading checkpoints with :meth:`~dace.sdfg.sdfg.SDFG.save` + followed by :meth:`~dace.sdfg.sdfg.SDFG.from_file`. + +Each of these is safe to do inside a live session (even a notebook), as the SDFG +is a pure Python data structure until you call :meth:`~dace.sdfg.sdfg.SDFG.compile`. + +Applying transformations interactively +-------------------------------------- + +The :ref:`transformation API ` provides three idioms for +firing a transformation from Python: + +1. **Pattern matching** with :meth:`~dace.sdfg.sdfg.SDFG.apply_transformations` + or :meth:`~dace.sdfg.sdfg.SDFG.apply_transformations_repeated` walks the + graph, finds every match for a transformation class, and applies it. + These methods accept a list of classes and an ``options`` dictionary, + making them convenient for batch experiments. + +2. **Direct application** via the ``apply_to`` class method targets a + specific subgraph by passing the matched nodes as keyword arguments. + This is the form used inside notebooks when you already know the nodes + you want to operate on: + + .. code-block:: python + + from dace.transformation.dataflow import MapFusion + MapFusion.apply_to(sdfg, first_map_exit=exit1, array=array_node, + second_map_entry=entry2) + +3. **Whole-program passes** are run with + :meth:`~dace.transformation.passes.pipeline.Pipeline.apply_pass` (or + :meth:`~dace.sdfg.sdfg.SDFG.apply_passes`) and operate on the SDFG + globally. Common examples include + :class:`~dace.transformation.passes.scalar_to_symbol.ScalarToSymbolPromotion` + and :class:`~dace.transformation.passes.transient_reuse.TransientReuse`. + +In all three cases the SDFG is mutated in place, so chaining operations is +just calling them in sequence. To revert an experiment, reload from the +saved baseline. + +Instrumentation feedback loop +----------------------------- + +Interactive optimization works best when paired with +:ref:`runtime instrumentation `. Set +:attr:`~dace.sdfg.nodes.Map.instrument` (or +:attr:`~dace.sdfg.sdfg.SDFG.instrument`) on the regions of interest, run +the program once, and inspect +:class:`~dace.codegen.instrumentation.report.InstrumentationReport` to +confirm whether the last transformation actually moved the needle. The +recommended workflow is: + +1. Identify the hottest map or library node from a profile. +2. Apply a transformation (or a small sequence) targeting it. +3. Re-run; compare the new instrumentation report against the previous + one. +4. If the change helped, save the SDFG; otherwise reload and try a + different sequence. + +Composing with the auto-optimizer +--------------------------------- + +The :ref:`automatic heuristics ` (``dace.transformation.auto.auto_optimize``) +are a useful starting point even when the goal is a fully manual schedule. +A common pattern is to call ``auto_optimize`` once to establish a strong +baseline, save the result, and then drive the remaining optimizations by +hand with the API described above. + +For more involved manipulations, such as scripting transformation searches, +see :ref:`sdfg-api` and the ``samples/optimization`` directory in the +repository. diff --git a/doc/optimization/optimization.rst b/doc/optimization/optimization.rst index 392d1895c4..8adaf067c0 100644 --- a/doc/optimization/optimization.rst +++ b/doc/optimization/optimization.rst @@ -29,7 +29,7 @@ When working with specific platforms, be sure to read the **Best Practices** doc recommended to read vendor-provided documentation on how to maximize performance on that platform. Finally, our experimental auto-tuning API allows for automatic optimization of SDFGs by searching over the set of possible -configurations. This is done by :ref:`evaluating the performance of each configuration ` and selecting the best one. +configurations. This is done by evaluating the performance of each configuration and selecting the best one. For example, :class:`~dace.optimization.map_permutation_tuner.MapPermutationTuner` automatically tunes the order of multi-dimensional maps for the best performance, and :class:`~dace.optimization.data_layout_tuner.DataLayoutTuner` globally tunes the data layout of arrays. @@ -49,7 +49,5 @@ The following subsections provide more information on the different types of opt blas vscode gpu - - -.. interactive -.. guidelines + interactive + guidelines diff --git a/doc/sdfg/ir.rst b/doc/sdfg/ir.rst index 8c02a618f6..7d74efbc32 100644 --- a/doc/sdfg/ir.rst +++ b/doc/sdfg/ir.rst @@ -214,7 +214,7 @@ that pertain to how it will be mapped to hardware (and to the generated code): ` **Storage location** refers to where the container will be allocated --- examples include :class:`~dace.dtypes.StorageType.CPU_Heap` for allocation using ``new[]`` and :class:`~dace.dtypes.StorageType.GPU_Global` for VRAM on the GPU (``{cuda,hip}Malloc``). The full built-in list can be found in the enumeration definition :class:`~dace.dtypes.StorageType`. The enumeration is -user-extensible (see :ref:`enums`), so adding new entries is easy. +user-extensible, so adding new entries is easy. **Allocation lifetime** refers to the allocation/deallocation scope of a data container. By default, :class:`~dace.dtypes.AllocationLifetime.Scope` is used, which specifies that the access nodes dictate the allocation lifetime --- the innermost common scope in which @@ -314,9 +314,6 @@ dimensions: example_with_undefined(a, b) -.. note:: - For more information about developing with symbolic expressions, read :ref:`symbolic`. - .. _connectors: Connectors diff --git a/doc/sdfg/schedule_tree.rst b/doc/sdfg/schedule_tree.rst new file mode 100644 index 0000000000..8f4989a6e5 --- /dev/null +++ b/doc/sdfg/schedule_tree.rst @@ -0,0 +1,177 @@ +.. _schedule_tree: + +Schedule Tree +============= + +The **schedule tree** is an alternative, tree-shaped intermediate representation of a +DaCe program. It is derived from an :ref:`SDFG ` and represents the same computation, +but in a form that resembles structured imperative code rather than a graph of states and +dataflow multigraphs. + +.. warning:: + The Schedule Tree representation is currently under development. Some of the tree node + types may change between versions, and conversion from schedule tree to SDFG is + experimental. + + +A schedule tree is composed of nested ``ScheduleTreeScope`` blocks (loops, branches, +maps, consume scopes, ...) and leaf ``ScheduleTreeNode`` statements (tasklets, copies, +library calls, symbol assignments, views, reference sets, ...). The root of the entire +program is a ``ScheduleTreeRoot``, which additionally carries the unified descriptor +repository (containers, symbols, constants, callback mapping, argument names). + +A textual rendering of a schedule tree looks like a Python-like program, for example:: + + for i = 0; (i < N); i = (i + 1): + map j in [0:M]: + B[i, j] = tasklet(A[i, j]) + C[i, j] = tasklet(...) + +The implementation lives in :mod:`dace.sdfg.analysis.schedule_tree`, with the node +hierarchy defined in ``treenodes.py`` and conversion routines in ``sdfg_to_tree.py`` and +``tree_to_sdfg.py``. + + +Why a Schedule Tree? +-------------------- + +The SDFG is a graph-based IR that exposes parallelism and data movement explicitly. This +is excellent for transformations that reason about dataflow, scopes, and memory, but it +can be cumbersome for analyses or transformations that are naturally expressed on +*structured control flow* — for example: + +* Reasoning about the **execution order** of statements without traversing states and + inter-state edges. +* Performing **control-flow transformations** such as loop-invariant code motion or + eliminating empty branches. +* Understanding the program at a **higher level** that closely mirrors the original + source code. +* Implementing analyses that benefit from a **lexical, tree-shaped view** (e.g., visitor + patterns analogous to Python's ``ast.NodeVisitor``). + +In a schedule tree: + +* Control flow is explicit and structured (``ForScope``, ``WhileScope``, ``DoWhileScope``, + ``IfScope`` / ``ElifScope`` / ``ElseScope``, ``GBlock`` for irreducible flow). +* Dataflow scopes (``MapScope``, ``ConsumeScope``) can be nested within control flow or + vice versa. +* Nested SDFGs are flattened: data containers from nested SDFGs are dealiased into a + unified, top-level descriptor repository, so a single naming system applies throughout + the tree. +* Each node provides ``input_memlets()`` / ``output_memlets()``, which return the set of + read/write memlets of the node, optionally with memlet propagation across loops and + maps. + + +When To Use It +-------------- + +Prefer the schedule tree representation when: + +* You want to develop a frontend that is more naturally expressed on structured control + flow graphs, for a smoother mapping from source code to IR. +* You want to write a transformation or analysis that is easier to express on + **structured control flow** than on a state graph. +* You need a **lexical traversal** of the program (the visitor / transformer base + classes ``ScheduleNodeVisitor`` and ``ScheduleNodeTransformer`` mirror Python's + ``ast`` API). +* You want to inspect or pretty-print a program in a **human-readable, code-like form** + (``ScheduleTreeRoot.as_string()``). +* You are reasoning about **input/output sets** of a region of code with memlet + propagation, taking loop ranges and map parameters into account. + +Prefer the SDFG representation when: + +* The transformation reasons primarily about **dataflow** between access nodes, + tasklets, and library nodes. +* You need to apply existing :ref:`pattern-matching transformations ` + or :ref:`passes `, most of which are written against the SDFG IR. +* You need to **generate code**: code generation operates on SDFGs (see + :ref:`codegen`). To generate code from a schedule tree, convert it back to an SDFG + first. + + +Converting Between SDFG and Schedule Tree +----------------------------------------- + +Every SDFG can be converted to a schedule tree, and a schedule tree can be converted +back to an SDFG. The two conversions are not bit-exact inverses (the round-trip may +introduce fewer states, simplify trivial control flow, etc.), but they preserve program +semantics. + +From SDFG to schedule tree:: + + from dace.sdfg.analysis.schedule_tree import sdfg_to_tree + + stree = sdfg_to_tree.as_schedule_tree(sdfg) + # Equivalent shortcut: + stree = sdfg.as_schedule_tree() + + print(stree.as_string()) + +By default, the SDFG is deep-copied before conversion. Pass ``in_place=True`` to avoid +the copy; note that an in-place conversion may leave the original SDFG in an +inconsistent state. + +From schedule tree back to SDFG:: + + from dace.sdfg.analysis.schedule_tree import tree_to_sdfg + + sdfg = tree_to_sdfg.from_schedule_tree(stree) + # Equivalent shortcut: + sdfg = stree.as_sdfg(validate=True, simplify=True) + +The reconstruction inserts state boundaries where required (e.g., write-after-write, +before labels, around irreducible control flow) according to the configured +``StateBoundaryBehavior``. + + +Working With Schedule Trees +--------------------------- + +The node API mirrors Python's ``ast`` module: + +* ``ScheduleNodeVisitor`` walks the tree with ``visit_ClassName`` dispatch. +* ``ScheduleNodeTransformer`` walks the tree and lets each ``visit_*`` method return a + replacement node, ``None`` to delete it, or a list of nodes to splice in. +* ``ScheduleTreeScope.preorder_traversal()`` yields all descendants in pre-order. + +A small example that removes empty ``IfScope`` nodes:: + + from dace.sdfg.analysis.schedule_tree import treenodes as tn + + class RemoveEmptyIfs(tn.ScheduleNodeTransformer): + def visit_IfScope(self, node: tn.IfScope): + if not node.children: + return None + return self.generic_visit(node) + + RemoveEmptyIfs().visit(stree) + +Some ready-made passes live in :mod:`dace.sdfg.analysis.schedule_tree.passes`, such as +``remove_unused_and_duplicate_labels`` and ``remove_empty_scopes``. + +To query memlets of a region with optional propagation:: + + inputs = stree.input_memlets() # Set of read memlets at the root scope + outputs = stree.output_memlets() # Set of written memlets at the root scope + + # Keep symbols local to inner scopes instead of propagating them outwards: + raw_inputs = stree.input_memlets(keep_locals=True) + + +Caveats +------- + +* The schedule tree is intended primarily for **analysis and structured transformation**, + not for code generation. Convert back to an SDFG before invoking the code generator. +* Conversion **flattens nested SDFGs** and renames their containers to the top-level + descriptors. If your analysis depends on nested SDFG identity, you must operate on + the SDFG directly. +* Some constructs that have no structured equivalent (e.g., irreducible control flow) + are represented inside a ``GBlock`` with explicit ``StateLabel`` and ``GotoNode`` + entries. +* Round-tripping ``SDFG -> schedule tree -> SDFG`` is **semantics-preserving but not + graph-identical**; expect benign differences such as added/removed empty states or + reorganized state transitions, especially when ``simplify=True`` is requested in + ``as_sdfg()``. diff --git a/doc/setup/integration.rst b/doc/setup/integration.rst index 300f172f9c..ca85091c5e 100644 --- a/doc/setup/integration.rst +++ b/doc/setup/integration.rst @@ -104,7 +104,7 @@ A compiled SDFG library contains three functions, which are named after the SDFG between invocations of the SDFG. The other functions take this handle as their first argument. The arguments to this function are only the symbols used in the SDFG, ordered by name. * ``__dace_exit_``: Deallocates all arrays and frees all data descriptors in the given handle. Returns - a value of 0 if finalized successfully, or another value on failure. + a value of 0 if finalized successfully, or another value on failure. * ``__program_``: The actual SDFG function, which takes the handle as its first argument, followed by the arguments to the SDFG, ordered by name, followed by the symbol arguments, also ordered by name. diff --git a/doc/setup/quickstart.rst b/doc/setup/quickstart.rst index 05e463c9a5..9221128f02 100644 --- a/doc/setup/quickstart.rst +++ b/doc/setup/quickstart.rst @@ -31,7 +31,7 @@ DaCe under the hood. Internally, DaCe creates a shared library (DLL/SO file) that can readily be used in any C ABI compatible language, such as C++ or FORTRAN (See :ref:`integration`). -From here on out, you can optimize (:ref:`interactively `, :ref:`programmatically `, or +From here on out, you can optimize (:ref:`interactively `, :ref:`programmatically `, or :ref:`automatically `), :ref:`instrument `, and distribute your code. diff --git a/doc/source/dace.cli.rst b/doc/source/dace.cli.rst index 4dce7c5bc2..3098564ffc 100644 --- a/doc/source/dace.cli.rst +++ b/doc/source/dace.cli.rst @@ -10,7 +10,7 @@ dace.cli.sdfgcc module :show-inheritance: dace.cli.sdfv module ---------------------- +-------------------- .. automodule:: dace.cli.sdfv :members: @@ -18,7 +18,7 @@ dace.cli.sdfv module :show-inheritance: dace.cli.daceprof module ----------------------- +------------------------ .. automodule:: dace.cli.daceprof :members: diff --git a/doc/source/dace.codegen.instrumentation.rst b/doc/source/dace.codegen.instrumentation.rst index 138ca40386..1f6acb118f 100644 --- a/doc/source/dace.codegen.instrumentation.rst +++ b/doc/source/dace.codegen.instrumentation.rst @@ -5,7 +5,7 @@ Submodules ---------- dace.codegen.instrumentation.gpu_tx_markers module ------------------------------------------------ +-------------------------------------------------- .. automodule:: dace.codegen.instrumentation.gpu_tx_markers :members: diff --git a/doc/source/dace.rst b/doc/source/dace.rst index d080513448..74b582927f 100644 --- a/doc/source/dace.rst +++ b/doc/source/dace.rst @@ -120,12 +120,3 @@ dace.symbolic module :members: :undoc-members: :show-inheritance: - - -Module contents ---------------- - -.. automodule:: dace - :members: - :undoc-members: - :show-inheritance: diff --git a/requirements.txt b/requirements.txt index bd5d29868e..6877c07e0e 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,11 +1,10 @@ astunparse==1.6.3 dill==0.3.9 -fparser==0.1.4 +fparser>0.1.3 mpmath==1.3.0 networkx==3.4.2 numpy==1.26.4 packaging==24.1 ply==3.11 PyYAML==6.0.2 -six==1.16.0 sympy==1.13.3 diff --git a/ruff.toml b/ruff.toml new file mode 100644 index 0000000000..7eaf19de0e --- /dev/null +++ b/ruff.toml @@ -0,0 +1,5 @@ +[lint] +select = ["F401"] + +[lint.per-file-ignores] +"__init__.py" = ["F401"] diff --git a/samples/codegen/tensor_cores.py b/samples/codegen/tensor_cores.py index f516f33d23..16edeafbcc 100644 --- a/samples/codegen/tensor_cores.py +++ b/samples/codegen/tensor_cores.py @@ -17,8 +17,6 @@ from dace.codegen.targets.cpp import cpp_array_expr, cpp_offset_expr # Frontend imports and helpers -from dace.frontend.common.op_repository import replaces -from dace.frontend.python.newast import ProgramVisitor # Transformations from dace.transformation.interstate import GPUTransformSDFG @@ -28,7 +26,6 @@ from dace.sdfg.state import ControlFlowRegion, StateSubgraphView from dace.codegen.prettycode import CodeIOStream from dace.codegen.dispatcher import DefinedType -from typing import Any, List # Other imports import itertools diff --git a/samples/explicit/histogram.py b/samples/explicit/histogram.py index c2c6f29ad7..6524dfb525 100644 --- a/samples/explicit/histogram.py +++ b/samples/explicit/histogram.py @@ -2,7 +2,6 @@ """ 2D histogram sample that showcases memlets with write-conflict resolution and unknown element. """ import argparse import dace -import math import numpy as np W = dace.symbol('W') diff --git a/samples/instrumentation/matmul_likwid.py b/samples/instrumentation/matmul_likwid.py index e9d0ae4938..a08d3c9bac 100644 --- a/samples/instrumentation/matmul_likwid.py +++ b/samples/instrumentation/matmul_likwid.py @@ -2,8 +2,6 @@ import dace import numpy as np -import dace.transformation.helpers as xfh - M = dace.symbol('M') K = dace.symbol('K') N = dace.symbol('N') diff --git a/samples/instrumentation/matmul_likwid_gpu.py b/samples/instrumentation/matmul_likwid_gpu.py index 55c113cca1..0eaf6c8ba0 100644 --- a/samples/instrumentation/matmul_likwid_gpu.py +++ b/samples/instrumentation/matmul_likwid_gpu.py @@ -2,8 +2,6 @@ import dace import numpy as np -import dace.transformation.helpers as xfh - M = dace.symbol('M') K = dace.symbol('K') N = dace.symbol('N') diff --git a/samples/optimization/matmul.py b/samples/optimization/matmul.py index 9353a5aae9..41edb8a270 100644 --- a/samples/optimization/matmul.py +++ b/samples/optimization/matmul.py @@ -6,11 +6,11 @@ import click import dace import numpy as np -from typing import List, Tuple +from typing import Tuple # For optimizations -from dace.transformation.dataflow import (DoubleBuffering, MapCollapse, MapExpansion, MapReduceFusion, StripMining, - InLocalStorage, AccumulateTransient, Vectorization) +from dace.transformation.dataflow import (DoubleBuffering, MapCollapse, MapReduceFusion, InLocalStorage, + AccumulateTransient, Vectorization) from dace.transformation import helpers as xfutil # For library node implementations diff --git a/samples/sdfg_api/stencil_boundaries.py b/samples/sdfg_api/stencil_boundaries.py index e2bc0c9429..f464829c9c 100644 --- a/samples/sdfg_api/stencil_boundaries.py +++ b/samples/sdfg_api/stencil_boundaries.py @@ -4,7 +4,6 @@ import dace from dace import subsets import numpy as np -from scipy import signal H = dace.symbol('H') W = dace.symbol('W') diff --git a/setup.py b/setup.py index 63678f89aa..90f8e8d945 100644 --- a/setup.py +++ b/setup.py @@ -44,61 +44,62 @@ with open(os.path.join(dace_path, "version.py"), "r") as fp: version = fp.read().strip().split(' ')[-1][1:-1] -setup(name='dace', - version=version, - url='https://github.com/spcl/dace', - author='SPCL @ ETH Zurich', - author_email='talbn@inf.ethz.ch', - description='Data-Centric Parallel Programming Framework', - long_description=long_description, - long_description_content_type='text/markdown', - classifiers=[ - "Programming Language :: Python :: 3", - "License :: OSI Approved :: BSD License", - "Operating System :: OS Independent", - ], - python_requires='>=3.10, <3.15', - packages=find_packages(exclude=["*.tests", "*.tests.*", "tests.*", "tests"]), - package_data={ - '': [ - '*.yml', 'codegen/CMakeLists.txt', 'codegen/tools/*.cpp', 'external/moodycamel/*.h', - 'external/moodycamel/LICENSE.md' - ] + runtime_files + cub_files + viewer_files + library_files + cmake_files - }, - include_package_data=True, - install_requires=[ - 'numpy', 'networkx >= 2.5, <= 3.5', 'astunparse', 'sympy >= 1.9', 'pyyaml', 'ply', 'fparser >= 0.1.3', 'dill', - 'pyreadline;platform_system=="Windows"', 'packaging', 'typing-extensions' - ] + cmake_requires, - extras_require={ - 'ml': ['onnx', 'torch', 'onnxsim', 'onnxscript', 'onnxruntime', 'protobuf', 'ninja'], - 'testing': [ - 'coverage', - 'pytest-cov', - 'scipy', - 'absl-py', - 'opt_einsum', - 'pymlir', - 'click', - 'ipykernel', - 'nbconvert', - 'pytest-timeout', - ], - 'ml-testing': [ - 'coverage', 'pytest-cov', 'scipy', 'absl-py', 'opt_einsum', 'pymlir', 'click', 'ipykernel', 'nbconvert', - 'pytest-timeout', 'transformers == 4.50', 'jax <= 0.6.2', 'efficientnet_pytorch' - ], - 'docs': ['jinja2<3.2.0', 'sphinx-autodoc-typehints', 'sphinx-rtd-theme>=0.5.1'], - 'linting': ['pre-commit==4.1.0', 'yapf==0.43.0'], - }, - entry_points={ - 'console_scripts': [ - 'dacelab = dace.cli.dacelab:main', - 'sdfv = dace.cli.sdfv:main', - 'sdfgcc = dace.cli.sdfgcc:main', - 'sdfg-diff = dace.cli.sdfg_diff:main', - 'fcfd = dace.cli.fcdc:main', - 'daceprof = dace.cli.daceprof:main', - 'dace-external-transformation-registry = dace.cli.external_transformation_registry:main', - ], - }) +setup( + name='dace', + version=version, + url='https://github.com/spcl/dace', + author='SPCL @ ETH Zurich', + author_email='talbn@inf.ethz.ch', + description='Data-Centric Parallel Programming Framework', + long_description=long_description, + long_description_content_type='text/markdown', + classifiers=[ + "Programming Language :: Python :: 3", + "License :: OSI Approved :: BSD License", + "Operating System :: OS Independent", + ], + python_requires='>=3.10, <3.15', + packages=find_packages(exclude=["*.tests", "*.tests.*", "tests.*", "tests"]), + package_data={ + '': [ + '*.yml', 'codegen/CMakeLists.txt', 'codegen/tools/*.cpp', 'external/moodycamel/*.h', + 'external/moodycamel/LICENSE.md' + ] + runtime_files + cub_files + viewer_files + library_files + cmake_files + }, + include_package_data=True, + install_requires=[ + 'numpy', 'networkx >= 2.5, <= 3.5', 'astunparse', 'sympy >= 1.9', 'pyyaml', 'ply', 'fparser >= 0.1.3, != 0.2.3', + 'dill', 'pyreadline;platform_system=="Windows"', 'packaging', 'typing-extensions' + ] + cmake_requires, + extras_require={ + 'ml': ['onnx', 'torch', 'onnxsim', 'onnxscript', 'onnxruntime', 'protobuf', 'ninja'], + 'testing': [ + 'coverage', + 'pytest-cov', + 'scipy', + 'absl-py', + 'opt_einsum', + 'pymlir', + 'click', + 'ipykernel', + 'nbconvert', + 'pytest-timeout', + ], + 'ml-testing': [ + 'coverage', 'pytest-cov', 'scipy', 'absl-py', 'opt_einsum', 'pymlir', 'click', 'ipykernel', 'nbconvert', + 'pytest-timeout', 'transformers == 4.50', 'jax <= 0.6.2', 'efficientnet_pytorch' + ], + 'docs': ['jinja2<3.2.0', 'sphinx-autodoc-typehints', 'sphinx-rtd-theme>=0.5.1'], + 'linting': ['pre-commit==4.1.0', 'yapf==0.43.0'], + }, + entry_points={ + 'console_scripts': [ + 'dacelab = dace.cli.dacelab:main', + 'sdfv = dace.cli.sdfv:main', + 'sdfgcc = dace.cli.sdfgcc:main', + 'sdfg-diff = dace.cli.sdfg_diff:main', + 'fcfd = dace.cli.fcdc:main', + 'daceprof = dace.cli.daceprof:main', + 'dace-external-transformation-registry = dace.cli.external_transformation_registry:main', + ], + }) diff --git a/tests/autodiff/torch_backward/test_pytorch.py b/tests/autodiff/torch_backward/test_pytorch.py index 05d2f6948a..d2028c728f 100644 --- a/tests/autodiff/torch_backward/test_pytorch.py +++ b/tests/autodiff/torch_backward/test_pytorch.py @@ -1,5 +1,4 @@ # Copyright 2019-2025 ETH Zurich and the DaCe authors. All rights reserved. -import numpy as np import pytest import copy diff --git a/tests/autodiff/torch_backward/test_training.py b/tests/autodiff/torch_backward/test_training.py index 0858288145..616f1bab01 100644 --- a/tests/autodiff/torch_backward/test_training.py +++ b/tests/autodiff/torch_backward/test_training.py @@ -1,10 +1,6 @@ # Copyright 2019-2025 ETH Zurich and the DaCe authors. All rights reserved. -import os -import copy import pytest -import numpy as np - pytest.importorskip("torch", reason="PyTorch not installed. Please install with: pip install dace[ml]") pytest.importorskip("transformers", reason="transformers not installed. Please install with: pip install dace[ml-testing]") diff --git a/tests/blas/nodes/axpy_test.py b/tests/blas/nodes/axpy_test.py index c7feb2029f..4a13da3b12 100755 --- a/tests/blas/nodes/axpy_test.py +++ b/tests/blas/nodes/axpy_test.py @@ -2,18 +2,13 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. import numpy as np -import pytest -import argparse import scipy -import random import dace from dace.memlet import Memlet import dace.libraries.blas as blas -from dace.transformation.interstate import InlineSDFG -from dace.transformation.dataflow import StreamingMemory from dace.libraries.standard.memory import aligned_ndarray diff --git a/tests/blas/nodes/dot_test.py b/tests/blas/nodes/dot_test.py index c06f3a9eea..19b600c4f1 100755 --- a/tests/blas/nodes/dot_test.py +++ b/tests/blas/nodes/dot_test.py @@ -11,9 +11,6 @@ import dace.libraries.blas as blas -from dace.transformation.interstate import InlineSDFG -from dace.config import set_temporary - def pure_graph(implementation, dtype, veclen): diff --git a/tests/blas/nodes/gemv_test.py b/tests/blas/nodes/gemv_test.py index e5f0916924..c028b633e9 100644 --- a/tests/blas/nodes/gemv_test.py +++ b/tests/blas/nodes/gemv_test.py @@ -5,7 +5,6 @@ from dace.memlet import Memlet import dace.libraries.blas as blas -from dace.transformation.interstate import InlineSDFG def pure_graph(dtype, transposed, expansion, veclen, alpha, beta, expansion_args=None): diff --git a/tests/blockreduce_cudatest.py b/tests/blockreduce_cudatest.py index 5249a15533..b98c9a1b58 100644 --- a/tests/blockreduce_cudatest.py +++ b/tests/blockreduce_cudatest.py @@ -5,37 +5,39 @@ from dace.transformation.interstate import GPUTransformSDFG from dace.memlet import Memlet -# Create SDFG -sdfg = dace.SDFG('block_reduction') -state = sdfg.add_state('a') - -# Nodes -A = state.add_array('A', (128, ), dace.float32) -B = state.add_array('B', (2, ), dace.float32) -me, mx = state.add_map('mymap', dict(bi='0:2')) -mei, mxi = state.add_map('mymap2', dict(i='0:32')) -red = state.add_reduce('lambda a,b: a+b', None, 0) -red.implementation = 'CUDA (block)' -tA = state.add_transient('tA', (2, ), dace.float32) -tB = state.add_transient('tB', (1, ), dace.float32) -write_tasklet = state.add_tasklet('writeout', {'inp'}, {'out'}, 'if i == 0: out = inp') - -# Edges -state.add_edge(A, None, me, None, Memlet.simple(A, '0:128')) -state.add_edge(me, None, mei, None, Memlet.simple(A, '(64*bi):(64*bi+64)')) -state.add_edge(mei, None, tA, None, Memlet.simple('A', '(64*bi+2*i):(64*bi+2*i+2)')) -state.add_edge(tA, None, red, None, Memlet.simple(tA, '0:2')) -state.add_edge(red, None, tB, None, Memlet.simple(tB, '0')) -state.add_edge(tB, None, write_tasklet, 'inp', Memlet.simple(tB, '0')) -state.add_edge(write_tasklet, 'out', mxi, None, Memlet.simple('B', 'bi', num_accesses=-1)) -state.add_edge(mxi, None, mx, None, Memlet.simple(B, 'bi')) -state.add_edge(mx, None, B, None, Memlet.simple(B, '0:2')) -sdfg.fill_scope_connectors() - @pytest.mark.gpu def test_blockreduce(): - print('Block reduction test') + # Create SDFG + sdfg = dace.SDFG('block_reduction') + sdfg.add_array('A', (128, ), dace.float32) + sdfg.add_array('B', (2, ), dace.float32) + sdfg.add_transient('tA', (2, ), dace.float32) + sdfg.add_transient('tB', (1, ), dace.float32) + state = sdfg.add_state('a') + + # Nodes + A = state.add_access('A') + B = state.add_access('B') + me, mx = state.add_map('mymap', dict(bi='0:2')) + mei, mxi = state.add_map('mymap2', dict(i='0:32')) + red = state.add_reduce('lambda a,b: a+b', None, 0) + red.implementation = 'CUDA (block)' + tA = state.add_access('tA') + tB = state.add_access('tB') + write_tasklet = state.add_tasklet('writeout', {'inp'}, {'out'}, 'if i == 0: out = inp') + + # Edges + state.add_edge(A, None, me, None, Memlet.simple(A, '0:128')) + state.add_edge(me, None, mei, None, Memlet.simple(A, '(64*bi):(64*bi+64)')) + state.add_edge(mei, None, tA, None, Memlet.simple('A', '(64*bi+2*i):(64*bi+2*i+2)')) + state.add_edge(tA, None, red, None, Memlet.simple(tA, '0:2')) + state.add_edge(red, None, tB, None, Memlet.simple(tB, '0')) + state.add_edge(tB, None, write_tasklet, 'inp', Memlet.simple(tB, '0')) + state.add_edge(write_tasklet, 'out', mxi, None, Memlet.simple('B', 'bi', num_accesses=-1)) + state.add_edge(mxi, None, mx, None, Memlet.simple(B, 'bi')) + state.add_edge(mx, None, B, None, Memlet.simple(B, '0:2')) + sdfg.fill_scope_connectors() Adata = np.random.rand(128).astype(np.float32) Bdata = np.random.rand(2).astype(np.float32) @@ -47,7 +49,6 @@ def test_blockreduce(): B_regression[1] = np.sum(Adata[64:]) diff = np.linalg.norm(B_regression - Bdata) / 128.0 - print("Difference:", diff) assert diff <= 1e-5 diff --git a/tests/canonicalize_memlet_tree_test.py b/tests/canonicalize_memlet_tree_test.py index 7a44f7788c..e048bf9c11 100644 --- a/tests/canonicalize_memlet_tree_test.py +++ b/tests/canonicalize_memlet_tree_test.py @@ -1,13 +1,10 @@ # Copyright 2019-2025 ETH Zurich and the DaCe authors. All rights reserved. -from typing import Tuple, Union +from typing import Tuple import dace -from dace import subsets as dace_sbs from dace.sdfg import nodes as dace_nodes from dace.sdfg.utils import canonicalize_memlet_trees -import pytest - from .transformations import utility diff --git a/tests/chained_nested_tasklet_test.py b/tests/chained_nested_tasklet_test.py index 40c6b80251..1de54a4ad7 100644 --- a/tests/chained_nested_tasklet_test.py +++ b/tests/chained_nested_tasklet_test.py @@ -4,12 +4,10 @@ import dace as dp from dace.sdfg import SDFG from dace.memlet import Memlet -from dace.data import Scalar # Constructs an SDFG with two consecutive tasklets def test_nested_map(): - print('SDFG consecutive tasklet (nested) test') # Externals (parameters, symbols) N = dp.symbol('N') n = 20 @@ -20,9 +18,11 @@ def test_nested_map(): # Construct SDFG mysdfg = SDFG('ctasklet_nested_map') + mysdfg.add_array('A', [N], dp.int32) + mysdfg.add_array('B', [N], dp.int32) state = mysdfg.add_state() - A_ = state.add_array('A', [N], dp.int32) - B_ = state.add_array('B', [N], dp.int32) + A_ = state.add_access('A') + B_ = state.add_access('B') omap_entry, omap_exit = state.add_map('omap', dict(k='0:2')) map_entry, map_exit = state.add_map('mymap', dict(i='0:N/2')) @@ -45,12 +45,10 @@ def test_nested_map(): mysdfg(A=input, B=output, N=n) diff = np.linalg.norm(10 * input - output) / n - print("Difference:", diff) assert diff <= 1e-5 def test_nested_sdfg(): - print('SDFG consecutive tasklet (nested SDFG) test') # Externals (parameters, symbols) N = dp.symbol('N') n = 20 @@ -61,15 +59,19 @@ def test_nested_sdfg(): # Construct outer SDFG mysdfg = SDFG('ctasklet_nested_sdfg') + mysdfg.add_array('A', [N], dp.int32) + mysdfg.add_array('B', [N], dp.int32) state = mysdfg.add_state() - A_ = state.add_array('A', [N], dp.int32) - B_ = state.add_array('B', [N], dp.int32) + A_ = state.add_access('A') + B_ = state.add_access('B') # Construct inner SDFG nsdfg = dp.SDFG('ctasklet_nested_sdfg_inner') + nsdfg.add_array('a', [N], dp.int32) + nsdfg.add_array('b', [N], dp.int32) nstate = nsdfg.add_state() - a = nstate.add_array('a', [N], dp.int32) - b = nstate.add_array('b', [N], dp.int32) + a = nstate.add_access('a') + b = nstate.add_access('b') map_entry, map_exit = nstate.add_map('mymap', dict(i='0:N/2')) tasklet = nstate.add_tasklet('mytasklet', {'aa'}, {'bb'}, 'bb = 5*aa') nstate.add_memlet_path(a, map_entry, tasklet, dst_conn='aa', memlet=Memlet('a[k*N/2+i]')) @@ -87,7 +89,6 @@ def test_nested_sdfg(): mysdfg(A=input, B=output, N=n) diff = np.linalg.norm(10 * input - output) / n - print("Difference:", diff) assert diff <= 1e-5 mysdfg.simplify() @@ -95,7 +96,6 @@ def test_nested_sdfg(): mysdfg(A=input, B=output, N=n) diff = np.linalg.norm(10 * input - output) / n - print("Difference:", diff) assert diff <= 1e-5 diff --git a/tests/chained_tasklet_test.py b/tests/chained_tasklet_test.py index 116db82613..a335f3a012 100644 --- a/tests/chained_tasklet_test.py +++ b/tests/chained_tasklet_test.py @@ -8,7 +8,6 @@ # Constructs an SDFG with two consecutive tasklets def test(): - print('SDFG consecutive tasklet test') # Externals (parameters, symbols) N = dp.symbol('N') n = 20 @@ -19,9 +18,11 @@ def test(): # Construct SDFG mysdfg = SDFG('ctasklet') + mysdfg.add_array('A', [N], dp.int32) + mysdfg.add_array('B', [N], dp.int32) state = mysdfg.add_state() - A_ = state.add_array('A', [N], dp.int32) - B_ = state.add_array('B', [N], dp.int32) + A_ = state.add_access('A') + B_ = state.add_access('B') map_entry, map_exit = state.add_map('mymap', dict(i='0:N')) tasklet = state.add_tasklet('mytasklet', {'a'}, {'b'}, 'b = 5*a') @@ -37,7 +38,6 @@ def test(): mysdfg(A=input, B=output, N=n) diff = np.linalg.norm(10 * input - output) / n - print("Difference:", diff) assert diff <= 1e-5 diff --git a/tests/codegen/allocation_lifetime_test.py b/tests/codegen/allocation_lifetime_test.py index 08b4c8d72c..b1278822a3 100644 --- a/tests/codegen/allocation_lifetime_test.py +++ b/tests/codegen/allocation_lifetime_test.py @@ -287,7 +287,7 @@ def test_alloc_multistate(): init = sdfg.add_state() end = sdfg.add_state() s2 = sdfg.add_state() - sdfg.add_loop(init, s2, end, 'i', '0', 'i < 5', 'i + 1') + sdfg.add_loop_state_machine(init, s2, end, 'i', '0', 'i < 5', 'i + 1') s1 = sdfg.add_state_before(s2) diff --git a/tests/codegen/argumet_signature_test.py b/tests/codegen/argument_signature_test.py similarity index 99% rename from tests/codegen/argumet_signature_test.py rename to tests/codegen/argument_signature_test.py index c29a63ac2e..e4b720a289 100644 --- a/tests/codegen/argumet_signature_test.py +++ b/tests/codegen/argument_signature_test.py @@ -1,5 +1,4 @@ import dace -import copy def test_argument_signature_test(): @@ -190,7 +189,7 @@ def make_sdfg() -> dace.SDFG: # If we have cupy we will also compile it. try: - import cupy as cp + import cupy as cp # noqa: F401 except ImportError: return diff --git a/tests/codegen/control_flow_generation_test.py b/tests/codegen/control_flow_generation_test.py index 0ee8d8f852..f337570780 100644 --- a/tests/codegen/control_flow_generation_test.py +++ b/tests/codegen/control_flow_generation_test.py @@ -4,7 +4,7 @@ import dace import numpy as np -from dace.sdfg.state import ConditionalBlock, ReturnBlock +from dace.sdfg.state import ConditionalBlock from dace.transformation.pass_pipeline import FixedPointPipeline from dace.transformation.passes.simplification.control_flow_raising import ControlFlowRaising @@ -39,7 +39,8 @@ def looptest(): sdfg: dace.SDFG = looptest.to_sdfg(simplify=True) if dace.Config.get_bool('optimizer', 'detect_control_flow'): - assert 'while (' in sdfg.generate_code()[0].code + code = sdfg.generate_code()[0].clean_code + assert 'while (' in code or 'for (' in code A = looptest() A_ref = np.array([0, 0, 2, 0, 4, 0, 6, 0, 8, 0], dtype=np.int32) diff --git a/tests/codegen/data_instrumentation_test.py b/tests/codegen/data_instrumentation_test.py index 8192679c9f..4a3e458336 100644 --- a/tests/codegen/data_instrumentation_test.py +++ b/tests/codegen/data_instrumentation_test.py @@ -1,5 +1,5 @@ # Copyright 2019-2022 ETH Zurich and the DaCe authors. All rights reserved. -from typing import Optional, Tuple +from typing import Optional import dace from dace import nodes from dace.properties import CodeBlock diff --git a/tests/codegen/gpu_memcpy_test.py b/tests/codegen/gpu_memcpy_test.py index c7a3525f95..b7ae974483 100644 --- a/tests/codegen/gpu_memcpy_test.py +++ b/tests/codegen/gpu_memcpy_test.py @@ -5,7 +5,6 @@ from dace.sdfg import nodes as dace_nodes import pytest -import copy import re # this test requires cupy module diff --git a/tests/codegen/gpu_min_warps_per_eu_test.py b/tests/codegen/gpu_min_warps_per_eu_test.py new file mode 100644 index 0000000000..9581b0d4d0 --- /dev/null +++ b/tests/codegen/gpu_min_warps_per_eu_test.py @@ -0,0 +1,27 @@ +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. + +import dace +import pytest + + +@pytest.mark.gpu +def test_min_warps_per_eu() -> None: + + @dace.program + def prog(a: dace.float64[100, 20] @ dace.StorageType.GPU_Global) -> None: + for i, j in dace.map[0:100, 0:20] @ dace.ScheduleType.GPU_Device: + a[i, j] = 1 + + sdfg = prog.to_sdfg() + for n, _ in sdfg.all_nodes_recursive(): + if isinstance(n, dace.nodes.MapEntry) and n.map.schedule == dace.ScheduleType.GPU_Device: + mapentry = n + + assert '__launch_bounds__(32)' in sdfg.generate_code()[1].code + mapentry.map.gpu_min_warps_per_eu = 4 + assert '__launch_bounds__(32,4)' in sdfg.generate_code( + )[1].code and '__launch_bounds__(32)' not in sdfg.generate_code()[1].code + + +if __name__ == '__main__': + test_min_warps_per_eu() diff --git a/tests/codegen/map_launch_test.py b/tests/codegen/map_launch_test.py index c58c550aeb..908751a1f3 100644 --- a/tests/codegen/map_launch_test.py +++ b/tests/codegen/map_launch_test.py @@ -2,7 +2,6 @@ import dace import numpy as np import copy -import sys import time import pytest diff --git a/tests/codegen/mpi_axpy.py b/tests/codegen/mpi_axpy.py index eb51030f21..9cfc29c39e 100644 --- a/tests/codegen/mpi_axpy.py +++ b/tests/codegen/mpi_axpy.py @@ -3,7 +3,6 @@ import argparse import dace import numpy as np -import scipy as sp from mpi4py import MPI from dace.transformation.dataflow import MPITransformMap diff --git a/tests/codegen/no_emit_copy_on_empty_memlet_test.py b/tests/codegen/no_emit_copy_on_empty_memlet_test.py index f521d7ac01..8630c580cd 100644 --- a/tests/codegen/no_emit_copy_on_empty_memlet_test.py +++ b/tests/codegen/no_emit_copy_on_empty_memlet_test.py @@ -1,6 +1,5 @@ # Copyright 2019-2025 ETH Zurich and the DaCe authors. All rights reserved. import dace -import pytest def _gen_sdfg() -> dace.SDFG: diff --git a/tests/codegen/sve/ast_test.py b/tests/codegen/sve/ast_test.py index 6cb896bd37..196f2a5f6d 100644 --- a/tests/codegen/sve/ast_test.py +++ b/tests/codegen/sve/ast_test.py @@ -4,7 +4,6 @@ import dace.dtypes from tests.codegen.sve.common import get_code import pytest -from dace.codegen.targets.sve.type_compatibility import IncompatibleTypeError N = dace.symbol('N') M = dace.symbol('M') diff --git a/tests/codegen/sve/common.py b/tests/codegen/sve/common.py index 922fe6791a..9445b8f9e8 100644 --- a/tests/codegen/sve/common.py +++ b/tests/codegen/sve/common.py @@ -1,6 +1,4 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. -from dace.dtypes import vector -import dace from dace.transformation.dataflow.sve.vectorization import SVEVectorization diff --git a/tests/codegen/sve/stream_test.py b/tests/codegen/sve/stream_test.py index 427059eca1..33d26306ad 100644 --- a/tests/codegen/sve/stream_test.py +++ b/tests/codegen/sve/stream_test.py @@ -1,7 +1,6 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. import dace from tests.codegen.sve.common import get_code -import pytest N = dace.symbol('N') diff --git a/tests/codegen/sve/wcr_test.py b/tests/codegen/sve/wcr_test.py index a311084b6f..97967c9807 100644 --- a/tests/codegen/sve/wcr_test.py +++ b/tests/codegen/sve/wcr_test.py @@ -1,8 +1,6 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. import dace from tests.codegen.sve.common import get_code -import pytest -from dace.codegen.targets.sve.util import NotSupportedError N = dace.symbol('N') diff --git a/tests/codegen/tasklet_generation_test.py b/tests/codegen/tasklet_generation_test.py index 79fb7664e2..fd38ee7191 100644 --- a/tests/codegen/tasklet_generation_test.py +++ b/tests/codegen/tasklet_generation_test.py @@ -1,5 +1,4 @@ import dace -import pytest def _gen_sdfg_with_a_print_tasklet_between_map_exits() -> dace.SDFG: diff --git a/tests/codegen/transient_same_name_test.py b/tests/codegen/transient_same_name_test.py index 7d799a33a7..15668399b2 100644 --- a/tests/codegen/transient_same_name_test.py +++ b/tests/codegen/transient_same_name_test.py @@ -1,7 +1,6 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. import numpy as np import dace -from dace.transformation.interstate import StateFusion @dace.program diff --git a/tests/codegen/wcr_atomic_test.py b/tests/codegen/wcr_atomic_test.py index f1495282b9..a9892ba957 100644 --- a/tests/codegen/wcr_atomic_test.py +++ b/tests/codegen/wcr_atomic_test.py @@ -1,7 +1,6 @@ # Copyright 2019-2023 ETH Zurich and the DaCe authors. All rights reserved. """ Tests atomic WCR detection in code generation. """ import dace -import numpy as np N = dace.symbol('N') diff --git a/tests/compile_folder_mode_test.py b/tests/compile_folder_mode_test.py new file mode 100644 index 0000000000..80b734ae89 --- /dev/null +++ b/tests/compile_folder_mode_test.py @@ -0,0 +1,240 @@ +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. +import numpy as np +import pytest +import uuid +import pathlib +import copy +import re + +import dace +from dace.codegen import compiler as sdfg_compiler + + +def _make_test_sdfg() -> dace.SDFG: + sdfg = dace.SDFG("test_sdfg_" + str(uuid.uuid1()).replace("-", "_")) + state = sdfg.add_state() + for name in "abc": + sdfg.add_array( + name, + shape=(10, ), + dtype=dace.float64, + transient=False, + ) + + state.add_mapped_tasklet( + "comp", + map_ranges={"__i": "0:10"}, + inputs={ + "__in1": dace.Memlet("a[__i]"), + "__in2": dace.Memlet("b[__i]"), + }, + outputs={ + "__out": dace.Memlet("c[__i]"), + }, + code="__out = __in1 + __in2", + external_edges=True, + ) + sdfg.validate() + + return sdfg + + +def _run_sdfg(csdfg): + res = {name: np.array(np.random.rand(10), copy=True, dtype=np.float64) for name in "abc"} + ref = copy.deepcopy(res) + ref["c"] = ref["a"] + ref["b"] + + csdfg(**res) + + assert all(np.allclose(ref[n], res[n]) for n in res.keys()) + + +def _load_and_run_sdfg(build_folder, sdfg): + csdfg = sdfg_compiler.load_precompiled_sdfg(build_folder, sdfg) + _run_sdfg(csdfg) + + +def test_development_folder_mode(): + with dace.config.temporary_config() as Config: + Config.set('compiler', 'build_folder_mode', value="development") + sdfg = _make_test_sdfg() + build_folder = pathlib.Path(sdfg.build_folder) + assert not build_folder.exists() + + not_csdfg = sdfg.compile(return_program_handle=False) + assert not_csdfg is None + + assert build_folder.exists() + assert sdfg_compiler.get_folder_mode(build_folder) == "development" + + expected_files = { + "build": pathlib.Path.is_dir, + "include": pathlib.Path.is_dir, + "map": pathlib.Path.is_dir, + "perf": pathlib.Path.is_dir, + "sample": pathlib.Path.is_dir, + "src": pathlib.Path.is_dir, + "CACHEDIR.TAG": pathlib.Path.is_file, + "FOLDER_MODE": pathlib.Path.is_file, + "dace.conf": pathlib.Path.is_file, + "dace_files.csv": pathlib.Path.is_file, + "dace_environments.csv": pathlib.Path.is_file, + "program.sdfgz": pathlib.Path.is_file, + } + + for found_path in build_folder.iterdir(): + found_file = found_path.name + assert found_file in expected_files + assert expected_files[found_file](found_path) + + # Now run it. + _load_and_run_sdfg(build_folder, sdfg) + + # Special case for development is that if there is no `FOLDER_MODE` it is still development. + version_file = build_folder / "FOLDER_MODE" + assert version_file.exists() + version_file.unlink() + assert not version_file.exists() + assert sdfg_compiler.get_folder_mode(build_folder) == "development" + + +def test_production_folder_mode(): + with dace.config.temporary_config() as Config: + Config.set('compiler', 'build_folder_mode', value="production") + sdfg = _make_test_sdfg() + build_folder = pathlib.Path(sdfg.build_folder) + assert not build_folder.exists() + + not_csdfg = sdfg.compile(return_program_handle=False) + assert not_csdfg is None + + assert build_folder.exists() + assert sdfg_compiler.get_folder_mode(build_folder) == "production" + + lib_path = sdfg_compiler.get_binary_name(build_folder, sdfg_name=sdfg.name, folder_mode="production") + libstub_path = sdfg_compiler._get_stub_library_path(lib_path.name) + + expected_files = { + "program.sdfgz": pathlib.Path.is_file, + "CACHEDIR.TAG": pathlib.Path.is_file, + "FOLDER_MODE": pathlib.Path.is_file, + lib_path.name: pathlib.Path.is_file, + libstub_path.name: pathlib.Path.is_file, + } + + for found_path in build_folder.iterdir(): + found_file = found_path.name + assert found_file in expected_files + assert expected_files[found_file](found_path) + + # Now run it. + _load_and_run_sdfg(build_folder, sdfg) + + # If we delete the FOLDER_MODE file we will get an error. + version_file = build_folder / "FOLDER_MODE" + assert version_file.exists() + version_file.unlink() + assert not version_file.exists() + + with pytest.raises(NotADirectoryError, + match=re.escape(f'``{build_folder}`` does not appear to be a valid build folder.')): + sdfg_compiler.get_folder_mode(build_folder) + + assert sdfg_compiler.get_folder_mode(build_folder, probe=True) is None + + +def _test_build_with_scheme_one_and_then_switch_impl( + version1: str, + version2: str, +) -> None: + with dace.config.temporary_config() as conf: + conf.set('compiler', 'build_folder_mode', value=version1) + + sdfg = _make_test_sdfg() + build_folder = pathlib.Path(sdfg.build_folder).resolve() + assert not build_folder.exists() + + csdfg1 = sdfg.compile() + + lib1_path = pathlib.Path(csdfg1.filename) + assert sdfg_compiler.get_folder_mode(build_folder) == version1 + assert lib1_path.exists() + + with dace.config.temporary_config() as conf: + conf.set('compiler', 'build_folder_mode', value=version1) + + # This is for ensuring that the code is actually regenerated. + conf.set('compiler', 'use_cache', value=False) + sdfg._recompile = True + sdfg.regenerate_code = True + + csdfg2 = sdfg.compile() + + lib2_path = pathlib.Path(csdfg2.filename) + + assert sdfg_compiler.get_folder_mode(build_folder) == version1 + assert csdfg1.sdfg.name != csdfg2.sdfg.name + assert lib1_path != lib2_path + assert lib1_path.exists() # Still existing. + assert lib2_path.exists() + assert str(lib1_path).startswith(str(build_folder)) + assert str(lib2_path).startswith(str(build_folder)) + + if version1 == "production": + assert lib1_path.parent == build_folder + assert lib2_path.parent == build_folder + libstub1_path = sdfg_compiler._get_stub_library_path(lib1_path.name) + libstub2_path = sdfg_compiler._get_stub_library_path(lib2_path.name) + + expected_files = { + "CACHEDIR.TAG": pathlib.Path.is_file, + "FOLDER_MODE": pathlib.Path.is_file, + "program.sdfgz": pathlib.Path.is_file, + lib1_path.name: pathlib.Path.is_file, + libstub1_path.name: pathlib.Path.is_file, + lib2_path.name: pathlib.Path.is_file, + libstub2_path.name: pathlib.Path.is_file, + } + for found_path in build_folder.iterdir(): + found_file = found_path.name + assert found_file in expected_files + assert expected_files[found_file](found_path) + + else: + assert str(lib1_path).startswith(str(build_folder / "build")) + assert str(lib2_path).startswith(str(build_folder / "build")) + assert (build_folder / "build").is_dir() + assert (build_folder / "src").is_dir() + assert (build_folder / "include").is_dir() + + _run_sdfg(csdfg1) + _run_sdfg(csdfg2) + + +def test_build_with_scheme_one_and_then_switch(): + _test_build_with_scheme_one_and_then_switch_impl( + version1="development", + version2="production", + ) + _test_build_with_scheme_one_and_then_switch_impl( + version1="production", + version2="development", + ) + + +def test_already_loaded_and_comple_again(): + _test_build_with_scheme_one_and_then_switch_impl( + version1="development", + version2="development", + ) + _test_build_with_scheme_one_and_then_switch_impl( + version1="production", + version2="production", + ) + + +if __name__ == '__main__': + test_development_folder_mode() + test_production_folder_mode() + test_already_loaded_and_comple_again() + test_build_with_scheme_one_and_then_switch() diff --git a/tests/compile_sdfg_test.py b/tests/compile_sdfg_test.py index 9a712ae856..9813ef77e3 100644 --- a/tests/compile_sdfg_test.py +++ b/tests/compile_sdfg_test.py @@ -19,15 +19,17 @@ def test(): # Construct SDFG mysdfg = SDFG('mysdfg') + mysdfg.add_array('A', [N], dp.int32) + mysdfg.add_array('B', [N], dp.int32) state = mysdfg.add_state() - A_ = state.add_array('A', [N], dp.int32) # NOTE: The names A and B are not - B_ = state.add_array('B', [N], dp.int32) # reserved, this is just to - # clarify that + # NOTE: The names A and B are not reserved, this is just to clarify that # variable name != array name + A_ = state.add_access('A') + B_ = state.add_access('B') # Easy way to add a tasklet - tasklet, map_entry, map_exit = state.add_mapped_tasklet('mytasklet', dict(i='0:N'), dict(a=Memlet.simple(A_, 'i')), - 'b = 5*a', dict(b=Memlet.simple(B_, 'i'))) + _, map_entry, map_exit = state.add_mapped_tasklet('mytasklet', dict(i='0:N'), dict(a=Memlet.simple(A_, 'i')), + 'b = 5*a', dict(b=Memlet.simple(B_, 'i'))) # Alternatively (the explicit way): #map_entry, map_exit = state.add_map('mymap', dict(i='0:N')) #tasklet = state.add_tasklet('mytasklet', {'a'}, {'b'}, 'b = 5*a') @@ -41,7 +43,6 @@ def test(): mysdfg(A=input, B=output, N=n) diff = np.linalg.norm(5 * input - output) / n - print("Difference:", diff) assert diff <= 1e-5 diff --git a/tests/config_test.py b/tests/config_test.py index 7a62d58a0e..9b2b217bd8 100644 --- a/tests/config_test.py +++ b/tests/config_test.py @@ -1,7 +1,6 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. from dace.config import Config, set_temporary, temporary_config import threading -import time def test_set_temporary(): diff --git a/tests/const_utilities_test.py b/tests/const_utilities_test.py index 8fe86f775c..4d0d37a38e 100644 --- a/tests/const_utilities_test.py +++ b/tests/const_utilities_test.py @@ -2,8 +2,7 @@ import copy import dace import dace.sdfg.utils as sdutils -import pytest -from dace.transformation.passes.analysis.scope_data_and_symbol_analysis import ScopeDataAndSymbolAnalysis, ScopeAnalysis +from dace.transformation.passes.analysis.scope_data_and_symbol_analysis import ScopeDataAndSymbolAnalysis def _add_shared_memory(sdfg: dace.SDFG, add_src_access_node: bool = False): diff --git a/tests/consume_chunk_cond_test.py b/tests/consume_chunk_cond_test.py index b2e252076c..5cfef81c91 100644 --- a/tests/consume_chunk_cond_test.py +++ b/tests/consume_chunk_cond_test.py @@ -2,21 +2,28 @@ import dace as dp import numpy as np -sdfg = dp.SDFG('fib_consume_cc') -state = sdfg.add_state('state') -nprocs = 4 -# Arrays -initial_value = state.add_array('iv', [1], dp.int32) -stream = state.add_stream('S', dp.int32, transient=True, buffer_size=256) -stream_init = state.add_stream('S', dp.int32, transient=True) -stream_out = state.add_stream('S', dp.int32, transient=True) -output = state.add_array('res', [1], dp.float32) +def test(): + """Fibonacci recursion using consume (with chunks, custom condition""" + + sdfg = dp.SDFG('fib_consume_cc') + sdfg.add_array('iv', [1], dp.int32) + sdfg.add_array('res', [1], dp.float32) + sdfg.add_stream('S', dp.int32, transient=True, buffer_size=256) + state = sdfg.add_state('state') + nprocs = 4 + + # Arrays + initial_value = state.add_access('iv') + stream = state.add_access('S') + stream_init = state.add_access('S') + stream_out = state.add_access('S') + output = state.add_access('res') -# Consume and tasklet -consume_entry, consume_exit = state.add_consume('cons', ('p', str(nprocs)), 'res[0] >= 44', chunksize=2) -tasklet = state.add_tasklet( - 'fibonacci', {'s'}, {'sout', 'val'}, """ + # Consume and tasklet + consume_entry, consume_exit = state.add_consume('cons', ('p', str(nprocs)), 'res[0] >= 44', chunksize=2) + tasklet = state.add_tasklet( + 'fibonacci', {'s'}, {'sout', 'val'}, """ for i in range(__dace_cons_numelems): if s[i] == 1: val = 1 @@ -25,32 +32,29 @@ sout = s[i] - 2 """) -# Edges -state.add_nedge(initial_value, stream_init, dp.Memlet.from_array(stream_init.data, stream_init.desc(sdfg))) -e = state.add_edge(stream, None, consume_entry, 'IN_stream', dp.Memlet.from_array(stream.data, stream.desc(sdfg))) + # Edges + state.add_nedge(initial_value, stream_init, dp.Memlet.from_array(stream_init.data, stream_init.desc(sdfg))) + e = state.add_edge(stream, None, consume_entry, 'IN_stream', dp.Memlet.from_array(stream.data, stream.desc(sdfg))) -# FIXME: Due to how memlets and propagation work, force access to stream to -# use an array instead of a scalar. -e.data.allow_oob = True -memlet = dp.Memlet.simple(stream, '0:2') -memlet.allow_oob = True + # FIXME: Due to how memlets and propagation work, force access to stream to + # use an array instead of a scalar. + e.data.allow_oob = True + memlet = dp.Memlet.simple(stream, '0:2') + memlet.allow_oob = True -state.add_edge(consume_entry, 'OUT_stream', tasklet, 's', memlet) -state.add_edge(tasklet, 'sout', consume_exit, 'IN_S', dp.Memlet.simple(stream_out, '0', num_accesses=-1)) -state.add_edge(consume_exit, 'OUT_S', stream_out, None, dp.Memlet.simple(stream_out, '0', num_accesses=-1)) -state.add_edge(tasklet, 'val', consume_exit, 'IN_V', - dp.Memlet.simple(output, '0', wcr_str='lambda a,b: a+b', num_accesses=-1)) -state.add_edge(consume_exit, 'OUT_V', output, None, - dp.Memlet.simple(output, '0', wcr_str='lambda a,b: a+b', num_accesses=-1)) + state.add_edge(consume_entry, 'OUT_stream', tasklet, 's', memlet) + state.add_edge(tasklet, 'sout', consume_exit, 'IN_S', dp.Memlet.simple(stream_out, '0', num_accesses=-1)) + state.add_edge(consume_exit, 'OUT_S', stream_out, None, dp.Memlet.simple(stream_out, '0', num_accesses=-1)) + state.add_edge(tasklet, 'val', consume_exit, 'IN_V', + dp.Memlet.simple(output, '0', wcr_str='lambda a,b: a+b', num_accesses=-1)) + state.add_edge(consume_exit, 'OUT_V', output, None, + dp.Memlet.simple(output, '0', wcr_str='lambda a,b: a+b', num_accesses=-1)) -consume_exit.add_in_connector('IN_S') -consume_exit.add_in_connector('IN_V') -consume_exit.add_out_connector('OUT_S') -consume_exit.add_out_connector('OUT_V') + consume_exit.add_in_connector('IN_S') + consume_exit.add_in_connector('IN_V') + consume_exit.add_out_connector('OUT_S') + consume_exit.add_out_connector('OUT_V') - -def test(): - print('Fibonacci recursion using consume (with chunks, custom condition)') input = np.ndarray([1], np.int32) output = np.ndarray([1], np.float32) input[0] = 10 @@ -60,7 +64,6 @@ def test(): sdfg(iv=input, res=output) diff = output[0] - regression - print('Difference:', diff) # Allowing for race conditions on quiescence condition assert not (diff < 0 or diff > nprocs) diff --git a/tests/consume_test.py b/tests/consume_test.py index 9abba4250d..c203016805 100644 --- a/tests/consume_test.py +++ b/tests/consume_test.py @@ -2,20 +2,25 @@ import dace as dp import numpy as np -sdfg = dp.SDFG('fib_consume') -state = sdfg.add_state('state') -# Arrays -initial_value = state.add_array('iv', [1], dp.int32) -stream = state.add_stream('S', dp.int32, transient=True) -stream_init = state.add_stream('S', dp.int32, transient=True) -stream_out = state.add_stream('S', dp.int32, transient=True) -output = state.add_array('res', [1], dp.float32) +def setup_sdfg() -> dp.SDFG: + sdfg = dp.SDFG('fib_consume') + sdfg.add_array('iv', [1], dp.int32) + sdfg.add_stream('S', dp.int32, transient=True) + sdfg.add_array('res', [1], dp.float32) -# Consume and tasklet -consume_entry, consume_exit = state.add_consume('cons', ('p', '4')) -tasklet = state.add_tasklet( - 'fibonacci', {'s'}, {'sout', 'val'}, """ + # Arrays + state = sdfg.add_state('state') + initial_value = state.add_access('iv') + stream = state.add_access('S') + stream_init = state.add_access('S') + stream_out = state.add_access('S') + output = state.add_access('res') + + # Consume and tasklet + consume_entry, consume_exit = state.add_consume('cons', ('p', '4')) + tasklet = state.add_tasklet( + 'fibonacci', {'s'}, {'sout', 'val'}, """ if s == 1: val = 1 elif s > 1: @@ -23,21 +28,23 @@ sout = s - 2 """) -# Edges -state.add_nedge(initial_value, stream_init, dp.Memlet.from_array(stream_init.data, stream_init.desc(sdfg))) -state.add_edge(stream, None, consume_entry, 'IN_stream', dp.Memlet.from_array(stream.data, stream.desc(sdfg))) -state.add_edge(consume_entry, 'OUT_stream', tasklet, 's', dp.Memlet.from_array(stream.data, stream.desc(sdfg))) -state.add_edge(tasklet, 'sout', consume_exit, 'IN_S', dp.Memlet.simple(stream_out, '0', num_accesses=-1)) -state.add_edge(consume_exit, 'OUT_S', stream_out, None, dp.Memlet.simple(stream_out, '0', num_accesses=-1)) -state.add_edge(tasklet, 'val', consume_exit, 'IN_V', - dp.Memlet.simple(output, '0', wcr_str='lambda a,b: a+b', num_accesses=-1)) -state.add_edge(consume_exit, 'OUT_V', output, None, - dp.Memlet.simple(output, '0', wcr_str='lambda a,b: a+b', num_accesses=-1)) + # Edges + state.add_nedge(initial_value, stream_init, dp.Memlet.from_array(stream_init.data, stream_init.desc(sdfg))) + state.add_edge(stream, None, consume_entry, 'IN_stream', dp.Memlet.from_array(stream.data, stream.desc(sdfg))) + state.add_edge(consume_entry, 'OUT_stream', tasklet, 's', dp.Memlet.from_array(stream.data, stream.desc(sdfg))) + state.add_edge(tasklet, 'sout', consume_exit, 'IN_S', dp.Memlet.simple(stream_out, '0', num_accesses=-1)) + state.add_edge(consume_exit, 'OUT_S', stream_out, None, dp.Memlet.simple(stream_out, '0', num_accesses=-1)) + state.add_edge(tasklet, 'val', consume_exit, 'IN_V', + dp.Memlet.simple(output, '0', wcr_str='lambda a,b: a+b', num_accesses=-1)) + state.add_edge(consume_exit, 'OUT_V', output, None, + dp.Memlet.simple(output, '0', wcr_str='lambda a,b: a+b', num_accesses=-1)) + + consume_exit.add_in_connector('IN_S') + consume_exit.add_in_connector('IN_V') + consume_exit.add_out_connector('OUT_S') + consume_exit.add_out_connector('OUT_V') -consume_exit.add_in_connector('IN_S') -consume_exit.add_in_connector('IN_V') -consume_exit.add_out_connector('OUT_S') -consume_exit.add_out_connector('OUT_V') + return sdfg def fibonacci(v): @@ -49,20 +56,19 @@ def fibonacci(v): return fibonacci(v - 1) + fibonacci(v - 2) -def test(): - print('Fibonacci recursion using consume') +def test_fibonacci_recursion_using_consume(): input = np.ndarray([1], np.int32) output = np.ndarray([1], np.float32) input[0] = 10 output[0] = 0 regression = fibonacci(input[0]) + sdfg = setup_sdfg() sdfg(iv=input, res=output) diff = (regression - output[0])**2 - print('Difference:', diff) assert diff <= 1e-5 if __name__ == '__main__': - test() + test_fibonacci_recursion_using_consume() diff --git a/tests/copy_serialization_test.py b/tests/copy_serialization_test.py new file mode 100644 index 0000000000..3d88b0da39 --- /dev/null +++ b/tests/copy_serialization_test.py @@ -0,0 +1,115 @@ +import dace +from dace.sdfg import nodes as dace_nodes +from dace.sdfg import graph as dace_graph +from dace.symbolic import pystr_to_symbolic as s2s +from dace.codegen.targets import cpp as dace_cpp + + +def _make_sdfg() -> tuple[dace.SDFG, dace.SDFGState, dace_nodes.AccessNode, dace_graph.MultiConnectorEdge[dace.Memlet], + dace_nodes.AccessNode]: + sdfg = dace.SDFG("copy_sdfg") + state = sdfg.add_state() + + for sname in [ + "__out_IDim_range_0", "__out_IDim_range_1", "__out_IDim_stride", "__out_JDim_range_0", "__out_JDim_range_1", + "__out_JDim_stride", "__out_KDim_range_0", "__out_KDim_range_1", "__out_KDim_stride" + ]: + sdfg.add_symbol(sname, dace.int32) + + sdfg.add_array( + "a", + dtype=dace.float64, + transient=False, + shape=( + s2s("max(0, -__out_IDim_range_0 + __out_IDim_range_1)"), + s2s("max(0, -__out_JDim_range_0 + __out_JDim_range_1)"), + s2s("max(0, -max(0, __out_KDim_range_0) + __out_KDim_range_1)"), + ), + strides=( + 1, + s2s("max(0, -__out_IDim_range_0 + __out_IDim_range_1)"), + s2s("max(0, -__out_IDim_range_0 + __out_IDim_range_1)*max(0, -__out_JDim_range_0 + __out_JDim_range_1)"), + ), + ) + sdfg.add_array( + "b", + dtype=dace.float64, + transient=False, + shape=( + s2s("max(0, -__out_IDim_range_0 + __out_IDim_range_1)"), + s2s("max(0, -__out_JDim_range_0 + __out_JDim_range_1)"), + s2s("max(0, -__out_KDim_range_0 + __out_KDim_range_1)"), + ), + strides=( + 1, + s2s("max(0, -__out_IDim_range_0 + __out_IDim_range_1)"), + s2s("max(0, -__out_IDim_range_0 + __out_IDim_range_1)*max(0, -__out_JDim_range_0 + __out_JDim_range_1)"), + ), + ) + + a = state.add_access("a") + b = state.add_access("b") + e = state.add_nedge( + a, b, + dace.Memlet( + data=a.data, + subset= + "0:max(0, __out_IDim_range_1 - __out_IDim_range_0), 0:max(0, __out_JDim_range_1 - __out_JDim_range_0), 0:max(0, __out_KDim_range_0, __out_KDim_range_1) - max(0, __out_KDim_range_0)", + other_subset= + "0:max(0, __out_IDim_range_1 - __out_IDim_range_0), 0:max(0, __out_JDim_range_1 - __out_JDim_range_0), -__out_KDim_range_0 + max(0, __out_KDim_range_0):max(0, __out_KDim_range_0, __out_KDim_range_1) - __out_KDim_range_0", + )) + + sdfg.validate() + + return sdfg, state, a, e, b + + +def test_copy_before_and_after_serialization(): + sdfg1, state1, a1, e1, b1 = _make_sdfg() + assert sdfg1.number_of_nodes() == 1 + assert state1.number_of_edges() == 1 + + copy_shape1, src_strides1, dst_strides1, _, _ = dace_cpp.memlet_copy_to_absolute_strides( + None, sdfg1, state1, e1, a1, b1) + print(f"|{copy_shape1}| = {len(copy_shape1)}") + assert len(copy_shape1) == 1 + + json = sdfg1.to_json() + + sdfg2 = dace.SDFG.from_json(json) + state2 = sdfg2.states()[0] + assert sdfg2.number_of_nodes() == 1 + assert state2.number_of_edges() == 1 + + e2 = state2.edges()[0] + a2 = e2.src + b2 = e2.dst + assert a2.data == a1.data + assert b2.data == b1.data + + copy_shape2, src_strides2, dst_strides2, _, _ = dace_cpp.memlet_copy_to_absolute_strides( + None, sdfg2, state2, e2, a2, b2) + + print(f"|{copy_shape2}| = {len(copy_shape2)}") + assert len(copy_shape2) == 1 + + return 0 + + +def test_memlet_copy_shape_roundtrip(): + sdfg1, state1, a1, e1, b1 = _make_sdfg() + + copy_shape1, *_ = dace_cpp.memlet_copy_to_absolute_strides(None, sdfg1, state1, e1, a1, b1) + + sdfg2 = dace.SDFG.from_json(sdfg1.to_json()) + state2 = sdfg2.states()[0] + e2 = state2.edges()[0] + + copy_shape2, *_ = dace_cpp.memlet_copy_to_absolute_strides(None, sdfg2, state2, e2, e2.src, e2.dst) + + assert copy_shape1 == copy_shape2 + + +if __name__ == "__main__": + test_memlet_copy_shape_roundtrip() + test_copy_before_and_after_serialization() diff --git a/tests/copynd_test.py b/tests/copynd_test.py index d805935515..3814a140e2 100644 --- a/tests/copynd_test.py +++ b/tests/copynd_test.py @@ -4,181 +4,96 @@ import pytest -def test(): - print("==== Program start ====") - print('Copy ND tests') - - N = dace.symbol('N') - n = 20 - - sdfg = dace.SDFG('copynd') - state = sdfg.add_state() - - arrays = [] - - # Copy 1: sub-2d array to sub-2d array - arrays.append(state.add_array('A_' + str(len(arrays)), [N, N], dace.float32)) - arrays.append(state.add_array('A_' + str(len(arrays)), [5, 7], dace.float32)) - state.add_edge(arrays[-2], None, arrays[-1], None, dace.memlet.Memlet.simple(arrays[-2], '5:10, N-7:N')) - - # Copy 2: 1d subset of a 4d array to a 1d array - arrays.append(state.add_array('A_' + str(len(arrays)), [N, N, N, N], dace.float32)) - arrays.append(state.add_array('A_' + str(len(arrays)), [N - 1], dace.float32)) - state.add_edge(arrays[-2], None, arrays[-1], None, dace.memlet.Memlet.simple(arrays[-2], '4,1,1:N,2')) - - # Copy 3: 5d array to 5d array - arrays.append(state.add_array('A_' + str(len(arrays)), [2, 3, 4, 5, 6], dace.float32)) - arrays.append(state.add_array('A_' + str(len(arrays)), [2, 3, 4, 5, 6], dace.float32)) - state.add_edge(arrays[-2], None, arrays[-1], None, dace.memlet.Memlet.simple(arrays[-2], '0:2,0:3,0:4,0:5,0:6')) - - # Copy 4: contiguous 1d subset of a 4d array to a 1d array - arrays.append(state.add_array('A_' + str(len(arrays)), [N, N, N, N], dace.float32)) - arrays.append(state.add_array('A_' + str(len(arrays)), [N - 1], dace.float32)) - state.add_edge(arrays[-2], None, arrays[-1], None, dace.memlet.Memlet.simple(arrays[-2], '4,1,2,1:N')) - - # Copy 5: 1d array to a 1d subset of a 4d array - arrays.append(state.add_array('A_' + str(len(arrays)), [N - 2], dace.float32)) - arrays.append(state.add_array('A_' + str(len(arrays)), [N, N, N, N], dace.float32)) - state.add_edge(arrays[-2], None, arrays[-1], None, dace.memlet.Memlet.simple(arrays[-1], '4,1:N-1,1,2')) - - # Copy 6: 4d array to a contiguous 1d subset of a 4d array - arrays.append(state.add_array('A_' + str(len(arrays)), [N - 2], dace.float32)) - arrays.append(state.add_array('A_' + str(len(arrays)), [N, N, N, N], dace.float32)) - state.add_edge(arrays[-2], None, arrays[-1], None, dace.memlet.Memlet.simple(arrays[-1], '4,1,2,1:N-1')) - - # Copy 7: True 4d copy (4d subarray to 4d array) - arrays.append(state.add_array('A_' + str(len(arrays)), [N, N, N, N], dace.float32)) - arrays.append(state.add_array('A_' + str(len(arrays)), [N - 5, N - 4, 3, N - 2], dace.float32)) - state.add_edge(arrays[-2], None, arrays[-1], None, dace.memlet.Memlet.simple(arrays[-2], - '5:N,2:N-2,N-10:N-7,1:N-1')) - - # Copy 8: 4d array with a stride to a 1d array - arrays.append(state.add_array('A_' + str(len(arrays)), [N, N, N, N], dace.float32)) - arrays.append(state.add_array('A_' + str(len(arrays)), [N / 2 - 1], dace.float32)) - state.add_edge(arrays[-2], None, arrays[-1], None, dace.memlet.Memlet.simple(arrays[-2], '4,1,2,1:N-1:2')) - - # Copy 9: 2d array to a 3d array with an offset - arrays.append(state.add_array('A_' + str(len(arrays)), [40, 40], dace.float32)) - arrays.append(state.add_array('A_' + str(len(arrays)), [3, 40, 40], dace.float32)) - state.add_edge(arrays[-2], None, arrays[-1], None, - dace.memlet.Memlet.simple(arrays[-2], '20:40, 10:30', other_subset_str='2, 10:30, 20:40')) - - # Copy 10: Copying from 2d array to another 2d array a 1d slice - # The src stride is 1, while the dst stride is N - arrays.append(state.add_array('A_' + str(len(arrays)), [40, 40], dace.float32)) - arrays.append(state.add_array('A_' + str(len(arrays)), [40, 40], dace.float32)) - state.add_edge(arrays[-2], None, arrays[-1], None, - dace.memlet.Memlet.simple(arrays[-2], '20, 10:30', other_subset_str='10:30, 10')) - - # Copy 11: Copying from 2d array to another 2d array a 1d slice - # The src stride is N, while the dst stride is 1 - arrays.append(state.add_array('A_' + str(len(arrays)), [40, 40], dace.float32)) - arrays.append(state.add_array('A_' + str(len(arrays)), [40, 40], dace.float32)) - state.add_edge(arrays[-2], None, arrays[-1], None, - dace.memlet.Memlet.simple(arrays[-2], '10:30, 20', other_subset_str='10, 10:30')) - - array_data = [ - np.random.rand(*[dace.symbolic.evaluate(s, {N: n}) - for s in a.desc(sdfg).shape]).astype(a.desc(sdfg).dtype.type) for a in arrays - ] - - args = {anode.label: adata for anode, adata in zip(arrays, array_data)} - args['N'] = n - sdfg(**args) - - N = n - - diffs = [ - np.linalg.norm(array_data[1] - array_data[0][5:10, N - 7:N]) / 5.0 * 7.0, - np.linalg.norm(array_data[3] - array_data[2][4, 1, 1:, 2]) / (N - 1), - np.linalg.norm(array_data[5] - array_data[4]) / 2.0 * 3 * 4 * 5 * 6, - np.linalg.norm(array_data[7] - array_data[6][4, 1, 2, 1:]) / (N - 1), - np.linalg.norm(array_data[9][4, 1:N - 1, 1, 2] - array_data[8]) / (N - 2), - np.linalg.norm(array_data[11][4, 1, 2, 1:N - 1] - array_data[10]) / (N - 2), - np.linalg.norm(array_data[13] - array_data[12][5:N, 2:N - 2, N - 10:N - 7, 1:N - 1]) / ((N - 5) * (N - 4) * 3 * - (N - 2)), - np.linalg.norm(array_data[15] - array_data[14][4, 1, 2, 1:(N - 1):2]) / (N / 2 - 1), - np.linalg.norm(array_data[17][2, 10:30, 20:40] - array_data[16][20:40, 10:30]) / 400, - np.linalg.norm(array_data[19][10:30, 10] - array_data[18][20, 10:30]) / 20, - np.linalg.norm(array_data[21][10, 10:30] - array_data[20][10:30, 20]) / 20 - ] - - print('Differences: ', diffs) - - assert all([diff < 1e-7 for diff in diffs]) - - -@pytest.mark.gpu -def test_gpu(): - print("==== Program start ====") - print('Copy ND tests') - +@pytest.mark.parametrize("device", ["cpu", pytest.param("gpu", marks=[pytest.mark.gpu])]) +def test(device: str): N = dace.symbol('N') n = 20 - sdfg = dace.SDFG('copynd_gpu') + sdfg = dace.SDFG(f'copynd_{device}') state = sdfg.add_state() arrays = [] # Copy 1: sub-2d array to sub-2d array - arrays.append(state.add_array('A_' + str(len(arrays)), [N, N], dace.float32)) - arrays.append(state.add_array('A_' + str(len(arrays)), [5, 7], dace.float32)) + sdfg.add_array('A_' + str(len(arrays)), [N, N], dace.float32) + arrays.append(state.add_access('A_' + str(len(arrays)))) + sdfg.add_array('A_' + str(len(arrays)), [5, 7], dace.float32) + arrays.append(state.add_access('A_' + str(len(arrays)))) state.add_edge(arrays[-2], None, arrays[-1], None, dace.memlet.Memlet.simple(arrays[-2], '5:10, N-7:N')) # Copy 2: 1d subset of a 4d array to a 1d array - arrays.append(state.add_array('A_' + str(len(arrays)), [N, N, N, N], dace.float32)) - arrays.append(state.add_array('A_' + str(len(arrays)), [N - 1], dace.float32)) + sdfg.add_array('A_' + str(len(arrays)), [N, N, N, N], dace.float32) + arrays.append(state.add_access('A_' + str(len(arrays)))) + sdfg.add_array('A_' + str(len(arrays)), [N - 1], dace.float32) + arrays.append(state.add_access('A_' + str(len(arrays)))) state.add_edge(arrays[-2], None, arrays[-1], None, dace.memlet.Memlet.simple(arrays[-2], '4,1,1:N,2')) # Copy 3: 5d array to 5d array - arrays.append(state.add_array('A_' + str(len(arrays)), [2, 3, 4, 5, 6], dace.float32)) - arrays.append(state.add_array('A_' + str(len(arrays)), [2, 3, 4, 5, 6], dace.float32)) + sdfg.add_array('A_' + str(len(arrays)), [2, 3, 4, 5, 6], dace.float32) + arrays.append(state.add_access('A_' + str(len(arrays)))) + sdfg.add_array('A_' + str(len(arrays)), [2, 3, 4, 5, 6], dace.float32) + arrays.append(state.add_access('A_' + str(len(arrays)))) state.add_edge(arrays[-2], None, arrays[-1], None, dace.memlet.Memlet.simple(arrays[-2], '0:2,0:3,0:4,0:5,0:6')) # Copy 4: contiguous 1d subset of a 4d array to a 1d array - arrays.append(state.add_array('A_' + str(len(arrays)), [N, N, N, N], dace.float32)) - arrays.append(state.add_array('A_' + str(len(arrays)), [N - 1], dace.float32)) + sdfg.add_array('A_' + str(len(arrays)), [N, N, N, N], dace.float32) + arrays.append(state.add_access('A_' + str(len(arrays)))) + sdfg.add_array('A_' + str(len(arrays)), [N - 1], dace.float32) + arrays.append(state.add_access('A_' + str(len(arrays)))) state.add_edge(arrays[-2], None, arrays[-1], None, dace.memlet.Memlet.simple(arrays[-2], '4,1,2,1:N')) # Copy 5: 1d array to a 1d subset of a 4d array - arrays.append(state.add_array('A_' + str(len(arrays)), [N - 2], dace.float32)) - arrays.append(state.add_array('A_' + str(len(arrays)), [N, N, N, N], dace.float32)) + sdfg.add_array('A_' + str(len(arrays)), [N - 2], dace.float32) + arrays.append(state.add_access('A_' + str(len(arrays)))) + sdfg.add_array('A_' + str(len(arrays)), [N, N, N, N], dace.float32) + arrays.append(state.add_access('A_' + str(len(arrays)))) state.add_edge(arrays[-2], None, arrays[-1], None, dace.memlet.Memlet.simple(arrays[-1], '4,1:N-1,1,2')) # Copy 6: 4d array to a contiguous 1d subset of a 4d array - arrays.append(state.add_array('A_' + str(len(arrays)), [N - 2], dace.float32)) - arrays.append(state.add_array('A_' + str(len(arrays)), [N, N, N, N], dace.float32)) + sdfg.add_array('A_' + str(len(arrays)), [N - 2], dace.float32) + arrays.append(state.add_access('A_' + str(len(arrays)))) + sdfg.add_array('A_' + str(len(arrays)), [N, N, N, N], dace.float32) + arrays.append(state.add_access('A_' + str(len(arrays)))) state.add_edge(arrays[-2], None, arrays[-1], None, dace.memlet.Memlet.simple(arrays[-1], '4,1,2,1:N-1')) # Copy 7: True 4d copy (4d subarray to 4d array) - arrays.append(state.add_array('A_' + str(len(arrays)), [N, N, N, N], dace.float32)) - arrays.append(state.add_array('A_' + str(len(arrays)), [N - 5, N - 4, 3, N - 2], dace.float32)) + sdfg.add_array('A_' + str(len(arrays)), [N, N, N, N], dace.float32) + arrays.append(state.add_access('A_' + str(len(arrays)))) + sdfg.add_array('A_' + str(len(arrays)), [N - 5, N - 4, 3, N - 2], dace.float32) + arrays.append(state.add_access('A_' + str(len(arrays)))) state.add_edge(arrays[-2], None, arrays[-1], None, dace.memlet.Memlet.simple(arrays[-2], '5:N,2:N-2,N-10:N-7,1:N-1')) # Copy 8: 4d array with a stride to a 1d array - arrays.append(state.add_array('A_' + str(len(arrays)), [N, N, N, N], dace.float32)) - arrays.append(state.add_array('A_' + str(len(arrays)), [N / 2 - 1], dace.float32)) + sdfg.add_array('A_' + str(len(arrays)), [N, N, N, N], dace.float32) + arrays.append(state.add_access('A_' + str(len(arrays)))) + sdfg.add_array('A_' + str(len(arrays)), [N / 2 - 1], dace.float32) + arrays.append(state.add_access('A_' + str(len(arrays)))) state.add_edge(arrays[-2], None, arrays[-1], None, dace.memlet.Memlet.simple(arrays[-2], '4,1,2,1:N-1:2')) # Copy 9: 2d array to a 3d array with an offset - arrays.append(state.add_array('A_' + str(len(arrays)), [40, 40], dace.float32)) - arrays.append(state.add_array('A_' + str(len(arrays)), [3, 40, 40], dace.float32)) + sdfg.add_array('A_' + str(len(arrays)), [40, 40], dace.float32) + arrays.append(state.add_access('A_' + str(len(arrays)))) + sdfg.add_array('A_' + str(len(arrays)), [3, 40, 40], dace.float32) + arrays.append(state.add_access('A_' + str(len(arrays)))) state.add_edge(arrays[-2], None, arrays[-1], None, dace.memlet.Memlet.simple(arrays[-2], '20:40, 10:30', other_subset_str='2, 10:30, 20:40')) # Copy 10: Copying from 2d array to another 2d array a 1d slice # The src stride is 1, while the dst stride is N - arrays.append(state.add_array('A_' + str(len(arrays)), [40, 40], dace.float32)) - arrays.append(state.add_array('A_' + str(len(arrays)), [40, 40], dace.float32)) + sdfg.add_array('A_' + str(len(arrays)), [40, 40], dace.float32) + arrays.append(state.add_access('A_' + str(len(arrays)))) + sdfg.add_array('A_' + str(len(arrays)), [40, 40], dace.float32) + arrays.append(state.add_access('A_' + str(len(arrays)))) state.add_edge(arrays[-2], None, arrays[-1], None, dace.memlet.Memlet.simple(arrays[-2], '20, 10:30', other_subset_str='10:30, 10')) # Copy 11: Copying from 2d array to another 2d array a 1d slice # The src stride is N, while the dst stride is 1 - arrays.append(state.add_array('A_' + str(len(arrays)), [40, 40], dace.float32)) - arrays.append(state.add_array('A_' + str(len(arrays)), [40, 40], dace.float32)) + sdfg.add_array('A_' + str(len(arrays)), [40, 40], dace.float32) + arrays.append(state.add_access('A_' + str(len(arrays)))) + sdfg.add_array('A_' + str(len(arrays)), [40, 40], dace.float32) + arrays.append(state.add_access('A_' + str(len(arrays)))) state.add_edge(arrays[-2], None, arrays[-1], None, dace.memlet.Memlet.simple(arrays[-2], '10:30, 20', other_subset_str='10, 10:30')) @@ -186,8 +101,10 @@ def test_gpu(): # padded[:, 1:H+1, 1:W+1, :] = conv2d(input, conv1) # output shape is (N, H + 2,W + 2, C2) # input shape is (N, H , W, C2) - arrays.append(state.add_array('A_' + str(len(arrays)), [5, 6, 7, 8], dace.float32)) - arrays.append(state.add_array('A_' + str(len(arrays)), [5, 8, 9, 8], dace.float32)) + sdfg.add_array('A_' + str(len(arrays)), [5, 6, 7, 8], dace.float32) + arrays.append(state.add_access('A_' + str(len(arrays)))) + sdfg.add_array('A_' + str(len(arrays)), [5, 8, 9, 8], dace.float32) + arrays.append(state.add_access('A_' + str(len(arrays)))) state.add_edge(arrays[-2], None, arrays[-1], None, dace.memlet.Memlet.simple(arrays[-2], '0:5, 0:6, 0:7 , 0:8', other_subset_str='0:5, 1:7, 1:8, 0:8')) @@ -198,7 +115,8 @@ def test_gpu(): args = {anode.label: adata for anode, adata in zip(arrays, array_data)} args['N'] = n - sdfg.apply_gpu_transformations() + if device == "gpu": + sdfg.apply_gpu_transformations() sdfg(**args) N = n @@ -219,11 +137,9 @@ def test_gpu(): np.linalg.norm(array_data[23][:, 1:-1, 1:-1, :] - array_data[22]) / 1680 ] - print('Differences: ', diffs) - assert all([diff < 1e-7 for diff in diffs]) if __name__ == "__main__": - test() - test_gpu() + test("cpu") + test("gpu") diff --git a/tests/cppunparse_test.py b/tests/cppunparse_test.py index 91e8c6b89b..ebe8548ebc 100644 --- a/tests/cppunparse_test.py +++ b/tests/cppunparse_test.py @@ -1,6 +1,5 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. from dace.codegen import cppunparse -import six def _test_py2cpp(func, expected_string): @@ -43,8 +42,7 @@ def test(): return (c * b); }""") - if six.PY3: - success &= _test_py2cpp("""def typed(a: int, b: float) -> float: + success &= _test_py2cpp("""def typed(a: int, b: float) -> float: c = a + b return c*b """, """float typed(int a, float b) { @@ -77,8 +75,7 @@ def lfunc(): }""") # void return value - if six.PY3: - success &= _test_py2cpp(""" + success &= _test_py2cpp(""" def lfunc() -> None: exit(1 >> 3) """, """void lfunc() { @@ -90,8 +87,7 @@ def lfunc() -> None: l = (l + 8);""") # Operations (augmented assignment) - if six.PY3: - success &= _test_py2cpp('l *= 3; l //= 8', """l *= 3; + success &= _test_py2cpp('l *= 3; l //= 8', """l *= 3; l = dace::math::ifloor(l / 8);""") success &= _test_pyexpr2cpp('a << 3', '(a << 3)') diff --git a/tests/custom_build_folder_test.py b/tests/custom_build_folder_test.py index 7b59345d66..d1d22fb3ac 100644 --- a/tests/custom_build_folder_test.py +++ b/tests/custom_build_folder_test.py @@ -18,7 +18,7 @@ def test_custom_build_folder(): csdfg = sdfg.compile() # Ensure files were generated in the right folder - assert os.path.isfile(os.path.join(sdfg.build_folder, 'program.sdfg')) + assert os.path.isfile(os.path.join(sdfg.build_folder, 'program.sdfgz')) # Ensure file is closed so it can be deleted del csdfg diff --git a/tests/dynamic_sdfg_functions_test.py b/tests/dynamic_sdfg_functions_test.py index 9a94f926d8..0ab5bef885 100644 --- a/tests/dynamic_sdfg_functions_test.py +++ b/tests/dynamic_sdfg_functions_test.py @@ -1,5 +1,4 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. -import math import numpy as np import dace as dp @@ -29,9 +28,11 @@ def test_dynamic_sdfg_with_math_functions(): # Construct SDFG mysdfg = SDFG('mymodexp') + mysdfg.add_array('A', [N], dp.float32) + mysdfg.add_array('B', [N], dp.float32) state = mysdfg.add_state() - A = state.add_array('A', [N], dp.float32) - B = state.add_array('B', [N], dp.float32) + A = state.add_access('A') + B = state.add_access('B') # Easy way to add a tasklet tasklet, map_entry, map_exit = state.add_mapped_tasklet('mytasklet', dict(i='0:N'), @@ -46,7 +47,6 @@ def test_dynamic_sdfg_with_math_functions(): #mymodexp_prog(input, output) diff = np.linalg.norm(np.exp(input) - output) / n - print("Difference:", diff) assert diff <= 1e-5 diff --git a/tests/fortran/allocate_test.py b/tests/fortran/allocate_test.py index e8ffd0a07e..8df87d2146 100644 --- a/tests/fortran/allocate_test.py +++ b/tests/fortran/allocate_test.py @@ -1,21 +1,8 @@ # Copyright 2019-2023 ETH Zurich and the DaCe authors. All rights reserved. -from fparser.common.readfortran import FortranStringReader -from fparser.common.readfortran import FortranFileReader -from fparser.two.parser import ParserFactory -import sys, os import numpy as np -import pytest -from dace import SDFG, SDFGState, nodes, dtypes, data, subsets, symbolic from dace.frontend.fortran import fortran_parser -from fparser.two.symbol_table import SymbolTable -from dace.sdfg import utils as sdutil - -import dace.frontend.fortran.ast_components as ast_components -import dace.frontend.fortran.ast_transforms as ast_transforms -import dace.frontend.fortran.ast_utils as ast_utils -import dace.frontend.fortran.ast_internal_classes as ast_internal_classes def test_fortran_frontend_basic_allocate(): diff --git a/tests/fortran/array_to_loop_offset.py b/tests/fortran/array_to_loop_offset.py index 7459e994b0..e5e7d4484a 100644 --- a/tests/fortran/array_to_loop_offset.py +++ b/tests/fortran/array_to_loop_offset.py @@ -2,7 +2,7 @@ import numpy as np -from dace.frontend.fortran import ast_transforms, fortran_parser +from dace.frontend.fortran import fortran_parser def test_fortran_frontend_arr2loop_without_offset(): diff --git a/tests/fortran/call_extract_test.py b/tests/fortran/call_extract_test.py index eaf83808ac..d3df2e8ffc 100644 --- a/tests/fortran/call_extract_test.py +++ b/tests/fortran/call_extract_test.py @@ -1,7 +1,6 @@ # Copyright 2019-2023 ETH Zurich and the DaCe authors. All rights reserved. import numpy as np -import pytest from dace.frontend.fortran import fortran_parser diff --git a/tests/fortran/dace_support_test.py b/tests/fortran/dace_support_test.py index 4e2036c1d6..9ad692ad0a 100644 --- a/tests/fortran/dace_support_test.py +++ b/tests/fortran/dace_support_test.py @@ -1,21 +1,8 @@ # Copyright 2023 ETH Zurich and the DaCe authors. All rights reserved. -from fparser.common.readfortran import FortranStringReader -from fparser.common.readfortran import FortranFileReader -from fparser.two.parser import ParserFactory -import sys, os import numpy as np -import pytest -from dace import SDFG, SDFGState, nodes, dtypes, data, subsets, symbolic from dace.frontend.fortran import fortran_parser -from fparser.two.symbol_table import SymbolTable -from dace.sdfg import utils as sdutil - -import dace.frontend.fortran.ast_components as ast_components -import dace.frontend.fortran.ast_transforms as ast_transforms -import dace.frontend.fortran.ast_utils as ast_utils -import dace.frontend.fortran.ast_internal_classes as ast_internal_classes def test_fortran_frontend_simplify(): diff --git a/tests/fortran/intrinsic_basic_test.py b/tests/fortran/intrinsic_basic_test.py index 039ad7c3e7..b65f288349 100644 --- a/tests/fortran/intrinsic_basic_test.py +++ b/tests/fortran/intrinsic_basic_test.py @@ -1,7 +1,6 @@ # Copyright 2019-2023 ETH Zurich and the DaCe authors. All rights reserved. import numpy as np -import pytest from dace.frontend.fortran import fortran_parser diff --git a/tests/fortran/intrinsic_math_test.py b/tests/fortran/intrinsic_math_test.py index 6fa7ab1e2b..71cbd753e2 100644 --- a/tests/fortran/intrinsic_math_test.py +++ b/tests/fortran/intrinsic_math_test.py @@ -1,7 +1,6 @@ # Copyright 2019-2023 ETH Zurich and the DaCe authors. All rights reserved. import numpy as np -import pytest from dace.frontend.fortran import fortran_parser diff --git a/tests/fortran/intrinsic_merge_test.py b/tests/fortran/intrinsic_merge_test.py index 174b0f4440..b81b826aef 100644 --- a/tests/fortran/intrinsic_merge_test.py +++ b/tests/fortran/intrinsic_merge_test.py @@ -2,7 +2,7 @@ import numpy as np -from dace.frontend.fortran import ast_transforms, fortran_parser +from dace.frontend.fortran import fortran_parser def test_fortran_frontend_merge_1d(): diff --git a/tests/fortran/intrinsic_minmaxval_test.py b/tests/fortran/intrinsic_minmaxval_test.py index b064166c43..8282691c21 100644 --- a/tests/fortran/intrinsic_minmaxval_test.py +++ b/tests/fortran/intrinsic_minmaxval_test.py @@ -2,7 +2,7 @@ import numpy as np -from dace.frontend.fortran import ast_transforms, fortran_parser +from dace.frontend.fortran import fortran_parser def test_fortran_frontend_minval_double(): diff --git a/tests/fortran/intrinsic_product_test.py b/tests/fortran/intrinsic_product_test.py index 095eb4485f..2fa7868006 100644 --- a/tests/fortran/intrinsic_product_test.py +++ b/tests/fortran/intrinsic_product_test.py @@ -3,7 +3,7 @@ import numpy as np import pytest -from dace.frontend.fortran import ast_transforms, fortran_parser +from dace.frontend.fortran import fortran_parser def test_fortran_frontend_product_array(): diff --git a/tests/fortran/intrinsic_sum_test.py b/tests/fortran/intrinsic_sum_test.py index a497271ff6..8a5d4bed70 100644 --- a/tests/fortran/intrinsic_sum_test.py +++ b/tests/fortran/intrinsic_sum_test.py @@ -2,7 +2,7 @@ import numpy as np -from dace.frontend.fortran import ast_transforms, fortran_parser +from dace.frontend.fortran import fortran_parser def test_fortran_frontend_sum2loop_1d_without_offset(): diff --git a/tests/fortran/offset_normalizer_test.py b/tests/fortran/offset_normalizer_test.py index b1635aac48..e2cb3e2c6d 100644 --- a/tests/fortran/offset_normalizer_test.py +++ b/tests/fortran/offset_normalizer_test.py @@ -2,7 +2,7 @@ import numpy as np -from dace.frontend.fortran import ast_transforms, fortran_parser +from dace.frontend.fortran import fortran_parser def test_fortran_frontend_offset_normalizer_1d(): diff --git a/tests/fortran/sum_to_loop_offset_test.py b/tests/fortran/sum_to_loop_offset_test.py index a497271ff6..8a5d4bed70 100644 --- a/tests/fortran/sum_to_loop_offset_test.py +++ b/tests/fortran/sum_to_loop_offset_test.py @@ -2,7 +2,7 @@ import numpy as np -from dace.frontend.fortran import ast_transforms, fortran_parser +from dace.frontend.fortran import fortran_parser def test_fortran_frontend_sum2loop_1d_without_offset(): diff --git a/tests/fortran/view_test.py b/tests/fortran/view_test.py index 2d4d14f8a0..5e96afce66 100644 --- a/tests/fortran/view_test.py +++ b/tests/fortran/view_test.py @@ -1,21 +1,8 @@ # Copyright 2023 ETH Zurich and the DaCe authors. All rights reserved. -from fparser.common.readfortran import FortranStringReader -from fparser.common.readfortran import FortranFileReader -from fparser.two.parser import ParserFactory -import sys, os import numpy as np -import pytest -from dace import SDFG, SDFGState, nodes, dtypes, data, subsets, symbolic from dace.frontend.fortran import fortran_parser -from fparser.two.symbol_table import SymbolTable -from dace.sdfg import utils as sdutil - -import dace.frontend.fortran.ast_components as ast_components -import dace.frontend.fortran.ast_transforms as ast_transforms -import dace.frontend.fortran.ast_utils as ast_utils -import dace.frontend.fortran.ast_internal_classes as ast_internal_classes def test_fortran_frontend_view_test(): diff --git a/tests/graph_test.py b/tests/graph_test.py index 9cee2d0be1..4a3ec780c1 100644 --- a/tests/graph_test.py +++ b/tests/graph_test.py @@ -110,7 +110,8 @@ def test_ordered_multidigraph(self): def test_dfs_edges(self): sdfg = dace.SDFG('test_dfs_edges') - before, _, _ = sdfg.add_loop(sdfg.add_state(), sdfg.add_state(), sdfg.add_state(), 'i', '0', 'i < 10', 'i + 1') + before, _, _ = sdfg.add_loop_state_machine(sdfg.add_state(), sdfg.add_state(), sdfg.add_state(), 'i', '0', + 'i < 10', 'i + 1') visited_edges = list(sdfg.dfs_edges(before)) assert len(visited_edges) == len(set(visited_edges)) diff --git a/tests/halfvec_cudatest.py b/tests/halfvec_cudatest.py index 824bddb092..1bf38e3573 100644 --- a/tests/halfvec_cudatest.py +++ b/tests/halfvec_cudatest.py @@ -5,7 +5,7 @@ import math import numpy as np import pytest -from dace.transformation.dataflow import MapFusionVertical, Vectorization +from dace.transformation.dataflow import Vectorization from dace.transformation.optimizer import Optimizer N = dace.symbol('N') diff --git a/tests/inline_nonsink_access_test.py b/tests/inline_nonsink_access_test.py index f3c29ace20..f97a3d32c4 100644 --- a/tests/inline_nonsink_access_test.py +++ b/tests/inline_nonsink_access_test.py @@ -1,6 +1,7 @@ -# Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. import dace import numpy as np +import pytest def make_sdfg(outer_shape, inner_shape, outer_index, inner_index): @@ -41,11 +42,10 @@ def make_sdfg(outer_shape, inner_shape, outer_index, inner_index): nsdfg_node = state.add_nested_sdfg(nsdfg, {}, {'C', 'D'}) state.add_edge(nsdfg_node, 'C', A, None, dace.Memlet('A')) state.add_edge(nsdfg_node, 'D', B, None, dace.Memlet('B')) - - sdfg.save('_dacegraphs/program.sdfg') return sdfg +@pytest.mark.skip(reason="This test should not pass due to the conservative nature of InlineSDFG") def test_same_shape(): A = np.random.rand(1).astype(np.float32) B = np.random.rand(1).astype(np.float32) @@ -53,7 +53,6 @@ def test_same_shape(): sdfg = make_sdfg([1], [1], '0', '0') sdfg.simplify() - sdfg.save('_dacegraphs/program.sdfg') sdfg(A=A, B=B) assert len(sdfg.node(0).nodes()) == 8 @@ -61,10 +60,10 @@ def test_same_shape(): expected = np.array([2**2, (2**2) + (2**6)], dtype=np.float32) result = np.array([A[0], B[0]], dtype=np.float32) diff = np.linalg.norm(expected - result) - print('Difference:', diff) assert diff <= 1e-6 +@pytest.mark.skip(reason="This test should not pass due to the conservative nature of InlineSDFG") def test_different_shape(): A = np.random.rand(20, 3).astype(np.float32) B = np.random.rand(20, 3).astype(np.float32) @@ -79,10 +78,11 @@ def test_different_shape(): expected = np.array([2**2, (2**2) + (2**6)], dtype=np.float32) result = np.array([A[1, 0], B[1, 0]], dtype=np.float32) diff = np.linalg.norm(expected - result) - print('Difference:', diff) assert diff <= 1e-6 if __name__ == "__main__": - test_same_shape() - test_different_shape() + # Tests should not pass due to the conservative nature of InlineSDFG + # test_same_shape() + # test_different_shape() + pass diff --git a/tests/inlining_test.py b/tests/inlining_test.py index 80c4157dba..36ede62955 100644 --- a/tests/inlining_test.py +++ b/tests/inlining_test.py @@ -1355,6 +1355,97 @@ def test_multistate_inline_nested_control_flow_blocks(): assert np.allclose(C, expected) +def test_inline_write_write_conflict(): + """ + Tests a regression where inlining an SDFG creates a write-write conflict on the same array. + """ + sdfg = dace.SDFG("tester") + sdfg.add_array("A", [10], dace.int32) + + state = sdfg.add_state("test_state", is_start_block=True) + access = state.add_access("A") + state.add_mapped_tasklet("set_A_to_5", + map_ranges={"__i": "0:10"}, + inputs={}, + code="__out = 5", + outputs={"__out": dace.Memlet("A[__i]")}, + output_nodes={"A": access}, + external_edges=True) + + n_sdfg = dace.SDFG("nested") + n_sdfg.add_array("A", [10], dace.int32) + n_state = n_sdfg.add_state("nested_state", is_start_block=True) + n_access = n_state.add_access("A") + n_state.add_mapped_tasklet("set_A_to_7", + map_ranges={"__i": "2:5"}, + inputs={}, + code="__out = 7", + outputs={"__out": dace.Memlet("A[__i]")}, + output_nodes={"A": n_access}, + external_edges=True) + + n_tasklet_2 = n_state.add_tasklet("plus_one", {"a"}, {"a1"}, "a1 = a + 1") + n_access_2 = n_state.add_access("A") + n_state.add_memlet_path(n_access, n_tasklet_2, dst_conn="a", memlet=dace.Memlet("A[3]")) + n_state.add_memlet_path(n_tasklet_2, n_access_2, src_conn="a1", memlet=dace.Memlet("A[3]")) + + nested = state.add_nested_sdfg(n_sdfg, inputs={"A"}, outputs={"A"}, name="nested_sdfg") + state.add_memlet_path(access, nested, dst_conn="A", memlet=dace.Memlet("A[3]")) + access_2 = state.add_access("A") + state.add_memlet_path(nested, access_2, src_conn="A", memlet=dace.Memlet("A[2:5]")) + + sdfg.validate() + + with pytest.raises(ValueError, match='Transformation cannot be applied'): + InlineSDFG().apply_to(sdfg, nested_sdfg=nested) + + +def test_inline_nested_accessnode(): + """ + Tests a use case where a non-transient access node is defined within a scope + """ + sdfg = dace.SDFG("tester") + sdfg.add_array("A", [10], dace.int32) + sdfg.add_array("B", [10], dace.int32) + + state = sdfg.add_state("test_state", is_start_block=True) + access = state.add_access("A") + state.add_mapped_tasklet("set_A_to_5", + map_ranges={"__i": "0:10"}, + inputs={}, + code="__out = 5", + outputs={"__out": dace.Memlet("A[__i]")}, + output_nodes={"A": access}, + external_edges=True) + + n_sdfg = dace.SDFG("nested") + n_sdfg.add_array("A", [10], dace.int32) + n_sdfg.add_array("B", [10], dace.int32) + n_state = n_sdfg.add_state("nested_state", is_start_block=True) + me, mx = n_state.add_map("map", {"__i": "2:5"}) + n_access = n_state.add_access("A") + n_access_b = n_state.add_access("B") + n_state.add_edge(me, None, n_access, None, dace.Memlet()) + tasklet = n_state.add_tasklet("set_A_to_7", {"__in"}, {"__out"}, "__out = 7") + n_state.add_edge(n_access, None, tasklet, "__in", dace.Memlet("A[__i]")) + n_state.add_edge(tasklet, "__out", n_access_b, None, dace.Memlet("B[__i]")) + n_tasklet_2 = n_state.add_tasklet("plus_one", {"a"}, {"a1"}, "a1 = a + 1") + n_access_2 = n_state.add_access("B") + n_state.add_memlet_path(n_access_b, n_tasklet_2, dst_conn="a", memlet=dace.Memlet("B[3]")) + n_state.add_memlet_path(n_tasklet_2, n_access_2, src_conn="a1", memlet=dace.Memlet("B[3]")) + n_state.add_edge(n_access_2, None, mx, None, dace.Memlet()) + + nested = state.add_nested_sdfg(n_sdfg, inputs={"A"}, outputs={"B"}, name="nested_sdfg") + state.add_memlet_path(access, nested, dst_conn="A", memlet=dace.Memlet("A[3]")) + access_2 = state.add_access("B") + state.add_memlet_path(nested, access_2, src_conn="B", memlet=dace.Memlet("B[2:5]")) + + sdfg.validate() + + with pytest.raises(ValueError, match='Transformation cannot be applied'): + InlineSDFG().apply_to(sdfg, nested_sdfg=nested) + + if __name__ == "__main__": test() # Skipped due to bug that cannot be reproduced outside CI @@ -1394,3 +1485,6 @@ def test_multistate_inline_nested_control_flow_blocks(): for outside_and_inner_symbol_have_same_meaning in [True, False]: test_singlestate_inline_with_symbol_mapping( outside_and_inner_symbol_have_same_meaning=outside_and_inner_symbol_have_same_meaning) + + test_inline_write_write_conflict() + test_inline_nested_accessnode() diff --git a/tests/kernel_fusion_cudatest.py b/tests/kernel_fusion_cudatest.py index 556957287f..8d6d6ce681 100644 --- a/tests/kernel_fusion_cudatest.py +++ b/tests/kernel_fusion_cudatest.py @@ -2,7 +2,6 @@ import pytest import dace import numpy as np -import sympy as sp def _construct_graph(tbsize_1=None, tbsize_2=None) -> dace.SDFG: diff --git a/tests/lib_reuse_test.py b/tests/lib_reuse_test.py index fe96df3d6a..f0076769bd 100644 --- a/tests/lib_reuse_test.py +++ b/tests/lib_reuse_test.py @@ -1,8 +1,11 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. +import pytest +import pathlib + import dace from dace.frontend.python.parser import DaceProgram from dace.codegen.exceptions import CompilationError -from dace.sdfg.utils import load_precompiled_sdfg +from dace.codegen.compiler import load_precompiled_sdfg import numpy as np @@ -22,8 +25,6 @@ def tasklet(i): def test_reload(): - print('Reloadable DaCe program test') - array_one = np.random.rand(10).astype(np.float64) array_two = np.random.rand(20).astype(np.float64) output_one = np.zeros(10, dtype=np.float64) @@ -32,16 +33,18 @@ def test_reload(): prog_one = program_generator(10, 2.0) prog_two = program_generator(20, 4.0) - # This should create two libraries for the two SDFGs, as they compile over - # the same folder + # This should create two libraries for the two SDFGs, as they compile over the same folder func1 = prog_one.compile() try: func2 = prog_two.compile() except CompilationError: - # On some systems (e.g., Windows), the file will be locked, so - # compilation will fail - print('Compilation failed due to locked file. Skipping test.') - return + # On some systems (e.g., Windows), the file will be locked, so compilation will fail + pytest.skip('Compilation failed due to locked file. Skipping test.') + + lib1_path = pathlib.Path(func1.filename) + lib2_path = pathlib.Path(func2.filename) + assert lib1_path != lib2_path + assert lib1_path.parent == lib2_path.parent func1(input=array_one, output=output_one) func2(input=array_two, output=output_two) @@ -53,10 +56,26 @@ def test_reload(): def test_load_precompiled(): + for folder_mode in ["development", "production"]: + with dace.config.temporary_config() as conf: + conf.set('compiler', 'build_folder_mode', value=folder_mode) + _load_precompiled_impl( + test_name="test_load_precompiled", + folder_mode=folder_mode, + ) + + +def _load_precompiled_impl(test_name: str, folder_mode: str) -> None: prog = program_generator(10, 2.0) sdfg = prog.to_sdfg() + sdfg.name = f'{test_name}_{sdfg.name}_{folder_mode}' + func1 = sdfg.compile() - func2 = load_precompiled_sdfg(sdfg.build_folder) + + if folder_mode == "production": + func2 = load_precompiled_sdfg(sdfg.build_folder, sdfg=sdfg) + elif folder_mode == "development": + func2 = load_precompiled_sdfg(sdfg.build_folder) inp = np.random.rand(10).astype(np.float64) output_one = np.zeros(10, dtype=np.float64) @@ -65,6 +84,11 @@ def test_load_precompiled(): func1(input=inp, output=output_one) func2(input=inp, output=output_two) + lib1_path = pathlib.Path(func1.filename) + lib2_path = pathlib.Path(func2.filename) + assert lib1_path != lib2_path + assert lib1_path.parent == lib2_path.parent + assert (np.allclose(output_one, output_two)) del func1 diff --git a/tests/library/blas_dot_test.py b/tests/library/blas_dot_test.py index 68e90b98c6..8a1c374466 100644 --- a/tests/library/blas_dot_test.py +++ b/tests/library/blas_dot_test.py @@ -3,11 +3,8 @@ import dace from dace.memlet import Memlet -from dace.codegen.exceptions import CompilerConfigurationError, CompilationError import dace.libraries.blas as blas import numpy as np -import sys -import warnings ############################################################################### diff --git a/tests/library/cutensor_ttranspose_test.py b/tests/library/cutensor_ttranspose_test.py new file mode 100644 index 0000000000..7c43680b0e --- /dev/null +++ b/tests/library/cutensor_ttranspose_test.py @@ -0,0 +1,181 @@ +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. +""" +Unit tests for the TensorTranspose library node with the cuTENSOR v2 expansion. +Tests float64 and int32 tensors using parametrization. +""" +import numpy as np +import pytest +import dace +from dace import dtypes +from dace.libraries.linalg import TensorTranspose + + +def _build_transpose_sdfg( + name: str, + inp_shape: tuple, + axes: list[int], + dtype: dace.typeclass = dace.float64, + implementation: str = "cuTENSOR", +) -> dace.SDFG: + out_shape = tuple(inp_shape[a] for a in axes) + + sdfg = dace.SDFG(name) + + # Host arrays (default storage = CPU) + sdfg.add_array("A_host", inp_shape, dtype) + sdfg.add_array("B_host", out_shape, dtype) + + # GPU arrays (transient) + sdfg.add_array("A_gpu", inp_shape, dtype, storage=dtypes.StorageType.GPU_Global, transient=True) + sdfg.add_array("B_gpu", out_shape, dtype, storage=dtypes.StorageType.GPU_Global, transient=True) + + state = sdfg.add_state("transpose_state", is_start_block=True) + + a_host_r = state.add_read("A_host") + a_gpu = state.add_access("A_gpu") + b_gpu = state.add_access("B_gpu") + b_host_w = state.add_write("B_host") + + tnode = TensorTranspose("_Transpose_", axes=axes) + tnode.implementation = implementation + state.add_node(tnode) + + state.add_edge(a_host_r, None, a_gpu, None, dace.Memlet.from_array("A_host", sdfg.arrays["A_host"])) + state.add_edge(a_gpu, None, tnode, "_inp_tensor", dace.Memlet.from_array("A_gpu", sdfg.arrays["A_gpu"])) + state.add_edge(tnode, "_out_tensor", b_gpu, None, dace.Memlet.from_array("B_gpu", sdfg.arrays["B_gpu"])) + state.add_edge(b_gpu, None, b_host_w, None, dace.Memlet.from_array("B_host", sdfg.arrays["B_host"])) + + return sdfg + + +def _run_transpose_test( + inp_shape: tuple, + axes: list[int], + dtype_np=np.float64, + dtype_dace=dace.float64, + implementation: str = "cuTENSOR", +): + """Helper to run a transpose test. For integers, uses random ints and tolerance 0.""" + rng = np.random.default_rng(42) + + if np.issubdtype(dtype_np, np.integer): + A = rng.integers(low=-100, high=100, size=inp_shape, dtype=dtype_np) + rtol, atol = 0, 0 + else: + A = rng.random(inp_shape).astype(dtype_np) + rtol, atol = 1e-12, 1e-14 + + expected = np.ascontiguousarray(np.transpose(A, axes)) + + name = f"test_transpose_{'_'.join(map(str, axes))}_{dtype_np.__name__}" + sdfg = _build_transpose_sdfg(name, inp_shape, axes, dtype_dace, implementation) + + B = np.zeros_like(expected) + compiled = sdfg.compile() + compiled(A_host=A, B_host=B) + + assert B.shape == expected.shape, f"Shape mismatch: got {B.shape}, expected {expected.shape}" + np.testing.assert_allclose(B, expected, rtol=rtol, atol=atol, err_msg=f"Transpose {axes} failed for {dtype_np}") + + +dtype_params = [ + (np.float64, dace.float64, "f64"), + (np.int32, dace.int32, "i32"), +] + + +@pytest.mark.gpu +@pytest.mark.parametrize("dtype_np,dtype_dace,type_name", dtype_params, ids=[p[2] for p in dtype_params]) +def test_transpose_3d_jik(dtype_np, dtype_dace, type_name): + """(i,j,k) -> (j,i,k)""" + _run_transpose_test( + inp_shape=(3, 5, 7), + axes=[1, 0, 2], + dtype_np=dtype_np, + dtype_dace=dtype_dace, + ) + + +@pytest.mark.gpu +@pytest.mark.parametrize("dtype_np,dtype_dace,type_name", dtype_params, ids=[p[2] for p in dtype_params]) +def test_transpose_3d_kji(dtype_np, dtype_dace, type_name): + """(i,j,k) -> (k,j,i)""" + _run_transpose_test( + inp_shape=(4, 6, 8), + axes=[2, 1, 0], + dtype_np=dtype_np, + dtype_dace=dtype_dace, + ) + + +@pytest.mark.gpu +@pytest.mark.parametrize("dtype_np,dtype_dace,type_name", dtype_params, ids=[p[2] for p in dtype_params]) +def test_transpose_4d_reverse(dtype_np, dtype_dace, type_name): + """(i,j,k,l) -> (l,k,j,i)""" + _run_transpose_test( + inp_shape=(2, 3, 5, 7), + axes=[3, 2, 1, 0], + dtype_np=dtype_np, + dtype_dace=dtype_dace, + ) + + +@pytest.mark.gpu +@pytest.mark.parametrize("dtype_np,dtype_dace,type_name", dtype_params, ids=[p[2] for p in dtype_params]) +def test_transpose_4d_cyclic(dtype_np, dtype_dace, type_name): + """(i,j,k,l) -> (j,k,l,i)""" + _run_transpose_test( + inp_shape=(2, 4, 6, 8), + axes=[1, 2, 3, 0], + dtype_np=dtype_np, + dtype_dace=dtype_dace, + ) + + +def _build_pure_sdfg(name, inp_shape, axes, dtype): + out_shape = tuple(inp_shape[a] for a in axes) + sdfg = dace.SDFG(name) + sdfg.add_array("A", inp_shape, dtype) + sdfg.add_array("B", out_shape, dtype) + + state = sdfg.add_state("s", is_start_block=True) + a_node = state.add_read("A") + b_node = state.add_write("B") + tnode = TensorTranspose("_T_", axes=axes) + tnode.implementation = "pure" + state.add_node(tnode) + state.add_edge(a_node, None, tnode, "_inp_tensor", dace.Memlet.from_array("A", sdfg.arrays["A"])) + state.add_edge(tnode, "_out_tensor", b_node, None, dace.Memlet.from_array("B", sdfg.arrays["B"])) + return sdfg + + +@pytest.mark.parametrize("dtype_np,dtype_dace,type_name", dtype_params, ids=[p[2] for p in dtype_params]) +def test_transpose_pure_3d_jik(dtype_np, dtype_dace, type_name): + """Pure expansion on CPU with different dtypes.""" + rng = np.random.default_rng(123) + inp_shape = (3, 5, 7) + axes = [1, 0, 2] + + if np.issubdtype(dtype_np, np.integer): + A = rng.integers(-100, 100, size=inp_shape, dtype=dtype_np) + rtol, atol = 0, 0 + else: + A = rng.random(inp_shape).astype(dtype_np) + rtol, atol = 1e-12, 1e-14 + + expected = np.ascontiguousarray(np.transpose(A, axes)) + + sdfg = _build_pure_sdfg(f"test_transpose_pure_jik_{type_name}", inp_shape, axes, dtype_dace) + B = np.zeros_like(expected) + compiled = sdfg.compile() + compiled(A=A, B=B) + + np.testing.assert_allclose(B, expected, rtol=rtol, atol=atol, err_msg=f"Pure transpose failed for {dtype_np}") + + +if __name__ == "__main__": + for (nptype, dacetype, strtype) in dtype_params: + test_transpose_pure_3d_jik(nptype, dacetype, strtype) + test_transpose_3d_jik(nptype, dacetype, strtype) + test_transpose_4d_reverse(nptype, dacetype, strtype) + test_transpose_4d_cyclic(nptype, dacetype, strtype) diff --git a/tests/library/gemm_test.py b/tests/library/gemm_test.py index 68f9fb355d..d7d43b1905 100644 --- a/tests/library/gemm_test.py +++ b/tests/library/gemm_test.py @@ -1,12 +1,10 @@ # Copyright 2019-2025 ETH Zurich and the DaCe authors. All rights reserved. import pytest -import warnings import itertools import sys import dace import random import numpy as np -from dace.codegen.exceptions import CompilerConfigurationError, CompilationError from dace.libraries.blas import Gemm M = dace.symbol('M') diff --git a/tests/library/include_test.py b/tests/library/include_test.py index 6b84632bae..11c6e13f09 100644 --- a/tests/library/include_test.py +++ b/tests/library/include_test.py @@ -21,7 +21,7 @@ def test_include(): assert_not_exists("FooLib") assert_not_exists("BarLib") - import foolib # foolib depends on barlib + import foolib # noqa: F401 (FooLib depends on BarLib) assert_exists("FooLib") assert_exists("BarLib") diff --git a/tests/library/lapack_getrf_test.py b/tests/library/lapack_getrf_test.py index ffe66f9c1d..0fea01d432 100644 --- a/tests/library/lapack_getrf_test.py +++ b/tests/library/lapack_getrf_test.py @@ -1,7 +1,7 @@ -# Copyright 2019-2023 ETH Zurich and the DaCe authors. All rights reserved. +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. import dace import dace.libraries.lapack as lapack -import dace.libraries.standard as std +import dace.libraries.linalg as linalg import numpy as np import pytest @@ -35,9 +35,9 @@ def make_sdfg(implementation, dtype, storage=dace.StorageType.Default): xo = state.add_access("x" + suffix) xin = state.add_access("xt" + suffix) xout = state.add_access("xt" + suffix) - transpose_in = std.Transpose("transpose_in", dtype=dtype) + transpose_in = linalg.Transpose("transpose_in", dtype=dtype) transpose_in.implementation = "cuBLAS" - transpose_out = std.Transpose("transpose_out", dtype=dtype) + transpose_out = linalg.Transpose("transpose_out", dtype=dtype) transpose_out.implementation = "cuBLAS" state.add_nedge(xhi, xi, Memlet.from_array(*xhost_arr)) state.add_nedge(xo, xho, Memlet.from_array(*xhost_arr)) diff --git a/tests/library/lapack_getrs_test.py b/tests/library/lapack_getrs_test.py index 97a2a35a6b..322ad005cd 100644 --- a/tests/library/lapack_getrs_test.py +++ b/tests/library/lapack_getrs_test.py @@ -1,7 +1,7 @@ -# Copyright 2019-2023 ETH Zurich and the DaCe authors. All rights reserved. +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. import dace import dace.libraries.lapack as lapack -import dace.libraries.standard as std +import dace.libraries.linalg as linalg import numpy as np import pytest @@ -40,7 +40,7 @@ def make_sdfg(implementation, dtype, storage=dace.StorageType.Default): Bho = state.add_read("B") Bin = state.add_access("B" + suffix) Bout = state.add_access("B" + suffix) - transpose_in = std.Transpose("transpose_in", dtype=dtype) + transpose_in = linalg.Transpose("transpose_in", dtype=dtype) transpose_in.implementation = "cuBLAS" state.add_nedge(Ahi, Ai, Memlet.from_array(*Ahost_arr)) state.add_nedge(Bhi, Bin, Memlet.from_array(*Bhost_arr)) diff --git a/tests/library/lapack_potrf_test.py b/tests/library/lapack_potrf_test.py index 7f7f6714d1..8ae45247b7 100644 --- a/tests/library/lapack_potrf_test.py +++ b/tests/library/lapack_potrf_test.py @@ -1,7 +1,7 @@ -# Copyright 2019-2023 ETH Zurich and the DaCe authors. All rights reserved. +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. import dace import dace.libraries.lapack as lapack -import dace.libraries.standard as std +import dace.libraries.linalg as linalg import numpy as np import pytest @@ -41,9 +41,9 @@ def make_sdfg(implementation, dtype, storage=dace.StorageType.Default): xo = state.add_access("x" + suffix) xin = state.add_access("xt" + suffix) xout = state.add_access("xt" + suffix) - transpose_in = std.Transpose("transpose_in", dtype=dtype) + transpose_in = linalg.Transpose("transpose_in", dtype=dtype) transpose_in.implementation = "cuBLAS" - transpose_out = std.Transpose("transpose_out", dtype=dtype) + transpose_out = linalg.Transpose("transpose_out", dtype=dtype) transpose_out.implementation = "cuBLAS" state.add_nedge(xhi, xi, Memlet.from_array(*xhost_arr)) state.add_nedge(xo, xho, Memlet.from_array(*xhost_arr)) diff --git a/tests/library/library_node_expand_test.py b/tests/library/library_node_expand_test.py index 5eac4053bd..1496484c0d 100644 --- a/tests/library/library_node_expand_test.py +++ b/tests/library/library_node_expand_test.py @@ -6,7 +6,7 @@ import os import numpy as np import copy -from dace import SDFGState, SDFG, library, dtypes +from dace import SDFGState, SDFG, library from dace.transformation.transformation import ExpandTransformation # Add the test library to the path diff --git a/tests/library/matmul_cudatest.py b/tests/library/matmul_cudatest.py index e7773087f2..8802ef0843 100644 --- a/tests/library/matmul_cudatest.py +++ b/tests/library/matmul_cudatest.py @@ -3,11 +3,9 @@ from dace.config import set_temporary from dace.library import change_default from dace.memlet import Memlet -from dace.codegen.exceptions import CompilerConfigurationError, CompilationError import dace.libraries.blas as blas import itertools import numpy as np -import sys import pytest ############################################################################### diff --git a/tests/library/mpi/distributed_descriptor_frontend_test.py b/tests/library/mpi/distributed_descriptor_frontend_test.py new file mode 100644 index 0000000000..76cb8b02b7 --- /dev/null +++ b/tests/library/mpi/distributed_descriptor_frontend_test.py @@ -0,0 +1,87 @@ +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. +import dace +import numpy as np +import pytest + +from dace.data.distributed import ProcessGrid +from dace.libraries.mpi.nodes.dummy import Dummy +from dace.libraries.mpi.nodes.bcast import Bcast + + +def _incoming_descriptor_name(sdfg, node, connector): + states = [state for state in sdfg.states() if node in state.nodes()] + assert len(states) == 1 + + edges = list(states[0].in_edges_by_connector(node, connector)) + assert len(edges) == 1 + return edges[0].data.data + + +@pytest.mark.mpi +def test_create_cart_bcast_uses_process_grid_descriptor(): + MPI = pytest.importorskip('mpi4py.MPI') + + @dace.program + def pgrid_bcast(A: dace.int32[10]): + pgrid = MPI.COMM_WORLD.Create_cart([1, 1]) + if pgrid != MPI.COMM_NULL: + pgrid.Bcast(A) + + sdfg = pgrid_bcast.to_sdfg() + process_grids = sdfg.process_grids + + assert not hasattr(sdfg, '_pgrids') + assert len(process_grids) == 1 + pgrid_name, pgrid = next(iter(process_grids.items())) + assert isinstance(pgrid, ProcessGrid) + assert pgrid.name == pgrid_name + assert sdfg.arrays[pgrid_name] is pgrid + + bcasts = [node for state in sdfg.states() for node in state.nodes() if isinstance(node, Bcast)] + assert len(bcasts) == 1 + assert _incoming_descriptor_name(sdfg, bcasts[0], '_grid') == pgrid_name + + +@pytest.mark.mpi +def test_create_cart_subgrid_bcast_uses_descriptor_name(): + MPI = pytest.importorskip('mpi4py.MPI') + + @dace.program + def subgrid_bcast(A: dace.int32[10], rank: dace.int32): + pgrid = MPI.COMM_WORLD.Create_cart([2, 1]) + if pgrid != MPI.COMM_NULL: + sgrid = pgrid.Sub([False, True]) + pgrid.Bcast(A) + B = np.empty_like(A) + B[:] = rank % 10 + if pgrid != MPI.COMM_NULL: + sgrid.Bcast(B) + A[:] = B + + sdfg = subgrid_bcast.to_sdfg() + process_grids = sdfg.process_grids + + assert len(process_grids) == 2 + assert all(pgrid.name == pgrid_name for pgrid_name, pgrid in process_grids.items()) + + init_code = sdfg.init_code['frame'].as_string + dummy_fields = '\n'.join(field for state in sdfg.states() for node in state.nodes() if isinstance(node, Dummy) + for field in node.fields) + for pgrid_name in process_grids: + assert f'__state->{pgrid_name}' in init_code + assert f'MPI_Comm {pgrid_name};' in dummy_fields + + bcasts = [node for state in sdfg.states() for node in state.nodes() if isinstance(node, Bcast)] + assert len(bcasts) == 2 + assert all(_incoming_descriptor_name(sdfg, bcast, '_grid') in process_grids for bcast in bcasts) + + sdfg.expand_library_nodes() + tasklet_code = '\n'.join(node.code.as_string for state in sdfg.states() for node in state.nodes() + if isinstance(node, dace.nodes.Tasklet)) + assert '_grid' in tasklet_code + assert '_comm' not in tasklet_code + + +if __name__ == "__main__": + test_create_cart_bcast_uses_process_grid_descriptor() + test_create_cart_subgrid_bcast_uses_descriptor_name() diff --git a/tests/library/mpi/mpi4py_test.py b/tests/library/mpi/mpi4py_test.py index c403a9dd2d..97dbc6a4a1 100644 --- a/tests/library/mpi/mpi4py_test.py +++ b/tests/library/mpi/mpi4py_test.py @@ -1,7 +1,6 @@ # Copyright 2019-2023 ETH Zurich and the DaCe authors. All rights reserved. import dace from dace.sdfg import utils -import dace.dtypes as dtypes import numpy as np import pytest @@ -321,8 +320,8 @@ def test_send_recv(): @dace.program def mpi4py_send_recv(rank: dace.int32, size: dace.int32): - src = np.full([1], (rank - 1) % size, dtype=np.int32) - dst = np.full([1], (rank + 1) % size, dtype=np.int32) + src = (rank - 1) % size + dst = (rank + 1) % size sbuf = np.full((1, ), rank, dtype=np.int32) commworld.Send(sbuf, dst, tag=0) rbuf = np.empty((1, ), dtype=np.int32) @@ -362,8 +361,7 @@ def mpi4py_alltoall(rank: dace.int32, size: dace.compiletime): val = func(rank=rank) ref = mpi4py_alltoall.f(rank, size) - if (not np.allclose(val, ref)): - raise (ValueError("The received values are not what I expected.")) + assert np.allclose(val, ref), "The received values are not what I expected." if __name__ == "__main__": diff --git a/tests/library/mpi/mpi_send_recv_test.py b/tests/library/mpi/mpi_send_recv_test.py index 9c8d78c042..63e435b6d3 100644 --- a/tests/library/mpi/mpi_send_recv_test.py +++ b/tests/library/mpi/mpi_send_recv_test.py @@ -39,6 +39,15 @@ def make_sdfg(dtype): return sdfg +def test_send_recv_validation_accepts_integer_rank_and_tag_descriptors(): + sdfg = make_sdfg(np.float64) + sdfg.arrays['src'].dtype = dace.int64 + sdfg.arrays['dest'].dtype = dace.uint32 + sdfg.arrays['tag'].dtype = dace.int16 + + sdfg.validate() + + ############################################################################### diff --git a/tests/library/mpi/process_grids_test.py b/tests/library/mpi/process_grids_test.py index 42fde912db..ca9bad96f2 100644 --- a/tests/library/mpi/process_grids_test.py +++ b/tests/library/mpi/process_grids_test.py @@ -1,4 +1,7 @@ # Copyright 2019-2022 ETH Zurich and the DaCe authors. All rights reserved. +from typing import Optional +from unittest import mock + import dace from dace.sdfg import utils import dace.dtypes as dtypes @@ -7,10 +10,36 @@ import pytest +class _MockProgramVisitor: + + def __init__(self): + self.globals = {} + self.scope_arrays = {} + self.scope_symbols = {} + self.current_lineinfo = None + self._target_name_counter = 0 + self._dispatcher = mock.MagicMock() + self._dispatcher.defined_vars = set() + + def get_target_name(self, output_index: Optional[int] = None, default: Optional[str] = None) -> str: + self._target_name_counter += 1 + return default or f'__pgrid{self._target_name_counter}' + + def __getattr__(self, name): + value = mock.MagicMock() + setattr(self, name, value) + return value + + +def _mock_program_visitor(): + return _MockProgramVisitor() + + @pytest.mark.mpi def test_process_grid(): P = dace.symbol('P', dace.int32) + pv = _mock_program_visitor() sdfg = dace.SDFG("process_grid_test") sdfg.add_symbol('P', dace.int32) @@ -20,17 +49,18 @@ def test_process_grid(): _, varr = sdfg.add_array("valid", (1, ), dtype=dace.bool_) state = sdfg.add_state("start") - pgrid_name = comm._cart_create(None, sdfg, state, [1, P]) + pgrid_name = comm._cart_create(pv, sdfg, state, [1, P]) state2 = sdfg.add_state("main") sdfg.add_edge(state, state2, dace.InterstateEdge()) - tasklet = state2.add_tasklet( - "MPI_Cart_get", {}, {'d', 'p', 'c', 'v'}, - f"MPI_Cart_get(__state->{pgrid_name}_comm, P, d, p, c);\nv = __state->{pgrid_name}_valid;", dtypes.Language.CPP) + tasklet = state2.add_tasklet("MPI_Cart_get", {'g'}, {'d', 'p', 'c', 'v'}, + "MPI_Cart_get(g, P, d, p, c);\nv = (g != MPI_COMM_NULL);", dtypes.Language.CPP) + pgrid = state2.add_read(pgrid_name) dims = state2.add_write("dims") periods = state2.add_write("periods") coords = state2.add_write("coords") valid = state2.add_write("valid") + state2.add_edge(pgrid, None, tasklet, 'g', dace.Memlet(data=pgrid_name)) state2.add_edge(tasklet, 'd', dims, None, dace.Memlet.from_array("dims", darr)) state2.add_edge(tasklet, 'p', periods, None, dace.Memlet.from_array("periods", parr)) state2.add_edge(tasklet, 'c', coords, None, dace.Memlet.from_array("coords", carr)) @@ -62,6 +92,7 @@ def test_process_grid(): def test_sub_grid(): P = dace.symbol('P', dace.int32) + pv = _mock_program_visitor() sdfg = dace.SDFG("sub_grid_test") sdfg.add_symbol('P', dace.int32) @@ -71,19 +102,19 @@ def test_sub_grid(): _, varr = sdfg.add_array("valid", (1, ), dtype=dace.bool_) state = sdfg.add_state("start") - parent_pgrid_name = comm._cart_create(None, sdfg, state, [1, P]) - pgrid_name = comm._cart_sub(None, sdfg, state, parent_pgrid_name, [False, True]) + parent_pgrid_name = comm._cart_create(pv, sdfg, state, [1, P]) + pgrid_name = comm._cart_sub(pv, sdfg, state, parent_pgrid_name, [False, True]) state2 = sdfg.add_state("main") sdfg.add_edge(state, state2, dace.InterstateEdge()) - tasklet = state2.add_tasklet( - "MPI_Cart_get", {}, {'d', 'p', 'c', 'v'}, - f"MPI_Cart_get(__state->{pgrid_name}_comm, P, &d, &p, &c);\nv = __state->{pgrid_name}_valid;", - dtypes.Language.CPP) + tasklet = state2.add_tasklet("MPI_Cart_get", {'g'}, {'d', 'p', 'c', 'v'}, + "MPI_Cart_get(g, P, &d, &p, &c);\nv = (g != MPI_COMM_NULL);", dtypes.Language.CPP) + pgrid = state2.add_read(pgrid_name) dims = state2.add_write("dims") periods = state2.add_write("periods") coords = state2.add_write("coords") valid = state2.add_write("valid") + state2.add_edge(pgrid, None, tasklet, 'g', dace.Memlet(data=pgrid_name)) state2.add_edge(tasklet, 'd', dims, None, dace.Memlet.from_array("dims", darr)) state2.add_edge(tasklet, 'p', periods, None, dace.Memlet.from_array("periods", parr)) state2.add_edge(tasklet, 'c', coords, None, dace.Memlet.from_array("coords", carr)) diff --git a/tests/library/stencil/stencil_node_test.py b/tests/library/stencil/stencil_node_test.py index d33814d3e9..3d1db45de8 100644 --- a/tests/library/stencil/stencil_node_test.py +++ b/tests/library/stencil/stencil_node_test.py @@ -2,7 +2,6 @@ import dace from dace.libraries.stencil import Stencil import numpy as np -from dace.transformation.interstate import InlineSDFG SIZE = dace.symbol("size") ROWS = dace.symbol("rows") diff --git a/tests/library/two_pkgs_test.py b/tests/library/two_pkgs_test.py index 7479be61c0..aa7c8d91ba 100644 --- a/tests/library/two_pkgs_test.py +++ b/tests/library/two_pkgs_test.py @@ -1,10 +1,8 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. import dace from dace import Memlet -from dace.codegen.exceptions import CompilerConfigurationError, CompilationError from dace.libraries.linalg import Inv import numpy as np -import warnings n = dace.symbol("n", dace.int64) diff --git a/tests/map_positive_step_validation_test.py b/tests/map_positive_step_validation_test.py new file mode 100644 index 0000000000..8c20e55d06 --- /dev/null +++ b/tests/map_positive_step_validation_test.py @@ -0,0 +1,115 @@ +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. +""" Map iteration must be ascending. + + A ``Map`` with a negative step fails SDFG validation. An empty + ``end < begin`` positive-step range is valid (it iterates zero times). + A symbolic step is checked at runtime by an ``assert(step > 0)`` the + CPU codegen emits, active only in Debug builds. + + The symbolic-step assertion runs in a subprocess: a failing C + ``assert`` raises ``SIGABRT`` rather than a Python exception, so the + Debug build type is set in the child's environment only. +""" +import os +import subprocess +import sys +import tempfile +import textwrap + +import pytest + +import dace + +# SDFG validation: negative-step maps are invalid; empty positive-step maps are valid. + + +def test_concrete_negative_step_map_is_invalid(): + """A map with a concrete negative step fails SDFG validation.""" + + @dace.program + def neg_step(a: dace.float64[8], b: dace.float64[8]): + for i in dace.map[5:0:-1]: + b[i] = a[i] + 1.0 + + with pytest.raises(Exception) as ei: + neg_step.to_sdfg(simplify=True).validate() + assert 'negative step' in str(ei.value) + + +def test_empty_positive_step_map_is_valid(): + """An empty ``end < begin`` positive-step map validates (a warning is + acceptable).""" + + @dace.program + def empty(a: dace.float64[16], b: dace.float64[16]): + for i in dace.map[10:3:1]: + b[i] = a[i] + 1.0 + + empty.to_sdfg(simplify=True).validate() # must not raise + + +def test_positive_step_map_is_valid(): + """An ascending map (unit or non-unit step) validates.""" + + @dace.program + def ok(a: dace.float64[32], b: dace.float64[32]): + for i in dace.map[2:30:2]: + b[i] = a[i] * 2.0 + + ok.to_sdfg(simplify=True).validate() # must not raise + + +# CPU codegen: a symbolic step is checked by a Debug-build runtime assertion. + +_CHILD = textwrap.dedent(''' + import numpy as np, dace + N = dace.symbol("N") + @dace.program + def symstep(a: dace.float64[N], b: dace.float64[N], S: dace.int64): + for i in dace.map[0:N:S]: + b[i] = a[i] + 1.0 + sdfg = symstep.to_sdfg(simplify=True) + a = np.random.rand(8); b = np.zeros(8) + sdfg(a=a, b=b, N=8, S={step}) + print("NO_ABORT") +''') + + +def _run_child(step: int): + """Run the symbolic-step kernel in a child Debug build (build type and + cache set only in the child env, never in the test process).""" + env = dict(os.environ) + env['DACE_compiler_build_type'] = 'Debug' + env['DACE_cache'] = 'unique' + # Must be a real file: the dace Python frontend cannot obtain source for + # a program defined via ``python -c``. + with tempfile.NamedTemporaryFile('w', suffix='.py', delete=False) as fh: + fh.write(_CHILD.format(step=step)) + path = fh.name + try: + return subprocess.run([sys.executable, path], env=env, capture_output=True, timeout=600) + finally: + os.unlink(path) + + +def test_symbolic_negative_step_aborts_in_debug_build(): + """A symbolic step that is negative at runtime trips the debug-only + ``assert(step > 0)`` -> the process aborts (SIGABRT).""" + proc = _run_child(-1) + assert proc.returncode != 0 and b'NO_ABORT' not in proc.stdout, \ + f"expected abort, rc={proc.returncode} out={proc.stdout!r}" + assert (proc.returncode == -6 or b'requires a positive step' in proc.stderr + or b'Assertion' in proc.stderr), \ + f"expected assertion failure, stderr={proc.stderr[-400:]!r}" + + +def test_symbolic_positive_step_runs_in_debug_build(): + """A positive symbolic step passes the assertion and runs normally in a + Debug build.""" + proc = _run_child(2) + assert proc.returncode == 0 and b'NO_ABORT' in proc.stdout, \ + f"expected clean run, rc={proc.returncode} stderr={proc.stderr[-400:]!r}" + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/memlet_lifetime_validation_test.py b/tests/memlet_lifetime_validation_test.py index db6ac9cc85..6417d71b1a 100644 --- a/tests/memlet_lifetime_validation_test.py +++ b/tests/memlet_lifetime_validation_test.py @@ -1,23 +1,24 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. -import numpy as np - import dace as dp from dace.sdfg import SDFG, InvalidSDFGError from dace.memlet import Memlet -from dace.data import Scalar + +import pytest def test(): - print('SDFG memlet lifetime validation test') # Externals (parameters, symbols) N = dp.symbol('N') # Construct SDFG 1 sdfg1 = SDFG('shouldntwork1') + sdfg1.add_array('A', [N], dp.int32) + sdfg1.add_array('B', [N], dp.int32) + sdfg1.add_transient('T', [1], dp.int32) state = sdfg1.add_state() - A = state.add_array('A', [N], dp.int32) - B = state.add_array('B', [N], dp.int32) - T = state.add_transient('T', [1], dp.int32) + A = state.add_access('A') + B = state.add_access('B') + T = state.add_access('T') tasklet_gen = state.add_tasklet('mytasklet', {'a'}, {'b'}, 'b = 5*a') map_entry, map_exit = state.add_map('mymap', dict(k='0:N')) @@ -33,18 +34,18 @@ def test(): state.add_edge(map_exit, 'OUT_1', tasklet_gen, 'a', Memlet.simple(B, '0')) state.add_edge(tasklet_gen, 'b', A, None, Memlet.simple(A, '0')) - try: + with pytest.raises(InvalidSDFGError): sdfg1.validate() - raise AssertionError("SDFG passed validation, test FAILED") - except InvalidSDFGError: - print("Test passed, exception successfully caught") # Construct SDFG 3 sdfg2 = SDFG('shouldntwork2') + sdfg2.add_array('A', [N], dp.int32) + sdfg2.add_array('B', [N], dp.int32) + sdfg2.add_transient('T', [N], dp.int32) state = sdfg2.add_state() - A = state.add_array('A', [N], dp.int32) - B = state.add_array('B', [N], dp.int32) - T = state.add_transient('T', [N], dp.int32) + A = state.add_access('A') + B = state.add_access('B') + T = state.add_access('T') tasklet_gen = state.add_tasklet('mytasklet', {'a'}, {'b'}, 'b = 5*a') map1_entry, map1_exit = state.add_map('mymap1', dict(k='0:N')) @@ -67,11 +68,8 @@ def test(): state.add_edge(T, None, map2_exit, 'IN_1', Memlet.simple(B, 'i')) state.add_edge(map2_exit, 'OUT_1', B, None, Memlet.simple(B, '0:N')) - try: + with pytest.raises(InvalidSDFGError): sdfg2.validate() - raise AssertionError("SDFG passed validation, test FAILED") - except (InvalidSDFGError, KeyError): - print("Test passed, exception successfully caught") if __name__ == '__main__': diff --git a/tests/memlet_propagation_squeezing_test.py b/tests/memlet_propagation_squeezing_test.py index ca2fe53088..265e345931 100644 --- a/tests/memlet_propagation_squeezing_test.py +++ b/tests/memlet_propagation_squeezing_test.py @@ -50,6 +50,90 @@ def make_sdfg(squeeze, name): return sdfg +def make_conditional_sdfg(): + M = dace.symbol('M') + sdfg = dace.SDFG('memlet_propagation_conditional') + sdfg.add_symbol('M', dace.int64) + sdfg.add_symbol('cond', dace.bool_) + sdfg.add_array('A', [4, M], dace.int64) + state = sdfg.add_state() + me, mx = state.add_map('map', dict(j='1:M')) + w = state.add_write('A') + + nsdfg = dace.SDFG('nested_conditional') + nsdfg.add_symbol('M', dace.int64) + nsdfg.add_symbol('cond', dace.bool_) + nsdfg.add_array('a', [4, M], dace.int64) + + cond_region = dace.sdfg.state.ConditionalBlock('if_region', sdfg=nsdfg) + + then_body = dace.sdfg.state.ControlFlowRegion('then_body', sdfg=nsdfg, parent=cond_region) + then_state = then_body.add_state('then_state', is_start_block=True) + then_write = then_state.add_write('a') + then_tasklet = then_state.add_tasklet('write_one', {}, {'out'}, 'out = 1') + then_state.add_edge(then_tasklet, 'out', then_write, None, dace.Memlet('a[0, 1]')) + + else_body = dace.sdfg.state.ControlFlowRegion('else_body', sdfg=nsdfg, parent=cond_region) + else_state = else_body.add_state('else_state', is_start_block=True) + else_write = else_state.add_write('a') + else_tasklet_1 = else_state.add_tasklet('write_two', {}, {'out'}, 'out = 2') + else_tasklet_2 = else_state.add_tasklet('write_three', {}, {'out'}, 'out = 3') + else_state.add_edge(else_tasklet_1, 'out', else_write, None, dace.Memlet('a[2, 0]')) + else_state.add_edge(else_tasklet_2, 'out', else_write, None, dace.Memlet('a[3, 0]')) + + cond_region.add_branch(dace.sdfg.state.CodeBlock('cond'), then_body) + cond_region.add_branch(dace.sdfg.state.CodeBlock('not cond'), else_body) + nsdfg.add_node(cond_region, is_start_block=True) + + nsdfg_node = state.add_nested_sdfg(nsdfg, {}, {'a'}, symbol_mapping=dict(j='j', M='M', cond='cond')) + state.add_nedge(me, nsdfg_node, dace.Memlet()) + state.add_memlet_path(nsdfg_node, mx, w, src_conn='a', memlet=dace.Memlet('A[0:4, j-1:j+1]')) + + propagation.propagate_memlets_sdfg(sdfg) + + return sdfg + + +def make_inverted_loop_sdfg(): + M = dace.symbol('M') + sdfg = dace.SDFG('memlet_propagation_inverted_loop') + sdfg.add_symbol('M', dace.int64) + sdfg.add_array('A', [5, M], dace.int64) + state = sdfg.add_state() + me, mx = state.add_map('map', dict(j='1:M')) + w = state.add_write('A') + + nsdfg = dace.SDFG('nested_inverted_loop') + nsdfg.using_explicit_control_flow = True + nsdfg.add_symbol('M', dace.int64) + nsdfg.add_symbol('i', dace.int64) + nsdfg.add_array('a', [5, M], dace.int64) + + loop_region = dace.sdfg.state.LoopRegion('loop_region', + condition_expr='i < 3', + loop_var='i', + initialize_expr='i = 0', + update_expr='i = i + 1', + inverted=True, + sdfg=nsdfg, + update_before_condition=False) + loop_state = loop_region.add_state('loop_state', is_start_block=True) + loop_write = loop_state.add_write('a') + loop_tasklet = loop_state.add_tasklet('write_value', {}, {'out'}, 'out = i + 7') + loop_state.add_edge(loop_tasklet, 'out', loop_write, None, dace.Memlet('a[i, 1]')) + loop_latch = loop_region.add_state('loop_latch') + loop_region.add_edge(loop_state, loop_latch, dace.InterstateEdge()) + nsdfg.add_node(loop_region, is_start_block=True) + + nsdfg_node = state.add_nested_sdfg(nsdfg, {}, {'a'}, symbol_mapping=dict(j='j', M='M')) + state.add_nedge(me, nsdfg_node, dace.Memlet()) + state.add_memlet_path(nsdfg_node, mx, w, src_conn='a', memlet=dace.Memlet('A[0:5, j-1]')) + + propagation.propagate_memlets_sdfg(sdfg) + + return sdfg + + def test_memlets_no_squeeze(): sdfg = make_sdfg(False, 'nonsqueezed') @@ -103,6 +187,56 @@ def test_memlets_squeeze(): assert out_memlet_2.subset[1] == (j - 1, j - 1, 1) +def test_memlets_conditional_upper_bound(): + sdfg = make_conditional_sdfg() + + M = 6 + A = np.zeros((4, M), dtype=np.int64) + expected_then = np.copy(A) + expected_then[0, 1:M] = 1 + + sdfg(A=A, M=M, cond=True) + assert np.allclose(A, expected_then) + + A = np.zeros((4, M), dtype=np.int64) + expected_else = np.copy(A) + expected_else[2, 0:M - 1] = 2 + expected_else[3, 0:M - 1] = 3 + + sdfg(A=A, M=M, cond=False) + assert np.allclose(A, expected_else) + + j = dace.symbolic.symbol('j') + main_state = sdfg.nodes()[0] + out_memlet = main_state.edges()[1].data + assert out_memlet.volume == 2 + assert out_memlet.dynamic == True + assert out_memlet.subset[0] == (0, 3, 1) + assert out_memlet.subset[1] == (j - 1, j, 1) + + +def test_memlets_inverted_loop(): + sdfg = make_inverted_loop_sdfg() + + M = 6 + A = np.zeros((5, M), dtype=np.int64) + expected = np.copy(A) + expected[0:4, 1:M] = np.reshape(np.arange(4) + 7, (4, 1)) + + sdfg(A=A, M=M) + assert np.allclose(A, expected) + + j = dace.symbolic.symbol('j') + main_state = sdfg.nodes()[0] + out_memlet = main_state.edges()[1].data + assert out_memlet.volume == 4 + assert out_memlet.dynamic == False + assert out_memlet.subset[0] == (0, 3, 1) + assert out_memlet.subset[1] == (j - 1, j, 1) + + if __name__ == '__main__': test_memlets_no_squeeze() test_memlets_squeeze() + test_memlets_conditional_upper_bound() + test_memlets_inverted_loop() diff --git a/tests/multiple_cr_test.py b/tests/multiple_cr_test.py index cf78c48242..cbadb3c1c2 100644 --- a/tests/multiple_cr_test.py +++ b/tests/multiple_cr_test.py @@ -1,5 +1,4 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. -import numpy as np import dace as dp from dace.sdfg import SDFG @@ -8,7 +7,6 @@ # Constructs an SDFG with multiple tasklets manually and runs it def test(): - print('SDFG multiple tasklet test') # Externals (parameters, symbols) N = dp.symbol('N') n = 20 @@ -21,10 +19,13 @@ def test(): # Construct SDFG mysdfg = SDFG('multiple_cr') + mysdfg.add_array('A', [N], dp.int64) + mysdfg.add_array('s', [1], dp.int64) + mysdfg.add_array('p', [1], dp.int64) state = mysdfg.add_state() - A = state.add_array('A', [N], dp.int64) - s = state.add_array('s', [1], dp.int64) - p = state.add_array('p', [1], dp.int64) + A = state.add_access('A') + s = state.add_access('s') + p = state.add_access('p') map_entry, map_exit = state.add_map('mymap', dict(i='0:N')) state.add_edge(A, None, map_entry, None, Memlet.simple(A, '0:N')) @@ -45,7 +46,6 @@ def test(): diff_sum = 5 * 20 - sum[0] diff_prod = 5**20 - product[0] - print("Difference:", diff_sum, '(sum)', diff_prod, '(product)') assert diff_sum <= 1e-5 and diff_prod <= 1e-5 diff --git a/tests/multiple_tasklet_test.py b/tests/multiple_tasklet_test.py index d7db9e1e36..9cfc7c4704 100644 --- a/tests/multiple_tasklet_test.py +++ b/tests/multiple_tasklet_test.py @@ -19,9 +19,11 @@ def test(): # Construct SDFG mysdfg = SDFG('multiple_tasklets') + mysdfg.add_array('A', [N], dp.int32) + mysdfg.add_array('B', [N], dp.int32) state = mysdfg.add_state() - A = state.add_array('A', [N], dp.int32) - B = state.add_array('B', [N], dp.int32) + A = state.add_access('A') + B = state.add_access('B') map_entry, map_exit = state.add_map('mymap', dict(i='0:N:2')) @@ -41,7 +43,6 @@ def test(): mysdfg(A=input, B=output, N=n) diff = np.linalg.norm(5 * input - output) / n - print("Difference:", diff) assert diff <= 1e-5 diff --git a/tests/multiprogram_cudatest.py b/tests/multiprogram_cudatest.py index 1b8dae3247..3a27cefc10 100644 --- a/tests/multiprogram_cudatest.py +++ b/tests/multiprogram_cudatest.py @@ -1,6 +1,5 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. import dace -from dace.transformation import optimizer from dace.transformation.dataflow import GPUTransformMap import numpy as np import pytest diff --git a/tests/nested_control_flow_test.py b/tests/nested_control_flow_test.py index 215a6afb90..b8d42f37a3 100644 --- a/tests/nested_control_flow_test.py +++ b/tests/nested_control_flow_test.py @@ -1,7 +1,6 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. import dace import numpy as np -import pytest @dace.program diff --git a/tests/nested_reduce_test.py b/tests/nested_reduce_test.py index fe46726ce6..aaf7372e83 100644 --- a/tests/nested_reduce_test.py +++ b/tests/nested_reduce_test.py @@ -3,26 +3,26 @@ import dace from dace.memlet import Memlet -# Create SDFG -sdfg = dace.SDFG('nested_reduction') -state = sdfg.add_state('a') -# Nodes -A = state.add_array('A', (40, ), dace.float32) -B = state.add_array('B', (20, ), dace.float32) -me, mx = state.add_map('mymap', dict(i='0:20')) -red = state.add_reduce('lambda a,b: a+b', None, 0) - -# Edges -state.add_edge(A, None, me, None, Memlet.simple(A, '0:40')) -state.add_edge(me, None, red, None, Memlet.simple(A, '(2*i):(2*i+2)')) -state.add_edge(red, None, mx, None, Memlet.simple(B, 'i')) -state.add_edge(mx, None, B, None, Memlet.simple(B, '0:20')) -sdfg.fill_scope_connectors() - - -def test(): - print('Nested reduction test') +def test_nested_reduction(): + # Create SDFG + sdfg = dace.SDFG('nested_reduction') + sdfg.add_array('A', (40, ), dace.float32) + sdfg.add_array('B', (20, ), dace.float32) + + # Nodes + state = sdfg.add_state('a') + A = state.add_access('A') + B = state.add_access('B') + me, mx = state.add_map('mymap', dict(i='0:20')) + red = state.add_reduce('lambda a,b: a+b', None, 0) + + # Edges + state.add_edge(A, None, me, None, Memlet.simple(A, '0:40')) + state.add_edge(me, None, red, None, Memlet.simple(A, '(2*i):(2*i+2)')) + state.add_edge(red, None, mx, None, Memlet.simple(B, 'i')) + state.add_edge(mx, None, B, None, Memlet.simple(B, '0:20')) + sdfg.fill_scope_connectors() Adata = np.random.rand(40).astype(np.float32) Bdata = np.random.rand(20).astype(np.float32) @@ -33,9 +33,8 @@ def test(): B_regression[:] += Adata[1::2] diff = np.linalg.norm(B_regression - Bdata) / 20.0 - print("Difference:", diff) assert diff <= 1e-5 if __name__ == "__main__": - test() + test_nested_reduction() diff --git a/tests/nested_sdfg_python_test.py b/tests/nested_sdfg_python_test.py index 5d27a3e702..7bdf9a0094 100644 --- a/tests/nested_sdfg_python_test.py +++ b/tests/nested_sdfg_python_test.py @@ -2,8 +2,6 @@ import numpy as np import dace as dp -from dace.sdfg import SDFG -from dace.memlet import Memlet N = dp.symbol('N') diff --git a/tests/nested_sdfg_test.py b/tests/nested_sdfg_test.py index 54af631da8..bda3936589 100644 --- a/tests/nested_sdfg_test.py +++ b/tests/nested_sdfg_test.py @@ -32,9 +32,11 @@ def do(): # Construct SDFG mysdfg = SDFG('outer_sdfg') + mysdfg.add_array('A', [N, N], dp.float32) + mysdfg.add_array('B', [N, N], dp.float32) state = mysdfg.add_state() - A = state.add_array('A', [N, N], dp.float32) - B = state.add_array('B', [N, N], dp.float32) + A = state.add_access('A') + B = state.add_access('B') map_entry, map_exit = state.add_map('elements', [('i', '0:N'), ('j', '0:N')]) nsdfg = state.add_nested_sdfg(sdfg_internal.to_sdfg(), {'input'}, {'output'}) @@ -53,7 +55,6 @@ def do(): mysdfg(A=input, B=output, N=N) diff = np.linalg.norm(output - np.power(input, 5)) / (N * N) - print("Difference:", diff) assert diff <= 1e-5 @@ -80,9 +81,11 @@ def do(): # Construct SDFG mysdfg = SDFG('outer_sdfg') + mysdfg.add_array('A', [N, N], dp.float32) + mysdfg.add_array('B', [N, N], dp.float32) state = mysdfg.add_state() - A = state.add_array('A', [N, N], dp.float32) - B = state.add_array('B', [N, N], dp.float32) + A = state.add_access('A') + B = state.add_access('B') map_entry, map_exit = state.add_map('elements', [('i', '0:N'), ('j', '0:N')]) internal = sdfg_internal.to_sdfg() @@ -104,7 +107,6 @@ def do(): mysdfg(A=input, B=output, N=N) diff = np.linalg.norm(output - np.power(input, 5)) / (N * N) - print("Difference:", diff) assert diff <= 1e-5 os.close(fd) diff --git a/tests/nested_vector_type_test.py b/tests/nested_vector_type_test.py index 990d96b1e0..f3983d6054 100644 --- a/tests/nested_vector_type_test.py +++ b/tests/nested_vector_type_test.py @@ -2,7 +2,6 @@ import dace from dace.memlet import Memlet -import dace.libraries.blas as blas import numpy as np import scipy diff --git a/tests/node_id_test.py b/tests/node_id_test.py new file mode 100644 index 0000000000..bd7db59aac --- /dev/null +++ b/tests/node_id_test.py @@ -0,0 +1,90 @@ +"""Test for node ID increment functionality.""" + +from copy import deepcopy +import pytest +from dace.sdfg.nodes import Node, AccessNode, reset_node_id_counter + + +def test_node_id_increment(): + """Test that creating new nodes increments the ID.""" + with reset_node_id_counter(): + node1 = Node() + node2 = Node() + node3 = Node() + + # IDs should be unique and increasing + assert node1.id == 0, "First node should have ID 0" + assert node2.id == 1, "Second node should have ID 1" + assert node3.id == 2, "Third node should have ID 2" + + +def test_node_deepcopy(): + """Test that deepcopy assigns a new ID.""" + with reset_node_id_counter(): + node1 = Node() + node1_copy = deepcopy(node1) + + # Copied node should have a different ID + assert node1.id == 0, "Original node should have ID 0" + assert node1_copy.id == 1, "Copied node should have ID 1" + assert node1_copy.id != node1.id, "Deepcopied node should have a different ID" + assert node1_copy.guid != node1.guid, "Deepcopied node should have a different GUID" + + +def test_access_node_id(): + """Test that AccessNode also gets incremented IDs.""" + with reset_node_id_counter(): + access1 = AccessNode("data1") + access2 = AccessNode("data2") + + # AccessNode IDs should be unique and increasing + assert access1.id == 0, "First AccessNode should have ID 0" + assert access2.id == 1, "Second AccessNode should have ID 1" + + +def test_access_node_deepcopy(): + """Test that deepcopy of AccessNode assigns a new ID.""" + with reset_node_id_counter(): + access1 = AccessNode("data1") + access1_copy = deepcopy(access1) + + # Copied AccessNode should have a different ID + assert access1.id == 0, "Original AccessNode should have ID 0" + assert access1_copy.id == 1, "Copied AccessNode should have ID 1" + assert access1_copy.id != access1.id, "Deepcopied AccessNode should have a different ID" + assert access1_copy.guid != access1.guid, "Deepcopied AccessNode should have a different GUID" + + +def test_multiple_nodes_unique_ids(): + """Test that multiple nodes have unique IDs.""" + with reset_node_id_counter(): + nodes = [Node() for _ in range(10)] + ids = [node.id for node in nodes] + + # All IDs should be unique + assert len(ids) == len(set(ids)), "All node IDs should be unique" + assert ids == list(range(10)), "IDs should be sequential from 0 to 9" + + +def test_context_manager_restores_counter(): + """Test that the context manager restores the old counter value.""" + with reset_node_id_counter(): + node1 = Node() + assert node1.id == 0 + + # Create a node outside the context manager + node2 = Node() + + # Now use the context manager again + with reset_node_id_counter(): + node3 = Node() + assert node3.id == 0, "Counter should be reset to 0" + + # Create another node outside + node4 = Node() + # node4 should continue from where node2 left off + assert node4.id == node2.id + 1, "Counter should be restored after context manager exits" + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/npbench/deep_learning/conv2d_bias_test.py b/tests/npbench/deep_learning/conv2d_bias_test.py index 5b776b84c9..a99e58b65c 100644 --- a/tests/npbench/deep_learning/conv2d_bias_test.py +++ b/tests/npbench/deep_learning/conv2d_bias_test.py @@ -5,9 +5,7 @@ import dace as dc import pytest import argparse -from dace.transformation.interstate import InlineSDFG from dace.transformation.auto.auto_optimize import auto_optimize -from dace.config import set_temporary from dace.autodiff import add_backward_pass C_in, C_out, H, K, N, W = (dc.symbol(s, dc.int64) for s in ('C_in', 'C_out', 'H', 'K', 'N', 'W')) diff --git a/tests/npbench/deep_learning/mlp_test.py b/tests/npbench/deep_learning/mlp_test.py index 0dd68fc467..13ea8ac240 100644 --- a/tests/npbench/deep_learning/mlp_test.py +++ b/tests/npbench/deep_learning/mlp_test.py @@ -43,7 +43,7 @@ def initialize(C_in, N, S0, S1, S2): mlp_sizes = [S0, S1, S2] # [300, 100, 10] # Inputs - input = np.random.rand(N, C_in).astype(np.float32) + input = rng.random((N, C_in), dtype=np.float32) # Weights w1 = rng.random((C_in, mlp_sizes[0]), dtype=np.float32) b1 = rng.random((mlp_sizes[0], ), dtype=np.float32) @@ -160,7 +160,7 @@ def autodiff_kernel(input: dc.float32[N, C_in], w1: dc.float32[C_in, S0], b1: dc jax_kernel = lambda input, w1, b1, w2, b2, w3, b3: mlp_jax_kernel(jnp, input, w1, b1, w2, b2, w3, b3) jax_grad = jax.jit(jax.grad(jax_kernel, argnums=0)) jax_grad_input = jax_grad(input, w1, b1, w2, b2, w3, b3) - np.testing.assert_allclose(gradient_input, jax_grad_input, rtol=1e-4, atol=1e-10) + np.testing.assert_allclose(gradient_input, jax_grad_input, rtol=1e-4, atol=1e-6) def test_cpu(): diff --git a/tests/npbench/misc/compute_test.py b/tests/npbench/misc/compute_test.py index 9bc220ddbf..31956b4cd0 100644 --- a/tests/npbench/misc/compute_test.py +++ b/tests/npbench/misc/compute_test.py @@ -6,7 +6,6 @@ import pytest import argparse from dace.transformation.auto.auto_optimize import auto_optimize -from dace.autodiff import add_backward_pass def relerror(val, ref): diff --git a/tests/npbench/misc/crc16_test.py b/tests/npbench/misc/crc16_test.py index 4b5d7b73e3..928a9e9d64 100644 --- a/tests/npbench/misc/crc16_test.py +++ b/tests/npbench/misc/crc16_test.py @@ -2,7 +2,6 @@ # Original application code: NPBench - https://github.com/spcl/npbench import dace.dtypes import numpy as np -import dace as dc import pytest import argparse from dace.transformation.auto.auto_optimize import auto_optimize diff --git a/tests/npbench/polybench/atax_test.py b/tests/npbench/polybench/atax_test.py index 8dbd812b82..93120c9517 100644 --- a/tests/npbench/polybench/atax_test.py +++ b/tests/npbench/polybench/atax_test.py @@ -5,9 +5,7 @@ import dace as dc import pytest import argparse -from dace.transformation.interstate import InlineSDFG from dace.transformation.auto.auto_optimize import auto_optimize -from dace.config import set_temporary from dace.autodiff import add_backward_pass # Data set sizes diff --git a/tests/npbench/polybench/bicg_test.py b/tests/npbench/polybench/bicg_test.py index e90b728609..d21815d6e3 100644 --- a/tests/npbench/polybench/bicg_test.py +++ b/tests/npbench/polybench/bicg_test.py @@ -5,7 +5,6 @@ import dace as dc import pytest import argparse -from dace.transformation.interstate import InlineSDFG from dace.transformation.auto.auto_optimize import auto_optimize from dace.autodiff import add_backward_pass diff --git a/tests/npbench/polybench/cholesky2_test.py b/tests/npbench/polybench/cholesky2_test.py index 50486f9126..96c6b8eb25 100644 --- a/tests/npbench/polybench/cholesky2_test.py +++ b/tests/npbench/polybench/cholesky2_test.py @@ -6,7 +6,6 @@ import dace as dc import pytest import argparse -from dace.transformation.interstate import InlineSDFG from dace.transformation.auto.auto_optimize import auto_optimize # Data set sizes diff --git a/tests/npbench/polybench/cholesky_test.py b/tests/npbench/polybench/cholesky_test.py index 309bb3197d..8e2d48fb05 100644 --- a/tests/npbench/polybench/cholesky_test.py +++ b/tests/npbench/polybench/cholesky_test.py @@ -6,7 +6,6 @@ import dace as dc import pytest import argparse -from dace.transformation.interstate import InlineSDFG from dace.transformation.auto.auto_optimize import auto_optimize from dace.autodiff import add_backward_pass diff --git a/tests/npbench/polybench/covariance_test.py b/tests/npbench/polybench/covariance_test.py index 1c8e2ed2a6..07d02c62c6 100644 --- a/tests/npbench/polybench/covariance_test.py +++ b/tests/npbench/polybench/covariance_test.py @@ -7,11 +7,7 @@ import dace as dc import pytest import argparse -from dace.transformation.interstate import InlineSDFG -from dace.transformation.dataflow import StreamingMemory, MapFusionVertical, StreamingComposition, PruneConnectors from dace.transformation.auto.auto_optimize import auto_optimize -from dace.libraries.standard import Reduce -from dace.libraries.blas import Gemv from dace.autodiff import add_backward_pass # Data set sizes diff --git a/tests/npbench/polybench/deriche_test.py b/tests/npbench/polybench/deriche_test.py index ac40cd14fc..606c22ed90 100644 --- a/tests/npbench/polybench/deriche_test.py +++ b/tests/npbench/polybench/deriche_test.py @@ -5,7 +5,6 @@ import dace as dc import pytest import argparse -from dace.transformation.interstate import InlineSDFG from dace.transformation.auto.auto_optimize import auto_optimize from dace.autodiff import add_backward_pass diff --git a/tests/npbench/polybench/doitgen_test.py b/tests/npbench/polybench/doitgen_test.py index 45da2d635f..a97931e4d0 100644 --- a/tests/npbench/polybench/doitgen_test.py +++ b/tests/npbench/polybench/doitgen_test.py @@ -5,7 +5,6 @@ import dace as dc import pytest import argparse -from dace.transformation.interstate import InlineSDFG from dace.transformation.auto.auto_optimize import auto_optimize from dace.autodiff import add_backward_pass diff --git a/tests/npbench/polybench/durbin_test.py b/tests/npbench/polybench/durbin_test.py index c119dd1361..df0860b5bd 100644 --- a/tests/npbench/polybench/durbin_test.py +++ b/tests/npbench/polybench/durbin_test.py @@ -5,9 +5,7 @@ import dace as dc import pytest import argparse -from dace.transformation.interstate import InlineSDFG from dace.transformation.auto.auto_optimize import auto_optimize -from dace.config import set_temporary from dace.autodiff import add_backward_pass # Data set sizes diff --git a/tests/npbench/polybench/fdtd_2d_test.py b/tests/npbench/polybench/fdtd_2d_test.py index ff0a9a9913..7012d346e7 100644 --- a/tests/npbench/polybench/fdtd_2d_test.py +++ b/tests/npbench/polybench/fdtd_2d_test.py @@ -5,8 +5,6 @@ import numpy as np import dace as dc import pytest -from dace.transformation.interstate import InlineSDFG -from dace.transformation.dataflow import StreamingMemory, StreamingComposition, MapFusionVertical from dace.transformation.auto.auto_optimize import auto_optimize import argparse from dace.autodiff import add_backward_pass diff --git a/tests/npbench/polybench/floyd_warshall_test.py b/tests/npbench/polybench/floyd_warshall_test.py index 48a8e4485f..d343f4b92f 100644 --- a/tests/npbench/polybench/floyd_warshall_test.py +++ b/tests/npbench/polybench/floyd_warshall_test.py @@ -6,8 +6,6 @@ import dace as dc import pytest import argparse -from dace.transformation.interstate import InlineSDFG, StateFusion -from dace.transformation.dataflow import StreamingMemory, MapFusionVertical, StreamingComposition, PruneConnectors from dace.transformation.auto.auto_optimize import auto_optimize # Data set sizes diff --git a/tests/npbench/polybench/gemm_npbench_test.py b/tests/npbench/polybench/gemm_npbench_test.py index 7c078d5e0d..e3722cdac3 100644 --- a/tests/npbench/polybench/gemm_npbench_test.py +++ b/tests/npbench/polybench/gemm_npbench_test.py @@ -5,7 +5,6 @@ import dace as dc import pytest import argparse -from dace.transformation.interstate import InlineSDFG from dace.transformation.auto.auto_optimize import auto_optimize from dace.autodiff import add_backward_pass diff --git a/tests/npbench/polybench/gemver_test.py b/tests/npbench/polybench/gemver_test.py index e49cf0c7d7..1c8ce29e22 100644 --- a/tests/npbench/polybench/gemver_test.py +++ b/tests/npbench/polybench/gemver_test.py @@ -5,7 +5,6 @@ import dace as dc import pytest import argparse -from dace.transformation.interstate import InlineSDFG from dace.transformation.auto.auto_optimize import auto_optimize from dace.autodiff import add_backward_pass diff --git a/tests/npbench/polybench/gesummv_test.py b/tests/npbench/polybench/gesummv_test.py index 62a13ba452..3198dc35ec 100644 --- a/tests/npbench/polybench/gesummv_test.py +++ b/tests/npbench/polybench/gesummv_test.py @@ -5,7 +5,6 @@ import dace as dc import pytest import argparse -from dace.transformation.interstate import InlineSDFG from dace.transformation.auto.auto_optimize import auto_optimize from dace.autodiff import add_backward_pass diff --git a/tests/npbench/polybench/gramschmidt_test.py b/tests/npbench/polybench/gramschmidt_test.py index 79e86b2da3..b94b86ecf7 100644 --- a/tests/npbench/polybench/gramschmidt_test.py +++ b/tests/npbench/polybench/gramschmidt_test.py @@ -5,9 +5,7 @@ import dace as dc import pytest import argparse -from dace.transformation.interstate import InlineSDFG from dace.transformation.auto.auto_optimize import auto_optimize -from dace.config import set_temporary from dace.autodiff import add_backward_pass # Data set sizes diff --git a/tests/npbench/polybench/heat_3d_test.py b/tests/npbench/polybench/heat_3d_test.py index 84b2d7a11c..fbd7f1306b 100644 --- a/tests/npbench/polybench/heat_3d_test.py +++ b/tests/npbench/polybench/heat_3d_test.py @@ -5,9 +5,7 @@ import dace as dc import pytest import argparse -from dace.transformation.interstate import InlineSDFG from dace.transformation.auto.auto_optimize import auto_optimize -from dace.config import set_temporary from dace.autodiff import add_backward_pass # Dataset sizes diff --git a/tests/npbench/polybench/jacobi_1d_test.py b/tests/npbench/polybench/jacobi_1d_test.py index ef324c2763..64f87c4693 100644 --- a/tests/npbench/polybench/jacobi_1d_test.py +++ b/tests/npbench/polybench/jacobi_1d_test.py @@ -5,7 +5,6 @@ import dace as dc import pytest import argparse -from dace.transformation.interstate import InlineSDFG from dace.transformation.auto.auto_optimize import auto_optimize from dace.autodiff import add_backward_pass diff --git a/tests/npbench/polybench/jacobi_2d_test.py b/tests/npbench/polybench/jacobi_2d_test.py index 6ad77c4163..05efa43243 100644 --- a/tests/npbench/polybench/jacobi_2d_test.py +++ b/tests/npbench/polybench/jacobi_2d_test.py @@ -5,8 +5,6 @@ import dace as dc import pytest import argparse -from dace.transformation.interstate import InlineSDFG -from dace.transformation.dataflow import StreamingMemory, MapFusionVertical from dace.transformation.auto.auto_optimize import auto_optimize from dace.autodiff import add_backward_pass diff --git a/tests/npbench/polybench/k2mm_test.py b/tests/npbench/polybench/k2mm_test.py index 7fedd26df1..78f281a2bd 100644 --- a/tests/npbench/polybench/k2mm_test.py +++ b/tests/npbench/polybench/k2mm_test.py @@ -5,7 +5,6 @@ import dace as dc import pytest import argparse -from dace.transformation.interstate import InlineSDFG from dace.transformation.auto.auto_optimize import auto_optimize from dace.autodiff import add_backward_pass diff --git a/tests/npbench/polybench/k3mm_test.py b/tests/npbench/polybench/k3mm_test.py index c726c3d9c4..25c3e140b3 100644 --- a/tests/npbench/polybench/k3mm_test.py +++ b/tests/npbench/polybench/k3mm_test.py @@ -5,7 +5,6 @@ import dace as dc import pytest import argparse -from dace.transformation.interstate import InlineSDFG from dace.transformation.auto.auto_optimize import auto_optimize from dace.autodiff import add_backward_pass diff --git a/tests/npbench/polybench/lu_test.py b/tests/npbench/polybench/lu_test.py index 6f812d0f25..d5ec2eb1e3 100644 --- a/tests/npbench/polybench/lu_test.py +++ b/tests/npbench/polybench/lu_test.py @@ -6,8 +6,6 @@ import dace as dc import pytest import argparse -from dace.transformation.interstate import InlineSDFG -from dace.transformation.dataflow import StreamingMemory, MapFusionVertical, StreamingComposition, PruneConnectors from dace.transformation.auto.auto_optimize import auto_optimize from dace.autodiff import add_backward_pass diff --git a/tests/npbench/polybench/ludcmp_test.py b/tests/npbench/polybench/ludcmp_test.py index 509ab75f98..2795aeb63c 100644 --- a/tests/npbench/polybench/ludcmp_test.py +++ b/tests/npbench/polybench/ludcmp_test.py @@ -6,9 +6,7 @@ import dace as dc import pytest import argparse -from dace.transformation.interstate import InlineSDFG from dace.transformation.auto.auto_optimize import auto_optimize -from dace.config import set_temporary from dace.autodiff import add_backward_pass # Dataset sizes diff --git a/tests/npbench/polybench/mvt_test.py b/tests/npbench/polybench/mvt_test.py index f3c35dfbd5..4dae81f0f4 100644 --- a/tests/npbench/polybench/mvt_test.py +++ b/tests/npbench/polybench/mvt_test.py @@ -5,7 +5,6 @@ import dace as dc import pytest import argparse -from dace.transformation.interstate import InlineSDFG from dace.transformation.auto.auto_optimize import auto_optimize from dace.autodiff import add_backward_pass diff --git a/tests/npbench/polybench/nussinov_test.py b/tests/npbench/polybench/nussinov_test.py index 6b96fd94ee..1c199cab4c 100644 --- a/tests/npbench/polybench/nussinov_test.py +++ b/tests/npbench/polybench/nussinov_test.py @@ -6,10 +6,7 @@ import dace as dc import pytest import argparse -from dace.transformation.interstate import InlineSDFG -from dace.transformation.dataflow import StreamingMemory, MapFusionVertical, StreamingComposition, PruneConnectors from dace.transformation.auto.auto_optimize import auto_optimize -from dace.config import set_temporary N = dc.symbol('N', dtype=dc.int32) diff --git a/tests/npbench/polybench/seidel_2d_test.py b/tests/npbench/polybench/seidel_2d_test.py index 0b4b0aeff0..eef5a86184 100644 --- a/tests/npbench/polybench/seidel_2d_test.py +++ b/tests/npbench/polybench/seidel_2d_test.py @@ -5,9 +5,6 @@ import dace as dc import pytest import argparse -from dace.transformation.interstate import InlineSDFG -from dace.transformation.auto.auto_optimize import auto_optimize -from dace.config import set_temporary from dace.autodiff import add_backward_pass # Dataset sizes diff --git a/tests/npbench/polybench/symm_test.py b/tests/npbench/polybench/symm_test.py index 17a4a5232b..2287359df8 100644 --- a/tests/npbench/polybench/symm_test.py +++ b/tests/npbench/polybench/symm_test.py @@ -5,7 +5,6 @@ import dace as dc import pytest import argparse -from dace.transformation.interstate import InlineSDFG from dace.transformation.auto.auto_optimize import auto_optimize from dace.autodiff import add_backward_pass diff --git a/tests/npbench/polybench/syr2k_test.py b/tests/npbench/polybench/syr2k_test.py index 411cc206e6..6b61bcc69f 100644 --- a/tests/npbench/polybench/syr2k_test.py +++ b/tests/npbench/polybench/syr2k_test.py @@ -5,7 +5,6 @@ import dace as dc import pytest import argparse -from dace.transformation.interstate import InlineSDFG from dace.transformation.auto.auto_optimize import auto_optimize from dace.autodiff import add_backward_pass diff --git a/tests/npbench/polybench/syrk_test.py b/tests/npbench/polybench/syrk_test.py index 33771b37be..359b58de5b 100644 --- a/tests/npbench/polybench/syrk_test.py +++ b/tests/npbench/polybench/syrk_test.py @@ -6,8 +6,6 @@ import dace as dc import pytest import argparse -from dace.transformation.interstate import InlineSDFG -from dace.transformation.dataflow import StreamingMemory, MapFusionVertical, StreamingComposition, PruneConnectors from dace.transformation.auto.auto_optimize import auto_optimize from dace.autodiff import add_backward_pass diff --git a/tests/npbench/polybench/trisolv_test.py b/tests/npbench/polybench/trisolv_test.py index af3be43b0a..c91c6242ad 100644 --- a/tests/npbench/polybench/trisolv_test.py +++ b/tests/npbench/polybench/trisolv_test.py @@ -5,7 +5,6 @@ import dace as dc import pytest import argparse -from dace.transformation.interstate import InlineSDFG from dace.transformation.auto.auto_optimize import auto_optimize from dace.autodiff import add_backward_pass diff --git a/tests/npbench/polybench/trmm_test.py b/tests/npbench/polybench/trmm_test.py index 110591d07e..d3cdbf9263 100644 --- a/tests/npbench/polybench/trmm_test.py +++ b/tests/npbench/polybench/trmm_test.py @@ -6,7 +6,6 @@ import dace as dc import pytest import argparse -from dace.transformation.interstate import InlineSDFG from dace.transformation.auto.auto_optimize import auto_optimize from dace.autodiff import add_backward_pass diff --git a/tests/numpy/attention_test.py b/tests/numpy/attention_test.py index 6b27af480c..015db6ddb8 100644 --- a/tests/numpy/attention_test.py +++ b/tests/numpy/attention_test.py @@ -1,6 +1,5 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. import dace -import numpy as np M = dace.symbol('M') N = dace.symbol('N') diff --git a/tests/numpy/augassign_test.py b/tests/numpy/augassign_test.py index 1ddeeaf9a8..932b535e1a 100644 --- a/tests/numpy/augassign_test.py +++ b/tests/numpy/augassign_test.py @@ -1,6 +1,5 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. import dace -import numpy as np from common import compare_numpy_output diff --git a/tests/numpy/binop_test.py b/tests/numpy/binop_test.py index 220d0930db..af1acbd120 100644 --- a/tests/numpy/binop_test.py +++ b/tests/numpy/binop_test.py @@ -1,7 +1,6 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. import dace import numpy as np -from copy import deepcopy as dc from common import compare_numpy_output diff --git a/tests/numpy/binop_with_scalars_test.py b/tests/numpy/binop_with_scalars_test.py index 79296ee539..70b3c4afaa 100644 --- a/tests/numpy/binop_with_scalars_test.py +++ b/tests/numpy/binop_with_scalars_test.py @@ -1,7 +1,6 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. import dace import numpy as np -import pytest from common import compare_numpy_output ### Left ##################################################################### diff --git a/tests/numpy/dTGL_test.py b/tests/numpy/dTGL_test.py index 647b67cfb2..b09d073829 100644 --- a/tests/numpy/dTGL_test.py +++ b/tests/numpy/dTGL_test.py @@ -1,6 +1,5 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. import dace -import numpy as np # Declaration of symbolic variables N, BS = (dace.symbol(name) for name in ['N', 'BS']) diff --git a/tests/numpy/flip_test.py b/tests/numpy/flip_test.py index 46d258e87f..bd3fc62f8d 100644 --- a/tests/numpy/flip_test.py +++ b/tests/numpy/flip_test.py @@ -1,6 +1,5 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. import numpy as np -from sympy.core.numbers import comp import dace from common import compare_numpy_output diff --git a/tests/numpy/linalg_test.py b/tests/numpy/linalg_test.py index d7d00a3807..15132fc04e 100644 --- a/tests/numpy/linalg_test.py +++ b/tests/numpy/linalg_test.py @@ -1,4 +1,4 @@ -# Copyright 2019-2023 ETH Zurich and the DaCe authors. All rights reserved. +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. import dace import numpy as np import pytest @@ -90,12 +90,10 @@ def tensordot_0(A: dace.float32[3, 3, 3, 3, 3, 3], B: dace.float32[3, 3, 3, 3, 3 assert (np.allclose(tensordot_0(A.copy(), B.copy()), tensordot_0.f(A, B))) -# TODO: Enable after fixing cuTENSOR in CI -#@pytest.mark.gpu -@pytest.mark.skip('CUTENSOR is not supported in CI') +@pytest.mark.gpu def test_tensordot_02(): - @dace.program(device=dace.dtypes.DeviceType.GPU) + @dace.program(device=dace.dtypes.DeviceType.GPU, auto_optimize=True) def tensordot_0(A: dace.float32[3, 3, 3, 3, 3, 3], B: dace.float32[3, 3, 3, 3, 3, 3]): return np.tensordot(A, B) @@ -129,12 +127,10 @@ def tensordot_1(A: dace.float32[3, 3, 3, 3, 3, 3], B: dace.float32[3, 3, 3, 3, 3 assert (np.allclose(tensordot_1(A.copy(), B.copy()), tensordot_1.f(A, B))) -# TODO: Enable after fixing cuTENSOR in CI -#@pytest.mark.gpu -@pytest.mark.skip('CUTENSOR is not supported in CI') +@pytest.mark.gpu def test_tensordot_12(): - @dace.program(device=dace.dtypes.DeviceType.GPU) + @dace.program(device=dace.dtypes.DeviceType.GPU, auto_optimize=True) def tensordot_1(A: dace.float32[3, 3, 3, 3, 3, 3], B: dace.float32[3, 3, 3, 3, 3, 3]): return np.tensordot(A, B, axes=([0, 3], [4, 2])) @@ -190,12 +186,10 @@ def tensordot_2b(A: dace.float32[3, 3, 3, 3, 3, 3], B: dace.float32[3, 3, 3, 3, assert (np.allclose(tensordot_2b(A.copy(), B.copy()), ref)) -# TODO: Enable after fixing cuTENSOR in CI -#@pytest.mark.gpu -@pytest.mark.skip('CUTENSOR is not supported in CI') +@pytest.mark.gpu def test_tensordot_22(): - @dace.program(device=dace.dtypes.DeviceType.GPU) + @dace.program(device=dace.dtypes.DeviceType.GPU, auto_optimize=True) def tensordot_2a(A: dace.float32[3, 3, 3, 3, 3, 3], B: dace.float32[3, 3, 3, 3, 3, 3]): return np.tensordot(A, B, axes=([0, 3], [4, 2]), out_axes=[7, 6, 5, 4, 3, 2, 1, 0]) @@ -205,7 +199,7 @@ def tensordot_2a(A: dace.float32[3, 3, 3, 3, 3, 3], B: dace.float32[3, 3, 3, 3, with dace.config.set_temporary('library', 'linalg', 'default_implementation', value='cuTENSOR'): assert (np.allclose(tensordot_2a(A.copy(), B.copy()), ref)) - @dace.program(device=dace.dtypes.DeviceType.GPU) + @dace.program(device=dace.dtypes.DeviceType.GPU, auto_optimize=True) def tensordot_2b(A: dace.float32[3, 3, 3, 3, 3, 3], B: dace.float32[3, 3, 3, 3, 3, 3]): return np.tensordot(A, B, axes=([0, 3], [4, 2]), out_axes=[0, 7, 1, 6, 2, 5, 3, 4]) diff --git a/tests/numpy/list_globals_test.py b/tests/numpy/list_globals_test.py index 9facde0e67..54f0378b89 100644 --- a/tests/numpy/list_globals_test.py +++ b/tests/numpy/list_globals_test.py @@ -5,7 +5,6 @@ ################## # Lists -from dace.frontend.python.common import DaceSyntaxError global_axes = [0, 2, 1] diff --git a/tests/numpy/npdot_test.py b/tests/numpy/npdot_test.py index 7fcd1805f3..3ae3d25079 100644 --- a/tests/numpy/npdot_test.py +++ b/tests/numpy/npdot_test.py @@ -1,7 +1,6 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. import dace import numpy as np -import pytest from common import compare_numpy_output diff --git a/tests/numpy/reductions_test.py b/tests/numpy/reductions_test.py index f90f1f21de..8a5b8e5b55 100644 --- a/tests/numpy/reductions_test.py +++ b/tests/numpy/reductions_test.py @@ -1,7 +1,6 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. import dace import numpy as np -from copy import deepcopy as dc from common import compare_numpy_output diff --git a/tests/numpy/rot90_test.py b/tests/numpy/rot90_test.py index 8a60460c43..c998edc19b 100644 --- a/tests/numpy/rot90_test.py +++ b/tests/numpy/rot90_test.py @@ -1,6 +1,5 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. import numpy as np -from sympy.core.numbers import comp import dace from common import compare_numpy_output diff --git a/tests/numpy/sequential_conditionals_test.py b/tests/numpy/sequential_conditionals_test.py index 7a50bbe9de..678271bf91 100644 --- a/tests/numpy/sequential_conditionals_test.py +++ b/tests/numpy/sequential_conditionals_test.py @@ -1,6 +1,5 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. import dace -import numpy as np # Declaration of symbolic variables N, BS = (dace.symbol(name) for name in ['N', 'BS']) diff --git a/tests/numpy/sse_test.py b/tests/numpy/sse_test.py index 280b190d53..ac2e459e69 100644 --- a/tests/numpy/sse_test.py +++ b/tests/numpy/sse_test.py @@ -1,6 +1,5 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. import dace -import numpy as np # Declaration of symbolic variables Nkz, NE, Nqz, Nw, N3D, NA, NB, Norb = (dace.symbol(name) diff --git a/tests/numpy/ufunc_nested_call_test.py b/tests/numpy/ufunc_nested_call_test.py index 2710071563..b85709d25b 100644 --- a/tests/numpy/ufunc_nested_call_test.py +++ b/tests/numpy/ufunc_nested_call_test.py @@ -1,7 +1,6 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. import dace import numpy as np -import pytest from common import compare_numpy_output diff --git a/tests/numpy/ufunc_replicated_input.py b/tests/numpy/ufunc_replicated_input.py index 276e751066..c3a9194787 100644 --- a/tests/numpy/ufunc_replicated_input.py +++ b/tests/numpy/ufunc_replicated_input.py @@ -1,8 +1,6 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. import dace -import math import numpy as np -import pytest from common import compare_numpy_output diff --git a/tests/numpy/ufunc_test.py b/tests/numpy/ufunc_test.py index aeb4e9ca3f..5e80d6ddac 100644 --- a/tests/numpy/ufunc_test.py +++ b/tests/numpy/ufunc_test.py @@ -1,8 +1,6 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. import dace -import math import numpy as np -import pytest from common import compare_numpy_output diff --git a/tests/numpy_bool_input_test.py b/tests/numpy_bool_input_test.py index 8140512d28..047e17d32c 100644 --- a/tests/numpy_bool_input_test.py +++ b/tests/numpy_bool_input_test.py @@ -1,5 +1,3 @@ -import numpy as np - import dace diff --git a/tests/octave/octave_test.py b/tests/octave/octave_test.py index b39301725d..56f4acfbad 100644 --- a/tests/octave/octave_test.py +++ b/tests/octave/octave_test.py @@ -1,4 +1,3 @@ -import dace from dace.cli import dacelab import os import pytest diff --git a/tests/offset_stride_test.py b/tests/offset_stride_test.py index b78ea6c43f..ca895945ca 100644 --- a/tests/offset_stride_test.py +++ b/tests/offset_stride_test.py @@ -7,7 +7,7 @@ def test(): - print('Multidimensional offset and stride test') + """Multidimensional offset and stride test""" # Externals (parameters, symbols) N = dp.symbol('N') n = 20 @@ -18,9 +18,12 @@ def test(): # Construct SDFG mysdfg = SDFG('offset_stride') + mysdfg.add_array('A', [6, 6], dp.float32, offset=[2, 3], strides=[N, 1], total_size=N * N) + mysdfg.add_array('B', [3, 2], dp.float32, offset=[-1, -1], strides=[3, 1], total_size=12) + state = mysdfg.add_state() - A_ = state.add_array('A', [6, 6], dp.float32, offset=[2, 3], strides=[N, 1], total_size=N * N) - B_ = state.add_array('B', [3, 2], dp.float32, offset=[-1, -1], strides=[3, 1], total_size=12) + A_ = state.add_access('A') + B_ = state.add_access('B') map_entry, map_exit = state.add_map('mymap', [('i', '1:4'), ('j', '1:3')]) tasklet = state.add_tasklet('mytasklet', {'a'}, {'b'}, 'b = a') @@ -34,7 +37,6 @@ def test(): mysdfg(A=input, B=output, N=n) diff = np.linalg.norm(output[0:3, 0:2] - input[3:6, 4:6]) / n - print("Difference:", diff) assert diff <= 1e-5 diff --git a/tests/parse_state_struct_test.py b/tests/parse_state_struct_test.py index f8553249ea..5cb248aa76 100644 --- a/tests/parse_state_struct_test.py +++ b/tests/parse_state_struct_test.py @@ -40,7 +40,7 @@ def _cuda_helper(): compiler.generate_program_folder(None, [program, dummy_cuda_target], BUILD_PATH) compiler.configure_and_compile(BUILD_PATH) - checker_dll = compiled_sdfg.ReloadableDLL(compiler.get_binary_name(BUILD_PATH, "cuda_helper"), "cuda_helper") + checker_dll = compiled_sdfg.ReloadableDLL(compiler.get_binary_name(BUILD_PATH, "cuda_helper")) class CudaHelper: diff --git a/tests/passes/access_ranges_test.py b/tests/passes/access_ranges_test.py index 3bab2e9ab0..220299e9b5 100644 --- a/tests/passes/access_ranges_test.py +++ b/tests/passes/access_ranges_test.py @@ -2,7 +2,6 @@ """ Tests the AccessRanges analysis pass. """ import dace from dace.transformation.passes.analysis import AccessRanges -import numpy as np N = dace.symbol('N') diff --git a/tests/passes/array_elimination_test.py b/tests/passes/array_elimination_test.py index f20428fda3..1b04a4754a 100644 --- a/tests/passes/array_elimination_test.py +++ b/tests/passes/array_elimination_test.py @@ -1,7 +1,5 @@ # Copyright 2019-2022 ETH Zurich and the DaCe authors. All rights reserved. -import pytest - import dace from dace.sdfg import utils as sdutil from dace.transformation.pass_pipeline import Pipeline diff --git a/tests/passes/empty_loop_elimination_test.py b/tests/passes/empty_loop_elimination_test.py index 2a3c3be5f6..881dabe076 100644 --- a/tests/passes/empty_loop_elimination_test.py +++ b/tests/passes/empty_loop_elimination_test.py @@ -1,6 +1,4 @@ # Copyright 2019-2025 ETH Zurich and the DaCe authors. All rights reserved. -import numpy as np -import pytest import dace from dace.transformation.passes.empty_loop_elimination import EmptyLoopElimination from dace.sdfg.state import LoopRegion diff --git a/tests/passes/lift_trivial_if_test.py b/tests/passes/lift_trivial_if_test.py new file mode 100644 index 0000000000..526ea979fd --- /dev/null +++ b/tests/passes/lift_trivial_if_test.py @@ -0,0 +1,286 @@ +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. +"""Tests for the ``LiftTrivialIf`` simplification pass.""" +import dace +from dace import InterstateEdge +from dace.sdfg.sdfg import CodeBlock, ConditionalBlock +from dace import ControlFlowRegion +from dace.transformation.passes.lift_trivial_if import LiftTrivialIf +import pytest + +# Conditions the pass must recognize as constant ``True``. +_ALWAYS_TRUE = [ + "True", + "1 == 1", + "2 > 1", + "5 >= 5", + "not False", + "2 + 2 == 4", + "abs(-5) == 5", + "max(1, 2, 3) == 3", +] + +# Conditions the pass must recognize as constant ``False``. +_ALWAYS_FALSE = [ + "False", + "1 == 2", + "5 < 3", + "10 <= 9", + "not True", + "2 + 2 == 5", + "abs(-5) == -5", +] + +# Conditions that reference unbound symbols -- the pass must leave these alone. +_CANT_EVAL = ["a < 5", "c == 0", "d >= 1"] + +# Conditions where pystr_to_symbolic yields an unevaluated symbolic expression +# (Function/Indexed/Symbol). bool() of such an expression is truthy, which would +# mis-classify dynamic dataflow conditions like ``A[0]`` as trivially true. +_DYNAMIC_RUNTIME_COND = ["A[0]", "tmp_r[0]", "x", "x[0] + 1", "A[i, j]"] + + +def _get_sdfg(condition: str): + """Build a one-state SDFG inside a single-branch ``ConditionalBlock``. + + :param condition: Python condition expression for the branch. + :returns: The SDFG with the conditional as its start block. + """ + sdfg = dace.SDFG("basic1") + _, A = sdfg.add_array(name="A", shape=[ + 5, + ], dtype=dace.float64, transient=False) + _, B = sdfg.add_array(name="B", shape=[ + 5, + ], dtype=dace.float64, transient=False) + cb = ConditionalBlock(label="cfb1", sdfg=sdfg, parent=sdfg) + sdfg.add_node(cb, is_start_block=True) + cfg = ControlFlowRegion(label="cfg1", sdfg=cb.sdfg, parent=cb) + cb.add_branch(condition=CodeBlock(condition), branch=cfg) + s1 = cfg.add_state(label="s1") + aA = s1.add_access("A") + aB = s1.add_access("B") + s1.add_edge(aA, None, aB, None, dace.memlet.Memlet.from_array("A", A)) + + return sdfg + + +def _get_if_else_sdfg(condition: str, body_in_else_branch: bool): + """Build an if/else SDFG with the body on the requested branch. + + :param condition: Python condition expression for the ``if`` arm. + :param body_in_else_branch: If ``True`` put the access-node body on the ``else`` + branch (otherwise on the ``if`` arm). + :returns: The SDFG with the conditional as its start block. + """ + sdfg = dace.SDFG("if_else_1") + _, A = sdfg.add_array(name="A", shape=[ + 5, + ], dtype=dace.float64, transient=False) + _, B = sdfg.add_array(name="B", shape=[ + 5, + ], dtype=dace.float64, transient=False) + cb = ConditionalBlock(label="cfb1", sdfg=sdfg, parent=sdfg) + sdfg.add_node(cb, is_start_block=True) + cfg_if_true = ControlFlowRegion(label="cfg_if_true", sdfg=cb.sdfg, parent=cb) + cfg_if_false = ControlFlowRegion(label="cfg_if_false", sdfg=cb.sdfg, parent=cb) + cb.add_branch(condition=CodeBlock(condition), branch=cfg_if_true) + cb.add_branch(condition=None, branch=cfg_if_false) + if body_in_else_branch: + s1 = cfg_if_true.add_state(label="s1", is_start_block=True) + s2 = cfg_if_false.add_state(label="fs1", is_start_block=True) + else: + s1 = cfg_if_false.add_state(label="s1", is_start_block=True) + s2 = cfg_if_true.add_state(label="ts1", is_start_block=True) + aA = s1.add_access("A") + aB = s1.add_access("B") + s1.add_edge(aA, None, aB, None, dace.memlet.Memlet.from_array("A", A)) + + return sdfg + + +def _get_nested_sdfg(condition1: str, condition2: str): + """Build a conditional nested inside another conditional. + + :param condition1: Outer condition. + :param condition2: Inner condition. + :returns: The SDFG with the outer conditional as its start block. + """ + sdfg = dace.SDFG("nested1") + sdfg = dace.SDFG("basic1") + _, A = sdfg.add_array(name="A", shape=[ + 5, + ], dtype=dace.float64, transient=False) + _, B = sdfg.add_array(name="B", shape=[ + 5, + ], dtype=dace.float64, transient=False) + cb = ConditionalBlock(label="cfb1", sdfg=sdfg, parent=sdfg) + sdfg.add_node(cb, is_start_block=True) + cfg1 = ControlFlowRegion(label="cfg1", sdfg=cb.sdfg, parent=cb) + cb.add_branch(condition=CodeBlock(condition1), branch=cfg1) + cb2 = ConditionalBlock(label="cfb2", sdfg=cfg1.sdfg, parent=cfg1) + cfg1.add_node(cb2, is_start_block=True) + cfg2 = ControlFlowRegion(label="cfg2", sdfg=cb2.sdfg, parent=cb2) + cb2.add_branch(condition=CodeBlock(condition2), branch=cfg2) + s1 = cfg2.add_state(label="s1", is_start_block=True) + aA = s1.add_access("A") + aB = s1.add_access("B") + s1.add_edge(aA, None, aB, None, dace.memlet.Memlet.from_array("A", A)) + + return sdfg + + +def _get_sdfg_with_many_states(): + """Build an SDFG with non-conditional states surrounding a nested conditional. + + :returns: An SDFG where the conditional sits between ``so1`` and ``so2``. + """ + sdfg = dace.SDFG("nested1") + sdfg = dace.SDFG("basic1") + _, A = sdfg.add_array(name="A", shape=[ + 5, + ], dtype=dace.float64, transient=False) + _, B = sdfg.add_array(name="B", shape=[ + 5, + ], dtype=dace.float64, transient=False) + cb = ConditionalBlock(label="cfb1", sdfg=sdfg, parent=sdfg) + so1 = sdfg.add_state("so1", is_start_block=True) + so2 = sdfg.add_state("so2") + sdfg.add_node(cb, is_start_block=False) + sdfg.add_edge(so1, cb, InterstateEdge()) + sdfg.add_edge(cb, so2, InterstateEdge()) + cfg1 = ControlFlowRegion(label="cfg1", sdfg=cb.sdfg, parent=cb) + cb.add_branch(condition=CodeBlock("1 == 1"), branch=cfg1) + cb2 = ConditionalBlock(label="cfb2", sdfg=cfg1.sdfg, parent=cfg1) + cfg1.add_node(cb2, is_start_block=True) + cfg2 = ControlFlowRegion(label="cfg2", sdfg=cb2.sdfg, parent=cb2) + cb2.add_branch(condition=CodeBlock("0 == 1"), branch=cfg2) + s1 = cfg2.add_state(label="s1", is_start_block=True) + aA = s1.add_access("A") + aB = s1.add_access("B") + s1.add_edge(aA, None, aB, None, dace.memlet.Memlet.from_array("A", A)) + return sdfg + + +@pytest.mark.parametrize("condition", _ALWAYS_TRUE) +def test_single_condition(condition: str): + """Trivially-true single-branch ``if`` is removed.""" + sdfg = _get_sdfg(condition) + sdfg.validate() + LiftTrivialIf().apply_pass(sdfg, {}) + sdfg.validate() + assert len({n for n in sdfg.all_control_flow_blocks() if isinstance(n, ConditionalBlock)}) == 0 + + +@pytest.mark.parametrize("condition", _CANT_EVAL) +def test_single_condition_cant_eval(condition: str): + """Free-symbol condition is left in place -- the pass can't prove it constant.""" + sdfg = _get_sdfg(condition) + sdfg.validate() + LiftTrivialIf().apply_pass(sdfg, {}) + sdfg.validate() + assert len({n for n in sdfg.all_control_flow_blocks() if isinstance(n, ConditionalBlock)}) == 1 + + +@pytest.mark.parametrize("condition1,condition2", + [(_ALWAYS_TRUE[i], _ALWAYS_TRUE[i + 1]) for i in range(len(_ALWAYS_TRUE) - 1)]) +def test_nested_condition(condition1: str, condition2: str): + """Trivially-true ``if`` inside trivially-true ``if`` collapses to plain dataflow.""" + sdfg = _get_nested_sdfg(condition1, condition2) + sdfg.validate() + LiftTrivialIf().apply_pass(sdfg, {}) + sdfg.validate() + assert len({n for n in sdfg.all_control_flow_blocks() if isinstance(n, ConditionalBlock)}) == 0 + + +@pytest.mark.parametrize("condition1,condition2", + [(_CANT_EVAL[i], _CANT_EVAL[i + 1]) for i in range(len(_CANT_EVAL) - 1)]) +def test_nested_condition_cant_eval(condition1: str, condition2: str): + """Nested unresolvable conditions are both preserved.""" + sdfg = _get_nested_sdfg(condition1, condition2) + sdfg.validate() + LiftTrivialIf().apply_pass(sdfg, {}) + sdfg.validate() + assert len({n for n in sdfg.all_control_flow_blocks() if isinstance(n, ConditionalBlock)}) == 2 + + +@pytest.mark.parametrize("condition", _ALWAYS_TRUE) +def test_if_else_cond_is_trivially_true(condition: str): + """``if/else`` with provably-true ``if`` keeps the ``if`` body.""" + sdfg = _get_if_else_sdfg(condition, False) + sdfg.validate() + LiftTrivialIf().apply_pass(sdfg, {}) + sdfg.validate() + assert len({n for n in sdfg.all_control_flow_blocks() if isinstance(n, ConditionalBlock)}) == 0 + + +@pytest.mark.parametrize("condition", _ALWAYS_FALSE) +def test_if_else_cond_is_trivially_false(condition: str): + """``if/else`` with provably-false ``if`` keeps the ``else`` body.""" + sdfg = _get_if_else_sdfg(condition, True) + sdfg.validate() + LiftTrivialIf().apply_pass(sdfg, {}) + sdfg.validate() + assert len({n for n in sdfg.all_control_flow_blocks() if isinstance(n, ConditionalBlock)}) == 0 + + +def test_cfg_is_a_middle_node(): + """The conditional being mid-graph (with predecessor + successor) doesn't break the splice.""" + sdfg = _get_sdfg_with_many_states() + sdfg.validate() + LiftTrivialIf().apply_pass(sdfg, {}) + sdfg.validate() + assert len({n for n in sdfg.all_control_flow_blocks() if isinstance(n, ConditionalBlock)}) == 0 + + +# Fortran frontends emit conditions of the form ``(x == y) == 0/1`` (a +# comparison of a comparison to an integer). Sympy fails to evaluate this due to mismatching types. +# If these ever stop folding to True/False the pass would silently miss real trivial-if cases in Fortran SDFGs. +@pytest.mark.parametrize("condition", ["(1 == 1) == 1", "(2 == 2) == 1", "(0 == 1) == 0"]) +def test_fortran_style_nested_comparison_true(condition: str): + """The boolean-arithmetic fallback handles the Fortran-frontend ``(x == y) == k`` shape.""" + sdfg = _get_sdfg(condition) + sdfg.validate() + LiftTrivialIf().apply_pass(sdfg, {}) + sdfg.validate() + assert len({n for n in sdfg.all_control_flow_blocks() if isinstance(n, ConditionalBlock)}) == 0 + + +@pytest.mark.parametrize("condition", _DYNAMIC_RUNTIME_COND) +def test_dynamic_runtime_cond_not_trivial(condition: str): + """Dynamic (runtime-dataflow) conditions must NOT be treated as trivially true/false.""" + sdfg = _get_sdfg(condition) + sdfg.validate() + LiftTrivialIf().apply_pass(sdfg, {}) + sdfg.validate() + # The conditional must survive: it depends on data we don't know at + # transform time. + assert len({n for n in sdfg.all_control_flow_blocks() if isinstance(n, ConditionalBlock)}) == 1 + + +def test_simplify_pipeline_includes_lift_trivial_if(): + """``SimplifyPass`` wires ``LiftTrivialIf`` in -- running simplify removes trivial ifs.""" + from dace.transformation.passes.simplify import SimplifyPass + sdfg = _get_sdfg("True") + sdfg.validate() + SimplifyPass().apply_pass(sdfg, {}) + sdfg.validate() + assert len({n for n in sdfg.all_control_flow_blocks() if isinstance(n, ConditionalBlock)}) == 0 + + +if __name__ == "__main__": + for c in _ALWAYS_TRUE: + test_single_condition(c) + + for i in range(len(_ALWAYS_TRUE) - 1): + c1 = _ALWAYS_TRUE[i] + c2 = _ALWAYS_TRUE[i + 1] + test_nested_condition(c1, c2) + + for c in _ALWAYS_TRUE: + test_if_else_cond_is_trivially_true(c) + + for c in _ALWAYS_FALSE: + test_if_else_cond_is_trivially_false(c) + + test_cfg_is_a_middle_node() + test_simplify_pipeline_includes_lift_trivial_if() diff --git a/tests/passes/loop_to_reduce_test.py b/tests/passes/loop_to_reduce_test.py new file mode 100644 index 0000000000..8a120c8d12 --- /dev/null +++ b/tests/passes/loop_to_reduce_test.py @@ -0,0 +1,406 @@ +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. +"""Tests for LoopToReduce.""" +import dace +from dace import memlet as mm +from dace.libraries.standard.nodes.reduce import Reduce +from dace.properties import CodeBlock +from dace.sdfg.state import ConditionalBlock, ControlFlowRegion, LoopRegion +from dace.transformation.passes.loop_to_reduce import LoopToReduce + +N = dace.symbol("N") +M = dace.symbol("M") + + +def _count_loops(sdfg: dace.SDFG) -> int: + return sum(1 for n, _ in sdfg.all_nodes_recursive() if isinstance(n, LoopRegion)) + + +def _assert_single_sum_reduce_identity_none(sdfg: dace.SDFG): + reduces = [n for n, _ in sdfg.all_nodes_recursive() if isinstance(n, Reduce)] + assert len(reduces) == 1, reduces + (red, ) = reduces + assert red.wcr == "lambda a, b: a + b" + assert red.identity is None + + +def test_sdfg_api_sum_reduction_is_lifted(): + """Hand-built ``for i in range(N): A[0] = A[0] + B[i]``.""" + sdfg = dace.SDFG("sum_reduce_api") + sdfg.add_array("A", [1], dace.float64) + sdfg.add_array("B", [N], dace.float64) + + loop = LoopRegion("loop", condition_expr="i < N", loop_var="i", initialize_expr="i = 0", update_expr="i = i + 1") + sdfg.add_node(loop, is_start_block=True) + + body = loop.add_state("body", is_start_block=True) + acc_r = body.add_read("A") + arr_r = body.add_read("B") + task = body.add_tasklet("add", {"in_a", "in_b"}, {"out"}, "out = in_a + in_b") + acc_w = body.add_write("A") + body.add_edge(acc_r, None, task, "in_a", mm.Memlet("A[0]")) + body.add_edge(arr_r, None, task, "in_b", mm.Memlet("B[i]")) + body.add_edge(task, "out", acc_w, None, mm.Memlet("A[0]")) + sdfg.validate() + assert _count_loops(sdfg) == 1 + + lifted = LoopToReduce().apply_pass(sdfg, {}) + sdfg.validate() + + assert lifted == 1 + assert _count_loops(sdfg) == 0 + _assert_single_sum_reduce_identity_none(sdfg) + + +@dace.program +def _frontend_augassign_len1(A: dace.float64[1], B: dace.float64[N]): + for i in range(N): + A[0] += B[i] + + +def test_frontend_augassign_length1_array_is_lifted(): + sdfg = _frontend_augassign_len1.to_sdfg(simplify=True) + sdfg.validate() + assert _count_loops(sdfg) >= 1 + + lifted = LoopToReduce().apply_pass(sdfg, {}) + sdfg.validate() + + assert lifted and lifted >= 1 + _assert_single_sum_reduce_identity_none(sdfg) + + +@dace.program +def _frontend_augassign_slice(C: dace.float64[M], B: dace.float64[N]): + for i in range(N): + C[3] += B[i] + + +def test_frontend_augassign_array_slice_is_lifted(): + """``C[k] += B[i]`` with ``C`` multi-element and ``k=3`` loop-invariant. + + Exercises the multi-element-accumulator path: the reduction's output + memlet must target the single slice ``C[3:4]``, not ``[0:0]``. + """ + sdfg = _frontend_augassign_slice.to_sdfg(simplify=True) + sdfg.validate() + assert _count_loops(sdfg) >= 1 + + lifted = LoopToReduce().apply_pass(sdfg, {}) + sdfg.validate() + + assert lifted and lifted >= 1 + _assert_single_sum_reduce_identity_none(sdfg) + + reduces = [(n, g) for n, g in sdfg.all_nodes_recursive() if isinstance(n, Reduce)] + (red, state) = reduces[0] + (out_edge, ) = state.out_edges(red) + assert out_edge.data.data == "C" + assert str(out_edge.data.subset) in {"3", "3:4", "3:3"} + + +def _build_interstate_reduction_sdfg(offset_expr: str): + """Loop body = 2 empty states + interstate edge ``{accum: accum + B[]}``.""" + sdfg = dace.SDFG(f"interstate_sum_{offset_expr.replace(' ', '').replace('-', 'm')}") + sdfg.add_symbol("accum", dace.float64) + sdfg.add_array("B", [N], dace.float64) + + pre = sdfg.add_state("pre", is_start_block=True) + loop = LoopRegion("loop", condition_expr="i < N", loop_var="i", initialize_expr="i = 0", update_expr="i = i + 1") + sdfg.add_node(loop) + sdfg.add_edge(pre, loop, dace.InterstateEdge()) + s1 = loop.add_state("s1", is_start_block=True) + s2 = loop.add_state("s2") + loop.add_edge(s1, s2, dace.InterstateEdge(assignments={"accum": f"accum + B[{offset_expr}]"})) + post = sdfg.add_state("post") + sdfg.add_edge(loop, post, dace.InterstateEdge()) + return sdfg + + +def test_interstate_edge_direct_index_is_lifted(): + sdfg = _build_interstate_reduction_sdfg("i") + sdfg.validate() + assert _count_loops(sdfg) == 1 + + lifted = LoopToReduce().apply_pass(sdfg, {}) + sdfg.validate() + + assert lifted == 1 + assert _count_loops(sdfg) == 0 + reduces = [n for n, _ in sdfg.all_nodes_recursive() if isinstance(n, Reduce)] + assert len(reduces) == 1 + (red, ) = reduces + assert red.wcr == "lambda a, b: a + b" + assert red.identity is None + + +def _build_conditional_minmax_sdfg(cond_expr: str): + """body = ConditionalBlock(cond_expr) > branch of (2 empty states + {accum: B[i]}).""" + safe = "".join(c if c.isalnum() else "_" for c in cond_expr) + sdfg = dace.SDFG(f"interstate_cond_{safe}") + sdfg.add_symbol("accum", dace.float64) + sdfg.add_array("B", [N], dace.float64) + + pre = sdfg.add_state("pre", is_start_block=True) + loop = LoopRegion("loop", condition_expr="i < N", loop_var="i", initialize_expr="i = 0", update_expr="i = i + 1") + sdfg.add_node(loop) + sdfg.add_edge(pre, loop, dace.InterstateEdge()) + + cb = ConditionalBlock("cb") + loop.add_node(cb, is_start_block=True) + branch = ControlFlowRegion("branch", sdfg=sdfg) + cb.add_branch(CodeBlock(cond_expr), branch) + s1 = branch.add_state("s1", is_start_block=True) + s2 = branch.add_state("s2") + branch.add_edge(s1, s2, dace.InterstateEdge(assignments={"accum": "B[i]"})) + + post = sdfg.add_state("post") + sdfg.add_edge(loop, post, dace.InterstateEdge()) + return sdfg + + +def _assert_single_reduce_with_wcr(sdfg: dace.SDFG, expected_wcr: str): + reduces = [n for n, _ in sdfg.all_nodes_recursive() if isinstance(n, Reduce)] + assert len(reduces) == 1, reduces + (red, ) = reduces + assert red.wcr == expected_wcr, red.wcr + assert red.identity is None + + +def test_conditional_interstate_gt_lifts_to_max(): + for cond in ("B[i] > accum", "accum < B[i]", "B[i] >= accum", "accum <= B[i]"): + sdfg = _build_conditional_minmax_sdfg(cond) + sdfg.validate() + assert _count_loops(sdfg) == 1 + lifted = LoopToReduce().apply_pass(sdfg, {}) + sdfg.validate() + assert lifted == 1, cond + assert _count_loops(sdfg) == 0 + _assert_single_reduce_with_wcr(sdfg, "lambda a, b: max(a, b)") + + +def test_conditional_interstate_lt_lifts_to_min(): + for cond in ("B[i] < accum", "accum > B[i]", "B[i] <= accum", "accum >= B[i]"): + sdfg = _build_conditional_minmax_sdfg(cond) + sdfg.validate() + assert _count_loops(sdfg) == 1 + lifted = LoopToReduce().apply_pass(sdfg, {}) + sdfg.validate() + assert lifted == 1, cond + assert _count_loops(sdfg) == 0 + _assert_single_reduce_with_wcr(sdfg, "lambda a, b: min(a, b)") + + +def test_conditional_interstate_unrelated_array_is_not_lifted(): + """Guard compares a different array than the assignment — reject.""" + sdfg = dace.SDFG("interstate_cond_mismatched") + sdfg.add_symbol("accum", dace.float64) + sdfg.add_array("B", [N], dace.float64) + sdfg.add_array("C", [N], dace.float64) + pre = sdfg.add_state("pre", is_start_block=True) + loop = LoopRegion("loop", condition_expr="i < N", loop_var="i", initialize_expr="i = 0", update_expr="i = i + 1") + sdfg.add_node(loop) + sdfg.add_edge(pre, loop, dace.InterstateEdge()) + cb = ConditionalBlock("cb") + loop.add_node(cb, is_start_block=True) + branch = ControlFlowRegion("branch", sdfg=sdfg) + cb.add_branch(CodeBlock("C[i] > accum"), branch) + s1 = branch.add_state("s1", is_start_block=True) + s2 = branch.add_state("s2") + branch.add_edge(s1, s2, dace.InterstateEdge(assignments={"accum": "B[i]"})) + post = sdfg.add_state("post") + sdfg.add_edge(loop, post, dace.InterstateEdge()) + sdfg.validate() + + assert LoopToReduce().apply_pass(sdfg, {}) is None + assert _count_loops(sdfg) == 1 + + +def _build_interstate_binop_sdfg(rhs_expr: str, accum_dtype=dace.int32, arr_dtype=dace.int32): + """Loop body = 2 empty states + interstate edge ``{accum: rhs_expr}``.""" + safe = "".join(c if c.isalnum() else "_" for c in rhs_expr) + sdfg = dace.SDFG(f"interstate_binop_{safe}") + sdfg.add_symbol("accum", accum_dtype) + sdfg.add_array("B", [N], arr_dtype) + pre = sdfg.add_state("pre", is_start_block=True) + loop = LoopRegion("loop", condition_expr="i < N", loop_var="i", initialize_expr="i = 0", update_expr="i = i + 1") + sdfg.add_node(loop) + sdfg.add_edge(pre, loop, dace.InterstateEdge()) + s1 = loop.add_state("s1", is_start_block=True) + s2 = loop.add_state("s2") + loop.add_edge(s1, s2, dace.InterstateEdge(assignments={"accum": rhs_expr})) + post = sdfg.add_state("post") + sdfg.add_edge(loop, post, dace.InterstateEdge()) + return sdfg + + +def test_interstate_edge_bitwise_or_is_lifted(): + """``{accum: accum | B[i]}`` -> Reduce with ``|`` WCR.""" + sdfg = _build_interstate_binop_sdfg("accum | B[i]") + sdfg.validate() + assert _count_loops(sdfg) == 1 + lifted = LoopToReduce().apply_pass(sdfg, {}) + sdfg.validate() + assert lifted == 1 + _assert_single_reduce_with_wcr(sdfg, "lambda a, b: a | b") + + +def test_interstate_edge_bitwise_and_is_lifted(): + """``{accum: accum & B[i]}`` -> Reduce with ``&`` WCR.""" + sdfg = _build_interstate_binop_sdfg("accum & B[i]") + sdfg.validate() + assert _count_loops(sdfg) == 1 + lifted = LoopToReduce().apply_pass(sdfg, {}) + sdfg.validate() + assert lifted == 1 + _assert_single_reduce_with_wcr(sdfg, "lambda a, b: a & b") + + +def test_interstate_edge_logical_or_is_lifted(): + """``{accum: accum or B[i]}`` -> Reduce with ``|`` WCR.""" + sdfg = _build_interstate_binop_sdfg("accum or B[i]") + sdfg.validate() + assert _count_loops(sdfg) == 1 + lifted = LoopToReduce().apply_pass(sdfg, {}) + sdfg.validate() + assert lifted == 1 + _assert_single_reduce_with_wcr(sdfg, "lambda a, b: a | b") + + +def test_interstate_edge_logical_and_is_lifted(): + """``{accum: accum and B[i]}`` -> Reduce with ``&`` WCR.""" + sdfg = _build_interstate_binop_sdfg("accum and B[i]") + sdfg.validate() + assert _count_loops(sdfg) == 1 + lifted = LoopToReduce().apply_pass(sdfg, {}) + sdfg.validate() + assert lifted == 1 + _assert_single_reduce_with_wcr(sdfg, "lambda a, b: a & b") + + +def test_tasklet_body_boolop_or_is_lifted(): + """Tasklet body ``out = a or b`` lifts to an OR reduction.""" + sdfg = dace.SDFG("tasklet_boolop_or") + sdfg.add_array("A", [1], dace.int32) + sdfg.add_array("B", [N], dace.int32) + loop = LoopRegion("loop", condition_expr="i < N", loop_var="i", initialize_expr="i = 0", update_expr="i = i + 1") + sdfg.add_node(loop, is_start_block=True) + state = loop.add_state("body", is_start_block=True) + r_a = state.add_read("A") + r_b = state.add_read("B") + w_a = state.add_write("A") + t = state.add_tasklet("or_t", {"a", "b"}, {"o"}, "o = a or b") + state.add_edge(r_a, None, t, "a", mm.Memlet("A[0]")) + state.add_edge(r_b, None, t, "b", mm.Memlet("B[i]")) + state.add_edge(t, "o", w_a, None, mm.Memlet("A[0]")) + sdfg.validate() + assert _count_loops(sdfg) == 1 + lifted = LoopToReduce().apply_pass(sdfg, {}) + sdfg.validate() + assert lifted == 1 + _assert_single_reduce_with_wcr(sdfg, "lambda a, b: a | b") + + +def _build_any_pattern_sdfg(assign_const: str = "1", guard: str = "B[i] == 1"): + """Mirrors FOR_l_600_c_600 (tmp_call_13): a LoopRegion whose body is a + ConditionalBlock on an array element, with a constant-RHS assignment to + a symbol.""" + sdfg = dace.SDFG(f"any_pattern_{assign_const}") + sdfg.add_symbol("tmp_call_13", dace.int32) + sdfg.add_array("B", [N], dace.int32) + pre = sdfg.add_state("pre", is_start_block=True) + loop = LoopRegion("loop", condition_expr="i < N", loop_var="i", initialize_expr="i = 0", update_expr="i = i + 1") + sdfg.add_node(loop) + sdfg.add_edge(pre, loop, dace.InterstateEdge()) + cb = ConditionalBlock("cb") + loop.add_node(cb, is_start_block=True) + branch = ControlFlowRegion("branch", sdfg=sdfg) + cb.add_branch(CodeBlock(guard), branch) + s1 = branch.add_state("s1", is_start_block=True) + s2 = branch.add_state("s2") + branch.add_edge(s1, s2, dace.InterstateEdge(assignments={"tmp_call_13": assign_const})) + post = sdfg.add_state("post") + sdfg.add_edge(loop, post, dace.InterstateEdge()) + return sdfg + + +def test_any_pattern_not_lifted_without_permissive(): + """Default (non-permissive) mode leaves the ``any`` pattern alone because + the lift assumes the guard array is 0/1-valued.""" + sdfg = _build_any_pattern_sdfg() + sdfg.validate() + assert _count_loops(sdfg) == 1 + assert LoopToReduce().apply_pass(sdfg, {}) is None + assert _count_loops(sdfg) == 1 + + +def test_any_pattern_lifts_to_or_in_permissive(): + """``{sym: "1"}`` gated by ``arr[f(i)] == 1`` lifts to a bitwise-OR + reduction over ``arr`` in permissive mode.""" + sdfg = _build_any_pattern_sdfg(assign_const="1", guard="B[i] == 1") + sdfg.validate() + lifted = LoopToReduce(permissive=True).apply_pass(sdfg, {}) + sdfg.validate() + assert lifted == 1 + assert _count_loops(sdfg) == 0 + _assert_single_reduce_with_wcr(sdfg, "lambda a, b: a | b") + + +def test_all_pattern_lifts_to_and_in_permissive(): + """``{sym: "0"}`` gated by ``arr[f(i)] == 0`` lifts to a bitwise-AND + reduction in permissive mode.""" + sdfg = _build_any_pattern_sdfg(assign_const="0", guard="B[i] == 0") + sdfg.validate() + lifted = LoopToReduce(permissive=True).apply_pass(sdfg, {}) + sdfg.validate() + assert lifted == 1 + assert _count_loops(sdfg) == 0 + _assert_single_reduce_with_wcr(sdfg, "lambda a, b: a & b") + + +def test_any_pattern_symbol_bridge_via_tmp_scalar(): + """When the accumulator is a symbol, ``_lift`` introduces a transient + ``_red_tmp_`` scalar, seeds it from the symbol, and assigns the + symbol back on the outgoing interstate edge.""" + from dace.libraries.standard.nodes.reduce import Reduce as _Reduce + sdfg = _build_any_pattern_sdfg() + LoopToReduce(permissive=True).apply_pass(sdfg, {}) + sdfg.validate() + + # Bridge scalar exists and is transient. + bridge_names = [k for k in sdfg.arrays if k.startswith("_red_tmp_tmp_call_13")] + assert len(bridge_names) == 1 + assert sdfg.arrays[bridge_names[0]].transient + + # Reduce writes to the bridge scalar. + reduces = [(n, g) for n, g in sdfg.all_nodes_recursive() if isinstance(n, _Reduce)] + assert len(reduces) == 1 + red, state = reduces[0] + (out_edge, ) = state.out_edges(red) + assert out_edge.data.data == bridge_names[0] + + # Outgoing interstate edge assigns the original symbol from the bridge. + for e in sdfg.all_interstate_edges(): + if "tmp_call_13" in (e.data.assignments or {}): + assert e.data.assignments["tmp_call_13"] == bridge_names[0] + break + else: + raise AssertionError("no interstate edge assigning tmp_call_13 from the bridge") + + +if __name__ == "__main__": + test_sdfg_api_sum_reduction_is_lifted() + test_frontend_augassign_length1_array_is_lifted() + test_frontend_augassign_array_slice_is_lifted() + test_interstate_edge_direct_index_is_lifted() + test_conditional_interstate_gt_lifts_to_max() + test_conditional_interstate_lt_lifts_to_min() + test_conditional_interstate_unrelated_array_is_not_lifted() + test_interstate_edge_bitwise_or_is_lifted() + test_interstate_edge_bitwise_and_is_lifted() + test_interstate_edge_logical_or_is_lifted() + test_interstate_edge_logical_and_is_lifted() + test_tasklet_body_boolop_or_is_lifted() + test_any_pattern_not_lifted_without_permissive() + test_any_pattern_lifts_to_or_in_permissive() + test_all_pattern_lifts_to_and_in_permissive() + test_any_pattern_symbol_bridge_via_tmp_scalar() diff --git a/tests/passes/pipeline_test.py b/tests/passes/pipeline_test.py index 05ec48d845..68f367f3cd 100644 --- a/tests/passes/pipeline_test.py +++ b/tests/passes/pipeline_test.py @@ -42,7 +42,7 @@ def test_pipeline_with_dependencies(): class PassA(MyPass): def depends_on(self): - return {MyPass} + return [MyPass] def apply_pass(self, sdfg, pipeline_results): res = super().apply_pass(sdfg, pipeline_results) @@ -70,7 +70,7 @@ def modifies(self) -> ppl.Modifies: class PassA(MyPass): def depends_on(self): - return {MyAnalysis} + return [MyAnalysis] def modifies(self) -> ppl.Modifies: return ppl.Modifies.Descriptors @@ -78,7 +78,7 @@ def modifies(self) -> ppl.Modifies: class PassB(MyPass): def depends_on(self): - return {MyAnalysis} + return [MyAnalysis] def modifies(self) -> ppl.Modifies: return ppl.Modifies.Symbols @@ -86,7 +86,7 @@ def modifies(self) -> ppl.Modifies: class PassC(MyPass): def depends_on(self): - return {MyAnalysis} + return [MyAnalysis] def modifies(self) -> ppl.Modifies: return ppl.Modifies.Everything @@ -104,7 +104,48 @@ def modifies(self) -> ppl.Modifies: assert result == {'MyAnalysis': 1, 'PassA': 1, 'PassB': 1, 'PassC': 1} +def test_pipeline_dependency_order(): + # Regression test for deterministic dependency ordering: when a pass returns multiple + # independent dependencies, they must be applied in the exact order they are listed. + order = [] + + class RecordingPass(ppl.Pass): + + def modifies(self) -> ppl.Modifies: + return ppl.Modifies.Nothing + + def should_reapply(self, _) -> bool: + return False + + def apply_pass(self, sdfg, _): + order.append(type(self).__name__) + return 1 + + class DepA(RecordingPass): + pass + + class DepB(RecordingPass): + pass + + class DepC(RecordingPass): + pass + + class Dependent(RecordingPass): + + def depends_on(self): + return [DepB, DepC, DepA] + + pipe = ppl.Pipeline([Dependent()]) + sdfg = empty.to_sdfg() + + pipe.apply_pass(sdfg, {}) + + # Dependencies run in the listed order, before the dependent pass itself. + assert order == ['DepB', 'DepC', 'DepA', 'Dependent'] + + if __name__ == '__main__': test_simple_pipeline() test_pipeline_with_dependencies() test_pipeline_modification_rerun() + test_pipeline_dependency_order() diff --git a/tests/passes/scalar_to_symbol_test.py b/tests/passes/scalar_to_symbol_test.py index 1565bb729c..8807fe330d 100644 --- a/tests/passes/scalar_to_symbol_test.py +++ b/tests/passes/scalar_to_symbol_test.py @@ -3,10 +3,8 @@ import dace from dace.sdfg.state import ConditionalBlock, LoopRegion from dace.transformation.passes import scalar_to_symbol -from dace.transformation import transformation as xf, interstate as isxf -from dace.transformation.interstate import loop_detection as ld +from dace.transformation import interstate as isxf -import collections from sympy import core as sympy_core import numpy as np import pytest diff --git a/tests/passes/simplification/continue_to_condition_test.py b/tests/passes/simplification/continue_to_condition_test.py index a2bfc32c94..2062715b4a 100644 --- a/tests/passes/simplification/continue_to_condition_test.py +++ b/tests/passes/simplification/continue_to_condition_test.py @@ -1,6 +1,5 @@ # Copyright 2019-2025 ETH Zurich and the DaCe authors. All rights reserved. import numpy as np -import pytest import dace from dace import nodes from dace.transformation.passes.simplification.continue_to_condition import ContinueToCondition diff --git a/tests/passes/split_tasklets_test.py b/tests/passes/split_tasklets_test.py index 922018de97..78a0a475e8 100644 --- a/tests/passes/split_tasklets_test.py +++ b/tests/passes/split_tasklets_test.py @@ -2,7 +2,6 @@ import typing import dace import re -import dace.sdfg.utils as sdutil import copy from dace.transformation.passes.split_tasklets import SplitTasklets import numpy diff --git a/tests/passes/symbol_propagation_test.py b/tests/passes/symbol_propagation_test.py index 554cea81bc..e085b20ae5 100644 --- a/tests/passes/symbol_propagation_test.py +++ b/tests/passes/symbol_propagation_test.py @@ -1,6 +1,5 @@ # Copyright 2019-2025 ETH Zurich and the DaCe authors. All rights reserved. import numpy as np -import pytest import dace from dace.sdfg import nodes diff --git a/tests/passes/symbol_ssa_test.py b/tests/passes/symbol_ssa_test.py index 5b0c7bd175..fa5b1c0c42 100644 --- a/tests/passes/symbol_ssa_test.py +++ b/tests/passes/symbol_ssa_test.py @@ -1,12 +1,9 @@ # Copyright 2019-2023 ETH Zurich and the DaCe authors. All rights reserved. """ Tests the symbol write scopes analysis pass. """ -import pytest - import dace from dace.transformation.pass_pipeline import FixedPointPipeline from dace.transformation.passes.symbol_ssa import StrictSymbolSSA -from dace.transformation.passes.analysis import SymbolScopeDict def test_loop_iter_symbol_reused_split(): diff --git a/tests/passes/writeset_underapproximation_test.py b/tests/passes/writeset_underapproximation_test.py index 19e9820def..d4bc48307c 100644 --- a/tests/passes/writeset_underapproximation_test.py +++ b/tests/passes/writeset_underapproximation_test.py @@ -605,7 +605,7 @@ def test_simple_loop_overwrite(): init = sdfg.add_state("init") end = sdfg.add_state("end") loop_body = sdfg.add_state("loop_body") - _, guard, _ = sdfg.add_loop(init, loop_body, end, "i", "0", "i < N", "i + 1") + _, guard, _ = sdfg.add_loop_state_machine(init, loop_body, end, "i", "0", "i < N", "i + 1") a0 = loop_body.add_access("A") loop_tasklet = loop_body.add_tasklet("overwrite", {}, {"a"}, "a = 0") loop_body.add_edge(loop_tasklet, "a", a0, None, dace.Memlet("A[i]")) @@ -629,8 +629,8 @@ def test_loop_2D_overwrite(): loop_body = sdfg.add_state("loop_body") loop_before_1 = sdfg.add_state("loop_before_1") loop_after_1 = sdfg.add_state("loop_after_1") - _, guard2, _ = sdfg.add_loop(loop_before_1, loop_body, loop_after_1, "i", "0", "i < N", "i + 1") - _, guard1, _ = sdfg.add_loop(init, loop_before_1, end, "j", "0", "j < M", "j + 1", loop_after_1) + _, guard2, _ = sdfg.add_loop_state_machine(loop_before_1, loop_body, loop_after_1, "i", "0", "i < N", "i + 1") + _, guard1, _ = sdfg.add_loop_state_machine(init, loop_before_1, end, "j", "0", "j < M", "j + 1", loop_after_1) a0 = loop_body.add_access("A") loop_tasklet = loop_body.add_tasklet("overwrite", {}, {"a"}, "a = 0") loop_body.add_edge(loop_tasklet, "a", a0, None, dace.Memlet("A[j,i]")) @@ -659,10 +659,12 @@ def test_loop_2D_propagation_gap_symbolic(): loop_after_1 = sdfg.add_state("loop_after_1") loop_before_2 = sdfg.add_state("loop_before_2") loop_after_2 = sdfg.add_state("loop_after_2") - _, guard3, _ = sdfg.add_loop(loop_before_1, loop_body, loop_after_1, "i", "0", "i < N", "i + 1") # inner-most loop - _, guard2, _ = sdfg.add_loop(loop_before_2, loop_before_1, loop_after_2, "k", "0", "k < K", "k + 1", - loop_after_1) # second-inner-most loop - _, guard1, _ = sdfg.add_loop(init, loop_before_2, end, "j", "0", "j < M", "j + 1", loop_after_2) # outer-most loop + _, guard3, _ = sdfg.add_loop_state_machine(loop_before_1, loop_body, loop_after_1, "i", "0", "i < N", + "i + 1") # inner-most loop + _, guard2, _ = sdfg.add_loop_state_machine(loop_before_2, loop_before_1, loop_after_2, "k", "0", "k < K", "k + 1", + loop_after_1) # second-inner-most loop + _, guard1, _ = sdfg.add_loop_state_machine(init, loop_before_2, end, "j", "0", "j < M", "j + 1", + loop_after_2) # outer-most loop a0 = loop_body.add_access("A") loop_tasklet = loop_body.add_tasklet("overwrite", {}, {"a"}, "a = 0") loop_body.add_edge(loop_tasklet, "a", a0, None, dace.Memlet("A[j,i]")) @@ -687,8 +689,8 @@ def test_2_loops_overwrite(): end = sdfg.add_state("end") loop_body_1 = sdfg.add_state("loop_body_1") loop_body_2 = sdfg.add_state("loop_body_2") - _, guard_1, after_state = sdfg.add_loop(init, loop_body_1, None, "i", "0", "i < N", "i + 1") - _, guard_2, _ = sdfg.add_loop(after_state, loop_body_2, end, "i", "0", "i < N", "i + 1") + _, guard_1, after_state = sdfg.add_loop_state_machine(init, loop_body_1, None, "i", "0", "i < N", "i + 1") + _, guard_2, _ = sdfg.add_loop_state_machine(after_state, loop_body_2, end, "i", "0", "i < N", "i + 1") a0 = loop_body_1.add_access("A") loop_tasklet_1 = loop_body_1.add_tasklet("overwrite", {}, {"a"}, "a = 0") loop_body_1.add_edge(loop_tasklet_1, "a", a0, None, dace.Memlet("A[i]")) @@ -720,9 +722,10 @@ def test_loop_2D_overwrite_propagation_gap_non_empty(): loop_after_1 = sdfg.add_state("loop_after_1") loop_before_2 = sdfg.add_state("loop_before_2") loop_after_2 = sdfg.add_state("loop_after_2") - _, guard3, _ = sdfg.add_loop(loop_before_1, loop_body, loop_after_1, "i", "0", "i < N", "i + 1") - _, guard2, _ = sdfg.add_loop(loop_before_2, loop_before_1, loop_after_2, "k", "0", "k < 10", "k + 1", loop_after_1) - _, guard1, _ = sdfg.add_loop(init, loop_before_2, end, "j", "0", "j < M", "j + 1", loop_after_2) + _, guard3, _ = sdfg.add_loop_state_machine(loop_before_1, loop_body, loop_after_1, "i", "0", "i < N", "i + 1") + _, guard2, _ = sdfg.add_loop_state_machine(loop_before_2, loop_before_1, loop_after_2, "k", "0", "k < 10", "k + 1", + loop_after_1) + _, guard1, _ = sdfg.add_loop_state_machine(init, loop_before_2, end, "j", "0", "j < M", "j + 1", loop_after_2) a0 = loop_body.add_access("A") loop_tasklet = loop_body.add_tasklet("overwrite", {}, {"a"}, "a = 0") loop_body.add_edge(loop_tasklet, "a", a0, None, dace.Memlet("A[j,i]")) @@ -751,9 +754,10 @@ def test_loop_nest_multiplied_indices(): loop_after_1 = sdfg.add_state("loop_after_1") loop_before_2 = sdfg.add_state("loop_before_2") loop_after_2 = sdfg.add_state("loop_after_2") - _, guard3, _ = sdfg.add_loop(loop_before_1, loop_body, loop_after_1, "i", "0", "i < N", "i + 1") - _, guard2, _ = sdfg.add_loop(loop_before_2, loop_before_1, loop_after_2, "k", "0", "k < 10", "k + 1", loop_after_1) - _, guard1, _ = sdfg.add_loop(init, loop_before_2, end, "j", "0", "j < M", "j + 1", loop_after_2) + _, guard3, _ = sdfg.add_loop_state_machine(loop_before_1, loop_body, loop_after_1, "i", "0", "i < N", "i + 1") + _, guard2, _ = sdfg.add_loop_state_machine(loop_before_2, loop_before_1, loop_after_2, "k", "0", "k < 10", "k + 1", + loop_after_1) + _, guard1, _ = sdfg.add_loop_state_machine(init, loop_before_2, end, "j", "0", "j < M", "j + 1", loop_after_2) a0 = loop_body.add_access("A") loop_tasklet = loop_body.add_tasklet("overwrite", {}, {"a"}, "a = 0") loop_body.add_edge(loop_tasklet, "a", a0, None, dace.Memlet("A[i,i*j]")) @@ -783,9 +787,10 @@ def test_loop_nest_empty_nested_loop(): loop_after_1 = sdfg.add_state("loop_after_1") loop_before_2 = sdfg.add_state("loop_before_2") loop_after_2 = sdfg.add_state("loop_after_2") - _, guard3, _ = sdfg.add_loop(loop_before_1, loop_body, loop_after_1, "i", "0", "i < N", "i + 1") - _, guard2, _ = sdfg.add_loop(loop_before_2, loop_before_1, loop_after_2, "k", "0", "k < 0", "k + 1", loop_after_1) - _, guard1, _ = sdfg.add_loop(init, loop_before_2, end, "j", "0", "j < M", "j + 1", loop_after_2) + _, guard3, _ = sdfg.add_loop_state_machine(loop_before_1, loop_body, loop_after_1, "i", "0", "i < N", "i + 1") + _, guard2, _ = sdfg.add_loop_state_machine(loop_before_2, loop_before_1, loop_after_2, "k", "0", "k < 0", "k + 1", + loop_after_1) + _, guard1, _ = sdfg.add_loop_state_machine(init, loop_before_2, end, "j", "0", "j < M", "j + 1", loop_after_2) a0 = loop_body.add_access("A") loop_tasklet = loop_body.add_tasklet("overwrite", {}, {"a"}, "a = 0") loop_body.add_edge(loop_tasklet, "a", a0, None, dace.Memlet("A[j,i]")) @@ -813,8 +818,8 @@ def test_loop_nest_inner_loop_conditional(): if_merge = sdfg.add_state("if_merge") loop_before_2 = sdfg.add_state("loop_before_2") loop_after_2 = sdfg.add_state("loop_after_2") - _, guard2, _ = sdfg.add_loop(loop_before_2, loop_body, loop_after_2, "k", "0", "k < N", "k + 1") - _, guard1, _ = sdfg.add_loop(init, if_guard, end, "j", "0", "j < M", "j + 1", if_merge) + _, guard2, _ = sdfg.add_loop_state_machine(loop_before_2, loop_body, loop_after_2, "k", "0", "k < N", "k + 1") + _, guard1, _ = sdfg.add_loop_state_machine(init, if_guard, end, "j", "0", "j < M", "j + 1", if_merge) sdfg.add_edge(if_guard, loop_before_2, dace.InterstateEdge(condition="j % 2 == 0")) sdfg.add_edge(if_guard, if_merge, dace.InterstateEdge(condition="j % 2 == 1")) sdfg.add_edge(loop_after_2, if_merge, dace.InterstateEdge()) @@ -909,7 +914,7 @@ def test_loop_break(): loop_body_0 = sdfg.add_state("loop_body_0") loop_body_1 = sdfg.add_state("loop_body_1") loop_after_1 = sdfg.add_state("loop_after_1") - _, guard3, _ = sdfg.add_loop(init, loop_body_0, loop_after_1, "i", "0", "i < N", "i + 1", loop_body_1) + _, guard3, _ = sdfg.add_loop_state_machine(init, loop_body_0, loop_after_1, "i", "0", "i < N", "i + 1", loop_body_1) sdfg.add_edge(loop_body_0, loop_after_1, dace.InterstateEdge(condition="i > 10")) sdfg.add_edge(loop_body_0, loop_body_1, dace.InterstateEdge(condition="not(i > 10)")) a0 = loop_body_1.add_access("A") diff --git a/tests/polybench/2mm.py b/tests/polybench/2mm.py index b9c7b81144..50cb960f11 100644 --- a/tests/polybench/2mm.py +++ b/tests/polybench/2mm.py @@ -1,5 +1,4 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. -import math import dace try: import polybench diff --git a/tests/polybench/3mm.py b/tests/polybench/3mm.py index e1e8c570f5..58073e22ca 100644 --- a/tests/polybench/3mm.py +++ b/tests/polybench/3mm.py @@ -1,5 +1,4 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. -import math import dace try: diff --git a/tests/polybench/atax.py b/tests/polybench/atax.py index a1e416360b..74ce2e80fa 100644 --- a/tests/polybench/atax.py +++ b/tests/polybench/atax.py @@ -1,5 +1,4 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. -import math import dace import polybench diff --git a/tests/polybench/bicg.py b/tests/polybench/bicg.py index 325cd25db2..463c3ddb04 100644 --- a/tests/polybench/bicg.py +++ b/tests/polybench/bicg.py @@ -1,5 +1,4 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. -import math import dace import polybench diff --git a/tests/polybench/doitgen.py b/tests/polybench/doitgen.py index 3c8bbf6a28..c7597b54ef 100644 --- a/tests/polybench/doitgen.py +++ b/tests/polybench/doitgen.py @@ -1,5 +1,4 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. -import math import dace import polybench diff --git a/tests/polybench/durbin.py b/tests/polybench/durbin.py index 8230dcfcc9..d0810bd090 100644 --- a/tests/polybench/durbin.py +++ b/tests/polybench/durbin.py @@ -1,5 +1,4 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. -import math import dace import polybench diff --git a/tests/polybench/floyd-warshall.py b/tests/polybench/floyd-warshall.py index a4546f0f3a..3cedb193e0 100644 --- a/tests/polybench/floyd-warshall.py +++ b/tests/polybench/floyd-warshall.py @@ -1,5 +1,4 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. -import math import dace try: import polybench diff --git a/tests/polybench/gemm.py b/tests/polybench/gemm.py index 3c0e2231fd..e274201df2 100644 --- a/tests/polybench/gemm.py +++ b/tests/polybench/gemm.py @@ -1,5 +1,4 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. -import math import dace try: import polybench diff --git a/tests/polybench/gemver.py b/tests/polybench/gemver.py index 591f08b1d7..998229a5a5 100644 --- a/tests/polybench/gemver.py +++ b/tests/polybench/gemver.py @@ -1,5 +1,4 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. -import math import dace import polybench diff --git a/tests/polybench/gesummv.py b/tests/polybench/gesummv.py index 13ea808360..245549daa1 100644 --- a/tests/polybench/gesummv.py +++ b/tests/polybench/gesummv.py @@ -1,5 +1,4 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. -import math import dace import polybench diff --git a/tests/polybench/gramschmidt.py b/tests/polybench/gramschmidt.py index ef293af361..f34d170adb 100644 --- a/tests/polybench/gramschmidt.py +++ b/tests/polybench/gramschmidt.py @@ -1,6 +1,5 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. import math -import numpy as np import dace import polybench diff --git a/tests/polybench/jacobi-1d.py b/tests/polybench/jacobi-1d.py index fac73173f2..8e65056d4d 100644 --- a/tests/polybench/jacobi-1d.py +++ b/tests/polybench/jacobi-1d.py @@ -1,7 +1,6 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. import dace import polybench -from absl import app, flags N = dace.symbol('N') tsteps = dace.symbol('tsteps') diff --git a/tests/polybench/lu.py b/tests/polybench/lu.py index 7696c08b56..4f3996e045 100644 --- a/tests/polybench/lu.py +++ b/tests/polybench/lu.py @@ -1,5 +1,4 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. -import math import numpy as np import dace import polybench diff --git a/tests/polybench/ludcmp.py b/tests/polybench/ludcmp.py index 651f903480..455a2eb135 100644 --- a/tests/polybench/ludcmp.py +++ b/tests/polybench/ludcmp.py @@ -1,5 +1,4 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. -import math import dace import polybench import numpy as np diff --git a/tests/polybench/mvt.py b/tests/polybench/mvt.py index 01f8ee5e22..c160b18808 100644 --- a/tests/polybench/mvt.py +++ b/tests/polybench/mvt.py @@ -1,5 +1,4 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. -import math import dace import polybench diff --git a/tests/polybench/nussinov.py b/tests/polybench/nussinov.py index 2a8d2c416a..2b82d52dbf 100644 --- a/tests/polybench/nussinov.py +++ b/tests/polybench/nussinov.py @@ -1,5 +1,4 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. -import math import dace import polybench diff --git a/tests/polybench/seidel-2d.py b/tests/polybench/seidel-2d.py index 1b14e9fe32..3c785389c9 100644 --- a/tests/polybench/seidel-2d.py +++ b/tests/polybench/seidel-2d.py @@ -1,7 +1,6 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. import dace import polybench -from absl import app, flags N = dace.symbol('N') tsteps = dace.symbol('tsteps') diff --git a/tests/polybench/symm.py b/tests/polybench/symm.py index 6c2a914c90..29385d37d7 100644 --- a/tests/polybench/symm.py +++ b/tests/polybench/symm.py @@ -1,5 +1,4 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. -import math import dace import polybench diff --git a/tests/polybench/syr2k.py b/tests/polybench/syr2k.py index 0a1e749518..1571801635 100644 --- a/tests/polybench/syr2k.py +++ b/tests/polybench/syr2k.py @@ -1,5 +1,4 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. -import math import dace import polybench diff --git a/tests/polybench/syrk.py b/tests/polybench/syrk.py index 26d28ba85e..03dee32f68 100644 --- a/tests/polybench/syrk.py +++ b/tests/polybench/syrk.py @@ -1,5 +1,4 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. -import math import dace import polybench diff --git a/tests/polybench/trisolv.py b/tests/polybench/trisolv.py index 007112cd16..931a2598e6 100644 --- a/tests/polybench/trisolv.py +++ b/tests/polybench/trisolv.py @@ -1,8 +1,6 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. -import math import dace import polybench -import numpy as np N = dace.symbol('N') diff --git a/tests/polybench/trmm.py b/tests/polybench/trmm.py index 7599668c3f..9aec061bb5 100644 --- a/tests/polybench/trmm.py +++ b/tests/polybench/trmm.py @@ -1,5 +1,4 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. -import math import dace import polybench diff --git a/tests/properties_test.py b/tests/properties_test.py index 124a557307..e987575dd8 100644 --- a/tests/properties_test.py +++ b/tests/properties_test.py @@ -4,7 +4,6 @@ import unittest import sympy as sp import dace -from collections import OrderedDict class PropertyTests(unittest.TestCase): diff --git a/tests/python_frontend/augassign_wcr_test.py b/tests/python_frontend/augassign_wcr_test.py index 1931d18577..a91277d7fe 100644 --- a/tests/python_frontend/augassign_wcr_test.py +++ b/tests/python_frontend/augassign_wcr_test.py @@ -1,7 +1,6 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. import dace import numpy as np -import pytest @dace.program diff --git a/tests/python_frontend/binop_replacements_test.py b/tests/python_frontend/binop_replacements_test.py index 40f5c6c16c..ab42c8f5d0 100644 --- a/tests/python_frontend/binop_replacements_test.py +++ b/tests/python_frontend/binop_replacements_test.py @@ -1,7 +1,6 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. import dace import numpy as np -import pytest N = dace.symbol('N') M = dace.symbol('M') diff --git a/tests/python_frontend/callee_autodetect_test.py b/tests/python_frontend/callee_autodetect_test.py index da1cc95f13..5f3211a326 100644 --- a/tests/python_frontend/callee_autodetect_test.py +++ b/tests/python_frontend/callee_autodetect_test.py @@ -211,7 +211,7 @@ def test_error_handling(): class NotConvertible(SDFGConvertible): def __call__(self, a): - import numpy as np + import numpy as np # noqa: F401 (acts as an SDFG conversion inhibitor) print('A very pythonic method', a) def __sdfg__(self, *args, **kwargs): @@ -251,7 +251,7 @@ class MaybeConvertible: @not_convertible def __call__(self, a): - import numpy as np + import numpy as np # noqa: F401 (acts as an SDFG conversion inhibitor) print('A very pythonic method', a) A = np.random.rand(20) diff --git a/tests/python_frontend/dace_program_test.py b/tests/python_frontend/dace_program_test.py index df5b80ef6f..704b8c5d6a 100644 --- a/tests/python_frontend/dace_program_test.py +++ b/tests/python_frontend/dace_program_test.py @@ -40,15 +40,15 @@ def very_unique_program_321(A: dace.float64[10]): a = np.random.rand(10) assert np.allclose(a + 1, very_unique_program_321(a)) - assert os.path.exists(os.path.join(build_folder, 'program.sdfg')) - sdfg = dace.SDFG.from_file(os.path.join(build_folder, 'program.sdfg')) + assert os.path.exists(os.path.join(build_folder, 'program.sdfgz')) + sdfg = dace.SDFG.from_file(os.path.join(build_folder, 'program.sdfgz')) # Replace the SDFG with the one we just created for node, _ in sdfg.all_nodes_recursive(): if isinstance(node, dace.nodes.Tasklet): node.code.as_string = re.sub(r'\b1\b', '2', node.code.as_string) - sdfg.save(os.path.join(build_folder, 'program.sdfg')) + sdfg.save(os.path.join(build_folder, 'program.sdfgz'), compress=True) # Now run the same program again, but this time with the SDFG already created @dace.program(recreate_sdfg=False) diff --git a/tests/python_frontend/debug_info_test.py b/tests/python_frontend/debug_info_test.py new file mode 100644 index 0000000000..0b84b606a0 --- /dev/null +++ b/tests/python_frontend/debug_info_test.py @@ -0,0 +1,55 @@ +from dace import SDFG, SDFGState +from dace import dtypes +from dace.config import Config, temporary_config + +import pytest + + +@pytest.fixture() +def state() -> SDFGState: + """Setup a simple SDFG for testing and return the state.""" + + sdfg = SDFG("tester") + sdfg.add_array("A", [10], dtypes.int64) + sdfg.add_array("B", [15], dtypes.int64) + + return sdfg.add_state("start_here", is_start_block=True) + + +def test_config_compiler_lineinfo_none(state: SDFGState) -> None: + with temporary_config(): + Config.set(*["compiler", "lineinfo"], value="none") + node = state.add_access("A") + assert node.debuginfo is None + + +def test_config_compiler_lineinfo_inspect(state: SDFGState) -> None: + # Ensure "inspect" is the default configuration + assert Config.get(*["compiler", "lineinfo"]) == "inspect" + + # Add an access and expect debug info to be added (from this file) + node = state.add_access("A") + assert node.debuginfo is not None + assert node.debuginfo.filename == __file__ + + +def test_using_default_lineinfo(state: SDFGState) -> None: + line = 42 + filename = "my/test/path.py" + state._default_lineinfo = dtypes.DebugInfo(start_line=line, filename=filename) + + read_A = state.add_read("A") + assert read_A.debuginfo is not None + assert read_A.debuginfo.filename == filename + assert read_A.debuginfo.start_line == line + + read_B = state.add_read("B", debuginfo=None) + assert read_B.debuginfo is not None + assert read_B.debuginfo.filename == filename + assert read_B.debuginfo.start_line == line + + +if __name__ == "__main__": + test_config_compiler_lineinfo_none() + test_config_compiler_lineinfo_inspect() + test_using_default_lineinfo() diff --git a/tests/python_frontend/device_annotations_test.py b/tests/python_frontend/device_annotations_test.py index 505591d9f4..21d636f338 100644 --- a/tests/python_frontend/device_annotations_test.py +++ b/tests/python_frontend/device_annotations_test.py @@ -3,7 +3,7 @@ import pytest import numpy as np -from dace.dtypes import StorageType, DeviceType, ScheduleType +from dace.dtypes import StorageType, ScheduleType from dace import dtypes try: diff --git a/tests/python_frontend/fields_and_global_arrays_test.py b/tests/python_frontend/fields_and_global_arrays_test.py index 628c478bc0..7f4b9e1ba6 100644 --- a/tests/python_frontend/fields_and_global_arrays_test.py +++ b/tests/python_frontend/fields_and_global_arrays_test.py @@ -6,7 +6,6 @@ import numpy as np from dataclasses import dataclass import pytest -import time from types import SimpleNamespace diff --git a/tests/python_frontend/identifiers_and_keywords_test.py b/tests/python_frontend/identifiers_and_keywords_test.py index b3ef2f875e..9be51a9a2c 100644 --- a/tests/python_frontend/identifiers_and_keywords_test.py +++ b/tests/python_frontend/identifiers_and_keywords_test.py @@ -71,7 +71,7 @@ def test_keyword_and(): @dace.program def keyword_assert(A: dace.float32[N], B: dace.float32[N], C: dace.bool, D: dace.bool): with C as A: - from dace import symbolic + from dace import symbolic # noqa: F401 (syntactic element that is part of the test) a = 5 del a assert (C == True) diff --git a/tests/python_frontend/indirection_in_map_test.py b/tests/python_frontend/indirection_in_map_test.py index 0377101cd1..cfa86ce944 100644 --- a/tests/python_frontend/indirection_in_map_test.py +++ b/tests/python_frontend/indirection_in_map_test.py @@ -1,7 +1,6 @@ # Copyright 2019-2023 ETH Zurich and the DaCe authors. All rights reserved. import dace import numpy as np -import scipy as sp import scipy.sparse as sparse M = dace.symbol('M') diff --git a/tests/python_frontend/loop_regions_test.py b/tests/python_frontend/loop_regions_test.py index 7d372141a5..cfc999f832 100644 --- a/tests/python_frontend/loop_regions_test.py +++ b/tests/python_frontend/loop_regions_test.py @@ -1,4 +1,7 @@ -# Copyright 2019-2024 ETH Zurich and the DaCe authors. All rights reserved. +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. +import json +import tempfile + import dace import numpy as np import pytest @@ -7,6 +10,17 @@ from dace.sdfg.state import LoopRegion +def _assert_roundtrip_json_stable(sdfg): + with tempfile.TemporaryDirectory() as tmpdir: + path1 = f'{tmpdir}/test.sdfg' + path2 = f'{tmpdir}/test2.sdfg' + sdfg.save(path1, hash=False) + dace.SDFG.from_file(path1).save(path2, hash=False) + + with open(path1, 'r') as fp1, open(path2, 'r') as fp2: + assert json.load(fp1) == json.load(fp2) + + @dace.program def for_loop(): A = dace.ndarray([10], dtype=dace.int32) @@ -104,6 +118,11 @@ def test_while_loop(): assert (np.array_equal(A, A_ref)) +def test_while_loop_roundtrip_is_stable(): + while_loop.use_explicit_cf = True + _assert_roundtrip_json_stable(while_loop.to_sdfg()) + + @dace.program def while_loop_with_break_continue(): A = dace.ndarray([10], dtype=dace.int32) diff --git a/tests/python_frontend/slice_subset_aliasing_test.py b/tests/python_frontend/slice_subset_aliasing_test.py new file mode 100644 index 0000000000..25f5ae3cbe --- /dev/null +++ b/tests/python_frontend/slice_subset_aliasing_test.py @@ -0,0 +1,55 @@ +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. +"""The Python frontend must give each memlet its own ``Subset`` object. + +Repeated reads of the same array slice hit the access cache and shared one +``Range`` object (``Memlet.simple`` stores the subset by reference), so an +in-place subset rewrite on one edge corrupted the other. +""" +import numpy as np + +import dace + +N = dace.symbol('N') + + +@dace.program +def two_sibling_slice_reads(out: dace.float64[N, N], arr: dace.float64[N, N]): + """Reads ``arr[i, k]`` in two sibling ``k`` loops; the second hits the cache.""" + for i in dace.map[0:N]: + for k in range(N): + out[i, k] = arr[i, k] + for k in range(N): + out[i, k] += arr[i, k] + + +def test_frontend_memlets_do_not_share_subset_objects(): + """No two memlets may share a subset (or other_subset) object.""" + sdfg = two_sibling_slice_reads.to_sdfg(simplify=False) + seen = {} + for sd in sdfg.all_sdfgs_recursive(): + for state in sd.all_states(): + for e in state.edges(): + if e.data is None: + continue + for sub in (e.data.subset, e.data.other_subset): + if sub is None: + continue + here = (state.label, e.data.data, str(sub)) + assert id(sub) not in seen, \ + f'subset object shared by two memlets: {seen[id(sub)]} and {here}' + seen[id(sub)] = here + + +def test_sibling_slice_reads_value_preserving(): + """End-to-end value check for the shape that triggered the aliasing.""" + n = 8 + rng = np.random.default_rng(0) + arr = rng.standard_normal((n, n)).astype(np.float64) + out = np.zeros((n, n)) + two_sibling_slice_reads(out=out, arr=arr, N=n) + assert np.allclose(out, 2.0 * arr) + + +if __name__ == '__main__': + test_frontend_memlets_do_not_share_subset_objects() + test_sibling_slice_reads_value_preserving() diff --git a/tests/python_frontend/type_statement_test.py b/tests/python_frontend/type_statement_test.py index 583b50f4eb..7f345453ed 100644 --- a/tests/python_frontend/type_statement_test.py +++ b/tests/python_frontend/type_statement_test.py @@ -1,6 +1,4 @@ # Copyright 2019-2023 ETH Zurich and the DaCe authors. All rights reserved. -import dace -import pytest # TODO: Investigate why pytest parses the DaCeProgram, even when the test is not supposed to run. # @pytest.mark.py312 diff --git a/tests/range_add_test.py b/tests/range_add_test.py index 51037b9306..f266cfe1f0 100644 --- a/tests/range_add_test.py +++ b/tests/range_add_test.py @@ -1,6 +1,5 @@ # Copyright 2019-2025 ETH Zurich and the DaCe authors. All rights reserved. -import dace from dace.subsets import Range import pytest diff --git a/tests/range_test.py b/tests/range_test.py new file mode 100644 index 0000000000..6005fa04fb --- /dev/null +++ b/tests/range_test.py @@ -0,0 +1,22 @@ +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. + +from dace import subsets + + +def test_range() -> None: + range = subsets.Range([(2, 9, 1, 4)]) + + assert range[0] == (2, 9, 1) + assert range.tile_sizes == [4] + + assert range.min_element() == [2] + assert range.max_element() == [9] + + +def test_range_reorder() -> None: + range = subsets.Range([(0, 4, 1, 1), (0, 9, 1, 5)]) + + range.reorder((1, 0)) + assert range[0] == (0, 9, 1) + assert range[1] == (0, 4, 1) + assert range.tile_sizes == [5, 1] diff --git a/tests/reduce_strided_disabledtest.py b/tests/reduce_strided_disabledtest.py index 6a6f1a24e0..635b466874 100644 --- a/tests/reduce_strided_disabledtest.py +++ b/tests/reduce_strided_disabledtest.py @@ -1,5 +1,4 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. -import copy import dace import dace.sdfg.nodes import numpy as np diff --git a/tests/reshape_test.py b/tests/reshape_test.py index eff6ef27a4..959134be8d 100644 --- a/tests/reshape_test.py +++ b/tests/reshape_test.py @@ -3,9 +3,7 @@ import dace from dace.sdfg import propagation as prop -from dace import nodes import dace.library -from dace.transformation import transformation as xf def test_unsqueeze(): diff --git a/tests/schedule_tree/naming_test.py b/tests/schedule_tree/naming_test.py index 8c39e3033f..420b0c642c 100644 --- a/tests/schedule_tree/naming_test.py +++ b/tests/schedule_tree/naming_test.py @@ -1,11 +1,10 @@ -# Copyright 2019-2023 ETH Zurich and the DaCe authors. All rights reserved. +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. import dace from dace.sdfg.analysis.schedule_tree import treenodes as tn from dace.sdfg.analysis.schedule_tree.sdfg_to_tree import as_schedule_tree from dace.transformation.passes.constant_propagation import ConstantPropagation import pytest -from typing import List def _irreducible_loop_to_loop(): @@ -18,11 +17,11 @@ def _irreducible_loop_to_loop(): # Add a loop l1 = sdfg.add_state() l2 = sdfg.add_state_after(l1) - sdfg.add_loop(s1, l1, s2, 'i', '0', 'i < 10', 'i + 1', loop_end_state=l2) + sdfg.add_loop_state_machine(s1, l1, s2, 'i', '0', 'i < 10', 'i + 1', loop_end_state=l2) l3 = sdfg.add_state() l4 = sdfg.add_state_after(l3) - sdfg.add_loop(s2, l3, e, 'i', '0', 'i < 10', 'i + 1', loop_end_state=l4) + sdfg.add_loop_state_machine(s2, l3, e, 'i', '0', 'i < 10', 'i + 1', loop_end_state=l4) # Irreducible part sdfg.add_edge(l3, l1, dace.InterstateEdge('i < 5')) @@ -171,7 +170,7 @@ def test_edgecase_symbol_mapping(): def _check_for_name_clashes(stree: tn.ScheduleTreeNode): - def _traverse(node: tn.ScheduleTreeScope, scopes: List[str]): + def _traverse(node: tn.ScheduleTreeScope, scopes: list[str]): for child in node.children: if isinstance(child, tn.LoopScope): itervar = child.loop.loop.loop_variable diff --git a/tests/schedule_tree/nesting_test.py b/tests/schedule_tree/nesting_test.py index 59512f88ab..5f4a102d2a 100644 --- a/tests/schedule_tree/nesting_test.py +++ b/tests/schedule_tree/nesting_test.py @@ -1,4 +1,4 @@ -# Copyright 2019-2023 ETH Zurich and the DaCe authors. All rights reserved. +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. """ Nesting and dealiasing tests for schedule trees. """ @@ -63,8 +63,8 @@ def tester(A: dace.float64[N, N]): simplified = dace.Config.get_bool('optimizer', 'automatic_simplification') if simplified: - assert [type(n) - for n in stree.preorder_traversal()][1:] == [tn.MapScope, tn.MapScope, tn.LoopScope, tn.TaskletNode] + node_types = [type(n) for n in stree.preorder_traversal()][1:] + assert node_types == [tn.MapScope, tn.MapScope, tn.ForScope, tn.TaskletNode] tasklet: tn.TaskletNode = list(stree.preorder_traversal())[-1] diff --git a/tests/schedule_tree/propagation_test.py b/tests/schedule_tree/propagation_test.py new file mode 100644 index 0000000000..f73ed037b8 --- /dev/null +++ b/tests/schedule_tree/propagation_test.py @@ -0,0 +1,109 @@ +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. +""" +Tests schedule tree input/output memlet computation +""" +import dace +from dace.sdfg import nodes +from dace.sdfg.analysis.schedule_tree import tree_to_sdfg as t2s, treenodes as tn +from dace.properties import CodeBlock + +import numpy as np +import pytest + + +def test_stree_propagation_forloop(): + N = dace.symbol('N') + + @dace.program + def tester(a: dace.float64[20]): + for i in range(1, N): + a[i] = 2 + a[1] = 1 + + stree = tester.to_sdfg().as_schedule_tree() + stree = t2s._insert_state_boundaries_to_tree(stree) + + node_types = [n for n in stree.preorder_traversal()] + assert isinstance(node_types[2], tn.ForScope) + memlet = dace.Memlet('a[1:N]') + memlet._is_data_src = False + assert list(node_types[2].output_memlets()) == [memlet] + + +@pytest.mark.skip("Suboptimal memlet propagation: https://github.com/spcl/dace/issues/2293") +def test_stree_propagation_symassign(): + # Manually create a schedule tree + N = dace.symbol('N') + stree = tn.ScheduleTreeRoot( + name='tester', + containers={ + 'A': dace.data.Array(dace.float64, [20]), + }, + symbols={ + 'N': N, + }, + children=[ + tn.MapScope(node=dace.nodes.MapEntry(dace.nodes.Map('map', ['i'], dace.subsets.Range([(1, N - 1, 1)]))), + children=[ + tn.AssignNode('j', CodeBlock('N + i'), dace.InterstateEdge(assignments=dict(j='N + i'))), + tn.TaskletNode(nodes.Tasklet('inner', {}, {'out'}, 'out = inp + 2'), + {'inp': dace.Memlet('A[j]')}, {'out': dace.Memlet('A[j]')}), + ]), + ], + ) + stree.children[0].parent = stree + for c in stree.children[0].children: + c.parent = stree.children[0] + + assert list(stree.children[0].input_memlets()) == [dace.Memlet('A[0:20]', volume=N - 1)] + + +@pytest.mark.skip("Suboptimal memlet propagation: https://github.com/spcl/dace/issues/2293") +def test_stree_propagation_dynset(): + H = dace.symbol('H') + nnz = dace.symbol('nnz') + W = dace.symbol('W') + + @dace.program + def spmv(A_row: dace.uint32[H + 1], A_col: dace.uint32[nnz], A_val: dace.float32[nnz], x: dace.float32[W]): + b = np.zeros([H], dtype=np.float32) + + for i in dace.map[0:H]: + for j in dace.map[A_row[i]:A_row[i + 1]]: + b[i] += A_val[j] * x[A_col[j]] + + return b + + sdfg = spmv.to_sdfg() + stree = sdfg.as_schedule_tree() + assert len(stree.children) == 2 + assert all(isinstance(c, tn.MapScope) for c in stree.children) + mapscope = stree.children[1] + _, _, dynrangemap = mapscope.children + assert isinstance(dynrangemap, tn.MapScope) + + # Check dynamic range map memlets + internal_memlets = list(dynrangemap.input_memlets()) + internal_memlet_data = [m.data for m in internal_memlets] + assert 'x' in internal_memlet_data + assert 'A_val' in internal_memlet_data + assert 'A_row' not in internal_memlet_data + for m in internal_memlets: + if m.data == 'A_val': + assert m.subset != dace.subsets.Range([(0, nnz - 1, 1)]) # Not propagated + + # Check top-level scope memlets + external_memlets = list(mapscope.input_memlets()) + assert dace.Memlet('A_row[0:H]') in external_memlets + assert dace.Memlet('A_row[1:H+1]') in external_memlets + assert dace.Memlet('x[0:W]', volume=0, dynamic=True) in external_memlets + assert dace.Memlet('A_val[0:nnz]', volume=0, dynamic=True) in external_memlets + for m in external_memlets: + if m.data == 'A_val': + assert m.subset == dace.subsets.Range([(0, nnz - 1, 1)]) # Propagated + + +if __name__ == '__main__': + test_stree_propagation_forloop() + test_stree_propagation_symassign() + test_stree_propagation_dynset() diff --git a/tests/schedule_tree/roundtrip_test.py b/tests/schedule_tree/roundtrip_test.py new file mode 100644 index 0000000000..bd6299a655 --- /dev/null +++ b/tests/schedule_tree/roundtrip_test.py @@ -0,0 +1,59 @@ +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. +""" +Tests conversion of schedule trees to SDFGs. +""" +import dace +import numpy as np + + +def test_implicit_inline_and_constants(): + """ + Tests implicit inlining upon roundtrip conversion, as well as constants with conflicting names. + """ + + @dace + def nester(A: dace.float64[20]): + A[:] = 12 + + @dace.program + def tester(A: dace.float64[20, 20]): + for i in dace.map[0:20]: + nester(A[:, i]) + + sdfg = tester.to_sdfg(simplify=False) + + # Inject constant into nested SDFG + assert len(list(sdfg.all_sdfgs_recursive())) > 1 + sdfg.add_constant('cst', 13) # Add an unused constant + sdfg.cfg_list[-1].add_constant('cst', 1, dace.data.Scalar(dace.float64)) + tasklet = next(n for n, _ in sdfg.all_nodes_recursive() if isinstance(n, dace.nodes.Tasklet)) + tasklet.code.as_string = tasklet.code.as_string.replace('12', 'cst') + + # Perform a roundtrip conversion + stree = sdfg.as_schedule_tree() + new_sdfg = stree.as_sdfg() + + assert len(list(new_sdfg.all_sdfgs_recursive())) == 1 + assert new_sdfg.constants['cst_0'].dtype == np.float64 + + # Test SDFG + a = np.random.rand(20, 20) + new_sdfg(A=a) # Tests arg_names + assert np.allclose(a, 1) + + +def test_name_propagation(): + name = "my_complicated_sdfg_test_name" + sdfg = dace.SDFG(name) + sdfg.add_state("empty", is_start_block=True) + + stree = sdfg.as_schedule_tree() + assert stree.name == name + + sdfg = stree.as_sdfg() + assert sdfg.name == name + + +if __name__ == '__main__': + test_implicit_inline_and_constants() + test_name_propagation() diff --git a/tests/schedule_tree/schedule_test.py b/tests/schedule_tree/schedule_test.py index c15eb99f88..442eb2de0b 100644 --- a/tests/schedule_tree/schedule_test.py +++ b/tests/schedule_tree/schedule_test.py @@ -1,4 +1,4 @@ -# Copyright 2019-2025 ETH Zurich and the DaCe authors. All rights reserved. +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. import dace from dace.sdfg.analysis.schedule_tree import treenodes as tn @@ -151,14 +151,14 @@ def test_irreducible_sub_sdfg(): sdfg.add_edge(s2, e, dace.InterstateEdge('b < 0')) # Add a loop following general block - sdfg.add_loop(e, sdfg.add_state(), None, 'i', '0', 'i < 10', 'i + 1') + sdfg.add_loop_state_machine(e, sdfg.add_state(), None, 'i', '0', 'i < 10', 'i + 1') FixedPointPipeline([ControlFlowRaising()]).apply_pass(sdfg, {}) stree = as_schedule_tree(sdfg) node_types = [type(n) for n in stree.preorder_traversal()] assert node_types.count(tn.GBlock) == 1 # Only one gblock - assert node_types.count(tn.LoopScope) == 1 # Check that the loop was detected + assert node_types.count(tn.ForScope) == 1 # Check that the for-loop was detected def test_irreducible_in_loops(): @@ -171,12 +171,11 @@ def test_irreducible_in_loops(): # Add a loop l1 = sdfg.add_state() l2 = sdfg.add_state_after(l1) - sdfg.add_loop(s1, l1, s2, 'i', '0', 'i < 10', 'i + 1', loop_end_state=l2) + sdfg.add_loop_state_machine(s1, l1, s2, 'i', '0', 'i < 10', 'i + 1', loop_end_state=l2) l3 = sdfg.add_state() l4 = sdfg.add_state_after(l3) - sdfg.add_loop(s2, l3, e, 'i', '0', 'i < 10', 'i + 1', loop_end_state=l4) - + sdfg.add_loop_state_machine(s2, l3, e, 'i', '0', 'i < 10', 'i + 1', loop_end_state=l4) # Irreducible part sdfg.add_edge(l3, l1, dace.InterstateEdge('i < 5')) diff --git a/tests/schedule_tree/to_sdfg_test.py b/tests/schedule_tree/to_sdfg_test.py new file mode 100644 index 0000000000..e6ae4658a1 --- /dev/null +++ b/tests/schedule_tree/to_sdfg_test.py @@ -0,0 +1,916 @@ +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. +""" +Tests components in conversion of schedule trees to SDFGs. +""" +import dace +from dace import data, subsets as sbs +from dace.codegen import control_flow as cf +from dace.properties import CodeBlock +from dace.sdfg import nodes +from dace.sdfg.analysis.schedule_tree import tree_to_sdfg as t2s, treenodes as tn +from dace.sdfg.state import ConditionalBlock, LoopRegion, SDFGState + +import pytest + + +def test_state_boundaries_none(): + # Manually create a schedule tree + stree = tn.ScheduleTreeRoot( + name='tester', + containers={ + 'A': data.Array(dace.float64, [20]), + }, + children=[ + tn.TaskletNode(nodes.Tasklet('bla', {}, {'out'}, 'out = 1'), {}, {'out': dace.Memlet('A[1]')}), + tn.TaskletNode(nodes.Tasklet('bla2', {'inp'}, {'out'}, 'out = inp + 1'), {'inp': dace.Memlet('A[1]')}, + {'out': dace.Memlet('A[1]')}), + ], + ) + + stree = t2s._insert_state_boundaries_to_tree(stree) + assert tn.StateBoundaryNode not in [type(n) for n in stree.children] + + +def test_state_boundaries_waw(): + # Manually create a schedule tree + stree = tn.ScheduleTreeRoot( + name='tester', + containers={ + 'A': data.Array(dace.float64, [20]), + }, + children=[ + tn.TaskletNode(nodes.Tasklet('bla', {}, {'out'}, 'out = 1'), {}, {'out': dace.Memlet('A[1]')}), + tn.TaskletNode(nodes.Tasklet('bla2', {}, {'out'}, 'out = 2'), {}, {'out': dace.Memlet('A[1]')}), + ], + ) + + stree = t2s._insert_state_boundaries_to_tree(stree) + assert [tn.TaskletNode, tn.StateBoundaryNode, tn.TaskletNode] == [type(n) for n in stree.children] + + +@pytest.mark.parametrize('overlap', (False, True)) +def test_state_boundaries_waw_ranges(overlap): + # Manually create a schedule tree + N = dace.symbol('N') + stree = tn.ScheduleTreeRoot( + name='tester', + containers={ + 'A': data.Array(dace.float64, [20]), + }, + symbols={'N': N}, + children=[ + tn.TaskletNode(nodes.Tasklet('bla', {}, {'out'}, 'pass'), {}, {'out': dace.Memlet('A[0:N/2]')}), + tn.TaskletNode(nodes.Tasklet('bla2', {}, {'out'}, 'pass'), {}, + {'out': dace.Memlet('A[1:N]' if overlap else 'A[N/2+1:N]')}), + ], + ) + + stree = t2s._insert_state_boundaries_to_tree(stree) + if overlap: + assert [tn.TaskletNode, tn.StateBoundaryNode, tn.TaskletNode] == [type(n) for n in stree.children] + else: + assert [tn.TaskletNode, tn.TaskletNode] == [type(n) for n in stree.children] + + +def test_state_boundaries_war(): + # Manually create a schedule tree + stree = tn.ScheduleTreeRoot( + name='tester', + containers={ + 'A': data.Array(dace.float64, [20]), + 'B': data.Array(dace.float64, [20]), + }, + children=[ + tn.TaskletNode(nodes.Tasklet('bla', {'inp'}, {'out'}, 'out = inp + 1'), {'inp': dace.Memlet('A[1]')}, + {'out': dace.Memlet('B[0]')}), + tn.TaskletNode(nodes.Tasklet('bla2', {}, {'out'}, 'out = 2'), {}, {'out': dace.Memlet('A[1]')}), + ], + ) + + stree = t2s._insert_state_boundaries_to_tree(stree) + assert [tn.TaskletNode, tn.StateBoundaryNode, tn.TaskletNode] == [type(n) for n in stree.children] + + +def test_state_boundaries_read_write_chain(): + # Manually create a schedule tree + stree = tn.ScheduleTreeRoot( + name='tester', + containers={ + 'A': data.Array(dace.float64, [20]), + 'B': data.Array(dace.float64, [20]), + }, + children=[ + tn.TaskletNode(nodes.Tasklet('bla1', {'inp'}, {'out'}, 'out = inp + 1'), {'inp': dace.Memlet('A[1]')}, + {'out': dace.Memlet('B[0]')}), + tn.TaskletNode(nodes.Tasklet('bla2', {'inp'}, {'out'}, 'out = inp + 1'), {'inp': dace.Memlet('B[0]')}, + {'out': dace.Memlet('A[1]')}), + tn.TaskletNode(nodes.Tasklet('bla3', {'inp'}, {'out'}, 'out = inp + 1'), {'inp': dace.Memlet('A[1]')}, + {'out': dace.Memlet('B[0]')}), + ], + ) + + stree = t2s._insert_state_boundaries_to_tree(stree) + assert [tn.TaskletNode, tn.TaskletNode, tn.TaskletNode] == [type(n) for n in stree.children] + + +def test_state_boundaries_data_race(): + # Manually create a schedule tree + stree = tn.ScheduleTreeRoot( + name='tester', + containers={ + 'A': data.Array(dace.float64, [20]), + 'B': data.Array(dace.float64, [20]), + }, + children=[ + tn.TaskletNode(nodes.Tasklet('bla1', {'inp'}, {'out'}, 'out = inp + 1'), {'inp': dace.Memlet('A[1]')}, + {'out': dace.Memlet('B[0]')}), + tn.TaskletNode(nodes.Tasklet('bla11', {'inp'}, {'out'}, 'out = inp + 1'), {'inp': dace.Memlet('A[1]')}, + {'out': dace.Memlet('B[1]')}), + tn.TaskletNode(nodes.Tasklet('bla2', {'inp'}, {'out'}, 'out = inp + 1'), {'inp': dace.Memlet('B[0]')}, + {'out': dace.Memlet('A[1]')}), + tn.TaskletNode(nodes.Tasklet('bla3', {'inp'}, {'out'}, 'out = inp + 1'), {'inp': dace.Memlet('A[1]')}, + {'out': dace.Memlet('B[0]')}), + ], + ) + + stree = t2s._insert_state_boundaries_to_tree(stree) + assert [tn.TaskletNode, tn.TaskletNode, tn.StateBoundaryNode, tn.TaskletNode, + tn.TaskletNode] == [type(n) for n in stree.children] + + +def test_state_boundaries_cfg(): + # Manually create a schedule tree + stree = tn.ScheduleTreeRoot( + name='tester', + containers={ + 'A': data.Array(dace.float64, [20]), + }, + children=[ + tn.TaskletNode(nodes.Tasklet('bla1', {}, {'out'}, 'out = 2'), {}, {'out': dace.Memlet('A[1]')}), + tn.ForScope(loop=cf.LoopRegion(label="for-loop", + condition_expr=CodeBlock("i < 20"), + loop_var="i", + initialize_expr=CodeBlock("i=0"), + update_expr=CodeBlock("i = i+1")), + children=[ + tn.TaskletNode(nodes.Tasklet('bla2', {}, {'out'}, 'out = i'), {}, + {'out': dace.Memlet('A[1]')}), + ]), + ], + ) + + stree = t2s._insert_state_boundaries_to_tree(stree) + assert [tn.TaskletNode, tn.StateBoundaryNode, tn.ForScope] == [type(n) for n in stree.children] + + +def test_state_boundaries_state_transition(): + # Manually create a schedule tree + stree = tn.ScheduleTreeRoot( + name='tester', + containers={ + 'A': data.Array(dace.float64, [20]), + }, + symbols={ + 'N': dace.symbol('N'), + }, + children=[ + tn.AssignNode('irrelevant', CodeBlock('N + 1'), dace.InterstateEdge(assignments=dict(irrelevant='N + 1'))), + tn.TaskletNode(nodes.Tasklet('bla', {}, {'out'}, 'out = 2'), {}, {'out': dace.Memlet('A[1]')}), + tn.AssignNode('relevant', CodeBlock('A[1] + 2'), + dace.InterstateEdge(assignments=dict(relevant='A[1] + 2'))), + ], + ) + + stree = t2s._insert_state_boundaries_to_tree(stree) + assert [ + tn.AssignNode, tn.StateBoundaryNode, tn.TaskletNode, tn.StateBoundaryNode, tn.AssignNode, tn.StateBoundaryNode + ] == [type(n) for n in stree.children] + + +@pytest.mark.parametrize('boundary', (False, True)) +def test_state_boundaries_propagation(boundary): + # Manually create a schedule tree + N = dace.symbol('N') + stree = tn.ScheduleTreeRoot( + name='tester', + containers={ + 'A': data.Array(dace.float64, [20]), + }, + symbols={ + 'N': N, + }, + children=[ + tn.MapScope(node=dace.nodes.MapEntry(dace.nodes.Map('map', ['i'], dace.subsets.Range([(1, N - 1, 1)]))), + children=[ + tn.TaskletNode(nodes.Tasklet('inner', {}, {'out'}, 'out = 2'), {}, + {'out': dace.Memlet('A[i]')}), + ]), + tn.TaskletNode(nodes.Tasklet('bla', {}, {'out'}, 'out = 2'), {}, + {'out': dace.Memlet('A[1]' if boundary else 'A[0]')}), + ], + ) + + stree = t2s._insert_state_boundaries_to_tree(stree) + + node_types = [type(n) for n in stree.preorder_traversal()] + if boundary: + assert [tn.MapScope, tn.TaskletNode, tn.StateBoundaryNode, tn.TaskletNode] == node_types[1:] + else: + assert [tn.MapScope, tn.TaskletNode, tn.TaskletNode] == node_types[1:] + + +@pytest.mark.parametrize("control_flow", (True, False)) +def test_create_state_boundary_state_transition(control_flow: bool): + sdfg = dace.SDFG("tester") + state = sdfg.add_state("start", is_start_block=True) + bnode = tn.StateBoundaryNode(control_flow) + + t2s._create_state_boundary(bnode, state) + new_label = "cf_state_boundary" if control_flow else "state_boundary" + assert ["start", new_label] == [state.label for state in sdfg.states()] + + +@pytest.mark.xfail(reason="Not yet implemented") +def test_create_state_boundary_empty_memlet(): + t2s._StreeToSDFG(boundary_behavior=t2s.StateBoundaryBehavior.EMPTY_MEMLET) + + +def test_create_tasklet_raw(): + stree = tn.ScheduleTreeRoot( + name='tester', + containers={ + 'A': data.Array(dace.float64, [20]), + }, + children=[ + tn.TaskletNode(nodes.Tasklet('bla', {}, {'out'}, 'out = 1'), {}, {'out': dace.Memlet('A[1]')}), + tn.TaskletNode(nodes.Tasklet('bla2', {'inp'}, {'out'}, 'out = inp + 1'), {'inp': dace.Memlet('A[1]')}, + {'out': dace.Memlet('A[1]')}), + ], + ) + + sdfg = stree.as_sdfg() + assert len(sdfg.states()) == 1 + state = sdfg.states()[0] + first_tasklet, write_read_node, second_tasklet, write_node = state.nodes() + + assert isinstance(first_tasklet, nodes.Tasklet) + assert first_tasklet.label == "bla" + assert not first_tasklet.in_connectors + assert first_tasklet.out_connectors.keys() == {"out"} + + assert isinstance(second_tasklet, nodes.Tasklet) + assert second_tasklet.label == "bla2" + assert second_tasklet.in_connectors.keys() == {"inp"} + assert second_tasklet.out_connectors.keys() == {"out"} + + assert [(first_tasklet, write_read_node), (write_read_node, second_tasklet), + (second_tasklet, write_node)] == [(edge.src, edge.dst) for edge in state.edges()] + + +def test_create_tasklet_waw(): + stree = tn.ScheduleTreeRoot( + name='tester', + containers={ + 'A': data.Array(dace.float64, [20]), + }, + children=[ + tn.TaskletNode(nodes.Tasklet('bla', {}, {'out'}, 'out = 1'), {}, {'out': dace.Memlet('A[1]')}), + tn.TaskletNode(nodes.Tasklet('bla2', {}, {'out'}, 'out = 2'), {}, {'out': dace.Memlet('A[1]')}), + ], + ) + + sdfg = stree.as_sdfg() + assert len(sdfg.states()) == 2 + s1, s2 = sdfg.states() + + s1_tasklet, s1_anode = s1.nodes() + assert [(s1_tasklet, s1_anode)] == [(edge.src, edge.dst) for edge in s1.edges()] + + s2_tasklet, s2_anode = s2.nodes() + assert [(s2_tasklet, s2_anode)] == [(edge.src, edge.dst) for edge in s2.edges()] + + +def test_create_tasklet_war(): + stree = tn.ScheduleTreeRoot( + name="tester", + containers={"A": data.Array(dace.float64, [20])}, + children=[ + tn.TaskletNode( + nodes.Tasklet("read_write", {"read"}, {"write"}, "write = read + 1"), + {"read": dace.Memlet("A[1]")}, + {"write": dace.Memlet("A[1]")}, + ) + ], + ) + + sdfg = stree.as_sdfg() + + sdfg_states = list(sdfg.states()) + assert len(sdfg_states) == 1 + + state_nodes = list(sdfg_states[0].nodes()) + assert [node.name for node in state_nodes + if isinstance(node, nodes.Tasklet)] == ["read_write"], "Expect one Tasklet node." + assert [node.data for node in state_nodes + if isinstance(node, nodes.AccessNode)] == ["A", "A"], "Expect two AccessNodes for A." + + +def test_create_loop_for(): + for_scope = tn.ForScope( + loop=LoopRegion(label="my_for_loop", + loop_var="i", + initialize_expr=CodeBlock("i = 0 "), + condition_expr=CodeBlock("i < 3"), + update_expr=CodeBlock("i = i+1")), + children=[ + tn.TaskletNode(nodes.Tasklet('assign_1', {}, {'out'}, 'out = 1'), {}, {'out': dace.Memlet('A[1]')}), + tn.TaskletNode(nodes.Tasklet('assign_2', {}, {'out'}, 'out = 2'), {}, {'out': dace.Memlet('A[1]')}), + ], + ) + + stree = tn.ScheduleTreeRoot(name='tester', containers={'A': data.Array(dace.float64, [20])}, children=[for_scope]) + sdfg = stree.as_sdfg() + + loops = list(filter(lambda x: isinstance(x, LoopRegion), sdfg.cfg_list)) + assert len(loops) == 1, "SDFG contains one LoopRegion" + + loop = loops[0] + assert isinstance(loop, LoopRegion) + assert loop.loop_variable == "i" + assert loop.init_statement == CodeBlock("i = 0") + assert loop.loop_condition == CodeBlock("i < 3") + assert loop.update_statement == CodeBlock("i = i+1") + + loop_states = list(filter(lambda x: isinstance(x, SDFGState), loop.nodes())) + assert len(loop_states) == 2, "Loop contains two states" + + tasklet_1: nodes.Tasklet = list(filter(lambda x: isinstance(x, nodes.Tasklet), loop_states[0].nodes()))[0] + assert tasklet_1.label == "assign_1" + tasklet_2: nodes.Tasklet = list(filter(lambda x: isinstance(x, nodes.Tasklet), loop_states[1].nodes()))[0] + assert tasklet_2.label == "assign_2" + + +def test_create_loop_while(): + while_scope = tn.WhileScope( + children=[ + tn.TaskletNode(nodes.Tasklet('assign_1', {}, {'out'}, 'out = 1'), {}, {'out': dace.Memlet('A[1]')}), + tn.TaskletNode(nodes.Tasklet('assign_2', {}, {'out'}, 'out = 2'), {}, {'out': dace.Memlet('A[1]')}), + ], + loop=LoopRegion( + label="my_while_loop", + condition_expr=CodeBlock("A[1] > 5"), + ), + ) + + stree = tn.ScheduleTreeRoot(name='tester', containers={'A': data.Array(dace.float64, [20])}, children=[while_scope]) + + sdfg = stree.as_sdfg() + + loops = list(filter(lambda x: isinstance(x, LoopRegion), sdfg.cfg_list)) + assert len(loops) == 1, "SDFG contains one LoopRegion" + + loop = loops[0] + assert isinstance(loop, LoopRegion) + assert loop.loop_variable == "" + assert loop.init_statement == None + assert loop.loop_condition == CodeBlock("A[1] > 5") + assert loop.update_statement == None + + loop_states = list(filter(lambda x: isinstance(x, SDFGState), loop.nodes())) + assert len(loop_states) == 2, "Loop contains two states" + + tasklet_1: nodes.Tasklet = list(filter(lambda x: isinstance(x, nodes.Tasklet), loop_states[0].nodes()))[0] + assert tasklet_1.label == "assign_1" + tasklet_2: nodes.Tasklet = list(filter(lambda x: isinstance(x, nodes.Tasklet), loop_states[1].nodes()))[0] + assert tasklet_2.label == "assign_2" + + +def test_create_if_else(): + stree = tn.ScheduleTreeRoot( + name="tester", + containers={'A': data.Array(dace.float64, [20])}, + children=[ + tn.IfScope( + condition=CodeBlock("A[0] > 0"), + children=[ + tn.TaskletNode(nodes.Tasklet("bla", {}, {"out"}, "out=1"), {}, {"out": dace.Memlet("A[1]")}), + ], + ), + tn.ElseScope(children=[ + tn.TaskletNode(nodes.Tasklet("blub", {}, {"out"}, "out=2"), {}, {"out": dace.Memlet("A[1]")}) + ]), + ]) + + sdfg = stree.as_sdfg() + + blocks = list(filter(lambda x: isinstance(x, ConditionalBlock), sdfg.cfg_list)) + assert len(blocks) == 1, "SDFG contains one ConditionalBlock" + + block = blocks[0] + assert isinstance(block, ConditionalBlock) + assert len(block.branches) == 2, "Block contains two branches" + + if_branch = list(filter(lambda x: x[0] is not None, block.branches))[0] + tasklets = list(filter(lambda x: isinstance(x, nodes.Tasklet), if_branch[1].nodes()[0].nodes())) + assert if_branch[0] == CodeBlock("A[0] > 0"), "If branch has condition" + assert tasklets[0].label == "bla", "If branch contains Tasklet('bla')" + + else_branch = list(filter(lambda x: x[0] is None, block.branches))[0] + tasklets = list(filter(lambda x: isinstance(x, nodes.Tasklet), else_branch[1].nodes()[0].nodes())) + assert else_branch[0] is None, "Else branch has no condition" + assert tasklets[0].label == "blub", "Else branch contains Tasklet('blub')" + + +@pytest.mark.xfail(reason="Not yet implemented") +def test_create_if_elif_else() -> None: + stree = tn.ScheduleTreeRoot( + name="tester", + containers={'A': data.Array(dace.float64, [20])}, + children=[ + tn.IfScope( + condition=CodeBlock("A[0] > 0"), + children=[ + tn.TaskletNode(nodes.Tasklet("bla", {}, {"out"}, "out=1"), {}, {"out": dace.Memlet("A[1]")}), + ], + ), + tn.ElifScope( + condition=CodeBlock("A[0] == 0"), + children=[ + tn.TaskletNode(nodes.Tasklet("blub", {}, {"out"}, "out=2"), {}, {"out": dace.Memlet("A[1]")}), + ], + ), + tn.ElseScope(children=[ + tn.TaskletNode(nodes.Tasklet("test", {}, {"out"}, "out=3"), {}, {"out": dace.Memlet("A[1]")}) + ]) + ]) + + sdfg = stree.as_sdfg() + + blocks = list(filter(lambda x: isinstance(x, ConditionalBlock), sdfg.cfg_list)) + assert len(blocks) == 1, "SDFG contains one ConditionalBlock" + + block = blocks[0] + assert isinstance(block, ConditionalBlock) + assert len(block.branches) == 3, "Block contains three branches" + + +def test_create_if_without_else(): + stree = tn.ScheduleTreeRoot( + name="tester", + containers={'A': data.Array(dace.float64, [20])}, + children=[ + tn.IfScope( + condition=CodeBlock("A[0] > 0"), + children=[ + tn.TaskletNode(nodes.Tasklet("bla", {}, {"out"}, "out=1"), {}, {"out": dace.Memlet("A[1]")}), + ], + ), + ], + ) + + sdfg = stree.as_sdfg() + + blocks = list(filter(lambda x: isinstance(x, ConditionalBlock), sdfg.cfg_list)) + assert len(blocks) == 1, "SDFG contains one ConditionalBlock" + + block = blocks[0] + assert isinstance(block, ConditionalBlock) + assert len(block.branches) == 1, "Block contains one branch" + + branch = list(filter(lambda x: x[0] is not None, block.branches))[0] + tasklets = list(filter(lambda x: isinstance(x, nodes.Tasklet), branch[1].nodes()[0].nodes())) + assert branch[0] == CodeBlock("A[0] > 0"), "Branch has condition" + assert tasklets[0].label == "bla", "Branch contains Tasklet('bla')" + + +def test_create_map_scope_write(): + stree = tn.ScheduleTreeRoot( + name="tester", + containers={'A': data.Array(dace.float64, [20])}, + children=[ + tn.MapScope( + node=nodes.MapEntry(nodes.Map("bla", "i", sbs.Range.from_string("0:20"))), + children=[ + tn.TaskletNode(nodes.Tasklet("assign_i", {}, {"out"}, "out = i"), {}, {"out": dace.Memlet("A[i]")}) + ], + ) + ], + ) + + sdfg = stree.as_sdfg() + sdfg.validate() + + +def test_create_map_scope_read_after_write(): + stree = tn.ScheduleTreeRoot( + name="tester", + containers={ + 'A': data.Array(dace.float64, [20]), + 'B': data.Array(dace.float64, [20], transient=True), + }, + children=[ + tn.MapScope( + node=nodes.MapEntry(nodes.Map("bla", "i", sbs.Range.from_string("0:20"))), + children=[ + tn.TaskletNode(nodes.Tasklet("write", {}, {"out"}, "out = i"), {}, {"out": dace.Memlet("B[i]")}), + tn.TaskletNode(nodes.Tasklet("read", {"in_field"}, {"out_field"}, "out_field = in_field"), + {"in_field": dace.Memlet("B[i]")}, {"out_field": dace.Memlet("A[i]")}) + ], + ) + ], + ) + + sdfg = stree.as_sdfg() + sdfg.validate() + + +def test_create_map_scope_write_after_read(): + stree = tn.ScheduleTreeRoot( + name="tester", + containers={"A": data.Array(dace.float64, [20])}, + children=[ + tn.MapScope( + node=nodes.MapEntry(nodes.Map("bla", "i", sbs.Range.from_string("0:20"))), + children=[ + tn.TaskletNode(nodes.Tasklet("read_write", {"read"}, {"write"}, "write = read+1"), + {"read": dace.Memlet("A[i]")}, {"write": dace.Memlet("A[i]")}) + ], + ) + ], + ) + + sdfg = stree.as_sdfg() + sdfg.validate() + + +def test_create_map_scope_copy(): + stree = tn.ScheduleTreeRoot( + name="tester", + containers={ + 'A': data.Array(dace.float64, [20]), + 'B': data.Array(dace.float64, [20]), + }, + children=[ + tn.MapScope( + node=nodes.MapEntry(nodes.Map("bla", "i", sbs.Range.from_string("0:20"))), + children=[ + tn.TaskletNode(nodes.Tasklet("copy", {"inp"}, {"out"}, "out = inp"), {"inp": dace.Memlet("A[i]")}, + {"out": dace.Memlet("B[i]")}) + ], + ) + ], + ) + + sdfg = stree.as_sdfg() + sdfg.validate() + + +def test_create_map_scope_double_memlet(): + stree = tn.ScheduleTreeRoot( + name="tester", + containers={ + 'A': data.Array(dace.float64, [20]), + 'B': data.Array(dace.float64, [20]), + }, + children=[ + tn.MapScope(node=nodes.MapEntry(nodes.Map("bla", "i", sbs.Range.from_string("0:10"))), + children=[ + tn.TaskletNode(nodes.Tasklet("sum", {"first", "second"}, {"out"}, "out = first + second"), { + "first": dace.Memlet("A[i]"), + "second": dace.Memlet("A[i+10]") + }, {"out": dace.Memlet("B[i]")}) + ]) + ]) + + sdfg = stree.as_sdfg() + sdfg.validate() + + +def test_create_nested_map_scope(): + stree = tn.ScheduleTreeRoot( + name="tester", + containers={'A': data.Array(dace.float64, [20])}, + children=[ + tn.MapScope( + node=nodes.MapEntry(nodes.Map("map_i", "i", sbs.Range.from_string("0:4"))), + children=[ + tn.MapScope( + node=nodes.MapEntry(nodes.Map("map_j", "j", sbs.Range.from_string("0:5"))), + children=[ + tn.TaskletNode(nodes.Tasklet("assign", {}, {"out"}, "out = i*5+j"), {}, + {"out": dace.Memlet("A[i*5+j]")}) + ], + ) + ], + ) + ], + ) + + sdfg = stree.as_sdfg() + sdfg.validate() + + +def test_double_map_with_for_loop(): + stree = tn.ScheduleTreeRoot( + name="tester", + containers={'A': data.Array(dace.float64, [60])}, + children=[ + tn.MapScope( + node=nodes.MapEntry(nodes.Map("map_i", "i", sbs.Range.from_string("0:4"))), + children=[ + tn.MapScope( + node=nodes.MapEntry(nodes.Map("map_j", "j", sbs.Range.from_string("0:5"))), + children=[ + tn.ForScope( + loop=LoopRegion( + label="loop_k", + loop_var="k", + initialize_expr=CodeBlock("k = 0 "), + condition_expr=CodeBlock("k < 3"), + update_expr=CodeBlock("k = k+1"), + ), + children=[ + tn.TaskletNode(nodes.Tasklet("assign", {}, {"out"}, "out = 1.0"), {}, + {"out": dace.Memlet("A[i*15+j*3+k]")}) + ], + ), + ], + ) + ], + ) + ], + ) + + sdfg = stree.as_sdfg() + assert sdfg.is_valid() + + +def test_triple_map_flat_if(): + stree = tn.ScheduleTreeRoot( + name="tester", + containers={'A': data.Array(dace.float64, [60])}, + children=[ + tn.MapScope( + node=nodes.MapEntry(nodes.Map("map_i", "i", sbs.Range.from_string("0:4"))), + children=[ + tn.MapScope( + node=nodes.MapEntry(nodes.Map("map_j", "j", sbs.Range.from_string("0:5"))), + children=[ + tn.MapScope( + node=nodes.MapEntry(nodes.Map("map_k", "k", sbs.Range.from_string("0:3"))), + children=[ + tn.IfScope( + condition=CodeBlock("A[0] > 0"), + children=[ + tn.TaskletNode(nodes.Tasklet("assign", {}, {"out"}, "out = 1"), {}, + {"out": dace.Memlet("A[i*15+j*3+k]")}) + ], + ), + tn.ElseScope(children=[ + tn.TaskletNode(nodes.Tasklet("assign", {}, {"out"}, "out = 2"), {}, + {"out": dace.Memlet("A[i*15+j*3+k]")}) + ], ), + ], + ) + ], + ) + ], + ) + ], + ) + + sdfg = stree.as_sdfg() + sdfg.validate() + + +def test_triple_map_nested_if(): + stree = tn.ScheduleTreeRoot( + name="tester", + containers={'A': data.Array(dace.float64, [60])}, + children=[ + tn.MapScope( + node=nodes.MapEntry(nodes.Map("map_i", "i", sbs.Range.from_string("0:4"))), + children=[ + tn.MapScope( + node=nodes.MapEntry(nodes.Map("map_j", "j", sbs.Range.from_string("0:5"))), + children=[ + tn.MapScope( + node=nodes.MapEntry(nodes.Map("map_k", "k", sbs.Range.from_string("0:3"))), + children=[ + tn.IfScope( + condition=CodeBlock("A[0] > 0"), + children=[ + tn.TaskletNode(nodes.Tasklet("assign", {}, {"out"}, "out = 1"), {}, + {"out": dace.Memlet("A[i*15+j*3+k]")}) + ], + ), + tn.ElseScope(children=[ + tn.IfScope( + condition=CodeBlock("A[1] > 0"), + children=[ + tn.TaskletNode(nodes.Tasklet("assign", {}, {"out"}, "out = 2"), {}, + {"out": dace.Memlet("A[i*15+j*3+k]")}) + ], + ), + tn.ElseScope(children=[ + tn.TaskletNode(nodes.Tasklet("assign", {}, {"out"}, "out = 3"), {}, + {"out": dace.Memlet("A[i*15+j*3+k]")}) + ], ) + ], ), + ], + ) + ], + ) + ], + ) + ], + ) + + sdfg = stree.as_sdfg() + sdfg.validate() + + +def test_create_nested_map_scope_multi_read(): + stree = tn.ScheduleTreeRoot( + name="tester", + containers={ + 'A': data.Array(dace.float64, [20]), + 'B': data.Array(dace.float64, [10]) + }, + children=[ + tn.MapScope( + node=nodes.MapEntry(nodes.Map("bla", "i", sbs.Range.from_string("0:2"))), + children=[ + tn.MapScope( + node=nodes.MapEntry(nodes.Map("blub", "j", sbs.Range.from_string("0:5"))), + children=[ + tn.TaskletNode(nodes.Tasklet("asdf", {"a_1", "a_2"}, {"out"}, "out = a_1 + a_2"), { + "a_1": dace.Memlet("A[i*5+j]"), + "a_2": dace.Memlet("A[10+i*5+j]"), + }, {"out": dace.Memlet("B[i*5+j]")}) + ], + ) + ], + ) + ], + ) + + sdfg = stree.as_sdfg() + sdfg.validate() + + +def test_map_with_state_boundary_inside(): + stree = tn.ScheduleTreeRoot( + name="tester", + containers={'A': data.Array(dace.float64, [20])}, + children=[ + tn.MapScope( + node=nodes.MapEntry(nodes.Map("bla", "i", sbs.Range.from_string("0:20"))), + children=[ + tn.TaskletNode(nodes.Tasklet('bla', {}, {'out'}, 'out = i'), {}, {'out': dace.Memlet('A[1]')}), + tn.TaskletNode(nodes.Tasklet('bla2', {}, {'out'}, 'out = 2*i'), {}, {'out': dace.Memlet('A[1]')}), + ], + ) + ], + ) + + sdfg = stree.as_sdfg() + sdfg.validate() + + +def test_map_calculate_temporary_in_two_loops(): + stree = tn.ScheduleTreeRoot( + name="tester", + containers={ + "A": data.Array(dace.float64, [20]), + "tmp": data.Array(dace.float64, [20], transient=True) + }, + children=[ + tn.MapScope( + node=nodes.MapEntry(nodes.Map("first_half", "i", sbs.Range.from_string("0:10"))), + children=[ + tn.TaskletNode(nodes.Tasklet("beginning", {}, {'out'}, 'out = i'), {}, + {'out': dace.Memlet("tmp[i]")}) + ], + ), + tn.MapScope( + node=nodes.MapEntry(nodes.Map("second_half", "i", sbs.Range.from_string("10:20"))), + children=[ + tn.TaskletNode(nodes.Tasklet("end", {}, {'out'}, 'out = i'), {}, {'out': dace.Memlet("tmp[i]")}) + ], + ), + tn.MapScope( + node=nodes.MapEntry(nodes.Map("read_tmp", "i", sbs.Range.from_string("0:20"))), + children=[ + tn.TaskletNode(nodes.Tasklet("read_temp", {"read"}, {"out"}, "out = read + 1"), + {"read": dace.Memlet("tmp[i]")}, {"out": dace.Memlet("A[i]")}) + ], + ) + ], + ) + + sdfg = stree.as_sdfg(simplify=True) + sdfg.validate() + + assert [node.name for node, _ in sdfg.all_nodes_recursive() + if isinstance(node, nodes.Tasklet)] == ["beginning", "end", "read_temp"] + + +def test_edge_assignment_read_after_write(): + stree = tn.ScheduleTreeRoot( + name="tester", + containers={}, + children=[ + tn.AssignNode("my_condition", CodeBlock("True"), dace.InterstateEdge()), + tn.AssignNode("condition", CodeBlock("my_condition"), dace.InterstateEdge()), + tn.StateBoundaryNode(), + ], + ) + + sdfg = stree.as_sdfg(simplify=False) + + assert [node.name for node in sdfg.nodes()] == ["tree_root", "state_boundary", "state_boundary_0"] + assert [edge.data.assignments for edge in sdfg.edges()] == [{"my_condition": "True"}, {"condition": "my_condition"}] + + +def test_assign_nodes_force_state_transition(): + stree = tn.ScheduleTreeRoot( + name='tester', + containers={ + 'A': data.Array(dace.float64, [20]), + }, + children=[ + tn.AssignNode("mySymbol", CodeBlock("1"), dace.InterstateEdge()), + tn.TaskletNode(nodes.Tasklet('bla', {}, {'out'}, 'out = mySymbol'), {}, {'out': dace.Memlet('A[1]')}), + ], + ) + + stree = t2s._insert_state_boundaries_to_tree(stree) + assert [type(child) for child in stree.children] == [tn.AssignNode, tn.StateBoundaryNode, tn.TaskletNode] + + +def test_assign_nodes_multiple_force_one_transition(): + stree = tn.ScheduleTreeRoot( + name='tester', + containers={ + 'A': data.Array(dace.float64, [20]), + }, + children=[ + tn.AssignNode("mySymbol", CodeBlock("1"), dace.InterstateEdge()), + tn.AssignNode("myOtherSymbol", CodeBlock("2"), dace.InterstateEdge()), + tn.TaskletNode(nodes.Tasklet('bla', {}, {'out'}, 'out = mySymbol + myOtherSymbol'), {}, + {'out': dace.Memlet('A[1]')}), + ], + ) + + stree = t2s._insert_state_boundaries_to_tree(stree) + assert [type(child) + for child in stree.children] == [tn.AssignNode, tn.AssignNode, tn.StateBoundaryNode, tn.TaskletNode] + + +def test_assign_nodes_avoid_duplicate_boundaries(): + stree = tn.ScheduleTreeRoot( + name='tester', + containers={ + 'A': data.Array(dace.float64, [20]), + }, + children=[ + tn.AssignNode("mySymbol", CodeBlock("1"), dace.InterstateEdge()), + tn.StateBoundaryNode(), + tn.TaskletNode(nodes.Tasklet('bla', {}, {'out'}, 'out = mySymbol + myOtherSymbol'), {}, + {'out': dace.Memlet('A[1]')}), + ], + ) + + stree = t2s._insert_state_boundaries_to_tree(stree) + assert [type(child) for child in stree.children] == [tn.AssignNode, tn.StateBoundaryNode, tn.TaskletNode] + + +if __name__ == '__main__': + test_state_boundaries_none() + test_state_boundaries_waw() + test_state_boundaries_waw_ranges(overlap=False) + test_state_boundaries_waw_ranges(overlap=True) + test_state_boundaries_war() + test_state_boundaries_read_write_chain() + test_state_boundaries_data_race() + test_state_boundaries_cfg() + test_state_boundaries_state_transition() + test_state_boundaries_propagation(boundary=False) + test_state_boundaries_propagation(boundary=True) + test_create_state_boundary_state_transition(control_flow=True) + test_create_state_boundary_state_transition(control_flow=False) + # test_create_state_boundary_empty_memlet() + test_create_tasklet_raw() + test_create_tasklet_waw() + test_create_loop_for() + test_create_loop_while() + test_create_if_else() + test_create_if_without_else() + test_create_map_scope_write() + test_create_map_scope_copy() + test_create_map_scope_double_memlet() + test_create_nested_map_scope() + test_create_nested_map_scope_multi_read() + test_map_with_state_boundary_inside() + test_edge_assignment_read_after_write() diff --git a/tests/schedule_tree/treenodes_test.py b/tests/schedule_tree/treenodes_test.py new file mode 100644 index 0000000000..31a0abbd21 --- /dev/null +++ b/tests/schedule_tree/treenodes_test.py @@ -0,0 +1,126 @@ +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. + +from dace.sdfg.analysis.schedule_tree import treenodes as tn +from dace import nodes + +import pytest + + +@pytest.fixture +def tasklet() -> nodes.Tasklet: + return tn.TaskletNode(nodes.Tasklet("noop", {}, {}, code="pass"), {}, {}) + + +@pytest.mark.parametrize('ScopeClass', ( + tn.ScheduleTreeScope, + tn.ControlFlowScope, + tn.GBlock, + tn.ElseScope, +)) +def test_schedule_tree_scope_children(ScopeClass: type[tn.ScheduleTreeScope], tasklet: nodes.Tasklet) -> None: + scope = ScopeClass(children=[tasklet]) + + for child in scope.children: + assert child.parent == scope + + scope = ScopeClass(children=[]) + scope.add_child(tasklet) + + for child in scope.children: + assert child.parent == scope + + scope = ScopeClass(children=[]) + scope.add_children([tasklet]) + + for child in scope.children: + assert child.parent == scope + + +@pytest.mark.parametrize('LoopScope', ( + tn.LoopScope, + tn.ForScope, + tn.WhileScope, + tn.DoWhileScope, +)) +def test_loop_scope_children(LoopScope: type[tn.LoopScope], tasklet: nodes.Tasklet) -> None: + scope = LoopScope(loop=None, children=[tasklet]) + + for child in scope.children: + assert child.parent == scope + + scope = LoopScope(loop=None, children=[]) + scope.add_child(tasklet) + + for child in scope.children: + assert child.parent == scope + + scope = LoopScope(loop=None, children=[]) + scope.add_children([tasklet]) + + for child in scope.children: + assert child.parent == scope + + +@pytest.mark.parametrize('IfScope', ( + tn.IfScope, + tn.StateIfScope, + tn.ElifScope, +)) +def test_if_scope_children(IfScope: type[tn.IfScope], tasklet: nodes.Tasklet) -> None: + scope = IfScope(condition=None, children=[tasklet]) + + for child in scope.children: + assert child.parent == scope + + scope = IfScope(condition=None, children=[]) + scope.add_child(tasklet) + + for child in scope.children: + assert child.parent == scope + + scope = IfScope(condition=None, children=[]) + scope.add_children([tasklet]) + + for child in scope.children: + assert child.parent == scope + + +@pytest.mark.parametrize('DataflowScope', ( + tn.DataflowScope, + tn.MapScope, + tn.ConsumeScope, +)) +def test_dataflow_scope_children(DataflowScope: type[tn.DataflowScope], tasklet: nodes.Tasklet) -> None: + scope = DataflowScope(node=None, children=[tasklet]) + + for child in scope.children: + assert child.parent == scope + + scope = DataflowScope(node=None, children=[]) + scope.add_child(tasklet) + + for child in scope.children: + assert child.parent == scope + + scope = DataflowScope(node=None, children=[]) + scope.add_children([tasklet]) + + for child in scope.children: + assert child.parent == scope + + +if __name__ == '__main__': + test_schedule_tree_scope_children(tn.ScheduleTreeScope, tasklet) + test_schedule_tree_scope_children(tn.ControlFlowScope, tasklet) + test_schedule_tree_scope_children(tn.GBlock, tasklet) + test_schedule_tree_scope_children(tn.ElseScope, tasklet) + test_loop_scope_children(tn.LoopScope, tasklet) + test_loop_scope_children(tn.ForScope, tasklet) + test_loop_scope_children(tn.WhileScope, tasklet) + test_loop_scope_children(tn.DoWhileScope, tasklet) + test_if_scope_children(tn.IfScope, tasklet) + test_if_scope_children(tn.StateIfScope, tasklet) + test_if_scope_children(tn.ElifScope, tasklet) + test_dataflow_scope_children(tn.DataflowScope, tasklet) + test_dataflow_scope_children(tn.MapScope, tasklet) + test_dataflow_scope_children(tn.ConsumeScope, tasklet) diff --git a/tests/scoped_symbol_serialization_test.py b/tests/scoped_symbol_serialization_test.py new file mode 100644 index 0000000000..b347b01088 --- /dev/null +++ b/tests/scoped_symbol_serialization_test.py @@ -0,0 +1,120 @@ +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. +"""Scope-driven symbol-dtype serialization. + +A symbol's serialized dtype is a deterministic function of the enclosing-scope +authority (the SDFG's registered symbols, overridden by any shadowing map/consume +scope), not of the symbol instance's own (SymPy-cache-stale) dtype. Outside SDFG +serialization the instance dtype is used. +""" +import json + +import dace +from dace import symbolic + + +def _resave(sdfg: dace.SDFG): + """Save -> load -> save; returns the two JSON strings (must match for a stable roundtrip).""" + s1 = json.dumps(sdfg.to_json()) + s2 = json.dumps(dace.SDFG.from_json(json.loads(s1)).to_json()) + return s1, s2 + + +def test_scope_authority_overrides_and_restores(): + """The active authority decides the dtype; opening a scope overrides and closing restores it.""" + sym = symbolic.symbol('x', dtype=dace.int64) + + # No scope active: the symbol's own (instance) dtype is used. + assert symbolic.serialize_symbolic(sym) == 'symbol($x, dtype=dace.int64)' + + with symbolic.serialization_symbol_dtypes({'x': dace.int32}): + assert symbolic.serialize_symbolic(sym) == '$x' # int32 == DEFAULT_SYMBOL_TYPE -> bare + with symbolic.serialization_symbol_dtypes({'x': dace.int64}): + assert symbolic.serialize_symbolic(sym) == 'symbol($x, dtype=dace.int64)' + assert symbolic.serialize_symbolic(sym) == '$x' # inner scope closed, outer restored + + assert symbolic.serialize_symbolic(sym) == 'symbol($x, dtype=dace.int64)' # back to no-scope + + +def test_scope_authority_only_overrides_named_symbols(): + """The authority overrides only the names it declares; an absent name keeps its own dtype.""" + sym = symbolic.symbol('y', dtype=dace.int64) + with symbolic.serialization_symbol_dtypes({'other': dace.int64}): + assert symbolic.serialize_symbolic(sym) == 'symbol($y, dtype=dace.int64)' # absent -> instance + + +def test_registered_sdfg_symbol_roundtrips_its_dtype(): + """A symbol registered on the SDFG keeps its declared dtype through a roundtrip.""" + sdfg = dace.SDFG('registered') + sdfg.add_symbol('N', dace.uint64) + sdfg.add_array('A', (symbolic.symbol('N', dace.uint64), ), dace.float64) + + s1, s2 = _resave(sdfg) + assert s1 == s2 + assert 'symbol($N, dtype=dace.uint64)' in s1 + + restored = dace.SDFG.from_json(json.loads(s1)) + assert restored.symbols['N'] == dace.uint64 + shape_sym = next(iter(restored.arrays['A'].shape[0].free_symbols)) + assert shape_sym.dtype == dace.uint64 + + +def test_map_iterator_dtype_follows_scope(): + """A map iterator's dtype is inferred from its range under the scope authority and round-trips.""" + sdfg = dace.SDFG('mapscope') + sdfg.add_symbol('N', dace.int64) + sdfg.add_array('A', (symbolic.symbol('N', dace.int64), ), dace.float64) + state = sdfg.add_state() + state.add_mapped_tasklet('m', + map_ranges={'i': '0:N'}, + inputs={}, + code='b = 1.0', + outputs={'b': dace.Memlet(data='A', subset='i')}, + external_edges=True) + + s1, s2 = _resave(sdfg) + assert s1 == s2 + # result_type_of(int32 literal, int64 N) == int64, so the iterator is int64 in scope. + assert 'symbol($i, dtype=dace.int64)' in s1 + + +def test_nested_symbol_mapping_referencing_outer_map_param_roundtrips(): + """The map_fission ``test_dependent_symbol`` shape: an outer map param used in a nested SDFG's + ``symbol_mapping`` serializes deterministically (regression for the flaky ``$i`` roundtrip).""" + sdfg = dace.SDFG('outer') + sdfg.add_symbol('fidx', dace.int32) + sdfg.add_symbol('lidx', dace.int32) + sdfg.add_array('A', (2, 10), dace.int32) + sdfg.add_array('B', (2, 10), dace.int32) + + inner = dace.SDFG('inner') + inner.add_symbol('first', dace.int32) + inner.add_symbol('last', dace.int32) + inner.add_array('A0', (10, ), dace.int32) + inner.add_array('B0', (10, ), dace.int32) + istate = inner.add_state('s', is_start_block=True) + istate.add_mapped_tasklet('plus', + map_ranges={'j': 'first:last'}, + inputs={'__a': dace.Memlet(data='A0', subset='j')}, + code='__b = __a + 1', + outputs={'__b': dace.Memlet(data='B0', subset='j')}, + external_edges=True) + + state = sdfg.add_state('outer', is_start_block=True) + a = state.add_access('A') + b = state.add_access('B') + me, mx = state.add_map('map', {'i': '0:2'}) + nsdfg = state.add_nested_sdfg(inner, {'A0'}, {'B0'}, + symbol_mapping={ + 'first': 'max(0, i - fidx)', + 'last': 'min(10, i + lidx)' + }) + state.add_memlet_path(a, me, nsdfg, memlet=dace.Memlet(data='A', subset='0, 0:10'), dst_conn='A0') + state.add_memlet_path(nsdfg, mx, b, memlet=dace.Memlet(data='B', subset='0, 0:10'), src_conn='B0') + + s1, s2 = _resave(sdfg) + assert s1 == s2 + + +if __name__ == '__main__': + import pytest + pytest.main([__file__, '-v']) diff --git a/tests/sdfg/cutout_test.py b/tests/sdfg/cutout_test.py index 99009ba82d..7db0559561 100644 --- a/tests/sdfg/cutout_test.py +++ b/tests/sdfg/cutout_test.py @@ -1,8 +1,7 @@ # Copyright 2019-2022 ETH Zurich and the DaCe authors. All rights reserved. import numpy as np import dace -from dace.sdfg.analysis.cutout import SDFGCutout, _reduce_in_configuration -import pytest +from dace.sdfg.analysis.cutout import SDFGCutout def test_cutout_onenode(): diff --git a/tests/sdfg/data/container_array_test.py b/tests/sdfg/data/container_array_test.py index 2773b5fec5..18a4fd1c92 100644 --- a/tests/sdfg/data/container_array_test.py +++ b/tests/sdfg/data/container_array_test.py @@ -1,7 +1,6 @@ # Copyright 2019-2023 ETH Zurich and the DaCe authors. All rights reserved. import ctypes -import pytest import dace import numpy as np @@ -126,16 +125,23 @@ def test_write_struct_array(): if_body.add_edge(indices, 'views', vcsr, None, dace.Memlet(data='vcsr.indices', subset='0:nnz')) if_body.add_edge(vcsr, 'views', B, None, dace.Memlet(data='B', subset='k')) # Make For Loop for j - j_before, j_guard, j_after = sdfg.add_loop(None, - if_before, - None, - 'j', - '0', - 'j < N', - 'j + 1', - loop_end_state=if_after) + j_before, _, j_after = sdfg.add_loop_state_machine(None, + if_before, + None, + 'j', + '0', + 'j < N', + 'j + 1', + loop_end_state=if_after) # Make For Loop for i - i_before, i_guard, i_after = sdfg.add_loop(None, j_before, None, 'i', '0', 'i < M', 'i + 1', loop_end_state=j_after) + i_before, i_guard, i_after = sdfg.add_loop_state_machine(None, + j_before, + None, + 'i', + '0', + 'i < M', + 'i + 1', + loop_end_state=j_after) sdfg.start_state = sdfg.node_id(i_before) i_before_guard = sdfg.edges_between(i_before, i_guard)[0] i_before_guard.data.assignments['idx'] = '0' @@ -154,7 +160,7 @@ def test_write_struct_array(): i_after.add_edge(indptr, 'views', vcsr, None, dace.Memlet(data='vcsr.indptr', subset='0:M+1')) i_after.add_edge(vcsr, 'views', B, None, dace.Memlet(data='B', subset='k')) - k_before, k_guard, k_after = sdfg.add_loop(None, i_before, None, 'k', '0', 'k < L', 'k + 1', loop_end_state=i_after) + sdfg.add_loop(None, i_before, None, 'k', '0', 'k < L', 'k + 1', loop_end_block=i_after) func = sdfg.compile() diff --git a/tests/sdfg/data/structure_test.py b/tests/sdfg/data/structure_test.py index caf32dfcfb..48bd1bf480 100644 --- a/tests/sdfg/data/structure_test.py +++ b/tests/sdfg/data/structure_test.py @@ -1,10 +1,7 @@ # Copyright 2019-2023 ETH Zurich and the DaCe authors. All rights reserved. import dace import numpy as np -import pytest -from dace import serialize -from dace.properties import make_properties from scipy import sparse @@ -98,16 +95,23 @@ def test_write_structure(): if_body.add_edge(t, '__out', indices, None, dace.Memlet(data='vindices', subset='idx')) if_body.add_edge(indices, 'views', B, None, dace.Memlet(data='B.indices', subset='0:nnz')) # Make For Loop for j - j_before, j_guard, j_after = sdfg.add_loop(None, - if_before, - None, - 'j', - '0', - 'j < N', - 'j + 1', - loop_end_state=if_after) + j_before, _, j_after = sdfg.add_loop_state_machine(None, + if_before, + None, + 'j', + '0', + 'j < N', + 'j + 1', + loop_end_state=if_after) # Make For Loop for i - i_before, i_guard, i_after = sdfg.add_loop(None, j_before, None, 'i', '0', 'i < M', 'i + 1', loop_end_state=j_after) + i_before, i_guard, i_after = sdfg.add_loop_state_machine(None, + j_before, + None, + 'i', + '0', + 'i < M', + 'i + 1', + loop_end_state=j_after) sdfg.start_state = sdfg.node_id(i_before) i_before_guard = sdfg.edges_between(i_before, i_guard)[0] i_before_guard.data.assignments['idx'] = '0' @@ -183,16 +187,23 @@ def test_local_structure(): if_body.add_edge(t, '__out', indices, None, dace.Memlet(data='tmp_vindices', subset='idx')) if_body.add_edge(indices, 'views', tmp, None, dace.Memlet(data='tmp.indices', subset='0:nnz')) # Make For Loop for j - j_before, j_guard, j_after = sdfg.add_loop(None, - if_before, - None, - 'j', - '0', - 'j < N', - 'j + 1', - loop_end_state=if_after) + j_before, _, j_after = sdfg.add_loop_state_machine(None, + if_before, + None, + 'j', + '0', + 'j < N', + 'j + 1', + loop_end_state=if_after) # Make For Loop for i - i_before, i_guard, i_after = sdfg.add_loop(None, j_before, None, 'i', '0', 'i < M', 'i + 1', loop_end_state=j_after) + i_before, i_guard, i_after = sdfg.add_loop_state_machine(None, + j_before, + None, + 'i', + '0', + 'i < M', + 'i + 1', + loop_end_state=j_after) sdfg.start_state = sdfg.node_id(i_before) i_before_guard = sdfg.edges_between(i_before, i_guard)[0] i_before_guard.data.assignments['idx'] = '0' @@ -347,16 +358,23 @@ def test_write_nested_structure(): if_body.add_edge(t, '__out', indices, None, dace.Memlet(data='vindices', subset='idx')) if_body.add_edge(indices, 'views', B, None, dace.Memlet(data='B.csr.indices', subset='0:nnz')) # Make For Loop for j - j_before, j_guard, j_after = sdfg.add_loop(None, - if_before, - None, - 'j', - '0', - 'j < N', - 'j + 1', - loop_end_state=if_after) + j_before, j_guard, j_after = sdfg.add_loop_state_machine(None, + if_before, + None, + 'j', + '0', + 'j < N', + 'j + 1', + loop_end_state=if_after) # Make For Loop for i - i_before, i_guard, i_after = sdfg.add_loop(None, j_before, None, 'i', '0', 'i < M', 'i + 1', loop_end_state=j_after) + i_before, i_guard, i_after = sdfg.add_loop_state_machine(None, + j_before, + None, + 'i', + '0', + 'i < M', + 'i + 1', + loop_end_state=j_after) sdfg.start_state = sdfg.node_id(i_before) i_before_guard = sdfg.edges_between(i_before, i_guard)[0] i_before_guard.data.assignments['idx'] = '0' @@ -465,16 +483,8 @@ def test_direct_read_structure_loops(): state.add_edge(data, None, t, '__val', dace.Memlet(data='A.data', subset='idx')) state.add_edge(t, '__out', B, None, dace.Memlet(data='B', subset='0:M, 0:N', volume=1)) - idx_before, idx_guard, idx_after = sdfg.add_loop(None, state, None, 'idx', 'A.indptr[i]', 'idx < A.indptr[i+1]', - 'idx + 1') - i_before, i_guard, i_after = sdfg.add_loop(None, - idx_before, - None, - 'i', - '0', - 'i < M', - 'i + 1', - loop_end_state=idx_after) + idx_loop = sdfg.add_loop(None, state, None, 'idx', 'A.indptr[i]', 'idx < A.indptr[i+1]', 'idx + 1') + sdfg.add_loop(None, idx_loop, None, 'i', '0', 'i < M', 'i + 1') func = sdfg.compile() diff --git a/tests/sdfg/data/tensor_test.py b/tests/sdfg/data/tensor_test.py index 7ceda5557f..2dfbb8fb6e 100644 --- a/tests/sdfg/data/tensor_test.py +++ b/tests/sdfg/data/tensor_test.py @@ -1,7 +1,6 @@ # Copyright 2019-2023 ETH Zurich and the DaCe authors. All rights reserved. import dace import numpy as np -import pytest from scipy import sparse diff --git a/tests/sdfg/loop_region_test.py b/tests/sdfg/loop_region_test.py index f09bf0004f..b6bfc2c227 100644 --- a/tests/sdfg/loop_region_test.py +++ b/tests/sdfg/loop_region_test.py @@ -1,4 +1,4 @@ -# Copyright 2019-2023 ETH Zurich and the DaCe authors. All rights reserved. +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. import dace import numpy as np from dace.sdfg.sdfg import SDFG @@ -311,7 +311,7 @@ def test_loop_to_stree_triple_nested_for(): stree = s2t.as_schedule_tree(sdfg) po_nodes = list(stree.preorder_traversal())[1:] - assert [type(n) for n in po_nodes] == [tn.LoopScope, tn.LoopScope, tn.LoopScope, tn.TaskletNode, tn.LibraryCall] + assert [type(n) for n in po_nodes] == [tn.ForScope, tn.ForScope, tn.ForScope, tn.TaskletNode, tn.LibraryCall] if __name__ == '__main__': diff --git a/tests/sdfg/memlet_utils_test.py b/tests/sdfg/memlet_utils_test.py index 3c85d72f21..8228481c36 100644 --- a/tests/sdfg/memlet_utils_test.py +++ b/tests/sdfg/memlet_utils_test.py @@ -1,10 +1,10 @@ -# Copyright 2019-2023 ETH Zurich and the DaCe authors. All rights reserved. +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. +import copy import dace import numpy as np import pytest -from dace import symbolic -from dace.sdfg import memlet_utils as mu +from dace.sdfg import graph, memlet_utils as mu import re from typing import Tuple, Optional @@ -12,6 +12,7 @@ def _replace_zero_with_one(memlet: dace.Memlet) -> dace.Memlet: if not isinstance(memlet.subset, dace.subsets.Range): return memlet + for i, (rb, re, rs) in enumerate(memlet.subset.ndrange()): if rb == 0: memlet.subset.ranges[i] = (1, 1, rs) @@ -19,7 +20,7 @@ def _replace_zero_with_one(memlet: dace.Memlet) -> dace.Memlet: @pytest.mark.parametrize('filter_type', ['none', 'same_array', 'different_array']) -def test_replace_memlet(filter_type): +def test_replace_memlet(filter_type: str) -> None: # Prepare SDFG sdfg = dace.SDFG('replace_memlet') sdfg.add_array('A', [2, 2], dace.float64) @@ -66,7 +67,7 @@ def test_replace_memlet(filter_type): assert B[0] == 1 -def _perform_non_lin_delin_test(sdfg: dace.SDFG, edge) -> bool: +def _perform_non_lin_delin_test(sdfg: dace.SDFG, edge: graph.MultiConnectorEdge) -> None: assert sdfg.number_of_nodes() == 1 state: dace.SDFGState = sdfg.states()[0] assert state.number_of_nodes() == 2 @@ -104,8 +105,6 @@ def _perform_non_lin_delin_test(sdfg: dace.SDFG, edge) -> bool: sdfg(a=a, b=b_opt) assert np.allclose(b_unopt, b_opt) - return True - def _make_non_lin_delin_sdfg( shape_a: Tuple[int, ...], @@ -211,6 +210,62 @@ def test_non_lin_delin_8(): _perform_non_lin_delin_test(sdfg, e) +def test_memlet_set() -> None: + empty_set = mu.MemletSet() + assert len(empty_set) == 0 + + memlet_set = mu.MemletSet([dace.Memlet("A[0:5]")]) + covered_set = mu.MemletSet([dace.Memlet("A[0:5]")], intersection_is_contained=False) + + assert dace.Memlet("A[0:2]") in memlet_set + assert dace.Memlet("A[0:2]") in covered_set + assert dace.Memlet("A[2:7]") in memlet_set + assert dace.Memlet("A[2:7]") not in covered_set + + assert dace.Memlet("B[0:2]") not in memlet_set + + before = copy.deepcopy(covered_set.internal_set) + covered_set.add(dace.Memlet("A[0:2]")) + assert covered_set.internal_set == before + + covered_set.add(dace.Memlet("A[4:9]")) + assert dace.Memlet("A[2:7]") in covered_set + assert covered_set.internal_set != before + + union = empty_set.union(dace.Memlet("A[0:3]"), dace.Memlet("A[2:10]")) + assert dace.Memlet("A[5:7]") not in empty_set + assert dace.Memlet("A[5:7]") in union + assert len(union.internal_set["A"]) == 1 + internal_memlet = list(union.internal_set["A"])[0] + assert internal_memlet.subset == dace.subsets.Range.from_string("0:10") + + +def test_memlet_dict() -> None: + A_01 = dace.Memlet("A[0:1]") + A_02 = dace.Memlet("A[0:2]") + A_34 = dace.Memlet("A[3:4]") + memlet_dict: mu.Memlet[list[int]] = mu.MemletDict() + assert len(memlet_dict) == 0 + assert A_02 not in memlet_dict + + memlet_dict[A_02] = [42] + assert A_02 in memlet_dict + assert A_01 in memlet_dict + assert A_34 not in memlet_dict + assert dace.Memlet("B[0:2]") not in memlet_dict + + memlet_dict[A_01].append(43) + assert memlet_dict[A_02] == [42, 43] + + memlet_dict.update({A_34: [0], A_01: [44]}) + assert A_34 in memlet_dict + assert memlet_dict[A_02] == [44] # @Tal this is expected, right? + assert memlet_dict[A_34] == [0] + + memlet_dict.clear() + assert len(memlet_dict) == 0 + + if __name__ == '__main__': test_replace_memlet('none') test_replace_memlet('same_array') @@ -224,3 +279,6 @@ def test_non_lin_delin_8(): test_non_lin_delin_6() test_non_lin_delin_7() test_non_lin_delin_8() + + test_memlet_set() + test_memlet_dict() diff --git a/tests/sdfg/nested_sdfg_deepcopy_test.py b/tests/sdfg/nested_sdfg_deepcopy_test.py index 3dbdcba2fd..4c59fbb7f6 100644 --- a/tests/sdfg/nested_sdfg_deepcopy_test.py +++ b/tests/sdfg/nested_sdfg_deepcopy_test.py @@ -2,7 +2,6 @@ """ Tests deepcopying (nested) SDFGs. """ import copy import dace -import numpy as np def test_deepcopy_same_state(): diff --git a/tests/sdfg/reference_test.py b/tests/sdfg/reference_test.py index 79b4ceff00..79079464e0 100644 --- a/tests/sdfg/reference_test.py +++ b/tests/sdfg/reference_test.py @@ -1,15 +1,30 @@ # Copyright 2019-2023 ETH Zurich and the DaCe authors. All rights reserved. """ Tests the use of Reference data descriptors. """ +import json +import tempfile + import dace from dace.sdfg import validation from dace.transformation.pass_pipeline import Pipeline -from dace.transformation.passes.analysis import FindReferenceSources +from dace.transformation.helpers import modified_symbols_between +from dace.transformation.passes.analysis import ControlFlowBlockReachability, FindReferenceSources, StateReachability from dace.transformation.passes.reference_reduction import ReferenceToView import numpy as np import pytest import networkx as nx +def _assert_roundtrip_json_stable(sdfg): + with tempfile.TemporaryDirectory() as tmpdir: + path1 = f'{tmpdir}/test.sdfg' + path2 = f'{tmpdir}/test2.sdfg' + sdfg.save(path1, hash=False) + dace.SDFG.from_file(path1).save(path2, hash=False) + + with open(path1, 'r') as fp1, open(path2, 'r') as fp2: + assert json.load(fp1) == json.load(fp2) + + def test_frontend_reference(): N = dace.symbol('N') M = dace.symbol('M') @@ -336,7 +351,7 @@ def _create_loop_reference_internal_use(): state = sdfg.add_state() after = sdfg.add_state() sdfg.add_edge(state, after, dace.InterstateEdge()) - sdfg.add_loop(istate, state, None, 'i', '0', 'i < 20', 'i + 1', loop_end_state=after) + sdfg.add_loop(istate, state, None, 'i', '0', 'i < 20', 'i + 1', loop_end_block=after) # Reference set inside loop state.add_edge(state.add_read('A'), None, state.add_write('ref'), 'set', dace.Memlet('A[i]')) @@ -363,11 +378,11 @@ def _create_loop_reference_nonfree_internal_use(): # First loop state1 = sdfg.add_state() - sdfg.add_loop(istate, state1, between_loops, 'i', '0', 'i < 20', 'i + 1') + sdfg.add_loop(istate, state1, between_loops, 'i', '0', 'i < 20', 'i + 1', label='set_loop') # Second loop state2 = sdfg.add_state() - sdfg.add_loop(between_loops, state2, None, 'i', '0', 'i < 20', 'i + 1') + sdfg.add_loop(between_loops, state2, None, 'i', '0', 'i < 20', 'i + 1', label='use_loop') # Reference set inside first loop state1.add_edge(state1.add_read('A'), None, state1.add_write('ref'), 'set', dace.Memlet('A[i]')) @@ -523,6 +538,10 @@ def test_scoped(reftoview): assert np.allclose(A, ref) +def test_scoped_roundtrip_is_stable(): + _assert_roundtrip_json_stable(_create_scoped_sdfg()) + + @pytest.mark.parametrize('reftoview', (False, True)) def test_scoped_empty_memlet(reftoview): sdfg = _create_scoped_empty_memlet_sdfg() @@ -646,6 +665,29 @@ def test_reference_loop_nonfree_internal_use(): assert np.allclose(ref, A) +def test_reachability_across_sibling_loops(): + sdfg = _create_loop_reference_nonfree_internal_use() + blocks = {block.label: block for block in sdfg.all_control_flow_blocks()} + + block_reach = ControlFlowBlockReachability().apply_pass(sdfg, {}) + assert blocks['block_0'] in block_reach[blocks['block_1'].parent_graph.cfg_id][blocks['block_1']] + assert blocks['block_2'] in block_reach[blocks['block_1'].parent_graph.cfg_id][blocks['block_1']] + + state_reach = StateReachability().apply_pass(sdfg, {}) + assert blocks['block_0'] in state_reach[sdfg.cfg_id][blocks['block_1']] + assert blocks['block_2'] in state_reach[sdfg.cfg_id][blocks['block_1']] + + +def test_modified_symbols_between_control_flow_blocks(): + sdfg = _create_loop_reference_internal_use() + blocks = {block.label: block for block in sdfg.all_control_flow_blocks()} + assert modified_symbols_between(blocks['block_0'], blocks['block_1']) == set() + + sdfg = _create_loop_reference_nonfree_internal_use() + blocks = {block.label: block for block in sdfg.all_control_flow_blocks()} + assert modified_symbols_between(blocks['block_1'], blocks['block_2']) == {'i'} + + @pytest.mark.parametrize(('array_outside_scope', 'depends_on_iterate'), ((False, True), (False, True))) def test_ref2view_refset_in_scope(array_outside_scope, depends_on_iterate): sdfg = dace.SDFG('reftest') @@ -762,6 +804,7 @@ def test_ref2view_reconnection(): test_multisubset(True) test_scoped(False) test_scoped(True) + test_scoped_roundtrip_is_stable() test_scoped_empty_memlet(False) test_scoped_empty_memlet(True) test_reference_neighbors(False) @@ -775,3 +818,4 @@ def test_ref2view_reconnection(): test_ref2view_refset_in_scope(True, False) test_ref2view_refset_in_scope(True, True) test_ref2view_reconnection() + test_modified_symbols_between_control_flow_blocks() diff --git a/tests/sdfg/scalar_return.py b/tests/sdfg/scalar_return.py index 3a4f325291..aba6151ae9 100644 --- a/tests/sdfg/scalar_return.py +++ b/tests/sdfg/scalar_return.py @@ -2,7 +2,6 @@ import dace import numpy as np import pytest -from typing import Tuple from dace.sdfg.validation import InvalidSDFGError diff --git a/tests/sdfg/schedule_inference_test.py b/tests/sdfg/schedule_inference_test.py index 8f4fcd6acb..f40d3a9acc 100644 --- a/tests/sdfg/schedule_inference_test.py +++ b/tests/sdfg/schedule_inference_test.py @@ -1,7 +1,6 @@ # Copyright 2019-2023 ETH Zurich and the DaCe authors. All rights reserved. """ Tests for default storage/schedule inference. """ import dace -from dace.sdfg.state import SDFGState from dace.sdfg.validation import InvalidSDFGNodeError from dace.sdfg.infer_types import set_default_schedule_and_storage_types from dace.transformation.helpers import get_parent_map diff --git a/tests/sdfg_validate_names_test.py b/tests/sdfg_validate_names_test.py index ac9f633496..a74627e874 100644 --- a/tests/sdfg_validate_names_test.py +++ b/tests/sdfg_validate_names_test.py @@ -1,117 +1,122 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. -import unittest +import pytest import dace - # Try to detect invalid names in SDFG -class NameValidationTests(unittest.TestCase): - # SDFG label - def test_sdfg_name1(self): - try: - sdfg = dace.SDFG(' ') - sdfg.validate() - self.fail('Failed to detect invalid SDFG') - except dace.sdfg.InvalidSDFGError as ex: - print('Exception caught:', ex) - - def test_sdfg_name2(self): - try: - sdfg = dace.SDFG('3sat') - sdfg.validate() - self.fail('Failed to detect invalid SDFG') - except dace.sdfg.InvalidSDFGError as ex: - print('Exception caught:', ex) - - # State - def test_state_duplication(self): - try: - sdfg = dace.SDFG('ok') - s1 = sdfg.add_state('also_ok') - s2 = sdfg.add_state('also_ok') - s2.label = 'also_ok' - sdfg.add_edge(s1, s2, dace.InterstateEdge()) - sdfg.validate() - self.fail('Failed to detect duplicate state') - except dace.sdfg.InvalidSDFGError as ex: - print('Exception caught:', ex) - - def test_state_name1(self): - try: - sdfg = dace.SDFG('ok') - sdfg.add_state('not ok') - sdfg.validate() - self.fail('Failed to detect invalid state') - except dace.sdfg.InvalidSDFGError as ex: - print('Exception caught:', ex) - - def test_state_name2(self): - try: - sdfg = dace.SDFG('ok') - sdfg.add_state('$5') - sdfg.validate() - self.fail('Failed to detect invalid state') - except dace.sdfg.InvalidSDFGError as ex: - print('Exception caught:', ex) - - # Array - def test_array(self): - try: - sdfg = dace.SDFG('ok') - state = sdfg.add_state('also_ok') - _8 = state.add_array('8', [1], dace.float32) - t = state.add_tasklet('tasklet', {'a'}, {}, 'print(a)') - state.add_edge(_8, None, t, 'a', dace.Memlet.from_array(_8.data, _8.desc(sdfg))) - sdfg.validate() - self.fail('Failed to detect invalid array name') - except (dace.sdfg.InvalidSDFGError, NameError) as ex: - print('Exception caught:', ex) - - # Tasklet - def test_tasklet(self): - try: - sdfg = dace.SDFG('ok') - state = sdfg.add_state('also_ok') - A = state.add_array('A', [1], dace.float32) - B = state.add_array('B', [1], dace.float32) - t = state.add_tasklet(' tasklet', {'a'}, {'b'}, 'b = a') - state.add_edge(A, None, t, 'a', dace.Memlet.from_array(A.data, A.desc(sdfg))) - state.add_edge(t, 'b', B, None, dace.Memlet.from_array(B.data, B.desc(sdfg))) - sdfg.validate() - self.fail('Failed to detect invalid tasklet name') - except dace.sdfg.InvalidSDFGNodeError as ex: - print('Exception caught:', ex) - - # Connector - def test_connector(self): - try: - sdfg = dace.SDFG('ok') - state = sdfg.add_state('also_ok') - A = state.add_array('A', [1], dace.float32) - B = state.add_array('B', [1], dace.float32) - t = state.add_tasklet('tasklet', {'$a'}, {' b'}, '') - state.add_edge(A, None, t, '$a', dace.Memlet.from_array(A.data, A.desc(sdfg))) - state.add_edge(t, ' b', B, None, dace.Memlet.from_array(B.data, B.desc(sdfg))) - sdfg.validate() - self.fail('Failed to detect invalid connectors') - except dace.sdfg.InvalidSDFGError as ex: - print('Exception caught:', ex) - - # Interstate edge - def test_interstate_edge(self): - try: - sdfg = dace.SDFG('ok') - state = sdfg.add_state('also_ok', is_start_block=True) - A = state.add_array('A', [1], dace.float32) - B = state.add_array('B', [1], dace.float32) - t = state.add_tasklet('tasklet', {'a'}, {'b'}, 'b = a') - state.add_edge(A, None, t, 'a', dace.Memlet.from_array(A.data, A.desc(sdfg))) - state.add_edge(t, 'b', B, None, dace.Memlet.from_array(B.data, B.desc(sdfg))) - sdfg.add_edge(state, state, dace.InterstateEdge(assignments={'%5': '1'})) - sdfg.validate() - self.fail('Failed to detect invalid interstate edge') - except dace.sdfg.InvalidSDFGInterstateEdgeError as ex: - print('Exception caught:', ex) + + +# SDFG label +def test_sdfg_name1(): + with pytest.raises(dace.sdfg.InvalidSDFGError): + dace.SDFG(' ') + + +def test_sdfg_name2(): + with pytest.raises(dace.sdfg.InvalidSDFGError): + dace.SDFG('3sat') + + +# State +def test_state_duplication(): + sdfg = dace.SDFG('ok') + s1 = sdfg.add_state('also_ok') + s2 = sdfg.add_state('also_ok') + s2.label = 'also_ok' + sdfg.add_edge(s1, s2, dace.InterstateEdge()) + + with pytest.raises(dace.sdfg.InvalidSDFGError): + sdfg.validate() + + +def test_state_name1(): + sdfg = dace.SDFG('ok') + sdfg.add_state('not ok') + + with pytest.raises(dace.sdfg.InvalidSDFGError): + sdfg.validate() + + +def test_state_name2(): + sdfg = dace.SDFG('ok') + sdfg.add_state('$5') + + with pytest.raises(dace.sdfg.InvalidSDFGError): + sdfg.validate() + + +# Array +def test_array(): + sdfg = dace.SDFG('ok') + sdfg.add_array('8', [1], dace.float32) + + state = sdfg.add_state('also_ok') + _8 = state.add_access('8') + t = state.add_tasklet('tasklet', {'a'}, {}, 'print(a)') + state.add_edge(_8, None, t, 'a', dace.Memlet.from_array(_8.data, _8.desc(sdfg))) + + with pytest.raises(dace.sdfg.InvalidSDFGError): + sdfg.validate() + + +# Tasklet +def test_tasklet(): + sdfg = dace.SDFG('ok') + sdfg.add_array('A', [1], dace.float32) + sdfg.add_array('B', [1], dace.float32) + + state = sdfg.add_state('also_ok') + A = state.add_access('A') + B = state.add_access('B') + t = state.add_tasklet(' tasklet', {'a'}, {'b'}, 'b = a') + state.add_edge(A, None, t, 'a', dace.Memlet.from_array(A.data, A.desc(sdfg))) + state.add_edge(t, 'b', B, None, dace.Memlet.from_array(B.data, B.desc(sdfg))) + + with pytest.raises(dace.sdfg.InvalidSDFGError): + sdfg.validate() + + +# Connector +def test_connector(): + sdfg = dace.SDFG('ok') + sdfg.add_array('A', [1], dace.float32) + sdfg.add_array('B', [1], dace.float32) + + state = sdfg.add_state('also_ok') + A = state.add_access('A') + B = state.add_access('B') + t = state.add_tasklet('tasklet', {'$a'}, {' b'}, '') + state.add_edge(A, None, t, '$a', dace.Memlet.from_array(A.data, A.desc(sdfg))) + state.add_edge(t, ' b', B, None, dace.Memlet.from_array(B.data, B.desc(sdfg))) + + with pytest.raises(dace.sdfg.InvalidSDFGError): + sdfg.validate() + + +# Interstate edge +def test_interstate_edge(): + sdfg = dace.SDFG('ok') + sdfg.add_array('A', [1], dace.float32) + sdfg.add_array('B', [1], dace.float32) + + state = sdfg.add_state('also_ok', is_start_block=True) + A = state.add_access('A') + B = state.add_access('B') + t = state.add_tasklet('tasklet', {'a'}, {'b'}, 'b = a') + state.add_edge(A, None, t, 'a', dace.Memlet.from_array(A.data, A.desc(sdfg))) + state.add_edge(t, 'b', B, None, dace.Memlet.from_array(B.data, B.desc(sdfg))) + sdfg.add_edge(state, state, dace.InterstateEdge(assignments={'%5': '1'})) + + with pytest.raises(dace.sdfg.InvalidSDFGError): + sdfg.validate() if __name__ == '__main__': - unittest.main() + test_sdfg_name1() + test_sdfg_name2() + test_state_duplication() + test_state_name1() + test_state_name2() + test_array() + test_tasklet() + test_connector() + test_interstate_edge() diff --git a/tests/sdfg_validate_scopes_test.py b/tests/sdfg_validate_scopes_test.py index db226dee80..99c4f496dd 100644 --- a/tests/sdfg_validate_scopes_test.py +++ b/tests/sdfg_validate_scopes_test.py @@ -1,30 +1,27 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. -import unittest +import pytest import dace -# Try to detect invalid scopes in SDFG -class ScopeValidationTests(unittest.TestCase): +def test_connector_mismatch(): + """Try to detect invalid scopes in SDFG""" + sdfg = dace.SDFG('a') + sdfg.add_array('A', [1], dace.float32) - def test_connector_mismatch(self): - try: - sdfg = dace.SDFG('a') - state = sdfg.add_state() - me, mx = state.add_map('b', dict(i="0:1")) - A = state.add_array('A', [1], dace.float32) - T = state.add_tasklet('T', {'a'}, {}, 'printf("%f", a)') + state = sdfg.add_state() + me, mx = state.add_map('b', dict(i="0:1")) + A = state.add_access('A') + T = state.add_tasklet('T', {'a'}, {}, 'printf("%f", a)') - me.add_in_connector('IN_a') - me.add_out_connector('OUT_b') - state.add_edge(A, None, me, 'IN_a', dace.Memlet.from_array(A.data, A.desc(sdfg))) - state.add_edge(me, 'OUT_b', T, 'a', dace.Memlet.simple(A, '0')) - state.add_edge(T, None, mx, None, dace.Memlet()) + me.add_in_connector('IN_a') + me.add_out_connector('OUT_b') + state.add_edge(A, None, me, 'IN_a', dace.Memlet.from_array(A.data, A.desc(sdfg))) + state.add_edge(me, 'OUT_b', T, 'a', dace.Memlet.simple(A, '0')) + state.add_edge(T, None, mx, None, dace.Memlet()) - sdfg.validate() - self.fail('Failed to detect invalid SDFG') - except dace.sdfg.InvalidSDFGError as ex: - print('Exception caught:', ex) + with pytest.raises(dace.sdfg.InvalidSDFGError): + sdfg.validate() if __name__ == '__main__': - unittest.main() + test_connector_mismatch() diff --git a/tests/simplify_cache_typed_test.py b/tests/simplify_cache_typed_test.py new file mode 100644 index 0000000000..f568f4c404 --- /dev/null +++ b/tests/simplify_cache_typed_test.py @@ -0,0 +1,110 @@ +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. +""" +Regressions tests to ensure LRU cache is type-aware +""" + +import sympy +from sympy.logic.boolalg import Boolean + +import dace + + +def _fresh_cache() -> None: + """Clear the ``simplify`` cache so each test is order independent.""" + dace.symbolic.simplify.cache_clear() + + +def test_simplify_bool_does_not_poison_integer_one(): + """``simplify(True)`` first must not turn ``simplify(Integer(1))`` boolean.""" + _fresh_cache() + # Caches sympy.true; on an untyped cache this poisons the key ``1``. + dace.symbolic.simplify(True) + + result = dace.symbolic.simplify(sympy.Integer(1)) + assert isinstance(result, (sympy.Integer, int)), \ + f"expected sympy.Integer, got {type(result).__name__}: {result!r}" + assert not isinstance(result, (Boolean, bool)) + assert result == 1 + + +def test_simplify_integer_one_does_not_poison_bool(): + """The reverse order: a cached ``Integer(1)`` must not leak into ``simplify(True)``.""" + _fresh_cache() + dace.symbolic.simplify(sympy.Integer(1)) + + result = dace.symbolic.simplify(True) + assert isinstance(result, (Boolean, bool)), \ + f"expected a sympy boolean, got {type(result).__name__}: {result!r}" + assert bool(result) is True + + +def test_simplify_bool_does_not_poison_integer_zero(): + """Same conflation exists for ``False``/``0`` -- integer direction stays typed.""" + _fresh_cache() + dace.symbolic.simplify(False) + + result = dace.symbolic.simplify(sympy.Integer(0)) + assert isinstance(result, (sympy.Integer, int)), \ + f"expected sympy.Integer, got {type(result).__name__}: {result!r}" + assert not isinstance(result, (Boolean, bool)) + assert result == 0 + + +def test_simplify_integer_zero_does_not_poison_bool(): + """``False``/``0`` -- boolean direction stays boolean.""" + _fresh_cache() + dace.symbolic.simplify(sympy.Integer(0)) + + result = dace.symbolic.simplify(False) + assert isinstance(result, (Boolean, bool)), \ + f"expected a sympy boolean, got {type(result).__name__}: {result!r}" + assert bool(result) is False + + +def test_simplify_symbolic_expressions_still_cache(): + """Normal symbolic inputs are unaffected and are still served from the cache.""" + _fresh_cache() + n = sympy.Symbol('N') + assert dace.symbolic.simplify(n + 1 - n) == 1 + assert dace.symbolic.simplify((n**2 - 1) / (n - 1)) == n + 1 + # Same expression again: served from the cache with the same result. + assert dace.symbolic.simplify(n + 1 - n) == 1 + assert dace.symbolic.simplify.cache_info().hits >= 1 + + +def test_volume_propagation_after_bool_simplify(): + """End-to-end: a ``bool`` fed to ``simplify`` must not break memlet-volume + propagation on an unrelated SDFG (the original order-dependent crash).""" + from dace.sdfg.propagation import propagate_memlets_sdfg + + _fresh_cache() + # The poisoning call pattern: a concrete shape comparison yields a Python + # ``bool``, which simplifies to ``sympy.true``. + poison = dace.symbolic.simplify(sympy.Integer(1) == 1) + assert bool(poison) is True + + sdfg = dace.SDFG('simplify_cache_volume_repro') + sdfg.add_array('A', [4], dace.float64) + sdfg.add_array('B', [4], dace.float64) + state = sdfg.add_state() + # A single-iteration map: the propagated volume is ``sympy.Integer(1)`` -- + # exactly the cache key that a cached ``simplify(True)`` poisons. + state.add_mapped_tasklet('copy', + dict(i='0:1'), + dict(inp=dace.Memlet('A[i]')), + 'out = inp', + dict(out=dace.Memlet('B[i]')), + external_edges=True) + # Without the fix this raised: TypeError: Property volume must be a literal + # or symbolic expression (the propagated volume was a BooleanTrue). + propagate_memlets_sdfg(sdfg) + sdfg.validate() + + +if __name__ == '__main__': + test_simplify_bool_does_not_poison_integer_one() + test_simplify_integer_one_does_not_poison_bool() + test_simplify_bool_does_not_poison_integer_zero() + test_simplify_integer_zero_does_not_poison_bool() + test_simplify_symbolic_expressions_still_cache() + test_volume_propagation_after_bool_simplify() diff --git a/tests/specialize_test.py b/tests/specialize_test.py index b9e91bb5df..5bc6db0065 100644 --- a/tests/specialize_test.py +++ b/tests/specialize_test.py @@ -1,5 +1,4 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. -import math import numpy as np import dace as dp @@ -7,11 +6,9 @@ from dace.memlet import Memlet -def test(): - print('Constant specialization test') - - N = dp.symbol('N') - M = dp.symbol('M') +def test_constant_specialization(): + N = dp.symbol('N', dtype=dp.int64) + M = dp.symbol('M', dtype=dp.int64) n = 20 m = 30 fullrange = '1:N-1,0:M' @@ -24,10 +21,14 @@ def test(): ########################################################################## spec_sdfg = SDFG('spectest') + spec_sdfg.add_array('A', [N, M], dp.float32) + spec_sdfg.add_transient('At', [N - 2, M], dp.float32) + spec_sdfg.add_array('B', [N, M], dp.float32) + state = spec_sdfg.add_state() - A = state.add_array('A', [N, M], dp.float32) - Atrans = state.add_transient('At', [N - 2, M], dp.float32) - B = state.add_array('B', [N, M], dp.float32) + A = state.add_access('A') + Atrans = state.add_access('At') + B = state.add_access('B') state.add_edge(A, None, Atrans, None, Memlet.simple(A, fullrange)) _, me, mx = state.add_mapped_tasklet('compute', dict(i=irange, j=jrange), dict(a=Memlet.simple(Atrans, 'i-1,j')), @@ -53,9 +54,8 @@ def test(): func(A=input, B=output, N=n, M=m) diff = np.linalg.norm(np.exp(input[1:(n - 1), 0:m]) - output[1:-1, :]) / n - print("Difference:", diff) assert diff <= 1e-5 if __name__ == "__main__": - test() + test_constant_specialization() diff --git a/tests/state_transition_test.py b/tests/state_transition_test.py index 6cc462dbb3..8429f8cc85 100644 --- a/tests/state_transition_test.py +++ b/tests/state_transition_test.py @@ -2,16 +2,16 @@ import dace import re -import sys import numpy as np +import pytest def test_state_transitions(): if not dace.config.Config.get_bool('optimizer', 'detect_control_flow'): - print("Control flow not enabled. Skipping test.") - return + pytest.skip("Control flow not enabled. Skipping test.") sdfg = dace.SDFG("Transitions") + sdfg.add_array("x", (1, ), int) start = sdfg.add_state("start") left = sdfg.add_state("left") @@ -32,8 +32,8 @@ def test_state_transitions(): s1_for_enter = sdfg.add_state("s1_for_enter") s1_for_body = sdfg.add_state("s1_for_body") - x0_for = s1_for_body.add_array("x", (1, ), int) - x1_for = s1_for_body.add_array("x", (1, ), int) + x0_for = s1_for_body.add_access("x") + x1_for = s1_for_body.add_access("x") tasklet_for = s1_for_body.add_tasklet("Update_x", {"x_in"}, {"x_out"}, "x_out = x_in + 1") s1_for_body.add_edge(x0_for, None, tasklet_for, "x_in", dace.memlet.Memlet.simple(x0_for, "0")) s1_for_body.add_edge(tasklet_for, "x_out", x1_for, None, dace.memlet.Memlet.simple(x1_for, "0")) @@ -56,8 +56,8 @@ def test_state_transitions(): s3_while_enter = sdfg.add_state("s3_while_enter") s3_while_body = sdfg.add_state("s3_while_body") - x0_while = s3_while_body.add_array("x", (1, ), int) - x1_while = s3_while_body.add_array("x", (1, ), int) + x0_while = s3_while_body.add_access("x") + x1_while = s3_while_body.add_access("x") tasklet_while = s3_while_body.add_tasklet("Update_x", {"x_in"}, {"x_out"}, "x_out = x_in * 2; i *= 2") s3_while_body.add_edge(x0_while, None, tasklet_while, "x_in", dace.memlet.Memlet.simple(x0_while, "0")) s3_while_body.add_edge(tasklet_while, "x_out", x1_while, None, dace.memlet.Memlet.simple(x1_while, "0")) @@ -81,7 +81,7 @@ def test_state_transitions(): s5_then = sdfg.add_state("s5_then") s6_else = sdfg.add_state("s6_else") - x1_else = s6_else.add_array("x", (1, ), int) + x1_else = s6_else.add_access("x") tasklet_else = s6_else.add_tasklet("Update_x", {}, {"x_out"}, "x_out = 42") s6_else.add_edge(tasklet_else, "x_out", x1_else, None, dace.memlet.Memlet.simple(x1_else, "0")) @@ -112,50 +112,47 @@ def test_state_transitions(): code = sdfg.generate_code()[0].code for_pattern = r"for.*i\s*=\s*0.*i\s*<\s*16" - if re.search(for_pattern, code) is None: - raise RuntimeError("For loop not detected in state transitions") + assert re.search(for_pattern, code) is not None, "For loop not detected in state transitions" - while_pattern = r"while.+i\s*<\s*128" - if re.search(while_pattern, code) is None: - raise RuntimeError("While loop not detected in state transitions") + while_pattern = r"(while|for).+i\s*<\s*128" + assert re.search(while_pattern, code) is not None, "While loop not detected in state transitions" if_pattern = r"if.+i\s*<\s*512" - if re.search(if_pattern, code) is None: - raise RuntimeError("If not detected in state transitions") + assert re.search(if_pattern, code) is not None, "If not detected in state transitions" else_pattern = r"}\s*else\s*{" - if re.search(else_pattern, code) is None: - raise RuntimeError("Else not detected in state transitions") + assert re.search(else_pattern, code) is not None, "Else not detected in state transitions" x_output = dace.ndarray([1], int) x_output[0] = 0 sdfg(x=x_output) x_output = x_output[0] - if x_output != 128: - raise RuntimeError("Expected x = 128, got {}".format(x_output)) + assert x_output == 128, f"Expected x = 128, got {x_output}" def test_state_transition_array(): """ Toplevel array usage in interstate edge """ sdfg = dace.SDFG('sta_test') + sdfg.add_array('inp', [1], dace.float32) + sdfg.add_array('A', [1], dace.float32) s0 = sdfg.add_state() s1 = sdfg.add_state() s2 = sdfg.add_state() # Arrays - inp = s0.add_array('inp', [1], dace.float32) - A = s0.add_array('A', [1], dace.float32) + inp = s0.add_access('inp') + A = s0.add_access('A') t = s0.add_tasklet('seta', {'a'}, {'b'}, 'b = a') s0.add_edge(inp, None, t, 'a', dace.Memlet.from_array(inp.data, inp.desc(sdfg))) s0.add_edge(t, 'b', A, None, dace.Memlet.from_array(A.data, A.desc(sdfg))) - A = s1.add_array('A', [1], dace.float32) + A = s1.add_access('A') t = s1.add_tasklet('geta', {'a'}, {}, 'printf("ok %f\\n", a + 1)') s1.add_edge(A, None, t, 'a', dace.Memlet.from_array(A.data, A.desc(sdfg))) - A = s2.add_array('A', [1], dace.float32) + A = s2.add_access('A') t = s2.add_tasklet('geta', {'a'}, {}, 'printf("BAD %f\\n", a - 1)') s2.add_edge(A, None, t, 'a', dace.Memlet.from_array(A.data, A.desc(sdfg))) diff --git a/tests/strided_range_copy_test.py b/tests/strided_range_copy_test.py index fa0031665d..65f7b1b33a 100644 --- a/tests/strided_range_copy_test.py +++ b/tests/strided_range_copy_test.py @@ -2,28 +2,35 @@ import dace import numpy as np -sr = dace.SDFG('stiledcopy') -s0 = sr.add_state('s0') -A = s0.add_array('A', [2, 16, 4], dace.float32) -B = s0.add_array('B', [4], dace.float32) -C = s0.add_array('C', [2, 16, 4], dace.float32) +def test(): + """Strided range copy tasklet test""" + sr = dace.SDFG('stiledcopy') + sr.add_array('A', [2, 16, 4], dace.float32) + sr.add_array('B', [4], dace.float32) + sr.add_array('C', [2, 16, 4], dace.float32) -D = s0.add_array('D', [128, 128], dace.float32) -E = s0.add_array('E', [8, 8], dace.float32) -F = s0.add_array('F', [128, 128], dace.float32) + sr.add_array('D', [128, 128], dace.float32) + sr.add_array('E', [8, 8], dace.float32) + sr.add_array('F', [128, 128], dace.float32) -# Reading A at [1, 0:8:8:2, 3] -s0.add_nedge(A, B, dace.Memlet.simple(A, '1, 0:10:8:2, 3')) -s0.add_nedge(B, C, dace.Memlet.simple(C, '1, 0:10:8:2, 3')) + s0 = sr.add_state('s0') + A = s0.add_access('A') + B = s0.add_access('B') + C = s0.add_access('C') -# Emulate a blocked tiled matrix multiplication pattern -s0.add_nedge(D, E, dace.Memlet.simple(D, '8:76:64:4,4:72:64:4')) -s0.add_nedge(E, F, dace.Memlet.simple(F, '8:76:64:4,4:72:64:4')) + D = s0.add_access('D') + E = s0.add_access('E') + F = s0.add_access('F') + # Reading A at [1, 0:8:8:2, 3] + s0.add_nedge(A, B, dace.Memlet.simple(A, '1, 0:10:8:2, 3')) + s0.add_nedge(B, C, dace.Memlet.simple(C, '1, 0:10:8:2, 3')) + + # Emulate a blocked tiled matrix multiplication pattern + s0.add_nedge(D, E, dace.Memlet.simple(D, '8:76:64:4,4:72:64:4')) + s0.add_nedge(E, F, dace.Memlet.simple(F, '8:76:64:4,4:72:64:4')) -def test(): - print('Strided range copy tasklet test') A = np.random.rand(2, 16, 4).astype(np.float32) B = np.random.rand(4).astype(np.float32) C = np.random.rand(2, 16, 4).astype(np.float32) @@ -48,7 +55,6 @@ def test(): E[4:8, 4:8] - F[72:76, 68:72], ] diff_array = [np.linalg.norm(d) for d in diffs] - print('Differences:', diff_array) assert np.average(np.array(diff_array)) <= 1e-5 diff --git a/tests/strided_range_test.py b/tests/strided_range_test.py index 50e9ccfefd..44c2041c0b 100644 --- a/tests/strided_range_test.py +++ b/tests/strided_range_test.py @@ -3,28 +3,29 @@ from dace.memlet import Memlet import numpy as np -sr = dace.SDFG('strided_range_test') -s0 = sr.add_state('s0') -A = s0.add_array('A', [2, 16, 4], dace.float32) -B = s0.add_array('B', [16], dace.float32) -tasklet = s0.add_tasklet('srtest', {'a'}, {'b'}, """ +def test_strided_range(): + sr = dace.SDFG('strided_range_test') + sr.add_array('A', [2, 16, 4], dace.float32) + sr.add_array('B', [16], dace.float32) + + s0 = sr.add_state('s0') + A = s0.add_access('A') + B = s0.add_access('B') + tasklet = s0.add_tasklet('srtest', {'a'}, {'b'}, """ b[0] = a[0,0] * 2 b[1] = a[0,1] * 2 b[2] = a[1,0] * 2 b[3] = a[1,1] * 2 """) -me, mx = s0.add_map('srmap', dict(i='0:4')) - -# Reading A at [1, 2i:2i+8:8:2, 3] -s0.add_memlet_path(A, me, tasklet, dst_conn='a', memlet=Memlet.simple(A, '1, 2*i:2*i+10:8:2, 3')) + me, mx = s0.add_map('srmap', dict(i='0:4')) -# Writing B at [4*i:4*i+4] -s0.add_memlet_path(tasklet, mx, B, src_conn='b', memlet=Memlet.simple(B, '4*i:4*i+4')) + # Reading A at [1, 2i:2i+8:8:2, 3] + s0.add_memlet_path(A, me, tasklet, dst_conn='a', memlet=Memlet.simple(A, '1, 2*i:2*i+10:8:2, 3')) + # Writing B at [4*i:4*i+4] + s0.add_memlet_path(tasklet, mx, B, src_conn='b', memlet=Memlet.simple(B, '4*i:4*i+4')) -def test(): - print('Strided range tasklet test') A = np.random.rand(2, 16, 4).astype(np.float32) B = np.random.rand(16).astype(np.float32) @@ -36,7 +37,6 @@ def test(): B[14:16] - 2 * A[1, 14:16, 3] ] diff = np.linalg.norm(np.array(diffs)) - print('Differences:', [np.linalg.norm(d) for d in diffs]) assert diff <= 1e-5 @@ -85,6 +85,6 @@ def padding(a): if __name__ == "__main__": - test() + test_strided_range() test_strided_view() test_strided_view_retval() diff --git a/tests/struct_test.py b/tests/struct_test.py index 245dbc14a1..0c6e615aee 100644 --- a/tests/struct_test.py +++ b/tests/struct_test.py @@ -1,5 +1,4 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. -import ctypes import dace import numpy as np diff --git a/tests/subset_covers_precise_test.py b/tests/subset_covers_precise_test.py index 6cdae75d17..154d756787 100644 --- a/tests/subset_covers_precise_test.py +++ b/tests/subset_covers_precise_test.py @@ -1,8 +1,5 @@ # Copyright 2019-2023 ETH Zurich and the DaCe authors. All rights reserved. -import pytest - -import dace from dace.config import Config from dace.subsets import Indices, Range diff --git a/tests/symbolic_roundtrip_test.py b/tests/symbolic_roundtrip_test.py new file mode 100644 index 0000000000..db0990183e --- /dev/null +++ b/tests/symbolic_roundtrip_test.py @@ -0,0 +1,358 @@ +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. +""" Round-trip tests for symbolic expressions through ``pystr_to_symbolic`` and ``symstr`` +(and through interstate-edge serialization), covering subscripts, operators, infinities, +NaN, booleans and float precision. """ +import sympy + +import dace +from dace import subsets, symbolic +from dace.symbolic import pystr_to_symbolic, symstr, arrays, free_symbols_and_functions, bitwise_or +from dace.frontend.python.newast import _subset_has_indirection + + +def _roundtrip(s, cpp_mode=False): + return symstr(pystr_to_symbolic(s), cpp_mode=cpp_mode) + + +def _idempotent(s): + once = symstr(pystr_to_symbolic(s)) + return once == symstr(pystr_to_symbolic(once)) + + +def test_operator_roundtrip_renders_operator(): + # The Python operators round-trip back to operators, not to ``func(a, b)``. + cases = {'~C': '~', 'a & b': '&', 'a | b': '|', 'a ^ b': '^', 'a << b': '<<', 'a >> b': '>>'} + for expr, op in cases.items(): + out = _roundtrip(expr) + assert op in out + assert 'bitwise' not in out and 'shift' not in out + assert _idempotent(expr) + + +def test_ternary_roundtrip(): + assert _roundtrip('b if c else d') == '(((b) if (c) else (d)))' + assert _idempotent('b if c else d') + + +def test_floordiv_roundtrip(): + # ``//`` round-trips to ``//`` in Python and lowers to ``/`` in C++; an explicit + # int_floor(a, b) keeps its function spelling in Python. + assert _roundtrip('a // b') == '(((a) // (b)))' + assert _roundtrip('a // b', cpp_mode=True) == '(((a) / (b)))' + assert _idempotent('a // b') + assert _roundtrip('int_floor(a, b)') == '(int_floor(a, b))' + assert _roundtrip('int_floor(a, b)', cpp_mode=True) == '(((a) / (b)))' + + +def test_cpp_lowering(): + # The operators must lower to the identical C++ operators. + assert _roundtrip('a << b', cpp_mode=True) == '(((a) << (b)))' + assert _roundtrip('a >> b', cpp_mode=True) == '(((a) >> (b)))' + assert _roundtrip('a & b', cpp_mode=True) == '(((a) & (b)))' + assert _roundtrip('a | b', cpp_mode=True) == '(((a) | (b)))' + assert _roundtrip('~c', cpp_mode=True) == '((~(c)))' + assert _roundtrip('b if c else d', cpp_mode=True) == '(((c) ? (b) : (d)))' + + +def test_operator_uses_internal_variant(): + # ``|`` parses to the ``__``-prefixed internal variant (still an instance of the + # bare class) so it round-trips to the operator. + op = pystr_to_symbolic('a | b') + assert str(op.func) == '__bitwise_or' + assert isinstance(op, bitwise_or) + + +def test_explicit_function_name_preserved(): + # An explicit ``bitwise_or(a, b)`` keeps its spelling in Python but still lowers to + # the operator in C++. + assert symstr(pystr_to_symbolic('bitwise_or(a, b)')) == '(bitwise_or(a, b))' + assert symstr(pystr_to_symbolic('bitwise_or(a, b)'), cpp_mode=True) == '(((a) | (b)))' + + +def test_subscript_roundtrip(): + assert _roundtrip('A[i]') == '(A[i])' + assert _roundtrip('sizes[i, j]') == '(sizes[i, j])' + assert _idempotent('A[i]') and _idempotent('sizes[i, j]') + + +def test_subscript_free_symbols_excludes_container(): + # Only the indices are free symbols; the container is a data access. + assert {str(s) for s in pystr_to_symbolic('A[i]').free_symbols} == {'i'} + assert {str(s) for s in pystr_to_symbolic('A[i, j]').free_symbols} == {'i', 'j'} + # Regression guard: a numeric-indexed access must not leak the container as a symbol + # (otherwise it gets mapped to its data descriptor and ``sympify`` raises). + assert pystr_to_symbolic('num[0]').free_symbols == set() + + +def test_arrays_helper(): + assert arrays('A[i] + sin(x)') == {'A'} + assert arrays('a[i, j, k]') == {'a'} + assert arrays('A[B[i]]') == {'A', 'B'} + assert arrays('a + b') == set() + + +def test_array_rename_through_subscript(): + # replace / replace_dict / subs must rename an array referenced via Subscript -- the + # GPU offloading pass relies on this (e.g. A_row -> gpu_A_row). + e = dace.InterstateEdge(assignments={'x': 'A_row[i + 1] - A_row[i]'}) + e.replace('A_row', 'gpu_A_row') + assert arrays(e.assignments['x']) == {'gpu_A_row'} + + e2 = dace.InterstateEdge(assignments={'x': 'A_row[i]'}) + e2.replace_dict({'A_row': 'gpu_A_row'}) + assert e2.assignments['x'] == 'gpu_A_row[i]' + + renamed = pystr_to_symbolic('A_row[i]').subs(symbolic.symbol('A_row'), symbolic.symbol('gpu_A_row')) + assert arrays(renamed) == {'gpu_A_row'} + + +def test_struct_member_subscript(): + # A struct-member array access (``a.b[i]``, nested ``A.B.C[0]``): only the indices + # are free symbols, the full member path is reported by ``arrays``, and the struct + # root must NOT leak into ``Range.free_symbols`` (which previously caused + # ``SympifyError`` when parsing map ranges). + e = pystr_to_symbolic('a.b[i]') + assert {str(s) for s in e.free_symbols} == {'i'} + assert arrays('a.b[i]') == {'a.b'} + assert {str(s) for s in subsets.Range([(0, e - 1, 1)]).free_symbols} == {'i'} + + # A is a struct, B a member struct, C a member array. + assert pystr_to_symbolic('A.B.C[0]').free_symbols == set() + assert arrays('A.B.C[0]') == {'A.B.C'} + assert {str(s) for s in pystr_to_symbolic('A.B.C[i]').free_symbols} == {'i'} + assert {str(s) for s in subsets.Range([(0, pystr_to_symbolic('A.B.C[i]'), 1)]).free_symbols} == {'i'} + + +def test_free_symbols_and_functions_excludes_arrays(): + # Clean split: ``arrays`` reports the container, ``free_symbols_and_functions`` does not. + assert free_symbols_and_functions('A[i] + sin(x)') == {'i', 'x'} + + +def test_scalar_versus_array_needs_descriptors(): + # A subscripted access is unambiguously a container; a bare name (possibly a rank-0 + # scalar) is indistinguishable from a symbol, so ``arrays`` reports only the former. + # Consumers recover scalars by intersecting ``free_symbols_and_functions`` with the + # SDFG's data descriptors. + assert arrays('A[i]') == {'A'} + assert arrays('s') == set() # ``s`` could be a scalar or a symbol; not decidable here + descriptors = {'A', 's'} # both are data containers in this hypothetical SDFG + referenced = (free_symbols_and_functions('A[i] + s') | arrays('A[i] + s')) & descriptors + assert referenced == {'A', 's'} # arrays() -> A; fsf & descriptors -> s + + +def test_scalar_to_symbol_promotion_moves_scalar_out_of_scalars(): + # A scalar read on an interstate edge is reported by scalars(); after scalar->symbol + # promotion it is a symbol instead, so scalars() no longer reports it (expression unchanged). + from dace.transformation.passes import scalar_to_symbol + + sdfg = dace.SDFG('promote_scalar') + sdfg.add_scalar('s', dace.int32, transient=True) + sdfg.add_array('A', [20], dace.int32) + init = sdfg.add_state('init', is_start_block=True) + init.add_edge(init.add_tasklet('set', {}, {'o'}, 'o = 7'), 'o', init.add_access('s'), None, dace.Memlet('s[0]')) + body = sdfg.add_state('body') + sdfg.add_edge(init, body, dace.InterstateEdge(assignments={'i': 's'})) + body.add_edge(body.add_tasklet('use', {}, {'o'}, 'o = i'), 'o', body.add_access('A'), None, dace.Memlet('A[0]')) + sdfg.add_symbol('i', dace.int32) + + assign = sdfg.edges()[0].data.assignments['i'] + assert symbolic.scalars(assign, sdfg.arrays) == {'s'} # 's' is a scalar read on the edge + assert 's' not in sdfg.symbols + + scalar_to_symbol.ScalarToSymbolPromotion().apply_pass(sdfg, {}) + + assert sdfg.edges()[0].data.assignments['i'] == 's' # expression unchanged + assert symbolic.scalars('s', sdfg.arrays) == set() # no longer a scalar descriptor + assert 's' in sdfg.symbols # now a symbol + + +def test_contains_sympy_functions_subscript(): + assert symbolic.contains_sympy_functions(pystr_to_symbolic('A[i]')) is True + assert symbolic.contains_sympy_functions(pystr_to_symbolic('a + b')) is False + assert symbolic.contains_sympy_functions(pystr_to_symbolic('a | b')) is False + + +def test_float_precision_preserved(): + # Floats print at the shortest round-tripping precision. + assert _roundtrip('3.14') == '3.14' + assert _roundtrip('0.1') == '0.1' + # A near-max double (Fortran ``HUGE``) stays finite and at full precision. + for huge in ('1.79769313486232e+308', '1.7976931348623157e+308', '-1.7976931348623157e+308'): + assert _roundtrip(huge) == huge + assert pystr_to_symbolic(huge) not in (sympy.oo, -sympy.oo) + + +def test_infinity_roundtrip(): + assert pystr_to_symbolic('inf') == sympy.oo + assert pystr_to_symbolic('-inf') == -sympy.oo + assert _roundtrip('inf') == 'inf' + assert _roundtrip('-inf') == '-inf' + # Infinity inside an expression resolves to ``oo`` (and ``a - inf`` to ``-oo``). + assert pystr_to_symbolic('a + inf').has(sympy.oo) + assert pystr_to_symbolic('a - inf').has(-sympy.oo) + + +def test_infinity_cpp_lowering(): + assert _roundtrip('inf', cpp_mode=True) == 'INFINITY' + assert _roundtrip('-inf', cpp_mode=True) == '-INFINITY' + + +def test_nan_roundtrip(): + assert pystr_to_symbolic('nan') is sympy.nan + assert _roundtrip('nan') == 'nan' + assert _roundtrip('nan', cpp_mode=True) == 'NAN' + + +def test_boolean_preserved_and_distinct_from_int(): + # Booleans round-trip as booleans and stay distinct from the integers 1/0. + assert _roundtrip('True') == 'True' + assert _roundtrip('False') == 'False' + assert _roundtrip('1') == '1' + assert _roundtrip('0') == '0' + assert isinstance(pystr_to_symbolic('True'), sympy.logic.boolalg.BooleanTrue) + assert _roundtrip('True', cpp_mode=True) == 'true' + + +def test_map_range_array_bound_is_hoisted(): + # An array access in a map-range bound must be hoisted to a dynamic-map-input symbol: + # the range then carries a symbol (not a raw Subscript) and the array is wired as a + # read. Without this, GPU offloading can't see the array and loop_to_map refuses it. + N = dace.symbol('N') + + @dace.program + def maprange(ptr: dace.uint32[N + 1], data: dace.float32[64], out: dace.float32[N]): + for i in dace.map[0:N]: + + @dace.map(_[ptr[i]:ptr[i + 1]]) + def inner(j): + d << data[j] + o >> out(1, lambda x, y: x + y)[i] + o = d + + sdfg = maprange.to_sdfg() + inner = [n for n, _ in sdfg.all_nodes_recursive() if isinstance(n, dace.nodes.MapEntry) and n.map.params == ['j']] + assert inner + for me in inner: + assert 'ptr[' not in str(me.map.range) # bound hoisted to a symbol, not a raw Subscript + assert me.map.range.free_symbols # the range is symbolic + + +def test_subset_indirection_detects_index_subscript(): + # An index-position array access (a gather, ``data[A_col[j]]``) is indirection; + # a plain index (``data[j]``) is not. With the head excluded from a Subscript's + # free_symbols, this detection depends on contains_sympy_functions being + # Subscript-aware. + gather = subsets.Range.from_indices([pystr_to_symbolic('A_col[j]')]) + plain = subsets.Range.from_indices([pystr_to_symbolic('j')]) + assert _subset_has_indirection(gather) + assert not _subset_has_indirection(plain) + assert pystr_to_symbolic('A_col[j]').free_symbols == {symbolic.symbol('j')} + + +def test_subset_indirection_detects_nested_and_multidim_subscript(): + # Nested gathers (``data[A[B[j]]]``) and a per-dimension gather (``data[i, A[j]]``) + # are both indirection; a fully plain multi-dim subset is not. + assert _subset_has_indirection(subsets.Range.from_indices([pystr_to_symbolic('A[B[j]]')])) + assert _subset_has_indirection(subsets.Range.from_indices([pystr_to_symbolic('i'), pystr_to_symbolic('A[j]')])) + assert not _subset_has_indirection(subsets.Range.from_indices([pystr_to_symbolic('i'), pystr_to_symbolic('j')])) + + +def test_interstate_edge_assignment_roundtrip(): + # Assignments survive save -> load -> save unchanged (idempotent serialization). + sdfg = dace.SDFG('iedge_roundtrip') + s0 = sdfg.add_state('s0', is_start_block=True) + s1 = sdfg.add_state('s1') + sdfg.add_edge( + s0, s1, dace.InterstateEdge(assignments={ + 'p': 'a | b', + 'q': 'True', + 'r': '1', + 's': '1.79769313486232e+308', + })) + + reloaded = dace.SDFG.from_json(sdfg.to_json()) + twice = dace.SDFG.from_json(reloaded.to_json()) + assert reloaded.edges()[0].data.assignments == twice.edges()[0].data.assignments + + assigns = reloaded.edges()[0].data.assignments + assert '|' in assigns['p'] and 'bitwise_or' not in assigns['p'] + assert assigns['q'] == 'True' # boolean preserved, not collapsed to '1' + assert assigns['r'] == '1' + assert assigns['s'] == '1.79769313486232e+308' # finite, full precision + + +def test_symbolic_expression_serialization_preserves_integerness(): + sdfg = dace.SDFG("roundtrip") + + sdfg.add_symbol("__out_IDim_range_0", dace.int32) + sdfg.add_symbol("__out_IDim_range_1", dace.int32) + + expr = symbolic.pystr_to_symbolic("max(0, -__out_IDim_range_0 + __out_IDim_range_1)") + + sdfg.add_array( + "A", + shape=[expr], + dtype=dace.float64, + ) + + reloaded = dace.SDFG.from_json(sdfg.to_json()) + + original = sdfg.arrays["A"].shape[0] + restored = reloaded.arrays["A"].shape[0] + + assert sympy.srepr(original) == sympy.srepr(restored) + assert original == restored + assert original.is_integer == restored.is_integer + + +def test_symbolic_roundtrip_preserves_integerness(): + expr = symbolic.pystr_to_symbolic("max(0, -__out_IDim_range_0 + __out_IDim_range_1)") + + rt = symbolic.pystr_to_symbolic(symbolic.symstr(expr)) + + assert expr.is_integer == rt.is_integer + + for s1, s2 in zip( + sorted(expr.free_symbols, key=str), + sorted(rt.free_symbols, key=str), + ): + assert s1.is_integer == s2.is_integer + + +def test_max_roundtrip_keeps_integer_assumptions(): + expr = symbolic.pystr_to_symbolic("max(0, -__out_IDim_range_0 + __out_IDim_range_1)") + + rt = symbolic.pystr_to_symbolic(symbolic.symstr(expr)) + + assert expr.is_integer + assert rt.is_integer + + +if __name__ == '__main__': + test_operator_roundtrip_renders_operator() + test_ternary_roundtrip() + test_floordiv_roundtrip() + test_cpp_lowering() + test_operator_uses_internal_variant() + test_explicit_function_name_preserved() + test_subscript_roundtrip() + test_subscript_free_symbols_excludes_container() + test_arrays_helper() + test_array_rename_through_subscript() + test_struct_member_subscript() + test_free_symbols_and_functions_excludes_arrays() + test_scalar_versus_array_needs_descriptors() + test_scalar_to_symbol_promotion_moves_scalar_out_of_scalars() + test_contains_sympy_functions_subscript() + test_float_precision_preserved() + test_infinity_roundtrip() + test_infinity_cpp_lowering() + test_nan_roundtrip() + test_boolean_preserved_and_distinct_from_int() + test_map_range_array_bound_is_hoisted() + test_subset_indirection_detects_index_subscript() + test_subset_indirection_detects_nested_and_multidim_subscript() + test_interstate_edge_assignment_roundtrip() + test_symbolic_expression_serialization_preserves_integerness() + test_symbolic_roundtrip_preserves_integerness() + test_max_roundtrip_keeps_integer_assumptions() diff --git a/tests/symbolic_serialization_test.py b/tests/symbolic_serialization_test.py new file mode 100644 index 0000000000..55b627b909 --- /dev/null +++ b/tests/symbolic_serialization_test.py @@ -0,0 +1,655 @@ +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. + +import numpy as np +import threading +import pytest + +import dace +from dace import subsets, symbolic +from dace.codegen.common import sym2cpp +from dace.properties import DictProperty, ListProperty +from dace.sdfg.infer_types import infer_connector_types +import sympy +import json + + +def test_symbolic_serialization_roundtrip_preserves_metadata(): + typed_sym = symbolic.symbol('test_sym', dtype=dace.uint64, nonnegative=True) + expr = symbolic.SymExpr(typed_sym + symbolic.TypedConstant(np.int16(2)), + typed_sym + symbolic.TypedConstant(np.int16(4))) + + serialized = symbolic.serialize_symbolic(expr) + restored = symbolic.deserialize_symbolic(serialized) + + assert serialized.startswith('SymExpr(') + assert '$test_sym' in serialized + assert isinstance(restored, symbolic.SymExpr) + + restored_sym = next(iter(restored.expr.free_symbols)) + assert isinstance(restored_sym, symbolic.symbol) + assert restored_sym.name == 'test_sym' + assert restored_sym.dtype == dace.uint64 + assert restored_sym.is_nonnegative + + constants = [atom for atom in restored.expr.atoms() if isinstance(atom, symbolic.TypedConstant)] + assert len(constants) == 1 + assert constants[0].dtype == dace.int16 + assert int(constants[0]) == 2 + + +def test_range_json_roundtrip_uses_symbolic_deserializer(): + sym = symbolic.symbol('N', dtype=dace.uint64, nonnegative=True) + rng = subsets.Range([(symbolic.TypedConstant(np.int16(2)), sym + symbolic.TypedConstant(np.int16(6)), + symbolic.TypedConstant(np.int16(2)), symbolic.TypedConstant(np.uint8(4)))]) + + restored = subsets.Range.from_json(rng.to_json(), {"version": dace.__version__}) + start, end, step = restored.ranges[0] + tile = restored.tile_sizes[0] + + assert isinstance(start, symbolic.TypedConstant) + assert start.dtype == dace.int16 + assert int(start) == 2 + + restored_sym = next(iter(end.free_symbols)) + assert isinstance(restored_sym, symbolic.symbol) + assert restored_sym.dtype == dace.uint64 + assert restored_sym.is_nonnegative + + assert isinstance(step, symbolic.TypedConstant) + assert step.dtype == dace.int16 + assert isinstance(tile, symbolic.TypedConstant) + assert tile.dtype == dace.uint8 + + +def test_symstr_codegen_for_typed_constants(): + expr = symbolic.deserialize_symbolic('2i16 + $N') + + python_expr = symbolic.symstr(expr) + cpp_expr = symbolic.symstr(expr, cpp_mode=True) + + assert 'dace.int16(2)' in python_expr + assert 'int16_t(2)' in cpp_expr + + +def test_symbol_name_clash_roundtrip(): + expr = symbolic.symbol('sin', dtype=dace.uint64, nonnegative=True) + symbolic.TypedConstant(np.uint64(1)) + + serialized = symbolic.serialize_symbolic(expr) + restored = symbolic.deserialize_symbolic(serialized) + + restored_sym = next(iter(restored.free_symbols)) + assert isinstance(restored_sym, symbolic.symbol) + assert restored_sym.name == 'sin' + assert restored_sym.dtype == dace.uint64 + + +def test_complex_symbol_roundtrip_preserves_dtype(): + expr = symbolic.symbol('c', dtype=dace.complex128) + + restored = symbolic.deserialize_symbolic(symbolic.serialize_symbolic(expr)) + + restored_sym = next(iter(restored.free_symbols)) + assert isinstance(restored_sym, symbolic.symbol) + assert restored_sym.name == 'c' + assert restored_sym.dtype == dace.complex128 + + +@pytest.mark.parametrize('dtype', [dace.complex64, dace.complex128]) +def test_suffixless_dtype_constant_roundtrips_via_cast_form(dtype): + """A constant whose dtype has no literal suffix (complex) serializes + via the parseable ``dace.(value)`` form and round-trips.""" + tc = symbolic.deserialize_symbolic(f'dace.{dtype.to_string()}(5)') + assert isinstance(tc, symbolic.TypedConstant) and tc.dtype == dtype + + serialized = symbolic.serialize_symbolic(tc) + restored = symbolic.deserialize_symbolic(serialized) + assert isinstance(restored, symbolic.TypedConstant) + assert restored.dtype == dtype and int(restored.value) == 5 + assert symbolic.serialize_symbolic(restored) == serialized + + +def test_complex_literal_parses_as_complex128(): + tc = symbolic.deserialize_symbolic('4j') + assert isinstance(tc, symbolic.TypedConstant) and tc.dtype == dace.complex128 + assert complex(sympy.re(tc.value), sympy.im(tc.value)) == complex(0, 4) + + +def test_complex_constant_parse_save_roundtrip(): + tc = symbolic.deserialize_symbolic('complex(3.0, 4.2)') + assert isinstance(tc, symbolic.TypedConstant) and tc.dtype == dace.complex128 + assert complex(sympy.re(tc.value), sympy.im(tc.value)) == complex(3.0, 4.2) + + serialized = symbolic.serialize_symbolic(tc) + assert serialized == '(3.0 + 4.2j)c128' + restored = symbolic.deserialize_symbolic(serialized) + assert restored.dtype == dace.complex128 + assert symbolic.serialize_symbolic(restored) == serialized + + c64 = symbolic.TypedConstant(complex(1, 2), dace.complex64) + s64 = symbolic.serialize_symbolic(c64) + assert s64 == '(1.0 + 2.0j)c64' + r64 = symbolic.deserialize_symbolic(s64) + assert r64.dtype == dace.complex64 and symbolic.serialize_symbolic(r64) == s64 + + +def test_sym2cpp_emits_uint64_literals(): + expr = symbolic.TypedConstant(np.uint64(1)) + symbolic.symbol('N', dtype=dace.uint64) + + cpp_expr = sym2cpp(expr) + + assert '1ULL' in cpp_expr + + +def test_sym2cpp_emits_float32_literals(): + expr = symbolic.TypedConstant(np.float32(1.5)) + symbolic.symbol('N', dtype=dace.float32) + + cpp_expr = sym2cpp(expr) + + assert '1.5f' in cpp_expr + + +def test_undefined_symbol_serialization_uses_dollar_prefix(): + undefined = symbolic.UndefinedSymbol() + + serialized = symbolic.serialize_symbolic(undefined) + restored = symbolic.deserialize_symbolic(serialized) + + assert serialized == '$?' + assert isinstance(restored, symbolic.UndefinedSymbol) + + +def test_symbol_assumption_roundtrip_preserves_bool_metadata(): + serialized = 'symbol($i, nonnegative=True)' + + restored = symbolic.deserialize_symbolic(serialized) + + assert symbolic.serialize_symbolic(restored) == serialized + + +def test_same_name_symbols_with_different_dtypes_serialize_independently(): + typed = symbolic.symbol('i', dtype=dace.int16) + default = symbolic.symbol('i') + + assert symbolic.serialize_symbolic(typed) == 'symbol($i, dtype=dace.int16)' + assert symbolic.serialize_symbolic(default) == '$i' + + +def test_typed_symbol_deserialization_does_not_strip_dtype(): + # Input may not be in canonical order, but dtype must survive a round-trip + original = '2*(1 + symbol($M, dtype=dace.int16))*symbol($M, dtype=dace.int16)' + restored = symbolic.deserialize_symbolic(original) + reserialized = symbolic.serialize_symbolic(restored) + + # Dtype still appears in the output + assert 'symbol($M, dtype=dace.int16)' in reserialized + + # The round-trip is now a fixed point (deterministic) + assert symbolic.serialize_symbolic(symbolic.deserialize_symbolic(reserialized)) == reserialized + + # Also verify the symbol object itself has the correct dtype + from dace import int16 + syms = {s.name: s for s in restored.free_symbols} + assert 'M' in syms and syms['M'].dtype == int16 + + +def test_rational_addition_roundtrip_preserves_serialization(): + serialized = '-10/3 + 1/3*$i' + + restored = symbolic.deserialize_symbolic(serialized) + + assert symbolic.serialize_symbolic(restored) == serialized + + +def test_pystr_to_symbolic_preserves_typed_symbols(): + # Prime the parser cache with an equal untyped SymPy expression. + cache_seed_expr = symbolic.deserialize_symbolic('-1 + $N') + cache_seed_restored = symbolic.pystr_to_symbolic(cache_seed_expr) + typed_expr = symbolic.deserialize_symbolic('-1 + symbol($N, dtype=dace.int16)') + + typed_restored = symbolic.pystr_to_symbolic(typed_expr) + + assert cache_seed_restored is cache_seed_expr + assert symbolic.serialize_symbolic(cache_seed_restored) == '-1 + $N' + assert typed_restored is typed_expr + assert symbolic.serialize_symbolic(typed_restored) == '-1 + symbol($N, dtype=dace.int16)' + + +def test_power_deserialization_preserves_typed_symbols_after_plain_power(): + plain_power = symbolic.deserialize_symbolic('$N**2') + + typed_power = symbolic.deserialize_symbolic('symbol($N, dtype=dace.int16)**2') + + assert symbolic.serialize_symbolic(plain_power) == '$N**2' + assert symbolic.serialize_symbolic(typed_power) == 'symbol($N, dtype=dace.int16)**2' + + +@pytest.mark.parametrize('simplify', [None, False]) +def test_pystr_to_symbolic_keeps_basic_unsimplified_by_default(simplify): + expr = sympy.Add(symbolic.symbol('N'), 1, 1, evaluate=False) + + restored = symbolic.pystr_to_symbolic(expr, simplify=simplify) + + assert restored is expr + assert len(restored.args) == 3 + + +def test_pystr_to_symbolic_simplifies_basic_when_requested(): + expr = sympy.Add(symbolic.symbol('N'), 1, 1, evaluate=False) + + restored = symbolic.pystr_to_symbolic(expr, simplify=True) + + assert restored == symbolic.symbol('N') + 2 + assert len(restored.args) == 2 + + +def test_range_json_roundtrip_preserves_typed_symbol_minus_one(): + json_range = { + 'type': 'Range', + 'ranges': [{ + 'start': '0', + 'end': '-1 + symbol($N, dtype=dace.int16)', + 'step': '1', + 'tile': '1', + }], + } + + rng = subsets.Range.from_json(json_range, {"version": dace.__version__}) + + assert symbolic.serialize_symbolic(rng.ranges[0][0]) == '0' + assert symbolic.serialize_symbolic(rng.ranges[0][1]) == '-1 + symbol($N, dtype=dace.int16)' + assert symbolic.serialize_symbolic(rng.ranges[0][2]) == '1' + assert symbolic.serialize_symbolic(rng.tile_sizes[0]) == '1' + + +def test_scalar_memlet_connector_type_after_symbolic_range_roundtrip(): + i = symbolic.symbol('i', dtype=dace.int64) + stencil_i = symbolic.symbol('stencil_i', dtype=dace.int64) + start = 2 * i - 2 * stencil_i + 1 + rng = subsets.Range([(start, start, 1)]) + restored = subsets.Range.from_json(rng.to_json(), {"version": dace.__version__}) + restored.replace({i: stencil_i}) + + assert 'i' not in restored.free_symbols + assert restored.ranges[0][0] == 1 + assert restored.ranges[0][1] == 1 + assert restored.num_elements() == 1 + + sdfg = dace.SDFG('scalar_memlet_connector_after_symbolic_roundtrip') + sdfg.add_array('A', [20], dace.float64) + sdfg.add_scalar('B', dace.float64) + state = sdfg.add_state() + read = state.add_read('A') + tasklet = state.add_tasklet('use_scalar', {'inp'}, {'out'}, 'out = inp') + write = state.add_write('B') + state.add_edge(read, None, tasklet, 'inp', dace.Memlet(data='A', subset=restored)) + state.add_edge(tasklet, 'out', write, None, dace.Memlet('B[0]')) + + infer_connector_types(sdfg) + + assert tasklet.in_connectors['inp'] == dace.float64 + + +def test_list_property_symbolic_type_json_roundtrip_supports_plain_names(): + prop = ListProperty(element_type=sympy.Basic) + + assert prop.to_json(['START']) == ['$START'] + + restored = prop.from_json(prop.to_json(['START']), {"version": dace.__version__}) + + assert len(restored) == 1 + assert restored[0] == symbolic.symbol('START') + + +def test_dict_property_symbolic_type_json_roundtrip_supports_plain_names(): + prop = DictProperty(key_type=str, value_type=sympy.Basic) + + assert prop.to_json({'N': 'N'}) == {'N': '$N'} + + restored = prop.from_json(prop.to_json({'N': 'N'}), {"version": dace.__version__}) + + assert restored == {'N': symbolic.symbol('N')} + + +@pytest.mark.parametrize('expr', [ + sympy.Add(symbolic.TypedConstant(np.int16(10)), + sympy.Mul(sympy.S.NegativeOne, symbolic.symbol('i'), evaluate=False), + evaluate=False), + sympy.Mul(symbolic.TypedConstant(np.int16(10)), symbolic.symbol('i'), evaluate=False), + sympy.Mul(symbolic.TypedConstant(np.int16(10)), sympy.Pow(symbolic.symbol('i'), -1, evaluate=False), + evaluate=False), + sympy.Pow(symbolic.TypedConstant(np.int16(10)), symbolic.symbol('i'), evaluate=False), + sympy.Mod(symbolic.symbol('i'), symbolic.TypedConstant(np.int16(3)), evaluate=False), +]) +def test_typed_binary_operator_roundtrip_preserves_serialization(expr): + serialized = symbolic.serialize_symbolic(expr) + restored = symbolic.deserialize_symbolic(serialized) + + assert symbolic.serialize_symbolic(restored) == serialized + + +def test_plain_integer_roundtrip_converts_to_sympy_integer(): + restored = symbolic.deserialize_symbolic(symbolic.serialize_symbolic(sympy.Integer(10))) + + assert restored == 10 + assert isinstance(restored, sympy.Integer) + + +def test_plain_sympy_integer_serializes_without_typed_suffix(): + assert symbolic.serialize_symbolic(sympy.Integer(42)) == '42' + + +@pytest.mark.parametrize('value', [-7, 2**80]) +def test_plain_python_integer_deserialization_uses_sympy_integer(value): + restored = symbolic.deserialize_symbolic(value) + + assert restored == value + assert isinstance(restored, sympy.Integer) + + +def test_untyped_literals_keep_plain_sympy_form(): + """Untyped literals deserialize to plain SymPy numbers (no implicit + DaCe-typing).""" + assert symbolic.deserialize_symbolic('5') == sympy.Integer(5) + assert isinstance(symbolic.deserialize_symbolic('5'), sympy.Integer) + assert isinstance(symbolic.deserialize_symbolic('5.0'), sympy.Float) + + +@pytest.mark.parametrize( + 'text,ctype', + [ + ('double(5)', 'double'), # float64 cast-wrapper form + ('int(5)', 'int'), # int32 cast-wrapper form + ('float(5)', 'float'), # float32 + ('int64_t(5)', 'int64_t'), # explicit-suffix dtype as a ctype cast + ]) +def test_cpp_ctype_cast_parses_to_typed_constant(text, ctype): + """``double(5)``/``int(5)`` (the C++ printer's cast fallback) must round-trip + into a TypedConstant.""" + restored = symbolic.deserialize_symbolic(text) + assert isinstance(restored, symbolic.TypedConstant), text + assert restored.dtype.ctype == ctype + + +@pytest.mark.parametrize('op', ['Min', 'Max']) +@pytest.mark.parametrize('literal', [ + '5.0', + '5.0f64', + '5f32', + '7f64', +]) +def test_minmax_with_float_literal_roundtrip(op, literal): + serialized = f'{op}({literal}, $N)' + restored = symbolic.deserialize_symbolic(serialized) + assert symbolic.serialize_symbolic(restored) == serialized + + +def test_minmax_with_ctype_cast_int_literal_normalizes_then_stable(): + first = symbolic.serialize_symbolic(symbolic.deserialize_symbolic('Max(double(7), 5.0)')) + assert first == '7f64' + second = symbolic.serialize_symbolic(symbolic.deserialize_symbolic(first)) + assert first == second + + +@pytest.mark.parametrize('serialized', [ + '5i16', + '7i64', + '5f32', + '5.0f64', + '(3.0 + 4.0j)c128', + '(3.0 - 4.0j)c128', + '(-3.0 + 4.0j)c128', + '(3.0 + 4.0j)c64', + '(4.0j)c128', + '(-4.0j)c128', + '(4.0j)c64', + '(-4.0j)c64', +]) +def test_typed_constant_canonical_form_roundtrips(serialized): + restored = symbolic.deserialize_symbolic(serialized) + assert isinstance(restored, symbolic.TypedConstant) + assert symbolic.serialize_symbolic(restored) == serialized + + +@pytest.mark.parametrize('input_form,canonical', [ + ('4j', '(4.0j)c128'), + ('-4j', '(-4.0j)c128'), + ('(0.0 + 4.0j)c128', '(4.0j)c128'), + ('complex(3.0, 4.0)', '(3.0 + 4.0j)c128'), + ('dace.complex128(complex(3.0, 4.0))', '(3.0 + 4.0j)c128'), + ('dace.complex64(complex(3.0, 4.0))', '(3.0 + 4.0j)c64'), + ('dace.complex128(complex(3.0, -4.0))', '(3.0 - 4.0j)c128'), +]) +def test_legacy_complex_form_normalizes_to_canonical_suffix(input_form, canonical): + first = symbolic.serialize_symbolic(symbolic.deserialize_symbolic(input_form)) + assert first == canonical + second = symbolic.serialize_symbolic(symbolic.deserialize_symbolic(first)) + assert first == second + + +def test_serialization_symbol_dtypes_isolation_multi_thread(): + first_thread_waits_to_enter_context = threading.Event() + first_thread_has_entered_context = threading.Event() + second_thread_has_entered_context = threading.Event() + first_thread_has_serialized = threading.Event() + second_thread_has_serialized = threading.Event() + failures = [] + lock = threading.Lock() + + def worker(native_dtype, override_dtype, wait_to_enter_context, sig_entered_context, do_serialize, sig_serialize, + do_non_scoped_serialize): + try: + sym = symbolic.symbol('N', dtype=native_dtype) + assert wait_to_enter_context.wait(timeout=10) + with symbolic.serialization_symbol_dtypes({'N': override_dtype}): + sig_entered_context.set() + assert do_serialize.wait(timeout=10) + assert symbolic.serialize_symbolic(sym) == f'symbol($N, dtype=dace.{override_dtype.to_string()})' + sig_serialize.set() + assert do_non_scoped_serialize.wait(timeout=10) + assert symbolic.serialize_symbolic(sym) == f'symbol($N, dtype=dace.{native_dtype.to_string()})' + except BaseException as ex: + with lock: + failures.append(ex) + + thread1 = threading.Thread( + target=worker, + args=( + dace.int64, # native_dtype + dace.int16, # override_dtype + first_thread_waits_to_enter_context, # wait_to_enter_context + first_thread_has_entered_context, # sig_entered_context + second_thread_has_entered_context, # do_serializing + first_thread_has_serialized, # sig_serialize + second_thread_has_serialized, # do_non_scoped_serialize + )) + thread2 = threading.Thread( + target=worker, + args=( + dace.uint64, # native_dtype + dace.uint32, # override_dtype + first_thread_has_entered_context, # wait_to_enter_context + second_thread_has_entered_context, # sig_entered_context + first_thread_has_serialized, # do_serialize + second_thread_has_serialized, # sig_serialize + second_thread_has_serialized, # do_non_scoped_serialize + )) + + thread1.start() + thread2.start() + + # To set it in motion we must allow the first thread to enter the context. + first_thread_waits_to_enter_context.set() + + thread1.join() + thread2.join() + + if failures: + raise failures[0] + + +def _roundtrip(expr): + """Return (first serialization, serialization after one parse round-trip).""" + s1 = symbolic.serialize_symbolic(expr) + # FIX: Use deserialize_symbolic instead of pystr_to_symbolic + s2 = symbolic.serialize_symbolic(symbolic.deserialize_symbolic(s1)) + return s1, s2 + + +def _ceiling_triangular(): + tN1 = symbolic.symbol('tN1') + tk1 = symbolic.symbol('tk1', dace.int64) + return sympy.ceiling(sympy.Rational(1, 2) * (1 + tk1) * tN1 - sympy.Rational(1, 2) * tN1 * tk1) + + +def _floor_triangular(): + tN2 = symbolic.symbol('tN2') + tk2 = symbolic.symbol('tk2', dace.int64) + return sympy.floor(sympy.Rational(1, 2) * (1 + tk2) * tN2 - sympy.Rational(1, 2) * tN2 * tk2) + + +def _mixed_unevaluated_sum(): + ta1, tb1, tc1 = symbolic.symbol('ta1'), symbolic.symbol('tb1'), symbolic.symbol('tc1') + return sympy.Rational(1, 2) * (ta1 + tb1) * tc1 - sympy.Rational(1, 2) * tc1 * tb1 + ta1 * (tb1 + 1) + + +def _nested_ceiling(): + tN1 = symbolic.symbol('tN3') + tk1 = symbolic.symbol('tk3', dace.int64) + inner = sympy.ceiling((tN1 - 1) * (1 + tk1) - tN1 * tk1) + return sympy.ceiling(sympy.Rational(1, 2) * inner + tN1) + + +def _stencil_bound_min(): + tN5, tb2 = symbolic.symbol('tN4'), symbolic.symbol('tb2') + return sympy.Min(tN5 - 1, tb2 + 31) + 1 + + +def _max_of_integers(): + ta3 = symbolic.symbol('t_out_range_0', integer=True) + tb3 = symbolic.symbol('t_out_range_1', integer=True) + return sympy.Max(ta3, tb3) + + +def _mul_with_transcendental_coeff(): + # An `is_number` factor (pi) must print as its own factor, not fold into the numeric + # coefficient -- folding it makes the Mul printer recurse forever. + return 3 * sympy.pi * symbolic.symbol('ttrc') + + +def _mul_with_imaginary_unit(): + # Exercises both the Mul printer (I is `is_number`) and `I` deserialization. + return 2 * sympy.I * symbolic.symbol('timg') + + +def _mul_with_sqrt_coeff(): + # sqrt(2) must deserialize back to the real Pow (not an undefined ``sqrt`` function), + # otherwise its sort key flips against the symbol on the round-trip. + return 3 * sympy.sqrt(2) * symbolic.symbol('tsqrt') + + +CEILING_AND_SIMILAR = { + 'ceiling_triangular': _ceiling_triangular, + 'floor_triangular': _floor_triangular, + 'mixed_unevaluated_sum': _mixed_unevaluated_sum, + 'nested_ceiling': _nested_ceiling, + 'stencil_bound_min': _stencil_bound_min, + 'max_of_integers': _max_of_integers, + 'mul_with_transcendental_coeff': _mul_with_transcendental_coeff, + 'mul_with_imaginary_unit': _mul_with_imaginary_unit, + 'mul_with_sqrt_coeff': _mul_with_sqrt_coeff, +} + + +@pytest.mark.parametrize('build', CEILING_AND_SIMILAR.values(), ids=list(CEILING_AND_SIMILAR)) +def test_serialization_is_fixed_point(build): + """serialize -> parse -> serialize must be byte-identical (and idempotent).""" + expr = build() + s1, s2 = _roundtrip(expr) + assert s2 == s1, f'round-trip changed the serialization:\n first: {s1!r}\n again: {s2!r}' + # A second pass must not drift either. + # FIX: Use deserialize_symbolic instead of pystr_to_symbolic + s3 = symbolic.serialize_symbolic(symbolic.deserialize_symbolic(s2)) + assert s3 == s1 + + +@pytest.mark.parametrize('build', CEILING_AND_SIMILAR.values(), ids=list(CEILING_AND_SIMILAR)) +def test_roundtrip_preserves_free_symbol_names(build): + expr = build() + # FIX: Use deserialize_symbolic instead of pystr_to_symbolic + reparsed = symbolic.deserialize_symbolic(symbolic.serialize_symbolic(expr)) + assert {s.name for s in reparsed.free_symbols} == {s.name for s in expr.free_symbols} + + +def test_standalone_fraction_deserializes_without_stray_identity(): + """A numeric quotient must fold to a reduced Number, not ``Mul(1, 1/3)``.""" + reparsed = symbolic.deserialize_symbolic('1/3') + assert not isinstance(reparsed, sympy.Mul) + assert reparsed == sympy.Rational(1, 3) + + +def test_imaginary_unit_roundtrips(): + """The serializer emits ``I`` for ``sympy.I``; it must deserialize back to it.""" + assert symbolic.deserialize_symbolic('I') is sympy.I + expr = 3 + 4 * sympy.I + assert symbolic.deserialize_symbolic(symbolic.serialize_symbolic(expr)) == expr + + +def test_sqrt_deserializes_to_real_sqrt(): + """``sqrt(2)`` must parse to the real Pow, not an opaque ``Function('sqrt')``.""" + reparsed = symbolic.deserialize_symbolic('sqrt(2)') + assert not isinstance(reparsed, sympy.core.function.AppliedUndef) + assert reparsed == sympy.sqrt(2) + + +def test_add_order_independent_of_arg_order(): + """Pinpoints the printer: term order must not depend on Add arg order.""" + N = symbolic.symbol('N') + k = symbolic.symbol('k', dace.int64) + t1 = sympy.Rational(1, 2) * (1 + k) * N + t2 = -sympy.Rational(1, 2) * N * k + assert (symbolic.serialize_symbolic(sympy.Add(t1, t2, evaluate=False)) == symbolic.serialize_symbolic( + sympy.Add(t2, t1, evaluate=False))) + + +def test_integer_symbol_assumptions_preserved(): + """No extra explicit assumptions (e.g. commutative) may leak in on reparse.""" + expr = _max_of_integers() + # FIX: Use deserialize_symbolic instead of pystr_to_symbolic + reparsed = symbolic.deserialize_symbolic(symbolic.serialize_symbolic(expr)) + assert sympy.srepr(reparsed) == sympy.srepr(expr) + assert all(s.is_integer for s in reparsed.free_symbols) + + +def test_int64_symbol_dtype_preserved(): + k = symbolic.symbol('k', dace.int64) + # FIX: Use deserialize_symbolic instead of pystr_to_symbolic + reparsed = symbolic.deserialize_symbolic(symbolic.serialize_symbolic(k)) + (rk, ) = reparsed.free_symbols + assert rk.name == 'k' + assert rk.dtype == dace.int64 + + +def test_sdfg_json_roundtrip_is_fixed_point(): + """Faithful reproduction of the original `test.sdfg` vs `test2.sdfg` mismatch. + + A symbolic array shape serializes through the same SymbolicProperty path as + memlet volumes / num_accesses. If `add_array` rejects the ceiling shape on + your branch, attach `volume` to a Memlet's `volume` instead and round-trip + that Memlet's `to_json` / `from_json`. + """ + N = symbolic.symbol('N', dace.int64) + k = symbolic.symbol('k', dace.int64) + volume = sympy.ceiling(sympy.Rational(1, 2) * (1 + k) * N - sympy.Rational(1, 2) * N * k) + + sdfg = dace.SDFG('roundtrip_determinism') + sdfg.add_symbol('N', dace.int64) + sdfg.add_symbol('k', dace.int64) + sdfg.add_array('A', [volume], dace.float64) + + j1 = sdfg.to_json() + j2 = dace.SDFG.from_json(j1).to_json() + assert json.dumps(j1, sort_keys=True) == json.dumps(j2, sort_keys=True) + + +if __name__ == '__main__': + pytest.main([__file__]) diff --git a/tests/tensorflow/callback_test.py b/tests/tensorflow/callback_test.py index 01b706e765..c2fd916405 100644 --- a/tests/tensorflow/callback_test.py +++ b/tests/tensorflow/callback_test.py @@ -1,6 +1,7 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. import pytest import numpy as np +from importlib.util import find_spec @pytest.mark.tensorflow @@ -35,8 +36,5 @@ def test_callback(): if __name__ == '__main__': - try: - import tensorflow + if find_spec("tensorflow"): test_callback() - except ImportError: - pass diff --git a/tests/tensorflow/compile_test.py b/tests/tensorflow/compile_test.py index 1fd00d32ed..0a9d3b6b6b 100644 --- a/tests/tensorflow/compile_test.py +++ b/tests/tensorflow/compile_test.py @@ -1,6 +1,7 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. import pytest import numpy as np +from importlib.util import find_spec @pytest.mark.tensorflow @@ -29,8 +30,5 @@ def test_compile(): if __name__ == '__main__': - try: - import tensorflow + if find_spec("tensorflow"): test_compile() - except ImportError: - pass diff --git a/tests/tensorflow/conv_test.py b/tests/tensorflow/conv_test.py index 41a6e384d7..4ccf345765 100644 --- a/tests/tensorflow/conv_test.py +++ b/tests/tensorflow/conv_test.py @@ -1,6 +1,7 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. import pytest import numpy as np +from importlib.util import find_spec @pytest.mark.tensorflow @@ -101,8 +102,5 @@ def test_conv(): if __name__ == '__main__': - try: - import tensorflow + if find_spec("tensorflow"): test_conv() - except ImportError: - pass diff --git a/tests/tensorflow/fbn_test.py b/tests/tensorflow/fbn_test.py index fcc5f56fde..09e4de2450 100644 --- a/tests/tensorflow/fbn_test.py +++ b/tests/tensorflow/fbn_test.py @@ -111,7 +111,6 @@ def test_fused_batch_norm(): if __name__ == '__main__': try: - import tensorflow test_fused_batch_norm() except ImportError: pass diff --git a/tests/tensorflow/ops_test.py b/tests/tensorflow/ops_test.py index a34e882967..1d36d1c374 100644 --- a/tests/tensorflow/ops_test.py +++ b/tests/tensorflow/ops_test.py @@ -1,6 +1,7 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. import pytest import numpy as np +from importlib.util import find_spec @pytest.mark.tensorflow @@ -104,11 +105,8 @@ def test_slice(): if __name__ == '__main__': - try: - import tensorflow + if find_spec("tensorflow"): test_shapen() test_mean() test_addn() test_slice() - except ImportError: - pass diff --git a/tests/tensorflow/pool_test.py b/tests/tensorflow/pool_test.py index b30b5f01fb..9b582fc65b 100644 --- a/tests/tensorflow/pool_test.py +++ b/tests/tensorflow/pool_test.py @@ -1,6 +1,7 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. import pytest import numpy as np +from importlib.util import find_spec @pytest.mark.tensorflow @@ -87,8 +88,5 @@ def test_pooling(): if __name__ == '__main__': - try: - import tensorflow + if find_spec("tensorflow"): test_pooling() - except ImportError: - pass diff --git a/tests/tensorflow/simple_test.py b/tests/tensorflow/simple_test.py index 31abdb8513..c3529d617b 100644 --- a/tests/tensorflow/simple_test.py +++ b/tests/tensorflow/simple_test.py @@ -1,6 +1,7 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. import pytest import numpy as np +from importlib.util import find_spec @pytest.mark.tensorflow @@ -26,8 +27,5 @@ def test_simple(): if __name__ == '__main__': - try: - import tensorflow + if find_spec("tensorflow"): test_simple() - except ImportError: - pass diff --git a/tests/threadlocal_stream_test.py b/tests/threadlocal_stream_test.py index ebefb9d04e..b26eee8782 100644 --- a/tests/threadlocal_stream_test.py +++ b/tests/threadlocal_stream_test.py @@ -1,38 +1,40 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. -import math import numpy as np import dace as dp from dace.sdfg import SDFG from dace.memlet import Memlet -N = dp.symbol('N') -sdfg = SDFG('tlstream') -state = sdfg.add_state('doit') -localarr = state.add_transient('la', [10], dp.float32) -localstream = state.add_stream('ls', dp.float32, 1, transient=True) -globalstream = state.add_stream('gs', dp.float32, 1, transient=True) -globalarr = state.add_array('ga', [N], dp.float32) +def test(): + """Thread-local stream test""" -me, mx = state.add_map('par', dict(i='0:N')) -tasklet = state.add_tasklet('arange', set(), {'a'}, 'a = i') + N = dp.symbol('N') + sdfg = SDFG('tlstream') + sdfg.add_transient('la', [10], dp.float32) + sdfg.add_stream('ls', dp.float32, 1, transient=True) + sdfg.add_stream('gs', dp.float32, 1, transient=True) + sdfg.add_array('ga', [N], dp.float32) + state = sdfg.add_state('doit') -state.add_nedge(me, tasklet, Memlet()) -state.add_edge(tasklet, 'a', localstream, None, Memlet.from_array(localstream.data, localstream.desc(sdfg))) -state.add_nedge(localstream, localarr, Memlet.from_array(localarr.data, localarr.desc(sdfg))) -state.add_nedge(localarr, mx, Memlet.from_array(globalstream.data, globalstream.desc(sdfg))) -state.add_nedge(mx, globalstream, Memlet.from_array(globalstream.data, globalstream.desc(sdfg))) -state.add_nedge(globalstream, globalarr, Memlet.from_array(globalarr.data, globalarr.desc(sdfg))) + localarr = state.add_access('la') + localstream = state.add_access('ls') + globalstream = state.add_access('gs') + globalarr = state.add_access('ga') -sdfg.fill_scope_connectors() + me, mx = state.add_map('par', dict(i='0:N')) + tasklet = state.add_tasklet('arange', set(), {'a'}, 'a = i') + state.add_nedge(me, tasklet, Memlet()) + state.add_edge(tasklet, 'a', localstream, None, Memlet.from_array(localstream.data, localstream.desc(sdfg))) + state.add_nedge(localstream, localarr, Memlet.from_array(localarr.data, localarr.desc(sdfg))) + state.add_nedge(localarr, mx, Memlet.from_array(globalstream.data, globalstream.desc(sdfg))) + state.add_nedge(mx, globalstream, Memlet.from_array(globalstream.data, globalstream.desc(sdfg))) + state.add_nedge(globalstream, globalarr, Memlet.from_array(globalarr.data, globalarr.desc(sdfg))) -def test(): - print('Thread-local stream test') + sdfg.fill_scope_connectors() N = 20 - output = np.ndarray([N], dtype=np.float32) output[:] = dp.float32(0) @@ -46,7 +48,6 @@ def test(): output = np.sort(output) diff = np.linalg.norm(output - np.arange(0, N, dtype=np.float32)) - print("Difference:", diff) assert diff <= 1e-5 diff --git a/tests/tile_test.py b/tests/tile_test.py index d1a3a46b95..0f2c98fa0b 100644 --- a/tests/tile_test.py +++ b/tests/tile_test.py @@ -1,7 +1,5 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. -import argparse import dace -import math import numpy as np W = dace.symbol('W') diff --git a/tests/tiling_test.py b/tests/tiling_test.py index fd14a15a63..0b653c5490 100644 --- a/tests/tiling_test.py +++ b/tests/tiling_test.py @@ -1,8 +1,6 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. -import argparse import dace from dace.transformation.dataflow import MapTiling -from dace.transformation.optimizer import SDFGOptimizer import numpy as np from scipy import ndimage diff --git a/tests/torch_forward/test_attn.py b/tests/torch_forward/test_attn.py index df673e057c..cf2aed3b75 100644 --- a/tests/torch_forward/test_attn.py +++ b/tests/torch_forward/test_attn.py @@ -6,7 +6,6 @@ from dace.ml import DaceModule -from dace.transformation.dataflow import RedundantSecondArray from tests.utils import torch_tensors_close diff --git a/tests/torch_forward/test_cpp_extension.py b/tests/torch_forward/test_cpp_extension.py index c8b1b624f4..8d16f6f5d2 100644 --- a/tests/torch_forward/test_cpp_extension.py +++ b/tests/torch_forward/test_cpp_extension.py @@ -4,7 +4,6 @@ pytest.importorskip("torch", reason="PyTorch not installed. Please install with: pip install dace[ml]") import os -import numpy as np import torch import torch.utils.cpp_extension from dace.codegen import targets, compiler diff --git a/tests/torch_forward/test_debug_transients.py b/tests/torch_forward/test_debug_transients.py index 995b986b8d..8704e4a953 100644 --- a/tests/torch_forward/test_debug_transients.py +++ b/tests/torch_forward/test_debug_transients.py @@ -4,7 +4,6 @@ pytest.importorskip("torch", reason="PyTorch not installed. Please install with: pip install dace[ml]") import torch from torch import nn -import numpy as np import dace from tests.utils import torch_tensors_close diff --git a/tests/torch_forward/test_img_op_implementations.py b/tests/torch_forward/test_img_op_implementations.py index 4ec120b0f9..7a15d4b36d 100644 --- a/tests/torch_forward/test_img_op_implementations.py +++ b/tests/torch_forward/test_img_op_implementations.py @@ -4,7 +4,6 @@ pytest.importorskip("torch", reason="PyTorch not installed. Please install with: pip install dace[ml]") import torch from torch import nn -import numpy as np from dace.ml import DaceModule from tests.utils import torch_tensors_close diff --git a/tests/transformations/add_threadblock_map_test.py b/tests/transformations/add_threadblock_map_test.py index aee328bae1..b215471494 100644 --- a/tests/transformations/add_threadblock_map_test.py +++ b/tests/transformations/add_threadblock_map_test.py @@ -3,7 +3,6 @@ import dace import pytest import numpy -import sys from dace.transformation.dataflow.add_threadblock_map import AddThreadBlockMap N = dace.symbol("N") diff --git a/tests/transformations/apply_transformations_once_test.py b/tests/transformations/apply_transformations_once_test.py index ba8a156e7d..9ce9a7ffaf 100644 --- a/tests/transformations/apply_transformations_once_test.py +++ b/tests/transformations/apply_transformations_once_test.py @@ -1,6 +1,5 @@ # Copyright 2019-2025 ETH Zurich and the DaCe authors. All rights reserved. import dace -import numpy as np from dace import transformation as dace_transformation, properties as dace_properties from dace.sdfg import nodes as dace_nodes from typing import Any, Union diff --git a/tests/transformations/copy_to_map_test.py b/tests/transformations/copy_to_map_test.py index 0ffff02c05..eecb1729e7 100644 --- a/tests/transformations/copy_to_map_test.py +++ b/tests/transformations/copy_to_map_test.py @@ -4,8 +4,6 @@ import copy import pytest import numpy as np -import re -from typing import Tuple, Optional def _copy_to_map(storage: dace.StorageType): @@ -111,7 +109,7 @@ def test_preprocess(): test_copy_to_map() test_flatten_to_map() try: - import cupy + import cupy # noqa: F401 test_copy_to_map_gpu() test_flatten_to_map_gpu() test_preprocess() diff --git a/tests/transformations/helpers_test.py b/tests/transformations/helpers_test.py new file mode 100644 index 0000000000..85b57ed463 --- /dev/null +++ b/tests/transformations/helpers_test.py @@ -0,0 +1,105 @@ +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. + +import dace + +from dace.transformation.helpers import all_isedges_between + + +def _make_consecutive_loops_sdfg(): + sdfg = dace.SDFG('consecutive_loops') + + before_first = sdfg.add_state('before_first', is_start_block=True) + first_body = sdfg.add_state('first_body') + between_loops = sdfg.add_state('between_loops') + sdfg.add_loop(before_first, first_body, between_loops, 'i', '0', 'i < 4', 'i + 1', label='first_loop') + + second_body = sdfg.add_state('second_body') + after_second = sdfg.add_state('after_second') + sdfg.add_loop(between_loops, second_body, after_second, 'j', '0', 'j < 4', 'j + 1', label='second_loop') + + return sdfg + + +def _make_nested_loops_sdfg(): + sdfg = dace.SDFG('nested_loops') + + before_outer = sdfg.add_state('before_outer', is_start_block=True) + outer_before_inner = sdfg.add_state('outer_before_inner') + inner_body = sdfg.add_state('inner_body') + after_inner = sdfg.add_state('after_inner') + after_outer = sdfg.add_state('after_outer') + + sdfg.add_loop(outer_before_inner, inner_body, after_inner, 'j', '0', 'j < 4', 'j + 1', label='inner_loop') + sdfg.add_loop(before_outer, + outer_before_inner, + after_outer, + 'i', + '0', + 'i < 4', + 'i + 1', + label='outer_loop', + loop_end_block=after_inner) + + return sdfg + + +def _edge_signature(edge): + return ( + edge.src.label, + type(edge.src).__name__, + getattr(edge.src.parent_graph, 'label', None), + edge.dst.label, + type(edge.dst).__name__, + getattr(edge.dst.parent_graph, 'label', None), + edge.data.condition.as_string if edge.data.condition else None, + tuple(sorted((str(k), str(v)) for k, v in edge.data.assignments.items())), + ) + + +def test_all_isedges_between_consecutive_loops(): + sdfg = _make_consecutive_loops_sdfg() + blocks = {block.label: block for block in sdfg.all_control_flow_blocks()} + + edges_to_after_first = { + _edge_signature(edge) + for edge in all_isedges_between(blocks['first_body'], blocks['between_loops']) + } + assert edges_to_after_first == { + ('first_loop', 'LoopRegion', 'consecutive_loops', 'between_loops', 'SDFGState', 'consecutive_loops', '1', ()), + } + + edges_to_second_body = { + _edge_signature(edge) + for edge in all_isedges_between(blocks['first_body'], blocks['second_body']) + } + assert edges_to_second_body == { + ('first_loop', 'LoopRegion', 'consecutive_loops', 'between_loops', 'SDFGState', 'consecutive_loops', '1', ()), + ('between_loops', 'SDFGState', 'consecutive_loops', 'second_loop', 'LoopRegion', 'consecutive_loops', '1', ()), + } + + +def test_all_isedges_between_nested_loops(): + sdfg = _make_nested_loops_sdfg() + blocks = {block.label: block for block in sdfg.all_control_flow_blocks()} + + edges_to_after_inner = { + _edge_signature(edge) + for edge in all_isedges_between(blocks['inner_body'], blocks['after_inner']) + } + assert edges_to_after_inner == { + ('inner_loop', 'LoopRegion', 'outer_loop', 'after_inner', 'SDFGState', 'outer_loop', '1', ()), + } + + edges_to_after_outer = { + _edge_signature(edge) + for edge in all_isedges_between(blocks['inner_body'], blocks['after_outer']) + } + assert edges_to_after_outer == { + ('inner_loop', 'LoopRegion', 'outer_loop', 'after_inner', 'SDFGState', 'outer_loop', '1', ()), + ('outer_loop', 'LoopRegion', 'nested_loops', 'after_outer', 'SDFGState', 'nested_loops', '1', ()), + } + + +if __name__ == '__main__': + test_all_isedges_between_consecutive_loops() + test_all_isedges_between_nested_loops() diff --git a/tests/transformations/interstate/condition_fusion_test.py b/tests/transformations/interstate/condition_fusion_test.py index 4d61bd992e..99062ddf07 100644 --- a/tests/transformations/interstate/condition_fusion_test.py +++ b/tests/transformations/interstate/condition_fusion_test.py @@ -1,8 +1,5 @@ # Copyright 2019-2025 ETH Zurich and the DaCe authors. All rights reserved. -import numpy as np -import pytest import dace -from dace import nodes from dace.transformation.interstate import ConditionFusion from dace.sdfg.state import ConditionalBlock diff --git a/tests/transformations/interstate/condition_map_interchange_test.py b/tests/transformations/interstate/condition_map_interchange_test.py index fe478f1773..7d0997c95a 100644 --- a/tests/transformations/interstate/condition_map_interchange_test.py +++ b/tests/transformations/interstate/condition_map_interchange_test.py @@ -2,10 +2,8 @@ """Tests the conditional map interchange transformation.""" import numpy as np -import pytest import dace from dace.transformation.interstate import ConditionMapInterchange -from dace.sdfg.state import LoopRegion from copy import deepcopy diff --git a/tests/transformations/interstate/loop_overwrite_elimination_test.py b/tests/transformations/interstate/loop_overwrite_elimination_test.py index 6d8fb08b52..c579942833 100644 --- a/tests/transformations/interstate/loop_overwrite_elimination_test.py +++ b/tests/transformations/interstate/loop_overwrite_elimination_test.py @@ -2,7 +2,6 @@ """Tests loop overwrite elimination transformations.""" import numpy as np -import pytest import dace from dace.transformation.interstate import LoopOverwriteElimination from dace.sdfg.state import LoopRegion diff --git a/tests/transformations/loop_to_map_test.py b/tests/transformations/interstate/loop_to_map_test.py similarity index 80% rename from tests/transformations/loop_to_map_test.py rename to tests/transformations/interstate/loop_to_map_test.py index ecb905c0f0..27f90c55c6 100644 --- a/tests/transformations/loop_to_map_test.py +++ b/tests/transformations/interstate/loop_to_map_test.py @@ -303,7 +303,7 @@ def test_need_for_tasklet(): aname, _ = sdfg.add_array('A', (10, ), dace.int32) bname, _ = sdfg.add_array('B', (10, ), dace.int32) body = sdfg.add_state('body') - _, _, _ = sdfg.add_loop(None, body, None, 'i', '0', 'i < 10', 'i + 1', None) + _, _, _ = sdfg.add_loop_state_machine(None, body, None, 'i', '0', 'i < 10', 'i + 1', None) anode = body.add_access(aname) bnode = body.add_access(bname) body.add_nedge(anode, bnode, dace.Memlet(data=aname, subset='i', other_subset='9 - i')) @@ -325,7 +325,7 @@ def test_need_for_transient(): aname, _ = sdfg.add_array('A', (10, 10), dace.int32) bname, _ = sdfg.add_array('B', (10, 10), dace.int32) body = sdfg.add_state('body') - _, _, _ = sdfg.add_loop(None, body, None, 'i', '0', 'i < 10', 'i + 1', None) + _, _, _ = sdfg.add_loop_state_machine(None, body, None, 'i', '0', 'i < 10', 'i + 1', None) anode = body.add_access(aname) bnode = body.add_access(bname) body.add_nedge(anode, bnode, dace.Memlet(data=aname, subset='0:10, i', other_subset='0:10, 9 - i')) @@ -390,7 +390,7 @@ def test_symbol_write_before_read(): body_start = sdfg.add_state() body = sdfg.add_state() body_end = sdfg.add_state() - sdfg.add_loop(init, body_start, None, 'i', '0', 'i < 20', 'i + 1', loop_end_state=body_end) + sdfg.add_loop_state_machine(init, body_start, None, 'i', '0', 'i < 20', 'i + 1', loop_end_state=body_end) # Internal loop structure sdfg.add_edge(body_start, body, dace.InterstateEdge(assignments=dict(j='0'))) @@ -410,7 +410,7 @@ def test_symbol_array_mix(overwrite): body = sdfg.add_state() body_end = sdfg.add_state() after = sdfg.add_state() - sdfg.add_loop(init, body_start, after, 'i', '0', 'i < 20', 'i + 1', loop_end_state=body_end) + sdfg.add_loop_state_machine(init, body_start, after, 'i', '0', 'i < 20', 'i + 1', loop_end_state=body_end) sdfg.out_edges(init)[0].data.assignments['sym'] = '0.0' @@ -438,7 +438,7 @@ def test_symbol_array_mix_2(parallel): body_start = sdfg.add_state() body_end = sdfg.add_state() after = sdfg.add_state() - sdfg.add_loop(init, body_start, after, 'i', '1', 'i < 20', 'i + 1', loop_end_state=body_end) + sdfg.add_loop_state_machine(init, body_start, after, 'i', '1', 'i < 20', 'i + 1', loop_end_state=body_end) sdfg.out_edges(init)[0].data.assignments['sym'] = '0.0' @@ -463,7 +463,7 @@ def test_internal_symbol_used_outside(overwrite): body = sdfg.add_state() body_end = sdfg.add_state() after = sdfg.add_state() - sdfg.add_loop(init, body_start, after, 'i', '0', 'i < 20', 'i + 1', loop_end_state=body_end) + sdfg.add_loop_state_machine(init, body_start, after, 'i', '0', 'i < 20', 'i + 1', loop_end_state=body_end) # Internal loop structure sdfg.add_edge(body_start, body, dace.InterstateEdge(assignments=dict(j='0'))) @@ -800,6 +800,146 @@ def tester(A: dace.float32[10, 10]): assert sdfg.apply_transformations_repeated(LoopToMap) == 2 +def test_stride_symbol_propagated_to_nested_sdfg(): + """A ``ConditionalBlock`` guard that indexes a 2D array whose stride uses + a free symbol ``S``. After two applications of ``LoopToMap``, both the + level-1 and level-2 ``NestedSDFG``s must carry ``S`` in their + ``symbol_mapping`` AND in the free-symbol set consumed by + ``generate_nsdfg_header`` -- otherwise codegen emits a ``loop_body_*`` + function using ``S`` without declaring it. Regression for a bug observed + on the velocity-tendencies pipeline where compile fails with + ``'S' was not declared in this scope``. + """ + from dace.properties import CodeBlock + from dace.sdfg.state import ConditionalBlock, ControlFlowRegion, LoopRegion + + sdfg = dace.SDFG("loop_to_map_missing_stride_symbol") + sdfg.add_symbol("S", dace.int32) + sdfg.add_array("A", (10, 10), dace.int32, strides=(1, dace.symbol("S"))) + sdfg.add_scalar("acc", dace.int32, transient=True) + + outer = LoopRegion("outer", condition_expr="k < 5", loop_var="k", initialize_expr="k = 0", update_expr="k = k + 1") + sdfg.add_node(outer, is_start_block=True) + + init = outer.add_state("init", is_start_block=True) + t0 = init.add_tasklet("t_init", {}, {"o"}, "o = 0") + init.add_edge(t0, "o", init.add_write("acc"), None, dace.Memlet("acc[0]")) + + inner = LoopRegion("inner", condition_expr="j < 5", loop_var="j", initialize_expr="j = 0", update_expr="j = j + 1") + outer.add_node(inner) + outer.add_edge(init, inner, dace.InterstateEdge()) + + cb = ConditionalBlock("cb") + inner.add_node(cb, is_start_block=True) + branch = ControlFlowRegion("branch", sdfg=sdfg) + cb.add_branch(CodeBlock("A[j, k] == 1"), branch) + + set_one = branch.add_state("set_one", is_start_block=True) + t1 = set_one.add_tasklet("t_set", {}, {"o"}, "o = 1") + set_one.add_edge(t1, "o", set_one.add_write("acc"), None, dace.Memlet("acc[0]")) + + sdfg.validate() + applied = sdfg.apply_transformations_repeated(LoopToMap, permissive=True) + assert applied == 2 + + # Every ``NestedSDFG`` the pass created must have ``S`` in both its + # ``symbol_mapping`` and its inner ``sdfg.symbols``. Otherwise the + # stride-using subscript is emitted without the declaration. + for node, _ in sdfg.all_nodes_recursive(): + if not isinstance(node, nodes.NestedSDFG): + continue + if any('S' in {str(s) for s in d.free_symbols} for d in node.sdfg.arrays.values()): + assert 'S' in node.symbol_mapping, node.symbol_mapping + assert 'S' in node.sdfg.symbols, sorted(node.sdfg.symbols) + + # End-to-end: the SDFG must compile cleanly. Before the fix this emits + # ``'S' was not declared in this scope`` inside ``loop_body_*``. + sdfg.compile() + + +def test_dynamic_write_slab_separated_by_iteration_var(): + """Loop whose body contains a Map that writes conditionally to a 2D + array at ``[nproma_i, level]``: per-lane dynamic, but each outer + iteration pins the ``level`` axis so no two iterations touch the same + slab. Earlier LoopToMap unconditionally rejected any dynamic-volume + write (``Dynamic memlet without WCR``); now it accepts the write when + the destination subset passes the same ``a*i+b`` check used for + non-dynamic writes. + + Shape is the ``_for_it_35`` level loop reduced to its essentials. + """ + from dace import memlet as mm + from dace.properties import CodeBlock + from dace.sdfg.state import ConditionalBlock, ControlFlowRegion + from dace.sdfg.propagation import propagate_memlets_sdfg + + sdfg = dace.SDFG("for_it_35_repro") + NLEV = dace.symbol("NLEV", dace.int32) + NPROMA = dace.symbol("NPROMA", dace.int32) + sdfg.add_array("src", [NPROMA, NLEV], dace.float64) + sdfg.add_array("cfl_clipping", [NPROMA, NLEV], dace.int32) + + outer = LoopRegion( + "level_loop", + condition_expr="level <= NLEV - 3", + loop_var="level", + initialize_expr="level = 3", + update_expr="level = level + 1", + ) + sdfg.add_node(outer, is_start_block=True) + body = outer.add_state("body", is_start_block=True) + + r_src = body.add_read("src") + w_cfl = body.add_write("cfl_clipping") + me, mx = body.add_map("nproma", {"i": "0:NPROMA"}) + me.add_in_connector("IN_src") + me.add_out_connector("OUT_src") + mx.add_in_connector("IN_cfl") + mx.add_out_connector("OUT_cfl") + + # Inner conditional write wrapped in a nested SDFG (this is what makes + # the aggregated MapExit memlet dynamic). + nsdfg = dace.SDFG("inner") + nsdfg.add_array("src", [1], dace.float64) + nsdfg.add_array("cfl", [1], dace.int32) + cb = ConditionalBlock("cb") + nsdfg.add_node(cb, is_start_block=True) + branch = ControlFlowRegion("branch", sdfg=nsdfg) + cb.add_branch(CodeBlock("(1 == 1)"), branch) + s = branch.add_state("s", is_start_block=True) + r_in = s.add_read("src") + w_out = s.add_write("cfl") + tk = s.add_tasklet("cmp", {"s"}, {"o"}, "o = (s > 0.5)") + s.add_edge(r_in, None, tk, "s", mm.Memlet("src[0]")) + s.add_edge(tk, "o", w_out, None, mm.Memlet("cfl[0]")) + + nested = body.add_nested_sdfg(nsdfg, {"src"}, {"cfl"}) + body.add_edge(r_src, None, me, "IN_src", mm.Memlet("src[0:NPROMA, level]")) + body.add_edge(me, "OUT_src", nested, "src", mm.Memlet("src[i, level]")) + body.add_edge(nested, "cfl", mx, "IN_cfl", mm.Memlet("cfl_clipping[i, level]")) + body.add_edge(mx, "OUT_cfl", w_cfl, None, mm.Memlet("cfl_clipping[0:NPROMA, level]")) + + sdfg.validate() + propagate_memlets_sdfg(sdfg) + + # The aggregated MapExit memlet is dynamic (conditional write), with + # dst_subset ``[0:NPROMA, level]``. + found_dyn = False + for state in sdfg.all_states(): + for n in state.nodes(): + if isinstance(n, nodes.MapExit): + for e in state.out_edges(n): + if e.data.dynamic: + found_dyn = True + assert found_dyn, "expected a dynamic memlet before LoopToMap" + + applied = sdfg.apply_transformations_repeated(LoopToMap, permissive=True) + assert applied == 1, f"LoopToMap should lift the level loop (got {applied})" + + # No LoopRegion should remain. + assert not any(isinstance(n, LoopRegion) for n, _ in sdfg.all_nodes_recursive()) + + if __name__ == "__main__": parser = argparse.ArgumentParser() @@ -840,3 +980,5 @@ def tester(A: dace.float32[10, 10]): test_rotated_loop_to_map(True) test_self_loop_to_map() test_nested_sdfg_nested_loop() + test_stride_symbol_propagated_to_nested_sdfg() + test_dynamic_write_slab_separated_by_iteration_var() diff --git a/tests/transformations/interstate/loop_unroll_test.py b/tests/transformations/interstate/loop_unroll_test.py index d744361d36..b725bff95b 100644 --- a/tests/transformations/interstate/loop_unroll_test.py +++ b/tests/transformations/interstate/loop_unroll_test.py @@ -1,4 +1,4 @@ -# Copyright 2019-2025 ETH Zurich and the DaCe authors. All rights reserved. +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. import re @@ -8,6 +8,7 @@ from dace.sdfg.sdfg import InterstateEdge from dace.sdfg.state import ConditionalBlock, ControlFlowRegion, LoopRegion from dace.transformation.interstate.loop_unroll import LoopUnroll +from dace.transformation.passes.prune_symbols import RemoveUnusedSymbols def _get_sdfg(add_state_before: bool, l: int) -> dace.SDFG: @@ -97,8 +98,114 @@ def test_compiler_unroll_pragma(): assert unroll_pragma, "Unroll pragma not found in generated code after setting unroll_pragma to True." +klev = dace.symbol('klev', dtype=dace.int32) +klon = dace.symbol('klon', dtype=dace.int32) +nclv = dace.symbol('nclv', dtype=dace.int32) + + +@dace.program +def melt_kernel( + zqxfg: dace.float64[nclv, klon], + zsolqa: dace.float64[nclv, nclv, klon], + zmeltmax: dace.float64[klon], + zicetot: dace.float64[klon], + imelt: dace.int32[nclv], +): + zepsec = 1e-14 + for jm in range(nclv): + for jl in range(klon): + if zmeltmax[jl] > zepsec and zicetot[jl] > zepsec: + zalfa2 = zqxfg[jm, jl] / zicetot[jl] + zmelt = min(zqxfg[jm, jl], zalfa2 * zmeltmax[jl]) + zqxfg[jm, jl] = zqxfg[jm, jl] - zmelt + zqxfg[imelt[jm] - 1, jl] = zqxfg[imelt[jm] - 1, jl] + zmelt + zsolqa[jm, imelt[jm] - 1, jl] = zsolqa[jm, imelt[jm] - 1, jl] + zmelt + zsolqa[imelt[jm] - 1, jm, jl] = zsolqa[imelt[jm] - 1, jm, jl] - zmelt + + +@dace.program +def triang_elim_kernel( + zqlhs: dace.float64[nclv, nclv, klon], + kidia: dace.int32, + kfdia: dace.int32, +): + for jn in range(1, nclv - 1 + 1): + for jm in range(jn + 1, nclv + 1): + for ik in range(jn + 1, nclv + 1): + for jl in range(kidia, kfdia + 1): + zqlhs[ik - 1, jm - 1, jl - 1] = (zqlhs[ik - 1, jm - 1, jl - 1] - + zqlhs[jn - 1, jm - 1, jl - 1] * zqlhs[ik - 1, jn - 1, jl - 1]) + + +def test_triang_elim(): + KLON, NCLV = 512, 5 + KIDIA, KFDIA = 1, KLON + + sdfg = triang_elim_kernel.to_sdfg() + sdfg.replace_dict({"nclv": NCLV}) + sdfg.validate() + + # Unroll the outer jn loop + loops = {n for n in sdfg.nodes() if isinstance(n, LoopRegion)} + assert len(loops) == 1 + loop = loops.pop() + assert loop.loop_variable == "jn" + + LoopUnroll().apply_to(sdfg=sdfg, loop=loop) + sdfg.validate() + + RemoveUnusedSymbols().apply_pass(sdfg, {}) + + free_syms = {str(s) for s in sdfg.free_symbols} + assert "jn" not in free_syms, f"Loop variable 'jn' still found in free symbols after unrolling: {free_syms}" + sdfg.compile() + + +def test_melt_kernel(): + sdfg = melt_kernel.to_sdfg() + sdfg.replace_dict({"nclv": 5}) + sdfg.validate() + + loops = {n for n in sdfg.nodes() if isinstance(n, LoopRegion)} + assert len(loops) == 1 + loop = loops.pop() + + LoopUnroll().apply_to(sdfg=sdfg, loop=loop) + sdfg.validate() + + # Test's aim to duplicate the case where a connection is missing + src_nodes = {n for n in sdfg.nodes() if sdfg.in_degree(n) == 0} + dst_nodes = {n for n in sdfg.nodes() if sdfg.out_degree(n) == 0} + assert len(src_nodes) == 1 + assert len(dst_nodes) == 1 + + other_nodes = set(sdfg.nodes()).difference((src_nodes.union(dst_nodes))) + assert all({(sdfg.in_degree(n) == 1 and sdfg.out_degree(n) == 1) for n in other_nodes}) + + +def test_replace_dict_inner_loop(): + nclv = dace.symbol('nclv', dtype=dace.int32) + klon = dace.symbol('klon', dtype=dace.int32) + + @dace.program + def nested_loop_kernel(A: dace.float64[nclv, klon], jn: dace.int32): + for jm in range(jn, nclv + 1): + for jl in range(klon): + A[jm, jl] = A[jm, jl] + A[jn, jl] + + sdfg = nested_loop_kernel.to_sdfg() + sdfg.replace_dict({"jn": "1"}) + sdfg.validate() + sdfg.remove_data("jn") + sdfg.validate() + assert "jn" not in {str(s) for s in sdfg.arglist()} + + if __name__ == "__main__": test_if_block_inside_for() test_empty_loop() test_top_level_for() test_compiler_unroll_pragma() + test_melt_kernel() + test_triang_elim() + test_replace_dict_inner_loop() diff --git a/tests/transformations/isolate_nested_sdfg_test.py b/tests/transformations/isolate_nested_sdfg_test.py index 78666a201c..86c5025c4c 100644 --- a/tests/transformations/isolate_nested_sdfg_test.py +++ b/tests/transformations/isolate_nested_sdfg_test.py @@ -1,7 +1,5 @@ # Copyright 2019-2025 ETH Zurich and the DaCe authors. All rights reserved. import dace -import numpy as np -import pytest from typing import Tuple, Union from dace import nodes as dace_nodes diff --git a/tests/transformations/map_collapse_test.py b/tests/transformations/map_collapse_test.py index 4d8f4d9e6a..02629a4c72 100644 --- a/tests/transformations/map_collapse_test.py +++ b/tests/transformations/map_collapse_test.py @@ -1,5 +1,4 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. -import copy import dace from dace.sdfg.utils import consolidate_edges from dace.transformation.dataflow import MapCollapse diff --git a/tests/transformations/mapfission_test.py b/tests/transformations/map_fission_test.py similarity index 66% rename from tests/transformations/mapfission_test.py rename to tests/transformations/map_fission_test.py index 2145918211..621ef61301 100644 --- a/tests/transformations/mapfission_test.py +++ b/tests/transformations/map_fission_test.py @@ -1,12 +1,12 @@ -# Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. import copy import dace +import pytest from dace.sdfg import nodes, utils as sdutils from dace.transformation.dataflow import MapFission from dace.transformation.interstate import InlineSDFG from dace.transformation.helpers import nest_state_subgraph import numpy as np -import unittest def mapfission_sdfg(): @@ -603,6 +603,363 @@ def test_dependent_symbol(): assert np.array_equal(val, ref) +@dace.program +def strided_two_ops(A: dace.float64[10], B: dace.float64[10]): + for i in dace.map[0:9:2]: + tmp = A[i] * 2.0 + B[i] = tmp + 1.0 + + +def test_strided_fission_step2(): + """ MapFission on a map with step=2 and a scalar border transient. """ + A = np.arange(10, dtype=np.float64) + B_ref = np.zeros(10, dtype=np.float64) + B_test = np.zeros(10, dtype=np.float64) + + strided_two_ops(A=A, B=B_ref) + + sdfg = strided_two_ops.to_sdfg() + sdfg.save("before.sdfg") + assert sdfg.apply_transformations(MapFission, validate=True, validate_all=True) > 0 + + sdfg(A=A, B=B_test) + sdfg.save("after.sdfg") + assert np.allclose(B_test, B_ref) + + +@dace.program +def strided_offset_ops(A: dace.float64[30], B: dace.float64[30]): + for i in dace.map[10:29:3]: + tmp = A[i] + 100.0 + B[i] = tmp * 0.5 + + +def test_strided_fission_offset_and_step(): + """ MapFission on a map with offset=10 and step=3. """ + A = np.arange(30, dtype=np.float64) + B_ref = np.zeros(30, dtype=np.float64) + B_test = np.zeros(30, dtype=np.float64) + + strided_offset_ops(A=A, B=B_ref) + + sdfg = strided_offset_ops.to_sdfg() + assert sdfg.apply_transformations(MapFission, validate=True, validate_all=True) > 0 + + sdfg(A=A, B=B_test) + assert np.allclose(B_test, B_ref) + + +_N = dace.symbol('N') +_START = dace.symbol('START') +_STEP = dace.symbol('STEP') +_STOP = dace.symbol('STOP') + + +@dace.program +def symbolic_strided(A: dace.float64[_N], B: dace.float64[_N]): + for i in dace.map[_START:_STOP:_STEP]: + B[i] = A[i] + 2.0 + A[i] = B[i] * 0.5 + + +def test_symbolic_strided_fission(): + """ MapFission on a symbolic strided map. """ + sdfg = symbolic_strided.to_sdfg() + assert sdfg.apply_transformations(MapFission, validate=True, validate_all=True) > 0 + + A_ref = np.arange(10, dtype=np.float64) + A_test = A_ref.copy() + B_ref = np.zeros(10, dtype=np.float64) + B_test = np.zeros(10, dtype=np.float64) + + symbolic_strided(A=A_ref, B=B_ref, N=10, START=1, STOP=9, STEP=2) + sdfg(A=A_test, B=B_test, N=10, START=1, STOP=9, STEP=2) + + assert np.allclose(B_test, B_ref) + assert np.allclose(A_test, A_ref) + + +@dace.program +def trivial_1d(A: dace.float64[16], B: dace.float64[16]): + for i in dace.map[0:16]: + tmp = A[i] + 1.0 + B[i] = tmp * 2.0 + + +def test_trivial_1d_step1(): + """ MapFission on a simple 1D map with step=1. """ + A = np.arange(16, dtype=np.float64) + B_ref = np.zeros(16, dtype=np.float64) + B_test = np.zeros(16, dtype=np.float64) + + trivial_1d(A=A, B=B_ref) + + sdfg = trivial_1d.to_sdfg() + assert sdfg.apply_transformations(MapFission, validate=True, validate_all=True) > 0 + + sdfg(A=A, B=B_test) + assert np.allclose(B_test, B_ref) + + +@dace.program +def trivial_3d(A: dace.float64[4, 5, 6], B: dace.float64[4, 5, 6]): + for i, j, k in dace.map[0:4, 0:5, 0:6]: + tmp = A[i, j, k] + 1.0 + B[i, j, k] = tmp + 2.0 + + +def test_trivial_3d_step1(): + """ MapFission on a 3D map with step=1. """ + A = np.arange(4 * 5 * 6, dtype=np.float64).reshape(4, 5, 6).copy() + B_ref = np.zeros((4, 5, 6), dtype=np.float64) + B_test = np.zeros((4, 5, 6), dtype=np.float64) + + trivial_3d(A=A, B=B_ref) + + sdfg = trivial_3d.to_sdfg() + assert sdfg.apply_transformations(MapFission, validate=True, validate_all=True) > 0 + + sdfg(A=A, B=B_test) + assert np.allclose(B_test, B_ref) + + +@dace.program +def map_with_data_cond(A: dace.float64[10]): + for i in dace.map[0:10]: + if A[i] > 0.0: + A[i] = A[i] + 1.9 + + +def test_map_with_if_nested_sdfg(): + """ MapFission must refuse maps whose body's interstate assignments depend on the map iterator. """ + # The number of applications depends on whether auto-opt is enabled. + # We only check numerical correctness. + sdfg = map_with_data_cond.to_sdfg() + + A_ref = np.array([-1.0, 2.0, -3.0, 4.0, 5.0, -6.0, 7.0, -8.0, 9.0, 10.0], dtype=np.float64) + A_test = A_ref.copy() + map_with_data_cond(A=A_ref) + sdfg(A=A_test) + assert np.allclose(A_test, A_ref) + + +N, M, P, K = (dace.symbol(s) for s in ('N', 'M', 'P', 'K')) + + +@dace.program +def _nested_parent(x: dace.float64[N, M], y: dace.float64[N, M]): + for j in dace.map[0:M]: + for jj in dace.map[0:P]: + for i in dace.map[0:N]: + x[i, j] = 1.0 + for i in dace.map[0:N]: + y[i, j] = 2.0 + + +@dace.program +def _toplevel_two_components(x: dace.float64[N, M], y: dace.float64[N, M]): + for j in dace.map[0:M]: + for i in dace.map[0:N]: + x[i, j] = 1.0 + for i in dace.map[0:N]: + y[i, j] = 2.0 + + +@dace.program +def _only_conditional(x: dace.float64[N]): + for i in dace.map[0:N]: + if i % 2 == 0: + x[i] = 1.0 + else: + x[i] = -1.0 + + +@dace.program +def _control_flow_in_scope(x: dace.float64[N, M], y: dace.float64[N, M]): + for j in dace.map[0:M]: + for i in dace.map[0:N]: + if j % 2 == 0: + x[i, j] = 1.0 + else: + x[i, j] = -1.0 + for i in dace.map[0:N]: + y[i, j] = 2.0 + + +@dace.program +def _five_set_five_cpy(s0: dace.float64[K], s1: dace.float64[K], s2: dace.float64[K], s3: dace.float64[K], + s4: dace.float64[K], a0: dace.float64[K], a1: dace.float64[K], a2: dace.float64[K], + a3: dace.float64[K], a4: dace.float64[K], c0: dace.float64[K], c1: dace.float64[K], + c2: dace.float64[K], c3: dace.float64[K], c4: dace.float64[K]): + for i in dace.map[0:K]: + s0[i] = 0.0 + s1[i] = 1.0 + s2[i] = 2.0 + s3[i] = 3.0 + s4[i] = 4.0 + c0[i] = a0[i] + c1[i] = a1[i] + c2[i] = a2[i] + c3[i] = a3[i] + c4[i] = a4[i] + + +@dace.program +def _three_set_two_cpy(s0: dace.float64[K], s1: dace.float64[K], s2: dace.float64[K], a0: dace.float64[K], + a1: dace.float64[K], c0: dace.float64[K], c1: dace.float64[K]): + for i in dace.map[0:K]: + s0[i] = 0.0 + s1[i] = 1.0 + s2[i] = 2.0 + c0[i] = a0[i] + c1[i] = a1[i] + + +def _find_map_entry(sdfg, param, nested): + """Find a single-parameter map entry by nesting. + + :param sdfg: The SDFG to search. + :param param: The map's sole iteration parameter name. + :param nested: Whether the map must be inside another scope. + :returns: A ``(state, map_entry)`` pair, or ``(None, None)``. + """ + for st in sdfg.all_states(): + for n in st.nodes(): + if isinstance(n, nodes.MapEntry) and n.map.params == [param]: + if (st.entry_node(n) is not None) == nested: + return st, n + return None, None + + +def _fission_maps(sdfg): + """Return the top-level ``(state, map_entry)`` pairs. + + :param sdfg: The SDFG to scan. + :returns: One pair per outermost map entry. + """ + return [(st, n) for st in sdfg.all_states() for n in st.nodes() + if isinstance(n, nodes.MapEntry) and st.entry_node(n) is None] + + +def _assert_no_spurious_connectors(sdfg, n_set, n_cpy): + """Assert each fissioned map carries exactly its component's connectors. + + :param sdfg: The fissioned SDFG. + :param n_set: Expected count of memset (zero-input) maps. + :param n_cpy: Expected count of memcpy (one-input) maps. + :raises AssertionError: On a wrong map count or an unused connector. + """ + maps = _fission_maps(sdfg) + assert len(maps) == n_set + n_cpy, f"expected {n_set + n_cpy} maps, got {len(maps)}" + n_in0 = n_in1 = 0 + for st, me in maps: + mx = st.exit_node(me) + in_c = {c for c in me.in_connectors if c.startswith('IN_')} + out_c = {c for c in me.out_connectors if c.startswith('OUT_')} + assert {e.dst_conn for e in st.in_edges(me) if e.dst_conn} == in_c + assert {e.src_conn for e in st.out_edges(me) if e.src_conn} == out_c + assert all(st.out_edges(mx)), "map exit has a dangling output" + deg = len(in_c) + assert deg in (0, 1), f"unexpected input arity {deg}" + n_in0 += deg == 0 + n_in1 += deg == 1 + assert n_in0 == n_set and n_in1 == n_cpy + + +def test_mapfission_handles_nested_parent_correctly(): + """Fission a map nested inside another map (regression: Leftover nodes).""" + n, m, p = 4, 3, 2 + sdfg = _nested_parent.to_sdfg(simplify=True) + _, jj = _find_map_entry(sdfg, 'jj', nested=True) + assert jj is not None + assert MapFission.can_be_applied_to(sdfg, map_entry=jj) is True + + x0, y0 = np.zeros((n, m)), np.zeros((n, m)) + copy.deepcopy(sdfg)(x=x0, y=y0, N=n, M=m, P=p) + + assert sdfg.apply_transformations_repeated(MapFission) >= 1 + sdfg.validate() + + x1, y1 = np.zeros((n, m)), np.zeros((n, m)) + sdfg(x=x1, y=y1, N=n, M=m, P=p) + assert np.allclose(x1, x0) and np.allclose(y1, y0) + assert np.allclose(x1, 1.0) and np.allclose(y1, 2.0) + + +def test_mapfission_still_applies_to_toplevel_map(): + """Top-level two-component fission is unaffected by the nested fix.""" + n, m = 6, 5 + sdfg = _toplevel_two_components.to_sdfg(simplify=True) + _, j = _find_map_entry(sdfg, 'j', nested=False) + assert j is not None and MapFission.can_be_applied_to(sdfg, map_entry=j) is True + + assert sdfg.apply_transformations_repeated(MapFission) >= 1 + sdfg.validate() + + x, y = np.zeros((n, m)), np.zeros((n, m)) + sdfg(x=x, y=y, N=n, M=m) + assert np.allclose(x, 1.0) and np.allclose(y, 2.0) + + +@pytest.mark.parametrize("prog,n_set,n_cpy", [(_five_set_five_cpy, 5, 5), (_three_set_two_cpy, 3, 2)]) +def test_memset_memcpy_fission_clean_connectors(prog, n_set, n_cpy): + """Many memset/memcpy components fission into maps with no spurious connectors.""" + k = 8 + sdfg = prog.to_sdfg(simplify=True) + names = list(sdfg.arg_names) + args = {nm: (np.zeros(k) if nm.startswith(('s', 'c')) else np.random.rand(k)) for nm in names} + ref = {nm: v.copy() for nm, v in args.items()} + copy.deepcopy(sdfg)(**ref, K=k) + + assert sdfg.apply_transformations_repeated(MapFission) >= 1 + sdfg.validate() + _assert_no_spurious_connectors(sdfg, n_set, n_cpy) + + out = {nm: v.copy() for nm, v in args.items()} + sdfg(**out, K=k) + for nm in names: + assert np.allclose(out[nm], ref[nm]), f"mismatch on {nm}" + + +def test_mapfission_does_not_apply_to_conditional_map(): + """A map whose sole body is a conditional: MapFission must not apply.""" + n = 6 + sdfg = _only_conditional.to_sdfg(simplify=True) + assert sdfg.apply_transformations_repeated(MapFission) == 0 + sdfg.validate() + + x = np.zeros(n) + sdfg(x=x, N=n) + assert np.allclose(x, np.where(np.arange(n) % 2 == 0, 1.0, -1.0)) + + +def test_mapfission_refuses_conditional_component_stays_valid(): + """A conditional component is refused; the rest stays valid (regression: Leftover nodes).""" + n, m = 5, 4 + sdfg = _control_flow_in_scope.to_sdfg(simplify=True) + + x0, y0 = np.zeros((n, m)), np.zeros((n, m)) + copy.deepcopy(sdfg)(x=x0, y=y0, N=n, M=m) + + for st in sdfg.all_states(): + for me in st.nodes(): + if not isinstance(me, nodes.MapEntry): + continue + has_cond = any( + isinstance(e.dst, nodes.NestedSDFG) and any( + type(b).__name__ == 'ConditionalBlock' for b in e.dst.sdfg.all_control_flow_regions(recursive=True)) + for e in st.out_edges(me)) + if has_cond: + assert MapFission.can_be_applied_to(sdfg, map_entry=me) is False + + sdfg.apply_transformations_repeated(MapFission) + sdfg.validate() + + x1, y1 = np.zeros((n, m)), np.zeros((n, m)) + sdfg(x=x1, y=y1, N=n, M=m) + assert np.allclose(x1, x0) and np.allclose(y1, y0) + + if __name__ == '__main__': test_subgraph() test_nested_sdfg() @@ -618,3 +975,15 @@ def test_dependent_symbol(): test_array_copy_outside_scope() test_single_data_multiple_connectors() test_dependent_symbol() + test_strided_fission_step2() + test_strided_fission_offset_and_step() + test_symbolic_strided_fission() + test_trivial_1d_step1() + test_trivial_3d_step1() + test_map_with_if_nested_sdfg() + test_mapfission_handles_nested_parent_correctly() + test_mapfission_still_applies_to_toplevel_map() + test_memset_memcpy_fission_clean_connectors(_five_set_five_cpy, 5, 5) + test_memset_memcpy_fission_clean_connectors(_three_set_two_cpy, 3, 2) + test_mapfission_does_not_apply_to_conditional_map() + test_mapfission_refuses_conditional_component_stays_valid() diff --git a/tests/transformations/map_fusion_horizontal_test.py b/tests/transformations/map_fusion_horizontal_test.py index b3c4509d60..a6cdc31c0a 100644 --- a/tests/transformations/map_fusion_horizontal_test.py +++ b/tests/transformations/map_fusion_horizontal_test.py @@ -1,11 +1,9 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. -from typing import Any, Union, Tuple, Type, Optional, List +from typing import Tuple -import numpy as np import pytest import dace -from dace import SDFG, SDFGState from dace.sdfg import nodes from dace.transformation import dataflow as dftrans diff --git a/tests/transformations/map_tiling_test.py b/tests/transformations/map_tiling_test.py index 802bf41d30..eac074afef 100644 --- a/tests/transformations/map_tiling_test.py +++ b/tests/transformations/map_tiling_test.py @@ -1,9 +1,7 @@ # Copyright 2019-2025 ETH Zurich and the DaCe authors. All rights reserved. from typing import List import dace -import numpy as np from dace.transformation.dataflow import MapTiling -import pytest N = dace.symbol('N') diff --git a/tests/transformations/nested_copies_among_connectors.py b/tests/transformations/nested_copies_among_connectors.py index 6e289b4fbc..170e960e5a 100644 --- a/tests/transformations/nested_copies_among_connectors.py +++ b/tests/transformations/nested_copies_among_connectors.py @@ -2,7 +2,6 @@ import dace import numpy as np -from dace.transformation import dataflow, interstate, subgraph from dace.transformation.interstate import InlineSDFG diff --git a/tests/transformations/prune_connectors_test.py b/tests/transformations/prune_connectors_test.py index 43c306e581..8a7be6c12c 100644 --- a/tests/transformations/prune_connectors_test.py +++ b/tests/transformations/prune_connectors_test.py @@ -3,7 +3,6 @@ import numpy as np import os import copy -import pytest from typing import Tuple import dace diff --git a/tests/transformations/redundant_copy_data_races_test.py b/tests/transformations/redundant_copy_data_races_test.py index 212fb3c89e..30a93d2cc0 100644 --- a/tests/transformations/redundant_copy_data_races_test.py +++ b/tests/transformations/redundant_copy_data_races_test.py @@ -1,6 +1,5 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. import dace -import numpy as np from dace import nodes diff --git a/tests/transformations/redundant_copy_test.py b/tests/transformations/redundant_copy_test.py index d9aa132a5d..00ca729f7d 100644 --- a/tests/transformations/redundant_copy_test.py +++ b/tests/transformations/redundant_copy_test.py @@ -1,4 +1,4 @@ -# Copyright 2019-2023 ETH Zurich and the DaCe authors. All rights reserved. +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. import numpy as np import pytest import copy @@ -6,7 +6,7 @@ import dace from dace import nodes, data as dace_data -from dace.libraries.standard import Transpose +from dace.libraries.linalg import Transpose from dace.transformation.dataflow import (RedundantArray, RedundantSecondArray, RedundantArrayCopying, RedundantArrayCopyingIn) diff --git a/tests/transformations/redundant_slices_test.py b/tests/transformations/redundant_slices_test.py index b8de6f0ba7..5085011739 100644 --- a/tests/transformations/redundant_slices_test.py +++ b/tests/transformations/redundant_slices_test.py @@ -6,8 +6,7 @@ import dace from dace import data, nodes -from dace.transformation.dataflow import RedundantReadSlice, RedundantWriteSlice, RemoveSliceView -from dace.transformation.interstate import InlineMultistateSDFG, InlineSDFG, StateFusion +from dace.transformation.dataflow import RedundantReadSlice, RedundantWriteSlice from dace.sdfg import utils as sdutil diff --git a/tests/transformations/state_fission_test.py b/tests/transformations/state_fission_test.py index ffd6625ee6..ef3629fccd 100644 --- a/tests/transformations/state_fission_test.py +++ b/tests/transformations/state_fission_test.py @@ -9,7 +9,7 @@ from dace.sdfg import nodes, graph from dace.transformation import helpers, dataflow -from .utility import count_nodes, unique_name, make_sdfg_args, compile_and_run_sdfg +from .utility import count_nodes, unique_name def make_vecAdd_sdfg(symbol_name: str, sdfg_name: str, access_nodes_dict: dict, dtype=dace.float32): diff --git a/tests/transformations/state_fusion_extended_test.py b/tests/transformations/state_fusion_extended_test.py index 20ad64acf8..ab8c9c06c0 100644 --- a/tests/transformations/state_fusion_extended_test.py +++ b/tests/transformations/state_fusion_extended_test.py @@ -1,3 +1,4 @@ +# Copyright 2019-2026 ETH Zurich and the DaCe authors. All rights reserved. from dace import SDFG, InterstateEdge, Memlet from dace import dtypes from dace.transformation.interstate import StateFusionExtended @@ -61,5 +62,52 @@ def test_extended_fusion(): assert sdfg.number_of_nodes() == 1 +def test_extended_fusion_refuses_unsafe_write_after_read(): + """A read state followed by an in-place write of the same array must + not be fused (write-after-read / anti-dependency). + + ``s1`` reads ``A[k]`` in two tasklets; the later ``s2`` does + ``A[k] = A[k] + 1``. Safely fusing would need a dependency edge to + every first-state sink reading ``A`` (arbitrarily fanned out), so + ``StateFusionExtended`` must refuse and leave the two states + separate; the interstate edge then keeps the write ordered after the + reads. Previously it fused without that ordering and the write + clobbered the still-pending reads. + """ + sdfg = SDFG('state_fusion_war_ordering') + sdfg.add_array('A', [8], dtypes.float64) + sdfg.add_array('B', [8], dtypes.float64) + sdfg.add_array('C', [8], dtypes.float64) + sdfg.add_symbol('k', dtypes.int64) + + s1 = sdfg.add_state('read_A', is_start_block=True) + s2 = sdfg.add_state('increment_A') + sdfg.add_edge(s1, s2, InterstateEdge()) + + ar_b = s1.add_read('A') + bw = s1.add_write('B') + tb = s1.add_tasklet('rb', {'_in'}, {'_out'}, '_out = _in') + s1.add_edge(ar_b, None, tb, '_in', Memlet('A[k]')) + s1.add_edge(tb, '_out', bw, None, Memlet('B[k]')) + + ar_c = s1.add_read('A') + cw = s1.add_write('C') + tc = s1.add_tasklet('rc', {'_in'}, {'_out'}, '_out = _in') + s1.add_edge(ar_c, None, tc, '_in', Memlet('A[k]')) + s1.add_edge(tc, '_out', cw, None, Memlet('C[k]')) + + ar2 = s2.add_read('A') + aw2 = s2.add_write('A') + ti = s2.add_tasklet('inc', {'_in'}, {'_out'}, '_out = _in + 1.0') + s2.add_edge(ar2, None, ti, '_in', Memlet('A[k]')) + s2.add_edge(ti, '_out', aw2, None, Memlet('A[k]')) + sdfg.validate() + + applied = sdfg.apply_transformations_repeated(StateFusionExtended) + assert applied == 0, 'write-after-read fusion must be refused' + assert sdfg.number_of_nodes() == 2, 'the two states must remain separate' + + if __name__ == '__main__': test_extended_fusion() + test_extended_fusion_refuses_unsafe_write_after_read() diff --git a/tests/transformations/state_fusion_test.py b/tests/transformations/state_fusion_test.py index aece80f574..a17d0d77e2 100644 --- a/tests/transformations/state_fusion_test.py +++ b/tests/transformations/state_fusion_test.py @@ -1,6 +1,5 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. import dace -from dace.transformation import transformation from dace.transformation.interstate import StateFusion import networkx as nx import numpy as np diff --git a/tests/transformations/strip_mining_test.py b/tests/transformations/strip_mining_test.py index 7d47b1e24f..6e7dbb271b 100644 --- a/tests/transformations/strip_mining_test.py +++ b/tests/transformations/strip_mining_test.py @@ -11,11 +11,13 @@ def test_strip_mining(): """ # 1. The program sdfg = dace.SDFG("assign") + sdfg.add_array("A", (64, ), dtype=dace.uint32) + sdfg.add_array("B", (1, ), dtype=dace.uint32) state = sdfg.add_state("main") # inputs - A = state.add_array("A", (64, ), dtype=dace.uint32) - B = state.add_array("B", (1, ), dtype=dace.uint32) + A = state.add_access("A") + B = state.add_access("B") # kernel map map_entry, map_exit = state.add_map("map", dict(i="0:64")) diff --git a/tests/transformations/subgraph_fusion/complex_test.py b/tests/transformations/subgraph_fusion/complex_test.py index 54ebcf1242..d3472a8e78 100644 --- a/tests/transformations/subgraph_fusion/complex_test.py +++ b/tests/transformations/subgraph_fusion/complex_test.py @@ -1,13 +1,9 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. import dace -import dace.transformation.subgraph.helpers as helpers -import dace.sdfg.nodes as nodes import numpy as np from dace.sdfg.graph import SubgraphView -import sys - from dace.transformation.subgraph import SubgraphFusion from util import expand_maps, expand_reduce, fusion diff --git a/tests/transformations/subgraph_fusion/create_out_transient_test.py b/tests/transformations/subgraph_fusion/create_out_transient_test.py index d331c92e1d..c3d33b4fde 100644 --- a/tests/transformations/subgraph_fusion/create_out_transient_test.py +++ b/tests/transformations/subgraph_fusion/create_out_transient_test.py @@ -1,14 +1,7 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. import dace -from dace.transformation.subgraph import SubgraphFusion -import dace.transformation.subgraph.helpers as helpers -import dace.sdfg.nodes as nodes import numpy as np -from dace.sdfg.graph import SubgraphView -from dace.transformation.interstate import StateFusion -from typing import List, Union -import sys from util import fusion N, M, O = [dace.symbol(s) for s in ['N', 'M', 'O']] diff --git a/tests/transformations/subgraph_fusion/disjoint_test.py b/tests/transformations/subgraph_fusion/disjoint_test.py index 7dfb015a72..3ea9757fb4 100644 --- a/tests/transformations/subgraph_fusion/disjoint_test.py +++ b/tests/transformations/subgraph_fusion/disjoint_test.py @@ -1,15 +1,9 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. -from typing import List, Union import numpy as np -from util import expand_maps, expand_reduce, fusion import dace -import dace.libraries.standard as stdlib -import dace.sdfg.nodes as nodes -import dace.transformation.subgraph.helpers as helpers from dace.sdfg.graph import SubgraphView -from dace.transformation.dataflow import ReduceExpansion from dace.transformation.subgraph import SubgraphFusion M = dace.symbol('M') diff --git a/tests/transformations/subgraph_fusion/expansion_test.py b/tests/transformations/subgraph_fusion/expansion_test.py index 4afa6b7793..a41abe27bc 100644 --- a/tests/transformations/subgraph_fusion/expansion_test.py +++ b/tests/transformations/subgraph_fusion/expansion_test.py @@ -1,15 +1,8 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. import dace -import dace.transformation.subgraph.helpers as helpers -import dace.sdfg.nodes as nodes import numpy as np -from dace.sdfg.graph import SubgraphView - -import sys - -from dace.transformation.subgraph import SubgraphFusion -from util import expand_maps, expand_reduce, fusion +from util import expand_maps N, M, O = [dace.symbol(s) for s in ['N', 'M', 'O']] diff --git a/tests/transformations/subgraph_fusion/fission_io_test.py b/tests/transformations/subgraph_fusion/fission_io_test.py index 8fae813d30..45d6ddd348 100644 --- a/tests/transformations/subgraph_fusion/fission_io_test.py +++ b/tests/transformations/subgraph_fusion/fission_io_test.py @@ -1,16 +1,11 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. from copy import deepcopy import dace -from dace.sdfg import nodes from dace.transformation.dataflow import MapFission -from dace.transformation.helpers import nest_state_subgraph import numpy as np -import unittest -from typing import Union, List from dace.sdfg.graph import SubgraphView from dace.transformation.subgraph import SubgraphFusion -import dace.transformation.subgraph.helpers as helpers from util import fusion diff --git a/tests/transformations/subgraph_fusion/fission_offsets1_test.py b/tests/transformations/subgraph_fusion/fission_offsets1_test.py index 6f3718ccb7..ffcbc1cc53 100644 --- a/tests/transformations/subgraph_fusion/fission_offsets1_test.py +++ b/tests/transformations/subgraph_fusion/fission_offsets1_test.py @@ -1,16 +1,10 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. -import copy import dace -from dace.sdfg import nodes from dace.transformation.dataflow import MapFission -from dace.transformation.helpers import nest_state_subgraph import numpy as np -import unittest -from typing import Union, List from dace.sdfg.graph import SubgraphView from dace.transformation.subgraph import SubgraphFusion -import dace.transformation.subgraph.helpers as helpers from util import fusion diff --git a/tests/transformations/subgraph_fusion/fission_offsets2_test.py b/tests/transformations/subgraph_fusion/fission_offsets2_test.py index e3d9d9490d..f6f929abec 100644 --- a/tests/transformations/subgraph_fusion/fission_offsets2_test.py +++ b/tests/transformations/subgraph_fusion/fission_offsets2_test.py @@ -1,16 +1,10 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. -import copy import dace -from dace.sdfg import nodes from dace.transformation.dataflow import MapFission -from dace.transformation.helpers import nest_state_subgraph import numpy as np -import unittest -from typing import Union, List from dace.sdfg.graph import SubgraphView from dace.transformation.subgraph import SubgraphFusion -import dace.transformation.subgraph.helpers as helpers from util import fusion diff --git a/tests/transformations/subgraph_fusion/fission_subgraph_test.py b/tests/transformations/subgraph_fusion/fission_subgraph_test.py index 43ec8f7d55..589098ba95 100644 --- a/tests/transformations/subgraph_fusion/fission_subgraph_test.py +++ b/tests/transformations/subgraph_fusion/fission_subgraph_test.py @@ -1,17 +1,11 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. -import copy from copy import deepcopy as dcpy import dace -from dace.sdfg import nodes from dace.transformation.dataflow import MapFission -from dace.transformation.helpers import nest_state_subgraph import numpy as np -import unittest -from typing import Union, List from dace.sdfg.graph import SubgraphView from dace.transformation.subgraph import SubgraphFusion -import dace.transformation.subgraph.helpers as helpers from util import fusion diff --git a/tests/transformations/subgraph_fusion/intermediate_mimo_test.py b/tests/transformations/subgraph_fusion/intermediate_mimo_test.py index d28e631469..fdee971346 100644 --- a/tests/transformations/subgraph_fusion/intermediate_mimo_test.py +++ b/tests/transformations/subgraph_fusion/intermediate_mimo_test.py @@ -1,16 +1,10 @@ # Copyright 2019-2023 ETH Zurich and the DaCe authors. All rights reserved. -import copy import dace -from dace.sdfg import nodes from dace.sdfg.graph import SubgraphView -from dace.transformation.helpers import nest_state_subgraph import numpy as np -import unittest -import sys from dace.transformation.subgraph import MultiExpansion, SubgraphFusion -from typing import Union, List from dace.sdfg.graph import SubgraphView N = dace.symbol('N') diff --git a/tests/transformations/subgraph_fusion/invariant_dimension_test.py b/tests/transformations/subgraph_fusion/invariant_dimension_test.py index c5653bfc86..f5c725cab9 100644 --- a/tests/transformations/subgraph_fusion/invariant_dimension_test.py +++ b/tests/transformations/subgraph_fusion/invariant_dimension_test.py @@ -2,17 +2,13 @@ import dace from dace.transformation.subgraph import SubgraphFusion import dace.sdfg.utils as utils -import dace.transformation.subgraph.helpers as helpers -import dace.sdfg.nodes as nodes import dace.subsets as subsets import numpy as np import itertools from dace.sdfg.graph import SubgraphView -from typing import Union, List -import sys from util import fusion N, M, O = [dace.symbol(s) for s in ['N', 'M', 'O']] diff --git a/tests/transformations/subgraph_fusion/parallel_test.py b/tests/transformations/subgraph_fusion/parallel_test.py index fae9cd2196..1cd5e51d93 100644 --- a/tests/transformations/subgraph_fusion/parallel_test.py +++ b/tests/transformations/subgraph_fusion/parallel_test.py @@ -1,10 +1,8 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. import dace from dace.transformation.subgraph import MultiExpansion, SubgraphFusion -import dace.sdfg.nodes as nodes import numpy as np -from typing import Union, List from dace.sdfg.graph import SubgraphView N, M, O, P, Q, R = [dace.symbol(s) for s in ['N', 'M', 'O', 'P', 'Q', 'R']] diff --git a/tests/transformations/subgraph_fusion/reduction_fuse_test.py b/tests/transformations/subgraph_fusion/reduction_fuse_test.py index 1e2597d510..0a5e9752e4 100644 --- a/tests/transformations/subgraph_fusion/reduction_fuse_test.py +++ b/tests/transformations/subgraph_fusion/reduction_fuse_test.py @@ -1,16 +1,10 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. -from typing import List, Union import numpy as np import pytest from util import expand_maps, expand_reduce, fusion import dace -import dace.libraries.standard as stdlib -import dace.sdfg.nodes as nodes -import dace.transformation.subgraph.helpers as helpers -from dace.sdfg.graph import SubgraphView -from dace.transformation.dataflow import ReduceExpansion M = dace.symbol('M') N = dace.symbol('N') diff --git a/tests/transformations/subgraph_fusion/reduction_test.py b/tests/transformations/subgraph_fusion/reduction_test.py index 4e4cdd7d7a..753fc71213 100644 --- a/tests/transformations/subgraph_fusion/reduction_test.py +++ b/tests/transformations/subgraph_fusion/reduction_test.py @@ -1,15 +1,10 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. -from typing import List, Union import numpy as np import pytest -from util import expand_maps, expand_reduce, fusion +from util import expand_reduce import dace -import dace.libraries.standard as stdlib -import dace.sdfg.nodes as nodes -import dace.transformation.subgraph.helpers as helpers -from dace.sdfg.graph import SubgraphView from dace.transformation.dataflow import ReduceExpansion M = dace.symbol('M') diff --git a/tests/transformations/subgraph_fusion/sequential1_test.py b/tests/transformations/subgraph_fusion/sequential1_test.py index c870475b0e..f87ba2d714 100644 --- a/tests/transformations/subgraph_fusion/sequential1_test.py +++ b/tests/transformations/subgraph_fusion/sequential1_test.py @@ -1,10 +1,5 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. import dace -import dace.transformation.subgraph.helpers as helpers -from dace.transformation.subgraph import SubgraphFusion -from dace.sdfg.graph import SubgraphView -import dace.sdfg.nodes as nodes -from typing import List, Union import numpy as np from util import fusion diff --git a/tests/transformations/subgraph_fusion/sequential2_cudatest.py b/tests/transformations/subgraph_fusion/sequential2_cudatest.py index 8f66146f59..07f308fa01 100644 --- a/tests/transformations/subgraph_fusion/sequential2_cudatest.py +++ b/tests/transformations/subgraph_fusion/sequential2_cudatest.py @@ -1,11 +1,6 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. import dace -from dace.transformation.subgraph import SubgraphFusion -import dace.transformation.subgraph.helpers as helpers -from dace.sdfg.graph import SubgraphView -import dace.sdfg.nodes as nodes import numpy as np -from typing import List, Union import pytest from util import fusion diff --git a/tests/transformations/subgraph_fusion/sequential2_test.py b/tests/transformations/subgraph_fusion/sequential2_test.py index a00664b5cb..de94e932b9 100644 --- a/tests/transformations/subgraph_fusion/sequential2_test.py +++ b/tests/transformations/subgraph_fusion/sequential2_test.py @@ -1,10 +1,5 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. import dace -from dace.transformation.subgraph import SubgraphFusion -from dace.sdfg.graph import SubgraphView -import dace.transformation.subgraph.helpers as helpers -from typing import List, Union -import dace.sdfg.nodes as nodes import numpy as np from util import fusion diff --git a/tests/transformations/subgraph_fusion/smax_test.py b/tests/transformations/subgraph_fusion/smax_test.py index be2b533036..c46cc9fa67 100644 --- a/tests/transformations/subgraph_fusion/smax_test.py +++ b/tests/transformations/subgraph_fusion/smax_test.py @@ -1,18 +1,12 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. -import sys -from typing import List, Union import numpy as np from util import expand_maps, expand_reduce, fusion import dace -import dace.dtypes as dtypes import dace.libraries.standard as stdlib import dace.sdfg.nodes as nodes -import dace.transformation.subgraph.helpers as helpers from dace.sdfg.graph import SubgraphView -from dace.transformation.dataflow import ReduceExpansion -from dace.transformation.subgraph import MultiExpansion, SubgraphFusion dace_dtype = dace.float32 H, B, SN, SM = (dace.symbol(s) for s in ('H', 'B', 'SN', 'SM')) diff --git a/tests/transformations/subgraph_fusion/tiling_pool_test.py b/tests/transformations/subgraph_fusion/tiling_pool_test.py index 6b9500c021..0de6a16aae 100644 --- a/tests/transformations/subgraph_fusion/tiling_pool_test.py +++ b/tests/transformations/subgraph_fusion/tiling_pool_test.py @@ -1,7 +1,6 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. import dace -from dace.sdfg import SDFG from dace.transformation.subgraph.stencil_tiling import StencilTiling from dace.transformation.subgraph import SubgraphFusion from dace.sdfg.graph import SubgraphView diff --git a/tests/transformations/subgraph_fusion/tiling_stencil_test.py b/tests/transformations/subgraph_fusion/tiling_stencil_test.py index 694e5c2460..0609530b06 100644 --- a/tests/transformations/subgraph_fusion/tiling_stencil_test.py +++ b/tests/transformations/subgraph_fusion/tiling_stencil_test.py @@ -3,7 +3,6 @@ import contextlib import dace -from dace.sdfg import SDFG from dace.transformation.subgraph.stencil_tiling import StencilTiling from dace.transformation.subgraph import SubgraphFusion from dace.sdfg.graph import SubgraphView diff --git a/tests/transformations/sve/sve_vectorization_test.py b/tests/transformations/sve/sve_vectorization_test.py index 6645aa3a0b..267c79acf3 100644 --- a/tests/transformations/sve/sve_vectorization_test.py +++ b/tests/transformations/sve/sve_vectorization_test.py @@ -1,7 +1,5 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. -from dace.sdfg.graph import NodeNotFoundError import dace -import numpy as np from dace.transformation.dataflow.sve.vectorization import SVEVectorization from dace import SDFG import dace.dtypes as dtypes diff --git a/tests/transformations/trivial_loop_elimination_test.py b/tests/transformations/trivial_loop_elimination_test.py index 6f2769f921..f08f0e91c7 100644 --- a/tests/transformations/trivial_loop_elimination_test.py +++ b/tests/transformations/trivial_loop_elimination_test.py @@ -1,8 +1,6 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. -from dace.sdfg.nodes import MapEntry import dace from dace.transformation.interstate import TrivialLoopElimination -from dace.symbolic import pystr_to_symbolic import unittest import numpy as np diff --git a/tests/transformations/utility.py b/tests/transformations/utility.py index 17efaf4ecd..9bf61b0d8a 100644 --- a/tests/transformations/utility.py +++ b/tests/transformations/utility.py @@ -1,6 +1,6 @@ # Copyright 2019-2025 ETH Zurich and the DaCe authors. All rights reserved. """Helper functions for compilation. """ -from typing import Any, Union, Tuple, Type, Optional, List, Dict +from typing import Any, Union, Tuple, Type, List import copy import numpy as np diff --git a/tests/trivial_map_elimination_test.py b/tests/trivial_map_elimination_test.py index b1bdeab77f..5a0ed1c6bd 100644 --- a/tests/trivial_map_elimination_test.py +++ b/tests/trivial_map_elimination_test.py @@ -1,6 +1,5 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. import dace -from dace.sdfg import nodes from dace.transformation.dataflow import TrivialMapElimination import unittest diff --git a/tests/typed_constant_type_inference_test.py b/tests/typed_constant_type_inference_test.py new file mode 100644 index 0000000000..1a5265c788 --- /dev/null +++ b/tests/typed_constant_type_inference_test.py @@ -0,0 +1,25 @@ +import numpy as np +import pytest + +from dace import dtypes, symbolic +from dace.sdfg import type_inference + + +def test_infer_expr_type_of_typed_constant(): + expr = symbolic.TypedConstant(np.int16(2)) + + inferred = type_inference.infer_expr_type(expr) + + assert inferred == dtypes.int16 + + +def test_infer_expr_type_with_typed_constant_expression(): + expr = symbolic.symbol('N', dtype=dtypes.int32) + symbolic.TypedConstant(np.uint64(1)) + + inferred = type_inference.infer_expr_type(expr, {'N': dtypes.int32}) + + assert inferred == dtypes.uint64 + + +if __name__ == '__main__': + pytest.main([__file__]) diff --git a/tests/undefined_symbol_test.py b/tests/undefined_symbol_test.py index bac734f155..e5246e0898 100644 --- a/tests/undefined_symbol_test.py +++ b/tests/undefined_symbol_test.py @@ -2,10 +2,8 @@ # Copyright 2019-2025 ETH Zurich and the DaCe authors. All rights reserved. import pytest -import numpy as np import sympy -import dace from dace import symbolic diff --git a/tests/unique_nested_sdfg_test.py b/tests/unique_nested_sdfg_test.py index 410b9829ba..3ab0551f9f 100644 --- a/tests/unique_nested_sdfg_test.py +++ b/tests/unique_nested_sdfg_test.py @@ -5,8 +5,6 @@ import dace import numpy as np -import argparse -import subprocess from dace.memlet import Memlet diff --git a/tests/utils.py b/tests/utils.py index b4ddb24314..28ff777eda 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -1,10 +1,8 @@ # Copyright 2019-2025 ETH Zurich and the DaCe authors. All rights reserved. import os -import typing import urllib.request, urllib.parse import pathlib -import pytest import dace import numpy as np diff --git a/tests/utils/array_dimension_utils.py b/tests/utils/array_dimension_utils.py index 903e53db76..f657c91fd9 100644 --- a/tests/utils/array_dimension_utils.py +++ b/tests/utils/array_dimension_utils.py @@ -1,7 +1,6 @@ # Copyright 2019-2025 ETH Zurich and the DaCe authors. All rights reserved. import dace -import pytest import tempfile import os diff --git a/tests/utils/is_contiguous_subset_test.py b/tests/utils/is_contiguous_subset_test.py index 1284d6907b..10fdb841c9 100644 --- a/tests/utils/is_contiguous_subset_test.py +++ b/tests/utils/is_contiguous_subset_test.py @@ -1,7 +1,6 @@ # Copyright 2019-2025 ETH Zurich and the DaCe authors. All rights reserved. import dace -import pytest import dace.subsets @@ -235,6 +234,21 @@ def test_non_standard_strides_returns_false(): assert result is False +def test_1d_slice(): + """Test a non packed 1D slice.""" + array = create_array((10, 10, 10), (1, 20, 400)) + subset_base = [(3, 8, 1), (4, 4, 1), (5, 5, 1)] + + for res, sbs_func in [ + (True, lambda x: x), # a[3:9, 4, 5] + (False, reversed), # a[5, 4, 3:9] + (False, lambda x: (x[0], (4, 5, 1), x[2])), # a[3:9, 4:6, 5] + ]: + subset = create_subset(sbs_func(subset_base)) + result = subset.is_contiguous_subset(array) + assert result is res + + if __name__ == "__main__": tests = [obj for name, obj in globals().items() if callable(obj) and name.startswith("test_")] for test_function in tests: diff --git a/tests/utils/specialize_scalar_test.py b/tests/utils/specialize_scalar_test.py index 8f99f154d9..96831cf1ce 100644 --- a/tests/utils/specialize_scalar_test.py +++ b/tests/utils/specialize_scalar_test.py @@ -3,7 +3,6 @@ from dace.codegen.control_flow import LoopRegion from dace.properties import CodeBlock from dace.sdfg.state import ConditionalBlock -import pytest import dace.sdfg.utils as sdutil diff --git a/tests/vector_min_test.py b/tests/vector_min_test.py index 9c526f1ff4..453ba2766f 100644 --- a/tests/vector_min_test.py +++ b/tests/vector_min_test.py @@ -1,5 +1,4 @@ # Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. -import math import numpy as np import dace as dp diff --git a/tests/wcr_cudatest.py b/tests/wcr_cudatest.py index b3604ae88b..0b1a419cd0 100644 --- a/tests/wcr_cudatest.py +++ b/tests/wcr_cudatest.py @@ -3,7 +3,6 @@ import dace from dace.transformation.interstate import GPUTransformSDFG -from typing import Dict, Tuple import pytest diff --git a/tutorials/getting_started.ipynb b/tutorials/getting_started.ipynb index b5d9c000f1..f2d2b0e41d 100644 --- a/tutorials/getting_started.ipynb +++ b/tutorials/getting_started.ipynb @@ -279,7 +279,7 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "A compiled SDFG, however, has to be invoked like an SDFG, with keyword arguments only:" + "A compiled SDFG can be invoked just like the original `@dace.program`. When the SDFG is generated by one of the frontends (such as the Python frontend used here), the original argument order is preserved in `sdfg.arg_names`, so positional arguments work as expected. If the SDFG was constructed directly with the SDFG API and `sdfg.arg_names` was not set, the compiled SDFG must instead be invoked with keyword arguments:" ] }, { diff --git a/tutorials/transformations.ipynb b/tutorials/transformations.ipynb index 931df79e18..ade02d65bc 100644 --- a/tutorials/transformations.ipynb +++ b/tutorials/transformations.ipynb @@ -685,7 +685,6 @@ "source": [ "from dace.transformation import transformation as xf\n", "from dace.sdfg import utils as sdutil\n", - "from dace import registry\n", "\n", "\n", "class RedundantArrayTasklet(xf.SingleStateTransformation):\n",