用于 MLIR 中自动微分的声学与模块化活动分析

《Proceedings of the ACM on Programming Languages》：Sound and Modular Activity Analysis for Automatic Differentiation in MLIR

本文提出一种基于抽象解释的自动微分活跃性分析优化方法，在MLIR编译器中实现并验证。通过形式化定义全局活跃性分析，结合函数总结的内联近似方法，在保持分析精度的同时提升性能，实验显示CPU和GPU场景下分别获得1.24和1.7倍加速比。

　　

摘要

计算导数在多个领域中都至关重要，从训练神经网络到精确的气候模拟都离不开这一过程。虽然可以通过自动微分（AD）工具生成导数，但通常需要采取积极的优化措施以避免影响程序性能。其中一项核心优化工作是识别那些对目标偏导数没有贡献的“非活跃”操作。

有多种工具可以为各种输入语言提供“活动性”分析功能，但这些工具通常仅提供非正式的正确性保证。本文正式定义了自动微分中的活动性分析方法，并证明了其正确性，同时将其实现了在MLIR编译器框架中。为了体现MLIR的通用性，本文首次对适用于自动微分的MLIR内部表示的一个子集进行了形式化描述。此外，本文还提出了一种基于函数摘要的全程序活动性分析的近似方法，并设计了一种能够从函数定义自动生成这些摘要的机制。

该实现方法在一个针对微分运算的基准测试套件上进行了评估。与不使用活动性分析的基线方案相比，使用该方法生成的程序在CPU上的运行时间提升了1.24倍，在GPU上的运行时间提升了1.7倍。评估结果还表明，基于函数摘要的程序内部分析能够100%地识别出全程序分析已经确定为非活跃的指令。