SparseTIR：面向深度学习稀疏编译的可组合抽象
机器学习
2023-02-22 v4 人工智能
编程语言
摘要
稀疏张量正迅速成为现代深度学习负载的关键组件。然而，开发高性能稀疏算子困难且繁琐，现有供应商库无法满足新算子日益增长的需求。稀疏张量编译器简化了算子的开发，但面向深度学习的高效稀疏编译仍具挑战，因为单一稀疏格式无法最大化硬件效率，且一次性编译器难以跟上最新硬件与系统进展。本文中，我们观察到应对这两大挑战的关键在于利用可组合格式与可组合变换。我们提出SparseTIR，一种为深度学习负载提供可组合格式与可组合变换的稀疏张量编译抽象。SparseTIR在这些可组合组件上构建用于性能调优的搜索空间。借助上述改进，SparseTIR在GPU上相对供应商库对单算子取得一致加速：GNN算子1.20-2.34倍，稀疏注意力算子1.05-2.98倍，稀疏卷积算子0.56-7.45倍。SparseTIR还将端到端GNN加速：GraphSAGE训练1.08-1.52倍，RGCN推理4.20-40.18倍。
引用
@article{arxiv.2207.04606,
title = {SparseTIR: Composable Abstractions for Sparse Compilation in Deep Learning},
author = {Zihao Ye and Ruihang Lai and Junru Shao and Tianqi Chen and Luis Ceze},
journal= {arXiv preprint arXiv:2207.04606},
year = {2023}
}
备注
To appear at ASPLOS 2023 (19 pages, 23 figures), source code available at https://github.com/uwsampl/sparsetir, artifact available at https://github.com/uwsampl/sparsetir-artifact