A framework for computing on sparse tensors based on operator properties

Abstract

Tensor operations have been traditionally limited to addition and multiplication operations. For operations of sparse tensors, theses semantics were extended to handle operating on zero values. However, there are many operators with a rich semantics of operator properties that can be used in dense and sparse tensor computations. This work addresses the problem of generating code for computing on a mix of sparse and dense tensors based on the properties of the operators on those tensors. I introduce the concept of a fill value to each tensor so that the data can be sparse on non-zeros. I show how to reason about the operator properties, along with the fill values of the input tensors in order to construct an IR describing how to iterate over these tensors. I show how we can take advantage of the operator properties to perform useful optimizations for both iterating over tensors and performing reductions. Lastly, I show how a user can leverage set notation to directly describe to a compiler how it should iterate over sparse tensors. The ideas discussed in this work have been prototyped in the open-source TACO system. The API used makes operator properties and tensor fill values have to be explicitly provided by the user. However, it makes the TACO system much more flexible. I show how the primitives exposed in this work allows one to efficiently perform several graph algorithms by drawing on the literature about GraphBLAS. In the evaluation section, we benchmark this system against the SuiteSparse implementation of GraphBLAS on a variety of graph algorithms to demonstrate its performance.