Machine Learning-Based SpGEMM Partitioning and Load Balancing for Multi-GPU Platforms

General sparse matrix-matrix multiplication (SpGEMM) is a main computational kernel in algorithms such as algebraic multigrid and multi-source breadth-first search. Despite modern computing systems frequently incorporating nodes with four or eight GPUs, most prior SpGEMM research has concentrated on single-GPU optimizations, leaving multi-GPU implementations largely underexplored. The irregular distribution of three sparse matrices presents challenges for optimizing SpGEMM on multi-GPU platforms, particularly in designing efficient task partitioning and scheduling algorithms. This study makes two significant contributions: First, we propose a machine learning-based performance prediction model to accurately predict the execution time of blocked SpGEMM operations. Second, we leverage this model to develop an innovative partitioning and scheduling algorithm for multi-GPU platforms, which achieves superior load balancing across devices. Experimental results show that our approach significantly outperforms traditional one-dimensional static partitioning methods in computational efficiency.