  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.
