Machine Learning Integration for DaCe (Autodiff - ONNX - PyTorch)

.github/workflows/copilot-setup-steps.yml

@@ -35,6 +35,6 @@ jobs:
 
       - name: Install DaCe in development mode
         run: |
-          python -m pip install --editable ".[testing,linting]"
+          python -m pip install --editable ".[testing,linting,ml]"
           pre-commit install
           pre-commit run

.github/workflows/fpga-ci.yml

@@ -32,7 +32,7 @@ jobs:
         python -m pip install --upgrade pip
         pip install pytest-xdist flake8 coverage click
         pip uninstall -y dace
-        pip install -e ".[testing]"
+        pip install -e ".[testing,ml]"
         curl -Os https://uploader.codecov.io/latest/linux/codecov
         chmod +x codecov
 

.github/workflows/general-ci.yml

@@ -59,7 +59,7 @@ jobs:
         else
             export DACE_optimizer_automatic_simplification=${{ matrix.simplify }}
         fi
-        pytest -n auto --cov-report=xml --cov=dace --tb=short --timeout_method thread --timeout=300 -m "not gpu and not verilator and not tensorflow and not mkl and not sve and not papi and not mlir and not lapack and not fpga and not mpi and not rtl_hardware and not scalapack and not datainstrument and not long and not sequential"
+        pytest -n auto --cov-report=xml --cov=dace --tb=short --timeout_method thread --timeout=300 -m "not gpu and not autodiff and not torch and not onnx and not verilator and not tensorflow and not mkl and not sve and not papi and not mlir and not lapack and not fpga and not mpi and not rtl_hardware and not scalapack and not datainstrument and not long and not sequential"
         ./codecov
 
     - name: Test OpenBLAS LAPACK

.github/workflows/gpu-ci.yml

@@ -37,7 +37,7 @@ jobs:
         pip install mpi4py
         pip install cupy
         pip uninstall -y dace
-        pip install -e ".[testing]"
+        pip install -e ".[testing,ml]"
         curl -Os https://uploader.codecov.io/latest/linux/codecov
         chmod +x codecov
 

.github/workflows/ml-ci.yml

@@ -0,0 +1,62 @@
+name: Machine Learning and Autodiff Tests
+
+on:
+  push:
+    branches: [ main, ci-fix ]
+  pull_request:
+    branches: [ main, ci-fix ]
+  merge_group:
+    branches: [ main, ci-fix ]
+
+concurrency:
+  group: ${{github.workflow}}-${{github.ref}}
+  cancel-in-progress: true
+
+jobs:
+  test:
+    if: "!contains(github.event.pull_request.labels.*.name, 'no-ci')"
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ['3.13']
+        simplify: [0,1,autoopt]
+
+    steps:
+    - uses: actions/checkout@v4
+      with:
+        submodules: 'recursive'
+    - name: Set up Python ${{ matrix.python-version }}
+      uses: actions/setup-python@v5
+      with:
+        python-version: ${{ matrix.python-version }}
+    - name: Install dependencies
+      run: |
+        sudo apt-get update
+        sudo apt-get install -y libyaml-dev cmake
+        sudo apt-get install -y libblas-dev libopenblas-dev liblapacke-dev
+        python -m pip install --upgrade pip
+        pip install flake8 pytest-xdist coverage
+        pip install -e ".[ml-testing,ml]"
+        curl -Os https://uploader.codecov.io/latest/linux/codecov
+        chmod +x codecov
+
+    - name: Test with pytest
+      run: |
+        export NOSTATUSBAR=1
+        export DACE_testing_serialization=1
+        export DACE_testing_deserialize_exception=1
+        export DACE_cache=unique
+        if [ "${{ matrix.simplify }}" = "autoopt" ]; then
+            export DACE_optimizer_automatic_simplification=1
+            export DACE_optimizer_autooptimize=1
+            echo "Auto-optimization heuristics"
+        else
+            export DACE_optimizer_automatic_simplification=${{ matrix.simplify }}
+        fi
+        pytest --cov-report=xml --cov=dace --tb=short --timeout_method thread --timeout=300 -v -m "(torch or onnx or autodiff) and not gpu"
+        ./codecov
+
+    - uses: codecov/codecov-action@v4
+      with:
+        token: ${{ secrets.CODECOV_TOKEN }}
+        verbose: true

.gitignore

@@ -195,3 +195,8 @@ _build/
 
 # Ignoring the test junk
 _all_tests/
+
+
+# Ignore downloaded ONNX models
+/*.onnx
+/*.bin

dace/__init__.py

@@ -35,6 +35,17 @@
     sys.path.insert(0, __external_transformations_path__)
 
 
+# Lazy loading for ml module to avoid eager TensorFlow/PyTorch imports
+def __getattr__(name):
+    if name == 'ml':
+        import importlib
+        ml_module = importlib.import_module('.ml', package='dace')
+        # Cache the module to avoid re-importing
+        globals()['ml'] = ml_module
+        return ml_module
+    raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
+
+
 # Hack that enables using @dace as a decorator
 # See https://stackoverflow.com/a/48100440/6489142
 class DaceModule(sys.modules[__name__].__class__):

dace/autodiff/__init__.py

@@ -0,0 +1,58 @@
+# Copyright 2019-2025 ETH Zurich and the DaCe authors. All rights reserved.
+"""
+DaCe Automatic Differentiation (AD) System.
+
+This module provides reverse-mode automatic differentiation for DaCe programs,
+enabling automatic computation of gradients for optimized numerical kernels.
+
+Main Components
+---------------
+- **add_backward_pass**: Main entry point for adding backward pass to an SDFG
+- **BackwardPassGenerator**: Core algorithm for generating backward passes
+- **BackwardImplementation**: ABC for implementing operation-specific backward rules
+- **BackwardContext**: Context information for backward pass generation
+- **BackwardResult**: Result of backward pass generation with forward/backward SDFGs
+- **AutoDiffException**: Base exception for autodiff errors
+
+Key Features
+------------
+- Support for control flow (loops, conditionals)
+- Data forwarding strategies (store vs recompute tradeoffs)
+- Extensible backward implementations for library nodes
+- Integration with PyTorch autograd
+- Automatic memory management for intermediate values
+
+
+"""
+
+from .base_abc import BackwardImplementation, BackwardContext, BackwardResult, AutoDiffException
+from .backward_pass_generator import BackwardPassGenerator
+from .autodiff import add_backward_pass
+
+try:
+    from .torch import make_backward_function
+    TORCH_INTEGRATION_AVAILABLE = True
+except ImportError:
+    make_backward_function = None
+    TORCH_INTEGRATION_AVAILABLE = False
+
+import sys
+from . import library
+
+__all__ = [
+    # Main API
+    "add_backward_pass",
+    # Core classes
+    "BackwardPassGenerator",
+    "BackwardContext",
+    "BackwardResult",
+    # Extension points
+    "BackwardImplementation",
+    # Exceptions
+    "AutoDiffException",
+    # Submodules
+    "library",
+]
+
+if TORCH_INTEGRATION_AVAILABLE:
+    __all__.append("make_backward_function")

dace/autodiff/analysis.py

@@ -0,0 +1,103 @@
+# Copyright 2019-2025 ETH Zurich and the DaCe authors. All rights reserved.
+"""
+Analysis helpers for autodiff
+"""
+from typing import Dict, Set, Tuple, Optional
+import collections
+
+import networkx as nx
+
+from dace.sdfg import SDFG, SDFGState, nodes, utils as sdfg_utils
+from dace.transformation.passes import analysis
+from dace.sdfg.state import FunctionCallRegion
+
+AccessSets = Dict[SDFGState, Tuple[Set[str], Set[str]]]
+
+
+def dependency_analysis(sdfg: SDFG) -> Dict[str, Set[str]]:
+    """
+    Analyze read dependencies of arrays in the SDFG.
+
+    :param sdfg: SDFG to analyze
+    :return: A dictionary mapping array names to a list of read dependencies.
+    """
+
+    # FIXME can be made more efficient
+    dependencies = nx.DiGraph()
+    for sdfg_node in sdfg.nodes():
+        if isinstance(sdfg_node, SDFGState):
+            for node in sdfg_node.data_nodes():
+                for edge in sdfg_node.edge_bfs(node, reverse=True):
+                    dependencies.add_edge(node.data, edge.data.data)
+        elif isinstance(sdfg_node, FunctionCallRegion):
+            for state in sdfg_node.nodes():
+                assert isinstance(state, SDFGState)
+                for node in state.data_nodes():
+                    for edge in state.edge_bfs(node, reverse=True):
+                        dependencies.add_edge(node.data, edge.data.data)
+
+    dependencies = nx.transitive_closure(dependencies)
+    result = {}
+    for array in dependencies:
+        result[array] = {nbr for nbr in dependencies.neighbors(array)}
+    return result
+
+
+def inverse_reachability(sdfg: SDFG) -> Dict[SDFGState, Set[SDFGState]]:
+
+    reachability = analysis.StateReachability().apply_pass(sdfg, {})
+    inverse_reachability = collections.defaultdict(set)
+    # iterate over cfg_ids
+    for cfg_id in reachability.keys():
+        for pred, successors in reachability[cfg_id].items():
+            for successor in successors:
+                inverse_reachability[successor].add(pred)
+
+    return inverse_reachability
+
+
+def is_previously_written(sdfg: SDFG,
+                          state: SDFGState,
+                          node: nodes.Node,
+                          array_name: str,
+                          access_sets: Optional[AccessSets] = None) -> bool:
+    """
+    Determine whether the given array name was written before the current node.
+
+    :param sdfg: the sdfg containing the node
+    :param state: the state containing the node
+    :param node: the node to check
+    :param array_name: the array name to check
+    :return: True if the array was written before the node, False otherwise.
+    """
+
+    if access_sets is None:
+        access_sets = analysis.AccessSets().apply_pass(sdfg, {})
+
+    reachable = inverse_reachability(sdfg)
+
+    # Check the current state
+    for subgraph in sdfg_utils.concurrent_subgraphs(state):
+        if node in subgraph.nodes():
+            # Get all the access nodes in the subgraph to the same data
+            for other_node in subgraph.data_nodes():
+                if other_node != node and other_node.data == array_name:
+                    # Check if this is a write node
+                    for in_edge in subgraph.in_edges(other_node):
+                        if in_edge.data.data == array_name:
+                            # Check if there's a path to our node,
+                            # since we only care about writes that happen before the current node
+                            if nx.has_path(state.nx, other_node, node):
+                                return True
+        else:
+            # This is not our current subgraph, check the write states
+            _, write_set = subgraph.read_and_write_sets()
+            if array_name in write_set:
+                return True
+
+    # Check other states
+    for predecessor in reachable[state]:
+        _, write_set = access_sets[predecessor]
+        if array_name in write_set:
+            return True
+    return False