Return
UniSparTa: A Unified Sparse Tensor Program Tuning Framework
DOI:10.1109/TCAD.2025.3636442.png)
Abstract
En 中文
Sparse tensor computation is widely used in deep learning and scientific computing. However, diverse sparse data and algorithmic characteristics at the application level, combined with the diversity of hardware platforms, pose significant challenges for efficient sparse tensor program optimization. Manually crafted operator libraries are time-consuming to develop and lack portability. To address this, we propose UniSparTa, a unified sparse tensor program tuning framework that automatically generates high-performance programs. First, we extract unified optimization principles for high-performance sparse tensor programs and propose a domain-specific language (DSL) to automatically generate a high-quality design space without manual intervention. Second, by analyzing the general distribution of the design space, we introduce an adaptive search strategy combining deep $Q$ -networks (DQNs) and simulated annealing (SA). Finally, to avoid the unacceptable time cost of real measurement during tuning, we propose a unified cost model based on multimodal fusion to accurately predict program performance. Furthermore, by leveraging data augmentation and transfer learning, we enable low-cost transfer prediction across different sparse data patterns, algorithms, and hardware platforms. Results show that, compared with the state-of-the-art operator library MKL, the manually optimized scheme ASpT, the tensor compiler TVM, and the sparse tensor tuning framework WACO, UniSparTa achieves average speedups of $1.98\times $ , $2.75\times $ , $6.13\times $ , and $1.75\times $ , respectively. Moreover, UniSparTa significantly accelerates the tuning process.
Keywords:
Autotuning
multimodal fusion
sparse tensor compiler
Journal
I
IF:
2.9
Papers:
586
Citations:
9.6K

