DaCe AD: High-Performance Automatic Differentiation

• DaCe AD is a high-performance automatic differentiation engine that integrates symbolic reverse-mode differentiation with ILP-based checkpointing for efficient memory–compute trade-offs.
• It captures unmodified Python, NumPy, PyTorch, or Fortran code into a data-centric SDFG intermediate representation to enable aggressive optimizations and kernel generation.
• The framework delivers significant speedups over JAX by leveraging techniques like loop fusion, vectorization, and parallelization to optimize scientific computing workloads.

DaCe AD is a high-performance automatic differentiation (AD) engine designed to bridge the gap between machine learning and scientific computing workloads. Built atop the DaCe framework and its Stateful DataFlow multiGraph (SDFG) intermediate representation, DaCe AD aims for zero-intrusion: it processes unmodified Python/NumPy, PyTorch, ONNX, or Fortran code, capturing it into an SDFG via DaCeML or @dace.program annotation. The framework employs a symbolic reverse-mode AD algorithm combined with a novel ILP-based checkpointing optimizer, enabling efficient utilization of compute and memory resources. DaCe AD delivers significant speedups over JAX, a leading Python AD framework, establishing itself as a reference tool for differentiating high-performance scientific workloads without imposing code rewrites or language restrictions (Boudaoud et al., 2 Sep 2025).

1. System Architecture and Workflow

DaCe AD operates as an extension to the DaCe data-centric programming framework, leveraging SDFG as its core IR. The workflow consists of the following stages:

• Frontend Capture: User code written in Python/NumPy, or functions decorated with @dace.program, are traced into SDFGs. DaCeML provides ONNX and PyTorch to SDFG conversion. No code modifications—such as enforcing immutability or modifying indexing—are required.
• Data-centric IR: SDFG nodes represent Access (data containers), Tasklet (scalar computations), Map (parallel loops), Library nodes, State containers, Memlets (explicit data movement), and Loop Regions (sequential loops). Edges capture explicit dataflow and memory subset relations.
• Reverse-mode AD Pass: The Critical Computation Subgraph (CCS) is determined via reverse BFS from the outputs, isolating only those nodes whose outputs influence the dependent variable. Each CCS node undergoes symbolic reversal to propagate gradients appropriately, and the reversed subgraph forms the backward-pass SDFG.
• Store/Recompute Optimizer: Given the SDFG pair (forward and backward), an ILP decides which forward intermediates to keep in memory (store) versus recompute during backpropagation, obeying a memory constraint.
• Code Generation and Tuning: Classical DaCe optimization passes—such as inlining, loop fusion, library matching, vectorization, and parallelization—are performed, emitting code for targets such as C++/OpenMP and CUDA.

The end-to-end workflow is: User Code → SDFG via DaCe/DaCeML → CCS Extraction → SDFG Reversal → ILP Checkpointing → Optimizations → Generated Kernel (Boudaoud et al., 2 Sep 2025).

2. Symbolic Differentiation Mechanism

DaCe AD implements symbolic reverse-mode AD at the granularity of SDFG Tasklets and subgraphs:

• For each computation node $z = \text{op}(x, y, \ldots)$, classical rules such as the chain and product rules are applied. For example:

$$\frac{\partial (u \cdot v)}{\partial x} = \frac{\partial u}{\partial x} \cdot v + u \cdot \frac{\partial v}{\partial x}$$

• Tasklets holding mathematical expressions are differentiated symbolically using a SymPy-like approach, generating new Tasklets for the gradient expressions in the backward SDFG, with Memlets updated to track correct dataflow.
• Sequential loops exploit reverse-iteration semantics, while parallel Map constructs are inverted with respect to their data dependencies but preserve the index space (Boudaoud et al., 2 Sep 2025).

This symbolic reverse-mode approach allows granular differentiation across complex scientific code, supporting a range of computational patterns without manual code scribbling.

3. ILP-Based Checkpointing for Memory–Compute Trade-offs

DaCe AD formulates the decision of storing or recomputing each forward-pass intermediate as a binary ILP problem to optimize performance under a given memory constraint. The formulation is as follows:

• Variables: $v_i \in \{0,1\}$, representing whether to store ($v_i = 1$) or recompute ($v_i = 0$) array $A_i$.
• Parameters:
  • $S_i$: memory size of $A_i$
  • $c_i$: recompute cost (e.g., MFLOP)
  • $R_i$: extra memory overhead for recomputation
• Objective:

$$\frac{\partial (u \cdot v)}{\partial x} = \frac{\partial u}{\partial x} \cdot v + u \cdot \frac{\partial v}{\partial x}$$0

• Constraints:

$$\frac{\partial (u \cdot v)}{\partial x} = \frac{\partial u}{\partial x} \cdot v + u \cdot \frac{\partial v}{\partial x}$$1

where $$\frac{\partial (u \cdot v)}{\partial x} = \frac{\partial u}{\partial x} \cdot v + u \cdot \frac{\partial v}{\partial x}$$2 is an affine memory-usage expression per allocation/deallocation event in the schedule.

For typical scientific workloads ($$\frac{\partial (u \cdot v)}{\partial x} = \frac{\partial u}{\partial x} \cdot v + u \cdot \frac{\partial v}{\partial x}$$3), the ILP solves in under 10 ms using commercial (e.g., Gurobi) or open-source (e.g., COIN-OR CBC) solvers. When candidate variable counts are large, the system employs heuristics such as grouping small temporaries or pre-filtering arrays by recompute cost to reduce problem size (Boudaoud et al., 2 Sep 2025).

4. Implementation Aspects

DaCe AD relies on structured SDFG graph representations with explicit annotation of program states, nodes, and dataflow (via Memlets). Each Tasklet includes both its native symbolic AST and the differentiated counterpart. A memory-timeline array tracks allocation and deallocation events, parameterized by the ILP variables.

The main implementation sequence is:

1. Tracing: User code is captured and lowered into SDFG.
2. CCS Extraction: Reverse BFS from output Access nodes (Algorithm 1 in (Boudaoud et al., 2 Sep 2025)).
3. Element Reversal: Library-node-specific reversals, symbolic differentiation for Tasklets, Map inversion, and Loop rewiring.
4. ILP Formulation and Solution: Constructs and solves the store/recompute optimization problem.
5. SDFG Optimizations: Includes fusion, inlining, vectorization, and library matching.
6. Code Generation: Produces optimized C++/OpenMP/CUDA kernels.

Python frontends are supported via the @dace.program decorator for NumPy/Python code and DaCeML interfaces for ONNX/PyTorch, leveraging jit-tracing or graph parsing for initial SDFG construction (Boudaoud et al., 2 Sep 2025).

5. Empirical Performance and Benchmark Results

DaCe AD was evaluated on the NPBench suite (52 kernels, with 38 AD-compatible) on a dual-socket Intel Xeon Gold 6154 (36 cores, AVX-512) system.

• Gradient-computation time, peak memory, and solver overhead were measured, with JAX JIT as the baseline.
• Speedups:
  • Vectorized kernels (12): 1.43× geomean over JAX (8/12 faster)
  • Non-vectorized kernels (26): 134× geomean (20/26 faster)
  • Overall (38 kernels): 4.1× geomean; 92× arithmetic mean speedup over JAX.
• Case study (Seidel2d kernel):
  • JAX JIT: 47 minutes on CPU (dynamic slices, immutability overhead)
  • DaCe AD: 1 second (via in-place updates, no bound checks, fused loops)
  • Result: 2,724× CPU speedup
• GPU results (NVIDIA V100, 9 non-vectorized kernels): arithmetic mean speedup >10×; for Seidel2d, 275× (Boudaoud et al., 2 Sep 2025).
• ILP Checkpointing microbenchmark: On a toy loop with three intermediates, the ILP (8 configurations) selects $$\frac{\partial (u \cdot v)}{\partial x} = \frac{\partial u}{\partial x} \cdot v + u \cdot \frac{\partial v}{\partial x}$$4 in 6.4 ms, achieving a 30% runtime reduction under a 500 MiB memory cap.

6. Current Limitations and Potential Extensions

Several expressiveness and feature limitations remain:

• No support for unstructured while-loops or loops with break/continue (requires iteration tracking).
• No handling of dynamic Python lists, recursion, or pointer-based indirection.
• No support for Tasklets with complex number arithmetic (applying a complex-step method is orthogonal).

Potential future directions include:

• Support for while-loops via runtime traced iteration counts and compact backward-pass generation.
• Enhancements in Memlet analysis to support indirect and sparse data access.
• Support for higher-order derivatives, such as vector-Jacobian products.
• Advanced MILP heuristics for checkpointing (e.g., variable grouping, warm-starting).
• Extension to distributed SDFGs for multi-node automatic differentiation (Boudaoud et al., 2 Sep 2025).

DaCe AD establishes a unified, code-intrusion-free platform for automatic differentiation in high-performance computing and machine learning, bringing together symbolic reverse-mode AD, ILP-based checkpointing and aggressive SDFG optimizations to significantly exceed the performance of contemporary AD systems in scientific code contexts.