UniSparTa: A Unified Sparse Tensor Program Tuning Framework

Z.M. Wang

doi:10.1109/TCAD.2025.3636442

UniSparTa: A Unified Sparse Tensor Program Tuning Framework

2025-11-24

0

PRE

AI

DOI：10.1109/TCAD.2025.3636442

原文链接

原文求助

摘要

En 中文

稀疏张量计算在深度学习和科学计算中应用广泛。然而，应用层面的多样化稀疏数据及算法特性，结合硬件平台的多样性，对高效稀疏张量程序优化构成显著挑战。手工设计的算子库开发耗时且缺乏可移植性。为解决此问题，我们提出UniSparTa，一个统一的稀疏张量程序调优框架，可自动生成高性能程序。首先，我们提取高性能稀疏张量程序的统一优化原则，并提出一种领域特定语言（DSL）以自动生成高质量设计空间，无需人工干预。其次，通过分析设计空间的普遍分布，我们引入一种自适应搜索策略，结合深度Q网络（DQNs）和模拟退火（SA）。最后，为避免调优过程中实测带来的不可接受的时间开销，我们提出一种基于多模态融合的统一成本模型，以准确预测程序性能。此外，通过利用数据增强和迁移学习，我们实现在不同稀疏数据模式、算法及硬件平台间的低成本性能迁移预测。结果表明，与当前最优算子库MKL、手工优化方案ASpT、张量编译器TVM以及稀疏张量调优框架WACO相比，UniSparTa分别实现了平均加速比1.98倍、2.75倍、6.13倍和1.75倍。此外，UniSparTa显著加速了调优过程。

Keyword:

Autotuning

multimodal fusion

sparse tensor compiler

UniSparTa: A Unified Sparse Tensor Program Tuning Framework

摘要

期刊

机构

引用论文