Reverse-Mode Automatic Differentiation and Optimization of GPU Kernels via Enzyme
Best Student Paper Finalist
Best Reproducibility Advancement Finalist
TimeWednesday, 17 November 20213:30pm - 4pm CST
DescriptionDerivatives 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.