Compiler Support for Sparse Tensor Computations in MLIR

Bixia Zheng
Fredrik Kjolstad
Nicolas Vasilache
Tatiana Shpeisman
ACM Transactions on Architecture and Code Optimization(2022) (to appear)
Google Scholar


Sparse tensors arise in problems in science, engineering, machine learning, and data analytics. Programs that operate on such tensors can exploit sparsity to reduce storage requirements and computational time. Developing and maintaining sparse software by hand, however, is a complex and error-prone task. Therefore, we propose to treat sparsity as a property, not a tedious implementation detail, and let a sparse compiler generate sparse code automatically from a sparsity-agnostic definition of the computation. This paper discusses the integration of this idea into MLIR.