[论文解读] Instead of Rewriting Foreign Code for Machine Learning, Automatically Synthesize Fast Gradients

William S. Moses, Valentin Churavy|arXiv (Cornell University)|Oct 4, 2020

Parallel Computing and Optimization Techniques参考文献 49被引用 10

一句话总结

该论文介绍了 Enzyme，这是一个针对 LLVM 的高性能自动微分（AD）编译器插件，可直接从优化后的 LLVM 中间表示（IR）合成快速梯度，从而在无需重写外部代码的情况下实现高效、语言无关的 AD。通过在编译器优化之后执行 AD，Enzyme 在未优化 IR 上相比传统 AD 实现了 4.5 倍的几何平均速度提升，在机器学习基准测试中优于当前最先进的工具。

ABSTRACT

Applying differentiable programming techniques and machine learning algorithms to foreign programs requires developers to either rewrite their code in a machine learning framework, or otherwise provide derivatives of the foreign code. This paper presents Enzyme, a high-performance automatic differentiation (AD) compiler plugin for the LLVM compiler framework capable of synthesizing gradients of statically analyzable programs expressed in the LLVM intermediate representation (IR). Enzyme synthesizes gradients for programs written in any language whose compiler targets LLVM IR including C, C++, Fortran, Julia, Rust, Swift, MLIR, etc., thereby providing native AD capabilities in these languages. Unlike traditional source-to-source and operator-overloading tools, Enzyme performs AD on optimized IR. On a machine-learning focused benchmark suite including Microsoft's ADBench, AD on optimized IR achieves a geometric mean speedup of 4.5x over AD on IR before optimization allowing Enzyme to achieve state-of-the-art performance. Packaging Enzyme for PyTorch and TensorFlow provides convenient access to gradients of foreign code with state-of-the art performance, enabling foreign code to be directly incorporated into existing machine learning workflows.

研究动机与目标

• 通过支持原生、高性能的自动微分，消除在机器学习框架中重写外部代码的需求。
• 解决在编译器优化之前应用 AD 所导致的性能瓶颈，该问题可能导致梯度代码效率低下。
• 实现在无源码访问的情况下对预编译或第三方库进行跨语言、跨框架的微分。
• 证明在低级、优化后的 IR 上实现高效 AD 是可行的，挑战了 AD 仅在高级语言中有效的假设。
• 为科学计算和仿真代码提供无缝集成到现代机器学习工作流中的路径。

提出的方法

• Enzyme 作为 LLVM 编译器框架中的插件运行，直接在优化后的 LLVM IR 上执行反向模式自动微分。
• 它在应用 AD 之前利用 LLVM 的完整优化套件（包括循环不变量代码移动等）来保留性能提升。
• 该系统通过分析优化后 IR 的数据流和控制流来合成梯度，生成高效的伴随代码。
• Enzyme 使用递归栈结构进行梯度计算，相比基于栈的方法显著降低了内存开销。
• 它通过统一的基于 IR 的微分流水线，支持多种目标 LLVM IR 的语言，包括 C、C++、Fortran、Julia、Rust 和 Swift。
• 该框架与 PyTorch 和 TensorFlow 集成，支持在现有机器学习工作流中直接对第三方代码进行微分。

实验结果

研究问题

• RQ1 是否可以在优化后的 LLVM IR 上高效执行自动微分，而非在未优化的源码或 IR 上？
• RQ2 在优化之后执行 AD 是否相比在未优化代码上使用传统 AD 显著提升梯度计算速度？
• RQ3 基于编译器插件的 AD 系统是否能在多样化的机器学习和科学计算工作负载中实现最先进性能？
• RQ4 Enzyme 在多大程度上能够实现对第三方代码（如物理模拟器或游戏引擎）的无缝集成，而无需修改源码？
• RQ5 是否可行构建一个高性能、跨语言的 AD 系统，使其在 IR 层次运行并支持 GPU 代码生成等高级功能？

主要发现

• Enzyme 在涵盖 Microsoft 的 ADBench 等多样化的基准测试套件中，相比未优化 IR 上的 AD（即 Reference 流水线），实现了 4.5 倍的几何平均速度提升。
• 性能提升主要源于后优化 AD，使得诸如循环不变量代码移动等优化能够消除梯度中的冗余计算。
• Enzyme 在 ADBench 的所有基准测试中均优于所评估的两个最快 C/C++ AD 工具，展现出最先进性能。
• Enzyme 所采用的递归栈结构相比基于栈的结构显著降低了内存开销，从而在 FFT 和 Bruss 基准测试中实现更优性能。
• Enzyme 支持对预编译库和头文件仅有的代码进行直接微分，而这是源到源 AD 工具无法实现的。
• 该框架成功对复杂的科学计算工作负载（如 Euler 和 RK4 求解器、LSTM 和 GMM）进行了微分，性能与专用工具相当或更优。