一句话：将等价饱和用于真实 ML 计算图，结合全局代价模型优化 XLA，在 CPU/GPU 上平均提速 3.45%。
ML 编译器依赖图级变换提升运行时性能，但局部重写可能改变后续数据布局、并行化、细粒度调度和内存管理等变换的合法性与收益，使编译器长期依赖易过时的手工启发式。本文提出用等价饱和替代这些局部启发式，并以考虑下游变换的全局性能模型评估变换收益，从而发现更有效的图优化组合。作者针对真实 ML 计算图和现代硬件，研究融合与布局优化中的代价建模方法，并解决宽范围代数优化带来的可扩展性问题。他们在 XLA 中设计并实现了 C++/Rust 的等价饱和 pass。在多个 CPU 和 GPU 平台基准上，该方法相对 XLA 优化流程平均提速 3.45%，NasRNN 在 CPU 上最高提速 56.26%。
Machine learning (ML) compilers rely on graph-level transformations to enhance the runtime performance of ML models. However, performing local transformations on individual operations can create effects far beyond the location of the rewrite. In particular, a local rewrite can change the profitability or legality of hard-to-predict downstream transformations, particularly regarding data layout, parallelization, fine-grained scheduling, and memory management. As a result, program transformations are often driven by manually-tuned compiler heuristics, which are quickly rendered obsolete by new hardware and model architectures. Instead of hand-written local heuristics, we propose the use of equality saturation. We replace such heuristics with a more robust global performance model, which accounts for downstream transformations. Equality saturation addresses the challenge of local optimizations inadvertently constraining or negating the benefits of subsequent transformations, thereby providing a solution that is inherently adaptable to newer workloads. While this approach still requires a global performance model to evaluate the profitability of transformations, it holds significant promise for increased automation and adaptability. This paper addresses challenges in applying equality saturation on real-world ML compute graphs and state-of-the-art hardware. By doing so, we present an improved method for discovering effective compositions of graph optimizations. We study different cost modeling approaches to deal with fusion and layout optimization, and tackle scalability issues that arise from considering a very wide range of algebraic optimizations. We design an equality saturation pass for the XLA compiler, with an implementation in C++ and Rust. We demonstrate an average speedup of 3.45% over XLA’s optimization flow across our benchmark suite on various CPU and GPU platforms, with a maximum speedup of 56.26% for NasRNN on CPU.
@article{paperbot2997,
title = {Mind the Abstraction Gap: Bringing Equality Saturation to Real-World ML Compilers},
author = {Arya Vohra and Leo Seojun Lee and Jakub Bachurski and Oleksandr Zinenko and Phitchaya Mangpo Phothilimthana and Albert Cohen and William S. Moses},
journal = {Proceedings of the ACM on Programming Languages},
volume = {9},
number = {OOPSLA2},
year = {2025},
doi = {10.1145/3763062}
}