Boosting the runtime performance of deep neural networks (DNNs) is critical due to their wide adoption in real-world tasks. Existing approaches to optimizing the tensor algebra expression of a DNN only consider expressions representable by a fixed set of predefined operators, missing possible optimization opportunities between general expressions. We propose OLLIE, the first derivation-based tensor program optimizer. OLLIE optimizes tensor programs by leveraging transformations between general tensor algebra expressions, enabling a significantly larger expression search space that includes those supported by prior work as special cases. OLLIE uses a hybrid derivation-based optimizer that effectively combines explorative and guided derivations to quickly discover highly optimized expressions. Evaluation on seven DNNs shows that OLLIE can outperform existing optimizers by up to 2.73× (1.46× on average) on an A100 GPU and up to 2.68× (1.51×) on a V100 GPU, respectively.
@article{arxiv.2208.02025,
title = {OLLIE: Derivation-based Tensor Program Optimizer},
author = {Liyan Zheng and Haojie Wang and Jidong Zhai and Muyan Hu and Zixuan Ma and Tuowei Wang and Shizhi Tang and Lei Xie and Kezhao Huang and Zhihao Jia},
journal= {arXiv preprint arXiv:2208.02025},
year = {2022}
}