一句话：提出 PolyJuice，用等价饱和生成语义等价的多样计算图，通过差异测试发现张量编译器的误编译缺陷。
张量编译器是深度学习应用跨硬件部署的关键，但其优化逻辑复杂，正确性难以保证。本文提出自动检测工具 PolyJuice，核心思想是构造语义等价的 computation graph 来验证张量编译器的正确性。主要挑战在于构造能高效探索编译期多样优化逻辑的等价图。作者从两方面解决：一是提出算术与结构等价的 rewrite rules 改写张量程序的数据流；二是设计高效的基于 equality saturation 的重写框架，为输入图寻找最简化与最复杂的语义等价计算图。由于这些图数据流不同，编译时会经历不同优化过程，从而暴露误编译。作者将工具应用于 PyTorch Inductor、OnnxRuntime、TVM、TensorRT、XLA 五个工业编译器以及 EinNet、Hidet 两个学术编译器，共检测出 84 个非崩溃误编译缺陷，其中 49 个被确认、20 个已修复。
Tensor compilers are essential for deploying deep learning applications across various hardware platforms. While powerful, they are inherently complex and present significant challenges in ensuring correctness. This paper introduces PolyJuice, an automatic detection tool for identifying mis-compilation bugs in tensor compilers. Its basic idea is to construct semantically-equivalent computation graphs to validate the correctness of tensor compilers. The main challenge is to construct equivalent graphs capable of efficiently exploring the diverse optimization logic during compilation. We approach it from two dimensions. First, we propose arithmetic and structural equivalent rewrite rules to modify the dataflow of a tensor program. Second, we design an efficient equality saturation based rewriting framework to identify the most simplified and the most complex equivalent computation graphs for an input graph. After that, the outcome computation graphs have different dataflow and will likely experience different optimization processes during compilation. We applied it to five well-tested industrial tensor compilers, namely PyTorch Inductor, OnnxRuntime, TVM, TensorRT, and XLA, as well as two well-maintained academic tensor compilers, EinNet and Hidet. In total, PolyJuice detected 84 non-crash mis-compilation bugs, out of which 49 were confirmed with 20 fixed.
@article{ZhouQGZ0024,
title = {PolyJuice: Detecting Mis-compilation Bugs in Tensor Compilers with Equality Saturation Based Rewriting},
author = {Chijin Zhou and Bingzhou Qian and Gwihwan Go and Quan Zhang and Shanshan Li and Yu Jiang},
journal = {Proceedings of the ACM on Programming Languages},
volume = {8},
number = {OOPSLA2},
year = {2024},
doi = {10.1145/3689757}
}