一句话:提出 SparseTIR,用可组合稀疏格式与变换构建搜索空间,提升深度学习稀疏算子编译性能。
稀疏张量在现代深度学习中日益重要,但高性能稀疏算子开发困难,现有厂商库难以满足新算子需求,而稀疏张量编译器又受限于单一格式与一次性编译。SparseTIR 提出可组合的稀疏格式与可组合变换作为编译抽象,并在此之上构建性能调优搜索空间,以自动探索不同格式和变换组合。评估在 GPU 上覆盖 GNN、稀疏注意力和稀疏卷积算子,相比厂商库取得 1.20–2.34x、1.05–2.98x 和 0.56–7.45x 加速;端到端 GraphSAGE 训练加速 1.08–1.52x,RGCN 推理加速 4.20–40.18x。这些结果显示可组合抽象能有效适应不同稀疏模式与硬件效率需求。
Sparse tensors are rapidly becoming critical components of modern deep learning workloads. However, developing high-performance sparse operators can be difficult and tedious, and existing vendor libraries cannot satisfy the escalating demands from new operators. Sparse tensor compilers simplify the development of operators, but efficient sparse compilation for deep learning remains challenging because a single sparse format cannot maximize hardware efficiency, and single-shot compilers cannot keep up with latest hardware and system advances. In this paper, we observe that the key to addressing both these challenges is to leverage composable formats and composable transformations. We propose SparseTIR, a sparse tensor compilation abstraction that offers composable formats and composable transformations for deep learning workloads. SparseTIR constructs a search space over these composable components for performance tuning. With these improvements, SparseTIR obtains consistent performance speedups vs vendor libraries on GPUs for single operators: 1.20-2.34x for GNN operators, 1.05-2.98x for sparse attention operators, and 0.56-7.45x for sparse convolution operators. SparseTIR also accelerates end-to-end GNNs by 1.08-1.52x for GraphSAGE training, and 4.20-40.18x for RGCN inference.
@inproceedings{paperbot10088,
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},
booktitle = {ASPLOS},
year = {2023},
doi = {10.1145/3582016.3582047}
}