PRISM: Processing-In-Memory Sparse MTTKRP for Tensor Decomposition Acceleration

2026-05-28

AI Summary

This work addresses the memory-bandwidth-bound nature of the sparse MTTKRP (spMTTKRP) operation in sparse tensor decomposition, which suffers from low efficiency on general-purpose processors. To overcome this limitation, the study introduces processing-in-memory (PIM) technology to accelerate spMTTKRP for the first time, leveraging the UPMEM PIM architecture. The authors devise an efficient sparse tensor tiling strategy, a customized numerical format, and specialized compute kernels tailored to the PIM platform, along with a CPU-PIM heterogeneous collaboration mechanism. Experimental results demonstrate that the pure PIM implementation achieves a 2.37× speedup over the state-of-the-art CPU baseline, while the heterogeneous approach yields a 2.64× improvement, both exhibiting substantially higher resource utilization efficiency compared to conventional CPU and GPU implementations.

Abstract

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.

Problem

• spMTTKRP
• tensor decomposition
• memory-bound
• sparse tensors
• Processing-In-Memory

Innovation

• Processing-In-Memory
• spMTTKRP
• Tensor Decomposition
• UPMEM
• Heterogeneous Computing