Operator Fusion Scheduling Optimization for TVM Deep Learning Compilers
Guangfei Zheng, Jianan Li, Wei Gao, Lin Han, Yingying Li, Jinlong Xu · 2023
This paper addresses the problem that the fusion of memory-intensive operators in the automatic scheduling module of TVM deep learning compiler is not fully considered, which leads to the inefficient operation of computational tasks defined by tensor expressions. In this paper, we design a scheduling rule for fusing memory-intensive arithmetic structures by analyzing the arithmetic characteristics, combining the characteristics of different hardware architectures, and assigning corresponding scheduling policies to arithmetic using scheduling primitives. The differences of the scheduling tree before and after the new scheduling rule are also analyzed to verify the effectiveness of the scheduling rule. The experimental results show that the new scheduling rule shows significant optimization effects on several test sets of different sizes, and the average speedup ratios reach 1.02x, 1.21x and 1.23x on CPU, GPU and DCU platforms, respectively; the highest speedup ratio is 1.95x on DCU platform.