Compilation of sparse array programming models
Rawn Henry, Olivia Hsu, Rohan Yadav, Stephen Chou, Kunle Olukotun, Saman P. Amarasinghe, Fredrik Kjolstad
Abstract
This paper shows how to compile sparse array programming languages. A sparse array programming language is an array programming language that supports element-wise application, reduction, and broadcasting of arbitrary functions over dense and sparse arrays with any fill value. Such a language has great expressive power and can express sparse and dense linear and tensor algebra, functions over images, exclusion and inclusion filters, and even graph algorithms.
Our compiler strategy generalizes prior work in the literature on sparse tensor algebra compilation to support any function applied to sparse arrays, instead of only addition and multiplication. To achieve this, we generalize the notion of sparse iteration spaces beyond intersections and unions. These iteration spaces are automatically derived by considering how algebraic properties annotated onto functions interact with the fill values of the arrays. We then show how to compile these iteration spaces to efficient code.
When compared with two widely-used Python sparse array packages, our evaluation shows that we generate built-in sparse array library features with a performance of 1.4× to 53.7× when measured against PyData/Sparse for user-defined functions and between 0.98× and 5.53× when measured against SciPy/Sparse for sparse array slicing. Our technique outperforms PyData/Sparse by 6.58× to 70.3×, and (where applicable) performs between 0.96× and 28.9× that of a dense NumPy implementation, on end-to-end sparse array applications. We also implement graph linear algebra kernels in our system with a performance of between 0.56× and 3.50× compared to that of the hand-optimized SuiteSparse:GraphBLAS library.
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Cited by top-tier papers17
- SparseTIR: Composable Abstractions for Sparse Compilation in Deep LearningZihao Ye, Ruihang Lai, Junru Shao, Tianqi Chen et al.ASPLOS 2023 · 86 citations
- The Sparse Abstract MachineOlivia Hsu, Maxwell Strange, Ritvik Sharma, Jaeyeon Won et al.ASPLOS 2023 · 37 citations
- Autoscheduling for sparse tensor algebra with an asymptotic cost modelWillow Ahrens, Fredrik Kjolstad, Saman P. AmarasinghePLDI 2022 · 30 citations
- Indexed Streams: A Formal Intermediate Representation for Fused Contraction ProgramsScott Kovach, Praneeth Kolichala, Tiancheng Gu, Fredrik KjolstadPLDI 2023 · 11 citations
- Compilation of dynamic sparse tensor algebraStephen Chou, Saman P. AmarasingheOOPSLA 2022 · 8 citations
Builds on2
- A sparse iteration space transformation framework for sparse tensor algebraRyan Senanayake, Changwan Hong, Ziheng Wang, Amalee Wilson et al.OOPSLA 2020 · 51 citations
- Automatic generation of efficient sparse tensor format conversion routinesStephen Chou, Fredrik Kjolstad, Saman P. AmarasinghePLDI 2020 · 26 citations
Related papers
- Compilation of Shape Operators on Sparse ArraysAlexander J. Root, Bobby Yan, Peiming Liu, Christophe Gyurgyik et al.OOPSLA 2024 · 3 citations
- Compilation of Modular and General Sparse WorkspacesGenghan Zhang, Olivia Hsu, Fredrik KjolstadPLDI 2024 · 6 citations
- SpDISTAL: Compiling Distributed Sparse Tensor ComputationsRohan Yadav, Alex Aiken, Fredrik KjolstadSC 2022 · 7 citations
- Compiler Support for Sparse Tensor ConvolutionsPeiming Liu, Alexander J. Root, Anlun Xu, Yinying Li et al.OOPSLA 2024 · 5 citations
- Mosaic: An Interoperable Compiler for Tensor AlgebraManya Bansal, Olivia Hsu, Kunle Olukotun, Fredrik KjolstadPLDI 2023 · 16 citations
