SABLE: Staging Blocked Evaluation of Sparse Matrix Computations (2407.00829v2)

Abstract: Sparse Matrices found in the real world often have some structure in their distribution of dense elements. While existing techniques specialize the generated code for the structure of matrices, their generality misses optimization opportunities. We propose a system that -- if the sparse matrix is stored in a blocked storage format -- can adapt its code generation strategy depending on the structure of the sparse matrix. Our system SABLE performs a specified computation over every element of {\em mostly} dense blocks instead of avoiding computing any sparse element and achieving regularity in generated code while having special treatment for hyper-sparse blocks (ie, blocks with very few dense elements). SABLE is extensible, providing a block iterator for users to express any computation over these non-empty blocks. We demonstrate that our approach can significantly accelerate SpMV and SpMM operations, surpassing the performance of state-of-the-art systems like Partially-Strided Codelets and Sparse Register Tiling.