Reverse Automatic Differentiation for Accelerate

Abstract

Purely functional array programming languages are powerful tools for expressing data-parallel computation in a way that is easily optimised to have good performance on parallel hardware like GPUs. Probabilistic inference and optimisation are fields that need this computational power, but also require reverse-mode automatic differentiation (AD) of code written in the programming language. Accelerate is a purely functional array programming language with an optimising compiler for multicore CPU and GPU backends that uses a strongly-typed abstract syntax tree to ensure that no compiler pass can perform type-incorrect program transformations. However, due to the second-order nature of the Accelerate language, existing algorithms for reverse-mode AD do not directly apply to Accelerate.

We define a new, pragmatic source-code transformation algorithm for reverse-mode AD on a significant subset of Accelerate; the supported subset excludes loops. We present an implementation of our algorithm in the Accelerate compiler and show that it performs reasonably competitively with other implementations of reverse-mode AD on benchmarks; as a compiler pass in the Accelerate compiler, our implementation cannot perform type-incorrect program transformations. Our reverse-mode AD algorithm can possibly be generalised to other second-order, purely functional array programming languages.