A Performance Portable Matrix Free Dense MTTKRP in GenTen

TL;DR

This work extends GenTen with portable, matrix-free dense MTTKRP variants that avoid explicitly forming the Khatri–Rao product, achieving memory complexity of $O(R\sum_{n=1}^d I_n)$ versus $O(R\prod_{n\neq k} I_n)$ for matrix-based methods. Among the proposed algorithms, MTTKRP-ELEM, MTTKRP-SLICE, and MTTKRP-TILE offer progressively cache-aware, tile-based strategies, with MTTKRP-TILE achieving substantial single-GPU gains and enabling rank-$2000$ decompositions on a single NVIDIA H100 compared to requiring six GPUs for GEMM-based MTTKRP. Analytical performance models (including $T_0$, $T_{\infty}$, and $T_{0,LM}$) and a tile-size heuristic guide algorithm selection and parameter tuning, validated by CPU and GPU experiments on large tensors (e.g., tear2tear1 and island). The results demonstrate meaningful memory and speedups (up to 50x memory reduction and 11x speedups) and highlight the practical viability of matrix-free MTTKRP for high-rank CP decompositions, while also outlining limitations and opportunities for improved caching and tiling strategies.

Abstract

We extend the GenTen tensor decomposition package by introducing an accelerated dense matricized tensor times Khatri-Rao product (MTTKRP), the workhorse kernel for canonical polyadic (CP) tensor decompositions, that is portable and performant on modern CPU and GPU architectures. In contrast to the state-of-the-art matrix multiply based MTTKRP kernels used by Tensor Toolbox, TensorLy, etc., that explicitly form Khatri-Rao matrices, we develop a matrix-free element-wise parallelization approach whose memory cost grows with the rank R like the sum of the tensor shape O(R(n+m+k)), compared to matrix-based methods whose memory cost grows like the product of the tensor shape O(R(mnk)). For the largest problem we study, a rank 2000 MTTKRP, the smaller growth rate yields a matrix-free memory cost of just 2% of the matrix-based methods, a 50x improvement. In practice, the reduced memory impact means our matrix-free MTTKRP can compute a rank 2000 tensor decomposition on a single NVIDIA H100 instead of six H100s using a matrix-based MTTKRP. We also compare our optimized matrix-free MTTKRP to baseline matrix-free implementations on different devices, showing a 3x single-device speedup on an Intel 8480+ CPU and an 11x speedup on a H100 GPU. In addition to numerical results, we provide fine grained performance models for an ideal multi-level cache machine, compare analytical performance predictions to empirical results, and provide a motivated heuristic selection for selecting an algorithmic hyperparameter.

Figures (5)

• Figure 1: Slicing of a 3-way tensor $\boldsymbol{\mathscr{Y}}\in \mathbb{R}^{I_1 \times I_2 \times I_3}$ along each mode.
• Figure 2: A tensor $\boldsymbol{\mathscr{Y}}\in \mathbb{R}^{I_1 \times I_2 \times I_3}$ decomposed into mode-$2$ slices $\boldsymbol{\mathscr{S}}_n^{(2)}$. Each slice in then assigned to a Kokkos league of two teams $p_1, p_2$, where each team computes a tile Hadamard product and atomically reduces their local contributions to $\bm{\mathbf{G}} {^{\intercal}} {^{\space}} ^{(k)} (n,j)$. In this instructive example we have one league per slice, but in practice we allow for more than one league per slice. However, the map between teams and tiles is always bijective.
• Figure 3: Worst case memory usage for the GEMM and TILE MTTKRP variants for the tearing mode tensor $\boldsymbol{\mathscr{A}}$ and the island coalescence tensor $\boldsymbol{\mathscr{B}}$ for different CP-ranks. The plot is annotated with lines showing the minimum number of NVIDIA H100 GPUs required incur the storage costs.
• Figure 4: Tearing mode tensor $\boldsymbol{\mathscr{A}}$ average (over all modes) MTTKRP performance (higher is better) for different algorithmic variants on different execution spaces.
• Figure 5: Time (lower is better) for CP-ALS on the island coalecence tensor $\boldsymbol{\mathscr{B}}$ with MTTKRP-TILE algorithm on a single H100 compared with MTTKRP-GEMM on one H100 per process. MTTKRP time is denoted with hatches, other CP-ALS operations (i.e., inner-product, Cholesky solve, etc.) by a grey fill, tensor redistribution time with a purple fill, and MTTKRP communication time with an orange fill.