Heuristic Adaptability to Input Dynamics for SpMM on GPUs

TL;DR

This paper introduces a novel auto-tuning approach for SpMM on GPUs that adapts to input data dynamics, significantly improving performance over static algorithms.

Contribution

It proposes a three-loop model for orthogonal design principles and introduces DA-SpMM, a heuristic GPU kernel that dynamically optimizes SpMM performance based on input data.

Findings

DA-SpMM achieves 1.26x to 1.37x speedup over NVIDIA cuSPARSE.

Up to 5.59x end-to-end speedup in Graph Neural Network applications.

The three-loop model covers existing and new SpMM designs.

Abstract

Sparse Matrix-Matrix Multiplication (SpMM) has served as fundamental components in various domains. Many previous studies exploit GPUs for SpMM acceleration because GPUs provide high bandwidth and parallelism. We point out that a static design does not always improve the performance of SpMM on different input data (e.g., >85\% performance loss with a single algorithm). In this paper, we consider the challenge of input dynamics from a novel auto-tuning perspective, while following issues remain to be solved: (1) Orthogonal design principles considering sparsity. Orthogonal design principles for such a sparse problem should be extracted to form different algorithms, and further used for performance tuning. (2) Nontrivial implementations in the algorithm space. Combining orthogonal design principles to create new algorithms needs to tackle with new challenges like thread race handling. (3) Heuristic adaptability to input dynamics. The heuristic adaptability is required to dynamically optimize code for input dynamics. To tackle these challenges, we first propose a novel three-loop model to extract orthogonal design principles for SpMM on GPUs. The model not only covers previous SpMM designs, but also comes up with new designs absent from previous studies. We propose techniques like conditional reduction to implement algorithms missing in previous studies. We further propose DA-SpMM, a Data-Aware heuristic GPU kernel for SpMM. DA-SpMM adaptively optimizes code considering input dynamics. Extensive experimental results show that, DA-SpMM achieves 1.26x~1.37x speedup compared with the best NVIDIA cuSPARSE algorithm on average, and brings up to 5.59x end-to-end speedup to applications like Graph Neural Networks.