PckGNN: Optimizing Aggregation Operators with Packing Strategies in Graph Neural Networks

Graph Neural Network (GNN) is one of the most prominent machine learning models. It involves a substantial amount of graph-based aggregation operators, which can be abstracted as sparse matrix computation kernels. Due to the irregularity of the graph adjacency matrix, the aggregation has long been the performance bottleneck of GNN. According to our measurements, existing GNN implementations fall short of achieving optimal performance due to their insufficient consideration of load balancing and data locality. To bridge these performance gaps, we propose PckGNN, which aims to accelerate GNN aggregation operators on GPUs with packing strategies. PckGNN categorizes graph matrices into two types based on different sparsity levels and conducts two packing strategies respectively. For sparser matrices, Neighbor Packing enhances load balancing through a moderate-grained non-zero grouping approach. For denser matrices, Bin Packing exposes more potentials of cache data reuse by bin partitioning, non-zero extracting, format converting and two-level scheduling. Experimental results on SpMM and SDDMM show that PckGNN achieves speedups of 1.46x ∼ 6.14x over existing implementations. When applied GNN inference of three typical models, it achieves speedups of more than 1.29x over the state-of-the-art frameworks.