SA‐Cost: Structure‐Aware Cost Model With Attention for Tensor Program Tuning

ABSTRACT Automatic scheduling optimization is a critical technology for improving the performance of deep learning compilers (DLCs). As neural network models grow in scale and hardware platforms diversify, the automatic generation of high‐performance tensor programs has become a focal point of research. A robust cost model is essential for automatic scheduling, as it enables the selection of the most optimal scheduling scheme from a multitude of possibilities. This paper introduces SA‐Cost, a novel cost model for DLCs. Unlike conventional approaches that treat tensor programs as a whole during feature extraction, our method reformulates feature extraction as an analysis of the inter‐relationships and semantic correlations among scheduling primitives. By capturing the semantic correlations among scheduling decisions via an attention mechanism, we construct a more precise evaluation model. The model performs a comprehensive assessment across four dimensions: decision space distribution, primitive significance, hardware awareness, and resource utilization efficiency. Furthermore, we incorporate an attention mechanism to identify key scheduling decisions and their contributions, improving the cost model's prediction accuracy. Experimental results demonstrate that on the GPU platform, SA‐Cost achieves an average speedup of over Ansor and over AMOS.