Reverse-Mode Automatic Differentiation and Optimization of GPU Kernels via Enzyme
Reverse-Mode Automatic Differentiation and Optimization of GPU Kernels via Enzyme
复制标题
DOI：
10.1145/3458817.3476165
复制
发表时间：
2021-11
期刊：
影响因子：
--
通讯作者：
William S. Moses;Valentin Churavy;Ludger Paehler;J. Hückelheim;S. Narayanan;Michel Schanen;J. Doerfert
中科院分区：
文献类型：
--
作者：
William S. Moses;Valentin Churavy;Ludger Paehler;J. Hückelheim;S. Narayanan;Michel Schanen;J. Doerfert
计算导数是科学计算和机器学习中许多算法的关键，例如优化、不确定性量化和稳定性分析。Enzyme是一个LLVM编译器插件，它执行反向模式自动微分（AD），从而为包括C/C++、Fortran、Julia和Rust在内的语言的程序生成高性能梯度。在这项工作之前，Enzyme和其他自动微分工具无法生成GPU内核的梯度。我们的论文提出了一系列新技术的组合，使Enzyme成为第一个能够生成GPU内核梯度的全自动反向模式自动微分工具。由于与其他工具不同，Enzyme在通用编译器内执行自动微分，我们能够引入几种新颖的GPU和特定于自动微分的优化。为了展示我们方法的通用性和效率，我们计算了基于GPU的高性能计算应用的梯度，这些应用在NVIDIA和AMD GPU上执行。所有基准测试都在原始程序执行时间的一个数量级内运行。如果没有GPU和特定于自动微分的优化，GPU内核的梯度要么由于资源不足而无法运行，要么具有不可行的开销。最后，我们证明通过增加线程数量或增加每个线程的工作量来增加问题规模，不会对微分的开销产生实质性影响。
Computing derivatives is key to many algorithms in scientific computing and machine learning such as optimization, uncertainty quantification, and stability analysis. Enzyme is a LL VM compiler plugin that performs reverse-mode automatic differentiation (AD) and thus generates high performance gradients of programs in languages including $\mathrm{C}/\mathrm{C}++$, Fortran, Julia, and Rust. Prior to this work, Enzyme and other AD tools were not capable of generating gradi-ents of GPU kernels. Our paper presents a combination of novel techniques that make Enzyme the first fully automatic reverse-mode AD tool to generate gradients of GPU kernels. Since unlike other tools Enzyme performs automatic differentiation within a general-purpose compiler, we are able to introduce several novel GPU and AD-specific optimizations. To show the generality and efficiency of our approach, we compute gradients office GPU-based HPC applications, executed on NVIDIA and AMD GPUs. All bench-marks run within an order of magnitude of the original program's execution time. Without GPU and AD-specific optimizations, gra-dients of GPU kernels either fail to run from a lack of resources or have infeasible overhead. Finally, we demonstrate that increasing the problem size by either increasing the number of threads or increasing the work per thread, does not substantially impact the overhead from differentiation.