GAS: A scheduling primitive dependency analysis-based cost model for tensor program optimization

Read full text
Authors: Yonghua Hu, Anxing Xie, Yaohua Wang, Zhe Li, Zenghua Cheng, Junyang Tang
Cited by: 1
References: 22

Abstract

Automatically generating high-performance tensor programs has become a promising approach for deploying deep neural networks. A key challenge lies in designing an effective cost model to navigate the vast scheduling search space. Existing approaches typically fall into two categories, each with limitations: offline learning cost models rely on large pre-collected datasets, which may be incomplete or device-specific, and online learning cost models depend on handcrafted features, requiring substantial manual effort and expertise. We propose GAS, a lightweight framework for generating tensor programs for deep learning applications. GAS reformulates feature extraction as a sequence-dependent analysis of scheduling primitives. Our cost model integrates three key factors to uncover performance-critical insights within scheduling sequences: (1) decision factors allocation, quantifying entropy and skewness of scheduling primitive factors to capture their dominance; (2) primitive contribution weights, measuring the relative impact of primitives on overall performance; and (3) structural semantic alignment, capturing correlations between scheduling primitive factors and hardware parallelism mechanisms. This approach reduces the complexity of handcrafted feature engineering and extensive pre-training datasets, significantly improving both efficiency and scalability. Experimental results on NVIDIA GPUs demonstrate that GAS achieves average speedups of 3.79 × over AMOS and 2.22 × over Ansor, while also consistently outperforming other state-of-the-art tensor compilers.

Fields of Study & Subjects

Tensor decomposition and applications Parallel Computing and Optimization Techniques Advanced Neural Network Applications

Publication details

Journal / source
Journal of Systems Architecture
Publisher
Elsevier BV
Volume
173
Pages
103721–103721
Publication date
21 January 2026
Language
EN