ParamSpMM: Adaptive and Efficient Sparse Matrix-Matrix Multiplication on GPUs for GNNs

TL;DR

ParamSpMM is a novel adaptive GPU-based sparse matrix multiplication method for GNNs, utilizing a new data structure and ML-based decision system to optimize performance across diverse inputs.

Contribution

It introduces ParamSpMM with a new data structure and ML-based configuration predictor, enabling highly adaptive and efficient SpMM computations for GNNs.

Findings

Outperforms Nvidia cuSPARSE with an average speedup of 1.92x

Effectively adapts to diverse input characteristics in GNNs

Enhances GNN training efficiency significantly

Abstract

Fueled by the ability to mine real-world graph data, GNN applications have experienced phenomenal growth. Sparse Matrix-Matrix Multiplication (SpMM) is a critical operator in GNNs. However, existing SpMM designs for GNNs struggle to adapt to diverse input characteristics. In this paper, we first conduct a comprehensive analysis of existing SpMM optimizations, revealing their limitations through statistical and empirical evidence. Based on this analysis, we introduce ParamSpMM, a parametric approach for highly adaptive and efficient SpMM computation in GNNs. It incorporates a new data structure, the Parameterized Compressed Sparse Row (PCSR), to flexibly integrate existing optimization techniques. ParamSpMM enables the configuration of these optimization techniques according to various input characteristics. Furthermore, we complement ParamSpMM with an ML-based SpMM-decider that predicts optimal configurations based on carefully crafted input features. Our evaluations demonstrate that ParamSpMM outperforms Nvidia cuSPARSE with an average speedup of 1.92x, significantly enhancing GNN training efficiency.