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

@@ -38,7 +38,7 @@ jobs:
         sudo apt-get install -y verilator # RTL simulation dependencies
         python -m pip install --upgrade pip
         pip install flake8 pytest-xdist coverage
-        pip install -e ".[testing]"
+        pip install -e ".[testing,ml]"
         curl -Os https://uploader.codecov.io/latest/linux/codecov
         chmod +x codecov
 

.github/workflows/gpu-ci.yml

@@ -36,7 +36,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/hardware_test.yml

@@ -15,7 +15,7 @@ jobs:
         python -m pip install --upgrade pip
         pip install pytest-xdist flake8
         pip uninstall -y dace
-        pip install -e ".[testing]"
+        pip install -e ".[testing,ml]"
 
     - name: Run FPGA Tests
       run: |

.github/workflows/heterogeneous-ci.yml

@@ -36,7 +36,7 @@ jobs:
         pip install flake8 pytest-xdist coverage
         pip install mpi4py pytest-mpi
         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
 

dace/autodiff/__init__.py

@@ -0,0 +1,57 @@
+"""
+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,96 @@
+"""
+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
+    :returns: 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
+    :returns: 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():
+            # this is our current subgraph, check if it was written before in this subgraph
+            for edge in state.edge_bfs(node, reverse=True):
+                if edge.data.data == array_name:
+                    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