Abstract

This paper extends prior work on sparse-tensor algebra compilers to generate asymptotically efficient code when sparse tensors are indexed by compound affine subscript expressions. Our technique enables compiler support for a wide range of sparse computations, including sparse convolutions and pooling that are widely used in ML and graphics applications. We propose an approach that gradually rewrites compound subscript expressions to simple subscript expressions with loops that exploit the sparsity pattern of the input sparse tensor. As a result, the time complexity of the generated kernel is bounded by the number of stored elements and not by the shape of the tensor. Our approach seamlessly incorporates into existing frameworks and is compatible with all recent advances in compilers for sparse computations, including the flexibility to efficiently handle arbitrary combinations on different sparse tensor formats. The implementation of our algorithm is open-sourced and upstreamed to the MLIR sparse compiler. Experimental results show that the generated code can outperform a dense implementation at about 80% sparsity and achieves 19.5x speedup when compared with the state-of-the-art compiler-based method at 99.9% sparsity.

Article

pdf

BibTeX

@article{peimingliu2024,
  title={Compiler Support for Sparse Tensor Convolutions},
  author={Peiming Liu and Alexander J Root and Anlun Xu and Yinying Li and Fredrik Kjolstad and Aart J.C Bik},
  journal={Proceedings of the ACM on Programming Languages},
  volume={8},
  issue={OOPSLA},
  year={2024},
  month={October}
}