SABLE: Staging Blocked Evaluation of Sparse Matrix Computations

Abstract: Structured sparsity, like regions of non-zero elements in sparse matrices, can offer optimization opportunities often overlooked by existing solutions that treat matrices as entirely dense or sparse. Block-based approaches, such as BCSR, partially address this issue by choosing between fixed-size blocks which results in wasted computation on zero elements. On the other hand, variable-sized blocks introduce overheads due to variable loop bounds unknown at compile time. We present SABLE, a novel staging framework that achieves the best of both approaches by generating region-specific code tailored for variable-sized blocks. SABLE partitions the matrix to identify profitable blocks and specializes generated code for vectorization. We evaluate SABLE on the SpMV kernel using the SuiteSparse collection. SABLE achieves a geomean of $1.07$, $2.73$ and $1.9$ speedup over the state of the art systems: Intel MKL, CSR5 and Partially-Strided Codelets, respectively, single threaded and even more when parallelized.