Swift: High-Performance Sparse Tensor Contraction for Scientific Applications

TL;DR

Swift is proposed, a novel algorithm for sparse tensor contraction that replaces the costly sorting with more efficient grouping, utilizes better data structures to represent tensors, and employs more memory-friendly hash table implementation.

Abstract

In scientific fields such as quantum computing, physics, chemistry, and machine learning, high dimensional data are typically represented using sparse tensors. Tensor contraction is a popular operation on tensors to exploit meaning or alter the input tensors. Tensor contraction is, however, computationally expensive and grows quadratically with the number of elements. For this reason, specialized algorithms have been created to only operate on the nonzero elements. Current sparse tensor contraction algorithms utilize sub-optimal data structures that perform unnecessary computations which increase execution time and the overall time complexity. We propose Swift, a novel algorithm for sparse tensor contraction that replaces the costly sorting with more efficient grouping, utilizes better data structures to represent tensors, and employs more memory-friendly hash table implementation. Swift is evaluated against the state-of-the-art sparse tensor contraction algorithm, demonstrating up to 20x speedup in various test cases and being able to handle imbalanced input tensors significantly better.

Figures (13)

• Figure 1: Example of converting standard tensor format to COO format in memory. i and j represent the rows and columns of the value denoted by v.
• Figure 2: Example of tensor contraction with 2-dimensional input tensors. In this case the operation is a matrix multiplication where the second mode of $\mathcal{X}$ is contracted with the first mode of $\mathcal{Y}$. As shown the dimensionality of the problem can be increased by adding more indices to match the dimension of the application.
• Figure 3: Visualization of Input Processing for Sparta Algorithm. (a) The X entries are sorted by the free modes and (b) the Y entries are hashed into a table where the key is the contracting mode and the value is the free mode with the corresponding entry value. Here, lower case letters (y1-y4) indicate indices of coordinates, and upper case letters ($Y_1$) indicates the length of the mode, which is 10 in this example.
• Figure 4: Visualization of Contraction and Accumulation process for Sparta Algorithm. If the contracting modes of an X entry match with the contracting modes of a Y entry, the output entry coordinate is the concatenation of free modes and the value is the multiplication of the individual values.
• Figure 5: Example of hash table search requiring unnecessary pointer chasing. If the requested contracting mode was 5, the search would cross contracting modes 3 and 2 which are a part of the same slot.