关键词: 算符融合, 图形重写, 张量虚拟机, 神经网络, 融合探索

Abstract: In addressing issues encountered in the process of operator fusion for compute-intensive neural networks using the Tensor Virtual Machine (TVM), such as excessive access counts and low memory resource utilization during layer-wise exploration of computational graphs, an optimization method for TVM operator fusion based on graph rewriting and fusion exploration is proposed. Initially, an analysis was conducted on the mapping types of operators. Subsequently, the computational graph was rewritten based on operation laws to simplify its structure, reducing the generation of intermediate results, lowering memory resource consumption, and enhancing fusion efficiency. Following that, a fusion exploration algorithm was employed to identify operators with lower fusion costs for prioritized fusion, thereby mitigating data redundancy and register overflow. Finally, neural network operator fusion was implemented on the CPU, and the fusion acceleration performance was tested. Experimental results indicated that the proposed method effectively reduced the number of computational graph layers and operators, decreased access frequency, and reduced data transfer volume. In comparison with the TVM operator fusion method, there was an average reduction of 18% in computational graph layers during the fusion process, and the fusion speed was increased by an average of 23%, confirming the effectiveness of the method in optimizing the computational graph fusion process.

Key words: operator fusion, graphic rewriting, tensor virtual machine, neural network, fusion exploration