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

Abstract

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.