Highly Efficient Load-Balanced Dataflow for SpGEMMs on Systolic Arrays

To enhance the efficiency of sparse neural network models, compression methods are commonly employed to store the non-zero elements in a sparse storage format. Sparse General Matrix Multiplication (SpGEMM) is a critical computation in deep neural networks. However, when utilizing systolic arrays for SpGEMM computations, a challenge arises due to the irregular flow of compressed, non-zero element activation data. This irregularity leads to varying lengths of activation data streams entering the systolic array per batch, potentially resulting in the underutilization of processing units. Our research focuses on repackaging compressed data streams using hardware-software co-design to minimize software pre-processing time. We also package unevenly sized sparse matrix rows post-compression into multiple groups of activation data streams with approximately equal lengths. This approach aims to evenly distribute the workload across the fixed output of the systolic array, thereby improving the utilization rate of Processing Elements (PE). Our evaluation demonstrates that our method achieves a 2.01x acceleration compared to uncompressed sparse data streams and a 3.63x average acceleration relative to TPUs. Furthermore, compared to the state-of-the-art SpGEMM accelerator SADD, our approach achieves an average of 2.01x acceleration.