Exploring Data Layout for Sparse Tensor Times Dense Matrix on GPUs

We focus on definitions that clarify the goals of the approach related to tensor computations we use. A tensor could be viewed as a group of fibers where a fiber is defined by fixing every index in a tensor but one. For example, Figure 2 shows a three-dimensional tensor A, which can be viewed as a group of fibers. When index k is the only unfixed index we get tube fibers, while fixing all other indices but i generates column fibers, and fixing all indices but j gives us row fibers. Another way of viewing a tensor is using the concept of slices, where slices are defined by fixing only one index variable. Figure 2 shows that tensor A can also be viewed as layers of slices. When we fix index k we get frontal slices, fixing index i generates horizontal slices, while fixing index j makes lateral slices. Slices and fibers help us explain semi-sparse tensors [18] which are tensors where each fiber at index (i, j) is dense. Figure 3 shows the difference between dense, semi-sparse, and sparse tensors.

Tensor matrix contraction on mode n, is called n-mode product, which is a product of multiplying tensor \(A \in \mathbb {R}^{I_{1} \times I_{2} \times ... \times I_{N}}\) by dense matrix \(B \in \mathbb {R}^{I_{n} \times R}\) along the nth dimension producing semi-sparse tensor \(C=A \times _{n} B\) . For example, we can represent tensor contraction on the 3^rd mode byHere we fix indices i and j in tensor A, which allows us to iterate over the tube fibers of A shown in the top right of Figure 2 and multiply them with their corresponding elements in dense matrix B.

Both fibers and slices are oriented based on the contraction mode. In this article, there are two main categories of modes we discuss. First, product mode is defined as the mode by which a tensor multiplies a matrix. Second, index mode is defined as all other modes of a tensor excluding the product mode. We need to take into consideration the tensor contraction mode since sparse tensor computations such as multiplication, transpose, or various point-wise operations require iterating over different modes in several orientations during the same computation. The simplest mode generic storage format for sparse tensors is COO format. However, the tradeoff is that COO is less compact than formats like CSF, which takes advantage of mode specificity to reduce indexing overhead cost per nonzero value.

Tensors can be unfolded in any mode. This is done by selecting a subset of the modes of a given tensor A as the rows and the other modes of A as the columns and then mapping the elements of A to corresponding entries of the resulting matrix. For example, in Figure 4, we unfold the sample dense tensor along a product mode, which creates a dense matrix whose rows correspond to the product mode of the original tensor and the columns correspond to all the remaining modes, i.e., index modes. Each column of the matrix becomes a mode-1 fiber of the original tensor.

We aim to simplify an N-dimensional sparse tensor representation for \(N\gt 2\) , reducing its dimensionality to 2 so that high-performance sparse linear algebra libraries such as cuSPARSE [24] or manually-written kernels can be invoked to solve sparse tensor contraction. In this way, we convert sparse tensor contraction into a more common and easier SpMM-like kernel to perform the computation. Matricization reshapes a tensor into an equivalent matrix by arranging all mode-n fibers to be the columns of a matrix. Mathematically, unfolding a tensor is a logical restructuring of the tensor.

Our approach to unfolding a tensor uses a hierarchical tree-like structure to store nonzero values and their corresponding coordinates. The first level corresponds to the first mode of the tensor; for example, the x coordinate of the first nonzero value. The second level corresponds to the second mode of the tensor; for example, the y coordinates of the nonzero values. Lastly, the leaf nodes contain the values.

To represent this hierarchical layout we used the tuple class template from the Thrust library [5] shown in Listing 1 to unfold tensors. Tuples can be instantiated with each argument specifying the type of element in the tuple. Consequently, tuples consist of a mix of tuple, integer, and float. Individual elements of a tuple may be accessed with the get function. For example, a (3,4,2) tensor can be unfolded into ((3,8),1), ((4,6),1), or ((2,12),1) matrices as shown in Figure 4 and using Equation (1).

We use lexicographic ordering [31] in tensor unfolding. Lexicographic ordering is a mapping from an n-dimensional space onto a linear space. One of the benefits of lexicographic ordering is that data close to each other in the folded N-dimensional tensor stays close to each other in the 2-D matrix. This helps with data reuse in the memory hierarchy and reduces total number of data movements required. Similarly, the row-major layout commonly used for CSR matrix formats can be done very efficiently by inducing the natural lexicographic ordering of the dimensions and iterating over the tensor in the natural order [23].

Lexicographical ordering sorts the nonzero values in a linearized manner. Use of tuples allows us to index nonzero values hierarchically. We sort the sparse tensor in order of modes depending on the product mode, and then we compare the indices in lexicographical order for each nonzero value. For example, if we look at Figure 4, it shows the dimension of each slice used to index into the coordinate values for each nonzero value; i.e., for nonzero value e located in the 1^st dimension, the coordinate values are \((0,0)\) .

After unfolding and reorganizing the tensor, we perform the contraction operation, as will be discussed in Section 4. Here we show in Figure 6 that for a single-mode contraction after unfolding, we decompose the product into a series of 1-mode rows in A which, when multiplied by B, produces dense 1-mode rows of C. For matrix multiplication, the index space of rows is 1-dimensional while for tensor contraction the index space of rows in A is n-dimensional. This allows us to treat n-dimensional contraction as a matrix multiplication operation. The result of an SpTM operation is a sparse tensor with one or more dense modes, which is a semi-sparse tensor.

A difference between our approach and PASTA’s mode-generic approach is that mode-generic approaches such as COO require inline search such as a binary search or a tree traversal to lookup the range of indices for nonzero values per row. On the other hand, mode-specific storage formats such as F-COO, CSR, and CSF require the reorganization of tensor A for every different mode of contraction. This makes lookup of the range of indices for nonzero values per row faster compared to mode-generic counterparts as it is pre-computed before the contraction kernel is called.

We describe the target architectures, input datasets, and the approach to data collection.

Figure 8 shows the SpTM kernel and cuSPARSE SpMM kernel performance compared to the current state-of-the-art implementation for sparse tensor contraction available from the PASTA suite [17]. N is the number of multi-vectors or columns in the dense matrix part of the contraction. The number of columns of the dense matrix is set to 16 which is the usual size in tensor decomposition’s lowrank property [18]. On the x-axis, each test case is designated by the tensor name followed by the contraction mode of the tensor.

The performance results in the previous subsections show that none of the implementations perform optimally across all sparse tensor cases. Moreover, the load balancing ratio is not a sufficient predictor of performance of the three kernels. To provide an automated scheme for selecting the best-performing implementation, we used MATLAB’s decision tree classifier to create a decision tree using available measurements. The classifier function call shown in Listing 3 identifies the required features that select the best-performing implementation. To train the decision tree, we used seven features of the sparse tensor, available at run time: number of nonzeros, number of rows multiplied by number of columns of the unfolded tensor, ratio of rows to columns of the unfolded tensor, log of the rows to column ratio, the load imbalance ratio of the unfolded tensor, load imbalance ratio divided by the log of rows to columns ratio, and unfolded tensor density information. We used the best-performing single precision kernel on the A100 GPU as the target for the classifier. We trained the classification decision tree using default values for tree depth control and cross-validated the model using 10-fold cross-validation. The decision tree algorithm selected only three out of the seven features shown in Figure 15 as the necessary features needed to identify the best-performing sparse tensor contraction kernel. The three features are load imbalance ratio of the unfolded tensor, number of nonzeros, and rows to column ratio of the unfolded tensor. Performance of kernels selected by the decision tree verses PASTA are shown in Figure 16.

The decision tree starts by using the load imbalance ratio to identify all the tensor contraction cases where the manually-written SpTM kernel outperforms both PASTA and cuSPARSE. Afterward, the decision tree uses the number of nonzero values of the tensor as an arbiter to identify the sparse tensor cases where cuSPARSE outperforms PASTA, and finally the shape of the unfolded sparse tensor is used to identify the sparse tensor cases where PASTA delivers the best performance. The decision tree mis-classified only three cases out of 55, which gives us an error rate of 5.45%.

Kolda et al. [15] wrote a comprehensive overview of tensor decomposition literature and tensor decomposition methods such as CP or Tucker. They also mathematically define unfolding of tensors. A tensor can be decomposed into the multiplication of matrices or vectors based on the nonzero entries. The function guiding the decomposition minimizes the error between the multiplication of the decomposed matrices or vectors and the nonzero values of the tensor. Tensor decomposition method finds acceptable values for the unknown values in a tensor through multiplying the tensor’s decomposed matrices or vectors.

Ma et al. [21] developed parallel algorithms for sparse tensor times dense matrix multiply (SpTTM) for both CPU and GPU. Their parallel algorithm shows 4.1× speedup on multicore Intel Core i7 and 18.8× speedup on NVIDIA K40c GPU over their sequential SpTTM implementation, respectively. They conclude that different input sparse tensors, multi-vectors, and operating on different modes all influence SpTTM’s performance.

Li et al. [19] also proposed hierarchical variation of the coordinate format called HiCOO that breaks down a sparse tensor into small sparse blocks. This approach reduces the memory requirements to store nonzero values from the tensor. They showcase their approach on the matricized tensor-times-Khatri-Rao product (MTTKRP) algorithm with small thread-count architectures such as CPUs. More recently Li et al. gathered essential sparse tensor operations from various real-world applications for multicore CPU and GPU architectures in a single parallel sparse tensor algorithm Benchmark suite [17]. These tensor operations are critical to the overall performance of tensor analysis algorithms such as tensor decomposition. Recent related work encodes multi-dimensional sparse tensors in a mode-agnostic way called linear format (ALTO) [9]. They developed their approach for MKTTRP and tested it on a CPU architecture.

The Surprisingly ParalleL spArse Tensor Toolkit (SPLATT) [28] is a C library for sparse tensor factorization on CPUs. It uses CSF format to perform SpMTTKRP and SPTTMc with options of either shared memory parallelization using OpenMP or distributed-memory parallelism using MPI.

TACO compiler [14] is a sparse tensor algebra compiler that can compile a basic tensor algebra expression with any number of additions and multiplications and operands can be stored in many different data structures. To date, high-performance GPU implementations have not been a focus of work on TACO.

There are several existing high-performance GPU implementations of SpMM. NVIDIA’s cuSPARSE library [24] is highly optimized for performance on NVIDIA GPUs, with SpMM performance 30–150× faster than CPU-only alternatives. It supports dense, COO, CSR, CSC, and Blocked CSR sparse matrix formats. Yang et al. [30] applied row-splitting and merge-based [22] algorithms to SpMM to efficiently hide global memory latency. Based on the matrix sparsity pattern, one of the algorithms is applied. Huang et al. [12] developed an acceleration method for general-purpose SpMM (GE-SpMM) algorithm for GPUs. Their work is tailored for Graph Neural Networks (GNNs) as it can expedite computations such as graph convolutional networks (GCN) and pooling-based GNNs. Hong et al. [10] solved SpMM using a dynamic load-balancing approach that takes advantage of tiling. They split the sparse matrix into heavy workload and light workload rows. The heavy work load rows are processed by CSR and the light work load rows are processed by doubly compressed sparse row (DCSR). Abu-Sufah et al. [1] designed and implemented two CUDA SpMM kernels targeted at structured matrices (DIA-SpMM) and matrices with uniform row lengths (ELL-SpMM). They evaluated the performance of their kernels with SpMV kernels in NVIDIA CUSP and cuSPARSE libraries, SpMM kernel available in cuSPARSE and SpMV kernels available in Intel MKL library executing on multi core processors. The author’s kernels showed superior performance because they; (i) used the most efficient storage schemes for the test matrices, (ii) multiplied each matrix by a block of vectors instead of one, and (iii) made the most of GPU registers to exploit data reuse in SpMM.

More broadly, Gale et al. [8] developed a library sparse linear algebra library targeted for deep neural networks on GPUs called Sputnik. It uses optimization techniques such as row reordering, memory alignment, and tiling. These optimizations allow it to outperform state-of-the-art baselines such as Adaptive sparse tiling for sparse matrix multiplication (ASpT) [11] and cuSPARSE, on deep learning inference models. Ahmad et al. [2] optimized an important sparse matrix multivector multiplication kernel within the Locally Optimal Block Preconditioned Conjugate Gradient (LOBPCG) solver and showed a 3% performance enhancement over the manually-tuned code for the same algorithm [3]. Using an inspector-executor approach, a data transformation converted the large, symmetric sparse matrix from a compressed sparse row format to a compressed sparse block (CSB) format. This representation was well suited for parallelization of the matrix and its transpose since it permitted storing only the upper triangular portion. In order to reduce the data movement associated with indices of the matrices, a short integer was used as the type for the matrices that pointed to the beginning of each CSB block.