Efficient Sparse Processing-in-Memory Architecture (ESPIM) for Machine Learning Inference (2404.04708v1)

Abstract: Emerging ML models (e.g., transformers) involve memory pin bandwidth-bound matrix-vector (MV) computation in inference. By avoiding pin crossings, processing in memory (PIM) can improve performance and energy for pin-bound workloads, as evidenced by recent commercial efforts in (digital) PIM. Sparse models can improve performance and energy of inference without losing much accuracy. However, unstructured sparse inference injects the key challenges of uncertainty, irregularity, and load imbalance into a dense PIM's operation across all the banks. The dense PIM reads the matrix cells from each bank and broadcasts the vector elements to all the banks exploiting DRAM organization. To address these challenges efficiently, we propose ESPIM which makes four contributions: (1) Because matrix sparsity increases the vector broadcast bandwidth demand per matrix column-read, ESPIM employs a fine-grained interleaving of the matrix cells so that each vector broadcast is shared among multiple rows in each bank, cutting the bandwidth demand. (2) ESPIM mostly avoids on-chip control's area and energy despite sparsity's uncertainties by exploiting the observation that the sparsity is data-dependent but static and known before inference. Accordingly, ESPIM employs static data-dependent scheduling (SDDS) (3) ESPIM decouples the matrix cell values and their indices, placing the indices ahead of the values to enable prefetching of the vector elements. We extend SDDS for performance and correctness with the decoupled prefetching. (4) Finally, we simplify the switch required to select the vector elements that match the matrix cells. We extend SDDS to improve performance by reducing conflicts in the simplified switch. In our simulations, ESPIM achieves 2x average (up to 4.2x) speedup over and 34% average (up to 63%) lower energy than Newton while incurring under 5% area.