EINNET: Optimizing Tensor Programs with Derivation-Based Transformations.

Liyan Zheng,Haojie Wang,Jidong Zhai,Muyan Hu,Zixuan Ma,Tuowei Wang,Shuhong Huang,Xupeng Miao,Shizhi Tang,Kezhao Huang,Zhihao Jia
2023-01-01
Abstract: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.72x (1.52x on average) on NVIDIA A100 and up to 2.68x (1.55x on average) on NVIDIA V100. EINNET is publicly available at https://github.com/InfiniTensor/InfiniTensor.
What problem does this paper attempt to address?