Introduction

Treating sparsity as a property instead of a (tedious) implementation detail lets the compiler's sparsifier1 generate sparse code automatically. The concept was pioneered for linear algebra by [Bik96] in MT1 and formalized to tensor algebra by [Kjolstad17,Kjolstad20] in the Sparse Tensor Algebra Compiler (TACO) project.

The SparseTensor dialect supports all the attributes, types, operations, and passes that are required to make sparse-tensor types first-class citizens within the MLIR compiler infrastructure. The dialect forms a bridge between high-level operations on sparse-tensor types and lower-level operations on the actual sparse-storage schemes consisting of positions, coordinates, and values. Lower-level support may consist of fully generated code or may be provided by means of a small sparse runtime support library.

The MLIR implementation [Biketal22] closely follows the 'sparse iteration theory' that forms the foundation of TACO. The compiler has a rewriting rule for each tensor expression in the Linalg dialect (MLIR's tensor index notation): the sparsity of each level of a tensor is indicated by the level-types (e.g., dense, compressed, singleton) and a specification of the order on the levels (see [Chou18] for an in-depth discussions and possible extensions to these level-types). For each tensor expression in the input, the compiler first constructs a topologically sorted iteration graph that reflects the ordering of coordinates with respect to the levels of each tensor; this ensures that all tensors are visited in the natural level-coordinate order. Next, iteration lattices are constructed for every index in topological order. Each iteration-lattice point consists of a conjunction of tensor coordinates and a tensor (sub)expression that needs to be evaluated for that conjunction. Within the lattice, iteration points are ordered according to the way coordinates are exhausted. As such, these iteration lattices drive actual sparse-code generation, a relatively straightforward one-to-one mapping from iteration lattices to combinations of for-loops, while-loops, and if-statements. Sparse-tensor outputs that materialize uninitialized are handled with either direct insertions, if all parallel loops are outermost, or insertions that indirectly go through a 1-dimensional access pattern expansion (a.k.a. workspace), where feasible [Gustavson72,Bik96,Kjolstad19].

  • [Bik96] Aart J.C. Bik. Compiler Support for Sparse Matrix Computations. PhD thesis, Leiden University, May 1996.

  • [Biketal22] Aart J.C. Bik, Penporn Koanantakool, Tatiana Shpeisman, Nicolas Vasilache, Bixia Zheng, and Fredrik Kjolstad. Compiler Support for Sparse Tensor Computations in MLIR. ACM Transactions on Architecture and Code Optimization, June, 2022. See: https://dl.acm.org/doi/10.1145/3544559

  • [Chou18] Stephen Chou, Fredrik Berg Kjolstad, and Saman Amarasinghe. Format Abstraction for Sparse Tensor Algebra Compilers. Proceedings of the ACM on Programming Languages, October 2018.

  • [Chou20] Stephen Chou, Fredrik Berg Kjolstad, and Saman Amarasinghe. Automatic Generation of Efficient Sparse Tensor Format Conversion Routines. Proceedings of the 41st ACM SIGPLAN Conference on Programming Language Design and Implementation, June, 2020.

  • [Gustavson72] Fred G. Gustavson. Some basic techniques for solving sparse systems of linear equations. In Sparse Matrices and Their Applications, pages 41–52. Plenum Press, New York, 1972.

  • [Kjolstad17] Fredrik Berg Kjolstad, Shoaib Ashraf Kamil, Stephen Chou, David Lugato, and Saman Amarasinghe. The Tensor Algebra Compiler. Proceedings of the ACM on Programming Languages, October 2017.

  • [Kjolstad19] Fredrik Berg Kjolstad, Peter Ahrens, Shoaib Ashraf Kamil, and Saman Amarasinghe. Tensor Algebra Compilation with Workspaces, Proceedings of the IEEE/ACM International Symposium on Code Generation and Optimization, 2019.

  • [Kjolstad20] Fredrik Berg Kjolstad. Sparse Tensor Algebra Compilation. PhD thesis, MIT, February, 2020.

Would you like to contribute? Here are some good first issues.

  1. Please note that we now prefer the term 'sparsifier' over the also commonly used 'sparse compiler' terminology to refer to such a pass to make it clear that the sparsifier pass is not a separate compiler, but should be an integral part of any compiler pipeline that is built with the MLIR compiler infrastructure. ↩