AD for an Array Language with Nested Parallelism

We present a technique for applying (forward and) reversemode automatic differentiation (AD) on a non-recursive secondorder functional array language that supports nested parallelism and is primarily aimed at efficient GPU execution.

The key idea is to eliminate the need for a "tape" by relying on redundant execution to bring into each new scope all program variables that may be needed by the differentiated code. Efficient execution is enabled by the observation that perfectly-nested scopes do not introduce re-execution, and such perfect nests are produced by known compiler transformations, e.g., flattening. Our technique differentiates loops and bulk-parallel operators-such as map, reduce, histogram, scan, scatter-by specific rewrite rules, and aggressively optimizes the resulting nested-parallel code. We report an experimental evaluation that compares with established AD solutions and demonstrates competitive performance on nine common benchmarks from recent applied AD literature.