MergePath-SpMM: Parallel Sparse Matrix-Matrix Algorithm for Graph Neural Network Acceleration

Mohsin Shan, Deniz Gurevin, Jared Nye, Caiwen Ding, Omer Khan

Abstract: Graph neural networks have seen tremendous adoption to perform complex predictive analytics on massive and unstructured real-world graphs. The trend in hardware accelerator designs has identified significant challenges with harnessing graph locality and workload imbalance due to ultra-sparse and irregular matrix computations at a massively parallel scale. This paper addresses the load imbalance challenge and identifies that state-of-the-art either introduces complex specialized hardware to auto-tune for load-balanced execution at runtime or relies on software-only approaches that exploit parallelism. We propose a novel software-only load-balancing sparse matrix-matrix (SpMM) algorithm that unlocks fine-grain parallelism while maintaining controlled need-based targeted synchronizations to achieve robust performance scaling. The MergePath-SpMM algorithm achieves superior performance using commercial offthe-shelf GPU processors when compared to state-of-the-art hardware accelerators and software-only implementations.