IMAX-SpMM: An Energy-Efficient and Dataflow-Optimized SpMM Kernel on CGLA for GNNs

Abstract

Sparse Matrix-Matrix Multiplication (SpMM) plays a central role in Graph Neural Networks (GNNs), yet it remains a significant performance bottleneck. Operations on highly sparse adjacency matrices suffer from poor cache locality and irregular memory access, which limits both performance and energy efficiency. Consequently, dataflow-centric optimizations that reduce redundant transfers and improve on-chip reuse are required. To address this issue, we propose IMAX-SpMM, a dataflow-optimized SpMM kernel designed primarily to improve energy efficiency by reducing DMA traffic and increasing on-chip data reuse. We also report speedups as supporting evidence that these optimizations shorten execution time, thereby contributing to the energy-efficiency gains. At the software level, IMAX-SpMM minimizes padding for each dataset. At the hardware level, it exploits the fast compilation capability of a coarse-grained linear array (CGLA) to perform dataset-specific dynamic optimization. These optimizations enable IMAX-SpMM to achieve a data reuse rate of 60.1% for Direct Memory Access (DMA)-loaded data, significantly improving memory transfer efficiency. Additionally, eliminating inter-thread data dependencies among the logical threads executed on each PE increases pipeline density and shortens computation time. Experimental evaluation on Graph Convolutional Network (GCN) demonstrates up to 3.07× speedup over an i9-10940X and 1.64× over an RTX3090. In terms of energy efficiency, the proposed method achieves up to 246.60× improvement compared to an i9-10940X and 25.60× compared to an RTX3090.