Groot: Graph-Centric Row Reordering with Tree for Sparse Matrix Multiplications on Tensor Cores

Sparse matrix multiplications are essential in scientific computing and machine learning applications. Recent researches offload sparse operations, such as sparse matrix-matrix multiplication (SpMM) and sampled dense-dense matrix multiplication (SDDMM), on Tensor Cores (TCs) for improved performance. However, their performance is often limited by the matrix's inherent sparsity and irregularity. In this paper, we find row reordering can potentially improve sparse operations on TCs, but existing reordering techniques exhibit limitations that hinder their effectiveness. To address the issues, we propose Groot, a graph-centric row reordering algorithm with tree. Groot aims to minimize row differences across the matrix, which is proved to be a NP-hard problem. To approximate the optimal solution, Groot firstly captures the local structure of the sparse matrix by constructing a k-nearest neighbor graph, where rows are represented as nodes. Then, it extracts a minimum spanning tree from the constructed graph for global structure optimization. Lastly, Groot traverses the extracted tree to obtain the final ordering. We evaluate Groot using real-world datasets in comparison with state-of-the-art reordering algorithms. Our results show that Groot significantly enhances the computational intensity of SpMM and SDDMM on TCs, delivering the average speedups of 1.8× and 2.0×, respectively. Furthermore, the performance gains extend broadly to sparse computations on CUDA cores and GNN systems.