SparseX-PIM: Sparse Linear Algebra on Processing-in-Memory

Abstract

Sparse linear algebra is a fundamental workload in scientific computing and machine learning, but its performance is severely constrained by memory bandwidth. Processing-in-memory (PIM) has shown great potential for sparse linear algebra, yet its efficiency is heavily bottlenecked by irregular computation and memory access patterns. Through a systematic analysis of sparse kernels, we identify three key challenges: inter-bank communication caused by poor data locality, workload imbalance caused by irregular sparsity, and computation stalls caused by complex task dependencies. Hardware-software co-design is key to resolving these tightly coupled challenges. In this work, we propose SparseX-PIM, a hardware-software co-designed 3D PIM architecture that unifies sparse matrix multiplication, factorization, and linear system solving, and enables the first PIM-based direct solver for sparse linear systems. On the software side, we develop an instruction set and an application programming interface (API) for code generation, together with a dependency-aware data layout and task scheduling scheme enabled by a PIM-oriented compiler, to improve load balance and reduce communication overhead. On the hardware side, we combine coarse-grained asynchronous execution with fine-grained pipelining, complemented by hierarchical scratchpads and grouped content-addressable memory (CAM) acceleration to effectively mask computation stalls and accelerate index matching. Evaluation results demonstrate that, compared with CPU and GPU baselines (PARDISO and cuDSS), SparseX-PIM achieves average speedups of 135.7× and 146.2× in factorization, and 5060.5× and 1987.0× in solving, alongside massive energy savings. Furthermore, SparseX-PIM delivers superior performance and energy efficiency compared with highly specialized state-of-the-art PIM accelerators across a broad spectrum of sparse kernels.