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

TL;DR

The paper addresses the memory pin bandwidth bottleneck in MV inference for modern ML models by extending digital processing-in-memory to handle unstructured sparsity efficiently. ESPIM builds on Newton with four coordinated innovations—fine-grained interleaving, static data-dependent scheduling, decoupled prefetching of indices and values, and a simplified, conflict-aware switch—to preserve the benefits of vector broadcasts while accommodating sparse matrices. Empirical results show ESPIM achieving about 2x average speedup over Newton (up to 4.2x) and around 34% lower energy (up to 63%) with under 5% area increase, while enabling robust performance across dense and sparse configurations and various bank counts. The work demonstrates that sparse PIM architectures can effectively accelerate ML inference at the edge, offering substantial gains over non-PIM systems and existing PIM approaches, particularly for transformer-style workloads with moderate sparsity.

Abstract

Emerging machine learning (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.

Figures (14)

• Figure 1: Newton's datapath for one bank
• Figure 2: Coarse-grained interleaving in dense matrix for one bank
• Figure 3: Newton's operation across all banks
• Figure 4: ESPIM unoptimized sparse-only datapath ($U$ is an execution unit comprising a MAC and other components)
• Figure 5: Fine-grained interleaving in compressed sparse matrix for one bank