BCB-SpTC: An Efficient Sparse High-Dimensional Tensor Contraction Employing Tensor Core Acceleration

• Published in: IEEE transactions on parallel and distributed systems 2024-12, Vol.35 (12), p.2435-2448
• Main Authors: Hu, Rong, Wang, Haotian, Yang, Wangdong, Ouyang, Renqiu, Li, Keqin, Li, Kenli
• Format: Article
• Language: English
• Subjects: Arithmetic; Bit optimization; Computer science; GPU; Graphics processing units; Indexes; Memory management; Optimization; Registers; Single instruction multiple data; Sparse matrices; sparse tensor contraction (SpTC); sparse tensor format; tensor core; Tensors
• ISSN: 1045-9219; 1558-2183
• DOI: 10.1109/TPDS.2024.3477746

Sparse tensor contraction (SpTC) is an important operator in tensor networks, which tends to generate a large amount of sparse high-dimensional data, placing higher demands on the computational performance and storage bandwidth of the processor. Using GPUs with powerful arithmetic characteristics is a reliable choice for accelerating SpTC, however, the high dimensionality and sparsity of tensor makes GPU-accelerated SpTC operators suffer from the difficulties of low computational intensity and high memory consumption. The recent introduction of Tensor Core Units (TCUs) on GPUs brings even more powerful arithmetic, which exacerbates the memory wall problem. To cope with the challenges, this paper proposes a new BCB format that linearizes the indices of multidimensional blocks to reduce block index accesses and uses a bitmap to store the distribution of non-zero elements in a block to reduce the storage overhead. A parallel blocking algorithm of BCB-SpTC is designed to divide the binary linear indices into free and contracted indexes to improve the pairing overhead of computational tasks. Then based on the characteristic computation method of TCUs, the proprietary filling method of TCUs is designed to overcome the inefficiency of parallel computation of sparse data on TCUs. Finally, experimental results on the A100 dataset show that BCB-SpTC improves the acceleration ratio by 1.1\times 1.1× to 21.3\times 21.3× over the existing SpTC GPU method.