EINNET: Optimizing Tensor Programs with Derivation-Based Transformations


Liyan Zheng, Haojie Wang, Jidong Zhai, Muyan Hu, Zixuan Ma, Tuowei Wang, and Shuhong Huang, Tsinghua University; Xupeng Miao, Carnegie Mellon University; Shizhi Tang and Kezhao Huang, Tsinghua University; Zhihao Jia, Carnegie Mellon University


Boosting the execution performance of deep neural networks (DNNs) is critical due to their wide adoption in real-world applications. However, existing approaches to optimizing the tensor computation of DNNs only consider transformations representable by a fixed set of predefined tensor operators, resulting in a highly restricted optimization space. To address this issue, we propose EinNet, a derivation-based tensor program optimizer. EinNet optimizes tensor programs by leveraging transformations between general tensor algebra expressions and automatically creating new operators desired by transformations, enabling a significantly larger search space that includes those supported by prior works as special cases. Evaluation on seven DNNs shows that EinNet outperforms existing tensor program optimizers by up to 2.72× (1.52× on average) on NVIDIA A100 and up to 2.68× (1.55× on average) on NVIDIA V100. EinNet is publicly available at https://github.com/InfiniTensor/InfiniTensor.

OSDI '23 Open Access Sponsored by
King Abdullah University of Science and Technology (KAUST)

Open Access Media

USENIX is committed to Open Access to the research presented at our events. Papers and proceedings are freely available to everyone once the event begins. Any video, audio, and/or slides that are posted after the event are also free and open to everyone. Support USENIX and our commitment to Open Access.

@inproceedings {288645,
author = {Liyan Zheng and Haojie Wang and Jidong Zhai and Muyan Hu and Zixuan Ma and Tuowei Wang and Shuhong Huang and Xupeng Miao and Shizhi Tang and Kezhao Huang and Zhihao Jia},
title = {{EINNET}: Optimizing Tensor Programs with {Derivation-Based} Transformations},
booktitle = {17th USENIX Symposium on Operating Systems Design and Implementation (OSDI 23)},
year = {2023},
isbn = {978-1-939133-34-2},
address = {Boston, MA},
pages = {739--755},
url = {https://www.usenix.org/conference/osdi23/presentation/zheng},
publisher = {USENIX Association},
month = jul

Presentation Video