Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Accelerating Sparse MTTKRP for Small Tensor Decomposition on GPU

Sasindu Wijeratne, Rajgopal Kannan, Viktor Prasanna

TL;DR

This paper tackles the bottleneck of sparse MTTKRP (spMTTKRP) in CPD for small tensors on GPUs by introducing a mode-specific tensor layout that stores multiple matricizations as separate tensor copies in GPU global memory, thereby improving data locality and eliminating inter-block communication overhead. It proposes an adaptive load-balancing scheme over GPU SMs, guided by the output-mode index counts and tensor sparsity, to achieve efficient workload distribution. The authors present a hypergraph-based tensor representation and a parallel algorithm that maps partitions to thread blocks, enabling mode-by-mode spMTTKRP with optimized local/global updates. Experimental results on an RTX 3090 show substantial geometric mean speedups (approximately 2.4×, 7.9×, and 8.9×) over state-of-the-art baselines, demonstrating the practical impact for small-tensor CPD in diverse domains. The work highlights how mode-specific data layouts plus adaptive scheduling can significantly reduce memory traffic and synchronization overhead in GPU-based tensor decompositions, with potential extensions to CPU and FPGA platforms.

Abstract

Sparse Matricized Tensor Times Khatri-Rao Product (spMTTKRP) is the bottleneck kernel of sparse tensor decomposition. In tensor decomposition, spMTTKRP is performed iteratively along all the modes of an input tensor. In this work, we propose a mode-specific tensor layout on GPU that uses multiple tensor copies, where each copy is optimized for a specific mode. The proposed tensor layout increases the data locality of external memory accesses and eliminates the intermediate values communicated between the GPU thread blocks and the GPU global memory. We also propose a tensor partitioning scheme to optimally distribute the total computations among GPU streaming multiprocessors based on the sparsity and the dimensions of the input tensor. Our approach achieves a geometric mean speedup of 2.4x, 7.9x, and 8.9x in total execution time compared with the state-of-the-art GPU baselines.

Accelerating Sparse MTTKRP for Small Tensor Decomposition on GPU

TL;DR

This paper tackles the bottleneck of sparse MTTKRP (spMTTKRP) in CPD for small tensors on GPUs by introducing a mode-specific tensor layout that stores multiple matricizations as separate tensor copies in GPU global memory, thereby improving data locality and eliminating inter-block communication overhead. It proposes an adaptive load-balancing scheme over GPU SMs, guided by the output-mode index counts and tensor sparsity, to achieve efficient workload distribution. The authors present a hypergraph-based tensor representation and a parallel algorithm that maps partitions to thread blocks, enabling mode-by-mode spMTTKRP with optimized local/global updates. Experimental results on an RTX 3090 show substantial geometric mean speedups (approximately 2.4×, 7.9×, and 8.9×) over state-of-the-art baselines, demonstrating the practical impact for small-tensor CPD in diverse domains. The work highlights how mode-specific data layouts plus adaptive scheduling can significantly reduce memory traffic and synchronization overhead in GPU-based tensor decompositions, with potential extensions to CPU and FPGA platforms.

Abstract

Sparse Matricized Tensor Times Khatri-Rao Product (spMTTKRP) is the bottleneck kernel of sparse tensor decomposition. In tensor decomposition, spMTTKRP is performed iteratively along all the modes of an input tensor. In this work, we propose a mode-specific tensor layout on GPU that uses multiple tensor copies, where each copy is optimized for a specific mode. The proposed tensor layout increases the data locality of external memory accesses and eliminates the intermediate values communicated between the GPU thread blocks and the GPU global memory. We also propose a tensor partitioning scheme to optimally distribute the total computations among GPU streaming multiprocessors based on the sparsity and the dimensions of the input tensor. Our approach achieves a geometric mean speedup of 2.4x, 7.9x, and 8.9x in total execution time compared with the state-of-the-art GPU baselines.

Paper Structure

This paper contains 30 sections, 2 equations, 5 figures, 3 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Background and Related Work
  3. Introduction to Tensors
  4. Tensor mode
  5. Indices of a nonzero tensor element
  6. Tensor matricization
  7. Small Tensors
  8. Canonical Poliyedic Tensor Decomposition (CPD)
  9. Elementwise computation
  10. Related Work
  11. Tensor Format
  12. Hypergraph Representation
  13. Adaptive Load Balancing Scheme
  14. Equal Distribution of Indices among Tensor Partitions (Load Balancing Scheme 1)
  15. Equal Distribution of Tensor Elements among Tensor Partitions (Load Balancing Scheme 2)
  16. ...and 15 more sections

Figures (5)

  • Figure 1: Elementwise computation 10.1145/3649153.3649187
  • Figure 2: Example hypergraph of a 3-mode tensor
  • Figure 3: Total execution time
  • Figure 4: Impact of the adaptive load balancing scheme
  • Figure 5: Total Memory Consumption