一句话:将 AD 活性分析形式化为抽象解释并证明其可靠性,在 MLIR 中实现含函数摘要的模块化版本,获得显著加速。
自动微分生成的导数代码常需激进优化以维持性能,其中核心优化是识别不影响目标偏导的非活跃操作。现有工具虽提供多种输入语言的活性分析,但通常只有非形式化的正确性保证。本文将 AD 活性分析形式定义为抽象解释,证明其可靠性,并在 MLIR 编译基础设施中实现。为应对 MLIR 的通用性,本文首次形式化了 MLIR 内部表示中适合 AD 的一个子集。此外,文章提出通过函数摘要对全程序活性分析进行可靠的过程内近似,并给出从函数定义自动推导这些摘要的机制。在面向微分的基准套件上评估,相比不使用活性分析的基线,生成程序在 CPU 上取得 1.24 倍几何平均加速,在 GPU 上取得 1.7 倍几何平均加速;函数摘要辅助的过程内分析可证明全程序分析所证明的全部非活跃指令,达到 100%。
Computing derivatives is paramount for multiple domains ranging from training neural networks to precise climate simulations. While derivatives can be generated by AD (Automatic Differentiation) tools, they often require aggressive optimization to avoid compromising program performance. One of the central optimizations consists of identifying inactive operations that do not contribute to the partial derivatives of interest. Multiple tools provide activity analyses for a variety of input languages, though often with only informal correctness guarantees. This paper formally defines activity analysis for AD as an abstract interpretation, proves its soundness, and implements it within the MLIR compiler infrastructure. To account for MLIR’s genericity, a subset of MLIR’s internal representation amenable to AD is formalized for the first time. Furthermore, the paper proposes a sound intraprocedural approximation of the whole-program activity analysis via function summaries along with a mechanism to automatically derive these summaries from function definitions. The implementation is evaluated on a differentiation-specific benchmark suite. It achieves a 1.24X geometric mean speedup on CPU and a 1.7X geometric mean speedup on GPU in the runtime of generated programs, when compared to a baseline that does not use activity analysis. The evaluation also demonstrates that the intraprocedural analysis with function summaries proves inactive 100% of instructions proven inactive by the whole-program analysis.
@article{paperbot2969,
title = {Sound and Modular Activity Analysis for Automatic Differentiation in MLIR},
author = {Mai Jacob Peng and William S. Moses and Oleksandr Zinenko and Christophe Dubach},
journal = {Proceedings of the ACM on Programming Languages},
volume = {9},
number = {OOPSLA2},
year = {2025},
doi = {10.1145/3763125}
}