ParDiff: Efficiently Parallelizing Reverse-Mode Automatic Differentiation with Direct Indexing

Automatic Differentiation (AD) is a technique that computes the derivatives of numerical programs by systematically applying the chain rule, playing a critical role in domains such as machine learning, simulation, and control systems. However, parallelizing differentiated programs remains a significant challenge due to the conflict between tapes (a data structure for intermediate variable storage) and summations: the differentiation process inherently introduces inter-thread summation patterns, which require prohibitively expensive atomic operations; and traditional tape designs tightly couple data retrieval with the program’s control flow, preventing code restructuring needed to eliminate these costly dependencies. To address these challenges, we present ParDiff, a novel AD system with a direct-indexed tape design, which enables summation-aware loop transformations and various parallel schemes for differentiated programs. This results in a higher degree of parallelization, less synchronization, and reduced inter-thread data movement. We conduct comprehensive experiments on both multi-core CPUs and GPUs. Results show that ParDiff delivers up to 483.21× (geometric mean: 30.88×) speedup over the state-of-the-art fully-AD system, Enzyme. It also achieves a speedup of 2.05× and 2.06× over PyTorch on CPU and GPU, respectively. The source code is publicly available at https://github.com/roastduck/FreeTensor.

• Publication year

  2026

• Venue

  ACM SIGPLAN Symposium on Principles & Practice of Parallel Programming

• Publication date

  2026-01-28

• Fields of study

  Computer Science

• Identifiers
  DOI 10.1145/3774934.3786418
• External record

  Open on Semantic Scholar

• Source metadata

  Semantic Scholar

• No claims are published for this paper.

• No concepts are published for this paper.

• Scalable Automatic Differentiation of Multiple Parallel Paradigms through Compiler Augmentation

  2022cited by this paper
• JAX-FEM: A differentiable GPU-accelerated 3D finite element solver for automatic inverse design and mechanistic data science

  2022cited by this paper
• FreeTensor: a free-form DSL with holistic optimizations for irregular tensor programs

  2022cited by this paper
• DR.JIT

  2022cited by this paper
• Advances in Neural Rendering

  2021cited by this paper
• Reverse-Mode Automatic Differentiation and Optimization of GPU Kernels via Enzyme

  2021cited by this paper
• Getting to the point: index sets and parallelism-preserving autodiff for pointful array programming

  2021cited by this paper
• Source-to-Source Automatic Differentiation of OpenMP Parallel Loops

  2021cited by this paper
• Dynamic Tensor Rematerialization

  2020cited by this paper
• Instead of Rewriting Foreign Code for Machine Learning, Automatically Synthesize Fast Gradients

  2020influential reference
• Towards Polyhedral Automatic Differentiation

  2019influential reference
• Soft Rasterizer: A Differentiable Renderer for Image-Based 3D Reasoning

  2019cited by this paper
• Adjoint computations by algorithmic differentiation of a parallel solver for time-dependent PDEs

  2019cited by this paper
• Reverse-mode algorithmic differentiation of an OpenMP-parallel compressible flow solver

  2019cited by this paper
• Don't Unroll Adjoint: Differentiating SSA-Form Programs

  2018influential reference
• Automatic differentiation in ML: Where we are and where we should be going

  2018cited by this paper
• Differentiable programming for image processing and deep learning in halide

  2018influential reference
• A benchmark of selected algorithmic differentiation tools on some problems in computer vision and machine learning

  2018cited by this paper
• SIMPLE adjoint message passing

  2018cited by this paper
• A Language and Compiler View on Differentiable Programming

  2018cited by this paper
• Compiling machine learning programs via high-level tracing

  2018cited by this paper
• Tangent: Automatic Differentiation Using Source Code Transformation in Python

  2017cited by this paper
• To Push or To Pull: On Reducing Communication and Synchronization in Graph Computations

  2017cited by this paper
• Automatic differentiation in PyTorch

  2017influential reference
• The Pluto+ Algorithm

  2016cited by this paper
• Presburger formulas and polyhedral compilation

  2016cited by this paper
• Clad — Automatic Differentiation Using Clang and LLVM

  2015cited by this paper
• Automatic differentiation in machine learning: a survey

  2015cited by this paper
• Fast Reverse-Mode Automatic Differentiation using Expression Templates in C++

  2014cited by this paper
• The Tapenade automatic differentiation tool: Principles, model, and specification

  2013cited by this paper
• Application of Automatic Differentiation to an Incompressible URANS Solver

  2012cited by this paper
• Automatic Differentiation for GPU-Accelerated 2D/3D Registration

  2008cited by this paper
• Reverse-mode AD in a functional framework: Lambda the ultimate backpropagator

  2008cited by this paper
• Explicit loop scheduling in OpenMP for parallel automatic differentiation

  2002cited by this paper
• Optimizing Compilers for Modern Architectures: A Dependence-based Approach

  2001cited by this paper
• Algorithm 799: revolve: an implementation of checkpointing for the reverse or adjoint mode of computational differentiation

  2000cited by this paper
• Achieving logarithmic growth of temporal and spatial complexity in reverse automatic differentiation

  1992cited by this paper
• Advanced compiler optimizations for supercomputers

  1986cited by this paper

• No citing papers are available for this paper.