RASSM: Residue-based Acceleration of Single Sparse Matrix Computation via Adaptive Tiling

Single-Sparse-Matrix Kernels (SSMKs) such as SpMM, SDDMM, SpMV, and SpTS form the backbone of applications such as data analytics, graph processing, finite-element analysis, machine learning (including GNNs and LLMs), etc. This paper introduces Residue-based Acceleration of Single Sparse Matrix Computation via Adaptive Tiling (RASSM), an input-dependent, adaptive 2-dimensional tiling technique for SSMKs. The adaptation leverages the concept of a residue matrix: a data structure that compactly captures the pattern of non-zeros in the sparse matrix. With residues, we show it is possible to make intelligent decisions on adaptive tile sizes, resulting in increased cache occupancy. Residues allow for optimizations across both spatial and temporal locality. RASSM improves data movement and overall performance as compared to prior techniques. For example, using spatial analysis for SpMM on commodity server CPUs, RASSM has 1.30X speedup over MKL, 1.32X over J-Stream, 1.20X over ASpT, 1.11X over CSF-4 uniform-shape, and 1.10X over CSF-4 uniform-occupancy. RASSM with temporal analysis improves this to 1.36X (vs. MKL), 1.38X (vs. J-Stream), 1.26X (vs. ASpT), 1.17X (vs. CSF-4 uniform-shape), and 1.16X (vs. CSF-4 uniform-occupancy).