A Graph-Based Method for AI Chip Operator Optimization Using Deep Learning

AI operators refer to reusable code programs encapsulated in AI chip frameworks that implement specific functions. To achieve high performance, it is necessary to combine hardware characteristics efficiently when programming. However, for operators with multiple computational steps, it is challenging to develop excellent scheduling strategies. To address this issue, this paper proposes a graph-based method for AI chip operator optimization. Firstly, establish bidirectional transformation relationship between operators and corresponding computation graphs. Then, use deepwalk and word2vec to convert operators’ computation graphs into embedding representations, and optimize corresponding operators by annotating the nodes using graph neural network. Also, simple operator fusion can be achieved by fusing graphs of multiple operators and optimizing the fused operator. Through the creation of an operator dataset and related experiments within the Cambricon community framework, this method demonstrates superior optimization and fusion of element-wise operators compared to other simple tuning methods.