A Novel Parallel Algorithm for Sparse Tensor Matrix Chain Multiplication via TCU-Acceleration
Haotian Wang, Wangdong Yang, 栄徳 宍戸, Renqiu Ouyang, Kenli Li, Keqin Li
Abstract
Analysis of multi-dimensional data, especially tensor decomposition, which extracts latent information, is becoming considerably popular. Although multi-dimensional sparse data is typically processed on multi-core processors, developing highly optimized GPU-based <b>Sp</b> arse <b>T</b> ensor <b>M</b> atrix <b>C</b> hain <b>M</b> ultiplication (SpTMCM) is challenging. The purpose of this paper is to investigate a novel approach named SpTMCM and to explore the discovery of SpTMCM coupled with the emerging computing core, Tensor Core Unit (TCU). In contrast to prior work, the proposed novel approach enables a uniform storage format and optimization approach for SpTMCM. We design a hybrid tensor format based on multi-dimensional tiling that divides the tensor depending on the tile threshold to address the inefficient memory accesses caused by the irregular nonzero distribution of the sparse tensor. Further, we develop a TCU-based tensor parallel algorithm with our novel approach to increase the memory bandwidth. Compared to state-of-the-art works, our method achieves <inline-formula><tex-math notation="LaTeX">$1.16\sim 24.12\times$</tex-math></inline-formula> speedup for SpMTTKRP and <inline-formula><tex-math notation="LaTeX">$5.07\sim 7.15\times$</tex-math></inline-formula> speedup for SpTTMChain across NVIDIA A100 GPU on a range of real-world sparse tensors.