PRISM: Processing-In-Memory Sparse MTTKRP for Tensor Decomposition Acceleration
For machine learning practitioners using tensor decomposition, this work addresses the memory bottleneck of spMTTKRP, but the speedup is limited to specific workloads and the approach shows performance degradation on some workloads.
PRISM is the first approach to accelerate the memory-bound spMTTKRP operation in tensor decomposition using Processing-In-Memory (PIM) technology, achieving up to 2.64x speedup over state-of-the-art CPU implementations.
Sparse tensors are the most used representation of sparse multidimensional data. Operations that decompose them, selecting their most important features while reducing their dimension, have become prevalent procedures in machine learning. One of the most used tensor decomposition algorithms is the Alternating Least Squares Canonical Polyadic Decomposition (CP-ALS), where the most time-consuming operation is the Sparse Matricized Tensor Times Khatri-Rao Product (spMTTKRP). This operation is strongly memory-bound, making it hard to implement efficiently on general-purpose processors. This work proposes PRISM, the first approach to tackle this operation using Processing-In-Memory (PIM) technology. We extensively characterize different partitioning strategies, number formats, and kernel optimizations that efficiently adapt this operation to UPMEM PIM, which is further boosted by heterogeneous collaboration with the CPU. The experimental results show that the proposed PIM-based and heterogeneous approaches achieve up to 2.37x and 2.64x speedup compared to state-of-the-art CPU implementations, respectively. However, the UPMEM distributed memory system can significantly hinder performance on certain workloads. Nonetheless, the efficiency of resource consumption for this approach, measured by peak performance fraction usage, is significantly higher than for both CPU and GPU.