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.
Thu 24 OctDisplayed time zone: Pacific Time (US & Canada) change
10:40 - 12:20 | Compilers and Optimisation 1OOPSLA 2024 at San Gabriel Chair(s): Emery D. Berger University of Massachusetts Amherst | ||
10:40 20mTalk | Compilation of Shape Operators on Sparse Arrays OOPSLA 2024 Alexander J Root Stanford University, Bobby Yan Stanford University, Peiming Liu Google Inc, Christophe Gyurgyik Stanford University, Aart Bik Google, Inc., Fredrik Kjolstad Stanford University DOI Pre-print | ||
11:00 20mTalk | Compiler Support for Sparse Tensor Convolutions OOPSLA 2024 Peiming Liu Google Inc, Alexander J Root Stanford University, Anlun Xu Google, Yinying Li Google, Fredrik Kjolstad Stanford University, Aart Bik Google, Inc. DOI | ||
11:20 20mTalk | Compiling Recurrences over Dense and Sparse Arrays OOPSLA 2024 Shiv Sundram Stanford University, Muhammad Usman Tariq Stanford University, Fredrik Kjolstad Stanford University DOI | ||
11:40 20mTalk | Fully Verified Instruction Scheduling OOPSLA 2024 Ziteng Yang Georgia Institute of Technology, Jun Shirako Georgia Institute of Technology, Vivek Sarkar Georgia Institute of Technology DOI Pre-print | ||
12:00 20mTalk | Homeostasis: Design and Implementation of a Self-Stabilizing Compiler (TOPLAS) OOPSLA 2024 Link to publication |