Dynamic Sparse Tensor Algebra Compilation

TL;DR

This paper introduces a compiler technique for generating efficient code to perform tensor algebra operations on dynamic sparse tensors with evolving structures, enabling faster computations compared to existing frameworks.

Contribution

It presents a new language for specifying recursive data structures and a compiler approach to generate optimized code for dynamic sparse tensor algebra.

Findings

Generated code for PageRank kernel is 1.05× faster than Aspen.

Outperforms PAM by 7.40× in dynamic matrix addition.

Efficiently handles dynamic data structures like linked lists and trees.

Abstract

This paper shows how to generate efficient tensor algebra code that compute on dynamic sparse tensors, which have sparsity structures that evolve over time. We propose a language for precisely specifying recursive, pointer-based data structures, and we show how this language can express a wide range of dynamic data structures that support efficient modification, such as linked lists, binary search trees, and B-trees. We then describe how, given high-level specifications of such data structures, a compiler can generate code to efficiently iterate over and compute with dynamic sparse tensors that are stored in the aforementioned data structures. Furthermore, we define an abstract interface that captures how nonzeros can be inserted into dynamic data structures, and we show how this abstraction guides a compiler to emit efficient code that store the results of sparse tensor algebra computations in dynamic data structures. We evaluate our technique and find that it generates efficient dynamic sparse tensor algebra kernels. Code that our technique emits to compute the main kernel of the PageRank algorithm is 1.05$\times$ as fast as Aspen, a state-of-the-art dynamic graph processing framework. Furthermore, our technique outperforms PAM, a parallel ordered (key-value) maps library, by 7.40$\times$ when used to implement element-wise addition of a dynamic sparse matrix to a static sparse matrix.