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

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$ to $21.3\times$ over the existing SpTC GPU method.

I. Introduction

Sparse tensor contraction (SpTC) has drawn a lot of interest from researchers in a variety of domains, such as quantum physics [1], [2], [3] and machine learning [4], [5]. For example, tensor contraction is a major operation in tensor networks and tensor decomposition algorithms [6]. The tensor contraction layer is a network structure with the ability to extract deep features in deep learning [7]. In these applications, tensors usually represent large datasets or high-dimensional physical models and need to be reshaped and reduced during computation. Tensor contraction can be traced back to matrix multiplication in linear algebra. Unlike matrix multiplication, tensor contraction can operate on tensors of any dimension. In tensor computation, the dimensions of tensors can sometimes be very large, so effectively reducing the dimensionality of tensors is crucial.