Swift: High-Performance Sparse Tensor Contraction for Scientific Applications

Tensor contraction can be expressed using the form shown in Equation 1, which can be expanded to Equation 2.

The tensor contraction operation takes input tensor $\mathcal{X}$, input tensor $\mathcal{Y}$, the list of contracting $\mathcal{X}$ modes (top), and the list of contracting $\mathcal{Y}$ modes (bottom), as exemplified in Equations 3 & 4. If a certain mode of tensor $\mathcal{X}$ is contracted with another mode of tensor $\mathcal{Y}$, both of these modes must have the same length, similar to how for matrix multiplication the number of columns in the first matrix must match the number of rows in the second matrix. The tensor contraction operation focuses on finding entries in $\mathcal{X}$ and $\mathcal{Y}$ that share contracting indices, multiplying the values together, and adding that result to the output tensor $\mathcal{Z}$. This can be seen in Equations 2 and 4 where a matching contracting index is denoted as c. All instances of matching contracting indices are summed with output coordinates which share the same free modes of $\mathcal{X}$ and $\mathcal{Y}$, hence the summation sign.

Tensor contraction is a generalization of matrix multiplication as shown in Figure 2. Note that the element at index $(i,j)$ in the output tensor $\mathcal{Z}$ is the dot product of the $i$th row of tensor $\mathcal{X}$ with the $j$th column of tensor $\mathcal{Y}$. The resulting element in tensor $\mathcal{Z}$ is the concatenation of the free modes (first mode from $\mathcal{X}$ and second mode from $\mathcal{Y}$, which correspond to indices 0 and 1, respectively), which mirrors Equation 2. The example shown in this figure can be applied to tensors with higher dimensionality which leads to the general concept of tensor contraction. Sparse tensor contraction occurs when one or both of the input tensors of a contraction operation are sparse [26].

Hash tables are used repeatedly in both related work and our proposed algorithm. After the element-wise contraction stage, which will be explained in section 3, each of the output entries must be combined with entries that share the same coordinate. Hash tables are a common solution for this problem as they operate on key-value pairs. In this instance the key would be coordinate of the entry and the value would be the value of the entry. The key is mapped to a table slot using a hash function, and the key-value pair is placed at that slot. The hash table achieves $O(1)$ search time derived from the fact that this hashing process described is theoretically invariant of how many elements are currently in the table [18]. A major design decision with this data structure is how to handle the scenario of two different keys being hashed to the same slot. The two most prominent solutions are probing and chaining.

Acceleration of sparse tensor contraction has been approached using hardware accelerators and high performance software algorithms. Hardware accelerators [13, 14, 15, 28, 29] almost always outperform software algorithms, but there is extreme commitment in the creation of hardware. In the realm of high performance software, Sparta is a state-of-the-art sparse tensor contraction algorithm which operates on COO formatted tensors. This algorithm addresses the main concerns with sparse contraction which are: irregular memory access patterns, unknown output tensor size, and high amounts of intermediate data. Machine learning libraries such as PyTorch and TensorFlow offer support for sparse tensor contraction but they do not have comparable contraction time to Sparta. The Sparta algorithm can be broken down into the Input Processing stage (Figure 3), Contraction & Accumulation (Figure 4), and Writeback, which are explained below.

After Contraction & Accumulation, because the algorithm was implemented using parallelism, each thread has a sub-collection of entries which belong to the output tensor. During Writeback, the thread-local collections of entries are iterated over and saved to a global data structure which represents the output tensor.

Time complexity of Sparta can be calculated by adding the time complexities of each individual stage in order to give an overview of how well the algorithm scales. Equation 3.1 shows the time complexity of the Sparta Algorithm where $nnz_{X}$ denotes the number of nonzero entries in tensor $\mathcal{X}$, and $cmode_{Yavg}$ denotes the average number of entries associated with each contracting mode of $\mathcal{Y}$.

The first term expresses the time complexity of the input processing stage. The input processing stage has time complexity $\mathcal{O}(nlog(n))$ with respect to, $nnz_{X}$, the number of nonzero entries in tensor $\mathcal{X}$ because its entries were sorted using quicksort. In contrast, the input processing stage has linear time complexity with respect to, $nnz_{Y}$, the number of nonzero entries in tensor $\mathcal{Y}$ because each was hashed into a table. The second term is the time complexity of the contraction and accumulation phases which must iterate over each nonzero $\mathcal{X}$ entry, and for each $\mathcal{X}$ entry the average mode length of $\mathcal{Y}$ must be iterated over to get the resulting output entry which is added to a hash table. Finally, the third term expresses the time complexity of the Writeback stage which is equal to the number of output entries because each entry must be saved to the output tensor.

The time complexity of the Contraction & Accumulation stages cannot be improved upon as it is necessary to iterate over each $\mathcal{X}$ entry and the corresponding $\mathcal{Y}$ entries to generate the correct output entries. Similarly, the time complexity of the Writeback stage cannot be improved upon as each entry must be exported. There is, however, potential to improve the time complexity of the input processing stage. It is also, of course, possible to improve the overall performance of each stage, even if the time complexity remains unchanged. Accessing the same number of elements sequentially in memory generally leads to performance improvement compared to random access, and two linear algorithms can have vastly different execution times based on the number of steps per element even if that number is constant.

Sparta represents the tensor $\mathcal{Y}$ as a hash table (HtY) which can be used to fetch a run of entries that share the same contracting mode ($C^{Y}$). Although the creation of the hash table can be done in $O(n)$ time and is generally fast, there is room for improvement when it comes to the time it takes to fetch elements during the contraction stage of the algorithm.

One aspect that makes this approach in Sparta slow compared to an array is the fact that chaining hash tables (the variant used in Sparta) leads to pointer chasing when traversing a linked list. This causes irregular memory accesses which harms cache performance. The second aspect that makes this approach slower is that during the hashing process the hash function commonly maps different keys to the same slot. This means if a contracting mode must be searched in the table, there will be undesired elements to filter through to access the desired entries. An example of this can be seen in Figure 5. In this case the requested contracting mode is 5, shown in top left, but two other nodes must be traversed to access that element because 3 entries were mapped to that slot of the table.

Combining the above two aspects, in a chaining implementation of a hash table, an additional pointer to the next entry must be fetched for each undesired contracting mode. This consumes a lot of time as there is a high probability a new line of memory must be fetched and placed into cache numerous times during this process.

Similarly, chaining hash tables are also used for the storage of output entries in the accumulation stage where an alternative data structure would have much better memory access time and cache-efficiency.

As previously stated, the input processing stage is composed mainly of the sorting of the $\mathcal{X}$ entries and the hashing of the $\mathcal{Y}$ entries. Our experiments show that the sorting of $\mathcal{X}$ vastly overshadows the hashing of $\mathcal{Y}$ which is reflected in the time complexity equation of Sparta. During the input processing stage of Sparta, tensor $\mathcal{X}$ is sorted by the free modes of its coordinates in order to get matching free modes next to each other. Sorting also has the additional property of putting these entries in ascending order based on the free modes. If there is a way to remove the reliance on the ascending order while still maintaining the correctness of contraction, much time can be saved from not needing to sort $\mathcal{X}$. This is achieved by using our proposed methods described in the following section.

If entries of tensor $\mathcal{X}$ are fetched in chunks where free modes are shared amongst entries, fetching is rendered more efficient due to locality. Therefore, the input processing stage should arrange the entries of tensor $\mathcal{X}$ such that entries which share free modes are placed contiguously in an array. This allows for a block of entries that share the same free mode to be fetched in constant time utilizing pointers to the start of runs. During the input processing stage, Swift utilizes grouping by free mode instead of sorting by free mode to manipulate the order of the input entries. Grouping by free mode simply means placing entries that share free mode next to each other without ensuring that entry free modes are in any particular order. While this might still take $\mathcal{O}({n}log(n))$ in general cases, a linear time version is possible here by assigning a small metric (numerical proxy) to represent each unique free mode and applying a less stringent sorting algorithm to this metric. In particular, we employ a relaxed counting sort in this work, which is an algorithm that can achieve linear time complexity when the sorting metric does not exceed the number of entries being sorted so it is ideal for this scenario.

A similar process to the one described in Section 4.1.3 above will be performed on the entries of tensor $\mathcal{Y}$ but with respect to the contracting modes instead of the free modes. The access time for a sequence of $\mathcal{Y}$ entries that share a desired contracting mode can be found in O(1) time. This can be done using the table of pointers to divisions in contracting modes which is given by Counting Sort results in an array of sorted elements with a corresponding vector of pointers to divisions of contracting modes, the access time for a sequence of $\mathcal{Y}$ entries that share a desired contracting mode can be found in O(1) time using the table of pointers to divisions in contracting modes.

When it comes to implementation, this data structure is significantly faster than the hash tables used in Sparta for three reasons. First, with a hash table it is very likely that two blocks of contracting modes will be hashed to the same slot. As a consequence, in order to get the desired run of contracting indices, the program must unnecessarily filter through the extraneous contracting modes that were hashed to the same slot. Second, the entries are stored contiguously in memory which greatly improves cache efficiency. Third, the performance of a hash table can vary based on its implementation. In the case of a chaining hash table, excessive time may be spent following pointers, while a probing hash table aiming for peak execution time is expected to consume around 1.3 times more memory compared to an array. For HtY, Sparta implements a chaining hash table whereas we use the faster probing hash table at the cost of slightly increased memory.

Chaining tables are required for Sparta’s accumulation stage because the $NNZ$ of the output tensor is unknown. However, the accumulation stage could be improved by using a probing hash table instead of a chaining hash table to store output entries from the contraction stage. Due to cache locality, it is much faster to iterate through a table in array format than to perform pointer chasing as would be done with a chaining table. The challenge lies in the fact that a probing hash table requires a predefined size, which can be difficult to estimate before the contraction process begins.

We address this challenge with the following clever method. During the grouping stage of tensor $\mathcal{Y}$, a vector is available that records the number of entries corresponding to each contracting index. One could go through every $\mathcal{X}$ entry that shares the same free mode and accumulate the number of $\mathcal{Y}$ entries associated with each of these operations. This can be used to assign an upper bound for the number of slots in the probing hash table such that dynamically allocated memory is not required. Using Figure 6 as an example, the contents of the array could represent the large number representation of the contracting modes of tensor $\mathcal{Y}$, thus the counts array will store the count of each contracting mode set.

As an example, given a set of 4 $\mathcal{X}$ entries that share the same free modes, the associated contracting modes could be $[1,2,3,3]$. Thus, the total number of $\mathcal{Y}$ entries to be fetched (and therefore output entries to be allocated), would be equal to $3+0+2+2=7$. For this sub-operation, a maximum of 9 slots is needed for the accumulation probing hash table and there is no need for dynamic allocation of additional memory.

The approach described has the following time complexity, which represents an improvement over that of Sparta. The input processing stage is linear with respect to the number of nonzero entries in tensor $\mathcal{X}$ and all remaining stages have unaltered time complexity to Sparta. The reduced time complexity of the Swift algorithm leads to a reduction in execution time while maintaining the correctness of the contraction results ²²2An even faster time complexity could have been achieved by simply leaving the entries of tensor $\mathcal{X}$ alone during the input processing stage. This would, however, force the accumulation stage to be performed after combining intermediate entries to a global array, which would cause a significant decrease in practical performance..

To validate the potential of our proposed method, Swift, we evaluated the execution time and peak memory usage for sparse tensor contraction operations on input tensors with varying $NNZ$ and order. For fair comparison, both Sparta and Swift were implemented in C and compiled using gcc with the optimization option -O3 and employed OpenMP for parallelism. Both programs used 8 threads maximum and were executed on a 11th Gen Intel(R) Core(TM) i7-11375H which runs @3.3GHz with 32GB of RAM. For execution time and peak memory measurements, two randomly generated tensor datasets were created with difference in mode length to observe the effect of change in sparsity. Additionally, we examined sparse tensor contraction operations between two tensors with a significant difference in $NNZ$ to assess how the algorithms handle load imbalance.

Figure 8 shows speedup in contraction time over the state-of-the-art as a function of the number of nonzero elements in the input tensors. The specifications of the tensors used in the operations can be seen in Table 1, Set 1, where tensors labeled $\mathcal{A}$ are of order 5 and tensors labeled $\mathcal{B}$ are of order 4. Testing different order combinations was done to show that both designs are flexible to shape. For these measurements, both $\mathcal{A}$ & $\mathcal{A}$ have all modes of equal length 50. Figure 9 similarly shows speedup as a function of $NNZ$ using the tensors described in Table 1, Set 2, where modes are of equal length 100. Tensors have randomly generated entries and values to ensure that there is no general structure that would bias the results.

For Set 1, the minimum speedup from employing Swift sits at a solid 1.8x for large $NNZ$ with small orders to a maximum of a blistering 7.7x observed speedup for small $NNZ$ with higher orders. The speedup is much more extreme for experiments involving tensors from Set 2 which are orders of magnitude more sparse. The fastest speedup is 19.6x for small sparse tensors, and the lowest speedup is still a substantial factor of 2.3x for more dense operations.

This shows the effect sparsity can have on speedup, regardless if the $NNZ$ of the input tensors is constant. Note the formulas for time complexity of Sparta, Equation 3.1, and Swift, Equation 4.4, both include the term $cmode_{Yavg}$ which represents the average number of Y entries assigned the same contracting modes. As sparsity increases, $cmode_{Yavg}$ will decrease which means a proportionately more amount of time is spent fetching entries from memory compared to processing them. Sparta and Swift perform the same amount of time processing entries, but Swift has much faster rate of fetching entries because of its cache-friendly storage of elements inside an array compared to a chaining hash-table. Keeping $NNZ$ and mode length constant, decreasing the order by 1 will have profound increase in density, thus self-contractions between $\mathcal{B}$ experience a lesser speedup. For reference,the sparsity for tensors from Set 1 range from 0.016% to 8% while tensors from Set 2 exhibit sparsity in the range of 0.0005% to 0.5%.

There is a visible decrease in speedup as the number of nonzero entries increases, despite having similar time complexities. As the number of output entries increases, most of the computation time will be concentrated in the Accumulation and Writeback stages, which are comparable between the two algorithms. Consequently, less time will be spent on the Input Processing and Contraction stages, where the primary differences between the algorithms lie. There are, however, many scientific applications that utilize tensors which fall into this range of $NNZ$. In such cases, Swift would provide significant speedup.

Figure 10 shows contractions between tensors found in scientific applications. Tensor specifications for these experiments can be found in Table 2. These examples further illustrate the strength of Swift when tensors are extremely sparse. When tensors have low density, a greater proportion of execution time will be devoted towards fetching entries compared to performing the contraction operation on the entries. Tensor LBNL has a significantly lower density ($3.7\times 10^{-15}$) which means modes being fetched from memory will have very few elements and thus fetching must happen more often. Swift utilizes an array to store elements which has much faster access time compared to a chaining hash table. Tensors with such few elements, such as Enron, have little room for improvement because there are so few operations being performed and thus a reduced time complexity will have minimal impact.

For each of the experiments performed in Section 5.2, the peak heap memory consumption was also measured using memusage on the executable files. The ratio of memory consumption from the proposed algorithm divided by that of Sparta is shown in Figures 11 and 12, for tensor contraction combinations found in Table 1 Set 1 and 2, respectively. There is a stable increase in memory consumption as the $NNZ$ for input tensors increases, where tensors contractions between lower order tensors result in a higher ratio of memory being used. Factors which cause Swift to require more memory are 1) the added array needed to store the counts of each contracting mode set for tensor $\mathcal{Y}$ during the contraction phase, and 2) the use of upper bounds for the accumulation hash table compared to the bare minimum being used by a chaining hash table in Sparta’s accumulation stage. The peak memory allocations to the stack from both programs is constant with respect to the number of nonzero entries, both allocating a negligible 1.9KB peak to the stack.

The third set of tensor contractions in Table 1 contains imbalanced operations in which the left tensor has vastly more non-zero elements compared to the right. Although both operations result in the same output tensor, this section compares the difference in contraction time when the order of the tensors in the operation are swapped. This can be important in the case of unknown tensor sizes such as the instance these algorithms are used in a chain of contractions where the output tensor of the previous layer has an unknown number of non-zero elements. The measurements are displayed in Figure 13. As can be seen, the operation $\mathcal{B}\otimes\mathcal{A}$ is similar for both algorithms, but the operation $\mathcal{A}\otimes\mathcal{B}$ is very imbalanced and Sparta has much worse performance. This imbalance should be due to the imbalanced time complexity and the difference in tensor representations based on the order of contraction for Sparta, versus Swift which represents both tensors in the same fashion and has a symmetric time complexity. For these extremely imbalanced cases, the Sparta execution time increases by a factor of 26x on average when the order of the operations is swapped, and Swift execution time increases by a factor of 0.96x when the order of operations is swapped.