SparseTIR is a tensor-level compiler for sparse/irregular operators in Deep Learning. The design goal of SparseTIR is to provide a general programming abstraction that can cover both sparse and irregular (e.g. Ragged Tensors) workloads in Deep Learning including Graph Neural Networks, Sparse Transformers, Sparse Convolutions, Network Pruning, etc. while generating high-performance code on heterogeneous hardware.
The key innovation of SparseTIR is composability:
- Format Composability: Decompose the computation on a single format to computation on hybrid formats.
- Transformation Composability: SparseTIR adopts multi-stage IR design on top of TVM's TensorIR, user can compose program transformations at different stages to bring together sparsity-aware optimizations such as format searching and low-level optimizations such as vectorization and tensorization.
Check out the Documentation site for installation instructions, tutorials, examples, and more. The Blitz Introduction to SparseTIR is a great place to get you familiar with SparseTIR's format annotations and compilation flow. You can check the tutorials on optimizing SpMM and RGMS to understand how to use composable formats and composable transformations to optimize sparse operators.
This repo is still experimental and documentations are under construction, the API will not be stable and we will frequently refactor the codebase. However, we believe making it public will benefit researchers and engineers in this field. Feel free to create an issue if you run into any problems or have any suggestions on SparseTIR. We plan to contribute SparseTIR to Apache TVM project as an optional module, and we will keep maintaining this repo until the core functionalities have been upstreamed. The codebase is mainly developed and maintained by Ruihang Lai and Zihao Ye, some early commits are rebased and squashed when syncing with Apache TVM upstream.
SparseTIR was built upon TensorIR in Apache TVM, we thank all the generous help from the TVM community. We also learned a lot from the TACO project: the format annotations and the concept of position space and coordinate space come from TACO. We thank the TACO team for their foundation work in sparse compilers.
You can also check out the following cool projects if you are interested in our work:
- MLIR-Sparse MLIR's sparse tensor dialect.
- SparTA SparTA is an abstraction for model sparsity in Deep Learning.