SparseTIR: Composable Abstractions for Sparse Compilation in Deep Learning

TL;DR

SparseTIR introduces a composable abstraction for sparse tensor compilation in deep learning, enabling flexible formats and transformations that significantly improve performance across various operators and end-to-end workloads.

Contribution

It proposes SparseTIR, a novel sparse tensor compilation framework with composable formats and transformations, addressing limitations of single-format and single-shot compilers.

Findings

Achieves 1.20-2.34x speedup on GPU for GNN operators

Attains 1.05-2.98x speedup for sparse attention operators

Provides 4.20-40.18x acceleration for RGCN inference

Abstract

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.