Graph-based Clustering Cost Model for Performance Optimization in Deep Learning Compilers

Deep learning (DL) compilers enable automated deployment of deep neural networks (DNNs) across diverse hardware platforms by leveraging cost models for efficient tensor program optimization. The effectiveness of these systems depends on the prediction accuracy and cross-model generalization capability of offlinetrained cost models. Existing approaches mainly focus on feature engineering and cost model architecture design, while largely overlooking structural similarity at the computational subgraph level, which limits further improvements in prediction performance. To address this limitation, we propose an unsupervised subgraph clustering framework based on graph autoencoders to enhance offline-trained cost models. Specifically, DL models are partitioned into computational subgraphs, and node features are extracted. A graph autoencoder learns high-dimensional structural embeddings of these subgraphs, followed by K-means clustering to group structurally similar patterns. Dedicated cost models are trained for different subgraph clusters to predict tensor program performance. This fine-grained, cluster-aware training strategy improves both prediction accuracy and generalization capability, thereby enhancing auto-tuning efficiency. Experimental results across multiple DL models demonstrate significant improvements in search efficiency and latency optimization on both CPU and GPU platforms. Compared with Tenset, our method achieves average search speedups of 9.46 & times; and 5.15 & times;, while achieving inference latency improvements of 1.18 & times; and 1.07 & times;.