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.

提出了 SparseTIR，一种为深度学习工作负载提供可组合格式和可组合转换的稀疏张量编译抽象。SparseTIR 对于单一操作符可以比供应商库获得 1.20-2.34 倍到 7.45 倍的性能加速。