arrow
Return

UniSparTa: A Unified Sparse Tensor Program Tuning Framework

delete2025-11-24
delete0
PRE
AI
Z
Z.M. Wang
宫磊 (Lei Gong)
X
Xiangjun Qu
C
Cheng Tang
W
Wenqi Lou
T
Teng Wang
Q
Qianyu Cheng
X
Xianglan Chen
王超 (Chao Wang)
周学海 (Xuehai Zhou)
DOI:10.1109/TCAD.2025.3636442delete
deleteOriginal
deleteOriginal request for help
deleteShare
deleteSave
Abstract

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
IEEE Transactions on Computer-Aided Design of Integrated Circuits and Systems
IF:
2.9
Papers:
586
Citations:
9.6K

Organization

U
University of Science and Technology of China
Scholars:
1.6W
Papers: 5.6K
Citations: 11.3W