A Unified Iteration Space Transformation Framework for Sparse and Dense Tensor Algebra
Abstract
We address the problem of optimizing mixed sparse and dense tensor algebra in a compiler. We show that standard loop transformations, such as stripmining, tiling, collapsing, parallelization and vectorization, can be applied to irregular loops over sparse iteration spaces. We also show how these transformations can be applied to the contiguous value arrays of sparse tensor data structures, which we call their position space, to unlock loadbalanced tiling and parallelism. We have prototyped these concepts in the opensource TACO system, where they are exposed as a scheduling API similar to the Halide domainspecific language for dense computations. Using this scheduling API, we show how to optimize mixed sparse/dense tensor algebra expressions, how to generate loadbalanced code by scheduling sparse tensor algebra in position space, and how to generate sparse tensor algebra GPU code. Our evaluation shows that our transformations let us generate good code that is competitive with many handoptimized implementations from the literature.
 Publication:

arXiv eprints
 Pub Date:
 December 2019
 arXiv:
 arXiv:2001.00532
 Bibcode:
 2020arXiv200100532S
 Keywords:

 Computer Science  Mathematical Software;
 Computer Science  Programming Languages