Cross-Platform Deep Learning Compilation with Model-Specific Cost Model Pre-training
摘要
Deep learning compilers have emerged as a paradigm to increase the efficiency of deployed neural network models. Their key component is a search-based method with a learned cost model, which enables the generation of high-performance tensor programs. Large-scale tensor program performance datasets can be used to pre-train a learned cost model, thereby speeding up search time. However, a pre-trained cost model is not as effective as the default one for cross-platform tensor program performance prediction. It remains a challenge to pre-train a cost model on other platforms using existing datasets. In this paper, we propose MSCPT, a cost model pre-training approach across platforms based on existing static datasets to reduce the search time and inference latency of neural network models. Compared to other approaches, MSCPT can accelerate the search time by an average of 3.77x and 2.15x on CPU and GPU platforms, respectively, while achieving lower inference latency.