张晨，清华大学计算机系高性能所博士生，导师为翟季冬老师，主要研究方向为面向人工智能和量子计算的高性能异构计算系统。在OSDI、SC、ATC、ICS会议上发表一作论文，并获得 ICS21 最佳学生论文。曾获得 SC19, SC20, ISC21 国际超级计算机竞赛冠军。获清华大学本科生特等奖学金、国家奖学金、北京市优秀毕业生、北京市优秀毕业设计等荣誉。2024 年 7 月，清华大学计算机系 PACMAN 实验室发布开源深度学习编译器 MagPy，可一键编译用户使用 Python 编写的深度学习程序，实现模型的自动加速。
尽管目前存在大量高性能的深度学习编译器，但是这些编译器均以计算图作为输入，需要由用户将编写的 Python 程序手动转化为计算图。为了避免这种不便性，该团队设计了 MagPy，直接面向用户编写的 Python+PyTorch 程序，自动将其转化为适用于深度学习编译器的计算图表示，从而充分发挥深度学习编译器的优化能力，避免用户使用复杂 Python 语法带来的性能下降，为用户带来易用性和效率的双丰收。
论文地址：https://www.usenix.org/system/files/atc24-zhang-chen.pdf
项目地址：https://github.com/heheda12345/MagPy
研究背景
近年来，深度学习在生物科学、天气预报和推荐系统等多个领域展示了其强大能力。为了简化编程过程，用户倾向于使用 Python 编写深度学习模型，并在需要进行张量操作时调用如 PyTorch 等的张量库。此时，用户程序会在调用张量库时立即执行张量操作，如此不加优化地直接执行程序性能较差。另一方面，为了提升深度学习模型的运行速度，深度学习编译器倾向于使用以算子图的格式表示的深度学习模型作为输入，在计算图上进行图级优化，如图替换和算子融合。当可以获取到模型的计算图时，代表性的深度学习编译器 TorchInductor 和 XLA 可以在 PyTorch 的基础上平均加速模型 1.47 倍和 1.40 倍。
近年来，深度学习在生物科学、天气预报和推荐系统等多个领域展示了其强大能力。为了简化编程过程，用户倾向于使用 Python 编写深度学习模型，并在需要进行张量操作时调用如 PyTorch 等的张量库。此时，用户程序会在调用张量库时立即执行张量操作，如此不加优化地直接执行程序性能较差。另一方面，为了提升深度学习模型的运行速度，深度学习编译器倾向于使用以算子图的格式表示的深度学习模型作为输入，在计算图上进行图级优化，如图替换和算子融合。当可以获取到模型的计算图时，代表性的深度学习编译器 TorchInductor 和 XLA 可以在 PyTorch 的基础上平均加速模型 1.47 倍和 1.40 倍。
由于 Python 程序具有极强的动态性，加之用户程序存在行为的不确定性，现有的计算图提取技术在处理较复杂的用户程序时无法取得最优的性能，如图 1 中的 TorchDynamo-Inductor（使用 TorchDynamo 提取计算图，使用 TorchInductor 编译）、 LazyTensor-XLA（使用 LazyTensor 追踪计算图，使用 XLA 编译）所示。
MagPy分析
MagPy 的核心思想是分析 Python 解释器中的执行状态信息，从而让编译器能够更好的理解用户程序。Python 解释器能够准确支持所有 Python 特性，并在运行时保留了高层次的执行状态信息，如各个变量的类型和值等等。通过有效利用解释器提供的信息，能够更全面地了解程序的行为，从而更好地提取程序计算图。
其次，只有外部值能影响程序行为。利用这一特性，可以更简易地检测出会导致计算图发生变化的因素。这里的 “程序行为” 包括计算图的结构和所有程序副作用（side effect）。只要程序从外部读取的所有值（如输入参数和全局变量）保持不变，且调用的函数的输出结果不具有随机性，程序行为就不会发生变化。因此，MagPy 只需验证所有从外部读取的值都不变，即可保证计算图结构不变。例如，尽管图 2 中的程序使用了许多复杂的 Python 特性，但只要所有从外部读取的值（如 x、dims、self.scale 和 self.dim，标记为粗体）与之前运行一致，计算图就保持不变。MagPy 会首先运行一个 “守卫函数” 对于这些值是否发生变化进行检查（Guards），当检查通过时，MagPy 将会运行一个 “模拟函数”(mock code)，用以调用经过深度学习编译器编译的计算图及模拟程序的所有副作用（如示例中的对 x 进行赋值）。
第三，守卫函数和模拟函数都可以通过分析程序执行状态来确定。守卫函数的作用是验证新一次执行的输入状态是否与之前运行匹配，模拟函数的目的是重现之前运行的最终状态。这两个部分仅基于运行时状态，而不是用户程序的逻辑。Python 解释器在解释执行程序的过程中，保留了所有需要的执行状态信息，因此不再需要具体分析 Python 复杂而动态的执行逻辑。守卫函数和模拟函数需要关注的变量包括显式读取或写入外部的值（如 self）以及被它们引用的值（如 self.dim）。因此，MagPy 设计了引用关系图来记录和分析程序行为。
基于上述观察，MagPy 提出了引用关系图（Reference Graph，简写为 RefGraph）来记录程序执行期间的程序状态。MagPy 定义了执行状态接口，用于在程序执行期间收集运行时信息，并使用基于标注的图更新规则来维护 RefGraph。MagPy 还提出了在 RefGraph 上进行遍历生成守卫函数和模拟函数的算法。具体细节可以阅读论文。
实验数据
MagPy 具有极高的 Python 语言特性覆盖率，其在对 ParityBench 中 1191 个静态的真实用户程序进行测试时，成功将 93.40% 的程序转化为完整的操作符图，大幅高于现有工作 TorchScript（35%）和 TorchDynamo（77.2%）
