SC21 Proceedings

The International Conference for High Performance Computing, Networking, Storage, and Analysis

Reverse-Mode Automatic Differentiation and Optimization of GPU Kernels via Enzyme

Authors: William S. Moses and Valentin Churavy (Massachusetts Institute of Technology (MIT)); Ludger Paehler (Technical University Munich); and Jan Hückelheim, Sri Hari Krishna Narayanan, Michel Schanen, and Johannes Doerfert (Argonne National Laboratory (ANL))

Abstract: Derivatives are key to algorithms in scientific computing and machine learning such as optimization, uncertainty quantification, and stability analysis. Enzyme is a LLVM compiler plugin for reverse-mode automatic differentiation (AD) and thus generates fast gradients of programs in a variety of languages, including C/C++, Fortran, Julia, and Rust. Our paper presents a combination of novel techniques that make Enzyme the first automatic reverse-mode AD tool to generate gradients of GPU kernels. As Enzyme differentiates within a general-purpose compiler, we are able to introduce novel GPU and AD-specific optimizations. We differentiate five GPU-based HPC applications, executed on NVIDIA and AMD GPUs. All benchmarks run within an order of magnitude of the original program's runtime. Without GPU and AD-specific optimizations, gradients of GPU kernels either fail to run from a lack of resources or have infeasible overhead. We show that increasing the problem size does not substantially impact the overhead from differentiation.

Back to Technical Papers Archive Listing